works

fix
close
2026-08-20 15:26:08 +00:00 · 2026-03-04 13:14:16 +08:00 · 2026-03-04 13:01:35 +08:00 · 2026-03-04 12:59:17 +08:00
119 changed files with 920 additions and 2253 deletions
@@ -45,10 +45,6 @@ inputs:
    description: "Install mesa"
    required: false
    default: 'false'
-  tinydreno:
-    description: "Install tinydreno"
-    required: false
-    default: 'false'
 runs:
  using: "composite"
  steps:
@@ -330,9 +326,3 @@ runs:
      if: inputs.mesa == 'true' && runner.os == 'macOS'
      shell: bash
      run: brew install sirhcm/tinymesa/tinymesa_cpu
-
-    # *** tinydreno ***
-    - name: Install tinydreno (linux)
-      if: inputs.tinydreno == 'true' && runner.os == 'Linux'
-      shell: bash
-      run: sudo curl -fL https://github.com/sirhcm/tinydreno/raw/refs/heads/master/libllvm-qcom.so -o /usr/lib/libllvm-qcom.so
@@ -332,7 +332,7 @@ jobs:
    # - name: Fuzz Padded Tensor Core GEMM (PTX)
    #   run: NV=1 NV_PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
    - name: HEVC Decode Benchmark
-      run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
+      run: VALIDATE=1 MAX_FRAMES=100 JITBEAM=1 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
    - name: Train MNIST
      run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
    - name: Run 10 CIFAR training steps
@@ -1,7 +1,7 @@
 name: Unit Tests
 env:
  # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '18'
+  CACHE_VERSION: '17'
  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
  PYTHONPATH: ${{ github.workspace }}
@@ -1011,26 +1011,3 @@ jobs:
          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
          python -m pytest -n=auto test/backend/test_ops.py --durations=20
-  qcomclcompiletests:
-    name: Compile-only (QCOM CL)
-    runs-on: ubuntu-24.04-arm
-    timeout-minutes: 15
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v4
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: compile-qcomcl
-          deps: testing_unit
-          tinydreno: 'true'
-          python-version: '3.12'
-      - name: Set env
-        shell: bash
-        run: printf "NULL=1\nNULL_ALLOW_COPYOUT=1\nNULL_QCOMCL=1" >> $GITHUB_ENV
-      - name: Run test_ops
-        shell: bash
-        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
-          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
-          python -m pytest -n=auto test/backend/test_ops.py --durations=20
@@ -1344,7 +1344,6 @@ def train_llama3():
  vocab_mask:Tensor = Tensor.arange(model_params['vocab_size']).reshape(1, 1, -1) >= real_vocab_size

  model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
-
  params = get_parameters(model)
  # weights are all bfloat16 for now
  assert params and all(p.dtype == dtypes.bfloat16 for p in params)
@@ -1382,20 +1381,16 @@ def train_llama3():

    vocab_mask.shard_(device, axis=2).realize()

-  is_offload_optim = bool(getenv("OFFLOAD_OPTIM"))
-  is_fake_offload = Device.DEFAULT == "NULL"
-  optim_device = ("CPU" if not is_fake_offload else "NULL:99") if is_offload_optim else None
+  optim_device = "CPU" if getenv("OFFLOAD_OPTIM") else None
  optim = GradAccClipAdamW(get_parameters(model), lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2,
                           eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)

  # init grads
-  if is_offload_optim:
-    for p in optim.params:
-      p.grad = Tensor.zeros(p.shape, dtype=p.dtype, device=optim_device, requires_grad=False).contiguous().realize()
-  else:
-    for p in optim.params:
-      p.grad = p.zeros_like().contiguous().realize()
+  for p in optim.params:
+    p.grad = p.empty_like().realize()
  grads: list[Tensor] = [p.grad for p in optim.params]
+  for p in optim.params:
+    p.grad.assign(p.grad.zeros_like()).realize()

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

@@ -1421,9 +1416,8 @@ def train_llama3():
    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
    loss.backward()
    assert all(p.grad is g for p,g in zip(optim.params, grads))
-    loss_cpu = loss.flatten().float().to("CPU")
-    Tensor.realize(loss_cpu, *grads)
-    return loss_cpu
+    Tensor.realize(loss, *grads)
+    return loss.flatten().float().to("CPU")

  @TinyJit
  def optim_step():
@@ -1431,13 +1425,12 @@ def train_llama3():
    scheduler.step()

    for g in grads:
-      g.assign(g.zeros_like())
+      g.assign(g.zeros_like()).realize()

-    lr_cpu = optim.lr.float().to("CPU")
-    grad_norm_cpu = grad_norm.float().to("CPU")
-    Tensor.realize(lr_cpu, grad_norm_cpu, *grads)
+    lr = optim.lr
+    Tensor.realize(lr, *grads)

-    return lr_cpu, grad_norm_cpu
+    return lr.float().to("CPU"), grad_norm.float().to("CPU")

  @TinyJit
  @Tensor.train(False)
@@ -1485,14 +1478,13 @@ def train_llama3():
  step_times = []
  while i < MAX_STEPS:
    GlobalCounters.reset()
-    actual_gbs = GBS if i >= 2 else BS
    if getenv("TRAIN", 1):
      profile_marker(f"train @ {i}")
      st = time.perf_counter()

      stopped = False
      losses, data_time, dev_time = [], 0, 0
-      for _ in range(grad_acc if i >= 2 else 1):
+      for _ in range(grad_acc if i >= 3 else 1):
        ist = time.perf_counter()
        try: tokens = next(train_iter)
        except StopIteration:
@@ -1517,7 +1509,7 @@ def train_llama3():
      if BENCHMARK: step_times.append(step_time)

      i += 1
-      sequences_seen += actual_gbs
+      sequences_seen += GBS

      mem_gb = GlobalCounters.mem_used / 1e9
      gflops = GlobalCounters.global_ops / 1e9 / dev_time
@@ -1560,7 +1552,7 @@ def train_llama3():
        print(f"epoch global_ops: {GlobalCounters.global_ops:_}, "
              f"epoch global_mem: {GlobalCounters.global_mem:_}")

-    if (sequences_seen // EVAL_FREQ != (sequences_seen - actual_gbs) // EVAL_FREQ and (i != 1 or EVAL_FREQ == 1)) or (BENCHMARK and i == BENCHMARK):
+    if (sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1)) or (BENCHMARK and i == BENCHMARK):
      if EVAL_BS == 0: return
      tqdm.write(f"evaluating after {sequences_seen} sequences")
      profile_marker(f"eval @ {i}")
@@ -52,6 +52,5 @@ class GradAccClipAdamW(Optimizer):
    return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]

  def _apply_update(self, t:Tensor, up:Tensor) -> Tensor:
-    wd = self.wd if t.ndim >= 2 else 0.0
-    up = up.shard_like(t) + self.lr.to(t.device) * wd * t.detach()
+    up = up.shard_like(t) + self.lr.to(t.device) * self.wd * t.detach()
    return t.detach() - up.cast(t.dtype)
@@ -1,37 +0,0 @@
-#!/usr/bin/env bash
-
-export PYTHONPATH="."
-export DEV=${DEV:-AMD}
-export EMULATE="AMD_CDNA4"
-export CHECK_OOB=0
-export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
-
-export DEBUG=${DEBUG:-2}
-export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
-export ALL2ALL=${ALL2ALL:-1}
-export USE_ATOMICS=${USE_ATOMICS:-0}
-export ASM_GEMM=${ASM_GEMM:-1}
-export WQKV=${WQKV:-1}
-export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
-
-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-1} MP=${MP:-8}
-export BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
-
-export MODEL="llama3"
-export BASEDIR="/raid/datasets/c4/"
-export LLAMA3_SIZE=${LLAMA3_SIZE:-"405B"}
-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=1 BENCHMARK=10
-if [ -z "$FULL_LAYERS" ]; then
-  export LLAMA_LAYERS=2
-fi
-
-python3 examples/mlperf/model_train.py
@@ -1,32 +0,0 @@
-#!/usr/bin/env bash
-
-export PYTHONPATH="."
-export DEV=${DEV:-AMD}
-export EMULATE="AMD_CDNA4"
-export CHECK_OOB=0
-export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
-
-export DEBUG=${DEBUG:-0}
-export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
-export ALL2ALL=${ALL2ALL:-1}
-export USE_ATOMICS=${USE_ATOMICS:-0}
-export ASM_GEMM=${ASM_GEMM:-1}
-export WQKV=${WQKV:-1}
-export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
-
-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-1} MP=${MP:-8}
-export BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1152}
-
-export MODEL="llama3"
-export BASEDIR="/raid/datasets/c4/"
-export LLAMA3_SIZE=${LLAMA3_SIZE:-"405B"}
-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
@@ -5,7 +5,6 @@ export DEV=${DEV:-AMD}
 export EMULATE="AMD_CDNA4"
 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}
@@ -15,7 +14,7 @@ export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-0}

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

 export MODEL="llama3"
@@ -23,7 +22,7 @@ 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 LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
 export SEQLEN=${SEQLEN:-8192}
@@ -1,43 +0,0 @@
-#!/usr/bin/env bash
-
-export PYTHONPATH="."
-export DEV=${DEV:-AMD}
-export EMULATE="AMD_CDNA4"
-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 USE_ATOMICS=${USE_ATOMICS:-0}
-export ASM_GEMM=${ASM_GEMM:-1}
-export WQKV=${WQKV:-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=1 BENCHMARK=10
-if [ -z "$FULL_LAYERS" ]; then
-  export LLAMA_LAYERS=2
-fi
-
-python3 examples/mlperf/model_train.py
@@ -15,7 +15,7 @@ export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-0}

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

 export MODEL="llama3"
@@ -23,7 +23,7 @@ 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 LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
 export SEQLEN=${SEQLEN:-8192}
@@ -1,38 +0,0 @@
-#!/usr/bin/env bash
-
-export PYTHONPATH="."
-export DEV=${DEV:-AMD}
-export EMULATE="AMD_CDNA4"
-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 USE_ATOMICS=${USE_ATOMICS:-0}
-export ASM_GEMM=${ASM_GEMM:-1}
-export WQKV=${WQKV:-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
@@ -31,7 +31,7 @@ def compile(onnx_file):
  for i in range(3):
    GlobalCounters.reset()
    print(f"run {i}")
-    with Context(DEBUG=max(DEBUG.value, 2 if i == 2 else 1), OPENPILOT_HACKS=1):
+    with Context(DEBUG=max(DEBUG.value, 2 if i == 2 else 1)):
      ret = run_onnx_jit(**inputs).numpy()
    # copy i == 1 so use of JITBEAM is okay
    if i == 1: test_val = np.copy(ret)
@@ -10,9 +10,9 @@ HEVC_ROUNDUP = getenv("DATA_ROUNDUP", 32)
@functools.cache
 def _hevc_jitted_decoder(out_image_size:tuple[int, int], max_hist:int, inplace:bool):
  def hevc_decode_frame(pos:Variable, hevc_tensor:Tensor, offset:Variable, sz:Variable, opaque:Tensor, i:Variable, *hist:Tensor, outbuf:Tensor|None=None):
-    x = hevc_tensor[offset:offset+sz*HEVC_ROUNDUP].decode_hevc_frame(pos, out_image_size, opaque[i], hist).realize()
+    x = hevc_tensor[offset:offset+sz*HEVC_ROUNDUP].decode_hevc_frame(pos, out_image_size, opaque[i], hist)
    if outbuf is not None: outbuf.assign(x).realize()
-    return x
+    return x.realize()
  return TinyJit(hevc_decode_frame)

 def hevc_decode(hevc_tensor:Tensor, opaque:Tensor, frame_info:list, luma_h:int, luma_w:int,
@@ -74,14 +74,10 @@ if __name__ == "__main__":
  Device.default.synchronize()

  # decode all frames using the iterator
-  tm = Timing("decoding whole file: ", on_exit=(lambda et: f", {len(frame_info)} frames, {len(frame_info)/(et/1e9):.2f} fps"))
-  with tm:
+  with Timing("decoding whole file: ", on_exit=(lambda et: f", {len(frame_info)} frames, {len(frame_info)/(et/1e9):.2f} fps")):
    images = list(hevc_decode(hevc_tensor, opaque_nv, frame_info, luma_h, luma_w, history=hist, preallocated_outputs=out_images))
    Device.default.synchronize()

-  fps = len(frame_info)/(tm.et/1e9)
-  assert fps >= getenv("ASSERT_FPS", 0), f"HEVC decode too slow: {fps:.2f} fps"
-
  # validation
  if getenv("VALIDATE", 0):
    import pickle
@@ -1,25 +1,25 @@
-import os, subprocess, sys, shlex
+import os, subprocess, sys
 from pathlib import Path
 from tinygrad.helpers import temp

 EXAMPLES_DIR = Path(__file__).parent
 PROFILE_PATH = Path(temp("profile.pkl", append_user=True))

-EXAMPLES = {
-  "empty":"test/backend/test_custom_kernel.py TestCustomKernel.test_empty",
-  "plus":"test/test_tiny.py TestTiny.test_plus",
-  "gemm":"-c \"from tinygrad import Tensor; (Tensor.empty(N:=64, N)@Tensor.empty(N, N)).realize()\"",
-  "ops":"extra/sqtt/examples/discover_ops.py"
-}
+EXAMPLES = [
+  "test/backend/test_custom_kernel.py TestCustomKernel.test_empty",
+  "test/test_tiny.py TestTiny.test_plus",
+  "test/test_tiny.py TestTiny.test_gemm",
+  "extra/sqtt/examples/discover_ops.py"
+]

 if __name__ == "__main__":
  arch = subprocess.check_output(["python", "-c", "from tinygrad import Device; print(Device['AMD'].arch)"], text=True,
                                 env={**os.environ, "DEBUG":"0"}).rstrip()
  (EXAMPLES_DIR/arch).mkdir(exist_ok=True)
-  for name,test in EXAMPLES.items():
+  for test in EXAMPLES:
    for i in range(2):
      # AM_RESET=1 gets a clear trace, does not work on mi300 machines
-      subprocess.run([sys.executable, *shlex.split(test)], cwd=EXAMPLES_DIR.parent.parent.parent,
+      subprocess.run([sys.executable, *test.split()], cwd=EXAMPLES_DIR.parent.parent.parent,
                     env={**os.environ, "AMD":"1", "AM_RESET":"1" if not arch.startswith("gfx9") else "0", "VIZ":"-2", "PYTHONPATH":"."})
-      PROFILE_PATH.rename(dest:=EXAMPLES_DIR/arch/f"profile_{name}_run_{i}.pkl")
+      PROFILE_PATH.rename(dest:=EXAMPLES_DIR/arch/f"profile_{test.split('.')[-1].replace('test_', '')}_run_{i}.pkl")
      print(f"saved SQTT trace to {dest}")
@@ -19,38 +19,6 @@ def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None
 def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
  return _sharded_empty(ref.shape, ref, axis)

-@functools.cache
-def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch):
-  def grad(dou:UOp, ker:UOp) -> tuple[None, None, UOp, UOp, UOp]:
-    do = Tensor(dou, device=dou.device)
-    attn = Tensor(ker.src[1].after(ker), device=ker.src[1].device)
-    l_vec = Tensor(ker.src[2].after(ker), device=ker.src[2].device)
-    xq = Tensor(ker.src[3], device=ker.src[3].device)
-    xk = Tensor(ker.src[4], device=ker.src[4].device)
-    xv = Tensor(ker.src[5], device=ker.src[5].device)
-
-    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)
-
-    # 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)
-    delta_vec, dq = Tensor.custom_kernel(delta_vec, dq, attn, do, fxn=functools.partial(custom_fa_backward_pre, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:2]
-
-    dq, dk_partial, dv_partial = Tensor.custom_kernel(dq, dk_partial, dv_partial, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:3]
-
-    # unshuffle dq: atomic_pk_add_bf16_with_warpid creates a shuffled layout within each 16x128 tile
-    # decompose each tile into (j=4, a=2, b=2, d=4, e=4, k=4, c=2) and permute to (e, k, j, a, d, b, c) = standard row-major
-    dq = dq.reshape(B, H, N//16, 4, 2, 2, 4, 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)
-
-    return None, None, dq.uop, dk.uop, dv.uop
-  return grad
-
 def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False):
  assert attn_mask is None, "attn_mask not supported"
  assert is_causal, "only causal attention supported"
@@ -77,7 +45,23 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
  attn = _sharded_empty_like(xq, axis=shard_axis)
  l_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)

-  grad = _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch)
+  def grad(dou:UOp, _) -> tuple[None, None, UOp, UOp, UOp]:
+    do = Tensor(dou, device=dou.device)
+    dq_in = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t)
+    dq = _sharded_empty_like(xq, axis=shard_axis)
+    dk = _sharded_empty_like(xk, axis=shard_axis)
+    dv = _sharded_empty_like(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)
+    delta_vec, dq_in = Tensor.custom_kernel(delta_vec, dq_in, attn, do, fxn=functools.partial(custom_fa_backward_pre, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:2]
+
+    dq_in, dk, dv = Tensor.custom_kernel(dq_in, dk, dv, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:3]
+
+    # unshuffle dq
+    dq = Tensor.custom_kernel(dq, dq_in, fxn=functools.partial(custom_fa_backward_post, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[0]
+
+    return None, None, dq.uop, dk.uop, dv.uop

  attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D), grad_fxn=grad)[:2]

@@ -105,6 +89,7 @@ def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, device:str, arch:st
                  arg=KernelInfo(name="custom_fa_forward", estimates=estimates))

  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
+
  lib = bytearray(lib)
  rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
  struct.pack_into('<I', lib, rodata_off, 160000)
@@ -135,6 +120,7 @@ def custom_fa_backward_pre(delta_vec:UOp, dq:UOp, o:UOp, do:UOp, device:str, arc
                  arg=KernelInfo(name="custom_fa_backward_pre", estimates=estimates))

  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
+
  lib = bytearray(lib)
  rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
  struct.pack_into('<I', lib, rodata_off, 160000)
@@ -152,7 +138,7 @@ def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_ve
  BLOCK_SIZE_KV = 256
  NUM_WARPS = 4
  NUM_THREADS = 64 * NUM_WARPS
-  gsz = (H, N // BLOCK_SIZE_KV, B)
+  gsz = (H_KV, 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")
@@ -165,6 +151,7 @@ def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_ve
                  arg=KernelInfo(name="custom_fa_backward", estimates=estimates))

  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
+
  lib = bytearray(lib)
  rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
  struct.pack_into('<I', lib, rodata_off, 160000)
@@ -195,6 +182,7 @@ def custom_fa_backward_post(dq_out:UOp, dq_in:UOp, device:str, arch:str, B:int,
                  arg=KernelInfo(name="custom_fa_backward_post", estimates=estimates))

  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
+
  lib = bytearray(lib)
  rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
  struct.pack_into('<I', lib, rodata_off, 160000)
@@ -37,7 +37,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, 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 {
@@ -47,7 +47,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_KV, (ATTN_N / BLOCK_SIZE_KV), ATTN_B); }
  dim3 block() { return dim3(NUM_THREADS); }
  size_t dynamic_shared_memory() { return MAX_SHARED_MEMORY; }
 };
@@ -55,12 +55,10 @@ 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 kv_head_idx = q_head_idx_fixed / GROUP_SIZE;
-  const int q_head_in_group = q_head_idx_fixed % GROUP_SIZE;
+  const int kv_head_idx = blockIdx.x;  // This is the KV head index
  const int seq_idx = blockIdx.y;
  const int batch_idx = blockIdx.z;
-  const int first_q_head = q_head_idx_fixed;
+  const int first_q_head = kv_head_idx * GROUP_SIZE;

  const int warpid = kittens::warpid();
  const int j = seq_idx * NUM_WARPS + warpid;
@@ -72,7 +70,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 * GROUP_SIZE;
  const int k_pos = j * WARP_SIZE_KV;

  constexpr float L_SCALE_FACTOR = 1.44269504089f;
@@ -3357,14 +3355,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, 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, 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);
@@ -66,7 +66,7 @@ template<int D, typename T=bf16, typename L=row_l, typename S=rt_32x16_s> using
 template<int D, typename T=bf16, typename L=col_l, typename S=rt_16x32_s> using qo_tile_transposed = rt<T, D, Q_BLOCK_SIZE, L, S>;
 template<int D, typename T=bf16, typename L=row_l, typename S=rt_32x16_s> using kv_tile = rt<T, KV_BLOCK_SIZE, D, L, S>;
 template<int D, typename T=bf16, typename L=col_l, typename S=rt_16x32_s> using kv_tile_transposed = rt<T, D, KV_BLOCK_SIZE, L, S>;
-template<typename T=float, typename L=col_l, typename S=rt_16x32_4_s> using attn_tile = rt<T, KV_BLOCK_SIZE, Q_BLOCK_SIZE, L, S>;
+template<int D, typename T=float, typename L=col_l, typename S=rt_16x32_4_s> using attn_tile = rt<T, KV_BLOCK_SIZE, Q_BLOCK_SIZE, L, S>;

 /**********************************************************/
 template<int THR_X, int THR_Y>
@@ -103,7 +103,7 @@ __device__ inline void mask_kv_tile(RT &dst, int q_abs, int k_abs, uint32_t neg_

    #pragma unroll
    for (int i = 0; i < dst.height; ++i) {
-        // Row base of the 32x* chunk produced by MFMA
+        // Row base of the 32x* chunk produced by MFMA 
        const int row_base = (i * 32) + ((lane >> 5) << 2); // multiplesof 4

        // Relative index of the FIRST element in this row-chunk w.r.t. q_pos
@@ -148,7 +148,7 @@ __device__ inline void mask_kv_tile(RT &dst, int q_abs, int k_abs, uint32_t neg_
 /**********************************************************/

 template<int D> struct attn_globals {
-    _gl_QKVO Qg, Kg, Vg, Og;
+    _gl_QKVO Qg, Kg, Vg, Og; 
    gl<float, -1, -1, -1, -1> L_vec;
    dim3 grid() { return dim3(ATTN_H, ((ATTN_N / Q_BLOCK_SIZE + NUM_WARPS - 1) / NUM_WARPS), ATTN_B); }
    dim3 block() { return dim3(NUM_THREADS); }
@@ -196,10 +196,10 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p

    kv_tile<D, bf16, col_l, rt_16x32_4_s> v_reg;
    qo_tile_transposed<D, float, col_l, rt_32x32_s> o_reg; // Output tile.
-    attn_tile<float, col_l, rt_32x32_s> att_block[2]; // attention tile, in float.
-    attn_tile<bf16, col_l, rt_32x32_s> att_block_bf16;
-    attn_tile<bf16, col_l, rt_16x32_4_s> att_block_bf16_in;
-    typename attn_tile<float, col_l, rt_32x32_s>::row_vec max_vec, norm_vec, max_vec_prev, scale_vec;
+    attn_tile<D, float, col_l, rt_32x32_s> att_block[2]; // attention tile, in float.
+    attn_tile<D, bf16, col_l, rt_32x32_s> att_block_bf16;
+    attn_tile<D, bf16, col_l, rt_16x32_4_s> att_block_bf16_in;
+    typename attn_tile<D, float, col_l, rt_32x32_s>::row_vec max_vec, norm_vec, max_vec_prev, scale_vec;

    zero(o_reg);
    zero(norm_vec);
@@ -241,8 +241,8 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    zero(att_block[0]);
    transpose(k_reg_transposed, k_reg);
    mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
-    __builtin_amdgcn_sched_barrier(0);
-    if constexpr (causal) {
+    __builtin_amdgcn_sched_barrier(0); 
+    if constexpr (causal) { 
        const int kv_end_pos = (1) * KV_BLOCK_SIZE;
        if (__builtin_expect(q_start_pos < kv_end_pos, 0)) {  // Only mask if needed
            mask_kv_tile(att_block[0], tile_idx, 0, neg_inf_v, lane);
@@ -269,7 +269,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    load(k_reg, k_smem[1]);
    // All warps then collaboratively load in the third slice of K (K2) into shared memory
    G::load<1, false>(k_smem[0], g.Kg, {batch_idx, 2, head_idx_kv, 0}, swizzled_offsets_K);
-    // All warps then collaboratively load in the second slice of V (V1) into shared memory
+    // All warps then collaboratively load in the second slice of V (V1) into shared memory 
    G::load<1, false>(v_smem[1], g.Vg, {batch_idx, 1, head_idx_kv, 0}, swizzled_offsets_V);
    asm volatile("s_waitcnt lgkmcnt(0)");
    asm volatile("s_waitcnt vmcnt(4)");
@@ -288,7 +288,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
        mul(norm_vec, norm_vec, scale_vec);
        col_sum(norm_vec, att_block[0], norm_vec);
        copy(att_block_bf16, att_block[0]);
-        att_block_bf16_in = *reinterpret_cast<attn_tile< bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+        att_block_bf16_in = *reinterpret_cast<attn_tile<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
        sched_barrier_exp_pairs<6, 3, 1>();
        sched_barrier_pairs<10, 5, 1>();
        __builtin_amdgcn_sched_barrier(0);
@@ -296,7 +296,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
        __builtin_amdgcn_sched_barrier(0);

        // Cluster 1:
-        //      Load K3 into shared
+        //      Load K3 into shared 
        G::load<1, false>(k_smem[1], g.Kg, {batch_idx, j, head_idx_kv, 0}, swizzled_offsets_K);
        //      Load V0 into registers
        load(v_reg, v_smem[0]);
@@ -348,7 +348,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
        mul(norm_vec, norm_vec, scale_vec);
        col_sum(norm_vec, att_block[1], norm_vec);
        copy(att_block_bf16, att_block[1]);
-        att_block_bf16_in = *reinterpret_cast<attn_tile<bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+        att_block_bf16_in = *reinterpret_cast<attn_tile<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
        sched_barrier_exp_pairs<6, 3, 3>();
        sched_barrier_pairs<10, 5, 3>();
        __builtin_amdgcn_s_setprio(0);
@@ -417,7 +417,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p

    col_sum(norm_vec, att_block[0], norm_vec);
    copy(att_block_bf16, att_block[0]);
-    att_block_bf16_in = *reinterpret_cast<attn_tile<bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+    att_block_bf16_in = *reinterpret_cast<attn_tile<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
    sched_barrier_exp_pairs<6, 3, 5>();
    sched_barrier_pairs<10, 5, 5>();
    __builtin_amdgcn_sched_barrier(0);
@@ -482,7 +482,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    mul(norm_vec, norm_vec, scale_vec);
    col_sum(norm_vec, att_block[1], norm_vec);
    copy(att_block_bf16, att_block[1]);
-    att_block_bf16_in = *reinterpret_cast<attn_tile<bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+    att_block_bf16_in = *reinterpret_cast<attn_tile<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
    sched_barrier_exp_pairs<6, 3, 7>();
    sched_barrier_pairs<10, 5, 7>();
    __builtin_amdgcn_sched_barrier(0);
@@ -544,7 +544,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    mul(norm_vec, norm_vec, scale_vec);
    col_sum(norm_vec, att_block[0], norm_vec);
    copy(att_block_bf16, att_block[0]);
-    att_block_bf16_in = *reinterpret_cast<attn_tile<bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+    att_block_bf16_in = *reinterpret_cast<attn_tile<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
    sched_barrier_exp_pairs<6, 3, 9>();
    sched_barrier_pairs<10, 5, 9>();
    __builtin_amdgcn_sched_barrier(0);
@@ -586,7 +586,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p

    col_sum(norm_vec, att_block[1], norm_vec);
    copy(att_block_bf16, att_block[1]);
-    att_block_bf16_in = *reinterpret_cast<attn_tile<bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+    att_block_bf16_in = *reinterpret_cast<attn_tile<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);

    __builtin_amdgcn_sched_barrier(0);
    mul_col(o_reg, o_reg, scale_vec);
@@ -505,9 +505,7 @@ tiny_backend_out = {**{f"aten.{x}.out":getattr(Tensor,x) for x in simple_tensor_
  "aten.lt.Tensor_out": Tensor.__lt__, "aten.lt.Scalar_out": Tensor.__lt__,
  "aten.le.Tensor_out": Tensor.__le__, "aten.le.Scalar_out": Tensor.__le__,
  "aten.clamp_max.Tensor_out": lambda input,max_: input.clamp(max_=max_),
-  "aten.clamp_max.out": lambda input,max_: input.clamp(max_=max_),
  "aten.clamp_min.Tensor_out": lambda input,min_: input.clamp(min_=min_),
-  "aten.clamp_min.out": lambda input,min_: input.clamp(min_=min_),
  "aten.fmod.Tensor_out": lambda input,other: input-input.div(other, rounding_mode="trunc")*other,
  # TODO: this might result in overflow issues
  "aten.round.decimals_out": lambda self,decimals: (self*10**decimals).round()/10**decimals,
@@ -590,7 +588,7 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
  "aten.repeat": lambda x,*repeats: Tensor.repeat(x,*repeats).contiguous(), # not a view
  "aten._softmax": lambda self,dim,half_to_float: self.softmax(dim),
  "aten._log_softmax": lambda self,dim,half_to_float: self.log_softmax(dim),
-  "aten.random_": lambda self: Tensor.randint(*self.shape, low=self.dtype.min, high=self.dtype.max, device=self.device, dtype=self.dtype),
+  "aten.random_": lambda self: Tensor.randint(*self.shape, low=dtypes.min(self.dtype), high=dtypes.max(self.dtype), device=self.device, dtype=self.dtype),
  "aten.random_.from": lambda self, from_, to: Tensor.randint(*self.shape, low=from_, high=to, device=self.device, dtype=self.dtype),
  "aten.uniform_": lambda self, low=0, high=1: Tensor.uniform(*self.shape, low=low, high=high, dtype=self.dtype),
  "aten.normal_": lambda self, mean=0, std=1: Tensor.normal(*self.shape, mean=mean, std=std, dtype=self.dtype),
@@ -3,7 +3,3 @@ set -e

 ditto -c -k --keepParent ./build/Release/TinyGPU.app ./build/Release/TinyGPU.zip
 xcrun notarytool submit ./build/Release/TinyGPU.zip --keychain-profile "hgwJFhdheiIEy82nDN" --wait
-
-rm ./build/Release/TinyGPU.zip
-xcrun stapler staple ./build/Release/TinyGPU.app
-ditto -c -k --keepParent ./build/Release/TinyGPU.app ./build/Release/TinyGPU.zip
@@ -48,16 +48,11 @@ def decode_profile(data:bytes) -> dict:
        name, ref, key, st, dur, fmt = u("<IIIIfI")
        v["events"].append({"name":strings[name], "ref":option(ref), "key":option(key), "st":st, "dur":dur, "fmt":strings[fmt]})
    else:
-      v["linear"] = u("<B")[0]
      v["peak"] = u("<Q")[0]
      for _ in range(event_count):
-        if v["linear"]:
-          ts, value = u("<IQ")
-          v["events"].append({"event":"freq", "ts":ts, "value":value})
-        else:
-          alloc, ts, key = u("<BII")
-          if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
-          else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IIIB") for _ in range(u("<I")[0])]}})
+        alloc, ts, key = u("<BII")
+        if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
+        else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IIIB") for _ in range(u("<I")[0])]}})
  return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}

 if __name__ == "__main__":
@@ -74,7 +74,7 @@ testing_minimal = [
  "hypothesis>=6.148.9",
  "z3-solver<4.15.4",  # 4.15.4 has a segfault when creating many z3.Context()
 ]
-testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate", "openai", "gguf>=0.18"]
+testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate", "openai", "gguf"]
 testing = [
  "tinygrad[testing_unit]",
  "pillow",
@@ -9,8 +9,8 @@ from tinygrad.renderer.amd import decode_inst
 from tinygrad.runtime.autogen.amd.rdna3.ins import SOPP
 from tinygrad.runtime.autogen.amd.rdna3.enum import SOPPOp
 from tinygrad.renderer.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_RDNA4, WAVEEND, INST, INST_RDNA4, VALUINST,
-                                     IMMEDIATE, IMMEDIATE_MASK, PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4, PACKET_TYPES_CDNA, CDNA_WAVESTART,
-                                     InstOp, InstOpRDNA4, print_packets, CDNA_WAVEEND)
+                                     IMMEDIATE, IMMEDIATE_MASK, PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4,
+                                     InstOp, InstOpRDNA4, print_packets)
 from test.amd.helpers import TARGET_TO_ARCH

 import tinygrad
@@ -21,7 +21,7 @@ OTHER_SIMD_OPS = {InstOp.OTHER_LDS_LOAD, InstOp.OTHER_LDS_STORE, InstOp.OTHER_LD
                  InstOp.OTHER_FLAT_STORE_128, InstOp.OTHER_GLOBAL_LOAD, InstOp.OTHER_GLOBAL_LOAD_VADDR,
                  InstOp.OTHER_GLOBAL_STORE_64, InstOp.OTHER_GLOBAL_STORE_96, InstOp.OTHER_GLOBAL_STORE_128,
                  InstOp.OTHER_GLOBAL_STORE_VADDR_128}
-OTHER_SIMD_OPS_RDNA4 = {InstOpRDNA4.OTHER_VMEM, InstOpRDNA4.OTHER_VMEM_5}
+OTHER_SIMD_OPS_RDNA4 = {InstOpRDNA4.OTHER_VMEM}

 # ═══════════════════════════════════════════════════════════════════════════════
 # ROCPROF DECODER
@@ -125,7 +125,7 @@ class SQTTExamplesTestBase(unittest.TestCase):
          self.assertIsInstance(packets[0], LAYOUT_HEADER, f"first packet should be LAYOUT_HEADER in {name}")

  def test_packet_types_valid(self):
-    all_classes = set(PACKET_TYPES_RDNA3.values()) | set(PACKET_TYPES_RDNA4.values()) | set(PACKET_TYPES_CDNA.values())
+    all_classes = set(PACKET_TYPES_RDNA3.values()) | set(PACKET_TYPES_RDNA4.values())
    for name, (events, *_) in self.examples.items():
      for i, event in enumerate(events):
        with self.subTest(example=name, event=i):
@@ -138,8 +138,8 @@ class SQTTExamplesTestBase(unittest.TestCase):
      if "empty" in name: continue
      with self.subTest(example=name):
        all_packets = [p for e in events for p in decode(e.blob)]
-        self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVESTART, WAVESTART_RDNA4, CDNA_WAVESTART))]), 0, f"no WAVESTART in {name}")
-        self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVEEND, CDNA_WAVEEND))]), 0, f"no WAVEEND in {name}")
+        self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVESTART, WAVESTART_RDNA4))]), 0, f"no WAVESTART in {name}")
+        self.assertGreater(len([p for p in all_packets if isinstance(p, WAVEEND)]), 0, f"no WAVEEND in {name}")

  def test_time_monotonic(self):
    for name, (events, *_) in self.examples.items():
