optimize openpilot vision on qcom

2026-08-16 18:18:25 +00:00 · 2026-07-18 00:48:58 +00:00
115 changed files with 1845 additions and 1851 deletions
@@ -171,7 +171,7 @@ jobs:
        llvm: 'true'
        amd: 'true'
    - name: Run NULL backend tests
-      run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --durations=20
+      run: DEV=NULL python -m pytest -n=auto test/null/ --durations=20
    - name: Run targeted tests on NULL backend
      run: |
        DEV=NULL python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
@@ -500,29 +500,6 @@ jobs:
      - name: Run LLVM test
        run: DEV=MOCKKFD+AMD:LLVM python test/device/test_amd_llvm.py

-  hcq2:
-    name: hcq2
-    runs-on: *linux
-    timeout-minutes: 5
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v6
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: hcq2
-          deps: testing_unit
-          amd: 'true'
-      - name: Run HCQ2 tests
-        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/test_tiny.py
-      - name: Run HCQ2 multi-device tests
-        run: |
-          HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_multitensor.py \
-            TestMultiTensor.test_simple_add TestMultiTensor.test_shard_reduce \
-            TestMultiTensor.test_backward_sum TestMultiTensor.test_matmul_shard_0_0
-      - name: Run HCQ2 JIT tests
-        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_jit.py
-
  testmockam:
    name: Linux (am)
    runs-on: *linux
@@ -643,7 +620,7 @@ jobs:
    - name: Run unit tests
      run: DEV=METAL python -m pytest -n=auto test/unit/ --durations=20
    - name: Run NULL backend tests
-      run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --durations=20
+      run: DEV=NULL python -m pytest -n=auto test/null/ --durations=20
    - name: Test tensor core ops (fake)
      run: DEV=METAL DEBUG=3 TC=2 python test/backend/test_ops.py TestOps.test_gemm
    - name: Test tensor core ops (real)
@@ -1753,12 +1753,11 @@ def train_gptoss():

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

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

  @TinyJit
-  @Context(TRAINING=1)
  def minibatch(tokens:Tensor):
    if is_dp: tokens = tokens.to(None).shard(device, 0)
    if not is_sharding: tokens = tokens.to(None)
@@ -154,7 +154,7 @@ class FlatTransformer:
    self.tok_embeddings = nn.Embedding(vocab_size, dim)
    self.tok_embeddings.weight = Tensor.normal(vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
    self.output = Tensor.normal(1, vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
-    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).clone().is_param_(False)
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().is_param_(False)

    def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().is_param_(False)
    names = ["xqkv", "xo", "x2"]
@@ -195,18 +195,18 @@ class FlatTransformer:
                                                                  next_grad_amax_state=next_grad_amax_xqkv)
    amaxs.append(new_amax)
    saves.extend([x_normed, rrms, *s, xqkv])
+    xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
+    xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
+    xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
+    xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
+
+    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
+    xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
    if getenv("HK_FLASH_ATTENTION"):
-      from extra.thunder.amd.fa import flash_attention, fused_qkv_rope
-      xq, xk, xv = fused_qkv_rope(xqkv, freqs_cis, self.n_heads, self.n_kv_heads, self.head_dim)
+      from extra.thunder.amd.fa import flash_attention
      attn, *save = flash_attention(xq, xk, xv, is_causal=True, write_flat=True)
      saves.extend(save)
    else:
-      xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
-      xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
-      xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
-      xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
-      xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
-      xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
      xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
      attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
    attn = attn.reshape(bsz, seqlen, -1)
@@ -313,8 +313,7 @@ class FlatTransformer:

  def __call__(self, tokens:Tensor, save:bool=True):
    h = self.tok_embeddings(tokens)
-    freqs_cis = self.freqs_cis.cast(h.dtype)
-    if not getenv("HK_FLASH_ATTENTION"): freqs_cis = freqs_cis[:, :tokens.shape[1], :, :, :]
+    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
    a, na, ga, nga, s = self._fp8_amax, self._fp8_next_amax, self._fp8_grad_amax, self._fp8_next_grad_amax, self._fp8_inv_scale
    for i in range(self.n_layers):
      attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i],
@@ -20,7 +20,7 @@ export ZERO_OPTIM=${ZERO_OPTIM:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}

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

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

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

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

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

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

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

+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_RPT", 1):
+    from extra.gemm.qcom_openpilot_vision_fp16 import patch_fp32_rpt
+    if (patched:=patch_fp32_rpt(run_onnx_jit)): print(f"repeat-packed {patched} QCOM vision kernels")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_SCHEDULE", 1):
+    from extra.gemm.qcom_openpilot_schedule_projection import patch_projection
+    if (patched:=patch_projection(run_onnx_jit)): print(f"rescheduled {patched} QCOM vision kernels")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_FULL_RPT", 1):
+    from extra.gemm.qcom_openpilot_inverse_full_rpt import patch_model as patch_full_rpt
+    if (patched:=patch_full_rpt(run_onnx_jit)): print(f"fully repeat-packed {patched} QCOM vision kernels")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_DEDUPE", 1):
+    from extra.gemm.qcom_openpilot_dedupe_head import dedupe_identical_calls
+    if (removed:=dedupe_identical_calls(run_onnx_jit)): print(f"deduplicated {len(removed)} QCOM kernels")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_PACK_CONV", 1):
+    from extra.gemm.qcom_openpilot_pack_conv_weights import patch_conv
+    if (patched:=patch_conv(run_onnx_jit)): print(f"packed weights for {patched} QCOM convolution kernels")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_LEVEL_SCHEDULE", 1):
+    from extra.gemm.qcom_openpilot_level_schedule import schedule_levels
+    if (moved:=schedule_levels(run_onnx_jit)): print(f"rescheduled {moved} QCOM kernels by dependency level")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_BATCH_HEAD", 1):
+    from extra.gemm.qcom_openpilot_batch_head import batch_head
+    if (combined:=batch_head(run_onnx_jit)): print(f"batched {combined} groups of QCOM head kernels")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_INPUT_PACK", 1):
+    from extra.gemm.qcom_openpilot_input_pack import patch_input_pack
+    if (patched:=patch_input_pack(run_onnx_jit)): print(f"vectorized {patched} QCOM input kernel")
  with open(OUTPUT, "wb") as f:
    pickle.dump(run_onnx_jit, f)
  mdl_sz = os.path.getsize(onnx_file)
@@ -72,7 +96,7 @@ def compile(onnx_file):
  print(f"mdl size is {mdl_sz/1e6:.2f}M")
  print(f"pkl size is {pkl_sz/1e6:.2f}M")
  print("**** compile done ****")
-  return inputs, test_val
+  return run_onnx_jit, inputs, test_val

 def test_vs_compile(run, inputs, test_val=None):

@@ -142,9 +166,10 @@ if __name__ == "__main__":
    test_vs_compile(pickle_loaded, inputs)
  else:
    onnx_file = fetch(OPENPILOT_MODEL)
-    inputs, outputs = compile(onnx_file)
+    pickle_loaded, inputs, outputs = compile(onnx_file)

-    with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
+    if OUTPUT != os.devnull:
+      with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)

    test_vs_compile(pickle_loaded, inputs, outputs)
    if getenv("SELFTEST"):
@@ -80,7 +80,7 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
      # NOTE: since this is part of K, these 2 can be anywhere in the frags and long as a and b match
      a_frag = a_frag.reshape(2, 8)[lane_m, :]
      b_frag = b_frag.reshape(2, 8)[lane_m, :]
-    wmma = UOp.wmma(a_frag, b_frag, acc_frag.after(k), (16, 16, 16), 'AMD', 32)
+    wmma = UOp.wmma(a_frag, b_frag, acc_frag.after(k), ((16, 16, 16), 'AMD', 32))
    acc_store = acc_frag.store(wmma).end(tile_m, tile_n)
  else:
    # registers for LOCAL -> REG
@@ -13,7 +13,7 @@ WMMA_ACC = WMMA_M // LANES_PER_WAVE_M
 THREADS_PER_BLOCK = WARP_SIZE * WAVES_M * WAVES_N
 LDS_PAD = 4  # pad LDS rows to reduce bank conflicts

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

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

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

  # end KV tile loop
  n_tile_end = acc_frag.store(pv).end(tm2, tn2).end(k_pv).barrier().end(n_tile)
@@ -1,5 +1,5 @@
 from tinygrad import UOp, dtypes
-from tinygrad.uop.ops import AxisType, KernelInfo, AddrSpace
+from tinygrad.uop.ops import AxisType, Ops, KernelInfo, AddrSpace
 from extra.gemm.amd_uop_matmul import test_matmul

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

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

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

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

 WARP_SIZE = 64

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

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

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

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

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

 WARP_SIZE = 64

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

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

  # store back the acc
  acc_store = acc[N_inner_loop, M_inner_loop].store(out)
