Merge remote-tracking branch 'origin/master' into delta_attention

* hcq2: c submitter

* x

* x

* x

* simpler

* simpler

* x

* x

* Dx

* revrt

* Dx

* x

* fst

* fix

.github/workflows/test.yml

@@ -86,6 +86,27 @@ jobs:
     - name: Custom tests
       run: DEV=CPU:LLVM GPUS=4 TINY_BACKEND=1 python3 -m pytest -nauto extra/torch_backend/test.py extra/torch_backend/test_inplace.py extra/torch_backend/test_multigpu.py extra/torch_backend/test_kernel_fusion.py --durations=20
 
+  torchbackendtrain:
+    name: Torch Backend Training
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    timeout-minutes: 15
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: torch-backend-pillow-torchvision-et-pt
+        deps: testing_unit
+        pydeps: "pillow torchvision expecttest"
+        llvm: 'true'
+    - name: Install ninja
+      run: |
+        sudo apt update || true
+        sudo apt install -y --no-install-recommends ninja-build
+    - name: Test beautiful_mnist in torch with TINY_BACKEND
+      run: STEPS=20 DEV=CPU TARGET_EVAL_ACC_PCT=90.0 MAX_BUFFER_SIZE=0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+
   bepython:
     name: Python Backend
     runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}

README.md

