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/fp8/fp8_linear.py

@@ -18,9 +18,9 @@ def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
   SEQ = inp.shape[1]
   OUT = weight.shape[0]
   IN = weight.shape[-1]
-  seq_idx = UOp.range(SEQ, 2, AxisType.LOOP)
-  out_idx = UOp.range(OUT, 3, AxisType.LOOP)
-  batch_idx = UOp.range(output.size//SEQ//OUT, 1, AxisType.LOOP)
+  seq_idx = UOp.range(SEQ, 2)
+  out_idx = UOp.range(OUT, 3)
+  batch_idx = UOp.range(output.size//SEQ//OUT, 1)
   reduce_idx = UOp.range(IN, 0, AxisType.REDUCE)
   product = (inp.index((seq_idx*IN+reduce_idx+batch_idx*IN*SEQ)) * weight.index((out_idx*IN+reduce_idx))).cast(dtypes.float)
   reduced = product.reduce(reduce_idx, arg=Ops.ADD)

extra/gemm/amd_copy_matmul.py

@@ -70,8 +70,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
 
   if use_wmma:
     k = UOp.range(BLOCK_K // WMMA_K, 101, AxisType.REDUCE)
-    tile_m = UOp.range(TM // WMMA_ACC, 200, AxisType.LOOP)
-    tile_n = UOp.range(TN, 201, AxisType.LOOP)
+    tile_m = UOp.range(TM // WMMA_ACC, 200)
+    tile_n = UOp.range(TN, 201)
 
     acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0,2,1)[tile_m, tile_n]
     a_frag = A_local.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_K // WMMA_K, WMMA_K)[wave_m, tile_m, lane_n, k]

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, AxisType.LOOP)
-  tn1 = UOp.range(TN, 201, AxisType.LOOP)
-  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, AxisType.LOOP)
-  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, AxisType.LOOP)
-  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, AxisType.LOOP)
-  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, AxisType.LOOP)
-  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, AxisType.LOOP)
-  tn2 = UOp.range(TD, 402, AxisType.LOOP)
-  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!")

extra/gemm/amd_uop_matmul.py

@@ -28,10 +28,10 @@ REG_TILES_PER_WAVE_M = BLOCK_M // (WAVES_PER_BLOCK_M * LANES_PER_WAVE_M * TM)
 assert WAVES_PER_BLOCK_M*REG_TILES_PER_WAVE_M*LANES_PER_WAVE_M*TM == BLOCK_M, "M reshape is wrong"
 assert WAVES_PER_BLOCK_N*REG_TILES_PER_WAVE_N*LANES_PER_WAVE_N*TN == BLOCK_N, "N reshape is wrong"
 
-def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.LOOP): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
+def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.WEAK): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
 def copy(dest:UOp, src:UOp, rng:int, upcast=False):
   assert dest.shape == src.shape
-  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.LOOP)
+  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.WEAK)
   return dest[*rngs].store(src[*rngs]).end(*rngs)
 
 def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:

extra/gemm/cdna_asm_gemm.py

@@ -171,8 +171,8 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
   M, K = A.shape[0]*A.shape[1], A.shape[2]
   K2, N = B.shape[(1 if B.ndim == 3 else 0):]
   assert K == K2
-  m = UOp.range(M, 1, AxisType.LOOP)
-  n = UOp.range(N, 2, AxisType.LOOP)
+  m = UOp.range(M, 1)
+  n = UOp.range(N, 2)
   k = UOp.range(K, 0, AxisType.REDUCE)
   mul = (A.flatten().index((m*UOp.const(dtypes.weakint, K)+k))*
          B.flatten().index((k*UOp.const(dtypes.weakint, N)+n))).cast(dtypes.float32)

extra/gemm/mi350x_uop_matmul_2.py

@@ -29,7 +29,7 @@ TID_SIZE = WARPGROUP_SIZE*WARP_SIZE
 
 def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=()):
   assert dest.shape == src.shape
-  rngs = [UOp.range(s, rng+i, AxisType.UPCAST if i in upcast else AxisType.LOOP) for i,s in enumerate(src.shape)]
+  rngs = [UOp.range(s, rng+i, AxisType.UPCAST if i in upcast else AxisType.WEAK) for i,s in enumerate(src.shape)]
   copy = dest[*rngs].store(src[*rngs]).end(*rngs)
   return dest.after(copy) if set else copy
 

extra/gemm_fragment.py

@@ -0,0 +1,195 @@
+"""
+tilelang-style matmul_relu written with tinygrad UOp APIs.
+
+Demonstrates that tilelang's T.alloc_fragment is expressible with existing
+tinygrad primitives: a per-thread REG buffer, wrapped in one Ops.UNSHARD per
+sharded axis over the LOCAL thread-grid ranges to form the full logical tile.
+Here the 64 threads are an 8x8 grid and each thread owns an 8x8 sub-tile --
+the 2-D fragment layout tilelang infers. The kernel is written against the
+full-tile UNSHARD view, and multi_pm (the same pass that lowers multi-device
+UNSHARDs) resolves it into per-thread shard code.
+
+Reference tilelang kernel:
+
+    @tilelang.jit
+    def matmul_relu(A, B, block_M=64, block_N=64, block_K=64,
+                    dtype=T.float16, accum_dtype=T.float32):
+        M, N, K = T.const('M, N, K')
+        C = T.empty([M, N], dtype)
+        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
+            A_shared = T.alloc_shared((block_M, block_K), dtype)
+            B_shared = T.alloc_shared((block_K, block_N), dtype)
+            C_local  = T.alloc_fragment((block_M, block_N), accum_dtype)
+            T.clear(C_local)
+            for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
+                T.copy(A[by * block_M, ko * block_K], A_shared)
+                T.copy(B[ko * block_K, bx * block_N], B_shared)
+                T.gemm(A_shared, B_shared, C_local)
+            for i, j in T.Parallel(block_M, block_N):
+                C_local[i, j] = T.max(C_local[i, j], 0)
+            T.copy(C_local, C[by * block_M, bx * block_N])
+        return C
+
+API mapping (tilelang -> tinygrad UOps, idioms from test/backend/test_custom_kernel.py):
+
+    T.Kernel(gx, gy, threads=T)    -> AxisType.GLOBAL ranges (blocks) + AxisType.LOCAL ranges (thread grid)
+    T.alloc_shared(shape, dtype)   -> UOp.placeholder(shape, dtype, slot, AddrSpace.LOCAL)
+    T.alloc_fragment(shape, dt)    -> per-thread REG placeholder, wrapped in one Ops.UNSHARD per sharded axis over
+                                      the AxisType.LOCAL ranges: fragment.unshard((axis_y, axis_x), (ty, tx)).
+                                      The full logical tile is the shard with each sharded axis multiplied by its
+                                      range size, exactly like device sharding, but the sharding axes are thread
+                                      axes carried by the RANGE metadata instead of a device tuple. C_local[i, j]
+                                      with [i, j] in this thread's shard is INDEX on the UNSHARD, which multi_pm
+                                      resolves into INDEX on the per-thread REG shard, axis by axis.
+    T.copy(gmem_slice, smem)       -> smem[thread_idx].set(gmem_slice[thread_idx], end=copy_rng). set returns the
+                                      smem tile AFTER the copy; the implicit-barrier pass turns the store->load
+                                      dependency of the loop that consumes it into a workgroup barrier
+    T.gemm (no WMMA)               -> C_local[..].set(C_local.after(k)[..] + a_shared[..] * b_shared[..], end=k)
+                                      with k a loop-carried LOOP range (codegen builds the register accumulator
+                                      from this self-referential store automatically)
+    T.copy(fragment, gmem)         -> gmem.index(gidx).store(C_local[..]).end(all_ranges)
+    UNSHARD lowering               -> multi_pm in codegen (full_rewrite_to_sink): INDEX/AFTER/STORE ops on the
+                                      full-tile view become per-thread shard ops, no UNSHARD survives into the program.
+"""
+
+from tinygrad.dtype import dtypes, AddrSpace, DType
+from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
+from tinygrad.helpers import cdiv, getenv
+from tinygrad.tensor import Tensor
+
+# ---------------------------------------------------------------------------
+# tilelang builtins, expressed with tinygrad UOp APIs
+# ---------------------------------------------------------------------------
+
+def alloc_shared(shape:tuple[int, ...], dtype:DType) -> UOp:
+  """T.alloc_shared: one LOCAL buffer shared by all threads in the block."""
+  return UOp.placeholder(tuple(shape), dtype, next(UOp.unique_num), AddrSpace.LOCAL)
+
+def alloc_fragment(shape:tuple[int, ...], dtype:DType, axes:tuple[int, ...], rngs:tuple[UOp, ...]) -> UOp:
+  """T.alloc_fragment: per-thread REG fragment + UNSHARD over the LOCAL thread grid.
+
+  Each thread privately owns shape[axis]//threads elements along every sharded
+  axis in a REG buffer. The UNSHARDs over the LOCAL thread ranges present the
+  full logical tile: full_shape = shard_shape with each sharded axis multiplied
+  by its range size. This is exactly how UNSHARD carries a DEVICE axis today,
+  except the sharding axes are thread axes carried by the RANGE metadata.
+  """
+  assert len(axes) == len(rngs)
+  assert all(tnum.op is Ops.RANGE and tnum.arg[-1] is AxisType.LOCAL for tnum in rngs), "fragments shard over LOCAL ranges"
+  assert all(shape[a] % (int(rng.vmax)+1) == 0 for a, rng in zip(axes, rngs))
+  by_axis = dict(zip(axes, rngs))
+  shard_shape = tuple(s // (int(by_axis[i].vmax)+1) if i in by_axis else s for i, s in enumerate(shape))
+  fragment = UOp.placeholder(shard_shape, dtype, next(UOp.unique_num), AddrSpace.REG)
+  return fragment.unshard(axes, rngs)
+
+# ---------------------------------------------------------------------------
+# GEMM kernel: C = relu(A @ B), float inputs (fp16 or fp32), fp32 fragment accumulator, no WMMA
+# ---------------------------------------------------------------------------
+
+# 64x64 output tile per block, 128 threads as an 8x16 grid; each thread owns an 8x4 fragment sub-tile
+# (the 2-D per-thread layout tilelang infers for this GEMM). The 4 contiguous columns (TN=4) are what
+# let codegen vectorize loads/stores to float4, matching tilelang's lowering exactly.
+BLOCK_M = BLOCK_N = BLOCK_K = 64
+TY = 8
+TX = 16
+THREADS = TY * TX
+TM = BLOCK_M // TY   # fragment rows per thread  (8)
+TN = BLOCK_N // TX   # fragment columns per thread (4)
+
+def matmul_relu_kernel(c:UOp, a:UOp, b:UOp) -> UOp:
+  """C[M, N] = relu(A[M, K] @ B[K, N]) -- one 64x64 tile per block, locals + a 2-D fragment."""
+  M, K = a.shape
+  K2, N = b.shape
+  assert K == K2 and a.dtype == b.dtype == c.dtype and not dtypes.is_int(a.dtype)
+  assert not (K % BLOCK_K or M % BLOCK_M or N % BLOCK_N), "test sizes must be multiples of the block sizes"
+
+  # with T.Kernel(T.ceildiv(N, BLOCK_N), T.ceildiv(M, BLOCK_M), threads=128) as (bx, by):
+  bx = UOp.range(cdiv(N, BLOCK_N), 0, AxisType.GLOBAL)
+  by = UOp.range(cdiv(M, BLOCK_M), 1, AxisType.GLOBAL)
+  # tx (N, 16) is the fast/inner LOCAL axis so a warp covers 16 cols x 2 rows --
+  # matching tilelang's (tidx>>4, tidx&15) warp composition. This keeps the 8 A_shared
+  # reads in a warp on only 2 row-groups (broadcast across 16 cols) instead of 8 rows
+  # (8-way bank conflict), since A_shared[row*512 + ...] all map to the same bank when 8
+  # distinct rows land in one warp.
+  tx = UOp.range(TX, 2, AxisType.LOCAL)
+  ty = UOp.range(TY, 3, AxisType.LOCAL)
+
+  # A_shared = T.alloc_shared((BLOCK_M, BLOCK_K), dtype)
+  # B_shared = T.alloc_shared((BLOCK_K, BLOCK_N), dtype)
+  A_shared = alloc_shared((BLOCK_M, BLOCK_K), a.dtype)
+  B_shared = alloc_shared((BLOCK_K, BLOCK_N), b.dtype)
+
+  # C_local = T.alloc_fragment((BLOCK_M, BLOCK_N), accum_dtype) -- an 8x4 REG tile per thread of the 8x16 grid
+  C_local = alloc_fragment((BLOCK_M, BLOCK_N), dtypes.float32, (0, 1), (ty, tx))
+
+  # T.clear(C_local) -- each thread zeroes its own fragment sub-tile
+  ic, jc = UOp.range(TM, 4, AxisType.LOOP), UOp.range(TN, 5, AxisType.UPCAST)
+  C_loc = C_local[ic*TM + ty, tx*TN + jc].set(0.0, end=(ic, jc))
+
+  # for ko in T.Pipelined(T.ceildiv(K, BLOCK_K), num_stages=3):
+  # (num_stages pipelining is async copy + multi-buffering; this is the synchronous single-buffer version)
+  ko = UOp.range(cdiv(K, BLOCK_K), 6, AxisType.LOOP)
+
+  # T.copy(A[by * BLOCK_M, ko * BLOCK_K], A_shared) -- each thread copies its own 8x4 sub-tile.
+  # Row index is iar*TM + ty (strided by TM across ty), matching tilelang's layout: thread ty owns
+  # rows {ty, ty+8, ..., ty+56} not {ty*8, ..., ty*8+7}.
+  iar, ka = UOp.range(TM, 7, AxisType.LOOP), UOp.range(TN, 8, AxisType.UPCAST)
+  A_store = A_shared[iar*TM + ty, tx*TN + ka].store(a[by*BLOCK_M + iar*TM + ty, ko*BLOCK_K + tx*TN + ka]).end(iar, ka)
+
+  # T.copy(B[ko * BLOCK_K, bx * BLOCK_N], B_shared)
+  kb, ibr = UOp.range(TM, 9, AxisType.LOOP), UOp.range(TN, 10, AxisType.UPCAST)
+  B_store = B_shared[kb*TM + ty, tx*TN + ibr].store(b[ko*BLOCK_K + kb*TM + ty, bx*BLOCK_N + tx*TN + ibr]).end(kb, ibr)
+
+  # get the shared after the stores (single barrier)
+  A_shared = A_shared.after(A_store, B_store)
+  B_shared = B_shared.after(A_store, B_store)
+
+  # T.gemm(A_shared, B_shared, C_local), no WMMA -- per-thread accumulate over its fragment sub-tile.
+  # identical to custom_gemm: a self-referential store over the loop-carried kk range,
+  # which codegen turns into a register accumulator
+  # kk is the outer compute loop (axis 11) so that for each kk we read all 8 A rows and reuse
+  # the B[kk] read across them -- matching tilelang's ko > kk > row > col access order exactly.
+  kk, ir = UOp.range(BLOCK_K, 11, AxisType.LOOP), UOp.range(TM, 12, AxisType.LOOP)
+  jj = UOp.range(TN, 13, AxisType.UPCAST)
+  acc = C_loc.after(kk)[ir*TM + ty, tx*TN + jj] + A_shared[ir*TM + ty, kk].cast(dtypes.float32) * B_shared[kk, tx*TN + jj].cast(dtypes.float32)
+  # closing the ko loop here too; codegen adds the barrier so no thread overwrites the tiles while others still read them
+  C_loc = C_loc[ir*TM + ty, tx*TN + jj].set(acc, end=(kk, ir, jj, ko))
+
+  # for i, j in T.Parallel(BLOCK_M, BLOCK_N): C_local[i, j] = T.max(C_local[i, j], 0)
+  # T.copy(C_local, C[by * BLOCK_M, bx * BLOCK_N]) -- per-thread store of the fragment shard (relu fused into it)
+  # LOOP: these loops are the per-thread output layout; convert_loop_to_global must not globalize them
+  ie, je = UOp.range(TM, 14, AxisType.LOOP), UOp.range(TN, 15, AxisType.UPCAST)
+  c_st = c[by*BLOCK_M + ie*TM + ty, bx*BLOCK_N + tx*TN + je].store(C_loc[ie*TM + ty, tx*TN + je].relu().cast(c.dtype))
+
+  # all open ranges are closed at the final store (ko was closed above).
+  # the fragment UNSHARDs go to codegen as is: multi_pm there resolves the full-tile view into per-thread shard code
+  return c_st.end(je, ie, tx, ty, bx, by).sink(arg=KernelInfo(name="matmul_relu", opts_to_apply=()))
+
+# ---------------------------------------------------------------------------
+# python wrapper: same signature as the tilelang function
+# ---------------------------------------------------------------------------
+
+def matmul_relu(a:Tensor, b:Tensor) -> Tensor:
+  """C = relu(A @ B), fp16 in/out with an fp32 fragment accumulator."""
+  c = Tensor.empty(a.shape[0], b.shape[1], dtype=a.dtype, device=a.device)
+  return c.custom_kernel(a, b, fxn=matmul_relu_kernel)[0]
+
+# ---------------------------------------------------------------------------
+# test
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+  from tinygrad import Device
+  assert Device[Device.DEFAULT].renderer.has_local, "this GPU-style kernel needs a backend with local memory (LOCAL ranges + barriers)"
+  M = K = N = getenv("N", 256)  # 4x4 grid of 64x64 tiles, 4 K chunks
+  dtype_in = dtypes.half if getenv("HALF") else dtypes.float
+
+  a = Tensor.randn(M, K, dtype=dtype_in).contiguous()
+  b = Tensor.randn(K, N, dtype=dtype_in).contiguous()
+  ref = (a @ b).relu().realize()
+
+  out = matmul_relu(a, b).realize()
+
+  import numpy as np
+  np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=1e-1, rtol=1e-2)
+  print("matmul_relu passed!")

extra/hcq2/ops_amd2.py

@@ -4,7 +4,7 @@ import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, co
 assert sys.platform != 'win32'
 from dataclasses import dataclass
 from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, encode_kernargs_clike, make_cmdbuf
-from tinygrad.runtime.support.hcq2 import make_binary_patch, make_patch
+from tinygrad.runtime.support.hcq2 import make_binary_patch, make_patches
 from tinygrad.uop.ops import sint, UOp
 from tinygrad.device import Compiled, BufferSpec, Buffer, Device
 from tinygrad.dtype import dtypes
