early lower cat

simpler
min
2026-03-07 11:39:20 +08:00 · 2026-03-07 10:11:21 +08:00 · 2026-03-07 10:00:29 +08:00 · 2026-03-06 18:51:25 +08:00 · 2026-03-06 18:46:28 +08:00 · 2026-03-06 00:54:39 +03:00
135 changed files with 2829 additions and 11249 deletions
@@ -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 JITBEAM=1 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
+      run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 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
@@ -617,6 +617,27 @@ jobs:
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

+  testcommausbgpubenchmark:
+    name: UsbGPU Benchmark (comma)
+    runs-on: [self-hosted, Linux, comma4]
+    timeout-minutes: 20
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: openpilot compile3 0.10.1 driving_vision
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." DEV=AMD AMD_LLVM=1 AMD_IFACE=USB ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot load_pickle 0.10.1 driving_vision
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." DEV=AMD AMD_IFACE=USB ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
+
  testreddriverbenchmark:
    name: AM Benchmark
    runs-on: [self-hosted, Linux, tinyboxrandom]
@@ -244,6 +244,37 @@ jobs:
    - name: Run TYPED=1
      run: CHECK_OOB=0 DEV=CPU TYPED=1 python test/test_tiny.py

+  nulltest:
+    name: Null Tests
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: unittest-13
+        pydeps: "pillow ftfy regex pre-commit"
+        deps: testing_unit
+        llvm: 'true'
+        amd: 'true'
+    - name: Run NULL backend tests
+      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
+    - name: Run targetted tests on NULL backend
+      run: NULL=1 python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
+    # TODO: too slow
+    # - name: Run SDXL on NULL backend
+    #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
+    - name: Run Clip tests for SD MLPerf on NULL backend
+      run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
+    - name: Run AMD emulated BERT training on NULL backend
+      run: EMULATE=AMD_RDNA4 NULL=1 NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+    # TODO: support fake weights
+    #- name: Run LLaMA 7B on 4 fake devices
+    #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
+
  unittest:
    name: Unit Tests
    runs-on: ubuntu-latest
@@ -268,20 +299,6 @@ jobs:
      run: |
        CPU=1 python test/null/test_device.py TestRunAsModule.test_module_runs
        CPU=1 python -m pytest -n=auto test/unit/ --durations=20
-    - name: Run NULL backend tests
-      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
-    - name: Run targetted tests on NULL backend
-      run: NULL=1 python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
-    # TODO: too slow
-    # - name: Run SDXL on NULL backend
-    #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
-    - name: Run Clip tests for SD MLPerf on NULL backend
-      run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
-    - name: Run AMD emulated BERT training on NULL backend
-      run: EMULATE=AMD_RDNA4 NULL=1 NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
-    # TODO: support fake weights
-    #- name: Run LLaMA 7B on 4 fake devices
-    #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
    - name: Run GC tests
      run: python test/external/external_uop_gc.py
    - name: External Benchmark Schedule
@@ -644,7 +661,7 @@ jobs:
          sudo apt-get update
          sudo apt-get install llvm-21 llvm-21-tools cloc
      - name: Install rocprof-trace-decoder
-        run: sudo PYTHONPATH="." ./extra/sqtt/install_sqtt_decoder.py
+        run: sudo PYTHONPATH="." ./extra/sqtt/install_rocprof_decoder.py
      - name: Run AMD renderer tests
        run: AMD_LLVM=0 python -m pytest -n=auto test/amd/ --durations 20
      - name: Run AMD renderer tests (AMD_LLVM=1)
@@ -3,7 +3,7 @@ from pathlib import Path
 import multiprocessing

 from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
-from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling, profile_marker
+from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling, profile_marker, DEBUG
 from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict, safe_load, safe_save
 from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW

@@ -1336,11 +1336,13 @@ def train_llama3():
  # vocab_size from the mixtral tokenizer
  if not SMALL: model_params |= {"vocab_size": 32000}
  real_vocab_size = model_params['vocab_size']
-  if (MP := getenv("MP", 1)) > 1: model_params['vocab_size'] = round_up(model_params['vocab_size'], 256 * MP)
-  vocab_mask:Tensor = Tensor.arange(model_params['vocab_size']).reshape(1, 1, -1) >= real_vocab_size
  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
  print(f"model parameters: {model_params}")

+  # pad vocab
+  if (MP := getenv("MP", 1)) > 1: model_params['vocab_size'] = round_up(model_params['vocab_size'], 256 * MP)
+  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
@@ -1385,8 +1387,10 @@ def train_llama3():

  # init grads
  for p in optim.params:
-    p.grad = p.zeros_like().contiguous().realize()
+    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)

@@ -1401,15 +1405,15 @@ def train_llama3():

  @TinyJit
  def minibatch(tokens:Tensor):
-    tokens = tokens.to(None)
    if (DP := getenv("DP", 1)) > 1:
      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-      tokens = tokens.shard(device, 0)
+      tokens = tokens.to(None).shard(device, 0)
    if (MP := getenv("MP", 1)) > 1:
      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
      tokens = tokens.shard(device)
+    if DP == 1 and MP == 1: tokens = tokens.to(None)
    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = vocab_mask.where(-float("inf"), logits).sparse_categorical_crossentropy(tokens[:, 1:])
+    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))
    Tensor.realize(loss, *grads)
@@ -1417,35 +1421,37 @@ def train_llama3():

  @TinyJit
  def optim_step():
-    optim.step()
+    grad_norm = optim.fstep(grads)
    scheduler.step()

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

    lr = optim.lr
    Tensor.realize(lr, *grads)

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

  @TinyJit
  @Tensor.train(False)
  def eval_step(tokens:Tensor):
-    tokens = tokens.to(None)
    if (DP := getenv("DP", 1)) > 1:
      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-      tokens = tokens.shard(device, 0)
+      tokens = tokens.to(None).shard(device, 0)
    if (MP := getenv("MP", 1)) > 1:
      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
      tokens = tokens.shard(device)
+    if DP == 1 and MP == 1: tokens = tokens.to(None)
    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = vocab_mask.where(-float("inf"), logits).sparse_categorical_crossentropy(tokens[:, 1:])
+    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
    return loss.flatten().float().to("CPU")

  # ** data iters **
  def fake_data(bs, samples):
+    import numpy as np
    for _ in range(samples // bs):
-      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=model_params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
+      fake_data_np = np.random.randint(0, model_params["vocab_size"], size=(bs, SEQLEN + 1), dtype=np.int32)
+      yield Tensor(fake_data_np, device="NPY")

  def get_train_iter():
    if getenv("FAKEDATA", 0):
@@ -1472,13 +1478,14 @@ 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):
+      for _ in range(grad_acc if i >= 2 else 1):
        ist = time.perf_counter()
        try: tokens = next(train_iter)
        except StopIteration:
@@ -1491,7 +1498,8 @@ def train_llama3():
      if stopped: break

      gt = time.perf_counter()
-      lr = optim_step().item()
+      ret = optim_step()
+      lr, grad_norm = ret[0].item(), ret[1].item()
      et = time.perf_counter()

      loss = sum(losses) / len(losses)
@@ -1502,18 +1510,21 @@ def train_llama3():
      if BENCHMARK: step_times.append(step_time)

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

      mem_gb = GlobalCounters.mem_used / 1e9
      gflops = GlobalCounters.global_ops / 1e9 / dev_time
      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * max(getenv("DP", 1), getenv("MP", 1)) * 2.3e15)) * 100
      tqdm.write(
          f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
-          f"{lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
+          f"{lr:.12f} LR, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
+      if DEBUG >= 1: tqdm.write("  mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))

      if WANDB:
        wandb.log({
-          "lr": lr, "train/loss": loss,
+          "train/loss": loss,
+          "train/lr": lr,
+          "train/grad_norm": grad_norm,
          "train/step_time": step_time,
          "train/gbs_time": gbs_time,
          "train/optim_time": optim_time,
@@ -1542,7 +1553,7 @@ def train_llama3():
        print(f"epoch global_ops: {GlobalCounters.global_ops:_}, "
              f"epoch global_mem: {GlobalCounters.global_mem:_}")

-    if (sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1)) or (BENCHMARK and i == BENCHMARK):
+    if (sequences_seen // EVAL_FREQ != (sequences_seen - actual_gbs) // EVAL_FREQ 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}")
@@ -1550,7 +1561,7 @@ def train_llama3():
      # run eval
      eval_losses = []
      eval_iter = get_eval_iter()
-      tqdm.write(f"evaluating {5760//EVAL_BS} batches of {EVAL_BS} sequences")
+      tqdm.write(f"evaluating {EVAL_SAMPLES//EVAL_BS} batches of {EVAL_BS} sequences")

      for j,tokens in tqdm(enumerate(eval_iter), total=EVAL_SAMPLES//EVAL_BS):
        eval_losses += eval_step(tokens).tolist()
@@ -7,41 +7,51 @@ class GradAccClipAdamW(Optimizer):
  def __init__(self, params:list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, grad_acc=1, clip_norm=1.0, device=None, fused=FUSE_OPTIM):
    super().__init__(params, lr, device, fused)
    self.b1, self.b2, self.eps, self.wd = b1, b2, eps, weight_decay
-    self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False).contiguous() for _ in [b1, b2])
+    self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False) for _ in [b1, b2])
    self.m = self._new_optim_param()
    self.v = self._new_optim_param()
    self.grad_acc, self.clip_norm = grad_acc, clip_norm

+  def fstep(self, grads:list[Tensor]):
+    if self.fused:
+      out, extra = self._step([], grads)
+      updates = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
+    else:
+      updates, extra = self._step([], grads)
+    for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i]))
+    to_realize = extra+self.params+self.buffers
+
+    Tensor.realize(*to_realize)
+    return extra[-1]
+
  def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
    for i in range(len(grads)):
-      if grads[i].device != self.m[i].device: grads[i] = grads[i].to(self.m[i].device)
+      if grads[i].device != self.m[i].device: grads[i].assign(grads[i].to(self.m[i].device))

    if self.fused:
-      grads[0] = grads[0] / self.grad_acc
+      grads[0].assign(grads[0] / self.grad_acc)
      total_norm = grads[0].float().square().sum().sqrt()
-      grads[0] = (grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype)
+      grads[0].assign((grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype))
    else:
      for i in range(len(grads)):
-        grads[i] = grads[i] / self.grad_acc
-      total_norm = Tensor.zeros((), dtype=dtypes.float32, device=self.device)
-      for g in grads:
-        total_norm += g.float().square().sum()
-      total_norm = total_norm.sqrt()
+        grads[i].assign(grads[i] / self.grad_acc).realize()
+      total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous().realize()
      for i in range(len(grads)):
-        grads[i] = (grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype)
+        grads[i].assign((grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype)).realize()

    ret = []
    self.b1_t *= self.b1
    self.b2_t *= self.b2
-    for i, (t, g) in enumerate(zip(params, grads)):
+    for i, g in enumerate(grads):
      self.m[i].assign((self.b1 * self.m[i] + (1.0 - self.b1) * g).cast(self.m[i].dtype))
      self.v[i].assign((self.b2 * self.v[i] + (1.0 - self.b2) * (g * g)).cast(self.v[i].dtype))
      m_hat = self.m[i] / (1.0 - self.b1_t)
      v_hat = self.v[i] / (1.0 - self.b2_t)
      up = m_hat / (v_hat.sqrt() + self.eps)
-      ret.append((self.lr * up).cast(t.dtype))
-    return ret, [self.b1_t, self.b2_t] + self.m + self.v
+      ret.append((self.lr * up).cast(g.dtype))
+    return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]

  def _apply_update(self, t:Tensor, up:Tensor) -> Tensor:
-    up = up.shard_like(t) + self.lr.to(t.device) * self.wd * t.detach()
+    wd = self.wd if t.ndim >= 2 else 0.0
+    up = up.shard_like(t) + self.lr.to(t.device) * wd * t.detach()
    return t.detach() - up.cast(t.dtype)
@@ -0,0 +1,36 @@
+#!/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 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
@@ -0,0 +1,31 @@
+#!/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 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,6 +5,7 @@ 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}
@@ -14,7 +15,7 @@ export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-0}

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

 export MODEL="llama3"
@@ -22,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="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
+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}
@@ -0,0 +1,43 @@
+#!/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} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export DP=${DP:-8} MP=${MP:-1} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-4}
 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="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
+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}
@@ -0,0 +1,38 @@
+#!/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
@@ -3,4 +3,4 @@ export BENCHMARK=5
 export EVAL_BS=0
 export VIZ=${VIZ:--1}
 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
-PYTHONPATH="." extra/viz/cli.py --profile --device "AMD" --top 20
+extra/viz/cli.py --profile --device "AMD" --top 20
@@ -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)):
+    with Context(DEBUG=max(DEBUG.value, 2 if i == 2 else 1), OPENPILOT_HACKS=1):
      ret = run_onnx_jit(**inputs).numpy()
    # copy i == 1 so use of JITBEAM is okay
    if i == 1: test_val = np.copy(ret)
@@ -0,0 +1,16 @@
+import sys, pickle
+from extra.bench_log import WallTimeEvent, BenchEvent
+from tinygrad.helpers import getenv
+
+PKL = sys.argv[1] if len(sys.argv) > 1 else "/tmp/openpilot.pkl"
+
+load_times = []
+
+for _ in range(10):
+  with WallTimeEvent(BenchEvent.STEP) as wte: pickle.load(open(PKL, 'rb'))
+  load_times.append(wte.time)
+  print(f"pickle load: {wte.time:6.2f} s")
+
+if (assert_time:=getenv("ASSERT_MIN_LOAD_TIME")):
+  min_time = min(load_times)
+  assert min_time < assert_time, f"Speed regression, expected min load time of < {assert_time} s but took: {min_time} s"
@@ -34,7 +34,8 @@ class WallTimeEvent:
    self.start = time.monotonic()
    return self
  def __exit__(self, *_):
-    _events[self.event]["wall"].append(time.monotonic() - self.start)
+    self.time = time.monotonic() - self.start
+    _events[self.event]["wall"].append(self.time)
    return False

 class KernelTimeEvent:
@@ -3,7 +3,7 @@ from tinygrad import Tensor, Device, dtypes
 from tinygrad.dtype import AddrSpace
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
 from tinygrad.renderer import Estimates
-from tinygrad.helpers import getenv, all_same, dedup
+from tinygrad.helpers import getenv, all_same, DEBUG
 from extra.gemm.asm.cdna.asm import build_kernel, TILE_M, TILE_N, TILE_K, NUM_WG

 # ** CDNA4 assembly gemm
@@ -26,17 +26,25 @@ def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:

 counters = {"used":0, "todos":[]}
 def todo(msg:str) -> bool: counters["todos"].append(msg); return False
-atexit.register(lambda: print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used'))
+def _asm_gemm_report():
+  print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used')
+  if DEBUG >= 2 and counters["todos"]:
+    from collections import Counter
+    for msg, cnt in Counter(counters["todos"]).most_common(): print(f'  {cnt:3d}x {msg}')
+atexit.register(_asm_gemm_report)

 def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
  if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}")
  if a.dtype not in {dtypes.bfloat16, dtypes.float16}: return todo(f"only bfloat16/float16, got {a.dtype}")
  batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
  N = b.shape[1]
-  # only sharding on the batch or K is tested, others might work too
  if isinstance(a.device, tuple):
-    if a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
+    if a.ndim == 2 and a.uop.axis == 0 and b.uop.axis is None: M //= len(a.device)
+    elif a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
+    elif a.ndim == 2 and a.uop.axis is None and b.uop.axis == 1: N //= len(a.device)
    elif a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None: batch //= len(a.device)
+    elif a.ndim == 3 and a.uop.axis is None and b.uop.axis == 1: N //= len(a.device)
+    elif a.ndim == 3 and a.uop.axis == 2 and b.uop.axis == 0: K //= len(a.device)
    else: return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
    dname = a.device[0]
  else: dname = a.device
@@ -78,6 +86,10 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp):
 def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
  assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
  counters["used"] += 1
+  unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0
+  if unfold_batch:
+    orig_batch = a.shape[0]
+    a = a.reshape(a.shape[0]*a.shape[1], a.shape[2])
  squeeze = a.ndim == 2
  if squeeze: a = a.unsqueeze(0)

@@ -85,9 +97,16 @@ def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
  N = b.shape[1]
  is_multi = isinstance(a.device, tuple)
  if (k_sharded:=is_multi and a.uop.axis == 2): K //= len(a.device)
+  if (m_sharded:=is_multi and a.uop.axis == 1): M //= len(a.device)
+  n_sharded = is_multi and b.uop.axis == 1

  if is_multi:
-    out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
+    if n_sharded:
+      out = Tensor(Tensor.empty(batch, M, N//len(a.device), dtype=a.dtype, device=a.device).uop.multi(2), device=a.device)
+    elif m_sharded:
+      out = Tensor(Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device).uop.multi(1), device=a.device)
+    else:
+      out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
  else:
    out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)

@@ -98,4 +117,6 @@ def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
  else:
    out = Tensor.custom_kernel(out, a, b, fxn=custom_uop_gemm, grad_fxn=custom_gemm_bw)[0]
  if k_sharded: out = out.sum(0)
-  return out.squeeze(0) if squeeze else out
+  out = out.squeeze(0) if squeeze else out
+  if unfold_batch: out = out.reshape(orig_batch, -1, out.shape[-1])
+  return out
@@ -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)
+    x = hevc_tensor[offset:offset+sz*HEVC_ROUNDUP].decode_hevc_frame(pos, out_image_size, opaque[i], hist).realize()
    if outbuf is not None: outbuf.assign(x).realize()
-    return x.realize()
+    return x
  return TinyJit(hevc_decode_frame)

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

  # decode all frames using the iterator
-  with Timing("decoding whole file: ", on_exit=(lambda et: f", {len(frame_info)} frames, {len(frame_info)/(et/1e9):.2f} fps")):
+  tm = Timing("decoding whole file: ", on_exit=(lambda et: f", {len(frame_info)} frames, {len(frame_info)/(et/1e9):.2f} fps"))
+  with tm:
    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
@@ -2,30 +2,13 @@

 ## Getting SQ Thread Trace

-SQTT is implemented on top of normal tinygrad profiling, `VIZ=1 SQTT=1` to get profile pickle with sqtt data embedded in it.
+`VIZ=2` to enable SQTT profiling.
+
+`SQTT_ITRACE_SE_MASK=X` to select shader engines for instruction tracing, -1 = all, 0 = disabled, >0 = SE bitmask, default 0b11.

 `SQTT_BUFFER_SIZE=X` to change size of SQTT buffer (per shader engine, 6 SEs on 7900xtx) in megabytes, default 256.

-`SQTT_ITRACE_SE_MASK=X` to select for which shader engines instruction tracing will be enabled, -1 is all, 0 is none (instruction tracing disabled), >0 is
-bitfield/mask for SEs to enable instruction tracing on. Masking shader engines will give smaller file sizes at a cost of less hits and kernels that
-don't have any wavefront on first simd of shader engine with instruction tracing enabled will not have instruction timings.
-The default is 2 (second shader engine only), only one for file size reasons, second instead of first because dispatch starts from it so there is
-greater chance that kernels with small global size will have instruction tracing data.
- 
-Note that instruction tracing might not be available for kernels with small global dims, this is not a bug, but it can be improved with various hacks
-to the point where it can reliably trace a kernel consisting of a single wavefront (am only, not quite reliable under amdgpu due to waves sometimes
-being dispatched starting from different simds). More info in comments in ops_amd.py
+## Viewing the traces

