minor bugfixes

extra/benchmark_llm.py

@@ -0,0 +1,93 @@
+#!/usr/bin/env python3
+"""Benchmark tinygrad LLM prefill and decode independently.
+
+Examples:
+  python -m extra.benchmark_llm --model qwen3:0.6b --max-context 32768
+  python -m extra.benchmark_llm --model /path/to/model.gguf --prompt-tokens 8192 --realize
+"""
+from __future__ import annotations
+
+import argparse, json, statistics, time
+from dataclasses import asdict, dataclass
+
+from tinygrad.helpers import fetch
+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
+
+
+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]:
+  # Avoid tokenizer and chat-template work while exercising the same embedding/model path.
+  # Changing token zero guarantees that Transformer.get_start_pos cannot reuse an earlier KV cache.
+  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)
+  begin = time.perf_counter()
+  next(gen)
+  ttft = time.perf_counter() - begin
+
+  decode_times: list[float] = []
+  for _ in range(decode_tokens):
+    begin = time.perf_counter()
+    next(gen)
+    decode_times.append(time.perf_counter() - begin)
+
+  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)
+
+
+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("--chunk-size", type=int, default=256)
+  parser.add_argument("--realize", action="store_true", help="Unpack model weights once at load time")
+  parser.add_argument("--json", action="store_true", help="Print machine-readable results")
+  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 max(args.prompt_tokens) + args.decode_tokens >= args.max_context:
+    parser.error("prompt plus decode tokens must fit within --max-context")
+
+  path = fetch(models.get(args.model, args.model))
+  model, kv = Transformer.from_gguf(path, args.max_context, realize=args.realize)
+  vocab_size = len(kv["tokenizer.ggml.tokens"])
+
+  model.warmup(args.chunk_size)
+
+  results = [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,
+                      "realize": args.realize, "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} realize={args.realize}")
+  print(f"{'prompt':>8} {'TTFT':>10} {'prefill':>14} {'decode':>14} {'decode p50':>12} {'decode p95':>12}")
+  for result in results:
+    print(f"{result.prompt_tokens:8d} {result.time_to_first_token_s:9.3f}s {result.prefill_tokens_per_s:11.1f} t/s "
+          f"{result.decode_tokens_per_s:11.1f} t/s {result.decode_p50_ms:9.2f} ms {result.decode_p95_ms:9.2f} ms")
+
+
+if __name__ == "__main__": main()

extra/gemm/amd_flash_attention.py

@@ -1,13 +1,16 @@
 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
+from tinygrad.helpers import GlobalCounters, Context
+import functools, math
 
-BLOCK_M, BLOCK_N = 64, 64
+BLOCK_M, BLOCK_N = 32, 32
+DECODE_BLOCK_N = 128
+DECODE_HEAD_TILE = 4
+DECODE_WAVES = 4
 WARP_SIZE = 32
 WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
-WAVES_M, WAVES_N = 4, 1
+WAVES_M, WAVES_N = 2, 2
 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
@@ -35,24 +38,137 @@ def warp_reduce_sum(val, lane):
     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}"
