enable rangeify const folding

tags on bufferize are the tensor tags (#12180 )
delete bufferized consts + output noops (#12163 )
2026-08-15 22:18:27 +00:00 · 2025-09-15 11:50:13 +08:00 · 2025-09-15 11:46:03 +08:00 · 2025-09-15 11:07:44 +08:00 · 2025-09-15 02:53:05 +03:00 · 2025-09-15 02:33:50 +03:00
394 changed files with 35612 additions and 234532 deletions
@@ -112,6 +112,15 @@ runs:
        fi

    # ******************* apt *******************
+    - name: Setup apt
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
+      shell: bash
+      run: |
+        sudo chown -R $USER:$USER /var/cache/apt/archives
+
+        echo 'Acquire::GzipIndexes "true";' | sudo tee /etc/apt/apt.conf.d/gzip
+        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
+        echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' | sudo tee -a /etc/apt/apt.conf.d/99keep-debs

    - name: Add OpenCL Repo
      if: inputs.opencl == 'true' && runner.os == 'Linux'
@@ -135,14 +144,11 @@ runs:
        wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
        echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-20 main" | sudo tee /etc/apt/sources.list.d/llvm.list

-    - name: apt-get update + install
+    - name: Compute Package List + Hash
      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
+      id: apt-pkgs
      shell: bash
      run: |
-        echo 'Acquire::GzipIndexes "true";' | sudo tee /etc/apt/apt.conf.d/gzip
-        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
-        sudo apt -qq update || true
-
        pkgs=""
        # **** OpenCL ****
        if [[ "${{ inputs.opencl }}" == "true" ]]; then
@@ -153,7 +159,7 @@ runs:
        fi
        # **** AMD ****
        if [[ "${{ inputs.amd }}" == "true" ]]; then
-          pkgs+=" hsa-rocr comgr hsa-rocr-dev liburing-dev libc6-dev"
+          pkgs+=" hsa-rocr comgr hsa-rocr-dev liburing-dev libibverbs-dev libc6-dev"
        fi
        # **** CUDA ****
        if [[ "${{ inputs.cuda }}" == "true" ]]; then
@@ -169,13 +175,30 @@ runs:
          pkgs+=" libllvm20 clang-20 lld-20"
        fi

+        echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
+        echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"
+
+    - name: Cache apt
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
+      uses: actions/cache@v4
+      with:
+        path: /var/cache/apt/archives/
+        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.APT_CACHE_VERSION }}
+
+    - name: Run apt Update + Install
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
+      shell: bash
+      run: |
+        sudo apt -qq update || true
+
        # ******** do install ********
-        if [[ -n "$pkgs" ]]; then
-          sudo apt-get -y --allow-unauthenticated --no-install-recommends install $pkgs
+        if [[ -n "${{ steps.apt-pkgs.outputs.pkgs }}" ]]; then
+          sudo apt-get -y --allow-unauthenticated --no-install-recommends install ${{ steps.apt-pkgs.outputs.pkgs }}
        fi

-    # **** AMD ****
+        sudo chown -R $USER:$USER /var/cache/apt/archives/

+    # **** AMD ****
    - name: Setup AMD (Linux)
      if: inputs.amd == 'true' && runner.os == 'Linux'
      shell: bash
@@ -202,16 +225,25 @@ runs:
    - name: Install gpuocelot dependencies (MacOS)
      if: inputs.ocelot == 'true' && runner.os == 'macOS'
      shell: bash
-      run: brew install --quiet cmake ninja llvm@15 zlib glew flex bison boost zstd ncurses
+      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
    - name: Cache gpuocelot
      if: inputs.ocelot == 'true'
      id: cache-build
      uses: actions/cache@v4
      env:
-        cache-name: cache-gpuocelot-build
+        cache-name: cache-gpuocelot-build-1
      with:
        path: ${{ github.workspace }}/gpuocelot/ocelot
-        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-0
+        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.BUILD_CACHE_VERSION }}
    - name: Clone/compile gpuocelot
      if: inputs.ocelot == 'true' && steps.cache-build.outputs.cache-hit != 'true'
      shell: bash
@@ -221,14 +253,20 @@ runs:
        git checkout b16039dc940dc6bc4ea0a98380495769ff35ed99
        mkdir build
        cd build
-        cmake .. -Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5
+
+        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
        ninja
    - name: Install gpuocelot
      if: inputs.ocelot == 'true'
      shell: bash
      run: |
        cd ${{ github.workspace }}/gpuocelot/ocelot/build
-        sudo cp libgpuocelot.${{ runner.os == 'macOS' && 'dylib' || 'so' }} /usr/${{ runner.os == 'macOS' && 'local/' || ''}}lib/
+        sudo cp libgpuocelot.${{ runner.os == 'macOS' && 'dylib' || 'so' }} /usr/${{ runner.os == 'macOS' && 'local/' || '' }}lib/

    # **** WebGPU ****

@@ -0,0 +1,91 @@
+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
@@ -52,26 +52,26 @@ 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 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=500 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 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
+      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=700 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
    - name: Run Stable Diffusion v2
-      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
+      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=1600 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
    # process replay can't capture this, the graph is too large
    - 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
+      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3000 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: Run huggingface_onnx test
-      run: METAL=1 python3.11 extra/huggingface_onnx/run_models.py test --debug FacebookAI/xlm-roberta-large
    - name: Test speed vs torch
-      run: BIG=2 MPS=1 python3.11 test/test_speed_v_torch.py | tee torch_speed.txt
+      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/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestKernelOpts.test_tensor_core_opts
+      run: METAL=1 python3.11 test/opt/test_tensor_cores.py
    - name: Test AMX tensor cores
      run: |
-        DEBUG=2 CPU=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
-        DEBUG=2 LLVM=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        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
    - 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)
@@ -99,7 +99,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 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2 JIT=1 ASSERT_MIN_STEP_TIME=8 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
@@ -109,21 +109,21 @@ jobs:
    - name: Train MNIST
      run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py | tee beautiful_mnist.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
+      run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=320 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
+      run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=385 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
    - 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
+      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: 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 AMD_LLVM=1 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)
@@ -189,22 +189,22 @@ jobs:
    - name: Run model inference benchmark
      run: NV=1 CAPTURE_PROCESS_REPLAY=0 NOCLANG=1 python3 test/external/external_model_benchmark.py
    - name: Test speed vs torch
-      run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/test_speed_v_torch.py | tee torch_speed.txt
+      run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
    - name: Test speed vs theoretical
      run: NV=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
    - name: Test benchmark allreduce
      run: NV=1 python test/external/external_benchmark_multitensor_allreduce.py
    - name: Test tensor cores
      run: |
-        NV=1 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestKernelOpts.test_tensor_core_opts
-        PTX=1 ALLOW_TF32=1 NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestKernelOpts.test_tensor_core_opts
+        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
    - 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 PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
+      run: NV=1 NV_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
@@ -214,7 +214,7 @@ jobs:
    - name: Run Stable Diffusion
      run: BENCHMARK_LOG=stable_diffusion NV=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 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+      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 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
@@ -238,9 +238,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 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2 NV=1 JIT=1 ASSERT_MIN_STEP_TIME=5 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 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=5 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
@@ -302,30 +302,30 @@ jobs:
    - 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
+      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: 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 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=85 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 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=68 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 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=75 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
+      run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=35 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 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.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 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+      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
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (NVIDIA Training)
@@ -391,13 +391,13 @@ jobs:
    #- name: Test speed vs torch
    #  run: |
    #    python3 -c "import torch; print(torch.__version__)"
-    #    LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/test_speed_v_torch.py | tee torch_speed.txt
+    #    LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
    - name: Test speed vs theoretical
      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 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops TestKernelOpts.test_tensor_core_opts
-        AMD=1 AMD_LLVM=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops TestKernelOpts.test_tensor_core_opts
+        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 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
@@ -415,9 +415,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 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=450 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
+      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=1400 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
@@ -443,9 +443,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 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.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
    - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half AMD=1 HALF=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      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
    - 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
@@ -508,19 +508,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 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=85 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 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=188 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
+      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
    - 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
+      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
    - name: Run full CIFAR training w 1 GPU
      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
+    #- 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)
@@ -570,13 +570,13 @@ 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 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+      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
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (AMD MLPerf)
@@ -605,24 +605,28 @@ 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
-    - name: validate openpilot 0.9.7
-      run: PYTHONPATH=. FLOAT16=0 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_7.txt
-    - name: benchmark openpilot 0.9.7
-      run: BENCHMARK_LOG=openpilot_0_9_7 PYTHONPATH=. QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_0_9_7.txt
-    - name: benchmark openpilot w IMAGE=2 0.9.7
-      run: BENCHMARK_LOG=openpilot_0_9_7_image 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.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_7.txt
-    - name: openpilot compile3 0.9.7
-      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx
-    - name: openpilot compile3 0.9.7+ tomb raider
-      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/e8bea2c78ffa92685ece511e9b554122aaf1a79d/selfdrive/modeld/models/supercombo.onnx
-    - name: openpilot dmonitoring compile3 0.9.7
-      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: benchmark openpilot 0.9.9 driving_vision
+      run: BENCHMARK_LOG=openpilot_0_9_9_vision ASSERT_MIN_STEP_TIME=30 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/driving_vision.onnx
+    - name: benchmark openpilot 0.9.9 driving_policy
+      run: BENCHMARK_LOG=openpilot_0_9_9_policy ASSERT_MIN_STEP_TIME=45 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/driving_policy.onnx
+    - name: benchmark openpilot 0.9.9 dmonitoring
+      run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring ASSERT_MIN_STEP_TIME=70 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="." 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="." 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="." 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="." 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 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 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
        # 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
@@ -637,3 +641,136 @@ jobs:
          openpilot_0_9_7.txt
          openpilot_image_0_9_4.txt
          openpilot_image_0_9_7.txt
+
+  testreddriverbenchmark:
+    name: AM Benchmark
+    runs-on: [self-hosted, Linux, tinyboxrandom]
+    timeout-minutes: 15
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Remove amd modules
+      run: ./extra/hcq/hcq_smi.py amd rmmod
+    - name: Kill stale pids
+      run: ./extra/hcq/hcq_smi.py amd kill_pids
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
+        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
+        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
+        ln -s /raid/weights/mixtral-8x7b-32kseqlen weights/mixtral-8x7b-32kseqlen
+        ln -s /raid/weights/LLaMA-2 weights/LLaMA-2
+        mkdir -p extra/datasets
+        ln -s /raid/datasets/imagenet extra/datasets/imagenet
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: reset process replay
+      run: test/external/process_replay/reset.py
+    - name: Test driver cold start time
+      run: time DEBUG=3 AMD=1 AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
+    - name: Test driver warm start time
+      run: time DEBUG=3 AMD=1 python3 test/test_tiny.py TestTiny.test_plus
+    # 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 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
+    - name: Test AMD=1
+      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 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
+    - name: Run 10 MLPerf Bert training steps (1 gpu)
+      # TODO: remove BERT_LAYERS once scheduler is fast
+      run: BENCHMARK_LOG=bert_10steps AMD=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 am_train_bert_one_gpu.txt
+    - uses: actions/upload-artifact@v4
+      with:
+        name: Speed (AM Driver)
+        path: |
+          am_matmul_amd.txt
+          am_train_cifar_one_gpu.txt
+          am_train_resnet_one_gpu.txt
+          am_train_bert_one_gpu.txt
+    - name: Run process replay tests
+      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
+
+  testgreendriverbenchmark:
+    name: NV Benchmark
+    runs-on: [self-hosted, Linux, tinyboxrandom]
+    timeout-minutes: 15
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Remove nv modules
+      run: ./extra/hcq/hcq_smi.py nv rmmod
+    - name: Kill stale pids
+      run: ./extra/hcq/hcq_smi.py nv kill_pids
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
+        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
+        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
+        ln -s /raid/weights/mixtral-8x7b-32kseqlen weights/mixtral-8x7b-32kseqlen
+        ln -s /raid/weights/LLaMA-2 weights/LLaMA-2
+        mkdir -p extra/datasets
+        ln -s /raid/datasets/imagenet extra/datasets/imagenet
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: reset process replay
+      run: test/external/process_replay/reset.py
+    - 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
+    - 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 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
+    - uses: actions/upload-artifact@v4
+      with:
+        name: Speed (NV Driver)
+        path: |
+          nv_llama3_beam.txt
+          nv_train_cifar_one_gpu.txt
+          nv_train_resnet_one_gpu.txt
+          nv_train_bert_one_gpu.txt
+    - name: Run process replay tests
+      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
@@ -10,6 +10,7 @@ jobs:
  run_script_job:
    runs-on: [self-hosted, Linux, tinybox]
    if: github.repository_owner == 'tinygrad'
+    timeout-minutes: 100

    steps:
    - name: Checkout Code
@@ -27,7 +28,7 @@ jobs:
        BENCHMARK_LOG=search_sdxl_cached PYTHONPATH=. AMD=1 JITBEAM=2 python examples/sdxl.py --noshow --timing --seed 0
    - name: Run winograd cifar with new search
      run: |
-        BENCHMARK_LOG=search_wino_cifar WINO=1 DEFAULT_FLOAT=HALF FUSE_ARANGE=1 JITBEAM=4 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 BS=1024 STEPS=500 python examples/hlb_cifar10.py
+        BENCHMARK_LOG=search_wino_cifar WINO=1 DEFAULT_FLOAT=HALF JITBEAM=4 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 BS=1024 STEPS=500 python examples/hlb_cifar10.py
    - name: Run winograd cifar with cached search
      run: |
-        BENCHMARK_LOG=search_wino_cifar_cached WINO=1 DEFAULT_FLOAT=HALF FUSE_ARANGE=1 JITBEAM=4 BS=1024 STEPS=500 python examples/hlb_cifar10.py
+        BENCHMARK_LOG=search_wino_cifar_cached WINO=1 DEFAULT_FLOAT=HALF JITBEAM=4 BS=1024 STEPS=500 python examples/hlb_cifar10.py
@@ -12,6 +12,7 @@ jobs:
  run_script_job:
    runs-on: [self-hosted, Linux, tinybox]
    if: github.repository_owner == 'tinygrad'
+    timeout-minutes: 360

    steps:
    - name: Checkout Code
@@ -26,4 +27,4 @@ jobs:
      run: |
        rm "~/.cache/tinygrad/cache_mlperf.db" || true
        BENCHMARK_LOG=mlpert_train_resnet LOGMLPERF=0 CACHEDB="~/.cache/tinygrad/cache_mlperf.db" examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
-        rm "~/.cache/tinygrad/cache_mlperf.db"
+        rm "~/.cache/tinygrad/cache_mlperf.db"
@@ -20,12 +20,6 @@ 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,11 +54,12 @@ 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
+disable=C,R,W0613,W0511,W0212,W0201,W0106,W0603,W0621,W0703,W1201,W1203,E1136,W1514,E1101,W0221,W0105,E0401,abstract-method,W0707
 # 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
@@ -39,9 +39,8 @@ Try a matmul. See how, despite the style, it is fused into one kernel with the p

 ```sh
 DEBUG=3 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())"
+N = 1024; a, b = Tensor.empty(N, N), Tensor.empty(N, N);
+(a.reshape(N, 1, N) * b.T.reshape(1, N, N)).sum(axis=2).realize()"
 ```

 And we can change `DEBUG` to `4` to see the generated code.
@@ -80,9 +79,8 @@ See [examples/beautiful_mnist.py](examples/beautiful_mnist.py) for the full vers

 tinygrad already supports numerous accelerators, including:

- [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] [OpenCL](tinygrad/runtime/ops_cl.py)
+- [x] [CPU](tinygrad/runtime/ops_cpu.py)
 - [x] [METAL](tinygrad/runtime/ops_metal.py)
 - [x] [CUDA](tinygrad/runtime/ops_cuda.py)
 - [x] [AMD](tinygrad/runtime/ops_amd.py)
@@ -118,7 +118,9 @@ generate_nv() {
  clang2py -k cdefstum \
    extra/nv_gpu_driver/clc6c0qmd.h \
    extra/nv_gpu_driver/clcec0qmd.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl0000.h \
    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl0080.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl2080.h \
    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl2080_notification.h \
    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc56f.h \
    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc86f.h \
@@ -147,6 +149,7 @@ generate_nv() {
    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrlc36f.h \
    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrlcb33.h \
    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrla06c.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl90f1.h \
    --clang-args="-include $NVKERN_SRC/src/common/sdk/nvidia/inc/nvtypes.h -I$NVKERN_SRC/src/common/inc -I$NVKERN_SRC/kernel-open/nvidia-uvm -I$NVKERN_SRC/kernel-open/common/inc -I$NVKERN_SRC/src/common/sdk/nvidia/inc -I$NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include -I$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl" \
    -o $BASE/nv_gpu.py
  fixup $BASE/nv_gpu.py
@@ -164,8 +167,30 @@ generate_nv() {
  sed -n '1i\
 nv_status_codes = {}
 /^NV_STATUS_CODE/ { s/^NV_STATUS_CODE(\([^,]*\), *\([^,]*\), *"\([^"]*\)") *.*$/\1 = \2\nnv_status_codes[\1] = "\3"/; p }' $NVKERN_SRC/src/common/sdk/nvidia/inc/nvstatuscodes.h >> $BASE/nv_gpu.py
-
  python3 -c "import tinygrad.runtime.autogen.nv_gpu"
+
+  clang2py -k cdefstum \
+    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/fsp/kern_fsp_cot_payload.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gspifpub.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gsp_fw_wpr_meta.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gsp_fw_sr_meta.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/gsp/gsp_init_args.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/gsp/gsp_init_args.h \
+    $NVKERN_SRC/src/common/uproc/os/common/include/libos_init_args.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/rmRiscvUcode.h \
+    $NVKERN_SRC/src/common/shared/msgq/inc/msgq/msgq_priv.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/vgpu/rpc_headers.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/vgpu/rpc_global_enums.h \
+    $NVKERN_SRC/src/nvidia/generated/g_rpc-structures.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/fsp/fsp_nvdm_format.h \
+    extra/nv_gpu_driver/g_rpc-message-header.h \
+    extra/nv_gpu_driver/gsp_static_config.h \
+    extra/nv_gpu_driver/vbios.h \
+    --clang-args="-DRPC_MESSAGE_STRUCTURES -DRPC_STRUCTURES -include $NVKERN_SRC/src/common/sdk/nvidia/inc/nvtypes.h -I$NVKERN_SRC/src/nvidia/generated -I$NVKERN_SRC/src/common/inc -I$NVKERN_SRC/src/nvidia/inc -I$NVKERN_SRC/src/nvidia/interface/ -I$NVKERN_SRC/src/nvidia/inc/kernel -I$NVKERN_SRC/src/nvidia/inc/libraries -I$NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc -I$NVKERN_SRC/kernel-open/nvidia-uvm -I$NVKERN_SRC/kernel-open/common/inc -I$NVKERN_SRC/src/common/sdk/nvidia/inc -I$NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include -I$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl" \
+    -o $BASE/nv/nv.py
+
+  fixup $BASE/nv/nv.py
+  python3 -c "import tinygrad.runtime.autogen.nv.nv"
 }

 generate_amd() {
@@ -173,11 +198,7 @@ generate_amd() {
  clang2py -k cdefstum \
    extra/hip_gpu_driver/sdma_registers.h \
    extra/hip_gpu_driver/nvd.h \
-    extra/hip_gpu_driver/kfd_pm4_headers_ai.h \
-    extra/hip_gpu_driver/soc21_enum.h \
-    extra/hip_gpu_driver/sdma_v6_0_0_pkt_open.h \
    extra/hip_gpu_driver/gc_11_0_0_offset.h \
-    extra/hip_gpu_driver/gc_10_3_0_offset.h \
    extra/hip_gpu_driver/sienna_cichlid_ip_offset.h \
    --clang-args="-I/opt/rocm/include -x c++" \
    -o $BASE/amd_gpu.py
@@ -215,6 +236,21 @@ generate_io_uring() {
  fixup $BASE/io_uring.py
 }

+generate_ib() {
+  clang2py -k cdefstum \
+    /usr/include/infiniband/verbs.h \
+    /usr/include/infiniband/verbs_api.h \
+    /usr/include/infiniband/ib_user_ioctl_verbs.h \
+    /usr/include/rdma/ib_user_verbs.h \
+    -o $BASE/ib.py
+
+  sed -i "s\import ctypes\import ctypes, ctypes.util\g" "$BASE/ib.py"
+  sed -i "s\FIXME_STUB\libibverbs\g" "$BASE/ib.py"
+  sed -i "s\FunctionFactoryStub()\ctypes.CDLL(ctypes.util.find_library('ibverbs'), use_errno=True)\g" "$BASE/ib.py"
+
+  fixup $BASE/ib.py
+}
+
 generate_libc() {
  clang2py -k cdefstum \
    $(dpkg -L libc6-dev | grep sys/mman.h) \
@@ -336,26 +372,6 @@ generate_am() {
    -o $BASE/am/pm4_nv.py
  fixup $BASE/am/pm4_nv.py

-  clang2py -k cdefstum \
-    $AMKERN_INC/vega10_enum.h \
-    -o $BASE/am/vega10.py
-  fixup $BASE/am/vega10.py
-
-  clang2py -k cdefstum \
-    $AMKERN_INC/navi10_enum.h \
-    -o $BASE/am/navi10.py
-  fixup $BASE/am/navi10.py
-
-  clang2py -k cdefstum \
-    $AMKERN_INC/soc21_enum.h \
-    -o $BASE/am/soc21.py
-  fixup $BASE/am/soc21.py
-
-  clang2py -k cdefstum \
-    $AMKERN_INC/soc24_enum.h \
-    -o $BASE/am/soc24.py
-  fixup $BASE/am/soc24.py
-
  clang2py -k cdefstum \
    extra/hip_gpu_driver/sdma_registers.h \
    $AMKERN_AMD/amdgpu/vega10_sdma_pkt_open.h \
@@ -419,7 +435,7 @@ generate_libusb() {
    -o $BASE/libusb.py

  fixup $BASE/libusb.py
-  sed -i "s\import ctypes\import ctypes, os\g" $BASE/libusb.py
+  sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/libusb.py
  sed -i "s/FIXME_STUB/libusb/g" "$BASE/libusb.py"
  sed -i "s/libusb_le16_to_cpu = libusb_cpu_to_le16//g" "$BASE/libusb.py"
  sed -i "s/FunctionFactoryStub()/None if (lib_path:=os.getenv('LIBUSB_PATH', ctypes.util.find_library('usb-1.0'))) is None else ctypes.CDLL(lib_path)/g" "$BASE/libusb.py"
@@ -436,9 +452,11 @@ elif [ "$1" == "kfd" ]; then generate_kfd
 elif [ "$1" == "nv" ]; then generate_nv
 elif [ "$1" == "amd" ]; then generate_amd
 elif [ "$1" == "am" ]; then generate_am
+elif [ "$1" == "nvdrv" ]; then generate_nvdrv
 elif [ "$1" == "sqtt" ]; then generate_sqtt
 elif [ "$1" == "qcom" ]; then generate_qcom
 elif [ "$1" == "io_uring" ]; then generate_io_uring
+elif [ "$1" == "ib" ]; then generate_ib
 elif [ "$1" == "libc" ]; then generate_libc
 elif [ "$1" == "llvm" ]; then generate_llvm
 elif [ "$1" == "kgsl" ]; then generate_kgsl
@@ -7,28 +7,30 @@

 print("******** first, the runtime ***********")

-from tinygrad.runtime.ops_cpu import ClangJITCompiler, MallocAllocator, CPUProgram
+from tinygrad.runtime.ops_cpu import ClangJITCompiler, CPUDevice, CPUProgram
+
+cpu = CPUDevice()

 # allocate some buffers
-out = MallocAllocator.alloc(4)
-a = MallocAllocator.alloc(4)
-b = MallocAllocator.alloc(4)
+out = cpu.allocator.alloc(4)
+a = cpu.allocator.alloc(4)
+b = cpu.allocator.alloc(4)

 # load in some values (little endian)
-MallocAllocator._copyin(a, memoryview(bytearray([2,0,0,0])))
-MallocAllocator._copyin(b, memoryview(bytearray([3,0,0,0])))
+cpu.allocator._copyin(a, memoryview(bytearray([2,0,0,0])))
+cpu.allocator._copyin(b, memoryview(bytearray([3,0,0,0])))

 # compile a program to a binary
 lib = ClangJITCompiler().compile("void add(int *out, int *a, int *b) { out[0] = a[0] + b[0]; }")

 # create a runtime for the program
-fxn = CPUProgram("add", lib)
+fxn = cpu.runtime("add", lib)

 # run the program
 fxn(out, a, b)

 # check the data out
-print(val := MallocAllocator._as_buffer(out).cast("I").tolist()[0])
+print(val := cpu.allocator._as_buffer(out).cast("I").tolist()[0])
 assert val == 5


@@ -46,7 +48,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
 out = Buffer(DEVICE, 1, dtypes.int32).allocate()
 a = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struct.pack("I", 2))))
 b = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struct.pack("I", 3))))
-# NOTE: a._buf is the same as the return from MallocAllocator.alloc
+# NOTE: a._buf is the same as the return from cpu.allocator.alloc

 # describe the computation
 buf_1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 1)