-## Converting pickled profile with SQTT data into RGP file
-
-```bash
-extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
-```
-
-Then load gpu0.rgp into Radeon GPU Profiler. It works just fine both in wine (macos, native version available for linux) and via ssh X forwarding
-
-If multiple gpus are used you can select which one to export with `-d` like this:
-
-```bash
-extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -d 'AMD:5' -o /tmp/gpu5.rgp
-```
+- Web UI: `tinygrad/viz/serve.py`
+- Command line: `python -m tinygrad.renderer.amd.sqtt`
@@ -1,152 +0,0 @@
-import os
-os.environ["PYTHONPATH"] = "."
-os.environ["SQTT"] = "1"
-if "DEV" not in os.environ: os.environ["DEV"] = "AMD"
-os.environ["PROFILE"] = "1"
-os.environ["AMD_LLVM"] = "0"
-
-from dataclasses import replace
-import atexit, contextlib
-from tinygrad import Tensor
-from tinygrad.helpers import system, OSX
-from tinygrad.runtime.ops_amd import AMDProgram
-from extra.sqtt.roc import decode, WaveExec, ProfileSQTTEvent
-from tinygrad.device import Device
-
-from extra.sqtt.attempt_sqtt_parse import parse_sqtt_print_packets
-
-dev = Device["AMD"]
-
-@contextlib.contextmanager
-def save_sqtt():
-  # clear the old traces
-  dev.profile_events.clear()
-  sqtt:dict[str, list[WaveExec]] = {}
-  yield sqtt
-  events = dev.profile_events
-
-  #rctx = decode(events)
-  #assert len(rctx.inst_execs) > 0, "empty sqtt output"
-  #sqtt.update(rctx.inst_execs)
-
-  for e in events:
-    if isinstance(e, ProfileSQTTEvent):
-      print(replace(e, blob=b''))
-      if e.se == 0:
-        parse_sqtt_print_packets(e.blob)
-
-template = """.text
-.globl matmul
-.p2align 8
-.type matmul,@function
-matmul:
-  INSTRUCTION
-
-.rodata
-.p2align 6
-.amdhsa_kernel matmul
-  .amdhsa_kernarg_size 8
-  .amdhsa_user_sgpr_kernarg_segment_ptr 1
-  .amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
-  .amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
-  .amdhsa_wavefront_size32 1
-.end_amdhsa_kernel
-
-.amdgpu_metadata
---
-amdhsa.version:
-  - 1
-  - 0
-amdhsa.kernels:
-  - .name: matmul
-    .symbol: matmul.kd
-    .group_segment_fixed_size: 0
-    .private_segment_fixed_size: 0
-    .wavefront_size: 32
-    .sgpr_count: 8
-    .vgpr_count: 8
-    .max_flat_workgroup_size: 1024
-    .kernarg_segment_align: 8
-    .kernarg_segment_size: 8
-    .args:
-      - .address_space:  global
-        .name:           a
-        .offset:         0
-        .size:           8
-        .type_name:      'float*'
-        .value_kind:     global_buffer
-...
-.end_amdgpu_metadata
-"""
-
-def run_asm(src, num_workgroups=1, num_waves=1):
-  WAVE_SIZE = 32
-  t = Tensor.empty(0x1000).realize()
-  buf = t.uop.buffer.ensure_allocated()
-  lib = dev.compiler.compile(template.replace("INSTRUCTION", '\n'.join(src)))
-  dev.compiler.disassemble(lib)
-  fxn = AMDProgram(dev, "matmul", lib)
-  fxn(buf._buf, global_size=(num_workgroups,1,1), local_size=(WAVE_SIZE*num_waves,1,1), wait=True)
-
-if __name__ == "__main__":
-  with save_sqtt() as sqtt:
-    run_asm([
-      "s_nop 100",
-      "s_nop 100",
-      "s_load_b64 s[0:1], s[0:1], null",
-      "s_waitcnt lgkmcnt(0)",
-      "s_nop 100",
-      "s_nop 100",
-      "s_add_i32 s2, s2, 10",
-      "s_add_i32 s2, s2, 10",
-      "s_nop 100",
-      "s_nop 100",
-      "v_mov_b32_e32 v0, 0",
-      "v_mov_b32_e32 v0, 0",
-      "s_nop 100",
-      "s_nop 100",
-      "v_dual_fmac_f32 v2, v48, v24 :: v_dual_fmac_f32 v9, v37, v51",
-      "v_dual_fmac_f32 v2, v48, v24 :: v_dual_fmac_f32 v9, v37, v51",
-      "s_nop 100",
-      "s_nop 100",
-      "global_load_b128 v[2:5], v0, s[0:1]",
-      "global_load_b128 v[2:5], v0, s[0:1]",
-      "s_nop 100",
-      "s_nop 100",
-      "s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)",
-      "s_endpgm",
-    ], num_workgroups=1, num_waves=1)
-  exit(0)
-
-  with save_sqtt() as sqtt:
-    #(Tensor.empty(16,16) @ Tensor.empty(16,16)).elu().realize()
-    #Tensor.empty(1, 64).sum(axis=1).realize()
-    Tensor.empty(1).log2().realize()
-  exit(0)
-
-  with save_sqtt() as sqtt:
-    # what's in v0?
-    run_asm([
-      "v_mov_b32_e32 v0, 0",
-      "v_mov_b32_e32 v1, 0",
-      "s_clause 0x1",
-      "s_load_b64 s[0:1], s[0:1], null",
-      "s_waitcnt lgkmcnt(0)",
-    ]+[
-      "global_load_b32 v1, v0, s[0:1]",
-    ]*10+[
-      "global_load_b32 v10, v1, s[0:1]",
-      "s_waitcnt vmcnt(0)",
-
-      #"v_rcp_f32 v1, v0"
-      #"v_add_f32_e32 v1 v0 v0",
-      #"v_add_f32_e32 v5 v4 v4",
-      #"v_add_f32_e32 v7 v6 v6",
-      #"v_add_f32_e32 v1 v0 v0",
-      #"v_add_f32_e32 v2 v1 v1",
-      #"s_nop 1"
-    ]*5+[
-      "v_add_f32_e32 v3 v2 v2",
-    ]*5+[
-      "v_mul_f32_e32 v3 v2 v2",
-    ]*7)
@@ -1,548 +0,0 @@
-import pickle, sys
-from tinygrad.helpers import getenv, Timing, colored
-from extra.sqtt.roc import decode, ProfileSQTTEvent
-
-# do these enums match fields in the packets?
-#from tinygrad.runtime.support.amd import import_soc
-#soc = import_soc([11])
-#perf_sel = {getattr(soc, k):k for k in dir(soc) if k.startswith("SQ_PERF_")}
-
-# Instruction packets (one per ISA op)
-# NOTE: these are bad guesses and may be wrong! feel free to update if you know better
-# some names were taken from SQ_TT_TOKEN_MASK_TOKEN_EXCLUDE_SHIFT
-
-# we see 18 opcodes
-# opcodes(18):  1  2  3  4  5  6  8  9  F 10 11 12 14 15 16 17 18 19
-# if you exclude everything, you are left with 6
-# opcodes( 6): 10 11 14 15 16 17
-# sometimes we see a lot of B, but not repeatable
-
-# not seen
-# 7 A C
-
-# NOTE: INST runs before EXEC
-
-OPCODE_COLORS = {
-  # dispatches are BLACK
-  0x1: "BLACK",
-  0x18: "BLACK",
-
-  # execs are yellow
-  0x2: "yellow",
-  0x3: "yellow",
-  0x4: "YELLOW",
-  0x5: "YELLOW",
-
-  # waves are blue
-  0x8: "blue",
-  0x9: "blue",
-  0x6: "cyan",
-  0xb: "cyan",
-}
-
-OPCODE_NAMES = {
-  # gated by SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT (but others must be enabled for it to show)
-  0x01: "VALUINST",
-  # gated by SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT
-  0x02: "VMEMEXEC",
-  # gated by SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT
-  0x03: "ALUEXEC",
-  # gated by SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT
-  0x04: "IMMEDIATE",
-  0x05: "IMMEDIATE_MASK",
-
-  # gated by SQ_TT_TOKEN_EXCLUDE_WAVERDY_SHIFT
-  0x06: "WAVERDY",
-  # gated by SQ_TT_TOKEN_EXCLUDE_WAVESTARTEND_SHIFT
-  0x08: "WAVEEND",
-  0x09: "WAVESTART",
-  # gated by SQ_TT_TOKEN_EXCLUDE_WAVEALLOC_SHIFT
-  0x0B: "WAVEALLOC",  # FFF00
-
-  # gated by NOT SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT
-  0x0D: "PERF",
-  # gated by SQ_TT_TOKEN_EXCLUDE_EVENT_SHIFT
-  0x12: "EVENT",
-  0x13: "EVENT_BIG",  # FFFFF800
-  # some gated by SQ_TT_TOKEN_EXCLUDE_REG_SHIFT, some always there. something is broken with the timing on this
-  0x14: "REG",
-  # gated by SQ_TT_TOKEN_EXCLUDE_INST_SHIFT
-  0x18: "INST",
-  # gated by SQ_TT_TOKEN_EXCLUDE_UTILCTR_SHIFT
-  0x19: "UTILCTR",
-
-  # this is the first (8 byte) packet in the bitstream
-  0x17: "LAYOUT_HEADER",       # layout/mode/group + selectors A/B (reversed)
-
-  # pure time (no extra bits)
-  0x0F: "TS_DELTA_SHORT",
-  0x10: "NOP",
-  0x11: "TS_WAVE_STATE",     # almost pure time, has a small flag
-
-  # not a good name, but seen and understood mostly
-  0x15: "SNAPSHOT",          # small delta + 50-ish bits of snapshot
-  0x16: "TS_DELTA_OR_MARK",  # 36-bit long delta or 36-bit marker
-
-  # packets we haven't seen / rarely see 0x0b
-  0x07: "TS_DELTA_S8_W3_7",         # shift=8,  width=3  (small delta)
-  0x0A: "TS_DELTA_S5_W2_A",         # shift=5,  width=2
-  0x0C: "TS_DELTA_S5_W3_B",         # shift=5,  width=3 (different consumer)
-}
-
-#  SALU    =  0x0 / s_mov_b32
-#  SMEM    =  0x1 / s_load_b*
-#  JUMP    =  0x3 / s_cbranch_scc0
-#  NEXT    =  0x4 / s_cbranch_execz
-#  MESSAGE =  0x9 / s_sendmsg
-#  VALU    =  0xb / v_(exp,log)_f32_e32
-#  VALU    =  0xd / v_lshlrev_b64
-#  VALU    =  0xe / v_mad_u64_u32
-#  VMEM    = 0x21 / global_load_b32
-#  VMEM    = 0x22 / global_load_b32
-#  VMEM    = 0x24 / global_store_b32
-#  VMEM    = 0x25 / global_store_b64
-#  VMEM    = 0x27 / global_store
-#  VMEM    = 0x28 / global_store_b64
-#  LDS     = 0x29 / ds_load_b128
-#  LDS     = 0x2b / ds_store_b32
-#  LDS     = 0x2e / ds_store_b128
-#  ????    = 0x5a / hidden global_load  instruction
-#  ????    = 0x5b / hidden global_load  instruction
-#  ????    = 0x5c / hidden global_store instruction
-#  VALU    = 0x73 / v_cmpx_eq_u32_e32 (not normal VALUINST)
-OPNAME = {
-  0x0: "SALU",
-  0x1: "SMEM",
-  0x3: "JUMP",
-  0x4: "NEXT",
-  0x9: "MESSAGE",
-  0xb: "VALU",
-  0xd: "VALU",
-  0xe: "VALU",
-  0x21: "VMEM_LOAD",
-  0x22: "VMEM_LOAD",
-  0x24: "VMEM_STORE",
-  0x25: "VMEM_STORE",
-  0x26: "VMEM_STORE",
-  0x27: "VMEM_STORE",
-  0x28: "VMEM_STORE",
-  0x29: "LDS_LOAD",
-  0x2b: "LDS_STORE",
-  0x2e: "LDS_STORE",
-  0x50: "__SIMD_LDS_LOAD",
-  0x51: "__SIMD_LDS_LOAD",
-  0x54: "__SIMD_LDS_STORE",
-  0x5a: "__SIMD_VMEM_LOAD",
-  0x5b: "__SIMD_VMEM_LOAD",
-  0x5c: "__SIMD_VMEM_STORE",
-  0x5d: "__SIMD_VMEM_STORE",
-  0x5e: "__SIMD_VMEM_STORE",
-  0x5f: "__SIMD_VMEM_STORE",
-  0x72: "SALU_OR",
-  0x73: "VALU_CMPX",
-}
-
-ALUSRC = {
-  1: "SALU",
-  2: "VALU",
-  3: "VALU_SALU",
-}
-
-MEMSRC = {
-  0: "LDS",
-  1: "__LDS",
-  2: "VMEM",
-  3: "__VMEM",
-}
-
-
-# these tables are from rocprof trace decoder
-# rocprof_trace_decoder_parse_data-0x11c6a0
-# parse_sqtt_180 = b *rocprof_trace_decoder_parse_data-0x11c6a0+0x110040
-
-# ---------- 1. local_138: 256-byte state->opcode table ----------
-
-STATE_TO_OPCODE: bytes = bytes([
-  0x10, 0x16, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x17, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x07, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x19, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x00, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x11, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x12, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x15, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x16, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x17, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x07, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x19, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x00, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x11, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x13, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x15, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-])
-
-# opcode mask (the bits used to determine the opcode, worked out by looking at the repeats in STATE_TO_OPCODE)
-
-opcode_mask = {
-  0x10: 0b1111,
-
-  0x16: 0b1111111,
-  0x17: 0b1111111,
-  0x07: 0b1111111,
-  0x19: 0b1111111,
-  0x11: 0b1111111,
-  0x12: 0b11111111,
-  0x13: 0b11111111,
-  0x15: 0b1111111,
-
-  0x18: 0b111,
-  0x1: 0b111,
-
-  0x5: 0b11111,
-  0x6: 0b11111,
-  0xb: 0b11111,
-  0x8: 0b11111,
-  0xc: 0b11111,
-  0xd: 0b11111,
-
-  0xf: 0b1111,
-  0x14: 0b1111,
-
-  0x9: 0b11111,
-  0xa: 0b11111,
-
-  0x4: 0b1111,
-  0x3: 0b1111,
-  0x2: 0b1111,
-}
-
-# ---------- 2. DAT_0012e280: nibble budget per opcode&0x1F ----------
-
-NIBBLE_BUDGET = [
-  0x08, 0x0C, 0x08, 0x08, 0x0C, 0x18, 0x18, 0x40, 0x14, 0x20, 0x30, 0x14, 0x34, 0x1C, 0x30, 0x08,
-  0x04, 0x18, 0x18, 0x20, 0x40, 0x40, 0x30, 0x40, 0x14, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-]
-
-# ---------- 3. delta_map from your hash nodes ----------
-
-# opcode -> (shift, width)
-DELTA_MAP_DEFAULT = {
-  0x01: (3,  3),   # shift=3,  end=6
-  0x02: (4,  2),   # shift=4,  end=6
-  0x03: (4,  2),   # shift=4,  end=6
-  0x04: (4,  3),   # shift=4,  end=7
-  0x05: (5,  3),   # shift=5,  end=8
-  0x06: (5,  3),   # shift=5,  end=8
-  0x07: (8,  3),   # shift=8,  end=11
-  0x08: (5,  3),   # shift=5,  end=8
-  0x09: (5,  2),   # shift=5,  end=7
-  0x0A: (5,  2),   # shift=5,  end=7
-  0x0B: (5,  3),   # shift=5,  end=8
-  0x0C: (5,  3),   # shift=5,  end=8
-  0x0D: (5,  3),   # shift=5,  end=8
-  # NOTE: 0x0e can never be decoded, it's not in the STATE_TO_OPCODE table
-  #0x0E: (7,  2),   # shift=7,  end=9
-  0x0F: (4,  4),   # shift=4,  end=8
-  0x10: (0,  0),   # shift=0,  end=0  (no delta)
-  0x11: (7,  9),   # shift=7,  end=16
-  0x12: (8,  3),   # shift=8,  end=11
-  0x13: (8,  3),   # shift=8,  end=11
-  0x14: (4,  3),   # shift=4,  end=7
-  0x15: (7,  3),   # shift=7,  end=10
-  0x16: (12, 36),  # shift=12, end=48 (36-bit field, matches the 0x16 special-case)
-  0x17: (0,  0),   # shift=0,  end=0  (no delta)
-  0x18: (4,  3),   # shift=4,  end=7
-  0x19: (7,  2),   # shift=7,  end=9
-}
-
-# ---------- 4. One-line-per-packet parser ----------
-
-def reg_mask(opcode):
-  nb_bits = NIBBLE_BUDGET[opcode & 0x1F]
-  shift, width = DELTA_MAP_DEFAULT[opcode]
-  delta_mask = ((1 << width) - 1) << shift
-  assert delta_mask & opcode_mask[opcode] == 0, "masks shouldn't overlap"
-  return ((1 << nb_bits) - 1) & ~(delta_mask | opcode_mask[opcode])
-
-def decode_packet_fields(opcode: int, reg: int) -> str:
-  """
-  Decode packet payloads conservatively, using:
-    - NIBBLE_BUDGET[opcode & 0x1F] to mask reg down to true width.
-    - DELTA_MAP_DEFAULT[opcode] to expose the "primary" field (often delta).
-    - Per-opcode layouts derived from rocprof's decompiled consumers.
-  """
-  # --- 0. Restrict to real packet bits not used in delta ---------------------------------
-  pkt = reg & reg_mask(opcode)
-  fields: list[str] = []
-
-  match opcode:
-    case 0x01: # VALUINST
-      # 6 bit field
-      flag = (pkt >> 6) & 1
-      wave = pkt >> 7
-      fields.append(f"wave={wave:x}")
-      if flag: fields.append("flag")
-    case 0x02: # VMEMEXEC
-      # 2 bit field (pipe is a guess)
-      src = pkt>>6
-      fields.append(f"src={src} [{MEMSRC.get(src, '')}]")
-    case 0x03: # ALUEXEC
-      # 2 bit field
-      src = pkt>>6
-      fields.append(f"src={src} [{ALUSRC.get(src, '')}]")
-    case 0x04: # IMMEDIATE_4
-      # 5 bit field (actually 4)
-      wave = pkt >> 7
-      fields.append(f"wave={wave:x}")
-    case 0x05: # IMMEDIATE_5
-      # 16 bit field
-      # 1 bit per wave
-      fields.append(f"mask={pkt>>8:016b}")
-    case 0x6:
-      # wave ready FFFF00
-      # 16 bit field
-      # 1 bit per wave
-      fields.append(f"mask={pkt>>8:016b}")
-    case 0x0d:
-      # 20 bit field
-      fields.append(f"arg = {pkt>>8:X}")
-    case 0x12:
-      fields.append(f"event = {pkt>>11:X}")
-    case 0x15:
-      fields.append(f"snap = {pkt>>10:X}")
-    case 0x19:
-      # wave end
-      fields.append(f"ctr = {pkt>>9:X}")
-    case 0xf:
-      extracted_delta = (reg >> 4) & 0xF
-      fields.append(f"strange_delta=0x{extracted_delta:x}")
-    case 0x11:
-      # DELTA_MAP_DEFAULT: shift=7, width=9 -> small delta.
-      # FF0000 is the mask
-      coarse    = pkt >> 16
-      fields.append(f"coarse=0x{coarse:02x}")
-      # From decomp:
-      #  - when layout<3 and coarse&1, it sets a "has interesting wave" flag
-      #  - when coarse&8, it marks all live waves as "terminated"
-      if coarse & 0x01:
-        fields.append("flag_wave_interest=1")
-      if coarse & 0x08:
-        fields.append("flag_terminate_all=1")
-    case 0x8:
-      # wave end, this is 20 bits (FFF00)
-      flag7   = (pkt >> 8) & 1
-      simd    = (pkt >> 9) & 3
-      cu      = ((pkt >> 11) & 0x7) | (flag7 << 3)
-      wave    = (pkt >> 15) & 0x1f
-      fields.append(f"wave={wave:x}")
-      fields.append(f"simd={simd}")
-      fields.append(f"cu={cu}")
-    case 0x9:
-      # From case 9 (WAVESTART) in multiple consumers:
-      #   flag7  = (w >> 7) & 1        (low bit of uVar41)
-      #   cls2   = (w >> 8) & 3        (class / group)
-      #   slot4  = (w >> 10) & 0xf     (slot / group index)
-      #   idx_lo = (w >> 0xd) & 0x1f   (low index, layout<4 path)
-      #   idx_hi = (w >> 0xf) & 0x1f   (high index, layout>=4 path)
-      #   id7    = (w >> 0x19) & 0x7f  (7-bit id)
-      flag7   = (pkt >> 7) & 1
-      simd    = (pkt >> 8) & 3
-      cu      = ((pkt >> 10) & 0x7) | (flag7 << 3)
-      wave    = (pkt >> 13) & 0x1F
-      id7     = (pkt >> 17)
-      fields.append(f"wave={wave:x}")
-      fields.append(f"simd={simd}")
-      fields.append(f"cu={cu}")
-      fields.append(f"id7=0x{id7:x}")
-    case 0x18:
-      # FFF88 is the mask
-      # From case 0x18:
-      #   low3   = w & 7
-      #   grp3   = (w >> 3) or (w >> 4) & 7   (layout-dependent)
-      #   flags  = bits 6 (B6) and 7 (B7)
-      #   hi8    = (w >> 0xc) & 0xff   (layout 4 path)
-      #   hi7    = (w >> 0xd) & 0x7f   (other layouts)
-      #   idx5   = (w >> 7) or (w >> 8) & 0x1f, used as wave index
-      flag1 = (pkt >> 3) & 1
-      flag2 = (pkt >> 7) & 1
-      wave = (pkt >> 8) & 0x1F
-      op = (pkt >> 13)
-      fields.append(f"wave={wave:x}")
-      fields.append(f"op=0x{op:02x} [{OPNAME.get(op, '')}]")
-      if flag1: fields.append("flag1")
-      if flag2: fields.append("flag2")
-    case 0x14:
-      subop   = (pkt >> 16) & 0xFFFF       # (short)(w >> 0x10)
-      val32   = (pkt >> 32) & 0xFFFFFFFF   # (uint)(w >> 0x20)
-      slot    = (pkt >> 7) & 0x7           # index in local_168[...] tables
-      hi_byte = (pkt >> 8) & 0xFF          # determines config vs marker
-
-      fields.append(f"subop=0x{subop:04x}")
-      fields.append(f"slot={slot}")
-      fields.append(f"val32=0x{val32:08x}")
-
-      if hi_byte & 0x80:
-        # Config flavour: writes config words into per-slot state arrays.
-        fields.append("kind=config")
-        if subop == 0x000C:
-          fields.append("slot=lo")
-        elif subop == 0x000D:
-          fields.append("slot=hi")
-      else:
-        # COR marker: subop 0xC342, payload "COR\0" → start of a COR region.
-        if subop == 0xC342:
-          fields.append("kind=cor_stream")
-          if val32 == 0x434F5200:
-            fields.append("cor_magic='COR\\0'")
-    case 0x16:
-      # Bits:
-      #   bit8  -> 0x100
-      #   bit9  -> 0x200
-      #   bits 12..47 -> 36-bit field used as delta or marker
-      bit8 = bool(pkt & 0x100)
-      bit9 = bool(pkt & 0x200)
-      if not bit9:
-        mode = "delta"
-      elif not bit8:
-        mode = "marker"
-      else:
-        mode = "other"
-      # need to use reg here
-      val36 = (reg >> 12) & ((1 << 36) - 1)
-      fields.append(f"mode={mode}")
-      if mode != "delta":
-        fields.append(f"val36=0x{val36:x}")
-    case 0x17:
-      # From decomp (two sites with identical logic):
-      #   layout = (w >> 7) & 0x3f
-      #   mode   = (w >> 0xd) & 3
-      #   group  = (w >> 0xf) & 7
-      #   sel_a  = (w >> 0x1c) & 0xf
-      #   sel_b  = (w >> 0x21) & 7
-      #   flag4  = (w >> 0x3b) & 1  (only meaningful when layout == 4)
-      layout = (pkt >> 7)  & 0x3F
-      simd   = (pkt >> 13) & 0x3  # you can change this by changing traced simd
-      group  = (pkt >> 15) & 0x7
-      sel_a  = (pkt >> 0x1C) & 0xF
-      sel_b  = (pkt >> 0x21) & 0x7
-      flag4  = (pkt >> 0x3B) & 0x1
-
-      fields.append(f"layout={layout}")
-      fields.append(f"group={group}")
-      fields.append(f"simd={simd}")
-      fields.append(f"sel_a={sel_a}")
-      fields.append(f"sel_b={sel_b}")
-      if layout == 4:
-        fields.append(f"layout4_flag={flag4}")
-    case _:
-      fields.append(f"{pkt:X} & {reg_mask(opcode):X}")
-  return ",".join(fields)
-
-FILTER_LEVEL = getenv("FILTER", 1)
-
-DEFAULT_FILTER: tuple[int, ...] = tuple()
-# NOP + pure time + "sample"
-if FILTER_LEVEL >= 0: DEFAULT_FILTER += (0x10, 0xf, 0x11)
-# reg + event + sample + marker
-# TODO: events are probably good
-if FILTER_LEVEL >= 1: DEFAULT_FILTER += (0x14, 0x12, 0x16)
-# instruction runs + valuinst
-if FILTER_LEVEL >= 2: DEFAULT_FILTER += (0x01, 0x02, 0x03)
-# instructions dispatch (inst, immed)
-if FILTER_LEVEL >= 3: DEFAULT_FILTER += (0x4, 0x5, 0x18)
-# waves
-if FILTER_LEVEL >= 4: DEFAULT_FILTER += (0x6, 0x8, 0x9)
-
-def parse_sqtt_print_packets(data: bytes, filter=DEFAULT_FILTER, verbose=True) -> None:
-  """
-  Minimal debug: print ONE LINE per decoded token (packet).
-
-  Now prints only the actual nibbles that belong to each packet, instead of
-  the full 64-bit shift register.
-  """
-  n = len(data)
-  time = 0
-  last_printed_time = 0
-  reg = 0          # shift register
-  offset = 0       # bit offset, in steps of 4 (one nibble)
-  nib_budget = 0x40
-  flags = 0
-  token_index = 0
-  opcodes_seen = set()
-
-  while (offset >> 3) < n:
-    # 1) Fill register with nibbles according to nib_budget
-    if nib_budget != 0:
-      target = offset + 4 + ((nib_budget - 1) & ~3)
-      while offset != target and (offset >> 3) < n:
-        byte = data[offset >> 3]
-        nib = (byte >> (offset & 4)) & 0xF
-        reg = ((reg >> 4) | (nib << 60)) & ((1 << 64) - 1)
-        offset += 4
-      if offset != target: break  # don't parse past the end
-
-    # 2) Decode token from low 8 bits
-    opcode = STATE_TO_OPCODE[reg & 0xFF]
-    opcodes_seen.add(opcode)
-
-    # 4) Set next nibble budget based on opcode
-    nib_budget = NIBBLE_BUDGET[opcode & 0x1F]
-
-    # 5) Get delta
-    shift, width = DELTA_MAP_DEFAULT[opcode]
-    delta = (reg >> shift) & ((1 << width) - 1)
-
-    # 6) Update time and handle special opcodes 0xF/0x16
-    if opcode == 0x16:
-      two_bits = (reg >> 8) & 0x3
-      if two_bits == 1:
-        flags |= 0x01
-
-      # Common 36-bit field at bits [12..47]
-      if (reg & 0x200) == 0:
-        # delta mode: add 36-bit delta to time
-        pass
-      elif (reg & 0x100) == 0:
-        # marker / other modes: no time advance
-        # real marker: bit9=1, bit8=0, non-zero payload
-        # "other" 0x16 variants, ignored for timing
-        delta = 0
-      else:
-        raise RuntimeError("unknown 0x16 delta")
-    elif opcode == 0x0F:
-      # opcode 0x0F has an offset of 4 to the delta
-      # update: it's actually computed to be 8 to match WAVESTART
-      delta = delta + 8
-
-    # Append extra decoded fields into the note string
-    note = decode_packet_fields(opcode, reg)
-
-    # this delta happens before the instruction
-    time += delta
-    token_index += 1
-
-    if verbose and (filter is None or opcode not in filter):
-      print(f"{time:8d} +{time-last_printed_time:8d} : "+colored(f"{OPCODE_NAMES[opcode]:18s} ", OPCODE_COLORS.get(opcode, "white"))+f"{note}")
-      last_printed_time = time
-
-  # Optional summary at the end
-  print(f"# done: tokens={token_index:_}, final_time={time}, flags=0x{flags:02x}")
-  if verbose:
-    print(f"opcodes({len(opcodes_seen):2d}):",
-          ' '.join([colored(f"{op:2X}", "WHITE" if op in opcodes_seen else "BLACK") for op in sorted(opcode_mask)]))
-
-
-def parse(fn:str):
-  with Timing(f"unpickle {fn}: "): dat = pickle.load(open(fn, "rb"))
-  #if getenv("ROCM", 0):
-  #  with Timing(f"decode {fn}: "): ctx = decode(dat)
-  dat_sqtt = [x for x in dat if isinstance(x, ProfileSQTTEvent)]
-  print(f"got {len(dat_sqtt)} SQTT events in {fn}")
-  return dat_sqtt
-
-if __name__ == "__main__":
-  fn = "extra/sqtt/examples/profile_gemm_run_0.pkl"
-  dat_sqtt = parse(sys.argv[1] if len(sys.argv) > 1 else fn)
-  for i,dat in enumerate(dat_sqtt):
-    with Timing(f"decode pkt {i} with len {len(dat.blob):_}: "):
-      parse_sqtt_print_packets(dat.blob, verbose=getenv("V", 1))
@@ -1,25 +1,25 @@
-import os, subprocess, sys
+import os, subprocess, sys, shlex
 from pathlib import Path
 from tinygrad.helpers import temp

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