@@ -140,7 +140,7 @@ Documentation along with a quick start guide can be found on the [docs website](
 ```python
 from tinygrad import Tensor
 
-x = Tensor.eye(3)
+x = Tensor.eye(3).clone()  # clone to make it a buffer
 y = Tensor([[2.0,0,-2.0]])
 z = y.matmul(x).sum()
 z.backward()

extra/hcq2/ops_amd2.py

@@ -561,9 +561,7 @@ class AMDDevice(HCQ2Compiled):
   def is_usb(self) -> bool: return False
 
   def __init__(self, device:str=""):
-    self.device_id = int(device.split(":")[1]) if ":" in device else 0
-
-    self.iface = self._select_iface()
+    self.iface = self._select_iface(device)
 
     self.target:tuple[int, ...] = ((trgt:=self.iface.props['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
     self.arch = "gfx%d%x%x" % self.target

extra/thunder/tiny/tk/tiles.py

@@ -209,7 +209,7 @@ class ST:
     return cls(uop, rows, cols, layout, base_shape, ker)
 
   def swizzle(self, row, col):
-    swizzled_offset = self.base_shape.swizzle(row, col, self._uop.dtype.scalar())
+    swizzled_offset = self.base_shape.swizzle(row, col, self._uop.dtype)
 
     row = swizzled_offset // self.base_shape.cols
     col = swizzled_offset % self.base_shape.cols

extra/torch_backend/backend.py

@@ -4,7 +4,7 @@
 # A006 Lambda argument `input` is shadowing a Python builtin
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.uop.ops import Ops, GroupOp
-from tinygrad.helpers import getenv, prod, strides_for_shape, argfix
+from tinygrad.helpers import getenv, prod, strides_for_shape
 import torch.lib
 TORCH_DEBUG = getenv("TORCH_DEBUG")
 import torch, pathlib, operator, functools, weakref
@@ -88,9 +88,6 @@ view_ops = {
   "aten.diagonal": Tensor.diagonal,
   }
 
-# torch 2.10 handles this natively
-if tuple(map(int, torch.__version__.split('.')[:2])) < (2, 10): view_ops.update({"aten.detach": Tensor.detach})
-
 for k,v in view_ops.items(): torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(wrap_view_op(v))
 
 def _get_view_ops(view): return getattr(view, "_view_ops", [])
@@ -99,46 +96,21 @@ def _apply_view_ops(target, ops):
   for fn, args, kwargs in ops: target = fn(target, *args, **kwargs)
   return target
 
-# similar to https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/InferSize.h
-def _reshape_target_shape(shape:tuple[int, ...], args) -> tuple[int, ...]|None:
-  if not (req := argfix(*args)): return None
-  new_shape, infer_idx = [], -1
-  for i, s in enumerate(req):
-    if s is None: s = shape[i] if i < len(shape) else None
-    if not isinstance(s, int): return None
-    if s == -1:
-      if infer_idx != -1: return None
-      infer_idx = len(new_shape)
-    new_shape.append(s)
-  total = prod(shape)
-  if infer_idx != -1:
-    known = prod(x for x in new_shape if x != -1)
-    if known == 0:
-      if total != 0: return None
-      new_shape[infer_idx] = 0
-    else: new_shape[infer_idx] = total // known
-  return tuple(new_shape) if prod(new_shape) == total else None
-
-# TODO: can we get rid of this? only for test_flatten_reshape_add
+# a chain of reshapes is undone by reshaping the value back to the base
 def _try_simple_reshape_view_write(base: Tensor, view: Tensor, val: Tensor) -> bool:
   if not (ops := _get_view_ops(view)): return False
-  shapes = [base.shape]
-  for fn, args, _ in ops:
-    if fn is Tensor.reshape:
-      if not (next_shape := _reshape_target_shape(shapes[-1], args)): return False
-      shapes.append(next_shape)
-  if shapes[-1] != view.shape: return False
-  for s in reversed(shapes[:-1]): val = val.reshape(s)
-  base.assign(val)
+  if any(fn is not Tensor.reshape for fn, _, _ in ops): return False
+  base.assign(val.reshape(base.shape))
   return True
 
 def _view_write(base: Tensor, view: Tensor, value: Tensor) -> None:
   val = value if value.dtype == base.dtype else value.cast(base.dtype)
-  if view.shape == base.shape: return base.assign(val)
   if _try_simple_reshape_view_write(base, view, val): return
   idx_base = Tensor.arange(base.numel(), dtype=dtypes.int32).reshape(base.shape)
   idx_view = _apply_view_ops(idx_base, _get_view_ops(view)).reshape(-1)
-  flat_base = base.reshape(base.numel()).contiguous()
+  # clone, not contiguous: contiguous() on a base that already owns its buffer returns the base itself, and scattering
+  # into that is an in-place write to a buffer other tensors still hold, which setitem refuses
+  flat_base = base.reshape(base.numel()).clone()
   flat_base[idx_view] = val.reshape(-1)
   base.assign(flat_base.reshape(base.shape))
 
@@ -301,6 +273,34 @@ def slice_tensor(self, dim=0, start=None, end=None, step=1):
   slices[dim] = slice(start, end, step)
   return self[slices]
 
+# the functional scatters. without an impl aten falls back to a path that assumes a real storage: "self.has_storage() INTERNAL ASSERT FAILED"
+def _scatter_into(self, src, dim, index):
+  out = unwrap(self).clone()
+  slices = [slice(None)] * out.ndim
+  slices[dim] = index
+  out[slices] = unwrap(src).cast(out.dtype)  # torch casts src to self's dtype, tinygrad setitem demands they already match
+  return wrap(out)
+
+@torch.library.impl("aten::slice_scatter", "privateuseone")
+def slice_scatter(self, src, dim=0, start=None, end=None, step=1): return _scatter_into(self, src, dim, slice(start, end, step))
+
+@torch.library.impl("aten::select_scatter", "privateuseone")
+def select_scatter(self, src, dim, index): return _scatter_into(self, src, dim, index)
+
+@torch.library.impl("aten::diagonal_scatter", "privateuseone")
+def diagonal_scatter(self, src, offset=0, dim1=0, dim2=1):
+  # a diagonal is not one axis, so scatter through the flat indices it picks out
+  base, out = unwrap(self), unwrap(self).clone().reshape(-1)
+  idx = Tensor.arange(base.numel(), dtype=dtypes.int32).reshape(base.shape).diagonal(offset, dim1, dim2).reshape(-1)
+  out[idx] = unwrap(src).cast(base.dtype).reshape(-1)
+  return wrap(out.reshape(base.shape))
+
+# the functional copy_. without an impl the fallback segfaults on a tensor with no storage
+@torch.library.impl("aten::copy", "privateuseone")
+def copy(self, src, non_blocking=False):
+  dest = unwrap(self)
+  return wrap(unwrap(src).cast(dest.dtype).to(dest.device).expand(dest.shape))
+
 @torch.library.impl("aten::slice_backward", "privateuseone")
 def slice_backward(grad_out, input_sizes, dim, start, end, step):
   grad_input = Tensor.zeros(input_sizes).contiguous()
@@ -341,7 +341,9 @@ for dim in [1, 2, 3]:
     torch.library.impl(f"aten::{pad_type}_pad{dim}d", "privateuseone")(functools.partial(pad_forward, mode=mode))
     torch.library.impl(f"aten::{pad_type}_pad{dim}d_backward", "privateuseone")(functools.partial(pad_backward, mode=mode))
 
-def upsample(self, size, align_corners=False, mode=None): return wrap(Tensor.interpolate(unwrap(self), size, mode=mode, align_corners=align_corners))
+# the schemas are all positional: (self, output_size, align_corners, *scales) for linear, (self, output_size, *scales) for nearest.
+def upsample(self, size, *args, mode=None):
+  return wrap(Tensor.interpolate(unwrap(self), size, mode=mode, align_corners=args[0] if mode == "linear" else False))
 for i,pre in enumerate(["", "bi", "tri"]):
   torch.library.impl(f"aten::upsample_{pre}linear{i+1}d", "privateuseone")(functools.partial(upsample, mode="linear"))
   torch.library.impl(f"aten::upsample_nearest{i+1}d", "privateuseone")(functools.partial(upsample, mode="nearest"))
@@ -413,6 +415,7 @@ def _linalg_svd(self, full_matrices=False):
 from torch._decomp import get_decompositions
 decomps = [
   aten.native_layer_norm_backward,
+  aten.native_group_norm_backward,
   aten.linalg_cross,
   aten.addmm,
   aten.addcmul,
@@ -447,12 +450,20 @@ decomps = [
   aten._softmax_backward_data, aten.embedding_dense_backward,
   aten.linalg_vector_norm,
   aten.binary_cross_entropy, aten.binary_cross_entropy_backward,
+  # the C++ mse/smooth_l1 kernels resize their out tensor, and a tiny tensor has no storage to resize
+  aten.mse_loss, aten.mse_loss_backward,
+  aten.smooth_l1_loss, aten.smooth_l1_loss_backward,
   aten.upsample_nearest2d.out,
+  # NOTE: only the "out" overload, the "vec" one is CompositeImplicitAutograd and overriding it loses the autograd kernel
+  aten.upsample_bicubic2d.out,
+  aten._adaptive_avg_pool2d,
   # activations
   aten.hardswish, aten.hardswish_backward,
   aten.hardtanh, aten.hardtanh_backward,
   aten.gelu, aten.gelu_backward,
-  aten.logical_and,
+  # NOTE: no aten.logical_or here, its decomposition reaches aten.bitwise_or through a path that checks aliasing by
+  # reading storage, which a tiny tensor has none of. it gets a direct impl below instead
+  aten.logical_and, aten.logical_xor,
   aten.randint,
   aten.eye,
   aten.hardsigmoid_backward,
@@ -495,7 +506,7 @@ simple_tensor_methods = [
   # reduce
   "all", "any", "argmax", "argmin", "cumsum", "cumprod",
   # complex
-  "avg_pool2d", "linspace"]
+  "linspace"]
 
 tiny_backend_out = {**{f"aten.{x}.out":getattr(Tensor,x) for x in simple_tensor_methods}, **{
   "aten.add.out": lambda input,other,alpha=1: input+alpha*other,
@@ -579,8 +590,8 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   # inplace ops using replace for fusion
   "aten.zero_": lambda x: x.const_like(0),
   "aten.fill_.Scalar": lambda x, y: x.const_like(y),
-  "aten.add_.Tensor": lambda self, other, alpha=1.0: self + other * alpha,
-  "aten.add_.Scalar": lambda self, other, alpha=1.0: self + other * alpha,
+  "aten.add_.Tensor": lambda self, other, alpha=1: self + other * alpha,
+  "aten.add_.Scalar": lambda self, other, alpha=1: self + other * alpha,
   "aten.mul_.Tensor": lambda self, other: self * other,
   "aten.mul_.Scalar": lambda self, other: self * other,
   # relu doesn't have an out form?
@@ -613,7 +624,9 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   # these don't work in out form, they have size 0
   "aten.abs": Tensor.abs,
   "aten.logical_not": Tensor.logical_not,
-  "aten.logical_or_": lambda x, y: x | y,
+  # compare against zero first: logical_* is bool-valued for any input dtype, while | is bitwise
+  "aten.logical_or": lambda x, y: (x != 0) | (y != 0),
+  "aten.logical_or_": lambda x, y: (x != 0) | (y != 0),
   "aten.multinomial": Tensor.multinomial,
   "aten.masked_fill_.Scalar": lambda self, mask, value: self.masked_fill(mask, value),
   "aten.masked_fill_.Tensor": lambda self, mask, value: self.masked_fill(mask, value),
@@ -622,14 +635,7 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   "aten.masked_select": Tensor.masked_select,
   "aten.all": Tensor.all,
   "aten.sgn": Tensor.sign,
-  "aten.acos": Tensor.acos,
   "aten.any": Tensor.any,
-  "aten.bitwise_not": Tensor.bitwise_not,
-  "aten.argmax": Tensor.argmax,
-  "aten.argmin": Tensor.argmin,
-  "aten.asinh": Tensor.asinh,
-  "aten.mul": Tensor.mul,
-  "aten.atanh": Tensor.atanh,
   "aten.fill_.Tensor": lambda self, value: self.const_like(value.reshape(()).item()),
   "aten.flip": Tensor.flip,
   "aten.scatter_reduce.two": Tensor.scatter_reduce,
@@ -642,8 +648,7 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
     Tensor.linspace(start, stop, steps, **({"dtype": _from_torch_dtype(dtype)} if dtype is not None else {})),
   "aten.topk": Tensor.topk,
   "aten.constant_pad_nd": lambda self, padding, value=0.0: self.pad(padding, mode="constant", value=value).contiguous(),
-  # TODO: input contiguous is needed to prevent CFGContext circular dependency assertion for shapes >512 (see test_cumsum_arange_large)
-  "aten.cumsum": lambda self, dim: self.contiguous().cumsum(dim),
+  "aten.cumsum": lambda self, dim: self.cumsum(dim),
   "aten.logsumexp": lambda self, axis, keepdim=False: self.logsumexp(axis[0], keepdim=keepdim),
   "aten.roll": Tensor.roll,
   "aten.logcumsumexp": Tensor.logcumsumexp,
@@ -652,6 +657,7 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
     self.ones_like(**{k: v for k, v in {"dtype": _from_torch_dtype(dtype) if dtype else None,
                                         "device": _from_torch_device(device) if device else None}.items() if v is not None}),
   "aten.max.dim": lambda self, dim, keepdim=False: (self.max(dim, keepdim), self.argmax(dim, keepdim).cast(dtype=dtypes.int64)),
+  "aten.min.dim": lambda self, dim, keepdim=False: (self.min(dim, keepdim), self.argmin(dim, keepdim).cast(dtype=dtypes.int64)),
   "aten.cummax": lambda self, dim: ((r := self.cummax(dim))[0], r[1].cast(dtypes.int64)),
   "aten.cummin": lambda self, dim: ((r := self.cummin(dim))[0], r[1].cast(dtypes.int64)),
   "aten.nonzero": Tensor.nonzero,

extra/torch_backend/test.py

@@ -166,6 +166,15 @@ class TestTorchBackend(unittest.TestCase):
     expected = np.array([[1.5, 5.2, 9.0], [13.2, 17.1, 18.4]], dtype=np.float32)
     np.testing.assert_equal(y3.cpu().numpy(), expected)
 
+  def test_argmax_argmin(self):
+    a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
+    c = a.cpu()
+    for got, want in [(a.argmax(), c.argmax()), (a.argmin(0), c.argmin(0)), (a.argmax(1, keepdim=True), c.argmax(1, keepdim=True)),
+                      (torch.min(a, 1).indices, torch.min(c, 1).indices), (torch.max(a, 1).indices, torch.max(c, 1).indices),
+                      (torch.min(a, 1).values, torch.min(c, 1).values), (torch.min(a, 1, keepdim=True).indices, torch.min(c, 1, keepdim=True).indices)]:
+      self.assertEqual(got.dtype, want.dtype)  # torch's arg reduces are int64, tinygrad's are int32
+      np.testing.assert_equal(got.cpu().numpy(), want.numpy())
+
   def test_isfinite(self):
     a = torch.ones(4, device=device)
     np.testing.assert_equal(torch.isfinite(a).cpu().numpy(), [True, True, True, True])
@@ -373,6 +382,22 @@ class TestTorchBackend(unittest.TestCase):
     for bwd_eps in [1e-5, 0.3]:
       for got, want in zip(run(device, bwd_eps), run("cpu", bwd_eps)): np.testing.assert_allclose(got, want, atol=1e-4, rtol=1e-3)
 
+  def test_groupnorm_backward(self):
+    def run(dev):
+      x = torch.arange(24., device=dev).reshape(2, 4, 3).requires_grad_()
+      w = torch.linspace(0.5, 2.0, 4).to(dev).requires_grad_()
+      torch.nn.functional.group_norm(x, 2, w, torch.zeros(4, device=dev)).square().sum().backward()
+      return x.grad.cpu().numpy(), w.grad.cpu().numpy()
+    for got, want in zip(run(device), run("cpu")): np.testing.assert_allclose(got, want, atol=1e-4, rtol=1e-3)
+
+  def test_mse_smooth_l1_loss_backward(self):
+    def run(dev, loss):
+      x = torch.arange(4., device=dev).requires_grad_()
+      loss(x, torch.ones(4, device=dev)).backward()
+      return x.grad.cpu().numpy()
+    for loss in [torch.nn.functional.mse_loss, torch.nn.functional.smooth_l1_loss]:
+      np.testing.assert_allclose(run(device, loss), run("cpu", loss), atol=1e-6)
+
   def test_batchnorm_unsqueeze(self):
     bn = torch.nn.BatchNorm2d(4).to(device)
     x = torch.randn(8, 4, 3, 3, device=device)
@@ -796,6 +821,86 @@ class TestTorchBackend(unittest.TestCase):
     np.testing.assert_allclose(w_tiny.grad.cpu().numpy(), w_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)
     np.testing.assert_allclose(b_tiny.grad.cpu().numpy(), b_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)
 
+  def test_write_through_detach_of_unrealized(self):
+    a = torch.empty(4, device=device)
+    a.detach().fill_(3)
+    np.testing.assert_equal(a.cpu().numpy(), [3, 3, 3, 3])
+
+  def test_square_transpose_inplace(self):
+    # a same-shape transpose is not a reshape: writing the transposed values straight back would scramble the base
+    a = torch.tensor([[0., 1., 2.], [3., 4., 5.], [6., 7., 8.]], device=device)
+    a.transpose(0, 1).add_(100)
+    np.testing.assert_equal(a.cpu().numpy(), [[100., 101., 102.], [103., 104., 105.], [106., 107., 108.]])
+
+  def test_interpolate(self):
+    a = torch.arange(4, dtype=torch.float32, device=device).reshape(1, 1, 2, 2)
+    nearest = torch.nn.functional.interpolate(a, scale_factor=2.0)
+    np.testing.assert_equal(nearest.cpu().numpy()[0, 0], [[0, 0, 1, 1], [0, 0, 1, 1], [2, 2, 3, 3], [2, 2, 3, 3]])
+    linear = torch.nn.functional.interpolate(a, size=(4, 4), mode="bilinear", align_corners=False)
+    ref = torch.nn.functional.interpolate(a.cpu(), size=(4, 4), mode="bilinear", align_corners=False)
+    np.testing.assert_allclose(linear.cpu().numpy(), ref.numpy(), rtol=1e-5)
+
+  def test_interpolate_bicubic_area(self):
+    a = torch.arange(32, dtype=torch.float32, device=device).reshape(1, 2, 4, 4)
+    for mode, scale in [("bicubic", 2.0), ("area", 0.5)]:
+      ref = torch.nn.functional.interpolate(a.cpu(), scale_factor=scale, mode=mode)
+      np.testing.assert_allclose(torch.nn.functional.interpolate(a, scale_factor=scale, mode=mode).cpu().numpy(), ref.numpy(), atol=1e-4)
+
+  @unittest.expectedFailure
+  def test_interpolate_bicubic_backward(self):
+    # the forward comes from a decomposition, but aten::upsample_bicubic2d_backward has none (nor does
+    # aten::_adaptive_avg_pool2d_backward, for area), so training through these modes needs a real kernel
+    x = torch.arange(32., dtype=torch.float32, device=device).reshape(1, 2, 4, 4).requires_grad_()
+    torch.nn.functional.interpolate(x, scale_factor=2.0, mode="bicubic").sum().backward()
+
+  @unittest.expectedFailure
+  def test_interpolate_inexact_scale(self):
+    # torch forwards the raw scale_factor, Tensor.interpolate recomputes it from output_size, and they disagree here
+    a = torch.arange(6, dtype=torch.float32, device=device).reshape(1, 1, 2, 3)
+    tiny = torch.nn.functional.interpolate(a, scale_factor=2.5, mode="bilinear")
+    ref = torch.nn.functional.interpolate(a.cpu(), scale_factor=2.5, mode="bilinear")
+    np.testing.assert_allclose(tiny.cpu().numpy(), ref.numpy(), rtol=1e-5)
+
+  def test_logical_or_xor(self):
+    a = torch.tensor([True, True, False, False], device=device)
+    b = torch.tensor([True, False, True, False], device=device)
+    np.testing.assert_equal(torch.logical_or(a, b).cpu().numpy(), [True, True, True, False])
+    np.testing.assert_equal(torch.logical_xor(a, b).cpu().numpy(), [False, True, True, False])
+    # bool-valued whatever the input dtype, so this is not | and ^
+    i, j = torch.tensor([2, 0, 5, 0], device=device), torch.tensor([0, 0, 1, 1], device=device)
+    np.testing.assert_equal(torch.logical_or(i, j).cpu().numpy(), [True, False, True, True])
+    np.testing.assert_equal(torch.logical_xor(i, j).cpu().numpy(), [True, False, False, True])
+
+  def test_slice_scatter(self):
+    # the scatters are functional: they return a new tensor and must leave the one they were given alone
+    a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
+    out = torch.slice_scatter(a, torch.ones(1, 4, device=device), 0, 0, 1)
+    np.testing.assert_equal(out.cpu().numpy(), [[1, 1, 1, 1], [4, 5, 6, 7], [8, 9, 10, 11]])
+    np.testing.assert_equal(a.cpu().numpy(), np.arange(12, dtype=np.float32).reshape(3, 4))
+
+  def test_slice_scatter_casts_src(self):
+    a = torch.zeros(3, 4, device=device)
+    out = torch.slice_scatter(a, torch.ones(1, 4, dtype=torch.int32, device=device), 0, 0, 1)
+    self.assertEqual(out.dtype, torch.float32)
+    np.testing.assert_equal(out.cpu().numpy()[0], np.ones(4, dtype=np.float32))
+
+  def test_select_scatter(self):
+    a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
+    out = torch.select_scatter(a, torch.ones(4, device=device), 0, 1)
+    np.testing.assert_equal(out.cpu().numpy(), [[0, 1, 2, 3], [1, 1, 1, 1], [8, 9, 10, 11]])
+
+  def test_diagonal_scatter(self):
+    a = torch.zeros(3, 3, device=device)
+    out = torch.diagonal_scatter(a, torch.arange(3, dtype=torch.float32, device=device))
+    np.testing.assert_equal(out.cpu().numpy(), np.diag([0., 1., 2.]))
+    np.testing.assert_equal(a.cpu().numpy(), np.zeros((3, 3), dtype=np.float32))
+
+  def test_copy_functional(self):
+    # without an impl this segfaults rather than fails: a regression here takes the whole run down
+    a = torch.arange(4, dtype=torch.float32, device=device)
+    out = torch.ops.aten.copy(a, torch.zeros(4, device=device))
+    np.testing.assert_equal(out.cpu().numpy(), [0., 0., 0., 0.])
+    np.testing.assert_equal(a.cpu().numpy(), [0., 1., 2., 3.])
 
 from tinygrad import Tensor
 class TestBackendHelpers(unittest.TestCase):

test/backend/test_const_folding.py

@@ -188,7 +188,6 @@ class TestTautologicalCompare(unittest.TestCase):
     np.testing.assert_equal((Tensor(True) < Tensor(False)).numpy(), False)
     np.testing.assert_equal((Tensor(True) < Tensor(True)).numpy(), False)
 
-  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support NaN comparison correctly")
   def test_a_eq_a(self):
     # self eq is always true for int or bool
     a = Tensor([1, 2, 3])

test/backend/test_dtype.py

@@ -340,6 +340,9 @@ class TestUint64DType(TestDType):
   DTYPE = dtypes.uint64
   def test_uint64_load(self):
     assert Tensor(2**64 - 1, dtype=dtypes.uint64).numpy() == 2**64 - 1
+  @unittest.skipIf(dtypes.double not in supported_dtypes, "needs float64")
+  def test_uint64_cast_double(self):
+    assert Tensor([2**32 + 1], dtype=dtypes.uint64).cast(dtypes.double).numpy() == 2**32 + 1
 
 @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
 class TestEmulatedUInt64DType(TestUint64DType):

test/backend/test_isel.py

@@ -7,7 +7,7 @@ from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops
 from tinygrad.renderer.isa import IselContext
 
 # INDEX on a register value with a constant index extracts a single element (the old GEP)
-def lane(y:UOp, i:int) -> UOp: return y.index(UOp.const(i, dtypes.int), dtype=y.dtype.scalar())
+def lane(y:UOp, i:int) -> UOp: return y.index(UOp.const(i, dtypes.int), dtype=y.dtype)
 
 @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only x86")
 class TestIselX86(unittest.TestCase):

test/backend/test_ops.py

@@ -720,10 +720,11 @@ class TestOps(unittest.TestCase):
       return torch.autograd.grad(t ** c, t)[0].item()
     for x in [-math.inf, 0, 1, math.inf]:
       for c in [-1, 0, 0.3, 1, 2]:
-        tiny_out = get_tiny_gradient(x, c)
         torch_out = get_torch_gradient(x, c)
+        # the pow backward routes through exp2/log2, whose 0/inf behavior is undefined on WEBGPU
+        if Device.DEFAULT == "WEBGPU" and not math.isfinite(torch_out): continue
+        tiny_out = get_tiny_gradient(x, c)
         if math.isnan(tiny_out):
-          if Device.DEFAULT == "WEBGPU": continue # TODO: WEBGPU issue with nan
           assert math.isnan(torch_out)
         else:
           self.assertAlmostEqual(tiny_out, torch_out, msg=f"{x}, {c}")
@@ -749,6 +750,7 @@ class TestOps(unittest.TestCase):
   def test_exp2_log2_zero_times_negative(self):
     # gallivm's exp2/log2 have "undefined behavior with infs, 0s and nans", so exp2(log2(0)*y) returns 0 instead of inf
     helper_test_op(None, lambda x,y: (x.log2()*y).exp2(), lambda x,y: (x.log2()*y).exp2(), vals=[[0.0], [-0.7]], forward_only=True)
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "pow at 0 routes through exp2/log2, whose 0/inf behavior is undefined on WEBGPU")
   def test_pow_zero_const(self):
     helper_test_op(None, lambda x: x**0.3, vals=[[0.0]])
     helper_test_op(None, lambda x: x**0.0, vals=[[0.0]])

test/backend/test_quantize_onnx.py

@@ -82,7 +82,7 @@ class TestQuantizeOnnxCPU(unittest.TestCase):
       linear = run_onnx({"input":inp})["output"].schedule_linear()
       prg = to_program(linear.src[-2].src[0], renderer=Device[Device.DEFAULT].renderer)
       daccs = [u for u in tuple(prg.src[1].src) if u.op is Ops.BUFFER and u.addrspace is AddrSpace.REG]
-      assert all(u.dtype.scalar() is dtypes.int for u in daccs)
+      assert all(u.dtype is dtypes.int for u in daccs)
 
 @unittest.skipIf(Device.DEFAULT != "DSP", "only tests for DSP")
 class TestQuantizeOnnx(unittest.TestCase):

test/backend/test_tensor.py

@@ -653,9 +653,19 @@ class TestZeroShapeTensor(unittest.TestCase):
     np.testing.assert_equal(Tensor([[1, 2]]).pad_to(2, 3).numpy(), [[1, 2, 0], [0, 0, 0]])
     np.testing.assert_equal(Tensor([[1, 2]]).pad_to(1, 3).numpy(), [[1, 2, 0]])
     np.testing.assert_equal(Tensor([[1, 2]]).pad_to(None, 3).numpy(), [[1, 2, 0]])
+    np.testing.assert_equal(Tensor([1, 2]).pad_to(4, value=2).numpy(), [1, 2, 2, 2])
+    np.testing.assert_equal(Tensor([[1, 2]]).pad_to(2, 3, value=-1).numpy(), [[1, 2, -1], [-1, -1, -1]])
+    np.testing.assert_equal(Tensor([1, 2]).pad_to(None, value=5).numpy(), [1, 2])  # no-op pad ignores the fill
     with self.assertRaises(ValueError): Tensor([1, 2]).pad_to(2, 3)
     with self.assertRaises(ValueError): Tensor([[1, 2]]).pad_to(3)
 
+  def test_max_shape(self):
+    from tinygrad import UOp
+    t = Tensor.empty(2, UOp.variable('v', 1, 32), 4)
+    self.assertEqual(t.max_shape, (2, 32, 4))
+    self.assertEqual(t.max_numel(), 2*32*4)
+    self.assertEqual(Tensor.empty(2, 3).max_shape, (2, 3))
+
   def test_shrink_into_zero(self):
     t = Tensor.rand(3, 4).realize()
     assert t.shrink((None, (2, 2))).realize().shape == (3, 0)

test/device/test_hcq2.py

@@ -10,5 +10,13 @@ class TestHCQ2(unittest.TestCase):
     with patch.object(Device[Device.DEFAULT], "has_copy_queue", False):
       np.testing.assert_equal(Tensor(np.arange(61, dtype=np.float32)).to(Device.DEFAULT).contiguous().realize().numpy(), np.arange(61))
 
+  def test_overlapping_device_tuples(self):
+    # an op on a wide device tuple followed by an op on an overlapping smaller tuple used to MMU-fault the smaller one
+    d4, d2 = tuple(f"{Device.DEFAULT}:{i}" for i in range(4)), tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
+    ref = Tensor.arange(16).contiguous().realize()
+    Tensor(ref.uop.copy_to_device(d4)).realize()
+    out = Tensor.ones(8).shard(d2, axis=0).contiguous().realize()
+    np.testing.assert_equal(out.numpy(), np.ones(8))
+
 if __name__ == "__main__":
   unittest.main()

test/external/external_uop_gc.py

@@ -44,7 +44,7 @@ def realized_matmul():
   z = y.matmul(x)
   Tensor.realize(z)
 def realized_gradient():
-  x = Tensor.eye(3)
+  x = Tensor.eye(3).clone()
   y = Tensor([[2.0,0,-2.0]])
   z = y.matmul(x).sum()
   z.backward()

test/helpers.py

@@ -86,7 +86,9 @@ def assert_jit_cache_len(fxn, expected_len):
   if linear is None or not linear.src:
     if expected_len != 0: raise KernelCountException(expected_len, 0)
     return
-  if expected_len and all(call_is_hcq(call) for call in linear.src): expected_len = 4 # HCQ2: fence + reset + merged same-queue calls + finalizer
+  if expected_len and all(call_is_hcq(call) for call in linear.src): # HCQ2: one batch submitter, or fence + reset + merged calls + finalizer
+    from tinygrad.runtime.support.hcq2 import HCQ_RUNTIME_DEV
+    expected_len = 1 if HCQ_RUNTIME_DEV.value == "CPU" else 4
   if call_is_graph(linear.src[0]):
     if len(linear.src) != 1: raise KernelCountException(1, len(linear.src))
     inner = linear.src[0].src[0].src[0]  # LINEAR UOp inside CUSTOM_FUNCTION

test/null/test_dtype_spec.py

@@ -51,10 +51,6 @@ class TestHelpers(unittest.TestCase):
     assert dtypes.is_float(dtypes.fp8e4m3)
     assert dtypes.is_float(dtypes.fp8e5m2)
 
-  @given(strat.sampled_from([d for d in DTYPES_DICT.values() if dtypes.is_float(d) or dtypes.is_int(d)]))
-  def test_scalar(self, dtype):
-    assert dtype.scalar() == dtype
-
   def test_from_py(self):
     assert dtypes.from_py(True) == dtypes.bool
     assert dtypes.from_py(Invalid) == dtypes.bool

test/null/test_tensor.py

@@ -69,7 +69,7 @@ class TestIdxUpcast(unittest.TestCase):
     if not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)):
       assert idx.op is Ops.INDEX
       idx_val = idx.src[1]
-      self.assertFalse(idx_val.overflows(idx_val.dtype.scalar()))
+      self.assertFalse(idx_val.overflows(idx_val.dtype))
 
   # use expand to generate kernel that uses large idx
   def do_op_then_assert(self, dtype: DType, dim1, dim2, dim3):

test/null/test_uop_symbolic.py

@@ -5,7 +5,8 @@ import z3
 from tinygrad.dtype import dtypes, ConstType, DType, Invalid
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad.uop.spec import spec_shared, type_verify
-from tinygrad.uop.symbolic import sym, pm_fold_cast_const, commutative, pm_simplify_valid, pm_move_where_on_load
+from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load
+from tinygrad.uop.weak import pm_cast_weak
 from tinygrad.uop.validate import uops_to_z3
 
 def check_uop_against_string(self, v:UOp, s:str):
@@ -35,7 +36,7 @@ class TestSymbolic(unittest.TestCase):
     self.assertEqual(solver.check(expr1 != expr2), z3.unsat, "simplified expression not equal to original")
 
   def helper_test_variable(self, v, n, m, s, test_z3:bool=True):
-    v_simplified = graph_rewrite(v, sym+pm_fold_cast_const, name="simplify symbolic uop")
+    v_simplified = graph_rewrite(v, sym+pm_cast_weak, name="simplify symbolic uop")
     if test_z3: self.check_equal_z3(v, v_simplified)
     nmin, nmax = v_simplified.vmin, v_simplified.vmax
     check_uop_against_string(self, v_simplified, s)

test/opt/test_gen_float4.py

@@ -10,12 +10,12 @@ from test.helpers import replace_opts
 class TestFloat4(unittest.TestCase):
   @staticmethod
   def count_float4(uops: list[UOp], n=4):
-    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype.scalar() == dtypes.float and uop.shape == (4,)]),
-            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype.scalar() == dtypes.float and uop.shape == (4,)]))
+    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype == dtypes.float and uop.shape == (4,)]),
+            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.float and uop.shape == (4,)]))
   @staticmethod
   def count_half4(uops: list[UOp]):
-    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype.scalar() == dtypes.half and uop.shape == (4,)]),
-            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype.scalar() == dtypes.half and uop.shape == (4,)]))
+    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype == dtypes.half and uop.shape == (4,)]),
+            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.half and uop.shape == (4,)]))
 
   def test_float4_basic(self):
     a = Tensor.empty(2, 8).realize()

test/unit/test_allreduce.py

@@ -64,7 +64,7 @@ class TestAllreduceCast(unittest.TestCase):
     with Context(ALLREDUCE_CAST=allreduce_cast, RING=0, SCACHE=0):
       t = Tensor.empty(4, 4, dtype=dtype).shard(ds, axis=0)
       linear = t.sum(0).linear_with_vars()[0]
-      return {si.src[1].buffer.dtype.scalar() for si in linear.src if si.src[0].op is Ops.COPY}
+      return {si.src[1].buffer.dtype for si in linear.src if si.src[0].op is Ops.COPY}
 
   def test_allreduce_cast_bf16(self):
     # with ALLREDUCE_CAST, allreduce copies stay in bfloat16 instead of promoting to float32

test/unit/test_attention.py

@@ -5,8 +5,6 @@ from tinygrad.llm.model import (
   GatedDeltaNetBlock, SSMConfig, TransformerBlock, TransformerConfig,
   apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk,
 )
-from tinygrad.llm.kernels import Linear, gated_delta_prefill
-from tinygrad.llm.gguf import ggml_data_to_tensor
 
 def apply_rope(x:Tensor, start_pos:int):
   B, H, T, Hd = x.shape
@@ -14,15 +12,6 @@ def apply_rope(x:Tensor, start_pos:int):
   freqs_cis = precompute_freqs_cis(Hd, start_pos+T)[start_pos:start_pos+T]
   return apply_rope_new(x, freqs_cis)
 
-class TestLinear(unittest.TestCase):
-  def test_recovers_packed_ggml_weight(self):
-    for ggml_type,packed_size,words in ((13, 176, 44), (14, 210, 210), (23, 136, 34)):
-      packed = Tensor.empty(packed_size+4, dtype=dtypes.uint8, device="CPU")[4:]
-      decoded = ggml_data_to_tensor(packed, 256, ggml_type).reshape(1, 256)
-      linear = Linear(256, 1, bias=False)
-      linear.set_quantized(decoded)
-      self.assertEqual((linear.ggml_type, linear.weight.numel()), (ggml_type, words))
-
 class TestAttention(unittest.TestCase):
   def test_apply_rope(self):
     x = Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32)
@@ -52,22 +41,6 @@ class TestAttention(unittest.TestCase):
     np.testing.assert_allclose(block.cache_kv[0, :, :, :seqlen, :].numpy(), expected.numpy(), rtol=1e-5, atol=1e-5)
 
 class TestGatedDeltaNetBlock(unittest.TestCase):
-  def test_gated_delta_rectangular_state_and_row_decay(self):
-    rng = np.random.default_rng(42)
-    q, k = (rng.normal(size=(1, 1, 3, 32)).astype(np.float32) for _ in range(2))
-    v, beta = rng.normal(size=(1, 1, 3, 4)).astype(np.float32), rng.uniform(size=(1, 1, 3)).astype(np.float32)
-    alpha, initial = rng.uniform(0.8, 1, size=(1, 1, 3, 4)).astype(np.float32), rng.normal(size=(1, 1, 4, 32)).astype(np.float32)
-    expected_state, expected_out = initial.copy(), np.empty_like(v)
-    for t in range(3):
-      previous, av = expected_state.copy(), alpha[:, :, t, :, None]
-      delta = (v[:, :, t] - (previous*k[:, :, t, None]).sum(-1)*alpha[:, :, t]) * beta[:, :, t, None]
-      expected_state = previous*av + delta[..., None]*k[:, :, t, None, :]
-      expected_out[:, :, t] = (previous*q[:, :, t, None]).sum(-1)*alpha[:, :, t] + delta*(q[:, :, t]*k[:, :, t]).sum(-1)
-    state = Tensor(initial).contiguous().realize()
-    out = gated_delta_prefill(Tensor(q), Tensor(k), Tensor(v), Tensor(beta), Tensor(alpha), state).realize()
-    np.testing.assert_allclose(out.numpy(), expected_out, rtol=1e-4, atol=1e-4)
-    np.testing.assert_allclose(state.numpy(), expected_state, rtol=1e-4, atol=1e-4)
-
   def _tensor_linspace(self, start:float, stop:float, shape:tuple[int, ...]) -> Tensor:
     return Tensor.linspace(start, stop, int(np.prod(shape)), dtype=dtypes.float32).reshape(*shape)
 

test/unit/test_dtype_weak.py

@@ -76,6 +76,11 @@ class TestWeakPromotion(unittest.TestCase):
       committed = graph_rewrite((UOp.const(1).cast(dtypes.int32) + UOp.const(1.0)).cast(dtypes.float32), pm_lower_index_dtype, ctx={})
     self.assertEqual([u.dtype for u in committed.toposort() if u.op is Ops.ADD], [dtypes.float32])
 
+  def test_div_sub_operand_kept_weak(self):
+    a = Tensor.empty(4, dtype=dtypes.float32)
+    for t in (a / 1, a - 0):
+      self.assertEqual(t.uop.src[1].dtype, dtypes.weakfloat)
+
   def test_cast_weak_expression_commits_at_cast_floor(self):
     # the floor never narrows: a cast BELOW the default does not pull the compute width down with it
     with Context(DEFAULT_FLOAT=dtypes.float32):
@@ -88,6 +93,13 @@ class TestWeakPromotion(unittest.TestCase):
       out = Tensor(1.0, dtype=dtypes.float32, device="CPU") / denom
       self.assertAlmostEqual(out.item(), 1 / (70000 + 1e-5), places=10)
 
+  def test_stacked_weak_casts_convert_each_kind(self):
+    # each weak cast is a kind conversion: weakint truncates before weakfloat re-lifts (neither is only a marker)
+    x = Tensor([2.5, -3.7], dtype=dtypes.float32, device="CPU")
+    stacked = x.cast(dtypes.weakint).cast(dtypes.weakfloat)
+    self.assertIs(stacked.dtype, dtypes.weakfloat)
+    self.assertEqual(stacked.tolist(), [2.0, -3.0])
+
   def test_uop_scalar_const_lifts_kind(self):
     for dtype, value, out_dtype, const_dtype in ((dtypes.weakint, 1, dtypes.weakint, dtypes.weakint),
                                                  (dtypes.int32, 1, dtypes.int32, dtypes.weakint),

test/unit/test_gradient.py

@@ -51,6 +51,10 @@ class TestTensorGradient(unittest.TestCase):
     with self.assertRaises(RuntimeError): x.sum().gradient(x)
     with self.assertRaises(RuntimeError): x.float().sum().gradient(x)
 
+  def test_const_target_raise(self):
+    t = Tensor(2.0)
+    with self.assertRaises(RuntimeError): (t * 2.0).gradient(t)
+
   def test_copy_to_device_gradient(self):
     t = Tensor([1.0, 2, 3]).realize()
     t.to("CPU:1").square().sum().backward()
@@ -100,7 +104,7 @@ class TestTensorGradient(unittest.TestCase):
 
   def test_implicit_broadcast_where_gradient(self):
     # WHERE with a bare ()-shape branch: the scalar's gradient counts the positions where it is selected
-    cond, x, w = Tensor([True, False, True]), Tensor([1.0, 2.0, 3.0]), Tensor(4.0)
+    cond, x, w = Tensor([True, False, True]), Tensor([1.0, 2.0, 3.0]), Tensor(4.0, dtype=dtypes.float32)
     dw = Tensor(cond.uop.alu(Ops.WHERE, x.uop, w.uop)).sum().gradient(w)[0]
     self.assertEqual(dw.shape, ())
     self.assertEqual(dw.item(), 1.0)
@@ -109,7 +113,7 @@ class TestTensorGradient(unittest.TestCase):
 
   def test_implicit_broadcast_alu_gradient(self):
     # MUL with a bare ()-shape src, no EXPAND in the graph
-    x, w = Tensor([1.0, 2.0, 3.0]), Tensor(2.0)
+    x, w = Tensor([1.0, 2.0, 3.0]), Tensor(2.0, dtype=dtypes.float32)
     m = x.uop.alu(Ops.MUL, w.uop)
     self.assertIs(m.src[1], w.uop)
     dw = Tensor(m).sum().gradient(w)[0]
@@ -118,7 +122,7 @@ class TestTensorGradient(unittest.TestCase):
 
   def test_implicit_broadcast_intermediate_accumulation(self):
     # s is used directly and through an implicit broadcast edge, each edge's gradient reduces to s's shape before they sum
-    x, p = Tensor([1.0, 2.0, 3.0]), Tensor(0.5)
+    x, p = Tensor([1.0, 2.0, 3.0]), Tensor(0.5, dtype=dtypes.float32)
     s = p.sin()
     z = Tensor(x.uop.alu(Ops.MUL, s.uop)).sum() + s
     dp = z.gradient(p)[0]

test/unit/test_llm_amd.py

@@ -1,91 +0,0 @@
-import unittest
-import numpy as np
-from tinygrad import Tensor, UOp, dtypes, nn
-from tinygrad.llm.kernels import Linear, amd_custom_kernels_supported
-from tinygrad.llm.kernels.amd import q8_quantize, quantized_attention
-from tinygrad.llm.gguf import ggml_data_to_tensor
-
-class TestQ8Quantize(unittest.TestCase):
-  def test_word_quant_weights_use_typed_buffer_view(self):
-    for ggml_type, type_size in ((13, 176), (23, 136)):
-      with self.subTest(ggml_type=ggml_type):
-        raw = Tensor(np.zeros(type_size + 4, dtype=np.uint8), device="CPU").contiguous().realize()[4:]
-        decoded = ggml_data_to_tensor(raw, 256, ggml_type).reshape(1, 256)
-        linear = Linear(256, 1, bias=False)
-        linear.set_quantized(decoded)
-        self.assertEqual(linear.ggml_type, ggml_type)
-        self.assertEqual(linear.weight.dtype, dtypes.uint32)
-        self.assertEqual(linear.weight.nbytes(), type_size)
-        self.assertEqual(linear.weight.uop.buf_uop.buffer.offset, 4)
-
-  def test_values_and_scales(self):
-    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
-    x = np.linspace(-3.1, 2.7, 64, dtype=np.float32).reshape(2, 32)
-    quant, scale = q8_quantize(Tensor(x), 2, 32)
-    scale_np = np.maximum(np.max(np.abs(x), axis=-1, keepdims=True) / 127, 1e-8)
-    expected = np.clip(np.rint(x / scale_np), -127, 127).astype(np.int8)
-    np.testing.assert_array_equal(quant.bitcast(dtypes.int8).reshape(2, 32).numpy(), expected)
-    np.testing.assert_allclose(scale.numpy(), scale_np, rtol=1e-6)
-
-  def test_q6_linear_compiles(self):
-    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
-    rng = np.random.default_rng(42)
-    packed = rng.integers(0, 256, 210, dtype=np.uint8)
-    packed[-2:] = np.array([0.01], dtype=np.float16).view(np.uint8)
-    raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:]
-    decoded = ggml_data_to_tensor(raw, 256, 14).reshape(1, 256)
-    linear = Linear(256, 1, bias=False)
-    nn.state.load_state_dict(linear, {"weight":decoded}, verbose=False, realize=False)
-    self.assertTrue(np.isfinite(linear(Tensor.randn(1, 256)).realize().item()))
-    self.assertEqual(linear.weight.uop.buf_uop.buffer.offset, 4)
-
-  def test_q6_linear_multiple_tokens(self):
-    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
-    rng = np.random.default_rng(42)
-    in_features, blocks = 2048, 16*2048//256
-    packed = rng.integers(0, 256, blocks*210, dtype=np.uint8)
-    for i in range(blocks): packed[i*210+208:i*210+210] = np.array([0.01], dtype=np.float16).view(np.uint8)
-    raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:]
-    decoded = ggml_data_to_tensor(raw, 16*in_features, 14).reshape(16, in_features)
-    weight = decoded.numpy()
-    linear = Linear(in_features, 16, bias=False)
-    nn.state.load_state_dict(linear, {"weight":decoded}, verbose=False, realize=False)
-    x = rng.normal(size=(3, in_features)).astype(np.float32)
-    scale = np.maximum(np.abs(x).reshape(3, in_features//32, 32).max(-1, keepdims=True) / 127, 1e-8)
-    xq = np.clip(np.rint(x.reshape(3, in_features//32, 32) / scale), -127, 127) * scale
-    np.testing.assert_allclose(linear(Tensor(x)).numpy(), xq.reshape(3, in_features) @ weight.T, rtol=2e-3, atol=2e-2)
-    self.assertEqual(linear.ggml_type, 14)
-
-    generic = Linear(in_features, 16, bias=False)
-    nn.state.load_state_dict(generic, {"weight":decoded}, verbose=False, realize=False)
-    generic(Tensor.randn(4, in_features)[:UOp.variable("tokens", 1, 4).bind(2)])
-    self.assertFalse(generic.use_custom_quant)
-    self.assertIsNone(generic.ggml_type)
-
-  def test_attention_uses_physical_cache_length(self):
-    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
-    q, k, v = Tensor.zeros(1, 2, 1, 32), Tensor.randn(1, 1, 1, 32), Tensor.randn(1, 1, 1, 32)
-    cache = Tensor.empty(2, 1, 1, 256, 32, dtype=dtypes.int8).contiguous()
-    scale = Tensor.empty(2, 1, 1, 256, dtype=dtypes.float16).contiguous()
-    out = quantized_attention(q, Tensor.stack(k, v), cache, scale, 0).realize()
-    np.testing.assert_allclose(out.numpy(), v.expand(1, 2, 1, 32).numpy(), rtol=2e-2, atol=2e-2)
-
-  def test_prefill_attention_unaligned_start(self):
-    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
-    rng = np.random.default_rng(42)
-    start_pos = 1718
-    q = Tensor.zeros(1, 8, 32, 128)
-    old_kv = rng.normal(size=(2, 1, 1, start_pos, 128)).astype(np.float32)
-    new_kv = rng.normal(size=(2, 1, 1, 32, 128)).astype(np.float32)
-    cache = Tensor.empty(2, 1, 1, 2048, 128, dtype=dtypes.int8).contiguous()
-    scale = Tensor.zeros(2, 1, 1, 2048, dtype=dtypes.float16).contiguous()
-    old_scale = np.maximum(np.max(np.abs(old_kv), axis=-1, keepdims=True) / 127, 1e-8).astype(np.float16)
-    Tensor.realize(cache[:, :, :, :start_pos].assign(Tensor(np.rint(old_kv / old_scale).astype(np.int8))),
-                   scale[:, :, :, :start_pos].assign(Tensor(old_scale.squeeze(-1))))
-    out = quantized_attention(q, Tensor(new_kv), cache, scale, UOp.variable("start_pos", 0, 2047).bind(start_pos)).realize()
-    values = cache[1, 0, 0, :start_pos+32].numpy().astype(np.float32) * \
-      scale[1, 0, 0, :start_pos+32].numpy().astype(np.float32)[:, None]
-    expected = np.stack([values[:start_pos+i+1].mean(0) for i in range(32)])[None, None].repeat(8, axis=1)
-    np.testing.assert_allclose(out.numpy(), expected, rtol=2e-3, atol=2e-3)
-
-if __name__ == "__main__": unittest.main()

test/unit/test_llm_server.py

@@ -42,7 +42,8 @@ class TestTransformerGenerate(unittest.TestCase):
       return Tensor([[42]])
     with patch.object(Transformer, '__call__', mock_call):
       next(model.generate([1, 2, 3, 4, 5, 42, 10]))
-    self.assertEqual(calls, [((1, 1), V_START_POS.bind(5)), ((1, 1), V_START_POS.bind(6))])
+    # recurrent blocks prefill chunks like attention blocks: the 2 new tokens go through one chunked call
+    self.assertEqual(calls, [((1, V_TOKS.bind(2)), V_START_POS.bind(5))])
 
   def test_recurrent_divergent_prompt_restarts(self):
     model, calls = Transformer(TEST_CONFIG), []

test/unit/test_schedule_cache.py

@@ -4,8 +4,13 @@ from tinygrad import Tensor, Variable, UOp, function
 from tinygrad.uop.ops import KernelInfo
 from tinygrad.schedule import schedule_cache
 
-def custom_set0_kernel(A:UOp, num:int) -> UOp:
-  return A[0].set(num).sink(arg=KernelInfo(f"custom_set0_{num}"))
+def custom_add_kernel(A:UOp, B:UOp, num:int=0) -> UOp:
+  return A[0].set(B[0] + num).sink(arg=KernelInfo(f"custom_add_{num}"))
+
+def custom_add_backward(grad_output:UOp, _) -> tuple[None, UOp]:
+  grad = Tensor.invalids(*grad_output.shape, dtype=grad_output.dtype, device=grad_output.device)
+  grad = Tensor.custom_kernel(grad, Tensor(grad_output, device=grad_output.device), fxn=functools.partial(custom_add_kernel, num=0))[0]
+  return None, grad.uop
 
 class TestScheduleCache(unittest.TestCase):
   def test_bound_variable_reuses_cache(self):
@@ -25,27 +30,27 @@ class TestScheduleCache(unittest.TestCase):
 
   def test_custom_kernel(self):
     for i in range(4):
-      a = Tensor.empty(1)
-      a = Tensor.custom_kernel(a, fxn=functools.partial(custom_set0_kernel, num=i))[0]
+      a, b = Tensor.empty(1), Tensor.ones(1)
+      a = Tensor.custom_kernel(a, b, fxn=functools.partial(custom_add_kernel, num=i))[0]
       a.realize()
-      self.assertEqual(a.item(), i)
+      self.assertEqual(a.item(), i+1)
 
   def test_same_custom_function_reuses_cache(self):
     schedule_cache.clear()
-    fxn = functools.partial(custom_set0_kernel, num=10)
+    fxn = functools.partial(custom_add_kernel, num=10)
 
     # first run
-    a = Tensor.empty(1)
-    a = Tensor.custom_kernel(a, fxn=fxn)[0]
+    a, x = Tensor.empty(1), Tensor.ones(1)
+    a = Tensor.custom_kernel(a, x, fxn=fxn)[0]
     a.realize()
-    self.assertEqual(a.item(), 10)
+    self.assertEqual(a.item(), 11)
     cache_size_after_first = len(schedule_cache)
 
     # second run with same function should reuse cache
-    b = Tensor.empty(1)
-    b = Tensor.custom_kernel(b, fxn=fxn)[0]
+    b, x = Tensor.empty(1), Tensor.ones(1)
+    b = Tensor.custom_kernel(b, x, fxn=fxn)[0]
     b.realize()
-    self.assertEqual(b.item(), 10)
+    self.assertEqual(b.item(), 11)
     self.assertEqual(len(schedule_cache), cache_size_after_first)
 
   def test_simple(self):
@@ -65,23 +70,28 @@ class TestScheduleCache(unittest.TestCase):
       print(num)
     self.assertEqual(len(schedule_cache), start_len_schedule_cache)
 
-  @unittest.expectedFailure
   def test_simple_precompile(self):
-    @function(precompile=True)
+    @function(precompile=True, precompile_backward=True)
     def f(x:Tensor) -> Tensor:
       out = Tensor.invalids(*x.shape, dtype=x.dtype, device=x.device)
-      out = Tensor.custom_kernel(out, fxn=functools.partial(custom_set0_kernel, num=10))[0]
+      out = Tensor.custom_kernel(out, x, fxn=functools.partial(custom_add_kernel, num=10), grad_fxn=custom_add_backward)[0]
       return out + x
 
     # warmup
     x = Tensor.ones(1).realize()
-    _ = f(x).realize()
+    out = f(x)
+    out.backward(x)
+    self.assertEqual(out.item(), 12)
+    self.assertEqual(x.grad.item(), 2)
 
     # use the cache next time function is called
     start_len_schedule_cache = len(schedule_cache)
     for _ in range(3):
-      num = f(x).realize()
-      self.assertEqual(num.item(), 11)
+      x = Tensor.ones(1).realize()
+      out = f(x)
+      out.backward(x)
+      self.assertEqual(out.item(), 12)
+      self.assertEqual(x.grad.item(), 2)
     self.assertEqual(len(schedule_cache), start_len_schedule_cache)
 
 if __name__ == "__main__":

tinygrad/codegen/__init__.py

@@ -12,7 +12,7 @@ from tinygrad.dtype import dtypes, AddrSpace
 
 # import all pattern matchers here
 from tinygrad.codegen.gpudims import pm_add_gpudims
-from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_fold_cast_const, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid
+from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid
 from tinygrad.uop.movement import mop_cleanup
 from tinygrad.codegen.decomp.dtype import pm_dtype_decomps
 from tinygrad.codegen.decomp.op import get_late_rewrite_patterns, get_simplifying_rewrite_patterns
@@ -346,7 +346,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
 
   # lower index dtype
   # NOTE: we need indexing_simplify to remove the cast to long using the Invalid
-  sink = graph_rewrite(sink, symbolic_simple+pm_fold_cast_const+pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes")
+  sink = graph_rewrite(sink, symbolic_simple+pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes")
 
   # final symbolic before decomp
   sink = graph_rewrite(sink, symbolic, name="final symbolic")

tinygrad/codegen/decomp/dtype.py

@@ -33,7 +33,8 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
       return (lo:=uops[0].cast(l2i_dt[dt])), (uops[0] / 2**32).cast(l2i_dt[dt]) - ((uops[0] < 0) & lo.ne(0))
     case Ops.CAST if dt in dtypes.floats:
       small = (a1.eq(0) & (a0 >= 0)) | (a1.eq(-1) & (a0 < 0))
-      return small.where(a0.cast(dt), ((a1.cast(dtypes.float32) * (2**32)) + a0.bitcast(dtypes.uint).cast(dtypes.float32)).cast(dt))
+      cdt = dt if dt == dtypes.float64 else dtypes.float32
+      return small.where(a0.cast(dt), ((a1.cast(cdt) * (2**32)) + a0.bitcast(dtypes.uint).cast(cdt)).cast(dt))
     case Ops.CAST: return a0.bitcast(dtypes.uint).cast(dt)
     case Ops.BITCAST: return a0.bitcast(dt), a1.bitcast(dt)
     case Ops.SHL:

tinygrad/device.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 from dataclasses import dataclass, replace
 from collections import defaultdict
-from typing import Any, Generic, TypeVar, Iterator, Generator, Self, TYPE_CHECKING
+from typing import Any, Callable, Generic, TypeVar, Iterator, Generator, Self, TYPE_CHECKING
 import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal
 from tinygrad.helpers import LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, PROFILE, temp, colored
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing
@@ -103,7 +103,7 @@ class Buffer:
   def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None, options:BufferSpec|None=None,
                initial_value:bytes|pickle.PickleBuffer|None=None, uop_refcount=0, base:Buffer|None=None, offset:int=0, preallocate=False):
     assert isinstance(dtype, DType)
-    self.device, self.size, self.dtype, self.options, self.offset, self.allocated_views = device, size, dtype, options, offset, 0
+    self.device, self.size, self.dtype, self.options, self.offset, self.allocated_views = Device.canonicalize(device), size, dtype, options, offset, 0
     self._bufs: dict[str, Any] = {}
     if base is None:
       assert offset == 0, "base buffers can't have offset"
@@ -116,7 +116,7 @@ class Buffer:
         if isinstance(initial_value, pickle.PickleBuffer): initial_value.release()
     else:
       assert base._base is None, "base can't have a base"
-      assert device == base.device, "base must have the same device"
+      assert self.device == base.device, "base must have the same device"
       self._base = base
     if preallocate: self.allocate()
   @property
@@ -133,7 +133,7 @@ class Buffer:
   # check if the underlying buffer is allocated, possibly from the base object
   def is_allocated(self) -> bool: return self.base.is_allocated() if self._base is not None else self.device in self._bufs
   def get_buf(self, device: str) -> Any:
-    if device not in self._bufs:
+    if (device:=Device.canonicalize(device)) not in self._bufs:
       allocator = Device[device].allocator
       if device == self.device: self.ensure_allocated()
       elif self._base is not None: self._bufs[device] = allocator._offset(self._base.get_buf(device), self.nbytes, self.offset)
@@ -331,17 +331,18 @@ class Program(Generic[DeviceType]):
                wait=False) -> float|None: pass
 
 class Compiled:
+  ifaces:list[Callable] = []
   profile_events:list[ProfileEvent] = [ProfileDeviceEvent("CPU")] # NOTE: CPU is the default device.
 
+  has_copy_queue:bool = True
+
   pm_lower:Any = None
   pm_bufferize:Any = None
 
-  has_copy_queue:bool = True
-
   def __init__(self, device:str, allocator:Allocator, renderers:list[type[Renderer]], runtime:type[Program[Self]]|None, graph=None, arch=None):
     from tinygrad.renderer import Renderer
     self.device, self.allocator, self.runtime_t, self.graph, self.renderers = device, allocator, runtime, graph, renderers or [Renderer]
-    self.arch = arch
+    self.device_id, self.arch = (int(idx) if ":" in device and (idx:=device.split(":")[1]).isdigit() else 0), arch
     self.cached_renderer:dict[Any, Renderer] = {}
 
   @property
@@ -364,11 +365,21 @@ class Compiled:
     return select_first_inited(select_by_name(self.renderers, self._renderer_name, t.renderer, f"{self.device} has no renderer {t.renderer!r}"),
                                f"No renderer for {self.device} is available", self.cached_renderer, t)
 
+  def _select_iface(self, device:str):
+    self.device_id = int(device.split(":")[1]) if ":" in device else 0
+    assert (v:=getenv(k:=f'{type(self).__name__[:-6].upper()}_IFACE', "")) == "",  \
+      f"{k}={v} is deprecated, use DEV={replace(DEV.target(type(self).__name__[:-6]), interface=v)} instead"
+    t = DEV.target(dev:=type(self).__name__[:-6])
+    filtered = select_by_name(self.ifaces, lambda i: i.__name__[:-5], t.interface, f"{dev} has no interface {t.interface!r}")
+    filtered = [i for i in filtered if t.interface.startswith("MOCK") or not i.__name__[:-5].startswith("MOCK")] # never fallback to mock ifaces
+    return select_first_inited([functools.partial(iface, self, self.device_id) for iface in filtered],
+                               f"No interface for {dev}:{self.device_id} is available")
+
   def count(self) -> int:
     """
     Returns the number of physical accelerators available to the runtime.
     """
-    return 1
+    return self.iface.count if hasattr(self, 'iface') else 1
 
   def synchronize(self):
     """
@@ -386,7 +397,7 @@ class Compiled:
     """
     Called at the end of process lifetime to allow the device to finalize.
     """
-    # override this in your device implementation
+    if hasattr(self, 'iface') and hasattr(self.iface, 'device_fini'): self.iface.device_fini()
 
 if PROFILE:
   @atexit.register
@@ -408,7 +419,7 @@ def enumerate_devices_str() -> Generator[str, None, None]:
     ren_results, iface_results = [], []
     try:
       d = Device[device]
-      for iface in [i for i in getattr(d, 'ifaces', []) if not i.__name__.startswith("MOCK")]:
+      for iface in [i for i in d.ifaces if not i.__name__.startswith("MOCK")]:
         try:
           name = iface.__name__[:-5]
           default_text, count = ("(default)", d.count()) if type(d.iface) is iface else (f"(DEV={name}+{device} to make default)", iface(d, 0).count) # type: ignore

tinygrad/dtype.py

@@ -66,7 +66,6 @@ class DType(metaclass=DTypeMetaClass):
   def __reduce__(self): return type(self), tuple(getattr(self, f.name) for f in fields(self))
   def __repr__(self): return f"dtypes.{INVERSE_DTYPES_DICT[self.name]}"
   def __lt__(self, o:DType): return (self.priority, self.bitsize, self.name, self.fmt) < (o.priority, o.bitsize, o.name, o.fmt)
-  def scalar(self) -> DType: return self
   @functools.cached_property
   def min(self):
     if dtypes.is_int(self): return 0 if dtypes.is_unsigned(self) else -2**(self.bitsize-1)

tinygrad/engine/jit.py

@@ -269,10 +269,14 @@ class _TinyJit(Generic[ReturnType]):
         big_linear, onetime_linear = prune_linear(big_linear, set(input_buf_uops))
         if DEBUG >= 1: print(f"pruned from {len(big_linear.src) + len(onetime_linear.src)} -> {len(big_linear.src)} kernels")
         run_linear(onetime_linear, var_vals)
+        del onetime_linear
 
       # hold all buffers reachable from live Tensors (e.g. lazy .grad created during capture), the memory planner can't suballocate those
       held_bufs = set(buffers) | {u for tref in list(all_tensors) if (t:=tref()) is not None for u in t.uop.toposort() if u.op is Ops.BUFFER}
       linear = jit_lower(big_linear, held_bufs, input_buf_uops)
+      # drop the pre-planning graph: it keeps the whole capture-time working set allocated (big_linear) or referenced (held_bufs).
+      # the planned linear only uses the arena/held buffers, so the intermediates must be freed before linking and first exec
+      del big_linear, held_bufs
       self.captured = CapturedJit(ret, linear, names, expected_input_info)
       ret = self.captured(input_buf_uops, var_vals)
     elif self.cnt >= 2:

tinygrad/engine/realize.py

@@ -3,9 +3,10 @@ from typing import cast, Iterator, Any, Sequence
 import time, random, itertools, math, contextlib, weakref, array
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, prod, flatten, Context, getenv, to_tuple
-from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, wait_cond
+from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite
 from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
+from tinygrad.dtype import dtypes
 from tinygrad.renderer import Estimates
 from tinygrad.codegen import to_program
 from tinygrad.codegen.opt.postrange import args_from_ast
@@ -13,7 +14,9 @@ from tinygrad.codegen.opt.postrange import args_from_ast
 # **************** Helpers ****************
 
 def get_call_arg_uops(call:UOp) -> tuple[UOp, ...]: return tuple(s for s in call.src[1:] if not s.is_bound_var)
-
+def get_call_var_uops(call:UOp, prg:UOp) -> list[UOp]:
+  bound = {s.src[0].expr: s.src[1].src[1] for s in call.src[1:] if s.is_bound_var}
+  return [bound.get(v.expr, v) for v in prg.arg.vars]
 def get_call_outs_ins(call:UOp) -> tuple[tuple[int, ...], tuple[int, ...]]:
   ast = call.src[0]
   if ast.op is Ops.PROGRAM: return tuple(ast.arg.outs), tuple(ast.arg.ins)
@@ -166,9 +169,10 @@ def exec_copy(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
 
 def exec_kernel(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
   et = None
-  for device, (bufs, device_vars) in zip(to_tuple(call.src[1].device), unwrap_multi(call, resolve_params(call, ctx.input_uops))):
+  resolved = resolve_params(call, ctx.input_uops)
+  for device, (bufs, device_vars) in zip(to_tuple(call.src[1].device), unwrap_multi(call, [resolved[i] for i in ast.arg.globals])):
     var_vals = {**ctx.var_vals, **device_vars}
-    prg_bufs = [bufs[i].ensure_allocated() for i in ast.arg.globals]
+    prg_bufs = [b.ensure_allocated() for b in bufs]
     rt = get_runtime(device, ast, cache=ctx.cache)
     global_size, local_size = ast.arg.launch_dims(var_vals)
     with track_stats(ctx, call, device, prg_bufs, var_vals) as tm:
@@ -200,28 +204,26 @@ def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
   return t[0]
 
 def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
-  if (info:=call.arg.aux).inputs is not None:
-    bufs = [_resolve(ctx.input_uops[i], ctx.input_uops).buffer for i in call.arg.aux.input_idxs]
-    table = call.src[1+info.inputs].buffer
-    for j,dev in enumerate(call.arg.aux.device):
-      addrs = array.array('Q', [(b.bufs[j] if isinstance(b, MultiBuffer) else b).get_buf(dev).va_addr for b in bufs])
-      mv = (table.bufs[j] if isinstance(table, MultiBuffer) else table).ensure_allocated()._buf.cpu_view().view(fmt='Q')
-      wait_cond(lambda: mv[0], value=0, timeout_ms=ctx.timeout or getenv("HCQDEV_WAIT_TIMEOUT_MS", 30000), msg=f"{dev} hang detected")
-      mv[:len(addrs)] = addrs
+  dev = cast(Any, Device[(info:= call.arg.aux).device[0]])
+  addrs = [(b.bufs[j] if isinstance(b:=_resolve(ctx.input_uops[k], ctx.input_uops).buffer, MultiBuffer) else b).get_buf(dev_name).va_addr
+           for devs, idxs in info.input_idxs for j, dev_name in enumerate(devs) for k in idxs]
+  dev.rt_buffer._buf.cpu_view().view(offset=(base:=dev.rt_allocator.alloc(len(addrs) * 8)), fmt='Q')[:len(addrs)] = array.array('Q', addrs)
 
-  exec_kernel(replace(ctx, update_stats=DEBUG>=3), call, ast)
+  tables = [UOp.from_buffer(dev.rt_buffer.view(len(idxs), dtypes.uint64, base + j*len(idxs)*8), HCQ_RUNTIME_DEV.value)
+            for devs, idxs in info.input_idxs for j in range(len(devs))]
+  if info.inputs is not None: call = call.substitute({call.src[1+info.inputs]: UOp.mstack(*tables)})
+  exec_kernel(replace(ctx, update_stats=DEBUG>=3, var_vals={**ctx.var_vals, "hcq_inputs_ptr": dev.rt_buffer._buf.va_addr + base}), call, ast)
 
   tms = []
-  for devices,name,estimates,prof in info.kernels:
+  for devices, stat_call, prof in info.kernels:
     for device in devices:
       tm = None
       if prof:
-        (d:=cast(Any, Device[device])).prof_ents[prof[0]] = ProfileGraphEntry(device, name, *prof)
+        (d:=cast(Any, Device[device])).prof_ents[prof[0]] = ProfileGraphEntry(device, stat_call.arg.name, *prof)
         if ctx.wait:
           d.synchronize(timeout=ctx.timeout)
           st, en = (d.signal(x)._buf.cpu_view().view(fmt='Q')[0] for x in prof)
           tms.append(tm:=float(en-st)/d.timestamp_divider/1e6)
-      stat_call = call.replace(arg=replace(call.arg, name=name, aux=replace(info, estimates=estimates, kernels=())))
       with track_stats(ctx, stat_call, device, [], ctx.var_vals) as et: et[0] = tm
   return max(tms) if tms else None
 
@@ -262,7 +264,7 @@ pm_exec = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="validate", name="ast"),), name="call", allow_any_len=True), exec_validate),
 ])
 
-if getenv("HCQ2"): from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link # noqa: E402 # down here, hcq2 imports the helpers above
+if getenv("HCQ2"): from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link, HCQ_RUNTIME_DEV # noqa: E402 # down here, hcq2 imports realize
 
 def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, profile:bool|None=None) -> UOp:
   if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)

tinygrad/function.py

@@ -1,4 +1,5 @@
 import functools, time
+from dataclasses import replace
 from typing import Generic, TypeVar, Callable, cast, overload
 from tinygrad.helpers import Context, dedup, getenv, DEBUG
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, PatternMatcher, UPat
@@ -23,6 +24,10 @@ def invalid_outputs(uret:UOp) -> set[UOp]:
   return {u.src[0].buf_uop for u in uret.backward_slice_with_self
           if u.op is Ops.STORE and u.src[1].base.is_invalid and not u.src[0].buf_uop.is_realized}
 
+def renumber_invalid_outputs(uret:UOp) -> UOp:
+  return uret.substitute({b:b.replace(arg=replace(b.arg, slot=i))
+                          for i,b in enumerate(x for x in uret.toposort(enter_calls=False) if x in invalid_outputs(uret))})
+
 ReturnType = TypeVar('ReturnType')
 class _function(Generic[ReturnType]):
   depth = 0
@@ -65,6 +70,7 @@ class _function(Generic[ReturnType]):
     # the BUFFERs that are left are the implicit inputs
     num_explicit = len(call_uops)
     uret = graph_rewrite(uret, pm_ctx, (call_uops, invalid_outputs(uret)), bottom_up=True, name="get_implicit_inputs")
+    uret = renumber_invalid_outputs(uret)
     name = getattr(self.fxn, '__qualname__', None) or type(self.fxn).__qualname__
     if not self.allow_implicit:
       implicit_buffers = [x for x in call_uops[num_explicit:] if x.op is Ops.BUFFER]

tinygrad/llm/kernels/__init__.py

@@ -1,92 +0,0 @@
-import functools
-from typing import cast
-from tinygrad import Tensor, UOp, nn, dtypes, Device, Context
-from tinygrad.device import Buffer
-from tinygrad.dtype import AddrSpace
-from tinygrad.helpers import prod
-from tinygrad.uop.ops import AxisType, KernelInfo, Ops
-
-def kernel_var(x:UOp) -> UOp:
-  # a Variable is a 0-d ALU BUFFER in the tensor graph; inside kernels it takes the ALU PARAM form (same name keeps the value binding)
-  return x.substitute({v: UOp.variable(v.expr, v.vmin, v.vmax, dtype=v.dtype, multiple_of=v.arg.multiple_of, param=True)
-                       for v in x.toposort() if v.is_variable})
-
-def amd_custom_kernels_supported(device:str|tuple[str, ...]|None) -> bool:
-  # the custom kernels are tuned for RDNA3 (gfx11): the WMMA register layouts don't match gfx12 (RDNA4)
-  # or CDNA (MFMA-only, wave64), and the dp4a builtins and 32-lane wave ops aren't portable either.
-  if isinstance(device, tuple): device = device[0]
-  if device is None or device.split(":")[0] != "AMD": return False
-  # Device[...] trips ALLOW_DEVICE_USAGE=0 in function contexts, the device is always open here anyway
-  with Context(ALLOW_DEVICE_USAGE=1):
-    return (t:=getattr(Device[device], "target", None)) is not None and t[0] == 11
-
-class Linear(nn.Linear):
-  ggml_type:int|None = None
-  def __init__(self, in_features:int, out_features:int, bias=True):
-    super().__init__(in_features, out_features, bias)
-    self.in_features, self.out_features = in_features, out_features
-    self.use_custom_quant = True
-  def set_quantized(self, decoded:Tensor):
-    packed_sizes = {decoded.numel() // 256 * type_size:typ for typ,type_size in ((13, 176), (14, 210), (23, 136))}
-    raw = next((u for u in decoded.uop.toposort() if u.op is Ops.SHRINK and u.dtype == dtypes.uint8 and prod(u.shape) in packed_sizes), None)
-    if raw is None: return
-    raw_offset = raw.contiguous_view_offset()
-    assert raw_offset is not None and raw_offset % 4 == 0 and raw.buf_uop.dtype == dtypes.uint8
-    self.ggml_type = packed_sizes[prod(raw.shape)]
-    # Q5_K and IQ4_XS kernels consume words. Store a typed buffer view directly: a lazy BITCAST is decomposed into
-    # byte-combining ALU before custom-kernel scheduling and would copy the entire packed weight on every JIT graph.
-    packed_dtype = dtypes.uint8 if self.ggml_type == 14 else dtypes.uint32
-    self.weight = Tensor(UOp.from_buffer(cast(Buffer, raw.buf_uop.buffer).view(raw.max_numel() * raw.dtype.itemsize // packed_dtype.itemsize,
-                                                                              packed_dtype, raw_offset)))
-  def __call__(self, x:Tensor) -> Tensor:
-    static = isinstance(x.numel(), int)
-    supported = self.use_custom_quant and amd_custom_kernels_supported(self.weight.device)
-    if self.ggml_type is None and not static: supported = self.use_custom_quant = False
-    if self.ggml_type is None and supported: self.set_quantized(self.weight)
-    if self.ggml_type in (13, 14, 23) and supported:
-      from tinygrad.llm.kernels.amd import q8_linear
-      return q8_linear(self, x)
-    return super().__call__(x)
-
-@functools.cache
-def _gated_delta_prefill_kernel(core:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:UOp, state:UOp, kq:UOp, start_pos:UOp|None=None) -> UOp:
-  batch, heads, tokens, value_dim = cast(tuple[int, int, int, int], core.shape)
-  key_dim, alpha_dim = cast(int, q.shape[-1]), cast(int, alpha.shape[-1]) if len(alpha.shape) == 4 else 1
-  core, v = (x.reshape(batch*heads, tokens, value_dim) for x in (core, v))
-  q, k = (x.reshape(batch*heads, tokens, key_dim) for x in (q, k))
-  beta, kq = (x.reshape(batch*heads, tokens) for x in (beta, kq))
-  alpha, state = alpha.reshape(batch*heads, tokens, alpha_dim), state.reshape(batch*heads, value_dim, key_dim)
-  bh, row, cols = UOp.range(batch*heads, 0, AxisType.GLOBAL), UOp.range(value_dim, 2), tuple(range(key_dim))
-  current = UOp.placeholder((key_dim,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
-  initial = None if start_pos is None else start_pos.eq(0)
-  current = current.after(UOp.group(*(current[col].store(state[bh, row, col].float() if initial is None else
-    initial.where(0, state[bh, row, col].float())) for col in cols)))
-  token = UOp.range(tokens, 1, AxisType.REDUCE)
-  previous = tuple(current.after(token)[col].load() for col in cols)
-  keys, queries = (tuple(x[bh, token, col].load() for col in cols) for x in (k, q))
-  av, bv = alpha[bh, token, row if alpha_dim > 1 else 0].load(), beta[bh, token].load()
-  state_k = sum((x*y for x,y in zip(previous, keys)), UOp.const(0, dtypes.float32))
-  state_q = sum((x*y for x,y in zip(previous, queries)), UOp.const(0, dtypes.float32))
-  delta = (v[bh, token, row].load() - state_k*av) * bv
-  step = UOp.group(core[bh, token, row].store(state_q*av + delta*kq[bh, token]),
-                   *(current[col].store(x*av + delta*y) for col,x,y in zip(cols, previous, keys))).end(token)
-  stores = (state[bh, row, col].store(current.after(step)[col].load().cast(state.dtype)) for col in cols)
-  return UOp.group(*stores).end(row, bh).sink(arg=KernelInfo(name="gated_delta_prefill", opts_to_apply=()))
-
-def gated_delta_prefill(q:Tensor, k:Tensor, v:Tensor, beta:Tensor, alpha:Tensor, state:Tensor, start_pos:Tensor|None=None) -> Tensor:
-  batch, heads, tokens, key_dim = q.shape
-  value_dim = v.shape[-1]
-  assert q.shape == k.shape and v.shape[:3] == q.shape[:3] and beta.shape == (batch, heads, tokens)
-  assert alpha.shape in ((batch, heads, tokens), (batch, heads, tokens, value_dim))
-  assert state.shape == (batch, heads, value_dim, key_dim)
-  kernel = _gated_delta_prefill_kernel
-  if amd_custom_kernels_supported(q.device) and key_dim % 32 == 0 and value_dim % 4 == 0:
-    from tinygrad.llm.kernels.amd import _gated_delta_prefill_kernel as kernel
-  core, kq = Tensor.empty_like(v), (q*k).sum(-1).contiguous()
-  srcs = (core, q.contiguous(), k.contiguous(), v.contiguous(), beta.contiguous(), alpha.contiguous(), state, kq)
-  if start_pos is None: return Tensor.custom_kernel(*srcs, fxn=kernel)[0]
-  contig = tuple(x.uop if x.uop.op is Ops.AFTER else x.uop.contiguous() for x in srcs)
-  params = tuple(UOp.placeholder_like(x, slot=i) for i,x in enumerate(contig))
-  assert start_pos.uop.is_bound_var
-  call = kernel(*params, kernel_var(start_pos.uop.src[0])).call(*contig, start_pos.uop)
-  return Tensor(contig[0].after(call))

tinygrad/llm/kernels/amd.py

@@ -1,451 +0,0 @@
-from __future__ import annotations
-import functools, math
-from typing import Callable, cast
-from tinygrad import Tensor, UOp
-from tinygrad.llm.kernels import Linear, kernel_var
-from tinygrad.uop.ops import AxisType, KernelInfo, Ops, resolve
-from tinygrad.dtype import AddrSpace, dtypes
-
-BLOCK_M, BLOCK_N, DECODE_HEAD_TILE, WARP_SIZE = 32, 32, 8, 32
-WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
-WAVES_M, WAVES_N, LANES_PER_WAVE_M, LANES_PER_WAVE_N = 2, 2, 2, 16
-WMMA_ACC, THREADS_PER_BLOCK = WMMA_M // LANES_PER_WAVE_M, WARP_SIZE * WAVES_M * WAVES_N
-LDS_PAD, WMMA_ARG, LOG2E = 4, ((WMMA_M, WMMA_N, WMMA_K), 'AMD', 32), math.log2(math.e)
-Q5_K, Q6_K, IQ4_XS, GGML_BLOCK_SIZE, Q8_GROUP_SIZE, Q5_WORDS, Q6_BYTES, IQ4_WORDS = 13, 14, 23, 256, 32, 44, 210, 34
-
-def warp_reduce(val:UOp, maximum:bool=False, full_wave:bool=False) -> UOp:
-  for offset in ((16, 8, 4, 2, 1) if full_wave else (8, 4, 2, 1)):
-    if val.op is Ops.INDEX and val.addrspace == AddrSpace.REG: val = val.load()
-    other = UOp(Ops.CUSTOM, dtypes.float, (val,), arg=
-      f"__builtin_bit_cast(float, __builtin_amdgcn_ds_swizzle(__builtin_bit_cast(int, {{0}}), {0x1f | offset<<10}))")
-    val = val.maximum(other) if maximum else val + other
-  return val
-
-def _reg(shape:tuple[int, ...], slot:int, value:float, dep:UOp|None=None) -> UOp:
-  ret = UOp.placeholder(shape, dtypes.float, slot=slot, addrspace=AddrSpace.REG)
-  return ret.after((ret if dep is None else ret.after(dep)).store(ret.const_like(value)))
-
-@functools.cache
-def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, cache_scale, valid_kv_len, max_kv_len, block_n):
-  if isinstance(valid_kv_len, UOp): valid_kv_len = kernel_var(valid_kv_len.unbind_all()[0])
-  _, B, H_KV, N, D = cast(tuple[int, int, int, int, int], cache_kv.shape)
-  _, H, M, _ = cast(tuple[int, int, int, int], q.shape)
-  assert M == 1 and H % H_KV == 0 and D % WARP_SIZE == 0 and max_kv_len <= N and max_kv_len % block_n == 0
-  G, CHUNK, DV, heads_per_wave = H // H_KV, block_n, D // WARP_SIZE, 2
-  head_tile = min(DECODE_HEAD_TILE, G)  # share each KV stream across two GQA heads per wave
-  assert G % head_tile == 0 and head_tile % heads_per_wave == 0
-  decode_waves, decode_group = head_tile // heads_per_wave, 4
-  block_bhkv = UOp.range(B*H_KV*(G//head_tile), 0, AxisType.GLOBAL)
-  valid_chunks = (valid_kv_len+CHUNK-1)//CHUNK
-  group_count = min(valid_chunks, out.shape[2]) if isinstance(valid_chunks, int) else valid_chunks.minimum(out.shape[2])
-  block_n = UOp.range(group_count, 1, AxisType.GLOBAL)
-  lane, wave = UOp.range(WARP_SIZE, 2, AxisType.LOCAL), UOp.range(decode_waves, 3, AxisType.LOCAL)
-  head_group, bhkv = block_bhkv % (G//head_tile), block_bhkv // (G//head_tile)
-  b, kv_head = bhkv // H_KV, bhkv % H_KV
-  dims = tuple(lane + i*WARP_SIZE for i in range(DV))
-  acc, row_max, row_sum = _reg((heads_per_wave, DV), 0, 0), _reg((heads_per_wave,), 1, -math.inf), _reg((heads_per_wave,), 2, 0)
-  groups_per_chunk, offset = CHUNK // decode_group, UOp.range(((valid_chunks+group_count-1)//group_count)*(CHUNK//decode_group), 100, AxisType.REDUCE)
-  chunk = block_n + (offset // groups_per_chunk) * group_count
-  keys = tuple(chunk*CHUNK + (offset % groups_per_chunk)*decode_group + i for i in range(decode_group))
-  valid = tuple(key < valid_kv_len for key in keys)
-  kvals, vvals = (tuple(tuple(cache_kv[kv, b, kv_head, key, d].float() *
-    is_valid.where(cache_scale[kv, b, kv_head, key].float(), UOp.const(0, dtypes.float)) for d in dims)
-    for key,is_valid in zip(keys, valid)) for kv in range(2))
-  q_heads = tuple(kv_head*G + head_group*head_tile + wave*heads_per_wave + head for head in range(heads_per_wave))
-  updates:list[UOp] = []
-  for head,q_head in enumerate(q_heads):
-    scores = tuple(warp_reduce(sum((q[b, q_head, 0, d].float()*k for d,k in zip(dims, key_kvals)),
-                                   UOp.const(0, dtypes.float)), full_wave=True) / math.sqrt(D) for key_kvals in kvals)
-    prev_acc, prev_max, prev_sum = acc.after(offset)[head], row_max.after(offset)[head], row_sum.after(offset)[head]
-    new_max = functools.reduce(lambda a,vs:a.maximum(vs[0].where(vs[1], UOp.const(-math.inf, dtypes.float))), zip(valid, scores), prev_max)
-    alpha = ((prev_max-new_max)*LOG2E).exp2()
-    betas = tuple(is_valid.where(((score-new_max)*LOG2E).exp2(), UOp.const(0, dtypes.float)) for is_valid,score in zip(valid, scores))
-    updates += [acc[head].store(prev_acc*alpha + sum((UOp.stack(*value)*beta for value,beta in zip(vvals, betas)), acc[head].const_like(0))),
-                row_sum[head].store(prev_sum*alpha + sum(betas, UOp.const(0, dtypes.float))), row_max[head].store(new_max)]
-  update = UOp.group(*updates).end(offset)
-  acc, row_max, row_sum = acc.after(update), row_max.after(update), row_sum.after(update)
-  stores = [out[b, q_head, block_n, d].store(acc[head, i]) for head,q_head in enumerate(q_heads) for i,d in enumerate(dims)] + \
-    [stats[b, q_head.valid(lane.eq(0)), block_n, i].store(x[head])
-     for head,q_head in enumerate(q_heads) for i,x in enumerate((row_max, row_sum))]
-  return UOp.group(*stores).end(lane, wave, block_n, block_bhkv).sink(arg=KernelInfo(name="flash_decode_partial", opts_to_apply=()))
-
-def amd_flash_attention_decode(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, cache_scale:Tensor, max_kv_len:int) -> Tensor:
-  B, H, D = cache_kv.shape[1], q.shape[1], cache_kv.shape[4]
-  block_n = 128
-  chunks = min(64, max_kv_len // block_n)
-  partial = Tensor.empty(B, H, chunks, D, dtype="float32", device=q.device)
-  stats = Tensor.empty(B, H, chunks, 2, dtype="float32", device=q.device)
-  decode_partial = functools.partial(_amd_flash_attention_decode_partial, valid_kv_len=valid_kv_len, max_kv_len=max_kv_len, block_n=block_n)
-  partial, stats = Tensor.custom_kernel(partial, stats, q, cache_kv, cache_scale, fxn=decode_partial)[:2]
-  live_chunks = (valid_kv_len+block_n-1)//block_n
-  live_chunks = min(live_chunks, chunks) if isinstance(live_chunks, int) else live_chunks.minimum(chunks)
-  partial, stats = partial[:, :, :live_chunks], stats[:, :, :live_chunks]
-  weights = ((stats[..., 0]-stats[..., 0].max(2, keepdim=True))*LOG2E).exp2()
-  return ((partial*weights.unsqueeze(-1)).sum(2) / (stats[..., 1]*weights).sum(2, keepdim=True)).unsqueeze(2)
-
-@functools.cache
-def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, kv_scale:UOp, valid_kv_len:int|UOp) -> UOp:
-  if isinstance(valid_kv_len, UOp): valid_kv_len = kernel_var(valid_kv_len.unbind_all()[0])
-  BH, M, D = q.shape
-  _, B, H_KV, physical_n, cache_dim = cache.shape
-  k, v = cache[0].reshape(B*H_KV, physical_n, cache_dim), cache[1].reshape(B*H_KV, physical_n, cache_dim)
-  kv_scale = kv_scale.reshape(2, B*H_KV, physical_n)
-  assert k.shape == v.shape and BH % k.shape[0] == 0 and k.shape[2] == D
-  gqa_group = BH // k.shape[0]
-  if isinstance(M, int) and isinstance(valid_kv_len, int):
-    assert M % BLOCK_M == 0 and valid_kv_len % BLOCK_N == 0
-  assert isinstance(D, int) and D % WMMA_K == 0 and D % LANES_PER_WAVE_N == 0
-  TM, TN, TD, SCALE = BLOCK_M//(WAVES_M*LANES_PER_WAVE_M), BLOCK_N//LANES_PER_WAVE_N, D//(WAVES_N*LANES_PER_WAVE_N), 1/math.sqrt(D)
-  block_bh, block_m = UOp.range(BH, 0, AxisType.GLOBAL), UOp.range(M // BLOCK_M, 1, AxisType.GLOBAL)
-  q = q.reshape(BH, M//BLOCK_M, BLOCK_M, D)[block_bh, block_m]
-  kv_head = block_bh // gqa_group
-  k, v = k[kv_head], v[kv_head]
-  o = o.reshape(BH, M//BLOCK_M, BLOCK_M, D)[block_bh, block_m]
-  wave_m, wave_n, lane = UOp.range(WAVES_M, 2, AxisType.LOCAL), UOp.range(WAVES_N, 3, AxisType.LOCAL), UOp.range(WARP_SIZE, -1, AxisType.WARP)
-  tid = (wave_m * WAVES_N + wave_n) * WARP_SIZE + lane
-  lane_m, lane_n = lane // LANES_PER_WAVE_N, lane % LANES_PER_WAVE_N
-  Q_ELEMS_PER_THREAD, KV_ELEMS_PER_THREAD = BLOCK_M * D // THREADS_PER_BLOCK, BLOCK_N * D // THREADS_PER_BLOCK
-  QP_lds = UOp.placeholder((BLOCK_M, D + LDS_PAD), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL)
-  KV_lds = UOp.placeholder((BLOCK_N, D + LDS_PAD), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL)[:, :D]
-  acc, m_i, l_i = _reg((TM, TD), 2, 0), _reg((TM,), 3, -math.inf), _reg((TM,), 4, 0)
-  n_tiles = (valid_kv_len - M + (block_m + 1) * BLOCK_M + BLOCK_N - 1) // BLOCK_N
-  n_tile = UOp.range(n_tiles, 100, AxisType.REDUCE)
-  Q_lds = QP_lds[:, :D]
-  Q_store = Q_lds.after(n_tile).reshape(THREADS_PER_BLOCK, Q_ELEMS_PER_THREAD)[tid].store(q.reshape(THREADS_PER_BLOCK, Q_ELEMS_PER_THREAD)[tid])
-  load_k = UOp.range(KV_ELEMS_PER_THREAD, 90, AxisType.WEAK)
-  kidx = n_tile*BLOCK_N*D + tid*KV_ELEMS_PER_THREAD + load_k
-  kval = k.reshape(physical_n*D)[kidx].float() * kv_scale[0, kv_head, kidx // D].float()
-  K_store = KV_lds.reshape(THREADS_PER_BLOCK, KV_ELEMS_PER_THREAD)[tid, load_k].store(kval).end(load_k)
-  qk_load_barrier = UOp.barrier(UOp.group(Q_store, K_store))
-  Q_lds, KV_lds_k = Q_lds.after(qk_load_barrier), KV_lds.after(qk_load_barrier)
-  S_reg = _reg((TM, TN), 6, 0, n_tile)
-  k_qk, tm1, tn1 = UOp.range(D//WMMA_K, 101, AxisType.REDUCE), UOp.range(TM//WMMA_ACC, 200), UOp.range(TN, 201)
-  S_frag = S_reg.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0, 2, 1)[tm1, tn1]
-  q_frag = Q_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, D // WMMA_K, WMMA_K)[wave_m, tm1, lane_n, k_qk]
-  k_frag = KV_lds_k.reshape(TN, WMMA_N, D // WMMA_K, WMMA_K)[tn1, lane_n, k_qk]
-  qk_done = S_frag.store(UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), *WMMA_ARG)).end(tm1, tn1).end(k_qk)
-  S_reg = S_reg.after(qk_done)
-  S_reg = S_reg.after(S_reg.store(S_reg * SCALE))
-  rm, rn = UOp.range(TM, 250, AxisType.WEAK), UOp.range(TN, 251, AxisType.WEAK)
-  q_idx = valid_kv_len - M + block_m * BLOCK_M + wave_m * WMMA_M + rm * LANES_PER_WAVE_M + lane_m
-  k_idx = n_tile * BLOCK_N + rn * LANES_PER_WAVE_N + lane_n
-  valid = k_idx <= q_idx
-  S_reg = S_reg.after(S_reg[rm, rn].store(valid.where(S_reg[rm, rn], S_reg[rm, rn].const_like(-math.inf))).end(rm, rn))
-  m_ij, rm2 = _reg((TM,), 7, -math.inf, n_tile), UOp.range(TN, 261, AxisType.REDUCE)
-  m_ij = m_ij.after(m_ij.store(m_ij.after(rm2).maximum(S_reg[:, rm2])).end(rm2))
-  ri_w = UOp.range(TM, 270)
-  m_ij = m_ij.after(m_ij[ri_w].store(warp_reduce(m_ij[ri_w], maximum=True)).end(ri_w))
-  tile_max = m_ij.reshape(TM, 1).expand(TM, TN).maximum(-1e30)
-  S_reg = S_reg.after(S_reg.store(((S_reg - tile_max) * LOG2E).exp2()))
-  p_local, ri_ws = _reg((TM,), 8, 0, n_tile), UOp.range(TM, 295, AxisType.WEAK)
-  p_sum = p_local.after(p_local[ri_ws].store(sum((warp_reduce(S_reg[ri_ws, rn]) for rn in range(TN)), S_reg.const_like(0))).end(ri_ws))
-  P_lds = QP_lds.flatten()[:WAVES_N * BLOCK_M * BLOCK_N].reshape(WAVES_N, BLOCK_M, BLOCK_N)
-  P_write = P_lds.reshape(WAVES_N, WAVES_M, TM, LANES_PER_WAVE_M, 1, TN, LANES_PER_WAVE_N, 1)
-  P_write = P_write.permute((1, 0, 3, 6, 2, 4, 5, 7)).reshape(THREADS_PER_BLOCK, TM, TN)
-  P_store = P_write[tid].store(S_reg.cast(dtypes.half))
-  beta_i, ri4 = UOp.placeholder((TM,), dtypes.float, slot=9, addrspace=AddrSpace.REG), UOp.range(TM, 330, AxisType.WEAK)
-  m_new_val = m_i[ri4].maximum(m_ij[ri4])
-  alpha_val = ((m_i[ri4] - m_new_val) * LOG2E).exp2()
-  beta_val = ((m_ij[ri4] - m_new_val) * LOG2E).exp2()
-  rj4 = UOp.range(TD, 331)
-  correction = UOp.group(acc[ri4, rj4].store(alpha_val * acc[ri4, rj4]).end(rj4),
-                         l_i[ri4].store(alpha_val * l_i[ri4] + beta_val * p_sum[ri4]),
-                         m_i[ri4].store(m_new_val), beta_i[ri4].store(beta_val)).end(ri4)
-  acc, l_i, m_i, beta_i = acc.after(correction), l_i.after(correction), m_i.after(correction), beta_i.after(correction)
-  V_lds = UOp.placeholder((D, BLOCK_N + LDS_PAD), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL)[:, :BLOCK_N]
-  V_copy, load_v = V_lds.after(qk_done).permute(1, 0), UOp.range(KV_ELEMS_PER_THREAD, 390, AxisType.WEAK)
-  vidx = n_tile*BLOCK_N*D + tid*KV_ELEMS_PER_THREAD + load_v
-  vval = v.reshape(physical_n*D)[vidx].float() * kv_scale[1, kv_head, vidx // D].float()
-  V_store = V_copy.reshape(THREADS_PER_BLOCK, KV_ELEMS_PER_THREAD)[tid, load_v].store(vval).end(load_v)
-  pv_barrier = UOp.barrier(UOp.group(P_store, V_store))
-  P_lds, V_lds = P_lds.after(pv_barrier), V_lds.after(pv_barrier)
-  pv_acc = _reg((TM, TD), 10, 0, n_tile).after(pv_barrier)
-  k_pv, tm2, tn2 = UOp.range(BLOCK_N//WMMA_K, 400, AxisType.REDUCE), UOp.range(TM//WMMA_ACC, 401, AxisType.WEAK), UOp.range(TD, 402, AxisType.WEAK)
-  pv_frag = pv_acc.reshape(TM // WMMA_ACC, WMMA_ACC, TD).permute(0, 2, 1)[tm2, tn2]
-  p_frag = P_lds[wave_n].reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_N // WMMA_K, WMMA_K)[wave_m, tm2, lane_n, k_pv]
-  v_frag = V_lds.reshape(WAVES_N, TD, WMMA_N, BLOCK_N // WMMA_K, WMMA_K)[wave_n, tn2, lane_n, k_pv]
-  pv_done = pv_frag.store(UOp.wmma(p_frag, v_frag, pv_frag.after(k_pv), *WMMA_ARG)).end(tm2, tn2).end(k_pv)
-  pv_acc = pv_acc.after(pv_done)
-  ri5, rj5 = UOp.range(TM, 410, AxisType.WEAK), UOp.range(TD, 411, AxisType.WEAK)
-  accumulate = acc[ri5, rj5].store(acc[ri5, rj5] + beta_i[ri5] * pv_acc[ri5, rj5]).end(ri5, rj5)
-  n_tile_end = accumulate.barrier().end(n_tile)
-  acc, l_i, m_i = acc.after(n_tile_end), l_i.after(n_tile_end), m_i.after(n_tile_end)
-  acc = acc.after(acc.store(acc * (1 / l_i).reshape(TM, 1).expand(TM, TD)))
-  o = o.reshape(WAVES_M, TM, LANES_PER_WAVE_M, 1, WAVES_N, TD, LANES_PER_WAVE_N, 1)
-  o = o.permute((0, 4, 2, 6, 1, 3, 5, 7)).reshape(THREADS_PER_BLOCK, TM, TD)
-  return o[tid].store(acc).end(wave_m, wave_n, lane).end(block_m, block_bh).sink(arg=KernelInfo(opts_to_apply=()))
-
-def flash_attention_causal_cached(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, cache_scale:Tensor) -> Tensor:
-  B, H, T, D = cast(tuple[int, int, int, int], q.shape)
-  out = Tensor.empty(B*H, T, D, dtype="float32", device=q.device)
-  flash_cached = functools.partial(_amd_flash_attention, valid_kv_len=valid_kv_len)
-  return Tensor.custom_kernel(out, q.reshape(B*H, T, D), cache_kv, cache_scale, fxn=flash_cached)[0].reshape(B, H, T, D)
-
-def quantized_attention(q:Tensor, stacked_kv:Tensor, cache_kv:Tensor, cache_scale:Tensor, start_pos:int|UOp) -> Tensor:
-  T = q.shape[2]
-  scale = (stacked_kv.float().abs().max(axis=-1, keepdim=True) / 127).maximum(1e-8).half()
-  packed_kv = (stacked_kv.float() / scale).round().clip(-127, 127).cast(dtypes.int8)
-  store_kv = cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(packed_kv.uop)
-  store_scale = cache_scale[:, :, :, start_pos:start_pos+T].uop.store(scale.squeeze(-1).uop)
-  # each store goes on its own buffer's AFTER: sharing both stores across both AFTERs leaves
-  # un-ended stores with open ranges in the kernel graph
-  assigned_kv, assigned_scale = Tensor(cache_kv.uop.after(store_kv)), Tensor(cache_scale.uop.after(store_scale))
-  # keep start_pos in its bound form at the graph level, the kernel builders unbind it to the kernel-side PARAM form
-  valid_end = start_pos+T
-  return amd_flash_attention_decode(q.half(), assigned_kv, valid_end, assigned_scale, cast(int, cache_kv.shape[3])) if resolve(T == 1) else \
-    flash_attention_causal_cached(q.half(), assigned_kv, valid_end, assigned_scale)
-
-def _amd_dp4a(a:UOp, b:UOp, c:UOp) -> UOp:
-  return UOp(Ops.CUSTOMI, dtypes.int32, (a.int(), b.int(), c), arg="__builtin_amdgcn_sudot4(true, {}, true, {}, {}, false)")
-
-def _amd_byte_perm(a:UOp, b:UOp, selectors:UOp) -> UOp:
-  return UOp(Ops.CUSTOMI, dtypes.uint32, tuple(x.cast(dtypes.uint32) for x in (a, b, selectors)), arg="__builtin_amdgcn_perm({}, {}, {})")
-
-def _amd_load(ptr:UOp, lanes:int|None=None) -> UOp:
-  assert ptr.op is Ops.INDEX
-  if lanes is None: return UOp(Ops.CUSTOMI, ptr.dtype, (ptr,), arg="__builtin_nontemporal_load({0})")
-  buf, coords = ptr.src[0], ptr.src[1:]
-  idx = sum((coord*math.prod(buf.shape[i+1:]) for i,coord in enumerate(coords)), UOp.const(0, dtypes.weakint))
-  return UOp(Ops.SHRINK, src=(buf.flatten(), idx, UOp.const(lanes, dtypes.weakint))).load(dtype=ptr.dtype)
-
-def _load_byte(raw:UOp, base:UOp, offset:UOp) -> UOp: return (raw[base + offset//4] >> ((offset&3)*8).cast(dtypes.uint32)) & 255
-def _half(value:UOp) -> UOp: return value.cast(dtypes.uint16).bitcast(dtypes.float16).float()
-
-def _iq4_bytes(packed:UOp, shift:int) -> UOp:
-  selectors = (packed >> shift) & 0x0f0f0f0f
-  low = _amd_byte_perm(UOp.const(0xf6eaddcf, dtypes.uint32), UOp.const(0xbfad9881, dtypes.uint32), selectors)
-  high = _amd_byte_perm(UOp.const(0x71594535, dtypes.uint32), UOp.const(0x26190d01, dtypes.uint32), selectors & 0x07070707)
-  return _amd_byte_perm(high, low, 0x03020100 | ((selectors & 0x08080808) >> 1))
-
-@functools.cache
-def _q8_quantize_kernel(q:UOp, scale:UOp, x:UOp, tokens:int, in_features:int) -> UOp:
-  groups = in_features//Q8_GROUP_SIZE
-  token_group, lane = UOp.range(tokens*groups, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
-  token, group = token_group//groups, token_group%groups
-  x = x.reshape(tokens, groups, 32)
-  group_scale = warp_reduce(x[token, group, lane].float().abs(), maximum=True, full_wave=True) / 127
-  group_scale = group_scale.maximum(1e-8)
-  word_lane = lane.minimum(7)
-  xs = tuple(x[token, group, word_lane*4+i].float() for i in range(4))
-  word = sum(((v/group_scale).round().clip(-127, 127).cast(dtypes.int8).cast(dtypes.uint8).cast(dtypes.uint32) << (i*8)
-              for i,v in enumerate(xs)), UOp.const(0, dtypes.uint32))
-  stores = (q[token, group, lane.valid(lane < 8)].store(word), scale[token, group.valid(lane.eq(0))].store(group_scale))
-  return UOp.group(*stores).end(token_group, lane).sink(arg=KernelInfo(name="q8_quantize", opts_to_apply=()))
-
-def q8_quantize(x:Tensor, tokens:int, in_features:int) -> tuple[Tensor, Tensor]:
-  groups = in_features//Q8_GROUP_SIZE
-  q, scale = Tensor.empty(tokens, groups, 8, dtype=dtypes.uint32, device=x.device), \
-    Tensor.empty(tokens, groups, dtype=dtypes.float32, device=x.device)
-  q, scale = Tensor.custom_kernel(q, scale, x, fxn=functools.partial(_q8_quantize_kernel, tokens=tokens, in_features=in_features))[:2]
-  return q, scale
-
-@functools.cache
-def _gated_delta_prefill_kernel(core:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:UOp, state:UOp, kq:UOp, start_pos:UOp|None=None) -> UOp:
-  batch, heads, tokens, value_dim, row_tile = *core.shape, 4
-  key_dim, alpha_dim = q.shape[-1], alpha.shape[-1] if len(alpha.shape) == 4 else 1
-  assert all(isinstance(x, int) for x in (batch, heads, tokens, value_dim, key_dim)) and key_dim % 32 == 0 and value_dim % row_tile == 0
-  batch, heads, tokens, value_dim, key_dim = cast(tuple[int, int, int, int, int], (batch, heads, tokens, value_dim, key_dim))
-  core, v = (x.reshape(batch*heads, tokens, value_dim) for x in (core, v))
-  q, k = (x.reshape(batch*heads, tokens, key_dim) for x in (q, k))
-  beta, kq = (x.reshape(batch*heads, tokens) for x in (beta, kq))
-  alpha, state = alpha.reshape(batch*heads, tokens, alpha_dim), state.reshape(batch*heads, value_dim, key_dim)
-  bh_row, lane = UOp.range(batch*heads*value_dim//row_tile, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
-  bh, row_base = bh_row // (value_dim//row_tile), (bh_row % (value_dim//row_tile))*row_tile
-  rows = tuple(row_base+i for i in range(row_tile))
-  cols = tuple(lane + i*32 for i in range(key_dim//32))
-  current = UOp.placeholder((row_tile*key_dim//32,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
-  initial = None if start_pos is None else start_pos.eq(0)
-  current = current.after(current.store(UOp.stack(*(state[bh, row, col].float() if initial is None else
-    initial.where(0, state[bh, row, col].float()) for row in rows for col in cols))))
-  token = UOp.range(tokens, 2, AxisType.REDUCE)
-  keys = tuple(k[bh, token, col].load() for col in cols)
-  queries = tuple(q[bh, token, col].load() for col in cols)
-  updates:list[UOp] = []
-  stores:list[UOp] = []
-  for row_idx,row in enumerate(rows):
-    previous = tuple(current.after(token)[row_idx*key_dim//32+i].load() for i in range(key_dim//32))
-    av, bv = alpha[bh, token, row if alpha_dim > 1 else 0].load(), beta[bh, token].load()
-    state_k = warp_reduce(sum((x*y for x,y in zip(previous, keys)), UOp.const(0, dtypes.float32)), full_wave=True)
-    state_q = warp_reduce(sum((x*y for x,y in zip(previous, queries)), UOp.const(0, dtypes.float32)), full_wave=True)
-    delta = (v[bh, token, row].load() - state_k*av) * bv
-    updates += [x*av + delta*y for x,y in zip(previous, keys)]
-    stores.append(core[bh, token, row.valid(lane.eq(0))].store(state_q*av + delta*kq[bh, token]))
-  step = UOp.group(*stores, current.store(UOp.stack(*updates))).end(token)
-  state_stores = (state[bh, row, col].store(current.after(step)[row_idx*key_dim//32+i].load().cast(state.dtype))
-                  for row_idx,row in enumerate(rows) for i,col in enumerate(cols))
-  return UOp.group(*state_stores).end(lane, bh_row).sink(arg=KernelInfo(name="gated_delta_prefill", opts_to_apply=()))
-
-def _wmma_layout(out:UOp, out_features:int, token_tile:int, output_tiles:int):
-  output_waves = 2 if out_features % (32*output_tiles) == 0 else 1
-  token_block, output_block = UOp.range(out.shape[0]//token_tile, 0), UOp.range(out_features//(16*output_tiles*output_waves), 1)
-  lane, wave = UOp.range(WARP_SIZE, 2, axis_type=AxisType.LOCAL), UOp.range(output_waves, 3, axis_type=AxisType.LOCAL)
-  hw_lane = UOp(Ops.CUSTOM, dtypes.int32, (lane.int(),), arg="__builtin_amdgcn_mbcnt_lo(-1, 0)").cast(dtypes.weakint)
-  col, half = hw_lane % 16, hw_lane // 16
-  outputs = tuple((output_block*output_waves+wave)*(16*output_tiles) + tile*16 + col for tile in range(output_tiles))
-  inputs = tuple(token_block*token_tile + tile*16 + col for tile in range(token_tile//16))
-  tokens = tuple(tuple(token_block*token_tile + tile*16 + half*8 + i for i in range(8)) for tile in range(token_tile//16))
-  return output_waves, token_block, output_block, lane, wave, half, outputs, inputs, tokens
-
-def _wmma_stores(out, outputs, tokens, accs, update, half):
-  def values(acc:UOp) -> tuple[UOp, ...]:
-    vals = tuple(acc.after(update)[i].load() for i in range(8))
-    swapped = tuple(UOp(Ops.CUSTOM, dtypes.float32, (value,),
-      arg="__builtin_bit_cast(float, __builtin_amdgcn_ds_swizzle(__builtin_bit_cast(int, {0}), 50688))") for value in vals)
-    low = half.eq(0)
-    return tuple(low.where(vals[i], swapped[i+4]) if j == 0 else low.where(swapped[i], vals[i+4]) for i in range(4) for j in range(2))
-  return [out[token, output].store(value) for output,output_accs in zip(outputs, accs)
-          for tile_tokens,acc in zip(tokens, output_accs) for token,value in zip(tile_tokens, values(acc))]
-
-def _decode_linear(out:UOp, out_features:int, group_count:int, group_dot, name:str) -> UOp:
-  chunks = (group_count+31)//32
-  token_output_chunk, lane = UOp.range(out.shape[0]*out_features*chunks, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
-  token, output, chunk = token_output_chunk // (out_features*chunks), (token_output_chunk//chunks) % out_features, token_output_chunk % chunks
-  group = lane+chunk*32
-  value = group_dot(token, output, group) if group_count % 32 == 0 else \
-    (group < group_count).where(group_dot(token, output, group.minimum(group_count-1)), UOp.const(0, dtypes.float32))
-  total = warp_reduce(value, full_wave=True)
-  return out[token, output, chunk.valid(lane.eq(0))].store(total.cast(out.dtype)).end(token_output_chunk, lane).sink(
-    arg=KernelInfo(name=name, opts_to_apply=()))
-
-def _q5_scales(raw:UOp, base:UOp, subgroup:UOp) -> tuple[UOp, UOp, UOp, UOp]:
-  scale = (subgroup < 4).where(_load_byte(raw, base, 4 + subgroup) & 63,
-    (_load_byte(raw, base, 8 + subgroup) & 15) | ((_load_byte(raw, base, subgroup) >> 6) << 4))
-  minimum = (subgroup < 4).where(_load_byte(raw, base, 8 + subgroup) & 63,
-    (_load_byte(raw, base, 8 + subgroup) >> 4) | ((_load_byte(raw, base, 4 + subgroup) >> 6) << 4))
-  d, dmin = (raw[base] & 0xffff).cast(dtypes.uint16), (raw[base] >> 16).cast(dtypes.uint16)
-  return _half(d), _half(dmin), scale.float(), minimum.float()
-
-@functools.cache
-def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, out_features:int, in_features:int, ggml_type:int) -> UOp:
-  group_count = in_features // Q8_GROUP_SIZE
-  def group_dot(token:UOp, output:UOp, group:UOp) -> UOp:
-    block, subgroup = group // 8, group % 8
-    xwords = _amd_load(xq[token, group, 0], 8)
-    if ggml_type == Q5_K:
-      base = (output * in_features//GGML_BLOCK_SIZE + block) * Q5_WORDS
-      qs_base, dot, qsum = base + 12 + (subgroup//2)*8, UOp.const(0, dtypes.int32), UOp.const(0, dtypes.int32)
-      for word_idx in range(8):
-        word = (raw[qs_base+word_idx] >> ((subgroup&1)*4).cast(dtypes.uint32)) & 0x0f0f0f0f
-        word |= ((raw[base+4+word_idx] >> subgroup.cast(dtypes.uint32)) & 0x01010101) << 4
-        dot, qsum = _amd_dp4a(word, xwords[word_idx], dot), _amd_dp4a(UOp.const(0x01010101, dtypes.uint32), xwords[word_idx], qsum)
-      d, dmin, scale, minimum = _q5_scales(raw, base, subgroup)
-      return (dot.float()*d*scale - qsum.float()*dmin*minimum) * xd[token, group]
-    if ggml_type == IQ4_XS:
-      base = (output * in_features//GGML_BLOCK_SIZE + block) * IQ4_WORDS
-      dot = UOp.const(0, dtypes.int32)
-      for word_idx in range(8):
-        packed = _amd_load(raw[base + 2 + subgroup*4 + word_idx%4])
-        dot = _amd_dp4a(_iq4_bytes(packed, 4*(word_idx//4)), xwords[word_idx], dot)
-      d, scale = _iq4_scales(raw, base, subgroup)
-      return dot.float() * xd[token, group] * d * scale
-    base = (output*in_features//GGML_BLOCK_SIZE+block)*Q6_BYTES
-    dots = [UOp.const(0, dtypes.int32), UOp.const(0, dtypes.int32)]
-    for word_idx in range(8):
-      pos, within = subgroup*32 + word_idx*4, (subgroup*32 + word_idx*4)%128
-      low = _amd_load(raw[base + (pos//128)*64 + within%64], 4) >> ((within//64)*4).cast(dtypes.uint8)
-      high = _amd_load(raw[base + 128 + (pos//128)*32 + within%32], 4) >> ((within//32)*2).cast(dtypes.uint8)
-      quant = ((low & 15) | ((high & 3) << 4)).bitcast(dtypes.int8) - 32
-      word = sum((quant[i].cast(dtypes.uint8).cast(dtypes.uint32) << (i*8) for i in range(4)), UOp.const(0, dtypes.uint32))
-      dots[word_idx//4] = _amd_dp4a(word, xwords[word_idx], dots[word_idx//4])
-    scales = [raw[base + 192 + subgroup*2+i].cast(dtypes.uint8).bitcast(dtypes.int8).float() for i in range(2)]
-    dbits = raw[base+208].cast(dtypes.uint16) | (raw[base+209].cast(dtypes.uint16) << 8)
-    return (dots[0].float()*scales[0] + dots[1].float()*scales[1]) * xd[token, group] * _half(dbits)
-  return _decode_linear(out, out_features, group_count, group_dot, {Q5_K:"linear_q5_k", IQ4_XS:"linear_iq4_xs", Q6_K:"linear_q6"}[ggml_type])
-
-def _quant_linear_wmma(out, x, out_features, in_features, type_words, layout, dequant, name):
-  x = x.reshape(out.shape[0], in_features)
-  _, token_block, output_block, lane, wave, physical_half, outputs, input_tokens, tokens = layout
-  token_tile, output_tiles = len(tokens)*16, len(outputs)
-  output_words = in_features // GGML_BLOCK_SIZE * type_words
-  accs = tuple(tuple(UOp.placeholder((8,), dtypes.float32, slot=ot*(token_tile//16)+tile, addrspace=AddrSpace.REG)
-                     for tile in range(token_tile // 16)) for ot in range(output_tiles))
-  accs = tuple(tuple(acc.after(acc.store(acc.const_like(0))) for acc in output_accs) for output_accs in accs)
-  group = UOp.range(in_features // Q8_GROUP_SIZE, 4, AxisType.REDUCE)
-  block, subgroup = group // 8, group % 8
-  wmma_accs = [list(output_accs) for output_accs in accs]
-  for half in range(2):
-    afrags = tuple(UOp.stack(*(x[input_token, group*32 + half*16 + i].cast(dtypes.float16) for i in range(16)))
-                   for input_token in input_tokens)
-    for output_tile,output in enumerate(outputs):
-      bfrag = UOp.stack(*dequant(output*output_words + block*type_words, subgroup, half))
-      for tile,afrag in enumerate(afrags):
-        previous = accs[output_tile][tile].after(group) if half == 0 else wmma_accs[output_tile][tile]
-        wmma_accs[output_tile][tile] = UOp.wmma(afrag, bfrag, previous, *WMMA_ARG)
-  update = UOp.group(*(acc.store(value) for output_accs,output_values in zip(accs, wmma_accs)
-                       for acc,value in zip(output_accs, output_values))).end(group)
-  return UOp.group(*_wmma_stores(out, outputs, tokens, accs, update, physical_half)).end(token_block, output_block, lane, wave).sink(
-    arg=KernelInfo(name=name, opts_to_apply=()))
-
-@functools.cache
-def _q5_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, out_features:int, in_features:int) -> UOp:
-  token_tile, output_tiles = (64, 1) if out_features <= 1024 and out.shape[0] % 64 == 0 else \
-    (64, 2) if out.shape[0] % 64 == 0 else (32 if out.shape[0] % 32 == 0 else 16, 2)
-  def dequant(base:UOp, subgroup:UOp, half:int) -> tuple[UOp, ...]:
-    d, dmin, scale, minimum = _q5_scales(raw, base, subgroup)
-    qs_base = base + 12 + (subgroup // 2)*8 + half*4
-    words = tuple((raw[qs_base+i] >> ((subgroup&1)*4).cast(dtypes.uint32) & 0x0f0f0f0f) |
-      ((raw[base+4+half*4+i] >> subgroup.cast(dtypes.uint32) & 0x01010101) << 4) for i in range(4))
-    return tuple(((word >> (byte*8) & 255).float()*d*scale-dmin*minimum).cast(dtypes.float16) for word in words for byte in range(4))
-  return _quant_linear_wmma(out, x, out_features, in_features, Q5_WORDS,
-                            _wmma_layout(out, out_features, token_tile, output_tiles), dequant, "linear_q5_k_f16_wmma")
-
-def _iq4_scales(raw:UOp, base:UOp, subgroup:UOp) -> tuple[UOp, UOp]:
-  low = _load_byte(raw, base, 4 + subgroup//2)
-  scale = ((low >> (4*(subgroup%2)).cast(dtypes.uint32)) & 15) | ((((raw[base] >> 16) >> (2*subgroup).cast(dtypes.uint32)) & 3) << 4)
-  return _half(raw[base] & 0xffff), (scale.cast(dtypes.uint8).bitcast(dtypes.int8)-32).float()
-
-@functools.cache
-def _iq4_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, lut:UOp, out_features:int, in_features:int) -> UOp:
-  token_tile = 32 if out_features <= 1024 and out.shape[0] % 32 == 0 else 64 if out.shape[0] % 64 == 0 and \
-    (out_features <= 6144 or out_features == 5120 and in_features > 8192) else 128 if out.shape[0] % 128 == 0 else \
-    32 if out.shape[0] % 32 == 0 else 16
-  output_tiles = 1 if out_features <= 1024 else 2 if out_features <= 6144 else 1 if out_features < 8192 else 2
-  layout = _wmma_layout(out, out_features, token_tile, output_tiles)
-  output_waves, _, _, lane, wave, _, _, _, _ = layout
-  local_lut = UOp.placeholder((256,), dtypes.uint32, slot=32, addrspace=AddrSpace.LOCAL)
-  tid, lut_items = wave*32+lane, 256//(32*output_waves)
-  lut = local_lut.after(UOp.group(*(local_lut[tid*lut_items+i].store(lut[tid*lut_items+i]) for i in range(lut_items))).barrier())
-  def dequant(base:UOp, subgroup:UOp, half:int) -> tuple[UOp, ...]:
-    d, scale = _iq4_scales(raw, base, subgroup)
-    scale = scale * d
-    if out_features <= 6144:
-      pairs = tuple(lut[((raw[base + 2 + subgroup*4 + word] >> (byte*8)) & 255).cast(dtypes.weakint)]
-                    for word in range(4) for byte in range(4))
-      return tuple((_half((pair >> (half*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in pairs)
-    def nibble(packed:UOp, index:int): return (packed >> (8*index+4*half)) & 15
-    lut_pairs = (lut[(nibble(packed, i) | nibble(packed, i+1)<<4).cast(dtypes.weakint)]
-                 for packed in (raw[base+2+subgroup*4+i] for i in range(4)) for i in (0, 2))
-    return tuple((_half((pair >> (i*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in lut_pairs for i in range(2))
-  return _quant_linear_wmma(out, x, out_features, in_features, IQ4_WORDS, layout, dequant, "linear_iq4_xs_f16_wmma")
-
-def q8_linear(layer:Linear, x:Tensor) -> Tensor:
-  assert layer.ggml_type in (Q5_K, Q6_K, IQ4_XS)
-  tokens = int(x.numel()) // layer.in_features
-  raw = layer.weight.uop.buf_uop
-  out_features, in_features = layer.out_features, layer.in_features
-  use_wmma = tokens % 16 == 0 and layer.out_features % 16 == 0
-  def run(fxn:Callable[..., UOp], out:UOp, *srcs:UOp) -> Tensor:
-    all_srcs = (out,)+srcs
-    params = tuple(UOp.placeholder_like(src, slot=i) for i,src in enumerate(all_srcs))
-    kernel = fxn(*params, out_features=out_features, in_features=in_features).call(*all_srcs)
-    result = Tensor(out.after(kernel))
-    if len(result.shape) == 3: result = result.sum(-1)
-    result = result.reshape(*x.shape[:-1], layer.out_features)
-    return result if layer.bias is None else result + layer.bias
-
-  out = Tensor.empty(tokens, layer.out_features, dtype=dtypes.float32, device=x.device).uop
-  if layer.ggml_type == Q5_K and use_wmma:
-    return run(_q5_linear_f16_wmma_kernel, out, raw, x.cast(dtypes.float16).contiguous().uop)
-  if layer.ggml_type == IQ4_XS and use_wmma:
-    return run(_iq4_linear_f16_wmma_kernel, out, raw, x.cast(dtypes.float16).contiguous().uop,
-               iq4_half_lut(str(x.device)).uop)
-  xq, xd = q8_quantize(x, tokens, layer.in_features)
-  decode = functools.partial(_quant_decode_kernel, ggml_type=layer.ggml_type)
-  out = Tensor.empty(tokens, layer.out_features, (layer.in_features+1023)//1024, dtype=dtypes.float32, device=x.device).uop
-  return run(decode, out, raw, xq.uop, xd.uop)
-
-@functools.cache
-def iq4_half_lut(device:str) -> Tensor:
-  from tinygrad.runtime.autogen.ggml_common import kvalues_iq4nl
-  return Tensor([x for j in range(16) for i in range(16) for x in (kvalues_iq4nl[i], kvalues_iq4nl[j])],
-                dtype=dtypes.float16, device=device).bitcast(dtypes.uint32).contiguous()

tinygrad/llm/model.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 import functools, itertools, pathlib
 from dataclasses import dataclass, replace
-from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, Context, dtypes
-from tinygrad.llm.kernels import Linear, gated_delta_prefill, amd_custom_kernels_supported
+from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, dtypes
+from tinygrad.nn import Linear
 from tinygrad.llm.gguf import gguf_load
 from tinygrad.uop.ops import resolve
 
@@ -167,38 +167,26 @@ class TransformerBlock(FFNBlock):
     k = apply_rope(k[..., :self.config.rope_dim], self.freqs_cis[start_pos:start_pos+T]).cat(k[..., self.config.rope_dim:], dim=-1)
 
     # NOTE: we don't want to change self.cache_kv, the function API doesn't support this well
-    stacked_kv = Tensor.stack(k, v)
-    if hasattr(self, "cache_kv_scale"):
-      from tinygrad.llm.kernels.amd import quantized_attention
-      attn = quantized_attention(q, stacked_kv, self.cache_kv, self.cache_kv_scale, start_pos)
-    else:
-      assigned_kv = Tensor(self.cache_kv.uop.after(self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(stacked_kv.uop)))
-      k = assigned_kv[0, :, :, 0:start_pos+T, :]
-      v = assigned_kv[1, :, :, 0:start_pos+T, :]
+    assigned_kv = Tensor(self.cache_kv.uop.after(self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(Tensor.stack(k, v).uop)))
+    k = assigned_kv[0, :, :, 0:start_pos+T, :]
+    v = assigned_kv[1, :, :, 0:start_pos+T, :]
 
-      #self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(stacked_kv)
-      #k = self.cache_kv[0, :, :, 0:start_pos+T, :]
-      #v = self.cache_kv[1, :, :, 0:start_pos+T, :]
+    #self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(Tensor.stack(k, v))
+    #k = self.cache_kv[0, :, :, 0:start_pos+T, :]
+    #v = self.cache_kv[1, :, :, 0:start_pos+T, :]
 
-      # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
-      # TODO: this if statement should be removed and it shouldn't generate extra kernels
-      mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, buffer=False).triu(start_pos+1) \
-        if resolve(T != 1) else None
-      attn = q.scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True)     # (B,H,T,Hd)
+    # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
+    # TODO: this if statement should be removed and it shouldn't generate extra kernels
+    mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, buffer=False).triu(start_pos+1) \
+      if resolve(T != 1) else None
+    attn = q.scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True)     # (B,H,T,Hd)
     attn = attn.transpose(1, 2).reshape(B, T, -1)                                    # back to (B,T,D)
     return self.attn_output(attn if not self.config.attn_output_gate else (attn * gate.sigmoid()))
 
   def _init_state(self, x:Tensor):
     if not hasattr(self, "cache_kv"):
-      # hybrid models use a quantized KV cache on AMD, sized in flash decode blocks of 256
-      quantize = amd_custom_kernels_supported(x.device) and self.config.ssm is not None
-      assert not quantize or self.config.max_context % 256 == 0, \
-        f"quantized KV cache needs max_context to be a multiple of 256, got {self.config.max_context}"
       self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim,
-                                   dtype=dtypes.int8 if quantize else dtypes.default_float, device=x.device)
-      if quantize:
-        self.cache_kv_scale = Tensor.zeros(2, x.shape[0], self.config.n_kv_heads, self.config.max_context,
-                                           dtype=dtypes.float16, device=x.device)
+                                   dtype=dtypes.default_float, device=x.device)
       self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta, device=x.device)
 
 class MLATransformerBlock(FFNBlock):
@@ -270,8 +258,12 @@ class GatedDeltaNetBlock(FFNBlock):
 
   def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
     B, T, _ = x.shape
+    # bind ints to a variable so the reset flag stays a runtime value (it toggles when generation restarts at position 0)
     start_pos = start_pos if isinstance(start_pos, UOp) else UOp.variable("start_pos", 0, self.config.max_context-1).bind(start_pos)
+    initial = Tensor(start_pos).eq(0)
     is_kda = hasattr(self, "ssm_g_a")
+    symbolic = isinstance(T, UOp)
+    T_pad = x.max_shape[1]  # symbolic chunks are padded to their max size: one graph serves every size
 
     # input processing
     x = x.half()
@@ -279,36 +271,57 @@ class GatedDeltaNetBlock(FFNBlock):
     out_gate = out_gate.reshape(B, T, self.num_v_heads, self.head_v_dim)
     beta = self.ssm_beta(x).sigmoid().reshape(B, T, self.num_v_heads)
     alpha = self.ssm_f_b(self.ssm_f_a(x)) if is_kda else self.ssm_alpha(x)
-    log_alpha = ((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, T, self.num_v_heads, -1) * self.ssm_a).squeeze(-1) \
-      if is_kda else ((alpha.float() + self.ssm_dt["bias"]).softplus() * self.ssm_a).reshape(B, T, self.num_v_heads)
+    log_alpha = ((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, T, self.num_v_heads, -1) *
+                 self.ssm_a.reshape(self.num_v_heads, -1))
 
     # qkv conv, conv_state is reset when starting from position 0
-    conv_state = Tensor(start_pos).eq(0).where(0, self.conv_state)
-    conv_window = conv_state.cat(self.attn_qkv(x), dim=1)
-    conv_out = ((conv_window * self.ssm_conv1d["weight"].T.unsqueeze(0)).sum(1) if is_kda and resolve(T == 1) else functools.reduce(lambda a,b: a+b,
-      (conv_window[:, i:i+T] * self.ssm_conv1d["weight"][:, i] for i in range(self.ssm_conv_kernel)))).silu()
+    conv_state = initial.where(0, self.conv_state)
+    # assemble the conv window in a static-size buffer: [conv_state | qkv rows | zero-pad].
+    # padded steps are exact no-ops: beta=0 (delta rule off), log_alpha=0 (decay 1 after exp)
+    win = Tensor.zeros(B, self.ssm_conv_kernel-1 + T_pad, self.conv_channels).uop
+    win = win.after(win[:, :self.ssm_conv_kernel-1].store(conv_state.cast(win.dtype).uop))
+    win = win.after(win[:, self.ssm_conv_kernel-1:self.ssm_conv_kernel-1+T].store(self.attn_qkv(x).cast(win.dtype).uop))
+    conv_window = Tensor(win)
+    # the last conv_kernel-1 columns of the window become the next conv state
+    conv_state_store = self.conv_state.uop.store(conv_window[:, T:T+self.ssm_conv_kernel-1].cast(self.conv_state.dtype).uop)
+
+    conv_out = functools.reduce(lambda a,b: a+b,
+      (conv_window[:, i:i+T_pad] * self.ssm_conv1d["weight"][:, i] for i in range(self.ssm_conv_kernel))).silu()
+    if symbolic:
+      out_gate = out_gate.pad_to((B, T_pad, self.num_v_heads, self.head_v_dim))
+      beta, log_alpha = beta.pad_to((B, T_pad, self.num_v_heads)), log_alpha.pad_to((B, T_pad, *log_alpha.shape[2:]))
     q, k, v = conv_out.split([self.q_dim, self.q_dim, self.conv_channels - 2*self.q_dim], dim=-1)
-    q = q.reshape(B, T, self.num_k_heads, self.head_k_dim).normalize(dim=-1, eps=1e-12 if is_kda else 1e-6)
-    k = k.reshape(B, T, self.num_k_heads, self.head_k_dim).normalize(dim=-1, eps=1e-12 if is_kda else 1e-6)
-    q, k = q.repeat(1, 1, self.num_v_heads//self.num_k_heads, 1), k.repeat(1, 1, self.num_v_heads//self.num_k_heads, 1)
-    v = v.reshape(B, T, self.num_v_heads, self.head_v_dim)
-    q, k, v, beta = [z.transpose(1, 2).float() for z in (q, k, v, beta)]
-    alpha = log_alpha.transpose(1, 2).exp()
+    qk_eps = 1e-12 if is_kda else 1e-6
+    q, k = (z.reshape(B, T_pad, self.num_k_heads, self.head_k_dim).normalize(dim=-1, eps=qk_eps)
+            .repeat(1, 1, self.num_v_heads//self.num_k_heads, 1) for z in (q, k))
+    v = v.reshape(B, T_pad, self.num_v_heads, self.head_v_dim)
+    # layout the per-step operands to broadcast against the (B, H, V, K) state
+    q, k, v, beta = (z.transpose(1, 2).float() for z in (q, k, v, beta))
+    q, k, v, beta = q.unsqueeze(-2) * self.head_k_dim**-0.5, k.unsqueeze(-2), v.unsqueeze(-1), beta.unsqueeze(-1).unsqueeze(-1)
+    alpha = log_alpha.transpose(1, 2).exp().unsqueeze(-1)  # per-channel decay for kda, per-head otherwise (B, H, T, V|1, 1)
 
-    # recurrent, the conv and recurrent states are updated in place
-    conv_state = conv_window[:, T:T+self.ssm_conv_kernel-1].cast(self.conv_state.dtype).contiguous()
-    state = Tensor(self.recurrent_state.uop.after(self.conv_state.uop.store(conv_state.uop)))
-    core = gated_delta_prefill(q * self.head_k_dim**-0.5, k, v, beta, alpha, state, Tensor(start_pos)).transpose(1, 2)
+    # recurrent: scan over the (padded) tokens, updating the recurrent state. collect the per-step outputs
+    state = Tensor(self.recurrent_state.uop.after(conv_state_store)).float()  # carry the conv write into this graph
+    state = initial.where(0, state)
+    outs = []
+    for t in range(T_pad):
+      s1 = state * alpha[:, :, t]  # decay the state
+      delta = (v[:, :, t] - (s1*k[:, :, t]).sum(-1, keepdim=True)) * beta[:, :, t]  # the delta rule update
+      state = s1 + delta * k[:, :, t]
+      outs.append((state * q[:, :, t]).sum(-1))
 
-    # output
-    core_attn_out = self.ssm_norm(core)
-    out_gate = out_gate.sigmoid() if is_kda else out_gate.silu()
-    return self.ssm_out((core_attn_out * out_gate).reshape(B, T, -1).cast(x.dtype)).contiguous()
+    # store the updated recurrent state in place, then read the stacked outputs after the write
+    core = Tensor(outs[0].stack(*outs[1:], dim=1).contiguous().uop.after(self.recurrent_state.uop.store(state.cast(self.recurrent_state.dtype).uop)))
+
+    # output; undo the padding before the output projection
+    z = (self.ssm_norm(core) * (out_gate.sigmoid() if is_kda else out_gate.silu())).cast(x.dtype).contiguous()
+    if symbolic: z = z[:, :T]
+    return self.ssm_out(z.reshape(B, T, -1))
 
   def _init_state(self, x):
     if not hasattr(self, "conv_state"):
-      self.conv_state = Tensor.empty(x.shape[0], self.ssm_conv_kernel-1, self.conv_channels, device=x.device)
-      self.recurrent_state = Tensor.empty(x.shape[0], self.num_v_heads, self.head_v_dim, self.head_k_dim, device=x.device)
+      self.conv_state = Tensor.zeros(x.shape[0], self.ssm_conv_kernel-1, self.conv_channels, device=x.device).clone()
+      self.recurrent_state = Tensor.zeros(x.shape[0], self.num_v_heads, self.head_v_dim, self.head_k_dim, device=x.device).clone()
 
 class Transformer:
   def __init__(self, config:TransformerConfig):
@@ -422,6 +435,9 @@ class Transformer:
       Tensor.realize(*params)
     return model, kv
 
+  def warmup(self):
+    for _ in range(2): list(zip(range(2), self.generate([0])))
+
   def get_start_pos(self, tokens:list[int]) -> int:
     # recurrent state can't be partially reused after divergence: reuse it only when tokens extend the cached prefix
     if self.has_recurrent_block:
@@ -430,19 +446,8 @@ class Transformer:
     prefix_len = sum(1 for _ in itertools.takewhile(lambda ab: ab[0] == ab[1], zip(tokens[:-1], self._cached_tokens)))
     return min(block._reusable_prefix_len(prefix_len, len(self._cached_tokens)) for block in self.blk)
 
-  def warmup(self, chunk_size:int=32):
-    prompt = [0] * (min(chunk_size, 256, self.max_context-1) if self.has_recurrent_block else 1)
-    if self.has_recurrent_block:
-      x = Tensor.empty(1, 1, self.blk[0].config.dim, device=self.token_embd.weight.device)
-      for block in self.blk: block._init_state(x)
-    for _ in range(2):
-      # NOTE: chunk_size must match what generate uses at serve time, otherwise the captured JIT rejects the new toks range
-      warm = self.generate(prompt, chunk_size=chunk_size)
-      with Context(JIT_BATCH_SIZE=getenv("PREFILL_JIT_BATCH_SIZE", 512) if self.has_recurrent_block else 0): next(warm)
-      with Context(JIT_BATCH_SIZE=0): next(warm)
-      self._cached_tokens = []
-
   def generate(self, tokens:list[int], chunk_size:int=32, temperature:float=0.0):
+    if self.has_recurrent_block: chunk_size = 1
     v_start_pos = UOp.variable("start_pos", 0, self.max_context-1)
     v_toks = UOp.variable("toks", 1, chunk_size)
     # TODO: use UOp.variable for temperature once float variables are supported
@@ -453,10 +458,8 @@ class Transformer:
     start_pos = self.get_start_pos(tokens)
     out, prompt_len = None, len(tokens)
     while len(tokens) < self.max_context:
-      # recurrent blocks prefill full chunks with a static shape, the tail of the prompt goes through the decode graph
-      remaining = len(tokens)-start_pos
-      n_toks = 1 if self.has_recurrent_block and remaining < chunk_size else min(chunk_size, remaining)
-      sp, nt = v_start_pos.bind(start_pos), n_toks if self.has_recurrent_block else v_toks.bind(n_toks)
+      n_toks = min(chunk_size, len(tokens) - start_pos)
+      sp, nt = v_start_pos.bind(start_pos), v_toks.bind(n_toks)
       out = self(t[:, sp:sp+nt] if start_pos < prompt_len or out is None else out, sp, temp).realize()
       start_pos += n_toks
       # chunked prefill: keep processing until all prompt tokens are consumed

tinygrad/mixin/elementwise.py

@@ -115,7 +115,8 @@ class ElementwiseMixin(CreationMixin):
     ```
     """
     a, b = self._broadcasted(x, reverse)
-    return a + (-b)
+    # alu, not +: _broadcasted already promoted these, and a second promote would cast -b (only a bare weak CONST is kept weak)
+    return a.alu(Ops.ADD, -b)
 
   def mul(self, x: Self | ConstType, reverse: bool = False) -> Self:
     """
@@ -245,8 +246,9 @@ class ElementwiseMixin(CreationMixin):
     if dtypes.is_int(a.dtype) and dtypes.is_int(b.dtype):
       if rounding_mode == "trunc": return a.alu(Ops.CDIV, b)
       if rounding_mode == "floor": return a.alu(Ops.FLOORDIV, b)
-      a = a.cast(dtypes.default_float)
-    d = a * b.reciprocal()
+    if dtypes.is_int(a.dtype) or a.dtype == dtypes.bool: a = a.cast(dtypes.default_float)
+    # alu, not *: _broadcasted already promoted these, and a second promote would cast 1/b (only a bare weak CONST is kept weak)
+    d = a.alu(Ops.MUL, b.reciprocal())
     if rounding_mode is None: return d
     if rounding_mode == "trunc": return d.trunc()
     if rounding_mode == "floor": return d.floor()

tinygrad/mixin/gradient.py

@@ -3,6 +3,7 @@ import math, dataclasses
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata, broadcast_axes
 from tinygrad.helpers import argsort
 from tinygrad.dtype import sum_acc_dtype
+from tinygrad.function import renumber_invalid_outputs
 
 def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
   if op == Ops.ADD: return (ctx._broadcast_to(ret.src[0].shape),)
@@ -32,7 +33,7 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
   params = {x.arg.slot:x for x in fxn.toposort(enter_calls=False) if x.op == Ops.PARAM}
   grad_args = ctx.src
   root_grad = UOp(Ops.TUPLE, src=tuple(UOp(Ops.NOOP) if g.op is Ops.NOOP else
-    g if g.base.op is Ops.CONST else g.param_like(len(args)+i) for i,g in enumerate(grad_args)))
+    g if g.device is None else g.param_like(len(args)+i) for i,g in enumerate(grad_args)))
   grads = compute_gradient(fxn, root_grad, set(params.values()))
   # for precompiled calls, substitute forward outputs with params so intermediates aren't recomputed
   fwd_subs = {src: src.param_like(len(args)+len(grad_args)+i) for i, src in enumerate(fxn.src)} if k.arg.precompile else {}