+def wave_reduce_sum(val, lane):
+  for offset in [16, 8, 4, 2, 1]: val = val + warp_shfl_xor(val, offset, lane)
+  return val
+
+@functools.cache
+def _amd_flash_attention_decode_partial(out:UOp, stats:UOp, q:UOp, cache_kv:UOp, valid_kv_len:int|UOp, max_kv_len:int) -> UOp:
+  _, B, H_KV, N, D = cache_kv.shape
+  _, H, M, _ = q.shape
+  assert M == 1 and H % H_KV == 0 and D % WARP_SIZE == 0 and max_kv_len <= N and max_kv_len % DECODE_BLOCK_N == 0
+  G, CHUNK, DV = H // H_KV, DECODE_BLOCK_N, D // WARP_SIZE
+  assert G % DECODE_HEAD_TILE == 0
+  block_bhkv = UOp.range(B*H_KV*(G//DECODE_HEAD_TILE), 0, AxisType.GLOBAL)
+  block_n = UOp.range((valid_kv_len+CHUNK-1)//CHUNK, 1, AxisType.GLOBAL)
+  lane, wave = UOp.range(WARP_SIZE, 2, AxisType.LOCAL), UOp.range(DECODE_WAVES, 3, AxisType.LOCAL)
+  head_group = block_bhkv % (G//DECODE_HEAD_TILE)
+  bhkv = block_bhkv // (G//DECODE_HEAD_TILE)
+  b, kv_head = bhkv // H_KV, bhkv % H_KV
+  dims = tuple(lane + i*WARP_SIZE for i in range(DV))
+
+  acc = UOp.placeholder((DECODE_HEAD_TILE, DV), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+  row_max = UOp.placeholder((DECODE_HEAD_TILE,), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+  row_sum = UOp.placeholder((DECODE_HEAD_TILE,), dtypes.float, slot=2, addrspace=AddrSpace.REG)
+  init = UOp.group(acc.store(acc.const_like(0)), row_max.store(row_max.const_like(-math.inf)), row_sum.store(row_sum.const_like(0)))
+  acc, row_max, row_sum = acc.after(init), row_max.after(init), row_sum.after(init)
+
+  offset = UOp.range(CHUNK//DECODE_WAVES, 100, AxisType.REDUCE)
+  key = block_n*CHUNK + wave*(CHUNK//DECODE_WAVES) + offset
+  valid = key < valid_kv_len
+  kvals = tuple(cache_kv[0, b, kv_head, key, d].float() for d in dims)
+  vvals = tuple(cache_kv[1, b, kv_head, key, d].float() for d in dims)
+  updates = []
+  for head in range(DECODE_HEAD_TILE):
+    q_head = kv_head*G + head_group*DECODE_HEAD_TILE + head
+    score = wave_reduce_sum(sum((q[b, q_head, 0, d].float()*k for d,k in zip(dims, kvals)), UOp.const(dtypes.float, 0)), lane) / math.sqrt(D)
+    new_max = valid.where(row_max[head].maximum(score), row_max[head])
+    alpha = valid.where(((row_max[head]-new_max)*LOG2E).exp2(), UOp.const(dtypes.float, 1))
+    beta = valid.where(((score-new_max)*LOG2E).exp2(), UOp.const(dtypes.float, 0))
+    updates += [acc[head].store(acc[head]*alpha + UOp.stack(*vvals)*beta),
+                row_sum[head].store(row_sum[head]*alpha + beta), row_max[head].store(new_max)]
+  update = UOp.group(*updates).end(offset)
+  acc, row_max, row_sum = acc.after(update), row_max.after(update), row_sum.after(update)
+
+  partial_acc = UOp.placeholder((DECODE_HEAD_TILE, DECODE_WAVES, D), dtypes.float, slot=3, addrspace=AddrSpace.LOCAL)
+  partial_stats = UOp.placeholder((DECODE_HEAD_TILE, DECODE_WAVES, 2), dtypes.float, slot=4, addrspace=AddrSpace.LOCAL)
+  partial_stores = []
+  for head in range(DECODE_HEAD_TILE):
+    partial_stores += [partial_acc[head, wave, d].store(acc[head, i]) for i,d in enumerate(dims)]
+    partial_stores += [partial_stats[head, wave.valid(lane.eq(0)), 0].store(row_max[head]),
+                       partial_stats[head, wave.valid(lane.eq(0)), 1].store(row_sum[head])]
+  merged = UOp.group(*partial_stores).barrier()
+  stores = []
+  for head in range(DECODE_HEAD_TILE):
+    q_head = kv_head*G + head_group*DECODE_HEAD_TILE + head
+    maximum = partial_stats.after(merged)[head, 0, 0].maximum(partial_stats.after(merged)[head, 1, 0])
+    scales = tuple(((partial_stats.after(merged)[head, w, 0]-maximum)*LOG2E).exp2() for w in range(DECODE_WAVES))
+    denominator = sum((partial_stats.after(merged)[head, w, 1]*scales[w] for w in range(DECODE_WAVES)), UOp.const(dtypes.float, 0))
+    stores += [out[b, q_head, block_n, d.valid(wave.eq(0))].store(
+      sum((partial_acc.after(merged)[head, w, d]*scales[w] for w in range(DECODE_WAVES)), UOp.const(dtypes.float, 0))) for d in dims]
+    stores += [stats[b, q_head.valid(lane.eq(0) & wave.eq(0)), block_n, 0].store(maximum),
+               stats[b, q_head.valid(lane.eq(0) & wave.eq(0)), block_n, 1].store(denominator)]
+  return UOp.group(*stores).end(lane, wave, block_n, block_bhkv).sink(arg=KernelInfo(name="flash_decode_partial", opts_to_apply=()))
+
+@functools.cache
+def _amd_flash_attention_decode_reduce(out:UOp, partial:UOp, stats:UOp, valid_chunks:int|UOp) -> UOp:
+  B, H, _, D = out.shape
+  assert D % WARP_SIZE == 0
+  DV = D // WARP_SIZE
+  block_bh, lane = UOp.range(B*H, 0, AxisType.GLOBAL), UOp.range(WARP_SIZE, 1, AxisType.LOCAL)
+  b, head = block_bh // H, block_bh % H
+  dims = tuple(lane + i*WARP_SIZE for i in range(DV))
+
+  row_max = UOp.placeholder((1,), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+  row_max = row_max.after(row_max.store(row_max.const_like(-math.inf)))
+  chunk_max = UOp.range(valid_chunks, 100, AxisType.REDUCE)
+  max_done = row_max.store(row_max.after(chunk_max).maximum(stats[b, head, chunk_max, 0])).end(chunk_max)
+  row_max = row_max.after(max_done)
+
+  numerator = UOp.placeholder((DV,), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+  denominator = UOp.placeholder((1,), dtypes.float, slot=2, addrspace=AddrSpace.REG)
+  init = UOp.group(numerator.store(numerator.const_like(0)), denominator.store(denominator.const_like(0)))
+  numerator, denominator = numerator.after(init), denominator.after(init)
+  chunk = UOp.range(valid_chunks, 101, AxisType.REDUCE)
+  scale = ((stats[b, head, chunk, 0]-row_max[0])*LOG2E).exp2()
+  update = UOp.group(numerator.store(numerator.after(chunk) + UOp.stack(*(partial[b, head, chunk, d] for d in dims))*scale),
+                     denominator.store(denominator.after(chunk) + stats[b, head, chunk, 1]*scale)).end(chunk)
+  numerator, denominator = numerator.after(update), denominator.after(update)
+  stores = [out[b, head, 0, d].store(numerator[i]/denominator[0]) for i,d in enumerate(dims)]
+  return UOp.group(*stores).end(lane, block_bh).sink(arg=KernelInfo(name="flash_decode_reduce", opts_to_apply=()))
+
+def amd_flash_attention_decode(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, max_kv_len:int|None=None) -> Tensor:
+  _, B, H_KV, N, D = cache_kv.shape
+  _, H, M, _ = q.shape
+  max_kv_len = N if max_kv_len is None else max_kv_len
+  assert M == 1 and max_kv_len <= N and max_kv_len % DECODE_BLOCK_N == 0
+  chunks = max_kv_len // DECODE_BLOCK_N
+  partial = Tensor.empty(B, H, chunks, D, dtype="float32", device=q.device)
+  stats = Tensor.empty(B, H, chunks, 2, dtype="float32", device=q.device)
+  partial, stats = Tensor.custom_kernel(partial, stats, q, cache_kv,
+    fxn=functools.partial(_amd_flash_attention_decode_partial, valid_kv_len=valid_kv_len, max_kv_len=max_kv_len))[:2]
+  live_chunks = (valid_kv_len+DECODE_BLOCK_N-1)//DECODE_BLOCK_N
+  out = Tensor.empty(B, H, 1, D, dtype="float32", device=q.device)
+  return Tensor.custom_kernel(out, partial, stats,
+    fxn=functools.partial(_amd_flash_attention_decode_reduce, valid_chunks=live_chunks))[0]
+
+@functools.cache
+def _amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp, causal:bool, valid_kv_len:int|UOp|None=None,
+                         key_limit:int|UOp|None=None) -> UOp:
+  # inputs are q=(B*H, M, D), k/v=(B*H, N, D). For causal attention q is the final M tokens of k/v.
+  BH, M, D = q.shape
+  physical_n = k.shape[1]
+  N = physical_n if valid_kv_len is None else valid_kv_len
+  assert k.shape == v.shape and BH % k.shape[0] == 0 and k.shape[2] == D
+  gqa_group = BH // k.shape[0]
+  if isinstance(M, int) and isinstance(N, int):
+    assert M % BLOCK_M == 0 and N % BLOCK_N == 0, \
+      f"M={M} and N={N} must be divisible by BLOCK_M={BLOCK_M} and BLOCK_N={BLOCK_N}"
+  assert isinstance(D, int) and 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)
+  # Each N wave computes the same score tile, then owns a disjoint slice of D for P@V.
+  TN = BLOCK_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)
+  block_m = UOp.range(M // 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]
+  q = q.reshape(BH, M//BLOCK_M, BLOCK_M, D)[block_bh, block_m]
+  k, v = k[block_bh // gqa_group], v[block_bh // gqa_group]
+  o = o.reshape(BH, M//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)
@@ -63,7 +179,8 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
 
   # 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
+  Q_ELEMS_PER_THREAD = BLOCK_M * D // THREADS_PER_BLOCK
+  KV_ELEMS_PER_THREAD = BLOCK_N * 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]
 
@@ -76,14 +193,18 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
   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)
+  # Causal blocks never need KV tiles strictly to their right. Besides saving work, this avoids an all
+  # -inf tile, whose online-softmax update would otherwise contain -inf - -inf.
+  n_tiles = (N - M + (block_m + 1) * BLOCK_M + BLOCK_N - 1) // BLOCK_N if causal else N // BLOCK_N
+  n_tile = UOp.range(n_tiles, 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])
+  Q_store = Q_lds.after(n_tile).reshape(THREADS_PER_BLOCK, Q_ELEMS_PER_THREAD)[tid].store(
+    q.reshape(THREADS_PER_BLOCK, Q_ELEMS_PER_THREAD)[tid])
+  load_k = UOp.range(KV_ELEMS_PER_THREAD, 90, AxisType.LOOP)
+  K_store = KV_lds.reshape(THREADS_PER_BLOCK, KV_ELEMS_PER_THREAD)[tid, load_k].store(
+    k.reshape(physical_n*D)[n_tile*BLOCK_N*D + tid*KV_ELEMS_PER_THREAD + load_k]).end(load_k)
   qk_load_barrier = UOp.barrier(UOp.group(Q_store, K_store))
   Q_lds = Q_lds.after(qk_load_barrier)
   KV_lds_k = KV_lds.after(qk_load_barrier)
@@ -96,7 +217,7 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
   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]
+  k_frag = KV_lds_k.reshape(TN, WMMA_N, D // WMMA_K, WMMA_K)[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)
@@ -104,6 +225,18 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
   # -- softmax in registers with warp shuffles --
   S_reg = S_reg.after(S_reg.store(S_reg * SCALE))
 
+  if causal:
+    # WMMA accumulator ownership: each lane owns an 8x4 fragment of the 64x64 score tile.
+    # q is aligned to the right of k, matching PyTorch's causal_lower_right mask.
+    rm = UOp.range(TM, 250, AxisType.LOOP)
+    rn = UOp.range(TN, 251, AxisType.LOOP)
+    q_idx = N - M + block_m * BLOCK_M + wave_m * WMMA_M + rm * LANES_PER_WAVE_M + lane_m
+    k_idx = n_tile * BLOCK_N + rn * LANES_PER_WAVE_N + lane_n
+    valid = k_idx <= q_idx
+    if key_limit is not None: valid = valid & (k_idx < key_limit)
+    masked = valid.where(S_reg[rm, rn], S_reg[rm, rn].const_like(-math.inf))
+    S_reg = S_reg.after(S_reg[rm, rn].store(masked).end(rm, rn))
+
   # 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)))
@@ -118,18 +251,19 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
 
   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))
+  # Reduce contiguous 16-key groups independently, matching the ordinary softmax reduction tree.
+  p_sum = p_local.after(p_local[ri_ws].store(
+    sum((warp_reduce_sum(S_reg[ri_ws, rn], lane) for rn in range(TN)), S_reg.const_like(0))).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)
+  # Store softmax weights in half for the WMMA P@V product; accumulation remains float.
+  P_lds = QP_lds.flatten()[:WAVES_N * BLOCK_M * BLOCK_N].reshape(WAVES_N, BLOCK_M, BLOCK_N)
+  P_write = P_lds.reshape(WAVES_N, WAVES_M, TM, LANES_PER_WAVE_M, 1, TN, LANES_PER_WAVE_N, 1)
+  P_write = P_write.permute((1, 0, 3, 6, 2, 4, 5, 7)).reshape(THREADS_PER_BLOCK, TM, TN)
   P_store = P_write[tid].store(S_reg.cast(dtypes.half))
 
   # -- online softmax correction --
+  beta_i = UOp.placeholder((TM,), dtypes.float, slot=9, addrspace=AddrSpace.REG)
   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()
@@ -139,29 +273,43 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
     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),
+    beta_i[ri4].store(beta_val),
   ).end(ri4)
   acc = acc.after(correction)
   l_i = l_i.after(correction)
   m_i = m_i.after(correction)
+  beta_i = beta_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])
+  # Load V transposed into LDS: PV's B operand is logically (D, BLOCK_N), while global V is (BLOCK_N, D).
+  # It reuses K's slot and must wait for QK WMMA to finish reading that slot.
+  V_lds = UOp.placeholder((D, BLOCK_N + LDS_PAD), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL)[:, :BLOCK_N]
+  V_copy = V_lds.after(qk_done).permute(1, 0)
+  load_v = UOp.range(KV_ELEMS_PER_THREAD, 390, AxisType.LOOP)
+  V_store = V_copy.reshape(THREADS_PER_BLOCK, KV_ELEMS_PER_THREAD)[tid, load_v].store(
+    v.reshape(physical_n*D)[n_tile*BLOCK_N*D + tid*KV_ELEMS_PER_THREAD + load_v]).end(load_v)
   pv_barrier = UOp.barrier(UOp.group(P_store, V_store))
   P_lds = P_lds.after(pv_barrier)
-  KV_lds_v = KV_lds.after(pv_barrier)
+  V_lds = V_lds.after(pv_barrier)
 
-  # -- acc += P @ V via WMMA --
+  # -- acc += beta * (P @ V) via WMMA --
+  pv_acc = UOp.placeholder((TM, TD), dtypes.float, slot=10, addrspace=AddrSpace.REG)
+  pv_acc = pv_acc.after(pv_acc.after(n_tile).store(pv_acc.const_like(0))).after(pv_barrier)
   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)
+  pv_frag = pv_acc.reshape(TM // WMMA_ACC, WMMA_ACC, TD).permute(0, 2, 1)[tm2, tn2]
+  p_frag = P_lds[wave_n].reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_N // WMMA_K, WMMA_K)[wave_m, tm2, lane_n, k_pv]
+  v_frag = V_lds.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, pv_frag.after(k_pv), *WMMA_ARG)
+  pv_done = pv_frag.store(pv).end(tm2, tn2).end(k_pv)
+  pv_acc = pv_acc.after(pv_done)
+
+  ri5 = UOp.range(TM, 410, AxisType.LOOP)
+  rj5 = UOp.range(TD, 411, AxisType.LOOP)
+  accumulate = acc[ri5, rj5].store(acc[ri5, rj5] + beta_i[ri5] * pv_acc[ri5, rj5]).end(ri5, rj5)
 
   # end KV tile loop
-  n_tile_end = acc_frag.store(pv).end(tm2, tn2).end(k_pv).barrier().end(n_tile)
+  n_tile_end = accumulate.barrier().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)
@@ -170,33 +318,49 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
   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)
