fixes

* more weak const without cast in const_like [PR]

* that?

.github/workflows/test.yml

@@ -219,8 +219,8 @@ jobs:
       run: python3 test/external/external_benchmark_schedule.py
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
-    - name: Repo line count < 25000 lines
-      run: MAX_LINE_COUNT=25000 python sz.py
+    - name: Repo line count <= 26000 lines
+      run: MAX_LINE_COUNT=26000 python sz.py
 
   spec:
     strategy:

examples/mlperf/model_train.py

@@ -1282,7 +1282,7 @@ def train_bert():
         previous_step = i
 
 def train_llama3():
-  from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8_DTYPE, MXFP8
+  from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8_DTYPE, MXFP8, MXFP4
   from examples.llama3 import MODEL_PARAMS
   from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
   from examples.mlperf.optim import GradAccClipAdamW, clip_grads
@@ -1434,9 +1434,9 @@ def train_llama3():
     load_state_dict(scheduler, safe_load(fn), realize=False)
 
   fp8_amax = [t for ts in model._fp8_amax.values() for t in ts]
-  fp8_next_amax = [t for ts in model._fp8_next_amax.values() for t in ts] if hasattr(model, "_fp8_next_amax") else []
-  fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts] if hasattr(model, "_fp8_grad_amax") else []
-  fp8_next_grad_amax = [t for ts in model._fp8_next_grad_amax.values() for t in ts] if hasattr(model, "_fp8_next_grad_amax") else []
+  fp8_next_amax = [t for ts in model._fp8_next_amax.values() for t in ts]
+  fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts]
+  fp8_next_grad_amax = [t for ts in model._fp8_next_grad_amax.values() for t in ts]
   fp8_inv_scales = list(model._fp8_inv_scale.values()) + list(model._fp8_next_inv_scale.values())
 
   from tinygrad.nn.state import get_state_dict
@@ -1462,8 +1462,7 @@ def train_llama3():
 
   @TinyJit
   def minibatch(tokens:Tensor):
-    for nxt in fp8_next_amax: nxt.assign(0)
-    for nxt in fp8_next_grad_amax: nxt.assign(0)
+    model.reset_amax()
     if is_dp: tokens = tokens.to(None).shard(device, 0)
     if is_mp: tokens = tokens.shard(device)
     if not is_sharding: tokens = tokens.to(None)
@@ -1487,8 +1486,7 @@ def train_llama3():
     scheduler.step()
 
     for g in grads: g.assign(0)
-    for cur, nxt in zip(fp8_amax, fp8_next_amax): cur.assign(nxt)
-    for cur, nxt in zip(fp8_grad_amax, fp8_next_grad_amax): cur.assign(nxt)
+    model.update_amax()
 
     lr_cpu = optim.lr.float().to("CPU")
     grad_norm_cpu = grad_norm.float().to("CPU")
@@ -1579,7 +1577,7 @@ def train_llama3():
 
       mem_gb = GlobalCounters.mem_used / 1e9
       gflops = GlobalCounters.global_ops / 1e9 / dev_time
-      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * 4.6e15)) * 100
+      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * (9.2e15 if MXFP4 else 4.6e15))) * 100
       tqdm.write(
           f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
           f"{lr:.12f} LR, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")

examples/mlperf/models/flat_llama.py

@@ -83,8 +83,8 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca
       return out, x_fp8
   return (x_fp8.dot(w.T, dtype=dtypes.float) * ((amax_x.float() + 1e-8) / FP8_MAX) * w_inv_scale).cast(dtypes.bfloat16), x_fp8
 
-def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
-                         next_amax_x:Tensor, grad_amax_state:Tensor, next_grad_amax_state:Tensor):
+def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor|None,
+                         next_amax_x:Tensor|None, grad_amax_state:Tensor|None, next_grad_amax_state:Tensor|None):
   if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
     from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_rmsnorm_mul_quantize_fp8
     x_fp8, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
@@ -96,8 +96,8 @@ def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, ep
                      next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
   return out, x_normed, rrms, ret
 
-def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
-                             next_amax_x:Tensor, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
+def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor|None,
+                             next_amax_x:Tensor|None, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
   if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
     from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_add_rmsnorm_mul_quantize_fp8
     x_fp8, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
@@ -111,9 +111,9 @@ def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w
   return out, h, x_normed, rrms, ret
 
 def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor,
-                             amax_x2:Tensor, next_amax_x2:Tensor,
-                             grad_amax_xw13:Tensor, next_grad_amax_xw13:Tensor,
-                             grad_amax_xout:Tensor, next_grad_amax_xout:Tensor):
+                             amax_x2:Tensor|None, next_amax_x2:Tensor|None,
+                             grad_amax_xw13:Tensor|None, next_grad_amax_xw13:Tensor|None,
+                             grad_amax_xout:Tensor|None, next_grad_amax_xout:Tensor|None):
   if FUSED_SILU_W13 and not MXFP4:
     from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
     x2_fp8 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13,
@@ -164,14 +164,15 @@ class FlatTransformer:
     self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).clone().is_param_(False)
 
     def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().is_param_(False)
+    n_amax = 0 if MXFP4 else n_layers
     names = ["xqkv", "xo", "x2"]
     names += ["x1", "x3"] if SPLIT_W13 else ["x13"]
-    self._fp8_amax = {name: [_amax() for _ in range(n_layers)] for name in names}
-    self._fp8_next_amax = {name: [_amax() for _ in range(n_layers)] for name in names}
+    self._fp8_amax = {name: [_amax() for _ in range(n_amax)] for name in names}
+    self._fp8_next_amax = {name: [_amax() for _ in range(n_amax)] for name in names}
     grad_names = ["xqkv", "xo", "xout"]
     grad_names += ["xw1", "xw3"] if SPLIT_W13 else ["xw13"]
-    self._fp8_grad_amax = {name: [_amax() for _ in range(n_layers)] for name in grad_names}
-    self._fp8_next_grad_amax = {name: [_amax() for _ in range(n_layers)] for name in grad_names}
+    self._fp8_grad_amax = {name: [_amax() for _ in range(n_amax)] for name in grad_names}
+    self._fp8_next_grad_amax = {name: [_amax() for _ in range(n_amax)] for name in grad_names}
     w_scales = [("wqkv", s_qkv), ("wo", s_o), ("w2", s_2)]
     w_scales += [("w1", s_1), ("w3", s_3)] if SPLIT_W13 else [("w13", s_13)]
     self._fp8_inv_scale = {name: (s if MXFP8 else s.float()).contiguous().is_param_(False) for name, s in w_scales}
@@ -195,9 +196,10 @@ class FlatTransformer:
     return (w * scale_b).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE), inv_scale
 
   def attention(self, x:Tensor, freqs_cis:Tensor, *, attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
-                amax_xqkv:Tensor, amax_xo:Tensor, s_qkv:Tensor, s_o:Tensor,
-                next_amax_xqkv:Tensor, next_amax_xo:Tensor,
-                grad_amax_xqkv:Tensor, grad_amax_xo:Tensor, next_grad_amax_xqkv:Tensor, next_grad_amax_xo:Tensor):
+                amax_xqkv:Tensor|None, amax_xo:Tensor|None, s_qkv:Tensor, s_o:Tensor,
+                next_amax_xqkv:Tensor|None, next_amax_xo:Tensor|None,
+                grad_amax_xqkv:Tensor|None, grad_amax_xo:Tensor|None,
+                next_grad_amax_xqkv:Tensor|None, next_grad_amax_xo:Tensor|None):
     bsz, seqlen, _ = x.shape
     saves = []
 
@@ -319,28 +321,33 @@ class FlatTransformer:
           for i in range(len(amax_dict[name])):
             amax_dict[name][i] = amax_dict[name][i].to(device).contiguous().is_param_(False)
 
+  def reset_amax(self):
+    for st in (self._fp8_next_amax, self._fp8_next_grad_amax):
+      for ts in st.values():
+        for t in ts: t.assign(0)
+
+  def update_amax(self):
+    for cur, nxt in ((self._fp8_amax, self._fp8_next_amax), (self._fp8_grad_amax, self._fp8_next_grad_amax)):
+      for name in cur:
+        for c, n in zip(cur[name], nxt[name]): c.assign(n)
+
   def __call__(self, tokens:Tensor, save:bool=True):
     h = self.tok_embeddings(tokens)
     freqs_cis = self.freqs_cis.cast(h.dtype)
     if not getenv("HK_FLASH_ATTENTION"): freqs_cis = freqs_cis[:, :tokens.shape[1], :, :, :]
     a, na, ga, nga, s = self._fp8_amax, self._fp8_next_amax, self._fp8_grad_amax, self._fp8_next_grad_amax, self._fp8_inv_scale
+    def amax_kwargs(i:int, act_names:tuple[str, ...], grad_names:tuple[str, ...]) -> dict[str, Tensor|None]:
+      specs = (("amax_", a, act_names), ("next_amax_", na, act_names), ("grad_amax_", ga, grad_names), ("next_grad_amax_", nga, grad_names))
+      if MXFP4: return dict.fromkeys(f"{prefix}{name}" for prefix, _, names in specs for name in names)
+      return {f"{prefix}{name}":val[name][i] for prefix, val, names in specs for name in names}
     for i in range(self.n_layers):
-      attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i],
-                         amax_xqkv=a["xqkv"][i], amax_xo=a["xo"][i], s_qkv=s["wqkv"][i], s_o=s["wo"][i],
-                         next_amax_xqkv=na["xqkv"][i], next_amax_xo=na["xo"][i],
-                         grad_amax_xqkv=ga["xqkv"][i], grad_amax_xo=ga["xo"][i],
-                         next_grad_amax_xqkv=nga["xqkv"][i], next_grad_amax_xo=nga["xo"][i])
-      ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], w2=self.w2[i],
-                        amax_x2=a["x2"][i], s_2=s["w2"][i], grad_amax_xout=ga["xout"][i], next_grad_amax_xout=nga["xout"][i],
-                        next_amax_x2=na["x2"][i])
+      attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i], s_qkv=s["wqkv"][i], s_o=s["wo"][i],
+                         **amax_kwargs(i, ("xqkv", "xo"), ("xqkv", "xo")))
+      ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], w2=self.w2[i], s_2=s["w2"][i], **amax_kwargs(i, ("x2",), ("xout",)))
       if SPLIT_W13:
-        ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i], amax_x1=a["x1"][i], amax_x3=a["x3"][i],
-                          next_amax_x1=na["x1"][i], next_amax_x3=na["x3"][i],
-                          s_1=s["w1"][i], s_3=s["w3"][i], grad_amax_xw1=ga["xw1"][i], grad_amax_xw3=ga["xw3"][i],
-                          next_grad_amax_xw1=nga["xw1"][i], next_grad_amax_xw3=nga["xw3"][i])
+        ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i], s_1=s["w1"][i], s_3=s["w3"][i], **amax_kwargs(i, ("x1", "x3"), ("xw1", "xw3")))
       else:
-        ffn_kwargs.update(w13=self.w13[i], amax_x13=a["x13"][i], s_13=s["w13"][i], grad_amax_xw13=ga["xw13"][i],
-                          next_grad_amax_xw13=nga["xw13"][i], next_amax_x13=na["x13"][i])
+        ffn_kwargs.update(w13=self.w13[i], s_13=s["w13"][i], **amax_kwargs(i, ("x13",), ("xw13",)))
       h, *_ = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)
 
     logits = matmul(self.norm(h), self.output[0], fp8=False)[0]
@@ -424,9 +431,7 @@ if __name__ == "__main__":
   @TinyJit
   def fwd_bwd(tokens:Tensor):
     with Timing("python forward: "):
-      for amax_dict in (model._fp8_next_amax, model._fp8_next_grad_amax):
-        for ts in amax_dict.values():
-          for nxt in ts: nxt.assign(0)
+      model.reset_amax()
       logits = model(tokens[:, :-1], save=llama_size=="8B")
       loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
     with Timing("python backward: "):

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/gpt_oss/implementations/tinybox_8xMI350X/dev_beam.sh

@@ -10,6 +10,7 @@ export DEVICE_IN_FUNCTION_BUG=1
 
 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
 export ALL2ALL=${ALL2ALL:-1}
 export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/gpt_oss/implementations/tinybox_8xMI350X/dev_run.sh

@@ -10,6 +10,7 @@ export DEVICE_IN_FUNCTION_BUG=1
 
 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
 export ALL2ALL=${ALL2ALL:-1}
 export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}

extra/benchmark_qwen_llm.py

@@ -0,0 +1,56 @@
+import argparse, time
+from tinygrad.llm.cli import models
+from tinygrad.llm.model import Transformer
+from tinygrad.helpers import fetch
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--model", default="qwen3.5:0.8b", help=f"Model choice ({', '.join(models.keys())}) or path to a local GGUF file")
+  parser.add_argument("--max-context", type=int, default=32768)
+  parser.add_argument("--prompt-tokens", type=int, default=3072)
+  parser.add_argument("--decode-tokens", type=int, default=16)
+  parser.add_argument("--chunk-size", type=int, default=256)
+  parser.add_argument("--expect-output", type=int, nargs="+", default=None, help="expected output tokens to assert on")
+  parser.add_argument("--skip-resume-check", action="store_true")
+  args = parser.parse_args()
+
+  startup_st = st = time.perf_counter()
+  model, _ = Transformer.from_gguf(fetch(models.get(args.model, args.model)), args.max_context)
+  print(f"load {time.perf_counter()-st:.3f}s", flush=True)
+  st = time.perf_counter()
+  model.warmup()
+  # warm up the chunked prefill JIT too, then invalidate the prompt cache (forces the state reset path)
+  if model.has_recurrent_block:
+    for _ in range(2):
+      warm = model.generate([0]*args.chunk_size, chunk_size=args.chunk_size)
+      next(warm), next(warm)
+      model._cached_tokens = [-1]
+  print(f"warm {time.perf_counter()-st:.3f}s", flush=True)
+  states = [getattr(block, name) for block in model.blk for name in ("cache_kv", "cache_k", "cache_v", "conv_state", "recurrent_state")
+            if hasattr(block, name)]
+  device = str(model.token_embd.weight.device)
+  assert all(str(state.device) == device and state.uop.is_realized for state in states)
+  assert all(block.cache_kv.shape[3] >= args.max_context for block in model.blk if hasattr(block, "cache_kv"))
+  assert model.prefill_jit.cnt >= 2 and model.rollout_jit.cnt >= 2
+  print(f"preallocated {sum(state.nbytes() for state in states)/2**30:.3f} GiB state on {device}", flush=True)
+
+  prompt = [257] + [1000+i%1000 for i in range(args.prompt_tokens-1)]
+  gen, st = model.generate(prompt, chunk_size=args.chunk_size), time.perf_counter()
+  output = [next(gen)]
+  pt = time.perf_counter()
+  print(f"prefill {args.prompt_tokens/(pt-st):.3f} tok/s", flush=True)
+  for _ in range(args.decode_tokens): output.append(next(gen))
+  print(f"decode {args.decode_tokens/(time.perf_counter()-pt):.3f} tok/s output {output}", flush=True)
+  if args.expect_output is not None: assert output == args.expect_output, f"expected {args.expect_output}, got {output}"
+
+  if not args.skip_resume_check:
+    follow = model._cached_tokens + [1234+i for i in range(8)]
+    full_prompt = list(follow)
+    resume_pos, gen, st = model.get_start_pos(follow), model.generate(follow, chunk_size=args.chunk_size), time.perf_counter()
+    resumed_token = next(gen)
+    print(f"resume {len(follow)-1} tokens from {resume_pos} in {time.perf_counter()-st:.3f}s token {resumed_token}", flush=True)
+    model._cached_tokens = [-1]
+    st = time.perf_counter()
+    full_token = next(model.generate(full_prompt, chunk_size=args.chunk_size))
+    print(f"full {time.perf_counter()-st:.3f}s token {full_token} match {resumed_token == full_token}", flush=True)
+    assert resumed_token == full_token

extra/gemm/amd_flash_attention.py

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

extra/gemm/moe_routing.py

@@ -5,7 +5,7 @@ BLOCK_ROW = 256
 
 def _sharded_invalids(shape:tuple[int, ...], dtype, device) -> Tensor:
   if isinstance(device, tuple):