@@ -158,7 +158,7 @@ def pm4_submit(ctx, lin):
 
   ib = UOp.placeholder((size_dw + 2,), dtypes.uint32, next(UOp.unique_num), device=devs, volatile=True).rtag("cmdbuf")
   done_idx, submit_idx = UOp.const(dtypes.int, size_dw + 0), UOp.const(dtypes.int, size_dw + 1)
-  submitted = (counter:=ib.after(*[make_patch(ib, (size_dw + i) * 4, UOp.const(dtypes.uint32, 0)) for i in range(2)]).index(submit_idx)).load()
+  submitted = (counter:=ib.after(make_patches(ib, [((size_dw + i) * 4, UOp.const(dtypes.uint32, 0)) for i in range(2)])).index(submit_idx)).load()
   completed = ib.after(loop:=UOp.loop(0)).index(done_idx).load()
   ib_free = completed.end(loop, completed != submitted)
 

extra/llama_kernels/quantize_fp8_delayed/__init__.py

@@ -16,7 +16,7 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:
 
   wg = UOp.range(NUM_WG, 0, AxisType.GLOBAL)
   tid = UOp.range(THREADS_PER_WG, 1, AxisType.LOCAL)
-  it = UOp.range((n_elems // VEC) // (NUM_WG * THREADS_PER_WG), 2, AxisType.LOOP)
+  it = UOp.range((n_elems // VEC) // (NUM_WG * THREADS_PER_WG), 2, AxisType.WEAK)
   lane = UOp.range(VEC, 3, AxisType.UNROLL)
 
   idx = (((it * NUM_WG + wg) * THREADS_PER_WG + tid) * VEC) + lane

extra/qcom_gpu_driver/opencl_ioctl.py

@@ -4,7 +4,7 @@ from hexdump import hexdump
 from copy import deepcopy
 import pathlib, sys
 from tinygrad.helpers import to_mv, getenv
-from tinygrad.runtime.autogen import adreno
+from tinygrad.runtime.autogen import mesa
 sys.path.append(pathlib.Path(__file__).parent.parent.parent.as_posix())
 
 IOCTL = getenv("IOCTL", 0)
@@ -23,7 +23,7 @@ for child in xml.getroot():
 CAPTURED_STATE = {}
 
 REGS = {}
-for k, v in adreno.__dict__.items():
+for k, v in mesa.__dict__.items():
   if k.startswith("REG_") and isinstance(v, int) and v > 1024: REGS[v] = k
 
 from extra.qcom_gpu_driver import msm_kgsl
@@ -42,7 +42,7 @@ def get_struct(argp, stype):
 
 def format_struct(s):
   sdats = []
-  for field_name, *_ in s._real_fields_:
+  for field_name, *_ in s._fields_:
     if field_name in {"__pad", "PADDING_0"}: continue
     dat = getattr(s, field_name)
     if isinstance(dat, int): sdats.append(f"{field_name}:0x{dat:X}")
@@ -96,9 +96,9 @@ def parse_cmd_buf(dat):
         CAPTURED_STATE['LOAD_FRAGS'].append((state_block, state_type, num_unit, dst_off))
 
         if state_block == SB6_CS_SHADER:
-          from extra.disassemblers.adreno import disasm_raw
+          from tinygrad.runtime.support.compiler_mesa import disas_adreno
           if state_type == ST6_SHADER and IOCTL > 3:
-            disasm_raw(get_mem(((vals[2] << 32) | vals[1]), num_unit * 128))
+            disas_adreno(get_mem(((vals[2] << 32) | vals[1]), num_unit * 128))
           if state_type == ST6_CONSTANTS:
             x = get_mem(((vals[2] << 32) | vals[1]), num_unit*4)
             CAPTURED_STATE['constants'] = x[:]
@@ -142,7 +142,7 @@ def parse_cmd_buf(dat):
       vals = struct.unpack("I"*size, dat[ptr+4:ptr+4+4*size])
       if IOCTL > 0: print(f"{ptr:3X} -- typ 4: {size=:3d}, {reg_name}", hprint(vals))
       for vi,v in enumerate(vals): CAPTURED_STATE[offset+vi] = v
-      if offset == adreno.REG_A6XX_SP_CS_CONFIG:
+      if offset == mesa.REG_A6XX_SP_CS_CONFIG:
         val = vals[0]
         if IOCTL > 0:
           print(f"\tBINDLESS_TEX={(val >> 0) & 0b1}")
@@ -215,79 +215,3 @@ def install_hook(c_function, python_function):
 
 libc = ctypes.CDLL(ctypes.util.find_library("libc"))
 install_hook(libc.ioctl, ioctl)
-
-def before_launch():
-  global CAPTURED_STATE
-  CAPTURED_STATE.clear()
-def collect_last_launch_state():
-  global CAPTURED_STATE
-  return deepcopy(CAPTURED_STATE)
-def compare_launch_state(state, good_state):
-  cmp = [
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_NTEX__MASK),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_NSAMP__MASK),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_NIBO__MASK),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_ENABLED),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_BINDLESS_TEX),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_BINDLESS_SAMP),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_BINDLESS_IBO),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_BINDLESS_UBO),
-
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_HALFREGFOOTPRINT__MASK),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_FULLREGFOOTPRINT__MASK),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_BRANCHSTACK__MASK),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_FULLREGFOOTPRINT__MASK),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_THREADMODE__MASK),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_EARLYPREAMBLE),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_MERGEDREGS),
-
-    (adreno.REG_A6XX_SP_CS_PVT_MEM_PARAM, adreno.A6XX_SP_CS_PVT_MEM_PARAM_MEMSIZEPERITEM__MASK),
-    (adreno.REG_A6XX_SP_CS_PVT_MEM_PARAM, adreno.A6XX_SP_CS_PVT_MEM_PARAM_HWSTACKSIZEPERTHREAD__MASK),
-
-    (adreno.REG_A6XX_SP_CS_UNKNOWN_A9B1, adreno.A6XX_SP_CS_UNKNOWN_A9B1_UNK5),
-    (adreno.REG_A6XX_SP_CS_UNKNOWN_A9B1, adreno.A6XX_SP_CS_UNKNOWN_A9B1_UNK6),
-
-    (adreno.REG_A6XX_SP_CS_BRANCH_COND, 0xffffffff),
-
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_0, adreno.A6XX_HLSQ_CS_NDRANGE_0_KERNELDIM__MASK),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_0, adreno.A6XX_HLSQ_CS_NDRANGE_0_LOCALSIZEX__MASK),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_0, adreno.A6XX_HLSQ_CS_NDRANGE_0_LOCALSIZEY__MASK),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_0, adreno.A6XX_HLSQ_CS_NDRANGE_0_LOCALSIZEZ__MASK),
-
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_1, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_2, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_3, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_4, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_5, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_6, 0xffffffff),
-
-    (adreno.REG_A6XX_HLSQ_CS_CNTL_0, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_CNTL_1, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_KERNEL_GROUP_X, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_KERNEL_GROUP_Y, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_KERNEL_GROUP_Z, 0xffffffff),
-  ]
-
-  for x,m in cmp:
-    print(f"Field {REGS[x]}, mask: 0x{m:X} cmp: {state.get(x, 0) & m} vs {good_state.get(x, 0) & m}")
-    if state.get(x, 0) & m != good_state.get(x, 0) & m:
-      return False, f"Field {REGS[x]}, mask: 0x{m:X} mismatch: {state.get(x, 0) & m} vs {good_state.get(x, 0) & m}"
-
-  for n in ['descriptors', 'ibos']:
-    if n not in good_state: continue
-    mv1, mv2 = state.get(n), good_state.get(n)
-
-    if len(mv1) != len(mv2): return False, f"{n}: len mismatch {len(mv1)} != {len(mv2)}"
-    mv1 = memoryview(bytearray(mv1)).cast('I')
-    mv2 = memoryview(bytearray(mv2)).cast('I')
-    for i in range(len(mv2)):
-      if i % 8 == 5 or i % 8 == 4: continue # addresses
-      if mv1[i]!=mv2[i]: return False, f"{n}: content mismatch {i} {mv1[i]} {mv2[i]}"
-
-  for n in ['samplers']:
-    if n not in good_state: continue
-    mv1, mv2 = state.get(n), good_state.get(n)
-    if len(mv1) != len(mv2): return False, f"{n}: len mismatch {len(mv1)} != {len(mv2)}"
-    if any(mv1[i]!=mv2[i] for i in range(len(mv1))): return False, f"{n}: content mismatch"
-
-  return True, "PASS"

extra/thunder/tiny/tk/kernel.py

@@ -48,14 +48,14 @@ class Kernel(AbstractContextManager):
   @property
   def warpgroup(self): return self.group(4)
 
-  def range(self, start:int, end:int=0, step:int=1, axis_type:AxisType=AxisType.LOOP, track:bool=True):
+  def range(self, start:int, end:int=0, step:int=1, axis_type:AxisType=AxisType.WEAK, track:bool=True):
     if end == 0: start, end = 0, start
     rng = _tk_range(start, end, step, axis_type, self.range_id)
     self.range_id += 1
     if track: self.range_stack.append(rng)
     return rng
 
-  def raw_range(self, end:int=0, axis_type:AxisType=AxisType.LOOP):
+  def raw_range(self, end:int=0, axis_type:AxisType=AxisType.WEAK):
     rng = UOp.range(end, self.range_id, axis_type=axis_type)
     self.range_id += 1
     return rng

spec/tinyspec.tex

@@ -80,7 +80,7 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op
 \op{Index} & $(T, i_0, i_1, \ldots)$ & --- & Index from left. $()$-shaped $i$ removes dim; $(k,)$-shaped makes it $k$. \\
 \op{Stack} & $(T_0, T_1, \ldots)$ & --- & Join along a newly created leading axis. All shapes must match. \\
 \op{Bitcast} & $(T,)$ & dtype & Reinterpret storage as target dtype; preserve total bytes. \\
-\op{Unshard} & $(T, R)$ & axis $a$ & Concatenate the shards indexed by \op{Range} $R$ along $a$; $R$ is outer. \\
+\op{Unshard} & $(T, R_0, R_1, \ldots)$ & axes $(a_0, a_1, \ldots)$ & Concatenate shards of \op{Range} $R_k$ along axis $a_k$; $R_k$ is outer. \\
 \bottomrule
 \end{tabular}
 
@@ -260,7 +260,7 @@ Every UOp has a \textbf{dtype}, \textbf{shape}, \textbf{device}, \textbf{addrspa
 \op{Const} & from arg & $()$ & \textsc{null} & $[v, v]$ \\
 \op{Param} & from arg & from $\mathrm{src}[0]$ & from arg & from src or dtype range \\[3pt]
 Movement ops & $\mathrm{src}[0].\mathrm{dtype}$ & (see op) & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
-\op{Unshard} & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0]$, axis $\times n$ & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
+\op{Unshard} & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0]$, each $a_k \times n_k$ & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
 \op{Reduce} & $\mathrm{src}[0].\mathrm{dtype}$ & remove first $n$ axes & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\[3pt]
 \op{Cast} & from arg & $\mathrm{src}[0].\mathrm{shape}$ & $\mathrm{src}[0].\mathrm{device}$ & clamped to dtype \\
 \op{Bitcast} & from arg & $\mathrm{src}[0].\mathrm{shape}$ & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\
@@ -286,9 +286,9 @@ $[a,A]$, $[b,B]$, $[c,C]$ denote min\_max of $\mathrm{src}[0]$, $\mathrm{src}[1]
 Default \emph{dtype range}: $[\mathrm{dtype\_min},\, \mathrm{dtype\_max}]$.
 
 \medskip
-\textbf{axis} tracks the multi-device sharding dimension. \op{Unshard} defines it (axis $=$ arg). \op{Buffer} with $n$-tuple device: axis $= 0$ (device dim).
-\op{Reshape} remaps axis to preserve the shard boundary. \op{Permute} follows the permutation. \op{Expand} shifts axis right by $|\mathbf{n}|$.
-\op{Reduce} on the shard axis $\to$ \textsc{null} (shard axis is among the first $n$ axes). \op{Replicated} on the shard axis $\to$ \textsc{null}. \op{Copy} $\to$ \textsc{null}. ALU ops inherit from sources. Default: \textsc{null}.
+\textbf{sharding} tracks multi-device sharding as a set of (axis, \op{Range}) pairs. \op{Unshard} defines it: arg is the tuple of sharded axes, one \op{Range} in src per axis (positional: the $k$-th \op{Range} shards the $k$-th axis). \op{Buffer} with $n$-tuple device: sharded on axis $0$ (device dim). The single-axis convenience \textbf{axis} is \textsc{null} unless exactly one axis is sharded.
+\op{Reshape} remaps each sharded axis to preserve its shard boundary. \op{Permute} follows the permutation. \op{Expand} shifts all sharded axes right by $|\mathbf{n}|$.
+\op{Reduce} on a sharded axis drops it. \op{Replicated} on the shard axis $\to$ \textsc{null}. \op{Copy} $\to$ \textsc{null}. ALU ops inherit from sources. Default: \textsc{null}.
 
 %% ============================================================
 \subsection*{Kernel Optimizations (OptOps) \normalfont\small--- schedule-level transforms on kernel ranges}
@@ -382,7 +382,7 @@ def scatter_add(T, idx, val):
 Let $D = (d_0, \ldots, d_{n-1})$ be an $n$-tuple device.
 \op{Copy} to an $n$-tuple device reshards with axis $= 0$. \op{Copy} never changes shape.
 
-\textbf{Sharding} splits a tensor along an axis across $n$ devices. It opens a \op{Range} of type \texttt{DEVICE} (a symbolic per-device index $d$), shrinks each device's view to its piece, then closes the range with \op{Unshard}$(T, R, a)$. The result is a logical tensor whose shape along axis $a$ is the full size; each device holds $1/n$ of it. \op{Unshard} is the inverse of sharding --- it marks the boundary between per-device computation and the logical multi-device tensor. The range need not be \texttt{DEVICE}; e.g.\ a \texttt{WARP} range closes the same way, concatenating per-lane shards along $a$ with the range as the outer factor.
+\textbf{Sharding} splits a tensor along an axis across $n$ devices. It opens a \op{Range} of type \texttt{DEVICE} (a symbolic per-device index $d$), shrinks each device's view to its piece, then closes the range with \op{Unshard}$(T, R, a)$. The result is a logical tensor whose shape along axis $a$ is the full size; each device holds $1/n$ of it. \op{Unshard} is the inverse of sharding --- it marks the boundary between per-device computation and the logical multi-device tensor. The range need not be \texttt{DEVICE}; e.g.\ a \texttt{WARP} range closes the same way, concatenating per-lane shards along $a$ with the range as the outer factor. A tensor may be sharded along several axes at once: \op{Unshard}$(T, R_0, R_1, \ldots;\; a_0, a_1, \ldots)$ carries one \op{Range} per sharded axis, and every movement op maps each sharded axis independently.
 
 \begin{lstlisting}
 # T has shape (s,) on a single device.

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_custom_kernel.py

@@ -2,6 +2,7 @@ import unittest
 from tinygrad import Tensor, UOp, GlobalCounters, Context, Device
 from tinygrad.dtype import AddrSpace, dtypes, Invalid
 from tinygrad.uop.ops import KernelInfo, AxisType, Ops
+from tinygrad.renderer.ptx import PTXRenderer
 
 # **** kernels ****
 
@@ -421,19 +422,76 @@ class TestCustomKernel(unittest.TestCase):
 
   @Context(DEV="CPU")
   def test_simple_from_source(self):
-    a = Tensor([0., 1., 2.]).realize()
-
-    src = "void test_src(float* restrict a) { a[0] = 1.0; }"
+    a = Tensor.arange(4).clone().realize()
+    src = "void test_src(int* restrict a) { a[0] = 1; }"
     # TODO: it currently requires a compiler for Ops.BINARY
     from tinygrad.device import Device
     binary = Device[a.device].renderer.compiler.compile(src)
     def custom_src_kernel(A:UOp) -> UOp:
       sink = UOp.sink(A, arg=KernelInfo(name="test_src"))
-      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())),
-                                   UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
+    a = Tensor.custom_kernel(a.reshape(2, 2).T, fxn=custom_src_kernel)[0]
+    self.assertEqual(a.tolist(), [[1, 2], [1, 3]])
 
-    a = Tensor.custom_kernel(a, fxn=custom_src_kernel)[0]
-    self.assertEqual(a.tolist(), [1., 1., 2.])
+class TestUnshardIndex(unittest.TestCase):
+  """Regression tests for INDEX on UNSHARD (fragment) resolution in schedule/multi.py.
+
+  A fragment is a per-thread REG buffer wrapped in UNSHARD over LOCAL thread ranges.
+  index_multi must resolve an INDEX on the UNSHARD view into an INDEX on the per-thread
+  shard. Two ownership patterns must work:
+    contiguous: idx = rng*shard_sz + local   (thread rng owns [rng*shard_sz, ...))
+    strided:    idx = rng + ir*shard_sz      (thread rng owns {rng, rng+shard_sz, ...})
+  """
+  def _run(self, kernel, shape=(8, 8)):
+    c = Tensor.empty(*shape)
+    out = Tensor.custom_kernel(c, fxn=kernel)[0]
+    try: return out.numpy()
+    except RuntimeError as e:
+      if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and "dynamic register indexing" in str(e):
+        self.skipTest("PTX does not support dynamic register indexing")
+      raise
+
+  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
+  def test_contiguous_fragment_index(self):
+    # thread ty owns rows [ty*8, ty*8+8) of a 64-row fragment -- contiguous ownership.
+    # This is the pre-existing case that index_multi always handled.
+    def kernel(C:UOp) -> UOp:
+      ty = UOp.range(8, 0, AxisType.LOCAL)
+      ir = UOp.range(8, 1, AxisType.LOOP)
+      j = UOp.range(8, 2, AxisType.LOOP)
+      # 8x8 fragment, 8 threads -> 64x8 full tile. thread ty owns rows [ty*8, ty*8+8).
+      frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
+      return C[ty*8 + ir, j].store(frag[ty*8 + ir, j]).end(j, ir, ty).sink(arg=KernelInfo(name="contig_frag"))
+    out = self._run(kernel, (64, 8))
+    assert out.shape == (64, 8)
+
+  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
+  def test_strided_fragment_index(self):
+    # thread ty owns rows {ty, ty+8, ty+16, ty+24, ..., ty+56} of a 64-row fragment --
+    # strided ownership. idx = ty + ir*8 where shard_sz=8 (8 threads, shard rows=8).
+    # The contiguous check (idx - rng*shard_sz) fails; the strided check
+    # (idx-rng) % shard_sz == 0 must succeed. This is the pattern the index_multi fix adds.
+    def kernel(C:UOp) -> UOp:
+      ty = UOp.range(8, 0, AxisType.LOCAL)
+      ir = UOp.range(8, 1, AxisType.LOOP)
+      j = UOp.range(8, 2, AxisType.LOOP)
+      # 8x8 fragment, 8 threads -> 64x8 full tile. thread ty owns rows {ty, ty+8, ..., ty+56}.
+      frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
+      return C[ty + ir*8, j].store(frag[ty + ir*8, j]).end(j, ir, ty).sink(arg=KernelInfo(name="strided_frag"))
+    out = self._run(kernel, (64, 8))
+    assert out.shape == (64, 8)
+
+  def test_fragment_index_cannot_shard(self):
+    # thread ty indexing rows [ty, ty+8) overlaps with other threads' rows -- this matches neither
+    # the contiguous nor the strided ownership pattern, so index_multi must raise.
+    def kernel(C:UOp) -> UOp:
+      ty = UOp.range(8, 0, AxisType.LOCAL)
+      ir = UOp.range(8, 1, AxisType.LOOP)
+      j = UOp.range(8, 2, AxisType.LOOP)
+      frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
+      return C[ty + ir, j].store(frag[ty + ir, j]).end(j, ir, ty).sink(arg=KernelInfo(name="bad_frag"))
+    with self.assertRaisesRegex(RuntimeError, "cannot shard index"):
+      self._run(kernel, (64, 8))
 
 class TestUOpReduce(unittest.TestCase):
   def test_uop_sum(self):

