dsp skip

* further clean up wmma

* comment

.github/workflows/test.yml

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

examples/mlperf/model_train.py

@@ -1753,11 +1753,12 @@ def train_gptoss():
 
   scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
 
-  # realize everything here
-  if optim.master_params: Tensor.realize(*optim.master_params)
+  if optim.master_params:
+    for m in optim.master_params: m.realize()
   Tensor.realize(*optim.params, *fp8_inv_scales)
 
   @TinyJit
+  @Context(TRAINING=1)
   def minibatch(tokens:Tensor):
     if is_dp: tokens = tokens.to(None).shard(device, 0)
     if not is_sharding: tokens = tokens.to(None)

examples/mlperf/models/flat_llama.py

@@ -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).contiguous().is_param_(False)
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).clone().is_param_(False)
 
     def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().is_param_(False)
     names = ["xqkv", "xo", "x2"]
@@ -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
+      from extra.thunder.amd.fa import flash_attention, fused_qkv_rope
+      xq, xk, xv = fused_qkv_rope(xqkv, freqs_cis, self.n_heads, self.n_kv_heads, self.head_dim)
       attn, *save = flash_attention(xq, xk, xv, is_causal=True, write_flat=True)
       saves.extend(save)
     else:
+      xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
+      xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
+      xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
+      xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
+      xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
+      xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
       xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
       attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
     attn = attn.reshape(bsz, seqlen, -1)
@@ -313,7 +313,8 @@ class FlatTransformer:
 
   def __call__(self, tokens:Tensor, save:bool=True):
     h = self.tok_embeddings(tokens)
-    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
+    freqs_cis = self.freqs_cis.cast(h.dtype)
+    if not getenv("HK_FLASH_ATTENTION"): freqs_cis = freqs_cis[:, :tokens.shape[1], :, :, :]
     a, na, ga, nga, s = self._fp8_amax, self._fp8_next_amax, self._fp8_grad_amax, self._fp8_next_grad_amax, self._fp8_inv_scale
     for i in range(self.n_layers):
       attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i],

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/gpt_oss/implementations/tinybox_8xMI350X/dev_beam.sh

@@ -20,7 +20,7 @@ export ZERO_OPTIM=${ZERO_OPTIM:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="gptoss"
@@ -34,7 +34,7 @@ export SEED=${SEED:-5760}
 export DATA_SEED=${DATA_SEED:-5760}
 
 export JITBEAM=${JITBEAM:-3}
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 
 export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/gpt_oss/implementations/tinybox_8xMI350X/dev_run.sh

@@ -20,7 +20,7 @@ export ZERO_OPTIM=${ZERO_OPTIM:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="gptoss"
@@ -34,6 +34,6 @@ export SEED=${SEED:-$RANDOM}
 export DATA_SEED=${DATA_SEED:-5760}
 
 export JITBEAM=${JITBEAM:-3}
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 
 python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh

@@ -2,5 +2,4 @@
 export BENCHMARK=${BENCHMARK:-5}
 export EVAL_BS=0
 VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=${DEBUG:--0} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
-SRC="AMD"; [[ $DEV == NULL* ]] && SRC="NULL"
-[ "$BENCHMARK" -le 3 ] || python -m tinygrad.viz.cli -s "$SRC" -t --interval "train @ 2" "train @ 3"
+[ "$BENCHMARK" -le 3 ] || [[ $DEV == NULL* ]] || python -m tinygrad.viz.cli -s AMD -t --interval "train @ 2" "train @ 3"

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/gpt_oss/implementations/tinybox_8xMI350X/dev_beam.sh

@@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=${DEV:-AMD}
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-2}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export MXFP8=${MXFP8:-1}
+export ZERO_OPTIM=${ZERO_OPTIM:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="gptoss"
+export BASEDIR="/raid/datasets/c4-8b/"
+export EVAL_TARGET=3.34 EVAL_FREQ=12288
+export END_LR="4e-5" WARMUP_STEPS=128 MAX_STEPS=1200000
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-5760}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+if [ -z "$FULL_LAYERS" ]; then
+  export LAYERS=${LAYERS:-2}
+fi
+
+python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/gpt_oss/implementations/tinybox_8xMI350X/dev_run.sh

@@ -0,0 +1,39 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=${DEV:-AMD}
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-0}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export MXFP8=${MXFP8:-1}
+export ZERO_OPTIM=${ZERO_OPTIM:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="gptoss"
+export BASEDIR="/raid/datasets/c4-8b/"
+export EVAL_TARGET=3.34 EVAL_FREQ=12288
+export END_LR="4e-5" WARMUP_STEPS=128 MAX_STEPS=1200000
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-$RANDOM}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh

@@ -0,0 +1,54 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=${DEV:-AMD}
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-2}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export FP8=${FP8:-1}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export FAST_CE=${FAST_CE:-1}
+export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
+export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
+export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
+export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
+export SPLIT_W13=${SPLIT_W13:-0}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-5760}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+if [ -z "$FULL_LAYERS" ]; then
+  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
+fi
+
+python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam_mp.sh

@@ -0,0 +1,54 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=${DEV:-AMD}
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-2}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export FP8=${FP8:-1}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export FAST_CE=${FAST_CE:-0}
+export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
+export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
+export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
+export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
+export SPLIT_W13=${SPLIT_W13:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-1} MP=${MP:-8} BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-5760}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+if [ -z "$FULL_LAYERS" ]; then
+  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
+fi
+
+python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh

@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=${DEV:-AMD}
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-0}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export FP8=${FP8:-1}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export FAST_CE=${FAST_CE:-1}
+export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
+export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
+export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
+export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
+export SPLIT_W13=${SPLIT_W13:-0}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-$RANDOM}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run_mp.sh

@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=${DEV:-AMD}
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-0}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
+export FP8=${FP8:-1}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
+export FAST_CE=${FAST_CE:-0}
+export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
+export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
+export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
+export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
+export SPLIT_W13=${SPLIT_W13:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-1} MP=${MP:-8} BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-32}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-$RANDOM}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+export BENCHMARK=${BENCHMARK:-5}
+export EVAL_BS=0
+VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=${DEBUG:--0} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
+[ "$BENCHMARK" -le 3 ] || [[ $DEV == NULL* ]] || python -m tinygrad.viz.cli -s AMD -t --interval "train @ 2" "train @ 3"

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/run_and_time.sh

@@ -0,0 +1,58 @@
+#!/usr/bin/env bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="."
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
+export DEV=AMD
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export HK_FLASH_ATTENTION=1
+export ALL2ALL=1
+export LATE_ALLREDUCE=0
+export USE_ATOMICS=1
+export ASM_GEMM=1
+export WQKV=1
+export MASTER_WEIGHTS=1
+export FP8=1
+export ALLREDUCE_CAST=1
+export FAST_CE=1
+export FUSED_INPUT_QUANTIZE=1
+export FUSED_GRAD_QUANTIZE=1
+export FUSED_ADD_NORM_MUL_QUANTIZE=1
+export FUSED_SILU_W13=1
+export SPLIT_W13=0
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=8 MP=1 BS=16 EVAL_BS=8 GRADIENT_ACC_STEPS=2
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=8B
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=8192
+
+export SEED=$RANDOM
+export DATA_SEED=$SEED
+
+export JITBEAM=3
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+export LOGMLPERF=1
+
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="llama31_8b_8xMI350x_${DATETIME}_${SEED}.log"
+
+# beam
+FAKEDATA=1 BENCHMARK=10 INITMLPERF=1 LLAMA_LAYERS=2 python3 examples/mlperf/model_train.py | tee "$LOGFILE"
+
+# run
+RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a "$LOGFILE"

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/startup_walltime.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+export BENCHMARK=5
+export EVAL_BS=0
+export FAKEDATA=1
+export NULL_ALLOW_COPYOUT=1
+export HIP_VISIBLE_DEVICES=""
+export DEV=NULL:HIP:gfx950
+export JITBEAM=0
+export LLAMA_LAYERS=${LLAMA_LAYERS:-"2"}
+time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh

examples/mlperf/training_submission_v6.1/tinycorp/systems/tinybox_8xMI300X.json

@@ -0,0 +1,38 @@
+{
+    "submitter": "tinycorp",
+    "division": "closed",
+    "status": "Available on-premise",
+    "system_name": "tinybox 8xMI300X",
+    "number_of_nodes": "1",
+    "host_processors_per_node": "2",
+    "host_processor_model_name": "AMD EPYC 9354",
+    "host_processor_core_count": "32",
+    "host_processor_vcpu_count": "64",
+    "host_processor_frequency": "",
+    "host_processor_caches": "",
+    "host_processor_interconnect": "",
+    "host_memory_capacity": "2304GB",
+    "host_storage_type": "NVMe SSD",
+    "host_storage_capacity": "3x 4TB raid array",
+    "host_networking": "",
+    "host_networking_topology": "",
+    "host_memory_configuration": "24x 96GB DDR5",
+    "accelerators_per_node": "8",
+    "accelerator_model_name": "AMD Instinct MI300X 192GB HBM3",
+    "accelerator_host_interconnect": "PCIe 5.0 x16",
+    "accelerator_frequency": "",
+    "accelerator_on-chip_memories": "",
+    "accelerator_memory_configuration": "HBM3",
+    "accelerator_memory_capacity": "192GB",
+    "accelerator_interconnect": "",
+    "accelerator_interconnect_topology": "",
+    "cooling": "air",
+    "hw_notes": "",
+    "framework": "tinygrad, branch mlperf_training_v5.0",
+    "other_software_stack": {
+        "python": "3.10.16",
+        "ROCm": "3.0.0+94441cb"
+    },
+    "operating_system": "Ubuntu 24.04.1 LTS",
+    "sw_notes": ""
+  }

examples/mlperf/training_submission_v6.1/tinycorp/systems/tinybox_8xMI350X.json

@@ -0,0 +1,38 @@
+{
+    "submitter": "tinycorp",
+    "division": "closed",
+    "status": "Available on-premise",
+    "system_name": "tinybox 8xMI350X",
+    "number_of_nodes": "1",
+    "host_processors_per_node": "2",
+    "host_processor_model_name": "AMD EPYC 9575F",
+    "host_processor_core_count": "32",
+    "host_processor_vcpu_count": "64",
+    "host_processor_frequency": "",
+    "host_processor_caches": "",
+    "host_processor_interconnect": "",
+    "host_memory_capacity": "3072 GiB",
+    "host_storage_type": "NVMe SSD",
+    "host_storage_capacity": "4TB",
+    "host_networking": "",
+    "host_networking_topology": "",
+    "host_memory_configuration": "24x 128GB DDR5",
+    "accelerators_per_node": "8",
+    "accelerator_model_name": "AMD Instinct MI350X 288GB HBM3e",
+    "accelerator_host_interconnect": "PCIe 5.0 x16",
+    "accelerator_frequency": "",
+    "accelerator_on-chip_memories": "",
+    "accelerator_memory_configuration": "HBM3",
+    "accelerator_memory_capacity": "288GB",
+    "accelerator_interconnect": "",
+    "accelerator_interconnect_topology": "",
+    "cooling": "air",
+    "hw_notes": "",
+    "framework": "tinygrad, branch mlperf_training_v6.0",
+    "other_software_stack": {
+        "python": "3.12.3",
+        "ROCm": "7.1.1"
+    },
+    "operating_system": "Ubuntu 24.04.3 LTS",
+    "sw_notes": ""
+  }

examples/mlperf/training_submission_v6.1/tinycorp/systems/tinybox_green.json

@@ -0,0 +1,38 @@
+{
+  "submitter": "tinycorp",
+  "division": "closed",
+  "status": "Available on-premise",
+  "system_name": "tinybox green",
+  "number_of_nodes": "1",
+  "host_processors_per_node": "1",
+  "host_processor_model_name": "AMD EPYC 7532",
+  "host_processor_core_count": "32",
+  "host_processor_vcpu_count": "64",
+  "host_processor_frequency": "",
+  "host_processor_caches": "",
+  "host_processor_interconnect": "",
+  "host_memory_capacity": "128GB",
+  "host_storage_type": "NVMe SSD",
+  "host_storage_capacity": "4 TB raid array + 1 TB boot",
+  "host_networking": "",
+  "host_networking_topology": "",
+  "host_memory_configuration": "8x 16GB DDR4",
+  "accelerators_per_node": "6",
+  "accelerator_model_name": "NVIDIA GeForce RTX 4090",
+  "accelerator_host_interconnect": "PCIe 4.0 x16",
+  "accelerator_frequency": "",
+  "accelerator_on-chip_memories": "",
+  "accelerator_memory_configuration": "GDDR6X",
+  "accelerator_memory_capacity": "24GB",
+  "accelerator_interconnect": "",
+  "accelerator_interconnect_topology": "",
+  "cooling": "air",
+  "hw_notes": "",
+  "framework": "tinygrad, branch mlperf_training_v5.0",
+  "other_software_stack": {
+    "python": "3.10.12",
+    "CUDA": "12.4"
+  },
+  "operating_system": "Ubuntu 22.04.4",
+  "sw_notes": ""
+}

examples/mlperf/training_submission_v6.1/tinycorp/systems/tinybox_red.json

@@ -0,0 +1,37 @@
+{
+  "submitter": "tinycorp",
+  "division": "closed",
+  "status": "Available on-premise",
+  "system_name": "tinybox red",
+  "number_of_nodes": "1",
+  "host_processors_per_node": "1",
+  "host_processor_model_name": "AMD EPYC 7532",
+  "host_processor_core_count": "32",
+  "host_processor_vcpu_count": "64",
+  "host_processor_frequency": "",
+  "host_processor_caches": "",
+  "host_processor_interconnect": "",
+  "host_memory_capacity": "128GB",
+  "host_storage_type": "NVMe SSD",
+  "host_storage_capacity": "4 TB raid array + 1 TB boot",
+  "host_networking": "",
+  "host_networking_topology": "",
+  "host_memory_configuration": "8x 16GB DDR4",
+  "accelerators_per_node": "6",
+  "accelerator_model_name": "AMD Radeon RX 7900 XTX",
+  "accelerator_host_interconnect": "PCIe 4.0 x16",
+  "accelerator_frequency": "",
+  "accelerator_on-chip_memories": "",
+  "accelerator_memory_configuration": "GDDR6",
+  "accelerator_memory_capacity": "24GB",
+  "accelerator_interconnect": "",
+  "accelerator_interconnect_topology": "",
+  "cooling": "air",
+  "hw_notes": "",
+  "framework": "tinygrad, branch mlperf_training_v5.0",
+  "other_software_stack": {
+    "python": "3.10.12"
+  },
+  "operating_system": "Ubuntu 22.04.4",
+  "sw_notes": ""
+}

extra/gemm/amd_copy_matmul.py

@@ -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

extra/gemm/amd_flash_attention.py

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

extra/gemm/metal_uop_matmul.py

@@ -1,5 +1,5 @@
 from tinygrad import UOp, dtypes
-from tinygrad.uop.ops import AxisType, Ops, KernelInfo, AddrSpace
+from tinygrad.uop.ops import AxisType, KernelInfo, AddrSpace
 from extra.gemm.amd_uop_matmul import test_matmul
 
 N = 2048
@@ -27,11 +27,8 @@ def hand_spec_tc_cores():
   acc = acc[0].set(0.0)
   acc = acc[1].set(0.0)
 
-  # TODO: make this simple
-  wmma_arg = ('WMMA_8_8_8_float_float', (8, 8, 8), dtypes.float, dtypes.float, 'METAL', 32, (((3, 2),), ((3, 2),), ((3, 2),)), ())
-
   acc_load = UOp.stack(acc.after(gk)[0], acc.after(gk)[1])
-  out = UOp(Ops.WMMA, dtypes.float, (a_tc, b_tc, acc_load), arg=wmma_arg)
+  out = UOp.wmma(a_tc, b_tc, acc_load, (8, 8, 8), 'METAL', 32)
 
   end_loop = UOp.group(*[acc[i].store(out.index(i)) for i in range(2)]).end(gk)
 

extra/gemm/mi350x_uop_matmul.py

@@ -6,7 +6,7 @@ os.environ["AMD_LLVM"] = "0"
 from tinygrad import Tensor, Context, dtypes, UOp, GlobalCounters
 from tinygrad.helpers import DEBUG, getenv
 from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import AxisType, KernelInfo, Ops
+from tinygrad.uop.ops import AxisType, KernelInfo
 
 WARP_SIZE = 64
 
@@ -137,8 +137,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
     acc_load = acc_after[N_inner_loop, M_inner_loop]
 
     # do WMMA
-    wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
-    out = UOp(Ops.WMMA, dtypes.float, (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg)
+    out = UOp.wmma(Ar[M_inner_loop], Br[N_inner_loop], acc_load, (16, 16, 32), 'AMD', 64)
 
     # store back the acc
     acc_store = acc[N_inner_loop, M_inner_loop].store(out)
@@ -193,8 +192,7 @@ acc = acc[init_l:=UOp.range(4, 1)].set(0.0, end=init_l)
 
 # do the wmma
 acc_load = UOp.stack(*[acc.after(K_loop)[i] for i in range(4)])
-wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
-out = UOp(Ops.WMMA, dtypes.float, (A_in, B_in, acc_load), arg=wmma_arg)
+out = UOp.wmma(A_in, B_in, acc_load, (16, 16, 32), 'AMD', 64)
 
 # store back the acc
 acc = acc.after(UOp.group(*[acc[i].store(out.index(i)) for i in range(4)]).end(K_loop))

extra/gemm/mi350x_uop_matmul_2.py

@@ -6,7 +6,7 @@ os.environ["AMD_LLVM"] = "0"
 from tinygrad import Tensor, Context, dtypes, UOp, GlobalCounters
 from tinygrad.helpers import DEBUG, getenv
 from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import sint, AxisType, KernelInfo, Ops
+from tinygrad.uop.ops import AxisType, KernelInfo
 
 WARP_SIZE = 64
 
@@ -60,8 +60,7 @@ def compute_on_locals(acc:UOp, Asl:UOp, Bsl:UOp, rng:int, afters:tuple[UOp, ...]
   acc_load = acc_after[N_inner_loop, M_inner_loop]
 
   # do WMMA
-  wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
-  out = UOp(Ops.WMMA, dtypes.float, (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg)
+  out = UOp.wmma(Ar[M_inner_loop], Br[N_inner_loop], acc_load, (16, 16, 32), 'AMD', 64)
 
   # store back the acc
   acc_store = acc[N_inner_loop, M_inner_loop].store(out)

extra/hcq2/hcq2.py

@@ -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
+from tinygrad.helpers import to_tuple, round_up, partition, data64_le, panic, ContextVar
 from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer
 from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, track_rewrites, GroupOp
 from tinygrad.uop.symbolic import symbolic
@@ -19,6 +19,8 @@ HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQ2Compiled')
 # *****************
 # 0. helpers
 
+HCQ_RUNTIME_DEV = ContextVar("HCQ_RUNTIME_DEV", "CPU")
+
 HCQ_DEVS = frozenset(("AMD",))
 HCQ_P2P_DEVS = HCQ_DEVS | frozenset(("CPU",))
 HCQ_CACHE_TAGS = frozenset(("program", "systems", "template"))
@@ -56,7 +58,7 @@ def make_patch(buf:UOp, off:sint, val:UOp, dtype=None) -> UOp:
   return buf.index(UOp.const(dtypes.int, off // buf.dtype.itemsize)).store(val.simplify().cast(dtype or buf.dtype))
 
 def make_binary_patch(buf:UOp, blob:bytes) -> UOp:
-  data = UOp(Ops.BITCAST, buf.dtype, (UOp(Ops.BINARY, dtypes.uint8, src=(), arg=blob),))
+  data = UOp(Ops.BINARY, src=(), arg=blob).bitcast(buf.dtype)
   r = UOp.range(len(blob) // buf.dtype.itemsize, 0, dtype=dtypes.int, src=(buf, data))
   return buf.index(r).store(data.index(r).load()).end(r)
 
@@ -219,7 +221,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).wait(make_signal_value(devs).index(UOp.const(dtypes.int, 0)) - 1)
+  wait = (make_signal(devs).index(zero:=UOp.const(dtypes.int, 0)).load() >= make_signal_value(devs).index(zero) - 1).wait()
   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)])
 
@@ -237,7 +239,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.wait(val + dep.tag))
+        new_src.append((sig.index(UOp.const(dtypes.int, 0)).load() >= val + dep.tag).wait())
       s = s.src[0]
     new_src.append(s)
   return submit.replace(src=(q.replace(src=tuple(new_src)),))
@@ -303,7 +305,7 @@ pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION
 
 # *****************
 
-def make_addr_table(call:UOp, gaddrs:list[UOp], name:str):
+def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[dict[UOp, UOp], tuple[UOp, ...]]:
   bare = {g: g.replace(src=(unwrap_after(g.src[0]),)) for g in gaddrs}
 
   order = sorted(dedup(bare.values()), key=lambda g: ((b:=unwrap_mstack(g.buf_uop)[0]).arg.slot, to_tuple(b.tag)))
@@ -312,26 +314,23 @@ def make_addr_table(call:UOp, gaddrs:list[UOp], name:str):
   reads = {g: table.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(UOp.const(dtypes.int, slots[bare[g]])).load() for g in gaddrs}
   return reads, (table.after(*[make_patch(table, i * table.dtype.itemsize, addr) for addr, i in slots.items()]),) if slots else ()
 
-def rm_rt_getaddrs(call:UOp) -> UOp|None:
-  if not (gaddrs:=[u for u in call.src[0].toposort() if u.op is Ops.GETADDR]): return None
+def make_blob_bufs(call:UOp, blobs:list[UOp]) -> tuple[dict[UOp, UOp], tuple[UOp, ...]]:
+  bufs = {b: make_placeholder(call.arg.aux.device, b.max_numel(), b.dtype, "template") for b in blobs}
+  return bufs, tuple(buf.after(make_binary_patch(buf, b.src[0].arg)) for b,buf in bufs.items())
+
+def rm_rt_uops(call:UOp) -> UOp|None:
+  if not (rt_uops:=[u for u in call.src[0].toposort() if u.op is Ops.GETADDR or (u.op is Ops.BITCAST and u.src[0].op is Ops.BINARY)]): return None
+  gaddrs, blobs = partition(rt_uops, lambda u: u.op is Ops.GETADDR)
   inputs, internals = partition(gaddrs, lambda g: all(x.op is Ops.PARAM and x.tag is None for x in unwrap_mstack(g.buf_uop)))
   runtimes, systems = partition(internals, lambda g: any(x.tag in {"program", "kernargs", "cmdbuf"} for x in unwrap_mstack(g.buf_uop)))
 
   # exec fills the inputs table with the input addresses every run, so it has no fill patches
-  (input_reads, _), (rt_reads, rt_fills), (sys_reads, sys_fills) = (make_addr_table(call, gs, name) for gs, name in
-    ((inputs, "inputs"), (runtimes, "runtime"), (systems, "systems")))
-  return call.replace(src=(call.src[0].substitute(input_reads | rt_reads | sys_reads), *call.src[1:], *rt_fills, *sys_fills),
+  (reads, _), *tables = [make_addr_table(call, gs, n) for gs,n in ((inputs, "inputs"), (runtimes, "runtime"), (systems, "systems"))] + \
+                        [make_blob_bufs(call, blobs)]
+  reads, fills = reads | {k:v for r,_ in tables for k,v in r.items()}, [f for _,fs in tables for f in fs]
+  return call.replace(src=(call.src[0].substitute(reads), *call.src[1:], *fills),
                       arg=replace(call.arg, aux=replace(call.arg.aux, input_idxs=tuple(sorted(dedup(g.buf_uop.arg.slot for g in inputs))))))
-pm_rm_rt_getaddrs = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), rm_rt_getaddrs)])
-
-# *****************
-
-def rm_rt_binaries(call:UOp) -> UOp|None:
-  if not (blobs:=[u for u in call.src[0].toposort() if u.op is Ops.BITCAST and u.src[0].op is Ops.BINARY]): return None
-  blob_bufs = {blob: make_placeholder(call.arg.aux.device, blob.max_numel(), blob.dtype, "template") for blob in blobs}
-  fills = [buf.after(make_binary_patch(buf, blob.src[0].arg)) for blob, buf in blob_bufs.items()]
-  return call.replace(src=(call.src[0].substitute(blob_bufs), *call.src[1:], *fills))
-pm_rm_rt_binaries = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), rm_rt_binaries)])
+pm_rm_rt_uops = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), rm_rt_uops)])
 
 # *****************
 
