minor bugfixes

.github/workflows/szdiff.yml

@@ -14,15 +14,12 @@ jobs:
     outputs:
       branchstat: ${{ steps.brstat.outputs.stat}}
     steps:
-      - name: Check code from PR branch
+      - name: Check code from PR branch 
         uses: actions/checkout@v6
         with:
           repository: ${{ github.event.pull_request.head.repo.full_name }}
           ref: ${{ github.event.pull_request.head.sha }}
           fetch-depth: 0
-          # PR code is only inspected with git rev-list, never executed
-          allow-unsafe-pr-checkout: true
-          persist-credentials: false
       - name: Check whether branch is up-to-date
         id: brstat
         run: |
@@ -54,9 +51,6 @@ jobs:
           repository: ${{ github.event.pull_request.head.repo.full_name }}
           ref: ${{ github.event.pull_request.head.sha }}
           path: pr
-          # PR code is only line-counted by master's sz.py, never executed
-          allow-unsafe-pr-checkout: true
-          persist-credentials: false
         # the base default to tinygrad master and cannot be other fork branch for security purpose
       - name: Checkout code from tinygrad master
         uses: actions/checkout@v6

examples/mlperf/model_train.py

@@ -1462,8 +1462,6 @@ 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)
     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)

examples/mlperf/models/flat_llama.py

@@ -37,8 +37,8 @@ def quantize_fp8(x:Tensor, amax_state:Tensor|None=None):
   return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal(), new_amax
 
 def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_scale:Tensor|None=None,
-           x_fp8:Tensor|None=None, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None, x_prequant_mx:tuple|None=None,
-           next_amax_x:Tensor|None=None) -> tuple[Tensor,...]:
+           x_fp8:Tensor|None=None, x_new_amax:Tensor|None=None,
+           grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None, x_prequant_mx:tuple|None=None) -> tuple[Tensor,...]:
   if not fp8:
     if ASM_GEMM:
       from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
@@ -56,14 +56,13 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca
     else:
       x_phys = (x_q.cast(dtypes.bfloat16) * _mx_block_scale(x_e8)).reshape(*l_shape, x_q.shape[-1])
       out = x_phys @ (w.cast(dtypes.bfloat16) * _mx_block_scale(w_inv_scale)).T
-    return out, x_q
+    return out, (amax_x.detach() if amax_x is not None else None), x_q
   if x_fp8 is None:
-    if FUSED_INPUT_QUANTIZE:
+    if FUSED_INPUT_QUANTIZE and amax_x is not None:
       from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed
-      x_fp8, _ = quantize_fp8_delayed(x, amax_x, next_amax_x, FP8_DTYPE)
+      x_fp8, _, x_new_amax, _ = quantize_fp8_delayed(x, amax_x, FP8_DTYPE)
     else:
-      x_fp8, _, new_amax_x = quantize_fp8(x, amax_state=amax_x)
-      next_amax_x.assign(new_amax_x)
+      x_fp8, _, x_new_amax = quantize_fp8(x, amax_state=amax_x)
   if ASM_GEMM:
     from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
     if can_use_asm_gemm(x_fp8, w.T):
@@ -74,51 +73,51 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca
       else:
         out = asm_gemm(x_fp8, w.T, x_scale=amax_x, w_scale=w_inv_scale, grad_amax_state=grad_amax_state,
                        next_grad_amax_state=next_grad_amax_state)
-      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
+      return out, x_new_amax, 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_new_amax, 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):
+                         grad_amax_state:Tensor, next_grad_amax_state:Tensor):
   if FUSED_ADD_NORM_MUL_QUANTIZE:
     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)
-    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
+    x_fp8, new_amax, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE)
+    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x, x_new_amax=new_amax,
                        grad_amax_state=grad_amax_state, next_grad_amax_state=next_grad_amax_state)
     return out, x_normed, rrms, ret
   x_normed, rrms = rmsnorm(x, eps)
   out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state,
-                     next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
+                     next_grad_amax_state=next_grad_amax_state)
   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):
+                             grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
   if FUSED_ADD_NORM_MUL_QUANTIZE:
     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)
-    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
+    x_fp8, new_amax, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE)
+    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x, x_new_amax=new_amax,
                        grad_amax_state=grad_amax_state, next_grad_amax_state=next_grad_amax_state)
     return out, h, x_normed, rrms, ret
   h = x + residual
   x_normed, rrms = rmsnorm(h, eps)
   out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state,
-                     next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
+                     next_grad_amax_state=next_grad_amax_state)
   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,
+                             amax_x2:Tensor,
                              grad_amax_xw13:Tensor, next_grad_amax_xw13:Tensor,
                              grad_amax_xout:Tensor, next_grad_amax_xout:Tensor):
   if FUSED_SILU_W13:
     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,
-                                                 next_grad_amax_state=next_grad_amax_xw13, amax_out=next_amax_x2)
-    out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, amax_x=amax_x2,
+    x2_fp8, new_amax_x2 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13,
+                                                 next_grad_amax_state=next_grad_amax_xw13)
+    out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, amax_x=amax_x2, x_new_amax=new_amax_x2,
                        grad_amax_state=grad_amax_xout, next_grad_amax_state=next_grad_amax_xout)
     return out, ret
   hidden = x_w13.shape[-1] // 2
   x_w1, x_w3 = x_w13[..., :hidden], x_w13[..., hidden:]
   out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout,
-                     next_grad_amax_state=next_grad_amax_xout, next_amax_x=next_amax_x2)
+                     next_grad_amax_state=next_grad_amax_xout)
   return out, ret
 
 class FlatTransformer:
@@ -187,14 +186,14 @@ class FlatTransformer:
 
   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):
     bsz, seqlen, _ = x.shape
-    saves = []
+    amaxs, saves = [], []
 
-    xqkv, x_normed, rrms, s = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
+    xqkv, x_normed, rrms, (new_amax, *s) = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
                                                                   amax_x=amax_xqkv, grad_amax_state=grad_amax_xqkv,
-                                                                  next_grad_amax_state=next_grad_amax_xqkv, next_amax_x=next_amax_xqkv)
+                                                                  next_grad_amax_state=next_grad_amax_xqkv)
+    amaxs.append(new_amax)
     saves.extend([x_normed, rrms, *s, xqkv])
     if getenv("HK_FLASH_ATTENTION"):
       from extra.thunder.amd.fa import flash_attention, fused_qkv_rope
@@ -212,62 +211,64 @@ class FlatTransformer:
       attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
     attn = attn.reshape(bsz, seqlen, -1)
 
-    out, *s = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo,
-                               next_grad_amax_state=next_grad_amax_xo, next_amax_x=next_amax_xo)
+    out, new_amax, *s = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo,
+                               next_grad_amax_state=next_grad_amax_xo)
+    amaxs.append(new_amax)
     saves.extend([*s, out])
-    return out, saves
+    return out, amaxs, saves
 
   def feed_forward(self, x:Tensor, residual:Tensor, **kwargs):
-    saves = []
+    amaxs, saves = [], []
 
     if SPLIT_W13:
       h = x + residual
       x_normed, rrms = rmsnorm(h, self.norm_eps)
       saves.extend([x_normed, rrms])
       inp = x_normed * kwargs["ffn_norm"]
-      x_w1, *s = matmul(inp, kwargs["w1"], amax_x=kwargs["amax_x1"], w_inv_scale=kwargs["s_1"],
-                                  grad_amax_state=kwargs["grad_amax_xw1"], next_grad_amax_state=kwargs["next_grad_amax_xw1"],
-                                  next_amax_x=kwargs["next_amax_x1"])
+      x_w1, new_amax, *s = matmul(inp, kwargs["w1"], amax_x=kwargs["amax_x1"], w_inv_scale=kwargs["s_1"],
+                                  grad_amax_state=kwargs["grad_amax_xw1"], next_grad_amax_state=kwargs["next_grad_amax_xw1"])
+      amaxs.append(new_amax)
       saves.extend([*s, x_w1])
-      x_w3, *s = matmul(inp, kwargs["w3"], amax_x=kwargs["amax_x3"], w_inv_scale=kwargs["s_3"],
-                                  grad_amax_state=kwargs["grad_amax_xw3"], next_grad_amax_state=kwargs["next_grad_amax_xw3"],
-                                  next_amax_x=kwargs["next_amax_x3"])
+      x_w3, new_amax, *s = matmul(inp, kwargs["w3"], amax_x=kwargs["amax_x3"], w_inv_scale=kwargs["s_3"],
+                                  grad_amax_state=kwargs["grad_amax_xw3"], next_grad_amax_state=kwargs["next_grad_amax_xw3"])
+      amaxs.append(new_amax)
       saves.extend([*s, x_w3])
       if FUSED_SILU_W13 and MXFP8:
         from extra.llama_kernels.fused_silu_mul_quantize_mxfp8 import fused_silu_mul_quantize_mxfp8
         aq, ae8, asi = fused_silu_mul_quantize_mxfp8(x_w1.reshape(-1, x_w1.shape[-1]), x_w3.reshape(-1, x_w3.shape[-1]))
-        out, *s = matmul(None, kwargs["w2"], x_prequant_mx=(aq, ae8, asi), amax_x=kwargs["amax_x2"],
-                         w_inv_scale=kwargs["s_2"], grad_amax_state=kwargs["grad_amax_xout"],
-                         next_grad_amax_state=kwargs["next_grad_amax_xout"], next_amax_x=kwargs["next_amax_x2"])
+        out, new_amax, *s = matmul(None, kwargs["w2"], x_prequant_mx=(aq, ae8, asi), amax_x=kwargs["amax_x2"],
+                                   w_inv_scale=kwargs["s_2"], grad_amax_state=kwargs["grad_amax_xout"],
+                                   next_grad_amax_state=kwargs["next_grad_amax_xout"])
         out = out.reshape(*x_w1.shape[:-1], kwargs["w2"].shape[0])
       else:
-        out, *s = matmul(x_w1.silu() * x_w3, kwargs["w2"], amax_x=kwargs["amax_x2"], w_inv_scale=kwargs["s_2"],
-                         grad_amax_state=kwargs["grad_amax_xout"], next_grad_amax_state=kwargs["next_grad_amax_xout"],
-                         next_amax_x=kwargs["next_amax_x2"])
+        out, new_amax, *s = matmul(x_w1.silu() * x_w3, kwargs["w2"], amax_x=kwargs["amax_x2"], w_inv_scale=kwargs["s_2"],
+                                   grad_amax_state=kwargs["grad_amax_xout"], next_grad_amax_state=kwargs["next_grad_amax_xout"])
+      amaxs.append(new_amax)
       saves.extend([*s, out])
     else:
-      x_w13, h, x_normed, rrms, s = add_norm_quantize_matmul(x, residual, kwargs["ffn_norm"], kwargs["w13"], kwargs["s_13"],
+      x_w13, h, x_normed, rrms, (new_amax, *s) = add_norm_quantize_matmul(x, residual, kwargs["ffn_norm"], kwargs["w13"], kwargs["s_13"],
                                                                           self.norm_eps, amax_x=kwargs["amax_x13"],
-                                                                          next_amax_x=kwargs["next_amax_x13"],
                                                                           grad_amax_state=kwargs["grad_amax_xw13"],
                                                                           next_grad_amax_state=kwargs["next_grad_amax_xw13"])
+      amaxs.append(new_amax)
       saves.extend([x_normed, rrms, *s, x_w13])
-      out, s = silu_w13_quantize_matmul(x_w13, kwargs["w2"], kwargs["s_2"], amax_x2=kwargs["amax_x2"],
-                                                     next_amax_x2=kwargs["next_amax_x2"],
+      out, (new_amax, *s) = silu_w13_quantize_matmul(x_w13, kwargs["w2"], kwargs["s_2"], amax_x2=kwargs["amax_x2"],
                                                      grad_amax_xw13=kwargs["grad_amax_xw13"],
                                                      next_grad_amax_xw13=kwargs["next_grad_amax_xw13"],
                                                      grad_amax_xout=kwargs["grad_amax_xout"],
                                                      next_grad_amax_xout=kwargs["next_grad_amax_xout"])
+      amaxs.append(new_amax)
       saves.extend([*s, out])
-    return out, h, saves
+    return out, h, amaxs, saves
 
   @function(precompile=True, precompile_backward=True)
   def run_layer(self, x:Tensor, freqs_cis:Tensor, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
-    attn, attn_saves = self.attention(x, freqs_cis, **attn_kwargs)
-    ffn, h, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs)
+    attn, attn_amaxs, attn_saves = self.attention(x, freqs_cis, **attn_kwargs)
+    ffn, h, ffn_amaxs, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs)
     h = h + ffn
-    if save: return (h, *attn_saves, *ffn_saves)
-    else: return (h,)
+    amaxs = tuple(a.detach() for a in (*attn_amaxs, *ffn_amaxs))
+    if save: return (h, *amaxs, *attn_saves, *ffn_saves)
+    else: return (h, *amaxs)
 
   def shard(self, device:tuple[str, ...], mp:bool=False):
     from tinygrad.nn.state import get_parameters
@@ -318,21 +319,21 @@ class FlatTransformer:
     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])
+                        amax_x2=a["x2"][i], s_2=s["w2"][i], grad_amax_xout=ga["xout"][i], next_grad_amax_xout=nga["xout"][i])
       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])
       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])
-      h, *_ = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)
+                          next_grad_amax_xw13=nga["xw13"][i])
+      h, *ret = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)
+      amax_names = ["xqkv", "xo"] + (["x1", "x3"] if SPLIT_W13 else ["x13"]) + ["x2"]
+      for name, new_val in zip(amax_names, ret[:len(amax_names)]):
+        na[name][i].assign(new_val)
 
     logits = matmul(self.norm(h), self.output[0], fp8=False)[0]
     return logits
@@ -415,9 +416,6 @@ 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)
       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/optim.py

@@ -96,7 +96,7 @@ class GradAccClipAdamW(Optimizer):
     up = up.float().shard_like(w) + self.lr.to(w.device) * wd * w.detach()
     new_w = w.detach() - up
     if master is not None: master.assign(new_w)
-    if self.zero and not (MXFP8 and t.dtype in dtypes.fp8s): new_w = self._zero_gather(new_w)
+    if self.zero: new_w = self._zero_gather(new_w)
     # when master is offloaded to a different device than the param, results are resharded back onto the param's (sharded) device
     offloaded = master is not None and master.device != t.device
     if STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16:
@@ -106,7 +106,6 @@ class GradAccClipAdamW(Optimizer):
       if MXFP8:
         from extra.gemm.cdna_asm_gemm import quantize_mxfp8
         w_q, w_e8, _ = quantize_mxfp8(new_w.reshape(-1, new_w.shape[-1]))
-        if self.zero: w_q, w_e8 = self._zero_gather(w_q), self._zero_gather(w_e8)
         new_e8 = w_e8.reshape(t._inv_scale.shape)
         t._inv_scale.assign(new_e8.shard_like(t._inv_scale) if offloaded else new_e8)
         ret = w_q.reshape(new_w.shape)

extra/benchmark_llm.py

@@ -0,0 +1,93 @@
+#!/usr/bin/env python3
+"""Benchmark tinygrad LLM prefill and decode independently.
+
+Examples:
+  python -m extra.benchmark_llm --model qwen3:0.6b --max-context 32768
+  python -m extra.benchmark_llm --model /path/to/model.gguf --prompt-tokens 8192 --realize
+"""
+from __future__ import annotations
+
+import argparse, json, statistics, time
+from dataclasses import asdict, dataclass
+
+from tinygrad.helpers import fetch
+from tinygrad.llm.cli import models
+from tinygrad.llm.model import Transformer
+
+
+@dataclass
+class Result:
+  prompt_tokens: int
+  decode_tokens: int
+  time_to_first_token_s: float
+  prefill_tokens_per_s: float
+  decode_tokens_per_s: float
+  decode_p50_ms: float
+  decode_p95_ms: float
+
+
+def percentile(values:list[float], percentile:float) -> float:
+  ordered = sorted(values)
+  return ordered[round((len(ordered) - 1) * percentile)]
+
+
+def synthetic_prompt(length:int, vocab_size:int, salt:int) -> list[int]:
+  # Avoid tokenizer and chat-template work while exercising the same embedding/model path.
+  # Changing token zero guarantees that Transformer.get_start_pos cannot reuse an earlier KV cache.
+  assert length > 0 and vocab_size > 256
+  return [256 + salt % (vocab_size - 256)] + [256 + (i * 7919) % (vocab_size - 256) for i in range(1, length)]
+
+
+def benchmark(model:Transformer, prompt:list[int], decode_tokens:int, chunk_size:int) -> Result:
+  gen = model.generate(prompt.copy(), chunk_size=chunk_size)
+  begin = time.perf_counter()
+  next(gen)
+  ttft = time.perf_counter() - begin
+
+  decode_times: list[float] = []
+  for _ in range(decode_tokens):
+    begin = time.perf_counter()
+    next(gen)
+    decode_times.append(time.perf_counter() - begin)
+
+  return Result(len(prompt), decode_tokens, ttft, len(prompt) / ttft, decode_tokens / sum(decode_times),
+                statistics.median(decode_times) * 1e3, percentile(decode_times, 0.95) * 1e3)
+
+
+def main() -> None:
+  parser = argparse.ArgumentParser(description="Measure LLM prefill and steady-state decode speed")
+  parser.add_argument("--model", default="qwen3:0.6b", help="Model preset or local GGUF path")
+  parser.add_argument("--max-context", type=int, default=32768)
+  parser.add_argument("--prompt-tokens", type=int, nargs="+", default=[128, 2048, 8192])
+  parser.add_argument("--decode-tokens", type=int, default=32)
+  parser.add_argument("--chunk-size", type=int, default=256)
+  parser.add_argument("--realize", action="store_true", help="Unpack model weights once at load time")
+  parser.add_argument("--json", action="store_true", help="Print machine-readable results")
+  args = parser.parse_args()
+
+  if args.decode_tokens < 1: parser.error("--decode-tokens must be positive")
+  if args.chunk_size < 1: parser.error("--chunk-size must be positive")
+  if max(args.prompt_tokens) + args.decode_tokens >= args.max_context:
+    parser.error("prompt plus decode tokens must fit within --max-context")
+
+  path = fetch(models.get(args.model, args.model))
+  model, kv = Transformer.from_gguf(path, args.max_context, realize=args.realize)
+  vocab_size = len(kv["tokenizer.ggml.tokens"])
+
+  model.warmup(args.chunk_size)
+
+  results = [benchmark(model, synthetic_prompt(n, vocab_size, salt=i+1), args.decode_tokens, args.chunk_size)
+             for i, n in enumerate(args.prompt_tokens)]
+  if args.json:
+    print(json.dumps({"model": args.model, "max_context": args.max_context, "chunk_size": args.chunk_size,
+                      "realize": args.realize, "results": [asdict(x) for x in results]}, indent=2))
+    return
+
+  print(f"model={args.model} max_context={args.max_context} chunk_size={args.chunk_size} realize={args.realize}")
+  print(f"{'prompt':>8} {'TTFT':>10} {'prefill':>14} {'decode':>14} {'decode p50':>12} {'decode p95':>12}")
+  for result in results:
+    print(f"{result.prompt_tokens:8d} {result.time_to_first_token_s:9.3f}s {result.prefill_tokens_per_s:11.1f} t/s "
+          f"{result.decode_tokens_per_s:11.1f} t/s {result.decode_p50_ms:9.2f} ms {result.decode_p95_ms:9.2f} ms")
+
+
+if __name__ == "__main__": main()