@@ -40,6 +41,7 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
   # collect needed gradient bodies, compact unused params, create a single backward CALL
   grad_bodies = [(i, grads[p]) for i in needed if (p:=params.get(i)) is not None and p in grads]
   bwd_body = UOp.maketuple(*(gb for _, gb in grad_bodies)).substitute(fwd_subs, walk=True)
+  bwd_body = renumber_invalid_outputs(bwd_body)
   bwd_body, compact_args = _compact_params(bwd_body, (*args, *grad_args, *fwd_outs))
   bwd_call = bwd_body.call(*compact_args, name=(k.arg.name or "")+"_backward", precompile=k.arg.precompile_backward)
   gb_map = {i: idx for idx, (i, _) in enumerate(grad_bodies)}

tinygrad/mixin/movement.py

@@ -46,6 +46,16 @@ class MovementMixin:
     """
     return prod(self.shape)
 
+  @property
+  def max_shape(self) -> tuple[int, ...]:
+    """The shape with every symbolic dimension replaced by its maximum."""
+    from tinygrad.uop.ops import to_max_shape  # deferred: ops.py imports the mixins
+    return to_max_shape(self.shape)
+
+  def max_numel(self) -> int:
+    """The number of elements in `max_shape`."""
+    return prod(self.max_shape)
+
   def size(self, dim:int|None=None) -> sint|tuple[sint, ...]:
     """
     Returns the size of the tensor. If `dim` is specified, return the length along dimension `dim`. Otherwise return the shape of the tensor.