-    return Tensor(Tensor.invalids(shape[0] // len(device), *shape[1:], dtype=dtype, device=device).uop.multi(0), device=device)
+    return Tensor.invalids(*shape, dtype=dtype, device=device[0]).shard(device, axis=0)
   return Tensor.invalids(*shape, dtype=dtype, device=device)
 
 def _atomic_add(device:str) -> str:

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_patches
+from tinygrad.runtime.support.hcq2 import make_binary_patch
 from tinygrad.uop.ops import sint, UOp
 from tinygrad.device import Compiled, BufferSpec, Buffer, Device
 from tinygrad.dtype import dtypes
@@ -152,19 +152,12 @@ def pm4_submit(ctx, lin):
   ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COMPUTE:0_{name}")
     for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
 
-  # two tail dwords coordinate safe IB reuse: GPU completions and host submits
-  size_dw = sum(len(ins.src) for ins in lin.src) + len(release_mem(ctx, 0, 0).src)
+  # the host fence at the start of the batch guarantees the ib is free to reuse
+  size_dw = sum(len(ins.src) for ins in lin.src)
   assert size_dw < (1 << 20), f"indirect buffer of {size_dw} dwords doesn't fit one packet"
 
-  ib = UOp.placeholder((size_dw + 2,), dtypes.uint32, next(UOp.unique_num), device=devs, volatile=True).rtag("cmdbuf")
-  done_idx, submit_idx = UOp.const(size_dw + 0, dtypes.int), UOp.const(size_dw + 1, dtypes.int)
-  init_counters = make_patches(ib, [((size_dw + i) * 4, UOp.const(0, dtypes.uint32)) for i in range(2)]).rtag("link")
-  submitted = (counter:=ib.after(init_counters).index(submit_idx)).load()
-  completed = ib.after(loop:=UOp.loop(0)).index(done_idx).load()
-  ib_free = completed.end(loop, completed != submitted)
-
-  bump_fence = pm4_store(ctx, UOp(Ops.SLICE, dtypes.uint32, (ib, UOp.const(size_dw)), 2), (submitted + 1).cast(dtypes.uint64))
-  cmdbuf = make_cmdbuf(lin.replace(src=lin.src + (bump_fence,)), devs, buf=ib, dep=ib_free)
+  ib = UOp.placeholder((size_dw,), dtypes.uint32, next(UOp.unique_num), device=devs, volatile=True).rtag("cmdbuf")
+  cmdbuf = make_cmdbuf(lin, devs, buf=ib)
 
   # the ring itself only carries a packet pointing at the ib, wrapping the ring
   put = put_ptr.index(zero:=UOp.const(0, dtypes.int))
@@ -174,7 +167,7 @@ def pm4_submit(ctx, lin):
   # advance the put/write pointers past the packet
   bump_put_ptr = put_ptr.index(zero).store(put + len(pkt))
   bump_wptr = wptr.index(zero).store(put + len(pkt))
-  flush = UOp.barrier(write_pkt, bump_put_ptr, bump_wptr, counter.store(submitted + 1))
+  flush = UOp.barrier(write_pkt, bump_put_ptr, bump_wptr)
   return doorbell.after(flush).index(zero).store(put + len(pkt))
 
 pm_pm4_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"), pm4_submit)])
@@ -518,8 +511,7 @@ class PCIIface(PCIIfaceBase):
       cq = d.compute_queue
       for b in (cq.put_value, cq.read_ptr, cq.write_ptr): b._buf.view.view(fmt='Q')[0] = 0
       d.iface.dev_impl.gfx.setup_ring(*cq.params)
-      d.timeline_signal('COMPUTE:0')._buf.cpu_view().mv.cast('Q')[0] = \
-        d.timeline_value('COMPUTE:0').as_memoryview(force_zero_copy=True).cast('Q')[0] - 1
+      d.signal('timeline')._buf.cpu_view().mv.cast('Q')[0] = d.signal('value', 1).as_memoryview(force_zero_copy=True).cast('Q')[0] - 1
 
   def sleep(self, timeout):
     if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
@@ -639,9 +631,6 @@ class AMDDevice(HCQ2Compiled):
     qname = f"{'COPY' if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA else 'COMPUTE'}:{idx}"
     self.pm_bufferize = PatternMatcher([
       (UPat(Ops.PARAM, tag=f"{qname}_{name}"), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"]
-    ] + [
-      (UPat(Ops.PARAM, tag=f"{qname}_timeline_signal"), lambda ctx, q=qname: ctx[0].timeline_signal(q)),
-      (UPat(Ops.PARAM, tag=f"{qname}_timeline_value"), lambda ctx, q=qname: ctx[0].timeline_value(q)),
     ]) + self.pm_bufferize
 
     return queue

extra/thunder/amd/fa_bwd_causal.cpp

@@ -533,6 +533,8 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
         if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
         if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
         mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
+        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
+        if constexpr (D == 64) asm volatile("s_nop 15");
         if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
         // Load K_j from shared memory to registers
         // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -791,6 +793,8 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
         if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
         if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
         mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
+        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
+        if constexpr (D == 64) asm volatile("s_nop 15");
         if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
         // Load K_j from shared memory to registers
         // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -1048,6 +1052,8 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
         if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
         if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
         mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
+        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
+        if constexpr (D == 64) asm volatile("s_nop 15");
         if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
         // Load K_j from shared memory to registers
         // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -1303,6 +1309,8 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
         if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}), Q_i_addr);
         if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}), Q_i_addr);
         mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
+        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
+        if constexpr (D == 64) asm volatile("s_nop 15");
         if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
         // Load K_j from shared memory to registers
         // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -1582,6 +1590,8 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
         if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
         if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
         mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
+        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
+        if constexpr (D == 64) asm volatile("s_nop 15");
         if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
         // Load K_j from shared memory to registers
         // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -1842,6 +1852,8 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
         if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
         if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
         mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
+        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
+        if constexpr (D == 64) asm volatile("s_nop 15");
         if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
         // Load K_j from shared memory to registers
         // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -2099,6 +2111,8 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
         if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
         if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
         mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
+        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
+        if constexpr (D == 64) asm volatile("s_nop 15");
         if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
         // Load K_j from shared memory to registers
         // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -2354,6 +2368,8 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
         if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}), Q_i_addr);
         if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}), Q_i_addr);
         mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
+        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
+        if constexpr (D == 64) asm volatile("s_nop 15");
         if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
         // Load K_j from shared memory to registers
         // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));

test/backend/test_custom_kernel.py

@@ -434,6 +434,16 @@ class TestCustomKernel(unittest.TestCase):
     a = Tensor.custom_kernel(a.reshape(2, 2).T, fxn=custom_src_kernel)[0]
     self.assertEqual(a.tolist(), [[1, 2], [1, 3]])
 
+  def test_inplace_transpose(self):
+    def custom_assign_row_max_kernel(A:UOp) -> UOp:
+      row = UOp.range(A.shape[0], 0)
+      col = UOp.range(A.shape[1], 1)
+      return A[row, col].store(A[row].max(axis=0)).end(col).end(row).sink(arg=KernelInfo(name=f"assign_row_max_{A.numel()}"))
+    a = Tensor.arange(4).clone().realize()
+    a = Tensor.custom_kernel(a.reshape(2, 2).T, fxn=custom_assign_row_max_kernel)[0]
+    self.assertEqual(a.flatten().tolist(), [2, 2, 3, 3])
+    self.assertEqual(a.shape, (2, 2))
+
 class TestCustomKernelInput(unittest.TestCase):
   def _test_mop(self, mop_fxn, max_kernels):
     # default: input is BUFFER

test/device/test_hcq.py

@@ -330,6 +330,7 @@ class TestHCQ(unittest.TestCase):
   # Test profile api
   def test_speed_exec_time(self):
     sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal()
+    st = time.perf_counter()
     TestHCQ.d0.hw_compute_queue_t().timestamp(sig_st) \
                                    .exec(TestHCQ.runtime, TestHCQ.kernargs_ba_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size) \
                                    .timestamp(sig_en) \
@@ -337,11 +338,13 @@ class TestHCQ(unittest.TestCase):
 
     TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
+    host_us = (time.perf_counter() - st) * 1e6
 
     et = float(sig_en.timestamp - sig_st.timestamp)
 
     print(f"exec kernel time: {et:.2f} us")
-    assert 0.1 <= et <= (3000000 if MOCKGPU or Device.DEFAULT in {"CPU"} else 100)
+    # emulated devices are only bounded by the host window around submit+wait
+    assert 0.1 <= et <= (host_us if MOCKGPU or Device.DEFAULT in {"CPU"} else 100)
 
   def test_speed_copy_bandwidth(self):
     if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
@@ -352,6 +355,7 @@ class TestHCQ(unittest.TestCase):
     b = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
 
     sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal()
+    st = time.perf_counter()
     TestHCQ.d0.hw_copy_queue_t().timestamp(sig_st) \
                                 .copy(a._buf, b._buf, SZ) \
                                 .timestamp(sig_en) \
@@ -359,13 +363,14 @@ class TestHCQ(unittest.TestCase):
 
     TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
+    host_ms = (time.perf_counter() - st) * 1e3
 
-    et = float(sig_en.timestamp - sig_st.timestamp)
-    et_ms = et / 1e3
+    et_ms = float(sig_en.timestamp - sig_st.timestamp) / 1e3
+    assert 0 < et_ms <= host_ms # timestamps are in us and cover only the copy
 
     gb_s = ((SZ / 1e9) / et_ms) * 1e3
     print(f"same device copy:  {et_ms:.2f} ms, {gb_s:.2f} GB/s")
-    assert (0.2 if MOCKGPU else 10) <= gb_s <= 1000
+    assert (0 if MOCKGPU else 10) <= gb_s <= 1000
 
   def test_speed_cross_device_copy_bandwidth(self):
     if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
@@ -379,6 +384,7 @@ class TestHCQ(unittest.TestCase):
     TestHCQ.d0.allocator._map(b._buf)
 
     sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal()
+    st = time.perf_counter()
     TestHCQ.d0.hw_copy_queue_t().timestamp(sig_st) \
                                 .copy(a._buf, b._buf, SZ) \
                                 .timestamp(sig_en) \
@@ -386,13 +392,14 @@ class TestHCQ(unittest.TestCase):
 
     TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
+    host_ms = (time.perf_counter() - st) * 1e3
 
-    et = float(sig_en.timestamp - sig_st.timestamp)
-    et_ms = et / 1e3
+    et_ms = float(sig_en.timestamp - sig_st.timestamp) / 1e3
+    assert 0 < et_ms <= host_ms # timestamps are in us and cover only the copy
 
     gb_s = ((SZ / 1e9) / et_ms) * 1e3
     print(f"cross device copy: {et_ms:.2f} ms, {gb_s:.2f} GB/s")
-    assert (0.2 if MOCKGPU else 2) <= gb_s <= 100
+    assert (0 if MOCKGPU else 2) <= gb_s <= 100
 
   def test_timeline_signal_rollover(self):
     for queue_type in [TestHCQ.d0.hw_compute_queue_t, TestHCQ.d0.hw_copy_queue_t]:

test/null/test_graph_rewrite.py

@@ -307,8 +307,8 @@ class TestRecurse(unittest.TestCase):
   def test_inf_loop(self):
     a = UOp.const(3)
     pm = PatternMatcher([
-      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
-      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)),
     ])
     with self.assertRaises(RuntimeError):
       graph_rewrite(a, pm)
@@ -316,8 +316,8 @@ class TestRecurse(unittest.TestCase):
   def test_inf_loop_bottom_up(self):
     a = UOp.const(3)
     pm = PatternMatcher([
-      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
-      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)),
     ])
     with self.assertRaises(RuntimeError):
       graph_rewrite(a, pm, bottom_up=True)
@@ -378,8 +378,8 @@ class TestWalkRewrite(unittest.TestCase):
     """A bouncing pattern applies once and stops instead of looping."""
     a = UOp.const(3)
     pm = PatternMatcher([
-      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
-      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)),
     ])
     with self.assertRaises(RuntimeError):
       graph_rewrite(a, pm, bottom_up=True)
@@ -456,8 +456,8 @@ class TestWalkRewrite(unittest.TestCase):
     """Bottom-up walk also applies once per node, no fixed-point iteration."""
     a = UOp.const(3)
     pm = PatternMatcher([
-      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
-      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)),
     ])
     ret = graph_rewrite(a, pm, bottom_up=True, walk=True)
     self.assertIs(ret, UOp.const(4))
@@ -511,7 +511,7 @@ class TestWalkRewrite(unittest.TestCase):
     def bpm_match(ctx, x):
       ctx.append((x.val if x.op is Ops.CONST else x.op, "bpm"))
       # rewrite const(1) -> const(10), short-circuiting its subtree
-      if x.op is Ops.CONST and x.val == 1: return x.replace(arg=10)
+      if x.op is Ops.CONST and x.val == 1: return UOp.const(10, x.dtype)
       return None
     def pm_match(ctx, x):
       ctx.append((x.val if x.op is Ops.CONST else x.op, "pm"))

test/null/test_llm_server.py

@@ -109,6 +109,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.llm.serve import Handler
+    source, handler = Mock(), Mock()
+    source.__iter__ = Mock(return_value=iter([{}]))
+    handler.wfile.write.side_effect = BrokenPipeError
+    Handler.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_uop_graph.py

@@ -593,7 +593,7 @@ class TestUOpTags(unittest.TestCase):
   def test_inc_by_one(self):
     g = UOp.const(1) + UOp.const(1)
     assert g.ssimplify() == 2
-    pm_plus_1 = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.replace(arg=x.val+1, tag=1) if x.tag is None else None)])
+    pm_plus_1 = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: UOp.const(x.val+1, x.dtype).rtag(1) if x.tag is None else None)])
     pm_strip_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
     g = graph_rewrite(g, pm_plus_1)
     assert g.ssimplify() == 4

test/null/test_uops.py

@@ -126,6 +126,12 @@ class TestConstFloatEq(unittest.TestCase):
     self.assertFalse(nan == Invalid)
     self.assertTrue(nan != Invalid)    # __ne__ must defer to the reflected eq, not swallow NotImplemented
 
+  def test_invalid_eq_defers_to_reflected(self):
+    class HoldsInvalid:  # a carrier that knows it holds Invalid. returning False for foreign types would silence its eq
+      def __eq__(self, other): return other is Invalid
+    self.assertTrue(Invalid == HoldsInvalid())
+    self.assertFalse(Invalid != HoldsInvalid())
+
   def test_matchers_agree_on_nan(self):
     n = UOp.const(math.nan, dtypes.float32)
     for compiled in (False, True):
@@ -447,7 +453,7 @@ class TestUPatHelpers(unittest.TestCase):
 
 class TestUopsObject(unittest.TestCase):
   def test_timing(self):
-    with Timing("create 10k uops:"): ret = [UOp(Ops.CONST, dtypes.int, arg=10000000+i) for i in range(10000)]
+    with Timing("create 10k uops:"): ret = [UOp.const(10000000+i, dtypes.int) for i in range(10000)]
     assert len(ret) == 10000
 
   def test_nested(self):

test/null/test_uops_stats.py

@@ -147,21 +147,21 @@ class TestUOpsStats(unittest.TestCase):
   #MULACC should have the same stats as MUL + ADD
   def test_mulacc(self):
     globl = UOp.param(0, dtypes.int, (3,))
-    o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1)
-    o2 = UOp(Ops.CONST, dtypes.int, tuple(), 2)
+    o1 = UOp.const(1, dtypes.int)
+    o2 = UOp.const(2, dtypes.int)
     u1 = globl.index(o1)
     u2 = globl.index(o2)
-    u3 = UOp(Ops.CONST, dtypes.int, tuple(), 3)
+    u3 = UOp.const(3, dtypes.int)
     u4 = UOp(Ops.MUL, src=(u1,u2))
     u5 = UOp(Ops.ADD, src=(u4,u3))
     uops = tuple(u5.toposort())
 
     globl = UOp.param(0, dtypes.int, (3,))
-    o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1)
-    o2 = UOp(Ops.CONST, dtypes.int, tuple(), 2)
+    o1 = UOp.const(1, dtypes.int)
+    o2 = UOp.const(2, dtypes.int)
     u1 = globl.index(o1)
     u2 = globl.index(o2)
-    u3 = UOp(Ops.CONST, dtypes.int, tuple(), 3)
+    u3 = UOp.const(3, dtypes.int)
     u4 = UOp(Ops.MULACC, src=(u1,u2,u3))
     uops_fma = tuple(u4.toposort())
 

test/null/test_viz.py

@@ -97,7 +97,7 @@ class TestViz(unittest.TestCase):
     # VIZ tracks rewrites up to and including the error
     def count_3(x:UOp):
       assert x.val <= 3
-      return x.replace(arg=x.val+1)
+      return UOp.const(x.val+1, x.dtype)
     err_pm = PatternMatcher([(UPat.cvar("x"), count_3),])
     a = UOp.const(1)
     with save_viz() as viz:
@@ -202,8 +202,8 @@ class TestViz(unittest.TestCase):
     a = UOp.const(3)
     b = UOp.const(4)
     pm = PatternMatcher([
-      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
-      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)),
     ])
     with save_viz() as viz:
       # use smaller stack limit for faster test (default is 250000)
@@ -224,7 +224,7 @@ class TestViz(unittest.TestCase):
     list(viz.get_details(0, 0))
 
   def test_enter_calls_rewrite(self):