@@ -0,0 +1,110 @@
+#!/usr/bin/env python3
+"""Batch adjacent independent openpilot head kernels into one QCOM launch."""
+import argparse, pickle, re
+from dataclasses import replace
+
+from tinygrad import Device
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops
+from extra.gemm.qcom_openpilot_ir3 import plain_name
+
+MAX_BATCH={"r_256_4_128_4":4,"r_128_16_4_16_4":4,"r_128_16_4_32_4":4,
+           "r_8_16_4_8_4":4,"r_8_4_8_4":4,"r_8_4_8_4n1":4}
+MAX_BATCH.update({"r_16_16_4_8_4":4,"r_4_16_4_8_4":4,"r_16_16_4_4":4,
+                  "r_4_4_4_4":4,"r_16_16_4_4n1":4,"r_4_4_4_4n1":4})
+
+
+def batched_source(source:str, name:str, batch_count:int) -> str:
+  match=re.search(r"__kernel void \w+\((.*?)\) \{",source,re.S)
+  if match is None: raise RuntimeError("kernel signature not found")
+  declarations=[x.strip() for x in match.group(1).split(",")]
+  arg_names=[x.rsplit(" ",1)[1] for x in declarations]
+  renamed=[]
+  bodies=[]
+  body=source[match.end():source.rfind("}")]
+  local_decls=re.findall(r"__attribute__\s*\(\(aligned \(\d+\)\)\)\s*__local\s+[^;]+;",body)
+  hoisted=[]
+  for batch in range(batch_count):
+    mapping={arg:f"{arg}_{batch}" for arg in arg_names}
+    renamed.extend(decl.rsplit(" ",1)[0]+" "+mapping[arg] for decl,arg in zip(declarations,arg_names))
+    branch=body
+    for declaration in local_decls:
+      local_match=re.search(r"(\w+)(\[[^;]+;)$",declaration)
+      if local_match is None: raise RuntimeError(f"local declaration not understood: {declaration}")
+      old=local_match.group(1)
+      new=f"{old}_{batch}"
+      hoisted.append(declaration[:local_match.start(1)]+new+local_match.group(2))
+      branch=branch.replace(declaration,"")
+      branch=re.sub(rf"\b{re.escape(old)}\b",new,branch)
+    for old,new in mapping.items(): branch=re.sub(rf"\b{re.escape(old)}\b",new,branch)
+    bodies.append(branch)
+  prefix=source[:match.start()]
+  count=len(declarations)
+  order=tuple(batch*count for batch in range(batch_count))+tuple(
+    batch*count+arg for batch in range(batch_count) for arg in range(1,count))
+  branches=" else ".join((f"if (get_group_id(1)=={batch}) " if batch < batch_count-1 else "")+f"{{{body}}}"
+                         for batch,body in enumerate(bodies))
+  return f"{prefix}__kernel void {name}_batch{batch_count}({','.join(renamed[i] for i in order)}) {{\n" \
+         f"{''.join(hoisted)}\n{branches}\n}}"
+
+
+def independent(calls:list) -> bool:
+  outputs={call.src[out+1] for call in calls for out in call.src[0].arg.outs}
+  return not any(arg in outputs for call in calls for i,arg in enumerate(call.src[1:]) if i not in call.src[0].arg.outs)
+
+
+def batch_head(model) -> int:
+  outer=model.captured.linear.src[0]
+  batch=list(outer.src[0].src[0].src)
+  new_batch=[]
+  combined=0
+  index=0
+  cache={}
+  while index < len(batch):
+    first=batch[index]
+    name=plain_name(first.src[0].arg.name) if first.op is Ops.CALL and first.src[0].op is Ops.PROGRAM else ""
+    if index+1 < len(batch) and name in MAX_BATCH:
+      calls=[first]
+      while index+len(calls) < len(batch) and len(calls) < MAX_BATCH[name]:
+        candidate=batch[index+len(calls)]
+        candidate_name=plain_name(candidate.src[0].arg.name) if candidate.op is Ops.CALL and candidate.src[0].op is Ops.PROGRAM else ""
+        if candidate_name != name or first.src[0].src[3].arg != candidate.src[0].src[3].arg: break
+        calls.append(candidate)
+      if len(calls) > 1 and independent(calls):
+        batch_count=len(calls)
+        program=first.src[0]
+        source=batched_source(program.src[2].arg,name,batch_count)
+        if source not in cache: cache[source]=Device["QCOM"].compiler.compile_cached(source)
+        aux0=program.arg.aux[0]
+        count=len(aux0)
+        ordered_aux=tuple(aux0[0] for _ in calls)+tuple(entry for _ in calls for entry in aux0[1:])
+        combined_aux=tuple(tuple((new_index,dtype,shape) for _old_index,dtype,shape in entry)
+                           for new_index,entry in enumerate(ordered_aux))
+        info=replace(program.arg,name=f"{name}_batch{batch_count}",global_size=(program.arg.global_size[0],batch_count,1),
+                     globals=tuple(range(count*batch_count)),outs=tuple(range(batch_count)),
+                     ins=tuple(range(batch_count,count*batch_count)),aux=(combined_aux,))
+        program=program.replace(arg=info,src=program.src[:2]+
+          (program.src[2].replace(arg=source),program.src[3].replace(arg=cache[source])))
+        new_batch.append(first.replace(src=(program,*[call.src[1] for call in calls],
+                                                *[arg for call in calls for arg in call.src[2:]])))
+        combined+=1
+        index+=batch_count
+        continue
+    new_batch.append(first)
+    index+=1
+  model.captured._linear=model.captured.linear.substitute({outer:create_graph_call(new_batch)},walk=True)
+  model.captured.__dict__.pop("linear",None)
+  return combined
+
+
+def main() -> None:
+  parser=argparse.ArgumentParser()
+  parser.add_argument("input")
+  parser.add_argument("output")
+  args=parser.parse_args()
+  with open(args.input,"rb") as f:model=pickle.load(f)
+  print("combined",batch_head(model))
+  with open(args.output,"wb") as f:pickle.dump(model,f)
+
+
+if __name__ == "__main__":main()
@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+"""Remove byte-identical duplicate linear chains in the driving-vision head."""
+import argparse, hashlib, pickle
+
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops, UOp
+from extra.gemm.qcom_openpilot_ir3 import plain_name
+
+TARGETS = {"r_128_16_4_32_4", "r_256_4_128_4", "r_128_16_4_16_4"}
+
+
+def dedupe_identical_calls(model, all_calls:bool=True) -> list[tuple[int, str]]:
+  """Alias calls with identical programs, inputs, and byte-identical constants."""
+  outer = model.captured.linear.src[0]
+  batch = outer.src[0].src[0].src
+  produced:dict[UOp, UOp] = {}
+  static_hash:dict[UOp, str] = {}
+  seen:dict[tuple, tuple[UOp, ...]] = {}
+  new_batch, removed = [], []
+
+  def representative(buf:UOp) -> UOp:
+    while buf in produced and produced[buf] is not buf: buf = produced[buf]
+    return buf
+
+  def content_hash(buf:UOp) -> str:
+    if buf not in static_hash:
+      static_hash[buf] = hashlib.sha256(memoryview(buf.buffer.numpy()).cast("B")).hexdigest()
+    return static_hash[buf]
+
+  for index, original in enumerate(batch):
+    call = original.replace(src=tuple(representative(x) if x in produced else x for x in original.src))
+    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM or (not all_calls and plain_name(call.src[0].arg.name) not in TARGETS):
+      new_batch.append(call)
+      if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM:
+        for out_index in call.src[0].arg.outs: produced[original.src[out_index+1]] = call.src[out_index+1]
+      continue
+    program = call.src[0]
+    output_indices = set(program.arg.outs)
+    signature_args = []
+    for arg_index, (before, after) in enumerate(zip(original.src[1:], call.src[1:])):
+      if arg_index in output_indices: continue
+      if before.op is Ops.PARAM:
+        signature_args.append(("param", before.arg))
+      elif before in produced:
+        signature_args.append(("dynamic", representative(before)))
+      else:
+        signature_args.append((str(after.dtype), after.buffer.size, content_hash(after)))
+    signature = (plain_name(program.arg.name), program.src[3].arg, tuple(signature_args))
+    outputs = tuple(original.src[i+1] for i in program.arg.outs)
+    if signature in seen:
+      canonical_outputs = seen[signature]
+      for output, canonical in zip(outputs, canonical_outputs): produced[output] = representative(canonical)
+      removed.append((index, plain_name(program.arg.name)))
+    else:
+      new_batch.append(call)
+      canonical_outputs = tuple(call.src[i+1] for i in program.arg.outs)
+      seen[signature] = canonical_outputs
+      for output, canonical in zip(outputs, canonical_outputs): produced[output] = canonical
+
+  # Apply aliases to consumers which occur after the duplicate chains.
+  new_batch = [call.replace(src=tuple(representative(x) if x in produced else x for x in call.src)) for call in new_batch]
+  model.captured._linear = model.captured.linear.substitute({outer:create_graph_call(new_batch)}, walk=True)
+  model.captured.__dict__.pop("linear", None)
+  return removed
+
+
+def main() -> None:
+  parser = argparse.ArgumentParser()
+  parser.add_argument("input")
+  parser.add_argument("output")
+  parser.add_argument("--all", action="store_true", help="deduplicate every program family, not only the head linears")
+  args = parser.parse_args()
+  with open(args.input, "rb") as f: model = pickle.load(f)
+  removed = dedupe_identical_calls(model, args.all)
+  with open(args.output, "wb") as f: pickle.dump(model, f)
+  print(f"removed {len(removed)} duplicate head calls: {removed}")
+
+
+if __name__ == "__main__": main()
@@ -0,0 +1,49 @@
+"""Vectorize the driving-vision uint8 input normalization kernel on QCOM."""
+from dataclasses import replace
+
+from tinygrad import Device
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops
+from extra.gemm.qcom_openpilot_ir3 import plain_name
+
+TARGET = "E_8192_3_4_2_4"
+SOURCE = r"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+__kernel void E_8192_3_4_2_4(write_only image2d_t O,__global uchar *A,__global uchar *B,
+                             __global half *MEAN,__global half *STD) {
+  int c=get_global_id(0),i=get_global_id(1),off=(c<<17)+(i<<2),mc=c<<2;
+  uchar4 a0=vload4(0,A+off),a1=vload4(0,A+off+32768);
+  uchar4 a2=vload4(0,A+off+65536),a3=vload4(0,A+off+98304);
+  uchar4 b0=vload4(0,B+off),b1=vload4(0,B+off+32768);
+  uchar4 b2=vload4(0,B+off+65536),b3=vload4(0,B+off+98304);
+  half4 ma=vload4(0,MEAN+mc),mb=vload4(0,MEAN+mc+12);
+  half4 ia=(half4)(1)/vload4(0,STD+mc),ib=(half4)(1)/vload4(0,STD+mc+12);
+  int x=c+(i&7)*24,y=i>>3;
+  write_imagef(O,(int2)(x,y),convert_float4(((half4)(a0.x,a1.x,a2.x,a3.x)-ma)*ia));
+  write_imagef(O,(int2)(x+3,y),convert_float4(((half4)(b0.x,b1.x,b2.x,b3.x)-mb)*ib));
+  write_imagef(O,(int2)(x+6,y),convert_float4(((half4)(a0.y,a1.y,a2.y,a3.y)-ma)*ia));
+  write_imagef(O,(int2)(x+9,y),convert_float4(((half4)(b0.y,b1.y,b2.y,b3.y)-mb)*ib));
+  write_imagef(O,(int2)(x+12,y),convert_float4(((half4)(a0.z,a1.z,a2.z,a3.z)-ma)*ia));
+  write_imagef(O,(int2)(x+15,y),convert_float4(((half4)(b0.z,b1.z,b2.z,b3.z)-mb)*ib));
+  write_imagef(O,(int2)(x+18,y),convert_float4(((half4)(a0.w,a1.w,a2.w,a3.w)-ma)*ia));
+  write_imagef(O,(int2)(x+21,y),convert_float4(((half4)(b0.w,b1.w,b2.w,b3.w)-mb)*ib));
+}"""
+
+
+def patch_input_pack(jit) -> int:
+  outer = jit.captured.linear.src[0]
+  batch = outer.src[0].src[0].src
+  lib, new_batch, replaced = None, [], 0
+  for call in batch:
+    name = plain_name(call.src[0].arg.name) if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM else ""
+    if name == TARGET:
+      if lib is None: lib = Device["QCOM"].compiler.compile_cached(SOURCE)
+      program = call.src[0]
+      program = program.replace(arg=replace(program.arg, global_size=(1, 64, 1), local_size=(3, 128, 1)),
+                                src=program.src[:2] +
+                                (program.src[2].replace(arg=SOURCE), program.src[3].replace(arg=lib)))
+      call, replaced = call.replace(src=(program, *call.src[1:])), replaced+1
+    new_batch.append(call)
+  if replaced:
+    jit.captured._linear = jit.captured.linear.substitute({outer:create_graph_call(new_batch)}, walk=True)
+    jit.captured.__dict__.pop("linear", None)
+  return replaced
@@ -0,0 +1,123 @@
+#!/usr/bin/env python3
+"""Pack all four FP32 accumulator vectors in the OpenPilot inverse projection."""
+import argparse, hashlib, pickle, struct
+
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops
+from extra.gemm.qcom_openpilot_ir3 import branch as BR, mad_f32 as MAD_F32, mov_f32 as MOV_F32, nop as NOP, plain_name
+
+TARGET="r_32_64_4_4_192_4"
+INVERSE_W_TARGETS={"r_512_16_4_4_48_4","r_128_32_4_4_96_4"}
+SAFE_DONORS={"d4c281a1","1fe26758","e34e7e58"}
+
+
+def replace_src2(ins:bytes, src2:int) -> bytes:
+  lo,hi=struct.unpack("<II",ins)
+  return struct.pack("<II",(lo&0xff00ffff)|(src2<<16),hi)
+
+
+def replace_low_src(ins:bytes, src:int) -> bytes:
+  lo,hi=struct.unpack("<II",ins)
+  return struct.pack("<II",(lo&0xffffff00)|src,hi)
+
+
+def pack_inverse_full(lib:bytes) -> bytes:
+  off,size=struct.unpack_from("<I",lib,0xc0)[0],struct.unpack_from("<I",lib,0x100)[0]
+  instrs=[lib[i:i+8] for i in range(off,off+size,8)]
+  if len(instrs)!=175: raise RuntimeError(f"expected 175 inverse instructions, got {len(instrs)}")
+  # Move loop control from r13.x into the existing r12.w zero register. This
+  # makes r13-r16 four contiguous accumulator vectors without growing the
+  # shader's declared register file.
+  out=instrs[:11]
+  for acc in ("r13.x","r14.x","r15.x","r16.x"): out.append(MOV_F32(acc,"r12.w",rpt=3))
+  loop_start=len(out)
+  body=list(instrs[16:32])
+  body[0]=replace_low_src(body[0],51)   # add r8.z, r12.w, 192
+  body[1]=replace_low_src(body[1],51)   # add r9.x, r12.w, 384
+  body[2]=replace_src2(body[2],51)      # add r9.z, c28.y, r12.w
+  body[3]=replace_low_src(body[3],51)   # mov r10.x, r12.w
+  out+=body
+  for component,weight in zip("xyzw",("r5.x","r2.x","r3.x","r4.x")):
+    out.append(MAD_F32("r13.x",f"r7.{component}",weight,"r13.x",rpt=3,r=True,sy=component=="x"))
+  out+=instrs[48:60]
+  control=list(instrs[60:65])
+  control[0]=replace_low_src(control[0],51)  # increment r12.w
+  control[2]=replace_low_src(control[2],51)  # compare r12.w
+  control[3]=MOV_F32("r12.w","r0.x")
+  out+=control
+  out.append(BR(loop_start-len(out)))
+  tail=list(instrs[66:131])
+  # The first residual moved from r12.w to r13.x; y/z/w were already in r13.
+  tail[90-66]=replace_src2(tail[90-66],52)
+  out+=tail
+  out += [NOP()]*(len(instrs)-len(out))
+  if len(out)!=len(instrs): raise RuntimeError(f"packed image has {len(out)} instructions")
+  return lib[:off]+b"".join(out)+lib[off+size:]
+
+
+def with_fregs(lib:bytes, count:int) -> bytes:
+  out=bytearray(lib)
+  regoff=struct.unpack_from("<I",out,0x34)[0]+0x14
+  regs=struct.unpack_from("<I",out,regoff)[0]
+  struct.pack_into("<I",out,regoff,(regs&0x80000000)|max(regs&0x7fffffff,count))
+  return bytes(out)
+
+
+def pack_inverse_w_full(lib:bytes) -> bytes:
+  off,size=struct.unpack_from("<I",lib,0xc0)[0],struct.unpack_from("<I",lib,0x100)[0]
+  ins=[lib[i:i+8] for i in range(off,off+size,8)]
+  if len(ins) not in (174,178): raise RuntimeError(f"expected 174/178 inverse-W instructions, got {len(ins)}")
+  out=ins[:19]+[MOV_F32("r17.x","r12.z",rpt=3)]
+  loop_start=len(out)
+  out+=ins[19:35]
+  for component,weight in zip("xyzw",("r5.x","r2.x","r3.x","r4.x")):
+    out.append(MAD_F32("r17.x",f"r7.{component}",weight,"r17.x",rpt=3,r=True,sy=component=="x"))
+  out+=ins[51:68]
+  out.append(BR(loop_start-len(out)))
+  tail=list(ins[69:])
+  mapping={50:68,52:69,53:70,54:71}
+  first_store=next(i for i,x in enumerate(tail) if struct.unpack_from("<I",x,4)[0]>>24==0xc0)
+  for i in range(first_store):
+    lo,_=struct.unpack("<II",tail[i])
+    src2=(lo>>16)&0xff
+    if src2 in mapping: tail[i]=replace_src2(tail[i],mapping[src2])
+  out+=tail
+  out += [NOP()]*(len(ins)-len(out))
+  return with_fregs(lib[:off]+b"".join(out)+lib[off+size:],18)
+
+
+def patch_model(model,names:set[str]|None=None) -> int:
+  outer=model.captured.linear.src[0]
+  batch=list(outer.src[0].src[0].src)
+  cache,patched={},0
+  for index,call in enumerate(batch):
+    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM: continue
+    name=plain_name(call.src[0].arg.name)
+    if name not in (names if names is not None else {TARGET}|INVERSE_W_TARGETS): continue
+    program=call.src[0]
+    old=program.src[3].arg
+    # These transforms relocate fixed compiler registers. A different QCOM compiler allocation can
+    # have the same instruction count but different live values and must not be patched by index.
+    if hashlib.sha1(old).hexdigest()[:8] not in SAFE_DONORS: continue
+    if old not in cache:
+      cache[old]=pack_inverse_full(old) if name==TARGET else pack_inverse_w_full(old)
+    program=program.replace(src=program.src[:3]+(program.src[3].replace(arg=cache[old]),))
+    batch[index]=call.replace(src=(program,*call.src[1:]))
+    patched+=1
+  model.captured._linear=model.captured.linear.substitute({outer:create_graph_call(batch)},walk=True)
+  model.captured.__dict__.pop("linear",None)
+  return patched
+
+
+def main() -> None:
+  ap=argparse.ArgumentParser()
+  ap.add_argument("input")
+  ap.add_argument("output")
+  ap.add_argument("--names",help="comma-separated program families")
+  args=ap.parse_args()
+  with open(args.input,"rb") as f:model=pickle.load(f)
+  print("patched",patch_model(model,set(args.names.split(",")) if args.names else None))
+  with open(args.output,"wb") as f:pickle.dump(model,f)
+
+
+if __name__=="__main__":main()
@@ -0,0 +1,34 @@
+"""Small IR3 encoding helpers used by the OpenPilot QCOM graph patches."""
+import re, struct
+
+ANSI_RE = re.compile(r"\x1b\[[0-9;]*m")
+
+def plain_name(name:str) -> str: return ANSI_RE.sub("", name)
+
+def _freg(name:str|int) -> int:
+  if isinstance(name, int): return name
+  register, component = name.replace("r", "").split(".")
+  return int(register) * 4 + "xyzw".index(component)
+
+def _pack(lo:int, hi:int) -> bytes: return struct.pack("<II", lo & 0xffffffff, hi & 0xffffffff)
+
+def nop() -> bytes: return _pack(0, 0)
+
+def branch(offset:int) -> bytes: return struct.pack("<iI", offset, 0x00900000)
+
+def mov_f32(dst:str, src:str, rpt:int=0, sy:bool=False, ss:bool=False, r:bool=False) -> bytes:
+  return _pack(_freg(src), (0x30044000 if sy else 0x20044000) | (0x1000 if ss else 0) |
+               (0x800 if r else 0) | ((rpt & 0x7f) << 8) | _freg(dst))
+
+def add_s(dst:str, src:str, imm:int, ss:bool=False) -> bytes:
+  lo = ((0x27 if imm < 0 else 0x20) << 24) | ((imm & 0xff) << 16) | _freg(src)
+  return _pack(lo, 0x42300000 | (0x1000 if ss else 0) | _freg(dst))
+
+def mad_f32(dst:str, src1:str, src2:str, src3:str, rpt:int=0, sy:bool=False, r:bool=False) -> bytes:
+  d, s1, s2, s3 = _freg(dst), _freg(src1), _freg(src2), _freg(src3)
+  hi = ((0x73 if sy else 0x63) << 24) | (0x80 << 16) | ((s2 >> 1) << 16) | (((s2 & 1) << 7 | (rpt & 0x7f)) << 8) | d
+  lo = (0x20000000 if r else 0) | (s3 << 16) | (0x8000 if r else 0) | s1
+  return _pack(lo, hi)
+
+def isam_f32(dst:str, coord:str, tex:int=0, samp:int=0) -> bytes:
+  return _pack((tex * 2) << 24 | ((samp & 7) << 21) | (_freg(coord) * 2 + 1), 0xa0001f00 | _freg(dst))
@@ -0,0 +1,47 @@
+#!/usr/bin/env python3
+"""Group ready OpenPilot graph calls by program family without crossing dependency levels."""
+import argparse, pickle
+from collections import defaultdict
+
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops
+from extra.gemm.qcom_openpilot_ir3 import plain_name
+
+
+def schedule_levels(model) -> int:
+  outer=model.captured.linear.src[0]
+  batch=list(outer.src[0].src[0].src)
+  writer,levels={},{}
+  grouped=defaultdict(list)
+  for sequence,call in enumerate(batch):
+    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM:
+      grouped[sequence].append((sequence,call))
+      continue
+    deps={writer[arg] for arg in call.src[1:] if arg in writer}
+    level=1+max((levels[dep] for dep in deps),default=-1)
+    levels[sequence]=level
+    grouped[level].append((sequence,call))
+    for output in call.src[0].arg.outs: writer[call.src[output+1]]=sequence
+  scheduled=[]
+  moved=0
+  for entries in grouped.values():
+    ordered=sorted(entries,key=lambda item:(plain_name(item[1].src[0].arg.name),item[0]))
+    scheduled.extend(call for _index,call in ordered)
+    moved+=sum(old_index!=entries[new_index][0] for new_index,(old_index,_call) in enumerate(ordered))
+  if scheduled != batch:
+    model.captured._linear=model.captured.linear.substitute({outer:create_graph_call(scheduled)},walk=True)
+    model.captured.__dict__.pop("linear",None)
+  return moved
+
+
+def main() -> None:
+  ap=argparse.ArgumentParser()
+  ap.add_argument("input")
+  ap.add_argument("output")
+  args=ap.parse_args()
+  with open(args.input,"rb") as f:model=pickle.load(f)
+  print("moved",schedule_levels(model))
+  with open(args.output,"wb") as f:pickle.dump(model,f)
+
+
+if __name__=="__main__":main()
@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+"""Prepack static weights for the slow stride-2 openpilot 7x7 convolutions."""
+import argparse, pickle, re
+
+import numpy as np
+from tinygrad import Device
+from tinygrad.device import Buffer
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops, UOp
+from extra.gemm.qcom_openpilot_ir3 import plain_name
+
+TARGETS={"r_16_8_16_2_4_4_7_7", "r_8_4_32_2_4_4_7_7", "r_4_2_64_2_4_4_7_7"}
+
+
+def packed_source(source:str) -> str:
+  start=source.index("    half val0")
+  end=source.index("    int alu20",start)
+  weight_name=re.search(r"__global half\* (data2_\d+)",source).group(1)  # type: ignore[union-attr]
+  block="    int wp=(((alu0*2+alu2)*7+Ridx0)*28);\n"+"\n".join(
+    f"    float4 w{i}=convert_float4(vload4(0,{weight_name}+wp+{i*4}));" for i in range(7))+"\n"
+  source=source[:start]+block+source[end:]
+  accum_start=source.index("    *(buf0+0)",start)
+  loop_prefix=source[start:accum_start]
+  casts=re.findall(r"    float (cast\d+) = \(\(float\)\(val(\d+)\)\);\n",loop_prefix)
+  assert len(casts) == 28
+  source=source[:start]+re.sub(r"    float cast\d+ = \(\(float\)\(val\d+\)\);\n", "", loop_prefix)+source[accum_start:]
+  mapping={cast:("w0.x" if int(val) == 27 else f"w{int(val)%7+1}.x" if int(val) < 6 else
+                       f"w{(int(val)-6)%7}.{'yzw'[(int(val)-6)//7]}") for cast,val in casts}
+  for old,new in sorted(mapping.items(),key=lambda item:-len(item[0])):
+    source=re.sub(rf"\b{old}\b",new,source)
+  return source
+
+
+def pack_weight_buffer(weight:UOp) -> UOp:
+  original=np.asarray(weight.buffer.numpy()).reshape(-1)
+  outputs=original.size//896
+  assert outputs*896 == original.size
+  packed=np.empty((outputs,2,7,7,4),dtype=np.float16)
+  for output in range(outputs):
+    for parity in range(2):
+      for row in range(7):
+        base=output*896+parity+row*28
+        for tap in range(7):
+          for component in range(4): packed[output,parity,row,tap,component]=original[base+component*224+tap*4]
+  buf=Buffer("QCOM",packed.size,weight.dtype,initial_value=bytearray(packed.tobytes()))
+  return UOp.from_buffer(buf)
+
+
+def patch_conv(model) -> int:
+  outer=model.captured.linear.src[0]
+  batch=outer.src[0].src[0].src
+  new_batch=[]
+  replaced=0
+  for call in batch:
+    name=plain_name(call.src[0].arg.name) if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM else ""
+    if name in TARGETS:
+      program=call.src[0]
+      source=packed_source(program.src[2].arg)
+      lib=Device["QCOM"].compiler.compile_cached(source)
+      program=program.replace(src=program.src[:2]+(program.src[2].replace(arg=source),program.src[3].replace(arg=lib)))
+      call=call.replace(src=(program,call.src[1],call.src[2],pack_weight_buffer(call.src[3]),*call.src[4:]))
+      replaced+=1
+    new_batch.append(call)
+  model.captured._linear=model.captured.linear.substitute({outer:create_graph_call(new_batch)},walk=True)
+  model.captured.__dict__.pop("linear",None)
+  return replaced
+
+
+def main() -> None:
+  parser=argparse.ArgumentParser()
+  parser.add_argument("input")
+  parser.add_argument("output")
+  args=parser.parse_args()
+  with open(args.input,"rb") as f:model=pickle.load(f)
+  print("patched",patch_conv(model))
+  with open(args.output,"wb") as f:pickle.dump(model,f)
+
+
+if __name__ == "__main__":main()
@@ -0,0 +1,111 @@
+#!/usr/bin/env python3
+"""Reschedule independent texture addresses in the dominant vision projection."""
+import argparse, pickle, struct
+
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops
+from extra.gemm.qcom_openpilot_ir3 import add_s as ADD_S, branch as BR, isam_f32 as ISAM_F32, mov_f32 as MOV_F32, nop as NOP, plain_name
+
+TARGET="r_32_192_4_4_64_4"
+FIRST_CONV_TARGET="r_64_32_16_4_4_6_3_3_4"
+FORWARD_STYLE={TARGET,"r_32_64_4_4_64_4","r_8_384_4_4_128_4"}
+GAP_STYLE={"r_512_16_4_4_16_4","r_512_48_4_4_16_4","r_128_32_4_4_32_4","r_128_96_4_4_32_4"}
+INVERSE_W_STYLE={"r_512_16_4_4_48_4","r_128_32_4_4_96_4"}
+INVERSE_STYLE={"r_32_64_4_4_192_4"}
+TARGETS=FORWARD_STYLE|GAP_STYLE|INVERSE_W_STYLE|INVERSE_STYLE|{FIRST_CONV_TARGET}
+
+
+def schedule_first_conv(lib:bytes) -> bytes:
+  image_offset, image_size=struct.unpack_from("<I",lib,0xc0)[0],struct.unpack_from("<I",lib,0x100)[0]
+  instrs=[lib[i:i+8] for i in range(image_offset,image_offset+image_size,8)]
+  if len(instrs) != 262: raise RuntimeError(f"expected 262 first-conv instructions, got {len(instrs)}")
+  # Use registers that the subsequent texture loads overwrite, allowing all
+  # eight independent input/weight coordinates to precede the texture reads.
+  addresses=[]
+  for index,(register,offset) in enumerate(zip(("r0","r2","r3","r4"),(-36,-24,-12,0))):
+    addresses.append(MOV_F32(f"{register}.x","r16.w",ss=index > 0) if offset == 0 else
+                     ADD_S(f"{register}.x","r16.w",offset,ss=index > 0))
+    addresses.append(MOV_F32(f"{register}.y","r16.z"))
+  addresses.extend(instrs[i] for i in (48,51,54,57))
+  loads=[ISAM_F32(dst,f"{coord}.x",tex=0) for dst,coord in zip(("r7.x","r6.x","r1.x","r0.x"),("r0","r2","r3","r4"))]
+  loads.extend(ISAM_F32(dst,coord,tex=1) for dst,coord in zip(("r2.x","r3.x","r4.x","r5.x"),("r8.x","r8.z","r9.x","r9.z")))
+  out=instrs[:32]+addresses+loads+instrs[60:86]
+  out.append(BR(31-len(out)))
+  out.extend(instrs[87:92])
+  out.append(BR(26-len(out)))
+  out.extend(instrs[93:99])
+  out.append(BR(24-len(out)))
+  out.extend(instrs[100:])
+  out.extend([NOP()]*(len(instrs)-len(out)))
+  if len(out) != len(instrs): raise RuntimeError(f"scheduled first conv has {len(out)} instructions")
+  return lib[:image_offset]+b"".join(out)+lib[image_offset+image_size:]
+
+
+def schedule_native_f16(lib:bytes) -> bytes:
+  image_offset, image_size=struct.unpack_from("<I",lib,0xc0)[0],struct.unpack_from("<I",lib,0x100)[0]
+  instrs=[lib[i:i+8] for i in range(image_offset,image_offset+image_size,8)]
+  if len(instrs) != 222: raise RuntimeError(f"expected 222 native-FP16 instructions, got {len(instrs)}")
+  addresses=(21,24,27,30,33,40,47,54)
+  loads=(23,26,29,32,35,42,49,56)
+  mads=tuple(range(36,40))+tuple(range(43,47))+tuple(range(50,54))+tuple(range(57,61))
+  out=instrs[:21]+[instrs[i] for i in addresses]+[instrs[i] for i in loads]+[instrs[i] for i in mads]+instrs[61:67]
+  out.append(BR(21-len(out)))
+  out.extend(instrs[68:])
+  out.extend([NOP()]*(len(instrs)-len(out)))
+  if len(out) != len(instrs): raise RuntimeError(f"scheduled native FP16 has {len(out)} instructions")
+  return lib[:image_offset]+b"".join(out)+lib[image_offset+image_size:]
+
+
+def schedule_loads(lib:bytes, name:str) -> bytes:
+  image_offset, image_size=struct.unpack_from("<I",lib,0xc0)[0],struct.unpack_from("<I",lib,0x100)[0]
+  instrs=[lib[i:i+8] for i in range(image_offset,image_offset+image_size,8)]
+  if name == TARGET and len(instrs) == 222: return schedule_native_f16(lib)
+  if len(instrs) < 160: raise RuntimeError(f"expected projection shader, got {len(instrs)} instructions")
+  # The compiler emits address, rpt5 nop, texture-read eight times. Calculate
+  # every independent address first, then issue the reads as one contiguous run.
+  if name in FORWARD_STYLE: start,body_end=20,66
+  elif name in GAP_STYLE: start,body_end=26,84
+  elif name in INVERSE_W_STYLE: start,body_end=19,76
+  elif name in INVERSE_STYLE: start,body_end=16,73
+  else: raise RuntimeError(f"unsupported projection {name}")
+  address_indices=tuple(start+3*i for i in range(8))
+  load_indices=tuple(start+3*i+2 for i in range(8))
+  out=instrs[:start]+[instrs[i] for i in address_indices]+[instrs[i] for i in load_indices]+instrs[start+24:body_end]
+  branch_index=len(out)
+  out.append(BR(start-branch_index))
+  out.extend(instrs[body_end+1:])
+  out.extend([NOP()]*(len(instrs)-len(out)))
+  if len(out) != len(instrs): raise RuntimeError(f"scheduled image has {len(out)} instructions")
+  return lib[:image_offset]+b"".join(out)+lib[image_offset+image_size:]
+
+
+def patch_projection(model) -> int:
+  outer=model.captured.linear.src[0]
+  batch=outer.src[0].src[0].src
+  new_batch=[]
+  cache={}
+  replaced=0
+  for call in batch:
+    name=plain_name(call.src[0].arg.name) if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM else ""
+    if name in TARGETS:
+      program=call.src[0]
+      old=program.src[3].arg
+      if old not in cache: cache[old]=schedule_first_conv(old) if name == FIRST_CONV_TARGET else schedule_loads(old,name)
+      program=program.replace(src=program.src[:3]+(program.src[3].replace(arg=cache[old]),))
+      call=call.replace(src=(program,*call.src[1:]))
+      replaced+=1
+    new_batch.append(call)
+  model.captured._linear=model.captured.linear.substitute({outer:create_graph_call(new_batch)},walk=True)
+  model.captured.__dict__.pop("linear",None)
+  return replaced
+def main() -> None:
+  parser=argparse.ArgumentParser()
+  parser.add_argument("input")
+  parser.add_argument("output")
+  args=parser.parse_args()
+  with open(args.input,"rb") as f:model=pickle.load(f)
+  print("patched",patch_projection(model))
+  with open(args.output,"wb") as f:pickle.dump(model,f)
+
+
+if __name__ == "__main__":main()
@@ -0,0 +1,165 @@
+#!/usr/bin/env python3
+"""Replace selected driving_vision 1x1 convolutions with vector FP16-acc kernels."""
+import struct
+from dataclasses import replace
+
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops
+from extra.gemm.qcom_openpilot_ir3 import branch as BR, mad_f32 as MAD_F32, nop as NOP, plain_name
+
+TARGET = "r_32_192_4_4_64_4"
+INVERSE_TARGET = "r_32_64_4_4_192_4"
+FIRST_CONV_TARGET = "r_64_32_16_4_4_6_3_3_4"
+FULL_Y_TARGETS = {TARGET, "r_32_64_4_4_64_4"}
+FULL_Z_TARGETS = {"r_8_384_4_4_128_4"}
+GAP_Y_TARGETS = {"r_512_16_4_4_16_4", "r_512_48_4_4_16_4", "r_128_32_4_4_32_4", "r_128_96_4_4_32_4"}
+INVERSE_W_TARGETS = {"r_512_16_4_4_48_4", "r_128_32_4_4_96_4"}
+OTHER_INVERSE_TARGETS: set[str] = set()
+FP32_TARGETS = FULL_Y_TARGETS | FULL_Z_TARGETS | GAP_Y_TARGETS | INVERSE_W_TARGETS | OTHER_INVERSE_TARGETS | {INVERSE_TARGET, FIRST_CONV_TARGET}
+
+def pack_fp32_mads(lib:bytes, component:str="y") -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  image = lib[image_offset:image_offset+image_size]
+  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
+  if len(instrs) < 116: raise RuntimeError(f"expected FP32 matmul loop through instruction 115, got {len(instrs)}")
+  out = instrs[:44]
+  for k_component, weight in zip("xyzw", ("r5.x", "r2.x", "r3.x", "r4.x")):
+    for acc, activation in zip(tuple(f"r{reg}.{component}" for reg in range(13, 17)), ("r7", "r6", "r1", "r0")):
+      out.append(MAD_F32(acc, f"{activation}.{k_component}", weight, acc, rpt=3,
+                         sy=len(out) == 44, r=True))
+  out += instrs[108:114]
+  branch_index = len(out)
+  out.append(BR(20-branch_index))
+  out += instrs[115:]
+  out += [NOP()] * (len(instrs)-len(out))
+  if len(out) != len(instrs): raise RuntimeError(f"packed FP32 image has {len(out)} instructions")
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def pack_gap_y_fp32_mads(lib:bytes) -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  image = lib[image_offset:image_offset+image_size]
+  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
+  if len(instrs) < 237: raise RuntimeError(f"expected at least 237 gap-Y instructions, got {len(instrs)}")
+  out = instrs[:66]
+  for component, weight in zip("xyzw", ("r5.x", "r2.x", "r3.x", "r4.x")):
+    for acc, activation in zip(("r14.y", "r15.y", "r16.y"), ("r6", "r1", "r0")):
+      out.append(MAD_F32(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True))
+  out += instrs[114:120]
+  branch_index = len(out)
+  out.append(BR(26-branch_index))
+  out += instrs[121:]
+  out += [NOP()] * (len(instrs)-len(out))
+  if len(out) != len(instrs): raise RuntimeError(f"packed gap-Y image has {len(out)} instructions")
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def pack_inverse_w_fp32_mads(lib:bytes) -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  image = lib[image_offset:image_offset+image_size]
+  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
+  if len(instrs) < 174: raise RuntimeError(f"expected at least 174 inverse-W instructions, got {len(instrs)}")
+  out = instrs[:59]
+  for component, weight in zip("xyzw", ("r5.x", "r2.x", "r3.x", "r4.x")):
+    for acc, activation in zip(("r13.w", "r14.w", "r15.w"), ("r6", "r1", "r0")):
+      out.append(MAD_F32(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True))
+  out += instrs[107:112]
+  branch_index = len(out)
+  out.append(BR(19-branch_index))
+  out += instrs[113:]
+  out += [NOP()] * (len(instrs)-len(out))
+  if len(out) != len(instrs): raise RuntimeError(f"packed inverse-W image has {len(out)} instructions")
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def pack_other_inverse_fp32_mads(lib:bytes) -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  image = lib[image_offset:image_offset+image_size]
+  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
+  if len(instrs) != 179: raise RuntimeError(f"expected 179 other-inverse instructions, got {len(instrs)}")
+  # The first output vector is split by loop-control registers. Keep it scalar;
+  # the remaining three vectors are contiguous from r14.y through r17.x.
+  out = instrs[:60]
+  for component, weight in zip("xyzw", ("r5.x", "r2.x", "r3.x", "r4.x")):
+    for acc, activation in zip(("r14.y", "r15.y", "r16.y"), ("r6", "r1", "r0")):
+      out.append(MAD_F32(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True))
+  out += instrs[108:114]
+  branch_index = len(out)
+  out.append(BR(20-branch_index))
+  out += instrs[115:]
+  out += [NOP()] * (len(instrs)-len(out))
+  if len(out) != len(instrs): raise RuntimeError(f"packed other-inverse image has {len(out)} instructions")
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def pack_first_conv_fp32_mads(lib:bytes) -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  image = lib[image_offset:image_offset+image_size]
+  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
+  if len(instrs) != 262: raise RuntimeError(f"expected 262 first-conv instructions, got {len(instrs)}")
+  out = instrs[:60]
+  first = True
+  for component, weight in zip("xyzw", ("r5", "r2", "r3", "r4")):
+    out.append(MAD_F32("r11.w", f"r7.{component}", f"{weight}.x", "r11.w", sy=first, r=True))
+    first = False
+    out.append(MAD_F32("r12.y", f"r7.{component}", f"{weight}.y", "r12.y", rpt=2, r=True))
+  for component, weight in zip("xyzw", ("r5.x", "r2.x", "r3.x", "r4.x")):
+    for acc, activation in zip(("r13.x", "r14.x", "r15.x"), ("r6", "r1", "r0")):
+      out.append(MAD_F32(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True))
+  out += instrs[124:130]
+  branch_index = len(out)
+  out.append(BR(31-branch_index))
+  out += instrs[131:]
+  # Compacting the innermost loop also relocates the two enclosing-loop branches.
+  # Their targets remain in the untouched prologue, so rebuild their relative offsets.
+  out[92] = BR(26-92)
+  out[99] = BR(24-99)
+  out += [NOP()] * (len(instrs)-len(out))
+  if len(out) != len(instrs): raise RuntimeError(f"packed first-conv image has {len(out)} instructions")
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def pack_inverse_fp32_mads(lib:bytes) -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  image = lib[image_offset:image_offset+image_size]
+  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
+  if len(instrs) != 175: raise RuntimeError(f"expected 175 inverse instructions, got {len(instrs)}")
+  # The first output vector straddles loop-control r13.x, so retain its scalar
+  # instructions. The remaining r14/r15/r16 accumulator vectors are contiguous.
+  out = instrs[:56]
+  for component, weight in zip("xyzw", ("r5.x", "r2.x", "r3.x", "r4.x")):
+    for acc, activation in zip(("r14.x", "r15.x", "r16.x"), ("r6", "r1", "r0")):
+      out.append(MAD_F32(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True))
+  out += instrs[104:109]
+  branch_index = len(out)
+  out.append(BR(16-branch_index))
+  out += instrs[110:]
+  out += [NOP()] * (len(instrs)-len(out))
+  if len(out) != len(instrs): raise RuntimeError(f"packed inverse image has {len(out)} instructions")
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def patch_fp32_rpt(jit, names:set[str]|None=None) -> int:
+  """Apply the verified FP32-accumulate repeat packing to a captured vision JIT."""
+  outer = jit.captured.linear.src[0]
+  batch = outer.src[0].src[0].src
+  new_batch, replaced = [], 0
+  for call in batch:
+    name = plain_name(call.src[0].arg.name) if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM else ""
+    if name in FP32_TARGETS and (names is None or name in names):
+      program = call.src[0]
+      patch = (pack_fp32_mads if name in FULL_Y_TARGETS else
+               (lambda lib:pack_fp32_mads(lib, "z")) if name in FULL_Z_TARGETS else
+               pack_gap_y_fp32_mads if name in GAP_Y_TARGETS else
+               pack_inverse_w_fp32_mads if name in INVERSE_W_TARGETS else
+               pack_other_inverse_fp32_mads if name in OTHER_INVERSE_TARGETS else
+               pack_first_conv_fp32_mads if name == FIRST_CONV_TARGET else pack_inverse_fp32_mads)
+      program = program.replace(src=program.src[:3] + (program.src[3].replace(arg=patch(program.src[3].arg)),))
+      if name == INVERSE_TARGET:
+        program = program.replace(arg=replace(program.arg, global_size=(8, 2, 1), local_size=(8, 16, 1)))
+      call, replaced = call.replace(src=(program, *call.src[1:])), replaced+1
+    new_batch.append(call)
+  if replaced:
+    jit.captured._linear = jit.captured.linear.substitute({outer:create_graph_call(new_batch)}, walk=True)
+    jit.captured.__dict__.pop("linear", None)
+  return replaced
@@ -3,7 +3,7 @@ from typing import cast, Callable, TypeVar, Generic, Any
 import struct, functools, time, collections, itertools
 from dataclasses import replace, dataclass
 from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, dedup, pluralize
-from tinygrad.helpers import to_tuple, round_up, partition, data64_le, panic, ContextVar
+from tinygrad.helpers import to_tuple, round_up, partition, data64_le, panic
 from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer
 from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, track_rewrites, GroupOp
 from tinygrad.uop.symbolic import symbolic
@@ -19,8 +19,6 @@ HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQ2Compiled')
 # *****************
 # 0. helpers

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

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

@@ -221,7 +219,7 @@ def add_global_sync(ctx:set[tuple[str, ...]], submit:UOp, q:UOp) -> UOp|None:
  ctx.add(devs)

  # some devices from a command buffer might be used for the first time this schedule, so we wait for their global timeline epoch.
-  wait = (make_signal(devs).index(zero:=UOp.const(dtypes.int, 0)).load() >= make_signal_value(devs).index(zero) - 1).wait()
+  wait = make_signal(devs).wait(make_signal_value(devs).index(UOp.const(dtypes.int, 0)) - 1)
  return submit.replace(src=(q.replace(src=(UOp(Ops.BARRIER, dtypes.void), wait, *q.src)),))
 pm_add_global_sync = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),), name="submit"), add_global_sync)])

@@ -239,7 +237,7 @@ def add_loads(ctx:set[int], submit:UOp, q:UOp) -> UOp|None:

        sig = make_mstack([make_signal(d if dl is None else devs[dl], queue=queue, sentinel=dl is None) for dl, d in zip(lanes, cur_devs)])
        val = make_mstack([make_signal_value(d if dl is None else devs[dl], queue=queue) for dl, d in zip(lanes, cur_devs)]).index(UOp.const(dtypes.int, 0))
-        new_src.append((sig.index(UOp.const(dtypes.int, 0)).load() >= val + dep.tag).wait())
+        new_src.append(sig.wait(val + dep.tag))
      s = s.src[0]
    new_src.append(s)
  return submit.replace(src=(q.replace(src=tuple(new_src)),))
@@ -305,7 +303,7 @@ pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION

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

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

  order = sorted(dedup(bare.values()), key=lambda g: ((b:=unwrap_mstack(g.buf_uop)[0]).arg.slot, to_tuple(b.tag)))
@@ -314,23 +312,26 @@ def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[dict[UOp, UOp
  reads = {g: table.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(UOp.const(dtypes.int, slots[bare[g]])).load() for g in gaddrs}
  return reads, (table.after(*[make_patch(table, i * table.dtype.itemsize, addr) for addr, i in slots.items()]),) if slots else ()

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

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

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

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

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

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

@@ -410,7 +411,8 @@ def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None, jit=False) -> UOp:
    # pie
    linear = graph_rewrite(linear, pm_split_patches, ctx=jit, walk=True, name="split rt/lt patches")
    linear = graph_rewrite(linear, pm_early_simplify + symbolic, bottom_up=False, name="simplify packed placeholders", enter_calls=True)
-    linear = graph_rewrite(linear, pm_rm_rt_uops, walk=True, name="replace rt uops")
+    linear = graph_rewrite(linear, pm_rm_rt_getaddrs, walk=True, name="replace rt getaddrs")
+    linear = graph_rewrite(linear, pm_rm_rt_binaries, walk=True, name="replace rt binaries")
    linear = graph_rewrite(linear, pm_replace_params, walk=True, name="replace with args")

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

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

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

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

 def resolve_getaddr(buf:UOp, g:UOp) -> UOp:
@@ -570,10 +571,6 @@ class HCQ2Buffer:
  def base(self) -> HCQ2Buffer: return self._base or self

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

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

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

 def pm4_barrier(ctx): return memory_barrier(ctx)

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

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

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

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

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

  timestamp_divider = 100.0  # AMD GPU clock: ticks/us

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

  def is_am(self) -> bool: return isinstance(self.iface, (PCIIface,))
  def is_usb(self) -> bool: return False
@@ -16,96 +16,6 @@ def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None
  axis = ref.uop.axis if axis is None else axis
  return Tensor(Tensor.invalids(*shape, dtype=dtype, device=ref.device).uop.multi(axis), dtype=dtype, device=ref.device)

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

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

    dq = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t)
    GROUP_SIZE = H_local // H_KV_local
-    HEADS_PER_WG = 2 if D == 128 and GROUP_SIZE % 2 == 0 else 1
-    dk_partial = _sharded_empty((B * GROUP_SIZE // HEADS_PER_WG, N, H_KV, D), xk, axis=shard_axis)
-    dv_partial = _sharded_empty((B * GROUP_SIZE // HEADS_PER_WG, N, H_KV, D), xv, axis=shard_axis)
+    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)
@@ -137,8 +46,8 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
      dq = dq.reshape(B, H, N//16, 4, 2, 2, D//32, 4, 4, 2).permute(0, 1, 2, 7, 8, 3, 4, 6, 5, 9).reshape(B, H, N, D).transpose(1, 2)

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

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

  BLOCK_SIZE_KV = 256
-  GROUP_SIZE = H // H_KV
-  HEADS_PER_WG = 2 if D == 128 and GROUP_SIZE % 2 == 0 else 1
  NUM_WARPS = 4
  NUM_THREADS = 64 * NUM_WARPS
-  gsz = (H // HEADS_PER_WG, 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")
@@ -28,7 +28,6 @@ constexpr int ATTN_H_KV = 8; // number of key/value heads (for GQA)
 #endif

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

 #ifndef ATTN_N
 constexpr int ATTN_N = 1024; // sequence length
@@ -54,7 +53,7 @@ using namespace kittens;
 using _gl_QdO  = gl<bf16, ATTN_B, ATTN_N, ATTN_H, ATTN_D>;
 using _gl_KV   = gl<bf16, ATTN_B, ATTN_N, ATTN_H_KV, ATTN_D>;
 using _gl_dQ   = gl<bf16, ATTN_B, ATTN_H, ATTN_N, ATTN_D>;
-using _gl_dKV  = gl<bf16, ATTN_B * (GROUP_SIZE / HEADS_PER_WG), 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 {
@@ -64,7 +63,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 / HEADS_PER_WG, (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; }
 };
@@ -72,7 +71,7 @@ template<int D> struct attn_bwd_combined_globals {
 template<int D> __launch_bounds__(NUM_THREADS, 1)
 __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr, bf16 *dO_ptr, bf16 *Q_ptr, bf16 *K_ptr, bf16 *V_ptr, float *L_vec_ptr, float *delta_vec_ptr) {

-  const int q_head_idx_fixed = blockIdx.x * HEADS_PER_WG;  // First query head handled by this workgroup.
+  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;
@@ -89,7 +88,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 * HEADS_PER_WG;
+  const int num_steps = num_steps_per_head;
  const int k_pos = j * WARP_SIZE_KV;

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

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

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

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

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

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

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

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

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

-  store<1>(g.dVg, dV_j, {batch_idx * (GROUP_SIZE / HEADS_PER_WG) + q_head_in_group / HEADS_PER_WG, 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 * (GROUP_SIZE / HEADS_PER_WG) + q_head_in_group / HEADS_PER_WG, 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);
@@ -3426,3 +3422,6 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
 }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -8,7 +8,7 @@ from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
 from tinygrad import Context, Device, Tensor, dtypes
 from hypothesis import given, settings, strategies as strat
-from test.helpers import rand_for_dtype, min_normal
+from test.helpers import rand_for_dtype
 from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX, FP8E4M3FNUZ_MAX, FP8E5M2FNUZ_MAX
 import pytest
 pytestmark = pytest.mark.filterwarnings("ignore")
@@ -25,10 +25,10 @@ def get_available_cast_dtypes(dtype: DType) -> List[DType]:
  if dtype not in supported_dtypes and dtype not in dtypes.fp8s+(dtypes.half,dtypes.bfloat16): return []
  return dts

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

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

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

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

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

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

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

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

 class TestUint16DType(TestDType):
@@ -1,9 +1,9 @@
 import numpy as np
-import functools, unittest
+import functools, unittest, ctypes

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

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

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

@unittest.skipUnless(Device[Device.DEFAULT].graph is not None, "graph support required")
 class TestGraph(unittest.TestCase):
  def skip_if_no_offset(self):
-    if Device.DEFAULT in {"WEBGPU", "CL"}: self.skipTest("device does not support _offset")
+    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
@@ -210,8 +213,8 @@ class TestGraph(unittest.TestCase):

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

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

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

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

 def run_fused_ce(bs:int, seqlen:int, vocab:int, label_smoothing:float=0.0) -> None:
  Tensor.manual_seed(0)
@@ -95,68 +92,5 @@ class TestLocalAmax(unittest.TestCase):
    self.assertEqual(GlobalCounters.kernel_count, 2)
    self.assertEqual(out.tolist(), [[0., 7., 14., 21.], [28., 35., 42., 49.], [120., 135., 150., 165.], [180., 195., 210., 225.]])

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

      vbuf.ensure_allocated()

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

 SDMA_MAX_COPY_SIZE = 0x400000

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

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

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

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

 # *** ASM24 Controller Mock ***

 _mock_usb_state: MockASM24State|None = None

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

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

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

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

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

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

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

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

  def test_from_py(self):
    assert dtypes.from_py(True) == dtypes.bool
-    assert dtypes.from_py(Invalid) == dtypes.bool
    assert dtypes.from_py(2) == dtypes.default_int
    assert dtypes.from_py(3.0) == dtypes.default_float
    assert dtypes.from_py([]) == dtypes.default_float
@@ -84,22 +84,22 @@ class TestUSBMMIOInterface(unittest.TestCase):
    self.mmio[2] = 0xFE
    self.assertEqual(full_view[2], 0xFE)

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

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

 if __name__ == "__main__":
  unittest.main()
@@ -51,8 +51,8 @@ class TestPatternMatcher(unittest.TestCase):
      ctx.append(True)
      assert len(x.src) == 0
      return x.replace(src=(UOp(Ops.NOOP),))
-    matcher = PatternMatcher([(UPat(Ops.NOOP, src=(), name="x"), fxn)])
-    c1 = UOp(Ops.NOOP)
+    matcher = PatternMatcher([(UPat(Ops.CONST, src=(), name="x"), fxn)])
+    c1 = UOp.const(dtypes.float, 1.0)
    # second rewrite shouldn't match anything
    ctx = []
    c1 = matcher.rewrite(c1, ctx)
@@ -309,6 +309,66 @@ class TestUOpGraph(unittest.TestCase):
        for uop, const in zip(uops, consts):
          self.assertEqual(uop, const)

+  @unittest.skip("no longer testable standalone")
+  def test_wmma_vectorize_fold(self):
+    for i in [2, 4, 8]:
+      vec = UOp(Ops.STACK, dtypes.half, tuple(UOp.const(dtypes.half, 0.0) for _ in range(i)))
+      var = UOp.variable("var", 0, 1, dtypes.half)
+      acc = UOp.variable('acc', 0, 1, dtypes.half)
+      wmma = UOp(Ops.WMMA, src=(vec, var, acc))
+      uops = to_uops_list([wmma])
+      self.assertEqual(uops[0], acc)
+      self.assertEqual(len(uops), 2)  # +1 for SINK
+
+    for i in [2, 4, 8]:
+      var = UOp.variable("var", 0, 1, dtypes.half)
+      vec = UOp(Ops.STACK, dtypes.half, tuple(UOp.const(dtypes.half, 0.0) for _ in range(i)))
+      acc = UOp.variable('acc', 0, 1, dtypes.half)
+      wmma = UOp(Ops.WMMA, src=(var, vec, acc))
+      uops = to_uops_list([wmma])
+      self.assertEqual(uops[0], acc)
+      self.assertEqual(len(uops), 2)  # +1 for SINK
+
+  @unittest.skip("wmma is wrong here, it needs an arg")
+  def test_wmma_vectorize_no_fold(self):
+    for i in [4, 8]:
+      vec = UOp(Ops.STACK, dtypes.half,
+                tuple(UOp.const(dtypes.half, 0.0) for _ in range(i//2)) +
+                tuple(UOp.variable(f'tmp{j}', 0, 1, dtypes.half) for j in range(i//2)))
+      var = UOp.variable(f'tmp{i}', 0, 1, dtypes.half)
+      acc = UOp.variable('acc', 0, 1, dtypes.half)
+      wmma = UOp(Ops.WMMA, src=(vec, var, acc))
+      uops = to_uops_list([wmma])
+      self.assertEqual(uops[-2], wmma)  # -2 to skip SINK
+
+    for i in [4, 8]:
+      var = UOp.variable(f'tmp{i}', 0, 1, dtypes.half)
+      vec = UOp(Ops.STACK, dtypes.half,
+                tuple(UOp.const(dtypes.half, 0.0) for _ in range(i//2)) +
+                tuple(UOp.variable(f'tmp{j}', 0, 1, dtypes.half) for j in range(i//2)))
+      acc = UOp.variable('acc', 0, 1, dtypes.half)
+      wmma = UOp(Ops.WMMA, src=(var, vec, acc))
+      uops = to_uops_list([wmma])
+      self.assertEqual(uops[-2], wmma)  # -2 to skip SINK
+
+    for i in [2, 4, 8]:
+      vec = UOp(Ops.STACK, dtypes.half,
+                tuple(UOp.const(dtypes.half, 1.0 if j == 0 else 0.0) for j in range(i)))
+      var = UOp.variable(f'tmp{i}', 0, 1, dtypes.half)
+      acc = UOp.variable('acc', 0, 1, dtypes.half)
+      wmma = UOp(Ops.WMMA, src=(vec, var, acc))
+      uops = to_uops_list([wmma])
+      self.assertEqual(uops[-2], wmma)  # -2 to skip SINK
+
+    for i in [2, 4, 8]:
+      var = UOp.variable(f'tmp{i}', 0, 1, dtypes.half)
+      vec = UOp(Ops.STACK, dtypes.half,
+                tuple(UOp.const(dtypes.half, 1.0 if j == 0 else 0.0) for j in range(i)))
+      acc = UOp.variable('acc', 0, 1, dtypes.half)
+      wmma = UOp(Ops.WMMA, src=(var, vec, acc))
+      uops = to_uops_list([wmma])
+      self.assertEqual(uops[-2], wmma)  # -2 to skip SINK
+
  def test_cast_alu_fold(self):
    d0 = UOp.param(0, dtypes.bool, (1,))
    d1 = UOp.param(1, dtypes.int, (1,))
@@ -3,50 +3,13 @@ import unittest
 import numpy as np
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import Timing, Context, cdiv
-from tinygrad.dtype import dtypes, ConstFloat, Invalid  # noqa: F401
+from tinygrad.dtype import dtypes, ConstFloat  # noqa: F401
 from tinygrad.device import Device
-from tinygrad.uop.ops import Ops, ParamArg, UOp, UPat, dtype_from_uop, exec_alu  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
-from tinygrad.uop.spec import spec_program, spec_shared, type_verify
+from tinygrad.uop.ops import Ops, ParamArg, UOp, UPat, exec_alu  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
+from tinygrad.uop.spec import spec_shared
 from tinygrad.uop.symbolic import sym
 from test.helpers import eval_uop, to_uops_list

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

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

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

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

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

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

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

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

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

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

-  def test_int_sqrt(self):
-    _assert_eq(Tensor([1, 4, 9, 16]).sqrt(), dtypes.default_float, [1, 2, 3, 4])
-
  @given(strat.sampled_from([d for d in core_dtypes if dtypes.is_int(d) and d in supported_dtypes]))
  def test_int_to_float_unary_func(self, dtype):
    for func in [
@@ -1,112 +0,0 @@
-import unittest
-from unittest.mock import patch
-
-from tinygrad import Tensor, dtypes
-from tinygrad.uop import Ops
-from tinygrad.uop.ops import UOp
-from tinygrad.uop.spec import spec_tensor
-
-
-class TestWeakPromotion(unittest.TestCase):
-  def test_rand_requires_concrete(self):
-    with self.assertRaises(ValueError): Tensor.rand(2, dtype=dtypes.weakfloat)
-    with self.assertRaises(ValueError): Tensor.const(dtypes.weakfloat, 1.0).rand_like()
-
-  def test_sum_stays_weak(self):
-    for weak, value in ((dtypes.weakint, 1), (dtypes.weakfloat, 1.0)):
-      self.assertEqual(Tensor.const(weak, value).expand(3).sum().dtype, weak)
-    self.assertEqual((Tensor.const(dtypes.weakfloat, 1.0).expand(3).sum() + Tensor([1], dtype=dtypes.float16)).dtype, dtypes.float16)
-
-  def test_storage_width(self):
-    t = Tensor.const(dtypes.weakint, 2)
-    for fn in (lambda: t.bitcast(dtypes.int32), lambda: Tensor.const(dtypes.int32, 2).bitcast(dtypes.weakint), t.element_size, t.nbytes):
-      with self.assertRaises(RuntimeError): fn()
-
-  def test_uop_scalar_const_unchanged(self):
-    for dtype, value in ((dtypes.index, 1), (dtypes.int32, 1), (dtypes.float32, 0.5)):
-      out = UOp.variable("x", 0.0 if dtype == dtypes.float32 else 0, 10.0 if dtype == dtypes.float32 else 10, dtype) + value
-      self.assertEqual((out.dtype, out.src[1].dtype), (dtype, dtype))
-
-  @unittest.expectedFailure  # TODO: a weak const defers to its consumer (JAX): these dtypes change once python scalars are weak consts
-  def test_changed_rows(self):
-    t_i8, t_f16, t_bf16 = Tensor([1], dtype=dtypes.int8), Tensor([1], dtype=dtypes.float16), Tensor([1], dtype=dtypes.bfloat16)
-    t_bool, t_u16 = Tensor([True]), Tensor([1], dtype=dtypes.uint16)
-    self.assertEqual((t_i8 + 0.5).dtype, dtypes.weakfloat)
-    self.assertEqual(((t_i8 + 0.5) + t_f16).dtype, dtypes.float16)
-    self.assertEqual(((t_i8 + 0.5) + t_bf16).dtype, dtypes.bfloat16)
-    self.assertEqual(((t_bool + 1) + t_i8).dtype, dtypes.int8)
-    self.assertEqual(((t_bool + 1) + t_u16).dtype, dtypes.uint16)
-    self.assertEqual((Tensor(3) + t_i8).dtype, dtypes.int8)
-    self.assertEqual(Tensor([2], dtype=dtypes.uint8).pad(((1, 1),), value=1).dtype, dtypes.uint8)
-    # zeros/ones are full with a python fill value, so they are weak too (jnp.zeros pins float32; deliberate divergence)
-    self.assertEqual((Tensor.zeros(3) + t_f16).dtype, dtypes.float16)
-
-  def test_unchanged_rows(self):
-    t_i8, t_f16, t_f32 = Tensor([1], dtype=dtypes.int8), Tensor([1], dtype=dtypes.float16), Tensor([1], dtype=dtypes.float32)
-    self.assertEqual((t_i8 + 1).dtype, dtypes.int8)
-    self.assertEqual((t_f16 + 0.5).dtype, dtypes.float16)
-    self.assertEqual((t_f32 + t_f16).dtype, dtypes.float32)
-
-  @unittest.expectedFailure  # TODO: dot of a weak const tensor defers to the other operand once python scalars are weak consts
-  def test_dot_defers_weak(self):
-    weak = Tensor([True, False]).where(Tensor(1), 2)
-    self.assertEqual(weak.dot(Tensor([1, 1], dtype=dtypes.int8)).dtype, dtypes.int8)
-
-  @unittest.expectedFailure  # TODO: Tensor(3).uop becomes CONST(weakint); Tensor.dtype is always uop.dtype; buffers lower to the default
-  def test_dtype_is_uop_dtype(self):
-    for value, weak, lowered in ((3, dtypes.weakint, dtypes.default_int), (0.5, dtypes.weakfloat, dtypes.default_float)):
-      t = Tensor(value)
-      self.assertEqual((t.uop.dtype, t.dtype), (weak, weak))
-      self.assertEqual(t.numpy().dtype.itemsize, lowered.itemsize)
-      realized = t.clone("CPU").realize()
-      self.assertEqual((realized.dtype, realized.uop.buffer.dtype), (lowered, lowered))
-    with patch.object(dtypes, "default_int", dtypes.int64):
-      self.assertEqual(Tensor(3).clone("CPU").realize().uop.buffer.dtype, dtypes.int64)
-
-  def test_integer_values(self):
-    x = Tensor.full((1,), 1, dtype=dtypes.int64, device="CPU")
-    self.assertEqual((x + 2**40).item(), 2**40 + 1)
-    self.assertEqual((x << 3).item(), 8)
-    self.assertTrue((x < 2**40).item())
-
-  def test_float64_precision(self):
-    value = 1.0 + 2**-40
-    x64 = Tensor.full((1,), 1.0, dtype=dtypes.float64, device="CPU")
-    self.assertEqual((x64 + value).item(), 2.0 + 2**-40)
-    x32 = Tensor.full((1,), 0.0, dtype=dtypes.float32, device="CPU")
-    self.assertEqual((x32 + value).item(), 1.0)
-
-  @unittest.expectedFailure  # TODO: exp/cos/sigmoid of a weak const stay weak instead of casting to a concrete float
-  def test_weak_transcendentals(self):
-    t_f16 = Tensor([1], dtype=dtypes.float16)
-    for out in (Tensor(2).exp(), Tensor(2).cos(), Tensor(2).sigmoid()):
-      self.assertEqual((out.dtype, (out + t_f16).dtype), (dtypes.weakfloat, dtypes.float16))
-
-  @unittest.expectedFailure  # TODO: where of weak consts stays weak and resolves per consumer
-  def test_where_and_shared_literal(self):
-    gate, weak = Tensor([True, False], device="CPU"), Tensor(2)
-    weak_where = gate.where(weak, 3)
-    self.assertEqual(weak_where.dtype, dtypes.weakint)
-    self.assertEqual((weak_where + Tensor([1, 1], dtype=dtypes.int64, device="CPU")).tolist(), [3, 4])
-    self.assertEqual((weak + Tensor([1], dtype=dtypes.int32, device="CPU")).item(), 3)
-    self.assertEqual((weak + Tensor([1], dtype=dtypes.int64, device="CPU")).item(), 3)
-
-  def test_null_lowering(self):
-    for t in (Tensor.full((1,), 1, dtype=dtypes.int64, device="NULL") + 2**40,
-              Tensor.full((1,), 1.0, dtype=dtypes.float64, device="NULL") + (1.0 + 2**-40)):
-      t.realize()
-      self.assertNotIn(t.uop.buffer.dtype, dtypes.weaks)
-
-
-class TestWeakSpec(unittest.TestCase):
-  def test_weak_operand_allowed(self):
-    x = UOp.variable("x", 0, 10, dtypes.int64)
-    weak = UOp.const(dtypes.weakint, 3)
-    for u in (x.alu(Ops.ADD, weak), x.alu(Ops.CMPLT, weak), x.alu(Ops.SHL, weak)):
-      self.assertIs(spec_tensor.rewrite(u), True)
-    gate = UOp.variable("gate", False, True, dtypes.bool)
-    self.assertIs(spec_tensor.rewrite(UOp(Ops.WHERE, dtypes.int8, (gate, UOp.const(dtypes.int8, 1), weak))), True)
-
-
-if __name__ == "__main__":
-  unittest.main()
@@ -1,6 +1,5 @@
 import unittest
 from tinygrad import Tensor
-from tinygrad.device import Buffer
 from tinygrad.dtype import Invalid, dtypes
 from tinygrad.engine.realize import run_linear

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

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

+  # tensor indexing uses reduce, so the entire result becomes invalid
+  @unittest.expectedFailure
  def test_tensor_index(self):
    idx = (Tensor.arange(4) < 2).where(Tensor([0, 1, 2, 3]), Invalid)
    out = Tensor([1.0, 2.0, 3.0, 4.0])[idx]
@@ -73,6 +73,11 @@ def contiguous_mops_to_view(c:UOp, src:UOp):
  # no symbolic shape
  if not all_int(c.shape): return None

+  # check if view is supported
+  from tinygrad.device import Device
+  devs = (src.device,) if isinstance(src.device, str) else src.device
+  if not all(hasattr(Device[d].allocator, "_offset") for d in devs): return None
+
  if buf.op is not Ops.MULTI and (view := _make_buffer_view(src)) is not None:
    view = (view.replace(dtype=c.dtype, arg=c.numel()) if c.op is Ops.BITCAST else view).reshape(c.shape)
    return c.replace(src=(view,)) if c.op is Ops.COPY else view
@@ -80,7 +80,7 @@ def unroll_axis(ctx:dict[int, int], u:UOp, arg):

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

@@ -260,13 +260,6 @@ pm_add_local_buffers = PatternMatcher([
  (UPat(Ops.STAGE, name="x"), add_local_buffer),
 ])+pm_mops

-# float ALUs need a float operand
-# make that cast explicit before the decomps, which expand SIN/LOG2/EXP2 into float polynomials and assert a float operand
-pm_cast_float_alu = PatternMatcher([
-  (UPat((Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL), src=(UPat(name="x"),), name="u"),
-   lambda u,x: u.replace(src=(x.cast(u.dtype),)) if x.dtype != u.dtype else None),
-])
-
 def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
  if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
  if DEBUG >= 5: print(pyrender(ast))
@@ -334,8 +327,6 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
  # final symbolic before decomp
  sink = graph_rewrite(sink, symbolic, name="final symbolic")

-  sink = graph_rewrite(sink, pm_cast_float_alu, name="cast float alu operands")
-
  # **** decomps ****

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

            # preserve extra reduces
            reduce_ranges = [x for x in UOp.sink(*reduceop.src[1:]).toposort() if x.op is Ops.RANGE and x.arg[0] not in tc_reduce_axes]
@@ -3,7 +3,7 @@ from dataclasses import dataclass, replace
 from collections import defaultdict
 from typing import Any, Generic, TypeVar, Iterator, Generator, TYPE_CHECKING
 import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal
-from tinygrad.helpers import LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, PROFILE, temp, colored
+from tinygrad.helpers import LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing
 from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize
 from tinygrad.dtype import DType, _to_np_dtype
@@ -112,7 +112,7 @@ class Buffer:
      if opaque is not None: self.allocate(opaque)
      if initial_value is not None:
        self.allocate()
-        self.copy_from(Buffer("PYTHON", self.size, self.dtype, opaque=memoryview(bytearray(initial_value))))
+        self.copyin(memoryview(initial_value))
    else:
      assert base._base is None, "base can't have a base"
      assert device == base.device, "base must have the same device"
@@ -135,8 +135,10 @@ class Buffer:
    if device not in self._bufs:
      allocator = Device[device].allocator
      if device == self.device: self.ensure_allocated()
-      elif self._base is not None: self._bufs[device] = allocator._offset(self._base.get_buf(device), self.nbytes, self.offset)
-      else: self._bufs[device] = allocator.map(self.ensure_allocated())
+      elif self._base is not None:
+        assert hasattr(allocator, "_offset"), "offset function required for view"
+        self._bufs[device] = allocator._offset(self._base.get_buf(device), self.nbytes, self.offset)
+      else: self._bufs[device] = allocator._map(self.ensure_allocated()._buf)
    return self._bufs[device]
  def ensure_allocated(self) -> Buffer: return self.allocate() if not self.is_initialized() else self
  def allocate(self, opaque=None, external_ptr=None) -> Buffer:
@@ -150,6 +152,7 @@ class Buffer:
    if self._base is not None:
      self._base.ensure_allocated()
      self._base.allocated_views += 1
+      assert hasattr(self.allocator, "_offset"), "offset function required for view"
      self._bufs[self.device] = self.allocator._offset(self.base._buf, self.nbytes, self.offset)
    else:
      self._bufs[self.device] = opaque if opaque is not None else self.allocator.alloc(self.nbytes, self.options)
@@ -176,7 +179,9 @@ class Buffer:
    if self._base is not None:
      return self.__class__, (self.device, self.size, self.dtype, None, None, None, 0, self.base, self.offset, self.is_allocated())
    if self.device == "NPY": return self.__class__, (self.device, self.size, self.dtype, self._buf, self.options, None, self.uop_refcount)
-    if self.is_allocated(): buf = bytearray(self.as_memoryview())
+    if self.is_allocated():
+      buf = bytearray(self.nbytes)
+      self.copyout(memoryview(buf))
    return self.__class__, (self.device, self.size, self.dtype, None, self.options, buf, self.uop_refcount)
  @property
  def trace_num(self) -> int:
@@ -191,22 +196,26 @@ class Buffer:
           (f" offset:{self.offset}" if self._base is not None else "") + (f" {self.options=}" if self.options is not None else "") + ">"
  def as_memoryview(self, allow_zero_copy=False, force_zero_copy=False) -> memoryview:
    # zero copy with as_memoryview (disabled by default due to use after free)
-    if (force_zero_copy or allow_zero_copy) and hasattr(self.allocator, '_as_buffer'): return self.allocator._as_buffer(self._buf)
+    if (force_zero_copy or allow_zero_copy) and hasattr(self.allocator, '_as_buffer') and self.options is None:
+      return self.allocator._as_buffer(self._buf)
    assert not force_zero_copy, "force zero copy was passed, but copy is required"
-    Buffer("PYTHON", self.size, self.dtype, opaque=(mv:=memoryview(bytearray(self.nbytes)))).copy_from(self)
-    return mv
+    return self.copyout(memoryview(bytearray(self.nbytes)))
  def numpy(self) -> 'np.ndarray': # type: ignore [name-defined] # noqa: F821
    import numpy as np
    assert _to_np_dtype(self.dtype) is not None, f"no np dtype for {self.dtype}"
    return np.frombuffer(self.as_memoryview(), dtype=_to_np_dtype(self.dtype))
-  def copy_from(self, src:Buffer) -> Buffer:
-    assert self.nbytes == src.nbytes, f"copy size mismatch, {self.nbytes} != {src.nbytes}"
-    assert self.is_initialized() and src.is_initialized(), "copy requires allocated buffers"
-    from tinygrad.engine.realize import run_linear
-    from tinygrad.uop.ops import UOp, Ops
-    du, su = UOp.from_buffer(self), UOp.from_buffer(src)
-    run_linear(UOp(Ops.LINEAR, src=(su.param_like(1).copy_to_device(self.device).call(du, su),)), update_stats=False)
+  def copyin(self, mv:memoryview):
+    mv = flat_mv(mv)
+    assert len(mv) == self.nbytes, f"size mismatch, {len(mv)=} != {self.dtype=} {self.size=}"
+    assert self.is_initialized(), "can't copyin to unallocated buffer"
+    self.allocator._copyin(self._buf, mv)
    return self
+  def copyout(self, mv:memoryview) -> memoryview:
+    mv = flat_mv(mv)
+    assert len(mv) == self.nbytes, f"size mismatch, {len(mv)=} != {self.dtype=} {self.size=}"
+    assert self.is_initialized(), "can't copyout unallocated buffer"
+    self.allocator._copyout(mv, self._buf)
+    return mv
  def view(self, size:int, dtype:DType, offset:int) -> Buffer:
    assert offset < self.nbytes, "offset must be less than nbytes"
    return Buffer(self.device, size, dtype, base=self.base, offset=self.offset+offset)
@@ -228,8 +237,6 @@ class Allocator(Generic[DeviceType]):
  def free(self, opaque, size:int, options:BufferSpec|None=None):
    self._free(opaque, options if options is not None else self.default_buffer_spec)

-  def map(self, buf:Buffer): return self._map(buf.ensure_allocated()._buf)
-
  # implemented by the runtime
  def _alloc(self, size:int, options:BufferSpec): raise NotImplementedError("need alloc")
  def _free(self, opaque, options:BufferSpec): pass  # if opaque is a Python object, you don't need a free
@@ -238,7 +245,7 @@ class Allocator(Generic[DeviceType]):
  def _map(self, buf): raise NotImplementedError("need map")
  def _unmap(self, mb): pass  # default no-op; override if _map allocates iface-side state
  # def _as_buffer(self, src) -> memoryview:
-  def _offset(self, buf, size:int, offset:int): raise NotImplementedError("need offset")
+  # def _offset(self, buf, size:int, offset:int):
  # def _transfer(self, dest, src, sz:int, src_dev, dest_dev):
  def _encode_decode(self, bufout, bufin, desc, hist:list, shape:tuple[int,...], frame_pos:int): raise NotImplementedError("need encdec") # optional

@@ -27,6 +27,8 @@ class InvalidType:
    if cls._instance is None: cls._instance = object.__new__(cls)
    return cls._instance
  def __eq__(self, other): return self is other
+  def __lt__(self, other): return self is not other
+  def __gt__(self, other): return self is not other
  def __hash__(self): return id(self)
  def __repr__(self): return "Invalid"
  def __reduce__(self): return (InvalidType, ())  # unpickle returns the singleton
@@ -98,7 +100,7 @@ class dtypes:
  @staticmethod
  def from_py(x) -> DType:
    # NOTE: isinstance(True, int) is True, so bool must be checked before int
-    if isinstance(x, (bool, InvalidType)): return dtypes.bool
+    if isinstance(x, bool): return dtypes.bool
    if isinstance(x, float): return dtypes.default_float
    if isinstance(x, int): return dtypes.default_int
    # put this in the last is faster because there are more items than lists/tuples to check
@@ -154,7 +156,6 @@ class dtypes:
  uints = (uint8, uint16, uint32, uint64)
  sints = (int8, int16, int32, int64)
  ints = uints + sints
-  weaks = (weakint, weakfloat)
  all = floats + ints + (bool,) # noqa: A003

 if (env_default_float := getenv("DEFAULT_FLOAT", "")):
@@ -207,7 +208,6 @@ def can_lossless_cast(dt0:DType, dt1:DType) -> bool:

 def sum_acc_dtype(dt:DType):
  # default acc dtype for sum
-  if dt in dtypes.weaks: return dt
  if dtypes.is_unsigned(dt): return least_upper_dtype(dt, dtypes.uint)
  if dtypes.is_int(dt) or dt == dtypes.bool: return least_upper_dtype(dt, dtypes.int)
  return least_upper_dtype(dt, to_dtype(getenv("SUM_DTYPE", "float32")))
@@ -166,8 +166,9 @@ def exec_copy(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
      elif src.device.startswith("DISK") and getattr(src.allocator.dev, 'fd', None) is not None \
           and hasattr(dest.allocator, 'copy_from_disk') and src.nbytes >= 4096 and dest.allocator.supports_copy_from_disk:
        dest.allocator.copy_from_disk(dest._buf, src._buf, src.nbytes)
-      elif hasattr(dest.allocator, '_as_buffer'): src.allocator._copyout(dest.allocator._as_buffer(dest._buf), src._buf)
-      else: dest.allocator._copyin(dest._buf, src.as_memoryview(allow_zero_copy=True))
+      elif src.device.startswith(("DISK", "TINYFS")) and hasattr(dest.allocator, '_as_buffer'):
+        src.allocator._copyout(dest.allocator._as_buffer(dest._buf), src._buf)
+      else: dest.copyin(src.as_memoryview(allow_zero_copy=True))
  return None

 def exec_kernel(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
@@ -215,11 +216,10 @@ def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:

  pm_exec.rewrite(call.replace(src=(ast,) + call.src[1:]), replace(ctx, update_stats=False))

-  st = time.perf_counter()
  for d in call.arg.aux.device:
    with track_stats(ctx, call, d, [], ctx.var_vals):
      if ctx.wait: Device[d].synchronize()
-  return time.perf_counter() - st
+  return None

 # flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src
 pm_flatten_linear = PatternMatcher([
@@ -77,8 +77,8 @@ class CreationMixin(DTypeMixin, MovementMixin):
    # if not buffer: assert device is None, "buffer=False does not support device specification"
    from tinygrad.uop.ops import UOp
    new_shape = argfix(shape)
-    dt = to_dtype(dtype) if dtype is not None else fill_value.dtype if isinstance(fill_value, UOp) else dtypes.from_py(fill_value)
-    val = cls.const(dt, fill_value)
+    dt = to_dtype(dtype) if dtype is not None else None
+    val = cls.const(dt or (fill_value.dtype if isinstance(fill_value, UOp) else dtypes.from_py(fill_value)), fill_value)
    val = val.reshape((1,)*len(new_shape)).expand(new_shape)
    return val.clone(device=device) if buffer else val

@@ -1,6 +1,5 @@
 from typing import TYPE_CHECKING, Self
 from tinygrad.dtype import DType, DTypeLike, dtypes, to_dtype
-from tinygrad.uop import Ops

 if TYPE_CHECKING:
  from tinygrad.uop.ops import UOp
@@ -30,7 +29,7 @@ class DTypeMixin:
    print(t.dtype, t.numpy())
    ```
    """