tinygrad/mixin/op.py

@@ -289,6 +289,12 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     if value == 0: return base
     return MovementMixin.pad(X.const_like(True, dtypes.bool), pads).where(base, value)
 
+  def pad_to(self, shape, *args, value:ConstType=0) -> Self:
+    # same mask trick as _pad_constant so the fill survives backends that realize PAD as 0-fill
+    ret = MovementMixin.pad_to(self, shape, *args)
+    if value == 0 or ret is self: return ret
+    return MovementMixin.pad_to(self.const_like(True, dtypes.bool), shape, *args).where(ret, value)
+
   def _pad_circular(self, pX:tuple[tuple[sint, sint], ...]) -> Self:
     # shrink first for negative pads, then wrap the non-negative remainder
     X = self.shrink(tuple((-smin(pB,0), smin(pA+sh,sh)) for (pB,pA),sh in zip(pX, self.shape)))
@@ -460,6 +466,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     """
     assert gradient is not None or self.shape == tuple(), "when no gradient is provided, backward must be called on a scalar tensor"
     if not (self.is_floating_point() and all(t.is_floating_point() for t in targets)): raise RuntimeError("only float Tensors have gradient")
+    if any(t.dtype in dtypes.weaks for t in targets): raise RuntimeError("cannot take gradient wrt a weak Tensor")
     from tinygrad.mixin.gradient import compute_gradient
     if gradient is None: gradient = self.const_like(1.0)
     target_uops = [t._uop for t in targets]

tinygrad/renderer/__init__.py

@@ -35,8 +35,8 @@ class Estimates:
         while len(buf.src) and buf.op is not Ops.PARAM: buf = buf.src[0]
         if buf.op is Ops.PARAM:
           # u.src[0] is INDEX, cap at buffer size for re-reads (e.g. matmul)
-          accessed = mem.get((buf, u.op), 0) + u.src[0].max_numel() * u.src[0].dtype.scalar().itemsize * mults
-          mem[(buf, u.op)] = smin(accessed, buf.max_numel() * buf.dtype.scalar().itemsize)
+          accessed = mem.get((buf, u.op), 0) + u.src[0].max_numel() * u.src[0].dtype.itemsize * mults
+          mem[(buf, u.op)] = smin(accessed, buf.max_numel() * buf.dtype.itemsize)
       if u.op is Ops.RANGE:
         mult_stack.append(mults)
         if u.dtype is not dtypes.void:  # unbounded loop, unknown trip count
@@ -47,9 +47,9 @@ class Estimates:
       elif u.op is Ops.SPECIAL: mults *= cast(sint, u.src[0].ssimplify()) # NOTE: we don't push to the mult_stack here, you can't end these
       elif u.op is Ops.PARAM and u.arg.addrspace == AddrSpace.ALU and u.expr == 'core_id': mults *= int(u.vmax) + 1
       elif u.op is Ops.LOAD and u.src[0].addrspace != AddrSpace.REG:
-        lds += u.max_numel() * u.dtype.scalar().itemsize * mults
+        lds += u.max_numel() * u.dtype.itemsize * mults
       elif u.op is Ops.STORE and u.src[0].addrspace != AddrSpace.REG:
-        lds += u.max_numel() * u.src[1].dtype.scalar().itemsize * mults
+        lds += u.max_numel() * u.src[1].dtype.itemsize * mults
       elif u.op in GroupOp.ALU and u not in excluded:
         flops += (mults * (2 if u.op is Ops.MULACC else 1)) * u.max_numel()
       elif u.op is Ops.WMMA and u not in excluded:

tinygrad/renderer/cstyle.py

@@ -107,11 +107,11 @@ def uops_to_dtypes(uops:list[UOp]) -> list[tuple[DType, int]]:
 
 def _wmma_name(u:UOp) -> str:
   # sanitize spaces in DType.name (int8 = "signed char")
-  return f"WMMA_{'_'.join(map(str, u.arg[0]))}_{u.arg[1].name}_{u.dtype.scalar().name}".replace(" ", "_")
+  return f"WMMA_{'_'.join(map(str, u.arg[0]))}_{u.arg[1].name}_{u.dtype.name}".replace(" ", "_")
 
 # (name, dims, dtype_in, dtype_out, device, threads, upcast_sizes)
 def wmma_args(uops:list[UOp]):
-  return dedup((_wmma_name(uop), uop.arg[0], uop.arg[1], uop.dtype.scalar(), *(uop.arg[2:4]),
+  return dedup((_wmma_name(uop), uop.arg[0], uop.arg[1], uop.dtype, *(uop.arg[2:4]),
                tuple(uop.src[i].shape[-1] for i in range(3)))
               for uop in uops if uop.op is Ops.WMMA)
 
@@ -182,8 +182,8 @@ class CStyleLanguage(Renderer):
     if addrspace in (AddrSpace.LOCAL, AddrSpace.GLOBAL) or override_ptr:
       suffix = "*"
     if sz > 1:
-      return prefix + self.type_map.get(scalar:=dtype.scalar(), scalar.name).replace(" ", "_") + str(sz) + suffix
-    return prefix + self.type_map.get(scalar:=dtype.scalar(), scalar.name) + suffix
+      return prefix + self.type_map.get(dtype, dtype.name).replace(" ", "_") + str(sz) + suffix
+    return prefix + self.type_map.get(dtype, dtype.name) + suffix
 
   def render_type(self, u:UOp): return self._render_dtype(u.dtype, u.max_numel(), u.addrspace, shape=u._shape)
   def render_access(self, u:UOp):
@@ -472,7 +472,7 @@ class CUDARenderer(CStyleLanguage):
 class NVCCRenderer(CUDARenderer):
   def __init__(self, target:Target): super().__init__(target, use_nvcc=True)
 
-def fp8_index(dtype: DType): return (dtypes.fp8e4m3, dtypes.fp8e5m2).index(dtype.scalar())
+def fp8_index(dtype: DType): return (dtypes.fp8e4m3, dtypes.fp8e5m2).index(dtype)
 def _ocml(op): return lambda x,dtype: f"__ocml_{op}_f{ {dtypes.half:16, dtypes.double:64}.get(dtype, 32)}({x})"
 
 class HIPRenderer(CStyleLanguage):
@@ -546,7 +546,7 @@ class HIPRenderer(CStyleLanguage):
       ockl = [(f"__ockl_get_{name}", "unsigned int", "size_t", "const") for name in ["local_id", "group_id", "local_size"]]
     ocml_ops = {Ops.EXP2: ("exp2", "pure"), Ops.LOG2: ("log2", "pure"), Ops.SQRT: ("sqrt", "const"), Ops.SIN: ("sin", ""), Ops.TRUNC: ("trunc", "")}
     ocml = [(f"__ocml_{ocml_ops[op][0]}_f{dt.bitsize}", dt.name, dt.name, ocml_ops[op][1])
-      for op, dt in dedup((u.op, u.dtype.scalar()) for u in uops) if op in ocml_ops and dt in (dtypes.half, dtypes.float, dtypes.double)]
+      for op, dt in dedup((u.op, u.dtype) for u in uops) if op in ocml_ops and dt in (dtypes.half, dtypes.float, dtypes.double)]
     if any(dt == dtypes.bfloat16 for dt, _ in used_dtypes):
       prefix.append(f"typedef {'__bf16' if self.is_cdna4(self.target.arch) else 'unsigned short'} hip_bfloat16;")
     if any(dt == dtypes.half for dt, _ in used_dtypes): prefix.append("#define half _Float16")

tinygrad/renderer/isa/x86.py

@@ -165,7 +165,7 @@ def scratch_buffer(elem_dt:DType, count:int, slot:int) -> UOp:
   return UOp.placeholder((count,), elem_dt, slot, AddrSpace.LOCAL)
 
 def gated_load(ctx, addr:UOp, alt:UOp, gate:UOp, x:UOp):
-  local = scratch_buffer(addr.src[0].dtype.scalar(), x.max_numel(), next(ctx))
+  local = scratch_buffer(addr.src[0].dtype, x.max_numel(), next(ctx))
   local_idx = local.index(UOp.const(0, dtypes.int32), dtype=dtypes.uint64)
   # the selected address is a 64bit value, the AFTER orders the load after the scratch store and carries the element dtype for the encoder
   sel = gate.where(addr.replace(dtype=dtypes.uint64), local_idx)
@@ -173,7 +173,7 @@ def gated_load(ctx, addr:UOp, alt:UOp, gate:UOp, x:UOp):
   return ptr.load(dtype=x.dtype)
 
 def gated_store(addr:UOp, gate:UOp, val:UOp):
-  local = scratch_buffer(addr.src[0].dtype.scalar(), val.max_numel(), -1)
+  local = scratch_buffer(addr.src[0].dtype, val.max_numel(), -1)
   sel = gate.where(addr.replace(dtype=dtypes.uint64), local.index(UOp.const(0, dtypes.int32), dtype=dtypes.uint64))
   return UOp(Ops.AFTER, addr.dtype, (sel,)).store(val)
 
@@ -237,7 +237,7 @@ def cmp(x:UOp) -> UOp:
   return x.ins(X86Ops.CMP, dtype=dtypes.void) if (i:=to_imm(x.src[1])) is None else x.ins(X86Ops.CMPi, dtype=dtypes.void, src=(x.src[0], i))
 def vcmp(x:UOp) -> UOp:
   v = imm(dtypes.uint8, {Ops.CMPLT: 1, Ops.CMPNE: 4, Ops.CMPEQ: 0}[x.op])
-  if x.dtype.scalar() is dtypes.float32: return x.ins(X86Ops.VCMPSS if x.max_numel() == 1 else X86Ops.VCMPPS, src=x.src + (v,))
+  if x.dtype is dtypes.float32: return x.ins(X86Ops.VCMPSS if x.max_numel() == 1 else X86Ops.VCMPPS, src=x.src + (v,))
   return x.ins(X86Ops.VCMPSD if x.max_numel() == 1 else X86Ops.VCMPPD, src=x.src + (v,))
 
 # vinsertps xmm2, xmm0, xmm1, imm
@@ -252,7 +252,7 @@ def vinsertps(x:UOp) -> UOp:
 # vpinsq xmm2, xmm0, rax, imm
 # inserts element in rax into any position in xmm0, result is written to xmm2 according to imm
 def vpins(x:UOp) -> UOp:
-  op = {1: X86Ops.VPINSRB, 2: X86Ops.VPINSRW, 4: X86Ops.VPINSRD, 8: X86Ops.VPINSRQ}[x.dtype.scalar().itemsize]
+  op = {1: X86Ops.VPINSRB, 2: X86Ops.VPINSRW, 4: X86Ops.VPINSRD, 8: X86Ops.VPINSRQ}[x.dtype.itemsize]
   return functools.reduce(lambda ret,i: x.ins(op, src=(ret, x.src[i], imm(dtypes.uint8, i))), range(len(x.src)), def_reg(x.dtype))
 
 # we don't call ctx.vreg on the srcs to avoid duplicates, a rewrite will assign the tuple of valid registers to a vreg

tinygrad/renderer/ptx.py

@@ -64,7 +64,7 @@ def render_wmma(ctx: "PTXRenderer", wmma: UOp):
 
   for src, regs in zip(wmma.src, ctx.wmma_r):
     for i, reg in enumerate(regs): # pack input and acc registers
-      if (elems_per_reg := 4 // src.dtype.scalar().itemsize) == 1: yield f"mov.b32 {reg}, {ctx.r[src][i]};"
+      if (elems_per_reg := 4 // src.dtype.itemsize) == 1: yield f"mov.b32 {reg}, {ctx.r[src][i]};"
       else: yield f"mov.b32 {reg}, {{{', '.join(ctx.r[src][i * elems_per_reg : (i+1) * elems_per_reg])}}};"
 
   dt_map_in, dt_map_out = {dtypes.float: "tf32", dtypes.half: "f16"}, {dtypes.float: "f32", dtypes.half: "f16"}
@@ -101,17 +101,17 @@ string_rewrite = PatternMatcher([
      if loc.addrspace == AddrSpace.REG else None),
   (UPat(Ops.STORE, src=(UPat((Ops.INDEX, Ops.SHRINK), name="loc"), UPat.var("var"))),
    lambda ctx, loc, var: f"st.{mem_type(loc)}" + \
-    f"{f'.v{cnt}' if ((cnt:=var.max_numel())>1) else ''}.{ctx.mem_types[var.dtype.scalar()]} " + \
+    f"{f'.v{cnt}' if ((cnt:=var.max_numel())>1) else ''}.{ctx.mem_types[var.dtype]} " + \
     f"[{ctx.r[loc]}+0], {('{' + ', '.join(ctx.r[var]) + '}') if var.max_numel() > 1 else ctx.r[var]};"),
   (UPat(Ops.LOAD, name="x", src=(UPat((Ops.INDEX, Ops.SHRINK), name="loc"), UPat.var("alt"), UPat.var("gate"))),
     lambda ctx, x, loc, alt, gate: flatten([
-    [f"mov.{ctx.mem_types[x.dtype.scalar()]} {v}, {render_val(0, x.dtype.scalar())};" for v in ctx.r[x]],
-    [f"@{ctx.r[gate]} ld.{mem_type(loc)}.v{x.max_numel()}.{ctx.mem_types[x.dtype.scalar()]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];"]
+    [f"mov.{ctx.mem_types[x.dtype]} {v}, {render_val(0, x.dtype)};" for v in ctx.r[x]],
+    [f"@{ctx.r[gate]} ld.{mem_type(loc)}.v{x.max_numel()}.{ctx.mem_types[x.dtype]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];"]
   ]) if alt.max_numel() > 1 else [
-    f"@{ctx.r[gate]} ld.{mem_type(loc)}.{ctx.mem_types[x.dtype.scalar()]} {ctx.r[x]}, [{ctx.r[loc]}+0];",
-    f"@!{ctx.r[gate]} mov.b{ctx.types[x.dtype.scalar()][1:]} {ctx.r[x]}, {ctx.r[alt]};"]),
+    f"@{ctx.r[gate]} ld.{mem_type(loc)}.{ctx.mem_types[x.dtype]} {ctx.r[x]}, [{ctx.r[loc]}+0];",
+    f"@!{ctx.r[gate]} mov.b{ctx.types[x.dtype][1:]} {ctx.r[x]}, {ctx.r[alt]};"]),
   (UPat(Ops.LOAD, name="x", src=(UPat((Ops.INDEX, Ops.SHRINK), name="loc"),)),
-    lambda ctx, x, loc: f"ld.{mem_type(loc)}.v{x.max_numel()}.{ctx.mem_types[x.dtype.scalar()]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];" \
+    lambda ctx, x, loc: f"ld.{mem_type(loc)}.v{x.max_numel()}.{ctx.mem_types[x.dtype]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];" \
      if x.max_numel() > 1 else f"ld.{mem_type(loc)}.{ctx.mem_types[x.dtype]} {ctx.r[x]}, [{ctx.r[loc]}+0];"),
   # simple
   (UPat(Ops.BUFFER, name="x"), lambda ctx, x: [] if x.addrspace == AddrSpace.REG else [
@@ -197,7 +197,7 @@ class PTXRenderer(Renderer):
         r[u] = [cast(str,r[x]) for x in u.src]
         continue
       if u.op is Ops.BUFFER and u.addrspace == AddrSpace.REG:
-        r[u] = [ssa("reg", u, self.types[u.dtype.scalar()]) for _ in range(u.max_numel())]
+        r[u] = [ssa("reg", u, self.types[u.dtype]) for _ in range(u.max_numel())]
         continue
       if u.op in {Ops.INDEX, Ops.SHRINK, Ops.LOAD} and u.src[0].addrspace in (AddrSpace.REG, AddrSpace.ALU):
         # on REG, INDEX/SHRINK pick the register (must be CONST) and LOAD is a noop
@@ -207,14 +207,14 @@ class PTXRenderer(Renderer):
         continue
       if u.op is Ops.SPECIAL: r[u] = "%" + u.arg
       elif u.op is Ops.LOAD:
-        r[u] = [ssa('val', dtype=self.types[u.dtype.scalar()]) for _ in range(u.max_numel())] if u.max_numel() > 1 else ssa('val', u)
+        r[u] = [ssa('val', dtype=self.types[u.dtype]) for _ in range(u.max_numel())] if u.max_numel() > 1 else ssa('val', u)
       elif u.op is Ops.PARAM: bufs.append((f"data{u.arg.slot}", u))
       elif u.op is Ops.WMMA:
         # registers for packing/unpacking input and acc
-        self.wmma_r = [[ssa("wmma_in", dtype="b32") for _ in range(0, len(r[u.src[0]]), 4 // u.src[0].dtype.scalar().itemsize)],
-                       [ssa("wmma_in", dtype="b32") for _ in range(0, len(r[u.src[1]]), 4 // u.src[0].dtype.scalar().itemsize)],
-                       [ssa("wmma_acc", dtype="b32") for _ in range(0, len(r[u.src[2]]), 4 // u.dtype.scalar().itemsize)]]
-        r[u] = [ssa("wmma", dtype=self.types[u.dtype.scalar()]) for _ in range(u.max_numel())]
+        self.wmma_r = [[ssa("wmma_in", dtype="b32") for _ in range(0, len(r[u.src[0]]), 4 // u.src[0].dtype.itemsize)],
+                       [ssa("wmma_in", dtype="b32") for _ in range(0, len(r[u.src[1]]), 4 // u.src[0].dtype.itemsize)],
+                       [ssa("wmma_acc", dtype="b32") for _ in range(0, len(r[u.src[2]]), 4 // u.dtype.itemsize)]]
+        r[u] = [ssa("wmma", dtype=self.types[u.dtype]) for _ in range(u.max_numel())]
       prefix, dtype = {Ops.CAST: ("cast", None), Ops.BITCAST: ("cast", None), Ops.END: ("pred", "pred"), Ops.RANGE: ("ridx", None),
         Ops.CONST: ("const", None), Ops.BUFFER: ("local", "u64"), Ops.INDEX: ("bidx", "u64"), Ops.SHRINK: ("bidx", "u64"),
         Ops.PARAM: ("dat", "u64" if u.addrspace is AddrSpace.GLOBAL else None), **{op: ("alu", None) for op in GroupOp.ALU}}.get(u.op, (None, None))

tinygrad/renderer/wgsl.py

@@ -50,8 +50,9 @@ wgsl_matcher = PatternMatcher([
   (UPat.store(UPat.var("b"), UPat.var("var"), name="s"), lambda b,var,s: packed_store(b,var) if is_packed(s) else None),
   (UPat.var("a") << UPat.var("b"),lambda a,b:(a.bitcast(dtypes.uint32)<<b.cast(dtypes.uint32)).bitcast(a.dtype) if b.dtype!=dtypes.uint32 else None),
   (UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
-  # fix nan check: 'a != a -> is_nan()'
-  (UPat.var("a") != UPat.var("a"), is_nan),
+  # fix nan check: 'a != a -> is_nan()'. the decomp rewrites (a != a).logical_not() to CMPEQ, so match both forms
+  (UPat.var("a", dtypes.floats) != UPat.var("a"), is_nan),
+  (UPat.var("a", dtypes.floats).alu(Ops.CMPEQ, UPat.var("a")), lambda a: is_nan(a).ne(True)),
   ])
 
 class WGSLRenderer(CStyleLanguage):

tinygrad/runtime/ops_amd.py

@@ -944,9 +944,7 @@ class AMDDevice(HCQCompiled):
   def is_usb(self) -> bool: return isinstance(self.iface, USBIface)
 
   def __init__(self, device:str=""):
-    self.device_id = int(device.split(":")[1]) if ":" in device else 0
-
-    self.iface = self._select_iface()
+    self.iface = self._select_iface(device)
 
     self.target:tuple[int, ...] = ((trgt:=self.iface.props['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
     self.arch = "gfx%d%x%x" % self.target

tinygrad/runtime/ops_cl.py

@@ -24,10 +24,10 @@ class CLCompiler(Compiler):
     super().__init__(f"compile_cl_{compile_key}")
   def compile(self, src:str) -> bytes:
     program = checked(cl.clCreateProgramWithSource(self.dev.context, 1, to_char_p_p([src.encode()]), None, status := ctypes.c_int32()), status)
-    build_status: int = cl.clBuildProgram(program, 1, self.dev.device_id, None, BP_CB(), None)
+    build_status: int = cl.clBuildProgram(program, 1, self.dev.cl_dev, None, BP_CB(), None)
     if build_status != 0:
-      cl.clGetProgramBuildInfo(program, self.dev.device_id, cl.CL_PROGRAM_BUILD_LOG, 0, None, log_size := ctypes.c_size_t())
-      cl.clGetProgramBuildInfo(program, self.dev.device_id, cl.CL_PROGRAM_BUILD_LOG,
+      cl.clGetProgramBuildInfo(program, self.dev.cl_dev, cl.CL_PROGRAM_BUILD_LOG, 0, None, log_size := ctypes.c_size_t())
+      cl.clGetProgramBuildInfo(program, self.dev.cl_dev, cl.CL_PROGRAM_BUILD_LOG,
                                log_size.value, mstr := ctypes.create_string_buffer(log_size.value), None)
       raise CompileError(f"OpenCL Compile Error\n\n{mstr.value.decode()}")
     check(cl.clGetProgramInfo(program, cl.CL_PROGRAM_BINARY_SIZES, ctypes.sizeof(ctypes.c_size_t), binary_sizes := (ctypes.c_size_t * 1)(), None))
@@ -39,11 +39,11 @@ class CLCompiler(Compiler):
 class CLProgram(Program['CLDevice']):
   def __init__(self, device:CLDevice, obj:TinyELF):
     self.dev, self.lib, self.signature = device, device.cl_compiler.compile_cached(obj.lib.decode()), obj.signature
-    self.program = checked(cl.clCreateProgramWithBinary(device.context, 1, device.device_id, (ctypes.c_size_t * 1)(len(self.lib)),
+    self.program = checked(cl.clCreateProgramWithBinary(device.context, 1, device.cl_dev, (ctypes.c_size_t * 1)(len(self.lib)),
                                                         to_char_p_p([self.lib], ctypes.c_ubyte), binary_status := ctypes.c_int32(),
                                                         errcode_ret := ctypes.c_int32()), errcode_ret)
     check(binary_status.value)
-    check(cl.clBuildProgram(self.program, 1, device.device_id, None, BP_CB(), None)) # NOTE: OSX requires this
+    check(cl.clBuildProgram(self.program, 1, device.cl_dev, None, BP_CB(), None)) # NOTE: OSX requires this
     self.kernel = checked(cl.clCreateKernel(self.program, obj.name.encode(), status := ctypes.c_int32()), status)
 
   def __del__(self):
@@ -101,17 +101,17 @@ class CLDevice(Compiled):
       CLDevice.device_ids = c.init_c_var((cl.cl_device_id * num_devices.value),
                                          lambda x: check(cl.clGetDeviceIDs(platform_ids[0], device_type, num_devices, x, None)))
 
-    self.device_id = CLDevice.device_ids[0 if ":" not in device else int(device.split(":")[1])]
-    self.device_name = (cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_NAME, 256,
+    self.cl_dev = CLDevice.device_ids[0 if ":" not in device else int(device.split(":")[1])]
+    self.device_name = (cl.clGetDeviceInfo(self.cl_dev, cl.CL_DEVICE_NAME, 256,
                                            buf:=ctypes.create_string_buffer(256), None), buf.value.decode())[1]
-    self.driver_version = (cl.clGetDeviceInfo(self.device_id, cl.CL_DRIVER_VERSION, 256,
+    self.driver_version = (cl.clGetDeviceInfo(self.cl_dev, cl.CL_DRIVER_VERSION, 256,
                                               buf:=ctypes.create_string_buffer(256), None), buf.value.decode())[1]
     if DEBUG >= 1: print(f"CLDevice: opening {self.device_name} with version {self.driver_version}")
-    self.context = checked(cl.clCreateContext(None, 1, self.device_id, CC_CB(), None, status := ctypes.c_int32()), status)
-    self.queue = checked(cl.clCreateCommandQueue(self.context, self.device_id, cl.CL_QUEUE_PROFILING_ENABLE, status), status)
+    self.context = checked(cl.clCreateContext(None, 1, self.cl_dev, CC_CB(), None, status := ctypes.c_int32()), status)
+    self.queue = checked(cl.clCreateCommandQueue(self.context, self.cl_dev, cl.CL_QUEUE_PROFILING_ENABLE, status), status)
     self.pending_copyin: list[memoryview] = []
-    check(cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_EXTENSIONS, 0, None, ctypes.byref(exts_len:=ctypes.c_size_t())))
-    self.device_exts = (cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_EXTENSIONS, exts_len.value,
+    check(cl.clGetDeviceInfo(self.cl_dev, cl.CL_DEVICE_EXTENSIONS, 0, None, ctypes.byref(exts_len:=ctypes.c_size_t())))
+    self.device_exts = (cl.clGetDeviceInfo(self.cl_dev, cl.CL_DEVICE_EXTENSIONS, exts_len.value,
                                            ctypes.byref(buf := ctypes.create_string_buffer(exts_len.value)), None),
                                            ctypes.string_at(buf).decode().split())[1]
 
@@ -119,7 +119,7 @@ class CLDevice(Compiled):
 
     arch = ",".join(self.device_exts)
     if "cl_khr_image2d_from_buffer" in self.device_exts:
-      check(cl.clGetDeviceInfo(self.device_id, cl.CL_DEVICE_IMAGE_PITCH_ALIGNMENT, 4, ctypes.byref(ipa := ctypes.c_uint32()), None))
+      check(cl.clGetDeviceInfo(self.cl_dev, cl.CL_DEVICE_IMAGE_PITCH_ALIGNMENT, 4, ctypes.byref(ipa := ctypes.c_uint32()), None))
       arch += f",IMAGE_PITCH_ALIGNMENT={ipa.value}"
     super().__init__(device, CLAllocator(self), [OpenCLRenderer], CLProgram, arch=arch)
 

tinygrad/runtime/ops_cpu.py

@@ -13,7 +13,7 @@ from tinygrad.renderer.isa.x86 import X86Renderer
 from tinygrad.runtime.support.elf import jit_loader
 from tinygrad.runtime.autogen import libc
 from tinygrad.codegen import do_to_program
-from tinygrad.engine.realize import pm_flatten_linear, get_call_arg_uops, get_runtime
+from tinygrad.engine.realize import pm_flatten_linear, get_call_arg_uops, get_call_var_uops, get_runtime
 from tinygrad import UOp, dtypes
 from tinygrad.dtype import AddrSpace
 from tinygrad.uop.ops import KernelInfo, Ops, UPat, PatternMatcher, graph_rewrite
@@ -64,7 +64,7 @@ def cpu_cmd(devs:tuple[str, ...], prog, *args:UOp) -> UOp:
   return UOp(Ops.INS, dtypes.void, words + (UOp.const(0, dtypes.uint64),) * (CMD_SIZE - len(words)), arg="cmd")
 
 def cpu_exec(ctx:tuple[str, ...], call:UOp, prg:UOp) -> UOp:
-  args = [get_call_arg_uops(call)[i].getaddr(ctx) for i in prg.arg.globals] + [v.cast(dtypes.uint64) for v in prg.arg.vars]
+  args = [get_call_arg_uops(call)[i].getaddr(ctx) for i in prg.arg.globals] + [v.cast(dtypes.uint64) for v in get_call_var_uops(call, prg)]
   if (core:=prg.arg.runtimevars.get('core_id')) is None: return cpu_cmd(ctx, prg, *args)
 
   la = [cpu_cmd(ctx,prg,*args[:(cid:=(len(prg.arg.globals)+core))],UOp.const(t, dtypes.uint64),*args[cid+1:]) for t in range(prg.arg.global_size[0])]
@@ -99,10 +99,11 @@ def encode_queue(q:UOp) -> UOp:
   e = UOp.range(cnt, next(UOp.unique_num), dtype=dtypes.int, src=(cmdbuf, ring))
   copy = UOp.group(*[ring.index((base + e*CMD_SIZE + w) % ring_words).store(cmdbuf.index(e*CMD_SIZE + w).load()) for w in range(CMD_SIZE)])
 
-  # wake the worker after each entry, keeping the post with the stores stops it from hoisting out of the loop
-  wake = copy.end(e) if WIN else make_signal(devs, tag="func:sem_post").after(copy).index(0).load().call(sem.index(0), ret_dtype=dtypes.void).end(e)
-  bumped = put.after(wake).index(0).store(put.index(0).load() + cnt)
-  return sysbuf.after(bumped).index(0).store(put.index(0).load() + cnt) if WIN else bumped
+  bumped = put.after(copy.end(e)).index(0).store(put.index(0).load() + cnt)
+  if WIN: return sysbuf.after(bumped).index(0).store(put.after(bumped).index(0).load())
+
+  e = UOp.range(cnt, next(UOp.unique_num), dtype=dtypes.int, src=(bumped,))
+  return make_signal(devs, tag="func:sem_post").after(e).index(0).load().call(sem.after(e).index(0), ret_dtype=dtypes.void).end(e)
 
 # *****************
 
@@ -196,11 +197,13 @@ class CPUDevice(HCQ2Compiled):
   pm_lower = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),)), encode_queue)])
 
   def __init__(self, device:str=""):
+    self.workers:list[CPUWorker] = []
     super().__init__(device, CPUAllocator(self), [ClangRenderer, CPULLVMRenderer, LVPRenderer, X86Renderer], CPUProgram,
       arch={'amd64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine().lower(), m)+",native")
 
     self.pm_bufferize = PatternMatcher(
-      [(UPat(Ops.PARAM, tag=f"COMPUTE:0_{n}"), lambda ctx, n=n: getattr(ctx[0].worker, n)) for n in ("ring", "put", "sem", "sys", "done")] +
+      [(UPat(Ops.PARAM, tag=f"{q}_{n}"), lambda ctx, q=q,n=n: getattr(ctx[0].worker(q), n))
+       for q in ("COMPUTE:0", "SUBMIT:0") for n in ("ring", "put", "sem", "sys", "done")] +
       [(UPat(Ops.PARAM, tag=f"func:{f}"), lambda ctx, f=f: ctx[0].func_ptr(f)) for f in FUNCS]) + self.pm_bufferize
 
     with Context(EMULATED_DTYPES="", TRACK_MATCH_STATS=0):
@@ -210,6 +213,12 @@ class CPUDevice(HCQ2Compiled):
 
   def func_ptr(self, name:str) -> Buffer: return self.func_table.view(1, dtypes.uint64, FUNCS.index(name)*8).ensure_allocated()
 
+  def synchronize(self, timeout:int|None=None):
+    for worker in self.workers:
+      put, done = (getattr(worker, x)._buf.cpu_view().view(fmt='Q') for x in ("put", "done"))
+      while done[0] < put[0]: self._wait_signal(done, put[0], timeout)
+    super().synchronize(timeout)
+
   @functools.cached_property
   def func_table(self) -> Buffer:
     lib = ctypes.windll.kernel32 if sys.platform == "win32" else libc.dll # type: ignore[attr-defined]
@@ -217,8 +226,8 @@ class CPUDevice(HCQ2Compiled):
       array.array('Q', [unwrap(ctypes.cast(getattr(lib, f), ctypes.c_void_p).value) for f in FUNCS])
     return ft
 
-  @functools.cached_property
-  def worker(self) -> CPUWorker:
+  @functools.cache
+  def worker(self, queue:str) -> CPUWorker:
     ring, put, sysbuf, done = (Buffer(self.device, sz, dtypes.uint64, preallocate=True) for sz in (RING_SLOTS*CMD_SIZE, 1, 1, 1))
     addr, hsem = 0, None
 
@@ -230,5 +239,6 @@ class CPUDevice(HCQ2Compiled):
     sem = Buffer(self.device, 1, dtypes.uint64, options=BufferSpec(external_ptr=addr), preallocate=True)
 
     worker_args = [ring._buf.va_addr, sysbuf._buf.va_addr if WIN else self.func_ptr('sem_wait')._buf.va_addr, done._buf.va_addr, addr]
-    (worker:=threading.Thread(target=self.prgs[worker_prog].fxn, daemon=True, args=[ctypes.c_uint64(x) for x in worker_args])).start()
-    return CPUWorker(ring, put, sem, sysbuf, done, worker)
+    (thread:=threading.Thread(target=self.prgs[worker_prog].fxn, daemon=True, args=[ctypes.c_uint64(x) for x in worker_args])).start()
+    self.workers.append(worker:=CPUWorker(ring, put, sem, sysbuf, done, thread))
+    return worker

tinygrad/runtime/ops_disk.py

@@ -49,7 +49,8 @@ class DiskDevice(Compiled):
     DiskDevice._tried_io_uring_init = True
 
     if sys.platform == 'linux' and not hasattr(sys, "getandroidapilevel"):
-      fd = libc.syscall(io_uring.NR_io_uring_setup, 4096, ctypes.byref(p:=io_uring.struct_io_uring_params()))
+      p = io_uring.struct_io_uring_params(flags=io_uring.IORING_SETUP_SQPOLL, sq_thread_idle=0xffffffff)
+      fd = libc.syscall(io_uring.NR_io_uring_setup, 4096, ctypes.byref(p))
       if fd < 0: return
 
       sq_ptr = libc.mmap(0, p.sq_off.array + p.sq_entries * 4, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | MAP_POPULATE, fd, 0)
@@ -67,6 +68,7 @@ class DiskDevice(Compiled):
         kring_mask=u32ptr(sq_ptr+p.cq_off.ring_mask), cqes=ctypes.cast(cq_ptr+p.cq_off.cqes, ctypes.POINTER(io_uring.struct_io_uring_cqe)))
 
       DiskDevice.io_uring = io_uring.struct_io_uring(ring_fd=fd, sq=sqdesc, cq=cqdesc) # type: ignore
+      libc.syscall(io_uring.NR_io_uring_enter, fd, 0, 0, io_uring.IORING_ENTER_SQ_WAKEUP)
 
 class DiskBuffer:
   def __init__(self, device:DiskDevice, size:int, offset=0):
@@ -124,7 +126,6 @@ class DiskAllocator(Allocator):
         # Send sqe
         DiskDevice.io_uring.sq.array[sqe_index] = sqe_index
         DiskDevice.io_uring.sq.ktail[0] = tail + 1
-        libc.syscall(io_uring.NR_io_uring_enter, DiskDevice.io_uring.ring_fd, 1, 1, io_uring.IORING_ENTER_GETEVENTS)
 
         reqs.append((copy_batch, copied_in, minor_offset, real_copy_size:=min(sqe.len - minor_offset, size - copied_in)))
         next_read_offset += sqe.len

tinygrad/runtime/ops_nv.py

@@ -588,8 +588,7 @@ class NVDevice(HCQCompiled[NVSignal]):
   def is_nvd(self) -> bool: return isinstance(self.iface, PCIIface)
 
   def __init__(self, device:str=""):
-    self.device_id = int(device.split(":")[1]) if ":" in device else 0
-    self.iface = self._select_iface()
+    self.iface = self._select_iface(device)
 
     device_params = nv_gpu.NV0080_ALLOC_PARAMETERS(deviceId=self.iface.gpu_instance, hClientShare=self.iface.root,
                                                    vaMode=nv_gpu.NV_DEVICE_ALLOCATION_VAMODE_OPTIONAL_MULTIPLE_VASPACES)

tinygrad/runtime/ops_rdma.py

@@ -101,6 +101,5 @@ class RDMAAllocator(HCQAllocatorBase):
 
 class RDMADevice(HCQCompiled):
   def __init__(self, device:str=""):
-    self.device_id = int(device.split(":")[1]) if ":" in device else 0
-    self.iface = MLXIface(self, self.device_id)
+    self.iface = MLXIface(self, int(device.split(":")[1]) if ":" in device else 0)
     super().__init__(device, RDMAAllocator(self), [], None, signal_t=None)

tinygrad/runtime/support/hcq.py

@@ -1,10 +1,9 @@
 from __future__ import annotations
 from typing import cast, Callable, Type, TypeVar, Generic, Any
-import contextlib, decimal, statistics, time, ctypes, array, os, struct, collections, functools, itertools
-from dataclasses import replace
+import contextlib, decimal, statistics, time, ctypes, array, os, struct, collections, itertools
 try: import fcntl # windows misses that
 except ImportError: fcntl = None #type:ignore[assignment]
-from tinygrad.helpers import DEV, PROFILE, getenv, to_mv, from_mv, cpu_profile, ProfileRangeEvent, select_first_inited, select_by_name, unwrap
+from tinygrad.helpers import DEV, PROFILE, getenv, to_mv, from_mv, cpu_profile, ProfileRangeEvent, unwrap
 from tinygrad.helpers import suppress_finalizing, pluralize, TracingKey
 from tinygrad.device import Device, BufferSpec, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent, Program, TinyELF
 from tinygrad.uop.ops import sym_infer, sint, UOp
@@ -393,8 +392,6 @@ class HCQCompiled(Compiled, Generic[SignalType]):
   def __init__(self, device:str, allocator:HCQAllocatorBase, compilers:list[type[Renderer]], runtime:type[Program]|None,
                signal_t:Type[SignalType]|None=None, comp_queue_t:Callable[..., HWQueue]|None=None, copy_queue_t:Callable[..., HWQueue]|None=None,
                kernargs_size=(16 << 20), sigalloc_size=0x1000, can_recover:bool=False, arch=None):
-    self.device_id:int = int(device.split(":")[1]) if ":" in device else 0
-
     from tinygrad.runtime.graph.hcq import HCQGraph
     super().__init__(device, allocator, compilers, runtime, HCQGraph, arch=arch)
 
@@ -424,8 +421,6 @@ class HCQCompiled(Compiled, Generic[SignalType]):
 
     if self._is_cpu(): HCQCompiled.cpu_devices.append(self)
 
-  def count(self) -> int: return self.iface.count if hasattr(self, 'iface') else 1
-
   def synchronize(self, timeout:int|None=None):
     if self.error_state is not None: raise self.error_state
     if not hasattr(self, 'timeline_signal'): return
@@ -491,16 +486,6 @@ class HCQCompiled(Compiled, Generic[SignalType]):
       buf, realloced = self.allocator.alloc(oldbuf.size if oldbuf is not None else new_size, options=options), False
     return buf, realloced
 
-  def _select_iface(self):
-    assert (v:=getenv(k:=f'{type(self).__name__[:-6].upper()}_IFACE', "")) == "",  \
-      f"{k}={v} is deprecated, use DEV={replace(DEV.target(type(self).__name__[:-6]), interface=v)} instead"
-    assert hasattr(self, "ifaces"), "must have ifaces to select an iface"
-    t = DEV.target(dev:=type(self).__name__[:-6])
-    filtered = select_by_name(self.ifaces, lambda i: i.__name__[:-5], t.interface, f"{dev} has no interface {t.interface!r}")
-    filtered = [i for i in filtered if t.interface.startswith("MOCK") or not i.__name__[:-5].startswith("MOCK")] # never fallback to mock ifaces
-    return select_first_inited([functools.partial(cast(Callable, iface), self, self.device_id) for iface in filtered],
-                               f"No interface for {dev}:{self.device_id} is available")
-
   def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"
 
   def rdma_dev(self):
@@ -512,9 +497,7 @@ class HCQCompiled(Compiled, Generic[SignalType]):
   def finalize(self):
     try: self.synchronize() # Try to finalize device in any case.
     except RuntimeError as e: print(f"{self.device} synchronization failed before finalizing: {e}")
-
-    # If the device has an interface, call its device_fini method to clean up resources.
-    if hasattr(self, 'iface') and hasattr(self.iface, 'device_fini'): self.iface.device_fini()
+    super().finalize()
 
 class HCQBuffer:
   def __init__(self, va_addr:sint, size:int, meta:Any=None, _base:HCQBuffer|None=None, view:MMIOInterface|None=None, owner:Any=None):

tinygrad/runtime/support/hcq2.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
-from typing import cast, Callable, TypeVar, Generic, Any, Sequence, Iterable
+from typing import cast, TypeVar, Generic, Any, Sequence, Iterable
 import struct, functools, time, collections, itertools, decimal, statistics
 from dataclasses import replace, dataclass
-from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, dedup, pluralize, JIT_BATCH_SIZE, unwrap, PROFILE
+from tinygrad.helpers import suppress_finalizing, dedup, pluralize, JIT_BATCH_SIZE, unwrap, PROFILE
 from tinygrad.helpers import to_tuple, round_up, partition, data64_le, panic, ContextVar, perf_counter_us, Context
 from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer, DepsTracker
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEntry, ProfileGraphEvent
@@ -23,17 +23,16 @@ HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQ2Compiled')
 HCQ_RUNTIME_DEV = ContextVar("HCQ_RUNTIME_DEV", "CPU")
 
 HCQ_DEVS = frozenset(("AMD", "CPU"))
-HCQ_P2P_DEVS = HCQ_DEVS | frozenset(("CPU",))
-HCQ_CACHE_TAGS = frozenset(("program", "systems", "template"))
+HCQ_CACHE_TAGS = frozenset(("program", "systems"))
 
 @dataclass(frozen=True)
 class HCQInfo:
   device:tuple[str, ...]
   estimates:Estimates = Estimates()
 
-  input_idxs:tuple[int, ...] = () # indexes into input_uops used by this call
-  inputs:int|None = None
-  kernels:tuple[tuple[tuple[str, ...], str, Estimates, tuple[int, ...]], ...] = ()
+  input_idxs:tuple[tuple[tuple[str, ...], tuple[int, ...]], ...] = () # per inputs table: (devices, indexes into input_uops)
+  inputs:int|None = None # index of the inputs table in call.src
+  kernels:tuple[tuple[tuple[str, ...], UOp, tuple[int, ...]], ...] = () # per kernel: (devices, a call carrying its name and estimates, timestamps)
 
 def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c
 
@@ -94,10 +93,10 @@ pm_replace_buffers = PatternMatcher([(UPat(Ops.CALL, name="call"), replace_call_
 # *****************
 # 1.1. prep: staging copies
 
-def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS - {"CPU"}) and not all_devices_in(b.device, HCQ_P2P_DEVS)
+def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS - {"CPU"}) and not all_devices_in(b.device, HCQ_DEVS)
 
 def _get_enqueue_devs(call:UOp) -> Any|None:
-  if not (bufs:=call.src[1:]) or not all(all_devices_in(b.device, HCQ_P2P_DEVS) for b in bufs): return None
+  if not (bufs:=call.src[1:]) or not all(all_devices_in(b.device, HCQ_DEVS) for b in bufs): return None
   devs = min(bufs, key=lambda b: to_tuple(b.device)[0].startswith("CPU")).device # prio to enqueue on not CPU device
   return devs if all_devices_in(devs, HCQ_DEVS) else None
 
@@ -216,7 +215,7 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> li
     # and make hcq call
     name, info = get_call_name(call, get_call_arg_uops(call)), HCQInfo(devices, estimate_uop(call))
     ts_ids = [next(UOp.unique_num) for _ in range(2)] if profile else []
-    kerns.append((devices, name, info.estimates, tuple(ts_ids)))
+    kerns.append((devices, make_call(name, call.src[0], info), tuple(ts_ids)))
 
     ts_ins = [UOp(Ops.INS, arg="timestamp", src=(make_signal(devices, s),)) for s in ts_ids]
     q += ts_ins[:1] + [call.replace(arg=replace(call.arg, aux=info))] + ts_ins[1:]
@@ -346,14 +345,11 @@ def split_patches(call:UOp) -> UOp|None:
   scatter = make_scatter_loops(input_patches, tables[0], lt_patches)
   body = body.substitute({p:p.substitute(scatter | reads) for p in rt_patches})
 
-  if inputs: # fence inputs
-    fills.append((t:=tables[0][0]).after(make_binary_patch(t, bytes(t.max_numel() * 8)))) # zeroed at link, slot 0 is the host fence
-    body = body.replace(src=(UOp.sink(*body.src[0].src, t.after(*body.src[0].src).index(0).store(0)),)) # open it once consumed
-
   lt_srcs = collections.defaultdict(list)
   for p in lt_patches: lt_srcs[p.buf_uop].append(p)
   return call.replace(src=(body, *call.src[1:], *[b.after(*ps) for b,ps in lt_srcs.items()], *fills),
-    arg=replace(call.arg, aux=replace(call.arg.aux, input_idxs=tuple(sorted(dedup(b.arg.slot for g in inputs for b in unwrap_mstack(g.buf_uop)))))))
+    arg=replace(call.arg, aux=replace(call.arg.aux, input_idxs=((call.arg.aux.device,
+      tuple(sorted(dedup(b.arg.slot for g in inputs for b in unwrap_mstack(g.buf_uop))))),) if inputs else call.arg.aux.input_idxs)))
 pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), split_patches)])
 
 # *****************
@@ -372,14 +368,13 @@ def replace_params(call:UOp) -> UOp|None:
 
   # keep buffers whose addresses become link-time constants alive and mapped
   held = args + [r.without_after for r in refhold]
-  addrs = dedup([g.src[0].without_after for x in call.src for g in x.toposort() if g.op is Ops.GETADDR])
+  addrs = dedup([g.src[0].without_after for g in call.toposort() if g.op is Ops.GETADDR])
   refhold += [a for a in addrs if a not in held and all(b.op is not Ops.PARAM or b.tag is not None for b in unwrap_mstack(a))]
 
   sub = {(b:=u.without_after): UOp.param(i, u.dtype, shape=b.shape, device=HCQ_RUNTIME_DEV.value, volatile=b.op is Ops.PARAM and b.arg.volatile)
          for i,u in enumerate(c_args)} | {v: v.replace(arg=replace(v.arg, slot=-1)) for v in variables if v.op is Ops.PARAM} | _rank_ranges(tops)
-  info = replace(call.arg.aux, inputs=next((i for i,u in enumerate(c_args) if u.without_after.tag == "inputs"), None))
-  return call.replace(src=(body.substitute(sub).replace(arg="hcq_args"), *c_args, *refhold),
-                      arg=replace(call.arg, aux=info)) # TODO: call.after(*refhold)?
+  info = replace(call.arg.aux, inputs=next((i for i,u in enumerate(c_args + refhold) if u.without_after.tag == "inputs"), None))
+  return call.replace(src=(body.substitute(sub).replace(arg="hcq_args"), *c_args, *refhold), arg=replace(call.arg, aux=info))
 pm_replace_params = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), replace_params)])
 
@@ -425,8 +420,44 @@ def callify_hcq(call:UOp, cf:UOp) -> UOp:
 pm_callify_hcq = PatternMatcher([(UPat(Ops.CALL, src=(
   UPat(Ops.CUSTOM_FUNCTION, arg="hcq_args", src=(UPat(Ops.SINK),), name="cf"),), name="call", allow_any_len=True), callify_hcq)])
 
+# *****************
+# 9. merge submitters
+
+def _lane_arg(a:UOp, lane:int, table:UOp) -> UOp: return table if a.tag == "inputs" else a.mselect(lane) if len(to_tuple(a.device)) > 1 else a
+
+def merge_batch(batch:list[UOp]) -> UOp:
+  tables = UOp.variable("hcq_inputs_ptr", 0, 2**64-1, dtypes.uint64, param=True)
+  lanes = [(c, j, sum(len(idxs) * 8 for _, idxs in c.arg.aux.input_idxs)) for c in batch for j in range(len(c.arg.aux.device))] # (call, lane, bytes)
+  offs = itertools.accumulate((table_bytes for _, _, table_bytes in lanes), initial=0) # every lane owns the next table of the region
+  cmds = [c.src[0].src[0].call(*[_lane_arg(a.without_after, j, tables + off) for a in c.src[1:]], UOp.variable("_device_num", 0, 1 << 30).bind(j))
+          for (c, j, _), off in zip(lanes, offs)]
+
+  info = HCQInfo((HCQ_RUNTIME_DEV.value,), sum((c.arg.aux.estimates for c in batch), start=Estimates()),
+                 input_idxs=tuple(x for c in batch for x in c.arg.aux.input_idxs), kernels=tuple(k for c in batch for k in c.arg.aux.kernels))
+  body = UOp.custom_function("hcq", make_submit(*cmds, devs=HCQ_RUNTIME_DEV.value, queue="SUBMIT:0").sink())
+  return body.call(*[s for c in batch for s in c.src[1:] if s.without_after.tag != "inputs"], name=f"hcq_submitter ({len(batch)})", aux=info)
+
+def merge_submitters(linear:UOp) -> UOp:
+  batches = [(k, list(g)) for k, g in itertools.groupby(linear.src, key=lambda c: isinstance(c.arg.aux, HCQInfo))]
+  return linear.replace(src=tuple(c for is_hcq, b in batches for c in ([merge_batch(b)] if is_hcq else b)))
+
+# *****************
+# hcq schedule
+
 hcq_compile_cache:dict[tuple[bytes, bool], UOp] = {}
 
+def hcq_lower(linear:UOp, pm_encode:PatternMatcher) -> UOp:
+  # lowering to hcq ir
+  linear = graph_rewrite(linear, pm_encode, walk=True, name="encode and pack", enter_calls=True)
+
+  # patches and runtime uops
+  linear = graph_rewrite(linear, pm_early_simplify+symbolic, bottom_up=False, name="simplify patches", enter_calls=True)
+  linear = graph_rewrite(linear, pm_split_patches, walk=True, name="split patches")
+
+  # and compile it
+  linear = graph_rewrite(linear, pm_replace_params, name="replace params")
+  return graph_rewrite(linear, pm_callify_hcq, name="callify hcq", enter_calls=True)
+
 @rewrite_group(lambda linear,input_uops,profile,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
 def hcq_compile(linear:UOp, input_uops:list[UOp]|None, profile:bool) -> UOp:
   if input_uops is not None:
@@ -441,16 +472,9 @@ def hcq_compile(linear:UOp, input_uops:list[UOp]|None, profile:bool) -> UOp:
     # schedule
     linear = graph_rewrite(linear, pm_schedule_and_merge, ctx=({s:p for p,s in back_map.items()}, profile), walk=True, name="schedule and merge hcq")
 
-    # lowering to hcq ir
-    linear = graph_rewrite(linear, pm_encode_cmdbufs+pm_pack_placeholders, walk=True, name="encode and pack", enter_calls=True)
-
-    # patches and runtime uops
-    linear = graph_rewrite(linear, pm_early_simplify+symbolic, bottom_up=False, name="simplify patches", enter_calls=True)
-    linear = graph_rewrite(linear, pm_split_patches, walk=True, name="split patches")
-
-    # and compile it
-    linear = graph_rewrite(linear, pm_replace_params, name="replace params")
-    final_linear = hcq_compile_cache[cache_key] = graph_rewrite(linear, pm_callify_hcq, name="callify hcq", enter_calls=True)
+    # lower to hcq programs, then pack the programs of every batch into one C submitter (needs a C runtime device for the program addresses)
+    linear = hcq_lower(linear, pm_encode_cmdbufs+pm_pack_placeholders)
+    final_linear = hcq_compile_cache[cache_key] = hcq_lower(merge_submitters(linear), pm_encode_cmdbufs) if HCQ_RUNTIME_DEV.value == "CPU" else linear
 
   return final_linear
 
@@ -467,7 +491,7 @@ pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, name="buf"), bufferize_buf)])
 # 7. resolve patches
 
 def push_stack(op, s): return UOp(Ops.STACK,
-  src=tuple(op.replace(dtype=op.dtype.scalar(), src=tuple(x if y is s else y for y in op.src)) for x in s.src))
+  src=tuple(op.replace(dtype=op.dtype, src=tuple(x if y is s else y for y in op.src)) for x in s.src))
 
 def fold_binary(buf:UOp, blob:UOp) -> UOp:
   for b in (m.bufs if isinstance(m:=buf.buffer, MultiBuffer) else (m,)):
@@ -532,7 +556,6 @@ class HCQ2Compiled(Compiled):
   wait_timeout_ms: float = 30000.0
 
   def __init__(self, device:str, allocator:HCQAllocator, compilers:list[type[Renderer]], runtime, can_recover:bool=False, arch=None):
-    self.device_id:int = int(device.split(":")[1]) if ":" in device else 0
     self.can_recover = can_recover
 
     self.pm_bufferize = PatternMatcher([
@@ -545,7 +568,6 @@ class HCQ2Compiled(Compiled):
 
     super().__init__(device, allocator, compilers, runtime, None, arch=arch)
 
-    self.rt_buffer = Buffer(self.device, 64 << 20, dtypes.uint8, options=BufferSpec(uncached=True, cpu_access=True))
     self.rt_allocator = BumpAllocator(64 << 20)
     self.prof_ents:dict[int, ProfileGraphEntry] = {}
 
@@ -569,6 +591,10 @@ class HCQ2Compiled(Compiled):
       tdiffs.append((st+perf_counter_us())/2 - gpu)
     Compiled.profile_events.append(ProfileDeviceEvent(self.device, statistics.median(tdiffs), self.device_props()))
 
+  @functools.cached_property
+  def rt_buffer(self) -> Buffer:
+    return Buffer(self.device, self.rt_allocator.size, dtypes.uint8, options=BufferSpec(uncached=True, cpu_access=True), preallocate=True)
+
   def new_buffer(self, b:UOp, cache:bool) -> Buffer:
     if cache 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))
@@ -580,42 +606,31 @@ class HCQ2Compiled(Compiled):
     buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = init_value
     return buf
 
+  def _wait_signal(self, sig:memoryview, value:int, timeout:int|None=None):
+    timeout = timeout if timeout is not None and self.can_recover else None
+    st, done = time.perf_counter(), sig[0]
+    while done < value:
+      if done != (done:=sig[0]): st = time.perf_counter()
+      elif time.perf_counter() - st > (timeout or self.wait_timeout_ms) / 1000: self.on_device_hang()
+
   def synchronize(self, timeout:int|None=None):
     if HCQ_RUNTIME_DEV.value != self.device: Device[HCQ_RUNTIME_DEV.value].synchronize()
 
     sig = self.signal("timeline").as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
     tl = self.signal("value", 1).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
-    timeout = timeout if timeout is not None and self.can_recover else None
-    st, done = time.perf_counter(), sig[0]
-    while done < tl[0] - 1:
-      if done != (done:=sig[0]): st = time.perf_counter()
-      elif time.perf_counter() - st > (timeout or self.wait_timeout_ms) / 1000: self.on_device_hang()
+    self._wait_signal(sig, tl[0] - 1, timeout)
     if self.prof_ents: self.collect_prof()
 
   def on_device_hang(self): raise RuntimeError(f"{self.device} hang detected")
 
   def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent.
 
-  def count(self) -> int: return self.iface.count if hasattr(self, 'iface') else 1
-
-  def _select_iface(self):
-    assert (v:=getenv(k:=f'{type(self).__name__[:-6].upper()}_IFACE', "")) == "",  \
-      f"{k}={v} is deprecated, use DEV={replace(DEV.target(type(self).__name__[:-6]), interface=v)} instead"
-    assert hasattr(self, "ifaces"), "must have ifaces to select an iface"
-    t = DEV.target(dev:=type(self).__name__[:-6])
-    filtered = select_by_name(self.ifaces, lambda i: i.__name__[:-5], t.interface, f"{dev} has no interface {t.interface!r}")
-    filtered = [i for i in filtered if t.interface.startswith("MOCK") or not i.__name__[:-5].startswith("MOCK")] # never fall back to mock ifaces
-    return select_first_inited([functools.partial(cast(Callable, iface), self, self.device_id) for iface in filtered],
-                               f"No interface for {dev}:{self.device_id} is available")
-
   def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"
 
   def finalize(self):
     try: self.synchronize() # try to finalize the device in any case
     except RuntimeError as e: print(f"{self.device} synchronization failed before finalizing: {e}")
-
-    # if the device has an interface, call device_fini to clean up resources
-    if hasattr(self, 'iface') and hasattr(self.iface, 'device_fini'): self.iface.device_fini()
+    super().finalize()
 
 @dataclass
 class HCQ2Buffer:

tinygrad/schedule/multi.py

@@ -22,14 +22,14 @@ def mstack_early_shrink(ms:UOp, shrink:UOp):
 
 def lower_broadcast_copy(c:UOp, x:UOp):
   if not (isinstance(c.device, tuple) and isinstance(x.device, str)): return None
-  if (sx:=x.simplify()).device is None and sx.base.op is Ops.CONST: return UOp(Ops.MSTACK, src=(sx,)*len(c.device))
+  if (sx:=x.simplify()).device is None: return UOp(Ops.MSTACK, src=(sx,)*len(c.device))
   return UOp(Ops.MSTACK, src=tuple(x.copy_to_device(d) for d in c.device))
 
 replace_allreduce = PatternMatcher([
   # BROADCAST: explicitly expand broadcast copies and combine with MSTACK
-  (UPat(Ops.COPY, name="c", src=(UPat(GroupOp.All-{Ops.CONST}, name="x"),)), lower_broadcast_copy),
+  (UPat(Ops.COPY, name="c", src=(UPat(name="x"),)), lower_broadcast_copy),
   # COPY_TO_ONE: if copying from multidevice to one, MSELECT the first (TODO: a little from each?)
-  (UPat(Ops.COPY, name="c", src=(UPat(GroupOp.All-{Ops.CONST}, name="x"),)), lambda c,x:
+  (UPat(Ops.COPY, name="c", src=(UPat(name="x"),)), lambda c,x:
     x.mselect(0).copy_to_device(c.device) if isinstance(c.device, str) and isinstance(x.device, tuple) else None),
   # MSELECT on MSTACK is replaced with nothing
   (UPat(Ops.MSELECT, src=(UPat(Ops.MSTACK, name="mstack"),), name="ms"), lambda mstack, ms: mstack.src[ms.arg]),

tinygrad/schedule/rangeify.py

@@ -4,7 +4,7 @@ import itertools
 from tinygrad.dtype import dtypes, AddrSpace, Invalid, to_dtype, strong_dtype
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg, shape_to_shape_arg
 from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, rewrite_group, identity_element
-from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const
+from tinygrad.uop.symbolic import symbolic
 from tinygrad.uop.movement import mop_cleanup
 from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS, SPEC
 from tinygrad.helpers import PCONTIG, FLOAT16, OPENPILOT_HACKS, argsort, partition, get_single_element
@@ -313,9 +313,9 @@ pm_const_buffer_folding = pm_mops+PatternMatcher([
    lambda idx,after: idx.const_like(Invalid) if after_all_invalid(after) else None),
   # hack if a noop turned to a const
   (UPat(Ops.NOOP, src=(UPat.cvar("c"),)), lambda c: c),
-  # mstack on CONST is CONST
-  (UPat(Ops.MSTACK, src=(UPat.var("s"),), allow_any_len=True).f(Ops.INDEX, allow_any_len=True),
-   lambda s: c if (c:=s.base).op is Ops.CONST else None),
+  # a deviceless MSTACK src is the same value on every device, so indexing the stack is just indexing that value
+  (UPat(Ops.MSTACK, src=(UPat.var("s"),), allow_any_len=True).f(Ops.INDEX, allow_any_len=True, name="idx"),
+   lambda s,idx: idx.replace(src=(s,)+idx.src[1:]) if s.device is None else None),
 ])
 
 pm_remove_bufferize = PatternMatcher([
@@ -584,7 +584,7 @@ def get_kernel_graph(sink:UOp) -> UOp:
   tsink, rctx = run_rangeify(tsink, bool(DEBUG_RANGEIFY))
 
   tsink = graph_rewrite(tsink,
-                        symbolic+pm_fold_cast_const+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize+pm_no_indexing_calls,
+                        symbolic+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize+pm_no_indexing_calls,
                         name="symbolic+reduce_collapse+debuf")
   tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")
 

tinygrad/tensor.py

@@ -665,13 +665,13 @@ class Tensor(RandMixin):
     ```
     """
     all_uops = self.uop.toposort()