-    pm = PatternMatcher([(UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4))])
+    pm = PatternMatcher([(UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype))])
     with save_viz() as viz:
       inner = UOp.const(3)
       call = UOp(Ops.CALL, src=(UOp(Ops.SINK, src=(inner,)),))

test/test_tiny.py

@@ -39,6 +39,10 @@ class TestTiny(unittest.TestCase):
     out = Tensor.ones(N).contiguous().sum()
     self.assertEqual(out.item(), N)
 
+  def test_eye(self):
+    out = Tensor.eye(3).flatten()
+    self.assertListEqual(out.tolist(), [1.0,0.0,0.0, 0.0,1.0,0.0, 0.0,0.0,1.0])
+
   def test_gemm(self, N=getenv("GEMM_N", 64), dtype=dtypes.float):
     a = Tensor.ones(N,N, dtype=dtype).contiguous()
     b = Tensor.eye(N, dtype=dtype).clone()

test/unit/test_allreduce.py

@@ -1,5 +1,5 @@
 import unittest
-from tinygrad import Tensor, dtypes
+from tinygrad import Tensor, UOp, dtypes
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import Ops
 
@@ -43,6 +43,13 @@ class TestRingAllReduce(unittest.TestCase):
     self.assertEqual(len(sinks), 2)
     self.assertTrue(all(dst != src for dst, src in pairs))
 
+  def test_symbolic_shape(self):
+    rows = UOp.variable("rows", 1, 4).bind(3)
+    t = Tensor.ones(4, 4).shard(("CPU:0", "CPU:1"), axis=1).realize()
+    out = t[:rows].sum(1).realize()
+    self.assertEqual(out.shape, (rows,))
+    self.assertTrue((out == 4).all().item())
+
   def test_correct_ring(self):
     with Context(RING=2):
       N = 4

test/unit/test_assign.py

