Merge branch 'master' into modern_scan

* llm: minor fixes + tests

* error

.github/workflows/benchmark.yml

@@ -521,6 +521,8 @@ jobs:
       run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
     - name: openpilot run_pickle 0.10.1 driving_vision
       run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py
+    - name: Test copy speeds
+      run: SIZE=64e6 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
 
   driverbenchmarks:
     name: PCI Driver Benchmark (DEV=${{ matrix.dev }})

.github/workflows/szdiff.yml

@@ -14,12 +14,15 @@ 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: |
@@ -51,6 +54,9 @@ 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,6 +1462,8 @@ 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, 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,...]:
+           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,...]:
   if not fp8:
     if ASM_GEMM:
       from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
@@ -56,13 +56,14 @@ 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, (amax_x.detach() if amax_x is not None else None), x_q
+    return out, x_q
   if x_fp8 is None:
-    if FUSED_INPUT_QUANTIZE and amax_x is not None:
+    if FUSED_INPUT_QUANTIZE:
       from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed
-      x_fp8, _, x_new_amax, _ = quantize_fp8_delayed(x, amax_x, FP8_DTYPE)
+      x_fp8, _ = quantize_fp8_delayed(x, amax_x, next_amax_x, FP8_DTYPE)
     else:
-      x_fp8, _, x_new_amax = quantize_fp8(x, amax_state=amax_x)
+      x_fp8, _, new_amax_x = quantize_fp8(x, amax_state=amax_x)
+      next_amax_x.assign(new_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):
@@ -73,51 +74,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_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
+      return out, x_fp8
+  return (x_fp8.dot(w.T, dtype=dtypes.float) * ((amax_x.float() + 1e-8) / FP8_MAX) * w_inv_scale).cast(dtypes.bfloat16), x_fp8
 
 def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
-                         grad_amax_state:Tensor, next_grad_amax_state:Tensor):
+                         next_amax_x: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, 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,
+    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,
                        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_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
   return out, x_normed, rrms, ret
 
 def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
-                             grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
+                             next_amax_x:Tensor, 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, 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,
+    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,
                        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_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
   return out, h, x_normed, rrms, ret
 
 def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor,
-                             amax_x2:Tensor,
+                             amax_x2:Tensor, next_amax_x2:Tensor,
                              grad_amax_xw13:Tensor, next_grad_amax_xw13:Tensor,
                              grad_amax_xout:Tensor, next_grad_amax_xout:Tensor):
   if FUSED_SILU_W13:
     from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
-    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,
+    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,
                        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_grad_amax_state=next_grad_amax_xout, next_amax_x=next_amax_x2)
   return out, ret
 
 class FlatTransformer:
@@ -186,14 +187,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
-    amaxs, saves = [], []
+    saves = []
 
-    xqkv, x_normed, rrms, (new_amax, *s) = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
+    xqkv, x_normed, rrms, 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)
-    amaxs.append(new_amax)
+                                                                  next_grad_amax_state=next_grad_amax_xqkv, next_amax_x=next_amax_xqkv)
     saves.extend([x_normed, rrms, *s, xqkv])
     if getenv("HK_FLASH_ATTENTION"):
       from extra.thunder.amd.fa import flash_attention, fused_qkv_rope
@@ -211,64 +212,62 @@ 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, 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)
+    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)
     saves.extend([*s, out])
-    return out, amaxs, saves
+    return out, saves
 
   def feed_forward(self, x:Tensor, residual:Tensor, **kwargs):
-    amaxs, saves = [], []
+    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, 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)
+      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"])
       saves.extend([*s, x_w1])
-      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)
+      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"])
       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, 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, *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 = out.reshape(*x_w1.shape[:-1], kwargs["w2"].shape[0])
       else:
-        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)
+        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"])
       saves.extend([*s, out])
     else:
-      x_w13, h, x_normed, rrms, (new_amax, *s) = add_norm_quantize_matmul(x, residual, kwargs["ffn_norm"], kwargs["w13"], kwargs["s_13"],
+      x_w13, h, x_normed, rrms, 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, (new_amax, *s) = silu_w13_quantize_matmul(x_w13, kwargs["w2"], kwargs["s_2"], amax_x2=kwargs["amax_x2"],
+      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"],
                                                      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, amaxs, saves
+    return out, h, 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_amaxs, attn_saves = self.attention(x, freqs_cis, **attn_kwargs)
-    ffn, h, ffn_amaxs, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs)
+    attn, attn_saves = self.attention(x, freqs_cis, **attn_kwargs)
+    ffn, h, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs)
     h = h + ffn
-    amaxs = tuple(a.detach() for a in (*attn_amaxs, *ffn_amaxs))
-    if save: return (h, *amaxs, *attn_saves, *ffn_saves)
-    else: return (h, *amaxs)
+    if save: return (h, *attn_saves, *ffn_saves)
+    else: return (h,)
 
   def shard(self, device:tuple[str, ...], mp:bool=False):
     from tinygrad.nn.state import get_parameters
@@ -319,21 +318,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])
+                        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])
       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])
-      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)
+                          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)
 
     logits = matmul(self.norm(h), self.output[0], fp8=False)[0]
     return logits
@@ -416,6 +415,9 @@ 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/models/gpt_oss.py

@@ -12,7 +12,7 @@ from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
 from extra.llama_kernels.rmsnorm import rmsnorm
-from extra.gemm.cdna_asm_gemm import _mx_block_scale, quantize_mxfp8
+from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8
 
 FP8_DTYPE = dtypes.fp8e4m3
 FP8_MAX = 448.0
@@ -39,8 +39,11 @@ def quant_dequant_mx(x:Tensor) -> Tensor:
   fxn = _quant_dequant_fwd_fxn(x.as_param(0).uop, x.device)
   return Tensor(UOp.maketuple(fxn.uop).call(x.uop, grad_fxn=_quant_dequant_bwd).gettuple(0))
 
+def _mx_scale(e8:Tensor) -> Tensor:
+  return _mx_block_scale(e8) if e8.ndim == 2 else _mx_block_scale_3d(e8)
+
 def _dequant_fwd(w_q:Tensor, w_scale:Tensor) -> Tensor:
-  return w_q.cast(dtypes.bfloat16) * _mx_block_scale(w_scale)
+  return w_q.cast(dtypes.bfloat16) * _mx_scale(w_scale)
 
 @functools.cache
 def _dequant_fwd_fxn(wq_p, ws_p, device):
@@ -48,7 +51,7 @@ def _dequant_fwd_fxn(wq_p, ws_p, device):
 
 def _dequant_bwd(grad:UOp, call:UOp) -> tuple:
   w_scale = Tensor(call.src[2])
-  return ((Tensor(grad).cast(dtypes.bfloat16) * _mx_block_scale(w_scale).cast(dtypes.bfloat16)).uop, None)
+  return ((Tensor(grad).cast(dtypes.bfloat16) * _mx_scale(w_scale).cast(dtypes.bfloat16)).uop, None)
 
 def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:
   fxn = _dequant_fwd_fxn(w_q.as_param(0).uop, w_scale.as_param(1).uop, w_q.device)

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: new_w = self._zero_gather(new_w)
+    if self.zero and not (MXFP8 and t.dtype in dtypes.fp8s): 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,6 +106,7 @@ 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/gemm/cdna_asm_gemm.py

@@ -128,6 +128,11 @@ def _mx_block_scale(e8:Tensor) -> Tensor:
   rows, scale_K = e8.shape
   return (e8.cast(dtypes.float32) - 127.0).exp2().reshape(rows, scale_K, 1).expand(rows, scale_K, 32).reshape(rows, scale_K*32)
 
+def _mx_block_scale_3d(e8:Tensor) -> Tensor:
+  # batched (E, rows, scale_K) dequant scale 2^(e8-127) broadcast to (E, rows, scale_K*32)
+  E, rows, scale_K = e8.shape
+  return (e8.cast(dtypes.float32) - 127.0).exp2().reshape(E, rows, scale_K, 1).expand(E, rows, scale_K, 32).reshape(E, rows, scale_K*32)
+
 counters = {"used":0, "todos":[]}
 def todo(msg:str) -> bool: counters["todos"].append(msg); return False
 def _asm_gemm_report():
@@ -275,10 +280,7 @@ 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, _, 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)
+        g_fp8, _ = quantize_fp8_delayed(g_t, g_amax, Tensor(next_grad_amax_state, 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
+from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, dedup, pluralize, JIT_BATCH_SIZE
 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,159 +102,138 @@ 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.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.
+# 2. deps
 
 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 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
-
+def _get_call_bufs_by_lane(call:UOp, devices:tuple[str, ...]) -> list[list[Any]]:
   refs = get_call_arg_uops(call)
-  outs, _ = get_call_outs_ins(call)
-  devices, queue = call.arg.aux.device, call.arg.aux.queue
+  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))]
 
-  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 _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
 
+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.arg.aux.device[dlane], dep.arg.aux.queue) != (devices[lane], queue)]
+    dep_lanes = [(dep, dlane, lane) for dep, dlane, lane in dep_lanes if (dep[0][dlane], dep[1]) != (devices[lane], queue)]
 
-  # 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)])
+  # 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)])
 
 # *****************
-# 2.3. merge into queues
+# 3. merge into queues
 
-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 _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 merge_queues(linear:UOp) -> UOp:
   new_src:list[UOp] = []
-  opened_qs:dict[tuple[tuple[str, ...], str], list[UOp]] = {} # (devs, queue) -> list of calls, kept in submit order
+  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)
 
   for call in linear.src:
-    if not isinstance(unwrap_after(call).arg.aux, HCQInfo):
+    if not isinstance(info:=call.arg.aux, HCQInfo) or info.name == "hcq_finalizer": # non-hcq call or finalizer: close all open queues
       new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in list(opened_qs)] + [call]
       continue
 
-    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]
+    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]
     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] == queue and set(k[0]) & set(devices)]
+      closing = [k for k in opened_qs if k[1] == info.queue and set(k[0]) & set(info.device)]
       new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in closing]
       new_rec = [call]