-    # backward fills .grad for every in-scope non-CONST float tensor
+    # backward fills .grad for every in-scope float tensor with a device
     tensors_need_grad: list[Tensor] = [t for tref in all_tensors if (t:=tref()) is not None and \
-                                       t.uop in all_uops and t.is_floating_point() and t.uop.op is not Ops.CONST]
+                                       t.uop in all_uops and t.is_floating_point() and t.device is not None]
     # clear contexts
     for t,g in zip(tensors_need_grad, self.gradient(*tensors_need_grad, gradient=gradient)):
       assert g.shape == t.shape, f"grad shape must match tensor shape, {g.shape!r} != {t.shape!r}"
-      if g.device is None and t.device is not None: g = g.clone(device=t.device)
+      if g.device is None: g = g.clone(device=t.device)
       if t.grad is None: t.grad = g
       else: t.grad.assign(t.grad + g.to(t.grad.device))
     return self

tinygrad/uop/ops.py

@@ -242,9 +242,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   arg:Any = None
   tag:Any = None
   def __del__(self):
-    if Ops is not None and self.op is Ops.BUFFER and (buffer:=buffers.get(self)) is not None: buffer.ref(-1)
+    # NOTE: getattr because this object may be partially constructed (e.g. if __init__ raised, like the BEAM timeout SIGALRM)
+    if Ops is not None and getattr(self, 'op', None) is Ops.BUFFER and (buffer:=buffers.get(self)) is not None: buffer.ref(-1)
     try: del UOpMetaClass.ucache[(self.op, self.dtype, self.src, self.arg, self.tag)]