+  o = o.reshape(WAVES_M, TM, LANES_PER_WAVE_M, 1, WAVES_N, TD, LANES_PER_WAVE_N, 1)
+  o = o.permute((0, 4, 2, 6, 1, 3, 5, 7)).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=()))
 
+def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
+  return _amd_flash_attention(o, q, k, v, causal=False)
+
+def amd_flash_attention_causal(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
+  return _amd_flash_attention(o, q, k, v, causal=True)
+
+def amd_flash_attention_causal_cached(o:UOp, q:UOp, cache_kv:UOp, *, valid_kv_len:int|UOp, key_limit:int|UOp|None=None) -> UOp:
+  _, B, H_KV, N, D = cache_kv.shape
+  k = cache_kv[0].reshape(B*H_KV, N, D)
+  v = cache_kv[1].reshape(B*H_KV, N, D)
+  return _amd_flash_attention(o, q, k, v, causal=True, valid_kv_len=valid_kv_len, key_limit=key_limit)
+
 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)
+  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)
   NUM_RUNS = getenv("CNT", 5)
   ets = []
   with Context(DEBUG=2):
     for _ in range(NUM_RUNS):
       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}")
+      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!")
     else:

test/external/external_test_llm_opencode.py

@@ -0,0 +1,300 @@
+"""End-to-end regression tests for tinygrad's OpenCode integration.
+
+Run against an existing server:
+  RUN_LLM_OPENCODE_REGRESSION=1 LLM_BASE_URL=http://127.0.0.1:9000/v1 \
+    python -m pytest test/external/external_test_llm_opencode.py -v
+
+Or let the test start the server:
+  RUN_LLM_OPENCODE_REGRESSION=1 \
+    LLM_GGUF=/raid/models/Qwen3.6-35B-A3B-UD-IQ4_XS.gguf \
+    python -m pytest test/external/external_test_llm_opencode.py -v
+"""
+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"
+DEFAULT_GGUF = "/raid/models/Qwen3.6-35B-A3B-UD-IQ4_XS.gguf"
+
+SORT_C = r"""#include <stdio.h>
+#include <stdlib.h>
+#include <time.h>
+#include <stdint.h>
+#include <string.h>
+
+#define N 1000000
+
+/* Optimized LSD Radix sort using 8-bit chunks (256 buckets) with loop unrolling and prefetching */
+static void radix_sort(int *a, int *buf, int n) {
+    const int BITS = 8;
+    const int MASK = (1 << BITS) - 1;
+    const int BUCKETS = (1 << BITS);
+    uint32_t count[BUCKETS];
+    uint32_t *u = (uint32_t *)a;
+    uint32_t *ubuf = (uint32_t *)buf;
+    uint32_t *src = u;
+    uint32_t *dst = ubuf;
+
+    /* Convert signed to unsigned by flipping sign bit */
+    for (int i = 0; i < n; i++)
+        u[i] ^= (uint32_t)1U << 31;
+
+    for (int shift = 0; shift < 32; shift += BITS) {
+        memset(count, 0, sizeof(count));
+
+        /* Counting pass - unrolled by 8 with prefetching */
+        int i = 0;
+        for (; i + 7 < n; i += 8) {
+            count[(src[i] >> shift) & MASK]++;
+            count[(src[i+1] >> shift) & MASK]++;
+            count[(src[i+2] >> shift) & MASK]++;
+            count[(src[i+3] >> shift) & MASK]++;
+            count[(src[i+4] >> shift) & MASK]++;
+            count[(src[i+5] >> shift) & MASK]++;
+            count[(src[i+6] >> shift) & MASK]++;
+            count[(src[i+7] >> shift) & MASK]++;
+        }
+        for (; i < n; i++)
+            count[(src[i] >> shift) & MASK]++;
+
+        /* Prefix sums in-place */
+        uint32_t total = 0;
+        for (int i = 0; i < BUCKETS; i++) {
+            uint32_t c = count[i];
+            count[i] = total;
+            total += c;
+        }
+
+        /* Distribution pass - unrolled by 8 with prefetching */
+        i = 0;
+        for (; i + 7 < n; i += 8) {
+            dst[count[(src[i] >> shift) & MASK]++] = src[i];
+            dst[count[(src[i+1] >> shift) & MASK]++] = src[i+1];
+            dst[count[(src[i+2] >> shift) & MASK]++] = src[i+2];
+            dst[count[(src[i+3] >> shift) & MASK]++] = src[i+3];
+            dst[count[(src[i+4] >> shift) & MASK]++] = src[i+4];
+            dst[count[(src[i+5] >> shift) & MASK]++] = src[i+5];
+            dst[count[(src[i+6] >> shift) & MASK]++] = src[i+6];
+            dst[count[(src[i+7] >> shift) & MASK]++] = src[i+7];
+        }
+        for (; i < n; i++)
+            dst[count[(src[i] >> shift) & MASK]++] = src[i];
+
+        /* Swap src/dst pointers */
+        uint32_t *tmp = src;
+        src = dst;
+        dst = tmp;
+    }
+
+    /* Copy back if needed */
+    if (src != u)
+        memcpy(u, src, n * sizeof(uint32_t));
+
+    /* Convert back to signed */
+    for (int i = 0; i < n; i++)
+        u[i] ^= (uint32_t)1U << 31;
+}
+
+int main() {
+    int *arr = malloc(N * sizeof(int));
+    int *buf = malloc(N * sizeof(int));
+    if (!arr || !buf) { perror("malloc"); return 1; }
+
+    srand(42);
+    for (int i = 0; i < N; i++)
+        arr[i] = rand() | (rand() << 15);
+
+    clock_t start = clock();
+    radix_sort(arr, buf, N);
+    clock_t end = clock();
+
+    double elapsed = (double)(end - start) / CLOCKS_PER_SEC * 1000;
+    printf("Sorted %d integers in %.3f ms\n", N, elapsed);
+
+    for (int i = 1; i < N; i++) {
+        if (arr[i] < arr[i - 1]) {
+            printf("ERROR: not sorted at index %d\n", i);
+            free(arr);
+            free(buf);
+            return 1;
+        }
+    }
+    printf("Verification passed.\n");
+
+    free(arr);
+    free(buf);
+    return 0;
+}
+"""
+
+
+def free_port() -> int:
+  with socket.socket() as sock:
+    sock.bind(("127.0.0.1", 0))
+    return sock.getsockname()[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, urllib.error.URLError):
+    return False
+
+
+@unittest.skipUnless(RUN_REGRESSION, "set RUN_LLM_OPENCODE_REGRESSION=1 to run the real-model 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")) is not None:
+      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.getenv("LLM_GGUF", DEFAULT_GGUF))
+    if not model.is_file(): raise unittest.SkipTest(f"model not found: {model}")
+    port = free_port()
+    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", "65536"],
+      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 setUp(self):
+    # The model and its KV/recurrent caches are stateful. Keep xdist workers from interleaving independent OpenCode
+    # conversations, which changes cache reuse and can hide precisely the incremental path this suite exercises.
+    import fcntl
+    self._fcntl = fcntl
+    self._server_lock = open("/tmp/tinygrad-llm-opencode-regression.lock", "w")
+    self._fcntl.flock(self._server_lock, self._fcntl.LOCK_EX)
+
+  def tearDown(self):
+    self._fcntl.flock(self._server_lock, self._fcntl.LOCK_UN)
+    self._server_lock.close()
+
+  def run_opencode(self, prompt:str, cwd:pathlib.Path, timeout:int=120) -> 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.copy()
+    env["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=timeout)
+    self.assertEqual(result.returncode, 0, result.stdout)
+    return re.sub(r"\x1b\[[0-?]*[ -/]*[@-~]", "", result.stdout)
+
+  def chat(self, messages:list[dict], max_tokens:int=4) -> dict:
+    request = urllib.request.Request(self.base_url + "/chat/completions", data=json.dumps({
+      "model":"tinygrad", "messages":messages, "max_tokens":max_tokens, "temperature":0,
+    }).encode(), headers={"Content-Type":"application/json"})
+    with urllib.request.urlopen(request, timeout=120) as response: return json.load(response)
+
+  def test_same_session_reuses_prompt_cache(self):
+    # Use a long stable prefix so reuse remains observable after the model rounds cache positions for its serving JIT.
+    messages = [{"role":"system", "content":"You are a concise assistant. " * 400}, {"role":"user", "content":"Reply OK."}]
+    first = self.chat(messages)
+    messages += [{"role":"assistant", "content":first["choices"][0]["message"].get("content") or ""},
+                 {"role":"user", "content":"Reply OK again."}]
+    second = self.chat(messages)
+    self.assertGreater(second["usage"]["prompt_tokens_details"]["cached_tokens"], 0)
+
+  def test_reads_and_correctly_explains_valid_c(self):
+    with tempfile.TemporaryDirectory() as directory:
+      cwd = pathlib.Path(directory)
+      source = cwd / "sort.c"
+      source.write_text(SORT_C)
+      subprocess.run(["cc", "-std=c11", "-Wall", "-Werror", "-fsyntax-only", str(source)], check=True)
+
+      output = self.run_opencode("sort.c?", cwd)
+      self.assertRegex(output, r"(?im)^\s*(?:→|>)\s*Read\s+sort\.c\s*$", "OpenCode did not execute the read tool")
+      self.assertRegex(output, r"(?i)radix sort")
+      self.assertNotRegex(output, r"(?i)(corrupt|mangled|not valid C|invalid C|syntax error|does not compile|won't compile)")
+      self.assertEqual(source.read_text(), SORT_C)
+      subprocess.run(["cc", "-std=c11", "-Wall", "-Werror", "-fsyntax-only", str(source)], check=True)
+
+  def test_read_tool_preserves_exact_contents(self):
+    with tempfile.TemporaryDirectory() as directory:
+      cwd = pathlib.Path(directory)
+      marker = "tinygrad-read-regression-7f3a91c2"
+      source = cwd / "exact.txt"
+      source.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*$", "OpenCode did not execute the read tool")
+      self.assertIn(marker, output)
+      self.assertEqual(source.read_text(), marker + "\n")
+
+  def test_executes_shell_tool(self):
+    with tempfile.TemporaryDirectory() as directory:
+      cwd = pathlib.Path(directory)
+      marker = cwd / "shell-regression.txt"
+      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",
+                       "OpenCode did not execute the shell tool")
+      self.assertTrue(marker.is_file(), output)
+      self.assertEqual(marker.read_text(), "tinygrad-shell-regression")
+
+  def test_does_not_repeat_identical_failed_shell_call(self):
+    with tempfile.TemporaryDirectory() as directory:
+      cwd = pathlib.Path(directory)
+      command = "clang -x c /dev/null -fsyntax-only -mllvm -tinygrad-definitely-invalid-option=1"
+      output = self.run_opencode(
+        f"Run `{command}` exactly once with the shell tool. After it fails, do not retry it; explain that the option is unsupported.", cwd)
+      self.assertIn("Unknown command line argument", output)
+      self.assertLessEqual(output.count(command), 1, output)
+
+  def test_stops_when_benchmark_goal_is_met(self):
+    with tempfile.TemporaryDirectory() as directory:
+      cwd = pathlib.Path(directory)
+      benchmark = cwd / "benchmark.sh"
+      benchmark.write_text("#!/bin/sh\necho 'Sorted 1000000 integers in 8.300 ms'\n")
+      benchmark.chmod(0o755)
+      output = self.run_opencode(
+        "Use the shell tool to run ./benchmark.sh. Keep optimizing until it reports under 10 ms, then stop immediately.", cwd)
+      self.assertIn("8.300 ms", output)
+      self.assertLessEqual(output.count("$ ./benchmark.sh"), 1, output)
+
+  def test_multiline_tool_argument_preserves_trailing_newline(self):
+    with tempfile.TemporaryDirectory() as directory:
+      cwd = pathlib.Path(directory)
+      target = cwd / "numbers.txt"
+      target.write_text("replace me\n")
+      output = self.run_opencode(
+        "Read numbers.txt, then use the write tool to replace it with the numbers 1 through 300, one number per line. Do not use bash.", cwd)
+      self.assertRegex(output, r"(?im)^\s*(?:←|→|>)\s*Write\s+numbers\.txt\s*$", "OpenCode did not execute the write tool")
+      self.assertEqual(target.read_text(), "".join(f"{i}\n" for i in range(1, 301)))
+
+
+if __name__ == "__main__": unittest.main()