@@ -383,7 +382,7 @@ pm_pack_placeholders = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNC
 # 8. callify hcq programs
 
 pm_callify_hcq = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="hcq", src=(UPat(Ops.SINK),), name="cf"),
-  lambda cf: cf.replace(src=(to_program(cf.src[0].replace(arg=KernelInfo("hcq_submit"), tag=1), Device["CPU"].renderer),)))])
+  lambda cf: cf.replace(src=(to_program(cf.src[0].replace(arg=KernelInfo("hcq_submit"), tag=1), Device[HCQ_RUNTIME_DEV.value].renderer),)))])
 
 hcq_compile_cache:dict[tuple[bytes, bool], UOp] = {}
 
@@ -411,8 +410,7 @@ 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_getaddrs, walk=True, name="replace rt getaddrs")
-    linear = graph_rewrite(linear, pm_rm_rt_binaries, walk=True, name="replace rt binaries")
+    linear = graph_rewrite(linear, pm_rm_rt_uops, walk=True, name="replace rt uops")
     linear = graph_rewrite(linear, pm_replace_params, walk=True, name="replace with args")
 
     # and compile it
@@ -425,7 +423,7 @@ def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None, jit=False) -> UOp:
 
 def bufferize_buf(ctx:bool, buf:UOp) -> UOp|None:
   if buf.tag is None: return None
-  return make_mstack(tuple(UOp.from_buffer((dv:=Device[dev]).pm_bufferize.rewrite(buf, ctx=(dv, ctx)), "CPU") for dev in to_tuple(buf.device)))
+  return make_mstack(tuple(UOp.from_buffer((dv:=Device[dev]).pm_bufferize.rewrite(buf, ctx=(dv, ctx)), HCQ_RUNTIME_DEV.value) for dev in to_tuple(buf.device)))
 pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, name="buf"), bufferize_buf)])
 
 # *****************
@@ -440,7 +438,8 @@ def fold_binary(buf:UOp, blob:UOp) -> UOp:
 
 def fold_const_store(buf:UOp, off:UOp, val:UOp) -> UOp:
   for b, v in zip((bs:=mb.bufs if isinstance((mb:=buf.buffer), MultiBuffer) else (mb,)), val.src if val.op is Ops.STACK else (val,)*len(bs)):
-    struct.pack_into(f'<{v.dtype.fmt}', b.ensure_allocated()._buf.cpu_view().mv.cast('B'), off.arg * buf.dtype.itemsize, truncate[v.dtype](v.arg))
+    data = struct.pack(f'<{v.dtype.fmt}', truncate[v.dtype](v.arg))
+    b.ensure_allocated()._buf.cpu_view().view(offset=off.arg * buf.dtype.itemsize, size=len(data), fmt='B')[:] = data
   return UOp(Ops.NOOP)
 
 def resolve_getaddr(buf:UOp, g:UOp) -> UOp:
@@ -571,6 +570,10 @@ class HCQ2Buffer:
   def base(self) -> HCQ2Buffer: return self._base or self
 
 class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
+  def _as_buffer(self, buf:HCQ2Buffer) -> memoryview:
+    self.dev.synchronize()
+    return buf.cpu_view().mv
+
   def _map(self, buf:HCQ2Buffer) -> HCQ2Buffer:
     if not hasattr(self, '_do_map'): raise NotImplementedError("map failed: no method implemented")
     return self._do_map(buf)
@@ -586,24 +589,3 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
     self.dev.iface.free(mb)
 
   def _offset(self, buf, size:int, offset:int) -> HCQ2Buffer: return buf.offset(offset=offset, size=size)
-
-  def _wrap(self, dev:str, sz:int, opaque:HCQ2Buffer) -> Buffer:
-    return Buffer(dev, sz, dtypes.uint8, opaque=opaque, options=BufferSpec(external_ptr=1))
-
-  def _copy(self, dst:Buffer, src:Buffer):
-    from tinygrad.engine.realize import run_linear
-    du, su = UOp.from_buffer(dst), UOp.from_buffer(src)
-    run_linear(UOp(Ops.LINEAR, src=(su.param_like(1).copy_to_device(dst.device).call(du, su),)), update_stats=True)
-
-  def _copyin(self, dest:HCQ2Buffer, src:memoryview):
-    s = Buffer(self.dev.device, len(src), dtypes.uint8, options=BufferSpec(host=True), preallocate=True)
-    s._buf.cpu_view()[:len(src)] = src
-    self._copy(self._wrap(self.dev.device, len(src), dest), s)
-
-  def _copyout(self, dest:memoryview, src:HCQ2Buffer):
-    d = Buffer(self.dev.device, len(dest), dtypes.uint8, options=BufferSpec(host=True), preallocate=True)
-    self._copy(d, self._wrap(self.dev.device, len(dest), src))
-    self.dev.synchronize()
-    dest[:] = d._buf.cpu_view()[:len(dest)]
-
-  # def _as_buffer(self, buf): return buf.cpu_view().mv

extra/hcq2/ops_amd2.py

@@ -90,7 +90,7 @@ def memory_barrier(ctx):
                  reg_done=getattr(ctx.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff),
     acquire_mem(ctx)))
 
-def pm4_wait(ctx, dst, val): return wait_reg_mem(ctx, val, mem=make_getaddr(dst, ctx.devs))
+def pm4_wait(ctx, x, y): return wait_reg_mem(ctx, y, mem=make_getaddr(x.buf_uop, ctx.devs))
 
 def pm4_barrier(ctx): return memory_barrier(ctx)
 
@@ -138,7 +138,7 @@ def pm4_program(ctx, call, prg):
 pm_pm4_opsel = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), pm4_program),
 
-  (UPat(Ops.WAIT, src=(UPat(name="dst"), UPat(name="val"))), pm4_wait),
+  (UPat(Ops.WAIT, src=(UPat.var("x") >= UPat.var("y"),)), pm4_wait),
   (UPat(Ops.BARRIER), pm4_barrier),
   (UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", src=(UPat(name="dst"),)), pm4_timestamp),
   (UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
@@ -184,10 +184,10 @@ def sdma_copy(ctx, call):
      ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz - off, ctx.max_copy_size) - 1), 0,
      *data64_le(src_addr + off), *data64_le(dst_addr + off)) for off in range(0, sz, ctx.max_copy_size)]))
 
-def sdma_wait(ctx, dst, val):
+def sdma_wait(ctx, x, y):
   op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \
      | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1)
-  return make_ins(SDMAOps.POLL_REGMEM, op, *data64_le(make_getaddr(dst, ctx.devs)), val, 0xffffffff,
+  return make_ins(SDMAOps.POLL_REGMEM, op, *data64_le(make_getaddr(x.buf_uop, ctx.devs)), y, 0xffffffff,
     ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))
 
 def sdma_store(ctx, dst, val):
@@ -203,7 +203,7 @@ pm_sdma_opsel = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy),
 
   (UPat(Ops.BARRIER), lambda: UOp(Ops.NOOP, dtypes.void, ())),
-  (UPat(Ops.WAIT, src=(UPat(name="dst"), UPat(name="val"))), sdma_wait),
+  (UPat(Ops.WAIT, src=(UPat.var("x") >= UPat.var("y"),)), sdma_wait),
   (UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", src=(UPat(name="dst"),)), sdma_timestamp),
   (UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), sdma_store),
 ])
@@ -536,7 +536,7 @@ class AMDDevice(HCQ2Compiled):
 
   timestamp_divider = 100.0  # AMD GPU clock: ticks/us
 
-  ifaces = [KFDIface, PCIIface]
+  ifaces = [KFDIface, PCIIface, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface)]
 
   def is_am(self) -> bool: return isinstance(self.iface, (PCIIface,))
   def is_usb(self) -> bool: return False

extra/thunder/amd/fa.py

@@ -16,6 +16,96 @@ def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None
   axis = ref.uop.axis if axis is None else axis
   return Tensor(Tensor.invalids(*shape, dtype=dtype, device=ref.device).uop.multi(axis), dtype=dtype, device=ref.device)
 