-    return self if self.dtype == (dt:=to_dtype(dtype)) else self._wrap_uop(self._uop.alu(Ops.CAST, arg=dt))
+    return self if self.dtype == (dt:=to_dtype(dtype)) else self._wrap_uop(self._uop.cast(dt))

  def bitcast(self, dtype:DTypeLike) -> Self:
    """
@@ -45,9 +44,7 @@ class DTypeMixin:
    print(t.dtype, t.numpy())
    ```
    """
-    dt = to_dtype(dtype)
-    if self.dtype in dtypes.weaks or dt in dtypes.weaks: raise RuntimeError(f"bitcast requires concrete dtypes, got {self.dtype} -> {dt}")
-    return self if self.dtype == dt else self._wrap_uop(self._uop.alu(Ops.BITCAST, arg=dt))
+    return self if self.dtype == (dt:=to_dtype(dtype)) else self._wrap_uop(self._uop.bitcast(dt))

  def element_size(self) -> int:
    """
@@ -58,7 +55,6 @@ class DTypeMixin:
    print(t.element_size())
    ```
    """
-    if self.dtype in dtypes.weaks: raise RuntimeError(f"element_size requires a concrete dtype, got {self.dtype}")
    return self.dtype.itemsize

  def is_floating_point(self) -> bool:
@@ -1,7 +1,7 @@
 import math, functools, operator
 from typing import TYPE_CHECKING, Literal, Self
 from tinygrad.uop import Ops
-from tinygrad.dtype import dtypes, ConstType, PyConst, least_upper_dtype
+from tinygrad.dtype import dtypes, ConstType, PyConst, least_upper_dtype, least_upper_float
 from tinygrad.helpers import argfix, polyN
 from tinygrad.mixin.creation import CreationMixin

@@ -49,9 +49,7 @@ class ElementwiseMixin(CreationMixin):
    """
    Returns a contiguous tensor.
    """
-    uop = self._uop
-    if uop.op is Ops.CONTIGUOUS or self.device is None or uop.has_buffer_identity(): return self._wrap_uop(uop)
-    return self._wrap_uop(uop.alu(Ops.CONTIGUOUS, **kwargs))
+    return self._wrap_uop(self._uop.contiguous(**kwargs))

  def contiguous_backward(self) -> Self:
    """
@@ -439,6 +437,9 @@ class ElementwiseMixin(CreationMixin):
  def threefry(self, seed: Self) -> Self:
    return self.alu(Ops.THREEFRY, seed)

+  def _ensure_float(self) -> Self:
+    return self if self.is_floating_point() else self.cast(least_upper_float(self.dtype))
+
  def reciprocal(self) -> Self:
    """
    Computes `1/x` element-wise.