-    opened_qs[(devices, queue)] = new_rec
+    opened_qs[(info.device, info.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
 
@@ -279,7 +258,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 runtime_reads and all(b.op is not Ops.PARAM or b.tag is not None for b in addressed_bufs)
+  return not val.variables() and not runtime_reads and all(b.op is not Ops.PARAM or b.tag is not None for b in addressed_bufs)
 
 def is_link_patch(p:UOp, jit:bool) -> bool:
   store = p.src[0] if (is_binary_patch:=p.op is Ops.END) else p
@@ -308,7 +287,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, to_tuple(b.tag)))
+  order = sorted(dedup(bare.values()), key=lambda g: ((b:=unwrap_mstack(g.buf_uop)[0]).arg.slot, repr(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}
@@ -391,19 +370,16 @@ 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:
-    # schedule
+    # prep
     linear = linear.substitute(back_map:={s.param_like(i): s for i,s in enumerate(input_uops)} if input_uops is not None else {}, walk=True)
     linear = graph_rewrite(linear, pm_insert_copy_staging + pm_flatten_linear, name="insert copy staging")
-    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)
+
+    # 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_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")
 
@@ -508,7 +484,8 @@ 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(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(uncached=True, 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/__init__.py

@@ -20,11 +20,6 @@ def local_abs_max(x:Tensor) -> Tensor:
   fxn = _local_abs_max_fxn(param.uop, x.device)
   return Tensor(fxn[0].uop.call(x.uop).gettuple(0))
 
-def scalar_amax(amax_buf:Tensor) -> Tensor:
-  if isinstance(amax_buf.device, tuple):
-    return local_abs_max(amax_buf).detach()
-  return amax_buf.max().detach()
-
 def shard_shape(shape:tuple, axis:int, ndev:int) -> list:
   s = list(shape)
   s[axis] //= ndev

extra/llama_kernels/cast_amax/__init__.py

@@ -3,7 +3,7 @@ import functools, pathlib
 from tinygrad import Tensor, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
-from extra.llama_kernels import NUM_WG, THREADS_PER_WG, compile_cpp, alloc_like, alloc_local, scalar_amax, dname_of
+from extra.llama_kernels import NUM_WG, THREADS_PER_WG, compile_cpp, alloc_like, dname_of
 
 # module-level mailbox: grad_xw13 UOp -> (grad_xw13_fp8 UOp, delayed amax UOp)
 # lets cdna_asm_gemm's bwd reuse the fp8 companion produced by the fused silu_mul bwd kernel
@@ -11,13 +11,13 @@ from extra.llama_kernels import NUM_WG, THREADS_PER_WG, compile_cpp, alloc_like,
 _grad_fp8_mailbox:dict[UOp, tuple[UOp, UOp]] = {}
 
 @functools.cache
-def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_buf:UOp, grad_amax:UOp,
+def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_next:UOp, grad_amax:UOp,
                           xw13:UOp, grad_x2:UOp, amax_state:UOp, grad_amax_state:UOp, dname:str) -> UOp:
   hidden = xw13.shape[2] // 2
   n_elems = xw13.shape[0] * xw13.shape[1] * hidden
   threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 * 3 + n_elems * 2 + NUM_WG * 4 + 4
-  sink = UOp.sink(grad_xw13_fp8.base, grad_amax_buf.base, grad_amax.base,
+  mem = n_elems * 2 * 3 + n_elems * 2 + 4 + 4
+  sink = UOp.sink(grad_xw13_fp8.base, grad_amax_next.base, grad_amax.base,
                   xw13.base, grad_x2.base, amax_state.base, grad_amax_state.base, threads, workgroups,
                   arg=KernelInfo(f"fused_silu_mul_bwd_w13_{n_elems}", estimates=Estimates(ops=10*n_elems, mem=mem)))
   src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_bwd_w13.cpp", n_elems, hidden)
@@ -25,14 +25,14 @@ def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_buf:UOp, grad_amax:UOp,
                                UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
 
 @functools.cache
-def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_buf:UOp, xw13:UOp, amax_state:UOp, grad_amax_state:UOp,
+def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_out:UOp, xw13:UOp, amax_state:UOp, grad_amax_state:UOp,
                                 next_grad_amax_state:UOp, dname:str) -> UOp:
   # NOTE: grad_amax_state is plumbed through as an unused fwd input so the bwd kernel can read it via kernel.src
   hidden = xw13.shape[2] // 2
   n_elems = xw13.shape[0] * xw13.shape[1] * hidden
   threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 * 2 + n_elems + NUM_WG * 4
-  sink = UOp.sink(fp8_out.base, amax_buf.base, xw13.base, amax_state.base, threads, workgroups,
+  mem = n_elems * 2 * 2 + n_elems + 4
+  sink = UOp.sink(fp8_out.base, amax_out.base, xw13.base, amax_state.base, threads, workgroups,
                   arg=KernelInfo(f"fused_silu_mul_cast_amax_w13_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=mem)))
   src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_fwd_w13.cpp", n_elems, hidden)
   return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
@@ -43,26 +43,23 @@ 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_buf = alloc_local((NUM_WG,), dtypes.float32,  device, axis)
+  grad_amax_next = Tensor(next_grad_amax_state, device=device)
   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()
-  grad_xw13_fp8, grad_amax_buf, grad_amax, *_ = Tensor.custom_kernel(
-    grad_xw13_fp8, grad_amax_buf, grad_amax,
+  grad_xw13_fp8, grad_amax_next, grad_amax, *_ = Tensor.custom_kernel(
+    grad_xw13_fp8, grad_amax_next, grad_amax,
     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)
-  new_grad_amax = scalar_amax(grad_amax_buf)
-  store_effect = next_grad_amax_state.store(new_grad_amax.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) -> tuple[Tensor, Tensor]:
-  # NOTE: silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns (fp8, new_amax)
+                           next_grad_amax_state:Tensor, amax_out:Tensor) -> Tensor:
+  # NOTE: silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns fp8.
   # 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
@@ -70,8 +67,7 @@ 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_buf = alloc_local((NUM_WG,),         dtypes.float32, xw13.device, axis)
   fxn = functools.partial(_custom_fused_cast_amax_w13, dname=dname_of(xw13.device))
-  fp8_out, amax_buf, *_ = Tensor.custom_kernel(fp8_out, amax_buf, xw13, amax_state, grad_amax_state, next_grad_amax_state,
+  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, scalar_amax(amax_buf)
+  return fp8_out

extra/llama_kernels/cast_amax/cast_amax_bwd_w13.cpp

@@ -23,14 +23,14 @@ static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC");
 
 // fused silu*mul backward, three outputs in a single HBM pass:
 //   1) fp8  grad_xw13_fp8  — delayed-scale quantize using grad_amax_state (mailbox to matmul bwd)
-//   2) fp32 grad_amax_buf  — per-WG partial |grad_xw13|, reduced into next step's grad_amax_state
+//   2) fp32 grad_amax_next — scalar |grad_xw13| via global atomic max
 //   3) fp32 grad_amax_out  — delayed grad amax used for quantize/GEMM epilogue scale
 // grad_amax_state is read for the fp8 scale. The store of new_grad_amax into grad_amax_state's
 // buffer is built in Python as a separate effect and threaded into grad_a via .after(store).
 extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
 fused_silu_mul_bwd_w13(
     __hip_fp8_storage_t*  __restrict__ grad_xw13_fp8_out,    // fp8,  2*N_ELEMS
-    float*                __restrict__ grad_amax_buf,        // fp32, NUM_WG per-WG partials
+    float*                __restrict__ grad_amax_next,       // fp32 scalar, initialized to 0 before launch
     float*                __restrict__ grad_amax_out,        // fp32 scalar delayed grad amax
     const __hip_bfloat16* __restrict__ xw13,                 // bf16, 2*N_ELEMS
     const __hip_bfloat16* __restrict__ grad_x2,              // bf16, N_ELEMS
@@ -92,5 +92,6 @@ fused_silu_mul_bwd_w13(
     if (tid < s) sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]);
     __syncthreads();
   }
-  if (tid == 0) grad_amax_buf[wg] = sdata[0];
+  if (tid == 0 && sdata[0] > *grad_amax_next)
+    atomicMax(reinterpret_cast<int32_t*>(grad_amax_next), __float_as_int(sdata[0]));
 }

extra/llama_kernels/cast_amax/cast_amax_fwd_w13.cpp

@@ -24,7 +24,7 @@ static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC (so VEC loads
 extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
 fused_silu_mul_cast_amax_w13(
     __hip_fp8_storage_t*  __restrict__ fp8_out,         // fp8, N_ELEMS
-    float*                __restrict__ amax_buf,        // fp32, NUM_WG (per-WG amaxes)
+    float*                __restrict__ amax_out,        // fp32 scalar, initialized to 0 before launch
     const __hip_bfloat16* __restrict__ xw13,            // bf16, 2*N_ELEMS
     const float*          __restrict__ amax_state)      // fp32 scalar
 {
@@ -67,7 +67,7 @@ fused_silu_mul_cast_amax_w13(
     *reinterpret_cast<uint64_t*>(&fp8_out[base]) = *reinterpret_cast<uint64_t*>(out);
   }
 
-  // LDS tree reduction: per-workgroup amax
+  // LDS tree reduction: per-workgroup amax, then global atomic into the scalar.
   sdata[tid] = local_max;
   __syncthreads();
   for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
@@ -75,5 +75,5 @@ fused_silu_mul_cast_amax_w13(
     __syncthreads();
   }
 
-  if (tid == 0) amax_buf[wg] = sdata[0];
+  if (tid == 0 && sdata[0] > *amax_out) atomicMax(reinterpret_cast<int32_t*>(amax_out), __float_as_int(sdata[0]));
 }

extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/__init__.py

@@ -3,19 +3,19 @@ import functools, pathlib
 from tinygrad import Tensor, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
-from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, scalar_amax, dname_of, compile_hip
+from extra.llama_kernels import NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, dname_of, compile_hip
 
 def _src() -> str: return (pathlib.Path(__file__).parent/"fused_rmsnorm_mul_quantize_fp8.cpp").read_text()
 def _src_bwd() -> str: return (pathlib.Path(__file__).parent/"fused_rmsnorm_mul_quantize_fp8_bwd.cpp").read_text()
 
 @functools.cache
-def _custom_fwd(fp8_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_buf:UOp,
+def _custom_fwd(fp8_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_out:UOp,
                 x:UOp, weight:UOp, amax_state:UOp, dname:str, eps_val:float) -> UOp:
   MBS, SEQ, HIDDEN = x.shape
   n_elems = MBS * SEQ * HIDDEN
   threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 + n_elems + MBS * SEQ * 4 + n_elems + HIDDEN * 2 + NUM_WG * 4 + 4
-  sink = UOp.sink(fp8_out.base, x_normed_out.base, rrms_out.base, amax_buf.base,
+  mem = n_elems * 2 + n_elems + MBS * SEQ * 4 + n_elems + HIDDEN * 2 + 4 + 4
+  sink = UOp.sink(fp8_out.base, x_normed_out.base, rrms_out.base, amax_out.base,
                   x.base, weight.base, amax_state.base, threads, workgroups,
                   arg=KernelInfo(f"fused_rmsnorm_mul_quantize_fp8_{n_elems}_h{HIDDEN}_eps{eps_val:.0e}",
                                  estimates=Estimates(ops=6*n_elems, mem=mem)))
@@ -26,13 +26,13 @@ def _custom_fwd(fp8_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_buf:UOp,
                                UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
 
 @functools.cache
-def _custom_fwd_add(fp8_out:UOp, h_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_buf:UOp,
+def _custom_fwd_add(fp8_out:UOp, h_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_out:UOp,
                     x:UOp, residual:UOp, weight:UOp, amax_state:UOp, dname:str, eps_val:float) -> UOp:
   MBS, SEQ, HIDDEN = x.shape
   n_elems = MBS * SEQ * HIDDEN
   threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 * 4 + MBS * SEQ * 4 + HIDDEN * 2 + NUM_WG * 4 + 4
-  sink = UOp.sink(fp8_out.base, h_out.base, x_normed_out.base, rrms_out.base, amax_buf.base,
+  mem = n_elems * 2 * 4 + MBS * SEQ * 4 + HIDDEN * 2 + 4 + 4
+  sink = UOp.sink(fp8_out.base, h_out.base, x_normed_out.base, rrms_out.base, amax_out.base,
                   x.base, residual.base, weight.base, amax_state.base, threads, workgroups,
                   arg=KernelInfo(f"fused_add_rmsnorm_mul_quantize_fp8_{n_elems}_h{HIDDEN}_eps{eps_val:.0e}",
                                  estimates=Estimates(ops=7*n_elems, mem=mem)))
@@ -85,7 +85,7 @@ def _bwd_common(fp8_grad_u, h_grad_u, x_u, x_normed_u, rrms_u, weight_u, amax_st
   return grad_total.uop, grad_weight_uop
 
 def _fused_bwd(gradient:UOp, kernel:UOp):
-  # NOTE: fwd inputs (fp8_out, x_normed_out, rrms_out, amax_buf, x, weight, amax_state)
+  # NOTE: fwd inputs (fp8_out, x_normed_out, rrms_out, amax_out, x, weight, amax_state)
   _, x_normed_u, rrms_u, _, x_u, weight_u, amax_state_u = kernel.src[1:]
   grad_x, grad_w = _bwd_common(gradient, None, x_u, x_normed_u, rrms_u, weight_u, amax_state_u, kernel)
   return (None, None, None, None, grad_x, grad_w, None)
@@ -112,8 +112,9 @@ 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) -> tuple[Tensor, Tensor, Tensor, Tensor]:
-  # NOTE: rmsnorm(x) * weight -> fp8 + amax. Returns (fp8, new_amax, x_normed, rrms).
+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).
   # 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}"
@@ -123,16 +124,15 @@ 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_buf     = alloc_local((NUM_WG,),         dtypes.float32,  x.device, axis)
   fxn = functools.partial(_custom_fwd, dname=dname_of(x.device), eps_val=eps)
-  fp8_out, x_normed_out, rrms_out, amax_buf, *_ = Tensor.custom_kernel(
-    fp8_out, x_normed_out, rrms_out, amax_buf, x, weight, amax_state, fxn=fxn, grad_fxn=_fused_bwd)
-  return fp8_out, scalar_amax(amax_buf), x_normed_out, rrms_out
+  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
 
 def fused_add_rmsnorm_mul_quantize_fp8(x:Tensor, residual:Tensor, weight:Tensor, amax_state:Tensor,
-                                       eps:float, fp8_dtype) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
+                                       eps:float, fp8_dtype, amax_out:Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]:
   # NOTE: h = x + residual; y_normed = rmsnorm(h); fp8 = quantize(y_normed * weight).
-  # Returns (fp8, new_amax, h, x_normed, rrms). h is also written so downstream can
+  # Returns (fp8, 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,9 +143,8 @@ 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_buf     = alloc_local((NUM_WG,),         dtypes.float32,  x.device, axis)
   fxn = functools.partial(_custom_fwd_add, dname=dname_of(x.device), eps_val=eps)
-  fp8_out, h_out, x_normed_out, rrms_out, amax_buf, *_ = Tensor.custom_kernel(
-    fp8_out, h_out, x_normed_out, rrms_out, amax_buf, x, residual, weight, amax_state,
+  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, scalar_amax(amax_buf), h_out, x_normed_out, rrms_out
+  return fp8_out, h_out, x_normed_out, rrms_out

extra/llama_kernels/fused_rmsnorm_mul_quantize_fp8/fused_rmsnorm_mul_quantize_fp8.cpp

@@ -7,7 +7,7 @@
 //   fp8 = fp8_sat(y * (FP8_MAX / amax_state))
 // Also writes:
 //   rrms[row]        — saved for the rmsnorm backward
-//   amax_buf[wg]     — per-WG |y| partials, reduced later to update amax_state
+//   amax_out         — scalar |y| via global atomic max
 //
 // Layout: one WG per row, ROWS_PER_WG rows per WG via grid-stride (ROWS = N_ELEMS / HIDDEN).
 // Each thread handles HIDDEN / THREADS_PER_WG elements per row.
@@ -48,7 +48,7 @@ fused_add_rmsnorm_mul_quantize_fp8(
     __hip_bfloat16*       __restrict__ h_out,           // bf16, ROWS*HIDDEN — x + residual (saved for downstream)
     __hip_bfloat16*       __restrict__ x_normed_out,    // bf16, ROWS*HIDDEN
     float*                __restrict__ rrms_out,        // fp32, ROWS
-    float*                __restrict__ amax_buf,        // fp32, NUM_WG
+    float*                __restrict__ amax_out,        // fp32 scalar, initialized to 0 before launch
     const __hip_bfloat16* __restrict__ x,               // bf16, ROWS*HIDDEN
     const __hip_bfloat16* __restrict__ residual,        // bf16, ROWS*HIDDEN — added into x before rmsnorm
     const __hip_bfloat16* __restrict__ weight,          // bf16, HIDDEN
@@ -60,7 +60,7 @@ fused_rmsnorm_mul_quantize_fp8(
     __hip_fp8_storage_t*  __restrict__ fp8_out,         // fp8, ROWS*HIDDEN
     __hip_bfloat16*       __restrict__ x_normed_out,    // bf16, ROWS*HIDDEN (saved for rmsnorm bwd)
     float*                __restrict__ rrms_out,        // fp32, ROWS (fp32 to match rmsnorm_bwd.cpp expectation)
-    float*                __restrict__ amax_buf,        // fp32, NUM_WG per-WG partials
+    float*                __restrict__ amax_out,        // fp32 scalar, initialized to 0 before launch
     const __hip_bfloat16* __restrict__ x,               // bf16, ROWS*HIDDEN
     const __hip_bfloat16* __restrict__ weight,          // bf16, HIDDEN (per-hidden scale)
     const float*          __restrict__ amax_state)      // fp32 scalar
@@ -144,12 +144,12 @@ fused_rmsnorm_mul_quantize_fp8(
     __syncthreads();  // before next row's sum_sq reduce reuses sdata
   }
 
-  // Final per-WG amax reduce.
+  // Final per-WG amax reduce, then global atomic into the scalar.
   sdata[tid] = local_max;
   __syncthreads();
   for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
     if (tid < s) sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]);
     __syncthreads();
   }
-  if (tid == 0) amax_buf[wg] = sdata[0];
+  if (tid == 0 && sdata[0] > *amax_out) atomicMax(reinterpret_cast<int32_t*>(amax_out), __float_as_int(sdata[0]));
 }

extra/llama_kernels/quantize_fp8_delayed/__init__.py

@@ -3,14 +3,13 @@ from tinygrad import Tensor, dtypes
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import prod
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
-from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, scalar_amax
+from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like
 
 @functools.cache
-def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_partial:UOp, x:UOp, amax_state:UOp) -> UOp:
+def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:UOp, device=None) -> UOp:
   VEC = 8
   n_elems = prod(x.shape)
   assert n_elems % (NUM_WG * THREADS_PER_WG * VEC) == 0
-  assert amax_partial.shape[0] == NUM_WG
 
   x = x.reshape(n_elems)
   fp8_out = fp8_out.reshape(n_elems)
@@ -46,8 +45,13 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_partial:UOp, x:UOp, amax_st
     lds = lds.after(lds[tid.valid(active)].store(lds[tid].maximum(other)).barrier())
     step //= 2
 
-  amax_store = amax_partial[tid.eq(0).where(wg, UOp.invalid())].store(lds[0])
-  return amax_store.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=()))
+  device = device[0].split(":")[0] if isinstance(device, tuple) else device.split(":")[0]
+  if device in {"AMD", "NULL"}: atomic_arg = "if ({2} > {3}) __hip_atomic_fetch_max((int*){0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);"
+  else: raise NotImplementedError(f"no atomic max for device {device}")
+  amax_idx = amax_out.reshape((1,)).index(UOp.const(dtypes.index, 0))
+  max_val = lds[0].load()
+  atomic = UOp(Ops.CUSTOM, dtypes.void, (amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=atomic_arg)
+  return atomic.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=()))
 
 @functools.cache
 def _custom_quantize_fp8_scalar(fp8_out:UOp, x:UOp, amax_state:UOp) -> UOp:
@@ -69,25 +73,19 @@ 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, 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 + per-WG |x| partials (then a small reduce produces scalar new_amax).
-  # 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.
+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.
+  # Fused kernel reads x once and writes fp8 + scalar amax via global atomic max.
   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_partial = alloc_local((NUM_WG,), dtypes.float32, x.device, axis)
-  fxn = _custom_quantize_fp8_with_amax
-  fp8_out, amax_partial, *_ = Tensor.custom_kernel(fp8_out, amax_partial, x, amax_state,
-                                                    fxn=fxn, grad_fxn=_quantize_fp8_delayed_bwd)
-  new_amax = scalar_amax(amax_partial)
+  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
-  store_effect = amax_state.uop.store(new_amax.uop)
-  return fp8_out, inv_scale, new_amax, store_effect
+  return fp8_out, inv_scale
 
 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 (B_local, N, H_local, H_KV_local, head_dim) == (2, 8192, 32, 8, 128)
+  assert H_local % H_KV_local == 0 and head_dim % 2 == 0 and head_dim <= 512
   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,8 +565,10 @@ 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,7 +161,6 @@ norecursedirs = [
   ".git",
 ]
 timeout = 300
-timeout_method = "thread"
 timeout_func_only = true
 testpaths = ["test"]
 filterwarnings = [

test/backend/test_dtype_alu.py

@@ -291,6 +291,33 @@ 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,9 +49,10 @@ def run_quantize_fp8(shape:tuple[int, ...], delayed:bool=True) -> None:
   with Context(DEBUG=0): Tensor.realize(x, amax_state)
 
   if delayed:
-    fp8, inv_scale, new_amax, _ = quantize_fp8_delayed(x, amax_state, FP8_DTYPE)
+    amax_out = Tensor.zeros((), dtype=dtypes.float32, device=x.device).realize()
+    fp8, inv_scale = quantize_fp8_delayed(x, amax_state, amax_out, FP8_DTYPE)
     ref_fp8, ref_inv_scale, ref_new_amax = quantize_fp8(x, amax_state=amax_state)
-    Tensor.realize(fp8, inv_scale, new_amax)
+    Tensor.realize(fp8, inv_scale)
     Tensor.realize(ref_fp8, ref_inv_scale, ref_new_amax)
   else:
     fp8 = quantize_fp8_scalar(x, amax_state, FP8_DTYPE)
@@ -63,9 +64,10 @@ 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 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())}"
+      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())}"
 