test/null/test_llm_server.py

@@ -17,6 +17,7 @@ class TestLLMServer(unittest.TestCase):
     cls.mock_tok.is_end = Mock(side_effect=lambda tid: tid in (999,))
 
     cls.mock_model = Mock()
+    cls.mock_model.max_context = 4
     cls.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 999]))
     cls.mock_model.get_start_pos = Mock(return_value=0)
 
@@ -129,6 +130,16 @@ class TestLLMServer(unittest.TestCase):
     self.assertIsNotNone(resp.usage.prompt_tokens)
     self.assertIsNotNone(resp.usage.completion_tokens)
 
+  def test_context_length_error(self):
+    from openai import BadRequestError
+    self.mock_tok.encode.return_value = [200, 201, 202, 203]
+    try:
+      with self.assertRaises(BadRequestError) as err:
+        self.client.chat.completions.create(model="test-model", messages=[{"role":"user", "content":"too long"}])
+      self.assertEqual(err.exception.code, "context_length_exceeded")
+    finally:
+      self.mock_tok.encode.return_value = [200, 201, 202]
+
   def test_max_tokens_streaming(self):
     self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 302, 303, 999]))
     stream = self.client.chat.completions.create(
@@ -170,13 +181,12 @@ class TestLLMToolCalls(unittest.TestCase):
     cls.mock_tok.is_end = Mock(return_value=False)
 
     cls.mock_model = Mock()
+    cls.mock_model.max_context = 4
     cls.mock_model.get_start_pos = Mock(return_value=0)
 
+    from tinygrad.llm.cli import FallbackTemplate
     from tinygrad.llm.serve import LLMServer
-    import jinja2
-    # .items() matches tool-aware templates and ensures OpenAI JSON argument strings are normalized before rendering the next turn.
-    template = jinja2.Template("""{% for m in messages %}{{ m.content or '' }}{% for tc in m.tool_calls or [] %}
-      {% for key, value in tc.function.arguments.items() %}{{ key }}={{ value }}{% endfor %}{% endfor %}{% endfor %}""")
+    template = FallbackTemplate(cls.mock_tok)
     cls.server = LLMServer(('127.0.0.1', 0), cls.mock_model, "tool-model", cls.mock_tok, template)
     cls.port = cls.server.server_address[1]
     cls.server_thread = threading.Thread(target=cls.server.serve_forever, daemon=True)
@@ -220,6 +230,20 @@ class TestLLMToolCalls(unittest.TestCase):
     self.assertEqual([json.loads(tc.function.arguments)["path"] for tc in response.choices[0].message.tool_calls], ["a", "b"])
     self.assertEqual(response.choices[0].finish_reason, "tool_calls")
 
+  def test_multiline_tool_argument_preserves_trailing_newline(self):
+    self.set_output("<tool_call>\n<function=write>\n<parameter=content>\nfirst\nsecond\n\n</parameter>\n"
+                    "<parameter=filePath>\nout.txt\n</parameter>\n</function>\n</tool_call>")
+    response = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Write out.txt"}], tools=self.tools())
+    args = json.loads(response.choices[0].message.tool_calls[0].function.arguments)
+    self.assertEqual(args, {"content":"first\nsecond\n", "filePath":"out.txt"})
+
+  def test_prefilled_reasoning_round_trip(self):
+    self.set_output("reasoning\n</think>\n\nanswer")
+    response = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Think"}],
+                                                   extra_body={"enable_thinking":True})
+    self.assertEqual(response.choices[0].message.reasoning_content, "reasoning\n")
+    self.assertEqual(response.choices[0].message.content, "\n\nanswer")
+
   def test_invalid_tool_call_becomes_content(self):
     self.set_output("<tool_call>not a call</tool_call>")
     response = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Hello"}], tools=self.tools())