extra/gemm/amd_flash_attention.py

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

extra/gemm/cdna_asm_gemm.py

@@ -280,7 +280,10 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp, n_scales:int=2, has_grad_amax:bool=
       elif getenv("FUSED_GRAD_QUANTIZE", 0):
         grad_amax_t = Tensor(grad_amax_state, device=a.device)
         g_amax = grad_amax_t
-        g_fp8, _ = quantize_fp8_delayed(g_t, g_amax, Tensor(next_grad_amax_state, device=a.device))
+        g_fp8, _, new_grad_amax, _ = quantize_fp8_delayed(g_t, g_amax)
+        store_effect = next_grad_amax_state.store(new_grad_amax.uop)
+        assert g_fp8.uop.op is Ops.AFTER, f"expected AFTER, got {g_fp8.uop.op}"
+        g_fp8 = Tensor(g_fp8.uop.replace(src=g_fp8.uop.src + (store_effect,)), device=a.device)
       else:
         grad_amax_t = Tensor(grad_amax_state, device=a.device)
         g_amax = grad_amax_t

extra/hcq2/hcq2.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 from typing import cast, Callable, TypeVar, Generic, Any
 import struct, functools, time, collections, itertools
 from dataclasses import replace, dataclass
-from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, dedup, pluralize, JIT_BATCH_SIZE
+from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, dedup, pluralize
 from tinygrad.helpers import to_tuple, round_up, partition, data64_le, panic, ContextVar
 from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer
 from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, track_rewrites, GroupOp
@@ -102,138 +102,159 @@ def stage_copy(dst:UOp, src:UOp) -> UOp|None:
 pm_insert_copy_staging = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy)])
 
 # *****************
-# 2. deps
+# 2.1. tag hcq calls
+
+def tag_hcq_call(ctx:itertools.count, call:UOp) -> UOp:
+  if (hcq_devs:=next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None)) is None: return call
+
+  queue = "COMPUTE:0" if call.src[0].op is Ops.PROGRAM else "COPY:0"
+  info = HCQInfo(get_call_name(call, get_call_arg_uops(call)), estimate_uop(call), to_tuple(hcq_devs), queue)
+  return call.replace(arg=replace(call.arg, aux=info)).rtag(next(ctx))
+pm_tag_hcq_calls = PatternMatcher([(UPat(Ops.LINEAR, name="l"), lambda ctx, l: l.replace(src=tuple(tag_hcq_call(ctx, s) for s in l.src)))])
+
+# *****************
+# 2.2. deps tracking
+# device.timeline_signal/value are the per-device schedule epoch. Before a schedule queue accesses memory owned by device N for the first time,
+# it waits for device[N].timeline_signal >= device[N].timeline_value - 1. This orders the schedule after all prior schedules that touched device N.
+#
+# queue.timeline_signal/value are per-queue progress counters used only inside a schedule.
+# Only the owner queue signals its queue.timeline_signal. Values are monotonic.
+#
+# At schedule end, one finalizer queue per touched device[N] waits for every active queue on device[N] to reach its schedule-local
+# final queue.timeline value, then signals device[N].timeline_signal with the schedule's reserved device epoch. After that, buffers/transients
+# for device N from this schedule are safe for the next schedule
+#
+# C programs reserve and bump timeline values, then patch command buffers with the concrete wait/signal values.
 
 class HCQDepsTracker(DepsTracker):
   @staticmethod
   def _key(buf:Any) -> tuple[Any, int, int]:
     return (buf.arg.slot, 0, buf.max_numel() * buf.dtype.itemsize) if isinstance(buf, UOp) else DepsTracker._key(buf)
 
-def _get_call_bufs_by_lane(call:UOp, devices:tuple[str, ...]) -> list[list[Any]]:
+def make_deps(u:UOp, dep_lanes:list[tuple[UOp, int, int]], nlanes:int) -> UOp:
+  deps:dict[UOp, list[int|None]] = collections.defaultdict(lambda: [None]*nlanes)
+  for dep, dlane, lane in dep_lanes: deps[dep][lane] = dlane
+  return u.after(*deps, arg=tuple(tuple(v) for v in deps.values()))
+
+def sched_sync(ctx:DepsTracker, call:UOp) -> UOp|None:
+  if not isinstance(call.arg.aux, HCQInfo): return None
+
   refs = get_call_arg_uops(call)
-  return [[b if b.op is Ops.PARAM else mb.bufs[lane] if isinstance(mb:=b.buffer, MultiBuffer) else mb for b in refs] for lane in range(len(devices))]
+  outs, _ = get_call_outs_ins(call)
+  devices, queue = call.arg.aux.device, call.arg.aux.queue
 
-def _get_deps(ctx:DepsTracker, bufs_by_lane:list[list[Any]], write, key:tuple[tuple[str, ...], str, int]) -> list[tuple[tuple, int, int]]:
-  dep_lanes:list[tuple[tuple, int, int]] = []
-  for lane, bufs in enumerate(bufs_by_lane):
-    dep_lanes += [(dep, dlane, lane) for dep, dlane in ctx.access_resources(bufs, write if write is not None else range(len(bufs)), (key, lane))]
-  return dep_lanes
+  dep_lanes:list[tuple[UOp, int, int]] = []
+  for lane, d in enumerate(devices):
+    lane_refs = [b if b.op is Ops.PARAM else mb.bufs[lane] if isinstance(mb:=b.buffer, MultiBuffer) else mb for b in refs]
+    for dep, dlane in ctx.access_resources(lane_refs, outs, (call, lane)): dep_lanes.append((dep, dlane, lane))
 
-def _build_wait_cmds(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)]
+    dep_lanes = [(dep, dlane, lane) for dep, dlane, lane in dep_lanes if (dep.arg.aux.device[dlane], dep.arg.aux.queue) != (devices[lane], queue)]
 
-  # opt2: keep latest dep per (dep device, queue, cur lane)
-  latest = {((dep[0][dlane], dep[1]), lane): (dep, dlane) for dep, dlane, lane in sorted(dep_lanes, key=lambda x: x[0][2])}
-  deps:dict[tuple, list[int|None]] = collections.defaultdict(lambda: [None]*len(devices))
-  for (_, lane), (dep, dlane) in latest.items(): deps[dep][lane] = dlane
-
-  waits = []
-  for (ddevs, dqueue, dtag), lanes in deps.items():
-    sig = make_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 = make_mstack([make_signal_value(d if dl is None else ddevs[dl], queue=dqueue) for dl, d in zip(lanes, devices)])
-    waits.append((sig.index(zero:=UOp.const(dtypes.int, 0)).load() >= val.index(zero) + dtag).wait())
-  return waits, {dtag for _, _, dtag in deps}
-
-def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[tuple[tuple[str, ...], str]],
-                      tracker:HCQDepsTracker) -> tuple[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, finalizers, waited = UOp.const(dtypes.int, 0), len(batch_info), [], set()
-  for _, devgroup in itertools.groupby(sorted(dedup([d for devs, _ in batch_info for d in devs])), 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")
-    waited |= cur_waited
-
-    # wait the syncs, store the device epoch; value bumps are a separate call: no lane may bump until every lane has patched its waits
-    submit = make_submit(*waits, make_signal(devs).store((tl:=make_signal_value(devs)).index(zero)), devs=devs, queue="COMPUTE:0")
-    upd = [(tl, 1)] + [(make_signal_value(devs, queue=qn), n) for qn in dedup([qn for bdevs, qn in batch_info if set(bdevs) & set(devs)])]
-    bump = UOp.barrier(*[s.index(zero, dtype=s.dtype).store(s.index(zero) + inc) for s, inc in upd])
-    finalizers += [UOp.custom_function("hcq", b.sink()).call(aux=HCQInfo("hcq_finalizer", Estimates(), devs, "COMPUTE:0")) for b in (submit, bump)]
-  return finalizers, 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()
-  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)
-    call_waits.append(cmds)
-    waited |= cur_waited
-
-  # build finalizers
-  finalizers, finalizer_waited = _build_finalizers(batch, batch_info, deps_tracker)
-  waited |= finalizer_waited
-
-  src = []
-  for tag, ((call, _), (devices, queue), cmds) in enumerate(zip(batch, batch_info, call_waits)):
-    # first queue use, sync prior device work with main signal
-    if batch_info.index((devices, queue)) == tag:
-      epoch = (make_signal(devices).index(0).load() >= make_signal_value(devices).index(0) - 1).wait()
-      cmds = [UOp(Ops.BARRIER), epoch] + cmds
-
-    # signal queue timeline if someone waits for us
-    store = make_signal(devices, queue=queue).store(make_signal_value(devices, queue=queue).index(0) + tag) if tag in waited else None
-
-    # and make hcq call
-    info = HCQInfo(get_call_name(call, get_call_arg_uops(call)), estimate_uop(call), devices, queue)
-    cmds = [*cmds, call.replace(arg=replace(call.arg, aux=info))] + ([store] if store is not None else [])
-    src.append(UOp.custom_function("hcq", make_submit(*cmds, devs=devices, queue=queue).sink()).call(name="hcq", aux=info))
-  return src + finalizers
-
-def sched_hcq_batches(l:UOp) -> UOp:
-  srcs:list[UOp] = []
-  batch:list[tuple[UOp, tuple[str, ...]]] = []
-  for call in l.src:
-    if (devs:=next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None)) is not None: batch.append((call, to_tuple(devs)))
-    else: srcs, batch = srcs + _finalize_batch(batch) + [call], []
-  return l.replace(src=tuple(srcs + _finalize_batch(batch)))
-pm_sched_hcq_batches = PatternMatcher([(UPat(Ops.LINEAR, name="l"), sched_hcq_batches)])
+  # keep latest dep per (dep device, queue, cur lane)
+  latest = {((dep.arg.aux.device[dlane], dep.arg.aux.queue), lane): (dep, dlane) for dep, dlane, lane in sorted(dep_lanes, key=lambda x: x[0].tag)}
+  return make_deps(call, [(dep, dlane, lane) for (_, lane), (dep, dlane) in latest.items()], len(devices))
+pm_sched_sync = PatternMatcher([(UPat(Ops.CALL, name="call"), sched_sync)])
 
 # *****************
-# 3. merge into queues
+# 2.3. merge into queues
 
-def _merged_hcq_call(calls:list[UOp]) -> UOp: # TODO: simplify?
-  if len(calls) == 1: return calls[0]
-  info = replace(calls[0].arg.aux, name=f"submit {calls[0].arg.aux.queue} ({len(calls)})",
-                 estimates=sum((c.arg.aux.estimates for c in calls), start=Estimates()))
-  cmds = [cmd for c in calls for cmd in get_submit(c).src[0].src]
-  return UOp.custom_function("hcq", make_submit(*cmds, devs=info.device, queue=info.queue).sink()).call(name="hcq", aux=info)
+def _merged_hcq_call(calls:list[UOp]):
+  info = replace(unwrap_after(calls[0]).arg.aux, estimates=sum((unwrap_after(c).arg.aux.estimates for c in calls), start=Estimates()))
+  cmdbuf = make_submit(*calls, devs=info.device, queue=info.queue)
+  return UOp.custom_function("hcq", cmdbuf.sink()).call(name="hcq", aux=info)
 
 def merge_queues(linear:UOp) -> UOp:
   new_src:list[UOp] = []
-  opened_qs:dict[tuple[tuple[str, ...], str], list[UOp]] = {} # (devs, queue) -> list of hcq calls, kept in submit order
-  limits = collections.defaultdict(lambda: JIT_BATCH_SIZE.value)
+  opened_qs:dict[tuple[tuple[str, ...], str], list[UOp]] = {} # (devs, queue) -> list of calls, kept in submit order
 
   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(unwrap_after(call).arg.aux, HCQInfo):
       new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in list(opened_qs)] + [call]
       continue
 
-    if (old:=opened_qs.pop(key:=(info.device, info.queue), None)) is not None:
-      if limits[key] and len(old) >= limits[key]: new_src, old, limits[key] = new_src + [_merged_hcq_call(old)], [], limits[key] * 2
-      new_rec = old + [call]
+    devices, queue = unwrap_after(call).arg.aux.device, unwrap_after(call).arg.aux.queue
+
+    if (old:=opened_qs.pop((devices, queue), None)) is not None: new_rec = old + [call]
     else:
       # no such queue opened: close every open submit on this queue that shares a device, so submit order is kept
-      closing = [k for k in opened_qs if k[1] == info.queue and set(k[0]) & set(info.device)]
+      closing = [k for k in opened_qs if k[1] == queue and set(k[0]) & set(devices)]
       new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in closing]
       new_rec = [call]
-    opened_qs[(info.device, info.queue)] = new_rec
+    opened_qs[(devices, queue)] = new_rec
   return linear.replace(src=tuple(new_src + [_merged_hcq_call(c) for c in opened_qs.values()]))
 pm_merge_queues = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), merge_queues)])
 
+# *****************
+# 2.4. finalizer
+
+def add_finalizer(ctx:itertools.count, linear:UOp) -> UOp:
+  # collect by device type
+  parts:dict[str, list[UOp]] = collections.defaultdict(list)
+  for call in linear.src:
+    if (c:=unwrap_after(call)).src[0].op is not Ops.CUSTOM_FUNCTION or c.src[0].arg != "hcq": continue
+    parts[c.arg.aux.device[0].split(':')[0]].append(unwrap_after(get_submit(call).src[0].src[0]))
+
+  nbump = next(ctx)
+  finalizers = []
+  for calls in parts.values():
+    devs = tuple(dedup(d for call in calls for d in unwrap_after(call).arg.aux.device))
+    zero = UOp.const(dtypes.int, 0)
+    tl = make_signal_value(devs)
+
+    # split each (multi-device) call into per-device deps, then store the device timeline value into the device signal after them
+    dep_lanes = [(call, dlane, devs.index(d)) for call in calls for dlane, d in enumerate(unwrap_after(call).arg.aux.device)]
+    store = make_deps(make_signal(devs).store(tl.index(zero)), dep_lanes, len(devs))
+    submit = make_submit(store, devs=devs, queue="COMPUTE:0")
+
+    upd = [(tl, 1)] + [(make_signal_value(devs, queue=qn), nbump) for qn in dedup([unwrap_after(call).arg.aux.queue for call in calls])]
+    patches = [s.after(submit).index(zero, dtype=s.dtype).store(s.index(zero) + inc) for s, inc in upd]
+    finalizers.append(UOp.custom_function("hcq", UOp.barrier(*patches).sink()).call(aux=HCQInfo("hcq finalizer", Estimates(), devs, "COMPUTE:0")))
+  return linear.replace(src=linear.src + tuple(finalizers))
+pm_add_finalizer = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), add_finalizer)])
+
+# *****************
+# 2.5. global sync
+
+def add_global_sync(ctx:set[tuple[str, ...]], submit:UOp, q:UOp) -> UOp|None:
+  if (devs:=q.arg[0]) in ctx: return None
+  ctx.add(devs)
+
+  # some devices from a command buffer might be used for the first time this schedule, so we wait for their global timeline epoch.
+  wait = (make_signal(devs).index(zero:=UOp.const(dtypes.int, 0)).load() >= make_signal_value(devs).index(zero) - 1).wait()
+  return submit.replace(src=(q.replace(src=(UOp(Ops.BARRIER, dtypes.void), wait, *q.src)),))
+pm_add_global_sync = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),), name="submit"), add_global_sync)])
+
+# *****************
+# 3.1. lower loads/stores
+
+def add_loads(ctx:set[int], submit:UOp, q:UOp) -> UOp|None:
+  cur_devs = q.arg[0]
+  new_src:list[UOp] = []
+  for s in q.src:
+    if s.op is Ops.AFTER:
+      for lanes, dep in zip(s.arg, s.src[1:]):
+        devs, queue = dep.arg.aux.device, dep.arg.aux.queue
+        ctx.add(dep.tag) # mark op to update signal.
+
+        sig = make_mstack([make_signal(d if dl is None else devs[dl], queue=queue, sentinel=dl is None) for dl, d in zip(lanes, cur_devs)])
+        val = make_mstack([make_signal_value(d if dl is None else devs[dl], queue=queue) for dl, d in zip(lanes, cur_devs)]).index(UOp.const(dtypes.int, 0))
+        new_src.append((sig.index(UOp.const(dtypes.int, 0)).load() >= val + dep.tag).wait())
+      s = s.src[0]
+    new_src.append(s)
+  return submit.replace(src=(q.replace(src=tuple(new_src)),))
+pm_add_inner_loads = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),), name="submit"), add_loads)])
+
+def add_stores(ctx:set[int], submit:UOp, q:UOp) -> UOp|None:
+  devs, queue = q.arg
+  new_src:list[UOp] = []
+  for op in q.src:
+    new_src.append(op)
+    if (sigval:=unwrap_after(op).tag) in ctx:
+      new_src.append(make_signal(devs, queue=queue).store(make_signal_value(devs, queue=queue).index(UOp.const(dtypes.int, 0)) + sigval))
+  return submit.replace(src=(q.replace(src=tuple(new_src)),))
+pm_add_inner_stores = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),), name="submit"), add_stores)])
+
 # *****************
 # 4.1. hcq lowering: programs
 