+@unittest.skipUnless(Device.DEFAULT == "AMD", "requires atomic max")
 class TestQuantizeFP8(unittest.TestCase):
   def setUp(self):
     ren = Device[Device.DEFAULT].renderer
@@ -81,10 +83,11 @@ 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()
-    fp8, _, new_amax, _ = quantize_fp8_delayed(x, amax_state, FP8_DTYPE)
-    Tensor.realize(fp8, new_amax)
+    amax_out = Tensor.zeros((), dtype=dtypes.float32, device=devs).realize()
+    fp8, _ = quantize_fp8_delayed(x, amax_state, amax_out, FP8_DTYPE)
+    Tensor.realize(fp8)
     assert fp8.uop.shape == x.uop.shape
-    assert new_amax.shape == ()
+    assert amax_out.shape == ()
 
 class TestLocalAmax(unittest.TestCase):
   def test_multi_tensor_local_shard_amax(self):

test/backend/test_ops.py

@@ -848,6 +848,8 @@ 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)
 
@@ -859,6 +861,8 @@ 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)
 
@@ -870,6 +874,7 @@ 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]]
@@ -879,6 +884,7 @@ 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_usb_asm24.py

@@ -1,31 +1,31 @@
 import unittest
-from tinygrad.helpers import Timing
+from tinygrad.helpers import Timing, getenv
 from tinygrad import Tensor, Device
 import numpy as np
 
 class TestDevCopySpeeds(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
-    cls.sz = 768
+    cls.sz = getenv("SIZE", 2e6)
     cls.dev = Device["AMD"]
     if not cls.dev.is_usb(): raise unittest.SkipTest("only test this on USB devices")
 
   def testCopyCPUtoDefault(self):
     for _ in range(10):
-      t = Tensor.ones(self.sz, self.sz, device="CPU").contiguous().realize()
+      t = Tensor.ones(self.sz, device="CPU", dtype='uchar').contiguous().realize()
       with Timing(f"copyin of {t.nbytes()/1e6:.2f} MB:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns * 1e3:.2f} MB/s"): # noqa: F821
         t.to(Device.DEFAULT).realize()
         Device[Device.DEFAULT].synchronize()
       del t
 
   def testCopyDefaulttoCPU(self):
-    t = Tensor.ones(self.sz, self.sz).contiguous().realize()
+    t = Tensor.ones(self.sz, dtype='uchar').contiguous().realize()
     for _ in range(10):
       with Timing(f"copyout of {t.nbytes()/1e6:.2f} MB:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns * 1e3:.2f} MB/s"):
         t.to('CPU').realize()
 
   def testValidateCopies(self):
-    t = Tensor.randn(self.sz, self.sz, device="CPU").contiguous().realize()
+    t = Tensor.randn(self.sz, device="CPU", dtype='uchar').contiguous().realize()
     x = t.to(Device.DEFAULT).realize()
     Device[Device.DEFAULT].synchronize()
 

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 = 2 # HCQ2 merges calls on the same queue
+  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 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
+import ctypes, time, os, builtins, fcntl, typing
 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 = {}
+tracked_fds: dict[int, typing.Any] = {}
 
 original_memoryview = builtins.memoryview
 class TrackedMemoryView:

test/null/test_helpers.py

@@ -1,6 +1,7 @@
 import ctypes, gzip, unittest, timeit, pickle
 from tinygrad import Variable
-from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, mv_address, count, all_same
+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 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

@@ -1,4 +1,4 @@
-import unittest, threading, time
+import unittest, threading, time, json
 from unittest.mock import Mock
 
 class TestLLMServer(unittest.TestCase):
@@ -7,12 +7,9 @@ class TestLLMServer(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
     cls.mock_tok = Mock()
-    cls.mock_tok.role = Mock(return_value=[100, 101])
     cls.mock_tok.encode = Mock(return_value=[200, 201, 202])
     cls.mock_tok.decode = Mock(return_value="Hello")
     cls.mock_tok.stream_decoder = Mock(return_value=lambda tid=None: "Hello" if tid is not None else "")
-    cls.mock_tok.end_turn = Mock(return_value=[998])
-    cls.mock_tok.prefix = Mock(return_value=[1])
     cls.mock_tok.preset = "llama3"
     cls.mock_tok.bos_id = 1
     cls.mock_tok.eos_id = 999
@@ -20,12 +17,14 @@ class TestLLMServer(unittest.TestCase):
     cls.mock_tok.is_end = Mock(side_effect=lambda tid: tid in (999,))
 
     cls.mock_model = Mock()
+    cls.mock_model.max_context = 4
     cls.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 999]))
     cls.mock_model.get_start_pos = Mock(return_value=0)
 