@@ -697,6 +697,40 @@ class TestAssignOrdering(unittest.TestCase):
   - Race conditions (concurrent access to same buffer)
   """
 
+  def test_packed_state_write_not_reordered_before_readers(self):
+    """A store to buffer B packed in another buffer's AFTER (rec.after(B.store(v))) must not be
+    reordered before producer kernels that read B. The store is tracked under the AFTER's base buffer
+    (rec), so without resolving the actual store targets the scheduler generates no WAR dependency for
+    B's readers and the write-back can run first, corrupting the producers' input."""
+    from tinygrad.llm.model import _gated_delta_prefill_kernel
+    def build(pre_realize:bool) -> np.ndarray:
+      def Tl(a, b, shape): return Tensor.linspace(a, b, int(np.prod(shape)), dtype=dtypes.float32).reshape(*shape)
+      x = Tensor.linspace(-1.0, 1.0, 12, dtype=dtypes.float32).reshape(1, 3, 4)
+      xh = (x / x.square().mean(-1, keepdim=True).sqrt()).half()
+      qkv = xh @ Tl(-0.15, 0.2, (6, 4)).T
+      conv_state = Tensor.zeros(1, 1, 6).clone()   # read by the conv below
+      rec_state = Tensor.zeros(1, 1, 2, 2).clone()
+      window = conv_state.cat(qkv, dim=1)
+      T = 3
+      w_conv = Tl(-0.05, 0.05, (6, 2))
+      conv_out = (window[:, 0:T] * w_conv[:, 0] + window[:, 1:T+1] * w_conv[:, 1]).silu()
+      q, k, v = conv_out.split([2, 2, 2], dim=-1)
+      q = q.reshape(1, T, 1, 2).normalize(dim=-1)
+      k = k.reshape(1, T, 1, 2).normalize(dim=-1)
+      v = v.reshape(1, T, 1, 2)
+      beta, alpha = Tensor.zeros(1, T, 1) + 0.5, Tensor.zeros(1, T, 1) + 0.9
+      q, k, v, beta = [z.transpose(1, 2).float() for z in (q, k, v, beta)]
+      alpha = alpha.transpose(1, 2).float().exp()
+      qs, kq = q * 2**-0.5, ((q * 2**-0.5)*k).sum(-1).contiguous()
+      # conv_state write-back packed into rec_state's AFTER, custom kernel consumes it
+      new_conv_state = window[:, T:T+1].contiguous()
+      state = Tensor(rec_state.uop.after(conv_state.uop.store(new_conv_state.uop)))
+      args = [Tensor.empty_like(v), qs, k, v, beta, alpha, state, kq]
+      if pre_realize: args = [a.realize() if i != 6 else a for i, a in enumerate(args)]
+      return Tensor.custom_kernel(*args, fxn=_gated_delta_prefill_kernel)[0].transpose(1, 2).realize().numpy()
+    # lazy execution (build the whole graph, then realize) must match eager (inputs realized up front)
+    np.testing.assert_allclose(build(False), build(True), rtol=1e-4, atol=1e-4)
+
   def test_overlapping_slice_assigns(self):
     """Overlapping slice assigns - later write should win for overlapping elements."""
     buf = Tensor.zeros(8).contiguous().realize()

test/unit/test_attention.py

@@ -176,6 +176,28 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
       np.testing.assert_allclose(recurrent_state, expected_recurrent[step], rtol=1e-3, atol=1e-3,
                                  err_msg=f"GatedDeltaNet reset recurrent cache mismatch at step {step}")
 
+  def test_gatedeltanet_prefill_matches_decode(self):
+    # chunked prefill (T>1, custom kernel) must produce the same output and state as sequential decode (T=1).
+    # uses lazy linspace weights, which exercise the scheduler's WAR tracking for the packed conv/recurrent
+    # state write-backs (a write to one buffer packed in another buffer's AFTER must not be reordered before
+    # the producers that read the target buffer)
+    config = self._make_config(max_context=3)
+    block = self._make_block(config)
+    x = Tensor.linspace(-1.0, 1.0, 3 * config.dim, dtype=dtypes.float32).reshape(1, 3, config.dim)
+
+    x_norm = block.attn_norm(x)
+    block._init_state(x_norm)
+    prefill = block._attention(x_norm, 0).realize().numpy()
+    prefill_conv, prefill_recurrent = self._cache_views(block)
+
+    block = self._make_block(config)
+    decode = np.concatenate([self._run_attention(block, x[:, t:t+1], t) for t in range(x.shape[1])], axis=1)
+    decode_conv, decode_recurrent = self._cache_views(block)
+
+    np.testing.assert_allclose(prefill, decode, rtol=1e-3, atol=1e-3, err_msg="prefill output mismatch")
+    np.testing.assert_allclose(prefill_conv, decode_conv, rtol=1e-3, atol=1e-3, err_msg="prefill conv cache mismatch")
+    np.testing.assert_allclose(prefill_recurrent, decode_recurrent, rtol=1e-3, atol=1e-3, err_msg="prefill recurrent cache mismatch")
+
   def test_kda_channel_decay(self):
     config = self._make_config(n_heads=2, ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=2, time_step_rank=2, inner_size=4, kda=True))
     block, x = GatedDeltaNetBlock(config, config.ssm), Tensor([[[1., 2., 0., 0.]]])

test/unit/test_llm_server.py

@@ -3,6 +3,7 @@ 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)
@@ -10,12 +11,44 @@ V_START_POS = UOp.variable("start_pos", 0, TEST_CONFIG.max_context-1)
 V_TOKS = UOp.variable("toks", 1, 32)  # 32 is the default chunk_size in generate
 
 class TestTransformerGenerate(unittest.TestCase):
+  def test_warmup(self):
+    model, calls = Transformer(TEST_CONFIG), []
+    def generate(tokens):
+      calls.append(tokens)
+      yield from (1, 2)
+    with patch.object(model, "generate", generate): model.warmup()
+    self.assertEqual(calls, [[0], [0]])
+
+  def test_first_recurrent_generate_before_state_init(self):
+    model = Transformer(TEST_CONFIG)
+    model.has_recurrent_block = True
+    with patch.object(Transformer, '__call__', return_value=Tensor([[42]])):
+      self.assertEqual(next(model.generate([0])), 42)
+
+  def test_recurrent_live_state_reuse(self):
+    model = Transformer(TEST_CONFIG)
+    model.has_recurrent_block = True
+    model._cached_tokens = [1, 2, 3, 4, 5]
+    self.assertEqual(model.get_start_pos([1, 2, 3, 4, 5, 42, 10]), 5)
+    calls = []
+    def mock_call(self, tokens, start_pos, temperature, **kwargs):
+      calls.append((tokens.shape, start_pos))
+      return Tensor([[42]])
+    with patch.object(Transformer, '__call__', mock_call):
+      next(model.generate([1, 2, 3, 4, 5, 42, 10]))
+    self.assertEqual(calls, [((1, 1), V_START_POS.bind(5)), ((1, 1), V_START_POS.bind(6))])
+
+  def test_template_starts_reasoning(self):
+    router = StreamRouter(reasoning=True)
+    self.assertEqual(list(router.route("reasoning</think>answer")),
+                     [("reasoning_content", "reasoning"), ("content", "answer")])
+
   def test_kv_cache_reuse(self):
     """Test that generate reuses the KV cache when tokens extend the cached prefix."""
     model = Transformer(TEST_CONFIG)
 
     captured_inputs = []
-    def mock_call(self, tokens, start_pos, temperature):
+    def mock_call(self, tokens, start_pos, temperature, **kwargs):
       captured_inputs.append((tokens.shape, start_pos))
       return Tensor([[42]])
 
@@ -40,7 +73,7 @@ class TestTransformerGenerate(unittest.TestCase):
     model = Transformer(TEST_CONFIG)
 
     captured_inputs = []
-    def mock_call(self, tokens, start_pos, temperature):
+    def mock_call(self, tokens, start_pos, temperature, **kwargs):
       captured_inputs.append((tokens.shape, start_pos))
       return Tensor([[42]])
 
@@ -88,7 +121,7 @@ class TestTransformerGenerate(unittest.TestCase):
 
     def get_prefill_flags(tokens, chunk_size):
       is_prefill = []
-      def mock_call(self, tokens, start_pos, temperature):
+      def mock_call(self, tokens, start_pos, temperature, **kwargs):
         is_prefill.append(resolve(tokens.shape[1] != 1))
         return Tensor([[42]])
       with patch.object(Transformer, '__call__', mock_call):
@@ -149,7 +182,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

@@ -125,8 +125,7 @@ def do_devectorize(b:UOp):
   # broadcasting needs to be already unpacked, Invalid matches any dtype and shape
   if not all(x.shape == b.shape or x.base.is_invalid for x in b.src): return None
   src = []
-  for idx in itertools.product(*[range(x) for x in b.shape]):
-    idx_c = [UOp.const(i) for i in idx]
+  for idx_c in itertools.product(*[[UOp.const(i) for i in range(x)] for x in b.shape]):
     src.append(b.replace(dtype=None, src=tuple(x.base if x.base.is_invalid else x.index(*idx_c) for x in b.src)))
   return UOp.stack(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src)
 
@@ -214,7 +213,7 @@ def reduce_ranges_to_acc(ctx:ReduceContext, r:UOp):
   topo = r.src[0].toposort()
   ended_ranges = flatten([x.ended_ranges for x in topo if x.op is Ops.END])
   input_ranges = tuple(x for x in topo if x.op is Ops.RANGE and x not in r.src[1:] and x not in ended_ranges)
-  acc_init = acc.after(*input_ranges).store(identity_element(r.arg[0], r.dtype))
+  acc_init = acc.after(*input_ranges).store(UOp.const(identity_element(r.arg[0], r.dtype)))
   acc_initted = acc.after(acc_init, *r.src[1:])
   inp = r.src[0].reduce(arg=r.arg) if r.arg[1] else r.src[0]
   acc_out = acc_initted.store(acc_initted.alu(r.arg[0], inp)).end(*r.src[1:]).rtag("mergeable")

tinygrad/codegen/decomp/dtype.py

@@ -178,7 +178,9 @@ pm_float_decomp = PatternMatcher([
    f2f(x.bitcast(f2f_dt[ctx[0]]), ctx[0], ctx[1]) if bc.dtype == ctx[0] else None),
   (UPat(Ops.CAST, dtypes.floats, src=(UPat.var("val"),), name="x"), lambda ctx,x,val:
    f2f_clamp(val.cast(ctx[1]), ctx[0]) if x.dtype == ctx[0] else None),
-  (UPat(GroupOp.All-{Ops.BITCAST}, dtypes.floats, name="x"), lambda ctx,x:
+  # a CONST has no srcs to cast, it restates its value at the emulating dtype
+  (UPat(Ops.CONST, dtypes.floats, name="x"), lambda ctx,x: UOp.const(x.val, ctx[1]) if x.dtype == ctx[0] else None),
+  (UPat(GroupOp.All-GroupOp.Defines-{Ops.CAST, Ops.BITCAST, Ops.CONST}, dtypes.floats, name="x"), lambda ctx,x:
    x.replace(dtype=ctx[1], src=tuple(s.cast(ctx[1]) if s.dtype == ctx[0] else s for s in x.src))
    if x.dtype == ctx[0] else None),
   (UPat(Ops.STORE, src=(UPat.var("idx"), UPat(Ops.BITCAST, dtypes.floats, name="val")), name='st'), lambda ctx,st,idx,val:

tinygrad/codegen/decomp/op.py

@@ -47,17 +47,17 @@ def fast_idiv(ren: Renderer, x: UOp, d: int, dont_cast=False) -> UOp|None:
 
 def threefry2x32(x: UOp, key: UOp):
   # split x and key from uint64 to two uint32
-  x0, x1 = (x & 0xffffffff).cast(dtypes.uint32), ((x // 2**32) & 0xffffffff).cast(dtypes.uint32)
-  key0, key1 = (key & 0xffffffff).cast(dtypes.uint32), ((key // 2**32) & 0xffffffff).cast(dtypes.uint32)
+  x0, x1 = x.cast(dtypes.uint32), (x >> 32).cast(dtypes.uint32)
+  key0, key1 = key.cast(dtypes.uint32), (key >> 32).cast(dtypes.uint32)
 
   rotations = [[13, 15, 26, 6], [17, 29, 16, 24]]
   ks = [key1, key0 ^ key1 ^ 0x1BD11BDA, key0]
   xr:list[UOp] = [x0 + ks[-1], x1 + ks[0]]
   for i in range(5):
-    for r in rotations[i % 2]: xr[0], xr[1] = (x0 := xr[0] + xr[1]), x0 ^ ((xr[1] * 2**r) + (xr[1] // 2**(32 - r)))
+    for r in rotations[i % 2]: xr[0], xr[1] = (x0 := xr[0] + xr[1]), x0 ^ ((xr[1] << r) + (xr[1] >> (32 - r)))
     xr = [(xr[0] + ks[i % 3]), (xr[1] + ks[(i + 1) % 3] + i + 1)]
 
-  return xr[1].cast(dtypes.uint64) * 2**32 | xr[0].cast(dtypes.uint64)
+  return (xr[1].cast(dtypes.uint64) << 32) | xr[0].cast(dtypes.uint64)
 
 # ***** decomposition patterns *****
 

tinygrad/codegen/late/regalloc.py

@@ -125,7 +125,7 @@ def regalloc_rewrite(ctx:LinearScanRegallocContext, x:UOp):
   # alloc/dealloc stack
   if ctx.stack_size > 0:
     sp = ctx.ren.stack_pointer()
-    offset = UOp(Ops.CONST, sp.dtype, arg=ctx.stack_size)
+    offset = UOp.const(ctx.stack_size, sp.dtype)
     if i == 0: before = [ctx.ren.isel_matcher.rewrite(UOp(Ops.SUB, src=(sp, offset), tag=sp.tag))] + before
     elif i == len(ctx.uops) - 2: before += [ctx.ren.isel_matcher.rewrite(UOp(Ops.ADD, src=(sp, offset), tag=sp.tag))]
 

tinygrad/dtype.py

@@ -27,7 +27,7 @@ class InvalidType:
   def __new__(cls):
     if cls._instance is None: cls._instance = object.__new__(cls)
     return cls._instance
-  def __eq__(self, other): return self is other
+  def __eq__(self, other): return self is other if isinstance(other, InvalidType) else NotImplemented  # foreign types get the reflected eq
   def __hash__(self): return id(self)
   def __repr__(self): return "Invalid"
   def __reduce__(self): return (InvalidType, ())  # unpickle returns the singleton
@@ -293,6 +293,12 @@ truncate: dict[DType, Callable] = {dtypes.bool: bool,
   **{getattr(dtypes, n): (lambda x, c=getattr(ctypes, f'c_{n}'): c(x).value)
      for n in ('float', 'double', 'int8', 'int16', 'int32', 'int64', 'uint8', 'uint16', 'uint32', 'uint64')}}
 
+def bitcast(x, in_dtype:DType, out_dtype:DType):
+  assert in_dtype.itemsize == out_dtype.itemsize, "bitcast itemsize mismatch"
+  packed = struct.pack(storage_fmt_for_dtype(in_dtype), to_storage_scalar(x, in_dtype))
+  out_val = struct.unpack(storage_fmt_for_dtype(out_dtype), packed)[0]
+  return from_storage_scalar(out_val, out_dtype)
+
 # numpy and torch dtype interop
 
 def _to_np_dtype(dtype:DType) -> type|None:

tinygrad/engine/jit.py

@@ -71,7 +71,7 @@ def jit_lower(linear:UOp, held_bufs:set[UOp], input_uops:list[UOp]) -> UOp:
   # parametrize input buffers: map each input buffer UOp to a PARAM with the correct slot index
   linear = linear.substitute({u: UOp.param(i, u.dtype, u.shape, u.device) for i,u in enumerate(input_uops)}, walk=True)
   linear = memory_plan_rewrite(linear, held_bufs)
-  linear = compile_linear(linear, beam=getenv("JITBEAM", BEAM.value), jit=True)
+  linear = compile_linear(linear, beam=getenv("JITBEAM", BEAM.value))
   if JIT < 2: linear = graph_split_rewrite(linear, max_batch_size=JIT_BATCH_SIZE.value)
   if VIZ: graph_rewrite(linear, PatternMatcher([]), name="View graphed linear")
   return linear
@@ -169,7 +169,7 @@ class CapturedJit(Generic[ReturnType]):
   expected_input_info: list[tuple[UOp, tuple[Variable, ...], DType, str]]  # (view, variables, dtype, device) per input
 
   @functools.cached_property
-  def linear(self) -> UOp: return link_linear(self._linear, jit=True)
+  def linear(self) -> UOp: return link_linear(self._linear)
 
   def __reduce__(self): return self.__class__, (self.ret, self._linear, self.expected_names, self.expected_input_info)
 

tinygrad/engine/realize.py

@@ -265,18 +265,18 @@ pm_exec = PatternMatcher([
 
 if getenv("HCQ2"): from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link # noqa: E402 # down here, hcq2 imports the helpers above
 
-def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, jit=False) -> UOp:
+def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None) -> 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)
   linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True)
-  if getenv("HCQ2"): linear = hcq_compile(linear, input_uops, jit=jit)
+  if getenv("HCQ2"): linear = hcq_compile(linear, input_uops)
   return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)
 
-def link_linear(linear:UOp, jit=False, cache=True) -> UOp: return hcq_link(linear, jit=jit, cache=cache) if getenv("HCQ2") else linear
+def link_linear(linear:UOp, cache=True) -> UOp: return hcq_link(linear, cache=cache) if getenv("HCQ2") else linear
 
 def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:Sequence[UOp]=(), update_stats=True, jit=False, wait=False):
   inputs = list(input_uops)
-  if not jit: linear = link_linear(compile_linear(linear, validate=VALIDATE_WITH_CPU, input_uops=inputs, jit=False))
+  if not jit: linear = link_linear(compile_linear(linear, validate=VALIDATE_WITH_CPU, input_uops=inputs))
   ctx = ExecContext(var_vals or {}, tuple(inputs), update_stats, jit, wait or DEBUG>=2)
   for call in linear.src: pm_exec.rewrite(call, ctx)
 
@@ -287,4 +287,5 @@ def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None
       from tinygrad.tensor import Tensor
       with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024, 1024).contiguous().realize(do_update_stats=False)
   ctx = ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False)
-  return pm_exec.rewrite(link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0), cache=ctx.cache).src[0], ctx)
+  linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0), cache=ctx.cache)
+  return max(pm_exec.rewrite(c, ctx) or 0.0 for c in linear.src)

tinygrad/llm/cli.py

@@ -113,7 +113,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, 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"])
@@ -164,9 +164,7 @@ 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)): model.warmup()
 
   # start server
   if args.serve: LLMServer(('', args.serve), model, model_name, tok, template).serve_forever()

tinygrad/llm/model.py

@@ -1,9 +1,12 @@
 from __future__ import annotations
 import functools, itertools, pathlib
 from dataclasses import dataclass, replace
-from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function
+from typing import cast
+from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, dtypes
+from tinygrad.nn import Linear
+from tinygrad.dtype import AddrSpace
 from tinygrad.llm.gguf import gguf_load
-from tinygrad.uop.ops import resolve
+from tinygrad.uop.ops import AxisType, KernelInfo, resolve
 
 @functools.cache
 def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, device:str|None=None) -> Tensor:
@@ -12,7 +15,7 @@ def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, device:str|
   return freqs.cos().cat(freqs.sin(), dim=-1).clone(device)
 
 class ExpertWeights:
-  """Like nn.Linear but with num_experts dimension. Weight shape: (num_experts, out_features, in_features)."""
+  """Like Linear but with num_experts dimension. Weight shape: (num_experts, out_features, in_features)."""
   def __init__(self, num_experts:int, in_features:int, out_features:int):
     self.weight = Tensor.zeros(num_experts, out_features, in_features)
   def __call__(self, sel:Tensor, x:Tensor) -> Tensor:
@@ -83,20 +86,20 @@ class FFNBlock:
 
     # --- feed-forward (MoE or dense) -------------------------------------
     if config.num_experts > 0:
-      self.ffn_gate_inp = nn.Linear(config.dim, config.num_experts, bias=False)  # router
+      self.ffn_gate_inp = Linear(config.dim, config.num_experts, bias=False)  # router
       if config.expert_bias: self.exp_probs_b = {"bias": Tensor.zeros(config.num_experts)}
       self.ffn_gate_exps = ExpertWeights(config.num_experts, config.dim, config.hidden_dim)
       self.ffn_up_exps = ExpertWeights(config.num_experts, config.dim, config.hidden_dim)
       self.ffn_down_exps = ExpertWeights(config.num_experts, config.hidden_dim, config.dim)
       if config.shared_expert_dim > 0:
-        self.ffn_gate_shexp = nn.Linear(config.dim, config.shared_expert_dim, bias=False)
-        self.ffn_up_shexp = nn.Linear(config.dim, config.shared_expert_dim, bias=False)
-        self.ffn_down_shexp = nn.Linear(config.shared_expert_dim, config.dim, bias=False)
+        self.ffn_gate_shexp = Linear(config.dim, config.shared_expert_dim, bias=False)
+        self.ffn_up_shexp = Linear(config.dim, config.shared_expert_dim, bias=False)
+        self.ffn_down_shexp = Linear(config.shared_expert_dim, config.dim, bias=False)
         if config.shared_expert_gate: self.ffn_gate_inp_shexp = {"weight": Tensor.zeros(config.dim)}
     else:
-      self.ffn_gate    = nn.Linear(config.dim, config.hidden_dim, bias=False)
-      self.ffn_up      = nn.Linear(config.dim, config.hidden_dim, bias=False)
-      self.ffn_down    = nn.Linear(config.hidden_dim, config.dim, bias=False)
+      self.ffn_gate    = Linear(config.dim, config.hidden_dim, bias=False)
+      self.ffn_up      = Linear(config.dim, config.hidden_dim, bias=False)
+      self.ffn_down    = Linear(config.hidden_dim, config.dim, bias=False)
 
   def _feed_forward(self, x:Tensor) -> Tensor:
     if hasattr(self, 'ffn_gate_exps'):
@@ -145,10 +148,10 @@ class TransformerBlock(FFNBlock):
     # --- attention projections (all linear, bias-free) ------------------
     q_proj_out       = config.head_dim * config.n_heads * (2 if config.attn_output_gate else 1)
     kv_proj_out      = config.head_dim * config.n_kv_heads
-    self.attn_q      = nn.Linear(config.dim, q_proj_out,  bias=config.qkv_bias)
-    self.attn_k      = nn.Linear(config.dim, kv_proj_out, bias=config.qkv_bias)
-    self.attn_v      = nn.Linear(config.dim, kv_proj_out, bias=config.qkv_bias)
-    self.attn_output = nn.Linear(config.head_dim * config.n_heads, config.dim, bias=False)
+    self.attn_q      = Linear(config.dim, q_proj_out,  bias=config.qkv_bias)
+    self.attn_k      = Linear(config.dim, kv_proj_out, bias=config.qkv_bias)
+    self.attn_v      = Linear(config.dim, kv_proj_out, bias=config.qkv_bias)
+    self.attn_output = Linear(config.head_dim * config.n_heads, config.dim, bias=False)
     if config.qk_norm: self.attn_q_norm, self.attn_k_norm = nn.RMSNorm(config.qk_norm, config.norm_eps), nn.RMSNorm(config.qk_norm, config.norm_eps)
 
   def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
@@ -195,16 +198,16 @@ class MLATransformerBlock(FFNBlock):
     super().__init__(config)
     qk_nope_head_dim = config.head_dim - config.rope_dim
     if config.q_lora_rank > 0:
-      self.attn_q_a = nn.Linear(config.dim, config.q_lora_rank, bias=False)
+      self.attn_q_a = Linear(config.dim, config.q_lora_rank, bias=False)
       self.attn_q_a_norm = nn.RMSNorm(config.q_lora_rank, config.norm_eps)
-      self.attn_q_b = nn.Linear(config.q_lora_rank, config.n_heads * config.head_dim, bias=False)
+      self.attn_q_b = Linear(config.q_lora_rank, config.n_heads * config.head_dim, bias=False)
     else:
-      self.attn_q = nn.Linear(config.dim, config.n_heads * config.head_dim, bias=False)
-    self.attn_kv_a_mqa = nn.Linear(config.dim, config.kv_lora_rank + config.rope_dim, bias=False)
+      self.attn_q = Linear(config.dim, config.n_heads * config.head_dim, bias=False)
+    self.attn_kv_a_mqa = Linear(config.dim, config.kv_lora_rank + config.rope_dim, bias=False)
     self.attn_kv_a_norm = nn.RMSNorm(config.kv_lora_rank, config.norm_eps)
     self.attn_k_b = {"weight": Tensor.zeros(config.n_heads, config.kv_lora_rank, qk_nope_head_dim)}
     self.attn_v_b = {"weight": Tensor.zeros(config.n_heads, config.v_head_dim, config.kv_lora_rank)}
-    self.attn_output = nn.Linear(config.n_heads * config.v_head_dim, config.dim, bias=False)
+    self.attn_output = Linear(config.n_heads * config.v_head_dim, config.dim, bias=False)
 
   def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
     B, T, _ = x.shape
@@ -237,6 +240,30 @@ class MLATransformerBlock(FFNBlock):
       self.cache_k = Tensor.empty(x.shape[0], 1, self.config.max_context, self.config.kv_lora_rank + self.config.rope_dim, device=x.device)
       self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta, device=x.device)
 
+# NOTE: this basic kernel works on any backend (CPU, AMD, ...); it's a sequential scan over all T tokens in a single kernel
+@functools.cache
+def _gated_delta_prefill_kernel(core:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:UOp, state:UOp, kq:UOp) -> UOp:
+  batch, heads, tokens, value_dim = cast(tuple[int, int, int, int], core.shape)
+  key_dim, alpha_dim = cast(int, q.shape[-1]), cast(int, alpha.shape[-1]) if len(alpha.shape) == 4 else 1
+  core, v = (x.reshape(batch*heads, tokens, value_dim) for x in (core, v))
+  q, k = (x.reshape(batch*heads, tokens, key_dim) for x in (q, k))
+  beta, kq = (x.reshape(batch*heads, tokens) for x in (beta, kq))
+  alpha, state = alpha.reshape(batch*heads, tokens, alpha_dim), state.reshape(batch*heads, value_dim, key_dim)
+  bh, row, cols = UOp.range(batch*heads, 0, AxisType.GLOBAL), UOp.range(value_dim, 2), tuple(range(key_dim))
+  current = UOp.placeholder((key_dim,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  current = current.after(UOp.group(*(current[col].store(state[bh, row, col].float()) for col in cols)))
+  token = UOp.range(tokens, 1, AxisType.REDUCE)
+  previous = tuple(current.after(token)[col].load() for col in cols)
+  keys, queries = (tuple(x[bh, token, col].load() for col in cols) for x in (k, q))
+  av, bv = alpha[bh, token, row if alpha_dim > 1 else 0].load(), beta[bh, token].load()
+  state_k = sum((x*y for x,y in zip(previous, keys)), UOp.const(0, dtypes.float32))
+  state_q = sum((x*y for x,y in zip(previous, queries)), UOp.const(0, dtypes.float32))
+  delta = (v[bh, token, row].load() - state_k*av) * bv
+  step = UOp.group(core[bh, token, row].store(state_q*av + delta*kq[bh, token]),
+                   *(current[col].store(x*av + delta*y) for col,x,y in zip(cols, previous, keys))).end(token)
+  stores = (state[bh, row, col].store(current.after(step)[col].load().cast(state.dtype)) for col in cols)
+  return UOp.group(*stores).end(row, bh).sink(arg=KernelInfo(name="gated_delta_prefill", opts_to_apply=()))
+
 class GatedDeltaNetBlock(FFNBlock):
   def __init__(self, config:TransformerConfig, ssm:SSMConfig):
     super().__init__(config)
@@ -244,54 +271,86 @@ class GatedDeltaNetBlock(FFNBlock):
     assert self.num_v_heads % self.num_k_heads == 0
     self.head_v_dim, self.ssm_conv_kernel = ssm.inner_size // ssm.time_step_rank, ssm.conv_kernel
     self.conv_channels, self.q_dim = ssm.inner_size + 2*ssm.group_count*ssm.state_size, ssm.state_size*ssm.group_count
-    self.attn_qkv = nn.Linear(config.dim, self.conv_channels, bias=False)
+    self.attn_qkv = Linear(config.dim, self.conv_channels, bias=False)
     if ssm.kda:
-      self.ssm_g_a, self.ssm_g_b = nn.Linear(config.dim, self.head_v_dim, bias=False), nn.Linear(self.head_v_dim, ssm.inner_size, bias=False)
-      self.ssm_f_a, self.ssm_f_b = nn.Linear(config.dim, self.head_k_dim, bias=False), nn.Linear(self.head_k_dim, ssm.inner_size, bias=False)
+      self.ssm_g_a, self.ssm_g_b = Linear(config.dim, self.head_v_dim, bias=False), Linear(self.head_v_dim, ssm.inner_size, bias=False)
+      self.ssm_f_a, self.ssm_f_b = Linear(config.dim, self.head_k_dim, bias=False), Linear(self.head_k_dim, ssm.inner_size, bias=False)
     else:
-      self.attn_gate = nn.Linear(config.dim, ssm.inner_size, bias=False)
-      self.ssm_alpha = nn.Linear(config.dim, self.num_v_heads, bias=False)
-    self.ssm_beta = nn.Linear(config.dim, self.num_v_heads, bias=False)
+      self.attn_gate = Linear(config.dim, ssm.inner_size, bias=False)
+      self.ssm_alpha = Linear(config.dim, self.num_v_heads, bias=False)
+    self.ssm_beta = Linear(config.dim, self.num_v_heads, bias=False)
     self.ssm_conv1d = {"weight": Tensor.zeros(self.conv_channels, self.ssm_conv_kernel)}
     self.ssm_dt = {"bias": Tensor.zeros(ssm.inner_size if ssm.kda else self.num_v_heads)}
     self.ssm_a = Tensor.zeros(self.num_v_heads, 1) if ssm.kda else Tensor.zeros(self.num_v_heads)
-    self.ssm_norm, self.ssm_out = nn.RMSNorm(self.head_v_dim, config.norm_eps), nn.Linear(ssm.inner_size, config.dim, bias=False)
+    self.ssm_norm, self.ssm_out = nn.RMSNorm(self.head_v_dim, config.norm_eps), Linear(ssm.inner_size, config.dim, bias=False)
 
   def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
     B, T, _ = x.shape
-    assert T == 1, "GatedDeltaNetBlock currently only supports T=1"
+    if resolve(T == 1):
+      # input processing
+      x = x.half()
+      out_gate = self.ssm_g_b(self.ssm_g_a(x)) if hasattr(self, "ssm_g_a") else self.attn_gate(x)
+      out_gate = out_gate.reshape(B, 1, self.num_v_heads, self.head_v_dim)
+      beta = self.ssm_beta(x).sigmoid().reshape(B, self.num_v_heads, 1, 1)
+      alpha = self.ssm_f_b(self.ssm_f_a(x)) if hasattr(self, "ssm_f_a") else self.ssm_alpha(x)
+      alpha = ((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, self.num_v_heads, -1) *
+               self.ssm_a.reshape(1, self.num_v_heads, -1)).exp().unsqueeze(-2)
+
+      # qkv conv
+      conv_window = self.conv_state.cat(self.attn_qkv(x), dim=1)
+      conv_out = (conv_window * self.ssm_conv1d["weight"].T.unsqueeze(0)).sum(1).silu()
+      q, k, v = conv_out.split([self.q_dim, self.q_dim, self.conv_channels - 2*self.q_dim], dim=-1)
+      q = q.reshape(B, self.num_k_heads, self.head_k_dim).normalize(dim=-1).repeat(1, self.num_v_heads//self.num_k_heads, 1)
+      k = k.reshape(B, self.num_k_heads, self.head_k_dim).normalize(dim=-1).repeat(1, self.num_v_heads//self.num_k_heads, 1)
+      v = v.reshape(B, self.num_v_heads, self.head_v_dim)
+      q, k, v = q.mul(self.head_k_dim**-0.5).unsqueeze(-1), k.unsqueeze(-1), v.unsqueeze(-1)
+
+      # recurrent
+      recurrent_state = self.recurrent_state * alpha
+      recurrent_state = recurrent_state + ((v - recurrent_state@k) * beta)@k.transpose(-1, -2)
+
+      # store the updated state
+      conv_state_store = self.conv_state.uop.store(conv_window[:, 1:, :].cast(self.conv_state.dtype).uop)
+      recurrent_state_store = self.recurrent_state.uop.store(recurrent_state.cast(self.recurrent_state.dtype).uop)
+      recurrent_state = Tensor(self.recurrent_state.uop.after(recurrent_state_store, conv_state_store))
+
+      # output
+      core_attn_out = self.ssm_norm((recurrent_state@q).squeeze(-1).reshape(B, 1, self.num_v_heads, self.head_v_dim))
+      out_gate = out_gate.sigmoid() if hasattr(self, "ssm_g_a") else out_gate.silu()
+      return self.ssm_out((core_attn_out * out_gate).reshape(B, 1, -1).cast(x.dtype))
+
+    is_kda = hasattr(self, "ssm_g_a")
 
     # input processing
     x = x.half()
-    out_gate = self.ssm_g_b(self.ssm_g_a(x)) if hasattr(self, "ssm_g_a") else self.attn_gate(x)
-    out_gate = out_gate.reshape(B, 1, self.num_v_heads, self.head_v_dim)
-    beta = self.ssm_beta(x).sigmoid().reshape(B, self.num_v_heads, 1, 1)
-    alpha = self.ssm_f_b(self.ssm_f_a(x)) if hasattr(self, "ssm_f_a") else self.ssm_alpha(x)
-    alpha = ((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, self.num_v_heads, -1) *
-             self.ssm_a.reshape(1, self.num_v_heads, -1)).exp().unsqueeze(-2)
+    out_gate = (self.ssm_g_b(self.ssm_g_a(x)) if is_kda else self.attn_gate(x)).reshape(B, T, self.num_v_heads, self.head_v_dim)
+    beta = self.ssm_beta(x).sigmoid().reshape(B, T, self.num_v_heads)
+    alpha = self.ssm_f_b(self.ssm_f_a(x)) if is_kda else self.ssm_alpha(x)
+    alpha = (((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, T, self.num_v_heads, -1) * self.ssm_a).squeeze(-1)
+             if is_kda else ((alpha.float() + self.ssm_dt["bias"]).softplus() * self.ssm_a).reshape(B, T, self.num_v_heads))
 
     # qkv conv
     conv_window = self.conv_state.cat(self.attn_qkv(x), dim=1)
-    conv_out = (conv_window * self.ssm_conv1d["weight"].T.unsqueeze(0)).sum(1).silu()
+    conv_out = (functools.reduce(lambda a,b: a+b,
+      (conv_window[:, i:i+T] * self.ssm_conv1d["weight"][:, i] for i in range(self.ssm_conv_kernel)))).silu()
     q, k, v = conv_out.split([self.q_dim, self.q_dim, self.conv_channels - 2*self.q_dim], dim=-1)
-    q = q.reshape(B, self.num_k_heads, self.head_k_dim).normalize(dim=-1).repeat(1, self.num_v_heads//self.num_k_heads, 1)
-    k = k.reshape(B, self.num_k_heads, self.head_k_dim).normalize(dim=-1).repeat(1, self.num_v_heads//self.num_k_heads, 1)
-    v = v.reshape(B, self.num_v_heads, self.head_v_dim)
-    q, k, v = q.mul(self.head_k_dim**-0.5).unsqueeze(-1), k.unsqueeze(-1), v.unsqueeze(-1)
+    q = q.reshape(B, T, self.num_k_heads, self.head_k_dim).normalize(dim=-1, eps=1e-12 if is_kda else 1e-6)
+    k = k.reshape(B, T, self.num_k_heads, self.head_k_dim).normalize(dim=-1, eps=1e-12 if is_kda else 1e-6)
+    q, k = q.repeat(1, 1, self.num_v_heads//self.num_k_heads, 1), k.repeat(1, 1, self.num_v_heads//self.num_k_heads, 1)
+    v = v.reshape(B, T, self.num_v_heads, self.head_v_dim)
+    q, k, v, beta = [z.transpose(1, 2).float() for z in (q, k, v, beta)]
+    alpha = alpha.transpose(1, 2).float().exp()
 
-    # recurrent
-    recurrent_state = self.recurrent_state * alpha
-    recurrent_state = recurrent_state + ((v - recurrent_state@k) * beta)@k.transpose(-1, -2)
-
-    # store the updated state
-    conv_state_store = self.conv_state.uop.store(conv_window[:, 1:, :].cast(self.conv_state.dtype).uop)
-    recurrent_state_store = self.recurrent_state.uop.store(recurrent_state.cast(self.recurrent_state.dtype).uop)
-    recurrent_state = Tensor(self.recurrent_state.uop.after(recurrent_state_store, conv_state_store))
+    # recurrent: run the gated delta rule over all T tokens in a single custom kernel, writing back the updated state
+    conv_state = conv_window[:, T:T+self.ssm_conv_kernel-1].cast(self.conv_state.dtype).contiguous()
+    state = Tensor(self.recurrent_state.uop.after(self.conv_state.uop.store(conv_state.uop)))
+    q, kq = q * self.head_k_dim**-0.5, ((q * self.head_k_dim**-0.5)*k).sum(-1).contiguous()
+    core = Tensor.custom_kernel(Tensor.empty_like(v), q, k, v, beta, alpha, state, kq,
+                                fxn=_gated_delta_prefill_kernel)[0].transpose(1, 2)
 
     # output
-    core_attn_out = self.ssm_norm((recurrent_state@q).squeeze(-1).reshape(B, 1, self.num_v_heads, self.head_v_dim))
-    out_gate = out_gate.sigmoid() if hasattr(self, "ssm_g_a") else out_gate.silu()
-    return self.ssm_out((core_attn_out * out_gate).reshape(B, 1, -1).cast(x.dtype))
+    gate = out_gate.sigmoid() if is_kda else out_gate.silu()
+    return self.ssm_out((self.ssm_norm(core) * gate).reshape(B, T, -1).cast(x.dtype)).contiguous()
 
   # recurrent state can't be partially reused after divergence, force a full rebuild
   def _state_reset_ops(self):
@@ -302,7 +361,7 @@ class GatedDeltaNetBlock(FFNBlock):
   def _init_state(self, x):
     if not hasattr(self, "conv_state"):
       self.conv_state = Tensor.zeros(x.shape[0], self.ssm_conv_kernel-1, self.conv_channels, device=x.device).clone()
-      self.recurrent_state = Tensor.zeros(x.shape[0], self.num_v_heads, self.head_v_dim, self.head_v_dim, device=x.device).clone()
+      self.recurrent_state = Tensor.zeros(x.shape[0], self.num_v_heads, self.head_v_dim, self.head_k_dim, device=x.device).clone()
 
 class Transformer:
   def __init__(self, config:TransformerConfig):
@@ -314,7 +373,7 @@ class Transformer:
                                block_cls(dense_config if i < config.leading_dense_blocks else config) for i in range(config.num_blocks)]
     self.token_embd  = nn.Embedding(config.vocab_size, config.dim)
     self.output_norm = nn.RMSNorm(config.dim, config.norm_eps)
-    self.output = nn.Linear(config.dim, config.vocab_size, bias=False)
+    self.output = Linear(config.dim, config.vocab_size, bias=False)
     self.max_context = config.max_context
     self.has_recurrent_block = any(isinstance(b, GatedDeltaNetBlock) for b in self.blk)
     self._cached_tokens: list[int] = []
@@ -415,12 +474,14 @@ class Transformer:
       Tensor.realize(*params)
     return model, kv
 
+  def warmup(self):
+    for _ in range(2): list(zip(range(2), self.generate([0])))
+
   def get_start_pos(self, tokens:list[int]) -> int:
     prefix_len = sum(1 for _ in itertools.takewhile(lambda ab: ab[0] == ab[1], zip(tokens[:-1], self._cached_tokens)))
     return min(block._reusable_prefix_len(prefix_len, len(self._cached_tokens)) for block in self.blk)
 
   def generate(self, tokens:list[int], chunk_size:int=32, temperature:float=0.0):
-    if self.has_recurrent_block: chunk_size = 1
     v_start_pos = UOp.variable("start_pos", 0, self.max_context-1)
     v_toks = UOp.variable("toks", 1, chunk_size)
     # TODO: use UOp.variable for temperature once float variables are supported
@@ -432,9 +493,12 @@ class Transformer:
     if start_pos < len(self._cached_tokens) and (resets := [r for b in self.blk for r in b._state_reset_ops()]): Tensor.realize(*resets)
     out, prompt_len = None, len(tokens)
     while len(tokens) < self.max_context:
-      n_toks = min(chunk_size, len(tokens) - start_pos)
+      # NOTE: Tensor.custom_kernel requires all-int input shapes (gated delta prefill), so recurrent blocks run the prompt
+      # in fixed-size chunks and roll out T=1, giving the JIT one prefill graph (per chunk size) and one rollout graph
+      n_toks = 1 if self.has_recurrent_block and len(tokens) - start_pos < chunk_size else min(chunk_size, len(tokens) - start_pos)
       sp, nt = v_start_pos.bind(start_pos), v_toks.bind(n_toks)
-      out = self(t[:, sp:sp+nt] if start_pos < prompt_len or out is None else out, sp, temp).realize()
+      sl = t[:, sp:sp+(n_toks if self.has_recurrent_block else nt)]
+      out = self(sl if start_pos < prompt_len or out is None else out, sp, temp).realize()
       start_pos += n_toks
       # chunked prefill: keep processing until all prompt tokens are consumed
       if start_pos < len(tokens): continue

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,7 @@ 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)
+      rendered = self.server.template.render(messages=body["messages"], tools=body.get("tools"), add_generation_prompt=True, 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:
@@ -129,7 +140,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/mixin/creation.py

@@ -1,5 +1,5 @@
 from typing import TYPE_CHECKING, Callable, Self
-from tinygrad.dtype import ConstType, DTypeLike, Invalid, dtypes, to_dtype
+from tinygrad.dtype import ConstType, DType, DTypeLike, Invalid, dtypes, to_dtype
 from tinygrad.helpers import argfix, prod
 from tinygrad.mixin.dtype import DTypeMixin
 from tinygrad.mixin.movement import MovementMixin
@@ -11,7 +11,7 @@ class CreationMixin(DTypeMixin, MovementMixin):
   @staticmethod
   def const(b, dtype=None): raise NotImplementedError
 
-  def const_like(self, b: ConstType) -> Self: return self._wrap_uop(self._uop.const_like(b))
+  def const_like(self, b: ConstType, dtype:DType|None=None) -> Self: return self._wrap_uop(self._uop.const_like(b, dtype))
 
   def _multi_like(self, fxn:'Callable[[tuple[sint, ...], str|None], Self]') -> Self:
     from tinygrad.uop.ops import UOp

tinygrad/mixin/op.py

@@ -287,7 +287,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     pads = tuple((smax(pB,0), smax(pA,0)) for pB,pA in pX) if has_neg else pX
     base = MovementMixin.pad(X, pads)
     if value == 0: return base
-    return MovementMixin.pad(X.const_like(1).cast(dtypes.bool), pads).where(base, value)
+    return MovementMixin.pad(X.const_like(True, dtypes.bool), pads).where(base, value)
 
   def _pad_circular(self, pX:tuple[tuple[sint, sint], ...]) -> Self:
     # shrink first for negative pads, then wrap the non-negative remainder
@@ -926,7 +926,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     ```
     """
     x, dim = self, self._resolve_dim(dim)