-    except AttributeError: pass
+    except (AttributeError, KeyError): pass
   def __reduce__(self):
     args = [self.op, self.dtype, self.src, self.arg, self.tag, self.metadata]
     if self.op is Ops.BUFFER and self.realized is not None: args.append(self.realized)
@@ -470,10 +471,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     if (ret:=self._shape) is None: raise RuntimeError(f"shape requested, but {self.op} doesn't have a shape")
     return ret
 
-  @property
-  def max_shape(self) -> tuple[int, ...]: return to_max_shape(self.shape)
-  def max_numel(self) -> int: return prod(self.max_shape)
-
   @property
   def shard_shape(self) -> tuple[sint, ...]:
     if not isinstance(self.device, tuple) or self.axis is None: return self.shape
@@ -587,9 +584,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   def const_like(self, b:ConstLike, dtype:DType|None=None):
     ret = UOp.const(b, dtype or self.dtype)
     return ret._mop(Ops.EXPAND, arg=self._shape) if self._shape and ret._shape != self._shape else ret
-  def vconst_like(self, b:ConstLike, dtype:DType|None=None):
+  def vconst_like(self, b:ConstLike):
     # for use after movement ops have been removed
-    return UOp.const(b, dtype or self.dtype).broadcast(self.max_numel())
+    return UOp.const(b, self.dtype).broadcast(self.max_numel())
   def ufix(self, x):
     if isinstance(x, UOp): return x
     return UOp.const(x)