test/backend/test_linearizer.py

@@ -13,6 +13,7 @@ from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.renderer.isa import ISARenderer
 from test.helpers import replace_opts
+from test.backend.test_softmax_fusion import single_kernel_softmax
 MOCKGPU = DEV.interface.startswith("MOCK")
 
 from tinygrad.uop.render import print_uops # noqa: F401 # pylint: disable=unused-import
@@ -392,6 +393,16 @@ class TestLinearizer(unittest.TestCase):
     # the global store doesn't change
     assert stores[1].src[1].dtype == dtypes.float
 
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
+  def test_two_grouped_stores_local(self):
+    # GROUP on both reduces puts two LOCAL buffers in one kernel, and the store to each needs its own barrier
+    a = Tensor.rand(32, 32).realize()
+    opts = [Opt(OptOps.GROUP, 1, 4), Opt(OptOps.GROUP, 2, 4)]
+    ast = helper_linearizer_opt(single_kernel_softmax(a), [opts])
+    uops = to_program(replace_opts(ast, opts), renderer=Device[Device.DEFAULT].renderer).src[1].src
+    self.assertEqual(len([u for u in uops if u.op is Ops.BARRIER]), 2)
+
 # *** helpers ***
 
 def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:

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/backend/test_multitensor.py

@@ -425,6 +425,49 @@ class TestMultiBufferView(unittest.TestCase):
     run_linear(linear, var_vals)
     np.testing.assert_equal(out.numpy(), ref[5].numpy())
 
+@unittest.skipIf(not_support_multi_device(), "need multi")
+class Test2DShard(unittest.TestCase):
+  def setUp(self):
+    self.devices_4 = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
+    self.rng = UOp.range(4, -1, AxisType.DEVICE)
+    self.rng0, self.rng1 = self.rng // 2, self.rng % 2
+
+  def _shard_2d(self, t:Tensor) -> Tensor:
+    u = t.uop.copy_to_device(self.devices_4)._shard(0, self.rng0)._shard(1, self.rng1).unshard((0, 1), (self.rng0, self.rng1))
+    return Tensor(u)
+
+  def test_2d_shard_basic(self):
+    ref = Tensor.arange(16).reshape(4, 4).contiguous().realize()
+    t = self._shard_2d(ref)
+    out = t.contiguous().realize()
+    np.testing.assert_equal(out.numpy(), ref.numpy())
+
+  def test_2d_shard_elementwise(self):
+    ref = Tensor.arange(16).reshape(4, 4).contiguous().realize()
+    t = self._shard_2d(ref)
+    out = (t + 1).contiguous().realize()
+    np.testing.assert_equal(out.numpy(), ref.numpy() + 1)
+
+  def test_2d_shard_sum_all(self):
+    ref = Tensor.arange(16).reshape(4, 4).contiguous().realize()
+    t = self._shard_2d(ref)
+    out = t.sum().contiguous().realize()
+    np.testing.assert_equal(out.numpy(), np.array(ref.numpy().sum()))
+
+  def test_2d_shard_sum_non_sharded_axis(self):
+    ref = Tensor.arange(4*4*2).reshape(4, 4, 2).contiguous().realize()
+    t = self._shard_2d(ref)
+    out = t.sum(axis=2).contiguous().realize()
+    np.testing.assert_equal(out.numpy(), ref.numpy().sum(axis=2))
+
+  def test_2d_shard_matmul(self):
+    a = Tensor.arange(16).reshape(4, 4).contiguous().realize()
+    b = Tensor.arange(16).reshape(4, 4).contiguous().realize()
+    a_s = self._shard_2d(a)
+    b_s = self._shard_2d(b)
+    out = (a_s @ b_s).contiguous().realize()
+    np.testing.assert_equal(out.numpy(), a.numpy() @ b.numpy())
+
 @unittest.skipIf(not_support_multi_device(), "need multi")
 class TestMultiTransformer(unittest.TestCase):
   @needs_second_gpu

test/backend/test_tensor_variable.py

@@ -1,6 +1,6 @@
 import unittest
 import numpy as np
-from tinygrad import Tensor, Variable, dtypes
+from tinygrad import Device, Tensor, Variable, TinyJit, dtypes
 from tinygrad.helpers import CHECK_OOB
 
 class TestTensorVariable(unittest.TestCase):
@@ -18,10 +18,18 @@ class TestTensorVariable(unittest.TestCase):
       self.assertListEqual((vv * t).tolist(), [2, 2, 2])
     except RuntimeError: pass
 
-  # TODO: a Variable PARAM lowers to int32, so a bound value that doesn't fit int32 truncates or fails to bind
-  @unittest.expectedFailure
+  @unittest.skipUnless(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires long support")
   def test_large_range_variable(self):
-    self.assertEqual(Tensor(Variable("b", 0, 2**40).bind(2**35)).item(), 2**35)
+    self.assertEqual(Tensor(Variable("b", 0, 2**40, dtype=dtypes.long).bind(2**35)).clone(Device.DEFAULT).item(), 2**35)
+
+  @unittest.skipUnless(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires long support")
+  def test_large_range_variable_jit(self):
+    @TinyJit
+    def f(a,b): return (Tensor(a+b).clone(Device.DEFAULT) * 2).realize()
+    for i in range(3):
+      a = Variable("a", 0, 2**10, dtype=dtypes.int).bind(i)
+      b = Variable("b", 0, 2**40, dtype=dtypes.long).bind(2**35)
+      self.assertEqual(f(a,b).item(), (2**35 + i) * 2)
 
   def test_variable_defers_like_a_literal(self):
     vv = Variable("a", 1, 10).bind(2)

test/external/external_benchmark_op_conv.py

@@ -13,9 +13,9 @@ from tinygrad.dtype import Invalid
 
 def vision_conv_143():
   c0 = UOp.param(0, dtypes.half, shape=(16, 1024, 4))
-  c2 = UOp.range(32, 3, AxisType.LOOP)
-  c5 = UOp.range(128, 4, AxisType.LOOP)
-  c8 = UOp.range(16, 2, AxisType.LOOP)
+  c2 = UOp.range(32, 3)
+  c5 = UOp.range(128, 4)
+  c8 = UOp.range(16, 2)
   c16 = UOp.range(7, 0, AxisType.REDUCE)
   c17 = c8*2+c16
   c24 = ((c17<3)!=True)&(c17<35)
@@ -39,9 +39,9 @@ def vision_conv_143():
 
 def vision_conv_153():
   c0 = UOp.param(0, dtypes.half, shape=(8, 1024, 4))
-  c2 = UOp.range(16, 3, AxisType.LOOP)
-  c5 = UOp.range(256, 4, AxisType.LOOP)
-  c8 = UOp.range(8, 2, AxisType.LOOP)
+  c2 = UOp.range(16, 3)
+  c5 = UOp.range(256, 4)
+  c8 = UOp.range(8, 2)
   c16 = UOp.range(7, 0, AxisType.REDUCE)
   c17 = c8*2+c16
   c24 = ((c17<3)!=True)&(c17<19)
@@ -65,7 +65,7 @@ def vision_conv_153():
 
 def dm_conv_172():
   c0 = UOp.param(0, dtypes.half, shape=(1, 240, 4))
-  c2 = UOp.range(960, 4, AxisType.LOOP)
+  c2 = UOp.range(960, 4)
   c5 = UOp.param(1, dtypes.half, shape=(8, 384, 4))
   c7 = UOp.range(32, 0, AxisType.REDUCE)
   c10 = UOp.range(4, 1, AxisType.REDUCE)

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/mockgpu/amd/emu.py

@@ -51,7 +51,7 @@ class _MXCSRContext:
     if lib is None or not hasattr(self, '_saved'): return
     lib.set_fpcr(self._saved)
 
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.device import Buffer, BufferSpec, Device
 from tinygrad.runtime.autogen import hsa
@@ -446,7 +446,7 @@ class _Ctx:
     """Create a lane range UOp with unique axis ID."""
     if n is None: n = self.wave_size
     self._axis_id += 1
-    return UOp.range(n, self._axis_id, AxisType.LOOP, dtype=dtypes.int)
+    return UOp.range(n, self._axis_id, dtype=dtypes.int)
 
   def unroll_lanes(self, get_lane_bit, exec_mask: UOp, apply_exec: bool = True) -> UOp:
     """Combine lane bits into a mask using RANGE+REDUCE (32-bit for RDNA, 64-bit for CDNA)."""

test/mockgpu/nv/nvgpu.py

@@ -91,7 +91,7 @@ class GPFIFO:
     args_cnt, vals_cnt = const0[80], const0[81]
     args_addr = qmd.constant_buffer_addr_lower_0 + (qmd.constant_buffer_addr_upper_0 << 32) + 0x160
     args = to_mv(args_addr, args_cnt*8).cast('Q')
-    vals = to_mv(args_addr + args_cnt*8, vals_cnt*4).cast('I')
+    vals = to_mv(args_addr + args_cnt*8, vals_cnt*8).cast('Q')
     cargs = [ctypes.cast(args[i], ctypes.c_void_p) for i in range(args_cnt)] + [ctypes.cast(vals[i], ctypes.c_void_p) for i in range(vals_cnt)]
     gx, gy, gz = qmd.cta_raster_width, qmd.cta_raster_height, qmd.cta_raster_depth
     lx, ly, lz = qmd.cta_thread_dimension0, qmd.cta_thread_dimension1, qmd.cta_thread_dimension2

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_linearizer_failures.py

@@ -8,8 +8,8 @@ from tinygrad.codegen import to_program
 class TestLinearizerFailures(unittest.TestCase):
   def test_fail_1(self):
     c0 = UOp.param(0, dtypes.float, (64,))
-    c1 = UOp.range(UOp.const(dtypes.weakint, 2), 1, AxisType.LOOP)
-    c2 = UOp.range(UOp.const(dtypes.weakint, 32), 2, AxisType.LOOP)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 2), 1, AxisType.WEAK)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 32), 2, AxisType.WEAK)
     c3 = ((c1*UOp.const(dtypes.weakint, 32))+c2)
     c4 = UOp.param(1, dtypes.float, (163840,))
     c5 = UOp.range(UOp.const(dtypes.weakint, 2560), 0, AxisType.REDUCE)

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

@@ -424,8 +424,8 @@ class TestUOpGraph(unittest.TestCase):
     # mnist indexing with split reduceop
     # Make sure we are not doign math on the loaded index, which would promote it to long
     c0 = UOp.param(0, dtypes.uchar, (128000,))
-    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.LOOP)
-    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.LOOP)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.WEAK)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.WEAK)
     c3 = UOp.param(1, dtypes.int, (512,))
     c4 = c3.index(c1)
     c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
@@ -441,8 +441,8 @@ class TestUOpGraph(unittest.TestCase):
   def test_load_idx_no_math_on_loaded(self):
     # test the (x+y)<c pattern where x has loads - we shouldn't do math on loaded indices
     c0 = UOp.param(0, dtypes.uchar, (128000,))
-    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.LOOP)
-    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.LOOP)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.WEAK)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.WEAK)
     c3 = UOp.param(1, dtypes.int, (512,))
     c4 = c3.index(c1)  # c4 is a load
     c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)

test/null/test_uop_symbolic.py

@@ -347,6 +347,12 @@ class TestSymbolic(unittest.TestCase):
   def test_mul_lt(self):
     self.helper_test_variable(Variable("a", 0, 5)*4 < 13, 0, 1, "(a<4)")
     self.helper_test_variable(Variable("a", 0, 5)*4 < 16, 0, 1, "(a<4)")
+    self.helper_test_variable(Variable("a", -5, 5)*4 < -13, 0, 1, "(a<-3)")
+    self.helper_test_variable(Variable("a", -5, 5)*-4 < 13, 0, 1, "((a*-1)<4)")
+    c0, c1 = 2, 2**54+1
+    self.helper_test_variable(Variable("a", 0, c1)*c0 < c1, 0, 1, f"(a<{2**53+1})")
+    c0, c1 = -2, -(2**54-1)
+    self.helper_test_variable(Variable("a", 0, -c1)*c0 < c1, 0, 1, f"((a*-1)<{-(2**53-1)})")
     self.helper_test_variable(Variable("a", 0, 5)*(-2) < 0, 0, 1, "((a*-1)<0)")
     self.helper_test_variable(Variable("a", 0, 5)*4 >= 12, 0, 1, "((a<3)!=True)")
     self.helper_test_variable(Variable("a", 0, 5)*4 >= 13, 0, 1, "((a<4)!=True)")

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_dtype_weak.py

@@ -117,6 +117,13 @@ class TestWeakPromotion(unittest.TestCase):
       t.realize()
       self.assertNotIn(t.uop.buffer.dtype, dtypes.weaks)
 
+  def test_computed_float_index_lowers(self):
+    # a half-pixel nearest index resolves its float-scaled range before the gather
+    idx = (Tensor.arange(8) + 0.5) / 4 - 0.5
+    idx = (idx.clip(0, 1) - 0.5).ceil().int()
+    out = Tensor([0, 1], device="NULL")[idx].contiguous().realize()
+    self.assertNotIn(out.uop.buffer.dtype, dtypes.weaks)
+
 
 class TestWeakStorageBoundary(unittest.TestCase):
   # weak has no storage: a weak assignment source casts when it defers to the destination, everything else raises

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)

test/unit/test_multitensor.py

@@ -597,7 +597,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
     t = Tensor.arange(64).reshape(8, 8).clone().realize()
     t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0)
 
-    with self.assertRaises(AssertionError):
+    with self.assertRaises(RuntimeError):
       # sharded axis shrink on non-device boundry is not allowed
       a = t.shrink(((0, 3), (0, 8))).contiguous()
       a.schedule_linear()

tinygrad/callify.py

@@ -83,7 +83,7 @@ def contiguous_mops_to_view(c:UOp, src:UOp):
     resolved = graph_rewrite(src, multi_pm, name="multi_buffer_view")
     if resolved.op is not Ops.UNSHARD: return None
     if (view := _make_buffer_view(resolved.src[0])) is None: return None
-    return view.reshape(resolved.src[0].shape).unshard(resolved.arg, resolved.src[1]).contiguous(tag=c.tag)
+    return view.reshape(resolved.src[0].shape).unshard(resolved.arg, resolved.src[1:]).contiguous(tag=c.tag)
 
   return None
 

tinygrad/codegen/__init__.py

@@ -21,6 +21,7 @@ from tinygrad.codegen.late.coalesce import indexing_simplify
 from tinygrad.codegen.opt.postrange import apply_opts
 from tinygrad.codegen.late.gater import pm_move_gates_from_index
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse
+from tinygrad.schedule.multi import multi_pm
 from tinygrad.schedule.rangeify import pm_mops
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
 from tinygrad.codegen.late.regalloc import LinearScanRegallocContext, pm_regalloc_rewrite
@@ -119,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 = []
@@ -186,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)
@@ -219,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)
@@ -261,13 +273,13 @@ def add_raw_barrier(after:UOp):
   # loads from a LOCAL buffer that depend (via AFTER) on stores to LOCAL memory need a workgroup barrier
   if after.addrspace is not AddrSpace.LOCAL: return None
   # one toposort over all the deps
-  deps = UOp.sink(*after.src[1:]).backward_slice
-  if not any(_is_local_store(x) for x in deps) or any(x.op is Ops.BARRIER for x in deps): return None
+  deps = UOp.sink(*after.src[1:]).toposort(gate=lambda x: x.op is not Ops.BARRIER)
+  if not any(_is_local_store(x) for x in deps): return None
   return after.src[0].after(UOp(Ops.BARRIER, src=after.src[1:]))
 
 def add_war_barrier(end:UOp):
   # a LOCAL buffer stored and loaded in the same loop needs a barrier at the end of the loop body
-  rngs = [r for r in end.src[1:] if r.op is Ops.RANGE and r.arg[1] in (AxisType.REDUCE, AxisType.LOOP) and r.vmax > 0]
+  rngs = [r for r in end.src[1:] if r.op is Ops.RANGE and r.arg[1] in (AxisType.REDUCE, AxisType.WEAK, AxisType.LOOP) and r.vmax > 0]
   if not rngs or end.src[0].op is Ops.BARRIER: return None
   sl = end.src[0].backward_slice_with_self
   # only stores that are inside this loop body (not in the backward slice through AFTER chains from other loops)
@@ -286,8 +298,11 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   if DEBUG >= 5: print(pyrender(ast))
   if SPEC: type_verify(ast, spec_tensor)
 
+  # resolve UNSHARDs (multi-device UNSHARDs are already resolved by the scheduler; this handles in-kernel shards, e.g. fragments)
+  sink = graph_rewrite(ast, multi_pm, name="multi_pm")
+
   # preprocess
-  sink = graph_rewrite(ast, pm_mops, name="early movement ops", bottom_up=True)
+  sink = graph_rewrite(sink, pm_mops, name="early movement ops", bottom_up=True)
 
   # first we optimize
   if optimize:
@@ -313,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")
@@ -326,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")
@@ -336,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")
@@ -464,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):
@@ -478,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/codegen/opt/heuristic.py