@@ -153,9 +153,7 @@ class SQTTExamplesTestBase(unittest.TestCase):
      if "gemm" not in name: continue
      with self.subTest(example=name):
        all_packets = [p for e in events for p in decode(e.blob)]
-        inst_names = [p.op.name for p in all_packets if isinstance(p, (INST, INST_RDNA4))]
-        self.assertGreater(len(inst_names), 0, f"no INST packets in {name}")
-        self.assertGreater(len([n for n in inst_names if n.startswith("JUMP")]), 0, f"no JUMP packets in {name}")
+        self.assertGreater(len([p for p in all_packets if isinstance(p, (INST, INST_RDNA4))]), 0, f"no INST packets in {name}")

  expected: dict[str, list[int]] = {}  # override in subclasses
  def test_packet_counts(self):
@@ -183,8 +181,8 @@ class SQTTExamplesTestBase(unittest.TestCase):
        for event in events:
          wave_starts: dict[tuple[int, int, int], int] = {}
          for p in decode(event.blob):
-            if isinstance(p, (WAVESTART, CDNA_WAVESTART, WAVESTART_RDNA4)): wave_starts[(p.wave, p.simd, p.cu)] = p._time
-            elif isinstance(p, (WAVEEND, CDNA_WAVEEND)) and (key := (p.wave, p.simd, p.cu)) in wave_starts:
+            if isinstance(p, (WAVESTART, WAVESTART_RDNA4)): wave_starts[(p.wave, p.simd, p.cu)] = p._time
+            elif isinstance(p, WAVEEND) and (key := (p.wave, p.simd, p.cu)) in wave_starts:
              our_waves.append((wave_starts[key], p._time))
        self.assertEqual(sorted(our_waves), sorted(roc_waves), f"wave times mismatch in {name}")

@@ -210,23 +208,17 @@ class SQTTExamplesTestBase(unittest.TestCase):
 class TestSQTTExamplesRDNA3(SQTTExamplesTestBase):
  target = "gfx1100"
  expected = {
-    "profile_empty_run_0": [1974, 1961, 2014, 2065, 2092, 1998],
-    "profile_empty_run_1": [1979, 1972, 2019, 2070, 2097, 2003],
-    "profile_gemm_run_0": [2038, 11076, 2324, 2129, 2156, 2062],
-    "profile_gemm_run_1": [2038, 11037, 2318, 2129, 2156, 2062],
-    "profile_ops_run_0": [2038, 5070, 2078, 2129, 2156, 2062],
-    "profile_ops_run_1": [2038, 5007, 2078, 2129, 2156, 2062],
-    "profile_plus_run_0": [1979, 1979, 2030, 2070, 2097, 2003],
-    "profile_plus_run_1": [1979, 2043, 2030, 2070, 2097, 2003],
+    "profile_empty_run_0": [1744, 1801, 1854, 1890, 1917, 1822],
+    "profile_empty_run_1": [1744, 1801, 1854, 1886, 1921, 1906],
+    "profile_gemm_run_0": [1800, 1867, 1899, 1898, 1914, 1895, 1694, 1779, 1819, 1872, 1877, 1858, 1750, 1834, 1866, 1834, 1911, 1796],
+    "profile_gemm_run_1": [1806, 1874, 1837, 1885, 1907, 1906, 1694, 1778, 1810, 1873, 1885, 1867, 1750, 1834, 1866, 1856, 1903, 1897],
+    "profile_plus_run_0": [1744, 1878, 1854, 1890, 1878, 1910],
+    "profile_plus_run_1": [1744, 1878, 1854, 1886, 1921, 1909],
  }

 class TestSQTTExamplesRDNA4(SQTTExamplesTestBase): target = "gfx1200"
-
-class TestSQTTExamplesCDNA(SQTTExamplesTestBase):
-  target = "gfx950"
-  def test_decode_all_examples(self): self.skipTest("TODO: correct deltas in the timestamp packet types, first packet is REGCS_CDNA")
-  def test_gemm_has_instructions(self): self.skipTest("TODO: decode CDNA inst packets")
-  def test_rocprof_wave_times_match(self): self.skipTest("TODO: requires timestamp patching")
+@unittest.skip("TODO: fix CDNA")
+class TestSQTTExamplesCDNA(SQTTExamplesTestBase): target = "gfx950"

 if __name__ == "__main__":
  unittest.main()
@@ -2,10 +2,9 @@
 import unittest, pickle
 from typing import Iterator
 from pathlib import Path
-from tinygrad.helpers import DEBUG, OSX, getenv, temp
+from tinygrad.helpers import DEBUG, OSX
 from tinygrad.renderer.amd.sqtt import print_packets, map_insts
 from tinygrad.runtime.autogen.amd.rdna3.ins import s_endpgm
-from tinygrad.viz.serve import sqtt_timeline
 from test.amd.disasm import disasm

 import tinygrad
@@ -54,7 +53,7 @@ class TestSQTTMapBase(unittest.TestCase):
  def setUpClass(cls):
    if cls is TestSQTTMapBase: raise unittest.SkipTest("base class")
    cls.examples = {}
-    for pkl_path in ([Path(temp("profile.pkl", append_user=True))] if getenv("LOAD_PROFILE") else sorted((EXAMPLES_DIR/cls.target).glob("*.pkl"))):
+    for pkl_path in sorted((EXAMPLES_DIR/cls.target).glob("*.pkl")):
      with open(pkl_path, "rb") as f:
        data = pickle.load(f)
      sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
@@ -69,33 +68,10 @@ class TestSQTTMapBase(unittest.TestCase):
        if event.kern not in kern_events: continue
        with self.subTest(example=name, kern=event.kern):
          # rocprof OSX has a bug for sopk decoding, linux rocprof works
-          pass_rocprof_err = OSX and target == "gfx1200" and name.startswith("profile_ops")
+          pass_rocprof_err = OSX and target == "gfx1200" and name.startswith("profile_py")
          passed_insts, n_waves, n_units = rocprof_inst_traces_match(event, kern_events[event.kern], target, pass_rocprof_err)
          if n_waves: print(f"{name}: passed for {passed_insts} instructions across {n_waves} waves scheduled on {n_units} wave units")

-  def test_sqtt_timeline(self):
-    for name, (events, kern_events, target) in self.examples.items():
-      for event in events:
-        if (p:=kern_events.get(event.kern)) is None: continue
-        with self.subTest(example=name, kern=event.kern):
-          if not (timeline:=sqtt_timeline(event.blob, p.lib, target)): continue
-          frequency = [e.key for e in timeline if type(e).__name__ == "ProfilePointEvent" and e.name == "freq_hz"]
-          mean = sum(frequency) / len(frequency)
-          variance = sum((v - mean) ** 2 for v in frequency) / len(frequency)
-          self.assertGreater(mean, 0)
-          self.assertGreater(variance, 0)
-          if DEBUG >= 2: print(f"{name:20s} SE:{event.se} {mean/1e9:.2f} GHz mean, {variance/1e18:.2f} GHz^2 variance")
-          events = [e for e in timeline if type(e).__name__ == "ProfileRangeEvent"]
-          insts, execs = 0, 0
-          for e in events:
-            if "EXEC" in e.device:
-              if "ALT" not in e.name.display_name: execs += 1
-            elif "WAVE" in e.device:
-              # sopk/immediates don't get ALU/MEM EXEC
-              if e.name.display_name not in {"IMMEDIATE", "IMMEDIATE_MASK", "JUMP", "JUMP_NO", "MESSAGE"}: insts += 1
-            else: raise Exception(f"timeline row must be INST or EXEC, got {e.device}")
-          self.assertEqual(execs, insts)
-
 class TestSQTTMapRDNA3(TestSQTTMapBase): target = "gfx1100"

 class TestSQTTMapRDNA4(TestSQTTMapBase): target = "gfx1200"
@@ -10,7 +10,7 @@ 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.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX, FP8E4M3FNUZ_MAX, FP8E5M2FNUZ_MAX
+from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX
 import pytest
 pytestmark = pytest.mark.filterwarnings("ignore")

@@ -101,14 +101,14 @@ class TestDType(unittest.TestCase):
  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "skip for now")
  def test_uint_overflow(self):
    if not dtypes.is_unsigned(self.DTYPE): raise unittest.SkipTest("only for unsigned")
-    v = self.DTYPE.max
+    v = dtypes.max(self.DTYPE)
    _test_to_np(Tensor(v, dtype=self.DTYPE)+2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))+2)
    _test_to_np(Tensor(v, dtype=self.DTYPE)*2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))*2)

  def test_dtypes_DTYPES_DICT(self):
    self.assertIn("float", DTYPES_DICT)
    self.assertIn("float32", DTYPES_DICT)
-    self.assertEqual(len(DTYPES_DICT), 28)
+    self.assertEqual(len(DTYPES_DICT), 26)
    self.assertTrue(all(isinstance(value, DType) for value in DTYPES_DICT.values()))
    self.assertTrue(all(issubclass(_to_np_dtype(value), np.generic) for value in DTYPES_DICT.values() if _to_np_dtype(value) is not None))

@@ -143,8 +143,6 @@ def _test_ops(a_dtype:DType, b_dtype:DType, target_dtype=None):
 class TestFp8s(unittest.TestCase):
  def test_fp8e4m3_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e4m3).dtype == dtypes.fp8e4m3
  def test_fp8e5m2_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e5m2).dtype == dtypes.fp8e5m2
-  def test_fp8e4m3fnuz_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e4m3fnuz).dtype == dtypes.fp8e4m3fnuz
-  def test_fp8e5m2fnuz_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e5m2fnuz).dtype == dtypes.fp8e5m2fnuz

 class TestFp8sConversions(unittest.TestCase):
  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3_MAX, max_value=FP8E4M3_MAX))
@@ -171,30 +169,6 @@ class TestFp8sConversions(unittest.TestCase):
  def test_fp8e5m2_to_float(self, x):
    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2).float().item())

-  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3FNUZ_MAX, max_value=FP8E4M3FNUZ_MAX))
-  def test_float_to_fp8e4m3fnuz(self, x):
-    np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.float8_e4m3fnuz).view(torch.uint8).item())
-
-  def test_float_to_fp8e4m3fnuz_extreme_values(self):
-    for x in [FP8E4M3FNUZ_MAX, FP8E4M3FNUZ_MAX*1.01, -FP8E4M3FNUZ_MAX, -FP8E4M3FNUZ_MAX*1.01, math.inf, -math.inf, math.nan, 0.0, -0.0]:
-      np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.float8_e4m3fnuz).view(torch.uint8).item())
-
-  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E5M2FNUZ_MAX, max_value=FP8E5M2FNUZ_MAX))
-  def test_float_to_fp8e5m2fnuz(self, x):
-    np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.float8_e5m2fnuz).view(torch.uint8).item())
-
-  def test_float_to_fp8e5m2fnuz_extreme_values(self):
-    for x in [FP8E5M2FNUZ_MAX, FP8E5M2FNUZ_MAX*1.01, -FP8E5M2FNUZ_MAX, -FP8E5M2FNUZ_MAX*1.01, math.inf, -math.inf, math.nan, 0.0, -0.0]:
-      np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.float8_e5m2fnuz).view(torch.uint8).item())
-
-  @given(strat.integers(min_value=0, max_value=255))
-  def test_fp8e4m3fnuz_to_float(self, x):
-    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e4m3fnuz).float().item())
-
-  @given(strat.integers(min_value=0, max_value=255))
-  def test_fp8e5m2fnuz_to_float(self, x):
-    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2fnuz).float().item())
-
@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "bfloat16 not supported")
 class TestBFloat16(unittest.TestCase):
  def test_bf16_creation_numpy(self):
@@ -516,3 +490,4 @@ class TestOpsBFloat16(unittest.TestCase):

 if __name__ == '__main__':
  unittest.main()
+
@@ -52,8 +52,6 @@ class ht:
 ht.bfloat16 = ht.uint16.filter(lambda x: ((x >> 7) & 0xFF) != 0)  # filter subnormal bfloat16
 ht.fp8e4m3 = ht.uint8
 ht.fp8e5m2 = ht.uint8
-ht.fp8e4m3fnuz = ht.uint8
-ht.fp8e5m2fnuz = ht.uint8

 def universal_test(a, b, dtype, op):
  if not isinstance(op, tuple): op = (op, op)
@@ -69,8 +67,7 @@ def universal_test(a, b, dtype, op):
    if not is_dtype_supported(dtype) or dtype in EMULATED_DTYPES.tolist(dtypes): # denormals are zero
      fe, fm = dtypes.finfo(dtype)
      atol, rtol = 2 ** (2 - (1 << (fe - 1))), 2 ** (-fm)
-    else: atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1),
-                        dtypes.fp8e4m3fnuz:(1e-1, 1e-1), dtypes.fp8e5m2fnuz:(5e-1, 5e-1)}.get(dtype, (1e-10, 1e-7))
+    else: atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1)}.get(dtype, (1e-10, 1e-7))
    np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
  else: np.testing.assert_equal(tensor_value, numpy_value)

@@ -90,8 +87,7 @@ def universal_test_unary(a, dtype, op):
  else: tensor_value, numpy_value = op[0](ta).numpy(), op[1](ta.numpy())
  if dtype in dtypes.floats:
    atol, rtol = { dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 2e-2),
-      dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2: (1.0, 5e-1),
-      dtypes.fp8e4m3fnuz:(1e-1, 1e-1), dtypes.fp8e5m2fnuz: (5e-1, 5e-1)}.get(dtype, (1e-6, 1e-5))
+      dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2: (1.0, 5e-1)}.get(dtype, (1e-6, 1e-5))
    np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
  else: np.testing.assert_equal(tensor_value, numpy_value)

@@ -159,26 +155,6 @@ class TestDTypeALU(unittest.TestCase):
  def test_emulated_fp8e5m2(self, a, b, op):
    universal_test(from_storage_scalar(a, dtypes.fp8e5m2), from_storage_scalar(b, dtypes.fp8e5m2), dtypes.fp8e5m2, op)

-  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3fnuz), f"no fp8e4m3fnuz on {Device.DEFAULT}")
-  @given(ht.fp8e4m3fnuz, ht.fp8e4m3fnuz, strat.sampled_from(binary_operations))
-  def test_fp8e4m3fnuz(self, a, b, op):
-    universal_test(from_storage_scalar(a, dtypes.fp8e4m3fnuz), from_storage_scalar(b, dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
-
-  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2fnuz), f"no fp8e5m2fnuz on {Device.DEFAULT}")
-  @given(ht.fp8e5m2fnuz, ht.fp8e5m2fnuz, strat.sampled_from(binary_operations))
-  def test_fp8e5m2fnuz(self, a, b, op):
-    universal_test(from_storage_scalar(a, dtypes.fp8e5m2fnuz), from_storage_scalar(b, dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
-
-  @given(ht.fp8e4m3fnuz, ht.fp8e4m3fnuz, strat.sampled_from(binary_operations))
-  @Context(EMULATED_DTYPES="fp8e4m3fnuz")
-  def test_emulated_fp8e4m3fnuz(self, a, b, op):
-    universal_test(from_storage_scalar(a, dtypes.fp8e4m3fnuz), from_storage_scalar(b, dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
-
-  @given(ht.fp8e5m2fnuz, ht.fp8e5m2fnuz, strat.sampled_from(binary_operations))
-  @Context(EMULATED_DTYPES="fp8e5m2fnuz")
-  def test_emulated_fp8e5m2fnuz(self, a, b, op):
-    universal_test(from_storage_scalar(a, dtypes.fp8e5m2fnuz), from_storage_scalar(b, dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
-
  @given(ht.float32, strat.sampled_from(unary_operations))
  def test_float32_unary(self, a, op): universal_test_unary(a, dtypes.float32, op)

@@ -222,30 +198,6 @@ class TestDTypeALU(unittest.TestCase):
    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2) != 0.0)
    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2), dtypes.fp8e5m2, op)

-  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3fnuz), f"no fp8e4m3fnuz on {Device.DEFAULT}")
-  @given(ht.fp8e4m3fnuz, strat.sampled_from(unary_operations))
-  def test_fp8e4m3fnuz_unary(self, a, op):
-    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz) != 0.0)
-    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
-
-  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2fnuz), f"no fp8e5m2fnuz on {Device.DEFAULT}")
-  @given(ht.fp8e5m2fnuz, strat.sampled_from(unary_operations))
-  def test_fp8e5m2fnuz_unary(self, a, op):
-    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz) != 0.0)
-    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
-
-  @given(ht.fp8e4m3fnuz, strat.sampled_from(unary_operations))
-  @Context(EMULATED_DTYPES="fp8e4m3fnuz")
-  def test_emulated_fp8e4m3fnuz_unary(self, a, op):
-    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz) != 0.0)
-    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
-
-  @given(ht.fp8e5m2fnuz, strat.sampled_from(unary_operations))
-  @Context(EMULATED_DTYPES="fp8e5m2fnuz")
-  def test_emulated_fp8e5m2fnuz_unary(self, a, op):
-    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz) != 0.0)
-    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
-
  @given(ht.uint8, ht.uint8, strat.sampled_from(integer_binary_operations))
  def test_uint8(self, a, b, op): universal_test(a, b, dtypes.uint8, op)

@@ -366,7 +318,7 @@ class TestDTypeALU(unittest.TestCase):

  @unittest.expectedFailure
  def test_unsafe_cast_float_to_int_failure(self):
-    val = float(dtypes.int32.max - 1)
+    val = float(dtypes.max(dtypes.int32) - 1)
    t1 = Tensor([val], dtype=dtypes.float32).cast(dtypes.int32)
    t2 = Tensor(val, dtype=dtypes.float32).cast(dtypes.int32)
    np.testing.assert_equal(t1.item(), t2.item())
@@ -15,7 +15,7 @@ from test.helpers import needs_second_gpu
 np.random.seed(1337)
 Tensor.manual_seed(1337)
 BUF_SIZE = 4096
-RUN_CNT = 5
+RUN_CNT = 4

 cached_prgs = {}
 def helper_exec_op(device, outbuf, inbufs):
@@ -47,17 +47,6 @@ def helper_create_offset_rawbuffer(base, offset=0):
  x = Buffer(base.device, base.size-offset, base.dtype, base=base, offset=offset)
  return x.ensure_allocated()

-def helper_alloc_rawbuffer_sized(device, size, fill=False):
-  rawbuf = Buffer(device, size, dtypes.int).ensure_allocated()
-  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())
-  return rawbuf
-
-def helper_make_view(base, offset_elems, size_elems):
-  return Buffer(base.device, size_elems, base.dtype, base=base, offset=offset_elems * base.dtype.itemsize).ensure_allocated()
-
 def helper_run_jit(jis, bufs, out_buffers):
  for rawbuf in out_buffers:
    mv = memoryview(bytearray(rawbuf.size * rawbuf.dtype.itemsize))
@@ -91,14 +80,6 @@ def helper_test_graphs(graph_impl, graphs, runs=RUN_CNT):

@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")
-
-  def skip_if_not_multigraph(self):
-    graph = g.func if isinstance(g:=(d:=Device[Device.DEFAULT]).graph, functools.partial) else g
-    if not issubclass(graph, MultiGraphRunner): self.skipTest("graph is not supported (not MultiGraphRunner)")
-    if not hasattr(d.allocator, '_transfer') or not d.allocator.supports_transfer: self.skipTest("device is not supported (no transfers)")
-
  def test_order_2_writes_to_same_buf(self):
    d0 = Device.DEFAULT
    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(5)]
@@ -129,6 +110,11 @@ class TestGraph(unittest.TestCase):

    helper_test_graphs(Device[d0].graph, graphs)

+  def skip_if_not_multigraph(self):
+    graph = g.func if isinstance(g:=(d:=Device[Device.DEFAULT]).graph, functools.partial) else g
+    if not issubclass(graph, MultiGraphRunner): self.skipTest("graph is not supported (not MultiGraphRunner)")
+    if not hasattr(d.allocator, '_transfer') or not d.allocator.supports_transfer: self.skipTest("device is not supported (no transfers)")
+
  def test_order_copy_writed(self):
    self.skip_if_not_multigraph()

@@ -279,58 +265,5 @@ class TestGraph(unittest.TestCase):

    helper_test_graphs(Device[d0].graph, graphs)

-  def test_partial_write_preserves_write_dep(self):
-    self.skip_if_not_multigraph()
-    self.skip_if_no_offset()
-    d0 = Device.DEFAULT
-
-    base = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 2, fill=True)
-    copy_src_full = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 2, fill=True)
-    copy_src_lo = helper_alloc_rawbuffer(d0, fill=True)
-    v_lo = helper_make_view(base, 0, BUF_SIZE)
-    v_hi = helper_make_view(base, BUF_SIZE, BUF_SIZE)
-    a, c = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(2)]
-
-    graphs = [
-      [helper_copy_op(d0, base, copy_src_full), helper_copy_op(d0, v_lo, copy_src_lo), helper_exec_op(d0, c, [v_hi, a])]
-    ]
-    helper_test_graphs(Device[d0].graph, graphs)
-
-  def test_partial_write_preserves_read_dep(self):
-    self.skip_if_not_multigraph()
-    self.skip_if_no_offset()
-    d0 = Device.DEFAULT
-
-    base = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 2, fill=True)
-    copy_dst = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 2, fill=True)
-    copy_src_lo = helper_alloc_rawbuffer(d0, fill=True)
-    v_lo = helper_make_view(base, 0, BUF_SIZE)
-    v_hi = helper_make_view(base, BUF_SIZE, BUF_SIZE)
-    a, b = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(2)]
-
-    graphs = [
-      [helper_copy_op(d0, copy_dst, base), helper_copy_op(d0, v_lo, copy_src_lo), helper_exec_op(d0, v_hi, [a, b])]
-    ]
-    helper_test_graphs(Device[d0].graph, graphs)
-
-  def test_middle_write_splits_write_dep(self):
-    self.skip_if_not_multigraph()
-    self.skip_if_no_offset()
-    d0 = Device.DEFAULT
-
-    base = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 3, fill=True)
-    copy_src_full = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 3, fill=True)
-    copy_src_mid = helper_alloc_rawbuffer(d0, fill=True)
-    v_lo = helper_make_view(base, 0, BUF_SIZE)
-    v_mid = helper_make_view(base, BUF_SIZE, BUF_SIZE)
-    v_hi = helper_make_view(base, BUF_SIZE * 2, BUF_SIZE)
-    a, c, e = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(3)]
-
-    graphs = [
-      [helper_copy_op(d0, base, copy_src_full), helper_copy_op(d0, v_mid, copy_src_mid),
-       helper_exec_op(d0, c, [v_lo, a]), helper_exec_op(d0, e, [v_hi, a])]
-    ]
-    helper_test_graphs(Device[d0].graph, graphs)
-
 if __name__ == '__main__':
  unittest.main()
@@ -6,7 +6,6 @@ from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, CPU_LLVM, AMD_LL
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
-from tinygrad.renderer.cstyle import QCOMCLRenderer
 from tinygrad.renderer.nir import NIRRenderer

 TINY_BACKEND = getenv("TINY_BACKEND")
@@ -437,7 +436,7 @@ class TestOps(unittest.TestCase):
    helper_test_op([(45,35), (45,35), (45,35)], lambda x,y,z: x.lerp(y,z))
    helper_test_op(None, lambda x,y,z: x.lerp(y,z), vals=[[1.,2.,3.], [4.,5.,6.], 0.5])

-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
+  @unittest.skipIf(Device.DEFAULT == "QCOM", "OpenCL fails to compile this (both on GPU(qcom)/QCOM backends)")
  def test_tril(self):
    helper_test_op([(3,3)], lambda x: x.tril())
    helper_test_op([(3,3)], lambda x: x.tril(1))
@@ -455,7 +454,7 @@ class TestOps(unittest.TestCase):
    helper_test_op([(5,3,3)], lambda x: x.tril(1))
    helper_test_op(None, lambda x: x.tril(), vals=[[[True] * 3] * 3], forward_only=True)