-    from tinygrad.llm.cli import LLMServer
+    from tinygrad.llm.cli import FallbackTemplate
+    from tinygrad.llm.serve import LLMServer
 
-    cls.server = LLMServer(('127.0.0.1', 0), cls.mock_model, "test-model", cls.mock_tok)
+    cls.server = LLMServer(('127.0.0.1', 0), cls.mock_model, "test-model", cls.mock_tok, FallbackTemplate(cls.mock_tok))
     cls.port = cls.server.server_address[1]
     cls.server_thread = threading.Thread(target=cls.server.serve_forever, daemon=True)
     cls.server_thread.start()
@@ -131,6 +130,16 @@ class TestLLMServer(unittest.TestCase):
     self.assertIsNotNone(resp.usage.prompt_tokens)
     self.assertIsNotNone(resp.usage.completion_tokens)
 
+  def test_context_length_error(self):
+    from openai import BadRequestError
+    self.mock_tok.encode.return_value = [200, 201, 202, 203]
+    try:
+      with self.assertRaises(BadRequestError) as err:
+        self.client.chat.completions.create(model="test-model", messages=[{"role":"user", "content":"too long"}])
+      self.assertEqual(err.exception.code, "context_length_exceeded")
+    finally:
+      self.mock_tok.encode.return_value = [200, 201, 202]
+
   def test_max_tokens_streaming(self):
     self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 302, 303, 999]))
     stream = self.client.chat.completions.create(
@@ -149,50 +158,6 @@ class TestLLMServer(unittest.TestCase):
     self.assertEqual(resp.choices[0].finish_reason, "length")
     self.assertEqual(resp.usage.completion_tokens, 2)
 
-  def test_assistant_prefill(self):
-    """Last assistant message should be treated as prefill (not a completed turn)."""
-    self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 999]))
-    captured_ids = []
-    orig_generate = self.mock_model.generate.side_effect
-    def capture_generate(ids, **kwargs):
-      captured_ids.extend(ids)
-      return orig_generate(ids, **kwargs)
-    self.mock_model.generate = Mock(side_effect=capture_generate)
-
-    resp = self.client.chat.completions.create(
-      model="test", messages=[
-        {"role": "user", "content": "Hello"},
-        {"role": "assistant", "content": "Sure"}
-      ], stream=False
-    )
-    # prefill tokens should be in ids: role("assistant") + encode("Sure") but NO end_turn after it
-    # and NO extra role("assistant") appended
-    role_tokens = self.mock_tok.role.call_args_list
-    # last role() call should be for "assistant" (the prefill message), not an extra one
-    self.assertEqual(role_tokens[-1], unittest.mock.call("assistant"))
-    # end_turn should be called once less than role() — the prefill assistant msg doesn't get end_turn
-    # NOTE: this is flaky in random order
-    #self.assertEqual(self.mock_tok.end_turn.call_count, self.mock_tok.role.call_count - 1)
-    self.assertIsNotNone(resp.choices[0].message.content)
-
-  def test_assistant_prefill_not_last(self):
-    """Assistant message that's NOT last should be a normal completed turn."""
-    self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 999]))
-    self.mock_tok.role.reset_mock()
-    self.mock_tok.end_turn.reset_mock()
-    self.client.chat.completions.create(
-      model="test", messages=[
-        {"role": "user", "content": "Hello"},
-        {"role": "assistant", "content": "Sure"},
-        {"role": "user", "content": "Continue"}
-      ], stream=False
-    )
-    # all messages get end_turn, plus an extra role("assistant") at the end
-    # roles: user, assistant, user, assistant(generation prompt) = 4 role calls
-    # end_turns: user, assistant, user = 3 end_turn calls (one per message)
-    self.assertEqual(self.mock_tok.end_turn.call_count, 3)
-    self.assertEqual(self.mock_tok.role.call_count, 4)
-
   def test_models_endpoint(self):
     import requests as req
     resp = req.get(f"http://127.0.0.1:{self.port}/v1/models")