@@ -169,7 +169,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
     else:
       # prioritize making expand axes local
       local_axis_ranking = [(any(k.rngs[axis] not in b.src[1].get_idx().backward_slice for b in k.bufs), axis) \
-                              for axis in k.axes_of(AxisType.GLOBAL, AxisType.LOOP) if k.rngs[axis].src[0].op is Ops.CONST]
+                              for axis in k.axes_of(AxisType.GLOBAL, AxisType.WEAK) if k.rngs[axis].src[0].op is Ops.CONST]
       to_local: list[tuple[int, int]] = []
       for _, axis in sorted(local_axis_ranking, key=lambda x: (-x[0], -x[1])):
         local_size = prod(sz for _, sz in to_local)
@@ -188,7 +188,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
     for threads in [32,16,12,8,6,5,4,3,2]:
       # Skip if too many threads. Heuristic: use about 128K ops per thread
       if threads > k.ren.global_max[0] or resolve(prod(k.full_shape) // (128 << 10) < threads): continue
-      for axis in k.axes_of(AxisType.LOOP):
+      for axis in k.axes_of(AxisType.WEAK):
         if k.full_shape[axis] % threads == 0:
           try: k.apply_opt(Opt(OptOps.THREAD, axis, threads))
           except KernelOptError: pass

tinygrad/codegen/opt/postrange.py

@@ -65,7 +65,7 @@ class Scheduler:
   def _output_rngs(self) -> list[UOp]:
     return flatten([[r for r in UOp.sink(*s.src[1:]).ranges if r.arg[-1] != AxisType.REDUCE] for s in self.ast.src if s.op is Ops.END])
   def _globalizable_rngs(self) -> list[UOp]:
-    ret = [r for r in self._output_rngs() if r.arg[-1] == AxisType.LOOP]
+    ret = [r for r in self._output_rngs() if r.arg[-1] == AxisType.WEAK]
     # exclude any output ranges from global that don't appear in all BUFFERIZE
     for x in self.ast.toposort():
       if x.op is Ops.STAGE:
@@ -86,8 +86,8 @@ class Scheduler:
     ret = []
     for x,r in zip(self.axis_types, self.rngs):
       if self.dont_use_locals and x == AxisType.GLOBAL: ret.append("BLUE")
-      elif r not in output_rngs and x == AxisType.LOOP: ret.append("BLACK")
-      elif r not in globalizible_rngs and x == AxisType.LOOP: ret.append("white")
+      elif r not in output_rngs and x == AxisType.WEAK: ret.append("BLACK")
+      elif r not in globalizible_rngs and x == AxisType.WEAK: ret.append("white")
       else: ret.append(axis_colors[x])
     return ret
   def colored_shape(self) -> str: return ' '.join([colored(f'{x.src[0].render():>4s}', color) for x,color in zip(self.rngs, self.colors())])
@@ -108,7 +108,7 @@ class Scheduler:
 
   # copied from kernel.py
   @property
-  def upcastable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP) \
+  def upcastable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GLOBAL, AxisType.LOCAL, AxisType.WEAK) \
                                                   if isinstance(s:=self.full_shape[i], int) and s > 1]
   @property
   def unrollable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE) \
@@ -161,10 +161,10 @@ class Scheduler:
         check(rng.arg[-1] in {AxisType.GROUP_REDUCE, AxisType.REDUCE}, "unroll is for GROUP_REDUCE/REDUCE")
       if opt.op is OptOps.UPCAST:
         check((self.ren is not None and self.ren.target.device == "DSP") or amt <= 16, "don't upcast more than 16")
-        check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP}, f"upcast is for GLOBAL/LOCAL/LOOP, not {rng.arg[-1]}")
+        check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOCAL, AxisType.WEAK}, f"upcast is for GLOBAL/LOCAL/LOOP, not {rng.arg[-1]}")
       if opt.op is OptOps.LOCAL:
         check(not self.dont_use_locals, "can't use locals")
-        check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOOP}, "local is for globals")
+        check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.WEAK}, "local is for globals")
       if opt.op is OptOps.THREAD:
         check(self.ren is not None and self.ren.has_threads, "target does not support threads")
         check(self.ren is not None and self.ren.global_max is not None and amt <= self.ren.global_max[0], "too many threads")

tinygrad/device.py

@@ -5,7 +5,7 @@ from typing import Any, Generic, TypeVar, Iterator, Generator, Self, TYPE_CHECKI
 import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal
 from tinygrad.helpers import LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, PROFILE, temp, colored
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing
-from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize, Target, unwrap
+from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize, Target, unwrap, round_up
 from tinygrad.dtype import DType, _to_np_dtype
 if TYPE_CHECKING: from tinygrad.renderer import Renderer
 
@@ -318,6 +318,13 @@ 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]:
+    for _,_,dt,_ in signature:
+      yield (offset:=round_up(offset, dt.itemsize)), dt
+      offset += dt.itemsize
 
 class Program(Generic[DeviceType]):
   def __init__(self, dev:DeviceType, obj:TinyELF): pass

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/nn/__init__.py

@@ -337,7 +337,7 @@ def _embedding_bwd(grad_emb:UOp, call:UOp) -> tuple:
     BLOCK_J = min(256, embed_size)
     n_j_blocks = (embed_size + BLOCK_J - 1) // BLOCK_J
     i = UOp.range(grad_emb_flat.shape[0], 0)         # batch_size * sequence_length -> GLOBAL
-    j_inner = UOp.range(BLOCK_J, 2, AxisType.LOOP if device in ("CPU", "NULL") else AxisType.LOCAL)  # BLOCK_J threads per workgroup
+    j_inner = UOp.range(BLOCK_J, 2, AxisType.WEAK if device in ("CPU", "NULL") else AxisType.LOCAL)  # BLOCK_J threads per workgroup
     j_outer = UOp.range(n_j_blocks, 1)
     j = j_outer * BLOCK_J + j_inner
     # mask padded embed

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/amd/elf.py

@@ -35,10 +35,9 @@ def assemble_linear(prg:UOp, lin:UOp, arch:str) -> bytes:
       elif val.offset < 106: max_sgpr = max(max_sgpr, val.offset + val.sz)
 
   # ** scan sink for metadata
-  sink, n_bufs, n_vars, lds_size, gids = prg.src[0], 0, 0, 0, set()
+  sink, param_sizes, lds_size, gids = prg.src[0], {}, 0, set()
   for u in sink.toposort():
-    if u.op is Ops.PARAM and u.addrspace is AddrSpace.ALU: n_vars += 1
-    elif u.op is Ops.PARAM: n_bufs += 1
+    if u.op is Ops.PARAM: param_sizes[u.arg.slot] = u.dtype.itemsize if u.addrspace is AddrSpace.ALU else 8
     elif u.op is Ops.BUFFER and u.addrspace is AddrSpace.LOCAL: lds_size += u.max_numel() * u.dtype.itemsize
     elif u.op is Ops.SPECIAL and u.arg.startswith("gidx"): gids.add(int(u.arg[-1]))
   code_bytes = b"".join(inst.to_bytes() for inst in insts)
@@ -60,7 +59,7 @@ def assemble_linear(prg:UOp, lin:UOp, arch:str) -> bytes:
   sgpr_granule = max(0, ceildiv(next_free_sgpr + 6, 8) - 1) if is_cdna else 0
   desc = amdgpu_kd.llvm_amdhsa_kernel_descriptor_t()
   desc.group_segment_fixed_size = lds_size
-  desc.kernarg_size = n_bufs * 8 + n_vars * 4
+  for sz in (param_sizes[i] for i in sorted(param_sizes)): desc.kernarg_size = round_up(desc.kernarg_size, sz) + sz
   desc.kernel_code_entry_byte_offset = -len(text)
 
   # https://llvm.org/docs/AMDGPUUsage.html#amdgpu-amdhsa-compute-pgm-rsrc1-gfx6-gfx12-table

tinygrad/renderer/cstyle.py

@@ -123,7 +123,8 @@ class CStyleLanguage(Renderer):
   smem_align: str = ""
   smem_prefix: str = ""
   smem_prefix_for_cast: bool = True
-  arg_int_prefix: str = "const int"
+  var_prefix: str = "const "
+  var_suffix: str = ""
   barrier: str = ""
   code_for_workitem: dict[Literal["g", "l", "i"], Callable] = {}
   extra_args: list[str] = []
@@ -149,9 +150,9 @@ class CStyleLanguage(Renderer):
     tmp = ""
     if any(is_image_shape(u._shape) for _,(u,_) in bufs):
       tmp = "const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;\n"
-    buftypes = [(name, ("volatile " if u.arg.volatile else "")+
-                       self._render_dtype(u.dtype, sz=1, addrspace=u.addrspace, mutable=mutable, shape=u._shape)+self.buffer_suffix \
-                 if u.addrspace == AddrSpace.GLOBAL else self.arg_int_prefix if u.dtype == dtypes.int else None) for name,(u,mutable) in bufs]
+    buftypes = [(name, ("volatile " if u.arg.volatile else "")+(self.var_prefix if u.addrspace == AddrSpace.ALU else "")+
+                       self._render_dtype(u.dtype, sz=1, addrspace=u.addrspace, mutable=mutable, shape=u._shape)+
+                       (self.var_suffix if u.addrspace == AddrSpace.ALU else self.buffer_suffix)) for name,(u,mutable) in bufs]
     local_dims = [u.src[0] for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]
     launch_bounds = prod([d.vmax for d in local_dims])
     prg = ''.join([f"{self.kernel_typedef.format(launch_bounds=launch_bounds)} {function_name}(",] +
@@ -258,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("")'
@@ -269,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
@@ -348,7 +459,8 @@ class MetalRenderer(CStyleLanguage):
   kernel_typedef = "kernel void"
   buffer_prefix = "device "
   smem_prefix = "threadgroup __attribute__((aligned(16))) "
-  arg_int_prefix = "constant int&"
+  var_prefix = "constant "
+  var_suffix = "&"
   barrier = "threadgroup_barrier(mem_flags::mem_threadgroup);"
   float4 = "float4"
   code_for_workitem = {"g": lambda x: f"gid.{chr(120+int(x))}", "l": lambda x: f"lid.{chr(120+int(x))}"}

tinygrad/renderer/nir.py

@@ -1,6 +1,6 @@
 from typing import Callable, Any
 from tinygrad.dtype import AddrSpace, DType, dtypes, truncate
-from tinygrad.helpers import DEBUG, OSX, unwrap, fromimport, Target, is_image_shape
+from tinygrad.helpers import DEBUG, OSX, unwrap, fromimport, Target, is_image_shape, round_up
 from tinygrad.renderer import Renderer
 from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str
@@ -246,9 +246,11 @@ class NIRRenderer(Renderer):
 
   def supported_dtypes(self): return {d for d in Renderer.supported_dtypes(self) if d not in dtypes.fp8s+(dtypes.bfloat16,)}
 
+def padded_idx(param_idx:int, size:int): return round_up(param_idx, size) + size
+
 class NAKRenderer(NIRRenderer):
-  param = nir_instr(nc=1, num_components=1, bs=lambda sz:sz*8, also=lambda self,sz: setattr(self, "param_idx", self.param_idx + sz),
-    intrins={"ALIGN_MUL":lambda sz:sz}, srcs=lambda self,b: [nsrc(nimm(b, 0, dtypes.int)), nsrc(nimm(b, self.param_idx, dtypes.int))])(
+  param = nir_instr(nc=1, num_components=1, bs=lambda sz:sz*8, also=lambda self,sz: setattr(self, "param_idx", padded_idx(self.param_idx, sz)),
+    intrins={"ALIGN_MUL":lambda sz:sz}, srcs=lambda self,b,sz: [nsrc(nimm(b,0,dtypes.int)), nsrc(nimm(b, round_up(self.param_idx,sz), dtypes.int))])(
        lambda self, b, x, sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_ldc_nv))
 
   def supported_dtypes(self): return {d for d in super().supported_dtypes() if (d != dtypes.half or int(self.target.arch[3:]) >= 53)}
@@ -263,12 +265,13 @@ class LVPRenderer(NIRRenderer):
   code_for_op = {k:v for k,v in NIRRenderer.code_for_op.items() if k != Ops.EXP2}
 
   param = nir_instr(nc=1, bs=lambda sz: sz * 8, num_components=1, intrins={"ALIGN_MUL":lambda sz: sz, "RANGE":lambda self: self.param_sz},
-    srcs=lambda b, self: [nsrc(nimm(b, 0, dtypes.int)), nsrc(nimm(b, self.param_idx, dtypes.int))], also=lambda self, sz:
-    setattr(self, "param_idx", self.param_idx+sz))(lambda self,b,x,sz: mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_ubo))
+    srcs=lambda b,self,sz: [nsrc(nimm(b, 0, dtypes.int)), nsrc(nimm(b, round_up(self.param_idx, sz), dtypes.int))], also=lambda self, sz:
+    setattr(self, "param_idx", padded_idx(self.param_idx, sz)))(lambda self,b,x,sz:
+                                                                mesa.nir_intrinsic_instr_create(b.shader, mesa.nir_intrinsic_load_ubo))
 
   def prerender(self, uops:list[UOp]):
     super().prerender(uops)
-    self.param_sz = sum([u.dtype.itemsize if u.addrspace is AddrSpace.ALU else 8 for u in uops if u.op is Ops.PARAM])
+    self.param_sz = functools.reduce(padded_idx, (u.element_size() if u.addrspace is AddrSpace.ALU else 8 for u in uops if u.op is Ops.PARAM), 0)
 
 def tovec(b, idx_y, idx_x): return nalu(b, "vec4", idx_x, idx_y, nundef(b, dtypes.int), nundef(b, dtypes.int))
 def nfloat(dtype): return mesa.nir_type_float16 if dtype == dtypes.half else mesa.nir_type_float32
@@ -306,7 +309,8 @@ class IR3Renderer(NIRRenderer):
     super().prerender(uops)
     self.texs:set[UOp] = set()
     self.img_idx = 0
-    self.param_sz = sum([u.dtype.itemsize if u.addrspace is AddrSpace.ALU else 8 for u in uops if u.op is Ops.PARAM])
+    self.param_sz = functools.reduce(padded_idx, (u.element_size() if u.addrspace is AddrSpace.ALU else 8
+                                                 for u in uops if u.op is Ops.PARAM and not is_image_shape(u._shape)), 0)
 
   def postrender(self, uops:list[UOp]):
     bufs = [u for u in uops if u.op is Ops.PARAM and u.addrspace is not AddrSpace.ALU]

tinygrad/renderer/ptx.py

@@ -145,6 +145,8 @@ class PTXRenderer(Renderer):
     from tinygrad.runtime.support.compiler_cuda import NVPTXCompiler, PTXCompiler
     self.compiler = (PTXCompiler if target.interface.startswith("MOCK") or target.device == "CUDA" else NVPTXCompiler)(target.arch)
     self.tensor_cores = PTXRenderer.tc_sm80 if (ver:=int(target.arch[3:])) >= 80 else tc.cuda_sm75 if ver >= 75 else []