+@functools.cache
+def custom_fused_qkv_rope_forward(q:UOp, k:UOp, v:UOp, xqkv:UOp, freqs_cis:UOp,
+                                  device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
+  code = (pathlib.Path(__file__).parent / "fused_qkv_rope.cpp").read_text()
+  threads = 256
+  thread_idx = UOp.special(threads, "lidx0")
+  block_idx_x, block_idx_y = UOp.special(B, "gidx0"), UOp.special(N, "gidx1")
+  sink = UOp.sink(q.base, k.base, v.base, xqkv.base, freqs_cis.base, thread_idx, block_idx_x, block_idx_y,
+                  arg=KernelInfo(name="fused_qkv_rope_forward"))
+  compile_args = ["-std=c++20", "-ffast-math", f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}",
+                  f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DTHREADS_PER_BLOCK={threads}"]
+  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
+
+@functools.cache
+def custom_fused_qkv_rope_backward(dxqkv:UOp, dq:UOp, dk:UOp, dv:UOp, freqs_cis:UOp,
+                                   device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
+  assert (B, N, H, H_KV, D) == (2, 8192, 32, 8, 128)
+  code = (pathlib.Path(__file__).parent / "fused_qkv_rope_bwd.cpp").read_text()
+  threads = 256
+  thread_idx = UOp.special(threads, "lidx0")
+  gsz = (B, N // 64, H + 2 * H_KV)
+  block_idx_x, block_idx_y, block_idx_z = (UOp.special(x, f"gidx{i}") for i, x in enumerate(gsz))
+  sink = UOp.sink(dxqkv.base, dq.base, dk.base, dv.base, freqs_cis.base, thread_idx, block_idx_x, block_idx_y, block_idx_z,
+                  arg=KernelInfo(name="fused_qkv_rope_backward"))
+  compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math", f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}",
+                  f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DTHREADS_PER_BLOCK={threads}"]
+  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
+
+def _fa_native_grads(dq:UOp, dk:UOp, dv:UOp) -> tuple[UOp, UOp, UOp]|None:
+  def unwrap_partial(x:UOp) -> UOp|None:
+    expected = (Ops.CAST, Ops.REDUCE, Ops.PERMUTE, Ops.CAST, Ops.RESHAPE, Ops.AFTER)
+    for op in expected:
+      if x.op is not op: return None
+      if op is not Ops.AFTER: x = x.src[0]
+    return x
+  dq_native, dk_partial, dv_partial = dq.base, unwrap_partial(dk), unwrap_partial(dv)
+  if dq_native.op is not Ops.AFTER or dk_partial is None or dv_partial is None: return None
+  B, N, H, D, H_KV = dq.shape[0], dq.shape[1], dq.shape[2], dq.shape[3], dk.shape[2]
+  heads_per_wg = 2 if D == 128 and (H // H_KV) % 2 == 0 else 1
+  partials = (H // H_KV) // heads_per_wg
+  if dq_native.shape != (B, H, N, D) or dk_partial.shape != (B * partials, N, H_KV, D) or dv_partial.shape != dk_partial.shape: return None
+  return dq_native, dk_partial, dv_partial
+
+def _fused_qkv_rope_grad(dq_u:UOp, dk_u:UOp, dv_u:UOp, call:UOp) -> tuple[None, None, None, UOp, None]:
+  dq, dk, dv = Tensor(dq_u, device=dq_u.device), Tensor(dk_u, device=dk_u.device), Tensor(dv_u, device=dv_u.device)
+  xqkv_u, freqs_u = call.src[4], call.src[5]
+  xqkv, freqs_cis = Tensor(xqkv_u, device=xqkv_u.device), Tensor(freqs_u, device=freqs_u.device)
+  B, N, _ = xqkv.shape
+  H, H_KV, D = dq.shape[2], dk.shape[2], dq.shape[3]
+  num_devices = len(xqkv.device) if isinstance(xqkv.device, tuple) else 1
+  is_dp, is_mp = xqkv.uop.axis == 0, xqkv.uop.axis == 2
+  B_local = B // num_devices if is_dp else B
+  H_local = H // num_devices if is_mp else H
+  H_KV_local = H_KV // num_devices if is_mp else H_KV
+  single_device = xqkv.device[0] if isinstance(xqkv.device, tuple) else xqkv.device
+  arch = Device[single_device].renderer.target.arch
+  fa_native = _fa_native_grads(dq_u, dk_u, dv_u)
+  assert fa_native is not None, "fused QKV RoPE backward requires native Flash Attention gradients"
+  dq, dk, dv = (Tensor(x, device=x.device) for x in fa_native)
+  dxqkv = _sharded_empty_like(xqkv, axis=xqkv.uop.axis if isinstance(xqkv.device, tuple) else None)
+  fxn = functools.partial(custom_fused_qkv_rope_backward, device=single_device, arch=arch,
+                          B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D)
+  dxqkv = Tensor.custom_kernel(dxqkv, dq, dk, dv, freqs_cis, fxn=fxn)[0]
+  return None, None, None, dxqkv.uop, None
+
+def fused_qkv_rope(xqkv:Tensor, freqs_cis:Tensor, n_heads:int, n_kv_heads:int, head_dim:int) -> tuple[Tensor, Tensor, Tensor]:
+  B, N, packed_dim = xqkv.shape
+  assert packed_dim == n_kv_heads * (n_heads // n_kv_heads + 2) * head_dim
+  assert freqs_cis.dtype == dtypes.bfloat16, f"fused QKV RoPE requires bfloat16 frequencies, got {freqs_cis.dtype}"
+  assert freqs_cis.shape == (1, freqs_cis.shape[1], 1, head_dim // 2, 2) and freqs_cis.shape[1] >= N, \
+    f"invalid RoPE frequency shape {freqs_cis.shape} for sequence length {N} and head dimension {head_dim}"
+  num_devices = len(xqkv.device) if isinstance(xqkv.device, tuple) else 1
+  is_dp, is_mp = xqkv.uop.axis == 0, xqkv.uop.axis == 2
+  B_local = B // num_devices if is_dp else B
+  H_local = n_heads // num_devices if is_mp else n_heads
+  H_KV_local = n_kv_heads // num_devices if is_mp else n_kv_heads
+  assert (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)
 
@@ -31,8 +121,9 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
 
     dq = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t)
     GROUP_SIZE = H_local // H_KV_local
-    dk_partial = _sharded_empty((B * GROUP_SIZE, N, H_KV, D), xk, axis=shard_axis)
-    dv_partial = _sharded_empty((B * GROUP_SIZE, N, H_KV, D), xv, axis=shard_axis)
+    HEADS_PER_WG = 2 if D == 128 and GROUP_SIZE % 2 == 0 else 1
+    dk_partial = _sharded_empty((B * GROUP_SIZE // HEADS_PER_WG, N, H_KV, D), xk, axis=shard_axis)
+    dv_partial = _sharded_empty((B * GROUP_SIZE // HEADS_PER_WG, N, H_KV, D), xv, axis=shard_axis)
 
     # delta_vec = (do * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
     delta_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
@@ -46,8 +137,8 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
       dq = dq.reshape(B, H, N//16, 4, 2, 2, D//32, 4, 4, 2).permute(0, 1, 2, 7, 8, 3, 4, 6, 5, 9).reshape(B, H, N, D).transpose(1, 2)
 
     # reduce partial dK/dV across GROUP_SIZE query heads
-    dk = dk_partial.reshape(B, GROUP_SIZE, N, H_KV, D).sum(1)
-    dv = dv_partial.reshape(B, GROUP_SIZE, N, H_KV, D).sum(1)
+    dk = dk_partial.reshape(B, GROUP_SIZE // HEADS_PER_WG, N, H_KV, D).sum(1)
+    dv = dv_partial.reshape(B, GROUP_SIZE // HEADS_PER_WG, N, H_KV, D).sum(1)
 
     if not has_sink: return None, None, dq.uop, dk.uop, dv.uop
     sinks = Tensor(ker.src[6], device=ker.src[6].device)
@@ -160,9 +251,11 @@ def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_ve
                   f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}"]
 
   BLOCK_SIZE_KV = 256
+  GROUP_SIZE = H // H_KV
+  HEADS_PER_WG = 2 if D == 128 and GROUP_SIZE % 2 == 0 else 1
   NUM_WARPS = 4
   NUM_THREADS = 64 * NUM_WARPS
-  gsz = (H, N // BLOCK_SIZE_KV, B)
+  gsz = (H // HEADS_PER_WG, N // BLOCK_SIZE_KV, B)
   lsz = (NUM_THREADS, 1, 1)
   threadIdx_x = UOp.special(lsz[0], "lidx0")
   blockIdx_x, blockIdx_y, blockIdx_z = UOp.special(gsz[0], "gidx0"), UOp.special(gsz[1], "gidx1"), UOp.special(gsz[2], "gidx2")

extra/thunder/amd/fa_bwd_causal.cpp

@@ -28,6 +28,7 @@ constexpr int ATTN_H_KV = 8; // number of key/value heads (for GQA)
 #endif
 
 constexpr int GROUP_SIZE = ATTN_H / ATTN_H_KV; // queries per KV head group
+constexpr int HEADS_PER_WG = (ATTN_D == 128 && GROUP_SIZE % 2 == 0) ? 2 : 1;
 
 #ifndef ATTN_N
 constexpr int ATTN_N = 1024; // sequence length
@@ -53,7 +54,7 @@ using namespace kittens;
 using _gl_QdO  = gl<bf16, ATTN_B, ATTN_N, ATTN_H, ATTN_D>;
 using _gl_KV   = gl<bf16, ATTN_B, ATTN_N, ATTN_H_KV, ATTN_D>;
 using _gl_dQ   = gl<bf16, ATTN_B, ATTN_H, ATTN_N, ATTN_D>;
-using _gl_dKV  = gl<bf16, ATTN_B * GROUP_SIZE, ATTN_N, ATTN_H_KV, ATTN_D>;
+using _gl_dKV  = gl<bf16, ATTN_B * (GROUP_SIZE / HEADS_PER_WG), ATTN_N, ATTN_H_KV, ATTN_D>;
 using _gl_Lvec = gl<float, ATTN_B, ATTN_H, 1, ATTN_N>;
 
 template<int D> struct attn_bwd_combined_globals {
@@ -63,7 +64,7 @@ template<int D> struct attn_bwd_combined_globals {
   _gl_dQ dQg;
   _gl_dKV dKg, dVg;
   _gl_Lvec L_vec, delta_vec;
-  dim3 grid() { return dim3(ATTN_H, (ATTN_N / BLOCK_SIZE_KV), ATTN_B); }
+  dim3 grid() { return dim3(ATTN_H / HEADS_PER_WG, (ATTN_N / BLOCK_SIZE_KV), ATTN_B); }
   dim3 block() { return dim3(NUM_THREADS); }
   size_t dynamic_shared_memory() { return MAX_SHARED_MEMORY; }
 };
@@ -71,7 +72,7 @@ template<int D> struct attn_bwd_combined_globals {
 template<int D> __launch_bounds__(NUM_THREADS, 1)
 __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr, bf16 *dO_ptr, bf16 *Q_ptr, bf16 *K_ptr, bf16 *V_ptr, float *L_vec_ptr, float *delta_vec_ptr) {
 
-  const int q_head_idx_fixed = blockIdx.x;  // This is the query head index [0, ATTN_H)
+  const int q_head_idx_fixed = blockIdx.x * HEADS_PER_WG;  // First query head handled by this workgroup.
   const int kv_head_idx = q_head_idx_fixed / GROUP_SIZE;
   const int q_head_in_group = q_head_idx_fixed % GROUP_SIZE;
   const int seq_idx = blockIdx.y;
@@ -88,7 +89,7 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
   // first Q step that can overlap this K_span:
   const int first_step = max(0, k_start_min / STEP_QO);
   const int num_steps_per_head = total_steps_per_head - first_step;
-  const int num_steps = num_steps_per_head;
+  const int num_steps = num_steps_per_head * HEADS_PER_WG;
   const int k_pos = j * WARP_SIZE_KV;
 
   constexpr float L_SCALE_FACTOR = 1.44269504089f;
@@ -270,12 +271,12 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
   if (num_steps > 1) {
     // Prologue
     {
-      const int q_head_idx = (0) / num_steps_per_head + first_q_head;
-      const int q_seq_idx = ((0) % num_steps_per_head) + first_step;
+      const int q_head_idx = first_q_head;
+      const int q_seq_idx = first_step;
       const int q_pos = q_seq_idx * STEP_QO;
 
-      const int next_q_head_idx = (0 + 1) / num_steps_per_head + first_q_head;
-      const int next_q_seq_idx = ((0 + 1) % num_steps_per_head) + first_step;
+      const int next_q_head_idx = first_q_head;
+      const int next_q_seq_idx = first_step + 1;
 
       // dot slice 0
       {
@@ -1332,15 +1333,18 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
 
     // 9. for 1 <= i <= T_r (1024 / 32 = 32)  
     for (int i = 1; i < num_steps - 1; ++i, tic ^= 1, toc ^= 1) {
-      const int last_q_head_idx = (i - 1) / num_steps_per_head + first_q_head;
-      const int last_q_seq_idx = ((i - 1) % num_steps_per_head) + first_step;
+      const int last_head_offset = (i - 1) >= num_steps_per_head;
+      const int last_q_head_idx = last_head_offset + first_q_head;
+      const int last_q_seq_idx = i - 1 - last_head_offset * num_steps_per_head + first_step;
 
-      const int q_head_idx = i / num_steps_per_head + first_q_head;
-      const int q_seq_idx = (i % num_steps_per_head) + first_step;
+      const int head_offset = i >= num_steps_per_head;
+      const int q_head_idx = head_offset + first_q_head;
+      const int q_seq_idx = i - head_offset * num_steps_per_head + first_step;
       const int q_pos = q_seq_idx * STEP_QO;
 
-      const int next_q_head_idx = (i + 1) / num_steps_per_head + first_q_head;
-      const int next_q_seq_idx = ((i + 1) % num_steps_per_head) + first_step;
+      const int next_head_offset = (i + 1) >= num_steps_per_head;
+      const int next_q_head_idx = next_head_offset + first_q_head;
+      const int next_q_seq_idx = i + 1 - next_head_offset * num_steps_per_head + first_step;
 
       // dot slice 0
       {
@@ -2378,11 +2382,11 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
     }
   }
 
-  const int last_q_head_idx = (num_steps - 2) / num_steps_per_head + first_q_head;
-  const int last_q_seq_idx = ((num_steps - 2) % num_steps_per_head) + first_step;
+  const int last_q_head_idx = first_q_head + HEADS_PER_WG - 1;
+  const int last_q_seq_idx = first_step + num_steps_per_head - 2;
 
-  const int q_head_idx = (num_steps - 1) / num_steps_per_head + first_q_head;
-  const int q_seq_idx = ((num_steps - 1) % num_steps_per_head) + first_step;
+  const int q_head_idx = first_q_head + HEADS_PER_WG - 1;
+  const int q_seq_idx = first_step + num_steps_per_head - 1;
   const int q_pos = q_seq_idx * STEP_QO;
   // Epilogue
   {
@@ -3407,14 +3411,14 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
     }
   }
 
-  store<1>(g.dVg, dV_j, {batch_idx * GROUP_SIZE + q_head_in_group, 0, kv_head_idx, 0}, {0, j, 0, 0});
+  store<1>(g.dVg, dV_j, {batch_idx * (GROUP_SIZE / HEADS_PER_WG) + q_head_in_group / HEADS_PER_WG, 0, kv_head_idx, 0}, {0, j, 0, 0});
   __builtin_amdgcn_s_waitcnt(0);
   __builtin_amdgcn_s_barrier();
 
   // We first copy dV_j_T from accumulator GPRs to vector GPRs and then perform the store
   accvgpr_read(dV_j_T, dK_j_T);
   mul(dV_j_T, dV_j_T, dP_SCALE_FACTOR);
-  store<1>(g.dKg, dV_j, {batch_idx * GROUP_SIZE + q_head_in_group, 0, kv_head_idx, 0}, {0, j, 0, 0});
+  store<1>(g.dKg, dV_j, {batch_idx * (GROUP_SIZE / HEADS_PER_WG) + q_head_in_group / HEADS_PER_WG, 0, kv_head_idx, 0}, {0, j, 0, 0});
 
   // Write out final dQ_i slice
   mul(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
@@ -3422,6 +3426,3 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
 }
 
 template __global__ void attend_bwd_combined_ker<ATTN_D>(bf16*, bf16*, bf16*, bf16*, bf16*, bf16*, bf16*, float*, float*);
-
-
-

extra/thunder/amd/fused_qkv_rope.cpp

@@ -0,0 +1,69 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+
+#ifndef ATTN_B
+#define ATTN_B 2
+#endif
+#ifndef ATTN_N
+#define ATTN_N 8192
+#endif
+#ifndef ATTN_H
+#define ATTN_H 32
+#endif
+#ifndef ATTN_H_KV
+#define ATTN_H_KV 8
+#endif
+#ifndef ATTN_D
+#define ATTN_D 128
+#endif
+#ifndef THREADS_PER_BLOCK
+#define THREADS_PER_BLOCK 256
+#endif
+
+constexpr int GROUP_SIZE = ATTN_H / ATTN_H_KV;
+constexpr int HALF_D = ATTN_D / 2;
+constexpr int PACKED_D = (GROUP_SIZE + 2) * ATTN_D;
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_BLOCK) void
+fused_qkv_rope_forward(
+    __hip_bfloat16*       __restrict__ q,
+    __hip_bfloat16*       __restrict__ k,
+    __hip_bfloat16*       __restrict__ v,
+    const __hip_bfloat16* __restrict__ xqkv,
+    const __hip_bfloat16* __restrict__ freqs_cis) {
+  const int b = blockIdx.x;
+  const int n = blockIdx.y;
+  const int bn = b * ATTN_N + n;
+  const int packed_bn = bn * ATTN_H_KV * PACKED_D;
+  const int q_bn = bn * ATTN_H * ATTN_D;
+  const int kv_bn = bn * ATTN_H_KV * ATTN_D;
+
+  if (threadIdx.x < HALF_D) {
+    const int pair = threadIdx.x;
+    const int even = pair << 1;
+    const float c = static_cast<float>(freqs_cis[((n * HALF_D + pair) * 2) + 0]);
+    const float s = static_cast<float>(freqs_cis[((n * HALF_D + pair) * 2) + 1]);
+
+    for (int kvh = 0; kvh < ATTN_H_KV; kvh++) {
+      const int base = packed_bn + kvh * PACKED_D;
+
+      for (int rep = 0; rep < GROUP_SIZE; rep++) {
+        const int qbase = base + rep * ATTN_D;
+        const int h = kvh * GROUP_SIZE + rep;
+        const float a = static_cast<float>(xqkv[qbase + even]);
+        const float bb = static_cast<float>(xqkv[qbase + even + 1]);
+        const int out = q_bn + h * ATTN_D + even;
+        q[out] = static_cast<__hip_bfloat16>(a * c - bb * s);
+        q[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c);
+      }
+
+      const float a = static_cast<float>(xqkv[base + GROUP_SIZE * ATTN_D + even]);
+      const float bb = static_cast<float>(xqkv[base + GROUP_SIZE * ATTN_D + even + 1]);
+      const int out = kv_bn + kvh * ATTN_D + even;
+      k[out] = static_cast<__hip_bfloat16>(a * c - bb * s);
+      k[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c);
+      v[out] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even];
+      v[out + 1] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even + 1];
+    }
+  }
+}

extra/thunder/amd/fused_qkv_rope_bwd.cpp

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

extra/thunder/amd/gemm_fp8.cpp

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

extra/thunder/amd/gemm_fp8_atb.cpp

@@ -134,7 +134,7 @@ __global__ __launch_bounds__(512, 2) void hk_fp8_atb_gemm(bf16 *C_ptr, fp8e4m3 *
 
     int wgid = blockIdx.x;
     const int WGM = 8;
-    wgid = chiplet_transform_chunked(wgid, total_blocks_needed, NUM_XCDS, 64);
+    wgid = chiplet_transform_chunked(wgid, total_blocks_needed, NUM_XCDS, 32);
 
     const int num_wgid_in_group = WGM * blocks_per_col;
     int group_id = wgid / num_wgid_in_group;

extra/thunder/tiny/tk/group.py

@@ -1,7 +1,6 @@
 import math
 from typing import cast, Callable
-from tinygrad import dtypes
-from tinygrad.uop.ops import AxisType, UOp, Ops
+from tinygrad.uop.ops import AxisType, UOp
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import prod
 
@@ -75,9 +74,9 @@ class Group:
 
     a_base_shape = cast(RT, a).base_shape
     if a_base_shape.cols == 16:
-      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_dims = (16, 16, 16)
     elif a_base_shape.cols == 32:
-      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_dims = (16, 16, 32)
     else: raise NotImplementedError(f"mma_AB not implemented for {a_base_shape.cols=}")
 
     for height in self.ker.range(c.shape[-3], track=False):
@@ -92,7 +91,7 @@ class Group:
           else: raise NotImplementedError(f"mma_AB not implemented for {a_base_shape.cols=}")
           d_in = UOp.stack(*[c[height, width, i] for i in range(4)])
 
-          out = UOp(Ops.WMMA, dtypes.float32, (a_in, b_in, d_in), arg=wmma_arg)
+          out = UOp.wmma(a_in, b_in, d_in, wmma_dims, 'AMD', 64)
           c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
           c_store = UOp.group(*c_i).end(height, width, inner)
 
@@ -105,9 +104,9 @@ class Group:
 
     a_base_shape = cast(RT, a).base_shape
     if a_base_shape.cols == 16:
-      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_dims = (16, 16, 16)
     elif a_base_shape.cols == 32:
-      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_dims = (16, 16, 32)
     else: raise NotImplementedError(f"mma_ABt not implemented for {a_base_shape.cols=}")
 
     for height in self.ker.range(c.shape[-3], track=False):
@@ -122,7 +121,7 @@ class Group:
           else: raise NotImplementedError(f"mma_ABt not implemented for {a_base_shape.cols=}")
           d_in = UOp.stack(*[c[height, width, i] for i in range(4)])
 
-          out = UOp(Ops.WMMA, dtypes.float32, (a_in, b_in, d_in), arg=wmma_arg)
+          out = UOp.wmma(a_in, b_in, d_in, wmma_dims, 'AMD', 64)
           c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
           c_store = UOp.group(*c_i).end(height, width, inner)
 
@@ -135,9 +134,9 @@ class Group:
 
     a_base_shape = cast(RT, a).base_shape
     if a_base_shape.cols == 16:
-      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_dims = (16, 16, 16)
     elif a_base_shape.cols == 32:
-      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_dims = (16, 16, 32)
     else: raise NotImplementedError(f"mma_AtB not implemented for {a_base_shape.cols=}")
 
     for height in self.ker.range(c.shape[-3], track=False):
@@ -152,7 +151,7 @@ class Group:
           else: raise NotImplementedError(f"mma_AtB not implemented for {a_base_shape.cols=}")
           d_in = UOp.stack(*[c[height, width, i] for i in range(4)])
 
-          out = UOp(Ops.WMMA, dtypes.float32, (a_in, b_in, d_in), arg=wmma_arg)
+          out = UOp.wmma(a_in, b_in, d_in, wmma_dims, 'AMD', 64)
           c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
           c_store = UOp.group(*c_i).end(height, width, inner)
 
@@ -165,9 +164,9 @@ class Group:
 
     a_base_shape = cast(RT, a).base_shape
     if a_base_shape.cols == 16:
-      wmma_arg = ('WMMA_16_16_16___bf16_float', (16, 16, 16), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2)), ((4, 2), (3, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_dims = (16, 16, 16)
     elif a_base_shape.cols == 32:
-      wmma_arg = ('WMMA_16_16_32___bf16_float', (16, 16, 32), dtypes.bfloat16, dtypes.float, 'AMD', 64, (((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2), (9, 2)), ((4, 2), (3, 2))), ()) # type: ignore
+      wmma_dims = (16, 16, 32)
     else: raise NotImplementedError(f"mma_AtBt not implemented for {a_base_shape.cols=}")
 
     for height in self.ker.range(c.shape[-3], track=False):
@@ -182,7 +181,7 @@ class Group:
           else: raise NotImplementedError(f"mma_AtBt not implemented for {a_base_shape.cols=}")
           d_in = UOp.stack(*[c[height, width, i] for i in range(4)])
 
-          out = UOp(Ops.WMMA, dtypes.float32, (a_in, b_in, d_in), arg=wmma_arg)
+          out = UOp.wmma(a_in, b_in, d_in, wmma_dims, 'AMD', 64)
           c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
           c_store = UOp.group(*c_i).end(height, width, inner)
 

test/backend/test_dtype.py

@@ -8,7 +8,7 @@ from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
 from tinygrad import Context, Device, Tensor, dtypes
 from hypothesis import given, settings, strategies as strat
-from test.helpers import rand_for_dtype
+from test.helpers import rand_for_dtype, min_normal
 from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX, FP8E4M3FNUZ_MAX, FP8E5M2FNUZ_MAX
 import pytest
 pytestmark = pytest.mark.filterwarnings("ignore")
@@ -25,10 +25,10 @@ def get_available_cast_dtypes(dtype: DType) -> List[DType]:
   if dtype not in supported_dtypes and dtype not in dtypes.fp8s+(dtypes.half,dtypes.bfloat16): return []
   return dts
 
-def _to_torch_storage_type(dtype:DType):
-  if dtype == dtypes.bfloat16: return torch.float32
-  if dtype in dtypes.fp8s: return torch.float32
-  return _to_torch_dtype(dtype)
+def _to_torch_storage(a:Tensor) -> torch.Tensor:
+  # tolist() of an fp8 Tensor gives floats, so convert and store in uint8
+  if a.dtype in dtypes.fp8s: return torch.tensor([float_to_fp8(x, a.dtype) for x in a.flatten().tolist()], dtype=torch.uint8).reshape(a.shape)
+  return torch.tensor(a.tolist(), dtype=_to_torch_dtype(a.dtype))
 
 def _test_to_np(a:Tensor, np_dtype, target):
   if DEBUG >= 2: print(a)
@@ -46,12 +46,15 @@ def _test_cast(a:Tensor, target_dtype:DType):
   if a.is_floating_point() and dtypes.is_unsigned(target_dtype):
     # converting negative float to unsigned integer is undefined
     a = a.abs()
+  if a.is_floating_point() and dtypes.is_float(target_dtype) and (mn:=min_normal(target_dtype)) >= min_normal(a.dtype):
+    # subnormals are zero, so an input below the target's min normal casts to 0
+    a = (a.abs() < mn).where(0, a)
 
   expected = list(a.numpy().astype(_to_np_dtype(target_dtype)))
   if target_dtype in dtypes.fp8s: expected = [truncate[target_dtype](x) for x in expected]
   _test_op(lambda: a.cast(target_dtype), target_dtype, expected)
 def _test_bitcast(a:Tensor, target_dtype:DType, target=None):
-  expected = torch.tensor(a.tolist(), dtype=_to_torch_storage_type(a.dtype)).view(_to_torch_dtype(target_dtype)).tolist()
+  expected = _to_torch_storage(a).view(_to_torch_dtype(target_dtype)).tolist()
   if target_dtype in dtypes.fp8s: expected = [fp8_to_float(x, target_dtype) for x in expected]
   _test_op(lambda: a.bitcast(target_dtype), target_dtype, target or expected)
 
@@ -61,10 +64,12 @@ class TestDType(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
     if cls.DTYPE is None: raise unittest.SkipTest("base class")
-    cls.DATA = rand_for_dtype(cls.DTYPE, 0x10, allow_subnormal=cls.DTYPE in supported_dtypes)
+    cls.DATA = rand_for_dtype(cls.DTYPE, 0x10, allow_subnormal=cls.DTYPE in supported_dtypes and cls.DTYPE not in dtypes.fp8s)
 
   def test_to_np(self):
-    _test_to_np(Tensor(self.DATA, dtype=self.DTYPE), _to_np_dtype(self.DTYPE), np.array(self.DATA, dtype=_to_np_dtype(self.DTYPE)))
+    a = Tensor(self.DATA, dtype=self.DTYPE)
+    self.assertEqual(a.dtype, self.DTYPE)
+    _test_to_np(a, _to_np_dtype(self.DTYPE), np.array(self.DATA, dtype=_to_np_dtype(self.DTYPE)))
 
   def test_casts_to(self):
     for dtype in get_available_cast_dtypes(self.DTYPE):
@@ -273,10 +278,11 @@ class TestBitCast(unittest.TestCase):
   @given(strat.sampled_from(dtype_ints + dtype_floats), strat.sampled_from(dtype_ints + dtype_floats))
   def test_shape_change_bitcast(self, dt1, dt2):
     data = rand_for_dtype(dt1, 32).reshape(2, 2, 8)
-    expected = torch.tensor(data.tolist(), dtype=_to_torch_storage_type(dt1)).view(_to_torch_dtype(dt2))
+    a = Tensor(data, dtype=dt1)
+    expected = _to_torch_storage(a).view(_to_torch_dtype(dt2))
     if dt2 in dtypes.fp8s:
       expected = torch.tensor([fp8_to_float(x, dt2) for x in expected.view(-1).tolist()]).view_as(expected)
-    _test_op(lambda: Tensor(data, dtype=dt1).bitcast(dt2), dt2, expected.tolist())
+    _test_op(lambda: a.bitcast(dt2), dt2, expected.tolist())
 
   def test_shape_change_bitcast_exceptions(self):
     with self.assertRaises(RuntimeError):
@@ -293,6 +299,11 @@ class TestBitCast(unittest.TestCase):
     b = a.bitcast(dtypes.float32)
     assert b.numpy()[0,0] == 1.
 
+  def test_bitcast_bf16_from_cast(self):
+    # a bfloat16 from a cast holds bfloat16 bits. 1.0 is 0x3f80 in bfloat16, which is 1.875 in half
+    a = Tensor([1.0], dtype=dtypes.float32).cast(dtypes.bfloat16)
+    assert a.bitcast(dtypes.half).numpy()[0] == 1.875
+
 class TestInt16DType(TestDType): DTYPE = dtypes.int16
 
 class TestUint16DType(TestDType):

test/backend/test_graph.py

@@ -1,9 +1,9 @@
 import numpy as np
-import functools, unittest, ctypes
+import functools, unittest
 
 from tinygrad.device import Device, Buffer
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import Context, from_mv
+from tinygrad.helpers import Context
 from tinygrad.dtype import dtypes
 from tinygrad.engine.jit import MultiGraphRunner
 from tinygrad.engine.realize import run_linear, compile_linear
@@ -31,7 +31,7 @@ def make_buffer(device, size=BUF_SIZE, fill=False):
   buf = Buffer(device, size, dtypes.int).ensure_allocated()
   if fill:
     with Context(DEBUG=0):
-      buf.copyin(Tensor(np.random.randint(-10000, 10000, size=size, dtype=np.int32)).realize().uop.base.realized.as_memoryview())
+      buf.copy_from(Tensor(np.random.randint(-10000, 10000, size=size, dtype=np.int32)).realize().uop.base.realized)
   return buf
 
 def make_view(base, offset_elems, size_elems):
@@ -55,15 +55,12 @@ def run_schedule(calls:list[UOp]):
   run_linear(UOp(Ops.LINEAR, src=tuple(calls)))
 
 def zero_bufs(bufs):
-  for b in bufs:
-    mv = memoryview(bytearray(b.nbytes))
-    ctypes.memset(from_mv(mv), 0, len(mv))
-    b.copyin(mv)
+  for b in bufs: b.copy_from(Buffer("PYTHON", b.size, b.dtype, opaque=memoryview(bytearray(b.nbytes))))
 
 @unittest.skipUnless(Device[Device.DEFAULT].graph is not None, "graph support required")
 class TestGraph(unittest.TestCase):
   def skip_if_no_offset(self):
-    if not hasattr(Device[Device.DEFAULT].allocator, "_offset"): self.skipTest("device does not support _offset")
+    if Device.DEFAULT in {"WEBGPU", "CL"}: self.skipTest("device does not support _offset")
 
   def skip_if_not_multigraph(self):
     graph = g.func if isinstance(g:=(d:=Device[Device.DEFAULT]).graph, functools.partial) else g
@@ -213,8 +210,8 @@ class TestGraph(unittest.TestCase):
 
   def test_graph_offset_bufs(self):
     self.skip_if_not_multigraph()
+    self.skip_if_no_offset()
     d0 = Device.DEFAULT
-    if not hasattr(Device[d0].allocator, "_offset"): self.skipTest("device does not support _offset")
 
     b0 = make_buffer(d0, fill=True)
     b1 = make_view(b0, 0, b0.size)

test/backend/test_linearizer.py

@@ -424,7 +424,7 @@ def copyout_outputs(outbufs:list[Buffer]) -> list[np.ndarray]:
   return [np.frombuffer(x.as_memoryview(), _to_np_dtype(x.dtype)) for x in outbufs]
 
 def reset_bufs(bufs:list[Buffer]):
-  for buf in bufs: buf.copyin(np.zeros((buf.size*buf.dtype.itemsize,), dtype=np.uint8).data)
+  for buf in bufs: buf.copy_from(Buffer("PYTHON", buf.size, buf.dtype, opaque=memoryview(bytearray(buf.nbytes))))
 
 def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[],
                                apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[]):

test/backend/test_llama_kernels.py

@@ -1,11 +1,14 @@
-import unittest
+import unittest, functools
 from tinygrad import Tensor, Device, dtypes, Context, GlobalCounters
 from tinygrad.helpers import getenv
 from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8
 from extra.llama_kernels.fused_ce import fused_ce_loss
 from extra.llama_kernels import local_abs_max
 from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed, quantize_fp8_scalar
+from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
+from extra.thunder.amd.fa import custom_fused_qkv_rope_backward, fused_qkv_rope
 from test.helpers import needs_second_gpu
+from test.backend.test_asm_gemm import has_hipcc
 
 def run_fused_ce(bs:int, seqlen:int, vocab:int, label_smoothing:float=0.0) -> None:
   Tensor.manual_seed(0)
@@ -92,5 +95,68 @@ class TestLocalAmax(unittest.TestCase):
     self.assertEqual(GlobalCounters.kernel_count, 2)
     self.assertEqual(out.tolist(), [[0., 7., 14., 21.], [28., 35., 42., 49.], [120., 135., 150., 165.], [180., 195., 210., 225.]])
 
+@unittest.skipUnless(has_hipcc() and Device.DEFAULT == "AMD", "requires hipcc to compile and amd device to run")
+class TestFusedQKVRoPE(unittest.TestCase):
+  SHAPE = (2, 8192, 32, 8, 128)
+
+  def rand_bf16(self, *shape:int) -> Tensor:
+    return (Tensor.randn(*shape) * 0.1).cast(dtypes.bfloat16).contiguous().realize()
+
+  def freqs_cis(self) -> Tensor:
+    _, N, _, _, D = self.SHAPE
+    return precompute_freqs_cis(D, N * 2).cast(dtypes.bfloat16).clone().realize()
+
+  def test_llama31_8b_forward(self):
+    Tensor.manual_seed(0)
+    B, N, H, H_KV, D = self.SHAPE
+    GROUP = H // H_KV
+    freqs_cis = self.freqs_cis()
+
+    x = self.rand_bf16(B, N, H_KV * (GROUP + 2) * D)
+    q, k, v = fused_qkv_rope(x, freqs_cis, H, H_KV, D)
+    Tensor.realize(q, k, v)
+    packed_ref = x.reshape(B, N, H_KV, GROUP + 2, D)
+    q_ref = packed_ref[:, :, :, :GROUP].reshape(B, N, H, D)
+    k_ref, v_ref = packed_ref[:, :, :, GROUP], packed_ref[:, :, :, GROUP+1]
+    q_ref, k_ref = apply_rotary_emb(q_ref, k_ref, freqs_cis[:, :N])
+    q_ref, k_ref, v_ref = q_ref.cast(dtypes.bfloat16), k_ref.cast(dtypes.bfloat16), v_ref.cast(dtypes.bfloat16)
+    Tensor.realize(q_ref, k_ref, v_ref)
+
+    with Context(DEBUG=0):
+      self.assertTrue(q.allclose(q_ref, atol=2e-2, rtol=0).item(), "Q forward mismatch")
+      self.assertTrue(k.allclose(k_ref, atol=2e-2, rtol=0).item(), "K forward mismatch")
+      self.assertTrue(v.allclose(v_ref, atol=0, rtol=0).item(), "V forward mismatch")
+
+  def test_llama31_8b_backward(self):
+    Tensor.manual_seed(1)
+    B, N, H, H_KV, D = self.SHAPE
+    PARTIALS = 2
+    GROUP = H // H_KV
+    freqs_cis = self.freqs_cis()
+    dq = self.rand_bf16(B, N, H, D)
+    dk_partial = self.rand_bf16(B * PARTIALS, N, H_KV, D)
+    dv_partial = self.rand_bf16(B * PARTIALS, N, H_KV, D)
+
+    # Invert Flash Attention's dQ layout transform to reproduce its native buffer.
+    dq_native = dq.transpose(1, 2).reshape(B, H, N//16, 4, 4, 4, 2, D//32, 2, 2) \
+      .permute(0, 1, 2, 5, 6, 8, 7, 3, 4, 9).reshape(B, H, N, D).contiguous().realize()
+    dx = Tensor.empty(B, N, H_KV * (GROUP + 2) * D, dtype=dtypes.bfloat16)
+    arch = Device[Device.DEFAULT].renderer.target.arch
+    fxn = functools.partial(custom_fused_qkv_rope_backward, device=Device.DEFAULT, arch=arch,
+                            B=B, N=N, H=H, H_KV=H_KV, D=D)
+    dx = Tensor.custom_kernel(dx, dq_native, dk_partial, dv_partial, freqs_cis, fxn=fxn)[0].realize()
+
+    def inverse_rope(x:Tensor) -> Tensor:
+      x = x.reshape(*x.shape[:-1], D//2, 2).float()
+      cs = freqs_cis[:, :N].float()
+      return Tensor.stack(x[..., 0] * cs[..., 0] + x[..., 1] * cs[..., 1],
+                          -x[..., 0] * cs[..., 1] + x[..., 1] * cs[..., 0], dim=-1).flatten(-2).cast(dtypes.bfloat16)
+
+    dq_ref = inverse_rope(dq).reshape(B, N, H_KV, GROUP, D)
+    dk_ref = inverse_rope(dk_partial.float().reshape(B, PARTIALS, N, H_KV, D).sum(1).cast(dtypes.bfloat16)).unsqueeze(3)
+    dv_ref = dv_partial.float().reshape(B, PARTIALS, N, H_KV, D).sum(1).cast(dtypes.bfloat16).unsqueeze(3)
+    ref = Tensor.cat(dq_ref, dk_ref, dv_ref, dim=3).reshape(*dx.shape).realize()
+    with Context(DEBUG=0): self.assertTrue(dx.allclose(ref, atol=2e-2, rtol=2e-2).item(), "backward mismatch")
+
 if __name__ == '__main__':
   unittest.main()

test/backend/test_multitensor.py

@@ -385,7 +385,7 @@ class TestMultiBufferView(unittest.TestCase):
     b_ref = view_fn(a_ref)
     b_multi = view_fn(a_multi).contiguous()
     linear, var_vals = b_multi.linear_with_vars()
-    if all(hasattr(Device[d].allocator, "_offset") for d in b_multi.device):
+    if all(not d.startswith(("WEBGPU", "CL")) for d in b_multi.device):
       compiled = [call for call in linear.src if call.src[0].op is Ops.SINK]
       self.assertEqual(len(compiled), 0, f"expected zero compiled kernels, got {len(compiled)}")
     run_linear(linear, var_vals)
@@ -417,7 +417,7 @@ class TestMultiBufferView(unittest.TestCase):
     a = Tensor.arange(8*12).reshape(8, 12).clone().shard(devices_4, axis=1).realize()
     out = a[5].contiguous()
     linear, var_vals = out.linear_with_vars()
-    if all(hasattr(Device[d].allocator, "_offset") for d in out.device):
+    if all(not d.startswith(("WEBGPU", "CL")) for d in out.device):
       compiled = [call for call in linear.src if call.src[0].op is Ops.SINK]
       self.assertEqual(len(compiled), 0)
     run_linear(linear, var_vals)

test/backend/test_profiler.py

@@ -70,9 +70,9 @@ class TestProfiler(unittest.TestCase):
     buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
 
     with helper_collect_profile(TestProfiler.d0) as profile:
-      buf1.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
+      buf1.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 1)))))
 
-    kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith(TestProfiler.d0.device)]
+    kernel_runs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith((TestProfiler.d0.device, "PYTHON"))]
     assert len(kernel_runs) == 1, "one kernel run is expected"
 
   def test_profile_multiops(self):
@@ -80,12 +80,12 @@ class TestProfiler(unittest.TestCase):
     buf1 = Buffer(Device.DEFAULT, 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
 
     with helper_collect_profile(TestProfiler.d0) as profile:
-      buf1.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
+      buf1.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 1)))))
       gs, ls = TestProfiler.prg.arg.launch_dims({})
       TestProfiler.runtime(buf1._buf, TestProfiler.a.uop.buffer._buf, global_size=gs, local_size=ls)
-      buf1.copyout(memoryview(bytearray(buf1.nbytes)))
+      buf1.as_memoryview()
 
-    evs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith(TestProfiler.d0.device)]
+    evs = [x for x in profile if isinstance(x, ProfileRangeEvent) and x.device.startswith((TestProfiler.d0.device, "PYTHON"))]
 
     assert len(evs) == 3, "3 kernel runs are expected"
     # NOTE: order of events does not matter, the tool is responsible for sorting them
@@ -103,12 +103,12 @@ class TestProfiler(unittest.TestCase):
     buf2 = Buffer(f"{Device.DEFAULT}:1", 2, dtypes.float, options=BufferSpec(nolru=True)).ensure_allocated()
 
     with helper_collect_profile(TestProfiler.d0, d1) as profile:
-      buf1.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
-      buf2.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
+      buf1.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 1)))))
+      buf2.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 1)))))
 
     for dev in [TestProfiler.d0.device, d1.device]:
       evs = [x for x in profile if isinstance(x, ProfileRangeEvent) and _dev_base(x.device) == dev]