@@ -1407,7 +1404,7 @@ class UPat(OpMixin):
     if self.is_any: return flatten([x.match(uop, store.copy()) for x in self.src[0]])
     if (self.op is not None and uop.op not in self.op) or \
        (self.name is not None and store.setdefault(self.name, uop) is not uop) or \
-       (self.match_dtype is not None and uop.dtype not in self.match_dtype and uop.dtype.scalar() not in self.match_dtype) or \
+       (self.match_dtype is not None and uop.dtype not in self.match_dtype) or \
        (self.arg is not None and self.arg != uop.arg) or \
        (self.match_tag is not None and uop.tag not in self.match_tag) or \
        (len(uop.src) < self.required_len) or \

tinygrad/uop/symbolic.py

@@ -100,10 +100,6 @@ pm_remove_invalid = PatternMatcher([
    if any(x.is_invalid for x in s.src) else None),
 ])
 
-# the one rule that collapses the pair CAST(dt, CONST(v)) into a typed CONST
-# TODO: delete this once CONST has no dtype
-pm_fold_cast_const = PatternMatcher([(UPat(Ops.CAST, name="root", src=(UPat.cvar("c"),)), lambda root, c: root.const_like(c.val))])
-
 symbolic_simple = pm_data_invalid + PatternMatcher([
   # ** self folding **
   (UPat.var("x") + 0, lambda x: x),    # x+0 -> x

tinygrad/uop/weak.py

@@ -16,10 +16,10 @@ def lower_weak_node(u:UOp) -> UOp|None:
 
 pm_lower_weak = PatternMatcher([
   (UPat(Ops.CONST, dtype=dtypes.weaks, name="u"), lambda u: UOp.const(u.val, select_dtype(u)).cast(u.dtype)),
-  # two stacked weak casts are a weakint value used as weakfloat (or vice versa): resolve the inner one at the outer kind's default.
+  # two stacked weak casts are two kind conversions: each resolves at its own kind's default
   # a SINGLE weak cast is never rewritten here, each consumer absorbs it on its own edge (see lower_weak_srcs)
   (UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat.var("x"),)),), name="u"),
-   lambda u,x: x.cast(select_dtype(u)).cast(u.dtype) if x.dtype not in dtypes.weaks else None),
+   lambda u,x: x.cast(select_dtype(u.src[0])).cast(select_dtype(u)).cast(u.dtype) if x.dtype not in dtypes.weaks else None),
   # Binary can widen from the bounds, all other nodes derive from the lowered sources.
   # a weakfloat Unary (sin/exp2/...) must resolve here, before the transcendental decomposition
   (UPat(GroupOp.Binary|GroupOp.Unary|{Ops.WHERE, Ops.RANGE, Ops.STACK, Ops.SPECIAL}, name="u"), lower_weak_node),
@@ -69,6 +69,9 @@ pm_cast_weak = PatternMatcher([
 ])
 
 pm_lower_index_dtype = pm_commit_weak+pm_cast_weak+PatternMatcher([
+  # a CAST between two concrete dtypes over a CONST is a value conversion: evaluate it once, at the width the CAST states
+  # TODO: delete this once CONST has no dtype
+  (UPat(Ops.CAST, dtypes.all, name="root", src=(UPat.cvar("c", dtypes.all),)), lambda root, c: root.const_like(c.val)),
   (UPat(GroupOp.All, name="u"),
    lambda ctx,u: lower_weak_srcs(ctx, u) if u.dtype not in dtypes.weaks and any(s.dtype in dtypes.weaks for s in u.src) else None),
   # a valid index into an n-element buffer lives in [0,n): a gated long index narrows when n-1 fits int32 (out-of-gate wraps, discarded)