+    if ver < 80: self.extra_matcher += PatternMatcher([(UPat((Ops.MAX, Ops.EXP2), dtype=dtypes.half, name="x"),
+      lambda x: UOp(x.op, src=tuple(vv.cast(dtypes.float32) for vv in x.src), arg=x.arg).cast(dtypes.half))])
 
   # language options
   kernel_prefix = """.version VERSION
@@ -199,6 +201,8 @@ class PTXRenderer(Renderer):
         continue
       if u.op in {Ops.INDEX, Ops.SHRINK, Ops.LOAD} and u.src[0].addrspace in (AddrSpace.REG, AddrSpace.ALU):
         # on REG, INDEX/SHRINK pick the register (must be CONST) and LOAD is a noop
+        if u.op is not Ops.LOAD and u.src[1].op is not Ops.CONST:
+          raise RuntimeError(f"PTX does not support dynamic register indexing: {u}")
         r[u] = r[u.src[0]] if u.op is Ops.LOAD else r[u.src[0]][u.src[1].arg]
         continue
       if u.op is Ops.SPECIAL: r[u] = "%" + u.arg

tinygrad/runtime/graph/cuda.py

@@ -20,7 +20,7 @@ class CUDAGraph(MultiGraphRunner):
         global_size, local_size = ast.arg.launch_dims({v: 0 for v in self.vars})
 
         c_deps, new_node = self.new_node([b.base for b in bufs], ast.arg.outs)
-        c_args, vargs = encode_args([b._buf for b in bufs], [device_vars.get(x.expr, 0) for x in ast.arg.vars])
+        c_args, vargs = encode_args([b._buf for b in bufs], [device_vars.get(x.expr, 0) for x in ast.arg.vars], runtime.signature)
         kern_params = cuda.CUDA_KERNEL_NODE_PARAMS_v1(runtime.prg, *global_size, *local_size, runtime.smem,
                                                       ctypes.cast(0, ctypes.POINTER(ctypes.c_void_p)), vargs)
         check(cuda.cuGraphAddKernelNode(ctypes.byref(new_node), self.graph, c_deps, len(c_deps or []), ctypes.byref(kern_params)))

tinygrad/runtime/graph/metal.py

@@ -1,7 +1,6 @@
 from typing import Any, cast
-import ctypes, decimal
-from tinygrad.dtype import dtypes
-from tinygrad.helpers import dedup, getenv, PROFILE
+import ctypes, decimal, struct
+from tinygrad.helpers import dedup, getenv, unwrap, PROFILE
 from tinygrad.device import Buffer, Device, ProfileGraphEntry, ProfileGraphEvent
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.engine.jit import GraphRunner, GraphException
@@ -25,9 +24,12 @@ class MetalGraph(GraphRunner):
     if self.icb.value is None: raise GraphException("create indirect command buffer failed, does your system support this?")
     self.needs_icb_fix = int(not self.dev.arch.startswith("Apple") or int(self.dev.arch[5:]) < 9)  # ICB fix not required on M3+ (Apple9+)
 
-    if len(self.vars): self.int_buf = self.dev.allocator.alloc(len(self.vars)*dtypes.int32.itemsize)
+    self.var_bind_data = []
+    if len(self.vars):
+      self.var_buf = self.dev.allocator.alloc(sum(dt.itemsize for r in self.runtimes for (_,_,dt,s) in unwrap(r).signature if s == ()))
+      self.var_buf_view, var_buf_offset = cast(MetalAllocator, self.dev.allocator)._as_buffer(self.var_buf), 0
 
-    all_pipelines, all_resources = [], [self.int_buf.buf] if len(self.vars) else []
+    all_pipelines, all_resources = [], [self.var_buf.buf] if len(self.vars) else []
     for j, ((_, ast, bufs, _), runtime, replace) in enumerate(zip(self.calls, self.runtimes, self.uop_replace)):
       assert runtime is not None
       icb_command = self.icb.indirectComputeCommandAtIndex(j).retained()
@@ -37,7 +39,10 @@ class MetalGraph(GraphRunner):
         if not any(pos == i for pos, _ in replace):
           icb_command.setKernelBuffer_offset_atIndex(b._buf.buf, b._buf.offset, i)
           all_resources.append(b._buf.buf)
-      for i, v in enumerate(ast.arg.vars): icb_command.setKernelBuffer_offset_atIndex(self.int_buf.buf, self.vars.index(v.expr)*4, len(bufs)+i)
+      for nm,i,dt,_ in runtime.signature[len(bufs):]:
+        icb_command.setKernelBuffer_offset_atIndex(self.var_buf.buf, var_buf_offset, i)
+        self.var_bind_data.append((nm, var_buf_offset, dt.fmt))
+        var_buf_offset += dt.itemsize
       global_size, local_size = ast.arg.launch_dims({v: 0 for v in self.vars})
       icb_command.concurrentDispatchThreadgroups_threadsPerThreadgroup(metal.MTLSize(*global_size), metal.MTLSize(*local_size))
       icb_command.setBarrier()
@@ -45,7 +50,6 @@ class MetalGraph(GraphRunner):
     self.all_resources = dedup(all_resources)
     self.all_pipelines = dedup(all_pipelines)
     self.command_buffer: Any = None
-    if len(self.vars): self.int_buf_view = cast(MetalAllocator, self.dev.allocator)._as_buffer(self.int_buf).cast('i')
     self.range = metal.NSRange(0, len(self.calls))
     self.updatable = sorted({j for j,r in enumerate(self.uop_replace) if r} | self.var_vals_replace.keys() | self.launch_dims_replace.keys())
 
@@ -66,7 +70,7 @@ class MetalGraph(GraphRunner):
     for j, global_dims, local_dims in self.updated_launch_dims(var_vals):
       self.icb.indirectComputeCommandAtIndex(j).concurrentDispatchThreadgroups_threadsPerThreadgroup(metal.MTLSize(*global_dims),
                                                                                                      metal.MTLSize(*local_dims))
-    for i, var in enumerate(self.vars): self.int_buf_view[i] = var_vals[var]
+    for nm,ofs,fmt in self.var_bind_data: struct.pack_into(fmt, self.var_buf_view, ofs, var_vals[nm])
 
     command_buffer = self.dev.mtl_queue.commandBuffer().retained()
     encoder = command_buffer.computeCommandEncoder().retained()

tinygrad/runtime/ops_amd.py

@@ -602,8 +602,7 @@ class AMDProgram(HCQProgram['AMDDevice']):
 
     if dev.sqtt_enabled: self.libhash: tuple[int, int] = struct.unpack('<Q', hashlib.md5(self.lib).digest()[:8])*2
 
-    super().__init__(CLikeArgsState, self.dev, self.name, kernargs_alloc_size=self.kernargs_segment_size+additional_alloc_sz, lib=self.lib,
-                     base=self.lib_gpu.va_addr)
+    super().__init__(CLikeArgsState, self.dev, obj, kernargs_alloc_size=self.kernargs_segment_size+additional_alloc_sz, base=self.lib_gpu.va_addr)
     weakref.finalize(self, self._fini, self.dev, self.lib_gpu, buf_spec)
 
   def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int|None, ...]=(),

tinygrad/runtime/ops_cl.py

@@ -55,7 +55,7 @@ class CLProgram(Program['CLDevice']):
   def __call__(self, *bufs:cl.cl_mem, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]|None=None, vals:tuple[int, ...]=(),
                wait=False, **kw) -> float|None:
     for i, (_, slot, dt, shape) in enumerate(self.signature):
-      b = bufs[slot] if slot < len(bufs) else ctypes.c_int32(vals[slot-len(bufs)])
+      b = bufs[slot] if slot < len(bufs) else getattr(ctypes, f"c_int{dt.bitsize}")(vals[slot-len(bufs)])
       if is_image_shape(shape):
         pitch = (round_up(shape[1], 256) if OSX else shape[1]) * 4 * dt.itemsize
         fmt = cl.cl_image_format(cl.CL_RGBA, {2:cl.CL_HALF_FLOAT, 4:cl.CL_FLOAT}[dt.itemsize])

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])
@@ -105,7 +226,8 @@ class CPUProgram(HCQProgram['CPUDevice']):
   except OSError: pass
 
   def __init__(self, dev:CPUDevice, obj:TinyELF):
-    self.runtimevars = {name:slot for name,slot,*_ in obj.signature if name == 'core_id'}
+    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
@@ -140,7 +262,7 @@ class CPUProgram(HCQProgram['CPUDevice']):
 
       self.fxn = ctypes.CFUNCTYPE(None)(self.addr)
 
-    super().__init__(LVPArgsState if LVP else HCQArgsState, dev, obj.name, kernargs_alloc_size=12+256 if LVP else 0)
+    super().__init__(LVPArgsState if LVP else HCQArgsState, dev, obj, kernargs_alloc_size=12+256 if LVP else 0)
 
   @suppress_finalizing
   def __del__(self):
@@ -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/ops_cuda.py

@@ -15,9 +15,10 @@ def check(status):
     error = ctypes.string_at(init_c_var(ctypes.POINTER(ctypes.c_char), lambda x: cuda.cuGetErrorString(status, ctypes.byref(x)))).decode()
     raise RuntimeError(f"CUDA Error {status}, {error}")
 
-def encode_args(args, vals) -> tuple[ctypes.Structure, ctypes.Array]:
-  c_args = init_c_struct_t(len(args) * 8 + len(vals) * 4, tuple([(f'f{i}', cuda.CUdeviceptr_v2, i*8) for i in range(len(args))] +
-                                                                [(f'v{i}', ctypes.c_int, len(args)*8 + i*4) for i in range(len(vals))]))(*args, *vals)
+def encode_args(args, vals, signature) -> tuple[ctypes.Structure, ctypes.Array]:
+  fields = ([(f'f{i}', cuda.CUdeviceptr_v2, i*8) for i in range(len(args))] +
+            [(f'v{i}', getattr(ctypes, f"c_int{dt.bitsize}"), off) for i,(off,dt) in enumerate(TinyELF.iter_sig(signature[len(args):], len(args)*8))])
+  c_args = init_c_struct_t(fields[-1][2] + ctypes.sizeof(fields[-1][1]) if len(fields) else 0, tuple(fields))(*args, *vals)
   vargs = (ctypes.c_void_p * 5)(ctypes.c_void_p(1), ctypes.cast(ctypes.byref(c_args), ctypes.c_void_p), ctypes.c_void_p(2),
                                 ctypes.cast(ctypes.pointer(ctypes.c_size_t(ctypes.sizeof(c_args))), ctypes.c_void_p), ctypes.c_void_p(0))
   return c_args, vargs
@@ -35,7 +36,7 @@ def cu_time_execution(cb, enable=False) -> float|None:
 
 class CUDAProgram(Program['CUDADevice']):
   def __init__(self, dev:CUDADevice, obj:TinyELF, smem:int=0):
-    self.dev, self.name, self.lib, self.smem = dev, obj.name, obj.lib, smem
+    self.dev, self.name, self.lib, self.signature, self.smem = dev, obj.name, obj.lib, obj.signature, smem
     if DEBUG >= 5: print("\n".join([f"{i+1:>3} {line}" for i, line in enumerate(pretty_ptx(obj.lib.decode('utf-8')).split("\n"))]))
 
     check(cuda.cuCtxSetCurrent(self.dev.context))
@@ -54,7 +55,7 @@ class CUDAProgram(Program['CUDADevice']):
   def __call__(self, *args, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
     check(cuda.cuCtxSetCurrent(self.dev.context))
     if not hasattr(self, "vargs"):
-      self.c_args, self.vargs = encode_args(args, vals)
+      self.c_args, self.vargs = encode_args(args, vals, self.signature)
 
       # HACK: For MOCKGPU send the args struct itself.
       if MOCKGPU: self.vargs = self.c_args # type: ignore[assignment]

tinygrad/runtime/ops_dsp.py

@@ -4,7 +4,7 @@ assert sys.platform != 'win32'
 from tinygrad.device import BufferSpec, Compiled, Allocator, Compiler, Program, TinyELF
 from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.uop.ops import Ops, UOp
-from tinygrad.helpers import getenv, round_up, mv_address, to_mv, cpu_objdump, system, DEBUG, suppress_finalizing, Target
+from tinygrad.helpers import getenv, round_up, mv_address, to_mv, cpu_objdump, system, DEBUG, suppress_finalizing, Target, unwrap
 from tinygrad.renderer.cstyle import ClangRenderer
 from tinygrad.runtime.autogen import libc, qcom_dsp
 if getenv("IOCTL"): import extra.dsp.run # noqa: F401 # pylint: disable=unused-import
@@ -38,7 +38,9 @@ class DSPRenderer(ClangRenderer):
             'struct dcvs_v2_req req = {.type=7, .dcvs_enable=0, .set_latency=1, .latency=100, .set_dcvs_params=1, .target_corner = 6 /* TURBO */};',
             'HAP_power_set((void*)handle, (void*)&req);']
     msrc += ['if ((sc>>24) != 2) return 0;']
-    msrc += [f'int sz_or_val_{i} = ((int*)pra[0].buf.pv)[{i}];' for i,b in enumerate(bufs)]
+    msrc += [f'{self._render_dtype(b[1][0].dtype) if b[1][0].addrspace == AddrSpace.ALU else "int"} sz_or_val_{i} = '
+             f'*({self._render_dtype(b[1][0].dtype) if b[1][0].addrspace == AddrSpace.ALU else "int"}*)((char*)pra[0].buf.pv+{i*8});'
+             for i,b in enumerate(bufs)]
     msrc += [f'int off{i} = ((int*)pra[1].buf.pv)[{i}];' for i,b in enumerate(bufs) if b[1][0].addrspace == AddrSpace.GLOBAL]
     msrc += [f'void *buf_{i} = HAP_mmap(0,sz_or_val_{i},3,0,pra[{i+3}].dma.fd,0)+off{i};'
              for i,b in enumerate(bufs) if b[1][0].addrspace == AddrSpace.GLOBAL]
@@ -64,15 +66,15 @@ def rpc_prep_args(ins=None, outs=None, in_fds=None):
   return pra, fds, attrs, (ins, outs)
 
 class DSPProgram(Program['DSPDevice']):
-  def __init__(self, dev:DSPDevice, obj:TinyELF):
-    self.dev, self.lib = dev, obj.lib
+  def __init__(self, dev:DSPDevice, obj:TinyELF): self.dev, self.lib, self.signature = dev, obj.lib, obj.signature
 
   def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
     if len(bufs) >= 16: raise RuntimeError(f"Too many buffers to execute: {len(bufs)}")
 
-    pra, fds, attrs, _ = rpc_prep_args(ins=[var_vals_mv:=memoryview(bytearray((len(bufs)+len(vals))*4)), off_mv:=memoryview(bytearray(len(bufs)*4))],
+    pra, fds, attrs, _ = rpc_prep_args(ins=[var_vals_mv:=memoryview(bytearray((len(bufs)+len(vals))*8)), off_mv:=memoryview(bytearray(len(bufs)*4))],
                                        outs=[timer:=memoryview(bytearray(8)).cast('Q')], in_fds=[b.share_info.fd for b in bufs])
-    var_vals_mv.cast('i')[:] = array.array('i', tuple(b.size for b in bufs) + vals)
+    for i,b in enumerate(bufs): struct.pack_into('i', var_vals_mv, i*8, b.size)
+    for i,(v,(_,_,dt,_)) in enumerate(zip(vals, self.signature[len(bufs):]), start=len(bufs)): struct.pack_into(unwrap(dt.fmt), var_vals_mv, i*8, v)
     off_mv.cast('I')[:] = array.array('I', tuple(b.offset for b in bufs))
     self.dev.exec_lib(self.lib, rpc_sc(method=2, ins=2, outs=1, fds=len(bufs)), pra, fds, attrs)
     return timer[0] / 1e6
@@ -266,7 +268,7 @@ class MockDSPRenderer(DSPRenderer):
         # for loop for big reads
         msrc.append(f"void *buf{i} = mmap2(0, {sz}, 3, 0x21, -1, 0); for(int rd = 0; rd < {sz}; rd += read(0, buf{i}+rd, {sz}-rd));")
       else:
-        msrc.append(f"unsigned int val{i}; read(0, &val{i}, 4);")
+        msrc.append(f"{self._render_dtype(b[1][0].dtype)} val{i}; read(0, &val{i}, {b[1][0].dtype.itemsize});")
     msrc.append("unsigned int st = inscount();")
     params = [(f'(void*)buf{i}' if b[1][0].addrspace == AddrSpace.GLOBAL else f'val{i}') for i,b in enumerate(bufs)]
     msrc.append(f"{function_name}({', '.join(params)});")
@@ -277,14 +279,16 @@ class MockDSPRenderer(DSPRenderer):
     return '\n'.join(msrc)
 
 class MockDSPProgram(Program[DSPDevice]):
-  def __init__(self, dev:DSPDevice, obj:TinyELF): self.lib = obj.lib
+  def __init__(self, dev:DSPDevice, obj:TinyELF): self.lib, self.signature = obj.lib, obj.signature
   def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
     with tempfile.NamedTemporaryFile(suffix=".out") as dsp_lib:
       dsp_lib.write(self.lib)
       dsp_lib.flush()
       os.chmod(dsp_lib.name, 0o0777)
       proc = subprocess.run(["qemu-hexagon-static", *(['-strace'] if DEBUG >= 5 else []), dsp_lib.name],
-        input=b''.join([bytes(to_mv(x.va_addr, x.size)) for x in bufs] + [struct.pack("I", x) for x in vals]), stdout=subprocess.PIPE, check=True)
+        input=b''.join([bytes(to_mv(x.va_addr, x.size)) for x in bufs] +
+                       [struct.pack(unwrap(dt.fmt), x) for x,(_,_,dt,_) in zip(vals, self.signature[len(bufs):])]),
+        stdout=subprocess.PIPE, check=True)
     offset = 4
     for x in bufs:
       to_mv(x.va_addr, x.size)[:] = proc.stdout[offset:offset+x.size]

tinygrad/runtime/ops_hip.py

@@ -25,7 +25,7 @@ class HIPDevice(Compiled):
 
 class HIPProgram(Program[HIPDevice]):
   def __init__(self, dev:HIPDevice, obj:TinyELF):
-    self.dev, self.name, self.lib = dev, obj.name, obj.lib
+    self.dev, self.name, self.lib, self.signature = dev, obj.name, obj.lib, obj.signature
     check(hip.hipSetDevice(self.dev.device_id))
     self.module = init_c_var(hip.hipModule_t, lambda x: check(hip.hipModuleLoadData(ctypes.byref(x), obj.lib)))
     self.prg = init_c_var(hip.hipFunction_t, lambda x: check(hip.hipModuleGetFunction(ctypes.byref(x), self.module, obj.name.encode("utf-8"))))
@@ -37,8 +37,9 @@ class HIPProgram(Program[HIPDevice]):
   def __call__(self, *args, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
     check(hip.hipSetDevice(self.dev.device_id))
     if not hasattr(self, "vargs"):
-      fields = [(f'f{i}', hip.hipDeviceptr_t, i*8) for i in range(len(args))] + [(f'v{i}', ctypes.c_int, len(args)*8+i*4) for i in range(len(vals))]
-      self.c_args = init_c_struct_t(len(args)*8+len(vals)*4, tuple(fields))(*args, *vals)
+      fields = ([(f'f{i}', hip.hipDeviceptr_t, i*8) for i in range(len(args))] +
+        [(f'v{i}', getattr(ctypes, f"c_int{dt.bitsize}"), o) for i,(o,dt) in enumerate(TinyELF.iter_sig(self.signature[len(args):], len(args)*8))])
+      self.c_args = init_c_struct_t(fields[-1][2] + ctypes.sizeof(fields[-1][1]) if len(fields) else 0, tuple(fields))(*args, *vals)
       self.vargs = (ctypes.c_void_p * 5)(1, ctypes.cast(ctypes.byref(self.c_args), ctypes.c_void_p), 2,
                                          ctypes.cast(ctypes.pointer(ctypes.c_size_t(ctypes.sizeof(self.c_args))), ctypes.c_void_p), 3)
 

tinygrad/runtime/ops_metal.py

@@ -113,7 +113,7 @@ class MetalCompiler(Compiler):
 
 class MetalProgram(Program[MetalDevice]):
   def __init__(self, dev:MetalDevice, obj:TinyELF):
-    self.dev, self.name, self.lib = dev, obj.name, obj.lib
+    self.dev, self.name, self.lib, self.signature = dev, obj.name, obj.lib, obj.signature
     data = objc.dispatch_data_create(obj.lib, len(obj.lib), None, None)
     self.library = self.dev.sysdevice.newLibraryWithData_error(data, ctypes.byref(error_lib:=metal.NSError().retained())).retained()
     error_check(error_lib)
@@ -138,7 +138,8 @@ class MetalProgram(Program[MetalDevice]):
     encoder = command_buffer.computeCommandEncoder().retained()
     encoder.setComputePipelineState(self.pipeline_state)
     for i,a in enumerate(bufs): encoder.setBuffer_offset_atIndex(a.buf, a.offset, i)
-    for i,a in enumerate(vals, start=len(bufs)): encoder.setBytes_length_atIndex(bytes(ctypes.c_int(a)), 4, i)
+    for a,(_,i,dt,_) in zip(vals, self.signature[len(bufs):]):
+      encoder.setBytes_length_atIndex(bytes(getattr(ctypes, f"c_int{dt.bitsize}")(a)), dt.itemsize, i)
     encoder.dispatchThreadgroups_threadsPerThreadgroup(metal.MTLSize(*global_size), metal.MTLSize(*local_size))
     encoder.endEncoding()
     command_buffer.setLabel(to_ns_str(self.name)) # TODO: is this always needed?

tinygrad/runtime/ops_nv.py

@@ -240,8 +240,10 @@ class NVVideoQueue(NVCommandQueue):
 
 class NVArgsState(CLikeArgsState):
   def __init__(self, buf:HCQBuffer, prg:NVProgram, bufs:tuple[HCQBuffer, ...], vals:tuple[int, ...]=()):
-    if isinstance(prg.dev.iface, MOCKIface): prg.cbuf_0[80:82] = [len(bufs), len(vals)]
-    super().__init__(buf, prg, bufs, vals=vals, prefix=prg.cbuf_0 or None)
+    if (is_mock:=isinstance(prg.dev.iface, MOCKIface)): prg.cbuf_0[80:82] = [len(bufs), len(vals)]
+    super().__init__(buf, prg, bufs, vals=() if is_mock else vals, prefix=prg.cbuf_0 or None)
+    # mock expects all vars to be 64 bit
+    if is_mock and vals: self.bind_sints_to_buf(*vals, buf=self.buf, fmt='q', offset=len(prg.cbuf_0)*4 + len(bufs)*8)
 
 class NVProgram(HCQProgram['NVDevice']):
   def __init__(self, dev:NVDevice, obj:TinyELF):
@@ -314,7 +316,7 @@ class NVProgram(HCQProgram['NVDevice']):
     self.max_threads = ((65536 // round_up(max(1, self.regs_usage) * 32, 256)) // 4) * 4 * 32
 
     # NV's kernargs is constbuffer, then arguments to the kernel follows. Kernargs also appends QMD at the end of the kernel.
-    super().__init__(NVArgsState, self.dev, self.name, kernargs_alloc_size=round_up(self.constbufs[0][1], 1 << 8) + (8 << 8))
+    super().__init__(NVArgsState, self.dev, obj, kernargs_alloc_size=round_up(self.constbufs[0][1], 1 << 8) + (8 << 8))
     weakref.finalize(self, self._fini, self.dev, self.lib_gpu, buf_spec)
 
   def _parse_elf_info(self, sh, start_off=0):

tinygrad/runtime/ops_qcom.py

@@ -10,7 +10,7 @@ from tinygrad.renderer.cstyle import QCOMCLRenderer
 from tinygrad.renderer.nir import IR3Renderer
 from tinygrad.helpers import getenv, mv_address, to_mv, round_up, data64_le, ceildiv, prod, cpu_profile, lo32, suppress_finalizing, is_image_shape
 from tinygrad.helpers import next_power2, flatten, PROFILE, IMAGE
-from tinygrad.dtype import dtypes
+from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.runtime.support.system import System
 if getenv("IOCTL"): import extra.qcom_gpu_driver.opencl_ioctl  # noqa: F401  # pylint: disable=unused-import
 
@@ -20,7 +20,7 @@ BUFTYPE_BUF, BUFTYPE_TEX, BUFTYPE_IBO = 0, 1, 2
 def dcache_flush():
   from tinygrad.uop.ops import UOp, Ops, KernelInfo
   from tinygrad.codegen import to_program
-  buf, n = UOp.param(0, dtypes.uint8, shape=(1,)), UOp.param(1, dtypes.int, shape=(1,), name="n", addrspace=None)
+  buf, n = UOp.param(0, dtypes.uint8, shape=(1,)), UOp.param(1, dtypes.int, shape=(), name="n", addrspace=AddrSpace.ALU)
   i = UOp.range(n, 0, dtype=dtypes.int)
   flush = UOp(Ops.CUSTOM, src=(buf.index(i * 64),), arg='__asm__ volatile("dc cvac, %0" :: "r"({0}) : "memory");')
   sink = UOp.sink(flush.end(i), UOp(Ops.CUSTOM, arg='__asm__ volatile("dsb sy" ::: "memory");'), arg=KernelInfo(name="dcache_flush"))
@@ -154,13 +154,15 @@ class QCOMComputeQueue(HWQueue):
                                                              state_block=mesa.SB6_CS_SHADER, num_unit=1024 // 4),
              *data64_le(args_state.buf.va_addr))
     self.cmd(mesa.CP_LOAD_STATE6_FRAG, qreg.cp_load_state6_0(state_type=mesa.ST_SHADER, state_src=mesa.SS6_INDIRECT,
-                                                             state_block=mesa.SB6_CS_SHADER, num_unit=round_up(prg.image_size, 128) // 128),
+                                                             state_block=mesa.SB6_CS_SHADER, num_unit=ceildiv(prg.image_size, 128)),
              *data64_le(prg.lib_gpu.va_addr))
 
     self.reg(mesa.REG_A6XX_SP_REG_PROG_ID_0, 0xfcfcfcfc, 0xfcfcfcfc, 0xfcfcfcfc, 0xfc, qreg.a6xx_sp_cs_const_config(constlen=1024 // 4, enabled=True))
 
     self.reg(mesa.REG_A6XX_SP_CS_PVT_MEM_STACK_OFFSET, qreg.a6xx_sp_cs_pvt_mem_stack_offset(prg.hw_stack_offset))
-    self.reg(mesa.REG_A6XX_SP_CS_INSTR_SIZE, qreg.a6xx_sp_cs_instr_size(prg.image_size // 4))
+    # image_size is in bytes, but INSTR_SIZE is measured in units of instruction groups (16 instructions, 8 bytes each)
+    # https://elixir.bootlin.com/mesa/mesa-26.1.5/source/src/freedreno/ir3/ir3_shader.h#L719-L723
+    self.reg(mesa.REG_A6XX_SP_CS_INSTR_SIZE, qreg.a6xx_sp_cs_instr_size(ceildiv(prg.image_size, 128)))
 
     if prg.samp_cnt > 0:
       self.cmd(mesa.CP_LOAD_STATE6_FRAG, qreg.cp_load_state6_0(state_type=mesa.ST_SHADER, state_src=mesa.SS6_INDIRECT,
@@ -210,10 +212,12 @@ class QCOMArgsState(HCQArgsState):
     if prg.samp_cnt > 0: to_mv(int(self.buf.va_addr) + prg.samp_off, len(prg.samplers) * 4).cast('I')[:] = array.array('I', prg.samplers)
     if prg.NIR:
       self.bind_sints_to_buf(*[b.va_addr for b in ubos], buf=self.buf, fmt='Q', offset=prg.buf_off)
-      self.bind_sints_to_buf(*vals, buf=self.buf, fmt='I', offset=prg.buf_off + len(ubos) * 8)
+      for v,(o,dt) in zip(vals, TinyELF.iter_sig(prg.signature[len(bufs):], len(ubos)*8)):
+        self.bind_sints_to_buf(v, buf=self.buf, fmt=dt.fmt, offset=prg.buf_off + o)
     else:
       for i, b in enumerate(ubos): self.bind_sints_to_buf(b.va_addr, buf=self.buf, fmt='Q', offset=prg.buf_offs[i])
-      for i, v in enumerate(vals): self.bind_sints_to_buf(v, buf=self.buf, fmt='I', offset=prg.buf_offs[i+len(ubos)])
+      for i,(v,(_,_,dt,_)) in enumerate(zip(vals, prg.signature[len(bufs):])):
+        self.bind_sints_to_buf(v, buf=self.buf, fmt=dt.fmt, offset=prg.buf_offs[i+len(ubos)])
 
     def _tex(b, ibo=False):
       imgdt, shape, buf = b
@@ -265,7 +269,7 @@ class QCOMProgram(HCQProgram['QCOMDevice']):
     dev._ensure_stack_size(self.hw_stack_offset * 4)
 
     kernargs_alloc_size = round_up(2048 + (self.tex_cnt + self.ibo_cnt) * 0x40 + len(self.samplers) * 4, 0x100)
-    super().__init__(QCOMArgsState, self.dev, self.name, kernargs_alloc_size=kernargs_alloc_size)
+    super().__init__(QCOMArgsState, self.dev, obj, kernargs_alloc_size=kernargs_alloc_size)
     weakref.finalize(self, self._fini, self.dev, self.lib_gpu, buf_spec)
 
   def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),

tinygrad/runtime/support/hcq.py

@@ -6,7 +6,7 @@ try: import fcntl # windows misses that
 except ImportError: fcntl = None #type:ignore[assignment]
 from tinygrad.helpers import DEV, PROFILE, getenv, to_mv, from_mv, cpu_profile, ProfileRangeEvent, select_first_inited, select_by_name, unwrap
 from tinygrad.helpers import suppress_finalizing, pluralize, TracingKey
-from tinygrad.device import Device, BufferSpec, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent, Program
+from tinygrad.device import Device, BufferSpec, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent, Program, TinyELF
 from tinygrad.uop.ops import sym_infer, sint, UOp
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.support.memory import BumpAllocator
@@ -326,14 +326,15 @@ class CLikeArgsState(HCQArgsState[ProgramType]):
     if prefix is not None: self.buf.cpu_view().view(size=len(prefix) * 4, fmt='I')[:] = array.array('I', prefix)
 
     self.bind_sints_to_buf(*[b.va_addr for b in bufs], buf=self.buf, fmt='Q', offset=len(prefix or []) * 4)
-    assert None not in vals
-    self.bind_sints_to_buf(*cast(tuple[sint, ...], vals), buf=self.buf, fmt='I', offset=len(prefix or []) * 4 + len(bufs) * 8)
+    for v,(val_offset,dt) in zip(vals, TinyELF.iter_sig(prg.signature[-len(vals):], len(bufs) * 8)):
+      assert v is not None
+      self.bind_sints_to_buf(v, buf=self.buf, fmt=dt.fmt, offset=len(prefix or []) * 4 + val_offset)
 
 class HCQProgram(Program[HCQDeviceType]):
-  def __init__(self, args_state_t:Type[HCQArgsState], dev:HCQDeviceType, name:str, kernargs_alloc_size:int, lib:bytes|None=None, base:int|None=None):
-    self.args_state_t, self.dev, self.name, self.kernargs_alloc_size = args_state_t, dev, name, kernargs_alloc_size
+  def __init__(self, args_state_t:Type[HCQArgsState], dev:HCQDeviceType, obj:TinyELF, kernargs_alloc_size:int, base:int|None=None):
+    self.args_state_t, self.dev, self.name, self.signature, self.kernargs_alloc_size = args_state_t, dev, obj.name, obj.signature, kernargs_alloc_size
     self.prof_prg_counter = next(self.dev.prof_prg_counter)
-    if PROFILE: Compiled.profile_events += [ProfileProgramEvent(dev.device, name, lib, base, self.prof_prg_counter)]
+    if PROFILE: Compiled.profile_events += [ProfileProgramEvent(dev.device, obj.name, obj.lib, base, self.prof_prg_counter)]
 
   @staticmethod
   def _fini(dev, buf, spec): dev.allocator.free(buf, buf.size, spec)

tinygrad/runtime/support/hcq2.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import cast, Callable, TypeVar, Generic, Any
+from typing import cast, Callable, TypeVar, Generic, Any, Sequence
 import struct, functools, time, collections, itertools
 from dataclasses import replace, dataclass
 from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, dedup, pluralize, JIT_BATCH_SIZE, unwrap
@@ -41,8 +41,9 @@ def unwrap_mstack(u):
   if u.op is Ops.MSTACK: return tuple(x for s in u.src for x in unwrap_mstack(s))
   return unwrap_mstack(u.src[0]) if u.op in {Ops.MSELECT, Ops.SLICE} else (u,)
 
-def make_patch(buf:UOp, off:sint, val:UOp) -> UOp:
-  return buf.index(UOp.const(dtypes.int, off // buf.dtype.itemsize)).store(val.simplify().cast(buf.dtype))
+def make_patches(buf:UOp, patches:Sequence[tuple[sint, UOp]]) -> UOp:
+  return buf.index(UOp.stack(*(UOp.const(dtypes.int, off // buf.dtype.itemsize) for off,_ in patches))) \
+            .store(UOp.stack(*(val.simplify().cast(buf.dtype) for _,val in patches)))
 
 def make_binary_patch(buf:UOp, blob:bytes) -> UOp:
   data = UOp(Ops.BINARY, src=(), arg=blob).bitcast(buf.dtype)
@@ -56,7 +57,7 @@ def make_cmdbuf(lin, devs, buf:UOp|None=None, dep:UOp|None=None):
     blob += struct.pack(f'<{ssimp.dtype.fmt}', ssimp.arg if ssimp.op is Ops.CONST else 0x0)
   cmdbuf = buf if buf is not None else UOp.placeholder((len(blob) // 4,), dtypes.uint32, next(UOp.unique_num), device=devs).rtag("cmdbuf")
   writable = cmdbuf.after(dep) if dep is not None else cmdbuf
-  return cmdbuf.after(make_binary_patch(writable, blob), *[make_patch(writable, off, s) for off, s in patches])
+  return cmdbuf.after(make_binary_patch(writable, blob), *((make_patches(writable, patches),) if patches else ()))
 
 def make_signal(devs, queue="COMPUTE:0", sentinel=False):
   return UOp.placeholder((1,), dtypes.uint64, 0, device=devs, volatile=True).rtag("sentinel_signal" if sentinel else f"{queue}_timeline_signal")
@@ -70,7 +71,7 @@ def encode_kernargs_clike(call:UOp, prg:UOp, devs:str|tuple[str, ...]) -> UOp:
   data, info = prg.arg
   buf = UOp.placeholder((data.kernargs_alloc_size // 4,), dtypes.uint32, next(UOp.unique_num), device=devs).rtag("kernargs")
   words = [w for gi in info.globals for w in data64_le(get_call_arg_uops(call)[gi].getaddr(devs))] + list(info.vars)
-  return buf.after(*[make_patch(buf, i * 4, w) for i, w in enumerate(words)])
+  return buf.after(*((make_patches(buf, [(i * 4, w) for i, w in enumerate(words)]),) if words else ()))
 
 # *****************
 # 0.1. prep: replace buffers with params
@@ -280,7 +281,7 @@ def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[dict[UOp, UOp
   table = UOp.placeholder((len(order),), dtypes.uint64, next(UOp.unique_num), device=call.arg.aux.device).rtag(name)
 
   reads = {g: table.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(UOp.const(dtypes.int, slots[bare[g]])).load() for g in gaddrs}
-  return reads, (table.after(*[make_patch(table, i * table.dtype.itemsize, addr) for addr, i in slots.items()]),) if slots else ()
+  return reads, (table.after(make_patches(table, [(i * table.dtype.itemsize, addr) for addr, i in slots.items()])),) if slots else ()
 
 def make_blob_bufs(call:UOp, blobs:list[UOp]) -> tuple[dict[UOp, UOp], tuple[UOp, ...]]:
   bufs = {b: UOp.placeholder((b.max_numel(),), b.dtype, next(UOp.unique_num), device=call.arg.aux.device).rtag("template") for b in blobs}
@@ -411,9 +412,10 @@ def fold_binary(buf:UOp, blob:UOp) -> UOp:
   return UOp(Ops.NOOP)
 
 def fold_const_store(buf:UOp, off:UOp, val:UOp) -> UOp:
-  for b, v in zip((bs:=mb.bufs if isinstance((mb:=buf.buffer), MultiBuffer) else (mb,)), val.src if val.op is Ops.STACK else (val,)*len(bs)):
-    data = struct.pack(f'<{v.dtype.fmt}', truncate[v.dtype](v.arg))
-    b.ensure_allocated().as_memoryview(force_zero_copy=True, no_sync=True).cast('B')[(byte_off:=off.arg*buf.dtype.itemsize):byte_off+len(data)] = data
+  for off,val in zip(off.src, val.src):
+    for b,v in zip((bs:=mb.bufs if isinstance((mb:=buf.buffer), MultiBuffer) else (mb,)), val.src if val.op is Ops.STACK else (val,)*len(bs)):
+      data = struct.pack(f'<{v.dtype.fmt}', truncate[v.dtype](v.arg))
+      b.ensure_allocated().as_memoryview(force_zero_copy=True, no_sync=True).cast('B')[(bo:=off.arg*buf.dtype.itemsize):bo+len(data)] = data
   return UOp(Ops.NOOP)
 
 def resolve_getaddr(buf:UOp, g:UOp) -> UOp:
@@ -437,8 +439,7 @@ pm_resolve_patches = PatternMatcher([
   (UPat(name="buf").index(UPat(Ops.RANGE), allow_any_len=True)
     .store(UPat.any(UPat(Ops.BINARY, name="blob"), UPat(Ops.BINARY, name="blob").bitcast()).index(UPat(Ops.RANGE), allow_any_len=True).load())
     .end(UPat(Ops.RANGE)), fold_binary),
-  (UPat({Ops.BUFFER, Ops.SLICE, Ops.MSTACK}, name="buf").index(UPat.cvar("off"))
-    .store(UPat.any(UPat.cvar("val"), UPat(Ops.STACK, name="val"))), fold_const_store),
+  (UPat({Ops.BUFFER, Ops.SLICE, Ops.MSTACK}, name="buf").index(UPat(Ops.STACK, name="off")).store(UPat(Ops.STACK, name="val")), fold_const_store),
 ])
 
 pm_assert_no_afters = PatternMatcher([(UPat(Ops.AFTER, name="a"), lambda a: panic(RuntimeError, f"AFTER left at hcq_link: {a.src[0].op}"))])

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/schedule/indexing.py

@@ -50,7 +50,7 @@ class IndexingContext:
 
   # create ranges
   range_idx: Iterator[int] = field(default_factory=itertools.count)
-  def new_range(self, s:sint, axistype:AxisType=AxisType.LOOP) -> UOp:
+  def new_range(self, s:sint, axistype:AxisType=AxisType.WEAK) -> UOp:
     if isinstance(s, UOp) and s.op is Ops.RANGE: return s
     # if a range has a 1 src, it's the same as UOp.const(dtypes.weakint, 0)
     return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(None, 0)

tinygrad/schedule/multi.py

@@ -1,5 +1,6 @@
 from tinygrad.helpers import all_same, prod, getenv, ALLREDUCE_CAST
-from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp, AxisType, graph_rewrite, broadcast_axes, _broadcast_shape
+from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp, AxisType, graph_rewrite, broadcast_axes, _broadcast_shape, sint_to_uop
+from tinygrad.uop.ops import sint, ssimplify
 from tinygrad.dtype import dtypes
 from tinygrad.schedule.allreduce import handle_allreduce
 
@@ -50,7 +51,12 @@ def shard_srcs(msrcs:tuple[UOp, ...], axis:int) -> list[UOp]:
   # normalize srcs to local shards on axis
   devices = [x.device for x in msrcs if x.device is not None]
   assert all_same(devices), f"all buffers must have the same device {devices}"
-  dcount = len(devices[0])
+  # without devices the sharding range comes from the UNSHARD itself (e.g. a LOCAL thread range);
+  # device shards range over the devices instead
+  if len(devices): sharding_rng = UOp.range(len(devices[0]), -1, AxisType.DEVICE)
+  else:
+    sharding_rng = next((m.src[1] for m in msrcs if m.op is Ops.UNSHARD), None)
+    assert sharding_rng is not None, "shard_srcs requires a device or a sharding range"
 
   out_shape = _broadcast_shape(*[x.shape for x in msrcs])
   srcs:list[UOp] = []
@@ -60,12 +66,19 @@ def shard_srcs(msrcs:tuple[UOp, ...], axis:int) -> list[UOp]:
       # same axis, just copy through
       srcs.append(mlb.src[0])
     else:
-      # otherwise every device gets the full copy, sharded iff this src has the axis (broadcast srcs stay whole)
+      # otherwise every shard gets the full copy, sharded iff this src has the axis (broadcast srcs stay whole)
       full = mlb if mlb.axis is None else copy_multi(mlb, mlb.device)
-      srcs.append(full if axis in broadcast_axes(mlb.shape, out_shape) else full._shard(src_axis, dcount))
+      srcs.append(full if axis in broadcast_axes(mlb.shape, out_shape) else full._shard(src_axis, sharding_rng))
   return srcs
 
 def alu_multi(root:UOp):
+  multis = [m for m in root.src if m.op is Ops.UNSHARD]
+  if not multis: return None
+  sharding = multis[0].sharding
+  if len(multis) == len(root.src) and all(m.sharding == sharding for m in multis):
+    srcs = [m.src[0] for m in root.src]
+    return srcs[0].alu(root.op, *srcs[1:]).unshard(multis[0].arg, multis[0].src[1:])
+  # resharding: single-axis fallback via shard_srcs
   axis = root.axis
   assert axis is not None
   srcs = shard_srcs(root.src, axis)
@@ -73,69 +86,156 @@ def alu_multi(root:UOp):
 
 def reduce_multi(root:UOp, multi:UOp):
   op, num_axes = root.arg
-  if multi.axis is not None and multi.axis < num_axes:
-    local = multi.src[0]._rop(op, tuple(range(num_axes)))
-    # allreduce in pre-cast dtype when sum_acc_dtype promoted from bf16/half
+  sharding = multi.sharding
+  reduced = [(ax, rng) for ax, rng in sharding if ax < num_axes]
+  remaining = [(ax, rng) for ax, rng in sharding if ax >= num_axes]
+  local = multi.src[0]._rop(op, tuple(range(num_axes)))
+  if reduced:
+    assert not remaining, f"partial allreduce not supported for multi-axis sharding {sharding}"
+    # all sharded axes are reduced: full allreduce
     if ALLREDUCE_CAST and multi.src[0].op is Ops.CAST and multi.src[0].src[0].dtype in (dtypes.bfloat16, dtypes.half):
       orig_dtype = multi.src[0].src[0].dtype
       return local.cast(orig_dtype).allreduce(op, multi.device).cast(local.dtype)
     return local.allreduce(op, multi.device)
-  # reduce on non sharded axes, piecewise is fine. if axis is None this is also correct
-  new_axis = multi.axis - num_axes if multi.axis is not None else None
-  return multi.src[0]._rop(op, tuple(range(num_axes))).unshard(new_axis, multi.src[1])
+  # no sharded axes reduced: piecewise, keep all remaining sharding
+  new_axes = tuple(ax - num_axes for ax, _ in remaining)
+  new_rngs = tuple(rng for _, rng in remaining)
+  return local.unshard(new_axes, new_rngs)
 
 def reshape_multi(root:UOp, multi:UOp):
   if prod(multi.shape) != prod(new_shape:=root.marg): raise RuntimeError("reshape must maintain prod(shape)")
-  if (new_axis:=root.axis) is not None: new_shape = tuple(s//len(multi.device) if a==new_axis else s for a,s in enumerate(new_shape))
-  return multi.src[0].reshape(new_shape).unshard(new_axis, multi.src[1])
+  # map every sharded axis through the reshape: the axis boundary must survive intact and stay divisible by its shard count
+  arg_acc:list[sint] = [1]
+  for s in new_shape: arg_acc.append(ssimplify(arg_acc[-1]*s))
+  new_shardings = []
+  for ax, rng in multi.sharding:
+    count = int(rng.vmax)+1
+    target = prod(multi.shape[:ax])
+    if target not in arg_acc: raise RuntimeError(f"reshape {multi.shape} -> {new_shape} moved items between shards")
+    new_ax = len(arg_acc) - arg_acc[::-1].index(target) - 1
+    if new_shape[new_ax] % count != 0: raise RuntimeError(f"reshape {multi.shape} -> {new_shape} moved items between shards")
+    new_shardings.append((new_ax, rng))
+  new_axs = {a for a, _ in new_shardings}
+  new_shape = tuple(s//(int(rng.vmax)+1) if a in new_axs else s for a,s in enumerate(new_shape))
+  return multi.src[0].reshape(new_shape).unshard(tuple(a for a,_ in new_shardings), tuple(r for _,r in new_shardings))
 
 def expand_multi(root:UOp, multi:UOp):
-  new_axis = None if multi.axis is None else multi.axis + len(root.marg)
-  return multi.src[0]._mop(Ops.EXPAND, arg=root.marg).unshard(new_axis, multi.src[1])
+  shift = len(root.marg)
+  return multi.src[0]._mop(Ops.EXPAND, arg=root.marg) \
+    .unshard(tuple(ax+shift for ax,_ in multi.sharding), tuple(r for _,r in multi.sharding))
 
 def pad_multi(root:UOp, multi:UOp):
-  assert multi.axis is None or root.marg[multi.axis] == (0, multi.shape[multi.axis]), f"padding not supported for {root.marg=}"
-  local_pad = tuple((0, multi.src[0].shape[multi.axis]) if a == multi.axis else s for a,s in enumerate(root.marg))
-  return multi.src[0]._mop(Ops.PAD, local_pad).unshard(multi.axis, multi.src[1])
+  for ax, _ in multi.sharding:
+    assert root.marg[ax] == (0, multi.shape[ax]), f"padding not supported for {root.marg=}"
+  counts = {a for a,_ in multi.sharding}
+  local_pad = tuple((0, multi.src[0].shape[a]) if a in counts else s for a,s in enumerate(root.marg))
+  return multi.src[0]._mop(Ops.PAD, local_pad).unshard(multi.arg, multi.src[1:])
 
 def permute_multi(root:UOp, multi:UOp):
   # all permutes supported!
-  return multi.src[0].permute(root.marg).unshard(root.axis, multi.src[1])
+  return multi.src[0].permute(root.marg) \
+    .unshard(tuple(root.marg.index(ax) for ax,_ in multi.sharding), tuple(r for _,r in multi.sharding))
 
 def shrink_multi(root:UOp, multi:UOp):
-  shard_bounds = tuple((s,e-s) for s,e in multi.bounds) if multi.axis is not None else ()
-  assert multi.axis is None or root.marg[multi.axis] == (0, multi.shape[multi.axis]) or root.marg[multi.axis] in shard_bounds, \
-    f"shrinking not supported for {root.marg=}"
-  if multi.axis is not None and root.marg[multi.axis] in shard_bounds and root.marg[multi.axis] != (0, multi.shape[multi.axis]):
-    # NOTE: shrink on the shard axis is only allowed when result is a single partition, denoted by the new real
-    # we just copy it to all the devices, no real. this will be optimized out later
-    non_shard_shrink = tuple((0, multi.src[0].shape[i]) if i == multi.axis else s for i, s in enumerate(root.marg))
-    return multi.src[0].copy_to_device(multi.device, arg=shard_bounds.index(root.marg[multi.axis]))._mop(Ops.SHRINK, non_shard_shrink)
-  local_shrink = tuple((0, multi.src[0].shape[multi.axis]) if a == multi.axis else s for a,s in enumerate(root.marg))
-  return multi.src[0]._mop(Ops.SHRINK, local_shrink).unshard(multi.axis, multi.src[1])
+  # resolve each sharded axis independently: a shrink to exactly this range's own shard resolves the UNSHARD along
+  # that axis (e.g. a fragment indexed by its LOCAL thread range becomes that thread's REG shard, no copy needed)
+  local_marg = list(root.marg)
+  remaining = list(multi.sharding)
+  for ax, rng in multi.sharding:
+    shard_sz = multi.src[0].shape[ax]
+    s, l = root.marg[ax]  # SHRINK marg is (start, length)
+    if sint_to_uop(l).ssimplify() == shard_sz and (sint_to_uop(s)-rng*shard_sz).ssimplify() == 0:
+      local_marg[ax] = (0, shard_sz)
+      remaining.remove((ax, rng))
+      continue
+    part_bounds = tuple((i*shard_sz, shard_sz) for i in range(int(rng.vmax)+1))
+    if (s, l) == (0, multi.shape[ax]): local_marg[ax] = (0, shard_sz)  # full axis stays sharded, shrink the other axes locally
+    else:
+      # NOTE: otherwise a shrink on the shard axis is only allowed on the legacy device path, selecting a single
+      # partition (which is copied to all the devices and optimized out later)
+      if len(multi.sharding) != 1 or not isinstance(multi.device, tuple) or (s, l) not in part_bounds:
+        raise RuntimeError(f"shrinking not supported for {root.marg=}")
+      non_shard_shrink = tuple((0, shard_sz) if i == ax else t for i, t in enumerate(root.marg))
+      return multi.src[0].copy_to_device(multi.device, arg=part_bounds.index((s, l)))._mop(Ops.SHRINK, non_shard_shrink)
+  val = multi.src[0]._mop(Ops.SHRINK, tuple(local_marg))
+  return val if not remaining else val.unshard(tuple(a for a,_ in remaining), tuple(r for _,r in remaining))
 
 def flip_multi(root:UOp, multi:UOp):
-  assert multi.axis is None or not root.marg[multi.axis], "flipping not supported on sharded axis"
-  return multi.src[0].flip([i for i,x in enumerate(root.marg) if x]).unshard(multi.axis, multi.src[1])
+  for ax, _ in multi.sharding:
+    if root.marg[ax]: raise RuntimeError(f"flipping not supported on sharded axis {ax}")
+  return multi.src[0].flip([i for i,x in enumerate(root.marg) if x]).unshard(multi.arg, multi.src[1:])
 
 def stack_multi(root:UOp):
   # STACK adds a leading axis: srcs are sharded one axis below the output
+  multis = [m for m in root.src if m.op is Ops.UNSHARD]
+  if not multis: return None
+  sharding = multis[0].sharding
+  if all(m.sharding == sharding for m in multis):
+    srcs = [m.src[0] if m.op is Ops.UNSHARD else m for m in root.src]
+    new_sharding = tuple((ax+1, rng) for ax, rng in sharding)
+    return UOp(Ops.STACK, src=tuple(srcs)).unshard(tuple(a for a,_ in new_sharding), tuple(r for _,r in new_sharding))
+  # resharding: single-axis fallback
   axis = root.axis
   assert axis is not None
   return UOp(Ops.STACK, src=tuple(shard_srcs(root.src, axis-1))).unshard(axis, next(m.src[1] for m in root.src if m.op is Ops.UNSHARD))
 
-def copy_multi(multi:UOp, device:str | tuple[str, ...]):
-  assert multi.axis is not None, "all multi ops have axis"
-  if isinstance(device, str):
-    pieces = [multi.src[0].mselect(i).copy_to_device(device) for i in range(len(multi.device))]
-    return pieces[0].cat(*pieces[1:], dim=multi.axis)
-  return multi.src[0]._unshard(multi.axis).allreduce(Ops.ADD, device)
+def index_multi(root:UOp, multi:UOp):
+  # INDEX on UNSHARD: resolve each sharded axis into this range's own shard.
+  # Two ownership patterns are supported:
+  #   contiguous: idx = rng*shard_sz + local   (thread rng owns [rng*shard_sz, ...))
+  #   strided:    idx = rng + ir*shard_sz      (thread rng owns {rng, rng+shard_sz, ...})
+  idxs = list(root.src[1:])
+  for ax, rng in multi.sharding:
+    shard_sz = multi.src[0].shape[ax]
+    local = (idxs[ax] - rng*shard_sz).simplify()
+    if local.vmin >= 0 and local.vmax < shard_sz:
+      idxs[ax] = local
+      continue
+    # strided ownership: idx ≡ rng (mod shard_sz), intra-shard position is (idx - rng) // shard_sz
+    diff = (idxs[ax] - rng).simplify()
+    if (mod:=(diff % shard_sz).simplify()).op is Ops.CONST and mod.arg == 0:
+      local = (diff // shard_sz).simplify()
+      if local.vmin >= 0 and local.vmax < shard_sz:
+        idxs[ax] = local
+        continue
+    raise RuntimeError(f"index_multi: cannot shard index {idxs[ax]} for UNSHARD axis {ax} with shard size {shard_sz}")
+  return multi.src[0].index(*idxs)
 
-def store_after_multi(dest:UOp, src:UOp): return dest.after(dest.store(src.src[0])).unshard(src.axis, src.src[1])
+def _shard_idx(rng:UOp, dev_idx:int) -> int:
+  drngs = [r for r in rng.ranges if r.arg[-1] is AxisType.DEVICE]
+  return 0 if not drngs else int(rng.substitute({drngs[0]: drngs[0].const_like(dev_idx)}).ssimplify())
+
+def copy_multi(multi:UOp, device:str | tuple[str, ...]):
+  sharding = multi.sharding
+  if isinstance(device, str):
+    # reconstruct by concatenating along each axis from last to first
+    piece_info: list[tuple[tuple, UOp]] = []
+    for i in range(len(multi.device)):
+      idxs = tuple(_shard_idx(r, i) for _, r in sharding)
+      piece_info.append((idxs, multi.src[0].mselect(i).copy_to_device(device)))
+    for j in range(len(sharding) - 1, -1, -1):
+      ax, rng = sharding[j]
+      groups: dict[tuple, list[tuple[int, UOp]]] = {}
+      for idxs, p in piece_info:
+        key = idxs[:j] + idxs[j+1:]
+        groups.setdefault(key, []).append((idxs[j], p))
+      piece_info = []
+      for key in sorted(groups):
+        grp = sorted(groups[key], key=lambda x: x[0])
+        piece_info.append((key, grp[0][1].cat(*[x[1] for x in grp[1:]], dim=ax)))
+    return piece_info[0][1]
+  # multi-device target: unshard all axes and allreduce
+  val = multi.src[0]
+  for ax, rng in sharding:
+    bsz = val.shape[ax]
+    val = val.pad(tuple((0,0) if a != ax else (bsz*rng, bsz*int(rng.vmax) - bsz*rng) for a in range(len(val.shape))))
+  return val.allreduce(Ops.ADD, device)
+
+def store_after_multi(dest:UOp, src:UOp): return dest.after(dest.store(src.src[0])).unshard(src.arg, src.src[1:])
 
 def passthrough_multi(root:UOp, multi:UOp):
   new_src = (multi.src[0],)+tuple(x.src[0] if x.op is Ops.UNSHARD else x for x in root.src[1:])
-  return UOp(root.op, root.dtype, src=new_src, arg=root.arg).unshard(multi.axis, multi.src[1])
+  return UOp(root.op, root.dtype, src=new_src, arg=root.arg).unshard(multi.arg, multi.src[1:])
 
 def rewrite_into_function(call:UOp):
   if call.arg.precompile: return None
@@ -145,7 +245,7 @@ def rewrite_into_function(call:UOp):
   assert new_body.op is Ops.TUPLE
   if any(s.op is Ops.UNSHARD for s in new_body.src):
     shard_call = call.replace(src=(UOp.maketuple(*[s.src[0] if s.op is Ops.UNSHARD else s for s in new_body.src]),)+new_args)
-    return UOp.maketuple(*[shard_call.gettuple(i).unshard(s.axis, s.src[1]) if s.op is Ops.UNSHARD else shard_call.gettuple(i)
+    return UOp.maketuple(*[shard_call.gettuple(i).unshard(s.arg, s.src[1:]) if s.op is Ops.UNSHARD else shard_call.gettuple(i)
                            for i, s in enumerate(new_body.src)])
   return call.replace(src=(new_body,)+new_args)
 
@@ -165,17 +265,17 @@ multi_pm = PatternMatcher([
   (UPat(Ops.PERMUTE, src=(UPat(Ops.UNSHARD, name="multi"), ), name="root"), permute_multi),
   (UPat(Ops.FLIP, src=(UPat(Ops.UNSHARD, name="multi"), ), name="root"), flip_multi),
   (UPat(Ops.STACK, name="root", custom_early_reject=set([Ops.UNSHARD])), stack_multi),
+  (UPat(Ops.INDEX, src=(UPat(Ops.UNSHARD, name="multi"),), name="root", allow_any_len=True), index_multi),
   (UPat(Ops.AFTER, src=(UPat(Ops.UNSHARD), UPat(Ops.STORE, src=(UPat(Ops.UNSHARD, name="dest"), UPat(Ops.UNSHARD, name="src"))))), store_after_multi),
   (UPat(Ops.COPY, src=(UPat(Ops.UNSHARD, name="multi"),), name="copy"), lambda multi,copy: copy_multi(multi, copy.arg)),
   (UPat(Ops.ALLREDUCE, src=(UPat(Ops.UNSHARD, name="multi"),), name="red"),
-    lambda multi,red: multi.src[0].allreduce(*red.arg).unshard(multi.axis, multi.src[1])),
+    lambda multi,red: multi.src[0].allreduce(*red.arg).unshard(multi.arg, multi.src[1:])),
 
   # resolve TUPLE+GETTUPLE (needed in multi)
   (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]),
   # GETTUPLE on UNSHARD: passthrough UNSHARD (e.g. when FUNCTION was replaced by UNSHARD(GETTUPLE(...)))
   (UPat(Ops.GETTUPLE, src=(UPat(Ops.UNSHARD, name="multi"),), name="g"),
-    lambda g, multi: multi.src[0].gettuple(g.arg).unshard(multi.axis, multi.src[1]) if multi.src[0].op in {Ops.FUNCTION, Ops.TUPLE}
-    else multi),
+    lambda g, multi: multi.src[0].gettuple(g.arg).unshard(multi.arg, multi.src[1:]) if multi.src[0].op in {Ops.FUNCTION, Ops.TUPLE} else multi),
   # rewrite into FUNCTION calls explicitly for UNSHARD (value-producing)
   (UPat(Ops.FUNCTION, name="call"), rewrite_into_function),
   (UPat((Ops.CALL, Ops.FUNCTION, Ops.AFTER), src=(UPat(Ops.UNSHARD, name="multi"), ), name="root", allow_any_len=True), passthrough_multi),

tinygrad/schedule/rangeify.py

@@ -342,9 +342,9 @@ def limit_bufs(ctx:IndexingContext, root:UOp):
     srcs = []
     for s in root.src:
       if s.op in GroupOp.Elementwise and s.device is not None:
-        # Insert bufferize: all AxisType.REDUCE before bufferize are AxisType.LOOP, the DEVICE range stays a launched axis
+        # Insert bufferize: all AxisType.REDUCE before bufferize are AxisType.WEAK, the DEVICE range stays a launched axis
         orig_ranges = s.ranges
-        end_ranges = [x.replace(arg=(next(ctx.range_idx), AxisType.LOOP)) if x.op is Ops.RANGE and x.arg[-1] is not AxisType.DEVICE else x
+        end_ranges = [x.replace(arg=(next(ctx.range_idx), AxisType.WEAK)) if x.op is Ops.RANGE and x.arg[-1] is not AxisType.DEVICE else x
                       for x in s.ranges]
         s = s.substitute(dict(zip(orig_ranges, end_ranges))).bufferize(*end_ranges, arg=BufferizeOpts(device=s.device)).index(*orig_ranges)
       srcs.append(s)

tinygrad/tensor.py

@@ -240,7 +240,7 @@ class Tensor(RandMixin):
     if capturing and not getenv("UNSAFE_ALLOW_JIT_BUFFER"):
       from tinygrad.engine.jit import JitError
       raise JitError("cannot access tensor data during JIT capture, the value will be baked in")
-    x = self.cast(strong_dtype(self.dtype)).contiguous()
+    x = self.contiguous()
     if self.uop.device is None or isinstance(self.device, tuple): x = x.clone("CPU")
     return cast(Buffer, x.realize().uop.buffer).ensure_allocated()
 
@@ -279,7 +279,6 @@ class Tensor(RandMixin):
     print(t.tolist())
     ```
     """