-    if (orig_len := int(x.shape[dim])) <= 1: return x, x.const_like(0).cast(dtypes.default_int)
+    if (orig_len := int(x.shape[dim])) <= 1: return x, x.const_like(0, dtypes.default_int)
     # pad to power of 2
     n_stages = (orig_len-1).bit_length()
     pads = tuple((0, 2**n_stages - orig_len) if i == dim else None for i in range(x.ndim))
@@ -1733,7 +1733,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     if Y.device is not None and self.device is not None and Y.device != self.device:
       raise RuntimeError(f"expected Y and self on the same device, {Y.device=}, {self.device=}")
     log_probs = self.log_softmax()
-    loss_mask = Y.ne(ignore_index) if ignore_index != -1 else Y.const_like(1).cast(dtypes.bool)
+    loss_mask = Y.ne(ignore_index) if ignore_index != -1 else Y.const_like(True, dtypes.bool)
     y = Y.unsqueeze(-1)._one_hot_along_dim(self.shape[-1], dim=-1) * loss_mask.unsqueeze(-1)
     smoothing = label_smoothing * (log_probs.mean(-1) * loss_mask)
     unreduced = ((1 - label_smoothing) * (log_probs * y).sum(-1) + smoothing)

tinygrad/mixin/rand.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 import math
 from typing import Self, cast
-from tinygrad.dtype import DType, DTypeLike, dtypes, least_upper_dtype, to_dtype
+from tinygrad.dtype import DType, DTypeLike, dtypes, least_upper_dtype, to_dtype, bitcast
 from tinygrad.helpers import all_int, argfix, ceildiv, prod, TRAINING
 from tinygrad.mixin.op import OpMixin
 from tinygrad.device import canonicalize_device
@@ -12,7 +12,7 @@ class RandMixin(OpMixin):
   def _threefry_random_bits(key, counts0, counts1):
     x = (counts1.cast(dtypes.uint64) << 32) | counts0.cast(dtypes.uint64)
     x = x.threefry((key[1].cast(dtypes.uint64) << 32) | key[0].cast(dtypes.uint64))
-    return (x & 0xffffffff).cast(dtypes.uint32).cat(((x >> 32) & 0xffffffff).cast(dtypes.uint32))
+    return x.cast(dtypes.uint32).cat((x >> 32).cast(dtypes.uint32))
 
   @classmethod
   def random_bits(cls, key:Self, counter:Self, num:int) -> Self:
@@ -33,7 +33,7 @@ class RandMixin(OpMixin):
     _, nmant = dtypes.finfo(dtype)
     uint_dtype = {1: dtypes.uint8, 2: dtypes.uint16, 4: dtypes.uint32, 8: dtypes.uint64}[dtype.itemsize]
     uint_bits = bits.bitcast(uint_dtype)
-    float_one_bits = uint_bits.const_like(1).cast(dtype).bitcast(uint_dtype)
+    float_one_bits = bitcast(1.0, dtype, uint_dtype)
     return uint_bits.rshift(dtype.bitsize - nmant).bitwise_or(float_one_bits).bitcast(dtype)[:prod(shape)].sub(1).reshape(shape)
 
   @classmethod
@@ -320,7 +320,7 @@ class RandMixin(OpMixin):
     # handle attention mask
     if is_causal:
       if attn_mask is not None: raise RuntimeError("cannot set attn_mask when is_causal=True")
-      attn_mask = qk.const_like(1).cast(dtypes.bool).tril()
+      attn_mask = qk.const_like(True, dtypes.bool).tril()
     if attn_mask is not None:
       if attn_mask.dtype == dtypes.bool: attn_mask = attn_mask.where(0, -float("inf"))
       qk = qk + attn_mask

tinygrad/nn/onnx.py

@@ -959,7 +959,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
   # Reimplemented here because you need legacy RNG for passing ONNX tests.
   def dropout_7(data:Tensor, ratio:float=0.5, training_mode:bool=False, seed:int|None=None):
     import numpy as np
-    if not training_mode: return data, data.const_like(True).cast(dtypes.bool)
+    if not training_mode: return data, data.const_like(True, dtypes.bool)
     if seed is not None:
       rand = Tensor(np.random.RandomState(seed).random(cast(tuple[int,...], data.shape)), dtype=data.dtype, device=data.device)
     else:

tinygrad/renderer/wgsl.py

@@ -6,7 +6,7 @@ from tinygrad.helpers import strip_parens
 def _mask(dt:DType): return 0xFF if dt.itemsize == 1 else 0xFFFF
 
 def sign_extend(val:UOp, sext_am:int):