@@ -203,5 +168,103 @@ class TestLLMServer(unittest.TestCase):
     self.assertEqual(data["data"][0]["id"], "test-model")
     self.assertEqual(data["data"][0]["object"], "model")
 
+class TestLLMToolCalls(unittest.TestCase):
+  """Tool calling through the OpenAI-compatible HTTP API."""
+
+  @classmethod
+  def setUpClass(cls):
+    cls.mock_tok = Mock()
+    cls.mock_tok.encode = Mock(return_value=[200, 201, 202])
+    cls.mock_tok.decode = Mock(return_value="")
+    cls.mock_tok.preset = "qwen2"
+    cls.mock_tok.bos_id, cls.mock_tok.eos_id, cls.mock_tok.eot_id = None, 999, None
+    cls.mock_tok.is_end = Mock(return_value=False)
+
+    cls.mock_model = Mock()
+    cls.mock_model.max_context = 4
+    cls.mock_model.get_start_pos = Mock(return_value=0)
+
+    from tinygrad.llm.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 %}""")
+    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)
+    cls.server_thread.start()
+    time.sleep(0.1)
+
+    from openai import OpenAI
+    cls.client = OpenAI(base_url=f"http://127.0.0.1:{cls.port}/v1", api_key="test")
+
+  @classmethod
+  def tearDownClass(cls):
+    cls.server.shutdown()
+    cls.server.server_close()
+
+  def set_output(self, text:str):
+    pieces = dict(enumerate(text, 1))
+    self.mock_tok.stream_decoder = Mock(return_value=lambda tid=None: pieces[tid] if tid is not None else "")
+    self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter(pieces))
+
+  @staticmethod
+  def tools():
+    return [{"type":"function", "function":{"name":"read", "description":"Read a file",
+      "parameters":{"type":"object", "properties":{"path":{"type":"string"}}, "required":["path"]}}}]
+
+  def test_streaming_tool_call(self):
+    self.set_output('before<tool_call>{"name":"read","arguments":{"path":"README.md"}}</tool_call>')
+    chunks = list(self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Read README.md"}],
+                                                     tools=self.tools(), stream=True))
+    self.assertEqual("".join(c.choices[0].delta.content or "" for c in chunks if c.choices), "before")
+    calls = [tc for c in chunks if c.choices for tc in c.choices[0].delta.tool_calls or []]
+    self.assertEqual(len(calls), 1)
+    self.assertEqual(calls[0].function.name, "read")
+    self.assertEqual(json.loads(calls[0].function.arguments), {"path":"README.md"})
+    self.assertEqual(chunks[-1].choices[0].finish_reason, "tool_calls")
+
+  def test_multiple_xml_tool_calls(self):
+    self.set_output("<tool_call><function=read><parameter=path>\"a\"</parameter></function></tool_call>"
+                    "<tool_call><function=read><parameter=path>\"b\"</parameter></function></tool_call>")
+    response = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Read a and b"}],
+                                                   tools=self.tools())
+    self.assertEqual([json.loads(tc.function.arguments)["path"] for tc in response.choices[0].message.tool_calls], ["a", "b"])
+    self.assertEqual(response.choices[0].finish_reason, "tool_calls")
+
+  def test_multiline_tool_argument_preserves_trailing_newline(self):
+    self.set_output("<tool_call>\n<function=write>\n<parameter=content>\nfirst\nsecond\n\n</parameter>\n"
+                    "<parameter=filePath>\nout.txt\n</parameter>\n</function>\n</tool_call>")
+    response = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Write out.txt"}], tools=self.tools())
+    args = json.loads(response.choices[0].message.tool_calls[0].function.arguments)
+    self.assertEqual(args, {"content":"first\nsecond\n", "filePath":"out.txt"})
+
+  def test_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())
+    self.assertEqual(response.choices[0].message.content, "<tool_call>not a call</tool_call>")
+    self.assertIsNone(response.choices[0].message.tool_calls)
+    self.assertEqual(response.choices[0].finish_reason, "stop")
+
+  def test_tool_call_in_reasoning_is_not_executed(self):
+    self.set_output('<think>draft <tool_call>{"name":"wrong","arguments":{}}</tool_call></think>answer')
+    response = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Hello"}], tools=self.tools())
+    self.assertEqual(response.choices[0].message.content, "answer")
+    self.assertIsNone(response.choices[0].message.tool_calls)
+    self.assertEqual(response.choices[0].finish_reason, "stop")
+
+  def test_tool_result_round_trip(self):
+    self.set_output('<tool_call>{"name":"read","arguments":{"path":"README.md"}}</tool_call>')
+    first = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Read README.md"}], tools=self.tools())
+    call = first.choices[0].message.tool_calls[0]
+    self.set_output("done")
+    second = self.client.chat.completions.create(model="tool-model", messages=[
+      {"role":"user", "content":"Read README.md"},
+      {"role":"assistant", "content":None, "tool_calls":[call.model_dump()]},
+      {"role":"tool", "tool_call_id":call.id, "content":"file contents"},
+    ], tools=self.tools())
+    self.assertEqual(second.choices[0].message.content, "done")
+    self.assertEqual(second.choices[0].finish_reason, "stop")
+
 if __name__ == '__main__':
   unittest.main()

test/null/test_llm_tokenizer.py

@@ -1,5 +1,5 @@
-import unittest, base64, functools, sys
-from tinygrad.llm.cli import SimpleTokenizer
+import unittest, base64, functools, re, sys, time, unicodedata
+from tinygrad.llm.cli import SimpleTokenizer, FallbackTemplate
 from tinygrad.helpers import fetch
 
 @unittest.skipIf(sys.platform == 'win32', "fetch race condition on Windows")
@@ -46,6 +46,41 @@ 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])
+    self._test_coding(self.llama_tok, "hello changed <|eot_id|>world again", [15339, 5614, 220, 128009, 14957, 1578])
+
+  def test_long_cached_prompt_matches_fresh_tokenization(self):
+    prefix = "system tools\n" * 700 + "<|eot_id|>"
+    first, changed = prefix + "run tower of hanoi", prefix + "run ls /"
+    expected = self.llama_tok.encode(changed)
+    self.llama_tok.encode(first)
+    self.assertEqual(self.llama_tok.encode(changed), expected)
+
   def test_tekken_from_gguf_kv(self):
     kv = {
       "tokenizer.ggml.tokens": ["<unk>", "<s>", "</s>", "[INST]", "[/INST]", "hello"],
@@ -54,10 +89,11 @@ class TestLLMTokenizer(unittest.TestCase):
       "tokenizer.ggml.eos_token_id": 2,
     }
     tok = SimpleTokenizer.from_gguf_kv(kv)
-    self.assertEqual(tok.role("user"), [3])
+    template = FallbackTemplate(tok)
+    self.assertEqual(template.role("user"), "[INST]")
     self.assertEqual(tok.encode("hello"), [5])
-    self.assertEqual(tok.end_turn(), [4])
-    self.assertEqual(tok.role("assistant"), [])
+    self.assertEqual(template.end_turn(), "[/INST]")
+    self.assertEqual(template.role("assistant"), "")
 
   def test_stream_decoder(self):
     """stream_decoder buffers incomplete UTF-8: token 25677 has 3/4 of emoji, token 138 completes it."""

test/null/test_schedule.py

@@ -1472,6 +1472,18 @@ 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,6 +127,13 @@ 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
@@ -364,6 +371,10 @@ 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
@@ -402,5 +413,15 @@ 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, dtypes
+from tinygrad import Device, Tensor, Variable, dtypes
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.codegen import to_program
 from tinygrad.codegen.opt import Opt, OptOps
@@ -114,5 +114,31 @@ 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,6 +18,9 @@ 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
@@ -25,7 +28,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 = Tensor.rand(M, K, dtype=dtype_in), Tensor.rand(K, N, dtype=dtype_in)
+  a, b = _tc_rand(M, K, dtype=dtype_in), _tc_rand(K, N, dtype=dtype_in)
   r = a.matmul(b, dtype=dtype_out)
   sched = r.schedule_linear()
   realized_ast = sched.src[-1].src[0]
@@ -44,7 +47,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 = Tensor.rand(M, K, dtype=dtype_in), Tensor.rand(K, N, dtype=dtype_in)
+  a, b = _tc_rand(M, K, dtype=dtype_in), _tc_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()
@@ -57,9 +60,11 @@ 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))
-  np.testing.assert_allclose(c, np_a @ np_b, atol=tc_atol, rtol=tc_rtol)
+  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)
 
 class TestTensorCores(unittest.TestCase):
   # TODO: don't skip bf16 for real device (METAL, AMD)
@@ -75,7 +80,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 = Tensor.rand(m, k, dtype=tc.dtype_in), Tensor.rand(k, n, dtype=tc.dtype_in)
+      a, b = _tc_rand(m, k, dtype=tc.dtype_in), _tc_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_dtype_weak.py

@@ -1,10 +1,11 @@
-import unittest
+import pathlib, tempfile, 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):
@@ -29,8 +30,7 @@ class TestWeakPromotion(unittest.TestCase):
       self.assertEqual(t.dtype, weak)
       self.assertEqual(t.data().itemsize, strong.itemsize)
       self.assertEqual(t.numpy().dtype.itemsize, strong.itemsize)
-      realized = t.clone("CPU").realize()
-      self.assertEqual((realized.dtype, realized.uop.buffer.dtype), (strong, strong))
+      with self.assertRaises(RuntimeError): t.clone("CPU")
     with patch.object(dtypes, "default_int", dtypes.int64):
       self.assertEqual(Tensor.const(dtypes.weakint, 3).numpy().dtype.itemsize, dtypes.int64.itemsize)
 
@@ -110,6 +110,65 @@ 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, is_numpy_ndarray, disable_gc
-from tinygrad.tensor import Tensor
+from tinygrad.helpers import polyN, disable_gc
+from tinygrad.tensor import Tensor, is_numpy_ndarray
 
 class TestPolyN(unittest.TestCase):
   def test_tensor(self):

test/unit/test_modern_scan.py

@@ -0,0 +1,36 @@
+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,8 +180,9 @@ 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, b.src[0].expr if b.op is Ops.BIND else None,
-                   b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL)
+                   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)
 
 pm_finalize_call = PatternMatcher([
   (UPat(Ops.AFTER, name="x"), finalize_after),

tinygrad/codegen/__init__.py

@@ -175,9 +175,13 @@ 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:])),
+   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),
 ])
 
 def fix_group_for_reduce(x:UOp):

tinygrad/codegen/decomp/dtype.py

@@ -33,11 +33,14 @@ 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:
-      lo, hi = shl(a0, b0_mod:=b0 & 31), shl(a1, b0_mod) | shr(shr(a0, 1), 31 - b0_mod)
+      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)
       return (b0 >= 32).where(zero, lo), (b0 >= 32).where(lo, hi)
     case Ops.SHR:
-      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)
+      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)
     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)]]
+  for di,do in [(dtypes.half,dtypes.float),(dtypes.half,dtypes.half),(dtypes.bfloat16,dtypes.float),(dtypes.int8,dtypes.int32)]]
 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/helpers.py

@@ -81,7 +81,6 @@ 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")
@@ -413,7 +412,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()
+_db_tables: set[str] = 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,10 +1,12 @@
 from __future__ import annotations
-import sys, argparse, codecs, typing, re, unicodedata, json, uuid, time, pathlib
+import sys, argparse, codecs, itertools, 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, stderr_log, colored, Context, fetch, profile_marker, getenv
-from tinygrad.viz.serve import TCPServerWithReuse, HTTPRequestHandler
+from tinygrad.helpers import partition, DEBUG, Timing, GlobalCounters, Context, fetch, profile_marker, getenv
 from tinygrad.llm.model import Transformer
+if TYPE_CHECKING:
+  import jinja2
 
 class SimpleTokenizer:
   def __init__(self, normal_tokens:dict[str, int], special_tokens:dict[str, int], preset:str="llama3",
@@ -18,7 +20,11 @@ class SimpleTokenizer:
 
     # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286
     # 0x323b0 is one past the max codepoint in unicode categories L/N/Z (0x323af is max L)
-    def ucat_range(pre: str): return "".join(re.escape(chr(cp)) for cp in range(0x323b0) if unicodedata.category(chr(cp)).startswith(pre))
+    # compact adjacent codepoints into ranges: listing them all makes re spend seconds on large prompts
+    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)
     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}]+")
@@ -64,25 +70,6 @@ class SimpleTokenizer:
     dec = codecs.getincrementaldecoder('utf-8')('replace')
     def _decode(tid:int|None=None) -> str: return dec.decode(self._tok2bytes[tid]) if tid is not None else dec.decode(b'', final=True)
     return _decode
-  def role(self, role:str):
-    if self.preset == 'olmo': return self.encode("<|" + role + "|>\n")  # OLMoE Instruct format
-    if self.preset == 'kimi-k2': return self.encode("<|im_" + role + "|>" + role + "<|im_middle|>")
-    if self.preset == 'qwen2': return self.encode("<|im_start|>" + role + "\n")
-    if self.preset == 'glm4': return self.encode("<|" + role + "|>")
-    if self.preset == 'tekken':
-      if role == 'user': return self.encode("[INST]")
-      if role == 'assistant': return []
-      raise ValueError(f"Unsupported role '{role}' for tokenizer preset '{self.preset}'")
-    return self.encode("<|start_header_id|>" + role + "<|end_header_id|>\n\n")
-  def end_turn(self):
-    if self.preset == 'olmo': return self.encode("\n")
-    if self.preset == 'kimi-k2': return [self.eos_id]
-    if self.preset == 'qwen2': return [self.eos_id] + self.encode("\n")
-    if self.preset == 'glm4': return []
-    if self.preset == 'tekken': return self.encode("[/INST]")
-    return [self.eos_id]
-  def prefix(self) -> list[int]:
-    return ([] if self.bos_id is None else [self.bos_id]) + (self.encode("<sop>") if self.preset == 'glm4' else [])
   def is_end(self, token_id:int) -> bool: return token_id in (self.eos_id, self.eot_id)
 
 models = {
@@ -105,81 +92,41 @@ models = {
   "glm-4.7-flash": "https://huggingface.co/unsloth/GLM-4.7-Flash-GGUF/resolve/main/GLM-4.7-Flash-Q4_K_M.gguf",
 }
 
-# *** simple OpenAI API compatible server with web interface on http://localhost:8000/ ***
+class FallbackTemplate:
+  # minimal jinja2.Template-compatible chat template without jinja2, no tool calling support
+  def __init__(self, tok:SimpleTokenizer): self.tok = tok
+  def role(self, role:str) -> str:
+    if self.tok.preset == 'olmo': return "<|" + role + "|>\n"  # OLMoE Instruct format
+    if self.tok.preset == 'kimi-k2': return "<|im_" + role + "|>" + role + "<|im_middle|>"
+    if self.tok.preset == 'qwen2': return "<|im_start|>" + role + "\n"
+    if self.tok.preset == 'glm4': return "<|" + role + "|>"
+    if self.tok.preset == 'tekken':
+      if role == 'user': return "[INST]"
+      if role == 'assistant': return ""
+      raise ValueError(f"Unsupported role '{role}' for tokenizer preset '{self.tok.preset}'")
+    return "<|start_header_id|>" + role + "<|end_header_id|>\n\n"
+  def end_turn(self) -> str:
+    if self.tok.preset == 'olmo': return "\n"
+    if self.tok.preset == 'kimi-k2': return self.tok.decode([self.tok.eos_id])
+    if self.tok.preset == 'qwen2': return self.tok.decode([self.tok.eos_id]) + "\n"
+    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:
+    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"])
+      content = msg.get("content")
+      if isinstance(content, str): out += content
+      elif isinstance(content, list):
+        for c in content:
+          if c["type"] == "text": out += c["text"]
+          else: raise RuntimeError(f"unhandled type: {c['type']}")
+      elif content is not None: raise RuntimeError(f"unknown content type: {type(content)}")
+      out += self.end_turn()
+    return out + self.role("assistant") if add_generation_prompt else out
 
-class Handler(HTTPRequestHandler):
-  server: LLMServer
-  def log_request(self, code='-', size='-'): pass
-  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):
-    model, tok = self.server.model, self.server.tok
-    cache_start_pos = model.get_start_pos(ids)
-    stderr_log(f"{self.path}  {colored('--', 'BLACK')}  "
-               f"in:{colored(f'{cache_start_pos:5d}', 'green')} +{len(ids)-cache_start_pos:5d}  {colored('--', 'BLACK')}  ")
-    tmpl = {"id":f"chatcmpl-{uuid.uuid4().hex[:24]}", "object":"chat.completion.chunk", "created":int(time.time()), "model":model_name}
-    yield {"choices": [{"index":0, "delta":{"role":"assistant","content":""}, "finish_reason":None}], **tmpl}
-    out: list[int] = []
-    finish_reason = "stop"
-    st = time.perf_counter()
-    dec = tok.stream_decoder()
-    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 tok.is_end(next_id): break
-      out.append(next_id)
-      yield {"choices": [{"index":0, "delta":{"content":dec(next_id)}, "finish_reason":None}], **tmpl}
-      if max_tokens is not None and len(out) >= max_tokens:
-        finish_reason = "length"
-        break
-    if (tail := dec()): yield {"choices": [{"index":0, "delta":{"content":tail}, "finish_reason":None}], **tmpl}
-    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}
-    et = time.perf_counter()
-    stderr_log(f"gen:{len(out)/(et-pt) if len(out) > 1 else 0:4.0f} tok/s  {colored('--', 'BLACK')}  "
-               f"out:{len(out):5d}  {colored('--', 'BLACK')}  total:{et-st:6.2f}s\n")
-
-  def do_POST(self):
-    tok = self.server.tok
-    raw_body = self.rfile.read(int(self.headers.get("Content-Length", "0")))
-    body: dict[str, typing.Any] = json.loads(raw_body.decode("utf-8"))
-    if DEBUG >= 1: print(json.dumps(body, indent=2))
-    if self.path == "/v1/chat/completions":
-      # extract tokens, last assistant message is treated as prefill
-      ids: list[int] = tok.prefix()
-      for i, msg in enumerate(body["messages"]):
-        ids += tok.role(msg["role"])
-        content = msg["content"]
-        if isinstance(content, str): ids += tok.encode(content)
-        elif isinstance(content, list):
-          for c in content:
-            if c["type"] == "text": ids += tok.encode(c["text"])
-            else: raise RuntimeError(f"unhandled type: {c['type']}")
-        else: raise RuntimeError(f"unknown content type: {type(content)}")
-        if msg["role"] == "assistant" and i == len(body["messages"]) - 1: break
-        ids += tok.end_turn()
-      else: ids += tok.role("assistant")
-
-      # 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)))
-      if body.get("stream"): self.stream_json(chunks)
-      else:
-        out, finish_reason = [], "stop"
-        for c in chunks:
-          if c["choices"] and c["choices"][0].get("delta", {}).get("content"): out.append(c["choices"][0]["delta"]["content"])
-          if c["choices"] and c["choices"][0].get("finish_reason"): finish_reason = c["choices"][0]["finish_reason"]
-        self.send_data(json.dumps({**c, "object":"chat.completion",
-          "choices":[{"index":0, "message":{"role":"assistant","content":"".join(out)}, "finish_reason":finish_reason}]}).encode())
-    else:
-      raise RuntimeError(f"unhandled path {self.path}")
-
-class LLMServer(TCPServerWithReuse):
-  def __init__(self, server_address:tuple, model:Transformer, model_name:str, tok:SimpleTokenizer):
-    self.model, self.model_name, self.tok = model, model_name, tok
-    super().__init__(server_address, Handler)
+from tinygrad.llm.serve import LLMServer
 
 def main():
   parser = argparse.ArgumentParser()
@@ -194,11 +141,26 @@ def main():
   model, kv = Transformer.from_gguf(fetch(models.get(args.model, args.model)), 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")
+  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}")
 
   # get tokenizer
   tok = SimpleTokenizer.from_gguf_kv(kv)
 
+  # use the model's chat template if jinja2 is available (enables model-specific formatting)
+  template: jinja2.Template|FallbackTemplate = FallbackTemplate(tok)
+  if (ct := kv.get('tokenizer.chat_template')) is not None:
+    try:
+      import jinja2
+      env = jinja2.Environment()
+      env.filters['tojson'] = lambda obj, **kwargs: json.dumps(obj, **kwargs)  # jinja2's tojson escapes <>& for HTML safety
+      env.globals['raise_exception'] = lambda msg: (_ for _ in ()).throw(RuntimeError(msg))
+      env.globals['strftime_now'] = lambda fmt: time.strftime(fmt)
+      env.globals['bos_token'] = tok.decode([tok.bos_id]) if tok.bos_id is not None else ""
+      env.globals['eos_token'] = tok.decode([tok.eos_id])
+      template = env.from_string(ct)
+    except ImportError: print("warning: jinja2 is not installed, the model's chat template is disabled")
+
   # warmup the JIT
   if args.warmup or args.serve:
     # run 2 tokens through the model twice to capture the JIT before serving
@@ -206,7 +168,7 @@ def main():
       for _ in range(2): list(zip(range(2), model.generate([0])))
 
   # start server
-  if args.serve: LLMServer(('', args.serve), model, model_name, tok).serve_forever()
+  if args.serve: LLMServer(('', args.serve), model, model_name, tok, template).serve_forever()
 
   # do benchmark
   if args.benchmark is not None:
@@ -224,16 +186,19 @@ def main():
     exit(0)
 
   # interactive chat
-  ids: list[int] = tok.prefix()
+  messages: list[dict] = []
   while 1:
-    try:
-      ids += tok.role("user") + tok.encode(input('>>> ')) + tok.end_turn() + tok.role("assistant")
-    except EOFError:
-      break
-    dec = tok.stream_decoder()
+    try: messages.append({"role":"user", "content":input('>>> ')})
+    except EOFError: break
+    ids = tok.encode(template.render(messages=messages, add_generation_prompt=True))
+    reply, dec = "", tok.stream_decoder()
     for next_id in model.generate(ids):
-      sys.stdout.write(dec(next_id) if not tok.is_end(next_id) else dec() + "\n\n")
+      if tok.is_end(next_id):
+        sys.stdout.write(dec() + "\n\n")
+        break
+      reply += (piece := dec(next_id))
+      sys.stdout.write(piece)
       sys.stdout.flush()
-      if tok.is_end(next_id): break
+    messages.append({"role":"assistant", "content":reply})
 
 if __name__ == "__main__": main()

tinygrad/llm/serve.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+import json, pathlib, re, time, typing, uuid
+from typing import TYPE_CHECKING
+from tinygrad.helpers import DEBUG, colored, stderr_log
+from tinygrad.viz.serve import TCPServerWithReuse, HTTPRequestHandler
+if TYPE_CHECKING:
+  from tinygrad.llm.cli import SimpleTokenizer
+  from tinygrad.llm.model import Transformer
+
+def parse_tool_call(s:str) -> tuple[str, typing.Any]|None:
+  s = s.strip()
+  if s.startswith("{"):  # hermes JSON format: {"name": ..., "arguments": {...}}
+    try:
+      call = json.loads(s)
+      return call["name"], call.get("arguments", call.get("parameters", {}))
+    except (json.JSONDecodeError, KeyError): return None
+  # XML format: <function=name>\n<parameter=key>\nvalue\n</parameter>...</function>
+  if (fm := re.match(r"<function=([^>]+)>\s*(.*?)\s*(?:</function>)?$", s, re.DOTALL)):
+    args = {}
+    for pm in re.finditer(r"<parameter=([^>]+)>(.*?)</parameter>", fm.group(2), re.DOTALL):
+      value = re.sub(r"^\r?\n|\r?\n\Z", "", pm.group(2))
+      try: args[pm.group(1)] = json.loads(value)
+      except json.JSONDecodeError: args[pm.group(1)] = value
+    return fm.group(1), args
+  return None
+
+def normalize_messages(messages:list[dict]) -> None:
+  # chat templates expect tool_call arguments as dicts (OpenAI clients send JSON strings)
+  for m in messages:
+    for tc in m.get("tool_calls") or []:
+      if "function" in tc and isinstance(args := tc["function"].get("arguments"), str):
+        try: tc["function"]["arguments"] = json.loads(args)
+        except json.JSONDecodeError: pass
+
+class StreamRouter:
+  # routes streamed output text to (field, text) deltas, keeping tool_call regions in .buf for the final parse
+  def __init__(self):
+    self.buf = ""
+    self.mode = "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:
+      before, self.buf = self.buf.split(tag, 1)
+      return before, True
+    hold = max((i for i in range(1, min(len(self.buf), len(tag))+1) if tag.startswith(self.buf[-i:])), default=0) if not final else 0
+    emit, self.buf = self.buf[:len(self.buf)-hold], self.buf[len(self.buf)-hold:]
+    return emit, False
+  def route(self, piece:str, final:bool=False) -> typing.Iterator[tuple[str, str]]:
+    self.buf += piece
+    if self.mode == "undecided":  # decide whether the output starts with a think block
+      if not final and len(self.buf) < len("<think>") and "<think>".startswith(self.buf): return
+      self.mode, self.buf = ("reasoning", self.buf[len("<think>"):]) if self.buf.startswith("<think>") else ("content", self.buf)
+    if self.mode == "reasoning":
+      emit, done = self.split("</think>", final)
+      if emit: yield "reasoning_content", emit
+      if not done: return
+      self.mode = "content"
+    if self.mode == "tool": return
+    emit, found = self.split("<tool_call>", final)
+    if emit: yield "content", emit
+    if found: self.mode, self.buf = "tool", "<tool_call>" + self.buf
+
+class Handler(HTTPRequestHandler):
+  server: LLMServer
+  def log_request(self, code='-', size='-'): pass
+  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):
+    model, tok = self.server.model, self.server.tok
+    cache_start_pos = model.get_start_pos(ids)
+    stderr_log(f"in:{colored(f'{cache_start_pos:5d}', 'green')} +{len(ids)-cache_start_pos:5d}  {colored('--', 'BLACK')}  ")
+    tmpl = {"id":f"chatcmpl-{uuid.uuid4().hex[:24]}", "object":"chat.completion.chunk", "created":int(time.time()), "model":model_name}
+    def chunk(d:dict): return {"choices": [{"index":0, "delta":d, "finish_reason":None}], **tmpl}
+    yield chunk({"role":"assistant", "content":""})
+    out: list[int] = []
+    finish_reason = "stop"
+    st = time.perf_counter()
+    dec = tok.stream_decoder()
+    router = StreamRouter()
+    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 tok.is_end(next_id): break
+      out.append(next_id)
+      for field, delta in router.route(dec(next_id)): yield chunk({field:delta})
+      if max_tokens is not None and len(out) >= max_tokens:
+        finish_reason = "length"
+        break
+    for field, delta in router.route(dec(), final=True): yield chunk({field:delta})
+    tool_calls: list[dict] = []
+    for m in re.finditer(r"<tool_call>\s*(.*?)\s*(?:</tool_call>|$)", router.buf, re.DOTALL):
+      if (parsed := parse_tool_call(m.group(1))) is None:
+        stderr_log(f"failed to parse tool call: {m.group(1)[:200]}")
+        yield chunk({"content":m.group(0)})  # don't silently drop output the client can't use
+      else:
+        name, args = parsed
+        tool_calls.append({"index":len(tool_calls), "id":f"call_{uuid.uuid4().hex[:24]}", "type":"function",
+                           "function":{"name":name, "arguments":args if isinstance(args, str) else json.dumps(args)}})
+    if tool_calls:
+      yield chunk({"tool_calls":tool_calls})
+      if finish_reason == "stop": finish_reason = "tool_calls"
+    yield {"choices": [{"index":0, "delta":{},"finish_reason":finish_reason}], **tmpl}
+    if include_usage:
+      yield {"choices": [], "usage": {"prompt_tokens": len(ids), "completion_tokens": len(out), "total_tokens": len(ids) + len(out)}, **tmpl}
+    et = time.perf_counter()
+    stderr_log(f"gen:{len(out)/(et-pt) if len(out) > 1 else 0:4.0f} tok/s  {colored('--', 'BLACK')}  "
+               f"out:{len(out):5d}  {colored('--', 'BLACK')}  total:{et-st:6.2f}s\n")
+
+  def do_POST(self):
+    request_st = time.perf_counter()
+    stderr_log(f"{self.path}  {colored('--', 'BLACK')}  ")
+    raw_body = self.rfile.read(int(self.headers.get("Content-Length", "0")))
+    body: dict[str, typing.Any] = json.loads(raw_body.decode("utf-8"))
+    if DEBUG >= 1: print(json.dumps(body, indent=2))
+    if self.path == "/v1/chat/completions":
+      # render and tokenize
+      normalize_messages(body["messages"])
+      rendered = self.server.template.render(messages=body["messages"], tools=body.get("tools"), add_generation_prompt=True)
+      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)
+
+      # 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)))
+      if body.get("stream"): self.stream_json(chunks)
+      else:
+        out, reasoning, tool_calls, finish_reason = [], [], [], "stop"
+        for c in chunks:
+          if not c["choices"]: continue
+          choice = c["choices"][0]
+          if (delta := choice.get("delta", {})):
+            if delta.get("content"): out.append(delta["content"])
+            if delta.get("reasoning_content"): reasoning.append(delta["reasoning_content"])
+            tool_calls += [{k:v for k, v in tc.items() if k != "index"} for tc in delta.get("tool_calls", [])]
+          if choice.get("finish_reason"): finish_reason = choice["finish_reason"]
+        message: dict[str, typing.Any] = {"role":"assistant", "content":"".join(out) or None}
+        if reasoning: message["reasoning_content"] = "".join(reasoning)
+        if tool_calls: message["tool_calls"] = tool_calls
+        self.send_data(json.dumps({**c, "object":"chat.completion",
+          "choices":[{"index":0, "message":message, "finish_reason":finish_reason}]}).encode())
+    else:
+      raise RuntimeError(f"unhandled path {self.path}")
+
+class LLMServer(TCPServerWithReuse):
+  def __init__(self, server_address:tuple, model:Transformer, model_name:str, tok:SimpleTokenizer, template:typing.Any):
+    self.model, self.model_name, self.tok, self.template = model, model_name, tok, template
+    super().__init__(server_address, Handler)

tinygrad/mixin/elementwise.py

@@ -49,6 +49,7 @@ 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)
+    m = self.max(axis=axis, keepdim=True).detach()
     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)
+    x_cummax = x.cummax(-1)[0].detach()
     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
+from tinygrad.tensor import Tensor, is_numpy_ndarray
 from tinygrad.mixin.op import ReductionStr
-from tinygrad.helpers import getenv, all_same, prod, flatten, make_tuple, argsort, is_numpy_ndarray, get_single_element, polyN, Context
+from tinygrad.helpers import getenv, all_same, prod, flatten, make_tuple, argsort, 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,6 +70,7 @@ 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,8 @@ def uops_to_dtypes(uops:list[UOp]) -> list[tuple[DType, int]]:
   return dedup((u.dtype, u.max_numel()) for u in uops if u.addrspace in (AddrSpace.ALU, None) and u.dtype != dtypes.void and u._shape is not None)
 
 def _wmma_name(u:UOp) -> str:
-  return f"WMMA_{'_'.join(map(str, u.arg[0]))}_{u.arg[1].name}_{u.dtype.scalar().name}"
+  # 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)
 def wmma_args(uops:list[UOp]):
@@ -565,6 +566,11 @@ class HIPRenderer(CStyleLanguage):
       # #define __WMMA_16_16_16_half_half __builtin_amdgcn_wmma_f16_16x16x16_f16_w32_gfx12
       elif self.tensor_cores == tc.amd_rdna4:
         prefix.append(f"#define __{name} __builtin_amdgcn_wmma_{type_map[dtype_out]}_16x16x16_{type_map[dtype_in]}_w32_gfx12")
+      elif dtype_out == dtypes.int32:
+        prefix.append("typedef int wmma_int4 __attribute__((ext_vector_type(4)));\n"+
+          f"static inline __attribute__((device)) int8 __{name}"+"""(signed_char16 a, signed_char16 b, int8 c) {
+  return __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32(true, __builtin_bit_cast(wmma_int4, a),
+    true, __builtin_bit_cast(wmma_int4, b), c, false);\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.SHL, X86Ops.SHR, X86Ops.SAR, X86Ops.IMUL, X86Ops.CMP,
+  ReadMem2nd = {X86Ops.ADD, X86Ops.SUB, X86Ops.AND, X86Ops.OR, X86Ops.XOR, 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,8 +99,9 @@ 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.SHLi, X86Ops.SHRi, 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.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 that read flags
   ReadFlags = {X86Ops.CMOVB, X86Ops.CMOVL, X86Ops.CMOVE, X86Ops.CMOVNE, X86Ops.SETB, X86Ops.SETL, X86Ops.SETE, X86Ops.SETNE, X86Ops.JB, X86Ops.JL,
@@ -272,6 +273,11 @@ 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]:
@@ -452,9 +458,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.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(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.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))),
@@ -651,7 +657,8 @@ 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
-    return _encode(rest[0], *address, *(None, *rest[1:])) if reg is None else _encode(None, *address, *(None, *rest[:1]))
+    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))
 
   if x.arg in X86GroupOp.Rm1st:
     if len(x.src) > 3: address, rest = x.src[:4], x.src[4:]