-    if self.dtype in dtypes.weaks: return self.cast(strong_dtype(self.dtype)).tolist()
     # TODO: remove half once minimum python supports it
     if self.dtype in (dtypes.half, dtypes.bfloat16, *dtypes.fp8s): return self.cast(dtypes.float32).tolist()
     if 0 in self.shape:

tinygrad/uop/ops.py

@@ -16,8 +16,8 @@ if TYPE_CHECKING:
 
 class AxisType(Enum):
   def __repr__(self): return str(self)
-  DEVICE = auto(); GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto() # noqa: E702
-  UNROLL = auto(); THREAD = auto(); PLACEHOLDER = auto() # noqa: E702
+  DEVICE = auto(); GLOBAL = auto(); WARP = auto(); LOCAL = auto(); WEAK = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto() # noqa: E702
+  UNROLL = auto(); THREAD = auto(); PLACEHOLDER = auto(); LOOP = auto() # noqa: E702
 
 @dataclass(frozen=True, order=True)
 class ParamArg:
@@ -35,14 +35,15 @@ class ParamArg:
               ("volatile", False))
     args = [repr(self.slot), repr(self.dtype)] + [f"{k}={v!r}" for k,default in fields if (v:=getattr(self, k)) != default]
     return f"ParamArg({', '.join(args)})"