-  return (UOp.where((val >> (sext_am - 1)) > 0, UOp.const(0xffffffff, dtypes.uint32) << sext_am, UOp.const(0, dtypes.uint32)) \
+  return (UOp.where((val >> (sext_am - 1)) > 0, UOp.const(0xffffffff << sext_am, dtypes.uint32), UOp.const(0, dtypes.uint32)) \
         | val.bitcast(dtypes.uint32)).bitcast(dtypes.int)
 
 # store for char: buf[idx/4] <- (var << (idx%4)*8))

tinygrad/runtime/ops_cpu.py

@@ -164,13 +164,13 @@ class CPUDevice(HCQCompiled):
     (UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),)), encode_host_queue)])
 
   pm_bufferize = PatternMatcher([
-    (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx[0].timeline("sentinel", (1 << 64) - 1)),
-    (UPat(Ops.PARAM, tag="COMPUTE:0_timeline_signal"), lambda ctx: ctx[0].timeline("signal", 0)),
-    (UPat(Ops.PARAM, tag="COMPUTE:0_timeline_value"), lambda ctx: ctx[0].timeline("value", 1)),
+    (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx[0].signal("sentinel", (1 << 64) - 1)),
+    (UPat(Ops.PARAM, tag="timeline_signal"), lambda ctx: ctx[0].signal("timeline")),
+    (UPat(Ops.PARAM, tag="timeline_value"), lambda ctx: ctx[0].signal("value", 1)),
   ])
 
   @functools.cache
-  def timeline(self, tag:str, init_value:int) -> Buffer:
+  def signal(self, name:str, init_value:int=0) -> Buffer:
     (buf:=Buffer(self.device, 1, dtypes.uint64, preallocate=True)).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = init_value
     return buf
 

tinygrad/runtime/ops_python.py

@@ -5,11 +5,11 @@
 from typing import Any, TYPE_CHECKING
 import pickle, base64, itertools, time, sys, functools
 from dataclasses import replace
-from tinygrad.dtype import DType, dtypes, AddrSpace, truncate, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar
+from tinygrad.dtype import bitcast, DType, dtypes, AddrSpace, truncate, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar
 from tinygrad.helpers import all_same, getenv, flatten, Target, IMAGE, is_image_shape, cpu_profile
 from tinygrad.device import Buffer, Compiled, Compiler, Allocator, Program, TinyELF
 from tinygrad.codegen.opt import tc
-from tinygrad.uop.ops import exec_alu, python_alu, Ops, UOp, GroupOp, bitcast
+from tinygrad.uop.ops import exec_alu, python_alu, Ops, UOp, GroupOp
 from tinygrad.renderer import Renderer
 
 def _load(m, i, dtype: DType):

tinygrad/runtime/support/hcq2.py

@@ -41,28 +41,33 @@ 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_patches(buf:UOp, patches:Sequence[tuple[sint, UOp]]) -> UOp:
-  offsets = UOp(Ops.STACK, dtypes.int, tuple(UOp.const(off // buf.dtype.itemsize, dtypes.int) for off,_ in patches))
-  values = UOp(Ops.STACK, buf.dtype, tuple(val.cast(buf.dtype) for _,val in patches))
-  return buf.index(offsets).store(values)
+def is_value_known_at_link(val:UOp) -> bool:
+  runtime_reads = [u for u in val.toposort() if u.op in (Ops.LOAD, Ops.INDEX)]
+  addressed_bufs = [b for g in val.toposort() if g.op is Ops.GETADDR for b in unwrap_mstack(g.buf_uop)]
+
+  # addr of input params is not known at link time
+  return not val.variables() and not runtime_reads and all(b.op is not Ops.PARAM or b.tag is not None for b in addressed_bufs)
+
+def make_patches(buf:UOp, patches:Sequence[tuple[sint, UOp]]) -> tuple[UOp, ...]:
+  return tuple(buf.index(UOp(Ops.STACK, dtypes.int, tuple(UOp.const(off // buf.dtype.itemsize, dtypes.int) for off,_ in ps)))
+               .store(UOp(Ops.STACK, buf.dtype, tuple(val.cast(buf.dtype) for _,val in ps))).rtag(tag)
+               for ps, tag in zip(partition(patches, lambda p: is_value_known_at_link(p[1])), ("link", None)) if ps)
 
 def make_binary_patch(buf:UOp, blob:bytes) -> UOp:
   data = UOp(Ops.BINARY, src=(), arg=blob).bitcast(buf.dtype)
   r = UOp.range(len(blob) // buf.dtype.itemsize, 0, dtype=dtypes.int, src=(buf, data))
-  return buf.index(r).store(data.index(r).load()).end(r)
+  return buf.index(r).store(data.index(r).load()).end(r).rtag("link")
 
-def make_cmdbuf(lin, devs, buf:UOp|None=None, dep:UOp|None=None):
+def make_cmdbuf(lin, devs, buf:UOp|None=None):
   blob, patches = bytearray(), []
   for s in (s for ins in lin.src for s in ins.src):
     if s.op is not Ops.CONST: patches.append((len(blob), s))
     blob.extend(struct.pack(f'<{s.dtype.fmt}', s.val if s.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, bytes(blob)), *((make_patches(writable, patches),) if patches else ()))
+  return cmdbuf.after(make_binary_patch(cmdbuf, bytes(blob)), *make_patches(cmdbuf, patches))
 
-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")
-def make_signal_value(devs, queue="COMPUTE:0"): return UOp.placeholder((1,), dtypes.uint64, 0, device=devs).rtag(f"{queue}_timeline_value")
+def make_signal(devs, slot:int=0, tag:str="signal") -> UOp:
+  return UOp.placeholder((1,), dtypes.uint64, slot, device=devs, volatile=True).rtag(tag)
 
 def make_submit(*cmds, devs:str|tuple[str, ...], queue:str) -> UOp:
   return UOp.custom_function("submit_cmdbuf", UOp(Ops.LINEAR, src=tuple(cmds), arg=(to_tuple(devs), queue)))
@@ -72,7 +77,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_patches(buf, [(i * 4, w) for i, w in enumerate(words)]),) if words else ()))
+  return buf.after(*make_patches(buf, [(i * 4, w) for i, w in enumerate(words)]))
 
 # *****************
 # 0.1. prep: replace buffers with params
@@ -115,7 +120,7 @@ def _get_deps(ctx:DepsTracker, bufs_by_lane:list[list[Any]], write, key:tuple[tu
     dep_lanes += [(dep, dlane, lane) for dep, dlane in ctx.access_resources(bufs, written, (key, lane))]
   return dep_lanes
 
-def _build_wait_cmds(dep_lanes:list[tuple[tuple, int, int]], devices:tuple[str, ...], queue:str) -> tuple[list[UOp], set[int]]:
+def _build_wait_cmds(slots:dict[str, int], dep_lanes:list[tuple[tuple, int, int]], devices:tuple[str, ...], queue:str) -> tuple[list[UOp], set[int]]:
   # opt1: same-queue ops are fifo-ordered
   if devices[0].split(":")[0] in {"AMD", "QCOM"} or queue.startswith("COPY"):
     dep_lanes = [(dep, dlane, lane) for dep, dlane, lane in dep_lanes if (dep[0][dlane], dep[1]) != (devices[lane], queue)]
@@ -127,70 +132,81 @@ def _build_wait_cmds(dep_lanes:list[tuple[tuple, int, int]], devices:tuple[str,
 
   waits = []
   for (ddevs, dqueue, dtag), lanes in deps.items():
-    sig = UOp.mstack(*[make_signal(d if dl is None else ddevs[dl], queue=dqueue, sentinel=dl is None) for dl, d in zip(lanes, devices)])
-    val = UOp.mstack(*[make_signal_value(d if dl is None else ddevs[dl], queue=dqueue) for dl, d in zip(lanes, devices)])
-    waits.append(UOp(Ops.INS, arg="wait", src=(sig, val.index(UOp.const(0, dtypes.int)) + dtag)))
+    sig = UOp.mstack(*[make_signal(d, tag="sentinel_signal") if dl is None else make_signal(ddevs[dl], slots[dqueue])
+                       for dl, d in zip(lanes, devices)])
+    waits.append(UOp(Ops.INS, arg="wait", src=(sig, UOp.const(dtag + 1, dtypes.uint64))))
   return waits, {dtag for _, _, dtag in deps}
 
+def make_fence(timeline:UOp, prev:UOp, sigs:list[UOp]) -> UOp:
+  free = (cur:=timeline.after(loop:=UOp.loop(0)).index(0).load()).end(loop, cur < prev.index(0).load())
+  return UOp.sink(*[s.after(free).index(0).store(0) for s in sigs])
+
+def _hcq_call(devs, name:str, body:UOp) -> UOp: return UOp.custom_function("hcq", body).call(aux=HCQInfo(name, Estimates(), devs, "COMPUTE:0"))
+
 def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[tuple[tuple[str, ...], str]],
-                      tracker:HCQDepsTracker) -> tuple[list[UOp], set[int]]:
+                      tracker:HCQDepsTracker, slots:dict[str, int]) -> tuple[list[UOp], list[UOp], set[int]]:
   # collect all buffers which belong to devices
   dev_bufs:dict[str, dict[int, Any]] = collections.defaultdict(dict)
   for call, devices in batch:
     for b in itertools.chain.from_iterable(_get_call_bufs_by_lane(call, devices)):
       for bd in to_tuple(b.device): dev_bufs[bd][id(b)] = b
 
-  zero, n, submits, bumps, waited = UOp.const(0, dtypes.int), len(batch_info), [], [], set()
+  n, fences, fins, waited = len(batch_info), [], [], set()
   for _, devgroup in itertools.groupby(sorted(dev_bufs), key=lambda d: d.split(":")[0]):
     devs = tuple(devgroup)
 
     # to finalize the batch, sync all accesses from other devices to buffers that belong to this device
     fin_deps = [dl for dl in _get_deps(tracker, [list(dev_bufs[d].values()) for d in devs], None, key=(devs, "COMPUTE:0", n)) if dl[0][2] < n]
-    waits, cur_waited = _build_wait_cmds(fin_deps, devs, "COMPUTE:0")
+    waits, cur_waited = _build_wait_cmds(slots, fin_deps, devs, "COMPUTE:0")
     waited |= cur_waited
 
-    # wait the syncs, store the device epoch
-    store = UOp(Ops.INS, arg="store", src=(make_signal(devs), (tl:=make_signal_value(devs)).index(zero) + n))
-    submits.append((devs, make_submit(*waits, store, devs=devs, queue="COMPUTE:0")))
-    upd = [(tl, n + 1)] + [(make_signal_value(devs, queue=qn), n)
-                           for qn in dedup([qn for bdevs, qn in batch_info if set(bdevs) & set(devs)]) if qn != "COMPUTE:0"]
-    bumps.append((devs, UOp.barrier(*[s.index(zero, dtype=s.dtype).store(s.index(zero) + inc) for s, inc in upd])))
+    # wait the syncs and signal the device epoch, then bump the timeline on the host
+    timeline, tl = make_signal(devs, tag="timeline_signal"), make_signal(devs, tag="timeline_value")
+    submit = make_submit(*waits, UOp(Ops.INS, arg="store", src=(timeline, tl.index(0))), devs=devs, queue="COMPUTE:0")
+    cur = (bump:=tl.after(submit).index(0)).load()
+    bumps = [bump.store(cur + 1)]
 
-  # NOTE: submit before bumps
-  fins = [UOp.custom_function("hcq", b.sink()).call(aux=HCQInfo("hcq_finalizer", Estimates(), devs, "COMPUTE:0")) for devs, b in submits + bumps]
-  return fins, waited
+    # devices running the batch reset their queue signals before each run, fencing on the epoch kept from the previous one
+    if qs:=dedup([qn for bdevs, qn in batch_info if set(bdevs) & set(devs)]):
+      prev = make_signal(devs, next(UOp.unique_num))
+      fences.append(_hcq_call(devs, "hcq_fence", make_fence(timeline, prev, [make_signal(devs, slots[q]) for q in qs])))
+      bumps.append(prev.after(submit).index(0).store(cur))
+    fins.append(_hcq_call(devs, "hcq_finalizer", UOp.sink(*bumps)))
+  return fences, fins, waited
 
 def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]]) -> list[UOp]:
   batch_info = [(devices, "COMPUTE:0" if call.src[0].op is Ops.PROGRAM else "COPY:0") for call, devices in batch]
 
   # schedule deps
   waited:set[int] = set()
+  slots:dict[str, int] = collections.defaultdict(lambda: next(UOp.unique_num))
   deps_tracker = HCQDepsTracker()
   call_waits:list[list[UOp]] = []
   for tag, ((call, _), (devices, queue)) in enumerate(zip(batch, batch_info)):
     deps = _get_deps(deps_tracker, _get_call_bufs_by_lane(call, devices), get_call_outs_ins(call)[0], key=(devices, queue, tag))
-    cmds, cur_waited = _build_wait_cmds(deps, devices, queue)
+    cmds, cur_waited = _build_wait_cmds(slots, deps, devices, queue)
     call_waits.append(cmds)
     waited |= cur_waited
 
-  # build finalizers
-  finalizers, finalizer_waited = _build_finalizers(batch, batch_info, deps_tracker)
+  # build fences and finalizers
+  fences, finalizers, finalizer_waited = _build_finalizers(batch, batch_info, deps_tracker, slots)
   waited |= finalizer_waited
 
   src = []
   for tag, ((call, _), (devices, queue), q) in enumerate(zip(batch, batch_info, call_waits)):
-    # first queue use, sync prior device work with main signal
+    # first queue use, sync prior device work with the device timeline
     if batch_info.index((devices, queue)) == tag:
-      q = [UOp(Ops.INS, arg="barrier", src=()), UOp(Ops.INS, arg="wait", src=(make_signal(devices), make_signal_value(devices).index(0) - 1))] + q
+      epoch = make_signal(devices, tag="timeline_value").index(0) - 1
+      q = [UOp(Ops.INS, arg="barrier", src=()), UOp(Ops.INS, arg="wait", src=(make_signal(devices, tag="timeline_signal"), epoch))] + q
 
     # and make hcq call
     info = HCQInfo(get_call_name(call, get_call_arg_uops(call)), estimate_uop(call), devices, queue)
     q += [call.replace(arg=replace(call.arg, aux=info))]
 
-    # signal queue timeline if someone waits for us
-    if tag in waited: q += [UOp(Ops.INS, arg="store", src=(make_signal(devices, queue), make_signal_value(devices, queue).index(0) + tag))]
+    # signal the queue if someone waits for us
+    if tag in waited: q += [UOp(Ops.INS, arg="store", src=(make_signal(devices, slots[queue]), UOp.const(tag + 1, dtypes.uint64)))]
     src.append(UOp.custom_function("hcq", make_submit(*q, devs=devices, queue=queue).sink()).call(name="hcq", aux=info))
-  return src + finalizers
+  return fences + src + finalizers
 
 def sched_hcq_batches(l:UOp) -> UOp:
   srcs:list[UOp] = []
@@ -216,7 +232,7 @@ def merge_queues(linear:UOp) -> UOp:
   limits:dict[tuple[tuple[str, ...], str], int] = collections.defaultdict(lambda: JIT_BATCH_SIZE.value)
 
   for call in linear.src:
-    if not isinstance(info:=call.arg.aux, HCQInfo) or info.name == "hcq_finalizer": # non-hcq call or finalizer: close all open queues
+    if not isinstance(info:=call.arg.aux, HCQInfo) or info.name.startswith("hcq_"): # non-hcq call, fence or finalizer: close all open queues
       new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in list(opened_qs)] + [call]
       continue
 
@@ -246,21 +262,11 @@ pm_encode_cmdbufs = PatternMatcher([
 
 # *****************
 
-def is_value_known_at_link(val:UOp) -> bool:
-  runtime_reads = [u for u in val.toposort() if u.op in (Ops.LOAD, Ops.INDEX)]
-  addressed_bufs = [b for g in val.toposort() if g.op is Ops.GETADDR for b in unwrap_mstack(g.buf_uop)]
+def get_getaddrs(p:UOp) -> list[UOp]: return [u for u in p.toposort(gate=lambda u: u.op is not Ops.AFTER) if u.op is Ops.GETADDR]
 
-  # addr of input params is not known at link time
-  return not val.variables() and not runtime_reads and all(b.op is not Ops.PARAM or b.tag is not None for b in addressed_bufs)
-
-def is_link_patch(p:UOp, jit:bool) -> bool:
-  if p.tag == "link": return True
-  store = p.src[0] if (is_binary_patch:=(p.op is Ops.END and p.src[0].op is Ops.STORE)) else p
-  if not jit: return store.buf_uop.tag == "program"
-  return is_binary_patch or (store.op is Ops.STORE and is_value_known_at_link(store.src[1]))
-
-def trim_link_patches(ctx:tuple[bool, list[UOp]], a:UOp) -> UOp|None:
-  links, kept = partition(a.src[1:], lambda p: is_link_patch(p, ctx[0]))
+def trim_link_patches(ctx:tuple[list[UOp], list[UOp]], a:UOp) -> UOp|None:
+  links, kept = partition(a.src[1:], lambda p: p.tag == "link")
+  ctx[0].extend(kept)
 
   # keep all patches from the link-time patches' subtrees in the C code
   afters = [u for u in UOp.sink(*links).toposort() if u.op is Ops.AFTER]
@@ -268,18 +274,7 @@ def trim_link_patches(ctx:tuple[bool, list[UOp]], a:UOp) -> UOp|None:
   return a.src[0].after(*kept, *[d for p in afters for d in p.src[1:]]) if links else None
 pm_trim_link_patches = PatternMatcher([(UPat(Ops.AFTER, src=(UPat((Ops.PARAM, Ops.MSTACK)),), allow_any_len=True, name="a"), trim_link_patches)])
 
-def split_patches(ctx:bool, call:UOp) -> UOp|None:
-  lt_patches:list[UOp] = []
-  body = graph_rewrite(call.src[0], pm_trim_link_patches, ctx=(ctx, lt_patches), name=f"trim link-time patches ({call.arg.aux.name})")
-
-  lt_srcs = collections.defaultdict(list)
-  for p in lt_patches: lt_srcs[p.buf_uop].append(p)
-  return call.replace(src=(body, *call.src[1:], *[b.after(*ps) for b,ps in lt_srcs.items()]))
-pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), split_patches)])
-
-# *****************
-
-def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[dict[UOp, UOp], tuple[UOp, ...]]:
+def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[UOp, dict[UOp, UOp], tuple[UOp, ...], dict[UOp, int]]:
   bare = {g: g.replace(src=(g.src[0].without_after,)) for g in gaddrs}
 
   order = sorted(dedup(bare.values()), key=lambda g: ((b:=unwrap_mstack(g.buf_uop)[0]).arg.slot, repr(b.tag)))
@@ -287,26 +282,47 @@ 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(slots[bare[g]], dtypes.int)).load() for g in gaddrs}
-  return reads, (table.after(make_patches(table, [(i * table.dtype.itemsize, addr) for addr, i in slots.items()])),) if slots else ()
+  fills = (table.after(*make_patches(table, [(i*table.dtype.itemsize, addr) for addr, i in slots.items()])),) if slots else ()
+  return table, reads, fills, {g:slots[bare[g]] for g in gaddrs}
 