@@ -258,7 +279,7 @@ def is_value_known_at_link(val:UOp) -> bool:
   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)
+  return 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:
   store = p.src[0] if (is_binary_patch:=p.op is Ops.END) else p
@@ -287,7 +308,7 @@ pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION
 def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[dict[UOp, UOp], tuple[UOp, ...]]:
   bare = {g: g.replace(src=(unwrap_after(g.src[0]),)) for g in gaddrs}
 
-  order = sorted(dedup(bare.values()), key=lambda g: ((b:=unwrap_mstack(g.buf_uop)[0]).arg.slot, repr(b.tag)))
+  order = sorted(dedup(bare.values()), key=lambda g: ((b:=unwrap_mstack(g.buf_uop)[0]).arg.slot, to_tuple(b.tag)))
   slots, table = {g:i for i,g in enumerate(order)}, make_placeholder(call.arg.aux.device, len(order), dtypes.uint64, name)
 
   reads = {g: table.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(UOp.const(dtypes.int, slots[bare[g]])).load() for g in gaddrs}
@@ -370,16 +391,19 @@ def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None, jit=False) -> UOp:
   if input_uops is not None: linear = graph_rewrite(linear, pm_replace_buffers, ctx=input_uops, walk=True, enter_calls=True, name="replace buffer")
 
   if (final_linear:=(hcq_compile_cache.get(cache_key:=(linear.key, jit)))) is None:
-    # prep
+    # schedule
     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")
-
-    # schedule
-    linear = graph_rewrite(linear, pm_sched_hcq_batches, walk=True, name="schedule hcq batches")
-    linear = linear.substitute({s: p for p, s in back_map.items()}, walk=True, enter_calls=True)
+    linear = graph_rewrite(linear, pm_tag_hcq_calls, ctx=(enumerator:=itertools.count(0)), walk=True, name="tag hcq calls")
+    linear = graph_rewrite(linear, pm_sched_sync, ctx=HCQDepsTracker(), walk=True, name="schedule sync")
+    linear = linear.substitute({s: p for p, s in back_map.items()}, walk=True)
     linear = graph_rewrite(linear, pm_merge_queues, walk=True, name="merge queues")
+    linear = graph_rewrite(linear, pm_add_finalizer, ctx=enumerator, walk=True, name="add finalizer")
+    linear = graph_rewrite(linear, pm_add_global_sync, ctx=set(), walk=True, name="add global sync", enter_calls=True)
 
     # lowering to hcq ir
+    linear = graph_rewrite(linear, pm_add_inner_loads, ctx=(waited:=set()), walk=True, name="add loads", enter_calls=True)
+    linear = graph_rewrite(linear, pm_add_inner_stores, ctx=waited, walk=True, name="add stores", enter_calls=True)
     linear = graph_rewrite(linear, pm_encode_cmdbufs, walk=True, name="encode cmdbufs", enter_calls=True)
     linear = graph_rewrite(linear, pm_pack_placeholders, walk=True, name="pack placeholders")
 
@@ -484,8 +508,7 @@ class HCQ2Compiled(Compiled):
     self.rt_allocator = BumpAllocator(64 << 20, wrap=False)
 
   def new_buffer(self, b:UOp, jit:bool) -> Buffer:
-    if jit or b.tag in HCQ_CACHE_TAGS:
-      return Buffer(self.device, b.max_numel(), b.dtype, options=BufferSpec(uncached=True, cpu_access=True, nolru=True))
+    if jit or b.tag in HCQ_CACHE_TAGS: return Buffer(self.device, b.max_numel(), b.dtype, options=BufferSpec(cpu_access=True, nolru=True))
     return self.rt_buffer.view(b.max_numel(), b.dtype, self.rt_allocator.alloc(b.max_numel() * b.dtype.itemsize, alignment=128))
 
   @functools.cache

extra/llama_kernels/cast_amax/__init__.py

@@ -43,7 +43,7 @@ def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
   device = xw13.device
   axis = xw13.axis if isinstance(device, tuple) else None
   grad_xw13_fp8 = alloc_like(xw13.shape, dtypes.fp8e4m3,  device, axis)
-  grad_amax_next = Tensor(next_grad_amax_state, device=device)
+  grad_amax_next = Tensor.zeros((), dtype=dtypes.float32, device=device).contiguous()
   grad_amax_state_t = Tensor(grad_amax_state, device=device)
   fxn = functools.partial(_custom_fused_bwd_w13, dname=dname_of(device))
   grad_amax = grad_amax_state_t.empty_like()
@@ -52,14 +52,16 @@ def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
     Tensor(xw13, device=device), Tensor(gradient, device=device).cast(dtypes.bfloat16),
     Tensor(amax_state, device=device), grad_amax_state_t, fxn=fxn)
   grad_xw13_uop = grad_xw13_fp8.uop.cast(dtypes.bfloat16)
+  store_effect = next_grad_amax_state.store(grad_amax_next.uop)
   assert grad_xw13_fp8.uop.op is Ops.AFTER, f"expected AFTER, got {grad_xw13_fp8.uop.op}"
+  grad_xw13_fp8_uop = grad_xw13_fp8.uop.replace(src=grad_xw13_fp8.uop.src + (store_effect,))
   # Stash fp8 companion for cdna_asm_gemm's bwd to attach to grad_a.
-  _grad_fp8_mailbox[grad_xw13_uop] = (grad_xw13_fp8.uop, grad_amax_state_t.uop)
+  _grad_fp8_mailbox[grad_xw13_uop] = (grad_xw13_fp8_uop, grad_amax_state_t.uop)
   return (None, None, grad_xw13_uop, None, None, None)
 
 def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_state:Tensor,
-                           next_grad_amax_state:Tensor, amax_out:Tensor) -> Tensor:
-  # NOTE: silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns fp8.
+                           next_grad_amax_state:Tensor) -> tuple[Tensor, Tensor]:
+  # NOTE: silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns (fp8, new_amax)
   # grad_amax_state: delayed amax for grad_xw13 fp8 quantization in the backward.
   assert xw13.dtype == dtypes.bfloat16, f"expected bf16, got {xw13.dtype}"
   MBS, SEQ, H2 = xw13.shape
@@ -67,7 +69,8 @@ def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_
   HIDDEN = H2 // 2
   axis = xw13.uop.axis if isinstance(xw13.device, tuple) else None
   fp8_out  = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype,      xw13.device, axis)
+  amax_out = Tensor.zeros((), dtype=dtypes.float32, device=xw13.device).contiguous()
   fxn = functools.partial(_custom_fused_cast_amax_w13, dname=dname_of(xw13.device))
   fp8_out, amax_out, *_ = Tensor.custom_kernel(fp8_out, amax_out, xw13, amax_state, grad_amax_state, next_grad_amax_state,
                                                 fxn=fxn, grad_fxn=_fused_quantize_bwd_w13)
-  return fp8_out
+  return fp8_out, amax_out

extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/__init__.py

@@ -112,9 +112,8 @@ def _fused_add_bwd(*args, **kwargs):
   grad_h, grad_w = _bwd_common(fp8_grad_u, h_grad_u, x_u, x_normed_u, rrms_u, weight_u, amax_state_u, kernel)
   return (None, None, None, None, None, grad_h, grad_h, grad_w, None)
 
-def fused_rmsnorm_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, eps:float, fp8_dtype,
-                                   amax_out:Tensor) -> tuple[Tensor, Tensor, Tensor]:
-  # NOTE: rmsnorm(x) * weight -> fp8 + amax. Returns (fp8, x_normed, rrms).
+def fused_rmsnorm_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, eps:float, fp8_dtype) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+  # NOTE: rmsnorm(x) * weight -> fp8 + amax. Returns (fp8, new_amax, x_normed, rrms).
   # x_normed + rrms are saved for the rmsnorm backward (also recomputed here from x regs).
   assert x.dtype == dtypes.bfloat16 and weight.dtype == dtypes.bfloat16
   assert x.shape[-1] == weight.shape[-1], f"HIDDEN mismatch: x={x.shape}, weight={weight.shape}"
@@ -124,15 +123,16 @@ def fused_rmsnorm_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, e
   fp8_out      = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype,       x.device, axis)
   x_normed_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
   rrms_out     = alloc_like((MBS, SEQ),         dtypes.float32,  x.device, axis)
+  amax_out     = Tensor.zeros((), dtype=dtypes.float32, device=x.device).contiguous()
   fxn = functools.partial(_custom_fwd, dname=dname_of(x.device), eps_val=eps)
   fp8_out, x_normed_out, rrms_out, amax_out, *_ = Tensor.custom_kernel(
     fp8_out, x_normed_out, rrms_out, amax_out, x, weight, amax_state, fxn=fxn, grad_fxn=_fused_bwd)
-  return fp8_out, x_normed_out, rrms_out
+  return fp8_out, amax_out, x_normed_out, rrms_out
 
 def fused_add_rmsnorm_mul_quantize_fp8(x:Tensor, residual:Tensor, weight:Tensor, amax_state:Tensor,
-                                       eps:float, fp8_dtype, amax_out:Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+                                       eps:float, fp8_dtype) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
   # NOTE: h = x + residual; y_normed = rmsnorm(h); fp8 = quantize(y_normed * weight).
-  # Returns (fp8, h, x_normed, rrms). h is also written so downstream can
+  # Returns (fp8, new_amax, h, x_normed, rrms). h is also written so downstream can
   # reuse it without recomputing x+residual — eliminates the separate residual-add kernel.
   assert x.dtype == dtypes.bfloat16 and residual.dtype == dtypes.bfloat16 and weight.dtype == dtypes.bfloat16
   assert x.shape == residual.shape
@@ -143,8 +143,9 @@ def fused_add_rmsnorm_mul_quantize_fp8(x:Tensor, residual:Tensor, weight:Tensor,
   h_out        = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
   x_normed_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
   rrms_out     = alloc_like((MBS, SEQ),         dtypes.float32,  x.device, axis)
+  amax_out     = Tensor.zeros((), dtype=dtypes.float32, device=x.device).contiguous()
   fxn = functools.partial(_custom_fwd_add, dname=dname_of(x.device), eps_val=eps)
   fp8_out, h_out, x_normed_out, rrms_out, amax_out, *_ = Tensor.custom_kernel(
     fp8_out, h_out, x_normed_out, rrms_out, amax_out, x, residual, weight, amax_state,
     fxn=fxn, grad_fxn=_fused_add_bwd)
-  return fp8_out, h_out, x_normed_out, rrms_out
+  return fp8_out, amax_out, h_out, x_normed_out, rrms_out

extra/llama_kernels/quantize_fp8_delayed/__init__.py

@@ -73,19 +73,24 @@ def _quantize_fp8_delayed_bwd(gradient:UOp, kernel:UOp):
   grad_x = (Tensor(gradient, device=device).float() * scale).cast(dtypes.bfloat16)
   return (None, None, grad_x.uop, None)
 
-def quantize_fp8_delayed(x:Tensor, amax_state:Tensor, amax_out:Tensor, fp8_dtype=dtypes.fp8e4m3) -> tuple[Tensor, Tensor]:
-  # NOTE: one-pass bf16 -> fp8 quantize with delayed scaling.
+def quantize_fp8_delayed(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3) -> tuple[Tensor, Tensor, Tensor, UOp]:
+  # NOTE: one-pass bf16 -> fp8 quantize with delayed scaling. Returns (fp8, inv_scale, new_amax, store_effect).
   # Fused kernel reads x once and writes fp8 + scalar amax via global atomic max.
+  # store_effect writes new_amax into amax_state's buffer — the caller must thread it into a realized
+  # output via `.after(store_effect)`. Calling `amax_state.assign(new_amax)` inside a grad_fxn does
+  # NOT work because .assign mutates only the temp Tensor's .uop, not the original layer-owned buffer.
   assert x.dtype == dtypes.bfloat16, f"expected bf16, got {x.dtype}"
   axis = x.uop.axis if isinstance(x.device, tuple) else None
   fp8_out      = alloc_like(x.shape,  fp8_dtype,      x.device, axis)
   n_elems = prod(x.uop.shard_shape)
   assert n_elems % NUM_WG == 0, f"{n_elems=} must divide over {NUM_WG=}"
+  amax_out = Tensor.zeros((), dtype=dtypes.float32, device=x.device).contiguous()
   fxn = functools.partial(_custom_quantize_fp8_with_amax, device=x.device)
   fp8_out, amax_out, *_ = Tensor.custom_kernel(fp8_out, amax_out, x, amax_state,
                                                 fxn=fxn, grad_fxn=_quantize_fp8_delayed_bwd)
   inv_scale = (amax_state.float() + 1e-8) / FP8_MAX
-  return fp8_out, inv_scale
+  store_effect = amax_state.uop.store(amax_out.uop)
+  return fp8_out, inv_scale, amax_out, store_effect
 
 def quantize_fp8_scalar(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3) -> Tensor:
   # NOTE: pure one-pass bf16 -> fp8 quantize with delayed scalar scale. No amax computation.

extra/thunder/amd/fa.py

@@ -94,7 +94,7 @@ def fused_qkv_rope(xqkv:Tensor, freqs_cis:Tensor, n_heads:int, n_kv_heads:int, h
   B_local = B // num_devices if is_dp else B
   H_local = n_heads // num_devices if is_mp else n_heads
   H_KV_local = n_kv_heads // num_devices if is_mp else n_kv_heads
-  assert H_local % H_KV_local == 0 and head_dim % 2 == 0 and head_dim <= 512
+  assert (B_local, N, H_local, H_KV_local, head_dim) == (2, 8192, 32, 8, 128)
   single_device = xqkv.device[0] if isinstance(xqkv.device, tuple) else xqkv.device
   arch = Device[single_device].renderer.target.arch
   axis = 0 if is_dp else 2 if is_mp else None

extra/torch_backend/backend.py

@@ -565,10 +565,8 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   "aten.floor_divide": lambda x,y: x//y,
   "aten.floor_divide_.Tensor": lambda x,y: x//y,
   "aten.__lshift__.Scalar": lambda x,y: x<<y,
-  "aten.__lshift__.Tensor": lambda x,y: x<<y,
   "aten.__ilshift__.Scalar": lambda x,y: x<<y,
   "aten.__rshift__.Scalar": lambda x,y: x>>y,
-  "aten.__rshift__.Tensor": lambda x,y: x>>y,
   "aten.__irshift__.Scalar": lambda x,y: x>>y,
   # inplace ops using replace for fusion
   "aten.zero_": lambda x: x.const_like(0),

pyproject.toml

@@ -161,6 +161,7 @@ norecursedirs = [
   ".git",
 ]
 timeout = 300
+timeout_method = "thread"
 timeout_func_only = true
 testpaths = ["test"]
 filterwarnings = [

test/backend/test_dtype_alu.py

@@ -291,33 +291,6 @@ class TestDTypeALU(unittest.TestCase):
   @Context(EMULATED_DTYPES="long")
   def test_emulated_int64(self, a, b, op): universal_test(a, b, dtypes.int64, op)
 
-  def _test_shl(self):
-    for dtype, values, distances in ((dtypes.int64, [-0x1234, 0x80000001, -1, 0x1234, 1], [0, 5, 31, 32, 62]),
-                                     (dtypes.uint64, [0x80000001, 0x80000001, 1, 0xFEDC, 1], [0, 5, 31, 32, 62]),
-                                     (dtypes.int8, [-3, 1, 7, -2, 1], [0, 1, 3, 5, 6]),
-                                     (dtypes.uint16, [3, 1, 0xFF, 7, 1], [0, 1, 7, 12, 15])):
-      with self.subTest(dtype=dtype):
-        result = Tensor(values, dtype=dtype) << Tensor(distances, dtype=dtype)
-        np.testing.assert_equal(result.numpy(), [x << d for x, d in zip(values, distances)])
-
-  def _test_shr(self):
-    for dtype, values, distances in ((dtypes.int64, [-(2**40), -1, -(2**50), -(2**40), 0x123456789ABCDEF], [0, 5, 31, 32, 63]),
-                                     (dtypes.uint64, [0xFEDCBA9876543210] * 5, [0, 5, 31, 32, 63]),
-                                     (dtypes.int8, [-128, -1, 64, -37, 1], [0, 1, 3, 5, 7]),
-                                     (dtypes.uint16, [0xFFFF] * 5, [0, 1, 8, 13, 15])):
-      with self.subTest(dtype=dtype):
-        result = Tensor(values, dtype=dtype) >> Tensor(distances, dtype=dtype)
-        np.testing.assert_equal(result.numpy(), [x >> d for x, d in zip(values, distances)])
-
-  def test_shl(self): self._test_shl()
-  def test_shr(self): self._test_shr()
-
-  @Context(EMULATED_DTYPES="long")
-  def test_emulated_shl(self): self._test_shl()
-
-  @Context(EMULATED_DTYPES="long")
-  def test_emulated_shr(self): self._test_shr()
-
   @given(ht.uint8, strat.sampled_from(integer_unary_operations))
   def test_uint8_unary(self, a, op): universal_test_unary(a, dtypes.uint8, op)
 

test/backend/test_llama_kernels.py

@@ -49,10 +49,9 @@ def run_quantize_fp8(shape:tuple[int, ...], delayed:bool=True) -> None:
   with Context(DEBUG=0): Tensor.realize(x, amax_state)
 
   if delayed:
-    amax_out = Tensor.zeros((), dtype=dtypes.float32, device=x.device).realize()
-    fp8, inv_scale = quantize_fp8_delayed(x, amax_state, amax_out, FP8_DTYPE)
+    fp8, inv_scale, new_amax, _ = quantize_fp8_delayed(x, amax_state, FP8_DTYPE)
     ref_fp8, ref_inv_scale, ref_new_amax = quantize_fp8(x, amax_state=amax_state)
-    Tensor.realize(fp8, inv_scale)
+    Tensor.realize(fp8, inv_scale, new_amax)
     Tensor.realize(ref_fp8, ref_inv_scale, ref_new_amax)
   else:
     fp8 = quantize_fp8_scalar(x, amax_state, FP8_DTYPE)
@@ -64,8 +63,8 @@ def run_quantize_fp8(shape:tuple[int, ...], delayed:bool=True) -> None:
     assert fp8.cast(dtypes.float).allclose(ref_fp8.cast(dtypes.float), atol=0, rtol=0).item(), "fp8 mismatch"
     if delayed:
       assert inv_scale.allclose(ref_inv_scale, atol=0, rtol=0).item(), "inv_scale mismatch"
-      assert amax_out.allclose(ref_new_amax, atol=0, rtol=0).item(), \
-        f"amax mismatch: got={amax_out.item()} ref={ref_new_amax.item()} diff={abs(amax_out.item()-ref_new_amax.item())}"
+      assert new_amax.allclose(ref_new_amax, atol=0, rtol=0).item(), \
+        f"amax mismatch: got={new_amax.item()} ref={ref_new_amax.item()} diff={abs(new_amax.item()-ref_new_amax.item())}"
 
 @unittest.skipUnless(Device.DEFAULT == "AMD", "requires atomic max")
 class TestQuantizeFP8(unittest.TestCase):
@@ -83,11 +82,10 @@ class TestQuantizeFP8(unittest.TestCase):
     x = Tensor.empty(2048*8, 1024, dtype=dtypes.bfloat16, device=devs).uop.multi(0)
     x = Tensor(x, device=devs)
     amax_state = Tensor.full((), 2.0, dtype=dtypes.float32, device=devs).contiguous()
-    amax_out = Tensor.zeros((), dtype=dtypes.float32, device=devs).realize()
-    fp8, _ = quantize_fp8_delayed(x, amax_state, amax_out, FP8_DTYPE)
-    Tensor.realize(fp8)
+    fp8, _, new_amax, _ = quantize_fp8_delayed(x, amax_state, FP8_DTYPE)
+    Tensor.realize(fp8, new_amax)
     assert fp8.uop.shape == x.uop.shape
-    assert amax_out.shape == ()
+    assert new_amax.shape == ()
 
 class TestLocalAmax(unittest.TestCase):
   def test_multi_tensor_local_shard_amax(self):

test/backend/test_ops.py

@@ -848,8 +848,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: tor << 0, lambda: (ten << 0).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor << 2, lambda: (ten << 2).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor << 31, lambda: (ten << 31).cast(dtypes.int32), forward_only=True)
-    helper_test_op([], lambda: tor << torch.tensor([0,2,4]).int(),
-                   lambda: (ten << Tensor([0,2,4], dtype=dtypes.uint32)).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor.__lshift__(2), lambda: ten.__lshift__(2).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor.bitwise_left_shift(2), lambda: ten.lshift(2).cast(dtypes.int32), forward_only=True)
 
@@ -861,8 +859,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: tor >> 0, lambda: (ten >> 0).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor >> 2, lambda: (ten >> 2).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor >> 31, lambda: (ten >> 31).cast(dtypes.int32), forward_only=True)
-    helper_test_op([], lambda: tor >> torch.tensor([0,2,4]).int(),
-                   lambda: (ten >> Tensor([0,2,4], dtype=dtypes.uint32)).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor.__rshift__(2), lambda: ten.__rshift__(2).cast(dtypes.int32), forward_only=True)
     helper_test_op([], lambda: tor.bitwise_right_shift(2), lambda: ten.rshift(2).cast(dtypes.int32), forward_only=True)
 
@@ -874,7 +870,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: tor << 2, lambda: ten << 2, forward_only=True)
     helper_test_op([], lambda: tor << 8, lambda: ten << 8, forward_only=True)
     helper_test_op([], lambda: tor << 31, lambda: ten << 31, forward_only=True)
