27b 935/47

# Conflicts:
#	extra/gemm/amd_flash_attention.py
#	tinygrad/codegen/__init__.py
#	tinygrad/runtime/ops_cpu.py
#	tinygrad/uop/ops.py

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

extra/gemm/amd_flash_attention.py

@@ -1,203 +1,33 @@
-from tinygrad import Tensor, UOp, getenv
-from tinygrad.uop.ops import AxisType, KernelInfo, Ops
-from tinygrad.dtype import AddrSpace, dtypes
-from tinygrad.helpers import DEBUG, GlobalCounters, Context
-import math
-
-BLOCK_M, BLOCK_N = 64, 64
-WARP_SIZE = 32
-WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
-WAVES_M, WAVES_N = 4, 1
-LANES_PER_WAVE_M, LANES_PER_WAVE_N = 2, 16
-WMMA_ACC = WMMA_M // LANES_PER_WAVE_M
-THREADS_PER_BLOCK = WARP_SIZE * WAVES_M * WAVES_N
-LDS_PAD = 4  # pad LDS rows to reduce bank conflicts
-
-WMMA_ARG = (WMMA_M, WMMA_N, WMMA_K), 'AMD', 32
-LOG2E = math.log2(math.e)
-
-def warp_shfl_xor(val, offset, lane):
-  """Read val from lane ^ offset using ds_bpermute."""
-  idx = ((lane ^ offset) * 4).cast(dtypes.int)
-  if val.op is Ops.INDEX and val.addrspace == AddrSpace.REG: val = val.load()
-  return UOp(Ops.CUSTOM, dtypes.float, (idx, val),
-             arg="__builtin_bit_cast(float, __builtin_amdgcn_ds_bpermute({0}, __builtin_bit_cast(int, {1})))")
-
-def warp_reduce_max(val, lane):
-  """Tree reduce MAX across LANES_PER_WAVE_N=16 lanes."""
-  for offset in [8, 4, 2, 1]:
-    val = UOp(Ops.MAX, dtypes.float, (val, warp_shfl_xor(val, offset, lane)))
-  return val
-
-def warp_reduce_sum(val, lane):
-  """Tree reduce SUM across LANES_PER_WAVE_N=16 lanes."""
-  for offset in [8, 4, 2, 1]:
-    val = val + warp_shfl_xor(val, offset, lane)
-  return val
-
-def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
-  # inputs are (B*H, N, D)
-  BH, N, D = q.shape
-  assert N % BLOCK_M == 0 and N % BLOCK_N == 0, f"N={N} must be divisible by BLOCK_M={BLOCK_M} and BLOCK_N={BLOCK_N}"
-  assert D % WMMA_K == 0 and D % LANES_PER_WAVE_N == 0, f"D={D} must be divisible by WMMA_K={WMMA_K} and LANES_PER_WAVE_N={LANES_PER_WAVE_N}"
-  assert BLOCK_M % (WAVES_M * WMMA_M) == 0 and BLOCK_N % LANES_PER_WAVE_N == 0
-  TM = BLOCK_M // (WAVES_M * LANES_PER_WAVE_M)
-  TN = BLOCK_N // (WAVES_N * LANES_PER_WAVE_N)
-  TD = D // (WAVES_N * LANES_PER_WAVE_N)
-  SCALE = 1.0 / math.sqrt(D)
-
-  block_bh = UOp.range(BH, 0, AxisType.GLOBAL)
-  block_m = UOp.range(N // BLOCK_M, 1, AxisType.GLOBAL)
-
-  q = q.reshape(BH, N//BLOCK_M, BLOCK_M, D)[block_bh, block_m]
-  k = k.reshape(BH, N//BLOCK_N, BLOCK_N, D)[block_bh]
-  v = v.reshape(BH, N//BLOCK_N, BLOCK_N, D)[block_bh]
-  o = o.reshape(BH, N//BLOCK_M, BLOCK_M, D)[block_bh, block_m]
-
-  wave_m = UOp.range(WAVES_M, 2, AxisType.LOCAL)
-  wave_n = UOp.range(WAVES_N, 3, AxisType.LOCAL)
-  lane = UOp.range(WARP_SIZE, -1, AxisType.WARP)
-  tid = (wave_m * WAVES_N + wave_n) * WARP_SIZE + lane
-  lane_m = lane // LANES_PER_WAVE_N
-  lane_n = lane % LANES_PER_WAVE_N
-
-  # LDS allocation: slot 0 = Q then P (shared), slot 1 = K then V
-  # TODO: the memory planner should be able to find this reuse
-  ELEMS_PER_THREAD = BLOCK_M * D // THREADS_PER_BLOCK
-  QP_lds = UOp.placeholder((BLOCK_M, D + LDS_PAD), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL)
-  KV_lds = UOp.placeholder((BLOCK_N, D + LDS_PAD), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL)[:, :D]
-
-  # register state
-  acc = UOp.placeholder((TM, TD), dtypes.float, slot=2, addrspace=AddrSpace.REG)
-  m_i = UOp.placeholder((TM,), dtypes.float, slot=3, addrspace=AddrSpace.REG)
-  l_i = UOp.placeholder((TM,), dtypes.float, slot=4, addrspace=AddrSpace.REG)
-  acc = acc.after(acc.store(acc.const_like(0)))
-  m_i = m_i.after(m_i.store(m_i.const_like(-math.inf)))
-  l_i = l_i.after(l_i.store(l_i.const_like(0)))
-
-  # ====== KV tile loop ======
-  n_tile = UOp.range(N // BLOCK_N, 100, AxisType.REDUCE)
-
-  # load Q + K into LDS (Q reloaded each iteration since P overwrites slot 0)
-  Q_lds = QP_lds[:, :D]
-  Q_store = Q_lds.after(n_tile).reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
-    q.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
-  K_store = KV_lds.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
-    k[n_tile].reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
-  # NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
-  Q_lds = Q_lds.after(UOp.group(Q_store, K_store))
-  KV_lds_k = KV_lds.after(UOp.group(Q_store, K_store))
-
-  # -- S = Q @ K^T via WMMA (re-init each n_tile) --
-  S_reg = UOp.placeholder((TM, TN), dtypes.float, slot=6, addrspace=AddrSpace.REG)
-  S_reg = S_reg.after(S_reg.after(n_tile).store(S_reg.const_like(0)))
-  k_qk = UOp.range(D // WMMA_K, 101, AxisType.REDUCE)
-  tm1 = UOp.range(TM // WMMA_ACC, 200)
-  tn1 = UOp.range(TN, 201)
-  S_frag = S_reg.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0, 2, 1)[tm1, tn1]
-  q_frag = Q_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, D // WMMA_K, WMMA_K)[wave_m, tm1, lane_n, k_qk]
-  k_frag = KV_lds_k.reshape(WAVES_N, TN, WMMA_N, D // WMMA_K, WMMA_K)[wave_n, tn1, lane_n, k_qk]
-  qk = UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), *WMMA_ARG)
-  qk_done = S_frag.store(qk).end(tm1, tn1).end(k_qk)
-  S_reg = S_reg.after(qk_done)
-
-  # -- softmax in registers with warp shuffles --
-  S_reg = S_reg.after(S_reg.store(S_reg * SCALE))
-
-  # per-thread local row max over TN=4 elements, then warp reduce across 16 lanes
-  m_ij = UOp.placeholder((TM,), dtypes.float, slot=7, addrspace=AddrSpace.REG)
-  m_ij = m_ij.after(m_ij.after(n_tile).store(m_ij.const_like(-math.inf)))
-  rm2 = UOp.range(TN, 261, AxisType.REDUCE)
-  m_ij = m_ij.after(m_ij.store(m_ij.after(rm2).maximum(S_reg[:, rm2])).end(rm2))
-  # warp reduce max (in-place)
-  ri_w = UOp.range(TM, 270)
-  m_ij = m_ij.after(m_ij[ri_w].store(warp_reduce_max(m_ij[ri_w], lane)).end(ri_w))
-
-  # compute P = exp(S - m_ij) in S_reg
-  S_reg = S_reg.after(S_reg.store(((S_reg - m_ij.reshape(TM, 1).expand(TM, TN)) * LOG2E).exp2()))
-
-  p_local = UOp.placeholder((TM,), dtypes.float, slot=8, addrspace=AddrSpace.REG)
-  p_local = p_local.after(p_local.after(n_tile).store(p_local.const_like(0)))
-  rp2 = UOp.range(TN, 291, AxisType.REDUCE)
-  p_local = p_local.after(p_local.store(p_local.after(rp2) + S_reg[:, rp2]).end(rp2))
-  ri_ws = UOp.range(TM, 295)
-  p_sum = p_local.after(p_local[ri_ws].store(warp_reduce_sum(p_local[ri_ws], lane)).end(ri_ws))
-
-  # write P = exp(S - m_ij) to P_lds (reuses slot 0, Q no longer needed)
-  P_lds = QP_lds[:, :BLOCK_N]
-  P_write = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_ACC, LANES_PER_WAVE_M, WAVES_N, TN, LANES_PER_WAVE_N)
-  P_write = P_write.permute((0, 4, 3, 6, 1, 2, 5)).reshape(THREADS_PER_BLOCK, TM, TN)
-  P_store = P_write[tid].store(S_reg.cast(dtypes.half))
-
-  # -- online softmax correction --
-  ri4 = UOp.range(TM, 330)
-  m_new_val = m_i[ri4].maximum(m_ij[ri4])
-  alpha_val = ((m_i[ri4] - m_new_val) * LOG2E).exp2()
-  beta_val = ((m_ij[ri4] - m_new_val) * LOG2E).exp2()
-  rj4 = UOp.range(TD, 331)
-  correction = UOp.group(
-    acc[ri4, rj4].store(alpha_val * acc[ri4, rj4]).end(rj4),
-    l_i[ri4].store(alpha_val * l_i[ri4] + beta_val * p_sum[ri4]),
-    m_i[ri4].store(m_new_val),
-  ).end(ri4)
-  acc = acc.after(correction)
-  l_i = l_i.after(correction)
-  m_i = m_i.after(correction)
-
-  # load V into KV_lds (must wait for QK WMMA to finish reading K from KV_lds)
-  V_store = KV_lds.after(qk_done).reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
-    v[n_tile].reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
-  # NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
-  P_lds = P_lds.after(UOp.group(P_store, V_store))
-  KV_lds_v = KV_lds.after(UOp.group(P_store, V_store))
-
-  # -- acc += P @ V via WMMA --
-  k_pv = UOp.range(BLOCK_N // WMMA_K, 400, AxisType.REDUCE)
-  tm2 = UOp.range(TM // WMMA_ACC, 401)
-  tn2 = UOp.range(TD, 402)
-  acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TD).permute(0, 2, 1)[tm2, tn2]
-  p_frag = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_N // WMMA_K, WMMA_K)[wave_m, tm2, lane_n, k_pv]
-  v_frag = KV_lds_v.reshape(WAVES_N, TD, WMMA_N, BLOCK_N // WMMA_K, WMMA_K)[wave_n, tn2, lane_n, k_pv]
-  pv = UOp.wmma(p_frag, v_frag, acc_frag.after(k_pv), *WMMA_ARG)
-
-  # end KV tile loop
-  n_tile_end = acc_frag.store(pv).end(tm2, tn2).end(k_pv).end(n_tile)
-  acc = acc.after(n_tile_end)
-  l_i = l_i.after(n_tile_end)
-  m_i = m_i.after(n_tile_end)
-
-  # normalize: acc /= l_i
-  acc = acc.after(acc.store(acc * (1 / l_i).reshape(TM, 1).expand(TM, TD)))
-
-  # store output
-  o = o.reshape(WAVES_M, TM // WMMA_ACC, WMMA_ACC, LANES_PER_WAVE_M, WAVES_N, TD, LANES_PER_WAVE_N)
-  o = o.permute((0, 4, 3, 6, 1, 2, 5)).reshape(THREADS_PER_BLOCK, TM, TD)
-  return o[tid].store(acc).end(wave_m, wave_n, lane).end(block_m, block_bh).sink(arg=KernelInfo(opts_to_apply=()))
+from tinygrad import Tensor, getenv
+from tinygrad.dtype import dtypes
+from tinygrad.helpers import GlobalCounters, Context
+from tinygrad.llm.kernels.amd import amd_flash_attention, amd_flash_attention_causal
 
 if __name__ == "__main__":
   B, H, N, D = getenv("B", 1), getenv("H", 32), getenv("N", 1024), getenv("D", 64)
-  q = Tensor.rand(B, H, N, D).cast(dtypes.half)
+  M, causal = getenv("M", N), getenv("CAUSAL", 0)
+  q = Tensor.rand(B, H, M, D).cast(dtypes.half)
   k = Tensor.rand(B, H, N, D).cast(dtypes.half)
   v = Tensor.rand(B, H, N, D).cast(dtypes.half)
-  o = Tensor.empty(B, H, N, D, dtype=dtypes.float)
+  o = Tensor.empty(B, H, M, D, dtype=dtypes.float)
   with Context(DEBUG=0): Tensor.realize(q, k, v)
 
-  q_flat, k_flat, v_flat, o_flat = q.reshape(B*H, N, D), k.reshape(B*H, N, D), v.reshape(B*H, N, D), o.reshape(B*H, N, D)
-  NUM_RUNS = getenv("CNT", 5)
+  q_flat, k_flat, v_flat, o_flat = q.reshape(B*H, M, D), k.reshape(B*H, N, D), v.reshape(B*H, N, D), o.reshape(B*H, M, D)
   ets = []
   with Context(DEBUG=2):
-    for _ in range(NUM_RUNS):
+    for _ in range(getenv("CNT", 5)):
       GlobalCounters.reset()
-      tst = Tensor.custom_kernel(o_flat, q_flat, k_flat, v_flat, fxn=amd_flash_attention)[0].realize()
+      tst = Tensor.custom_kernel(o_flat, q_flat, k_flat, v_flat,
+                                 fxn=amd_flash_attention_causal if causal else amd_flash_attention)[0].realize()
       ets.append(GlobalCounters.time_sum_s)
   print(f"best time: {min(ets)*1e3:.2f}ms")
 
   if getenv("VERIFY", 1):
     with Context(DEBUG=0):
-      ref = q.float().scaled_dot_product_attention(k.float(), v.float()).reshape(B*H, N, D).realize()
-      err = (ref - tst).square().mean().item()
-    print(f"mean squared error {err}")
-    if err > 1e-2:
-      raise RuntimeError("flash attention is wrong!")
-    else:
-      print("flash attention is correct!")
+      mask = Tensor.full((1, 1, M, N), float("-inf"), buffer=False).triu(N-M+1) if causal else None
+      ref = q.float().scaled_dot_product_attention(k.float(), v.float(), attn_mask=mask).reshape(B*H, M, D).realize()
+      diff = (ref - tst).abs()
+      err, max_err = diff.square().mean().item(), diff.max().item()
+    print(f"mean squared error {err}, max error {max_err}")
+    if err > 1e-2: raise RuntimeError("flash attention is wrong!")
+    print("flash attention is correct!")

test/backend/test_amd_flash_attention.py

@@ -0,0 +1,65 @@
+import unittest
+import numpy as np
+
+from tinygrad import Device, Tensor, TinyJit
+from tinygrad.llm.kernels.amd import amd_flash_attention_decode, flash_attention_causal_cached
+
+
+@unittest.skipUnless(Device.DEFAULT.startswith("AMD"), "AMD flash attention required")
+class TestAMDFlashAttention(unittest.TestCase):
+  def _test_decode(self, max_kv_len:int, valid_kv_len:int, n_heads:int=16, n_kv_heads:int=2, quantized:bool=False):
+    rng = np.random.default_rng(1)
+    q_np = rng.standard_normal((1, n_heads, 1, 256)).astype(np.float16)
+    kv_np = rng.standard_normal((2, 1, n_kv_heads, max_kv_len, 256)).astype(np.float16)
+    scale_np = np.maximum(np.max(np.abs(kv_np.astype(np.float32)), axis=-1), 1e-8) / 127
+    if quantized:
+      kv_np = np.clip(np.rint(kv_np.astype(np.float32) / scale_np[..., None]), -127, 127).astype(np.int8)
+    q, kv = Tensor(q_np).realize(), Tensor(kv_np).realize()
+    scale = Tensor(scale_np.astype(np.float16)).realize() if quantized else None
+
+    @TinyJit
+    def decode(q:Tensor, kv:Tensor): return amd_flash_attention_decode(q, kv, valid_kv_len, max_kv_len, scale).realize()
+
+    out = None
+    for _ in range(3): out = decode(q, kv).numpy()
+    assert out is not None
+    q_ref = q_np[0, :, 0].astype(np.float32)
+    kv_ref = kv_np.astype(np.float32) * scale_np[..., None] if quantized else kv_np.astype(np.float32)
+    k_ref, v_ref = kv_ref[:, 0, :, :valid_kv_len]
+    expected = np.empty((n_heads, 256), dtype=np.float32)
+    for head in range(n_heads):
+      scores = q_ref[head] @ k_ref[head // (n_heads // n_kv_heads)].T / np.sqrt(256)
+      probs = np.exp(scores - scores.max())
+      expected[head] = probs @ v_ref[head // (n_heads // n_kv_heads)] / probs.sum()
+
+    self.assertTrue(np.isfinite(out).all())
+    np.testing.assert_allclose(out[0, :, 0], expected, rtol=2e-3, atol=2e-3)
+
+  def test_short_decode_is_finite_and_matches_reference(self): self._test_decode(8192, 25)
+
+  def test_q8_cache_matches_dequantized_reference(self): self._test_decode(8192, 25, quantized=True)
+
+  def test_q8_cached_prefill_matches_dequantized_reference(self):
+    rng = np.random.default_rng(2)
+    heads, kv_heads, tokens, dim = 16, 2, 32, 256
+    q = rng.standard_normal((1, heads, tokens, dim)).astype(np.float16)
+    kv = rng.standard_normal((2, 1, kv_heads, tokens, dim)).astype(np.float16)
+    scale = np.maximum(np.max(np.abs(kv.astype(np.float32)), axis=-1), 1e-8) / 127
+    packed = np.clip(np.rint(kv.astype(np.float32) / scale[..., None]), -127, 127).astype(np.int8)
+    got = flash_attention_causal_cached(Tensor(q).realize(), Tensor(packed).realize(), tokens, tokens,
+                                        Tensor(scale.astype(np.float16)).realize()).numpy()
+    dequant = packed.astype(np.float32) * scale.astype(np.float16).astype(np.float32)[..., None]
+    expected = np.empty_like(got)
+    for head in range(heads):
+      scores = q[0, head].astype(np.float32) @ dequant[0, 0, head // (heads // kv_heads)].T / np.sqrt(dim)
+      scores[np.triu_indices(tokens, 1)] = -np.inf
+      probs = np.exp(scores - scores.max(axis=-1, keepdims=True))
+      expected[0, head] = probs @ dequant[1, 0, head // (heads // kv_heads)] / probs.sum(axis=-1, keepdims=True)
+    np.testing.assert_allclose(got, expected, rtol=2e-3, atol=2e-3)
+
+  def test_six_query_heads_per_kv_head(self): self._test_decode(8192, 25, n_heads=12, n_kv_heads=2)
+
+  def test_hierarchical_decode_matches_reference(self): self._test_decode(16384, 4097)
+
+
+if __name__ == "__main__": unittest.main()

test/backend/test_llama_kernels.py

@@ -131,6 +131,7 @@ class TestFusedQKVRoPE(unittest.TestCase):
       self.assertTrue(v.allclose(v_ref, atol=0, rtol=0).item(), "V forward mismatch")
 
   def test_llama31_8b_backward(self):
+    if not Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950"): self.skipTest("requires CDNA4")
     Tensor.manual_seed(1)
     B, N, H, H_KV, D = self.SHAPE
     PARTIALS = 2

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/null/test_hcq_iface.py

@@ -1,9 +1,26 @@
 import unittest, array, time
 from tinygrad.helpers import mv_address
 from tinygrad.runtime.support.hcq import MMIOInterface
+from tinygrad.runtime.support.memory import VirtMapping
+from tinygrad.runtime.support.system import PCIIfaceBase
 from tinygrad.runtime.support.usb import USBMMIOInterface
 from test.mockgpu.usb import MockUSB
 
+class TestPCIIface(unittest.TestCase):
+  def test_sysmem_mapping_respects_uncached(self):
+    class MM:
+      def alloc_vaddr(self, size, align): return 0x10000
+      def map_range(self, vaddr, size, paddrs, aspace, uncached=False, snooped=False):
+        return VirtMapping(vaddr, size, paddrs, aspace, uncached, snooped)
+    class PCI:
+      def bar_info(self, bar): return 0, 256 << 20
+      def alloc_sysmem(self, size, **kwargs): return memoryview(bytearray(size)), [0x20000]
+    iface = PCIIfaceBase.__new__(PCIIfaceBase)
+    iface.dev, iface.vram_bar, iface.pci_dev = None, 0, PCI()
+    iface.dev_impl = type("DevImpl", (), {"mm": MM()})()
+    for uncached in (False, True):
+      with self.subTest(uncached=uncached): self.assertEqual(iface.alloc(4096, host=True, uncached=uncached).meta.mapping.uncached, uncached)
+
 class TestHCQIface(unittest.TestCase):
   def setUp(self):
     self.size = 4 << 10

test/null/test_llm_server.py

@@ -95,6 +95,17 @@ class TestLLMServer(unittest.TestCase):
     self.assertGreater(len(chunks), 0)
     self.assertEqual(chunks[-1].choices[0].finish_reason, "stop")
 
+  def test_chat_template_kwargs(self):
+    import jinja2
+    with patch.object(self.server, "template", jinja2.Template(
+        "{% for m in messages %}{% if m.role == 'assistant' and preserve_thinking %}<think>{{ m.reasoning_content }}</think>"
+        "{% endif %}{{ m.content }}|{% endfor %}")):
+      list(self.client.chat.completions.create(model="test", messages=[
+        {"role":"user", "content":"first"}, {"role":"assistant", "content":"answer", "reasoning_content":"reason"},
+        {"role":"user", "content":"next"}], stream=True,
+        extra_body={"chat_template_kwargs":{"preserve_thinking":True, "messages":[]}}))
+    self.mock_tok.encode.assert_called_with("first|<think>reason</think>answer|next|")
+
   def test_content_is_streamed(self):
     stream = self.client.chat.completions.create(
       model="test",
@@ -109,6 +120,28 @@ class TestLLMServer(unittest.TestCase):
 
     self.assertGreater(len(contents), 0)
 
+  def test_interrupted_stream_logs_tokens(self):
+    with patch.object(self.mock_model, "generate", side_effect=lambda ids, **kwargs: iter([300, 301, 999])), \
+         patch("tinygrad.llm.serve.stderr_log") as log, patch("tinygrad.llm.serve.colored", side_effect=lambda text, color: text) as color:
+      stream = self.server.RequestHandlerClass.run_model(Mock(server=self.server), [200, 201, 202], "test")
+      next(stream)
+      next(stream)
+      stream.close()
+    interrupt = log.call_args.args[0]
+    self.assertFalse(interrupt.startswith("\n"))
+    self.assertTrue(interrupt.endswith("\n"))
+    self.assertIn("gen:", interrupt)
+    self.assertIn("out:    1", interrupt)
+    self.assertTrue(any(args[0].startswith("total:") and args[1] == "red" for args, _ in color.call_args_list))
+
+  def test_stream_disconnect_closes_source(self):
+    from tinygrad.viz.serve import HTTPRequestHandler
+    source, handler = Mock(), Mock()
+    source.__iter__ = Mock(return_value=iter([{}]))
+    handler.wfile.write.side_effect = BrokenPipeError
+    HTTPRequestHandler.stream_json(handler, source)
+    source.close.assert_called_once()
+
   def test_non_streaming(self):
     resp = self.client.chat.completions.create(
       model="test-model",

test/null/test_schedule_cache.py

@@ -1,10 +1,12 @@
-import unittest
-from tinygrad import Tensor, Variable, Context
+import os, unittest
+from unittest.mock import patch
+from tinygrad import Tensor, Variable, Context, UOp
+from tinygrad.callify import transform_to_call
 from tinygrad.helpers import cpu_events
-from tinygrad.schedule import schedule_cache
+from tinygrad.schedule import lower_sink_to_linear, schedule_cache
 
 def schedule_one():
-  Tensor([1]).schedule_linear()
+  (Tensor.empty(1) + 1).schedule_linear()
 
 class TestScheduleCache(unittest.TestCase):
   def test_bound_variable_var_vals(self):
@@ -37,5 +39,22 @@ class TestScheduleCache(unittest.TestCase):
       num_events_cache = len(cpu_events)
     self.assertLess(num_events_cache, num_events_no_cache)
 
+  def test_disk_schedule_cache(self):
+    function = transform_to_call(UOp.sink((Tensor.empty(1) + 1).uop))[0].src[0]
+    schedule_cache.clear()
+    with patch.dict(os.environ, {"DISK_SCACHE":"1"}), \
+         patch("tinygrad.schedule.diskcache_get", return_value=None), \
+         patch("tinygrad.schedule.diskcache_put") as cache_put:
+      lower_sink_to_linear(function)
+      cached = cache_put.call_args.args[2]
+
+    schedule_cache.clear()
+    with patch.dict(os.environ, {"DISK_SCACHE":"1"}), \
+         patch("tinygrad.schedule.diskcache_get", return_value=cached) as cache_get, \
+         patch("tinygrad.schedule.diskcache_put") as cache_put:
+      self.assertIs(lower_sink_to_linear(function), cached)
+      cache_get.assert_called_once()
+      cache_put.assert_not_called()
+
 if __name__ == "__main__":
   unittest.main()

test/unit/test_attention.py

@@ -1,9 +1,9 @@
 import unittest
 import numpy as np
-from tinygrad import Tensor, dtypes
+from tinygrad import Device, Tensor, dtypes
 from tinygrad.llm.model import (
   GatedDeltaNetBlock, SSMConfig, TransformerBlock, TransformerConfig,
-  apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk,
+  apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk, topk_softmax,
 )
 
 def apply_rope(x:Tensor, start_pos:int):
@@ -67,7 +67,10 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
   def _run_attention(self, block:GatedDeltaNetBlock, x:Tensor, start_pos:int):
     x_norm = block.attn_norm(x)
     block._init_state(x_norm)
-    return block._attention(x_norm, start_pos).realize().numpy()
+    out = block._attention(x_norm, start_pos).realize()
+    assert block.pending_state is not None
+    Tensor.realize(block.conv_state.assign(block.pending_state[0]), block.recurrent_state.assign(block.pending_state[1]))
+    return out.numpy()
 
   def _cache_views(self, block:GatedDeltaNetBlock) -> tuple[np.ndarray, np.ndarray]:
     if hasattr(block, 'conv_state'):
@@ -86,8 +89,8 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
     x_float = x.astype(np.float32)
     return (x_float / np.sqrt((x_float * x_float).mean(axis=-1, keepdims=True) + eps)) * weight.astype(np.float32)
 
-  def _normalize_np(self, x:np.ndarray, eps:float=1e-12) -> np.ndarray:
-    return x / np.maximum(np.sqrt((x * x).sum(axis=-1, keepdims=True)), eps)
+  def _normalize_np(self, x:np.ndarray, eps:float=1e-6) -> np.ndarray:
+    return x / np.sqrt((x * x).sum(axis=-1, keepdims=True) + eps)
 
   def _softplus_np(self, x:np.ndarray) -> np.ndarray:
     return np.log1p(np.exp(-np.abs(x))) + np.maximum(x, 0)
@@ -199,5 +202,27 @@ class TestPairwiseTopk(unittest.TestCase):
         self.assertEqual(set(sel.numpy()[b, t].tolist()), expected)
         np.testing.assert_allclose(vals.numpy()[b, t], data[b, t][sel.numpy()[b, t]])
 
+  def test_256_experts_matches_numpy(self):
+    rng = np.random.default_rng(42)
+    data = rng.standard_normal((256, 256) if Device.DEFAULT.startswith("AMD") else (4, 3, 256), dtype=np.float32)
+    # Include ties crossing wave boundaries to cover deterministic expert selection.
+    data[..., [7, 39, 71, 103, 135, 167, 199, 231]] = 10.0
+    expected = np.apply_along_axis(lambda row:np.lexsort((-np.arange(256), row))[-8:], -1, data)
+    x = Tensor(data)
+    for _ in range(5 if Device.DEFAULT.startswith("AMD") else 1):
+      vals, sel = pairwise_topk(x, 8)
+      np.testing.assert_equal(sel.numpy(), expected)
+      np.testing.assert_allclose(vals.numpy(), np.take_along_axis(data, expected, axis=-1))
+
+  def test_256_experts_softmax_matches_reference(self):
+    rng = np.random.default_rng(123)
+    data = rng.standard_normal((256, 256) if Device.DEFAULT.startswith("AMD") else (4, 3, 256), dtype=np.float32)
+    data[..., [7, 39, 71, 103, 135, 167, 199, 231]] = 10.0
+    probs, sel = topk_softmax(Tensor(data), 8)
+    selected = np.take_along_axis(data, sel.numpy(), axis=-1)
+    expected = np.exp(selected - selected.max(axis=-1, keepdims=True))
+    expected /= expected.sum(axis=-1, keepdims=True)
+    np.testing.assert_allclose(probs.numpy(), expected, rtol=2e-6, atol=2e-7)
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_cpu.py

@@ -1,13 +1,47 @@
-import unittest, io
+import unittest, io, os, subprocess, sys
 from contextlib import redirect_stdout
-from tinygrad import Tensor, Device
+from tinygrad import Tensor, Device, UOp
 from tinygrad.helpers import Target
 from tinygrad.renderer.nir import LVPRenderer
 from tinygrad.renderer.isa.x86 import X86Renderer
 from tinygrad.codegen import to_program
+from tinygrad.runtime.ops_cpu import RING_SLOTS
+from tinygrad.uop.ops import AxisType, KernelInfo
 
 @unittest.skipIf(Device.DEFAULT != "CPU", "only run on CPU")
 class TestCPU(unittest.TestCase):
+  def test_parallel_workers_exit_cleanly(self):
+    env = os.environ.copy()
+    env.update(DEV="CPU", CPU_PARALLEL_UOPS="1")
+    proc = subprocess.run([sys.executable, "-c",
+      "from tinygrad import Tensor; assert Tensor.arange(32).sum().item() == 496"], env=env, capture_output=True, text=True)
+    self.assertEqual(proc.returncode, 0, proc.stderr)
+
+  def test_32_buffer_kernel(self):
+    def add_inputs(out:UOp, *inputs:UOp) -> UOp:
+      return out[0].store(sum((x[0] for x in inputs), start=UOp.const(out.dtype, 0))).sink(
+        arg=KernelInfo(name="add_31_inputs", opts_to_apply=()))
+    inputs = [Tensor([i], device="CPU").realize() for i in range(31)]
+    out = Tensor.custom_kernel(Tensor.empty(1, device="CPU"), *inputs, fxn=add_inputs)[0]
+    self.assertEqual(out.item(), sum(range(31)))
+
+  def test_command_ring_backpressure(self):
+    dev, count = Device["CPU"], RING_SLOTS + 257
+    signal, queue = dev.new_signal(value=0), dev.hw_compute_queue_t()
+    for value in range(1, count + 1): queue.signal(signal, value)
+    queue.submit(dev)
+    signal.wait(count, timeout=10000)
+    self.assertEqual(signal.value, count)
+
+  def test_parallel_launch(self):
+    def fill(out:UOp) -> UOp:
+      idx = UOp.range(67, 0, AxisType.GLOBAL)
+      return out[idx].store(idx).end(idx).sink(arg=KernelInfo(name="parallel_launch", optimize=False, parallel=True))
+    probe = Tensor.custom_kernel(Tensor.empty(67, device="CPU"), fxn=fill)[0]
+    self.assertTrue(to_program(probe.schedule_linear().src[-1].src[0], Device["CPU"].renderer).arg.parallel)
+    out = Tensor.custom_kernel(Tensor.empty(67, device="CPU"), fxn=fill)[0]
+    self.assertEqual(out.tolist(), list(range(67)))
+
   def test_arch_feats(self):
     ast = (Tensor.empty(16) + Tensor.empty(16)).schedule_linear().src[-1].src[0]
     for ren in Device[Device.DEFAULT].renderers:

test/unit/test_function.py

@@ -557,6 +557,14 @@ class TestFunctionTuple(unittest.TestCase):
     def f(a:Tensor): return Tensor.custom_kernel(Tensor.empty(*a.shape, dtype=a.dtype, device=a.device), a, fxn=inplace_add)[0]
     with self.assertRaisesRegex(RuntimeError, "implicit buffer"): f(Tensor([1., 2., 3., 4.]).contiguous().realize())
 
+  def test_shrink_load_is_program_input(self):
+    out, inp = UOp.param(0, dtypes.float, (1,)), UOp.param(1, dtypes.float, (8,))
+    values = UOp(Ops.SHRINK, src=(inp, UOp.const(dtypes.weakint, 0), UOp.const(dtypes.weakint, 8))).load()
+    sink = out[0].store(values.index(0)).sink(arg=KernelInfo(name="vector_load"))
+    info = ProgramInfo.from_sink(sink)
+    self.assertEqual(info.outs, (0,))
+    self.assertEqual(info.ins, (1,))
+
   def test_custom_kernel_write_only_persistent_output_is_implicit(self):
     # a write-only custom_kernel output that is a realized buffer must be captured
     def write(C:UOp, A:UOp) -> UOp:

test/unit/test_jit.py

@@ -2,6 +2,7 @@ import unittest, numpy as np
 from test.helpers import assert_jit_cache_len
 from tinygrad import Tensor, TinyJit, Context, UOp, dtypes
 from tinygrad.engine.jit import JitError
+from tinygrad.uop.ops import Ops
 
 def _simple_test(add, extract=lambda x: x, N=10):
   for _ in range(5):
@@ -12,6 +13,19 @@ def _simple_test(add, extract=lambda x: x, N=10):
   assert_jit_cache_len(add, 1)
 
 class TestJit(unittest.TestCase):
+  def test_parallel_compile(self):
+    from tinygrad.codegen import to_program_cache
+    from tinygrad.engine.realize import compile_linear
+    calls = [call for i in range(16) for call in (Tensor.empty(1, device="CPU") + i).schedule_linear().src]
+    cache = to_program_cache.copy()
+    try:
+      to_program_cache.clear()
+      with Context(PARALLEL_COMPILE=2): linear = compile_linear(UOp(Ops.LINEAR, src=tuple(calls)), jit=True)
+      self.assertTrue(all(call.op is not Ops.CALL or call.src[0].op is Ops.PROGRAM for call in linear.src))
+    finally:
+      to_program_cache.clear()
+      to_program_cache.update(cache)
+
   def test_jitbeam_triggers_beam(self):
     from unittest.mock import patch
     from tinygrad.helpers import getenv as _getenv

test/unit/test_llm_quant_cpu.py

@@ -0,0 +1,752 @@
+import functools, sys, unittest
+import numpy as np
+
+from tinygrad import Device, Tensor, TinyJit, UOp, dtypes, nn
+from tinygrad.llm.gguf import _GGML_QUANT, ggml_data_to_tensor
+from tinygrad.llm.kernels import amd as llm_amd
+from tinygrad.llm.kernels.cpu import (attention_decode, attention_prefill, causal_conv_silu, expert_pair, expert_silu,
+                              expert_weighted_sum, f16_linear,
+                              f16_matvec, gated_delta, gated_delta_prefill, gated_delta_q8, gdn_qkv, iq3_repack, moe_ffn, q6_argmax, q8_batched_pair,
+                              q8_gdn_norm_projections, q8_gdn_projections, q8_linear_pair, q8_repack, q8_silu_linear,
+                              rmsnorm, rmsnorm_f16_linear, shared_gate,
+                              silu, silu_mul, uop_attention_prefill, uop_f16_matvec, uop_linear, uop_moe_ffn, uop_q8_linear_pair,
+                              uop_q8_prequant_linear, uop_expert_silu_weighted, weighted_sum)
+from tinygrad.llm.kernels.cpu import _dot_bytes_ptr, _dot_nibbles_ptr
+from tinygrad.llm.model import biased_sigmoid_topk, pairwise_topk, Embedding, ExpertWeights, FFNBlock, Linear, Transformer, TransformerConfig
+from tinygrad.uop.ops import KernelInfo
+
+
+def q8_activation(x:np.ndarray) -> np.ndarray:
+  grouped = x.reshape(*x.shape[:-1], -1, 32)
+  scale = np.maximum(np.max(np.abs(grouped), axis=-1, keepdims=True) / 127, 1e-8)
+  return (np.clip(np.rint(grouped / scale), -127, 127) * scale).reshape(x.shape)
+
+def q8k_activation(x:np.ndarray) -> np.ndarray:
+  grouped = x.reshape(*x.shape[:-1], -1, 256)
+  signed_max = np.take_along_axis(grouped, np.argmax(np.abs(grouped), axis=-1, keepdims=True), axis=-1)
+  scale = -signed_max / 127
+  inverse = np.divide(1, scale, out=np.zeros_like(scale), where=scale != 0)
+  quantized = np.sign(grouped * inverse) * np.floor(np.abs(grouped * inverse) + 0.5)
+  return (np.minimum(quantized, 127) * scale).reshape(x.shape)
+
+
+def random_packed(rng:np.random.Generator, ggml_type:int, elements:int) -> np.ndarray:
+  block_size, type_size = _GGML_QUANT[ggml_type]
+  blocks = rng.integers(0, 256, size=(elements // block_size, type_size), dtype=np.uint8)
+  scales = rng.uniform(0.001, 0.02, size=len(blocks)).astype(np.float16).view(np.uint8).reshape(-1, 2)
+  blocks[:, :2] = scales
+  if ggml_type in (12, 13): blocks[:, 2:4] = scales
+  if ggml_type == 14: blocks[:, -2:] = scales
+  return blocks.flatten()
+
+
+@unittest.skipUnless(Device.DEFAULT == "AMD", "requires DEV=AMD")
+class TestLLMQuantAMD(unittest.TestCase):
+  @staticmethod
+  def assert_q8_equal(result:tuple[Tensor, Tensor, Tensor], expected:np.ndarray):
+    grouped = expected.reshape(expected.shape[0], -1, 32)
+    scale = np.maximum(np.max(np.abs(grouped), axis=-1) / 127, 1e-8)
+    quant = np.clip(np.rint(grouped / scale[..., None]), -127, 127).astype(np.int8)
+    np.testing.assert_equal(result[0].numpy().view(np.int8).reshape(grouped.shape), quant)
+    np.testing.assert_allclose(result[1].numpy(), scale, rtol=1e-6, atol=1e-8)
+    np.testing.assert_equal(result[2].numpy(), quant.astype(np.int32).sum(-1))
+
+  def test_gated_delta_decode_matches_reference(self):
+    rng = np.random.default_rng(39)
+    batch, heads, dim = 1, 2, 128
+    q, k, v = [rng.standard_normal((batch, heads, dim), dtype=np.float32) for _ in range(3)]
+    beta, alpha = rng.random((batch, heads), dtype=np.float32), rng.uniform(0.9, 1, (batch, heads)).astype(np.float32)
+    state = rng.standard_normal((batch, heads, dim, dim), dtype=np.float32).astype(np.float16)
+    state_k, state_q = np.einsum("bhij,bhj->bhi", state, k), np.einsum("bhij,bhj->bhi", state, q)
+    delta = (v - state_k * alpha[..., None]) * beta[..., None]
+    expected_core = state_q * alpha[..., None] + delta * np.sum(k*q, axis=-1)[..., None]
+    expected_state = (state * alpha[..., None, None] + delta[..., None] * k[..., None, :]).astype(np.float16)
+    core, next_state = llm_amd.gated_delta_decode(
+      *(Tensor(x, device="AMD") for x in (q, k, v, beta, alpha)), Tensor(state, device="AMD"))
+    Tensor.realize(core, next_state)
+    np.testing.assert_allclose(core.numpy(), expected_core, rtol=2e-4, atol=1e-3)
+    np.testing.assert_allclose(next_state.numpy(), expected_state, rtol=1e-3, atol=4e-3)
+
+  def test_f16_matvec_matches_reference(self):
+    rng, in_features, out_features = np.random.default_rng(38), 512, 13
+    x = rng.standard_normal((1, in_features), dtype=np.float32).astype(np.float16)
+    weight = rng.standard_normal((out_features, in_features), dtype=np.float32).astype(np.float16)
+    got = llm_amd.f16_matvec(Tensor(x, device="AMD"), Tensor(weight, device="AMD")).numpy()
+    np.testing.assert_allclose(got, x.astype(np.float32) @ weight.astype(np.float32).T, rtol=2e-5, atol=2e-4)
+
+  def test_fused_rmsnorm_quantization_matches_reference(self):
+    rng, eps = np.random.default_rng(37), 1e-6
+    x = rng.standard_normal((1, 256), dtype=np.float32)
+    weight = rng.standard_normal((256,), dtype=np.float32).astype(np.float16)
+    normalized = x / np.sqrt(np.mean(x*x, axis=-1, keepdims=True) + eps) * weight
+    self.assert_q8_equal(llm_amd.q8_rmsnorm(Tensor(x, device="AMD"), Tensor(weight, device="AMD"), eps), normalized)
+
+    core = rng.standard_normal((1, 2, 128), dtype=np.float32)
+    gate = rng.standard_normal((1, 1, 2, 128), dtype=np.float32)
+    head_norm = core / np.sqrt(np.mean(core*core, axis=-1, keepdims=True) + eps) * weight[:128]
+    expected = head_norm * gate.reshape(1, 2, 128) / (1 + np.exp(-gate.reshape(1, 2, 128)))
+    self.assert_q8_equal(llm_amd.q8_gated_rmsnorm(*(Tensor(v, device="AMD") for v in (core, gate, weight[:128])), eps),
+                         expected.reshape(1, -1))
+
+  def test_gated_delta_prefill_matches_sequential_reference(self):
+    rng = np.random.default_rng(36)
+    batch, heads, tokens, dim = 1, 2, 5, 128
+    q, k, v = [rng.standard_normal((batch, heads, tokens, dim), dtype=np.float32) for _ in range(3)]
+    q = q / np.linalg.norm(q, axis=-1, keepdims=True) / np.float32(np.sqrt(dim))
+    k = k / np.linalg.norm(k, axis=-1, keepdims=True)
+    beta, alpha = rng.random((batch, heads, tokens), dtype=np.float32), rng.uniform(0.9, 1, (batch, heads, tokens)).astype(np.float32)
+    state = rng.standard_normal((batch, heads, dim, dim), dtype=np.float32).astype(np.float16)
+    expected_core, expected_state = np.empty_like(q), state.astype(np.float32)
+    for token in range(tokens):
+      state_k = np.einsum("bhij,bhj->bhi", expected_state, k[:, :, token])
+      state_q = np.einsum("bhij,bhj->bhi", expected_state, q[:, :, token])
+      delta = (v[:, :, token] - state_k * alpha[:, :, token, None]) * beta[:, :, token, None]
+      expected_core[:, :, token] = state_q * alpha[:, :, token, None] + delta * np.sum(k[:, :, token] * q[:, :, token], axis=-1)[..., None]
+      expected_state = expected_state * alpha[:, :, token, None, None] + delta[..., None] * k[:, :, token, None, :]
+    core, next_state = llm_amd.gated_delta_prefill(
+      *(Tensor(x, device="AMD") for x in (q, k, v, beta, alpha)), Tensor(state, device="AMD"))
+    np.testing.assert_allclose(core.numpy(), expected_core, rtol=2e-4, atol=1e-3)
+    np.testing.assert_allclose(next_state.numpy(), expected_state.astype(np.float16), rtol=2e-4, atol=1e-3)
+
+  def test_q8_quantize_matches_reference(self):
+    rng, tokens, in_features = np.random.default_rng(35), 3, 256
+    x = rng.standard_normal((tokens, in_features), dtype=np.float32)
+    grouped = x.reshape(tokens, -1, 32)
+    expected_scale = np.maximum(np.max(np.abs(grouped), axis=-1) / 127, 1e-8)
+    expected_quant = np.clip(np.rint(grouped / expected_scale[..., None]), -127, 127).astype(np.int8)
+    quant, scale, group_sum = llm_amd.q8_quantize_sum(Tensor(x, device="AMD"), tokens, in_features)
+    np.testing.assert_equal(quant.numpy().view(np.int8).reshape(grouped.shape), expected_quant)
+    np.testing.assert_allclose(scale.numpy(), expected_scale, rtol=1e-7, atol=0)
+    np.testing.assert_equal(group_sum.numpy(), expected_quant.astype(np.int32).sum(-1))
+
+  def test_q4_embedding_matches_reference(self):
+    rng, vocab_size, embed_size = np.random.default_rng(34), 16, 256
+    raw = random_packed(rng, 12, vocab_size * embed_size)
+    expected = ggml_data_to_tensor(Tensor(raw), vocab_size * embed_size, 12).reshape(vocab_size, embed_size).half()
+    storage = Tensor(np.concatenate((np.zeros(68, dtype=np.uint8), raw)), dtype=dtypes.uint8, device="AMD").realize()
+    embedding = Embedding(vocab_size, embed_size)
+    embedding.set_quantized(storage[68:], 12)
+    idx = np.array([[7, 1, 15], [0, 4, 7]], dtype=np.int32)
+    np.testing.assert_equal(embedding(Tensor(idx, device="AMD")).numpy(), expected.numpy()[idx])
+
+  def test_iq4_lut_is_ready_for_jit_capture(self):
+    rng, in_features, out_features = np.random.default_rng(33), 256, 16
+    raw = random_packed(rng, 23, out_features * in_features)
+    weight = ggml_data_to_tensor(Tensor(raw), out_features * in_features, 23).numpy().reshape(out_features, in_features)
+    llm_amd.iq4_half_lut.cache_clear()
+    layer = Linear(in_features, out_features, bias=False)
+    layer.set_quantized(Tensor(raw, dtype=dtypes.uint8, device="AMD").realize(), 23)
+    @TinyJit
+    def run(x:Tensor): return layer(x).realize()
+    x = rng.standard_normal((16, in_features), dtype=np.float32)
+    expected = x.astype(np.float16).astype(np.float32) @ weight.astype(np.float16).astype(np.float32).T
+    for _ in range(2): np.testing.assert_allclose(run(Tensor(x, device="AMD")).numpy(), expected, rtol=1e-5, atol=2e-3)
+
+  def test_packed_linear_offset_matches_reference(self):
+    rng = np.random.default_rng(32)
+    for ggml_type,in_features in ((8, 256), (12, 256), (13, 256), (14, 256), (23, 256)):
+      for tokens in ((1, 16, 32, 64, 128) if ggml_type == 23 else (1, 16, 128) if ggml_type in (12, 13) else
+                     (1, 16) if ggml_type == 14 else (1,)):
+        raw, out_features = random_packed(rng, ggml_type, 64 * in_features), 64
+        weight = ggml_data_to_tensor(Tensor(raw), out_features * in_features, ggml_type).numpy().reshape(out_features, in_features)
+        storage = Tensor(np.concatenate((np.zeros(68, dtype=np.uint8), raw)), dtype=dtypes.uint8, device="AMD").realize()
+        layer = Linear(in_features, out_features, bias=False)
+        layer.set_quantized(storage[68:], ggml_type)
+        x = rng.standard_normal((tokens, in_features), dtype=np.float32)
+        expected = x.astype(np.float16).astype(np.float32) @ weight.astype(np.float16).astype(np.float32).T \
+          if ggml_type in (12, 13, 23) and tokens > 1 else q8_activation(x) @ weight.T
+        np.testing.assert_allclose(layer(Tensor(x, device="AMD")).numpy(), expected, rtol=1e-5, atol=2e-3)
+
+  def test_iq3_expert_prefill_and_decode_match_reference(self):
+    rng = np.random.default_rng(31)
+    num_experts, in_features, out_features = 2, 256, 16
+    raw = random_packed(rng, 21, num_experts * out_features * in_features)
+    weight = ggml_data_to_tensor(Tensor(raw, device="CPU"), num_experts * out_features * in_features,
+                                 21).numpy().reshape(num_experts, out_features, in_features)
+    experts = ExpertWeights(num_experts, in_features, out_features)
+    experts.set_quantized(Tensor.empty(num_experts, out_features, in_features),
+                          Tensor(raw, dtype=dtypes.uint8, device="AMD").realize(), 21)
+    x = rng.standard_normal((2, 1, in_features), dtype=np.float32)
+    for sel in (np.array([[1, 0], [0, 1]], dtype=np.int32), np.array([1, 0], dtype=np.int32)):
+      activation = x if sel.ndim == 2 else x[:1]
+      expected = np.stack([q8_activation(activation).reshape(-1, in_features)[route // 2] @ weight[expert].T
+                           for route,expert in enumerate(sel.reshape(-1))]).reshape(*sel.shape, out_features)
+      got = experts(Tensor(sel, device="AMD"), Tensor(activation, device="AMD")).numpy()
+      np.testing.assert_allclose(got, expected, rtol=1e-5, atol=5e-4)
+
+
+@unittest.skipUnless(sys.platform.startswith("linux") and Device.DEFAULT == "CPU", "requires DEV=CPU on Linux")
+class TestLLMQuantCPU(unittest.TestCase):
+  def test_grouped_byte_dot_uop(self):
+    rng = np.random.default_rng(20)
+    a = rng.integers(-128, 128, 32, dtype=np.int8)
+    b = rng.integers(-127, 128, 32, dtype=np.int8)
+    out = Tensor.empty(8, dtype=dtypes.int32, device="CPU")
+    ta, tb = Tensor(a, device="CPU").realize(), Tensor(b, device="CPU").realize()
+    def dot_kernel(out:UOp, a:UOp, b:UOp) -> UOp:
+      parts = _dot_bytes_ptr(a[0], b[0])
+      return UOp.group(*(out[i].store(parts.index(i)) for i in range(8))).sink(arg=KernelInfo("grouped_byte_dot"))
+    got = Tensor.custom_kernel(out, ta, tb, fxn=dot_kernel)[0].numpy()
+    expected = (a.astype(np.int32) * b.astype(np.int32)).reshape(8, 4).sum(axis=1)
+    np.testing.assert_equal(got, expected)
+
+  def test_scaled_grouped_byte_dot_uop(self):
+    rng = np.random.default_rng(22)
+    a = rng.integers(-128, 128, 32, dtype=np.int8)
+    b = rng.integers(-127, 128, 32, dtype=np.int8)
+    out = Tensor.empty(8, dtype=dtypes.int32, device="CPU")
+    ta, tb = Tensor(a, device="CPU").realize(), Tensor(b, device="CPU").realize()
+    def dot_kernel(out:UOp, a:UOp, b:UOp) -> UOp:
+      parts = _dot_bytes_ptr(a[0], b[0]) * 7
+      return UOp.group(*(out[i].store(parts.index(i)) for i in range(8))).sink(arg=KernelInfo("scaled_grouped_byte_dot"))
+    got = Tensor.custom_kernel(out, ta, tb, fxn=dot_kernel)[0].numpy()
+    expected = (a.astype(np.int32) * b.astype(np.int32)).reshape(8, 4).sum(axis=1) * 7
+    np.testing.assert_equal(got, expected)
+
+  def test_unpack_lut_dot_uop(self):
+    rng = np.random.default_rng(21)
+    packed, x = rng.integers(0, 256, 16, dtype=np.uint8), rng.integers(-127, 128, 32, dtype=np.int8)
+    values = (-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113)
+    out = Tensor.empty(8, dtype=dtypes.int32, device="CPU")
+    tp, tx = Tensor(packed, device="CPU").realize(), Tensor(x, device="CPU").realize()
+    def dot_kernel(out:UOp, packed:UOp, x:UOp) -> UOp:
+      parts = _dot_nibbles_ptr(packed[0], x[0], values)
+      return UOp.group(*(out[i].store(parts.index(i)) for i in range(8))).sink(arg=KernelInfo("unpack_lut_dot"))
+    got = Tensor.custom_kernel(out, tp, tx, fxn=dot_kernel)[0].numpy()
+    decoded = np.array(values, dtype=np.int8)[np.concatenate((packed & 15, packed >> 4))]
+    expected = (decoded.astype(np.int32) * x.astype(np.int32)).reshape(8, 4).sum(axis=1)
+    np.testing.assert_equal(got, expected)
+
+  def test_generate_accepts_different_recurrent_prefill_shapes(self):
+    class TinyRecurrentTransformer(Transformer):
+      def __init__(self):
+        self.max_context, self.has_recurrent_block = 32, True
+        self.token_embd = nn.Embedding(4, 1)
+        self.blk, self._cached_tokens = [], []
+        self._state_checkpoints, self._state_checkpoint_pos = [], 0
+        self._save_state_jit = self._restore_state_jit = None
+        self._warming_up = False
+        self.prefill_jit = TinyJit(self.forward)
+        self.flash_prefill_jit = TinyJit(functools.partial(self.forward, use_flash=True))
+        self.sample_prefill_jit = TinyJit(functools.partial(self.forward, sample=True))
+        self.recurrent_prefill_jits = {}
+        self.rollout_jits, self.sample_rollout_jits = {}, {}
+
+      def forward(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor, use_flash:bool=False,
+                  kv_len:int|UOp|None=None, valid_len:int|UOp|None=None, sample:bool=False) -> Tensor:
+        return tokens[:, -1:] + 1
+
+    model = TinyRecurrentTransformer()
+    for _ in range(3): self.assertEqual(next(model.generate([1] * 8, chunk_size=8)), 2)
+    self.assertEqual(next(model.generate([1] * 3, chunk_size=8)), 2)
+
+  def test_attention_decode_matches_causal_gqa_reference(self):
+    rng = np.random.default_rng(0)
+    batch, heads, kv_heads, cache_len, head_dim = 1, 16, 2, 4113, 32
+    q = rng.standard_normal((batch, heads, 1, head_dim), dtype=np.float32)
+    cache = rng.standard_normal((2, batch, kv_heads, cache_len, head_dim), dtype=np.float32).astype(np.float16)
+    for pos in (0, 11, 4096):
+      with self.subTest(pos=pos):
+        start_pos = UOp.variable("start_pos", 0, cache_len-1).bind(pos)
+        got = attention_decode(Tensor(q, device="CPU").contiguous(),
+                                    Tensor(cache, device="CPU").contiguous(), start_pos).numpy()
+        expected = np.empty_like(got)
+        for head in range(heads):
+          kv_head = head // (heads // kv_heads)
+          keys = cache[0, 0, kv_head, :pos+1].astype(np.float32)
+          values = cache[1, 0, kv_head, :pos+1].astype(np.float32)
+          scores = q[0, head, 0] @ keys.T / np.sqrt(head_dim)
+          probs = np.exp(scores - scores.max())
+          expected[0, head, 0] = probs / probs.sum() @ values
+        np.testing.assert_allclose(got, expected, rtol=2e-5, atol=2e-5)
+
+  def test_attention_prefill_matches_causal_gqa_reference(self):
+    rng = np.random.default_rng(18)
+    batch, heads, tokens, kv_heads, cache_len, head_dim, pos = 1, 4, 5, 2, 19, 256, 3
+    q = rng.standard_normal((batch, heads, tokens, head_dim), dtype=np.float32)
+    cache = rng.standard_normal((2, batch, kv_heads, cache_len, head_dim), dtype=np.float32).astype(np.float16)
+    tq, tcache = Tensor(q, device="CPU").contiguous(), Tensor(cache, device="CPU").contiguous()
+    start_pos = UOp.variable("start_pos", 0, cache_len-1).bind(pos)
+    got = attention_prefill(tq, tcache, start_pos).numpy()
+    expected = np.empty_like(got)
+    for head in range(heads):
+      kv_head = head // (heads // kv_heads)
+      for token in range(tokens):
+        keys = cache[0, 0, kv_head, :pos+token+1].astype(np.float32)
+        values = cache[1, 0, kv_head, :pos+token+1].astype(np.float32)
+        scores = q[0, head, token] @ keys.T / np.sqrt(head_dim)
+        probs = np.exp(scores - scores.max())
+        expected[0, head, token] = probs / probs.sum() @ values
+    np.testing.assert_allclose(got, expected, rtol=2e-5, atol=2e-5)
+    np.testing.assert_allclose(uop_attention_prefill(tq, tcache, start_pos).numpy(), expected, rtol=2e-5, atol=2e-5)
+    np.testing.assert_allclose(uop_attention_prefill(tq[:, :, :4], tcache, start_pos).numpy(), expected[:, :, :4],
+                               rtol=2e-5, atol=2e-5)
+
+  def test_gated_delta_matches_reference(self):
+    rng = np.random.default_rng(4)
+    batch, heads, dim = 2, 3, 16
+    q, k, v = (rng.standard_normal((batch, heads, dim), dtype=np.float32) for _ in range(3))
+    beta, alpha = rng.random((batch, heads), dtype=np.float32), rng.random((batch, heads), dtype=np.float32)
+    state = rng.standard_normal((batch, heads, dim, dim), dtype=np.float32)
+    for state_dtype in (np.float32, np.float16):
+      with self.subTest(state_dtype=state_dtype):
+        typed_state = state.astype(state_dtype)
+        args = (*(Tensor(x, device="CPU") for x in (q, k, v, beta, alpha)), Tensor(typed_state, device="CPU"))
+        core, next_state = gated_delta(*args)
+        typed_state_k, typed_state_q = (np.einsum("bhij,bhj->bhi", typed_state.astype(np.float32), x) for x in (k, q))
+        typed_delta = (v - typed_state_k * alpha[..., None]) * beta[..., None]
+        typed_core = typed_state_q * alpha[..., None] + typed_delta * np.sum(k * q, axis=-1, keepdims=True)
+        typed_next = typed_state.astype(np.float32) * alpha[..., None, None] + typed_delta[..., None] * k[..., None, :]
+        np.testing.assert_allclose(core.numpy(), typed_core, rtol=2e-5, atol=2e-5)
+        np.testing.assert_allclose(next_state.numpy(), typed_next.astype(state_dtype), rtol=2e-5, atol=2e-5)
+        norm = nn.RMSNorm(dim, eps=1e-6)
+        norm.weight = Tensor(rng.standard_normal(dim, dtype=np.float32), device="CPU").half().realize()
+        normalized, _ = gated_delta(*args, norm_weight=norm.weight, norm_eps=norm.eps)
+        np.testing.assert_allclose(normalized.numpy(), rmsnorm(norm, core).numpy(), rtol=2e-5, atol=2e-5)
+        inplace_state = Tensor(typed_state, device="CPU").realize()
+        inplace_core, inplace_next = gated_delta(*(Tensor(x, device="CPU") for x in (q, k, v, beta, alpha)),
+                                                       inplace_state, inplace=True)
+        Tensor.realize(inplace_core, inplace_next)
+        np.testing.assert_allclose(inplace_core.numpy(), typed_core, rtol=2e-5, atol=2e-5)
+        np.testing.assert_allclose(inplace_next.numpy(), typed_next.astype(state_dtype), rtol=2e-5, atol=2e-5)
+        np.testing.assert_equal(inplace_state.numpy(), inplace_next.numpy())
+
+  def test_gated_delta_prefill_matches_sequential_reference(self):
+    rng = np.random.default_rng(17)
+    batch, heads, tokens, dim = 1, 2, 5, 128
+    q, k, v = [rng.standard_normal((batch, heads, tokens, dim), dtype=np.float32) for _ in range(3)]
+    beta = rng.random((batch, heads, tokens), dtype=np.float32)
+    alpha = rng.random((batch, heads, tokens), dtype=np.float32)
+    state = rng.standard_normal((batch, heads, dim, dim), dtype=np.float32).astype(np.float16)
+    weight = rng.standard_normal(dim, dtype=np.float32).astype(np.float16)
+    eps = 1e-6
+    expected_core = np.empty_like(q)
+    expected_state = state.astype(np.float32)
+    for token in range(tokens):
+      for b in range(batch):
+        for h in range(heads):
+          qq, kk, vv = q[b, h, token], k[b, h, token], v[b, h, token]
+          aa, bb, current = alpha[b, h, token], beta[b, h, token], expected_state[b, h]
+          delta = (vv - (current @ kk) * aa) * bb
+          core = (current @ qq) * aa + delta * (kk @ qq)
+          expected_state[b, h] = current * aa + delta[:, None] * kk[None, :]
+          expected_core[b, h, token] = core / np.sqrt(np.mean(core * core) + eps) * weight
+    got_core, got_state = gated_delta_prefill(
+      *[Tensor(z, device="CPU") for z in (q, k, v, beta, alpha)], Tensor(state, device="CPU"), Tensor(weight, device="CPU"), eps)
+    np.testing.assert_allclose(got_core.numpy(), expected_core, rtol=2e-5, atol=2e-5)
+    np.testing.assert_allclose(got_state.numpy(), expected_state.astype(np.float16), rtol=1e-3, atol=1e-3)
+
+  def test_gated_delta_q8_feeds_projection(self):
+    rng = np.random.default_rng(18)
+    batch, heads, dim = 1, 2, 32
+    q, k, v = [Tensor(rng.standard_normal((batch, heads, dim), dtype=np.float32), device="CPU").realize() for _ in range(3)]
+    beta, alpha = [Tensor(rng.random((batch, heads), dtype=np.float32), device="CPU").realize() for _ in range(2)]
+    state = Tensor(rng.standard_normal((batch, heads, dim, dim), dtype=np.float32), device="CPU").half().realize()
+    gate = Tensor(rng.standard_normal((batch, heads, dim), dtype=np.float32), device="CPU").half().realize()
+    norm_weight = Tensor(rng.standard_normal(dim, dtype=np.float32), device="CPU").half().realize()
+    layer = Linear(heads * dim, 17, bias=False)
+    layer.set_quantized(Tensor(random_packed(rng, 8, layer.out_features * layer.in_features),
+                               dtype=dtypes.uint8, device="CPU").realize(), 8)
+
+    core, expected_state = gated_delta(q, k, v, beta, alpha, state, norm_weight=norm_weight, norm_eps=1e-6)
+    expected = q8_silu_linear(layer, gate.reshape(batch, 1, -1), core.reshape(batch, 1, -1))
+    xq, xd, got_state = gated_delta_q8(q, k, v, beta, alpha, state, gate, norm_weight, 1e-6)
+    got = uop_q8_prequant_linear(layer, xq, xd).reshape(batch, 1, -1)
+    np.testing.assert_allclose(got_state.numpy(), expected_state.numpy(), rtol=2e-5, atol=2e-5)
+    np.testing.assert_allclose(got.numpy(), expected.numpy(), rtol=1e-3, atol=1e-3)
+
+  def test_gdn_qkv_matches_normal_api(self):
+    rng = np.random.default_rng(19)
+    batch, tokens, k_heads, v_heads, dim = 2, 5, 2, 4, 16
+    conv = Tensor(rng.standard_normal((batch, tokens, (2*k_heads+v_heads)*dim), dtype=np.float32), device="CPU")
+    q, k, v = conv.split([k_heads*dim, k_heads*dim, v_heads*dim], dim=-1)
+    q = (q.reshape(batch, tokens, k_heads, dim) *
+         (q.reshape(batch, tokens, k_heads, dim).square().sum(-1, keepdim=True) + 1e-6).rsqrt()).repeat(1, 1, v_heads//k_heads, 1)
+    k = (k.reshape(batch, tokens, k_heads, dim) *
+         (k.reshape(batch, tokens, k_heads, dim).square().sum(-1, keepdim=True) + 1e-6).rsqrt()).repeat(1, 1, v_heads//k_heads, 1)
+    expected = (q.transpose(1, 2) * dim**-0.5, k.transpose(1, 2), v.reshape(batch, tokens, v_heads, dim).transpose(1, 2))
+    for got, ref in zip(gdn_qkv(conv, k_heads, v_heads, dim), expected):
+      np.testing.assert_allclose(got.numpy(), ref.numpy(), rtol=2e-5, atol=2e-5)
+
+  def test_decode_rmsnorm_matches_reference(self):
+    rng = np.random.default_rng(5)
+    for rows in (1, 32, 128):
+      for dtype in (dtypes.float16, dtypes.float32):
+        for weight_dtype in (dtypes.float16, dtypes.float32):
+          with self.subTest(rows=rows, dtype=dtype, weight_dtype=weight_dtype):
+            norm = nn.RMSNorm(64, eps=1e-6)
+            norm.weight = Tensor(rng.standard_normal(64, dtype=np.float32), device="CPU").cast(weight_dtype).realize()
+            x = Tensor(rng.standard_normal((rows, 64), dtype=np.float32), device="CPU").cast(dtype).realize()
+            tol = 5e-4 if dtype == dtypes.float16 else 2e-6
+            np.testing.assert_allclose(rmsnorm(norm, x).numpy(), norm(x).numpy(), rtol=tol, atol=tol)
+
+  def test_causal_conv_silu_matches_reference(self):
+    rng = np.random.default_rng(15)
+    batch, tokens, channels, kernel = 2, 7, 64, 4
+    for dtype in (dtypes.float16, dtypes.float32):
+      for weight_dtype in (dtypes.float16, dtypes.float32):
+        if dtype == weight_dtype == dtypes.float16: continue
+        with self.subTest(dtype=dtype, weight_dtype=weight_dtype):
+          state = Tensor(rng.standard_normal((batch, kernel - 1, channels), dtype=np.float32), device="CPU").realize()
+          x = Tensor(rng.standard_normal((batch, tokens, channels), dtype=np.float32), device="CPU").cast(dtype).realize()
+          weight = Tensor(rng.standard_normal((channels, kernel), dtype=np.float32),
+                          device="CPU").cast(weight_dtype).realize()
+          window = state.cat(x, dim=1)
+          expected = functools.reduce(lambda a,b: a+b, (window[:, i:i+tokens] * weight[:, i] for i in range(kernel))).silu()
+          np.testing.assert_allclose(causal_conv_silu(state, x, weight).numpy(), expected.numpy(), rtol=2e-6, atol=2e-6)
+          np.testing.assert_allclose(causal_conv_silu(state, x, weight.T.contiguous()).numpy(), expected.numpy(), rtol=2e-6, atol=2e-6)
+
+    # This is the vectorized path used by Qwen's prefill.
+    state = Tensor(rng.standard_normal((1, kernel - 1, channels), dtype=np.float32), device="CPU").realize()
+    x = Tensor(rng.standard_normal((1, tokens, channels), dtype=np.float32), device="CPU").half().realize()
+    weight = Tensor(rng.standard_normal((channels, kernel), dtype=np.float32), device="CPU").half().realize()
+    expected = causal_conv_silu(state, x, weight).numpy()
+    np.testing.assert_allclose(causal_conv_silu(state, x, weight.T.contiguous()).numpy(), expected, rtol=2e-6, atol=2e-6)
+
+  def test_shared_gate_matches_reference(self):
+    rng = np.random.default_rng(6)
+    for dtype in (dtypes.float16, dtypes.float32):
+      with self.subTest(dtype=dtype):
+        x = Tensor(rng.standard_normal((3, 64), dtype=np.float32), device="CPU").cast(dtype).realize()
+        weight = Tensor(rng.standard_normal(64, dtype=np.float32), device="CPU").half().realize()
+        expected = (x * weight).sum(axis=-1, keepdim=True).sigmoid()
+        np.testing.assert_allclose(shared_gate(x, weight).numpy(), expected.numpy(), rtol=2e-5, atol=2e-5)
+
+  def test_silu_mul_matches_reference(self):
+    rng = np.random.default_rng(10)
+    for dtype in (dtypes.float16, dtypes.float32):
+      with self.subTest(dtype=dtype):
+        gate = Tensor(rng.standard_normal((2, 64), dtype=np.float32), device="CPU").cast(dtype).realize()
+        up = Tensor(rng.standard_normal((2, 64), dtype=np.float32), device="CPU").cast(dtype).realize()
+        np.testing.assert_equal(silu(gate).numpy(), gate.silu().numpy())
+        np.testing.assert_allclose(silu_mul(gate, up).numpy(), (gate.silu() * up).numpy(), rtol=2e-5, atol=2e-5)
+    gate = Tensor(rng.standard_normal((2, 64), dtype=np.float32), device="CPU").half().realize()
+    up = Tensor(rng.standard_normal((2, 64), dtype=np.float32), device="CPU").realize()
+    np.testing.assert_equal(silu_mul(gate, up).numpy(), (gate.silu() * up).numpy())
+    gate = Tensor(rng.standard_normal(4096, dtype=np.float32), device="CPU").half().realize()
+    up = Tensor(rng.standard_normal(4096, dtype=np.float32), device="CPU").realize()
+    np.testing.assert_allclose(silu_mul(gate, up).numpy(), (gate.silu() * up).numpy(), rtol=2e-6, atol=2e-6)
+
+  def test_biased_topk_matches_reference(self):
+    rng = np.random.default_rng(7)
+    logits = Tensor(rng.standard_normal((1, 2, 256), dtype=np.float32), device="CPU").half().realize()
+    bias = Tensor(rng.standard_normal(256, dtype=np.float32), device="CPU").half().realize()
+    probs = logits.sigmoid()
+    _, expected_sel = pairwise_topk(probs + bias, 8)
+    expected = probs.gather(-1, expected_sel)
+    expected = expected / expected.sum(axis=-1, keepdim=True)
+    got, got_sel = biased_sigmoid_topk(logits, bias, 8, normalize=True)
+    np.testing.assert_equal(got_sel.numpy(), expected_sel.numpy().reshape(2, 8))
+    np.testing.assert_allclose(got.numpy(), expected.numpy().reshape(2, 8), rtol=5e-4, atol=5e-4)
+
+  def test_packed_linear_matches_q8_activation_reference(self):
+    rng = np.random.default_rng(1)
+    for ggml_type, in_features in ((8, 64), (14, 256)):
+      for tokens in (1, 3, 8):
+        with self.subTest(ggml_type=ggml_type, tokens=tokens):
+          out_features = 7
+          raw = random_packed(rng, ggml_type, out_features * in_features)
+          layer = Linear(in_features, out_features, bias=False)
+          layer.set_quantized(Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), ggml_type)
+          x = rng.standard_normal((tokens, in_features), dtype=np.float32)
+          got = layer(Tensor(x, device="CPU")).numpy()
+          weight = ggml_data_to_tensor(Tensor(raw), out_features * in_features, ggml_type).numpy().reshape(out_features, in_features)
+          np.testing.assert_allclose(got, q8_activation(x) @ weight.T, rtol=1e-5, atol=5e-4)
+
+  def test_q8_linear_pair_matches_separate(self):
+    rng = np.random.default_rng(11)
+    in_features = 64
+    layers = []
+    for out_features in (7, 11):
+      raw = random_packed(rng, 8, out_features * in_features)
+      layer = Linear(in_features, out_features, bias=False)
+      layer.set_quantized(Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), 8)
+      layers.append(layer)
+    for dtype in (dtypes.float16, dtypes.float32):
+      with self.subTest(dtype=dtype):
+        x = Tensor(rng.standard_normal((1, in_features), dtype=np.float32), device="CPU").cast(dtype).realize()
+        got = q8_linear_pair(*layers, x)
+        for paired,layer in zip(got, layers): np.testing.assert_allclose(paired.numpy(), layer(x).numpy(), rtol=2e-5, atol=2e-5)
+        for paired,layer in zip(q8_linear_pair(*layers, x.reshape(1, 1, in_features)), layers):
+          self.assertEqual(paired.shape, (1, 1, layer.out_features))
+          np.testing.assert_allclose(paired.numpy(), layer(x).numpy().reshape(1, 1, -1), rtol=2e-5, atol=2e-5)
+        for layer in layers: layer.cpu_repacked = q8_repack(layer.weight, layer.out_features, layer.in_features).realize()
+        for repacked,original in zip(q8_linear_pair(*layers, x), got): np.testing.assert_equal(repacked.numpy(), original.numpy())
+        original_weights = [layer.weight for layer in layers]
+        for layer in layers: layer.weight = Tensor.zeros_like(layer.weight).realize()
+        for repacked,original in zip(uop_q8_linear_pair(*layers, x), got):
+          np.testing.assert_allclose(repacked.numpy(), original.numpy(), rtol=2e-6, atol=1e-5)
+        for layer,weight in zip(layers, original_weights): layer.weight = weight
+        for layer in layers: layer.cpu_repacked = None
+
+  def test_large_q8_uop_linear_repacked_matches_raw(self):
+    rng = np.random.default_rng(20)
+    in_features, out_features = 1024, 7
+    raw = random_packed(rng, 8, out_features * in_features)
+    layer = Linear(in_features, out_features, bias=False)
+    layer.set_quantized(Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), 8)
+    x = Tensor(rng.standard_normal((1, in_features), dtype=np.float32), device="CPU").realize()
+    expected = uop_linear(layer, x).numpy()
+    layer.cpu_repacked = q8_repack(layer.weight, out_features, in_features).realize()
+    np.testing.assert_allclose(uop_linear(layer, x).numpy(), expected, rtol=2e-6, atol=1e-5)
+
+  def test_q8_batched_pair_matches_separate(self):
+    rng = np.random.default_rng(15)
+    in_features = 64
+    layers = []
+    for out_features in (7, 11):
+      raw = random_packed(rng, 8, out_features * in_features)
+      layer = Linear(in_features, out_features, bias=False)
+      layer.set_quantized(Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), 8)
+      layers.append(layer)
+    for dtype in (dtypes.float16, dtypes.float32):
+      with self.subTest(dtype=dtype):
+        x = Tensor(rng.standard_normal((2, 3, in_features), dtype=np.float32), device="CPU").cast(dtype).realize()
+        got = q8_batched_pair(*layers, x)
+        for paired,layer in zip(got, layers): np.testing.assert_allclose(paired.numpy(), layer(x).numpy(), rtol=2e-5, atol=2e-5)
+        for layer in layers: layer.cpu_repacked = q8_repack(layer.weight, layer.out_features, layer.in_features).realize()
+        for repacked,original in zip(q8_batched_pair(*layers, x), got): np.testing.assert_equal(repacked.numpy(), original.numpy())
+        for layer in layers: layer.cpu_repacked = None
+
+  def test_q8_silu_linear_matches_separate(self):
+    rng = np.random.default_rng(16)
+    in_features, out_features = 64, 11
+    layer = Linear(in_features, out_features, bias=False)
+    layer.set_quantized(Tensor(random_packed(rng, 8, out_features * in_features), dtype=dtypes.uint8, device="CPU").realize(), 8)
+    gate = Tensor(rng.standard_normal((2, 4, in_features), dtype=np.float32), device="CPU").half().realize()
+    up = Tensor(rng.standard_normal(gate.shape, dtype=np.float32), device="CPU").realize()
+    fused = q8_silu_linear(layer, gate, up)
+    separate = layer(silu_mul(gate, up).half())
+    linear = layer(gate[:, 0])
+    single = q8_silu_linear(layer, gate[:1, :1], up[:1, :1])
+    np.testing.assert_equal(fused.numpy(), separate.numpy())
+    layer.cpu_repacked = q8_repack(layer.weight, layer.out_features, layer.in_features).realize()
+    np.testing.assert_equal(q8_silu_linear(layer, gate, up).numpy(), fused.numpy())
+    np.testing.assert_equal(q8_silu_linear(layer, gate[:1, :1], up[:1, :1]).numpy(), single.numpy())
+    np.testing.assert_equal(layer(gate[:, 0]).numpy(), linear.numpy())
+
+  def test_q8_gdn_projections_match_separate(self):
+    rng = np.random.default_rng(10)
+    in_features = 256
+    layers = []
+    for out_features in (1, 1):
+      raw = random_packed(rng, 8, out_features * in_features)
+      layer = Linear(in_features, out_features, bias=False)
+      layer.set_quantized(Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), 8)
+      layers.append(layer)
+    x = Tensor(rng.standard_normal((1, in_features), dtype=np.float32), device="CPU").half().realize()
+    weight = Tensor(rng.standard_normal((16, in_features), dtype=np.float32), device="CPU").half().realize()
+    got = q8_gdn_projections(*layers, weight, x)
+    expected = (*q8_linear_pair(*layers, x), x @ weight.T)
+    for fused,separate in zip(got, expected): np.testing.assert_allclose(fused.numpy(), separate.numpy(), rtol=1e-3, atol=1e-3)
+    for layer in layers: layer.cpu_repacked = q8_repack(layer.weight, layer.out_features, layer.in_features).realize()
+    for repacked,original in zip(q8_gdn_projections(*layers, weight, x), got):
+      np.testing.assert_equal(repacked.numpy(), original.numpy())
+    norm = nn.RMSNorm(in_features, eps=1e-6)
+    norm.weight = Tensor(rng.standard_normal(in_features, dtype=np.float32), device="CPU").half().realize()
+    raw_x = Tensor(rng.standard_normal((1, in_features), dtype=np.float32), device="CPU").realize()
+    expected = q8_gdn_projections(*layers, weight, rmsnorm(norm, raw_x).half())
+    for fused,separate in zip(q8_gdn_norm_projections(*layers, weight, raw_x, norm), expected):
+      np.testing.assert_equal(fused.numpy(), separate.numpy())
+
+  def test_f16_linear_matches_standard(self):
+    rng = np.random.default_rng(13)
+    layer = Linear(256, 37, bias=False)
+    layer.weight = Tensor(rng.standard_normal((37, 256), dtype=np.float32), device="CPU").half().realize()
+    for dtype in (dtypes.float16, dtypes.float32):
+      for tokens in (1, 3):
+        with self.subTest(dtype=dtype, tokens=tokens):
+          x = Tensor(rng.standard_normal((tokens, 256), dtype=np.float32), device="CPU").cast(dtype).realize()
+          np.testing.assert_allclose(f16_linear(layer, x).numpy(), layer(x).numpy(), rtol=1e-5, atol=2e-5)
+          np.testing.assert_allclose(f16_matvec(x, layer.weight).numpy(), layer(x).numpy(), rtol=1e-5, atol=2e-5)
+          np.testing.assert_allclose(uop_f16_matvec(x, layer.weight).numpy(), layer(x).numpy(), rtol=1e-5, atol=2e-5)
+    norm = nn.RMSNorm(256, eps=1e-6)
+    norm.weight = Tensor(rng.standard_normal(256, dtype=np.float32), device="CPU").half().realize()
+    x = Tensor(rng.standard_normal((1, 256), dtype=np.float32), device="CPU").realize()
+    normalized, out = rmsnorm_f16_linear(norm, layer, x)
+    expected = rmsnorm(norm, x)
+    np.testing.assert_allclose(normalized.numpy(), expected.numpy(), rtol=2e-6, atol=2e-6)
+    np.testing.assert_allclose(out.numpy(), f16_linear(layer, expected).numpy(), rtol=1e-5, atol=2e-5)
+
+  def test_q6_argmax_matches_materialized_logits(self):
+    rng = np.random.default_rng(8)
+    in_features, out_features = 256, 37
+    raw = random_packed(rng, 14, out_features * in_features)
+    weight = ggml_data_to_tensor(Tensor(raw), out_features * in_features, 14).numpy().reshape(out_features, in_features)
+    layer = Linear(in_features, out_features, bias=False)
+    layer.set_quantized(Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), 14)
+    for _ in range(3):
+      x = Tensor(rng.standard_normal((1, in_features), dtype=np.float32), device="CPU").realize()
+      expected = int(np.argmax(weight @ q8k_activation(x.numpy()).reshape(-1)))
+      self.assertEqual(q6_argmax(layer, x).item(), expected)
+
+  def test_packed_experts_match_reference(self):
+    rng = np.random.default_rng(2)
+    num_experts, in_features, out_features = 2, 256, 5
+    sel, x = np.array([1, 0, 1, 0], dtype=np.int32), rng.standard_normal((4, in_features), dtype=np.float32)
+    for ggml_type in (14, 21, 23):
+      with self.subTest(ggml_type=ggml_type):
+        raw = random_packed(rng, ggml_type, num_experts * out_features * in_features)
+        weight = ggml_data_to_tensor(Tensor(raw), num_experts * out_features * in_features,
+                                     ggml_type).numpy().reshape(num_experts, out_features, in_features)
+        experts = ExpertWeights(num_experts, in_features, out_features)
+        experts.set_quantized(Tensor(weight), Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), ggml_type)
+        got = experts(Tensor(sel, device="CPU"), Tensor(x, device="CPU")).numpy()
+        activation = q8k_activation(x) if ggml_type in (21, 23) else q8_activation(x)
+        expected = np.stack([activation[i] @ weight[expert].T for i,expert in enumerate(sel)])
+        np.testing.assert_allclose(got, expected, rtol=1e-5, atol=5e-4)
+        if ggml_type == 21:
+          experts.cpu_repacked = iq3_repack(experts.weight, num_experts * out_features, in_features).realize()
+          np.testing.assert_allclose(experts(Tensor(sel, device="CPU"), Tensor(x, device="CPU")).numpy(),
+                                     expected, rtol=1e-5, atol=5e-4)
+        direct_sel = np.array([1, 0], dtype=np.int32)
+        direct = experts(Tensor(direct_sel, device="CPU"), Tensor(x[:1], device="CPU")).numpy()
+        direct_activation = q8k_activation(x[:1]) if ggml_type in (21, 23) else q8_activation(x[:1])
+        direct_expected = np.stack([direct_activation[0] @ weight[expert].T for expert in direct_sel])
+        np.testing.assert_allclose(direct, direct_expected, rtol=1e-5, atol=5e-4)
+  def test_weighted_expert_sum_matches_reference(self):
+    rng = np.random.default_rng(16)
+    x = rng.standard_normal((2, 4, 257), dtype=np.float32)
+    probs = rng.random((2, 4), dtype=np.float32)
+    got = weighted_sum(Tensor(x, device="CPU"), Tensor(probs, device="CPU")).numpy()
+    np.testing.assert_allclose(got, (x * probs[..., None]).sum(axis=1), rtol=1e-6, atol=1e-6)
+
+  def test_quantized_expert_weighted_sum_matches_separate(self):
+    rng = np.random.default_rng(29)
+    num_experts, inputs, routes_per_input, in_features, out_features = 16, 4, 8, 256, 64
+    sel = Tensor(rng.integers(0, num_experts, (1, inputs, routes_per_input), dtype=np.int32), device="CPU").realize()
+    x = Tensor(rng.standard_normal((*sel.shape, in_features), dtype=np.float32), device="CPU").realize()
+    probs = Tensor(rng.random(sel.shape, dtype=np.float32), device="CPU").realize()
+    for ggml_type in (14, 23):
+      with self.subTest(ggml_type=ggml_type):
+        layer = ExpertWeights(num_experts, in_features, out_features)
+        layer.set_quantized(Tensor.empty(num_experts, out_features, in_features),
+                            Tensor(random_packed(rng, ggml_type, num_experts * in_features * out_features),
+                                   dtype=dtypes.uint8, device="CPU").realize(), ggml_type)
+        np.testing.assert_allclose(expert_weighted_sum(layer, sel, x, probs).numpy(), weighted_sum(layer(sel, x), probs).numpy(),
+                                   rtol=5e-6, atol=1e-5)
+
+  def test_fused_expert_silu_matches_separate(self):
+    rng = np.random.default_rng(12)
+    num_experts, in_features, out_features = 3, 256, 7
+    sel = Tensor(np.array([2, 0], dtype=np.int32), device="CPU")
+    x = Tensor(rng.standard_normal((1, in_features), dtype=np.float32), device="CPU").realize()
+    for ggml_type in (14, 21, 23):
+      with self.subTest(ggml_type=ggml_type):
+        experts = []
+        for _ in range(2):
+          raw = random_packed(rng, ggml_type, num_experts * out_features * in_features)
+          weight = ggml_data_to_tensor(Tensor(raw), num_experts * out_features * in_features,
+                                       ggml_type).reshape(num_experts, out_features, in_features)
+          expert = ExpertWeights(num_experts, in_features, out_features)
+          expert.set_quantized(weight, Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), ggml_type)
+          experts.append(expert)
+        gate, up = expert_pair(*experts, sel, x)
+        np.testing.assert_allclose(expert_silu(*experts, sel, x).numpy(), silu_mul(gate, up).numpy(), rtol=1e-4, atol=1e-4)
+        if ggml_type == 21:
+          direct_expected = expert_silu(*experts, sel, x).numpy()
+          batch_sel = Tensor(np.array([[2, 2], [1, 2]], dtype=np.int32), device="CPU")
+          batch_x = Tensor(rng.standard_normal((2, in_features), dtype=np.float32), device="CPU").realize()
+          expected = expert_silu(*experts, batch_sel, batch_x).numpy()
+          batch_gate, batch_up = expert_pair(*experts, batch_sel, batch_x)
+          np.testing.assert_allclose(expected, silu_mul(batch_gate, batch_up).numpy(), rtol=1e-4, atol=1e-4)
+          for expert in experts:
+            expert.cpu_repacked = iq3_repack(expert.weight, expert.num_experts * expert.out_features, expert.in_features).realize()
+          np.testing.assert_allclose(expert_silu(*experts, sel, x).numpy(), direct_expected, rtol=1e-5, atol=1e-4)
+          np.testing.assert_allclose(expert_silu(*experts, batch_sel, batch_x).numpy(), expected, rtol=1e-5, atol=1e-4)
+
+  def test_expert_silu_weighted_reuses_routes(self):
+    rng = np.random.default_rng(34)
+    num_experts, in_features, hidden, out_features = 3, 256, 256, 64
+    layers = []
+    for ggml_type,layer_in,layer_out in ((21, in_features, hidden), (21, in_features, hidden), (23, hidden, out_features)):
+      layer = ExpertWeights(num_experts, layer_in, layer_out)
+      raw = random_packed(rng, ggml_type, num_experts * layer_in * layer_out)
+      weight = ggml_data_to_tensor(Tensor(raw), num_experts * layer_in * layer_out, ggml_type).reshape(
+        num_experts, layer_out, layer_in)
+      layer.set_quantized(weight, Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), ggml_type)
+      if ggml_type == 21: layer.cpu_repacked = iq3_repack(layer.weight, num_experts * layer_out, layer_in).realize()
+      layers.append(layer)
+    sel = Tensor(np.array([[2, 0], [1, 2]], dtype=np.int32), device="CPU")
+    x = Tensor(rng.standard_normal((2, in_features), dtype=np.float32), device="CPU").realize()
+    probs = Tensor(rng.random(sel.shape, dtype=np.float32), device="CPU").realize()
+    expected = expert_weighted_sum(layers[2], sel, expert_silu(layers[0], layers[1], sel, x), probs).numpy()
+    got = uop_expert_silu_weighted(layers[0], layers[1], layers[2], sel, x, probs).numpy()
+    np.testing.assert_allclose(got, expected, rtol=5e-6, atol=1e-5)
+
+  def test_fused_moe_matches_separate_quantized_layers(self):
+    rng = np.random.default_rng(15)
+    dim = hidden = 256
+    config = TransformerConfig(1, dim, hidden, 1, 1, 1e-6, 32, dim, 1e6, dim, dim,
+                               num_experts=3, num_experts_per_tok=2, shared_expert_dim=hidden)
+    block = FFNBlock(config)
+    routed_weights = {}
+    for name,expert,ggml_type in (("gate", block.ffn_gate_exps, 21), ("up", block.ffn_up_exps, 21),
+                                  ("down", block.ffn_down_exps, 23)):
+      elements = expert.num_experts * expert.in_features * expert.out_features
+      raw = random_packed(rng, ggml_type, elements)
+      weight = ggml_data_to_tensor(Tensor(raw), elements, ggml_type).reshape(
+        expert.num_experts, expert.out_features, expert.in_features)
+      routed_weights[name] = weight.numpy().astype(np.float32)
+      expert.set_quantized(weight, Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), ggml_type)
+      if ggml_type == 21: expert.cpu_repacked = iq3_repack(expert.weight, expert.num_experts * expert.out_features, expert.in_features).realize()
+    for layer in (block.ffn_gate_shexp, block.ffn_up_shexp, block.ffn_down_shexp):
+      elements = layer.in_features * layer.out_features
+      raw = random_packed(rng, 8, elements)
+      layer.set_quantized(Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), 8)
+      layer.cpu_repacked = q8_repack(layer.weight, layer.out_features, layer.in_features).realize()
+    block.ffn_gate_inp_shexp["weight"] = Tensor(rng.standard_normal(dim, dtype=np.float32), device="CPU").half().realize()
+
+    x = Tensor(rng.standard_normal((1, 2, dim), dtype=np.float32), device="CPU").realize()
+    probs = Tensor(np.array([[[0.7, 0.3], [0.4, 0.6]]], dtype=np.float32), device="CPU").realize()
+    sel = Tensor(np.array([[[2, 0], [1, 2]]], dtype=np.int32), device="CPU").realize()
+    selected = sel.numpy().reshape(-1)
+    quantized_x = q8k_activation(x.numpy()).reshape(2, dim)
+    gate = np.stack([routed_weights["gate"][expert] @ quantized_x[route // 2] for route,expert in enumerate(selected)])
+    up = np.stack([routed_weights["up"][expert] @ quantized_x[route // 2] for route,expert in enumerate(selected)])
+    routed_hidden = silu_mul(Tensor(gate, device="CPU"), Tensor(up, device="CPU")).numpy()
+    quantized_hidden = q8k_activation(routed_hidden)
+    routed = np.stack([routed_weights["down"][expert] @ quantized_hidden[route] for route,expert in enumerate(selected)])
+    routed = Tensor((routed.reshape(2, 2, dim) * probs.numpy().reshape(2, 2, 1)).sum(1).reshape(1, 2, dim), device="CPU")
+    shared_gate_out, shared_up = block.ffn_gate_shexp(x), block.ffn_up_shexp(x)
+    shared = block.ffn_down_shexp(silu_mul(shared_gate_out, shared_up))
+    expected = routed + shared * shared_gate(x, block.ffn_gate_inp_shexp["weight"])
+    original = moe_ffn(block, x, probs, sel).numpy()
+    np.testing.assert_allclose(original, expected.numpy(), rtol=1e-4, atol=2e-2)
+    for expert in (block.ffn_gate_exps, block.ffn_up_exps):
+      expert.cpu_repacked = iq3_repack(expert.weight, expert.num_experts * expert.out_features, expert.in_features).realize()
+    np.testing.assert_equal(moe_ffn(block, x, probs, sel).numpy(), original)
+    np.testing.assert_allclose(uop_moe_ffn(block, x[:, :1], probs[:, :1], sel[:, :1]).numpy(), original[:, :1],
+                               rtol=1e-4, atol=2e-2)
+
+  def test_fused_moe_q6_down_matches_separate(self):
+    rng = np.random.default_rng(16)
+    dim = hidden = 256
+    config = TransformerConfig(1, dim, hidden, 1, 1, 1e-6, 32, dim, 1e6, dim, dim,
+                               num_experts=3, num_experts_per_tok=2, shared_expert_dim=hidden)
+    block = FFNBlock(config)
+    for expert,ggml_type in ((block.ffn_gate_exps, 21), (block.ffn_up_exps, 21), (block.ffn_down_exps, 14)):
+      elements = expert.num_experts * expert.in_features * expert.out_features
+      raw = random_packed(rng, ggml_type, elements)
+      weight = ggml_data_to_tensor(Tensor(raw), elements, ggml_type).reshape(
+        expert.num_experts, expert.out_features, expert.in_features)
+      expert.set_quantized(weight, Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), ggml_type)
+    for expert in (block.ffn_gate_exps, block.ffn_up_exps):
+      expert.cpu_repacked = iq3_repack(expert.weight, expert.num_experts * expert.out_features, expert.in_features).realize()
+    for layer in (block.ffn_gate_shexp, block.ffn_up_shexp, block.ffn_down_shexp):
+      elements = layer.in_features * layer.out_features
+      raw = random_packed(rng, 8, elements)
+      layer.set_quantized(Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), 8)
+    block.ffn_gate_inp_shexp["weight"] = Tensor(rng.standard_normal(dim, dtype=np.float32), device="CPU").half().realize()
+
+    x = Tensor(rng.standard_normal((1, 1, dim), dtype=np.float32), device="CPU").realize()
+    probs = Tensor(np.array([[[0.7, 0.3]]], dtype=np.float32), device="CPU").realize()
+    sel = Tensor(np.array([[[2, 0]]], dtype=np.int32), device="CPU").realize()
+    hidden = expert_silu(block.ffn_gate_exps, block.ffn_up_exps, sel, x.unsqueeze(2))
+    routed = weighted_sum(block.ffn_down_exps(sel, hidden), probs)
+    gate, up = block.ffn_gate_shexp(x), block.ffn_up_shexp(x)
+    shared = block.ffn_down_shexp(silu_mul(gate, up)) * shared_gate(x, block.ffn_gate_inp_shexp["weight"])
+    np.testing.assert_allclose(moe_ffn(block, x, probs, sel).numpy(), (routed + shared).numpy(), rtol=3e-4, atol=1e-4)
+
+if __name__ == "__main__":
+  unittest.main()

test/unit/test_llm_server.py

@@ -3,10 +3,26 @@ from unittest.mock import patch
 from tinygrad import Tensor, UOp
 from tinygrad.schedule import schedule_cache
 from tinygrad.llm.model import Transformer, TransformerConfig
+from tinygrad.llm.serve import StreamRouter
 
 TEST_CONFIG = TransformerConfig(num_blocks=1, dim=64, hidden_dim=128, n_heads=2, n_kv_heads=2,
                            norm_eps=1e-5, vocab_size=100, head_dim=32, rope_theta=10000.0, rope_dim=32, v_head_dim=32, max_context=32)
 
+class TestStreamRouter(unittest.TestCase):
+  @staticmethod
+  def route(router:StreamRouter, *pieces:str) -> dict[str, str]:
+    routed = [x for piece in pieces for x in router.route(piece)]
+    routed += list(router.route("", final=True))
+    return {field:"".join(text for f, text in routed if f == field) for field in {x[0] for x in routed}}
+
+  def test_generated_reasoning_tag(self):
+    self.assertEqual(self.route(StreamRouter(), "<thi", "nk>reason", "</thi", "nk>answer"),
+                     {"reasoning_content":"reason", "content":"answer"})
+
+  def test_prompt_opened_reasoning(self):
+    self.assertEqual(self.route(StreamRouter(reasoning=True), "reason", "</thi", "nk>answer"),
+                     {"reasoning_content":"reason", "content":"answer"})
+
 class TestTransformerGenerate(unittest.TestCase):
   def test_kv_cache_reuse(self):
     """Test that generate reuses the KV cache when tokens extend the cached prefix."""
@@ -151,8 +167,9 @@ class TestTransformerGenerate(unittest.TestCase):
     """Temperature from generate should be passed through to __call__."""
     model = Transformer(TEST_CONFIG)
     captured_temps = []
-    def mock_call(self, tokens, start_pos, temperature):
+    def mock_call(_self, tokens, start_pos, temperature, **kwargs):
       captured_temps.append(float(temperature.item()))
+      self.assertTrue(kwargs["sample"])
       return Tensor([[42]])
     with patch.object(Transformer, '__call__', mock_call):
       gen = model.generate([1, 2, 3], temperature=0.6)

tinygrad/codegen/__init__.py

@@ -120,8 +120,17 @@ pm_expand_broadcast = pm_wmma_add+PatternMatcher([
   (UPat(Ops.WMMA, name="b"), broadcast_and_devec_wmma),
 ])
 
-def do_devectorize(b:UOp):
-  if b.shape == (): return None
+@functools.cache
+def _uses_shrink_memory(x:UOp) -> bool:
+  if x.op in (Ops.LOAD, Ops.STORE): return x.src[0].op is Ops.SHRINK
+  if x.op is Ops.AFTER: return _uses_shrink_memory(x.src[0])
+  if x.op in GroupOp.Elementwise or x.op in (Ops.STACK, Ops.RESHAPE, Ops.PERMUTE):
+    return any(_uses_shrink_memory(y) for y in x.src)
+  return False
+
+def do_devectorize(ctx:Renderer|tuple[dict, Renderer, set[UOp]], b:UOp):
+  preserved = ctx[2] if isinstance(ctx, tuple) and len(ctx) == 3 else set()
+  if b.shape == () or b in preserved or _uses_shrink_memory(b): return None
   # broadcasting needs to be already unpacked, Invalid matches any dtype and shape
   if not all(x.shape == b.shape or x.base.arg is Invalid for x in b.src): return None
   src = []
@@ -187,6 +196,7 @@ def fix_group_for_reduce(x:UOp):
 @dataclass
 class ReduceContext:
   acc_num: int = 0
+  renderer: Renderer|None = None
 
 def merge_reduce_ends(sink:UOp):
   # merge ENDs that share the same range and nesting context (only those created by reduce_to_acc)
@@ -220,7 +230,8 @@ def reduce_ranges_to_acc(ctx:ReduceContext, r:UOp):
   acc_out = acc_initted.store(acc_initted.alu(r.arg[0], inp)).end(*r.src[1:]).rtag("mergeable")
   return acc.after(acc_out)
 
-def expand_horizontal_reduce(r:UOp):
+def expand_horizontal_reduce(ctx:ReduceContext, r:UOp):
+  if ctx.renderer is not None and ctx.renderer.has_native_reduce(r): return None
   inp = r.src[0]
   vals = [inp.index(*idx) for idx in itertools.product(*[range(inp.max_shape[a]) for a in range(r.arg[1])])]
   return functools.reduce(lambda x,y: x.alu(r.arg[0], y), vals)
@@ -317,7 +328,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   sink = graph_rewrite(sink, expander2, ctx=build_range_map(sink), name="expander")
 
   # remove reduce
-  sink = graph_rewrite(sink, mop_cleanup+pm_reduce_local, ctx=ReduceContext(), name="remove reduces")
+  sink = graph_rewrite(sink, mop_cleanup+pm_reduce_local, ctx=ReduceContext(renderer=ren), name="remove reduces")
 
   # add locals
   sink = graph_rewrite(sink, pm_add_local_buffers, ctx=itertools.count(0), name="add local buffers")
@@ -330,7 +341,12 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   sink = graph_rewrite(sink, symbolic_simple+pm_expand_broadcast+pm_add_loads, name="*** expand broadcast / add loads")
 
   # devectorize
-  sink = graph_rewrite(sink, symbolic_simple+devectorizer2, ctx=ren, name="devectorize2")
+  native_reduce_uops = {u for r in sink.toposort() if ren.has_native_reduce(r) for u in r.src[0].backward_slice_with_self}
+  if native_reduce_uops:
+    sink = graph_rewrite(sink, symbolic_simple, name="pre-devectorize symbolic")
+    native_reduce_uops = {u for r in sink.toposort() if ren.has_native_reduce(r) for u in r.src[0].backward_slice_with_self}
+  sink = graph_rewrite(sink, devectorizer2 if native_reduce_uops else symbolic_simple+devectorizer2,
+                       ctx=({}, ren, native_reduce_uops) if native_reduce_uops else ren, name="devectorize2")
 
   # simplify indexing
   sink = graph_rewrite(sink, indexing_simplify, name="simplify load/store indexing")
@@ -340,7 +356,12 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
 
   # do memory coalescing (late)
   sink = memory_coalescing(sink, ren)
-  sink = graph_rewrite(sink, symbolic_simple+ew_devectorizer+pm_simplify_add_image, name="add images", ctx=({}, ren), bottom_up=True)
+  native_reduce_uops = {u for r in sink.toposort() if ren.has_native_reduce(r) for u in r.src[0].backward_slice_with_self}
+  if native_reduce_uops:
+    sink = graph_rewrite(sink, symbolic_simple, name="pre-image symbolic", bottom_up=True)
+    native_reduce_uops = {u for r in sink.toposort() if ren.has_native_reduce(r) for u in r.src[0].backward_slice_with_self}
+  sink = graph_rewrite(sink, (ew_devectorizer if native_reduce_uops else symbolic_simple+ew_devectorizer)+pm_simplify_add_image,
+                       name="add images", ctx=({}, ren, native_reduce_uops), bottom_up=True)
 
   # extra symbolic before decomp. crashes without this?
   sink = graph_rewrite(sink, sym, name="extra symbolic")
@@ -468,7 +489,7 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
   if ast.op is Ops.PROGRAM: prg = ast
   elif ast.op is Ops.SINK:
     assert isinstance(ast.arg, KernelInfo), "requires KernelInfo on arg to to_program"
-    full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None)
+    full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None and ast.arg.optimize)
     prog_info = ProgramInfo.from_sink(full_sink, renderer.target)
     # instruction selection
     if isinstance(renderer, ISARenderer):
@@ -482,8 +503,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/codegen/late/coalesce.py

@@ -68,7 +68,7 @@ def image_valid_dims(base:DType, size:int, arch:str) -> list[tuple[int,int]]:
   return [(pxls//ALIGN//k, ALIGN*k) for k in range(ceildiv(pxls//ALIGN, MAXW), min(pxls//ALIGN, MAXW//ALIGN)+1) if (pxls//ALIGN)%k == 0]
 
 def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
-  shapes, ren = ctx
+  shapes, ren = ctx[:2]
   if not IMAGE or ren.target.device not in {"QCOM", "CL", "PYTHON", "NULL"}: return None
   valid, x = x.get_valid(), x.get_idx()
   # search for dims that drop the most valid statements
@@ -106,6 +106,7 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
     # TODO: this should handle images too, it's just memory coalescing
     if u.op in {Ops.LOAD, Ops.STORE}:
       assert len(u.src) == (2 if u.op is Ops.STORE else 1), "memory coalescing does not support gated loads/stores"
+      if u.src[0].op is Ops.SHRINK: continue
       assert u.src[0].op is Ops.INDEX, f"memory coalescing should be on INDEX, not {u.src[0].op}"
       buf, idx_u = u.src[0].src
       if buf.addrspace == AddrSpace.REG: continue

tinygrad/codegen/late/linearizer.py

@@ -9,7 +9,9 @@ def linearize(sink:UOp) -> list[UOp]:
   # this is a toposort with priority
   lst = list(sink.toposort())
   out_degree:defaultdict[UOp, int] = defaultdict(int)
-  priorities:dict[UOp, tuple[int, int, Any]] = {}
+  priorities:dict[UOp, tuple[int, int, int, Any]] = {}
+  wmma_depth:dict[UOp, int] = {}
+  for u in lst: wmma_depth[u] = max((wmma_depth[s] for s in u.src), default=0) + (u.op is Ops.WMMA)
 
   # get consumers and assign priorities
   # NOTE: this requires the lst be locally toposorted
@@ -30,7 +32,7 @@ def linearize(sink:UOp) -> list[UOp]:
       case Ops.RANGE: priority = 5    # placing RANGE is good
       case Ops.END: priority = -5     # placing END is bad
       case _: priority = 0            # everything else has priority 0
-    priorities[u] = (run_count, priority, extra)
+    priorities[u] = (run_count, priority, wmma_depth[u], extra)
 
   # number the uops in "ideal" order
   nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: priorities[x]+(x.tuplize if TUPLE_ORDER else ())))}

tinygrad/device.py

@@ -318,6 +318,7 @@ class TinyELF:
   target: Target
   # tuple of (name, slot, dtype, shape)
   signature: tuple[tuple[str|None, int, DType, tuple], ...]
+  parallel: bool = False
 
   @staticmethod
   def iter_sig(signature:tuple[tuple[str|None, int, DType, tuple], ...], offset:int=0) -> Generator[tuple[int, DType], None, None]:

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, multiprocessing
+from concurrent.futures import ProcessPoolExecutor
 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 PARALLEL_COMPILE, NUM_CPU_THREADS
 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,17 @@ 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, NUM_CPU_THREADS.value, len(pending))
+      with ProcessPoolExecutor(workers, mp_context=multiprocessing.get_context("spawn")) as pool:
+        for key,program in pool.map(parallel_to_program, pending.values()): to_program_cache[key] = program
   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

@@ -1,5 +1,5 @@
 from __future__ import annotations
-import sys, argparse, codecs, itertools, typing, re, unicodedata, json, time
+import sys, os, argparse, codecs, itertools, typing, re, unicodedata, json, time
 from typing import TYPE_CHECKING
 from tinygrad import nn
 from tinygrad.uop.ops import UOp, Ops
@@ -112,7 +112,7 @@ class FallbackTemplate:
     if self.tok.preset == 'glm4': return ""
     if self.tok.preset == 'tekken': return "[/INST]"
     return self.tok.decode([self.tok.eos_id])
-  def render(self, messages:list[dict], tools=None, add_generation_prompt:bool=True) -> str:
+  def render(self, messages:list[dict], tools=None, add_generation_prompt:bool=True, **kwargs) -> str:
     out = self.tok.decode([] if self.tok.bos_id is None else [self.tok.bos_id]) + ("<sop>" if self.tok.preset == 'glm4' else "")
     for msg in messages:
       out += self.role(msg["role"])
@@ -136,6 +136,7 @@ def main():
   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()
+  if args.warmup or args.serve: os.environ.setdefault("DISK_SCACHE", "1")
 
   # load the model
   model, kv = Transformer.from_gguf(fetch(models.get(args.model, args.model)), args.max_context)
@@ -163,9 +164,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/gguf.py

@@ -1,9 +1,11 @@
-import functools, io, pathlib, re, struct
+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.
@@ -20,7 +22,25 @@ _GGML_NATIVE = {0: dtypes.float32, 1: dtypes.float16, 24: dtypes.int8, 25: dtype
 _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:
+_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
+  packed, ggml_type = meta
+  return Tensor(packed), ggml_type
+
+def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int, contiguous:bool=True) -> Tensor:
   """
   Converts ggml tensor data to a tinygrad tensor.
 
@@ -42,7 +62,8 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
 
   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()
+    blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1]))
+    if contiguous: blocks = blocks.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 ])
@@ -130,8 +151,7 @@ readers: dict[int, Callable[[io.BufferedIOBase], Any]] = { 8: read_str, 9: read_
 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()
+  tensor = tensor.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!")
@@ -145,7 +165,14 @@ def _gguf_parse(tensor: Tensor) -> tuple[dict, dict[str, Tensor]]:
   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}
+  state_dict = {}
+  for name, dims, typ, off in t_infos:
+    n, shape = prod(dims), tuple(reversed(dims))
+    decoded = ggml_data_to_tensor(data:=tensor[data_start + off:], n, typ).reshape(*shape)
+    if typ in _GGML_QUANT:
+      block_size, type_size = _GGML_QUANT[typ]
+      _quantized_tensors[decoded.uop] = (data[:n//block_size*type_size].uop, typ)
+    state_dict[name] = decoded
   return kv_data, state_dict
 
 def _gguf_split_paths(path: pathlib.Path, kv: dict) -> list[pathlib.Path]:
@@ -169,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)))
+  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))[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/kernels/__init__.py

@@ -0,0 +1 @@
+"""Custom kernels used by tinygrad.llm."""

tinygrad/llm/kernels/generic.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+import functools
+from tinygrad import Tensor, UOp, dtypes
+from tinygrad.dtype import AddrSpace
+from tinygrad.uop.ops import AxisType, KernelInfo
+
+def _topk_256_sort(values:UOp, indices:UOp, next_values:UOp, next_indices:UOp, ready:UOp, lane:UOp) -> tuple[UOp, UOp, UOp]:
+  for size in (2, 4, 8, 16, 32, 64, 128, 256):
+    for stride in (128, 64, 32, 16, 8, 4, 2, 1)[9-size.bit_length():]:
+      partner = lane ^ stride
+      a_value, b_value = values.after(ready)[lane], values.after(ready)[partner]
+      a_index, b_index = indices.after(ready)[lane], indices.after(ready)[partner]
+      a_first = (a_value < b_value) | (a_value.eq(b_value) & (a_index > b_index))
+      want_first = (lane & stride).eq(0).eq((lane & size).eq(0))
+      take_a = want_first.eq(a_first)
+      ready = UOp.group(next_values.after(ready)[lane].store(take_a.where(a_value, b_value)),
+                        next_indices.after(ready)[lane].store(take_a.where(a_index, b_index))).barrier()
+      values, next_values, indices, next_indices = next_values, values, next_indices, indices
+  return values, indices, ready
+
+@functools.cache
+def _topk_256_kernel(out:UOp, sel:UOp, x:UOp, k:int, softmax:bool=False) -> UOp:
+  outer, lane = UOp.range(out.shape[0], 0), UOp.range(256, 1, axis_type=AxisType.LOCAL)
+  values = UOp.placeholder((256,), x.dtype, 0, addrspace=AddrSpace.LOCAL)
+  indices = UOp.placeholder((256,), dtypes.int32, 1, addrspace=AddrSpace.LOCAL)
+  next_values = UOp.placeholder((256,), x.dtype, 2, addrspace=AddrSpace.LOCAL)
+  next_indices = UOp.placeholder((256,), dtypes.int32, 3, addrspace=AddrSpace.LOCAL)
+  ready = UOp.group(values.after(outer)[lane].store(x[outer, lane]),
+                    indices.after(outer)[lane].store(lane.int())).barrier()
+  values, indices, ready = _topk_256_sort(values, indices, next_values, next_indices, ready, lane)
+  valid = lane < k
+  src = 256 - k + lane
+  value = values.after(ready)[src]
+  if softmax:
+    max_value = values.after(ready)[255]
+    value = (value - max_value).exp() / sum(((values.after(ready)[256-k+i] - max_value).exp() for i in range(k)),
+                                            UOp.const(x.dtype, 0))
+  stores = (out[outer, lane.valid(valid)].store(value),
+            sel[outer, lane.valid(valid)].store(indices.after(ready)[src]))
+  return UOp.group(*stores).end(outer, lane).sink(arg=KernelInfo(name=f"topk_256{'_softmax' if softmax else ''}", opts_to_apply=()))
+
+def topk_256(x:Tensor, k:int, softmax:bool=False) -> tuple[Tensor, Tensor]:
+  outer = int(x.numel()) // 256
+  values = Tensor.empty(outer, k, dtype=x.dtype, device=x.device)
+  indices = Tensor.empty(outer, k, dtype=dtypes.int32, device=x.device)
+  values, indices = Tensor.custom_kernel(values, indices, x.reshape(outer, 256),
+    fxn=lambda out,sel,x:_topk_256_kernel(out, sel, x, k, softmax))[:2]
+  return values.reshape(*x.shape[:-1], k), indices.reshape(*x.shape[:-1], k)
+
+@functools.cache
+def _inverse_unit_lower_kernel(out:UOp, x:UOp, n:int) -> UOp:
+  outer_count = 1
+  for dim in out.shape[:-2]:
+    assert isinstance(dim, int)
+    outer_count *= dim
+  outer, lane = UOp.range(outer_count, 0), UOp.range(n, 1, axis_type=AxisType.LOCAL)
+  raw = UOp.placeholder((n*n,), x.dtype, 0, addrspace=AddrSpace.LOCAL)
+  solved = UOp.placeholder((n*n,), x.dtype, 1, addrspace=AddrSpace.LOCAL)
+  ready = UOp.group(*(raw[row*n+lane].store(x.flatten()[outer*n*n+row*n+lane]) for row in range(n))).barrier()
+  for row in range(n):
+    base, previous = raw.after(ready), solved.after(ready)
+    value = base[row*n+lane] + sum((base[row*n+i] * previous[i*n+lane] for i in range(row)), UOp.const(x.dtype, 0))
+    ready = solved.after(ready)[row*n+lane].store((lane < row).where(value, UOp.const(x.dtype, 0))).barrier()
+  result = solved.after(ready)
+  stores = [out.flatten()[outer*n*n+row*n+lane].store(lane.eq(row).where(UOp.const(x.dtype, 1), result[row*n+lane]))
+            for row in range(n)]
+  return UOp.group(*stores).end(outer, lane).sink(arg=KernelInfo(name="inverse_unit_lower", opts_to_apply=()))
+
+def inverse_unit_lower(x:Tensor) -> Tensor:
+  """Reference-ordered inverse of I-x for a strictly lower-triangular x."""
+  n = x.shape[-1]
+  assert isinstance(n, int)
+  if n == 64 and str(x.device).startswith("AMD"):
+    out = Tensor.empty(*x.shape, dtype=x.dtype, device=x.device)
+    return Tensor.custom_kernel(out, x, fxn=lambda out,x:_inverse_unit_lower_kernel(out, x, n))[0]
+  rows = [x[..., 0, :].const_like(0)]
+  for i in range(1, n):
+    prefix = x[..., i, :i]
+    previous = Tensor.stack(*rows, dim=-2)[..., :, :i]
+    rows.append((prefix + (prefix.unsqueeze(-1) * previous).sum(-2)).pad((0, n-i)))
+  return Tensor.stack(*rows, dim=-2) + Tensor.eye(n, dtype=x.dtype).to(x.device)

tinygrad/llm/serve.py

@@ -34,9 +34,9 @@ def normalize_messages(messages:list[dict]) -> None:
 
 class StreamRouter:
   # routes streamed output text to (field, text) deltas, keeping tool_call regions in .buf for the final parse
-  def __init__(self):
+  def __init__(self, reasoning:bool=False):
     self.buf = ""
-    self.mode = "undecided"  # output inside a think block is sent as reasoning_content
+    self.mode = "reasoning" if reasoning else "undecided"  # output inside a think block is sent as reasoning_content
   def split(self, tag:str, final:bool) -> tuple[str, bool]:
     # split buf on the first full tag, holding back a partial tag at the end unless final
     if tag in self.buf:
@@ -66,47 +66,58 @@ class Handler(HTTPRequestHandler):
   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):
+  def run_model(self, ids:list[int], model_name:str, include_usage=False, max_tokens:int|None=None, temperature:float=0.0,
+                reasoning:bool=False):
     model, tok = self.server.model, self.server.tok
     prompt_tokens = len(ids)
     cache_start_pos = model.get_start_pos(ids)
     stderr_log(f"in:{colored(f'{cache_start_pos:5d}', 'green')} +{len(ids)-cache_start_pos:5d}  {colored('--', 'BLACK')}  ")
     tmpl = {"id":f"chatcmpl-{uuid.uuid4().hex[:24]}", "object":"chat.completion.chunk", "created":int(time.time()), "model":model_name}
     def chunk(d:dict): return {"choices": [{"index":0, "delta":d, "finish_reason":None}], **tmpl}
-    yield chunk({"role":"assistant", "content":""})
     out: list[int] = []
     finish_reason = "stop"
-    st = time.perf_counter()
+    st = pt = time.perf_counter()
     dec = tok.stream_decoder()
-    router = StreamRouter()
-    for next_id in model.generate(ids, temperature=temperature):
-      if len(out) == 0: stderr_log(f"prefill:{(prompt_tokens-cache_start_pos)/((pt:=time.perf_counter())-st):4.0f} tok/s  {colored('--', 'BLACK')}  ")
-      if tok.is_end(next_id): break
-      out.append(next_id)
-      for field, delta in router.route(dec(next_id)): yield chunk({field:delta})
-      if max_tokens is not None and len(out) >= max_tokens:
-        finish_reason = "length"
-        break
-    for field, delta in router.route(dec(), final=True): yield chunk({field:delta})
-    tool_calls: list[dict] = []
-    for m in re.finditer(r"<tool_call>\s*(.*?)\s*(?:</tool_call>|$)", router.buf, re.DOTALL):
-      if (parsed := parse_tool_call(m.group(1))) is None:
-        stderr_log(f"failed to parse tool call: {m.group(1)[:200]}")
-        yield chunk({"content":m.group(0)})  # don't silently drop output the client can't use
-      else:
-        name, args = parsed
-        tool_calls.append({"index":len(tool_calls), "id":f"call_{uuid.uuid4().hex[:24]}", "type":"function",
-                           "function":{"name":name, "arguments":args if isinstance(args, str) else json.dumps(args)}})
-    if tool_calls:
-      yield chunk({"tool_calls":tool_calls})
-      if finish_reason == "stop": finish_reason = "tool_calls"
-    yield {"choices": [{"index":0, "delta":{},"finish_reason":finish_reason}], **tmpl}
-    if include_usage:
-      yield {"choices": [], "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": len(out),
-                                      "total_tokens": prompt_tokens + len(out)}, **tmpl}
-    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")
+    router = StreamRouter(reasoning)
+    def log_stats(interrupted:bool=False):
+      et = time.perf_counter()
+      total = f"total:{et-st:6.2f}s"
+      stderr_log(f"gen:{len(out)/(et-pt) if len(out) > 1 else 0:4.0f} tok/s  {colored('--', 'BLACK')}  "
+                 f"out:{len(out):5d}  {colored('--', 'BLACK')}  {colored(total, 'red') if interrupted else total}\n")
+    completed = False
+    try:
+      yield chunk({"role":"assistant", "content":""})
+      for next_id in model.generate(ids, temperature=temperature):
+        if len(out) == 0:
+          stderr_log(f"prefill:{(prompt_tokens-cache_start_pos)/((pt:=time.perf_counter())-st):4.0f} tok/s  {colored('--', 'BLACK')}  ")
+        if tok.is_end(next_id): break
+        out.append(next_id)
+        for field, delta in router.route(dec(next_id)): yield chunk({field:delta})
+        if max_tokens is not None and len(out) >= max_tokens:
+          finish_reason = "length"
+          break
+      for field, delta in router.route(dec(), final=True): yield chunk({field:delta})
+      tool_calls: list[dict] = []
+      for m in re.finditer(r"<tool_call>\s*(.*?)\s*(?:</tool_call>|$)", router.buf, re.DOTALL):
+        if (parsed := parse_tool_call(m.group(1))) is None:
+          stderr_log(f"failed to parse tool call: {m.group(1)[:200]}")
+          yield chunk({"content":m.group(0)})  # don't silently drop output the client can't use
+        else:
+          name, args = parsed
+          tool_calls.append({"index":len(tool_calls), "id":f"call_{uuid.uuid4().hex[:24]}", "type":"function",
+                             "function":{"name":name, "arguments":args if isinstance(args, str) else json.dumps(args)}})
+      if tool_calls:
+        yield chunk({"tool_calls":tool_calls})
+        if finish_reason == "stop": finish_reason = "tool_calls"
+      completed = True
+      yield {"choices": [{"index":0, "delta":{},"finish_reason":finish_reason}], **tmpl}
+      if include_usage:
+        yield {"choices": [], "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": len(out),
+                                        "total_tokens": prompt_tokens + len(out)}, **tmpl}
+      log_stats()
+    except GeneratorExit:
+      if not completed: log_stats(interrupted=True)
+      raise
 
   def do_POST(self):
     request_st = time.perf_counter()
@@ -117,7 +128,10 @@ class Handler(HTTPRequestHandler):
     if self.path == "/v1/chat/completions":
       # render and tokenize
       normalize_messages(body["messages"])
-      rendered = self.server.template.render(messages=body["messages"], tools=body.get("tools"), add_generation_prompt=True)
+      template_kwargs = body.get("chat_template_kwargs") or {}
+      if not isinstance(template_kwargs, dict): raise TypeError("chat_template_kwargs must be an object")
+      rendered = self.server.template.render(**(template_kwargs | {
+        "messages":body["messages"], "tools":body.get("tools"), "add_generation_prompt":True}))
       ids: list[int] = self.server.tok.encode(rendered)
       stderr_log(f"prep:{(time.perf_counter()-request_st)*1e3:5.0f} ms  {colored('--', 'BLACK')}  ")
       if len(ids) >= self.server.model.max_context:
@@ -129,7 +143,8 @@ class Handler(HTTPRequestHandler):
       # 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)))
+                              max_tokens=max_tokens, temperature=float(body.get("temperature", 0.0)),
+                              reasoning=rendered.rstrip().endswith("<think>"))
       if body.get("stream"): self.stream_json(chunks)
       else:
         out, reasoning, tool_calls, finish_reason = [], [], [], "stop"

tinygrad/renderer/__init__.py

@@ -61,6 +61,7 @@ class Renderer:
   suffix: str = ""
   # TODO: make this generic with a list of supported types
   supports_float4: bool = True
+  def has_native_reduce(self, x:UOp) -> bool: return False
   has_local: bool = True
   has_threads: bool = False
   has_shared: bool = True

tinygrad/renderer/cstyle.py

@@ -259,6 +259,22 @@ class ClangRenderer(CStyleLanguage):
   gep_arr_threshold = 0
   has_local = False
   has_threads = bool(getenv("THREADS", 1))
+  @staticmethod
+  def _byte_dot_reduce(x:UOp) -> tuple[UOp, UOp, UOp|None]|None:
+    if x.op is not Ops.REDUCE or x.arg != (Ops.ADD, 1) or len(x.src) != 1: return None
+    permute = x.src[0]
+    if permute.op is not Ops.PERMUTE or permute.arg != (1, 0) or permute.src[0].op is not Ops.RESHAPE: return None
+    product, scale = permute.src[0].src[0], None
+    if product.op is Ops.MUL:
+      for base, candidate_scale in (product.src, product.src[::-1]):
+        if base.op is Ops.MUL and all(s.op is Ops.CAST and s.dtype is dtypes.int32 and s.src[0].dtype is dtypes.int8 for s in base.src):
+          product, scale = base, candidate_scale
+          break
+    if product.op is not Ops.MUL or not all(s.op is Ops.CAST and s.dtype is dtypes.int32 and s.src[0].dtype is dtypes.int8 for s in product.src):
+      return None
+    a, b = (s.src[0] for s in product.src)
+    return (a, b, scale) if a.shape == b.shape and a.shape in ((16,), (32,)) else None
+  def has_native_reduce(self, x:UOp) -> bool: return self._byte_dot_reduce(x) is not None
   global_max = (NUM_CPU_THREADS.value, 0, 0)
   infinity = "__builtin_inff()"
   nan = '__builtin_nanf("")'
@@ -270,6 +286,100 @@ class ClangRenderer(CStyleLanguage):
                  Ops.SQRT: lambda x,dtype: f"__builtin_sqrt({x})" if dtype == dtypes.float64 else f"__builtin_sqrtf({x})",
                  Ops.TRUNC: lambda x,dtype: f"__builtin_trunc({x})" if dtype == dtypes.float64 else f"__builtin_truncf({x})",
                  Ops.FDIV: lambda a,b,dtype: f"({a}/{b})"}
+  @staticmethod
+  def _render_byte_dot(ctx, a:UOp, b:UOp, scale:UOp|None=None):
+    if a.shape != b.shape or a.shape not in ((16,), (32,)): return None
+    lanes, suffix = a.shape[0] // 2, "128" if a.shape[0] == 16 else "256"
+    if scale is None: multiplier = "1"
+    elif scale.max_numel() == 1: multiplier = ctx[scale]
+    elif scale.op is Ops.STACK and all(x is scale.src[0] for x in scale.src): multiplier = ctx[scale.src[0]]
+    else: multiplier = f"({ctx[scale]})[0]"
+    return f"__builtin_ia32_pmaddwd{suffix}(__builtin_bit_cast(short __attribute__((ext_vector_type({lanes}))), " \
+           f"__builtin_ia32_pmaddubsw{suffix}(__builtin_elementwise_abs({ctx[a]}), __builtin_ia32_psignb{suffix}({ctx[b]}, {ctx[a]}))), " \
+           f"(short __attribute__((ext_vector_type({lanes})))){{{','.join([multiplier] * lanes)}}})"
+  @staticmethod
+  def _uniform_vector_const(x:UOp):
+    return x.src[0].arg if x.op is Ops.STACK and x.src and all(v.op is Ops.CONST and v.arg == x.src[0].arg for v in x.src) else None
+  @staticmethod
+  def _render_lut_lookup(ctx, x:UOp):
+    values:dict[int, int] = {}
+    index:UOp|None = None
+    node = x
+    while node.op is Ops.WHERE and node.src[0].op in (Ops.CMPEQ, Ops.CMPNE):
+      cond, true_value, false_value = node.src
+      left, right = cond.src
+      key = ClangRenderer._uniform_vector_const(right)
+      if key is None: left, right, key = right, left, ClangRenderer._uniform_vector_const(left)
+      if not isinstance(key, int) or not 0 <= key < 16 or (index is not None and left is not index): return None
+      index = left
+      selected, node = (true_value, false_value) if cond.op is Ops.CMPEQ else (false_value, true_value)
+      value = ClangRenderer._uniform_vector_const(selected)
+      if not isinstance(value, int): return None
+      values[key] = value
+    default = ClangRenderer._uniform_vector_const(node)
+    if index is None or not isinstance(default, int) or index.shape not in ((16,), (32,)): return None
+    lut_values = tuple(values.get(i, default) for i in range(16))
+    lanes, suffix = index.shape[0], "128" if index.shape == (16,) else "256"
+    charv = f"signed char __attribute__((ext_vector_type({lanes})))"
+    lut = f"({charv}){{{','.join(map(str, lut_values * (lanes // 16)))}}}"
+    return f"__builtin_ia32_pshufb{suffix}({lut}, __builtin_bit_cast({charv}, {ctx[index]}))"
+  @staticmethod
+  def _render_contiguous_stack(ctx, x:UOp):
+    if not x.src or not all(v.op is Ops.INDEX and len(v.src) == 2 and v.src[0] is x.src[0] and
+                            v.src[1].op is Ops.CONST for v in x.src): return None
+    indices = tuple(v.src[1].arg for v in x.src)
+    if indices != tuple(range(indices[0], indices[0] + len(indices))): return None
+    return f"__builtin_shufflevector({ctx[x.src[0]]}, {ctx[x.src[0]]}, {','.join(map(str, indices))})"
+  @staticmethod
+  def _render_concat_stack(ctx, x:UOp):
+    if len(x.src) != 2 or x.src[0].shape != x.src[1].shape or len(x.src[0].shape) != 1: return None
+    lanes = x.src[0].shape[0]
+    return f"__builtin_shufflevector({ctx[x.src[0]]}, {ctx[x.src[1]]}, {','.join(map(str, range(lanes*2)))})"
+  @staticmethod
+  def _render_vector_permute(ctx, x:UOp):
+    old_shape, order = x.src[0].shape, x.arg
+    new_shape = tuple(old_shape[i] for i in order)
+    indices = []
+    for flat in range(x.max_numel()):
+      coord, rem = [], flat
+      for size in reversed(new_shape):
+        coord.append(rem % size)
+        rem //= size
+      new_coord = tuple(reversed(coord))
+      old_coord = tuple(new_coord[order.index(i)] for i in range(len(order)))
+      old_flat = sum(c * math.prod(old_shape[i+1:]) for i,c in enumerate(old_coord))
+      indices.append(old_flat)
+    return f"__builtin_shufflevector({ctx[x.src[0]]}, {ctx[x.src[0]]}, {','.join(map(str, indices))})"
+  @staticmethod
+  def _render_vector_load(ctx, x:UOp, address:UOp):
+    vec_type = ctx._render_dtype(x.dtype, x.max_numel(), AddrSpace.REG)
+    # GGUF block fields are not necessarily vector aligned. memcpy expresses an unaligned load without
+    # aliasing or alignment assumptions and Clang folds the fixed-size copy to a single vector load.
+    return f"({{ {vec_type} _v; __builtin_memcpy(&_v, {ctx[address]}, sizeof(_v)); _v; }})"
+  @staticmethod
+  def _render_vector_store(ctx, address:UOp, value:UOp):
+    vec_type = ctx._render_dtype(value.dtype, value.max_numel(), AddrSpace.REG)
+    # Explicit vector stores can target unaligned packed data or register-backed arrays.
+    return f"do {{ {vec_type} _v = {ctx[value]}; __builtin_memcpy({ctx[address]}, &_v, sizeof(_v)); }} while (0);"
+  string_rewrite = PatternMatcher([
+    (UPat(Ops.WHERE, name="x"), lambda ctx,x: ClangRenderer._render_lut_lookup(ctx, x)),
+    (UPat(Ops.MUL, dtypes.int32, src=(UPat(Ops.REDUCE, name="dot"), UPat.var("scale"))),
+     lambda ctx,dot,scale: ClangRenderer._render_byte_dot(ctx, pair[0], pair[1], scale)
+     if (pair:=ClangRenderer._byte_dot_reduce(dot)) is not None and pair[2] is None and scale.vmin >= -32768 and scale.vmax <= 32767 else None),
+    (UPat(Ops.REDUCE, name="x"), lambda ctx,x: ClangRenderer._render_byte_dot(ctx, *pair)
+     if (pair:=ClangRenderer._byte_dot_reduce(x)) is not None else None),
+    (UPat(Ops.CAST, name="x"), lambda ctx,x: f"__builtin_convertvector({ctx[x.src[0]]}, {ctx.render_type(x)})"
+     if x.max_numel() > 1 else None),
+    (UPat(Ops.RESHAPE, name="x"), lambda ctx,x: ctx[x.src[0]]),
+    (UPat(Ops.PERMUTE, name="x"), lambda ctx,x: ClangRenderer._render_vector_permute(ctx, x)),
+    (UPat(Ops.STACK, name="x"), lambda ctx,x: ClangRenderer._render_concat_stack(ctx, x)),
+    (UPat(Ops.STACK, name="x"), lambda ctx,x: ClangRenderer._render_contiguous_stack(ctx, x)),
+    (UPat(Ops.STORE, src=(UPat(Ops.SHRINK, name="address"), UPat.var("value")), name="x"),
+     lambda ctx,x,address,value: ClangRenderer._render_vector_store(ctx, address, value)
+     if value.max_numel() > 1 else None),
+    (UPat(Ops.LOAD, src=(UPat(Ops.SHRINK, name="address"),), name="x"),
+     lambda ctx,x,address: ClangRenderer._render_vector_load(ctx, x, address)),
+  ]) + base_rewrite
 
   # LLVM legalizes double => half/bf16 cast on systems that don't support it natively (like x86 cpus without AVX512-FP16) into a compiler-rt libcall.
   # there is also no native bfl16 <-> fp16 conversion on those CPUs

tinygrad/runtime/ops_cpu.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
-import platform, sys, os, ctypes, functools, mmap, threading, array, itertools
+import platform, sys, os, ctypes, functools, mmap, threading, array, itertools, pathlib
 from dataclasses import replace
 from typing import cast
-from tinygrad.helpers import to_mv, OSX, WIN, Context, mv_address, suppress_finalizing, unwrap, data64_le, partition
+from tinygrad.helpers import to_mv, OSX, WIN, Context, mv_address, suppress_finalizing, unwrap, data64_le, partition, getenv
 from tinygrad.device import Buffer, BufferSpec, TinyELF
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, HCQArgsState, HCQSignal, HCQProgram, MMIOInterface
 from tinygrad.runtime.support.hcq import CLikeArgsState
@@ -17,7 +17,9 @@ from tinygrad import UOp, dtypes
 from tinygrad.dtype import AddrSpace
 from tinygrad.uop.ops import sint, KernelInfo, Ops, UPat, PatternMatcher, graph_rewrite
 
-MAX_ARGS, CMD_SIZE, RING_SLOTS = 63, 64, (16 << 10)
+MAX_ARGS, CMD_SIZE, RING_SLOTS = 32, 33, (16 << 10)
+CPU_CORES = getenv("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)))
+PARALLEL_WORKERS, PARALLEL_PARTICIPANTS = min(31, max(1, CPU_CORES-1)), min(32, max(2, CPU_CORES))
 
 def signal_prog():
   val = UOp.param(1, dtypes.int, (), vmin_vmax=(0, dtypes.int.max), name="value", addrspace=AddrSpace.ALU)
@@ -43,9 +45,16 @@ def quit_prog():
   close = fn[2].load().call(sem[0], ret_dtype=dtypes.void) # sem_close(sem)
   return fn.after(close)[0].load().call(UOp.const(dtypes.uint64, 0), ret_dtype=dtypes.void) # pthread_exit(0)
 
+def post_many_prog():
+  sem, post_fn = UOp.param(0, dtypes.uint64, (1,)), UOp.param(1, dtypes.uint64, (1,))
+  count = UOp.param(2, dtypes.int, (), vmin_vmax=(1, RING_SLOTS), name="count", addrspace=AddrSpace.ALU)
+  post = UOp.range(count, 0)
+  return post_fn.after(post)[0].load().call(sem.after(post)[0], ret_dtype=dtypes.void).end(post)
+
 def worker_prog():
   ring = UOp.param(0, dtypes.uint64, (RING_SLOTS * CMD_SIZE,), volatile=True)
   wait, sem = UOp.param(1, dtypes.uint64, (1,), volatile=True), UOp.param(2, dtypes.uint64, (1,))
+  progress = UOp.param(3, dtypes.uint64, (1,), volatile=True)
   cur = UOp.range(2**64-1, 0, dtype=dtypes.uint64)
 
   # spin on windows, sem_wait to sleep on posix
@@ -53,7 +62,73 @@ def worker_prog():
   else: ready = wait.after(cur)[0].load().call(sem.after(cur)[0], ret_dtype=dtypes.void)
 
   entry = [ring.after(ready).index((cur % RING_SLOTS) * CMD_SIZE + i).load() for i in range(CMD_SIZE)]
-  return entry[0].call(*entry[1:], ret_dtype=dtypes.void).end(cur)
+  return progress.after(entry[0].call(*entry[1:], ret_dtype=dtypes.void), cur)[0].store(cur + 1).end(cur)
+
+def parallel_wait_prog():
+  generation = UOp.param(0, dtypes.uint64, (1,), volatile=True)
+  completed = UOp.param(1, dtypes.uint64, (1,), volatile=True)
+  ready = UOp.param(2, dtypes.uint64, (1,), volatile=True)
+  wait_fn, sem = UOp.param(3, dtypes.uint64, (1,)), UOp.param(4, dtypes.uint64, (1,))
+  relax = "__builtin_ia32_pause();" if platform.machine().lower() in ("x86_64", "amd64") else \
+          '__asm__ __volatile__("yield");' if platform.machine().lower() in ("aarch64", "arm64") else ""
+  value = UOp(Ops.CUSTOMI, dtypes.uint64, (generation.index(0), completed.index(0), wait_fn[0].load(), sem.index(0)), arg=
+    "({{ unsigned long _seen = *((volatile unsigned long *){1}), _v; "
+    "do {{ int _i = 0; do {{ _v = *((volatile unsigned long *){0}); if (_v > _seen) break; __RELAX__ }} while (++_i < __SPIN__); "
+    "if (_v <= _seen) while (((int (*)(unsigned long)){2})((unsigned long){3}) != 0) {{}}; }} while (_v <= _seen); _v; }})"
+    .replace("__RELAX__", relax).replace("__SPIN__", str(getenv("CPU_UOP_SPIN", 10000000))))
+  return ready[0].store(value)
+
+def parallel_worker_prog():
+  generation = UOp.param(0, dtypes.uint64, (1,), volatile=True)
+  group_count = UOp.param(1, dtypes.uint64, (1,), volatile=True)
+  ring_addr = UOp.param(2, dtypes.uint64, (1,), volatile=True)
+  completed = UOp.param(3, dtypes.uint64, (PARALLEL_WORKERS,), volatile=True)
+  worker_id = UOp.param(4, dtypes.int, (), vmin_vmax=(1, PARALLEL_WORKERS), name="worker_id", addrspace=AddrSpace.ALU)
+  wait_fn, sem = UOp.param(5, dtypes.uint64, (1,)), UOp.param(6, dtypes.uint64, (1,))
+  ready_values, helper_fn = UOp.param(7, dtypes.uint64, (PARALLEL_WORKERS,), volatile=True), UOp.param(8, dtypes.uint64, (1,))
+  cur = UOp.loop(0)
+  worker_idx = worker_id-1
+  ready = helper_fn.after(cur)[0].load().call(generation.after(cur).index(0), completed.after(cur).index(worker_idx),
+                                               ready_values.after(cur).index(worker_idx), wait_fn.after(cur).index(0),
+                                               sem.after(cur).index(0), ret_dtype=dtypes.void)
+  seen = ready_values.after(ready)[worker_idx].load()
+  worker = worker_id.cast(dtypes.uint64)
+  count = group_count.after(seen)[0].load()
+  work_count = (worker < count).where((count - worker + PARALLEL_PARTICIPANTS - 1) // PARALLEL_PARTICIPANTS, 0)
+  work = UOp.range(work_count, 2, dtype=dtypes.uint64)
+  command = worker + work * UOp.const(dtypes.uint64, PARALLEL_PARTICIPANTS)
+  address = ring_addr.after(seen)[0].load()
+  entry = [UOp(Ops.CUSTOM, dtypes.uint64, (address, command * CMD_SIZE + i),
+               arg="*((volatile unsigned long *){0} + {1})") for i in range(CMD_SIZE)]
+  finished = entry[0].call(*entry[1:], ret_dtype=dtypes.void).end(work)
+  return completed.after(finished, seen)[worker_id-1].store(seen).end(cur, count.ne(0))
+
+def parallel_dispatch_prog():
+  commands = UOp.param(0, dtypes.uint64, (RING_SLOTS * CMD_SIZE,), volatile=True)
+  generation = UOp.param(1, dtypes.uint64, (1,), volatile=True)
+  group_count = UOp.param(2, dtypes.uint64, (1,), volatile=True)
+  ring_addr = UOp.param(3, dtypes.uint64, (1,), volatile=True)
+  completed = UOp.param(4, dtypes.uint64, (PARALLEL_WORKERS,), volatile=True)
+  count = UOp.param(5, dtypes.int, (), vmin_vmax=(1, 2**31-1), name="count", addrspace=AddrSpace.ALU)
+  workers = count.minimum(PARALLEL_PARTICIPANTS) - 1
+  post_fn, sems = UOp.param(6, dtypes.uint64, (1,)), UOp.param(7, dtypes.uint64, (PARALLEL_WORKERS,))
+  address = UOp(Ops.CUSTOM, dtypes.uint64, (commands.index(0),), arg="(unsigned long){0}")
+  publish = UOp.group(ring_addr[0].store(address), group_count[0].store(count.cast(dtypes.uint64)))
+  current = generation.after(publish)[0].load()
+  next_generation = current + 1
+  signal = generation[0].store(next_generation)
+  if WIN: wake = signal
+  else:
+    wake_worker = UOp.range(workers, 3)
+    wake = post_fn.after(signal)[0].load().call(sems[wake_worker].load(), ret_dtype=dtypes.void).end(wake_worker)
+  work = UOp.range((count + PARALLEL_PARTICIPANTS - 1) // PARALLEL_PARTICIPANTS, 2)
+  command = work * PARALLEL_PARTICIPANTS
+  entry = [commands.after(wake).index(command * CMD_SIZE + i).load() for i in range(CMD_SIZE)]
+  own_done = entry[0].call(*entry[1:], ret_dtype=dtypes.void).end(work)
+  worker = UOp.range(workers, 0)
+  wait = UOp.loop(1)
+  done = completed.after(own_done, wait).index(worker).load()
+  return done.end(wait, done < next_generation).end(worker).sink(arg=KernelInfo("parallel_dispatch_prog"), tag=1)
 
 def host_wait(ctx, dst:UOp, val:UOp) -> UOp:
   return (cur:=dst.after(loop:=UOp.loop(next(ctx))).index(UOp.const(dtypes.int, 0)).load()).end(loop, cur < val)
@@ -70,28 +145,74 @@ class CPUComputeQueue(HWQueue):
   def __init__(self, dev):
     super().__init__()
     self.dev = dev
+    self._encoded:array.array|None = None
+    self._exec_groups:list[tuple[int, int, bool]] = []
   def _cmd(self, prog, args=(), vals=()): return self.exec(prg:=self.dev.prgs[prog], prg.fill_kernargs(args, vals), None, None)
   def memory_barrier(self): return self
   def exec(self, prg:CPUProgram, args_state:HCQArgsState, global_size, local_size):
     if (lvp:=isinstance(args_state, LVPArgsState)): self.bind_args_state(args_state)
     args:list[sint|None] = [args_state.buf.va_addr] if lvp else [*[x.va_addr for x in args_state.bufs], *args_state.vals]
-    assert len(args) <= MAX_ARGS, f"CPU programs support at most {MAX_ARGS} arguments, got {len(args)}"
+    assert len(args) <= MAX_ARGS, f"CPU program {prg.name!r} supports at most {MAX_ARGS} arguments, got {len(args)}"
+    start = len(self._q) // CMD_SIZE
     for tid in range(1 if lvp else (global_size or (1,))[0]):
       if not lvp and 'core_id' in prg.runtimevars: args[prg.runtimevars['core_id']] = tid
       self.q(prg, *[unwrap(x) for x in args], *([0] * (MAX_ARGS - len(args))))
+    self._exec_groups.append((start, len(self._q) // CMD_SIZE, prg.parallel))
     return self
   def wait(self, signal, value=0): return self._cmd(wait_prog, (signal.base_buf,), (value,))
   def timestamp(self, signal): return self._cmd(timestamp_prog, (signal.base_buf.offset(8, 8), self.dev.func_table._buf.offset(0, 8)))
   def signal(self, signal, value:sint=0): return self._cmd(signal_prog, (signal.base_buf,), (value,))
   def _submit(self, dev):
     dev.ensure_worker()
+    if self._encoded is None:
+      self._encoded = array.array('Q', ((x.addr if i % CMD_SIZE == 0 else int(x)) & ((1<<64)-1) for i,x in enumerate(self._q)))
+    else:
+      for off, _ in self.q_sints: self._encoded[off] = int(self._q[off]) & ((1<<64)-1)
+    encoded = self._encoded
+    parallel_buf = None
+    if dev.parallel_uops and not dev.parallel_shutdown:
+      completed = dev.progress_view[0]
+      still_inflight = []
+      for end_pos,buf in dev.parallel_command_inflight:
+        if end_pos <= completed: dev.parallel_command_pool.append(buf)
+        else: still_inflight.append((end_pos, buf))
+      dev.parallel_command_inflight = still_inflight
+      parallel_words = sum((end-start) * CMD_SIZE for start,end,parallel in self._exec_groups if parallel and end-start > 1)
+      if parallel_words:
+        required_bytes = parallel_words * 8
+        parallel_buf = next((buf for buf in dev.parallel_command_pool if buf.nbytes >= required_bytes), None)
+        if parallel_buf is not None: dev.parallel_command_pool.remove(parallel_buf)
+        else: parallel_buf = Buffer(dev.device, required_bytes, dtypes.uint8, preallocate=True)
+        parallel_view = parallel_buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
+        transformed, parallel_offset = array.array('Q'), 0
+        dispatch = dev.prgs[parallel_dispatch_prog]
+        for start,end,parallel in self._exec_groups:
+          if not parallel or end-start == 1:
+            transformed.extend(encoded[start*CMD_SIZE:end*CMD_SIZE])
+            continue
+          words = (end-start) * CMD_SIZE
+          parallel_view[parallel_offset:parallel_offset+words] = encoded[start*CMD_SIZE:end*CMD_SIZE]
+          args = [parallel_buf._buf.va_addr + parallel_offset * 8, dev.parallel_state._buf.va_addr,
+                  dev.parallel_state._buf.va_addr + 8, dev.parallel_state._buf.va_addr + 16,
+                  dev.parallel_state._buf.va_addr + 24, end-start, dev.func_table._buf.va_addr + 32,
+                  dev.parallel_sems._buf.va_addr]
+          transformed.extend(array.array('Q', [dispatch.addr, *args, *([0] * (MAX_ARGS-len(args)))]))
+          parallel_offset += words
+        encoded = transformed
     ring_view = dev.ring.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
-    for off in range(0, len(self._q), CMD_SIZE):
-      entry = [self._q[off].addr, *self._q[off+1:off+CMD_SIZE]]
-      ring_view[(base:=(dev.ring_pos % RING_SLOTS) * CMD_SIZE):base+CMD_SIZE] = array.array('Q', (int(x) & ((1<<64)-1) for x in entry))
-      dev.ring_pos += 1
+    cmds, submitted = len(encoded) // CMD_SIZE, 0
+    while submitted < cmds:
+      consumed = dev.progress_view[0]
+      if (available:=RING_SLOTS - (dev.ring_pos - consumed)) == 0: continue
+      start = dev.ring_pos % RING_SLOTS
+      count = min(cmds - submitted, available, RING_SLOTS - start)
+      src = submitted * CMD_SIZE
+      ring_view[start*CMD_SIZE:(start+count)*CMD_SIZE] = encoded[src:src+count*CMD_SIZE]
+      dev.ring_pos += count
       if WIN: dev.sys.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = dev.ring_pos
-      else: assert libc.sem_post(dev.sem) == 0
+      else: dev.post_many(dev.sem_addr, dev.func_table._buf.va_addr + 32, count)
+      submitted += count
+    if parallel_buf is not None: dev.parallel_command_inflight.append((dev.ring_pos, parallel_buf))
 
 class LVPArgsState(CLikeArgsState):
   def __init__(self, buf, prg, bufs, vals=()): super().__init__(buf, prg, bufs, vals, [*data64_le(buf.va_addr + 12), (len(bufs) + len(vals)) * 2])
@@ -106,6 +227,7 @@ class CPUProgram(HCQProgram['CPUDevice']):
 
   def __init__(self, dev:CPUDevice, obj:TinyELF):
     self.signature, self.runtimevars = obj.signature, {name:slot for name,slot,*_ in obj.signature if name == 'core_id'}
+    self.parallel = obj.parallel
 
     LVP = obj.target.renderer == "LVP"
     if sys.platform == "win32": # mypy doesn't understand when WIN is used here
@@ -189,32 +311,103 @@ class CPUDevice(HCQCompiled):
 
     # TODO: move to hcq2
     with Context(EMULATED_DTYPES="", TRACK_MATCH_STATS=0):
-      prgs = {f: f().sink(arg=KernelInfo(f.__name__), tag=1) for f in (signal_prog, wait_prog, timestamp_prog, quit_prog, worker_prog)}
-      self.prgs = {f: self.runtime(do_to_program(v, ClangRenderer(replace(self.renderer.target, renderer="CLANG"))).to_elf()) for f,v in prgs.items()}
+      helpers = (signal_prog, wait_prog, timestamp_prog, quit_prog, worker_prog, post_many_prog,
+                 parallel_wait_prog, parallel_worker_prog, parallel_dispatch_prog)
+      prgs = {f: f().sink(arg=KernelInfo(f.__name__), tag=1) for f in helpers}
+      renderer = ClangRenderer(replace(self.renderer.target, renderer="CLANG"))
+      self.prgs = {f: self.runtime(do_to_program(v, renderer).to_elf()) for f,v in prgs.items()}
+    if not WIN:
+      self.post_many = ctypes.CFUNCTYPE(None, ctypes.c_uint64, ctypes.c_uint64, ctypes.c_int)(cast(CPUProgram, self.prgs[post_many_prog]).addr)
+
+    self.worker:threading.Thread|None = None
+    self.parallel_uops = bool(getenv("CPU_PARALLEL_UOPS", 1)) and not WIN
+    self.parallel_shutdown = False
+    self._physical_affinity = False
 
   @functools.cached_property
   def ring(self) -> Buffer: return Buffer(self.device, RING_SLOTS * CMD_SIZE, dtypes.uint64, preallocate=True)
   @functools.cached_property
   def sys(self) -> Buffer: return Buffer(self.device, 1, dtypes.uint64, preallocate=True)
   @functools.cached_property
+  def progress(self) -> Buffer: return Buffer(self.device, 1, dtypes.uint64, preallocate=True)
+  @functools.cached_property
   def sem_buf(self) -> Buffer: return Buffer(self.device, 1, dtypes.uint8, options=BufferSpec(external_ptr=self.sem_addr), preallocate=True)
 
   # TODO: move to hcq2 infra
   @functools.cached_property
   def func_table(self) -> Buffer:
-    fns = ([0, ctypes.windll.kernel32.ExitThread, 0, 0] if WIN else  # type: ignore[attr-defined]
-           [libc.dll.clock_gettime, libc.dll.pthread_exit, libc.dll.sem_wait, libc.dll.sem_close])
+    fns = ([0, ctypes.windll.kernel32.ExitThread, 0, 0, 0] if WIN else  # type: ignore[attr-defined]
+           [libc.dll.clock_gettime, libc.dll.pthread_exit, libc.dll.sem_wait, libc.dll.sem_close, libc.dll.sem_post])
     addrs = array.array('Q', [unwrap(ctypes.cast(f, ctypes.c_void_p).value) if f else 0 for f in fns])
     (ft:=Buffer(self.device, len(fns), dtypes.uint64, preallocate=True)).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[:] = addrs
     return ft
 
   @functools.cache
   def ensure_worker(self):
-    threading.Thread(target=cast(CPUProgram, self.prgs[worker_prog]).fxn, daemon=True, args=[ctypes.c_uint64(x) for x in
-      [self.ring._buf.va_addr, self.sys._buf.va_addr if WIN else self.func_table._buf.va_addr+16, self.sem_addr]]).start()
+    self.progress_view = self.progress.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
+    self.worker = threading.Thread(target=cast(CPUProgram, self.prgs[worker_prog]).fxn, daemon=True, args=[ctypes.c_uint64(x) for x in
+      [self.ring._buf.va_addr, self.sys._buf.va_addr if WIN else self.func_table._buf.va_addr+16, self.sem_addr,
+       self.progress._buf.va_addr]])
+    self.worker.start()
+
+    if self.parallel_uops:
+      self.parallel_state = Buffer(self.device, 3 + 2 * PARALLEL_WORKERS, dtypes.uint64, preallocate=True)
+      self.parallel_command_pool:list[Buffer] = []
+      self.parallel_command_inflight:list[tuple[int, Buffer]] = []
+      self.parallel_sems = Buffer(self.device, PARALLEL_WORKERS, dtypes.uint64, preallocate=True)
+      parallel_sem_addrs, self.parallel_sem_handles = [], []
+      for i in range(PARALLEL_WORKERS):
+        sem = libc.sem_open(sem_name:=f"/tinygrad-{os.getpid()}-{id(self):x}-p{i}".encode(),
+                            os.O_CREAT|os.O_EXCL, 0o600, 0)  # type: ignore[call-arg]
+        if (sem_addr:=unwrap(ctypes.cast(sem, ctypes.c_void_p).value)) == ctypes.c_void_p(-1).value or libc.sem_unlink(sem_name):
+          raise OSError(ctypes.get_errno(), "parallel semaphore")
+        self.parallel_sem_handles.append(sem)
+        parallel_sem_addrs.append(sem_addr)
+      self.parallel_sems.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[:] = array.array('Q', parallel_sem_addrs)
+      state_view = self.parallel_state.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
+      self.parallel_generation_view, self.parallel_count_view = state_view[0:1], state_view[1:2]
+      self.parallel_ring_addr_view = state_view[2:3]
+      self.parallel_completed_view = state_view[3:3+PARALLEL_WORKERS]
+      self.parallel_helper_fn = Buffer(self.device, 1, dtypes.uint64, preallocate=True)
+      self.parallel_helper_fn.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = cast(CPUProgram, self.prgs[parallel_wait_prog]).addr
+      parallel_ready_addr = self.parallel_state._buf.va_addr + (3 + PARALLEL_WORKERS) * 8
+      self.parallel_workers = [threading.Thread(target=cast(CPUProgram, self.prgs[parallel_worker_prog]).fxn,
+        args=(ctypes.c_uint64(self.parallel_state._buf.va_addr), ctypes.c_uint64(self.parallel_state._buf.va_addr + 8),
+              ctypes.c_uint64(self.parallel_state._buf.va_addr + 16), ctypes.c_uint64(self.parallel_state._buf.va_addr + 24),
+              ctypes.c_uint64(i+1), ctypes.c_uint64(self.func_table._buf.va_addr + 16),
+              ctypes.c_uint64(parallel_sem_addrs[i]), ctypes.c_uint64(parallel_ready_addr),
+              ctypes.c_uint64(self.parallel_helper_fn._buf.va_addr)),
+        daemon=True) for i in range(PARALLEL_WORKERS)]
+      for worker in self.parallel_workers: worker.start()
+
+  def use_physical_cores(self):
+    self.ensure_worker()
+    if self._physical_affinity or not sys.platform.startswith("linux") or self.worker is None or self.worker.native_id is None: return
+    allowed = os.sched_getaffinity(self.worker.native_id)
+    physical:dict[tuple[str, str], int] = {}
+    try:
+      for cpu in sorted(allowed):
+        topology = pathlib.Path(f"/sys/devices/system/cpu/cpu{cpu}/topology")
+        key = ((topology / "physical_package_id").read_text().strip(), (topology / "core_id").read_text().strip())
+        physical.setdefault(key, cpu)
+    except (OSError, ValueError): return
+    if physical:
+      cpus = tuple(physical.values())
+      os.sched_setaffinity(self.worker.native_id, set(cpus))
+      for i, worker in enumerate(getattr(self, "parallel_workers", ())):
+        if worker.native_id is not None: os.sched_setaffinity(worker.native_id, {cpus[(i+1) % len(cpus)]})
+    self._physical_affinity = True
 
   def finalize(self):
-    if self.ring_pos == 0: return # the worker starts with the first submit
+    if self.worker is None: return
+    self.synchronize()
+    if self.parallel_uops:
+      self.parallel_shutdown = True
+      self.parallel_count_view[0] = 0
+      self.parallel_generation_view[0] += 1
+      for sem in self.parallel_sem_handles: assert libc.sem_post(sem) == 0
+      for worker in self.parallel_workers: worker.join()
+      for sem in self.parallel_sem_handles: assert libc.sem_close(sem) == 0
     ft = self.func_table._buf
     CPUComputeQueue(self)._cmd(quit_prog, (ft.offset(8, 8),) if WIN else (ft.offset(8, 24), self.sem_buf._buf)).submit(self)
-    self.ring_pos = 0
+    self.worker = None

tinygrad/runtime/support/system.py

@@ -269,7 +269,8 @@ class PCIIfaceBase:
     if should_use_sysmem:
       vaddr = self.dev_impl.mm.alloc_vaddr(size:=round_up(size, mmap.PAGESIZE), align=mmap.PAGESIZE)
       memview, paddrs = self.pci_dev.alloc_sysmem(size, vaddr=vaddr, contiguous=contiguous)
-      mapping = self.dev_impl.mm.map_range(vaddr, size, [(paddr, 0x1000) for paddr in paddrs], aspace=AddrSpace.SYS, snooped=True, uncached=True)
+      mapping = self.dev_impl.mm.map_range(vaddr, size, [(paddr, 0x1000) for paddr in paddrs], aspace=AddrSpace.SYS,
+                                           snooped=True, uncached=uncached)
       return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(mapping, has_cpu_mapping=True, hMemory=paddrs[0]), view=memview, owner=self.dev)
 
     mapping = self.dev_impl.mm.valloc(size:=round_up(size, 0x1000), uncached=uncached, contiguous=cpu_access)

tinygrad/schedule/__init__.py

@@ -82,7 +82,7 @@ def create_schedule(sched_sink:UOp) -> UOp:
 from tinygrad.schedule.memory import memory_plan_rewrite
 from tinygrad.engine.realize import capturing, pm_flatten_linear
 from tinygrad.schedule.rangeify import get_kernel_graph
-from tinygrad.helpers import CAPTURING
+from tinygrad.helpers import CAPTURING, diskcache_get, diskcache_put, getenv
 from tinygrad.uop.ops import PatternMatcher, UPat, ParamArg
 from tinygrad.dtype import AddrSpace
 
@@ -110,14 +110,19 @@ def lower_sink_to_linear(function:UOp) -> UOp|None:
   st = time.perf_counter()
   if isinstance(function.arg, KernelInfo): return None
   cache_key = function.key
-  if not SCACHE or (sc_ret:=schedule_cache.get(cache_key, None)) is None:
+  disk_scache = bool(getenv("DISK_SCACHE", 0))
+  sc_ret = schedule_cache.get(cache_key, None) if SCACHE else None
+  if sc_ret is None and SCACHE and disk_scache: sc_ret = diskcache_get("schedule", {"key":cache_key})
+  if sc_ret is None:
     if SPEC: type_verify(function, spec_tensor)
     # support recursive CALLs
     linear = create_schedule(get_kernel_graph(function))
-    if SCACHE: schedule_cache[cache_key] = linear
+    if SCACHE:
+      schedule_cache[cache_key] = linear
+      if disk_scache: diskcache_put("schedule", {"key":cache_key}, linear)
   else:
     # schedule cache hit
-    linear = sc_ret
+    schedule_cache[cache_key] = linear = sc_ret
   if (DEBUG >= 1 and len(linear.src) > 1) or DEBUG >= 3:
     for frm in inspect.stack():
       if frm.filename == "<string>": continue

tinygrad/uop/ops.py

@@ -1180,7 +1180,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     assert self.op is Ops.PROGRAM and isinstance(self.arg, ProgramInfo), "to_elf should only be called on a PROGRAM ast"
     sig = tuple((u.arg.name, u.arg.slot, u.dtype, u._shape)
                 for u in tuple(filter(lambda u: u.op is Ops.PARAM and u.addrspace != AddrSpace.ALU, self.src[1].src)) + self.arg.vars)
-    return TinyELF(self.src[3].arg, self.arg.function_name, self.arg.target, sig)
+    return TinyELF(self.src[3].arg, self.arg.function_name, self.arg.target, sig, self.arg.parallel)
 
 @dataclass(frozen=True)
 class KernelInfo:
@@ -1191,6 +1191,8 @@ class KernelInfo:
   opts_to_apply: tuple|None = None
   estimates: Estimates|None = None
   beam: int = 0
+  optimize: bool = True
+  parallel: bool = False
   @property
   def function_name(self): return to_function_name(self.name)
 
@@ -1204,6 +1206,7 @@ class ProgramInfo:
   outs: tuple[int, ...] = ()
   ins: tuple[int, ...] = ()
   target: Target = Target()
+  parallel: bool = False
 
   @property
   def function_name(self): return to_function_name(self.name)
@@ -1241,7 +1244,8 @@ class ProgramInfo:
       if u.op is Ops.PARAM and u in _vars and u.expr == 'core_id': global_size[0] = int(u.vmax) + 1
     return ProgramInfo(sink.arg.name if isinstance(sink.arg, KernelInfo) else "test", tuple(global_size),
                        tuple(local_size) if local_size is not None else None, tuple(sorted(dedup(_vars), key=lambda v: v.arg.slot)),
-                       tuple(sorted(dedup(_globals))), tuple(sorted(dedup(outs))), tuple(sorted(dedup(ins))), target)
+                       tuple(sorted(dedup(_globals))), tuple(sorted(dedup(outs))), tuple(sorted(dedup(ins))),
+                       target, sink.arg.parallel if isinstance(sink.arg, KernelInfo) else False)
 
 @dataclass(frozen=True)
 class CallInfo:

tinygrad/uop/spec.py

@@ -207,6 +207,9 @@ spec_program = PatternMatcher([
   # allow special SHRINK
   (UPat(Ops.SHRINK, src=(UPat((Ops.PARAM, Ops.BUFFER, Ops.AFTER)), UPat(), UPat(Ops.CONST))), lambda: True),
 
+  (UPat((Ops.RESHAPE, Ops.PERMUTE), name="x"), lambda x: x.addrspace is AddrSpace.ALU),
+  (UPat(Ops.REDUCE, src=(UPat(name="x"),), arg=(Ops.ADD, 1)), lambda x: x.addrspace is AddrSpace.ALU),
+
   # movement ops are not allowed in programs
   (UPat(GroupOp.Movement), lambda: False),
 

tinygrad/viz/serve.py

@@ -37,7 +37,7 @@ class HTTPRequestHandler(BaseHTTPRequestHandler):
         self.wfile.flush()
       self.wfile.write("data: [DONE]\n\n".encode("utf-8"))
     # pass if client closed connection
-    except (BrokenPipeError, ConnectionResetError): return
+    except (BrokenPipeError, ConnectionResetError): source.close()
 
 from tinygrad.uop.ops import TrackedGraphRewrite, RewriteTrace, UOp, Ops, GroupOp, srender, sint, sym_infer, range_str, range_start, multirange_str
 from tinygrad.uop.ops import KernelInfo