@@ -247,5 +271,69 @@ class TestLLMToolCalls(unittest.TestCase):
     self.assertEqual(second.choices[0].message.content, "done")
     self.assertEqual(second.choices[0].finish_reason, "stop")
 
+  def test_tool_turn_remains_a_reusable_prefix_after_next_user_message(self):
+    class Tokenizer:
+      eos_id, eot_id = 0x110000, None
+      def encode(self, text): return [ord(c) for c in text]
+      def stream_decoder(self): return lambda tid=None: "" if tid is None else chr(tid)
+      def is_end(self, token_id): return token_id == self.eos_id
+    class Template:
+      def render(self, messages, tools=None, add_generation_prompt=True, enable_thinking=False, preserve_thinking=False):
+        out = ""
+        for m in messages:
+          content = (m.get("content") or "").strip()
+          if m["role"] == "assistant":
+            reasoning = (m.get("reasoning_content") or "").strip()
+            out += f"<assistant><think>\n{reasoning}\n</think>\n\n{content}"
+            for tc in m.get("tool_calls") or []:
+              fn = tc["function"]
+              out += ("\n\n" if content else "") + f"<tool_call>\n<function={fn['name']}>\n"
+              for name,value in fn["arguments"].items(): out += f"<parameter={name}>\n{value}\n</parameter>\n"
+              out += "</function>\n</tool_call>"
+          else: out += f"<{m['role']}>{content}"
+          out += "</turn>"
+        if add_generation_prompt: out += "<assistant><think>\n" if enable_thinking else "<assistant><think>\n\n</think>\n\n"
+        return out
+    class Model:
+      max_context = 10000
+      def __init__(self):
+        self.cached = []
+        self.outputs = iter(("reasoning\n</think>\n\nChecking now.\n\n<tool_call>\n<function=read>\n"
+                             "<parameter=path>\nsort.c\n</parameter>\n</function>\n</tool_call>",
+                             "finished reasoning\n</think>\n\nfinished."))
+      def get_start_pos(self, ids):
+        return next((i for i,(a,b) in enumerate(zip(ids, self.cached)) if a != b), min(len(ids), len(self.cached)))
+      def generate(self, ids, **kwargs):
+        output = [ord(c) for c in next(self.outputs)]
+        self.cached = ids + output
+        yield from output
+        yield Tokenizer.eos_id
+
+    from tinygrad.llm.serve import LLMServer
+    model = Model()
+    server = LLMServer(('127.0.0.1', 0), model, "prefix-model", Tokenizer(), Template())
+    thread = threading.Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    from openai import OpenAI
+    client = OpenAI(base_url=f"http://127.0.0.1:{server.server_address[1]}/v1", api_key="test")
+    try:
+      messages = [{"role":"user", "content":"Read sort.c"}]
+      first = client.chat.completions.create(model="prefix-model", messages=messages, tools=self.tools(),
+                                             extra_body={"enable_thinking":True})
+      self.assertEqual(first.choices[0].message.reasoning_content, "reasoning\n")
+      self.assertEqual(first.choices[0].message.content, "\n\nChecking now.\n\n")
+      cached_len = len(model.cached)
+      call = first.choices[0].message.tool_calls[0]
+      messages += [{"role":"assistant", "content":first.choices[0].message.content,
+                    "reasoning_content":first.choices[0].message.reasoning_content, "tool_calls":[call.model_dump()]},
+                   {"role":"tool", "tool_call_id":call.id, "content":"file contents"}]
+      second = client.chat.completions.create(model="prefix-model", messages=messages, tools=self.tools(),
+                                              extra_body={"enable_thinking":True})
+      self.assertEqual(second.choices[0].message.content, "\n\nfinished.")
+      self.assertEqual(second.usage.prompt_tokens_details.cached_tokens, cached_len)
+    finally:
+      server.shutdown()
+      server.server_close()
+
 if __name__ == '__main__':
   unittest.main()