@@ -447,7 +448,7 @@ class ElementwiseMixin(CreationMixin):
    print(Tensor([1., 2., 3., 4.]).reciprocal().numpy())
    ```
    """
-    return self.alu(Ops.RECIPROCAL)
+    return self._ensure_float().alu(Ops.RECIPROCAL)

  def trunc(self) -> Self:
    """
@@ -467,7 +468,7 @@ class ElementwiseMixin(CreationMixin):
    print(Tensor([1., 2., 3., 4.]).sqrt().numpy())
    ```
    """
-    return self.alu(Ops.SQRT)
+    return self._ensure_float().alu(Ops.SQRT)

  def sin(self) -> Self:
    """
@@ -477,7 +478,7 @@ class ElementwiseMixin(CreationMixin):
    print(Tensor([0., math.pi/2, math.pi, 3*math.pi/2, 2*math.pi]).sin().numpy())
    ```
    """
-    return self.alu(Ops.SIN)
+    return self._ensure_float().alu(Ops.SIN)

  def cos(self) -> Self:
    """
@@ -514,7 +515,7 @@ class ElementwiseMixin(CreationMixin):
    print(Tensor([1., 2., 4., 8.]).log2().numpy())
    ```
    """
-    return self.alu(Ops.LOG2)
+    return self._ensure_float().alu(Ops.LOG2)

  def exp2(self) -> Self:
    """