-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}
-  return bufs, tuple(buf.after(make_binary_patch(buf, b.src[0].arg)) for b,buf in bufs.items())
+def make_scatter_loop(patches:list[UOp], inputs_table:tuple, lt_patches:list[UOp]) -> dict[UOp, UOp]:
+  (table, _, _, slots), dst, data, subs = inputs_table, patches[0].buf_uop, [], {}
+  for p in patches:
+    words = [(off, val, get_getaddrs(val)) for off,val in zip(p.src[0].src[1].src, p.src[1].src)]
+    data += [off.val << 32 | slots[gaddrs[0]] for off,_,gaddrs in words if gaddrs][::2]
+    scalars = [(off.val*dst.dtype.itemsize, val) for off,val,gaddrs in words if not gaddrs]
+    subs[p] = UOp.group(*make_patches(dst, scalars)) if scalars else UOp(Ops.NOOP)
 
-def rm_rt_uops(call:UOp) -> UOp|None:
-  if not (rt_uops:=[u for u in call.src[0].toposort() if u.op is Ops.GETADDR or (u.op is Ops.BITCAST and u.src[0].op is Ops.BINARY)]): return None
-  gaddrs, blobs = partition(rt_uops, lambda u: u.op is Ops.GETADDR)
-  inputs, internals = partition(gaddrs, lambda g: all(x.op is Ops.PARAM and x.tag is None for x in unwrap_mstack(g.buf_uop)))
+  # plan entry: dst word offset << 32 | addr table slot
+  plan = UOp.placeholder((len(data),), dtypes.uint64, next(UOp.unique_num), device=dst.device).rtag("systems")
+  entry = plan.index(ridx:=UOp.range(len(data), next(UOp.unique_num), dtype=dtypes.int, src=(plan, dst))).load()
+  slot, widx = ((entry & 0xffffffff) % table.max_numel()).cast(dtypes.int), ((entry >> 32) % (dst.max_numel()-1)).cast(dtypes.int) # CHECK_OOB bounds
+  loop = UOp.group(*[dst.index(widx+i).store((table.index(slot).load() >> 32*i).cast(dtypes.uint32)) for i in range(2)]).end(ridx)
+  lt_patches.append(make_binary_patch(plan, struct.pack(f'<{len(data)}Q', *data)))
+  subs[patches[0]] = UOp.group(loop, subs[patches[0]])
+  return subs
+
+def is_input_addr(g:UOp) -> bool: return all(x.op is Ops.PARAM and x.tag is None for x in unwrap_mstack(g.buf_uop))
+
+def split_patches(call:UOp) -> UOp|None:
+  rt_patches:list[UOp] = []
+  lt_patches:list[UOp] = []
+  body = graph_rewrite(call.src[0], pm_trim_link_patches, ctx=(rt_patches, lt_patches), name=f"trim link-time patches ({call.arg.aux.name})")
+
+  # split patches
+  inputs, internals = partition(dedup(g for p in rt_patches for g in get_getaddrs(p)), is_input_addr)
   runtimes, systems = partition(internals, lambda g: any(x.tag in {"program", "kernargs", "cmdbuf"} for x in unwrap_mstack(g.buf_uop)))
+  tables = [make_addr_table(call, gs, n) for gs,n in ((inputs, "inputs"), (runtimes, "runtime"), (systems, "systems"))]
+  reads, fills = {k:v for _,r,_,_ in tables for k,v in r.items()}, [f for t in tables[1:] for f in t[2]] # inputs table is filled by exec
+  input_patches = [p for p in rt_patches if (gs:=get_getaddrs(p)) and all(map(is_input_addr, gs))]
+  scatter = make_scatter_loop(input_patches, tables[0], lt_patches) if input_patches else {}
+  body = body.substitute({p:p.substitute(scatter | reads) for p in rt_patches})
 
-  # exec fills the inputs table with the input addresses every run, so it has no fill patches
-  (reads, _), *tables = [make_addr_table(call, gs, n) for gs,n in ((inputs, "inputs"), (runtimes, "runtime"), (systems, "systems"))] + \
-                        [make_blob_bufs(call, blobs)]
-  reads, fills = reads | {k:v for r,_ in tables for k,v in r.items()}, [f for _,fs in tables for f in fs]
-  return call.replace(src=(call.src[0].substitute(reads), *call.src[1:], *fills),
-                      arg=replace(call.arg, aux=replace(call.arg.aux, input_idxs=tuple(sorted(dedup(g.buf_uop.arg.slot for g in inputs))))))
-pm_rm_rt_uops = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), rm_rt_uops)])
+  lt_srcs = collections.defaultdict(list)
+  for p in lt_patches: lt_srcs[p.buf_uop].append(p)
+  return call.replace(src=(body, *call.src[1:], *[b.after(*ps) for b,ps in lt_srcs.items()], *fills),
+    arg=replace(call.arg, aux=replace(call.arg.aux, input_idxs=tuple(sorted(dedup(b.arg.slot for g in inputs for b in unwrap_mstack(g.buf_uop)))))))
+pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), split_patches)])
 
 # *****************
 
@@ -372,15 +388,15 @@ def callify_hcq(call:UOp, cf:UOp) -> UOp:
 pm_callify_hcq = PatternMatcher([(UPat(Ops.CALL, src=(
   UPat(Ops.CUSTOM_FUNCTION, arg="hcq_args", src=(UPat(Ops.SINK),), name="cf"),), name="call", allow_any_len=True), callify_hcq)])
 
-hcq_compile_cache:dict[tuple[bytes, bool], UOp] = {}
+hcq_compile_cache:dict[bytes, UOp] = {}
 
-@track_rewrites(lambda linear,input_uops,jit,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
-def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None, jit=False) -> UOp:
+@track_rewrites(lambda linear,input_uops,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
+def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None) -> UOp:
   if input_uops is not None:
     slots = {u:i for i,u in reversed(tuple(enumerate(input_uops)))}
     linear = graph_rewrite(linear, pm_replace_buffers, ctx=(input_uops, slots), walk=True, name="replace buffer")
 
-  if (final_linear:=(hcq_compile_cache.get(cache_key:=(linear.key, jit)))) is None:
+  if (final_linear:=(hcq_compile_cache.get(cache_key:=linear.key))) is None:
     # prep
     linear = linear.substitute(back_map:={s.param_like(i): s for i,s in enumerate(input_uops)} if input_uops is not None else {}, walk=True)
     linear = graph_rewrite(linear, pm_insert_copy_staging+pm_flatten_linear, name="insert copy staging")
@@ -391,11 +407,12 @@ def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None, jit=False) -> UOp:
     # lowering to hcq ir
     linear = graph_rewrite(linear, pm_encode_cmdbufs+pm_pack_placeholders, walk=True, name="encode and pack", enter_calls=True)
 
-    # patches
-    linear = graph_rewrite(linear, pm_split_patches+pm_early_simplify+symbolic, ctx=jit, bottom_up=False, name="simplify patches", enter_calls=True)
+    # patches and runtime uops
+    linear = graph_rewrite(linear, pm_early_simplify+symbolic, bottom_up=False, name="simplify patches", enter_calls=True)
+    linear = graph_rewrite(linear, pm_split_patches, walk=True, name="split patches")
 
     # and compile it
-    linear = graph_rewrite(linear, pm_replace_params, bpm=pm_rm_rt_uops, name="replace rt uops and params")
+    linear = graph_rewrite(linear, pm_replace_params, name="replace params")
     final_linear = hcq_compile_cache[cache_key] = graph_rewrite(linear, pm_callify_hcq, name="callify hcq", enter_calls=True)
 
   return final_linear