-      assert len(evs) == 1, "one kernel runs are expected"
+      assert len(evs) == (0 if hasattr(TestProfiler.d0.allocator, '_as_buffer') else 1), "one kernel runs are expected"
 
   def test_profile_multidev_transfer(self):
     try: d1 = Device[f"{Device.DEFAULT}:1"]

test/backend/test_renderer_failures.py

@@ -1,6 +1,6 @@
 import unittest
 import numpy as np
-from tinygrad.device import Device
+from tinygrad.device import Device, Buffer
 from tinygrad.dtype import dtypes, ConstType
 from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
@@ -10,13 +10,13 @@ from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.wgsl import WGSLRenderer
 from tinygrad.runtime.ops_python import PythonRenderer
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, python_alu
-from tinygrad.tensor import Tensor, _to_np_dtype
+from tinygrad.tensor import Tensor
 
 def _test_uop_result(inputs:list[Tensor], sink:UOp, local_size=None):
   for x in inputs: x.realize()
   sz = 1 if local_size is None else prod(local_size)
   outs = [UOp.new_buffer(Device.DEFAULT, sz, u.src[1].dtype) for u in sink.src if u.op is Ops.STORE]
-  for u in outs: u.buffer.allocate().copyin(np.zeros(sz, dtype=_to_np_dtype(u.dtype)).data)
+  for u in outs: u.buffer.allocate().copy_from(Buffer("PYTHON", sz, u.dtype, opaque=memoryview(bytearray(u.buffer.nbytes))))
   run_linear(UOp(Ops.LINEAR, src=(sink.call(*outs, *(x.uop.base for x in inputs)),)))
   return [u.buffer.numpy() for u in outs]
 

test/backend/test_schedule.py

@@ -2,7 +2,7 @@
 # schedule confirms the right things are capable of fusing
 # NOTE: this has overlap with external_test_opt.py
 
-import unittest
+import unittest, time
 import numpy as np
 
 from tinygrad import nn, dtypes, Device, Tensor, Variable
@@ -197,6 +197,20 @@ class TestLimitBufs(unittest.TestCase):
         base = (idx >= i).where(a + b, base)
       assert all(x > 0 for x in base.tolist())
 
+  def test_limit_bufs_linear_scaling(self):
+    def sched_time(n):
+      with Context(TRACK_MATCH_STATS=0, DEBUG=0):
+        bufs = [Tensor.ones(16).contiguous().realize() for _ in range(4)]
+        root = bufs[0]
+        for i in range(n): root = root + bufs[i % 4]
+        with Context(MAX_KERNEL_BUFFERS=8, SCACHE=0):
+          st = time.perf_counter()
+          root.schedule_linear()
+          return time.perf_counter() - st
+    sched_time(400)
+    t1, t2 = min(sched_time(400) for _ in range(3)), min(sched_time(1600) for _ in range(3))
+    self.assertLess(t2/t1, 8, f"{t1*1e3:.1f}ms -> {t2*1e3:.1f}ms")
+
 class TestSwizzle(unittest.TestCase):
   def test_swizzle_simple(self):
     Tensor.manual_seed(0)

test/backend/test_subbuffer.py

@@ -4,11 +4,10 @@ from tinygrad.device import Buffer
 from tinygrad.helpers import Context, DEV
 from test.helpers import needs_second_gpu
 
-@unittest.skipUnless(hasattr(Device[Device.DEFAULT].allocator, "_offset"), "subbuffer not supported")
+@unittest.skipIf(Device.DEFAULT in {"WEBGPU", "CL"}, "subbuffer not supported")
 class TestSubBuffer(unittest.TestCase):
   def setUp(self):
-    self.buf = Buffer(Device.DEFAULT, 10, dtypes.uint8).ensure_allocated()
-    self.buf.copyin(memoryview(bytearray(range(10))))
+    self.buf = Buffer(Device.DEFAULT, 10, dtypes.uint8, initial_value=bytes(range(10)))
     self.buf_unalloc = Buffer(Device.DEFAULT, 10, dtypes.uint8)
 
   def test_subbuffer(self):
@@ -59,7 +58,7 @@ class TestSubBuffer(unittest.TestCase):
       _ = Buffer(Device.DEFAULT, 10, dtypes.uint8).ensure_allocated()
 
       self.buf.ensure_allocated()
-      self.buf.copyin(memoryview(bytearray(range(10, 20))))
+      self.buf.copy_from(Buffer("PYTHON", 10, dtypes.uint8, opaque=memoryview(bytearray(range(10, 20)))))
 
       vbuf.ensure_allocated()
 
@@ -109,15 +108,15 @@ class TestSubBuffer(unittest.TestCase):
   def test_subbuffer_copy_in_out(self):
     sub_buf = self.buf.view(3, dtypes.uint8, offset=3).ensure_allocated() # [3:6]
     data_out_sub = bytearray([0]*3)
-    sub_buf.copyout(memoryview(data_out_sub))
+    data_out_sub[:] = sub_buf.as_memoryview()
     assert data_out_sub == bytearray(range(3, 6))
-    sub_buf.copyin(memoryview(bytearray(range(3))))
+    sub_buf.copy_from(Buffer("PYTHON", 3, dtypes.uint8, opaque=memoryview(bytearray(range(3)))))
     assert sub_buf.as_memoryview().tolist() == list(range(3))
     assert self.buf.as_memoryview().tolist()[3:6] == list(range(3))
-    sub_buf.copyout(memoryview(data_out_sub))
+    data_out_sub[:] = sub_buf.as_memoryview()
     assert data_out_sub == bytearray(range(3))
     data_out_base = bytearray([0]*10)
-    self.buf.copyout(memoryview(data_out_base))
+    data_out_base[:] = self.buf.as_memoryview()
     assert data_out_base[0:3] == bytearray(range(0, 3))
     assert data_out_base[3:6] == data_out_sub
     assert data_out_base[6:10] == bytearray(range(6, 10))
@@ -129,27 +128,27 @@ class TestSubBuffer(unittest.TestCase):
     self.assertTrue(view2.is_allocated())
 
     data_in = bytearray([7, 8, 9])
-    view2.copyin(memoryview(data_in))
+    view2.copy_from(Buffer("PYTHON", 3, view2.dtype, opaque=memoryview(data_in)))
     data_out_v2 = bytearray([0]*3)
-    view2.copyout(memoryview(data_out_v2))
+    data_out_v2[:] = view2.as_memoryview()
     assert data_in == data_out_v2
 
     expected_base_data = memoryview(bytearray(range(10)))
     expected_base_data[4:7] = data_in
 
     data_out_base = bytearray([0]*10)
-    self.buf.copyout(memoryview(data_out_base))
+    data_out_base[:] = self.buf.as_memoryview()
     assert expected_base_data == data_out_base
 
   def test_subbuffer_alloc(self):
     sub_buf = self.buf.view(4, dtypes.int8, offset=3)
     sub_buf.allocate()
-    sub_buf.copyin(memoryview(bytearray(range(10, 14))))
+    sub_buf.copy_from(Buffer("PYTHON", 4, dtypes.int8, opaque=memoryview(bytearray(range(10, 14)))))
     assert self.buf.as_memoryview().tolist()[3:7] == sub_buf.as_memoryview().tolist()
 
     sub_buf = self.buf_unalloc.view(4, dtypes.int8, offset=3)
     sub_buf.allocate()
-    sub_buf.copyin(memoryview(bytearray(range(10, 14))))
+    sub_buf.copy_from(Buffer("PYTHON", 4, dtypes.int8, opaque=memoryview(bytearray(range(10, 14)))))
     assert self.buf_unalloc.as_memoryview().tolist()[3:7] == sub_buf.as_memoryview().tolist()
 
   def test_subbuffer_dealloc(self):

test/backend/test_uops.py

@@ -33,11 +33,9 @@ def _test_single_value(vals, op, dts):
   alu = uop(uops, op, output_dtype, loads)
   out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), alu))
   buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
-  buf2 = [Buffer(Device.DEFAULT, 1, dtype).allocate().copyin(np.array([a], dtype=_to_np_dtype(dtype)).data) for a,dtype in zip(vals, dts)]
+  buf2 = [Buffer(Device.DEFAULT, 1, dtype, initial_value=np.array([a], dtype=_to_np_dtype(dtype)).tobytes()) for a,dtype in zip(vals, dts)]
   run_uops([out], [buf]+buf2)
-  ret = np.empty(1, _to_np_dtype(output_dtype))
-  buf.copyout(ret.data)
-  return ret[0]
+  return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0]
 
 def _test_single_value_const(vals, op, dts):
   uops = []
@@ -48,9 +46,7 @@ def _test_single_value_const(vals, op, dts):
   out = buf_store[UOp.const(dtypes.int32, 0)].store(alu)
   buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
   run_uops([out], [buf])
-  ret = np.empty(1, _to_np_dtype(output_dtype))
-  buf.copyout(ret.data)
-  return ret[0]
+  return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0]
 
 def _test_uops_result(output_dtype, uops, res):
   # uops = []
@@ -59,9 +55,7 @@ def _test_uops_result(output_dtype, uops, res):
   out = uop(uops, Ops.STORE, dtypes.void, (buf_store.index(uop(uops, Ops.CONST, dtypes.int32, (), 0)), res))
   buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
   run_uops([out], [buf])
-  ret = np.empty(1, _to_np_dtype(output_dtype))
-  buf.copyout(ret.data)
-  return ret[0]
+  return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0]
 
 class TestUOps(unittest.TestCase):
   def _equal(self, v1, v2):

test/device/test_hcq.py

@@ -31,8 +31,8 @@ class TestHCQ(unittest.TestCase):
 
   def setUp(self):
     TestHCQ.d0.synchronize()
-    TestHCQ.a.uop.buffer.copyin(memoryview(bytearray(struct.pack("ff", 0, 1))))
-    TestHCQ.b.uop.buffer.copyin(memoryview(bytearray(struct.pack("ff", 0, 0))))
+    TestHCQ.a.uop.buffer.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 1)))))
+    TestHCQ.b.uop.buffer.copy_from(Buffer("PYTHON", 2, dtypes.float, opaque=memoryview(bytearray(struct.pack("ff", 0, 0)))))
     TestHCQ.d0.synchronize() # wait for copyins to complete
 
   # Test signals