test/null/test_llm_tokenizer.py

@@ -46,6 +46,18 @@ class TestLLMTokenizer(unittest.TestCase):
   def test_llama_repeat(self): self._test_coding(self.llama_tok, "00000000000000000", [ 931, 931, 931, 931, 931, 410 ])
   def test_llama_pat(self): self._test_coding(self.llama_tok, "today\n  \n", [ 31213, 14211 ])
 
+  def test_llama_continued_conversation(self):
+    self._test_coding(self.llama_tok, "hello <|eot_id|>world", [15339, 220, 128009, 14957])
+    self._test_coding(self.llama_tok, "hello <|eot_id|>world again", [15339, 220, 128009, 14957, 1578])
+    self._test_coding(self.llama_tok, "hello changed <|eot_id|>world again", [15339, 5614, 220, 128009, 14957, 1578])
+
+  def test_long_cached_prompt_matches_fresh_tokenization(self):
+    prefix = "system tools\n" * 700 + "<|eot_id|>"
+    first, changed = prefix + "run tower of hanoi", prefix + "run ls /"
+    expected = self.llama_tok.encode(changed)
+    self.llama_tok.encode(first)
+    self.assertEqual(self.llama_tok.encode(changed), expected)
+
   def test_tekken_from_gguf_kv(self):
     kv = {
       "tokenizer.ggml.tokens": ["<unk>", "<s>", "</s>", "[INST]", "[/INST]", "hello"],

test/unit/test_attention.py

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

test/unit/test_llm_server.py

@@ -151,7 +151,7 @@ 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()))
       return Tensor([[42]])
     with patch.object(Transformer, '__call__', mock_call):

tinygrad/codegen/__init__.py

@@ -1,7 +1,7 @@
 from dataclasses import replace, dataclass
-import itertools, functools
+import itertools, functools, hashlib, pickle
 from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC
-from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, panic
+from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, panic, CCACHE, diskcache_get, diskcache_put
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, GroupOp
 from tinygrad.uop.ops import AxisType
 from tinygrad.uop.render import pyrender
@@ -466,5 +466,10 @@ to_program_cache: dict[tuple, UOp] = {}
 def to_program(ast:UOp, renderer:Renderer) -> UOp:
   config = (NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32)
   key = (ast.key, type(renderer), renderer.target, *[x.value for x in config])
-  if (prg:=to_program_cache.get(key)) is None: to_program_cache[key] = prg = do_to_program(ast, renderer)
+  if (prg:=to_program_cache.get(key)) is None:
+    disk_key = hashlib.sha256(pickle.dumps(key)).digest()
+    if not CCACHE or (prg:=diskcache_get("program", {"key":disk_key})) is None:
+      prg = do_to_program(ast, renderer)
+      if CCACHE: diskcache_put("program", {"key":disk_key}, prg)
+    to_program_cache[key] = prg
   return prg

tinygrad/engine/jit.py

@@ -39,7 +39,6 @@ def graph_split_rewrite(linear:UOp, max_batch_size:int=0) -> UOp:
     if len(current_batch) <= 1 and not getenv("GRAPH_ONE_KERNEL"): new_src.extend(current_batch)
     else:
       new_src.append(create_graph_call(current_batch))
-      max_batch_size *= 2
       if DEBUG >= 2: print(f"JIT GRAPHing batch with {len(current_batch)} kernels")
     current_batch, current_batch_devs = [], []
 

tinygrad/llm/cli.py

@@ -1,9 +1,7 @@
 from __future__ import annotations
 import sys, argparse, codecs, typing, re, unicodedata, json, time
 from typing import TYPE_CHECKING
-from tinygrad import nn
-from tinygrad.uop.ops import UOp, Ops
-from tinygrad.helpers import partition, DEBUG, Timing, GlobalCounters, Context, fetch, profile_marker, getenv
+from tinygrad.helpers import BEAM, DEBUG, JIT_BATCH_SIZE, Timing, GlobalCounters, Context, fetch, profile_marker, getenv
 from tinygrad.llm.model import Transformer
 if TYPE_CHECKING:
   import jinja2
@@ -20,24 +18,36 @@ class SimpleTokenizer:
 
     # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286
     # 0x323b0 is one past the max codepoint in unicode categories L/N/Z (0x323af is max L)