-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"
-]
+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"
+}

 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 test in EXAMPLES:
+  for name,test in EXAMPLES.items():
    for i in range(2):
      # AM_RESET=1 gets a clear trace, does not work on mi300 machines
-      subprocess.run([sys.executable, *test.split()], cwd=EXAMPLES_DIR.parent.parent.parent,
+      subprocess.run([sys.executable, *shlex.split(test)], 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_{test.split('.')[-1].replace('test_', '')}_run_{i}.pkl")
+      PROFILE_PATH.rename(dest:=EXAMPLES_DIR/arch/f"profile_{name}_run_{i}.pkl")
      print(f"saved SQTT trace to {dest}")
@@ -118,7 +118,7 @@ def decode(sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, Inst]])
    nonlocal exc
    try: rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
    except AttributeError as e:
-      exc = RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install")
+      exc = RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_rocprof_decoder.py to install")
      exc.__cause__ = e
  (t:=threading.Thread(target=worker, daemon=True)).start()
  t.join()
@@ -19,6 +19,38 @@ 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"
@@ -45,23 +77,7 @@ 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)

-  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
+  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)

  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]

@@ -89,7 +105,6 @@ 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)
@@ -120,7 +135,6 @@ 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)
@@ -138,7 +152,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_KV, N // BLOCK_SIZE_KV, B)
+  gsz = (H, 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")
@@ -151,7 +165,6 @@ 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)
@@ -182,7 +195,6 @@ 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, ATTN_N, ATTN_H_KV, ATTN_D>;
+using _gl_dKV  = gl<bf16, ATTN_B * GROUP_SIZE, 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_KV, (ATTN_N / BLOCK_SIZE_KV), ATTN_B); }
+  dim3 grid() { return dim3(ATTN_H, (ATTN_N / BLOCK_SIZE_KV), ATTN_B); }
  dim3 block() { return dim3(NUM_THREADS); }
  size_t dynamic_shared_memory() { return MAX_SHARED_MEMORY; }
 };
@@ -55,10 +55,12 @@ 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 kv_head_idx = blockIdx.x;  // This is the KV head index
+  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 seq_idx = blockIdx.y;
  const int batch_idx = blockIdx.z;
-  const int first_q_head = kv_head_idx * GROUP_SIZE;
+  const int first_q_head = q_head_idx_fixed;

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

  constexpr float L_SCALE_FACTOR = 1.44269504089f;
@@ -3355,14 +3357,14 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
    }
  }