@@ -526,7 +527,7 @@ class ElementwiseMixin(CreationMixin):
    print(Tensor([0., 1., 2., 3.]).exp2().numpy())
    ```
    """
-    return self.alu(Ops.EXP2)
+    return self._ensure_float().alu(Ops.EXP2)

  def pow(self, x: Self | ConstType, reverse: bool = False) -> Self:
    """
@@ -60,7 +60,7 @@ class RandMixin(OpMixin):
    ```
    """
    dt = to_dtype(dtype or dtypes.default_float)
-    if not dtypes.is_float(dt) or dt in dtypes.weaks: raise ValueError(f"rand only supports concrete float dtypes, got {dt}")
+    if not dtypes.is_float(dt): raise ValueError(f"rand only supports float dtypes, got {dt}")
    if not all_int(shape:=argfix(*shape)) or not all(s >= 0 for s in shape): raise ValueError(f"invalid input {shape=}")
    if device is not None and not isinstance(device, str): raise ValueError(f"rand only supports single device, got {device=}")
    device = cast(str, canonicalize_device(device))
@@ -423,6 +423,9 @@ class OnnxRunner:
    if not eligible_ops: raise NotImplementedError(f"{op=} is not supported for domain {required_opset.domain} and version {required_opset.version}")
    return eligible_ops[max(eligible_ops.keys())]

+  def get_empty_input_data(self, device:str|None=None, dtype:DType|None=None) -> dict[str, Tensor]:
+    return {name:Tensor.empty(*spec.shape, device=device, dtype=dtype or spec.dtype) for name, spec in self.graph_inputs.items()}
+
  def to(self, device:str|None):
    self.graph_values = {k: (v.to(device) if isinstance(v, Tensor) else v) for k,v in self.graph_values.items()}
    self.graph_nodes = tuple(OnnxNode(n.op, n.opset_id, tuple(n.inputs), tuple(n.outputs),
@@ -31,8 +31,7 @@ class TensorIO(io.RawIOBase, BinaryIO):
  def writelines(self, lines: Iterable[Any]): raise io.UnsupportedOperation("TensorIO.writelines not supported")

 safe_dtypes = {"BOOL":dtypes.bool, "I8":dtypes.int8, "U8":dtypes.uint8, "I16":dtypes.int16, "U16":dtypes.uint16, "I32":dtypes.int, "U32":dtypes.uint,
-               "I64":dtypes.int64, "U64":dtypes.uint64, "F8_E4M3":dtypes.fp8e4m3, "F8_E5M2":dtypes.fp8e5m2,
-               "F16":dtypes.float16, "BF16":dtypes.bfloat16, "F32":dtypes.float32, "F64":dtypes.float64}
+               "I64":dtypes.int64, "U64":dtypes.uint64, "F16":dtypes.float16, "BF16":dtypes.bfloat16, "F32":dtypes.float32, "F64":dtypes.float64}
 inverse_safe_dtypes = {v:k for k,v in safe_dtypes.items()}

 def accept_filename(func: Callable[[Tensor], T]) -> Callable[[Tensor|str|pathlib.Path], T]:
@@ -52,7 +52,7 @@ class Estimates:
      elif u.op in GroupOp.ALU and u not in excluded:
        flops += (mults * (2 if u.op is Ops.MULACC else 1)) * u.max_numel()
      elif u.op is Ops.WMMA and u not in excluded:
-        flops += 2 * prod(u.arg[0]) // u.arg[3] * mults
+        flops += 2 * prod(u.arg[1]) // u.arg[5] * mults
    return Estimates(flops, lds, sum(mem.values()))

 class Renderer:
@@ -59,7 +59,7 @@ base_rewrite = PatternMatcher([
  (UPat(Ops.STORE, src=(UPat.var('bidx'), UPat.var("var"))), lambda ctx,bidx,var: f"{ctx.render_access(bidx)} = {ctx[var]};"),

  # alu/gep
-  (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{_wmma_name(x)}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]})"),
+  (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{x.arg[0]}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]})"),
  (UPat(GroupOp.ALU, name="x"), lambda ctx,x: ctx.code_for_op[x.op](
    *([strip_parens(ctx[v]) if v.op == x.op and x.op in {Ops.ADD, Ops.MUL, Ops.XOR, Ops.OR, Ops.AND} else ctx[v] for v in x.src]), x.dtype)),

@@ -98,13 +98,9 @@ pm_manual_bf16_cast = PatternMatcher([
 def uops_to_dtypes(uops:list[UOp]) -> list[tuple[DType, int]]:
  return dedup((u.dtype, u.max_numel()) for u in uops if u.addrspace in (AddrSpace.ALU, None) and u.dtype != dtypes.void and u._shape is not None)

-def _wmma_name(u:UOp) -> str:
-  return f"WMMA_{'_'.join(map(str, u.arg[0]))}_{u.arg[1].name}_{u.dtype.scalar().name}"
-
-# (name, dims, dtype_in, dtype_out, device, threads, upcast_axes)
+# (name, dims, dtype_in, dtype_out, device, threads, upcast_axes, reduce_axes)
 def wmma_args(uops:list[UOp]):
-  return dedup((_wmma_name(uop), uop.arg[0], uop.arg[1], uop.dtype.scalar(), *(uop.arg[2:5]))
-                for uop in uops if uop.op is Ops.WMMA)
+  return dedup((uop.arg[0], uop.arg[1], uop.arg[2], uop.dtype.scalar(), *(uop.arg[4:8])) for uop in uops if uop.op is Ops.WMMA)

 class CStyleLanguage(Renderer):
  kernel_typedef: str = "void"
@@ -362,7 +358,7 @@ class MetalRenderer(CStyleLanguage):
    # NOTE: this is copied from PTX
    (UPat((Ops.SQRT, Ops.EXP2, Ops.LOG2, Ops.SIN), dtype=dtypes.bfloat16, name="x"),
      lambda x: (UOp(x.op, src=tuple(vv.cast(dtypes.float) for vv in x.src), arg=x.arg).cast(dtypes.bfloat16))),
-  ]) + pm_manual_bf16_cast
+  ])

  string_rewrite = PatternMatcher([
    (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"as_type<{ctx.render_dtype(x.dtype)}>(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"),
@@ -370,7 +366,7 @@ class MetalRenderer(CStyleLanguage):

  def render_kernel(self, function_name, kernel, bufs, uops, prefix=None):
    prefix = ["#include <metal_stdlib>","using namespace metal;"]
-    deduped_wmma_args = dedup([(name, dtype_in, dtype_out) for name, _, dtype_in, dtype_out, _, _, _ in wmma_args(uops)])
+    deduped_wmma_args = dedup([(name, dtype_in, dtype_out) for name, _, dtype_in, dtype_out, _, _, _, _ in wmma_args(uops)])
    for name, dtype_in, dtype_out in deduped_wmma_args:
      dstr_out, dstr_in = self._render_dtype(dtype_out, 2, AddrSpace.REG), self._render_dtype(dtype_in, 2, AddrSpace.REG)
      prefix.append(
@@ -442,7 +438,7 @@ class CUDARenderer(CStyleLanguage):
      or (count in (2,4,8,16) and dt in dtypes.fp8s)]
    dt_map_in = { dtypes.float: "tf32", dtypes.half: "f16", dtypes.bfloat16: "bf16", dtypes.fp8e4m3: "e4m3", dtypes.fp8e5m2: "e5m2" }
    dt_map_out = { dtypes.float: "f32", dtypes.half: "f16" }
-    for name, (N, M, K), dtype_in, dtype_out, _, _, upcast_axes in wmma_args(uops):
+    for name, (N, M, K), dtype_in, dtype_out, _, _, upcast_axes, _ in wmma_args(uops):
      upcast_sizes = [prod(size for _, size in upcast) for upcast in upcast_axes]
      wmma_dtypes = [self._render_dtype(dtype, size, AddrSpace.REG) for dtype, size in zip([dtype_in, dtype_in, dtype_out], upcast_sizes)]
      n_operands = [size*dtype.itemsize//4 for dtype, size in zip([dtype_in, dtype_in, dtype_out], upcast_sizes)] # 4 => CUDA reg size in bytes
@@ -488,9 +484,9 @@ class HIPRenderer(CStyleLanguage):
    if not self.is_cdna4(target.arch): self.extra_matcher += pm_manual_bf16_cast
    if self.is_cdna(target.arch):
      self.string_rewrite = PatternMatcher([
-        (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{_wmma_name(x)}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]},"
-          f" {fp8_index(x.src[0].dtype)}, {fp8_index(x.src[0].dtype)}, 0, 0, 0, 0)" if x.arg[0][2] == 128 else None),
-        (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{_wmma_name(x)}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]}, 0, 0, 0)"),
+        (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{x.arg[0]}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]},"
+          f" {fp8_index(x.src[0].dtype)}, {fp8_index(x.src[0].dtype)}, 0, 0, 0, 0)" if x.arg[1][2] == 128 else None),
+        (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{x.arg[0]}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]}, 0, 0, 0)"),
        (UPat(Ops.CONST, dtypes.fp8s, name="x"), lambda ctx,x: f"f32_to_fp8({ctx.nan}, {fp8_index(x.dtype)})" if math.isnan(x.arg) else None),
        (UPat(Ops.CONST, dtypes.fp8s, arg=math.inf, name="x"), lambda ctx,x: f"f32_to_fp8({ctx.infinity}, {fp8_index(x.dtype)})"),
        (UPat(Ops.CONST, dtypes.fp8s, arg=-math.inf, name="x"), lambda ctx,x: f"f32_to_fp8(-{ctx.infinity}, {fp8_index(x.dtype)})"),
@@ -516,8 +512,8 @@ class HIPRenderer(CStyleLanguage):
  type_map = {dtypes.bfloat16: "hip_bfloat16", dtypes.fp8e4m3: "hip_fp8", dtypes.fp8e5m2: "hip_bf8"}
  extra_matcher = create_non_native_float_pats((dtypes.bfloat16, *dtypes.fp8s)) + PatternMatcher([
    (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
-      lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2]))
-      if x.src[0].max_numel() == 8 and x.src[0].dtype in dtypes.fp8_ocp else None),
+      lambda x: UOp(Ops.WMMA, src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64),
+        x.src[2]), arg=(*x.arg,)) if x.src[0].max_numel() == 8 and x.src[0].dtype in dtypes.fp8_ocp else None),
    # bfloat16 constant casting
    (UPat.cvar('x', dtypes.bfloat16), lambda x: cast_float_to_bf16(UOp.const(dtypes.float, x.arg))),
  ])
@@ -556,7 +552,7 @@ class HIPRenderer(CStyleLanguage):
    prefix += [f'extern "C" __attribute__((device{f", {atr}" if atr else ""})) {dto} {meth}({dti});' for meth,dti,dto,atr in ockl+ocml]
    prefix += [self.render_vector_prefix(dt, count) for dt, count in used_dtypes if count > 1]

-    for name, (N, M, K), dtype_in, dtype_out, _, _, _ in wmma_args(uops): # TODO: handle TCs f32_bf16 and bf16_bf16 w/ wrapper
+    for name, (N, M, K), dtype_in, dtype_out, _, _, _, _ in wmma_args(uops): # TODO: handle TCs f32_bf16 and bf16_bf16 w/ wrapper
      if self.is_cdna(self.target.arch):
        if (N, M, K) == (16, 16, 16): type_map[dtypes.bfloat16] = 'bf16_1k'
        elif (N, M, K) == (16, 16, 32): type_map = {**type_map, dtypes.bfloat16: "_bf16", dtypes.half: "_f16"}
@@ -37,11 +37,11 @@ def render_wmma_amd(ctx, wmma: UOp, cdna=False) -> str:
  dt_map = {dtypes.half: "f16", dtypes.float: "f32", dtypes.ushort: "bf16.1k" if cdna else "bf16", dtypes.bfloat16: "bf16.1k" if cdna else "bf16",
            dtypes.fp8e4m3: ".fp8.fp8", dtypes.fp8e5m2: ".bf8.bf8"}
  # https://github.com/llvm/llvm-project/blob/main/clang/test/CodeGenOpenCL/builtins-amdgcn-mfma.cl
-  N,M,K = wmma.arg[0]
+  N,M,K = wmma.arg[1]
  if cdna:
    if K == 32: dt_map.update({dtypes.half: ".f16", dtypes.bfloat16: ".bf16"})
    return f"  {ctx[wmma]} = call {ldt(wmma.dtype, wmma.max_numel())} @llvm.amdgcn.mfma.{dt_map[wmma.src[-1].dtype]}" + \
-           f".{N}x{M}x{K}{dt_map[wmma.arg[1]]}(" + ", ".join([f"{ldt(w.dtype, w.max_numel())} {ctx[w]}" for w in wmma.src]) + ", i32 0, i32 0, i32 0)"
+           f".{N}x{M}x{K}{dt_map[wmma.arg[2]]}(" + ", ".join([f"{ldt(w.dtype, w.max_numel())} {ctx[w]}" for w in wmma.src]) + ", i32 0, i32 0, i32 0)"
  # https://github.com/llvm/llvm-project/blob/main/llvm/test/CodeGen/AMDGPU/GlobalISel/llvm.amdgcn.wmma_32.ll
  # example: %wmma0 = call <8 x float> @llvm.amdgcn.wmma.f32.16x16x16.f16(<16 x half> %v99,<16 x half> %v100,<8 x float> %v101)
  return f"  {ctx[wmma]} = call {ldt(wmma.dtype, wmma.max_numel())} @llvm.amdgcn.wmma.{dt_map[wmma.src[-1].dtype]}.16x16x16." + \
@@ -253,30 +253,29 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n  %trunc
    if self.is_cdna:
      self.extra_matcher += PatternMatcher([
        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
-          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
+          lambda x: UOp(Ops.WMMA, src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]), arg=x.arg)
          if x.max_numel() == 4 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 4 else None),
        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
-          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2]))
-          if x.max_numel() == 4 and x.src[0].dtype in dtypes.fp8_ocp and x.src[0].max_numel() == 8 else None),
+          lambda x: UOp(Ops.WMMA, src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64),
+            x.src[2]), arg=x.arg) if x.max_numel() == 4 and x.src[0].dtype in dtypes.fp8_ocp and x.src[0].max_numel() == 8 else None),
      ])
    if target.arch in {"gfx1100", "gfx1151"}:
      self.extra_matcher += PatternMatcher([
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(x.replace(
+        (UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(UOp(Ops.WMMA,
          src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(j//2) if j%2 == 0 else UOp.const(x.src[2].dtype, 0.0)
-            for j in range(x.max_numel()*2)))),
-          arg=(*x.arg[:4], (*x.arg[4][:2], ((0, x.max_numel()*2),)))).index(i*2)
+            for j in range(x.max_numel()*2)))), arg=(*x.arg[:6], (*x.arg[6][:2], ((0, x.max_numel()*2),)), *x.arg[7:])).index(i*2)
          for i in range(x.max_numel()))) if x.max_numel() == 8 else None),
-        (UPat(Ops.WMMA, name="x"), lambda x: x.replace(
-          src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
+        (UPat(Ops.WMMA, name="x"), lambda x: UOp(Ops.WMMA,
+          src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]), arg=x.arg)
          if x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 16 else None),
      ])
    if target.arch in {"gfx1200", "gfx1201"}:
      self.extra_matcher += PatternMatcher([
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16), lambda x: x.replace(
-          src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2].bitcast(dtypes.uint16)))
+        (UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16), lambda x: UOp(Ops.WMMA,
+          src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2].bitcast(dtypes.uint16)), arg=x.arg)
            .bitcast(dtypes.bfloat16) if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None),
        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
-          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
+          lambda x: UOp(Ops.WMMA, src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]), arg=x.arg)
          if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None)
      ])

@@ -60,7 +60,7 @@ def mem_type(x:UOp) -> str: return 'shared' if x.addrspace == AddrSpace.LOCAL el

 def render_wmma(ctx: "PTXRenderer", wmma: UOp):
  assert ctx.wmma_r, "registry values for wmma must be populated"
-  (N, M, K), dtype_in, dtype_out = wmma.arg[0], wmma.arg[1], wmma.dtype
+  (N, M, K), dtype_in, dtype_out = wmma.arg[1], wmma.arg[2], wmma.arg[3]

  for src, regs in zip(wmma.src, ctx.wmma_r):
    for i, reg in enumerate(regs): # pack input and acc registers
@@ -200,7 +200,7 @@ class HCQGraph(MultiGraphRunner):
        uop_replace_j = dict(self.uop_replace[j])
        for bufid in range(len(bufs)):
          if (replace_iidx:=uop_replace_j.get(bufid)) is not None: self.input_replace_map[enqueue_dev].add((replace_iidx, dev_idx))
-          else: cast(HCQAllocator, enqueue_dev.allocator)._map(self.hcq_bufs[j][bufid])
+          else: cast(HCQAllocator, enqueue_dev.allocator).map(self.hcq_bufs[j][bufid])
        enqueue_queue.copy(self.hcq_bufs[j][0], self.hcq_bufs[j][1], dest.nbytes)
        self.copy_to_devs[cast(HCQCompiled, Device[dest.device])].add(cast(HCQCompiled, Device[src.device]))

@@ -265,7 +265,7 @@ class HCQGraph(MultiGraphRunner):
    for dev in self.devices:
      for iidx, dev_idx in self.input_replace_map[dev]:
        buf = b.bufs[dev_idx] if isinstance(b:=input_uops[iidx].buffer, MultiBuffer) else b
-        cast(HCQAllocator, dev.allocator)._map(buf._buf)
+        cast(HCQAllocator, dev.allocator).map(buf._buf)

    # Wait and restore signals
    self.kickoff_value += 1
@@ -386,7 +386,7 @@ class AMDComputeQueue(HWQueue):
    with self.pred_exec(xcc_mask=0b1):
      # NOTE: this needs an EOP buffer on the queue or it will NULL pointer
      self.release_mem(signal.value_addr, value, self.pm4.data_sel__mec_release_mem__send_32_bit_low,
-                       self.pm4.int_sel__mec_release_mem__none, cache_flush=True)
+                       self.pm4.int_sel__mec_release_mem__send_interrupt_after_write_confirm, cache_flush=True)

      if (dev:=signal.owner) is not None and signal.is_timeline and not dev.is_am():
        self.release_mem(dev.queue_event_mailbox_ptr, dev.queue_event.event_id, self.pm4.data_sel__mec_release_mem__send_32_bit_low,
@@ -490,6 +490,8 @@ class AMDCopyQueue(HWQueue):
    if (dev:=signal.owner) is not None and signal.is_timeline and not dev.is_am():
      self.q(self.sdma.SDMA_OP_FENCE | fence_flags, *data64_le(dev.queue_event_mailbox_ptr), dev.queue_event.event_id)
      self.q(self.sdma.SDMA_OP_TRAP, self.sdma.SDMA_PKT_TRAP_INT_CONTEXT_INT_CONTEXT(dev.queue_event.event_id))
+    elif dev is not None and dev.is_am(): self.q(self.sdma.SDMA_OP_TRAP, 0)
+
    return self

  def wait(self, signal:AMDSignal, value:sint=0):
@@ -640,7 +642,7 @@ class AMDProgram(HCQProgram):

 class AMDAllocator(HCQAllocator['AMDDevice']):
  def __init__(self, dev:AMDDevice):
-    super().__init__(dev, copy_bufs=getattr(dev.iface, 'copy_bufs', None),
+    super().__init__(dev, copy_bufs=getattr(dev.iface, 'copy_bufs', None), max_copyout_size=0x1000 if dev.is_usb() else None,
                     supports_copy_from_disk=dev.has_sdma_queue, supports_transfer=dev.has_sdma_queue and not dev.is_usb())

  def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
@@ -648,10 +650,11 @@ class AMDAllocator(HCQAllocator['AMDDevice']):

  def _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque)

-  def _do_map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+  def _map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)

  def _copyout(self, dest:memoryview, src:HCQBuffer):
    if not self.dev.is_usb(): return super()._copyout(dest, src)
+    if not self.dev.iface.pci_dev.usb.usb.is_custom: return super()._copyout(dest, src)
    self.dev.synchronize()

    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=TracingKey(f"{self.dev.device} -> TINY", ret=dest.nbytes), enabled=PROFILE,
@@ -677,12 +680,11 @@ class AMDQueueDesc:
    try:
      self.write_ptr[0] = self.put_value

-      # No mapped access with usb
-      if dev.is_am() and not dev.is_usb():
-        # Ensure all prior writes are visible to the GPU.
-        System.memory_barrier()
-        dev.iface.dev_impl.gmc.flush_hdp()
+      # Ensure all prior writes are visible to the GPU.
+      System.memory_barrier()

+      # Flush hdp if queue is in dev mem.
+      if dev.is_am() and not dev.is_usb(): dev.iface.dev_impl.gmc.flush_hdp()
      self.doorbell[0] = self.put_value if doorbell_value is None else doorbell_value
    except Exception as e:
      dev.error_state = e
@@ -915,10 +917,11 @@ class USBIface(PCIIface):
    self.dev, self.pci_dev, self.vram_bar, self.count = dev, USBPCIDevice("AM", *visible[dev_id]), 0, len(visible)
    self.dev_impl = AMDev(self.pci_dev)
    self._compute_props()
+    self.pci_dev.usb._pci_cacheable += [self.pci_dev.bar_info(2)] # doorbell region is cacheable

    # special regions
    self.copy_bufs = [self._dma_region(ctrl_addr=0xf000, sys_addr=0x200000, size=0x80000)]
-    self.sys_buf, self.sys_next_off = self._dma_region(ctrl_addr=0xa000, sys_addr=0x820000, size=0x1000), 0x200
+    self.sys_buf, self.sys_next_off = self._dma_region(ctrl_addr=0xa000, sys_addr=0x820000, size=0x1000), 0x800
    self.cq_buf = self._dma_region(ctrl_addr=0xb800, sys_addr=0x822000, size=0x1000)

  def _dma_region(self, ctrl_addr, sys_addr, size):
@@ -926,14 +929,19 @@ class USBIface(PCIIface):
    return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(region, has_cpu_mapping=False), view=self.pci_dev.dma_view(ctrl_addr, size), owner=self.dev)

  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, **kwargs) -> HCQBuffer:
-    # usb allocates uncached and cpu_access in vram. vram writes are faster than sram writes
-    if host and self.sys_next_off + size < self.sys_buf.size:
+    # custom usb allocates uncached and cpu_access in vram. vram writes are faster than sram writes
+    if (host or (not self.pci_dev.usb.usb.is_custom and uncached and cpu_access)) and self.sys_next_off + size < self.sys_buf.size:
      self.sys_next_off += size
      return self.sys_buf.offset(self.sys_next_off - size, size)

    # force devmem
    return super().alloc(size, host=False, uncached=uncached, cpu_access=cpu_access, contiguous=contiguous, force_devmem=True, **kwargs)

+  def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
+                   xcc_id=0, idx=0):
+    if queue_type == kfd.KFD_IOC_QUEUE_TYPE_COMPUTE: self.pci_dev.usb._pci_cacheable += [(ring.cpu_view().addr, ring.size)]
+    return super().create_queue(queue_type, ring, gart, rptr, wptr, eop_buffer, cwsr_buffer, ctl_stack_size, ctx_save_restore_size, xcc_id, idx)
+
  def sleep(self, timeout): pass

 def _mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface,), {})
@@ -129,10 +129,10 @@ class CPUAllocator(HCQAllocator):
  def _as_buffer(self, src) -> memoryview:
    self.dev.synchronize()
    return to_mv(src.va_addr, src.size)
-  def _do_map(self, buf:HCQBuffer):
+  def _map(self, buf:HCQBuffer):
    if buf.view is None or not isinstance(buf.view, MMIOInterface): raise RuntimeError("Cannot map buffer without view to cpu")
    return HCQBuffer(buf.view.addr, buf.size, view=buf.view, owner=buf.owner)
-  def _unmap(self, mb): pass  # CPU _do_map returns a view wrapper, nothing to release
+  def _unmap(self, mb): pass  # CPU _map returns a view wrapper, nothing to release

 class CPUDevice(HCQCompiled):
  def __init__(self, device:str=""):
@@ -66,4 +66,3 @@ class HIPAllocator(LRUAllocator[HIPDevice]):
  def _copyout(self, dest:memoryview, src):
    self.dev.synchronize()
    check(hip.hipMemcpy(mv_address(dest), src, len(dest), hip.hipMemcpyDeviceToHost))
-  def _offset(self, buf, size:int, offset:int): return hip.hipDeviceptr_t(buf.value + offset)
@@ -28,7 +28,7 @@ class NullAllocator(Allocator['NullDevice']):
    if not NULL_ALLOW_COPYOUT: raise RuntimeError("no copyout on NULL")
  def _transfer(self, dest, src, sz:int, src_dev, dest_dev):
    with cpu_profile(f"{src_dev.device} -> {dest_dev.device}", f"{src_dev.device}:SDMA:0"): pass
-  def _offset(self, buf, size:int, offset:int): pass
+  def _offset(self, buf, offset:int, size:int): pass

 class NullGraph(MultiGraphRunner):
  def __call__(self, input_uops:tuple[UOp, ...], var_vals:dict[str, int], wait=False) -> float|None:
@@ -342,7 +342,7 @@ class NVAllocator(HCQAllocator['NVDevice']):

  def _do_free(self, opaque:HCQBuffer, options:BufferSpec): self.dev.iface.free(opaque)

-  def _do_map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+  def _map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)

  def _encode_decode(self, bufout:HCQBuffer, bufin:HCQBuffer, desc_buf:HCQBuffer, hist:list[HCQBuffer], shape:tuple[int,...], frame_pos:int):
    assert all(h.va_addr % 0x100 == 0 for h in hist + [bufin, bufout, desc_buf]), "all buffers must be 0x100 aligned"
@@ -6,8 +6,8 @@ from typing import Any, TYPE_CHECKING
 import pickle, base64, itertools, time, sys, functools
 from dataclasses import replace
 from tinygrad.dtype import DType, dtypes, AddrSpace, truncate, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar
-from tinygrad.helpers import all_same, getenv, flatten, Target, IMAGE, is_image_shape, cpu_profile
-from tinygrad.device import Buffer, Compiled, Compiler, Allocator
+from tinygrad.helpers import all_same, getenv, flatten, Target, IMAGE, is_image_shape
+from tinygrad.device import Compiled, Compiler, Allocator
 from tinygrad.codegen.opt import tc
 from tinygrad.uop.ops import exec_alu, python_alu, Ops, UOp, GroupOp, bitcast
 from tinygrad.renderer import Renderer
@@ -134,7 +134,7 @@ class PythonProgram:
        elif u.op is Ops.WMMA:
          first_src_dtype = u.src[0].dtype
          assert isinstance(first_src_dtype, DType) # mypy
-          dims, dtype_in, device, threads = u.arg[0], first_src_dtype, u.arg[2], u.arg[3]
+          dims, dtype_in, device, threads = u.arg[1], first_src_dtype, u.arg[4], u.arg[5]
          wmma_helper = functools.partial(generic_wmma_helper, src_values, warp_size)
          # TODO: refactor these to a shared TensorCoreLayout
          if device == "METAL":
@@ -223,13 +223,8 @@ class PythonRenderer(Renderer):

 class PythonAllocator(Allocator['PythonDevice']):
  def _alloc(self, size, options): return memoryview(bytearray(size))
-  def _as_buffer(self, src) -> memoryview: return src
-  def _copyin(self, dest, src:memoryview):
-    with cpu_profile("TINY -> PYTHON", f"{self.dev.device}:COPY"): dest[:] = src
-  def _copyout(self, dest:memoryview, src):
-    with cpu_profile("PYTHON -> TINY", f"{self.dev.device}:COPY"): dest[:] = src
-  def map(self, buf:Buffer): return buf.as_memoryview(force_zero_copy=True)
-  def _offset(self, buf:memoryview, size:int, offset:int): return buf[offset:offset+size]
+  def _copyin(self, dest, src:memoryview): dest[:] = src
+  def _copyout(self, dest:memoryview, src): dest[:] = src

 class PythonDevice(Compiled):
  def __init__(self, device:str):
@@ -84,7 +84,6 @@ class QCOMComputeQueue(HWQueue):
    self.cmd(mesa.CP_WAIT_FOR_IDLE)
    if self.dev.gpu_id[:2] < (7, 3):
      self.cmd(mesa.CP_EVENT_WRITE, qreg.cp_event_write_0(event=mesa.CACHE_FLUSH_TS), *data64_le(signal.value_addr), lo32(value))
-      self._cache_flush(write_back=True, invalidate=False, sync=False, memsync=False)
    else:
      # TODO: support devices starting with 8 Gen 1. Also, 700th series have convenient CP_GLOBAL_TIMESTAMP and CP_LOCAL_TIMESTAMP
      raise RuntimeError('CP_EVENT_WRITE7 is not supported')
@@ -135,7 +134,7 @@ class QCOMComputeQueue(HWQueue):
    self.reg(mesa.REG_A6XX_SP_PERFCTR_SHADER_MASK, qreg.a6xx_sp_perfctr_shader_mask(cs=True))
    self.reg(mesa.REG_A6XX_TPL1_MODE_CNTL, qreg.a6xx_tpl1_mode_cntl(isammode=mesa.ISAMMODE_GL if prg.NIR else mesa.ISAMMODE_CL))
    self.reg(mesa.REG_A6XX_TPL1_DBG_ECO_CNTL, 0)
-    self.cmd(mesa.CP_WAIT_FOR_IDLE)
+    # CP_RUN_OPENCL snapshots the programmed compute state, so graph command streams can program the next dispatch without draining prior shader work.

    self.reg(mesa.REG_A6XX_SP_CS_NDRANGE_0,
             qreg.a6xx_sp_cs_ndrange_0(kerneldim=3, localsizex=local_size[0] - 1, localsizey=local_size[1] - 1, localsizez=local_size[2] - 1),
@@ -192,7 +191,7 @@ class QCOMComputeQueue(HWQueue):
               qreg.cp_exec_cs_1(ngroups_x=global_size[0]), qreg.cp_exec_cs_2(ngroups_y=global_size[1]), qreg.cp_exec_cs_3(_ngroups_z=global_size[2]))
    else: self.cmd(mesa.CP_RUN_OPENCL, 0)

-    self._cache_flush(write_back=True, invalidate=False, sync=False, memsync=False)
+    # The final queue signal writes back data needed by host and cross-queue consumers.
    return self

 class QCOMArgsState(HCQArgsState):
--- a/Show More
+++ b/Show More