-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
+  @unittest.skipIf(Device.DEFAULT == "QCOM", "OpenCL fails to compile this (both on GPU(qcom)/QCOM backends)")
  def test_triu(self):
    helper_test_op([(3,3)], lambda x: x.triu())
    helper_test_op([(3,3)], lambda x: x.triu(1))
@@ -479,9 +478,9 @@ class TestOps(unittest.TestCase):
    helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[1., 0., 3., -4.], 3.])
    helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[1., 0., 3., -4.], [-1., -2., 3., 0.]])
    helper_test_op(None, torch.maximum, Tensor.maximum,
-                   vals=[[-1234, 0, 1234, dtypes.int.max, dtypes.int.min], dtypes.int.max], forward_only=True)
+                   vals=[[-1234, 0, 1234, dtypes.max(dtypes.int), dtypes.min(dtypes.int)], dtypes.max(dtypes.int)], forward_only=True)
    helper_test_op(None, torch.maximum, Tensor.maximum,
-                   vals=[[-1234, 0, 1234, dtypes.int.max, dtypes.int.min], dtypes.int.min], forward_only=True)
+                   vals=[[-1234, 0, 1234, dtypes.max(dtypes.int), dtypes.min(dtypes.int)], dtypes.min(dtypes.int)], forward_only=True)
    helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[True, False, False], True], forward_only=True)
    helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[True, False, False], [True, True, False]], forward_only=True)

@@ -496,9 +495,9 @@ class TestOps(unittest.TestCase):
    helper_test_op(None, torch.minimum, Tensor.minimum, vals=[[1., 0., 3., -4.], 3.])
    helper_test_op(None, torch.minimum, Tensor.minimum, vals=[[1., 0., 3., -4.], [-1., -2., 3., 0.]])
    helper_test_op(None, torch.minimum, Tensor.minimum,
-                   vals=[[-1234, 0, 1234, dtypes.int.max, dtypes.int.min], dtypes.int.max], forward_only=True)
+                   vals=[[-1234, 0, 1234, dtypes.max(dtypes.int), dtypes.min(dtypes.int)], dtypes.max(dtypes.int)], forward_only=True)
    helper_test_op(None, torch.minimum, Tensor.minimum,
-                   vals=[[-1234, 0, 1234, dtypes.int.max, dtypes.int.min], dtypes.int.min], forward_only=True)
+                   vals=[[-1234, 0, 1234, dtypes.max(dtypes.int), dtypes.min(dtypes.int)], dtypes.min(dtypes.int)], forward_only=True)
    helper_test_op(None, torch.minimum, Tensor.minimum, vals=[[True, False, False], True], forward_only=True)
    helper_test_op(None, torch.minimum, Tensor.minimum, vals=[[True, False, False], [True, True, False]], forward_only=True)

@@ -766,7 +765,6 @@ class TestOps(unittest.TestCase):

    self.helper_test_exception([(4), (4)], lambda x,y: x.bitwise_xor(y), expected=RuntimeError)

-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
  def test_and(self):
    data = [[1,-8,1],[32,1,6]]
    tor = torch.tensor(data, dtype=torch.int)
@@ -784,7 +782,6 @@ class TestOps(unittest.TestCase):

    self.helper_test_exception([(4), (4)], lambda x,y: x.bitwise_and(y), expected=RuntimeError)

-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
  def test_or(self):
    data = [[1,-8,1],[32,1,6]]
    tor = torch.tensor(data, dtype=torch.int)
@@ -1173,7 +1170,6 @@ class TestOps(unittest.TestCase):
    helper_test_op(None, lambda x: x.type(torch.int32).argmax().type(torch.int32), lambda x: x.argmax(), forward_only=True, vals=[[False, True]])
    helper_test_op(None, lambda x: x.type(torch.int32).argmax().type(torch.int32), lambda x: x.argmax(), forward_only=True, vals=[[True, False]])

-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
  def test_argmin(self):
    # check if it returns the first index for multiple occurrences
    helper_test_op(None, lambda x: x.argmin().type(torch.int32), lambda x: x.argmin(), forward_only=True, vals=[[2, 2]])
@@ -1479,7 +1475,6 @@ class TestOps(unittest.TestCase):
  def test_prod_dtype_arg(self):
    with self.assertRaises(AttributeError): Tensor([1.0, 2.0]).prod(dtype="")

-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
  def test_min(self):
    helper_test_op([(3,3)], lambda x: x.min())
    helper_test_op([(45,3)], lambda x: x.min())
@@ -1508,6 +1503,7 @@ class TestOps(unittest.TestCase):
    helper_test_op([(3,3)], lambda x: torch.full_like(x, 2).prod(), lambda x: (x.full_like(2)).prod(), forward_only=True)
    helper_test_op([(3,3)], lambda x: torch.full_like(x, 2).max(), lambda x: (x.full_like(2)).max(), forward_only=True)

+  @unittest.skipIf(Device.DEFAULT == "QCOM", "OpenCL fails to compile this (both on GPU(qcom)/QCOM backends)")
  def test_any(self):
    helper_test_op([(3,4,5,6)], lambda x: x.any(), forward_only=True)
    helper_test_op(None, lambda x: x.any(), vals=[[True, True]], forward_only=True)
@@ -1519,7 +1515,7 @@ class TestOps(unittest.TestCase):
  def test_any_zero_axis(self):
    helper_test_op([(1,0,3,0,5)], lambda x: x.any(axis=(1,3)), forward_only=True)

-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
+  @unittest.skipIf(Device.DEFAULT == "QCOM", "OpenCL fails to compile this (both on GPU(qcom)/QCOM backends)")
  def test_all(self):
    helper_test_op([(3,4,5,6)], lambda x: x.all(), forward_only=True)
    helper_test_op(None, lambda x: x.all(), vals=[[True, True]], forward_only=True)
@@ -1669,15 +1665,6 @@ class TestOps(unittest.TestCase):
    helper_test_op([(10,10,10)], lambda x: x.log_softmax(1), atol=1e-7, grad_atol=1e-7)
    helper_test_op([(10,10,10)], lambda x: x.log_softmax(2), atol=1e-7, grad_atol=1e-7)

-  def test_normalize(self):
-    helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x), lambda x: x.normalize(), atol=1e-7, grad_atol=1e-7)
-    helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x, dim=0), lambda x: x.normalize(dim=0), atol=1e-7, grad_atol=1e-7)
-    helper_test_op([(10,10,10)], lambda x: torch.nn.functional.normalize(x, dim=2), lambda x: x.normalize(dim=2), atol=1e-7, grad_atol=1e-7)
-    helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x, p=1), lambda x: x.normalize(p=1), atol=1e-7, grad_atol=1e-7)
-    helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x, p=3, dim=0), lambda x: x.normalize(p=3, dim=0), atol=1e-7, grad_atol=1e-7)
-    helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x, p=0), lambda x: x.normalize(p=0), atol=1e-7, grad_atol=1e-7)
-    helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x, p=-1), lambda x: x.normalize(p=-1), atol=1e-7, grad_atol=1e-7)
-
  def test_logsumexp(self):
    helper_test_op([(45,65)], lambda x: torch.logsumexp(x, dim=0), lambda x: x.logsumexp(0), atol=1e-7, grad_atol=1e-7)
    helper_test_op([(45,65)], lambda x: torch.logsumexp(x, dim=0, keepdim=True), lambda x: x.logsumexp(0, True), atol=1e-7, grad_atol=1e-7)
@@ -2893,7 +2880,6 @@ class TestOps(unittest.TestCase):
    helper_test_op([(2,5,6,5,3,4)], lambda x: x[...,c,:,e], lambda x: x[...,k,:,p])

  @slow_test
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
  def test_slice_fancy_indexing_dim_collapse_int(self):
    a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
    # dim collapse from int
@@ -2904,7 +2890,6 @@ class TestOps(unittest.TestCase):
    helper_test_op([(2,5,6,5,3,4)], lambda x: x[1,:,3:11:2,d,0:2], lambda x: x[1,:,3:11:2,o,0:2])

  @slow_test
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
  def test_slice_fancy_indexing_dim_inject_none(self):
    a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
    # dim injection from None
@@ -2939,7 +2924,6 @@ class TestOps(unittest.TestCase):
                            lambda x: x[Tensor([[0,1,-1],[-1,-2,0]]), Tensor([2,1,-1])])

  @slow_test
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
  def test_slice_fancy_indexing_list_indices(self):
    a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
    helper_test_op([(2,5,6,5,3,4)], lambda x: x[((0,),)])
@@ -2951,7 +2935,6 @@ class TestOps(unittest.TestCase):
    helper_test_op([(2,5,6,5,3,4)], lambda x: x[a,(2,1,0),c,(-2,1,0),e], lambda x: x[i,(2,1,0),k,(-2,1,0),p])

  @slow_test
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, QCOMCLRenderer), "QCOM CL vectorized bool bug")
  def test_slice_fancy_indexing_tuple_indices(self):
    a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
    helper_test_op([(2,5,6,5,3,4)], lambda x: x[(((0,),),)], lambda x: x[(((0,),),)])
@@ -3293,6 +3276,7 @@ class TestOps(unittest.TestCase):
    helper_test_op([(20,)], lambda x: (x>0.5).nonzero().int(), lambda x: (x>0.5).nonzero(), forward_only=True)
    helper_test_op([(10, 5, 3)], lambda x: (x>0.5).nonzero().int(), lambda x: (x>0.5).nonzero(), forward_only=True)

+  @unittest.skipIf(Device.DEFAULT == "QCOM", "OpenCL fails to compile this (both on GPU(qcom)/QCOM backends)")
  def test_cast(self):
    helper_test_op([(3, 3)], lambda x: x.float())
    helper_test_op(None, lambda x: x.float(), vals=[[0, 1, 2, 3]], forward_only=True)
@@ -204,13 +204,13 @@ class TestQuantizeOnnx(unittest.TestCase):
    W = Tensor(m2:=(np.random.uniform(0, 255, size=(N,N)).astype(wi))).realize()
    tg_dtype = dtypes.int8 if xi == np.int8 else dtypes.uint8
    out = (X.int().matmul(W.int())//1000)
-    if clip: out = out.clip(tg_dtype.min, tg_dtype.max)
+    if clip: out = out.clip(dtypes.min(tg_dtype),dtypes.max(tg_dtype))
    out = out.cast(tg_dtype)
    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)] if opts is None else opts
    sexec(out, opts, replace_src, run_count=1)
    tout = out.numpy()
    mout = ((m1.astype(np.int32) @ m2.astype(np.int32)) // 1000)
-    if clip: mout = mout.clip(tg_dtype.min, tg_dtype.max)
+    if clip: mout = mout.clip(dtypes.min(tg_dtype),dtypes.max(tg_dtype))
    mout = mout.astype(xi)
    print(tout)
    print(mout)
@@ -62,7 +62,7 @@ class TestRendererFailures(unittest.TestCase):
 class TestCStyleFailures(unittest.TestCase):
  def test_inline_const_alu(self):
    # CPU doesn't use the max function
-    ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int, dtypes.int.min+1))
+    ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int, dtypes.min(dtypes.int)+1))
    self.assertEqual(ret[0], 1)

  def _test_src_strip_paren(self, op: Ops, should_strip_paren:bool=True):
@@ -797,7 +797,7 @@ class TestSchedule(unittest.TestCase):

  @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
  def test_image_dot_f16_fusion(self):
-    with Context(FLOAT16=1, OPENPILOT_HACKS=1):
+    with Context(FLOAT16=1):
      def cnt():
        x, y, z = Tensor.empty((64, 64), dtype='float'), Tensor.empty((64, 64), dtype='float'), Tensor.empty((64, 64), dtype='float')
        a = (x @ y).relu()
@@ -811,42 +811,21 @@ class TestSchedule(unittest.TestCase):
      self.assertEqual(cnt1, 5)
      self.assertEqual(cnt2, 5)

-  @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
-  def test_image_f16_residual_fusion(self):
-    with Context(FLOAT16=1, OPENPILOT_HACKS=1):
-      def cnt():
-        inp = Tensor.empty((512,), dtype='float')
-        b1, b2 = Tensor.empty((512, 1024), dtype='float'), Tensor.empty((1024, 512), dtype='float')
-        c1, c2 = Tensor.empty((1024,), dtype='float'), Tensor.empty((512,), dtype='float')
-        rb = (((((inp @ b1) + c1).relu() @ b2) + c2).relu() + inp).relu()
-        b16, c16 = Tensor.empty((512, 16), dtype='float'), Tensor.empty((16,), dtype='float')
-        b32, c32 = Tensor.empty((512, 32), dtype='float'), Tensor.empty((32,), dtype='float')
-        sched = Tensor.schedule((rb @ b16 + c16).relu(), (rb @ b32 + c32).relu())
-        for si in sched: si.lower()
-        return len([si for si in sched if isinstance(si.prg, CompiledRunner)])
-
-      with Context(IMAGE=1): cnt1 = cnt()
-      with Context(IMAGE=2): cnt2 = cnt()
-
-      self.assertEqual(cnt1, 9)
-      self.assertEqual(cnt2, 9)
-
  @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
  @unittest.expectedFailure
  def test_image_conv_fusion(self):
-    with Context(OPENPILOT_HACKS=1):
-      def cnt():
-        x, y, z = Tensor.empty((1, 4, 3, 3)), Tensor.empty((4, 1, 3, 3)), Tensor.empty((4, 1, 7, 7))
-        a = x.conv2d(y, Tensor.empty(4), groups=4, padding=1)
-        b = a.conv2d(z, groups=4, padding=3)
-        sched = (a + b).schedule()
-        for si in sched: si.lower()
-        return len([si for si in sched if isinstance(si.prg, CompiledRunner)])
+    def cnt():
+      x, y, z = Tensor.empty((1, 4, 3, 3)), Tensor.empty((4, 1, 3, 3)), Tensor.empty((4, 1, 7, 7))
+      a = x.conv2d(y, Tensor.empty(4), groups=4, padding=1)
+      b = a.conv2d(z, groups=4, padding=3)
+      sched = (a + b).schedule()
+      for si in sched: si.lower()
+      return len([si for si in sched if isinstance(si.prg, CompiledRunner)])

-      with Context(IMAGE=1): cnt1 = cnt()
-      with Context(IMAGE=2): cnt2 = cnt()
+    with Context(IMAGE=1): cnt1 = cnt()
+    with Context(IMAGE=2): cnt2 = cnt()

-      self.assertEqual(cnt1, cnt2)
+    self.assertEqual(cnt1, cnt2)

  def _test_fusion(self, shapes, f, cnt):
    with Context(DEBUG=0, TRACK_MATCH_STATS=0): args = [Tensor.randn(s).realize() for s in shapes]
@@ -251,11 +251,8 @@ class TestSetitem(unittest.TestCase):
    s1 = t.sum()
    t[3:].assign(2.0)
    s2 = t.sum()
-    try:
-      np.testing.assert_allclose([s0.item(), s1.item(), s2.item()], [0.0, 3.0, 9.0])
-    except AssertionError:
-      # TODO: broken now, lazy sums all see final buffer state
-      np.testing.assert_allclose([s0.item(), s1.item(), s2.item()], [9.0, 9.0, 9.0])
+    # TODO: s0 and s1 see final buffer state, should be [0.0, 3.0, 9.0]
+    np.testing.assert_allclose([s0.item(), s1.item(), s2.item()], [9.0, 9.0, 9.0])

    # eager version
    t = Tensor.zeros(6).contiguous().realize()
@@ -276,11 +273,8 @@ class TestSetitem(unittest.TestCase):
    a.assign(new_a)
    b.assign(new_b)
    np.testing.assert_allclose(a.numpy(), [4, 6, 8, 10])
-    try:
-      np.testing.assert_allclose(b.numpy(), [0, 2, 4, 6])
-    except AssertionError:
-      # TODO: broken now, new_b sees mutated a
-      np.testing.assert_allclose(b.numpy(), [8, 12, 16, 20])
+    # TODO: new_b sees mutated a, should be [0, 2, 4, 6]
+    np.testing.assert_allclose(b.numpy(), [8, 12, 16, 20])

    # eager version
    a = Tensor.arange(4, dtype=dtypes.float).contiguous().realize()
@@ -271,7 +271,6 @@ class SDMAExecutor(AMDQueue):
      elif op == amd_gpu.SDMA_OP_GCR: self._execute_gcr()
      elif op == amd_gpu.SDMA_OP_COPY: self._execute_copy()
      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}")
    return self.rptr[0] - prev_rptr

@@ -75,20 +75,20 @@ class TestHelpers(unittest.TestCase):
  def test_dtype_range(self):
    for dt in core_dtypes:
      if dtypes.is_float(dt):
-        np.testing.assert_equal(dt.min, -math.inf)
-        np.testing.assert_equal(dt.max, math.inf)
+        np.testing.assert_equal(dtypes.min(dt), -math.inf)
+        np.testing.assert_equal(dtypes.max(dt), math.inf)
        np.testing.assert_equal(dt.min, -math.inf)
        np.testing.assert_equal(dt.max, math.inf)
      elif dtypes.is_int(dt):
        info = np.iinfo(_to_np_dtype(dt))
-        np.testing.assert_equal(dt.min, info.min)
-        np.testing.assert_equal(dt.max, info.max)
+        np.testing.assert_equal(dtypes.min(dt), info.min)
+        np.testing.assert_equal(dtypes.max(dt), info.max)
        np.testing.assert_equal(dt.min, info.min)
        np.testing.assert_equal(dt.max, info.max)
      else:
        assert dt == dtypes.bool, dt
-        np.testing.assert_equal(dt.min, False)
-        np.testing.assert_equal(dt.max, True)
+        np.testing.assert_equal(dtypes.min(dt), False)
+        np.testing.assert_equal(dtypes.max(dt), True)
        np.testing.assert_equal(dt.min, False)
        np.testing.assert_equal(dt.max, True)

@@ -14,7 +14,6 @@ class TestLLMServer(unittest.TestCase):

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

    cls.bos_id = 1
    cls.eos_id = 999
@@ -11,8 +11,8 @@ def b(i, base=None, offset=0, pin=False, size=16):
  if pin: global_map[i].ref(1)
  return global_map[i]

-def check_assign(buffers:list[list[Buffer]|tuple[Buffer, ...]], copies:list[tuple[Buffer, Buffer]]|None=None):
-  assigned = _internal_memory_planner(buffers, copies=copies)
+def check_assign(buffers:list[list[Buffer]|tuple[Buffer, ...]]):
+  assigned = _internal_memory_planner(buffers, noopt_buffers=None)

  taken_parts = set()
  first_appearance, last_appearance = {}, {}
@@ -134,75 +134,5 @@ class TestMemoryPlanner(unittest.TestCase):
    ]
    check_assign(bs)

-  def test_copy_bufs_separate_from_compute(self):
-    bs = [
-      [b(0), b(1)],
-      [b(1), b(2)],
-      [b(3), b(2)],
-    ]
-    assigned = _internal_memory_planner(bs, copies=[(b(1), b(0))])
-    r1, r2 = assigned.get(b(1), b(1)), assigned.get(b(2), b(2))
-    assert r1.base != r2.base
-
-  def test_copy_bufs_reuse_among_copies(self):
-    bs = [
-      [b(0), b(1)],
-      [b(2), b(1)],
-      [b(3), b(2)],
-    ]
-    assigned = _internal_memory_planner(bs, copies=[(b(1), b(0)), (b(2), b(1))])
-    r1, r2 = assigned.get(b(1), b(1)), assigned.get(b(2), b(2))
-    assert r1.base == r2.base
-
-  def test_compute_bufs_reuse_among_compute(self):
-    bs = [
-      [b(0), b(1)],
-      [b(2), b(1)],
-      [b(3), b(2)],
-      [b(4), b(3)],
-    ]
-    assigned = _internal_memory_planner(bs, copies=[(b(1), b(0))])
-    r2, r3 = assigned.get(b(2), b(2)), assigned.get(b(3), b(3))
-    assert r2.base == r3.base
-
-  def test_copy_and_compute_no_cross_reuse(self):
-    bs = [
-      [b(0), b(1)],
-      [b(2), b(1)],
-      [b(3), b(2)],
-    ]
-    assigned = _internal_memory_planner(bs, copies=[(b(2), b(1))])
-    r0, r2 = assigned.get(b(0), b(0)), assigned.get(b(2), b(2))
-    assert r0.base != r2.base
-
-  def test_multiple_copy_bufs_with_offsets(self):
-    bs = [
-      [b(0, pin=True), b(1), b(2)],
-      [b(3, base=0, offset=1, size=8), b(1), b(2)],
-      [b(4), b(3)],
-      [b(5), b(4)],
-    ]
-    check_assign(bs, copies=[(b(1), b(0)), (b(2), b(0))])
-
-  def test_copy_bufs_pinned_mixed(self):
-    bs = [
-      [b(0, pin=True), b(1), b(2)],
-      [b(1), b(3), b(2)],
-      [b(4), b(3)],
-      [b(5), b(4), b(0)],
-    ]
-    check_assign(bs, copies=[(b(1), b(0)), (b(3), b(1))])
-
-  def test_deferred_copy_frees_chain(self):
-    bs = []
-    copies = []
-    for i in range(6):
-      copy_buf, compute_buf = b(i * 2 + 1), b(i * 2 + 2)
-      bs.append([copy_buf, b(0, pin=True)])
-      bs.append([compute_buf, copy_buf])
-      copies.append((copy_buf, b(0, pin=True)))
-    bs.append([b(100, pin=True)])
-    check_assign(bs, copies=copies)
-
 if __name__ == "__main__":
  unittest.main()
@@ -1,6 +1,5 @@
 import gc, unittest
 from tinygrad import Tensor, GlobalCounters, dtypes
-from tinygrad.engine.jit import TinyJit

 class TestMultiRamUsage(unittest.TestCase):
  def setUp(self):
@@ -108,37 +107,6 @@ class TestMultiRamUsage(unittest.TestCase):
  def test_matmul_half(self): self._test_matmul_half(dev_count=2)
  def test_matmul_half_alt(self): self._test_matmul_half(dev_count=4)

-  def test_multi_layer_allreduce(self):
-    N = 32
-    devices_2 = ("NULL:1", "NULL:2")
-
-    def make_inp():
-      x = Tensor.zeros(N, N).contiguous().shard(devices_2, axis=None).realize()
-      w1 = Tensor.zeros(N, N).contiguous().shard(devices_2, axis=1).realize()
-      w2 = Tensor.zeros(N, N).contiguous().shard(devices_2, axis=0).realize()
-      return x, w1, w2
-
-    def run_layers(n_layers):
-      GlobalCounters.reset()
-
-      @TinyJit
-      def f(x, w1, w2):
-        for _ in range(n_layers):
-          x = (x @ w1 @ w2)
-        return x.contiguous()
-
-      for _ in range(3):
-        a = make_inp()
-        r = f(*a)
-        del a, r
-
-      gc.collect()
-      return GlobalCounters.mem_used
-
-    mem_2 = run_layers(2)
-    mem_4 = run_layers(4)
-    self.assertEqual(mem_2, mem_4, f"graph memory should not grow with layers: 2 layers={mem_2}, 4 layers={mem_4}")
-
 class TestMultiAxis(unittest.TestCase):
  def test_reshape_shard_invalid(self):
    devices = ("NULL:0", "NULL:1")
@@ -580,22 +580,21 @@ class TestSchedule(unittest.TestCase):

  # this is the failing case in openpilot...it's very simple like this
  def test_image_conv_fusion(self):
-    with Context(OPENPILOT_HACKS=1):
-      w1 = Tensor.empty(16, 16, 1, 1)
-      b1 = Tensor.empty(16)
-      w2 = Tensor.empty(16, 16, 1, 1)
-      b2 = Tensor.empty(16)
-      w3 = Tensor.empty(16, 16, 1, 1)
-      b3 = Tensor.empty(16)
+    w1 = Tensor.empty(16, 16, 1, 1)
+    b1 = Tensor.empty(16)
+    w2 = Tensor.empty(16, 16, 1, 1)
+    b2 = Tensor.empty(16)
+    w3 = Tensor.empty(16, 16, 1, 1)
+    b3 = Tensor.empty(16)

-      x = Tensor.empty(1, 16, 32, 32)
-      x = base = x.image_conv2d(w1, b1)
-      x = x.image_conv2d(w2, b2) + base
-      x = x.image_conv2d(w3, b3)
+    x = Tensor.empty(1, 16, 32, 32)
+    x = base = x.image_conv2d(w1, b1)
+    x = x.image_conv2d(w2, b2) + base
+    x = x.image_conv2d(w3, b3)

-      # NOOP, 3 convs, contiguous
-      #check_schedule(x, 5)
-      check_schedule(x, 7)
+    # NOOP, 3 convs, contiguous
+    #check_schedule(x, 5)
+    check_schedule(x, 7)

  def test_image_conv_fusion_minimal(self):
    b1 = Tensor.empty(16)
@@ -1189,11 +1188,5 @@ class TestBufferView(unittest.TestCase):
    b = a.shrink(((200, 800),)).shrink(((0, 300),)).reshape((30, 10)).shrink(((20, 25), (0, 10))).contiguous()
    run_schedule(check_schedule(b, 0))

-class TestInvalidTensor(unittest.TestCase):
-  def test_full_invalid_is_zero_kernels(self):
-    from tinygrad.dtype import Invalid
-    t = Tensor.full((4,), Invalid, dtype=dtypes.float)
-    check_schedule(t, 0)
-
 if __name__ == '__main__':
  unittest.main(verbosity=2)
@@ -351,7 +351,7 @@ class TestImageSimplification(unittest.TestCase):

    self.check(load,
               "((((idx2*2)+r0)<11)&((((idx1*8)+r1)<3)!=True))",
-               "(idx0+(idx1*512+r1*64)+-192)",
+               "(((idx0+((idx1*512)+(r1*64)))+832)%1024)",
               "((((idx2*2)+r0)+(((idx1+((r1+5)//8))+1)//2))+-4)")

  def test_simplify1(self):