-  store<1>(g.dVg, dV_j, {batch_idx, 0, kv_head_idx, 0}, {0, j, 0, 0});
+  store<1>(g.dVg, dV_j, {batch_idx * GROUP_SIZE + q_head_in_group, 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, 0, kv_head_idx, 0}, {0, j, 0, 0});
+  store<1>(g.dKg, dV_j, {batch_idx * GROUP_SIZE + q_head_in_group, 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<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<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<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;
+    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;

    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<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+        att_block_bf16_in = *reinterpret_cast<attn_tile< 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<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+        att_block_bf16_in = *reinterpret_cast<attn_tile<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<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+    att_block_bf16_in = *reinterpret_cast<attn_tile<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<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+    att_block_bf16_in = *reinterpret_cast<attn_tile<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<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+    att_block_bf16_in = *reinterpret_cast<attn_tile<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<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+    att_block_bf16_in = *reinterpret_cast<attn_tile<bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);

    __builtin_amdgcn_sched_barrier(0);
    mul_col(o_reg, o_reg, scale_vec);
@@ -505,7 +505,9 @@ 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,
@@ -1,7 +1,6 @@
 # simple tests
 import unittest
 import torch
-import warnings
 from tinygrad.helpers import getenv, GlobalCounters
 if getenv("TINY_BACKEND2"):
  import extra.torch_backend.backend2
@@ -18,9 +17,7 @@ class TestKernelFusionRegression(unittest.TestCase):
    torch.manual_seed(42)
    GlobalCounters.reset()
    fn().detach().cpu().numpy()
-    expectation = f"{GlobalCounters.kernel_count} vs {expected_kernels} expected."
-    if GlobalCounters.kernel_count < expected_kernels: warnings.warn(f"{expectation} Expectation can be lowered.", UserWarning)
-    self.assertLessEqual(GlobalCounters.kernel_count, expected_kernels, f"{expectation}")
+    self.assertEqual(GlobalCounters.kernel_count, expected_kernels)

  def test_elementwise_fusion(self):
    def fn():
@@ -34,7 +31,7 @@ class TestKernelFusionRegression(unittest.TestCase):
      conv = torch.nn.Conv2d(3, 16, 3, padding=1).to(device)
      with torch.no_grad():
        return torch.nn.functional.relu(conv(x))
-    self._check_kernel_count(fn, 8)
+    self._check_kernel_count(fn, 6)

  def test_batchnorm_fusion(self):
    def fn():
@@ -44,7 +41,7 @@ class TestKernelFusionRegression(unittest.TestCase):
      bn.eval()
      with torch.no_grad():
        return torch.nn.functional.relu(bn(conv(x)))
-    self._check_kernel_count(fn, 16)
+    self._check_kernel_count(fn, 10)

  def test_reduce_fusion(self):
    def fn():
@@ -92,7 +89,7 @@ class TestKernelFusionRegression(unittest.TestCase):
        out = bn(conv(x))
        out += identity
        return torch.nn.functional.relu(out)
-    self._check_kernel_count(fn, 17)
+    self._check_kernel_count(fn, 12)

  def test_multiple_inplace_ops_fusion(self):
    def fn():
@@ -117,7 +114,7 @@ class TestKernelFusionRegression(unittest.TestCase):
      bn.train()
      with torch.no_grad():
        return bn(x)
-    self._check_kernel_count(fn, 10)
+    self._check_kernel_count(fn, 8)

  # this is a minimal extra/other_mnist/beautiful_mnist_torch.py to cover fusion for training with optimizer
  def test_mnist_training_fusion(self):
@@ -138,7 +135,7 @@ class TestKernelFusionRegression(unittest.TestCase):
      loss.backward()
      optimizer.step()
      return loss
-    self._check_kernel_count(fn, 28)
+    self._check_kernel_count(fn, 24)

 if __name__ == "__main__":
  unittest.main()
@@ -0,0 +1,17 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>com.apple.application-identifier</key>
+	<string>9YG3G8543N.org.tinygrad.tinygpu.edriver</string>
+	<key>com.apple.developer.driverkit</key>
+	<true/>
+	<key>com.apple.developer.driverkit.transport.pci</key>
+	<array>
+		<dict>
+			<key>IOPCIPrimaryMatch</key>
+			<string>0x000010de&amp;0x0000FFFF</string>
+		</dict>
+	</array>
+</dict>
+</plist>
@@ -0,0 +1,33 @@
+#!/bin/bash
+set -e
+
+xcodebuild clean build CODE_SIGN_IDENTITY="" CODE_SIGNING_REQUIRED=NO -alltargets -configuration Release build
+
+cp "../profiles/edriver_rel_2.provisionprofile" "./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext/embedded.provisionprofile"
+cp "../profiles/installer_provisioning.provisionprofile" "./build/Release/TinyGPU.app/Contents/embedded.provisionprofile"
+
+codesign \
+    --sign "Developer ID Application: tinygrad, Corp. (9YG3G8543N)" \
+    --entitlements ./TinyGPUDriverExtension/TinyGPUDriver.NV.Release.entitlements \
+    --verbose \
+    --options runtime \
+    --timestamp \
+    --force \
+    ./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext
+
+codesign \
+    --sign "Developer ID Application: tinygrad, Corp. (9YG3G8543N)" \
+    --entitlements ./macOS/macOS.entitlements \
+    --options runtime \
+    --verbose \
+    --timestamp \
+    --force \
+    ./build/Release/TinyGPU.app
+
+codesign --verify --deep --strict --verbose=4 ./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext
+
+codesign --verify --deep --strict --verbose=4 ./build/Release/TinyGPU.app
+
+spctl -a -vv ./build/Release/TinyGPU.app
+
+spctl -a -vv ./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext
@@ -1,17 +1,17 @@
 A command line tool for exploring the VIZ trace.

-After running with VIZ=-1, use `PYTHONPATH=. extra/viz/cli.py` to explore the saved trace files.
+After running with VIZ=-1, use `extra/viz/cli.py` to explore the saved trace files.

 ## Inspect runtime profiling

-Use `PYTHONPATH=. extra/viz/cli.py --profile` to list all traced devices.
+Use `extra/viz/cli.py --profile` to list all traced devices.

 List top slowest kernels on a device: `--profile --device "AMD"`
 List samples of a kernel on a device: `--profile --device "AMD" --kernel E_3`

 ## Inspect codegen and PatternMatcher

-Use `PYTHONPATH=. extra/viz/cli.py --rewrites` to list all traced kernels.
+Use `extra/viz/cli.py --rewrites` to list all traced kernels.

 List all codegen steps for a kernel: `--rewrites --kernel E_3`
 Get source code: `--rewrites --kernel E_3 --select "View Source"`
@@ -1,44 +1,73 @@
 #!/usr/bin/env python3
 import os
 os.environ["VIZ"] = "0"
-import argparse, pathlib
+import argparse, pathlib, sys, struct, json
 from typing import Iterator
 from tinygrad.viz import serve as viz
 from tinygrad.uop.ops import RewriteTrace
 from tinygrad.helpers import temp, ansistrip, colored, time_to_str, ansilen
-from test.null.test_viz import load_profile
+
+# ** generic helpers

 def optional_eq(val:dict, arg:str|None) -> bool: return arg is None or ansistrip(val["name"]) == arg

 def print_data(data:dict) -> None:
  if isinstance(data.get("value"), Iterator):
    for m in data["value"]:
-      if m.get("uop"):
-        print("Input UOp:")
-        print(m["uop"])
-      if not m["diff"]: continue
-      print("Rewrites:")
-      fp = pathlib.Path(m["upat"][0][0])
-      print(f"{fp.parent.name}/{fp.name}:{m['upat'][0][1]}")
-      print(m["upat"][1])
-      for line in m["diff"]:
-        color = "red" if line.startswith("-") else "green" if line.startswith("+") else None
-        print(colored(line, color))
+      if m.get("uop"): print(f"Input UOp:\n{m['uop']}")
+      if m.get("diff"):
+        loc = pathlib.Path(m["upat"][0][0])
+        print(f"Rewrite at {loc.parent.name}/{loc.name}:{m['upat'][0][1]}\n{m['upat'][1]}")
+        for line in m["diff"]: print(colored(line, "red" if line.startswith("-") else "green" if line.startswith("+") else None))
  if data.get("src") is not None: print(data["src"])

+# ** Profiler trace decoder
+
+# 0 means None, otherwise it's an enum value
+def option(i:int) -> int|None: return None if i == 0 else i-1
+
+def decode_profile(data:bytes) -> dict:
+  ret, off = data, 0
+  def u(fmt:str) -> tuple:
+    nonlocal off
+    vals = struct.unpack_from(fmt, ret, off)
+    off += struct.calcsize(fmt)
+    return vals
+  total_dur, global_peak, index_len, layout_len = u("<IQII")
+  strings, dtypes, markers = json.loads(ret[off:off+index_len]).values()
+  off += index_len
+  layout:dict[str, dict] = {}
+  for _ in range(layout_len):
+    klen = u("<B")[0]
+    k = ret[off:off+klen].decode()
+    off += klen
+    layout[k] = v = {"events":[]}
+    event_type, event_count = u("<BI")
+    if event_type == 0:
+      for _ in range(event_count):
+        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["peak"] = u("<Q")[0]
+      for _ in range(event_count):
+        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__":
  parser = argparse.ArgumentParser()
  g_mode = parser.add_argument_group("mode")
  g_mode.add_argument("--profile", action="store_true", help="View profile trace")
  g_mode.add_argument("--rewrites", action="store_true", help="View rewrites trace")
+  g_common = parser.add_argument_group("common options")
+  g_common.add_argument("--kernel", type=str, default=None, metavar="NAME", help="Select a kernel by name (optional name, default: only list names)")
  g_profile = parser.add_argument_group("profile options")
  g_profile.add_argument("--device", type=str, default=None, metavar="NAME", help="Select a device (optional name, default: only list names)")
  g_profile.add_argument("--top", type=int, default=10, metavar="N", help="Number of top kernels to show (-1 for all, default: 10)")
  g_rewrites = parser.add_argument_group("rewrites options")
  g_rewrites.add_argument("--select", type=str, default=None, metavar="NAME",
                          help="Select an item within the chosen kernel (optional name, default: only list names)")
-  g_common = parser.add_argument_group("common options")
-  g_common.add_argument("--kernel", type=str, default=None, metavar="NAME", help="Select a kernel by name (optional name, default: only list names)")
  parser.add_argument("--profile-path", type=pathlib.Path, metavar="PATH", help="Path to profile (optional file, default: latest profile)",
                        default=pathlib.Path(temp("profile.pkl", append_user=True)))
  parser.add_argument("--rewrites-path", type=pathlib.Path, metavar="PATH", help="Path to rewrites (optional file, default: latest rewrites)",
@@ -46,14 +75,14 @@ if __name__ == "__main__":
  args = parser.parse_args()
  if not args.profile and not args.rewrites:
    parser.print_help()
-    exit(0)
+    sys.exit(0)

  viz.trace = viz.load_pickle(args.rewrites_path, default=RewriteTrace([], [], {}))
  viz.ctxs = viz.get_rewrites(viz.trace)

  if args.profile:
    from tabulate import tabulate
-    profile = load_profile(viz.load_pickle(args.profile_path, default=[]))
+    profile = decode_profile(viz.get_profile(viz.load_pickle(args.profile_path, default=[])))
    agg, total, n = {}, 0, 0
    if args.device is None: print("Select a device:")
    for k,v in profile["layout"].items():
@@ -63,7 +92,7 @@ if __name__ == "__main__":
      for e in v.get("events", []):
        et = e["dur"]*1e-6
        if args.kernel is not None:
-          if ansistrip(e["name"]) == args.kernel and n < 10:
+          if optional_eq(e, args.kernel) and n < 10:
            ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
            name = e["name"]+(" " * (46 - ansilen(e["name"])))
            print(f"{name} {ptm}/{(et or 0)*1e3:9.2f}ms  "+e['fmt'].replace('\n', ' | ')+"  ")
@@ -81,7 +110,7 @@ if __name__ == "__main__":
        other_t = total-sum(t for _, (t, _) in sel)
        table.append([f"Other ({len(other)} unique)", time_to_str(other_t, w=9), sum(c for _,(_,c) in other), f"{other_t/total*100.0:.2f}%"])
      print(tabulate(table, headers=["name", "total", "count", "pct"], tablefmt="github"))
-    exit(0)
+    sys.exit(0)

  for k in viz.ctxs:
    if not optional_eq(k, args.kernel): continue
@@ -104,6 +104,34 @@ class TestCmpClass(unittest.TestCase):
    st = run_program(instructions, n_lanes=1)
    self.assertEqual(st.vcc & 1, 0, "Signaling NaN should not match quiet mask")

+  def test_v_cmp_lg_f32_nan(self):
+    """v_cmp_lg_f32 is ordered not-equal (<>): NaN <> x should be False per IEEE 754."""
+    quiet_nan = 0x7fc00000
+    one_f32 = 0x3f800000  # 1.0f
+    instructions = [
+      s_mov_b32(s[0], quiet_nan),
+      v_mov_b32_e32(v[0], s[0]),
+      s_mov_b32(s[1], one_f32),
+      v_mov_b32_e32(v[1], s[1]),
+      v_cmp_lg_f32_e32(v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vcc & 1, 0, "v_cmp_lg_f32(NaN, 1.0) should be 0")
+
+  def test_v_cmp_neq_f32_nan(self):
+    """v_cmp_neq_f32 is unordered not-equal (!=): NaN != x should be True per IEEE 754."""
+    quiet_nan = 0x7fc00000
+    one_f32 = 0x3f800000  # 1.0f
+    instructions = [
+      s_mov_b32(s[0], quiet_nan),
+      v_mov_b32_e32(v[0], s[0]),
+      s_mov_b32(s[1], one_f32),
+      v_mov_b32_e32(v[1], s[1]),
+      v_cmp_neq_f32_e32(v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vcc & 1, 1, "v_cmp_neq_f32(NaN, 1.0) should be 1")
+
  def test_v_cmp_sets_vcc_bits(self):
    """V_CMP_EQ sets VCC bits based on per-lane comparison."""
    instructions = [
@@ -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_STORE}
+OTHER_SIMD_OPS_RDNA4 = {InstOpRDNA4.OTHER_VMEM, InstOpRDNA4.OTHER_VMEM_5}

 # ═══════════════════════════════════════════════════════════════════════════════
 # ROCPROF DECODER
@@ -153,7 +153,9 @@ 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)]
-        self.assertGreater(len([p for p in all_packets if isinstance(p, (INST, INST_RDNA4))]), 0, f"no INST packets in {name}")
+        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}")

  expected: dict[str, list[int]] = {}  # override in subclasses
  def test_packet_counts(self):
@@ -208,12 +210,14 @@ class SQTTExamplesTestBase(unittest.TestCase):
 class TestSQTTExamplesRDNA3(SQTTExamplesTestBase):
  target = "gfx1100"
  expected = {
-    "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],
+    "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],
  }

 class TestSQTTExamplesRDNA4(SQTTExamplesTestBase): target = "gfx1200"
@@ -1,186 +0,0 @@
-"""Tests comparing sqtt.py PACKET_TYPES_RDNA3/RDNA4 against AMD's rocprof-trace-decoder binary."""
-import unittest, struct, ctypes, pickle
-from pathlib import Path
-
-ROCPROF_LIB = Path("/usr/lib/librocprof-trace-decoder.so")
-import tinygrad
-EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"
-
-# CDNA pkt_fmt -> size in bytes (extracted from rocprof hash table)
-CDNA_PKT_SIZES = {0: 2, 1: 8, 2: 8, 3: 4, 4: 2, 5: 6, 6: 2, 7: 2, 8: 2, 9: 2, 10: 2, 11: 8, 12: 6, 13: 4, 14: 8, 15: 6}
-
-def _find_segment(perms: str):
-  """Find a segment of the loaded library with given permissions (e.g. 'rw-p', 'r--p')."""
-  with open('/proc/self/maps', 'r') as f:
-    for line in f:
-      if 'librocprof-trace-decoder.so' in line and f' {perms} ' in line:
-        parts = line.split()
-        return int(parts[0].split('-')[0], 16), int(parts[2], 16)
-  return None, None
-
-def _read_array(file_offset: int, count: int):
-  """Read an array of uint8 at file_offset from the loaded library."""
-  base, seg_offset = _find_segment('rw-p')
-  if base is None: return None
-  return list((ctypes.c_uint8 * count).from_address(base + (file_offset - seg_offset)))
-
-def _load_lib():
-  if not ROCPROF_LIB.exists(): return False
-  ctypes.CDLL(str(ROCPROF_LIB))
-  return True
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# RDNA EXTRACTION (nibble-based format)
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def extract_bit_tables():
-  """Extract bit budget tables. Returns (layout2, layout3, layout4) or None."""
-  if not _load_lib(): return None
-  return _read_array(0x2d220, 32), _read_array(0x2d280, 32), _read_array(0x2d2c0, 32)
-
-def extract_delta_fields():
-  """Extract delta bitfield tables. Returns (layout2, layout3, layout4) dicts mapping type_id -> (lo, hi)."""
-  if not _load_lib(): return None
-  ro_base, ro_offset = _find_segment('r--p')
-  if ro_base is None: return None
-
-  def read_table(file_offset, num_entries):
-    addr = ro_base + (file_offset - ro_offset)
-    data = bytes((ctypes.c_uint8 * (num_entries * 12)).from_address(addr))
-    return {type_id: (lo, hi) for j in range(0, len(data), 12)
-            for type_id, lo, hi in [struct.unpack('<III', data[j:j+12])] if type_id < 32}
-
-  return read_table(0x26800, 24), read_table(0x26dc0, 25), read_table(0x27300, 27)
-
-def extract_packet_encodings():
-  """Extract packet encodings. Returns (L2, L3, L4) dicts mapping type_id -> (mask, value)."""
-  if not _load_lib(): return None
-  rw_base, rw_offset = _find_segment('rw-p')
-  if rw_base is None: return None
-
-  # Read base encodings from registration vector at 0x2d340
-  vec_start = ctypes.c_void_p.from_address(rw_base + (0x2d340 - rw_offset)).value
-  vec_end = ctypes.c_void_p.from_address(rw_base + (0x2d348 - rw_offset)).value
-  base = {}
-  if vec_start and vec_end:
-    for i in range((vec_end - vec_start) // 32):
-      addr = vec_start + i * 32
-      type_id = ctypes.c_uint8.from_address(addr).value
-      pat_start = ctypes.c_void_p.from_address(addr + 8).value
-      pat_end = ctypes.c_void_p.from_address(addr + 16).value
-      if pat_start and pat_end and 0 < (n := pat_end - pat_start) <= 8:
-        pat = list((ctypes.c_uint8 * n).from_address(pat_start))
-        base[type_id] = (sum(1 << j for j in range(n)), sum(b << j for j, b in enumerate(pat)))
-
-  return {**base, 17: (0x7f, 0x51), 25: (0x7f, 0x31)}, base, {**base}  # L2 has overrides
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# CDNA EXTRACTION (16-bit header format)
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def extract_cdna_packet_sizes():
-  """Extract CDNA pkt_fmt -> size mapping by running rocprof decoder to populate its hash table."""
-  if not _load_lib(): return None
-  from test.amd.test_sqtt_examples import run_rocprof_decoder
-
-  if not (pkl_path := next((EXAMPLES_DIR / "gfx950").glob("*.pkl"), None)): return None
-  with open(pkl_path, "rb") as f: data = pickle.load(f)
-  sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
-  prg = next((e for e in data if type(e).__name__ == "ProfileProgramEvent"), None)
-  if not sqtt_events or not prg: return None
-
-  # Run decoder to trigger hash table initialization
-  run_rocprof_decoder([e.blob for e in sqtt_events], prg.lib, prg.base, "gfx950")
-
-  # Extract hash table: head at 0x2d4f0, nodes are 16 bytes (next[8], key[4], value[4])
-  rw_base, rw_offset = _find_segment('rw-p')
-  if not (head := ctypes.c_void_p.from_address(rw_base + (0x2d4f0 - rw_offset)).value if rw_base else None): return None
-
-  pkt_sizes: dict[int, int] = {}
-  node, seen = head, set()
-  while node and node not in seen and len(pkt_sizes) < 20:
-    seen.add(node)
-    key, val = ctypes.c_uint32.from_address(node + 8).value, ctypes.c_uint32.from_address(node + 12).value
-    if key < 16 and val in (0x10, 0x20, 0x30, 0x40): pkt_sizes[key] = {0x10: 2, 0x20: 4, 0x30: 6, 0x40: 8}[val]
-    node = ctypes.c_void_p.from_address(node).value  # type: ignore[assignment]
-  return pkt_sizes if len(pkt_sizes) == 16 else None
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# TESTS
-# ═══════════════════════════════════════════════════════════════════════════════
-
-class TestSQTTMatchesBinary(unittest.TestCase):
-  def test_bit_counts_match_layout3(self): self._test_bit_counts(3)
-  def test_bit_counts_match_layout4(self): self._test_bit_counts(4)
-  def test_encodings_match_layout3(self): self._test_encodings(3)
-  def test_encodings_match_layout4(self): self._test_encodings(4)
-  def test_delta_fields_match_layout3(self): self._test_delta_fields(3)
-  def test_delta_fields_match_layout4(self): self._test_delta_fields(4)
-
-  def test_cdna_packet_sizes(self):
-    """Extract and verify CDNA pkt_fmt -> size mapping from rocprof's hash table."""
-    if not (EXAMPLES_DIR / "gfx950").exists(): self.skipTest("no CDNA examples")
-    if not (pkt_sizes := extract_cdna_packet_sizes()): self.skipTest("rocprof-trace-decoder not installed")
-    for pkt_fmt, size in CDNA_PKT_SIZES.items():
-      with self.subTest(pkt_fmt=pkt_fmt): self.assertEqual(pkt_sizes.get(pkt_fmt), size)
-
-  def test_cdna_packet_definitions(self):
-    from tinygrad.renderer.amd.sqtt import PACKET_TYPES_CDNA
-    for pkt_fmt, pkt_cls in PACKET_TYPES_CDNA.items():
-      with self.subTest(packet=pkt_cls.__name__):
-        self.assertEqual(pkt_cls.encoding.default, pkt_fmt)
-        self.assertEqual(CDNA_PKT_SIZES[pkt_fmt] * 2, pkt_cls._size_nibbles)  # type: ignore[attr-defined]
-
-  def _test_bit_counts(self, layout: int):
-    if not (tables := extract_bit_tables()): self.skipTest("rocprof-trace-decoder not installed")
-    from tinygrad.renderer.amd.sqtt import PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4
-    # rocprof's bit table says L4 type 7 (TS_DELTA_S8_W3) is 72 bits, but the actual decoder uses 64 bits
-    skip = {(4, 7)}
-    for type_id, pkt_cls in {3: PACKET_TYPES_RDNA3, 4: PACKET_TYPES_RDNA4}[layout].items():
-      if (layout, type_id) in skip: continue
-      with self.subTest(packet=pkt_cls.__name__):
-        self.assertEqual(pkt_cls._size_nibbles * 4, tables[layout - 2][type_id])  # type: ignore[attr-defined]
-
-  def _test_encodings(self, layout: int):
-    if not (encodings := extract_packet_encodings()): self.skipTest("rocprof-trace-decoder not installed")
-    from tinygrad.renderer.amd.sqtt import PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4
-    for type_id, pkt_cls in {3: PACKET_TYPES_RDNA3, 4: PACKET_TYPES_RDNA4}[layout].items():
-      with self.subTest(packet=pkt_cls.__name__):
-        self.assertEqual((pkt_cls.encoding.mask, pkt_cls.encoding.default), encodings[layout - 2][type_id])
-
-  def _test_delta_fields(self, layout: int):
-    if not (deltas := extract_delta_fields()): self.skipTest("rocprof-trace-decoder not installed")
-    from tinygrad.renderer.amd.sqtt import PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4
-    for type_id, pkt_cls in {3: PACKET_TYPES_RDNA3, 4: PACKET_TYPES_RDNA4}[layout].items():
-      if type_id not in deltas[layout - 2]: continue
-      delta = getattr(pkt_cls, 'delta', None)
-      actual = (0, 0) if delta is None else (delta.lo, delta.hi + 1)
-      with self.subTest(packet=pkt_cls.__name__): self.assertEqual(actual, deltas[layout - 2][type_id])
-
-if __name__ == "__main__":
-  tables = extract_bit_tables()
-  encodings = extract_packet_encodings()
-  deltas = extract_delta_fields()
-
-  TYPE_NAMES = {1: 'VALUINST', 2: 'VMEMEXEC', 3: 'ALUEXEC', 4: 'IMMEDIATE', 5: 'IMMEDIATE_MASK', 6: 'WAVERDY',
-    7: 'TS_DELTA_S8_W3', 8: 'WAVEEND', 9: 'WAVESTART', 10: 'TS_DELTA_S5_W2', 11: 'WAVEALLOC', 12: 'TS_DELTA_S5_W3',
-    13: 'PERF', 14: 'UTILCTR', 15: 'TS_DELTA_SHORT', 16: 'NOP', 17: 'TS_WAVE_STATE', 18: 'EVENT', 19: 'EVENT_BIG',
-    20: 'REG', 21: 'SNAPSHOT', 22: 'TS_DELTA_OR_MARK', 23: 'LAYOUT_HEADER', 24: 'INST', 25: 'UNK_25'}
-
-  print("L2:", tables[0], "\nL3:", tables[1], "\nL4:", tables[2])
-  if encodings and tables:
-    print(f"\n{'TypeID':>6} {'Name':>18} {'L2 enc':>12} {'L3 enc':>12} {'L4 enc':>12}"
-          f" {'L2':>4} {'L3':>4} {'L4':>4} {'L2 delta':>12} {'L3 delta':>12} {'L4 delta':>12}")
-    print("-" * 140)
-    for type_id in sorted(set(encodings[0]) | set(encodings[1]) | set(encodings[2])):
-      name = TYPE_NAMES.get(type_id, f'UNK_{type_id}')
-      bits = [tables[i][type_id] if type_id < len(tables[i]) else 0 for i in range(3)]
-      enc_strs = [f"0x{encodings[i][type_id][0]:02x}/0x{encodings[i][type_id][1]:02x}" if type_id in encodings[i] else "-" for i in range(3)]
-      delta_strs = [f"[{d[1]-1}:{d[0]}]" if (d := deltas[i].get(type_id, (0, 0)))[1] > d[0] else "-" for i in range(3)]
-      print(f"{type_id:6d} {name:>18} {enc_strs[0]:>12} {enc_strs[1]:>12} {enc_strs[2]:>12}"
-            f" {bits[0]:4d} {bits[1]:4d} {bits[2]:4d} {delta_strs[0]:>12} {delta_strs[1]:>12} {delta_strs[2]:>12}")
-
-  cdna = extract_cdna_packet_sizes()
-  if cdna: print(f"\nCDNA packet sizes: {cdna}")
-
-  unittest.main()
@@ -24,7 +24,7 @@ def rocprof_inst_traces_match(sqtt, prg, target, pass_rocprof_err=False):

  passed_insts = 0
  for pkt, info in map_insts(sqtt.blob, prg.lib, target):
-    if DEBUG >= 2: print_packets([pkt])
+    if DEBUG >= 2: print_packets([(pkt, info)])
    if info is None: continue
    if DEBUG >= 2: print(f"{' '*29}{disasm(info.inst)}")
    rocprof_inst = next(rwaves_iter[info.wave][0])
@@ -68,7 +68,7 @@ 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_py")
+          pass_rocprof_err = OSX and target == "gfx1200" and name.startswith("profile_ops")
          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")

@@ -47,6 +47,18 @@ def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:i
 def verify_asm_gemm_k_sharded(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=8) -> None:
  run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=1, b_shard=0, gpus=gpus)

+def verify_asm_gemm_n_sharded(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
+  run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=None, b_shard=1, gpus=gpus)
+
+def verify_asm_gemm_m_sharded(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
+  run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=0, b_shard=None, gpus=gpus)
+
+def verify_asm_gemm_n_sharded_2d(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
+  run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=None, b_shard=1, gpus=gpus)
+
+def verify_asm_gemm_k_sharded_3d(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
+  run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=2, b_shard=0, gpus=gpus)
+
 # 128x smaller than usual
 # uses the UOp GEMM, runs on non CDNA4 and CI
@unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
@@ -60,6 +72,14 @@ class TestGemm(unittest.TestCase):
  def test_gemm_multi(self): verify_asm_gemm(2, 64, 32, 32, gpus=2)
  @needs_second_gpu
  def test_gemm_k_sharded(self): verify_asm_gemm_k_sharded(64, 64, 2*64, gpus=2)
+  @needs_second_gpu
+  def test_gemm_m_sharded(self): verify_asm_gemm_m_sharded(2*64, 64, 32, gpus=2)
+  @needs_second_gpu
+  def test_gemm_n_sharded(self): verify_asm_gemm_n_sharded(1, 64, 64, 32, gpus=2)
+  @needs_second_gpu
+  def test_gemm_n_sharded_2d(self): verify_asm_gemm_n_sharded_2d(64, 2*64, 32, gpus=2)
+  @needs_second_gpu
+  def test_gemm_k_sharded_3d(self): verify_asm_gemm_k_sharded_3d(1, 64, 32, 2*64, gpus=2)

 # uses the Asm GEMM on CDNA4 only for speed reasons
 class TestGemmLarge(unittest.TestCase):
@@ -101,6 +121,20 @@ class TestGemmLarge(unittest.TestCase):
      verify_asm_gemm(3, 256, 256, 256)
  def test_gemm_previously_unsupported(self): verify_asm_gemm(8, 1024, 1024, 4096, gpus=8)

+  # M-sharded 2D
+  def test_m_sharded_1(self): verify_asm_gemm_m_sharded(8*8192, 4096, 4096, dtype=dtypes.bfloat16, gpus=8)
+  def test_m_sharded_2(self): verify_asm_gemm_m_sharded(8*4096, 14336, 4096, dtype=dtypes.bfloat16, gpus=8)
+
+  # N-sharded 2D
+  def test_n_sharded_2d_1(self): verify_asm_gemm_n_sharded_2d(8192, 8*4096, 4096, dtype=dtypes.bfloat16, gpus=8)
+  def test_n_sharded_2d_2(self): verify_asm_gemm_n_sharded_2d(4096, 8*14336, 4096, dtype=dtypes.bfloat16, gpus=8)
+
+  # tensor parallel shapes (Llama 8B, MP=8)
+  def test_tp_n_sharded_wq(self): verify_asm_gemm_n_sharded(1, 8192, 4096, 4096, dtype=dtypes.bfloat16, gpus=8)
+  def test_tp_n_sharded_w1(self): verify_asm_gemm_n_sharded(1, 8192, 14336, 4096, dtype=dtypes.bfloat16, gpus=8)
+  def test_tp_k_sharded_wo(self): verify_asm_gemm_k_sharded_3d(1, 8192, 4096, 4096, dtype=dtypes.bfloat16, gpus=8)
+  def test_tp_k_sharded_w2(self): verify_asm_gemm_k_sharded_3d(1, 8192, 4096, 14336, dtype=dtypes.bfloat16, gpus=8)
+
  # more shapes: vary M, N, K independently
  def test_shape_small_square(self): verify_asm_gemm(1, 256, 256, 256)
  def test_shape_small_rect_m(self): verify_asm_gemm(1, 512, 256, 256)
@@ -30,8 +30,7 @@ class TestMovedConstFolding(unittest.TestCase):
  def test_copy_padded_const(self):
    schedule = Tensor.ones(4, device="CPU:0").pad(((1, 1),)).to("CPU:1").schedule()
    assert not any(si.ast.op is Ops.COPY for si in schedule), "const copy should be folded"
-    # TODO: this is wrong, should be [0, 1, 1, 1, 1, 0]
-    np.testing.assert_equal(Tensor.ones(4, device="CPU:0").pad(((1, 1),)).to("CPU:1").numpy(), [1, 1, 1, 1, 1, 1])
+    np.testing.assert_equal(Tensor.ones(4, device="CPU:0").pad(((1, 1),)).to("CPU:1").numpy(), [0, 1, 1, 1, 1, 0])

  def test_cast_padded(self):
    # NOTE: it's always 1 kernel when calling .numpy, limitation of _check_ast_count
@@ -15,7 +15,7 @@ from test.helpers import needs_second_gpu
 np.random.seed(1337)
 Tensor.manual_seed(1337)
 BUF_SIZE = 4096
-RUN_CNT = 4
+RUN_CNT = 5

 cached_prgs = {}
 def helper_exec_op(device, outbuf, inbufs):
@@ -47,6 +47,17 @@ 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))
@@ -80,6 +91,14 @@ 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)]
@@ -110,11 +129,6 @@ 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()