@@ -376,7 +376,7 @@ class TestHCQ(unittest.TestCase):
     SZ = 200_000_000
     b = Buffer(f"{Device.DEFAULT}:1", SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
     a = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
-    TestHCQ.d0.allocator.map(b._buf)
+    TestHCQ.d0.allocator._map(b._buf)
 
     sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal()
     TestHCQ.d0.hw_copy_queue_t().timestamp(sig_st) \
@@ -454,7 +454,7 @@ class TestHCQ(unittest.TestCase):
     buf1 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
     buf2 = Buffer(f"{Device.DEFAULT}:1", 1, dtypes.int8, options=BufferSpec(nolru=True)).ensure_allocated()
     buf3 = Buffer(Device.DEFAULT, 1, dtypes.int8, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
-    TestHCQ.d0.allocator.map(buf2._buf)
+    TestHCQ.d0.allocator._map(buf2._buf)
 
     for i in range(256):
       ctypes.memset(buf3._buf.va_addr, i, 1)
@@ -569,7 +569,7 @@ class TestHCQ(unittest.TestCase):
 
       local_buf = Buffer(f"{Device.DEFAULT}:{devid}", sz, dtypes.uint8, options=BufferSpec(cpu_access=True)).ensure_allocated()
 
-      d.allocator.map(cpu_buffer._buf)
+      d.allocator._map(cpu_buffer._buf)
 
       d.hw_copy_queue_t().wait(d.timeline_signal, d.timeline_value - 1) \
                          .copy(local_buf._buf, cpu_buffer._buf, sz) \

test/device/test_ocl.py

@@ -34,7 +34,7 @@ class TestCLError(unittest.TestCase):
     data = list(range(65))
     unaligned = memoryview(bytearray(data))[1:]
     buffer = Buffer("CL", 64, dtypes.uint8).allocate()
-    buffer.copyin(unaligned)
+    buffer.copy_from(Buffer("PYTHON", 64, dtypes.uint8, opaque=unaligned))
     result = memoryview(bytearray(len(data) - 1))
-    buffer.copyout(result)
+    result[:] = buffer.as_memoryview()
     assert unaligned == result, "Unaligned data copied in must be equal to data copied out."

test/external/external_test_hcq.py

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

test/external/external_test_usb_asm24.py

@@ -1,47 +1,12 @@
-import unittest, time
-from tinygrad.runtime.support.usb import ASM24Controller
+import unittest
 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 = 512
+    cls.sz = 768
     cls.dev = Device["AMD"]
     if not cls.dev.is_usb(): raise unittest.SkipTest("only test this on USB devices")
 
@@ -70,4 +35,4 @@ class TestDevCopySpeeds(unittest.TestCase):
     del x, y, t
 
 if __name__ == "__main__":
-  unittest.main()
+  unittest.main()

test/external/fuzz_graph.py

@@ -1,7 +1,7 @@
-import random, ctypes
+import random
 import numpy as np
 from tinygrad.device import Buffer, Device
-from tinygrad.helpers import Context, getenv, from_mv
+from tinygrad.helpers import Context, getenv
 from tinygrad.dtype import dtypes
 from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.engine.realize import BufferXfer, get_runner, ExecItem
@@ -29,7 +29,7 @@ def alloc_rawbuffer(device, fill=False):
   if fill:
     with Context(DEBUG=0):
       data = np.random.randint(-10000, 10000, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
-      rawbuf.copyin(Tensor(data).realize().uop.base.realized.as_memoryview())
+      rawbuf.copy_from(Tensor(data).realize().uop.base.realized)
   return rawbuf
 
 def gen_kernel_ji(device, deps):
@@ -84,9 +84,7 @@ def run_jit(jis, all_buffers, input_buffers, var_vals):
   with Context(DEBUG=0):
     for rawbuf in all_buffers:
       if rawbuf in input_buffers: continue
-      mv = memoryview(bytearray(rawbuf.nbytes))
-      ctypes.memset(from_mv(mv), 0, len(mv))
-      rawbuf.copyin(mv)
+      rawbuf.copy_from(Buffer("PYTHON", rawbuf.size, rawbuf.dtype, opaque=memoryview(bytearray(rawbuf.nbytes))))
 
   for ei in jis: ei.run(var_vals, jit=True)
 

test/helpers.py

@@ -6,7 +6,7 @@ from tinygrad import Tensor, dtypes, Device
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.codegen import to_program
-from tinygrad.dtype import DType
+from tinygrad.dtype import DType, truncate
 from tinygrad.nn.state import get_parameters
 from tinygrad.helpers import T, Target, DEV
 from tinygrad.renderer import Renderer
@@ -38,6 +38,10 @@ def call_is_graph(call:UOp) -> bool:
   ast = call.src[0]
   return ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph"
 
+def call_is_hcq(call:UOp) -> bool:
+  ast = call.src[0]
+  return ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq"
+
 def jit_cache_count(linear:UOp) -> int:
   n = 0
   for call in linear.src:
@@ -51,6 +55,7 @@ def assert_jit_cache_len(fxn, expected_len):
   if linear is None or not linear.src:
     assert expected_len == 0, expected_len
     return
+  if expected_len and all(call_is_hcq(call) for call in linear.src): expected_len = 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
@@ -58,6 +63,8 @@ def assert_jit_cache_len(fxn, expected_len):
   else:
     assert len(linear.src) == expected_len, f"expected {expected_len}, got {len(linear.src)}"
 
+def min_normal(dt:DType) -> float: return 2.0 ** (2 - (1 << (dtypes.finfo(dt)[0] - 1)))
+
 def rand_for_dtype(dt:DType, size:int, allow_subnormal=True):
   if dtypes.is_unsigned(dt):
     return np.random.randint(0, 100, size=size, dtype=_to_np_dtype(dt))
@@ -66,9 +73,8 @@ def rand_for_dtype(dt:DType, size:int, allow_subnormal=True):
   elif dt == dtypes.bool:
     return np.random.choice([True, False], size=size)
   ret = np.random.uniform(-10, 10, size=size).astype(_to_np_dtype(dt))
-  if not allow_subnormal:
-    min_normal = 2.0 ** (2 - (1 << (dtypes.finfo(dt)[0] - 1)))
-    ret = np.where(np.abs(ret) < min_normal, 0, ret)
+  if dt == dtypes.bfloat16 or dt in dtypes.fp8s: ret = np.array([truncate[dt](x) for x in ret], dtype=ret.dtype)
+  if not allow_subnormal: ret = np.where(np.abs(ret) < min_normal(dt), 0, ret)
   return ret
 
 def timeit(fxn:Callable[..., T], *args, **kwargs) -> tuple[T, float]:

test/mockgpu/amd/amdgpu.py

@@ -9,7 +9,7 @@ MOCKGPU_ARCH = "cdna4" if DEV.arch == "gfx950" else "rdna4" if DEV.arch.startswi
 assert (ma:=getenv("MOCKGPU_ARCH", "")) == "", "MOCKGPU_ARCH is deprecated, use DEV=" + \
   str(replace(DEV.value, arch={"cdna4":"gfx950", "rdna4":"gfx1201"}.get(ma, "gfx1100"))) # type: ignore
 GFX_TARGET_VERSION = {"rdna3": 110000, "rdna4": 120000, "cdna4": 90500}[MOCKGPU_ARCH]
-import tinygrad.runtime.autogen.amd_gpu as amd_gpu, tinygrad.runtime.autogen.am.pm4_nv as pm4
+import tinygrad.runtime.autogen.amd_gpu as amd_gpu, tinygrad.runtime.autogen.am.pm4_nv as pm4, tinygrad.runtime.autogen.am.sdma_6_0_0 as sdma
 
 SDMA_MAX_COPY_SIZE = 0x400000
 
@@ -275,6 +275,7 @@ class SDMAExecutor(AMDQueue):
       elif op == amd_gpu.SDMA_OP_POLL_REGMEM: cont = self._execute_poll_regmem()
       elif op == amd_gpu.SDMA_OP_GCR: self._execute_gcr()
       elif op == amd_gpu.SDMA_OP_COPY: self._execute_copy()
+      elif op == sdma.SDMA_OP_WRITE: self._execute_write()
       elif op == amd_gpu.SDMA_OP_TIMESTAMP: self._execute_timestamp()
       elif op == 32: self.rptr[0] += 4  # SDMA_OP_DUMMY_TRAP: pipeline flush, no interrupt
       else: raise RuntimeError(f"Unknown SDMA op {op}")
@@ -289,6 +290,12 @@ class SDMAExecutor(AMDQueue):
     struct = sdma_pkts.trap.from_address(self.base + self.rptr[0] % self.size)
     self.rptr[0] += ctypes.sizeof(struct)
 
+  def _execute_write(self):
+    packet = to_mv(self.base + self.rptr[0] % self.size, 16).cast('I')
+    addr, count = packet[1] | packet[2] << 32, packet[3] + 1
+    ctypes.memmove(self.gpu.translate_addr(addr), self.base + self.rptr[0] % self.size + 16, count * 4)
+    self.rptr[0] += (4 + count) * 4
+
   def _execute_poll_regmem(self):
     struct = sdma_pkts.poll_regmem.from_address(self.base + self.rptr[0] % self.size)
 

test/mockgpu/amd/emu.py

@@ -803,7 +803,7 @@ def _compile_sopp(inst: ir3.SOPP | ir4.SOPP, ctx: _Ctx) -> UOp:
     pcode = get_pcode(inst.op)
     pc_bytes = ctx.rpc()  # PC is already 64-bit byte address
     vcc, exec_val = ctx.rmask(_c(VCC_LO.offset)), ctx.rexec()
-    srcs = {'PC': pc_bytes.cast(dtypes.int64), 'SIMM16': simm16, 'SCC': ctx.rsgpr_dyn(_c(SCC.offset)), 'VCC': vcc,
+    srcs: dict[str, UOp|int] = {'PC': pc_bytes.cast(dtypes.int64), 'SIMM16': simm16, 'SCC': ctx.rsgpr_dyn(_c(SCC.offset)), 'VCC': vcc,
             'VCCZ': vcc.eq(UOp.const(vcc.dtype, 0)).cast(dtypes.uint32),
             'EXECZ': exec_val.eq(UOp.const(exec_val.dtype, 0)).cast(dtypes.uint32)}
     for dest, val in parse_pcode(pcode, srcs)[1]:
@@ -858,7 +858,7 @@ def _compile_sop(inst: ir3.SOP1|ir3.SOP2|ir3.SOPC|ir3.SOPK|ir4.SOP1|ir4.SOP2|ir4
     if isinstance(inst, ir4.SOPK): s0 = simm16
     elif isinstance(inst, irc.SOPK) and 'CMPK' not in op_name and 'SETREG' not in op_name: s0 = simm16_sext
     else: s0 = ctx.rsgpr_dyn(sdst_off)
-    srcs = {'S0': s0, 'S1': simm16_sext, 'SIMM16': simm16_sext, 'D0': ctx.rsgpr_dyn(sdst_off)}
+    srcs: dict[str, UOp|int] = {'S0': s0, 'S1': simm16_sext, 'SIMM16': simm16_sext, 'D0': ctx.rsgpr_dyn(sdst_off)}
     dst_off, dst_size = sdst_off, 1
     # S_GETREG_B32: extract bits from HW register. Handle as special case since HW_REGISTERS is not a normal variable.
     # HW register values are stored at SGPR[SGPR_COUNT-16 + hwRegId] by _init_wave.

test/mockgpu/amd/pcode.py

@@ -688,10 +688,10 @@ class Parser:
         return _extract_bits(base, hi, lo)
       # Dynamic bit slice: (base >> lo) & ((1 << (hi - lo + 1)) - 1)
       dt = dtypes.uint64 if base.dtype in (dtypes.uint64, dtypes.int64) else dtypes.uint32
-      hi, lo = first.cast(dt), second.cast(dt)
-      width = hi - lo + _const(dt, 1)
+      hi_u, lo_u = first.cast(dt), second.cast(dt)
+      width = hi_u - lo_u + _const(dt, 1)
       mask = (_const(dt, 1) << width) - _const(dt, 1)
-      return (base.cast(dt) >> lo) & mask
+      return (base.cast(dt) >> lo_u) & mask
     self.eat('RBRACKET')
     dt_suffix = None
     if self.try_eat('DOT'):
@@ -1123,11 +1123,11 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
           val = parse_tokens(toks[j:], env, funcs)
           lo_dt, hi_dt = DTYPES.get(lo_type, dtypes.uint64), DTYPES.get(hi_type, dtypes.uint32)
           lo_bits = 64 if lo_dt in (dtypes.uint64, dtypes.int64) else 32
-          lo_val = val.cast(lo_dt) if val.dtype.itemsize * 8 <= lo_bits else (val & _const(val.dtype, (1 << lo_bits) - 1)).cast(lo_dt)
-          hi_val = (val >> _const(val.dtype, lo_bits)).cast(hi_dt)
-          block_assigns[lo_var] = env[lo_var] = lo_val
-          block_assigns[hi_var] = env[hi_var] = hi_val
-          if assigns is not None: assigns.extend([(f'{lo_var}.{lo_type}', lo_val), (f'{hi_var}.{hi_type}', hi_val)])
+          lo_u = val.cast(lo_dt) if val.dtype.itemsize * 8 <= lo_bits else (val & _const(val.dtype, (1 << lo_bits) - 1)).cast(lo_dt)
+          hi_u = (val >> _const(val.dtype, lo_bits)).cast(hi_dt)
+          block_assigns[lo_var] = env[lo_var] = lo_u
+          block_assigns[hi_var] = env[hi_var] = hi_u
+          if assigns is not None: assigns.extend([(f'{lo_var}.{lo_type}', lo_u), (f'{hi_var}.{hi_type}', hi_u)])
           i += 1
           continue
 

test/mockgpu/usb.py

@@ -6,38 +6,25 @@ class MockUSB:
   def __init__(self, mem):
     self.mem = mem
   def read(self, address, size): return bytes(self.mem[address:address+size])
-  def write(self, address, data, ignore_cache=False): self.mem[address:address+len(data)] = data
-  def pcie_mem_req(self, address, value=None, size=1):
-    if value is None: return int.from_bytes(self.mem[address:address+size], "little")
-    else: self.mem[address:address+size] = value.to_bytes(size, "little")
-  def pcie_mem_write(self, address, values, size):
-    for i, value in enumerate(values): self.pcie_mem_req(address + i * size, value, size)
+  def write(self, address, data): self.mem[address:address+len(data)] = data
+  def pcie_mem_read(self, address, nbytes): return bytes(self.mem[address:address+nbytes])
+  def pcie_mem_write(self, address, data): self.mem[address:address+len(data)] = data
 
 # *** ASM24 Controller Mock ***
 
 _mock_usb_state: MockASM24State|None = None
 
 class MockASM24State:
-  """Mock ASM24 controller: XRAM memory map, DMA windows, TLP engine, PCI config space.
+  """Mock custom ASM24 controller: XRAM, DMA windows, PCI config space, and GPU BARs.
 
   Memory map (64KB XRAM):
     0xA000-0xAFFF: DMA window -> sys 0x820000
     0xB000-0xB1FF: DMA window -> sys 0x800000
-    0xB200-0xB7FF: PCI MMIO (TLP engine)
+    0xB200-0xB7FF: controller PCI MMIO
     0xF000-0xFFFF: DMA window -> sys 0x200000 (512KB)
   """
   XRAM_SIZE = 0x10000
 
-  TLP_FMT_TYPE = 0xB210
-  TLP_BYTE_EN = 0xB217
-  TLP_ADDR_LO = 0xB218
-  TLP_ADDR_HI = 0xB21C
-  TLP_DATA = 0xB220
-  TLP_COMPL = 0xB22A
-  TLP_TRIGGER = 0xB254
-  TLP_LINK_STATUS = 0xB284
-  TLP_STATUS = 0xB296
-
   def __init__(self, gpu, driver, vram_size:int, doorbell_size:int, mmio_size:int):
     self.gpu, self.driver = gpu, driver
     self._xram = bytearray(self.XRAM_SIZE)
@@ -93,53 +80,7 @@ class MockASM24State:
       if ctrl_addr <= addr < ctrl_addr + dma_size:
         (ctypes.c_ubyte * 1).from_address(host_addr + (addr - ctrl_addr))[0] = value
         return
-    if addr == self.TLP_STATUS:
-      self._xram[addr] &= ~value & 0xFF
-      return
     self._xram[addr] = value
-    if addr == self.TLP_TRIGGER and value == 0x0F: self._process_tlp()
-
-  # --- TLP engine ---
-
-  def _process_tlp(self):
-    fmt_type, byte_en = self._xram[self.TLP_FMT_TYPE], self._xram[self.TLP_BYTE_EN]
-    addr_lo = int.from_bytes(self._xram[self.TLP_ADDR_LO:self.TLP_ADDR_LO+4], 'big')
-    addr_hi = int.from_bytes(self._xram[self.TLP_ADDR_HI:self.TLP_ADDR_HI+4], 'big')
-    address = addr_lo | (addr_hi << 32)
-
-    size, offset, tmp = 0, 0, byte_en
-    while tmp and not (tmp & 1):
-      offset += 1
-      tmp >>= 1
-    while tmp:
-      size += tmp & 1
-      tmp >>= 1
-
-    is_write, is_cfg = bool(fmt_type & 0x40), (fmt_type & 0xbe) == 0x04
-
-    if is_cfg:
-      bus, dev, fn, byte_addr = (address >> 24) & 0xFF, (address >> 19) & 0x1F, (address >> 16) & 0x7, address & 0xFFC
-      if is_write:
-        data = int.from_bytes(self._xram[self.TLP_DATA:self.TLP_DATA+4], 'big')
-        self._cfg_write(bus, dev, fn, byte_addr + offset, (data >> (8 * offset)) & ((1 << (8 * size)) - 1), size)
-      else:
-        self._xram[self.TLP_DATA:self.TLP_DATA+4] = int.from_bytes(self._get_cfg(bus, dev, fn)[byte_addr:byte_addr+4], 'little').to_bytes(4, 'big')
-      self._xram[self.TLP_COMPL:self.TLP_COMPL+2] = (4).to_bytes(2, 'big')
-      self._xram[self.TLP_LINK_STATUS] = 0x01 if not is_write else 0x00
-      self._xram[self.TLP_STATUS] = 0x02
-      return
-
-    if is_write:
-      data = int.from_bytes(self._xram[self.TLP_DATA:self.TLP_DATA+4], 'big')
-      self._pcie_dispatch(address + offset, (data >> (8 * offset)) & ((1 << (8 * size)) - 1), size)
-    else:
-      result = self._pcie_dispatch(address + offset, None, size)
-      if result is not None:
-        self._xram[self.TLP_DATA:self.TLP_DATA+4] = ((result << (8 * offset)) & 0xFFFFFFFF).to_bytes(4, 'big')
-
-    self._xram[self.TLP_COMPL:self.TLP_COMPL+2] = (size & 0xFFF).to_bytes(2, 'big')
-    self._xram[self.TLP_LINK_STATUS] = 0x01 if not is_write else 0x00
-    self._xram[self.TLP_STATUS] = 0x02
 
   def _cfg_write(self, bus:int, dev:int, fn:int, byte_addr:int, val:int, size:int):
     cfg = self._get_cfg(bus, dev, fn)
@@ -168,49 +109,104 @@ class MockASM24State:
     # Generic config write
     for i in range(size): cfg[byte_addr + i] = (val >> (8 * i)) & 0xFF
 
-  def _pcie_dispatch(self, address:int, value:int|None, size:int) -> int|None:
+  def _find_bar(self, address:int, size:int) -> tuple[int, int]:
     for reg_off, (bar_addr, bar_size) in self._bar_addrs.items():
-      if bar_addr <= address < bar_addr + bar_size:
-        offset = address - bar_addr
-        if reg_off == 0x10:  # BAR0 - VRAM
-          if value is None: return int.from_bytes(bytes(self.gpu.vram[offset:offset+size]), "little")
-          self.gpu.vram[offset:offset+size] = list(value.to_bytes(size, "little"))
-          return None
-        if reg_off == 0x18:  # BAR2 - Doorbell
-          if value is None: return int.from_bytes(bytes(self._doorbell[offset:offset+size]), "little")
-          for i, b in enumerate(value.to_bytes(size, "little")): self._doorbell[offset + i] = b
-          self.driver._emulate_execute()
-          return None
-        if reg_off == 0x24:  # BAR5 - MMIO
-          if value is None: return self.gpu.mmio[offset // 4]
-          self.gpu.mmio[offset // 4] = value
-          return None
-    raise ValueError(f"PCIe address {address:#x} not mapped to any BAR")
+      if bar_addr <= address and address + size <= bar_addr + bar_size: return reg_off, address - bar_addr
+    raise ValueError(f"PCIe range {address:#x}+{size:#x} not mapped to any BAR")
 
-  # --- CDB processing (called by MockUSB3.send_batch) ---
+  def _pcie_read(self, address:int, size:int) -> bytes:
+    reg_off, offset = self._find_bar(address, size)
+    if reg_off == 0x10: return bytes(self.gpu.vram[offset:offset+size])
+    if reg_off == 0x18: return bytes(self._doorbell[offset:offset+size])
+    if reg_off == 0x24: return bytes((self.gpu.mmio[(offset+i)//4] >> (8*((offset+i)&3))) & 0xFF for i in range(size))
+    raise RuntimeError(f"unsupported BAR register {reg_off:#x}")
 
-  def process_cdb(self, cdb:bytes, rlen:int, send_data:bytes|None) -> bytes|None:
-    op = cdb[0]
-    if op == 0xE5:  # write byte
-      self._xram_write_byte(((cdb[2] << 16) | (cdb[3] << 8) | cdb[4]) & 0xFFFF, cdb[1])
-      return None
-    if op == 0xE4:  # read
-      return self._xram_read(((cdb[2] << 16) | (cdb[3] << 8) | cdb[4]) & 0xFFFF, cdb[1])
-    if op == 0x8A and send_data is not None and 0xF000 in self._dma_regions:  # SCSI write
-      host_addr, dma_size = self._dma_regions[0xF000]
-      ctypes.memmove(host_addr, send_data, min(len(send_data), dma_size))
+  def _pcie_write(self, address:int, data:bytes):
+    reg_off, offset = self._find_bar(address, len(data))
+    if reg_off == 0x10: self.gpu.vram[offset:offset+len(data)] = list(data)
+    elif reg_off == 0x18:
+      self._doorbell[offset:offset+len(data)] = list(data)
+      self.driver._emulate_execute()
+    elif reg_off == 0x24:
+      updates: dict[int, int] = {}
+      for i, byte in enumerate(data):
+        idx, shift = (offset+i)//4, 8*((offset+i)&3)
+        updates[idx] = (updates.get(idx, self.gpu.mmio[idx]) & ~(0xFF << shift)) | (byte << shift)
+      for idx, val in updates.items(): self.gpu.mmio[idx] = val
+    else: raise RuntimeError(f"unsupported BAR register {reg_off:#x}")
+
+  def _pcie_dispatch(self, address:int, value:int|None, size:int) -> int|None:
+    if value is None: return int.from_bytes(self._pcie_read(address, size), 'little')
+    self._pcie_write(address, value.to_bytes(size, 'little'))
     return None
 
 class MockUSB3:
   @classmethod
   def list_devices(cls, vendor, dev): return [(0, "usb:mock")]
   def __init__(self, *args, **kwargs):
-    self.product, self.is_custom = "", False
-  def send_batch(self, cdbs:list[bytes], idata:list[int]|None=None, odata:list[bytes|None]|None=None) -> list[bytes|None]:
+    self.product = "custom mock"
+    self._bulk_read_op: tuple[str, int, int]|None = None
+    self._bulk_write_op: tuple[str, int, int]|None = None
+    self._f0_reply = bytes(8)
+
+  @property
+  def state(self) -> MockASM24State:
     assert _mock_usb_state is not None
-    idata, odata = idata or [0] * len(cdbs), odata or [None] * len(cdbs)
-    results: list[bytes|None] = []
-    for cdb, rlen, sdata in zip(cdbs, idata, odata):
-      result = _mock_usb_state.process_cdb(cdb, rlen, sdata)
-      results.append(result if rlen > 0 else None)
-    return results
+    return _mock_usb_state
+
+  def control_write(self, request:int, value:int=0, index:int=0, data:bytes=b'', timeout:int=1000):
+    if request == 0xF3:
+      self.state._xram[0xB450] = 0x78 if value else 0
+    elif request == 0xE5:
+      self.state._xram_write_byte(value, index)
+    elif request == 0xF2:
+      op = ("sram_read" if value & 0x8000 else "sram_write", 0xF000, (value & 0x7FFF) * 512)
+      if value & 0x8000: self._bulk_read_op = op
+      else: self._bulk_write_op = op
+    elif request == 0xF0:
+      address_lo, address_hi, payload = struct.unpack('<III', data)
+      address, fmt_type, byte_en = address_lo | (address_hi << 32), value & 0xFF, value >> 8
+      if index == 1: self._bulk_write_op = ("pcie_write", address, payload * 4)
+      elif index == 2: self._bulk_read_op = ("pcie_read", address, payload * 4)
+      else:
+        assert index == 0 and byte_en
+        offset = (byte_en & -byte_en).bit_length() - 1
+        size, is_write, is_cfg = byte_en.bit_count(), bool(fmt_type & 0x40), (fmt_type & 0xBE) == 0x04
+        if is_cfg:
+          bus, dev, fn, byte_addr = (address >> 24) & 0xFF, (address >> 19) & 0x1F, (address >> 16) & 0x7, address & 0xFFC
+          if is_write: self.state._cfg_write(bus, dev, fn, byte_addr + offset, (payload >> (8 * offset)) & ((1 << (8 * size))-1), size)
+          else: payload = int.from_bytes(self.state._get_cfg(bus, dev, fn)[byte_addr:byte_addr+4], 'little')
+        elif is_write:
+          self.state._pcie_dispatch(address + offset, (payload >> (8 * offset)) & ((1 << (8 * size))-1), size)
+        else: payload = (self.state._pcie_dispatch(address + offset, None, size) or 0) << (8 * offset)
+        self._f0_reply = struct.pack('<I', payload & 0xFFFFFFFF) + bytes(4)
+    else: raise ValueError(f"unsupported control OUT request 0x{request:02X}")
+
+  def control_read(self, request:int, length:int, value:int=0, index:int=0, timeout:int=1000) -> memoryview:
+    if request == 0xE4: data = self.state._xram_read(value, length)
+    elif request == 0xF0: data = self._f0_reply
+    else: raise ValueError(f"unsupported control IN request 0x{request:02X}")
+    return memoryview(data[:length])
+
+  def bulk_write(self, data:bytes, timeout:int=1000):
+    assert self._bulk_write_op is not None
+    op, address, size = self._bulk_write_op
+    assert len(data) == size
+    if op == "sram_write":
+      host_addr, region_size = self.state._dma_regions[address]
+      ctypes.memmove(host_addr, data, min(len(data), region_size))
+    elif op == "pcie_write": self.state._pcie_write(address, data)
+    else: raise RuntimeError(f"cannot bulk write for {op}")
+    self._bulk_write_op = None
+
+  def bulk_read(self, length:int, timeout:int=1000) -> memoryview:
+    assert self._bulk_read_op is not None
+    op, address, size = self._bulk_read_op
+    assert length == size
+    if op == "sram_read":
+      host_addr, region_size = self.state._dma_regions[address]
+      data = bytes((ctypes.c_ubyte * min(length, region_size)).from_address(host_addr))
+    elif op == "pcie_read": data = self.state._pcie_read(address, length)
+    else: raise RuntimeError(f"cannot bulk read for {op}")
+    self._bulk_read_op = None
+    return memoryview(data)

test/null/test_const_folding.py

@@ -1,6 +1,6 @@
 import unittest, itertools, math
 from tinygrad import Tensor, dtypes, Context
-from tinygrad.dtype import DType, ConstType
+from tinygrad.dtype import DType, ConstType, Invalid
 from tinygrad.uop.ops import Ops, UOp
 from test.helpers import full_rewrite
 import numpy as np
@@ -34,6 +34,24 @@ class TestUnaryOpsConstFolding(unittest.TestCase):
     x = x.clip(0, 1).realize()
     _check_ast_count(1, x.neg())
 
+class TestWeakConstFolding(unittest.TestCase):
+  def test_weakint_math(self):
+    out = (UOp.const(dtypes.weakint, 2**40) + UOp.const(dtypes.weakint, 2**40)).simplify()
+    self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakint, 2**41))
+
+  def test_float_unaries(self):
+    for dtype in (dtypes.weakint, dtypes.weakfloat):
+      for op in (Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL):
+        out = UOp.const(dtype, 4).alu(op).simplify()
+        self.assertEqual((out.op, out.dtype), (Ops.CONST, dtypes.weakfloat))
+
+  def test_weakfloat_math(self):
+    out = (UOp.const(dtypes.weakfloat, 1.25) + UOp.const(dtypes.weakfloat, 2.5)).simplify()
+    self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakfloat, 3.75))
+
+  def test_invalid_poison(self):
+    self.assertIs(UOp.const(dtypes.weakint, Invalid).alu(Ops.CDIV, UOp.const(dtypes.weakint, 0)).simplify().arg, Invalid)
+
 class TestBinaryOpsConstFolding(unittest.TestCase):
   def test_add_literal_zero(self):
     _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + 0)

test/null/test_dtype_spec.py

@@ -1,6 +1,6 @@
 import unittest, math, struct, operator
 from tinygrad import Tensor, Device
-from tinygrad.dtype import DTYPES_DICT, dtypes, truncate, float_to_fp16, float_to_bf16, _to_np_dtype, least_upper_dtype, least_upper_float
+from tinygrad.dtype import DTYPES_DICT, dtypes, Invalid, truncate, float_to_fp16, float_to_bf16, _to_np_dtype, least_upper_dtype, least_upper_float
 
 from tinygrad.helpers import getenv
 from hypothesis import given, settings, strategies as strat
@@ -57,6 +57,7 @@ class TestHelpers(unittest.TestCase):
 
   def test_from_py(self):
     assert dtypes.from_py(True) == dtypes.bool
+    assert dtypes.from_py(Invalid) == dtypes.bool
     assert dtypes.from_py(2) == dtypes.default_int
     assert dtypes.from_py(3.0) == dtypes.default_float
     assert dtypes.from_py([]) == dtypes.default_float

test/null/test_hcq_iface.py

@@ -84,22 +84,22 @@ class TestUSBMMIOInterface(unittest.TestCase):
     self.mmio[2] = 0xFE
     self.assertEqual(full_view[2], 0xFE)
 
-  def test_pcimem_byte(self):
+  def test_pcimem_dword(self):
     usb2 = MockUSB(bytearray(self.size))
-    mmio_pci = USBMMIOInterface(usb2, 0, self.size, fmt='B', pcimem=True)
-    mmio_pci[3] = 0x11
-    self.assertEqual(mmio_pci[3], 0x11)
-    self.assertEqual(usb2.mem[3], 0x11)
+    mmio_pci = USBMMIOInterface(usb2, 0, self.size, fmt='I', pcimem=True)
+    mmio_pci[3] = 0x11223344
+    self.assertEqual(mmio_pci[3], 0x11223344)
+    self.assertEqual(usb2.mem[12:16], b'\x44\x33\x22\x11')
 
   def test_pcimem_slice(self):
     usb3 = MockUSB(bytearray(self.size))
     mmio_pci = USBMMIOInterface(usb3, 0, self.size, fmt='B', pcimem=True)
-    values = [2, 3, 4]
-    mmio_pci[4:7] = values
-    raw = mmio_pci[4:7]
+    values = [2, 3, 4, 5]
+    mmio_pci[4:8] = values
+    raw = mmio_pci[4:8]
     self.assertIsInstance(raw, bytes)
     self.assertEqual(list(raw), values)
-    self.assertEqual([mmio_pci[i] for i in range(4, 7)], values)
+    self.assertEqual(list(usb3.mem[4:8]), values)
 
 if __name__ == "__main__":
   unittest.main()

test/null/test_pattern_matcher.py

@@ -51,8 +51,8 @@ class TestPatternMatcher(unittest.TestCase):
       ctx.append(True)
       assert len(x.src) == 0
       return x.replace(src=(UOp(Ops.NOOP),))
-    matcher = PatternMatcher([(UPat(Ops.CONST, src=(), name="x"), fxn)])
-    c1 = UOp.const(dtypes.float, 1.0)
+    matcher = PatternMatcher([(UPat(Ops.NOOP, src=(), name="x"), fxn)])
+    c1 = UOp(Ops.NOOP)
     # second rewrite shouldn't match anything
     ctx = []
     c1 = matcher.rewrite(c1, ctx)

test/null/test_uop_graph.py

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

test/null/test_uops.py

@@ -3,13 +3,50 @@ import unittest
 import numpy as np
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import Timing, Context, cdiv
-from tinygrad.dtype import dtypes, ConstFloat  # noqa: F401
+from tinygrad.dtype import dtypes, ConstFloat, Invalid  # noqa: F401
 from tinygrad.device import Device
-from tinygrad.uop.ops import Ops, ParamArg, UOp, UPat, exec_alu  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
-from tinygrad.uop.spec import spec_shared
+from tinygrad.uop.ops import Ops, ParamArg, UOp, UPat, dtype_from_uop, exec_alu  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
+from tinygrad.uop.spec import spec_program, spec_shared, type_verify
 from tinygrad.uop.symbolic import sym
 from test.helpers import eval_uop, to_uops_list
 
+class TestDTypeFromUOp(unittest.TestCase):
+  def test_broadcastable_promotion(self):
+    self.assertEqual(dtype_from_uop(Ops.ADD, (UOp.const(dtypes.float32, 1.0), UOp.const(dtypes.float16, 1.0)), None), dtypes.float32)
+    self.assertEqual(dtype_from_uop(Ops.MUL, (UOp.const(dtypes.int8, 1), UOp.const(dtypes.int32, 1)), None), dtypes.int32)
+    self.assertEqual(dtype_from_uop(Ops.ADD, (UOp.const(dtypes.weakint, 1), UOp.const(dtypes.int8, 1)), None), dtypes.int8)
+
+  def test_same_dtype_fast_path(self):
+    src = (UOp.const(dtypes.index, 1), UOp.const(dtypes.index, 2))
+    self.assertEqual(dtype_from_uop(Ops.ADD, src, None), dtypes.index)
+
+  def test_where_promotion(self):
+    cond = UOp.const(dtypes.bool, True)
+    self.assertEqual(dtype_from_uop(Ops.WHERE, (cond, UOp.const(dtypes.float32, 1.0), UOp.const(dtypes.float16, 1.0)), None), dtypes.float32)
+    idx = UOp.range(4, 0)
+    self.assertEqual(idx.valid(idx < 4).dtype, dtypes.index)
+
+  def test_const_dtype_from_value(self):
+    self.assertEqual(dtype_from_uop(Ops.CONST, (), True), dtypes.bool)
+    self.assertEqual(dtype_from_uop(Ops.CONST, (), 3), dtypes.weakint)
+    self.assertEqual(dtype_from_uop(Ops.CONST, (), ConstFloat(3.0)), dtypes.weakfloat)
+    self.assertEqual(dtype_from_uop(Ops.CONST, (), Invalid), dtypes.bool)
+    self.assertRaises(TypeError, dtype_from_uop, Ops.CONST, (), (1, 2))
+
+  @Context(SPEC=2)
+  def test_const_default_dtype_is_derived(self):
+    self.assertEqual(UOp(Ops.CONST, arg=3).dtype, dtypes.weakint)
+    self.assertEqual(UOp(Ops.CONST, arg=ConstFloat(3.0)).dtype, dtypes.weakfloat)
+    self.assertEqual(UOp(Ops.CONST, arg=True).dtype, dtypes.bool)
+    self.assertEqual(UOp(Ops.CONST, arg=Invalid).dtype, dtypes.bool)
+    # an explicit (strong) const dtype is legal until the field is removed
+    self.assertEqual(UOp.const(dtypes.int32, 3).dtype, dtypes.int32)
+
+  def test_weak_dtype_rejected_by_program_spec(self):
+    for weak, concrete, value in ((dtypes.weakint, dtypes.int32, 1), (dtypes.weakfloat, dtypes.float32, 1.0)):
+      with self.assertRaises(RuntimeError): type_verify(UOp.const(weak, value).sink(), spec_program)
+      type_verify(UOp.const(concrete, value).sink(), spec_program)
+
 class TestSafeCast(unittest.TestCase):
   def test_cast_folds(self):
     a = UOp.variable("a", 1, 10, dtype=dtypes.int32)
@@ -40,6 +77,12 @@ class TestExecALU(unittest.TestCase):
   def test_sqrt(self):
     self.assertEqual(exec_alu(Ops.SQRT, dtypes.float, (0.0,)), 0.0)
 
+  def test_invalid_poison(self):
+    # Invalid poisons any binary op regardless of result dtype: a comparison must not fold to a boolean
+    self.assertIs(exec_alu(Ops.CMPLT, dtypes.bool, (Invalid, 1)), Invalid)
+    self.assertIs(exec_alu(Ops.CMPNE, dtypes.bool, (Invalid, 1)), Invalid)
+    self.assertIs(exec_alu(Ops.ADD, dtypes.index, (Invalid, 1)), Invalid)
+
   def test_div(self):
     self.assertEqual(exec_alu(Ops.CDIV, dtypes.int8, (8, 2)), 4)
     self.assertEqual(exec_alu(Ops.CDIV, dtypes.int8, (7, 3)), 2)
@@ -368,5 +411,16 @@ class TestUOpRender(unittest.TestCase):
     u = UOp(Ops.STACK, dtype=dtypes.int, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)))
     self.assertEqual(u.render(), "{0,1,2}")
 
+class TestContiguousViewOffset(unittest.TestCase):
+  def _check(self, u, expected): self.assertEqual(u.contiguous_view_offset(), expected)
+
+  def test_simple(self): self._check(UOp.empty(10), 0)
+  def test_shrink(self): self._check(UOp.empty(10)[1:8], 1)
+  def test_2d(self): self._check(UOp.empty(2,5)[1, 2:4], 7)
+  def test_shrink_to_one(self): self._check(UOp.empty(10)[1], 1)
+  def test_expand_is_none(self): self._check(UOp.empty(1).expand(2), None)
+  def test_shrink_invalid(self): self._check(UOp.empty(4).pad((2,2))[0], None)
+  def test_strided(self): self._check(UOp.empty(4)[::2], None)
+
 if __name__ == '__main__':
   unittest.main()

test/null/test_viz.py

@@ -227,9 +227,9 @@ class TestViz(unittest.TestCase):
     pm = PatternMatcher([(UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4))])
     with save_viz() as viz:
       inner = UOp.const(dtypes.int, 3)
-      func = UOp(Ops.FUNCTION, src=(UOp(Ops.SINK, src=(inner,)),))
-      call = UOp(Ops.CALL, src=(func,))
-      graph_rewrite(call, TrackedPatternMatcher(pm.patterns), enter_calls=True)
+      call = UOp(Ops.CALL, src=(UOp(Ops.SINK, src=(inner,)),))
+      func = UOp(Ops.FUNCTION, src=(UOp(Ops.TUPLE, src=(call,)),))
+      graph_rewrite(func, TrackedPatternMatcher(pm.patterns), enter_calls=True)
     details = list(viz.get_details(0, 0))
     self.assertTrue(details[-1]["change"], "viz replay should detect change inside CALL")
 

test/opt/test_tensor_cores.py

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

test/test_tiny.py

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

test/unit/test_disk_tensor.py

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

test/unit/test_dtype_spec.py

@@ -149,6 +149,9 @@ class TestAutoCastType(unittest.TestCase):
   def tearDown(self):
     dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float
 
+  def test_int_sqrt(self):
+    _assert_eq(Tensor([1, 4, 9, 16]).sqrt(), dtypes.default_float, [1, 2, 3, 4])
+
   @given(strat.sampled_from([d for d in core_dtypes if dtypes.is_int(d) and d in supported_dtypes]))
   def test_int_to_float_unary_func(self, dtype):
     for func in [

test/unit/test_dtype_weak.py

@@ -0,0 +1,112 @@
+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()

test/unit/test_helpers.py

@@ -1,6 +1,6 @@
-import unittest
+import unittest, gc
 import numpy as np
-from tinygrad.helpers import polyN, is_numpy_ndarray
+from tinygrad.helpers import polyN, is_numpy_ndarray, disable_gc
 from tinygrad.tensor import Tensor
 
 class TestPolyN(unittest.TestCase):
@@ -11,5 +11,19 @@ class TestIsNumpyNdarray(unittest.TestCase):
   def test_tensor_numpy(self):
     self.assertTrue(is_numpy_ndarray(Tensor([1, 2, 3]).numpy()))
 
+class TestDisableGC(unittest.TestCase):
+  def test_recursive_decorator(self):
+    was_enabled = gc.isenabled()
+    @disable_gc()
+    def recurse(depth:int):
+      self.assertFalse(gc.isenabled())
+      if depth: recurse(depth-1)
+      self.assertFalse(gc.isenabled())
+    try:
+      recurse(2)
+      self.assertEqual(gc.isenabled(), was_enabled)
+    finally:
+      (gc.enable if was_enabled else gc.disable)()
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_invalid_tensor.py

@@ -1,5 +1,6 @@
 import unittest
 from tinygrad import Tensor
+from tinygrad.device import Buffer
 from tinygrad.dtype import Invalid, dtypes
 from tinygrad.engine.realize import run_linear
 
@@ -9,7 +10,7 @@ class TestInvalidTensor(unittest.TestCase):
     buf = out.uop.buffer
     buf.allocate()
     sentinel = memoryview(bytearray(b'\x42' * buf.nbytes))
-    buf.copyin(sentinel)
+    buf.copy_from(Buffer("PYTHON", buf.size, buf.dtype, opaque=sentinel))
     before = buf.as_memoryview().cast(out.dtype.fmt).tolist()
     run_linear(linear, var_vals)
     ret = buf.as_memoryview().cast(out.dtype.fmt).tolist()
@@ -64,6 +65,17 @@ class TestInvalidTensor(unittest.TestCase):
     out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid) > 1
     self._invalid_test_helper(out, [False, True, None, None])
 
+  def test_where_invalid_condition(self):
+    a, x = Tensor.arange(4), Tensor([0, 1, 2, 3])
+    bad = (a < 2).where(a, Invalid)
+    out = (bad < 1).logical_not().where(x + 10, x + 20)
+    self._invalid_test_helper(out, [20, 11, None, None])
+
+  def test_where_invalid_condition_bare(self):
+    cond = Tensor.full((4,), Invalid, dtype=dtypes.bool, buffer=False)
+    out = cond.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor([10.0, 20.0, 30.0, 40.0]))
+    self._invalid_test_helper(out, [None, None, None, None])
+
   def test_where_unary(self):
     mask = Tensor.arange(4) < 2
     out = mask.where(Tensor([1.0, 4.0, 9.0, 16.0]), Invalid).sqrt()
@@ -115,8 +127,6 @@ class TestInvalidTensor(unittest.TestCase):
     out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.int, buffer=False)).bitcast(dtypes.int)
     self._invalid_test_helper(out, [0x3f800000, 0x40000000, None, None])
 
-  # tensor indexing uses reduce, so the entire result becomes invalid
-  @unittest.expectedFailure
   def test_tensor_index(self):
     idx = (Tensor.arange(4) < 2).where(Tensor([0, 1, 2, 3]), Invalid)
     out = Tensor([1.0, 2.0, 3.0, 4.0])[idx]

tinygrad/callify.py

@@ -73,11 +73,6 @@ def contiguous_mops_to_view(c:UOp, src:UOp):
   # no symbolic shape
   if not all_int(c.shape): return None
 
-  # check if view is supported
-  from tinygrad.device import Device
-  devs = (src.device,) if isinstance(src.device, str) else src.device
-  if not all(hasattr(Device[d].allocator, "_offset") for d in devs): return None
-
   if buf.op is not Ops.MULTI and (view := _make_buffer_view(src)) is not None:
     view = (view.replace(dtype=c.dtype, arg=c.numel()) if c.op is Ops.BITCAST else view).reshape(c.shape)
     return c.replace(src=(view,)) if c.op is Ops.COPY else view

tinygrad/codegen/__init__.py

@@ -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[6]
+  in0, in1, out0 = u.arg[4]
   wmma = u.replace(src=(contract_axis(ctx, u.src[0], in0), contract_axis(ctx, u.src[1], in1), u.src[2]), tag=None)
   return unroll_axis(ctx, wmma, out0)
 
@@ -260,6 +260,13 @@ pm_add_local_buffers = PatternMatcher([
   (UPat(Ops.STAGE, name="x"), add_local_buffer),
 ])+pm_mops
 
+# float ALUs need a float operand
+# make that cast explicit before the decomps, which expand SIN/LOG2/EXP2 into float polynomials and assert a float operand
+pm_cast_float_alu = PatternMatcher([
+  (UPat((Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL), src=(UPat(name="x"),), name="u"),
+   lambda u,x: u.replace(src=(x.cast(u.dtype),)) if x.dtype != u.dtype else None),
+])
+
 def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
   if DEBUG >= 5: print(pyrender(ast))
@@ -327,6 +334,8 @@ 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)

tinygrad/codegen/opt/postrange.py

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

tinygrad/device.py

@@ -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, flat_mv, PROFILE, temp, colored
+from tinygrad.helpers import LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, PROFILE, temp, colored
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing
 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.copyin(memoryview(initial_value))
+        self.copy_from(Buffer("PYTHON", self.size, self.dtype, opaque=memoryview(bytearray(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,10 +135,8 @@ 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:
-        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)
+      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())
     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:
@@ -152,7 +150,6 @@ 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)
@@ -179,9 +176,7 @@ 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.nbytes)
-      self.copyout(memoryview(buf))
+    if self.is_allocated(): buf = bytearray(self.as_memoryview())
     return self.__class__, (self.device, self.size, self.dtype, None, self.options, buf, self.uop_refcount)
   @property
   def trace_num(self) -> int:
@@ -196,26 +191,22 @@ 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') and self.options is None:
-      return self.allocator._as_buffer(self._buf)
+    if (force_zero_copy or allow_zero_copy) and hasattr(self.allocator, '_as_buffer'): return self.allocator._as_buffer(self._buf)
     assert not force_zero_copy, "force zero copy was passed, but copy is required"
-    return self.copyout(memoryview(bytearray(self.nbytes)))
+    Buffer("PYTHON", self.size, self.dtype, opaque=(mv:=memoryview(bytearray(self.nbytes)))).copy_from(self)
+    return mv
   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 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)
+  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)
     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)
@@ -237,6 +228,8 @@ 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
@@ -245,7 +238,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):
+  def _offset(self, buf, size:int, offset:int): raise NotImplementedError("need offset")
   # 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
 

tinygrad/dtype.py

@@ -27,8 +27,6 @@ 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
@@ -100,7 +98,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): return dtypes.bool
+    if isinstance(x, (bool, InvalidType)): 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
@@ -156,6 +154,7 @@ 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", "")):
@@ -208,6 +207,7 @@ 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")))

tinygrad/engine/realize.py

@@ -166,9 +166,8 @@ 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 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))
+      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))
   return None
 
 def exec_kernel(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
@@ -216,10 +215,11 @@ 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 None
+  return time.perf_counter() - st
 
 # flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src
 pm_flatten_linear = PatternMatcher([

tinygrad/helpers.py

@@ -572,6 +572,8 @@ class tqdm(Generic[T]):
 def trange(n:int, **kwargs) -> tqdm[int]: return tqdm(range(n), total=n, **kwargs)
 
 class disable_gc(contextlib.ContextDecorator):
+  # ContextDecorator otherwise reuses self, so recursive calls overwrite _was_enabled.
+  def _recreate_cm(self): return type(self)()
   def __enter__(self):
     self._was_enabled = gc.isenabled()
     if self._was_enabled: gc.disable()

tinygrad/mixin/creation.py

@@ -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 None
-    val = cls.const(dt or (fill_value.dtype if isinstance(fill_value, UOp) else dtypes.from_py(fill_value)), fill_value)
+    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)
     val = val.reshape((1,)*len(new_shape)).expand(new_shape)
     return val.clone(device=device) if buffer else val
 

tinygrad/mixin/dtype.py

@@ -1,5 +1,6 @@
 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
@@ -29,7 +30,7 @@ class DTypeMixin:
     print(t.dtype, t.numpy())
     ```
     """
-    return self if self.dtype == (dt:=to_dtype(dtype)) else self._wrap_uop(self._uop.cast(dt))
+    return self if self.dtype == (dt:=to_dtype(dtype)) else self._wrap_uop(self._uop.alu(Ops.CAST, arg=dt))
 
   def bitcast(self, dtype:DTypeLike) -> Self:
     """
@@ -44,7 +45,9 @@ class DTypeMixin:
     print(t.dtype, t.numpy())
     ```
     """
-    return self if self.dtype == (dt:=to_dtype(dtype)) else self._wrap_uop(self._uop.bitcast(dt))
+    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))
 
   def element_size(self) -> int:
     """
@@ -55,6 +58,7 @@ 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:

tinygrad/mixin/elementwise.py

@@ -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, least_upper_float
+from tinygrad.dtype import dtypes, ConstType, PyConst, least_upper_dtype
 from tinygrad.helpers import argfix, polyN
 from tinygrad.mixin.creation import CreationMixin
 
@@ -49,7 +49,9 @@ class ElementwiseMixin(CreationMixin):
     """
     Returns a contiguous tensor.
     """
-    return self._wrap_uop(self._uop.contiguous(**kwargs))
+    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))
 
   def contiguous_backward(self) -> Self:
     """
@@ -437,9 +439,6 @@ 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.
@@ -448,7 +447,7 @@ class ElementwiseMixin(CreationMixin):
     print(Tensor([1., 2., 3., 4.]).reciprocal().numpy())
     ```
     """
-    return self._ensure_float().alu(Ops.RECIPROCAL)
+    return self.alu(Ops.RECIPROCAL)
 
   def trunc(self) -> Self:
     """
@@ -468,7 +467,7 @@ class ElementwiseMixin(CreationMixin):
     print(Tensor([1., 2., 3., 4.]).sqrt().numpy())
     ```
     """
-    return self._ensure_float().alu(Ops.SQRT)
+    return self.alu(Ops.SQRT)
 
   def sin(self) -> Self:
     """
@@ -478,7 +477,7 @@ class ElementwiseMixin(CreationMixin):
     print(Tensor([0., math.pi/2, math.pi, 3*math.pi/2, 2*math.pi]).sin().numpy())
     ```
     """
-    return self._ensure_float().alu(Ops.SIN)
+    return self.alu(Ops.SIN)
 
   def cos(self) -> Self:
     """
@@ -515,7 +514,7 @@ class ElementwiseMixin(CreationMixin):
     print(Tensor([1., 2., 4., 8.]).log2().numpy())
     ```
     """
-    return self._ensure_float().alu(Ops.LOG2)
+    return self.alu(Ops.LOG2)
 
   def exp2(self) -> Self:
     """
@@ -527,7 +526,7 @@ class ElementwiseMixin(CreationMixin):
     print(Tensor([0., 1., 2., 3.]).exp2().numpy())
     ```
     """
-    return self._ensure_float().alu(Ops.EXP2)
+    return self.alu(Ops.EXP2)
 
   def pow(self, x: Self | ConstType, reverse: bool = False) -> Self:
     """

tinygrad/mixin/rand.py

@@ -60,7 +60,7 @@ class RandMixin(OpMixin):
     ```
     """
     dt = to_dtype(dtype or dtypes.default_float)
-    if not dtypes.is_float(dt): raise ValueError(f"rand only supports float dtypes, got {dt}")
+    if not dtypes.is_float(dt) or dt in dtypes.weaks: raise ValueError(f"rand only supports concrete 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))

tinygrad/nn/onnx.py

@@ -423,9 +423,6 @@ 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),

tinygrad/nn/state.py

@@ -31,7 +31,8 @@ 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, "F16":dtypes.float16, "BF16":dtypes.bfloat16, "F32":dtypes.float32, "F64":dtypes.float64}
+               "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}
 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]:

tinygrad/renderer/__init__.py

@@ -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[1]) // u.arg[5] * mults
+        flops += 2 * prod(u.arg[0]) // u.arg[3] * mults
     return Estimates(flops, lds, sum(mem.values()))
 
 class Renderer:

tinygrad/renderer/cstyle.py

@@ -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"__{x.arg[0]}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]})"),
+  (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(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,9 +98,13 @@ 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)
 
-# (name, dims, dtype_in, dtype_out, device, threads, upcast_axes, reduce_axes)
+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)
 def wmma_args(uops:list[UOp]):
-  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)
+  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)
 
 class CStyleLanguage(Renderer):
   kernel_typedef: str = "void"
@@ -358,7 +362,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]]}))"),
@@ -366,7 +370,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(
@@ -438,7 +442,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
@@ -484,9 +488,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"__{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.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.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)})"),
@@ -512,8 +516,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: 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),
+      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),
     # bfloat16 constant casting
     (UPat.cvar('x', dtypes.bfloat16), lambda x: cast_float_to_bf16(UOp.const(dtypes.float, x.arg))),
   ])
@@ -552,7 +556,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"}

tinygrad/renderer/llvmir.py

@@ -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[1]
+  N,M,K = wmma.arg[0]
   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[2]]}(" + ", ".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[1]]}(" + ", ".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,29 +253,30 @@ 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: UOp(Ops.WMMA, src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]), arg=x.arg)
+          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
           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: 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),
+          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),
       ])
     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(UOp(Ops.WMMA,
+        (UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(x.replace(
           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[:6], (*x.arg[6][:2], ((0, x.max_numel()*2),)), *x.arg[7:])).index(i*2)
+            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 i in range(x.max_numel()))) if x.max_numel() == 8 else None),
-        (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)
+        (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]))
           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: 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)
+        (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)))
             .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: UOp(Ops.WMMA, src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]), arg=x.arg)
+          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
           if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None)
       ])
 

tinygrad/renderer/ptx.py

@@ -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[1], wmma.arg[2], wmma.arg[3]
+  (N, M, K), dtype_in, dtype_out = wmma.arg[0], wmma.arg[1], wmma.dtype
 
   for src, regs in zip(wmma.src, ctx.wmma_r):
     for i, reg in enumerate(regs): # pack input and acc registers

tinygrad/runtime/graph/hcq.py

@@ -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

tinygrad/runtime/ops_amd.py

@@ -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__send_interrupt_after_write_confirm, cache_flush=True)
+                       self.pm4.int_sel__mec_release_mem__none, 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,8 +490,6 @@ class AMDCopyQueue(HWQueue):
     if (dev:=signal.owner) is not None and signal.is_timeline and not dev.is_am():
       self.q(self.sdma.SDMA_OP_FENCE | fence_flags, *data64_le(dev.queue_event_mailbox_ptr), dev.queue_event.event_id)
       self.q(self.sdma.SDMA_OP_TRAP, self.sdma.SDMA_PKT_TRAP_INT_CONTEXT_INT_CONTEXT(dev.queue_event.event_id))
-    elif dev is not None and dev.is_am(): self.q(self.sdma.SDMA_OP_TRAP, 0)
-
     return self
 
   def wait(self, signal:AMDSignal, value:sint=0):
@@ -642,7 +640,7 @@ class AMDProgram(HCQProgram):
 
 class AMDAllocator(HCQAllocator['AMDDevice']):
   def __init__(self, dev:AMDDevice):
-    super().__init__(dev, copy_bufs=getattr(dev.iface, 'copy_bufs', None), max_copyout_size=0x1000 if dev.is_usb() else None,
+    super().__init__(dev, copy_bufs=getattr(dev.iface, 'copy_bufs', 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:
@@ -650,11 +648,10 @@ class AMDAllocator(HCQAllocator['AMDDevice']):
 
   def _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque)
 
-  def _map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+  def _do_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,
@@ -680,11 +677,12 @@ class AMDQueueDesc:
     try:
       self.write_ptr[0] = self.put_value
 
-      # Ensure all prior writes are visible to the GPU.
-      System.memory_barrier()
+      # 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()
 
-      # 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
@@ -917,11 +915,10 @@ 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), 0x800
+    self.sys_buf, self.sys_next_off = self._dma_region(ctrl_addr=0xa000, sys_addr=0x820000, size=0x1000), 0x200
     self.cq_buf = self._dma_region(ctrl_addr=0xb800, sys_addr=0x822000, size=0x1000)
 
   def _dma_region(self, ctrl_addr, sys_addr, size):
@@ -929,19 +926,14 @@ 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:
-    # 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:
+    # 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:
       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,), {})

tinygrad/runtime/ops_cpu.py

@@ -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 _map(self, buf:HCQBuffer):
+  def _do_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 _map returns a view wrapper, nothing to release
+  def _unmap(self, mb): pass  # CPU _do_map returns a view wrapper, nothing to release
 
 class CPUDevice(HCQCompiled):
   def __init__(self, device:str=""):

tinygrad/runtime/ops_hip.py

@@ -66,3 +66,4 @@ 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)

tinygrad/runtime/ops_null.py

@@ -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, offset:int, size:int): pass
+  def _offset(self, buf, size:int, offset:int): pass
 
 class NullGraph(MultiGraphRunner):
   def __call__(self, input_uops:tuple[UOp, ...], var_vals:dict[str, int], wait=False) -> float|None:

tinygrad/runtime/ops_nv.py

@@ -342,7 +342,7 @@ class NVAllocator(HCQAllocator['NVDevice']):
 
   def _do_free(self, opaque:HCQBuffer, options:BufferSpec): self.dev.iface.free(opaque)
 
-  def _map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+  def _do_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"

tinygrad/runtime/ops_python.py

@@ -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
-from tinygrad.device import Compiled, Compiler, Allocator
+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.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[1], first_src_dtype, u.arg[4], u.arg[5]
+          dims, dtype_in, device, threads = u.arg[0], first_src_dtype, u.arg[2], u.arg[3]
           wmma_helper = functools.partial(generic_wmma_helper, src_values, warp_size)
           # TODO: refactor these to a shared TensorCoreLayout
           if device == "METAL":
@@ -223,8 +223,13 @@ class PythonRenderer(Renderer):
 
 class PythonAllocator(Allocator['PythonDevice']):
   def _alloc(self, size, options): return memoryview(bytearray(size))
-  def _copyin(self, dest, src:memoryview): dest[:] = src
-  def _copyout(self, dest:memoryview, src): dest[:] = src
+  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]
 
 class PythonDevice(Compiled):
   def __init__(self, device:str):

tinygrad/runtime/ops_rdma.py

@@ -14,11 +14,11 @@ class RDMACopyQueue(HWQueue):
     super().__init__()
 
   def _wqe_data(self, buf:HCQBuffer, sz:int, nic:RDMADevice) -> bytes:
-    cast(HCQAllocatorBase, nic.allocator).map(buf)
+    cast(HCQAllocatorBase, nic.allocator)._map(buf)
     return struct.pack('>IIQ', sz, buf.mappings[nic].meta, buf.mappings[nic].va_addr + (buf.va_addr - buf.base.va_addr))
 
   def encode_ring(self, hwq:HWQueue, dev:HCQCompiled, iface:MLXIface, qp:MLXQP, cq_buf:HCQBuffer, head:sint, ring_uar:bool=False):
-    for buf in [iface.dbr_buf, cq_buf] + ([iface.uar_buf] if ring_uar else []): cast(HCQAllocator, dev.allocator).map(buf)
+    for buf in [iface.dbr_buf, cq_buf] + ([iface.uar_buf] if ring_uar else []): cast(HCQAllocator, dev.allocator)._map(buf)
     hwq.write(iface.dbr_buf.offset(qp.qp_dbr + (4 if ring_uar else 0)), to_be32(head + 1))
     if ring_uar: hwq.write(iface.uar_buf.offset(0x800), to_be64(((head << 8) | 0x0a) << 32 | ((qp.qp_info['qpn'] << 8) | 2)), b64=True)
     hwq.poll_bit(cq_buf.offset((head & (qp.cq_size - 1)) * 64 + 60, 4), ((head >> (qp.cq_size.bit_length() - 1)) & 1) << 24, mask=0x01000000)
@@ -72,7 +72,7 @@ class MLXIface(PCIIfaceBase):
 class RDMAAllocator(HCQAllocatorBase):
   def __init__(self, dev:RDMADevice): super().__init__(dev, batch_cnt=0)
 
-  def _map(self, buf:HCQBuffer) -> HCQBuffer:
+  def _do_map(self, buf:HCQBuffer) -> HCQBuffer:
     owner = unwrap(buf.base.owner)
     bar, paddrs = owner.iface.pci_dev.bar_info(owner.iface.vram_bar)[0], buf.base.meta.mapping.paddrs  # type: ignore[attr-defined]
     page_sz = (2 << 20) if min(sz for _, sz in paddrs) >= (2 << 20) else (4 << 10)

tinygrad/runtime/support/c.py

@@ -97,7 +97,7 @@ class DLL(ctypes.CDLL):
       libpaths = {"posix": [d for d in os.environ.get('LD_LIBRARY_PATH', '').split(os.pathsep) if d] + ["/usr/lib64", "/usr/lib", "/usr/local/lib"],
                   "nt": os.environ['PATH'].split(os.pathsep),
                   "darwin": ["/opt/homebrew/lib", f"/System/Library/Frameworks/{p}.framework", f"/System/Library/PrivateFrameworks/{p}.framework"],
-                  'linux': ['/lib', '/lib64', f"/lib/{sysconfig.get_config_var('MULTIARCH')}", "/usr/lib/wsl/lib/"]}
+                  'linux': ["/usr/lib/wsl/lib/", '/lib', '/lib64', f"/lib/{sysconfig.get_config_var('MULTIARCH')}"]}
       if (pth:=pathlib.Path(p)).is_absolute():
         if pth.is_file(): return p
         else: continue

tinygrad/runtime/support/hcq.py

@@ -409,7 +409,7 @@ class HCQCompiled(Compiled, Generic[SignalType]):
 
     if signal_t is not None:
       # Map signals if any
-      for sig_page in HCQCompiled.signal_pages[self.peer_group]: cast(HCQAllocator, self.allocator).map(sig_page)
+      for sig_page in HCQCompiled.signal_pages[self.peer_group]: cast(HCQAllocator, self.allocator)._map(sig_page)
 
       self.sigalloc_size = sigalloc_size
       self.timeline_signal, self._shadow_timeline_signal = self.new_signal(value=0, is_timeline=True), self.new_signal(value=0, is_timeline=True)
@@ -453,7 +453,7 @@ class HCQCompiled(Compiled, Generic[SignalType]):
     if not HCQCompiled.signal_pool[pg:=self.peer_group]:
       HCQCompiled.signal_pages[pg].append(alc:=self.allocator.alloc(self.sigalloc_size, BufferSpec(host=True, uncached=True, cpu_access=True)))
       HCQCompiled.signal_pool[pg] += [alc.offset(offset=off, size=16) for off in range(0, alc.size, 16)]
-      for dev in HCQCompiled.peer_groups[pg]: cast(HCQAllocator, dev.allocator).map(alc)
+      for dev in HCQCompiled.peer_groups[pg]: cast(HCQAllocator, dev.allocator)._map(alc)
     return self.signal_t(base_buf=HCQCompiled.signal_pool[pg].pop(), owner=self, **kwargs)
 
   def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent.
@@ -546,20 +546,21 @@ class HCQAllocatorBase(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
   This class implements basic copy operations following the HCQ API, utilizing both types of `HWQueue`.
   """
 
-  def __init__(self, dev:HCQDeviceType, batch_size:int=(2 << 20), batch_cnt:int=32, copy_bufs=None, max_copyout_size:int|None=None, **kwargs):
+  def __init__(self, dev:HCQDeviceType, batch_size:int=(2 << 20), batch_cnt:int=32, copy_bufs=None, **kwargs):
     super().__init__(dev, **kwargs)
     self.b = copy_bufs or [self._alloc(batch_size, BufferSpec(host=True)) for _ in range(batch_cnt)]
-    self.b_timeline, self.b_next, self.max_copyout_size = [0] * len(self.b), 0, max_copyout_size
+    self.b_timeline, self.b_next = [0] * len(self.b), 0
 
-  def map(self, buf:HCQBuffer):
-    if self.dev in buf.mapped_devs: return
+  def _map(self, buf:HCQBuffer) -> HCQBuffer:
+    if self.dev in buf.mapped_devs: return buf
     if buf.owner is None: raise RuntimeError(f"map failed: buffer {buf.va_addr} has no owner, it's a virtual buffer")
-    if not hasattr(self, '_map'): raise NotImplementedError("map failed: no method implemented")
+    if not hasattr(self, '_do_map'): raise NotImplementedError("map failed: no method implemented")
 
     # Since it's unified memory space, any buffer mapping is valid for all devices after successful map.
     # Devices can save mappings and internal metadata as a new buffer.
-    if (mb:=self._map(buf)) is not None: buf.mappings[self.dev] = mb
+    if (mb:=self._do_map(buf)) is not None: buf.mappings[self.dev] = mb
     buf.mapped_devs.append(self.dev)
+    return buf
 
   @suppress_finalizing
   def _free(self, buf:HCQBuffer, options:BufferSpec|None=None):
@@ -617,7 +618,7 @@ class HCQAllocator(HCQAllocatorBase, Generic[HCQDeviceType]):
 
     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,
                      dev_suff="SDMA:0"):
-      for i in range(0, dest.nbytes, cp_size:=(self.max_copyout_size or self.b[0].size)):
+      for i in range(0, dest.nbytes, cp_size:=self.b[0].size):
         self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \
                                   .copy(self.b[0], src.offset(i), lsize:=min(cp_size, dest.nbytes-i)) \
                                   .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
@@ -627,7 +628,7 @@ class HCQAllocator(HCQAllocatorBase, Generic[HCQDeviceType]):
   def _transfer(self, dest:HCQBuffer, src:HCQBuffer, sz:int, src_dev:HCQDeviceType, dest_dev:HCQDeviceType):
     if src_dev.peer_group != dest_dev.peer_group: return src_dev.rdma_dev().allocator._transfer(dest, src, sz, src_dev, dest_dev)
 
-    cast(HCQAllocator, src_dev.allocator).map(dest)
+    cast(HCQAllocator, src_dev.allocator)._map(dest)
 
     assert src_dev.hw_copy_queue_t is not None
     with hcq_profile(src_dev, queue_type=src_dev.hw_copy_queue_t, desc=TracingKey(f"{src_dev.device} -> {dest_dev.device}", ret=sz), enabled=PROFILE,

tinygrad/runtime/support/nv/nvdev.py

@@ -102,7 +102,7 @@ class NVDev:
     self.include("dev_fb", "tu102")
     self.include("dev_gc6_island", "ga102")
 
-    if (needs_reset:=self.reg("NV_PFB_PRI_MMU_WPR2_ADDR_HI").read() != 0):
+    if self.reg("NV_PFB_PRI_MMU_WPR2_ADDR_HI").read() != 0:
       self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) & ~pci.PCI_COMMAND_MASTER, 2)
       if DEBUG >= 2: print(f"nv {self.devfmt}: WPR2 is up. Issuing a full reset.", flush=True)
       self.pci_dev.reset()
@@ -118,7 +118,7 @@ class NVDev:
     self.flcn:NV_FLCN|NV_FLCN_COT = NV_FLCN_COT(self) if self.fmc_boot else NV_FLCN(self)
     self.gsp:NV_GSP = NV_GSP(self)
 
-    if needs_reset: self.flcn.wait_for_reset()
+    self.flcn.wait_for_reset()
 
   def _early_mmu_init(self):
     self.include("dev_vm", "tu102")

tinygrad/runtime/support/system.py

@@ -4,7 +4,7 @@ from tinygrad.helpers import round_up, getenv, OSX, temp, ceildiv, unwrap, fetch
 from tinygrad.runtime.autogen import libc, pci, vfio, iokit, corefoundation
 from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface, HCQBuffer, hcq_filter_visible_devices
 from tinygrad.runtime.support.memory import VirtMapping, AddrSpace, BumpAllocator
-from tinygrad.runtime.support.usb import USB3, CustomASM24Controller, ASM24Controller, USBMMIOInterface
+from tinygrad.runtime.support.usb import USB3, CustomASM24Controller, USBMMIOInterface
 
 MAP_FIXED, MAP_FIXED_NOREPLACE = 0x10, 0x100000
 MAP_LOCKED, MAP_POPULATE, MAP_NORESERVE = 0 if OSX else 0x2000, getattr(mmap, "MAP_POPULATE", 0 if OSX else 0x008000), 0x400
@@ -85,7 +85,7 @@ class _System:
     except IndexError: raise RuntimeError(f"{device}:{dev_id} does not exist ({pluralize('device', len(ds))} available)")
     return cl(device[:2], pcibus)
 
-  def pci_setup_usb_bars(self, usb:CustomASM24Controller|ASM24Controller, gpu_bus:int, mem_base:int, pref_mem_base:int) -> dict[int, tuple[int, int]]:
+  def pci_setup_usb_bars(self, usb:CustomASM24Controller, gpu_bus:int, mem_base:int, pref_mem_base:int) -> dict[int, tuple[int, int]]:
     for bus in range(gpu_bus):
       # All 3 values must be written at the same time.
       buses = (0 << 0) | ((bus+1) << 8) | ((gpu_bus) << 16)
@@ -221,9 +221,9 @@ class USBPCIDevice(PCIDevice):
   def __init__(self, devpref:str, dev, pcibus):
     self.pcibus, self.peer_group = pcibus, f"USBPCIDevice_{pcibus}"
     self.lock_fd = System.flock_acquire(f"{devpref.lower()}_{pcibus.lower()}.lock")
-    usb = USB3(dev, 0x81, 0x83, 0x02, 0x04)
+    usb = USB3(dev)
     if DEBUG >= 1: print(f"am {self.pcibus}: product string: {usb.product!r}")
-    self.usb: CustomASM24Controller | ASM24Controller = CustomASM24Controller(usb) if usb.is_custom else ASM24Controller(usb)
+    self.usb: CustomASM24Controller = CustomASM24Controller(usb)
     self._bar_info = System.pci_setup_usb_bars(self.usb, gpu_bus=4, mem_base=0x10000000, pref_mem_base=(32 << 30))
     self.sram = BumpAllocator(size=0x80000, wrap=False) # asm24 controller sram
 

tinygrad/runtime/support/usb.py

@@ -1,7 +1,6 @@
-import ctypes, struct, dataclasses, array, itertools, time, functools
-from typing import Sequence
+import ctypes, struct, time, functools, itertools
 from tinygrad.runtime.autogen import libusb
-from tinygrad.helpers import DEBUG, DEV, to_mv, round_up, OSX, getenv, ceildiv
+from tinygrad.helpers import DEBUG, DEV, to_mv, round_up, ceildiv
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support import c
 
@@ -32,13 +31,10 @@ class USB3:
     libusb.libusb_free_device_list(devs, 1)
     return ret
 
-  def __init__(self, dev:c.POINTER[libusb.struct_libusb_device], ep_data_in:int, ep_stat_in:int, ep_data_out:int, ep_cmd_out:int,
-               max_streams:int=31, use_bot=False):
-    self.ep_data_in, self.ep_stat_in, self.ep_data_out, self.ep_cmd_out = ep_data_in, ep_stat_in, ep_data_out, ep_cmd_out
-    self.max_streams, self.use_bot = max_streams, use_bot
-    self._transferred = ctypes.c_int(0)
-    self._bulk_in_buf, self._bulk_in_mv = alloc_cbuffer(4 << 20)
-    self._bulk_out_buf, self._bulk_out_mv = alloc_cbuffer(4 << 20)
+  def __init__(self, dev:c.POINTER[libusb.struct_libusb_device], *args, **kwargs):
+    self._tags, self._transferred = itertools.count(1), ctypes.c_int(0)
+    self._bulk_buf, self._bulk_mv = alloc_cbuffer(4 << 20)
+    self._ctrl_buf, self._ctrl_mv = alloc_cbuffer(0x1000)
 
     self.handle = c.init_c_var(c.POINTER[libusb.struct_libusb_device_handle], lambda x: checked(libusb.libusb_open)(dev, x))
 
@@ -48,8 +44,7 @@ class USB3:
     checked(libusb.libusb_get_device_descriptor)(libusb.libusb_get_device(self.handle), ctypes.byref(_desc))
     _ret = checked(libusb.libusb_get_string_descriptor_ascii)(self.handle, _desc.iProduct, _buf, 256)
     self.product = bytes(_buf[:_ret]).decode("ascii", errors="replace")
-    self.is_custom = self.product.startswith("custom")
-    if self.is_custom: self.use_bot = use_bot = True
+    assert self.product.startswith("custom") or self.product.startswith("AS2462")
 
     # Detach kernel driver if needed
     if checked(libusb.libusb_kernel_driver_active)(self.handle, 0):
@@ -59,146 +54,41 @@ class USB3:
     # Set configuration and claim interface
     checked(libusb.libusb_set_configuration)(self.handle, 1)
     checked(libusb.libusb_claim_interface)(self.handle, 0)
+    checked(libusb.libusb_set_interface_alt_setting)(self.handle, 0, 0)
 
-    if use_bot:
-      checked(libusb.libusb_set_interface_alt_setting)(self.handle, 0, 0)
-      self._tag = 0
-    else:
-      checked(libusb.libusb_set_interface_alt_setting)(self.handle, 0, 1)
+  def control_write(self, request:int, value:int=0, index:int=0, data:bytes=b'', timeout:int=1000):
+    assert len(data) <= len(self._ctrl_mv)
+    self._ctrl_mv[:len(data)] = data
+    assert checked(libusb.libusb_control_transfer)(self.handle, 0x40, request, value, index, self._ctrl_buf, len(data), timeout) == len(data)
 
-      # Clear any stalled endpoints
-      all_eps = (self.ep_data_out, self.ep_data_in, self.ep_stat_in, self.ep_cmd_out)
-      for ep in all_eps: checked(libusb.libusb_clear_halt)(self.handle, ep)
+  def control_read(self, request:int, length:int, value:int=0, index:int=0, timeout:int=1000) -> memoryview:
+    assert length <= len(self._ctrl_mv)
+    assert checked(libusb.libusb_control_transfer)(self.handle, 0xC0, request, value, index, self._ctrl_buf, length, timeout) == length
+    return self._ctrl_mv[:length]
 
-      # Allocate streams
-      stream_eps = (ctypes.c_uint8 * 3)(self.ep_data_out, self.ep_data_in, self.ep_stat_in)
-      checked(libusb.libusb_alloc_streams)(self.handle, self.max_streams * len(stream_eps), stream_eps, len(stream_eps))
+  def bulk_write(self, payload:bytes, timeout:int=1000):
+    if len(payload) > len(self._bulk_mv): self._bulk_buf, self._bulk_mv = alloc_cbuffer(len(payload))
+    self._bulk_mv[:len(payload)] = payload
+    checked(libusb.libusb_bulk_transfer, "bulk OUT 0x02 failed") \
+      (self.handle, 0x02, self._bulk_buf, len(payload), self._transferred, timeout)
+    assert self._transferred.value == len(payload), f"bulk OUT short write: {self._transferred.value}/{len(payload)} bytes"
 
-      # Base cmd
-      cmd_template = bytes([0x01, 0x00, 0x00, 0x01, *([0] * 12), 0xE4, 0x24, 0x00, 0xB2, 0x1A, 0x00, 0x00, 0x00, *([0] * 8)])
+  def bulk_read(self, length:int, timeout:int=1000) -> memoryview:
+    if length > len(self._bulk_mv): self._bulk_buf, self._bulk_mv = alloc_cbuffer(length)
+    checked(libusb.libusb_bulk_transfer, "bulk IN 0x81 failed")(self.handle, 0x81, self._bulk_buf, length, self._transferred, timeout)
+    return self._bulk_mv[:self._transferred.value]
 
-      # Init pools
-      self.tr = {ep: [libusb.libusb_alloc_transfer(0) for _ in range(self.max_streams)] for ep in all_eps}
-
-      self.buf_cmd = [(ctypes.c_uint8 * len(cmd_template))(*cmd_template) for _ in range(self.max_streams)]
-      self.buf_stat = [(ctypes.c_uint8 * 64)() for _ in range(self.max_streams)]
-      self.buf_data_in = [(ctypes.c_uint8 * 0x1000)() for _ in range(self.max_streams)]
-      self.buf_data_out = [(ctypes.c_uint8 * 0x80000)() for _ in range(self.max_streams)]
-      self.buf_data_out_mvs = [to_mv(ctypes.addressof(self.buf_data_out[i]), 0x80000) for i in range(self.max_streams)]
-
-      for slot in range(self.max_streams): struct.pack_into(">B", self.buf_cmd[slot], 3, slot + 1)
-
-  def _prep_transfer(self, tr, ep, stream_id, buf, length):
-    tr.contents.dev_handle, tr.contents.endpoint, tr.contents.length, tr.contents.buffer = self.handle, ep, length, buf
-    tr.contents.status, tr.contents.flags, tr.contents.timeout, tr.contents.num_iso_packets = 0xff, 0, 1000, 0
-    tr.contents.type = (libusb.LIBUSB_TRANSFER_TYPE_BULK_STREAM if stream_id is not None else libusb.LIBUSB_TRANSFER_TYPE_BULK)
-    if stream_id is not None: libusb.libusb_transfer_set_stream_id(tr, stream_id)
-    return tr
-
-  def _submit_and_wait(self, cmds):
-    for tr in cmds: checked(libusb.libusb_submit_transfer)(tr)
-
-    running = len(cmds)
-    while running:
-      checked(libusb.libusb_handle_events)(USB3.ctx())
-      running = len(cmds)
-      for tr in cmds:
-        if tr.contents.status == libusb.LIBUSB_TRANSFER_COMPLETED: running -= 1
-        elif tr.contents.status != 0xFF: raise RuntimeError(f"EP 0x{tr.contents.endpoint:02X} error: {tr.contents.status}")
-
-  def _bulk_out(self, ep: int, payload: bytes, timeout: int = 1000):
-    if len(payload) > len(self._bulk_out_mv): self._bulk_out_buf, self._bulk_out_mv = alloc_cbuffer(len(payload))
-    self._bulk_out_mv[:len(payload)] = payload
-    checked(libusb.libusb_bulk_transfer, f"bulk OUT 0x{ep:02X} failed")(self.handle, ep, self._bulk_out_buf, len(payload), self._transferred, timeout)
-    assert self._transferred.value == len(payload), f"bulk OUT short write on 0x{ep:02X}: {self._transferred.value}/{len(payload)} bytes"
-
-  def _bulk_in(self, ep: int, length: int, timeout: int = 1000) -> memoryview:
-    if length > len(self._bulk_in_mv): self._bulk_in_buf, self._bulk_in_mv = alloc_cbuffer(length)
-    checked(libusb.libusb_bulk_transfer, f"bulk IN 0x{ep:02X} failed")(self.handle, ep, self._bulk_in_buf, length, self._transferred, timeout)
-    return self._bulk_in_mv[:self._transferred.value]
-
-  def send_batch(self, cdbs:list[bytes], idata:list[int]|None=None, odata:list[bytes|None]|None=None) -> list[bytes|None]:
-    idata, odata = idata or [0] * len(cdbs), odata or [None] * len(cdbs)
-    results:list[bytes|None] = []
-    tr_window, op_window = [], []
-
-    for idx, (cdb, rlen, send_data) in enumerate(zip(cdbs, idata, odata)):
-      if self.use_bot:
-        dir_in = rlen > 0
-        data_len = rlen if dir_in else (len(send_data) if send_data is not None else 0)
-        assert not (rlen > 0 and send_data is not None), "BOT mode only supports either read or write per command"
-
-        # CBW
-        self._tag += 1
-        flags = 0x80 if dir_in else 0x00
-        cbw = struct.pack("<IIIBBB", 0x43425355, self._tag, data_len, flags, 0, len(cdb)) + cdb + b"\x00" * (16 - len(cdb))
-        self._bulk_out(self.ep_data_out, cbw)
-
-        # DAT
-        if dir_in:
-          results.append(bytes(self._bulk_in(self.ep_data_in, rlen)))
-        else:
-          if send_data is not None:
-            self._bulk_out(self.ep_data_out, send_data)
-          results.append(None)
-
-        # CSW
-        sig, rtag, residue, status = struct.unpack("<IIIB", self._bulk_in(self.ep_data_in, 13, timeout=2000))
-        assert sig == 0x53425355, f"Bad CSW signature 0x{sig:08X}, expected 0x53425355"
-        assert rtag == self._tag, f"CSW tag mismatch: got {rtag}, expected {self._tag}"
-        assert status == 0, f"SCSI command failed, CSW status=0x{status:02X}, residue={residue}"
-      else:
-        # allocate slot and stream. stream is 1-based
-        slot, stream = idx % self.max_streams, (idx % self.max_streams) + 1
-
-        # build cmd packet
-        self.buf_cmd[slot][16:16+len(cdb)] = list(cdb)
-
-        # cmd + stat transfers
-        tr_window.append(self._prep_transfer(self.tr[self.ep_cmd_out][slot], self.ep_cmd_out, None, self.buf_cmd[slot], len(self.buf_cmd[slot])))
-        tr_window.append(self._prep_transfer(self.tr[self.ep_stat_in][slot], self.ep_stat_in, stream, self.buf_stat[slot], 64))
-
-        if rlen:
-          if rlen > len(self.buf_data_in[slot]): self.buf_data_in[slot] = (ctypes.c_uint8 * round_up(rlen, 0x1000))()
-          tr_window.append(self._prep_transfer(self.tr[self.ep_data_in][slot], self.ep_data_in, stream, self.buf_data_in[slot], rlen))
-
-        if send_data is not None:
-          if len(send_data) > len(self.buf_data_out[slot]):
-            self.buf_data_out[slot] = (ctypes.c_uint8 * len(send_data))()
-            self.buf_data_out_mvs[slot] = to_mv(ctypes.addressof(self.buf_data_out[slot]), len(send_data))
-
-          self.buf_data_out_mvs[slot][:len(send_data)] = bytes(send_data)
-          tr_window.append(self._prep_transfer(self.tr[self.ep_data_out][slot], self.ep_data_out, stream, self.buf_data_out[slot], len(send_data)))
-
-        op_window.append((idx, slot, rlen))
-        if (idx + 1 == len(cdbs)) or len(op_window) >= self.max_streams:
-          self._submit_and_wait(tr_window)
-          for idx, slot, rlen in op_window: results.append(bytes(self.buf_data_in[slot][:rlen]) if rlen else None)
-          tr_window = []
-
-    return results
-
-@dataclasses.dataclass(frozen=True)
-class WriteOp: addr:int; data:bytes; ignore_cache:bool=True # noqa: E702
-
-@dataclasses.dataclass(frozen=True)
-class ReadOp: addr:int; size:int # noqa: E702
-
-@dataclasses.dataclass(frozen=True)
-class ScsiWriteOp: data:bytes; lba:int=0 # noqa: E702
+  # NOTE: keep it for flash.py
+  def send_batch(self, cdbs:list[bytes], odata:list[bytes|None]|None=None):
+    for cdb, data in zip(cdbs, odata or [None] * len(cdbs)):
+      self.bulk_write(struct.pack("<IIIBBB16s", 0x43425355, tag:=next(self._tags), len(data) if data is not None else 0, 0, 0, len(cdb), cdb))
+      if data is not None: self.bulk_write(data)
+      sig, rtag, _, status = struct.unpack("<IIIB", self.bulk_read(13, timeout=2000))
+      assert (sig, rtag, status) == (0x53425355, tag, 0)
 
 class CustomASM24Controller:
-  def __init__(self, usb:USB3|None=None):
-    if not usb:
-      devs = USB3.list_devices(0xADD1, 0x0001)
-      assert len(devs), "no ASM24 controller found"
-      self.usb = USB3(devs[0][0], 0x81, 0x83, 0x02, 0x04, use_bot=True)
-    else: self.usb = usb
-    self._pci_cacheable: list[tuple[int, int]] = []
-    self._pci_cache: dict[int, int|None] = {}
-
-    self._f0_out_buf, self._f0_out_mv = alloc_cbuffer(0x1000) # for f0 and e4, allocate big enough for e4
-    self._f0_in_buf, _ = alloc_cbuffer(8)
+  def __init__(self, usb:USB3):
+    self.usb = usb
 
     # Custom firmware now boots with PCIe off. Power it on before probing the link.
     ltssm = self.read(0xB450, 1)[0]
@@ -206,44 +96,31 @@ class CustomASM24Controller:
     ltssm = self.read(0xB450, 1)[0]
     if ltssm != 0x78: raise RuntimeError(f"PCIe link not up (LTSSM=0x{ltssm:02X}), custom firmware not ready")
 
-  def set_pcie_power(self, enabled:bool, timeout:int=10000):
-    checked(libusb.libusb_control_transfer,
-            f"F3 PCIe power {'on' if enabled else 'off'} failed")(self.usb.handle, 0x40, 0xF3, int(enabled), 0, None, 0, timeout)
-
-  # === PCIe TLP via 0xF0 vendor command ===
+  def set_pcie_power(self, enabled:bool, timeout:int=10000): self.usb.control_write(0xF3, value=int(enabled), timeout=timeout)
 
   def _f0_out(self, fmt_type:int, byte_en:int, address:int, value:int, mode:int=0):
-    struct.pack_into('<III', self._f0_out_mv, 0, address & 0xFFFFFFFF, address >> 32, value)
-    ret = libusb.libusb_control_transfer(self.usb.handle, 0x40, 0xF0, fmt_type | (byte_en << 8), mode & 0x03, self._f0_out_buf, 12, 5000)
-    assert ret == 12, f"F0 OUT failed: {ret}"
+    self.usb.control_write(0xF0, fmt_type | (byte_en << 8), mode & 0x03, struct.pack('<III', address & 0xFFFFFFFF, address >> 32, value), 5000)
 
   def _f0_in(self) -> tuple[int, int, int]:
-    ret = libusb.libusb_control_transfer(self.usb.handle, 0xC0, 0xF0, 0, 0, self._f0_in_buf, 8, 5000)
-    assert ret == 8, f"F0 IN failed: {ret}"
-    return struct.unpack_from('<I', self._f0_in_buf, 0)[0], (self._f0_in_buf[4] >> 5) & 0x7, self._f0_in_buf[7]
-
-  def _is_pci_cacheable(self, addr:int) -> bool: return any(x <= addr <= x + sz for x, sz in self._pci_cacheable)
+    data = self.usb.control_read(0xF0, 8, timeout=5000)
+    return struct.unpack_from('<I', data)[0], (data[4] >> 5) & 0x7, data[7]
 
   def pcie_request(self, fmt_type:int, address:int, value:int|None=None, size:int=4, cnt:int=10):
-    if fmt_type == 0x60 and size == 4 and self._is_pci_cacheable(address) and self._pci_cache.get(address) == value: return
     assert size > 0 and size <= 4, f"Invalid size {size}"
     if DEBUG >= 5: print("pcie_request", hex(fmt_type), hex(address), value, size)
 
     offset = address & 0x3
     byte_en = ((1 << size) - 1) << offset
-    self._pci_cache[address] = value if size == 4 and fmt_type == 0x60 else None
-
     self._f0_out(fmt_type, byte_en, address & ~0x3, (value << (8 * offset)) if value is not None else 0)
 
-    # Fast path: memory writes and messages don't return completions (same logic as ASM24Controller).
+    # Fast path: memory writes and messages don't return completions.
     if ((fmt_type & 0b11011111) == 0b01000000) or ((fmt_type & 0b10111000) == 0b00110000): return
 
     # Read TLPs and config writes: read completion via 0xF0 IN. Retry on error/timeout.
     data, cpl_status, ret_status = self._f0_in()
     if ret_status != 0:
       time.sleep(0.001)  # TODO: this sleep is very picky
-      if cnt > 0:
-        return self.pcie_request(fmt_type, address, value, size, cnt=cnt-1)
+      if cnt > 0: return self.pcie_request(fmt_type, address, value, size, cnt=cnt-1)
       raise RuntimeError(f"TLP error after retries: ret_status={ret_status}, address={address:#x}")
 
     if cpl_status:
@@ -258,219 +135,69 @@ class CustomASM24Controller:
     address = (bus << 24) | (dev << 19) | (fn << 16) | (byte_addr & 0xfff)
     return self.pcie_request(fmt_type, address, value, size)
 
-  def pcie_mem_req(self, address:int, value:int|None=None, size:int=4):
-    return self.pcie_request(0x60 if value is not None else 0x20, address, value, size)
-
-  def pcie_mem_write(self, address:int, values:list[int], size:int):
+  def pcie_mem_write(self, address:int, data:bytes):
     """Streaming PCIe memory write via 0xF0 mode 1 + bulk OUT. Data is little-endian dwords on the wire."""
-    if not values: return
-    self._f0_out(0x60, 0x0F, address, len(values), mode=1)
-    self.usb._bulk_out(0x02, struct.pack(f'<{len(values)}I', *values))
+    if not data: return
+    assert len(data) % 4 == 0, f"pcie_mem_write requires 4-byte aligned size, got {len(data)}"
+    self._f0_out(0x60, 0x0F, address, len(data) // 4, mode=1)
+    self.usb.bulk_write(data)
 
-  def pcie_mem_read(self, address:int, nbytes:int) -> bytes:
+  def pcie_mem_read(self, address:int, nbytes:int) -> memoryview:
     """Streaming PCIe memory read via 0xF0 mode 2 + bulk IN. Returns little-endian bytes."""
     assert nbytes % 4 == 0, f"pcie_mem_read requires 4-byte aligned size, got {nbytes}"
     self._f0_out(0x20, 0x0F, address, nbytes // 4, mode=2)
-    return self.usb._bulk_in(0x81, nbytes, timeout=30000)
+    return self.usb.bulk_read(nbytes, timeout=30000)
 
-  # === XDATA read/write (0xE4/0xE5 vendor control transfers) ===
-
-  def read(self, base_addr:int, length:int, **kwargs) -> bytes:
+  def read(self, base_addr:int, length:int) -> bytes:
     """Read from chip XDATA via vendor control IN (bRequest=0xE4). wValue=addr, wLength=size."""
     result = b''
     for off in range(0, length, 0xFF):
       chunk = min(0xFF, length - off)
-      ret = libusb.libusb_control_transfer(self.usb.handle, 0xC0, 0xE4, base_addr + off, 0, self._f0_out_buf, chunk, 1000)
-      assert ret == chunk, f"read(0x{base_addr + off:04X}, {chunk}) failed: {ret}"
-      result += bytes(self._f0_out_buf[:ret])
-    return result[:length]
+      result += self.usb.control_read(0xE4, chunk, value=base_addr + off)
+    return result
 
-  def write(self, base_addr:int, data:bytes, **kwargs):
+  def write(self, base_addr:int, data:bytes):
     """Write to chip XDATA via vendor control OUT (bRequest=0xE5). wValue=addr, wIndex=val."""
-    for off, val in enumerate(data):
-      checked(libusb.libusb_control_transfer,
-              f"write(0x{base_addr + off:04X}, 0x{val:02X}) failed")(self.usb.handle, 0x40, 0xE5, base_addr + off, val, None, 0, 1000)
+    for off, val in enumerate(data): self.usb.control_write(0xE5, value=base_addr + off, index=val)
 
-  def scsi_write(self, buf:bytes, lba:int=0):
+  def scsi_write(self, buf:bytes):
     """Write to SRAM via 0xF2 vendor command + bulk OUT."""
     buf_padded = buf + b'\x00' * (round_up(len(buf), 512) - len(buf))
     sectors = len(buf_padded) // 512
-    num_slots = round_up(len(buf_padded), 0x4000) // 0x4000  # 16KB per slot
-    # 0xF2 OUT: wValue=sectors, wIndex=start_slot|(num_slots<<8)
+    num_slots = ceildiv(len(buf_padded), 0x4000)  # 16KB per slot
     windex = (num_slots & 0xFF) << 8
-    checked(libusb.libusb_control_transfer, "F2 setup failed")(self.usb.handle, 0x40, 0xF2, sectors, windex, None, 0, 1000)
-    self.usb._bulk_out(0x02, buf_padded)
+    self.usb.control_write(0xF2, value=sectors, index=windex)
+    self.usb.bulk_write(buf_padded)
 
   def scsi_read_arm(self, size:int):
     windex = (ceildiv(size, 0x4000) & 0xFF) << 8
-    checked(libusb.libusb_control_transfer,
-            "F2 read arm failed")(self.usb.handle, 0x40, 0xF2, (ceildiv(size, 512) & 0x7FFF) | 0x8000, windex, None, 0, 1000)
+    self.usb.control_write(0xF2, value=(ceildiv(size, 512) & 0x7FFF) | 0x8000, index=windex)
 
-  def scsi_read(self, size:int) -> memoryview: return self.usb._bulk_in(0x81, round_up(size, 512), timeout=10000)[:size]
-
-class ASM24Controller:
-  def __init__(self, usb:USB3|None=None):
-    if not usb:
-      devs = USB3.list_devices(0xADD1, 0x0001)
-      assert len(devs), "no ASM24 controller found"
-      self.usb = USB3(devs[0][0], 0x81, 0x83, 0x02, 0x04, use_bot=bool(getenv("USE_BOT", 0)))
-    else: self.usb = usb
-    self._cache: dict[int, int|None] = {}
-    self._pci_cacheable: list[tuple[int, int]] = []
-    self._pci_cache: dict[int, int|None] = {}
-
-    # Init controller.
-    self.exec_ops([WriteOp(0x54b, b' '), WriteOp(0x54e, b'\x04'), WriteOp(0x5a8, b'\x02'), WriteOp(0x5f8, b'\x04'),
-      WriteOp(0x7ec, b'\x01\x00\x00\x00'), WriteOp(0xc422, b'\x02'), WriteOp(0x0, b'\x33')])
-
-  def exec_ops(self, ops:Sequence[WriteOp|ReadOp|ScsiWriteOp]):
-    cdbs:list[bytes] = []
-    idata:list[int] = []
-    odata:list[bytes|None] = []
-
-    def _add_req(cdb:bytes, i:int, o:bytes|None):
-      nonlocal cdbs, idata, odata
-      cdbs, idata, odata = cdbs + [cdb], idata + [i], odata + [o]
-
-    for op in ops:
-      if isinstance(op, WriteOp):
-        for off, value in enumerate(op.data):
-          addr = ((op.addr + off) & 0x1FFFF) | 0x500000
-          if not op.ignore_cache and self._cache.get(addr) == value: continue
-          _add_req(struct.pack('>BBBHB', 0xE5, value, addr >> 16, addr & 0xFFFF, 0), 0, None)
-          self._cache[addr] = value
-      elif isinstance(op, ReadOp):
-        assert op.size <= 0xff
-        addr = (op.addr & 0x1FFFF) | 0x500000
-        _add_req(struct.pack('>BBBHB', 0xE4, op.size, addr >> 16, addr & 0xFFFF, 0), op.size, None)
-        for i in range(op.size): self._cache[addr + i] = None
-      elif isinstance(op, ScsiWriteOp):
-        sectors = round_up(len(op.data), 512) // 512
-        _add_req(struct.pack('>BBQIBB', 0x8A, 0, op.lba, sectors, 0, 0), 0, op.data+b'\x00'*((sectors*512)-len(op.data)))
-
-    return self.usb.send_batch(cdbs, idata, odata)
-
-  def write(self, base_addr:int, data:bytes, ignore_cache:bool=True): return self.exec_ops([WriteOp(base_addr, data, ignore_cache)])
-
-  def scsi_write(self, buf:bytes, lba:int=0):
-    if len(buf) > 0x4000: buf += b'\x00' * (round_up(len(buf), 0x10000) - len(buf))
-
-    for i in range(0, len(buf), 0x10000):
-      self.exec_ops([ScsiWriteOp(buf[i:i+0x10000], lba), WriteOp(0x171, b'\xff\xff\xff', ignore_cache=True)])
-      self.exec_ops([WriteOp(0xce6e, b'\x00\x00', ignore_cache=True)])
-
-    if len(buf) > 0x4000:
-      for i in range(4): self.exec_ops([WriteOp(0xce40 + i, b'\x00', ignore_cache=True)])
-
-  def read(self, base_addr:int, length:int, stride:int=0xff) -> bytes:
-    parts = self.exec_ops([ReadOp(base_addr + off, min(stride, length - off)) for off in range(0, length, stride)])
-    return b''.join(p or b'' for p in parts)[:length]
-
-  def _is_pci_cacheable(self, addr:int) -> bool: return any(x <= addr <= x + sz for x, sz in self._pci_cacheable)
-  def pcie_prep_request(self, fmt_type:int, address:int, value:int|None=None, size:int=4) -> list[WriteOp]:
-    if fmt_type == 0x60 and size == 4 and self._is_pci_cacheable(address) and self._pci_cache.get(address) == value: return []
-
-    assert fmt_type >> 8 == 0 and size > 0 and size <= 4, f"Invalid fmt_type {fmt_type} or size {size}"
-    if DEBUG >= 5: print("pcie_request", hex(fmt_type), hex(address), value, size)
-
-    masked_address, offset = address & 0xFFFFFFFC, address & 0x3
-    assert size + offset <= 4 and (value is None or value >> (8 * size) == 0)
-    self._pci_cache[address] = value if size == 4 and fmt_type == 0x60 else None
-
-    return ([WriteOp(0xB220, struct.pack('>I', value << (8 * offset)), ignore_cache=False)] if value is not None else []) + \
-      [WriteOp(0xB218, struct.pack('>I', masked_address), ignore_cache=False), WriteOp(0xB21c, struct.pack('>I', address>>32), ignore_cache=False),
-       WriteOp(0xB217, bytes([((1 << size) - 1) << offset]), ignore_cache=False), WriteOp(0xB210, bytes([fmt_type]), ignore_cache=False),
-       WriteOp(0xB254, b"\x0f", ignore_cache=True), WriteOp(0xB296, b"\x04", ignore_cache=True)]
-
-  def pcie_request(self, fmt_type, address, value=None, size=4, cnt=10):
-    self.exec_ops(self.pcie_prep_request(fmt_type, address, value, size))
-
-    # Fast path for write requests
-    if ((fmt_type & 0b11011111) == 0b01000000) or ((fmt_type & 0b10111000) == 0b00110000): return
-
-    while (stat:=self.read(0xB296, 1)[0]) & 2 == 0:
-      if stat & 1:
-        self.write(0xB296, bytes([0x01]))
-        if cnt > 0: return self.pcie_request(fmt_type, address, value, size, cnt=cnt-1)
-    assert stat == 2, f"stat read 2 was {stat}"
-
-    # Retrieve completion data from Link Status (0xB22A, 0xB22B)
-    b284 = self.read(0xB284, 1)[0]
-    completion = struct.unpack('>H', self.read(0xB22A, 2))
-
-    # Validate completion status based on PCIe request typ
-    # Completion TLPs for configuration requests always have a byte count of 4.
-    assert completion[0] & 0xfff == (4 if (fmt_type & 0xbe == 0x04) else size)
-
-    # Extract completion status field
-    status = (completion[0] >> 13) & 0x7
-
-    # Handle completion errors or inconsistencies
-    if status or ((fmt_type & 0xbe == 0x04) and (((value is None) and (not (b284 & 0x01))) or ((value is not None) and (b284 & 0x01)))):
-      status_map = {0b001: f"Unsupported Request: invalid address/function (target might not be reachable): {address:#x}",
-                    0b100: "Completer Abort: abort due to internal error", 0b010: "Configuration Request Retry Status: configuration space busy"}
-      raise RuntimeError(f"TLP status: {status_map.get(status, 'Reserved (0b{:03b})'.format(status))}")
-
-    if value is None: return (struct.unpack('>I', self.read(0xB220, 4))[0] >> (8 * (address & 0x3))) & ((1 << (8 * size)) - 1)
-
-  def pcie_cfg_req(self, byte_addr, bus=1, dev=0, fn=0, value=None, size=4):
-    assert byte_addr >> 12 == 0 and bus >> 8 == 0 and dev >> 5 == 0 and fn >> 3 == 0, f"Invalid byte_addr {byte_addr}, bus {bus}, dev {dev}, fn {fn}"
-
-    fmt_type = (0x44 if value is not None else 0x4) | int(bus > 0)
-    address = (bus << 24) | (dev << 19) | (fn << 16) | (byte_addr & 0xfff)
-    return self.pcie_request(fmt_type, address, value, size)
-
-  def pcie_mem_req(self, address, value=None, size=4): return self.pcie_request(0x60 if value is not None else 0x20, address, value, size)
-
-  def pcie_mem_write(self, address, values, size):
-    ops = [self.pcie_prep_request(0x60, address + i * size, value, size) for i, value in enumerate(values)]
-
-    # Send in batches of 4 for OSX and 16 for Linux (benchmarked values)
-    for i in range(0, len(ops), bs:=(4 if OSX else 16)): self.exec_ops(list(itertools.chain.from_iterable(ops[i:i+bs])))
+  def scsi_read(self, size:int) -> memoryview: return self.usb.bulk_read(round_up(size, 512), timeout=10000)[:size]
 
 class USBMMIOInterface(MMIOInterface):
   def __init__(self, usb, addr, size, fmt, pcimem=True): # pylint: disable=super-init-not-called
-    self.usb, self.addr, self.nbytes, self.fmt, self.pcimem, self.el_sz = usb, addr, size, fmt, pcimem, struct.calcsize(fmt)
+    self.usb, self.addr, self.nbytes, self.fmt, self.el_sz, self.pcimem = usb, addr, size, fmt, struct.calcsize(fmt), pcimem
 
-  def __getitem__(self, index): return self._access_items(index)
-  def __setitem__(self, index, val): self._access_items(index, val)
+  def _off_from_index(self, index):
+    if isinstance(index, slice): return ((index.start or 0) * self.el_sz, ((index.stop or len(self))-(index.start or 0)) * self.el_sz)
+    return (index * self.el_sz, self.el_sz)
 
-  def _access_items(self, index, val=None):
-    if isinstance(index, slice): return self._acc((index.start or 0) * self.el_sz, ((index.stop or len(self))-(index.start or 0)) * self.el_sz, val)
-    return self._acc_one(index * self.el_sz, self.el_sz, val) if self.pcimem else self._acc(index * self.el_sz, self.el_sz, val)
+  def __getitem__(self, index):
+    off, sz = self._off_from_index(index)
+    if self.pcimem:
+      assert sz % 4 == 0 and off % 4 == 0, f"pcie_mem_read requires 4-byte aligned access, got off={off}, sz={sz}"
+      data = self.usb.pcie_mem_read(self.addr + off, sz)
+    else: data = self.usb.scsi_read(sz) if self.addr == 0xf000 else self.usb.read(self.addr + off, sz)
+    return int.from_bytes(data, "little") if sz == self.el_sz else data
+
+  def __setitem__(self, index, data):
+    off, _ = self._off_from_index(index)
+    data = struct.pack(self.fmt, data) if isinstance(data, int) else bytes(data)
+    if not self.pcimem: self.usb.scsi_write(data) if self.addr == 0xf000 else self.usb.write(self.addr + off, data)
+    else: self.usb.pcie_mem_write(self.addr+off, data)
 
   def view(self, offset:int=0, size:int|None=None, fmt=None):
-    return USBMMIOInterface(self.usb, self.addr+offset, size or (self.nbytes - offset), fmt=fmt or self.fmt, pcimem=self.pcimem)
-
-  def _acc_size(self, sz): return next(x for x in [('I', 4), ('H', 2), ('B', 1)] if sz % x[1] == 0)
-
-  def _acc_one(self, off, sz, val=None):
-    upper = 0 if sz < 8 else self.usb.pcie_mem_req(self.addr + off + 4, val if val is None else (val >> 32), 4)
-    lower = self.usb.pcie_mem_req(self.addr + off, val if val is None else val & 0xffffffff, min(sz, 4))
-    if val is None: return lower | (upper << 32)
-
-  def _acc(self, off, sz, data=None):
-    if data is None: # read op
-      if not self.pcimem:
-        if self.addr == 0xf000 and hasattr(self.usb, 'scsi_read'): return self.usb.scsi_read(sz)
-        return int.from_bytes(self.usb.read(self.addr + off, sz), "little") if sz == self.el_sz else self.usb.read(self.addr + off, sz)
-
-      # Fast path: streaming PCIe read if controller supports it
-      if hasattr(self.usb, 'pcie_mem_read') and sz >= 4 and sz % 4 == 0:
-        return self.usb.pcie_mem_read(self.addr + off, sz)
-
-      acc, acc_size = self._acc_size(sz)
-      return bytes(array.array(acc, [self._acc_one(off + i * acc_size, acc_size) for i in range(sz // acc_size)]))
-
-    # write op
-    data = struct.pack(self.fmt, data) if isinstance(data, int) else bytes(data)
-
-    if not self.pcimem:
-      # Fast path for writing into buffer 0xf000
-      use_cache = 0xa800 <= self.addr <= 0xb000
-      return self.usb.scsi_write(bytes(data)) if self.addr == 0xf000 else self.usb.write(self.addr + off, bytes(data), ignore_cache=not use_cache)
-
-    _, acc_sz = self._acc_size(len(data) * struct.calcsize(self.fmt))
-    self.usb.pcie_mem_write(self.addr+off, [int.from_bytes(data[i:i+acc_sz], "little") for i in range(0, len(data), acc_sz)], acc_sz)
+    return USBMMIOInterface(self.usb, self.addr+offset, self.nbytes-offset if size is None else size, fmt=fmt or self.fmt, pcimem=self.pcimem)
 
 if DEV.interface.startswith("MOCK"): from test.mockgpu.usb import MockUSB3 as USB3  # type: ignore  # noqa: F811

tinygrad/schedule/indexing.py

@@ -45,6 +45,8 @@ class BufferizeOpts:
 class IndexingContext:
   realize_map: dict[UOp, None|list[int]] = field(default_factory=dict)
   range_map: dict[UOp, tuple[tuple[UOp, ...], tuple[UOp, ...]]] = field(default_factory=dict)
+  # loads reachable from each UOp memoized across matches
+  buf_cache: dict[UOp, frozenset[UOp]] = field(default_factory=dict)
 
   # create ranges
   range_idx: Iterator[int] = field(default_factory=itertools.count)

tinygrad/schedule/memory.py

@@ -1,5 +1,4 @@
 from collections import defaultdict
-from tinygrad.device import Device
 from tinygrad.helpers import NO_MEMORY_PLANNER, DEBUG, round_up
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.dtype import dtypes
@@ -13,7 +12,8 @@ def _collect_bufs(u:UOp) -> list[UOp]:
 def _can_plan(b:UOp, held_bufs:set[UOp]) -> bool:
   if b in held_bufs: return False
   devs = (b.device,) if isinstance(b.device, str) else b.device
-  return all(not d.startswith(("DISK", "TINYFS")) and hasattr(Device[d].allocator, "_offset") for d in devs)
+  # CL and WEBGPU do not support views, see explanation in contiguous_view_offset
+  return all(not d.startswith(("DISK", "TINYFS", "CL", "WEBGPU")) for d in devs)
 
 LaneKey = tuple[str, int]
 

tinygrad/schedule/rangeify.py

@@ -329,12 +329,11 @@ def limit_bufs(ctx:IndexingContext, root:UOp):
   device = device if isinstance(device, str) else device[0].split(":")[0]
   if not (MAX_BUFS:=MAX_KERNEL_BUFFERS.value or DEVICE_MAX_BUFS.get(device, 0)): return None
 
-  bufs: set[UOp] = set()
-  def gate_input(u:UOp):
-    # TODO: add cache to fix n^2
-    if is_load:=(u.op in {Ops.STAGE, Ops.AFTER, Ops.PARAM, Ops.MSELECT, Ops.MSTACK}): bufs.add(u)
-    return not is_load
-  root.toposort(gate=gate_input)
+  def visitor(u:UOp) -> frozenset[UOp]:
+    if u.op in {Ops.STAGE, Ops.AFTER, Ops.PARAM, Ops.MSELECT, Ops.MSTACK}: return frozenset((u,))
+    if len(u.src) == 1: return ctx.buf_cache[u.src[0]]
+    return frozenset().union(*[ctx.buf_cache[s] for s in u.src])
+  bufs = root.topovisit(visitor, ctx.buf_cache)
 
   if len(bufs) > MAX_BUFS - 1: # NOTE: this -1 is for the output buffer
     srcs = []