-    def ucat_range(pre: str): return "".join(re.escape(chr(cp)) for cp in range(0x323b0) if unicodedata.category(chr(cp)).startswith(pre))
+    # Compact adjacent codepoints into ranges. Build L/N/Z together: scanning Unicode three times is measurable at server startup.
+    runs: dict[str, list[tuple[int, int]]] = {pre:[] for pre in "LNZ"}
+    for cp in range(0x323b0):
+      if (pre:=unicodedata.category(chr(cp))[0]) not in runs: continue
+      if runs[pre] and cp == runs[pre][-1][1]+1: runs[pre][-1] = (runs[pre][-1][0], cp)
+      else: runs[pre].append((cp, cp))
+    def ucat_range(pre:str) -> str:
+      def esc(cp:int) -> str: return f"\\U{cp:08x}"
+      return "".join(esc(st) if st == en else f"{esc(st)}-{esc(en)}" for st,en in runs[pre])
     r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + ucat_range("Z"), ucat_range("N"), ucat_range("L")
     self._split_to_word = re.compile("(?i:'s|'t|'re|'ve|'m|'ll|'d)|" + \
       f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+")
     self._split_to_sentence = re.compile("|".join(re.escape(tok) for tok in special_tokens.keys()) if special_tokens else r"(?!)")
 
-    self._normal_tokens = {bytes(self._byte_decoder[c] for c in tok): tid for tok, tid in normal_tokens.items()}
+    byte_translation = str.maketrans(self._byte_decoder)
+    self._normal_tokens = {tok.translate(byte_translation).encode("latin1"): tid for tok, tid in normal_tokens.items()}
     self._special_tokens = special_tokens
     self._tok2bytes = {tid: tok for tok, tid in self._normal_tokens.items()} | {tid: tok.encode() for tok, tid in self._special_tokens.items()}
+    self._encode_cache: tuple[str, tuple[int, ...], list[tuple[int, int]]]|None = None
     self.preset = preset
     self.bos_id, self.eos_id, self.eot_id = bos_id, eos_id, eot_id
 
   @staticmethod
   def from_gguf_kv(kv:dict):
     # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L1818-L1820
-    vocab: typing.Iterable[tuple[str, int]] = ((tok, idx) for idx, tok in enumerate(kv["tokenizer.ggml.tokens"]))
-    normal_tokens, special_tokens = partition(vocab, lambda e: kv["tokenizer.ggml.token_type"][e[1]] == 1)
-    return SimpleTokenizer(dict(normal_tokens), dict(special_tokens), kv["tokenizer.ggml.pre"],
+    normal_tokens: dict[str, int] = {}
+    special_tokens: dict[str, int] = {}
+    for idx,(tok,token_type) in enumerate(zip(kv["tokenizer.ggml.tokens"], kv["tokenizer.ggml.token_type"])):
+      (normal_tokens if token_type == 1 else special_tokens)[tok] = idx
+    return SimpleTokenizer(normal_tokens, special_tokens, kv["tokenizer.ggml.pre"],
       bos_id=kv.get('tokenizer.ggml.bos_token_id') if kv.get('tokenizer.ggml.add_bos_token', True) else None,
       eos_id=kv.get('tokenizer.ggml.eos_token_id', 0), eot_id=kv.get('tokenizer.ggml.eot_token_id'))
 
@@ -56,10 +66,23 @@ class SimpleTokenizer:
   def encode(self, text:str) -> list[int]:
     tokens: list[int] = []
     pos = 0
-    for match in self._split_to_sentence.finditer(text):
+    checkpoints: list[tuple[int, int]] = []
+    if self._encode_cache is not None:
+      old_text, old_tokens, old_checkpoints = self._encode_cache
+      if text == old_text: return list(old_tokens)
+      common, limit = 0, min(len(text), len(old_text))
+      while common+4096 <= limit and text[common:common+4096] == old_text[common:common+4096]: common += 4096
+      common += next((i for i,(a,b) in enumerate(zip(text[common:limit], old_text[common:limit])) if a != b), limit-common)
+      if (checkpoint := next((x for x in reversed(old_checkpoints) if x[0] <= common), None)) is not None:
+        pos, token_pos = checkpoint
+        tokens, checkpoints = list(old_tokens[:token_pos]), [x for x in old_checkpoints if x[0] <= pos]
+    for match in self._split_to_sentence.finditer(text, pos):
       tokens.extend(self._encode_sentence(text[pos:match.start(0)]) + [self._special_tokens[text[match.start(0):match.end(0)]]])
       pos = match.end(0)
-    return tokens + self._encode_sentence(text[pos:])
+      checkpoints.append((pos, len(tokens)))
+    tokens += self._encode_sentence(text[pos:])
+    self._encode_cache = text, tuple(tokens), checkpoints
+    return tokens
 
   def decode(self, ids:list[int]) -> str: return b''.join(self._tok2bytes[tid] for tid in ids).decode(errors='replace')
   def stream_decoder(self) -> typing.Callable[..., str]:
@@ -108,7 +131,8 @@ 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, enable_thinking:bool=False,
+             preserve_thinking:bool=False) -> 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"])
@@ -130,15 +154,16 @@ def main():
   parser.add_argument("--max_context", type=int, default=4096, help="Max Context Length")
   parser.add_argument("--serve", nargs='?', type=int, const=8000, metavar="PORT", help="Run OpenAI compatible API (optional port, default 8000)")
   parser.add_argument("--warmup", action="store_true", help="warmup the JIT")
+  parser.add_argument("--beam", type=int, help="Kernel optimization beam width")
   parser.add_argument("--benchmark", nargs='?', type=int, const=20, metavar="COUNT", help="Benchmark tok/s (optional count, default 20)")
   args = parser.parse_args()
 
   # load the model
-  model, kv = Transformer.from_gguf(fetch(models.get(args.model, args.model)), args.max_context)
+  model_path = fetch(models.get(args.model, args.model))
+  model, kv = Transformer.from_gguf(model_path, args.max_context)
   model_name = kv.get('general.name') or kv.get('general.basename') or args.model
-  file_sizes = [y.nbytes() for y in UOp.sink(*[x.uop for x in nn.state.get_parameters(model)]).toposort() if y.op is Ops.BUFFER]
-  print(f"using model \"{model_name}\" with {sum(file_sizes):,} bytes and {sum(x.numel() for x in nn.state.get_parameters(model)):,} params, "
-        f"max context {args.max_context} on {nn.state.get_parameters(model)[0].device}")
+  print(f"using model \"{model_name}\" with {model_path.stat().st_size:,} bytes and {model.parameter_count:,} params, "
+        f"max context {model.max_context} on {model.token_embd.weight.device}")
 
   # get tokenizer
   tok = SimpleTokenizer.from_gguf_kv(kv)
@@ -159,9 +184,12 @@ 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])))
+    amd_server = bool(args.serve) and str(model.token_embd.weight.device).startswith("AMD")
+    beam = args.beam if args.beam is not None else 2 if amd_server else BEAM.value
+    print(f"warming serving JITs with BEAM={beam}")
+    batch_size = 448 if amd_server else JIT_BATCH_SIZE.value
+    with Context(DEBUG=DEBUG.value, BEAM=beam, JIT_BATCH_SIZE=batch_size):
+      model.warmup()
 
   # start server
   if args.serve: LLMServer(('', args.serve), model, model_name, tok, template).serve_forever()

tinygrad/llm/gguf.py

@@ -1,7 +1,8 @@
-import functools, io, pathlib, re, struct
+import functools, io, pathlib, re, struct, weakref
 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.nn.state import TensorIO
@@ -20,7 +21,14 @@ _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()
+
+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.
 
@@ -35,14 +43,14 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
   if (dtype := _GGML_NATIVE.get(ggml_type)) is not None:
     return t[:dtype.itemsize * n].contiguous().bitcast(dtype)
 