@@ -265,5 +279,58 @@ 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()
@@ -199,6 +199,13 @@ class TestMultiTensor(unittest.TestCase):
      run_schedule(sched)
      np.testing.assert_equal(xt.numpy(), X_np[i*2:i*2+2])

+  def test_cat_on_non_shard_axis(self):
+    # cat must be lowered to PAD/ADD before multi_pm runs, otherwise MULTI nodes are not handled
+    X = Tensor.arange(8).reshape(4, 2).realize().shard_(devices_2, 0)
+    Y = Tensor.arange(8, 16).reshape(4, 2).realize().shard_(devices_2, 0)
+    Z = X.cat(Y, dim=1)
+    np.testing.assert_equal(Z.numpy(), np.concatenate([np.arange(8).reshape(4, 2), np.arange(8, 16).reshape(4, 2)], axis=1))
+
  @given(strat.sampled_from((devices_2, devices_3)),
         strat.sampled_from((Ops.ADD, Ops.MUL, Ops.MAX)),
         strat.sampled_from((None, 0, 1)), strat.sampled_from((None, 0, 1)))
@@ -1665,6 +1665,15 @@ 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)
@@ -795,6 +795,39 @@ class TestSchedule(unittest.TestCase):
      self.assertIsNotNone(out.uop.base.realized)
      self.assertIsInstance(out.uop.base.realized.dtype, ImageDType)

+  @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
+  def test_image_dot_f16_fusion(self):
+    with Context(FLOAT16=1, OPENPILOT_HACKS=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()
+        sched = ((a @ z).relu() + a).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()
+
+      self.assertEqual(cnt1, 5)
+      self.assertEqual(cnt2, 5)
+
+  @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)])
+
+      with Context(IMAGE=1): cnt1 = cnt()
+      with Context(IMAGE=2): cnt2 = cnt()
+
+      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]
    run_schedule(check_schedule(compare:=f(*args), cnt))
@@ -251,8 +251,11 @@ class TestSetitem(unittest.TestCase):
    s1 = t.sum()
    t[3:].assign(2.0)
    s2 = t.sum()
-    # 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])
+    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])

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

    # eager version
    a = Tensor.arange(4, dtype=dtypes.float).contiguous().realize()
@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.device import Buffer
-from tinygrad.helpers import Context
+from tinygrad.helpers import Context, getenv
 from test.helpers import needs_second_gpu

@unittest.skipUnless(hasattr(Device[Device.DEFAULT].allocator, "_offset"), "subbuffer not supported")
@@ -42,7 +42,7 @@ class TestSubBuffer(unittest.TestCase):
    assert out == [102, 103]

  @needs_second_gpu
-  @unittest.skipIf(Device.DEFAULT not in {"CUDA", "NV", "AMD"}, "only NV, AMD, CUDA")
+  @unittest.skipIf(Device.DEFAULT not in {"CUDA", "NV", "AMD"} or getenv("MOCKGPU"), "only NV, AMD, CUDA")
  def test_subbuffer_transfer(self):
    t = Tensor.arange(0, 10, dtype=dtypes.uint8).realize()
    vt = t[2:5].contiguous().realize()
@@ -69,6 +69,58 @@ class TestSymbolicOps(unittest.TestCase):
        # symbolic shape dropout is not supported
        self.test_attention(dropout_p=0.5)

+  def test_sdpa_symbolic_seq_len(self):
+    # symbolic seq_len on all of q/k/v (dim -2 after transpose)
+    q = Tensor.rand(2, 10, 4, 8)
+    k = Tensor.rand(2, 10, 4, 8)
+    v = Tensor.rand(2, 10, 4, 8)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      Tensor.realize(q, k, v)
+      symbolic = q[:, :vi].transpose(1, 2).scaled_dot_product_attention(
+        k[:, :vi].transpose(1, 2), v[:, :vi].transpose(1, 2)).realize()[:2, :4, :i, :8].numpy()
+      expected = q[:, :i].transpose(1, 2).scaled_dot_product_attention(
+        k[:, :i].transpose(1, 2), v[:, :i].transpose(1, 2)).realize().numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_sdpa_symbolic_seq_len_query_only(self):
+    # symbolic seq_len on query only (dim -2 after transpose)
+    q = Tensor.rand(2, 10, 4, 8)
+    k = Tensor.rand(2, 5, 4, 8)
+    v = Tensor.rand(2, 5, 4, 8)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      Tensor.realize(q, k, v)
+      symbolic = q[:, :vi].transpose(1, 2).scaled_dot_product_attention(
+        k.transpose(1, 2), v.transpose(1, 2)).realize()[:2, :4, :i, :8].numpy()
+      expected = q[:, :i].transpose(1, 2).scaled_dot_product_attention(
+        k.transpose(1, 2), v.transpose(1, 2)).realize().numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_sdpa_symbolic_batch(self):
+    # symbolic batch dim (dim 0)
+    q = Tensor.rand(10, 4, 3, 8)
+    k = Tensor.rand(10, 4, 3, 8)
+    v = Tensor.rand(10, 4, 3, 8)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      Tensor.realize(q, k, v)
+      symbolic = q[:vi].scaled_dot_product_attention(k[:vi], v[:vi]).realize()[:i, :4, :3, :8].numpy()
+      expected = q[:i].scaled_dot_product_attention(k[:i], v[:i]).realize().numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
+  def test_sdpa_symbolic_heads(self):
+    # symbolic heads dim (dim -3)
+    q = Tensor.rand(2, 10, 3, 8)
+    k = Tensor.rand(2, 10, 3, 8)
+    v = Tensor.rand(2, 10, 3, 8)
+    for i in range(1, 5):
+      vi = Variable("i", 1, 10).bind(i)
+      Tensor.realize(q, k, v)
+      symbolic = q[:, :vi].scaled_dot_product_attention(k[:, :vi], v[:, :vi]).realize()[:2, :i, :3, :8].numpy()
+      expected = q[:, :i].scaled_dot_product_attention(k[:, :i], v[:, :i]).realize().numpy()
+      np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
+
  def test_attention_pos_0_sz_0(self):
    Attention(128, 8)(Tensor.ones(1, 0, 128), Variable("start_pos", 0, 128).bind(0), None)

@@ -136,6 +136,30 @@ class TestTensorVariable(unittest.TestCase):
    with self.assertRaises(AssertionError):
      t.chunk(2, dim=0)  # can't split along symbolic dim

+  def test_symbolic_var_sum(self, var_name="u"):
+    t = Variable("t", 1, 10).bind(4)
+    v = Variable(var_name, 1, 5).bind(1)
+    mask = (Tensor.full((1, 1, t, v+t), 1) + 1).contiguous()
+    mask.shrink(((0, 1), (0, 1), (0, 4), (0, 4))).numpy()
+  def test_symbolic_var_sum_alt_name(self): self.test_symbolic_var_sum("s")
+
+  def test_symbolic_triu(self):
+    t = Variable("t", 1, 10).bind(4)
+    for start_pos in (0, 1, 3):
+      var_start_pos = Variable("start_pos", 0, 5).bind(start_pos)
+      mask = Tensor.full((1, 1, t, var_start_pos+t), float("-inf")).triu(var_start_pos+1)
+      out = mask.shrink(((0, 1), (0, 1), (0, 4), (0, start_pos+4))).numpy()
+      expected = np.triu(np.full((1, 1, 4, start_pos+4), float("-inf")), k=start_pos+1)
+      np.testing.assert_equal(out, expected)
+
+  def test_symbolic_tril(self):
+    t = Variable("t", 1, 10).bind(4)
+    for start_pos in (0, 1, 3):
+      var_start_pos = Variable("start_pos", 0, 5).bind(start_pos)
+      mask = Tensor.full((1, 1, t, var_start_pos+t), float("-inf")).tril(var_start_pos+1)
+      out = mask.shrink(((0, 1), (0, 1), (0, 4), (0, start_pos+4))).numpy()
+      expected = np.tril(np.full((1, 1, 4, start_pos+4), float("-inf")), k=start_pos+1)
+      np.testing.assert_equal(out, expected)

 if __name__ == '__main__':
  unittest.main()