@@ -60,7 +62,7 @@ s = UOp(Ops.SINK, dtypes.void, (st_0,))

 # convert the computation to a "linearized" format (print the format)
 from tinygrad.engine.realize import get_program, CompiledRunner
-program = get_program(Device[DEVICE].renderer, s)
+program = get_program(s, Device[DEVICE].renderer)

 # compile a program (and print the source)
 fxn = CompiledRunner(program)
@@ -78,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.engine.kernelize import get_kernelize_map
+from tinygrad.schedule.kernelize import get_kernelize_map

 # allocate some values + load in values
 a = UOp.new_buffer(DEVICE, 1, dtypes.int32)
@@ -52,7 +52,7 @@ Signals are device-dependent structures used for synchronization and timing in H
 The following Python code demonstrates the usage of signals:

 ```python
-signal = your_device.signal_t()
+signal = your_device.new_signal(value=0)

 HWQueue().timestamp(signal) \
         .signal(signal, value_to_fire) \
@@ -0,0 +1,60 @@
+# tinygrad directory layout
+
+This explains the flow of a big graph down to programs.
+
+Directories are listed in order of how they are processed.
+
+---
+
+## tinygrad/schedule
+
+Group UOps into kernels.
+
+::: tinygrad.schedule.kernelize.get_kernelize_map
+    options:
+        members: false
+        show_labels: false
+        show_source: false
+
+---
+
+## tinygrad/codegen/opt
+
+Transforms the ast into an optimized ast. This is where BEAM search and heuristics live.
+
+---
+
+## tinygrad/codegen
+
+Transform the optimized ast into a linearized list of UOps.
+
+::: tinygrad.codegen.full_rewrite
+    options:
+        members: false
+        show_labels: false
+        show_source: false
+
+---
+
+## tinygrad/renderer
+
+Transform the linearized list of UOps into a program, represented as a string.
+
+::: tinygrad.renderer.Renderer
+    options:
+        members:
+            - render
+        show_labels: false
+        show_source: false
+
+---
+
+## tinygrad/engine
+
+Abstracted high level interface to the runtimes.
+
+::: tinygrad.engine.realize.get_program
+    options:
+        members: false
+        show_labels: false
+        show_source: false
@@ -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: `GPU=1 DEBUG=4 python3 -m pytest`
+Example: `CL=1 DEBUG=4 python3 -m pytest`

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

@@ -31,13 +31,12 @@ 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)
-GPU                 | [1]        | enable the GPU (OpenCL) backend
+CL                  | [1]        | enable 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 (Clang) backend
-LLVM                | [1]        | enable LLVM backend
+CPU                 | [1]        | enable CPU 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
@@ -126,7 +126,7 @@ print(t_log_grad.uop)
 """
 void E_(float* restrict data0, float* restrict data1) {
  float val0 = *(data1+0);
-  *(data0+0) = (0.6931471805599453f*(1/(val0*0.6931471805599453f)));
+  *(data0+0) = (1/val0);
 }
 """
 # the derivative is close to 1/3
@@ -9,10 +9,10 @@ tinygrad supports various runtimes, enabling your code to scale across a wide ra
 | [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 |
+| [OpenCL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cl.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.1.6). |
+| [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

@@ -35,6 +35,7 @@ Elementwise ops operate on a per element basis. They don't change the shape of t

 ::: tinygrad.Tensor.relu
 ::: tinygrad.Tensor.sigmoid
+::: tinygrad.Tensor.logsigmoid
 ::: tinygrad.Tensor.hardsigmoid
 ::: tinygrad.Tensor.elu
 ::: tinygrad.Tensor.celu
@@ -77,6 +78,7 @@ 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

@@ -26,5 +26,6 @@
 ::: tinygrad.Tensor.transpose
 ::: tinygrad.Tensor.flatten
 ::: tinygrad.Tensor.unflatten
+::: tinygrad.Tensor.diag
 ::: tinygrad.Tensor.roll
 ::: tinygrad.Tensor.rearrange
@@ -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, and the green box includes six 4090 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, 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.

 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.

@@ -47,8 +47,8 @@ Reboot after making these changes or restart the `displayservice.service` servic

 The [default tinybox image](https://github.com/tinygrad/tinyos) ships with tinygrad and PyTorch. While we develop tinygrad, the box is universal hardware. Use whatever framework you desire, run notebooks, download demos, install more things, train, inference, live, laugh, love, you aren't paying per hour for this box so the only limit is your imagination.

-## tinychat
+## Building the OS image

-Since LLMs are so popular, we ship with a built in tinygrad based chatbot using a LLaMA-3 finetune. Visit the IP (not the BMC IP) of your tinybox in a web browser on your computer or phone, and you'll find a friendly looking chat interface. This chatbot also provides an OpenAI compatible LLM API on that port, so you can script it.
+The OS image is built using `ubuntu-image` from <https://github.com/tinygrad/tinyos>.

-The conversations you have with this chatbot are between you and your tinybox. Also, the history in the web app is saved on the client, not the tinybox.
+After cloning, run `make green` or `make red` to build a tinybox green or tinybox red image respectively.
@@ -2,7 +2,6 @@ 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
@@ -150,13 +149,12 @@ 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()

-  np.random.seed(1337)
+  Tensor.manual_seed(1337)
+  num_train_samples = X_train.shape[0]
+
  for epoch in range(math.ceil(hyp['misc']['train_epochs'])):
-    # TODO: move to tinygrad
    gst = time.perf_counter()
-    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
+    tidxs = Tensor.randperm(num_train_samples, dtype='int')[:num_steps_per_epoch*batchsize].reshape(num_steps_per_epoch, batchsize)
    train_loss:float = 0
    for epoch_step in (t:=trange(num_steps_per_epoch)):
      st = time.perf_counter()
@@ -1,12 +1,12 @@
 # model based off https://medium.com/data-science/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
-from typing import List, Callable
+from typing import Callable
 from tinygrad import Tensor, TinyJit, nn, GlobalCounters
 from tinygrad.helpers import getenv, colored, trange
 from tinygrad.nn.datasets import mnist

 class Model:
  def __init__(self):
-    self.layers: List[Callable[[Tensor], Tensor]] = [
+    self.layers: list[Callable[[Tensor], Tensor]] = [
      nn.Conv2d(1, 32, 5), Tensor.relu,
      nn.Conv2d(32, 32, 5), Tensor.relu,
      nn.BatchNorm(32), Tensor.max_pool2d,
@@ -21,17 +21,15 @@ if __name__ == "__main__":
  X_train, Y_train, X_test, Y_test = mnist(fashion=getenv("FASHION"))

  model = Model()
-  opt = nn.optim.Adam(nn.state.get_parameters(model))
+  opt = (nn.optim.Adam if not getenv("MUON") else nn.optim.Muon)(nn.state.get_parameters(model))

  @TinyJit
  @Tensor.train()
  def train_step() -> Tensor:
    opt.zero_grad()
    samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
-    # TODO: this "gather" of samples is very slow. will be under 5s when this is fixed
    loss = model(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
-    opt.step()
-    return loss
+    return loss.realize(*opt.schedule_step())

  @TinyJit
  def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100
@@ -1,11 +1,10 @@
-import sys, time, pickle
+import sys, time
 from tinygrad import TinyJit, GlobalCounters, fetch, getenv
-from tinygrad.frontend.onnx import OnnxRunner, onnx_load
+from tinygrad.frontend.onnx import OnnxRunner
 from extra.onnx_helpers import get_example_inputs, validate

 def load_onnx_model(onnx_file):
-  onnx_model = onnx_load(onnx_file)
-  run_onnx = OnnxRunner(onnx_model)
+  run_onnx = OnnxRunner(onnx_file)
  run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(None) for k,v in kwargs.items()}).values())), prune=True, optimize=True)
  return run_onnx_jit, run_onnx.graph_inputs

@@ -4,7 +4,7 @@ sys.path.append(os.getcwd())

 from io import StringIO
 from contextlib import redirect_stdout
-from tinygrad import Tensor, nn, Device, dtypes
+from tinygrad import Tensor, nn
 from tinygrad.helpers import Timing, colored, getenv, fetch
 from extra.models.llama import Transformer, convert_from_huggingface, fix_bf16
 from sentencepiece import SentencePieceProcessor
@@ -10,6 +10,7 @@ import tensorflow as tf
 import tf2onnx
 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

 def get_uncompiled_model2(dataset_size=32, output_size=4):
@@ -25,7 +26,7 @@ class TinyOnnx:
  def __init__(self, keras_model):
    input_signature = [tf.TensorSpec([1,32], tf.float32, name='x')]
    onnx_model, _ = tf2onnx.convert.from_keras(keras_model, input_signature, opset=13)
-    self.run_onnx = OnnxRunner(onnx_model)
+    self.run_onnx = OnnxRunner(Tensor(onnx_model.SerializeToString(), device="PYTHON"))

  def forward(self, x):
    return self.run_onnx({"x": x}, debug=False)['predictions']
@@ -47,8 +48,8 @@ def compile_onnx_model(onnx_model):
  cprog.append("void initialize(float *weights) {")
  weights = bytes()
  for name,cl in bufs_to_save.items():
-    cprog.append(f"memcpy({name}, weights + {len(weights)//4}, {len(cl._buf)*4});")
-    weights += bytes(cl._buf)
+    cprog.append(f"memcpy({name}, weights + {len(weights)//4}, {cl._buf.size});")
+    weights += bytes(to_mv(cl._buf.va_addr, cl._buf.size))
  cprog.append("}")

  # write the weights to disk
@@ -181,6 +181,7 @@ 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
@@ -188,7 +189,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 GPU" if DEBUG>=2 else "")+
+      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 "")+
                  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):
@@ -197,8 +198,13 @@ 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 ****
@@ -1,133 +0,0 @@
-from extra.models.resnet import ResNet50
-from extra.mcts_search import mcts_search
-from examples.mlperf.helpers import get_mlperf_bert_model
-from tinygrad import Tensor, Device, dtypes, nn
-from tinygrad.opt.kernel import Kernel
-from tinygrad.opt.heuristic import hand_coded_optimizations
-from tinygrad.uop.ops import Ops, sym_infer
-from tinygrad.device import Compiled
-from tinygrad.opt.search import beam_search, bufs_from_lin
-from tinygrad.helpers import DEBUG, ansilen, getenv, colored, TRACEMETA
-from extra.optimization.helpers import time_linearizer
-
-def get_sched_resnet():
-  mdl = ResNet50()
-  optim = (nn.optim.LARS if getenv("LARS") else nn.optim.SGD)(nn.state.get_parameters(mdl))
-  BS = getenv("BS", 64)
-
-  # run model twice to get only what changes, these are the kernels of the model
-  for _ in range(2):
-    out = mdl(Tensor.empty(BS, 3, 224, 224))
-    targets = [out]
-    if getenv("BACKWARD"):
-      optim.zero_grad()
-      out.sparse_categorical_crossentropy(Tensor.empty(BS, dtype=dtypes.int)).backward()
-      targets += [x for x in optim.schedule_step()]
-    sched = Tensor.schedule(*targets)
-    print(f"schedule length {len(sched)}")
-  return sched
-
-def get_sched_bert():
-  mdl = get_mlperf_bert_model()
-  optim = nn.optim.LAMB(nn.state.get_parameters(mdl))
-
-  # fake data
-  BS = getenv("BS", 9)
-  input_ids = Tensor.empty((BS, 512), dtype=dtypes.float32)
-  segment_ids = Tensor.empty((BS, 512), dtype=dtypes.float32)
-  attention_mask = Tensor.empty((BS, 512), dtype=dtypes.default_float)
-  masked_positions = Tensor.empty((BS, 76), dtype=dtypes.float32)
-  masked_lm_ids = Tensor.empty((BS, 76), dtype=dtypes.float32)
-  masked_lm_weights = Tensor.empty((BS, 76), dtype=dtypes.float32)
-  next_sentence_labels = Tensor.empty((BS, 1), dtype=dtypes.float32)
-
-  # run model twice to get only what changes, these are the kernels of the model
-  for _ in range(2):
-    lm_logits, seq_relationship_logits = mdl(input_ids, attention_mask, masked_positions, segment_ids)
-    targets = [lm_logits, seq_relationship_logits]
-    if getenv("BACKWARD"):
-      optim.zero_grad()
-      loss = mdl.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
-      # ignore grad norm and loss scaler for now
-      loss.backward()
-      targets += [x for x in optim.schedule_step()]
-    sched = Tensor.schedule(*targets)
-    print(f"schedule length {len(sched)}")
-  return sched
-
-if __name__ == "__main__":
-  if getenv("HALF", 1):
-    dtypes.default_float = dtypes.half
-
-  # the device we are optimizing for
-  device: Compiled = Device[Device.DEFAULT]
-  if getenv("BACKWARD"): Tensor.training = True
-  print(f"optimizing for {Device.DEFAULT}")
-
-  sched = globals()[f"get_sched_{getenv('MODEL', 'resnet')}"]()
-  sched = [x for x in sched if x.ast.op is Ops.SINK]
-
-  # focus on one kernel
-  if getenv("KERNEL", -1) >= 0: sched = sched[getenv("KERNEL", -1):getenv("KERNEL", -1)+1]
-
-  # work with the schedule
-  total_tm = 0
-  running_gflops = 0
-  usage = {}
-  for i,si in enumerate(sched):
-    if DEBUG >= 3: print(si.ast)
-
-    rawbufs = bufs_from_lin(Kernel(si.ast))
-
-    # "linearize" the op into uops in different ways
-    lins: list[tuple[Kernel, str]] = []
-
-    # always try hand coded opt
-    lin = Kernel(si.ast, opts=device.renderer)
-    lin.apply_opts(hand_coded_optimizations(lin))
-    lins.append((lin, "HC"))
-
-    # maybe try tensor cores
-    lin = Kernel(si.ast, opts=device.renderer)
-    if lin.apply_tensor_cores():
-      lins.append((lin, "TC"))
-
-    # try a beam search
-    if beam:=getenv("BEAM"):
-      lin = Kernel(si.ast, opts=device.renderer)
-      lin = beam_search(lin, rawbufs, beam, bool(getenv("BEAM_ESTIMATE", 1)))
-      lins.append((lin, "BEAM"))
-
-    # try MCTS
-    if mcts:=getenv("MCTS"):
-      lin = Kernel(si.ast, opts=device.renderer)
-      lin = mcts_search(lin, rawbufs, mcts)
-      lins.append((lin, "MCTS"))
-
-    # benchmark the programs
-    choices = []
-    for lin, nm in lins:
-      tm = time_linearizer(lin, rawbufs, allow_test_size=False, cnt=10, disable_cache=True)
-      ops = (prg:=lin.to_program()).estimates.ops
-      gflops = sym_infer(ops, {k:k.min for k in lin.ast.variables()})*1e-9/tm
-      choices.append((tm, gflops, lin, prg, nm))
-
-    sorted_choices = sorted(choices, key=lambda x: x[0])
-    if DEBUG >= 1: # print all kernels
-      for tm, gflops, lin, prg, nm in choices:
-        print(f"                 kernel {i:2d} {lin.name+' '*(37-ansilen(lin.name))} {str(prg.global_size):18s} {str(prg.local_size):12s} takes {tm*1000:7.2f} ms, {gflops:6.0f} GFLOPS -- {colored(nm, 'green') if lin is sorted_choices[0][2] else nm}")
-
-    tm, gflops, lin, prg, nm = sorted_choices[0]
-    if getenv("SRC"):
-      print(si.ast)
-      print(lin.applied_opts)
-      print(lin.to_program().src)
-    total_tm += tm
-    running_gflops += gflops * tm
-    if (key := str([str(m) for m in si.metadata])) not in usage: usage[key] = (0, 0)
-    usage[key] = (usage[key][0] + tm, usage[key][1] + 1)
-    print(f"*** {total_tm*1000:7.2f} ms : kernel {i:2d} {lin.name+' '*(37-ansilen(lin.name))} {str(prg.global_size):18s} {str(prg.local_size):12s} takes {tm*1000:7.2f} ms, {gflops:6.0f} GFLOPS {[repr(m) if TRACEMETA >= 2 else str(m) for m in si.metadata]}")
-  print(f"******* total {total_tm*1000:.2f} ms, {running_gflops/total_tm:6.0f} GFLOPS")
-  print("usage:")
-  for k in sorted(usage, key=lambda x: -usage[x][0])[:10]:
-    print(f"{usage[k][0]*1000:.2f} ms: {k} ({usage[k][1]} times)")
@@ -7,8 +7,8 @@ import random, time
 import numpy as np
 from typing import Optional
 from extra.lr_scheduler import OneCycleLR
-from tinygrad import nn, dtypes, Tensor, Device, GlobalCounters, TinyJit
-from tinygrad.nn.state import get_state_dict, get_parameters
+from tinygrad import nn, dtypes, Tensor, Device, GlobalCounters, TinyJit, Variable
+from tinygrad.nn.state import get_state_dict
 from tinygrad.nn import optim
 from tinygrad.helpers import Context, BEAM, WINO, getenv, colored, prod
 from extra.bench_log import BenchEvent, WallTimeEvent
@@ -118,7 +118,7 @@ class SpeedyResNet:
 # hyper-parameters were exactly the same as the original repo
 bias_scaler = 58
 hyp = {
-  'seed' : 200,
+  'seed' : 201,
  '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):
@@ -201,24 +202,37 @@ def train_cifar():
    idx_y = Tensor.arange(H, dtype=dtypes.int32).reshape((1,1,H,1))
    return (idx_x >= low_x) * (idx_x < (low_x + mask_size)) * (idx_y >= low_y) * (idx_y < (low_y + mask_size))

-  def random_crop(X:Tensor, crop_size=32):
-    mask = make_square_mask(X.shape, crop_size)
-    mask = mask.expand((-1,3,-1,-1))
-    X_cropped = Tensor(X.numpy()[mask.numpy()])
-    return X_cropped.reshape((-1, 3, crop_size, crop_size))
+  # Similar, but different enough.
+  def make_random_crop_indices(shape, mask_size) -> Tensor:
+    BS, _, H, W = shape
+    low_x = Tensor.randint(BS, low=0, high=W-mask_size).reshape(BS,1,1,1)
+    low_y = Tensor.randint(BS, low=0, high=H-mask_size).reshape(BS,1,1,1)
+    idx_x = Tensor.arange(mask_size, dtype=dtypes.int32).reshape((1,1,1,mask_size))
+    idx_y = Tensor.arange(mask_size, dtype=dtypes.int32).reshape((1,1,mask_size,1))
+    return low_x, low_y, idx_x, idx_y

-  def cutmix(X:Tensor, Y:Tensor, mask_size=3):
-    # fill the square with randomly selected images from the same batch
+  def random_crop(X:Tensor, crop_size=32):
+    Xs, Ys, Xi, Yi = make_random_crop_indices(X.shape, crop_size)
+    return X.gather(-1, (Xs + Xi).expand(-1, 3, X.shape[2], -1)).gather(-2, ((Ys+Yi).expand(-1, 3, crop_size, crop_size)))
+
+  def cutmix(X, Y, order, mask_size=3):
    mask = make_square_mask(X.shape, mask_size)
-    order = list(range(0, X.shape[0]))
-    random.shuffle(order)
-    X_patch = Tensor(X.numpy()[order], device=X.device, dtype=X.dtype)
-    Y_patch = Tensor(Y.numpy()[order], device=Y.device, dtype=Y.dtype)
+    X_patch, Y_patch = X[order], Y[order]
    X_cutmix = mask.where(X_patch, X)
    mix_portion = float(mask_size**2)/(X.shape[-2]*X.shape[-1])
    Y_cutmix = mix_portion * Y_patch + (1. - mix_portion) * Y
    return X_cutmix, Y_cutmix

+  @TinyJit
+  def augmentations(X:Tensor, Y:Tensor):
+    perms = Tensor.randperm(X.shape[0], device=X.device) # We reuse perms for cutmix, because they are expensivne to generate
+    if getenv("RANDOM_CROP", 1):
+      X = random_crop(X, crop_size=32)
+    if getenv("RANDOM_FLIP", 1):
+      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'])
+
  # the operations that remain inside batch fetcher is the ones that involves random operations
  def fetch_batches(X_in:Tensor, Y_in:Tensor, BS:int, is_train:bool):
    step, epoch = 0, 0
@@ -226,28 +240,16 @@ def train_cifar():
      st = time.monotonic()
      X, Y = X_in, Y_in
      if is_train:
-        # TODO: these are not jitted
-        if getenv("RANDOM_CROP", 1):
-          X = random_crop(X, crop_size=32)
-        if getenv("RANDOM_FLIP", 1):
-          X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X) # flip LR
-        if getenv("CUTMIX", 1):
-          if step >= hyp['net']['cutmix_steps']:
-            X, Y = cutmix(X, Y, mask_size=hyp['net']['cutmix_size'])
-        order = list(range(0, X.shape[0]))
-        random.shuffle(order)
-        X, Y = X.numpy()[order], Y.numpy()[order]
-      else:
-        X, Y = X.numpy(), Y.numpy()
+        X, Y, X_cm, Y_cm = augmentations(X, Y)
+        if getenv("CUTMIX", 1) and step >= hyp['net']['cutmix_steps']: X, Y = X_cm, Y_cm
      et = time.monotonic()
      print(f"shuffling {'training' if is_train else 'test'} dataset in {(et-st)*1e3:.2f} ms ({epoch=})")
-      for i in range(0, X.shape[0], BS):
-        # pad the last batch  # TODO: not correct for test
-        batch_end = min(i+BS, Y.shape[0])
-        x = Tensor(X[batch_end-BS:batch_end], device=X_in.device, dtype=X_in.dtype)
-        y = Tensor(Y[batch_end-BS:batch_end], device=Y_in.device, dtype=Y_in.dtype)
+
+      vi = Variable("i", 0, (full_batches := (X.shape[0] // BS) * BS) - BS)
+      for i in range(0, full_batches, BS):
        step += 1
-        yield x, y
+        vib = vi.bind(i)
+        yield X[vib:vib+BS], Y[vib:vib+BS]
      epoch += 1
      if not is_train: break

@@ -353,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
@@ -411,12 +413,17 @@ 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 GPU" if DEBUG>=2 else "")+
+            with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" 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)
@@ -248,7 +248,7 @@ if __name__ == "__main__":
  parser.add_argument("--port", type=int, default=7776, help="Web server port")
  parser.add_argument("--debug", action="store_true", help="Enable debug mode")
  parser.add_argument("--seed", type=int, help="Random seed")
-  parser.add_argument("--temperature", type=int, default=0.85, help="Temperature")
+  parser.add_argument("--temperature", type=float, default=0.85, help="Temperature")
  parser.add_argument("--benchmark", action="store_true", help="Run a benchmark")
  parser.add_argument("--timing", action="store_true", help="Print timing per token")
  parser.add_argument("--profile", action="store_true", help="Output profile data")
@@ -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 GPU" if DEBUG>=2 else "")+
+            with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" 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 GPU" if DEBUG>=2 else "")+
+            with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" 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):

@@ -1,11 +1,11 @@
 #!/usr/bin/env python3
 import os
 if "NOOPT" not in os.environ: os.environ["NOOPT"] = "1"
-from tinygrad import Device, nn, Tensor, dtypes, Variable
+from tinygrad import Device, nn, Tensor, dtypes
 Device.DEFAULT = "CPU"
 from train_gpt2 import GPT, GPTConfig
-from tinygrad.helpers import dedup, to_function_name, flatten, getenv, GlobalCounters, ansilen, to_function_name
-from tinygrad.engine.realize import get_kernel, run_schedule
+from tinygrad.helpers import dedup, flatten, getenv, GlobalCounters, to_function_name
+from tinygrad.engine.realize import get_kernel
 from tinygrad.engine.memory import memory_planner
 from tinygrad.uop.ops import Ops

@@ -279,9 +279,15 @@ 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:
-      tok_Tens = (logits/temp).softmax().multinomial()
+      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()
    else:
      tok_Tens = logits.argmax(axis=-1).unsqueeze(0)
    tok = tok_Tens.item()
@@ -298,6 +304,7 @@ 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()
@@ -308,8 +315,9 @@ 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)
+  tinyoutput = generate(model, tokenizer, prompt, n_tokens_to_gen=num_toks, sample=sample, temp=temp, top_k=top_k)
  print(tinyoutput)
  print('TIME: ', time.time() - s)
  TORCHOUTPUT = "Why is gravity \nso important?\nBecause it's the only"
@@ -1,4 +1,4 @@
-import os, random, pickle, queue
+import os, random, pickle, queue, struct, math, functools, hashlib, time
 from typing import List
 from pathlib import Path
 from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu_count
@@ -6,6 +6,7 @@ from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu
 import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX
+from tinygrad.nn.state import TensorIO

 ### ResNet

@@ -510,6 +511,274 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
      # happens with BENCHMARK set
      pass

+# llama3
+
+class BinIdxDataset:
+  def __init__(self, base_path:Path):
+    self.idx_t = Tensor(base_path.with_name(f"{base_path.name}.idx"))
+    self.idx = TensorIO(self.idx_t)
+
+    # parse idx file
+    magic = self.idx.read(9)
+    assert magic == b"MMIDIDX\x00\x00", "invalid index file format"
+    version, = struct.unpack("<Q", self.idx.read(8))
+    assert version == 1, "unsupported index version"
+    dtype_code, = struct.unpack("<B", self.idx.read(1))
+    self.dtype = {1:dtypes.uint8, 2:dtypes.int8, 3:dtypes.int16, 4:dtypes.int32, 5:dtypes.int64, 6:dtypes.float64, 7:dtypes.double, 8:dtypes.uint16}[dtype_code]
+    self.count, = struct.unpack("<Q", self.idx.read(8))
+    doc_count, = struct.unpack("<Q", self.idx.read(8))
+
+    start = self.idx.tell()
+    end = start + self.count * dtypes.int32.itemsize
+    self.sizes = self.idx_t[start:end].bitcast(dtypes.int32).numpy()
+
+    start = end
+    end = start + self.count * dtypes.int64.itemsize
+    self.pointers = self.idx_t[start:end].bitcast(dtypes.int64).numpy()
+
+    start = end
+    end = start + doc_count * dtypes.int64.itemsize
+    self.doc_idx = self.idx_t[start:end].bitcast(dtypes.int64).numpy()
+
+    # bin file
+    self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin"))
+
+  def _index(self, idx) -> tuple[int, int]:
+    return int(self.pointers[idx]), int(self.sizes[idx])
+
+  def get(self, idx, offset:int=0, length:int|None=None):
+    ptr, size = self._index(idx)
+    if length is None: length = size - offset
+    ptr += offset * self.dtype.itemsize
+    return self.bin_t[ptr:ptr+length*self.dtype.itemsize].bitcast(self.dtype).to(None)
+
+# https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/datasets.html
+class GPTDataset:
+  def __init__(self, base_path:Path, samples:int, seqlen:int, seed:int, shuffle:bool):
+    self.samples, self.seqlen = samples, seqlen
+    self.shuffle = shuffle
+    self.rng = np.random.RandomState(seed)
+
+    self.indexed_dataset = BinIdxDataset(base_path)
+
+    # check for cache
+    cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
+    cache_path = base_path.with_name(f"{base_path.name}.{cache_hash}.index_cache")
+    print(f"try loading GPTDataset from {cache_path}...")
+    if cache_path.exists():
+      print("cache found, loading...")
+      with open(cache_path, "rb") as f:
+        self.doc_idx, self.sample_idx, self.shuffle_idx = pickle.load(f)
+    else:
+      print("cache not found, building index...")
+      self.doc_idx = self._build_doc_idx()
+      self.sample_idx = self._build_sample_idx()
+      self.shuffle_idx = self._build_shuffle_idx()
+      # save cache
+      with open(cache_path, "wb") as f:
+        pickle.dump((self.doc_idx, self.sample_idx, self.shuffle_idx), f)
+
+  def __getitem__(self, idx):
+    if idx is None:
+      text = self._get(0)
+    else:
+      text = self._get(idx)
+
+    return text
+
+  def _get(self, idx):
+    idx = self.shuffle_idx[idx]
+
+    doc_idx_beg, doc_idx_beg_offset = self.sample_idx[idx]
+    doc_idx_end, doc_idx_end_offset = self.sample_idx[idx + 1]
+
+    doc_ids, sample_parts = [], []
+
+    if doc_idx_beg == doc_idx_end:
+      doc_ids.append(self.doc_idx[doc_idx_beg])
+
+      sample_parts.append(
+          self.indexed_dataset.get(
+            int(self.doc_idx[doc_idx_beg]), offset=int(doc_idx_beg_offset), length=int(doc_idx_end_offset - doc_idx_beg_offset + 1)))
+    else:
+      for i in range(doc_idx_beg, doc_idx_end + 1):
+        doc_ids.append(self.doc_idx[i])
+
+        offset = 0 if i > doc_idx_beg else doc_idx_beg_offset
+        length = None if i < doc_idx_end else int(doc_idx_end_offset + 1)
+        sample_parts.append(self.indexed_dataset.get(int(self.doc_idx[i]), offset=int(offset), length=length))
+
+    # concat all parts
+    text = Tensor.cat(*sample_parts)
+
+    return text
+
+  @functools.cached_property
+  def tokens_per_epoch(self) -> int:
+    return sum(self.indexed_dataset.sizes.tolist())
+
+  @functools.cached_property
+  def num_epochs(self) -> int:
+    # we need enough epochs to cover the requested amount of tokens
+    num_epochs = 1
+    num_tokens = self.tokens_per_epoch
+    while num_tokens < self.samples * self.seqlen:
+      num_epochs += 1
+      num_tokens += self.tokens_per_epoch
+    return num_epochs
+
+  # https://github.com/NVIDIA/Megatron-LM/blob/94bd476bd840c2fd4c3ebfc7448c2af220f4832b/megatron/core/datasets/gpt_dataset.py#L558
+  def _build_doc_idx(self):
+    print(f"building doc_idx for {self.num_epochs=}, {self.indexed_dataset.count=}")
+    st = time.perf_counter()
+    # doc_idx = np.mgrid[:self.num_epochs, :self.indexed_dataset.count][1]
+    doc_idx = np.arange(self.indexed_dataset.count).reshape(1, -1).repeat(self.num_epochs, axis=0).flatten()
+    doc_idx = doc_idx.astype(np.int32)
+    at = time.perf_counter()
+    if self.shuffle: self.rng.shuffle(doc_idx)
+    print(f"doc_idx built in {at - st:.3f}s, shuffled in {time.perf_counter() - at:.3f}s")
+    return doc_idx
+
+  def _build_sample_idx(self):
+    print(f"building sample_idx for {self.samples=}, {self.seqlen=}, {self.doc_idx.shape[0]=}")
+    sample_idx_max = max(self.doc_idx.shape[0], self.indexed_dataset.sizes.max())
+    sample_idx = np.empty((self.samples + 1, 2), dtype=np.int64 if sample_idx_max > dtypes.int32.max else np.int32)
+
+    sample_idx_idx, doc_idx_idx, doc_offset = 0, 0, 0
+    sample_idx[sample_idx_idx, 0], sample_idx[sample_idx_idx, 1] = doc_idx_idx, doc_offset
+    sample_idx_idx += 1
+
+    for _ in tqdm(range(1, self.samples + 1)):
+      remaining_seqlen = self.seqlen + 1
+      while remaining_seqlen > 0:
+        doc_idx = int(self.doc_idx[doc_idx_idx])
+        doc_len = int(self.indexed_dataset.sizes[doc_idx]) - doc_offset
+        remaining_seqlen -= doc_len
+        if remaining_seqlen <= 0:
+          doc_offset += remaining_seqlen + doc_len - 1
+          remaining_seqlen = 0
+        else:
+          if doc_idx_idx == len(self.doc_idx) - 1:
+            assert sample_idx_idx == self.samples
+            doc_idx = int(self.doc_idx[doc_idx_idx])
+            doc_offset = int(self.indexed_dataset.sizes[doc_idx]) - 1
+            break
+          doc_idx_idx += 1
+          doc_offset = 0
+
+      sample_idx[sample_idx_idx, 0], sample_idx[sample_idx_idx, 1] = doc_idx_idx, doc_offset
+      sample_idx_idx += 1
+
+    return sample_idx
+
+  def _build_shuffle_idx(self):
+    print(f"building shuffle_idx for {self.samples=}")
+    st = time.perf_counter()
+    shuffle_idx = np.arange(self.samples, dtype=np.int32)
+    at = time.perf_counter()
+    if self.shuffle: self.rng.shuffle(shuffle_idx)
+    print(f"shuffle_idx built in {at - st:.3f}s, shuffled in {time.perf_counter() - at:.3f}s")
+    return shuffle_idx
+
+class BlendedGPTDataset:
+  def __init__(self, paths:list[Path], weights:list[float], samples:int, seqlen:int, seed:int, shuffle:bool):
+    self.shuffle = shuffle
+    self.rng = np.random.RandomState(seed)
+
+    # normalize weights
+    total_weight = sum(weights)
+    self.weights = [w / total_weight for w in weights]
+
+    self.samples = samples
+    surplus = 0.005
+    samples_per_blend = [math.ceil(math.ceil(self.samples * w) * (1 + surplus)) for w in self.weights]
+
+    self.datasets = [GPTDataset(path, samples_per_blend[i], seqlen, seed + i, shuffle) for i,path in enumerate(paths)]
+
+    # check for cache
+    cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
+    cache_path = paths[0].with_name(f"{paths[0].name}.{cache_hash}.blend_cache")
+    print(f"try loading BlendedGPTDataset from {cache_path}...")
+    if cache_path.exists():
+      print("cache found, loading...")
+      with open(cache_path, "rb") as f:
+        self.dataset_idx, self.dataset_sample_idx = pickle.load(f)
+    else:
+      print("cache not found, building index...")
+      self.dataset_idx, self.dataset_sample_idx = self._build_blend_idx()
+      # save cache
+      with open(cache_path, "wb") as f:
+        pickle.dump((self.dataset_idx, self.dataset_sample_idx), f)
+
+  def get(self, idx:int):
+    tokens = self.datasets[self.dataset_idx[idx]][self.dataset_sample_idx[idx]]
+    return tokens
+
+  def _build_blend_idx(self):
+    dataset_idx = np.zeros(self.samples, dtype=np.int16)
+    dataset_sample_idx = np.zeros(self.samples, dtype=np.int64)
+
+    unspent_datasets = set(range(len(self.datasets)))
+    dataset_sample_counts = [0] * len(self.datasets)
+
+    for i in tqdm(range(self.samples)):
+      error_argmax, error_max = 0, 0.0
+      for di in unspent_datasets:
+        error = self.weights[di] * max(i, 1) - dataset_sample_counts[di]
+        if error > error_max:
+          error_max = error
+          error_argmax = di
+
+      dataset_idx[i] = error_argmax
+      dataset_sample_idx[i] = dataset_sample_counts[error_argmax]
+
+      dataset_sample_counts[error_argmax] += 1
+
+    return dataset_idx, dataset_sample_idx
+
+def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
+  if val:
+    dataset = BlendedGPTDataset([
+      base_dir / "validation" / "c4-validationn-91205-samples.en_text_document",
+    ], [
+      1.0
+    ], samples, seqlen, seed, False)
+  else:
+    dataset = BlendedGPTDataset([
+      base_dir / "c4-train.en_6_text_document",
+      base_dir / "c4-train.en_7_text_document",
+    ], [
+      1.0, 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)
+
+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"
@@ -538,6 +807,18 @@ if __name__ == "__main__":
      for x in batch_load_retinanet(dataset, val, base_dir):
        pbar.update(x[0].shape[0])

+  def load_llama3(val):
+    bs = 24
+    samples = 5760 if val else 1_200_000 * 1152
+    seqlen = 8192
+
+    max_, min_ = 0, math.inf
+    for tokens in tqdm(batch_load_llama3(bs, samples, seqlen, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=5760, val=bool(val)), total=samples//bs):
+      max_ = max(max_, tokens.shape[1])
+      min_ = min(min_, tokens.shape[1])
+    print(f"max seq length: {max_}")
+    print(f"min seq length: {min_}")
+
  load_fn_name = f"load_{getenv('MODEL', 'resnet')}"
  if load_fn_name in globals():
    globals()[load_fn_name](getenv("VAL", 1))
@@ -212,7 +212,7 @@ def get_mlperf_bert_model():
  from examples.mlperf.initializers import LinearBert, EmbeddingBert, LayerNormBert

  bert.Linear = LinearBert
-  bert.Embedding = EmbeddingBert 
+  bert.Embedding = EmbeddingBert
  bert.LayerNorm = LayerNormBert

  from extra.models.bert import BertForPretraining
@@ -39,7 +39,7 @@ class LinearBert(nn.Linear):
  def __init__(self, in_features, out_features, bias=True, std=0.02):
    self.weight = std * rand_truncn(out_features, in_features, dtype=dtypes.float32)
    self.bias = Tensor.zeros(out_features, dtype=dtypes.float32) if bias else None
-  
+
  def __call__(self, x:Tensor):
    return x.cast(dtypes.default_float).linear(self.weight.cast(dtypes.default_float).transpose(), self.bias.cast(dtypes.default_float) if self.bias is not None else None)

@@ -1,4 +1,5 @@
-from tinygrad import Tensor, dtypes
+import math
+from tinygrad import dtypes
 from tinygrad.nn.optim import Optimizer

 from extra.lr_scheduler import LR_Scheduler
@@ -20,3 +21,19 @@ class PolynomialDecayWithWarmup(LR_Scheduler):
    warmup_lr = (self.epoch_counter * (1.0 / self.warmup)) * self.initial_lr
    x = (1 - (self.epoch_counter - self.warmup) / (self.epochs - self.warmup + 1))
    return (self.epoch_counter <= self.warmup).where(warmup_lr, (self.initial_lr - self.end_lr) * x ** self.power + self.end_lr).cast(self.optimizer.lr.dtype)
+
+class CosineAnnealingLRWithWarmup(LR_Scheduler):
+  def __init__(self, optimizer:Optimizer, base_lr, end_lr, warmup_steps:int, decay_steps:int):
+    assert warmup_steps > 0 and decay_steps > 0
+    super().__init__(optimizer)
+    self.base_lr = base_lr
+    self.end_lr = end_lr
+    self.warmup_steps = warmup_steps
+    self.decay_steps = decay_steps
+    # set lr for first warmup step
+    self.optimizer.lr.assign(self.get_lr()).realize()
+
+  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)
@@ -1,6 +1,6 @@
-import re
-import string
+import re, string
 from collections import Counter
+from tinygrad import Tensor

 def levenshtein(a, b):
  n, m = len(a), len(b)
@@ -59,3 +59,11 @@ def f1_score(x, y):
  p = ns / len(xt)
  r = ns / len(yt)
  return 2 * p * r / (p + r)
+
+def log_perplexity(logit:Tensor, target:Tensor, ignore_index:int|None=None):
+  # logit has shape (n_samples, seq_len, vocab_size), target has shape (n_samples, seq_len)
+  assert logit.ndim == 3, logit.ndim
+  assert target.ndim == 2, target.ndim
+  assert logit.shape[:2] == target.shape, f"{logit.shape[:2]=}, {target.shape=}"
+  log_prob = logit.log_softmax(axis=-1)
+  return log_prob.transpose(1, 2).nll_loss(target, ignore_index=ignore_index)
@@ -1,4 +1,4 @@
-import time
+import time, math
 start = time.perf_counter()
 from pathlib import Path
 import numpy as np
@@ -241,6 +241,52 @@ def eval_mrcnn():
  evaluate_predictions_on_coco(bbox_output, iou_type='bbox')
  evaluate_predictions_on_coco(mask_output, iou_type='segm')

+def eval_llama3():
+  from extra.models.llama import Transformer
+  from examples.llama3 import MODEL_PARAMS, load, convert_from_huggingface
+  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/"))
+
+  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)
+
+  @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()
+
+  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)
+
+  losses = []
+  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}")
+
 if __name__ == "__main__":
  # inference only
  Tensor.training = False
@@ -4,8 +4,8 @@ import multiprocessing

 from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
 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
+from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict, safe_load, safe_save
+from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW

 from extra.lr_scheduler import LRSchedulerGroup
 from examples.mlperf.helpers import get_training_state, load_training_state
@@ -252,6 +252,10 @@ 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:
@@ -344,6 +348,8 @@ 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
@@ -933,7 +939,7 @@ def train_step_bert(model, optimizer, scheduler, loss_scaler:float, GPUS, grad_a
    # TODO: OOM without this realize with large grad_acc
    Tensor.realize(*[p.grad for p in optimizer.params])

-  global_norm = Tensor([0.0], dtype=dtypes.float32, device=optimizer[0].device)
+  global_norm = Tensor(0.0, dtype=dtypes.float32, device=optimizer[0].device)
  for p in optimizer.params:
    p.grad = p.grad / loss_scaler
    global_norm += p.grad.float().square().sum()
@@ -1284,6 +1290,209 @@ def train_bert():
        MLLOGGER.start(key=mllog_constants.BLOCK_START, value=None, metadata={"first_epoch_num": 1, "epoch_num": 1, "epoch_count": 1, "samples_count": i * GBS, "step_num": i, "first_step_num": i+1})
        previous_step = i

+def train_llama3():
+  from extra.models.llama import Transformer
+  from examples.llama3 import MODEL_PARAMS
+  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 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
+  opt_adamw_beta_2 = 0.95
+  opt_adamw_epsilon = 1e-5
+  opt_adamw_weight_decay = 0.1
+
+  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_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152)  # NOTE: cannot change for benchmark
+  opt_end_learning_rate = getenv("END_LR", 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
+  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
+  model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
+
+  if getenv("FAKEDATA"):
+    for v in get_parameters(model):
+      v = v.assign(Tensor.empty(v.shape))
+
+  if (DP := getenv("DP", 1)) > 1:
+    device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
+    for v in get_parameters(model):
+      v.shard_(device, axis=None)
+
+  if (MP := getenv("MP", 1)) > 1:
+    device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
+    for k,v in get_state_dict(model).items():
+      if 'scale' in k: v.shard_(device, axis=None)  # from quantized
+      elif '.attention.wq' in k: v.shard_(device, axis=0)
+      elif '.attention.wk' in k: v.shard_(device, axis=0)
+      elif '.attention.wv' in k: v.shard_(device, axis=0)
+      elif '.attention.wo' in k: v.shard_(device, axis=1)
+      elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
+      elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
+      elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
+      elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
+      elif 'output.weight' in k: v.shard_(device, axis=0)
+      else:
+        # attention_norm, ffn_norm, norm
+        v.shard_(device, axis=None)
+      # prevents memory spike on device 0
+      v.realize()
+
+  optim = AdamW(get_parameters(model), lr=0.0,
+                b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
+  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):
+    optim.zero_grad()
+    # grad acc
+    for batch in tokens.split(tokens.shape[0]//grad_acc):
+      if (DP := getenv("DP", 1)) > 1:
+        device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
+        batch = batch.shard(device, 0)
+      if (MP := getenv("MP", 1)) > 1:
+        device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
+        batch = batch.shard(device)
+      logits:Tensor = model(batch[:, :-1], start_pos=0, temperature=math.nan)
+      loss = logits.sparse_categorical_crossentropy(batch[:, 1:])
+      loss.backward()
+      Tensor.realize(*[p.grad for p in optim.params])
+    # L2 norm grad clip
+    # https://github.com/NVIDIA/NeMo/blob/3368c3fc0b4a186ab33a1d68a504315100c0b2a6/nemo/collections/nlp/modules/common/megatron/clip_grads.py#L57
+    # https://docs.pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html
+    if not getenv("DISABLE_GRAD_CLIP_NORM"):
+      total_norm = Tensor(0.0, dtype=dtypes.float32, device=optim.params[0].device)
+      for p in optim.params:
+        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)
+
+    optim.step()
+    scheduler.step()
+
+    lr = optim.lr
+    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()
+
+  # ** 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
+  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]
+
+    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)):
+      tqdm.write("saving checkpoint")
+      if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
+      fn = f"{ckpt_dir}/llama3_{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
+
 if __name__ == "__main__":
  multiprocessing.set_start_method('spawn')

@@ -4,6 +4,8 @@ export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128

+export IGNORE_OOB=1
+
 export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
 # export BEAM_LOG_SURPASS_MAX=1
@@ -5,6 +5,8 @@ export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
 export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1

+export IGNORE_OOB=1
+
 export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
 export BASEDIR="/raid/datasets/wiki"
@@ -8,6 +8,8 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
 export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
 export TRAIN_STEPS=3900

+export IGNORE_OOB=1
+
 export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
 export BASEDIR="/raid/datasets/wiki"
@@ -11,6 +11,8 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
 export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
 export TRAIN_STEPS=3900

+export IGNORE_OOB=1
+
 export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
 export BASEDIR="/raid/datasets/wiki"
@@ -2,9 +2,9 @@

 export PYTHONPATH="." NV=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90

-export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+export IGNORE_OOB=1

 export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
@@ -2,9 +2,9 @@

 export PYTHONPATH="." NV=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90

-export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+export IGNORE_OOB=1

 export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
@@ -5,9 +5,9 @@ set -o pipefail  # Make pipeline fail if any command fails
 export PYTHONPATH="." NV=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_green"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90

-export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+export IGNORE_OOB=1

 export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
@@ -2,9 +2,9 @@

 export PYTHONPATH="." AMD=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90

-export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+export IGNORE_OOB=1

 export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
@@ -2,9 +2,9 @@

 export PYTHONPATH="." AMD=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90

-export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+export IGNORE_OOB=1

 export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
@@ -5,9 +5,9 @@ set -o pipefail  # Make pipeline fail if any command fails
 export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_red"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90

-export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+export IGNORE_OOB=1

 export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
@@ -1,8 +1,7 @@
 # https://arxiv.org/pdf/2409.02060
-import time
+import time, functools
 import numpy as np
 np.set_printoptions(suppress=True, linewidth=1000)
-import functools
 from tinygrad import Tensor, nn, Device, GlobalCounters
 from tinygrad.helpers import Timing, getenv
 from extra.models.llama import Transformer, convert_from_huggingface
@@ -17,7 +16,7 @@ class MixtureFeedForward:
  def __call__(self, x:Tensor) -> Tensor:
    assert x.shape[0] == 1, "only BS=1"
    assert x.shape[1] == 1, "only length=1"
-    g = self.gate(x).float().softmax(-1)
+    g = self.gate(x).softmax(-1)

    g = g.squeeze() # (BS, length, num_experts) -> (num_experts,)
    probs, sel = g.topk(self.activated_experts)
@@ -25,7 +24,7 @@ class MixtureFeedForward:
    # run MoE
    x_up_gate = x.dot(self.gate_proj[sel].permute(0,2,1)).silu() * x.dot(self.up_proj[sel].permute(0,2,1))
    x_down = x_up_gate.dot(self.down_proj[sel].permute(0,2,1))
-    return (x_down.float() * probs.reshape(self.activated_experts, 1, 1)).sum(axis=0)
+    return (x_down * probs.reshape(self.activated_experts, 1, 1)).sum(axis=0)

 # model is bf16, 1.3B active, 6.9B total
 # M3 Max is 400 GB/s, so 400/2.6 = ~154 tok/s
@@ -5,29 +5,26 @@ if "IMAGE" not in os.environ: os.environ["IMAGE"] = "2"
 if "NOLOCALS" not in os.environ: os.environ["NOLOCALS"] = "1"
 if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"

-from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device
+from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device, dtypes
 from tinygrad.helpers import DEBUG, getenv
-from tinygrad.tensor import _from_np_dtype
 from tinygrad.engine.realize import CompiledRunner

 import onnx
-from onnx.helper import tensor_dtype_to_np_dtype
-from tinygrad.frontend.onnx import OnnxRunner, onnx_load
+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"

 def compile(onnx_file):
-  onnx_model = onnx_load(onnx_file)
-  run_onnx = OnnxRunner(onnx_model)
+  run_onnx = OnnxRunner(onnx_file)
  print("loaded model")

-  input_shapes = {inp.name:tuple(x.dim_value for x in inp.type.tensor_type.shape.dim) for inp in onnx_model.graph.input}
-  input_types = {inp.name: tensor_dtype_to_np_dtype(inp.type.tensor_type.elem_type) for inp in onnx_model.graph.input}
+  input_shapes = {name: spec.shape for name, spec in run_onnx.graph_inputs.items()}
+  input_types = {name: spec.dtype for name, spec in run_onnx.graph_inputs.items()}
  # Float inputs and outputs to tinyjits for openpilot are always float32
-  input_types = {k:(np.float32 if v==np.float16 else v) for k,v in input_types.items()}
+  input_types = {k:(dtypes.float32 if v is dtypes.float16 else v) for k,v in input_types.items()}
  Tensor.manual_seed(100)
-  new_inputs = {k:Tensor.randn(*shp, dtype=_from_np_dtype(input_types[k])).mul(8).realize() for k,shp in sorted(input_shapes.items())}
+  new_inputs = {k:Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize() for k,shp in sorted(input_shapes.items())}
  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
  print("created tensors")

@@ -57,11 +54,11 @@ def compile(onnx_file):
      gated_read_image_count += ei.prg.p.src.count("?read_image")
  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"too many kernels! {kernel_count=}, {allowed_kernel_count=}"
+    assert kernel_count == allowed_kernel_count, f"different kernels! {kernel_count=}, {allowed_kernel_count=}"
  if (allowed_read_image:=getenv("ALLOWED_READ_IMAGE", -1)) != -1:
    assert read_image_count == allowed_read_image, f"different read_image! {read_image_count=}, {allowed_read_image=}"
  if (allowed_gated_read_image:=getenv("ALLOWED_GATED_READ_IMAGE", -1)) != -1:
-    assert gated_read_image_count <= allowed_gated_read_image, f"too many gated read_image! {gated_read_image_count=}, {allowed_gated_read_image=}"
+    assert gated_read_image_count == allowed_gated_read_image, f"different gated read_image! {gated_read_image_count=}, {allowed_gated_read_image=}"

  with open(OUTPUT, "wb") as f:
    pickle.dump(run_onnx_jit, f)
@@ -1,23 +1,21 @@
-import sys, onnx
-from tinygrad import Tensor, fetch, GlobalCounters
+import sys
+from tinygrad import Tensor, fetch, GlobalCounters, dtypes
 from tinygrad.uop.ops import UOp
 from tinygrad.frontend.onnx import OnnxRunner
-from tinygrad.engine.kernelize import get_kernelize_map
+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 GPU=1 IMAGE=2 FLOAT16=1 VIZ=1 DEBUG=2 python3 examples/openpilot/compile4.py
+# NOLOCALS=1 CL=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"

 if __name__ == "__main__":
-  fn = fetch(OPENPILOT_MODEL)
  onnx_file = fetch(OPENPILOT_MODEL)
-  onnx_model = onnx.load(onnx_file)
-  run_onnx = OnnxRunner(onnx_model)
+  run_onnx = OnnxRunner(onnx_file)

-  inputs = run_onnx.get_empty_input_data("npy")
+  inputs = run_onnx.get_empty_input_data("npy", dtypes.float32)
  out: Tensor = next(iter(run_onnx({k:v.to(None) for k,v in inputs.items()}).values())).to('cpu')
  root = out.uop
  targets = [x.uop for x in inputs.values()]
@@ -27,7 +27,7 @@ class Model(nn.Module):

 if __name__ == "__main__":
  if getenv("TINY_BACKEND"):
-    import tinygrad.frontend.torch
+    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
+from tinygrad import nn, Tensor, Device
 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 GPU" if DEBUG>=2 else "") +
+      with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" 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)
@@ -5,8 +5,8 @@

 from tinygrad import Tensor, TinyJit, dtypes, GlobalCounters
 from tinygrad.nn import Conv2d, GroupNorm
-from tinygrad.nn.state import safe_load, load_state_dict, get_state_dict
-from tinygrad.helpers import fetch, trange, colored, Timing
+from tinygrad.nn.state import safe_load, load_state_dict
+from tinygrad.helpers import fetch, trange, colored, Timing, getenv
 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
+import argparse, tempfile, time
 from abc import ABC, abstractmethod
 from pathlib import Path
 from PIL import Image
@@ -342,11 +342,13 @@ 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,
@@ -358,8 +360,13 @@ 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


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

  # validation!
-  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:
+  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:
    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,7 +2,7 @@
 # https://github.com/ekagra-ranjan/huggingface-blog/blob/main/stable_diffusion.md
 import tempfile
 from pathlib import Path
-import argparse
+import argparse, time
 from collections import namedtuple
 from typing import Dict, Any

@@ -266,17 +266,23 @@ 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)
@@ -7,7 +7,7 @@
 from examples.beautiful_mnist import Model
 from tinygrad import Tensor, nn, getenv, GlobalCounters, Variable
 from tinygrad.nn.datasets import mnist
-from tinygrad.helpers import trange, DEBUG
+from tinygrad.helpers import trange

 # STEPS=70 python3 examples/stunning_mnist.py
 # NOTE: it's broken with STACK=1, why?
@@ -2,10 +2,9 @@
 #!POPCORN gpu A100
 # not a stable API, but works

-import torch, functools
+import torch
 from tinygrad import Tensor, TinyJit, Device
-from tinygrad.engine.realize import CompiledRunner
-from tinygrad.helpers import get_single_element, Context, OSX
+from tinygrad.helpers import Context, OSX
 from tinygrad.dtype import _from_torch_dtype

@TinyJit
@@ -2,8 +2,6 @@ import sys
 import random
 import json
 import numpy
-from pathlib import Path
-from PIL import Image
 from tinygrad.tensor import Tensor
 from tinygrad.nn.optim import SGD
 from tinygrad.nn.state import safe_save, safe_load, get_state_dict, load_state_dict
@@ -5,7 +5,7 @@ from typing import Optional, Union, Literal, List

 from tinygrad import Tensor, TinyJit, Variable, nn
 from tinygrad.nn.state import torch_load, load_state_dict
-from tinygrad.helpers import getenv, DEBUG, fetch
+from tinygrad.helpers import getenv, fetch

 import numpy as np
 import librosa
@@ -321,7 +321,7 @@ if __name__ == "__main__":
        log_spec = prep_audio(total.reshape(1, -1), model.batch_size, truncate=True)
        encoded_audio = model.encoder.encode(Tensor(log_spec))
      # pass the previously inferred tokens as 'prefix' - https://github.com/openai/whisper/discussions/117#discussioncomment-3727051
-      out = model.decoder(Tensor([lst]), 0, encoded_audio, streaming=True).realize()
+      out = model.decoder(Tensor([lst]), 0, encoded_audio).realize()
      idx = int(out[0,-1].argmax().numpy().item())
      lst.append(idx)
      dec = enc.decode(lst)
@@ -71,8 +71,8 @@ def bbox_iou(box1, box2):
  # get the coordinates of the intersection rectangle
  inter_rect_x1 = np.maximum(b1_x1, b2_x1)
  inter_rect_y1 = np.maximum(b1_y1, b2_y1)
-  inter_rect_x2 = np.maximum(b1_x2, b2_x2)
-  inter_rect_y2 = np.maximum(b1_y2, b2_y2)
+  inter_rect_x2 = np.minimum(b1_x2, b2_x2)
+  inter_rect_y2 = np.minimum(b1_y2, b2_y2)
  #Intersection area
  inter_area = np.clip(inter_rect_x2 - inter_rect_x1 + 1, 0, 99999) * np.clip(inter_rect_y2 - inter_rect_y1 + 1, 0, 99999)
  #Union Area
@@ -297,13 +297,13 @@ class Darknet:
          # Get the number of weights of batchnorm
          num_bn_biases = math.prod(bn.bias.shape)
          # Load weights
-          bn_biases = Tensor(weights[ptr:ptr + num_bn_biases])
+          bn_biases = Tensor(weights[ptr:ptr + num_bn_biases].astype(np.float32))
          ptr += num_bn_biases
-          bn_weights = Tensor(weights[ptr:ptr+num_bn_biases])
+          bn_weights = Tensor(weights[ptr:ptr+num_bn_biases].astype(np.float32))
          ptr += num_bn_biases
-          bn_running_mean = Tensor(weights[ptr:ptr+num_bn_biases])
+          bn_running_mean = Tensor(weights[ptr:ptr+num_bn_biases].astype(np.float32))
          ptr += num_bn_biases
-          bn_running_var = Tensor(weights[ptr:ptr+num_bn_biases])
+          bn_running_var = Tensor(weights[ptr:ptr+num_bn_biases].astype(np.float32))
          ptr += num_bn_biases
          # Cast the loaded weights into dims of model weights
          bn_biases = bn_biases.reshape(shape=tuple(bn.bias.shape))
@@ -319,7 +319,7 @@ class Darknet:
          # load biases of the conv layer
          num_biases = math.prod(conv.bias.shape)
          # Load weights
-          conv_biases = Tensor(weights[ptr: ptr+num_biases])
+          conv_biases = Tensor(weights[ptr: ptr+num_biases].astype(np.float32))
          ptr += num_biases
          # Reshape
          conv_biases = conv_biases.reshape(shape=tuple(conv.bias.shape))
@@ -327,7 +327,7 @@ class Darknet:
          conv.bias = conv_biases
        # Load weighys for conv layers
        num_weights = math.prod(conv.weight.shape)
-        conv_weights = Tensor(weights[ptr:ptr+num_weights])
+        conv_weights = Tensor(weights[ptr:ptr+num_weights].astype(np.float32))
        ptr += num_weights
        conv_weights = conv_weights.reshape(shape=tuple(conv.weight.shape))
        conv.weight = conv_weights
@@ -371,7 +371,7 @@ class Darknet:
 if __name__ == "__main__":
  model = Darknet(fetch('https://raw.githubusercontent.com/pjreddie/darknet/master/cfg/yolov3.cfg').read_bytes())
  print("Loading weights file (237MB). This might take a while…")
-  model.load_weights('https://pjreddie.com/media/files/yolov3.weights')
+  model.load_weights('https://github.com/shadiakiki1986/yolov3.weights/releases/download/3.0.1/yolov3.weights')
  if len(sys.argv) > 1:
    url = sys.argv[1]
  else:
@@ -2,14 +2,12 @@
 import os
 from ultralytics import YOLO
 from pathlib import Path
-from tinygrad.frontend.onnx import OnnxRunner, onnx_load
+from tinygrad.frontend.onnx import OnnxRunner
 from extra.onnx_helpers import get_example_inputs
-from tinygrad.tensor import Tensor

 os.chdir("/tmp")
 if not Path("yolov8n-seg.onnx").is_file():
  model = YOLO("yolov8n-seg.pt")
  model.export(format="onnx", imgsz=[480,640])
-onnx_model = onnx_load(open("yolov8n-seg.onnx", "rb"))
-run_onnx = OnnxRunner(onnx_model)
+run_onnx = OnnxRunner("yolov8n-seg.onnx")
 run_onnx(get_example_inputs(run_onnx.graph_inputs), debug=True)
@@ -1,7 +1,5 @@
 from tinygrad.nn import Conv2d, BatchNorm2d
 from tinygrad.tensor import Tensor
-from tinygrad.device import is_dtype_supported
-from tinygrad import dtypes
 import numpy as np
 from itertools import chain
 from pathlib import Path
@@ -8,8 +8,6 @@ from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager, AMPageTableEntry
 from tinygrad.runtime.support.am.ip import AM_SOC, AM_GMC, AM_IH, AM_PSP, AM_SMU, AM_GFX, AM_SDMA

-AM_VERSION = 0xA0000005
-
 def bold(s): return f"\033[1m{s}\033[0m"

 def trim(s:str, length:int) -> str:
@@ -73,21 +71,12 @@ class AMSMI(AMDev):
    self._run_discovery()
    self._build_regs()

-    if self.reg("regSCRATCH_REG7").read() != AM_VERSION:
+    if self.reg("regSCRATCH_REG7").read() != AMDev.Version:
      raise Exception(f"Unsupported AM version: {self.reg('regSCRATCH_REG7').read():x}")

-    self.is_booting, self.smi_dev = True, True
+    self.is_booting = True
+    self.init_sw(smi_dev=True)
    self.partial_boot = True # do not init anything
-    self.mm = AMMemoryManager(self, self.vram_size)
-
-    # Initialize IP blocks
-    self.soc:AM_SOC = AM_SOC(self)
-    self.gmc:AM_GMC = AM_GMC(self)
-    self.ih:AM_IH = AM_IH(self)
-    self.psp:AM_PSP = AM_PSP(self)
-    self.smu:AM_SMU = AM_SMU(self)
-
-    for ip in [self.soc, self.gmc, self.ih, self.psp, self.smu]: ip.init_sw()

  def read_pci_state(self):
    with open(f"/sys/bus/pci/devices/{self.pcibus}/power_state", "r") as f: return f.read().strip().rstrip()
@@ -136,7 +125,7 @@ class SMICtx:
        if d.pci_state == "D0": d._init_from_d0()
        os.system('clear')

-      if d.pci_state == "D0" and d.reg("regSCRATCH_REG7").read() != AM_VERSION:
+      if d.pci_state == "D0" and d.reg("regSCRATCH_REG7").read() != AMDev.Version:
        self.devs.remove(d)
        self.opened_pcidevs.remove(d.pcibus)
        os.system('clear')
@@ -295,8 +284,8 @@ if __name__ == "__main__":
          while True:
            try: pid = subprocess.check_output(['sudo', 'lsof', '-t', dev]).decode('utf-8').split('\n')[0]
            except subprocess.CalledProcessError: break
-            if stopped_pids[pid] > 0: time.sleep(0.5)
-            if stopped_pids[pid] == 10:
+            if stopped_pids[pid] > 0: time.sleep(0.1)
+            if stopped_pids[pid] == 64:
              print(f"{dev[8:-5]}: can't stop process {pid}, exitting")
              exit(1)

@@ -1,6 +1,16 @@
-import re, ctypes, sys
+import re, ctypes, sys, importlib

-from tinygrad.runtime.autogen.am import am, mp_11_0, mp_13_0_0, nbio_4_3_0, mmhub_3_0_0, gc_11_0_0, osssys_6_0_0
+from tinygrad.runtime.support.am.amdev import AMDev, AMRegister
+class AMDFake(AMDev):
+  def __init__(self, devfmt, vram, doorbell, mmio, dma_regions=None):
+    self.devfmt, self.vram, self.doorbell64, self.mmio, self.dma_regions = devfmt, vram, doorbell, mmio, dma_regions
+    self._run_discovery()
+    self._build_regs()
+
+amdev = importlib.import_module("tinygrad.runtime.support.am.amdev")
+amdev.AMDev = AMDFake
+
+from tinygrad.runtime.ops_amd import PCIIface

 def parse_amdgpu_logs(log_content, register_names=None):
  register_map = register_names
@@ -23,16 +33,11 @@ def parse_amdgpu_logs(log_content, register_names=None):
  return processed_log

 def main():
-  regs_offset = {13: {0: [3072, 37784576]}, 28: {0: [93184, 37754880], 1: [201327616, 201461760], 2: [209716224, 209850368], 3: [218104832, 218238976], 4: [226493440, 226627584], 5: [234882048, 235016192], 6: [243270656, 243404800]}, 21: {0: [28672, 12582912, 37795840, 130023424, 306184192], 1: [201326592, 201463808, 201465856, 204210176, 204472320], 2: [209715200, 209852416, 209854464, 212598784, 212860928], 3: [218103808, 218241024, 218243072, 220987392, 221249536], 4: [226492416, 226629632, 226631680, 229376000, 229638144], 5: [234881024, 235018240, 235020288, 237764608, 238026752], 6: [243269632, 243406848, 243408896, 246153216, 246415360]}, 22: {0: [18, 192, 13504, 36864, 37764096]}, 1: {0: [4704, 40960, 114688, 37760000]}, 2: {0: [3872, 37790720]}, 11: {0: [70656, 38103040]}, 12: {0: [106496, 37783552]}, 15: {0: [90112, 14417920, 14680064, 14942208, 38009856]}, 16: {0: [90112, 14417920, 14680064, 14942208, 38009856]}, 14: {0: [0, 20, 3360, 66560, 37859328, 67371008]}, 26: {0: [0, 20, 3360, 66560, 37859328, 67371008]}, 23: {0: [4256, 37789696]}, 33: {0: [0, 20, 3360, 66560, 37859328, 67371008]}, 25: {0: []}, 3: {0: [4704, 40960, 114688, 37760000]}, 4: {0: [4704, 40960, 114688, 37760000]}, 24: {0: [92160, 92672, 37752832, 54788096]}, 27: {0: [91648, 37751808], 1: [201339904, 201458176], 2: [209728512, 209846784], 3: [218117120, 218235392], 4: [226505728, 226624000], 5: [234894336, 235012608], 6: [243282944, 243401216]}, 29: {0: [201342976, 201344000, 205520896, 205537280], 1: [209731584, 209732608, 213909504, 213925888], 2: [218120192, 218121216, 222298112, 222314496], 3: [226508800, 226509824, 230686720, 230703104], 4: [234897408, 234898432, 239075328, 239091712], 5: [243286016, 243287040, 247463936, 247480320]}, 17: {0: [30720, 32256], 1: [31488, 73728]}}
-
  reg_names = {}
-  def _prepare_registers(modules):
-    for base, m in modules:
-      for k, regval in m.__dict__.items():
-        if k.startswith("reg") and not k.endswith("_BASE_IDX") and (base_idx:=getattr(m, f"{k}_BASE_IDX", None)) is not None:
-          reg_names[regs_offset[am.__dict__.get(f"{base}_HWIP")][0][base_idx] + regval] = k
-
-  _prepare_registers([("MP0", mp_13_0_0), ("NBIO", nbio_4_3_0), ("MMHUB", mmhub_3_0_0), ("GC", gc_11_0_0), ("OSSSYS", osssys_6_0_0)])
+  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}"

  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_gpu import CLProgram, CLBuffer
+from tinygrad.runtime.ops_cl import CLProgram, CLBuffer
 from tinygrad import dtypes
 from tqdm import trange, tqdm
 from matplotlib import pyplot as plt
@@ -1,5 +1,5 @@
 from typing import Tuple, List, NamedTuple, Any, Dict, Optional, Union, DefaultDict, cast
-from tinygrad.opt.kernel import Ops, MemOp, UOp
+from tinygrad.codegen.opt.kernel import Ops, MemOp, UOp
 from tinygrad.uop.ops import BinaryOps, UnaryOps
 from tinygrad.dtype import DType, dtypes
 from tinygrad.helpers import DEBUG
@@ -156,7 +156,7 @@ def uops_to_asmstyle(lang, function_name:str, uops:List[UOp]):
        lang.ins.append(AssemblyInstruction(Ops.ALU, out, [tmp], args))
      else:
        lang.ins.append(AssemblyInstruction(Ops.ALU, out, [lang.tor[x] for x in vin], args))
-    elif uop == Ops.DEFINE_ACC:
+    elif uop == Ops.DEFINE_REG:
      reg = lang.newreg(u, dtype=dtype)
      lang.ins.append(AssemblyInstruction(Ops.LOAD, reg, [], args))
    elif uop == Ops.SPECIAL:
@@ -3,7 +3,7 @@ from platform import system
 from typing import Tuple, Dict, List, Optional
 from tinygrad import dtypes
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
-from tinygrad.opt.kernel import Ops, UOp
+from tinygrad.codegen.opt.kernel import Ops, UOp
 from tinygrad.helpers import CI
 from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage

@@ -1,7 +1,7 @@
 from typing import List
 import struct
 from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
-from tinygrad.opt.kernel import Ops, UOp
+from tinygrad.codegen.opt.kernel import Ops, UOp
 from tinygrad import dtypes
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
 from tinygrad.runtime.ops_cuda import arch
@@ -2,9 +2,9 @@ import yaml
 from typing import Tuple, Set, Dict
 from tinygrad import dtypes
 from tinygrad.codegen.assembly import AssemblyCodegen, Register
-from tinygrad.opt.kernel import Ops
+from tinygrad.codegen.opt.kernel import Ops
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
-from tinygrad.runtime.ops_gpu import ROCM_LLVM_PATH
+from tinygrad.runtime.ops_cl 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_gpu import CLProgram, CLBuffer, ROCM_LLVM_PATH
+from tinygrad.runtime.ops_cl 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[Variable, int]):
+  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[str, 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.expr}" for v in var_vals])
+    args += sorted([f"int {v}" 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[Variable, int], wait=False):
+  def __call__(self, rawbufs: List[Buffer], var_vals: Dict[str, 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].expr)]), 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])]), 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[Variable, int]):
+  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[str, 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]])
+      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])

    # 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[Variable, int], wait=False) -> Optional[float]:
+  def __call__(self, input_rawbuffers: List[Buffer], var_vals: Dict[str, 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])
+        self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[v.expr])

    # Update launch dims
    for j in self.jc_idx_with_updatable_launch_dims:
@@ -15,8 +15,8 @@ def uops_to_rdna(function_name:str, uops:UOpGraph) -> str:
        u.vin = tuple(n if x == o else x for x in u.vin)
    # pointer indexing
    if u.uop in {UOps.LOAD, UOps.STORE} and u.vin[0].dtype.itemsize > 1:
-      val = UOp(UOps.CONST, dtypes.int, tuple(), arg=u.vin[0].dtype.itemsize, insert_before=uops.uops.index(u))
-      ptr = UOp(UOps.ALU, dtypes.int, (u.vin[1], val), arg=BinaryOps.MUL, insert_before=uops.uops.index(u))
+      val = UOp(UOps.CONST, dtypes.int, tuple(), arg=u.vin[0].dtype.itemsize, insert_at=uops.uops.index(u))
+      ptr = UOp(UOps.ALU, dtypes.int, (u.vin[1], val), arg=BinaryOps.MUL, insert_at=uops.uops.index(u))
      u.vin = (u.vin[0], ptr) + u.vin[2:]
  #uops.print()

@@ -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[1].startswith("lidx"):
-        r[u] = f'v{u.arg[0]}'
-      elif u.arg[1].startswith("gidx"):
-        r[u] = f's{2+u.arg[0]}'
+      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}'
      else:
        raise NotImplementedError
    elif u.uop == UOps.CONST:
@@ -2,7 +2,7 @@ from typing import Dict, List, Final, Callable, DefaultDict
 from collections import defaultdict
 from tinygrad.uop.ops import UnaryOps, BinaryOps, TernaryOps, Op
 from tinygrad.helpers import DType, PtrDType, dtypes, ImageDType, DEBUG, getenv
-from tinygrad.opt.kernel import  UOp, Ops
+from tinygrad.codegen.opt.kernel import  UOp, Ops
 from triton.compiler import compile as triton_compile
 import linecache
 import math
@@ -88,7 +88,7 @@ def uops_to_triton(function_name:str, uops:List[UOp]):
      assert dtype is not None
      if len(vin) == 2: kk(f"{ssa(u, 'val')} = {render_cast(f'tl.load({r[vin[0]]} + { fill_dims_for_idx(r[vin[1]], dims)}, mask = {render_valid(valid)})', dtype)}")
      else: kk(f"{ssa(u, 'val')} = {render_cast(f'tl.where({r[vin[2]]}, tl.load({r[vin[0]]}+{fill_dims_for_idx(r[vin[1]],dims)} , mask={render_valid(valid+[r[vin[2]]])}), 0.0)', dtype)}")
-    elif uop == Ops.DEFINE_ACC: kk(f"{ssa(u, 'acc')} = {define_scalar(local_size, dtype, args).replace('//', '/')}")
+    elif uop == Ops.DEFINE_REG: kk(f"{ssa(u, 'acc')} = {define_scalar(local_size, dtype, args).replace('//', '/')}")
    elif uop == Ops.CONST: r[u] = define_scalar([], dtype, args)
    elif uop == Ops.ASSIGN:
      kk(f"{r[vin[0]]} = {r[vin[1]].replace('//', '/')}")
@@ -4,13 +4,13 @@ from tinygrad.renderer import ProgramSpec
 from tinygrad.tensor import Device, Tensor
 from tinygrad.engine.jit import TinyJit
 from tinygrad.nn.state import get_state_dict
-from tinygrad.helpers import Context
+from tinygrad.helpers import Context, to_mv
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import Ops
 import json
 from collections import OrderedDict

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

 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,11 +67,12 @@ 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(cl._buf)])
+      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)});" for (name, args, _global_size, _local_size) in statements] + ["}"]
+    cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)}, {thread_id});" for (name, args, _global_size, _local_size) in statements] + ["}"]
    return '\n'.join(headers + cprog)
  else:
    if bufs_to_save:
@@ -239,7 +240,9 @@ 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"
-  with Context(JIT=2): run,special_names = jit_model(model, *inputs)
+
+  # NOTE: CPU_COUNT=1, since export does not support threading
+  with Context(JIT=2, CPU_COUNT=1): 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()}
@@ -1,96 +1,42 @@
 # kernel8_batched_gmem.s from https://seb-v.github.io/optimization/update/2025/01/20/Fast-GPU-Matrix-multiplication.html
+# sudo PATH=/opt/homebrew/Cellar/llvm/20.1.6/bin:$PATH AMD_LLVM=0 AMD=1 DEBUG=2 python3 extra/gemm/amd_matmul.py
 import pathlib
-import numpy as np
 from dataclasses import replace
-from tinygrad import Tensor, Device, Context
+from tinygrad import Tensor, Device, Context, GlobalCounters
 from tinygrad.helpers import getenv
-from tinygrad.opt.kernel import Kernel, Opt, OptOps
-from tinygrad.engine.realize import CompiledRunner, ExecItem
-from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops, UOp
-
-# TODO: on METAL for `DEBUG=4 python3 extra/gemm/amd_matmul.py`
-#  * fix load grouping (like float4). idk why it's not working, need new devectorizer (this is a Monday project)
-#  * DONE - remove extra barrier
-#  * DONE (moved Ops.ADD) - fix load order to be in order (the +0 one is last!)
-#  * explore async (fast) global load -> local store
-#  * why is TC=3 broken for 4096x4096?
-#  * write syntactic sugar for these local additions + use it in tensor core kernel.py
+from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program

 N = 4096
-LN = 16
 run_count = 5

-from tinygrad.shape.shapetracker import ShapeTracker, View
-def transform_load(ctx:tuple[Kernel, set[UOp]], x:UOp):
-  if x.src[0].op is not Ops.DEFINE_GLOBAL: return None
-  if x in ctx[1]: return None
-  print(ctx[0].colored_shape())
-  ctx[1].add(x)
-  input_st: ShapeTracker = x.src[1].arg
-  #strides = input_st.real_strides()
-  #strides = (0,0)+strides[2:]
-  if input_st.real_strides()[2] == 0:
-    perm = (0,1,5,3,4,2)
-    strides = (0,0,LN*4,4,0,0,1,0)
-  elif input_st.real_strides()[3] == 0:
-    perm = (0,1,2,5,4,3)
-    strides = (0,0,LN*4,4,0,0,0,1)
-  else:
-    return None
-  if len(input_st.shape) == 8:
-    local_st = ShapeTracker(views=(View.create((1,1,LN,LN,1,1,4,4), strides),))
-    perm = perm + (6,7)
-  else:
-    local_st = ShapeTracker(views=(View.create((1,1,LN,LN,1,1)),))
-  #local_st = ShapeTracker(views=(View.create((1,1,LN,LN,1,1)),))
-  load_st = local_st.permute(perm)
-  input_st = input_st.permute(perm)
-  lcl = UOp(Ops.DEFINE_LOCAL, x.dtype.ptr(local_st.real_size(), local=True), (), f"temp{x.src[0].arg}")
-  global_load = x.replace(src=(x.src[0], input_st.to_uop()))
-  ret = UOp(Ops.STORE, src=(lcl, local_st.to_uop(), global_load))
-  return UOp(Ops.LOAD, x.dtype, src=(lcl, load_st.to_uop(), ret))
-
-local_loads_pm = PatternMatcher([
-  (UPat(Ops.LOAD, name="x"), transform_load),
-])
-
-def ast_transform(k, ast):
-  #return ast
-  ast = graph_rewrite(ast, local_loads_pm, ctx=(k, set()))
-  #ast = ast.replace(arg=replace(ast.arg, upcasted=0))
-  print(ast)
-  return ast
-
 if __name__ == "__main__":
-  rng = np.random.default_rng()
-  a = Tensor(na:=rng.random((4096, 4096), dtype=np.float32)).realize()
-  b = Tensor(nb:=rng.random((4096, 4096), dtype=np.float32)).realize()
-  c = a @ b
-  si = c.schedule()[-1]
-  k = Kernel(si.ast, opts=Device[Device.DEFAULT].renderer)
-  #opts = [Opt(op=OptOps.LOCAL, axis=1, arg=16),
-  #        Opt(op=OptOps.LOCAL, axis=0, arg=8),
-  #        Opt(op=OptOps.UPCAST, axis=2, arg=4),
-  #        Opt(op=OptOps.UPCAST, axis=1, arg=4),
-  #        Opt(op=OptOps.UPCAST, axis=0, arg=2)]
-  #opts = [Opt(op=OptOps.UPCAST, axis=1, arg=4),
-  #        Opt(op=OptOps.UPCAST, axis=0, arg=4),
-  #        Opt(op=OptOps.LOCAL, axis=1, arg=8),
-  #        Opt(op=OptOps.LOCAL, axis=0, arg=4)]
-  opts = [Opt(op=OptOps.UNROLL, axis=0, arg=LN),
-          #Opt(op=OptOps.UPCAST, axis=0, arg=4),
-          #Opt(op=OptOps.UPCAST, axis=1, arg=4),
-          Opt(op=OptOps.LOCAL, axis=1, arg=LN),
-          Opt(op=OptOps.LOCAL, axis=0, arg=LN)]
-  k.apply_opts(opts)
-  prg = k.to_program(ast_transform=ast_transform)
-  if getenv("FAST", 1) and Device.DEFAULT == "AMD":
-    #src = (pathlib.Path(__file__).parent / "fp32_sgemm_amd" / "src" / "kernel8_batched_gmem.s").read_text()
-    src = (pathlib.Path(__file__).parent / "kernel8_batched_gmem.s").read_text()
-    prg = replace(prg, src=src, global_size=[N//128, N//128, 1], local_size=[128, 1, 1])
-  print(prg.global_size, prg.local_size)
-  ei = ExecItem(CompiledRunner(prg), [x.ensure_allocated() for x in si.bufs], si.metadata)
+  ast = (Tensor.empty(N, N)@Tensor.empty(N, N)).schedule()[-1].ast
+  prg = get_program(ast, Device.default.renderer)
+
+  if getenv("ASM") == 1:
+    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel8_batched_gmem.s").read_text()
+    prgfast = replace(prg, name="kernel", src=src, global_size=[N//128, N//128, 1], local_size=[128, 1, 1])
+  elif getenv("ASM") == -1:
+    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel3_registers.cpp").read_text()
+    prgfast = replace(prg, name="kernel3_registers", src=src, global_size=[N//128, N//128, 1], local_size=[256, 1, 1])
+  elif getenv("ASM") == -2:
+    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel4_gmem_df.cpp").read_text()
+    prgfast = replace(prg, name="kernel4_gmem_db", src=src, global_size=[N//128, N//128, 1], local_size=[256, 1, 1])
+  else:
+    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel5_lds_optim.cpp").read_text()
+    prgfast = replace(prg, name="kernel5_lds_optim", src=src, global_size=[N//128, N//128, 1], local_size=[128, 1, 1])
+  runner = CompiledRunner(prgfast)
+
+  a = Tensor.randn(N, N).realize()
+  b = Tensor.randn(N, N).realize()
+  c = Tensor.zeros(N, N).contiguous().realize()
+
+  GlobalCounters.reset()
+  with Context(DEBUG=2):
+    for _ in range(run_count): tc = (a@b).realize()
+
+  GlobalCounters.reset()
+  ei = ExecItem(runner, [a.uop.buffer, b.uop.buffer, c.uop.buffer])
  with Context(DEBUG=2):
    for _ in range(run_count): ei.run(wait=True)
-  nc = c.numpy()
-  np.testing.assert_allclose(na@nb, nc, rtol=1e-5)
+  print(f"custom  {(c-tc).square().mean().item()}")
@@ -0,0 +1,143 @@
+typedef long unsigned int size_t;
+extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
+extern "C" __attribute__((device, const)) size_t __ockl_get_group_id(unsigned int);
+struct Dim3 { size_t x, y, z; };
+#define __shared__ __attribute__((shared, aligned(16)))
+__attribute__((device)) inline void __syncthreads() {
+  __builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup");
+  __builtin_amdgcn_s_barrier();
+  __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");
+}
+
+#define BLOCK_SIZE 256
+extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, BLOCK_SIZE)))
+kernel3_registers(float *a, float *b, float *c)
+{
+  constexpr int N = 4096;
+  constexpr float alpha = 1.0;
+  constexpr float beta = 0.0;
+
+  const Dim3 blockIdx{ __ockl_get_group_id(0), __ockl_get_group_id(1), __ockl_get_group_id(2) };
+  const Dim3 threadIdx{ __ockl_get_local_id(0), __ockl_get_local_id(1), __ockl_get_local_id(2) };
+
+  // Block Tile size
+  constexpr int BN = 128;
+  constexpr int BM = 128;
+  // Number of Row or column we read per batch
+  constexpr int BK = 8;
+
+  // Thread Tile size
+  constexpr int TN = 4;
+  constexpr int TM = 4;
+
+  constexpr int nbWaves = BLOCK_SIZE / 32;
+  // Wave Tile size
+  constexpr int WN = 64;
+  constexpr int WM = BN * BM / nbWaves / WN;
+
+  // Number of wave on X & Y axis in the Block tile
+  constexpr int nbWaveX = BN / WN;
+  constexpr int nbWaveY = BM / WM;
+
+  const int waveIndex = threadIdx.x / 32;
+  const int waveIdx = waveIndex % nbWaveX;
+  const int waveIdy = waveIndex / nbWaveX;
+  const int indexInWave = threadIdx.x % 32;
+
+  // A wave is a block of 8x4 of the output matrix
+  constexpr int nbThreadXPerWave = 8;
+  constexpr int nbThreadYPerWave = 4;
+
+  // Thread coordinates in Wave
+  const int idxInWave = indexInWave % nbThreadXPerWave;
+  const int idyInWave = indexInWave / nbThreadXPerWave;
+
+  constexpr int nbIterWaveN = WN / (nbThreadXPerWave * TN);
+  constexpr int nbIterWaveM = WM / (nbThreadYPerWave * TM);
+
+  // Wave Sub-tile size
+  constexpr int SUBWN = WN / nbIterWaveN;
+  constexpr int SUBWM = WM / nbIterWaveM;
+
+  // Thread mapping to read BKxBN block from A
+  int rAIdx = threadIdx.x % BK;
+  int rAIdy = threadIdx.x / BK;
+  // Thread mapping to read BNxBK block from B
+  int rBIdx = threadIdx.x % BN;
+  int rBIdy = threadIdx.x / BN;
+
+  constexpr int strideReadB = BLOCK_SIZE / BN;
+  constexpr int strideReadA = BLOCK_SIZE / BK;
+  constexpr int nbReadsB = BN * BK / BLOCK_SIZE;
+  constexpr int nbReadsA = BM * BK / BLOCK_SIZE;
+
+  float A_col[nbIterWaveM * TM];
+  float B_row[nbIterWaveN * TN];
+
+  __shared__ float As[BK][BM];
+  __shared__ float Bs[BK][BN];
+
+  float c_regs[TM * nbIterWaveM * TN * nbIterWaveN] = {0.0f};
+
+  // Iteration over BK blocks.
+  for (int kId = 0; kId < N; kId += BK) {
+    __syncthreads();
+
+    // We populate the Shared Memory with Ks row and columns
+    for (int i = 0; i < nbReadsB; i++) {
+      int index_x = BN * blockIdx.x + rBIdx;
+      int index_y = rBIdy + i * strideReadB + kId;
+      Bs[index_y % BK][index_x % BN] = b[N * index_y + index_x];
+    }
+
+    for (int i = 0; i < nbReadsA; i++) {
+      int index_x = rAIdx + kId;
+      int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
+      As[(index_x % BK)][(index_y % BM)] = a[N * index_y + index_x];
+    }
+
+    __syncthreads();
+    for (int k = 0; k < BK; k++) {
+      // we cache A & B for the entire Wave tile
+      for (int iterWave = 0; iterWave < nbIterWaveN; iterWave++) {
+        for (int i = 0; i < TN; i++) {
+          int index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i;
+          B_row[iterWave * TN + i] = Bs[k][index];
+        }
+      }
+
+      for (int iterWave = 0; iterWave < nbIterWaveM; iterWave++) {
+        for (int i = 0; i < TM; i++) {
+          int index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i;
+          A_col[iterWave * TM + i] = As[k][index];
+        }
+      }
+
+      // we accumulate to C_regs
+      for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
+        for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
+          for (int yt = 0; yt < TM; yt++) {
+            for (int xt = 0; xt < TN; xt++) {
+              const int x = iterWaveN * TN + xt;
+              const int y = iterWaveM * TM + yt;
+              c_regs[y * TN * nbIterWaveN + x] += A_col[y] * B_row[x];
+            }
+          }
+        }
+      }
+    }
+  }
+
+  for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
+    for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
+      int xOut = blockIdx.x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave;
+      int yOut = blockIdx.y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave;
+      for (int yt = 0; yt < TM; yt++) {
+        for (int xt = 0; xt < TN; xt++) {
+          int indexC = N * (yOut + yt) + xOut + xt;
+          c[indexC] = beta * c[indexC] + alpha * c_regs[TN * nbIterWaveN * (iterWaveM * TM + yt) + (iterWaveN * TN + xt)];
+        }
+      }
+    }
+  }
+}
@@ -0,0 +1,172 @@
+typedef long unsigned int size_t;
+extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
+extern "C" __attribute__((device, const)) size_t __ockl_get_group_id(unsigned int);
+struct Dim3 { size_t x, y, z; };
+#define __shared__ __attribute__((shared, aligned(16)))
+__attribute__((device)) inline void __syncthreads() {
+  __builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup");
+  __builtin_amdgcn_s_barrier();
+  __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");
+}
+
+#define BLOCK_SIZE 256
+extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, BLOCK_SIZE)))
+kernel4_gmem_db(float *a, float *b, float *c)
+{
+  constexpr int N = 4096;
+  constexpr float alpha = 1.0;
+  constexpr float beta = 0.0;
+
+  const Dim3 blockIdx{ __ockl_get_group_id(0), __ockl_get_group_id(1), __ockl_get_group_id(2) };
+  const Dim3 threadIdx{ __ockl_get_local_id(0), __ockl_get_local_id(1), __ockl_get_local_id(2) };
+
+  // Block Tile size
+  constexpr int BN = 128;
+  constexpr int BM = 128;
+  // Number of Row or column we read per batch
+  constexpr int BK = 8;
+
+  // Thread Tile size
+  constexpr int TN = 4;
+  constexpr int TM = 4;
+
+  constexpr int nbWaves = BLOCK_SIZE / 32;
+  // Wave Tile size
+  constexpr int WN = 64;
+  constexpr int WM = BN * BM / nbWaves / WN;
+
+  // Number of wave on X & Y axis in the Block tile
+  constexpr int nbWaveX = BN / WN;
+  constexpr int nbWaveY = BM / WM;
+
+  const int waveIndex = threadIdx.x / 32;
+  const int waveIdx = waveIndex % nbWaveX;
+  const int waveIdy = waveIndex / nbWaveX;
+  const int indexInWave = threadIdx.x % 32;
+
+  // A wave is a block of 8x4 of the output matrix
+  constexpr int nbThreadXPerWave = 8;
+  constexpr int nbThreadYPerWave = 4;
+
+  // Thread coordinates in Wave
+  const int idxInWave = indexInWave % nbThreadXPerWave;
+  const int idyInWave = indexInWave / nbThreadXPerWave;
+
+  constexpr int nbIterWaveN = WN / (nbThreadXPerWave * TN);
+  constexpr int nbIterWaveM = WM / (nbThreadYPerWave * TM);
+
+  // Wave Sub-tile size
+  constexpr int SUBWN = WN / nbIterWaveN;
+  constexpr int SUBWM = WM / nbIterWaveM;
+
+  // Thread mapping to read BKxBN block from A
+  int rAIdx = threadIdx.x % BK;
+  int rAIdy = threadIdx.x / BK;
+  // Thread mapping to read BNxBK block from B
+  int rBIdx = threadIdx.x % BN;
+  int rBIdy = threadIdx.x / BN;
+
+  constexpr int strideReadB = BLOCK_SIZE / BN;
+  constexpr int strideReadA = BLOCK_SIZE / BK;
+  constexpr int nbReadsB = BN * BK / BLOCK_SIZE;
+  constexpr int nbReadsA = BM * BK / BLOCK_SIZE;
+
+  float A_col[nbIterWaveM * TM];
+  float B_row[nbIterWaveN * TN];
+
+  __shared__ float As[BK][BM];
+  __shared__ float Bs[BK][BN];
+
+  float c_regs[TM * nbIterWaveM * TN * nbIterWaveN] = {0.0f};
+
+  for (int i = 0; i < nbReadsB; i++) {
+    int index_x = BN * blockIdx.x + rBIdx;
+    int index_y = rBIdy + i * strideReadB;
+    Bs[index_y % BK][index_x % BN] = b[N * index_y + index_x];
+  }
+
+  for (int i = 0; i < nbReadsA; i++) {
+    int index_x = rAIdx;
+    int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
+    As[(index_x % BK)][(index_y % BM)] = a[N * index_y + index_x];
+  }
+
+  __syncthreads();
+  // Iteration over BK blocks.
+  for (int kId = 0; kId < N; kId += BK) {
+    float regA[nbReadsA];
+    float regB[nbReadsB];
+    if (kId < N - BK) {
+      // We populate the Shared Memory with Ks row and columns
+      for (int i = 0; i < nbReadsB; i++) {
+        int index_x = BN * blockIdx.x + rBIdx;
+        int index_y = rBIdy + i * strideReadB + kId + BK;
+        regB[i] = b[N * index_y + index_x];
+      }
+
+      for (int i = 0; i < nbReadsA; i++) {
+        int index_x = rAIdx + kId + BK;
+        int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
+        regA[i] = a[N * index_y + index_x];
+      }
+    }
+
+    for (int k = 0; k < BK; k++) {
+      // we cache A & B for the entire Wave tile
+      for (int iterWave = 0; iterWave < nbIterWaveN; iterWave++) {
+        for (int i = 0; i < TN; i++) {
+          int index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i;
+          B_row[iterWave * TN + i] = Bs[k][index];
+        }
+      }
+
+      for (int iterWave = 0; iterWave < nbIterWaveM; iterWave++) {
+        for (int i = 0; i < TM; i++) {
+          int index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i;
+          A_col[iterWave * TM + i] = As[k][index];
+        }
+      }
+
+      // we accumulate to C_regs
+      for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
+        for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
+          for (int yt = 0; yt < TM; yt++) {
+            for (int xt = 0; xt < TN; xt++) {
+              const int x = iterWaveN * TN + xt;
+              const int y = iterWaveM * TM + yt;
+              c_regs[y * TN * nbIterWaveN + x] += A_col[y] * B_row[x];
+            }
+          }
+        }
+      }
+    }
+    __syncthreads();
+    if (kId < N - BK) {
+      for (int i = 0; i < nbReadsB; i++) {
+        int index_x = BN * blockIdx.x + rBIdx;
+        int index_y = rBIdy + i * strideReadB + kId + BK;
+        Bs[index_y % BK][index_x % BN] = regB[i]; // row
+      }
+
+      for (int i = 0; i < nbReadsA; i++) {
+        int index_x = rAIdx + kId + BK;
+        int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
+        As[(index_x % BK)][(index_y % BM)] = regA[i];
+      }
+      __syncthreads();
+    }
+  }
+
+  for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
+    for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
+      int xOut = blockIdx.x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave;
+      int yOut = blockIdx.y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave;
+      for (int yt = 0; yt < TM; yt++) {
+        for (int xt = 0; xt < TN; xt++) {
+          int indexC = N * (yOut + yt) + xOut + xt;
+          c[indexC] = beta * c[indexC] + alpha * c_regs[TN * nbIterWaveN * (iterWaveM * TM + yt) + (iterWaveN * TN + xt)];
+        }
+      }
+    }
+  }
+}
@@ -0,0 +1,172 @@
+typedef long unsigned int size_t;
+extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
+extern "C" __attribute__((device, const)) size_t __ockl_get_group_id(unsigned int);
+struct Dim3 { size_t x, y, z; };
+#define __shared__ __attribute__((shared, aligned(16)))
+__attribute__((device)) inline void __syncthreads() {
+  __builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup");
+  __builtin_amdgcn_s_barrier();
+  __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");
+}
+
+#define BLOCK_SIZE 128
+extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, BLOCK_SIZE)))
+kernel5_lds_optim(float *a, float *b, float *c)
+{
+  constexpr int N = 4096;
+  constexpr float alpha = 1.0;
+  constexpr float beta = 0.0;
+
+  const Dim3 blockIdx{ __ockl_get_group_id(0), __ockl_get_group_id(1), __ockl_get_group_id(2) };
+  const Dim3 threadIdx{ __ockl_get_local_id(0), __ockl_get_local_id(1), __ockl_get_local_id(2) };
+
+  // Block Tile size
+  constexpr int BN = 128;
+  constexpr int BM = 128;
+  // Number of Row or column we read per batch
+  constexpr int BK = 8;
+
+  // Thread Tile size
+  constexpr int TN = 4;
+  constexpr int TM = 4;
+
+  constexpr int nbWaves = BLOCK_SIZE / 32;
+  // Wave Tile size
+  constexpr int WN = 128;
+  constexpr int WM = BN * BM / nbWaves / WN;
+
+  // Number of wave on X & Y axis in the Block tile
+  constexpr int nbWaveX = BN / WN;
+  constexpr int nbWaveY = BM / WM;
+
+  const int waveIndex = threadIdx.x / 32;
+  const int waveIdx = waveIndex % nbWaveX;
+  const int waveIdy = waveIndex / nbWaveX;
+  const int indexInWave = threadIdx.x % 32;
+
+  // A wave is a block of 8x4 of the output matrix
+  constexpr int nbThreadXPerWave = 8;
+  constexpr int nbThreadYPerWave = 4;
+
+  // Thread coordinates in Wave
+  const int idxInWave = indexInWave % nbThreadXPerWave;
+  const int idyInWave = indexInWave / nbThreadXPerWave;
+
+  constexpr int nbIterWaveN = WN / (nbThreadXPerWave * TN);
+  constexpr int nbIterWaveM = WM / (nbThreadYPerWave * TM);
+
+  // Wave Sub-tile size
+  constexpr int SUBWN = WN / nbIterWaveN;
+  constexpr int SUBWM = WM / nbIterWaveM;
+
+  // Thread mapping to read BKxBN block from A
+  int rAIdx = threadIdx.x % BK;
+  int rAIdy = threadIdx.x / BK;
+  // Thread mapping to read BNxBK block from B
+  int rBIdx = threadIdx.x % BN;
+  int rBIdy = threadIdx.x / BN;
+
+  constexpr int strideReadB = BLOCK_SIZE / BN;
+  constexpr int strideReadA = BLOCK_SIZE / BK;
+  constexpr int nbReadsB = BN * BK / BLOCK_SIZE;
+  constexpr int nbReadsA = BM * BK / BLOCK_SIZE;
+
+  float A_col[nbIterWaveM * TM];
+  float B_row[nbIterWaveN * TN];
+
+  __shared__ float As[BK][BM+4]; // 4 padding to avoid bank conflicts
+  __shared__ float Bs[BK][BN];
+
+  float c_regs[TM * nbIterWaveM * TN * nbIterWaveN] = {0.0f};
+
+  // initial copy into shared memory
+  for (int i = 0; i < nbReadsB; i++) {
+    int index_x = BN * blockIdx.x + rBIdx;
+    int index_y = rBIdy + i * strideReadB;
+    Bs[index_y % BK][index_x % BN] = b[N * index_y + index_x];
+  }
+  for (int i = 0; i < nbReadsA; i++) {
+    int index_x = rAIdx;
+    int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
+    As[(index_x % BK)][(index_y % BM)] = a[N * index_y + index_x];
+  }
+
+  __syncthreads();
+  // Iteration over BK blocks.
+  for (int kId = 0; kId < N; kId += BK) {
+    float regA[nbReadsA];
+    float regB[nbReadsB];
+    if (kId < N - BK) {
+      // We populate the Shared Memory with Ks row and columns
+      for (int i = 0; i < nbReadsB; i++) {
+        int index_x = BN * blockIdx.x + rBIdx;
+        int index_y = rBIdy + i * strideReadB + kId + BK;
+        regB[i] = b[N * index_y + index_x];
+      }
+
+      for (int i = 0; i < nbReadsA; i++) {
+        int index_x = rAIdx + kId + BK;
+        int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
+        regA[i] = a[N * index_y + index_x];
+      }
+    }
+
+    for (int k = 0; k < BK; k++) {
+      // we cache A & B for the entire Wave tile
+      for (int iterWave = 0; iterWave < nbIterWaveN; iterWave++) {
+        for (int i = 0; i < TN; i++) {
+          int index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i;
+          B_row[iterWave * TN + i] = Bs[k][index];
+        }
+      }
+
+      for (int iterWave = 0; iterWave < nbIterWaveM; iterWave++) {
+        for (int i = 0; i < TM; i++) {
+          int index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i;
+          A_col[iterWave * TM + i] = As[k][index];
+        }
+      }
+
+      // we accumulate to C_regs
+      for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
+        for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
+          for (int yt = 0; yt < TM; yt++) {
+            for (int xt = 0; xt < TN; xt++) {
+              const int x = iterWaveN * TN + xt;
+              const int y = iterWaveM * TM + yt;
+              c_regs[y * TN * nbIterWaveN + x] += A_col[y] * B_row[x];
+            }
+          }
+        }
+      }
+    }
+    __syncthreads();
+    if (kId < N - BK) {
+      for (int i = 0; i < nbReadsB; i++) {
+        int index_x = BN * blockIdx.x + rBIdx;
+        int index_y = rBIdy + i * strideReadB + kId + BK;
+        Bs[index_y % BK][index_x % BN] = regB[i]; // row
+      }
+
+      for (int i = 0; i < nbReadsA; i++) {
+        int index_x = rAIdx + kId + BK;
+        int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
+        As[(index_x % BK)][(index_y % BM)] = regA[i];
+      }
+      __syncthreads();
+    }
+  }
+
+  for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
+    for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
+      int xOut = blockIdx.x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave;
+      int yOut = blockIdx.y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave;
+      for (int yt = 0; yt < TM; yt++) {
+        for (int xt = 0; xt < TN; xt++) {
+          int indexC = N * (yOut + yt) + xOut + xt;
+          c[indexC] = beta * c[indexC] + alpha * c_regs[TN * nbIterWaveN * (iterWaveM * TM + yt) + (iterWaveN * TN + xt)];
+        }
+      }
+    }
+  }
+}
@@ -1,6 +1,5 @@
 	.text
 	.amdgcn_target "amdgcn-amd-amdhsa--gfx1100"
-	;.amdhsa_code_object_version 5
 	.protected	kernel                  ; -- Begin function kernel
 	.globl	kernel
 	.p2align	8
@@ -9,7 +8,7 @@ kernel:                                 ; @kernel
 ; %bb.0:                                ; %.preheader193

 	;; Init code for matrix A and B buffer Loads  - START
-	s_load_b128 s[20:23], s[0:1], 0x8 ; Matrix A and B
+	s_load_b128 s[20:23], s[0:1], 0x0 ; Matrix A and B
 	s_waitcnt lgkmcnt(0)

 	; Matrix B offsets:
@@ -76,14 +75,12 @@ kernel:                                 ; @kernel


 	s_clause 0x1
-	;s_load_b128 s[4:7], s[0:1], 0x18 ; N, alpha, beta, ???
-	s_load_b128 s[8:11], s[0:1], 0x8 ; Matrix A and B
-
-	s_mov_b32 s4, 4096 ; hardcode 4096
-	s_mov_b32 s5, 0x3f800000 ; alpha
-	s_mov_b32 s6, 0    ; beta
-	s_mov_b32 s7, 0
-
+	; s_load_b128 s[4:7], s[0:1], 0x18
+	; N=4096, alpha=1.0, beta=0.0
+	s_mov_b32 s4, 4096
+	s_mov_b32 s5, 0x3F800000
+	s_mov_b32 s6, 0
+	s_load_b128 s[8:11], s[0:1], 0x0
 	s_lshl_b32 s2, s14, 7
 	v_lshrrev_b32_e32 v4, 3, v0
 	v_or_b32_e32 v1, s2, v0
@@ -93,7 +90,7 @@ kernel:                                 ; @kernel
 	v_or_b32_e32 v22, s3, v4
 	v_ashrrev_i32_e32 v2, 31, v1
 	s_lshr_b32 s12, s12, 25
-	s_load_b64 s[0:1], s[0:1], 0  ; Matrix C
+	s_load_b64 s[0:1], s[0:1], 0x10
 	v_lshlrev_b32_e32 v135, 2, v118
 	s_delay_alu instid0(VALU_DEP_2) | instskip(SKIP_3) | instid1(VALU_DEP_3)
 	v_lshlrev_b64 v[5:6], 2, v[1:2]
@@ -463,7 +460,7 @@ kernel:                                 ; @kernel

 	v_mov_b32_e32 v5, 0
 	v_mov_b32_e32 v3, 0
-	s_add_i32 s7, s4, -1
+	s_add_i32 s7, s4, -8
 	s_add_u32 s8, s8, 32
 	s_addc_u32 s9, s9, 0
 	s_mov_b32 s12, 0
@@ -2398,18 +2395,9 @@ amdhsa.kernels:
        .offset:         16
        .size:           8
        .value_kind:     global_buffer
-      - .offset:         24
-        .size:           4
-        .value_kind:     by_value
-      - .offset:         28
-        .size:           4
-        .value_kind:     by_value
-      - .offset:         32
-        .size:           4
-        .value_kind:     by_value
    .group_segment_fixed_size: 8320
    .kernarg_segment_align: 8
-    .kernarg_segment_size: 36
+    .kernarg_segment_size: 24
    .language:       OpenCL C
    .language_version:
      - 2
@@ -0,0 +1,355 @@
+from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, graph_rewrite, AxisType, PatternMatcher, UPat
+from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
+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.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)])
+
+N = 4096
+run_count = 5
+
+BN = 128
+BM = 128
+BK = 8
+
+TN = 4
+TM = 4
+
+# NOTE: this is from testgrad
+# change reduceop axes and input ShapeTrackers, view gets replaced with a reshape.
+# src->r->view  -->   src->view->r
+def swizzle_reduceop(src:UOp, r:UOp, view:UOp):
+  if r.tag is not None: return None
+  # confirm the input is in order
+  # TODO: replace this with a UOp that allows for nothing else then remove this
+  permute = tuple(i for i in range(len(src.shape)) if i not in r.axis_arg)+r.axis_arg
+  assert permute == tuple(range(len(permute))), f"reduce axis must already be in order, {permute} isn't"
+
+  # append the reduce shape to each of the views
+  prshape = prod(rshape:=src.shape[-len(r.axis_arg):])
+  rstrides = strides_for_shape(rshape)
+  nv = [View.create(v.shape+rshape, tuple(x*prshape for x in v.strides)+rstrides, v.offset*prshape,
+                    v.mask+tuple((0,s) for s in rshape) if v.mask is not None else None) for v in unwrap(view.st).views]
+
+  # no reshape required with shrinking REDUCE_AXIS
+  return UOp(Ops.REDUCE_AXIS, r.dtype, (src.view(ShapeTracker(tuple(nv))),),
+             (r.arg[0], tuple(range(len(view.shape), len(view.shape) + len(r.axis_arg)))))
+
+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()
+  with Context(RANGEIFY=1):
+    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 = 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))
+
+def hl_spec_kernel3():
+  nbIterWaveM = 2
+  nbIterWaveN = 2
+
+  # define buffers
+  # TODO: remove these views once the defines have a shape
+  a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1).view(ShapeTracker.from_shape((N,N)))
+  b = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2).view(ShapeTracker.from_shape((N,N))).permute((1,0))
+  c = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0).view(ShapeTracker.from_shape((N,N)))
+  As = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BM, AddrSpace.LOCAL), arg=0).view(ShapeTracker.from_shape((BK, BM))).permute((1,0))
+  Bs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BN, AddrSpace.LOCAL), arg=1).view(ShapeTracker.from_shape((BK, BN))).permute((1,0))
+  A_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveM * TM, AddrSpace.REG), arg=0).view(ShapeTracker.from_shape((nbIterWaveM * TM,)))
+  B_row = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveN * TN, AddrSpace.REG), arg=1).view(ShapeTracker.from_shape((nbIterWaveN * TN,)))
+
+  # shape buffers. TODO: permutes
+  full_shape = (N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, N//BK, BK)
+  a = a.reshape((N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, 1, 1, 1, 1, N//BK, BK)).expand(full_shape)
+  b = b.reshape((1, 1, 1, 1, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, N//BK, BK)).expand(full_shape)
+  c = c.reshape((N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, 1, 1))
+  As = As.reshape((1, nbIterWaveM, BM//(nbIterWaveM * TM), TM, 1, 1, 1, 1, 1, BK)).expand(full_shape)
+  Bs = Bs.reshape((1, 1, 1, 1, 1, nbIterWaveN, BN//(nbIterWaveN * TN), TN, 1, BK)).expand(full_shape)
+  A_col = A_col.reshape((1, nbIterWaveM, 1, TM, 1, 1, 1, 1, 1, 1)).expand(full_shape)
+  B_row = B_row.reshape((1, 1, 1, 1, 1, nbIterWaveN, 1, TN, 1, 1)).expand(full_shape)
+
+  #                     U1   L2 L3 L4 L5   U6 U7      U9   L10 L11 L12 L13   U14 U15      U17  U18  U19
+  expanded_shape = (32, 2,   2, 2, 2, 2,   2, 2,  32, 2,   2,  2,  2,  2,    2,  2,  512, 2,   2,   2)
+  assert len(expanded_shape) == 20
+  permute_a = list(range(len(expanded_shape)))
+  permute_b = permute_a[:]
+
+  # this makes all the global loads match
+  # this can also be more simply done by rebinding the RANGEs
+  # but sadly, rebinding the RANGEs doesn't work to change the order of the local axes
+  permute_a[17:20] = [11,12,13]
+  permute_a[11:14] = [17,18,19]
+  permute_a[7], permute_a[10] = permute_a[10], permute_a[7]
+  permute_a[2:7] = [3,4,5,6,2]
+
+  permute_b[2:16] = [19,9,10,11,17,18,8,2,12,13,14,15,3,4]
+  permute_b[17:20] = [5,6,7]
+
+  a_permute   = a.reshape(expanded_shape).permute(tuple(permute_a)).reshape(full_shape)
+  As_permute = As.reshape(expanded_shape).permute(tuple(permute_a)).reshape(full_shape)
+
+  b_permute   = b.reshape(expanded_shape).permute(tuple(permute_b)).reshape(full_shape)
+  Bs_permute = Bs.reshape(expanded_shape).permute(tuple(permute_b)).reshape(full_shape)
+
+  #out = (a.load() * b.load()).r(Ops.ADD, (8, 9))
+  out = (As.load(As_permute.store(a_permute.load())) * Bs.load(Bs_permute.store(b_permute.load()))).r(Ops.ADD, (8, 9))
+  #out = (A_col.load(A_col.store(As.load(As.store(a.load())))) * B_row.load(B_row.store(Bs.load(Bs.store(b.load()))))).r(Ops.ADD, (8, 9))
+
+  axis_types = (
+    AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.UPCAST,
+    AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.UPCAST,
+    AxisType.REDUCE, AxisType.REDUCE)
+
+  sink = c.store(out).sink(arg=KernelInfo(name="tg_"+to_colored(full_shape, axis_types), axis_types=axis_types))
+  sink = graph_rewrite(sink, merge_views)
+  return sink
+
+def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
+  BLOCK_SIZE = 128 if kernel5 else 256
+
+  nbWaves = BLOCK_SIZE // 32
+  WN = 128 if kernel5 else 64
+  WM = BN * BM // nbWaves // WN
+
+  nbWaveX = BN // WN
+  nbWaveY = BM // WM
+
+  threadIdx_x = UOp(Ops.SPECIAL, dtypes.int, arg=("lidx0", BLOCK_SIZE))
+  waveIndex = threadIdx_x // 32
+  waveIdx = waveIndex % nbWaveX
+  waveIdy = waveIndex // nbWaveX
+  indexInWave = threadIdx_x % 32
+
+  nbThreadXPerWave = 8
+  nbThreadYPerWave = 4
+
+  idxInWave = indexInWave % nbThreadXPerWave
+  idyInWave = indexInWave // nbThreadXPerWave
+
+  nbIterWaveN = WN // (nbThreadXPerWave * TN)
+  nbIterWaveM = WM // (nbThreadYPerWave * TM)
+
+  SUBWN = WN // nbIterWaveN
+  SUBWM = WM // nbIterWaveM
+
+  # Thread mapping to read BKxBN block from A
+  rAIdx = threadIdx_x % BK
+  rAIdy = threadIdx_x // BK
+  # Thread mapping to read BNxBK block from B
+  rBIdx = threadIdx_x % BN
+  rBIdy = threadIdx_x // BN
+
+  strideReadB = BLOCK_SIZE // BN
+  strideReadA = BLOCK_SIZE // BK
+  nbReadsB = BN * BK // BLOCK_SIZE
+  nbReadsA = BM * BK // BLOCK_SIZE
+
+  blockIdx_x = UOp(Ops.SPECIAL, dtypes.int, arg=("gidx0", N//BN))
+  blockIdx_y = UOp(Ops.SPECIAL, dtypes.int, arg=("gidx1", N//BM))
+
+  a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1)
+  b = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2)
+  c = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0)
+
+  A_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveM * TM, AddrSpace.REG), arg=0)
+  B_row = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveN * TN, AddrSpace.REG), arg=1)
+
+  BM_As_stride = (BM+4) if kernel5 else BM
+  As = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BM_As_stride, AddrSpace.LOCAL), arg=0)
+  Bs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BN, AddrSpace.LOCAL), arg=1)
+
+  c_regs = UOp(Ops.DEFINE_REG, dtypes.float.ptr(TM * nbIterWaveM * TN * nbIterWaveN), arg=2)
+
+  i = UOp.range(c_regs.dtype.size, 16)
+  init_store = c_regs[i].store(UOp.const(dtypes.float, 0.0), i)
+
+  if kernel4:
+    regA = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbReadsA, AddrSpace.REG), arg=3)
+    regB = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbReadsB, AddrSpace.REG), arg=4)
+
+    # initial load from globals into locals (0)
+    kId = 0
+
+    # load from globals into locals
+    i = UOp.range(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)
+    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 = kId_range*BK
+
+    barrier = UOp.barrier(As_store, Bs_store)
+
+    # load from globals into registers (next round)
+    i = UOp.range(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)
+    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)
+
+      # load from locals into registers
+      iterWave = UOp.range(nbIterWaveN, first_range+1)
+      i = UOp.range(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)
+      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)
+      x = iterWaveN * TN + xt
+      y = iterWaveM * TM + yt
+      c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
+      # sketchy, this should end the kId_range but it doesn't
+      sink = c_regs_idx.store(c_regs_idx.load(init_store) + A_col[y].load(A_col_store) * B_row[x].load(B_row_store),
+                              iterWaveM, iterWaveN, yt, xt, k)
+      return sink
+
+    # TODO: kId_range should endrange after a barrier
+    sink = inner_loop(5, (barrier, regB_store, regA_store)).barrier()
+
+    # load from registers into locals
+    i = UOp.range(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)
+    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)
+
+    # final iteration without the copy
+    sink = inner_loop(16, (UOp.barrier(Bs_store, As_store),))
+  else:
+    kId_range = UOp.range(N//BK, 0)
+    kId = kId_range*BK
+
+    # load from globals into locals
+    i = UOp.range(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)
+    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)
+
+    # load from locals into registers
+    iterWave = UOp.range(nbIterWaveN, 4)
+    i = UOp.range(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)
+    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)
+    x = iterWaveN * TN + xt
+    y = iterWaveM * TM + yt
+    c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
+    sink = c_regs_idx.store(c_regs_idx.load(init_store) + A_col[y].load(A_col_store) * B_row[x].load(B_row_store),
+                            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)
+  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
+  sink = c[indexC].store(c_regs[TN * nbIterWaveN * (iterWaveM * TM + yt) + (iterWaveN * TN + xt)].load(sink),
+                         iterWaveM, iterWaveN, yt, xt)
+
+  return sink.sink(arg=KernelInfo(name="tinygemm"))
+
+if __name__ == "__main__":
+  HL = getenv("HL")
+  if HL == 3: hprg = rangeify_kernel3()
+  elif HL == 2: hprg = top_spec_kernel3()
+  elif HL == 1: hprg = hl_spec_kernel3()
+  else: hprg = hand_spec_kernel3()
+  if HL == 3:
+    with Context(RANGEIFY=1, BLOCK_REORDER=0):
+      prg = get_program(hprg, Device.default.renderer)
+  else:
+    prg = get_program(hprg, Device.default.renderer)
+  print(prg.src)
+  if getenv("SRC"): exit(0)
+  hrunner = CompiledRunner(prg)
+
+  a = Tensor.randn(N, N).realize()
+  b = Tensor.randn(N, N).realize()
+  hc = Tensor.zeros(N, N).contiguous().realize()
+
+  GlobalCounters.reset()
+  with Context(DEBUG=2):
+    for _ in range(run_count): tc = (a@b).realize()
+
+  GlobalCounters.reset()
+  buffers = [hc.uop.buffer, a.uop.buffer, b.uop.buffer]
+  ei = ExecItem(hrunner, buffers)
+  with Context(DEBUG=2):
+    for _ in range(run_count): ei.run(wait=True)
+  err = (hc-tc).square().mean().item()
+  print(f"hrunner {err}")
+  if err > 1e-06: raise RuntimeError("matmul is wrong!")
@@ -0,0 +1,90 @@
+import numpy as np
+import halide as hl
+from tinygrad.helpers import Timing, getenv
+
+# HL_DEBUG_CODEGEN=1
+N = getenv("N", 1024)
+
+def gemm_pipeline(gpu=False):
+  # ---------------- Vars & Parameters ----------------
+  i, j = hl.Var("i"), hl.Var("j")  # output tile coordinates
+
+  A = hl.InputBuffer(hl.Float(32), 2)  # [M, K]
+  B = hl.InputBuffer(hl.Float(32), 2)  # [K, N]
+
+  A.dim(0).set_bounds(0, N)
+  A.dim(1).set_bounds(0, N)
+  B.dim(0).set_bounds(0, N)
+  B.dim(1).set_bounds(0, N)
+
+  # ---------------- Definition ----------------
+
+  k = hl.RDom([(0, N)])
+
+  partial = hl.Func("partial")
+  partial[i, j] = 0.0
+  partial[i, j] += A[i, k] * B[k, j]
+
+  C = hl.Func("C")
+  C[i, j] = partial[i, j]
+
+  if not gpu:
+    # ---------------- Schedule ----------------
+    VEC = 16
+    TILE_I = 64
+    TILE_J = 64
+
+    io, jo, ii, ji = hl.Var("io"), hl.Var("jo"), hl.Var("ii"), hl.Var("ji")
+    C.update().tile(i, j, io, jo, ii, ji, TILE_I, TILE_J).fuse(io, jo, io).parallel(io).vectorize(ji, VEC)
+  else:
+    # ---------------- Schedule ----------------
+    GRP_I    = 8     # output tile size
+    GRP_J    = 16
+
+    #partial.store_in(hl.MemoryType.Register)
+    #partial.update().unroll(k, 4)
+
+    io, jo, ii, ji = hl.Var(), hl.Var(), hl.Var(), hl.Var()
+    C.gpu_tile(i, j, io, jo, ii, ji, GRP_I, GRP_J, hl.TailStrategy.RoundUp)
+
+  return C, A, B
+
+if __name__ == "__main__":
+  pipe, A, B = gemm_pipeline(gpu=True)
+
+  # NOTE: meteal does nothing
+  target = hl.get_host_target().with_feature(hl.TargetFeature.Metal)
+
+  a_np = np.random.randn(N, N).astype(np.float32)
+  b_np = np.random.randn(N, N).astype(np.float32)
+
+  # reverse order is correct!
+  a_hal = hl.Buffer(b_np)
+  b_hal = hl.Buffer(a_np)
+  A.set(a_hal)
+  B.set(b_hal)
+
+  pipe.compile_to_lowered_stmt("/tmp/my_function.html", [A, B], hl.StmtOutputFormat.HTML, target=target)
+  #exit(0)
+
+  c_hal = hl.Buffer(hl.Float(32), [N,N])
+  with Timing("halide gemm "):
+    pipe.realize(c_hal, target)
+    c_hal.copy_to_host()
+    c_out = np.array(c_hal)
+  print(c_out)
+
+  # tinygrad gets 60 ms with no BEAM, 20 ms with BEAM on CPU
+  with Timing("halide gemm "):
+    pipe.realize(c_hal, target)
+    c_hal.copy_to_host()
+
+  # Check correctness
+  with Timing("numpy gemm "):
+    ref = a_np @ b_np
+  max_err = np.abs(ref - c_out).max()
+  print("Max absolute error:", max_err)
+  assert max_err < 1e-4, "GEMM result incorrect!"
+
+  print("Pipeline ran on", target)
+  print("Success - GEMM Halide-Python output matches NumPy.")
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import numpy as np
-from tinygrad.runtime.ops_gpu import CLProgram, CLCompiler
+from tinygrad.runtime.ops_cl 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["GPU"]
+device = Device["CL"]

 # 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("GPU", 8, dtypes.float32).allocate()
-b = Buffer("GPU", 0x10, dtypes.float16).allocate()
-c = Buffer("GPU", 8*0x10, dtypes.float16).allocate()
+a = Buffer("CL", 8, dtypes.float32).allocate()
+b = Buffer("CL", 0x10, dtypes.float16).allocate()
+c = Buffer("CL", 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)
@@ -2,11 +2,12 @@ import numpy as np, os
 from tinygrad.helpers import getenv, flat_mv
 from tinygrad import dtypes
 from typing import Optional, List, Tuple, cast, Dict, Final, DefaultDict, Self
+from tinygrad.engine.realize import get_program

 # for copied uops
-from tinygrad.opt.kernel import Kernel, KernelOptError
+from tinygrad.codegen.opt.kernel import Kernel, KernelOptError
 from tinygrad.uop.ops import UOp, Ops, BinaryOps, UnaryOps, TernaryOps, KernelInfo
-from tinygrad.opt.search import Opt, OptOps
+from tinygrad.codegen.opt.search import Opt, OptOps
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.dtype import PtrDType, DType, DTYPES_DICT
 from tinygrad.shape.shapetracker import ShapeTracker
@@ -55,7 +56,7 @@ def randoms():
 def ast_to_cuda_prog(compiler, ast, opts):
  k = Kernel(ast)
  k.apply_opts(opts)
-  p = k.to_program()
+  p = get_program(k.ast, k.opts, k.applied_opts)
  return CUDAProgram(device, p.function_name, compiler.compile(p.src))

 if __name__ == "__main__":
@@ -74,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 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 CUDA_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.opt.kernel import OptOps
+from tinygrad.codegen.opt 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
@@ -1,7 +1,7 @@
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.helpers import getenv, DEBUG
-from tinygrad.opt.kernel import Kernel, Opt, OptOps
-from tinygrad.engine.realize import CompiledRunner, ExecItem
+from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
 from dataclasses import replace

 N = 4096
@@ -29,7 +29,7 @@ if __name__ == "__main__":
      Opt(op=OptOps.LOCAL, axis=0, amt=2),
    ]
  k.apply_opts(opts)
-  prg = k.to_program()
+  prg = get_program(k.ast, k.opts, k.applied_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 PTX=1 python3 extra/gemm/triton_nv_matmul.py
+# CUDA=1 CUDA_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
@@ -37,10 +37,10 @@ B = Tensor.rand(K, N, device="CPU")
 C = (A.reshape(M, 1, K) * B.permute(1,0).reshape(1, N, K)).sum(axis=2)

 sched = C.schedule()
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Kernel
 from tinygrad.device import CompilerOptions
 lin = Kernel(sched[-1].ast, CompilerOptions(has_local=False, supports_float4=False))
-lin.linearize()
+lin.to_program()
 from tinygrad.runtime.ops_cpu import renderer
 src = renderer("mmult", lin.uops)
 print(src)
@@ -0,0 +1,122 @@
+#!/usr/bin/env python3
+
+from tinygrad.runtime.support.system import System
+import argparse, glob, os, re, time, subprocess, sys
+
+def scan_devs_based_on_lock(prefix:str, args) -> list[str]:
+  target_dev = args.pci_bus if 'pci_bus' in args.__dir__() else ""
+
+  devs = []
+  for dev in glob.glob(f'/tmp/{prefix}_*.lock'):
+    dev_id = dev[8:-5]
+    if os.path.exists(f"/sys/bus/pci/devices/{dev_id}") and dev_id.startswith(target_dev): devs.append(dev_id)
+  return devs
+
+def _do_reset_device(pci_bus): System.pci_reset(pci_bus)
+def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")
+
+def cmd_remove_module(args):
+  modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia", "ast"] if args.backend == "nv" else ["amdgpu"]
+  to_unload = [m for m in modules if _is_module_loaded(m)]
+  if not to_unload: print("Kernel modules are not loaded")
+  else:
+    print("Removing kernel modules:", ", ".join(to_unload))
+    try: subprocess.run(["sudo", "modprobe", "-r", *to_unload], check=True)
+    except subprocess.CalledProcessError as e:
+      print("Failed to unload all modules — they may be in use.", file=sys.stderr)
+      sys.exit(e.returncode)
+
+def cmd_insert_module(args):
+  cmd_remove_module(args)
+  cmd_reset_devices(args)
+
+  module = "nvidia" if args.backend == "nv" else "amdgpu"
+  if _is_module_loaded(module):
+    print(f"{module} kernel module already loaded")
+    return
+
+  print(f"Inserting kernel module: {module}")
+  if args.backend == "nv":
+    subprocess.run(["nvidia-smi"], check=True)
+  elif args.backend == "amd":
+    subprocess.run(["sudo", "modprobe", "amdgpu"], check=True)
+
+def cmd_reset_devices(args):
+  devs = scan_devs_based_on_lock({"amd":"am", "nv":"nv"}[args.backend], args)
+
+  for dev in devs:
+    print(f"Resetting device {dev}")
+    if args.backend != "amd": _do_reset_device(dev)
+    time.sleep(0.2)
+
+def cmd_show_pids(args):
+  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
+
+  for dev in devs:
+    try:
+      pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
+      print(f"{dev}: {pid}")
+    except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")
+
+def cmd_kill_pids(args):
+  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
+
+  for dev in devs:
+    try:
+      pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
+      print(f"{dev}: {pid}")
+    except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")
+
+def cmd_kill_pids(args):
+  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
+
+  for dev in devs:
+    for i in range(128):
+      if i > 0: time.sleep(0.2)
+
+      try:
+        try: pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
+        except subprocess.CalledProcessError: break
+
+        print(f"Killing process {pid} (which uses {dev})")
+        subprocess.run(['sudo', 'kill', '-9', pid], check=True)
+      except subprocess.CalledProcessError as e:
+        print(f"Failed to kill process for device {dev}: {e}", file=sys.stderr)
+
+def add_common_commands(parent_subparsers):
+  p_insmod = parent_subparsers.add_parser("insmod", help="Insert a kernel module")
+  p_insmod.set_defaults(func=cmd_insert_module)
+
+  p_rmmod = parent_subparsers.add_parser("rmmod", help="Remove a kernel module")
+  p_rmmod.set_defaults(func=cmd_remove_module)
+
+  p_reset = parent_subparsers.add_parser("reset", help="Reset a device")
+  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device to reset")
+  p_reset.set_defaults(func=cmd_reset_devices)
+
+  p_reset = parent_subparsers.add_parser("pids", help="Show pids of processes using the device")
+  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device")
+  p_reset.set_defaults(func=cmd_show_pids)
+
+  p_reset = parent_subparsers.add_parser("kill_pids", help="Kill pids of processes using the device")
+  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device")
+  p_reset.set_defaults(func=cmd_kill_pids)
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  backend_subparsers = parser.add_subparsers(dest="backend", required=True, metavar="{nv,amd}", help="Hardware backend to target")
+
+  nv_parser = backend_subparsers.add_parser("nv", help="NVIDIA GPUs")
+  nv_commands = nv_parser.add_subparsers(dest="command", required=True)
+  add_common_commands(nv_commands)
+
+  amd_parser = backend_subparsers.add_parser("amd", help="AMD GPUs")
+  amd_commands = amd_parser.add_subparsers(dest="command", required=True)
+  add_common_commands(amd_commands)
+
+  args = parser.parse_args()
+  if args.command is None:
+    parser.print_help(sys.stderr)
+    sys.exit(1)
+
+  args.func(args)
@@ -4,9 +4,9 @@ To add a new test, define a `TestSpec`-based class in a file in the `tests/` fol

 You can choose which tests to load from which file:
 ```bash