-  def q_to_uint8(t: Tensor, b: int) -> Tensor:
-    # TODO: rewrite with arange?
-    shift_tensor, bitmask = Tensor.stack(*[ Tensor(2**(i*b), device=t.device, dtype=t.dtype) for i in range(8//b) ]), 0xff >> (8 - b)
-    return t.unsqueeze(-1).expand((*t.shape,8//b)).div(shift_tensor, rounding_mode="trunc").bitwise_and(bitmask).transpose(-1, -2).flatten(-2)
+  def q_to_uint8(t:Tensor, b:int) -> Tensor:
+    shift_tensor, bitmask = Tensor.stack(*[Tensor(2**(i*b), device=t.device, dtype=t.dtype) for i in range(8//b)]), 0xff >> (8-b)
+    return t.unsqueeze(-1).expand((*t.shape, 8//b)).div(shift_tensor, rounding_mode="trunc").bitwise_and(bitmask).transpose(-1, -2).flatten(-2)
 
   if (nelements_nbytes := _GGML_QUANT.get(ggml_type)) is not None:
     from tinygrad.runtime.autogen import ggml_common as _ggml
-    blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1])).contiguous()
+    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 ])
@@ -146,7 +154,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]:

tinygrad/llm/serve.py

@@ -17,9 +17,14 @@ def parse_tool_call(s:str) -> tuple[str, typing.Any]|None:
   # XML format: <function=name>\n<parameter=key>\nvalue\n</parameter>...</function>
   if (fm := re.match(r"<function=([^>]+)>\s*(.*?)\s*(?:</function>)?$", s, re.DOTALL)):
     args = {}
-    for pm in re.finditer(r"<parameter=([^>]+)>\s*(.*?)\s*</parameter>", fm.group(2), re.DOTALL):
-      try: args[pm.group(1)] = json.loads(pm.group(2))
-      except json.JSONDecodeError: args[pm.group(1)] = pm.group(2)
+    for pm in re.finditer(r"<parameter=([^>]+)>(.*?)</parameter>", fm.group(2), re.DOTALL):
+      value = pm.group(2)
+      if value.startswith("\r\n"): value = value[2:]
+      elif value.startswith("\n"): value = value[1:]
+      if value.endswith("\r\n"): value = value[:-2]
+      elif value.endswith("\n"): value = value[:-1]
+      try: args[pm.group(1)] = json.loads(value)
+      except json.JSONDecodeError: args[pm.group(1)] = value
     return fm.group(1), args
   return None
 
@@ -33,9 +38,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, thinking:bool=False):
     self.buf = ""
-    self.mode = "undecided"  # output inside a think block is sent as reasoning_content
+    self.mode = "reasoning" if thinking 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:
@@ -65,11 +70,11 @@ 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, thinking: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"{self.path}  {colored('--', 'BLACK')}  "
-               f"in:{colored(f'{cache_start_pos:5d}', 'green')} +{len(ids)-cache_start_pos:5d}  {colored('--', 'BLACK')}  ")
+    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":""})
@@ -77,9 +82,9 @@ class Handler(HTTPRequestHandler):
     finish_reason = "stop"
     st = time.perf_counter()
     dec = tok.stream_decoder()
-    router = StreamRouter()
+    router = StreamRouter(thinking)
     for next_id in model.generate(ids, temperature=temperature):
-      if len(out) == 0: stderr_log(f"prefill:{(len(ids)-cache_start_pos)/((pt:=time.perf_counter())-st):4.0f} tok/s  {colored('--', 'BLACK')}  ")
+      if 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})
@@ -101,25 +106,35 @@ class Handler(HTTPRequestHandler):
       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": len(ids), "completion_tokens": len(out), "total_tokens": len(ids) + len(out)}, **tmpl}
+      yield {"choices": [], "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": len(out),
+                                        "total_tokens": prompt_tokens + len(out),
+                                        "prompt_tokens_details":{"cached_tokens":cache_start_pos}}, **tmpl}
     et = time.perf_counter()
     stderr_log(f"gen:{len(out)/(et-pt) if len(out) > 1 else 0:4.0f} tok/s  {colored('--', 'BLACK')}  "
                f"out:{len(out):5d}  {colored('--', 'BLACK')}  total:{et-st:6.2f}s\n")
 
   def do_POST(self):
+    request_st = time.perf_counter()
+    stderr_log(f"{self.path}  {colored('--', 'BLACK')}  ")
     raw_body = self.rfile.read(int(self.headers.get("Content-Length", "0")))
     body: dict[str, typing.Any] = json.loads(raw_body.decode("utf-8"))
     if DEBUG >= 1: print(json.dumps(body, indent=2))
     if self.path == "/v1/chat/completions":
       # render and tokenize
       normalize_messages(body["messages"])
-      rendered = self.server.template.render(messages=body["messages"], tools=body.get("tools"), add_generation_prompt=True)
+      rendered = self.server.template.render(messages=body["messages"], tools=body.get("tools"), add_generation_prompt=True,
+                                             enable_thinking=body.get("enable_thinking", False), preserve_thinking=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:
+        return self.send_data(json.dumps({"error":{"message":f"prompt has {len(ids)} tokens, but the model context is "
+          f"{self.server.model.max_context}", "type":"invalid_request_error", "param":"messages", "code":"context_length_exceeded"}}).encode(),
+          status_code=400)
 
       # 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)), thinking=body.get("enable_thinking", False))
       if body.get("stream"): self.stream_json(chunks)
       else:
         out, reasoning, tool_calls, finish_reason = [], [], [], "stop"

tinygrad/renderer/cstyle.py

@@ -99,7 +99,7 @@ def uops_to_dtypes(uops:list[UOp]) -> list[tuple[DType, int]]:
   return dedup((u.dtype, u.max_numel()) for u in uops if u.addrspace in (AddrSpace.ALU, None) and u.dtype != dtypes.void and u._shape is not None)
 
 def _wmma_name(u:UOp) -> str:
-  return f"WMMA_{'_'.join(map(str, u.arg[0]))}_{u.arg[1].name}_{u.dtype.scalar().name}"
+  return f"WMMA_{'_'.join(map(str, u.arg[0]))}_{u.arg[1].name}_{u.dtype.scalar().name}".replace(" ", "_")
 
 # (name, dims, dtype_in, dtype_out, device, threads, upcast_sizes)
 def wmma_args(uops:list[UOp]):
@@ -565,6 +565,11 @@ class HIPRenderer(CStyleLanguage):
       # #define __WMMA_16_16_16_half_half __builtin_amdgcn_wmma_f16_16x16x16_f16_w32_gfx12
       elif self.tensor_cores == tc.amd_rdna4:
         prefix.append(f"#define __{name} __builtin_amdgcn_wmma_{type_map[dtype_out]}_16x16x16_{type_map[dtype_in]}_w32_gfx12")
+      elif dtype_out == dtypes.int32:
+        prefix.append("typedef int wmma_int4 __attribute__((ext_vector_type(4)));\n"+
+          f"static inline __attribute__((device)) int8 __{name}"+"""(signed_char16 a, signed_char16 b, int8 c) {
+  return __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, __builtin_bit_cast(wmma_int4, a),
+    true, __builtin_bit_cast(wmma_int4, b), c, true);\n}""")
       elif dtype_out == dtypes.float:
         prefix.append(f"#define __{name} __builtin_amdgcn_wmma_f32_16x16x16_{'f16' if dtype_in == dtypes.half else 'bf16'}_w32")
       else: prefix.append(f"static inline __attribute__((device)) half8 __{name}"+"""(half16 a, half16 b, half8 c) {