-    helper_test_op([], lambda: tor << torch.tensor([0,2,8,31]).int(), lambda: ten << Tensor([0,2,8,31], dtype=dtypes.int), forward_only=True)
 
   def test_rshift_signed(self):
     data = [[-1, -3, 1, 7], [0, -2147483648, 2147483647, -1]]
@@ -884,7 +879,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: tor >> 2, lambda: ten >> 2, forward_only=True)
     helper_test_op([], lambda: tor >> 8, lambda: ten >> 8, forward_only=True)
     helper_test_op([], lambda: tor >> 31, lambda: ten >> 31, forward_only=True)
-    helper_test_op([], lambda: tor >> torch.tensor([0,2,8,31]).int(), lambda: ten >> Tensor([0,2,8,31], dtype=dtypes.int), forward_only=True)
 
   def test_idiv_shift_rewrite_negative(self):
     a = Tensor(-5).div(2, rounding_mode="trunc").item()

test/external/external_test_llm_opencode.py

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

test/helpers.py

@@ -55,7 +55,7 @@ def assert_jit_cache_len(fxn, expected_len):
   if linear is None or not linear.src:
     assert expected_len == 0, expected_len
     return
-  if expected_len and all(call_is_hcq(call) for call in linear.src): expected_len = 3 # HCQ2: merged same-queue calls + finalizer + bumps
+  if expected_len and all(call_is_hcq(call) for call in linear.src): expected_len = 2 # HCQ2 merges calls on the same queue
   if call_is_graph(linear.src[0]):
     assert len(linear.src) == 1, len(linear.src)
     inner = linear.src[0].src[0].src[0]  # LINEAR UOp inside CUSTOM_FUNCTION

test/mockgpu/mockgpu.py

@@ -1,4 +1,4 @@
-import ctypes, time, os, builtins, fcntl, typing
+import ctypes, time, os, builtins, fcntl
 from tinygrad.helpers import DEV
 from tinygrad.runtime.support.hcq import FileIOInterface
 from tinygrad.runtime.autogen import libc
@@ -9,7 +9,7 @@ start = time.perf_counter()
 
 drivers = [cls() for t in DEV.value if (cls:={"MOCKPCI+AMD": AMDriver, "MOCKKFD+AMD": AMDDriver, "MOCK+AMD": AMDDriver, "MOCKUSB+AMD": AMUSBDriver,
                                               "MOCK+NV": NVDriver}.get(f"{t.interface}+{t.device}"))]
-tracked_fds: dict[int, typing.Any] = {}
+tracked_fds = {}
 
 original_memoryview = builtins.memoryview
 class TrackedMemoryView:

test/null/test_helpers.py

@@ -1,7 +1,6 @@
 import ctypes, gzip, unittest, timeit, pickle
 from tinygrad import Variable
-from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, mv_address, count, all_same
-from tinygrad.tensor import is_numpy_ndarray
+from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, mv_address, count, all_same
 from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
 from tinygrad.helpers import ceildiv, ansistrip, get_shape
 from tinygrad.tensor import Tensor

test/null/test_llm_server.py

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

test/null/test_llm_tokenizer.py

@@ -1,4 +1,4 @@
-import unittest, base64, functools, re, sys, time, unicodedata
+import unittest, base64, functools, sys
 from tinygrad.llm.cli import SimpleTokenizer, FallbackTemplate
 from tinygrad.helpers import fetch
 
@@ -46,29 +46,6 @@ class TestLLMTokenizer(unittest.TestCase):
   def test_llama_repeat(self): self._test_coding(self.llama_tok, "00000000000000000", [ 931, 931, 931, 931, 931, 410 ])
   def test_llama_pat(self): self._test_coding(self.llama_tok, "today\n  \n", [ 31213, 14211 ])
 
-  def test_split_regex_matches_naive_listing(self):
-    # the compacted codepoint ranges must match the same text as listing every codepoint
-    def naive(pre): return "".join(re.escape(chr(cp)) for cp in range(0x323b0) if unicodedata.category(chr(cp)).startswith(pre))
-    r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + naive("Z"), naive("N"), naive("L")
-    naive_re = re.compile("(?i:'s|'t|'re|'ve|'m|'ll|'d)|" +
-      f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+")
-    sample = "hello world 한국어 中文 текст ١٢٣ 123 😊\n  \ttoday\n'équivalent ²³№ "
-    self.assertEqual(SimpleTokenizer({}, {})._split_to_word.findall(sample), naive_re.findall(sample))
-
-  def test_split_regex_speed(self):
-    # the naive listing compiles a 429KB pattern that takes 10+s to match a 225KB prompt; ranges keep it small and fast
-    tok = SimpleTokenizer({}, {})
-    self.assertLess(len(tok._split_to_word.pattern), 100_000)
-    text = "The quick brown fox jumps over the lazy dog. " * 5000
-    tok._split_to_word.findall(text)  # warmup
-    tms = []
-    for _ in range(5):
-      st = time.perf_counter()
-      words = tok._split_to_word.findall(text)
-      tms.append(time.perf_counter() - st)
-    self.assertLess(min(tms), 4)  # best-of-5 is robust to CI scheduling pauses; new code takes ~60ms
-    self.assertEqual(len(words), 50001)
-
   def test_llama_continued_conversation(self):
     self._test_coding(self.llama_tok, "hello <|eot_id|>world", [15339, 220, 128009, 14957])
     self._test_coding(self.llama_tok, "hello <|eot_id|>world again", [15339, 220, 128009, 14957, 1578])

test/null/test_schedule.py

@@ -1472,18 +1472,6 @@ class TestSchedule(unittest.TestCase):
     x.softmax().sum().backward()
     run_linear(*check_schedule(x.grad, 4))
 
-  def test_logsumexp_backward(self):
-    Tensor.manual_seed(0)
-    x = Tensor.randn(4, 12, 64, 64).realize()
-    x.logsumexp(-1).sum().backward()
-    run_linear(*check_schedule(x.grad, 3))
-
-  def test_logcumsumexp_backward(self):
-    Tensor.manual_seed(0)
-    x = Tensor.randn(4, 512).realize()
-    x.logcumsumexp(-1).sum().backward()
-    run_linear(*check_schedule(x.grad, 3))
-
   def test_scaled_dot_product_attention_fusion(self):
     x, y, z, m = (Tensor.empty(32, 8, 16, 16) for _ in range(4))
     out = Tensor.scaled_dot_product_attention(x, y, z, attn_mask=m)

test/null/test_uop_vmin_vmax.py

@@ -127,13 +127,6 @@ class TestVminVmaxProperties(unittest.TestCase):
     self.assertEqual(x.vmin, 0)
     self.assertEqual(x.vmax, 10 >> 2)
 
-  def test_vmin_vmax_cast_unsigned(self):
-    # a fitting source keeps exact bounds: no wrap can occur
-    self.assertEqual(UOp.variable('x', 5, 10).cast(dtypes.uint8)._min_max, (5, 10))
-    # a possibly-negative or too-large source can wrap: conservative
-    self.assertEqual(UOp.variable('x', -1, 10).cast(dtypes.uint8)._min_max, (0, 255))
-    self.assertEqual(UOp.variable('x', 250, 260).cast(dtypes.uint8)._min_max, (0, 255))
-
   def test_vmin_vmax_xor_neg1(self):
     x = UOp.variable('x', 3, 7)
     uop = x ^ -1
@@ -371,10 +364,6 @@ class TestConstFactor(unittest.TestCase):
     uop = (x * 3) * 5
     self.assertEqual(uop.const_factor(), 15)  # Constant multipliers are combined (3 * 5 = 15)
 
-  def test_const_factor_variable_multiple_of(self):
-    x = UOp.variable('x', 16, 32, multiple_of=4)
-    self.assertEqual(x.const_factor(), 4)
-
 class TestDivides(unittest.TestCase):
   def test_divides_constant_exact(self):
     # Divides a constant by an exact divisor
@@ -413,15 +402,5 @@ class TestDivides(unittest.TestCase):
     result = uop.divides(3)
     self.assertIsNone(result)  # Cannot divide by 3, since 4 is not divisible by 3
 
-  def test_divides_variable_multiple_of_exact(self):
-    x = UOp.variable('x', 16, 32, multiple_of=4)
-    result = x.divides(4)
-    self.assertIsNotNone(result)
-
-  def test_divides_variable_multiple_of_factor(self):
-    x = UOp.variable('x', 16, 32, multiple_of=4)
-    result = x.divides(2)
-    self.assertIsNotNone(result)
-
 if __name__ == '__main__':
   unittest.main()

test/opt/test_gen_float4.py

@@ -1,5 +1,5 @@
 import unittest
-from tinygrad import Device, Tensor, Variable, dtypes
+from tinygrad import Device, Tensor, dtypes
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.codegen import to_program
 from tinygrad.codegen.opt import Opt, OptOps
@@ -114,31 +114,5 @@ class TestFloat4(unittest.TestCase):
 
     assert TestFloat4.count_float4(uops) == (1, 1)
 
-  def test_float4_aligned_variable(self):
-    x = Variable('x', 0, 4, multiple_of=4).bind(4)
-    a = Tensor.empty(4).realize()
-    b = Tensor.empty(12).realize().shrink(((x, x+4),))
-    c = a + b
-
-    # should float4 both
-
-    s = c.linear_with_vars()[0].src[0]
-    uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[1].src)
-
-    assert TestFloat4.count_float4(uops) == (2, 1)
-
-  def test_float4_unaligned_variable(self):
-    x = Variable('x', 0, 4, multiple_of=2).bind(4)
-    a = Tensor.empty(4).realize()
-    b = Tensor.empty(12).realize().shrink(((x, x+4),))
-    c = a + b
-
-    # should float4 a but not b
-
-    s = c.linear_with_vars()[0].src[0]
-    uops = tuple(to_program(replace_opts(s.src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).src[1].src)
-
-    assert TestFloat4.count_float4(uops) == (1, 1)
-
 if __name__ == '__main__':
   unittest.main()

test/opt/test_tensor_cores.py

@@ -18,9 +18,6 @@ from test.backend.test_linearizer import helper_realized_ast, helper_linearizer_
 
 # NOTE: to_program always passes in Device[Device.DEFAULT].renderer explicitly for process_replay!!!
 
-def _tc_rand(*shape, dtype:DType) -> Tensor:
-  return Tensor.randint(*shape, low=dtype.min, high=dtype.max+1, dtype=dtype) if dtypes.is_int(dtype) else Tensor.rand(*shape, dtype=dtype)
-
 def run_program(prg:UOp, bufs:list[Buffer]):
   buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in bufs]
   for u,b in zip(buf_uops, bufs): buffers[u] = b
@@ -28,7 +25,7 @@ def run_program(prg:UOp, bufs:list[Buffer]):
 
 def helper_tc_ensure_uops_and_opts_count(N: int, M:int, K:int, dtype_in:DType, dtype_out:DType, axis:int=0, tc_select:int=-1, tc_opt:int=0,
                                          ensure_triggered:bool=True):
-  a, b = _tc_rand(M, K, dtype=dtype_in), _tc_rand(K, N, dtype=dtype_in)
+  a, b = Tensor.rand(M, K, dtype=dtype_in), Tensor.rand(K, N, dtype=dtype_in)
   r = a.matmul(b, dtype=dtype_out)
   sched = r.schedule_linear()
   realized_ast = sched.src[-1].src[0]
