error

llm: minor fixes + tests
pytest: use timeout_method signal (#17094 )
2026-08-21 09:46:07 +00:00 · 2026-07-20 14:26:07 -07:00 · 2026-07-20 14:21:30 -07:00 · 2026-07-20 15:19:24 -04:00 · 2026-07-20 15:18:48 -04:00 · 2026-07-20 15:08:54 -04:00
151 changed files with 3013 additions and 3867 deletions
@@ -521,6 +521,8 @@ jobs:
      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
    - name: openpilot run_pickle 0.10.1 driving_vision
      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py
+    - name: Test copy speeds
+      run: SIZE=64e6 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds

  driverbenchmarks:
    name: PCI Driver Benchmark (DEV=${{ matrix.dev }})
@@ -14,12 +14,15 @@ jobs:
    outputs:
      branchstat: ${{ steps.brstat.outputs.stat}}
    steps:
-      - name: Check code from PR branch 
+      - name: Check code from PR branch
        uses: actions/checkout@v6
        with:
          repository: ${{ github.event.pull_request.head.repo.full_name }}
          ref: ${{ github.event.pull_request.head.sha }}
          fetch-depth: 0
+          # PR code is only inspected with git rev-list, never executed
+          allow-unsafe-pr-checkout: true
+          persist-credentials: false
      - name: Check whether branch is up-to-date
        id: brstat
        run: |
@@ -51,6 +54,9 @@ jobs:
          repository: ${{ github.event.pull_request.head.repo.full_name }}
          ref: ${{ github.event.pull_request.head.sha }}
          path: pr
+          # PR code is only line-counted by master's sz.py, never executed
+          allow-unsafe-pr-checkout: true
+          persist-credentials: false
        # the base default to tinygrad master and cannot be other fork branch for security purpose
      - name: Checkout code from tinygrad master
        uses: actions/checkout@v6
@@ -171,7 +171,7 @@ jobs:
        llvm: 'true'
        amd: 'true'
    - name: Run NULL backend tests
-      run: DEV=NULL python -m pytest -n=auto test/null/ --durations=20
+      run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --durations=20
    - name: Run targeted tests on NULL backend
      run: |
        DEV=NULL python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
@@ -500,6 +500,29 @@ jobs:
      - name: Run LLVM test
        run: DEV=MOCKKFD+AMD:LLVM python test/device/test_amd_llvm.py

+  hcq2:
+    name: hcq2
+    runs-on: *linux
+    timeout-minutes: 5
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v6
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: hcq2
+          deps: testing_unit
+          amd: 'true'
+      - name: Run HCQ2 tests
+        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/test_tiny.py
+      - name: Run HCQ2 multi-device tests
+        run: |
+          HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_multitensor.py \
+            TestMultiTensor.test_simple_add TestMultiTensor.test_shard_reduce \
+            TestMultiTensor.test_backward_sum TestMultiTensor.test_matmul_shard_0_0
+      - name: Run HCQ2 JIT tests
+        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_jit.py
+
  testmockam:
    name: Linux (am)
    runs-on: *linux
@@ -620,7 +643,7 @@ jobs:
    - name: Run unit tests
      run: DEV=METAL python -m pytest -n=auto test/unit/ --durations=20
    - name: Run NULL backend tests
-      run: DEV=NULL python -m pytest -n=auto test/null/ --durations=20
+      run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --durations=20
    - name: Test tensor core ops (fake)
      run: DEV=METAL DEBUG=3 TC=2 python test/backend/test_ops.py TestOps.test_gemm
    - name: Test tensor core ops (real)
@@ -3,7 +3,7 @@
 # tinygrad implementation of https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
 # https://myrtle.ai/learn/how-to-train-your-resnet-8-bag-of-tricks/
 # https://siboehm.com/articles/22/CUDA-MMM
-import random, time, math
+import random, time
 import numpy as np
 from typing import Optional
 from extra.lr_scheduler import OneCycleLR
@@ -49,13 +49,14 @@ class UnsyncedBatchNorm:
      # https://github.com/pytorch/pytorch/blob/c618dc13d2aa23625cb0d7ada694137532a4fa33/aten/src/ATen/native/cuda/Normalization.cuh
      # There's "online" algorithms that fix this, like https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_Online_algorithm
      batch_mean = x.mean(axis=(1,3,4))
-      batch_var = (x*x).mean(axis=(1,3,4)) - batch_mean*batch_mean
+      y = (x - batch_mean.detach().reshape(shape=[batch_mean.shape[0], 1, -1, 1, 1]))  # d(var)/d(mean) = 0
+      batch_var = (y*y).mean(axis=(1,3,4))
      batch_invstd = batch_var.add(self.eps).pow(-0.5)

      # NOTE: wow, this is done all throughout training in most PyTorch models
      if self.track_running_stats:
        self.running_mean.assign((1-self.momentum) * self.running_mean + self.momentum * batch_mean.detach().cast(self.running_mean.dtype))
-        batch_var_adjust = prod(x.shape[1:])/(prod(x.shape[1:])-x.shape[2])
+        batch_var_adjust = prod(y.shape[1:])/(prod(y.shape[1:])-y.shape[2])
        self.running_var.assign((1-self.momentum) * self.running_var + self.momentum * batch_var_adjust * batch_var.detach().cast(self.running_var.dtype))
        self.num_batches_tracked += 1
    else:
@@ -69,37 +70,25 @@ class BatchNorm(nn.BatchNorm2d if getenv("SYNCBN") else UnsyncedBatchNorm):
    super().__init__(num_features, track_running_stats=False, eps=1e-12, momentum=0.85, affine=True)
    self.weight.is_param_(False)

-class MatmulConv2d(nn.Conv2d):
-  def __call__(self, x:Tensor) -> Tensor:
-    if not getenv("MATMUL_CONV", 1): return super().__call__(x)
-    assert self.groups == 1 and self.stride == self.dilation == self.padding == 1 and self.bias is None
-    bs, cin, _, _ = x.shape
-    cout, _, ky, kx = self.weight.shape
-    patches = x.pad((1, 1, 1, 1))._pool((ky, kx), 1, 1)
-    oy, ox = patches.shape[2:4]
-    patches = patches.permute(0, 2, 3, 1, 4, 5).reshape(bs*oy*ox, cin*ky*kx).contiguous().contiguous_backward()
-    out = (patches @ self.weight.reshape(cout, cin*ky*kx).T).contiguous().contiguous_backward()
-    return out.reshape(bs, oy, ox, cout).permute(0, 3, 1, 2)
-
 class ConvGroup:
  def __init__(self, channels_in, channels_out):
-    self.conv1 = MatmulConv2d(channels_in,  channels_out, kernel_size=3, padding=1, bias=False)
-    self.conv2 = MatmulConv2d(channels_out, channels_out, kernel_size=3, padding=1, bias=False)
+    self.conv1 = nn.Conv2d(channels_in,  channels_out, kernel_size=3, padding=1, bias=False)
+    self.conv2 = nn.Conv2d(channels_out, channels_out, kernel_size=3, padding=1, bias=False)

    self.norm1 = BatchNorm(channels_out)
    self.norm2 = BatchNorm(channels_out)

  def __call__(self, x):
-    x = self.conv1(x).contiguous()
+    x = self.conv1(x)
    x = x.max_pool2d(2)
    x = x.float()
-    x = self.norm1(x).contiguous()
+    x = self.norm1(x)
    x = x.cast(dtypes.default_float)
-    x = x.quick_gelu().contiguous()
+    x = x.quick_gelu()
    residual = x
-    x = self.conv2(x).contiguous()
+    x = self.conv2(x)
    x = x.float()
-    x = self.norm2(x).contiguous()
+    x = self.norm2(x)
    x = x.cast(dtypes.default_float)
    x = x.quick_gelu()

@@ -122,10 +111,7 @@ class SpeedyResNet:
  def __call__(self, x, training=True):
    # pad to 32x32 because whitening conv creates 31x31 images that are awfully slow to compute with
    # TODO: remove the pad but instead let the kernel optimize itself
-    def forward(x):
-      x = x.conv2d(self.whitening).pad((1,0,0,1)).contiguous()
-      for layer in self.net: x = layer(x).contiguous()
-      return x
+    forward = lambda x: x.conv2d(self.whitening).pad((1,0,0,1)).sequential(self.net)
    return forward(x) if training else (forward(x) + forward(x[..., ::-1])) / 2.

 # hyper-parameters were exactly the same as the original repo
@@ -230,31 +216,15 @@ def train_cifar():
    Y_cutmix = mix_portion * Y_patch + (1. - mix_portion) * Y
    return X_cutmix, Y_cutmix

-  def random_permutation(rows:int, cols:int) -> Tensor:
-    size = rows * cols
-    # An affine map is a permutation when its stride is coprime to the domain size.
-    while math.gcd(stride:=random.randrange(1, size), size) != 1: pass
-    return (Tensor.arange(size) * stride + random.randrange(size)) % size
-
-  def shuffled_augmentations(X:Tensor, Y:Tensor):
-    perms = random_permutation(X.shape[0] // BS, BS)
-    X, Y = X[:perms.shape[0]], Y[:perms.shape[0]]
+  @TinyJit
+  def augmentations(X:Tensor, Y:Tensor):
+    perms = Tensor.randperm(X.shape[0], device=X.device) # We reuse perms for cutmix, because they are expensive to generate
    if getenv("RANDOM_CROP", 1):
      X = random_crop(X, crop_size=32)
    if getenv("RANDOM_FLIP", 1):
      # NOTE: RANGEIFY=1 needs this contiguous or the X[perms] is very slow
      X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X).contiguous() # flip LR
    X, Y = X[perms], Y[perms]
-    return X, Y, perms
-
-  @TinyJit
-  def augmentations(X:Tensor, Y:Tensor):
-    X, Y, _ = shuffled_augmentations(X, Y)
-    return X, Y
-
-  @TinyJit
-  def augmentations_cutmix(X:Tensor, Y:Tensor):
-    X, Y, perms = shuffled_augmentations(X, Y)
    return X, Y, *cutmix(X, Y, perms, mask_size=hyp['net']['cutmix_size'])

  # the operations that remain inside batch fetcher is the ones that involves random operations
@@ -264,10 +234,8 @@ def train_cifar():
      st = time.monotonic()
      X, Y = X_in, Y_in
      if is_train:
-        if getenv("CUTMIX", 1) and step >= hyp['net']['cutmix_steps']:
-          _, _, X, Y = augmentations_cutmix(X, Y)
-        else:
-          X, Y = augmentations(X, Y)
+        X, Y, X_cm, Y_cm = augmentations(X, Y)
+        if getenv("CUTMIX", 1) and step >= hyp['net']['cutmix_steps']: X, Y = X_cm, Y_cm
      et = time.monotonic()
      print(f"shuffling {'training' if is_train else 'test'} dataset in {(et-st)*1e3:.2f} ms ({epoch=})")

@@ -287,6 +255,7 @@ def train_cifar():

  class modelEMA():
    def __init__(self, w, net):
+      # self.model_ema = copy.deepcopy(net) # won't work for opencl due to unpickeable pyopencl._cl.Buffer
      self.net_ema = SpeedyResNet(w)
      for net_ema_param, net_param in zip(get_state_dict(self.net_ema).values(), get_state_dict(net).values()):
        net_ema_param.assign(net_param.numpy())
@@ -311,15 +280,6 @@ def train_cifar():

  # initialize model weights
  model = SpeedyResNet(W)
-  model_state = get_state_dict(model)
-  random_params = [x for name, x in model_state.items() if x.is_param and "bias" not in name]
-  Tensor.manual_seed(getenv('SEED', hyp['seed']))
-  random_values = Tensor.rand(sum(x.numel() for x in random_params))
-  offset = 0
-  for param in random_params:
-    bound = prod(param.shape[1:]) ** -0.5
-    param.replace(((random_values[offset:offset+param.numel()] * (2 * bound)) - bound).reshape(param.shape).cast(param.dtype))
-    offset += param.numel()

  # padding is not timed in the original repo since it can be done all at once
  X_train = pad_reflect(X_train, size=hyp['net']['pad_amount'])
@@ -336,7 +296,7 @@ def train_cifar():
        x.to_(GPUS)

  # parse the training params into bias and non-bias
-  params_dict = model_state
+  params_dict = get_state_dict(model)
  params_bias = []
  params_non_bias = []
  for params in params_dict:
@@ -360,7 +320,7 @@ def train_cifar():
  lr_sched_non_bias = OneCycleLR(opt_non_bias, max_lr=hyp['opt']['non_bias_lr'], pct_start=pct_start, div_factor=initial_div_factor, final_div_factor=1./(initial_div_factor*final_lr_ratio), total_steps=STEPS)

  def train_step(model, optimizer, lr_scheduler, X, Y):
-    out = model(X).contiguous()
+    out = model(X)
    loss_batchsize_scaler = 512/BS
    loss = cross_entropy(out, Y, reduction='none', label_smoothing=hyp['opt']['label_smoothing']).mul(hyp['opt']['loss_scale_scaler']*loss_batchsize_scaler).sum().div(hyp['opt']['loss_scale_scaler'])

@@ -371,19 +331,15 @@ def train_cifar():
      return loss.realize(*optimizer.schedule_step(), *lr_scheduler[0].schedule_step(), *lr_scheduler[1].schedule_step())
    return loss.realize()

-  train_step_jitted = TinyJit(train_step, warmup=False)
+  train_step_jitted = TinyJit(train_step)

-  def eval_forward(model, X):
-    return model(X).realize()
-
-  def eval_step(out, out_flipped, Y):
-    out = (out + out_flipped) / 2.
+  def eval_step(model, X, Y):
+    out = model(X, training=False)
+    loss = cross_entropy(out, Y, reduction='mean')
    correct = out.argmax(axis=1) == Y.argmax(axis=1)
-    return correct.sum().realize()
-  eval_forward_jitted = TinyJit(eval_forward, warmup=False)
-  eval_forward_ema_jitted = TinyJit(eval_forward, warmup=False)
-  eval_step_jitted = TinyJit(eval_step, warmup=False)
-  eval_step_ema_jitted = TinyJit(eval_step, warmup=False)
+    return correct.realize(), loss.realize()
+  eval_step_jitted     = TinyJit(eval_step)
+  eval_step_ema_jitted = TinyJit(eval_step)

  # 97 steps in 2 seconds = 20ms / step
  # step is 1163.42 GOPS = 56 TFLOPS!!!, 41% of max 136
@@ -404,37 +360,31 @@ def train_cifar():
    while i <= STEPS:
      if i % getenv("EVAL_STEPS", STEPS) == 0 and i > 1 and not getenv("DISABLE_BACKWARD"):
        # Using Context(TRAINING=0) here actually bricks batchnorm, even with track_running_stats=True
-        correct_sum = correct_sum_ema = None
-        correct_len = correct_len_ema = 0
+        corrects = []
+        corrects_ema = []
+        losses = []
+        losses_ema = []
        for Xt, Yt in fetch_batches(X_test, Y_test, BS=EVAL_BS, is_train=False):
          if len(GPUS) > 1:
            Xt.shard_(GPUS, axis=0)
            Yt.shard_(GPUS, axis=0)

-          Xt_contiguous = Xt.contiguous().realize()
-          out = eval_forward_jitted(model, Xt_contiguous).clone().realize()
-          out_flipped = eval_forward_jitted(model, Xt_contiguous[..., ::-1].contiguous().realize())
-          batch_correct = eval_step_jitted(out, out_flipped, Yt)
-          correct_sum = batch_correct if correct_sum is None else correct_sum + batch_correct
-          correct_len += Yt.shape[0]
+          correct, loss = eval_step_jitted(model, Xt, Yt)
+          losses.append(loss.numpy().tolist())
+          corrects.extend(correct.numpy().tolist())
          if model_ema:
-            out_ema = eval_forward_ema_jitted(model_ema.net_ema, Xt_contiguous).clone().realize()
-            out_flipped_ema = eval_forward_ema_jitted(model_ema.net_ema, Xt_contiguous[..., ::-1].contiguous().realize())
-            batch_correct_ema = eval_step_ema_jitted(out_ema, out_flipped_ema, Yt)
-            correct_sum_ema = batch_correct_ema if correct_sum_ema is None else correct_sum_ema + batch_correct_ema
-            correct_len_ema += Yt.shape[0]
+            correct_ema, loss_ema = eval_step_ema_jitted(model_ema.net_ema, Xt, Yt)
+            losses_ema.append(loss_ema.numpy().tolist())
+            corrects_ema.extend(correct_ema.numpy().tolist())

        # collect accuracy across ranks
-        assert correct_sum is not None
-        correct_count = correct_sum.item()
-        if model_ema:
-          assert correct_sum_ema is not None
-          correct_count_ema = correct_sum_ema.item()
+        correct_sum, correct_len = sum(corrects), len(corrects)
+        if model_ema: correct_sum_ema, correct_len_ema = sum(corrects_ema), len(corrects_ema)

-        eval_acc_pct = correct_count/correct_len*100.0
-        if model_ema: acc_ema = correct_count_ema/correct_len_ema*100.0
-        print(f"eval     {correct_count}/{correct_len} {eval_acc_pct:.2f}% STEP={i} (in {(time.monotonic()-st)*1e3:.2f} ms)")
-        if model_ema: print(f"eval ema {correct_count_ema}/{correct_len_ema} {acc_ema:.2f}% STEP={i}")
+        eval_acc_pct = correct_sum/correct_len*100.0
+        if model_ema: acc_ema = correct_sum_ema/correct_len_ema*100.0
+        print(f"eval     {correct_sum}/{correct_len} {eval_acc_pct:.2f}%, {(sum(losses)/len(losses)):7.2f} val_loss STEP={i} (in {(time.monotonic()-st)*1e3:.2f} ms)")
+        if model_ema: print(f"eval ema {correct_sum_ema}/{correct_len_ema} {acc_ema:.2f}%, {(sum(losses_ema)/len(losses_ema)):7.2f} val_loss STEP={i}")

      if STEPS == 0 or i == STEPS: break

@@ -1462,6 +1462,8 @@ def train_llama3():

  @TinyJit
  def minibatch(tokens:Tensor):
+    for nxt in fp8_next_amax: nxt.assign(0)
+    for nxt in fp8_next_grad_amax: nxt.assign(0)
    if is_dp: tokens = tokens.to(None).shard(device, 0)
    if is_mp: tokens = tokens.shard(device)
    if not is_sharding: tokens = tokens.to(None)
@@ -1753,11 +1755,12 @@ def train_gptoss():

  scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)

-  # realize everything here
-  if optim.master_params: Tensor.realize(*optim.master_params)
+  if optim.master_params:
+    for m in optim.master_params: m.realize()
  Tensor.realize(*optim.params, *fp8_inv_scales)

  @TinyJit
+  @Context(TRAINING=1)
  def minibatch(tokens:Tensor):
    if is_dp: tokens = tokens.to(None).shard(device, 0)
    if not is_sharding: tokens = tokens.to(None)
@@ -37,8 +37,8 @@ def quantize_fp8(x:Tensor, amax_state:Tensor|None=None):
  return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal(), new_amax

 def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_scale:Tensor|None=None,
-           x_fp8:Tensor|None=None, x_new_amax:Tensor|None=None,
-           grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None, x_prequant_mx:tuple|None=None) -> tuple[Tensor,...]:
+           x_fp8:Tensor|None=None, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None, x_prequant_mx:tuple|None=None,
+           next_amax_x:Tensor|None=None) -> tuple[Tensor,...]:
  if not fp8:
    if ASM_GEMM:
      from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
@@ -56,13 +56,14 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca
    else:
      x_phys = (x_q.cast(dtypes.bfloat16) * _mx_block_scale(x_e8)).reshape(*l_shape, x_q.shape[-1])
      out = x_phys @ (w.cast(dtypes.bfloat16) * _mx_block_scale(w_inv_scale)).T
-    return out, (amax_x.detach() if amax_x is not None else None), x_q
+    return out, x_q
  if x_fp8 is None:
-    if FUSED_INPUT_QUANTIZE and amax_x is not None:
+    if FUSED_INPUT_QUANTIZE:
      from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed
-      x_fp8, _, x_new_amax, _ = quantize_fp8_delayed(x, amax_x, FP8_DTYPE)
+      x_fp8, _ = quantize_fp8_delayed(x, amax_x, next_amax_x, FP8_DTYPE)
    else:
-      x_fp8, _, x_new_amax = quantize_fp8(x, amax_state=amax_x)
+      x_fp8, _, new_amax_x = quantize_fp8(x, amax_state=amax_x)
+      next_amax_x.assign(new_amax_x)
  if ASM_GEMM:
    from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
    if can_use_asm_gemm(x_fp8, w.T):
@@ -73,51 +74,51 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca
      else:
        out = asm_gemm(x_fp8, w.T, x_scale=amax_x, w_scale=w_inv_scale, grad_amax_state=grad_amax_state,
                       next_grad_amax_state=next_grad_amax_state)
-      return out, x_new_amax, x_fp8
-  return (x_fp8.dot(w.T, dtype=dtypes.float) * ((amax_x.float() + 1e-8) / FP8_MAX) * w_inv_scale).cast(dtypes.bfloat16), x_new_amax, x_fp8
+      return out, x_fp8
+  return (x_fp8.dot(w.T, dtype=dtypes.float) * ((amax_x.float() + 1e-8) / FP8_MAX) * w_inv_scale).cast(dtypes.bfloat16), x_fp8

 def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
-                         grad_amax_state:Tensor, next_grad_amax_state:Tensor):
+                         next_amax_x:Tensor, grad_amax_state:Tensor, next_grad_amax_state:Tensor):
  if FUSED_ADD_NORM_MUL_QUANTIZE:
    from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_rmsnorm_mul_quantize_fp8
-    x_fp8, new_amax, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE)
-    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x, x_new_amax=new_amax,
+    x_fp8, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
+    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
                       grad_amax_state=grad_amax_state, next_grad_amax_state=next_grad_amax_state)
    return out, x_normed, rrms, ret
  x_normed, rrms = rmsnorm(x, eps)
  out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state,
-                     next_grad_amax_state=next_grad_amax_state)
+                     next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
  return out, x_normed, rrms, ret

 def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
-                             grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
+                             next_amax_x:Tensor, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
  if FUSED_ADD_NORM_MUL_QUANTIZE:
    from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_add_rmsnorm_mul_quantize_fp8
-    x_fp8, new_amax, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE)
-    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x, x_new_amax=new_amax,
+    x_fp8, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
+    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
                       grad_amax_state=grad_amax_state, next_grad_amax_state=next_grad_amax_state)
    return out, h, x_normed, rrms, ret
  h = x + residual
  x_normed, rrms = rmsnorm(h, eps)
  out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state,
-                     next_grad_amax_state=next_grad_amax_state)
+                     next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
  return out, h, x_normed, rrms, ret

 def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor,
-                             amax_x2:Tensor,
+                             amax_x2:Tensor, next_amax_x2:Tensor,
                             grad_amax_xw13:Tensor, next_grad_amax_xw13:Tensor,
                             grad_amax_xout:Tensor, next_grad_amax_xout:Tensor):
  if FUSED_SILU_W13:
    from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
-    x2_fp8, new_amax_x2 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13,
-                                                 next_grad_amax_state=next_grad_amax_xw13)
-    out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, amax_x=amax_x2, x_new_amax=new_amax_x2,
+    x2_fp8 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13,
+                                                 next_grad_amax_state=next_grad_amax_xw13, amax_out=next_amax_x2)
+    out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, amax_x=amax_x2,
                       grad_amax_state=grad_amax_xout, next_grad_amax_state=next_grad_amax_xout)
    return out, ret
  hidden = x_w13.shape[-1] // 2
  x_w1, x_w3 = x_w13[..., :hidden], x_w13[..., hidden:]
  out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout,
-                     next_grad_amax_state=next_grad_amax_xout)
+                     next_grad_amax_state=next_grad_amax_xout, next_amax_x=next_amax_x2)
  return out, ret

 class FlatTransformer:
@@ -154,7 +155,7 @@ class FlatTransformer:
    self.tok_embeddings = nn.Embedding(vocab_size, dim)
    self.tok_embeddings.weight = Tensor.normal(vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
    self.output = Tensor.normal(1, vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
-    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().is_param_(False)
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).clone().is_param_(False)

    def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().is_param_(False)
    names = ["xqkv", "xo", "x2"]
@@ -186,89 +187,87 @@ class FlatTransformer:

  def attention(self, x:Tensor, freqs_cis:Tensor, *, attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
                amax_xqkv:Tensor, amax_xo:Tensor, s_qkv:Tensor, s_o:Tensor,
+                next_amax_xqkv:Tensor, next_amax_xo:Tensor,
                grad_amax_xqkv:Tensor, grad_amax_xo:Tensor, next_grad_amax_xqkv:Tensor, next_grad_amax_xo:Tensor):
    bsz, seqlen, _ = x.shape
-    amaxs, saves = [], []
+    saves = []

-    xqkv, x_normed, rrms, (new_amax, *s) = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
+    xqkv, x_normed, rrms, s = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
                                                                  amax_x=amax_xqkv, grad_amax_state=grad_amax_xqkv,
-                                                                  next_grad_amax_state=next_grad_amax_xqkv)
-    amaxs.append(new_amax)
+                                                                  next_grad_amax_state=next_grad_amax_xqkv, next_amax_x=next_amax_xqkv)
    saves.extend([x_normed, rrms, *s, xqkv])
-    xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
-    xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
-    xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
-    xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
-
-    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
-    xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
    if getenv("HK_FLASH_ATTENTION"):
-      from extra.thunder.amd.fa import flash_attention
+      from extra.thunder.amd.fa import flash_attention, fused_qkv_rope
+      xq, xk, xv = fused_qkv_rope(xqkv, freqs_cis, self.n_heads, self.n_kv_heads, self.head_dim)
      attn, *save = flash_attention(xq, xk, xv, is_causal=True, write_flat=True)
      saves.extend(save)
    else:
+      xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
+      xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
+      xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
+      xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
+      xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
+      xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
      xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
      attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
    attn = attn.reshape(bsz, seqlen, -1)

-    out, new_amax, *s = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo,
-                               next_grad_amax_state=next_grad_amax_xo)
-    amaxs.append(new_amax)
+    out, *s = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo,
+                               next_grad_amax_state=next_grad_amax_xo, next_amax_x=next_amax_xo)
    saves.extend([*s, out])
-    return out, amaxs, saves
+    return out, saves

  def feed_forward(self, x:Tensor, residual:Tensor, **kwargs):
-    amaxs, saves = [], []
+    saves = []

    if SPLIT_W13:
      h = x + residual
      x_normed, rrms = rmsnorm(h, self.norm_eps)
      saves.extend([x_normed, rrms])
      inp = x_normed * kwargs["ffn_norm"]
-      x_w1, new_amax, *s = matmul(inp, kwargs["w1"], amax_x=kwargs["amax_x1"], w_inv_scale=kwargs["s_1"],
-                                  grad_amax_state=kwargs["grad_amax_xw1"], next_grad_amax_state=kwargs["next_grad_amax_xw1"])
-      amaxs.append(new_amax)
+      x_w1, *s = matmul(inp, kwargs["w1"], amax_x=kwargs["amax_x1"], w_inv_scale=kwargs["s_1"],
+                                  grad_amax_state=kwargs["grad_amax_xw1"], next_grad_amax_state=kwargs["next_grad_amax_xw1"],
+                                  next_amax_x=kwargs["next_amax_x1"])
      saves.extend([*s, x_w1])
-      x_w3, new_amax, *s = matmul(inp, kwargs["w3"], amax_x=kwargs["amax_x3"], w_inv_scale=kwargs["s_3"],
-                                  grad_amax_state=kwargs["grad_amax_xw3"], next_grad_amax_state=kwargs["next_grad_amax_xw3"])
-      amaxs.append(new_amax)
+      x_w3, *s = matmul(inp, kwargs["w3"], amax_x=kwargs["amax_x3"], w_inv_scale=kwargs["s_3"],
+                                  grad_amax_state=kwargs["grad_amax_xw3"], next_grad_amax_state=kwargs["next_grad_amax_xw3"],
+                                  next_amax_x=kwargs["next_amax_x3"])
      saves.extend([*s, x_w3])
      if FUSED_SILU_W13 and MXFP8:
        from extra.llama_kernels.fused_silu_mul_quantize_mxfp8 import fused_silu_mul_quantize_mxfp8
        aq, ae8, asi = fused_silu_mul_quantize_mxfp8(x_w1.reshape(-1, x_w1.shape[-1]), x_w3.reshape(-1, x_w3.shape[-1]))
-        out, new_amax, *s = matmul(None, kwargs["w2"], x_prequant_mx=(aq, ae8, asi), amax_x=kwargs["amax_x2"],
-                                   w_inv_scale=kwargs["s_2"], grad_amax_state=kwargs["grad_amax_xout"],
-                                   next_grad_amax_state=kwargs["next_grad_amax_xout"])
+        out, *s = matmul(None, kwargs["w2"], x_prequant_mx=(aq, ae8, asi), amax_x=kwargs["amax_x2"],
+                         w_inv_scale=kwargs["s_2"], grad_amax_state=kwargs["grad_amax_xout"],
+                         next_grad_amax_state=kwargs["next_grad_amax_xout"], next_amax_x=kwargs["next_amax_x2"])
        out = out.reshape(*x_w1.shape[:-1], kwargs["w2"].shape[0])
      else:
-        out, new_amax, *s = matmul(x_w1.silu() * x_w3, kwargs["w2"], amax_x=kwargs["amax_x2"], w_inv_scale=kwargs["s_2"],
-                                   grad_amax_state=kwargs["grad_amax_xout"], next_grad_amax_state=kwargs["next_grad_amax_xout"])
-      amaxs.append(new_amax)
+        out, *s = matmul(x_w1.silu() * x_w3, kwargs["w2"], amax_x=kwargs["amax_x2"], w_inv_scale=kwargs["s_2"],
+                         grad_amax_state=kwargs["grad_amax_xout"], next_grad_amax_state=kwargs["next_grad_amax_xout"],
+                         next_amax_x=kwargs["next_amax_x2"])
      saves.extend([*s, out])
    else:
-      x_w13, h, x_normed, rrms, (new_amax, *s) = add_norm_quantize_matmul(x, residual, kwargs["ffn_norm"], kwargs["w13"], kwargs["s_13"],
+      x_w13, h, x_normed, rrms, s = add_norm_quantize_matmul(x, residual, kwargs["ffn_norm"], kwargs["w13"], kwargs["s_13"],
                                                                          self.norm_eps, amax_x=kwargs["amax_x13"],
+                                                                          next_amax_x=kwargs["next_amax_x13"],
                                                                          grad_amax_state=kwargs["grad_amax_xw13"],
                                                                          next_grad_amax_state=kwargs["next_grad_amax_xw13"])
-      amaxs.append(new_amax)
      saves.extend([x_normed, rrms, *s, x_w13])
-      out, (new_amax, *s) = silu_w13_quantize_matmul(x_w13, kwargs["w2"], kwargs["s_2"], amax_x2=kwargs["amax_x2"],
+      out, s = silu_w13_quantize_matmul(x_w13, kwargs["w2"], kwargs["s_2"], amax_x2=kwargs["amax_x2"],
+                                                     next_amax_x2=kwargs["next_amax_x2"],
                                                     grad_amax_xw13=kwargs["grad_amax_xw13"],
                                                     next_grad_amax_xw13=kwargs["next_grad_amax_xw13"],
                                                     grad_amax_xout=kwargs["grad_amax_xout"],
                                                     next_grad_amax_xout=kwargs["next_grad_amax_xout"])
-      amaxs.append(new_amax)
      saves.extend([*s, out])
-    return out, h, amaxs, saves
+    return out, h, saves

  @function(precompile=True, precompile_backward=True)
  def run_layer(self, x:Tensor, freqs_cis:Tensor, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
-    attn, attn_amaxs, attn_saves = self.attention(x, freqs_cis, **attn_kwargs)
-    ffn, h, ffn_amaxs, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs)
+    attn, attn_saves = self.attention(x, freqs_cis, **attn_kwargs)
+    ffn, h, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs)
    h = h + ffn
-    amaxs = tuple(a.detach() for a in (*attn_amaxs, *ffn_amaxs))
-    if save: return (h, *amaxs, *attn_saves, *ffn_saves)
-    else: return (h, *amaxs)
+    if save: return (h, *attn_saves, *ffn_saves)
+    else: return (h,)

  def shard(self, device:tuple[str, ...], mp:bool=False):
    from tinygrad.nn.state import get_parameters
@@ -313,26 +312,27 @@ class FlatTransformer:

  def __call__(self, tokens:Tensor, save:bool=True):
    h = self.tok_embeddings(tokens)
-    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
+    freqs_cis = self.freqs_cis.cast(h.dtype)
+    if not getenv("HK_FLASH_ATTENTION"): freqs_cis = freqs_cis[:, :tokens.shape[1], :, :, :]
    a, na, ga, nga, s = self._fp8_amax, self._fp8_next_amax, self._fp8_grad_amax, self._fp8_next_grad_amax, self._fp8_inv_scale
    for i in range(self.n_layers):
      attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i],
                         amax_xqkv=a["xqkv"][i], amax_xo=a["xo"][i], s_qkv=s["wqkv"][i], s_o=s["wo"][i],
+                         next_amax_xqkv=na["xqkv"][i], next_amax_xo=na["xo"][i],
                         grad_amax_xqkv=ga["xqkv"][i], grad_amax_xo=ga["xo"][i],
                         next_grad_amax_xqkv=nga["xqkv"][i], next_grad_amax_xo=nga["xo"][i])
      ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], w2=self.w2[i],
-                        amax_x2=a["x2"][i], s_2=s["w2"][i], grad_amax_xout=ga["xout"][i], next_grad_amax_xout=nga["xout"][i])
+                        amax_x2=a["x2"][i], s_2=s["w2"][i], grad_amax_xout=ga["xout"][i], next_grad_amax_xout=nga["xout"][i],
+                        next_amax_x2=na["x2"][i])
      if SPLIT_W13:
        ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i], amax_x1=a["x1"][i], amax_x3=a["x3"][i],
+                          next_amax_x1=na["x1"][i], next_amax_x3=na["x3"][i],
                          s_1=s["w1"][i], s_3=s["w3"][i], grad_amax_xw1=ga["xw1"][i], grad_amax_xw3=ga["xw3"][i],
                          next_grad_amax_xw1=nga["xw1"][i], next_grad_amax_xw3=nga["xw3"][i])
      else:
        ffn_kwargs.update(w13=self.w13[i], amax_x13=a["x13"][i], s_13=s["w13"][i], grad_amax_xw13=ga["xw13"][i],
-                          next_grad_amax_xw13=nga["xw13"][i])
-      h, *ret = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)
-      amax_names = ["xqkv", "xo"] + (["x1", "x3"] if SPLIT_W13 else ["x13"]) + ["x2"]
-      for name, new_val in zip(amax_names, ret[:len(amax_names)]):
-        na[name][i].assign(new_val)
+                          next_grad_amax_xw13=nga["xw13"][i], next_amax_x13=na["x13"][i])
+      h, *_ = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)

    logits = matmul(self.norm(h), self.output[0], fp8=False)[0]
    return logits
@@ -415,6 +415,9 @@ if __name__ == "__main__":
  @TinyJit
  def fwd_bwd(tokens:Tensor):
    with Timing("python forward: "):
+      for amax_dict in (model._fp8_next_amax, model._fp8_next_grad_amax):
+        for ts in amax_dict.values():
+          for nxt in ts: nxt.assign(0)
      logits = model(tokens[:, :-1], save=llama_size=="8B")
      loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
    with Timing("python backward: "):
@@ -12,7 +12,7 @@ from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
 from extra.llama_kernels.rmsnorm import rmsnorm
-from extra.gemm.cdna_asm_gemm import _mx_block_scale, quantize_mxfp8
+from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8

 FP8_DTYPE = dtypes.fp8e4m3
 FP8_MAX = 448.0
@@ -39,8 +39,11 @@ def quant_dequant_mx(x:Tensor) -> Tensor:
  fxn = _quant_dequant_fwd_fxn(x.as_param(0).uop, x.device)
  return Tensor(UOp.maketuple(fxn.uop).call(x.uop, grad_fxn=_quant_dequant_bwd).gettuple(0))

+def _mx_scale(e8:Tensor) -> Tensor:
+  return _mx_block_scale(e8) if e8.ndim == 2 else _mx_block_scale_3d(e8)
+
 def _dequant_fwd(w_q:Tensor, w_scale:Tensor) -> Tensor:
-  return w_q.cast(dtypes.bfloat16) * _mx_block_scale(w_scale)
+  return w_q.cast(dtypes.bfloat16) * _mx_scale(w_scale)

@functools.cache
 def _dequant_fwd_fxn(wq_p, ws_p, device):
@@ -48,7 +51,7 @@ def _dequant_fwd_fxn(wq_p, ws_p, device):

 def _dequant_bwd(grad:UOp, call:UOp) -> tuple:
  w_scale = Tensor(call.src[2])
-  return ((Tensor(grad).cast(dtypes.bfloat16) * _mx_block_scale(w_scale).cast(dtypes.bfloat16)).uop, None)
+  return ((Tensor(grad).cast(dtypes.bfloat16) * _mx_scale(w_scale).cast(dtypes.bfloat16)).uop, None)

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

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
 export GBS=$((BS * GRADIENT_ACC_STEPS))

 export MODEL="gptoss"
@@ -34,7 +34,7 @@ export SEED=${SEED:-5760}
 export DATA_SEED=${DATA_SEED:-5760}

 export JITBEAM=${JITBEAM:-3}
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0

 export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
@@ -20,7 +20,7 @@ export ZERO_OPTIM=${ZERO_OPTIM:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
 export GBS=$((BS * GRADIENT_ACC_STEPS))

 export MODEL="gptoss"
@@ -34,6 +34,6 @@ export SEED=${SEED:-$RANDOM}
 export DATA_SEED=${DATA_SEED:-5760}

 export JITBEAM=${JITBEAM:-3}
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0

 python3 examples/mlperf/model_train.py
@@ -2,5 +2,4 @@
 export BENCHMARK=${BENCHMARK:-5}
 export EVAL_BS=0
 VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=${DEBUG:--0} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
-SRC="AMD"; [[ $DEV == NULL* ]] && SRC="NULL"
-[ "$BENCHMARK" -le 3 ] || python -m tinygrad.viz.cli -s "$SRC" -t --interval "train @ 2" "train @ 3"
+[ "$BENCHMARK" -le 3 ] || [[ $DEV == NULL* ]] || python -m tinygrad.viz.cli -s AMD -t --interval "train @ 2" "train @ 3"
@@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=${DEV:-AMD}
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-2}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export MXFP8=${MXFP8:-1}
+export ZERO_OPTIM=${ZERO_OPTIM:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="gptoss"
+export BASEDIR="/raid/datasets/c4-8b/"
+export EVAL_TARGET=3.34 EVAL_FREQ=12288
+export END_LR="4e-5" WARMUP_STEPS=128 MAX_STEPS=1200000
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-5760}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+if [ -z "$FULL_LAYERS" ]; then
+  export LAYERS=${LAYERS:-2}
+fi
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,39 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=${DEV:-AMD}
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-0}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export MXFP8=${MXFP8:-1}
+export ZERO_OPTIM=${ZERO_OPTIM:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="gptoss"
+export BASEDIR="/raid/datasets/c4-8b/"
+export EVAL_TARGET=3.34 EVAL_FREQ=12288
+export END_LR="4e-5" WARMUP_STEPS=128 MAX_STEPS=1200000
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-$RANDOM}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,54 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=${DEV:-AMD}
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-2}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export FP8=${FP8:-1}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export FAST_CE=${FAST_CE:-1}
+export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
+export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
+export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
+export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
+export SPLIT_W13=${SPLIT_W13:-0}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-5760}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+if [ -z "$FULL_LAYERS" ]; then
+  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
+fi
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,54 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=${DEV:-AMD}
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-2}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export FP8=${FP8:-1}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export FAST_CE=${FAST_CE:-0}
+export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
+export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
+export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
+export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
+export SPLIT_W13=${SPLIT_W13:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-1} MP=${MP:-8} BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-5760}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+if [ -z "$FULL_LAYERS" ]; then
+  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
+fi
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=${DEV:-AMD}
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-0}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export FP8=${FP8:-1}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export FAST_CE=${FAST_CE:-1}
+export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
+export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
+export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
+export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
+export SPLIT_W13=${SPLIT_W13:-0}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-$RANDOM}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=${DEV:-AMD}
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-0}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export FP8=${FP8:-1}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export FAST_CE=${FAST_CE:-0}
+export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
+export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
+export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
+export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
+export SPLIT_W13=${SPLIT_W13:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-1} MP=${MP:-8} BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-32}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-$RANDOM}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,5 @@
+#!/bin/bash
+export BENCHMARK=${BENCHMARK:-5}
+export EVAL_BS=0
+VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=${DEBUG:--0} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
+[ "$BENCHMARK" -le 3 ] || [[ $DEV == NULL* ]] || python -m tinygrad.viz.cli -s AMD -t --interval "train @ 2" "train @ 3"
@@ -0,0 +1,58 @@
+#!/usr/bin/env bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=AMD
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export HK_FLASH_ATTENTION=1
+export ALL2ALL=1
+export LATE_ALLREDUCE=0
+export USE_ATOMICS=1
+export ASM_GEMM=1
+export WQKV=1
+export MASTER_WEIGHTS=1
+export FP8=1
+export ALLREDUCE_CAST=1
+export FAST_CE=1
+export FUSED_INPUT_QUANTIZE=1
+export FUSED_GRAD_QUANTIZE=1
+export FUSED_ADD_NORM_MUL_QUANTIZE=1
+export FUSED_SILU_W13=1
+export SPLIT_W13=0
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=8 MP=1 BS=16 EVAL_BS=8 GRADIENT_ACC_STEPS=2
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=8B
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=8192
+
+export SEED=$RANDOM
+export DATA_SEED=$SEED
+
+export JITBEAM=3
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+export LOGMLPERF=1
+
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="llama31_8b_8xMI350x_${DATETIME}_${SEED}.log"
+
+# beam
+FAKEDATA=1 BENCHMARK=10 INITMLPERF=1 LLAMA_LAYERS=2 python3 examples/mlperf/model_train.py | tee "$LOGFILE"
+
+# run
+RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a "$LOGFILE"
@@ -0,0 +1,10 @@
+#!/bin/bash
+export BENCHMARK=5
+export EVAL_BS=0
+export FAKEDATA=1
+export NULL_ALLOW_COPYOUT=1
+export HIP_VISIBLE_DEVICES=""
+export DEV=NULL:HIP:gfx950
+export JITBEAM=0
+export LLAMA_LAYERS=${LLAMA_LAYERS:-"2"}
+time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh
@@ -0,0 +1,38 @@
+{
+    "submitter": "tinycorp",
+    "division": "closed",
+    "status": "Available on-premise",
+    "system_name": "tinybox 8xMI300X",
+    "number_of_nodes": "1",
+    "host_processors_per_node": "2",
+    "host_processor_model_name": "AMD EPYC 9354",
+    "host_processor_core_count": "32",
+    "host_processor_vcpu_count": "64",
+    "host_processor_frequency": "",
+    "host_processor_caches": "",
+    "host_processor_interconnect": "",
+    "host_memory_capacity": "2304GB",
+    "host_storage_type": "NVMe SSD",
+    "host_storage_capacity": "3x 4TB raid array",
+    "host_networking": "",
+    "host_networking_topology": "",
+    "host_memory_configuration": "24x 96GB DDR5",
+    "accelerators_per_node": "8",
+    "accelerator_model_name": "AMD Instinct MI300X 192GB HBM3",
+    "accelerator_host_interconnect": "PCIe 5.0 x16",
+    "accelerator_frequency": "",
+    "accelerator_on-chip_memories": "",
+    "accelerator_memory_configuration": "HBM3",
+    "accelerator_memory_capacity": "192GB",
+    "accelerator_interconnect": "",
+    "accelerator_interconnect_topology": "",
+    "cooling": "air",
+    "hw_notes": "",
+    "framework": "tinygrad, branch mlperf_training_v5.0",
+    "other_software_stack": {
+        "python": "3.10.16",
+        "ROCm": "3.0.0+94441cb"
+    },
+    "operating_system": "Ubuntu 24.04.1 LTS",
+    "sw_notes": ""
+  }
@@ -0,0 +1,38 @@
+{
+    "submitter": "tinycorp",
+    "division": "closed",
+    "status": "Available on-premise",
+    "system_name": "tinybox 8xMI350X",
+    "number_of_nodes": "1",
+    "host_processors_per_node": "2",
+    "host_processor_model_name": "AMD EPYC 9575F",
+    "host_processor_core_count": "32",
+    "host_processor_vcpu_count": "64",
+    "host_processor_frequency": "",
+    "host_processor_caches": "",
+    "host_processor_interconnect": "",
+    "host_memory_capacity": "3072 GiB",
+    "host_storage_type": "NVMe SSD",
+    "host_storage_capacity": "4TB",
+    "host_networking": "",
+    "host_networking_topology": "",
+    "host_memory_configuration": "24x 128GB DDR5",
+    "accelerators_per_node": "8",
+    "accelerator_model_name": "AMD Instinct MI350X 288GB HBM3e",
+    "accelerator_host_interconnect": "PCIe 5.0 x16",
+    "accelerator_frequency": "",
+    "accelerator_on-chip_memories": "",
+    "accelerator_memory_configuration": "HBM3",
+    "accelerator_memory_capacity": "288GB",
+    "accelerator_interconnect": "",
+    "accelerator_interconnect_topology": "",
+    "cooling": "air",
+    "hw_notes": "",
+    "framework": "tinygrad, branch mlperf_training_v6.0",
+    "other_software_stack": {
+        "python": "3.12.3",
+        "ROCm": "7.1.1"
+    },
+    "operating_system": "Ubuntu 24.04.3 LTS",
+    "sw_notes": ""
+  }
@@ -0,0 +1,38 @@
+{
+  "submitter": "tinycorp",
+  "division": "closed",
+  "status": "Available on-premise",
+  "system_name": "tinybox green",
+  "number_of_nodes": "1",
+  "host_processors_per_node": "1",
+  "host_processor_model_name": "AMD EPYC 7532",
+  "host_processor_core_count": "32",
+  "host_processor_vcpu_count": "64",
+  "host_processor_frequency": "",
+  "host_processor_caches": "",
+  "host_processor_interconnect": "",
+  "host_memory_capacity": "128GB",
+  "host_storage_type": "NVMe SSD",
+  "host_storage_capacity": "4 TB raid array + 1 TB boot",
+  "host_networking": "",
+  "host_networking_topology": "",
+  "host_memory_configuration": "8x 16GB DDR4",
+  "accelerators_per_node": "6",
+  "accelerator_model_name": "NVIDIA GeForce RTX 4090",
+  "accelerator_host_interconnect": "PCIe 4.0 x16",
+  "accelerator_frequency": "",
+  "accelerator_on-chip_memories": "",
+  "accelerator_memory_configuration": "GDDR6X",
+  "accelerator_memory_capacity": "24GB",
+  "accelerator_interconnect": "",
+  "accelerator_interconnect_topology": "",
+  "cooling": "air",
+  "hw_notes": "",
+  "framework": "tinygrad, branch mlperf_training_v5.0",
+  "other_software_stack": {
+    "python": "3.10.12",
+    "CUDA": "12.4"
+  },
+  "operating_system": "Ubuntu 22.04.4",
+  "sw_notes": ""
+}
@@ -0,0 +1,37 @@
+{
+  "submitter": "tinycorp",
+  "division": "closed",
+  "status": "Available on-premise",
+  "system_name": "tinybox red",
+  "number_of_nodes": "1",
+  "host_processors_per_node": "1",
+  "host_processor_model_name": "AMD EPYC 7532",
+  "host_processor_core_count": "32",
+  "host_processor_vcpu_count": "64",
+  "host_processor_frequency": "",
+  "host_processor_caches": "",
+  "host_processor_interconnect": "",
+  "host_memory_capacity": "128GB",
+  "host_storage_type": "NVMe SSD",
+  "host_storage_capacity": "4 TB raid array + 1 TB boot",
+  "host_networking": "",
+  "host_networking_topology": "",
+  "host_memory_configuration": "8x 16GB DDR4",
+  "accelerators_per_node": "6",
+  "accelerator_model_name": "AMD Radeon RX 7900 XTX",
+  "accelerator_host_interconnect": "PCIe 4.0 x16",
+  "accelerator_frequency": "",
+  "accelerator_on-chip_memories": "",
+  "accelerator_memory_configuration": "GDDR6",
+  "accelerator_memory_capacity": "24GB",
+  "accelerator_interconnect": "",
+  "accelerator_interconnect_topology": "",
+  "cooling": "air",
+  "hw_notes": "",
+  "framework": "tinygrad, branch mlperf_training_v5.0",
+  "other_software_stack": {
+    "python": "3.10.12"
+  },
+  "operating_system": "Ubuntu 22.04.4",
+  "sw_notes": ""
+}
@@ -80,7 +80,7 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
      # NOTE: since this is part of K, these 2 can be anywhere in the frags and long as a and b match
      a_frag = a_frag.reshape(2, 8)[lane_m, :]
      b_frag = b_frag.reshape(2, 8)[lane_m, :]
-    wmma = UOp.wmma(a_frag, b_frag, acc_frag.after(k), ((16, 16, 16), 'AMD', 32))
+    wmma = UOp.wmma(a_frag, b_frag, acc_frag.after(k), (16, 16, 16), 'AMD', 32)
    acc_store = acc_frag.store(wmma).end(tile_m, tile_n)
  else:
    # registers for LOCAL -> REG
@@ -13,7 +13,7 @@ WMMA_ACC = WMMA_M // LANES_PER_WAVE_M
 THREADS_PER_BLOCK = WARP_SIZE * WAVES_M * WAVES_N
 LDS_PAD = 4  # pad LDS rows to reduce bank conflicts

-WMMA_ARG = ((WMMA_M, WMMA_N, WMMA_K), 'AMD', 32)
+WMMA_ARG = (WMMA_M, WMMA_N, WMMA_K), 'AMD', 32
 LOG2E = math.log2(math.e)

 def warp_shfl_xor(val, offset, lane):
@@ -97,7 +97,7 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
  S_frag = S_reg.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0, 2, 1)[tm1, tn1]
  q_frag = Q_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, D // WMMA_K, WMMA_K)[wave_m, tm1, lane_n, k_qk]
  k_frag = KV_lds_k.reshape(WAVES_N, TN, WMMA_N, D // WMMA_K, WMMA_K)[wave_n, tn1, lane_n, k_qk]
-  qk = UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), WMMA_ARG)
+  qk = UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), *WMMA_ARG)
  qk_done = S_frag.store(qk).end(tm1, tn1).end(k_qk)
  S_reg = S_reg.after(qk_done)

@@ -158,7 +158,7 @@ def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
  acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TD).permute(0, 2, 1)[tm2, tn2]
  p_frag = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_N // WMMA_K, WMMA_K)[wave_m, tm2, lane_n, k_pv]
  v_frag = KV_lds_v.reshape(WAVES_N, TD, WMMA_N, BLOCK_N // WMMA_K, WMMA_K)[wave_n, tn2, lane_n, k_pv]
-  pv = UOp.wmma(p_frag, v_frag, acc_frag.after(k_pv), WMMA_ARG)
+  pv = UOp.wmma(p_frag, v_frag, acc_frag.after(k_pv), *WMMA_ARG)

  # end KV tile loop
  n_tile_end = acc_frag.store(pv).end(tm2, tn2).end(k_pv).barrier().end(n_tile)
@@ -128,6 +128,11 @@ def _mx_block_scale(e8:Tensor) -> Tensor:
  rows, scale_K = e8.shape
  return (e8.cast(dtypes.float32) - 127.0).exp2().reshape(rows, scale_K, 1).expand(rows, scale_K, 32).reshape(rows, scale_K*32)

+def _mx_block_scale_3d(e8:Tensor) -> Tensor:
+  # batched (E, rows, scale_K) dequant scale 2^(e8-127) broadcast to (E, rows, scale_K*32)
+  E, rows, scale_K = e8.shape
+  return (e8.cast(dtypes.float32) - 127.0).exp2().reshape(E, rows, scale_K, 1).expand(E, rows, scale_K, 32).reshape(E, rows, scale_K*32)
+
 counters = {"used":0, "todos":[]}
 def todo(msg:str) -> bool: counters["todos"].append(msg); return False
 def _asm_gemm_report():
@@ -275,10 +280,7 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp, n_scales:int=2, has_grad_amax:bool=
      elif getenv("FUSED_GRAD_QUANTIZE", 0):
        grad_amax_t = Tensor(grad_amax_state, device=a.device)
        g_amax = grad_amax_t
-        g_fp8, _, new_grad_amax, _ = quantize_fp8_delayed(g_t, g_amax)
-        store_effect = next_grad_amax_state.store(new_grad_amax.uop)
-        assert g_fp8.uop.op is Ops.AFTER, f"expected AFTER, got {g_fp8.uop.op}"
-        g_fp8 = Tensor(g_fp8.uop.replace(src=g_fp8.uop.src + (store_effect,)), device=a.device)
+        g_fp8, _ = quantize_fp8_delayed(g_t, g_amax, Tensor(next_grad_amax_state, device=a.device))
      else:
        grad_amax_t = Tensor(grad_amax_state, device=a.device)
        g_amax = grad_amax_t
@@ -1,5 +1,5 @@
 from tinygrad import UOp, dtypes
-from tinygrad.uop.ops import AxisType, Ops, KernelInfo, AddrSpace
+from tinygrad.uop.ops import AxisType, KernelInfo, AddrSpace
 from extra.gemm.amd_uop_matmul import test_matmul

 N = 2048
@@ -27,11 +27,8 @@ def hand_spec_tc_cores():
  acc = acc[0].set(0.0)
  acc = acc[1].set(0.0)

-  # TODO: make this simple
-  wmma_arg = ('WMMA_8_8_8_float_float', (8, 8, 8), dtypes.float, dtypes.float, 'METAL', 32, (((3, 2),), ((3, 2),), ((3, 2),)), ())
-
  acc_load = UOp.stack(acc.after(gk)[0], acc.after(gk)[1])
-  out = UOp(Ops.WMMA, dtypes.float, (a_tc, b_tc, acc_load), arg=wmma_arg)
+  out = UOp.wmma(a_tc, b_tc, acc_load, (8, 8, 8), 'METAL', 32)

  end_loop = UOp.group(*[acc[i].store(out.index(i)) for i in range(2)]).end(gk)

@@ -6,7 +6,7 @@ os.environ["AMD_LLVM"] = "0"
 from tinygrad import Tensor, Context, dtypes, UOp, GlobalCounters
 from tinygrad.helpers import DEBUG, getenv
 from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import AxisType, KernelInfo, Ops
+from tinygrad.uop.ops import AxisType, KernelInfo

 WARP_SIZE = 64

@@ -137,8 +137,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
    acc_load = acc_after[N_inner_loop, M_inner_loop]

    # do WMMA
-    wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
-    out = UOp(Ops.WMMA, dtypes.float, (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg)
+    out = UOp.wmma(Ar[M_inner_loop], Br[N_inner_loop], acc_load, (16, 16, 32), 'AMD', 64)

    # store back the acc
    acc_store = acc[N_inner_loop, M_inner_loop].store(out)
@@ -193,8 +192,7 @@ acc = acc[init_l:=UOp.range(4, 1)].set(0.0, end=init_l)

 # do the wmma
 acc_load = UOp.stack(*[acc.after(K_loop)[i] for i in range(4)])
-wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
-out = UOp(Ops.WMMA, dtypes.float, (A_in, B_in, acc_load), arg=wmma_arg)
+out = UOp.wmma(A_in, B_in, acc_load, (16, 16, 32), 'AMD', 64)

 # store back the acc
 acc = acc.after(UOp.group(*[acc[i].store(out.index(i)) for i in range(4)]).end(K_loop))
@@ -6,7 +6,7 @@ os.environ["AMD_LLVM"] = "0"
 from tinygrad import Tensor, Context, dtypes, UOp, GlobalCounters
 from tinygrad.helpers import DEBUG, getenv
 from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import sint, AxisType, KernelInfo, Ops
+from tinygrad.uop.ops import AxisType, KernelInfo

 WARP_SIZE = 64

@@ -60,8 +60,7 @@ def compute_on_locals(acc:UOp, Asl:UOp, Bsl:UOp, rng:int, afters:tuple[UOp, ...]
  acc_load = acc_after[N_inner_loop, M_inner_loop]

  # do WMMA
-  wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
-  out = UOp(Ops.WMMA, dtypes.float, (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg)
+  out = UOp.wmma(Ar[M_inner_loop], Br[N_inner_loop], acc_load, (16, 16, 32), 'AMD', 64)

  # store back the acc
  acc_store = acc[N_inner_loop, M_inner_loop].store(out)
@@ -2,8 +2,8 @@ from __future__ import annotations
 from typing import cast, Callable, TypeVar, Generic, Any
 import struct, functools, time, collections, itertools
 from dataclasses import replace, dataclass
-from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, dedup, pluralize
-from tinygrad.helpers import to_tuple, round_up, partition, data64_le, panic
+from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, dedup, pluralize, JIT_BATCH_SIZE
+from tinygrad.helpers import to_tuple, round_up, partition, data64_le, panic, ContextVar
 from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer
 from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, track_rewrites, GroupOp
 from tinygrad.uop.symbolic import symbolic
@@ -19,6 +19,8 @@ HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQ2Compiled')
 # *****************
 # 0. helpers

+HCQ_RUNTIME_DEV = ContextVar("HCQ_RUNTIME_DEV", "CPU")
+
 HCQ_DEVS = frozenset(("AMD",))
 HCQ_P2P_DEVS = HCQ_DEVS | frozenset(("CPU",))
 HCQ_CACHE_TAGS = frozenset(("program", "systems", "template"))
@@ -56,7 +58,7 @@ def make_patch(buf:UOp, off:sint, val:UOp, dtype=None) -> UOp:
  return buf.index(UOp.const(dtypes.int, off // buf.dtype.itemsize)).store(val.simplify().cast(dtype or buf.dtype))

 def make_binary_patch(buf:UOp, blob:bytes) -> UOp:
-  data = UOp(Ops.BITCAST, buf.dtype, (UOp(Ops.BINARY, dtypes.uint8, src=(), arg=blob),))
+  data = UOp(Ops.BINARY, src=(), arg=blob).bitcast(buf.dtype)
  r = UOp.range(len(blob) // buf.dtype.itemsize, 0, dtype=dtypes.int, src=(buf, data))
  return buf.index(r).store(data.index(r).load()).end(r)

@@ -100,159 +102,138 @@ def stage_copy(dst:UOp, src:UOp) -> UOp|None:
 pm_insert_copy_staging = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy)])

 # *****************
-# 2.1. tag hcq calls
-
-def tag_hcq_call(ctx:itertools.count, call:UOp) -> UOp:
-  if (hcq_devs:=next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None)) is None: return call
-
-  queue = "COMPUTE:0" if call.src[0].op is Ops.PROGRAM else "COPY:0"
-  info = HCQInfo(get_call_name(call, get_call_arg_uops(call)), estimate_uop(call), to_tuple(hcq_devs), queue)
-  return call.replace(arg=replace(call.arg, aux=info)).rtag(next(ctx))
-pm_tag_hcq_calls = PatternMatcher([(UPat(Ops.LINEAR, name="l"), lambda ctx, l: l.replace(src=tuple(tag_hcq_call(ctx, s) for s in l.src)))])
-
-# *****************
-# 2.2. deps tracking
-# device.timeline_signal/value are the per-device schedule epoch. Before a schedule queue accesses memory owned by device N for the first time,
-# it waits for device[N].timeline_signal >= device[N].timeline_value - 1. This orders the schedule after all prior schedules that touched device N.
-#
-# queue.timeline_signal/value are per-queue progress counters used only inside a schedule.
-# Only the owner queue signals its queue.timeline_signal. Values are monotonic.
-#
-# At schedule end, one finalizer queue per touched device[N] waits for every active queue on device[N] to reach its schedule-local
-# final queue.timeline value, then signals device[N].timeline_signal with the schedule's reserved device epoch. After that, buffers/transients
-# for device N from this schedule are safe for the next schedule
-#
-# C programs reserve and bump timeline values, then patch command buffers with the concrete wait/signal values.
+# 2. deps

 class HCQDepsTracker(DepsTracker):
  @staticmethod
  def _key(buf:Any) -> tuple[Any, int, int]:
    return (buf.arg.slot, 0, buf.max_numel() * buf.dtype.itemsize) if isinstance(buf, UOp) else DepsTracker._key(buf)

-def make_deps(u:UOp, dep_lanes:list[tuple[UOp, int, int]], nlanes:int) -> UOp:
-  deps:dict[UOp, list[int|None]] = collections.defaultdict(lambda: [None]*nlanes)
-  for dep, dlane, lane in dep_lanes: deps[dep][lane] = dlane
-  return u.after(*deps, arg=tuple(tuple(v) for v in deps.values()))
-
-def sched_sync(ctx:DepsTracker, call:UOp) -> UOp|None:
-  if not isinstance(call.arg.aux, HCQInfo): return None
-
+def _get_call_bufs_by_lane(call:UOp, devices:tuple[str, ...]) -> list[list[Any]]:
  refs = get_call_arg_uops(call)
-  outs, _ = get_call_outs_ins(call)
-  devices, queue = call.arg.aux.device, call.arg.aux.queue
+  return [[b if b.op is Ops.PARAM else mb.bufs[lane] if isinstance(mb:=b.buffer, MultiBuffer) else mb for b in refs] for lane in range(len(devices))]

-  dep_lanes:list[tuple[UOp, int, int]] = []
-  for lane, d in enumerate(devices):
-    lane_refs = [b if b.op is Ops.PARAM else mb.bufs[lane] if isinstance(mb:=b.buffer, MultiBuffer) else mb for b in refs]
-    for dep, dlane in ctx.access_resources(lane_refs, outs, (call, lane)): dep_lanes.append((dep, dlane, lane))
+def _get_deps(ctx:DepsTracker, bufs_by_lane:list[list[Any]], write, key:tuple[tuple[str, ...], str, int]) -> list[tuple[tuple, int, int]]:
+  dep_lanes:list[tuple[tuple, int, int]] = []
+  for lane, bufs in enumerate(bufs_by_lane):
+    dep_lanes += [(dep, dlane, lane) for dep, dlane in ctx.access_resources(bufs, write if write is not None else range(len(bufs)), (key, lane))]
+  return dep_lanes

+def _build_wait_cmds(dep_lanes:list[tuple[tuple, int, int]], devices:tuple[str, ...], queue:str) -> tuple[list[UOp], set[int]]:
+  # opt1: same-queue ops are fifo-ordered
  if devices[0].split(":")[0] in {"AMD", "QCOM"} or queue.startswith("COPY"):
-    dep_lanes = [(dep, dlane, lane) for dep, dlane, lane in dep_lanes if (dep.arg.aux.device[dlane], dep.arg.aux.queue) != (devices[lane], queue)]
+    dep_lanes = [(dep, dlane, lane) for dep, dlane, lane in dep_lanes if (dep[0][dlane], dep[1]) != (devices[lane], queue)]

-  # keep latest dep per (dep device, queue, cur lane)
-  latest = {((dep.arg.aux.device[dlane], dep.arg.aux.queue), lane): (dep, dlane) for dep, dlane, lane in sorted(dep_lanes, key=lambda x: x[0].tag)}
-  return make_deps(call, [(dep, dlane, lane) for (_, lane), (dep, dlane) in latest.items()], len(devices))
-pm_sched_sync = PatternMatcher([(UPat(Ops.CALL, name="call"), sched_sync)])
+  # opt2: keep latest dep per (dep device, queue, cur lane)
+  latest = {((dep[0][dlane], dep[1]), lane): (dep, dlane) for dep, dlane, lane in sorted(dep_lanes, key=lambda x: x[0][2])}
+  deps:dict[tuple, list[int|None]] = collections.defaultdict(lambda: [None]*len(devices))
+  for (_, lane), (dep, dlane) in latest.items(): deps[dep][lane] = dlane
+
+  waits = []
+  for (ddevs, dqueue, dtag), lanes in deps.items():
+    sig = make_mstack([make_signal(d if dl is None else ddevs[dl], queue=dqueue, sentinel=dl is None) for dl, d in zip(lanes, devices)])
+    val = make_mstack([make_signal_value(d if dl is None else ddevs[dl], queue=dqueue) for dl, d in zip(lanes, devices)])
+    waits.append((sig.index(zero:=UOp.const(dtypes.int, 0)).load() >= val.index(zero) + dtag).wait())
+  return waits, {dtag for _, _, dtag in deps}
+
+def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[tuple[tuple[str, ...], str]],
+                      tracker:HCQDepsTracker) -> tuple[list[UOp], set[int]]:
+  # collect all buffers which belong to devices
+  dev_bufs:dict[str, dict[int, Any]] = collections.defaultdict(dict)
+  for call, devices in batch:
+    for b in itertools.chain.from_iterable(_get_call_bufs_by_lane(call, devices)):
+      for bd in to_tuple(b.device): dev_bufs[bd][id(b)] = b
+
+  zero, n, finalizers, waited = UOp.const(dtypes.int, 0), len(batch_info), [], set()
+  for _, devgroup in itertools.groupby(sorted(dedup([d for devs, _ in batch_info for d in devs])), key=lambda d: d.split(":")[0]):
+    devs = tuple(devgroup)
+
+    # to finalize the batch, sync all accesses from other devices to buffers that belong to this device
+    fin_deps = [dl for dl in _get_deps(tracker, [list(dev_bufs[d].values()) for d in devs], None, key=(devs, "COMPUTE:0", n)) if dl[0][2] < n]
+    waits, cur_waited = _build_wait_cmds(fin_deps, devs, "COMPUTE:0")
+    waited |= cur_waited
+
+    # wait the syncs, store the device epoch; value bumps are a separate call: no lane may bump until every lane has patched its waits
+    submit = make_submit(*waits, make_signal(devs).store((tl:=make_signal_value(devs)).index(zero)), devs=devs, queue="COMPUTE:0")
+    upd = [(tl, 1)] + [(make_signal_value(devs, queue=qn), n) for qn in dedup([qn for bdevs, qn in batch_info if set(bdevs) & set(devs)])]
+    bump = UOp.barrier(*[s.index(zero, dtype=s.dtype).store(s.index(zero) + inc) for s, inc in upd])
+    finalizers += [UOp.custom_function("hcq", b.sink()).call(aux=HCQInfo("hcq_finalizer", Estimates(), devs, "COMPUTE:0")) for b in (submit, bump)]
+  return finalizers, waited
+
+def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]]) -> list[UOp]:
+  batch_info = [(devices, "COMPUTE:0" if call.src[0].op is Ops.PROGRAM else "COPY:0") for call, devices in batch]
+
+  # schedule deps
+  waited:set[int] = set()
+  deps_tracker = HCQDepsTracker()
+  call_waits:list[list[UOp]] = []
+  for tag, ((call, _), (devices, queue)) in enumerate(zip(batch, batch_info)):
+    deps = _get_deps(deps_tracker, _get_call_bufs_by_lane(call, devices), get_call_outs_ins(call)[0], key=(devices, queue, tag))
+    cmds, cur_waited = _build_wait_cmds(deps, devices, queue)
+    call_waits.append(cmds)
+    waited |= cur_waited
+
+  # build finalizers
+  finalizers, finalizer_waited = _build_finalizers(batch, batch_info, deps_tracker)
+  waited |= finalizer_waited
+
+  src = []
+  for tag, ((call, _), (devices, queue), cmds) in enumerate(zip(batch, batch_info, call_waits)):
+    # first queue use, sync prior device work with main signal
+    if batch_info.index((devices, queue)) == tag:
+      epoch = (make_signal(devices).index(0).load() >= make_signal_value(devices).index(0) - 1).wait()
+      cmds = [UOp(Ops.BARRIER), epoch] + cmds
+
+    # signal queue timeline if someone waits for us
+    store = make_signal(devices, queue=queue).store(make_signal_value(devices, queue=queue).index(0) + tag) if tag in waited else None
+
+    # and make hcq call
+    info = HCQInfo(get_call_name(call, get_call_arg_uops(call)), estimate_uop(call), devices, queue)
+    cmds = [*cmds, call.replace(arg=replace(call.arg, aux=info))] + ([store] if store is not None else [])
+    src.append(UOp.custom_function("hcq", make_submit(*cmds, devs=devices, queue=queue).sink()).call(name="hcq", aux=info))
+  return src + finalizers
+
+def sched_hcq_batches(l:UOp) -> UOp:
+  srcs:list[UOp] = []
+  batch:list[tuple[UOp, tuple[str, ...]]] = []
+  for call in l.src:
+    if (devs:=next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None)) is not None: batch.append((call, to_tuple(devs)))
+    else: srcs, batch = srcs + _finalize_batch(batch) + [call], []
+  return l.replace(src=tuple(srcs + _finalize_batch(batch)))
+pm_sched_hcq_batches = PatternMatcher([(UPat(Ops.LINEAR, name="l"), sched_hcq_batches)])

 # *****************
-# 2.3. merge into queues
+# 3. merge into queues

-def _merged_hcq_call(calls:list[UOp]):
-  info = replace(unwrap_after(calls[0]).arg.aux, estimates=sum((unwrap_after(c).arg.aux.estimates for c in calls), start=Estimates()))
-  cmdbuf = make_submit(*calls, devs=info.device, queue=info.queue)
-  return UOp.custom_function("hcq", cmdbuf.sink()).call(name="hcq", aux=info)
+def _merged_hcq_call(calls:list[UOp]) -> UOp: # TODO: simplify?
+  if len(calls) == 1: return calls[0]
+  info = replace(calls[0].arg.aux, name=f"submit {calls[0].arg.aux.queue} ({len(calls)})",
+                 estimates=sum((c.arg.aux.estimates for c in calls), start=Estimates()))
+  cmds = [cmd for c in calls for cmd in get_submit(c).src[0].src]
+  return UOp.custom_function("hcq", make_submit(*cmds, devs=info.device, queue=info.queue).sink()).call(name="hcq", aux=info)

 def merge_queues(linear:UOp) -> UOp:
  new_src:list[UOp] = []
-  opened_qs:dict[tuple[tuple[str, ...], str], list[UOp]] = {} # (devs, queue) -> list of calls, kept in submit order
+  opened_qs:dict[tuple[tuple[str, ...], str], list[UOp]] = {} # (devs, queue) -> list of hcq calls, kept in submit order
+  limits = collections.defaultdict(lambda: JIT_BATCH_SIZE.value)

  for call in linear.src:
-    if not isinstance(unwrap_after(call).arg.aux, HCQInfo):
+    if not isinstance(info:=call.arg.aux, HCQInfo) or info.name == "hcq_finalizer": # non-hcq call or finalizer: close all open queues
      new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in list(opened_qs)] + [call]
      continue

-    devices, queue = unwrap_after(call).arg.aux.device, unwrap_after(call).arg.aux.queue
-
-    if (old:=opened_qs.pop((devices, queue), None)) is not None: new_rec = old + [call]
+    if (old:=opened_qs.pop(key:=(info.device, info.queue), None)) is not None:
+      if limits[key] and len(old) >= limits[key]: new_src, old, limits[key] = new_src + [_merged_hcq_call(old)], [], limits[key] * 2
+      new_rec = old + [call]
    else:
      # no such queue opened: close every open submit on this queue that shares a device, so submit order is kept
-      closing = [k for k in opened_qs if k[1] == queue and set(k[0]) & set(devices)]
+      closing = [k for k in opened_qs if k[1] == info.queue and set(k[0]) & set(info.device)]
      new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in closing]
      new_rec = [call]
-    opened_qs[(devices, queue)] = new_rec
+    opened_qs[(info.device, info.queue)] = new_rec
  return linear.replace(src=tuple(new_src + [_merged_hcq_call(c) for c in opened_qs.values()]))
 pm_merge_queues = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), merge_queues)])

-# *****************
-# 2.4. finalizer
-
-def add_finalizer(ctx:itertools.count, linear:UOp) -> UOp:
-  # collect by device type
-  parts:dict[str, list[UOp]] = collections.defaultdict(list)
-  for call in linear.src:
-    if (c:=unwrap_after(call)).src[0].op is not Ops.CUSTOM_FUNCTION or c.src[0].arg != "hcq": continue
-    parts[c.arg.aux.device[0].split(':')[0]].append(unwrap_after(get_submit(call).src[0].src[0]))
-
-  nbump = next(ctx)
-  finalizers = []
-  for calls in parts.values():
-    devs = tuple(dedup(d for call in calls for d in unwrap_after(call).arg.aux.device))
-    zero = UOp.const(dtypes.int, 0)
-    tl = make_signal_value(devs)
-
-    # split each (multi-device) call into per-device deps, then store the device timeline value into the device signal after them
-    dep_lanes = [(call, dlane, devs.index(d)) for call in calls for dlane, d in enumerate(unwrap_after(call).arg.aux.device)]
-    store = make_deps(make_signal(devs).store(tl.index(zero)), dep_lanes, len(devs))
-    submit = make_submit(store, devs=devs, queue="COMPUTE:0")
-
-    upd = [(tl, 1)] + [(make_signal_value(devs, queue=qn), nbump) for qn in dedup([unwrap_after(call).arg.aux.queue for call in calls])]
-    patches = [s.after(submit).index(zero, dtype=s.dtype).store(s.index(zero) + inc) for s, inc in upd]
-    finalizers.append(UOp.custom_function("hcq", UOp.barrier(*patches).sink()).call(aux=HCQInfo("hcq finalizer", Estimates(), devs, "COMPUTE:0")))
-  return linear.replace(src=linear.src + tuple(finalizers))
-pm_add_finalizer = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), add_finalizer)])
-
-# *****************
-# 2.5. global sync
-
-def add_global_sync(ctx:set[tuple[str, ...]], submit:UOp, q:UOp) -> UOp|None:
-  if (devs:=q.arg[0]) in ctx: return None
-  ctx.add(devs)
-
-  # some devices from a command buffer might be used for the first time this schedule, so we wait for their global timeline epoch.
-  wait = make_signal(devs).wait(make_signal_value(devs).index(UOp.const(dtypes.int, 0)) - 1)
-  return submit.replace(src=(q.replace(src=(UOp(Ops.BARRIER, dtypes.void), wait, *q.src)),))
-pm_add_global_sync = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),), name="submit"), add_global_sync)])
-
-# *****************
-# 3.1. lower loads/stores
-
-def add_loads(ctx:set[int], submit:UOp, q:UOp) -> UOp|None:
-  cur_devs = q.arg[0]
-  new_src:list[UOp] = []
-  for s in q.src:
-    if s.op is Ops.AFTER:
-      for lanes, dep in zip(s.arg, s.src[1:]):
-        devs, queue = dep.arg.aux.device, dep.arg.aux.queue
-        ctx.add(dep.tag) # mark op to update signal.
-
-        sig = make_mstack([make_signal(d if dl is None else devs[dl], queue=queue, sentinel=dl is None) for dl, d in zip(lanes, cur_devs)])
-        val = make_mstack([make_signal_value(d if dl is None else devs[dl], queue=queue) for dl, d in zip(lanes, cur_devs)]).index(UOp.const(dtypes.int, 0))
-        new_src.append(sig.wait(val + dep.tag))
-      s = s.src[0]
-    new_src.append(s)
-  return submit.replace(src=(q.replace(src=tuple(new_src)),))
-pm_add_inner_loads = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),), name="submit"), add_loads)])
-
-def add_stores(ctx:set[int], submit:UOp, q:UOp) -> UOp|None:
-  devs, queue = q.arg
-  new_src:list[UOp] = []
-  for op in q.src:
-    new_src.append(op)
-    if (sigval:=unwrap_after(op).tag) in ctx:
-      new_src.append(make_signal(devs, queue=queue).store(make_signal_value(devs, queue=queue).index(UOp.const(dtypes.int, 0)) + sigval))
-  return submit.replace(src=(q.replace(src=tuple(new_src)),))
-pm_add_inner_stores = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),), name="submit"), add_stores)])
-
 # *****************
 # 4.1. hcq lowering: programs

@@ -277,7 +258,7 @@ def is_value_known_at_link(val:UOp) -> bool:
  addressed_bufs = [b for g in val.toposort() if g.op is Ops.GETADDR for b in unwrap_mstack(g.buf_uop)]

  # addr of input params is not known at link time
-  return not runtime_reads and all(b.op is not Ops.PARAM or b.tag is not None for b in addressed_bufs)
+  return not val.variables() and not runtime_reads and all(b.op is not Ops.PARAM or b.tag is not None for b in addressed_bufs)

 def is_link_patch(p:UOp, jit:bool) -> bool:
  store = p.src[0] if (is_binary_patch:=p.op is Ops.END) else p
@@ -303,35 +284,32 @@ pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION

 # *****************

-def make_addr_table(call:UOp, gaddrs:list[UOp], name:str):
+def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[dict[UOp, UOp], tuple[UOp, ...]]:
  bare = {g: g.replace(src=(unwrap_after(g.src[0]),)) for g in gaddrs}

-  order = sorted(dedup(bare.values()), key=lambda g: ((b:=unwrap_mstack(g.buf_uop)[0]).arg.slot, to_tuple(b.tag)))
+  order = sorted(dedup(bare.values()), key=lambda g: ((b:=unwrap_mstack(g.buf_uop)[0]).arg.slot, repr(b.tag)))
  slots, table = {g:i for i,g in enumerate(order)}, make_placeholder(call.arg.aux.device, len(order), dtypes.uint64, name)

  reads = {g: table.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(UOp.const(dtypes.int, slots[bare[g]])).load() for g in gaddrs}
  return reads, (table.after(*[make_patch(table, i * table.dtype.itemsize, addr) for addr, i in slots.items()]),) if slots else ()

-def rm_rt_getaddrs(call:UOp) -> UOp|None:
-  if not (gaddrs:=[u for u in call.src[0].toposort() if u.op is Ops.GETADDR]): return None
+def make_blob_bufs(call:UOp, blobs:list[UOp]) -> tuple[dict[UOp, UOp], tuple[UOp, ...]]:
+  bufs = {b: make_placeholder(call.arg.aux.device, b.max_numel(), b.dtype, "template") for b in blobs}
+  return bufs, tuple(buf.after(make_binary_patch(buf, b.src[0].arg)) for b,buf in bufs.items())
+
+def rm_rt_uops(call:UOp) -> UOp|None:
+  if not (rt_uops:=[u for u in call.src[0].toposort() if u.op is Ops.GETADDR or (u.op is Ops.BITCAST and u.src[0].op is Ops.BINARY)]): return None
+  gaddrs, blobs = partition(rt_uops, lambda u: u.op is Ops.GETADDR)
  inputs, internals = partition(gaddrs, lambda g: all(x.op is Ops.PARAM and x.tag is None for x in unwrap_mstack(g.buf_uop)))
  runtimes, systems = partition(internals, lambda g: any(x.tag in {"program", "kernargs", "cmdbuf"} for x in unwrap_mstack(g.buf_uop)))

  # exec fills the inputs table with the input addresses every run, so it has no fill patches
-  (input_reads, _), (rt_reads, rt_fills), (sys_reads, sys_fills) = (make_addr_table(call, gs, name) for gs, name in
-    ((inputs, "inputs"), (runtimes, "runtime"), (systems, "systems")))
-  return call.replace(src=(call.src[0].substitute(input_reads | rt_reads | sys_reads), *call.src[1:], *rt_fills, *sys_fills),
+  (reads, _), *tables = [make_addr_table(call, gs, n) for gs,n in ((inputs, "inputs"), (runtimes, "runtime"), (systems, "systems"))] + \
+                        [make_blob_bufs(call, blobs)]
+  reads, fills = reads | {k:v for r,_ in tables for k,v in r.items()}, [f for _,fs in tables for f in fs]
+  return call.replace(src=(call.src[0].substitute(reads), *call.src[1:], *fills),
                      arg=replace(call.arg, aux=replace(call.arg.aux, input_idxs=tuple(sorted(dedup(g.buf_uop.arg.slot for g in inputs))))))
-pm_rm_rt_getaddrs = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), rm_rt_getaddrs)])
-
-# *****************
-
-def rm_rt_binaries(call:UOp) -> UOp|None:
-  if not (blobs:=[u for u in call.src[0].toposort() if u.op is Ops.BITCAST and u.src[0].op is Ops.BINARY]): return None
-  blob_bufs = {blob: make_placeholder(call.arg.aux.device, blob.max_numel(), blob.dtype, "template") for blob in blobs}
-  fills = [buf.after(make_binary_patch(buf, blob.src[0].arg)) for blob, buf in blob_bufs.items()]
-  return call.replace(src=(call.src[0].substitute(blob_bufs), *call.src[1:], *fills))
-pm_rm_rt_binaries = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), rm_rt_binaries)])
+pm_rm_rt_uops = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), rm_rt_uops)])

 # *****************

@@ -383,7 +361,7 @@ pm_pack_placeholders = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNC
 # 8. callify hcq programs

 pm_callify_hcq = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="hcq", src=(UPat(Ops.SINK),), name="cf"),
-  lambda cf: cf.replace(src=(to_program(cf.src[0].replace(arg=KernelInfo("hcq_submit"), tag=1), Device["CPU"].renderer),)))])
+  lambda cf: cf.replace(src=(to_program(cf.src[0].replace(arg=KernelInfo("hcq_submit"), tag=1), Device[HCQ_RUNTIME_DEV.value].renderer),)))])

 hcq_compile_cache:dict[tuple[bytes, bool], UOp] = {}

@@ -392,27 +370,23 @@ def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None, jit=False) -> UOp:
  if input_uops is not None: linear = graph_rewrite(linear, pm_replace_buffers, ctx=input_uops, walk=True, enter_calls=True, name="replace buffer")

  if (final_linear:=(hcq_compile_cache.get(cache_key:=(linear.key, jit)))) is None:
-    # schedule
+    # prep
    linear = linear.substitute(back_map:={s.param_like(i): s for i,s in enumerate(input_uops)} if input_uops is not None else {}, walk=True)
    linear = graph_rewrite(linear, pm_insert_copy_staging + pm_flatten_linear, name="insert copy staging")
-    linear = graph_rewrite(linear, pm_tag_hcq_calls, ctx=(enumerator:=itertools.count(0)), walk=True, name="tag hcq calls")
-    linear = graph_rewrite(linear, pm_sched_sync, ctx=HCQDepsTracker(), walk=True, name="schedule sync")
-    linear = linear.substitute({s: p for p, s in back_map.items()}, walk=True)
+
+    # schedule
+    linear = graph_rewrite(linear, pm_sched_hcq_batches, walk=True, name="schedule hcq batches")
+    linear = linear.substitute({s: p for p, s in back_map.items()}, walk=True, enter_calls=True)
    linear = graph_rewrite(linear, pm_merge_queues, walk=True, name="merge queues")
-    linear = graph_rewrite(linear, pm_add_finalizer, ctx=enumerator, walk=True, name="add finalizer")
-    linear = graph_rewrite(linear, pm_add_global_sync, ctx=set(), walk=True, name="add global sync", enter_calls=True)

    # lowering to hcq ir
-    linear = graph_rewrite(linear, pm_add_inner_loads, ctx=(waited:=set()), walk=True, name="add loads", enter_calls=True)
-    linear = graph_rewrite(linear, pm_add_inner_stores, ctx=waited, walk=True, name="add stores", enter_calls=True)
    linear = graph_rewrite(linear, pm_encode_cmdbufs, walk=True, name="encode cmdbufs", enter_calls=True)
    linear = graph_rewrite(linear, pm_pack_placeholders, walk=True, name="pack placeholders")

    # pie
    linear = graph_rewrite(linear, pm_split_patches, ctx=jit, walk=True, name="split rt/lt patches")
    linear = graph_rewrite(linear, pm_early_simplify + symbolic, bottom_up=False, name="simplify packed placeholders", enter_calls=True)
-    linear = graph_rewrite(linear, pm_rm_rt_getaddrs, walk=True, name="replace rt getaddrs")
-    linear = graph_rewrite(linear, pm_rm_rt_binaries, walk=True, name="replace rt binaries")
+    linear = graph_rewrite(linear, pm_rm_rt_uops, walk=True, name="replace rt uops")
    linear = graph_rewrite(linear, pm_replace_params, walk=True, name="replace with args")

    # and compile it
@@ -425,7 +399,7 @@ def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None, jit=False) -> UOp:

 def bufferize_buf(ctx:bool, buf:UOp) -> UOp|None:
  if buf.tag is None: return None
-  return make_mstack(tuple(UOp.from_buffer((dv:=Device[dev]).pm_bufferize.rewrite(buf, ctx=(dv, ctx)), "CPU") for dev in to_tuple(buf.device)))
+  return make_mstack(tuple(UOp.from_buffer((dv:=Device[dev]).pm_bufferize.rewrite(buf, ctx=(dv, ctx)), HCQ_RUNTIME_DEV.value) for dev in to_tuple(buf.device)))
 pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, name="buf"), bufferize_buf)])

 # *****************
@@ -440,7 +414,8 @@ def fold_binary(buf:UOp, blob:UOp) -> UOp:

 def fold_const_store(buf:UOp, off:UOp, val:UOp) -> UOp:
  for b, v in zip((bs:=mb.bufs if isinstance((mb:=buf.buffer), MultiBuffer) else (mb,)), val.src if val.op is Ops.STACK else (val,)*len(bs)):
-    struct.pack_into(f'<{v.dtype.fmt}', b.ensure_allocated()._buf.cpu_view().mv.cast('B'), off.arg * buf.dtype.itemsize, truncate[v.dtype](v.arg))
+    data = struct.pack(f'<{v.dtype.fmt}', truncate[v.dtype](v.arg))
+    b.ensure_allocated()._buf.cpu_view().view(offset=off.arg * buf.dtype.itemsize, size=len(data), fmt='B')[:] = data
  return UOp(Ops.NOOP)

 def resolve_getaddr(buf:UOp, g:UOp) -> UOp:
@@ -509,7 +484,8 @@ class HCQ2Compiled(Compiled):
    self.rt_allocator = BumpAllocator(64 << 20, wrap=False)

  def new_buffer(self, b:UOp, jit:bool) -> Buffer:
-    if jit or b.tag in HCQ_CACHE_TAGS: return Buffer(self.device, b.max_numel(), b.dtype, options=BufferSpec(cpu_access=True, nolru=True))
+    if jit or b.tag in HCQ_CACHE_TAGS:
+      return Buffer(self.device, b.max_numel(), b.dtype, options=BufferSpec(uncached=True, cpu_access=True, nolru=True))
    return self.rt_buffer.view(b.max_numel(), b.dtype, self.rt_allocator.alloc(b.max_numel() * b.dtype.itemsize, alignment=128))

  @functools.cache
@@ -571,6 +547,10 @@ class HCQ2Buffer:
  def base(self) -> HCQ2Buffer: return self._base or self

 class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
+  def _as_buffer(self, buf:HCQ2Buffer) -> memoryview:
+    self.dev.synchronize()
+    return buf.cpu_view().mv
+
  def _map(self, buf:HCQ2Buffer) -> HCQ2Buffer:
    if not hasattr(self, '_do_map'): raise NotImplementedError("map failed: no method implemented")
    return self._do_map(buf)
@@ -586,24 +566,3 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
    self.dev.iface.free(mb)

  def _offset(self, buf, size:int, offset:int) -> HCQ2Buffer: return buf.offset(offset=offset, size=size)
-
-  def _wrap(self, dev:str, sz:int, opaque:HCQ2Buffer) -> Buffer:
-    return Buffer(dev, sz, dtypes.uint8, opaque=opaque, options=BufferSpec(external_ptr=1))
-
-  def _copy(self, dst:Buffer, src:Buffer):
-    from tinygrad.engine.realize import run_linear
-    du, su = UOp.from_buffer(dst), UOp.from_buffer(src)
-    run_linear(UOp(Ops.LINEAR, src=(su.param_like(1).copy_to_device(dst.device).call(du, su),)), update_stats=True)
-
-  def _copyin(self, dest:HCQ2Buffer, src:memoryview):
-    s = Buffer(self.dev.device, len(src), dtypes.uint8, options=BufferSpec(host=True), preallocate=True)
-    s._buf.cpu_view()[:len(src)] = src
-    self._copy(self._wrap(self.dev.device, len(src), dest), s)
-
-  def _copyout(self, dest:memoryview, src:HCQ2Buffer):
-    d = Buffer(self.dev.device, len(dest), dtypes.uint8, options=BufferSpec(host=True), preallocate=True)
-    self._copy(d, self._wrap(self.dev.device, len(dest), src))
-    self.dev.synchronize()
-    dest[:] = d._buf.cpu_view()[:len(dest)]
-
-  # def _as_buffer(self, buf): return buf.cpu_view().mv
@@ -90,7 +90,7 @@ def memory_barrier(ctx):
                 reg_done=getattr(ctx.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff),
    acquire_mem(ctx)))

-def pm4_wait(ctx, dst, val): return wait_reg_mem(ctx, val, mem=make_getaddr(dst, ctx.devs))
+def pm4_wait(ctx, x, y): return wait_reg_mem(ctx, y, mem=make_getaddr(x.buf_uop, ctx.devs))

 def pm4_barrier(ctx): return memory_barrier(ctx)

@@ -138,7 +138,7 @@ def pm4_program(ctx, call, prg):
 pm_pm4_opsel = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), pm4_program),

-  (UPat(Ops.WAIT, src=(UPat(name="dst"), UPat(name="val"))), pm4_wait),
+  (UPat(Ops.WAIT, src=(UPat.var("x") >= UPat.var("y"),)), pm4_wait),
  (UPat(Ops.BARRIER), pm4_barrier),
  (UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", src=(UPat(name="dst"),)), pm4_timestamp),
  (UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
@@ -184,10 +184,10 @@ def sdma_copy(ctx, call):
     ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz - off, ctx.max_copy_size) - 1), 0,
     *data64_le(src_addr + off), *data64_le(dst_addr + off)) for off in range(0, sz, ctx.max_copy_size)]))

-def sdma_wait(ctx, dst, val):
+def sdma_wait(ctx, x, y):
  op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \
     | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1)
-  return make_ins(SDMAOps.POLL_REGMEM, op, *data64_le(make_getaddr(dst, ctx.devs)), val, 0xffffffff,
+  return make_ins(SDMAOps.POLL_REGMEM, op, *data64_le(make_getaddr(x.buf_uop, ctx.devs)), y, 0xffffffff,
    ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))

 def sdma_store(ctx, dst, val):
@@ -203,7 +203,7 @@ pm_sdma_opsel = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy),

  (UPat(Ops.BARRIER), lambda: UOp(Ops.NOOP, dtypes.void, ())),
-  (UPat(Ops.WAIT, src=(UPat(name="dst"), UPat(name="val"))), sdma_wait),
+  (UPat(Ops.WAIT, src=(UPat.var("x") >= UPat.var("y"),)), sdma_wait),
  (UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", src=(UPat(name="dst"),)), sdma_timestamp),
  (UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), sdma_store),
 ])
@@ -536,7 +536,7 @@ class AMDDevice(HCQ2Compiled):

  timestamp_divider = 100.0  # AMD GPU clock: ticks/us

-  ifaces = [KFDIface, PCIIface]
+  ifaces = [KFDIface, PCIIface, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface)]

  def is_am(self) -> bool: return isinstance(self.iface, (PCIIface,))
  def is_usb(self) -> bool: return False
@@ -20,11 +20,6 @@ def local_abs_max(x:Tensor) -> Tensor:
  fxn = _local_abs_max_fxn(param.uop, x.device)
  return Tensor(fxn[0].uop.call(x.uop).gettuple(0))

-def scalar_amax(amax_buf:Tensor) -> Tensor:
-  if isinstance(amax_buf.device, tuple):
-    return local_abs_max(amax_buf).detach()
-  return amax_buf.max().detach()
-
 def shard_shape(shape:tuple, axis:int, ndev:int) -> list:
  s = list(shape)
  s[axis] //= ndev
@@ -3,7 +3,7 @@ import functools, pathlib
 from tinygrad import Tensor, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
-from extra.llama_kernels import NUM_WG, THREADS_PER_WG, compile_cpp, alloc_like, alloc_local, scalar_amax, dname_of
+from extra.llama_kernels import NUM_WG, THREADS_PER_WG, compile_cpp, alloc_like, dname_of

 # module-level mailbox: grad_xw13 UOp -> (grad_xw13_fp8 UOp, delayed amax UOp)
 # lets cdna_asm_gemm's bwd reuse the fp8 companion produced by the fused silu_mul bwd kernel
@@ -11,13 +11,13 @@ from extra.llama_kernels import NUM_WG, THREADS_PER_WG, compile_cpp, alloc_like,
 _grad_fp8_mailbox:dict[UOp, tuple[UOp, UOp]] = {}

@functools.cache
-def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_buf:UOp, grad_amax:UOp,
+def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_next:UOp, grad_amax:UOp,
                          xw13:UOp, grad_x2:UOp, amax_state:UOp, grad_amax_state:UOp, dname:str) -> UOp:
  hidden = xw13.shape[2] // 2
  n_elems = xw13.shape[0] * xw13.shape[1] * hidden
  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 * 3 + n_elems * 2 + NUM_WG * 4 + 4
-  sink = UOp.sink(grad_xw13_fp8.base, grad_amax_buf.base, grad_amax.base,
+  mem = n_elems * 2 * 3 + n_elems * 2 + 4 + 4
+  sink = UOp.sink(grad_xw13_fp8.base, grad_amax_next.base, grad_amax.base,
                  xw13.base, grad_x2.base, amax_state.base, grad_amax_state.base, threads, workgroups,
                  arg=KernelInfo(f"fused_silu_mul_bwd_w13_{n_elems}", estimates=Estimates(ops=10*n_elems, mem=mem)))
  src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_bwd_w13.cpp", n_elems, hidden)
@@ -25,14 +25,14 @@ def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_buf:UOp, grad_amax:UOp,
                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))

@functools.cache
-def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_buf:UOp, xw13:UOp, amax_state:UOp, grad_amax_state:UOp,
+def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_out:UOp, xw13:UOp, amax_state:UOp, grad_amax_state:UOp,
                                next_grad_amax_state:UOp, dname:str) -> UOp:
  # NOTE: grad_amax_state is plumbed through as an unused fwd input so the bwd kernel can read it via kernel.src
  hidden = xw13.shape[2] // 2
  n_elems = xw13.shape[0] * xw13.shape[1] * hidden
  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 * 2 + n_elems + NUM_WG * 4
-  sink = UOp.sink(fp8_out.base, amax_buf.base, xw13.base, amax_state.base, threads, workgroups,
+  mem = n_elems * 2 * 2 + n_elems + 4
+  sink = UOp.sink(fp8_out.base, amax_out.base, xw13.base, amax_state.base, threads, workgroups,
                  arg=KernelInfo(f"fused_silu_mul_cast_amax_w13_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=mem)))
  src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_fwd_w13.cpp", n_elems, hidden)
  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
@@ -43,26 +43,23 @@ def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
  device = xw13.device
  axis = xw13.axis if isinstance(device, tuple) else None
  grad_xw13_fp8 = alloc_like(xw13.shape, dtypes.fp8e4m3,  device, axis)
-  grad_amax_buf = alloc_local((NUM_WG,), dtypes.float32,  device, axis)
+  grad_amax_next = Tensor(next_grad_amax_state, device=device)
  grad_amax_state_t = Tensor(grad_amax_state, device=device)
  fxn = functools.partial(_custom_fused_bwd_w13, dname=dname_of(device))
  grad_amax = grad_amax_state_t.empty_like()
-  grad_xw13_fp8, grad_amax_buf, grad_amax, *_ = Tensor.custom_kernel(
-    grad_xw13_fp8, grad_amax_buf, grad_amax,
+  grad_xw13_fp8, grad_amax_next, grad_amax, *_ = Tensor.custom_kernel(
+    grad_xw13_fp8, grad_amax_next, grad_amax,
    Tensor(xw13, device=device), Tensor(gradient, device=device).cast(dtypes.bfloat16),
    Tensor(amax_state, device=device), grad_amax_state_t, fxn=fxn)
  grad_xw13_uop = grad_xw13_fp8.uop.cast(dtypes.bfloat16)
-  new_grad_amax = scalar_amax(grad_amax_buf)
-  store_effect = next_grad_amax_state.store(new_grad_amax.uop)
  assert grad_xw13_fp8.uop.op is Ops.AFTER, f"expected AFTER, got {grad_xw13_fp8.uop.op}"
-  grad_xw13_fp8_uop = grad_xw13_fp8.uop.replace(src=grad_xw13_fp8.uop.src + (store_effect,))
  # Stash fp8 companion for cdna_asm_gemm's bwd to attach to grad_a.
-  _grad_fp8_mailbox[grad_xw13_uop] = (grad_xw13_fp8_uop, grad_amax_state_t.uop)
+  _grad_fp8_mailbox[grad_xw13_uop] = (grad_xw13_fp8.uop, grad_amax_state_t.uop)
  return (None, None, grad_xw13_uop, None, None, None)

 def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_state:Tensor,
-                           next_grad_amax_state:Tensor) -> tuple[Tensor, Tensor]:
-  # NOTE: silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns (fp8, new_amax)
+                           next_grad_amax_state:Tensor, amax_out:Tensor) -> Tensor:
+  # NOTE: silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns fp8.
  # grad_amax_state: delayed amax for grad_xw13 fp8 quantization in the backward.
  assert xw13.dtype == dtypes.bfloat16, f"expected bf16, got {xw13.dtype}"
  MBS, SEQ, H2 = xw13.shape
@@ -70,8 +67,7 @@ def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_
  HIDDEN = H2 // 2
  axis = xw13.uop.axis if isinstance(xw13.device, tuple) else None
  fp8_out  = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype,      xw13.device, axis)
-  amax_buf = alloc_local((NUM_WG,),         dtypes.float32, xw13.device, axis)
  fxn = functools.partial(_custom_fused_cast_amax_w13, dname=dname_of(xw13.device))
-  fp8_out, amax_buf, *_ = Tensor.custom_kernel(fp8_out, amax_buf, xw13, amax_state, grad_amax_state, next_grad_amax_state,
+  fp8_out, amax_out, *_ = Tensor.custom_kernel(fp8_out, amax_out, xw13, amax_state, grad_amax_state, next_grad_amax_state,
                                                fxn=fxn, grad_fxn=_fused_quantize_bwd_w13)
-  return fp8_out, scalar_amax(amax_buf)
+  return fp8_out
@@ -23,14 +23,14 @@ static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC");

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

-  // LDS tree reduction: per-workgroup amax
+  // LDS tree reduction: per-workgroup amax, then global atomic into the scalar.
  sdata[tid] = local_max;
  __syncthreads();
  for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
@@ -75,5 +75,5 @@ fused_silu_mul_cast_amax_w13(
    __syncthreads();
  }

-  if (tid == 0) amax_buf[wg] = sdata[0];
+  if (tid == 0 && sdata[0] > *amax_out) atomicMax(reinterpret_cast<int32_t*>(amax_out), __float_as_int(sdata[0]));
 }
@@ -3,19 +3,19 @@ import functools, pathlib
 from tinygrad import Tensor, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
-from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, scalar_amax, dname_of, compile_hip
+from extra.llama_kernels import NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, dname_of, compile_hip

 def _src() -> str: return (pathlib.Path(__file__).parent/"fused_rmsnorm_mul_quantize_fp8.cpp").read_text()
 def _src_bwd() -> str: return (pathlib.Path(__file__).parent/"fused_rmsnorm_mul_quantize_fp8_bwd.cpp").read_text()

@functools.cache
-def _custom_fwd(fp8_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_buf:UOp,
+def _custom_fwd(fp8_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_out:UOp,
                x:UOp, weight:UOp, amax_state:UOp, dname:str, eps_val:float) -> UOp:
  MBS, SEQ, HIDDEN = x.shape
  n_elems = MBS * SEQ * HIDDEN
  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 + n_elems + MBS * SEQ * 4 + n_elems + HIDDEN * 2 + NUM_WG * 4 + 4
-  sink = UOp.sink(fp8_out.base, x_normed_out.base, rrms_out.base, amax_buf.base,
+  mem = n_elems * 2 + n_elems + MBS * SEQ * 4 + n_elems + HIDDEN * 2 + 4 + 4
+  sink = UOp.sink(fp8_out.base, x_normed_out.base, rrms_out.base, amax_out.base,
                  x.base, weight.base, amax_state.base, threads, workgroups,
                  arg=KernelInfo(f"fused_rmsnorm_mul_quantize_fp8_{n_elems}_h{HIDDEN}_eps{eps_val:.0e}",
                                 estimates=Estimates(ops=6*n_elems, mem=mem)))
@@ -26,13 +26,13 @@ def _custom_fwd(fp8_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_buf:UOp,
                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))

@functools.cache
-def _custom_fwd_add(fp8_out:UOp, h_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_buf:UOp,
+def _custom_fwd_add(fp8_out:UOp, h_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_out:UOp,
                    x:UOp, residual:UOp, weight:UOp, amax_state:UOp, dname:str, eps_val:float) -> UOp:
  MBS, SEQ, HIDDEN = x.shape
  n_elems = MBS * SEQ * HIDDEN
  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 * 4 + MBS * SEQ * 4 + HIDDEN * 2 + NUM_WG * 4 + 4
-  sink = UOp.sink(fp8_out.base, h_out.base, x_normed_out.base, rrms_out.base, amax_buf.base,
+  mem = n_elems * 2 * 4 + MBS * SEQ * 4 + HIDDEN * 2 + 4 + 4
+  sink = UOp.sink(fp8_out.base, h_out.base, x_normed_out.base, rrms_out.base, amax_out.base,
                  x.base, residual.base, weight.base, amax_state.base, threads, workgroups,
                  arg=KernelInfo(f"fused_add_rmsnorm_mul_quantize_fp8_{n_elems}_h{HIDDEN}_eps{eps_val:.0e}",
                                 estimates=Estimates(ops=7*n_elems, mem=mem)))
@@ -85,7 +85,7 @@ def _bwd_common(fp8_grad_u, h_grad_u, x_u, x_normed_u, rrms_u, weight_u, amax_st
  return grad_total.uop, grad_weight_uop

 def _fused_bwd(gradient:UOp, kernel:UOp):
-  # NOTE: fwd inputs (fp8_out, x_normed_out, rrms_out, amax_buf, x, weight, amax_state)
+  # NOTE: fwd inputs (fp8_out, x_normed_out, rrms_out, amax_out, x, weight, amax_state)
  _, x_normed_u, rrms_u, _, x_u, weight_u, amax_state_u = kernel.src[1:]
  grad_x, grad_w = _bwd_common(gradient, None, x_u, x_normed_u, rrms_u, weight_u, amax_state_u, kernel)
  return (None, None, None, None, grad_x, grad_w, None)
@@ -112,8 +112,9 @@ def _fused_add_bwd(*args, **kwargs):
  grad_h, grad_w = _bwd_common(fp8_grad_u, h_grad_u, x_u, x_normed_u, rrms_u, weight_u, amax_state_u, kernel)
  return (None, None, None, None, None, grad_h, grad_h, grad_w, None)

-def fused_rmsnorm_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, eps:float, fp8_dtype) -> tuple[Tensor, Tensor, Tensor, Tensor]:
-  # NOTE: rmsnorm(x) * weight -> fp8 + amax. Returns (fp8, new_amax, x_normed, rrms).
+def fused_rmsnorm_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, eps:float, fp8_dtype,
+                                   amax_out:Tensor) -> tuple[Tensor, Tensor, Tensor]:
+  # NOTE: rmsnorm(x) * weight -> fp8 + amax. Returns (fp8, x_normed, rrms).
  # x_normed + rrms are saved for the rmsnorm backward (also recomputed here from x regs).
  assert x.dtype == dtypes.bfloat16 and weight.dtype == dtypes.bfloat16
  assert x.shape[-1] == weight.shape[-1], f"HIDDEN mismatch: x={x.shape}, weight={weight.shape}"
@@ -123,16 +124,15 @@ def fused_rmsnorm_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, e
  fp8_out      = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype,       x.device, axis)
  x_normed_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
  rrms_out     = alloc_like((MBS, SEQ),         dtypes.float32,  x.device, axis)
-  amax_buf     = alloc_local((NUM_WG,),         dtypes.float32,  x.device, axis)
  fxn = functools.partial(_custom_fwd, dname=dname_of(x.device), eps_val=eps)
-  fp8_out, x_normed_out, rrms_out, amax_buf, *_ = Tensor.custom_kernel(
-    fp8_out, x_normed_out, rrms_out, amax_buf, x, weight, amax_state, fxn=fxn, grad_fxn=_fused_bwd)
-  return fp8_out, scalar_amax(amax_buf), x_normed_out, rrms_out
+  fp8_out, x_normed_out, rrms_out, amax_out, *_ = Tensor.custom_kernel(
+    fp8_out, x_normed_out, rrms_out, amax_out, x, weight, amax_state, fxn=fxn, grad_fxn=_fused_bwd)
+  return fp8_out, x_normed_out, rrms_out

 def fused_add_rmsnorm_mul_quantize_fp8(x:Tensor, residual:Tensor, weight:Tensor, amax_state:Tensor,
-                                       eps:float, fp8_dtype) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
+                                       eps:float, fp8_dtype, amax_out:Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]:
  # NOTE: h = x + residual; y_normed = rmsnorm(h); fp8 = quantize(y_normed * weight).
-  # Returns (fp8, new_amax, h, x_normed, rrms). h is also written so downstream can
+  # Returns (fp8, h, x_normed, rrms). h is also written so downstream can
  # reuse it without recomputing x+residual — eliminates the separate residual-add kernel.
  assert x.dtype == dtypes.bfloat16 and residual.dtype == dtypes.bfloat16 and weight.dtype == dtypes.bfloat16
  assert x.shape == residual.shape
@@ -143,9 +143,8 @@ def fused_add_rmsnorm_mul_quantize_fp8(x:Tensor, residual:Tensor, weight:Tensor,
  h_out        = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
  x_normed_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
  rrms_out     = alloc_like((MBS, SEQ),         dtypes.float32,  x.device, axis)
-  amax_buf     = alloc_local((NUM_WG,),         dtypes.float32,  x.device, axis)
  fxn = functools.partial(_custom_fwd_add, dname=dname_of(x.device), eps_val=eps)
-  fp8_out, h_out, x_normed_out, rrms_out, amax_buf, *_ = Tensor.custom_kernel(
-    fp8_out, h_out, x_normed_out, rrms_out, amax_buf, x, residual, weight, amax_state,
+  fp8_out, h_out, x_normed_out, rrms_out, amax_out, *_ = Tensor.custom_kernel(
+    fp8_out, h_out, x_normed_out, rrms_out, amax_out, x, residual, weight, amax_state,
    fxn=fxn, grad_fxn=_fused_add_bwd)
-  return fp8_out, scalar_amax(amax_buf), h_out, x_normed_out, rrms_out
+  return fp8_out, h_out, x_normed_out, rrms_out
@@ -7,7 +7,7 @@
 //   fp8 = fp8_sat(y * (FP8_MAX / amax_state))
 // Also writes:
 //   rrms[row]        — saved for the rmsnorm backward
-//   amax_buf[wg]     — per-WG |y| partials, reduced later to update amax_state
+//   amax_out         — scalar |y| via global atomic max
 //
 // Layout: one WG per row, ROWS_PER_WG rows per WG via grid-stride (ROWS = N_ELEMS / HIDDEN).
 // Each thread handles HIDDEN / THREADS_PER_WG elements per row.
@@ -48,7 +48,7 @@ fused_add_rmsnorm_mul_quantize_fp8(
    __hip_bfloat16*       __restrict__ h_out,           // bf16, ROWS*HIDDEN — x + residual (saved for downstream)
    __hip_bfloat16*       __restrict__ x_normed_out,    // bf16, ROWS*HIDDEN
    float*                __restrict__ rrms_out,        // fp32, ROWS
-    float*                __restrict__ amax_buf,        // fp32, NUM_WG
+    float*                __restrict__ amax_out,        // fp32 scalar, initialized to 0 before launch
    const __hip_bfloat16* __restrict__ x,               // bf16, ROWS*HIDDEN
    const __hip_bfloat16* __restrict__ residual,        // bf16, ROWS*HIDDEN — added into x before rmsnorm
    const __hip_bfloat16* __restrict__ weight,          // bf16, HIDDEN
@@ -60,7 +60,7 @@ fused_rmsnorm_mul_quantize_fp8(
    __hip_fp8_storage_t*  __restrict__ fp8_out,         // fp8, ROWS*HIDDEN
    __hip_bfloat16*       __restrict__ x_normed_out,    // bf16, ROWS*HIDDEN (saved for rmsnorm bwd)
    float*                __restrict__ rrms_out,        // fp32, ROWS (fp32 to match rmsnorm_bwd.cpp expectation)
-    float*                __restrict__ amax_buf,        // fp32, NUM_WG per-WG partials
+    float*                __restrict__ amax_out,        // fp32 scalar, initialized to 0 before launch
    const __hip_bfloat16* __restrict__ x,               // bf16, ROWS*HIDDEN
    const __hip_bfloat16* __restrict__ weight,          // bf16, HIDDEN (per-hidden scale)
    const float*          __restrict__ amax_state)      // fp32 scalar
@@ -144,12 +144,12 @@ fused_rmsnorm_mul_quantize_fp8(
    __syncthreads();  // before next row's sum_sq reduce reuses sdata
  }

-  // Final per-WG amax reduce.
+  // Final per-WG amax reduce, then global atomic into the scalar.
  sdata[tid] = local_max;
  __syncthreads();
  for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
    if (tid < s) sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]);
    __syncthreads();
  }
-  if (tid == 0) amax_buf[wg] = sdata[0];
+  if (tid == 0 && sdata[0] > *amax_out) atomicMax(reinterpret_cast<int32_t*>(amax_out), __float_as_int(sdata[0]));
 }
@@ -3,14 +3,13 @@ from tinygrad import Tensor, dtypes
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import prod
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
-from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, scalar_amax
+from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like

@functools.cache
-def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_partial:UOp, x:UOp, amax_state:UOp) -> UOp:
+def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:UOp, device=None) -> UOp:
  VEC = 8
  n_elems = prod(x.shape)
  assert n_elems % (NUM_WG * THREADS_PER_WG * VEC) == 0
-  assert amax_partial.shape[0] == NUM_WG

  x = x.reshape(n_elems)
  fp8_out = fp8_out.reshape(n_elems)
@@ -46,8 +45,13 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_partial:UOp, x:UOp, amax_st
    lds = lds.after(lds[tid.valid(active)].store(lds[tid].maximum(other)).barrier())
    step //= 2

-  amax_store = amax_partial[tid.eq(0).where(wg, UOp.invalid())].store(lds[0])
-  return amax_store.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=()))
+  device = device[0].split(":")[0] if isinstance(device, tuple) else device.split(":")[0]
+  if device in {"AMD", "NULL"}: atomic_arg = "if ({2} > {3}) __hip_atomic_fetch_max((int*){0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);"
+  else: raise NotImplementedError(f"no atomic max for device {device}")
+  amax_idx = amax_out.reshape((1,)).index(UOp.const(dtypes.index, 0))
+  max_val = lds[0].load()
+  atomic = UOp(Ops.CUSTOM, dtypes.void, (amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=atomic_arg)
+  return atomic.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=()))

@functools.cache
 def _custom_quantize_fp8_scalar(fp8_out:UOp, x:UOp, amax_state:UOp) -> UOp:
@@ -69,25 +73,19 @@ def _quantize_fp8_delayed_bwd(gradient:UOp, kernel:UOp):
  grad_x = (Tensor(gradient, device=device).float() * scale).cast(dtypes.bfloat16)
  return (None, None, grad_x.uop, None)

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

 def quantize_fp8_scalar(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3) -> Tensor:
  # NOTE: pure one-pass bf16 -> fp8 quantize with delayed scalar scale. No amax computation.
@@ -16,6 +16,96 @@ def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None
  axis = ref.uop.axis if axis is None else axis
  return Tensor(Tensor.invalids(*shape, dtype=dtype, device=ref.device).uop.multi(axis), dtype=dtype, device=ref.device)

+@functools.cache
+def custom_fused_qkv_rope_forward(q:UOp, k:UOp, v:UOp, xqkv:UOp, freqs_cis:UOp,
+                                  device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
+  code = (pathlib.Path(__file__).parent / "fused_qkv_rope.cpp").read_text()
+  threads = 256
+  thread_idx = UOp.special(threads, "lidx0")
+  block_idx_x, block_idx_y = UOp.special(B, "gidx0"), UOp.special(N, "gidx1")
+  sink = UOp.sink(q.base, k.base, v.base, xqkv.base, freqs_cis.base, thread_idx, block_idx_x, block_idx_y,
+                  arg=KernelInfo(name="fused_qkv_rope_forward"))
+  compile_args = ["-std=c++20", "-ffast-math", f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}",
+                  f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DTHREADS_PER_BLOCK={threads}"]
+  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
+
+@functools.cache
+def custom_fused_qkv_rope_backward(dxqkv:UOp, dq:UOp, dk:UOp, dv:UOp, freqs_cis:UOp,
+                                   device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
+  assert (B, N, H, H_KV, D) == (2, 8192, 32, 8, 128)
+  code = (pathlib.Path(__file__).parent / "fused_qkv_rope_bwd.cpp").read_text()
+  threads = 256
+  thread_idx = UOp.special(threads, "lidx0")
+  gsz = (B, N // 64, H + 2 * H_KV)
+  block_idx_x, block_idx_y, block_idx_z = (UOp.special(x, f"gidx{i}") for i, x in enumerate(gsz))
+  sink = UOp.sink(dxqkv.base, dq.base, dk.base, dv.base, freqs_cis.base, thread_idx, block_idx_x, block_idx_y, block_idx_z,
+                  arg=KernelInfo(name="fused_qkv_rope_backward"))
+  compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math", f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}",
+                  f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DTHREADS_PER_BLOCK={threads}"]
+  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
+
+def _fa_native_grads(dq:UOp, dk:UOp, dv:UOp) -> tuple[UOp, UOp, UOp]|None:
+  def unwrap_partial(x:UOp) -> UOp|None:
+    expected = (Ops.CAST, Ops.REDUCE, Ops.PERMUTE, Ops.CAST, Ops.RESHAPE, Ops.AFTER)
+    for op in expected:
+      if x.op is not op: return None
+      if op is not Ops.AFTER: x = x.src[0]
+    return x
+  dq_native, dk_partial, dv_partial = dq.base, unwrap_partial(dk), unwrap_partial(dv)
+  if dq_native.op is not Ops.AFTER or dk_partial is None or dv_partial is None: return None
+  B, N, H, D, H_KV = dq.shape[0], dq.shape[1], dq.shape[2], dq.shape[3], dk.shape[2]
+  heads_per_wg = 2 if D == 128 and (H // H_KV) % 2 == 0 else 1
+  partials = (H // H_KV) // heads_per_wg
+  if dq_native.shape != (B, H, N, D) or dk_partial.shape != (B * partials, N, H_KV, D) or dv_partial.shape != dk_partial.shape: return None
+  return dq_native, dk_partial, dv_partial
+
+def _fused_qkv_rope_grad(dq_u:UOp, dk_u:UOp, dv_u:UOp, call:UOp) -> tuple[None, None, None, UOp, None]:
+  dq, dk, dv = Tensor(dq_u, device=dq_u.device), Tensor(dk_u, device=dk_u.device), Tensor(dv_u, device=dv_u.device)
+  xqkv_u, freqs_u = call.src[4], call.src[5]
+  xqkv, freqs_cis = Tensor(xqkv_u, device=xqkv_u.device), Tensor(freqs_u, device=freqs_u.device)
+  B, N, _ = xqkv.shape
+  H, H_KV, D = dq.shape[2], dk.shape[2], dq.shape[3]
+  num_devices = len(xqkv.device) if isinstance(xqkv.device, tuple) else 1
+  is_dp, is_mp = xqkv.uop.axis == 0, xqkv.uop.axis == 2
+  B_local = B // num_devices if is_dp else B
+  H_local = H // num_devices if is_mp else H
+  H_KV_local = H_KV // num_devices if is_mp else H_KV
+  single_device = xqkv.device[0] if isinstance(xqkv.device, tuple) else xqkv.device
+  arch = Device[single_device].renderer.target.arch
+  fa_native = _fa_native_grads(dq_u, dk_u, dv_u)
+  assert fa_native is not None, "fused QKV RoPE backward requires native Flash Attention gradients"
+  dq, dk, dv = (Tensor(x, device=x.device) for x in fa_native)
+  dxqkv = _sharded_empty_like(xqkv, axis=xqkv.uop.axis if isinstance(xqkv.device, tuple) else None)
+  fxn = functools.partial(custom_fused_qkv_rope_backward, device=single_device, arch=arch,
+                          B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D)
+  dxqkv = Tensor.custom_kernel(dxqkv, dq, dk, dv, freqs_cis, fxn=fxn)[0]
+  return None, None, None, dxqkv.uop, None
+
+def fused_qkv_rope(xqkv:Tensor, freqs_cis:Tensor, n_heads:int, n_kv_heads:int, head_dim:int) -> tuple[Tensor, Tensor, Tensor]:
+  B, N, packed_dim = xqkv.shape
+  assert packed_dim == n_kv_heads * (n_heads // n_kv_heads + 2) * head_dim
+  assert freqs_cis.dtype == dtypes.bfloat16, f"fused QKV RoPE requires bfloat16 frequencies, got {freqs_cis.dtype}"
+  assert freqs_cis.shape == (1, freqs_cis.shape[1], 1, head_dim // 2, 2) and freqs_cis.shape[1] >= N, \
+    f"invalid RoPE frequency shape {freqs_cis.shape} for sequence length {N} and head dimension {head_dim}"
+  num_devices = len(xqkv.device) if isinstance(xqkv.device, tuple) else 1
+  is_dp, is_mp = xqkv.uop.axis == 0, xqkv.uop.axis == 2
+  B_local = B // num_devices if is_dp else B
+  H_local = n_heads // num_devices if is_mp else n_heads
+  H_KV_local = n_kv_heads // num_devices if is_mp else n_kv_heads
+  assert H_local % H_KV_local == 0 and head_dim % 2 == 0 and head_dim <= 512
+  single_device = xqkv.device[0] if isinstance(xqkv.device, tuple) else xqkv.device
+  arch = Device[single_device].renderer.target.arch
+  axis = 0 if is_dp else 2 if is_mp else None
+  q = _sharded_empty((B, N, n_heads, head_dim), xqkv, axis=axis, dtype=dtypes.bfloat16)
+  k = _sharded_empty((B, N, n_kv_heads, head_dim), xqkv, axis=axis, dtype=dtypes.bfloat16)
+  v = _sharded_empty((B, N, n_kv_heads, head_dim), xqkv, axis=axis, dtype=dtypes.bfloat16)
+  fxn = functools.partial(custom_fused_qkv_rope_forward, device=single_device, arch=arch,
+                          B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=head_dim)
+  q, k, v, *_ = Tensor.custom_kernel(q, k, v, xqkv, freqs_cis, fxn=fxn, grad_fxn=_fused_qkv_rope_grad)
+  return q, k, v
+
 def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
  return _sharded_empty(ref.shape, ref, axis)

@@ -31,8 +121,9 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha

    dq = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t)
    GROUP_SIZE = H_local // H_KV_local
-    dk_partial = _sharded_empty((B * GROUP_SIZE, N, H_KV, D), xk, axis=shard_axis)
-    dv_partial = _sharded_empty((B * GROUP_SIZE, N, H_KV, D), xv, axis=shard_axis)
+    HEADS_PER_WG = 2 if D == 128 and GROUP_SIZE % 2 == 0 else 1
+    dk_partial = _sharded_empty((B * GROUP_SIZE // HEADS_PER_WG, N, H_KV, D), xk, axis=shard_axis)
+    dv_partial = _sharded_empty((B * GROUP_SIZE // HEADS_PER_WG, N, H_KV, D), xv, axis=shard_axis)

    # delta_vec = (do * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
    delta_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
@@ -46,8 +137,8 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
      dq = dq.reshape(B, H, N//16, 4, 2, 2, D//32, 4, 4, 2).permute(0, 1, 2, 7, 8, 3, 4, 6, 5, 9).reshape(B, H, N, D).transpose(1, 2)

    # reduce partial dK/dV across GROUP_SIZE query heads
-    dk = dk_partial.reshape(B, GROUP_SIZE, N, H_KV, D).sum(1)
-    dv = dv_partial.reshape(B, GROUP_SIZE, N, H_KV, D).sum(1)
+    dk = dk_partial.reshape(B, GROUP_SIZE // HEADS_PER_WG, N, H_KV, D).sum(1)
+    dv = dv_partial.reshape(B, GROUP_SIZE // HEADS_PER_WG, N, H_KV, D).sum(1)

    if not has_sink: return None, None, dq.uop, dk.uop, dv.uop
    sinks = Tensor(ker.src[6], device=ker.src[6].device)
@@ -160,9 +251,11 @@ def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_ve
                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}"]

  BLOCK_SIZE_KV = 256
+  GROUP_SIZE = H // H_KV
+  HEADS_PER_WG = 2 if D == 128 and GROUP_SIZE % 2 == 0 else 1
  NUM_WARPS = 4
  NUM_THREADS = 64 * NUM_WARPS
-  gsz = (H, N // BLOCK_SIZE_KV, B)
+  gsz = (H // HEADS_PER_WG, N // BLOCK_SIZE_KV, B)
  lsz = (NUM_THREADS, 1, 1)
  threadIdx_x = UOp.special(lsz[0], "lidx0")
  blockIdx_x, blockIdx_y, blockIdx_z = UOp.special(gsz[0], "gidx0"), UOp.special(gsz[1], "gidx1"), UOp.special(gsz[2], "gidx2")
@@ -28,6 +28,7 @@ constexpr int ATTN_H_KV = 8; // number of key/value heads (for GQA)
 #endif

 constexpr int GROUP_SIZE = ATTN_H / ATTN_H_KV; // queries per KV head group
+constexpr int HEADS_PER_WG = (ATTN_D == 128 && GROUP_SIZE % 2 == 0) ? 2 : 1;

 #ifndef ATTN_N
 constexpr int ATTN_N = 1024; // sequence length
@@ -53,7 +54,7 @@ using namespace kittens;
 using _gl_QdO  = gl<bf16, ATTN_B, ATTN_N, ATTN_H, ATTN_D>;
 using _gl_KV   = gl<bf16, ATTN_B, ATTN_N, ATTN_H_KV, ATTN_D>;
 using _gl_dQ   = gl<bf16, ATTN_B, ATTN_H, ATTN_N, ATTN_D>;
-using _gl_dKV  = gl<bf16, ATTN_B * GROUP_SIZE, ATTN_N, ATTN_H_KV, ATTN_D>;
+using _gl_dKV  = gl<bf16, ATTN_B * (GROUP_SIZE / HEADS_PER_WG), ATTN_N, ATTN_H_KV, ATTN_D>;
 using _gl_Lvec = gl<float, ATTN_B, ATTN_H, 1, ATTN_N>;

 template<int D> struct attn_bwd_combined_globals {
@@ -63,7 +64,7 @@ template<int D> struct attn_bwd_combined_globals {
  _gl_dQ dQg;
  _gl_dKV dKg, dVg;
  _gl_Lvec L_vec, delta_vec;
-  dim3 grid() { return dim3(ATTN_H, (ATTN_N / BLOCK_SIZE_KV), ATTN_B); }
+  dim3 grid() { return dim3(ATTN_H / HEADS_PER_WG, (ATTN_N / BLOCK_SIZE_KV), ATTN_B); }
  dim3 block() { return dim3(NUM_THREADS); }
  size_t dynamic_shared_memory() { return MAX_SHARED_MEMORY; }
 };
@@ -71,7 +72,7 @@ template<int D> struct attn_bwd_combined_globals {
 template<int D> __launch_bounds__(NUM_THREADS, 1)
 __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr, bf16 *dO_ptr, bf16 *Q_ptr, bf16 *K_ptr, bf16 *V_ptr, float *L_vec_ptr, float *delta_vec_ptr) {

-  const int q_head_idx_fixed = blockIdx.x;  // This is the query head index [0, ATTN_H)
+  const int q_head_idx_fixed = blockIdx.x * HEADS_PER_WG;  // First query head handled by this workgroup.
  const int kv_head_idx = q_head_idx_fixed / GROUP_SIZE;
  const int q_head_in_group = q_head_idx_fixed % GROUP_SIZE;
  const int seq_idx = blockIdx.y;
@@ -88,7 +89,7 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
  // first Q step that can overlap this K_span:
  const int first_step = max(0, k_start_min / STEP_QO);
  const int num_steps_per_head = total_steps_per_head - first_step;
-  const int num_steps = num_steps_per_head;
+  const int num_steps = num_steps_per_head * HEADS_PER_WG;
  const int k_pos = j * WARP_SIZE_KV;

  constexpr float L_SCALE_FACTOR = 1.44269504089f;
@@ -270,12 +271,12 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
  if (num_steps > 1) {
    // Prologue
    {
-      const int q_head_idx = (0) / num_steps_per_head + first_q_head;
-      const int q_seq_idx = ((0) % num_steps_per_head) + first_step;
+      const int q_head_idx = first_q_head;
+      const int q_seq_idx = first_step;
      const int q_pos = q_seq_idx * STEP_QO;

-      const int next_q_head_idx = (0 + 1) / num_steps_per_head + first_q_head;
-      const int next_q_seq_idx = ((0 + 1) % num_steps_per_head) + first_step;
+      const int next_q_head_idx = first_q_head;
+      const int next_q_seq_idx = first_step + 1;

      // dot slice 0
      {
@@ -1332,15 +1333,18 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr

    // 9. for 1 <= i <= T_r (1024 / 32 = 32)  
    for (int i = 1; i < num_steps - 1; ++i, tic ^= 1, toc ^= 1) {
-      const int last_q_head_idx = (i - 1) / num_steps_per_head + first_q_head;
-      const int last_q_seq_idx = ((i - 1) % num_steps_per_head) + first_step;
+      const int last_head_offset = (i - 1) >= num_steps_per_head;
+      const int last_q_head_idx = last_head_offset + first_q_head;
+      const int last_q_seq_idx = i - 1 - last_head_offset * num_steps_per_head + first_step;

-      const int q_head_idx = i / num_steps_per_head + first_q_head;
-      const int q_seq_idx = (i % num_steps_per_head) + first_step;
+      const int head_offset = i >= num_steps_per_head;
+      const int q_head_idx = head_offset + first_q_head;
+      const int q_seq_idx = i - head_offset * num_steps_per_head + first_step;
      const int q_pos = q_seq_idx * STEP_QO;

-      const int next_q_head_idx = (i + 1) / num_steps_per_head + first_q_head;
-      const int next_q_seq_idx = ((i + 1) % num_steps_per_head) + first_step;
+      const int next_head_offset = (i + 1) >= num_steps_per_head;
+      const int next_q_head_idx = next_head_offset + first_q_head;
+      const int next_q_seq_idx = i + 1 - next_head_offset * num_steps_per_head + first_step;

      // dot slice 0
      {
@@ -2378,11 +2382,11 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
    }
  }

-  const int last_q_head_idx = (num_steps - 2) / num_steps_per_head + first_q_head;
-  const int last_q_seq_idx = ((num_steps - 2) % num_steps_per_head) + first_step;
+  const int last_q_head_idx = first_q_head + HEADS_PER_WG - 1;
+  const int last_q_seq_idx = first_step + num_steps_per_head - 2;

-  const int q_head_idx = (num_steps - 1) / num_steps_per_head + first_q_head;
-  const int q_seq_idx = ((num_steps - 1) % num_steps_per_head) + first_step;
+  const int q_head_idx = first_q_head + HEADS_PER_WG - 1;
+  const int q_seq_idx = first_step + num_steps_per_head - 1;
  const int q_pos = q_seq_idx * STEP_QO;
  // Epilogue
  {
@@ -3407,14 +3411,14 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
    }
  }

-  store<1>(g.dVg, dV_j, {batch_idx * GROUP_SIZE + q_head_in_group, 0, kv_head_idx, 0}, {0, j, 0, 0});
+  store<1>(g.dVg, dV_j, {batch_idx * (GROUP_SIZE / HEADS_PER_WG) + q_head_in_group / HEADS_PER_WG, 0, kv_head_idx, 0}, {0, j, 0, 0});
  __builtin_amdgcn_s_waitcnt(0);
  __builtin_amdgcn_s_barrier();

  // We first copy dV_j_T from accumulator GPRs to vector GPRs and then perform the store
  accvgpr_read(dV_j_T, dK_j_T);
  mul(dV_j_T, dV_j_T, dP_SCALE_FACTOR);
-  store<1>(g.dKg, dV_j, {batch_idx * GROUP_SIZE + q_head_in_group, 0, kv_head_idx, 0}, {0, j, 0, 0});
+  store<1>(g.dKg, dV_j, {batch_idx * (GROUP_SIZE / HEADS_PER_WG) + q_head_in_group / HEADS_PER_WG, 0, kv_head_idx, 0}, {0, j, 0, 0});

  // Write out final dQ_i slice
  mul(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
@@ -3422,6 +3426,3 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
 }

 template __global__ void attend_bwd_combined_ker<ATTN_D>(bf16*, bf16*, bf16*, bf16*, bf16*, bf16*, bf16*, float*, float*);
-
-
-
@@ -0,0 +1,69 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+
+#ifndef ATTN_B
+#define ATTN_B 2
+#endif
+#ifndef ATTN_N
+#define ATTN_N 8192
+#endif
+#ifndef ATTN_H
+#define ATTN_H 32
+#endif
+#ifndef ATTN_H_KV
+#define ATTN_H_KV 8
+#endif
+#ifndef ATTN_D
+#define ATTN_D 128
+#endif
+#ifndef THREADS_PER_BLOCK
+#define THREADS_PER_BLOCK 256
+#endif
+
+constexpr int GROUP_SIZE = ATTN_H / ATTN_H_KV;
+constexpr int HALF_D = ATTN_D / 2;
+constexpr int PACKED_D = (GROUP_SIZE + 2) * ATTN_D;
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_BLOCK) void
+fused_qkv_rope_forward(
+    __hip_bfloat16*       __restrict__ q,
+    __hip_bfloat16*       __restrict__ k,
+    __hip_bfloat16*       __restrict__ v,
+    const __hip_bfloat16* __restrict__ xqkv,
+    const __hip_bfloat16* __restrict__ freqs_cis) {
+  const int b = blockIdx.x;
+  const int n = blockIdx.y;
+  const int bn = b * ATTN_N + n;
+  const int packed_bn = bn * ATTN_H_KV * PACKED_D;
+  const int q_bn = bn * ATTN_H * ATTN_D;
+  const int kv_bn = bn * ATTN_H_KV * ATTN_D;
+
+  if (threadIdx.x < HALF_D) {
+    const int pair = threadIdx.x;
+    const int even = pair << 1;
+    const float c = static_cast<float>(freqs_cis[((n * HALF_D + pair) * 2) + 0]);
+    const float s = static_cast<float>(freqs_cis[((n * HALF_D + pair) * 2) + 1]);
+
+    for (int kvh = 0; kvh < ATTN_H_KV; kvh++) {
+      const int base = packed_bn + kvh * PACKED_D;
+
+      for (int rep = 0; rep < GROUP_SIZE; rep++) {
+        const int qbase = base + rep * ATTN_D;
+        const int h = kvh * GROUP_SIZE + rep;
+        const float a = static_cast<float>(xqkv[qbase + even]);
+        const float bb = static_cast<float>(xqkv[qbase + even + 1]);
+        const int out = q_bn + h * ATTN_D + even;
+        q[out] = static_cast<__hip_bfloat16>(a * c - bb * s);
+        q[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c);
+      }
+
+      const float a = static_cast<float>(xqkv[base + GROUP_SIZE * ATTN_D + even]);
+      const float bb = static_cast<float>(xqkv[base + GROUP_SIZE * ATTN_D + even + 1]);
+      const int out = kv_bn + kvh * ATTN_D + even;
+      k[out] = static_cast<__hip_bfloat16>(a * c - bb * s);
+      k[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c);
+      v[out] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even];
+      v[out + 1] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even + 1];
+    }
+  }
+}
@@ -0,0 +1,161 @@
+#include "kittens.cuh"
+
+using namespace kittens;
+
+#ifndef ATTN_B
+#define ATTN_B 2
+#endif
+#ifndef ATTN_N
+#define ATTN_N 8192
+#endif
+#ifndef ATTN_H
+#define ATTN_H 32
+#endif
+#ifndef ATTN_H_KV
+#define ATTN_H_KV 8
+#endif
+#ifndef ATTN_D
+#define ATTN_D 128
+#endif
+#ifndef THREADS_PER_BLOCK
+#define THREADS_PER_BLOCK 256
+#endif
+constexpr int GROUP_SIZE = ATTN_H / ATTN_H_KV;
+constexpr int HALF_D = ATTN_D / 2;
+constexpr int PACKED_H = ATTN_H_KV * (GROUP_SIZE + 2);
+constexpr int HEADS_PER_WG = ATTN_D == 128 && GROUP_SIZE % 2 == 0 ? 2 : 1;
+constexpr int KV_PARTIALS = GROUP_SIZE / HEADS_PER_WG;
+constexpr int NUM_WARPS = 4;
+constexpr int TILE_N = 16;
+
+template<typename T> using grad_tile = rt<T, TILE_N, ATTN_D, row_l, rt_16x32_s>;
+
+template<int axis, ducks::rt::row_layout RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<RT>>
+__device__ __forceinline__ void load_fa_shuffled(RT &dst, const GL &src, const COORD &idx) {
+  using U = typename GL::dtype;
+  using U2 = base_types::packing<U>::packed_type;
+  U *src_ptr = (U*)&src[(idx.template unit_coord<axis, 3>())];
+  const int row_stride = src.template stride<axis>();
+  const int lane = kittens::laneid();
+  const int tile_row_stride = row_stride * dst.base_tile_rows;
+  const int tile_stride = dst.base_tile_rows * dst.base_tile_cols;
+  const uint32_t buffer_size = src.batch() * src.depth() * src.rows() * src.cols() * sizeof(U);
+  const buffer_resource br = make_buffer_resource(reinterpret_cast<uintptr_t>(src_ptr), buffer_size, 0x00020000);
+
+  #pragma unroll
+  for (int i = 0; i < dst.height; i++) {
+    #pragma unroll
+    for (int j = 0; j < dst.width; j++) {
+      const float4 loaded = std::bit_cast<float4>(llvm_amdgcn_raw_buffer_load_b128(
+        std::bit_cast<i32x4>(br), (i * tile_row_stride + j * tile_stride + lane * 8) * sizeof(U), 0, 0));
+      const U2 *packed = reinterpret_cast<const U2*>(&loaded);
+      #pragma unroll
+      for (int k = 0; k < dst.packed_per_base_tile; k++) dst.tiles[i][j].data[k] = packed[k];
+    }
+  }
+}
+
+template<int axis, ducks::rt::row_layout RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<RT>>
+__device__ __forceinline__ void store_fa_shuffled(const GL &dst, const RT &src, const COORD &idx) {
+  using U = typename GL::dtype;
+  U *dst_ptr = (U*)&dst[(idx.template unit_coord<axis, 3>())];
+  const int row_stride = dst.template stride<axis>();
+  const int lane = kittens::laneid();
+  const int row_offset = (lane % 4) * 4;
+  const int col_offset = ((lane / 32) * 16) + (((lane % 32) / 16) * 2) + (((lane % 16) / 4) * 4);
+  const uint32_t buffer_size = dst.batch() * dst.depth() * dst.rows() * dst.cols() * sizeof(U);
+  const buffer_resource br = make_buffer_resource(reinterpret_cast<uintptr_t>(dst_ptr), buffer_size, 0x00020000);
+
+  #pragma unroll
+  for (int i = 0; i < src.height; i++) {
+    const int row = src.base_tile_rows * i + row_offset;
+    #pragma unroll
+    for (int j = 0; j < src.width; j++) {
+      const int col = src.base_tile_cols * j + col_offset;
+      #pragma unroll
+      for (int k = 0; k < src.packed_per_base_tile; k++) llvm_amdgcn_raw_buffer_store_b32(
+        *reinterpret_cast<const uint32_t*>(&src.tiles[i][j].data[k]), std::bit_cast<i32x4>(br),
+        ((row + k) * row_stride + col) * sizeof(U), 0, 0);
+    }
+  }
+}
+
+template<ducks::rt::row_layout RT>
+__device__ __forceinline__ void inverse_rope_fa(RT &tile, const bf16_2 *freqs, const int n_base) {
+  const int lane = kittens::laneid();
+  const int row_offset = (lane % 4) * 4;
+  const int col_offset = ((lane / 32) * 16) + (((lane % 32) / 16) * 2) + (((lane % 16) / 4) * 4);
+  #pragma unroll
+  for (int i = 0; i < tile.height; i++) {
+    #pragma unroll
+    for (int j = 0; j < tile.width; j++) {
+      const int col = tile.base_tile_cols * j + col_offset;
+      #pragma unroll
+      for (int k = 0; k < tile.packed_per_base_tile; k++) {
+        const int row = tile.base_tile_rows * i + row_offset + k;
+        const float2 cs = __bfloat1622float2(freqs[(n_base + row) * HALF_D + col / 2]);
+        const float2 g = __bfloat1622float2(tile.tiles[i][j].data[k]);
+        tile.tiles[i][j].data[k] = __float22bfloat162_rn(make_float2(g.x * cs.x + g.y * cs.y, -g.x * cs.y + g.y * cs.x));
+      }
+    }
+  }
+}
+
+template<ducks::rt::row_layout RT>
+__device__ __forceinline__ void inverse_rope(RT &tile, const bf16_2 *freqs, const int n_base) {
+  const int lane = kittens::laneid();
+  #pragma unroll
+  for (int i = 0; i < tile.height; i++) {
+    const int row = tile.base_tile_rows * i + lane % tile.base_tile_rows;
+    #pragma unroll
+    for (int j = 0; j < tile.width; j++) {
+      #pragma unroll
+      for (int k = 0; k < tile.packed_per_base_tile; k++) {
+        const int col = tile.base_tile_cols * j + tile.base_tile_stride * (lane / tile.base_tile_rows) + 2 * k;
+        const float2 cs = __bfloat1622float2(freqs[(n_base + row) * HALF_D + col / 2]);
+        const float2 g = __bfloat1622float2(tile.tiles[i][j].data[k]);
+        tile.tiles[i][j].data[k] = __float22bfloat162_rn(make_float2(g.x * cs.x + g.y * cs.y, -g.x * cs.y + g.y * cs.x));
+      }
+    }
+  }
+}
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_BLOCK) void
+fused_qkv_rope_backward(
+    bf16*       __restrict__ dxqkv,
+    const bf16* __restrict__ dq,
+    const bf16* __restrict__ dk,
+    const bf16* __restrict__ dv,
+    const bf16* __restrict__ freqs_cis) {
+  gl<bf16, -1, -1, -1, -1> out{dxqkv, ATTN_B, ATTN_N, PACKED_H, ATTN_D};
+  gl<bf16, -1, -1, -1, -1> dqg{const_cast<bf16*>(dq), ATTN_B, ATTN_H, ATTN_N, ATTN_D};
+  gl<bf16, -1, -1, -1, -1> dkg{const_cast<bf16*>(dk), ATTN_B * KV_PARTIALS, ATTN_N, ATTN_H_KV, ATTN_D};
+  gl<bf16, -1, -1, -1, -1> dvg{const_cast<bf16*>(dv), ATTN_B * KV_PARTIALS, ATTN_N, ATTN_H_KV, ATTN_D};
+  const int b = blockIdx.x, n_tile = blockIdx.y * NUM_WARPS + kittens::warpid(), n_base = n_tile * TILE_N;
+  const int field = blockIdx.z;
+
+  if (field < ATTN_H) {
+    grad_tile<bf16> tile;
+    load_fa_shuffled<2>(tile, dqg, {b, field, n_tile, 0});
+    inverse_rope_fa(tile, reinterpret_cast<const bf16_2*>(freqs_cis), n_base);
+    const int out_head = (field / GROUP_SIZE) * (GROUP_SIZE + 2) + field % GROUP_SIZE;
+    store_fa_shuffled<1>(out, tile, {b, n_tile, out_head, 0});
+  } else {
+    const bool is_k = field < ATTN_H + ATTN_H_KV;
+    const int kvh = field - ATTN_H - (is_k ? 0 : ATTN_H_KV);
+    const auto &src = is_k ? dkg : dvg;
+    grad_tile<bf16> partial, tile;
+    grad_tile<float> partial_f, sum;
+    zero(sum);
+    #pragma unroll
+    for (int p = 0; p < KV_PARTIALS; p++) {
+      load<1>(partial, src, {b * KV_PARTIALS + p, n_tile, kvh, 0});
+      copy(partial_f, partial);
+      add(sum, sum, partial_f);
+    }
+    copy(tile, sum);
+    if (is_k) inverse_rope(tile, reinterpret_cast<const bf16_2*>(freqs_cis), n_base);
+    const int out_head = kvh * (GROUP_SIZE + 2) + GROUP_SIZE + !is_k;
+    store<1>(out, tile, {b, n_tile, out_head, 0});
+  }
+}
@@ -148,12 +148,28 @@ __global__ __launch_bounds__(512, 2) void hk_fp8_gemm(bf16 *C_ptr, fp8e4m3 *A_pt
    RT_C cC;
    RT_C cD;

-    // Calculate which block this threadblock should work on
-    int global_block_id = blockIdx.x;
-
-    // Convert linear block ID to 2D coordinates
-    int block_row = global_block_id / blocks_per_col;
-    int block_col = global_block_id % blocks_per_col;
+    int block_row, block_col;
+    if constexpr (N > M) {
+        // Wide outputs repeatedly consume the same A rows. Keep a short strip
+        // resident on each XCD while walking N to improve local cache reuse.
+        int wgid = chiplet_transform_chunked(int(blockIdx.x), total_blocks_needed, NUM_XCDS, 64);
+        constexpr int WGM = 3;
+        const int num_wgid_in_group = WGM * blocks_per_col;
+        const int group_id = wgid / num_wgid_in_group;
+        const int first_block_row = group_id * WGM;
+        const int group_size_m = min(blocks_per_row - first_block_row, WGM);
+        block_row = first_block_row + ((wgid % num_wgid_in_group) % group_size_m);
+        block_col = (wgid % num_wgid_in_group) / group_size_m;
+    } else {
+        int wgid = chiplet_transform_chunked(int(blockIdx.x), total_blocks_needed, NUM_XCDS, 64);
+        constexpr int WGM = 8;
+        const int num_wgid_in_group = WGM * blocks_per_col;
+        const int group_id = wgid / num_wgid_in_group;
+        const int first_block_row = group_id * WGM;
+        const int group_size_m = min(blocks_per_row - first_block_row, WGM);
+        block_row = first_block_row + ((wgid % num_wgid_in_group) % group_size_m);
+        block_col = (wgid % num_wgid_in_group) / group_size_m;
+    }
    int block_m = block_row * BLOCK_SIZE_ROW;
    int block_n = block_col * BLOCK_SIZE_COL;

@@ -134,7 +134,7 @@ __global__ __launch_bounds__(512, 2) void hk_fp8_atb_gemm(bf16 *C_ptr, fp8e4m3 *

    int wgid = blockIdx.x;
    const int WGM = 8;
-    wgid = chiplet_transform_chunked(wgid, total_blocks_needed, NUM_XCDS, 64);
+    wgid = chiplet_transform_chunked(wgid, total_blocks_needed, NUM_XCDS, 32);

    const int num_wgid_in_group = WGM * blocks_per_col;
    int group_id = wgid / num_wgid_in_group;
@@ -1,7 +1,6 @@
 import math
 from typing import cast, Callable
-from tinygrad import dtypes
-from tinygrad.uop.ops import AxisType, UOp, Ops
+from tinygrad.uop.ops import AxisType, UOp
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import prod

@@ -75,9 +74,9 @@ class Group:

    a_base_shape = cast(RT, a).base_shape
    if a_base_shape.cols == 16:
-      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_dims = (16, 16, 16)
    elif a_base_shape.cols == 32:
-      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_dims = (16, 16, 32)
    else: raise NotImplementedError(f"mma_AB not implemented for {a_base_shape.cols=}")

    for height in self.ker.range(c.shape[-3], track=False):
@@ -92,7 +91,7 @@ class Group:
          else: raise NotImplementedError(f"mma_AB not implemented for {a_base_shape.cols=}")
          d_in = UOp.stack(*[c[height, width, i] for i in range(4)])

-          out = UOp(Ops.WMMA, dtypes.float32, (a_in, b_in, d_in), arg=wmma_arg)
+          out = UOp.wmma(a_in, b_in, d_in, wmma_dims, 'AMD', 64)
          c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
          c_store = UOp.group(*c_i).end(height, width, inner)

@@ -105,9 +104,9 @@ class Group:

    a_base_shape = cast(RT, a).base_shape
    if a_base_shape.cols == 16:
-      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_dims = (16, 16, 16)
    elif a_base_shape.cols == 32:
-      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_dims = (16, 16, 32)
    else: raise NotImplementedError(f"mma_ABt not implemented for {a_base_shape.cols=}")

    for height in self.ker.range(c.shape[-3], track=False):
@@ -122,7 +121,7 @@ class Group:
          else: raise NotImplementedError(f"mma_ABt not implemented for {a_base_shape.cols=}")
          d_in = UOp.stack(*[c[height, width, i] for i in range(4)])

-          out = UOp(Ops.WMMA, dtypes.float32, (a_in, b_in, d_in), arg=wmma_arg)
+          out = UOp.wmma(a_in, b_in, d_in, wmma_dims, 'AMD', 64)
          c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
          c_store = UOp.group(*c_i).end(height, width, inner)

@@ -135,9 +134,9 @@ class Group:

    a_base_shape = cast(RT, a).base_shape
    if a_base_shape.cols == 16:
-      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_dims = (16, 16, 16)
    elif a_base_shape.cols == 32:
-      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_dims = (16, 16, 32)
    else: raise NotImplementedError(f"mma_AtB not implemented for {a_base_shape.cols=}")

    for height in self.ker.range(c.shape[-3], track=False):
@@ -152,7 +151,7 @@ class Group:
          else: raise NotImplementedError(f"mma_AtB not implemented for {a_base_shape.cols=}")
          d_in = UOp.stack(*[c[height, width, i] for i in range(4)])

-          out = UOp(Ops.WMMA, dtypes.float32, (a_in, b_in, d_in), arg=wmma_arg)
+          out = UOp.wmma(a_in, b_in, d_in, wmma_dims, 'AMD', 64)
          c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
          c_store = UOp.group(*c_i).end(height, width, inner)

@@ -165,9 +164,9 @@ class Group:

    a_base_shape = cast(RT, a).base_shape
    if a_base_shape.cols == 16:
-      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_dims = (16, 16, 16)
    elif a_base_shape.cols == 32:
-      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_dims = (16, 16, 32)
    else: raise NotImplementedError(f"mma_AtBt not implemented for {a_base_shape.cols=}")

    for height in self.ker.range(c.shape[-3], track=False):
@@ -182,7 +181,7 @@ class Group:
          else: raise NotImplementedError(f"mma_AtBt not implemented for {a_base_shape.cols=}")
          d_in = UOp.stack(*[c[height, width, i] for i in range(4)])

-          out = UOp(Ops.WMMA, dtypes.float32, (a_in, b_in, d_in), arg=wmma_arg)
+          out = UOp.wmma(a_in, b_in, d_in, wmma_dims, 'AMD', 64)
          c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
          c_store = UOp.group(*c_i).end(height, width, inner)

@@ -565,8 +565,10 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
  "aten.floor_divide": lambda x,y: x//y,
  "aten.floor_divide_.Tensor": lambda x,y: x//y,
  "aten.__lshift__.Scalar": lambda x,y: x<<y,
+  "aten.__lshift__.Tensor": lambda x,y: x<<y,
  "aten.__ilshift__.Scalar": lambda x,y: x<<y,
  "aten.__rshift__.Scalar": lambda x,y: x>>y,
+  "aten.__rshift__.Tensor": lambda x,y: x>>y,
  "aten.__irshift__.Scalar": lambda x,y: x>>y,
  # inplace ops using replace for fusion
  "aten.zero_": lambda x: x.const_like(0),
@@ -161,7 +161,6 @@ norecursedirs = [
  ".git",
 ]
 timeout = 300
-timeout_method = "thread"
 timeout_func_only = true
 testpaths = ["test"]
 filterwarnings = [
@@ -8,7 +8,7 @@ from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
 from tinygrad import Context, Device, Tensor, dtypes
 from hypothesis import given, settings, strategies as strat
-from test.helpers import rand_for_dtype
+from test.helpers import rand_for_dtype, min_normal
 from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX, FP8E4M3FNUZ_MAX, FP8E5M2FNUZ_MAX
 import pytest
 pytestmark = pytest.mark.filterwarnings("ignore")
@@ -25,10 +25,10 @@ def get_available_cast_dtypes(dtype: DType) -> List[DType]:
  if dtype not in supported_dtypes and dtype not in dtypes.fp8s+(dtypes.half,dtypes.bfloat16): return []
  return dts

-def _to_torch_storage_type(dtype:DType):
-  if dtype == dtypes.bfloat16: return torch.float32
-  if dtype in dtypes.fp8s: return torch.float32
-  return _to_torch_dtype(dtype)
+def _to_torch_storage(a:Tensor) -> torch.Tensor:
+  # tolist() of an fp8 Tensor gives floats, so convert and store in uint8
+  if a.dtype in dtypes.fp8s: return torch.tensor([float_to_fp8(x, a.dtype) for x in a.flatten().tolist()], dtype=torch.uint8).reshape(a.shape)
+  return torch.tensor(a.tolist(), dtype=_to_torch_dtype(a.dtype))

 def _test_to_np(a:Tensor, np_dtype, target):
  if DEBUG >= 2: print(a)
@@ -46,12 +46,15 @@ def _test_cast(a:Tensor, target_dtype:DType):
  if a.is_floating_point() and dtypes.is_unsigned(target_dtype):
    # converting negative float to unsigned integer is undefined
    a = a.abs()
+  if a.is_floating_point() and dtypes.is_float(target_dtype) and (mn:=min_normal(target_dtype)) >= min_normal(a.dtype):
+    # subnormals are zero, so an input below the target's min normal casts to 0
+    a = (a.abs() < mn).where(0, a)

  expected = list(a.numpy().astype(_to_np_dtype(target_dtype)))
  if target_dtype in dtypes.fp8s: expected = [truncate[target_dtype](x) for x in expected]
  _test_op(lambda: a.cast(target_dtype), target_dtype, expected)
 def _test_bitcast(a:Tensor, target_dtype:DType, target=None):
-  expected = torch.tensor(a.tolist(), dtype=_to_torch_storage_type(a.dtype)).view(_to_torch_dtype(target_dtype)).tolist()
+  expected = _to_torch_storage(a).view(_to_torch_dtype(target_dtype)).tolist()
  if target_dtype in dtypes.fp8s: expected = [fp8_to_float(x, target_dtype) for x in expected]
  _test_op(lambda: a.bitcast(target_dtype), target_dtype, target or expected)

@@ -61,10 +64,12 @@ class TestDType(unittest.TestCase):
  @classmethod
  def setUpClass(cls):
    if cls.DTYPE is None: raise unittest.SkipTest("base class")
-    cls.DATA = rand_for_dtype(cls.DTYPE, 0x10, allow_subnormal=cls.DTYPE in supported_dtypes)
+    cls.DATA = rand_for_dtype(cls.DTYPE, 0x10, allow_subnormal=cls.DTYPE in supported_dtypes and cls.DTYPE not in dtypes.fp8s)

  def test_to_np(self):
-    _test_to_np(Tensor(self.DATA, dtype=self.DTYPE), _to_np_dtype(self.DTYPE), np.array(self.DATA, dtype=_to_np_dtype(self.DTYPE)))
+    a = Tensor(self.DATA, dtype=self.DTYPE)
+    self.assertEqual(a.dtype, self.DTYPE)
+    _test_to_np(a, _to_np_dtype(self.DTYPE), np.array(self.DATA, dtype=_to_np_dtype(self.DTYPE)))

  def test_casts_to(self):
    for dtype in get_available_cast_dtypes(self.DTYPE):
@@ -273,10 +278,11 @@ class TestBitCast(unittest.TestCase):
  @given(strat.sampled_from(dtype_ints + dtype_floats), strat.sampled_from(dtype_ints + dtype_floats))
  def test_shape_change_bitcast(self, dt1, dt2):
    data = rand_for_dtype(dt1, 32).reshape(2, 2, 8)
-    expected = torch.tensor(data.tolist(), dtype=_to_torch_storage_type(dt1)).view(_to_torch_dtype(dt2))
+    a = Tensor(data, dtype=dt1)
+    expected = _to_torch_storage(a).view(_to_torch_dtype(dt2))
    if dt2 in dtypes.fp8s:
      expected = torch.tensor([fp8_to_float(x, dt2) for x in expected.view(-1).tolist()]).view_as(expected)
-    _test_op(lambda: Tensor(data, dtype=dt1).bitcast(dt2), dt2, expected.tolist())
+    _test_op(lambda: a.bitcast(dt2), dt2, expected.tolist())

  def test_shape_change_bitcast_exceptions(self):
    with self.assertRaises(RuntimeError):
@@ -293,6 +299,11 @@ class TestBitCast(unittest.TestCase):
    b = a.bitcast(dtypes.float32)
    assert b.numpy()[0,0] == 1.

+  def test_bitcast_bf16_from_cast(self):
+    # a bfloat16 from a cast holds bfloat16 bits. 1.0 is 0x3f80 in bfloat16, which is 1.875 in half
+    a = Tensor([1.0], dtype=dtypes.float32).cast(dtypes.bfloat16)
+    assert a.bitcast(dtypes.half).numpy()[0] == 1.875
+
 class TestInt16DType(TestDType): DTYPE = dtypes.int16

 class TestUint16DType(TestDType):
@@ -291,6 +291,33 @@ class TestDTypeALU(unittest.TestCase):
  @Context(EMULATED_DTYPES="long")
  def test_emulated_int64(self, a, b, op): universal_test(a, b, dtypes.int64, op)

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

@@ -1,9 +1,9 @@
 import numpy as np
-import functools, unittest, ctypes
+import functools, unittest

 from tinygrad.device import Device, Buffer
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import Context, from_mv
+from tinygrad.helpers import Context
 from tinygrad.dtype import dtypes
 from tinygrad.engine.jit import MultiGraphRunner
 from tinygrad.engine.realize import run_linear, compile_linear
@@ -31,7 +31,7 @@ def make_buffer(device, size=BUF_SIZE, fill=False):
  buf = Buffer(device, size, dtypes.int).ensure_allocated()
  if fill:
    with Context(DEBUG=0):
-      buf.copyin(Tensor(np.random.randint(-10000, 10000, size=size, dtype=np.int32)).realize().uop.base.realized.as_memoryview())
+      buf.copy_from(Tensor(np.random.randint(-10000, 10000, size=size, dtype=np.int32)).realize().uop.base.realized)
  return buf

 def make_view(base, offset_elems, size_elems):
@@ -55,15 +55,12 @@ def run_schedule(calls:list[UOp]):
  run_linear(UOp(Ops.LINEAR, src=tuple(calls)))

 def zero_bufs(bufs):
-  for b in bufs:
-    mv = memoryview(bytearray(b.nbytes))
-    ctypes.memset(from_mv(mv), 0, len(mv))
-    b.copyin(mv)
+  for b in bufs: b.copy_from(Buffer("PYTHON", b.size, b.dtype, opaque=memoryview(bytearray(b.nbytes))))

@unittest.skipUnless(Device[Device.DEFAULT].graph is not None, "graph support required")
 class TestGraph(unittest.TestCase):
  def skip_if_no_offset(self):
-    if not hasattr(Device[Device.DEFAULT].allocator, "_offset"): self.skipTest("device does not support _offset")
+    if Device.DEFAULT in {"WEBGPU", "CL"}: self.skipTest("device does not support _offset")

  def skip_if_not_multigraph(self):
    graph = g.func if isinstance(g:=(d:=Device[Device.DEFAULT]).graph, functools.partial) else g
@@ -213,8 +210,8 @@ class TestGraph(unittest.TestCase):

  def test_graph_offset_bufs(self):
    self.skip_if_not_multigraph()
+    self.skip_if_no_offset()
    d0 = Device.DEFAULT
-    if not hasattr(Device[d0].allocator, "_offset"): self.skipTest("device does not support _offset")

    b0 = make_buffer(d0, fill=True)
    b1 = make_view(b0, 0, b0.size)
@@ -424,7 +424,7 @@ def copyout_outputs(outbufs:list[Buffer]) -> list[np.ndarray]:
  return [np.frombuffer(x.as_memoryview(), _to_np_dtype(x.dtype)) for x in outbufs]

 def reset_bufs(bufs:list[Buffer]):
-  for buf in bufs: buf.copyin(np.zeros((buf.size*buf.dtype.itemsize,), dtype=np.uint8).data)
+  for buf in bufs: buf.copy_from(Buffer("PYTHON", buf.size, buf.dtype, opaque=memoryview(bytearray(buf.nbytes))))

 def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[],
                               apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[]):
@@ -1,11 +1,14 @@
-import unittest
+import unittest, functools
 from tinygrad import Tensor, Device, dtypes, Context, GlobalCounters
 from tinygrad.helpers import getenv
 from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8
 from extra.llama_kernels.fused_ce import fused_ce_loss
 from extra.llama_kernels import local_abs_max
 from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed, quantize_fp8_scalar
+from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
+from extra.thunder.amd.fa import custom_fused_qkv_rope_backward, fused_qkv_rope
 from test.helpers import needs_second_gpu
+from test.backend.test_asm_gemm import has_hipcc

 def run_fused_ce(bs:int, seqlen:int, vocab:int, label_smoothing:float=0.0) -> None:
  Tensor.manual_seed(0)
@@ -46,9 +49,10 @@ def run_quantize_fp8(shape:tuple[int, ...], delayed:bool=True) -> None:
  with Context(DEBUG=0): Tensor.realize(x, amax_state)

  if delayed:
-    fp8, inv_scale, new_amax, _ = quantize_fp8_delayed(x, amax_state, FP8_DTYPE)
+    amax_out = Tensor.zeros((), dtype=dtypes.float32, device=x.device).realize()
+    fp8, inv_scale = quantize_fp8_delayed(x, amax_state, amax_out, FP8_DTYPE)
    ref_fp8, ref_inv_scale, ref_new_amax = quantize_fp8(x, amax_state=amax_state)
-    Tensor.realize(fp8, inv_scale, new_amax)
+    Tensor.realize(fp8, inv_scale)
    Tensor.realize(ref_fp8, ref_inv_scale, ref_new_amax)
  else:
    fp8 = quantize_fp8_scalar(x, amax_state, FP8_DTYPE)
@@ -60,9 +64,10 @@ def run_quantize_fp8(shape:tuple[int, ...], delayed:bool=True) -> None:
    assert fp8.cast(dtypes.float).allclose(ref_fp8.cast(dtypes.float), atol=0, rtol=0).item(), "fp8 mismatch"
    if delayed:
      assert inv_scale.allclose(ref_inv_scale, atol=0, rtol=0).item(), "inv_scale mismatch"
-      assert new_amax.allclose(ref_new_amax, atol=0, rtol=0).item(), \
-        f"amax mismatch: got={new_amax.item()} ref={ref_new_amax.item()} diff={abs(new_amax.item()-ref_new_amax.item())}"
+      assert amax_out.allclose(ref_new_amax, atol=0, rtol=0).item(), \
+        f"amax mismatch: got={amax_out.item()} ref={ref_new_amax.item()} diff={abs(amax_out.item()-ref_new_amax.item())}"

+@unittest.skipUnless(Device.DEFAULT == "AMD", "requires atomic max")
 class TestQuantizeFP8(unittest.TestCase):
  def setUp(self):
    ren = Device[Device.DEFAULT].renderer
@@ -78,10 +83,11 @@ class TestQuantizeFP8(unittest.TestCase):
    x = Tensor.empty(2048*8, 1024, dtype=dtypes.bfloat16, device=devs).uop.multi(0)
    x = Tensor(x, device=devs)
    amax_state = Tensor.full((), 2.0, dtype=dtypes.float32, device=devs).contiguous()
-    fp8, _, new_amax, _ = quantize_fp8_delayed(x, amax_state, FP8_DTYPE)
-    Tensor.realize(fp8, new_amax)
+    amax_out = Tensor.zeros((), dtype=dtypes.float32, device=devs).realize()
+    fp8, _ = quantize_fp8_delayed(x, amax_state, amax_out, FP8_DTYPE)
+    Tensor.realize(fp8)
    assert fp8.uop.shape == x.uop.shape
-    assert new_amax.shape == ()
+    assert amax_out.shape == ()

 class TestLocalAmax(unittest.TestCase):
  def test_multi_tensor_local_shard_amax(self):
@@ -92,5 +98,68 @@ class TestLocalAmax(unittest.TestCase):
    self.assertEqual(GlobalCounters.kernel_count, 2)
    self.assertEqual(out.tolist(), [[0., 7., 14., 21.], [28., 35., 42., 49.], [120., 135., 150., 165.], [180., 195., 210., 225.]])

+@unittest.skipUnless(has_hipcc() and Device.DEFAULT == "AMD", "requires hipcc to compile and amd device to run")
+class TestFusedQKVRoPE(unittest.TestCase):
+  SHAPE = (2, 8192, 32, 8, 128)
+
+  def rand_bf16(self, *shape:int) -> Tensor:
+    return (Tensor.randn(*shape) * 0.1).cast(dtypes.bfloat16).contiguous().realize()
+
+  def freqs_cis(self) -> Tensor:
+    _, N, _, _, D = self.SHAPE
+    return precompute_freqs_cis(D, N * 2).cast(dtypes.bfloat16).clone().realize()
+
+  def test_llama31_8b_forward(self):
+    Tensor.manual_seed(0)
+    B, N, H, H_KV, D = self.SHAPE
+    GROUP = H // H_KV
+    freqs_cis = self.freqs_cis()
+
+    x = self.rand_bf16(B, N, H_KV * (GROUP + 2) * D)
+    q, k, v = fused_qkv_rope(x, freqs_cis, H, H_KV, D)
+    Tensor.realize(q, k, v)
+    packed_ref = x.reshape(B, N, H_KV, GROUP + 2, D)
+    q_ref = packed_ref[:, :, :, :GROUP].reshape(B, N, H, D)
+    k_ref, v_ref = packed_ref[:, :, :, GROUP], packed_ref[:, :, :, GROUP+1]
+    q_ref, k_ref = apply_rotary_emb(q_ref, k_ref, freqs_cis[:, :N])
+    q_ref, k_ref, v_ref = q_ref.cast(dtypes.bfloat16), k_ref.cast(dtypes.bfloat16), v_ref.cast(dtypes.bfloat16)
+    Tensor.realize(q_ref, k_ref, v_ref)
+
+    with Context(DEBUG=0):
+      self.assertTrue(q.allclose(q_ref, atol=2e-2, rtol=0).item(), "Q forward mismatch")
+      self.assertTrue(k.allclose(k_ref, atol=2e-2, rtol=0).item(), "K forward mismatch")
+      self.assertTrue(v.allclose(v_ref, atol=0, rtol=0).item(), "V forward mismatch")
+
+  def test_llama31_8b_backward(self):
+    Tensor.manual_seed(1)
+    B, N, H, H_KV, D = self.SHAPE
+    PARTIALS = 2
+    GROUP = H // H_KV
+    freqs_cis = self.freqs_cis()
+    dq = self.rand_bf16(B, N, H, D)
+    dk_partial = self.rand_bf16(B * PARTIALS, N, H_KV, D)
+    dv_partial = self.rand_bf16(B * PARTIALS, N, H_KV, D)
+
+    # Invert Flash Attention's dQ layout transform to reproduce its native buffer.
+    dq_native = dq.transpose(1, 2).reshape(B, H, N//16, 4, 4, 4, 2, D//32, 2, 2) \
+      .permute(0, 1, 2, 5, 6, 8, 7, 3, 4, 9).reshape(B, H, N, D).contiguous().realize()
+    dx = Tensor.empty(B, N, H_KV * (GROUP + 2) * D, dtype=dtypes.bfloat16)
+    arch = Device[Device.DEFAULT].renderer.target.arch
+    fxn = functools.partial(custom_fused_qkv_rope_backward, device=Device.DEFAULT, arch=arch,
+                            B=B, N=N, H=H, H_KV=H_KV, D=D)
+    dx = Tensor.custom_kernel(dx, dq_native, dk_partial, dv_partial, freqs_cis, fxn=fxn)[0].realize()
+
+    def inverse_rope(x:Tensor) -> Tensor:
+      x = x.reshape(*x.shape[:-1], D//2, 2).float()
+      cs = freqs_cis[:, :N].float()
+      return Tensor.stack(x[..., 0] * cs[..., 0] + x[..., 1] * cs[..., 1],
+                          -x[..., 0] * cs[..., 1] + x[..., 1] * cs[..., 0], dim=-1).flatten(-2).cast(dtypes.bfloat16)
+
+    dq_ref = inverse_rope(dq).reshape(B, N, H_KV, GROUP, D)
+    dk_ref = inverse_rope(dk_partial.float().reshape(B, PARTIALS, N, H_KV, D).sum(1).cast(dtypes.bfloat16)).unsqueeze(3)
+    dv_ref = dv_partial.float().reshape(B, PARTIALS, N, H_KV, D).sum(1).cast(dtypes.bfloat16).unsqueeze(3)
+    ref = Tensor.cat(dq_ref, dk_ref, dv_ref, dim=3).reshape(*dx.shape).realize()
+    with Context(DEBUG=0): self.assertTrue(dx.allclose(ref, atol=2e-2, rtol=2e-2).item(), "backward mismatch")
+
 if __name__ == '__main__':
  unittest.main()
@@ -385,7 +385,7 @@ class TestMultiBufferView(unittest.TestCase):
    b_ref = view_fn(a_ref)
    b_multi = view_fn(a_multi).contiguous()
    linear, var_vals = b_multi.linear_with_vars()
-    if all(hasattr(Device[d].allocator, "_offset") for d in b_multi.device):
+    if all(not d.startswith(("WEBGPU", "CL")) for d in b_multi.device):
      compiled = [call for call in linear.src if call.src[0].op is Ops.SINK]
      self.assertEqual(len(compiled), 0, f"expected zero compiled kernels, got {len(compiled)}")
    run_linear(linear, var_vals)
@@ -417,7 +417,7 @@ class TestMultiBufferView(unittest.TestCase):
    a = Tensor.arange(8*12).reshape(8, 12).clone().shard(devices_4, axis=1).realize()
    out = a[5].contiguous()
    linear, var_vals = out.linear_with_vars()
-    if all(hasattr(Device[d].allocator, "_offset") for d in out.device):
+    if all(not d.startswith(("WEBGPU", "CL")) for d in out.device):
      compiled = [call for call in linear.src if call.src[0].op is Ops.SINK]
      self.assertEqual(len(compiled), 0)
    run_linear(linear, var_vals)
@@ -848,6 +848,8 @@ class TestOps(unittest.TestCase):
    helper_test_op([], lambda: tor << 0, lambda: (ten << 0).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor << 2, lambda: (ten << 2).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor << 31, lambda: (ten << 31).cast(dtypes.int32), forward_only=True)
+    helper_test_op([], lambda: tor << torch.tensor([0,2,4]).int(),
+                   lambda: (ten << Tensor([0,2,4], dtype=dtypes.uint32)).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor.__lshift__(2), lambda: ten.__lshift__(2).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor.bitwise_left_shift(2), lambda: ten.lshift(2).cast(dtypes.int32), forward_only=True)

@@ -859,6 +861,8 @@ class TestOps(unittest.TestCase):
    helper_test_op([], lambda: tor >> 0, lambda: (ten >> 0).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor >> 2, lambda: (ten >> 2).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor >> 31, lambda: (ten >> 31).cast(dtypes.int32), forward_only=True)
+    helper_test_op([], lambda: tor >> torch.tensor([0,2,4]).int(),
+                   lambda: (ten >> Tensor([0,2,4], dtype=dtypes.uint32)).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor.__rshift__(2), lambda: ten.__rshift__(2).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor.bitwise_right_shift(2), lambda: ten.rshift(2).cast(dtypes.int32), forward_only=True)

@@ -870,6 +874,7 @@ class TestOps(unittest.TestCase):
    helper_test_op([], lambda: tor << 2, lambda: ten << 2, forward_only=True)
    helper_test_op([], lambda: tor << 8, lambda: ten << 8, forward_only=True)
    helper_test_op([], lambda: tor << 31, lambda: ten << 31, forward_only=True)
+    helper_test_op([], lambda: tor << torch.tensor([0,2,8,31]).int(), lambda: ten << Tensor([0,2,8,31], dtype=dtypes.int), forward_only=True)

  def test_rshift_signed(self):
    data = [[-1, -3, 1, 7], [0, -2147483648, 2147483647, -1]]
@@ -879,6 +884,7 @@ class TestOps(unittest.TestCase):
    helper_test_op([], lambda: tor >> 2, lambda: ten >> 2, forward_only=True)
    helper_test_op([], lambda: tor >> 8, lambda: ten >> 8, forward_only=True)
    helper_test_op([], lambda: tor >> 31, lambda: ten >> 31, forward_only=True)
+    helper_test_op([], lambda: tor >> torch.tensor([0,2,8,31]).int(), lambda: ten >> Tensor([0,2,8,31], dtype=dtypes.int), forward_only=True)

  def test_idiv_shift_rewrite_negative(self):
    a = Tensor(-5).div(2, rounding_mode="trunc").item()
@@ -70,9 +70,9 @@ class TestProfiler(unittest.TestCase):
    buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()

    with helper_collect_profile(TestProfiler.d0) as profile:
-      buf1.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
+      buf1.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 1)))))

-    kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith(TestProfiler.d0.device)]
+    kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith((TestProfiler.d0.device, "PYTHON"))]
    assert len(kernel_runs) == 1, "one kernel run is expected"

  def test_profile_multiops(self):
@@ -80,12 +80,12 @@ class TestProfiler(unittest.TestCase):
    buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()

    with helper_collect_profile(TestProfiler.d0) as profile:
-      buf1.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
+      buf1.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 1)))))
      gs, ls = TestProfiler.prg.arg.launch_dims({})
      TestProfiler.runtime(buf1._buf, TestProfiler.a.uop.buffer._buf, global_size=gs, local_size=ls)
-      buf1.copyout(memoryview(bytearray(buf1.nbytes)))
+      buf1.as_memoryview()

-    evs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith(TestProfiler.d0.device)]
+    evs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith((TestProfiler.d0.device, "PYTHON"))]

    assert len(evs) == 3, "3 kernel runs are expected"
    # NOTE: order of events does not matter, the tool is responsible for sorting them
@@ -103,12 +103,12 @@ class TestProfiler(unittest.TestCase):
    buf2 = Buffer(f"{Device.DEFAULT}:1", 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()

    with helper_collect_profile(TestProfiler.d0, d1) as profile:
-      buf1.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
-      buf2.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
+      buf1.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 1)))))
+      buf2.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 1)))))

    for dev in [TestProfiler.d0.device, d1.device]:
      evs = [x for x in profile if isinstance(x, ProfileRangeEvent) and _dev_base(x.device) == dev]
-      assert len(evs) == 1, "one kernel runs are expected"
+      assert len(evs) == (0 if hasattr(TestProfiler.d0.allocator, '_as_buffer') else 1), "one kernel runs are expected"

  def test_profile_multidev_transfer(self):
    try: d1 = Device[f"{Device.DEFAULT}:1"]
@@ -1,6 +1,6 @@
 import unittest
 import numpy as np
-from tinygrad.device import Device
+from tinygrad.device import Device, Buffer
 from tinygrad.dtype import dtypes, ConstType
 from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
@@ -10,13 +10,13 @@ from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.wgsl import WGSLRenderer
 from tinygrad.runtime.ops_python import PythonRenderer
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, python_alu
-from tinygrad.tensor import Tensor, _to_np_dtype
+from tinygrad.tensor import Tensor

 def _test_uop_result(inputs:list[Tensor], sink:UOp, local_size=None):
  for x in inputs: x.realize()
  sz = 1 if local_size is None else prod(local_size)
  outs = [UOp.new_buffer(Device.DEFAULT, sz, u.src[1].dtype) for u in sink.src if u.op is Ops.STORE]
-  for u in outs: u.buffer.allocate().copyin(np.zeros(sz, dtype=_to_np_dtype(u.dtype)).data)
+  for u in outs: u.buffer.allocate().copy_from(Buffer("PYTHON", sz, u.dtype, opaque=memoryview(bytearray(u.buffer.nbytes))))
  run_linear(UOp(Ops.LINEAR, src=(sink.call(*outs, *(x.uop.base for x in inputs)),)))
  return [u.buffer.numpy() for u in outs]

@@ -2,7 +2,7 @@
 # schedule confirms the right things are capable of fusing
 # NOTE: this has overlap with external_test_opt.py

-import unittest
+import unittest, time
 import numpy as np

 from tinygrad import nn, dtypes, Device, Tensor, Variable
@@ -197,6 +197,20 @@ class TestLimitBufs(unittest.TestCase):
        base = (idx >= i).where(a + b, base)
      assert all(x > 0 for x in base.tolist())

+  def test_limit_bufs_linear_scaling(self):
+    def sched_time(n):
+      with Context(TRACK_MATCH_STATS=0, DEBUG=0):
+        bufs = [Tensor.ones(16).contiguous().realize() for _ in range(4)]
+        root = bufs[0]
+        for i in range(n): root = root + bufs[i % 4]
+        with Context(MAX_KERNEL_BUFFERS=8, SCACHE=0):
+          st = time.perf_counter()
+          root.schedule_linear()
+          return time.perf_counter() - st
+    sched_time(400)
+    t1, t2 = min(sched_time(400) for _ in range(3)), min(sched_time(1600) for _ in range(3))
+    self.assertLess(t2/t1, 8, f"{t1*1e3:.1f}ms -> {t2*1e3:.1f}ms")
+
 class TestSwizzle(unittest.TestCase):
  def test_swizzle_simple(self):
    Tensor.manual_seed(0)
@@ -4,11 +4,10 @@ from tinygrad.device import Buffer
 from tinygrad.helpers import Context, DEV
 from test.helpers import needs_second_gpu

-@unittest.skipUnless(hasattr(Device[Device.DEFAULT].allocator, "_offset"), "subbuffer not supported")
+@unittest.skipIf(Device.DEFAULT in {"WEBGPU", "CL"}, "subbuffer not supported")
 class TestSubBuffer(unittest.TestCase):
  def setUp(self):
-    self.buf = Buffer(Device.DEFAULT, 10, dtypes.uint8).ensure_allocated()
-    self.buf.copyin(memoryview(bytearray(range(10))))
+    self.buf = Buffer(Device.DEFAULT, 10, dtypes.uint8, initial_value=bytes(range(10)))
    self.buf_unalloc = Buffer(Device.DEFAULT, 10, dtypes.uint8)

  def test_subbuffer(self):
@@ -59,7 +58,7 @@ class TestSubBuffer(unittest.TestCase):
      _ = Buffer(Device.DEFAULT, 10, dtypes.uint8).ensure_allocated()

      self.buf.ensure_allocated()
-      self.buf.copyin(memoryview(bytearray(range(10, 20))))
+      self.buf.copy_from(Buffer("PYTHON", 10, dtypes.uint8, opaque=memoryview(bytearray(range(10, 20)))))

      vbuf.ensure_allocated()

@@ -109,15 +108,15 @@ class TestSubBuffer(unittest.TestCase):
  def test_subbuffer_copy_in_out(self):
    sub_buf = self.buf.view(3, dtypes.uint8, offset=3).ensure_allocated() # [3:6]
    data_out_sub = bytearray([0]*3)
-    sub_buf.copyout(memoryview(data_out_sub))
+    data_out_sub[:] = sub_buf.as_memoryview()
    assert data_out_sub == bytearray(range(3, 6))
-    sub_buf.copyin(memoryview(bytearray(range(3))))
+    sub_buf.copy_from(Buffer("PYTHON", 3, dtypes.uint8, opaque=memoryview(bytearray(range(3)))))
    assert sub_buf.as_memoryview().tolist() == list(range(3))
    assert self.buf.as_memoryview().tolist()[3:6] == list(range(3))
-    sub_buf.copyout(memoryview(data_out_sub))
+    data_out_sub[:] = sub_buf.as_memoryview()
    assert data_out_sub == bytearray(range(3))
    data_out_base = bytearray([0]*10)
-    self.buf.copyout(memoryview(data_out_base))
+    data_out_base[:] = self.buf.as_memoryview()
    assert data_out_base[0:3] == bytearray(range(0, 3))
    assert data_out_base[3:6] == data_out_sub
    assert data_out_base[6:10] == bytearray(range(6, 10))
@@ -129,27 +128,27 @@ class TestSubBuffer(unittest.TestCase):
    self.assertTrue(view2.is_allocated())

    data_in = bytearray([7, 8, 9])
-    view2.copyin(memoryview(data_in))
+    view2.copy_from(Buffer("PYTHON", 3, view2.dtype, opaque=memoryview(data_in)))
    data_out_v2 = bytearray([0]*3)
-    view2.copyout(memoryview(data_out_v2))
+    data_out_v2[:] = view2.as_memoryview()
    assert data_in == data_out_v2

    expected_base_data = memoryview(bytearray(range(10)))
    expected_base_data[4:7] = data_in

    data_out_base = bytearray([0]*10)
-    self.buf.copyout(memoryview(data_out_base))
+    data_out_base[:] = self.buf.as_memoryview()
    assert expected_base_data == data_out_base

  def test_subbuffer_alloc(self):
    sub_buf = self.buf.view(4, dtypes.int8, offset=3)
    sub_buf.allocate()
-    sub_buf.copyin(memoryview(bytearray(range(10, 14))))
+    sub_buf.copy_from(Buffer("PYTHON", 4, dtypes.int8, opaque=memoryview(bytearray(range(10, 14)))))
    assert self.buf.as_memoryview().tolist()[3:7] == sub_buf.as_memoryview().tolist()

    sub_buf = self.buf_unalloc.view(4, dtypes.int8, offset=3)
    sub_buf.allocate()
-    sub_buf.copyin(memoryview(bytearray(range(10, 14))))
+    sub_buf.copy_from(Buffer("PYTHON", 4, dtypes.int8, opaque=memoryview(bytearray(range(10, 14)))))
    assert self.buf_unalloc.as_memoryview().tolist()[3:7] == sub_buf.as_memoryview().tolist()

  def test_subbuffer_dealloc(self):
@@ -33,11 +33,9 @@ def _test_single_value(vals, op, dts):
  alu = uop(uops, op, output_dtype, loads)
  out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), alu))
  buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
-  buf2 = [Buffer(Device.DEFAULT, 1, dtype).allocate().copyin(np.array([a], dtype=_to_np_dtype(dtype)).data) for a,dtype in zip(vals, dts)]
+  buf2 = [Buffer(Device.DEFAULT, 1, dtype, initial_value=np.array([a], dtype=_to_np_dtype(dtype)).tobytes()) for a,dtype in zip(vals, dts)]
  run_uops([out], [buf]+buf2)
-  ret = np.empty(1, _to_np_dtype(output_dtype))
-  buf.copyout(ret.data)
-  return ret[0]
+  return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0]

 def _test_single_value_const(vals, op, dts):
  uops = []
@@ -48,9 +46,7 @@ def _test_single_value_const(vals, op, dts):
  out = buf_store[UOp.const(dtypes.int32, 0)].store(alu)
  buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
  run_uops([out], [buf])
-  ret = np.empty(1, _to_np_dtype(output_dtype))
-  buf.copyout(ret.data)
-  return ret[0]
+  return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0]

 def _test_uops_result(output_dtype, uops, res):
  # uops = []
@@ -59,9 +55,7 @@ def _test_uops_result(output_dtype, uops, res):
  out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), res))
  buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
  run_uops([out], [buf])
-  ret = np.empty(1, _to_np_dtype(output_dtype))
-  buf.copyout(ret.data)
-  return ret[0]
+  return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0]

 class TestUOps(unittest.TestCase):
  def _equal(self, v1, v2):
@@ -31,8 +31,8 @@ class TestHCQ(unittest.TestCase):

  def setUp(self):
    TestHCQ.d0.synchronize()
-    TestHCQ.a.uop.buffer.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
-    TestHCQ.b.uop.buffer.copyin(memoryview(bytearray(struct.pack("ff", 0, 0))))
+    TestHCQ.a.uop.buffer.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 1)))))
+    TestHCQ.b.uop.buffer.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 0)))))
    TestHCQ.d0.synchronize() # wait for copyins to complete

  # Test signals
@@ -376,7 +376,7 @@ class TestHCQ(unittest.TestCase):
    SZ = 200_000_000
    b = Buffer(f"{Device.DEFAULT}:1", SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
    a = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
-    TestHCQ.d0.allocator.map(b._buf)
+    TestHCQ.d0.allocator._map(b._buf)

    sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal()
    TestHCQ.d0.hw_copy_queue_t().timestamp(sig_st) \
@@ -454,7 +454,7 @@ class TestHCQ(unittest.TestCase):
    buf1 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
    buf2 = Buffer(f"{Device.DEFAULT}:1", 1, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
    buf3 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
-    TestHCQ.d0.allocator.map(buf2._buf)
+    TestHCQ.d0.allocator._map(buf2._buf)

    for i in range(256):
      ctypes.memset(buf3._buf.va_addr, i, 1)
@@ -569,7 +569,7 @@ class TestHCQ(unittest.TestCase):

      local_buf = Buffer(f"{Device.DEFAULT}:{devid}", sz, dtypes.uint8, options=BufferSpec(cpu_access=True)).ensure_allocated()

-      d.allocator.map(cpu_buffer._buf)
+      d.allocator._map(cpu_buffer._buf)

      d.hw_copy_queue_t().wait(d.timeline_signal, d.timeline_value - 1) \
                         .copy(local_buf._buf, cpu_buffer._buf, sz) \
@@ -34,7 +34,7 @@ class TestCLError(unittest.TestCase):
    data = list(range(65))
    unaligned = memoryview(bytearray(data))[1:]
    buffer = Buffer("CL", 64, dtypes.uint8).allocate()
-    buffer.copyin(unaligned)
+    buffer.copy_from(Buffer("PYTHON", 64, dtypes.uint8, opaque=unaligned))
    result = memoryview(bytearray(len(data) - 1))
-    buffer.copyout(result)
+    result[:] = buffer.as_memoryview()
    assert unaligned == result, "Unaligned data copied in must be equal to data copied out."
@@ -47,8 +47,8 @@ class TestHCQ(unittest.TestCase):

  def setUp(self):
    TestHCQ.d0.synchronize()
-    TestHCQ.a.uop.buffer.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
-    TestHCQ.b.uop.buffer.copyin(memoryview(bytearray(struct.pack("ff", 0, 0))))
+    TestHCQ.a.uop.buffer.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 1)))))
+    TestHCQ.b.uop.buffer.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 0)))))
    TestHCQ.d0.synchronize() # wait for copyins to complete

  def test_run_1000_times_one_submit(self):
@@ -1,66 +1,31 @@
-import unittest, time
-from tinygrad.runtime.support.usb import ASM24Controller
-from tinygrad.helpers import Timing
+import unittest
+from tinygrad.helpers import Timing, getenv
 from tinygrad import Tensor, Device
 import numpy as np

-class TestASMController(unittest.TestCase):
-  @classmethod
-  def setUpClass(cls):
-    cls.ctrl = ASM24Controller()
-
-  def test_write_and_read(self):
-    base = 0xF000
-    data = b"hello!"
-    self.ctrl.write(base, data)
-    out = self.ctrl.read(base, len(data))
-    self.assertEqual(out, data)
-
-  def test_scsi_write_and_read_from_f000(self):
-    payload = bytes([0x5B]) * 4096
-    self.ctrl.scsi_write(payload, lba=0)
-    back = self.ctrl.read(0xF000, len(payload))
-    self.assertEqual(back, payload)
-
-  def test_scsi_write_speed_4k(self):
-    payload = bytes([0x5A]) * 4096
-    start = time.perf_counter()
-    self.ctrl.scsi_write(payload, lba=0)
-    dur_ms = (time.perf_counter() - start) * 1000
-    print(f"scsi_write 4K took {dur_ms:.3f} ms")
-
-  def test_read_speed_4k(self):
-    payload = bytes([0xA5]) * 4096
-    self.ctrl.write(0xF000, payload)
-    start = time.perf_counter()
-    out = self.ctrl.read(0xF000, 4096)
-    dur_ms = (time.perf_counter() - start) * 1000
-    print(f"read 4K took {dur_ms:.3f} ms")
-    self.assertEqual(out, payload)
-
 class TestDevCopySpeeds(unittest.TestCase):
  @classmethod
  def setUpClass(cls):
-    cls.sz = 512
+    cls.sz = getenv("SIZE", 2e6)
    cls.dev = Device["AMD"]
    if not cls.dev.is_usb(): raise unittest.SkipTest("only test this on USB devices")

  def testCopyCPUtoDefault(self):
    for _ in range(10):
-      t = Tensor.ones(self.sz, self.sz, device="CPU").contiguous().realize()
+      t = Tensor.ones(self.sz, device="CPU", dtype='uchar').contiguous().realize()
      with Timing(f"copyin of {t.nbytes()/1e6:.2f} MB:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns * 1e3:.2f} MB/s"): # noqa: F821
        t.to(Device.DEFAULT).realize()
        Device[Device.DEFAULT].synchronize()
      del t

  def testCopyDefaulttoCPU(self):
-    t = Tensor.ones(self.sz, self.sz).contiguous().realize()
+    t = Tensor.ones(self.sz, dtype='uchar').contiguous().realize()
    for _ in range(10):
      with Timing(f"copyout of {t.nbytes()/1e6:.2f} MB:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns * 1e3:.2f} MB/s"):
        t.to('CPU').realize()

  def testValidateCopies(self):
-    t = Tensor.randn(self.sz, self.sz, device="CPU").contiguous().realize()
+    t = Tensor.randn(self.sz, device="CPU", dtype='uchar').contiguous().realize()
    x = t.to(Device.DEFAULT).realize()
    Device[Device.DEFAULT].synchronize()

@@ -70,4 +35,4 @@ class TestDevCopySpeeds(unittest.TestCase):
    del x, y, t

 if __name__ == "__main__":
-  unittest.main()
+  unittest.main()
@@ -1,7 +1,7 @@
-import random, ctypes
+import random
 import numpy as np
 from tinygrad.device import Buffer, Device
-from tinygrad.helpers import Context, getenv, from_mv
+from tinygrad.helpers import Context, getenv
 from tinygrad.dtype import dtypes
 from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.engine.realize import BufferXfer, get_runner, ExecItem
@@ -29,7 +29,7 @@ def alloc_rawbuffer(device, fill=False):
  if fill:
    with Context(DEBUG=0):
      data = np.random.randint(-10000, 10000, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
-      rawbuf.copyin(Tensor(data).realize().uop.base.realized.as_memoryview())
+      rawbuf.copy_from(Tensor(data).realize().uop.base.realized)
  return rawbuf

 def gen_kernel_ji(device, deps):
@@ -84,9 +84,7 @@ def run_jit(jis, all_buffers, input_buffers, var_vals):
  with Context(DEBUG=0):
    for rawbuf in all_buffers:
      if rawbuf in input_buffers: continue
-      mv = memoryview(bytearray(rawbuf.nbytes))
-      ctypes.memset(from_mv(mv), 0, len(mv))
-      rawbuf.copyin(mv)
+      rawbuf.copy_from(Buffer("PYTHON", rawbuf.size, rawbuf.dtype, opaque=memoryview(bytearray(rawbuf.nbytes))))

  for ei in jis: ei.run(var_vals, jit=True)

@@ -6,7 +6,7 @@ from tinygrad import Tensor, dtypes, Device
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.codegen import to_program
-from tinygrad.dtype import DType
+from tinygrad.dtype import DType, truncate
 from tinygrad.nn.state import get_parameters
 from tinygrad.helpers import T, Target, DEV
 from tinygrad.renderer import Renderer
@@ -38,6 +38,10 @@ def call_is_graph(call:UOp) -> bool:
  ast = call.src[0]
  return ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph"

+def call_is_hcq(call:UOp) -> bool:
+  ast = call.src[0]
+  return ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq"
+
 def jit_cache_count(linear:UOp) -> int:
  n = 0
  for call in linear.src:
@@ -51,6 +55,7 @@ def assert_jit_cache_len(fxn, expected_len):
  if linear is None or not linear.src:
    assert expected_len == 0, expected_len
    return
+  if expected_len and all(call_is_hcq(call) for call in linear.src): expected_len = 3 # HCQ2: merged same-queue calls + finalizer + bumps
  if call_is_graph(linear.src[0]):
    assert len(linear.src) == 1, len(linear.src)
    inner = linear.src[0].src[0].src[0]  # LINEAR UOp inside CUSTOM_FUNCTION
@@ -58,6 +63,8 @@ def assert_jit_cache_len(fxn, expected_len):
  else:
    assert len(linear.src) == expected_len, f"expected {expected_len}, got {len(linear.src)}"

+def min_normal(dt:DType) -> float: return 2.0 ** (2 - (1 << (dtypes.finfo(dt)[0] - 1)))
+
 def rand_for_dtype(dt:DType, size:int, allow_subnormal=True):
  if dtypes.is_unsigned(dt):
    return np.random.randint(0, 100, size=size, dtype=_to_np_dtype(dt))
@@ -66,9 +73,8 @@ def rand_for_dtype(dt:DType, size:int, allow_subnormal=True):
  elif dt == dtypes.bool:
    return np.random.choice([True, False], size=size)
  ret = np.random.uniform(-10, 10, size=size).astype(_to_np_dtype(dt))
-  if not allow_subnormal:
-    min_normal = 2.0 ** (2 - (1 << (dtypes.finfo(dt)[0] - 1)))
-    ret = np.where(np.abs(ret) < min_normal, 0, ret)
+  if dt == dtypes.bfloat16 or dt in dtypes.fp8s: ret = np.array([truncate[dt](x) for x in ret], dtype=ret.dtype)
+  if not allow_subnormal: ret = np.where(np.abs(ret) < min_normal(dt), 0, ret)
  return ret

 def timeit(fxn:Callable[..., T], *args, **kwargs) -> tuple[T, float]:
@@ -9,7 +9,7 @@ MOCKGPU_ARCH = "cdna4" if DEV.arch == "gfx950" else "rdna4" if DEV.arch.startswi
 assert (ma:=getenv("MOCKGPU_ARCH", "")) == "", "MOCKGPU_ARCH is deprecated, use DEV=" + \
  str(replace(DEV.value, arch={"cdna4":"gfx950", "rdna4":"gfx1201"}.get(ma, "gfx1100"))) # type: ignore
 GFX_TARGET_VERSION = {"rdna3": 110000, "rdna4": 120000, "cdna4": 90500}[MOCKGPU_ARCH]
-import tinygrad.runtime.autogen.amd_gpu as amd_gpu, tinygrad.runtime.autogen.am.pm4_nv as pm4
+import tinygrad.runtime.autogen.amd_gpu as amd_gpu, tinygrad.runtime.autogen.am.pm4_nv as pm4, tinygrad.runtime.autogen.am.sdma_6_0_0 as sdma

 SDMA_MAX_COPY_SIZE = 0x400000

@@ -275,6 +275,7 @@ class SDMAExecutor(AMDQueue):
      elif op == amd_gpu.SDMA_OP_POLL_REGMEM: cont = self._execute_poll_regmem()
      elif op == amd_gpu.SDMA_OP_GCR: self._execute_gcr()
      elif op == amd_gpu.SDMA_OP_COPY: self._execute_copy()
+      elif op == sdma.SDMA_OP_WRITE: self._execute_write()
      elif op == amd_gpu.SDMA_OP_TIMESTAMP: self._execute_timestamp()
      elif op == 32: self.rptr[0] += 4  # SDMA_OP_DUMMY_TRAP: pipeline flush, no interrupt
      else: raise RuntimeError(f"Unknown SDMA op {op}")
@@ -289,6 +290,12 @@ class SDMAExecutor(AMDQueue):
    struct = sdma_pkts.trap.from_address(self.base + self.rptr[0] % self.size)
    self.rptr[0] += ctypes.sizeof(struct)

+  def _execute_write(self):
+    packet = to_mv(self.base + self.rptr[0] % self.size, 16).cast('I')
+    addr, count = packet[1] | packet[2] << 32, packet[3] + 1
+    ctypes.memmove(self.gpu.translate_addr(addr), self.base + self.rptr[0] % self.size + 16, count * 4)
+    self.rptr[0] += (4 + count) * 4
+
  def _execute_poll_regmem(self):
    struct = sdma_pkts.poll_regmem.from_address(self.base + self.rptr[0] % self.size)

@@ -803,7 +803,7 @@ def _compile_sopp(inst: ir3.SOPP | ir4.SOPP, ctx: _Ctx) -> UOp:
    pcode = get_pcode(inst.op)
    pc_bytes = ctx.rpc()  # PC is already 64-bit byte address
    vcc, exec_val = ctx.rmask(_c(VCC_LO.offset)), ctx.rexec()
-    srcs = {'PC': pc_bytes.cast(dtypes.int64), 'SIMM16': simm16, 'SCC': ctx.rsgpr_dyn(_c(SCC.offset)), 'VCC': vcc,
+    srcs: dict[str, UOp|int] = {'PC': pc_bytes.cast(dtypes.int64), 'SIMM16': simm16, 'SCC': ctx.rsgpr_dyn(_c(SCC.offset)), 'VCC': vcc,
            'VCCZ': vcc.eq(UOp.const(vcc.dtype, 0)).cast(dtypes.uint32),
            'EXECZ': exec_val.eq(UOp.const(exec_val.dtype, 0)).cast(dtypes.uint32)}
    for dest, val in parse_pcode(pcode, srcs)[1]:
@@ -858,7 +858,7 @@ def _compile_sop(inst: ir3.SOP1|ir3.SOP2|ir3.SOPC|ir3.SOPK|ir4.SOP1|ir4.SOP2|ir4
    if isinstance(inst, ir4.SOPK): s0 = simm16
    elif isinstance(inst, irc.SOPK) and 'CMPK' not in op_name and 'SETREG' not in op_name: s0 = simm16_sext
    else: s0 = ctx.rsgpr_dyn(sdst_off)
-    srcs = {'S0': s0, 'S1': simm16_sext, 'SIMM16': simm16_sext, 'D0': ctx.rsgpr_dyn(sdst_off)}
+    srcs: dict[str, UOp|int] = {'S0': s0, 'S1': simm16_sext, 'SIMM16': simm16_sext, 'D0': ctx.rsgpr_dyn(sdst_off)}
    dst_off, dst_size = sdst_off, 1
    # S_GETREG_B32: extract bits from HW register. Handle as special case since HW_REGISTERS is not a normal variable.
    # HW register values are stored at SGPR[SGPR_COUNT-16 + hwRegId] by _init_wave.
@@ -688,10 +688,10 @@ class Parser:
        return _extract_bits(base, hi, lo)
      # Dynamic bit slice: (base >> lo) & ((1 << (hi - lo + 1)) - 1)
      dt = dtypes.uint64 if base.dtype in (dtypes.uint64, dtypes.int64) else dtypes.uint32
-      hi, lo = first.cast(dt), second.cast(dt)
-      width = hi - lo + _const(dt, 1)
+      hi_u, lo_u = first.cast(dt), second.cast(dt)
+      width = hi_u - lo_u + _const(dt, 1)
      mask = (_const(dt, 1) << width) - _const(dt, 1)
-      return (base.cast(dt) >> lo) & mask
+      return (base.cast(dt) >> lo_u) & mask
    self.eat('RBRACKET')
    dt_suffix = None
    if self.try_eat('DOT'):
@@ -1123,11 +1123,11 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
          val = parse_tokens(toks[j:], env, funcs)
          lo_dt, hi_dt = DTYPES.get(lo_type, dtypes.uint64), DTYPES.get(hi_type, dtypes.uint32)
          lo_bits = 64 if lo_dt in (dtypes.uint64, dtypes.int64) else 32
-          lo_val = val.cast(lo_dt) if val.dtype.itemsize * 8 <= lo_bits else (val & _const(val.dtype, (1 << lo_bits) - 1)).cast(lo_dt)
-          hi_val = (val >> _const(val.dtype, lo_bits)).cast(hi_dt)
-          block_assigns[lo_var] = env[lo_var] = lo_val
-          block_assigns[hi_var] = env[hi_var] = hi_val
-          if assigns is not None: assigns.extend([(f'{lo_var}.{lo_type}', lo_val), (f'{hi_var}.{hi_type}', hi_val)])
+          lo_u = val.cast(lo_dt) if val.dtype.itemsize * 8 <= lo_bits else (val & _const(val.dtype, (1 << lo_bits) - 1)).cast(lo_dt)
+          hi_u = (val >> _const(val.dtype, lo_bits)).cast(hi_dt)
+          block_assigns[lo_var] = env[lo_var] = lo_u
+          block_assigns[hi_var] = env[hi_var] = hi_u
+          if assigns is not None: assigns.extend([(f'{lo_var}.{lo_type}', lo_u), (f'{hi_var}.{hi_type}', hi_u)])
          i += 1
          continue

@@ -1,4 +1,4 @@
-import ctypes, time, os, builtins, fcntl
+import ctypes, time, os, builtins, fcntl, typing
 from tinygrad.helpers import DEV
 from tinygrad.runtime.support.hcq import FileIOInterface
 from tinygrad.runtime.autogen import libc
@@ -9,7 +9,7 @@ start = time.perf_counter()

 drivers = [cls() for t in DEV.value if (cls:={"MOCKPCI+AMD": AMDriver, "MOCKKFD+AMD": AMDDriver, "MOCK+AMD": AMDDriver, "MOCKUSB+AMD": AMUSBDriver,
                                              "MOCK+NV": NVDriver}.get(f"{t.interface}+{t.device}"))]
-tracked_fds = {}
+tracked_fds: dict[int, typing.Any] = {}

 original_memoryview = builtins.memoryview
 class TrackedMemoryView:
@@ -6,38 +6,25 @@ class MockUSB:
  def __init__(self, mem):
    self.mem = mem
  def read(self, address, size): return bytes(self.mem[address:address+size])
-  def write(self, address, data, ignore_cache=False): self.mem[address:address+len(data)] = data
-  def pcie_mem_req(self, address, value=None, size=1):
-    if value is None: return int.from_bytes(self.mem[address:address+size], "little")
-    else: self.mem[address:address+size] = value.to_bytes(size, "little")
-  def pcie_mem_write(self, address, values, size):
-    for i, value in enumerate(values): self.pcie_mem_req(address + i * size, value, size)
+  def write(self, address, data): self.mem[address:address+len(data)] = data
+  def pcie_mem_read(self, address, nbytes): return bytes(self.mem[address:address+nbytes])
+  def pcie_mem_write(self, address, data): self.mem[address:address+len(data)] = data

 # *** ASM24 Controller Mock ***

 _mock_usb_state: MockASM24State|None = None

 class MockASM24State:
-  """Mock ASM24 controller: XRAM memory map, DMA windows, TLP engine, PCI config space.
+  """Mock custom ASM24 controller: XRAM, DMA windows, PCI config space, and GPU BARs.

  Memory map (64KB XRAM):
    0xA000-0xAFFF: DMA window -> sys 0x820000
    0xB000-0xB1FF: DMA window -> sys 0x800000
-    0xB200-0xB7FF: PCI MMIO (TLP engine)
+    0xB200-0xB7FF: controller PCI MMIO
    0xF000-0xFFFF: DMA window -> sys 0x200000 (512KB)
  """
  XRAM_SIZE = 0x10000

-  TLP_FMT_TYPE = 0xB210
-  TLP_BYTE_EN = 0xB217
-  TLP_ADDR_LO = 0xB218
-  TLP_ADDR_HI = 0xB21C
-  TLP_DATA = 0xB220
-  TLP_COMPL = 0xB22A
-  TLP_TRIGGER = 0xB254
-  TLP_LINK_STATUS = 0xB284
-  TLP_STATUS = 0xB296
-
  def __init__(self, gpu, driver, vram_size:int, doorbell_size:int, mmio_size:int):
    self.gpu, self.driver = gpu, driver
    self._xram = bytearray(self.XRAM_SIZE)
@@ -93,53 +80,7 @@ class MockASM24State:
      if ctrl_addr <= addr < ctrl_addr + dma_size:
        (ctypes.c_ubyte * 1).from_address(host_addr + (addr - ctrl_addr))[0] = value
        return
-    if addr == self.TLP_STATUS:
-      self._xram[addr] &= ~value & 0xFF
-      return
    self._xram[addr] = value
-    if addr == self.TLP_TRIGGER and value == 0x0F: self._process_tlp()
-
-  # --- TLP engine ---
-
-  def _process_tlp(self):
-    fmt_type, byte_en = self._xram[self.TLP_FMT_TYPE], self._xram[self.TLP_BYTE_EN]
-    addr_lo = int.from_bytes(self._xram[self.TLP_ADDR_LO:self.TLP_ADDR_LO+4], 'big')
-    addr_hi = int.from_bytes(self._xram[self.TLP_ADDR_HI:self.TLP_ADDR_HI+4], 'big')
-    address = addr_lo | (addr_hi << 32)
-
-    size, offset, tmp = 0, 0, byte_en
-    while tmp and not (tmp & 1):
-      offset += 1
-      tmp >>= 1
-    while tmp:
-      size += tmp & 1
-      tmp >>= 1
-
-    is_write, is_cfg = bool(fmt_type & 0x40), (fmt_type & 0xbe) == 0x04
-
-    if is_cfg:
-      bus, dev, fn, byte_addr = (address >> 24) & 0xFF, (address >> 19) & 0x1F, (address >> 16) & 0x7, address & 0xFFC
-      if is_write:
-        data = int.from_bytes(self._xram[self.TLP_DATA:self.TLP_DATA+4], 'big')
-        self._cfg_write(bus, dev, fn, byte_addr + offset, (data >> (8 * offset)) & ((1 << (8 * size)) - 1), size)
-      else:
-        self._xram[self.TLP_DATA:self.TLP_DATA+4] = int.from_bytes(self._get_cfg(bus, dev, fn)[byte_addr:byte_addr+4], 'little').to_bytes(4, 'big')
-      self._xram[self.TLP_COMPL:self.TLP_COMPL+2] = (4).to_bytes(2, 'big')
-      self._xram[self.TLP_LINK_STATUS] = 0x01 if not is_write else 0x00
-      self._xram[self.TLP_STATUS] = 0x02
-      return
-
-    if is_write:
-      data = int.from_bytes(self._xram[self.TLP_DATA:self.TLP_DATA+4], 'big')
-      self._pcie_dispatch(address + offset, (data >> (8 * offset)) & ((1 << (8 * size)) - 1), size)
-    else:
-      result = self._pcie_dispatch(address + offset, None, size)
-      if result is not None:
-        self._xram[self.TLP_DATA:self.TLP_DATA+4] = ((result << (8 * offset)) & 0xFFFFFFFF).to_bytes(4, 'big')
-
-    self._xram[self.TLP_COMPL:self.TLP_COMPL+2] = (size & 0xFFF).to_bytes(2, 'big')
-    self._xram[self.TLP_LINK_STATUS] = 0x01 if not is_write else 0x00
-    self._xram[self.TLP_STATUS] = 0x02

  def _cfg_write(self, bus:int, dev:int, fn:int, byte_addr:int, val:int, size:int):
    cfg = self._get_cfg(bus, dev, fn)
@@ -168,49 +109,104 @@ class MockASM24State:
    # Generic config write
    for i in range(size): cfg[byte_addr + i] = (val >> (8 * i)) & 0xFF

-  def _pcie_dispatch(self, address:int, value:int|None, size:int) -> int|None:
+  def _find_bar(self, address:int, size:int) -> tuple[int, int]:
    for reg_off, (bar_addr, bar_size) in self._bar_addrs.items():
-      if bar_addr <= address < bar_addr + bar_size:
-        offset = address - bar_addr
-        if reg_off == 0x10:  # BAR0 - VRAM
-          if value is None: return int.from_bytes(bytes(self.gpu.vram[offset:offset+size]), "little")
-          self.gpu.vram[offset:offset+size] = list(value.to_bytes(size, "little"))
-          return None
-        if reg_off == 0x18:  # BAR2 - Doorbell
-          if value is None: return int.from_bytes(bytes(self._doorbell[offset:offset+size]), "little")
-          for i, b in enumerate(value.to_bytes(size, "little")): self._doorbell[offset + i] = b
-          self.driver._emulate_execute()
-          return None
-        if reg_off == 0x24:  # BAR5 - MMIO
-          if value is None: return self.gpu.mmio[offset // 4]
-          self.gpu.mmio[offset // 4] = value
-          return None
-    raise ValueError(f"PCIe address {address:#x} not mapped to any BAR")
+      if bar_addr <= address and address + size <= bar_addr + bar_size: return reg_off, address - bar_addr
+    raise ValueError(f"PCIe range {address:#x}+{size:#x} not mapped to any BAR")

-  # --- CDB processing (called by MockUSB3.send_batch) ---
+  def _pcie_read(self, address:int, size:int) -> bytes:
+    reg_off, offset = self._find_bar(address, size)
+    if reg_off == 0x10: return bytes(self.gpu.vram[offset:offset+size])
+    if reg_off == 0x18: return bytes(self._doorbell[offset:offset+size])
+    if reg_off == 0x24: return bytes((self.gpu.mmio[(offset+i)//4] >> (8*((offset+i)&3))) & 0xFF for i in range(size))
+    raise RuntimeError(f"unsupported BAR register {reg_off:#x}")

-  def process_cdb(self, cdb:bytes, rlen:int, send_data:bytes|None) -> bytes|None:
-    op = cdb[0]
-    if op == 0xE5:  # write byte
-      self._xram_write_byte(((cdb[2] << 16) | (cdb[3] << 8) | cdb[4]) & 0xFFFF, cdb[1])
-      return None
-    if op == 0xE4:  # read
-      return self._xram_read(((cdb[2] << 16) | (cdb[3] << 8) | cdb[4]) & 0xFFFF, cdb[1])
-    if op == 0x8A and send_data is not None and 0xF000 in self._dma_regions:  # SCSI write
-      host_addr, dma_size = self._dma_regions[0xF000]
-      ctypes.memmove(host_addr, send_data, min(len(send_data), dma_size))
+  def _pcie_write(self, address:int, data:bytes):
+    reg_off, offset = self._find_bar(address, len(data))
+    if reg_off == 0x10: self.gpu.vram[offset:offset+len(data)] = list(data)
+    elif reg_off == 0x18:
+      self._doorbell[offset:offset+len(data)] = list(data)
+      self.driver._emulate_execute()
+    elif reg_off == 0x24:
+      updates: dict[int, int] = {}
+      for i, byte in enumerate(data):
+        idx, shift = (offset+i)//4, 8*((offset+i)&3)
+        updates[idx] = (updates.get(idx, self.gpu.mmio[idx]) & ~(0xFF << shift)) | (byte << shift)
+      for idx, val in updates.items(): self.gpu.mmio[idx] = val
+    else: raise RuntimeError(f"unsupported BAR register {reg_off:#x}")
+
+  def _pcie_dispatch(self, address:int, value:int|None, size:int) -> int|None:
+    if value is None: return int.from_bytes(self._pcie_read(address, size), 'little')
+    self._pcie_write(address, value.to_bytes(size, 'little'))
    return None

 class MockUSB3:
  @classmethod
  def list_devices(cls, vendor, dev): return [(0, "usb:mock")]
  def __init__(self, *args, **kwargs):
-    self.product, self.is_custom = "", False
-  def send_batch(self, cdbs:list[bytes], idata:list[int]|None=None, odata:list[bytes|None]|None=None) -> list[bytes|None]:
+    self.product = "custom mock"
+    self._bulk_read_op: tuple[str, int, int]|None = None
+    self._bulk_write_op: tuple[str, int, int]|None = None
+    self._f0_reply = bytes(8)
+
+  @property
+  def state(self) -> MockASM24State:
    assert _mock_usb_state is not None
-    idata, odata = idata or [0] * len(cdbs), odata or [None] * len(cdbs)
-    results: list[bytes|None] = []
-    for cdb, rlen, sdata in zip(cdbs, idata, odata):
-      result = _mock_usb_state.process_cdb(cdb, rlen, sdata)
-      results.append(result if rlen > 0 else None)
-    return results
+    return _mock_usb_state
+
+  def control_write(self, request:int, value:int=0, index:int=0, data:bytes=b'', timeout:int=1000):
+    if request == 0xF3:
+      self.state._xram[0xB450] = 0x78 if value else 0
+    elif request == 0xE5:
+      self.state._xram_write_byte(value, index)
+    elif request == 0xF2:
+      op = ("sram_read" if value & 0x8000 else "sram_write", 0xF000, (value & 0x7FFF) * 512)
+      if value & 0x8000: self._bulk_read_op = op
+      else: self._bulk_write_op = op
+    elif request == 0xF0:
+      address_lo, address_hi, payload = struct.unpack('<III', data)
+      address, fmt_type, byte_en = address_lo | (address_hi << 32), value & 0xFF, value >> 8
+      if index == 1: self._bulk_write_op = ("pcie_write", address, payload * 4)
+      elif index == 2: self._bulk_read_op = ("pcie_read", address, payload * 4)
+      else:
+        assert index == 0 and byte_en
+        offset = (byte_en & -byte_en).bit_length() - 1
+        size, is_write, is_cfg = byte_en.bit_count(), bool(fmt_type & 0x40), (fmt_type & 0xBE) == 0x04
+        if is_cfg:
+          bus, dev, fn, byte_addr = (address >> 24) & 0xFF, (address >> 19) & 0x1F, (address >> 16) & 0x7, address & 0xFFC
+          if is_write: self.state._cfg_write(bus, dev, fn, byte_addr + offset, (payload >> (8 * offset)) & ((1 << (8 * size))-1), size)
+          else: payload = int.from_bytes(self.state._get_cfg(bus, dev, fn)[byte_addr:byte_addr+4], 'little')
+        elif is_write:
+          self.state._pcie_dispatch(address + offset, (payload >> (8 * offset)) & ((1 << (8 * size))-1), size)
+        else: payload = (self.state._pcie_dispatch(address + offset, None, size) or 0) << (8 * offset)
+        self._f0_reply = struct.pack('<I', payload & 0xFFFFFFFF) + bytes(4)
+    else: raise ValueError(f"unsupported control OUT request 0x{request:02X}")
+
+  def control_read(self, request:int, length:int, value:int=0, index:int=0, timeout:int=1000) -> memoryview:
+    if request == 0xE4: data = self.state._xram_read(value, length)
+    elif request == 0xF0: data = self._f0_reply
+    else: raise ValueError(f"unsupported control IN request 0x{request:02X}")
+    return memoryview(data[:length])
+
+  def bulk_write(self, data:bytes, timeout:int=1000):
+    assert self._bulk_write_op is not None
+    op, address, size = self._bulk_write_op
+    assert len(data) == size
+    if op == "sram_write":
+      host_addr, region_size = self.state._dma_regions[address]
+      ctypes.memmove(host_addr, data, min(len(data), region_size))
+    elif op == "pcie_write": self.state._pcie_write(address, data)
+    else: raise RuntimeError(f"cannot bulk write for {op}")
+    self._bulk_write_op = None
+
+  def bulk_read(self, length:int, timeout:int=1000) -> memoryview:
+    assert self._bulk_read_op is not None
+    op, address, size = self._bulk_read_op
+    assert length == size
+    if op == "sram_read":
+      host_addr, region_size = self.state._dma_regions[address]
+      data = bytes((ctypes.c_ubyte * min(length, region_size)).from_address(host_addr))
+    elif op == "pcie_read": data = self.state._pcie_read(address, length)
+    else: raise RuntimeError(f"cannot bulk read for {op}")
+    self._bulk_read_op = None
+    return memoryview(data)
@@ -1,6 +1,6 @@
 import unittest, itertools, math
 from tinygrad import Tensor, dtypes, Context
-from tinygrad.dtype import DType, ConstType
+from tinygrad.dtype import DType, ConstType, Invalid
 from tinygrad.uop.ops import Ops, UOp
 from test.helpers import full_rewrite
 import numpy as np
@@ -34,6 +34,24 @@ class TestUnaryOpsConstFolding(unittest.TestCase):
    x = x.clip(0, 1).realize()
    _check_ast_count(1, x.neg())

+class TestWeakConstFolding(unittest.TestCase):
+  def test_weakint_math(self):
+    out = (UOp.const(dtypes.weakint, 2**40) + UOp.const(dtypes.weakint, 2**40)).simplify()
+    self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakint, 2**41))
+
+  def test_float_unaries(self):
+    for dtype in (dtypes.weakint, dtypes.weakfloat):
+      for op in (Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL):
+        out = UOp.const(dtype, 4).alu(op).simplify()
+        self.assertEqual((out.op, out.dtype), (Ops.CONST, dtypes.weakfloat))
+
+  def test_weakfloat_math(self):
+    out = (UOp.const(dtypes.weakfloat, 1.25) + UOp.const(dtypes.weakfloat, 2.5)).simplify()
+    self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakfloat, 3.75))
+
+  def test_invalid_poison(self):
+    self.assertIs(UOp.const(dtypes.weakint, Invalid).alu(Ops.CDIV, UOp.const(dtypes.weakint, 0)).simplify().arg, Invalid)
+
 class TestBinaryOpsConstFolding(unittest.TestCase):
  def test_add_literal_zero(self):
    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + 0)
@@ -1,6 +1,6 @@
 import unittest, math, struct, operator
 from tinygrad import Tensor, Device
-from tinygrad.dtype import DTYPES_DICT, dtypes, truncate, float_to_fp16, float_to_bf16, _to_np_dtype, least_upper_dtype, least_upper_float
+from tinygrad.dtype import DTYPES_DICT, dtypes, Invalid, truncate, float_to_fp16, float_to_bf16, _to_np_dtype, least_upper_dtype, least_upper_float

 from tinygrad.helpers import getenv
 from hypothesis import given, settings, strategies as strat
@@ -57,6 +57,7 @@ class TestHelpers(unittest.TestCase):

  def test_from_py(self):
    assert dtypes.from_py(True) == dtypes.bool
+    assert dtypes.from_py(Invalid) == dtypes.bool
    assert dtypes.from_py(2) == dtypes.default_int
    assert dtypes.from_py(3.0) == dtypes.default_float
    assert dtypes.from_py([]) == dtypes.default_float
@@ -348,19 +348,6 @@ class TestStopEarly(unittest.TestCase):
    ret = (c+d).substitute({c:cn}, extra_pm=pm_cvisit)
    assert ret == cn+d

-class TestFastSubstitute(unittest.TestCase):
-  def test_replacement_tree_is_substituted(self):
-    a, b, c, d = [UOp.variable(x, 0, 10) for x in "abcd"]
-    self.assertIs((a+4).substitute({a:b+c, b:d}), (d+c)+4)
-
-  def test_rebuilt_node_is_substituted(self):
-    a, b, c, d = [UOp.variable(x, 0, 10) for x in "abcd"]
-    self.assertIs(((a+b)*2).substitute({a:c, c+b:d}), d*2)
-
-  def test_mapping_cycle(self):
-    a, b = [UOp.variable(x, 0, 10) for x in "ab"]
-    with self.assertRaises(RuntimeError): (a+1).substitute({a:b, b:a})
-
 class TestWalkRewrite(unittest.TestCase):
  """Tests for graph_rewrite with walk=True (MLIR Walk Pattern Rewrite Driver semantics).
  walk=True gives a single-pass traversal that does NOT revisit or re-traverse into rewritten subtrees.
@@ -84,22 +84,22 @@ class TestUSBMMIOInterface(unittest.TestCase):
    self.mmio[2] = 0xFE
    self.assertEqual(full_view[2], 0xFE)

-  def test_pcimem_byte(self):
+  def test_pcimem_dword(self):
    usb2 = MockUSB(bytearray(self.size))
-    mmio_pci = USBMMIOInterface(usb2, 0, self.size, fmt='B', pcimem=True)
-    mmio_pci[3] = 0x11
-    self.assertEqual(mmio_pci[3], 0x11)
-    self.assertEqual(usb2.mem[3], 0x11)
+    mmio_pci = USBMMIOInterface(usb2, 0, self.size, fmt='I', pcimem=True)
+    mmio_pci[3] = 0x11223344
+    self.assertEqual(mmio_pci[3], 0x11223344)
+    self.assertEqual(usb2.mem[12:16], b'\x44\x33\x22\x11')

  def test_pcimem_slice(self):
    usb3 = MockUSB(bytearray(self.size))
    mmio_pci = USBMMIOInterface(usb3, 0, self.size, fmt='B', pcimem=True)
-    values = [2, 3, 4]
-    mmio_pci[4:7] = values
-    raw = mmio_pci[4:7]
+    values = [2, 3, 4, 5]
+    mmio_pci[4:8] = values
+    raw = mmio_pci[4:8]
    self.assertIsInstance(raw, bytes)
    self.assertEqual(list(raw), values)
-    self.assertEqual([mmio_pci[i] for i in range(4, 7)], values)
+    self.assertEqual(list(usb3.mem[4:8]), values)

 if __name__ == "__main__":
  unittest.main()
@@ -1,6 +1,7 @@
 import ctypes, gzip, unittest, timeit, pickle
 from tinygrad import Variable
-from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, mv_address, count, all_same
+from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, mv_address, count, all_same
+from tinygrad.tensor import is_numpy_ndarray
 from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
 from tinygrad.helpers import ceildiv, ansistrip, get_shape
 from tinygrad.tensor import Tensor
@@ -1,4 +1,4 @@
-import unittest, threading, time
+import unittest, threading, time, json
 from unittest.mock import Mock

 class TestLLMServer(unittest.TestCase):
@@ -7,12 +7,9 @@ class TestLLMServer(unittest.TestCase):
  @classmethod
  def setUpClass(cls):
    cls.mock_tok = Mock()
-    cls.mock_tok.role = Mock(return_value=[100, 101])
    cls.mock_tok.encode = Mock(return_value=[200, 201, 202])
    cls.mock_tok.decode = Mock(return_value="Hello")
    cls.mock_tok.stream_decoder = Mock(return_value=lambda tid=None: "Hello" if tid is not None else "")
-    cls.mock_tok.end_turn = Mock(return_value=[998])
-    cls.mock_tok.prefix = Mock(return_value=[1])
    cls.mock_tok.preset = "llama3"
    cls.mock_tok.bos_id = 1
    cls.mock_tok.eos_id = 999
@@ -20,12 +17,14 @@ class TestLLMServer(unittest.TestCase):
    cls.mock_tok.is_end = Mock(side_effect=lambda tid: tid in (999,))

    cls.mock_model = Mock()
+    cls.mock_model.max_context = 4
    cls.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 999]))
    cls.mock_model.get_start_pos = Mock(return_value=0)

-    from tinygrad.llm.cli import LLMServer
+    from tinygrad.llm.cli import FallbackTemplate
+    from tinygrad.llm.serve import LLMServer

-    cls.server = LLMServer(('127.0.0.1', 0), cls.mock_model, "test-model", cls.mock_tok)
+    cls.server = LLMServer(('127.0.0.1', 0), cls.mock_model, "test-model", cls.mock_tok, FallbackTemplate(cls.mock_tok))
    cls.port = cls.server.server_address[1]
    cls.server_thread = threading.Thread(target=cls.server.serve_forever, daemon=True)
    cls.server_thread.start()
@@ -131,6 +130,16 @@ class TestLLMServer(unittest.TestCase):
    self.assertIsNotNone(resp.usage.prompt_tokens)
    self.assertIsNotNone(resp.usage.completion_tokens)

+  def test_context_length_error(self):
+    from openai import BadRequestError
+    self.mock_tok.encode.return_value = [200, 201, 202, 203]
+    try:
+      with self.assertRaises(BadRequestError) as err:
+        self.client.chat.completions.create(model="test-model", messages=[{"role":"user", "content":"too long"}])
+      self.assertEqual(err.exception.code, "context_length_exceeded")
+    finally:
+      self.mock_tok.encode.return_value = [200, 201, 202]
+
  def test_max_tokens_streaming(self):
    self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 302, 303, 999]))
    stream = self.client.chat.completions.create(
@@ -149,50 +158,6 @@ class TestLLMServer(unittest.TestCase):
    self.assertEqual(resp.choices[0].finish_reason, "length")
    self.assertEqual(resp.usage.completion_tokens, 2)

-  def test_assistant_prefill(self):
-    """Last assistant message should be treated as prefill (not a completed turn)."""
-    self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 999]))
-    captured_ids = []
-    orig_generate = self.mock_model.generate.side_effect
-    def capture_generate(ids, **kwargs):
-      captured_ids.extend(ids)
-      return orig_generate(ids, **kwargs)
-    self.mock_model.generate = Mock(side_effect=capture_generate)
-
-    resp = self.client.chat.completions.create(
-      model="test", messages=[
-        {"role": "user", "content": "Hello"},
-        {"role": "assistant", "content": "Sure"}
-      ], stream=False
-    )
-    # prefill tokens should be in ids: role("assistant") + encode("Sure") but NO end_turn after it
-    # and NO extra role("assistant") appended
-    role_tokens = self.mock_tok.role.call_args_list
-    # last role() call should be for "assistant" (the prefill message), not an extra one
-    self.assertEqual(role_tokens[-1], unittest.mock.call("assistant"))
-    # end_turn should be called once less than role() — the prefill assistant msg doesn't get end_turn
-    # NOTE: this is flaky in random order
-    #self.assertEqual(self.mock_tok.end_turn.call_count, self.mock_tok.role.call_count - 1)
-    self.assertIsNotNone(resp.choices[0].message.content)
-
-  def test_assistant_prefill_not_last(self):
-    """Assistant message that's NOT last should be a normal completed turn."""
-    self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 999]))
-    self.mock_tok.role.reset_mock()
-    self.mock_tok.end_turn.reset_mock()
-    self.client.chat.completions.create(
-      model="test", messages=[
-        {"role": "user", "content": "Hello"},
-        {"role": "assistant", "content": "Sure"},
-        {"role": "user", "content": "Continue"}
-      ], stream=False
-    )
-    # all messages get end_turn, plus an extra role("assistant") at the end
-    # roles: user, assistant, user, assistant(generation prompt) = 4 role calls
-    # end_turns: user, assistant, user = 3 end_turn calls (one per message)
-    self.assertEqual(self.mock_tok.end_turn.call_count, 3)
-    self.assertEqual(self.mock_tok.role.call_count, 4)
-
  def test_models_endpoint(self):
    import requests as req
    resp = req.get(f"http://127.0.0.1:{self.port}/v1/models")
@@ -203,5 +168,103 @@ class TestLLMServer(unittest.TestCase):
    self.assertEqual(data["data"][0]["id"], "test-model")
    self.assertEqual(data["data"][0]["object"], "model")

+class TestLLMToolCalls(unittest.TestCase):
+  """Tool calling through the OpenAI-compatible HTTP API."""
+
+  @classmethod
+  def setUpClass(cls):
+    cls.mock_tok = Mock()
+    cls.mock_tok.encode = Mock(return_value=[200, 201, 202])
+    cls.mock_tok.decode = Mock(return_value="")
+    cls.mock_tok.preset = "qwen2"
+    cls.mock_tok.bos_id, cls.mock_tok.eos_id, cls.mock_tok.eot_id = None, 999, None
+    cls.mock_tok.is_end = Mock(return_value=False)
+
+    cls.mock_model = Mock()
+    cls.mock_model.max_context = 4
+    cls.mock_model.get_start_pos = Mock(return_value=0)
+
+    from tinygrad.llm.serve import LLMServer
+    import jinja2
+    # .items() matches tool-aware templates and ensures OpenAI JSON argument strings are normalized before rendering the next turn.
+    template = jinja2.Template("""{% for m in messages %}{{ m.content or '' }}{% for tc in m.tool_calls or [] %}
+      {% for key, value in tc.function.arguments.items() %}{{ key }}={{ value }}{% endfor %}{% endfor %}{% endfor %}""")
+    cls.server = LLMServer(('127.0.0.1', 0), cls.mock_model, "tool-model", cls.mock_tok, template)
+    cls.port = cls.server.server_address[1]
+    cls.server_thread = threading.Thread(target=cls.server.serve_forever, daemon=True)
+    cls.server_thread.start()
+    time.sleep(0.1)
+
+    from openai import OpenAI
+    cls.client = OpenAI(base_url=f"http://127.0.0.1:{cls.port}/v1", api_key="test")
+
+  @classmethod
+  def tearDownClass(cls):
+    cls.server.shutdown()
+    cls.server.server_close()
+
+  def set_output(self, text:str):
+    pieces = dict(enumerate(text, 1))
+    self.mock_tok.stream_decoder = Mock(return_value=lambda tid=None: pieces[tid] if tid is not None else "")
+    self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter(pieces))
+
+  @staticmethod
+  def tools():
+    return [{"type":"function", "function":{"name":"read", "description":"Read a file",
+      "parameters":{"type":"object", "properties":{"path":{"type":"string"}}, "required":["path"]}}}]
+
+  def test_streaming_tool_call(self):
+    self.set_output('before<tool_call>{"name":"read","arguments":{"path":"README.md"}}</tool_call>')
+    chunks = list(self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Read README.md"}],
+                                                     tools=self.tools(), stream=True))
+    self.assertEqual("".join(c.choices[0].delta.content or "" for c in chunks if c.choices), "before")
+    calls = [tc for c in chunks if c.choices for tc in c.choices[0].delta.tool_calls or []]
+    self.assertEqual(len(calls), 1)
+    self.assertEqual(calls[0].function.name, "read")
+    self.assertEqual(json.loads(calls[0].function.arguments), {"path":"README.md"})
+    self.assertEqual(chunks[-1].choices[0].finish_reason, "tool_calls")
+
+  def test_multiple_xml_tool_calls(self):
+    self.set_output("<tool_call><function=read><parameter=path>\"a\"</parameter></function></tool_call>"
+                    "<tool_call><function=read><parameter=path>\"b\"</parameter></function></tool_call>")
+    response = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Read a and b"}],
+                                                   tools=self.tools())
+    self.assertEqual([json.loads(tc.function.arguments)["path"] for tc in response.choices[0].message.tool_calls], ["a", "b"])
+    self.assertEqual(response.choices[0].finish_reason, "tool_calls")
+
+  def test_multiline_tool_argument_preserves_trailing_newline(self):
+    self.set_output("<tool_call>\n<function=write>\n<parameter=content>\nfirst\nsecond\n\n</parameter>\n"
+                    "<parameter=filePath>\nout.txt\n</parameter>\n</function>\n</tool_call>")
+    response = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Write out.txt"}], tools=self.tools())
+    args = json.loads(response.choices[0].message.tool_calls[0].function.arguments)
+    self.assertEqual(args, {"content":"first\nsecond\n", "filePath":"out.txt"})
+
+  def test_invalid_tool_call_becomes_content(self):
+    self.set_output("<tool_call>not a call</tool_call>")
+    response = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Hello"}], tools=self.tools())
+    self.assertEqual(response.choices[0].message.content, "<tool_call>not a call</tool_call>")
+    self.assertIsNone(response.choices[0].message.tool_calls)
+    self.assertEqual(response.choices[0].finish_reason, "stop")
+
+  def test_tool_call_in_reasoning_is_not_executed(self):
+    self.set_output('<think>draft <tool_call>{"name":"wrong","arguments":{}}</tool_call></think>answer')
+    response = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Hello"}], tools=self.tools())
+    self.assertEqual(response.choices[0].message.content, "answer")
+    self.assertIsNone(response.choices[0].message.tool_calls)
+    self.assertEqual(response.choices[0].finish_reason, "stop")
+
+  def test_tool_result_round_trip(self):
+    self.set_output('<tool_call>{"name":"read","arguments":{"path":"README.md"}}</tool_call>')
+    first = self.client.chat.completions.create(model="tool-model", messages=[{"role":"user", "content":"Read README.md"}], tools=self.tools())
+    call = first.choices[0].message.tool_calls[0]
+    self.set_output("done")
+    second = self.client.chat.completions.create(model="tool-model", messages=[
+      {"role":"user", "content":"Read README.md"},
+      {"role":"assistant", "content":None, "tool_calls":[call.model_dump()]},
+      {"role":"tool", "tool_call_id":call.id, "content":"file contents"},
+    ], tools=self.tools())
+    self.assertEqual(second.choices[0].message.content, "done")
+    self.assertEqual(second.choices[0].finish_reason, "stop")
+
 if __name__ == '__main__':
  unittest.main()
@@ -1,5 +1,5 @@
-import unittest, base64, functools, sys
-from tinygrad.llm.cli import SimpleTokenizer
+import unittest, base64, functools, re, sys, time, unicodedata
+from tinygrad.llm.cli import SimpleTokenizer, FallbackTemplate
 from tinygrad.helpers import fetch

@unittest.skipIf(sys.platform == 'win32', "fetch race condition on Windows")
@@ -46,6 +46,41 @@ class TestLLMTokenizer(unittest.TestCase):
  def test_llama_repeat(self): self._test_coding(self.llama_tok, "00000000000000000", [ 931, 931, 931, 931, 931, 410 ])
  def test_llama_pat(self): self._test_coding(self.llama_tok, "today\n  \n", [ 31213, 14211 ])

+  def test_split_regex_matches_naive_listing(self):
+    # the compacted codepoint ranges must match the same text as listing every codepoint
+    def naive(pre): return "".join(re.escape(chr(cp)) for cp in range(0x323b0) if unicodedata.category(chr(cp)).startswith(pre))
+    r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + naive("Z"), naive("N"), naive("L")
+    naive_re = re.compile("(?i:'s|'t|'re|'ve|'m|'ll|'d)|" +
+      f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+")
+    sample = "hello world 한국어 中文 текст ١٢٣ 123 😊\n  \ttoday\n'équivalent ²³№ "
+    self.assertEqual(SimpleTokenizer({}, {})._split_to_word.findall(sample), naive_re.findall(sample))
+
+  def test_split_regex_speed(self):
+    # the naive listing compiles a 429KB pattern that takes 10+s to match a 225KB prompt; ranges keep it small and fast
+    tok = SimpleTokenizer({}, {})
+    self.assertLess(len(tok._split_to_word.pattern), 100_000)
+    text = "The quick brown fox jumps over the lazy dog. " * 5000
+    tok._split_to_word.findall(text)  # warmup
+    tms = []
+    for _ in range(5):
+      st = time.perf_counter()
+      words = tok._split_to_word.findall(text)
+      tms.append(time.perf_counter() - st)
+    self.assertLess(min(tms), 4)  # best-of-5 is robust to CI scheduling pauses; new code takes ~60ms
+    self.assertEqual(len(words), 50001)
+
+  def test_llama_continued_conversation(self):
+    self._test_coding(self.llama_tok, "hello <|eot_id|>world", [15339, 220, 128009, 14957])
+    self._test_coding(self.llama_tok, "hello <|eot_id|>world again", [15339, 220, 128009, 14957, 1578])
+    self._test_coding(self.llama_tok, "hello changed <|eot_id|>world again", [15339, 5614, 220, 128009, 14957, 1578])
+
+  def test_long_cached_prompt_matches_fresh_tokenization(self):
+    prefix = "system tools\n" * 700 + "<|eot_id|>"
+    first, changed = prefix + "run tower of hanoi", prefix + "run ls /"
+    expected = self.llama_tok.encode(changed)
+    self.llama_tok.encode(first)
+    self.assertEqual(self.llama_tok.encode(changed), expected)
+
  def test_tekken_from_gguf_kv(self):
    kv = {
      "tokenizer.ggml.tokens": ["<unk>", "<s>", "</s>", "[INST]", "[/INST]", "hello"],
@@ -54,10 +89,11 @@ class TestLLMTokenizer(unittest.TestCase):
      "tokenizer.ggml.eos_token_id": 2,
    }
    tok = SimpleTokenizer.from_gguf_kv(kv)
-    self.assertEqual(tok.role("user"), [3])
+    template = FallbackTemplate(tok)
+    self.assertEqual(template.role("user"), "[INST]")
    self.assertEqual(tok.encode("hello"), [5])
-    self.assertEqual(tok.end_turn(), [4])
-    self.assertEqual(tok.role("assistant"), [])
+    self.assertEqual(template.end_turn(), "[/INST]")
+    self.assertEqual(template.role("assistant"), "")

  def test_stream_decoder(self):
    """stream_decoder buffers incomplete UTF-8: token 25677 has 3/4 of emoji, token 138 completes it."""
@@ -51,8 +51,8 @@ class TestPatternMatcher(unittest.TestCase):
      ctx.append(True)
      assert len(x.src) == 0
      return x.replace(src=(UOp(Ops.NOOP),))
-    matcher = PatternMatcher([(UPat(Ops.CONST, src=(), name="x"), fxn)])
-    c1 = UOp.const(dtypes.float, 1.0)
+    matcher = PatternMatcher([(UPat(Ops.NOOP, src=(), name="x"), fxn)])
+    c1 = UOp(Ops.NOOP)
    # second rewrite shouldn't match anything
    ctx = []
    c1 = matcher.rewrite(c1, ctx)
@@ -1472,6 +1472,18 @@ class TestSchedule(unittest.TestCase):
    x.softmax().sum().backward()
    run_linear(*check_schedule(x.grad, 4))

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

-  @unittest.skip("no longer testable standalone")
-  def test_wmma_vectorize_fold(self):
-    for i in [2, 4, 8]:
-      vec = UOp(Ops.STACK, dtypes.half, tuple(UOp.const(dtypes.half, 0.0) for _ in range(i)))
-      var = UOp.variable("var", 0, 1, dtypes.half)
-      acc = UOp.variable('acc', 0, 1, dtypes.half)
-      wmma = UOp(Ops.WMMA, src=(vec, var, acc))
-      uops = to_uops_list([wmma])
-      self.assertEqual(uops[0], acc)
-      self.assertEqual(len(uops), 2)  # +1 for SINK
-
-    for i in [2, 4, 8]:
-      var = UOp.variable("var", 0, 1, dtypes.half)
-      vec = UOp(Ops.STACK, dtypes.half, tuple(UOp.const(dtypes.half, 0.0) for _ in range(i)))
-      acc = UOp.variable('acc', 0, 1, dtypes.half)
-      wmma = UOp(Ops.WMMA, src=(var, vec, acc))
-      uops = to_uops_list([wmma])
-      self.assertEqual(uops[0], acc)
-      self.assertEqual(len(uops), 2)  # +1 for SINK
-
-  @unittest.skip("wmma is wrong here, it needs an arg")
-  def test_wmma_vectorize_no_fold(self):
-    for i in [4, 8]:
-      vec = UOp(Ops.STACK, dtypes.half,
-                tuple(UOp.const(dtypes.half, 0.0) for _ in range(i//2)) +
-                tuple(UOp.variable(f'tmp{j}', 0, 1, dtypes.half) for j in range(i//2)))
-      var = UOp.variable(f'tmp{i}', 0, 1, dtypes.half)
-      acc = UOp.variable('acc', 0, 1, dtypes.half)
-      wmma = UOp(Ops.WMMA, src=(vec, var, acc))
-      uops = to_uops_list([wmma])
-      self.assertEqual(uops[-2], wmma)  # -2 to skip SINK
-
-    for i in [4, 8]:
-      var = UOp.variable(f'tmp{i}', 0, 1, dtypes.half)
-      vec = UOp(Ops.STACK, dtypes.half,
-                tuple(UOp.const(dtypes.half, 0.0) for _ in range(i//2)) +
-                tuple(UOp.variable(f'tmp{j}', 0, 1, dtypes.half) for j in range(i//2)))
-      acc = UOp.variable('acc', 0, 1, dtypes.half)
-      wmma = UOp(Ops.WMMA, src=(var, vec, acc))
-      uops = to_uops_list([wmma])
-      self.assertEqual(uops[-2], wmma)  # -2 to skip SINK
-
-    for i in [2, 4, 8]:
-      vec = UOp(Ops.STACK, dtypes.half,
-                tuple(UOp.const(dtypes.half, 1.0 if j == 0 else 0.0) for j in range(i)))
-      var = UOp.variable(f'tmp{i}', 0, 1, dtypes.half)
-      acc = UOp.variable('acc', 0, 1, dtypes.half)
-      wmma = UOp(Ops.WMMA, src=(vec, var, acc))
-      uops = to_uops_list([wmma])
-      self.assertEqual(uops[-2], wmma)  # -2 to skip SINK
-
-    for i in [2, 4, 8]:
-      var = UOp.variable(f'tmp{i}', 0, 1, dtypes.half)
-      vec = UOp(Ops.STACK, dtypes.half,
-                tuple(UOp.const(dtypes.half, 1.0 if j == 0 else 0.0) for j in range(i)))
-      acc = UOp.variable('acc', 0, 1, dtypes.half)
-      wmma = UOp(Ops.WMMA, src=(var, vec, acc))
-      uops = to_uops_list([wmma])
-      self.assertEqual(uops[-2], wmma)  # -2 to skip SINK
-
  def test_cast_alu_fold(self):
    d0 = UOp.param(0, dtypes.bool, (1,))
    d1 = UOp.param(1, dtypes.int, (1,))
@@ -127,6 +127,13 @@ class TestVminVmaxProperties(unittest.TestCase):
    self.assertEqual(x.vmin, 0)
    self.assertEqual(x.vmax, 10 >> 2)

+  def test_vmin_vmax_cast_unsigned(self):
+    # a fitting source keeps exact bounds: no wrap can occur
+    self.assertEqual(UOp.variable('x', 5, 10).cast(dtypes.uint8)._min_max, (5, 10))
+    # a possibly-negative or too-large source can wrap: conservative
+    self.assertEqual(UOp.variable('x', -1, 10).cast(dtypes.uint8)._min_max, (0, 255))
+    self.assertEqual(UOp.variable('x', 250, 260).cast(dtypes.uint8)._min_max, (0, 255))
+
  def test_vmin_vmax_xor_neg1(self):
    x = UOp.variable('x', 3, 7)
    uop = x ^ -1
@@ -3,13 +3,50 @@ import unittest
 import numpy as np
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import Timing, Context, cdiv
-from tinygrad.dtype import dtypes, ConstFloat  # noqa: F401
+from tinygrad.dtype import dtypes, ConstFloat, Invalid  # noqa: F401
 from tinygrad.device import Device
-from tinygrad.uop.ops import Ops, ParamArg, UOp, UPat, exec_alu  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
-from tinygrad.uop.spec import spec_shared
+from tinygrad.uop.ops import Ops, ParamArg, UOp, UPat, dtype_from_uop, exec_alu  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
+from tinygrad.uop.spec import spec_program, spec_shared, type_verify
 from tinygrad.uop.symbolic import sym
 from test.helpers import eval_uop, to_uops_list

+class TestDTypeFromUOp(unittest.TestCase):
+  def test_broadcastable_promotion(self):
+    self.assertEqual(dtype_from_uop(Ops.ADD, (UOp.const(dtypes.float32, 1.0), UOp.const(dtypes.float16, 1.0)), None), dtypes.float32)
+    self.assertEqual(dtype_from_uop(Ops.MUL, (UOp.const(dtypes.int8, 1), UOp.const(dtypes.int32, 1)), None), dtypes.int32)
+    self.assertEqual(dtype_from_uop(Ops.ADD, (UOp.const(dtypes.weakint, 1), UOp.const(dtypes.int8, 1)), None), dtypes.int8)
+
+  def test_same_dtype_fast_path(self):
+    src = (UOp.const(dtypes.index, 1), UOp.const(dtypes.index, 2))
+    self.assertEqual(dtype_from_uop(Ops.ADD, src, None), dtypes.index)
+
+  def test_where_promotion(self):
+    cond = UOp.const(dtypes.bool, True)
+    self.assertEqual(dtype_from_uop(Ops.WHERE, (cond, UOp.const(dtypes.float32, 1.0), UOp.const(dtypes.float16, 1.0)), None), dtypes.float32)
+    idx = UOp.range(4, 0)
+    self.assertEqual(idx.valid(idx < 4).dtype, dtypes.index)
+
+  def test_const_dtype_from_value(self):
+    self.assertEqual(dtype_from_uop(Ops.CONST, (), True), dtypes.bool)
+    self.assertEqual(dtype_from_uop(Ops.CONST, (), 3), dtypes.weakint)
+    self.assertEqual(dtype_from_uop(Ops.CONST, (), ConstFloat(3.0)), dtypes.weakfloat)
+    self.assertEqual(dtype_from_uop(Ops.CONST, (), Invalid), dtypes.bool)
+    self.assertRaises(TypeError, dtype_from_uop, Ops.CONST, (), (1, 2))
+
+  @Context(SPEC=2)
+  def test_const_default_dtype_is_derived(self):
+    self.assertEqual(UOp(Ops.CONST, arg=3).dtype, dtypes.weakint)
+    self.assertEqual(UOp(Ops.CONST, arg=ConstFloat(3.0)).dtype, dtypes.weakfloat)
+    self.assertEqual(UOp(Ops.CONST, arg=True).dtype, dtypes.bool)
+    self.assertEqual(UOp(Ops.CONST, arg=Invalid).dtype, dtypes.bool)
+    # an explicit (strong) const dtype is legal until the field is removed
+    self.assertEqual(UOp.const(dtypes.int32, 3).dtype, dtypes.int32)
+
+  def test_weak_dtype_rejected_by_program_spec(self):
+    for weak, concrete, value in ((dtypes.weakint, dtypes.int32, 1), (dtypes.weakfloat, dtypes.float32, 1.0)):
+      with self.assertRaises(RuntimeError): type_verify(UOp.const(weak, value).sink(), spec_program)
+      type_verify(UOp.const(concrete, value).sink(), spec_program)
+
 class TestSafeCast(unittest.TestCase):
  def test_cast_folds(self):
    a = UOp.variable("a", 1, 10, dtype=dtypes.int32)
@@ -40,6 +77,12 @@ class TestExecALU(unittest.TestCase):
  def test_sqrt(self):
    self.assertEqual(exec_alu(Ops.SQRT, dtypes.float, (0.0,)), 0.0)

+  def test_invalid_poison(self):
+    # Invalid poisons any binary op regardless of result dtype: a comparison must not fold to a boolean
+    self.assertIs(exec_alu(Ops.CMPLT, dtypes.bool, (Invalid, 1)), Invalid)
+    self.assertIs(exec_alu(Ops.CMPNE, dtypes.bool, (Invalid, 1)), Invalid)
+    self.assertIs(exec_alu(Ops.ADD, dtypes.index, (Invalid, 1)), Invalid)
+
  def test_div(self):
    self.assertEqual(exec_alu(Ops.CDIV, dtypes.int8, (8, 2)), 4)
    self.assertEqual(exec_alu(Ops.CDIV, dtypes.int8, (7, 3)), 2)
@@ -368,5 +411,16 @@ class TestUOpRender(unittest.TestCase):
    u = UOp(Ops.STACK, dtype=dtypes.int, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)))
    self.assertEqual(u.render(), "{0,1,2}")

+class TestContiguousViewOffset(unittest.TestCase):
+  def _check(self, u, expected): self.assertEqual(u.contiguous_view_offset(), expected)
+
+  def test_simple(self): self._check(UOp.empty(10), 0)
+  def test_shrink(self): self._check(UOp.empty(10)[1:8], 1)
+  def test_2d(self): self._check(UOp.empty(2,5)[1, 2:4], 7)
+  def test_shrink_to_one(self): self._check(UOp.empty(10)[1], 1)
+  def test_expand_is_none(self): self._check(UOp.empty(1).expand(2), None)
+  def test_shrink_invalid(self): self._check(UOp.empty(4).pad((2,2))[0], None)
+  def test_strided(self): self._check(UOp.empty(4)[::2], None)
+
 if __name__ == '__main__':
  unittest.main()
@@ -227,9 +227,9 @@ class TestViz(unittest.TestCase):
    pm = PatternMatcher([(UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4))])
    with save_viz() as viz:
      inner = UOp.const(dtypes.int, 3)
-      func = UOp(Ops.FUNCTION, src=(UOp(Ops.SINK, src=(inner,)),))
-      call = UOp(Ops.CALL, src=(func,))
-      graph_rewrite(call, TrackedPatternMatcher(pm.patterns), enter_calls=True)
+      call = UOp(Ops.CALL, src=(UOp(Ops.SINK, src=(inner,)),))
+      func = UOp(Ops.FUNCTION, src=(UOp(Ops.TUPLE, src=(call,)),))
+      graph_rewrite(func, TrackedPatternMatcher(pm.patterns), enter_calls=True)
    details = list(viz.get_details(0, 0))
    self.assertTrue(details[-1]["change"], "viz replay should detect change inside CALL")

@@ -149,7 +149,8 @@ class TestTensorCores(unittest.TestCase):

        # TODO: support this even if numpy doesn't
        if _to_np_dtype(real_bufs[0].dtype) is None: continue
-        real_bufs[0].copyin(np.zeros((real_bufs[0].size, ), dtype=_to_np_dtype(real_bufs[0].dtype)).data) # Zero to check that all values are filled
+        # Zero to check that all values are filled
+        real_bufs[0].copy_from(Buffer("PYTHON", real_bufs[0].size, real_bufs[0].dtype, opaque=memoryview(bytearray(real_bufs[0].nbytes))))
        run_program(ast, real_bufs)
        result = np.frombuffer(real_bufs[0].as_memoryview(), _to_np_dtype(real_bufs[0].dtype))

@@ -39,14 +39,17 @@ class TestTiny(unittest.TestCase):
    out = Tensor.ones(N).contiguous().sum()
    self.assertEqual(out.item(), N)

-  def test_gemm(self, N=getenv("GEMM_N", 64)):
-    a = Tensor.ones(N,N).contiguous()
-    b = Tensor.eye(N).clone()
+  def test_gemm(self, N=getenv("GEMM_N", 64), dtype=dtypes.float):
+    a = Tensor.ones(N,N, dtype=dtype).contiguous()
+    b = Tensor.eye(N, dtype=dtype).clone()
    lst = (out:=a@b).tolist()
    for y in range(N):
      for x in range(N):
        self.assertEqual(lst[y][x], 1.0, msg=f"mismatch at ({y},{x})")
-    self.assertEqual(out.dtype, dtypes.float)
+    self.assertEqual(out.dtype, dtype)
+
+  @unittest.skipIf(Device.DEFAULT == "DSP", "half is broken on DSP")
+  def test_hgemm(self): self.test_gemm(dtype=dtypes.half)

  def test_gemv(self, N=getenv("GEMV_N", 64), out_dtype=dtypes.float):
    a = Tensor.ones(1,N).contiguous()
@@ -180,11 +180,15 @@ class TestSafetensors(TempDirTestCase):

  def test_save_all_dtypes(self):
    for dtype in dedup(DTYPES_DICT.values()):
-      if dtype in [dtypes.bfloat16]: continue # not supported in numpy
+      if dtype in dtypes.fp8_fnuz: continue # not supported by safetensors
      path = self.tmp(f"ones.{dtype}.safetensors")
      ones = Tensor(np.random.rand(10,10), dtype=dtype)
      safe_save(get_state_dict(ones), path)
-      np.testing.assert_equal(ones.numpy(), list(safe_load(path).values())[0].numpy())
+      loaded = list(safe_load(path).values())[0]
+      # numpy has no fp8 or bfloat16, compare the stored bytes
+      if dtype == dtypes.bfloat16 or dtype in dtypes.fp8s:
+        np.testing.assert_equal(ones.bitcast(dtypes.uint8).numpy(), loaded.bitcast(dtypes.uint8).numpy())
+      else: np.testing.assert_equal(ones.numpy(), loaded.numpy())

  def test_load_supported_types(self):
    import torch
@@ -149,6 +149,9 @@ class TestAutoCastType(unittest.TestCase):
  def tearDown(self):
    dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float

+  def test_int_sqrt(self):
+    _assert_eq(Tensor([1, 4, 9, 16]).sqrt(), dtypes.default_float, [1, 2, 3, 4])
+
  @given(strat.sampled_from([d for d in core_dtypes if dtypes.is_int(d) and d in supported_dtypes]))
  def test_int_to_float_unary_func(self, dtype):
    for func in [
@@ -0,0 +1,183 @@
+import pathlib, tempfile, unittest
+from unittest.mock import patch
+
+from tinygrad import Tensor, dtypes
+from tinygrad.uop import Ops
+from tinygrad.uop.ops import UOp
+from tinygrad.uop.spec import spec_tensor
+from tinygrad.nn.state import safe_save
+
+
+class TestWeakPromotion(unittest.TestCase):
+  def test_rand_requires_concrete(self):
+    with self.assertRaises(ValueError): Tensor.rand(2, dtype=dtypes.weakfloat)
+    with self.assertRaises(ValueError): Tensor.const(dtypes.weakfloat, 1.0).rand_like()
+    with self.assertRaises(ValueError): Tensor.const(dtypes.weakfloat, 1.0).randn_like()
+
+  def test_sum_stays_weak(self):
+    for weak, value in ((dtypes.weakint, 1), (dtypes.weakfloat, 1.0)):
+      self.assertEqual(Tensor.const(weak, value).expand(3).sum().dtype, weak)
+    self.assertEqual((Tensor.const(dtypes.weakfloat, 1.0).expand(3).sum() + Tensor([1], dtype=dtypes.float16)).dtype, dtypes.float16)
+
+  def test_storage_width(self):
+    t = Tensor.const(dtypes.weakint, 2)
+    for fn in (lambda: t.bitcast(dtypes.int32), lambda: Tensor.const(dtypes.int32, 2).bitcast(dtypes.weakint), t.element_size, t.nbytes):
+      with self.assertRaises(RuntimeError): fn()
+
+  def test_materialize_at_default_dtype(self):
+    for weak, value, strong in ((dtypes.weakint, 3, dtypes.default_int), (dtypes.weakfloat, 0.5, dtypes.default_float)):
+      t = Tensor.const(weak, value)
+      self.assertEqual(t.dtype, weak)
+      self.assertEqual(t.data().itemsize, strong.itemsize)
+      self.assertEqual(t.numpy().dtype.itemsize, strong.itemsize)
+      with self.assertRaises(RuntimeError): t.clone("CPU")
+    with patch.object(dtypes, "default_int", dtypes.int64):
+      self.assertEqual(Tensor.const(dtypes.weakint, 3).numpy().dtype.itemsize, dtypes.int64.itemsize)
+
+  def test_uop_scalar_const_unchanged(self):
+    for dtype, value in ((dtypes.index, 1), (dtypes.int32, 1), (dtypes.float32, 0.5)):
+      out = UOp.variable("x", 0.0 if dtype == dtypes.float32 else 0, 10.0 if dtype == dtypes.float32 else 10, dtype) + value
+      self.assertEqual((out.dtype, out.src[1].dtype), (dtype, dtype))
+
+  @unittest.expectedFailure  # TODO: a weak const defers to its consumer (JAX): these dtypes change once python scalars are weak consts
+  def test_changed_rows(self):
+    t_i8, t_f16, t_bf16 = Tensor([1], dtype=dtypes.int8), Tensor([1], dtype=dtypes.float16), Tensor([1], dtype=dtypes.bfloat16)
+    t_bool, t_u16 = Tensor([True]), Tensor([1], dtype=dtypes.uint16)
+    self.assertEqual((t_i8 + 0.5).dtype, dtypes.weakfloat)
+    self.assertEqual(((t_i8 + 0.5) + t_f16).dtype, dtypes.float16)
+    self.assertEqual(((t_i8 + 0.5) + t_bf16).dtype, dtypes.bfloat16)
+    self.assertEqual(((t_bool + 1) + t_i8).dtype, dtypes.int8)
+    self.assertEqual(((t_bool + 1) + t_u16).dtype, dtypes.uint16)
+    self.assertEqual((Tensor(3) + t_i8).dtype, dtypes.int8)
+    self.assertEqual(Tensor([2], dtype=dtypes.uint8).pad(((1, 1),), value=1).dtype, dtypes.uint8)
+    # zeros/ones are full with a python fill value, so they are weak too (jnp.zeros pins float32; deliberate divergence)
+    self.assertEqual((Tensor.zeros(3) + t_f16).dtype, dtypes.float16)
+
+  def test_unchanged_rows(self):
+    t_i8, t_f16, t_f32 = Tensor([1], dtype=dtypes.int8), Tensor([1], dtype=dtypes.float16), Tensor([1], dtype=dtypes.float32)
+    self.assertEqual((t_i8 + 1).dtype, dtypes.int8)
+    self.assertEqual((t_f16 + 0.5).dtype, dtypes.float16)
+    self.assertEqual((t_f32 + t_f16).dtype, dtypes.float32)
+
+  @unittest.expectedFailure  # TODO: dot of a weak const tensor defers to the other operand once python scalars are weak consts
+  def test_dot_defers_weak(self):
+    weak = Tensor([True, False]).where(Tensor(1), 2)
+    self.assertEqual(weak.dot(Tensor([1, 1], dtype=dtypes.int8)).dtype, dtypes.int8)
+
+  @unittest.expectedFailure  # TODO: Tensor(3).uop becomes CONST(weakint); Tensor.dtype is always uop.dtype; buffers lower to the default
+  def test_dtype_is_uop_dtype(self):
+    for value, weak, lowered in ((3, dtypes.weakint, dtypes.default_int), (0.5, dtypes.weakfloat, dtypes.default_float)):
+      t = Tensor(value)
+      self.assertEqual((t.uop.dtype, t.dtype), (weak, weak))
+      self.assertEqual(t.numpy().dtype.itemsize, lowered.itemsize)
+      realized = t.clone("CPU").realize()
+      self.assertEqual((realized.dtype, realized.uop.buffer.dtype), (lowered, lowered))
+    with patch.object(dtypes, "default_int", dtypes.int64):
+      self.assertEqual(Tensor(3).clone("CPU").realize().uop.buffer.dtype, dtypes.int64)
+
+  def test_integer_values(self):
+    x = Tensor.full((1,), 1, dtype=dtypes.int64, device="CPU")
+    self.assertEqual((x + 2**40).item(), 2**40 + 1)
+    self.assertEqual((x << 3).item(), 8)
+    self.assertTrue((x < 2**40).item())
+
+  def test_float64_precision(self):
+    value = 1.0 + 2**-40
+    x64 = Tensor.full((1,), 1.0, dtype=dtypes.float64, device="CPU")
+    self.assertEqual((x64 + value).item(), 2.0 + 2**-40)
+    x32 = Tensor.full((1,), 0.0, dtype=dtypes.float32, device="CPU")
+    self.assertEqual((x32 + value).item(), 1.0)
+
+  @unittest.expectedFailure  # TODO: exp/cos/sigmoid of a weak const stay weak instead of casting to a concrete float
+  def test_weak_transcendentals(self):
+    t_f16 = Tensor([1], dtype=dtypes.float16)
+    for out in (Tensor(2).exp(), Tensor(2).cos(), Tensor(2).sigmoid()):
+      self.assertEqual((out.dtype, (out + t_f16).dtype), (dtypes.weakfloat, dtypes.float16))
+
+  @unittest.expectedFailure  # TODO: where of weak consts stays weak and resolves per consumer
+  def test_where_and_shared_literal(self):
+    gate, weak = Tensor([True, False], device="CPU"), Tensor(2)
+    weak_where = gate.where(weak, 3)
+    self.assertEqual(weak_where.dtype, dtypes.weakint)
+    self.assertEqual((weak_where + Tensor([1, 1], dtype=dtypes.int64, device="CPU")).tolist(), [3, 4])
+    self.assertEqual((weak + Tensor([1], dtype=dtypes.int32, device="CPU")).item(), 3)
+    self.assertEqual((weak + Tensor([1], dtype=dtypes.int64, device="CPU")).item(), 3)
+
+  def test_null_lowering(self):
+    for t in (Tensor.full((1,), 1, dtype=dtypes.int64, device="NULL") + 2**40,
+              Tensor.full((1,), 1.0, dtype=dtypes.float64, device="NULL") + (1.0 + 2**-40)):
+      t.realize()
+      self.assertNotIn(t.uop.buffer.dtype, dtypes.weaks)
+
+
+class TestWeakStorageBoundary(unittest.TestCase):
+  # weak has no storage: a weak assignment source casts when it defers to the destination, everything else raises
+  def test_weak_source(self):
+    w3, w05 = Tensor.const(dtypes.weakint, 3).reshape(1).expand(2), Tensor.const(dtypes.weakfloat, 0.5).reshape(1)
+    dst = Tensor.zeros(2, dtype=dtypes.int8, device="CPU").contiguous().realize()
+    self.assertEqual(dst.assign(w3).realize().tolist(), [3, 3])                       # weakint defers to int8
+    with self.assertRaises(RuntimeError): dst.assign(w05.expand(2))                   # weakfloat into int does not defer
+    with self.assertRaises(RuntimeError): dst[0:1] = w05
+    fdst = Tensor.zeros(2, dtype=dtypes.float32, device="CPU").contiguous().realize()
+    fdst[0:1] = w05                                                                    # weakfloat defers to float
+    self.assertEqual(fdst.tolist(), [0.5, 0.0])
+    with tempfile.TemporaryDirectory() as td:                                          # the DISK path checks the same
+      ddst = Tensor.empty(2, dtype=dtypes.int32, device=f"DISK:{td}/t")
+      self.assertEqual(ddst.assign(w3).tolist(), [3, 3])
+      with self.assertRaises(RuntimeError): ddst.assign(w05.expand(2))
+
+  def test_weak_has_no_storage(self):
+    w = Tensor.const(dtypes.weakint, 3)
+    with self.assertRaises(RuntimeError): w.assign(Tensor([1], device="CPU"))
+    with self.assertRaises(RuntimeError): w.reshape(1)[0] = 1
+    with tempfile.TemporaryDirectory() as td:
+      with self.assertRaises(ValueError): safe_save({"x": w.reshape(1).expand(2)}, f"{td}/w.safetensors")
+    with self.assertRaises(RuntimeError): Tensor.empty(2, dtype=dtypes.weakint)
+    with self.assertRaises(RuntimeError): UOp.new_buffer("CPU", 2, dtypes.weakint)  # the one storage boundary
+    with self.assertRaises(RuntimeError): Tensor([1], dtype=dtypes.weakint)
+    import numpy as np
+    with self.assertRaises(RuntimeError): Tensor(np.ones(2, dtype=np.int32), dtype=dtypes.weakint)
+    self.assertEqual(Tensor(np.array(3), dtype=dtypes.weakint).dtype, dtypes.weakint)  # a 0-D ndarray is a const, not storage
+    with self.assertRaises(RuntimeError): Tensor(np.ones(2, dtype=np.float32), dtype=dtypes.weakfloat)
+    with self.assertRaises(RuntimeError): Tensor(bytes(8), dtype=dtypes.weakfloat)
+    with self.assertRaises(RuntimeError): Tensor(bytes(8), dtype=dtypes.weakint)
+    with tempfile.NamedTemporaryFile(suffix=".bin") as f:
+      f.write(bytes(8))
+      f.flush()
+      with self.assertRaises(RuntimeError): Tensor(pathlib.Path(f.name), dtype=dtypes.weakint)
+
+class TestWeakMaterializationEntries(unittest.TestCase):
+  # everything that creates storage from a weak value raises
+  def test_reads_commit_storage_raises(self):
+    for weak, value, strong in ((dtypes.weakint, 3, dtypes.default_int), (dtypes.weakfloat, 0.5, dtypes.default_float)):
+      def weak_val():
+        return Tensor([True], device="CPU").where(Tensor.const(weak, value), Tensor.const(weak, value))
+      self.assertEqual(weak_val().dtype, weak)
+      self.assertEqual(weak_val().to("CPU").dtype, weak)
+      self.assertEqual(weak_val().data().format, strong.fmt)
+      self.assertEqual(weak_val().numpy().dtype.itemsize, strong.itemsize)
+      self.assertEqual(weak_val().tolist(), [value])
+      self.assertEqual(weak_val().cast(strong).realize().uop.buffer.dtype, strong)
+      for entry in (lambda t: t.contiguous(), lambda t: t.realize(), lambda t: t.clone(),
+                    lambda t: t.to("CPU:1").realize(), lambda t: t.as_param(0)):
+        with self.assertRaises(RuntimeError): entry(weak_val())
+
+  def test_empty_reads_commit(self):
+    for weak, strong in ((dtypes.weakint, dtypes.default_int), (dtypes.weakfloat, dtypes.default_float)):
+      empty = Tensor.const(weak, 0).reshape(1).shrink(((0, 0),))
+      self.assertEqual(empty.data().format, strong.fmt)
+      self.assertEqual(empty.numpy().dtype.itemsize, strong.itemsize)
+      self.assertEqual(empty.tolist(), [])
+
+class TestWeakSpec(unittest.TestCase):
+  def test_weak_operand_allowed(self):
+    x = UOp.variable("x", 0, 10, dtypes.int64)
+    weak = UOp.const(dtypes.weakint, 3)
+    for u in (x.alu(Ops.ADD, weak), x.alu(Ops.CMPLT, weak), x.alu(Ops.SHL, weak)):
+      self.assertIs(spec_tensor.rewrite(u), True)
+    gate = UOp.variable("gate", False, True, dtypes.bool)
+    self.assertIs(spec_tensor.rewrite(UOp(Ops.WHERE, dtypes.int8, (gate, UOp.const(dtypes.int8, 1), weak))), True)
+
+
+if __name__ == "__main__":
+  unittest.main()
@@ -1,7 +1,7 @@
-import unittest
+import unittest, gc
 import numpy as np
-from tinygrad.helpers import polyN, is_numpy_ndarray
-from tinygrad.tensor import Tensor
+from tinygrad.helpers import polyN, disable_gc
+from tinygrad.tensor import Tensor, is_numpy_ndarray

 class TestPolyN(unittest.TestCase):
  def test_tensor(self):
@@ -11,5 +11,19 @@ class TestIsNumpyNdarray(unittest.TestCase):
  def test_tensor_numpy(self):
    self.assertTrue(is_numpy_ndarray(Tensor([1, 2, 3]).numpy()))

+class TestDisableGC(unittest.TestCase):
+  def test_recursive_decorator(self):
+    was_enabled = gc.isenabled()
+    @disable_gc()
+    def recurse(depth:int):
+      self.assertFalse(gc.isenabled())
+      if depth: recurse(depth-1)
+      self.assertFalse(gc.isenabled())
+    try:
+      recurse(2)
+      self.assertEqual(gc.isenabled(), was_enabled)
+    finally:
+      (gc.enable if was_enabled else gc.disable)()
+
 if __name__ == '__main__':
  unittest.main()
@@ -1,5 +1,6 @@
 import unittest
 from tinygrad import Tensor
+from tinygrad.device import Buffer
 from tinygrad.dtype import Invalid, dtypes
 from tinygrad.engine.realize import run_linear

@@ -9,7 +10,7 @@ class TestInvalidTensor(unittest.TestCase):
    buf = out.uop.buffer
    buf.allocate()
    sentinel = memoryview(bytearray(b'\x42' * buf.nbytes))
-    buf.copyin(sentinel)
+    buf.copy_from(Buffer("PYTHON", buf.size, buf.dtype, opaque=sentinel))
    before = buf.as_memoryview().cast(out.dtype.fmt).tolist()
    run_linear(linear, var_vals)
    ret = buf.as_memoryview().cast(out.dtype.fmt).tolist()
@@ -64,6 +65,17 @@ class TestInvalidTensor(unittest.TestCase):
    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid) > 1
    self._invalid_test_helper(out, [False, True, None, None])

+  def test_where_invalid_condition(self):
+    a, x = Tensor.arange(4), Tensor([0, 1, 2, 3])
+    bad = (a < 2).where(a, Invalid)
+    out = (bad < 1).logical_not().where(x + 10, x + 20)
+    self._invalid_test_helper(out, [20, 11, None, None])
+
+  def test_where_invalid_condition_bare(self):
+    cond = Tensor.full((4,), Invalid, dtype=dtypes.bool, buffer=False)
+    out = cond.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor([10.0, 20.0, 30.0, 40.0]))
+    self._invalid_test_helper(out, [None, None, None, None])
+
  def test_where_unary(self):
    mask = Tensor.arange(4) < 2
    out = mask.where(Tensor([1.0, 4.0, 9.0, 16.0]), Invalid).sqrt()
@@ -115,8 +127,6 @@ class TestInvalidTensor(unittest.TestCase):
    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.int, buffer=False)).bitcast(dtypes.int)
    self._invalid_test_helper(out, [0x3f800000, 0x40000000, None, None])

-  # tensor indexing uses reduce, so the entire result becomes invalid
-  @unittest.expectedFailure
  def test_tensor_index(self):
    idx = (Tensor.arange(4) < 2).where(Tensor([0, 1, 2, 3]), Invalid)
    out = Tensor([1.0, 2.0, 3.0, 4.0])[idx]
@@ -83,5 +83,10 @@ class TestTensorData(unittest.TestCase):
    assert dat.shape == (2,2)
    # NOTE: python can't deref float16

+  def test_tolist_empty_shapes(self):
+    for shape, expected in (((0,), []), ((2, 0), [[], []]), ((0, 2), []),
+                            ((2, 0, 3), [[], []]), ((2, 3, 0), [[[], [], []], [[], [], []]])):
+      self.assertEqual(Tensor.ones(*shape).tolist(), expected)
+
 if __name__ == '__main__':
  unittest.main()
@@ -73,11 +73,6 @@ def contiguous_mops_to_view(c:UOp, src:UOp):
  # no symbolic shape
  if not all_int(c.shape): return None

-  # check if view is supported
-  from tinygrad.device import Device
-  devs = (src.device,) if isinstance(src.device, str) else src.device
-  if not all(hasattr(Device[d].allocator, "_offset") for d in devs): return None
-
  if buf.op is not Ops.MULTI and (view := _make_buffer_view(src)) is not None:
    view = (view.replace(dtype=c.dtype, arg=c.numel()) if c.op is Ops.BITCAST else view).reshape(c.shape)
    return c.replace(src=(view,)) if c.op is Ops.COPY else view
@@ -1,19 +1,18 @@
 from dataclasses import replace, dataclass
 import itertools, functools
-from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, PROFILE, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC
+from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC
 from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, panic
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, GroupOp
-from tinygrad.uop.ops import TRACK_MATCH_STATS
 from tinygrad.uop.ops import AxisType
 from tinygrad.uop.render import pyrender
 from tinygrad.uop.spec import type_verify, spec_tensor, spec_program
 from tinygrad.renderer import Renderer, Estimates
 from tinygrad.renderer.isa import ISARenderer, IselContext, PreRegAllocContext
-from tinygrad.dtype import dtypes, AddrSpace, Invalid
+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_simplify_valid, 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
@@ -23,7 +22,6 @@ from tinygrad.codegen.opt.postrange import apply_opts
 from tinygrad.codegen.late.gater import pm_move_gates_from_index
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse
 from tinygrad.schedule.rangeify import pm_mops
-from tinygrad.schedule.indexing import apply_movement_op
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
 from tinygrad.codegen.late.regalloc import LinearScanRegallocContext, pm_regalloc_rewrite
 from tinygrad.codegen.late.coalese import memory_coalesing, pm_simplify_add_image
@@ -31,8 +29,6 @@ from tinygrad.helpers import all_same, flatten, argsort, partition
 from tinygrad.uop.ops import _align_left, _broadcast_shape, identity_element
 from tinygrad.schedule.rangeify import BufferizeOpts

-empty_matcher = PatternMatcher([])
-
 def do_number_param(ctx:list[int], x:UOp):
  if x.arg.slot != -1: return None
  ctx[0] += 1
@@ -83,9 +79,10 @@ def unroll_axis(ctx:dict[int, int], u:UOp, arg):
  return out.permute(argsort(permute_head+permute_tail))

 def expand_wmma(ctx:dict[int, int], u:UOp):
-  if u.tag != 1: return None
-  in0, in1, out0 = u.arg[6]
-  wmma = u.replace(src=(contract_axis(ctx, u.src[0], in0), contract_axis(ctx, u.src[1], in1), u.src[2]), tag=None)
+  if u.arg[4] is None: return None
+  in0, in1, out0 = u.arg[4]
+  wmma = u.replace(src=(contract_axis(ctx, u.src[0], in0), contract_axis(ctx, u.src[1], in1), u.src[2]),
+                   arg=(*u.arg[:4], None))
  return unroll_axis(ctx, wmma, out0)

 expander2 = PatternMatcher([
@@ -117,10 +114,6 @@ def broadcast_and_devec_wmma(b:UOp):
    src.append(b.replace(src=tuple([x.index(*idx_c) for x in src_reshaped])))
  return UOp.stack(*src).reshape(b.shape)

-@functools.cache
-def shape_indexes(shape:tuple[int, ...]) -> tuple[tuple[UOp, ...], ...]:
-  return tuple(tuple(UOp.const(dtypes.index, i) for i in idx) for idx in itertools.product(*map(range, shape)))
-
 pm_wmma_add = PatternMatcher([
  (UPat(Ops.WMMA, name="wmma") + UPat.var("add"),
   lambda add, wmma: UOp(wmma.op, src=(wmma.src[0], wmma.src[1], wmma.src[2]+add), arg=wmma.arg)),
@@ -136,48 +129,15 @@ unbroadcast = pm_wmma_add+PatternMatcher([
  (UPat(Ops.WMMA, name="b"), broadcast_and_devec_wmma),
 ])

-def do_devectorize(ctx, b:UOp):
-  ren = ctx.ren if isinstance(ctx, DevectorizeContext) else (ctx[-1] if isinstance(ctx, tuple) else ctx)
-  if b.op in GroupOp.Elementwise and b.dtype in dtypes.floats and ren.supports_float4: return None
-  if (shape:=b._shape) is None or shape == (): return None
+def do_devectorize(b:UOp):
+  if b.shape == (): return None
  # broadcasting needs to be already unpacked
-  if any(x._shape != shape for x in b.src): return None
+  if not all_same([x.shape for x in b.src]): return None
  src = []
-  for idx_c in shape_indexes(shape):
-    new_src = tuple(index_lane(ctx, x, idx_c) if isinstance(ctx, DevectorizeContext) else
-                    (x.src[idx_c[0].arg] if len(idx_c) == 1 and x.op is Ops.STACK else UOp(Ops.INDEX, x.dtype, (x,)+idx_c)) for x in b.src)
-    src.append(UOp(b.op, b.dtype, new_src, b.arg, b.tag))
-  return UOp(Ops.STACK, b.dtype, tuple(src)).reshape(shape) if b.op is not Ops.STORE else UOp.group(*src)
-
-@dataclass
-class DevectorizeContext:
-  ren: Renderer
-  lanes: dict[tuple[UOp, tuple[UOp, ...]], UOp]
-  @property
-  def rewrite_cache_key(self): return (type(self.ren), self.ren.target)
-
-def index_lane(ctx:DevectorizeContext, x:UOp, idxs:tuple[UOp, ...]) -> UOp:
-  key = (x, idxs)
-  if (ret:=ctx.lanes.get(key)) is not None: return ret
-  if x.op is Ops.STACK and idxs and idxs[0].op is Ops.CONST:
-    ret = index_lane(ctx, x.src[idxs[0].arg], idxs[1:]) if len(idxs) > 1 else x.src[idxs[0].arg]
-  elif x.op in GroupOp.Movement and len(idxs) == len(x.shape):
-    ret = index_lane(ctx, x.src[0], apply_movement_op(x.op, x.src[0].shape, x.marg, idxs))
-  elif x.op is Ops.INDEX:
-    ret = index_lane(ctx, x.src[0], x.src[1:]+idxs)
-  elif x.op in GroupOp.Elementwise:
-    ret = UOp(x.op, x.dtype, tuple(index_lane(ctx, s, idxs) if s._shape else s for s in x.src), x.arg, x.tag)
-  else: ret = UOp(Ops.INDEX, x.dtype, (x,)+idxs)
-  ctx.lanes[key] = ret
-  return ret
-
-def index_elementwise(x:UOp, idx:UOp):
-  indexes = idx.src[1:]
-  return UOp(x.op, x.dtype, tuple(UOp(Ops.INDEX, s.dtype, (s,)+indexes) if s._shape else s for s in x.src), x.arg, x.tag)
-
-def index_elementwise_lane(ctx:DevectorizeContext, x:UOp, idx:UOp):
-  indexes = idx.src[1:]
-  return UOp(x.op, x.dtype, tuple(index_lane(ctx, s, indexes) if s._shape else s for s in x.src), x.arg, x.tag)
+  for idx in itertools.product(*[range(x) for x in b.shape]):
+    idx_c = [UOp.const(dtypes.index, i) for i in idx]
+    src.append(b.replace(src=tuple([x.index(*idx_c) for x in b.src])))
+  return UOp.stack(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src)

 def do_stack_wmma(u:UOp):
  if all(x.op in (Ops.STACK, Ops.WMMA) for x in u.src): return None
@@ -193,13 +153,11 @@ def do_stack_wmma(u:UOp):
 ew_devectorizer = PatternMatcher([
  # unpack broadcasting
  (UPat(GroupOp.Elementwise, name="b"), do_devectorize),
-  (UPat(GroupOp.Elementwise, name="x").f(Ops.INDEX, allow_any_len=True, name="idx"), index_elementwise),
 ])

 devectorizer2 = mop_cleanup+pm_mops+PatternMatcher([
  # unpack broadcasting
  (UPat(GroupOp.Elementwise|{Ops.LOAD,Ops.STORE}, name="b"), do_devectorize),
-  (UPat(GroupOp.Elementwise, name="x").f(Ops.INDEX, allow_any_len=True, name="idx"), index_elementwise_lane),
  # INDEX without src is nothing (TODO: this should be in mop_cleanup)
  (UPat(Ops.INDEX, src=(UPat.var('x'),)), lambda x: x),
  # unpack WMMA
@@ -303,6 +261,13 @@ pm_add_local_buffers = PatternMatcher([
  (UPat(Ops.STAGE, name="x"), add_local_buffer),
 ])+pm_mops

+# float ALUs need a float operand
+# make that cast explicit before the decomps, which expand SIN/LOG2/EXP2 into float polynomials and assert a float operand
+pm_cast_float_alu = PatternMatcher([
+  (UPat((Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL), src=(UPat(name="x"),), name="u"),
+   lambda u,x: u.replace(src=(x.cast(u.dtype),)) if x.dtype != u.dtype else None),
+])
+
 def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
  if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
  if DEBUG >= 5: print(pyrender(ast))
@@ -341,54 +306,55 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
  sink = graph_rewrite(sink, pm_add_local_buffers, ctx=itertools.count(0), name="add local buffers")

  # add gpu dims (late). this works after devectorize, but it's faster here
-  if VIZ or PROFILE or TRACK_MATCH_STATS: sink = graph_rewrite(sink, pm_add_gpudims, ctx=ren, name="add gpudims")
-  elif (gpu_sink:=pm_add_gpudims.rewrite(sink, ren)) is not None: sink = gpu_sink
+  sink = graph_rewrite(sink, pm_add_gpudims, ctx=ren, name="add gpudims")

  # **** optimizations are done, now we lower to actual code ****

  sink = graph_rewrite(sink, symbolic_simple+unbroadcast+pm_add_loads, name="*** unbroadcast / add loads")

  # devectorize
-  sink = graph_rewrite(sink, symbolic_simple+devectorizer2, ctx=DevectorizeContext(ren, {}), name="devectorize2")
+  sink = graph_rewrite(sink, symbolic_simple+devectorizer2, ctx=ren, name="devectorize2")
+
+  # simplify indexing
+  sink = graph_rewrite(sink, indexing_simplify, name="simplify load/store indexing")

  # some coalesing misses without this
  sink = graph_rewrite(sink, sym, name="early symbolic")

  # do memory coalesing (late)
  sink = memory_coalesing(sink, ren)
-  if IMAGE: sink = graph_rewrite(sink, symbolic_simple+ew_devectorizer+pm_simplify_add_image, name="add images", ctx=({}, ren), bottom_up=True)
+  sink = graph_rewrite(sink, symbolic_simple+ew_devectorizer+pm_simplify_add_image, name="add images", ctx=({}, ren), bottom_up=True)

-  has_invalid = (sink.op is Ops.CONST and sink.arg is Invalid) or any(u.op is Ops.CONST and u.arg is Invalid for u in sink.backward_slice)
-  if has_invalid:
-    sink = graph_rewrite(sink, pm_simplify_valid, name="simplify valid after coalescing")
-    has_invalid = (sink.op is Ops.CONST and sink.arg is Invalid) or any(u.op is Ops.CONST and u.arg is Invalid for u in sink.backward_slice)
+  # extra symbolic before decomp. crashes without this?
+  sink = graph_rewrite(sink, sym, name="extra symbolic")

  # lower index dtype
  # NOTE: we need indexing_simplify to remove the cast to long using the Invalid
  sink = graph_rewrite(sink, pm_lower_index_dtype+indexing_simplify, name="lower all index dtypes")

+  # final symbolic before decomp
+  sink = graph_rewrite(sink, symbolic, name="final symbolic")
+
+  sink = graph_rewrite(sink, pm_cast_float_alu, name="cast float alu operands")
+
  # **** decomps ****

-  # final symbolic + floordiv/mod + dtype decomp
+  # floordiv+mod / dtype decomp (early)
  supported_ops = tuple(ren.code_for_op.keys())
-  pm_decomp = symbolic+get_simplifying_rewrite_patterns(supported_ops)
+  pm_decomp = symbolic_simple+get_simplifying_rewrite_patterns(supported_ops)
+  sink = graph_rewrite(sink, pm_decomp, name="early decompositions")

  # late decomps + move gates from unrenderable INVALID where
-  candidate_dtypes = {*dtypes.fp8s, dtypes.bfloat16, dtypes.half, dtypes.long, dtypes.ulong}
-  emulated_dtypes = set(EMULATED_DTYPES.tolist(dtypes)) | (candidate_dtypes - ren.supported_dtypes())
-  needs_dtype_decomp = sink.dtype in emulated_dtypes or any(u.dtype in emulated_dtypes for u in sink.backward_slice)
-  if needs_dtype_decomp:
-    sink = graph_rewrite(sink, pm_decomp, name="early decompositions")
-    sink = graph_rewrite(sink, pm_dtype_decomps, ctx=(set(), ren), name="decomp dtypes")
+  sink = graph_rewrite(sink, pm_dtype_decomps, ctx=(set(), ren), name="decomp dtypes")
  pm_decomp = pm_decomp+\
    get_late_rewrite_patterns(supported_ops, bool(DISABLE_FAST_IDIV))+\
    get_transcendental_patterns(supported_ops, TRANSCENDENTAL>=2)
  sink = graph_rewrite(sink, pm_decomp, ctx=ren, name="late decompositions")
-  if has_invalid: sink = graph_rewrite(sink, pm_move_gates_from_index, name="move gates from index")
+  sink = graph_rewrite(sink, pm_move_gates_from_index, name="move gates from index")

  # final rules for the renderer (without sym)
-  extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else empty_matcher
-  pm_final_rewrite = symbolic_simple+extra_matcher+pm_split_ends+pm_no_index
+  extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([])
+  pm_final_rewrite = pm_decomp+extra_matcher+pm_split_ends+pm_no_index
  sink = graph_rewrite(sink, pm_final_rewrite+pm_remove_invalid, ctx=ren, name="final rewrite")

  # this was the linearizer
@@ -469,7 +435,7 @@ pm_to_program = PatternMatcher([

@track_rewrites(name=lambda ast,renderer,ret,**kwargs: TracingKey(ret.src[0].arg.name,(ret.src[0].arg.function_name, ast), ret=renderer), replay=True)
@Context(ALLOW_DEVICE_USAGE=0)
-def do_to_program(ast:UOp, renderer:Renderer, compile_binary=True) -> UOp:
+def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
  """
  Transform an AST into a compiled PROGRAM. May trigger BEAM search.

@@ -480,50 +446,25 @@ def do_to_program(ast:UOp, renderer:Renderer, compile_binary=True) -> UOp:
  Returns:
    The Ops.PROGRAM with SINK/LINEAR/SOURCE/BINARY.
  """
-  from tinygrad.codegen.opt.gemm import cooperative_gemm_program, direct_conv_bwd_activation_program
-  from tinygrad.codegen.opt.reduce import activation_var_grad_program, bn_grad_512_program, channel_reduce_program, col2im_program
-  from tinygrad.codegen.opt.reduce import im2col_program, moments_512_program
-  from tinygrad.codegen.opt.reduce import maxpool_backward_program, maxpool_program
-  if ast.op is Ops.SINK and (prg:=direct_conv_bwd_activation_program(ast, renderer, compile_binary)) is not None: return prg
-  if ast.op is Ops.SINK and (prg:=cooperative_gemm_program(ast, renderer, compile_binary)) is not None: return prg
-  if ast.op is Ops.SINK and (prg:=activation_var_grad_program(ast, renderer, compile_binary)) is not None: return prg
-  if ast.op is Ops.SINK and (prg:=moments_512_program(ast, renderer, compile_binary)) is not None: return prg
-  if ast.op is Ops.SINK and (prg:=bn_grad_512_program(ast, renderer, compile_binary)) is not None: return prg
-  if ast.op is Ops.SINK and (prg:=channel_reduce_program(ast, renderer, compile_binary)) is not None: return prg
-  if ast.op is Ops.SINK and (prg:=col2im_program(ast, renderer, compile_binary)) is not None: return prg
-  if ast.op is Ops.SINK and (prg:=im2col_program(ast, renderer, compile_binary)) is not None: return prg
-  if ast.op is Ops.SINK and (prg:=maxpool_backward_program(ast, renderer, compile_binary)) is not None: return prg
-  if ast.op is Ops.SINK and (prg:=maxpool_program(ast, renderer, compile_binary)) is not None: return prg
  if ast.op is Ops.PROGRAM: prg = ast
  elif ast.op is Ops.SINK:
    assert isinstance(ast.arg, KernelInfo), "requires KernelInfo on arg to to_program"
    full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None)
+    prog_info = ProgramInfo.from_sink(full_sink)
    # instruction selection
    if isinstance(renderer, ISARenderer):
      full_sink = graph_rewrite(full_sink, renderer.pre_isel_matcher, ctx=itertools.count(-1, -1), name="pre instruction selection", bottom_up=True)
      full_sink = graph_rewrite(full_sink, renderer.isel_matcher, ctx=IselContext(full_sink), name="instruction selection", bottom_up=True)
-    prg = UOp(Ops.PROGRAM, src=(full_sink,))
+    prg = UOp(Ops.PROGRAM, src=(full_sink,), arg=prog_info)
  else: raise RuntimeError(f"can't call to_program on {ast.op}")
-  if VIZ:
-    if not isinstance(prg.arg, ProgramInfo): prg = prg.replace(arg=ProgramInfo.from_sink(prg.src[0]))
-    prg = graph_rewrite(prg, pm_to_program, ctx=renderer, name="linearize/render")
-    graph_rewrite(prg, PatternMatcher([]), name="View Program")
-    return prg
-  # PROGRAM lowering is a linear root-only pipeline. Driving it through graph_rewrite
-  # needlessly walks the full SINK and LINEAR graphs between each stage.
-  if len(prg.src) == 1: prg = do_linearize(renderer, prg, prg.src[0])
  if not isinstance(prg.arg, ProgramInfo): prg = prg.replace(arg=ProgramInfo.from_sink(prg.src[0]))
-  if prg.src[0].arg.estimates is None and (estimated:=do_estimates(prg, prg.src[0], prg.src[1])) is not None: prg = estimated
-  if len(prg.src) == 2:
-    prg = do_assemble(renderer, prg, prg.src[1]) if isinstance(renderer, ISARenderer) else do_render(renderer, prg, prg.src[1])
-  if compile_binary and len(prg.src) == 3 and (compiled:=do_compile(renderer, prg, prg.src[2])) is not None: prg = compiled
+  prg = graph_rewrite(prg, pm_to_program, ctx=renderer, name="linearize/render")
+  if VIZ: graph_rewrite(prg, PatternMatcher([]), name="View Program")
  return prg

 to_program_cache: dict[tuple, UOp] = {}
-def to_program(ast:UOp, renderer:Renderer, compile_binary=True) -> UOp:
+def to_program(ast:UOp, renderer:Renderer) -> UOp:
  config = (NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32)
-  # UOps are structurally interned, so identity is already a collision-free structural
-  # cache key within this process and avoids recursively hashing every kernel graph.
-  key = (ast, type(renderer), renderer.target, compile_binary, *[x.value for x in config])
-  if (prg:=to_program_cache.get(key)) is None: to_program_cache[key] = prg = do_to_program(ast, renderer, compile_binary=compile_binary)
+  key = (ast.key, type(renderer), renderer.target, *[x.value for x in config])
+  if (prg:=to_program_cache.get(key)) is None: to_program_cache[key] = prg = do_to_program(ast, renderer)
  return prg
@@ -33,11 +33,14 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
    case Ops.CAST: return a0.bitcast(dtypes.uint).cast(dt)
    case Ops.BITCAST: return a0.bitcast(dt), a1.bitcast(dt)
    case Ops.SHL:
-      lo, hi = shl(a0, b0_mod:=b0 & 31), shl(a1, b0_mod) | shr(shr(a0, 1), 31 - b0_mod)
+      a0u, a1u, n = a0.bitcast(dtypes.uint), a1.bitcast(dtypes.uint), (b0 & 31).cast(dtypes.uint)
+      lo, hi = (a0u << n).bitcast(dt), ((a1u << n) | ((a0u >> 1) >> (31 - n))).bitcast(dt)
      return (b0 >= 32).where(zero, lo), (b0 >= 32).where(lo, hi)
    case Ops.SHR:
-      lo, hi = shr(a0, b0_mod:=b0 & 31) | shl(shl(a1, 1), 31 - b0_mod), shr(a1, b0_mod)
-      return (b0 >= 32).where(hi, lo), (b0 >= 32).where(zero, hi)
+      a0u, a1u, n = a0.bitcast(dtypes.uint), a1.bitcast(dtypes.uint), (b0 & 31).cast(dtypes.uint)
+      lo, hi = ((a0u >> n) | ((a1u << 1) << (31 - n))).bitcast(dt), a1 >> (b0 & 31)
+      fill = a1 >> 31 if dt == dtypes.int else zero  # vacated high word: sign bits when signed, else 0
+      return (b0 >= 32).where(hi, lo), (b0 >= 32).where(fill, hi)
    case Ops.ADD: return (low:=a0+b0), (a1 + b1).replace(dtype=dt) + (low.bitcast(dtypes.uint) < a0.bitcast(dtypes.uint)).cast(dt)
    case Ops.SUB: return a0 - b0, a1 - b1 - (a0.bitcast(dtypes.uint) < b0.bitcast(dtypes.uint)).cast(dt)
    case Ops.MUL:
@@ -43,22 +43,19 @@ def fast_idiv(ren: Renderer, x: UOp, d: int, dont_cast=False) -> UOp|None:

 # ***** threefry *****

-def threefry2x32(x: UOp, key: UOp, ctx:Renderer|None=None):
-  native_shifts = ctx is not None and Ops.SHL in ctx.code_for_op and Ops.SHR in ctx.code_for_op
-  shl = (lambda v, s: v << s) if native_shifts else (lambda v, s: v * 2**s)
-  shr = (lambda v, s: v >> s) if native_shifts else (lambda v, s: v // 2**s)
+def threefry2x32(x: UOp, key: UOp):
  # split x and key from uint64 to two uint32
-  x0, x1 = x.cast(dtypes.uint32), shr(x, 32).cast(dtypes.uint32)
-  key0, key1 = key.cast(dtypes.uint32), shr(key, 32).cast(dtypes.uint32)
+  x0, x1 = (x & 0xffffffff).cast(dtypes.uint32), ((x // 2**32) & 0xffffffff).cast(dtypes.uint32)
+  key0, key1 = (key & 0xffffffff).cast(dtypes.uint32), ((key // 2**32) & 0xffffffff).cast(dtypes.uint32)

  rotations = [[13, 15, 26, 6], [17, 29, 16, 24]]
  ks = [key1, key0 ^ key1 ^ 0x1BD11BDA, key0]
  xr:list[UOp] = [x0 + ks[-1], x1 + ks[0]]
  for i in range(5):
-    for r in rotations[i % 2]: xr[0], xr[1] = (x0 := xr[0] + xr[1]), x0 ^ (shl(xr[1], r) + shr(xr[1], 32-r))
+    for r in rotations[i % 2]: xr[0], xr[1] = (x0 := xr[0] + xr[1]), x0 ^ ((xr[1] * 2**r) + (xr[1] // 2**(32 - r)))
    xr = [(xr[0] + ks[i % 3]), (xr[1] + ks[(i + 1) % 3] + i + 1)]

-  return shl(xr[1].cast(dtypes.uint64), 32) | xr[0].cast(dtypes.uint64)
+  return xr[1].cast(dtypes.uint64) * 2**32 | xr[0].cast(dtypes.uint64)

 # ***** decomposition patterns *****

@@ -42,7 +42,7 @@ def simplify_valid_image_load(buf:UOp, idx_y:UOp, idx_x:UOp, valid:UOp) -> UOp|N
  if not is_image_shape(buf._shape): return None
  if idx_x.dtype != idx_y.dtype: idx_x, idx_y = idx_x.cast(dtypes.int), idx_y.cast(dtypes.int)
  start_idx = idx_x.stack(idx_y)
-  idx = uop_given_valid(valid, start_idx, try_simplex=True)
+  idx = uop_given_valid(valid, start_idx)
  drop_stmt = _drop_valid_stmts(valid, idx, buf._shape[0], buf._shape[1])

  if not drop_stmt and idx is start_idx: return None
@@ -74,7 +74,7 @@ def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
  # search for dims that drop the most valid statements
  best_drop, cands = -1, []
  for ch, cw in [shapes[buf.arg.slot]] if buf.arg.slot in shapes else image_valid_dims(buf.dtype, buf.max_numel(), ren.target.arch):
-    cidx = uop_given_valid(valid, ((x//4)%cw).stack(x//(4*cw)), try_simplex=True)
+    cidx = uop_given_valid(valid, ((x//4)%cw).stack(x//(4*cw)))
    dropped = len(_drop_valid_stmts(valid, cidx, ch, cw))
    if dropped > best_drop: best_drop, cands = dropped, [(ch, cw, cidx)]
    elif dropped == best_drop: cands.append((ch, cw, cidx))
@@ -1,461 +0,0 @@
-from dataclasses import dataclass, replace
-from math import prod
-from tinygrad.dtype import dtypes
-from tinygrad.renderer import Estimates, Renderer
-from tinygrad.uop.ops import AxisType, Ops, ProgramInfo, UOp, ssimplify
-
-WMMA_M = WMMA_N = WMMA_K = 16
-BLOCK_M = BLOCK_N = 128
-BLOCK_K = 32
-THREADS = 128
-
-@dataclass(frozen=True)
-class GemmMatch:
-  c:UOp
-  a:UOp
-  b:UOp
-  m:int
-  n:int
-  k:int
-  old:UOp|None = None
-  scale:float = 0.0
-  a_kxm:bool = False
-  b_kxn:bool = False
-
-@dataclass(frozen=True)
-class BatchedGemmMatch:
-  c:UOp
-  a:UOp
-  b:UOp
-  m:int
-  n:int
-  k:int
-  batch:int
-
-@dataclass(frozen=True)
-class DirectConvBwdActivationMatch:
-  m:int
-  cin:int
-  cout:int
-  spatial:int
-  residual:bool = False
-
-@dataclass(frozen=True)
-class GemmOutputNCHW:
-  spatial:int
-
-def _match_gemm(ast:UOp, device:str, arch:str) -> GemmMatch|None:
-  if device != "AMD" or not arch.startswith("gfx11") or len(ast.src) != 1: return None
-  end = ast.src[0]
-  if end.op is not Ops.END or len(end.src) < 3 or end.src[0].op is not Ops.STORE: return None
-  store, ranges = end.src[0], end.src[1:]
-  if any(x.op is not Ops.RANGE or x.arg[1] is not AxisType.LOOP for x in ranges) or store.src[0].op is not Ops.INDEX: return None
-  value = store.src[1]
-  old, old_index, scale = None, None, 0.0
-  if value.op is Ops.CAST and value.dtype == dtypes.half and value.src[0].op is Ops.REDUCE: reduce = value.src[0]
-  elif value.op is Ops.CAST and value.dtype == dtypes.float and value.src[0].op is Ops.ADD:
-    add_lhs, add_rhs = value.src[0].src
-    if add_lhs.op is not Ops.CAST or add_lhs.dtype != dtypes.half or add_lhs.src[0].op is not Ops.REDUCE:
-      add_lhs, add_rhs = add_rhs, add_lhs
-    if add_lhs.op is not Ops.CAST or add_lhs.dtype != dtypes.half or add_lhs.src[0].op is not Ops.REDUCE: return None
-    if add_rhs.op is not Ops.MUL: return None
-    old_idx = next((x for x in add_rhs.src if x.op is Ops.INDEX and x.dtype == dtypes.half), None)
-    scale_uop = next((x for x in add_rhs.src if x.op is Ops.CONST and x.dtype == dtypes.half), None)
-    if old_idx is None or scale_uop is None: return None
-    reduce, old, old_index, scale = add_lhs.src[0], old_idx.src[0], old_idx.src[1], float(scale_uop.arg)
-  else: return None
-  if reduce.arg != (Ops.ADD, 0) or len(reduce.src) != 2 or reduce.src[1].op is not Ops.RANGE: return None
-  k, product = reduce.src[1], reduce.src[0]
-  if product.op is not Ops.CAST or product.dtype != dtypes.float or product.src[0].op is not Ops.MUL: return None
-  lhs, rhs = product.src[0].src
-  if lhs.op is not Ops.INDEX or rhs.op is not Ops.INDEX: return None
-  if any(x.dtype != dtypes.half for x in (lhs.src[0], rhs.src[0])): return None
-  if store.src[0].src[0].dtype not in (dtypes.half, dtypes.float): return None
-
-  out_idx = ssimplify(store.src[0].src[1].get_idx())
-  if old_index is not None and ssimplify(old_index.get_idx()) is not out_idx: return None
-  k_size, total_size = int(k.vmax)+1, prod(int(x.vmax)+1 for x in ranges)
-  lhs_idx, rhs_idx = ssimplify(lhs.src[1].get_idx()), ssimplify(rhs.src[1].get_idx())
-  for m in ranges:
-    m_size, n_size = int(m.vmax)+1, total_size//(int(m.vmax)+1)
-    n = ssimplify(out_idx%n_size)
-    if ssimplify(out_idx//n_size) is not m or int(n.vmin) != 0 or int(n.vmax) != n_size-1: continue
-    n64 = old is None and n_size == 64
-    a_mxk_idx, a_kxm_idx = ssimplify(m*k_size+k), ssimplify(k*m_size+m)
-    b_nxk_idx, b_kxn_idx = ssimplify(n*k_size+k), ssimplify(k*n_size+n)
-    for a_idx, a_buf, b_idx, b_buf in ((lhs_idx, lhs.src[0], rhs_idx, rhs.src[0]), (rhs_idx, rhs.src[0], lhs_idx, lhs.src[0])):
-      if not (a_idx is a_mxk_idx or a_idx is a_kxm_idx): continue
-      if not (b_idx is b_nxk_idx or b_idx is b_kxn_idx): continue
-      if n64 and (a_idx is not a_mxk_idx or b_idx is not b_nxk_idx): continue
-      g = GemmMatch(store.src[0].src[0], a_buf, b_buf, m_size, n_size, k_size, old, scale,
-                    a_idx is a_kxm_idx, b_idx is b_kxn_idx)
-      bm, bn = _gemm_block_m(g), _gemm_block_n(g)
-      if m_size % bm or n_size % bn or k_size % BLOCK_K: continue
-      if (m_size//bm)*(n_size//bn) < (32 if old is not None else 512): continue
-      return g
-  return None
-
-def _match_batched_gemm(ast:UOp, device:str, arch:str) -> BatchedGemmMatch|None:
-  if device != "AMD" or not arch.startswith("gfx11") or len(ast.src) != 1: return None
-  end = ast.src[0]
-  if end.op is not Ops.END or len(end.src) != 4 or end.src[0].op is not Ops.STORE: return None
-  store, ranges = end.src[0], end.src[1:]
-  if any(x.op is not Ops.RANGE or x.arg[1] is not AxisType.LOOP for x in ranges) or store.src[0].op is not Ops.INDEX: return None
-  reduce = store.src[1]
-  if reduce.op is not Ops.REDUCE or reduce.dtype != dtypes.float or reduce.arg != (Ops.ADD, 0) or len(reduce.src) != 2: return None
-  k, product = reduce.src[1], reduce.src[0]
-  if k.op is not Ops.RANGE or product.op is not Ops.CAST or product.dtype != dtypes.float or product.src[0].op is not Ops.MUL: return None
-  lhs, rhs = product.src[0].src
-  if lhs.op is not Ops.INDEX or rhs.op is not Ops.INDEX or lhs.dtype != dtypes.half or rhs.dtype != dtypes.half: return None
-  out_idx, lhs_idx, rhs_idx = (ssimplify(x.src[1].get_idx()) for x in (store.src[0], lhs, rhs))
-  k_size = int(k.vmax)+1
-  for m in ranges:
-    for n in ranges:
-      if n is m: continue
-      batch = next(x for x in ranges if x is not m and x is not n)
-      m_size, n_size, batch_size = int(m.vmax)+1, int(n.vmax)+1, int(batch.vmax)+1
-      if ssimplify(out_idx//(n_size*batch_size)) is not m or ssimplify((out_idx//batch_size)%n_size) is not n or \
-         ssimplify(out_idx%batch_size) is not batch: continue
-      if m_size % 64 or n_size % 32 or k_size % BLOCK_K: continue
-      a_idx, b_idx = ssimplify((batch*k_size+k)*m_size+m), ssimplify((batch*k_size+k)*n_size+n)
-      if lhs_idx is a_idx and rhs_idx is b_idx: a, b = lhs.src[0], rhs.src[0]
-      elif rhs_idx is a_idx and lhs_idx is b_idx: a, b = rhs.src[0], lhs.src[0]
-      else: continue
-      return BatchedGemmMatch(store.src[0].src[0], a, b, m_size, n_size, k_size, batch_size)
-  return None
-
-def _gemm_block_m(g:GemmMatch) -> int:
-  if g.old is None and (g.m,g.n,g.k) in ((393216,576,256), (24576,512,4608)): return 64
-  return 32 if g.old is not None and (g.m < 512 or (g.a_kxm and g.b_kxn and g.k >= 65536)) else \
-    64 if g.old is not None else BLOCK_M
-def _gemm_block_n(g:GemmMatch) -> int:
-  if g.old is not None and g.a_kxm and g.b_kxn and (g.m,g.n,g.k) == (256,2304,65536): return 128
-  if g.old is not None and g.a_kxm and g.b_kxn and (g.m,g.n,g.k) == (256,2304,98304): return 64
-  if g.n == 288 and g.old is None and not g.a_kxm and g.b_kxn: return 96
-  if g.n == 576 and g.old is None and not g.a_kxm and g.b_kxn: return 192
-  if g.old is not None and g.a_kxm and g.b_kxn and g.m == 256 and g.k >= 65536: return BLOCK_K
-  return 64 if g.n == 64 and g.old is None and not g.a_kxm and not g.b_kxn else BLOCK_N
-def _gemm_block_k(g:GemmMatch) -> int:
-  if g.old is not None and g.a_kxm and g.b_kxn and g.m == 256 and g.k >= 65536: return 64
-  if g.old is not None and (g.m < 512 or (g.a_kxm and g.b_kxn and g.k >= 65536)) and g.k % 64 == 0: return 64
-  return 144 if _gemm_block_n(g) == 64 and g.k in (288, 576) else BLOCK_K
-def _batched_block_k(g:BatchedGemmMatch) -> int:
-  if (g.batch, g.m, g.n, g.k) == (192, 64, 288, 8192): return 16
-  if (g.batch, g.m, g.n, g.k) == (96, 64, 576, 4096): return BLOCK_K
-  return 128 if g.m == 64 and g.batch >= 96 and g.k % 128 == 0 else BLOCK_K
-def _batched_block_n(g:BatchedGemmMatch) -> int:
-  if (g.batch, g.m, g.n, g.k) == (192, 64, 288, 8192): return 288
-  return 192 if g.n == 576 else BLOCK_N
-def _batched_threads(g:BatchedGemmMatch) -> int:
-  if _batched_block_n(g) == 288: return 192
-  return _batched_block_n(g)
-def _batched_family_name(g:BatchedGemmMatch) -> str:
-  return f"coop_bgemm_bm64_bn{_batched_block_n(g)}_bk{_batched_block_k(g)}_t{_batched_threads(g)}" \
-    f"{'_partial_n' if g.n % _batched_block_n(g) else ''}"
-def _gemm_threads(g:GemmMatch) -> int:
-  if _gemm_block_n(g) == 192: return 192
-  return THREADS
-
-def _gemm_family_name(g:GemmMatch, output_nchw:GemmOutputNCHW|None) -> str:
-  return f"coop_gemm_bm{_gemm_block_m(g)}_bn{_gemm_block_n(g)}_bk{_gemm_block_k(g)}_t{_gemm_threads(g)}" \
-    f"{'_kxm' if g.a_kxm else ''}{'_kxn' if g.b_kxn else ''}{'_acc' if g.old is not None else ''}" \
-    f"{'_nchw' if output_nchw is not None else ''}"
-
-def _render_gemm(g:GemmMatch, name:str, output_nchw:GemmOutputNCHW|None=None) -> str:
-  bm, bn_size, bk, threads = _gemm_block_m(g), _gemm_block_n(g), _gemm_block_k(g), _gemm_threads(g)
-  pack_b = g.b_kxn and not g.a_kxm
-  as_stride, bs_stride = (bm if g.a_kxm else bk+8), (bn_size if g.b_kxn else bk+8)
-  waves_m, waves_n = (2, 3) if threads == 192 else (1, 2) if threads == 64 else \
-                     (((1, 8) if bm == 32 else (2, 4) if bm == 64 else (4, 2)) if threads == 256 else
-                      ((1, 4) if bm <= 64 else (2, 2)))
-  tiles_m, tiles_n = bm//(waves_m*WMMA_M), bn_size//(waves_n*WMMA_N)
-  cslot, aslot, bslot = g.c.arg.slot, g.a.arg.slot, g.b.arg.slot
-  buffers = (g.c, g.a, g.b) + ((g.old,) if g.old is not None else ())
-  params = ', '.join(f'{"float" if x.dtype == dtypes.float else "half"}* p{x.arg.slot}' for x in sorted(buffers, key=lambda x:x.arg.slot))
-  params += ', int M, int N, int K' + (', int S' if output_nchw is not None else '')
-  lines = [
-    '#define half _Float16',
-    'typedef half half16 __attribute__((ext_vector_type(16)));',
-    'typedef float float8 __attribute__((ext_vector_type(8)));',
-    'typedef unsigned uint8 __attribute__((ext_vector_type(8)));',
-    'typedef unsigned uint16 __attribute__((ext_vector_type(16)));',
-    'typedef unsigned short ushort16 __attribute__((ext_vector_type(16)));',
-    '#define HALF_BITS(x) (unsigned short)(x), (unsigned short)((x)>>16)',
-    '#define WMMA __builtin_amdgcn_wmma_f32_16x16x16_f16_w32',
-    f'extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size({threads}, {threads}))) {name}({params}) {{',
-    f'  __attribute__((shared, aligned(32))) half As[{bk*bm if g.a_kxm else bm*as_stride}], '
-    f'Bs[{bk*bn_size if g.b_kxn else bn_size*bs_stride}];',
-    '  int bn=__builtin_amdgcn_workgroup_id_x(), bm=__builtin_amdgcn_workgroup_id_y(), tid=__builtin_amdgcn_workitem_id_x();',
-    f'  int wave=tid>>5, lane=tid&31, wm=wave/{waves_n}, wn=wave%{waves_n}, row=lane&15, halfrow=lane>>4;',
-  ]
-  lines.append(f'  float8 c[{tiles_m}][{tiles_n}]={{}};')
-  lines += [f'  for (int kt=0; kt<K/{bk}; kt++) {{']
-  if g.a_kxm:
-    aseg_count = bm*bk//16
-    lines += ['    #pragma unroll', f'    for (int q=0; q<{(aseg_count+threads-1)//threads}; q++) {{',
-              f'      int aseg=tid+q*{threads};']
-    if aseg_count % threads: lines.append(f'      if (aseg<{aseg_count}) {{')
-    lines += [f'      int ak=aseg/{bm//16}, am=(aseg%{bm//16})*16;',
-              f'      *((half16*)(As+ak*{bm}+am))=*((half16*)(p{aslot}+((long)(kt*{bk}+ak)*M)+bm*{bm}+am));']
-    if aseg_count % threads: lines.append('      }')
-    lines.append('    }')
-  else:
-    prefix, suffix = (f'    if (tid<{bm}) {{ ', ' }') if threads != bm else ('    ', '')
-    lines += [f'{prefix}long ao=((long)(bm*{bm}+tid)*K)+kt*{bk};']
-    lines += ['      #pragma unroll', f'      for (int q=0; q<{bk//16}; q++) '
-              f'*((half16*)(As+tid*{as_stride}+q*16))=*((half16*)(p{aslot}+ao+q*16));']
-    lines[-1] += suffix
-  if pack_b:
-    bsegs = bn_size//16
-    if threads != bsegs*16: lines.append(f'    if (tid<{bsegs*16}) {{')
-    lines += ['    #pragma unroll', f'    for (int q=0; q<{bk//32}; q++) {{',
-              f'      int bkp=tid/{bsegs}+q*16, bn0=(tid%{bsegs})*16;',
-              f'      half16 bv0=*((half16*)(p{bslot}+((long)(kt*{bk}+bkp*2)*N)+bn*{bn_size}+bn0));',
-              f'      half16 bv1=*((half16*)(p{bslot}+((long)(kt*{bk}+bkp*2+1)*N)+bn*{bn_size}+bn0));',
-              '      ushort16 pb0=__builtin_bit_cast(ushort16,bv0), pb1=__builtin_bit_cast(ushort16,bv1);',
-              f'      *((uint16*)(((unsigned*)Bs)+bkp*{bn_size}+bn0))=__builtin_convertvector(pb0,uint16)|'
-              '(__builtin_convertvector(pb1,uint16)<<16);', '    }']
-    if threads != bsegs*16: lines.append('    }')
-  elif g.b_kxn:
-    bsegs = bn_size//16
-    lines += ['    #pragma unroll', f'    for (int q=0; q<{bn_size*bk//(threads*16)}; q++) {{',
-              f'      int bseg=tid+q*{threads}, bki=bseg/{bsegs}, bni=(bseg%{bsegs})*16;',
-              f'      *((half16*)(Bs+bki*{bn_size}+bni))=*((half16*)(p{bslot}+'
-              f'((long)(kt*{bk}+bki)*N)+bn*{bn_size}+bni));', '    }']
-  else:
-    prefix, suffix = (f'    if (tid<{bn_size}) {{ ', ' }') if threads != bn_size else ('    ', '')
-    lines += [f'{prefix}long bo=((long)(bn*{bn_size}+tid)*K)+kt*{bk};']
-    lines += ['      #pragma unroll', f'      for (int q=0; q<{bk//16}; q++) '
-              f'*((half16*)(Bs+tid*{bs_stride}+q*16))=*((half16*)(p{bslot}+bo+q*16));']
-    lines[-1] += suffix
-  lines += ['    __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier(); '
-            '__builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");']
-  lines += ['    #pragma unroll', f'    for (int ki=0; ki<{bk//WMMA_K}; ki++) {{']
-  lines += [f'      half16 av[{tiles_m}], bv[{tiles_n}];', '      #pragma unroll', f'      for (int im=0; im<{tiles_m}; im++) {{']
-  if g.a_kxm:
-    lines += ['        half16 v;', '        #pragma unroll',
-              f'        for (int e=0; e<16; e++) v[e]=As[(ki*16+e)*{bm}+(wm*{tiles_m}+im)*16+row];', '        av[im]=v;']
-  else:
-    lines += [f'        av[im]=*((half16*)(As+(((wm*{tiles_m}+im)*16+row)*{as_stride}+ki*16)));']
-  lines += ['      }', '      #pragma unroll', f'      for (int jn=0; jn<{tiles_n}; jn++) {{']
-  if pack_b:
-    half_bits = ','.join(f'HALF_BITS(bp[{e}])' for e in range(8))
-    lines += ['        uint8 bp;', '        #pragma unroll',
-              f'        for (int e=0; e<8; e++) bp[e]=((unsigned*)Bs)[(ki*8+e)*{bn_size}+(wn*{tiles_n}+jn)*16+row];',
-              f'        ushort16 bb=(ushort16){{{half_bits}}};', '        bv[jn]=__builtin_bit_cast(half16,bb);']
-  elif g.b_kxn:
-    lines += ['        half16 v;', '        #pragma unroll',
-              f'        for (int e=0; e<16; e++) v[e]=Bs[(ki*16+e)*{bn_size}+(wn*{tiles_n}+jn)*16+row];', '        bv[jn]=v;']
-  else:
-    lines += [f'        bv[jn]=*((half16*)(Bs+(((wn*{tiles_n}+jn)*16+row)*{bs_stride}+ki*16)));']
-  lines += ['      }', '      #pragma unroll', f'      for (int im=0; im<{tiles_m}; im++) {{', '        #pragma unroll',
-            f'        for (int jn=0; jn<{tiles_n}; jn++) c[im][jn]=WMMA(av[im],bv[jn],c[im][jn]);', '      }', '    }']
-  lines += ['    __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier(); '
-            '__builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");', '  }']
-  lines += ['  #pragma unroll', f'  for (int im=0; im<{tiles_m}; im++) {{', '    #pragma unroll',
-            f'    for (int jn=0; jn<{tiles_n}; jn++) {{', '      #pragma unroll', '      for (int e=0; e<8; e++) {',
-            f'        int om=bm*{bm}+(wm*{tiles_m}+im)*16+e*2+halfrow, on=bn*{bn_size}+(wn*{tiles_n}+jn)*16+row;']
-  if output_nchw is not None:
-    lines.append('        int spatial_size=S*S;')
-    lines.append('        long oi=((long)(om/spatial_size)*N+on)*spatial_size+om%spatial_size;')
-  else: lines.append('        long oi=(long)om*N+on;')
-  if g.old is None: lines.append(f'        p{cslot}[oi]=(half)c[im][jn][e];')
-  else: lines.append(f'        p{cslot}[oi]=(float)((half)c[im][jn][e]+(half){g.scale}*p{g.old.arg.slot}[oi]);')
-  lines += ['      }', '    }', '  }']
-  lines += ['}']
-  return '\n'.join(lines)
-
-def _render_batched_gemm(g:BatchedGemmMatch, name:str) -> tuple[str, int]:
-  bm, bn_size, bk, threads, waves_m = 64, _batched_block_n(g), _batched_block_k(g), _batched_threads(g), 1
-  waves_n = threads//32
-  exact_n = g.n % bn_size == 0
-  tiles_m, tiles_n = bm//(waves_m*WMMA_M), bn_size//(waves_n*WMMA_N)
-  cslot, aslot, bslot = g.c.arg.slot, g.a.arg.slot, g.b.arg.slot
-  params = ', '.join(f'{"float" if x.dtype == dtypes.float else "half"}* p{x.arg.slot}'
-                     for x in sorted((g.c, g.a, g.b), key=lambda x:x.arg.slot))
-  params += ', int B, int M, int N, int K'
-  lines = [
-    '#define half _Float16',
-    'typedef half half16 __attribute__((ext_vector_type(16)));',
-    'typedef float float8 __attribute__((ext_vector_type(8)));',
-    '#define WMMA __builtin_amdgcn_wmma_f32_16x16x16_f16_w32',
-    f'extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size({threads}, {threads}))) {name}({params}) {{',
-    f'  __attribute__((shared, aligned(32))) half As[{bm*bk}], Bs[{bn_size*bk}];',
-    '  int bn=__builtin_amdgcn_workgroup_id_x(), bm=__builtin_amdgcn_workgroup_id_y(), '
-    'batch=__builtin_amdgcn_workgroup_id_z(), tid=__builtin_amdgcn_workitem_id_x();',
-    f'  int wave=tid>>5, lane=tid&31, wm=wave/{waves_n}, wn=wave%{waves_n}, row=lane&15, halfrow=lane>>4;',
-  ]
-  lines.append(f'  float8 c[{tiles_m}][{tiles_n}]={{}};')
-  lines.append(f'  for (int kt=0; kt<K/{bk}; kt++) {{')
-  aseg_count = bm*bk//16
-  lines += ['    #pragma unroll', f'    for (int q=0; q<{(aseg_count+threads-1)//threads}; q++) {{',
-            f'      int aseg=tid+q*{threads};']
-  if aseg_count % threads: lines.append(f'      if (aseg<{aseg_count}) {{')
-  lines += [f'      int ak=aseg/{bm//16}, am=(aseg%{bm//16})*16;',
-            f'      *((half16*)(As+ak*{bm}+am))=*((half16*)(p{aslot}+'
-            f'((long)(batch*K+kt*{bk}+ak)*M)+bm*{bm}+am));']
-  if aseg_count % threads: lines.append('      }')
-  lines.append('    }')
-  bsegs = bn_size//16
-  lines += ['    #pragma unroll', f'    for (int q=0; q<{bn_size*bk//(threads*16)}; q++) {{',
-            f'      int bseg=tid+q*{threads}, bki=bseg/{bsegs}, bni=(bseg%{bsegs})*16;',
-            f'      *((half16*)(Bs+bki*{bn_size}+bni))=' + ('' if exact_n else f'bn*{bn_size}+bni<N ? ') +
-            f'*((half16*)(p{bslot}+((long)(batch*K+kt*{bk}+bki)*N)+bn*{bn_size}+bni))' +
-            (';' if exact_n else ' : (half16){0};'), '    }']
-  lines += ['    __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier(); '
-            '__builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");']
-  lines += ['    #pragma unroll', f'    for (int ki=0; ki<{bk//WMMA_K}; ki++) {{',
-            f'      half16 av[{tiles_m}], bv[{tiles_n}];', '      #pragma unroll',
-            f'      for (int im=0; im<{tiles_m}; im++) {{', '        half16 v;', '        #pragma unroll',
-            f'        for (int e=0; e<16; e++) v[e]=As[(ki*16+e)*{bm}+(wm*{tiles_m}+im)*16+row];',
-            '        av[im]=v;', '      }', '      #pragma unroll', f'      for (int jn=0; jn<{tiles_n}; jn++) {{',
-            '        half16 v;', '        #pragma unroll',
-            f'        for (int e=0; e<16; e++) v[e]=Bs[(ki*16+e)*{bn_size}+(wn*{tiles_n}+jn)*16+row];',
-            '        bv[jn]=v;', '      }', '      #pragma unroll', f'      for (int im=0; im<{tiles_m}; im++) {{',
-            '        #pragma unroll', f'        for (int jn=0; jn<{tiles_n}; jn++) c[im][jn]=WMMA(av[im],bv[jn],c[im][jn]);',
-            '      }', '    }']
-  lines += ['    __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier(); '
-            '__builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");', '  }']
-  lines += ['  #pragma unroll', f'  for (int im=0; im<{tiles_m}; im++) {{', '    #pragma unroll',
-            f'    for (int jn=0; jn<{tiles_n}; jn++) {{',
-            f'      int om=bm*{bm}+(wm*{tiles_m}+im)*16+halfrow, on=bn*{bn_size}+(wn*{tiles_n}+jn)*16+row;',
-            ('      {' if exact_n else '      if (on<N) {'), '        #pragma unroll',
-            f'        for (int e=0; e<8; e++) p{cslot}[((long)(om+e*2)*N+on)*B+batch]=c[im][jn][e];',
-            '      }', '    }', '  }']
-  lines += ['}']
-  return '\n'.join(lines), bm
-
-def _render_direct_conv_bwd_activation(g:DirectConvBwdActivationMatch, name:str) -> str:
-  bm, bn, bk, threads = 128, min(g.cin, 64), 32, 128
-  persist_weights = g.cin <= 64
-  tiles_m = 2
-  tiles_n, aslot, bslot = bn//16, (3 if g.residual else 2), (4 if g.residual else 3)
-  params = 'float* p0, float* p1, half* p2, half* p3, half* p4' if g.residual else 'float* p0, half* p1, half* p2, half* p3'
-  lines = [
-    '#define half _Float16',
-    'typedef half half16 __attribute__((ext_vector_type(16)));',
-    'typedef half half8 __attribute__((ext_vector_type(8)));',
-    'typedef float float8 __attribute__((ext_vector_type(8)));',
-    'typedef unsigned uint8 __attribute__((ext_vector_type(8)));',
-    'typedef unsigned short ushort16 __attribute__((ext_vector_type(16)));',
-    '#define HALF_BITS(x) (unsigned short)(x), (unsigned short)((x)>>16)',
-    '#define WMMA __builtin_amdgcn_wmma_f32_16x16x16_f16_w32',
-    f'extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size({threads},{threads}))) {name}({params}) {{',
-    f'  __attribute__((shared, aligned(32))) half As[{bm*max(bk,bn)}], Bs[{bn*(g.cout if persist_weights else bk)}];',
-    '  int bn=__builtin_amdgcn_workgroup_id_x(), bm=__builtin_amdgcn_workgroup_id_y(), tid=__builtin_amdgcn_workitem_id_x();',
-    '  int wave=tid>>5, lane=tid&31, wm=wave, row=lane&15, halfrow=lane>>4;',
-    f'  float8 total[{tiles_m}][{tiles_n}]={{}};',
-    '  for (int patch=0; patch<9; patch++) {',
-    f'    float8 c[{tiles_m}][{tiles_n}]={{}};',
-  ]
-  if persist_weights:
-    lines += ['    #pragma unroll', f'    for (int q=0; q<{bn*(g.cout//2)//threads}; q++) {{',
-              f'      int be=tid+q*{threads}, bp=be/{bn}, ci=be%{bn};',
-              f'      half blo=p{bslot}[(bp*2)*{g.cin*9}+(bn*{bn}+ci)*9+patch];',
-              f'      half bhi=p{bslot}[(bp*2+1)*{g.cin*9}+(bn*{bn}+ci)*9+patch];',
-              f'      ((unsigned*)Bs)[bp*{bn}+ci]=__builtin_bit_cast(unsigned short,blo)|'
-              '((unsigned)__builtin_bit_cast(unsigned short,bhi)<<16);', '    }']
-    lines += ['    __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier(); '
-              '__builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");']
-  lines += [f'    for (int ct=0; ct<{g.cout//bk}; ct++) {{', f'      int co0=ct*{bk};']
-  if not persist_weights:
-    lines += ['      #pragma unroll', f'      for (int q=0; q<{bn*(bk//2)//threads}; q++) {{',
-              f'        int be=tid+q*{threads}, bp=be/{bn}, ci=be%{bn};',
-              f'        half blo=p{bslot}[(co0+bp*2)*{g.cin*9}+(bn*{bn}+ci)*9+patch];',
-              f'        half bhi=p{bslot}[(co0+bp*2+1)*{g.cin*9}+(bn*{bn}+ci)*9+patch];',
-              f'        ((unsigned*)Bs)[bp*{bn}+ci]=__builtin_bit_cast(unsigned short,blo)|'
-              '((unsigned)__builtin_bit_cast(unsigned short,bhi)<<16);', '      }']
-    lines += ['      __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier(); '
-              '__builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");']
-  lines += [
-            '      #pragma unroll', f'      for (int ki=0; ki<{bk//16}; ki++) {{',
-            f'        half16 av[{tiles_m}], bv[{tiles_n}];', '        #pragma unroll', f'        for (int im=0; im<{tiles_m}; im++) {{',
-            f'          int ami=bm*{bm}+(wm*{tiles_m}+im)*16+row, ab=ami/{g.spatial*g.spatial}, apos=ami%{g.spatial*g.spatial};',
-            f'          int ay=apos/{g.spatial}+1-patch/3, ax=apos%{g.spatial}+1-patch%3;',
-            f'          av[im]=(ay>=0 && ay<{g.spatial} && ax>=0 && ax<{g.spatial}) ? '
-            f'*((half16*)(p{aslot}+((long)ab*{g.spatial*g.spatial}+ay*{g.spatial}+ax)*{g.cout}+co0+ki*16)) : (half16){{0}};',
-            '        }',
-            '        #pragma unroll', f'        for (int jn=0; jn<{tiles_n}; jn++) {{', '          uint8 bp;', '          #pragma unroll',
-            f'          for (int e=0; e<8; e++) bp[e]=((unsigned*)Bs)[({"ct*16+" if persist_weights else ""}ki*8+e)*{bn}+jn*16+row];',
-            '          ushort16 bb=(ushort16){HALF_BITS(bp[0]),HALF_BITS(bp[1]),HALF_BITS(bp[2]),HALF_BITS(bp[3]),'
-            'HALF_BITS(bp[4]),HALF_BITS(bp[5]),HALF_BITS(bp[6]),HALF_BITS(bp[7])};',
-            '          bv[jn]=__builtin_bit_cast(half16,bb);', '        }', '        #pragma unroll',
-            f'        for (int im=0; im<{tiles_m}; im++) {{',
-            '          #pragma unroll', f'          for (int jn=0; jn<{tiles_n}; jn++) c[im][jn]=WMMA(av[im],bv[jn],c[im][jn]);',
-            '        }', '      }']
-  if not persist_weights:
-    lines += ['      __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier(); '
-              '__builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");']
-  lines += ['    }',
-            '    __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier(); '
-            '__builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");',
-            '    #pragma unroll', f'    for (int im=0; im<{tiles_m}; im++) {{', '      #pragma unroll',
-            f'      for (int jn=0; jn<{tiles_n}; jn++) total[im][jn]+=__builtin_convertvector('
-            '__builtin_convertvector(c[im][jn],half8),float8);', '    }', '  }',
-            '  #pragma unroll', f'  for (int im=0; im<{tiles_m}; im++) {{', '    #pragma unroll',
-            f'    for (int jn=0; jn<{tiles_n}; jn++) {{', '      #pragma unroll', '      for (int e=0; e<8; e++) {',
-            f'        int lm=(wm*{tiles_m}+im)*16+e*2+halfrow, lc=jn*16+row;',
-            f'        As[lc*{bm}+lm]=(half)total[im][jn][e];', '      }', '    }', '  }',
-            '  __builtin_amdgcn_fence(__ATOMIC_RELEASE,"workgroup"); __builtin_amdgcn_s_barrier(); '
-            '__builtin_amdgcn_fence(__ATOMIC_ACQUIRE,"workgroup");']
-  lines += ['  #pragma unroll', f'  for (int q=0; q<{bm*bn//(threads*8)}; q++) {{',
-            f'    int oe=tid*{bm*bn//threads}+q*8, lc=oe/{bm}, lm=oe%{bm};',
-            f'    int gm=bm*{bm}+lm, ob=gm/{g.spatial*g.spatial}, pos=gm%{g.spatial*g.spatial};',
-            f'    long oo=((long)ob*{g.cin}+bn*{bn}+lc)*{g.spatial*g.spatial}+pos;',
-            '    half8 grad=*((half8*)(As+oe));']
-  if g.residual: lines += ['    grad+=*((half8*)(p2+oo));', '    half8 z=__builtin_convertvector(*((float8*)(p1+oo)),half8);']
-  else: lines.append('    half8 z=*((half8*)(p1+oo));')
-  lines += ['    half8 sig=(half8)1.0/((half8)1.0+__builtin_elementwise_exp2(z*(half)-2.4554669595930156));',
-            '    *((float8*)(p0+oo))=__builtin_convertvector(sig*grad+(half)1.702*z*grad*sig*((half)1.0-sig),float8);', '  }']
-  lines.append('}')
-  return '\n'.join(lines)
-
-def direct_conv_bwd_activation_program(ast:UOp, renderer:Renderer, compile_binary:bool) -> UOp|None:
-  if not isinstance(g:=ast.tag, DirectConvBwdActivationMatch) or renderer.target.device != "AMD" or \
-     not renderer.target.arch.startswith("gfx11"): return None
-  name = f"coop_direct_conv_bwd_activation_{g.m}_{g.cin}_{g.cout}{'_res' if g.residual else ''}"
-  source = _render_direct_conv_bwd_activation(g, name)
-  sink = ast.replace(arg=replace(ast.arg, name=name, estimates=Estimates(2*g.m*g.cin*g.cout*9, 0, 0)))
-  slots = (0,1,2,3,4) if g.residual else (0,1,2,3)
-  bn = min(g.cin, 64)
-  info = ProgramInfo(name=name, global_size=((g.cin+bn-1)//bn, g.m//128, 1), local_size=(128,1,1),
-                     globals=slots, outs=(0,), ins=slots[1:])
-  src:tuple[UOp, ...] = (sink, UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=source))
-  if compile_binary: src += (UOp(Ops.BINARY, arg=renderer.compiler.compile_cached(source)),)
-  return UOp(Ops.PROGRAM, src=src, arg=info)
-
-
-def cooperative_gemm_program(ast:UOp, renderer:Renderer, compile_binary:bool) -> UOp|None:
-  if (g:=_match_gemm(ast, renderer.target.device, renderer.target.arch)) is not None:
-    output_nchw = ast.tag if isinstance(ast.tag, GemmOutputNCHW) else None
-    name = _gemm_family_name(g, output_nchw)
-    source = _render_gemm(g, name, output_nchw)
-    global_size = (g.n//_gemm_block_n(g), g.m//_gemm_block_m(g), 1)
-    local_size = (_gemm_threads(g), 1, 1)
-    estimates = Estimates(2*g.m*g.n*g.k, 2*(g.m*g.k+g.n*g.k+g.m*g.n), 2*(g.m*g.k+g.n*g.k+g.m*g.n))
-    slots = tuple(sorted((g.c.arg.slot, g.a.arg.slot, g.b.arg.slot) + ((g.old.arg.slot,) if g.old is not None else ())))
-    out_slot = g.c.arg.slot
-    suffix = f"_{g.m}_{g.n}_{g.k}"
-    variables:tuple[UOp, ...] = (UOp.variable(f"M{suffix}", g.m, g.m, dtypes.int),
-                                 UOp.variable(f"N{suffix}", g.n, g.n, dtypes.int),
-                                 UOp.variable(f"K{suffix}", g.k, g.k, dtypes.int))
-    if output_nchw is not None:
-      variables += (UOp.variable(f"S{suffix}_{output_nchw.spatial}", output_nchw.spatial, output_nchw.spatial, dtypes.int),)
-  elif (bg:=_match_batched_gemm(ast, renderer.target.device, renderer.target.arch)) is not None:
-    name = _batched_family_name(bg)
-    source, bm = _render_batched_gemm(bg, name)
-    global_size = ((bg.n+_batched_block_n(bg)-1)//_batched_block_n(bg), bg.m//bm, bg.batch)
-    local_size = (_batched_threads(bg), 1, 1)
-    estimates = Estimates(2*bg.batch*bg.m*bg.n*bg.k, 2*bg.batch*(bg.m*bg.k+bg.n*bg.k+2*bg.m*bg.n),
-                          2*bg.batch*(bg.m*bg.k+bg.n*bg.k+2*bg.m*bg.n))
-    slots, out_slot = tuple(sorted((bg.c.arg.slot, bg.a.arg.slot, bg.b.arg.slot))), bg.c.arg.slot
-    suffix = f"_{bg.batch}_{bg.m}_{bg.n}_{bg.k}"
-    variables = (UOp.variable(f"B{suffix}", bg.batch, bg.batch, dtypes.int), UOp.variable(f"M{suffix}", bg.m, bg.m, dtypes.int),
-                 UOp.variable(f"N{suffix}", bg.n, bg.n, dtypes.int), UOp.variable(f"K{suffix}", bg.k, bg.k, dtypes.int))
-  else: return None
-  sink = ast.replace(arg=replace(ast.arg, name=name, estimates=estimates))
-  info = ProgramInfo(name=name, global_size=global_size, local_size=local_size, vars=variables, globals=slots,
-                     outs=(out_slot,), ins=tuple(x for x in slots if x != out_slot))
-  src:tuple[UOp, ...] = (sink, UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=source))
-  if compile_binary: src += (UOp(Ops.BINARY, arg=renderer.compiler.compile_cached(source)),)
-  return UOp(Ops.PROGRAM, src=src, arg=info)
@@ -35,54 +35,13 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
      pass
    if good_tc_opt:
      if rngs is not None:
-        tc_sizes = [r.src[0] for r in rngs]
-        skinny_output = any(resolve(sz < 2, False) for sz in tc_sizes[:2])
-        very_long_reduce = resolve(tc_sizes[2] >= 4096, False)
-        long_reduce = resolve(tc_sizes[2] >= 1024, False)
-        small_m = resolve(tc_sizes[0] <= 32, False)
-        tiny_m = resolve(tc_sizes[0] <= 16, False)
-        if resolve(tc_sizes[0] >= 32, False) and resolve(tc_sizes[0] < 33, False) and resolve(tc_sizes[1] >= 1536, False) and \
-           resolve(tc_sizes[2] >= 128, False) and resolve(tc_sizes[2] <= 512, False):
-          upcast_n = 8 if resolve(tc_sizes[1] >= 4096, False) else 6
-          rngs[1] = tk.apply_opt(Opt(OptOps.UPCAST, tk.rngs.index(rngs[1]), upcast_n))[0]
-          rngs[0] = tk.apply_opt(Opt(OptOps.UPCAST, tk.rngs.index(rngs[0]), 2))[0]
-          rngs[0] = tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[0]), 8))[0]
-          rngs[1] = tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[1]), 2))[0]
-          return tk
        for tc_dim in [1,0]: # attempt to upcast M and N
-          if skinny_output or resolve(tc_sizes[tc_dim] >= 32768, False): continue
-          short_conv_n = tc_dim == 1 and resolve(tc_sizes[0] >= 65536, False) and resolve(tc_sizes[1] == 18, False) and \
-                         resolve(tc_sizes[2] <= 4, False)
-          upcast_sizes = [2] if short_conv_n or (very_long_reduce and (tiny_m or (tc_dim == 0 and small_m))) else [5,4,3,2]
-          szs = [sz for sz in upcast_sizes if rngs[tc_dim].src[0].divides(sz) is not None]
+          szs = [sz for sz in [5,4,3,2] if rngs[tc_dim].src[0].divides(sz) is not None]
          if szs:
            # set it to the replaced range
            rngs[tc_dim] = tk.apply_opt(Opt(OptOps.UPCAST, tk.rngs.index(rngs[tc_dim]), szs[0]))[0]
-        if skinny_output:
-          outer_rngs = [r for r in tk.rngs if r.arg[-1] is AxisType.GLOBAL and r not in rngs[:2]]
-          if outer_rngs and outer_rngs[0].src[0].divides(2) is not None:
-            tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(outer_rngs[0]), 2))
-          if very_long_reduce:
-            outer_rngs = [r for r in tk.rngs if r.arg[-1] is AxisType.GLOBAL and r not in rngs[:2]]
-            if outer_rngs and (outer_local:=next((x for x in (16, 4, 2) if outer_rngs[0].src[0].divides(x) is not None), None)):
-              tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(outer_rngs[0]), outer_local))
-          return tk
-        local_sizes = [2] if long_reduce and small_m else [4,2]
-        if (szs := [sz for sz in local_sizes if rngs[0].src[0].divides(sz) is not None]):
-          rngs[0] = tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[0]), szs[0]))[0]
-        if long_reduce and small_m and resolve(tc_sizes[1] >= 128, False) and \
-           rngs[1].arg[-1] is AxisType.GLOBAL and rngs[1].src[0].divides(3) is not None:
-          rngs[1] = tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[1]), 3))[0]
-        if tk.applied_opts[-1] == Opt(OptOps.LOCAL, 1, 4):
-          outer_rngs = [r for r in tk.rngs if r.arg[-1] is AxisType.GLOBAL]
-          if outer_rngs and resolve(outer_rngs[0].src[0] >= 384, False) and \
-             (outer_local:=next((x for x in (4, 3, 2) if outer_rngs[0].src[0].divides(x) is not None), None)):
-            tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(outer_rngs[0]), outer_local))
-        elif resolve(tc_sizes[0] <= 4, False) and resolve(tc_sizes[2] >= 256, False):
-          outer_rngs = [r for r in tk.rngs if r.arg[-1] is AxisType.GLOBAL]
-          outer_local = 16 if resolve(tc_sizes[2] >= 512, False) else 2
-          if (outer_rng:=next((r for r in reversed(outer_rngs) if r.src[0].divides(outer_local) is not None), None)) is not None:
-            tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(outer_rng), outer_local))
+        if (szs := [sz for sz in [4,2] if rngs[0].src[0].divides(sz) is not None]): # attempt to local N
+          tk.apply_opt(Opt(OptOps.LOCAL, tk.rngs.index(rngs[0]), szs[0]))
      return tk

  # make a copy so it does not mutate the input
@@ -123,8 +82,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:

  # are we grouping? (requires local shape support)
  if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= (240 if NOLOCALS else 2048), False):
-    group_sizes = (8, 16) if len(k.bufs) >= 7 else (16,)
-    for axis, sz in itertools.product((0, 1, 2), group_sizes):
+    for axis, sz in itertools.product((0, 1, 2), (16,)):
      try:
        k.apply_opt(Opt(OptOps.GROUPTOP, axis, sz))
        break
@@ -156,11 +114,10 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
  # potentially do more upcasts of non reduce axes based on a heuristic
  is_dsp = k.ren is not None and k.ren.target.device == "DSP"
  upcasted_axis: set[int] = set()
-  while resolve(prod(k.output_shape[i] for i in k.upcastable_dims) >= 1024) and (k.upcast_size() < 2):
+  while resolve(prod(k.output_shape[i] for i in k.upcastable_dims) >= 1024) and (k.upcast_size() < 32):
    xb_choices = []
    # consider all upcastable axes with 3 or 4 upcast (128 on the DSP)
-    upcast_amounts = ([128] if not len(upcasted_axis) else []) if is_dsp else ([3,4] if k.reduceop is not None else [2,3,4])
-    for axis, upcast_amount in itertools.product(k.upcastable_dims, upcast_amounts):
+    for axis, upcast_amount in itertools.product(k.upcastable_dims, ([128] if not len(upcasted_axis) else []) if is_dsp else [3,4]):
      # if we haven't upcasted it, it mods, and buffer has stride 0 on axis while having no stride 0 in the upcasted axis already
      if axis in upcasted_axis or k.full_shape[axis]%upcast_amount != 0: continue
      rng = k.rngs[axis]
@@ -184,17 +141,15 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:

  # if last reduce dim is small(ish), loop unroll the reduce
  # NOTE: this can fail on multireduce with mismatching dimensions, this is okay
-  four_by_four_reduce = len(k.unrollable_dims) >= 2 and all(resolve(k.full_shape[x] == 4, False) for x in k.unrollable_dims[-2:])
  try:
    if k.unrollable_dims and (k.upcast_size() <= 4 or not k.axes_of(AxisType.UNROLL)) and (k.upcast_size() < 64):
-      if (s:=k.full_shape[k.unrollable_dims[-1]]) <= 4:
+      if (s:=k.full_shape[k.unrollable_dims[-1]]) <= 32:
        k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, 0))
        # if it's small, upcast a second reduce dimension too
        if k.unrollable_dims and s <= 3 and k.full_shape[k.unrollable_dims[-1]] <= 3:
          k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, 0))
      else:
-        bn_spatial_reduce = len(k.unrollable_dims) >= 2 and resolve(k.full_shape[k.unrollable_dims[-2]] == 6, False)
-        for splits in ([8, 4] if bn_spatial_reduce else [4]):
+        for splits in [4]:
          if k.full_shape[axis:=k.unrollable_dims[-1]]%splits == 0:
            k.apply_opt(Opt(OptOps.UNROLL, len(k.unrollable_dims)-1, splits))
            break
@@ -211,44 +166,20 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
    if NOLOCALS:
      k.apply_opt(Opt(OptOps.NOLOCALS))
    else:
-      special_local = False
-      if four_by_four_reduce and len(global_axes:=k.axes_of(AxisType.GLOBAL, AxisType.LOOP)) >= 3:
-        lk, local_rngs = k.copy(), [k.rngs[x] for x in global_axes[-3:]]
-        try:
-          for rng, sz in zip(local_rngs, (16, 4, 4)):
-            lk.apply_opt(Opt(OptOps.LOCAL, lk.rngs.index(rng), sz))
-          k, special_local = lk, True
-        except KernelOptError: pass
      # prioritize making expand axes local
-      if not special_local:
-        local_axis_ranking = [(any(k.rngs[axis] not in b.src[1].get_idx().backward_slice for b in k.bufs), axis) \
-                                for axis in k.axes_of(AxisType.GLOBAL, AxisType.LOOP) if k.rngs[axis].src[0].op is Ops.CONST]
-        to_local: list[tuple[int, int]] = []
-        for _, axis in sorted(local_axis_ranking, key=lambda x: (-x[0], -x[1])):
-          local_size = prod(sz for _, sz in to_local)
-          local_sz: int|None = next((x for x in ([32] * (axis == 0) + [16,8,4,3,2]) if k.full_shape[axis] % x == 0 and local_size * x <= 128), None)
-          if local_sz is not None: to_local.append((axis, local_sz))
-        deleted_shape = 0
-        for axis, local_sz in sorted(to_local[:3]):
-          axis = axis - deleted_shape
-          will_delete_shape = local_sz == k.full_shape[axis]
-          k.apply_opt(Opt(OptOps.LOCAL, axis, local_sz))
-          if will_delete_shape: deleted_shape += 1
-
-      # Both 3x3 reduce axes are already fully unrolled above. Tile the exposed
-      # spatial axes without changing reduction order.
-      unroll_sizes = [k.full_shape[x] for x in k.axes_of(AxisType.UNROLL)]
-      global_axes = k.axes_of(AxisType.GLOBAL, AxisType.LOOP)
-      if unroll_sizes == [3, 3] and len(global_axes) >= 2:
-        axis_size = k.full_shape[global_axes[1]]
-        if axis_size <= 16: k.apply_opt(Opt(OptOps.UPCAST, 1, 0))
-        elif axis_size % 4 == 0: k.apply_opt(Opt(OptOps.UPCAST, 1, 4))
-        if len(global_axes) >= 4 and k.full_shape[global_axes[2]] == 4: k.apply_opt(Opt(OptOps.LOCAL, 2, 4))
-
-      remaining_reduces = [k.full_shape[x] for x in k.unrollable_dims]
-      if len(remaining_reduces) == 2 and remaining_reduces[0] == 6 and remaining_reduces[1] >= 8 and \
-         remaining_reduces[1] % 4 == 0 and k.upcast_size() <= 16:
-        k.apply_opt(Opt(OptOps.UNROLL, 1, 4))
+      local_axis_ranking = [(any(k.rngs[axis] not in b.src[1].get_idx().backward_slice for b in k.bufs), axis) \
+                              for axis in k.axes_of(AxisType.GLOBAL, AxisType.LOOP) if k.rngs[axis].src[0].op is Ops.CONST]
+      to_local: list[tuple[int, int]] = []
+      for _, axis in sorted(local_axis_ranking, key=lambda x: (-x[0], -x[1])):
+        local_size = prod(sz for _, sz in to_local)
+        local_sz: int|None = next((x for x in ([32] * (axis == 0) + [16,8,4,3,2]) if k.full_shape[axis] % x == 0 and local_size * x <= 128), None)
+        if local_sz is not None: to_local.append((axis, local_sz))
+      deleted_shape = 0
+      for axis, local_sz in sorted(to_local[:3]):
+        axis = axis - deleted_shape
+        will_delete_shape = local_sz == k.full_shape[axis]
+        k.apply_opt(Opt(OptOps.LOCAL, axis, local_sz))
+        if will_delete_shape: deleted_shape += 1

  # **** threading ****

@@ -301,9 +301,8 @@ class Scheduler:
            # TODO: remove tc_upcast_axes from the arg
            # do the reduce_axes always disappear? i think they don't
            # they need to be moved into the WMMA srcs
-            wmma_arg = (str(tc), tc.dims, tc.dtype_in, tc.dtype_out, self.ren.target.device, tc.threads, tc_upcast_axes, ()) #, tc_reduce_axes)
-            tc_uop = UOp(Ops.WMMA, src=(
-              srcs[0], srcs[1], UOp.const(tc.dtype_out, (0.0,)*tc.elements_per_thread[2])), arg=wmma_arg, tag=1)
+            tc_uop = UOp.wmma(srcs[0], srcs[1], UOp.const(tc.dtype_out, (0.0,)*tc.elements_per_thread[2]),
+                              tc.dims, self.ren.target.device, tc.threads, tc_upcast_axes=tc_upcast_axes)

            # preserve extra reduces
            reduce_ranges = [x for x in UOp.sink(*reduceop.src[1:]).toposort() if x.op is Ops.RANGE and x.arg[0] not in tc_reduce_axes]
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	02bb555162	error	2026-07-20 14:26:07 -07:00
geohot	1156623444	llm: minor fixes + tests	2026-07-20 14:21:30 -07:00
sirhcmandGitHub	b1cbd1a43f	pytest: use timeout_method signal (#17094 )	2026-07-20 15:19:24 -04:00
chenyuandGitHub	dbb0f6067e	clean up ALU rules in spec.py (#17095 )	2026-07-20 15:18:48 -04:00
chenyuandGitHub	8481eba866	allow-unsafe-pr-checkout for szdiff.yml (#17096 ) it uses sz.py on master to parse the change, should be safe	2026-07-20 15:08:54 -04:00
nimlgenandGitHub	2b96d64496	hcq2: tiny opts and fixes (#17092 )	2026-07-20 18:46:52 +03:00
Pol Puigdemont PlanaandGitHub	ef77963cfd	derivative of logsumexp is independent of max (#17088 ) same as #7009 but for logsumexp and logcumsumexp. fwd+bwd kernel count 5 -> 3 for both. gradients unchanged (ties, -inf masks, torch-compared at grad_atol=1e-7).	2026-07-20 06:52:16 -07:00
qazalandGitHub	abba2aebda	llama: correct fused qkv shape assert (#17086 )	2026-07-20 15:51:21 +09:00
qazalandGitHub	1cf8f2f68c	llama: inplace amax update (#17064 ) * llama: inplace amax update * remove amax_out return * work * fit * work * work * keep * diff cleanup	2026-07-20 15:05:41 +09:00
chenyuandGitHub	ac3f56a1a2	more shift tests (#17083 )	2026-07-19 16:05:13 -04:00
chenyuandGitHub	89117d8b9e	use real shift in l2i decomp [pr] (#17080 ) works for variable shift distace too, also fixed signed arithmetic fill	2026-07-19 13:15:06 -04:00
chenyuandGitHub	9970a0aad0	fix Tensor << Tensor for x86 (#17082 ) * fix Tensor << Tensor for x86 * torch	2026-07-19 12:31:05 -04:00
chenyuandGitHub	0146a30125	improve cast to unsign min_max [pr] (#17078 )	2026-07-18 21:58:41 -04:00
George HotzandGitHub	b53cd35cff	llm: make tokenizer fast (kimi) (#17077 ) * llm: make tokenizer fast * simpler * re.escape + qcom mypy fix	2026-07-18 17:31:59 -07:00
Rick WierengaandGitHub	82debb4557	only allow x86_64 target arch on X86Renderer (#17076 )	2026-07-18 19:51:01 -04:00
wozeparrotandGitHub	ee290b3e39	optim: mxfp8 zero 1 allgathers in fp8 (#17073 )	2026-07-18 07:44:50 -07:00
nimlgenandGitHub	232529ce88	hcq2: simpler sync (#17069 ) * x * y * n	2026-07-18 16:27:44 +03:00
qazalandGitHub	24d8681be7	viz: better sidebar collapse ux (#17072 )	2026-07-18 18:07:58 +09:00
chenyuandGitHub	47629f4bcf	more weak dtype materialization raise (#17071 )	2026-07-17 23:15:14 -04:00
chenyuandGitHub	f315df29a0	no weak Tensor from and to real buffer (#17067 ) * no weak Tensor from and to real buffer creation, assign, safe_save * is_numpy_ndarray to tensor * one more	2026-07-17 16:09:10 -04:00
George HotzandGitHub	86a6ad8ed2	llm: split cli.py into serve.py with the HTTP server (#17065 ) * llm: split cli.py into serve.py with the HTTP server * min edit	2026-07-17 10:45:37 -07:00
George HotzandGitHub	3ee2baf71d	llm: add tool calling support (kimi) (#17061 ) * llm: add tool calling support * simpler * cls * gpt cleanup * more gpt cleanups * tests for tools calling	2026-07-17 10:20:01 -07:00
qazalandGitHub	7dd3422c63	llama: replace two stage amax with atomics (#17063 ) * atomic amax in c kernels * quantize fp8 UOp kernel * diff	2026-07-17 19:27:10 +09:00
wozeparrotandGitHub	a836c3822a	gptoss: 3d mx block scale (#17062 )	2026-07-16 23:30:24 -07:00
sirhcmandGitHub	6f1176ea90	benchmarks: test usbgpu copy speeds on comma (#17060 )	2026-07-17 02:02:21 -04:00
George HotzandGitHub	46172bb7c7	llm: add optional jinja template support (kimi) (#17058 ) * add jinja template support (kimi) * fix tests * lil * more crap to fallback	2026-07-16 19:02:58 -07:00
chenyuandGitHub	88826a6f35	no weak dtype for randn_like either (#17055 )	2026-07-16 18:29:12 -04:00
chenyuandGitHub	3bfd62e915	fix 0 size tolist to match numpy (#17054 )	2026-07-16 17:42:52 -04:00
nimlgenandGitHub	709babb97c	system: remove sibling functions of PCIDevice (#17052 )	2026-07-17 00:15:32 +03:00
George HotzandGitHub	d8b83daac6	set tc_upcast_axes to None when done with it (#17053 ) * set tc_upcast_axes to None when done with it * no tag needed	2026-07-16 14:15:21 -07:00
stylishvoidandGitHub	c74149c973	avoid repeated parsing and toposort in _valid_priority [PR] (#17049 ) * avoid repeated parsing and toposort in _valid_priority * use backward_slice_with_self instead	2026-07-16 16:24:08 -04:00
chenyuandGitHub	6fa0b2b19e	materialize weak dtype casts to default (#17051 ) in clone and _buffer	2026-07-16 16:12:33 -04:00
George HotzandGitHub	4d8c3d3fc9	add test_hgemm to test_tiny (#17050 ) * add test_hgemm to test_tiny * dsp skip	2026-07-16 13:12:10 -07:00
George HotzandGitHub	2b1146b3f4	further clean up wmma (#17048 ) * further clean up wmma * comment	2026-07-16 11:43:23 -07:00
chenyuandGitHub	f6a92d0a16	sum_acc_dtype(weak) is weak (#17047 ) also no explicit weak for rand	2026-07-16 14:32:37 -04:00
George HotzandGitHub	61e104bdfb	use UOp.wmma everywhere (#17045 ) * use UOp.wmma everywhere * fix	2026-07-16 10:40:48 -07:00
chenyuandGitHub	5a4156c5d1	bitcast and element_size raise for weak dtypes (#17046 )	2026-07-16 13:07:45 -04:00
nimlgenandGitHub	7eb197b1bb	nv: always wait for reset (#17043 ) * nv: always wait for reset * x	2026-07-16 16:35:12 +03:00
chenyuandGitHub	dba8b6b505	allow weak alu operands (#17044 )	2026-07-16 09:33:20 -04:00
nimlgenandGitHub	e33e96415f	hcq2: tiny cleanupg (#17042 )	2026-07-16 16:14:54 +03:00
jk20342 GitHub Jacob Kitchen	810d8732f9	fix n^2 in limit_bufs by memoizing reachable loads [PR] (#17017 ) * fix n^2 in limit_bufs by memoizing reachable loads [pr] * Update test_schedule.py --------- Co-authored-by: Jacob Kitchen <[email protected]>	2026-07-15 23:54:04 -07:00
Théo H.GitHub George Hotz	1c74e044a4	search /usr/lib/wsl/lib first for linux (#17027 ) Co-authored-by: George Hotz <[email protected]>	2026-07-15 23:28:17 -07:00
George HotzandGitHub	8b0dd870ce	use wmma helper (#17038 )	2026-07-15 23:25:17 -07:00
qazalandGitHub	783042d216	viz: graph stays in place when sidebars resize (#17037 ) * viz: sidebars can resize independent of main graph * both sidebars * fix device-list * more work * no variables * raw 15% * fix custom view * minor detail	2026-07-16 11:47:22 +09:00
chenyuandGitHub	e8d3047a50	dtype_from_uop cleanup [PR] (#17036 )	2026-07-15 21:52:21 -04:00
chenyuandGitHub	6b7fee7d9f	minor lower_alu_dtype cleanup [PR] (#17034 )	2026-07-15 17:42:23 -04:00
chenyuandGitHub	be075b200a	weak dtypes in dtype_from_uop [PR] (#17032 ) * weak dtypes in dtype_from_uop [PR] * no weak in spec_program * weak const fold tests	2026-07-15 16:54:31 -04:00
chenyuandGitHub	3ffb4dc4bc	unify lower index in lower_alu_dtype [PR] (#17033 ) will work for weak types too	2026-07-15 16:38:42 -04:00
nimlgenandGitHub	d6fddb066f	usb: keep only custom (#17029 ) * usb: keep only custom * mockgpu by gpt * gpt said sorry * revert * reset * fix * flash	2026-07-15 22:40:52 +03:00
wozeparrotandGitHub	0d30f97584	mlperf: make v6.1 dir (#17031 )	2026-07-15 10:40:55 -07:00
chenyuandGitHub	c23d8188e1	remove _ensure_float [pr] (#17030 ) do this cast late. allow `SQRT(int)`	2026-07-15 11:19:20 -04:00
chenyuandGitHub	0d19970edc	least_upper_dtype in dtype_from_uop [PR] (#17028 )	2026-07-15 09:27:53 -04:00
chenyuandGitHub	ebe26420a7	update where Invalid rules [pr] (#17026 ) fixed TestInvalidTensor.test_tensor_index	2026-07-15 00:00:30 -04:00
wozeparrotandGitHub	06169f5013	gptoss: small fixes (#17025 )	2026-07-14 20:40:23 -07:00
chenyuandGitHub	47ddf94f17	remove InvalidType lt and gt (#17023 ) not really used	2026-07-14 21:59:15 -04:00
sirhcmandGitHub	c9baa2ef79	use pattern matcher in contiguous_view_offset [PR] (#17022 )	2026-07-14 19:37:36 -04:00
nimlgenandGitHub	4257939e50	remove copyin/copyout from Buffer (#17020 ) * remove copyin/copyout from Buffer * x * x * x * x	2026-07-14 19:47:22 +03:00
qazalandGitHub	939f28d571	fused qkv rope custom kernel (#17021 ) * work * fused qkv_norm * work * speed * not that yet * test cleanup * just clone * remove .realize() * cleanup tests	2026-07-15 01:08:42 +09:00
chenyuandGitHub	82fbca43c5	fix Tensor(np) dtype and support fp8 safetensor (#17019 )	2026-07-14 09:31:21 -04:00
chenyuandGitHub	872225e47d	update dtype tests for small dtypes (#17016 )	2026-07-14 08:07:00 -04:00
qazalandGitHub	edfef062ed	skip viz.cli -t in null device (#17018 )	2026-07-14 19:24:58 +09:00
chenyuandGitHub	55bb251130	add pm_manual_bf16_cast to Metal [pr] (#17015 ) mitigate metal compiler bug for `as_type<half>( (bfloat)(const) )`	2026-07-13 21:53:47 -04:00
sirhcmandGitHub	a9fbc7db7b	expect _offset support, CL and WEBGPU are outliers (#17014 )	2026-07-13 18:53:32 -04:00
chenyuandGitHub	9ce96c2628	fix subnormal in test_dtype (#17013 ) * fix subnormal in test_dtype should fix flaky test/backend/test_dtype.py::TestFp8e4m3::test_casts_from * better	2026-07-13 18:53:13 -04:00
chenyuandGitHub	c898dfe150	remove UOp.contiguous override [PR] (#17012 ) also cleaned up max_shard_shape	2026-07-13 16:10:58 -04:00
chenyuandGitHub	681a5e0cfd	remove UOp cast and bitcast override [PR] (#17011 )	2026-07-13 14:23:11 -04:00
chenyuandGitHub	0410c9325d	make test/null follow the SPEC (#17010 )	2026-07-13 14:01:41 -04:00
nimlgenandGitHub	e4bdc529c4	hcq2 ci (#17008 ) * hcq2 ci * x	2026-07-13 19:29:08 +03:00
nimlgenandGitHub	4536a57f79	hcq rename map (#17009 ) * hcq rename map * x	2026-07-13 19:23:12 +03:00
qazalandGitHub	62ad646d1c	llama: gemm/fa backward speedups (gpt 5.6) (#17007 ) * fp8 atb gemm speedup * work * revert * fa bw faster	2026-07-14 00:27:00 +09:00
nimlgenandGitHub	4d2becddf8	hcq2: spec=2 (#17006 ) * hcq2: spec=2 * hcq: isolate HCQ spec rules * chq * move	2026-07-13 18:13:33 +03:00
nimlgenandGitHub	ab9dde04a9	amd: do not spam with traps (#17004 )	2026-07-13 16:36:01 +03:00
chenyuandGitHub	223c6d74c3	remove unused get_empty_input_data (#17002 )	2026-07-12 22:13:30 -04:00
George HotzandGitHub	dde2e736e5	fix disable_gc decorator reentrancy (#16999 )	2026-07-12 15:32:54 -07:00