@@ -704,8 +711,9 @@ 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.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.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.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),
@@ -784,6 +792,7 @@ 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.fp8e4m3: ".fp8.fp8", dtypes.fp8e5m2: ".bf8.bf8", dtypes.int8: "iu8", dtypes.int32: "i32"}
   # https://github.com/llvm/llvm-project/blob/main/clang/test/CodeGenOpenCL/builtins-amdgcn-mfma.cl
   N,M,K = wmma.arg[0]
   if cdna:
@@ -44,9 +44,10 @@ 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.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 ")")
+    f"{dt_map[wmma.arg[1]]}(" + ", ".join(args) + (", 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",
@@ -261,6 +262,9 @@ 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 = bytearray(lib[(image_offset:=_read_lib(lib, 0xc0)):image_offset+self.image_size])
+    self.image = 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.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.multiple_of, 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,6 +434,7 @@ 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)
@@ -470,7 +471,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)
+   UOp.variable(v.arg.name, v.arg.vmin_vmax[0], v.arg.vmin_vmax[1], v.dtype, multiple_of=v.arg.multiple_of)
    if v.arg.name is not None and v.arg.vmin_vmax is not None else None),
 
   # this renumbers the params
@@ -522,7 +523,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]):
+  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]):
     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, TypeVar, Generic, TYPE_CHECKING