@@ -388,17 +388,18 @@ class TestImageSimplification(unittest.TestCase):
    alu8 = (idx0//8%32//4)
    alu9 = idx0<256

+    # TODO: can this be simplified further?
    load = get_load_image_uop(shape, alu9, (((alu8+(alu2*8))%64),(alu2//8)))
-    self.check(load, "(idx0<256)", "(idx0%2*32+idx0//32+8)", "(idx0//2%4)")
+    self.check(load, "(idx0<256)", "(((((idx0%8)*32)+(idx0//32))+8)%64)", "((idx0%8)//2)")

    load = get_load_image_uop(shape, alu9, (((alu8+(alu3*8))%64),(alu3//8)))
-    self.check(load, "(idx0<256)", "(idx0%2*32+idx0//32+16)", "(idx0//2%4)")
+    self.check(load, "(idx0<256)", "(((((idx0%8)*32)+(idx0//32))+16)%64)", "((idx0%8)//2)")

    load = get_load_image_uop(shape, alu9, (((alu8+(alu4*8))%64),(alu4//8)))
-    self.check(load, "(idx0<256)", "(idx0%2*32+idx0//32+24)", "(idx0//2%4)")
+    self.check(load, "(idx0<256)", "(((((idx0%8)*32)+(idx0//32))+24)%64)", "((idx0%8)//2)")

    load = get_load_image_uop(shape, alu9, (((alu8+(alu5*8))%64),(alu5//8)))
-    self.check(load, "(idx0<256)", "(idx0%2*32+idx0//32)", "(idx0//2%4)")
+    self.check(load, "(idx0<256)", "((((idx0%8)*32)+(idx0//32))%64)", "((idx0%8)//2)")

  def test_simplify5(self):
    # openpilot 0.9.7, chunk replacement to simplify
@@ -413,7 +414,7 @@ class TestImageSimplification(unittest.TestCase):
    valid = alu3<640

    load = get_load_image_uop(shape, valid, idx)
-    self.check(load, None, "((idx0+((idx1//3)*16))+128)", "((idx1%3)*4)")
+    self.check(load, "(((idx0+(idx1*64))%192)<160)", "((idx0+((idx1//3)*16))+128)", "(((idx0+(idx1*64))%192)//16)")

  def test_simplify6(self):
    # from openpilot
@@ -315,7 +315,7 @@ class TestProgressBar(unittest.TestCase):
    for _ in tinytqdm(range(100)): pass
    tinytqdm_time = time.perf_counter() - st

-    assert tinytqdm_time < 5 * tqdm_time
+    assert tinytqdm_time < 2 * tqdm_time

  def test_tqdm_perf_high_iter(self):
    st = time.perf_counter()
@@ -756,7 +756,7 @@ class TestLoadStoreFolding(unittest.TestCase):
    self.assertEqual(len(gated_load.src), 2)  # PTRCAT + alt
    result = graph_rewrite(gated_load, load_store_folding, name='test')
    # After rewrite, should be CAT of LOADs, each preserving alt
-    self.assertEqual(result.op, Ops.VCAT)
+    self.assertEqual(result.op, Ops.CAT)
    for inner_load in result.src:
      self.assertEqual(inner_load.op, Ops.LOAD)
      self.assertEqual(len(inner_load.src), 2)  # INDEX + alt
@@ -220,7 +220,7 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable(usum([Variable("a", 0, 7)*4, Variable("b", 0, 3)*4]) % 2, 0, 0, "0")

  def test_sum_div_some_factor(self):
-    self.helper_test_variable(usum([Variable("a", 0, 7)*5, Variable("b", 0, 3)*4]) // 2, 0, 23, "((a*2)+(b*2)+(a//2))")
+    self.helper_test_variable(usum([Variable("a", 0, 7)*5, Variable("b", 0, 3)*4]) // 2, 0, 23, "(((a*5)//2)+(b*2))")

  def test_sum_div_trim_const(self):
    self.helper_test_variable((Variable("a", 0, 7)*4 + Variable("b", 0, 3)*4 + 7) // 16, 0, 2, "(((a+b)+1)//4)")
@@ -228,10 +228,10 @@ class TestSymbolic(unittest.TestCase):
  def test_sum_div_some_partial_factor(self):
    self.helper_test_variable(usum([Variable("a", 0, 7)*6, Variable("b", 0, 7)*6]) // 16, 0, 5, "(((a*3)+(b*3))//8)")
    self.helper_test_variable(usum([uconst(16), Variable("a", 0, 7)*6, Variable("b", 0, 7)*6]) // 16, 1, 6, "((((a*3)+(b*3))//8)+1)")
-    self.helper_test_variable((Variable("a", 0, 7)*30+20)//20, 1, 11, "((a+(a//2))+1)")
+    self.helper_test_variable((Variable("a", 0, 7)*30+20)//20, 1, 11, "(((a*3)//2)+1)")

  def test_sum_div_no_factor(self):
-    self.helper_test_variable(usum([Variable("a", 0, 7)*5, Variable("b", 0, 3)*5]) // 2, 0, 25, "((a*2)+(b*2)+((a+b)//2))")
+    self.helper_test_variable(usum([Variable("a", 0, 7)*5, Variable("b", 0, 3)*5]) // 2, 0, 25, "(((a*5)+(b*5))//2)")

  def test_mod_min_max(self):
    self.helper_test_variable(Variable("x", 0, 10)%Variable("y", 1, 10), 0, 9, "(x%y)")
@@ -286,11 +286,6 @@ class TestSymbolic(unittest.TestCase):
                              "(((z+(x*-1))+(y*-1))+7)")
    self.helper_test_variable((10+12*Variable("x",0,2)+Variable("y", 0, 4)%3)%13, 8, 12, "(((x*-1)+(y%3))+10)")

-  def test_mod_congruence_tied_remainder(self):
-    # when f%c == c/2, both r and r-c have equal abs — try both signs
-    self.helper_test_variable((3+2*Variable("x",0,1)+3*Variable("y",0,1))%4, 0, 3, "((x*-2)+(y*-1)+3)")
-    self.helper_test_variable((3+6*Variable("x",0,1)+7*Variable("y",0,1))%4, 0, 3, "((x*-2)+(y*-1)+3)")
-
  def test_div_congruence(self):
    self.helper_test_variable((3+3*Variable("a",0,3))//4, 0, 3, "a")
    self.helper_test_variable((18+17*Variable("a",0,2)+17)//18, 1, 3, "(a+1)")
@@ -302,9 +297,6 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable((3+Variable("a",0,1))%4, 0, 3, "((a*-3)+3)")
    self.helper_test_variable((3+Variable("a",4,5))%4, 0, 3, "((a*-3)+15)")

-  def test_div_binary_expression(self):
-    self.helper_test_variable((3+Variable("a",0,1))//4, 0, 1, "a")
-
  def test_sum_div_const(self):
    self.helper_test_variable(usum([Variable("a", 0, 7)*4, uconst(3)]) // 4, 0, 7, "a")

@@ -553,13 +545,6 @@ class TestSymbolic(unittest.TestCase):
  def test_div_into_mod(self):
    self.helper_test_variable((Variable("idx", 0, 16)*4)%8//4, 0, 1, "(idx%2)")

-  def test_mod_div_reorder(self):
-    # (x % (a*b)) // a -> (x // a) % b, enables div-mod recombine
-    x = Variable("x", 0, 23)
-    self.helper_test_variable(x % 6 // 3, 0, 1, "(x//3%2)")
-    self.helper_test_variable(x % 12 // 4, 0, 2, "(x//4%3)")
-    self.helper_test_variable(x%12//4*4 + x%4 + x//12*12, 0, 23, "x")
-
  def test_div_neg_cancel(self):
    self.helper_test_variable((-Variable("idx", 0, 100)+199)//-4 + 50, 1, 26, "((idx//4)+1)")
    self.helper_test_variable((-Variable("idx", 0, 100)+200)//-4 + 50, 0, 25, "((idx+3)//4)")
@@ -603,7 +588,8 @@ class TestSymbolic(unittest.TestCase):
    gidx0 = Variable("gidx0", 0, 2)
    lidx2 = Variable("lidx2", 0, 12)
    lidx3 = Variable("lidx3", 0, 12)
-    self.helper_test_variable((gidx0*3+lidx2*19+lidx3*38)//(3*19), 0, 12, "((lidx2+(lidx3*2))//3)")
+    # TODO: improve nest_div_by_smallest_factor to get ((lidx2+(lidx3*2))//3)
+    self.helper_test_variable((gidx0*3+lidx2*19+lidx3*38)//(3*19), 0, 12, "((gidx0+(lidx2*19+lidx3*38)//3)//19)")

  def test_sum_mul_distribute(self):
    gidx0 = Variable("gidx0", 0, 7)
@@ -620,12 +606,6 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable((idx0*v+idx1)//v, 0, 2, "(idx0)")
    self.helper_test_variable((idx0*v+idx1)%v, 0, start_pos, "idx1")

-  def test_mod_variable_denom_factor_remainder(self):
-    d = Variable("d", 2, 5)
-    a = Variable("a", 0, 3)
-    b = Variable("b", 0, 1)
-    self.helper_test_variable((d*a+b)%d, 0, 1, "b")
-
  def test_divmod_variable_denom_fold_to_const(self):
    x = Variable("x", 20, 23)
    y = Variable("y", 8, 10)
@@ -675,7 +655,8 @@ class TestSymbolic(unittest.TestCase):
    a = Variable("a", 0, 2)
    b = Variable("b", 0, 100)
    self.helper_test_variable((31 * a + 1) % 30 + ((31 * a + 1) // 30) * 30, 1, 63, "((a*31)+1)")
-    self.helper_test_variable((31 * b + 1) % 18 + ((31 * b + 1) // 18) * 18, 1, 3101, "((b*31)+1)")
+    with self.assertRaises(AssertionError):
+      self.helper_test_variable((31 * b + 1) % 18 + ((31 * b + 1) // 18) * 18, 1, 3101, "((b*31)+1)")

  def test_div_mod_recombine_3level(self):
    gidx = Variable("gidx", 0, 150527)
@@ -699,112 +680,8 @@ class TestSymbolic(unittest.TestCase):
    b = Variable("b", 0, 100)
    exp = (16 * b + 2) % 18 + ((16 * b + 2) // 18) * 18
    self.helper_test_variable(exp, 2, 1602, "((b*16)+2)")
-    self.helper_test_variable((30 * b + 1) % 18 + ((30 * b + 1) // 18) * 18, 1, 3001, "((b*30)+1)")
-
-  def test_div_partial_quotient(self):
-    # IDIV should extract partial quotients when const_factor > divisor, matching what MOD already does
-    # (f*x+c)//d -> (f%d*x+c)//d + (f//d)*x  when f >= d
-    b = Variable("b", 0, 100)
-    self.helper_test_variable((31*b+1)//18, 0, 172, "(((b*13)+1)//18+b)")
-    self.helper_test_variable((19*b+3)//7, 0, 271, "(((b*5)+3)//7+(b*2))")
-
-  def test_gcd_with_remainder(self):
-    # gcd_with_remainder: factor GCD out of non-constant terms and denominator
-    a = Variable("a", 0, 2)
-    self.helper_test_variable((a*4)//6, 0, 1, "(a*2//3)")
-    self.helper_test_variable((a*4+1)//6, 0, 1, "(a*2//3)")
-    self.helper_test_variable((a*4+2)//6, 0, 1, "((a*2+1)//3)")
-    self.helper_test_variable((a*4+3)//6, 0, 1, "((a*2+1)//3)")
-    self.helper_test_variable((a*4)%6, 0, 4, "(a*2%3*2)")
-    self.helper_test_variable((a*4+1)%6, 1, 5, "(a*2%3*2+1)")
-    self.helper_test_variable((a*4+2)%6, 0, 4, "((a*2+1)%3*2)")
-    self.helper_test_variable((a*4+3)%6, 1, 5, "((a*2+1)%3*2+1)")
-
-  def test_div_by_factor_tie_break(self):
-    a = Variable("a", 0, 1)
-    b = Variable("b", 0, 1)
-    with Context(CORRECT_DIVMOD_FOLDING=1):
-      self.helper_test_variable((a*2+b*3+2)//6, 0, 1, "((a+b+1)//3)")
-
-  def test_div_mod_recombine_large_coeff(self):
-    # recombine must work even when coeff > divisor: both mod and div reduce the coeff the same way
-    b = Variable("b", 0, 100)
-    self.helper_test_variable((19*b+3)%7 + ((19*b+3)//7)*7, 3, 1903, "((b*19)+3)")
-    a = Variable("a", 0, 10)
-    self.helper_test_variable((25*a+3)%10 + ((25*a+3)//10)*10, 3, 253, "((a*25)+3)")
-
-  def test_mod_nest_by_factor(self):
-    # (a*f+b) % (f*k) = (a%k)*f + b when 0<=b<f — mirrors nest_div_by_factor for MOD
-    gidx0 = Variable("gidx0", 0, 15)
-    lidx0 = Variable("lidx0", 0, 3)
-    # f=4, k=2, c=8: (gidx0*4+lidx0)%8 = (gidx0%2)*4 + lidx0
-    self.helper_test_variable((gidx0*4+lidx0)%8, 0, 7, "(lidx0+gidx0%2*4)")
-    # f=2, k=4: (gidx0*2+lidx0)%8 where lidx0 in [0,1]
-    lidx1 = Variable("lidx1", 0, 1)
-    self.helper_test_variable((gidx0*2+lidx1)%8, 0, 7, "(lidx1+gidx0%4*2)")
-    # f=3, k=3: (a*3+b)%9 where b in [0,2]
-    a = Variable("a", 0, 10)
-    b = Variable("b", 0, 2)
-    self.helper_test_variable((a*3+b)%9, 0, 8, "(b+a%3*3)")
-
-  def test_mod_nest_by_factor_with_const(self):
-    # nest_by_factor MOD with non-zero constant offset: (a*f+b+const) % (f*k) = (a%k)*f + b + const when 0<=b+const<f
-    a = Variable("a", 0, 7)
-    b = Variable("b", 0, 1)
-    # f=4, k=2, const=2: (a*4+b+2)%8 = (a%2)*4 + b + 2
-    self.helper_test_variable((a*4+b+2)%8, 2, 7, "(b+a%2*4+2)")
-    # f=6, k=2, const=3: (a*6+b+3)%12 = (a%2)*6 + b + 3
-    b2 = Variable("b", 0, 2)
-    self.helper_test_variable((a*6+b2+3)%12, 3, 11, "(b+a%2*6+3)")
-    # f=3, k=2, const=1: (a*3+b+1)%6 = (a%2)*3 + b + 1
-    self.helper_test_variable((a*3+b+1)%6, 1, 5, "(b+a%2*3+1)")
-
-  def test_div_nest_by_factor_with_const(self):
-    # nest_by_factor IDIV: (160*a + 5*b + 4*c + K) // 60 should pick div=5 (clean) over div=4 (dirty)
-    a = Variable("a", 0, 2)
-    b = Variable("b", 0, 31)
-    c = Variable("c", 0, 1)
-    self.helper_test_variable((160*a + 5*b + 4*c) // 60, 0, 7, "(a*2+(b+a*8)//12)")
-    self.helper_test_variable((160*a + 5*b + 4*c + 1) // 60, 0, 8, "(a*2+(b+c+a*8)//12)")
-    self.helper_test_variable((160*a + 5*b + 4*c + 2) // 60, 0, 8, "(a*2+(b+c+a*8)//12)")
-    self.helper_test_variable((160*a + 5*b + 4*c + 3) // 60, 0, 8, "(a*2+(b+c+a*8)//12)")
-    self.helper_test_variable((160*a + 5*b + 4*c + 59) // 60, 0, 8, "(a*2+(b+c+a*8+11)//12)")
-
-  def test_div_mod_recombine_after_nesting(self):
-    # when nest_div_by_factor simplifies the div, the mod must also nest so recombine can fire
-    gidx0 = Variable("gidx0", 0, 15)
-    lidx0 = Variable("lidx0", 0, 3)
-    x = gidx0*4+lidx0
-    # div nests: x//8 -> gidx0//2, mod nests: x%8 -> (gidx0%2)*4+lidx0, then recombine gives x back
-    self.helper_test_variable((x//8)*8 + x%8, 0, 63, "(lidx0+gidx0*4)")
-    # with a scaling factor: recombine gives x*2
-    self.helper_test_variable((x//8)*16 + (x%8)*2, 0, 126, "(gidx0*8+lidx0*2)")
-    # two variables with different factors
-    a = Variable("a", 0, 7)
-    b = Variable("b", 0, 1)
-    y = a*6+b
-    # div nests: y//12 -> a//2, mod nests: y%12 -> (a%2)*6+b, recombine
-    self.helper_test_variable((y//12)*12 + y%12, 0, 43, "(b+a*6)")
-
-  def test_div_mod_recombine_in_additive_sum(self):
-    x = Variable("x", 0, 31)
-    y = Variable("y", 0, 5)
-    # recombine should work inside larger additive sums, not just in the two special y+... tree shapes
-    self.helper_test_variable((x//8)*4 + y + (x//2)%4, 0, 20, "(y+x//2)")
-    self.helper_test_variable(y + (x//8)*4 + (x//2)%4, 0, 20, "(y+x//2)")
-
-  def test_div_mod_recompose_low_order_remainder(self):
-    x = Variable("x", 0, 127)
-    self.helper_test_variable((x//2)%4*2 + x%2, 0, 7, "(x%8)")
-
-  def test_reshape_index_roundtrip(self):
-    # simulate reshape index decompose then recompose — the core pattern this enables
-    # (8,8) decomposed for (16,4): combined=r0*8+r1, div and mod by 4
-    r0 = Variable("r0", 0, 7)
-    r1 = Variable("r1", 0, 7)
-    combined = r0*8+r1
-    src_idx = (combined//4)*4 + combined%4
-    self.helper_test_variable(src_idx, 0, 63, "(r1+r0*8)")
+    with self.assertRaises(AssertionError):
+      self.helper_test_variable((30 * b + 1) % 18 + ((30 * b + 1) // 18) * 18, 1, 3001, "((b*30)+1)")

  def test_gated_load(self):
    idx = Variable("idx", 0, 24)
@@ -952,12 +829,6 @@ class TestSymbolic(unittest.TestCase):
    self.assertIn((a.cast(dtypes.long)+b.cast(dtypes.long)).render(), "(long)((a+b))")
    self.assertIn((a.cast(dtypes.long)*b.cast(dtypes.long)).render(), "(long)((a*b))")

-  def test_nested_mod_negative_range(self):
-    # (x%(k*c))%c = x%c holds for cmod regardless of signs since sign(x%(k*c)) = sign(x)
-    x = Variable("x", 0, 1575)
-    self.helper_test_variable(((x + (-1064)) % 512) % 4, -3, 3, "((x+-1064)%4)")
-    self.helper_test_variable(((x + (-1064)) % 512) % 128, -127, 127, "((x+-1064)%128)")
-
 class TestSymbolicNumeric(unittest.TestCase):
  def helper_test_numeric(self, f):
    MIN, MAX = 0, 10
@@ -1049,17 +920,6 @@ class TestSymInfer(unittest.TestCase):

    assert sym_infer(UOp.const(dtypes.float, 1.5).bitcast(dtypes.uint), {}) == 1069547520

-  def test_sym_infer_deeply_nested(self):
-    # build an expression that exceeds Python's nested parentheses limit for eval
-    # max(x, negative_const) can't be simplified when x can be negative, so nesting compounds
-    a = Variable("a", 1, 8192)
-    b = Variable("b", 0, 8191)
-    expr = a
-    for _ in range(200):
-      expr = (expr * (b + a)).maximum(uconst(-33554432)) * uconst(-1) + a
-    result = sym_infer(expr, {"a": 1, "b": 0})
-    assert isinstance(result, int)
-
 """
@unittest.skip("not supported on uops yet")
 class TestSymbolicSymbolicOps(unittest.TestCase):
@@ -64,8 +64,8 @@ class TestVminVmaxProperties(unittest.TestCase):

    # negative mask: x & -1 could be anything since -1 has all bits set
    uop = x & -1
-    self.assertEqual(uop.vmin, dtypes.int32.min)
-    self.assertEqual(uop.vmax, dtypes.int32.max)
+    self.assertEqual(uop.vmin, dtypes.min(dtypes.int32))
+    self.assertEqual(uop.vmax, dtypes.max(dtypes.int32))

  def test_vmin_vmax_multiplication_with_variable(self):
    # vmin and vmax for multiplication with a variable
@@ -136,8 +136,8 @@ class TestVminVmaxProperties(unittest.TestCase):
    self.assertEqual(x_bool.vmin, False)
    self.assertEqual(x_bool.vmax, True)
    x_uint = x.cast(dtypes.uint)
-    self.assertEqual(x_uint.vmin, dtypes.uint.min)
-    self.assertEqual(x_uint.vmax, dtypes.uint.max)
+    self.assertEqual(x_uint.vmin, dtypes.min(dtypes.uint))
+    self.assertEqual(x_uint.vmax, dtypes.max(dtypes.uint))

  def test_vmin_vmax_invalid(self):
    i = UOp.invalid()
@@ -1,4 +1,4 @@
-import unittest, decimal, sys, json
+import unittest, decimal, sys
 from dataclasses import dataclass
 from typing import Generator

@@ -509,15 +509,9 @@ class TestVizProfiler(BaseTestViz):

  def test_calltrace(self):
    def fxn(): return Tensor.empty(10).mul(2).realize()
-    with cpu_profile(TracingKey("test_fxn"), "CUSTOM"):
-      fxn()
-    codegen_trace = get_viz_list()[0]["steps"][0]["trace"]
-    assert any(fxn.__code__.co_filename == f and fxn.__code__.co_firstlineno == l for f,l,*_ in codegen_trace), str(codegen_trace)
-    profile_ret = load_profile(cpu_events)
-    e = profile_ret["layout"]["CUSTOM"]["events"][0]
-    self.assertEqual(e["name"], "test_fxn")
-    runtime_trace = json.loads(e["fmt"].replace("TB:", ""))
-    assert any(fxn.__code__.co_filename == f and fxn.__code__.co_firstlineno+1 == l for f,l,*_ in runtime_trace), str(runtime_trace)
+    fxn()
+    trace = get_viz_list()[0]["steps"][0]["trace"]
+    assert any(fxn.__code__.co_filename == f and fxn.__code__.co_firstlineno == l for f,l,*_ in trace), str(trace)

  # can pack up to 1hr 11 min of trace events
  def test_trace_duration(self):
@@ -499,11 +499,7 @@ class TestAssign(unittest.TestCase):
    # assign to a shape-changing bitcast view (only works on DISK currently)
    a = Tensor([0]*8, dtype=dtypes.uint8).realize()
    a.bitcast(dtypes.int64).assign(Tensor([12345], dtype=dtypes.int64)).realize()
-    try:
-      np.testing.assert_equal(a.numpy(), [57, 48, 0, 0, 0, 0, 0, 0])
-    except AssertionError:
-      # TODO: broken now
-      np.testing.assert_equal(a.numpy(), [0]*8)
+    np.testing.assert_equal(a.numpy(), [0]*8)  # TODO: should be [57, 48, 0, 0, 0, 0, 0, 0] (little-endian 12345)

  @unittest.skip("don't use output buffer, and mismatch dtype no longer supported")
  def test_cast_assignment(self):
@@ -695,11 +691,7 @@ class TestAssignOrdering(unittest.TestCase):
    right = buf[4:8].clone()  # lazy - not captured yet
    buf[0:4].assign(right).realize()  # this works
    buf[4:8].assign(left).realize()  # left now reads from modified buf!
-    try:
-      np.testing.assert_equal(buf.numpy(), [5, 6, 7, 8, 1, 2, 3, 4])
-    except AssertionError:
-      # TODO: broken now
-      np.testing.assert_equal(buf.numpy(), [5, 6, 7, 8, 5, 6, 7, 8])
+    np.testing.assert_equal(buf.numpy(), [5, 6, 7, 8, 5, 6, 7, 8])  # TODO: wrong! should be [5,6,7,8,1,2,3,4]

    # with .realize() on temps: values captured before writes
    buf = Tensor([1, 2, 3, 4, 5, 6, 7, 8]).contiguous().realize()
@@ -817,55 +809,40 @@ class TestAssignToUnrealizedView(unittest.TestCase):
    c = t.to("CPU:1")  # unrealized COPY
    self.assertIs(c.uop.base.op, Ops.COPY)
    c[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).to("CPU:1").contiguous().realize())
-    try:
-      self.assertEqual(c.tolist(), [[0,1],[0,1]])
-    except AssertionError:
-      # TODO: broken now
-      self.assertEqual(c.tolist(), [[0,0],[0,0]])
+    # TODO: should be [[0,1],[0,1]]
+    self.assertEqual(c.tolist(), [[0,0],[0,0]])

  def test_contiguous(self):
    t = Tensor([[1,2],[3,4]]).contiguous().realize()
    c = t.permute(1,0).contiguous()  # unrealized CONTIGUOUS
    self.assertIs(c.uop.base.op, Ops.CONTIGUOUS)
    c[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).contiguous().realize())
-    try:
-      self.assertEqual(c.tolist(), [[1,1],[2,1]])
-    except AssertionError:
-      # TODO: broken now
-      self.assertEqual(c.tolist(), [[1,3],[2,4]])
+    # TODO: should be [[1,1],[2,1]]
+    self.assertEqual(c.tolist(), [[1,3],[2,4]])

  def test_contiguous_backward(self):
    t = Tensor([[1,2],[3,4]]).contiguous().realize()
    cb = t.contiguous_backward()  # unrealized CONTIGUOUS_BACKWARD
    self.assertIs(cb.uop.base.op, Ops.CONTIGUOUS_BACKWARD)
    cb[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).contiguous().realize())
-    try:
-      self.assertEqual(cb.tolist(), [[1,1],[3,1]])
-    except AssertionError:
-      # TODO: broken now
-      self.assertEqual(cb.tolist(), [[1,2],[3,4]])
+    # TODO: should be [[1,1],[3,1]]
+    self.assertEqual(cb.tolist(), [[1,2],[3,4]])

  def test_detach_copy(self):
    t = Tensor.zeros(2,2, dtype=dtypes.int).to("CPU:0").contiguous().realize()
    d = t.to("CPU:1").detach()  # DETACH(unrealized COPY)
    self.assertIs(d.uop.base.op, Ops.COPY)
    d[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).to("CPU:1").contiguous().realize())
-    try:
-      self.assertEqual(d.tolist(), [[0,1],[0,1]])
-    except AssertionError:
-      # TODO: broken now
-      self.assertEqual(d.tolist(), [[0,0],[0,0]])
+    # TODO: should be [[0,1],[0,1]]
+    self.assertEqual(d.tolist(), [[0,0],[0,0]])

  def test_detach_contiguous(self):
    t = Tensor([[1,2],[3,4]]).contiguous().realize()
    d = t.permute(1,0).contiguous().detach()  # DETACH(unrealized CONTIGUOUS)
    self.assertIs(d.uop.base.op, Ops.CONTIGUOUS)
    d[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).contiguous().realize())
-    try:
-      self.assertEqual(d.tolist(), [[1,1],[2,1]])
-    except AssertionError:
-      # TODO: broken now
-      self.assertEqual(d.tolist(), [[1,3],[2,4]])
+    # TODO: should be [[1,1],[2,1]]
+    self.assertEqual(d.tolist(), [[1,3],[2,4]])

  def test_alu(self):
    a = Tensor([1,2,3,4]).contiguous().realize()
@@ -873,74 +850,31 @@ class TestAssignToUnrealizedView(unittest.TestCase):
    c = a + b  # unrealized ADD
    self.assertIs(c.uop.base.op, Ops.ADD)
    c[:2].assign(Tensor([99, 99]).realize())
-    try:
-      self.assertEqual(c.tolist(), [99,99,10,12])
-    except AssertionError:
-      # TODO: broken now, silently dropped
-      self.assertEqual(c.tolist(), [6,8,10,12])
+    # TODO: silently dropped, should be [99,99,10,12] or raise an error
+    self.assertEqual(c.tolist(), [6,8,10,12])

  def test_reduce(self):
    a = Tensor([[1,2],[3,4]]).contiguous().realize()
    r = a.sum(axis=0)  # unrealized REDUCE_AXIS
    self.assertIs(r.uop.base.op, Ops.REDUCE_AXIS)
    r[:1].assign(Tensor([99]).realize())
-    try:
-      self.assertEqual(r.tolist(), [99,6])
-    except AssertionError:
-      # TODO: broken now, silently dropped
-      self.assertEqual(r.tolist(), [4,6])
+    # TODO: silently dropped, should be [99,6] or raise an error
+    self.assertEqual(r.tolist(), [4,6])

  def test_cast(self):
    a = Tensor([1,2,3,4]).contiguous().realize()
    c = a.float()  # unrealized CAST
    self.assertIs(c.uop.base.op, Ops.CAST)
    c[:2].assign(Tensor([99, 99], dtype=dtypes.float).realize())
-    try:
-      self.assertEqual(c.tolist(), [99,99,3,4])
-    except AssertionError:
-      # TODO: broken now, silently dropped
-      self.assertEqual(c.tolist(), [1,2,3,4])
+    # TODO: silently dropped, should be [99,99,3,4] or raise an error
+    self.assertEqual(c.tolist(), [1,2,3,4])

  def test_const(self):
    c = Tensor(5).reshape(1, 1).expand(2, 2)
    self.assertIs(c.uop.base.op, Ops.CONST)
    c[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).contiguous().realize())
-    try:
-      self.assertEqual(c.tolist(), [[5,1],[5,1]])
-    except AssertionError:
-      # TODO: broken now, silently dropped
-      self.assertEqual(c.tolist(), [[5,5],[5,5]])
-
-class TestPartialAssignToSharedBuffer(unittest.TestCase):
-  def test_five_slices(self):
-    big = Tensor.zeros(50).contiguous().realize()
-    views = [big[i*10:(i+1)*10].reshape(2, 5) for i in range(5)]
-    for v in views: v.assign(v + 1)
-    Tensor.realize(*views)
-    for v in views:
-      np.testing.assert_allclose(v.numpy(), np.ones((2, 5)))
-
-  def test_many_slices(self):
-    n_params = 10
-    big = Tensor.zeros(n_params * 12).contiguous().realize()
-    grads = [big[i*12:(i+1)*12].reshape(3, 4) for i in range(n_params)]
-    for g in grads: g.assign(g + 1)
-    Tensor.realize(*grads)
-    for g in grads:
-      np.testing.assert_allclose(g.numpy(), np.ones((3, 4)))
-
-  def test_mixed_shapes(self):
-    big = Tensor.zeros(100).contiguous().realize()
-    shapes = [(3, 4), (4, 6), (6, 4), (2, 5), (4, 3)]
-    pos, views = 0, []
-    for s in shapes:
-      n = s[0] * s[1]
-      views.append(big[pos:pos+n].reshape(*s))
-      pos += n
-    for v in views: v.assign(v + 1)
-    Tensor.realize(*views)
-    for v, s in zip(views, shapes):
-      np.testing.assert_allclose(v.numpy(), np.ones(s))
+    # TODO: silently dropped, should be [[5,1],[5,1]] or raise an error
+    self.assertEqual(c.tolist(), [[5,5],[5,5]])

 if __name__ == "__main__":
  unittest.main()
@@ -2,7 +2,7 @@ import unittest
 import numpy as np
 from tinygrad import Tensor, function
 from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import UOp, Ops
+from tinygrad.uop.ops import UOp

 class TestCall(unittest.TestCase):
  def test_call_plus(self):
@@ -100,42 +100,6 @@ class TestCall(unittest.TestCase):
    c = Tensor.call(a, b, fxn=a.as_param(0) + b.as_param(1))
    np.testing.assert_equal(c.numpy(), 2 * np.ones((10, 10)))

-class TestCallShape(unittest.TestCase):
-  def test_call_shape_int(self):
-    # fixed-shape function: shape passes through unchanged
-    @function
-    def f(x:Tensor) -> Tensor: return x * 2
-    self.assertEqual(f(Tensor.empty(4, 8)).shape, (4, 8))
-
-  def test_call_shape_param_substitution(self):
-    # symbolic shape dimension is substituted: inner PARAM replaced with the BIND arg
-    @function
-    def f(x:Tensor) -> Tensor: return x * 2
-    sz = UOp.variable("sz", 1, 8)
-    shape = f(Tensor.empty(8)[:sz.bind(5)]).shape
-    # the PARAM should be gone, replaced with the BIND from the call arg
-    self.assertIsInstance(shape[0], UOp)
-    self.assertNotEqual(shape[0].op, Ops.PARAM)
-    self.assertEqual(shape[0], sz.bind(5))
-
-  def test_call_shape_expr_substitution(self):
-    # expression containing PARAMs in shape gets fully substituted
-    @function
-    def f(x:Tensor) -> Tensor: return x + 1
-    sz = UOp.variable("sz", 1, 10)
-    shape = f(Tensor.empty(10, 4)[:sz.bind(3)]).shape
-    self.assertIsInstance(shape[0], UOp)
-    self.assertNotEqual(shape[0].op, Ops.PARAM)
-    self.assertEqual(shape[1], 4)
-
-  def test_call_shape_no_param_passthrough(self):
-    # a non-PARAM UOp shape element passes through unchanged
-    @function
-    def f(x:Tensor) -> Tensor: return x * 3
-    sz = UOp.variable("sz", 1, 8)
-    shape = f(Tensor.empty(8)[:sz.bind(5)]).shape
-    self.assertEqual(shape[0], sz.bind(5))
-
 class TestCallSchedule(unittest.TestCase):
  def test_reshape_precompile(self):
    a = Tensor.empty(4, 8).realize()
@@ -174,68 +138,5 @@ class TestCallSchedule(unittest.TestCase):
    (out-ref).square().mean().backward()
    out.realize(a.grad, b.grad, c.grad)

-  def test_precompile_symbolic_shape(self):
-    """precompile with a symbolic-shaped input produces correct values and shape"""
-    @function(precompile=True)
-    def f(x:Tensor) -> Tensor: return x * 2
-    sz = UOp.variable("sz", 1, 8)
-    a = Tensor([1., 2., 3., 4., 5., 6., 7., 8.])[:sz.bind(5)]
-    out = f(a)
-    self.assertIsInstance(out.shape[0], UOp)
-    np.testing.assert_allclose(out[:5].numpy(), [2., 4., 6., 8., 10.])
-
-  def test_precompile_symbolic_shape_contiguous(self):
-    """precompile with a .contiguous() inside the function body on a symbolic-shaped input"""
-    @function(precompile=True)
-    def f(x:Tensor) -> Tensor: return (x * 2).contiguous() + 1
-    sz = UOp.variable("sz", 1, 8)
-    a = Tensor([1., 2., 3., 4., 5., 6., 7., 8.])[:sz.bind(3)]
-    out = f(a)
-    self.assertIsInstance(out.shape[0], UOp)
-    np.testing.assert_allclose(out[:3].numpy(), [3., 5., 7.])
-
-  def test_precompile_symbolic_shape_chain(self):
-    """precompiled symbolic result used in downstream ops (tests AFTER has correct symbolic shape)"""
-    @function(precompile=True)
-    def f(x:Tensor) -> Tensor: return x * 2
-    sz = UOp.variable("sz", 1, 8)
-    a = Tensor([1., 2., 3., 4., 5., 6., 7., 8.])[:sz.bind(4)]
-    out = f(a) + 10  # downstream op on the precompiled result
-    self.assertIsInstance(out.shape[0], UOp)
-    np.testing.assert_allclose(out[:4].numpy(), [12., 14., 16., 18.])
-
-  def test_precompile_bind_arg(self):
-    """precompile with a BIND (scalar variable) as a function argument"""
-    @function(precompile=True)
-    def f(x:Tensor, scale:UOp) -> Tensor: return x * scale
-    v = UOp.variable("scale", 1, 100)
-    a = Tensor([1., 2., 3.])
-    out = f(a, v.bind(5))
-    np.testing.assert_allclose(out.numpy(), [5., 10., 15.])
-
-  def test_precompile_schedule_cache_hit(self):
-    """two instances of the same @function should produce identical function body keys (schedule cache hit)"""
-    @function(precompile=True)
-    def f(x:Tensor) -> Tensor: return x + Tensor.full(x.shape, -1.0)
-    a = Tensor.empty(4, 8)
-    b = Tensor.empty(4, 8)
-    r0, r1 = f(a), f(b)
-    # find the CALL nodes
-    c0 = next(u for u in r0.uop.toposort() if u.op is Ops.CALL)
-    c1 = next(u for u in r1.uop.toposort() if u.op is Ops.CALL)
-    # the function bodies (src[0]) should have identical keys — unique consts must not leak through
-    self.assertEqual(c0.src[0].key, c1.src[0].key)
-
-  def test_precompile_symbolic_2d(self):
-    """precompile with symbolic shapes in 2D (tests debuf reshape with symbolic PARAM)"""
-    @function(precompile=True)
-    def f(x:Tensor) -> Tensor: return x * 2 + 1
-    sz = UOp.variable("sz", 1, 16)
-    a = Tensor.arange(16*4).reshape(16, 4).float()[:sz.bind(5)]
-    out = f(a)
-    # result shape should have the symbolic dim, not the max
-    self.assertIsInstance(out.shape[0], UOp)
-    np.testing.assert_allclose(out[:5].numpy(), (np.arange(16*4).reshape(16, 4)[:5] * 2 + 1).astype(np.float32))
-
 if __name__ == '__main__':
  unittest.main()
@@ -269,16 +269,22 @@ class TestDiskTensor(TempDirTestCase):
    assert tout == list([(x+1,x) for x in range(32,64,2)])

  def test_strided_read(self):
-    # test non-contiguous (strided) read raises
+    # test non-contiguous (strided) read - should read elements at indices 0, 2, 4
    dt = Tensor([0, 1, 2, 3, 4, 5]).to(f"disk:{self.tmp('dt_strided_read')}")
-    with self.assertRaisesRegex(RuntimeError, "non-contiguous view is not supported"):
-      dt[::2].tolist()
+    with self.assertRaises(RuntimeError):
+      result = dt[::2].tolist()
+      # TODO: dt[::2] selects indices 0, 2, 4, so result should be [0, 2, 4]
+      # self.assertEqual(result, [0, 2, 4])
+      self.assertEqual(result, [0, 1, 2])  # wrong!

  def test_permuted_read(self):
-    # test non-contiguous (permuted) read raises
+    # test non-contiguous (permuted) read - should read transposed
    dt = Tensor([[0, 1, 2], [3, 4, 5]]).to(f"disk:{self.tmp('dt_permuted_read')}")
-    with self.assertRaisesRegex(RuntimeError, "non-contiguous view is not supported"):
-      dt.T.tolist()
+    with self.assertRaises(RuntimeError):
+      result = dt.T.tolist()
+      # TODO: transpose should give [[0, 3], [1, 4], [2, 5]]
+      # self.assertEqual(result, [[0, 3], [1, 4], [2, 5]])
+      self.assertEqual(result, [[0, 1], [2, 3], [4, 5]])  # wrong!

  def test_write_ones(self):
    out = Tensor.ones(10, 10, device="CPU").contiguous()
@@ -304,10 +310,13 @@ class TestDiskTensor(TempDirTestCase):
    self.assertEqual(dt.tolist(), [[1], [3]])

  def test_strided_setitem(self):
-    # test non-contiguous (strided) setitem raises
+    # test non-contiguous (strided) setitem - should set elements at indices 0, 2, 4
    dt = Tensor([1, 2, 3, 4, 5, 6]).to(f"disk:{self.tmp('dt_strided_setitem')}")
-    with self.assertRaisesRegex(RuntimeError, "non-contiguous view is not supported"):
+    with self.assertRaises(RuntimeError):
      dt[::2] = Tensor([10, 20, 30])
+      # TODO: dt[::2] selects indices 0, 2, 4, so result should be [10, 2, 20, 4, 30, 6]
+      # self.assertEqual(dt.tolist(), [10, 2, 20, 4, 30, 6])
+      self.assertEqual(dt.tolist(), [10, 20, 30, 4, 5, 6])  # wrong!

  def test_advanced_setitem_not_supported(self):
    dt = Tensor.arange(12).reshape(3, 4).to(f"disk:{self.tmp('dt_advanced_setitem')}")
@@ -17,15 +17,13 @@ dtype_floats = [dt for dt in core_dtypes if dtypes.is_float(dt) and is_dtype_sup

 FP8E4M3_MAX = 448.0
 FP8E5M2_MAX = 57344.0
-FP8E4M3FNUZ_MAX = 240.0
-FP8E5M2FNUZ_MAX = 57344.0

 def _assert_eq(tensor:Tensor, target_dtype:DType, target, tol_target_dtype:float=1e-7):
  if DEBUG >= 2: print(tensor.numpy())
  try:
    assert tensor.dtype == target_dtype
    np.testing.assert_allclose(tensor.numpy(), target, rtol={dtypes.float16:1e-3, dtypes.bfloat16:1e-2,
-      dtypes.fp8e4m3:1e-1, dtypes.fp8e5m2:5e-1, dtypes.fp8e4m3fnuz:1e-1, dtypes.fp8e5m2fnuz:5e-1}.get(target_dtype, tol_target_dtype))
+      dtypes.fp8e4m3:1e-1, dtypes.fp8e5m2:5e-1}.get(target_dtype, tol_target_dtype))

  except AssertionError as e:
    raise AssertionError(f"\ntensor {tensor.numpy()} dtype {tensor.dtype} does not match target {target} with dtype {target_dtype}") from e
@@ -26,16 +26,16 @@ class TestGGUF(unittest.TestCase):
    # codes 0-7 = [0, 1, 2, 3, 4, 6, 8, 12], codes 8-15 are their negatives
    block = np.array([0x80] + list(range(16)), dtype=np.uint8)  # E=128, nibbles 0-15 in low, zeros in high
    expected = np.array([0., 1., 2., 3., 4., 6., 8., 12., -0., -1., -2., -3., -4., -6., -8., -12.] + [0.]*16, dtype=np.float32)
-    np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 32, GGMLQuantizationType.MXFP4.value).numpy().flatten(), expected)
+    np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 32, 39).numpy().flatten(), expected)

  def test_dequantization_q4_0(self): self._test_dequantization(GGMLQuantizationType.Q4_0)
  def test_dequantization_q4_1(self): self._test_dequantization(GGMLQuantizationType.Q4_1)
  def test_dequantization_q8_0(self): self._test_dequantization(GGMLQuantizationType.Q8_0)
  def test_dequantization_q4_k(self): self._test_dequantization(GGMLQuantizationType.Q4_K)
-  def test_dequantization_q5_k(self): self._test_dequantization(GGMLQuantizationType.Q5_K)
  def test_dequantization_q6_k(self): self._test_dequantization(GGMLQuantizationType.Q6_K)
-  def test_dequantization_mxfp4(self): self._test_dequantization(GGMLQuantizationType.MXFP4)
-  def test_dequantization_mxfp4_old(self):
+  def test_dequantization_mxfp4(self):
+    MXFP4 = 39
+
    def encode(nibbles, E):
      packed = [(low & 0xF) | ((high & 0xF) << 4) for low, high in zip(nibbles[:16], nibbles[16:])]
      return np.array([E] + packed, dtype=np.uint8)
@@ -56,10 +56,12 @@ class TestGGUF(unittest.TestCase):
      blocks.append(encode(codes, E))
      expected.extend(decode(c, E) for c in codes)
    tensor = Tensor(np.concatenate(blocks))
-    out = ggml_data_to_tensor(tensor, len(expected), GGMLQuantizationType.MXFP4.value)
-    np.testing.assert_equal(out.numpy(), expected)
+    out = ggml_data_to_tensor(tensor, len(expected), MXFP4)
+    # TODO: should this be exact equal? somehow failed on CI
+    np.testing.assert_allclose(out.numpy(), expected, atol=0.0, rtol=1e-6)

  def test_dequantization_mxfp4_block(self):
+    MXFP4 = 39
    # https://gist.github.com/Ananta-Ranganathan/3317b6ed51a3b033e9c2564fafb4e043
    # used the above script to download the first block of blk.0.attn_k_b.weight from
    # https://huggingface.co/unsloth/GLM-4.7-Flash-GGUF/blob/main/GLM-4.7-Flash-MXFP4_MOE.gguf
@@ -74,8 +76,9 @@ class TestGGUF(unittest.TestCase):
                        0.03125000,  0.00000000, 0.06250000,  0.01562500,
                        -0.06250000,  0.00000000, 0.00000000, -0.01562500,
                        0.04687500,  0.00000000, 0.00000000,  0.01562500], dtype=np.float32)
-    out = ggml_data_to_tensor(Tensor(block), 32, GGMLQuantizationType.MXFP4.value)
-    np.testing.assert_equal(out.numpy(), expected)
+    out = ggml_data_to_tensor(Tensor(block), 32, MXFP4)
+    # TODO: similar to previous test fails on Mac CI with assert_equal for unclear reason
+    np.testing.assert_allclose(out.numpy(), expected, atol=0.0, rtol=1e-6)

  def test_expected_failure_unknown_type(self):
    with self.assertRaises(ValueError):
@@ -127,9 +130,8 @@ class TestGGUFGEMV(unittest.TestCase):
    q_data = rng.integers(0, 256, size=n_blocks * type_size, dtype=np.uint8).reshape(n_blocks, type_size)
    scales = np.float16(rng.standard_normal(n_blocks * 4)).view(np.uint8).reshape(n_blocks, -1)
    if qtype == GGMLQuantizationType.Q8_0: q_data[:, :2] = scales[:, :2]                    # d at offset 0
-    elif qtype in (GGMLQuantizationType.Q4_K, GGMLQuantizationType.Q5_K): q_data[:, :4] = scales[:, :4] # d, dmin at offset 0
+    elif qtype == GGMLQuantizationType.Q4_K: q_data[:, :4] = scales[:, :4]                  # d, dmin at offset 0
    elif qtype == GGMLQuantizationType.Q6_K: q_data[:, -2:] = scales[:, :2]                 # d at end
-    elif qtype == GGMLQuantizationType.MXFP4: q_data[:, 0] = rng.integers(120, 136, size=n_blocks, dtype=np.uint8) # constrain byte0
    q_data = q_data.flatten()
    ref = dequantize(q_data, qtype).reshape(rows, cols)

@@ -149,13 +151,10 @@ class TestGGUFGEMV(unittest.TestCase):
    x = rng.standard_normal(cols).astype(np.float32)
    np.testing.assert_allclose((tensors["weight"] @ Tensor(x)).numpy(), ref @ x, atol=1e-2, rtol=1e-2)
    np.testing.assert_equal(tensors["weight"].numpy(), ref)
-    assert np.isfinite(ref).all() and np.isfinite(tensors["weight"].numpy()).all(), f"{qtype.name} has NaN/Inf"

  def test_gguf_gemv_q8_0(self): self._test_gguf_gemv(GGMLQuantizationType.Q8_0)
  def test_gguf_gemv_q4_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q4_K)
-  def test_gguf_gemv_q5_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q5_K)
  def test_gguf_gemv_q6_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q6_K)
-  def test_gguf_gemv_mxfp4(self): self._test_gguf_gemv(GGMLQuantizationType.MXFP4)

 if __name__ == '__main__':
  unittest.main()
@@ -1,126 +0,0 @@
-import unittest
-from tinygrad import Tensor
-from tinygrad.dtype import Invalid, dtypes
-from tinygrad.engine.realize import run_schedule
-
-class TestInvalidTensor(unittest.TestCase):
-  def _invalid_test_helper(self, out, expected):
-    sched = out.schedule()
-    buf = out.uop.buffer
-    buf.allocate()
-    sentinel = memoryview(bytearray(b'\x42' * buf.nbytes))
-    buf.copyin(sentinel)
-    before = buf.as_memoryview().cast(out.dtype.fmt).tolist()
-    run_schedule(sched)
-    ret = buf.as_memoryview().cast(out.dtype.fmt).tolist()
-
-    for i,v in enumerate(expected): self.assertEqual(ret[i], before[i] if v is None else v)
-
-  def test_where_x_invalid(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid)
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_where_invalid_x(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Invalid, Tensor([1.0, 2.0, 3.0, 4.0]))
-    self._invalid_test_helper(out, [None, None, 3.0, 4.0])
-
-  def test_where_invalid_2d(self):
-    mask = Tensor.arange(6).reshape(2, 3) < 3
-    vals = Tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
-    out = mask.where(vals, Invalid)
-    self._invalid_test_helper(out, [1.0, 2.0, 3.0, None, None, None])
-
-  def test_where_invalid_int(self):
-    mask = Tensor.arange(3) < 2
-    out = mask.where(Tensor([10, 20, 30]), Invalid)
-    self._invalid_test_helper(out, [10, 20, None])
-
-  def test_where_invalid_add(self):
-    mask = Tensor.arange(3) < 2
-    mixed = mask.where(Tensor([10.0, 20.0, 30.0]), Invalid)
-    out = mixed + Tensor([1.0, 2.0, 3.0])
-    self._invalid_test_helper(out, [11.0, 22.0, None])
-
-  def test_where_invalid_add_left(self):
-    mask = Tensor.arange(3) < 2
-    mixed = mask.where(Tensor([10.0, 20.0, 30.0]), Invalid)
-    out = Tensor([1.0, 2.0, 3.0]) + mixed
-    self._invalid_test_helper(out, [11.0, 22.0, None])
-
-  def test_where_always_true(self):
-    mask = Tensor.arange(3) < 10
-    out = mask.where(Tensor([10.0, 20.0, 30.0]), Invalid)
-    self._invalid_test_helper(out, [10.0, 20.0, 30.0])
-
-  def test_where_cast(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid).cast(dtypes.int)
-    self._invalid_test_helper(out, [1, 2, None, None])
-
-  def test_where_compare(self):
-    mask = Tensor.arange(4) < 2
-    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_unary(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 4.0, 9.0, 16.0]), Invalid).sqrt()
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_where_where(self):
-    mask1 = Tensor.arange(4) < 2
-    mask2 = Tensor.arange(4) > 0
-    out = mask2.where(mask1.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid), Invalid)
-    self._invalid_test_helper(out, [None, 2.0, None, None])
-
-  def test_where_reduce_always_true(self):
-    mask = Tensor.arange(4) < 9
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid).sum()
-    self._invalid_test_helper(out, [10.0])
-
-  def test_invalid_unary(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.float).sqrt())
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_invalid_binary(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.float) + 2)
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_invalid_binary_left(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), 2 + Tensor.full((4,), Invalid, dtype=dtypes.float))
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_invalid_reshape(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid).reshape(2,2)
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_invalid_cast(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.int).cast(dtypes.float))
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_invalid_bitcast(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.int).bitcast(dtypes.float))
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-  def test_where_bitcast(self):
-    mask = Tensor.arange(4) < 2
-    out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.int)).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]
-    self._invalid_test_helper(out, [1.0, 2.0, None, None])
-
-if __name__ == '__main__':
-  unittest.main()
@@ -1,112 +1,29 @@
 import unittest
 from unittest.mock import patch
-from tinygrad import Tensor, UOp
-from tinygrad.engine.schedule import schedule_cache
+from tinygrad import Tensor

 class TestTransformerGenerate(unittest.TestCase):
-  def test_kv_cache_reuse(self):
-    """Test that generate reuses the KV cache when tokens extend the cached prefix."""
+  def test_start_pos_parameter_is_used(self):
+    """Test that start_pos parameter is not ignored (regression test for always resetting to 0)."""
    from tinygrad.apps.llm import Transformer
+    # Create a minimal transformer
    model = Transformer(num_blocks=1, dim=64, hidden_dim=128, n_heads=2, n_kv_heads=2,
                        norm_eps=1e-5, vocab_size=100, head_dim=32, rope_theta=10000.0, max_context=32)

    captured_inputs = []
    def mock_call(self, tokens, start_pos):
-      captured_inputs.append((tokens.shape, start_pos if isinstance(start_pos, int) else start_pos.val))
-      return Tensor([[42]])
+      captured_inputs.append((tokens.shape, start_pos if isinstance(start_pos, int) else start_pos.bind_val))
+      return Tensor([[42]])  # return a fake next token

    with patch.object(Transformer, '__call__', mock_call):
-      # first conversation: prefill 5 tokens + 1 decode
      tokens = [1, 2, 3, 4, 5]
-      gen = model.generate(tokens)
-      next(gen)  # prefill
-      next(gen)  # decode
+      gen = model.generate(tokens, start_pos=3)
+      next(gen)  # get first token

-      # second call extends the conversation — cached prefix should be reused
-      captured_inputs.clear()
-      tokens = [1, 2, 3, 4, 5, 42, 42, 10, 11, 12]
-      gen = model.generate(tokens)
-      next(gen)
-
-    # should only process tokens[7:] = [10, 11, 12] since first 7 are cached
-    toks_shape = captured_inputs[0][0][-1]
-    self.assertEqual(toks_shape.val if isinstance(toks_shape, UOp) else toks_shape, 3)
-    self.assertEqual(captured_inputs[0][1], 7)
-
-  def test_kv_cache_invalidation(self):
-    """Test that generate invalidates the KV cache when tokens diverge from the cached prefix."""
-    from tinygrad.apps.llm import Transformer
-    model = Transformer(num_blocks=1, dim=64, hidden_dim=128, n_heads=2, n_kv_heads=2,
-                        norm_eps=1e-5, vocab_size=100, head_dim=32, rope_theta=10000.0, max_context=32)
-
-    captured_inputs = []
-    def mock_call(self, tokens, start_pos):
-      captured_inputs.append((tokens.shape, start_pos if isinstance(start_pos, int) else start_pos.val))
-      return Tensor([[42]])
-
-    with patch.object(Transformer, '__call__', mock_call):
-      # first conversation
-      gen = model.generate([1, 2, 3, 4, 5])
-      next(gen)
-
-      # completely different prompt — KV cache should be invalidated
-      captured_inputs.clear()
-      gen = model.generate([10, 20, 30])
-      next(gen)
-
-    # should process all 3 tokens from start
-    toks_shape = captured_inputs[0][0][-1]
-    self.assertEqual(toks_shape.val if isinstance(toks_shape, UOp) else toks_shape, 3)
-    self.assertEqual(captured_inputs[0][1], 0)
-
-  def test_two_prompts_schedule_cache(self):
-    """Third prompt should hit the schedule cache, not miss (first two warm up both jits: prefill + decode)."""
-    from tinygrad.apps.llm import Transformer
-    model = Transformer(num_blocks=1, dim=64, hidden_dim=128, n_heads=2, n_kv_heads=2,
-                        norm_eps=1e-5, vocab_size=100, head_dim=32, rope_theta=10000.0, max_context=64)
-
-    # first two prompts warm up both jits (prefill + decode)
-    ids = list(range(1, 6))
-    gen = model.generate(ids)
-    for _ in range(3): next(gen)
-
-    ids += list(range(10, 15))
-    gen = model.generate(ids)
-    for _ in range(3): next(gen)
-    cache_size_after_warmup = len(schedule_cache)
-
-    # third prompt should reuse the same schedule cache entries, not create new ones
-    ids += list(range(20, 25))
-    gen = model.generate(ids)
-    for _ in range(3): next(gen)
-
-    self.assertEqual(cache_size_after_warmup, len(schedule_cache),
-      f"third prompt added {len(schedule_cache) - cache_size_after_warmup} new schedule cache entries (expected 0)")
-
-  def test_chunked_prefill(self):
-    """When prompt > chunk_size, all chunks should be prefill"""
-    from tinygrad.apps.llm import Transformer
-    from tinygrad.uop.ops import resolve
-    model = Transformer(num_blocks=1, dim=64, hidden_dim=128, n_heads=2, n_kv_heads=2,
-                        norm_eps=1e-5, vocab_size=100, head_dim=32, rope_theta=10000.0, max_context=64)
-
-    def get_prefill_flags(tokens, chunk_size):
-      is_prefill = []
-      def mock_call(self, tokens, start_pos):
-        is_prefill.append(resolve(tokens.shape[1] != 1))
-        return Tensor([[42]])
-      with patch.object(Transformer, '__call__', mock_call):
-        gen = model.generate(tokens, chunk_size=chunk_size)
-        for _ in range(3): next(gen)
-      model._cached_tokens = []
-      return is_prefill
-
-    # 8 tokens, chunk_size=4 -> 2 prefill chunks
-    self.assertEqual(get_prefill_flags(list(range(8)), 4), [True, True, False, False])
-    # 9 tokens, chunk_size=4 -> 3 prefill chunks (4+4+1)
-    self.assertEqual(get_prefill_flags(list(range(9)), 4), [True, True, True, False, False])
-    # 4 tokens, chunk_size=4 -> 1 prefill chunk
-    self.assertEqual(get_prefill_flags(list(range(4)), 4), [True, False, False])
+    # With start_pos=3, the initial tensor should only have tokens[3:] = [4, 5] (length 2)
+    # If the bug existed (start_pos always reset to 0), it would have all 5 tokens
+    self.assertEqual(captured_inputs[0][0][-1], 2)  # shape should be (1, 2)
+    self.assertEqual(captured_inputs[0][1], 3)  # start_pos should be 3, not 0

 if __name__ == '__main__':
  unittest.main()
@@ -1,34 +0,0 @@
-import unittest
-from unittest.mock import MagicMock
-from tinygrad import Device
-from tinygrad.engine.realize import CompiledRunner
-
-@unittest.skipUnless(Device.DEFAULT == "METAL", "Metal device required to run")
-class TestMetalGraph(unittest.TestCase):
-  def setUp(self):
-    from tinygrad.runtime.graph.metal import MetalGraph
-    from tinygrad.runtime.ops_metal import MetalBuffer
-    self.MetalGraph = MetalGraph
-    self.MetalBuffer = MetalBuffer
-    self.dev = Device[Device.DEFAULT]
-
-  def metal_buf(self, offset): return MagicMock(_buf=self.MetalBuffer(MagicMock(), 4, offset))
-
-  def ei(self, *bufs):
-    ei = MagicMock()
-    ei.prg = MagicMock(spec=CompiledRunner)
-    ei.bufs = list(bufs)
-    return ei
-
-  def test_supports_exec_item_normal_offset(self):
-    assert self.MetalGraph.supports_exec_item([self.dev], self.ei(self.metal_buf(0), self.metal_buf(100), self.metal_buf(0xFFFFFFFF))) is True
-
-  def test_supports_exec_item_overflow_offset(self):
-    assert self.MetalGraph.supports_exec_item([self.dev], self.ei(self.metal_buf(0), self.metal_buf(0x100000000))) is False
-
-  def test_supports_exec_item_nonmetal_buf(self):
-    # HCQBuffer.offset is a method, not an int — must not crash
-    self.MetalGraph.supports_exec_item([self.dev], self.ei(MagicMock(**{"_buf.offset": lambda: 0})))
-
-if __name__ == "__main__":
-  unittest.main()
@@ -1,8 +1,7 @@
 from __future__ import annotations
-import sys, argparse, typing, re, unicodedata, json, uuid, time, functools, itertools
+import sys, argparse, typing, re, unicodedata, json, uuid, time, functools
 from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function
-from tinygrad.uop.ops import resolve
-from tinygrad.helpers import partition, DEBUG, Timing, GlobalCounters, stderr_log, colored, Context
+from tinygrad.helpers import partition, DEBUG, Timing, GlobalCounters, stderr_log, colored
 from tinygrad.viz.serve import TCPServerWithReuse, HTTPRequestHandler

 class SimpleTokenizer:
@@ -117,7 +116,7 @@ class TransformerBlock:
      self.ffn_up      = nn.Linear(dim, hidden_dim, bias=False)
      self.ffn_down    = nn.Linear(hidden_dim, dim, bias=False)

-  @function(precompile=bool(getenv("PRECOMPILE", 0)))
+  @function
  def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
    x_norm = self.attn_norm(x)                       # (B,T,D)
    q, k, v = self.attn_q(x_norm), self.attn_k(x_norm), self.attn_v(x_norm)
@@ -145,7 +144,7 @@ class TransformerBlock:

    # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
    # TODO: this if statement should be removed and it shouldn't generate extra kernels
-    mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device).triu(start_pos+1) if resolve(T != 1) else None
+    mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device).triu(start_pos+1) if T > 1 else None
    attn = q.scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True)     # (B,H,T,Hd)
    attn = attn.transpose(1, 2).reshape(B, T, -1)                                    # back to (B,T,D)
    attn = self.attn_output(attn)
@@ -166,8 +165,7 @@ class TransformerBlock:
  def __call__(self, x: Tensor, start_pos: int|UOp):
    if not hasattr(self, "cache_kv"):
      # TODO: how is the dtype of this determined?
-      # NOTE: clone is used to promise the creation of a specific buffer
-      self.cache_kv = Tensor.zeros(2, x.shape[0], self.n_kv_heads, self.max_context, self.head_dim, device=x.device).clone()
+      self.cache_kv = Tensor.zeros(2, x.shape[0], self.n_kv_heads, self.max_context, self.head_dim, device=x.device).contiguous().realize()
    return self._feed_forward(self._attention(x, start_pos)).contiguous()

 class Transformer:
@@ -179,10 +177,8 @@ class Transformer:
    self.output_norm = nn.RMSNorm(dim, norm_eps)
    self.output = nn.Linear(dim, vocab_size, bias=False)
    self.max_context = max_context
-    self._cached_tokens: list[int] = []
-    # we specialize the JIT for prefill and rollout
-    self.prefill_jit = TinyJit(self.forward)
-    self.rollout_jit = TinyJit(self.forward)
+    # JIT is used if T=1 and start_pos is a UOp. TODO: make this not needed by including T in the JIT and making start_pos always a UOp
+    self.forward_jit = TinyJit(self.forward)

  def forward(self, tokens:Tensor, start_pos:int|UOp) -> Tensor:
    x = self.token_embd(tokens)                           # (B, T, D)
@@ -191,10 +187,10 @@ class Transformer:
    return self.output(self.output_norm(x))[:, -1, :].softmax(-1, dtype="float").argmax(-1, keepdim=True)

  def __call__(self, tokens:Tensor, start_pos:int|UOp=0) -> Tensor:
-    return (self.prefill_jit if resolve(tokens.shape[1] != 1) else self.rollout_jit)(tokens, start_pos)
+    return (self.forward_jit if getenv("JIT", 1) and tokens.shape[1] == 1 and isinstance(start_pos, UOp) else self.forward)(tokens, start_pos)

  @staticmethod
-  def from_gguf(gguf:Tensor, max_context:int|None=None, realize=bool(getenv("REALIZE", 0))) -> tuple[Transformer, dict]:
+  def from_gguf(gguf:Tensor, max_context:int|None=None, realize=bool(getenv("REALIZE", 1))) -> tuple[Transformer, dict]:
    # TODO: remove the need for copy to default device
    kv, state_dict = nn.state.gguf_load(gguf.to(None).realize())

@@ -229,26 +225,15 @@ class Transformer:
      Tensor.realize(*params)
    return model, kv

-  def get_start_pos(self, tokens:list[int]):
-    return sum(1 for _ in itertools.takewhile(lambda ab: ab[0] == ab[1], zip(tokens[:-1], self._cached_tokens)))
-
-  def generate(self, tokens:list[int], chunk_size:int=32):
-    v_start_pos = UOp.variable("start_pos", 0, self.max_context-1)
-    v_toks = UOp.variable("toks", 1, chunk_size)
-    # assign all input tokens once, then slice from start_pos for the model call
-    t = Tensor(tokens + [0] * (self.max_context - len(tokens)), dtype="int32").reshape(1, self.max_context)
-    # recompute start_pos from what's currently valid in the kv cache
-    start_pos = self.get_start_pos(tokens)
-    out, prompt_len = None, len(tokens)
+  def generate(self, tokens:list[int], start_pos=0):
+    v_start_pos = UOp.variable("start_pos", 1, self.max_context-1)
+    t = Tensor([tokens[start_pos:]], dtype="int32")
    while len(tokens) < self.max_context:
-      sp, nt = v_start_pos.bind(start_pos), v_toks.bind(min(chunk_size, len(tokens) - start_pos))
-      out = self(t[:, sp:sp+nt] if start_pos < prompt_len or out is None else out, sp).realize()
-      start_pos += nt.val
-      # chunked prefill: keep processing until all prompt tokens are consumed
-      if start_pos < len(tokens): continue
-      tokens.append(int(out.item()))
-      self._cached_tokens = tokens[:]
-      yield tokens[-1]
+      t = self(t, v_start_pos.bind(start_pos) if getenv("SYM", 1) and start_pos != 0 and t.shape[-1] == 1 else start_pos)
+      next_id = int(t.item())
+      tokens.append(next_id)
+      start_pos = len(tokens) - 1
+      yield next_id

 models = {
  "llama3.2:1b": "https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q6_K.gguf",
@@ -277,7 +262,7 @@ CHAT_HTML = b'''<!DOCTYPE html><html><head><title>tinygrad chat</title><style>
           background: #2f2f2f; color: inherit; font: inherit;
           border: none; outline: none; resize: none; border-radius: 24px; field-sizing: content }
 </style></head><body><div id="chat"></div>
-<textarea id="input" rows="1" placeholder="Ask anything" autofocus></textarea>
+<textarea id="input" rows="1" placeholder="Ask anything"></textarea>
 <script>
  input.onkeydown = (e) => { if (e.key === 'Enter' && !e.shiftKey && !e.isComposing) { e.preventDefault(); send() } }
  const msgs = [];
@@ -305,22 +290,20 @@ class Handler(HTTPRequestHandler):
  def log_request(self, code='-', size='-'): pass
  def do_GET(self): self.send_data(CHAT_HTML, content_type="text/html")
  def run_model(self, ids:list[int], model_name:str, include_usage=False):
-    cache_start_pos = model.get_start_pos(ids)
-    stderr_log(f"{self.path}  {colored('--', 'BLACK')}  "
-               f"in:{colored(f'{cache_start_pos:5d}', 'green')} +{len(ids)-cache_start_pos:5d}  {colored('--', 'BLACK')}  ")
+    stderr_log(f"{self.path}  {colored('--', 'BLACK')}  in:{len(ids):5d}  {colored('--', 'BLACK')}  ")
    tmpl = {"id":f"chatcmpl-{uuid.uuid4().hex[:24]}", "object":"chat.completion.chunk", "created":int(time.time()), "model":model_name}
    yield {"choices": [{"index":0, "delta":{"role":"assistant","content":""}, "finish_reason":None}], **tmpl}
    out: list[int] = []
    st = time.perf_counter()
    for next_id in model.generate(ids):
-      if len(out) == 0: stderr_log(f"prefill:{(len(ids)-cache_start_pos)/((pt:=time.perf_counter())-st):4.0f} tok/s  {colored('--', 'BLACK')}  ")
+      if len(out) == 0: stderr_log(f"prefill:{len(ids)/((pt:=time.perf_counter())-st):4.0f} tok/s  {colored('--', 'BLACK')}  ")
      if next_id == eos_id: break
      out.append(next_id)
      yield {"choices": [{"index":0, "delta":{"content":tok.decode([next_id])}, "finish_reason":None}], **tmpl}
    yield {"choices": [{"index":0, "delta":{},"finish_reason":"stop"}], **tmpl}
    if include_usage:
      yield {"choices": [], "usage": {"prompt_tokens": len(ids), "completion_tokens": len(out), "total_tokens": len(ids) + len(out)}, **tmpl}
-    stderr_log(f"gen:{len(out)/(time.perf_counter()-pt):4.0f} tok/s  {colored('--', 'BLACK')}  out:{len(out):5d}\n")
+    stderr_log(f"out:{len(out):5d}  {colored('--', 'BLACK')}  gen: {len(out)/(time.perf_counter()-pt):4.0f} tok/s\n")

  def do_POST(self):
    raw_body = self.rfile.read(int(self.headers.get("Content-Length", "0")))
@@ -379,7 +362,7 @@ if __name__ == "__main__":

  # do benchmark
  if args.benchmark:
-    gen = model.generate(toks:=[bos_id or 0])
+    gen = model.generate(toks:=[bos_id or 0], 0)
    for _ in range(args.benchmark):
      GlobalCounters.reset()
      with Timing(on_exit=lambda x: f", {1e9/x:6.2f} tok/s, {GlobalCounters.global_mem/x:7.2f} GB/s,"
@@ -388,20 +371,17 @@ if __name__ == "__main__":
    exit(0)

  # start server
-  if args.serve:
-    # warmup: run 2 tokens through the model twice to capture the JIT before serving
-    with Context(DEBUG=max(DEBUG.value, 1)):
-      for _ in range(2): list(zip(range(2), model.generate([0])))
-    TCPServerWithReuse(('', args.serve), Handler).serve_forever()
+  if args.serve: TCPServerWithReuse(('', args.serve), Handler).serve_forever()

  # interactive chat
  ids: list[int] = [bos_id] if bos_id is not None else []
  while 1:
+    start_pos = max(len(ids) - 1, 0)
    try:
      ids += tok.role("user") + tok.encode(input('>>> ')) + tok.end_turn(eos_id) + tok.role("assistant")
    except EOFError:
      break
-    for next_id in model.generate(ids):
+    for next_id in model.generate(ids, start_pos):
      sys.stdout.write(tok.decode([next_id]) if next_id != eos_id else "\n\n")
      sys.stdout.flush()
      if next_id == eos_id: break
@@ -1,15 +1,15 @@
 import math
 from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType, sint_to_uop
-from tinygrad.helpers import dedup, get_contraction
+from tinygrad.helpers import all_int, dedup, get_contraction
 from tinygrad.dtype import dtypes, AddrSpace, Invalid
 from tinygrad.renderer import Renderer

-def _dim_max(d:sint) -> int: return d if isinstance(d, int) else int(d.vmax)
-
 def _group_dims(dims:tuple[sint, ...], max_sizes:tuple[int, ...]):
+  # TODO: symbolic shape
+  if not all_int(dims): return dims
  while len(dims) > len(max_sizes) or any(d > m for d,m in zip(dims, max_sizes)):
    for i,m in enumerate(max_sizes):
-      if i < (len(dims)-1) and _dim_max(dims[i]) * _dim_max(dims[i+1]) <= m:
+      if i < (len(dims)-1) and dims[i] * dims[i+1] <= m:
        dims = dims[:i] + (dims[i]*dims[i+1],) + dims[i+2:]
        break
    else: return None
@@ -130,7 +130,7 @@ load_store_folding = PatternMatcher([
  (UPat(Ops.STORE, src=(UPat(Ops.GEP, name="gep"), UPat.var("st")), name="sto"), gep_on_store),
  # put PTRCAT after LOAD
  (UPat(Ops.LOAD, src=(UPat(Ops.PTRCAT, name="cat"),), name="ld", allow_any_len=True),
-   lambda cat,ld: UOp(Ops.VCAT, cat.dtype.base.vec(cat.dtype.vcount), tuple(ld.replace(dtype=x.dtype.base, src=(x,)+ld.src[1:]) for x in cat.src))),
+   lambda cat,ld: UOp(Ops.CAT, cat.dtype.base.vec(cat.dtype.vcount), tuple(ld.replace(dtype=x.dtype.base, src=(x,)+ld.src[1:]) for x in cat.src))),
  # put PTRCAT after STORE
  (UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data")), name="sto"), cat_after_store),
 ])
@@ -181,7 +181,7 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):

  # if it wasn't split, we return None. otherwise we CAT them
  if len(ret) <= 1: return None
-  return UOp(Ops.VCAT, ls.dtype, tuple(ret)) if ls.op is Ops.LOAD else UOp.group(*ret)
+  return UOp(Ops.CAT, ls.dtype, tuple(ret)) if ls.op is Ops.LOAD else UOp.group(*ret)

 def _do_image_fixup(dt:ImageDType, idx:UOp) -> tuple[UOp, UOp, int, int]:
  buf = idx.src[0]
@@ -189,10 +189,7 @@ def _do_image_fixup(dt:ImageDType, idx:UOp) -> tuple[UOp, UOp, int, int]:
  h, w = dt.shape[0], dt.shape[1]
  if IMAGE == 1 and valid is not None:
    h, w = max(ImageDType.valid_dims(dt), key=lambda hw:
-                # maximize number of valids removed
-               (len(_drop_valid_stmts(valid, idx:=uop_given_valid(valid, UOp.vectorize((x//4)%hw[1], x//(4*hw[1]))), *hw)),
-                # and minimize idx complexity (number of nodes)
-                -len(idx.backward_slice)))
+               (len(_drop_valid_stmts(valid, idx:=uop_given_valid(valid, UOp.vectorize((x//4)%hw[1], x//(4*hw[1]))), *hw)), -len(idx.backward_slice)))
    buf = buf.replace(dtype=(dtypes.imageh if dt.itemsize == 2 else dtypes.imagef)((h, w, 4), w * 4 * dt.itemsize))
  oidx = UOp(Ops.VECTORIZE, dtypes.index.vec(2), ((x // 4) % w, (x // (4*w))))
  return x, idx.replace(src=(buf, oidx.valid(valid))), w, h
@@ -51,7 +51,7 @@ def do_expand(root:UOp):
        new_srcs.append(src)
      elif src.dtype.count > 1:
        # put any input dtype > 1 grouped together
-        new_srcs.append(UOp(Ops.VCAT, src.dtype.scalar().vec(expand_sz*src.dtype.count), (src,)*expand_sz))
+        new_srcs.append(UOp(Ops.CAT, src.dtype.scalar().vec(expand_sz*src.dtype.count), (src,)*expand_sz))
      else:
        # repeat the arg
        new_srcs.append(src.broadcast(expand_sz))
@@ -16,7 +16,7 @@ pm_flatten_range = PatternMatcher([
  (UPat((Ops.REDUCE, Ops.STORE, Ops.END), name="r"), flatten_range),
 ])

-def count_divmod(x:UOp) -> int: return sum(u.op in {Ops.IDIV, Ops.MOD} for u in x.backward_slice)
+def count_divmod(x:UOp) -> int: return len([u for u in x.toposort() if u.op in {Ops.IDIV, Ops.MOD}])
 def simplify_merge_adjacent(u:UOp) -> UOp|None:
  reduce_ranges = [x.ranges for x in u.backward_slice_with_self if x.op is Ops.REDUCE]
  # on END we only want to merge adjacent ranges, on REDUCE we want to try all combinations
@@ -6,7 +6,7 @@ import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re
 from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup, ContextVar
 from tinygrad.helpers import unwrap_class_type, suppress_finalizing, select_first_inited, VIZ, CPU_LLVM, CPU_LVP, NV_PTX, CUDA_PTX, NV_NAK
-from tinygrad.helpers import EMULATED_DTYPES, NULL_IR3, NULL_QCOMCL, TracingKey
+from tinygrad.helpers import EMULATED_DTYPES, TracingKey
 from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
 if TYPE_CHECKING: from tinygrad.renderer import Renderer

@@ -353,12 +353,11 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
    if device == "NV": return not CI and not NV_PTX and not NV_NAK
    if device in {"CPU"}: return not CI and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"} and not CPU_LVP
    return device in {"AMD", "CL", "PYTHON", "NULL"}
-  if dtype in dtypes.fp8_ocp:
+  if dtype in dtypes.fp8s:
    if device == "CUDA": return not CI and not CUDA_PTX
    if device == "NV": return not CI and not NV_PTX and not NV_NAK
-    if device == "AMD": return not CI and getattr(Device["AMD"], "target") == (9,5,0)
+    if device == "AMD": return not CI and getattr(Device["AMD"], "target") in {(9,4,2), (9,5,0)}
    return device in {"PYTHON", "NULL"}
-  if dtype in dtypes.fp8_fnuz: return device in {"PYTHON", "NULL"}
  if device == "WEBGPU": return dtype in [dtypes.bool, dtypes.char, dtypes.uchar, dtypes.short,
                                          dtypes.ushort, dtypes.float, dtypes.int32, dtypes.uint32, dtypes.half]
  # for CI GPU and OSX, cl_khr_fp16 isn't supported
@@ -372,7 +371,7 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
    if device in ["CUDA", "NV"]: return not CI
    if device == "CPU" and CPU_LLVM: return OSX
    if device == "PYTHON": return sys.version_info >= (3, 12)
-  if dtype == dtypes.float64: return (device not in {"METAL", "QCOM"} and not (OSX and device == "CL") and not NULL_IR3 and not NULL_QCOMCL
+  if dtype == dtypes.float64: return (device not in {"METAL", "QCOM"} and not (OSX and device == "CL") and not getenv("NULL_IR3")
                                      and dtypes.long not in EMULATED_DTYPES.tolist(dtypes))
  return True

@@ -30,7 +30,6 @@ class InvalidType:
  def __hash__(self): return id(self)
  def __repr__(self): return "Invalid"
  def __reduce__(self): return (InvalidType, ())  # unpickle returns the singleton
-  def __format__(self, spec): return "Invalid"

 Invalid = InvalidType()

@@ -79,14 +78,10 @@ class DType(metaclass=DTypeMetaClass):
    return PtrDType(self.priority, self.bitsize, self.name, self.fmt, self.count, None, self, addrspace, 1, size)
  def scalar(self) -> DType: return self._scalar if self._scalar is not None else self
  def nbytes(self) -> int: raise RuntimeError("only ptr types have nbytes")
-  @functools.cached_property
-  def min(self):
-    if dtypes.is_int(self): return 0 if dtypes.is_unsigned(self) else -2**(self.scalar().bitsize-1)
-    return -float("inf") if dtypes.is_float(self) else False
-  @functools.cached_property
-  def max(self):
-    if dtypes.is_int(self): return 2**(self.scalar().bitsize)-1+self.min
-    return float("inf") if dtypes.is_float(self) else True
+  @property
+  def min(self): return dtypes.min(self)
+  @property
+  def max(self): return dtypes.max(self)

@dataclass(frozen=True, eq=False)
 class PtrDType(DType):
@@ -176,11 +171,21 @@ class dtypes:
    # int is the default. wrap floats in ConstFloat to distinguish -0.0 from 0.0 in cache
    return ConstFloat(float(val)) if dtypes.is_float(dtype) else bool(val) if dtypes.is_bool(dtype) else int(val)
  @staticmethod
+  @functools.cache
+  def min(dtype:DType):
+    if dtypes.is_int(dtype): return 0 if dtypes.is_unsigned(dtype) else -2**(dtype.scalar().bitsize-1)
+    return -float("inf") if dtypes.is_float(dtype) else False
+  @staticmethod
+  @functools.cache
+  def max(dtype:DType):
+    if dtypes.is_int(dtype): return 2**(dtype.scalar().bitsize)-1+dtypes.min(dtype)
+    return float("inf") if dtypes.is_float(dtype) else True
+  @staticmethod
  def finfo(dtype:DType) -> tuple[int, int]:
    """(exponent, mantissa)"""
    if not dtypes.is_float(dtype): raise ValueError(f"{dtype} is not a floating point type")
    return {dtypes.float16: (5, 10), dtypes.bfloat16: (8, 7), dtypes.float32: (8, 23), dtypes.float64: (11, 52),
-            dtypes.fp8e4m3: (4, 3), dtypes.fp8e5m2: (5, 2), dtypes.fp8e4m3fnuz: (4, 3), dtypes.fp8e5m2fnuz: (5, 2)}[dtype]
+            dtypes.fp8e5m2: (5, 2), dtypes.fp8e4m3: (4, 3)}[dtype]
  void: Final[DType] = DType.new(-1, 0, "void", None)
  index: Final[DType] = DType.new(-1, 800, "index", None)
  bool: Final[DType] = DType.new(0, 1, "bool", '?')
@@ -196,8 +201,6 @@ class dtypes:
  _uint256: Final[DType] = DType.new(8, 256, "uint256", None)
  fp8e4m3: Final[DType] = DType.new(9, 8, "float8_e4m3", None)
  fp8e5m2: Final[DType] = DType.new(10, 8, "float8_e5m2", None)
-  fp8e4m3fnuz: Final[DType] = DType.new(9, 8, "float8_e4m3fnuz", None)
-  fp8e5m2fnuz: Final[DType] = DType.new(10, 8, "float8_e5m2fnuz", None)
  float16: Final[DType] = DType.new(11, 16, "half", 'e')
  # bfloat16 has higher priority than float16, so least_upper_dtype(dtypes.int64, dtypes.uint64) = dtypes.float16
  bfloat16: Final[DType] = DType.new(12, 16, "__bf16", None)
@@ -218,9 +221,7 @@ class dtypes:
  default_float: ClassVar[DType] = float32
  default_int: ClassVar[DType] = int32

-  fp8_ocp = (fp8e4m3, fp8e5m2)
-  fp8_fnuz = (fp8e4m3fnuz, fp8e5m2fnuz)
-  fp8s = fp8_ocp + fp8_fnuz
+  fp8s = (fp8e4m3, fp8e5m2)
  floats = fp8s + (float16, bfloat16, float32, float64)
  int8s = (uint8, int8)
  int16s = (uint16, int16)
@@ -242,9 +243,8 @@ def to_dtype(dtype:DTypeLike) -> DType: return dtype if isinstance(dtype, DType)
 # we don't support weak type and complex type
 promo_lattice = { dtypes.bool: [dtypes.int8, dtypes.uint8], dtypes.int8: [dtypes.int16], dtypes.int16: [dtypes.int32], dtypes.int32: [dtypes.int64],
  dtypes.int64: [dtypes.uint64], dtypes.uint8: [dtypes.int16, dtypes.uint16], dtypes.uint16: [dtypes.int32, dtypes.uint32],
-  dtypes.uint32: [dtypes.int64, dtypes.uint64], dtypes.uint64: [dtypes.fp8e4m3, dtypes.fp8e5m2, dtypes.fp8e4m3fnuz, dtypes.fp8e5m2fnuz],
-  dtypes.fp8e4m3: [dtypes.float16, dtypes.bfloat16], dtypes.fp8e5m2: [dtypes.float16, dtypes.bfloat16],
-  dtypes.fp8e4m3fnuz: [dtypes.float16, dtypes.bfloat16], dtypes.fp8e5m2fnuz: [dtypes.float16, dtypes.bfloat16],
+  dtypes.uint32: [dtypes.int64, dtypes.uint64], dtypes.uint64: [dtypes.fp8e4m3, dtypes.fp8e5m2],
+  dtypes.fp8e5m2: [dtypes.float16, dtypes.bfloat16], dtypes.fp8e4m3: [dtypes.float16, dtypes.bfloat16],
  dtypes.float16: [dtypes.float32], dtypes.bfloat16: [dtypes.float32], dtypes.float32: [dtypes.float64], }

@functools.cache
@@ -299,14 +299,10 @@ def float_to_bf16(x):
 _fp8_cfg = {
  dtypes.fp8e4m3: (7, 4, 0x7, 0x3F50000000000000, 0x407D000000000000, 0x7E, 0x3F90000000000000),
  dtypes.fp8e5m2: (15, 3, 0x3, 0x3EE0000000000000, 0x40EE000000000000-1, 0x7B, 0x3F10000000000000),
-  dtypes.fp8e4m3fnuz: (8, 4, 0x7, 0x3F40000000000000, 0x406F000000000000-1, 0x7F, 0x3F80000000000000),
-  dtypes.fp8e5m2fnuz: (16, 3, 0x3, 0x3ED0000000000000, 0x40EE000000000000-1, 0x7F, 0x3F00000000000000),
 }

 def float_to_fp8(x: float, dtype: DType) -> int:
  assert dtype in dtypes.fp8s, "Only for fp8s"
-  if dtype in dtypes.fp8_fnuz and not math.isfinite(x): return 0x80
-  if dtype in dtypes.fp8_fnuz and x == 0.0: return 0x00
  # e4m3 don't support inf, return 0x7f(+NaN) and 0xff(-NaN) to match jax
  # NaN is unordered, can't compare with zero, use math.copysign to get sign
  if dtype == dtypes.fp8e4m3 and not math.isfinite(x): return 0x7f if math.copysign(1, x) > 0 else 0xff
@@ -326,17 +322,16 @@ def float_to_fp8(x: float, dtype: DType) -> int:
    res, half = mantissa >> shift, half_ulp << shift
    round_bits = (xbits | (1 << 52)) & ((half << 1) - 1)
    if round_bits > half or (round_bits == half and res & 1): res += 1
-  return 0 if dtype in dtypes.fp8_fnuz and res == 0 else int(res | sign)  # fnuz has no negative zero
+  return int(res | sign)

 def fp8_to_float(x: int, dtype: DType) -> float:
  assert dtype in dtypes.fp8s, "Only for fp8s"
-  if dtype in dtypes.fp8_fnuz and x == 0x80: return math.nan
  if (x & 0x7F) == 0: return -0.0 if x & 0x80 else 0.0
  bias, sig_bits, *_ = _fp8_cfg[dtype]
  mant_bits, exp_bits = sig_bits - 1, 8 - sig_bits
  exp_max, mant_max = (1 << exp_bits) - 1, (1 << mant_bits) - 1
  sign, exp, mantissa = (x >> 7) & 1, (x >> mant_bits) & exp_max, x & mant_max
-  if dtype not in dtypes.fp8_fnuz and exp == exp_max:
+  if exp == exp_max:
    if dtype == dtypes.fp8e5m2: return math.copysign(math.nan if mantissa else math.inf, -1 if sign else 1)
    if mantissa == mant_max: return math.nan
  val = (mantissa / (mant_max + 1)) * 2 ** (1 - bias) if exp == 0 else (1 + mantissa / (mant_max + 1)) * 2 ** (exp - bias)
@@ -1,7 +1,7 @@
 from dataclasses import dataclass, field
 from tinygrad.uop.ops import UOp, UPat, PatternMatcher, Ops, GroupOp, graph_rewrite, track_rewrites
 from tinygrad.dtype import dtypes, ImageDType
-from tinygrad.helpers import prod, DEBUG, VIZ, pluralize
+from tinygrad.helpers import prod, DEBUG, argsort, VIZ, pluralize, FLOAT16

@dataclass
 class AllocCtx:
@@ -32,8 +32,8 @@ def apply_after(ctx:AllocCtx, u:UOp):
 # CONTIGUOUS and ASSIGN + parents are the only nodes that get updated
 add_tags = PatternMatcher([
  (UPat(Ops.COPY, name="u"), disk_copy_is_buffer),
-  # no tag on copies that are assigned
-  (UPat(Ops.ASSIGN, src=(UPat(), UPat(Ops.COPY, name="c")), name="a"),
+  # no tag on copies/allreduces that are assigned
+  (UPat(Ops.ASSIGN, src=(UPat(), UPat((Ops.COPY, Ops.ALLREDUCE), name="c")), name="a"),
   lambda a,c: a.replace(src=(a.src[0], c.rtag(())), tag=a.tag+c.tag) if a.tag and c.tag else None),
  (UPat(Ops.AFTER, name="u"), apply_after),
  (UPat({Ops.CONTIGUOUS, Ops.ASSIGN}, name="x"), tag_uop),
@@ -51,8 +51,6 @@ def _buffer_like(u:UOp) -> UOp:
  return buffer

 def replace_contig_with_assign(u:UOp):
-  # can't allocate a buffer without a device (e.g., inside a CALL function body with only PARAMs)
-  if u._device is None: return None
  # if size is 0, remove the contig
  if u.size == 0: return u.src[0]
  # no real contig for DISK/TINYFS tensors, they are left alone
@@ -65,6 +63,15 @@ def replace_assign_with_contig(u:UOp):
  if assigned_to.op is not Ops.BUFFER:
    return u.src[1].contiguous(tag=u.tag)

+def found_contiguous(ctx:dict[UOp, UOp], contig:UOp, src:UOp):
+  if (x:=src).op is Ops.CAST and x.dtype == dtypes.half and FLOAT16: x, contig = x.src[0], contig.cast(dtypes.float)
+  while x is not x.base:
+    if x.op is Ops.PERMUTE: contig = contig.permute(argsort(x.marg))
+    elif x.op is Ops.RESHAPE: contig = contig.reshape(x.src[0].shape)
+    else: return None
+    x = x.src[0]
+  ctx[x] = contig
+
 def contiguous_mops_to_view(c:UOp):
  """CONTIGUOUS(MOPS(BUFFER)) → CONTIGUOUS(BUFFER_VIEW) when movement ops collapse to a contiguous range."""
  src = c.src[0]
@@ -94,12 +101,8 @@ def transform_precompiled_call(c:UOp) -> UOp|None:
  if not c.arg.precompile: return None
  if c.src[0].op is Ops.SINK: return None
  out = _buffer_like(c)
-  input_buffers = tuple(x.contiguous() if x.op not in {Ops.AFTER, Ops.BIND} else x for x in c.src[1:])
  fxn = out.param_like(len(c.src)-1).assign(c.src[0]).sink()
-  ret = out.after(c.replace(src=(fxn, *input_buffers, out), dtype=dtypes.void, tag=None))
-  # if the CALL has symbolic shapes, shrink the max-sized output to the actual symbolic shape
-  if any(isinstance(s, UOp) for s in c.shape): ret = ret.shrink(tuple((0, s) for s in c.shape))
-  return ret
+  return out.after(c.replace(src=(fxn,)+tuple(x.contiguous() if x.op is not Ops.AFTER else x for x in c.src[1:])+(out,), dtype=dtypes.void, tag=None))

 # NOTE: adding rules to here is bad. these all need to run before the schedule cache
 pm_early_transform_tensor_graph = PatternMatcher([
@@ -109,11 +112,15 @@ pm_early_transform_tensor_graph = PatternMatcher([
  # CONTIGUOUS(MOPS(BUFFER/BUFFER_VIEW)) → CONTIGUOUS(BUFFER_VIEW) when movement ops collapse to contiguous range
  (UPat(Ops.CONTIGUOUS, src=(UPat(GroupOp.Movement),), name="c"), contiguous_mops_to_view),

+  # *** CONTIGUOUS replacement hack for openpilot ***
+  (UPat(Ops.CONTIGUOUS, src=(UPat((*GroupOp.Movement, Ops.CAST), name="src"),), name="contig"), found_contiguous),
+  # replace ALU sources with contiguous versions found above
+  (UPat(GroupOp.ALU, name="alu"), lambda ctx,alu: alu.replace(src=new_src) if (new_src:=tuple(ctx.get(s, s) for s in alu.src)) != alu.src else None),
+
  # add CONTIGUOUS to tagged UOps
-  (UPat(GroupOp.All-{Ops.CONTIGUOUS, Ops.ASSIGN, Ops.AFTER}, name="x"),
-   lambda x: x.rtag(None).contiguous(tag=x.tag) if x.tag else x.replace(tag=None)),
-  # remove extra CONTIGUOUS on AFTER (only when after target is contiguous)
-  (UPat(Ops.CONTIGUOUS, src=(UPat(Ops.AFTER, name="a"),), name="c"),
+  (UPat(GroupOp.All-{Ops.CONTIGUOUS, Ops.ASSIGN}, name="x"), lambda x: x.rtag(None).contiguous(tag=x.tag) if x.tag else x.replace(tag=None)),
+  # remove extra CONTIGUOUS on ASSIGN (only when assign target is contiguous)
+  (UPat(Ops.CONTIGUOUS, src=(UPat(Ops.ASSIGN, name="a"),), name="c"),
   lambda a,c: a.replace(tag=(a.tag or ())+(c.tag or ())) if a.src[0].has_buffer_identity() else None),
  # replace ASSIGN with CONTIGUOUS
  (UPat(Ops.ASSIGN, name="u"), replace_assign_with_contig),
@@ -129,7 +136,7 @@ def untag_and_append(ctx:AllocCtx, x:UOp):
  for t in x.tag:
    original_uop: UOp = ctx.uop_list[t]
    replace_uop = ret
-    while replace_uop.op is Ops.AFTER: replace_uop = replace_uop.src[0]
+    while replace_uop.op is Ops.ASSIGN: replace_uop = replace_uop.src[0]
    ctx.buffer_map[original_uop] = replace_uop.shrink_to(original_uop.shape)
  ctx.assigns.append(ret)
  return ret
@@ -143,10 +150,10 @@ def replace_input_buffer(ctx:AllocCtx, b:UOp):
                   b._min_max if b.op is Ops.BIND else None, b.src[0].arg[0] if b.op is Ops.BIND else None)

 pm_finalize_call = PatternMatcher([
-  (UPat(Ops.AFTER, name="x"), untag_and_append),
+  (UPat(Ops.ASSIGN, name="x"), untag_and_append),
  (UPat(Ops.AFTER, name="x"), append_after),
  (UPat(Ops.COPY, name="x"), lambda ctx,x: append_after(ctx,x) if isinstance(x.device, str) and x.device.startswith(("DISK", "TINYFS")) else None),
-  # remove unique from const. TODO: this is copied in function.py
+  # replace UNIQUE with LUNIQUE for CONST cache key normalization
  (UPat(Ops.CONST, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE, name="d")), name="b"), lambda b,d: b.replace(src=(d,))),
 ])

@@ -172,7 +179,7 @@ def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
  big_sink = graph_rewrite(big_sink, add_tags, ctx=ctx, bottom_up=True, name="number the uops")

  # here we can break the tensor graph. this is the only place you need to maintain numbered tags
-  big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, name="early transform tensor graph")
+  big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, ctx={}, name="early transform tensor graph")

  # here we construct the final buffer_map. this is everything that will go into the tensor map
  graph_rewrite(big_sink, pm_finalize_call, ctx=ctx, name="finalize call")
@@ -114,9 +114,9 @@ class GraphRunner(Runner):
          assert ji.prg.p.local_size is not None
          self.launch_dims_base[j] = (tuple(ji.prg.p.global_size), tuple(ji.prg.p.local_size))

-    # used in MultiGraphRunner. tracks (offset, end, dep) ranges per base buffer id to handle suballocated buffers correctly.
-    self.w_dependency_map: dict[int, list[tuple[int, int, Any]]] = collections.defaultdict(list)
-    self.r_dependency_map: dict[int, list[tuple[int, int, Any]]] = collections.defaultdict(list)
+    # used in MultiGraphRunner. the ints are id() of _bufs
+    self.w_dependency_map: dict[int, Any] = {}
+    self.r_dependency_map: dict[int, list[Any]] = collections.defaultdict(list)

    assert jit_cache[0].prg is not None
    super().__init__(colored(f"<batched {len(jit_cache)}>", "cyan"), jit_cache[0].prg.device.split(":")[0], estimates.simplify())
@@ -132,22 +132,19 @@ class GraphRunner(Runner):
      yield j, (dims[gl] if gl is not None else self.launch_dims_base[j][0]), (dims[lc] if lc is not None else self.launch_dims_base[j][1])

  def _access_resources(self, bufs:list[Buffer], write:list[int], new_dependency:Any):
+    # To synchronize access to resources, we monitor the necessary prerequisites for accessing each resource,
+    # whether for write or read operations. A resource can be accessed by either a single writer or multiple readers.
    wait_nodes = []
+
    for i,buf in enumerate(bufs):
-      key, s, e = id(buf.base._buf), buf.offset, buf.offset + buf.nbytes
-      wait_nodes += [dep for st,en,dep in self.w_dependency_map[key] if st < e and s < en]
-      if i in write: wait_nodes += [dep for st,en,dep in self.r_dependency_map[key] if st < e and s < en]
-    for i,buf in enumerate(bufs):
-      key, s, e = id(buf.base._buf), buf.offset, buf.offset + buf.nbytes
+      if id(buf.base._buf) in self.w_dependency_map: wait_nodes.append(self.w_dependency_map[id(buf.base._buf)])
      if i in write:
-        for dmap in [self.w_dependency_map, self.r_dependency_map]:
-          kept = []
-          for st,en,dep in dmap[key]:
-            if st < min(s, en): kept.append((st, min(s, en), dep))
-            if max(e, st) < en: kept.append((max(e, st), en, dep))
-          dmap[key] = kept
-        self.w_dependency_map[key].append((s, e, new_dependency))
-      else: self.r_dependency_map[key].append((s, e, new_dependency))
+        if id(buf.base._buf) in self.r_dependency_map: wait_nodes.extend(self.r_dependency_map.pop(id(buf.base._buf)))
+
+    for i,buf in enumerate(bufs):
+      if i in write: self.w_dependency_map[id(buf.base._buf)] = new_dependency
+      else: self.r_dependency_map[id(buf.base._buf)].append(new_dependency)
+
    return list({id(x):x for x in wait_nodes}.values())

  @staticmethod
@@ -360,8 +357,9 @@ class TinyJit(Generic[ReturnType]):
        jit_cache = pruned

      # memory planning (optional)
-      copies = [(cast(Buffer,ji.bufs[0]),cast(Buffer,ji.bufs[1])) for ji in jit_cache if isinstance(ji.prg, (BufferXfer, BufferCopy, EncDec))]
-      assigned = _internal_memory_planner([cast(list[Buffer], item.bufs) for item in jit_cache], copies, debug_prefix="JIT ")
+      # Exclude buffers involved in transfer ops to preserve parallelism.
+      noopt_buffers = {b for ji in jit_cache if isinstance(ji.prg, (BufferXfer, BufferCopy, EncDec)) for b in ji.bufs}
+      assigned = _internal_memory_planner([cast(list[Buffer], item.bufs) for item in jit_cache], noopt_buffers, debug_prefix="JIT ")
      jit_cache = [replace(item, bufs=[assigned.get(b,b).ensure_allocated() for b in item.bufs if b is not None]) for item in jit_cache]

      input_replace = get_input_replace(jit_cache, input_buffers)
@@ -7,50 +7,43 @@ from tinygrad.uop.ops import Ops
 from tinygrad.dtype import dtypes, ImageDType
 from tinygrad.runtime.support.memory import TLSFAllocator

-LaneKey = tuple[str, int]
-
 # **************** memory planning ****************

-def _internal_memory_planner(buffers:list[list[Buffer]], copies:list[tuple[Buffer, Buffer]]|None=None,
-                             ignore_checks=False, debug_prefix="") -> dict[Buffer, Buffer]:
+def _internal_memory_planner(buffers:list[list[Buffer]], noopt_buffers=None, ignore_checks=False, debug_prefix="") -> dict[Buffer, Buffer]:
  if NO_MEMORY_PLANNER: return {}
  first_appearance, last_appearance, buf_to_opt = {}, {}, set()
  for i,u in enumerate(buffers):
    for buf in u:
-      if not ignore_checks and (buf.is_allocated() or buf.base.is_allocated() or buf.uop_refcount > 0): continue
+      should_skip = buf.is_allocated() or buf.base.is_allocated() or buf.uop_refcount > 0 or (noopt_buffers is not None and buf.base in noopt_buffers)
+      if not ignore_checks and should_skip: continue
      if buf.base not in first_appearance: first_appearance[buf.base] = i
      last_appearance[buf.base] = i
      buf_to_opt.add(buf)

-  # Separate copy and compute buffers into different lanes and defer cross-queue frees to avoid introducing dependencies (copy->compute->copy)
-  copy_dsts, copy_srcs = ({dst.base for dst,_ in copies}, {src.base for _,src in copies}) if copies else (set(), set())
-  def _key(buf) -> LaneKey: return (buf.device, 1 if buf in copy_dsts or buf in copy_srcs else 0)
-  buf_hold = {buf: last_appearance[buf] - first_appearance[buf] + 1 for buf in first_appearance if buf in copy_dsts or buf in copy_srcs}
-
  # Sort buffer operations in timeline order. Two events: buffer is allocated or buffer is freed.
  buffer_requests = sorted([((first_appearance[buf], True), buf) for buf in first_appearance.keys()] + \
-    [((last_appearance[buf] + 1 + buf_hold.get(buf, 0), False), buf) for buf in first_appearance.keys()], key=lambda x: x[0])
-  total_memory = sum(round_up(buf.nbytes, BLK:=0x1000) for buf in first_appearance.keys()) * 2 # *2 for fragmentation (which is about 15%)
+                           [((last_appearance[buf] + 1, False), buf) for buf in first_appearance.keys()], key=lambda x: x[0])
+  total_memory = sum(round_up(buf.nbytes, min_block_size:=0x1000) for buf in first_appearance.keys()) * 2 # *2 for fragmentation (which is about 15%)

  # Try to suballocate from a shared buffer managed by global_planner using TLSFAllocator.
  # Also track buffer replacements for buffers that do not support suballocation.
  buffer_replace:dict[Buffer, tuple[Buffer|None, int|None]] = {}
  reuse_buffers:dict[tuple, list[Buffer]] = defaultdict(list)
-  global_planner:dict[LaneKey, tuple[int, TLSFAllocator]] = defaultdict(lambda: (0, TLSFAllocator(total_memory, block_size=BLK, lv2_cnt=32)))
+  global_planner:dict[str, tuple[int, TLSFAllocator]] = defaultdict(lambda: (0, TLSFAllocator(total_memory, block_size=min_block_size, lv2_cnt=32)))
  for (_, is_open_ev), buf in buffer_requests:
    # Check if suballocation is possible for the given buffer and device.
    if hasattr(Device[buf.device].allocator, "_offset") and not isinstance(buf.dtype, ImageDType):
-      if is_open_ev: buffer_replace[buf] = (None, global_planner[_key(buf)][1].alloc(round_up(buf.nbytes, BLK)))
-      else: global_planner[_key(buf)][1].free(cast(int, buffer_replace[buf][1]))
-      global_planner[_key(buf)] = (max(global_planner[_key(buf)][0], buffer_replace[buf][1] + buf.nbytes), global_planner[_key(buf)][1])
+      if is_open_ev: buffer_replace[buf] = (None, global_planner[buf.device][1].alloc(round_up(buf.nbytes, 0x1000)))
+      else: global_planner[buf.device][1].free(cast(int, buffer_replace[buf][1]))
+      global_planner[buf.device] = (max(global_planner[buf.device][0], buffer_replace[buf][1] + buf.nbytes), global_planner[buf.device][1])
    else:
-      key = (_key(buf), buf.dtype, buf.options, buf.nbytes)
+      key = (buf.device, buf.dtype, buf.options, buf.nbytes)
      if is_open_ev: buffer_replace[buf] = (reuse_buffers[key].pop(), None) if key in reuse_buffers and len(reuse_buffers[key]) > 0 else (buf, None)
      else: reuse_buffers[key].append(cast(Buffer, buffer_replace[buf][0]))

  # Allocate global buffers based on the memory planner.
-  global_buffers = {key: Buffer(key[0], round_up(sz, BLK), dtypes.int8) for key, (sz, _) in global_planner.items()}
-  buffer_resolve:dict[Buffer, tuple[Buffer, int|None]] = {buf: (base or global_buffers[_key(buf)], off) for buf,(base,off) in buffer_replace.items()}
+  global_buffers = {dev: Buffer(dev, round_up(sz, 0x1000), dtypes.int8) for dev, (sz, _) in global_planner.items()}
+  buffer_resolve:dict[Buffer, tuple[Buffer, int|None]] = {buf: (base or global_buffers[buf.device], off) for buf,(base,off) in buffer_replace.items()}

  # Assign buffers. First, assign full buffers (not sub-buffers).
  assigned:dict[Buffer, Buffer] = {}
@@ -73,5 +66,5 @@ def _internal_memory_planner(buffers:list[list[Buffer]], copies:list[tuple[Buffe
 def memory_planner(schedule:list[ExecItem]) -> list[ExecItem]:
  # Exclude buffers involved in load ops (e.g transfers) to preserve parallelism in graphs.
  assigned = _internal_memory_planner([[b for b in si.bufs if b is not None] for si in schedule],
-                                      copies=[(cast(Buffer,si.bufs[0]),cast(Buffer,si.bufs[1])) for si in schedule if si.ast.op is Ops.COPY])
+                                      noopt_buffers={b for si in schedule if si.ast.op is not Ops.SINK for b in si.bufs if b is not None})
  return [ExecItem(si.ast, [assigned.get(x, x) if x is not None else None for x in si.bufs], si.metadata, si.fixedvars) for si in schedule]
@@ -93,8 +93,8 @@ class BufferXfer(BufferCopy):
  def copy(self, dest, src): dest.allocator._transfer(dest._buf, src._buf, dest.nbytes, src_dev=src.allocator.dev, dest_dev=dest.allocator.dev)

 class EncDec(Runner):
-  def __init__(self, cf:UOp, total_sz:int, device:str):
-    self.shape, self.pos_var = tuple(s.arg for s in cf.src if s.op is Ops.CONST), cf.variables()[0].expr
+  def __init__(self, encdec:UOp, total_sz:int, device:str):
+    self.shape, self.pos_var = encdec.arg[0], encdec.variables()[0].expr
    name = f"enc/dec {total_sz/1e6:7.2f}M, HEVC" if total_sz >= 1e6 else f"enc/dec {total_sz:8d}, HEVC"
    super().__init__(colored(name, "yellow"), device, Estimates(lds=total_sz, mem=total_sz))
  def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int], wait=False):
@@ -130,7 +130,7 @@ si_lowerer = PatternMatcher([
  (UPat(Ops.COPY, name="copy"), lambda ctx,copy: (BufferXfer(ctx[0].nbytes, ctx[0].device, ctx[1].device) \
      if hasattr(alc:=Device[ctx[0].device].allocator, '_transfer') and alc.supports_transfer and all_same([x.device.split(":")[0] for x in ctx]) \
      else BufferCopy(ctx[0].nbytes, ctx[0].device, ctx[1].device))),
-  (UPat(Ops.CUSTOM_FUNCTION, arg="encdec", name="cf"), lambda ctx,cf: EncDec(cf, ctx[0].nbytes, ctx[0].device)),
+  (UPat(Ops.ENCDEC, name="encdec"), lambda ctx,encdec: EncDec(encdec, ctx[0].nbytes, ctx[1].device)),
 ])

@dataclass
@@ -71,7 +71,7 @@ def linear_to_schedule(linear:UOp) -> list[ExecItem]:
      base = buf_uops[1].buffer
      assert isinstance(base, Buffer), "base can't be MultiBuffer"
      buffers[buf_uops[0]] = base.view(buf_uops[0].arg, ast.dtype, ast.arg[1]*base.dtype.itemsize)
-    ubufs = [b.buffer for b in buf_uops if b.op is not Ops.BIND]
+    ubufs = [b.buffer for b in buf_uops]
    metadata = si.arg.metadata
    if any(isinstance(x, MultiBuffer) for x in ubufs):
      assert all(isinstance(x, MultiBuffer) for x in ubufs), "kernel must all be multibuffer"
@@ -131,8 +131,17 @@ def lower_sink_to_linear(function:UOp) -> UOp|None:
          f" | {len(UOpMetaClass.ucache):7d} uops in cache"+("" if frm is None else f" | {frm.filename}:{frm.lineno}"))
  return linear

+def soft_allreduce(c:UOp, a:UOp):
+  from tinygrad.schedule.multi import handle_allreduce
+  to = c.src[1].param_like(0)
+  src = c.src[2].param_like(1)
+  red = UOp(Ops.ALLREDUCE, dtype=a.arg, src=(src, a.src[1]), arg=a.arg)
+  return to.assign(handle_allreduce(src, red)).sink().call(*c.src[1:])
+
 pm_schedule = PatternMatcher([
  (UPat(Ops.SINK, name="function"), lower_sink_to_linear),
+  # soft handler of allreduce
+  (UPat(Ops.CALL, src=(UPat(Ops.ALLREDUCE, name="a"),), allow_any_len=True, name="c"), soft_allreduce),
 ])

@track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[0]))}")
@@ -13,8 +13,6 @@ pm_ctx = PatternMatcher([
  (UPat((Ops.BUFFER, Ops.BIND), name="x"), add_to_ctx),
  (UPat((Ops.ASSIGN, Ops.CONTIGUOUS), name="x"),
   lambda ctx,x: add_to_ctx(ctx,x) if not x.op_in_backward_slice_with_self(Ops.PARAM) else None),
-  # strip UNIQUE from unique consts — they don't need buffer identity inside function bodies
-  (UPat(Ops.CONST, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="x"), lambda ctx,x: x.replace(src=(x.src[1],))),
 ])

 ReturnType = TypeVar('ReturnType')
@@ -1,7 +1,5 @@
 from __future__ import annotations
-import time
-START_TIME = time.perf_counter()
-import os, functools, platform, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc
+import os, functools, platform, time, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc
 from collections import defaultdict
 import subprocess, shutil, math, types, copyreg, inspect, importlib, decimal, itertools
 from dataclasses import dataclass, field
@@ -176,7 +174,7 @@ class ContextVar(Generic[T]):
    return [getattr(obj, x) if obj else x for x in self.value.split(',') if x]

 DEBUG, BEAM, NOOPT = ContextVar("DEBUG", 0), ContextVar("BEAM", 0), ContextVar("NOOPT", 0)
-IMAGE, FLOAT16, OPENPILOT_HACKS = ContextVar("IMAGE", 0), ContextVar("FLOAT16", 0), ContextVar("OPENPILOT_HACKS", 0)
+IMAGE, FLOAT16 = ContextVar("IMAGE", 0), ContextVar("FLOAT16", 0)
 JIT, JIT_BATCH_SIZE = ContextVar("JIT", 2 if OSX and ARCH_X86 else 1), ContextVar("JIT_BATCH_SIZE", 32)
 WINO, CAPTURING, TRACEMETA = ContextVar("WINO", 0), ContextVar("CAPTURING", 1), ContextVar("TRACEMETA", 1)
 USE_TC, TC_SELECT, TC_OPT, AMX = ContextVar("TC", 1), ContextVar("TC_SELECT", -1), ContextVar("TC_OPT", 0), ContextVar("AMX", 0)
@@ -195,8 +193,7 @@ CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasat
 CPU_CC, CPU_LLVM, CPU_LVP = ContextVar("CPU_CC", ""), ContextVar("CPU_LLVM", 0), ContextVar("CPU_LVP", 0)
 NV_CC, NV_PTX, NV_NAK, NV_NVCC = ContextVar("NV_CC", ""), ContextVar("NV_PTX", 0), ContextVar("NV_NAK", 0), ContextVar("NV_NVCC", 0)
 CUDA_CC, CUDA_PTX, CUDA_NVCC = ContextVar("CUDA_CC", ""), ContextVar("CUDA_PTX", 0), ContextVar("CUDA_NVCC", 0)
-NULL_QCOMCL, NULL_IR3, NULL_NAK = ContextVar("NULL_QCOMCL", 0), ContextVar("NULL_IR3", 0), ContextVar("NULL_NAK", 0)
-NULL_ALLOW_COPYOUT = ContextVar("NULL_ALLOW_COPYOUT", 0)
+NULL_IR3, NULL_NAK, NULL_ALLOW_COPYOUT = ContextVar("NULL_IR3", 0), ContextVar("NULL_NAK", 0), ContextVar("NULL_ALLOW_COPYOUT", 0)
 AMD_CC, AMD_LLVM, AMD_HIPCC  = ContextVar("AMD_CC", ""), ContextVar("AMD_LLVM", 0), ContextVar("AMD_HIPCC", 0)
 QCOM_CC, QCOM_IR3 = ContextVar("QCOM_CC", ""), ContextVar("QCOM_IR3", 0)
 # VIZ implies PROFILE, but you can run PROFILE without VIZ
@@ -1,5 +1,5 @@
 from __future__ import annotations
-import math, functools
+import math
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import prod, make_tuple, flatten, USE_ATOMICS
@@ -343,10 +343,6 @@ def _embedding_fwd(weight:Tensor, idx:Tensor) -> Tensor:
  arange = Tensor.arange(weight.shape[0], requires_grad=False, device=weight.device)
  return (arange == idx.unsqueeze(-1)).unsqueeze(-1).where(weight, 0).sum(-2, dtype=weight.dtype)

-@functools.cache
-def _embedding_fwd_fxn(wp, ip, device):
-  return _embedding_fwd(Tensor(wp, device=device), Tensor(ip, device=device))
-
 class Embedding:
  """
  A simple lookup table that stores embeddings of a fixed dictionary and size.
@@ -363,9 +359,7 @@ class Embedding:

  def __call__(self, idx:Tensor) -> Tensor:
    if not dtypes.is_int(idx.dtype): raise TypeError(f"Expected integer dtype for index in embedding, got {idx.dtype}")
-    if USE_ATOMICS:
-      fxn = _embedding_fwd_fxn(self.weight.as_param(0).uop, idx.as_param(1).uop, self.weight.device)
-      return Tensor.call(self.weight, idx, fxn=fxn, grad_fxn=_embedding_bwd)
+    if USE_ATOMICS: return Tensor.call(self.weight, idx, fxn=_embedding_fwd(self.weight.as_param(0), idx.as_param(1)), grad_fxn=_embedding_bwd)
    return _embedding_fwd(self.weight, idx)

 class LSTMCell:
@@ -153,19 +153,21 @@ class OnnxPBParser:

  def _parse_ModelProto(self) -> dict:
    """Entry point for parsing the ONNX model."""
-    graph: dict|None = None
-    opset_imports: list[OpSetId] = []
+    obj: dict[str, Any] = {"opset_import": []}
    for fid, wire_type in self._parse_message(self.reader.len):
      match fid:
-        case 7: graph = self._parse_GraphProto()
-        case 8: opset_imports.append(self._parse_OperatorSetIdProto())
+        case 4: obj["domain"] = self.reader.read_string()
+        case 5: obj["model_version"] = self.reader.read_int64()
+        case 7: obj["graph"] = self._parse_GraphProto()
+        case 8: obj["opset_import"].append(self._parse_OperatorSetIdProto())
        case _: self.reader.skip_field(wire_type)
-    assert graph is not None
+
    # update opset version
-    versions = {opset.domain: opset.version for opset in opset_imports}
-    graph["node"] = [OnnxNode(n.op, OpSetId(n.opset_id.domain, versions.get(n.opset_id.domain, 1)), n.inputs, n.outputs, n.opts)
-                     for n in graph["node"]]
-    return graph
+    opset_imports = {Domain.from_onnx(x.get('domain')):x.get('version', 1) for x in obj["opset_import"]}
+    for n in obj["graph"]["node"]:
+      n_ = n["parsed_node"]
+      n["parsed_node"] = OnnxNode(n_.op, OpSetId(n_.opset_id.domain, opset_imports.get(n_.opset_id.domain, 1)), n_.inputs, n_.outputs, n_.opts)
+    return obj

  def _parse_GraphProto(self) -> dict:
    obj: dict[str, Any] = {"node": [], "initializer": [], "input": [], "output": []}
@@ -179,23 +181,26 @@ class OnnxPBParser:
        case _: self.reader.skip_field(wire_type)
    return obj

-  def _parse_NodeProto(self) -> OnnxNode:
-    inputs: list[str] = []
-    outputs: list[str] = []
-    attributes: list[tuple[str, Any]] = []
-    domain: str|None = None
-    op_type = ""
+  def _parse_NodeProto(self) -> dict:
+    obj: dict[str, Any] = {"input": [], "output": [], "attribute": [], "domain": None}
    for fid, wire_type in self._parse_message(self._decode_end_pos()):
      match fid:
-        case 1: inputs.append(self.reader.read_string())
-        case 2: outputs.append(self.reader.read_string())
-        case 4: op_type = self.reader.read_string()
-        case 5: attributes.append(self._parse_AttributeProto())
-        case 7: domain = self.reader.read_string()
+        case 1: obj["input"].append(self.reader.read_string())
+        case 2: obj["output"].append(self.reader.read_string())
+        case 3: obj["name"] = self.reader.read_string()
+        case 4: obj["op_type"] = self.reader.read_string()
+        case 5: obj["attribute"].append(self._parse_AttributeProto())
+        case 6: obj["doc_string"] = self.reader.read_string()
+        case 7: obj["domain"] = self.reader.read_string()
        case _: self.reader.skip_field(wire_type)
-    return OnnxNode(op_type, OpSetId(Domain.from_onnx(domain), 1), tuple(inputs), tuple(outputs), dict(attributes))

-  def _parse_TensorProto(self) -> tuple[str, Tensor]:
+    # parse node
+    attributes = {attr_dict["name"]: attr_dict[AttributeType(attr_dict["type"]).to_field_name()] for attr_dict in obj["attribute"]}
+    opset_id = OpSetId(Domain.from_onnx(obj.get('domain')), 1)  # default version, to be updated later in _parse_ModelProto
+    obj["parsed_node"] = OnnxNode(obj["op_type"], opset_id, tuple(obj["input"]), tuple(obj["output"]), attributes)
+    return obj
+
+  def _parse_TensorProto(self) -> dict:
    obj: dict[str, Any] = {"dims": []}
    for fid, wire_type in self._parse_message(self._decode_end_pos()):
      match fid:
@@ -215,16 +220,18 @@ class OnnxPBParser:
    # load external data
    if self.load_external_data and obj.get("data_location", 0) == 1:
      if "external_data" not in obj: raise ValueError("no external_data")
-      ext = dict(obj["external_data"])
-      if "location" not in ext: raise ValueError("no location in external_data")
-      offset = int(ext.get("offset", "0"))
-      length = int(ext["length"]) if "length" in ext else None
+      location, length, offset = None, None, 0
+      for kv in obj["external_data"]:
+        if kv["key"] == "location": location = kv["value"]
+        elif kv["key"] == "offset": offset = int(kv["value"])
+        elif kv["key"] == "length": length = int(kv["value"])
+      if location is None: raise ValueError("no location in external_data")

      if self.file_path is None:
        if isinstance(self.tensor.device, str) and self.tensor.device.startswith("DISK:"):
          self.file_path = pathlib.Path(self.tensor.device[5:])
        else: raise ValueError("onnx external_data needs the origin file path, try passing onnx file path to onnx_load")
-      ext_path = self.file_path.parent.joinpath(ext["location"])
+      ext_path = self.file_path.parent.joinpath(location)
      if not ext_path.exists(): raise FileNotFoundError(f"external location not exists: {ext_path}")

      ext_tensor = Tensor(ext_path)
@@ -234,20 +241,23 @@ class OnnxPBParser:
    # parse tensor
    to_dtype = dtype_fallback(true_dtype := OnnxDataType(obj['data_type']).to_dtype(), "buffer parse")
    shape = tuple(obj['dims'])
-    data_fields = [f for f in ('float_data','int32_data','int64_data','double_data','uint64_data','raw_data') if f in obj]
-    data = obj[get_single_element(data_fields)]
-    name = obj.get("name", "")
-    if not isinstance(data, Tensor): return name, Tensor(data, dtype=to_dtype).reshape(shape)
-    assert data.dtype == dtypes.uint8, data
+    present_fields = [field for field in ['float_data', 'int32_data', 'int64_data', 'double_data', 'uint64_data', 'raw_data'] if field in obj]
+    assert len(present_fields) == 1, f"only 1 data field is allowed from {obj=}"
+    data = obj[present_fields[0]]
+    if not isinstance(data, Tensor):
+      obj["parsed_tensor"] = Tensor(data, dtype=to_dtype).reshape(shape)
+      return obj
+    assert isinstance(data, Tensor) and data.dtype == dtypes.uint8, data
    data = data.bitcast(true_dtype).reshape(shape)
    data = data.to(Device.DEFAULT) if true_dtype is to_dtype else data.to("cpu").cast(to_dtype).to(Device.DEFAULT)
    # const folding
    if shape == ():
      if data.dtype == dtypes.float16 and sys.version_info < (3, 12): data = data.cast(dtypes.float32)
      data = Tensor(data.item(), dtype=to_dtype).reshape(shape)
-    return name, data
+    obj["parsed_tensor"] = data
+    return obj

-  def _parse_AttributeProto(self) -> tuple[str, Any]:
+  def _parse_AttributeProto(self) -> dict:
    obj: dict[str, Any] = {"floats": [], "ints": [], "strings": []}
    for fid, wire_type in self._parse_message(self._decode_end_pos()):
      match fid:
@@ -255,7 +265,7 @@ class OnnxPBParser:
        case 2: obj["f"] = self.reader.read_float()
        case 3: obj["i"] = self.reader.read_int64()
        case 4: obj["s"] = self.reader.read_bytes().data().tobytes().decode("utf8")
-        case 5: obj["t"] = self._parse_TensorProto()[1]
+        case 5: obj["t"] = self._parse_TensorProto()['parsed_tensor']
        case 6: obj["g"] = OnnxRunner._from_subgraph(self._parse_GraphProto())
        case 7: obj["floats"].append(self.reader.read_float())
        case 8: obj["ints"].append(self.reader.read_int64())
@@ -263,22 +273,26 @@ class OnnxPBParser:
        case 20: obj["type"] = self.reader.read_int64()
        case _: self.reader.skip_field(wire_type)
    obj["floats"], obj["ints"], obj["strings"] = tuple(obj["floats"]), tuple(obj["ints"]), tuple(obj["strings"])
-    return obj["name"], obj[AttributeType(obj["type"]).to_field_name()]
+    return obj

-  def _parse_ValueInfoProto(self) -> tuple[str, OnnxValue|None]:
-    name, type_obj = "", None
+  def _parse_ValueInfoProto(self) -> dict:
+    obj: dict[str, Any] = {}
    for fid, wire_type in self._parse_message(self._decode_end_pos()):
      match fid:
-        case 1: name = self.reader.read_string()
-        case 2: type_obj = self._parse_TypeProto()
+        case 1: obj["name"] = self.reader.read_string()
+        case 2: obj["type"] = self._parse_TypeProto()
        case _: self.reader.skip_field(wire_type)
-    if type_obj is None: return name, None
+
+    # parse type
+    if "type" not in obj: return {**obj, "parsed_type": None}
+    type_obj = obj["type"]
    if is_optional := "optional_type" in type_obj: type_obj = type_obj["optional_type"]["elem_type"]
    if is_sequence := "sequence_type" in type_obj: type_obj = type_obj["sequence_type"]["elem_type"]
    assert "tensor_type" in type_obj, type_obj
    shape_dims = type_obj['tensor_type'].get('shape', {}).get('dim', [])
-    return name, OnnxValue(tuple(d.get('dim_param') or d.get('dim_value') for d in shape_dims),
-                           OnnxDataType(type_obj['tensor_type']['elem_type']).to_dtype(), is_optional, is_sequence)
+    obj['parsed_type'] = OnnxValue(tuple(d.get('dim_param') or d.get('dim_value') for d in shape_dims),
+                                   OnnxDataType(type_obj['tensor_type']['elem_type']).to_dtype(), is_optional, is_sequence)
+    return obj

  def _parse_TypeProto(self) -> dict:
    obj: dict[str, Any] = {}
@@ -324,24 +338,23 @@ class OnnxPBParser:
        case _: self.reader.skip_field(wire_type)
    return obj

-  def _parse_StringStringEntryProto(self) -> tuple[str, str]:
-    key, value = "", ""
+  def _parse_StringStringEntryProto(self) -> dict:
+    obj: dict[str, Any] = {}
    for fid, wire_type in self._parse_message(self._decode_end_pos()):
      match fid:
-        case 1: key = self.reader.read_string()
-        case 2: value = self.reader.read_string()
+        case 1: obj["key"] = self.reader.read_string()
+        case 2: obj["value"] = self.reader.read_string()
        case _: self.reader.skip_field(wire_type)
-    return key, value
+    return obj

-  def _parse_OperatorSetIdProto(self) -> OpSetId:
-    domain: str|None = None
-    version = 1
+  def _parse_OperatorSetIdProto(self) -> dict:
+    obj: dict[str, Any] = {}
    for fid, wire_type in self._parse_message(self._decode_end_pos()):
      match fid:
-        case 1: domain = self.reader.read_string()
-        case 2: version = self.reader.read_int64()
+        case 1: obj["domain"] = self.reader.read_string()
+        case 2: obj["version"] = self.reader.read_int64()
        case _: self.reader.skip_field(wire_type)
-    return OpSetId(Domain.from_onnx(domain), version)
+    return obj

 # ***** python const *****
 required_input_python_consts: dict[str, tuple[int, ...]] = {
@@ -367,15 +380,16 @@ class OnnxRunner:
    model_path: The ONNX model, provided as a file path (a string or Path object) or a Tensor.
  """
  def __init__(self, model_path: Tensor | str | pathlib.Path):
-    self._init_from_graph(OnnxPBParser(model_path, load_external_data=True).parse())
+    model = OnnxPBParser(model_path, load_external_data=True).parse()
+    self._init_from_graph(model["graph"])

  def _init_from_graph(self, graph: dict, is_subgraph: bool = False):
-    self.is_training = any(n.opset_id.domain in {Domain.AI_ONNX_TRAINING, Domain.AI_ONNX_PREVIEW_TRAINING} for n in graph["node"])
+    self.is_training = any(n['parsed_node'].opset_id.domain in {Domain.AI_ONNX_TRAINING, Domain.AI_ONNX_PREVIEW_TRAINING} for n in graph["node"])
    self.graph_name = graph["name"] if is_subgraph else ""
-    self.graph_values: dict[str, Any] = {"": None, **dict(graph["initializer"])}
-    self.graph_inputs = {name: typ for name, typ in graph["input"] if name not in self.graph_values}
-    self.graph_outputs = tuple(name for name, _ in graph["output"])
-    self.graph_nodes = tuple(graph["node"])
+    self.graph_values = {"": None, **{i["name"]: i["parsed_tensor"] for i in graph["initializer"]}}
+    self.graph_inputs = {i["name"]: i["parsed_type"] for i in graph["input"] if i["name"] not in self.graph_values}
+    self.graph_outputs = tuple(o["name"] for o in graph["output"])
+    self.graph_nodes = tuple(n["parsed_node"] for n in graph["node"])
    # track names from initializers and Constant nodes for fast path optimizations
    self.const_names: set[str] = set(self.graph_values.keys()) | {o for n in self.graph_nodes if n.op == "Constant" for o in n.outputs}

@@ -497,7 +511,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    o_ = [((i - 1) // s + 1) for i,s in zip(i_, s_)]
    return _onnx_pads_to_tiny_pads(_auto_pad([(o-1)*s+k-i for o,i,k,s in zip(o_, i_, k_, s_)], auto_pad))

-  def _clamp_cast(x:Tensor, dtype:DType): return x.clamp(dtype.min, dtype.max).cast(dtype)
+  def _clamp_cast(x:Tensor, dtype:DType): return x.clamp(dtypes.min(dtype), dtypes.max(dtype)).cast(dtype)

  def _prepare_quantize(x:Tensor, scale:Tensor, zero_point:Tensor|int, axis=1, block_size=0):
    if axis < 0: axis += x.ndim
@@ -1209,7 +1223,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT

  def DynamicQuantizeLinear(x: Tensor):
    # only support uint8
-    qmin, qmax = dtypes.uint8.min, dtypes.uint8.max
+    qmin, qmax = dtypes.min(dtypes.uint8), dtypes.max(dtypes.uint8)
    scale = (x.max().maximum(0) + ((-x).max()).maximum(0)) / (qmax - qmin)
    zero_point = _clamp_cast((qmin - x.min() / scale).round(), dtypes.uint8)
    y = _clamp_cast((x / scale).round() + zero_point, dtypes.uint8)
@@ -298,7 +298,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
  Converts ggml tensor data to a tinygrad tensor.

  Supported native types: float32 (id: 0), float16 (id: 1), int8 (id: 16), int16 (id: 17), int32 (id: 18)
-  Supported quantized types: Q4_0 (id: 2), Q4_1 (id: 3), Q8_0 (id: 8), Q4_K (id: 12), Q5_K (id: 13), Q6_K (id: 14), MXFP4 (id: 39)
+  Supported quantized types: Q4_0 (id: 2), Q4_1 (id: 3), Q8_0 (id: 8), Q4_K (id: 12), Q6_K (id: 14), MXFP4 (id: 39)
  """
  # https://github.com/ggerganov/ggml/blob/323951f1bdcdfbd5b5ff3a9a7c3770e63b1a560e/include/ggml.h#L356

@@ -312,23 +312,19 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
    return t.unsqueeze(-1).expand((*t.shape,8//b)).idiv(shift_tensor).bitwise_and(bitmask).transpose(-1, -2).flatten(-2)

  # map to (number of elements, number of bytes)
-  if (nelements_nbytes := { 2:(32,18), 3:(32,20), 8:(32,34), 12:(256,144), 13:(256,176), 14:(256,210), 39:(32,17) }.get(ggml_type)) is not None:
+  if (nelements_nbytes := { 2: (32, 18), 3: (32, 20), 8: (32, 34), 12: (256, 144), 14: (256, 210), 39: (32, 17) }.get(ggml_type)) is not None:
    blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1])).contiguous()
    if ggml_type == 2: return (q_to_uint8(blocks[:,2:], 4).bitcast(dtypes.int8) - 8) * blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32)
    if ggml_type == 3:
      d, m = (blocks[:,s:s+2].bitcast(dtypes.float16).cast(dtypes.float32) for s in [ 0, 2 ])
      return q_to_uint8(blocks[:,4:], 4).bitcast(dtypes.int8) * d + m
    if ggml_type == 8: return blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32) * blocks[:,2:].bitcast(dtypes.int8)
-     # Q4_K: 256 elements per 144-byte block (d:2, dmin:2, scales:12, qs:128)
-     # Q5_K: 256 elements per 176-byte block (d:2, dmin:2, scales:12, qh:32, qs:128)
-    if ggml_type in (12, 13):
+    if ggml_type == 12:  # Q4_K: 256 elements per 144-byte block (d:2, dmin:2, scales:12, qs:128)
      d, dmin = (blocks[:,i:i+2].bitcast(dtypes.float16).cast(dtypes.float32).unsqueeze(-1) for i in [0, 2])
      s = blocks[:,4:16]  # 12 bytes: 6-bit scales[0-3], 6-bit mins[0-3], high bits[4-7]
      sc = s[:,0:4].bitwise_and(63).cat(s[:,8:12].bitwise_and(0xF).bitwise_or(s[:,0:4].rshift(6).lshift(4)), dim=-1)
      mn = s[:,4:8].bitwise_and(63).cat(s[:,8:12].rshift(4).bitwise_or(s[:,4:8].rshift(6).lshift(4)), dim=-1)
-      qs_off = 48 if ggml_type == 13 else 16
-      q = Tensor.stack((qs:=blocks[:,qs_off:qs_off+128].reshape(-1,4,32)).bitwise_and(0xF), qs.rshift(4), dim=2).reshape(-1,8,32)
-      if ggml_type == 13: q = q + q_to_uint8(blocks[:,16:48], 1).reshape(-1, 8, 32) * 16
+      q = Tensor.stack((qs:=blocks[:,16:144].reshape(-1,4,32)).bitwise_and(0xF), qs.rshift(4), dim=2).reshape(-1,8,32).cast(dtypes.float32)
      return (d * sc.unsqueeze(-1) * q - dmin * mn.unsqueeze(-1)).flatten(-2)
    if ggml_type == 14:
      xl, xh = q_to_uint8(blocks[:,:128].reshape((-1, 2, 64)), 4), q_to_uint8(blocks[:,128:192].reshape((-1, 2, 32)), 2).lshift(4)
@@ -336,14 +332,15 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
      d = blocks[:,-2:].bitcast(dtypes.float16).cast(dtypes.float32).expand((-1, 256))
      return d * (xl.bitwise_or(xh).bitcast(dtypes.int8) - 32).flatten(-2) * scales
    if ggml_type == 39:
-      e = blocks[:, 0].cast(dtypes.uint32)
-      small_bits = Tensor([0x00200000, 0x00400000], dtype=dtypes.uint32, device=t.device)[e.clip(0, 1).cast(dtypes.int32)] # e = 0 or e = 1 case
-      d = (e < 2).where(small_bits, ((e - 1) * 0x00800000).cast(dtypes.uint32)).bitcast(dtypes.float32).unsqueeze(-1)
+      e_int = blocks[:, 0].cast(dtypes.int32)
+      d = ((e_int >= 2).cast(dtypes.float32) * (e_int.cast(dtypes.float32) - 128).exp2() +
+           (e_int == 1).cast(dtypes.float32) * 2.0**(-127) +
+           (e_int == 0).cast(dtypes.float32) * 2.0**(-128)).unsqueeze(-1)
      codes = q_to_uint8(blocks[:, 1:17], 4)
-      fp4_lut = Tensor([0.0, 1.0, 2.0, 3.0, 4.0, 6.0, 8.0, 12.0,
-                       -0.0,-1.0,-2.0,-3.0,-4.0,-6.0,-8.0,-12.0],
-                       dtype=dtypes.float32, device=t.device)
-      fp4_val = fp4_lut[codes]
+      sign = 1.0 - codes.rshift(3).cast(dtypes.float32) * 2.0
+      exp, mant = codes.rshift(1).bitwise_and(0x3).cast(dtypes.float32), codes.bitwise_and(0x1).cast(dtypes.float32)
+      fp4_val = sign * 2.0 * ((exp != 0).cast(dtypes.float32) * (1.0 + 0.5 * mant) * (exp - 1.0).exp2() +
+                              (exp == 0).cast(dtypes.float32) * 0.5 * mant)
      return (fp4_val * d).flatten(-2)[:n]
  raise ValueError(f"GGML type '{ggml_type}' is not supported!")

@@ -106,15 +106,13 @@ class InstOpRDNA4(Enum):
  SMEM = 0x1
  JUMP = 0x3
  JUMP_NO = 0x4
-  CALL = 0x5
+  JUMP_UNCOND = 0x5
  MESSAGE = 0x9
  VALU_TRANS = 0xb
  VALU_B2 = 0xd
  VALU_B4 = 0xe
  VINTERP = 0x12
  VMEM_RD_1 = 0x21
-  VMEM_RD_2 = 0x22
-  VMEM_WR_1 = 0x23
  VMEM_WR_2 = 0x24
  VMEM_WR_3 = 0x25
  VMEM_WR_4 = 0x26
@@ -133,8 +131,7 @@ class InstOpRDNA4(Enum):
  VALU_SCL_TRANS = 0x99
  SALU_2 = 0x9b
  SALU_5 = 0x9c
-  OTHER_VMEM = 0xbd
-  OTHER_VMEM_5 = 0xc1
+  OTHER_VMEM = 0xc1

 # ═══════════════════════════════════════════════════════════════════════════════
 # PACKET TYPE BASE CLASS
@@ -187,16 +184,19 @@ class TS_DELTA_SHORT(PacketType):
 class TS_DELTA_OR_MARK(PacketType):
  encoding = bits[6:0] == 0b0000001
  delta = bits[47:12]
-  pl = bits[8:8]
-  rt = bits[9:9]
+  bit8 = bits[8:8]
+  bit9 = bits[9:9]
  @property
-  def is_marker(self) -> bool: return bool(self.rt and not self.pl)
+  def is_marker(self) -> bool: return bool(self.bit9 and not self.bit8)

-class TS_DELTA_OR_MARK_RDNA4(TS_DELTA_OR_MARK):
+class TS_DELTA_OR_MARK_RDNA4(PacketType):  # Layout 4: 48->64 bits
+  encoding = bits[6:0] == 0b0000001
  delta = bits[63:12]
-  rt = bits[7:7]
-  pl = bits[8:8]
-  tl = bits[9:9]
+  bit7 = bits[7:7]
+  bit8 = bits[8:8]
+  bit9 = bits[9:9]
+  @property
+  def is_marker(self) -> bool: return bool((self.bit9 and not self.bit8) or self.bit7)

 class TS_DELTA_S5_W2(PacketType):
  encoding = bits[4:0] == 0b11100
@@ -402,17 +402,16 @@ class CDNA_PKT_2(PacketType):
  unk_padding = bits[63:8]

 class CDNA_WAVESTART(PacketType):
-  """type 3: 32-bit wave start (Wave/group_id)"""
+  """pkt_fmt=3: 32-bit WAVESTART packet (case 0x8)"""
  encoding = bits[3:0] == 3
-  sh = bits[5:5]
-  cu = bits[9:6]
-  wave = bits[13:10]
-  simd = bits[15:14]
-  pipe = bits[17:16]
-  me = bits[19:18]
-  _gap = bits[21:20]
-  count = bits[28:22]
-  _padding = bits[31:29]
+  unk_0 = bits[5:5]       # (data >> 5) & 1
+  unk_1 = bits[9:6]       # (data >> 6) & 0xf
+  wave = bits[13:10]      # (data >> 10) & 0xf
+  simd = bits[15:14]      # (data >> 0xe) & 3
+  cu = bits[17:16]        # (data >> 0x10) & 3
+  unk_5 = bits[19:18]     # (data >> 0x12) & 3
+  unk_6 = bits[28:22]     # (data >> 0x16) & 0x7f
+  unk_padding = bits[31:29]

 class CDNA_PKT_4(PacketType):
  """pkt_fmt=4: 16-bit packet (case 0xc, same as 0x8/0x14)"""
@@ -422,21 +421,21 @@ class CDNA_PKT_4(PacketType):
  unk_2 = bits[13:10]     # (data_word >> 10) & 0xf
  unk_3 = bits[15:14]     # (data_word >> 0xe)

-class REGCS_CDNA(PacketType):
-  """type 5: 48-bit register CS write (RegCs)"""
+class CDNA_PKT_5(PacketType):
+  """pkt_fmt=5: 48-bit packet (case 0x10)"""
  encoding = bits[3:0] == 5
-  pipe = bits[6:5]
-  _me_raw = bits[8:7]
-  regaddr = bits[15:9]
-  regdata = bits[47:16]
+  unk_0 = bits[6:5]       # (data >> 5) & 3
+  unk_1 = bits[7:7]       # (data >> 7) + 1 & 1
+  unk_2 = bits[15:9]      # (data >> 9) & 0x7f
+  unk_padding = bits[47:16]

 class CDNA_WAVEEND(PacketType):
-  """type 6: 16-bit wave end (group_id)"""
+  """pkt_fmt=6: 16-bit WAVEEND packet (case 0x14, same as 0x8/0xc)"""
  encoding = bits[3:0] == 6
-  sh = bits[5:5]
-  cu = bits[9:6]
-  wave = bits[13:10]
-  simd = bits[15:14]
+  unk_0 = bits[5:5]       # (data_word >> 5) & 1
+  unk_1 = bits[9:6]       # (data_word >> 6) & 0xf
+  wave = bits[13:10]      # (data_word >> 10) & 0xf
+  simd = bits[15:14]      # (data_word >> 0xe)

 class CDNA_EXEC(PacketType):
  """pkt_fmt=10: 16-bit EXEC packet (case 0x24)"""
@@ -509,7 +508,7 @@ class CDNA_PKT_15(PacketType):
  unk_padding = bits[47:16]

 PACKET_TYPES_CDNA: dict[int, type[PacketType]] = {
-  0: CDNA_DELTA, 1: CDNA_TIMESTAMP, 2: CDNA_PKT_2, 3: CDNA_WAVESTART, 4: CDNA_PKT_4, 5: REGCS_CDNA, 6: CDNA_WAVEEND,
+  0: CDNA_DELTA, 1: CDNA_TIMESTAMP, 2: CDNA_PKT_2, 3: CDNA_WAVESTART, 4: CDNA_PKT_4, 5: CDNA_PKT_5, 6: CDNA_WAVEEND,
  7: CDNA_PKT_7, 8: CDNA_PKT_8, 9: CDNA_PKT_9, 10: CDNA_EXEC, 11: CDNA_PKT_11, 12: CDNA_PKT_12,
  13: CDNA_INST, 14: CDNA_PKT_14, 15: CDNA_PKT_15,
 }
@@ -596,6 +595,7 @@ def map_insts(data:bytes, lib:bytes, target:str) -> Iterator[tuple[PacketType, I
  # map pcs to insts
  from tinygrad.viz.serve import amd_decode
  pc_map = amd_decode(lib, target)
+
  wave_pc:dict[int, int] = {}
  # only processing packets on one [CU, SIMD] unit
  def simd_select(p) -> bool: return getattr(p, "cu", 0) == 0 and getattr(p, "simd", 0) == 0
@@ -604,12 +604,14 @@ def map_insts(data:bytes, lib:bytes, target:str) -> Iterator[tuple[PacketType, I
    if isinstance(p, (WAVESTART, WAVESTART_RDNA4)):
      assert p.wave not in wave_pc, "only one inflight wave per unit"
      wave_pc[p.wave] = next(iter(pc_map))
-    elif isinstance(p, WAVEEND):
+      continue
+    if isinstance(p, WAVEEND):
      pc = wave_pc.pop(p.wave)
      yield (p, InstructionInfo(pc, p.wave, s_endpgm()))
+      continue
    # skip OTHER_ instructions, they don't belong to this unit
-    elif isinstance(p, (INST, INST_RDNA4)) and p.op.name.startswith("OTHER_"): pass
-    elif isinstance(p, IMMEDIATE_MASK):
+    if isinstance(p, (INST, INST_RDNA4)) and p.op.name.startswith("OTHER_"): continue
+    if isinstance(p, IMMEDIATE_MASK):
      # immediate mask may yield multiple times per packet
      for wave in range(16):
        if p.mask & (1 << wave):
@@ -618,24 +620,30 @@ def map_insts(data:bytes, lib:bytes, target:str) -> Iterator[tuple[PacketType, I
          assert type(inst).__name__ == "SOPP", f"IMMEDIATE_MASK packet must map to SOPP, got {inst}"
          wave_pc[wave] += inst.size()
          yield (p, InstructionInfo(pc, wave, inst))
-    elif isinstance(p, (VALUINST, INST, INST_RDNA4, IMMEDIATE)):
+      continue
+    if isinstance(p, (VALUINST, INST, INST_RDNA4, IMMEDIATE)):
      inst = pc_map[pc:=wave_pc[p.wave]]
-      # s_delay_alu and s_wait_alu instructions are skipped
+      # s_delay_alu doesn't get a packet?
      while (inst_op:=getattr(inst, 'op_name', '')) in {"S_DELAY_ALU", "S_WAIT_ALU"}:
        wave_pc[p.wave] += inst.size()
        inst = pc_map[pc:=wave_pc[p.wave]]
-      # assert branch always has a JUMP packet
-      if "BRANCH" in inst_op and not (isinstance(p, (INST, INST_RDNA4)) and p.op.name.startswith("JUMP")):
-        raise AssertionError(f"{inst_op} can only be followed by JUMP, got {p}")
+      # identify a branch instruction, only used for asserts
+      branch_inst = inst if "BRANCH" in inst_op else None
+      if branch_inst is not None:
+        assert isinstance(p, (INST, INST_RDNA4)) and p.op.name in {"JUMP_NO", "JUMP", "JUMP_UNCOND"}, f"branch can only be folowed by JUMP, got {p}"
      # JUMP handling
-      if isinstance(p, (INST, INST_RDNA4)) and p.op in {InstOp.JUMP, InstOpRDNA4.JUMP}:
-        x = getattr(inst, 'simm16') & 0xffff
-        wave_pc[p.wave] += inst.size() + (x - 0x10000 if x & 0x8000 else x)*4
+      if (isinstance(p, INST) and p.op is InstOp.JUMP) or (isinstance(p, INST_RDNA4) and p.op is InstOpRDNA4.JUMP):
+        simm16 = getattr(branch_inst, 'simm16')
+        assert branch_inst is not None and simm16 is not None, f"JUMP packet must map to a branch instruction, got {inst}"
+        x = simm16 & 0xffff
+        wave_pc[p.wave] += branch_inst.size() + (x - 0x10000 if x & 0x8000 else x)*4
      else:
+        if branch_inst is not None: assert inst_op != "S_BRANCH", f"S_BRANCH must have a JUMP packet, got {p}"
        wave_pc[p.wave] += inst.size()
      yield (p, InstructionInfo(pc, p.wave, inst))
+      continue
    # for all other packets (VMEMEXEC, ALUEXEC, etc.), yield with None
-    else: yield (p, None)
+    yield (p, None)

 # ═══════════════════════════════════════════════════════════════════════════════
 # PRINTER
@@ -566,11 +566,4 @@ class AMDHIPCCRenderer(AMDHIPRenderer):
    super().__init__(arch)
    self.compiler = HIPCCCompiler(arch)

-class QCOMCLRenderer(OpenCLRenderer):
-  device = "QCOM"
-
-  def __init__(self, chip_id):
-    from tinygrad.runtime.support.compiler_qcom import QCOMCompiler
-    self.chip_id, self.compiler = chip_id, QCOMCompiler(chip_id)
-
-  def __reduce__(self): return self.__class__, (self.chip_id,)
+class QCOMRenderer(OpenCLRenderer): device = "QCOM"
@@ -5,7 +5,7 @@ from tinygrad.helpers import dedup, getenv, merge_dicts, PROFILE
 from tinygrad.device import Buffer, ProfileGraphEntry, ProfileGraphEvent
 from tinygrad.engine.realize import ExecItem, CompiledRunner
 from tinygrad.engine.jit import GraphRunner, GraphException
-from tinygrad.runtime.ops_metal import wait_check, to_ns_str, MetalBuffer
+from tinygrad.runtime.ops_metal import wait_check, to_ns_str
 from tinygrad.runtime.autogen import metal
 from tinygrad.runtime.support import objc

@@ -13,6 +13,7 @@ class MetalGraph(GraphRunner):
  def __init__(self, jit_cache: list[ExecItem], input_buffers: list[Buffer], var_vals: dict[str, int],
               orig_valid_positions: dict[int, set[int]]|None = None):
    super().__init__(jit_cache, input_buffers, var_vals, orig_valid_positions)
+    if not all(isinstance(ji.prg, CompiledRunner) for ji in jit_cache): raise GraphException

    # create metal batch exec
    icb_descriptor = metal.MTLIndirectCommandBufferDescriptor.new()
@@ -108,9 +109,3 @@ class MetalGraph(GraphRunner):
    if PROFILE and self.command_buffer is not None:
      wait_check(self.command_buffer)
      self.collect_timestamps()
-
-  @staticmethod
-  def supports_exec_item(devs, ei:ExecItem) -> bool:
-    # Metal ICB replay encodes offsets as uint32; reject if any Metal buffer offset exceeds 32-bit range.
-    if any(b is not None and isinstance(b._buf, MetalBuffer) and b._buf.offset > 0xFFFFFFFF for b in ei.bufs): return False
-    return GraphRunner.supports_exec_item(devs, ei)
@@ -485,7 +485,6 @@ class AMDCopyQueue(HWQueue):
    if (dev:=signal.owner) is not None and signal.is_timeline and not dev.is_am():
      self.q(self.sdma.SDMA_OP_FENCE | fence_flags, *data64_le(dev.queue_event_mailbox_ptr), dev.queue_event.event_id)
      self.q(self.sdma.SDMA_OP_TRAP, self.sdma.SDMA_PKT_TRAP_INT_CONTEXT_INT_CONTEXT(dev.queue_event.event_id))
-    elif dev is not None and dev.is_am(): self.q(self.sdma.SDMA_OP_TRAP, 0)

    return self

@@ -867,12 +866,10 @@ class PCIIface(PCIIfaceBase):
    return AMDQueueDesc(ring=ring.cpu_view().view(fmt='I'), doorbell=self.dev_impl.doorbell64.view(doorbell_index * 8, 8, fmt='Q'), put_value=0,
      read_ptr=gart.cpu_view().view(offset=rptr, size=8, fmt='Q'), write_ptr=gart.cpu_view().view(offset=wptr, size=8, fmt='Q'), params=rcvr_params)

-  def _collect_interrupts(self, reset=False, drain_only=False):
+  def _collect_faults(self, reset=False):
    devs:list[AMDDevice] = [d for pg in HCQCompiled.peer_groups.values() for d in pg if isinstance(d, AMDDevice) and d.is_am()]
    for d in devs:
-      if drain_only: d.iface.dev_impl.ih.drain()
-      else: d.iface.dev_impl.ih.interrupt_handler()
-
+      d.iface.dev_impl.ih.interrupt_handler()
      if reset and d.iface.dev_impl.recover(force=d.error_state is not None):
        d.compute_queue.put_value = d.compute_queue.read_ptr[0] = d.compute_queue.write_ptr[0] = 0
        d.iface.dev_impl.gfx.setup_ring(*d.compute_queue.params)
@@ -882,11 +879,11 @@ class PCIIface(PCIIfaceBase):
  def sleep(self, timeout):
    if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
      self.pci_dev.irq_fd.read(8 * events_cnt)
-    self._collect_interrupts()
+    self._collect_faults()
    if self.dev_impl.is_err_state: raise RuntimeError("Device is in error state")

  def on_device_hang(self):
-    self._collect_interrupts(reset=True)
+    self._collect_faults(reset=True)
    raise RuntimeError("Device hang detected")

  def device_fini(self): self.dev_impl.fini()
@@ -1079,10 +1076,6 @@ class AMDDevice(HCQCompiled):
    self.hw_compute_queue_t().memory_barrier().signal(self.timeline_signal, self.next_timeline()).submit(self)
    self.synchronize()

-  def synchronize(self, timeout:int|None=None):
-    super().synchronize(timeout)
-    if self.is_am() and self.error_state is None: self.iface._collect_interrupts(reset=False, drain_only=True)
-
  def on_device_hang(self): self.iface.on_device_hang()

  def device_props(self): return self.iface.props
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	4662fb413f	works	2026-03-04 13:14:16 +08:00
geohot	ccb5dcf3b8	fix	2026-03-04 13:01:35 +08:00
geohot	6b82b51759	close	2026-03-04 12:59:17 +08:00