@@ -47,7 +44,7 @@ def helper_tc_ensure_uops_and_opts_count(N: int, M:int, K:int, dtype_in:DType, d
     except KernelOptError: pass
 
 def helper_tc_allclose(N:int, M:int, K:int, dtype_in:DType, dtype_out:DType, axis:int=0, tc_select:int=-1, tc_opt:int=0, use_tensor_cores:int=1):
-  a, b = _tc_rand(M, K, dtype=dtype_in), _tc_rand(K, N, dtype=dtype_in)
+  a, b = Tensor.rand(M, K, dtype=dtype_in), Tensor.rand(K, N, dtype=dtype_in)
   np_a, np_b = a.numpy(), b.numpy()
   r = a.matmul(b, dtype=dtype_out)
   if dtype_in == dtypes.bfloat16: r = r.float()
@@ -60,11 +57,9 @@ def helper_tc_allclose(N:int, M:int, K:int, dtype_in:DType, dtype_out:DType, axi
   run_program(ast, bufs)
   if dtype_in == dtypes.half: tc_atol, tc_rtol = 1e-2, 1e-3
   elif dtype_in == dtypes.bfloat16: tc_atol, tc_rtol = (1e-1, 2e-2) if dtype_out == dtypes.bfloat16 else (1e-2, 1e-2)
-  elif not dtypes.is_float(dtype_in): tc_atol, tc_rtol = 0, 0
   else: tc_atol, tc_rtol = 5e-3, 1e-4
   c = bufs[0].numpy().reshape((M,N))
-  ref = (np_a.astype(np.int32) @ np_b.astype(np.int32)) if not dtypes.is_float(dtype_in) else (np_a @ np_b)
-  np.testing.assert_allclose(c, ref, atol=tc_atol, rtol=tc_rtol)
+  np.testing.assert_allclose(c, np_a @ np_b, atol=tc_atol, rtol=tc_rtol)
 
 class TestTensorCores(unittest.TestCase):
   # TODO: don't skip bf16 for real device (METAL, AMD)
@@ -80,7 +75,7 @@ class TestTensorCores(unittest.TestCase):
   def test_tensor_cores_codegen(self):
     for tc in Device[Device.DEFAULT].renderer.tensor_cores:
       n, m, k = tc.dims
-      a, b = _tc_rand(m, k, dtype=tc.dtype_in), _tc_rand(k, n, dtype=tc.dtype_in)
+      a, b = Tensor.rand(m, k, dtype=tc.dtype_in), Tensor.rand(k, n, dtype=tc.dtype_in)
       r = a.matmul(b, dtype=tc.dtype_out)
       prg = to_program(replace_opts(r.schedule_linear().src[-1].src[0],
                         [Opt(op=OptOps.TC, axis=0, arg=(-1, 2, 1))]), Device[Device.DEFAULT].renderer)

test/unit/test_attention.py

@@ -67,7 +67,10 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
   def _run_attention(self, block:GatedDeltaNetBlock, x:Tensor, start_pos:int):
     x_norm = block.attn_norm(x)
     block._init_state(x_norm)
-    return block._attention(x_norm, start_pos).realize().numpy()
+    out = block._attention(x_norm, start_pos).realize()
+    assert block.pending_state is not None
+    Tensor.realize(block.conv_state.assign(block.pending_state[0]), block.recurrent_state.assign(block.pending_state[1]))
+    return out.numpy()
 
   def _cache_views(self, block:GatedDeltaNetBlock) -> tuple[np.ndarray, np.ndarray]:
     if hasattr(block, 'conv_state'):
@@ -86,8 +89,8 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
     x_float = x.astype(np.float32)
     return (x_float / np.sqrt((x_float * x_float).mean(axis=-1, keepdims=True) + eps)) * weight.astype(np.float32)
 
-  def _normalize_np(self, x:np.ndarray, eps:float=1e-12) -> np.ndarray:
-    return x / np.maximum(np.sqrt((x * x).sum(axis=-1, keepdims=True)), eps)
+  def _normalize_np(self, x:np.ndarray, eps:float=1e-6) -> np.ndarray:
+    return x / np.sqrt((x * x).sum(axis=-1, keepdims=True) + eps)
 
   def _softplus_np(self, x:np.ndarray) -> np.ndarray:
     return np.log1p(np.exp(-np.abs(x))) + np.maximum(x, 0)

test/unit/test_dtype_weak.py

@@ -1,11 +1,10 @@
-import pathlib, tempfile, unittest
+import unittest
 from unittest.mock import patch
 
 from tinygrad import Tensor, dtypes
 from tinygrad.uop import Ops
 from tinygrad.uop.ops import UOp
 from tinygrad.uop.spec import spec_tensor
-from tinygrad.nn.state import safe_save
 
 
 class TestWeakPromotion(unittest.TestCase):
@@ -30,7 +29,8 @@ class TestWeakPromotion(unittest.TestCase):
       self.assertEqual(t.dtype, weak)
       self.assertEqual(t.data().itemsize, strong.itemsize)
       self.assertEqual(t.numpy().dtype.itemsize, strong.itemsize)
-      with self.assertRaises(RuntimeError): t.clone("CPU")
+      realized = t.clone("CPU").realize()
+      self.assertEqual((realized.dtype, realized.uop.buffer.dtype), (strong, strong))
     with patch.object(dtypes, "default_int", dtypes.int64):
       self.assertEqual(Tensor.const(dtypes.weakint, 3).numpy().dtype.itemsize, dtypes.int64.itemsize)
 
@@ -110,65 +110,6 @@ class TestWeakPromotion(unittest.TestCase):
       self.assertNotIn(t.uop.buffer.dtype, dtypes.weaks)
 
 
-class TestWeakStorageBoundary(unittest.TestCase):
-  # weak has no storage: a weak assignment source casts when it defers to the destination, everything else raises
-  def test_weak_source(self):
-    w3, w05 = Tensor.const(dtypes.weakint, 3).reshape(1).expand(2), Tensor.const(dtypes.weakfloat, 0.5).reshape(1)
-    dst = Tensor.zeros(2, dtype=dtypes.int8, device="CPU").contiguous().realize()
-    self.assertEqual(dst.assign(w3).realize().tolist(), [3, 3])                       # weakint defers to int8
-    with self.assertRaises(RuntimeError): dst.assign(w05.expand(2))                   # weakfloat into int does not defer
-    with self.assertRaises(RuntimeError): dst[0:1] = w05
-    fdst = Tensor.zeros(2, dtype=dtypes.float32, device="CPU").contiguous().realize()
-    fdst[0:1] = w05                                                                    # weakfloat defers to float
-    self.assertEqual(fdst.tolist(), [0.5, 0.0])
-    with tempfile.TemporaryDirectory() as td:                                          # the DISK path checks the same
-      ddst = Tensor.empty(2, dtype=dtypes.int32, device=f"DISK:{td}/t")
-      self.assertEqual(ddst.assign(w3).tolist(), [3, 3])
-      with self.assertRaises(RuntimeError): ddst.assign(w05.expand(2))
-
-  def test_weak_has_no_storage(self):
-    w = Tensor.const(dtypes.weakint, 3)
-    with self.assertRaises(RuntimeError): w.assign(Tensor([1], device="CPU"))
-    with self.assertRaises(RuntimeError): w.reshape(1)[0] = 1
-    with tempfile.TemporaryDirectory() as td:
-      with self.assertRaises(ValueError): safe_save({"x": w.reshape(1).expand(2)}, f"{td}/w.safetensors")
-    with self.assertRaises(RuntimeError): Tensor.empty(2, dtype=dtypes.weakint)
-    with self.assertRaises(RuntimeError): UOp.new_buffer("CPU", 2, dtypes.weakint)  # the one storage boundary
-    with self.assertRaises(RuntimeError): Tensor([1], dtype=dtypes.weakint)
-    import numpy as np
-    with self.assertRaises(RuntimeError): Tensor(np.ones(2, dtype=np.int32), dtype=dtypes.weakint)
-    self.assertEqual(Tensor(np.array(3), dtype=dtypes.weakint).dtype, dtypes.weakint)  # a 0-D ndarray is a const, not storage
-    with self.assertRaises(RuntimeError): Tensor(np.ones(2, dtype=np.float32), dtype=dtypes.weakfloat)
-    with self.assertRaises(RuntimeError): Tensor(bytes(8), dtype=dtypes.weakfloat)
-    with self.assertRaises(RuntimeError): Tensor(bytes(8), dtype=dtypes.weakint)
-    with tempfile.NamedTemporaryFile(suffix=".bin") as f:
-      f.write(bytes(8))
-      f.flush()
-      with self.assertRaises(RuntimeError): Tensor(pathlib.Path(f.name), dtype=dtypes.weakint)
-
-class TestWeakMaterializationEntries(unittest.TestCase):
-  # everything that creates storage from a weak value raises
-  def test_reads_commit_storage_raises(self):
-    for weak, value, strong in ((dtypes.weakint, 3, dtypes.default_int), (dtypes.weakfloat, 0.5, dtypes.default_float)):
-      def weak_val():
-        return Tensor([True], device="CPU").where(Tensor.const(weak, value), Tensor.const(weak, value))
-      self.assertEqual(weak_val().dtype, weak)
-      self.assertEqual(weak_val().to("CPU").dtype, weak)
-      self.assertEqual(weak_val().data().format, strong.fmt)
-      self.assertEqual(weak_val().numpy().dtype.itemsize, strong.itemsize)
-      self.assertEqual(weak_val().tolist(), [value])
-      self.assertEqual(weak_val().cast(strong).realize().uop.buffer.dtype, strong)
-      for entry in (lambda t: t.contiguous(), lambda t: t.realize(), lambda t: t.clone(),
-                    lambda t: t.to("CPU:1").realize(), lambda t: t.as_param(0)):
-        with self.assertRaises(RuntimeError): entry(weak_val())
-
-  def test_empty_reads_commit(self):
-    for weak, strong in ((dtypes.weakint, dtypes.default_int), (dtypes.weakfloat, dtypes.default_float)):
-      empty = Tensor.const(weak, 0).reshape(1).shrink(((0, 0),))
-      self.assertEqual(empty.data().format, strong.fmt)
-      self.assertEqual(empty.numpy().dtype.itemsize, strong.itemsize)
-      self.assertEqual(empty.tolist(), [])
-
 class TestWeakSpec(unittest.TestCase):
   def test_weak_operand_allowed(self):
     x = UOp.variable("x", 0, 10, dtypes.int64)

test/unit/test_helpers.py

@@ -1,7 +1,7 @@
 import unittest, gc
 import numpy as np
-from tinygrad.helpers import polyN, disable_gc
-from tinygrad.tensor import Tensor, is_numpy_ndarray
+from tinygrad.helpers import polyN, is_numpy_ndarray, disable_gc
+from tinygrad.tensor import Tensor
 
 class TestPolyN(unittest.TestCase):
   def test_tensor(self):

test/unit/test_llm_server.py

@@ -151,7 +151,7 @@ class TestTransformerGenerate(unittest.TestCase):
     """Temperature from generate should be passed through to __call__."""
     model = Transformer(TEST_CONFIG)
     captured_temps = []
-    def mock_call(self, tokens, start_pos, temperature):
+    def mock_call(self, tokens, start_pos, temperature, **kwargs):
       captured_temps.append(float(temperature.item()))
       return Tensor([[42]])
     with patch.object(Transformer, '__call__', mock_call):

test/unit/test_modern_scan.py

@@ -1,36 +0,0 @@
-import unittest
-from tinygrad import Tensor
-from tinygrad.uop.ops import UOp, AddrSpace
-
-class TestModernScan(unittest.TestCase):
-  def test_copy_local(self):
-    N = 256
-    state = Tensor.empty(N)
-    tmp = UOp.placeholder((N,), state.dtype, slot=-1, addrspace=AddrSpace.LOCAL)
-    tmp = tmp.after(tmp.store(state.uop))
-    state.assign(tmp)
-    state.realize()
-
-  """
-  def test_scan_gemv(self):
-    N = 256
-    gemvs = Tensor.empty(3, N, N)
-    state = Tensor.empty(N)
-    Tensor.realize(gemvs, state)
-
-    #tmp = UOp.placeholder((N,), state.dtype, slot=-1, addrspace=AddrSpace.REG)
-    tmp = Tensor.empty(N, dtype=state.dtype).uop
-    tmp = tmp.after(tmp.store(state.uop))
-    #rng = UOp.range(3, -1)
-    #tmp = tmp.after(tmp.store(state.uop, rng))
-    #tmp = tmp.after(tmp.store(tmp @ gemvs.uop[rng]).end(rng))
-    state.assign(tmp)
-
-    state.realize()
-  """
-
-if __name__ == '__main__':
-  unittest.main()
-
-
-

tinygrad/callify.py

@@ -180,9 +180,8 @@ def finalize_after(ctx:AllocCtx, x:UOp):
 def replace_input_buffer(ctx:AllocCtx, b:UOp):
   ctx.replacements.append(b)
   return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b.device,
-                   b._min_max if b.op is Ops.BIND else None, name=b.src[0].expr if b.op is Ops.BIND else None,
-                   addrspace=b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL,
-                   multiple_of=b.src[0].arg.multiple_of if b.op is Ops.BIND else None)
+                   b._min_max if b.op is Ops.BIND else None, b.src[0].expr if b.op is Ops.BIND else None,
+                   b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL)
 
 pm_finalize_call = PatternMatcher([
   (UPat(Ops.AFTER, name="x"), finalize_after),

tinygrad/codegen/__init__.py

@@ -1,7 +1,7 @@
 from dataclasses import replace, dataclass
-import itertools, functools
+import itertools, functools, hashlib, pickle
 from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC
-from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, panic
+from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, panic, CCACHE, diskcache_get, diskcache_put
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, GroupOp
 from tinygrad.uop.ops import AxisType
 from tinygrad.uop.render import pyrender
@@ -175,13 +175,9 @@ devectorizer2 = mop_cleanup+pm_mops+PatternMatcher([
   # EXPAND on scalar -> STACK
   (UPat(Ops.EXPAND, src=(UPat.var("x"), UPat()), name="out"),
    lambda x,out: UOp.stack(*([x]*out.max_numel())) if x.shape == () and out.shape == (out.max_numel(),) else None),
-  # TODO: make this all generic
   # INDEX on INDEX is INDEX
   (UPat(Ops.INDEX, src=(UPat(Ops.INDEX, name="idx1", allow_any_len=True),), allow_any_len=True, name="idx2"),
-   lambda idx1,idx2: idx1.src[0].index(*idx1.src[1:], *idx2.src[1:]) if all(x.shape == () for x in idx1.src[1:]+idx2.src[1:]) else None),
-  # INDEX on shaped INDEX
-  (UPat(Ops.INDEX, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx1_arg"))),), allow_any_len=True, name="idx2"),
-   lambda buf,idx1_arg,idx2: buf.index(idx1_arg.index(*idx2.src[1:])) if len(idx1_arg.shape) == len(idx2.src[1:]) else None),
+   lambda idx1, idx2: idx1.src[0].index(*idx1.src[1:], *idx2.src[1:])),
 ])
 
 def fix_group_for_reduce(x:UOp):
@@ -470,5 +466,10 @@ to_program_cache: dict[tuple, UOp] = {}
 def to_program(ast:UOp, renderer:Renderer) -> UOp:
   config = (NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32)
   key = (ast.key, type(renderer), renderer.target, *[x.value for x in config])
-  if (prg:=to_program_cache.get(key)) is None: to_program_cache[key] = prg = do_to_program(ast, renderer)
+  if (prg:=to_program_cache.get(key)) is None:
+    disk_key = hashlib.sha256(pickle.dumps(key)).digest()
+    if not CCACHE or (prg:=diskcache_get("program", {"key":disk_key})) is None:
+      prg = do_to_program(ast, renderer)
+      if CCACHE: diskcache_put("program", {"key":disk_key}, prg)
+    to_program_cache[key] = prg
   return prg

tinygrad/codegen/decomp/dtype.py

@@ -33,14 +33,11 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
     case Ops.CAST: return a0.bitcast(dtypes.uint).cast(dt)
     case Ops.BITCAST: return a0.bitcast(dt), a1.bitcast(dt)
     case Ops.SHL:
-      a0u, a1u, n = a0.bitcast(dtypes.uint), a1.bitcast(dtypes.uint), (b0 & 31).cast(dtypes.uint)
-      lo, hi = (a0u << n).bitcast(dt), ((a1u << n) | ((a0u >> 1) >> (31 - n))).bitcast(dt)
+      lo, hi = shl(a0, b0_mod:=b0 & 31), shl(a1, b0_mod) | shr(shr(a0, 1), 31 - b0_mod)
       return (b0 >= 32).where(zero, lo), (b0 >= 32).where(lo, hi)
     case Ops.SHR:
-      a0u, a1u, n = a0.bitcast(dtypes.uint), a1.bitcast(dtypes.uint), (b0 & 31).cast(dtypes.uint)
-      lo, hi = ((a0u >> n) | ((a1u << 1) << (31 - n))).bitcast(dt), a1 >> (b0 & 31)
-      fill = a1 >> 31 if dt == dtypes.int else zero  # vacated high word: sign bits when signed, else 0
-      return (b0 >= 32).where(hi, lo), (b0 >= 32).where(fill, hi)
+      lo, hi = shr(a0, b0_mod:=b0 & 31) | shl(shl(a1, 1), 31 - b0_mod), shr(a1, b0_mod)
+      return (b0 >= 32).where(hi, lo), (b0 >= 32).where(zero, hi)
     case Ops.ADD: return (low:=a0+b0), (a1 + b1).replace(dtype=dt) + (low.bitcast(dtypes.uint) < a0.bitcast(dtypes.uint)).cast(dt)
     case Ops.SUB: return a0 - b0, a1 - b1 - (a0.bitcast(dtypes.uint) < b0.bitcast(dtypes.uint)).cast(dt)
     case Ops.MUL:

tinygrad/codegen/opt/tc.py

@@ -103,7 +103,7 @@ amd_rdna3 = [TensorCore(dims=(16,16,16), threads=32, elements_per_thread=(16,16,
   opts=("l0","l0","l0","l0","l1","u1","u1","u1"),
   swizzle=((('l4', 'u0', 'u1', 'u2', 'l0'), ('r1', 'r2', 'r3'), ('l1', 'l2', 'l3', 'r0')),
            (('l0', 'l1', 'l2', 'l3', 'l4'), ('r1', 'r2', 'r3'), ('u0', 'u1', 'u2', 'r0'))))
-  for di,do in [(dtypes.half,dtypes.float),(dtypes.half,dtypes.half),(dtypes.bfloat16,dtypes.float),(dtypes.int8,dtypes.int32)]]
+  for di,do in [(dtypes.half,dtypes.float),(dtypes.half,dtypes.half),(dtypes.bfloat16,dtypes.float)]]
 amd_rdna4 = [TensorCore(dims=(16,16,16), threads=32, elements_per_thread=(8,8,8), dtype_in=di, dtype_out=do,
   opts=("l0","l0","l0","l0","u1","u1","u1","l1"),
   swizzle=((('u0', 'u1', 'u2', 'l4', 'r2'), ('r0', 'r1', 'r3'), ('l0', 'l1', 'l2', 'l3')),

tinygrad/engine/jit.py

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

tinygrad/helpers.py

@@ -81,6 +81,7 @@ def to_be32(val:Any) -> Any: return ((val & 0xFF) << 24) | (((val >> 8) & 0xFF)
 def to_be64(val:Any) -> Any: return to_be32(val >> 32) | (to_be32(val & 0xFFFFFFFF) << 32)
 def getbits(value: int, start: int, end: int): return (value >> start) & ((1 << (end - start + 1)) - 1)
 def i2u(bits: int, value: int): return value if value >= 0 else (1<<bits)+value
+def is_numpy_ndarray(x) -> bool: return str(type(x)) == "<class 'numpy.ndarray'>"
 def merge_dicts(ds:Iterable[dict[T,U]]) -> dict[T,U]:
   kvs = set([(k,v) for d in ds for k,v in d.items()])
   if len(kvs) != len(set(kv[0] for kv in kvs)): raise RuntimeError(f"{kvs} contains different values for the same key")
@@ -412,7 +413,7 @@ def diskcache_get(table:str, key:dict|str|int) -> Any:
   if (val:=res.fetchone()) is not None: return pickle.loads(val[0])
   return None
 
-_db_tables: set[str] = set()
+_db_tables = set()
 def diskcache_put(table:str, key:dict|str|int, val:Any, prepickled=False):
   if CACHELEVEL < 1: return val
   if isinstance(key, (str,int)): key = {"key": key}

tinygrad/llm/cli.py

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

tinygrad/llm/gguf.py

@@ -1,7 +1,8 @@
-import functools, io, pathlib, re, struct
+import functools, io, pathlib, re, struct, weakref
 from typing import Any, Callable
 
 from tinygrad.tensor import Tensor
+from tinygrad.uop.ops import UOp
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import prod, round_up
 from tinygrad.nn.state import TensorIO
@@ -20,7 +21,14 @@ _GGML_NATIVE = {0: dtypes.float32, 1: dtypes.float16, 24: dtypes.int8, 25: dtype
 _GGML_QUANT = {2:(32,18), 3:(32,20), 6:(32,22), 7:(32,24), 8:(32,34),
                12:(256,144), 13:(256,176), 14:(256,210), 18:(256,98), 21:(256,110), 22:(256,82), 23:(256,136), 39:(32,17), 41:(128,18)}
 
-def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
+_quantized_tensors:weakref.WeakKeyDictionary[UOp, tuple[UOp, int]] = weakref.WeakKeyDictionary()
+
+def get_ggml_quantization(tensor:Tensor) -> tuple[Tensor, int]|None:
+  if (meta:=_quantized_tensors.get(tensor.uop)) is None: return None
+  packed, ggml_type = meta
+  return Tensor(packed), ggml_type
+
+def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int, contiguous:bool=True) -> Tensor:
   """
   Converts ggml tensor data to a tinygrad tensor.
 
@@ -35,14 +43,14 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
   if (dtype := _GGML_NATIVE.get(ggml_type)) is not None:
     return t[:dtype.itemsize * n].contiguous().bitcast(dtype)
 
-  def q_to_uint8(t: Tensor, b: int) -> Tensor:
-    # TODO: rewrite with arange?
-    shift_tensor, bitmask = Tensor.stack(*[ Tensor(2**(i*b), device=t.device, dtype=t.dtype) for i in range(8//b) ]), 0xff >> (8 - b)
-    return t.unsqueeze(-1).expand((*t.shape,8//b)).div(shift_tensor, rounding_mode="trunc").bitwise_and(bitmask).transpose(-1, -2).flatten(-2)
+  def q_to_uint8(t:Tensor, b:int) -> Tensor:
+    shift_tensor, bitmask = Tensor.stack(*[Tensor(2**(i*b), device=t.device, dtype=t.dtype) for i in range(8//b)]), 0xff >> (8-b)
+    return t.unsqueeze(-1).expand((*t.shape, 8//b)).div(shift_tensor, rounding_mode="trunc").bitwise_and(bitmask).transpose(-1, -2).flatten(-2)
 
   if (nelements_nbytes := _GGML_QUANT.get(ggml_type)) is not None:
     from tinygrad.runtime.autogen import ggml_common as _ggml
-    blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1])).contiguous()
+    blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1]))
+    if contiguous: blocks = blocks.contiguous()
     if ggml_type == 2: return (q_to_uint8(blocks[:,2:], 4).bitcast(dtypes.int8) - 8) * blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32)
     if ggml_type == 3:
       d, m = (blocks[:,s:s+2].bitcast(dtypes.float16).cast(dtypes.float32) for s in [ 0, 2 ])
@@ -146,7 +154,14 @@ def _gguf_parse(tensor: Tensor) -> tuple[dict, dict[str, Tensor]]:
   alignment, pos = kv_data.get("general.alignment", 32), r.tell()
   data_start = round_up(pos, alignment)
 
-  state_dict = {name: ggml_data_to_tensor(tensor[data_start + off:], prod(dims), typ).reshape(*reversed(dims)) for name, dims, typ, off in t_infos}
+  state_dict = {}
+  for name, dims, typ, off in t_infos:
+    n, shape = prod(dims), tuple(reversed(dims))
+    decoded = ggml_data_to_tensor(data:=tensor[data_start + off:], n, typ).reshape(*shape)
+    if typ in _GGML_QUANT:
+      block_size, type_size = _GGML_QUANT[typ]
+      _quantized_tensors[decoded.uop] = (data[:n//block_size*type_size].uop, typ)
+    state_dict[name] = decoded
   return kv_data, state_dict
 
 def _gguf_split_paths(path: pathlib.Path, kv: dict) -> list[pathlib.Path]:

tinygrad/llm/serve.py

@@ -18,7 +18,11 @@ def parse_tool_call(s:str) -> tuple[str, typing.Any]|None:
   if (fm := re.match(r"<function=([^>]+)>\s*(.*?)\s*(?:</function>)?$", s, re.DOTALL)):
     args = {}
     for pm in re.finditer(r"<parameter=([^>]+)>(.*?)</parameter>", fm.group(2), re.DOTALL):
-      value = re.sub(r"^\r?\n|\r?\n\Z", "", pm.group(2))
+      value = pm.group(2)
+      if value.startswith("\r\n"): value = value[2:]
+      elif value.startswith("\n"): value = value[1:]
+      if value.endswith("\r\n"): value = value[:-2]
+      elif value.endswith("\n"): value = value[:-1]
       try: args[pm.group(1)] = json.loads(value)
       except json.JSONDecodeError: args[pm.group(1)] = value
     return fm.group(1), args
@@ -34,9 +38,9 @@ def normalize_messages(messages:list[dict]) -> None:
 
 class StreamRouter:
   # routes streamed output text to (field, text) deltas, keeping tool_call regions in .buf for the final parse
-  def __init__(self):
+  def __init__(self, thinking:bool=False):
     self.buf = ""
-    self.mode = "undecided"  # output inside a think block is sent as reasoning_content
+    self.mode = "reasoning" if thinking else "undecided"  # output inside a think block is sent as reasoning_content
   def split(self, tag:str, final:bool) -> tuple[str, bool]:
     # split buf on the first full tag, holding back a partial tag at the end unless final
     if tag in self.buf:
@@ -66,8 +70,9 @@ class Handler(HTTPRequestHandler):
   def do_GET(self):
     if self.path == "/v1/models": self.send_data(json.dumps({"object":"list","data":[{"id":self.server.model_name,"object":"model"}]}).encode())
     else: self.send_data((pathlib.Path(__file__).parent / "chat.html").read_bytes(), content_type="text/html")
-  def run_model(self, ids:list[int], model_name:str, include_usage=False, max_tokens:int|None=None, temperature:float=0.0):
+  def run_model(self, ids:list[int], model_name:str, include_usage=False, max_tokens:int|None=None, temperature:float=0.0, thinking:bool=False):
     model, tok = self.server.model, self.server.tok
+    prompt_tokens = len(ids)
     cache_start_pos = model.get_start_pos(ids)
     stderr_log(f"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}
@@ -77,9 +82,9 @@ class Handler(HTTPRequestHandler):
     finish_reason = "stop"
     st = time.perf_counter()
     dec = tok.stream_decoder()
-    router = StreamRouter()
+    router = StreamRouter(thinking)
     for next_id in model.generate(ids, temperature=temperature):
-      if len(out) == 0: stderr_log(f"prefill:{(len(ids)-cache_start_pos)/((pt:=time.perf_counter())-st):4.0f} tok/s  {colored('--', 'BLACK')}  ")
+      if len(out) == 0: stderr_log(f"prefill:{(prompt_tokens-cache_start_pos)/((pt:=time.perf_counter())-st):4.0f} tok/s  {colored('--', 'BLACK')}  ")
       if tok.is_end(next_id): break
       out.append(next_id)
       for field, delta in router.route(dec(next_id)): yield chunk({field:delta})
@@ -101,7 +106,9 @@ class Handler(HTTPRequestHandler):
       if finish_reason == "stop": finish_reason = "tool_calls"
     yield {"choices": [{"index":0, "delta":{},"finish_reason":finish_reason}], **tmpl}
     if include_usage:
-      yield {"choices": [], "usage": {"prompt_tokens": len(ids), "completion_tokens": len(out), "total_tokens": len(ids) + len(out)}, **tmpl}
+      yield {"choices": [], "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": len(out),
+                                        "total_tokens": prompt_tokens + len(out),
+                                        "prompt_tokens_details":{"cached_tokens":cache_start_pos}}, **tmpl}
     et = time.perf_counter()
     stderr_log(f"gen:{len(out)/(et-pt) if len(out) > 1 else 0:4.0f} tok/s  {colored('--', 'BLACK')}  "
                f"out:{len(out):5d}  {colored('--', 'BLACK')}  total:{et-st:6.2f}s\n")
@@ -115,11 +122,11 @@ 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,
+                                             enable_thinking=body.get("enable_thinking", False), preserve_thinking=True)
       ids: list[int] = self.server.tok.encode(rendered)
       stderr_log(f"prep:{(time.perf_counter()-request_st)*1e3:5.0f} ms  {colored('--', 'BLACK')}  ")
       if len(ids) >= self.server.model.max_context:
-        stderr_log(f"{colored('context length exceeded', 'red')}  in:{len(ids):5d}  max:{self.server.model.max_context:5d}\n")
         return self.send_data(json.dumps({"error":{"message":f"prompt has {len(ids)} tokens, but the model context is "
           f"{self.server.model.max_context}", "type":"invalid_request_error", "param":"messages", "code":"context_length_exceeded"}}).encode(),
           status_code=400)
@@ -127,7 +134,7 @@ 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)), thinking=body.get("enable_thinking", False))
       if body.get("stream"): self.stream_json(chunks)
       else:
         out, reasoning, tool_calls, finish_reason = [], [], [], "stop"

tinygrad/mixin/elementwise.py

@@ -49,7 +49,6 @@ class ElementwiseMixin(CreationMixin):
     """
     Returns a contiguous tensor.
     """
-    if self.dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {self.dtype}")
     uop = self._uop
     if uop.op is Ops.CONTIGUOUS or self.device is None or uop.has_buffer_identity(): return self._wrap_uop(uop)
     return self._wrap_uop(uop.alu(Ops.CONTIGUOUS, **kwargs))

tinygrad/mixin/op.py

@@ -658,7 +658,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     print(t.logsumexp(axis=1).numpy())
     ```
     """
-    m = self.max(axis=axis, keepdim=True).detach()
+    m = self.max(axis=axis, keepdim=True)
     return (self - m).exp().sum(axis=axis, keepdim=keepdim).log() + (m if keepdim else m.squeeze(axis))
 
   def _softmax(self, axis, dtype:DTypeLike|None=None) -> tuple[Self, Self, Self]:
@@ -841,7 +841,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     x = self.transpose(axis, -1)
     last_dim_size = x.shape[-1]
     x_unsqueezed = x.unsqueeze(-2).expand((None,)*(self.ndim-1)+(last_dim_size, None))
-    x_cummax = x.cummax(-1)[0].detach()
+    x_cummax, _ = x.cummax(-1)
     mask = type(self).ones(last_dim_size, last_dim_size, buffer=False).tril()
     ret = mask.where(x_unsqueezed - x_cummax.unsqueeze(-1), self.dtype.min).exp().sum(-1).log() + x_cummax
     return ret.transpose(-1, axis)

tinygrad/nn/onnx.py

@@ -2,9 +2,9 @@
 from typing import Any, Sequence, cast, Literal, NamedTuple, Generator
 import dataclasses, functools, io, math, types, warnings, pathlib, sys, os, struct, enum
 from tinygrad.nn.state import TensorIO
-from tinygrad.tensor import Tensor, is_numpy_ndarray
+from tinygrad.tensor import Tensor
 from tinygrad.mixin.op import ReductionStr
-from tinygrad.helpers import getenv, all_same, prod, flatten, make_tuple, argsort, get_single_element, polyN, Context
+from tinygrad.helpers import getenv, all_same, prod, flatten, make_tuple, argsort, is_numpy_ndarray, get_single_element, polyN, Context
 from tinygrad.dtype import DType, ConstType, dtypes, _from_np_dtype, truncate, least_upper_dtype, DTYPES_DICT
 from tinygrad.device import Device
 from tinygrad.uop.ops import sint, _broadcast_shape

tinygrad/nn/state.py

@@ -70,7 +70,6 @@ def safe_save(tensors:dict[str, Tensor], fn:str, metadata:dict[str, Any]|None=No
   nn.state.safe_save({'t':t}, "test.safetensor")
   ```
   """
-  if any(v.dtype in dtypes.weaks for v in tensors.values()): raise ValueError("safe_save requires concrete dtypes")
   headers, offset = {}, 0
   if metadata: headers['__metadata__'] = metadata
   for k,v in tensors.items():

tinygrad/renderer/cstyle.py

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

tinygrad/renderer/isa/x86.py

@@ -86,7 +86,7 @@ class X86GroupOp:
                 X86Ops.CMPi, X86Ops.IMULi, X86Ops.LEA}
 
   # X86Ops whose second src can read from memory NOTE: some of these are TwoAddress so the second src is actually the first
-  ReadMem2nd = {X86Ops.ADD, X86Ops.SUB, X86Ops.AND, X86Ops.OR, X86Ops.XOR, X86Ops.IMUL, X86Ops.CMP,
+  ReadMem2nd = {X86Ops.ADD, X86Ops.SUB, X86Ops.AND, X86Ops.OR, X86Ops.XOR, X86Ops.SHL, X86Ops.SHR, X86Ops.SAR, X86Ops.IMUL, X86Ops.CMP,
                 X86Ops.VADDSS, X86Ops.VADDSD, X86Ops.VADDPS, X86Ops.VADDPD, X86Ops.VSUBSS, X86Ops.VSUBSD, X86Ops.VSUBPS, X86Ops.VSUBPD,
                 X86Ops.VMULSS, X86Ops.VMULSD, X86Ops.VMULPS, X86Ops.VMULPD, X86Ops.VDIVSS, X86Ops.VDIVSD, X86Ops.VDIVPS, X86Ops.VDIVPD,
                 X86Ops.VPADDB, X86Ops.VPADDW, X86Ops.VPADDD, X86Ops.VPADDQ, X86Ops.VPSUBB, X86Ops.VPSUBW, X86Ops.VPSUBD, X86Ops.VPSUBQ,
@@ -99,9 +99,8 @@ class X86GroupOp:
 
   # X86Ops that can write to memory
   WriteMem = {X86Ops.MOVm, X86Ops.MOVi, X86Ops.VMOVSSm, X86Ops.VMOVSDm, X86Ops.VMOVUPSm, X86Ops.VMOVDm, X86Ops.VMOVQm,
-              X86Ops.ADDi, X86Ops.SUBi, X86Ops.ANDi, X86Ops.ORi, X86Ops.XORi, X86Ops.SHL, X86Ops.SHLi, X86Ops.SHR, X86Ops.SHRi, X86Ops.SAR,
-              X86Ops.SARi, X86Ops.SETNE, X86Ops.SETE, X86Ops.SETL, X86Ops.SETB,
-              X86Ops.VCVTPS2PH, X86Ops.VPEXTRB, X86Ops.VPEXTRW, X86Ops.VPEXTRD, X86Ops.VPEXTRQ}
+              X86Ops.ADDi, X86Ops.SUBi, X86Ops.ANDi, X86Ops.ORi, X86Ops.XORi, X86Ops.SHLi, X86Ops.SHRi, X86Ops.SARi, X86Ops.SETNE,
+              X86Ops.SETE, X86Ops.SETL, X86Ops.SETB, X86Ops.VCVTPS2PH, X86Ops.VPEXTRB, X86Ops.VPEXTRW, X86Ops.VPEXTRD, X86Ops.VPEXTRQ}
 
   # X86Ops that read flags
   ReadFlags = {X86Ops.CMOVB, X86Ops.CMOVL, X86Ops.CMOVE, X86Ops.CMOVNE, X86Ops.SETB, X86Ops.SETL, X86Ops.SETE, X86Ops.SETNE, X86Ops.JB, X86Ops.JL,
@@ -273,11 +272,6 @@ def idiv(ctx:IselContext, x:UOp) -> UOp:
   # this move "cleanses" the register constraints (rax/rdx) of idiv as that only applies on definition and not on the uses of idiv
   return x.ins(X86Ops.MOV, src=(idiv,))
 
-# a variable shift count implicitly reads cl so it goes in rcx, the shifted value can't be in rcx
-def shift(x:UOp, op:X86Ops) -> UOp:
-  val = x.ins(X86Ops.MOV, src=(x.src[0],), tag=tuple(r for r in WGPR if r is not RCX))
-  return x.ins(op, src=(val, x.ins(X86Ops.MOV, src=(x.src[1],), tag=(RCX,))))
-
 # a memory address operand is (base, index, displacement, size). size is the element size, it scales the index and is the memory operand width.
 # it is materialized as an immediate so the address stays correct if the base register is ever spilled and refilled
 def fold_address(x:UOp) -> tuple[UOp, UOp, UOp, UOp]:
@@ -458,9 +452,9 @@ isel_matcher = PatternMatcher([
   (UPat.var("a", dtypes.ints+(dtypes.bool,)) ^ UPat.cvar("c"), lambda a,c: a.ins(X86Ops.XORi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
   (UPat(Ops.SUB, dtypes.ints, (UPat.var("a"), UPat.cvar("c"))), lambda a,c: a.ins(X86Ops.SUBi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
   # scalar int binary with register
-  ((UPat(dtype=dtypes.ints) << UPat()).named("x"), lambda x: shift(x, X86Ops.SHL)),
-  ((UPat(dtype=dtypes.uints) >> UPat()).named("x"), lambda x: shift(x, X86Ops.SHR)),
-  ((UPat(dtype=dtypes.sints) >> UPat()).named("x"), lambda x: shift(x, X86Ops.SAR)),
+  (UPat.var("a", dtypes.ints) << UPat.var("b"), lambda a,b: a.ins(X86Ops.SHL, src=(a, b))),
+  (UPat.var("a", dtypes.uints) >> UPat.var("b"), lambda a,b: a.ins(X86Ops.SHR, src=(a, b))),
+  (UPat.var("a", dtypes.sints) >> UPat.var("b"), lambda a,b: a.ins(X86Ops.SAR, src=(a, b))),
   (UPat.var("a", dtypes.ints) + UPat.var("b"), lambda a,b: a.ins(X86Ops.ADD, src=(a, b))),
   (UPat.var("a", dtypes.ints) * UPat.var("b"), lambda a,b: a.ins(X86Ops.IMUL, src=(a, b))),
   (UPat.var("a", dtypes.ints+(dtypes.bool,)) & UPat.var("b"), lambda a,b: a.ins(X86Ops.AND, src=(a, b))),
@@ -657,8 +651,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
   if x.arg in X86GroupOp.WriteMem:
     if len(x.src) > 4: address, rest = x.src[:4], x.src[4:]
     else: address, rest = (x, None, None, None), x.src
-    imm_uop = rest[:1] if rest and rest[0].op is Ops.CONST else (None,)
-    return _encode(rest[0], *address, *(None, *rest[1:])) if reg is None else _encode(None, *address, *(None, *imm_uop))
+    return _encode(rest[0], *address, *(None, *rest[1:])) if reg is None else _encode(None, *address, *(None, *rest[:1]))
 
   if x.arg in X86GroupOp.Rm1st:
     if len(x.src) > 3: address, rest = x.src[:4], x.src[4:]
@@ -711,9 +704,8 @@ encodings = {
   # int division
   X86Ops.IDIV: lambda x: encode(x, 0xF7, reg=7), X86Ops.DIV: lambda x: encode(x, 0xF7, reg=6),
   # scalar int binary
-  X86Ops.SHL: lambda x: encode(x, 0xD3, reg=4), X86Ops.SHLi: lambda x: encode(x, 0xC1, reg=4),
-  X86Ops.SHR: lambda x: encode(x, 0xD3, reg=5), X86Ops.SHRi: lambda x: encode(x, 0xC1, reg=5),
-  X86Ops.SAR: lambda x: encode(x, 0xD3, reg=7), X86Ops.SARi: lambda x: encode(x, 0xC1, reg=7),
+  X86Ops.SHLi: lambda x: encode(x, 0xC1, reg=4),
+  X86Ops.SHRi: lambda x: encode(x, 0xC1, reg=5), X86Ops.SARi: lambda x: encode(x, 0xC1, reg=7),
   X86Ops.ADD: lambda x: encode(x, 0x03), X86Ops.ADDi: lambda x: encode(x, 0x81, reg=0),
   X86Ops.SUB: lambda x: encode(x, 0x2B), X86Ops.SUBi: lambda x: encode(x, 0x81, reg=5),
   X86Ops.AND: lambda x: encode(x, 0x23), X86Ops.ANDi: lambda x: encode(x, 0x81, reg=4),
@@ -792,7 +784,6 @@ class X86Renderer(ISARenderer):
   post_regalloc_matcher = post_regalloc_matcher
   code_for_op = {x: lambda: None for x in (Ops.SQRT, Ops.AND, Ops.OR, Ops.SHL, Ops.SHR, Ops.NEG, Ops.SUB, Ops.FDIV, Ops.CMPLT, Ops.CMPEQ)}
   def __init__(self, target:Target):
-    if target.arch.split(",")[0] != "x86_64": raise RuntimeError(f"X86Renderer only supports x86_64, got {target.arch}")
     super().__init__(target)
     from tinygrad.runtime.support.compiler_cpu import X86Compiler
     self.compiler = X86Compiler()

tinygrad/renderer/llvmir.py

@@ -35,7 +35,7 @@ def lcast(input_type:DType, output_type:DType):
 
 def render_wmma_amd(ctx, wmma: UOp, cdna=False) -> str:
   dt_map = {dtypes.half: "f16", dtypes.float: "f32", dtypes.ushort: "bf16.1k" if cdna else "bf16", dtypes.bfloat16: "bf16.1k" if cdna else "bf16",
-            dtypes.fp8e4m3: ".fp8.fp8", dtypes.fp8e5m2: ".bf8.bf8", dtypes.int8: "iu8", dtypes.int32: "i32"}
+            dtypes.fp8e4m3: ".fp8.fp8", dtypes.fp8e5m2: ".bf8.bf8"}
   # https://github.com/llvm/llvm-project/blob/main/clang/test/CodeGenOpenCL/builtins-amdgcn-mfma.cl
   N,M,K = wmma.arg[0]
   if cdna:
@@ -44,10 +44,9 @@ def render_wmma_amd(ctx, wmma: UOp, cdna=False) -> str:
            f".{N}x{M}x{K}{dt_map[wmma.arg[1]]}(" + ", ".join([f"{ldt(w.dtype, w.max_numel())} {ctx[w]}" for w in wmma.src]) + ", i32 0, i32 0, i32 0)"
   # https://github.com/llvm/llvm-project/blob/main/llvm/test/CodeGen/AMDGPU/GlobalISel/llvm.amdgcn.wmma_32.ll
   # example: %wmma0 = call <8 x float> @llvm.amdgcn.wmma.f32.16x16x16.f16(<16 x half> %v99,<16 x half> %v100,<8 x float> %v101)
-  args = [f"{ldt(w.dtype, w.max_numel())} {ctx[w]}" for w in wmma.src]
-  if wmma.arg[1] == dtypes.int8: args = ["i1 true", args[0], "i1 true", args[1], args[2]]  # iu8 flags A/B signed
   return f"  {ctx[wmma]} = call {ldt(wmma.dtype, wmma.max_numel())} @llvm.amdgcn.wmma.{dt_map[wmma.src[-1].dtype]}.16x16x16." + \
-    f"{dt_map[wmma.arg[1]]}(" + ", ".join(args) + (", i1 false)" if wmma.dtype != dtypes.float else ")")
+    f"{dt_map[wmma.src[0].dtype]}(" + ", ".join([f"{ldt(w.dtype, w.max_numel())} {ctx[w]}" for w in wmma.src]) + (", i1 false)" \
+      if wmma.dtype != dtypes.float else ")")
 
 # llvm ops, lop[<dtype>][<op>]
 unsigned_lop = { Ops.ADD: "add", Ops.MUL: "mul", Ops.CDIV: "udiv", Ops.CMOD: "urem",
@@ -262,9 +261,6 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n  %trunc
       ])
     if target.arch in {"gfx1100", "gfx1151"}:
       self.extra_matcher += PatternMatcher([
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.int32), lambda x: x.replace(
-          src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
-          if x.src[0].dtype == dtypes.int8 and x.src[0].max_numel() == 16 else None),
         (UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(x.replace(
           src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(j//2) if j%2 == 0 else UOp.const(x.src[2].dtype, 0.0)
             for j in range(x.max_numel()*2)))),

tinygrad/runtime/ops_qcom.py

@@ -278,7 +278,7 @@ class QCOMProgram(HCQProgram):
   def _parse_lib(self, lib):
     # Extract image binary
     self.image_size = _read_lib(lib, 0x100)
-    self.image = lib[(image_offset:=_read_lib(lib, 0xc0)):image_offset+self.image_size]
+    self.image = bytearray(lib[(image_offset:=_read_lib(lib, 0xc0)):image_offset+self.image_size])
 
     # Parse image descriptors
     image_desc_off = _read_lib(lib, 0x110)

tinygrad/schedule/multi.py

@@ -148,7 +148,7 @@ def rewrite_into_function(call:UOp):
 
 def param_to_multi(p:UOp):
   if p.axis is None: return None
-  return UOp.param(p.arg.slot, p.dtype, p.shard_shape, p.device, p.arg.vmin_vmax, p.arg.multiple_of, p.arg.name, p.arg.addrspace).multi(p.axis)
+  return UOp.param(p.arg.slot, p.dtype, p.shard_shape, p.device, p.arg.vmin_vmax, p.arg.name, p.arg.addrspace).multi(p.axis)
 
 # NOTE: this is the same pattern as unrolled ranges
 multi_pm = PatternMatcher([

tinygrad/schedule/rangeify.py

@@ -434,7 +434,6 @@ class LocalAddBufferContext:
   opts:tuple|None = None
 
 def debuf(ctx:LocalAddBufferContext, buf:UOp):
-  if buf.addrspace != AddrSpace.GLOBAL: return None
   param = UOp(Ops.PARAM, src=(UOp.const(dtypes.int, prod(buf.max_shape)),),
               arg=ParamArg(ctx.dg, buf.dtype, addrspace=buf.addrspace, device=buf.device))
   ret = param.reshape(buf.max_shape)
@@ -471,7 +470,7 @@ to_define_global = PatternMatcher([
   (UPat(Ops.STORE, name="x"), find_bufs),
   (UPat((Ops.BUFFER, Ops.MSTACK, Ops.MSELECT), name="buf"), debuf),
   (UPat(Ops.PARAM, name="v"), lambda v:
-   UOp.variable(v.arg.name, v.arg.vmin_vmax[0], v.arg.vmin_vmax[1], v.dtype, multiple_of=v.arg.multiple_of)
+   UOp.variable(v.arg.name, v.arg.vmin_vmax[0], v.arg.vmin_vmax[1], v.dtype)
    if v.arg.name is not None and v.arg.vmin_vmax is not None else None),
 
   # this renumbers the params
@@ -523,7 +522,7 @@ def split_store(x:UOp) -> UOp|None:
   else: ret = ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts))
 
   kernel = ret.call(*lctx.map.values(), *lctx.vars.keys())
-  if ret.op is Ops.SINK and not all_same([x.device for x in kernel.src[1:] if x.op is not Ops.BIND and x.device is not None]):
+  if ret.op is Ops.SINK and not all_same([x.device for x in kernel.src[1:] if x.op is not Ops.BIND]):
     raise RuntimeError(f"all buffers must be on the same device: {tuple(b.buf_uop for b in kernel.src[1:])}")
   return kernel
 

tinygrad/tensor.py

@@ -1,10 +1,10 @@
 # inspired by https://github.com/karpathy/micrograd/blob/master/micrograd/engine.py
 from __future__ import annotations
 import time, functools, sys, inspect, pathlib, hashlib, weakref
-from typing import Any, Callable, cast, get_args, ParamSpec, TypeGuard, TypeVar, Generic, TYPE_CHECKING
+from typing import Any, Callable, cast, get_args, ParamSpec, TypeVar, Generic, TYPE_CHECKING
 if TYPE_CHECKING: import numpy
-from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, least_upper_dtype, to_dtype, strong_dtype, _from_np_dtype, _to_np_dtype, PyConst
-from tinygrad.helpers import all_int, getenv, fully_flatten, fetch, Metadata, TRACEMETA, TracingKey
+from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, to_dtype, strong_dtype, _from_np_dtype, _to_np_dtype, PyConst
+from tinygrad.helpers import all_int, getenv, fully_flatten, fetch, Metadata, TRACEMETA, is_numpy_ndarray, TracingKey
 from tinygrad.helpers import cpu_profile, suppress_finalizing, disable_gc
 from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, _index_to_concrete_int, Variable, _broadcast_shape
 from tinygrad.mixin.rand import RandMixin
@@ -34,8 +34,6 @@ def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str) -> None:
 
 # **** Tensor helper functions ****
 
-def is_numpy_ndarray(x) -> "TypeGuard[numpy.ndarray]": return str(type(x)) == "<class 'numpy.ndarray'>"
-
 def _fromnp(x: 'numpy.ndarray') -> UOp:
   ret = UOp.new_buffer("NPY", x.size, _from_np_dtype(x.dtype))
   # fake realize
@@ -77,21 +75,22 @@ class Tensor(RandMixin):
       data = UOp.const(_dtype or dtypes.default_float, 0)
     elif isinstance(data, get_args(ConstType)):
       data = UOp.const(_dtype or dtypes.from_py(data), data)
-    elif is_numpy_ndarray(data) and data.shape == ():
-      data = UOp.const(_dtype or _from_np_dtype(data.dtype), data.item())
-    else:
-      if _dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {_dtype}")
-      if isinstance(data, bytes): data = UOp._frompy(data, _dtype or dtypes.uint8, _device)
-      elif isinstance(data, (list, tuple)):
-        if _dtype is None:
-          if (d := fully_flatten(data)) and all(isinstance(s, bool) for s in d): _dtype = dtypes.bool
-          else: _dtype = dtypes.default_int if d and all_int(d) else dtypes.default_float  # NOTE: this works because all_int([True, False]) is True
-        data = UOp._frompy(data, _dtype, _device)
-      elif is_numpy_ndarray(data):
+    elif isinstance(data, bytes): data = UOp._frompy(data, _dtype or dtypes.uint8, _device)
+    elif isinstance(data, (list, tuple)):
+      if _dtype is None:
+        if (d := fully_flatten(data)) and all(isinstance(s, bool) for s in d): _dtype = dtypes.bool
+        else: _dtype = dtypes.default_int if d and all_int(d) else dtypes.default_float  # NOTE: this works because all_int([True, False]) is True
+      data = UOp._frompy(data, _dtype, _device)
+    elif is_numpy_ndarray(data):
+      import numpy as np
+      assert isinstance(data, np.ndarray), f"expected np.ndarray, got {data}"
+      if data.shape == ():
+        data = UOp.const(_dtype or _from_np_dtype(data.dtype), data.item())
+      else:
         data = _fromnp(data.astype(npdtype) if _dtype is not None and (npdtype:=_to_np_dtype(_dtype)) is not None else data)
-      elif isinstance(data, pathlib.Path):
-        _dtype = _dtype or dtypes.uint8
-        data = UOp.new_buffer(f"DISK:{data.resolve()}", data.stat().st_size // _dtype.itemsize, _dtype)
+    elif isinstance(data, pathlib.Path):
+      _dtype = _dtype or dtypes.uint8
+      data = UOp.new_buffer(f"DISK:{data.resolve()}", data.stat().st_size // _dtype.itemsize, _dtype)
 
     # by this point, it has to be a UOp
     if not isinstance(data, UOp): raise RuntimeError(f"can't create Tensor from {data!r} with type {type(data)}")
@@ -179,7 +178,6 @@ class Tensor(RandMixin):
 
   def linear_with_vars(self, *lst:Tensor) -> tuple[UOp, dict[str, int]]:
     """Creates the LINEAR UOp needed to realize these Tensor(s), with Variables."""
-    if any(t.dtype in dtypes.weaks for t in (self,)+lst): raise RuntimeError("cannot realize a weak dtype; cast to a concrete dtype first")
     big_sink, becomes_map = transform_to_call(UOp.sink(*[x.uop for x in (self,)+lst]))
     _apply_map_to_tensors(becomes_map, name="buffers")
     return create_linear_with_vars(big_sink)
@@ -207,16 +205,14 @@ class Tensor(RandMixin):
     return self
 
   def assign(self, x:Tensor|PyConst|list|tuple) -> Tensor:
-    if self.dtype in dtypes.weaks: raise RuntimeError("cannot assign into a weak tensor; it has no storage")
     is_disk = isinstance(self.device, str) and self.device.startswith(("DISK", "TINYFS"))
     if not isinstance(x, Tensor): x = Tensor(x, device="CPU" if is_disk else self.device, dtype=self.dtype)
     if self.uop is x.uop: return self  # a self assign is a NOOP
     # broadcast x (shape only, dtype must match)
     x = x._broadcast_to(self.shape)
-    if x.dtype in dtypes.weaks: x = x.cast(least_upper_dtype(self.dtype, x.dtype))
-    if x.dtype != self.dtype: raise RuntimeError(f"assign dtype mismatch {self.dtype} != {x.dtype}")
     if not is_disk and x.uop.device is not None and self.device is not None and self.device != x.device:
       raise RuntimeError(f"assign device mismatch {self.device} != {x.device}")
+    if not is_disk and self.dtype != x.dtype: raise RuntimeError(f"assign dtype mismatch {self.dtype} != {x.dtype}")
     if isinstance(self.device, tuple) and x.uop.device is not None and self.uop.axis != x.uop.axis:
       raise RuntimeError(f"multi axis mismatch {self.uop.axis} != {x.uop.axis}")
 
@@ -257,7 +253,6 @@ class Tensor(RandMixin):
     print(np.frombuffer(t.data(), dtype=np.int32))
     ```
     """
-    if self.dtype in dtypes.weaks: return self.cast(strong_dtype(self.dtype)).data()
     if 0 in self.shape: return memoryview(bytearray(0)).cast(self.dtype.fmt)  # type: ignore[arg-type,return-value]
     assert all_int(self.shape), f"no data if shape is symbolic, {self.shape=}"
     buf = self._buffer()
@@ -281,7 +276,6 @@ class Tensor(RandMixin):
     print(t.tolist())
     ```
     """
-    if self.dtype in dtypes.weaks: return self.cast(strong_dtype(self.dtype)).tolist()
     # TODO: remove half once minimum python supports it
     if self.dtype in (dtypes.half, dtypes.bfloat16, *dtypes.fp8s): return self.cast(dtypes.float32).tolist()
     if 0 in self.shape:
@@ -299,7 +293,6 @@ class Tensor(RandMixin):
     print(repr(t.numpy()))
     ```
     """
-    if self.dtype in dtypes.weaks: return self.cast(strong_dtype(self.dtype)).numpy()
     assert all_int(self.shape), f"no data if shape is symbolic, {self.shape=}"
     import numpy as np
     if self.dtype in { dtypes.bfloat16, *dtypes.fp8s }: return self.float().numpy()
@@ -457,10 +450,7 @@ class Tensor(RandMixin):
   def _rop(self, op:Ops, axis:tuple[int, ...]) -> Tensor: return self._apply_uop(UOp._rop, op=op, axis=axis)
 
   def __setitem__(self, indices, v:Tensor|PyConst|list|tuple) -> None:
-    if self.dtype in dtypes.weaks: raise RuntimeError("cannot setitem into a weak tensor; it has no storage")
-    if isinstance(v, Tensor):
-      if v.dtype in dtypes.weaks: v = v.cast(least_upper_dtype(self.dtype, v.dtype))
-      if v.dtype != self.dtype: raise RuntimeError(f"setitem dtype mismatch: {self.dtype=} != {v.dtype=}")
+    if isinstance(v, Tensor) and v.dtype != self.dtype: raise RuntimeError(f"setitem dtype mismatch: {self.dtype=} != {v.dtype=}")
     # raise if mutation would diverge from eager (allow only pure views of a realized buffer; exclude +=/-= RHS via v_uop/v_bw)
     v_uop, v_bw = (v.uop, v.uop.backward_slice) if isinstance(v, Tensor) else (None, {})
     if self.uop.op_in_backward_slice_with_self(Ops.BUFFER):

tinygrad/uop/divandmod.py

@@ -11,8 +11,6 @@ def fold_divmod_general(d: UOp) -> UOp|None:
   if y.vmin==y.vmax==0: raise ZeroDivisionError(f"{'Division' if d.op is Ops.FLOORDIV else 'Mod'} by zero trying to rewrite {x.alu(d.op, y)}")
   # x//y is constant
   if (xdiv:=x//y).vmin == xdiv.vmax: return x - xdiv.vmin*y if d.op is Ops.FLOORMOD else xdiv.const_like(xdiv.vmin)
-  # PARAM // c is irreducible
-  if x.op is Ops.PARAM and y.op is Ops.CONST and x.arg.multiple_of % y.arg == 0: return d.const_like(0) if d.op is Ops.FLOORMOD else None
 
   # split uops for the rest of the processing
   x_peeled, const = x.pop_const()

tinygrad/uop/ops.py

@@ -4,7 +4,7 @@ import sys, time, functools, itertools, math, operator, hashlib, os, types, pick
 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
+from tinygrad.dtype import ConstType, dtypes, DType, DTypeLike, truncate, least_upper_dtype, least_upper_float, strong_dtype, Invalid, AddrSpace
 from tinygrad.dtype import ConstFloat, PyConst, InvalidType, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar
 from tinygrad.device import Buffer, MultiBuffer, canonicalize_device
 from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA
@@ -24,13 +24,12 @@ class ParamArg:
   slot: int
   dtype: DType
   vmin_vmax: tuple[PyConst, PyConst]|None = None
-  multiple_of: int|None = None
   name: str|None = None
   addrspace: AddrSpace|None = AddrSpace.GLOBAL
   axis: int|None = None
   device: str|tuple[str, ...]|None = None
   def __repr__(self):
-    fields = (("vmin_vmax", None), ("multiple_of", None), ("name", None), ("addrspace", AddrSpace.GLOBAL), ("axis", None), ("device", None))
+    fields = (("vmin_vmax", None), ("name", None), ("addrspace", AddrSpace.GLOBAL), ("axis", None), ("device", None))
     args = [repr(self.slot), repr(self.dtype)] + [f"{k}={v!r}" for k,default in fields if (v:=getattr(self, k)) != default]
     return f"ParamArg({', '.join(args)})"
 axis_letters = {AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L", AxisType.UPCAST: "u",
@@ -705,7 +704,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   def copy_to_device(self, device:str|tuple[str, ...], arg=None):
     assert arg is None or isinstance(self.device, tuple)
     inp = self if arg is None else UOp(Ops.MSELECT, src=(self,), arg=arg)
-    if inp.dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {inp.dtype}")
     return UOp(Ops.COPY, src=(inp,), arg=device)
   def mselect(self, arg:int) -> UOp: return UOp(Ops.MSELECT, src=(self,), arg=arg)
   def mstack(self, *srcs: UOp) -> UOp: return UOp(Ops.MSTACK, src=(self,)+srcs)
@@ -760,7 +758,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   @staticmethod
   def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType, num=None):
-    if dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {dtype}")
     slot = next(UOp.unique_num) if num is None else num
     return UOp(Ops.BUFFER, src=(shape_to_shape_arg((size,)),), arg=ParamArg(slot, dtype, device=device))
   @staticmethod
@@ -788,7 +785,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     return ret if ret.device == device else ret.copy_to_device(device)
   def clone(self, device=None) -> UOp:
     device = device or self.device
-    ret = self.empty_like(device=device)
+    ret = self.empty_like(dtype=strong_dtype(self.dtype), device=device)
     src = self if self.device is None or self.device == device else self.copy_to_device(device)
     return ret.after(ret.store(src.cast(ret.dtype)))
   @recursive_property
@@ -812,10 +809,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     if self.op is Ops.BUFFER: return self.arg.addrspace
     if self.op in {Ops.SPECIAL, Ops.RANGE}: return AddrSpace.ALU
     if self.op is Ops.LOAD: return AddrSpace.ALU # LOAD brings things into the ALU
-    if self.op in {Ops.INDEX, Ops.CAST, Ops.AFTER, Ops.REDUCE, Ops.STORE, Ops.MSTACK, Ops.MSELECT, Ops.END}:
+    if self.op in {Ops.INDEX, Ops.CAST, Ops.AFTER, Ops.REDUCE, Ops.STORE, Ops.MSTACK, Ops.MSELECT}:
       return self.src[0].addrspace
     if self.op in GroupOp.Movement: return self.src[0].addrspace
-    if self.op in {Ops.STACK, Ops.WMMA, Ops.GROUP} or self.op in GroupOp.Elementwise:
+    if self.op in {Ops.STACK, Ops.WMMA} or self.op in GroupOp.Elementwise:
       ad = [x.addrspace for x in self.src if x.addrspace is not None]
       if not len(ad) or not all_same(ad): return None
       return ad[0]
@@ -916,9 +913,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   # *** uop Variable stuff ***
 
   @staticmethod
-  def variable(name:str, min_val:PyConst, max_val:PyConst, dtype:DType=dtypes.index, multiple_of:int=1) -> UOp:
+  def variable(name:str, min_val:PyConst, max_val:PyConst, dtype:DType=dtypes.index) -> UOp:
     return UOp(Ops.PARAM, src=(shape_to_shape_arg(()),),
-               arg=ParamArg(-1, dtype, name=name, vmin_vmax=(min_val, max_val), multiple_of=multiple_of, addrspace=AddrSpace.ALU))
+               arg=ParamArg(-1, dtype, name=name, vmin_vmax=(min_val, max_val), addrspace=AddrSpace.ALU))
   @property
   def expr(self) -> str:
     assert self.op is Ops.PARAM
@@ -927,7 +924,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     assert self.op is Ops.PARAM and self.addrspace is AddrSpace.ALU, f"op is {self.op}, need PARAM"
     uval = self.const_like(val) if isinstance(val, int) else val
     assert self.vmin <= uval.vmin and uval.vmax <= self.vmax, f"bind {val} not in range [{self.vmin}, {self.vmax}]"
-    assert uval.divides(self.arg.multiple_of) is not None, f"bind {val} not divisible by {self.arg.multiple_of}"
     return UOp(Ops.BIND, src=(self, uval))
   def unbind(self) -> tuple[Variable, int]:
     assert self.op is Ops.BIND and self.src[0].op is Ops.PARAM and self.src[1].op is Ops.CONST, f"can't unbind {self}"
@@ -950,7 +946,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     if self.op is Ops.STACK: return math.gcd(*[x.const_factor() for x in self.src])
     if self.op is Ops.ADD: return math.gcd(self.src[0].const_factor(), self.src[1].const_factor())
     if self.op is Ops.MUL: return self.src[0].arg if self.src[0].op is Ops.CONST else self.src[1].arg if self.src[1].op is Ops.CONST else 1
-    if self.op is Ops.PARAM and self.arg.multiple_of is not None: return self.arg.multiple_of
     return 1
   def divides(self, v:int) -> UOp|None:
     if v==1: return self
@@ -962,7 +957,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     if self.op is Ops.MUL:
       if (d0:=self.src[0].divides(v)) is not None: return d0 * self.src[1]
       if (d1:=self.src[1].divides(v)) is not None: return self.src[0] * d1
-    if self.op is Ops.PARAM and self.arg.multiple_of is not None: return self // v if self.arg.multiple_of%v == 0 else None
     return None # generic None if we aren't sure
   def pop_const(self, op=Ops.ADD) -> tuple[UOp, PyConst]:  # NOTE: assume Invalid ALU is resolved
     return (self.src[0], self.src[1].arg) if self.op is op and self.src[1].op is Ops.CONST else (self, identity_element(op, self.dtype))
@@ -1033,12 +1027,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     if self.op is Ops.STACK: return min(x.vmin for x in self.src), max(x.vmax for x in self.src)
     if self.op is Ops.CONST and self.arg is not Invalid: return self.arg, self.arg
     if self.op is Ops.INDEX: return self.src[0]._min_max
-    if self.op is Ops.CAST:
-      # a cast to unsigned keeps exact bounds when the source fits
-      # TODO: can do more based on new dtype window
-      if dtypes.is_unsigned(self.dtype) and 0 <= self.src[0].vmin and self.src[0].vmax <= self.dtype.max: return self.src[0]._min_max
-      if self.dtype in dtypes.floats+dtypes.sints+(dtypes.index,):
-        return max(self.dtype.min, self.src[0].vmin), min(self.src[0].vmax, self.dtype.max)
+    # TODO: CAST to bool/unsigned is not monotone, still some case can be simplified
+    if self.op is Ops.CAST and self.dtype in dtypes.floats+dtypes.sints+(dtypes.index,):
+      return max(self.dtype.min, self.src[0].vmin), min(self.src[0].vmax, self.dtype.max)
     return self.dtype.min, self.dtype.max
 
   @functools.cached_property
@@ -1091,16 +1082,16 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   # TODO: this should replace placeholder
   @staticmethod
-  def param(slot:int, dtype:DType, shape:tuple[sint, ...]|None=None, device=None, vmin_vmax:tuple[PyConst, PyConst]|None=None,
-            multiple_of:int|None=None, name=None, addrspace=AddrSpace.GLOBAL, axis:int|None=None):
-    if dtype in dtypes.weaks: raise RuntimeError(f"cannot create param for weak dtype {dtype}")
+  def param(slot:int, dtype:DType, shape:tuple[sint, ...]|None=None, device=None, vmin_vmax:tuple[PyConst, PyConst]|None=None, name=None,
+            addrspace=AddrSpace.GLOBAL, axis:int|None=None):
     if shape is not None and axis is not None and isinstance(device, tuple):
       shape = tuple(s*len(device) if i == axis else s for i,s in enumerate(shape))
     src: tuple[UOp, ...] = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),)
-    return UOp(Ops.PARAM, src=src, arg=ParamArg(slot, dtype, vmin_vmax, multiple_of, name, addrspace, axis, device))
+    return UOp(Ops.PARAM, src=src, arg=ParamArg(slot, dtype, vmin_vmax, name, addrspace, axis, device))
   def param_like(self, slot:int):
     addrspace = self.addrspace if self.addrspace is not None else AddrSpace.GLOBAL
-    if self.op is Ops.BIND: return self.src[0].replace(arg=replace(self.src[0].arg, slot=slot, addrspace=addrspace))
+    if self.op is Ops.BIND:
+      return UOp.param(slot, self.dtype, self._shape, self.device, cast(tuple[int, int], self._min_max), self.src[0].expr, addrspace)
     return UOp.param(slot, self.dtype, self.shard_shape if self.axis is not None else self._shape, self.device, addrspace=addrspace, axis=self.axis)
 
   @staticmethod

tinygrad/uop/spec.py

@@ -61,13 +61,15 @@ spec_shared = PatternMatcher([
   (UPat(Ops.STACK, src=(UPat(),), allow_any_len=True, name="s"),
    lambda s: all_same([x.shape for x in s.src]) and all(x.dtype == s.dtype for x in s.src)),
 
-  # ALUs: operands match the result dtype, except comparisons/WHERE; renderer-lowered shifts may use a uint32 count
+  # ALUs: most ALUs have all matching dtypes, except CMPLT, CMPNE, and WHERE
   # a weak dtype matches any dtype (TODO: make python scalars weak consts)
-  (UPat(Ops.WHERE, name="w", src=(UPat(dtype=dtypes.bool), UPat(), UPat())),
-   lambda w: all(s.dtype == w.dtype or s.dtype in dtypes.weaks for s in w.src[1:])),
-  (UPat(GroupOp.Comparison, dtype=dtypes.bool, src=(UPat.var("x"), UPat.var("y"))),
+  (UPat(Ops.WHERE, name="w", src=(UPat(dtype=dtypes.bool), UPat.var("x"), UPat.var("y"))),
+   lambda w,x,y: all(s.dtype == w.dtype or s.dtype in dtypes.weaks for s in (x,y))),
+  (UPat((Ops.CMPLT, Ops.CMPNE, Ops.CMPEQ), dtype=dtypes.bool, src=(UPat.var("x"), UPat.var("y"))),
    lambda x,y: x.dtype == y.dtype or x.dtype in dtypes.weaks or y.dtype in dtypes.weaks),
-  (UPat((Ops.SHL, Ops.SHR), src=(UPat.var("x"), UPat(dtype=dtypes.uint)), name="a"), lambda a,x: a.dtype == x.dtype or None),
+  # and SHL/SHR, the shift distance can be an int
+  (UPat((Ops.SHL, Ops.SHR), src=(UPat.var("x"), UPat.var("y")), name="a"),
+   lambda a,x,y: a.dtype == x.dtype and (y.dtype in (x.dtype, dtypes.uint) or y.dtype in dtypes.weaks)),
   (UPat((Ops.CDIV, Ops.CMOD, Ops.FLOORDIV, Ops.FLOORMOD), name="x"), lambda x: None if dtypes.is_int(x.dtype) else False),
   (UPat(GroupOp.ALU, name="x"), lambda x: all(y.dtype == x.dtype or y.dtype in dtypes.weaks for y in x.src)),
 

tinygrad/viz/index.html

@@ -420,7 +420,7 @@
   <div class="main-container">
     <div class="floating-container">
       <button class="btn collapse-btn">
-        <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="1.5" width="20"><rect x="3" y="4" width="18" height="16" rx="2"/><path d="M8 4v16M16 4v16"/></svg>
+        <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" width="20"><path d="M15 19l-7-7 7-7"/></svg>
       </button>
       <button class="btn" id="zoom-to-fit-btn" aria-label="Fit graph">
         <svg xmlns="http://www.w3.org/2000/svg" fill="none" viewBox="0 0 24 24" stroke-width="1.5" stroke="currentColor" width="20">

tinygrad/viz/js/index.js

@@ -735,7 +735,6 @@ async function renderProfiler(path, opts) {
     }
   }
 
-  let lastCanvasRect = null;
   function resize() {
     const [width, height] = canvasDims();
     if (canvas.width === width*dpr && canvas.height === height*dpr) return;
@@ -744,11 +743,6 @@ async function renderProfiler(path, opts) {
     canvas.style.height = `${height}px`;
     canvas.style.width = `${width}px`;
     ctx.scale(dpr, dpr);
-    const newRect = rect(canvas);
-    if (lastCanvasRect != null && lastCanvasRect.width > 0) {
-      zoomLevel = d3.zoomIdentity.translate(zoomLevel.x+lastCanvasRect.left-newRect.left, 0).scale(zoomLevel.k*lastCanvasRect.width/width);
-    }
-    lastCanvasRect = { left:newRect.left, width };
     d3.select(canvas).call(canvasZoom.transform, zoomLevel);
   }