+from typing import Any, Callable, cast, get_args, ParamSpec, TypeGuard, TypeVar, Generic, TYPE_CHECKING
 if TYPE_CHECKING: import numpy
-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.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.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,6 +34,8 @@ 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
@@ -75,22 +77,21 @@ 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 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:
+    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):
         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)}")
@@ -178,6 +179,7 @@ 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)
@@ -205,14 +207,16 @@ 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}")
 
@@ -253,6 +257,7 @@ 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()
@@ -276,6 +281,7 @@ 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:
@@ -293,6 +299,7 @@ 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()
@@ -450,7 +457,10 @@ 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 isinstance(v, Tensor) and v.dtype != self.dtype: raise RuntimeError(f"setitem dtype mismatch: {self.dtype=} != {v.dtype=}")
+    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=}")
     # 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,6 +11,8 @@ 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, strong_dtype, Invalid, AddrSpace
+from tinygrad.dtype import ConstType, dtypes, DType, DTypeLike, truncate, least_upper_dtype, least_upper_float, 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,12 +24,13 @@ 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), ("name", None), ("addrspace", AddrSpace.GLOBAL), ("axis", None), ("device", None))
+    fields = (("vmin_vmax", None), ("multiple_of", 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",
@@ -704,6 +705,7 @@ 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)
@@ -758,6 +760,7 @@ 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
@@ -785,7 +788,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(dtype=strong_dtype(self.dtype), device=device)
+    ret = self.empty_like(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
@@ -809,10 +812,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}:
+    if self.op in {Ops.INDEX, Ops.CAST, Ops.AFTER, Ops.REDUCE, Ops.STORE, Ops.MSTACK, Ops.MSELECT, Ops.END}:
       return self.src[0].addrspace
     if self.op in GroupOp.Movement: return self.src[0].addrspace
-    if self.op in {Ops.STACK, Ops.WMMA} or self.op in GroupOp.Elementwise:
+    if self.op in {Ops.STACK, Ops.WMMA, Ops.GROUP} 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]
@@ -913,9 +916,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   # *** uop Variable stuff ***
 
   @staticmethod