@@ -455,11 +472,11 @@ pm_assert_no_afters = PatternMatcher([(UPat(Ops.AFTER, name="a"), lambda a: pani
 
 def link_buf_key(a:UOp): return a.key, to_tuple(a.device)
 link_buf_cache:dict[tuple[bytes, tuple[str, ...]], UOp] = {}
-link_linear_cache:dict[tuple[bytes, bool], UOp] = {}
+link_linear_cache:dict[bytes, UOp] = {}
 
-@track_rewrites(lambda _,jit,cache,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}")
-def hcq_link(linear:UOp, jit=False, cache=True) -> UOp:
-  if (linked:=link_linear_cache.get(linear_key:=(linear.key, jit))) is not None: return linked
+@track_rewrites(lambda _,cache,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}")
+def hcq_link(linear:UOp, cache=True) -> UOp:
+  if (linked:=link_linear_cache.get(linear_key:=linear.key)) is not None: return linked
 
   bufs = {(j,i):a for j,c in enumerate(linear.src) for i,a in enumerate(c.src[1:], 1)
           if a.op is Ops.AFTER and unwrap_mstack(a.src[0])[0].tag in HCQ_CACHE_TAGS}
@@ -480,7 +497,10 @@ class HCQ2Compiled(Compiled):
     self.device_id:int = int(device.split(":")[1]) if ":" in device else 0
 
     self.pm_bufferize = PatternMatcher([
-      (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx[0].timeline_signal("sentinel", (1 << 64) - 1)),
+      (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx[0].signal("sentinel", (1 << 64) - 1)),
+      (UPat(Ops.PARAM, tag="timeline_signal"), lambda ctx: ctx[0].signal("timeline")),
+      (UPat(Ops.PARAM, tag="timeline_value"), lambda ctx: ctx[0].signal("value", 1)),
+      (UPat(Ops.PARAM, tag="signal", name="b"), lambda ctx, b: ctx[0].signal(b.arg.slot)),
       (UPat(Ops.PARAM, name="b"), lambda ctx, b: None if b.tag is None else ctx[0].new_buffer(b, cache=ctx[1]))
     ])
 
@@ -495,21 +515,15 @@ class HCQ2Compiled(Compiled):
     return self.rt_buffer.view(b.max_numel(), b.dtype, self.rt_allocator.alloc(b.max_numel() * b.dtype.itemsize, alignment=128))
 
   @functools.cache
-  def timeline_signal(self, queue:str, init_value:int=0) -> Buffer:
+  def signal(self, name:str|int, init_value:int=0) -> Buffer:
     buf = Buffer(self.device, 1, dtypes.uint64, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
     buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = init_value
     return buf
 
-  @functools.cache
-  def timeline_value(self, queue:str, init_value:int=1) -> Buffer:
-    buf = Buffer("CPU", 1, dtypes.uint64, preallocate=True)
-    buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = init_value
-    return buf
-
   def synchronize(self, timeout:int|None=None):
     if not hasattr(self, 'iface'): return
-    sig = self.timeline_signal("COMPUTE:0").as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
-    tl = self.timeline_value("COMPUTE:0").as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
+    sig = self.signal("timeline").as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
+    tl = self.signal("value", 1).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
     st = time.perf_counter()
     while sig[0] < tl[0] - 1:
       if time.perf_counter() - st > (timeout or 3000) / 1000: self.on_device_hang()

tinygrad/schedule/__init__.py

@@ -26,6 +26,15 @@ def _split_after(after: UOp) -> tuple[tuple[UOp, ...], tuple[UOp, ...]]:
     raise AssertionError(f"AFTER source should be CALL, END, STORE, or AFTER, not {invalid[0].op}")
   return tuple(kernels), tuple(deps)
 
+def _kernel_write_targets(k:UOp) -> list[tuple[UOp, UOp]]:
+  # the buffers a kernel stores to, each with the buffer state that the write supersedes (resolved from the call args)
+  call = k.src[0] if k.op is Ops.END else k
+  out: list[tuple[UOp, UOp]] = []
+  for s in call.src[0].toposort():
+    if s.op is Ops.STORE and s.src[0].buf_uop.op is Ops.PARAM and s.src[0].buf_uop.arg.slot >= 0:
+      out.extend((st.buf_uop, st) for st in _states(call.src[s.src[0].buf_uop.arg.slot+1]))
+  return out
+
 def create_schedule(sched_sink:UOp) -> UOp:
   with cpu_profile(TracingKey("toposort sched_sink")):
     # build kernel dependency graph: edges from producer kernel to consumer kernels
@@ -38,7 +47,13 @@ def create_schedule(sched_sink:UOp) -> UOp:
       kernels, after_deps = _split_after(u)
       prev_state = _unwrap_src(u.src[0])
       prev_kernels = set(_split_after(prev_state)[0]) if prev_state.op is Ops.AFTER else set()
-      writes.setdefault(u.buf_uop, []).append((u, prev_state, tuple(k for k in kernels if k not in prev_kernels)))
+      new_kernels = tuple(k for k in kernels if k not in prev_kernels)
+      writes.setdefault(u.buf_uop, []).append((u, prev_state, new_kernels))
+      # a kernel may store to buffers other than the AFTER's base buffer (e.g. state writes packed into another
+      # buffer's AFTER); register those writes under the buffer they actually target so readers get WAR deps
+      for k in new_kernels:
+        for buf, pstate in _kernel_write_targets(k):
+          if buf is not u.buf_uop: writes.setdefault(buf, []).append((u, pstate, (k,)))
       for k in kernels:
         in_degree.setdefault(k, 0)
         if k.op is Ops.END: assert k.src[0].op is Ops.CALL, f"END src[0] should be KERNEL, not {k.src[0].op}"
@@ -196,4 +211,6 @@ def create_linear_with_vars(big_sink:UOp) -> tuple[UOp, dict[str, int]]:
     return UOp(Ops.LINEAR, src=()), var_vals
 
   held_bufs = ({b for b in linear_call.src[1:] if b.op is Ops.BUFFER} if linear_call.op is Ops.CALL else set())
+  # buffers that already hold data can't be suballocated by the memory planner, custom kernels write them in place
+  held_bufs |= {b for b in big_sink.toposort(gate_kernel_sink) if b.op is Ops.BUFFER and b.buffer.is_allocated()}
   return memory_plan_rewrite(linear, held_bufs), var_vals

tinygrad/schedule/allreduce.py

@@ -5,24 +5,27 @@ from tinygrad.uop.ops import UOp
 # *** allreduce implementation ***
 def handle_allreduce(buf:UOp, red:UOp) -> UOp|None:
   if not isinstance(buf.device, tuple): return None
-  assert all_int(buf.shape), f"does not support symbolic shape {buf.shape}"
   ndev, shape, numel = len(buf.device), buf.shape, prod(buf.shape)
   op, device = red.arg
 
   # ring allreduce doesn't provide a benefit with only 2 nodes or where number of elements is less than 256k (empirically)
   # fallback to naive allreduce to save on kernel dispatch, chunking and reassembling chunks.
-  use_all2all = (ALL2ALL >= 2 or (ndev > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and ALL2ALL >= 1))
-  use_ring = not use_all2all and (RING >= 2 or (ndev > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and RING >= 1))
+  concrete = all_int(shape)
+  use_all2all = concrete and (ALL2ALL >= 2 or (ndev > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and ALL2ALL >= 1))
+  use_ring = concrete and not use_all2all and (RING >= 2 or (ndev > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and RING >= 1))
   if DEBUG >= 2: print(f"{'ALL2ALL' if use_all2all else 'RING' if use_ring else 'NAIVE'} ALLREDUCE {ndev}x{numel} | {buf.dtype}")
 
+  if not concrete: buf = buf.pad_to(buf.max_shape)
   # contiguous before we copy it
   buf = buf.contiguous()
 
   # naive: copy to all devices. if you shrink later, that'll be handled
   if not use_ring and not use_all2all:
-    return functools.reduce(lambda x,y: x.alu(op, y), [buf.mselect(i).copy_to_device(device) for i in range(ndev)])
+    out = functools.reduce(lambda x,y: x.alu(op, y), [buf.mselect(i).copy_to_device(device) for i in range(ndev)])
+    return out if concrete else out.shrink_to(shape)
 
   # chunk data into ndev pieces
+  assert isinstance(numel, int)
   factor = next((f for f in [32, 16, 8, 4, 2] if numel % f == 0), 1)
   base, left = divmod(numel // factor,  ndev)
   chunks = list(itertools.pairwise(itertools.accumulate([(base + 1) * factor] * left + [base * factor] * (ndev - left), initial=0)))

tinygrad/schedule/indexing.py

@@ -25,11 +25,21 @@ def realize_store_after_src(ctx:dict[UOp, None], dest:UOp, src:UOp):
   # you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce
   if dest.base in src.backward_slice_with_self: ctx[src] = None
 
+# the inputs of a custom kernel resolve to whole buffers (one per PARAM slot), so they have to materialize.
+# buffer states (AFTER/BUFFER/PARAM) and views of buffers already materialize, anything else must be realized.
+def realize_custom_kernel_srcs(ctx:dict[UOp, None], c:UOp) -> None:
+  for s in c.src[1:]:
+    t = s
+    while t.op in GroupOp.Movement or t.op is Ops.SLICE: t = t.src[0]
+    if t.op not in {Ops.AFTER, Ops.BUFFER, Ops.PARAM, Ops.MSELECT, Ops.MSTACK}: ctx[s] = None
+
 pm_generate_realize_map = PatternMatcher([
   # always realize
   (UPat({Ops.CONTIGUOUS, Ops.STORE}, name="tr"), realize),
   # realize srcs of these
   (UPat((Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
+  # realize the inputs of custom kernel calls
+  (UPat(Ops.CALL, src=(UPat(Ops.SINK),), name="c", allow_any_len=True), realize_custom_kernel_srcs),
   # sometimes we need to realize the src of STORE if there's a self-access
   (UPat(Ops.STORE, src=(UPat.var("dest"), UPat.var("src"))), realize_store_after_src),
 ])

tinygrad/schedule/rangeify.py

@@ -445,7 +445,7 @@ class LocalAddBufferContext:
   opts:tuple|None = None
 
 def debuf(ctx:LocalAddBufferContext, buf:UOp):
-  param = UOp(Ops.PARAM, src=(UOp.const(prod(buf.max_shape), dtypes.int),),
+  param = UOp(Ops.PARAM, src=(UOp.const(prod(buf.max_shape)),),
               arg=ParamArg(ctx.dg, buf.dtype, addrspace=buf.addrspace, device=buf.device))
   ret = param.reshape(buf.max_shape)
   # if the buffer has symbolic shape, shrink the max-sized view to the actual shape

tinygrad/uop/ops.py

@@ -5,12 +5,12 @@ from dataclasses import dataclass, replace
 from enum import Enum, auto
 from tinygrad.uop import Ops, GroupOp
 from tinygrad.dtype import ConstType, dtypes, DType, DTypeLike, truncate, least_upper_dtype, least_upper_float, Invalid, AddrSpace, strong_dtype
-from tinygrad.dtype import ConstFloat, PyConst, InvalidType, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar, weak_dtype
+from tinygrad.dtype import PyConst, InvalidType, weak_dtype, bitcast
 from tinygrad.device import Buffer, MultiBuffer, canonicalize_device, TinyELF
 from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA
 from tinygrad.helpers import PROFILE, dedup, cdiv, cmod, floordiv, floormod, diskcache_put, to_function_name, cpu_profile, TracingKey
 from tinygrad.helpers import VIZ, SPEC, CAPTURE_PROCESS_REPLAY, DISALLOW_BROADCAST, get_shape, fully_flatten, to_tuple
-from tinygrad.helpers import colored, ansilen, printable, Target
+from tinygrad.helpers import colored, ansilen, printable, Target, is_image_shape
 if TYPE_CHECKING:
   from tinygrad.renderer import Estimates
 
@@ -125,13 +125,18 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None:
       # a CALL of an opaque body is void, a CALL of an address can return a value
       return dtypes.void if src[0].dtype is dtypes.void else None
     case Ops.CUSTOM | Ops.CUSTOMI | Ops.PYLITERAL:
-      return dtypes.void
+      return None
     case Ops.INS:
       return None
     case Ops.NOOP:
       # NOOP can be void or carry any dtype (e.g. x.f(Ops.NOOP) or substitute base with NOOP)
       return None
-    case Ops.LOAD | Ops.INDEX | Ops.UNSHARD | Ops.REDUCE | Ops.AFTER | Ops.RANGE | \
+    case Ops.INDEX:
+      # an image access is always float, no matter the storage dtype
+      # TODO: should there be a CAST so src[0].dtype just work?
+      if (b:=src[0]).op is Ops.PARAM and is_image_shape(b.shape): return dtypes.float
+      return b.dtype
+    case Ops.LOAD | Ops.UNSHARD | Ops.REDUCE | Ops.AFTER | Ops.RANGE | \
          Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.COPY | Ops.STAGE | Ops.DETACH | \
          Ops.MSTACK | Ops.MSELECT | Ops.ALLREDUCE | Ops.SPECIAL:
       # pass through first
@@ -272,7 +277,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     return pretty_print(self)
   def argstr(self):
     if self.op is Ops.REDUCE: return f'({", ".join(map(str, self.arg))})'
-    return f"ConstFloat({float.__repr__(self.arg)})" if isinstance(self.arg, ConstFloat) else repr(self.arg)
+    return repr(self.arg)
   def tagstr(self): return f", tag={self.tag}" if self.tag is not None else ""
 
   def f(self, op, **kwargs): return UOp(op, dtype=kwargs.pop("dtype", self.dtype), src=(self,), **kwargs)
@@ -1292,12 +1297,6 @@ def exec_alu(op:Ops, dtype:DType, operands, truncate_output=True):
   if truncate_output and (truncate_fxn:=truncate.get(dtype)) is not None: return truncate_fxn(alu)
   return alu
 
-def bitcast(x, in_dtype:DType, out_dtype:DType):
-  assert in_dtype.itemsize == out_dtype.itemsize, "bitcast itemsize mismatch"
-  packed = struct.pack(storage_fmt_for_dtype(in_dtype), to_storage_scalar(x, in_dtype))
-  out_val = struct.unpack(storage_fmt_for_dtype(out_dtype), packed)[0]
-  return from_storage_scalar(out_val, out_dtype)
-
 # ***** pattern matcher *****
 
 def get_location() -> tuple[str, int]:
@@ -1386,7 +1385,6 @@ class UPat(OpMixin):
   def after(self, *src:UPat, **kwargs): return UPat(Ops.AFTER, self.match_dtype, (self,)+src, **kwargs)
   def end(self, *src:UPat, **kwargs): return UPat(Ops.END, src=(self,)+src, **kwargs)
 
-  def const_like(self, b:ConstLike): return UPat.const(cast(ConstType, b), self.match_dtype)
   def _broadcasted(self, y, reverse=False) -> tuple[UPat, UPat]:
     y = self.ufix(y)
     return (y, self) if reverse else (self, y)
@@ -1746,9 +1744,7 @@ def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=N
 
 def _rebuild_dtype(n:UOp, new_src:tuple[UOp,...]) -> DType:
   # TODO: delete this once the dtype field is removed, every rebuild will re-derive
-  # TODO: these ops keep their stored dtype until dtype_from_uop works
-  if n.op in {Ops.INDEX, Ops.CUSTOM, Ops.CUSTOMI, Ops.PYLITERAL} or \
-     all(a.dtype is b.dtype or b.base.is_invalid for a,b in zip(n.src, new_src)): return n.dtype
+  if all(a.dtype is b.dtype or b.base.is_invalid for a,b in zip(n.src, new_src)): return n.dtype
   return dtype_from_uop(n.op, new_src, n.arg) or n.dtype
 
 def sint_to_uop(x:sint, dtype=dtypes.weakint) -> UOp: return UOp.const(x, dtype)
@@ -1764,7 +1760,7 @@ def lower_weak_node(u:UOp) -> UOp|None:
                     else unwrap(dtype_from_uop(u.op, src, u.arg)))
   return u.replace(dtype=None, src=src[:start]+tuple(s.cast(dt) for s in src[start:])).cast(u.dtype)
 pm_lower_weak = PatternMatcher([
-  (UPat(Ops.CONST, dtype=dtypes.weaks, name="u"), lambda u: u.replace(dtype=select_dtype(u)).cast(u.dtype)),
+  (UPat(Ops.CONST, dtype=dtypes.weaks, name="u"), lambda u: UOp.const(u.val, select_dtype(u)).cast(u.dtype)),
   # two stacked weak casts are a weakint value used as weakfloat (or vice versa): resolve the inner one at the outer kind's default.
   # a SINGLE weak cast is never rewritten here, each consumer absorbs it on its own edge (see lower_weak_srcs)
   (UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat.var("x"),)),), name="u"),

tinygrad/uop/symbolic.py

@@ -1,8 +1,8 @@
 # all of symbolic lives here now
-import math, struct
+import math
 from collections import defaultdict
 from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu
-from tinygrad.dtype import PyConst, ConstType, dtypes, can_lossless_cast, Invalid
+from tinygrad.dtype import PyConst, ConstType, dtypes, can_lossless_cast, Invalid, bitcast
 from tinygrad.helpers import partition, all_same, prod, flatten, unwrap, IMAGE, dedup
 from tinygrad.uop.divandmod import div_and_mod_symbolic
 from tinygrad.uop.movement import mop_cleanup
@@ -13,17 +13,15 @@ from tinygrad.codegen.decomp.transcendental import xpow
 # ******** phase 1 of symbolic used to live in ops, it's the most generic folding rules ********
 
 def simplify_pow(x:UOp, c:UOp) -> UOp|None:
-  if c.val < 0: return x.reciprocal().pow(-c)
+  if c.val < 0: return x.reciprocal().pow(-c.val)
   if c.val == 0: return x.const_like(1)
-  if int(c.val-0.5)+0.5 == c.val: return x.pow(c.const_like(c.val-0.5)) * x.sqrt()
-  if int(c.val) == c.val: return (y := x.pow(c.const_like(c.val//2))) * y * (x if c.val%2 == 1 else 1)
+  if int(c.val-0.5)+0.5 == c.val: return x.pow(c.val-0.5) * x.sqrt()
+  if int(c.val) == c.val: return (y := x.pow(c.val//2)) * y * (x if c.val%2 == 1 else 1)
   return None
 
 def fold_bitcast(root:UOp, c:UOp) -> UOp|None:
-  if (from_fmt:=c.dtype.fmt) is None or (to_fmt:=root.dtype.fmt) is None: return None
-  if c.dtype.itemsize != root.dtype.itemsize: return None
-  def convert(v:ConstType) -> ConstType: return struct.unpack(to_fmt, struct.pack(from_fmt, v))[0]
-  return root.const_like(convert(c.val))
+  if c.dtype.fmt is None or root.dtype.fmt is None or c.dtype.itemsize != root.dtype.itemsize: return None
+  return root.const_like(bitcast(c.val, c.dtype, root.dtype))
 
 def const_arg(u:UOp) -> ConstType|tuple[ConstType, ...]|None:
   if u.op is Ops.CONST: return u.val
@@ -127,7 +125,6 @@ symbolic_simple = pm_data_invalid + PatternMatcher([
   (UPat.var("x") ^ UPat.var("x"), lambda x: x.const_like(0)), # x^x -> 0
   (UPat.var("x") & 0, lambda x: x.const_like(0)), # x&0 -> 0
   # (x&mask)>>k -> x>>k when mask only clears bits below k
-  # TODO: combine this with "# rules for threefry" below
   ((UPat.var("x") & UPat.cvar("mask")) >> UPat.cvar("k"),
    lambda x,mask,k: x >> k.val if mask.val | ((1 << k.val) - 1) == -1 else None),
   ((UPat.var("x") & UPat.cvar("mask")) // UPat.cvar("c"),
@@ -168,13 +165,10 @@ symbolic_simple = pm_data_invalid + PatternMatcher([
   (UPat.var("x").alu(Ops.POW, UPat.cvar("c")), simplify_pow),
   # positive const ** x
   (UPat.cvar("c").alu(Ops.POW, UPat.var("x")), lambda c,x: c if c.val == 1 else (x*math.log2(c.val)).exp2() if c.val > 0 else None),
-  # rules for threefry
-  ((UPat.var('x', dtypes.uint64)&0xFFFFFFFF).cast(dtypes.uint32), lambda x: x.cast(dtypes.uint32)),
-  (((UPat.var(None, dtypes.uint64)*(1<<32)) | UPat.var('y',  dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y),
-  (((UPat.var('x',  dtypes.uint64)*(1<<32)) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))//(1<<32), lambda x: x),
-  (((UPat.var(None, dtypes.uint64)<<32) | UPat.var('y',  dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y),
-  (((UPat.var('x',  dtypes.uint64)<<32) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))//(1<<32), lambda x: x),
-  (((UPat.var('x',  dtypes.uint64)<<32) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))>>32, lambda x: x),
+  # unpack a uint64 packed from two uint32 (threefry)
+  (((UPat.var(None, dtypes.uint64)<<32) | UPat.var('y', dtypes.uint32).cast(dtypes.uint64)).cast(dtypes.uint32), lambda y: y),
+  (((UPat.var('x', dtypes.uint32).cast(dtypes.uint64)<<32) | UPat.var(None, dtypes.uint32).cast(dtypes.uint64))>>32,
+   lambda x: x.cast(dtypes.uint64)),
   # ** simple where folding **
   # a conditional with the same results either way is a noop, also fold const conditionals
   (UPat.var().where(UPat.var("val"), UPat.var("val")), lambda val: val),
@@ -286,9 +280,11 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
   (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x.cast(b.dtype) if can_lossless_cast(x.dtype, a.dtype) else None),
   (UPat.var('x', dtypes.ints+(dtypes.weakint,)).cast(dtypes.ints+(dtypes.weakint,), name="a").cast(name="b"),
     lambda x,a,b: x.cast(b.dtype) if a.dtype.min<=x.vmin and x.vmax<=a.dtype.max else None),
-  # try to do math in int instead of long
+  # try to do math in int instead of long, keep weak const weak
   (UPat(GroupOp.Binary, src=(UPat.var("x", dtypes.long), UPat.var("y", dtypes.long)), name="u"), lambda u,x,y:
-    x.cast(dtypes.int).alu(u.op, y.cast(dtypes.int)).cast(u.dtype) if not any(v.overflows(dtypes.int) for v in (u,x,y)) else None),
+    (UOp.const(x.val) if x.op is Ops.CONST else x.cast(dtypes.int)).alu(u.op,
+     UOp.const(y.val) if y.op is Ops.CONST else y.cast(dtypes.int)).cast(u.dtype)
+    if not any(v.overflows(dtypes.int) for v in (u,x,y)) else None),
   ((UPat.var("x", dtypes.weakint) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+c.cast(cast.dtype)),
   # only RANGE/IF/STORE/KERNEL have side effects
   (UPat(Ops.AFTER, name="x"), lambda x: x.replace(src=(x.src[0],)+

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
@@ -47,7 +47,7 @@ from tinygrad.dtype import dtypes, AddrSpace
 
 uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B",
                Ops.RANGE: "#c8a0e0", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff",
-               Ops.INDEX: "#D8F9E4", Ops.STACK: "#D8F9E4",
+               Ops.INDEX: "#CEF9B7", Ops.STACK: "#D8F9E4",
                Ops.WMMA: "#efefc0", Ops.UNSHARD: "#f6ccff", Ops.INS: "#eec4ff",
                **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
                Ops.SLICE: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.GETADDR: "#9DB1F0", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6",