-axis_letters = {AxisType.DEVICE: "d", AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L",
-                AxisType.UPCAST: "u", AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
+axis_letters = {AxisType.DEVICE: "d", AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.WEAK: "L",
+                AxisType.LOOP: "L", AxisType.UPCAST: "u", AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
 axis_colors = {AxisType.DEVICE: "green", AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN",
-               AxisType.LOOP: "WHITE", AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"}
+               AxisType.WEAK: "WHITE", AxisType.LOOP: "WHITE", AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red",
+               AxisType.UNROLL: "magenta"}
 
 # NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
-axis_to_pos = {AxisType.DEVICE: -2, AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2,
-               AxisType.UPCAST: 3, AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
+axis_to_pos = {AxisType.DEVICE: -2, AxisType.WEAK: -1, AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1,
+               AxisType.LOCAL: 2, AxisType.UPCAST: 3, AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
 
 range_start = {Ops.STAGE: 1, Ops.REDUCE: 1, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.FUNCTION: 1,
                Ops.SLICE: 2, Ops.LINEAR: 0}
@@ -438,7 +439,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
         case Ops.FLIP:
           if len(ps) != len(self.marg) or not all(isinstance(x, bool) for x in self.marg): raise ValueError(f"bad flip on {ps}, {self.marg}")
           return ps
-        case Ops.UNSHARD: return tuple(s*(int(self.src[1].vmax)+1) if a == self.axis else s for a,s in enumerate(ps))
+        case Ops.UNSHARD: return tuple(s*(int(self.src[1:][self.arg.index(a)].vmax)+1) if a in self.arg else s for a,s in enumerate(ps))
         case Ops.REDUCE:
           num_axes = self.arg[1]
           if not isinstance(num_axes, int) or num_axes < 0 or num_axes > len(ps):
@@ -616,10 +617,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
       ret = UOp(Ops.CONST, dtype, arg=dtype.const(b), src=())
     return ret._mop(Ops.EXPAND, arg=shape) if shape is not None and shape != () and ret.shape != shape else ret
   @staticmethod
-  def range(end:sint, axis_id, axis_type=AxisType.LOOP, *arg, dtype=dtypes.weakint, src=(), **kwargs):
+  def range(end:sint, axis_id, axis_type=AxisType.WEAK, *arg, dtype=dtypes.weakint, src=(), **kwargs):
     return UOp(Ops.RANGE, src=(sint_to_uop(end, dtype),)+src, arg=(axis_id, axis_type)+arg, **kwargs)
   @staticmethod
-  def loop(axis_id:int, *arg): return UOp(Ops.RANGE, src=(UOp(Ops.NOOP),), arg=(axis_id, AxisType.LOOP)+arg)
+  def loop(axis_id:int, *arg): return UOp(Ops.RANGE, src=(UOp(Ops.NOOP),), arg=(axis_id, AxisType.WEAK)+arg)
   @staticmethod
   def special(end:sint, name:str, dtype=dtypes.weakint): return UOp(Ops.SPECIAL, src=(sint_to_uop(end, dtype),), arg=name)
   @staticmethod
@@ -664,24 +665,38 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   # *** multi-device helpers ***
 
-  def unshard(self, axis:int|None, device_range:UOp|None=None):
-    assert isinstance(self.device, tuple), f"multi device must be tuple, {self.device} isn't"
+  def unshard(self, axis:int|tuple[int, ...]|None, device_range:UOp|tuple[UOp, ...]|None=None):
     assert axis is not None, "multi None is no longer supported"
-    # an UNSHARD always has two srcs: the value and the DEVICE range it ends (defaults to a DEVICE range over the devices)
-    if device_range is None: device_range = UOp.range(len(self.device), -1, AxisType.DEVICE)
-    assert device_range.op is Ops.RANGE and device_range.arg[-1] is AxisType.DEVICE
-    return UOp(Ops.UNSHARD, src=(self, device_range), arg=axis)
+    # an UNSHARD carries the value and one sharding range per sharded axis (arg is the tuple of sharded axes,
+    # sorted). the single-axis axis form defaults the range to a DEVICE range over the devices; a range need not
+    # be DEVICE, e.g. a LOCAL range shards a kernel tile into per-thread fragments
+    if isinstance(axis, int): axis = (axis,)
+    if device_range is None:
+      assert isinstance(self.device, tuple), f"multi device must be tuple, {self.device} isn't"
+      device_range = (UOp.range(len(self.device), -1, AxisType.DEVICE),)
+    if isinstance(device_range, UOp): device_range = (device_range,)
+    assert isinstance(device_range, tuple) and len(axis) == len(device_range) and len(set(axis)) == len(axis)
+    axis, device_range = map(tuple, zip(*sorted(zip(axis, device_range))))
+    return UOp(Ops.UNSHARD, src=(self, *device_range), arg=axis)
+
+  @property
+  def sharding(self) -> tuple[tuple[int, UOp], ...]:
+    """(axis, RANGE) pairs this value is sharded over (the source of truth for shard bounds/counts)."""
+    return tuple(zip(self.arg, self.src[1:])) if self.op is Ops.UNSHARD else ()
 
   @property
   def bounds(self):
     if self.axis is None: raise RuntimeError("bounds is not defined when axis is None")
-    return tuple(itertools.pairwise(itertools.accumulate([self.src[0].shape[self.axis] for _ in self.device], initial=0)))
+    dcount = int(self.src[1].vmax)+1 if self.op is Ops.UNSHARD else len(self.device)
+    return tuple(itertools.pairwise(itertools.accumulate([self.src[0].shape[self.axis] for _ in range(dcount)], initial=0)))
 
   @functools.cached_property
   def axis(self) -> int|None:
     # COPY removes axis. TODO: add more tests for this, and consider MSELECT/MSTACK
     if self.op is Ops.COPY: return None
-    if self.op is Ops.UNSHARD: return self.arg
+    if self.op is Ops.UNSHARD:
+      if len(self.arg) != 1: raise RuntimeError(f"UOp is sharded on multiple axes {self.arg}, use .sharding")
+      return self.arg[0]
     # GETTUPLE: axis comes from the specific TUPLE element, not src[0]
     if self.op is Ops.GETTUPLE:
       in_tuple = self.src[0].src[0] if self.src[0].op is Ops.FUNCTION else self.src[0]
@@ -705,7 +720,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
       target = ssimplify(prod(self.src[0].shape[:src_axis]))
       if target not in arg_acc: raise RuntimeError(f"reshape {self.src[0].shape} -> {self.shape} moved items between shards")
       new_axis = len(arg_acc) - arg_acc[::-1].index(target) - 1
-      if self.shape[new_axis] % len(self.device) != 0: raise RuntimeError(f"reshape {self.src[0].shape} -> {self.shape} moved items between shards")
+      dcount = len(self.device) if isinstance(self.device, tuple) else \
+        int(next(u.src[1] for u in self.src[0].toposort() if u.op is Ops.UNSHARD).vmax)+1
+      if self.shape[new_axis] % dcount != 0: raise RuntimeError(f"reshape {self.src[0].shape} -> {self.shape} moved items between shards")
       return new_axis
     if self.op is Ops.PERMUTE: return self.marg.index(src_axis) if src_axis is not None else None
     if self.op is Ops.EXPAND: return src_axis + len(self.marg) if src_axis is not None else None
@@ -716,15 +733,15 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     dnum = UOp.range(dcount, -1, AxisType.DEVICE)
     return self.pad(tuple((0,0) if a != axis else (bsz*dnum, bsz*(dcount-1) - bsz*dnum) for a in range(len(self.shape))))
 
-  def _shard(self, axis:int, dcount:int) -> UOp:
+  def _shard(self, axis:int, rng:UOp) -> UOp:
     if len(self.shape) == 0: return self  # scalars broadcast, no sharding needed
-    dnum = UOp.range(dcount, -1, AxisType.DEVICE)
+    dcount = int(rng.vmax)+1
     if self.shape[axis] % dcount != 0: raise RuntimeError(f"multi axis uneven: {self.shape[axis]=} {axis=} {dcount=}")
     sz = self.shape[axis] // dcount
-    return self.shrink(tuple((0,s) if i != axis else (dnum*sz,dnum*sz+sz) for i,s in enumerate(self.shape)))
+    return self.shrink(tuple((0,s) if i != axis else (rng*sz,rng*sz+sz) for i,s in enumerate(self.shape)))
   def shard(self, devices:tuple[str, ...], axis:int|None=None) -> UOp:
     copied = self.copy_to_device(devices)
-    return copied if axis is None else copied._shard(axis, len(devices)).unshard(axis)
+    return copied if axis is None else copied._shard(axis, UOp.range(len(devices), -1, AxisType.DEVICE)).unshard(axis)
 
   def copy_to_device(self, device:str|tuple[str, ...], arg=None):
     assert arg is None or isinstance(self.device, tuple)
@@ -844,7 +861,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     if self.op is Ops.BUFFER: return self.arg.addrspace
     if self.op in {Ops.SPECIAL, Ops.RANGE}: return AddrSpace.ALU
     if self.op is Ops.LOAD: return AddrSpace.ALU # LOAD brings things into the ALU
-    if self.op in {Ops.INDEX, Ops.CAST, Ops.AFTER, Ops.REDUCE, Ops.STORE, Ops.MSTACK, Ops.MSELECT, Ops.END}:
+    if self.op in {Ops.INDEX, Ops.CAST, Ops.AFTER, Ops.REDUCE, Ops.STORE, Ops.MSTACK, Ops.MSELECT, Ops.END, Ops.UNSHARD}:
       return self.src[0].addrspace
     if self.op in GroupOp.Movement: return self.src[0].addrspace
     if self.op in {Ops.STACK, Ops.WMMA, Ops.GROUP} or self.op in GroupOp.Elementwise:
@@ -872,7 +889,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     # CL 1.1 provides the clCreateSubBuffer API, but at the time of writing, relevant CL runtimes (rusticl, adreno, nvidia, amd) do not provide
     # reasonable values for CL_DEVICE_MEM_BASE_ADDR_ALIGN. cl_ext_buffer_device_address could potentially help, but this extension is not provided
     # by relevant CL runtimes at time of writing.
-    if any(d.startswith(("WEBGPU", "CL")) for d in ((self.device,) if isinstance(self.device, str) else self.device)): return None
+    if (dev:=self.device) is not None and any(d.startswith(("WEBGPU", "CL")) for d in ((dev,) if isinstance(dev, str) else dev)): return None
 
     idx = self.flatten().index(UOp.range(self.numel(), 0))
     out = graph_rewrite(idx, pm_mops+symbolic+pm_contiguous_view_offset, ctx=self, name="contiguous_view_offset")
@@ -1163,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:
@@ -1174,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)
 
@@ -1187,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)
@@ -1224,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:
@@ -1743,7 +1764,7 @@ pm_lower_weak = PatternMatcher([
   # a weakfloat Unary (sin/exp2/...) must resolve here, before the transcendental decomposition
   (UPat(GroupOp.Binary|GroupOp.Unary|{Ops.WHERE, Ops.RANGE, Ops.STACK, Ops.SPECIAL}, name="u"), lower_weak_node),
   (UPat(Ops.PARAM, dtype=dtypes.weakint, name="u"),
-    lambda u: u.replace(dtype=None, arg=replace(u.arg, dtype=dtypes.int)).cast(dtypes.weakint) if u.addrspace == AddrSpace.ALU else None),
+    lambda u: u.replace(dtype=None, arg=replace(u.arg, dtype=select_dtype(u))).cast(dtypes.weakint) if u.addrspace == AddrSpace.ALU else None),
 ])
 def lower_weak_srcs(ctx:dict[UOp, UOp]|None, u:UOp) -> UOp|None:
   if ctx is None: ctx = {}

tinygrad/uop/spec.py

@@ -142,7 +142,8 @@ spec_tensor = PatternMatcher([
    if isinstance(buf.arg, ParamArg) and buf.addrspace is AddrSpace.GLOBAL else None),
 
   # Tensor variable bindings
-  (UPat(Ops.BIND, (dtypes.int, dtypes.weakint,), (UPat(Ops.PARAM), UPat.cvar(dtype=(dtypes.int,dtypes.weakint,))), arg=None), lambda: True),
+  (UPat(Ops.BIND, (dtypes.int, dtypes.long, dtypes.weakint,), (UPat(Ops.PARAM), UPat.cvar(dtype=(dtypes.int,dtypes.long,dtypes.weakint,))), arg=None),
+   lambda: True),
 
   # custom function
   (UPat(Ops.CUSTOM_FUNCTION, name="x"), lambda x: isinstance(x.arg, str)),
@@ -175,9 +176,9 @@ spec_tensor = PatternMatcher([
    len(red.arg) == 2 and red.arg[0] in GroupOp.Reduce and is_device(red.arg[1])),
 
   # UNSHARD/MSELECT/MSTACK
-  # an UNSHARD always has two srcs: the value and the DEVICE range it ends
-  (UPat(Ops.UNSHARD, name="multi"), lambda multi: len(multi.src) == 2 and matches_dtype(multi.src[0], multi.dtype)
-    and isinstance(multi.arg, int) and multi.src[1].op is Ops.RANGE and multi.src[1].arg[-1] is AxisType.DEVICE),
+  # an UNSHARD carries the value and one sharding range per sharded axis (usually a DEVICE RANGE, but can be a derived expression)
+  (UPat(Ops.UNSHARD, name="multi"), lambda multi: len(multi.src) == 1+len(multi.arg) and matches_dtype(multi.src[0], multi.dtype)
+    and all(isinstance(a, int) for a in multi.arg) and all(r.dtype in dtypes.weaks for r in multi.src[1:])),
   (UPat(Ops.MSELECT, name="x"), lambda x: isinstance(x.src[0].device, tuple) and x.arg < len(x.src[0].device)),
   (UPat(Ops.MSTACK, name="x"), lambda x: all(isinstance(s.device, str) for s in x.src) or (all_same(x.src) and x.src[0].device is None)),
 
@@ -206,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/uop/symbolic.py

@@ -262,12 +262,9 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
   # (x//c1)//c2 -> x//(c1*c2) for c2>0
   ((UPat.var("x") // UPat.cvar("c1")) // UPat.cvar("c2"), lambda x,c1,c2: x//(c1*c2) if c2.vmin>0 else None),
   # ** lt **
-  # c0*x<c1 for positive int c0,c1
+  # c0*x<c1 -> sign(c0)*x < ceil(c1/abs(c0))
   ((UPat.cvar("c0")*UPat.var("x", dtype=dtypes.weakint))<UPat.cvar("c1"),
-   lambda x,c0,c1: x<math.ceil(c1.arg/c0.arg) if c0.arg > 0 and c1.arg > 0 else None),
-  # c0*x<c1 for negative int c0 and non-positive c1
-  ((UPat.cvar("c0")*UPat.var("x", dtype=dtypes.weakint))<UPat.cvar("c1"),
-   lambda x,c0,c1: (-x)<(-(math.floor(-c1.arg/-c0.arg))) if c0.arg < 0 and c0.arg != -1 and c1.arg <= 0 else None),
+   lambda x,c0,c1: (x if c0.arg > 0 else -x)<-(-c1.arg//abs(c0.arg)) if abs(c0.arg) > 1 else None),
   # x//d<c -> x<c*d for d>0, and -> c*d<x for d<0
   ((UPat.var("x", dtype=dtypes.weakint)//UPat.cvar("d"))<UPat.cvar("c"),
    lambda x,d,c: (x<c.arg*d.arg) if d.arg > 0 else (x>c.arg*d.arg) if d.arg < 0 else None),

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