-  def variable(name:str, min_val:PyConst, max_val:PyConst, dtype:DType=dtypes.index) -> UOp:
+  def variable(name:str, min_val:PyConst, max_val:PyConst, dtype:DType=dtypes.index, multiple_of:int=1) -> UOp:
     return UOp(Ops.PARAM, src=(shape_to_shape_arg(()),),
-               arg=ParamArg(-1, dtype, name=name, vmin_vmax=(min_val, max_val), addrspace=AddrSpace.ALU))
+               arg=ParamArg(-1, dtype, name=name, vmin_vmax=(min_val, max_val), multiple_of=multiple_of, addrspace=AddrSpace.ALU))
   @property
   def expr(self) -> str:
     assert self.op is Ops.PARAM
@@ -924,6 +927,7 @@ 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}"
@@ -946,6 +950,7 @@ 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
@@ -957,6 +962,7 @@ 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))
@@ -1027,9 +1033,12 @@ 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
-    # 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)
+    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)
     return self.dtype.min, self.dtype.max
 
   @functools.cached_property
@@ -1082,16 +1091,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, name=None,
-            addrspace=AddrSpace.GLOBAL, axis:int|None=None):
+  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}")
     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, name, addrspace, axis, device))
+    return UOp(Ops.PARAM, src=src, arg=ParamArg(slot, dtype, vmin_vmax, multiple_of, 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 UOp.param(slot, self.dtype, self._shape, self.device, cast(tuple[int, int], self._min_max), self.src[0].expr, addrspace)
+    if self.op is Ops.BIND: return self.src[0].replace(arg=replace(self.src[0].arg, slot=slot, addrspace=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,15 +61,13 @@ 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: most ALUs have all matching dtypes, except CMPLT, CMPNE, and WHERE
+  # ALUs: operands match the result dtype, except comparisons/WHERE; renderer-lowered shifts may use a uint32 count
   # a weak dtype matches any dtype (TODO: make python scalars weak consts)
-  (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"))),
+  (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"))),
    lambda x,y: x.dtype == y.dtype or x.dtype in dtypes.weaks or y.dtype in dtypes.weaks),
-  # 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.SHL, Ops.SHR), src=(UPat.var("x"), UPat(dtype=dtypes.uint)), name="a"), lambda a,x: a.dtype == x.dtype or None),
   (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="2" width="20"><path d="M15 19l-7-7 7-7"/></svg>
+        <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>
       </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,6 +735,7 @@ async function renderProfiler(path, opts) {
     }
   }
 
+  let lastCanvasRect = null;
   function resize() {
     const [width, height] = canvasDims();
     if (canvas.width === width*dpr && canvas.height === height*dpr) return;
@@ -743,6 +744,11 @@ 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);
   }