start disabling rules there's no tests for

2025-08-19 23:36:49 -07:00
499 changed files with 8208 additions and 48109 deletions
@@ -225,22 +225,13 @@ runs:
    - name: Install gpuocelot dependencies (MacOS)
      if: inputs.ocelot == 'true' && runner.os == 'macOS'
      shell: bash
-      run: |
-        pkgs=(cmake ninja llvm@15 zlib glew flex bison [email protected] zstd ncurses)
-        for f in "${pkgs[@]}"; do
-          brew ls --versions "$f" >/dev/null 2>&1 || brew install --quiet "$f"
-        done
-
-        # Fix boost 1.85 for gpuocelot
-        ln -s /opt/homebrew/opt/[email protected] /opt/homebrew/opt/boost || true
-        ln -s /opt/homebrew/opt/boost/lib/libboost_atomic-mt.dylib /opt/homebrew/opt/boost/lib/libboost_atomic.dylib || true
-        ln -s /opt/homebrew/opt/boost/lib/libboost_thread-mt.dylib /opt/homebrew/opt/boost/lib/libboost_thread.dylib || true
+      run: brew install --quiet cmake ninja llvm@15 zlib glew flex bison boost zstd ncurses
    - name: Cache gpuocelot
      if: inputs.ocelot == 'true'
      id: cache-build
      uses: actions/cache@v4
      env:
-        cache-name: cache-gpuocelot-build-1
+        cache-name: cache-gpuocelot-build
      with:
        path: ${{ github.workspace }}/gpuocelot/ocelot
        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.BUILD_CACHE_VERSION }}
@@ -253,13 +244,7 @@ runs:
        git checkout b16039dc940dc6bc4ea0a98380495769ff35ed99
        mkdir build
        cd build
-
-        CMAKE_ARGS="-Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5"
-        if [[ "${{ runner.os }}" == "macOS" ]]; then
-          CMAKE_ARGS="$CMAKE_ARGS -DBoost_INCLUDE_DIR=$(brew --prefix boost)/include -DBoost_LIBRARY_DIR=$(brew --prefix boost)/lib"
-        fi
-
-        cmake .. $CMAKE_ARGS
+        cmake .. -Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5
        ninja
    - name: Install gpuocelot
      if: inputs.ocelot == 'true'
@@ -1,91 +0,0 @@
-name: Autogen
-env:
-  # increment this when downloads substantially change to avoid the internet
-  DOWNLOAD_CACHE_VERSION: '12'
-  PYTHON_CACHE_VERSION: '3'
-  APT_CACHE_VERSION: '1'
-  BUILD_CACHE_VERSION: '1'
-  CAPTURE_PROCESS_REPLAY: 1
-  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-  PYTHONPATH: ${{ github.workspace }}
-
-on:
-  push:
-    branches:
-      - master
-  pull_request:
-    paths: 
-    - 'tinygrad/runtime/autogen/**/*'
-  workflow_dispatch:
-    paths: 
-    - 'tinygrad/runtime/autogen/**/*'
-
-jobs:
-  autogen:
-    name: Autogen
-    runs-on: ubuntu-24.04
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        opencl: 'true'
-        amd: 'true'
-        cuda: 'true'
-        webgpu: 'true'
-        llvm: 'true'
-    - name: Install autogen support packages
-      run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev
-    - name: Verify OpenCL autogen
-      run: |
-        cp tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
-        ./autogen_stubs.sh opencl
-        diff /tmp/opencl.py.bak tinygrad/runtime/autogen/opencl.py
-    - name: Verify CUDA autogen
-      run: |
-        cp tinygrad/runtime/autogen/cuda.py /tmp/cuda.py.bak
-        cp tinygrad/runtime/autogen/nv_gpu.py /tmp/nv_gpu.py.bak
-        ./autogen_stubs.sh cuda
-        ./autogen_stubs.sh nv
-        diff /tmp/cuda.py.bak tinygrad/runtime/autogen/cuda.py
-        diff /tmp/nv_gpu.py.bak tinygrad/runtime/autogen/nv_gpu.py
-    - name: Verify AMD autogen
-      run: |
-        cp tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
-        cp tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
-        cp tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
-        cp tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
-        cp tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
-        ./autogen_stubs.sh hsa
-        ./autogen_stubs.sh kfd
-        ./autogen_stubs.sh comgr
-        ./autogen_stubs.sh amd
-        ./autogen_stubs.sh sqtt
-        diff /tmp/hsa.py.bak tinygrad/runtime/autogen/hsa.py
-        diff /tmp/kfd.py.bak tinygrad/runtime/autogen/kfd.py
-        diff /tmp/comgr.py.bak tinygrad/runtime/autogen/comgr.py
-        diff /tmp/amd_gpu.py.bak tinygrad/runtime/autogen/amd_gpu.py
-        diff /tmp/sqtt.py.bak tinygrad/runtime/autogen/sqtt.py
-    - name: Verify Linux autogen
-      run: |
-        cp tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
-        cp tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
-        cp tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
-        ./autogen_stubs.sh libc
-        ./autogen_stubs.sh io_uring
-        ./autogen_stubs.sh ib
-        diff /tmp/libc.py.bak tinygrad/runtime/autogen/libc.py
-        diff /tmp/io_uring.py.bak tinygrad/runtime/autogen/io_uring.py
-        diff /tmp/ib.py.bak tinygrad/runtime/autogen/ib.py
-    - name: Verify WebGPU autogen
-      run: |
-        cp tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak
-        ./autogen_stubs.sh webgpu
-        diff /tmp/webgpu.py.bak tinygrad/runtime/autogen/webgpu.py
-    - name: Verify LLVM autogen
-      run: |
-        cp tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
-        ./autogen_stubs.sh llvm
-        diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
@@ -28,7 +28,7 @@ jobs:
      # since sudo is required for usbgpu on macos, move the cache to a new location, as some of the files are owned by root
      PYTHONPYCACHEPREFIX: /tmp/tiny_python_pycache
    runs-on: [self-hosted, macOS]
-    timeout-minutes: 60
+    timeout-minutes: 20
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -52,28 +52,24 @@ jobs:
    - name: reset process replay
      run: python3.11 test/external/process_replay/reset.py
    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion JIT=1 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
    - name: Run Stable Diffusion without fp16
-      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=900 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
+      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
    - name: Run Stable Diffusion v2
-      # TODO: very slow step time
-      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=10000 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
+      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
    # process replay can't capture this, the graph is too large
-    # TODO: too slow
-    # - name: Run SDXL
-    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+    - name: Run SDXL
+      run: BENCHMARK_LOG=stable_diffusion_xl CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run model inference benchmark
      run: METAL=1 python3.11 test/external/external_model_benchmark.py
    - name: Test speed vs torch
      run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
    - name: Test tensor cores
-      run: METAL=1 python3.11 test/opt/test_tensor_cores.py
+      run: METAL=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Test AMX tensor cores
      run: |
-        DEBUG=2 CPU=1 CPU_LLVM=0 AMX=1 python3.11 test/opt/test_tensor_cores.py
-        DEBUG=2 CPU=1 CPU_LLVM=1 AMX=1 python3.11 test/opt/test_tensor_cores.py
-        DEBUG=2 CPU=1 CPU_LLVM=0 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
-        DEBUG=2 CPU=1 CPU_LLVM=1 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
+        DEBUG=2 CPU=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
+        DEBUG=2 LLVM=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
    - name: Run Tensor Core GEMM (float)
      run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
    - name: Run Tensor Core GEMM (half)
@@ -101,7 +97,7 @@ jobs:
    - name: Run GPT2
      run: |
        BENCHMARK_LOG=gpt2_nojit JIT=0 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        BENCHMARK_LOG=gpt2 JIT=1 ASSERT_MIN_STEP_TIME=13 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2 JIT=1 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
    - name: Run GPT2 w HALF
      run: BENCHMARK_LOG=gpt2_half HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
    - name: Run GPT2 w HALF/BEAM
@@ -110,27 +106,22 @@ jobs:
      run: BENCHMARK_LOG=olmoe python3.11 examples/olmoe.py
    - name: Train MNIST
      run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py | tee beautiful_mnist.txt
-
-    # NOTE: this is failing in CI. it is not failing on my machine and I don't really have a way to debug it
-    # the error is "RuntimeError: Internal Error (0000000e:Internal Error)"
-    #- name: Run 10 CIFAR training steps
-    #  run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=3000 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
-    #- name: Run 10 CIFAR training steps w HALF
-    #  run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=3000 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
-
+    - name: Run 10 CIFAR training steps
+      run: BENCHMARK_LOG=cifar_10steps JIT=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
+    - name: Run 10 CIFAR training steps w HALF
+      run: BENCHMARK_LOG=cifar_10steps_half JIT=2 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
    #- name: Run 10 CIFAR training steps w BF16
    #  run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
-    # TODO: too slow
-    # - name: Run 10 CIFAR training steps w winograd
-    #   run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+    - name: Run 10 CIFAR training steps w winograd
+      run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
    - name: UsbGPU boot time
      run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
    - name: UsbGPU tiny tests
      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
    - name: UsbGPU copy speeds
      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
-    #- name: UsbGPU openpilot test
-    #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB NOLOCALS=0 IMAGE=0 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+    - name: UsbGPU openpilot test
+      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB NOLOCALS=0 IMAGE=0 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (Mac)
@@ -167,7 +158,7 @@ jobs:
  testnvidiabenchmark:
    name: tinybox green Benchmark
    runs-on: [self-hosted, Linux, tinyboxgreen]
-    timeout-minutes: 60
+    timeout-minutes: 30
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -203,15 +194,15 @@ jobs:
      run: NV=1 python test/external/external_benchmark_multitensor_allreduce.py
    - name: Test tensor cores
      run: |
-        NV=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
-        NV=1 NV_PTX=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
+        NV=1 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        PTX=1 ALLOW_TF32=1 NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Run Tensor Core GEMM (CUDA)
      run: |
        CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
        CUDA=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
        CUDA=1 SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_tf32.txt
    - name: Run Tensor Core GEMM (PTX)
-      run: NV=1 NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
+      run: NV=1 PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
    - name: Run Tensor Core GEMM (NV)
      run: NV=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_nv.txt
    - name: Test NV=1
@@ -220,9 +211,8 @@ jobs:
      run: DEBUG=2 CUDA=1 python -m pytest -rA test/test_tiny.py
    - name: Run Stable Diffusion
      run: BENCHMARK_LOG=stable_diffusion NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
-    # TODO: too slow
-    # - name: Run SDXL
-    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+    - name: Run SDXL
+      run: BENCHMARK_LOG=stable_diffusion_xl CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run LLaMA
      run: |
        BENCHMARK_LOG=llama_nojit NV=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
@@ -246,9 +236,9 @@ jobs:
    - name: Run GPT2
      run: |
        BENCHMARK_LOG=gpt2_nojit NV=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        BENCHMARK_LOG=gpt2 NV=1 JIT=1 ASSERT_MIN_STEP_TIME=4 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2 NV=1 JIT=1 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
    - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=6 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
    - name: Run GPT2 w HALF/BEAM
      run: BENCHMARK_LOG=gpt2_half_beam NV=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
    - uses: actions/upload-artifact@v4
@@ -282,7 +272,7 @@ jobs:
  testmorenvidiabenchmark:
    name: tinybox green Training Benchmark
    runs-on: [self-hosted, Linux, tinyboxgreen]
-    timeout-minutes: 60
+    timeout-minutes: 20
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -307,33 +297,30 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: test/external/process_replay/reset.py
-    # TODO: too slow
-    # - name: Fuzz Padded Tensor Core GEMM (NV)
-    #   run: NV=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
-    # TODO: too slow
-    # - name: Fuzz Padded Tensor Core GEMM (PTX)
-    #   run: NV=1 NV_PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
+    - name: Fuzz Padded Tensor Core GEMM (NV)
+      run: NV=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
+    - name: Fuzz Padded Tensor Core GEMM (PTX)
+      run: NV=1 PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
    - name: Train MNIST
      run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
    - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: BENCHMARK_LOG=cifar_10steps NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=310 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
    - name: Run 10 CIFAR training steps w BF16
-      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=310 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
-    # TODO: too slow
-    # - name: Run 10 CIFAR training steps w winograd
-    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+      run: BENCHMARK_LOG=cifar_10steps_bf16 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+    - name: Run 10 CIFAR training steps w winograd
+      run: BENCHMARK_LOG=cifar_10steps_half_wino NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
+      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
    - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
+      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
    - name: Run MLPerf resnet eval on training data
      run: time BENCHMARK_LOG=resnet_eval NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
-    #- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-    #  run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
-    #- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-    #  run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
+    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+      run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
+    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
+      run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
    - name: Run 10 MLPerf Bert training steps (6 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
      run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
@@ -357,7 +344,7 @@ jobs:
  testamdbenchmark:
    name: tinybox red Benchmark
    runs-on: [self-hosted, Linux, tinybox]
-    timeout-minutes: 60
+    timeout-minutes: 20
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -407,8 +394,8 @@ jobs:
      run: AMD=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
    - name: Test tensor cores
      run: |
-        AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
-        AMD=1 AMD_LLVM=1 python3 test/opt/test_tensor_cores.py
+        AMD=1 AMD_LLVM=0 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+        AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
        AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
    - name: Run Tensor Core GEMM (AMD)
      run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
@@ -426,10 +413,9 @@ jobs:
    - name: Test AM warm start time
      run: time AMD=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
-    # TODO: too slow
-    # - name: Run SDXL
-    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+      run: BENCHMARK_LOG=stable_diffusion AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+    - name: Run SDXL
+      run: BENCHMARK_LOG=stable_diffusion_xl CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run LLaMA 7B
      run: |
        BENCHMARK_LOG=llama_nojit AMD=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
@@ -455,9 +441,9 @@ jobs:
    - name: Run GPT2
      run: |
        BENCHMARK_LOG=gpt2_nojit AMD=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        BENCHMARK_LOG=gpt2 AMD=1 JIT=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2 AMD=1 JIT=1 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
    - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half AMD=1 HALF=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      run: BENCHMARK_LOG=gpt2_half AMD=1 HALF=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
    - name: Run GPT2 w HALF/BEAM
      run: BENCHMARK_LOG=gpt2_half_beam AMD=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
    - uses: actions/upload-artifact@v4
@@ -488,7 +474,7 @@ jobs:
  testmoreamdbenchmark:
    name: tinybox red Training Benchmark
    runs-on: [self-hosted, Linux, tinybox]
-    timeout-minutes: 60
+    timeout-minutes: 30
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -520,20 +506,19 @@ jobs:
    - name: Train MNIST
      run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
    - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=330 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: BENCHMARK_LOG=cifar_10steps AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=330 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
-    # - name: Run 10 CIFAR training steps w BF16
-    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
-    # TODO: too slow
-    # - name: Run 10 CIFAR training steps w winograd
-    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+      run: BENCHMARK_LOG=cifar_10steps_half AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+    - name: Run 10 CIFAR training steps w BF16
+      run: BENCHMARK_LOG=cifar_10steps_bf16 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+    - name: Run 10 CIFAR training steps w winograd
+      run: BENCHMARK_LOG=cifar_10steps_half_wino AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
-    #- name: Run full CIFAR training steps w 6 GPUS
-    #  run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
-    #- name: Run full CIFAR training steps w 6 GPUS (REMOTE)
-    #  run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu_remote.txt
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
+    - name: Run full CIFAR training steps w 6 GPUS
+      run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
+    - name: Run full CIFAR training steps w 6 GPUS (REMOTE)
+      run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu_remote.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (AMD Training)
@@ -552,7 +537,7 @@ jobs:
  testmlperfamdbenchmark:
    name: tinybox red MLPerf Benchmark
    runs-on: [self-hosted, Linux, tinybox]
-    timeout-minutes: 60
+    timeout-minutes: 30
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -583,10 +568,10 @@ jobs:
      run: test/external/process_replay/reset.py
    - name: Run MLPerf resnet eval
      run: time BENCHMARK_LOG=resnet_eval AMD=1 MODEL=resnet python3 examples/mlperf/model_eval.py
-    #- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-    #  run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
-    #- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-    #  run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
+    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+      run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
+    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
+      run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
    - name: Run 10 MLPerf Bert training steps (6 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
      run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
@@ -625,21 +610,21 @@ jobs:
    - name: benchmark openpilot 0.9.9 dmonitoring
      run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: openpilot compile3 0.9.9 driving_vision
-      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=18 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
+      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.9.9 driving_policy
-      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=7 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
+      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
    - name: openpilot compile3 0.9.9 dmonitoring
-      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=12 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
+      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: openpilot compile3 Space Lab policy + vision
      run: |
-        PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
-        PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b
+        PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
+        PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b
    - name: benchmark MobileNetV2 on DSP
      run: |
        # generate quantized weights
        ln -s /data/home/tiny/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
        ln -s /data/home/tiny/tinygrad/testsig-*.so .
-        PYTHONPATH=. CC=clang-19 CPU=1 CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
+        PYTHONPATH=. CC=clang-19 CPU=1 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
        # benchmark on DSP with NOOPT=1, the devectorizer has issues
        PYTHONPATH=. CC=clang-19 DSP=1 DONT_REALIZE_EXPAND=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
    - name: Run process replay tests
@@ -658,7 +643,7 @@ jobs:
  testreddriverbenchmark:
    name: AM Benchmark
    runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 20
+    timeout-minutes: 15
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -694,8 +679,8 @@ jobs:
    # Fails on 9070
    # - name: Test tensor cores
    #   run: |
-    #     AMD=1 AMD_LLVM=0 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
-    #     AMD=1 AMD_LLVM=1 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
+    #     AMD=1 AMD_LLVM=0 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+    #     AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
    #     AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
    - name: Run Tensor Core GEMM (AMD)
      run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee am_matmul_amd.txt
@@ -703,12 +688,8 @@ jobs:
      run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
    - name: Test DISK copy time
      run: AMD=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
-    - name: Test CPU copy time
-      run: |
-        AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
-        AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
    # TODO: enable
    # - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
    #   run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
@@ -729,7 +710,7 @@ jobs:
  testgreendriverbenchmark:
    name: NV Benchmark
    runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 20
+    timeout-minutes: 15
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -761,19 +742,15 @@ jobs:
    - name: Test driver start time
      run: time DEBUG=3 NV=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Test tensor cores
-      run: NV=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
+      run: NV=1 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Test DISK copy time
      run: NV=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
-    - name: Test CPU copy time
-      run: |
-        NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
-        NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
    - name: Test LLAMA-3
      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0 | tee nv_llama3_beam.txt
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
-    #- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-    #  run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
+      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
+    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+      run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
    - name: Run 10 MLPerf Bert training steps (1 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
      run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee nv_train_bert_one_gpu.txt
@@ -7,7 +7,6 @@ env:
  BUILD_CACHE_VERSION: '1'
  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-  PYTHONPATH: ${{ github.workspace }}

 on:
  push:
@@ -30,10 +29,12 @@ jobs:
        key: llvm-speed
        deps: testing_minimal
        llvm: 'true'
+    - name: External Benchmark Schedule
+      run: PYTHONPATH="." python3 test/external/external_benchmark_schedule.py
    - name: Speed Test
-      run: CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py
+      run: LLVM=1 python3 test/speed/external_test_speed_v_torch.py
    - name: Speed Test (BEAM=2)
-      run: BEAM=2 CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py
+      run: BEAM=2 LLVM=1 python3 test/speed/external_test_speed_v_torch.py

  docs:
    name: Docs
@@ -46,7 +47,7 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        deps: docs
-        pydeps: "capstone torch"
+        pydeps: "capstone"
    - name: Build wheel and show size
      run: |
        pip install build
@@ -70,29 +71,98 @@ jobs:
        source venv/bin/activate
        pip install $GITHUB_WORKSPACE
        cp $GITHUB_WORKSPACE/examples/beautiful_mnist.py .
-        BS=2 STEPS=10 python beautiful_mnist.py
+        PYTHONPATH=$GITHUB_WORKSPACE BS=2 STEPS=10 python beautiful_mnist.py
    - name: Test Docs Build
      run: python -m mkdocs build --strict
    - name: Test Docs
      run: |
        python docs/abstractions2.py
        python docs/abstractions3.py
-    - name: Test README
-      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' README.md > README.py &&  python README.py
    - name: Test Quickstart
-      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' docs/quickstart.md > quickstart.py &&  python quickstart.py
+      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' docs/quickstart.md > quickstart.py &&  PYTHONPATH=. python quickstart.py
    - name: Test DEBUG
      run: DEBUG=100 python3 -c "from tinygrad import Tensor; N = 1024; a, b = Tensor.rand(N, N), Tensor.rand(N, N); c = (a.reshape(N, 1, N) * b.T.reshape(1, N, N)).sum(axis=2); print((c.numpy() - (a.numpy() @ b.numpy())).mean())"
    - name: Compile EfficientNet to C and test it
      run: |
-        CPU=1 CPU_LLVM=0 python examples/compile_efficientnet.py > recognize.c
+        CPU=1 PYTHONPATH="." python examples/compile_efficientnet.py > recognize.c
        clang -O2 recognize.c -lm -o recognize
        cat test/models/efficientnet/Chicken.jpg | ./recognize | grep cock

+  autogen:
+    name: Autogen
+    runs-on: ubuntu-24.04
+    timeout-minutes: 15
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        opencl: 'true'
+        amd: 'true'
+        cuda: 'true'
+        webgpu: 'true'
+        llvm: 'true'
+    - name: Install autogen support packages
+      run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev
+    - name: Verify OpenCL autogen
+      run: |
+        cp tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
+        ./autogen_stubs.sh opencl
+        diff /tmp/opencl.py.bak tinygrad/runtime/autogen/opencl.py
+    - name: Verify CUDA autogen
+      run: |
+        cp tinygrad/runtime/autogen/cuda.py /tmp/cuda.py.bak
+        cp tinygrad/runtime/autogen/nv_gpu.py /tmp/nv_gpu.py.bak
+        ./autogen_stubs.sh cuda
+        ./autogen_stubs.sh nv
+        diff /tmp/cuda.py.bak tinygrad/runtime/autogen/cuda.py
+        diff /tmp/nv_gpu.py.bak tinygrad/runtime/autogen/nv_gpu.py
+    - name: Verify AMD autogen
+      run: |
+        cp tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
+        cp tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
+        cp tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
+        cp tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
+        cp tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
+        ./autogen_stubs.sh hsa
+        ./autogen_stubs.sh kfd
+        ./autogen_stubs.sh comgr
+        ./autogen_stubs.sh amd
+        ./autogen_stubs.sh sqtt
+        diff /tmp/hsa.py.bak tinygrad/runtime/autogen/hsa.py
+        diff /tmp/kfd.py.bak tinygrad/runtime/autogen/kfd.py
+        diff /tmp/comgr.py.bak tinygrad/runtime/autogen/comgr.py
+        diff /tmp/amd_gpu.py.bak tinygrad/runtime/autogen/amd_gpu.py
+        diff /tmp/sqtt.py.bak tinygrad/runtime/autogen/sqtt.py
+    - name: Verify Linux autogen
+      run: |
+        cp tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
+        cp tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
+        cp tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
+        ./autogen_stubs.sh libc
+        ./autogen_stubs.sh io_uring
+        ./autogen_stubs.sh ib
+        diff /tmp/libc.py.bak tinygrad/runtime/autogen/libc.py
+        diff /tmp/io_uring.py.bak tinygrad/runtime/autogen/io_uring.py
+        diff /tmp/ib.py.bak tinygrad/runtime/autogen/ib.py
+    - name: Verify WebGPU autogen
+      run: |
+        cp tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak
+        ./autogen_stubs.sh webgpu
+        diff /tmp/webgpu.py.bak tinygrad/runtime/autogen/webgpu.py
+    - name: Verify LLVM autogen
+      run: |
+        cp tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
+        ./autogen_stubs.sh llvm
+        diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
+
  torchbackend:
    name: Torch Backend Tests
    runs-on: ubuntu-latest
    timeout-minutes: 15
+    env:
+      IGNORE_OOB: 0
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -112,24 +182,26 @@ jobs:
        pip3 install --upgrade --force-reinstall ruff==0.11.0
        python3 -m ruff check extra/torch_backend/backend.py
    - name: Test one op
-      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
+      run: PYTHONPATH=. FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
    - name: Test ResNet-18
-      run: DEBUG=2 python3 extra/torch_backend/example.py
+      run: PYTHONPATH=. DEBUG=2 python3 extra/torch_backend/example.py
    - name: My (custom) tests
-      run: python3 extra/torch_backend/test.py
+      run: PYTHONPATH=. python3 extra/torch_backend/test.py
    - name: Test one op in torch tests
-      run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
+      run: PYTHONPATH=. DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
    - name: Test Ops with TINY_BACKEND
-      run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
+      run: PYTHONPATH=. LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
    - name: Test in-place operations on views
-      run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
+      run: PYTHONPATH=. TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
    - name: Test multi-gpu
-      run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
+      run: PYTHONPATH=. LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py

  torchbackendmore:
    name: Torch Backend Tests More
    runs-on: ubuntu-latest
    timeout-minutes: 15
+    env:
+      IGNORE_OOB: 0
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -144,14 +216,86 @@ jobs:
        sudo apt update || true
        sudo apt install -y --no-install-recommends ninja-build
    - name: Test beautiful_mnist in torch with TINY_BACKEND
-      run: CPU=1 CPU_LLVM=1 TARGET_EVAL_ACC_PCT=96.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+      run: SPLIT_REDUCEOP=0 FUSE_ARANGE=1 PYTHONPATH=. LLVM=1 TARGET_EVAL_ACC_PCT=96.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
    - name: Test some torch tests (expect failure)
-      run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
+      run: PYTHONPATH=. python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
+
+  tc:
+    name: Tensor Core tests
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    env:
+      IGNORE_OOB: 0
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: uops-minimal
+        deps: testing_minimal
+    - name: Test IMAGE=2 support
+      run: |
+        IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
+        IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_simple_conv2d
+    - name: Test emulated METAL tensor cores
+      run: |
+        DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_big_gemm
+        PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
+        PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+    - name: Test emulated AMX tensor cores
+      run: PYTHONPATH=. DEBUG=2 AMX=1 EMULATE_AMX=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
+    - name: Test emulated AMD tensor cores
+      run: |
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+    - name: Test emulated AMD MFMA tensor cores
+      run: |
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_MFMA=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_MFMA=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_MFMA=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+    - name: Test emulated AMD RDNA4 tensor cores
+      run: |
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+    - name: Test emulated CUDA tensor cores
+      run: |
+        DEBUG=2 EMULATE_CUDA=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
+        DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
+        DEBUG=2 EMULATE_CUDA_SM75=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
+        PYTHONPATH="." DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
+        PYTHONPATH="." DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+    - name: Test emulated INTEL OpenCL tensor cores
+      run: DEBUG=2 EMULATE_INTEL=1 FORWARD_ONLY=1 PYTHON=1 HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
+    - name: Full test tensor cores
+      run: |
+        PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
+        PYTHONPATH=. DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
+        PYTHONPATH=. DEBUG=2 EMULATE_INTEL=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
+        PYTHONPATH=. DEBUG=2 AMX=1 EMULATE_AMX=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
+    - name: Test device flop counts
+      run: |
+        PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
+        PYTHONPATH=. DEBUG=2 EMULATE_CUDA=1 PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
+        PYTHONPATH=. DEBUG=2 EMULATE_INTEL=1 PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
+        PYTHONPATH=. DEBUG=2 AMX=1 EMULATE_AMX=1 PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStats.test_simple_matmul

  bepython:
    name: Python Backend
    runs-on: ubuntu-latest
-    timeout-minutes: 15
+    timeout-minutes: 10
+    env:
+      IGNORE_OOB: 0
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -161,60 +305,15 @@ jobs:
        key: be-minimal
        deps: testing_minimal
    - name: Test dtype with Python emulator
-      run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
+      run: DEBUG=1 PYTHONPATH=. PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
    - name: Test ops with Python emulator
-      run: DEBUG=2 SKIP_SLOW_TEST=1 PYTHON=1 python3 -m pytest -n=auto test/test_ops.py --durations=20
+      run: DEBUG=2 PYTHON=1 python3 -m pytest -n=auto test/test_ops.py -k "not (test_split or test_simple_cumsum or test_cumsum or test_einsum or test_dot or test_dot_1d or test_big_gemm or test_broadcastdot or test_multidot or test_var_axis or test_std_axis or test_broadcast_full or test_broadcast_partial or test_simple_conv3d or test_dilated_conv_transpose2d or test_simple_conv_transpose3d or test_large_input_conv2d or test_max_pool2d or test_max_pool2d_simple or test_max_pool2d_bigger_stride or test_avg_pool2d or test_cat or test_scaled_product_attention or test_scaled_product_attention_causal or test_slice_fancy_indexing_dim_inject_none or test_slice_fancy_indexing_list_indices or test_slice_fancy_indexing_no_dim_collapse or test_slice_fancy_indexing_tuple_indices or test_slice_fancy_indexing_list_with_tensors or test_slice_fancy_indexing_dim_collapse_int or test_interpolate_bilinear or test_interpolate_bilinear_corners_aligned or test_scaled_dot_product_attention or test_cummax or test_simple_cummax or test_logcumsumexp or test_sort or test_cumprod)" --durations=20
    - name: Test uops with Python emulator
      run: PYTHON=1 python3 -m pytest test/test_uops.py --durations=20
    - name: Test symbolic with Python emulator
-      run: PYTHON=1 python3 test/test_symbolic_ops.py
+      run: PYTHONPATH=. PYTHON=1 python3 test/test_symbolic_ops.py
    - name: test_renderer_failures with Python emulator
-      run: PYTHON=1 python3 -m pytest -rA test/test_renderer_failures.py::TestRendererFailures
-    - name: Test IMAGE=2 support
-      run: |
-        IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
-        IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_simple_conv2d
-    - name: Test emulated METAL tensor cores
-      run: |
-        DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_big_gemm
-        DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
-    - name: Test emulated AMX tensor cores
-      run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
-    - name: Test emulated AMD tensor cores
-      run: |
-        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
-    - name: Test emulated AMD MFMA tensor cores
-      run: |
-        DEBUG=2 EMULATE=AMD_MFMA FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 EMULATE=AMD_MFMA FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
-    - name: Test emulated AMD RDNA4 tensor cores
-      run: |
-        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
-    - name: Test emulated CUDA tensor cores
-      run: |
-        DEBUG=2 EMULATE=CUDA FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
-        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
-        DEBUG=2 EMULATE=CUDA_SM75 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
-        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
-    - name: Test emulated INTEL OpenCL tensor cores
-      run: DEBUG=2 EMULATE=INTEL FORWARD_ONLY=1 PYTHON=1 HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
-    - name: Test emulated AMX tensor cores
-      run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
-    - name: Test device flop counts
-      run: |
-        DEBUG=2 EMULATE=METAL PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 EMULATE=AMD PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 EMULATE=CUDA PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 EMULATE=INTEL PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 AMX=1 EMULATE=AMX PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStats.test_simple_matmul
+      run: PYTHONPATH=. PYTHON=1 python3 -m pytest -rA test/test_renderer_failures.py::TestRendererFailures

  linter:
    name: Linters
@@ -244,8 +343,6 @@ jobs:
      run: |
        python -m mypy --strict-equality --lineprecision-report .
        cat lineprecision.txt
-    - name: Run TYPED=1
-      run: TYPED=1 python -c "import tinygrad"

  unittest:
    name: Unit Tests
@@ -259,39 +356,32 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: unittest-12
-        pydeps: "pillow numpy ftfy regex"
+        pydeps: "pillow"
        deps: testing_unit
-    - name: Check Device.DEFAULT
-      run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
+    - name: Test README
+      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' README.md > README.py &&  PYTHONPATH=. python README.py
    - name: Run unit tests
-      run: CPU=1 python -m pytest -n=auto test/unit/ --durations=20
-    - name: Check SPEC=1
-      run: SPEC=1 python3 test/test_tiny.py
+      run: PYTHONPATH="." python -m pytest -n=auto test/unit/ --durations=20
    - name: Run targetted tests on NULL backend
-      run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
-    # TODO: too slow
-    # - name: Run SDXL on NULL backend
-    #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
-    - name: Run Clip tests for SD MLPerf on NULL backend
-      run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
+      run: PYTHONPATH="." NULL=1 python3 test/test_multitensor.py TestMultiTensor.test_data_parallel_resnet_train_step
+    - name: Run SDXL on NULL backend
+      run: MAX_BUFFER_SIZE=0 PYTHONPATH="." NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
    # TODO: support fake weights
    #- name: Run LLaMA 7B on 4 fake devices
    #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
    - name: Run GC tests
-      run: python test/external/external_uop_gc.py
-    - name: External Benchmark Schedule
-      run: python3 test/external/external_benchmark_schedule.py
+      run: PYTHONPATH="." python test/external/external_uop_gc.py
    - name: Run process replay tests
      uses: ./.github/actions/process-replay
    - name: Regen dataset on test_tiny
      run: |
        test/external/process_replay/reset.py
        CAPTURE_PROCESS_REPLAY=1 python test/test_tiny.py TestTiny.test_plus
-        python extra/optimization/extract_dataset.py
+        PYTHONPATH=. python extra/optimization/extract_dataset.py
        gzip -c /tmp/sops > extra/datasets/sops.gz
-        #DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py
-    - name: Repo line count < 18000 lines
-      run: MAX_LINE_COUNT=18000 python sz.py
+        DEBUG=1 MIN_ASTS=1 PYTHONPATH=. python extra/optimization/get_action_space.py
+    - name: Repo line count < 17000 lines
+      run: MAX_LINE_COUNT=17000 python sz.py

  fuzzing:
    name: Fuzzing
@@ -311,15 +401,17 @@ jobs:
      run: python test/external/fuzz_fast_idiv.py
    - name: Fuzz Test shapetracker
      run: |
-        python test/external/fuzz_shapetracker.py
-        python test/external/fuzz_shapetracker_math.py
+        PYTHONPATH="." python test/external/fuzz_shapetracker.py
+        PYTHONPATH="." python test/external/fuzz_shapetracker_math.py
    - name: Fuzz Test shape ops
      run: python test/external/fuzz_shape_ops.py

-  testopenclimage:
-    name: CL IMAGE Tests
+  testgpuimage:
+    name: 'GPU IMAGE Tests'
    runs-on: ubuntu-22.04
-    timeout-minutes: 15
+    timeout-minutes: 10
+    env:
+      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -329,17 +421,25 @@ jobs:
          key: gpu-image
          deps: testing_minimal
          opencl: 'true'
-      - name: Test CL IMAGE=2 ops + training
+      - name: Run Kernel Count Test
+        run: PYTHONPATH="." GPU=1 python -m pytest -n=auto test/external/external_test_opt.py
+      - name: Test WINO=1
+        run: GPU=1 DEBUG=2 WINO=1 python3 test/test_ops.py TestOps.test_simple_conv2d
+      - name: Test GPU IMAGE=2 ops + training
        run: |
-          CL=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20
-          CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
+          PYTHONPATH="." GPU=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20
+          PYTHONPATH="." GPU=1 IMAGE=2 python3 test/models/test_end2end.py TestEnd2End.test_linear_mnist
+      - name: Run fused optimizer tests
+        run: PYTHONPATH="." GPU=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

-  testgpumisc:
-    name: CL Misc tests
+  testgendataset:
+    name: 'GPU Generate Kernel Dataset'
    runs-on: ubuntu-22.04
    timeout-minutes: 10
+    env:
+      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -350,11 +450,7 @@ jobs:
          deps: testing_minimal
          opencl: 'true'
      - name: Generate Dataset
-        run: CL=1 extra/optimization/generate_dataset.sh
-      - name: Run Kernel Count Test
-        run: CL=1 python -m pytest -n=auto test/external/external_test_opt.py
-      - name: Run fused optimizer tests
-        run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py
+        run: PYTHONPATH="." extra/optimization/generate_dataset.sh
      - name: Upload artifact
        uses: actions/upload-artifact@v4
        with:
@@ -362,9 +458,11 @@ jobs:
          path: /tmp/sops.gz

  testopenpilot:
-    name: openpilot Compile Tests
+    name: 'openpilot Compile Tests'
    runs-on: ubuntu-22.04
    timeout-minutes: 15
+    env:
+      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -377,26 +475,26 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          ALLOWED_KERNEL_COUNT=190 ALLOWED_READ_IMAGE=2041 ALLOWED_GATED_READ_IMAGE=41 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
+          PYTHONPATH="." ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2134 ALLOWED_GATED_READ_IMAGE=13 FLOAT16=0 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot alt model correctness (float32)
-        run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
+        run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot fastvits model correctness (float32)
-        run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+        run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
      # - name: Test openpilot simple_plan vision model correctness (float32)
-      #   run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/35ff4f4577002f2685e50c8346addae33fe8da27a41dd4d6a0f14d1f4b1af81b
+      #   run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/35ff4f4577002f2685e50c8346addae33fe8da27a41dd4d6a0f14d1f4b1af81b
      - name: Test openpilot LLVM compile
-        run: CPU=1 CPU_LLVM=1 LLVMOPT=1 JIT=2 BEAM=0 IMAGE=0 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+        run: PYTHONPATH="." LLVM=1 LLVMOPT=1 JIT=2 BEAM=0 IMAGE=0 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot compile4
-        run: NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 DEBUG=2 python3 examples/openpilot/compile4.py
+        run: PYTHONPATH="." NOLOCALS=1 GPU=1 IMAGE=2 FLOAT16=1 DEBUG=2 python3 examples/openpilot/compile4.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

-# ****** ONNX Tests ******
-
  testonnxcpu:
-    name: ONNX (CPU) Tests
+    name: 'ONNX (CPU) Tests'
    runs-on: ubuntu-22.04
    timeout-minutes: 20
+    env:
+      IGNORE_OOB: 0

    steps:
      - name: Checkout Code
@@ -409,22 +507,25 @@ jobs:
          python-version: '3.11'
          llvm: 'true'
      - name: Test ONNX (CPU)
-        run: CPU=1 CPU_LLVM=0 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+        run: CPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
      - name: Test ONNX (LLVM)
-        run: CPU=1 CPU_LLVM=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+        run: LLVM=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
      - name: Test ONNX Runner (CPU)
-        run: CPU=1 CPU_LLVM=0 python3 test/external/external_test_onnx_runner.py
+        run: CPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_runner.py
      - name: Test Additional ONNX Ops (CPU)
-        run: CPU=1 CPU_LLVM=0 python3 test/external/external_test_onnx_ops.py
+        run: CPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_ops.py
      - name: Test Quantize ONNX
-        run: CPU=1 CPU_LLVM=0 python3 test/test_quantize_onnx.py
+        run: CPU=1 PYTHONPATH=. python3 test/test_quantize_onnx.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

  testopencl:
-    name: ONNX (CL)+Optimization Tests
+    name: 'ONNX (GPU)+Optimization Tests'
    runs-on: ubuntu-22.04
    timeout-minutes: 20
+    env:
+      IGNORE_OOB: 0
+
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -436,22 +537,18 @@ jobs:
          pydeps: "tensorflow==2.15.1 tensorflow_addons"
          python-version: '3.11'
          opencl: 'true'
-      - name: Test ONNX (CL)
-        run: CL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
-      #- name: Test Optimization Helpers
-      #  run: DEBUG=1 python3 extra/optimization/test_helpers.py
+      - name: Test ONNX (GPU)
+        run: GPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+      - name: Test Optimization Helpers
+        run: PYTHONPATH="." DEBUG=1 python3 extra/optimization/test_helpers.py
      #- name: Test Action Space
-      #  run: DEBUG=1 CL=1 python3 extra/optimization/get_action_space.py
+      #  run: PYTHONPATH="." DEBUG=1 GPU=1 python3 extra/optimization/get_action_space.py
      - name: Test Beam Search
-        run: CL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
+        run: PYTHONPATH="." GPU=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
      - name: Test MLPerf stuff
-        run: CL=1 python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
-      - name: NULL=1 beautiful_mnist_multigpu
-        run: NULL=1 python examples/beautiful_mnist_multigpu.py
-      - name: Test Bert training
-        run: NULL=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+        run: GPU=1 python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
      - name: Test llama 3 training
-        run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
+        run: MAX_BUFFER_SIZE=0 PYTHONPATH="." DEV=NULL SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -469,12 +566,12 @@ jobs:
      - name: Test 1B LLM
        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm | grep -i rooster

-# ****** Models Tests ******
-
  testmodels:
    name: Models (llvm+cpu+gpu)
    runs-on: ubuntu-22.04
    timeout-minutes: 15
+    env:
+      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -486,38 +583,20 @@ jobs:
          opencl: 'true'
          llvm: 'true'
      - name: Test models (llvm)
-        run: CPU=1 CPU_LLVM=1 python -m pytest -n=auto test/models --durations=20
-      - name: Test models (opencl)
-        run: CL=1 python -m pytest -n=auto test/models --durations=20
+        run: LLVM=1 python -m pytest -n=auto test/models --durations=20
+      - name: Test models (gpu)
+        run: GPU=1 python -m pytest -n=auto test/models --durations=20
      - name: Test models (cpu)
-        run: CPU=1 CPU_LLVM=0 python -m pytest -n=auto test/models --durations=20
+        run: CPU=1 python -m pytest -n=auto test/models --durations=20
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

-  testmetalmodels:
-    name: Models (metal)
-    runs-on: macos-14
-    timeout-minutes: 20
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: metal
-        deps: testing
-        python-version: '3.11'
-    - name: Test models (Metal)
-      run: METAL=1 python -m pytest -n=auto test/models --durations=20
-    - name: Test LLaMA compile speed
-      run: METAL=1 python test/external/external_test_speed_llama.py
-
-# ****** Feature Tests ******
-
  testdevectorize:
    name: Linux (devectorize)
    runs-on: ubuntu-24.04
    timeout-minutes: 15
+    env:
+      IGNORE_OOB: 0
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -529,16 +608,18 @@ jobs:
        pydeps: "pillow"
        llvm: "true"
    - name: Test LLVM=1 DEVECTORIZE=0
-      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
+      run: LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
    - name: Test LLVM=1 DEVECTORIZE=0 for model
-      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
+      run: PYTHONPATH="." LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
    - name: Test CPU=1 DEVECTORIZE=0
-      run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
+      run: CPU=1 DEVECTORIZE=0 FUSE_ARANGE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"

  testdsp:
    name: Linux (DSP)
    runs-on: ubuntu-24.04
    timeout-minutes: 15
+    env:
+      IGNORE_OOB: 0
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -565,9 +646,9 @@ jobs:
    - name: Run test_tiny on DSP
      run: DEBUG=2 DSP=1 python test/test_tiny.py
    - name: Test transcendentals
-      run: CC=clang-20 DEBUG=2 DSP=1 python test/test_transcendental.py TestTranscendentalVectorized
+      run: CC=clang-20 PYTHONPATH="." DEBUG=2 DSP=1 python test/test_transcendental.py TestTranscendentalVectorized
    - name: Test quantize onnx
-      run: DEBUG=2 DSP=1 python3 test/test_quantize_onnx.py
+      run: PYTHONPATH="." DEBUG=2 DSP=1 python3 test/test_quantize_onnx.py

  testwebgpu:
    name: Linux (WebGPU)
@@ -605,10 +686,12 @@ jobs:
    runs-on: ubuntu-22.04
    timeout-minutes: 20
    env:
+      IGNORE_OOB: 0
      AMD: 1
      MOCKGPU: 1
      FORWARD_ONLY: 1
      AMD_LLVM: ${{ matrix.backend == 'amdllvm' && '1' || matrix.backend != 'amdllvm' && '0' }}
+      PYTHONPATH: ${{ github.workspace }}
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -634,7 +717,7 @@ jobs:
        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
      - name: Run TestOps.test_add with SQTT
        run: |
-          VIZ=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
+          PROFILE=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
          extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
      - name: Run process replay tests
        uses: ./.github/actions/process-replay
@@ -648,9 +731,7 @@ jobs:
    name: Linux (${{ matrix.backend }})
    runs-on: ubuntu-22.04
    timeout-minutes: 20
-    env:
-      MOCKGPU: 1
-      FORWARD_ONLY: 1
+
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -662,26 +743,29 @@ jobs:
          cuda: 'true'
          ocelot: 'true'
      - name: Set env
-        run: printf "${{ matrix.backend == 'PTX' && 'CUDA=1\nCUDA_PTX=1' || matrix.backend == 'nv' && 'NV=1\nSKIP_SLOW_TEST=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'PTX' && 'FORWARD_ONLY=1\nJIT=1\nOPT=2\nCUDA=1\nPTX=1\nMOCKGPU=1' || matrix.backend == 'nv' && 'NV=1\nMOCKGPU=1\nFORWARD_ONLY=1' }}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
-          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CUDA','NV'], Device.DEFAULT"
-          DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+          PYTHONPATH=${{ github.workspace }} python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CUDA','NV'], Device.DEFAULT"
+          DEBUG=5 PYTHONPATH=${{ github.workspace }} FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
      - name: Run pytest (cuda)
        # skip multitensor because it's slow
        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore test/test_gc.py --ignore test/test_multitensor.py --durations=20
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

-  testcpuopencl:
+  tests:
    strategy:
      fail-fast: false
      matrix:
-        backend: [llvm, cpu, opencl]
+        backend: [llvm, cpu, gpu]

    name: Linux (${{ matrix.backend }})
    runs-on: ubuntu-22.04
    timeout-minutes: 20
+    env:
+      IGNORE_OOB: 0
+
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -690,124 +774,65 @@ jobs:
        with:
          key: ${{ matrix.backend }}-minimal
          deps: testing_minimal
-          opencl: ${{ matrix.backend == 'opencl' && 'true' }}
+          opencl: ${{ matrix.backend == 'gpu' && 'true' }}
          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'CL=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'gpu' && 'GPU=1' }}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
-          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT"
-          DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
-      - name: Run pytest (${{ matrix.backend }})
+          PYTHONPATH=${{ github.workspace }} python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['LLVM','CPU','GPU'], Device.DEFAULT"
+          DEBUG=5 PYTHONPATH=${{ github.workspace }} FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+      - name: Run pytest (not cuda)
        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
      - name: Run TRANSCENDENTAL math
        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

-  amdremote:
-    name: Linux (remote)
-    runs-on: ubuntu-22.04
-    timeout-minutes: 20
-    env:
-      REMOTE: 1
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v4
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: linux-remote
-          deps: testing_minimal
-          amd: 'true'
-          llvm: 'true'
-          opencl: 'true'
-      - name: Start remote server
-        run: |
-          start_server() {
-            systemd-run --user \
-              --unit="$1" \
-              --setenv=REMOTEDEV="$2" \
-              --setenv=MOCKGPU=1 \
-              --setenv=PYTHONPATH=. \
-              --setenv=PORT="$3" \
-              --working-directory="$(pwd)" \
-              python tinygrad/runtime/ops_remote.py
-          }
-
-          start_server "remote-server-amd-1" "AMD" 6667
-          start_server "remote-server-amd-2" "AMD" 6668
-          start_server "remote-server-gpu" "CL" 7667
-          start_server "remote-server-cpu" "CPU" 8667
-      - name: Check Device.DEFAULT and print some source
-        env:
-          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
-        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
-          python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'AMD', Device.default.properties.real_device"
-          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
-      - name: Run REMOTE=1 Test (AMD)
-        env:
-          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
-        run: |
-          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_remote.py test/test_tensor_variable.py --durations 20
-      - name: Run REMOTE=1 Test (CL)
-        env:
-          HOST: 127.0.0.1:7667*6
-        run: |
-          python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py --durations 20
-          IMAGE=2 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
-      - name: Run REMOTE=1 Test (CPU)
-        env:
-          HOST: 127.0.0.1:8667*6
-        run: |
-          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py --durations 20
-      - name: Show remote server logs
-        if: always()
-        run: |
-          journalctl --user -u remote-server-amd-1 --no-pager
-          journalctl --user -u remote-server-amd-2 --no-pager
-          journalctl --user -u remote-server-gpu --no-pager
-          journalctl --user -u remote-server-cpu --no-pager
-
 # ****** OSX Tests ******

-  testmetal:
+  testmetal2:
    name: MacOS (unit)
    runs-on: macos-14
    timeout-minutes: 20
+    env:
+      IGNORE_OOB: 0
+
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
-        key: metal
+        key: metal2
        deps: testing
        python-version: '3.11'
        amd: 'true'
        cuda: 'true'
        ocelot: 'true'
        llvm: 'true'
-    - name: Run unit tests
-      run: METAL=1 python -m pytest -n=auto test/unit/ --durations=20
+    - name: Run real world test
+      run: METAL=1 python -m pytest -n=auto test/models/test_real_world.py --durations=20
+    - name: Test models (Metal)
+      run: METAL=1 python -m pytest -n=auto test/models -v --durations=20
    - name: Run ONNX
      run: METAL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
    - name: Test tensor core ops (fake)
-      run: METAL=1 DEBUG=3 TC=2 python test/test_ops.py TestOps.test_gemm
+      run: TC=2 METAL=1 DEBUG=3 python test/test_ops.py TestOps.test_gemm
    - name: Test tensor core ops (real)
      run: METAL=1 DEBUG=3 python test/test_ops.py TestOps.test_big_gemm
+    - name: Test LLaMA compile speed
+      run: PYTHONPATH="." METAL=1 python test/external/external_test_speed_llama.py
    - name: Test Beam Search
-      run: METAL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
+      run: PYTHONPATH="." METAL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
    #- name: Fuzz Test linearizer
-    #  run: METAL=1 DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
+    #  run: PYTHONPATH="." METAL=1 DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
    - name: Run TRANSCENDENTAL math
-      run: METAL=1 TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
+      run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
    - name: Run pytest (amd)
      env:
        MOCKGPU: 1
        AMD: 1
-        AMD_LLVM: 0
        FORWARD_ONLY: 1
      run: |
        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
@@ -815,14 +840,13 @@ jobs:
      env:
        MOCKGPU: 1
        AMD: 1
-        AMD_LLVM: 1
        FORWARD_ONLY: 1
      run: |
        python -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py test/device/test_amd_llvm.py --durations=20
    - name: Run pytest (ptx)
      env:
        MOCKGPU: 1
-        NV_PTX: 1
+        PTX: 1
        NV: 1
        FORWARD_ONLY: 1
      run: |
@@ -865,7 +889,7 @@ jobs:
    #    cp $GITHUB_WORKSPACE/test/web/test_viz.js .
    #    node test_viz.js
    - name: Test ONNX Runner (WEBGPU)
-      run: WEBGPU=1 python3 test/external/external_test_onnx_runner.py
+      run: WEBGPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_runner.py

  osxremote:
   name: MacOS (remote metal)
@@ -891,6 +915,72 @@ jobs:
       run: |
         python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py

+  amdremote:
+    name: Linux (remote)
+    runs-on: ubuntu-22.04
+    timeout-minutes: 20
+    env:
+      REMOTE: 1
+      PYTHONPATH: ${{ github.workspace }}
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: linux-remote
+          deps: testing_minimal
+          amd: 'true'
+          llvm: 'true'
+          opencl: 'true'
+      - name: Start remote server
+        run: |
+          start_server() {
+            systemd-run --user \
+              --unit="$1" \
+              --setenv=REMOTEDEV="$2" \
+              --setenv=MOCKGPU=1 \
+              --setenv=PYTHONPATH=. \
+              --setenv=PORT="$3" \
+              --working-directory="$(pwd)" \
+              python tinygrad/runtime/ops_remote.py
+          }
+
+          start_server "remote-server-amd-1" "AMD" 6667
+          start_server "remote-server-amd-2" "AMD" 6668
+          start_server "remote-server-gpu" "GPU" 7667
+          start_server "remote-server-cpu" "CPU" 8667
+      - name: Check Device.DEFAULT and print some source
+        env:
+          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
+        run: |
+          python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
+          python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'AMD', Device.default.properties.real_device"
+          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
+      - name: Run REMOTE=1 Test (AMD)
+        env:
+          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
+        run: |
+          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_remote.py test/test_tensor_variable.py --durations 20
+      - name: Run REMOTE=1 Test (GPU)
+        env:
+          HOST: 127.0.0.1:7667*6
+        run: |
+          python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py --durations 20
+          IMAGE=2 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
+      - name: Run REMOTE=1 Test (CPU)
+        env:
+          HOST: 127.0.0.1:8667*6
+        run: |
+          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py --durations 20
+      - name: Show remote server logs
+        if: always()
+        run: |
+          journalctl --user -u remote-server-amd-1 --no-pager
+          journalctl --user -u remote-server-amd-2 --no-pager
+          journalctl --user -u remote-server-gpu --no-pager
+          journalctl --user -u remote-server-cpu --no-pager
+
  osxtests:
    strategy:
      fail-fast: false
@@ -899,6 +989,8 @@ jobs:
    name: MacOS (${{ matrix.backend }})
    runs-on: macos-15
    timeout-minutes: 20
+    env:
+      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -910,10 +1002,10 @@ jobs:
          pydeps: "capstone"
          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'metal' && 'METAL=1'}}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'metal' && 'METAL=1'}}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
+          python -c "from tinygrad import Device; assert Device.DEFAULT == '${{ matrix.backend }}'.upper(), Device.DEFAULT"
          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
      - name: Run pytest (${{ matrix.backend }})
        run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
@@ -934,6 +1026,8 @@ jobs:
    name: Windows (${{ matrix.backend }})
    runs-on: windows-latest
    timeout-minutes: 15
+    env:
+      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -945,13 +1039,12 @@ jobs:
          pydeps: ${{ matrix.backend == 'webgpu' && 'dawn-python' || '' }}
      - name: Set env
        shell: bash
-        run:  printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'webgpu' && 'WEBGPU=1'}}" >> $GITHUB_ENV
+        run:  printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'webgpu' && 'WEBGPU=1'}}" >> $GITHUB_ENV
      - name: Run unit tests
        if: matrix.backend=='llvm'
-        # test_newton_schulz hits RecursionError
-        run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_elf.py --ignore=test/unit/test_tar.py --ignore=test/unit/test_linalg.py --durations=20
+        run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_elf.py --ignore=test/unit/test_tar.py
      - name: Run pytest (${{ matrix.backend }})
        shell: bash
        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
+          python -c "from tinygrad import Device; assert Device.DEFAULT == '${{ matrix.backend }}'.upper(), Device.DEFAULT"
          python -m pytest -n=auto test/test_tiny.py test/test_ops.py --durations=20
@@ -20,6 +20,12 @@ repos:
        language: system
        always_run: true
        pass_filenames: false
+      - id: devicetests
+        name: select GPU tests
+        entry: env GPU=1 PYTHONPATH="." python3 -m pytest test/test_uops.py test/test_search.py
+        language: system
+        always_run: true
+        pass_filenames: false
      - id: tests
        name: subset of tests
        entry: env PYTHONPATH="." python3 -m pytest -n=4 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
@@ -54,12 +54,11 @@ confidence=
 # --enable=similarities". If you want to run only the classes checker, but have
 # no Warning level messages displayed, use"--disable=all --enable=classes
 # --disable=W"
-disable=C,R,W0613,W0511,W0212,W0201,W0106,W0603,W0621,W0703,W1201,W1203,E1136,W1514,E1101,W0221,W0105,E0401,abstract-method,W0707
+disable=C,R,W0613,W0511,W0212,W0201,W0106,W0603,W0621,W0703,W1201,W1203,E1136,W1514,E1101,W0221,W0105,E0401,abstract-method
 # E1101 for function binding
 # W0221 for Function class
 # W0105 for comment strings
 # E0401 for missing imports
-# W0707 for not reraising

 # Enable the message, report, category or checker with the given id(s). You can
 # either give multiple identifier separated by comma (,) or put this option
@@ -79,8 +79,9 @@ See [examples/beautiful_mnist.py](examples/beautiful_mnist.py) for the full vers

 tinygrad already supports numerous accelerators, including:

- [x] [OpenCL](tinygrad/runtime/ops_cl.py)
- [x] [CPU](tinygrad/runtime/ops_cpu.py)
+- [x] [GPU (OpenCL)](tinygrad/runtime/ops_gpu.py)
+- [x] [CPU (C Code)](tinygrad/runtime/ops_cpu.py)
+- [x] [LLVM](tinygrad/runtime/ops_llvm.py)
 - [x] [METAL](tinygrad/runtime/ops_metal.py)
 - [x] [CUDA](tinygrad/runtime/ops_cuda.py)
 - [x] [AMD](tinygrad/runtime/ops_amd.py)
@@ -414,29 +414,10 @@ generate_sqtt() {
  clang2py -k cdefstum \
    extra/sqtt/sqtt.h \
    -o $BASE/sqtt.py
+
  fixup $BASE/sqtt.py
  sed -i "s\import ctypes\import ctypes, os\g" $BASE/sqtt.py
  python3 -c "import tinygrad.runtime.autogen.sqtt"
-
-  ROCPROF_COMMIT_HASH=dd0485100971522cc4cd8ae136bdda431061a04d
-  ROCPROF_SRC=/tmp/rocprof-trace-decoder-$ROCPROF_COMMIT_HASH
-  if [ ! -d "$ROCPROF_SRC" ]; then
-    git clone https://github.com/ROCm/rocprof-trace-decoder $ROCPROF_SRC
-    pushd .
-    cd $ROCPROF_SRC
-    git reset --hard $ROCPROF_COMMIT_HASH
-    popd
-  fi
-
-  clang2py -k cdefstum \
-    $ROCPROF_SRC/include/rocprof_trace_decoder.h \
-    $ROCPROF_SRC/include/trace_decoder_instrument.h \
-    $ROCPROF_SRC/include/trace_decoder_types.h \
-    -o extra/sqtt/rocprof/rocprof.py
-  fixup extra/sqtt/rocprof/rocprof.py
-  sed -i '1s/^/# pylint: skip-file\n/' extra/sqtt/rocprof/rocprof.py
-  sed -i "s/import ctypes/import ctypes\nfrom tinygrad.helpers import fetch/g" extra/sqtt/rocprof/rocprof.py
-  sed -i "s|FunctionFactoryStub()|ctypes.CDLL(str(fetch('https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so')))|g" extra/sqtt/rocprof/rocprof.py
 }

 generate_webgpu() {
@@ -42,6 +42,7 @@ import struct
 from tinygrad.dtype import dtypes
 from tinygrad.device import Buffer, Device
 from tinygrad.uop.ops import UOp, Ops
+from tinygrad.shape.shapetracker import ShapeTracker

 # allocate some buffers + load in values
 out = Buffer(DEVICE, 1, dtypes.int32).allocate()
@@ -50,14 +51,13 @@ b = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struc
 # NOTE: a._buf is the same as the return from cpu.allocator.alloc

 # describe the computation
-idx = UOp.const(dtypes.index, 0)
 buf_1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 1)
 buf_2 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 2)
-ld_1 = UOp(Ops.LOAD, dtypes.int32, (buf_1.index(idx),))
-ld_2 = UOp(Ops.LOAD, dtypes.int32, (buf_2.index(idx),))
+ld_1 = UOp(Ops.LOAD, dtypes.int32, (buf_1.view(ShapeTracker.from_shape((1,))),))
+ld_2 = UOp(Ops.LOAD, dtypes.int32, (buf_2.view(ShapeTracker.from_shape((1,))),))
 alu = ld_1 + ld_2
 output_buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 0)
-st_0 = UOp(Ops.STORE, dtypes.void, (output_buf.index(idx), alu))
+st_0 = UOp(Ops.STORE, dtypes.void, (output_buf.view(ShapeTracker.from_shape((1,))), alu))
 s = UOp(Ops.SINK, dtypes.void, (st_0,))

 # convert the computation to a "linearized" format (print the format)
@@ -80,7 +80,7 @@ print("******** third, the UOp ***********")

 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.schedule import create_schedule_with_vars
-from tinygrad.schedule.rangeify import get_rangeify_map
+from tinygrad.schedule.kernelize import get_kernelize_map

 # allocate some values + load in values
 a = UOp.new_buffer(DEVICE, 1, dtypes.int32)
@@ -93,10 +93,10 @@ out = a + b
 s = UOp(Ops.SINK, dtypes.void, (out,))

 # group the computation into kernels
-becomes_map = get_rangeify_map(s)
+becomes_map = get_kernelize_map(s)

 # the compute maps to an assign
-assign = becomes_map[a+b].base
+assign = becomes_map[a+b]

 # the first source is the output buffer (data)
 assert assign.src[0].op is Ops.BUFFER
@@ -10,7 +10,7 @@ Directories are listed in order of how they are processed.

 Group UOps into kernels.

-::: tinygrad.schedule.rangeify.get_rangeify_map
+::: tinygrad.schedule.kernelize.get_kernelize_map
    options:
        members: false
        show_labels: false
@@ -22,6 +22,12 @@ Group UOps into kernels.

 Transforms the ast into an optimized ast. This is where BEAM search and heuristics live.

+::: tinygrad.codegen.opt.get_optimized_ast
+    options:
+        members: false
+        show_labels: false
+        show_source: false
+
 ---

 ## tinygrad/codegen
@@ -3,7 +3,7 @@
 This is a list of environment variable that control the runtime behavior of tinygrad and its examples.
 Most of these are self-explanatory, and are usually used to set an option at runtime.

-Example: `CL=1 DEBUG=4 python3 -m pytest`
+Example: `GPU=1 DEBUG=4 python3 -m pytest`

 However you can also decorate a function to set a value only inside that function.

@@ -31,16 +31,19 @@ These control the behavior of core tinygrad even when used as a library.
 Variable | Possible Value(s) | Description
 ---|---|---
 DEBUG               | [1-7]      | enable debugging output (operations, timings, speed, generated code and more)
-CL                  | [1]        | enable OpenCL backend
+GPU                 | [1]        | enable the GPU (OpenCL) backend
 CUDA                | [1]        | enable CUDA backend
 AMD                 | [1]        | enable AMD backend
 NV                  | [1]        | enable NV backend
 METAL               | [1]        | enable Metal backend (for Mac M1 and after)
-CPU                 | [1]        | enable CPU backend
+CPU                 | [1]        | enable CPU (Clang) backend
+LLVM                | [1]        | enable LLVM backend
 BEAM                | [#]        | number of beams in kernel beam search
 DEFAULT_FLOAT       | [HALF, ...]| specify the default float dtype (FLOAT32, HALF, BFLOAT16, FLOAT64, ...), default to FLOAT32
 IMAGE               | [1-2]      | enable 2d specific optimizations
 FLOAT16             | [1]        | use float16 for images instead of float32
+PTX                 | [1]        | enable the specialized [PTX](https://docs.nvidia.com/cuda/parallel-thread-execution/) assembler for Nvidia GPUs. If not set, defaults to generic CUDA codegen backend.
+PROFILE             | [1]        | enable profiling. This feature is supported in NV, AMD, QCOM and METAL backends.
 VISIBLE_DEVICES     | [list[int]]| restricts the NV/AMD devices that are available. The format is a comma-separated list of identifiers (indexing starts with 0).
 JIT                 | [0-2]      | 0=disabled, 1=[jit enabled](quickstart.md#jit) (default), 2=jit enabled, but graphs are disabled
 VIZ                 | [1]        | 0=disabled, 1=[viz enabled](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/viz)
@@ -2,17 +2,17 @@

 tinygrad supports various runtimes, enabling your code to scale across a wide range of devices. The default runtime can be automatically selected based on the available hardware, or you can force a specific runtime to be default using environment variables (e.g., `CPU=1`).

-| Runtime | Description | Compiler Options | Requirements |
-|---------|-------------|------------------|--------------|
-| [NV](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_nv.py) | Provides acceleration for NVIDIA GPUs | nvrtc (default)<br>PTX (`NV_PTX=1`) | Ampere/Ada/Blackwell series GPUs.<br>You can select an interface via `NV_IFACE=(NVK\|PCI)`. See [NV interfaces](#nv-interfaces) for details. |
-| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | LLVM (`AMD_LLVM=1`)<br>HIP/COMGR (`AMD_HIP=1`) | RDNA2 or newer GPUs.<br>You can select an interface via `AMD_IFACE=(KFD\|PCI\|USB)`. See [AMD interfaces](#amd-interfaces) for details. |
-| [QCOM](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_qcom.py) | Provides acceleration for QCOM GPUs | - | 6xx series GPUs |
-| [METAL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_metal.py) | Utilizes Metal for acceleration on Apple devices | - | M1+ Macs; Metal 3.0+ for `bfloat` support |
-| [CUDA](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cuda.py) | Utilizes CUDA for acceleration on NVIDIA GPUs | nvrtc (default)<br> PTX (`CUDA_PTX=1`) | NVIDIA GPU with CUDA support |
-| [CL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cl.py) | Accelerates computations using OpenCL on GPUs | - | OpenCL 2.0 compatible device |
-| [CPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang or llvm compiler | Clang JIT (default)<br>LLVM IR (`CPU_LLVM=1`) | `clang` compiler in system `PATH` |
-| [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | - | Dawn library installed and discoverable. Binaries: [pydawn v0.3.0](https://github.com/wpmed92/pydawn/releases/tag/v0.3.0) |
-
+| Runtime | Description | Requirements |
+|---------|-------------|--------------|
+| [NV](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_nv.py) | Provides acceleration for NVIDIA GPUs | Ampere/Ada series GPUs |
+| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | RDNA2/RDNA3/RDNA4 series GPUs. You can select one of the interfaces for communication by setting `AMD_IFACE=(KFD|PCI)`. See [AMD interfaces](#amd-interfaces) for more details. |
+| [QCOM](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_qcom.py) | Provides acceleration for QCOM GPUs | 6xx series GPUs |
+| [METAL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_metal.py) | Utilizes Metal for acceleration on Apple devices | M1+ Macs; Metal 3.0+ for `bfloat` support |
+| [CUDA](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cuda.py) | Utilizes CUDA for acceleration on NVIDIA GPUs | NVIDIA GPU with CUDA support |
+| [GPU (OpenCL)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_gpu.py) | Accelerates computations using OpenCL on GPUs | OpenCL 2.0 compatible device |
+| [CPU (C Code)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang compiler | `clang` compiler in system `PATH` |
+| [LLVM (LLVM IR)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_llvm.py) | Runs on CPU using the LLVM compiler infrastructure | llvm libraries installed and findable |
+| [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | Dawn library installed and findable. Download binaries [here](https://github.com/wpmed92/pydawn/releases/tag/v0.3.0). |

 ## Interoperability

@@ -70,12 +70,5 @@ AMD backend supports several interfaces for communicating with devices:

 * `KFD`: uses the amdgpu driver
 * `PCI`: uses the [AM driver](developer/am.md)
-* `USB`: USB3 interafce for asm24xx chips.

 You can force an interface by setting `AMD_IFACE` to one of these values. In the case of `AMD_IFACE=PCI`, this may unbind your GPU from the amdgpu driver.
-
-## NV Interfaces
-NV backend supports several interfaces for communicating with devices:
-
-* `NVK`: uses the nvidia driver
-* `PCI`: uses the [NV driver](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/support/nv/nvdev.py)
@@ -78,7 +78,6 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.minimum
 ::: tinygrad.Tensor.where
 ::: tinygrad.Tensor.copysign
-::: tinygrad.Tensor.logaddexp

 ## Casting Ops

@@ -6,7 +6,7 @@ If you don't have a tinybox and you want one, see [tinygrad.org](https://tinygra

 ## Welcome

-Welcome to your tinybox! The tinybox is the universal system purpose-built for all AI infrastructure and workloads, from training to inference. The red box includes six 7900XTX GPUs, the green box includes six 4090 GPUs, and the green v2 box includes four 5090 GPUs. Whether you bought a red one or a green one, we want you to love it.
+Welcome to your tinybox! The tinybox is the universal system purpose-built for all AI infrastructure and workloads, from training to inference. The red box includes six 7900XTX GPUs, and the green box includes six 4090 GPUs. Whether you bought a red one or a green one, we want you to love it.

 We don't have a stupid cloud service, you don't have to create a tiny account to set it up, and we aren't tracking how you use the box. We're just happy you bought one. This petaflop is your petaflop.

@@ -2,6 +2,7 @@ import time
 start_tm = time.perf_counter()
 import math
 from typing import Tuple, cast
+import numpy as np
 from tinygrad import Tensor, nn, GlobalCounters, TinyJit, dtypes, Device
 from tinygrad.helpers import partition, trange, getenv, Context
 from extra.lr_scheduler import OneCycleLR
@@ -10,7 +11,7 @@ GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))]

 # override tinygrad defaults
 dtypes.default_float = dtypes.half
-Context(FUSE_OPTIM=1).__enter__()
+Context(FUSE_ARANGE=1, FUSE_OPTIM=1).__enter__()

 # from https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
 batchsize = getenv("BS", 1024)
@@ -149,12 +150,13 @@ if __name__ == "__main__":
      acc.append((out.argmax(-1) == Y).sum() / eval_batchsize)
    return Tensor.stack(*loss).mean() / (batchsize*loss_batchsize_scaler), Tensor.stack(*acc).mean()

-  Tensor.manual_seed(1337)
-  num_train_samples = X_train.shape[0]
-
+  np.random.seed(1337)
  for epoch in range(math.ceil(hyp['misc']['train_epochs'])):
+    # TODO: move to tinygrad
    gst = time.perf_counter()
-    tidxs = Tensor.randperm(num_train_samples, dtype='int')[:num_steps_per_epoch*batchsize].reshape(num_steps_per_epoch, batchsize)
+    idxs = np.arange(X_train.shape[0])
+    np.random.shuffle(idxs)
+    tidxs = Tensor(idxs, dtype='int')[:num_steps_per_epoch*batchsize].reshape(num_steps_per_epoch, batchsize)  # NOTE: long doesn't fold
    train_loss:float = 0
    for epoch_step in (t:=trange(num_steps_per_epoch)):
      st = time.perf_counter()
@@ -1,6 +1,6 @@
 import sys, time
 from tinygrad import TinyJit, GlobalCounters, fetch, getenv
-from tinygrad.nn.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner
 from extra.onnx_helpers import get_example_inputs, validate

 def load_onnx_model(onnx_file):
@@ -8,7 +8,7 @@ import numpy as np
 import subprocess
 import tensorflow as tf
 import tf2onnx
-from tinygrad.nn.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import to_mv
 from extra.export_model import export_model_clang, compile_net, jit_model
@@ -26,8 +26,8 @@ class Attention:
      start_pos = start_pos.val

    if HALF: x = x.half()
-    xqkv = self.c_attn(x).reshape(None, None, 3, self.n_heads, self.head_dim)
-    xq, xk, xv = [xqkv[:, :, i, :, :] for i in range(3)]
+    xqkv = self.c_attn(x)
+    xq, xk, xv = [xqkv.shrink((None, None, (i*self.dim, (i+1)*self.dim))).reshape(None, None, self.n_heads, self.head_dim) for i in range(3)]
    bsz, seqlen, _, _ = xq.shape

    # create kv cache
@@ -35,11 +35,11 @@ class Attention:
      self.cache_kv = Tensor.zeros(2, bsz, MAX_CONTEXT, self.n_heads, self.head_dim, dtype=x.dtype).contiguous().realize()

    # update the cache
-    self.cache_kv[:, :, start_pos:start_pos+seqlen, :, :].assign(Tensor.stack(xk, xv)).realize()
+    self.cache_kv.shrink((None, None,(start_pos,start_pos+seqlen),None,None)).assign(Tensor.stack(xk, xv)).realize()

    if start_pos > 0:
-      keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
-      values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
+      keys = self.cache_kv[0].shrink((None, (0, start_pos+seqlen), None, None))
+      values = self.cache_kv[1].shrink((None, (0, start_pos+seqlen), None, None))
    else:
      keys = xk
      values = xv
@@ -64,7 +64,7 @@ class TransformerBlock:

  def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]):
    h = x + self.attn(self.ln_1(x), start_pos, mask).float()
-    return (h + self.mlp(self.ln_2(h))).contiguous()
+    return (h + self.mlp(self.ln_2(h)))

 class Transformer:
  def __init__(self, dim, n_heads, n_layers, norm_eps, vocab_size, max_seq_len=1024):
@@ -181,7 +181,6 @@ class GPT2:
    self.tokenizer = tokenizer

  def generate(self, prompt:str, max_length:int, temperature:float, timing:bool=False, batch_size:int=1):
-    step_times = []
    prompt_tokens = self.tokenizer.encode(prompt, allowed_special={"<|endoftext|>"})
    toks = [prompt_tokens[:] for _ in range(batch_size)]
    start_pos = 0
@@ -189,7 +188,7 @@ class GPT2:
      GlobalCounters.reset()
      if timing: print("")
      st = GlobalCounters.time_sum_s
-      with Timing("ran model in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "")+
+      with Timing("ran model in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+
                  f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
                  (f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=timing):
        with WallTimeEvent(BenchEvent.STEP):
@@ -198,13 +197,8 @@ class GPT2:
          else:
            tokens = Tensor([x[start_pos:] for x in toks])
          tok = self.model(tokens, Variable("start_pos", 1 if start_pos else 0, MAX_CONTEXT-1).bind(start_pos), temperature).tolist()
-      step_times.append((GlobalCounters.time_sum_s-st)*1e3)
      start_pos = len(toks[0])
      for i,t in enumerate(tok): toks[i].append(t)
-
-    if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
-      min_time = min(step_times)
-      assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
    return [self.tokenizer.decode(x) for x in toks]

 # **** main code ****
@@ -118,7 +118,7 @@ class SpeedyResNet:
 # hyper-parameters were exactly the same as the original repo
 bias_scaler = 58
 hyp = {
-  'seed' : 201,
+  'seed' : 200,
  'opt': {
    'bias_lr':            1.76 * bias_scaler/512,
    'non_bias_lr':        1.76 / 512,
@@ -145,6 +145,7 @@ hyp = {
  },
 }

+@Context(FUSE_ARANGE=getenv("FUSE_ARANGE", 1))
 def train_cifar():

  def set_seed(seed):
@@ -228,8 +229,7 @@ def train_cifar():
    if getenv("RANDOM_CROP", 1):
      X = random_crop(X, crop_size=32)
    if getenv("RANDOM_FLIP", 1):
-      # NOTE: RANGEIFY=1 needs this contiguous or the X[perms] is very slow
-      X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X).contiguous() # flip LR
+      X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X) # flip LR
    X, Y = X[perms], Y[perms]
    return X, Y, *cutmix(X, Y, perms, mask_size=hyp['net']['cutmix_size'])

@@ -355,7 +355,7 @@ def train_cifar():

  # https://www.anandtech.com/show/16727/nvidia-announces-geforce-rtx-3080-ti-3070-ti-upgraded-cards-coming-in-june
  # 136 TFLOPS is the theoretical max w float16 on 3080 Ti
-  step_times = []
+
  model_ema: Optional[modelEMA] = None
  projected_ema_decay_val = hyp['ema']['decay_base'] ** hyp['ema']['every_n_steps']
  i = 0
@@ -413,17 +413,12 @@ def train_cifar():
          model_ema.update(model, Tensor([projected_ema_decay_val*(i/STEPS)**hyp['ema']['decay_pow']]))

      cl = time.monotonic()
-      step_times.append((cl-st)*1000.0)
      device_str = loss.device if isinstance(loss.device, str) else f"{loss.device[0]} * {len(loss.device)}"
      #  53  221.74 ms run,    2.22 ms python,  219.52 ms CL,  803.39 loss, 0.000807 LR, 4.66 GB used,   3042.49 GFLOPS,    674.65 GOPS
      print(f"{i:3d} {(cl-st)*1000.0:7.2f} ms run, {(et-st)*1000.0:7.2f} ms python, {(cl-et)*1000.0:7.2f} ms {device_str}, {loss_cpu:7.2f} loss, {opt_non_bias.lr.numpy()[0]:.6f} LR, {GlobalCounters.mem_used/1e9:.2f} GB used, {GlobalCounters.global_ops*1e-9/(cl-st):9.2f} GFLOPS, {GlobalCounters.global_ops*1e-9:9.2f} GOPS")
      st = cl
      i += 1

-  if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
-    min_time = min(step_times)
-    assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
-
  # verify eval acc
  if target := getenv("TARGET_EVAL_ACC_PCT", 0.0):
    if eval_acc_pct >= target:
@@ -478,7 +478,7 @@ After you are done speaking, output [EOS]. You are not Chad.
      with Profiling(enabled=args.profile):
        with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/s, {GlobalCounters.global_mem/x:.2f} GB/s, param {param_bytes/x:.2f} GB/s"):
          with WallTimeEvent(BenchEvent.STEP):
-            with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "")+
+            with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+
                        f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
                        (f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=args.timing):
              tok_tensor = llama.model(next_tok, start_pos, args.temperature)
@@ -441,7 +441,7 @@ if __name__ == "__main__":
      with Profiling(enabled=args.profile):
        with Timing("total ", on_exit=lambda x: f", {1e9/x:.2f} tok/s, {GlobalCounters.global_mem/x:.2f} GB/s, param {param_bytes/x:.2f} GB/s"):
          with WallTimeEvent(BenchEvent.STEP):
-            with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "")+
+            with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+
                        f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
                        (f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None):
              tok = model(Tensor([[last_tok]], device=device), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P)
@@ -479,7 +479,7 @@ if __name__ == "__main__":
        st = GlobalCounters.time_sum_s
        with Profiling(enabled=args.profile):
          with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/s, {GlobalCounters.global_mem/x:.2f} GB/s, param {param_bytes/x:.2f} GB/s"):
-            with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "")+
+            with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+
                        f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
                        (f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=args.timing):

@@ -279,15 +279,9 @@ def generate(model, tokenizer, prompt: str, n_tokens_to_gen: int = 10, temp: boo
  # Loading in the prompt tokens
  logits = model.forward(Tensor([tks]))[:, -1, :]
  for _ in tqdm(range(n_tokens_to_gen), desc="Speed Gen"):
+    # TODO: topk
    if sample:
-      scaled_logits = logits / temp
-      if top_k is not None:
-        topk_values, topk_indices = scaled_logits.topk(top_k)
-        filtered_logits = Tensor.full_like(scaled_logits, -float("inf"))
-        filtered_logits = filtered_logits.scatter(dim=-1, index=topk_indices, src=topk_values)
-        tok_Tens = filtered_logits.softmax().multinomial()
-      else:
-        tok_Tens = scaled_logits.softmax().multinomial()
+      tok_Tens = (logits/temp).softmax().multinomial()
    else:
      tok_Tens = logits.argmax(axis=-1).unsqueeze(0)
    tok = tok_Tens.item()
@@ -304,7 +298,6 @@ if __name__ == "__main__":
  parser.add_argument("--size", type=str, default="370m",
                      help=f"Size of model to use [{', '.join([k for k in MODELS.keys()])}]")
  parser.add_argument("--n_tokens", type=int, default=10, help="Number of tokens to generate")
-  parser.add_argument("--top_k", type=int, help="Limit sampling to the top k most likely tokens")
  parser.add_argument("--sample", dest="sample", action="store_true", help="Sample flag")
  parser.add_argument("--temp", type=float, default=1.0, help="Sampling temp has to be <=1.0")
  args = parser.parse_args()
@@ -315,9 +308,8 @@ if __name__ == "__main__":
  num_toks = args.n_tokens
  sample = args.sample
  temp = args.temp
-  top_k = args.top_k
  s = time.time()
-  tinyoutput = generate(model, tokenizer, prompt, n_tokens_to_gen=num_toks, sample=sample, temp=temp, top_k=top_k)
+  tinyoutput = generate(model, tokenizer, prompt, n_tokens_to_gen=num_toks, sample=sample, temp=temp)
  print(tinyoutput)
  print('TIME: ', time.time() - s)
  TORCHOUTPUT = "Why is gravity \nso important?\nBecause it's the only"
@@ -511,33 +511,6 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
      # happens with BENCHMARK set
      pass

-# stable diffusion callbacks to match mlperf ref; declared here because they're pickled
-def filter_dataset(sample:dict): return {k:v for k,v in sample.items() if k in {'npy', 'txt'}}
-def collate(batch:list[dict]):
-  ret = {"npy": [], "txt": [], "__key__": []}
-  for sample in batch:
-    for k,v in sample.items():
-      ret[k].append(v)
-  return ret
-def collate_fn(batch): return batch
-
-# Reference (code): https://github.com/mlcommons/training/blob/2f4a93fb4888180755a8ef55f4b977ef8f60a89e/stable_diffusion/ldm/data/webdatasets.py, Line 55
-# Reference (params): https://github.com/mlcommons/training/blob/ab4ae1ca718d7fe62c369710a316dff18768d04b/stable_diffusion/configs/train_01x08x08.yaml, Line 107
-def batch_load_train_stable_diffusion(urls:str, BS:int):
-  import webdataset
-  dataset = webdataset.WebDataset(urls=urls, resampled=True, cache_size=-1, cache_dir=None)
-  dataset = dataset.shuffle(size=1000)
-  dataset = dataset.decode()
-  dataset = dataset.map(filter_dataset)
-  dataset = dataset.batched(BS, partial=False, collation_fn=collate)
-  dataset = webdataset.WebLoader(dataset, batch_size=None, shuffle=False, num_workers=1, persistent_workers=True, collate_fn=collate_fn)
-
-  for x in dataset:
-    assert isinstance(x, dict) and all(isinstance(k, str) for k in x.keys()) and all(isinstance(v, list) for v in x.values())
-    assert all(isinstance(moment_mean_logvar, np.ndarray) and moment_mean_logvar.shape==(1,8,64,64) for moment_mean_logvar in x["npy"])
-    assert all(isinstance(caption, str) for caption in x["txt"])
-    yield x
-
 # llama3

 class BinIdxDataset:
@@ -785,27 +758,6 @@ def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0
      batch.append(tokens)
    yield Tensor.stack(batch, dim=0)

-def batch_load_llama3_small(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
-  if val:
-    dataset = BlendedGPTDataset([
-      base_dir / "c4-validation-91205-samples.en_text_document",
-    ], [
-      1.0
-    ], samples, seqlen, seed, False)
-  else:
-    dataset = BlendedGPTDataset([
-      base_dir / "c4-train.en_6_text_document",
-    ], [
-      1.0
-    ], samples, seqlen, seed, True)
-
-  for b in range(math.ceil(samples / bs)):
-    batch = []
-    for i in range(bs):
-      tokens = dataset.get(b * bs + i)
-      batch.append(tokens)
-    yield Tensor.stack(batch, dim=0)
-
 if __name__ == "__main__":
  def load_unet3d(val):
    assert not val, "validation set is not supported due to different sizes on inputs"
@@ -2,9 +2,7 @@ import math
 from typing import Union

 from tinygrad import Tensor, nn, dtypes
-from tinygrad.helpers import prod, argfix, Context
-from tinygrad.nn.state import get_parameters
-from extra.models.unet import UNetModel
+from tinygrad.helpers import prod, argfix

 # rejection sampling truncated randn
 def rand_truncn(*shape, dtype=None, truncstds=2, **kwargs) -> Tensor:
@@ -19,10 +17,6 @@ def he_normal(*shape, a: float = 0.00, **kwargs) -> Tensor:
  std = math.sqrt(2.0 / (1 + a ** 2)) / math.sqrt(prod(argfix(*shape)[1:])) / 0.87962566103423978
  return std * rand_truncn(*shape, **kwargs)

-# Stable Diffusion v2 training uses default torch gelu, which doesn't use tanh approximation
-def gelu_erf(x:Tensor) -> Tensor:
-  return 0.5 * x * (1.0 + (x / 1.4142135623730951).erf())
-
 class Conv2dHeNormal(nn.Conv2d):
  def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True):
    super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
@@ -133,59 +127,3 @@ class Conv2dRetinaNet(nn.Conv2d):
  def __call__(self, x:Tensor) -> Tensor:
    return x.conv2d(self.weight.cast(dtypes.default_float), self.bias.cast(dtypes.default_float) if self.bias is not None else None,
                    groups=self.groups, stride=self.stride, dilation=self.dilation, padding=self.padding)
-
-# copy torch AMP: isolate mixed precision to just the below autocast ops, instead of using dtypes.default_float which affects all new Tensors
-class AutocastLinear(nn.Linear):
-  cast_dtype=dtypes.bfloat16 # enable monkeypatching of the mixed precision dtype
-  def __call__(self, x:Tensor) -> Tensor:
-    dtype = type(self).cast_dtype
-    return x.cast(dtype).linear(self.weight.cast(dtype).transpose(), self.bias.cast(dtype) if self.bias is not None else None)
-
-class AutocastConv2d(nn.Conv2d):
-  cast_dtype=dtypes.bfloat16
-  def __call__(self, x:Tensor) -> Tensor:
-    dtype = type(self).cast_dtype
-    return x.cast(dtype).conv2d(self.weight.cast(dtype), self.bias.cast(dtype), self.groups, self.stride, self.dilation, self.padding)
-
-# copy torch AMP: upcast to float32 before GroupNorm and LayerNorm
-class AutocastGroupNorm(nn.GroupNorm):
-  def __call__(self, x:Tensor) -> Tensor:
-    return super().__call__(x.cast(dtypes.float32))
-
-class AutocastLayerNorm(nn.LayerNorm):
-  def __call__(self, x:Tensor) -> Tensor:
-    return super().__call__(x.cast(dtypes.float32))
-
-def zero_module(module):
-  for p in get_parameters(module): p.assign(Tensor.zeros_like(p).contiguous())
-
-# Stable Diffusion mlperf reference doesn't call scaled_dot_product_attention
-# copy torch AMP: upcast to float32 before softmax on CUDA
-def attn_f32_softmax(q:Tensor, k:Tensor, v:Tensor) -> Tensor:
-  return (q.matmul(k.transpose(-2,-1), dtype=dtypes.float32) / math.sqrt(q.shape[-1])).softmax(-1).cast(q.dtype) @ v
-
-def init_stable_diffusion(version:str, pretrained:str, devices:list[str]):
-  from examples.stable_diffusion import StableDiffusion
-  from tinygrad.nn.state import safe_load, safe_save, load_state_dict, get_state_dict
-  from tempfile import TemporaryDirectory
-  model = StableDiffusion(version=version, pretrained=pretrained)
-  unet:UNetModel = model.model.diffusion_model
-
-  # this prevents extra consumption of memory, enabling much larger BS
-  Tensor.realize(*get_parameters(unet))
-  with TemporaryDirectory(prefix="unet_init") as tmp:
-    safe_save(get_state_dict(unet), init_fn:=f"{tmp}/init_model.safetensors")
-    load_state_dict(unet, safe_load(init_fn))
-
-  sqrt_alphas_cumprod = model.alphas_cumprod.sqrt().realize()
-  sqrt_one_minus_alphas_cumprod = (1 - model.alphas_cumprod).sqrt().realize()
-
-  if len(devices) > 1:
-    to_move = [sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod]
-    if version == "v2-mlperf-train": to_move += get_parameters(unet) + get_parameters(model.cond_stage_model)
-    for p in to_move:
-      p.to_(devices)
-    with Context(BEAM=0):
-      Tensor.realize(*to_move)
-
-  return model, unet, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod
@@ -1,9 +1,8 @@
 import math
-from tinygrad import dtypes, Tensor
+from tinygrad import dtypes
 from tinygrad.nn.optim import Optimizer

 from extra.lr_scheduler import LR_Scheduler
-from typing import Callable

 # https://github.com/mlcommons/training/blob/e237206991d10449d9675d95606459a3cb6c21ad/image_classification/tensorflow2/lars_util.py
 class PolynomialDecayWithWarmup(LR_Scheduler):
@@ -37,24 +36,4 @@ class CosineAnnealingLRWithWarmup(LR_Scheduler):
  def get_lr(self):
    warmup_lr = ((self.epoch_counter+1) / self.warmup_steps) * self.base_lr
    decay_lr = self.end_lr + 0.5 * (self.base_lr-self.end_lr) * (1 + (((self.epoch_counter+1-self.warmup_steps)/self.decay_steps) * math.pi).cos())
-    return (self.epoch_counter < self.warmup_steps).where(warmup_lr, decay_lr).cast(self.optimizer.lr.dtype)
-
-# Reference: https://github.com/mlcommons/training/blob/64b14a9abc74e08779a175abca7d291f8c957632/stable_diffusion/ldm/lr_scheduler.py, Lines 36-97
-class LambdaLinearScheduler:
-  def __init__(self, warm_up_steps:int, f_min:float, f_max:float, f_start:float, cycle_lengths:int):
-    self.lr_warm_up_steps, self.f_min, self.f_max, self.f_start, self.cycle_lengths = warm_up_steps, f_min, f_max, f_start, cycle_lengths
-
-  def schedule(self, n:Tensor) -> Tensor:
-    warm_up = (n < self.lr_warm_up_steps)
-    f_warm_up = (self.f_max - self.f_start) / self.lr_warm_up_steps * n + self.f_start
-    return warm_up.where(f_warm_up, self.f_min + (self.f_max - self.f_min) * (self.cycle_lengths - n) / (self.cycle_lengths))
-
-# based on torch.optim.lr_scheduler.LambdaLR
-class LambdaLR(LR_Scheduler):
-  def __init__(self, optimizer:Optimizer, base_lr:Tensor, lr_lambda:Callable):
-    super().__init__(optimizer)
-    self.base_lr, self.lr_lambda = base_lr, lr_lambda
-    self.step()
-
-  def get_lr(self):
-    return self.base_lr * self.lr_lambda(self.epoch_counter - 1)
+    return (self.epoch_counter < self.warmup_steps).where(warmup_lr, decay_lr).cast(self.optimizer.lr.dtype)
@@ -1,10 +1,10 @@
-import time, math, os
+import time, math
 start = time.perf_counter()
 from pathlib import Path
 import numpy as np
 from tinygrad import Tensor, Device, dtypes, GlobalCounters, TinyJit
 from tinygrad.nn.state import get_parameters, load_state_dict, safe_load
-from tinygrad.helpers import getenv, Context, prod
+from tinygrad.helpers import getenv
 from extra.bench_log import BenchEvent, WallTimeEvent
 def tlog(x): print(f"{x:25s}  @ {time.perf_counter()-start:5.2f}s")

@@ -243,299 +243,31 @@ def eval_mrcnn():

 def eval_llama3():
  from extra.models.llama import Transformer
-  from examples.llama3 import MODEL_PARAMS, load, convert_from_huggingface
+  from examples.llama3 import MODEL_PARAMS
  from tinygrad.helpers import tqdm

-  BASEDIR = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
-  BS = getenv("BS", 4)
-  SMALL = getenv("SMALL", 0)
-  SEQLEN = getenv("SEQLEN", 8192)
-  MODEL_PATH = Path(getenv("MODEL_PATH", "/raid/weights/llama31_8b/"))
+  bs = 4
+  sequence_length = 512

-  params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
-  params = params | {"vocab_size": 32000} if not SMALL else params
-  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
-  model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
-
-  # load weights
-  weights = load(str(MODEL_PATH / "model.safetensors.index.json"))
-  if "model.embed_tokens.weight" in weights:
-    print("converting from huggingface format")
-    weights = convert_from_huggingface(weights, params["n_layers"], params["n_heads"], params["n_kv_heads"])
-
-  load_state_dict(model, weights, strict=False, consume=True)
+  model = Transformer(**(MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}), max_context=sequence_length, jit=False, disable_kv_cache=True)

  @TinyJit
  def eval_step(model, tokens):
    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-    return loss.flatten().float()
+    return loss.flatten()

-  if SMALL:
-    from examples.mlperf.dataloader import batch_load_llama3_small
-    iter = batch_load_llama3_small(BS, 5760, SEQLEN, BASEDIR, val=True)
-  else:
-    from examples.mlperf.dataloader import batch_load_llama3
-    iter = batch_load_llama3(BS, 5760, SEQLEN, BASEDIR, val=True)
+  from examples.mlperf.dataloader import batch_load_llama3
+  iter = batch_load_llama3(bs, 5760, sequence_length, Path(getenv("BASEDIR", "/raid/datasets/c4/")), True)

  losses = []
-  for tokens in tqdm(iter, total=5760//BS):
+  for tokens in tqdm(iter, total=5760//bs):
    GlobalCounters.reset()
    losses += eval_step(model, tokens).tolist()
    tqdm.write(f"loss: {np.mean(losses)}")

-  log_perplexity = np.mean(losses)
-  print(f"Log Perplexity: {log_perplexity}")
-
-# NOTE: BEAM hangs on 8xmi300x with DECODE_BS=384 in final realize below; function is declared here for external testing
-@TinyJit
-def vae_decode(x:Tensor, vae, disable_beam=False) -> Tensor:
-  from examples.stable_diffusion import AutoencoderKL
-  assert isinstance(vae, AutoencoderKL)
-  x = vae.post_quant_conv(1./0.18215 * x)
-
-  x = vae.decoder.conv_in(x)
-  x = vae.decoder.mid(x)
-  for i, l in enumerate(vae.decoder.up[::-1]):
-    print("decode", x.shape)
-    for b in l['block']: x = b(x)
-    if 'upsample' in l:
-      bs,c,py,px = x.shape
-      x = x.reshape(bs, c, py, 1, px, 1).expand(bs, c, py, 2, px, 2).reshape(bs, c, py*2, px*2)
-      x = l['upsample']['conv'](x)
-    if i == len(vae.decoder.up) - 1 and disable_beam:
-      with Context(BEAM=0): x.realize()
-    else: x.realize()
-  x = vae.decoder.conv_out(vae.decoder.norm_out(x).swish())
-
-  x = ((x + 1.0) / 2.0).clip(0.0, 1.0)
-  return x
-
-def eval_stable_diffusion():
-  import csv, PIL, sys
-  from tqdm import tqdm
-  from examples.mlperf.initializers import init_stable_diffusion, gelu_erf
-  from examples.stable_diffusion import AutoencoderKL
-  from extra.models.unet import UNetModel
-  from tinygrad.nn.state import load_state_dict, torch_load
-  from tinygrad.helpers import BEAM
-  from extra.models import clip
-  from extra.models.clip import FrozenOpenClipEmbedder
-  from extra.models.clip import OpenClipEncoder
-  from extra.models.inception import FidInceptionV3
-
-  config = {}
-  GPUS               = config["GPUS"]                   = [f"{Device.DEFAULT}:{i}" for i in range(getenv("GPUS", 1))]
-  for x in GPUS: Device[x]
-  print(f"running eval on {GPUS}")
-  seed               = config["seed"]                   = getenv("SEED", 12345)
-  CKPTDIR            = config["CKPTDIR"]                = Path(getenv("CKPTDIR", "./checkpoints"))
-  DATADIR            = config["DATADIR"]                = Path(getenv("DATADIR", "./datasets"))
-  CONTEXT_BS         = config["CONTEXT_BS"]             = getenv("CONTEXT_BS", 1 * len(GPUS))
-  DENOISE_BS         = config["DENOISE_BS"]             = getenv("DENOISE_BS", 1 * len(GPUS))
-  DECODE_BS          = config["DECODE_BS"]              = getenv("DECODE_BS", 1 * len(GPUS))
-  INCEPTION_BS       = config["INCEPTION_BS"]           = getenv("INCEPTION_BS", 1 * len(GPUS))
-  CLIP_BS            = config["CLIP_BS"]                = getenv("CLIP_BS", 1 * len(GPUS))
-  EVAL_CKPT_DIR      = config["EVAL_CKPT_DIR"]          = getenv("EVAL_CKPT_DIR", "")
-  STOP_IF_CONVERGED  = config["STOP_IF_CONVERGED"]      = getenv("STOP_IF_CONVERGED", 0)
-
-  if (WANDB := getenv("WANDB", "")):
-    import wandb
-    wandb.init(config=config, project="MLPerf-Stable-Diffusion")
-
-  assert EVAL_CKPT_DIR != "", "provide a directory with checkpoints to be evaluated"
-  print(f"running eval on checkpoints in {EVAL_CKPT_DIR}\nSEED={seed}")
-  eval_queue:list[tuple[int, Path]] = []
-  for p in Path(EVAL_CKPT_DIR).iterdir():
-    if p.name.endswith(".safetensors"):
-      ckpt_iteration = p.name.split(".safetensors")[0]
-      assert ckpt_iteration.isdigit(), f"invalid checkpoint name: {p.name}, expected <digits>.safetensors"
-      eval_queue.append((int(ckpt_iteration), p))
-  assert len(eval_queue), f'no files ending with ".safetensors" were found in {EVAL_CKPT_DIR}'
-  print(sorted(eval_queue, reverse=True))
-
-  Tensor.manual_seed(seed)  # seed for weight initialization
-  model, unet, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod = init_stable_diffusion("v2-mlperf-eval", CKPTDIR / "sd" / "512-base-ema.ckpt", GPUS)
-
-  # load prompts for generating images for validation; 2 MB of data total
-  with open(DATADIR / "coco2014" / "val2014_30k.tsv") as f:
-    reader = csv.DictReader(f, delimiter="\t")
-    eval_inputs:list[dict] = [{"image_id": int(row["image_id"]), "id": int(row["id"]), "caption": row["caption"]} for row in reader]
-  assert len(eval_inputs) == 30_000
-  # NOTE: the clip weights are the same between model.cond_stage_model and clip_encoder
-  eval_timesteps = list(reversed(range(1, 1000, 20)))
-
-  original_device, Device.DEFAULT = Device.DEFAULT, "CPU"
-  # The choice of alphas_prev[0] = alphas_cumprod[0] seems arbitrary, but it's how the mlperf ref does it:
-  #   alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())
-  eval_alphas_prev = model.alphas_cumprod[0:1].cat(model.alphas_cumprod[list(range(1, 1000, 20))[:-1]]).to(GPUS).realize()
-  inception = FidInceptionV3().load_from_pretrained(CKPTDIR / "inception" / "pt_inception-2015-12-05-6726825d.pth")
-  vision_cfg = {'width': 1280, 'layers': 32, 'd_head': 80, 'image_size': 224, 'patch_size': 14}
-  text_cfg = {'width': 1024, 'n_heads': 16, 'layers': 24, 'vocab_size': 49408, 'ctx_length': 77}
-  clip.gelu = gelu_erf
-  clip_encoder = OpenClipEncoder(1024, text_cfg, vision_cfg)
-  loaded = torch_load(CKPTDIR / "clip" / "open_clip_pytorch_model.bin")
-  loaded.update({"attn_mask": clip_encoder.attn_mask, "mean": clip_encoder.mean, "std": clip_encoder.std})
-  load_state_dict(clip_encoder, loaded)
-  Device.DEFAULT=original_device
-
-  @TinyJit
-  def denoise_step(x:Tensor, x_x:Tensor, t_t:Tensor, uc_c:Tensor, sqrt_alphas_cumprod_t:Tensor, sqrt_one_minus_alphas_cumprod_t:Tensor,
-                    alpha_prev:Tensor, unet:UNetModel, GPUS) -> Tensor:
-    out_uncond, out = unet(x_x, t_t, uc_c).to("CPU").reshape(-1, 2, 4, 64, 64).chunk(2, dim=1)
-    out_uncond = out_uncond.squeeze(1).shard(GPUS,axis=0)
-    out = out.squeeze(1).shard(GPUS,axis=0)
-    v_t = out_uncond + 8.0 * (out - out_uncond)
-    e_t = sqrt_alphas_cumprod_t * v_t + sqrt_one_minus_alphas_cumprod_t * x
-    pred_x0 = sqrt_alphas_cumprod_t * x - sqrt_one_minus_alphas_cumprod_t * v_t
-    dir_xt = (1. - alpha_prev).sqrt() * e_t
-    x_prev = alpha_prev.sqrt() * pred_x0 + dir_xt
-    return x_prev.realize()
-
-  def shard_tensor(t:Tensor) -> Tensor: return t.shard(GPUS, axis=0) if len(GPUS) > 1 else t.to(GPUS[0])
-  def get_batch(whole:Tensor, i:int, bs:int) -> tuple[Tensor, int]:
-    batch = whole[i: i + bs].to("CPU")
-    if (unpadded_bs:=batch.shape[0]) < bs:
-      batch = batch.cat(batch[-1:].expand(bs - unpadded_bs, *batch[-1].shape))
-    return batch, unpadded_bs 
-
-  @Tensor.train(mode=False)
-  def eval_unet(eval_inputs:list[dict], unet:UNetModel, cond_stage:FrozenOpenClipEmbedder, first_stage:AutoencoderKL,
-                inception:FidInceptionV3, clip:OpenClipEncoder) -> tuple[float, float]:
-    # Eval is divided into 5 jits, one per model
-    # It doesn't make sense to merge these jits, e.g. unet repeats 50 times in isolation; images fork to separate inception/clip
-    # We're generating and scoring 30,000 images per eval, and all the data can flow through one jit at a time
-    # To maximize throughput for each jit, we have only one model/jit on the GPU at a time, and pool outputs from each jit off-GPU
-    for model in (unet, first_stage, inception, clip):
-      Tensor.realize(*[p.to_("CPU") for p in get_parameters(model)])
-
-    uc_written = False
-    models = (cond_stage, unet, first_stage, inception, clip)
-    jits = (jit_context:=TinyJit(cond_stage.embed_tokens), denoise_step, vae_decode, jit_inception:=TinyJit(inception),
-            jit_clip:=TinyJit(clip.get_clip_score))
-    all_bs = (CONTEXT_BS, DENOISE_BS, DECODE_BS, INCEPTION_BS, CLIP_BS)
-    if (EVAL_SAMPLES:=getenv("EVAL_SAMPLES", 0)) and EVAL_SAMPLES > 0:
-      eval_inputs = eval_inputs[0:EVAL_SAMPLES]
-    output_shapes = [(ns:=len(eval_inputs),77), (ns,77,1024), (ns,4,64,64), (ns,3,512,512), (ns,2048), (ns,)]
-    # Writing progress to disk lets us resume eval if we crash
-    stages = ["tokens", "embeds", "latents", "imgs", "inception", "clip"]
-    disk_tensor_names, disk_tensor_shapes = stages + ["end", "uc"], output_shapes + [(6,), (1,77,1024)]
-    if not all(os.path.exists(f"{EVAL_CKPT_DIR}/{name}.bytes") for name in disk_tensor_names):
-      for name, shape in zip(disk_tensor_names, disk_tensor_shapes):
-        file = Path(f"{EVAL_CKPT_DIR}/{name}.bytes")
-        file.unlink(missing_ok=True)
-        with file.open("wb") as f: f.truncate(prod(shape) * 4)
-    progress = {name: Tensor.empty(*shape, device=f"disk:{EVAL_CKPT_DIR}/{name}.bytes", dtype=dtypes.int if name in {"tokens", "end"} else dtypes.float)
-                for name, shape in zip(disk_tensor_names, disk_tensor_shapes)}
-
-    def embed_tokens(tokens:Tensor) -> Tensor:
-      nonlocal uc_written
-      if not uc_written:
-        with Context(BEAM=0): progress["uc"].assign(cond_stage.embed_tokens(cond_stage.tokenize("").to(GPUS)).to("CPU").realize()).realize()
-        uc_written = True
-      return jit_context(shard_tensor(tokens))
-
-    def generate_latents(embeds:Tensor) -> Tensor:
-      uc_c = Tensor.stack(progress["uc"].to("CPU").expand(bs, 77, 1024), embeds, dim=1).reshape(-1, 77, 1024)
-      uc_c = shard_tensor(uc_c)
-      x = shard_tensor(Tensor.randn(bs,4,64,64))
-      for step_idx, timestep in enumerate(tqdm(eval_timesteps)):
-        reversed_idx = Tensor([50 - step_idx - 1], device=GPUS)
-        alpha_prev = eval_alphas_prev[reversed_idx]
-        ts = Tensor.full(bs, fill_value=timestep, dtype=dtypes.int, device="CPU")
-        ts_ts = shard_tensor(ts.cat(ts))
-        ts = shard_tensor(ts)
-        sqrt_alphas_cumprod_t = sqrt_alphas_cumprod[ts].reshape(bs, 1, 1, 1)
-        sqrt_one_minus_alphas_cumprod_t = sqrt_one_minus_alphas_cumprod[ts].reshape(bs, 1, 1, 1)
-        x_x = shard_tensor(Tensor.stack(x.to("CPU"), x.to("CPU"), dim=1).reshape(-1, 4, 64, 64))
-        x.assign(denoise_step(x, x_x, ts_ts, uc_c, sqrt_alphas_cumprod_t, sqrt_one_minus_alphas_cumprod_t, alpha_prev, unet, GPUS)).realize()
-      return x
-
-    def decode_latents(latents:Tensor) -> Tensor: return vae_decode(shard_tensor(latents), first_stage, disable_beam=True)
-    def generate_inception(imgs:Tensor) -> Tensor: return jit_inception(shard_tensor(imgs))[:,:,0,0]
-
-    def calc_clip_scores(batch:Tensor, batch_tokens:Tensor) -> Tensor:
-      # Tensor.interpolate does not yet support bicubic, so we use PIL
-      batch = (batch.to(GPUS[0]).permute(0,2,3,1) * 255).clip(0, 255).cast(dtypes.uint8).numpy()
-      batch = [np.array(PIL.Image.fromarray(batch[i]).resize((224,224), PIL.Image.BICUBIC)) for i in range(bs)]
-      batch = shard_tensor(Tensor(np.stack(batch, axis=0).transpose(0,3,1,2), device="CPU").realize())
-      batch = batch.cast(dtypes.float) / 255
-      batch = (batch - model.mean) / model.std
-      batch = jit_clip(shard_tensor(batch_tokens), batch)
-      return batch
-
-    callbacks = (embed_tokens, generate_latents, decode_latents, generate_inception, calc_clip_scores)
-
-    # save every forward pass output to disk; NOTE: this needs ~100 GB disk space because 30k images are large
-    def stage_progress(stage_idx:int) -> int: return progress["end"].to("CPU")[stage_idx].item()
-    if stage_progress(0) < len(eval_inputs):
-      tokens = []
-      for i in tqdm(range(0, len(eval_inputs), CONTEXT_BS)):
-        subset = [cond_stage.tokenize(row["caption"], device="CPU") for row in eval_inputs[i: i+CONTEXT_BS]]
-        tokens.append(Tensor.cat(*subset, dim=0).realize())
-      progress["tokens"].assign(Tensor.cat(*tokens, dim=0).realize()).realize()
-      progress["end"][0:1].assign(Tensor([len(eval_inputs)], dtype=dtypes.int)).realize()
-    prev_stage = "tokens"
-    tokens = progress["tokens"]
-
-    # wrapper code for every model
-    for stage_idx, model, jit, bs, callback in zip(range(1,6), models, jits, all_bs, callbacks):
-      stage = stages[stage_idx]
-      if stage_progress(stage_idx) >= len(eval_inputs):
-        prev_stage = stage
-        continue # use cache
-      t0 = time.perf_counter()
-      print(f"starting eval with model: {model}")
-      if stage_idx == 1: inputs = tokens
-      elif stage_idx == 5: inputs = progress["imgs"]
-      else: inputs = progress[prev_stage]
-
-      Tensor.realize(*[p.to_(GPUS) for p in get_parameters(model)])
-      for batch_idx in tqdm(range(stage_progress(stage_idx), inputs.shape[0], bs)):
-        t1 = time.perf_counter()
-        batch, unpadded_bs = get_batch(inputs, batch_idx, bs)
-        if isinstance(model, OpenClipEncoder): batch = callback(batch, get_batch(tokens, batch_idx, bs)[0].realize())
-        else: batch = callback(batch)
-        # to(GPUS[0]) is necessary for this to work, without that the result is still on GPUS, probably due to a bug
-        batch = batch.to(GPUS[0]).to("CPU")[0:unpadded_bs].realize()
-        progress[stage][batch_idx: batch_idx + bs].assign(batch).realize()
-        # keep track of what our last output was, so we can resume from there if we crash in this loop
-        progress["end"][stage_idx: stage_idx + 1].assign(Tensor([batch_idx + bs], dtype=dtypes.int)).realize()
-        print(f"model: {model}, batch_idx: {batch_idx}, elapsed: {(time.perf_counter() - t1):.2f}")
-      del batch
-        
-      jit.reset()
-      Tensor.realize(*[p.to_("CPU") for p in get_parameters(model)])
-      print(f"done with model: {model}, elapsed: {(time.perf_counter() - t0):.2f}")
-      prev_stage = stage
-
-    inception_stats_fn = str(DATADIR / "coco2014" / "val2014_30k_stats.npz")
-    fid_score = inception.compute_score(progress["inception"].to("CPU"), inception_stats_fn)
-    clip_score = progress["clip"].to(GPUS[0]).mean().item()
-    for name in disk_tensor_names:
-      Path(f"{EVAL_CKPT_DIR}/{name}.bytes").unlink(missing_ok=True)
-    
-    if EVAL_SAMPLES and BEAM:
-      print("BEAM COMPLETE", flush=True) # allows wrapper script to detect BEAM search completion and retry if it failed
-      sys.exit() # Don't eval additional models; we don't care about clip/fid scores when running BEAM on eval sample subset
-
-    return clip_score, fid_score
-
-  # evaluate checkpoints in reverse chronological order
-  for ckpt_iteration, p in sorted(eval_queue, reverse=True):
-    unet_ckpt = safe_load(p)
-    load_state_dict(unet, unet_ckpt)
-    clip_score, fid_score = eval_unet(eval_inputs, unet, model.cond_stage_model, model.first_stage_model, inception, clip_encoder)
-    converged = True if clip_score >= 0.15 and fid_score <= 90 else False
-    print(f"eval results for {EVAL_CKPT_DIR}/{p.name}: clip={clip_score}, fid={fid_score}, converged={converged}")
-    if WANDB:
-      wandb.log({"eval/ckpt_iteration": ckpt_iteration, "eval/clip_score": clip_score, "eval/fid_score": fid_score})
-    if converged and STOP_IF_CONVERGED:
-      print(f"Convergence detected, exiting early before evaluating other checkpoints due to STOP_IF_CONVERGED={STOP_IF_CONVERGED}")
-      sys.exit()
-
-  # for testing
-  return clip_score, fid_score, ckpt_iteration
+  log_perplexity = Tensor(losses).mean()
+  print(f"Log Perplexity: {log_perplexity.item()}")

 if __name__ == "__main__":
  # inference only
@@ -3,8 +3,8 @@ from pathlib import Path
 import multiprocessing

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

 from extra.lr_scheduler import LRSchedulerGroup
@@ -252,10 +252,6 @@ def train_resnet():
        print(f"epoch global_ops: {steps_in_train_epoch * GlobalCounters.global_ops:_}, "
              f"epoch global_mem: {steps_in_train_epoch * GlobalCounters.global_mem:_}")
        # if we are doing beam search, run the first eval too
-        if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
-          min_time = min(step_times)
-          assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
-
        if (TRAIN_BEAM or EVAL_BEAM) and e == start_epoch: break
        return
    if MLLOGGER and RUNMLPERF:
@@ -348,8 +344,6 @@ def train_resnet():
        print(f"saving ckpt to {fn}")
        safe_save(get_training_state(model, optimizer_group, scheduler_group), fn)

-
-
 def train_retinanet():
  from contextlib import redirect_stdout
  from examples.mlperf.dataloader import batch_load_retinanet
@@ -707,7 +701,7 @@ def train_unet3d():
  ```BASEDIR=<folder_path> ./examples/mlperf/scripts/setup_kits19_dataset.sh```

  2) To start training the model, run the following:
-  ```time PYTHONPATH=. WANDB=1 TRAIN_BEAM=3 GPUS=6 BS=6 MODEL=unet3d python3 examples/mlperf/model_train.py```
+  ```time PYTHONPATH=. WANDB=1 TRAIN_BEAM=3 FUSE_CONV_BW=1 GPUS=6 BS=6 MODEL=unet3d python3 examples/mlperf/model_train.py```
  """
  from examples.mlperf.losses import dice_ce_loss
  from examples.mlperf.metrics import dice_score
@@ -749,6 +743,7 @@ def train_unet3d():
    "train_beam": TRAIN_BEAM,
    "eval_beam": EVAL_BEAM,
    "wino": WINO.value,
+    "fuse_conv_bw": FUSE_CONV_BW.value,
    "gpus": GPUS,
    "default_float": dtypes.default_float.name
  }
@@ -1295,20 +1290,15 @@ def train_llama3():
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup

  config = {}
-  BASEDIR            = config["BASEDIR"]                = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
  BS                 = config["BS"]                     = getenv("BS", 16)
  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
  SEED               = config["SEED"]                   = getenv("SEED", 5760)
  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
  TRAIN_ON_VAL       = config["TRAIN_ON_VAL"]           = getenv("TRAIN_ON_VAL", 0)
-  SMALL              = config["SMALL"]                  = getenv("SMALL", 0)
  SAMPLES            = config["SAMPLES"]                = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000 * 1152)
-  EVAL_FREQ          = config["EVAL_FREQ"]              = getenv("EVAL_FREQ", 46080)
-  EVAL_BS            = config["EVAL_BS"]                = getenv("EVAL_BS", 16)
-  EVAL_TARGET        = config["EVAL_TARGET"]            = getenv("EVAL_TARGET", 5.6)

-  # LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
+  # LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 FUSE_ARANGE=1 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
  # trains to 7

  opt_adamw_beta_1 = 0.9
@@ -1318,14 +1308,13 @@ def train_llama3():

  opt_gradient_clip_norm = 1.0
  opt_learning_rate_warmup_steps = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
-  opt_learning_rate_decay_steps = getenv("MAX_STEPS", math.ceil(1_200_000 * 1152 / GBS)) - opt_learning_rate_warmup_steps
+  opt_learning_rate_decay_steps = getenv("DECAY_STEPS", math.ceil(1_200_000 * 1152 / GBS) - opt_learning_rate_warmup_steps)
  opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152)  # NOTE: cannot change for benchmark
-  opt_end_learning_rate = getenv("END_LR", 8e-7)
+  opt_end_learning_rate = 8e-7

  # TODO: confirm weights are in bf16
  # vocab_size from the mixtral tokenizer
-  params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
-  params = params | {"vocab_size": 32000} if not SMALL else params
+  params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}
  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
  model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)

@@ -1361,15 +1350,6 @@ def train_llama3():
                b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
  scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)

-  if resume_ckpt := getenv("RESUME_CKPT"):
-    fn = f"./ckpts/llama3_{resume_ckpt}.safe"
-    print(f"loading initial checkpoint from {fn}")
-    load_state_dict(model, safe_load(fn), realize=False)
-
-    fn = f"./ckpts/llama3_{resume_ckpt}_optim.safe"
-    print(f"loading optim checkpoint from {fn}")
-    load_state_dict(scheduler, safe_load(fn), realize=False)
-
  @TinyJit
  @Tensor.train()
  def train_step(model, tokens:Tensor, grad_acc:int):
@@ -1395,7 +1375,7 @@ def train_llama3():
        total_norm += p.grad.float().square().sum()
      total_norm = total_norm.sqrt().contiguous()
      for p in optim.params:
-        p.grad = p.grad * (opt_gradient_clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)
+        p.grad = p.grad * opt_gradient_clip_norm / (total_norm + 1e-6)

    optim.step()
    scheduler.step()
@@ -1404,231 +1384,33 @@ def train_llama3():
    loss.realize(lr)
    return loss, lr

-  @TinyJit
-  @Tensor.train(False)
-  def eval_step(model, tokens:Tensor):
-    if (DP := getenv("DP", 1)) > 1:
-      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-      tokens = tokens.shard(device, 0)
-    if (MP := getenv("MP", 1)) > 1:
-      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
-      tokens = tokens.shard(device)
-    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-    return loss.flatten().float()
+  if getenv("FAKEDATA", 0):
+    def fake_data():
+      for _ in range(SAMPLES // GBS):
+        yield Tensor.randint(GBS, SEQLEN + 1, low=0, high=32000, dtype=dtypes.int32, device=Device.DEFAULT)
+    iter = fake_data()
+  else:
+    from examples.mlperf.dataloader import batch_load_llama3
+    iter = batch_load_llama3(GBS, SAMPLES, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=bool(TRAIN_ON_VAL))

-  # ** data iters **
-  def fake_data(bs, samples):
-    for _ in range(samples // bs):
-      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
-
-  def get_train_iter():
-    if getenv("FAKEDATA", 0):
-      return fake_data(GBS, SAMPLES)
-    else:
-      if SMALL:
-        from examples.mlperf.dataloader import batch_load_llama3_small
-        return batch_load_llama3_small(GBS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
-      else:
-        from examples.mlperf.dataloader import batch_load_llama3
-        return batch_load_llama3(GBS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
-
-  def get_eval_iter():
-    if getenv("FAKEDATA", 0):
-      return fake_data(EVAL_BS, 5760)
-    else:
-      if SMALL:
-        from examples.mlperf.dataloader import batch_load_llama3_small
-        return batch_load_llama3_small(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
-      else:
-        from examples.mlperf.dataloader import batch_load_llama3
-        return batch_load_llama3(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
-
-  iter = get_train_iter()
-  i, sequences_seen = resume_ckpt, 0
+  i = 0
  for tokens in tqdm(iter, total=SAMPLES//GBS):
    t = time.perf_counter()
    GlobalCounters.reset()
    loss, lr = train_step(model, tokens, grad_acc)
    loss = loss.float().item()
-
-    i += 1
-    sequences_seen += tokens.shape[0]
-
+    # above as tqdm.write f-string
    tqdm.write(f"{loss:.4f} loss, {lr.item():.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, {time.perf_counter()-t:.2f} s")
    if (fname:=getenv("LOSS_FILE", "")):
      with open(fname, "a") as f:
        f.write(f"{i} {loss:.4f} {lr.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n")

-    if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
+    if getenv("CKPT") and (i % 200 == 0 or i == 10):
      tqdm.write("saving checkpoint")
      if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
-      fn = f"{ckpt_dir}/llama3_{i}.safe"
+      fn = f"{ckpt_dir}/{i}.safe"
      safe_save(get_state_dict(model), fn)
-
-      tqdm.write("saving optim checkpoint")
-      fn = f"{ckpt_dir}/llama3_{i}_optim.safe"
-      safe_save(get_state_dict(scheduler), fn)
-
-    if sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1):
-      tqdm.write(f"evaluating after {sequences_seen} sequences")
-
-      # run eval
-      eval_losses = []
-      eval_iter = get_eval_iter()
-      tqdm.write(f"evaluating {5760//EVAL_BS} batches of {EVAL_BS} sequences")
-
-      for tokens in tqdm(eval_iter, total=5760//EVAL_BS):
-        eval_losses += eval_step(model, tokens).tolist()
-      log_perplexity = Tensor(eval_losses).mean().float().item()
-
-      tqdm.write(f"eval log perplexity: {log_perplexity:.4f}")
-
-      if log_perplexity < EVAL_TARGET:
-        tqdm.write(f"target achieved after {sequences_seen} sequences")
-        if getenv("CKPT"):
-          if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
-          fn = f"{ckpt_dir}/llama3.safe"
-          safe_save(get_state_dict(model), fn)
-        break
-
-def train_stable_diffusion():
-  from extra.models.unet import UNetModel
-  from examples.mlperf.dataloader import batch_load_train_stable_diffusion
-  from examples.mlperf.lr_schedulers import LambdaLR, LambdaLinearScheduler
-  from examples.mlperf.initializers import init_stable_diffusion
-  from examples.mlperf.helpers import get_training_state
-  import numpy as np
-
-  config = {}
-  GPUS               = config["GPUS"]                   = [f"{Device.DEFAULT}:{i}" for i in range(getenv("GPUS", 1))]
-  seed               = config["seed"]                   = getenv("SEED", 12345)
-  # ** hyperparameters **
-  BS                 = config["BS"]                     = getenv("BS", 1 * len(GPUS))
-  BASE_LR            = config["LEARNING_RATE"]          = getenv("LEARNING_RATE", 2.5e-7)
-  # https://github.com/mlcommons/training_policies/blob/cfa99da479b8d5931f7a3c67612d021dfb47510a/training_rules.adoc#benchmark_specific_rules
-  # "Checkpoint must be collected every 512,000 images. CEIL(512000 / global_batch_size) if 512000 is not divisible by GBS."
-  # NOTE: It's inferred that "steps" is the unit for the output of the CEIL formula, based on all other cases of CEIL in the rules
-  CKPT_STEP_INTERVAL = config["CKPT_STEP_INTERVAL"]     = getenv("CKPT_STEP_INTERVAL", math.ceil(512_000 / BS))
-  CKPTDIR            = config["CKPTDIR"]                = Path(getenv("CKPTDIR", "./checkpoints"))
-  DATADIR            = config["DATADIR"]                = Path(getenv("DATADIR", "./datasets"))
-  UNET_CKPTDIR       = config["UNET_CKPTDIR"]           = Path(getenv("UNET_CKPTDIR", "./checkpoints"))
-  TOTAL_CKPTS        = config["TOTAL_CKPTS"]            = getenv("TOTAL_CKPTS", 0)
-
-  print(f"training on {GPUS}")
-  lr = BS * BASE_LR
-  print(f"BS={BS}, BASE_LR={BASE_LR}, lr={lr}")
-  print(f"CKPT_STEP_INTERVAL = {CKPT_STEP_INTERVAL}")
-  for x in GPUS: Device[x]
-  if (WANDB := getenv("WANDB", "")):
-    import wandb
-    wandb.init(config=config, project="MLPerf-Stable-Diffusion")
-
-  Tensor.manual_seed(seed)  # seed for weight initialization
-  model, unet, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod = init_stable_diffusion("v2-mlperf-train", CKPTDIR / "sd" / "512-base-ema.ckpt", GPUS)
-
-  optimizer = AdamW(get_parameters(unet))
-  lambda_lr_callback = LambdaLinearScheduler(1000, 1.0, 1.0, 1e-06, 10000000000000).schedule
-  lr_scheduler = LambdaLR(optimizer, Tensor(lr, dtype=dtypes.float, device=optimizer.device), lambda_lr_callback)
-
-  @TinyJit
-  def train_step(mean:Tensor, logvar:Tensor, tokens:Tensor, unet:UNetModel, optimizer:LAMB, lr_scheduler:LambdaLR) -> Tensor:
-    optimizer.zero_grad()
-
-    timestep = Tensor.randint(BS, low=0, high=model.alphas_cumprod.shape[0], dtype=dtypes.int, device=GPUS[0])
-    latent_randn = Tensor.randn(*mean.shape, device=GPUS[0])
-    noise = Tensor.randn(*mean.shape, device=GPUS[0])
-    for t in (mean, logvar, tokens, timestep, latent_randn, noise):
-      t.shard_(GPUS, axis=0)
-
-    std = Tensor.exp(0.5 * logvar.clamp(-30.0, 20.0))
-    latent = (mean + std * latent_randn) * 0.18215
-
-    sqrt_alphas_cumprod_t = sqrt_alphas_cumprod[timestep].reshape(timestep.shape[0], 1, 1, 1)
-    sqrt_one_minus_alphas_cumprod_t = sqrt_one_minus_alphas_cumprod[timestep].reshape(timestep.shape[0], 1, 1, 1)
-    latent_with_noise = sqrt_alphas_cumprod_t * latent + sqrt_one_minus_alphas_cumprod_t * noise
-    v_true = sqrt_alphas_cumprod_t * noise - sqrt_one_minus_alphas_cumprod_t * latent
-
-    context = model.cond_stage_model.embed_tokens(tokens)
-
-    out = unet(latent_with_noise, timestep, context)
-    loss = ((out - v_true) ** 2).mean()
-    del mean, logvar, std, latent, noise, sqrt_alphas_cumprod_t, sqrt_one_minus_alphas_cumprod_t
-    del out, v_true, context, latent_randn, tokens, timestep
-    loss.backward()
-
-    optimizer.step()
-    lr_scheduler.step()
-    loss, out_lr = loss.detach().to("CPU"), optimizer.lr.to("CPU")
-    Tensor.realize(loss, out_lr)
-    return loss, out_lr
-    
-  # checkpointing takes ~9 minutes without this, and ~1 minute with this
-  @TinyJit
-  def ckpt_to_cpu():
-    ckpt = get_training_state(unet, optimizer, lr_scheduler)
-    # move to CPU first so more GPU bufs aren't created (can trigger OOM)
-    for k,v in ckpt.items(): ckpt[k] = v.detach().to("CPU")
-    Tensor.realize(*[v for v in ckpt.values()])
-    for k,v in ckpt.items(): ckpt[k] = v.cast(v.dtype.base).contiguous()
-    Tensor.realize(*[v for v in ckpt.values()])
-    return ckpt
-
-  # training loop
-  dl = batch_load_train_stable_diffusion(f'{DATADIR}/laion-400m/webdataset-moments-filtered/{{00000..00831}}.tar', BS)
-  # for tests
-  saved_checkpoints = []
-
-  train_start_time = time.perf_counter()
-  t0 = t6 = time.perf_counter()
-  for i, batch in enumerate(dl, start=1):
-    loop_time = time.perf_counter() - t0
-    t0 = time.perf_counter()
-    dl_time = t0 - t6
-    GlobalCounters.reset()
-
-    mean, logvar = np.split(np.concatenate(batch["npy"], axis=0), 2, axis=1)
-    mean, logvar = Tensor(mean, dtype=dtypes.float32, device="CPU"), Tensor(logvar, dtype=dtypes.float32, device="CPU")
-    tokens = []
-    for text in batch['txt']: tokens += model.cond_stage_model.tokenizer.encode(text, pad_with_zeros=True)
-    tokens = Tensor(tokens, dtype=dtypes.int32, device="CPU").reshape(-1, 77)
-
-    t1 = time.perf_counter()
-    loss, lr = train_step(mean, logvar, tokens, unet, optimizer, lr_scheduler)
-    loss_item, lr_item = loss.item(), lr.item()
-    t2 = time.perf_counter()
-
-    if i == 3:
-      for _ in range(3): ckpt_to_cpu() # do this at the beginning of run to prevent OOM surprises when checkpointing
-      print("BEAM COMPLETE", flush=True) # allows wrapper script to detect BEAM search completion and retry if it failed
-      
-    total_train_time = time.perf_counter() - train_start_time
-    if WANDB:
-      wandb.log({"train/loss": loss_item, "train/lr": lr_item, "train/loop_time_prev": loop_time, "train/dl_time": dl_time, "train/step": i,
-                    "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (t2-t1), "train/input_prep_time": t1-t0,
-                    "train/train_step_time": t2-t1, "train/total_time": total_train_time})
-
-      if i == 1 and wandb.run is not None:
-        with open(f"{UNET_CKPTDIR}/wandb_run_id_{wandb.run.id}", "w") as f:
-          f.write(f"wandb.run.id = {wandb.run.id}")
-
-    if i % CKPT_STEP_INTERVAL == 0:
-      # https://github.com/mlcommons/training_policies/blob/cfa99da479b8d5931f7a3c67612d021dfb47510a/training_rules.adoc#benchmark_specific_rules
-      # "evaluation is done offline, the time is not counted towards the submission time."
-      fn = f"{UNET_CKPTDIR}/{i}.safetensors"
-      print(f"saving unet checkpoint at {fn}")
-      saved_checkpoints.append(fn)
-      safe_save({k.replace("model.", ""):v for k,v in ckpt_to_cpu().items() if k.startswith("model.")}, fn)
-      if TOTAL_CKPTS and i == TOTAL_CKPTS * CKPT_STEP_INTERVAL:
-        print(f"ending run after {i} steps ({TOTAL_CKPTS} checkpoints collected)")
-        return saved_checkpoints
-
-    t3 = time.perf_counter()
-    print(f"""step {i}: {GlobalCounters.global_ops * 1e-9 / (t2-t1):9.2f} GFLOPS, mem_used: {GlobalCounters.mem_used / 1e9:.2f} GB,
-  loop_time_prev: {loop_time:.2f}, dl_time: {dl_time:.2f}, input_prep_time: {t1-t0:.2f}, train_step_time: {t2-t1:.2f},
-  t3-t2: {t3-t2:.4f}, loss:{loss_item:.5f}, lr:{lr_item:.3e}, total_train_time:{total_train_time:.2f}
-  """)
-    t6 = time.perf_counter()
+    i += 1

 if __name__ == "__main__":
  multiprocessing.set_start_method('spawn')
@@ -1638,7 +1420,7 @@ if __name__ == "__main__":
  else: bench_log_manager = contextlib.nullcontext()

  with Tensor.train():
-    for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,maskrcnn,stable_diffusion").split(","):
+    for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,maskrcnn").split(","):
      nm = f"train_{m}"
      if nm in globals():
        print(f"training {m}")
@@ -1,57 +0,0 @@
-#!/usr/bin/env bash
-# adapted from https://github.com/mlcommons/training/blob/4bdf5c8ed218ad76565a2ba1ac27c919ccc6d689/stable_diffusion/README.md
-
-# setup dirs
-
-DATA=/raid/datasets/stable_diffusion
-
-LAION=$DATA/laion-400m/webdataset-moments-filtered 
-COCO=$DATA/coco2014
-mkdir -p $LAION $COCO
-
-CKPT=/raid/weights/stable_diffusion
-mkdir -p $CKPT/clip $CKPT/sd $CKPT/inception
-
-# download data
-
-# if rclone isn't installed system-wide / in your PATH, put the executable path in quotes below
-#RCLONE=""
-RCLONE="rclone"
-
-## VAE-encoded image latents, from 6.1M image subset of laion-400m
-## about 1 TB for whole download
-$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
-$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/laion-400m/moments-webdataset-filtered/ ${LAION} --include="*.tar" -P
-$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/laion-400m/moments-webdataset-filtered/sha512sums.txt ${LAION} -P
-cd $LAION && grep -E '\.tar$' sha512sums.txt | sha512sum -c --quiet - && \
-  echo "All .tar files verified" || { echo "Checksum failure when validating downloaded Laion moments"; exit 1; }
-
-## prompts and FID statistics from 30k image subset of coco2014
-## 33 MB
-$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
-$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/coco2014/val2014_30k.tsv ${COCO} -P
-
-$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
-$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/coco2014/val2014_30k_stats.npz ${COCO} -P
-
-# download checkpoints
-
-## clip (needed for text and vision encoders for validation)
-CLIP_WEIGHTS_URL="https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin"
-CLIP_WEIGHTS_SHA256="9a78ef8e8c73fd0df621682e7a8e8eb36c6916cb3c16b291a082ecd52ab79cc4"
-CLIP_CONFIG_URL="https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/raw/main/open_clip_config.json"
-wget -N -P ${CKPT}/clip ${CLIP_WEIGHTS_URL}
-wget -N -P ${CKPT}/clip ${CLIP_CONFIG_URL}
-echo "${CLIP_WEIGHTS_SHA256}  ${CKPT}/clip/open_clip_pytorch_model.bin"                    | sha256sum -c
-
-## sd (needed for latent->image decoder for validation, also has clip text encoder for training)
-SD_WEIGHTS_URL='https://huggingface.co/stabilityai/stable-diffusion-2-base/resolve/main/512-base-ema.ckpt'
-SD_WEIGHTS_SHA256="d635794c1fedfdfa261e065370bea59c651fc9bfa65dc6d67ad29e11869a1824"
-wget -N -P ${CKPT}/sd ${SD_WEIGHTS_URL}
-echo "${SD_WEIGHTS_SHA256}  ${CKPT}/sd/512-base-ema.ckpt"                    | sha256sum -c
-
-## inception (needed for validation)
-FID_WEIGHTS_URL='https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth'
-FID_WEIGHTS_SHA1="bd836944fd6db519dfd8d924aa457f5b3c8357ff"
-wget -N -P ${CKPT}/inception ${FID_WEIGHTS_URL}
-echo "${FID_WEIGHTS_SHA1}  ${CKPT}/inception/pt_inception-2015-12-05-6726825d.pth"                    | sha1sum -c
@@ -1,72 +0,0 @@
-#!/usr/bin/env bash
-
-DATETIME=${2:-$(date "+%m%d%H%M")}
-LOGFILE="${HOME}/logs/sd_mi300x_${DATETIME}.log"
-# UNET_CKPTDIR must be set: training saves checkpoints to this path, then a separate eval process scans this path to know which checkpoints to eval
-export UNET_CKPTDIR="${HOME}/stable_diffusion/training_checkpoints/${DATETIME}"
-mkdir -p "${HOME}/logs" "$UNET_CKPTDIR"
-
-# run this script in isolation when using the --bg flag
-if [[ "${1:-}" == "--bg" ]]; then
-  echo "logging output to $LOGFILE"
-  echo "saving UNet checkpoints to $UNET_CKPTDIR"
-  script_path="$(readlink -f "${BASH_SOURCE[0]}")"
-  nohup bash "$script_path" run "$DATETIME" >"$LOGFILE" 2>&1 & disown $!
-  exit 0
-fi
-
-# venv management
-if [[ -d .venv-sd-mlperf ]]; then
-  . .venv-sd-mlperf/bin/activate
-else
-  python3 -m venv .venv-sd-mlperf && . .venv-sd-mlperf/bin/activate
-  pip install --index-url https://download.pytorch.org/whl/cpu torch && pip install tqdm numpy ftfy regex pillow scipy wandb webdataset
-fi
-pip list
-apt list --installed | grep amdgpu
-rocm-smi --version
-modinfo amdgpu | grep version
-
-export BEAM=2 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 IGNORE_JIT_FIRST_BEAM=1 HCQDEV_WAIT_TIMEOUT_MS=300000
-export AMD_LLVM=0 # bf16 seems to require this
-export DATADIR="/raid/datasets/stable_diffusion"
-export CKPTDIR="/raid/weights/stable_diffusion"
-export EVAL_CKPT_DIR=$UNET_CKPTDIR
-export MODEL="stable_diffusion" PYTHONPATH="."
-export GPUS=8 BS=304
-export CONTEXT_BS=816 DENOISE_BS=600 DECODE_BS=384 INCEPTION_BS=560 CLIP_BS=240
-export WANDB=1
-export PARALLEL=4
-export PYTHONUNBUFFERED=1
-sudo rocm-smi -d 0 1 2 3 4 5 6 7 --setperfdeterminism 1500 || exit 1
-
-# Retry BEAM search if script fails before BEAM COMPLETE is printed, but don't retry after that
-run_retry(){ local try=0 max=5 code tmp py pgid kids
-  while :; do
-    tmp=$(mktemp)
-    setsid bash -c 'exec env "$@"' _ "$@" > >(tee -a "$LOGFILE" | tee "$tmp") 2>&1 &
-    py=$!; pgid=$(ps -o pgid= -p "$py" | tr -d ' ')
-    wait "$py"; code=$?
-    [[ -n "$pgid" ]] && { kill -TERM -"$pgid" 2>/dev/null; sleep 1; kill -KILL -"$pgid" 2>/dev/null; }
-    kids=$(pgrep -P "$py" || true)
-    while [[ -n "$kids" ]]; do
-      kill -TERM $kids 2>/dev/null; sleep 0.5
-      kids=$(for k in $kids; do pgrep -P "$k" || true; done)
-    done
-    grep -q 'BEAM COMPLETE' "$tmp" && { rm -f "$tmp"; return 1; }
-    rm -f "$tmp"
-    ((code==0)) && return 0
-    ((try>=max)) && return 2
-    ((try++)); sleep 90; echo "try = ${try}"
-  done
-}
-
-# Power limiting to 400W is only needed if GPUs fall out of sync (causing 2.2x increased train time) at higher power, which has been observed at 450W
-sudo rocm-smi -d 0 1 2 3 4 5 6 7 --setpoweroverdrive 750 && \
-run_retry TOTAL_CKPTS=7 python3 examples/mlperf/model_train.py; (( $? == 2 )) && { echo "training failed before BEAM completion"; exit 2; }
-sleep 90
-
-run_retry EVAL_SAMPLES=600 python3 examples/mlperf/model_eval.py; (( $? == 2 )) && { echo "eval failed before BEAM completion"; exit 2; }
-# Checkpoints will be evaluated in reverse chronological order, even if above training crashed early
-# STOP_IF_CONVERGED=1: Stop the eval after the first time convergence is detected; no more checkpoints will be evaluated after that.
-STOP_IF_CONVERGED=1 python3 examples/mlperf/model_eval.py
@@ -1,4 +1,4 @@
-import os, sys, pickle, time, re
+import os, sys, pickle, time
 import numpy as np
 if "FLOAT16" not in os.environ: os.environ["FLOAT16"] = "1"
 if "IMAGE" not in os.environ: os.environ["IMAGE"] = "2"
@@ -10,7 +10,7 @@ from tinygrad.helpers import DEBUG, getenv
 from tinygrad.engine.realize import CompiledRunner

 import onnx
-from tinygrad.nn.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner

 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
 OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
@@ -52,8 +52,6 @@ def compile(onnx_file):
      kernel_count += 1
      read_image_count += ei.prg.p.src.count("read_image")
      gated_read_image_count += ei.prg.p.src.count("?read_image")
-      for v in [m.group(1) for m in re.finditer(r'(val\d+)\s*=\s*read_imagef\(', ei.prg.p.src)]:
-        if len(re.findall(fr'[\?\:]{v}\.[xyzw]', ei.prg.p.src)) > 0: gated_read_image_count += 1
  print(f"{kernel_count=},  {read_image_count=}, {gated_read_image_count=}")
  if (allowed_kernel_count:=getenv("ALLOWED_KERNEL_COUNT", -1)) != -1:
    assert kernel_count == allowed_kernel_count, f"different kernels! {kernel_count=}, {allowed_kernel_count=}"
@@ -79,20 +77,13 @@ def test_vs_compile(run, new_inputs, test_val=None):
            **{k:Tensor(v, device="NPY").realize() for k,v in new_inputs_numpy.items() if 'img' not in k}}

  # run 20 times
-  step_times = []
  for _ in range(20):
    st = time.perf_counter()
    out = run(**inputs)
    mt = time.perf_counter()
    val = out.numpy()
    et = time.perf_counter()
-    step_times.append((et-st)*1e3)
-    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")
-
-  if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
-    min_time = min(step_times)
-    assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
-
+    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {(et-st)*1e3:6.2f} ms")
  print(out, val.shape, val.dtype)
  if test_val is not None: np.testing.assert_equal(test_val, val)
  print("**** test done ****")
@@ -1,12 +1,12 @@
 import sys
 from tinygrad import Tensor, fetch, GlobalCounters, dtypes
 from tinygrad.uop.ops import UOp
-from tinygrad.nn.onnx import OnnxRunner
-from tinygrad.schedule.rangeify import get_rangeify_map
+from tinygrad.frontend.onnx import OnnxRunner
+from tinygrad.schedule.kernelize import get_kernelize_map
 from tinygrad.engine.schedule import create_schedule_with_vars
 from tinygrad.engine.realize import run_schedule

-# NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 VIZ=1 DEBUG=2 python3 examples/openpilot/compile4.py
+# NOLOCALS=1 GPU=1 IMAGE=2 FLOAT16=1 VIZ=1 DEBUG=2 python3 examples/openpilot/compile4.py

 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
 OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
@@ -33,7 +33,7 @@ if __name__ == "__main__":
        if not in_target_path[s]:
          independent_set[s] = None
  independent = UOp.sink(*independent_set.keys())
-  kernelized = get_rangeify_map(independent)
+  kernelized = get_kernelize_map(independent)
  independent = independent.substitute(kernelized)
  schedule, var_vals = create_schedule_with_vars(independent)
  run_schedule(schedule)
@@ -27,7 +27,7 @@ class Model(nn.Module):

 if __name__ == "__main__":
  if getenv("TINY_BACKEND"):
-    import tinygrad.nn.torch  # noqa: F401
+    import tinygrad.frontend.torch  # noqa: F401
    device = torch.device("tiny")
  else:
    device = torch.device({"METAL":"mps","NV":"cuda"}.get(Device.DEFAULT, "cpu"))
@@ -8,7 +8,7 @@ from typing import Dict, Union

 from extra.models.llama import Transformer, convert_from_huggingface, fix_bf16
 from examples.llama3 import load
-from tinygrad import nn, Tensor, Device
+from tinygrad import nn, Tensor
 from tinygrad.helpers import fetch, colored, GlobalCounters, Timing, DEBUG
 from tinygrad.nn.state import load_state_dict, get_parameters

@@ -80,7 +80,7 @@ if __name__ == "__main__":
    st = GlobalCounters.time_sum_s
    next_tok = Tensor([toks[start_pos:]]) if tok_tensor is None or (len(toks)-start_pos) > 1 else tok_tensor.reshape(1, 1)
    with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/s, {GlobalCounters.global_mem/x:.2f} GB/s, param {param_bytes/x:.2f} GB/s"):
-      with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "") +
+      with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "") +
                  f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB" +
                  (f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=args.timing):
        tok_tensor = transformer(next_tok, start_pos, args.temperature)
@@ -6,7 +6,7 @@
 from tinygrad import Tensor, TinyJit, dtypes, GlobalCounters
 from tinygrad.nn import Conv2d, GroupNorm
 from tinygrad.nn.state import safe_load, load_state_dict
-from tinygrad.helpers import fetch, trange, colored, Timing, getenv
+from tinygrad.helpers import fetch, trange, colored, Timing
 from extra.models.clip import Embedder, FrozenClosedClipEmbedder, FrozenOpenClipEmbedder
 from extra.models.unet import UNetModel, Upsample, Downsample, timestep_embedding
 from extra.bench_log import BenchEvent, WallTimeEvent
@@ -14,7 +14,7 @@ from examples.stable_diffusion import ResnetBlock, Mid
 import numpy as np

 from typing import Dict, List, Callable, Optional, Any, Set, Tuple, Union, Type
-import argparse, tempfile, time
+import argparse, tempfile
 from abc import ABC, abstractmethod
 from pathlib import Path
 from PIL import Image
@@ -342,13 +342,11 @@ class DPMPP2MSampler:
    sigmas = self.discretization(num_steps).to(x.device)
    x *= Tensor.sqrt(1.0 + sigmas[0] ** 2.0)
    num_sigmas = len(sigmas)
-    step_times = []

    old_denoised = None
    for i in trange(num_sigmas - 1):
      with Timing("step in ", enabled=timing, on_exit=lambda _: f", using {GlobalCounters.mem_used/1e9:.2f} GB"):
        GlobalCounters.reset()
-        st = time.perf_counter_ns()
        with WallTimeEvent(BenchEvent.STEP):
          x, old_denoised = self.sampler_step(
            old_denoised=old_denoised,
@@ -360,13 +358,8 @@ class DPMPP2MSampler:
            c=c,
            uc=uc,
          )
-          step_times.append(t:=(time.perf_counter_ns() - st)*1e-6)
          x.realize(old_denoised)

-    if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
-      min_time = min(step_times)
-      assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
-
    return x


@@ -437,8 +430,8 @@ if __name__ == "__main__":
    im.show()

  # validation!
-  is_default = args.prompt == default_prompt and args.steps == 10 and args.seed == 0 and args.guidance == 6.0 and args.width == args.height == 1024
-  if is_default and not args.weights and not args.fakeweights:
+  if args.prompt == default_prompt and args.steps == 10 and args.seed == 0 and args.guidance == 6.0 and args.width == args.height == 1024 \
+    and not args.weights:
    ref_image = Tensor(np.array(Image.open(Path(__file__).parent / "sdxl_seed0.png")))
    distance = (((x.cast(dtypes.float) - ref_image.cast(dtypes.float)) / ref_image.max())**2).mean().item()
    assert distance < 4e-3, colored(f"validation failed with {distance=}", "red")
@@ -2,20 +2,18 @@
 # https://github.com/ekagra-ranjan/huggingface-blog/blob/main/stable_diffusion.md
 import tempfile
 from pathlib import Path
-import argparse, time
+import argparse
 from collections import namedtuple
 from typing import Dict, Any

 from PIL import Image
 import numpy as np
 from tinygrad import Device, GlobalCounters, dtypes, Tensor, TinyJit
-from tinygrad.helpers import Timing, Context, getenv, fetch, colored, tqdm, flatten
+from tinygrad.helpers import Timing, Context, getenv, fetch, colored, tqdm
 from tinygrad.nn import Conv2d, GroupNorm
 from tinygrad.nn.state import torch_load, load_state_dict, get_state_dict
-from extra.models.clip import Closed, Tokenizer, FrozenOpenClipEmbedder
-from extra.models import unet, clip
+from extra.models.clip import Closed, Tokenizer
 from extra.models.unet import UNetModel
-from examples.mlperf.initializers import AutocastLinear, AutocastConv2d, AutocastGroupNorm, AutocastLayerNorm, zero_module, attn_f32_softmax, gelu_erf
 from extra.bench_log import BenchEvent, WallTimeEvent

 class AttnBlock:
@@ -156,46 +154,12 @@ unet_params: Dict[str,Any] = {
  "use_linear": False,
 }

-mlperf_params: Dict[str,Any] = {"adm_in_ch": None, "in_ch": 4, "out_ch": 4, "model_ch": 320, "attention_resolutions": [4, 2, 1], "num_res_blocks": 2,
-                                "channel_mult": [1, 2, 4, 4], "d_head": 64, "transformer_depth": [1, 1, 1, 1], "ctx_dim": 1024, "use_linear": True,
-                                "num_groups":16, "st_norm_eps":1e-6}
-
 class StableDiffusion:
-  def __init__(self, version:str|None=None, pretrained:str|None=None):
+  def __init__(self):
    self.alphas_cumprod = get_alphas_cumprod()
-    if version != "v2-mlperf-train":
-      self.first_stage_model = AutoencoderKL() # only needed for decoding generated latents to images; not needed in mlperf training from preprocessed moments
-
-    if not version:
-      self.cond_stage_model = namedtuple("CondStageModel", ["transformer"])(transformer = namedtuple("Transformer", ["text_model"])(text_model = Closed.ClipTextTransformer()))
-      unet_init_params = unet_params
-    elif version in {"v2-mlperf-train", "v2-mlperf-eval"}:
-      unet_init_params = mlperf_params
-      clip.gelu = gelu_erf
-      self.cond_stage_model = FrozenOpenClipEmbedder(**{"dims": 1024, "n_heads": 16, "layers": 24, "return_pooled": False, "ln_penultimate": True,
-                                                        "clip_tokenizer_version": "sd_mlperf_v5_0"})
-      unet.Linear, unet.Conv2d, unet.GroupNorm, unet.LayerNorm = AutocastLinear, AutocastConv2d, AutocastGroupNorm, AutocastLayerNorm
-      unet.attention, unet.gelu, unet.mixed_precision_dtype = attn_f32_softmax, gelu_erf, dtypes.bfloat16
-      if pretrained:
-        print("loading text encoder")
-        weights: dict[str,Tensor] = {k.replace("cond_stage_model.", "", 1):v for k,v in torch_load(pretrained)["state_dict"].items() if k.startswith("cond_stage_model.")}
-        weights["model.attn_mask"] = Tensor.full((77, 77), fill_value=float("-inf")).triu(1)
-        load_state_dict(self.cond_stage_model, weights)
-        # only the eval model needs the decoder
-        if version == "v2-mlperf-eval":
-          print("loading image latent encoder")
-          weights = {k.replace("first_stage_model.", "", 1):v for k,v in torch_load(pretrained)["state_dict"].items() if k.startswith("first_stage_model.")}
-          load_state_dict(self.first_stage_model, weights)
-
-    self.model = namedtuple("DiffusionModel", ["diffusion_model"])(diffusion_model = UNetModel(**unet_init_params))
-    if version == "v2-mlperf-train":
-      # the mlperf reference inits certain weights as zeroes
-      for bb in flatten(self.model.diffusion_model.input_blocks) + self.model.diffusion_model.middle_block + flatten(self.model.diffusion_model.output_blocks):
-        if isinstance(bb, unet.ResBlock):
-          zero_module(bb.out_layers[3])
-        elif isinstance(bb, unet.SpatialTransformer):
-          zero_module(bb.proj_out)
-      zero_module(self.model.diffusion_model.out[2])
+    self.model = namedtuple("DiffusionModel", ["diffusion_model"])(diffusion_model = UNetModel(**unet_params))
+    self.first_stage_model = AutoencoderKL()
+    self.cond_stage_model = namedtuple("CondStageModel", ["transformer"])(transformer = namedtuple("Transformer", ["text_model"])(text_model = Closed.ClipTextTransformer()))

  def get_x_prev_and_pred_x0(self, x, e_t, a_t, a_prev):
    temperature = 1
@@ -269,14 +233,12 @@ if __name__ == "__main__":

  # load in weights
  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
-    load_state_dict(model, torch_load(fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt'))['state_dict'], verbose=False, strict=False, realize=False)
+    load_state_dict(model, torch_load(fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt'))['state_dict'], strict=False)

    if args.fp16:
      for k,v in get_state_dict(model).items():
        if k.startswith("model"):
-          v.replace(v.cast(dtypes.float16))
-
-    Tensor.realize(*get_state_dict(model).values())
+          v.replace(v.cast(dtypes.float16).realize())

  # run through CLIP to get context
  tokenizer = Tokenizer.ClipTokenizer()
@@ -304,23 +266,17 @@ if __name__ == "__main__":
  def run(model, *x): return model(*x).realize()

  # this is diffusion
-  step_times = []
  with Context(BEAM=getenv("LATEBEAM")):
    for index, timestep in (t:=tqdm(list(enumerate(timesteps))[::-1])):
      GlobalCounters.reset()
-      st = time.perf_counter_ns()
      t.set_description("%3d %3d" % (index, timestep))
      with Timing("step in ", enabled=args.timing, on_exit=lambda _: f", using {GlobalCounters.mem_used/1e9:.2f} GB"):
        with WallTimeEvent(BenchEvent.STEP):
          tid = Tensor([index])
          latent = run(model, unconditional_context, context, latent, Tensor([timestep]), alphas[tid], alphas_prev[tid], Tensor([args.guidance]))
          if args.timing: Device[Device.DEFAULT].synchronize()
-      step_times.append((time.perf_counter_ns() - st)*1e-6)
    del run

-  if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
-    min_time = min(step_times)
-    assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
  # upsample latent space to image with autoencoder
  x = model.decode(latent)
  print(x.shape)
@@ -32,7 +32,7 @@ if __name__ == "__main__":

  lr = 5e-3
  transform = ComposeTransforms([
-    lambda x: [Image.fromarray(xx).resize((64, 64)) for xx in x],
+    lambda x: [Image.fromarray(xx, mode='L').resize((64, 64)) for xx in x],
    lambda x: np.stack([np.asarray(xx) for xx in x], 0),
    lambda x: x / 255.0,
    lambda x: np.tile(np.expand_dims(x, 1), (1, 3, 1, 1)).astype(np.float32),
@@ -109,7 +109,7 @@ class TextDecoder:

  def forward(self, x:Tensor, pos:Union[Variable, Literal[0]], encoded_audio:Tensor):
    seqlen = x.shape[-1]
-    x = self.token_embedding(x) + self.positional_embedding.shrink(((pos, pos+seqlen), None))
+    x = self.token_embedding(x) + self.positional_embedding.shrink(((pos, pos+seqlen), None, None))
    for block in self.blocks: x = block(x, xa=encoded_audio, mask=self.mask, len=pos)
    return self.output_tok(x)

@@ -2,7 +2,7 @@
 import os
 from ultralytics import YOLO
 from pathlib import Path
-from tinygrad.nn.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner
 from extra.onnx_helpers import get_example_inputs

 os.chdir("/tmp")
@@ -37,7 +37,7 @@ def main():
  dev = PCIIface(None, 0)
  for x, y in dev.dev_impl.__dict__.items():
    if isinstance(y, AMRegister):
-      for inst, addr in y.addr.items(): reg_names[addr] = f"{x}, xcc={inst}"
+      for inst, addr in y.addr.keys(): reg_names[addr] = f"{x}, xcc={inst}"

  with open(sys.argv[1], 'r') as f:
    log_content = log_content_them = f.read()
@@ -1,7 +1,7 @@
 # copying the kernels from https://github.com/microsoft/ArchProbe into Python
 import numpy as np
 import pickle
-from tinygrad.runtime.ops_cl import CLProgram, CLBuffer
+from tinygrad.runtime.ops_gpu import CLProgram, CLBuffer
 from tinygrad import dtypes
 from tqdm import trange, tqdm
 from matplotlib import pyplot as plt
@@ -4,7 +4,7 @@ from tinygrad import dtypes
 from tinygrad.codegen.assembly import AssemblyCodegen, Register
 from tinygrad.codegen.opt.kernel import Ops
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
-from tinygrad.runtime.ops_cl import ROCM_LLVM_PATH
+from tinygrad.runtime.ops_gpu import ROCM_LLVM_PATH

 # ugh, is this really needed?
 from extra.helpers import enable_early_exec
@@ -5,7 +5,7 @@ from tinygrad.helpers import colored
 from extra.helpers import enable_early_exec
 early_exec = enable_early_exec()

-from tinygrad.runtime.ops_cl import CLProgram, CLBuffer, ROCM_LLVM_PATH
+from tinygrad.runtime.ops_gpu import CLProgram, CLBuffer, ROCM_LLVM_PATH

 ENABLE_NON_ASM = False

@@ -10,13 +10,13 @@ from tinygrad.renderer.cstyle import ClangRenderer
 render_dtype = ClangRenderer().render_dtype

 class ClangGraph(GraphRunner):
-  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[str, int]):
+  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]):
    super().__init__(jit_cache, input_rawbuffers, var_vals)
    if not all(isinstance(ji.prg, CompiledRunner) for ji in jit_cache): raise GraphException

    prgs = '\n'.join(dedup([cast(CompiledRunner, ji.prg).p.src for ji in jit_cache]))
    args = [f"{render_dtype(x.dtype)}* arg{i}" for i,x in enumerate(input_rawbuffers)]
-    args += sorted([f"int {v}" for v in var_vals])
+    args += sorted([f"int {v.expr}" for v in var_vals])
    code = ["void batched("+','.join(args)+") {"]
    for ji in jit_cache:
      args = []
@@ -34,6 +34,6 @@ class ClangGraph(GraphRunner):
    assert compiler is not None
    self._prg = ClangProgram("batched", compiler.compile(prgs+"\n"+"\n".join(code))) # no point in caching the pointers

-  def __call__(self, rawbufs: List[Buffer], var_vals: Dict[str, int], wait=False):
+  def __call__(self, rawbufs: List[Buffer], var_vals: Dict[Variable, int], wait=False):
    return cpu_time_execution(
-    lambda: self._prg(*[x._buf for x in rawbufs], *[x[1] for x in sorted(var_vals.items(), key=lambda x: x[0])]), enable=wait)
+    lambda: self._prg(*[x._buf for x in rawbufs], *[x[1] for x in sorted(var_vals.items(), key=lambda x: x[0].expr)]), enable=wait)
@@ -26,7 +26,7 @@ class VirtAQLQueue(AQLQueue):
    self.available_packet_slots -= 1

 class HSAGraph(MultiGraphRunner):
-  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[str, int]):
+  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]):
    super().__init__(jit_cache, input_rawbuffers, var_vals)

    # Check all jit items are compatible.
@@ -53,7 +53,7 @@ class HSAGraph(MultiGraphRunner):
      self.ji_kargs_structs[j] = ji.prg._prg.args_struct_t.from_address(kernargs_ptrs[ji.prg.dev])
      kernargs_ptrs[ji.prg.dev] += round_up(ctypes.sizeof(ji.prg._prg.args_struct_t), 16)
      for i in range(len(ji.bufs)): self.ji_kargs_structs[j].__setattr__(f'f{i}', cast(Buffer, ji.bufs[i])._buf)
-      for i in range(len(ji.prg.p.vars)): self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[ji.prg.p.vars[i].expr])
+      for i in range(len(ji.prg.p.vars)): self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[ji.prg.p.vars[i]])

    # Build queues.
    self.virt_aql_queues: Dict[Compiled, VirtAQLQueue] = {dev:VirtAQLQueue(dev, 2*len(self.jit_cache)+16) for dev in self.devices}
@@ -106,7 +106,7 @@ class HSAGraph(MultiGraphRunner):
    for sig in self.signals_to_reset: hsa.hsa_signal_silent_store_relaxed(sig, 0)
    hsa.hsa_signal_silent_store_relaxed(self.finish_signal, 0)

-  def __call__(self, input_rawbuffers: List[Buffer], var_vals: Dict[str, int], wait=False) -> Optional[float]:
+  def __call__(self, input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int], wait=False) -> Optional[float]:
    # Wait and restore signals
    hsa.hsa_signal_wait_scacquire(self.finish_signal, hsa.HSA_SIGNAL_CONDITION_LT, 1, (1 << 64) - 1, hsa.HSA_WAIT_STATE_ACTIVE)
    for sig in self.signals_to_reset: hsa.hsa_signal_silent_store_relaxed(sig, 1)
@@ -123,7 +123,7 @@ class HSAGraph(MultiGraphRunner):
    # Update var_vals
    for j in self.jc_idx_with_updatable_var_vals:
      for i,v in enumerate(cast(CompiledRunner, self.jit_cache[j].prg).p.vars):
-        self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[v.expr])
+        self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[v])

    # Update launch dims
    for j in self.jc_idx_with_updatable_launch_dims:
@@ -29,10 +29,10 @@ def uops_to_rdna(function_name:str, uops:UOpGraph) -> str:
  r: Dict[UOp, str] = {}
  for u in uops:
    if u.uop == UOps.SPECIAL:
-      if u.arg.startswith("lidx"):
-        r[u] = f'v{u.src[0].arg}'
-      elif u.arg.startswith("gidx"):
-        r[u] = f's{2+u.src[0].arg}'
+      if u.arg[1].startswith("lidx"):
+        r[u] = f'v{u.arg[0]}'
+      elif u.arg[1].startswith("gidx"):
+        r[u] = f's{2+u.arg[0]}'
      else:
        raise NotImplementedError
    elif u.uop == UOps.CONST:
@@ -10,7 +10,7 @@ from tinygrad.uop.ops import Ops
 import json
 from collections import OrderedDict

-EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CPU", "CUDA", "CL"]
+EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CPU", "CUDA", "GPU"]

 def compile_net(run:TinyJit, special_names:Dict[int,str]) -> Tuple[Dict[str,str],List[Tuple[str,List[str],List[int]]],Dict[str,Tuple[int,DType,int]],Dict[str,Tensor]]:
  functions, bufs, bufs_to_save, statements, bufnum = {}, {}, {}, [], 0
@@ -67,12 +67,11 @@ def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,in
  forward_args = ",".join(f"{dtype}{'*' if name not in symbolic_vars.values() else ''} {name}" for name,dtype,_ in (outputs+inputs if wasm else inputs+outputs))

  if not wasm:
-    thread_id = 0 # NOTE: export does not support threading, thread_id is always 0
    for name,cl in bufs_to_save.items():
      weight = ''.join(["\\x%02X"%x for x in bytes(to_mv(cl._buf.va_addr, cl._buf.size))])
      cprog.append(f"unsigned char {name}_data[] = \"{weight}\";")
    cprog += [f"{dtype_map[dtype]} {name}[{len}];" if name not in bufs_to_save else f"{dtype_map[dtype]} *{name} = ({dtype_map[dtype]} *){name}_data;" for name,(len,dtype,_key) in bufs.items() if name not in input_names+output_names]
-    cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)}, {thread_id});" for (name, args, _global_size, _local_size) in statements] + ["}"]
+    cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)});" for (name, args, _global_size, _local_size) in statements] + ["}"]
    return '\n'.join(headers + cprog)
  else:
    if bufs_to_save:
@@ -240,9 +239,7 @@ export default {model_name};

 def export_model(model, target:str, *inputs, model_name: Optional[str] = "model", stream_weights=False):
  assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, f"only {', '.join(EXPORT_SUPPORTED_DEVICE)} are supported"
-
-  # NOTE: CPU_COUNT=1, since export does not support threading
-  with Context(JIT=2, CPU_COUNT=1): run,special_names = jit_model(model, *inputs)
+  with Context(JIT=2): run,special_names = jit_model(model, *inputs)
  functions, statements, bufs, bufs_to_save = compile_net(run, special_names)
  state = get_state_dict(model)
  weight_names = {id(x.uop.base.realized): name for name, x in state.items()}
@@ -5,7 +5,7 @@ from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import getenv, colored, prod, unwrap
 from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad.shape.view import strides_for_shape
-from tinygrad.codegen.opt.kernel import axis_colors, Opt, OptOps
+from tinygrad.codegen.opt.kernel import axis_colors
 from tinygrad.codegen.opt.swizzler import merge_views, view_left

 def to_colored(full_shape, axis_types): return '_'.join([colored(str(s), axis_colors[at]) for s,at in zip(full_shape, axis_types)])
@@ -44,27 +44,13 @@ pm = PatternMatcher([
  (UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), swizzle_reduceop),
 ])

-def rangeify_kernel3():
-  a = Tensor.empty(N,N)
-  b = Tensor.empty(N,N)
-  c = a@b
-  #c = c.reshape((32,2,16,4,32,2,16,4)).contiguous()
-  sink = c.schedule()[-1].ast
-  #print(sink)
-
-  opts  = [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.UPCAST, 0, 2)]
-  opts += [Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 1, 16), Opt(OptOps.UPCAST, 1, 2)]
-  opts += [Opt(OptOps.UNROLL, 0, 8)]
-
-  return sink.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
-
 def top_spec_kernel3():
  a = Tensor.empty(N,N)
  b = Tensor.empty(N,N)
  c = a@b
  sink = c.schedule()[-1].ast
  L = 16
-  sink = sink.reshape((N//L, L, N//L, L)) #.lift({0:UOp.range(N//BM, 0), 2:UOp.range(N//BN, 1)})
+  sink = sink.reshape((N//L, L, N//L, L)) #.lift({0:UOp.range(dtypes.int, N//BM, 0), 2:UOp.range(dtypes.int, N//BN, 1)})
  sink = graph_rewrite(sink, view_left+pm)
  axis_types = (AxisType.GLOBAL, AxisType.LOCAL, AxisType.GLOBAL, AxisType.LOCAL, AxisType.REDUCE)
  return sink.replace(arg=KernelInfo(name="top_"+to_colored(sink.full_shape, axis_types), axis_types=axis_types))
@@ -185,7 +171,7 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):

  c_regs = UOp(Ops.DEFINE_REG, dtypes.float.ptr(TM * nbIterWaveM * TN * nbIterWaveN), arg=2)

-  i = UOp.range(c_regs.dtype.size, 16)
+  i = UOp.range(dtypes.int, c_regs.dtype.size, 16)
  init_store = c_regs[i].store(UOp.const(dtypes.float, 0.0), i)

  if kernel4:
@@ -196,53 +182,53 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
    kId = 0

    # load from globals into locals
-    i = UOp.range(nbReadsB, 0)
+    i = UOp.range(dtypes.int, nbReadsB, 0)
    index_x = BN * blockIdx_x + rBIdx
    index_y = rBIdy + i * strideReadB + kId
    Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(b[N * index_y + index_x].load(), i)

-    i = UOp.range(nbReadsA, 1)
+    i = UOp.range(dtypes.int, nbReadsA, 1)
    index_x = rAIdx + kId
    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
    As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(a[N * index_y + index_x].load(), i)

    # iterate over the middle chunk
-    kId_range = UOp.range(N//BK-1, 2)
+    kId_range = UOp.range(dtypes.int, N//BK-1, 2)
    kId = kId_range*BK

    barrier = UOp.barrier(As_store, Bs_store)

    # load from globals into registers (next round)
-    i = UOp.range(nbReadsB, 3)
+    i = UOp.range(dtypes.int, nbReadsB, 3)
    index_x = BN * blockIdx_x + rBIdx
    index_y = rBIdy + i * strideReadB + kId + BK
    regB_store = regB[i].store(b[N * index_y + index_x].load(), i)

-    i = UOp.range(nbReadsA, 4)
+    i = UOp.range(dtypes.int, nbReadsA, 4)
    index_x = rAIdx + kId + BK
    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
    regA_store = regA[i].store(a[N * index_y + index_x].load(), i)

    def inner_loop(first_range, inp_dep=()):
      # inner unroll
-      k = UOp.range(BK, first_range+0)
+      k = UOp.range(dtypes.int, BK, first_range+0)

      # load from locals into registers
-      iterWave = UOp.range(nbIterWaveN, first_range+1)
-      i = UOp.range(TN, first_range+2)
+      iterWave = UOp.range(dtypes.int, nbIterWaveN, first_range+1)
+      i = UOp.range(dtypes.int, TN, first_range+2)
      index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
      B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(*inp_dep), iterWave, i)

-      iterWave = UOp.range(nbIterWaveM, first_range+3)
-      i = UOp.range(TM, first_range+4)
+      iterWave = UOp.range(dtypes.int, nbIterWaveM, first_range+3)
+      i = UOp.range(dtypes.int, TM, first_range+4)
      index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
      A_col_store = A_col[iterWave*TM + i].store(As[k*BM_As_stride + index].load(*inp_dep), iterWave, i)

      # do the GEMM math
-      iterWaveM = UOp.range(nbIterWaveM, first_range+5)
-      yt = UOp.range(TM, first_range+6)
-      iterWaveN = UOp.range(nbIterWaveN, first_range+7)
-      xt = UOp.range(TN, first_range+8)
+      iterWaveM = UOp.range(dtypes.int, nbIterWaveM, first_range+5)
+      yt = UOp.range(dtypes.int, TM, first_range+6)
+      iterWaveN = UOp.range(dtypes.int, nbIterWaveN, first_range+7)
+      xt = UOp.range(dtypes.int, TN, first_range+8)
      x = iterWaveN * TN + xt
      y = iterWaveM * TM + yt
      c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
@@ -255,12 +241,12 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
    sink = inner_loop(5, (barrier, regB_store, regA_store)).barrier()

    # load from registers into locals
-    i = UOp.range(nbReadsB, 14)
+    i = UOp.range(dtypes.int, nbReadsB, 14)
    index_x = BN * blockIdx_x + rBIdx
    index_y = rBIdy + i * strideReadB + kId + BK
    Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(regB[i].load(sink), i, kId_range)

-    i = UOp.range(nbReadsA, 15)
+    i = UOp.range(dtypes.int, nbReadsA, 15)
    index_x = rAIdx + kId + BK
    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
    As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(regA[i].load(sink), i, kId_range)
@@ -268,40 +254,40 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
    # final iteration without the copy
    sink = inner_loop(16, (UOp.barrier(Bs_store, As_store),))
  else:
-    kId_range = UOp.range(N//BK, 0)
+    kId_range = UOp.range(dtypes.int, N//BK, 0)
    kId = kId_range*BK

    # load from globals into locals
-    i = UOp.range(nbReadsB, 1)
+    i = UOp.range(dtypes.int, nbReadsB, 1)
    index_x = BN * blockIdx_x + rBIdx
    index_y = rBIdy + i * strideReadB + kId
    Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(b[N * index_y + index_x].load(), i)

-    i = UOp.range(nbReadsA, 2)
+    i = UOp.range(dtypes.int, nbReadsA, 2)
    index_x = rAIdx + kId
    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
    As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(a[N * index_y + index_x].load(), i)

    barrier = UOp.barrier(As_store, Bs_store)

-    k = UOp.range(BK, 3)
+    k = UOp.range(dtypes.int, BK, 3)

    # load from locals into registers
-    iterWave = UOp.range(nbIterWaveN, 4)
-    i = UOp.range(TN, 5)
+    iterWave = UOp.range(dtypes.int, nbIterWaveN, 4)
+    i = UOp.range(dtypes.int, TN, 5)
    index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
    B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(barrier), iterWave, i)

-    iterWave = UOp.range(nbIterWaveM, 6)
-    i = UOp.range(TM, 7)
+    iterWave = UOp.range(dtypes.int, nbIterWaveM, 6)
+    i = UOp.range(dtypes.int, TM, 7)
    index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
    A_col_store = A_col[iterWave*TM + i].store(As[k*BM_As_stride + index].load(barrier), iterWave, i)

    # do the GEMM math
-    iterWaveM = UOp.range(nbIterWaveM, 8)
-    yt = UOp.range(TM, 9)
-    iterWaveN = UOp.range(nbIterWaveN, 10)
-    xt = UOp.range(TN, 12)
+    iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 8)
+    yt = UOp.range(dtypes.int, TM, 9)
+    iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 10)
+    xt = UOp.range(dtypes.int, TN, 12)
    x = iterWaveN * TN + xt
    y = iterWaveM * TM + yt
    c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
@@ -309,10 +295,10 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
                            iterWaveM, iterWaveN, yt, xt, k, kId_range)

  # store c_regs into c
-  iterWaveM = UOp.range(nbIterWaveM, 1000)
-  yt = UOp.range(TM, 1001)
-  iterWaveN = UOp.range(nbIterWaveN, 1002)
-  xt = UOp.range(TN, 1003)
+  iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 1000)
+  yt = UOp.range(dtypes.int, TM, 1001)
+  iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 1002)
+  xt = UOp.range(dtypes.int, TN, 1003)
  xOut = blockIdx_x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave
  yOut = blockIdx_y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave
  indexC = N * (yOut + yt) + xOut + xt
@@ -323,15 +309,10 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):

 if __name__ == "__main__":
  HL = getenv("HL")
-  if HL == 3: hprg = rangeify_kernel3()
-  elif HL == 2: hprg = top_spec_kernel3()
+  if HL == 2: hprg = top_spec_kernel3()
  elif HL == 1: hprg = hl_spec_kernel3()
  else: hprg = hand_spec_kernel3()
-  if HL == 3:
-    with Context(BLOCK_REORDER=0):
-      prg = get_program(hprg, Device.default.renderer)
-  else:
-    prg = get_program(hprg, Device.default.renderer)
+  prg = get_program(hprg, Device.default.renderer)
  print(prg.src)
  if getenv("SRC"): exit(0)
  hrunner = CompiledRunner(prg)
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import numpy as np
-from tinygrad.runtime.ops_cl import CLProgram, CLCompiler
+from tinygrad.runtime.ops_gpu import CLProgram, CLCompiler
 from tinygrad import Device, dtypes
 from tinygrad.device import Buffer
 from hexdump import hexdump
@@ -11,7 +11,7 @@ from hexdump import hexdump
 # https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_split_matrix_multiply_accumulate.html
 # https://hc34.hotchips.org/assets/program/conference/day1/GPU%20HPC/Intel_s%20Ponte%20Vecchio%20GPU%20-%20Architecture%20Systems%20and%20Software%20FINAL.pdf

-device = Device["CL"]
+device = Device["GPU"]

 # NOTE: only the subgroup type 8 ones work
 prog = CLProgram(device, "test", CLCompiler(device, "test").compile(f"""
@@ -26,9 +26,9 @@ __kernel void test(__global float* data0, const __global int* data1, const __glo
 """))
 #with open("/tmp/test.elf", "wb") as f: f.write(prog.lib)

-a = Buffer("CL", 8, dtypes.float32).allocate()
-b = Buffer("CL", 0x10, dtypes.float16).allocate()
-c = Buffer("CL", 8*0x10, dtypes.float16).allocate()
+a = Buffer("GPU", 8, dtypes.float32).allocate()
+b = Buffer("GPU", 0x10, dtypes.float16).allocate()
+c = Buffer("GPU", 8*0x10, dtypes.float16).allocate()

 row = np.array([1,2,3,4,5,6,7,8,1,2,3,4,5,6,7,8], np.float16)
 mat = np.random.random((8, 0x10)).astype(np.float16)
@@ -56,7 +56,7 @@ def randoms():
 def ast_to_cuda_prog(compiler, ast, opts):
  k = Kernel(ast)
  k.apply_opts(opts)
-  p = get_program(k.ast, k.opts, k.applied_opts)
+  p = get_program(k.get_optimized_ast(), k.opts)
  return CUDAProgram(device, p.function_name, compiler.compile(p.src))

 if __name__ == "__main__":
@@ -75,7 +75,7 @@ if __name__ == "__main__":

    if GEMM_VARIATION == "max" and (M%64)==0 and (N%128)==0 and (K%64)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.float and DTYPE_ACC == dtypes.float:
      print("Using CUDA and triton-generated kernel")
-      # See nv_triton_gemm.annotated.ptx for PTX code which was generated from `PYTHONPATH=. DEBUG=6 CUDA=1 CUDA_PTX=1 python3 extra/gemm/triton_nv_matmul.py`
+      # See nv_triton_gemm.annotated.ptx for PTX code which was generated from `PYTHONPATH=. DEBUG=6 CUDA=1 PTX=1 python3 extra/gemm/triton_nv_matmul.py`
      # this kernel with M=N=K=4096 does 162TFLOPS, vs torch at 144TFLOPS and BEAM=8 tinygrad at 138TFLOPS.  theo max is 165TFLOPS.

      # WMMA element size is (M, N, K) = (16, 8, 16)
@@ -2,7 +2,7 @@ import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.helpers import getenv, get_single_element
 from tinygrad.dtype import _to_np_dtype
-from tinygrad.codegen.opt import OptOps
+from tinygrad.codegen.opt.kernel import OptOps
 from tinygrad.engine.realize import lower_schedule

 dtype_in = dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else dtypes.float
@@ -29,7 +29,7 @@ if __name__ == "__main__":
      Opt(op=OptOps.LOCAL, axis=0, amt=2),
    ]
  k.apply_opts(opts)
-  prg = get_program(k.ast, k.opts, k.applied_opts)
+  prg = get_program(k.get_optimized_ast(), k.opts)
  new_src = prg.src
  # can mod source here
  prg = replace(prg, src=new_src)
@@ -43,7 +43,7 @@ def matmul_kernel(c_ptr, a_ptr, b_ptr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N:
  c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
  tl.store(c_ptrs, c)

-# CUDA=1 CUDA_PTX=1 python3 extra/gemm/triton_nv_matmul.py
+# CUDA=1 PTX=1 python3 extra/gemm/triton_nv_matmul.py
 if __name__ == "__main__":
  BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 64, 128, 64
  M, N, K = 4096, 4096, 4096
@@ -7,6 +7,7 @@ bert_train_params = {
  "GPUS": 6,
  "BS": 96,
  "EVAL_BS": 96,
+  "FUSE_ARANGE": 1,
  "BASEDIR": "/raid/datasets/wiki",
 }

@@ -50,7 +50,7 @@ def ioctls_from_header():
  hdr = (pathlib.Path(__file__).parent / "kfd_ioctl.h").read_text().replace("\\\n", "")
  pattern = r'#define\s+(AMDKFD_IOC_[A-Z0-9_]+)\s+AMDKFD_IOW?R?\((0x[0-9a-fA-F]+),\s+struct\s([A-Za-z0-9_]+)\)'
  matches = re.findall(pattern, hdr, re.MULTILINE)
-  return {int(nr, 0x10):(name, getattr(kfd_ioctl, "struct_"+sname, None)) for name, nr, sname in matches}
+  return {int(nr, 0x10):(name, getattr(kfd_ioctl, "struct_"+sname)) for name, nr, sname in matches}
 nrs = ioctls_from_header()

@ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_int, ctypes.c_ulong, ctypes.c_void_p)
@@ -1,7 +1,7 @@
 import onnx, yaml, tempfile, time, argparse, json
 from pathlib import Path
 from typing import Any
-from tinygrad.nn.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner
 from extra.onnx_helpers import validate, get_example_inputs
 from extra.huggingface_onnx.huggingface_manager import DOWNLOADS_DIR, snapshot_download_with_retry

@@ -88,7 +88,7 @@ def mcts_search(lin:Kernel, rawbufs:List[Buffer], amt:int) -> Kernel:
    return ret

  rawbufs = _ensure_buffer_alloc(rawbufs)
-  var_vals = {k.expr:(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
+  var_vals = {k:(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
  dev = Device[lin.opts.device]
  root = MCTSNode(lin)

@@ -9,9 +9,6 @@ from PIL import Image
 import numpy as np
 import re, gzip

-# Allow for monkeypatching for mlperf.
-gelu = Tensor.gelu
-
@lru_cache()
 def default_bpe():
  # Clip tokenizer, taken from https://github.com/openai/CLIP/blob/main/clip/simple_tokenizer.py (MIT license)
@@ -56,8 +53,8 @@ class Tokenizer:
    cs = [chr(n) for n in cs]
    return dict(zip(bs, cs))
  class ClipTokenizer:
-    def __init__(self, version=None):
-      self.byte_encoder, self.version = Tokenizer.bytes_to_unicode(), version
+    def __init__(self):
+      self.byte_encoder = Tokenizer.bytes_to_unicode()
      merges = gzip.open(default_bpe()).read().decode("utf-8").split('\n')
      merges = merges[1:49152-256-2+1]
      merges = [tuple(merge.split()) for merge in merges]
@@ -65,17 +62,11 @@ class Tokenizer:
      vocab = vocab + [v+'</w>' for v in vocab]
      for merge in merges:
        vocab.append(''.join(merge))
-      if self.version == "sd_mlperf_v5_0":
-        import regex
-        vocab.extend(['<start_of_text>', '<end_of_text>'])
-        self.cache = {'<start_of_text>': '<start_of_text>', '<end_of_text>': '<end_of_text>'}
-        self.pat = regex.compile(r"""<start_of_text>|<end_of_text>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", regex.IGNORECASE)
-      else:
-        vocab.extend(['<|startoftext|>', '<|endoftext|>'])
-        self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
-        self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[^\s]+""", re.IGNORECASE)
+      vocab.extend(['<|startoftext|>', '<|endoftext|>'])
      self.encoder = dict(zip(vocab, range(len(vocab))))
      self.bpe_ranks = dict(zip(merges, range(len(merges))))
+      self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
+      self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[^\s]+""", re.IGNORECASE)

    def bpe(self, token):
      if token in self.cache:
@@ -119,17 +110,8 @@ class Tokenizer:

    def encode(self, text:str, pad_with_zeros:bool=False) -> List[int]:
      bpe_tokens: List[int] = []
-      if self.version == "sd_mlperf_v5_0":
-        import regex, ftfy, html
-        text = ftfy.fix_text(text)
-        text = html.unescape(html.unescape(text)).strip()
-        text = Tokenizer.whitespace_clean(text).lower()
-        re_module = regex
-      else:
-        text = Tokenizer.whitespace_clean(text.strip()).lower()
-        re_module = re
-
-      for token in re_module.findall(self.pat, text):
+      text = Tokenizer.whitespace_clean(text.strip()).lower()
+      for token in re.findall(self.pat, text):
        token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
        bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
      # Truncation, keeping two slots for start and end tokens.
@@ -270,8 +252,10 @@ class Open:
      q,k,v = [y.reshape(T, B*self.n_heads, self.d_head).transpose(0, 1).reshape(B, self.n_heads, T, self.d_head) for y in proj.chunk(3)]

      attn_output = Tensor.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
-      attn_output = attn_output.permute(2, 0, 1, 3).reshape(T, B, C)
+      attn_output = attn_output.permute(2, 0, 1, 3).reshape(T*B, C)
+
      attn_output = self.out_proj(attn_output)
+      attn_output = attn_output.reshape(T, B, C)

      return attn_output

@@ -279,10 +263,9 @@ class Open:
    def __init__(self, dims, hidden_dims):
      self.c_fc   = Linear(dims, hidden_dims)
      self.c_proj = Linear(hidden_dims, dims)
-      self.gelu = gelu

    def __call__(self, x:Tensor) -> Tensor:
-      return x.sequential([self.c_fc, self.gelu, self.c_proj])
+      return x.sequential([self.c_fc, Tensor.gelu, self.c_proj])

  # https://github.com/mlfoundations/open_clip/blob/58e4e39aaabc6040839b0d2a7e8bf20979e4558a/src/open_clip/transformer.py#L210
  class ResidualAttentionBlock:
@@ -367,15 +350,15 @@ class Open:
 # https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/modules/encoders/modules.py#L396
 # https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/modules/encoders/modules.py#L498
 class FrozenOpenClipEmbedder(Embedder):
-  def __init__(self, dims:int, n_heads:int, layers:int, return_pooled:bool, ln_penultimate:bool=False, clip_tokenizer_version=None):
-    self.tokenizer = Tokenizer.ClipTokenizer(version=clip_tokenizer_version)
+  def __init__(self, dims:int, n_heads:int, layers:int, return_pooled:bool, ln_penultimate:bool=False):
+    self.tokenizer = Tokenizer.ClipTokenizer()
    self.model = Open.ClipTextTransformer(dims, n_heads, layers)
    self.return_pooled = return_pooled
    self.input_key = "txt"
    self.ln_penultimate = ln_penultimate

  def tokenize(self, text:str, device:Optional[str]=None) -> Tensor:
-    return Tensor(self.tokenizer.encode(text, pad_with_zeros=True), dtype=dtypes.int32, device=device).reshape(1,-1)
+    return Tensor(self.tokenizer.encode(text, pad_with_zeros=True), dtype=dtypes.int64, device=device).reshape(1,-1)

  def text_transformer_forward(self, x:Tensor, attn_mask:Optional[Tensor]=None):
    for r in self.model.transformer.resblocks:
@@ -466,7 +449,7 @@ class OpenClipEncoder:
    x = x + self.positional_embedding
    x = self.transformer(x, attn_mask=self.attn_mask)
    x = self.ln_final(x)
-    x = x[Tensor.arange(x.shape[0], device=x.device), tokens.argmax(axis=-1)]
+    x = x[:, tokens.argmax(axis=-1)]
    x = x @ self.text_projection
    return x

@@ -270,10 +270,8 @@ class FidInceptionV3:
    self.Mixed_7b = inception.Mixed_7b
    self.Mixed_7c = inception.Mixed_7c

-  def load_from_pretrained(self, path=None):
-    if path is None:
-      path = fetch("https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth", "pt_inception-2015-12-05-6726825d.pth")
-    state_dict = torch_load(str(path))
+  def load_from_pretrained(self):
+    state_dict = torch_load(str(fetch("https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth", "pt_inception-2015-12-05-6726825d.pth")))
    for k,v in state_dict.items():
      if k.endswith(".num_batches_tracked"):
        state_dict[k] = v.reshape(1)
@@ -249,5 +249,8 @@ def convert_from_gguf(weights:dict[str, Tensor], n_layers:int):
  return sd

 def fix_bf16(weights:dict[Any, Tensor]):
-  # TODO: without casting to float16, 70B llama OOM on tinybox.
-  return {k:v.cast(dtypes.float32).cast(dtypes.float16) if v.dtype == dtypes.bfloat16 else v for k,v in weights.items()}
+  if getenv("SUPPORT_BF16", 1):
+    # TODO: without casting to float16, 70B llama OOM on tinybox.
+    return {k:v.cast(dtypes.float32).cast(dtypes.float16) if v.dtype == dtypes.bfloat16 else v for k,v in weights.items()}
+  # TODO: check if device supports bf16
+  return {k:v.llvm_bf16_cast(dtypes.half).to(v.device) if v.dtype == dtypes.bfloat16 else v for k,v in weights.items()}
@@ -1,24 +1,21 @@
-from tinygrad import Tensor, dtypes, nn
+from tinygrad import Tensor, dtypes
+from tinygrad.nn import Linear, Conv2d, GroupNorm, LayerNorm
 from tinygrad.device import is_dtype_supported
-from typing import Optional, Union, List, Any, Tuple, Callable
+from typing import Optional, Union, List, Any, Tuple
 import math

-# allow for monkeypatching
-Linear, Conv2d, GroupNorm, LayerNorm = nn.Linear, nn.Conv2d, nn.GroupNorm, nn.LayerNorm
-attention, gelu, mixed_precision_dtype = Tensor.scaled_dot_product_attention, Tensor.gelu, dtypes.float16
-
 # https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/modules/diffusionmodules/util.py#L207
 def timestep_embedding(timesteps:Tensor, dim:int, max_period=10000):
  half = dim // 2
  freqs = (-math.log(max_period) * Tensor.arange(half, device=timesteps.device) / half).exp()
  args = timesteps.unsqueeze(1) * freqs.unsqueeze(0)
  out = Tensor.cat(args.cos(), args.sin(), dim=-1)
-  return out.cast(mixed_precision_dtype) if is_dtype_supported(mixed_precision_dtype) else out
+  return out.cast(dtypes.float16) if is_dtype_supported(dtypes.float16) else out

 class ResBlock:
-  def __init__(self, channels:int, emb_channels:int, out_channels:int, num_groups:int=32):
+  def __init__(self, channels:int, emb_channels:int, out_channels:int):
    self.in_layers = [
-      GroupNorm(num_groups, channels),
+      GroupNorm(32, channels),
      Tensor.silu,
      Conv2d(channels, out_channels, 3, padding=1),
    ]
@@ -27,7 +24,7 @@ class ResBlock:
      Linear(emb_channels, out_channels),
    ]
    self.out_layers = [
-      GroupNorm(num_groups, out_channels),
+      GroupNorm(32, out_channels),
      Tensor.silu,
      lambda x: x,  # needed for weights loading code to work
      Conv2d(out_channels, out_channels, 3, padding=1),
@@ -48,37 +45,35 @@ class CrossAttention:
    self.to_v = Linear(ctx_dim,   n_heads*d_head, bias=False)
    self.num_heads = n_heads
    self.head_size = d_head
-    self.attn = attention
    self.to_out = [Linear(n_heads*d_head, query_dim)]

  def __call__(self, x:Tensor, ctx:Optional[Tensor]=None) -> Tensor:
    ctx = x if ctx is None else ctx
    q,k,v = self.to_q(x), self.to_k(ctx), self.to_v(ctx)
    q,k,v = [y.reshape(x.shape[0], -1, self.num_heads, self.head_size).transpose(1,2) for y in (q,k,v)]
-    attention = self.attn(q, k, v).transpose(1,2)
+    attention = Tensor.scaled_dot_product_attention(q, k, v).transpose(1,2)
    h_ = attention.reshape(x.shape[0], -1, self.num_heads * self.head_size)
    return h_.sequential(self.to_out)

 class GEGLU:
  def __init__(self, dim_in:int, dim_out:int):
    self.proj = Linear(dim_in, dim_out * 2)
-    self.gelu = gelu
    self.dim_out = dim_out

  def __call__(self, x:Tensor) -> Tensor:
    x, gate = self.proj(x).chunk(2, dim=-1)
-    return x * self.gelu(gate)
+    return x * gate.gelu()

 class FeedForward:
  def __init__(self, dim:int, mult:int=4):
-    self.net: tuple[GEGLU, Callable, nn.Linear] = (
+    self.net = [
      GEGLU(dim, dim*mult),
      lambda x: x,  # needed for weights loading code to work
      Linear(dim*mult, dim)
-    )
+    ]

  def __call__(self, x:Tensor) -> Tensor:
-    return x.sequential(list(self.net))
+    return x.sequential(self.net)

 class BasicTransformerBlock:
  def __init__(self, dim:int, ctx_dim:int, n_heads:int, d_head:int):
@@ -97,13 +92,12 @@ class BasicTransformerBlock:

 # https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/modules/attention.py#L619
 class SpatialTransformer:
-  def __init__(self, channels:int, n_heads:int, d_head:int, ctx_dim:Union[int,List[int]], use_linear:bool, depth:int=1,
-               norm_eps:float=1e-5):
+  def __init__(self, channels:int, n_heads:int, d_head:int, ctx_dim:Union[int,List[int]], use_linear:bool, depth:int=1):
    if isinstance(ctx_dim, int):
      ctx_dim = [ctx_dim]*depth
    else:
      assert isinstance(ctx_dim, list) and depth == len(ctx_dim)
-    self.norm = GroupNorm(32, channels, eps=norm_eps)
+    self.norm = GroupNorm(32, channels)
    assert channels == n_heads * d_head
    self.proj_in  = Linear(channels, channels) if use_linear else Conv2d(channels, channels, 1)
    self.transformer_blocks = [BasicTransformerBlock(channels, ctx_dim[d], n_heads, d_head) for d in range(depth)]
@@ -140,9 +134,7 @@ class Upsample:

 # https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/modules/diffusionmodules/openaimodel.py#L472
 class UNetModel:
-  def __init__(self, adm_in_ch:Optional[int], in_ch:int, out_ch:int, model_ch:int, attention_resolutions:List[int], num_res_blocks:int,
-               channel_mult:List[int], transformer_depth:List[int], ctx_dim:Union[int,List[int]], use_linear:bool=False, d_head:Optional[int]=None,
-               n_heads:Optional[int]=None, num_groups:int=32, st_norm_eps:float=1e-5):
+  def __init__(self, adm_in_ch:Optional[int], in_ch:int, out_ch:int, model_ch:int, attention_resolutions:List[int], num_res_blocks:int, channel_mult:List[int], transformer_depth:List[int], ctx_dim:Union[int,List[int]], use_linear:bool=False, d_head:Optional[int]=None, n_heads:Optional[int]=None):
    self.model_ch = model_ch
    self.num_res_blocks = [num_res_blocks] * len(channel_mult)

@@ -182,12 +174,12 @@ class UNetModel:
    for idx, mult in enumerate(channel_mult):
      for _ in range(self.num_res_blocks[idx]):
        layers: List[Any] = [
-          ResBlock(ch, time_embed_dim, model_ch*mult, num_groups),
+          ResBlock(ch, time_embed_dim, model_ch*mult),
        ]
        ch = mult * model_ch
        if ds in attention_resolutions:
          d_head, n_heads = get_d_and_n_heads(ch)
-          layers.append(SpatialTransformer(ch, n_heads, d_head, ctx_dim, use_linear, depth=transformer_depth[idx], norm_eps=st_norm_eps))
+          layers.append(SpatialTransformer(ch, n_heads, d_head, ctx_dim, use_linear, depth=transformer_depth[idx]))

        self.input_blocks.append(layers)
        input_block_channels.append(ch)
@@ -201,9 +193,9 @@ class UNetModel:

    d_head, n_heads = get_d_and_n_heads(ch)
    self.middle_block: List = [
-      ResBlock(ch, time_embed_dim, ch, num_groups),
-      SpatialTransformer(ch, n_heads, d_head, ctx_dim, use_linear, depth=transformer_depth[-1], norm_eps=st_norm_eps),
-      ResBlock(ch, time_embed_dim, ch, num_groups),
+      ResBlock(ch, time_embed_dim, ch),
+      SpatialTransformer(ch, n_heads, d_head, ctx_dim, use_linear, depth=transformer_depth[-1]),
+      ResBlock(ch, time_embed_dim, ch),
    ]

    self.output_blocks = []
@@ -211,13 +203,13 @@ class UNetModel:
      for i in range(self.num_res_blocks[idx] + 1):
        ich = input_block_channels.pop()
        layers = [
-          ResBlock(ch + ich, time_embed_dim, model_ch*mult, num_groups),
+          ResBlock(ch + ich, time_embed_dim, model_ch*mult),
        ]
        ch = model_ch * mult

        if ds in attention_resolutions:
          d_head, n_heads = get_d_and_n_heads(ch)
-          layers.append(SpatialTransformer(ch, n_heads, d_head, ctx_dim, use_linear, depth=transformer_depth[idx], norm_eps=st_norm_eps))
+          layers.append(SpatialTransformer(ch, n_heads, d_head, ctx_dim, use_linear, depth=transformer_depth[idx]))

        if idx > 0 and i == self.num_res_blocks[idx]:
          layers.append(Upsample(ch))
@@ -225,7 +217,7 @@ class UNetModel:
        self.output_blocks.append(layers)

    self.out = [
-      GroupNorm(num_groups, ch),
+      GroupNorm(32, ch),
      Tensor.silu,
      Conv2d(model_ch, out_ch, 3, padding=1),
    ]
@@ -238,10 +230,10 @@ class UNetModel:
      assert y.shape[0] == x.shape[0]
      emb = emb + y.sequential(self.label_emb[0])

-    if is_dtype_supported(mixed_precision_dtype):
-      emb = emb.cast(mixed_precision_dtype)
-      ctx = ctx.cast(mixed_precision_dtype)
-      x   = x  .cast(mixed_precision_dtype)
+    if is_dtype_supported(dtypes.float16):
+      emb = emb.cast(dtypes.float16)
+      ctx = ctx.cast(dtypes.float16)
+      x   = x  .cast(dtypes.float16)

    def run(x:Tensor, bb) -> Tensor:
      if isinstance(bb, ResBlock): x = bb(x, emb)
@@ -272,4 +272,4 @@ def compare_launch_state(states, good_states):

  return True, "PASS"

-# IOCTL=1 CUDA=1 CUDA_PTX=1 python3 test/test_ops.py TestOps.test_tiny_add
+# IOCTL=1 PTX=1 CUDA=1 python3 test/test_ops.py TestOps.test_tiny_add
@@ -1,6 +1,6 @@
 from tinygrad import Tensor
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.nn.onnx import OnnxRunner, OnnxValue
+from tinygrad.frontend.onnx import OnnxRunner, OnnxValue
 import numpy as np
 import onnxruntime as ort

@@ -7,7 +7,7 @@ rm $LOGOPS
 test/external/process_replay/reset.py

 CI=1 python3 -m pytest -n=auto test/test_ops.py test/test_nn.py test/test_winograd.py test/models/test_real_world.py --durations=20
-CL=1 python3 -m pytest test/test_tiny.py
+GPU=1 python3 -m pytest test/test_tiny.py

 # extract, sort and uniq
 extra/optimization/extract_dataset.py
@@ -1,6 +1,6 @@
 # stuff needed to unpack a kernel
 from tinygrad import Variable
-from tinygrad.codegen.opt import Opt, OptOps
+from tinygrad.codegen.opt.kernel import Opt, OptOps
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.dtype import dtypes, PtrDType
 from tinygrad.shape.shapetracker import ShapeTracker
@@ -81,7 +81,7 @@ def lin_to_feats(lin:Kernel, use_sts=True):
  ret = [float(x) for x in ret]

  if use_sts:
-    my_sts = dedup([(x.shape == lin.full_shape, x.is_expanded(), any(v.mask is not None for v in x.views), len(x.views)) for x in lin.sts])
+    my_sts = dedup([(x.shape == lin.full_shape, x.real_strides(), any(v.mask is not None for v in x.views), len(x.views)) for x in lin.sts])
    assert len(my_sts) < MAX_BUFS
    sts_len = 3 + 5*MAX_DIMS
    for s in my_sts:
@@ -115,7 +115,7 @@ def time_linearizer(lin:Kernel, rawbufs:list[Buffer], allow_test_size=True, max_
  assert dev.compiler is not None

  rawbufs = _ensure_buffer_alloc(rawbufs)
-  var_vals: dict[str, int] = {k.expr:int(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
+  var_vals: dict[Variable, int] = {k:int(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
  p = get_program(lin.get_optimized_ast(), lin.opts)
  tms = _time_program(p, dev.compiler.compile(p.src), var_vals, rawbufs,
                      max_global_size=max_global_size if allow_test_size else None, clear_l2=clear_l2, cnt=cnt, name=to_function_name(lin.name))
@@ -16,9 +16,9 @@ class TestBeamSearch(unittest.TestCase):
    BEAM.value = self.old_beam

  def test_variable_ast_beam(self):
-    vi = Variable("a", 1, 10).bind(3)
-    a = rand(10, 3)[:vi]
-    a = (a+1).realize()
+    with Context(IGNORE_OOB=1):
+      a = rand(3, 3).reshape((Variable("a", 1, 10).bind(3), 3))
+      a = (a+1).realize()

  def test_big_prime_number(self):
    a = rand(367, 367)
@@ -42,16 +42,18 @@ class TestBeamSearch(unittest.TestCase):

  def test_variable_big_prime_number(self):
    v = Variable("v", 1, 400).bind(367)
-    a = rand(367, 400)
-    b = rand(400, 367)
-    c = (a[:, :v] @ b[:v, :]).realize()
-    np.testing.assert_allclose(c.numpy(), a[:, :367].numpy() @ b[:367, :].numpy(), atol=1e-4, rtol=1e-4)
+    a = rand(367, 367)
+    b = rand(367, 367)
+    with Context(IGNORE_OOB=1):
+      c = (a.reshape(367, v) @ b.reshape(v, 367)).realize()
+      np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), atol=1e-4, rtol=1e-4)

  def test_variable_shrink_prime_number(self):
    v = Variable("v", 1, 400).bind(367)
    a = rand(400, 367)
-    b = (a.shrink(((0,v), None))+1)[:367,:367].realize()
-    np.testing.assert_allclose(b.numpy(), a.numpy()[:367]+1, atol=1e-4, rtol=1e-4)
+    with Context(IGNORE_OOB=1):
+      b = (a.shrink(((0,v), None))+1).reshape(367,367).realize()
+      np.testing.assert_allclose(b.numpy(), a.numpy()[:367]+1, atol=1e-4, rtol=1e-4)

  def test_no_mutate_rawbuffers(self):
    a = rand(3, 3).realize()
@@ -1,6 +1,6 @@
 import ctypes, array
 from hexdump import hexdump
-from tinygrad.runtime.ops_cl import CLDevice
+from tinygrad.runtime.ops_gpu import GPUDevice
 from tinygrad.helpers import getenv, to_mv, mv_address
 from tinygrad.dtype import dtypes
 from tinygrad import Tensor, TinyJit
@@ -8,7 +8,7 @@ from tinygrad.runtime.autogen import opencl as cl
 if getenv("IOCTL"): import extra.qcom_gpu_driver.opencl_ioctl  # noqa: F401  # pylint: disable=unused-import

 # create raw opencl buffer.
-gdev = CLDevice()
+gdev = GPUDevice()
 cl_buf = cl.clCreateBuffer(gdev.context, cl.CL_MEM_READ_WRITE, 0x100, None, status := ctypes.c_int32())
 assert status.value == 0

@@ -673,7 +673,6 @@ impl<'a> Thread<'a> {
                                        39 => f32::log2(s0),
                                        42 => 1.0 / s0,
                                        43 => 1.0 / s0,
-                                        46 => 1.0 / f32::sqrt(s0),
                                        51 => f32::sqrt(s0),
                                        _ => todo_instr!(instruction)?,
                                    }
@@ -930,7 +929,7 @@ impl<'a> Thread<'a> {

            let op = ((instr >> 16) & 0x3ff) as u32;
            match op {
-                764 | 765 | 288 | 289 | 290 | 766 | 767 | 768 | 769 => {
+                764 | 765 | 288 | 289 | 290 | 766 | 768 | 769 => {
                    let vdst = (instr & 0xff) as usize;
                    let sdst = ((instr >> 8) & 0x7f) as usize;
                    let f = |i: u32| -> usize { ((instr >> i) & 0x1ff) as usize };
@@ -944,16 +943,6 @@ impl<'a> Thread<'a> {
                    assert_eq!(clmp, 0);

                    let vcc = match op {
-                        767 => {
-                            let (s0, s1, s2): (u32, u32, u64) = (self.val(s0), self.val(s1), self.val(s2));
-                            let (mul_result, overflow_mul) = (s0 as i64).overflowing_mul(s1 as i64);
-                            let (ret, overflow_add) = mul_result.overflowing_add(s2 as i64);
-                            let overflowed = overflow_mul || overflow_add;
-                            if self.exec.read() {
-                                self.vec_reg.write64(vdst, ret as u64);
-                            }
-                            overflowed
-                        },
                        766 => {
                            let (s0, s1, s2): (u32, u32, u64) = (self.val(s0), self.val(s1), self.val(s2));
                            let (mul_result, overflow_mul) = (s0 as u64).overflowing_mul(s1 as u64);
@@ -1257,7 +1246,7 @@ impl<'a> Thread<'a> {
                            }

                            let ret = match op {
-                                257 | 259 | 299 | 260 | 261 | 264 | 272 | 392 | 426 | 430 | 531 | 537 | 540 | 551 | 567 | 796 => {
+                                257 | 259 | 299 | 260 | 261 | 264 | 272 | 392 | 426 | 531 | 537 | 540 | 551 | 567 | 796 => {
                                    let s0 = f32::from_bits(s0).negate(0, neg).absolute(0, abs);
                                    let s1 = f32::from_bits(s1).negate(1, neg).absolute(1, abs);
                                    let s2 = f32::from_bits(s2).negate(2, neg).absolute(2, abs);
@@ -1269,7 +1258,6 @@ impl<'a> Thread<'a> {
                                        272 => f32::max(s0, s1),
                                        299 => f32::mul_add(s0, s1, f32::from_bits(self.vec_reg[vdst])),
                                        426 => s0.recip(),
-                                        430 => 1.0 / f32::sqrt(s0),
                                        531 => f32::mul_add(s0, s1, s2),
                                        537 => f32::min(f32::min(s0, s1), s2),
                                        540 => f32::max(f32::max(s0, s1), s2),
@@ -2637,14 +2625,6 @@ mod test_vop1 {
        assert_eq!(thread.vec_reg[3], 1071644672);
    }

-    #[test]
-    fn test_v_rsq_f32() {
-        let mut thread = _helper_test_thread();
-        thread.vec_reg[0] = f32::to_bits(4.0);
-        r(&vec![0x7E005D00, END_PRG], &mut thread);
-        assert_eq!(f32::from_bits(thread.vec_reg[0]), 0.5);
-    }
-
    #[test]
    fn test_v_frexp_exp_i32_f64() {
        [(3573412790272.0, 42), (69.0, 7), (2.0, 2), (f64::NEG_INFINITY, 0)]
@@ -58,7 +58,7 @@ if __name__ == "__main__":
            GlobalCounters.kernel_count -= 1

        if not getenv("NOOPT"): k.apply_opts(hand_coded_optimizations(k))
-        p2 = get_program(k.ast, k.opts, k.applied_opts)
+        p2 = get_program(k.get_optimized_ast(), k.opts)
        new_ei = replace(ei, prg=CompiledRunner(p2))
        new_ei.run()
        new_jit.append(new_ei)
@@ -4,7 +4,7 @@

 Only supported on 7900XTX, requires either AM (`rmmod amdgpu`) or disabling power gating on AMD (`ppfeaturemask=0xffff3fff`, don't forget to rebuild initramfs)

-SQTT is implemented on top of normal tinygrad profiling, `VIZ=1 SQTT=1` to get profile pickle with sqtt data embedded in it.
+SQTT is implemented on top of normal tinygrad PROFILE=1, `PROFILE=1 SQTT=1` to get profile pickle with sqtt data embedded in it.

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

@@ -1,68 +0,0 @@
-import ctypes
-from dataclasses import dataclass
-import tinygrad.runtime.autogen.comgr as comgr
-from tinygrad.runtime.support.compiler_amd import check
-
-@dataclass
-class InstrCtx:
-  pc:int=0
-  inst:str=""
-
-@comgr.amd_comgr_create_disassembly_info.argtypes[2]
-def instr_cb(text, user_data):
-  c = ctypes.cast(user_data, ctypes.POINTER(ctypes.py_object)).contents.value
-  c.inst = ctypes.string_at(text).decode("utf-8","replace").strip()
-  return comgr.AMD_COMGR_STATUS_SUCCESS
-
-# nop callback
-@comgr.amd_comgr_create_disassembly_info.argtypes[3]
-def addr_cb(*args): return comgr.AMD_COMGR_STATUS_SUCCESS
-
-def comgr_get_address_table(lib:bytes) -> dict[int, tuple[str, int]]:
-  check(comgr.amd_comgr_create_data(comgr.AMD_COMGR_DATA_KIND_EXECUTABLE, ctypes.byref(data_src:=comgr.amd_comgr_data_t())))
-  lib_buf = ctypes.create_string_buffer(lib, len(lib))
-  check(comgr.amd_comgr_set_data(data_src, len(lib), lib_buf))
-  check(comgr.amd_comgr_get_data_isa_name(data_src, isa_sz:=ctypes.c_size_t(128), isa:=(ctypes.c_char*isa_sz.value)()))
-
-  @comgr.amd_comgr_create_disassembly_info.argtypes[1]
-  def memory_cb(from_addr, to, size, _):
-    base, buf_len = ctypes.addressof(lib_buf), len(lib_buf)
-    start = int(from_addr) - base
-    if start < 0 or start >= buf_len: return 0
-    ctypes.memmove(to, base + start, n:=min(int(size), buf_len - start))
-    return n
-
-  info_src = comgr.amd_comgr_disassembly_info_t()
-  check(comgr.amd_comgr_create_disassembly_info(ctypes.cast(isa, ctypes.POINTER(ctypes.c_char)), memory_cb, instr_cb, addr_cb, info_src))
-
-  @comgr.amd_comgr_iterate_symbols.argtypes[1]
-  def sym_callback(sym, udata):
-    check(comgr.amd_comgr_symbol_get_info(sym, comgr.AMD_COMGR_SYMBOL_INFO_TYPE, ctypes.byref(sym_type:=ctypes.c_int())))
-    if sym_type.value != comgr.AMD_COMGR_SYMBOL_TYPE_FUNC: return comgr.AMD_COMGR_STATUS_SUCCESS
-    check(comgr.amd_comgr_symbol_get_info(sym, comgr.AMD_COMGR_SYMBOL_INFO_VALUE, ctypes.byref(vaddr:=ctypes.c_uint64())))
-    check(comgr.amd_comgr_symbol_get_info(sym, comgr.AMD_COMGR_SYMBOL_INFO_SIZE, ctypes.byref(size:=ctypes.c_uint64())))
-    check(comgr.amd_comgr_map_elf_virtual_address_to_code_object_offset(data_src, vaddr.value, ctypes.byref(offset:=ctypes.c_uint64()),
-                                                                        ctypes.byref(ctypes.c_uint64()), ctypes.byref(nobits:=ctypes.c_bool())))
-    check(nobits.value)
-    base = ctypes.addressof(lib_buf)
-    pc = base + offset.value
-    end = pc + size.value
-    addr_table = ctypes.cast(udata, ctypes.POINTER(ctypes.py_object)).contents.value
-    instr_ref = ctypes.py_object(ctx:=InstrCtx())
-    instr_ptr = ctypes.cast(ctypes.pointer(instr_ref), ctypes.c_void_p)
-    while pc < end:
-      size_read = ctypes.c_uint64(0)
-      ctx.pc = pc
-      st = comgr.amd_comgr_disassemble_instruction(info_src, ctypes.c_uint64(pc), instr_ptr, ctypes.byref(size_read))
-      if st == comgr.AMD_COMGR_STATUS_SUCCESS and size_read.value:
-        rel = (pc - base) - offset.value
-        addr_table[vaddr.value + rel] = (ctx.inst, int(size_read.value))
-        pc += size_read.value
-      else: # don't inf loop if comgr fails
-        b = ctypes.c_ubyte.from_buffer(lib_buf, pc - base).value
-        addr_table[vaddr.value + (pc - base - offset.value)] = (f"DISASSEMBLER ISSUE 0x{b:02x}", 1)
-        pc += 1
-    return comgr.AMD_COMGR_STATUS_SUCCESS
-  addr_table:dict[int, tuple[str, int]] = {}
-  check(comgr.amd_comgr_iterate_symbols(data_src, sym_callback, ctypes.cast(ctypes.pointer(ctypes.py_object(addr_table)), ctypes.c_void_p)))
-  return addr_table
@@ -155,10 +155,6 @@ class RGP:
      device_event = device_events[device]
    sqtt_events = [x for x in profile if isinstance(x, ProfileSQTTEvent) and x.device == device_event.device]
    if len(sqtt_events) == 0: raise RuntimeError(f"Device {device_event.device} doesn't contain SQTT data")
-    device_props = sqtt_events[0].props
-    gfx_ver = device_props['gfx_target_version'] // 10000
-    gfx_iplvl = getattr(sqtt, f"SQTT_GFXIP_LEVEL_GFXIP_{device_props['gfx_target_version']//10000}_{(device_props['gfx_target_version']//100)%100}",
-                        getattr(sqtt, f"SQTT_GFXIP_LEVEL_GFXIP_{device_props['gfx_target_version']//10000}", None))
    sqtt_itrace_enabled = any([event.itrace for event in sqtt_events])
    sqtt_itrace_masked = not all_same([event.itrace for event in sqtt_events])
    sqtt_itrace_se_mask = functools.reduce(lambda a,b: a|b, [int(event.itrace) << event.se for event in sqtt_events], 0) if sqtt_itrace_masked else 0
@@ -196,21 +192,21 @@ class RGP:
        flags=0,
        trace_shader_core_clock=0x93f05080,
        trace_memory_clock=0x4a723a40,
-        device_id={110000: 0x744c, 110003: 0x7480, 120001: 0x7550}[device_props['gfx_target_version']],
+        device_id=0x744c,
        device_revision_id=0xc8,
        vgprs_per_simd=1536,
        sgprs_per_simd=128*16,
-        shader_engines=device_props['array_count'] // device_props['simd_arrays_per_engine'],
-        compute_unit_per_shader_engine=device_props['simd_count'] // device_props['simd_per_cu'] // (device_props['array_count'] // device_props['simd_arrays_per_engine']),
-        simd_per_compute_unit=device_props['simd_per_cu'],
-        wavefronts_per_simd=device_props['max_waves_per_simd'],
+        shader_engines=6,
+        compute_unit_per_shader_engine=16,
+        simd_per_compute_unit=2,
+        wavefronts_per_simd=16,
        minimum_vgpr_alloc=4,
        vgpr_alloc_granularity=8,
        minimum_sgpr_alloc=128,
        sgpr_alloc_granularity=128,
        hardware_contexts=8,
        gpu_type=sqtt.SQTT_GPU_TYPE_DISCRETE,
-        gfxip_level=gfx_iplvl,
+        gfxip_level=sqtt.SQTT_GFXIP_LEVEL_GFXIP_11_0,
        gpu_index=0,
        gds_size=0,
        gds_per_shader_engine=0,
@@ -222,7 +218,7 @@ class RGP:
        vram_bus_width=384, # 384-bit
        l2_cache_size=6 * 1024 * 1024, # 6 MB
        l1_cache_size=32 * 1024, # 32 KB per SIMD (?)
-        lds_size=device_props['lds_size_in_kb'] * 1024,
+        lds_size=65536, # 64 KB per CU
        gpu_name=b'NAVI31',
        alu_per_clock=0,
        texture_per_clock=0,
@@ -261,7 +257,7 @@ class RGP:
            major_version=0, minor_version=2,
          ),
          shader_engine_index=sqtt_event.se,
-          sqtt_version={11: sqtt.SQTT_VERSION_3_2, 12: sqtt.SQTT_VERSION_3_3}.get(gfx_ver),
+          sqtt_version=sqtt.SQTT_VERSION_3_2,
          _0=sqtt.union_sqtt_file_chunk_sqtt_desc_0(
            v1=sqtt.struct_sqtt_file_chunk_sqtt_desc_0_v1(
              instrumentation_spec_version=1,
@@ -1,96 +0,0 @@
-import ctypes, pathlib, argparse, pickle, re, functools, dataclasses
-from extra.sqtt.rocprof import rocprof
-from extra.sqtt.disasm import comgr_get_address_table
-from tinygrad.helpers import temp, DEBUG
-from tinygrad.device import ProfileEvent, ProfileProgramEvent
-from tinygrad.runtime.ops_amd import ProfileSQTTEvent
-
-@dataclasses.dataclass
-class InstInfo:
-  typ:str=""
-  inst:str=""
-  hit:int=0
-  lat:int=0
-  stall:int=0
-  def __str__(self): return f"{self.inst:>20} hits:{self.typ:>6} hits:{self.hit:>6} latency:{self.lat:>6} stall:{self.stall:>6}"
-
-  def on_ev(self, ev):
-    self.hit, self.lat, self.stall = self.hit + 1, self.lat + ev.duration, self.stall + ev.stall
-
-class _ROCParseCtx:
-  def __init__(self, sqtt_evs:list[ProfileSQTTEvent], prog_evs:list[ProfileProgramEvent]):
-    self.sqtt_evs, self.prog_evs = iter(sqtt_evs), prog_evs
-    self.wave_events, self.disasms, self.addr2prg = {}, {}, {}
-
-    for prog in prog_evs:
-      for addr, info in comgr_get_address_table(prog.lib).items():
-        self.disasms[prog.base + addr] = info
-        self.addr2prg[prog.base + addr] = prog
-
-  def next_sqtt(self): return next(self.sqtt_evs, None)
-  def find_program(self, addr): return self.addr2prg[addr]
-
-  def on_occupancy_ev(self, ev):
-    if DEBUG >= 4: print("OCC", ev.time, ev.cu, ev.simd, ev.wave_id, ev.start)
-
-  def on_wave_ev(self, ev):
-    if DEBUG >= 4: print("WAVE", ev.wave_id, ev.cu, ev.simd, ev.contexts, ev.begin_time, ev.end_time)
-
-    asm = {}
-    for j in range(ev.instructions_size):
-      inst_ev = ev.instructions_array[j]
-      inst_typ = rocprof.rocprofiler_thread_trace_decoder_inst_category_t__enumvalues[inst_ev.category]
-      asm.setdefault(inst_ev.pc.address, InstInfo(typ=inst_typ, inst=self.disasms[inst_ev.pc.address][0]))
-      asm[inst_ev.pc.address].on_ev(inst_ev)
-
-    self.wave_events[(self.find_program(ev.instructions_array[0].pc.address).name, ev.wave_id, ev.cu, ev.simd)] = asm
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser()
-  parser.add_argument('--profile', type=pathlib.Path, help='Path to profile', default=pathlib.Path(temp("profile.pkl", append_user=True)))
-  args = parser.parse_args()
-
-  with args.profile.open("rb") as f: profile = pickle.load(f)
-  sqtt_events:list[ProfileSQTTEvent] = []
-  prog_events:list[ProfileProgramEvent] = []
-  for e in profile:
-    if isinstance(e, ProfileSQTTEvent): sqtt_events.append(e)
-    if isinstance(e, ProfileProgramEvent) and e.device.startswith("AMD"): prog_events.append(e)
-
-  ROCParseCtx = _ROCParseCtx(sqtt_events, prog_events)
-
-  @rocprof.rocprof_trace_decoder_se_data_callback_t
-  def copy_cb(buf, buf_size, data_ptr):
-    if (prof:=ROCParseCtx.next_sqtt()) is None: return 0
-    buf[0] = ctypes.cast((ctypes.c_ubyte * len(prof.blob)).from_buffer_copy(prof.blob), ctypes.POINTER(ctypes.c_ubyte))
-    buf_size[0] = len(prof.blob)
-    return len(prof.blob)
-
-  @rocprof.rocprof_trace_decoder_trace_callback_t
-  def trace_cb(record_type, events_ptr, n, data_ptr):
-    match record_type:
-      case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY:
-        for ev in (rocprof.rocprofiler_thread_trace_decoder_occupancy_t * n).from_address(events_ptr): ROCParseCtx.on_occupancy_ev(ev)
-      case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE:
-        for ev in (rocprof.rocprofiler_thread_trace_decoder_wave_t * n).from_address(events_ptr): ROCParseCtx.on_wave_ev(ev)
-      case _:
-        if DEBUG >= 2: print(rocprof.rocprofiler_thread_trace_decoder_record_type_t__enumvalues[record_type], events_ptr, n)
-    return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
-
-  @rocprof.rocprof_trace_decoder_isa_callback_t
-  def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, data_ptr):
-    instr, mem_size_ptr[0] = ROCParseCtx.disasms[pc.address]
-
-    # this is the number of bytes to next instruction, set to 0 for end_pgm
-    if instr == "s_endpgm": mem_size_ptr[0] = 0
-    if (max_sz:=size_ptr[0]) == 0: return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES
-
-    # truncate the instr if it doesn't fit
-    if (str_sz:=len(instr_bytes:=instr.encode()))+1 > max_sz: str_sz = max_sz
-    ctypes.memmove(instr_ptr, instr_bytes, str_sz)
-    size_ptr[0] = str_sz
-
-    return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
-
-  rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
-  print(ROCParseCtx.wave_events.keys())
@@ -1,657 +0,0 @@
-# pylint: skip-file
-# mypy: ignore-errors
-# -*- coding: utf-8 -*-
-#
-# TARGET arch is: []
-# WORD_SIZE is: 8
-# POINTER_SIZE is: 8
-# LONGDOUBLE_SIZE is: 16
-#
-import ctypes
-from tinygrad.helpers import fetch
-
-
-class AsDictMixin:
-    @classmethod
-    def as_dict(cls, self):
-        result = {}
-        if not isinstance(self, AsDictMixin):
-            # not a structure, assume it's already a python object
-            return self
-        if not hasattr(cls, "_fields_"):
-            return result
-        # sys.version_info >= (3, 5)
-        # for (field, *_) in cls._fields_:  # noqa
-        for field_tuple in cls._fields_:  # noqa
-            field = field_tuple[0]
-            if field.startswith('PADDING_'):
-                continue
-            value = getattr(self, field)
-            type_ = type(value)
-            if hasattr(value, "_length_") and hasattr(value, "_type_"):
-                # array
-                if not hasattr(type_, "as_dict"):
-                    value = [v for v in value]
-                else:
-                    type_ = type_._type_
-                    value = [type_.as_dict(v) for v in value]
-            elif hasattr(value, "contents") and hasattr(value, "_type_"):
-                # pointer
-                try:
-                    if not hasattr(type_, "as_dict"):
-                        value = value.contents
-                    else:
-                        type_ = type_._type_
-                        value = type_.as_dict(value.contents)
-                except ValueError:
-                    # nullptr
-                    value = None
-            elif isinstance(value, AsDictMixin):
-                # other structure
-                value = type_.as_dict(value)
-            result[field] = value
-        return result
-
-
-class Structure(ctypes.Structure, AsDictMixin):
-
-    def __init__(self, *args, **kwds):
-        # We don't want to use positional arguments fill PADDING_* fields
-
-        args = dict(zip(self.__class__._field_names_(), args))
-        args.update(kwds)
-        super(Structure, self).__init__(**args)
-
-    @classmethod
-    def _field_names_(cls):
-        if hasattr(cls, '_fields_'):
-            return (f[0] for f in cls._fields_ if not f[0].startswith('PADDING'))
-        else:
-            return ()
-
-    @classmethod
-    def get_type(cls, field):
-        for f in cls._fields_:
-            if f[0] == field:
-                return f[1]
-        return None
-
-    @classmethod
-    def bind(cls, bound_fields):
-        fields = {}
-        for name, type_ in cls._fields_:
-            if hasattr(type_, "restype"):
-                if name in bound_fields:
-                    if bound_fields[name] is None:
-                        fields[name] = type_()
-                    else:
-                        # use a closure to capture the callback from the loop scope
-                        fields[name] = (
-                            type_((lambda callback: lambda *args: callback(*args))(
-                                bound_fields[name]))
-                        )
-                    del bound_fields[name]
-                else:
-                    # default callback implementation (does nothing)
-                    try:
-                        default_ = type_(0).restype().value
-                    except TypeError:
-                        default_ = None
-                    fields[name] = type_((
-                        lambda default_: lambda *args: default_)(default_))
-            else:
-                # not a callback function, use default initialization
-                if name in bound_fields:
-                    fields[name] = bound_fields[name]
-                    del bound_fields[name]
-                else:
-                    fields[name] = type_()
-        if len(bound_fields) != 0:
-            raise ValueError(
-                "Cannot bind the following unknown callback(s) {}.{}".format(
-                    cls.__name__, bound_fields.keys()
-            ))
-        return cls(**fields)
-
-
-class Union(ctypes.Union, AsDictMixin):
-    pass
-
-
-
-c_int128 = ctypes.c_ubyte*16
-c_uint128 = c_int128
-void = None
-if ctypes.sizeof(ctypes.c_longdouble) == 16:
-    c_long_double_t = ctypes.c_longdouble
-else:
-    c_long_double_t = ctypes.c_ubyte*16
-
-def string_cast(char_pointer, encoding='utf-8', errors='strict'):
-    value = ctypes.cast(char_pointer, ctypes.c_char_p).value
-    if value is not None and encoding is not None:
-        value = value.decode(encoding, errors=errors)
-    return value
-
-
-def char_pointer_cast(string, encoding='utf-8'):
-    if encoding is not None:
-        try:
-            string = string.encode(encoding)
-        except AttributeError:
-            # In Python3, bytes has no encode attribute
-            pass
-    string = ctypes.c_char_p(string)
-    return ctypes.cast(string, ctypes.POINTER(ctypes.c_char))
-
-
-
-class FunctionFactoryStub:
-    def __getattr__(self, _):
-      return ctypes.CFUNCTYPE(lambda y:y)
-
-# libraries['FIXME_STUB'] explanation
-# As you did not list (-l libraryname.so) a library that exports this function
-# This is a non-working stub instead.
-# You can either re-run clan2py with -l /path/to/library.so
-# Or manually fix this by comment the ctypes.CDLL loading
-_libraries = {}
-_libraries['FIXME_STUB'] = ctypes.CDLL(str(fetch('https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so'))) #  ctypes.CDLL('FIXME_STUB')
-
-
-
-# values for enumeration 'rocprofiler_thread_trace_decoder_info_t'
-rocprofiler_thread_trace_decoder_info_t__enumvalues = {
-    0: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_NONE',
-    1: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_DATA_LOST',
-    2: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_STITCH_INCOMPLETE',
-    3: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_WAVE_INCOMPLETE',
-    4: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_LAST',
-}
-ROCPROFILER_THREAD_TRACE_DECODER_INFO_NONE = 0
-ROCPROFILER_THREAD_TRACE_DECODER_INFO_DATA_LOST = 1
-ROCPROFILER_THREAD_TRACE_DECODER_INFO_STITCH_INCOMPLETE = 2
-ROCPROFILER_THREAD_TRACE_DECODER_INFO_WAVE_INCOMPLETE = 3
-ROCPROFILER_THREAD_TRACE_DECODER_INFO_LAST = 4
-rocprofiler_thread_trace_decoder_info_t = ctypes.c_uint32 # enum
-class struct_rocprofiler_thread_trace_decoder_pc_t(Structure):
-    pass
-
-struct_rocprofiler_thread_trace_decoder_pc_t._pack_ = 1 # source:False
-struct_rocprofiler_thread_trace_decoder_pc_t._fields_ = [
-    ('address', ctypes.c_uint64),
-    ('code_object_id', ctypes.c_uint64),
-]
-
-rocprofiler_thread_trace_decoder_pc_t = struct_rocprofiler_thread_trace_decoder_pc_t
-class struct_rocprofiler_thread_trace_decoder_perfevent_t(Structure):
-    pass
-
-struct_rocprofiler_thread_trace_decoder_perfevent_t._pack_ = 1 # source:False
-struct_rocprofiler_thread_trace_decoder_perfevent_t._fields_ = [
-    ('time', ctypes.c_int64),
-    ('events0', ctypes.c_uint16),
-    ('events1', ctypes.c_uint16),
-    ('events2', ctypes.c_uint16),
-    ('events3', ctypes.c_uint16),
-    ('CU', ctypes.c_ubyte),
-    ('bank', ctypes.c_ubyte),
-    ('PADDING_0', ctypes.c_ubyte * 6),
-]
-
-rocprofiler_thread_trace_decoder_perfevent_t = struct_rocprofiler_thread_trace_decoder_perfevent_t
-class struct_rocprofiler_thread_trace_decoder_occupancy_t(Structure):
-    pass
-
-struct_rocprofiler_thread_trace_decoder_occupancy_t._pack_ = 1 # source:False
-struct_rocprofiler_thread_trace_decoder_occupancy_t._fields_ = [
-    ('pc', rocprofiler_thread_trace_decoder_pc_t),
-    ('time', ctypes.c_uint64),
-    ('reserved', ctypes.c_ubyte),
-    ('cu', ctypes.c_ubyte),
-    ('simd', ctypes.c_ubyte),
-    ('wave_id', ctypes.c_ubyte),
-    ('start', ctypes.c_uint32, 1),
-    ('_rsvd', ctypes.c_uint32, 31),
-]
-
-rocprofiler_thread_trace_decoder_occupancy_t = struct_rocprofiler_thread_trace_decoder_occupancy_t
-
-# values for enumeration 'rocprofiler_thread_trace_decoder_wstate_type_t'
-rocprofiler_thread_trace_decoder_wstate_type_t__enumvalues = {
-    0: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EMPTY',
-    1: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_IDLE',
-    2: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EXEC',
-    3: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_WAIT',
-    4: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_STALL',
-    5: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_LAST',
-}
-ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EMPTY = 0
-ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_IDLE = 1
-ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EXEC = 2
-ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_WAIT = 3
-ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_STALL = 4
-ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_LAST = 5
-rocprofiler_thread_trace_decoder_wstate_type_t = ctypes.c_uint32 # enum
-class struct_rocprofiler_thread_trace_decoder_wave_state_t(Structure):
-    pass
-
-struct_rocprofiler_thread_trace_decoder_wave_state_t._pack_ = 1 # source:False
-struct_rocprofiler_thread_trace_decoder_wave_state_t._fields_ = [
-    ('type', ctypes.c_int32),
-    ('duration', ctypes.c_int32),
-]
-
-rocprofiler_thread_trace_decoder_wave_state_t = struct_rocprofiler_thread_trace_decoder_wave_state_t
-
-# values for enumeration 'rocprofiler_thread_trace_decoder_inst_category_t'
-rocprofiler_thread_trace_decoder_inst_category_t__enumvalues = {
-    0: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_NONE',
-    1: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_SMEM',
-    2: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_SALU',
-    3: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_VMEM',
-    4: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_FLAT',
-    5: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_LDS',
-    6: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_VALU',
-    7: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_JUMP',
-    8: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_NEXT',
-    9: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_IMMED',
-    10: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_CONTEXT',
-    11: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_MESSAGE',
-    12: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_BVH',
-    13: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_LAST',
-}
-ROCPROFILER_THREAD_TRACE_DECODER_INST_NONE = 0
-ROCPROFILER_THREAD_TRACE_DECODER_INST_SMEM = 1
-ROCPROFILER_THREAD_TRACE_DECODER_INST_SALU = 2
-ROCPROFILER_THREAD_TRACE_DECODER_INST_VMEM = 3
-ROCPROFILER_THREAD_TRACE_DECODER_INST_FLAT = 4
-ROCPROFILER_THREAD_TRACE_DECODER_INST_LDS = 5
-ROCPROFILER_THREAD_TRACE_DECODER_INST_VALU = 6
-ROCPROFILER_THREAD_TRACE_DECODER_INST_JUMP = 7
-ROCPROFILER_THREAD_TRACE_DECODER_INST_NEXT = 8
-ROCPROFILER_THREAD_TRACE_DECODER_INST_IMMED = 9
-ROCPROFILER_THREAD_TRACE_DECODER_INST_CONTEXT = 10
-ROCPROFILER_THREAD_TRACE_DECODER_INST_MESSAGE = 11
-ROCPROFILER_THREAD_TRACE_DECODER_INST_BVH = 12
-ROCPROFILER_THREAD_TRACE_DECODER_INST_LAST = 13
-rocprofiler_thread_trace_decoder_inst_category_t = ctypes.c_uint32 # enum
-class struct_rocprofiler_thread_trace_decoder_inst_t(Structure):
-    pass
-
-struct_rocprofiler_thread_trace_decoder_inst_t._pack_ = 1 # source:False
-struct_rocprofiler_thread_trace_decoder_inst_t._fields_ = [
-    ('category', ctypes.c_uint32, 8),
-    ('stall', ctypes.c_uint32, 24),
-    ('duration', ctypes.c_int32),
-    ('time', ctypes.c_int64),
-    ('pc', rocprofiler_thread_trace_decoder_pc_t),
-]
-
-rocprofiler_thread_trace_decoder_inst_t = struct_rocprofiler_thread_trace_decoder_inst_t
-class struct_rocprofiler_thread_trace_decoder_wave_t(Structure):
-    pass
-
-struct_rocprofiler_thread_trace_decoder_wave_t._pack_ = 1 # source:False
-struct_rocprofiler_thread_trace_decoder_wave_t._fields_ = [
-    ('cu', ctypes.c_ubyte),
-    ('simd', ctypes.c_ubyte),
-    ('wave_id', ctypes.c_ubyte),
-    ('contexts', ctypes.c_ubyte),
-    ('_rsvd1', ctypes.c_uint32),
-    ('_rsvd2', ctypes.c_uint32),
-    ('_rsvd3', ctypes.c_uint32),
-    ('begin_time', ctypes.c_int64),
-    ('end_time', ctypes.c_int64),
-    ('timeline_size', ctypes.c_uint64),
-    ('instructions_size', ctypes.c_uint64),
-    ('timeline_array', ctypes.POINTER(struct_rocprofiler_thread_trace_decoder_wave_state_t)),
-    ('instructions_array', ctypes.POINTER(struct_rocprofiler_thread_trace_decoder_inst_t)),
-]
-
-rocprofiler_thread_trace_decoder_wave_t = struct_rocprofiler_thread_trace_decoder_wave_t
-class struct_rocprofiler_thread_trace_decoder_realtime_t(Structure):
-    pass
-
-struct_rocprofiler_thread_trace_decoder_realtime_t._pack_ = 1 # source:False
-struct_rocprofiler_thread_trace_decoder_realtime_t._fields_ = [
-    ('shader_clock', ctypes.c_int64),
-    ('realtime_clock', ctypes.c_uint64),
-    ('reserved', ctypes.c_uint64),
-]
-
-rocprofiler_thread_trace_decoder_realtime_t = struct_rocprofiler_thread_trace_decoder_realtime_t
-
-# values for enumeration 'rocprofiler_thread_trace_decoder_shaderdata_flags_t'
-rocprofiler_thread_trace_decoder_shaderdata_flags_t__enumvalues = {
-    0: 'ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_IMM',
-    1: 'ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_PRIV',
-}
-ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_IMM = 0
-ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_PRIV = 1
-rocprofiler_thread_trace_decoder_shaderdata_flags_t = ctypes.c_uint32 # enum
-class struct_rocprofiler_thread_trace_decoder_shaderdata_t(Structure):
-    pass
-
-struct_rocprofiler_thread_trace_decoder_shaderdata_t._pack_ = 1 # source:False
-struct_rocprofiler_thread_trace_decoder_shaderdata_t._fields_ = [
-    ('time', ctypes.c_int64),
-    ('value', ctypes.c_uint64),
-    ('cu', ctypes.c_ubyte),
-    ('simd', ctypes.c_ubyte),
-    ('wave_id', ctypes.c_ubyte),
-    ('flags', ctypes.c_ubyte),
-    ('reserved', ctypes.c_uint32),
-]
-
-rocprofiler_thread_trace_decoder_shaderdata_t = struct_rocprofiler_thread_trace_decoder_shaderdata_t
-
-# values for enumeration 'rocprofiler_thread_trace_decoder_record_type_t'
-rocprofiler_thread_trace_decoder_record_type_t__enumvalues = {
-    0: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_GFXIP',
-    1: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY',
-    2: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_PERFEVENT',
-    3: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE',
-    4: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_INFO',
-    5: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_DEBUG',
-    6: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_SHADERDATA',
-    7: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_REALTIME',
-    8: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_RT_FREQUENCY',
-    9: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_LAST',
-}
-ROCPROFILER_THREAD_TRACE_DECODER_RECORD_GFXIP = 0
-ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY = 1
-ROCPROFILER_THREAD_TRACE_DECODER_RECORD_PERFEVENT = 2
-ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE = 3
-ROCPROFILER_THREAD_TRACE_DECODER_RECORD_INFO = 4
-ROCPROFILER_THREAD_TRACE_DECODER_RECORD_DEBUG = 5
-ROCPROFILER_THREAD_TRACE_DECODER_RECORD_SHADERDATA = 6
-ROCPROFILER_THREAD_TRACE_DECODER_RECORD_REALTIME = 7
-ROCPROFILER_THREAD_TRACE_DECODER_RECORD_RT_FREQUENCY = 8
-ROCPROFILER_THREAD_TRACE_DECODER_RECORD_LAST = 9
-rocprofiler_thread_trace_decoder_record_type_t = ctypes.c_uint32 # enum
-
-# values for enumeration 'c__EA_rocprofiler_thread_trace_decoder_status_t'
-c__EA_rocprofiler_thread_trace_decoder_status_t__enumvalues = {
-    0: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS',
-    1: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR',
-    2: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES',
-    3: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_ARGUMENT',
-    4: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_SHADER_DATA',
-    5: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_LAST',
-}
-ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS = 0
-ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR = 1
-ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES = 2
-ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_ARGUMENT = 3
-ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_SHADER_DATA = 4
-ROCPROFILER_THREAD_TRACE_DECODER_STATUS_LAST = 5
-c__EA_rocprofiler_thread_trace_decoder_status_t = ctypes.c_uint32 # enum
-rocprofiler_thread_trace_decoder_status_t = c__EA_rocprofiler_thread_trace_decoder_status_t
-rocprofiler_thread_trace_decoder_status_t__enumvalues = c__EA_rocprofiler_thread_trace_decoder_status_t__enumvalues
-rocprof_trace_decoder_trace_callback_t = ctypes.CFUNCTYPE(c__EA_rocprofiler_thread_trace_decoder_status_t, rocprofiler_thread_trace_decoder_record_type_t, ctypes.POINTER(None), ctypes.c_uint64, ctypes.POINTER(None))
-rocprof_trace_decoder_isa_callback_t = ctypes.CFUNCTYPE(c__EA_rocprofiler_thread_trace_decoder_status_t, ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_uint64), ctypes.POINTER(ctypes.c_uint64), struct_rocprofiler_thread_trace_decoder_pc_t, ctypes.POINTER(None))
-rocprof_trace_decoder_se_data_callback_t = ctypes.CFUNCTYPE(ctypes.c_uint64, ctypes.POINTER(ctypes.POINTER(ctypes.c_ubyte)), ctypes.POINTER(ctypes.c_uint64), ctypes.POINTER(None))
-try:
-    rocprof_trace_decoder_parse_data = _libraries['FIXME_STUB'].rocprof_trace_decoder_parse_data
-    rocprof_trace_decoder_parse_data.restype = rocprofiler_thread_trace_decoder_status_t
-    rocprof_trace_decoder_parse_data.argtypes = [rocprof_trace_decoder_se_data_callback_t, rocprof_trace_decoder_trace_callback_t, rocprof_trace_decoder_isa_callback_t, ctypes.POINTER(None)]
-except AttributeError:
-    pass
-try:
-    rocprof_trace_decoder_get_info_string = _libraries['FIXME_STUB'].rocprof_trace_decoder_get_info_string
-    rocprof_trace_decoder_get_info_string.restype = ctypes.POINTER(ctypes.c_char)
-    rocprof_trace_decoder_get_info_string.argtypes = [rocprofiler_thread_trace_decoder_info_t]
-except AttributeError:
-    pass
-try:
-    rocprof_trace_decoder_get_status_string = _libraries['FIXME_STUB'].rocprof_trace_decoder_get_status_string
-    rocprof_trace_decoder_get_status_string.restype = ctypes.POINTER(ctypes.c_char)
-    rocprof_trace_decoder_get_status_string.argtypes = [rocprofiler_thread_trace_decoder_status_t]
-except AttributeError:
-    pass
-rocprofiler_thread_trace_decoder_debug_callback_t = ctypes.CFUNCTYPE(None, ctypes.c_int64, ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), ctypes.POINTER(None))
-uint64_t = ctypes.c_uint64
-try:
-    rocprof_trace_decoder_dump_data = _libraries['FIXME_STUB'].rocprof_trace_decoder_dump_data
-    rocprof_trace_decoder_dump_data.restype = rocprofiler_thread_trace_decoder_status_t
-    rocprof_trace_decoder_dump_data.argtypes = [ctypes.POINTER(ctypes.c_char), uint64_t, rocprofiler_thread_trace_decoder_debug_callback_t, ctypes.POINTER(None)]
-except AttributeError:
-    pass
-class union_rocprof_trace_decoder_gfx9_header_t(Union):
-    pass
-
-class struct_rocprof_trace_decoder_gfx9_header_t_0(Structure):
-    pass
-
-struct_rocprof_trace_decoder_gfx9_header_t_0._pack_ = 1 # source:False
-struct_rocprof_trace_decoder_gfx9_header_t_0._fields_ = [
-    ('legacy_version', ctypes.c_uint64, 13),
-    ('gfx9_version2', ctypes.c_uint64, 3),
-    ('DSIMDM', ctypes.c_uint64, 4),
-    ('DCU', ctypes.c_uint64, 5),
-    ('reserved1', ctypes.c_uint64, 1),
-    ('SEID', ctypes.c_uint64, 6),
-    ('reserved2', ctypes.c_uint64, 32),
-]
-
-union_rocprof_trace_decoder_gfx9_header_t._pack_ = 1 # source:False
-union_rocprof_trace_decoder_gfx9_header_t._anonymous_ = ('_0',)
-union_rocprof_trace_decoder_gfx9_header_t._fields_ = [
-    ('_0', struct_rocprof_trace_decoder_gfx9_header_t_0),
-    ('raw', ctypes.c_uint64),
-]
-
-rocprof_trace_decoder_gfx9_header_t = union_rocprof_trace_decoder_gfx9_header_t
-class union_rocprof_trace_decoder_instrument_enable_t(Union):
-    pass
-
-class struct_rocprof_trace_decoder_instrument_enable_t_0(Structure):
-    pass
-
-struct_rocprof_trace_decoder_instrument_enable_t_0._pack_ = 1 # source:False
-struct_rocprof_trace_decoder_instrument_enable_t_0._fields_ = [
-    ('char1', ctypes.c_uint32, 8),
-    ('char2', ctypes.c_uint32, 8),
-    ('char3', ctypes.c_uint32, 8),
-    ('char4', ctypes.c_uint32, 8),
-]
-
-union_rocprof_trace_decoder_instrument_enable_t._pack_ = 1 # source:False
-union_rocprof_trace_decoder_instrument_enable_t._anonymous_ = ('_0',)
-union_rocprof_trace_decoder_instrument_enable_t._fields_ = [
-    ('_0', struct_rocprof_trace_decoder_instrument_enable_t_0),
-    ('u32All', ctypes.c_uint32),
-]
-
-rocprof_trace_decoder_instrument_enable_t = union_rocprof_trace_decoder_instrument_enable_t
-class union_rocprof_trace_decoder_packet_header_t(Union):
-    pass
-
-class struct_rocprof_trace_decoder_packet_header_t_0(Structure):
-    pass
-
-struct_rocprof_trace_decoder_packet_header_t_0._pack_ = 1 # source:False
-struct_rocprof_trace_decoder_packet_header_t_0._fields_ = [
-    ('opcode', ctypes.c_uint32, 8),
-    ('type', ctypes.c_uint32, 4),
-    ('data20', ctypes.c_uint32, 20),
-]
-
-union_rocprof_trace_decoder_packet_header_t._pack_ = 1 # source:False
-union_rocprof_trace_decoder_packet_header_t._anonymous_ = ('_0',)
-union_rocprof_trace_decoder_packet_header_t._fields_ = [
-    ('_0', struct_rocprof_trace_decoder_packet_header_t_0),
-    ('u32All', ctypes.c_uint32),
-]
-
-rocprof_trace_decoder_packet_header_t = union_rocprof_trace_decoder_packet_header_t
-
-# values for enumeration 'rocprof_trace_decoder_packet_opcode_t'
-rocprof_trace_decoder_packet_opcode_t__enumvalues = {
-    4: 'ROCPROF_TRACE_DECODER_PACKET_OPCODE_CODEOBJ',
-    5: 'ROCPROF_TRACE_DECODER_PACKET_OPCODE_RT_TIMESTAMP',
-    6: 'ROCPROF_TRACE_DECODER_PACKET_OPCODE_AGENT_INFO',
-}
-ROCPROF_TRACE_DECODER_PACKET_OPCODE_CODEOBJ = 4
-ROCPROF_TRACE_DECODER_PACKET_OPCODE_RT_TIMESTAMP = 5
-ROCPROF_TRACE_DECODER_PACKET_OPCODE_AGENT_INFO = 6
-rocprof_trace_decoder_packet_opcode_t = ctypes.c_uint32 # enum
-
-# values for enumeration 'rocprof_trace_decoder_agent_info_type_t'
-rocprof_trace_decoder_agent_info_type_t__enumvalues = {
-    0: 'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_RT_FREQUENCY_KHZ',
-    1: 'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_COUNTER_INTERVAL',
-    2: 'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_LAST',
-}
-ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_RT_FREQUENCY_KHZ = 0
-ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_COUNTER_INTERVAL = 1
-ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_LAST = 2
-rocprof_trace_decoder_agent_info_type_t = ctypes.c_uint32 # enum
-class union_rocprof_trace_decoder_codeobj_marker_tail_t(Union):
-    pass
-
-class struct_rocprof_trace_decoder_codeobj_marker_tail_t_0(Structure):
-    pass
-
-struct_rocprof_trace_decoder_codeobj_marker_tail_t_0._pack_ = 1 # source:False
-struct_rocprof_trace_decoder_codeobj_marker_tail_t_0._fields_ = [
-    ('isUnload', ctypes.c_uint32, 1),
-    ('bFromStart', ctypes.c_uint32, 1),
-    ('legacy_id', ctypes.c_uint32, 30),
-]
-
-union_rocprof_trace_decoder_codeobj_marker_tail_t._pack_ = 1 # source:False
-union_rocprof_trace_decoder_codeobj_marker_tail_t._anonymous_ = ('_0',)
-union_rocprof_trace_decoder_codeobj_marker_tail_t._fields_ = [
-    ('_0', struct_rocprof_trace_decoder_codeobj_marker_tail_t_0),
-    ('raw', ctypes.c_uint32),
-]
-
-rocprof_trace_decoder_codeobj_marker_tail_t = union_rocprof_trace_decoder_codeobj_marker_tail_t
-
-# values for enumeration 'rocprof_trace_decoder_codeobj_marker_type_t'
-rocprof_trace_decoder_codeobj_marker_type_t__enumvalues = {
-    0: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_TAIL',
-    1: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_LO',
-    2: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_LO',
-    3: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_HI',
-    4: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_HI',
-    5: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_LO',
-    6: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_HI',
-    7: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_LAST',
-}
-ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_TAIL = 0
-ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_LO = 1
-ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_LO = 2
-ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_HI = 3
-ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_HI = 4
-ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_LO = 5
-ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_HI = 6
-ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_LAST = 7
-rocprof_trace_decoder_codeobj_marker_type_t = ctypes.c_uint32 # enum
-__all__ = \
-    ['ROCPROFILER_THREAD_TRACE_DECODER_INFO_DATA_LOST',
-    'ROCPROFILER_THREAD_TRACE_DECODER_INFO_LAST',
-    'ROCPROFILER_THREAD_TRACE_DECODER_INFO_NONE',
-    'ROCPROFILER_THREAD_TRACE_DECODER_INFO_STITCH_INCOMPLETE',
-    'ROCPROFILER_THREAD_TRACE_DECODER_INFO_WAVE_INCOMPLETE',
-    'ROCPROFILER_THREAD_TRACE_DECODER_INST_BVH',
-    'ROCPROFILER_THREAD_TRACE_DECODER_INST_CONTEXT',
-    'ROCPROFILER_THREAD_TRACE_DECODER_INST_FLAT',
-    'ROCPROFILER_THREAD_TRACE_DECODER_INST_IMMED',
-    'ROCPROFILER_THREAD_TRACE_DECODER_INST_JUMP',
-    'ROCPROFILER_THREAD_TRACE_DECODER_INST_LAST',
-    'ROCPROFILER_THREAD_TRACE_DECODER_INST_LDS',
-    'ROCPROFILER_THREAD_TRACE_DECODER_INST_MESSAGE',
-    'ROCPROFILER_THREAD_TRACE_DECODER_INST_NEXT',
-    'ROCPROFILER_THREAD_TRACE_DECODER_INST_NONE',
-    'ROCPROFILER_THREAD_TRACE_DECODER_INST_SALU',
-    'ROCPROFILER_THREAD_TRACE_DECODER_INST_SMEM',
-    'ROCPROFILER_THREAD_TRACE_DECODER_INST_VALU',
-    'ROCPROFILER_THREAD_TRACE_DECODER_INST_VMEM',
-    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_DEBUG',
-    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_GFXIP',
-    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_INFO',
-    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_LAST',
-    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY',
-    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_PERFEVENT',
-    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_REALTIME',
-    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_RT_FREQUENCY',
-    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_SHADERDATA',
-    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE',
-    'ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_IMM',
-    'ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_PRIV',
-    'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR',
-    'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_ARGUMENT',
-    'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_SHADER_DATA',
-    'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES',
-    'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_LAST',
-    'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS',
-    'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EMPTY',
-    'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EXEC',
-    'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_IDLE',
-    'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_LAST',
-    'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_STALL',
-    'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_WAIT',
-    'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_COUNTER_INTERVAL',
-    'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_LAST',
-    'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_RT_FREQUENCY_KHZ',
-    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_HI',
-    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_LO',
-    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_HI',
-    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_LO',
-    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_LAST',
-    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_HI',
-    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_LO',
-    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_TAIL',
-    'ROCPROF_TRACE_DECODER_PACKET_OPCODE_AGENT_INFO',
-    'ROCPROF_TRACE_DECODER_PACKET_OPCODE_CODEOBJ',
-    'ROCPROF_TRACE_DECODER_PACKET_OPCODE_RT_TIMESTAMP',
-    'c__EA_rocprofiler_thread_trace_decoder_status_t',
-    'rocprof_trace_decoder_agent_info_type_t',
-    'rocprof_trace_decoder_codeobj_marker_tail_t',
-    'rocprof_trace_decoder_codeobj_marker_type_t',
-    'rocprof_trace_decoder_dump_data',
-    'rocprof_trace_decoder_get_info_string',
-    'rocprof_trace_decoder_get_status_string',
-    'rocprof_trace_decoder_gfx9_header_t',
-    'rocprof_trace_decoder_instrument_enable_t',
-    'rocprof_trace_decoder_isa_callback_t',
-    'rocprof_trace_decoder_packet_header_t',
-    'rocprof_trace_decoder_packet_opcode_t',
-    'rocprof_trace_decoder_parse_data',
-    'rocprof_trace_decoder_se_data_callback_t',
-    'rocprof_trace_decoder_trace_callback_t',
-    'rocprofiler_thread_trace_decoder_debug_callback_t',
-    'rocprofiler_thread_trace_decoder_info_t',
-    'rocprofiler_thread_trace_decoder_inst_category_t',
-    'rocprofiler_thread_trace_decoder_inst_t',
-    'rocprofiler_thread_trace_decoder_occupancy_t',
-    'rocprofiler_thread_trace_decoder_pc_t',
-    'rocprofiler_thread_trace_decoder_perfevent_t',
-    'rocprofiler_thread_trace_decoder_realtime_t',
-    'rocprofiler_thread_trace_decoder_record_type_t',
-    'rocprofiler_thread_trace_decoder_shaderdata_flags_t',
-    'rocprofiler_thread_trace_decoder_shaderdata_t',
-    'rocprofiler_thread_trace_decoder_status_t',
-    'rocprofiler_thread_trace_decoder_status_t__enumvalues',
-    'rocprofiler_thread_trace_decoder_wave_state_t',
-    'rocprofiler_thread_trace_decoder_wave_t',
-    'rocprofiler_thread_trace_decoder_wstate_type_t',
-    'struct_rocprof_trace_decoder_codeobj_marker_tail_t_0',
-    'struct_rocprof_trace_decoder_gfx9_header_t_0',
-    'struct_rocprof_trace_decoder_instrument_enable_t_0',
-    'struct_rocprof_trace_decoder_packet_header_t_0',
-    'struct_rocprofiler_thread_trace_decoder_inst_t',
-    'struct_rocprofiler_thread_trace_decoder_occupancy_t',
-    'struct_rocprofiler_thread_trace_decoder_pc_t',
-    'struct_rocprofiler_thread_trace_decoder_perfevent_t',
-    'struct_rocprofiler_thread_trace_decoder_realtime_t',
-    'struct_rocprofiler_thread_trace_decoder_shaderdata_t',
-    'struct_rocprofiler_thread_trace_decoder_wave_state_t',
-    'struct_rocprofiler_thread_trace_decoder_wave_t', 'uint64_t',
-    'union_rocprof_trace_decoder_codeobj_marker_tail_t',
-    'union_rocprof_trace_decoder_gfx9_header_t',
-    'union_rocprof_trace_decoder_instrument_enable_t',
-    'union_rocprof_trace_decoder_packet_header_t']
@@ -43,7 +43,6 @@ enum sqtt_version
   SQTT_VERSION_2_3 = 0x6, /* GFX9 */
   SQTT_VERSION_2_4 = 0x7, /* GFX10+ */
   SQTT_VERSION_3_2 = 0xb, /* GFX11+ */
-   SQTT_VERSION_3_3 = 0xc, /* GFX12+ */
 };

 enum sqtt_file_chunk_type
@@ -145,8 +144,6 @@ enum sqtt_gfxip_level
   SQTT_GFXIP_LEVEL_GFXIP_10_1 = 0x7,
   SQTT_GFXIP_LEVEL_GFXIP_10_3 = 0x9,
   SQTT_GFXIP_LEVEL_GFXIP_11_0 = 0xc,
-   SQTT_GFXIP_LEVEL_GFXIP_11_5 = 0xd,
-   SQTT_GFXIP_LEVEL_GFXIP_12 = 0x10,
 };

 enum sqtt_memory_type
@@ -430,8 +427,6 @@ enum elf_gfxip_level
   EF_AMDGPU_MACH_AMDGCN_GFX1010 = 0x033,
   EF_AMDGPU_MACH_AMDGCN_GFX1030 = 0x036,
   EF_AMDGPU_MACH_AMDGCN_GFX1100 = 0x041,
-   EF_AMDGPU_MACH_AMDGCN_GFX1150 = 0x043,
-   EF_AMDGPU_MACH_AMDGCN_GFX1200 = 0x04e,
 };

 struct sqtt_file_chunk_spm_db {
@@ -1,20 +0,0 @@
-from extra.optimization.helpers import load_worlds, ast_str_to_ast
-from tinygrad.helpers import tqdm
-from tinygrad.uop.ops import pyrender, UOp, Ops
-from tinygrad import dtypes
-from tinygrad.shape.shapetracker import ShapeTracker, View
-inf, nan = float('inf'), float('nan')
-
-if __name__ == "__main__":
-  ast_strs = load_worlds()
-  for i, ast_str in enumerate(tqdm(ast_strs)):
-    good_ast = ast_str_to_ast(ast_str)
-    code = '\n'.join(pyrender(good_ast))
-    print("\n***************\n\n"+code)
-    exec(code)
-    if str(good_ast) != str(ast):
-      print(code)
-      print("MISMATCH")
-      print(good_ast)
-      print(ast)
-      break
@@ -4,13 +4,13 @@ import struct
 import json
 import traceback
 import numpy as np
-from tinygrad.runtime.ops_cl import CLProgram, compile_gpu
+from tinygrad.runtime.ops_gpu import CLProgram, compile_gpu
 from tinygrad.device import Device
 from tinygrad.helpers import DEBUG, getenv
 from collections import defaultdict
 import pyopencl as cl
-from tinygrad.runtime.ops_cl import OSX_TIMING_RATIO
-CL = Device["CL"]
+from tinygrad.runtime.ops_gpu import OSX_TIMING_RATIO
+CL = Device["GPU"]

 DEBUGCL = getenv("DEBUGCL", 0)
 FLOAT16 = getenv("FLOAT16", 0)
@@ -110,7 +110,7 @@ class Thneed:
    prgs = {}
    for o in jdat['binaries']:
      nptr = ptr + o['length']
-      prgs[o['name']] = CLProgram(Device["CL"], o['name'], weights[ptr:nptr])
+      prgs[o['name']] = CLProgram(Device["GPU"], o['name'], weights[ptr:nptr])
      ptr = nptr

    # populate the cl_cache
@@ -267,7 +267,7 @@ class Thneed:
    for prg, args in self.cl_cache:
      events.append(prg.clprg(CL.queue, *args))
    mt = time.monotonic()
-    Device["CL"].synchronize()
+    Device["GPU"].synchronize()
    et = time.monotonic() - st
    print(f"submit in {(mt-st)*1000.0:.2f} ms, total runtime is {et*1000.0:.2f} ms")

@@ -1,400 +0,0 @@
-/**
- * @file
- * @brief Basic operations on generic types.
- */
-
-#pragma once
-
-#include <cuda_bf16.h>
-#include <limits>
-#include "base_types.cuh"
-
-namespace kittens {
-
-/**
- * @namespace base_ops
- *
- * @brief A namespace for operations on basic data types.
- */
-namespace base_ops {
-
-/* ----------  CONST OPS  ---------- */
-
-/**
- * @brief Represents the zero constant operation.
- *
- * This operation returns the zero value of the specified type.
- *
- * @tparam T The data type for which to return the zero value.
- * @return The zero value of type T.
- */
-struct zero {
-    template<typename T, typename... args> __device__ static inline constexpr T op(args... _) { return base_types::constants<T>::zero();      }
-};
-/**
- * @brief Represents the one constant operation.
- *
- * This operation returns the one value of the specified type.
- *
- * @tparam T The data type for which to return the one value.
- * @return The one value of type T.
- */
-struct one {
-    template<typename T, typename... args> __device__ static inline constexpr T op(args... _) { return base_types::constants<T>::one();       }
-};
-/**
- * @brief Represents the positive infinity constant operation.
- *
- * This operation returns the positive infinity value of the specified type.
- *
- * @tparam T The data type for which to return the positive infinity value.
- * @return The positive infinity value of type T.
- */
-struct pos_infty {
-    template<typename T, typename... args> __device__ static inline constexpr T op(args... _) { return base_types::constants<T>::pos_infty(); }
-};
-/**
- * @brief Represents the negative infinity constant operation.
- *
- * This operation returns the negative infinity value of the specified type.
- *
- * @tparam T The data type for which to return the negative infinity value.
- * @return The negative infinity value of type T.
- */
-struct neg_infty {
-    template<typename T, typename... args> __device__ static inline constexpr T op(args... _) { return base_types::constants<T>::neg_infty(); }
-};
-
-
-/* ----------  UNARY OPS  ---------- */
-
-/**
- * @brief Exponential function operation.
- *
- * This operation calculates the exponential of the input value.
- *
- * @tparam T The data type of the input and output values.
- * @param x[in] The input value.
- * @return The exponential of the input value.
- */
-struct exp {
-    template<typename T> static __device__ inline T op(const T &x) { return exp(x); }
-};
-template<> __device__ inline float  exp::op<float> (const float &x ) { return __expf(x);                        }
-template<> __device__ inline float2 exp::op<float2>(const float2 &x) { return float2{__expf(x.x), __expf(x.y)}; }
-template<> __device__ inline bf16   exp::op<bf16>  (const bf16 &x  ) { return hexp(x);                          }
-template<> __device__ inline bf16_2 exp::op<bf16_2>(const bf16_2 &x) { return h2exp(x);                         }
-template<> __device__ inline half   exp::op<half>  (const half &x  ) { return hexp(x);                          }
-template<> __device__ inline half_2 exp::op<half_2>(const half_2 &x) { return h2exp(x);                         }
-
-/**
- * @brief Exponential function operation, in base 2
- *
- * This operation calculates the exponential of the input value, in base 2.
- *
- * @tparam T The data type of the input and output values.
- * @param x[in] The input value.
- * @return The exponential of the input value.
- */
-struct exp2 {
-    template<typename T> static __device__ inline T op(const T &x) { return exp2f(x); }
-};
-template<> __device__ inline float  exp2::op<float> (const float &x ) { return exp2f(x);                        }
-template<> __device__ inline float2 exp2::op<float2>(const float2 &x) { return float2{exp2f(x.x), exp2f(x.y)}; }
-template<> __device__ inline bf16   exp2::op<bf16>  (const bf16 &x  ) { return hexp2(x);                          }
-template<> __device__ inline bf16_2 exp2::op<bf16_2>(const bf16_2 &x) { return h2exp2(x);                         }
-template<> __device__ inline half   exp2::op<half>  (const half &x  ) { return hexp2(x);                          }
-template<> __device__ inline half_2 exp2::op<half_2>(const half_2 &x) { return h2exp2(x);                         }
-/**
- * @brief Natural log function operation.
- *
- * This operation calculates the natural logarithm of the input value.
- *
- * @tparam T The data type of the input and output values.
- * @param x[in] The input value.
- * @return The natural logarithm of the input value.
- */
-struct log {
-    template<typename T> static __device__ inline T op(const T &x) { return log(x); }
-};
-template<> __device__ inline float  log::op<float> (const float &x ) { return __logf(x);                        }
-template<> __device__ inline float2 log::op<float2>(const float2 &x) { return float2{__logf(x.x), __logf(x.y)}; }
-template<> __device__ inline bf16   log::op<bf16>  (const bf16 &x  ) { return hlog(x);                          }
-template<> __device__ inline bf16_2 log::op<bf16_2>(const bf16_2 &x) { return h2log(x);                         }
-template<> __device__ inline half   log::op<half>  (const half &x  ) { return hlog(x);                          }
-template<> __device__ inline half_2 log::op<half_2>(const half_2 &x) { return h2log(x);                         }
-/**
- * @brief Logarithm base 2 operation.
- *
- * This operation calculates the logarithm base 2 of the input value.
- *
- * @tparam T The data type of the input and output values.
- * @param x[in] The input value.
- * @return The logarithm base 2 of the input value.
- */
-struct log2 {
-    template<typename T> static __device__ inline T op(const T &x) { return log2(x); }
-};
-template<> __device__ inline float  log2::op<float> (const float &x ) { return __log2f(x);                        }
-template<> __device__ inline float2 log2::op<float2>(const float2 &x) { return float2{__log2f(x.x), __log2f(x.y)}; }
-template<> __device__ inline bf16   log2::op<bf16>  (const bf16 &x  ) { return hlog2(x);                          }
-template<> __device__ inline bf16_2 log2::op<bf16_2>(const bf16_2 &x) { return h2log2(x);                         }
-template<> __device__ inline half   log2::op<half>  (const half &x  ) { return hlog2(x);                          }
-template<> __device__ inline half_2 log2::op<half_2>(const half_2 &x) { return h2log2(x);                         }
-/**
- * @brief Absolute value operation.
- *
- * This operation calculates the absolute value of the input.
- *
- * @tparam T The data type of the input and output values.
- * @param x[in] The input value.
- * @return The absolute value of the input.
- */
-struct abs {
-    template<typename T> static __device__ inline T op(const T &x) { return abs(x); }
-};
-template<> __device__ inline float  abs::op<float> (const float &x ) { return fabsf(x);                       }
-template<> __device__ inline float2 abs::op<float2>(const float2 &x) { return float2{fabsf(x.x), fabsf(x.y)}; }
-template<> __device__ inline bf16   abs::op<bf16>  (const bf16 &x  ) { return __habs(x);                      }
-template<> __device__ inline bf16_2 abs::op<bf16_2>(const bf16_2 &x) { return __habs2(x);                     }
-template<> __device__ inline half   abs::op<half>  (const half &x  ) { return __habs(x);                      }
-template<> __device__ inline half_2 abs::op<half_2>(const half_2 &x) { return __habs2(x);                     }
-/**
- * @brief Rectified Linear Unit (ReLU) operation.
- *
- * This operation applies the ReLU function to the input, which is the
- * maximum of zero and the input value.
- *
- * @tparam T The data type of the input and output values.
- * @param x[in] The input value.
- * @return The result of ReLU function applied to the input.
- */
-struct relu {
-    template<typename T> static __device__ inline T op(const T &x) { return max(x, base_types::constants<T>::zero()); }
-};
-template<> __device__ inline float  relu::op<float> (const float &x ) { return max(x, 0.f);                                  }
-template<> __device__ inline float2 relu::op<float2>(const float2 &x) { return float2{max(x.x, 0.f), max(x.y, 0.f)};         }
-template<> __device__ inline bf16   relu::op<bf16>  (const bf16 &x  ) { return __hmax(x, base_types::constants<bf16>::zero());    }
-template<> __device__ inline bf16_2 relu::op<bf16_2>(const bf16_2 &x) { return __hmax2(x, base_types::constants<bf16_2>::zero()); }
-template<> __device__ inline half   relu::op<half>  (const half &x  ) { return __hmax(x, base_types::constants<half>::zero());    }
-template<> __device__ inline half_2 relu::op<half_2>(const half_2 &x) { return __hmax2(x, base_types::constants<half_2>::zero()); }
-/**
- * @brief Copy operation.
- *
- * This operation returns the input value unchanged.
- *
- * @tparam T The data type of the input and output values.
- * @param a[in] The input value.
- * @return The same value as the input.
- */
-struct copy { // for non-compile-time setters.
-    template<typename T> static __device__ inline T op(const T &a) { return a; }
-};
-
-
-/* ----------  BINARY OPS  ---------- */
-
-/**
- * @brief Copy2 operation.
- *
- * This operation returns the second input value unchanged.
- *
- * @tparam T The data type of the input and output values.
- * @param a[in] The first input value (ignored).
- * @param b[in] The second input value.
- * @return The same value as the second input.
- */
-struct copy2 { // this turns out to be a slightly hacky op that makes some code cleaner :/
-    template<typename T> static __device__ inline T op(const T &a, const T &b) { return b; }
-};
-/**
- * @brief Sum operation.
- *
- * This operation calculates the sum of two input values.
- *
- * @tparam T The data type of the input and output values.
- * @param a[in] The first input value.
- * @param b[in] The second input value.
- * @return The sum of the input values.
- */
-struct sum {
-    template<typename T> static __device__ inline T op(const T &a, const T &b) { return a+b; }
-};
-template<> __device__ inline float2 sum::op<float2>(const float2 &a, const float2 &b) {
-#ifdef KITTENS_BLACKWELL
-    float2 c;
-    asm volatile("add.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b));
-    return c;
-#else
-    return float2{a.x+b.x, a.y+b.y};
-#endif
-}
-template<> __device__ inline bf16   sum::op<bf16>  (const bf16   &a, const bf16   &b) { return __hadd(a, b);             }
-template<> __device__ inline bf16_2 sum::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hadd2(a, b);            }
-template<> __device__ inline half   sum::op<half>  (const half   &a, const half   &b) { return __hadd(a, b);             }
-template<> __device__ inline half_2 sum::op<half_2>(const half_2 &a, const half_2 &b) { return __hadd2(a, b);            }
-/**
- * @brief Subtraction operation.
- *
- * This operation calculates the difference between two input values.
- *
- * @tparam T The data type of the input and output values.
- * @param a[in] The first input value.
- * @param b[in] The second input value.
- * @return The difference between the input values.
- */
-struct sub {
-    template<typename T> static __device__ inline T op(const T &a, const T &b) { return a-b; }
-};
-template<> __device__ inline float2 sub::op<float2>(const float2 &a, const float2 &b) { 
-#ifdef KITTENS_BLACKWELL
-    float2 c;
-    asm volatile("sub.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b));
-    return c;
-#else
-    return float2{a.x-b.x, a.y-b.y}; 
-#endif
-}
-template<> __device__ inline bf16   sub::op<bf16>  (const bf16   &a, const bf16   &b) { return __hsub(a, b);             }
-template<> __device__ inline bf16_2 sub::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hsub2(a, b);            }
-template<> __device__ inline half   sub::op<half>  (const half   &a, const half   &b) { return __hsub(a, b);             }
-template<> __device__ inline half_2 sub::op<half_2>(const half_2 &a, const half_2 &b) { return __hsub2(a, b);            }
-/**
- * @brief Multiplication operation.
- *
- * This operation calculates the product of two input values.
- *
- * @tparam T The data type of the input and output values.
- * @param a[in] The first input value.
- * @param b[in] The second input value.
- * @return The product of the input values.
- */
-struct mul {
-    template<typename T> static __device__ inline T op(const T &a, const T &b) { return a*b; }
-};
-template<> __device__ inline float2 mul::op<float2>(const float2 &a, const float2 &b) { 
-#ifdef KITTENS_BLACKWELL
-    float2 c;
-    asm volatile("mul.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b));
-    return c;
-#else
-    return float2{a.x*b.x, a.y*b.y}; 
-#endif
-}
-template<> __device__ inline bf16   mul::op<bf16>  (const bf16   &a, const bf16   &b) { return __hmul(a, b);             }
-template<> __device__ inline bf16_2 mul::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hmul2(a, b);            }
-template<> __device__ inline half   mul::op<half>  (const half   &a, const half   &b) { return __hmul(a, b);             }
-template<> __device__ inline half_2 mul::op<half_2>(const half_2 &a, const half_2 &b) { return __hmul2(a, b);            }
-/**
- * @brief Division operation.
- *
- * This operation calculates the quotient of two input values.
- *
- * @tparam T The data type of the input and output values.
- * @param a[in] The first input value.
- * @param b[in] The second input value.
- * @return The quotient of the input values.
- */
-struct div {
-    template<typename T> static __device__ inline T op(const T &a, const T &b) { return a/b; }
-};
-template<> __device__ inline float2 div::op<float2>(const float2 &a, const float2 &b) { return float2{a.x/b.x, a.y/b.y}; }
-template<> __device__ inline bf16   div::op<bf16>  (const bf16   &a, const bf16   &b) { return __hdiv(a, b);             }
-template<> __device__ inline bf16_2 div::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __h2div(a, b);            } // this op is a special snowflake
-template<> __device__ inline half   div::op<half>  (const half   &a, const half   &b) { return __hdiv(a, b);             }
-template<> __device__ inline half_2 div::op<half_2>(const half_2 &a, const half_2 &b) { return __h2div(a, b);            }
-/**
- * @brief Maximum operation.
- *
- * This operation calculates the maximum of two input values.
- *
- * @tparam T The data type of the input and output values.
- * @param a[in] The first input value.
- * @param b[in] The second input value.
- * @return The maximum of the input values.
- */
- struct max {
-    template<typename T> static __device__ inline T op(const T &a, const T &b) { return ::max(a, b); }
-};
-template<>  __device__ inline float2 max::op<float2>(const float2 &a, const float2 &b) { return float2{::max(a.x, b.x), ::max(a.y, b.y)}; }
-template<>  __device__ inline bf16   max::op<bf16>  (const bf16   &a, const bf16   &b) { return __hmax(a, b);                             }
-template<>  __device__ inline bf16_2 max::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hmax2(a, b);                            }
-template<>  __device__ inline half   max::op<half>  (const half   &a, const half   &b) { return __hmax(a, b);                             }
-template<>  __device__ inline half_2 max::op<half_2>(const half_2 &a, const half_2 &b) { return __hmax2(a, b);                            }
-/**
- * @brief Minimum operation.
- *
- * This operation calculates the minimum of two input values.
- *
- * @tparam T The data type of the input and output values.
- * @param a[in] The first input value.
- * @param b[in] The second input value.
- * @return The minimum of the input values.
- */
-struct min {
-    template<typename T> static __device__ inline T op(const T &a, const T &b) { return ::min(a, b); }
-};
-template<>  __device__ inline float2 min::op<float2>(const float2 &a, const float2 &b) { return float2{::min(a.x, b.x), ::min(a.y, b.y)}; }
-template<>  __device__ inline bf16   min::op<bf16>  (const bf16   &a, const bf16   &b) { return __hmin(a, b);                         }
-template<>  __device__ inline bf16_2 min::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hmin2(a, b);                        }
-template<>  __device__ inline half   min::op<half>  (const half   &a, const half   &b) { return __hmin(a, b);                         }
-template<>  __device__ inline half_2 min::op<half_2>(const half_2 &a, const half_2 &b) { return __hmin2(a, b);                        }
-
-
-/* ----------  TERNARY OPS  ---------- */
-
-/**
- * @brief Fused multiply-add operation A * B + C.
- *
- * This operation performs a fused multiply-add, computing (A * B) + C with only one rounding.
- *
- * @tparam T The data type of the input and output values.
- * @param a[in] The first input value.
- * @param b[in] The second input value.
- * @param c[in] The third input value to be added.
- * @return The result of the fused multiply-add operation.
- */
-struct fma_AxBtC {
-    template<typename T> static __device__ inline T op(const T &a, const T &b, const T &c) {
-        return sum::op<T>(mul::op<T>(a, b), c);
-    }
-};
-template<> __device__ inline float2 fma_AxBtC::op<float2>(const float2 &a, const float2 &b, const float2 &c) {
-#ifdef KITTENS_BLACKWELL
-    float2 d;
-    asm volatile("fma.rn.f32x2 %0, %1, %2, %3;" : "=l"(*(uint64_t*)&d) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b), "l"(*(uint64_t*)&c));
-    return d;
-#else
-    return float2{a.x*b.x+c.x, a.y*b.y+c.y};
-#endif
-}
-/**
- * @brief Fused multiply-add operation A * C + B.
- *
- * This operation performs a fused multiply-add, computing (A * C) + B with only one rounding.
- * This is particularly useful for attention mechanisms in neural networks.
- *
- * @tparam T The data type of the input and output values.
- * @param a[in] The first input value.
- * @param b[in] The third input value to be added.
- * @param c[in] The second input value.
- * @return The result of the fused multiply-add operation.
- */
-struct fma_AxCtB { // this is the one needed for attention
-    template<typename T> static __device__ inline T op(const T &a, const T &b, const T &c) {
-        return sum::op<T>(mul::op<T>(a, c), b);
-    }
-};
-template<> __device__ inline float2 fma_AxCtB::op<float2>(const float2 &a, const float2 &b, const float2 &c) {
-#ifdef KITTENS_BLACKWELL
-    float2 d;
-    asm volatile("fma.rn.f32x2 %0, %1, %2, %3;" : "=l"(*(uint64_t*)&d) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&c), "l"(*(uint64_t*)&b));
-    return d;
-#else
-    return float2{a.x*c.x+b.x, a.y*c.y+b.y};
-#endif
-}
-
-} // namespace base_ops
-
-} // namespace kittens
@@ -1,519 +0,0 @@
-/**
- * @file
- * @brief Declarations, manipulations, and wrappers for basic types.
- * 
- * This file is a bunch of utilities for going back and forth between different types.
- * 
- * Many of them are for the compiler, so as to clean up the code. It unfortunately
- * seems necessary when we have types we really care about that are less than word width.
- */
-
-#pragma once
-
-#ifdef KITTENS_HOPPER
-#include <cuda_fp8.h>
-#endif
-
-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
-#include <string>
-#include <bit>
-
-
-namespace kittens {
-
-/**
- * @brief Bfloat16 floating-point type.
- */
-using bf16 = __nv_bfloat16;
-/**
- * @brief Half-precision floating-point type.
- */
-using half = __half;
-/**
- * @brief Packed word of two bfloat16 floating-point values.
- */
-using bf16_2 = __nv_bfloat162;
-/**
- * @brief Packed word of two half-precision floating-point values.
- */
-using half_2 = __half2;
-#ifdef KITTENS_HOPPER
-/**
- * @brief float8 floating-point type.
- */
-using fp8e4m3 = __nv_fp8_e4m3;
-using fp8e5m2 = __nv_fp8_e5m2;
-#ifdef KITTENS_BLACKWELL
-using fp8e8m0 = __nv_fp8_e8m0;
-#endif
-/**
- * @brief 2-packed float8 floating-point type.
- */
-using fp8e4m3_2 = __nv_fp8x2_e4m3;
-using fp8e5m2_2 = __nv_fp8x2_e5m2;
-#ifdef KITTENS_BLACKWELL
-using fp8e8m0_2 = __nv_fp8x2_e8m0;
-#endif
-/**
- * @brief 4-packed float8 floating-point type.
- */
-using fp8e4m3_4 = __nv_fp8x4_e4m3;
-using fp8e5m2_4 = __nv_fp8x4_e5m2;
-#ifdef KITTENS_BLACKWELL
-using fp8e8m0_4 = __nv_fp8x4_e8m0;
-#endif
-#endif
-
-namespace ducks {
-/**
- * @namespace base_types
- *
- * @brief A namespace for concepts for basic data types.
- */
-namespace base_types {
-
-#ifdef KITTENS_HOPPER
-#ifdef KITTENS_BLACKWELL
-template<typename T>
-concept T2 = std::is_same_v<T, float2> || std::is_same_v<T, bf16_2> || std::is_same_v<T, half_2> || std::is_same_v<T, fp8e4m3_4> || std::is_same_v<T, fp8e5m2_4> || std::is_same_v<T, fp8e8m0_4>; // could add half_2 later if implemented.
-template<typename T>
-concept T1 = std::is_same_v<T, float>  || std::is_same_v<T, bf16  > || std::is_same_v<T, half> || std::is_same_v<T, fp8e4m3> || std::is_same_v<T, fp8e5m2> || std::is_same_v<T, fp8e8m0>; // could add half_2 later if implemented.
-#else
-template<typename T>
-concept T2 = std::is_same_v<T, float2> || std::is_same_v<T, bf16_2> || std::is_same_v<T, half_2> || std::is_same_v<T, fp8e4m3_4> || std::is_same_v<T, fp8e5m2_4>;
-template<typename T>
-concept T1 = std::is_same_v<T, float>  || std::is_same_v<T, bf16  > || std::is_same_v<T, half> || std::is_same_v<T, fp8e4m3> || std::is_same_v<T, fp8e5m2>;
-#endif
-#else
-template<typename T>
-concept T2 = std::is_same_v<T, float2> || std::is_same_v<T, bf16_2> || std::is_same_v<T, half_2>;
-template<typename T>
-concept T1 = std::is_same_v<T, float>  || std::is_same_v<T, bf16  > || std::is_same_v<T, half>;
-#endif
-
-} // namespace base_types
-} // namespace ducks
-
-/**
- * @namespace base_types
- *
- * @brief A namespace for ThunderKittens basic data types.
- */
-namespace base_types {
-
-/**
- * @brief Provides compile-time constants for different types.
- *
- * @tparam T The type for which to provide constants.
- */
-template<typename T> struct constants {
-    /**
-     * @brief Zero
-     * @return Constexpr zero with type T
-     */
-    static __device__ inline constexpr T zero()      { return T{0}; }
-    /**
-     * @brief One
-     * @return Constexpr one with type T
-     */
-    static __device__ inline constexpr T one()       { return T{1}; }
-    /**
-     * @brief Positive infinity. Particularly useful for initializing before a min op.
-     * @return Constexpr positive infinity with type T
-     */
-    static __device__ inline constexpr T pos_infty() { return T{INFINITY}; } // I'll find a better way at some point but this appears to work.
-    /**
-     * @brief Negative infinity. Particularly useful for initializing before a max op.
-     * @return Constexpr negative infinity with type T
-     */
-    static __device__ inline constexpr T neg_infty() { return T{-INFINITY}; }
-};
-template<> struct constants<float2> {
-    static __device__ inline constexpr float2 zero()      { return float2{0.f, 0.f}; }
-    static __device__ inline constexpr float2 one()       { return float2{1.f, 1.f}; }
-    static __device__ inline constexpr float2 pos_infty() { return float2{constants<float>::pos_infty(), constants<float>::pos_infty()}; }
-    static __device__ inline constexpr float2 neg_infty() { return float2{constants<float>::neg_infty(), constants<float>::neg_infty()}; }
-};
-template<> struct constants<bf16> {
-    static __device__ inline constexpr bf16 zero()      { return std::bit_cast<__nv_bfloat16>(uint16_t(0x0000)); } // unfortunately __float2bf16_rn is not constexpr
-    static __device__ inline constexpr bf16 one()       { return std::bit_cast<__nv_bfloat16>(uint16_t(0x3F80)); }
-    static __device__ inline constexpr bf16 pos_infty() { return std::bit_cast<__nv_bfloat16>(uint16_t(0x7F80)); }
-    static __device__ inline constexpr bf16 neg_infty() { return std::bit_cast<__nv_bfloat16>(uint16_t(0xFF80)); }
-};
-template<> struct constants<bf16_2> {
-    static __device__ inline constexpr bf16_2 zero()      { return bf16_2{constants<bf16>::zero(),      constants<bf16>::zero()};      }
-    static __device__ inline constexpr bf16_2 one()       { return bf16_2{constants<bf16>::one(),       constants<bf16>::one()};       }
-    static __device__ inline constexpr bf16_2 pos_infty() { return bf16_2{constants<bf16>::pos_infty(), constants<bf16>::pos_infty()}; }
-    static __device__ inline constexpr bf16_2 neg_infty() { return bf16_2{constants<bf16>::neg_infty(), constants<bf16>::neg_infty()}; }
-};
-template<> struct constants<half> {
-    static __device__ inline constexpr half zero()      { return std::bit_cast<__half>(uint16_t(0x0000)); }
-    static __device__ inline constexpr half one()       { return std::bit_cast<__half>(uint16_t(0x3C00)); }
-    static __device__ inline constexpr half pos_infty() { return std::bit_cast<__half>(uint16_t(0x7C00)); }
-    static __device__ inline constexpr half neg_infty() { return std::bit_cast<__half>(uint16_t(0xFC00)); }
-};
-template<> struct constants<half_2> {
-    static __device__ inline constexpr half_2 zero()      { return half_2{constants<half>::zero(),      constants<half>::zero()};      }
-    static __device__ inline constexpr half_2 one()       { return half_2{constants<half>::one(),       constants<half>::one()};       }
-    static __device__ inline constexpr half_2 pos_infty() { return half_2{constants<half>::pos_infty(), constants<half>::pos_infty()}; }
-    static __device__ inline constexpr half_2 neg_infty() { return half_2{constants<half>::neg_infty(), constants<half>::neg_infty()}; }
-};
-#ifdef KITTENS_HOPPER
-template<> struct constants<fp8e4m3> {
-    static __device__ inline constexpr fp8e4m3 zero() { return std::bit_cast<__nv_fp8_e4m3>(uint8_t(0x00)); }
-    static __device__ inline constexpr fp8e4m3 one() { return std::bit_cast<__nv_fp8_e4m3>(uint8_t(0x38)); }
-};
-template<> struct constants<fp8e4m3_2> {
-    static __device__ inline constexpr fp8e4m3_2 zero() { return std::bit_cast<fp8e4m3_2>(uint16_t(0x0000)); }
-    static __device__ inline constexpr fp8e4m3_2 one() { return std::bit_cast<fp8e4m3_2>(uint16_t(0x3838)); }
-};
-template<> struct constants<fp8e4m3_4> {
-    static __device__ inline constexpr fp8e4m3_4 zero() { return std::bit_cast<fp8e4m3_4>(uint32_t(0x00000000)); }
-    static __device__ inline constexpr fp8e4m3_4 one() { return std::bit_cast<fp8e4m3_4>(uint32_t(0x38383838)); }
-};
-template<> struct constants<fp8e5m2> {
-    static __device__ inline constexpr fp8e5m2 zero() { return std::bit_cast<__nv_fp8_e5m2>(uint8_t(0x00)); }
-    static __device__ inline constexpr fp8e5m2 one() { return std::bit_cast<__nv_fp8_e5m2>(uint8_t(0x3C)); }
-};
-template<> struct constants<fp8e5m2_2> {
-    static __device__ inline constexpr fp8e5m2_2 zero() { return std::bit_cast<fp8e5m2_2>(uint16_t(0x0000)); }
-    static __device__ inline constexpr fp8e5m2_2 one() { return std::bit_cast<fp8e5m2_2>(uint16_t(0x3C3C)); }
-};
-template<> struct constants<fp8e5m2_4> {
-    static __device__ inline constexpr fp8e5m2_4 zero() { return std::bit_cast<fp8e5m2_4>(uint32_t(0x00000000)); }
-    static __device__ inline constexpr fp8e5m2_4 one() { return std::bit_cast<fp8e5m2_4>(uint32_t(0x3C3C3C3C)); }
-};
-#endif
-
-template<> struct constants<int> {
-    static __device__ inline constexpr int zero()      { return 0; }
-    static __device__ inline constexpr int one()       { return 1; }
-};
-template<> struct constants<int2> {
-    static __device__ inline constexpr int2 zero()      { return int2{0, 0}; }
-    static __device__ inline constexpr int2 one()       { return int2{1, 1}; }
-};
-
-/**
- * @brief Provides information about packing of elements for a given type.
- *
- * @tparam T The type for which to provide packing information.
- */
-template<typename T> struct packing {
-    /**
-     * @brief The number of elements packed together.
-     *
-     * @return constexpr int representing number of elements within the type.
-     */
-    static __device__ inline constexpr int num() { return 1; }
-    /**
-     * @brief Packs a single T element twice (replicated) into its packed type.
-     *
-     * @param i[in] The element to pack.
-     * @return The packed type.
-     */
-    static __device__ inline constexpr T pack(const bf16 &i);
-};
-template<> struct packing<bf16> {
-    static __device__ inline constexpr int num() { return 1; }
-    using unpacked_type = bf16;
-    using packed_type = bf16_2;
-    static __device__ inline constexpr bf16_2 pack(const bf16 &i) { return bf16_2{i, i}; }
-};
-template<> struct packing<bf16_2> {
-    static __device__ inline constexpr int num() { return 2; }
-    using unpacked_type = bf16;
-    using packed_type = bf16_2;
-    static __device__ inline constexpr bf16_2 pack(const bf16 &i) { return bf16_2{i, i}; } // this replication makes code cleaner later.
-};
-template<> struct packing<half> {
-    static __device__ inline constexpr int num() { return 1; }
-    using unpacked_type = half;
-    using packed_type = half_2;
-    static __device__ inline constexpr half_2 pack(const half &i) { return half_2{i, i}; }
-};
-template<> struct packing<half_2> {
-    static __device__ inline constexpr int num() { return 2; }
-    using unpacked_type = half;
-    using packed_type = half_2;
-    static __device__ inline constexpr half_2 pack(const half &i) { return half_2{i, i}; } // this replication makes code cleaner later.
-};
-template<> struct packing<float> {
-    static __device__ inline constexpr int num() { return 1; }
-    using unpacked_type = float;
-    using packed_type = float2;
-    static __device__ inline constexpr float2 pack(const float &i) { return float2{i, i}; }
-};
-template<> struct packing<float2> {
-    static __device__ inline constexpr int num() { return 2; }
-    using unpacked_type = float;
-    using packed_type = float2;
-    static __device__ inline constexpr float2 pack(const float &i) { return float2{i, i}; } // this replication makes code cleaner later.
-};
-template<> struct packing<char> {
-    static __device__ inline constexpr int num() { return 1; }
-    using unpacked_type = char;
-    using packed_type = char2;
-    static __device__ inline constexpr char2 pack(const char &i) { return char2{i, i}; } // this replication makes code cleaner later.
-};
-template<> struct packing<char2> {
-    static __device__ inline constexpr int num() { return 2; }
-    using unpacked_type = char;
-    using packed_type = char2;
-    static __device__ inline constexpr char2 pack(const char &i) { return char2{i, i}; } // this replication makes code cleaner later.
-};
-template<> struct packing<int> {
-    static __device__ inline constexpr int num() { return 1; }
-    using unpacked_type = int;
-    using packed_type = int2;
-    static __device__ inline constexpr int2 pack(const int &i) { return int2{i, i}; } // this replication makes code cleaner later.
-};
-template<> struct packing<int2> {
-    static __device__ inline constexpr int num() { return 2; }
-    using unpacked_type = int;
-    using packed_type = int2;
-    static __device__ inline constexpr int2 pack(const int &i) { return int2{i, i}; } // this replication makes code cleaner later.
-};
-template<> struct packing<uint> {
-    static __device__ inline constexpr int num() { return 1; }
-    using unpacked_type = uint;
-    using packed_type = uint2;
-    static __device__ inline constexpr uint2 pack(const uint &i) { return uint2{i, i}; } // this replication makes code cleaner later.
-};
-template<> struct packing<uint2> {
-    static __device__ inline constexpr int num() { return 2; }
-    using unpacked_type = uint;
-    using packed_type = uint2;
-    static __device__ inline constexpr uint2 pack(const uint &i) { return uint2{i, i}; } // this replication makes code cleaner later.
-};
-struct uint64_2 { uint64_t x, y; };
-template<> struct packing<uint64_t> {
-    static __device__ inline constexpr int num() { return 1; }
-    using unpacked_type = uint64_t;
-    using packed_type = uint64_2;
-    static __device__ inline constexpr uint64_2 pack(const uint64_t &i) { return uint64_2{i, i}; } // this replication makes code cleaner later.
-};
-template<> struct packing<uint64_2> {
-    static __device__ inline constexpr int num() { return 2; }
-    using unpacked_type = uint64_t;
-    using packed_type = uint64_2;
-    static __device__ inline constexpr uint64_2 pack(const uint64_t &i) { return uint64_2{i, i}; } // this replication makes code cleaner later.
-};
-template<> struct packing<float4> {
-    static __device__ inline constexpr int num() { return 4; }
-};
-template<> struct packing<int4> {
-    static __device__ inline constexpr int num() { return 4; }
-};
-#ifdef KITTENS_HOPPER
-template<> struct packing<fp8e4m3> {
-    static __device__ inline constexpr int num() { return 1; }
-    using unpacked_type = fp8e4m3;
-    using packed_type = fp8e4m3_4;
-};
-template<> struct packing<fp8e4m3_4> {
-    static __device__ inline constexpr int num() { return 4; }
-    using unpacked_type = fp8e4m3;
-    using packed_type = fp8e4m3_4;
-};
-template<> struct packing<fp8e5m2> {
-    static __device__ inline constexpr int num() { return 1; }
-    using unpacked_type = fp8e5m2;
-    using packed_type = fp8e5m2_4;
-};
-template<> struct packing<fp8e5m2_4> {
-    static __device__ inline constexpr int num() { return 4; }
-    using unpacked_type = fp8e5m2;
-    using packed_type = fp8e5m2_4;
-};
-#ifdef KITTENS_BLACKWELL
-template<> struct packing<fp8e8m0> {
-    static __device__ inline constexpr int num() { return 1; }
-    using unpacked_type = fp8e8m0;
-    using packed_type = fp8e8m0_4;
-};
-template<> struct packing<fp8e8m0_4> {
-    static __device__ inline constexpr int num() { return 4; }
-    using unpacked_type = fp8e8m0;
-    using packed_type = fp8e8m0_4;
-};
-#endif
-#endif
-
-
-/**
- * @brief Provides templated functionality to convert between different types.
- *
- * @tparam T The target type for conversion.
- * @tparam U The source type for conversion.
- */
-template<typename T, typename U> struct convertor {
-    /**
-     * @brief Converts a value of type U to type T.
-     *
-     * @param u[in] The value of type U to convert.
-     * @return T The converted value of type T.
-     */
-    static __host__ __device__ inline T convert(const U & u) {
-        return (T)u;
-    }
-};
-template<> struct convertor<float, bf16> {
-    static __host__ __device__ inline float convert(const bf16 & u) {
-        return 	__bfloat162float(u);
-    }
-};
-template<> struct convertor<bf16, float> {
-    static __host__ __device__ inline bf16 convert(const float & u) {
-        return 	__float2bfloat16_rn(u);
-    }
-};
-template<> struct convertor<float2, bf16_2> {
-    static __host__ __device__ inline float2 convert(const bf16_2 & u) {
-        return 	__bfloat1622float2(u);
-    }
-};
-template<> struct convertor<bf16_2, float2> {
-    static __host__ __device__ inline bf16_2 convert(const float2 & u) {
-        return 	__float22bfloat162_rn(u);
-    }
-};
-template<> struct convertor<float, half> {
-    static __host__ __device__ inline float convert(const half & u) {
-        return __half2float(u);
-    }
-};
-template<> struct convertor<half, float> {
-    static __host__ __device__ inline half convert(const float & u) {
-        return __float2half(u);
-    }
-};
-template<> struct convertor<float2, half_2> {
-    static __host__ __device__ inline float2 convert(const half_2 & u) {
-        return __half22float2(u);
-    }
-};
-template<> struct convertor<half_2, float2> {
-    static __host__ __device__ inline half_2 convert(const float2 & u) {
-        return __float22half2_rn(u);
-    }
-};
-template<> struct convertor<bf16, half> {
-    static __host__ __device__ inline bf16 convert(const half & u) {
-        return __float2bfloat16_rn(__half2float(u));
-    }
-};
-template<> struct convertor<half, bf16> {
-    static __host__ __device__ inline half convert(const bf16 & u) {
-        return __float2half(__bfloat162float(u));
-    }
-};
-template<> struct convertor<bf16_2, half_2> {
-    static __host__ __device__ inline bf16_2 convert(const half_2 & u) {
-        return __float22bfloat162_rn(__half22float2(u));
-    }
-};
-template<> struct convertor<half_2, bf16_2> {
-    static __host__ __device__ inline half_2 convert(const bf16_2 & u) {
-        return __float22half2_rn(__bfloat1622float2(u));
-    }
-};
-#ifdef KITTENS_HOPPER
-// fp8e4m3
-template<> struct convertor<fp8e4m3_4, float4> {
-    static __host__ __device__ inline fp8e4m3_4 convert(const float4& u) {
-        return __nv_fp8x4_e4m3(u); 
-    }
-};
-template<> struct convertor<float4, fp8e4m3_4> {
-    static __host__ __device__ inline float4 convert(const fp8e4m3_4& u) {
-        __nv_fp8_e4m3 *vals = reinterpret_cast<__nv_fp8_e4m3*>(const_cast<__nv_fp8x4_e4m3*>(&u));
-        return make_float4(float(vals[0]), float(vals[1]), float(vals[2]), float(vals[3]));
-    }
-};
-template<> struct convertor<fp8e4m3_2, float2> {
-    static __host__ __device__ inline fp8e4m3_2 convert(const float2& u) {
-        return __nv_fp8x2_e4m3(u); 
-    }
-};
-template<> struct convertor<float2, fp8e4m3_2> {
-    static __host__ __device__ inline float2 convert(const fp8e4m3_2& u) {
-        __nv_fp8_e4m3 *vals = reinterpret_cast<__nv_fp8_e4m3*>(const_cast<__nv_fp8x2_e4m3*>(&u));
-        return make_float2(float(vals[0]), float(vals[1]));
-    }
-};
-template<> struct convertor<fp8e4m3, float> {
-    static __host__ __device__ inline fp8e4m3 convert(const float & u) {
-        return __nv_fp8_e4m3(u);
-    }
-};
-template<> struct convertor<float, fp8e4m3> {
-    static __host__ __device__ inline float convert(const fp8e4m3 & u) {
-        return float(u);
-    }
-};
-template<> struct convertor<bf16_2, fp8e4m3_4> {
-    static __host__ __device__ inline bf16_2 convert(const fp8e4m3_4 & u) {
-        float4 f4 = convertor<float4, fp8e4m3_4>::convert(u);
-        float2 f2 = make_float2(f4.x, f4.y);
-        return __float22bfloat162_rn(f2);
-    }
-};
-template<> struct convertor<fp8e4m3_4, bf16_2> {
-    static __host__ __device__ inline fp8e4m3_4 convert(const bf16_2 & u) {
-        float2 f2 = __bfloat1622float2(u);
-        float4 f4 = make_float4(f2.x, f2.y, 0.0f, 0.0f);
-        return __nv_fp8x4_e4m3(f4);
-    }
-};
-// fp8e5m2
-template<> struct convertor<fp8e5m2_4, float4> {
-    static __host__ __device__ inline fp8e5m2_4 convert(const float4& u) {
-        return __nv_fp8x4_e5m2(u); 
-    }
-};
-template<> struct convertor<float4, fp8e5m2_4> {
-    static __host__ __device__ inline float4 convert(const fp8e5m2_4& u) {
-        __nv_fp8_e5m2 *vals = reinterpret_cast<__nv_fp8_e5m2*>(const_cast<__nv_fp8x4_e5m2*>(&u));
-        return make_float4(float(vals[0]), float(vals[1]), float(vals[2]), float(vals[3]));
-    }
-};
-template<> struct convertor<fp8e5m2_2, float2> {
-    static __host__ __device__ inline fp8e5m2_2 convert(const float2& u) {
-        return __nv_fp8x2_e5m2(u); 
-    }
-};
-template<> struct convertor<float2, fp8e5m2_2> {
-    static __host__ __device__ inline float2 convert(const fp8e5m2_2& u) {
-        __nv_fp8_e5m2 *vals = reinterpret_cast<__nv_fp8_e5m2*>(const_cast<__nv_fp8x2_e5m2*>(&u));
-        return make_float2(float(vals[0]), float(vals[1]));
-    }
-};
-template<> struct convertor<fp8e5m2, float> {
-    static __host__ __device__ inline fp8e5m2 convert(const float & u) {
-        return __nv_fp8_e5m2(u);
-    }
-};
-template<> struct convertor<float, fp8e5m2> {
-    static __host__ __device__ inline float convert(const fp8e5m2 & u) {
-        return float(u);
-    }
-};
-template<> struct convertor<bf16_2, fp8e5m2_4> {
-    static __host__ __device__ inline bf16_2 convert(const fp8e5m2_4 & u) {
-        float4 f4 = convertor<float4, fp8e5m2_4>::convert(u);
-        float2 f2 = make_float2(f4.x, f4.y);
-        return __float22bfloat162_rn(f2);
-    }
-};
-template<> struct convertor<fp8e5m2_4, bf16_2> {
-    static __host__ __device__ inline fp8e5m2_4 convert(const bf16_2 & u) {
-        float2 f2 = __bfloat1622float2(u);
-        float4 f4 = make_float4(f2.x, f2.y, 0.0f, 0.0f);
-        return __nv_fp8x4_e5m2(f4);
-    }
-};
-#endif
-}
-}
@@ -1,11 +0,0 @@
-/**
- * @file
- * @brief A collection of common resources on which ThunderKittens depends.
- */
- 
-
-#pragma once
-
-#include "util.cuh"
-#include "base_types.cuh"
-#include "base_ops.cuh"
@@ -1,56 +0,0 @@
-#pragma once
-
-// Reset
-#define TK_RESET "\033[0m"
-
-// Foreground colors
-#define TK_FG_BLACK "\033[30m"
-#define TK_FG_RED "\033[31m"
-#define TK_FG_GREEN "\033[32m"
-#define TK_FG_YELLOW "\033[33m"
-#define TK_FG_BLUE "\033[34m"
-#define TK_FG_MAGENTA "\033[35m"
-#define TK_FG_CYAN "\033[36m"
-#define TK_FG_WHITE "\033[37m"
-
-// Background colors
-#define TK_BG_BLACK "\033[40m"
-#define TK_BG_RED "\033[41m"
-#define TK_BG_GREEN "\033[42m"
-#define TK_BG_YELLOW "\033[43m"
-#define TK_BG_BLUE "\033[44m"
-#define TK_BG_MAGENTA "\033[45m"
-#define TK_BG_CYAN "\033[46m"
-#define TK_BG_WHITE "\033[47m"
-
-// Bright foreground colors
-#define TK_FG_BRIGHT_BLACK "\033[90m"
-#define TK_FG_BRIGHT_RED "\033[91m"
-#define TK_FG_BRIGHT_GREEN "\033[92m"
-#define TK_FG_BRIGHT_YELLOW "\033[93m"
-#define TK_FG_BRIGHT_BLUE "\033[94m"
-#define TK_FG_BRIGHT_MAGENTA "\033[95m"
-#define TK_FG_BRIGHT_CYAN "\033[96m"
-#define TK_FG_BRIGHT_WHITE "\033[97m"
-
-// Bright background colors
-#define TK_BG_BRIGHT_BLACK "\033[100m"
-#define TK_BG_BRIGHT_RED "\033[101m"
-#define TK_BG_BRIGHT_GREEN "\033[102m"
-#define TK_BG_BRIGHT_YELLOW "\033[103m"
-#define TK_BG_BRIGHT_BLUE "\033[104m"
-#define TK_BG_BRIGHT_MAGENTA "\033[105m"
-#define TK_BG_BRIGHT_CYAN "\033[106m"
-#define TK_BG_BRIGHT_WHITE "\033[107m"
-
-// Text styles
-#define TK_BOLD "\033[1m"
-#define TK_DIM "\033[2m"
-#define TK_ITALIC "\033[3m"
-#define TK_UNDERLINE "\033[4m"
-#define TK_BLINK "\033[5m"
-#define TK_REVERSE "\033[7m"
-#define TK_HIDDEN "\033[8m"
-
-// Macro to combine styles
-#define TK_STYLE(...)  "\033[" #__VA_ARGS__ "m"
@@ -1,314 +0,0 @@
-/**
- * @file
- * @brief General utilities for ThunderKittens.
- */
-
-#pragma once
-
-#include <stdint.h>
-#include <type_traits>
-#include <concepts>
-#include <memory>
-
-// CUDA driver API
-#define CUCHECK(cmd) do {                                     \
-    CUresult err = cmd;                                       \
-    if (err != CUDA_SUCCESS) {                                \
-        const char *errStr;                                   \
-        cuGetErrorString(err, &errStr);                       \
-        fprintf(stderr, "Failed: CUDA error %s:%d '%s'\n",    \
-            __FILE__, __LINE__, errStr);                      \
-        exit(EXIT_FAILURE);                                   \
-    }                                                         \
-} while(0)
-
-// CUDA runtime API
-#define CUDACHECK(cmd) do {                                   \
-    cudaError_t err = cmd;                                    \
-    if (err != cudaSuccess) {                                 \
-        fprintf(stderr, "Failed: CUDA error %s:%d '%s'\n",    \
-            __FILE__, __LINE__, cudaGetErrorString(err));     \
-        exit(EXIT_FAILURE);                                   \
-    }                                                         \
-} while(0)
-
-/**
- * @namespace kittens
- *
- * @brief The main namespace of ThunderKittens.
- */
-namespace kittens {
-
-/* ----------  GENERAL CONSTANTS FOR KITTENS  ---------- */
-
-/**
- * @brief Tile dimension constant.
- */
-template<typename T> constexpr int TILE_COL_DIM = sizeof(T) == 1 ? 32 : 16;
-template<typename T> constexpr int TILE_ROW_DIM = 16;
-/**
- * @brief Tile num elements constant calculated as TILE_DIM squared.
- */
-template<typename T> constexpr int TILE_ELEMENTS{TILE_COL_DIM<T>*TILE_ROW_DIM<T>};
-/**
- * @brief Constant representing number of threads in a warp.
- */
-constexpr int WARP_THREADS{32};
-/**
- * @brief Constant representing number of threads in a warpgroup of four warps.
- */
-constexpr int WARPGROUP_THREADS{128};
-/**
-
- * @brief Constant representing number of warps in a warpgroup of four warps.
- */
-constexpr int WARPGROUP_WARPS{4};
-/**
-
- * @brief Get the warp ID of the current thread.
- * @return The warp ID.
- */
-__device__ static __forceinline__ int warpid() {
-    // uint32_t wid;
-    // asm volatile("mov.u32 %0, %warpid;" : "=r"(wid));
-    // return wid;
-    return threadIdx.x >> 5;
-}
-/**
- * @brief Get the warpgroup ID of the current thread.
- * @return The warpgroup ID.
- */
-__device__ static __forceinline__ int warpgroupid() { return warpid() >> 2; }
-/**
- * @brief Get the lane ID of the current thread within its warp.
- * @return The lane ID.
- */
-__device__ static __forceinline__ int laneid() {
-    // uint32_t lid;
-    // asm volatile("mov.u32 %0, %laneid;" : "=r"(lid));
-    // return lid;
-    return threadIdx.x & 31;
-}
-
-#if defined(KITTENS_HOPPER)
-constexpr int MAX_SHARED_MEMORY = 227000;
-#elif defined(KITTENS_A100)
-constexpr int MAX_SHARED_MEMORY = 164000;
-#elif defined(KITTENS_4090)
-constexpr int MAX_SHARED_MEMORY = 100000;
-#endif
-
-struct transpose {
-    static constexpr int N = 0; // not transposed
-    static constexpr int T = 1; // transposed
-};
-struct axis {
-    static constexpr int ROW = 0; // row axis of a tile
-    static constexpr int COL = 1; // column axis of a tile
-};
-
-/* ----------  TYPE HELPERS  ---------- */
-
-/**
- * @namespace ducks
- *
- * @brief ThunderKittens' namespace for template metaprogramming..
- * 
- * This includes primarily dummy types and concept wrappers, along
- * with a few additional utilities.
- */
-namespace ducks {
-
-/**
- * @brief A type representing an empty default for a template.
- */
-struct default_type {};
-
-// This macro can't be done as a template, so it doesn't really have a location in kittens.
-#define typeof(A) typename std::remove_const<typename std::remove_reference<decltype(A)>::type>::type
-
-}
-
-/* ----------  SHUFFLE UTILS  ---------- */
-
-/**
- * @brief Mask constant for all active threads in a warp.
- */
-static constexpr uint32_t MASK_ALL = 0xFFFFFFFF;
-
-/**
- * @brief Perform a shuffle down operation on a packed type synchronously across a warp.
- * @tparam T The type of the value to be shuffled.
- * @param mask[in] The mask of active threads.
- * @param f[in] The value to be shuffled.
- * @param delta[in] The number of positions to shuffle down.
- * @return The result of the shuffle operation.
- */
-template<typename T>
-__device__ static inline T packed_shfl_down_sync(uint32_t mask, const T &f, int delta) {
-    return __shfl_down_sync(mask, f, delta);
-}
-template<>
-__device__ inline float2 packed_shfl_down_sync<float2>(uint32_t mask, const float2 &f, int delta) {
-    float2 r;
-    r.x = __shfl_down_sync(mask, f.x, delta);
-    r.y = __shfl_down_sync(mask, f.y, delta);
-    return r;
-}
-/**
- * @brief Perform a packed shuffle operation synchronously across a warp.
- * @tparam T The type of the value to be shuffled.
- * @param mask[in] The mask of active threads.
- * @param f[in] The value to be shuffled.
- * @param src[in] The source lane from which to shuffle.
- * @return The result of the shuffle operation.
- */
-template<typename T>
-__device__ static inline T packed_shfl_sync(uint32_t mask, const T &f, int src) {
-    return __shfl_sync(mask, f, src);
-}
-template<>
-__device__ inline float2 packed_shfl_sync<float2>(uint32_t mask, const float2 &f, int src) {
-    float2 r;
-    r.x = __shfl_sync(mask, f.x, src);
-    r.y = __shfl_sync(mask, f.y, src);
-    return r;
-}
-
-/* ----------  SHARED MEMORY UTILS  ---------- */
-
-// namespace ducks {
-// namespace sb {
-// struct identifier {};
-// }
-// }
-
-// template<typename Args...>
-// struct sb {
-//     using identifier = ducks::sb::identifier;
-//     Args... args;
-// };
-
-// namespace ducks {
-// namespace sb {
-// template<typename T> concept all = requires {
-//     typename T::identifier;
-// } && std::is_same_v<T::identifier, identifier>;
-// }
-// }
-
-// Joyously stolen from https://github.com/NVIDIA/cutlass/blob/5c447dd84f8ae0e1d48ff9a2eae26ce8c4958101/include/cute/container/alignment.hpp#L51
-#if defined(__CUDACC__)
-#define KITTENS_ALIGN_AS(n) __align__(n)
-#else
-#define KITTENS_ALIGN_AS(n) alignas(n)
-#endif
-
-#ifdef KITTENS_HOPPER
-#define KITTENS_DEFAULT_ALIGN KITTENS_ALIGN_AS(128)
-#else
-#define KITTENS_DEFAULT_ALIGN KITTENS_ALIGN_AS(16)
-#endif
-
-/**
- * @brief Dummy structure for alignment purposes. Needed for WGMMA and TMA calls.
- */
-struct KITTENS_DEFAULT_ALIGN alignment_dummy { int dummy; };
-/**
- * @brief Very simple allocator for dynamic shared memory. Advances pointer and tracks alignments.
- * @tparam default_alignment The default alignment this allocator will enforce. If <=0 (default -1) it will not align.
- */
-#ifdef KITTENS_HOPPER
-template<int default_alignment=1024> 
-#else
-template<int default_alignment=16> 
-#endif
-struct shared_allocator {
-    int *ptr;
-
-    private:
-        // Recursive template to generate N-dimensional array type
-        template<typename A, size_t... dims>
-        struct variadic_array;
-        template<typename A, size_t first_dim, size_t... rest_dims>
-        struct variadic_array<A, first_dim, rest_dims...> {
-            using type = typename variadic_array<A, rest_dims...>::type[first_dim];
-        };
-        template<typename A>
-        struct variadic_array<A> {
-            using type = A;
-        };
-        template<typename A, size_t... dims> 
-        using variadic_array_t = typename variadic_array<A, dims...>::type;
-
-        template<int alignment>
-        __device__ inline void align_ptr() {
-            if constexpr (alignment > 0) {
-                uint64_t p = reinterpret_cast<uint64_t>(ptr);
-                if(p % alignment != 0) {
-                    ptr = (int*)(p + (alignment-(p%alignment)));
-                }
-            }
-        }
-
-    public:
-        /**
-        * @brief Construct a new shared allocator using a pointer to extern shared memory.
-        * @param[in] _ptr Pointer to the start of the extern shared memory.
-        */
-        __device__ shared_allocator(int *_ptr): ptr(_ptr) {}
-        /**
-        * @brief Allocate shared memory for a single instance or N-dimensional array of type A.
-        * @tparam A The type of the object to allocate.
-        * @tparam dims... A list of dimensions for the N-dimensional array.
-        * @return Reference to the allocated object.
-        */
-        template<typename A, size_t... dims> 
-        __device__ inline variadic_array_t<A, dims...>& allocate() {
-            // static_assert(sizeof(A) % default_alignment == 0, "Type is not aligned properly for array allocation");
-            align_ptr<default_alignment>();
-            using at = variadic_array_t<A, dims...>;
-            at*p = reinterpret_cast<at*>(ptr);
-            ptr += sizeof(at)/sizeof(int);
-            return *p;
-        }
-        /**
-        * @brief Allocate shared memory for a single instance or N-dimensional array of type A.
-        * @tparam alignment An alignment to enforce for this particular object.
-        * @tparam A The type of the object to allocate.
-        * @tparam dims... A list of dimensions for the N-dimensional array.
-        * @return Reference to the allocated object.
-        */
-        template<int alignment, typename A, size_t... dims> 
-        __device__ inline variadic_array_t<A, dims...>& allocate() {
-            // static_assert(sizeof(A) % alignment == 0, "Type is not aligned properly for array allocation");
-            align_ptr<alignment>();
-            using at = variadic_array_t<A, dims...>;
-            at*p = reinterpret_cast<at*>(ptr);
-            ptr += sizeof(at)/sizeof(int);
-            return *p;
-        }
-};
-#if (defined(KITTENS_HOPPER) || defined(KITTENS_BLACKWELL))
-/**
- * @brief A wrapper for an allocator that enforces sufficient alignment to be used for TMA loads and stores.
- */
-using tma_allocator = shared_allocator<1024>;
-using tma_swizzle_allocator = tma_allocator; // swizzled TMA modes require up to 1024 byte alignments :/
-
-/* Get CTA ID within a cluster */
-__device__ static inline int3 clusterIdx() {
-    int3 cluster_idx;
-    asm volatile("mov.u32 %0, %clusterid.x;\n" : "=r"(cluster_idx.x));
-    asm volatile("mov.u32 %0, %clusterid.y;\n" : "=r"(cluster_idx.y));
-    asm volatile("mov.u32 %0, %clusterid.z;\n" : "=r"(cluster_idx.z));
-    return cluster_idx;
-}
-__device__ static inline int cluster_ctarank() {
-    uint32_t ctarank;
-    asm volatile("mov.u32 %0, %cluster_ctarank;\n" : "=r"(ctarank));
-    return ctarank;
-}
-#endif
-
-} // namespace kittens
@@ -1,12 +0,0 @@
-/**
- * @file
- * @brief The master header file of ThunderKittens. This file includes everything you need!
- */
-
-#pragma once
-
-#include "common/common.cuh"
-#include "types/types.cuh"
-#include "ops/ops.cuh"
-#include "pyutils/util.cuh"
-// #include "pyutils/pyutils.cuh" // for simple binding without including torch
@@ -1,51 +0,0 @@
-/**
- * @file
- * @brief An aggregate header of all device (multi-GPU) operations defined by ThunderKittens
- */
-
-#pragma once
-
-#include "../../types/types.cuh"
-
-namespace kittens {
-
-template<int _NUM_DEVICES>
-struct device {
-
-static_assert(_NUM_DEVICES >= 0 && _NUM_DEVICES <= 72, "Invalid number of devices");
-static constexpr int NUM_DEVICES = _NUM_DEVICES;
-
-#ifdef KITTENS_HOPPER
-
-using barrier_t = pgl<gl<int, 1, 1, 1, -1>, NUM_DEVICES, true>;
-
-/**
- * @brief Multi-GPU synchronization barrier for coordinated kernel exit
- * 
- * Performs a synchronization across all devices to ensure all GPUs complete 
- * their work before any kernel exits. Does not synchronize intra-node threads
- * or threadblocks.
- * 
- * @param barrier Pre-allocated barrier structure, must be initialized to 0
- * @param dev_idx Current device index (0 to NUM_DEVICES - 1)
- * @param id Synchronization point identifier (default: 0). 0 is fine for most cases
- * 
- */
-__device__ static inline void sync_on_exit(const barrier_t &barrier, const int dev_idx, const int id = 0) {
-    if (blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0 &&
-        threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0) {
-        cuda::atomic_ref<int, cuda::thread_scope_system> barrier_uc(barrier[dev_idx][{id}]);
-
-        // Inter-note check-in
-        multimem<int>::red<reduce_op::ADD>(barrier.mc_ptr_at({id}), 1);
-        asm volatile ("{fence.proxy.alias;}" ::: "memory");
-        while (barrier_uc.load(cuda::memory_order_acquire) < NUM_DEVICES);
-        barrier_uc.fetch_sub(NUM_DEVICES, cuda::memory_order_release);
-    }
-}
-
-#endif
-
-};
-
-} // namespace kittens
@@ -1,96 +0,0 @@
-/**
- * @file
- * @brief An aggregate header of all group (multi-warp) operations defined by ThunderKittens
- */
-
-#pragma once
-
-#include <cuda/pipeline>
-
-#include "../../common/common.cuh"
-#include "../../types/types.cuh"
-#include "../thread/thread.cuh" // several group memory ops rely on underlying warp-scope ops
-
-#define KITTENS_CHECK_WARP static_assert(GROUP_WARPS==1, "Warp (GROUP_WARPS=1) function called from a non-warp group.");
-// A "warpgroup" is a special group of 4 consecutive warps defined by NVIDIA for certain SM_90+ operations.
-#define KITTENS_CHECK_WARPGROUP static_assert(GROUP_WARPS==4, "Warpgroup (GROUP_WARPS=4) function called from a non-warpgroup group.");
-
-// WGMMA relies on some template structures that cannot be specialized within the group struct, so we declare them in advance.
-#ifdef KITTENS_HOPPER
-#include "mma/warpgroup/base/base.cuh"
-#endif
-
-namespace kittens {
-/*
-This is meant to be used with a `using group_N = kittens::group<NUM_WORKERS>;` at the start of every kernel.
-*/
-template<int _GROUP_WARPS>
-struct group {
-static constexpr int GROUP_WARPS = _GROUP_WARPS; // This alias produces nice parallelism.
-static constexpr int GROUP_THREADS = GROUP_WARPS * kittens::WARP_THREADS; // This alias produces nice parallelism.
-__device__ static inline int laneid() { return threadIdx.x % GROUP_THREADS; }
-__device__ static inline int warpid() { return laneid() / kittens::WARP_THREADS; }
-__device__ static inline int groupid() { return threadIdx.x / GROUP_THREADS; }
-
-__device__ static inline void sync(int id) {
-    asm volatile("bar.sync %0, %1;\n" :: "r"(id), "n"(GROUP_THREADS));
-}
-template<uint32_t MASK=0xFFFFFFFF> __device__ static inline void sync() {
-    static_assert(GROUP_WARPS==1, "barrier-less sync() can only be called by a single warp!");
-    asm volatile("bar.warp.sync %0;\n" :: "n"(MASK));
-}
-__device__ static inline void arrive(int id) {
-    asm volatile("bar.arrive %0, %1;\n" :: "r"(id), "n"(GROUP_THREADS));
-}
-
-#include "memory/memory.cuh"
-#include "shared/shared.cuh"
-#include "register/register.cuh"
-
-#ifdef KITTENS_HOPPER
-#include "mma/mma.cuh"
-
-template<int n_reg> __device__ static inline void increase_registers() {
-    static_assert(n_reg % 8 == 0, "n_reg must be a multiple of 8");
-    asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" :: "n"(n_reg));
-}
-template<int n_reg> __device__ static inline void decrease_registers() {
-    static_assert(n_reg % 8 == 0, "n_reg must be a multiple of 8");
-    asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" :: "n"(n_reg));
-}
-__device__ static inline void producer_registers() { decrease_registers<24>(); }
-template<int NCWG> __device__ static inline void consumer_registers() { increase_registers<480/NCWG - 8*(NCWG>3) - 224*(NCWG==1)>(); }
-
-#endif
-
-};
-
-namespace everyone {
-
-// Block-level synchronization
-__device__ static inline void sync(int id) {
-    asm volatile("bar.sync %0;\n" :: "r"(id));
-}
-
-// Cluster-level synchronization functions
-namespace tma {
-namespace cluster {
-__device__ static inline void arrive_aligned() { // All threads in the cluster must call this
-    asm volatile ("barrier.cluster.arrive.release.aligned;\n");
-}
-__device__ static inline void wait_aligned() {
-    asm volatile ("barrier.cluster.wait.acquire.aligned;\n");
-}
-__device__ static inline void sync() {
-    arrive_aligned();
-    wait_aligned();
-}
-}
-}
-
-};
-
-using warp = group<1>;      // scope used by most pre-Hopper GPUs, and also for most register operations.
-using warpgroup = group<4>; // special scope commonly used by Hopper and later.
-
-}
@@ -1,21 +0,0 @@
-/**
- * @file
- * @brief An aggregate header of colaborative group memory movement operations
- */
-
-#include "util/util.cuh"
-#include "tile/tile.cuh"
-#include "vec/vec.cuh"
-
-#ifdef KITTENS_HOPPER
-struct tma {
-#include "util/tma.cuh"
-#include "tile/tma.cuh"
-#include "vec/tma.cuh"
-struct cluster {
-#include "util/tma_cluster.cuh"
-#include "tile/tma_cluster.cuh"
-#include "vec/tma_cluster.cuh"
-};
-};
-#endif
@@ -1,42 +0,0 @@
-/**
- * @file
- * @brief Functions for a group to collaboratively transfer data directly between global memory and registers and back.
- */
-
-/**
- * @brief Collaboratively loads data from a source array into register tiles.
- *
- * @tparam RT The register tile type.
- * @tparam U The data type of the source array.
- * @param dst[out] The destination tile to load data into.
- * @param src[in] The source array to load data from.
- * @param row_stride[in] The stride in elements between rows in the source array.
- */
-template<int axis, ducks::crt::all CRT, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<crt<typename CRT::T, GROUP_WARPS*CRT::rows, CRT::cols, typename CRT::layout>>>
-__device__ inline static void load(CRT &dst, const CGL &src, const COORD &idx) {
-    load<axis, CRT::component, CGL::component, COORD>(dst.real, src.real, idx);
-    load<axis, CRT::component, CGL::component, COORD>(dst.imag, src.imag, idx);
-}
-template<ducks::crt::all CRT, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<crt<typename CRT::T, GROUP_WARPS*CRT::rows, CRT::cols, typename CRT::layout>>>
-__device__ inline static void load(CRT &dst, const CGL &src, const COORD &idx) {
-    load<2, CRT, CGL>(dst, src, idx);
-}
-
-/**
- * @brief Collaboratively stores data from register tiles to a destination array in global memory.
- *
- * @tparam RT The register tile type.
- * @tparam U The data type of the destination array.
- * @param[out] dst The destination array in global memory to store data into.
- * @param[in] src The source register tile to store data from.
- * @param row_stride[in] The stride in elements between rows in the destination array.
- */
-template<int axis, ducks::crt::all CRT, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<crt<typename CRT::T, GROUP_WARPS*CRT::rows, CRT::cols, typename CRT::layout>>>
-__device__ inline static void store(CGL &dst, const CRT &src, const COORD &idx) {
-    store<axis, typename CRT::component, typename CGL::component>(dst.real, src.real, idx);
-    store<axis, typename CRT::component, typename CGL::component>(dst.imag, src.imag, idx);
-}
-template<ducks::crt::all CRT, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<crt<typename CRT::T, GROUP_WARPS*CRT::rows, CRT::cols, typename CRT::layout>>>
-__device__ inline static void store(CGL &dst, const CRT &src, const COORD &idx) {
-    store<2, CRT, CGL>(dst, src, idx);
-}
@@ -1,37 +0,0 @@
-/**
- * @file
- * @brief Group (collaborative warp) ops for loading shared tiles from and storing to global memory. 
- */
-
-template<int axis, bool assume_aligned, ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
-__device__ static inline void load(CST &dst, const CGL &src, const COORD &idx) {
-    load<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
-    load<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
-}
-template<ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
-__device__ static inline void load(CST &dst, const CGL &src, const COORD &idx) {
-    load<2, false, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
-    load<2, false, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
-}
-
-template<int axis, bool assume_aligned, ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
-__device__ static inline void store(CGL &dst, const CST &src, const COORD &idx) {
-    store<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
-    store<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
-}
-template<ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
-__device__ static inline void store(CGL &dst, const CST &src, const COORD &idx) {
-    store<2, false, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
-    store<2, false, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
-}
-
-template<int axis, bool assume_aligned, ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
-__device__ static inline void load_async(CST &dst, const CGL &src, const COORD &idx) {
-    load_async<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
-    load_async<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
-}
-template<ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
-__device__ static inline void load_async(CST &dst, const CGL &src, const COORD &idx) {
-    load_async<2, false, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
-    load_async<2, false, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
-}
@@ -1,34 +0,0 @@
-/**
- * @file
- * @brief Functions for a warpgroup to collaboratively transfer data directly between shared memory and registers and back.
- */
-
-/**
- * @brief Collaboratively load data from a shared tile into register tiles split across a warpgroup.
- *
- * @tparam RT The register tile type
- * @tparam ST The shared tile type
- * @param dst[out] The destination register tile.
- * @param src[in]  The source shared tile.
- */
-template<ducks::crt::all RT, ducks::cst::all ST>
-__device__ inline static void load(RT &dst, const ST &src) {
-    load(dst.real, src.real);
-    load(dst.imag, src.imag);
-}
-
-
-/**
- * @brief Collaboratively store data into a shared tile from register tiles split across a warpgroup.
- *
- * @tparam RT The register tile type
- * @tparam ST The shared tile type
- * @param dst[out] The destination shared tile.
- * @param src[in]  The source register tile.
- */
-template<ducks::cst::all ST, ducks::crt::all RT>
-__device__ inline static void store(ST &dst, const RT &src) {
-    store(dst.real, src.real);
-    store(dst.imag, src.imag);
-}
-
@@ -1,207 +0,0 @@
-/**
- * @file
- * @brief Functions for a group to collaboratively transfer data directly between global memory and registers and back.
- */
-
-/**
- * @brief Collaboratively loads data from a source array into row-major layout tiles.
- *
- * @tparam RT The row-major layout tile type.
- * @tparam U The data type of the source array.
- * @param dst[out] The destination tile to load data into.
- * @param src[in] The source array to load data from.
- * @param row_stride[in] The stride in elements between rows in the source array.
- */
-template<int axis, ducks::rt::row_layout RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
-__device__ inline static void load(RT &dst, const GL &src, const COORD &idx) {
-    using T2 = RT::dtype;
-    using U = typename GL::dtype;
-
-    #ifdef KITTENS_HOPPER
-    static_assert(!std::is_same_v<T2, fp8e4m3_4> && !std::is_same_v<T2, fp8e5m2_4>, "Unsupported type for load/store");
-    #endif
-
-    U *src_ptr = (U*)&src[(idx.template unit_coord<axis, 3>())];
-    const int row_stride = src.template stride<axis>();
-    using U2 = base_types::packing<U>::packed_type;
-    int warp_laneid = threadIdx.x % WARP_THREADS;
-    int local_warpid;
-    if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
-    else local_warpid = warpid();
-    const int row_offset = dst.rows*local_warpid;
-    #pragma unroll
-    for(int i = 0; i < dst.height; i++) {
-        int row = row_offset + i*dst.tile_size_row + (warp_laneid / 4);
-        #pragma unroll
-        for(int j = 0; j < dst.width; j++) {
-            int col = j*dst.tile_size_col + 2*(warp_laneid % 4);
-            dst.tiles[i][j].data[0] = base_types::convertor<T2, U2>::convert(*(U2*)(&src_ptr[(row+0)*row_stride + (col+0)]));
-            dst.tiles[i][j].data[2] = base_types::convertor<T2, U2>::convert(*(U2*)(&src_ptr[(row+0)*row_stride + (col+8)]));
-        }
-        #pragma unroll
-        for(int j = 0; j < dst.width; j++) {
-            int col = j*dst.tile_size_col + 2*(warp_laneid % 4);
-            dst.tiles[i][j].data[1] = base_types::convertor<T2, U2>::convert(*(U2*)(&src_ptr[(row+8)*row_stride + (col+0)]));
-            dst.tiles[i][j].data[3] = base_types::convertor<T2, U2>::convert(*(U2*)(&src_ptr[(row+8)*row_stride + (col+8)]));
-        }
-    }
-}
-/**
- * @brief Collaboratively loads data from a source array into column-major layout tiles.
- *
- * @tparam RT The column-major layout tile type.
- * @tparam U The data type of the source array.
- * @param dst[out] The destination tile to load data into.
- * @param src[in] The source array to load data from.
- * @param row_stride[in] The stride in elements between rows in the source array.
- */
-template<int axis, ducks::rt::col_layout RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
-__device__ inline static void load(RT &dst, const GL &src, const COORD &idx) {
-    using T = typename RT::T;
-    using U = typename GL::dtype;
-
-    #ifdef KITTENS_HOPPER
-    static_assert(!std::is_same_v<T, fp8e4m3> && !std::is_same_v<T, fp8e5m2>, "Unsupported type for load/store");
-    #endif
-
-    U *src_ptr = (U*)&src[(idx.template unit_coord<axis, 3>())];
-    const int row_stride = src.template stride<axis>();
-    int warp_laneid = threadIdx.x % WARP_THREADS;
-    int local_warpid;
-    if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
-    else local_warpid = warpid();
-    const int row_offset = dst.rows*local_warpid;
-    #pragma unroll
-    for(int i = 0; i < dst.height; i++) {
-        int row = row_offset + i*dst.tile_size_row + 2*(warp_laneid % 4);
-        #pragma unroll
-        for(int j = 0; j < dst.width; j++) {
-            int col = j*dst.tile_size_col + (warp_laneid / 4);
-            dst.tiles[i][j].data[0].x = base_types::convertor<T, U>::convert(src_ptr[(row+0)*row_stride + (col+0)]);
-            dst.tiles[i][j].data[1].x = base_types::convertor<T, U>::convert(src_ptr[(row+0)*row_stride + (col+8)]);
-        }
-        #pragma unroll
-        for(int j = 0; j < dst.width; j++) {
-            int col = j*dst.tile_size_col + (warp_laneid / 4);
-            dst.tiles[i][j].data[0].y = base_types::convertor<T, U>::convert(src_ptr[(row+1)*row_stride + (col+0)]);
-            dst.tiles[i][j].data[1].y = base_types::convertor<T, U>::convert(src_ptr[(row+1)*row_stride + (col+8)]);
-        }
-        #pragma unroll
-        for(int j = 0; j < dst.width; j++) {
-            int col = j*dst.tile_size_col + (warp_laneid / 4);
-            dst.tiles[i][j].data[2].x = base_types::convertor<T, U>::convert(src_ptr[(row+8)*row_stride + (col+0)]);
-            dst.tiles[i][j].data[3].x = base_types::convertor<T, U>::convert(src_ptr[(row+8)*row_stride + (col+8)]);
-        }
-        #pragma unroll
-        for(int j = 0; j < dst.width; j++) {
-            int col = j*dst.tile_size_col + (warp_laneid / 4);
-            dst.tiles[i][j].data[2].y = base_types::convertor<T, U>::convert(src_ptr[(row+9)*row_stride + (col+0)]);
-            dst.tiles[i][j].data[3].y = base_types::convertor<T, U>::convert(src_ptr[(row+9)*row_stride + (col+8)]);
-        }
-    }
-}
-template<ducks::rt::all RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
-__device__ inline static void load(RT &dst, const GL &src, const COORD &idx) {
-    load<2>(dst, src, idx);
-}
-/**
- * @brief Collaboratively stores data from register tiles to a destination array in global memory with a row-major layout.
- *
- * @tparam RT The register tile type with a row-major layout.
- * @tparam U The data type of the destination array.
- * @param[out] dst The destination array in global memory to store data into.
- * @param[in] src The source register tile to store data from.
- * @param row_stride[in] The stride in elements between rows in the destination array.
- */
-template<int axis, ducks::rt::row_layout RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
-__device__ inline static void store(const GL &dst, const RT &src, const COORD &idx) {
-    using T2 = RT::dtype;
-    using U = typename GL::dtype;
-
-    #ifdef KITTENS_HOPPER
-    static_assert(!std::is_same_v<T2, fp8e4m3_4> && !std::is_same_v<T2, fp8e5m2_4>, "Unsupported type for load/store");
-    #endif
-
-    U *dst_ptr = (U*)&dst[(idx.template unit_coord<axis, 3>())];
-    const int row_stride = dst.template stride<axis>();
-    using U2 = base_types::packing<U>::packed_type;
-    int warp_laneid = threadIdx.x % WARP_THREADS;
-    int local_warpid;
-    if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
-    else local_warpid = warpid();
-    const int row_offset = src.rows*local_warpid;
-    #pragma unroll
-    for(int i = 0; i < src.height; i++) {
-        int row = row_offset + i*src.tile_size_row + (warp_laneid / 4);
-        #pragma unroll
-        for(int j = 0; j < src.width; j++) {
-            int col = j*src.tile_size_col + 2*(warp_laneid % 4);
-            *(U2*)(&dst_ptr[(row+0)*row_stride + (col+0)]) = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[0]);
-            *(U2*)(&dst_ptr[(row+0)*row_stride + (col+8)]) = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[2]);
-        }
-        #pragma unroll
-        for(int j = 0; j < src.width; j++) {
-            int col = j*src.tile_size_col + 2*(warp_laneid % 4);
-            *(U2*)(&dst_ptr[(row+8)*row_stride + (col+0)]) = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[1]);
-            *(U2*)(&dst_ptr[(row+8)*row_stride + (col+8)]) = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[3]);
-        }
-    }
-}
-/**
- * @brief Collaboratively stores data from register tiles to a destination array in global memory with a column-major layout.
- *
- * @tparam RT The register tile type with a column-major layout.
- * @tparam U The data type of the destination array.
- * @param[out] dst The destination array in global memory to store data into.
- * @param[in] src The source register tile to store data from.
- * @param row_stride[in] The stride in elements between rows in the destination array.
- */
-template<int axis, ducks::rt::col_layout RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
-__device__ inline static void store(const GL &dst, const RT &src, const COORD &idx) {
-    using T = base_types::packing<typename RT::dtype>::unpacked_type;
-    using U = typename GL::dtype;
-
-    #ifdef KITTENS_HOPPER
-    static_assert(!std::is_same_v<T, fp8e4m3_4> && !std::is_same_v<T, fp8e5m2_4>, "Unsupported type for load/store");
-    #endif
-    
-    U *dst_ptr = (U*)&dst[(idx.template unit_coord<axis, 3>())];
-    const int row_stride = dst.template stride<axis>();
-    int warp_laneid = threadIdx.x % WARP_THREADS;
-    int local_warpid;
-    if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
-    else local_warpid = warpid();
-    const int row_offset = src.rows*local_warpid;
-    #pragma unroll
-    for(int i = 0; i < src.height; i++) {
-        int row = row_offset + i*src.tile_size_row + 2*(warp_laneid % 4);
-        #pragma unroll
-        for(int j = 0; j < src.width; j++) {
-            int col = j*src.tile_size_col + (warp_laneid / 4);
-            dst_ptr[(row+0)*row_stride + (col+0)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[0].x);
-            dst_ptr[(row+0)*row_stride + (col+8)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[1].x);
-        }
-        #pragma unroll
-        for(int j = 0; j < src.width; j++) {
-            int col = j*src.tile_size_col + (warp_laneid / 4);
-            dst_ptr[(row+1)*row_stride + (col+0)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[0].y);
-            dst_ptr[(row+1)*row_stride + (col+8)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[1].y);
-        }
-        #pragma unroll
-        for(int j = 0; j < src.width; j++) {
-            int col = j*src.tile_size_col + (warp_laneid / 4);
-            dst_ptr[(row+8)*row_stride + (col+0)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[2].x);
-            dst_ptr[(row+8)*row_stride + (col+8)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[3].x);
-        }
-        #pragma unroll
-        for(int j = 0; j < src.width; j++) {
-            int col = j*src.tile_size_col + (warp_laneid / 4);
-            dst_ptr[(row+9)*row_stride + (col+0)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[2].y);
-            dst_ptr[(row+9)*row_stride + (col+8)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[3].y);
-        }
-    }
-}
-template<ducks::rt::all RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
-__device__ inline static void store(const GL &dst, const RT &src, const COORD &idx) {
-    store<2>(dst, src, idx);
-}
@@ -1,168 +0,0 @@
-/**
- * @file
- * @brief Group (collaborative warp) ops for loading shared tiles from and storing to global memory. 
- */
- 
-
-/**
- * @brief Loads data from global memory into a shared memory tile.
- *
- * @tparam ST The type of the shared tile.
- * @param[out] dst The destination shared memory tile.
- * @param[in] src The source global memory array.
- * @param[in] idx The coordinate of the tile in the global memory array.
- */
-template<int axis, bool assume_aligned, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void load(ST &dst, const GL &src, const COORD &idx) {
-    using T = typename ST::dtype;
-    const int row_stride = src.template stride<axis>();
-    // we can handle this many rows each time we run a memcpy_async
-    constexpr int elem_per_memcpy = sizeof(float4)/sizeof(typename ST::dtype);
-    constexpr int memcpy_per_row = dst.cols / elem_per_memcpy;
-    constexpr int total_calls = (dst.height*dst.width * kittens::TILE_ROW_DIM<T>*kittens::TILE_COL_DIM<T> + GROUP_THREADS*elem_per_memcpy-1) / (GROUP_THREADS*elem_per_memcpy); // round up
-    constexpr int total_rows = dst.height*dst.width;
-
-    coord<> unit_coord = idx.template unit_coord<axis, 3>();
-    typename GL::dtype *src_ptr = (typename GL::dtype*)&src[unit_coord];
-    uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
-    int laneid = threadIdx.x % GROUP_THREADS;
-
-    #pragma unroll
-    for(int i = 0; i < total_calls; i++) {
-
-        int load_idx = i * GROUP_THREADS + laneid;
-        
-        int row = load_idx / memcpy_per_row;
-        int col = (load_idx*elem_per_memcpy) % dst.cols;
-
-        if constexpr (assume_aligned) {
-            float4 tmp;
-            move<float4>::ldg(tmp, (float4*)&src_ptr[row*row_stride + col]);
-            move<float4>::sts(dst.idx(dst_ptr, {row, col}), tmp);
-        }
-        else {
-            if (row + unit_coord.template dim<axis>() < src.template shape<axis>()) {
-                float4 tmp;
-                move<float4>::ldg(tmp, (float4*)&src_ptr[row*row_stride + col]);
-                move<float4>::sts(dst.idx(dst_ptr, {row, col}), tmp);
-            }
-            else {
-                float4 zeros = {0.f,0.f,0.f,0.f};
-                move<float4>::sts(dst.idx(dst_ptr, {row, col}), zeros); // use the default value
-            }
-        }
-    }
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void load(ST &dst, const GL &src, const COORD &idx) {
-    load<2, false, ST, GL, COORD>(dst, src, idx);
-}
-
-/**
- * @brief Stores data from a shared memory tile into global memory.
- *
- * @tparam ST The type of the shared tile.
- * @param[out] dst The destination global memory array.
- * @param[in] src The source shared memory tile.
- * @param row_stride[in] The stride between rows in the destination array.
- */
-template<int axis, bool assume_aligned, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store(const GL &dst, const ST &src, const COORD &idx) {
-    using T = typename ST::dtype;
-    const int row_stride = dst.template stride<axis>();
-    // we can handle this many rows each time we run a memcpy_async
-    constexpr int elem_per_memcpy = sizeof(float4)/sizeof(typename ST::dtype);
-    constexpr int memcpy_per_row = src.cols / elem_per_memcpy;
-    constexpr int total_calls = (src.height*src.width * kittens::TILE_ROW_DIM<T>*kittens::TILE_COL_DIM<T> + GROUP_THREADS*elem_per_memcpy-1) / (GROUP_THREADS*elem_per_memcpy); // round up
-
-    coord<> unit_coord = idx.template unit_coord<axis, 3>();
-    typename GL::dtype *dst_ptr = (typename GL::dtype*)&dst[unit_coord];
-    uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src.data[0]));
-    int laneid = threadIdx.x % GROUP_THREADS;
-
-    #pragma unroll
-    for(int i = 0; i < total_calls; i++) {
-
-        int load_idx = i * GROUP_THREADS + laneid;
-        
-        int row = load_idx / memcpy_per_row;
-        int col = (load_idx*elem_per_memcpy) % src.cols;
-
-        if constexpr (assume_aligned) {
-            float4 tmp;
-            move<float4>::lds(tmp, src.idx(src_ptr, {row, col}));
-            move<float4>::stg((float4*)&dst_ptr[row*row_stride + col], tmp);
-        }
-        else {
-            if (row + unit_coord.template dim<axis>() < dst.template shape<axis>()) {
-                float4 tmp;
-                move<float4>::lds(tmp, src.idx(src_ptr, {row, col}));
-                move<float4>::stg((float4*)&dst_ptr[row*row_stride + col], tmp);
-            }
-        }
-    }
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store(const GL &dst, const ST &src, const COORD &idx) {
-    store<2, false, ST, GL, COORD>(dst, src, idx);
-}
-
-/**
- * @brief Asynchronously loads data from global memory into a shared memory tile.
- *
- * @tparam ST The type of the shared tile.
- * @param[out] dst The destination shared memory tile.
- * @param[in] src The source global memory array.
- *
- * @note This function expects 16-byte alignments. Otherwise, behavior is undefined.
- */
-template<int axis, bool assume_aligned, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx) {
-    using T = typename ST::dtype;
-    const int row_stride = src.template stride<axis>();
-    // we can handle this many rows each time we run a memcpy_async
-    constexpr int elem_per_memcpy = sizeof(float4)/sizeof(typename ST::dtype);
-    constexpr int memcpy_per_row = dst.cols / elem_per_memcpy;
-    constexpr int total_calls = (dst.height*dst.width * kittens::TILE_ROW_DIM<T>*kittens::TILE_COL_DIM<T> + GROUP_THREADS*elem_per_memcpy-1) / (GROUP_THREADS*elem_per_memcpy); // round up
-
-    coord<> unit_coord = idx.template unit_coord<axis, 3>();
-    typename GL::dtype *src_ptr = (typename GL::dtype*)&src[unit_coord];
-    uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
-    int laneid = threadIdx.x % GROUP_THREADS;
-
-    #pragma unroll
-    for(int i = 0; i < total_calls; i++) {
-
-        int load_idx = i * GROUP_THREADS + laneid;
-        
-        int row = load_idx / memcpy_per_row;
-        int col = (load_idx*elem_per_memcpy) % dst.cols;
-
-        if constexpr (assume_aligned) {
-            asm volatile(
-                "cp.async.cg.shared.global.L2::128B [%0], [%1], 16;\n"
-                :: "r"(dst.idx(dst_ptr, {row, col})), "l"(&src_ptr[row*row_stride + col])
-                : "memory"
-            );
-        }
-        else {
-            if (row + unit_coord.template dim<axis>() < src.template shape<axis>()) {
-                asm volatile(
-                    "cp.async.cg.shared.global.L2::128B [%0], [%1], 16;\n"
-                    :: "r"(dst.idx(dst_ptr, {row, col})), "l"(&src_ptr[row*row_stride + col])
-                    : "memory"
-                );
-            }
-            else {
-                // printf("thread %d skipping async load on row %d, col %d\n", threadIdx.x, row + unit_coord.template dim<axis>(), col);
-                float4 zeros = {0.f,0.f,0.f,0.f};
-                move<float4>::sts(dst.idx(dst_ptr, {row, col}), zeros); // use the default value
-            }
-        }
-    }
-    asm volatile("cp.async.commit_group;\n" ::: "memory");
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx) {
-    load_async<2, false, ST, GL, COORD>(dst, src, idx);
-}
@@ -1,323 +0,0 @@
-/**
- * @file
- * @brief Functions for a warpgroup to collaboratively transfer data directly between shared memory and registers and back.
- */
-
-/**
- * @brief Collaboratively load data from a shared tile into register tiles split across a warpgroup.
- *
- * @tparam RT The register tile type
- * @tparam ST The shared tile type
- * @param dst[out] The destination register tile.
- * @param src[in]  The source shared tile.
- */
-template<ducks::rt::all RT, ducks::st::all ST>
-__device__ inline static void load(RT &dst, const ST &src) {
-    constexpr int height = ST::height;
-    constexpr int warp_height = RT::height;
-    static_assert(height%GROUP_WARPS == 0, "Group load / store requires tile height to be a multiple of GROUP_WARPS.");
-    static_assert(height%warp_height == 0, "Group load / store requires tile height to be a multiple of the RT height.");
-    static_assert(ST::width==RT::width, "Group load / store requires tile widths to match.");
-    int local_warpid;
-    if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
-    else local_warpid = warpid();
-    using T2 = RT::dtype;
-    using U  = ST::dtype;
-    using T  = base_types::packing<T2>::unpacked_type;
-    using U2 = base_types::packing<U>::packed_type;
-    int warp_laneid = ::kittens::laneid();
-
-    // convert to shared state space
-    uint32_t shared_addr = static_cast<uint32_t>(__cvta_generic_to_shared(&src.data[0]));
-
-    #pragma unroll
-    for(int i = 0; i < dst.height; i++) {
-        #pragma unroll
-        for(int j = 0; j < dst.width; j++) {
-            if constexpr (sizeof(typename ST::dtype) == 2) {
-                // handle the row-major layout for 16-bit types
-                U2 tmp[4];
-                int row = (local_warpid*warp_height + i)*dst.tile_size_row + (warp_laneid % 16);
-                int col = j*dst.tile_size_col + (warp_laneid / 16) * 8;
-                if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row>) {
-                    move<U2>::ldsm4(tmp[0], tmp[1], tmp[2], tmp[3], src.idx(shared_addr, {row, col}));
-                }
-                else {
-                    move<U2>::ldsm4t(tmp[0], tmp[2], tmp[1], tmp[3], src.idx(shared_addr, {row, col}));
-                }
-                dst.tiles[i][j].data[0] = base_types::convertor<T2, U2>::convert(tmp[0]);
-                dst.tiles[i][j].data[1] = base_types::convertor<T2, U2>::convert(tmp[1]);
-                dst.tiles[i][j].data[2] = base_types::convertor<T2, U2>::convert(tmp[2]);
-                dst.tiles[i][j].data[3] = base_types::convertor<T2, U2>::convert(tmp[3]);
-            }
-            else if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row> && sizeof(typename ST::dtype) == 1) {
-                // handle the row-major layout for 8-bit types
-                int warp_group_16 = (warp_laneid / 16);  // divide each warp into two groups of 16 threads
-                int lane_in_16 = warp_laneid % 16;       // position in group of 16 threads
-                int row = (local_warpid*warp_height + i)*dst.tile_size_row + (lane_in_16 % 16); // find base row for warp in warpgroup and then distribute the 16 threads in the warp across the rows
-                int col = j*dst.tile_size_col + warp_group_16 * 16; // find base column and then *16 for second half of the warp
-
-                U2 tmp[4];
-                if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row>) {
-                    move<U2>::ldsm4(tmp[0], tmp[1], tmp[2], tmp[3], src.idx(shared_addr, {row, col}));
-                }
-                else {
-                    move<U2>::ldsm4t(tmp[0], tmp[2], tmp[1], tmp[3], src.idx(shared_addr, {row, col}));
-                }
-                dst.tiles[i][j].data[0] = base_types::convertor<T2, U2>::convert(tmp[0]);
-                dst.tiles[i][j].data[1] = base_types::convertor<T2, U2>::convert(tmp[1]);
-                dst.tiles[i][j].data[2] = base_types::convertor<T2, U2>::convert(tmp[2]);
-                dst.tiles[i][j].data[3] = base_types::convertor<T2, U2>::convert(tmp[3]);
-            }
-            else if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row> && sizeof(typename ST::dtype) == 4) {
-                // handle the row-major layout for 32-bit types
-                int row = (local_warpid*warp_height + i)*dst.tile_size_row + (warp_laneid / 4);
-                int col = j*dst.tile_size_col + 2*(warp_laneid % 4);
-                if constexpr (ST::rows != ST::underlying_rows || ST::cols != ST::underlying_cols) { // subtile case
-                    row += src.row_offset;
-                    col += src.col_offset;
-                }
-                int blit = sizeof(typename ST::dtype) * ((warp_laneid%4) / 2);
-                U2 tmp[4];
-                static constexpr int swizzle_repeat = ST::swizzle_bytes * 8;
-                static constexpr int subtile_cols   = ST::swizzle_bytes / sizeof(U);
-                const int outer_idx = col/subtile_cols;
-                const uint32_t addr_1 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+0)*subtile_cols + col%subtile_cols);
-                const uint32_t addr_2 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+8)*subtile_cols + col%subtile_cols);
-                const int swizzle_1 = blit ^ ((addr_1 % swizzle_repeat) >> 7) << 4;
-                const int swizzle_2 = blit ^ ((addr_2 % swizzle_repeat) >> 7) << 4;
-                move<U>::lds(tmp[0].x, (addr_1+ 0)^swizzle_1);
-                move<U>::lds(tmp[0].y, (addr_1+ 4)^swizzle_1);
-                move<U>::lds(tmp[2].x, (addr_1+32)^swizzle_1);
-                move<U>::lds(tmp[2].y, (addr_1+36)^swizzle_1);
-                move<U>::lds(tmp[1].x, (addr_2+ 0)^swizzle_2);
-                move<U>::lds(tmp[1].y, (addr_2+ 4)^swizzle_2);
-                move<U>::lds(tmp[3].x, (addr_2+32)^swizzle_2);
-                move<U>::lds(tmp[3].y, (addr_2+36)^swizzle_2);
-                dst.tiles[i][j].data[0] = base_types::convertor<T2, U2>::convert(tmp[0]);
-                dst.tiles[i][j].data[1] = base_types::convertor<T2, U2>::convert(tmp[1]);
-                dst.tiles[i][j].data[2] = base_types::convertor<T2, U2>::convert(tmp[2]);
-                dst.tiles[i][j].data[3] = base_types::convertor<T2, U2>::convert(tmp[3]);
-                if(blit) {
-                    #pragma unroll
-                    for(int k = 0; k < 4; k++) {
-                        dst.tiles[i][j].data[k] = T2{dst.tiles[i][j].data[k].y, dst.tiles[i][j].data[k].x};
-                    }
-                }
-            }
-            else {
-                // handle the column-major layout
-                int row = (local_warpid*warp_height + i)*dst.tile_size_row + 2*(warp_laneid % 4);
-                int col = j*dst.tile_size_col + (warp_laneid / 4);
-                U2 tmp[4];
-                move<U>::lds(tmp[0].x, src.idx(shared_addr, {row+0, col+0}));
-                move<U>::lds(tmp[0].y, src.idx(shared_addr, {row+1, col+0}));
-                move<U>::lds(tmp[1].x, src.idx(shared_addr, {row+0, col+8}));
-                move<U>::lds(tmp[1].y, src.idx(shared_addr, {row+1, col+8}));
-                move<U>::lds(tmp[2].x, src.idx(shared_addr, {row+8, col+0}));
-                move<U>::lds(tmp[2].y, src.idx(shared_addr, {row+9, col+0}));
-                move<U>::lds(tmp[3].x, src.idx(shared_addr, {row+8, col+8}));
-                move<U>::lds(tmp[3].y, src.idx(shared_addr, {row+9, col+8}));
-                dst.tiles[i][j].data[0] = base_types::convertor<T2, U2>::convert(tmp[0]);
-                dst.tiles[i][j].data[1] = base_types::convertor<T2, U2>::convert(tmp[1]);
-                dst.tiles[i][j].data[2] = base_types::convertor<T2, U2>::convert(tmp[2]);
-                dst.tiles[i][j].data[3] = base_types::convertor<T2, U2>::convert(tmp[3]);
-            }
-        }
-    }
-}
-
-
-/**
- * @brief Collaboratively store data into a shared tile from register tiles split across a warpgroup.
- *
- * @tparam RT The register tile type
- * @tparam ST The shared tile type
- * @param dst[out] The destination shared tile.
- * @param src[in]  The source register tile.
- */
-template<ducks::st::all ST, ducks::rt::all RT>
-__device__ inline static void store(ST &dst, const RT &src) {
-    constexpr int height = ST::height;
-    constexpr int warp_height = RT::height;
-    static_assert(height%GROUP_WARPS == 0, "Group load / store requires tile height to be a multiple of GROUP_WARPS.");
-    static_assert(height%warp_height == 0, "Group load / store requires tile height to be a multiple of the RT height.");
-    static_assert(ST::width==RT::width, "Group load / store requires tile widths to match.");
-    int local_warpid;
-    if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
-    else local_warpid = warpid();
-    using T2 = RT::dtype;
-    using U  = ST::dtype;
-    using T  = base_types::packing<T2>::unpacked_type;
-    using U2 = base_types::packing<U>::packed_type;
-    int warp_laneid = ::kittens::laneid();
-
-    // convert to shared state space
-    uint32_t shared_addr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
-
-    #pragma unroll
-    for(int i = 0; i < warp_height; i++) {
-        #pragma unroll
-        for(int j = 0; j < src.width; j++) {
-            if constexpr (sizeof(typename ST::dtype) == 2) {
-                // handle the row-major layout
-                U2 tmp[4];
-                tmp[0] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[0]);
-                tmp[1] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[1]);
-                tmp[2] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[2]);
-                tmp[3] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[3]);
-#ifdef KITTENS_HOPPER
-                int row = (local_warpid*warp_height + i)*src.tile_size_row + (warp_laneid % 16);
-                int col = j*src.tile_size_col + (warp_laneid / 16) * 8;
-                if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row>) {
-                    move<U2>::stsm4(dst.idx(shared_addr, {row, col}), tmp[0], tmp[1], tmp[2], tmp[3]);
-                }
-                else {
-                    move<U2>::stsm4t(dst.idx(shared_addr, {row, col}), tmp[0], tmp[2], tmp[1], tmp[3]);
-                }
-#else
-                if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row>) {
-                    int row = (local_warpid*warp_height + i)*src.tile_size_row + (warp_laneid / 4);
-                    int col = j*src.tile_size_col + 2*(warp_laneid % 4);
-                    move<U2>::sts(dst.idx(shared_addr, {row+0, col+0}), tmp[0]);
-                    move<U2>::sts(dst.idx(shared_addr, {row+8, col+0}), tmp[1]);
-                    move<U2>::sts(dst.idx(shared_addr, {row+0, col+8}), tmp[2]);
-                    move<U2>::sts(dst.idx(shared_addr, {row+8, col+8}), tmp[3]);
-                }
-                else {
-                    int row = (local_warpid*warp_height + i)*src.tile_size_row + 2*(warp_laneid % 4);
-                    int col = j*src.tile_size_col + (warp_laneid / 4);
-                    move<U>::sts(dst.idx(shared_addr, {row+0, col+0}), tmp[0].x);
-                    move<U>::sts(dst.idx(shared_addr, {row+1, col+0}), tmp[0].y);
-                    move<U>::sts(dst.idx(shared_addr, {row+0, col+8}), tmp[1].x);
-                    move<U>::sts(dst.idx(shared_addr, {row+1, col+8}), tmp[1].y);
-                    move<U>::sts(dst.idx(shared_addr, {row+8, col+0}), tmp[2].x);
-                    move<U>::sts(dst.idx(shared_addr, {row+9, col+0}), tmp[2].y);
-                    move<U>::sts(dst.idx(shared_addr, {row+8, col+8}), tmp[3].x);
-                    move<U>::sts(dst.idx(shared_addr, {row+9, col+8}), tmp[3].y);
-                }
-#endif
-            }
-            else if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row> && sizeof(typename ST::dtype) == 1) { 
-                // handle the row-major layout for 8-bit types
-
-                int warp_group_16 = (warp_laneid / 16);  // divide each warp into two groups of 16 threads
-                int lane_in_16 = warp_laneid % 16;       // position in group of 16 threads
-                int row = (local_warpid*warp_height + i)*src.tile_size_row + (lane_in_16 % 16); // find base row for warp in warpgroup and then distribute the 16 threads in the warp across the rows
-                int col = j*src.tile_size_col + warp_group_16 * 16; // find base column and then *16 for second half of the warp
-
-                U2 tmp[4];
-                tmp[0] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[0]);
-                tmp[1] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[1]);
-                tmp[2] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[2]);
-                tmp[3] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[3]);
-                if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row>) {
-                    move<U2>::stsm4(dst.idx(shared_addr, {row, col}), tmp[0], tmp[1], tmp[2], tmp[3]);
-                }
-                else {
-                    move<U2>::stsm4t(dst.idx(shared_addr, {row, col}), tmp[0], tmp[2], tmp[1], tmp[3]);
-                }
-            }
-            else if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row> && sizeof(typename ST::dtype) == 4) {
-                // handle the row-major layout for 32-bit types
-                int row = (local_warpid*warp_height + i)*src.tile_size_row + (warp_laneid / 4);
-                int col = j*src.tile_size_col + 2*(warp_laneid % 4);
-                if constexpr (ST::rows != ST::underlying_rows || ST::cols != ST::underlying_cols) { // subtile case
-                    row += dst.row_offset;
-                    col += dst.col_offset;
-                }
-                int blit = sizeof(typename ST::dtype) * ((warp_laneid%4) / 2);
-                T2 reg_tmp[4];
-                if(blit) {
-                    #pragma unroll
-                    for(int k = 0; k < 4; k++) {
-                        reg_tmp[k] = T2{src.tiles[i][j].data[k].y, src.tiles[i][j].data[k].x};
-                    }
-                }
-                else {
-                    #pragma unroll
-                    for(int k = 0; k < 4; k++) {
-                        reg_tmp[k] = src.tiles[i][j].data[k];
-                    }
-                }
-                U2 tmp[4];
-                tmp[0] = base_types::convertor<U2, T2>::convert(reg_tmp[0]);
-                tmp[1] = base_types::convertor<U2, T2>::convert(reg_tmp[1]);
-                tmp[2] = base_types::convertor<U2, T2>::convert(reg_tmp[2]);
-                tmp[3] = base_types::convertor<U2, T2>::convert(reg_tmp[3]);
-                static constexpr int swizzle_repeat = ST::swizzle_bytes * 8;
-                static constexpr int subtile_cols   = ST::swizzle_bytes / sizeof(U);
-                const int outer_idx = col/subtile_cols;
-                const uint32_t addr_1 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+0)*subtile_cols + col%subtile_cols);
-                const uint32_t addr_2 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+8)*subtile_cols + col%subtile_cols);
-                const int swizzle_1 = blit ^ ((addr_1 % swizzle_repeat) >> 7) << 4;
-                const int swizzle_2 = blit ^ ((addr_2 % swizzle_repeat) >> 7) << 4;
-                move<U>::sts((addr_1+ 0)^swizzle_1, tmp[0].x);
-                move<U>::sts((addr_1+ 4)^swizzle_1, tmp[0].y);
-                move<U>::sts((addr_1+32)^swizzle_1, tmp[2].x);
-                move<U>::sts((addr_1+36)^swizzle_1, tmp[2].y);
-                move<U>::sts((addr_2+ 0)^swizzle_2, tmp[1].x);
-                move<U>::sts((addr_2+ 4)^swizzle_2, tmp[1].y);
-                move<U>::sts((addr_2+32)^swizzle_2, tmp[3].x);
-                move<U>::sts((addr_2+36)^swizzle_2, tmp[3].y);
-            }
-            else {
-                // handle the column-major layout
-                int row = (local_warpid*warp_height + i)*src.tile_size_row + 2*(warp_laneid % 4);
-                int col = j*src.tile_size_col + (warp_laneid / 4);
-                U2 tmp[4];
-                tmp[0] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[0]);
-                tmp[1] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[1]);
-                tmp[2] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[2]);
-                tmp[3] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[3]);
-                move<U>::sts(dst.idx(shared_addr, {row+0, col+0}), tmp[0].x);
-                move<U>::sts(dst.idx(shared_addr, {row+1, col+0}), tmp[0].y);
-                move<U>::sts(dst.idx(shared_addr, {row+0, col+8}), tmp[1].x);
-                move<U>::sts(dst.idx(shared_addr, {row+1, col+8}), tmp[1].y);
-                move<U>::sts(dst.idx(shared_addr, {row+8, col+0}), tmp[2].x);
-                move<U>::sts(dst.idx(shared_addr, {row+9, col+0}), tmp[2].y);
-                move<U>::sts(dst.idx(shared_addr, {row+8, col+8}), tmp[3].x);
-                move<U>::sts(dst.idx(shared_addr, {row+9, col+8}), tmp[3].y);
-            }
-        }
-    }
-}
-
-// Load and store of vectors from/to shared tiles.
-
-template<ducks::rv::naive_layout RV, ducks::st::all ST>
-__device__ inline static auto load(RV &dst, const ST &src, int2 row_col) {
-    KITTENS_CHECK_WARP;
-    static_assert(ST::cols>=RV::length, "Shared tile must be at least as wide as the vector.");
-    using T = RV::T;
-    using U = ST::T;
-    int warp_laneid = ::kittens::laneid();
-
-    // convert to shared state space
-    uint32_t shared_addr = static_cast<uint32_t>(__cvta_generic_to_shared(&src.data[0]));
-
-    #pragma unroll
-    for(int col = warp_laneid; col < dst.length; col+=WARP_THREADS) {
-        U tmp;
-        move<U>::lds(tmp, src.idx(shared_addr, {row_col.x, row_col.y + col}));
-        dst.data[col/WARP_THREADS][0] = base_types::convertor<T, U>::convert(tmp);
-    }
-}
-
-template<ducks::rv::naive_layout RV, ducks::st::all ST>
-__device__ inline static auto store(ST &dst, const RV &src, int2 row_col) {
-    KITTENS_CHECK_WARP;
-    static_assert(ST::cols>=RV::length, "Shared tile must be at least as wide as the vector.");
-    using T = RV::T;
-    using U = ST::T;
-    int warp_laneid = ::kittens::laneid();
-
-    // convert to shared state space
-    uint32_t shared_addr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
-
-    #pragma unroll
-    for(int col = warp_laneid; col < src.length; col+=WARP_THREADS) {
-        U tmp = base_types::convertor<U, T>::convert(src.data[col/WARP_THREADS][0]);
-        move<U>::sts(dst.idx(shared_addr, {row_col.x, row_col.y + col}), tmp);
-    }
-}
@@ -1,325 +0,0 @@
-/**
- * @file
- * @brief Group (collaborative warp) ops for loading tensor tiles into register tiles.
- */
-
-/**
- * @brief Load data from a tensor tile into a register tile.
- *
- * @tparam RT The register tile type
- * @tparam TM The tensor memory tile type
- * @param dst[out] The destination register tile.
- * @param src[in]  The source tensor tile.
- */
-template<ducks::rt::row_layout RT, ducks::tt::all TM>
-__device__ inline static void load_async(RT &dst, const TM &src) {
-    if constexpr (GROUP_WARPS == 1) {
-        static_assert(RT::height == TM::height, "register tile and tensor tile must match height");
-        static_assert(RT::width == TM::width, "register tile and tensor tile must match width");
-
-        using T2 = RT::dtype;
-        using U  = typename TM::dtype;
-        using U2 = base_types::packing<typename TM::dtype>::packed_type;
-
-        if constexpr (sizeof(typename TM::dtype) == 1) {
-            #pragma unroll
-            for(int i = 0; i < dst.height; i++) {
-                #pragma unroll
-                for(int j = 0; j < dst.width; j++) {
-                    asm volatile(
-                        "tcgen05.ld.sync.aligned.16x128b.x2.pack::16b.b32 {%0, %1, %2, %3}, [%4];\n"
-                        : "=r"(*(uint32_t*) &dst.tiles[i][j].data[0]),
-                            "=r"(*(uint32_t*) &dst.tiles[i][j].data[1]),
-                            "=r"(*(uint32_t*) &dst.tiles[i][j].data[2]),
-                            "=r"(*(uint32_t*) &dst.tiles[i][j].data[3])
-                        : "r"(src.addr + ((i * dst.tile_size_row) << 16) + (j * dst.tile_size_col)/(4/(uint32_t)sizeof(U)))
-                    );
-                }
-            }
-        } else if constexpr (sizeof(typename TM::dtype) == 2) {
-            #pragma unroll
-            for(int i = 0; i < dst.height; i++) {
-                #pragma unroll
-                for(int j = 0; j < dst.width; j++) {
-                    asm volatile(
-                        "tcgen05.ld.sync.aligned.16x128b.x2.pack::16b.b32 {%0, %1, %2, %3}, [%4];\n"
-                        : "=r"(*(uint32_t*) &dst.tiles[i][j].data[0]),
-                            "=r"(*(uint32_t*) &dst.tiles[i][j].data[1]),
-                            "=r"(*(uint32_t*) &dst.tiles[i][j].data[2]),
-                            "=r"(*(uint32_t*) &dst.tiles[i][j].data[3])
-                        : "r"(src.addr + ((i * dst.tile_size_row) << 16) + (j * dst.tile_size_col))
-                    );
-                }
-            }
-        }
-        else if constexpr (sizeof(typename TM::dtype) == 4) {
-            #pragma unroll
-            for(int i = 0; i < dst.height; i++) {
-                if constexpr (dst.width%4 == 0) {
-                    #pragma unroll
-                    for(int j = 0; j < dst.width; j+=4) {
-                        U2 data[16];
-                        asm volatile(
-                            "tcgen05.ld.sync.aligned.16x256b.x8.b32 {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];\n"
-                            : "=f"(data[0].x), "=f"(data[0].y),
-                            "=f"(data[1].x), "=f"(data[1].y),
-                            "=f"(data[2].x), "=f"(data[2].y),
-                            "=f"(data[3].x), "=f"(data[3].y),
-                            "=f"(data[4].x), "=f"(data[4].y),
-                            "=f"(data[5].x), "=f"(data[5].y),
-                            "=f"(data[6].x), "=f"(data[6].y),
-                            "=f"(data[7].x), "=f"(data[7].y),
-                            "=f"(data[8].x), "=f"(data[8].y),
-                            "=f"(data[9].x), "=f"(data[9].y),
-                            "=f"(data[10].x), "=f"(data[10].y),
-                            "=f"(data[11].x), "=f"(data[11].y),
-                            "=f"(data[12].x), "=f"(data[12].y),
-                            "=f"(data[13].x), "=f"(data[13].y),
-                            "=f"(data[14].x), "=f"(data[14].y),
-                            "=f"(data[15].x), "=f"(data[15].y)
-                            : "r"(src.addr + ((i * dst.tile_size_row) << 16) + (j * dst.tile_size_col)/(4/(uint32_t)sizeof(U)))
-                        );
-                        #pragma unroll
-                        for(int k = 0; k < 4; k++) {
-                            dst.tiles[i][j+0].data[k] = base_types::convertor<T2, U2>::convert(data[k]);
-                            dst.tiles[i][j+1].data[k] = base_types::convertor<T2, U2>::convert(data[k+4]);
-                            dst.tiles[i][j+2].data[k] = base_types::convertor<T2, U2>::convert(data[k+8]);
-                            dst.tiles[i][j+3].data[k] = base_types::convertor<T2, U2>::convert(data[k+12]);
-                        }
-                    }
-                }
-                else if constexpr (dst.width%2 == 0) {
-                    #pragma unroll
-                    for(int j = 0; j < dst.width; j+=2) {
-                        U2 data[8];
-                        asm volatile(
-                            "tcgen05.ld.sync.aligned.16x256b.x4.b32 {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16];\n"
-                            : "=f"(data[0].x), "=f"(data[0].y),
-                            "=f"(data[1].x), "=f"(data[1].y),
-                            "=f"(data[2].x), "=f"(data[2].y),
-                            "=f"(data[3].x), "=f"(data[3].y),
-                            "=f"(data[4].x), "=f"(data[4].y),
-                            "=f"(data[5].x), "=f"(data[5].y),
-                            "=f"(data[6].x), "=f"(data[6].y),
-                            "=f"(data[7].x), "=f"(data[7].y)
-                            : "r"(src.addr + ((i * dst.tile_size_row) << 16) + (j * dst.tile_size_col)/(4/(uint32_t)sizeof(U)))
-                        );
-                        #pragma unroll
-                        for(int k = 0; k < 4; k++) {
-                            dst.tiles[i][j+0].data[k] = base_types::convertor<T2, U2>::convert(data[k]);
-                            dst.tiles[i][j+1].data[k] = base_types::convertor<T2, U2>::convert(data[k+4]);
-                        }
-                    }
-                }
-                else {
-                    #pragma unroll
-                    for(int j = 0; j < dst.width; j++) {
-                        U2 data[4];
-                        asm volatile(
-                            "tcgen05.ld.sync.aligned.16x256b.x2.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n"
-                            : "=f"(data[0].x), "=f"(data[0].y),
-                            "=f"(data[1].x), "=f"(data[1].y),
-                            "=f"(data[2].x), "=f"(data[2].y),
-                            "=f"(data[3].x), "=f"(data[3].y)
-                            : "r"(src.addr + ((i * dst.tile_size_row) << 16) + (j * dst.tile_size_col)/(4/(uint32_t)sizeof(U)))
-                        );
-                        #pragma unroll
-                        for(int k = 0; k < 4; k++) {
-                            dst.tiles[i][j].data[k] = base_types::convertor<T2, U2>::convert(data[k]);
-                        }
-                    }
-                }
-            }
-        }
-    }
-    else {
-        static_assert(GROUP_WARPS==4 || GROUP_WARPS==8);
-        constexpr int warp_rows = TM::rows/GROUP_WARPS;
-        static_assert(TM::cols==RT::cols);
-        static_assert(warp_rows==RT::rows);
-        if constexpr (GROUP_WARPS == 4) {
-            auto src_subtile = src.template subtile<tt<typename TM::dtype, warp_rows, TM::cols>>(32*warpid(), 0);
-            ::kittens::group<1>::load_async(dst, src_subtile);
-        }
-        else {
-            auto src_subtile = src.template subtile<tt<typename TM::dtype, warp_rows, TM::cols>>(32*(warpid()%4)+16*(warpid()/4), 0);
-            ::kittens::group<1>::load_async(dst, src_subtile);
-        }
-    }
-}
-
-
-/**
- * @brief Store data into a tensor tile from a register tile.
- *
- * @tparam RT The register tile type
- * @tparam TM The tensor memory tile type
- * @param dst[out] The destination tensor tile.
- * @param src[in]  The source register tile.
- */
-template<ducks::rt::all RT, ducks::tt::all TM>
-__device__ inline static void store_async(TM &dst, const RT &src) {
-    if constexpr (GROUP_WARPS == 1) {
-        static_assert(RT::height == TM::height, "register tile and tensor tile must match height");
-        static_assert(RT::width == TM::width, "register tile and tensor tile must match width");
-
-        using T2 = RT::dtype;
-        using T = base_types::packing<T2>::unpacked_type;
-        using U = TM::dtype;
-        using U2 = base_types::packing<U>::packed_type;
-
-        if constexpr (sizeof(typename TM::dtype) == 2) {
-            #pragma unroll
-            for(int i = 0; i < src.height; i++) {
-                if constexpr (src.width%4 == 0) {
-                    #pragma unroll
-                    for(int j = 0; j < src.width; j+=4) {
-                        asm volatile(
-                            "tcgen05.st.sync.aligned.16x128b.x8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};\n"
-                            :: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
-                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[0]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[1]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[2]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[3]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[0]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[1]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[2]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[3]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+2].data[0]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+2].data[1]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+2].data[2]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+2].data[3]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+3].data[0]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+3].data[1]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+3].data[2]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+3].data[3])
-                        );
-                    }
-                }
-                else if constexpr (src.width%2 == 0) {
-                    #pragma unroll
-                    for(int j = 0; j < src.width; j+=2) {
-                        asm volatile(
-                            "tcgen05.st.sync.aligned.16x128b.x4.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};\n"
-                            :: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
-                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[0]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[1]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[2]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[3]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[0]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[1]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[2]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[3])
-                        );
-                    }
-                }
-                else {
-                    #pragma unroll
-                    for(int j = 0; j < src.width; j++) {
-                        asm volatile(
-                            "tcgen05.st.sync.aligned.16x128b.x2.b32 [%0], {%1, %2, %3, %4};\n"
-                            :: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
-                            "r"(*(uint32_t*)&src.tiles[i][j].data[0]),
-                            "r"(*(uint32_t*)&src.tiles[i][j].data[1]),
-                            "r"(*(uint32_t*)&src.tiles[i][j].data[2]),
-                            "r"(*(uint32_t*)&src.tiles[i][j].data[3])
-                        );
-                    }
-                }
-            }
-        }
-        else if constexpr (sizeof(typename TM::dtype) == 4) {
-            #pragma unroll
-            for(int i = 0; i < src.height; i++) {
-                if constexpr(src.width%4 == 0) {
-                    #pragma unroll
-                    for(int j = 0; j < src.width; j+=4) {
-                        U2 data[16];
-                        #pragma unroll
-                        for(int k = 0; k < 4; k++) {
-                            data[k] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[k]);
-                            data[k+4] = base_types::convertor<U2, T2>::convert(src.tiles[i][j+1].data[k]);
-                            data[k+8] = base_types::convertor<U2, T2>::convert(src.tiles[i][j+2].data[k]);
-                            data[k+12] = base_types::convertor<U2, T2>::convert(src.tiles[i][j+3].data[k]);
-                        }
-                        asm volatile(
-                            "tcgen05.st.sync.aligned.16x256b.x8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32};\n"
-                            :: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
-                            "f"(data[0].x), "f"(data[0].y),
-                            "f"(data[1].x), "f"(data[1].y),
-                            "f"(data[2].x), "f"(data[2].y),
-                            "f"(data[3].x), "f"(data[3].y),
-                            "f"(data[4].x), "f"(data[4].y),
-                            "f"(data[5].x), "f"(data[5].y),
-                            "f"(data[6].x), "f"(data[6].y),
-                            "f"(data[7].x), "f"(data[7].y),
-                            "f"(data[8].x), "f"(data[8].y),
-                            "f"(data[9].x), "f"(data[9].y),
-                            "f"(data[10].x), "f"(data[10].y),
-                            "f"(data[11].x), "f"(data[11].y),
-                            "f"(data[12].x), "f"(data[12].y),
-                            "f"(data[13].x), "f"(data[13].y),
-                            "f"(data[14].x), "f"(data[14].y),
-                            "f"(data[15].x), "f"(data[15].y)
-                        );
-                    }
-                }
-                else if constexpr(src.width%2 == 0) {
-                    #pragma unroll
-                    for(int j = 0; j < src.width; j+=2) {
-                        U2 data[8];
-                        #pragma unroll
-                        for(int k = 0; k < 4; k++) {
-                            data[k] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[k]);
-                            data[k+4] = base_types::convertor<U2, T2>::convert(src.tiles[i][j+1].data[k]);
-                        }
-                        asm volatile(
-                            "tcgen05.st.sync.aligned.16x256b.x4.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};\n"
-                            :: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
-                            "f"(data[0].x), "f"(data[0].y),
-                            "f"(data[1].x), "f"(data[1].y),
-                            "f"(data[2].x), "f"(data[2].y),
-                            "f"(data[3].x), "f"(data[3].y),
-                            "f"(data[4].x), "f"(data[4].y),
-                            "f"(data[5].x), "f"(data[5].y),
-                            "f"(data[6].x), "f"(data[6].y),
-                            "f"(data[7].x), "f"(data[7].y)
-                        );
-                    }
-                }
-                else {
-                    #pragma unroll
-                    for(int j = 0; j < src.width; j++) {
-                        U2 data[4];
-                        #pragma unroll
-                        for(int k = 0; k < 4; k++) {
-                            data[k] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[k]);
-                        }
-                        asm volatile(
-                            "tcgen05.st.sync.aligned.16x256b.x2.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};\n"
-                            :: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
-                            "f"(data[0].x), "f"(data[0].y),
-                            "f"(data[1].x), "f"(data[1].y),
-                            "f"(data[2].x), "f"(data[2].y),
-                            "f"(data[3].x), "f"(data[3].y)
-                        );
-                    }
-                }
-            }
-        }
-    }
-    else {
-        static_assert(GROUP_WARPS==4 || GROUP_WARPS==8);
-        constexpr int warp_rows = TM::rows/GROUP_WARPS;
-        static_assert(TM::cols==RT::cols);
-        static_assert(warp_rows==RT::rows);
-        if constexpr (GROUP_WARPS == 4) {
-            auto dst_subtile = dst.template subtile<tt<typename TM::dtype, warp_rows, TM::cols>>(32*warpid(), 0);
-            ::kittens::group<1>::store_async(dst_subtile, src);
-        }
-        else {
-            auto dst_subtile = dst.template subtile<tt<typename TM::dtype, warp_rows, TM::cols>>(32*(warpid()%4)+16*(warpid()/4), 0);
-            ::kittens::group<1>::store_async(dst_subtile, src);
-        }
-    }
-}
@@ -1,16 +0,0 @@
-/**
- * @file
- * @brief An aggregate header of group memory operations on tiles.
- */
-
-#include "shared_to_register.cuh"
-#include "global_to_register.cuh"
-#include "global_to_shared.cuh"
-#ifdef KITTENS_BLACKWELL
-#include "tensor_to_register.cuh"
-#endif
-
-#include "complex/complex_shared_to_register.cuh"
-#include "complex/complex_global_to_register.cuh"
-#include "complex/complex_global_to_shared.cuh"
-
@@ -1,134 +0,0 @@
-/**
- * @file
- * @brief Functions for a group scope to call tile TMA functions.
- */
-
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void prefetch(ST &dst, const GL &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::prefetch<axis, policy, ST, GL, COORD>(dst, src, idx); // Don't do the mask
-    }
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void prefetch(ST &dst, const GL &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::prefetch<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx); // Don't do the mask
-    }
-}
-
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_async(const GL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_async<axis, policy, ST, GL, COORD>(dst, src, idx); // Don't do the mask
-    }
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_async(const GL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
-    }
-}
-
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_async(const PGL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_async<axis, policy, ST, PGL, COORD>(dst, src, idx); // Don't do the mask
-    }
-}
-template<ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_async(const PGL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_async<dim::ROW, cache_policy::NORMAL, ST, PGL, COORD>(dst, src, idx);
-    }
-}
-
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_add_async(const GL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_add_async<axis, policy, ST, GL, COORD>(dst, src, idx); // Don't do the mask
-    }
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_add_async(const GL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_add_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
-    }
-}
-
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_add_async(const PGL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_add_async<axis, policy, ST, PGL, COORD>(dst, src, idx); // Don't do the mask
-    }
-}
-template<ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_add_async(const PGL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_add_async<dim::ROW, cache_policy::NORMAL, ST, PGL, COORD>(dst, src, idx);
-    }
-}
-
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_min_async(const GL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_min_async<axis, policy, ST, GL, COORD>(dst, src, idx); // Don't do the mask
-    }
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_min_async(const GL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_min_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
-    }
-}
-
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_min_async(const PGL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_min_async<axis, policy, ST, PGL, COORD>(dst, src, idx); // Don't do the mask
-    }
-}
-template<ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_min_async(const PGL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_min_async<dim::ROW, cache_policy::NORMAL, ST, PGL, COORD>(dst, src, idx);
-    }
-}
-
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_max_async(const GL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_max_async<axis, policy, ST, GL, COORD>(dst, src, idx); // Don't do the mask
-    }
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_max_async(const GL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_max_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
-    }
-}
-
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_max_async(const PGL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_max_async<axis, policy, ST, PGL, COORD>(dst, src, idx); // Don't do the mask
-    }
-}
-template<ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_max_async(const PGL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_max_async<dim::ROW, cache_policy::NORMAL, ST, PGL, COORD>(dst, src, idx);
-    }
-}
-
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar) {
-    if(laneid() == 0) {
-        ::kittens::tma::load_async<axis, policy, ST, GL, COORD>(dst, src, idx, bar); // Don't do the mask
-    }
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar) {
-    if(laneid() == 0) {
-        ::kittens::tma::load_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx, bar);
-    }
-}
@@ -1,33 +0,0 @@
-/**
- * @file
- * @brief Functions for a group scope to call tile TMA cluster functions.
- */
-
-
-#ifdef KITTENS_BLACKWELL
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) {
-    if(laneid() == 0) {
-        ::kittens::tma::cluster::load_async<axis, policy, ST, GL, COORD>(dst, src, idx, bar, cluster_mask, dst_mbar_cta);
-    }
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) {
-    if(laneid() == 0) {
-        ::kittens::tma::cluster::load_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx, bar, cluster_mask, dst_mbar_cta);
-    }
-}
-#else
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask) {
-    if(laneid() == 0) {
-        ::kittens::tma::cluster::load_async<axis, policy, ST, GL, COORD>(dst, src, idx, bar, cluster_mask);
-    }
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask) {
-    if(laneid() == 0) {
-        ::kittens::tma::cluster::load_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx, bar, cluster_mask);
-    }
-}
-#endif
@@ -1,68 +0,0 @@
-/**
- * @file
- * @brief Various utilities for group TMA memory operations.
- */
-
-/* ----------   Barrier functions for async load  ---------- */
-
-/**
-* @brief Sets the number of bytes expected at the semaphore.
-*
-* This function sets the number of bytes expected at the semaphore for the first thread in the warp.
-* It converts the semaphore pointer to a generic shared memory pointer and uses an inline assembly
-* instruction to set the expected number of bytes.
-*
-* @param semaphore Reference to the semaphore variable.
-* @param bytes The number of bytes expected at the semaphore.
-*/
-__device__ static inline void expect_bytes(semaphore& bar, uint32_t bytes) {
-    if(laneid() == 0) {
-        ::kittens::tma::expect_bytes(bar, bytes);
-    }
-}
-/**
-* @brief Sets the number of bytes expected at the semaphore.
-*
-* This function sets the number of bytes expected at the mbarrier before the transaction arrives.
-*/
-template<typename T, typename... args>
-__device__ static inline void expect(semaphore& bar, const T& _1, const args&... _2) {
-    expect_bytes(bar, size_bytes<T, args...>);
-}
-
-/* ----------   Synchronization functions for async store  ---------- */
-
-/**
- * @brief Commits previous asynchronous TMA stores to a group and performs them.
-*/
-__device__ static inline void store_commit_group() {
-    asm volatile("cp.async.bulk.commit_group;");
-}
-/**
- * @brief Waits for previous committed TMA store groups to complete.
- *
- * @tparam N The maximum number of remaining TMA store groups. Defaults to 0.
-*/
-template <int N=0>
-__device__ static inline void store_async_wait() {
-    asm volatile (
-        "cp.async.bulk.wait_group %0;"
-        :
-        : "n"(N)
-        : "memory"
-    );
-}
-/**
- * @brief Waits for previous committed TMA store groups to finish reading from shared memory.
- *
- * @tparam N The maximum number of remaining TMA store groups. Defaults to 0.
-*/
-template <int N=0>
-__device__ static inline void store_async_read_wait() {
-    asm volatile (
-        "cp.async.bulk.wait_group.read %0;"
-        :
-        : "n"(N)
-        : "memory"
-    );
-}
@@ -1,90 +0,0 @@
-
-/**
-* @brief Waits for the requested semaphore phase, at cluster scope
-*
-* @param semaphore Reference to the semaphore variable.
-* @param kPhaseBit The phase bit used for the semaphore.
-*/
-__device__ static inline void wait(semaphore& bar, int kPhaseBit) {
-    void const* const ptr = &bar;
-    uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr)); 
-
-    asm volatile (
-        "{\n"
-        ".reg .pred                P1;\n"
-        "LAB_WAIT:\n"
-        "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%0], %1;\n"
-        "@P1                       bra.uni DONE;\n"
-        "bra.uni                   LAB_WAIT;\n"
-        "DONE:\n"
-        "}\n"
-        :: "r"(mbar_ptr),
-        "r"(kPhaseBit)
-    );
-}
-
-/**
-* @brief Sets the number of bytes expected at the semaphore, assuming a multicast instruction.
-*
-* This function sets the number of bytes expected at the semaphore for the first thread in the warp.
-* It converts the semaphore pointer to a generic shared memory pointer and uses an inline assembly
-* instruction to set the expected number of bytes.
-* 
-* It's worth being aware that this function is particularly necessary for multicast loads, and
-* distributed shared memory can actually be done with a normal tma::expect followed by wait. See
-* the unit tests of dsmem for an example.
-*
-* @param semaphore Reference to the semaphore variable.
-* @param bytes The number of bytes expected at the semaphore.
-*/
-__device__ static inline void expect_bytes(semaphore& bar, uint32_t bytes, int dst_cta) {
-    if(laneid() == 0) {
-        ::kittens::tma::cluster::expect_bytes(bar, bytes, dst_cta);
-    }
-}
-/**
-* @brief Sets the number of bytes expected at the semaphore.
-*
-* This function sets the number of bytes expected at the semaphore for the first thread in the warp.
-* It converts the semaphore pointer to a generic shared memory pointer and uses an inline assembly
-* instruction to set the expected number of bytes.
-*
-* @tparam T The type of the data to be stored at the semaphore.
-* @param semaphore Reference to the semaphore variable.
-*/
-/**
-* @brief Sets the number of bytes expected at the semaphore.
-*
-* This function sets the number of bytes expected at the mbarrier before the transaction arrives.
-*/
-template<typename T, typename... args>
-__device__ static inline void expect(semaphore& bar, int dst_cta, const T& _1, const args&... _2) {
-    expect_bytes(bar, size_bytes<T, args...>, dst_cta);
-}
-
-/**
-* @brief Arrives at a semaphore in cluster scope.
-*
-* Marks a thread arrival at an mbarrier
-*
-* @param semaphore Reference to the semaphore variable.
-* @param kPhaseBit The phase bit used for the semaphore.
-*/
-__device__ static inline void arrive(semaphore& bar, int dst_cta, uint32_t count=1) {
-    if(laneid() == 0) {
-        ::kittens::tma::cluster::arrive(bar, dst_cta, count);
-    }
-}
-
-// Generic transfer
-__device__ static inline void store_async(void *dst, void *src, int dst_cta, uint32_t size_bytes, semaphore& bar) {
-    if(laneid() == 0) {
-        ::kittens::tma::cluster::store_async(dst, src, dst_cta, size_bytes, bar);
-    }
-}
-
-// Templated transfer for convenience
-template<typename T>
-__device__ static inline void store_async(T &dst_, T &src_, int dst_cta, semaphore& bar) {
-    store_async((void*)&dst_, (void*)&src_, dst_cta, size_bytes<T>, bar);
-}
@@ -1,168 +0,0 @@
-/**
- * @file
- * @brief Various utilities for group memory operations.
- */
-
-
-template<int N=0> __device__ static inline void load_async_wait(int bar_id) { // for completing (non-TMA) async loads
-    asm volatile("cp.async.wait_group %0;\n" : : "n"(N) : "memory");
-    sync(bar_id);
-}
-template<int N=0> __device__ static inline void load_async_wait() { // for completing (non-TMA) async loads
-    KITTENS_CHECK_WARP
-    asm volatile("cp.async.wait_group %0;\n" : : "n"(N) : "memory");
-    __syncwarp();
-}
-
-__device__ static inline void arrive(barrier<GROUP_WARPS> bar) {
-    asm volatile("bar.arrive %0, %1;\n" :: "r"(bar.barrier_id), "n"(GROUP_WARPS*WARP_THREADS) : "memory");
-}
-__device__ static inline void arrive_and_wait(barrier<GROUP_WARPS> bar) {
-    asm volatile("bar.sync %0, %1;\n" :: "r"(bar.barrier_id), "n"(GROUP_WARPS*WARP_THREADS) : "memory");
-}
-
-/**
- * @brief Initializes a synchronization semaphore with a transaction count and sets the expected number of bytes.
- *
- * This function sets up a semaphore that is used to synchronize threads within a block during asynchronous operations.
- * It initializes the semaphore with a thread count semaphore.
- *
- * Additionally, if it is given a shared tile type, it will also call `set_bytes` to prepare for the memory transaction.
- *
- * @param[out] semaphore The semaphore variable to initialize.
- * @param[in] tc The thread counter for the semaphore.
- */
-__device__ static inline void init_semaphore(semaphore& bar, int thread_count, int transaction_count=0) {
-    if (laneid() == 0) {
-        void const* const ptr = &bar;
-        uint32_t bar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr)); 
-
-        asm volatile (
-            "mbarrier.init.shared::cta.b64 [%0], %1;\n"
-            :: "r"(bar_ptr), "r"(thread_count+transaction_count)
-        );
-    }
-}
-/**
- * @brief Invalidate an mbarrier
- *
- * @param[out] semaphore The semaphore variable to initialize.
- * @param[in] tc The thread counter for the semaphore.
- */
-__device__ static inline void invalidate_semaphore(semaphore& bar) {
-    if (laneid() == 0) {
-        void const* const ptr = &bar;
-        uint32_t bar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr)); 
-        asm volatile (
-            "mbarrier.inval.shared::cta.b64 [%0];\n"
-            :: "r"(bar_ptr)
-        );
-    }
-}
-
-/**
-* @brief Arrives at a semaphore.
-*
-* Marks a warp arrival at an mbarrier
-*
-* @param semaphore Reference to the semaphore variable.
-* @param kPhaseBit The phase bit used for the semaphore.
-*/
-__device__ static inline void arrive(semaphore& sem) {
-    if(laneid() == 0) {
-            uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&sem)); 
-            asm volatile (
-                "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];\n"
-            :
-            : "r"(mbar_ptr)
-            : "memory"
-        );
-    }
-}
-template<int num_warps> __device__ static inline void arrive(barrier<num_warps> bar) {
-    asm volatile("bar.arrive %0, %1;\n" :: "r"(bar.barrier_id), "n"(num_warps*WARP_THREADS) : "memory");
-}
-
-#ifdef KITTENS_HOPPER
-/**
-* @brief Arrives at a semaphore.
-*
-* Marks a warp arrival at an mbarrier
-*
-* @param semaphore Reference to the semaphore variable.
-* @param kPhaseBit The phase bit used for the semaphore.
-*/
-__device__ static inline void arrive(semaphore& sem, uint32_t count) {
-    if(laneid() == 0) {
-        uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&sem));
-        asm volatile (
-            "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0], %1;\n"
-            :
-            : "r"(mbar_ptr), "r"(count)
-            : "memory"
-        );
-    }
-}
-#endif
-
-/**
-* @brief Waits for the requested semaphore phase.
-*
-* @param semaphore Reference to the semaphore variable.
-* @param kPhaseBit The phase bit used for the semaphore.
-*/
-__device__ static inline void wait(semaphore& sem, int kPhaseBit) {
-    void const* const ptr = &sem;
-    uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr)); 
-
-#ifdef KITTENS_HOPPER
-    asm volatile (
-        "{\n"
-        ".reg .pred                P1;\n"
-        "LAB_WAIT:\n"
-        "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
-        "@P1                       bra.uni DONE;\n"
-        "bra.uni                   LAB_WAIT;\n"
-        "DONE:\n"
-        "}\n"
-        :: "r"(mbar_ptr),
-        "r"(kPhaseBit)
-    );
-#else
-    asm volatile (
-        "{\n"
-        ".reg .pred                P1;\n"
-        "LAB_WAIT:\n"
-        "mbarrier.test_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
-        "@P1                       bra.uni DONE;\n"
-        "nanosleep.u32 5;\n" // wait a few nanoseconds on pre-Hopper architectures to save instruction issue slots
-        "bra.uni                   LAB_WAIT;\n"
-        "DONE:\n"
-        "}\n"
-        :: "r"(mbar_ptr),
-        "r"(kPhaseBit)
-    );
-#endif
-}
-
-/**
-* @brief Checks if the requested semaphore phase is ready.
-*
-* @param semaphore Reference to the semaphore variable.
-* @param kPhaseBit The phase bit used for the semaphore.
-*/
-__device__ static inline int test_wait(semaphore& sem, int kPhaseBit) {
-    void const* const ptr = &sem;
-    uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
-    int result;
-    asm volatile (
-        "{\n"
-        ".reg .pred P1;\n"
-        "mbarrier.test_wait.parity.shared::cta.b64 P1, [%1], %2;\n"
-        "selp.u32 %0,1,0,P1;"
-        "}\n"
-        : "=r"(result)
-        : "r"(mbar_ptr), "r"(kPhaseBit)
-    );
-    return result;
-}
@@ -1,138 +0,0 @@
-/**
- * @file
- * @brief Functions for a warpgroup to collaboratively transfer  data directly between global memory and registers and back.
- */
-
-/**
- * @brief Collaboratively loads data into register vectors from a source array in global memory.
- *
- * @tparam RV The register vector type.
- * @tparam U The data type of the source array.
- * @param[out] dst The destination register vector to load data into.
- * @param[in] src The source array in global memory to load data from.
- */
-template<ducks::rv::all RV, ducks::gl::all GL>
-__device__ inline static void load(RV &dst, const GL &src, const coord<rv<typename RV::T, GROUP_WARPS*RV::length, typename RV::layout>> &idx) {
-    if constexpr (GROUP_WARPS == 1) {
-        using T2 = RV::dtype;
-        using U = typename GL::dtype;
-        using U2 = base_types::packing<U>::packed_type;
-        using T = base_types::packing<T2>::unpacked_type;
-
-        U *src_ptr = (U*)&src[(idx.template unit_coord<-1, 3>())];
-        int laneid = ::kittens::laneid();
-        
-        if constexpr (std::is_same_v<typename RV::layout, align_l>) {
-            #pragma unroll
-            for(auto w = 0; w < (dst.outer_dim+3)/4; w++) {
-                int idx = w*64 + (laneid/4)*8 + 2*(laneid%4);
-                int o_dim = w*4 + (laneid/4) / 2;
-                int i_dim = (laneid/4) % 2;
-                // this should be a maximally coalesced load.
-                if(idx < dst.outer_dim*16)
-                    dst[o_dim][i_dim] = base_types::convertor<T2, U2>::convert(*(U2*)&src_ptr[idx]);
-            }
-            // now we need to do a bunch of shuffle_sync's to make sure everyone has everything they need.
-            #pragma unroll
-            for(auto w = 0; w < dst.outer_dim; w++) {
-                int leader = 8*(w%4) + (laneid%4); // repeats every 64 columns
-                dst[w][0] = packed_shfl_sync(MASK_ALL, dst[w][0], leader);
-                dst[w][1] = packed_shfl_sync(MASK_ALL, dst[w][1], leader+4);
-            }
-        }
-        else if constexpr (std::is_same_v<typename RV::layout, ortho_l>) {
-            // really hoping https://stackoverflow.com/questions/15029765/is-coalescing-triggered-for-accessing-memory-in-reverse-order is still true
-            // otherwise there will be some pain :/
-            #pragma unroll
-            for(auto w = 0; w < (dst.outer_dim+1)/2; w++) {
-                int idx = w*32 + (laneid%4)*8 + (laneid/4);
-                int o_dim = w*2 + (laneid%4) / 2;
-                // this should be a maximally coalesced load.
-                if(idx < dst.outer_dim*16) {
-                    T tmp = base_types::convertor<T, U>::convert(src_ptr[idx]);
-                    if(laneid%2==0) dst[o_dim][0].x = tmp;
-                    else dst[o_dim][0].y = tmp;
-                }
-            }
-            // now we need to do a bunch of shuffle_sync's to make sure everyone has everything they need.
-            #pragma unroll
-            for(auto w = 0; w < dst.outer_dim; w++) {
-                int leader = (laneid/4)*4 + 2*(w%2); // repeats every 64 columns
-                dst[w][0].x = __shfl_sync(MASK_ALL, dst[w][0].x, leader);
-                dst[w][0].y = __shfl_sync(MASK_ALL, dst[w][0].y, leader+1);
-            }
-        }
-        else if constexpr (std::is_same_v<typename RV::layout, naive_l>) {
-            #pragma unroll
-            for(auto w = 0; w < dst.outer_dim; w++) {
-                if(w < dst.outer_dim-1 || dst.length%32 == 0 || laneid<16) {
-                    dst[w][0] = base_types::convertor<T, U>::convert(src_ptr[w*32 + laneid]);
-                }
-            }
-        }
-    }
-    else {
-        // Call warp level load
-        ::kittens::group<1>::load(dst, src, coord<RV>(idx.b, idx.d, idx.r, idx.c*GROUP_WARPS+warpid()));
-    }
-}
-/**
- * @brief Collaboratively stores data from register vectors to a destination array in global memory.
- *
- * @tparam RV The register vector type.
- * @tparam U The data type of the destination array.
- * @param[out] dst The destination array in global memory to store data into.
- * @param[in] src The source register vector to store data from.
- */
-template<ducks::rv::all RV, ducks::gl::all GL>
-__device__ inline static void store(GL &dst, const RV &src, const coord<rv<typename RV::T, GROUP_WARPS*RV::length, typename RV::layout>> &idx) {
-    if constexpr (GROUP_WARPS == 1) {
-        using T2 = RV::dtype;
-        using U = typename GL::dtype;
-        using U2 = base_types::packing<U>::packed_type;
-        using T = base_types::packing<T2>::unpacked_type;
-        
-        U *dst_ptr = (U*)&dst[(idx.template unit_coord<-1, 3>())];
-        int laneid = ::kittens::laneid();
-        
-        if constexpr (std::is_same_v<typename RV::layout, align_l>) {
-            #pragma unroll
-            for(auto w = 0; w < (src.outer_dim+3)/4; w++) {
-                int idx = w*64 + (laneid/4)*8 + 2*(laneid%4);
-                int o_dim = w*4 + (laneid/4) / 2;
-                int i_dim = (laneid/4) % 2;
-                // this should be a maximally coalesced store. I hope!
-                if(idx < src.outer_dim*16)
-                    *(U2*)&dst_ptr[idx] = base_types::convertor<U2, T2>::convert(src[o_dim][i_dim]);
-            }
-        }
-        else if constexpr (std::is_same_v<typename RV::layout, ortho_l>) {
-            // really hoping https://stackoverflow.com/questions/15029765/is-coalescing-triggered-for-accessing-memory-in-reverse-order is still true
-            // otherwise there will be some pain :/
-            #pragma unroll
-            for(auto w = 0; w < (src.outer_dim+1)/2; w++) {
-                int idx = w*32 + (laneid%4)*8 + (laneid/4);
-                int o_dim = w*2 + (laneid%4) / 2;
-                // this should be a maximally coalesced load.
-                if(idx < src.outer_dim*16) {
-                    U tmp;
-                    if(laneid%2==0) tmp = base_types::convertor<U, T>::convert(src[o_dim][0].x);
-                    else tmp = base_types::convertor<U, T>::convert(src[o_dim][0].y);
-                    dst_ptr[idx] = tmp;
-                }
-            }
-        }
-        else if constexpr (std::is_same_v<typename RV::layout, naive_l>) {
-            #pragma unroll
-            for(auto w = 0; w < src.outer_dim; w++) {
-                if(w < src.outer_dim-1 || src.length%32 == 0 || laneid<16) {
-                    dst_ptr[w*32 + laneid] = base_types::convertor<U, T>::convert(src[w][0]);
-                }
-            }
-        }
-    }
-    else {
-        // Call warp level store
-        ::kittens::group<1>::store(dst, src, coord<RV>(idx.b, idx.d, idx.r, idx.c*GROUP_WARPS+warpid()));
-    }
-}
--- a/Show More
+++ b/Show More