@@ -0,0 +1,31 @@
+#!/usr/bin/env python3
+"""
+Stress test for beam timeout + device recovery on AM devices.
+
+Usage:
+  AMD=1 python test/external/external_test_beam_timeout_recovery.py
+"""
+from tinygrad import Tensor, Device
+from tinygrad.helpers import Context
+from tinygrad.runtime.ops_amd import AMDDevice
+
+if __name__ == "__main__":
+  dev = Device["AMD"]
+  assert isinstance(dev, AMDDevice) and dev.is_am(), "not am"
+
+  N = 10000
+  for i in range(N):
+    with Context(DEBUG=0, BEAM=0):
+      a = Tensor.rand(4096, 4096, device="AMD").contiguous().realize()
+      b = Tensor.rand(4096, 4096, device="AMD").contiguous().realize()
+      c = a.matmul(b)
+      c.realize()
+    try: dev.synchronize(timeout=1)
+    except RuntimeError as e: print(e)
+    with Context(DEBUG=0, BEAM=0):
+      a = Tensor.ones(512, 512, device="AMD").contiguous().realize()
+      b = Tensor.ones(512, 512, device="AMD").contiguous().realize()
+      result = a.matmul(b).realize()[0, 0].item()
+    assert result == 512.0, f"iter {i}: got {result}"
+    print(f"  iter {i+1}/{N}: ok")
+  print(f"=== All {N} iterations passed ===")
@@ -0,0 +1,55 @@
+import subprocess, sys, os, random
+
+CHILD_SCRIPT = """
+import os, random
+import numpy as np
+from tinygrad import Tensor, Device
+from tinygrad.runtime.ops_amd import AMDDevice
+
+dev = Device["AMD"]
+for i in range({N}):
+  sz = random.randint(1, {MAX_SZ})
+  data = np.random.randint(0, 256, sz, dtype=np.uint8)
+  t = Tensor(data, device="AMD").contiguous().realize()
+  dev.synchronize()
+  result = t.numpy()
+  assert (result == data).all(), f"Data mismatch at iter {{i}}"
+""".strip()
+
+def run_child(n_ops, max_sz, timeout):
+  env = os.environ.copy()
+  env.setdefault("SDMA_RING_SIZE", "4096")
+
+  script = CHILD_SCRIPT.format(N=n_ops, MAX_SZ=max_sz)
+  p = subprocess.Popen([sys.executable, "-c", script], stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env)
+
+  try:
+    _, stderr = p.communicate(timeout=timeout)
+    return ("ok" if p.returncode == 0 else "fail"), stderr.decode(errors='replace')
+  except subprocess.TimeoutExpired:
+    p.kill()
+    p.communicate()
+    return "timeout", "TIMEOUT: SDMA ring likely stuck"
+
+if __name__ == "__main__":
+  n_iters = int(os.environ.get("FUZZ_ITERS", "10000"))
+  timeout = int(os.environ.get("FUZZ_TIMEOUT", "10"))
+  max_sz = int(os.environ.get("FUZZ_MAX_SZ", "65536"))
+
+  timeouts = 0
+  failures = 0
+
+  for i in range(n_iters):
+    # Run child with many ops to stress the small sdma ring buffer across warm starts
+    n_ops = random.randint(20, 100)
+    status, stderr = run_child(n_ops=n_ops, max_sz=max_sz, timeout=timeout)
+    if status == "timeout":
+      timeouts += 1
+      print(f"\tstderr: {stderr[:500]}")
+    elif status == "fail":
+      failures += 1
+      print(f"\tstderr: {stderr[:500]}")
+    else:
+      print(f"iter {i}: ok (n_ops={n_ops})")
+
+  print(f"\n=== Results: {n_iters} iterations, {timeouts} timeouts, {failures} failures ===")
@@ -4,13 +4,19 @@
 These tests intentionally cause GPU faults to verify error handling.
 Run with: AMD=1 python -m pytest test/external/external_test_gpu_crash.py -v
 """
-import unittest, re
+import unittest, re, importlib
 from tinygrad.device import Device
-from tinygrad.runtime.autogen.amd.rdna3.ins import *  # noqa: F403
 from tinygrad.renderer.amd.dsl import s, v, Inst, NULL

-def assemble(code:str, name:str="test") -> str:
-  kd = {"next_free_vgpr": 8, "next_free_sgpr": 8, "wavefront_size32": 1, "user_sgpr_kernarg_segment_ptr": 1, "kernarg_size": 8}
+RDNA3_CDNA3_MAP = {"v_mov_b32_e32": "v_mov_b32_e32", "s_mov_b32": "s_mov_b32", "s_waitcnt": "s_waitcnt", "s_endpgm": "s_endpgm",
+                   "global_load_b32": "global_load_dword", "global_store_b32": "global_store_dword",
+                   "global_atomic_add_u32": "global_atomic_add", "flat_load_b32": "flat_load_dword",
+                   "flat_store_b32": "flat_store_dword", "flat_atomic_add_u32": "flat_atomic_add", "s_load_b32": "s_load_dword"}
+
+def assemble(code:str, name:str="test", is_cdna:bool=False) -> str:
+  kd = {"next_free_vgpr": 8, "next_free_sgpr": 8, "user_sgpr_kernarg_segment_ptr": 1, "kernarg_size": 8}
+  if is_cdna: kd["accum_offset"] = 8
+  else: kd["wavefront_size32"] = 1
  return f".text\n.globl {name}\n.p2align 8\n.type {name},@function\n{name}:\n{code}\n.rodata\n.p2align 6\n.amdhsa_kernel {name}\n" + \
         "\n".join(f".amdhsa_{k} {v}" for k,v in kd.items()) + "\n.end_amdhsa_kernel"

@@ -21,6 +27,10 @@ class TestGPUCrash(unittest.TestCase):
    from tinygrad.runtime.support.compiler_amd import HIPCompiler
    cls.dev = Device["AMD"]
    cls.compiler = HIPCompiler(cls.dev.arch)
+    cls.is_cdna = cls.dev.target[0] < 10
+    ins = importlib.import_module('tinygrad.runtime.autogen.amd.' + ('cdna' if cls.is_cdna else 'rdna3') + '.ins')
+    for rdna3_name, cdna3_name in RDNA3_CDNA3_MAP.items():
+      setattr(cls, rdna3_name, getattr(ins, cdna3_name if cls.is_cdna else rdna3_name))

  def setUp(self):
    # Verify device works before each test
@@ -33,7 +43,7 @@ class TestGPUCrash(unittest.TestCase):

  def _run(self, code: str):
    from tinygrad.runtime.ops_amd import AMDProgram
-    prg = AMDProgram(self.dev, "test", self.compiler.compile(assemble(code)))
+    prg = AMDProgram(self.dev, "test", self.compiler.compile(assemble(code, is_cdna=self.is_cdna)))
    prg(self.dev.allocator.alloc(64), global_size=(1,1,1), local_size=(1,1,1), wait=True)

  def _run_insts(self, insts: list[Inst]):
@@ -57,32 +67,32 @@ class TestOutOfBoundsMemoryAccess(TestGPUCrash):

  def test_global_load_null_ptr(self):
    """Global load from NULL pointer."""
-    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0),
-             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0),
+             self.global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
    self._assert_gpu_fault(lambda: self._run_insts(insts))

  def test_global_store_null_ptr(self):
    """Global store to NULL pointer."""
-    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 0xDEADBEEF),
-             global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 0xDEADBEEF),
+             self.global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
    self._assert_gpu_fault(lambda: self._run_insts(insts))

  def test_global_load_unmapped_high_address(self):
    """Global load from high unmapped address (0xDEAD00000000)."""
-    insts = [v_mov_b32_e32(v[0], 0x00000000), v_mov_b32_e32(v[1], 0xDEAD),
-             global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    insts = [self.v_mov_b32_e32(v[0], 0x00000000), self.v_mov_b32_e32(v[1], 0xDEAD),
+             self.global_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
    self._assert_gpu_fault(lambda: self._run_insts(insts))

  def test_global_store_unmapped_high_address(self):
    """Global store to high unmapped address."""
-    insts = [v_mov_b32_e32(v[0], 0x00000000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 0x12345678),
-             global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    insts = [self.v_mov_b32_e32(v[0], 0x00000000), self.v_mov_b32_e32(v[1], 0xDEAD), self.v_mov_b32_e32(v[2], 0x12345678),
+             self.global_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
    self._assert_gpu_fault(lambda: self._run_insts(insts))

  def test_global_atomic_unmapped(self):
    """Atomic operation on unmapped memory."""
-    insts = [v_mov_b32_e32(v[0], 0xBEEF0000), v_mov_b32_e32(v[1], 0xDEAD), v_mov_b32_e32(v[2], 1),
-             global_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    insts = [self.v_mov_b32_e32(v[0], 0xBEEF0000), self.v_mov_b32_e32(v[1], 0xDEAD), self.v_mov_b32_e32(v[2], 1),
+             self.global_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
    self._assert_gpu_fault(lambda: self._run_insts(insts))


@@ -91,14 +101,14 @@ class TestSMEMFaults(TestGPUCrash):

  def test_smem_load_null(self):
    """SMEM load from NULL base."""
-    insts = [s_mov_b32(s[2], 0), s_mov_b32(s[3], 0),
-             s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(0), s_endpgm()]
+    insts = [self.s_mov_b32(s[2], 0), self.s_mov_b32(s[3], 0),
+             self.s_load_b32(s[4], s[2:3], 0, soffset=NULL), self.s_waitcnt(0), self.s_endpgm()]
    self._assert_gpu_fault(lambda: self._run_insts(insts))

  def test_smem_load_unmapped(self):
    """SMEM load from unmapped address."""
-    insts = [s_mov_b32(s[2], 0xBEEF0000), s_mov_b32(s[3], 0xDEAD),
-             s_load_b32(s[4], s[2:3], 0, soffset=NULL), s_waitcnt(0), s_endpgm()]
+    insts = [self.s_mov_b32(s[2], 0xBEEF0000), self.s_mov_b32(s[3], 0xDEAD),
+             self.s_load_b32(s[4], s[2:3], 0, soffset=NULL), self.s_waitcnt(0), self.s_endpgm()]
    self._assert_gpu_fault(lambda: self._run_insts(insts))


@@ -107,20 +117,20 @@ class TestFlatMemoryFaults(TestGPUCrash):

  def test_flat_load_null(self):
    """FLAT load from NULL address."""
-    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0),
-             flat_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0),
+             self.flat_load_b32(v[2], addr=v[0:1], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
    self._assert_gpu_fault(lambda: self._run_insts(insts))

  def test_flat_store_null(self):
    """FLAT store to NULL address."""
-    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 0xDEADBEEF),
-             flat_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 0xDEADBEEF),
+             self.flat_store_b32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
    self._assert_gpu_fault(lambda: self._run_insts(insts))

  def test_flat_atomic_null(self):
    """FLAT atomic on NULL address."""
-    insts = [v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_mov_b32_e32(v[2], 1),
-             flat_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), s_waitcnt(0), s_endpgm()]
+    insts = [self.v_mov_b32_e32(v[0], 0), self.v_mov_b32_e32(v[1], 0), self.v_mov_b32_e32(v[2], 1),
+             self.flat_atomic_add_u32(addr=v[0:1], data=v[2], saddr=NULL, offset=0), self.s_waitcnt(0), self.s_endpgm()]
    self._assert_gpu_fault(lambda: self._run_insts(insts))


@@ -10,7 +10,7 @@ from tinygrad.helpers import Profiling

 class FakeProgram:
  def __init__(self, name:str, prg:bytes, **kwargs): pass
-  def __call__(self, *bufs, global_size, local_size, vals=(), wait=False): pass
+  def __call__(self, *bufs, global_size, local_size, vals=(), wait=False, **kw): pass

 class FakeAllocator(Allocator[Compiled]):
  def _alloc(self, sz, options): return None
@@ -416,10 +416,10 @@ class Parser:
      case '||' | '|': return left | right
      case '&&' | '&': return left & right
      case '^': return left ^ right
-      case '==' | '<>': return left.eq(right) if op == '==' else left.ne(right)
+      case '==': return left.eq(right)
      case '!=': return left.ne(right)
-      case '>=' | '<=' | '>' | '<':
-        ops = {'>=':(lambda a,b:a>=b),'<=':(lambda a,b:a<=b),'>':(lambda a,b:a>b),'<':(lambda a,b:a<b)}
+      case '>=' | '<=' | '>' | '<' | '<>':
+        ops = {'>=':(lambda a,b:a>=b),'<=':(lambda a,b:a<=b),'>':(lambda a,b:a>b),'<':(lambda a,b:a<b),'<>':(lambda a,b:a.ne(b))}
        return self._cmp_nan(left, right, ops[op])
      case '>>' | '<<': return (left >> right) if op == '>>' else (left << right)
      case '+' | '-':
@@ -87,7 +87,7 @@ class TestHuggingFaceOnnxModels(unittest.TestCase):
      "input_ids": np.random.randint(0, 250002, (1, 11), dtype=np.int64),
      "attention_mask": np.ones((1, 11), dtype=np.int64),
    }
-    self._validate(repo_id, model_file, custom_inputs)
+    self._validate(repo_id, model_file, custom_inputs, atol=1e-3)

 if __name__ == "__main__":
  unittest.main()
@@ -14,6 +14,7 @@ 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
@@ -8,7 +8,7 @@ class TestDataset(unittest.TestCase):
    X_train[0].contiguous().realize()
    GlobalCounters.reset()
    X_train[0].contiguous().realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
+    self.assertLessEqual(GlobalCounters.kernel_count, 1)  # 0 if BUFFER_VIEW (zero-copy), 1 otherwise

 if __name__ == '__main__':
  unittest.main()
@@ -1,10 +1,12 @@
 import gc, unittest
 from tinygrad import Tensor, GlobalCounters, dtypes
+from tinygrad.engine.jit import TinyJit

 class TestMultiRamUsage(unittest.TestCase):
  def setUp(self):
    gc.collect()
    self.baseline = GlobalCounters.mem_used
+    self.baseline_per_device = dict(GlobalCounters.mem_used_per_device)
    self.N = 100
  def assertUsed(self, amt, strict=True):
    gc.collect()
@@ -12,6 +14,11 @@ class TestMultiRamUsage(unittest.TestCase):
    print(f"used {used} bytes")
    if strict: self.assertEqual(used, amt)
    else: self.assertLessEqual(used, amt)
+  def assertDeviceUsed(self, expected:dict[str, int]):
+    gc.collect()
+    for dev, amt in expected.items():
+      used = GlobalCounters.mem_used_per_device[dev] - self.baseline_per_device.get(dev, 0)
+      self.assertEqual(used, amt, f"device {dev}: expected {amt} bytes used, got {used}")

  def test_zeros(self):
    _ = Tensor.zeros(self.N, self.N).contiguous().realize()
@@ -59,6 +66,33 @@ class TestMultiRamUsage(unittest.TestCase):
    X.shard_(devices_4, axis=0).realize()
    self.assertUsed(256 * 4)  # TODO: can be zero

+  def test_zeros_per_device(self):
+    _ = Tensor.zeros(self.N, self.N, device="NULL").contiguous().realize()
+    self.assertDeviceUsed({"NULL": self.N*self.N*4})
+
+  def test_zeros_del_per_device(self):
+    _ = Tensor.zeros(self.N, self.N, device="NULL").contiguous().realize()
+    del _
+    self.assertDeviceUsed({"NULL": 0})
+
+  def test_zeros_copy_per_device(self):
+    devices_2 = ("NULL:1", "NULL:2")
+    _ = Tensor.zeros(self.N, self.N).contiguous().to(devices_2).realize()
+    self.assertDeviceUsed({"NULL:1": self.N*self.N*4, "NULL:2": self.N*self.N*4})
+
+  def test_zeros_shard_per_device(self):
+    devices_2 = ("NULL:1", "NULL:2")
+    _ = Tensor.zeros(self.N, self.N).contiguous().shard(devices_2, axis=0).realize()
+    self.assertDeviceUsed({"NULL:1": self.N*(self.N//2)*4, "NULL:2": self.N*(self.N//2)*4})
+
+  def test_sharded_memory_replicated_per_device(self):
+    devices_4 = tuple(f"NULL:{i+1}" for i in range(4))
+    X = Tensor.ones(256, device="NULL").contiguous().realize()
+    self.assertDeviceUsed({"NULL": 256*4})
+    X.shard_(devices_4).realize()
+    for d in devices_4:
+      self.assertDeviceUsed({d: 256*4})
+
  def _test_matmul_half(self, dev_count:int):
    N = 32
    total_mem = {}
@@ -74,6 +108,38 @@ 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)

+  @unittest.expectedFailure
+  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")
@@ -117,7 +117,7 @@ class TestContiguous(unittest.TestCase):
  def test_size_change_buffer_view(self):
    a = Tensor.empty(4)
    b = a.reshape((1, 1, 4)).shrink(((0, 1), (0, 1), (0, 3))).contiguous()
-    check_schedule(b, 1)
+    check_schedule(b, 0)  # contiguous shrink of a realized buffer is a zero-copy BUFFER_VIEW

  def test_double_contiguous_realizes_once(self):
    a = Tensor.empty(4, 1)
@@ -234,6 +234,18 @@ class TestSchedule(unittest.TestCase):
    d = Tensor.empty(1).assign(c)
    check_schedule(d, 1)

+  def test_detach_assign(self):
+    a = Tensor.ones(4, 4).contiguous().realize()
+    buf1, buf2 = Tensor.empty(4, 4).contiguous(), Tensor.empty(4, 4).contiguous()
+    r = buf2.assign(buf1.assign(a + 1.0) * 2.0)
+    check_schedule(r.detach().contiguous(), 2)
+
+  def test_contiguous_backward_assign(self):
+    a = Tensor.ones(4, 4).contiguous().realize()
+    buf1, buf2 = Tensor.empty(4, 4).contiguous(), Tensor.empty(4, 4).contiguous()
+    r = buf2.assign(buf1.assign(a + 1.0) * 2.0)
+    check_schedule(r.contiguous_backward().contiguous(), 2)
+
  def test_mulacc_relu_fusion(self):
    a = Tensor.empty(10)
    b = Tensor.empty(10)
@@ -568,21 +580,22 @@ class TestSchedule(unittest.TestCase):

  # this is the failing case in openpilot...it's very simple like this
  def test_image_conv_fusion(self):
-    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)
+    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)

-    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)
@@ -1158,5 +1171,23 @@ class TestFusionOp(unittest.TestCase):
    self.assertEqual(len(sched), 1)
    self.assertLess(time.perf_counter()-st, 2.0)

+# NOTE: the NULL backend supports BUFFER_VIEW
+class TestBufferView(unittest.TestCase):
+  def test_shrink_contiguous_is_buffer_view(self):
+    # simple 1D shrink of a realized buffer should be BUFFER_VIEW, not a copy kernel
+    a = Tensor.arange(100).contiguous().realize()
+    b = a.shrink(((10, 50),)).contiguous()
+    run_schedule(check_schedule(b, 0))
+
+  def test_shrink_2d_contiguous_is_buffer_view(self):
+    a = Tensor.arange(100).reshape(10,10).contiguous().realize()
+    b = a.shrink(((1, 5),None)).contiguous()
+    run_schedule(check_schedule(b, 0))
+
+  def test_chained_shrink_is_buffer_view(self):
+    a = Tensor.arange(1000).contiguous().realize()
+    b = a.shrink(((200, 800),)).shrink(((0, 300),)).reshape((30, 10)).shrink(((20, 25), (0, 10))).contiguous()
+    run_schedule(check_schedule(b, 0))
+
 if __name__ == '__main__':
  unittest.main(verbosity=2)
@@ -390,16 +390,16 @@ class TestImageSimplification(unittest.TestCase):

    # 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%8)*32)+(idx0//32))+8)%64)", "((idx0%8)//2)")
+    self.check(load, "(idx0<256)", "(((((idx0%8)*32)+(idx0//32))+8)%64)", "(idx0//2%4)")

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

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

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

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

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

  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 < 2 * tqdm_time
+    assert tinytqdm_time < 5 * 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.CAT)
+    self.assertEqual(result.op, Ops.VCAT)
    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*5)//2)+(b*2))")
+    self.helper_test_variable(usum([Variable("a", 0, 7)*5, Variable("b", 0, 3)*4]) // 2, 0, 23, "((a*2)+(b*2)+(a//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*3)//2)+1)")
+    self.helper_test_variable((Variable("a", 0, 7)*30+20)//20, 1, 11, "((a+(a//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*5)+(b*5))//2)")
+    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))")

  def test_mod_min_max(self):
    self.helper_test_variable(Variable("x", 0, 10)%Variable("y", 1, 10), 0, 9, "(x%y)")
@@ -297,6 +297,9 @@ 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")

@@ -545,6 +548,13 @@ 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)")
@@ -588,8 +598,7 @@ class TestSymbolic(unittest.TestCase):
    gidx0 = Variable("gidx0", 0, 2)
    lidx2 = Variable("lidx2", 0, 12)
    lidx3 = Variable("lidx3", 0, 12)
-    # 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)")
+    self.helper_test_variable((gidx0*3+lidx2*19+lidx3*38)//(3*19), 0, 12, "((lidx2+(lidx3*2))//3)")

  def test_sum_mul_distribute(self):
    gidx0 = Variable("gidx0", 0, 7)
@@ -606,6 +615,12 @@ 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)
@@ -655,8 +670,7 @@ 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)")
-    with self.assertRaises(AssertionError):
-      self.helper_test_variable((31 * b + 1) % 18 + ((31 * b + 1) // 18) * 18, 1, 3101, "((b*31)+1)")
+    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)
@@ -680,8 +694,21 @@ 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)")
-    with self.assertRaises(AssertionError):
-      self.helper_test_variable((30 * b + 1) % 18 + ((30 * b + 1) // 18) * 18, 1, 3001, "((b*30)+1)")
+    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_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_gated_load(self):
    idx = Variable("idx", 0, 24)
@@ -920,6 +947,17 @@ 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):
@@ -1,4 +1,4 @@
-import unittest, decimal, json, struct, sys
+import unittest, decimal, sys, json
 from dataclasses import dataclass
 from typing import Generator

@@ -357,41 +357,9 @@ class TestVizIntegration(BaseTestViz):

 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry
 from tinygrad.viz.serve import get_profile
+from extra.viz.cli import decode_profile

-class TinyUnpacker:
-  def __init__(self, buf): self.buf, self.offset = buf, 0
-  def __call__(self, fmt:str) -> tuple:
-    ret = struct.unpack_from(fmt, self.buf, self.offset)
-    self.offset += struct.calcsize(fmt)
-    return ret
-
-# 0 means None, otherwise it's an enum value
-def option(i:int) -> int|None: return None if i == 0 else i-1
-
-def load_profile(lst:list[ProfileEvent]) -> dict:
-  ret = get_profile(lst)
-  u = TinyUnpacker(ret)
-  total_dur, global_peak, index_len, layout_len = u("<IQII")
-  strings, dtypes, markers = json.loads(ret[u.offset:u.offset+index_len]).values()
-  u.offset += index_len
-  layout:dict[str, dict] = {}
-  for _ in range(layout_len):
-    klen = u("<B")[0]
-    k = ret[u.offset:u.offset+klen].decode()
-    u.offset += klen
-    layout[k] = v = {"events":[]}
-    event_type, event_count = u("<BI")
-    if event_type == 0:
-      for _ in range(event_count):
-        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["peak"] = u("<Q")[0]
-      for _ in range(event_count):
-        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}
+def load_profile(lst:list[ProfileEvent]) -> dict: return decode_profile(get_profile(lst))

 class TestVizProfiler(BaseTestViz):
  def test_transfer_uses_copy_device(self):
@@ -443,7 +411,7 @@ class TestVizProfiler(BaseTestViz):
    j = load_profile(prof)
    event = j['layout']['NV:SDMA:0']['events'][0]
    gbs = sz/(dur*1e-6)*1e-9
-    self.assertEqual(event['fmt'], f"{gbs:.0f} GB/s")
+    self.assertTrue(event['fmt'].startswith(f"{gbs:.0f} GB/s"))

  def test_graph(self):
    prof = [ProfileDeviceEvent(device='NV', tdiff=decimal.Decimal(-1000)),
@@ -485,7 +453,9 @@ class TestVizProfiler(BaseTestViz):
    j = load_profile(prof)
    sdma_events = j['layout']['NV:1:SDMA:0']['events']
    gbs = sz/(dur*1e-6)*1e-9
-    self.assertEqual(sdma_events[0]['fmt'], f"{gbs:.0f} GB/s")
+    timing_txt, trace_txt = sdma_events[0]['fmt'].split("\nTB:")
+    self.assertEqual(timing_txt, f"{gbs:.0f} GB/s")
+    self.assertEqual(json.loads(trace_txt)[0][0], __file__)

  def test_block_ordering(self):
    prof = [ProfileDeviceEvent(device='NV', tdiff=decimal.Decimal(-1000)),
@@ -499,7 +499,11 @@ 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()
-    np.testing.assert_equal(a.numpy(), [0]*8)  # TODO: should be [57, 48, 0, 0, 0, 0, 0, 0] (little-endian 12345)
+    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)

  @unittest.skip("don't use output buffer, and mismatch dtype no longer supported")
  def test_cast_assignment(self):
@@ -687,16 +691,20 @@ class TestAssignOrdering(unittest.TestCase):
    """Swap two non-overlapping slices - requires reading both before writing."""
    # without .realize() on temps: values not captured before overwriting
    buf = Tensor([1, 2, 3, 4, 5, 6, 7, 8]).contiguous().realize()
-    left = buf[0:4].contiguous()  # lazy - not captured yet
-    right = buf[4:8].contiguous()  # lazy - not captured yet
+    left = buf[0:4].clone()  # lazy - not captured yet
+    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!
-    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]
+    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])

    # with .realize() on temps: values captured before writes
    buf = Tensor([1, 2, 3, 4, 5, 6, 7, 8]).contiguous().realize()
-    left = buf[0:4].contiguous().realize()
-    right = buf[4:8].contiguous().realize()
+    left = buf[0:4].clone().realize()
+    right = buf[4:8].clone().realize()
    buf[0:4].assign(right).realize()
    buf[4:8].assign(left).realize()
    np.testing.assert_equal(buf.numpy(), [5, 6, 7, 8, 1, 2, 3, 4])
@@ -809,40 +817,55 @@ 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())
-    # TODO: should be [[0,1],[0,1]]
-    self.assertEqual(c.tolist(), [[0,0],[0,0]])
+    try:
+      self.assertEqual(c.tolist(), [[0,1],[0,1]])
+    except AssertionError:
+      # TODO: broken now
+      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())
-    # TODO: should be [[1,1],[2,1]]
-    self.assertEqual(c.tolist(), [[1,3],[2,4]])
+    try:
+      self.assertEqual(c.tolist(), [[1,1],[2,1]])
+    except AssertionError:
+      # TODO: broken now
+      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())
-    # TODO: should be [[1,1],[3,1]]
-    self.assertEqual(cb.tolist(), [[1,2],[3,4]])
+    try:
+      self.assertEqual(cb.tolist(), [[1,1],[3,1]])
+    except AssertionError:
+      # TODO: broken now
+      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())
-    # TODO: should be [[0,1],[0,1]]
-    self.assertEqual(d.tolist(), [[0,0],[0,0]])
+    try:
+      self.assertEqual(d.tolist(), [[0,1],[0,1]])
+    except AssertionError:
+      # TODO: broken now
+      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())
-    # TODO: should be [[1,1],[2,1]]
-    self.assertEqual(d.tolist(), [[1,3],[2,4]])
+    try:
+      self.assertEqual(d.tolist(), [[1,1],[2,1]])
+    except AssertionError:
+      # TODO: broken now
+      self.assertEqual(d.tolist(), [[1,3],[2,4]])

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

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

 class TestCall(unittest.TestCase):
  def test_call_plus(self):
@@ -100,5 +100,142 @@ 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()
+    a = a.reshape(4,4,2).assign(Tensor.empty(4,4,2)).reshape(8,4)
+    @function(precompile=True)
+    def s(x): return x.sum(axis=0)
+    (s(a)*3).realize()
+
+  def test_call_precompiled(self):
+    a = Tensor.empty(4, 8)
+    @function(precompile=True)
+    def s(x): return x*2
+    (s(a)*3).realize()
+
+  def test_double_call(self):
+    a = Tensor.empty(4, 8)
+    @function(precompile=True)
+    def s(x): return x*2
+    s(s(a)).realize()
+
+  def test_double_call_contiguous(self):
+    a = Tensor.empty(4, 8)
+    @function(precompile=True)
+    def s(x): return x*2
+    s(s(a).contiguous()).realize()
+
+  def test_call_double_gemm(self):
+    a = Tensor.randn(4, 8, requires_grad=True)
+    b = Tensor.randn(8, 12, requires_grad=True)
+    c = Tensor.randn(12, 16, requires_grad=True)
+    ref = Tensor.randn(4, 16)
+    Tensor.realize(a,b,c,ref)
+    @function(precompile=True)
+    def gemm(a:Tensor, b:Tensor, c:Tensor) -> Tensor: return (a@b)@c
+    out = gemm(a,b,c)
+    (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,22 +269,16 @@ 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 - should read elements at indices 0, 2, 4
+    # test non-contiguous (strided) read raises
    dt = Tensor([0, 1, 2, 3, 4, 5]).to(f"disk:{self.tmp('dt_strided_read')}")
-    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!
+    with self.assertRaisesRegex(RuntimeError, "non-contiguous view is not supported"):
+      dt[::2].tolist()

  def test_permuted_read(self):
-    # test non-contiguous (permuted) read - should read transposed
+    # test non-contiguous (permuted) read raises
    dt = Tensor([[0, 1, 2], [3, 4, 5]]).to(f"disk:{self.tmp('dt_permuted_read')}")
-    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!
+    with self.assertRaisesRegex(RuntimeError, "non-contiguous view is not supported"):
+      dt.T.tolist()

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

  def test_strided_setitem(self):
-    # test non-contiguous (strided) setitem - should set elements at indices 0, 2, 4
+    # test non-contiguous (strided) setitem raises
    dt = Tensor([1, 2, 3, 4, 5, 6]).to(f"disk:{self.tmp('dt_strided_setitem')}")
-    with self.assertRaises(RuntimeError):
+    with self.assertRaisesRegex(RuntimeError, "non-contiguous view is not supported"):
      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')}")
@@ -70,6 +70,14 @@ class TestFunction(unittest.TestCase):
    b = Tensor([4,5,6])
    np.testing.assert_equal(f(a, b).numpy(), [5,7,9])

+  def test_contiguous_backward(self):
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor: return (a + b).contiguous_backward()
+
+    a = Tensor([1,2,3])
+    b = Tensor([4,5,6])
+    np.testing.assert_equal(f(a, b).numpy(), [5,7,9])
+
  def test_method(self):
    class Foo:
      def __init__(self): self.w = Tensor([10,20,30])
@@ -32,6 +32,7 @@ class TestGGUF(unittest.TestCase):
  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):
    MXFP4 = 39
@@ -130,7 +131,7 @@ 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 == GGMLQuantizationType.Q4_K: q_data[:, :4] = scales[:, :4]                  # d, dmin 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.Q6_K: q_data[:, -2:] = scales[:, :2]                 # d at end
    q_data = q_data.flatten()
    ref = dequantize(q_data, qtype).reshape(rows, cols)
@@ -151,9 +152,11 @@ 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)

 if __name__ == '__main__':
@@ -83,6 +83,15 @@ class TestLinAlg(unittest.TestCase):
      s_diag = (S.unsqueeze(-2) * Tensor.eye(2))
      reconstruction_helper([U, s_diag, V], a)

+  def test_svd_identity_4x4(self):
+    a = Tensor.eye(4)
+    U,S,V = a.svd()
+    assert not np.isnan(U.numpy()).any()
+    assert not np.isnan(S.numpy()).any()
+    assert not np.isnan(V.numpy()).any()
+    s_diag = (S.unsqueeze(-2) * Tensor.eye(4))
+    reconstruction_helper([U, s_diag, V], a)
+
  def test_svd_rank1(self):
    a = Tensor([[1.0, 1.0], [2.0, 2.0]]).realize()
    U, S, V = a.svd()
@@ -1,29 +1,87 @@
 import unittest
 from unittest.mock import patch
-from tinygrad import Tensor
+from tinygrad import Tensor, UOp
+from tinygrad.engine.schedule import schedule_cache

 class TestTransformerGenerate(unittest.TestCase):
-  def test_start_pos_parameter_is_used(self):
-    """Test that start_pos parameter is not ignored (regression test for always resetting to 0)."""
+  def test_kv_cache_reuse(self):
+    """Test that generate reuses the KV cache when tokens extend the cached prefix."""
    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.bind_val))
-      return Tensor([[42]])  # return a fake next token
+      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: prefill 5 tokens + 1 decode
      tokens = [1, 2, 3, 4, 5]
-      gen = model.generate(tokens, start_pos=3)
-      next(gen)  # get first token
+      gen = model.generate(tokens)
+      next(gen)  # prefill
+      next(gen)  # decode

-    # 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
+      # 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)")

 if __name__ == '__main__':
  unittest.main()
@@ -0,0 +1,34 @@
+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,7 +1,8 @@
 from __future__ import annotations
-import sys, argparse, typing, re, unicodedata, json, uuid, time, functools
+import sys, argparse, typing, re, unicodedata, json, uuid, time, functools, itertools
 from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function
-from tinygrad.helpers import partition, DEBUG, Timing, GlobalCounters, stderr_log, colored
+from tinygrad.uop.ops import resolve
+from tinygrad.helpers import partition, DEBUG, Timing, GlobalCounters, stderr_log, colored, Context
 from tinygrad.viz.serve import TCPServerWithReuse, HTTPRequestHandler

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

-  @function
+  @function(precompile=bool(getenv("PRECOMPILE", 0)))
  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)
@@ -143,13 +144,14 @@ class TransformerBlock:
    #v = self.cache_kv[1, :, :, 0:start_pos+T, :]

    # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
-    mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device).triu(int(start_pos)+1) if T > 1 else None
+    # 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
    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)
    return x + attn

-  @function
+  @function(precompile=bool(getenv("PRECOMPILE", 0)))
  def _feed_forward(self, h: Tensor) -> Tensor:
    h_norm = self.ffn_norm(h)
    if hasattr(self, 'ffn_gate_exps'):
@@ -164,7 +166,8 @@ 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?
-      self.cache_kv = Tensor.zeros(2, x.shape[0], self.n_kv_heads, self.max_context, self.head_dim, device=x.device).contiguous().realize()
+      # 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()
    return self._feed_forward(self._attention(x, start_pos)).contiguous()

 class Transformer:
@@ -176,8 +179,10 @@ class Transformer:
    self.output_norm = nn.RMSNorm(dim, norm_eps)
    self.output = nn.Linear(dim, vocab_size, bias=False)
    self.max_context = max_context
-    # 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)
+    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)

  def forward(self, tokens:Tensor, start_pos:int|UOp) -> Tensor:
    x = self.token_embd(tokens)                           # (B, T, D)
@@ -186,12 +191,12 @@ 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.forward_jit if getenv("JIT", 1) and tokens.shape[1] == 1 and isinstance(start_pos, UOp) else self.forward)(tokens, start_pos)
+    return (self.prefill_jit if resolve(tokens.shape[1] != 1) else self.rollout_jit)(tokens, start_pos)

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

    # all state items should be float16, not float32
    state_dict = {k:v.cast('float16') if getenv("HALF", 1) else v for k,v in state_dict.items()}
@@ -224,15 +229,26 @@ class Transformer:
      Tensor.realize(*params)
    return model, kv

-  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")
+  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 = None
    while len(tokens) < self.max_context:
-      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
+      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 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]

 models = {
  "llama3.2:1b": "https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q6_K.gguf",
@@ -261,9 +277,9 @@ 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"></textarea>
+<textarea id="input" rows="1" placeholder="Ask anything" autofocus></textarea>
 <script>
-  input.onkeydown = (e) => { if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); send() } }
+  input.onkeydown = (e) => { if (e.key === 'Enter' && !e.shiftKey && !e.isComposing) { e.preventDefault(); send() } }
  const msgs = [];
  async function send() {
    if (!input.value.trim()) return;
@@ -289,20 +305,22 @@ 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):
-    stderr_log(f"{self.path}  {colored('--', 'BLACK')}  in:{len(ids):5d}  {colored('--', 'BLACK')}  ")
+    cache_start_pos = model.get_start_pos(ids)
+    stderr_log(f"{self.path}  {colored('--', 'BLACK')}  "
+               f"in:{colored(f'{cache_start_pos:5d}', 'green')} +{len(ids)-cache_start_pos:5d}  {colored('--', 'BLACK')}  ")
    tmpl = {"id":f"chatcmpl-{uuid.uuid4().hex[:24]}", "object":"chat.completion.chunk", "created":int(time.time()), "model":model_name}
    yield {"choices": [{"index":0, "delta":{"role":"assistant","content":""}, "finish_reason":None}], **tmpl}
    out: list[int] = []
    st = time.perf_counter()
    for next_id in model.generate(ids):
-      if len(out) == 0: stderr_log(f"prefill:{len(ids)/((pt:=time.perf_counter())-st):4.0f} tok/s  {colored('--', 'BLACK')}  ")
+      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 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"out:{len(out):5d}  {colored('--', 'BLACK')}  gen: {len(out)/(time.perf_counter()-pt):4.0f} tok/s\n")
+    stderr_log(f"gen:{len(out)/(time.perf_counter()-pt):4.0f} tok/s  {colored('--', 'BLACK')}  out:{len(out):5d}\n")

  def do_POST(self):
    raw_body = self.rfile.read(int(self.headers.get("Content-Length", "0")))
@@ -354,31 +372,36 @@ if __name__ == "__main__":
  import gc
  gc.collect()

-  # do benchmark
-  if args.benchmark:
-    gen = model.generate([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,"
-                  f" {GlobalCounters.global_mem//1000000}/{GlobalCounters.mem_used//1000000} MB"): next(gen)
-    exit(0)
-
  # extract some metadata
  tok = SimpleTokenizer.from_gguf_kv(kv)
  bos_id: int|None = kv.get('tokenizer.ggml.bos_token_id') if kv.get('tokenizer.ggml.add_bos_token', True) else None
  eos_id: int = kv['tokenizer.ggml.eos_token_id']

-  # start server
-  if args.serve: TCPServerWithReuse(('', args.serve), Handler).serve_forever()
+  # do benchmark
+  if args.benchmark:
+    gen = model.generate(toks:=[bos_id or 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,"
+                  f" {GlobalCounters.global_mem//1000000}/{GlobalCounters.mem_used//1000000} MB  --  "+\
+                  tok.decode(toks).replace("\n", "\\n")): next(gen)
+    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()
+
+  # 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, start_pos):
+    for next_id in model.generate(ids):
      sys.stdout.write(tok.decode([next_id]) if next_id != eos_id else "\n\n")
      sys.stdout.flush()
      if next_id == eos_id: break
@@ -41,15 +41,15 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
    # split ranges
    sink = graph_rewrite(sink, pm_split_ranges+pm_flatten_range, ctx={}, name="split ranges")

+    # create image buffers
+    if IMAGE == 1 and ren.device in {"QCOM", "CL"}: sink = graph_rewrite(sink, pm_make_images, name="create image buffers", bottom_up=True)
+
    # symbolic (NOTE: this is a requirement for pm_simplify_ranges to be correct)
    sink = graph_rewrite(sink, sym+pm_flatten_range, name="initial symbolic")

    # optimize (schedule) the AST
    sink = graph_rewrite(sink, pm_simplify_ranges, name="simplify ranges")

-    # create image buffers
-    if IMAGE == 1 and ren.device in {"QCOM", "CL"}: sink = graph_rewrite(sink, pm_make_images, name="create image buffers", bottom_up=True)
-
    # do postrange optimization, BEAM or hand_coded_optimizations
    sink = apply_opts(sink, ren)

@@ -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 all_int, dedup, get_contraction
+from tinygrad.helpers import 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 dims[i] * dims[i+1] <= m:
+      if i < (len(dims)-1) and _dim_max(dims[i]) * _dim_max(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.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))),
+   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))),
  # 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.CAT, ls.dtype, tuple(ret)) if ls.op is Ops.LOAD else UOp.group(*ret)
+  return UOp(Ops.VCAT, 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,7 +189,10 @@ 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:
-               (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)))
+                # 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.simplify().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.CAT, src.dtype.scalar().vec(expand_sz*src.dtype.count), (src,)*expand_sz))
+        new_srcs.append(UOp(Ops.VCAT, src.dtype.scalar().vec(expand_sz*src.dtype.count), (src,)*expand_sz))
      else:
        # repeat the arg
        new_srcs.append(src.broadcast(expand_sz))
@@ -36,7 +36,8 @@ def get_test_global_size(global_size, max_global_size, var_vals):
  return test_global_size, input_size / prod(test_global_size)

 def _time_program(p:ProgramSpec, lib:bytes, var_vals:dict[str, int], rawbufs:list[Buffer], early_stop:float|None=None,
-                  allow_test_size:int=True, max_global_size:int|None=65536, clear_l2=False, cnt=3, name="test") -> list[float]:
+                  allow_test_size:int=True, max_global_size:int|None=65536, clear_l2=False, cnt=3, name="test", dev_timeout=False) -> list[float]:
+  timeout = int(early_stop * 1e3) if dev_timeout and early_stop is not None and early_stop < math.inf else None
  factor = 1
  if allow_test_size and max_global_size is not None:
    global_size, factor = get_test_global_size(p.global_size, max_global_size, var_vals)
@@ -50,7 +51,7 @@ def _time_program(p:ProgramSpec, lib:bytes, var_vals:dict[str, int], rawbufs:lis
      if hasattr(dev:=Device[p.device], 'invalidate_caches'): dev.invalidate_caches()
      else:
        with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024,1024).contiguous().realize(do_update_stats=False)
-    tms.append(unwrap(car(input_bufs, var_vals, wait=True))*factor)
+    tms.append(unwrap(car(input_bufs, var_vals, wait=True, timeout=timeout))*factor)
    if early_stop is not None and early_stop < min(tms): break
  return tms

@@ -161,7 +162,8 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True
          continue
        seen_libs.add(lib)
        try: tms = _time_program(p, lib, var_vals, rawbufs, early_stop=beam[0][1]*3 if len(beam) else 1.0,
-                                 allow_test_size=allow_test_size, clear_l2=hasattr(dev, 'invalidate_caches'))
+                                 allow_test_size=allow_test_size, clear_l2=hasattr(dev, 'invalidate_caches'),
+                                 dev_timeout=getenv("BEAM_DEV_TIMEOUT", 1))
        except Exception as e:
          if BEAM_DEBUG: print(f"BEAM failed for opts: {candidates[i].applied_opts}\n{e}")
          if isinstance(e, RuntimeError): continue
@@ -141,6 +141,7 @@ class Buffer:
      self._buf = opaque if opaque is not None else self.allocator.alloc(self.nbytes, self.options)
      if not self.device.startswith("DISK") and (self.options is None or self.options.external_ptr is None):
        GlobalCounters.mem_used += self.nbytes
+        GlobalCounters.mem_used_per_device[self.device] += self.nbytes
      if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", self.trace_num, {"dtype":self.dtype, "sz":self.size}))
    return self
  def deallocate(self):
@@ -149,6 +150,7 @@ class Buffer:
    if self._base is None:
      if GlobalCounters is not None and not self.device.startswith("DISK") and (self.options is None or self.options.external_ptr is None):
        GlobalCounters.mem_used -= self.nbytes
+        GlobalCounters.mem_used_per_device[self.device] -= self.nbytes
      if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", self.trace_num))
      self.allocator.free(self._buf, self.nbytes, self.options)
    elif self._base is not None: self._base.allocated_views -= 1
@@ -1,7 +1,7 @@
 from dataclasses import dataclass, field
-from tinygrad.uop.ops import UOp, UPat, PatternMatcher, Ops, GroupOp, graph_rewrite, identity_element, track_rewrites
-from tinygrad.dtype import ImageDType
-from tinygrad.helpers import prod, DEBUG, argsort, VIZ, pluralize
+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

@dataclass
 class AllocCtx:
@@ -40,11 +40,7 @@ add_tags = PatternMatcher([
  (UPat(GroupOp.All, name="x"), lambda ctx,x: tag_uop(ctx,x) if x in ctx.bases else None),
 ])

-def replace_contig_with_assign(u:UOp):
-  # 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
-  if isinstance(u._device, str) and u._device.startswith(("DISK", "TINYFS")): return u.rtag(None)
+def _buffer_like(u:UOp) -> UOp:
  dtype = u.dtype
  if isinstance(dtype, ImageDType):
    if prod(dtype.shape) != prod(u.max_shard_shape) or ([x for x in u.max_shard_shape if x != 1] or [1])[-1] % 4 != 0:
@@ -52,7 +48,16 @@ def replace_contig_with_assign(u:UOp):
      dtype = dtype.base
  buffer = UOp.new_buffer(u.device, u.shard_size, dtype).reshape(u.max_shard_shape)
  if isinstance(u.device, tuple) and u.axis is not None: buffer = buffer.multi(u.axis)
-  return buffer.assign(u.src[0]).rtag(u.tag)
+  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
+  if isinstance(u._device, str) and u._device.startswith(("DISK", "TINYFS")): return u.rtag(None)
+  return _buffer_like(u).assign(u.src[0]).rtag(u.tag)

 def replace_assign_with_contig(u:UOp):
  assigned_to = u
@@ -60,38 +65,61 @@ 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):
-  x = src
-  while x is not src.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[src.base] = 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]
+  buf = src.base
+  if buf.op not in {Ops.BUFFER, Ops.BUFFER_VIEW}: return None
+  if src.op is Ops.RESHAPE and src.src[0].op in {Ops.BUFFER, Ops.BUFFER_VIEW}: return None

+  # no symbolic shape
+  if not all(isinstance(x, int) for x in c.shape): return None
+
+  # check if view is supported
+  if not isinstance(c.device, str): return None
+  from tinygrad.device import Device
+  if not hasattr(Device[c.device].allocator, "_offset"): return None
+
+  # see if this can be a view
+  offset = src.contiguous_view_offset()
+  if offset is None: return None
+
+  # merge BUFFER_VIEWs
+  if buf.op is Ops.BUFFER_VIEW: offset, buf = offset + buf.arg[1], buf.src[0]
+
+  # NOTE: this contiguous is removed because this BUFFER_VIEW/RESHAPE has_buffer_identity
+  return UOp(Ops.BUFFER_VIEW, src.dtype, (buf,), (src.size, offset)).reshape(src.shape).contiguous(tag=c.tag)
+
+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
+
+# NOTE: adding rules to here is bad. these all need to run before the schedule cache
 pm_early_transform_tensor_graph = PatternMatcher([
-  # CONTIGUOUS replacement hack for openpilot
-  (UPat(Ops.CONTIGUOUS, src=(UPat(GroupOp.Movement, 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),
+  # transform precompiled CALLs
+  (UPat(Ops.CALL, name="c"), transform_precompiled_call),
+
+  # 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),
+
  # add CONTIGUOUS to tagged UOps
  (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+c.tag) if a.src[0].has_buffer_identity() else None),
+   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),
  # replace CONTIGUOUS with ASSIGNs
  (UPat(Ops.CONTIGUOUS, name="u"), replace_contig_with_assign),
-  # remove DETACH/CONTIGUOUS_BACKWARD
+  # remove DETACH/CONTIGUOUS_BACKWARD (allows more contiguous removal)
  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),
-  # reduce of size 0 is the identity element
-  (UPat(Ops.REDUCE_AXIS, name="reduce", src=(UPat.var("x"),)),
-   lambda reduce,x: reduce.const_like(identity_element(reduce.arg[0], reduce.dtype)) if x.size == 0 and reduce.size != 0 else None),
-  # handle size 0
-  (UPat(GroupOp.All-{Ops.SINK}, name="x"), lambda x: x.const_like(0).rtag(x.tag) if x._shape is not None and x.size == 0 else None),
-  # early fixup const copy (TODO: is this wrong if there's a pad?)
-  (UPat(Ops.COPY, src=(UPat.var("s"), UPat()), name="c"), lambda c,s: c.const_like(ss.arg) if (ss:=s.base).op is Ops.CONST else None),
 ])

 def untag_and_append(ctx:AllocCtx, x:UOp):
@@ -117,13 +145,15 @@ pm_finalize_call = PatternMatcher([
  (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),
-  # replace UNIQUE with LUNIQUE for CONST cache key normalization
+  # remove unique from const. TODO: this is copied in function.py
  (UPat(Ops.CONST, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE, name="d")), name="b"), lambda b,d: b.replace(src=(d,))),
 ])

 pm_replace_buf = PatternMatcher([
  # replace BUFFER with PARAM for cache key normalization
  (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="b"), replace_input_buffer),
+  # replace BUFFER_VIEW with PARAM. this rewrite is bottom up so BUFFERs we don't need won't be in the input
+  (UPat(Ops.BUFFER_VIEW, src=(UPat(Ops.BUFFER),), name="b"), replace_input_buffer),
  # strip value from BIND for cache key normalization, so different values hit same cache
  (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR), UPat(Ops.CONST)), name="b"), replace_input_buffer),
 ])
@@ -141,10 +171,10 @@ 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, ctx={}, name="early transform tensor graph")
+  big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, 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")
-  ret = graph_rewrite(UOp.sink(*ctx.assigns), pm_replace_buf, ctx=ctx, name="replace bufs").call(*ctx.replacements)
+  ret = graph_rewrite(UOp.sink(*ctx.assigns), pm_replace_buf, ctx=ctx, bottom_up=True, name="replace bufs").call(*ctx.replacements)
  if VIZ: graph_rewrite(ret, PatternMatcher([]), name="View Call")
  return ret, ctx.buffer_map
@@ -50,7 +50,7 @@ class CompiledRunner(Runner):

  def __reduce__(self): return self.__class__, (self.p,)

-  def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int]|None=None, wait=False) -> float|None:
+  def __call__(self, rawbufs:list[Buffer], var_vals:dict[str, int]|None=None, wait=False, timeout:int|None=None) -> float|None:
    if var_vals is None: var_vals = {}
    global_size, local_size = self.p.launch_dims(var_vals)
    if Device[self.p.device].renderer.has_local and local_size is None and all_int(self.p.global_size):
@@ -58,7 +58,7 @@ class CompiledRunner(Runner):
      global_size = [g//l if g%l == 0 else g/l for g,l in zip(global_size, local_size)]
      self.p = replace(self.p, global_size=global_size, local_size=local_size)
    return self._prg(*[x._buf for x in rawbufs], global_size=tuple(global_size), local_size=tuple(local_size) if local_size else None,
-                     vals=tuple(var_vals[k.expr] if k.expr not in self.p.runtimevars else None for k in self.p.vars), wait=wait)
+                     vals=tuple(var_vals[k.expr] if k.expr not in self.p.runtimevars else None for k in self.p.vars), wait=wait, timeout=timeout)

 class ViewOp(Runner):
  def __init__(self, buf:Buffer): super().__init__(colored(f"view {buf.nbytes:8d} @ {buf.offset:<10d}", "yellow"), buf.device)
@@ -2,10 +2,9 @@ import time, inspect
 from typing import cast
 from collections import deque
 from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites, graph_rewrite, gate_kernel_sink, KernelInfo
-from tinygrad.uop.ops import _remove_all_tags
 from tinygrad.uop.spec import type_verify, tensor_spec
 from tinygrad.device import Buffer, MultiBuffer
-from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR
+from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR, flatten
 from tinygrad.engine.realize import ExecItem

 # **** schedule linearizer
@@ -23,7 +22,7 @@ def create_schedule(sched_sink:UOp) -> UOp:
    for u in sched_sink.toposort(gate_kernel_sink):
      if u.op is not Ops.AFTER: continue
      k = u.src[1]
-      assert k.op in {Ops.CALL, Ops.END, Ops.LINEAR}, f"AFTER src[1] should be CALL or END, not {k.op}"
+      assert k.op in {Ops.CALL, Ops.END}, f"AFTER src[1] should be CALL or END, not {k.op}"
      in_degree.setdefault(k, 0)
      if k.op is Ops.END: assert k.src[0].op is Ops.CALL, f"END src[0] should be KERNEL, not {k.src[0].op}"
      # WAR deps from rangeify are stored in AFTER src[2:]
@@ -92,27 +91,31 @@ def create_new_buffer(ctx:tuple[dict[UOp, UOp], tuple[UOp, ...]], b:UOp):
  return ret

 pm_post_sched_cache = PatternMatcher([
-  (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[1][x.arg].rtag() if x.tag is None else None),
+  (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[1][x.arg]),
  # create new BUFFERs for LUNIQUE BUFFERs from rangeify
  (UPat(Ops.BUFFER, src=(UPat(Ops.LUNIQUE), UPat(Ops.DEVICE)), name="b"), create_new_buffer),
 ])

-# the AFTER structure is already in LINEAR
-pm_collapse_after = PatternMatcher([
-  (UPat(Ops.AFTER, name="x"), lambda x: x.src[0])
+pm_resolve_linear_call = PatternMatcher([
+  # call LINEAR is resolved here
+  (UPat(Ops.CALL, src=(UPat(Ops.LINEAR),), name="linear_call", allow_any_len=True), lambda linear_call:
+   graph_rewrite(linear_call.src[0], pm_post_sched_cache, ctx=({}, linear_call.src[1:]), walk=True, name="params to buffers")),
+  # LINEAR on LINEAR
+  (UPat(Ops.LINEAR, custom_early_reject={Ops.LINEAR}, name="x"),
+   lambda x: x.replace(src=tuple(flatten(x.src if x.op is Ops.LINEAR else (x,) for x in x.src)))),
 ])

 schedule_cache: dict[bytes, UOp] = {}
-def lower_schedule_to_linear(big_sink:UOp) -> UOp|None:
+# ctx is just for DEBUG on inner
+def lower_sink_to_linear(function:UOp) -> UOp|None:
  st = time.perf_counter()
-  function = big_sink.src[0]
  if isinstance(function.arg, KernelInfo): return None
-  if not SCACHE or (sc_ret:=schedule_cache.get(function.key, None)) is None:
-    if SPEC: type_verify(big_sink, tensor_spec)
+  cache_key = function.key
+  if not SCACHE or (sc_ret:=schedule_cache.get(cache_key, None)) is None:
+    if SPEC: type_verify(function, tensor_spec)
    # support recursive CALLs
-    function = graph_rewrite(function, pm_schedule, name="inner schedule to linear")
    linear = create_schedule(get_kernel_graph(function))
-    if SCACHE: schedule_cache[function.key] = linear
+    if SCACHE: schedule_cache[cache_key] = linear
  else:
    # schedule cache hit
    linear = sc_ret
@@ -124,20 +127,21 @@ def lower_schedule_to_linear(big_sink:UOp) -> UOp|None:
    else:
      frm = None
    print(f"scheduled {len(linear.src):5d} kernels in {(time.perf_counter()-st)*1000:8.2f} ms"+\
-          f" | {' cache hit' if SCACHE and sc_ret is not None else 'CACHE MISS'} {function.key.hex()[:8]}"+\
+          f" | {' cache hit' if SCACHE and sc_ret is not None else 'CACHE MISS'} {cache_key.hex()[:8]}"+\
          f" | {len(UOpMetaClass.ucache):7d} uops in cache"+("" if frm is None else f" | {frm.filename}:{frm.lineno}"))
-  # TODO: use walk and avoid the remove tags
-  linear = graph_rewrite(linear, pm_post_sched_cache, ctx=({}, big_sink.src[1:]), walk=True, name="params to buffers")
-  return graph_rewrite(linear, pm_collapse_after+_remove_all_tags, name="remove tags/after")
+  return linear

 pm_schedule = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.SINK),), allow_any_len=True, name="big_sink"), lower_schedule_to_linear),
+  (UPat(Ops.SINK, name="function"), lower_sink_to_linear),
 ])

@track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[0]))}")
 def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[list[ExecItem], dict[str, int]]:
  # big_sink srcs are all the Tensors
-  linear = graph_rewrite(big_sink, pm_schedule, name="schedule to linear")
+  linear_call = graph_rewrite(big_sink, pm_schedule, name="schedule to linear", enter_calls=True)
+
+  # this recursively resolves the linear_call and allocates buffers
+  linear = graph_rewrite(linear_call, pm_resolve_linear_call, name="resolve linear call")

  # vars used in the schedule
  used_vars = set().union(*[{v.expr for v in si.src[0].variables()} for si in linear.src])
@@ -1,5 +1,5 @@
 import functools
-from typing import Generic, TypeVar, Callable, cast
+from typing import Generic, TypeVar, Callable, cast, overload
 from tinygrad.helpers import Context, dedup, getenv
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, PatternMatcher, UPat
 from tinygrad.tensor import Tensor
@@ -13,12 +13,15 @@ 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')
-class function(Generic[ReturnType]):
-  def __init__(self, fxn:Callable[..., ReturnType]):
+class _function(Generic[ReturnType]):
+  def __init__(self, fxn:Callable[..., ReturnType], *, precompile:bool=False):
    self.fxn = fxn
+    self.precompile = precompile

  def __get__(self, obj, objtype=None): return functools.partial(self.__call__, obj) if obj is not None else self

@@ -57,6 +60,14 @@ class function(Generic[ReturnType]):
    #call = assigned.call(*call_uops, buffer, name=name)
    #ret = buffer.after(call)

-    ret = uret.call(*call_uops, name=name)
+    ret = uret.call(*call_uops, name=name, precompile=self.precompile)
    return cast(ReturnType, Tensor(ret, device=ret.device))

+# overload signatures support both @function and @function(precompile=True) syntax
+@overload
+def function(fxn:Callable[..., ReturnType], *, precompile:bool=False) -> _function[ReturnType]: ...
+@overload
+def function(fxn:None=None, *, precompile:bool=False) -> Callable[[Callable[..., ReturnType]], _function[ReturnType]]: ...
+def function(fxn=None, *, precompile:bool=False):
+  if fxn is None: return lambda f: _function(f, precompile=precompile)
+  return _function(fxn, precompile=precompile)
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	37a40bf975	early lower cat	2026-03-07 11:39:20 +08:00
geohot	af1db22b25	simpler	2026-03-07 10:11:21 +08:00
geohot	be0f9d1055	min	2026-03-07 10:00:29 +08:00
geohot	5b9a6c5520	Add Ops.CAT movement op (ai slop)	2026-03-06 18:51:25 +08:00
George HotzandGitHub	6fd18ef875	rename CAT to VCAT (#15167 )	2026-03-06 18:46:28 +08:00
Roelof van DijkandGitHub	059c6326c0	metal uint32 icb offset overflow (#15156 ) * metal uint32 icb offset overflow fix: diff supports_exec_item GraphRunner.supports_exec_item tests fix: can't import on non-metal stricter * also test the non-metal buffer case * imports on non-mac	2026-03-06 00:54:39 +03:00
chenyuandGitHub	da61088ca4	more divmod recombine (#15162 )	2026-03-05 12:53:22 -05:00
chenyuandGitHub	167a1d56a6	improve divmod folding (#15148 ) canonicalize to div than mod which enables more simplifcation	2026-03-05 10:07:36 -05:00
sirhcmandGitHub	b824579e4d	simplify image_conv2d pitch alignment hacks (#15158 )	2026-03-05 07:17:34 -05:00
qazalandGitHub	5bf542469d	viz: python traceback for USER device (#15160 ) * start * ux * unittests	2026-03-05 20:22:09 +09:00
Roelof van DijkandGitHub	d65923bda5	tensor.py: add normalize function (#15159 ) * tensor.py: add normalize function * p==0 should match torch	2026-03-05 18:55:53 +08:00
wozeparrotandGitHub	4544da1c54	llama3 fixes part3 (#15152 )	2026-03-05 01:17:54 -08:00
Roelof van DijkandGitHub	fc0534910c	q5k is like q4k (#15155 )	2026-03-05 17:02:49 +08:00
Ananta RanganathanandGitHub	8ef656324e	FIXED TEST Q5_K GGUF dequant (#15147 ) * q5_k gguf support as separate pr * fix the problematic gemv test for q5_k * add assert to make sure the gemv test cant fail with warning instead of error	2026-03-05 16:32:36 +08:00
George HotzandGitHub	e97922a57c	LLM speedup with two jits, prefill/rollout (#15153 ) * START_TIME * print cleanup * fix tests	2026-03-05 16:21:09 +08:00
wozeparrotandGitHub	be23772d43	llama3 fixes part2 (#15150 )	2026-03-04 23:43:50 -08:00
wozeparrotandGitHub	0c769289eb	llama3: more scripts (#15107 )	2026-03-04 22:18:03 -08:00
George HotzandGitHub	fb43b415f9	fix symbolic shape call + chunked prefill (#15149 ) * fix precompile for symbolic shape * chunked prefill * cleaner * test that	2026-03-05 14:02:26 +08:00
George HotzandGitHub	8a82b26522	llm: print the prefill cache size (#15146 ) * print the llm prefill cache size * mock that too	2026-03-05 12:13:28 +08:00
chenyuandGitHub	b5370fd52d	use copy_multi in alu_multi [pr] (#15143 ) * use copy_multi in alu_multi [pr] * copy to anything	2026-03-04 22:53:00 -05:00
George HotzandGitHub	72a9ed6e23	fix render depth bug + add warmup to serve + no realize default (#15144 ) * fix render depth bug + add warmup to serve * make realize not the default	2026-03-05 11:21:16 +08:00
George HotzandGitHub	ac1847cbf7	fully symbolic llm (#15097 ) * work * llm symbolic (almost) * work * revert that * llm sym * works * cleanups * cache tokens with the kv cache * cleanups * cleanups	2026-03-05 10:22:11 +08:00
qazalandGitHub	33a1970045	sqtt: simplify inst mapping, validate JUMP processing in CI (#15139 ) * jump cleanup * assert there's a JUMP * new example for JUMP * regenerate examples * rdna4 work * new packets * work * less for branch handling * less verbose * fix err message	2026-03-05 09:53:12 +09:00
chenyuandGitHub	04da527a7a	minor div_and_mod_symbolic cleanups (#15138 )	2026-03-04 19:05:44 -05:00
chenyuandGitHub	106d18b792	use UOp methods in allreduce.py [pr] (#15137 ) except the one line with Ops.BUFFER and Ops.NOOP, not sure what that's for	2026-03-04 17:15:33 -05:00
chenyuandGitHub	34594bcaaf	Revert "bug in metal: offset is stored as uint32, overflow (#15129 )" (#15136 ) This reverts commit `9c58db16fa`.	2026-03-04 16:54:42 -05:00
Roelof van DijkandGitHub	9c58db16fa	bug in metal: offset is stored as uint32, overflow (#15129 ) * metal uint32 icb offset overflow * fix: diff * supports_exec_item * GraphRunner.supports_exec_item * tests * fix: can't import on non-metal	2026-03-04 22:52:12 +03:00
chenyuandGitHub	4cce283790	relax test_tqdm_perf (#15134 )	2026-03-04 12:58:47 -05:00
chenyuandGitHub	fae400d300	update assign tests to also test the expected behavior (#15132 )	2026-03-04 11:34:43 -05:00
chenyuandGitHub	1f96cc2b51	update non-contiguous buffer error message [pr] (#15131 ) * update non-contiguous buffer error message [pr] also cleaned up the tests * order	2026-03-04 11:13:26 -05:00
nimlgenandGitHub	563d5c3211	more graph tests (#15130 )	2026-03-04 19:01:12 +03:00
nimlgenandGitHub	cdc48da9cd	hevc: assert and speed (#15122 ) * hevc: assert and speed * simpler	2026-03-04 19:01:02 +03:00
wozeparrotandGitHub	4e9b85ecfd	fa: pull inputs out of call (#15127 )	2026-03-04 03:15:49 -08:00
geohot	47faa2d7b4	hotfix: llm kv cache uses clone instead of realize to avoid many realize	2026-03-04 19:07:03 +08:00
George Hotz GitHub Woze Parrot	8ebd24637b	fix fa forward building with clang 22 (#15124 ) * fix fa forward building with clang 22 * fix: override rocm path --------- Co-authored-by: Woze Parrot <[email protected]>	2026-03-04 02:32:25 -08:00
sirhcmandGitHub	592f9bf6c6	set OPENPILOT_HACKS=1 to enable replace assign (#15123 )	2026-03-04 05:26:04 -05:00
wozeparrotandGitHub	df23057984	fa: change bwd grid dim + unshuffle using mops (#15068 )	2026-03-04 01:23:40 -08:00
sirhcmandGitHub	5623cea7b1	move openpilot contiguous hacks to schedule (#15120 )	2026-03-04 03:04:06 -05:00
wozeparrotandGitHub	759c7fc81c	failing test for allreduce memory usage (#15106 )	2026-03-03 23:38:38 -08:00
George HotzandGitHub	5ecfe549e7	allreduce is a function with LATE_ALLREDUCE=1 (#15119 ) * allreduce as a function * allreduce function * support allreduce function * LATE_ALLREDUCE	2026-03-04 15:17:58 +08:00
sirhcmandGitHub	e7e70a3c95	simplify idx before counting backward_slice (#15117 )	2026-03-03 23:53:50 -05:00
George HotzandGitHub	2d72a4a90c	fix copying padded const (#15116 ) * fix const padding cpu * remove comment	2026-03-04 10:39:45 +08:00
chenyuandGitHub	b5ebb4d06d	contiguous_view_offset returns only offset [pr] (#15113 ) size is always input.size	2026-03-03 15:23:39 -05:00
nimlgenandGitHub	abd830b260	am: setup_rinf returns only doorbell (#15112 )	2026-03-03 19:27:41 +03:00
nimlgenandGitHub	4b42bb54aa	am: reset sdma to start from 0 (#15109 )	2026-03-03 18:14:46 +03:00
George HotzandGitHub	01ddb4c267	add precompile to call (#15099 ) * add precompile to call * put get back * something * after structure * alt * keep it call * resolve call * resolve linear call * precompile works with llm * revert rangeify * color for debugging * getenv PRECOMPILE * clean up deco pattern * fully recursive sink scheduling * revert llama * fix SPEC=2	2026-03-03 22:32:42 +08:00
qazalandGitHub	c7f908b788	sqtt: fix rdna4 structs (#15111 ) * work * DEBUG=2	2026-03-03 23:32:14 +09:00
qazalandGitHub	8dd691761d	sqtt: remove old files (#15108 )	2026-03-03 22:43:24 +09:00
sirhcm GitHub Comma Device	de043226ba	benchmark comma usbgpu driving_vision step and load time (#15103 ) Co-authored-by: Comma Device <[email protected]>	2026-03-03 06:08:03 -05:00
sirhcmandGitHub	5f6b610da1	FLOAT16 logic for IMAGE==1 goes back to image_conv2d (#15105 )	2026-03-03 05:37:57 -05:00
wozeparrotandGitHub	529318259c	fix: fix null tests to actually use null device (#15104 )	2026-03-03 02:05:47 -08:00
George HotzandGitHub	7d025089e3	no after removal (#15102 ) * no after removal * we are using walk * null schedule test * pytest deps * Revert "pytest deps" This reverts commit `5e1c5304ec`. * Revert "null schedule test" This reverts commit `02da66053e`. * clean null tests	2026-03-03 17:50:31 +08:00
wozeparrotandGitHub	92c16810ac	feat: per device mem_used (#15100 )	2026-03-03 01:31:28 -08:00
qazalandGitHub	e3a0598d0b	viz: the whole pc should be in view (#15101 )	2026-03-03 17:17:53 +09:00
b1tgandGitHub	a9ea36de79	assembly/amd: v_cmp_lg_f32 is ordered not-equal (#14982 )	2026-03-03 15:37:48 +08:00
wozeparrotandGitHub	c35de9bd68	asm_gemm: support more sharding (#15002 )	2026-03-02 23:16:37 -08:00
wozeparrotandGitHub	824ba4386a	llama3 dp fix (#15098 )	2026-03-02 22:43:07 -08:00
chenyuandGitHub	5dcf29b1a0	use clone in test_swap_slices (#15096 )	2026-03-02 22:05:12 -05:00
sirhcmandGitHub	c70e8af068	move IMAGE FLOAT16 logic to allocations (#15095 ) * FLOAT16 logic in allocations * cleanup * separate that * only apply when IMAGE == 1 * test passing now * create image buffers earlier	2026-03-02 22:00:05 -05:00
George HotzandGitHub	d483e4153a	buffer view is like buffer (#15082 ) * buffer view is like buffer * fix * swap_reshape_shrink * contiguous on gguf, fix overlap * revert that * _device_supports_view * this * fix that test * 0 buffers * that test was wrong * this * check correct size * contig BUFFER_VIEW * this * fix tests * buffer view tests * om * fix torch * no MOCKGPU * skip	2026-03-03 09:52:33 +08:00
qazalandGitHub	62ee976c1b	gemm/asm: cleanup repeated patterns to helper functions (#15094 )	2026-03-03 08:14:47 +09:00
qazalandGitHub	848f5cea96	viz: sqtt instruction packet trace (#15065 )	2026-03-03 07:55:04 +09:00
chenyuandGitHub	14d1c5fdfd	assign fusion tests on detach and contiguous_backward (#15092 )	2026-03-02 15:21:51 -05:00
nimlgenandGitHub	dfa180413d	tbgpu: sign nv (#15087 )	2026-03-02 22:58:30 +03:00
chenyuandGitHub	71f228f80f	test exact kernel count in torch_backend/test_kernel_fusion (#15091 )	2026-03-02 14:26:32 -05:00
chenyuandGitHub	f80b1033c5	simpler Tensor.all (#15089 ) same generated kernel	2026-03-02 11:08:55 -05:00
chenyuandGitHub	4008f7d4e8	move Tensor.one_hot +1 to python (#15088 )	2026-03-02 10:56:41 -05:00
nimlgenandGitHub	dafbe9733a	am: cleanup (#15086 )	2026-03-02 17:06:21 +03:00
qazalandGitHub	f7aeff6061	viz: cli.py cleanups, do not require PYTHONPATH (#15085 ) * cleanup the print * sys.exit * equal check * cleanup unpacker * cli doesn't need PYTHONPATH * no semicolons * %s/PYTHONPATH=. //g	2026-03-02 19:24:38 +09:00
George HotzandGitHub	5ff278446c	add contiguous_view_offset (#15084 ) * add contiguous_view_offset * no int	2026-03-02 18:05:04 +08:00
sirhcmandGitHub	977c270774	IMAGE=1 kernel count failing tests (#15083 )	2026-03-02 04:35:26 -05:00
George HotzandGitHub	3539693555	Support triu variable on diagonal + SDPA symbolic (#15081 ) * triu variable * fails * dumbbb * no commutative in reshape * real fix * revert that * sdpa symbolic tests	2026-03-02 12:19:48 +08:00
wozeparrotandGitHub	a4f6365929	llama3: fstep takes grads (#15069 )	2026-03-01 20:05:07 -08:00
Nick GitHub George Hotz	8e8e9f6ff6	assert removal for _tri() + tests (#15073 ) * assert removal for _tri() and tests * removed import * tests triu/tril like in prefill --------- Co-authored-by: George Hotz <[email protected]>	2026-03-02 10:34:28 +08:00
nimlgenandGitHub	ccbbca05ef	beam: add dev_timeout for am (#15063 ) * beam: add dev_timeout for am * all covered * fk * x * fuzz * reset * f	2026-03-01 16:57:29 +03:00
chenyuandGitHub	8cb4368967	delete unused END NOOP rule [pr] (#15077 )	2026-03-01 00:09:05 -05:00
chenyuandGitHub	efce99adc9	skip isComposing key press in llm.py (#15076 ) for the CJK input user	2026-02-28 20:31:53 -05:00
chenyuandGitHub	103ea16ec0	add contiguous back to svd (#15074 ) can cause infinite loop	2026-02-28 16:49:26 -05:00
chenyuandGitHub	fe0fa8333b	Revert "improve Tensor.sort indices (#15070 )" (#15072 ) This reverts commit `e3003631f2`.	2026-02-28 14:40:30 -05:00
chenyuandGitHub	e3003631f2	improve Tensor.sort indices (#15070 ) * improve Tensor.sort indices instead of N^2 match at the end, have an arange to start and go through the same N(logN)^2 path * contiguous	2026-02-28 14:16:16 -05:00
wozeparrotandGitHub	cfc5cf65ad	llama3: vocab padding fix + jit copies on fakedata (#15067 )	2026-02-28 08:44:55 -08:00
chenyuandGitHub	76170d035a	relax atol for test_xlm_roberta_large (#15066 )	2026-02-28 11:22:35 -05:00
qazalandGitHub	cfb8e6922d	viz: arrow keys move through time (#15064 ) * work * automatic zoom, keeping scale * the whole shape should be out of view	2026-02-28 23:52:36 +09:00
nimlgenandGitHub	9b3450c9da	test gpu crash on cdna (#15062 )	2026-02-28 13:17:59 +03:00