-RUN_FILES="hcq,allocator" python3 extra/hcqfuzz/fuzzer.py
+PYTHONPATH=. RUN_FILES="hcq,allocator" python3 extra/hcqfuzz/fuzzer.py
 ```
 Or skip tests from any file:
 ```bash
-SKIP_FILES="allocator" python3 extra/hcqfuzz/fuzzer.py
+PYTHONPATH=. SKIP_FILES="allocator" python3 extra/hcqfuzz/fuzzer.py
 ```
@@ -8,7 +8,6 @@ bert_train_params = {
  "BS": 96,
  "EVAL_BS": 96,
  "FUSE_ARANGE": 1,
-  "FUSE_ARANGE_UINT": 0,
  "BASEDIR": "/raid/datasets/wiki",
 }

@@ -5,7 +5,7 @@ start = time.perf_counter()
 # *** ioctl lib ***
 libc = ctypes.CDLL(ctypes.util.find_library("c"))
 # platform.processor calls `uname -p` which can return `unknown` on some systems
-processor = os.getenv("IOCTL_PROCESSOR") or platform.processor()
+processor = os.getenv("IOCTL_PROCESSOR") or platform.processor() or platform.machine()
 IOCTL_SYSCALL = {"aarch64": 0x1d, "x86_64":16}[processor]

 def get_struct(argp, stype):
@@ -0,0 +1,61 @@
+# HuggingFace ONNX
+
+Tool for discovering, downloading, and validating ONNX models from HuggingFace.
+
+## Extra Dependencies
+
+```bash
+pip install huggingface_hub pyyaml requests onnx onnxruntime numpy
+```
+
+## Huggingface Manager (discovering and downloading)
+
+The `huggingface_manager.py` script discovers top ONNX models from HuggingFace, collects metadata, and optionally downloads them.
+
+```bash
+# Download top 50 models sorted by downloads
+python huggingface_manager.py --limit 50 --download
+
+# Just collect metadata (no download)
+python huggingface_manager.py --limit 100
+
+# Sort by likes instead of downloads
+python huggingface_manager.py --limit 20 --sort likes --download
+
+# Custom output file
+python huggingface_manager.py --limit 10 --output my_models.yaml
+```
+
+### Output Format
+
+The tool generates a YAML file with the following structure:
+
+```yaml
+repositories:
+  "model-name":
+    url: "https://huggingface.co/model-name"
+    download_path: "/path/to/models/..."  # when --download used
+    files:
+      - file: "model.onnx"
+        size: "90.91MB"
+total_size: "2.45GB"
+created_at: "2024-01-15T10:30:00Z"
+```
+
+## Run Models (validation)
+
+The `run_models.py` script validates ONNX models against ONNX Runtime for correctness.
+
+```bash
+# Validate models from a YAML configuration file
+python run_models.py --validate huggingface_repos.yaml
+
+# Debug specific repository (downloads and validates all ONNX models)
+python run_models.py --debug sentence-transformers/all-MiniLM-L6-v2
+
+# Debug specific model file
+python run_models.py --debug sentence-transformers/all-MiniLM-L6-v2/onnx/model.onnx
+
+# Debug with model truncation for debugging and validating intermediate results
+DEBUGONNX=1 python run_models.py --debug sentence-transformers/all-MiniLM-L6-v2/onnx/model.onnx --truncate 10
+```
@@ -1,85 +0,0 @@
-import yaml, time, requests, argparse
-from pathlib import Path
-from huggingface_hub import list_models, HfApi
-from tinygrad.helpers import tqdm
-
-HUGGINGFACE_URL = "https://huggingface.co"
-SKIPPED_FILES = [
-  "fp16", "int8", "uint8", "quantized",      # numerical accuracy issues
-  "avx2", "arm64", "avx512", "avx512_vnni",  # numerical accuracy issues
-  "q4", "q4f16", "bnb4",                     # unimplemented quantization
-  "model_O4",                                # requires non cpu ort runner and MemcpyFromHost op
-  "merged",                                  # TODO implement attribute with graph type and Loop op
-]
-SKIPPED_REPO_PATHS = [
-  # Invalid model-index
-  "AdamCodd/vit-base-nsfw-detector",
-  # TODO: implement attribute with graph type and Loop op
-  "minishlab/potion-base-8M", "minishlab/M2V_base_output", "minishlab/potion-retrieval-32M",
-  # TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, GroupQueryAttention
-  "HuggingFaceTB/SmolLM2-360M-Instruct",
-  # TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, RotaryEmbedding, MultiHeadAttention
-  "HuggingFaceTB/SmolLM2-1.7B-Instruct",
-  # TODO: implmement RandomNormalLike
-  "stabilityai/stable-diffusion-xl-base-1.0", "stabilityai/sdxl-turbo", 'SimianLuo/LCM_Dreamshaper_v7',
-  # TODO: implement NonZero
-  "mangoapps/fb_zeroshot_mnli_onnx",
-  # TODO huge Concat in here with 1024 (1, 3, 32, 32) Tensors, and maybe a MOD bug with const folding
-  "briaai/RMBG-2.0",
-]
-
-def get_top_repos(n: int, sort: str) -> list[str]: # list["FacebookAI/xlm-roberta-large", ...]
-  print(f"** Getting top {n} models sorted by {sort} **")
-  repos = []
-  i = 0
-  for model in list_models(filter="onnx", sort=sort):
-    if model.id in SKIPPED_REPO_PATHS: continue
-    print(f"{i+1}/{n}: {model.id} ({getattr(model, sort)})")
-    repos.append(model.id)
-    i += 1
-    if i == n: break
-  return repos
-
-def get_metadata(repos:list[str]) -> dict:
-  api = HfApi()
-  repos_metadata = {"repositories": {}}
-  total_size = 0
-
-  # TODO: speed head requests up with async?
-  for repo in tqdm(repos, desc="Getting metadata"):
-    files_metadata = []
-    model_info = api.model_info(repo)
-
-    for file in model_info.siblings:
-      filename = file.rfilename
-      if not (filename.endswith('.onnx') or filename.endswith('.onnx_data')): continue
-      if any(skip_str in filename for skip_str in SKIPPED_FILES): continue
-      head = requests.head(f"{HUGGINGFACE_URL}/{repo}/resolve/main/{filename}", allow_redirects=True)
-      file_size = file.size or int(head.headers.get('Content-Length', 0))
-      files_metadata.append({"file": filename, "size": f"{file_size/1e6:.2f}MB"})
-      total_size += file_size
-
-    repos_metadata["repositories"][repo] = {
-      "url": f"{HUGGINGFACE_URL}/{repo}",
-      "download_path": None,
-      "files": files_metadata,
-    }
-  repos_metadata['total_size'] = f"{total_size/1e9:.2f}GB"
-  repos_metadata['created_at'] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
-  return repos_metadata
-
-if __name__ == "__main__":
-  sort = "downloads" # recent 30 days downloads
-  huggingface_onnx_dir = Path(__file__).parent
-
-  parser = argparse.ArgumentParser(description="Produces a YAML file with metadata of top huggingface onnx models")
-  parser.add_argument("--limit", type=int, required=True, help="Number of top repositories to process (e.g., 100)")
-  parser.add_argument("--output", type=str, default="huggingface_repos.yaml", help="Output YAML file name to save the report")
-  args = parser.parse_args()
-
-  top_repos = get_top_repos(args.limit, sort)
-  metadata = get_metadata(top_repos)
-  yaml_path = huggingface_onnx_dir / args.output
-  with open(yaml_path, 'w') as f:
-    yaml.dump(metadata, f, sort_keys=False)
-    print(f"YAML saved to: {str(yaml_path)}")
@@ -1,29 +0,0 @@
-import yaml, argparse
-from pathlib import Path
-from huggingface_hub import snapshot_download
-
-def download_models(yaml_file: str, download_dir: str) -> None:
-  with open(yaml_file, 'r') as f: metadata = yaml.safe_load(f)
-  n = len(metadata["repositories"])
-
-  for i, (model_id, model_data) in enumerate(metadata["repositories"].items()):
-    print(f"Downloading {i+1}/{n}: {model_id}...")
-    allow_patterns = [file_info["file"] for file_info in model_data["files"]]
-    root_path = Path(snapshot_download(repo_id=model_id, allow_patterns=allow_patterns, cache_dir=download_dir))
-    # download configs too (the sizes are small)
-    snapshot_download(repo_id=model_id, allow_patterns=["*config.json"], cache_dir=download_dir)
-    print(f"Downloaded model files to: {root_path}")
-    model_data["download_path"] = str(root_path)
-
-  # Save the updated metadata back to the YAML file
-  with open(yaml_file, 'w') as f: yaml.dump(metadata, f, sort_keys=False)
-  print("Download completed according to YAML file.")
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser(description="Download models from Huggingface Hub based on a YAML configuration file.")
-  parser.add_argument("input", type=str, help="Path to the input YAML configuration file containing model information.")
-  args = parser.parse_args()
-
-  models_folder = Path(__file__).parent / "models"
-  models_folder.mkdir(parents=True, exist_ok=True)
-  download_models(args.input, str(models_folder))
@@ -0,0 +1,230 @@
+import yaml
+import time
+import requests
+import argparse
+from pathlib import Path
+from huggingface_hub import list_models, HfApi, snapshot_download
+from tinygrad.helpers import _ensure_downloads_dir
+DOWNLOADS_DIR = _ensure_downloads_dir() / "models"
+from tinygrad.helpers import tqdm
+
+def snapshot_download_with_retry(*, repo_id: str, allow_patterns: list[str]|tuple[str, ...]|None=None, cache_dir: str|Path|None=None,
+                                 tries: int=2, **kwargs) -> Path:
+  for attempt in range(tries):
+    try:
+      return Path(snapshot_download(
+        repo_id=repo_id,
+        allow_patterns=allow_patterns,
+        cache_dir=str(cache_dir) if cache_dir is not None else None,
+        **kwargs
+      ))
+    except Exception as e:
+      if attempt == tries-1: raise
+      time.sleep(1)
+
+# Constants for filtering models
+HUGGINGFACE_URL = "https://huggingface.co"
+SKIPPED_FILES = [
+  "fp16", "int8", "uint8", "quantized",      # numerical accuracy issues
+  "avx2", "arm64", "avx512", "avx512_vnni",  # numerical accuracy issues
+  "q4", "q4f16", "bnb4",                     # unimplemented quantization
+  "model_O4",                                # requires non cpu ort runner and MemcpyFromHost op
+  "merged",                                  # TODO implement attribute with graph type and Loop op
+]
+
+SKIPPED_REPO_PATHS = [
+  # Invalid model-index
+  "AdamCodd/vit-base-nsfw-detector",
+  # TODO: implement attribute with graph type and Loop op
+  "minishlab/potion-base-8M", "minishlab/M2V_base_output", "minishlab/potion-retrieval-32M",
+  # TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, GroupQueryAttention
+  "HuggingFaceTB/SmolLM2-360M-Instruct",
+  # TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, RotaryEmbedding, MultiHeadAttention
+  "HuggingFaceTB/SmolLM2-1.7B-Instruct",
+  # TODO: implement RandomNormalLike
+  "stabilityai/stable-diffusion-xl-base-1.0", "stabilityai/sdxl-turbo", 'SimianLuo/LCM_Dreamshaper_v7',
+  # TODO: implement NonZero
+  "mangoapps/fb_zeroshot_mnli_onnx",
+  # TODO huge Concat in here with 1024 (1, 3, 32, 32) Tensors, and maybe a MOD bug with const folding
+  "briaai/RMBG-2.0",
+]
+
+
+class HuggingFaceONNXManager:
+  def __init__(self):
+    self.base_dir = Path(__file__).parent
+    self.models_dir = DOWNLOADS_DIR
+    self.api = HfApi()
+
+  def discover_models(self, limit: int, sort: str = "downloads") -> list[str]:
+    print(f"Discovering top {limit} ONNX models sorted by {sort}...")
+    repos = []
+    i = 0
+
+    for model in list_models(filter="onnx", sort=sort):
+      if model.id in SKIPPED_REPO_PATHS:
+        continue
+
+      print(f"  {i+1}/{limit}: {model.id} ({getattr(model, sort)})")
+      repos.append(model.id)
+      i += 1
+      if i == limit:
+        break
+
+    print(f"Found {len(repos)} suitable ONNX models")
+    return repos
+
+  def collect_metadata(self, repos: list[str]) -> dict:
+    print(f"Collecting metadata for {len(repos)} repositories...")
+    metadata = {"repositories": {}}
+    total_size = 0
+
+    for repo in tqdm(repos, desc="Collecting metadata"):
+      try:
+        files_metadata = []
+        model_info = self.api.model_info(repo)
+
+        for file in model_info.siblings:
+          filename = file.rfilename
+          if not (filename.endswith('.onnx') or filename.endswith('.onnx_data')):
+            continue
+          if any(skip_str in filename for skip_str in SKIPPED_FILES):
+            continue
+
+          # Get file size from API or HEAD request
+          try:
+            head = requests.head(
+              f"{HUGGINGFACE_URL}/{repo}/resolve/main/{filename}",
+              allow_redirects=True,
+              timeout=10
+            )
+            file_size = file.size or int(head.headers.get('Content-Length', 0))
+          except requests.RequestException:
+            file_size = file.size or 0
+
+          files_metadata.append({
+            "file": filename,
+            "size": f"{file_size/1e6:.2f}MB"
+          })
+          total_size += file_size
+
+        if files_metadata:  # Only add repos with valid ONNX files
+          metadata["repositories"][repo] = {
+            "url": f"{HUGGINGFACE_URL}/{repo}",
+            "download_path": None,
+            "files": files_metadata,
+          }
+
+      except Exception as e:
+        print(f"WARNING: Failed to collect metadata for {repo}: {e}")
+        continue
+
+    metadata['total_size'] = f"{total_size/1e9:.2f}GB"
+    metadata['created_at'] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
+
+    print(f"Collected metadata for {len(metadata['repositories'])} repositories")
+    print(f"Total estimated download size: {metadata['total_size']}")
+
+    return metadata
+
+  def download_models(self, metadata: dict) -> dict:
+    self.models_dir.mkdir(parents=True, exist_ok=True)
+
+    repos = metadata["repositories"]
+    n = len(repos)
+
+    print(f"Downloading {n} repositories to {self.models_dir}...")
+
+    for i, (model_id, model_data) in enumerate(repos.items()):
+      print(f"  Downloading {i+1}/{n}: {model_id}...")
+
+      try:
+        # Download ONNX model files
+        allow_patterns = [file_info["file"] for file_info in model_data["files"]]
+        root_path = snapshot_download_with_retry(
+          repo_id=model_id,
+          allow_patterns=allow_patterns,
+          cache_dir=str(self.models_dir)
+        )
+
+        # Download config files (usually small)
+        snapshot_download_with_retry(
+          repo_id=model_id,
+          allow_patterns=["*config.json"],
+          cache_dir=str(self.models_dir)
+        )
+
+        model_data["download_path"] = str(root_path)
+        print(f"    Downloaded to: {root_path}")
+
+      except Exception as e:
+        print(f"    ERROR: Failed to download {model_id}: {e}")
+        model_data["download_path"] = None
+        continue
+
+    successful_downloads = sum(1 for repo in repos.values() if repo["download_path"] is not None)
+    print(f"Successfully downloaded {successful_downloads}/{n} repositories")
+    print(f"All models saved to: {self.models_dir}")
+
+    return metadata
+
+  def save_metadata(self, metadata: dict, output_file: str):
+    yaml_path = self.base_dir / output_file
+    with open(yaml_path, 'w') as f:
+      yaml.dump(metadata, f, sort_keys=False)
+    print(f"Metadata saved to: {yaml_path}")
+
+  def discover_and_download(self, limit: int, output_file: str = "huggingface_repos.yaml",
+                          sort: str = "downloads", download: bool = True):
+    print(f"Starting HuggingFace ONNX workflow...")
+    print(f"   Limit: {limit} models")
+    print(f"   Sort by: {sort}")
+    print(f"   Download: {'Yes' if download else 'No'}")
+    print(f"   Output: {output_file}")
+    print("-" * 50)
+
+    repos = self.discover_models(limit, sort)
+
+    metadata = self.collect_metadata(repos)
+
+    if download:
+      metadata = self.download_models(metadata)
+
+    self.save_metadata(metadata, output_file)
+
+    print("-" * 50)
+    print("Workflow completed successfully!")
+    if download:
+      successful = sum(1 for repo in metadata["repositories"].values()
+                     if repo["download_path"] is not None)
+      print(f"{successful}/{len(metadata['repositories'])} models downloaded")
+
+    return metadata
+
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser(
+    description="HuggingFace ONNX Model Manager - Discover, collect metadata, and download ONNX models",
+  )
+
+  parser.add_argument("--limit", type=int, help="Number of top repositories to process")
+  parser.add_argument("--output", type=str, default="huggingface_repos.yaml",
+                     help="Output YAML file name (default: huggingface_repos.yaml)")
+  parser.add_argument("--sort", type=str, default="downloads",
+                     choices=["downloads", "likes", "created", "modified"],
+                     help="Sort criteria for model discovery (default: downloads)")
+
+  parser.add_argument("--download", action="store_true", default=False,
+                     help="Download models after collecting metadata")
+
+  args = parser.parse_args()
+
+  if not args.limit: parser.error("--limit is required")
+
+  manager = HuggingFaceONNXManager()
+  manager.discover_and_download(
+    limit=args.limit,
+    output_file=args.output,
+    sort=args.sort,
+    download=args.download
+  )
@@ -1,10 +1,11 @@
-import onnx, yaml, tempfile, time, collections, pprint, argparse, json
+import onnx, yaml, tempfile, time, argparse, json
 from pathlib import Path
-from tinygrad.frontend.onnx import OnnxRunner, onnx_load
-from extra.onnx import get_onnx_ops
+from typing import Any
+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

-def get_config(root_path: Path):
+def get_config(root_path: Path) -> dict[str, Any]:
  ret = {}
  for path in root_path.rglob("*config.json"):
    config = json.load(path.open())
@@ -12,20 +13,19 @@ def get_config(root_path: Path):
      ret.update(config)
  return ret

-def run_huggingface_validate(onnx_model_path, config, rtol, atol):
-  onnx_model = onnx_load(onnx_model_path)
-  onnx_runner = OnnxRunner(onnx_model)
-  inputs = get_example_inputs(onnx_runner.graph_inputs, config)
-  validate(onnx_model_path, inputs, rtol=rtol, atol=atol)
-
-def get_tolerances(file_name): # -> rtol, atol
+def get_tolerances(file_name: str) -> tuple[float, float]:
  # TODO very high rtol atol
  if "fp16" in file_name: return 9e-2, 9e-2
  if any(q in file_name for q in ["int8", "uint8", "quantized"]): return 4, 4
  return 4e-3, 3e-2

+def run_huggingface_validate(onnx_model_path: str | Path, config: dict[str, Any], rtol: float, atol: float):
+  onnx_runner = OnnxRunner(onnx_model_path)
+  inputs = get_example_inputs(onnx_runner.graph_inputs, config)
+  validate(onnx_model_path, inputs, rtol=rtol, atol=atol)
+
 def validate_repos(models:dict[str, tuple[Path, Path]]):
-  print(f"** Validating {len(model_paths)} models **")
+  print(f"** Validating {len(models)} models **")
  for model_id, (root_path, relative_path) in models.items():
    print(f"validating model {model_id}")
    model_path = root_path / relative_path
@@ -37,25 +37,6 @@ def validate_repos(models:dict[str, tuple[Path, Path]]):
    et = time.time() - st
    print(f"passed, took {et:.2f}s")

-def retrieve_op_stats(models:dict[str, tuple[Path, Path]]) -> dict:
-  ret = {}
-  op_counter = collections.Counter()
-  unsupported_ops = collections.defaultdict(set)
-  supported_ops = get_onnx_ops()
-  print(f"** Retrieving stats from {len(model_paths)} models **")
-  for model_id, (root_path, relative_path) in models.items():
-    print(f"examining {model_id}")
-    model_path = root_path / relative_path
-    onnx_runner = OnnxRunner(onnx.load(model_path))
-    for node in onnx_runner.graph_nodes:
-      op_counter[node.op] += 1
-      if node.op not in supported_ops:
-        unsupported_ops[node.op].add(model_id)
-    del onnx_runner
-  ret["unsupported_ops"] = {k:list(v) for k, v in unsupported_ops.items()}
-  ret["op_counter"] = op_counter.most_common()
-  return ret
-
 def debug_run(model_path, truncate, config, rtol, atol):
  if truncate != -1:
    model = onnx.load(model_path)
@@ -72,12 +53,9 @@ def debug_run(model_path, truncate, config, rtol, atol):
    run_huggingface_validate(model_path, config, rtol, atol)

 if __name__ == "__main__":
-  parser = argparse.ArgumentParser(description="Huggingface ONNX Model Validator and Ops Checker")
-  parser.add_argument("input", type=str, help="Path to the input YAML configuration file containing model information.")
-  parser.add_argument("--check_ops", action="store_true", default=False,
-                      help="Check support for ONNX operations in models from the YAML file")
-  parser.add_argument("--validate", action="store_true", default=False,
-                      help="Validate correctness of models from the YAML file")
+  parser = argparse.ArgumentParser(description="Huggingface ONNX Model Validator")
+  parser.add_argument("--validate", type=str, default="",
+                      help="Validate correctness of models from the specified YAML configuration file")
  parser.add_argument("--debug", type=str, default="",
                      help="""Validates without explicitly needing a YAML or models pre-installed.
                      provide repo id (e.g. "minishlab/potion-base-8M") to validate all onnx models inside the repo
@@ -86,13 +64,13 @@ if __name__ == "__main__":
  parser.add_argument("--truncate", type=int, default=-1, help="Truncate the ONNX model so intermediate results can be validated")
  args = parser.parse_args()

-  if not (args.check_ops or args.validate or args.debug):
-    parser.error("Please provide either --validate, --check_ops, or --debug.")
+  if not (args.validate or args.debug):
+    parser.error("Please provide either --validate <yaml_file> or --debug <repo_id>.")
  if args.truncate != -1 and not args.debug:
    parser.error("--truncate and --debug should be used together for debugging")

-  if args.check_ops or args.validate:
-    with open(args.input, 'r') as f:
+  if args.validate:
+    with open(args.validate, 'r') as f:
      data = yaml.safe_load(f)
      assert all(repo["download_path"] is not None for repo in data["repositories"].values()), "please run `download_models.py` for this yaml"
      model_paths = {
@@ -102,22 +80,16 @@ if __name__ == "__main__":
        if model["file"].endswith(".onnx")
      }

-    if args.check_ops:
-      pprint.pprint(retrieve_op_stats(model_paths))
-
-    if args.validate:
-      validate_repos(model_paths)
+    validate_repos(model_paths)

  if args.debug:
-    from huggingface_hub import snapshot_download
-    download_dir = Path(__file__).parent / "models"
    path:list[str] = args.debug.split("/")
    if len(path) == 2:
      # repo id
      # validates all onnx models inside repo
      repo_id = "/".join(path)
-      root_path = Path(snapshot_download(repo_id=repo_id, allow_patterns=["*.onnx", "*.onnx_data"], cache_dir=download_dir))
-      snapshot_download(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=download_dir)
+      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*.onnx", "*.onnx_data"], cache_dir=DOWNLOADS_DIR)
+      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=DOWNLOADS_DIR)
      config = get_config(root_path)
      for onnx_model in root_path.rglob("*.onnx"):
        rtol, atol = get_tolerances(onnx_model.name)
@@ -129,8 +101,8 @@ if __name__ == "__main__":
      onnx_model = path[-1]
      assert path[-1].endswith(".onnx")
      repo_id, relative_path = "/".join(path[:2]), "/".join(path[2:])
-      root_path = Path(snapshot_download(repo_id=repo_id, allow_patterns=[relative_path], cache_dir=download_dir))
-      snapshot_download(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=download_dir)
+      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=[relative_path], cache_dir=DOWNLOADS_DIR)
+      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=DOWNLOADS_DIR)
      config = get_config(root_path)
      rtol, atol = get_tolerances(onnx_model)
      print(f"validating {relative_path} with truncate={args.truncate}, {rtol=}, {atol=}")
@@ -4,9 +4,10 @@ import numpy as np
 np.set_printoptions(suppress=True)
 import math, functools, time, random, statistics
 from tinygrad.helpers import DEBUG, getenv, CACHELEVEL, diskcache_get, diskcache_put, colored, Profiling
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.opt.kernel import Kernel
 from tinygrad.device import Buffer, Device, CompileError
-from tinygrad.opt.search import _ensure_buffer_alloc, get_kernel_actions, _time_program
+from tinygrad.codegen.opt.search import _ensure_buffer_alloc, get_kernel_actions, _time_program
+from tinygrad.engine.realize import get_program

 class MCTSNode:
  def __init__(self, kernel:Kernel, parent=None):
@@ -87,7 +88,7 @@ def mcts_search(lin:Kernel, rawbufs:List[Buffer], amt:int) -> Kernel:
    return ret

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

@@ -110,7 +111,7 @@ def mcts_search(lin:Kernel, rawbufs:List[Buffer], amt:int) -> Kernel:
      seen_asts[opt_ast.key] = node

      # lowering (50% of the time)
-      p = node.kernel.to_program(name_override="test")
+      p = get_program(node.kernel.get_optimized_ast(name_override="test"), node.kernel.opts)

      # rollout
      tm1 = time.perf_counter()
@@ -1,5 +1,5 @@
 from typing import Union, Optional, Any
-import collections
+import collections, math
 from tinygrad import Tensor, Variable, TinyJit, dtypes, nn, Device
 from tinygrad.helpers import getenv, DEBUG

@@ -99,7 +99,9 @@ class FeedForward:
    self.w3 = linear(dim, hidden_dim, bias=False) # the gate in Gated Linear Unit

  def __call__(self, x:Tensor) -> Tensor:
-    return self.w2(self.w1(x).silu() * self.w3(x)) # SwiGLU [arxiv/2002.05202, eq (5)]
+    w1 = self.w1(x).silu()
+    w3 = self.w3(x.contiguous_backward())  # this fixes a strange fusion that makes tensor cores miss
+    return self.w2(w1 * w3)

 class TransformerBlock:
  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int, norm_eps:float, max_context:int, linear=nn.Linear,
@@ -111,7 +113,7 @@ class TransformerBlock:

  def __call__(self, x:Tensor, start_pos:Union[Variable,int], freqs_cis:Tensor, mask:Optional[Tensor]):
    h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)
-    return (h + self.feed_forward(self.ffn_norm(h))).contiguous()
+    return (h + self.feed_forward(self.ffn_norm(h))).contiguous().contiguous_backward()

 # standard openai sampling
 def sample(logits: Tensor, temp: float, k: int, p: float, af: float, ap: float):
@@ -166,27 +168,27 @@ def sample(logits: Tensor, temp: float, k: int, p: float, af: float, ap: float):

 class Transformer:
  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size, linear=nn.Linear, embedding=nn.Embedding,
-               n_kv_heads=None, rope_theta=10000, max_context=1024, jit=True, feed_forward=FeedForward, qk_norm=None):
-    self.layers = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, max_context, linear, feed_forward=feed_forward, qk_norm=qk_norm) for _ in range(n_layers)]
+               n_kv_heads=None, rope_theta=10000, max_context=1024, jit=True, feed_forward=FeedForward, qk_norm=None, disable_kv_cache=False):
+    self.layers = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, 0 if disable_kv_cache else max_context,
+                                    linear, feed_forward=feed_forward, qk_norm=qk_norm) for _ in range(n_layers)]
    self.norm = nn.RMSNorm(dim, norm_eps)
    self.tok_embeddings = embedding(vocab_size, dim)
    self.output = nn.Linear(dim, vocab_size, bias=False) if embedding == nn.Embedding else linear(dim, vocab_size, bias=False)
    self.max_context = max_context
-    self.freqs_cis = precompute_freqs_cis(dim // n_heads, self.max_context * 2, rope_theta).contiguous()
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, self.max_context * 2, rope_theta).contiguous().requires_grad_(False)
    self.forward_jit = TinyJit(self.forward) if jit else None

  def forward(self, tokens:Tensor, start_pos:Union[Variable,int], temperature:float, top_k:int, top_p:float, alpha_f:float, alpha_p:float):
    _bsz, seqlen = tokens.shape
    h = self.tok_embeddings(tokens)
+    freqs_cis = self.freqs_cis.cast(h.dtype)[:, start_pos:start_pos+seqlen, :, :, :]

-    self.freqs_cis = self.freqs_cis.cast(h.dtype).kernelize()
-    freqs_cis = self.freqs_cis[:, start_pos:start_pos+seqlen, :, :, :]
-
-    mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1).kernelize() if seqlen > 1 else None
+    mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1) if seqlen > 1 else None
    for layer in self.layers: h = layer(h, start_pos, freqs_cis, mask)
-    logits = self.output(self.norm(h)).float()[:, -1, :]
+    logits = self.output(self.norm(h))
+    if math.isnan(temperature): return logits

-    return sample(logits.flatten(), temperature, top_k, top_p, alpha_f, alpha_p).kernelize()
+    return sample(logits[:, -1, :].flatten(), temperature, top_k, top_p, alpha_f, alpha_p)

  def __call__(self, tokens:Tensor, start_pos:int, temperature:float=0.0, top_k:int=0, top_p:float=0.8, alpha_f:float=0.0, alpha_p:float=0.0):
    # TODO: better way to handle the first call v.s. the rest?
@@ -247,8 +249,5 @@ def convert_from_gguf(weights:dict[str, Tensor], n_layers:int):
  return sd

 def fix_bf16(weights:dict[Any, Tensor]):
-  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()}
+  # 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()}
--- a/Show More
+++ b/Show More