amd isel renderer

viz: wave color is locally scoped (#14728 )
am: reset mi300 with pm4 (#14727 )
2026-02-13 17:50:39 +08:00 · 2026-02-13 18:22:20 +09:00 · 2026-02-13 11:22:32 +03:00 · 2026-02-13 16:07:06 +08:00 · 2026-02-13 16:06:58 +08:00 · 2026-02-13 00:01:05 -05:00
611 changed files with 225391 additions and 139588 deletions
@@ -11,5 +11,5 @@ runs:
        git fetch origin $CURRENT_SHA
        export COMMIT_MESSAGE=$(git show -s --format=%B "$CURRENT_SHA")
        export CURRENT_HEAD=$(git rev-parse HEAD)
-        cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && IGNORE_OOB=1 PYTHONPATH=. python3 process_replay.py
+        cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && CHECK_OOB=0 PYTHONPATH=. python3 process_replay.py
        git checkout $CURRENT_HEAD  # restore to branch
@@ -56,7 +56,15 @@ runs:

    # **** Caching packages ****

+    - name: Cache Python packages (PR)
+      if: github.event_name == 'pull_request'
+      id: restore-venv-pr
+      uses: actions/cache/restore@v4
+      with:
+        path: ${{ github.workspace }}/.venv
+        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
    - name: Cache Python packages
+      if: github.event_name != 'pull_request'
      id: restore-venv
      uses: actions/cache@v4
      with:
@@ -65,23 +73,23 @@ runs:

    # **** Caching downloads ****

-    - name: Cache downloads (Linux)
-      if: inputs.key != '' && runner.os == 'Linux'
-      uses: actions/cache@v4
+    - name: Cache downloads (PR)
+      if: inputs.key != '' && github.event_name == 'pull_request'
+      uses: actions/cache/restore@v4
      with:
-        path: ~/.cache/tinygrad/downloads/
+        path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
        key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
-    - name: Cache downloads (macOS)
-      if: inputs.key != '' && runner.os == 'macOS'
+    - name: Cache downloads
+      if: inputs.key != '' && github.event_name != 'pull_request'
      uses: actions/cache@v4
      with:
-        path: ~/Library/Caches/tinygrad/downloads/
+        path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
        key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}

    # **** Python deps ****

    - name: Install dependencies in venv (with extra)
-      if: inputs.deps != '' && steps.restore-venv.outputs.cache-hit != 'true'
+      if: inputs.deps != '' && steps.restore-venv-pr.outputs.cache-hit != 'true' && steps.restore-venv.outputs.cache-hit != 'true'
      shell: bash
      run: |
        python -m venv .venv
@@ -92,7 +100,7 @@ runs:
        fi
        python -m pip install -e ".[${{ inputs.deps }}]" ${{ inputs.pydeps }} --extra-index-url https://download.pytorch.org/whl/cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
    - name: Install dependencies in venv (without extra)
-      if: inputs.deps == '' && steps.restore-venv.outputs.cache-hit != 'true'
+      if: inputs.deps == '' && steps.restore-venv-pr.outputs.cache-hit != 'true' && steps.restore-venv.outputs.cache-hit != 'true'
      shell: bash
      run: |
        python -m venv .venv
@@ -182,8 +190,14 @@ runs:
        echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
        echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"

+    - name: Cache apt (PR)
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true') && github.event_name == 'pull_request'
+      uses: actions/cache/restore@v4
+      with:
+        path: /var/cache/apt/archives/
+        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
    - name: Cache apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true') && github.event_name != 'pull_request'
      uses: actions/cache@v4
      with:
        path: /var/cache/apt/archives/
@@ -239,8 +253,17 @@ runs:
        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 (PR)
+      if: inputs.ocelot == 'true' && github.event_name == 'pull_request'
+      id: cache-build-pr
+      uses: actions/cache/restore@v4
+      env:
+        cache-name: cache-gpuocelot-build-1
+      with:
+        path: ${{ github.workspace }}/gpuocelot/ocelot
+        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
    - name: Cache gpuocelot
-      if: inputs.ocelot == 'true'
+      if: inputs.ocelot == 'true' && github.event_name != 'pull_request'
      id: cache-build
      uses: actions/cache@v4
      env:
@@ -249,7 +272,7 @@ runs:
        path: ${{ github.workspace }}/gpuocelot/ocelot
        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
    - name: Clone/compile gpuocelot
-      if: inputs.ocelot == 'true' && steps.cache-build.outputs.cache-hit != 'true'
+      if: inputs.ocelot == 'true' && steps.cache-build-pr.outputs.cache-hit != 'true' && steps.cache-build.outputs.cache-hit != 'true'
      shell: bash
      run: |
        git clone --recurse-submodules https://github.com/gpuocelot/gpuocelot.git ${{ github.workspace }}/gpuocelot
@@ -14,10 +14,12 @@ on:
    paths:
    - 'tinygrad/runtime/autogen/**/*'
    - 'tinygrad/runtime/support/autogen.py'
+    - '.github/workflows/autogen.yml'
  workflow_dispatch:
    paths:
    - 'tinygrad/runtime/autogen/**/*'
    - 'tinygrad/runtime/support/autogen.py'
+    - '.github/workflows/autogen.yml'

 jobs:
  autogen:
@@ -38,103 +40,36 @@ jobs:
        mesa: 'true'
        pydeps: 'pyyaml mako'
    - name: Install autogen support packages
-      run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev
-    - name: Verify OpenCL autogen
+      run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev
+    - name: Regenerate autogen files
      run: |
-        mv tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
+        find tinygrad/runtime/autogen -type f -name "*.py" -not -path "*/amd/*" -not -name "__init__.py" -not -name "comgr_3.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete
        python3 -c "from tinygrad.runtime.autogen import opencl"
-        diff /tmp/opencl.py.bak tinygrad/runtime/autogen/opencl.py
-    - name: Verify CUDA autogen
-      run: |
-        mv tinygrad/runtime/autogen/cuda.py /tmp/cuda.py.bak
-        mv tinygrad/runtime/autogen/nvrtc.py /tmp/nvrtc.py.bak
-        mv tinygrad/runtime/autogen/nvjitlink.py /tmp/nvjitlink.py.bak
-        mv tinygrad/runtime/autogen/nv_570.py /tmp/nv_570.py.bak
-        mv tinygrad/runtime/autogen/nv.py /tmp/nv.py.bak
-        python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv"
-        diff /tmp/cuda.py.bak tinygrad/runtime/autogen/cuda.py
-        diff /tmp/nvrtc.py.bak tinygrad/runtime/autogen/nvrtc.py
-        diff /tmp/nvjitlink.py.bak tinygrad/runtime/autogen/nvjitlink.py
-        diff /tmp/nv_570.py.bak tinygrad/runtime/autogen/nv_570.py
-        diff /tmp/nv.py.bak tinygrad/runtime/autogen/nv.py
-    - name: Verify AMD autogen
-      run: |
-        mv tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
-        mv tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
-        mv tinygrad/runtime/autogen/hip.py /tmp/hip.py.bak
-        mv tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
-        mv tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
-        mv tinygrad/runtime/autogen/rocprof.py /tmp/rocprof.py.bak
-        mv tinygrad/runtime/autogen/am/am.py /tmp/am_am.py.bak
-        mv tinygrad/runtime/autogen/am/pm4_soc15.py /tmp/am_pm4_soc15.py.bak
-        mv tinygrad/runtime/autogen/am/pm4_nv.py /tmp/am_pm4_nv.py.bak
-        mv tinygrad/runtime/autogen/am/sdma_4_0_0.py /tmp/am_sdma_4_0_0.py.bak
-        mv tinygrad/runtime/autogen/am/sdma_5_0_0.py /tmp/am_sdma_5_0_0.py.bak
-        mv tinygrad/runtime/autogen/am/sdma_6_0_0.py /tmp/am_sdma_6_0_0.py.bak
-        mv tinygrad/runtime/autogen/am/smu_v13_0_0.py /tmp/am_smu_v13_0_0.py.bak
-        mv tinygrad/runtime/autogen/am/smu_v14_0_2.py /tmp/am_smu_v14_0_2.py.bak
-        python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof; from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v14_0_2"
-        diff /tmp/comgr.py.bak tinygrad/runtime/autogen/comgr.py
-        diff /tmp/hsa.py.bak tinygrad/runtime/autogen/hsa.py
-        diff /tmp/hip.py.bak tinygrad/runtime/autogen/hip.py
-        diff /tmp/amd_gpu.py.bak tinygrad/runtime/autogen/amd_gpu.py
-        diff /tmp/sqtt.py.bak tinygrad/runtime/autogen/sqtt.py
-        diff /tmp/rocprof.py.bak tinygrad/runtime/autogen/rocprof.py
-        diff /tmp/am_am.py.bak tinygrad/runtime/autogen/am/am.py
-        diff /tmp/am_pm4_soc15.py.bak tinygrad/runtime/autogen/am/pm4_soc15.py
-        diff /tmp/am_pm4_nv.py.bak tinygrad/runtime/autogen/am/pm4_nv.py
-        diff /tmp/am_sdma_4_0_0.py.bak tinygrad/runtime/autogen/am/sdma_4_0_0.py
-        diff /tmp/am_sdma_5_0_0.py.bak tinygrad/runtime/autogen/am/sdma_5_0_0.py
-        diff /tmp/am_sdma_6_0_0.py.bak tinygrad/runtime/autogen/am/sdma_6_0_0.py
-        diff /tmp/am_smu_v13_0_0.py.bak tinygrad/runtime/autogen/am/smu_v13_0_0.py
-        diff /tmp/am_smu_v14_0_2.py.bak tinygrad/runtime/autogen/am/smu_v14_0_2.py
-    - name: Verify Linux autogen
-      run: |
-        mv tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
-        mv tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
-        mv tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
-        mv tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
-        mv tinygrad/runtime/autogen/pci.py /tmp/pci.py.bak
-        mv tinygrad/runtime/autogen/vfio.py /tmp/vfio.py.bak
+        python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv"
+        python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
+        python3 -c "from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v13_0_6, smu_v14_0_2"
        python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
-        diff /tmp/libc.py.bak tinygrad/runtime/autogen/libc.py
-        diff /tmp/kfd.py.bak tinygrad/runtime/autogen/kfd.py
-        diff /tmp/io_uring.py.bak tinygrad/runtime/autogen/io_uring.py
-        diff /tmp/ib.py.bak tinygrad/runtime/autogen/ib.py
-        diff /tmp/pci.py.bak tinygrad/runtime/autogen/pci.py
-        diff /tmp/vfio.py.bak tinygrad/runtime/autogen/vfio.py
-    - name: Verify LLVM autogen
-      run: |
-        mv tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
        python3 -c "from tinygrad.runtime.autogen import llvm"
-        diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
-    - name: Verify WebGPU autogen
-      run: |
-        mv tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak
        python3 -c "from tinygrad.runtime.autogen import webgpu"
-        diff /tmp/webgpu.py.bak tinygrad/runtime/autogen/webgpu.py
-    - name: Verify Qualcomm autogen
-      run: |
-        mv tinygrad/runtime/autogen/kgsl.py /tmp/kgsl.py.bak
-        mv tinygrad/runtime/autogen/qcom_dsp.py /tmp/qcom_dsp.py.bak
        python3 -c "from tinygrad.runtime.autogen import kgsl, qcom_dsp"
-        diff /tmp/kgsl.py.bak tinygrad/runtime/autogen/kgsl.py
-        diff /tmp/qcom_dsp.py.bak tinygrad/runtime/autogen/qcom_dsp.py
-    - name: Verify libusb autogen
-      run: |
-        mv tinygrad/runtime/autogen/libusb.py /tmp/libusb.py.bak
        python3 -c "from tinygrad.runtime.autogen import libusb"
-        diff /tmp/libusb.py.bak tinygrad/runtime/autogen/libusb.py
-    - name: Verify mesa autogen
-      run: |
-        mv tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
        python3 -c "from tinygrad.runtime.autogen import mesa"
-        diff /tmp/mesa.py.bak tinygrad/runtime/autogen/mesa.py
-    - name: Verify libclang autogen
-      run: |
-        cp tinygrad/runtime/autogen/libclang.py /tmp/libclang.py.bak
+        python3 -c "from tinygrad.runtime.autogen import avcodec"
        REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
-        diff /tmp/libclang.py.bak tinygrad/runtime/autogen/libclang.py
+    - name: Check for differences
+      run: |
+        if ! git diff --quiet; then
+          git diff > autogen-ubuntu.patch
+          echo "Autogen files out of date. Apply patch from: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
+          exit 1
+        fi
+    - name: Upload patch artifact
+      if: failure()
+      uses: actions/upload-artifact@v4
+      with:
+        name: autogen-ubuntu-patch
+        path: autogen-ubuntu.patch
+
  autogen-mac:
    name: In-tree Autogen (macos)
    runs-on: macos-14
@@ -146,11 +81,24 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        llvm: 'true'
-    - name: Verify macos autogen
+    - name: Regenerate autogen files
      run: |
-        mv tinygrad/runtime/autogen/metal.py /tmp/metal.py.bak
-        LIBCLANG_PATH=/opt/homebrew/opt/llvm@20/lib/libclang.dylib python3 -c "from tinygrad.runtime.autogen import metal"
-        diff /tmp/metal.py.bak tinygrad/runtime/autogen/metal.py
+        rm tinygrad/runtime/autogen/metal.py tinygrad/runtime/autogen/iokit.py tinygrad/runtime/autogen/corefoundation.py
+        python3 -c "from tinygrad.runtime.autogen import metal, iokit, corefoundation"
+    - name: Check for differences
+      run: |
+        if ! git diff --quiet; then
+          git diff > autogen-macos.patch
+          echo "Autogen files out of date. Apply patch from: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
+          exit 1
+        fi
+    - name: Upload patch artifact
+      if: failure()
+      uses: actions/upload-artifact@v4
+      with:
+        name: autogen-macos-patch
+        path: autogen-macos.patch
+
  autogen-comgr-3:
    name: In-tree Autogen (comgr 3)
    runs-on: ubuntu-24.04
@@ -169,8 +117,20 @@ jobs:
        echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
        sudo apt -qq update || true
        sudo apt-get install -y --no-install-recommends libclang-20-dev comgr
-    - name: Verify comgr (3) autogen
+    - name: Regenerate autogen files
      run: |
-        mv tinygrad/runtime/autogen/comgr_3.py /tmp/comgr_3.py.bak
+        rm tinygrad/runtime/autogen/comgr_3.py
        python3 -c "from tinygrad.runtime.autogen import comgr_3"
-        diff /tmp/comgr_3.py.bak tinygrad/runtime/autogen/comgr_3.py
+    - name: Check for differences
+      run: |
+        if ! git diff --quiet; then
+          git diff > autogen-comgr3.patch
+          echo "Autogen files out of date. Apply patch from: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
+          exit 1
+        fi
+    - name: Upload patch artifact
+      if: failure()
+      uses: actions/upload-artifact@v4
+      with:
+        name: autogen-comgr3-patch
+        path: autogen-comgr3.patch
@@ -16,6 +16,48 @@ on:
  workflow_dispatch:

 jobs:
+  # the goal of this test is to replicate a normal person on a laptop running the test
+  # no process replay, no benchmarks, no CI, just a normal laptop person
+  # the 3 minute timeout should not be raised
+  testmacpytest:
+    name: Mac pytest
+    runs-on: [self-hosted, macOS]
+    timeout-minutes: 3
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    # brew install uv
+    - name: setup python environment
+      run: |
+        rm -rf /tmp/tinygrad_pytest_ci
+        uv venv /tmp/tinygrad_pytest_ci
+        source /tmp/tinygrad_pytest_ci/bin/activate
+        uv pip install .[testing]
+    - name: setup staging db
+      run: |
+        echo "CACHEDB=/tmp/pytest-db-ci.db" >> $GITHUB_ENV
+        rm -f /tmp/pytest-db-ci*
+    # TODO: remove this step once all old caches are migrated
+    - name: Migrate old huggingface cache (symlinks break onnxruntime 1.24+)
+      run: |
+        cd ~/Library/Caches/tinygrad/downloads/models 2>/dev/null || exit 0
+        for old_dir in models--*; do
+          [ -d "$old_dir" ] || continue
+          repo_id=$(echo "$old_dir" | sed 's/models--//; s/--/\//g')
+          snapshot=$(ls -1 "$old_dir/snapshots" 2>/dev/null | head -1)
+          [ -n "$snapshot" ] || continue
+          mkdir -p "$repo_id"
+          cp -RLn "$old_dir/snapshots/$snapshot/"* "$repo_id/" 2>/dev/null || true
+        done
+    - name: Run pytest -nauto
+      run: |
+        source /tmp/tinygrad_pytest_ci/bin/activate
+        pytest -nauto --durations=20
+
  testmacbenchmark:
    name: Mac Benchmark
    env:
@@ -49,19 +91,19 @@ jobs:
    - name: Print macOS version
      run: sw_vers
    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 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=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
    - name: Run Stable Diffusion without fp16
-      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=720 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=720 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing
    - name: Run Stable Diffusion v2
      # TODO: very slow step time
-      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 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=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing
    # process replay can't capture this, the graph is too large
    - name: Run SDXL
-      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing
    - name: Run model inference benchmark
      run: METAL=1 NOCLANG=1 python3.11 test/external/external_model_benchmark.py
    - name: Test speed vs torch
-      run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
+      run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py
    - name: Test tensor cores
      run: METAL=1 python3.11 test/opt/test_tensor_cores.py
    - name: Test AMX tensor cores
@@ -71,84 +113,59 @@ jobs:
        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
+      run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py
    - name: Run Tensor Core GEMM (half)
-      run: DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_half.txt
+      run: DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3.11 extra/gemm/simple_matmul.py
    - name: Run Tensor Core GEMM (bfloat16)
-      run: DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
+      run: DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py
    - name: Fuzz Padded Tensor Core GEMM
      run: METAL=1 M_START=6 M_STOP=10 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=6 K_STOP=24 K_STEP=1 TC_OPT=2 DEBUG=2 python3.11 ./extra/gemm/fuzz_matmul.py
    - name: Run LLaMA
      run: |
-        BENCHMARK_LOG=llama_nojit JIT=0 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
-        BENCHMARK_LOG=llama JIT=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
+        BENCHMARK_LOG=llama_nojit JIT=0 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=llama JIT=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
    - name: Run LLaMA with BEAM
-      run: BENCHMARK_LOG=llama_beam JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
+      run: BENCHMARK_LOG=llama_beam JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
    - name: Run quantized LLaMA
      run: |
-        BENCHMARK_LOG=llama_int8 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize int8 | tee llama_int8.txt
-        BENCHMARK_LOG=llama_nf4 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize nf4 | tee llama_nf4.txt
+        BENCHMARK_LOG=llama_int8 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize int8
+        BENCHMARK_LOG=llama_nf4 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize nf4
    - name: Run quantized LLaMA3
      run: |
-        BENCHMARK_LOG=llama3_int8 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize int8 | tee llama3_int8.txt
-        BENCHMARK_LOG=llama3_nf4 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize nf4 | tee llama3_nf4.txt
+        BENCHMARK_LOG=llama3_int8 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize int8
+        BENCHMARK_LOG=llama3_nf4 python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize nf4
    #- name: Run LLaMA 7B on 4 (virtual) GPUs
-    #  run: python3.11 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_four_gpu.txt
+    #  run: python3.11 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing
    - name: Run GPT2
      run: |
-        BENCHMARK_LOG=gpt2_nojit JIT=0 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        BENCHMARK_LOG=gpt2 JIT=1 ASSERT_MIN_STEP_TIME=13 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2_nojit JIT=0 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=gpt2 JIT=1 ASSERT_MIN_STEP_TIME=13 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
    - 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
+      run: BENCHMARK_LOG=gpt2_half HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing
    - name: Run GPT2 w HALF/BEAM
-      run: BENCHMARK_LOG=gpt2_half_beam HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
+      run: BENCHMARK_LOG=gpt2_half_beam HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing
    - name: Run OLMoE
      run: BENCHMARK_LOG=olmoe python3.11 examples/olmoe.py
    - name: Train MNIST
-      run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py | tee beautiful_mnist.txt
+      run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py

    # NOTE: this is failing in CI. it is not failing on my machine and I don't really have a way to debug it
    # the error is "RuntimeError: Internal Error (0000000e:Internal Error)"
    #- name: Run 10 CIFAR training steps
-    #  run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=3000 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
+    #  run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=3000 STEPS=10 python3.11 examples/hlb_cifar10.py
    #- name: Run 10 CIFAR training steps w HALF
-    #  run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=3000 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
+    #  run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=3000 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py

    #- 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
+    #  run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py
    # TODO: too slow
    # - name: Run 10 CIFAR training steps w winograd
-    #   run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+    #   run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (Mac)
        path: |
          onnx_inference_speed.csv
-          torch_speed.txt
-          llama_unjitted.txt
-          llama_jitted.txt
-          llama_beam.txt
-          llama_int8.txt
-          llama_nf4.txt
-          llama3_int8.txt
-          llama3_nf4.txt
-          llama_four_gpu.txt
-          gpt2_unjitted.txt
-          gpt2_jitted.txt
-          gpt2_half.txt
-          gpt2_half_beam.txt
-          matmul.txt
-          matmul_half.txt
-          matmul_bfloat16.txt
-          sd.txt
-          sd_no_fp16.txt
-          sdv2.txt
-          sdxl.txt
-          beautiful_mnist.txt
-          train_cifar.txt
-          train_cifar_half.txt
-          train_cifar_bf16.txt
-          train_cifar_wino.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.11 process_replay.py

@@ -170,6 +187,10 @@ jobs:
      run: |
        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: Kill stale pids
+      run: |
+        PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
+        PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
    - 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
@@ -215,7 +236,7 @@ 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/speed/external_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
    - name: Test speed vs theoretical
      run: NV=1 IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
    - name: Test benchmark allreduce
@@ -226,79 +247,58 @@ jobs:
        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
-        CUDA=1 SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_fp8.txt
+        CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+        CUDA=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+        CUDA=1 SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
+        CUDA=1 SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
    - name: Run Tensor Core GEMM (PTX)
-      run: NV=1 NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
+      run: NV=1 NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
    - 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
+      run: NV=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
    - name: Test NV=1
      run: DEBUG=2 NV=1 python -m pytest -rA test/test_tiny.py
    - name: Test CUDA=1
      run: DEBUG=2 CUDA=1 python -m pytest -rA test/test_tiny.py
    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
    # TODO: too slow
    # - name: Run SDXL
-    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+    #   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
    - 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
-        BENCHMARK_LOG=llama NV=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
+        BENCHMARK_LOG=llama_nojit NV=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=llama NV=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
    - name: Run LLaMA with BEAM
-      run: BENCHMARK_LOG=llama_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
+      run: BENCHMARK_LOG=llama_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
    # - name: Run LLaMA 7B on 4 GPUs
-    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_four_gpu.txt
+    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing
    # - name: Run LLaMA 7B on 6 GPUs
-    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_six_gpu.txt
+    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
    - name: Run LLaMA-3 8B BEAM
-      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
+      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
-      run: BENCHMARK_LOG=llama3_beam_4gpu NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
+      run: BENCHMARK_LOG=llama3_beam_4gpu NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
    - name: Run quantized LLaMA3
-      run: BENCHMARK_LOG=llama3_fp8 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --temperature 0 --benchmark --quantize fp8 | tee llama3_fp8.txt
+      run: BENCHMARK_LOG=llama3_fp8 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --temperature 0 --benchmark --quantize fp8
    # - name: Run LLaMA-3 8B on 6 GPUs
-    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
+    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
    # - name: Run LLaMA-2 70B
-    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 MAX_CONTEXT=256 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_2_70B.txt
+    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 MAX_CONTEXT=256 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
    - name: Run Mixtral 8x7B
-      run: time BENCHMARK_LOG=mixtral NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
+      run: time BENCHMARK_LOG=mixtral NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing
    - name: Run GPT2
      run: |
-        BENCHMARK_LOG=gpt2_nojit NV=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        BENCHMARK_LOG=gpt2 NV=1 JIT=1 ASSERT_MIN_STEP_TIME=4 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2_nojit NV=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=gpt2 NV=1 JIT=1 ASSERT_MIN_STEP_TIME=4 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
    - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=6 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=6 python3 examples/gpt2.py --count 10 --temperature 0 --timing
    - 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
+      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
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (NVIDIA)
        path: |
          onnx_inference_speed.csv
-          torch_speed.txt
-          matmul.txt
-          matmul_bfloat16.txt
-          matmul_tf32.txt
-          matmul_ptx.txt
-          matmul_nv.txt
-          sd.txt
-          sdxl.txt
-          llama_unjitted.txt
-          llama_jitted.txt
-          llama_beam.txt
-          llama3_beam.txt
-          llama3_four_gpu.txt
-          llama3_six_gpu.txt
-          llama3_fp8.txt
-          llama_2_70B.txt
-          mixtral.txt
-          gpt2_unjitted.txt
-          gpt2_jitted.txt
-          gpt2_half.txt
-          gpt2_half_beam.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

@@ -337,44 +337,30 @@ jobs:
    # - name: Fuzz Padded Tensor Core GEMM (PTX)
    #   run: NV=1 NV_PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
    - name: HEVC Decode Benchmark
-      run: VALIDATE=1 MAX_FRAMES=100 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
+      run: VALIDATE=1 MAX_FRAMES=100 JITBEAM=1 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
    - name: Train MNIST
-      run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
+      run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
    - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 python3 examples/hlb_cifar10.py
    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=110 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=110 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
    - name: Run 10 CIFAR training steps w BF16
-      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=120 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=120 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
    # - name: Run 10 CIFAR training steps w winograd
-    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
+      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
+      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    - 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
+      run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
    - 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
+      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
    - 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=72 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)
-        path: |
-          beautiful_mnist.txt
-          train_cifar.txt
-          train_cifar_half.txt
-          train_cifar_bf16.txt
-          train_cifar_wino.txt
-          train_cifar_one_gpu.txt
-          train_cifar_six_gpu.txt
-          train_resnet.txt
-          train_resnet_one_gpu.txt
-          train_bert.txt
+      run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    - 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

@@ -389,10 +375,12 @@ jobs:
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
-    - name: Remove amdgpu
-      run: sudo rmmod amdgpu || true
-    - name: Cleanup running AM processes
-      run: python extra/amdpci/am_smi.py --pids --kill
+    - name: Setcap to python
+      run: ./extra/amdpci/setup_python_cap.sh
+    - name: Remove amd modules
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
+    - name: Kill stale pids
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
    #- name: Insert amdgpu
    #  run: sudo modprobe amdgpu
    - name: Symlink models and datasets
@@ -426,16 +414,18 @@ 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/speed/external_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
    - name: Test speed vs theoretical
      run: AMD=1 IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
-    - name: Test tensor cores
-      run: |
-        AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
-        AMD=1 AMD_LLVM=1 python3 test/opt/test_tensor_cores.py
-        AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+    - name: Test tensor cores AMD_LLVM=0
+      run: AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
+    # TODO: this is flaky
+    # - name: Test tensor cores AMD_LLVM=1
+    #   run: AMD=1 AMD_LLVM=1 python3 test/opt/test_tensor_cores.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
+      run: |
+        AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+        AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
    - name: Test AMD=1
      run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
    #- name: Test HIP=1
@@ -450,61 +440,39 @@ jobs:
    - name: Test AM warm start time
      run: time AMD=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
    - name: Run SDXL
-      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing
    - 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
-        BENCHMARK_LOG=llama AMD=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_jitted.txt
+        BENCHMARK_LOG=llama_nojit AMD=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=llama AMD=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
    - name: Run LLaMA 7B with BEAM
-      run: BENCHMARK_LOG=llama_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
+      run: BENCHMARK_LOG=llama_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
    # - name: Run LLaMA 7B on 4 GPUs
-    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_four_gpu.txt
+    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing
    # - name: Run LLaMA 7B on 6 GPUs
-    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_six_gpu.txt
+    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
    - name: Run LLaMA-3 8B BEAM
-      run: BENCHMARK_LOG=llama3_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
+      run: BENCHMARK_LOG=llama3_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
-      run: BENCHMARK_LOG=llama3_beam_4gpu AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
+      run: BENCHMARK_LOG=llama3_beam_4gpu AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
    # - name: Run LLaMA-3 8B on 6 GPUs
-    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
+    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
    #- name: Restore amdgpu
    #  run: sudo modprobe amdgpu
    # - name: Run LLaMA-2 70B
-    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_2_70B.txt
+    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
    - name: Run Mixtral 8x7B
-      run: time BENCHMARK_LOG=mixtral AMD=1 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
+      run: time BENCHMARK_LOG=mixtral AMD=1 python3 examples/mixtral.py --temperature 0 --count 10 --timing
    - name: Run GPT2
      run: |
-        BENCHMARK_LOG=gpt2_nojit AMD=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        BENCHMARK_LOG=gpt2 AMD=1 JIT=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2_nojit AMD=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=gpt2 AMD=1 JIT=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
    - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half AMD=1 HALF=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      run: BENCHMARK_LOG=gpt2_half AMD=1 HALF=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --count 10 --temperature 0 --timing
    - 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
-      with:
-        name: Speed (AMD)
-        path: |
-          onnx_inference_speed.csv
-          torch_speed.txt
-          llama_unjitted.txt
-          llama_jitted.txt
-          llama_beam.txt
-          llama3_beam.txt
-          llama3_four_gpu.txt
-          llama3_six_gpu.txt
-          llama_2_70B.txt
-          gpt2_unjitted.txt
-          gpt2_jitted.txt
-          gpt2_half.txt
-          gpt2_half_beam.txt
-          matmul.txt
-          matmul_amd.txt
-          sd.txt
-          sdxl.txt
-          mixtral.txt
+      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
    - 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

@@ -519,10 +487,12 @@ jobs:
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
-    - name: Remove amdgpu
-      run: sudo rmmod amdgpu || true
-    - name: Cleanup running AM processes
-      run: python extra/amdpci/am_smi.py --pids --kill
+    - name: Setcap to python
+      run: ./extra/amdpci/setup_python_cap.sh
+    - name: Remove amd modules
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
+    - name: Kill stale pids
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -541,31 +511,22 @@ jobs:
    - name: reset process replay
      run: test/external/process_replay/reset.py
    - name: Train MNIST
-      run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
+      run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
    - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py
    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=200 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=200 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
    # - name: Run 10 CIFAR training steps w BF16
-    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
    # TODO: too slow
    # - name: Run 10 CIFAR training steps w winograd
-    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
-    - uses: actions/upload-artifact@v4
-      with:
-        name: Speed (AMD Training)
-        path: |
-          beautiful_mnist.txt
-          train_cifar.txt
-          train_cifar_half.txt
-          train_cifar_bf16.txt
-          train_cifar_wino.txt
-          train_cifar_one_gpu.txt
-          train_cifar_six_gpu.txt
+      run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    - name: Test full tinyfs load
+      run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
    - 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

@@ -580,10 +541,12 @@ jobs:
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
-    - name: Remove amdgpu
-      run: sudo rmmod amdgpu || true
-    - name: Cleanup running AM processes
-      run: python extra/amdpci/am_smi.py --pids --kill
+    - name: Setcap to python
+      run: ./extra/amdpci/setup_python_cap.sh
+    - name: Remove amd modules
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
+    - name: Kill stale pids
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -604,19 +567,12 @@ jobs:
    - 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
+      run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
    - 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
+      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
    - 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=72 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)
-        path: |
-          train_resnet.txt
-          train_resnet_one_gpu.txt
-          train_bert.txt
+      run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    - 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

@@ -639,19 +595,21 @@ jobs:
    - name: reset process replay
      run: test/external/process_replay/reset.py
    - name: openpilot compile3 0.10.0 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
    - name: openpilot compile3 0.10.0 dmonitoring
      run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: DEBUG=2 openpilot compile3 0.10.1 driving_vision
      run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
    - name: DEBUG=2 IMAGE=1 openpilot compile3 0.10.1 driving_vision
      run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+    - name: IMAGE=1 openpilot compile3 0.10.1 driving_vision
+      run: BENCHMARK_LOG=image_1_openpilot_0_10_1_vision PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.10.1 driving_vision
      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.10.1 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
    - name: openpilot compile3 0.10.1 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=10 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: benchmark MobileNetV2 on DSP
      run: |
        # generate quantized weights
@@ -674,10 +632,12 @@ jobs:
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
+    - name: Setcap to python
+      run: ./extra/amdpci/setup_python_cap.sh
    - name: Remove amd modules
-      run: ./extra/hcq/hcq_smi.py amd rmmod
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
    - name: Kill stale pids
-      run: ./extra/hcq/hcq_smi.py amd kill_pids
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -706,7 +666,7 @@ jobs:
    #     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
+      run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
    - name: Test AMD=1
      run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
    - name: Test DISK copy time
@@ -716,20 +676,12 @@ jobs:
        AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
        AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    # - 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
+    #   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
    - 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
+      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
    - 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

@@ -744,10 +696,12 @@ jobs:
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
+    - name: Setcap to python
+      run: ./extra/amdpci/setup_python_cap.sh
    - name: Remove nv modules
-      run: ./extra/hcq/hcq_smi.py nv rmmod
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py nv rmmod
    - name: Kill stale pids
-      run: ./extra/hcq/hcq_smi.py nv kill_pids
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -776,21 +730,13 @@ jobs:
        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
+      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
+      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
+      run: 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
    - 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
+      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
    - 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
@@ -1,10 +1,11 @@
 name: Unit Tests
 env:
  # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '15'
+  CACHE_VERSION: '16'
  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
  PYTHONPATH: ${{ github.workspace }}
+  CHECK_OOB: 1

 on:
  push:
@@ -25,17 +26,19 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: llvm-speed
-        deps: testing_minimal
+        deps: testing_unit
        llvm: 'true'
    - name: Speed Test
-      run: CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py
+      run: CPU=1 CPU_LLVM=1 THREADS=0 python3 test/speed/external_test_speed_v_torch.py
    - name: Speed Test (BEAM=2)
-      run: BEAM=2 CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py
+      run: BEAM=2 CPU=1 CPU_LLVM=1 THREADS=0 python3 test/speed/external_test_speed_v_torch.py

  docs:
    name: Docs
    runs-on: ubuntu-22.04
    timeout-minutes: 10
+    env:
+      CHECK_OOB: 0
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -95,27 +98,23 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: torch-backend-pillow-torchvision-et-pt
-        deps: testing_minimal
+        deps: testing_unit
        pydeps: "pillow torchvision expecttest"
        llvm: 'true'
    - name: Install ninja
      run: |
        sudo apt update || true
        sudo apt install -y --no-install-recommends ninja-build
-    - name: Lint with ruff
-      run: |
-        pip3 install --upgrade --force-reinstall ruff==0.11.0
-        python3 -m ruff check extra/torch_backend/backend.py
    - name: Test one op
-      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
+      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Test ResNet-18
      run: DEBUG=2 python3 extra/torch_backend/example.py
-    - name: My (custom) tests
-      run: python3 extra/torch_backend/test.py
+    - name: custom tests
+      run: python3 -m pytest -n auto extra/torch_backend/test.py --durations=20
    - name: Test one op in torch tests
      run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
    - name: Test Ops with TINY_BACKEND
-      run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
+      run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/backend/test_ops.py --durations=20
    - name: Test in-place operations on views
      run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
    - name: Test multi-gpu
@@ -135,7 +134,7 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: torch-backend-pillow-torchvision-et-pt
-        deps: testing_minimal
+        deps: testing_unit
        llvm: 'true'
    - name: Install ninja
      run: |
@@ -157,27 +156,27 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: be-minimal
-        deps: testing_minimal
+        deps: testing_unit
    - name: Test dtype with Python emulator
-      run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
+      run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/backend/test_dtype.py test/backend/test_dtype_alu.py
    - name: Test ops with Python emulator
-      run: DEBUG=2 SKIP_SLOW_TEST=1 PYTHON=1 python3 -m pytest -n=auto test/test_ops.py --durations=20
+      run: DEBUG=2 SKIP_SLOW_TEST=1 PYTHON=1 python3 -m pytest -n=auto test/backend/test_ops.py --durations=20
    - name: Test uops with Python emulator
-      run: PYTHON=1 python3 -m pytest test/test_uops.py --durations=20
+      run: PYTHON=1 python3 -m pytest test/backend/test_uops.py --durations=20
    - name: Test symbolic with Python emulator
-      run: PYTHON=1 python3 test/test_symbolic_ops.py
+      run: PYTHON=1 python3 test/backend/test_symbolic_ops.py
    - name: test_renderer_failures with Python emulator
-      run: PYTHON=1 python3 -m pytest -rA test/test_renderer_failures.py::TestRendererFailures
+      run: PYTHON=1 python3 -m pytest -rA test/backend/test_renderer_failures.py::TestRendererFailures
    - name: Test IMAGE=2 support
      run: |
-        IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
-        IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_simple_conv2d
+        IMAGE=2 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
+        IMAGE=2 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_simple_conv2d
    - name: Test emulated METAL tensor cores
      run: |
-        DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_big_gemm
+        DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_big_gemm
        DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
    - name: Test emulated AMX tensor cores
-      run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
+      run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
    - name: Test emulated AMD tensor cores
      run: |
        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
@@ -198,9 +197,9 @@ jobs:
        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
    - name: Test emulated CUDA tensor cores
      run: |
-        DEBUG=2 EMULATE=CUDA FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
-        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
-        DEBUG=2 EMULATE=CUDA_SM75 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
+        DEBUG=2 EMULATE=CUDA FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm_fp16
+        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
+        DEBUG=2 EMULATE=CUDA_SM75 FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm_fp16
        DEBUG=2 EMULATE=CUDA_SM89 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
    - name: Test emulated INTEL OpenCL tensor cores
      run: DEBUG=2 EMULATE=INTEL FORWARD_ONLY=1 PYTHON=1 HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
@@ -208,18 +207,17 @@ jobs:
      run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
    - name: Test device flop counts
      run: |
-        DEBUG=2 EMULATE=METAL PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 EMULATE=AMD PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 EMULATE=CUDA PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 EMULATE=INTEL PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 AMX=1 EMULATE=AMX PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStats.test_simple_matmul
+        DEBUG=2 EMULATE=METAL PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=AMD PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=CUDA PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=INTEL PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 AMX=1 EMULATE=AMX PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStats.test_simple_matmul

  linter:
    name: Linters
    runs-on: ubuntu-latest
    timeout-minutes: 10

-    # TODO: run the pre-commit hook to replace a lot of this
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -231,19 +229,20 @@ jobs:
        deps: linting
    - name: Lint bad-indentation and trailing-whitespace with pylint
      run: python -m pylint --disable=all -e W0311 -e C0303 --jobs=0 --indent-string='  ' --recursive=y .
-    - name: Lint with ruff
+    - name: Run pre-commit linting hooks
+      run: SKIP=tiny,tests,example pre-commit run --all-files
+    - name: Lint additional files with ruff
      run: |
-        pip3 install --upgrade --force-reinstall ruff==0.11.0
-        python3 -m ruff check .
        python3 -m ruff check examples/mlperf/ --ignore E501
        python3 -m ruff check extra/thunder/tiny/ --ignore E501 --ignore F841 --ignore E722
-    - name: Run mypy
+        python3 -m ruff check extra/torch_backend/backend.py
+    - name: Run mypy with lineprecision report
      run: |
-        python -m mypy --strict-equality --lineprecision-report .
+        python -m mypy --lineprecision-report .
+        grep -v autogen lineprecision.txt | awk 'NR>2 {lines+=$2; precise+=$3; imprecise+=$4; any+=$5; empty+=$6} END {t=lines-empty; printf "TOTAL: %d lines, %d precise (%.1f%%), %d imprecise (%.1f%%), %d any (%.1f%%)\n", t, precise, 100*precise/t, imprecise, 100*imprecise/t, any, 100*any/t}'
        cat lineprecision.txt
-    # broken because of UPatAny
-    #- name: Run TYPED=1
-    #  run: TYPED=1 python -c "import tinygrad"
+    - name: Run TYPED=1
+      run: CHECK_OOB=0 DEV=CPU TYPED=1 python test/test_tiny.py

  unittest:
    name: Unit Tests
@@ -256,24 +255,30 @@ jobs:
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
-        key: unittest-12
-        pydeps: "pillow numpy ftfy regex"
+        key: unittest-13
+        pydeps: "pillow ftfy regex pre-commit"
        deps: testing_unit
+        llvm: 'true'
+        amd: 'true'
+    - name: Run pre-commit test hooks
+      run: SKIP=ruff,mypy pre-commit run --all-files
    - name: Check Device.DEFAULT
      run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
    - name: Run unit tests
      run: |
-        CPU=1 python test/unit/test_device.py TestRunAsModule.test_module_runs
-        CPU=1 python -m pytest -n=auto test/unit/ --durations=20 --deselect=test/unit/test_device.py::TestRunAsModule::test_module_runs
+        CPU=1 python test/null/test_device.py TestRunAsModule.test_module_runs
+        CPU=1 python -m pytest -n=auto test/unit/ --durations=20
+    - name: Run NULL backend tests
+      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
    - name: Run targetted tests on NULL backend
-      run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
+      run: NULL=1 python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
    # TODO: too slow
    # - name: Run SDXL on NULL backend
    #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
    - name: Run Clip tests for SD MLPerf on NULL backend
      run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
    - name: Run AMD emulated BERT training on NULL backend
-      run: EMULATE=AMD_RDNA4 NULL=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+      run: EMULATE=AMD_RDNA4 NULL=1 NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    # TODO: support fake weights
    #- name: Run LLaMA 7B on 4 fake devices
    #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
@@ -290,8 +295,8 @@ jobs:
        python extra/optimization/extract_dataset.py
        gzip -c /tmp/sops > extra/datasets/sops.gz
        #DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py
-    - name: Repo line count < 20000 lines
-      run: MAX_LINE_COUNT=20000 python sz.py
+    - name: Repo line count < 24000 lines
+      run: MAX_LINE_COUNT=24000 python sz.py

  spec:
    strategy:
@@ -311,7 +316,7 @@ jobs:
        deps: testing_unit
        python-version: '3.14'
    - name: Test SPEC=2
-      run: IGNORE_OOB=0 SPEC=2 PYTHONPATH="." pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
+      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 test/unit test/backend test/opt --ignore test/backend/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}

  fuzzing:
    name: Fuzzing
@@ -345,11 +350,11 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: gpu-image
-          deps: testing_minimal
+          deps: testing_unit
          opencl: 'true'
      - name: Test CL IMAGE=2 ops
        run: |
-          CL=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20
+          CL=1 IMAGE=2 python -m pytest -n=auto test/backend/test_ops.py --durations=20
        # TODO: training is broken
        # CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
      - name: Run process replay tests
@@ -366,14 +371,14 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: gen-dataset
-          deps: testing_minimal
+          deps: testing
          opencl: 'true'
      - name: Generate Dataset
        run: CL=1 extra/optimization/generate_dataset.sh
      - name: Run Kernel Count Test
        run: CL=1 python -m pytest -n=auto test/external/external_test_opt.py
      - name: Run fused optimizer tests
-        run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py test/test_optim.py -k "not muon"
+        run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py test/backend/test_optim.py -k "not muon"
      - name: Upload artifact
        uses: actions/upload-artifact@v4
        with:
@@ -421,7 +426,7 @@ jobs:
        with:
          key: onnxoptc
          deps: testing
-          python-version: '3.11'
+          python-version: '3.12'
          llvm: 'true'
      - name: Test ONNX (CPU)
        run: CPU=1 CPU_LLVM=0 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
@@ -432,7 +437,7 @@ jobs:
      - name: Test Additional ONNX Ops (CPU)
        run: CPU=1 CPU_LLVM=0 python3 test/external/external_test_onnx_ops.py
      - name: Test Quantize ONNX
-        run: CPU=1 CPU_LLVM=0 python3 test/test_quantize_onnx.py
+        run: CPU=1 CPU_LLVM=0 python3 test/backend/test_quantize_onnx.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -449,7 +454,7 @@ jobs:
          key: onnxoptl
          deps: testing
          pydeps: "tensorflow==2.19"
-          python-version: '3.11'
+          python-version: '3.12'
          opencl: 'true'
      - name: Test ONNX (CL)
        run: CL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
@@ -462,11 +467,11 @@ jobs:
      - name: Test MLPerf stuff
        run: CL=1 python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
      - name: NULL=1 beautiful_mnist_multigpu
-        run: NULL=1 python examples/beautiful_mnist_multigpu.py
+        run: NULL=1 NULL_ALLOW_COPYOUT=1 python examples/beautiful_mnist_multigpu.py
      - name: Test Bert training
-        run: NULL=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+        run: NULL=1 NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
      - name: Test llama 3 training
-        run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
+        run: NULL=1 NULL_ALLOW_COPYOUT=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -474,6 +479,8 @@ jobs:
    name: Test LLM
    runs-on: ubuntu-24.04
    timeout-minutes: 15
+    env:
+      CHECK_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -521,7 +528,7 @@ jobs:
      with:
        key: metal
        deps: testing
-        python-version: '3.11'
+        python-version: '3.12'
    - name: Test models (Metal)
      run: METAL=1 python -m pytest -n=auto test/models --durations=20
    - name: Test LLaMA compile speed
@@ -540,15 +547,15 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: devectorize-minimal
-        deps: testing_minimal
+        deps: testing_unit
        pydeps: "pillow"
        llvm: "true"
    - name: Test LLVM=1 DEVECTORIZE=0
-      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py
+      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/backend/test_ops.py
    - name: Test LLVM=1 DEVECTORIZE=0 for model
      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
    - name: Test CPU=1 DEVECTORIZE=0
-      run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py
+      run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/backend/test_ops.py

  testdsp:
    name: Linux (DSP)
@@ -561,8 +568,8 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: dsp-minimal
-        deps: testing_minimal
-        pydeps: "onnx==1.18.0 onnxruntime pillow"
+        deps: testing_unit
+        pydeps: "onnx==1.18.0 onnxruntime ml_dtypes"
        llvm: "true"
    - name: Set up Docker Buildx
      uses: docker/setup-buildx-action@v3
@@ -574,15 +581,15 @@ jobs:
        load: true
        tags: qemu-hexagon:latest
        cache-from: type=gha
-        cache-to: type=gha,mode=min
+        cache-to: ${{ github.event_name != 'pull_request' && 'type=gha,mode=min' || '' }}
    - name: Set MOCKDSP env
      run: printf "MOCKDSP=1" >> $GITHUB_ENV
    - name: Run test_tiny on DSP
      run: DEBUG=2 DSP=1 python test/test_tiny.py
    - name: Test transcendentals
-      run: CC=clang-20 DEBUG=2 DSP=1 python test/test_transcendental.py TestTranscendentalVectorized
+      run: CC=clang-20 DEBUG=2 DSP=1 python test/backend/test_transcendental.py TestTranscendentalVectorized
    - name: Test quantize onnx
-      run: DEBUG=2 DSP=1 python3 test/test_quantize_onnx.py
+      run: DEBUG=2 DSP=1 python3 test/backend/test_quantize_onnx.py

  testwebgpu:
    name: Linux (WebGPU)
@@ -595,18 +602,16 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: webgpu-minimal
-        deps: testing_minimal
-        python-version: '3.11'
+        deps: testing_unit
+        python-version: '3.12'
        webgpu: 'true'
    - name: Check Device.DEFAULT (WEBGPU) and print some source
      run: |
        WEBGPU=1 python -c "from tinygrad import Device; assert Device.DEFAULT == 'WEBGPU', Device.DEFAULT"
-        WEBGPU=1 DEBUG=4 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+        WEBGPU=1 DEBUG=4 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Run selected webgpu tests
      run: |
-          WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit \
-          --ignore=test/test_copy_speed.py --ignore=test/test_rearrange_einops.py \
-          --ignore=test/test_fuzz_shape_ops.py --durations=20
+          WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/backend --durations=20
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -631,30 +636,81 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: ${{ matrix.backend }}-minimal
-          deps: testing_minimal
+          deps: testing_unit
          amd: 'true'
          llvm: ${{ matrix.backend == 'amdllvm' && 'true' }}
      - name: Check Device.DEFAULT and print some source
        run: |
          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['AMD'], Device.DEFAULT"
-          DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+          DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
      - name: Run LLVM test
        if: matrix.backend=='amdllvm'
        run: python test/device/test_amd_llvm.py
      - name: Run pytest (amd)
-        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/device/test_hcq.py --durations=20
+        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/testextra/test_cfg_viz.py --durations=20
      - name: Run pytest (amd)
        run: python -m pytest test/external/external_test_am.py --durations=20
      - name: Run TRANSCENDENTAL math
-        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
+        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
      - name: Run TestOps.test_add with SQTT
        run: |
-          VIZ=1 PMC=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
-          VIZ=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
+          VIZ=-2 DEBUG=5 python3 test/backend/test_ops.py TestOps.test_add
          extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
+      - name: Run AMD emulated mmapeak on NULL backend
+        env:
+          AMD: 0
+        run: PYTHONPATH=. NULL=1 EMULATE=AMD python extra/mmapeak/mmapeak.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

+  testamdasm:
+    name: AMD ASM IDE
+    runs-on: ubuntu-24.04
+    timeout-minutes: 20
+    env:
+      AMD: 1
+      PYTHON_REMU: 1
+      MOCKGPU: 1
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: rdna3-emu
+          deps: testing_unit
+          amd: 'true'
+          python-version: '3.14'
+      - name: Verify AMD autogen is up to date
+        run: |
+          python -m tinygrad.renderer.amd.generate
+          git diff --exit-code tinygrad/runtime/autogen/amd/
+      - name: Install LLVM 21
+        run: |
+          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)-21 main" | sudo tee /etc/apt/sources.list.d/llvm.list
+          sudo apt-get update
+          sudo apt-get install llvm-21 llvm-21-tools cloc
+      - name: RDNA3 Line Count
+        run: cloc --by-file tinygrad/renderer/amd/*.py
+      - name: Install rocprof-trace-decoder
+        run: sudo PYTHONPATH="." ./extra/sqtt/install_sqtt_decoder.py
+      - name: Run RDNA3 emulator tests
+        run: AMD_LLVM=0 python -m pytest -n=auto test/amd/ --durations 20
+      - name: Run RDNA3 emulator tests (AMD_LLVM=1)
+        run: AMD_LLVM=1 python -m pytest -n=auto test/amd/ --durations 20
+      - name: Run RDNA3 dtype tests
+        run: AMD_LLVM=0 pytest -n=auto test/backend/test_dtype_alu.py test/backend/test_dtype.py --durations 20
+      - name: Run RDNA3 dtype tests (AMD_LLVM=1)
+        run: AMD_LLVM=1 pytest -n=auto test/backend/test_dtype_alu.py test/backend/test_dtype.py --durations 20
+      # TODO: run all once emulator is faster
+      - name: Run RDNA3 ops tests
+        run: SKIP_SLOW_TEST=1 AMD_LLVM=0 pytest -n=auto test/backend/test_ops.py -k "test_sparse_categorical_crossentropy or test_tril or test_nonzero or test_softmax_argmax" --durations 20
+      - name: Run RDNA4 emulator tests
+        run: MOCKGPU_ARCH=rdna4 python -m pytest test/test_tiny.py -v --durations 20
+      - name: Run CDNA4 emulator tests
+        run: AMD_LLVM=1 MOCKGPU_ARCH=cdna4 python -m pytest test/test_tiny.py -v --durations 20
+
  testnvidia:
    strategy:
      fail-fast: false
@@ -674,7 +730,7 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: ${{ matrix.backend }}-minimal
-          deps: testing_minimal
+          deps: testing_unit
          cuda: 'true'
          ocelot: 'true'
      - name: Set env
@@ -682,10 +738,12 @@ jobs:
      - name: Check Device.DEFAULT and print some source
        run: |
          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CUDA','NV'], Device.DEFAULT"
-          DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+          DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
      - name: Run pytest (cuda)
        # skip multitensor because it's slow
-        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore test/test_gc.py --ignore test/test_multitensor.py --durations=20
+        run: python -m pytest -n=auto test/backend --ignore test/backend/test_multitensor.py --durations=20
+      - name: Run TestOps.test_add with PMA
+        run: VIZ=-1 PMA=1 DEBUG=5 python3 test/backend/test_ops.py TestOps.test_add
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -705,7 +763,7 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: ${{ matrix.backend }}-minimal
-          deps: testing_minimal
+          deps: testing_unit
          opencl: ${{ matrix.backend == 'opencl' && 'true' }}
          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
          mesa: ${{ matrix.backend == 'lvp' && 'true' }}
@@ -714,11 +772,11 @@ jobs:
      - name: Check Device.DEFAULT and print some source
        run: |
          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT"
-          DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+          DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
      - name: Run pytest (${{ matrix.backend }})
-        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
+        run: python -m pytest -n=auto test/backend --durations=20
      - name: Run TRANSCENDENTAL math
-        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
+        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -736,25 +794,27 @@ jobs:
      with:
        key: metal
        deps: testing
-        python-version: '3.11'
+        python-version: '3.12'
        amd: 'true'
        cuda: 'true'
        ocelot: 'true'
        llvm: 'true'
    - name: Run unit tests
      run: METAL=1 python -m pytest -n=auto test/unit/ --durations=20
+    - name: Run NULL backend tests
+      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
    - name: Run ONNX
      run: METAL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
    - name: Test tensor core ops (fake)
-      run: METAL=1 DEBUG=3 TC=2 python test/test_ops.py TestOps.test_gemm
+      run: METAL=1 DEBUG=3 TC=2 python test/backend/test_ops.py TestOps.test_gemm
    - name: Test tensor core ops (real)
-      run: METAL=1 DEBUG=3 python test/test_ops.py TestOps.test_big_gemm
+      run: METAL=1 DEBUG=3 python test/backend/test_ops.py TestOps.test_big_gemm
    - name: Test Beam Search
      run: METAL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
    #- name: Fuzz Test linearizer
    #  run: METAL=1 DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
    - name: Run TRANSCENDENTAL math
-      run: METAL=1 TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
+      run: METAL=1 TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
    - name: Run pytest (amd)
      env:
        MOCKGPU: 1
@@ -795,14 +855,14 @@ jobs:
        key: osx-webgpu
        deps: testing
        webgpu: 'true'
-    - name: Test infinity math in WGSL
-      run: WEBGPU=1 python -m pytest -n=auto test/test_renderer_failures.py::TestWGSLFailures::test_multiply_infinity --durations=20
    - name: Build WEBGPU Efficientnet
      run: WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m examples.compile_efficientnet
-    - name: Clean npm cache
-      run: npm cache clean --force
-    - name: Install Puppeteer
-      run: npm install puppeteer
+    - name: Run selected webgpu tests
+      run: WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m pytest -n=auto test/backend --durations=20
+    #- name: Clean npm cache
+    #  run: npm cache clean --force
+    #- name: Install Puppeteer
+    #  run: npm install puppeteer
    # this is also flaky
    #- name: Run WEBGPU Efficientnet
    #  run: node test/web/test_webgpu.js
@@ -834,8 +894,7 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: macos-${{ matrix.backend }}-minimal
-          deps: testing_minimal
-          pydeps: "capstone"
+          deps: testing_unit
          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
          mesa: ${{ matrix.backend == 'lvp' && 'true' }}
      - name: Set env
@@ -845,7 +904,7 @@ jobs:
          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU','LVP':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
      - name: Run pytest (${{ matrix.backend }})
-        run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
+        run: python3 -m pytest -n=auto test/backend --durations=20
      - name: Run process replay tests
        uses: ./.github/actions/process-replay
      - name: Run macOS-specific unit test
@@ -878,12 +937,16 @@ jobs:
      - name: Run unit tests
        if: matrix.backend=='llvm'
        # test_newton_schulz hits RecursionError
-        run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_elf.py --ignore=test/unit/test_tar.py --ignore=test/unit/test_linalg.py --durations=20
+        run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_tar.py --ignore=test/unit/test_linalg.py --durations=20
+      - name: Run NULL backend tests
+        if: matrix.backend=='llvm'
+        shell: bash
+        run: CPU=0 CPU_LLVM=0 NULL=1 python -m pytest -n=auto test/null/ --ignore=test/null/test_elf.py --durations=20
      - name: Run pytest (${{ matrix.backend }})
        shell: bash
        run: |
          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
-          python -m pytest -n=auto test/test_tiny.py test/test_ops.py --durations=20
+          python -m pytest -n=auto test/test_tiny.py test/backend/test_ops.py --durations=20

 # ****** Compile-only Tests ******

@@ -902,15 +965,15 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: compile-${{ matrix.backend }}
-          deps: testing_minimal
+          deps: testing_unit
          mesa: ${{ (matrix.backend == 'ir3' || matrix.backend == 'nak') && 'true' }}
-          python-version: '3.14'
+          python-version: '3.12'
      - name: Set env
        shell: bash
-        run: printf "NULL=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
+        run: printf "NULL=1\nNULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
      - name: Run test_ops
        shell: bash
        run: |
          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
-          DEBUG=4 python3 test/test_ops.py TestOps.test_add
-          python -m pytest -n=auto test/test_ops.py --durations=20
+          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20
@@ -58,10 +58,11 @@ weights
 *.lprof
 comgr_*
 *.pkl
+!extra/sqtt/examples/**/*.pkl
 site/
 profile_stats
 *.log
 target
 .mypy_cache
 mutants
-.mutmut-cache
+.mutmut-cache
@@ -16,7 +16,7 @@ repos:
        pass_filenames: false
      - id: mypy
        name: mypy
-        entry: python3 -m mypy tinygrad/ --strict-equality
+        entry: python3 -m mypy
        language: system
        always_run: true
        pass_filenames: false
@@ -28,7 +28,7 @@ repos:
        pass_filenames: false
      - id: tests
        name: comprehensive test suite
-        entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_schedule.py test/test_assign.py test/test_tensor.py test/test_jit.py test/unit/test_schedule_cache.py test/unit/test_pattern_matcher.py test/unit/test_uop_symbolic.py test/unit/test_helpers.py
+        entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/backend/test_ops.py test/backend/test_schedule.py test/unit/test_assign.py test/backend/test_tensor.py test/backend/test_jit.py test/unit/test_schedule_cache.py test/null/test_pattern_matcher.py test/null/test_uop_symbolic.py test/unit/test_helpers.py
        language: system
        always_run: true
        pass_filenames: false
@@ -34,33 +34,6 @@ result = graph_rewrite(uop, pm)
 ### Schedule Cache
 Schedules are cached by graph structure. BIND nodes (variables with bound values) are unbound before cache key computation so different values hit the same cache.

-## Directory Structure
-
-```
-tinygrad/
-├── tensor.py          # Tensor class, user API
-├── device.py          # Buffer, device management
-├── dtype.py           # Data types
-├── helpers.py         # Utilities, environment vars
-├── uop/
-│   ├── ops.py         # UOp class, Ops enum, PatternMatcher
-│   ├── spec.py        # UOp type verification
-│   └── symbolic.py    # Symbolic math simplification
-├── engine/
-│   ├── schedule.py    # Schedule creation, caching
-│   ├── realize.py     # Tensor realization
-│   ├── jit.py         # JIT compilation
-│   └── memory.py      # Memory planning
-├── schedule/
-│   ├── rangeify.py    # Convert movements to ranges
-│   └── indexing.py    # Index calculations
-├── codegen/
-│   ├── kernel.py      # Kernel optimization
-│   └── uopgraph.py    # UOp graph transformations
-├── renderer/          # Code generation (CUDA, Metal, etc.)
-└── runtime/           # Device backends
-```
-
 ## Testing

 ```bash
@@ -68,10 +41,10 @@ tinygrad/
 python -m pytest test/unit/test_schedule_cache.py -xvs

 # Run with timeout
-python -m pytest test/test_symbolic_ops.py -x --timeout=60
+python -m pytest test/backend/test_symbolic_ops.py -x --timeout=60

 # Debug with print
-DEBUG=2 python -m pytest test/test_schedule.py::test_name -xvs
+DEBUG=2 python -m pytest test/backend/test_schedule.py::test_name -xvs

 # Visualize UOp graphs
 VIZ=1 python -c "from tinygrad import Tensor; Tensor.ones(10).sum().realize()"
@@ -79,7 +52,7 @@ VIZ=1 python -c "from tinygrad import Tensor; Tensor.ones(10).sum().realize()"

 ## Common Environment Variables

- `DEBUG=1-4` - Increasing verbosity
+- `DEBUG=1-7` - Increasing verbosity (7 shows assembly output)
 - `VIZ=1` - Enable graph visualization
 - `SPEC=1` - Enable UOp spec verification
 - `NOOPT=1` - Disable optimizations
@@ -100,6 +73,16 @@ VIZ=1 python -c "from tinygrad import Tensor; Tensor.ones(10).sum().realize()"
 - Run tests before proposing commits
 - Test with `SPEC=2` when modifying UOp-related code

+## Auto-generated Files (DO NOT EDIT)
+
+The following files are auto-generated and should never be edited manually:
+- `tinygrad/runtime/autogen/amd/{arch}/__init__.py` - Generated by `python -m tinygrad.renderer.amd.dsl --arch {arch}`
+- `tinygrad/runtime/autogen/amd/{arch}/gen_pcode.py` - Generated by `python -m tinygrad.renderer.amd.pcode --arch {arch}`
+
+Where `{arch}` is one of: `rdna3`, `rdna4`, `cdna`
+
+To add missing instruction implementations, add them to `tinygrad/renderer/amd/emu.py` instead.
+
 ## Style Notes

 - 2-space indentation, 150 char line limit
@@ -209,9 +192,12 @@ When optimizing tinygrad internals:

 9. **Avoid creating intermediate objects in hot paths** - For example, `any(x.op in ops for x in self.backward_slice)` is faster than `any(x.op in ops for x in {self:None, **self.backward_slice})` because it avoids dict creation.

-## Pattern Matching Profiling
+## Pattern Matching Analysis

-Use `TRACK_MATCH_STATS=2` to identify expensive patterns:
+**Use the right tool:**
+
+- `TRACK_MATCH_STATS=2` - **Profiling**: identify expensive patterns
+- `VIZ=-1` - **Inspection**: see all transformations, what every match pattern does, the before/after diffs

 ```bash
 TRACK_MATCH_STATS=2 PYTHONPATH="." python3 test/external/external_benchmark_schedule.py
@@ -225,3 +211,17 @@ Key patterns to watch (from ResNet50 benchmark):
 - `vmin==vmax folding`: ~55ms, 0.33% match rate - checks 52K ops but rarely matches

 Patterns with 0% match rate are workload-specific overhead. They may be useful in other workloads, so don't remove them without understanding their purpose.
+
+```bash
+# Save the trace
+VIZ=-1 python test/test_tiny.py TestTiny.test_gemm
+
+# Explore it
+./extra/viz/cli.py --help
+```
+
+## AMD Performance Counter Profiling
+
+Set VIZ to `-2` to save performance counters traces for the AMD backend.
+
+Use the CLI in `./extra/sqtt/roc.py` to explore the trace.
@@ -192,7 +192,7 @@ For more examples on how to run the full test suite please refer to the [CI work
 Some examples of running tests locally:
 ```sh
 python3 -m pip install -e '.[testing]'  # install extra deps for testing
-python3 test/test_ops.py                # just the ops tests
+python3 test/backend/test_ops.py        # just the ops tests
 python3 -m pytest test/                 # whole test suite
 ```

@@ -28,7 +28,7 @@ Transforms the ast into an optimized ast. This is where BEAM search and heuristi

 Transform the optimized ast into a linearized and rendered program.

-::: tinygrad.codegen.full_rewrite_to_program
+::: tinygrad.codegen.get_program
    options:
        members: false
        show_labels: false
@@ -70,7 +70,7 @@ AMD backend supports several interfaces for communicating with devices:

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

 You can force an interface by setting `AMD_IFACE` to one of these values. In the case of `AMD_IFACE=PCI`, this may unbind your GPU from the amdgpu driver.

@@ -6,6 +6,7 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.neg
 ::: tinygrad.Tensor.log
 ::: tinygrad.Tensor.log2
+::: tinygrad.Tensor.log10
 ::: tinygrad.Tensor.exp
 ::: tinygrad.Tensor.exp2
 ::: tinygrad.Tensor.sqrt
@@ -87,4 +88,8 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.float
 ::: tinygrad.Tensor.half
 ::: tinygrad.Tensor.int
-::: tinygrad.Tensor.bool
+::: tinygrad.Tensor.bool
+::: tinygrad.Tensor.bfloat16
+::: tinygrad.Tensor.double
+::: tinygrad.Tensor.long
+::: tinygrad.Tensor.short
@@ -27,5 +27,6 @@
 ::: tinygrad.Tensor.flatten
 ::: tinygrad.Tensor.unflatten
 ::: tinygrad.Tensor.diag
+::: tinygrad.Tensor.diagonal
 ::: tinygrad.Tensor.roll
 ::: tinygrad.Tensor.rearrange
@@ -7,6 +7,7 @@
 ::: tinygrad.Tensor.any
 ::: tinygrad.Tensor.all
 ::: tinygrad.Tensor.isclose
+::: tinygrad.Tensor.allclose
 ::: tinygrad.Tensor.mean
 ::: tinygrad.Tensor.var
 ::: tinygrad.Tensor.var_mean
@@ -30,7 +31,9 @@
 ::: tinygrad.Tensor.matmul
 ::: tinygrad.Tensor.einsum
 ::: tinygrad.Tensor.cumsum
+::: tinygrad.Tensor.cumprod
 ::: tinygrad.Tensor.cummax
+::: tinygrad.Tensor.cummin
 ::: tinygrad.Tensor.triu
 ::: tinygrad.Tensor.tril
 ::: tinygrad.Tensor.interpolate
@@ -38,7 +41,9 @@
 ::: tinygrad.Tensor.scatter_reduce
 ::: tinygrad.Tensor.masked_select
 ::: tinygrad.Tensor.masked_fill
+::: tinygrad.Tensor.nonzero
 ::: tinygrad.Tensor.sort
+::: tinygrad.Tensor.argsort
 ::: tinygrad.Tensor.topk
 ::: tinygrad.Tensor.multinomial

@@ -56,3 +61,8 @@
 ::: tinygrad.Tensor.sparse_categorical_crossentropy
 ::: tinygrad.Tensor.cross_entropy
 ::: tinygrad.Tensor.nll_loss
+
+## Linear Algebra
+
+::: tinygrad.Tensor.qr
+::: tinygrad.Tensor.svd
@@ -0,0 +1,196 @@
+from tinygrad import Tensor, dtypes, Context, getenv, UOp, fetch
+from tinygrad.uop.ops import Ops, PatternMatcher, UPat
+from tinygrad.uop.symbolic import symbolic
+from tinygrad.codegen import Renderer
+from tinygrad.codegen.opt import Opt, OptOps
+
+# ************************* implementation of the problem ************************
+
+def myhash(a: Tensor) -> Tensor:
+  a = (a + 0x7ED55D16) + (a << 12)
+  a = (a ^ 0xC761C23C) ^ (a >> 19)
+  a = (a + 0x165667B1) + (a << 5)
+  a = (a + 0xD3A2646C) ^ (a << 9)
+  a = (a + 0xFD7046C5) + (a << 3)
+  a = (a ^ 0xB55A4F09) ^ (a >> 16)
+  return a
+
+def select_with_where_tree(values: Tensor, relative_idx: Tensor) -> Tensor:
+  n = values.shape[0]
+  if n == 1: return values[0].expand(relative_idx.shape)
+
+  mid = n // 2
+  left = select_with_where_tree(values[:mid], relative_idx)
+  right = select_with_where_tree(values[mid:], relative_idx - mid)
+
+  go_left = relative_idx < mid
+  return go_left.where(left, right)
+
+def tree_traversal(forest: Tensor, val: Tensor, height: int, rounds: int, where_tree_threshold=3) -> Tensor:
+  # All walkers start at idx=0
+  idx = Tensor.zeros(val.shape, device=val.device, dtype=dtypes.uint32)
+
+  for r in range(rounds):
+    level = r % (height + 1)
+    level_start = (1 << level) - 1
+    level_size = 1 << level
+
+    if level == 0:
+      # At root (level 0), all walkers are at idx=0
+      # No gather needed, just broadcast the root value
+      node_val = forest[0].expand(val.shape)
+      idx = idx * 0  # Reset to 0
+    elif level <= where_tree_threshold:
+      # Small level: use where-tree
+      level_values = forest[level_start : level_start + level_size]
+      relative_idx = (idx - level_start)
+      node_val = select_with_where_tree(level_values, relative_idx)
+    else:
+      # Large level: use gather
+      node_val = forest.gather(0, idx)
+
+    val = myhash(val ^ node_val)
+    idx = (idx << 1) + (1 + (val & 1))
+
+    # No wrap check needed! At round 10 (level becomes 0), we reset idx above.
+
+  return val.contiguous(arg=(Opt(OptOps.UPCAST, 0, 8),))
+
+# ************************* renderer for VLIW machine *************************
+
+def loop_unrolling(sink:UOp):
+  rng = [x for x in sink.toposort() if x.op is Ops.RANGE]
+  if len(rng) == 0: return None
+  print(f"unrolling loop with size {rng[0].vmax+1}")
+  unrolled_sinks = [sink.substitute({rng[0]:rng[0].const_like(i)}).src[0] for i in range(rng[0].vmax+1)]
+  return UOp.sink(*unrolled_sinks, arg=sink.arg)
+
+global_addrs = []
+vliw_prepare = PatternMatcher([
+  # loop unrolling (should be a part of tinygrad)
+  (UPat(Ops.SINK, name="sink"), loop_unrolling),
+  # cast is fake
+  (UPat(Ops.CAST, name="c"), lambda c: c.src[0]),
+  # rewrites to hardcode the addresses in memory
+  (UPat(Ops.PARAM, name="dg"), lambda dg: UOp.const(dtypes.uint, global_addrs[dg.arg])),
+  # INDEX is just plus
+  (UPat(Ops.INDEX, name="i"), lambda i: i.src[0]+i.src[1]),
+])+symbolic
+
+class VLIWRenderer(Renderer):
+  has_local = False  # TODO: this should be the default / cleaned up
+  # this says this backend supports MULACC + more. decompositions uses this
+  code_for_op: dict = {Ops.MULACC: None, Ops.ADD: "+", Ops.MUL: "*",
+                       Ops.XOR: "^", Ops.AND: "&", Ops.OR: "|",
+                       Ops.SHL: "<<", Ops.SHR: ">>", Ops.CMPLT: "<"}
+  # this matcher runs while still in graph form
+  pre_matcher = vliw_prepare
+
+  def render(self, uops:list[UOp]):
+
+    # TODO: this is a minimal renderer. for low cycle count, make it good
+    # to get speed, you need to add VLIW packing
+    # to get under 1536 regs, you need to add a register allocator
+    # we left the fun parts to you
+
+    print(f"rendering with {len(uops)} uops")
+    reg, inst = 0, []
+    r: dict[UOp, int] = {}
+    for u in uops:
+      assert u.dtype.count in (1,8), "dtype count must be 1 or 8"
+
+      # dumb register allocator
+      if u.op not in {Ops.STORE, Ops.SINK, Ops.GEP}:
+        r[u] = reg
+        reg += u.dtype.count
+
+      # render UOps to instructions
+      match u.op:
+        case Ops.SINK:
+          inst.append({"flow": [("halt",)]})
+        case Ops.CONST:
+          inst.append({"load": [("const", r[u], u.arg)]})
+        case Ops.GEP:
+          # a GEP is just an alias to a special register in the vector
+          r[u] = r[u.src[0]] + u.arg[0]
+        case Ops.VECTORIZE:
+          if all(s == u.src[0] for s in u.src):
+            # if all sources are the same, we can broadcast
+            inst.append({"valu": [("vbroadcast", r[u], r[u.src[0]])]})
+          else:
+            # this is a copy into a contiguous chunk of registers
+            inst.extend({"flow": [("add_imm", r[u]+i, r[s], 0)]} for i,s in enumerate(u.src) if r[s] != r[u]+i)
+        case Ops.LOAD:
+          op = "vload" if u.dtype.count > 1 else "load"
+          inst.append({"load": [(op, r[u], r[u.src[0]])]})
+        case Ops.STORE:
+          op = "vstore" if u.src[1].dtype.count > 1 else "store"
+          inst.append({"store": [(op, r[u.src[0]], r[u.src[1]])]})
+        case Ops.MULACC:
+          assert u.dtype.count == 8
+          inst.append({"valu": [("multiply_add", r[u], r[u.src[0]], r[u.src[1]], r[u.src[2]])]})
+        case Ops.WHERE:
+          assert u.dtype.count == 8
+          inst.append({"flow": [("vselect", r[u], r[u.src[0]], r[u.src[1]], r[u.src[2]])]})
+        case _ if u.op in self.code_for_op:
+          cat = "valu" if u.dtype.count > 1 else "alu"
+          inst.append({cat: [(self.code_for_op[u.op], r[u], r[u.src[0]], r[u.src[1]])]})
+        case _:
+          raise NotImplementedError(f"unhandled op {u.op}")
+    return repr(inst)
+
+# ************************* test and render *************************
+
+import sys, types
+PROBLEM_URL = "https://raw.githubusercontent.com/anthropics/original_performance_takehome/refs/heads/main/tests/frozen_problem.py"
+sys.modules["problem"] = problem = types.ModuleType("problem")
+exec(fetch(PROBLEM_URL).read_text(), problem.__dict__)
+
+if __name__ == "__main__":
+  batch_size = getenv("BS", 256)
+  height = 10
+  rounds = getenv("ROUNDS", 16)
+
+  # build problem
+  tree = problem.Tree.generate(height)
+  inp = problem.Input.generate(tree, batch_size, rounds)
+  mem = problem.build_mem_image(tree, inp)
+  global_addrs.extend([mem[6], mem[6], mem[4]])  # output, input, forest
+
+  # *** verify the kernel in tinygrad compared to reference ***
+
+  forest_t = Tensor(tree.values, dtype=dtypes.uint32)
+  val_t = Tensor(inp.values, dtype=dtypes.uint32)
+
+  if getenv("VERIFY", 1):
+    # verify on normal tinygrad device
+    with Context(PCONTIG=2):
+      out = tree_traversal(forest_t, val_t, height, rounds)
+      val_out = out.tolist()
+    problem.reference_kernel(tree, inp)
+    assert val_out == inp.values
+    print("verification passed")
+
+  # *** render to device ***
+
+  from tinygrad.codegen import get_program
+  with Context(PCONTIG=2, DEVECTORIZE=2, SPEC=0):
+    out = tree_traversal(forest_t, val_t, height, rounds)
+    sink = out.schedule()[-1].ast
+    prg = get_program(sink, VLIWRenderer())
+
+  # *** run on Machine and compare ***
+
+  # NOTE: the scratch size needs to be reduced to 1536 when you have a register allocator
+  src = eval(prg.src)
+  max_regs = max(t[1] for instr in src for v in instr.values() for t in v if len(t) > 1) + 8
+  print(f"{max_regs:5d} regs used" + ("" if max_regs <= 1536 else "       <-- WARNING: TOO MANY REGISTERS, MUST BE <= 1536"))
+  machine = problem.Machine(mem, src, problem.DebugInfo(scratch_map={}), n_cores=1, trace=False, scratch_size=max_regs)
+  machine.run()
+  print(f"ran for {machine.cycle:5d} cycles" + ("" if machine.cycle <= 1363 else "  <-- EVEN CLAUDE GOT 1363"))
+
+  # compare to reference
+  ref_mem = mem.copy()
+  for _ in problem.reference_kernel2(ref_mem, {}): pass
+  assert machine.mem[mem[6]:mem[6]+mem[2]] == ref_mem[mem[6]:mem[6]+mem[2]]
+  print("compare passed!")
@@ -0,0 +1,79 @@
+from typing import Optional
+from tinygrad import Tensor
+from tinygrad.dtype import DTypeLike, dtypes
+import math
+
+# rewritten from numpy
+def rfftfreq(n: int, d: float = 1.0, device=None) -> Tensor:
+  val = 1.0 / (n * d)
+  N = n // 2 + 1
+  results = Tensor.arange(N, device=device)
+  return results * val
+
+# just like in librosa
+def fft_frequencies(sr: float, n_fft: int) -> Tensor:
+  return rfftfreq(n=n_fft, d=1.0 / sr)
+
+def hz_to_mel(freq: Tensor) -> Tensor:
+  # linear part
+  f_min = 0.0
+  f_sp = 200.0 / 3
+  mels = (freq - f_min) / f_sp
+
+  # log-scale part
+  min_log_hz = 1000.0  # beginning of log region (Hz)
+  mask = freq >= min_log_hz
+  return mask.where(((min_log_hz - f_min) / f_sp) + (freq / min_log_hz).log() / (math.log(6.4) / 27.0), mels)
+
+def mel_to_hz(mels: Tensor) -> Tensor:
+  # linear scale
+  f_min = 0.0
+  f_sp = 200.0 / 3
+  freqs = f_min + f_sp * mels
+
+  # nonlinear scale
+  min_log_hz = 1000.0  # beginning of log region (Hz)
+  min_log_mel = (min_log_hz - f_min) / f_sp  # same (Mels)
+  logstep = math.log(6.4) / 27.0  # step size for log region
+
+  log_t = mels >= min_log_mel
+  freqs = log_t.where(min_log_hz * ((logstep * (mels - min_log_mel)).exp()), freqs)
+  return freqs
+
+def mel_frequencies(n_mels: int = 128, *, fmin: float = 0.0, fmax: float = 11025.0) -> Tensor:
+  # center freqs of mel bands - uniformly spaced between limits
+  min_max_mel = hz_to_mel(Tensor([fmin, fmax]))
+
+  mels = Tensor.linspace(min_max_mel[0], min_max_mel[1], n_mels)
+  hz = mel_to_hz(mels)
+  return hz
+
+def mel(
+  *,
+  sr: float,
+  n_fft: int,
+  n_mels: int = 128,
+  fmin: float = 0.0,
+  fmax: Optional[float] = None,
+  dtype: DTypeLike = dtypes.default_float,
+) -> Tensor:
+  if fmax is None:
+    fmax = float(sr) / 2
+
+  n_mels = int(n_mels)
+
+  fftfreqs = fft_frequencies(sr=sr, n_fft=n_fft)  # center freqs of each FFT bin
+  mel_f = mel_frequencies(n_mels + 2, fmin=fmin, fmax=fmax)  # center freqs of mel bands
+
+  fdiff = mel_f[1:] - mel_f[:-1]
+  ramps = mel_f[None].T.expand(-1, fftfreqs.shape[-1]) - fftfreqs
+
+  lower = -ramps[:n_mels] / fdiff[:n_mels][None].T
+  upper = ramps[2 : n_mels + 2] / fdiff[1 : n_mels + 1][None].T
+  weights = lower.minimum(upper).maximum(0)
+
+  # Slaney-style mel is scaled to be approx constant energy per channel
+  enorm = 2.0 / (mel_f[2 : n_mels + 2] - mel_f[:n_mels])
+  weights *= enorm[:, None]
+
+  return weights
@@ -72,7 +72,7 @@ def loader_process(q_in, q_out, X:Tensor, seed):
      #storage_tensor._copyin(img_tensor.numpy())

      # faster
-      X[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
+      X[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = img.tobytes()

      # ideal
      #X[idx].assign(img.tobytes())   # NOTE: this is slow!
@@ -213,12 +213,13 @@ class InterleavedDataset:
    self.queues[queue_index].queue.extend(load_file(file))

 # Reference: https://github.com/mlcommons/training/blob/1c8a098ae3e70962a4f7422c0b0bd35ae639e357/language_model/tensorflow/bert/run_pretraining.py, Line 394
-def batch_load_train_bert(BS:int):
+def batch_load_train_bert(BS:int, seed:int|None=None):
  from extra.datasets.wikipedia import get_wiki_train_files
+  rng = random.Random(seed)
  fs = sorted(get_wiki_train_files())
  train_files = []
  while fs: # TF shuffle
-    random.shuffle(fs)
+    rng.shuffle(fs)
    train_files.append(fs.pop(0))

  cycle_length = min(getenv("NUM_CPU_THREADS", min(os.cpu_count(), 8)), len(train_files))
@@ -263,8 +264,8 @@ def load_unet3d_data(preprocessed_dataset_dir, seed, queue_in, queue_out, X:Tens
      x = random_brightness_augmentation(x)
      x = gaussian_noise(x)

-    X[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = x.tobytes()
-    Y[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = y.tobytes()
+    X[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = x.tobytes()
+    Y[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = y.tobytes()

    queue_out.put(idx)
  queue_out.put(None)
@@ -378,12 +379,12 @@ def load_retinanet_data(base_dir:Path, val:bool, queue_in:Queue, queue_out:Queue
      clipped_match_idxs = np.clip(match_idxs, 0, None)
      clipped_boxes, clipped_labels = tgt["boxes"][clipped_match_idxs], tgt["labels"][clipped_match_idxs]

-      boxes[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_boxes.tobytes()
-      labels[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_labels.tobytes()
-      matches[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = match_idxs.tobytes()
-      anchors[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = anchor.tobytes()
+      boxes[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = clipped_boxes.tobytes()
+      labels[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = clipped_labels.tobytes()
+      matches[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = match_idxs.tobytes()
+      anchors[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = anchor.tobytes()

-    imgs[idx].contiguous().realize().uop.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
+    imgs[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = img.tobytes()

    queue_out.put(idx)
  queue_out.put(None)
@@ -551,7 +552,7 @@ class BinIdxDataset:
    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.dtype = {1:np.dtype(np.uint8), 2:np.dtype(np.int8), 3:np.dtype(np.int16), 4:np.dtype(np.int32), 5:np.dtype(np.int64), 6:np.dtype(np.float64), 7:np.dtype(np.double), 8:np.dtype(np.uint16)}[dtype_code]
    self.count, = struct.unpack("<Q", self.idx.read(8))
    doc_count, = struct.unpack("<Q", self.idx.read(8))

@@ -568,7 +569,7 @@ class BinIdxDataset:
    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"))
+    self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin")).numpy()

  def _index(self, idx) -> tuple[int, int]:
    return int(self.pointers[idx]), int(self.sizes[idx])
@@ -577,7 +578,7 @@ class BinIdxDataset:
    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)
+    return self.bin_t[ptr:ptr+length*self.dtype.itemsize].view(self.dtype)

 # https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/datasets.html
 class GPTDataset:
@@ -636,7 +637,7 @@ class GPTDataset:
        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)
+    text = np.concatenate(sample_parts, axis=0)

    return text

@@ -763,48 +764,27 @@ class BlendedGPTDataset:

    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):
+def get_llama3_dataset(samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False) -> BlendedGPTDataset:
+  if small:
+    if val:
+      return BlendedGPTDataset(
+        [base_dir / "c4-validation-91205-samples.en_text_document"], [1.0], samples, seqlen, seed, shuffle=False)
+    return BlendedGPTDataset(
+      [base_dir / "c4-train.en_6_text_document"], [1.0], samples, seqlen, seed, shuffle=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)
+    return BlendedGPTDataset(
+      [base_dir / "validation" / "c4-validationn-91205-samples.en_text_document"], [1.0], samples, seqlen, seed, shuffle=False)
+  return BlendedGPTDataset(
+    [base_dir / "c4-train.en_6_text_document", base_dir / "c4-train.en_7_text_document"], [1.0, 1.0], samples, seqlen, seed, shuffle=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 iterate_llama3_dataset(dataset:BlendedGPTDataset, bs:int):
+  for b in range(math.ceil(dataset.samples / bs)):
+    batch = [dataset.get(b * bs + i) for i in range(bs)]
+    stacked = np.stack(batch, axis=0)
+    yield Tensor(stacked, device="NPY")

-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)
+def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False):
+  return iterate_llama3_dataset(get_llama3_dataset(samples, seqlen, base_dir, seed, val, small), bs)

 if __name__ == "__main__":
  def load_unet3d(val):
@@ -219,7 +219,18 @@ def get_mlperf_bert_model():
  config = get_mlperf_bert_config()
  if getenv("DISABLE_DROPOUT", 0):
    config["hidden_dropout_prob"] = config["attention_probs_dropout_prob"] = 0.0
-  return BertForPretraining(**config)
+  model = BertForPretraining(**config)
+  if getenv("FP8_TRAIN"):
+    from extra.fp8.fp8_linear import convert_to_float8_training
+    def module_filter_fn(mod, fqn):
+      if isinstance(mod, LinearBert):
+        skip_layers = [] if (ln:=config["num_hidden_layers"]) <= 2 else ["bert.encoder.layer.0.", f"bert.encoder.layer.{ln-1}"]
+        if mod.weight.shape[-1] >= 1024 and "encoder" in fqn and not any(name in fqn for name in skip_layers):
+          print(f"replacing linear with fp8: {fqn} {mod.weight.shape}")
+          return True
+      return False
+    convert_to_float8_training(model, module_filter_fn)
+  return model

 def get_fake_data_bert(BS:int):
  return {
@@ -234,12 +234,9 @@ def eval_llama3():
    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)
+  from examples.mlperf.dataloader import get_llama3_dataset, iterate_llama3_dataset
+  eval_dataset = get_llama3_dataset(5760, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
+  iter = iterate_llama3_dataset(eval_dataset, BS)

  losses = []
  for tokens in tqdm(iter, total=5760//BS):
@@ -3,7 +3,7 @@ from pathlib import Path
 import multiprocessing

 from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
-from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling
+from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling, profile_marker
 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

@@ -1008,6 +1008,7 @@ def train_bert():
  config["DISABLE_DROPOUT"] = getenv("DISABLE_DROPOUT", 0)
  config["TRAIN_BEAM"]    = TRAIN_BEAM = getenv("TRAIN_BEAM", BEAM.value)
  config["EVAL_BEAM"]     = EVAL_BEAM  = getenv("EVAL_BEAM", BEAM.value)
+  config["FP8_TRAIN"]     = getenv("FP8_TRAIN", 0)

  Tensor.manual_seed(seed)  # seed for weight initialization

@@ -1085,7 +1086,7 @@ def train_bert():
  if RUNMLPERF:
    # only load real data with RUNMLPERF
    eval_it = iter(batch_load_val_bert(EVAL_BS))
-    train_it = iter(tqdm(batch_load_train_bert(BS), total=train_steps, disable=BENCHMARK))
+    train_it = iter(tqdm(batch_load_train_bert(BS, seed=seed), total=train_steps, disable=BENCHMARK))
    for _ in range(start_step): next(train_it) # Fast forward
  else:
    # repeat fake data
@@ -1147,7 +1148,7 @@ def train_bert():

        device_str = parameters[0].device if isinstance(parameters[0].device, str) else f"{parameters[0].device[0]} * {len(parameters[0].device)}"
        loss = loss.item()
-        assert not math.isnan(loss)
+        if not getenv("FP8_TRAIN"): assert not math.isnan(loss)
        lr = lr.item()

      cl = time.perf_counter()
@@ -1160,7 +1161,7 @@ def train_bert():
      if WANDB:
        wandb.log({"lr": lr, "train/loss": loss, "train/global_norm": global_norm.item(), "train/step_time": cl - st,
                    "train/python_time": pt - st, "train/data_time": dt - pt, "train/cl_time": cl - dt,
-                    "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*GBS})
+                    "train/mem":GlobalCounters.mem_used / 1e9, "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*GBS})

      train_data, next_data = next_data, None
      i += 1
@@ -1285,17 +1286,24 @@ def train_llama3():
  from examples.llama3 import MODEL_PARAMS
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup

+  BENCHMARK = getenv("BENCHMARK")
+
  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)
-  assert grad_acc == 1, f"{grad_acc=} is not supported"
  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
  SEED               = config["SEED"]                   = getenv("SEED", 5760)
+  DATA_SEED          = config["DATA_SEED"]              = getenv("DATA_SEED", SEED)
  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_SAMPLES       = config["EVAL_SAMPLES"]           = getenv("EVAL_SAMPLES", 5760 if not SMALL else 1024)
+  MAX_STEPS          = config["MAX_STEPS"]              = getenv("MAX_STEPS", math.ceil(1_200_000 * 1152 / GBS))
+  WARMUP_STEPS       = config["WARMUP_STEPS"]           = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
+  LR                 = config["LR"]                     = getenv("LR", 8e-5 * GBS / 1152)
+  END_LR             = config["END_LR"]                 = getenv("END_LR", 8e-7)
  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)
@@ -1309,15 +1317,26 @@ def train_llama3():
  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)
+  opt_learning_rate_warmup_steps = WARMUP_STEPS
+  opt_learning_rate_decay_steps = MAX_STEPS - opt_learning_rate_warmup_steps
+  opt_base_learning_rate = LR
+  opt_end_learning_rate = END_LR
+
+  Tensor.manual_seed(SEED)  # seed for weight initialization
+
+  # ** init wandb **
+  WANDB = getenv("WANDB")
+  if WANDB:
+    import wandb
+    wandb_args = {"id": wandb_id, "resume": "must"} if (wandb_id := getenv("WANDB_RESUME", "")) else {}
+    wandb.init(config=config, **wandb_args, project="MLPerf-LLaMA3")

  model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
  # vocab_size from the mixtral tokenizer
  if not SMALL: model_params |= {"vocab_size": 32000}
  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
+  print(f"model parameters: {model_params}")
+
  model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
  params = get_parameters(model)
  # weights are all bfloat16 for now
@@ -1353,6 +1372,12 @@ def train_llama3():

  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)
+
+  # init grads
+  for p in optim.params:
+    p.grad = p.zeros_like().contiguous().realize()
+  grads = [p.grad for p in optim.params]
+
  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"):
@@ -1365,9 +1390,8 @@ def train_llama3():
    load_state_dict(scheduler, safe_load(fn), realize=False)

  @TinyJit
-  @Tensor.train()
-  def train_step(model, tokens:Tensor):
-    optim.zero_grad()
+  def minibatch(tokens:Tensor):
+    tokens = tokens.to(None)
    if (DP := getenv("DP", 1)) > 1:
      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
      tokens = tokens.shard(device, 0)
@@ -1377,27 +1401,41 @@ def train_llama3():
    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
    loss.backward()
+    assert all(p.grad is g for p,g in zip(optim.params, grads))
+    Tensor.realize(loss, *grads)
+    return loss.flatten().float().to("CPU")
+
+  @TinyJit
+  def optim_step():
+    for p in optim.params:
+      p.grad.assign(p.grad / grad_acc)
+
    # 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)
+      for g in grads:
+        total_norm += g.float().square().sum()
+      total_norm = total_norm.sqrt().contiguous().realize()
+      for g in grads:
+        g.assign((g * (opt_gradient_clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(g.dtype)).realize()

    optim.step()
    scheduler.step()

+    for g in grads:
+      g.assign(g.zeros_like().contiguous()).realize()
+
    lr = optim.lr
-    loss.realize(lr)
-    return loss, lr
+    Tensor.realize(lr, *grads)
+
+    return lr.float().to("CPU")

  @TinyJit
  @Tensor.train(False)
-  def eval_step(model, tokens:Tensor):
+  def eval_step(tokens:Tensor):
+    tokens = tokens.to(None)
    if (DP := getenv("DP", 1)) > 1:
      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
      tokens = tokens.shard(device, 0)
@@ -1406,7 +1444,7 @@ def train_llama3():
      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()
+    return loss.flatten().float().to("CPU")

  # ** data iters **
  def fake_data(bs, samples):
@@ -1417,45 +1455,80 @@ def train_llama3():
    if getenv("FAKEDATA", 0):
      return fake_data(BS, SAMPLES)
    else:
-      if SMALL:
-        from examples.mlperf.dataloader import batch_load_llama3_small
-        return batch_load_llama3_small(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
-      else:
-        from examples.mlperf.dataloader import batch_load_llama3
-        return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
+      from examples.mlperf.dataloader import batch_load_llama3
+      return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=DATA_SEED, val=bool(TRAIN_ON_VAL), small=bool(SMALL))
+
+  if getenv("FAKEDATA", 0):
+    eval_dataset = None
+  else:
+    from examples.mlperf.dataloader import get_llama3_dataset
+    eval_dataset = get_llama3_dataset(EVAL_SAMPLES, SEQLEN, BASEDIR, val=True, small=bool(SMALL))

  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)
+    if eval_dataset is None:
+      return fake_data(EVAL_BS, EVAL_SAMPLES)
+    from examples.mlperf.dataloader import iterate_llama3_dataset
+    return iterate_llama3_dataset(eval_dataset, EVAL_BS)

-  iter = get_train_iter()
+  num_params = sum(p.numel() for p in params) - model_params["vocab_size"]*model_params["dim"]
+  train_iter = get_train_iter()
  i, sequences_seen = resume_ckpt, 0
-  for tokens in tqdm(iter, total=SAMPLES//GBS):
+  step_times = []
+  while i < MAX_STEPS:
    GlobalCounters.reset()
    if getenv("TRAIN", 1):
-      t = time.perf_counter()
-      loss, lr = train_step(model, tokens)
+      profile_marker(f"train @ {i}")
+      st = time.perf_counter()
+
+      stopped = False
+      for _ in range(grad_acc):
+        ist = time.perf_counter()
+        try: tokens = next(train_iter)
+        except StopIteration:
+          stopped = True
+          break
+        dt = time.perf_counter()
+        loss = minibatch(tokens)
+      if stopped: break
+
+      gt = time.perf_counter()
+      lr = optim_step()
+      ot = time.perf_counter()
+
      loss = loss.float().item()
      lr = lr.item()

+      et = time.perf_counter()
+      step_time = et - st
+      gbs_time = gt - st
+      optim_time = ot - gt
+      data_time = dt - ist
+      dev_time = step_time - data_time * grad_acc
+      if BENCHMARK: step_times.append(step_time)
+
      i += 1
-      sequences_seen += tokens.shape[0]
+      sequences_seen += GBS

-      sec = time.perf_counter()-t
+      mem_gb = GlobalCounters.mem_used / 1e9
+      gflops = GlobalCounters.global_ops / 1e9 / dev_time
+      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * max(getenv("DP", 1), getenv("MP", 1)) * 2.3e15)) * 100
      tqdm.write(
-        f"{i:5} {sec:.2f} s run, {loss:.4f} loss, {lr:.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, "
-        f"{GlobalCounters.global_ops * 1e-9 / sec:9.2f} GFLOPS")
+          f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
+          f"{lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")

-      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 WANDB:
+        wandb.log({
+          "lr": lr, "train/loss": loss,
+          "train/step_time": step_time,
+          "train/gbs_time": gbs_time,
+          "train/optim_time": optim_time,
+          "train/dev_time": dev_time,
+          "train/data_time": data_time,
+          "train/mem": mem_gb,
+          "train/GFLOPS": gflops,
+          "train/MFU": mfu,
+          "train/sequences_seen": sequences_seen
+        })

      if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
        tqdm.write("saving checkpoint")
@@ -1467,20 +1540,36 @@ def train_llama3():
        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):
+      if i == BENCHMARK:
+        median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]
+        estimated_total_minutes = int(median_step_time * (SAMPLES // GBS) / 60)
+        print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
+        print(f"epoch global_ops: {GlobalCounters.global_ops:_}, "
+              f"epoch global_mem: {GlobalCounters.global_mem:_}")
+
+    if (sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1)) or (BENCHMARK and i == BENCHMARK):
+      if EVAL_BS == 0: return
      tqdm.write(f"evaluating after {sequences_seen} sequences")
+      profile_marker(f"eval @ {i}")

      # 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()
+      for j,tokens in tqdm(enumerate(eval_iter), total=EVAL_SAMPLES//EVAL_BS):
+        eval_losses += eval_step(tokens).tolist()
+
+        if BENCHMARK and (j+1) == min(BENCHMARK, EVAL_SAMPLES//EVAL_BS):
+          return
+
      log_perplexity = Tensor(eval_losses).mean().float().item()

      tqdm.write(f"eval log perplexity: {log_perplexity:.4f}")

+      if WANDB:
+        wandb.log({"eval/log_perplexity": log_perplexity, "eval/sequences_seen": sequences_seen})
+
      if log_perplexity < EVAL_TARGET:
        tqdm.write(f"target achieved after {sequences_seen} sequences")
        if getenv("CKPT"):
@@ -1559,7 +1648,7 @@ def train_stable_diffusion():
    loss, out_lr = loss.detach().to("CPU"), optimizer.lr.to("CPU")
    Tensor.realize(loss, out_lr)
    return loss, out_lr
-    
+
  # checkpointing takes ~9 minutes without this, and ~1 minute with this
  @TinyJit
  def ckpt_to_cpu():
@@ -1598,7 +1687,7 @@ def train_stable_diffusion():
    if i == 3:
      for _ in range(3): ckpt_to_cpu() # do this at the beginning of run to prevent OOM surprises when checkpointing
      print("BEAM COMPLETE", flush=True) # allows wrapper script to detect BEAM search completion and retry if it failed
-      
+
    total_train_time = time.perf_counter() - train_start_time
    if WANDB:
      wandb.log({"train/loss": loss_item, "train/lr": lr_item, "train/loop_time_prev": loop_time, "train/dl_time": dl_time, "train/step": i,
@@ -4,7 +4,7 @@ export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128

-export IGNORE_OOB=1
+export CHECK_OOB=0

 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
@@ -5,7 +5,7 @@ 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 CHECK_OOB=0
 export REWRITE_STACK_LIMIT=500000

 export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -8,7 +8,7 @@ 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 CHECK_OOB=0
 export REWRITE_STACK_LIMIT=500000

 export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -11,7 +11,7 @@ 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 CHECK_OOB=0
 export REWRITE_STACK_LIMIT=500000

 export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -0,0 +1,20 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1 DEBUG=0 JIT=1 FLASH_ATTENTION=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
+
+# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
+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 CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000
+
+export BEAM=0 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"
+
+export WANDB=1 PARALLEL=0
+
+RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -0,0 +1,24 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
+
+# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
+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 CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000
+
+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"
+export BEAM_TIMEOUT_SEC=15
+export FP8_TRAIN=1
+# search
+IGNORE_BEAM_CACHE=1 BENCHMARK=10 BERT_LAYERS=2 RUNMLPERF=0 python3 examples/mlperf/model_train.py
+
+export WANDB=1 PARALLEL=0
+
+RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -11,7 +11,7 @@ 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 CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000

 export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -4,7 +4,7 @@ export PYTHONPATH="." NV=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72

-export IGNORE_OOB=1
+export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=500000

 export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -4,7 +4,7 @@ export PYTHONPATH="." NV=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72

-export IGNORE_OOB=1
+export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=500000

 export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -7,7 +7,7 @@ export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_green"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72

-export IGNORE_OOB=1
+export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=500000

 export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -4,7 +4,7 @@ export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

-export IGNORE_OOB=1
+export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=500000

 export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -4,7 +4,7 @@ export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

-export IGNORE_OOB=1
+export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=500000

 export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -7,7 +7,7 @@ export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_red"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

-export IGNORE_OOB=1
+export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=500000

 export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
@@ -0,0 +1,37 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+
+export DEBUG=${DEBUG:-2}
+export FLASH_ATTENTION=${FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-5760}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+
+export FAKEDATA=1 BENCHMARK=10 LLAMA_LAYERS=2
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,35 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+
+export DEBUG=${DEBUG:-0}
+export FLASH_ATTENTION=${FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export USE_ATOMICS=${USE_ATOMICS:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-8} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-$RANDOM}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,6 @@
+#!/bin/bash
+export BENCHMARK=5
+export EVAL_BS=0
+export VIZ=${VIZ:--1}
+examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
+PYTHONPATH="." extra/viz/cli.py --profile --device "AMD" --top 20
@@ -0,0 +1,10 @@
+#!/bin/bash
+export BENCHMARK=5
+export EVAL_BS=0
+export FAKEDATA=1
+export NULL_ALLOW_COPYOUT=1
+export HIP_VISIBLE_DEVICES=""
+export DEV=NULL
+export JITBEAM=0
+export LLAMA_LAYERS=${LLAMA_LAYERS:-"2"}
+time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
@@ -6,7 +6,6 @@ import argparse, time
 from collections import namedtuple
 from typing import Dict, Any

-from PIL import Image
 import numpy as np
 from tinygrad import Device, GlobalCounters, dtypes, Tensor, TinyJit
 from tinygrad.helpers import Timing, Context, getenv, fetch, colored, tqdm, flatten, profile_marker
@@ -336,6 +335,7 @@ if __name__ == "__main__":
  print(x.shape)

  profile_marker("save image")
+  from PIL import Image
  im = Image.fromarray(x.numpy())
  print(f"saving {args.out}")
  im.save(args.out)
@@ -7,7 +7,7 @@ if __name__ == "__main__":
  with open(fetch(sys.argv[1]), "rb") as f:
    run_onnx_jit = pickle.load(f)
  input_name = run_onnx_jit.captured.expected_names[0]
-  device = run_onnx_jit.captured.expected_st_vars_dtype_device[0][-1]
+  device = run_onnx_jit.captured.expected_input_info[0][-1]
  print(f"input goes into {input_name=} on {device=}")
  hit = 0
  for i,(img,y) in enumerate(imagenet_dataloader(cnt=getenv("CNT", 100))):
@@ -48,7 +48,7 @@ def prepare_browser_chunks(model):
        weight_metadata = metadata.get(name, default)
        weight_metadata["parts"][part_num] = {"file": i, "file_start_pos": cursor, "size": size}
        metadata[name] = weight_metadata
-        data = bytes(state_dict[name].uop.base.realized.as_buffer())
+        data = bytes(state_dict[name].uop.base.realized.as_memoryview())
        data = data if not offsets else data[offsets[0]:offsets[1]]
        writer.write(data)
        cursor += size
@@ -93,7 +93,7 @@ if __name__ == "__main__":
    forward: Any = None

  sub_steps = [
-    Step(name = "textModel", input = [Tensor.randn(1, 77)], forward = model.cond_stage_model.transformer.text_model),
+    Step(name = "textModel", input = [Tensor.randint(1, 77, low=0, high=49408, dtype=dtypes.int32)], forward = model.cond_stage_model.transformer.text_model),
    Step(name = "diffusor", input = [Tensor.randn(1, 77, 768), Tensor.randn(1, 77, 768), Tensor.randn(1,4,64,64), Tensor.rand(1), Tensor.randn(1), Tensor.randn(1), Tensor.randn(1)], forward = model),
    Step(name = "decoder", input = [Tensor.randn(1,4,64,64)], forward = model.decode),
    Step(name = "f16tof32", input = [Tensor.randn(2097120, dtype=dtypes.uint32)], forward = u32_to_f16)
@@ -7,6 +7,7 @@ from tinygrad import Tensor, TinyJit, Variable, nn, dtypes
 from tinygrad.nn.state import torch_load, load_state_dict
 from tinygrad.helpers import getenv, fetch

+from examples.audio_helpers import mel
 import numpy as np
 import librosa

@@ -159,7 +160,7 @@ def prep_audio(waveforms: List[np.ndarray], batch_size: int, truncate=False) ->

  stft = librosa.stft(waveforms, n_fft=N_FFT, hop_length=HOP_LENGTH, window='hann', dtype=np.csingle)
  magnitudes = np.absolute(stft[..., :-1]) ** 2
-  mel_spec = librosa.filters.mel(sr=RATE, n_fft=N_FFT, n_mels=N_MELS) @ magnitudes
+  mel_spec = mel(sr=RATE, n_fft=N_FFT, n_mels=N_MELS).numpy() @ magnitudes

  log_spec = np.log10(np.clip(mel_spec, 1e-10, None))
  log_spec = np.maximum(log_spec, log_spec.max((1,2), keepdims=True) - 8.0)
@@ -92,7 +92,7 @@ class SMICtx:
    self.prev_terminal_width = 0
    self.prev_terminal_height = 0

-    remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:"]
+    remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:", "Processing accelerators:"]
    lspci = subprocess.check_output(["lspci"]).decode("utf-8").splitlines()
    self.lspci = {l.split()[0]: l.split(" ", 1)[1] for l in lspci}
    for k,v in self.lspci.items():
@@ -153,7 +153,7 @@ class SMICtx:
    tables = {}
    for dev in self.devs:
      match dev.ip_ver[am.MP1_HWIP]:
-        case (13,0,6): table_t = dev.smu.smu_mod.MetricsTableX_t
+        case (13,0,6): table_t = dev.smu.smu_mod.MetricsTableV0_t
        case (13,0,12): table_t = dev.smu.smu_mod.MetricsTableV2_t
        case _: table_t = dev.smu.smu_mod.SmuMetricsExternal_t
      tables[dev] = dev.smu.read_table(table_t, dev.smu.smu_mod.SMU_TABLE_SMU_METRICS) if dev.pci_state == "D0" else None
@@ -165,17 +165,17 @@ class SMICtx:

  def get_gfx_activity(self, dev, metrics):
    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return max(0, min(100, self._smuq10_round(metrics.SocketGfxBusy)))
+      case (13,0,6)|(13,0,12): return max(0, min(100, self._smuq10_round(metrics.SocketGfxBusy)))
      case _: return metrics.SmuMetrics.AverageGfxActivity

  def get_mem_activity(self, dev, metrics):
    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return max(0, min(100, self._smuq10_round(metrics.DramBandwidthUtilization)))
+      case (13,0,6)|(13,0,12): return max(0, min(100, self._smuq10_round(metrics.DramBandwidthUtilization)))
      case _: return metrics.SmuMetrics.AverageUclkActivity

  def get_temps(self, dev, metrics, compact=False):
    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6):
+      case (13,0,6)|(13,0,12):
        temps = {
          "Hotspot": self._smuq10_round(metrics.MaxSocketTemperature),
          "HBM": self._smuq10_round(metrics.MaxHbmTemperature),
@@ -191,7 +191,7 @@ class SMICtx:

  def get_voltage(self, dev, metrics, compact=False):
    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return {}
+      case (13,0,6)|(13,0,12): return {}
      case _:
        voltage_keys = [(k, name) for k, name in dev.smu.smu_mod.SVI_PLANE_e.items()
                        if k < dev.smu.smu_mod.SVI_PLANE_COUNT and metrics.SmuMetrics.AvgVoltage[k] != 0]
@@ -205,33 +205,33 @@ class SMICtx:
  def get_gfx_freq(self, dev, metrics):
    if metrics is None: return 0
    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return self._smuq10_round(metrics.GfxclkFrequency[0])
+      case (13,0,6)|(13,0,12): return self._smuq10_round(metrics.GfxclkFrequency[0])
      case _:
        return metrics.SmuMetrics.AverageGfxclkFrequencyPostDs if self.get_gfx_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
               metrics.SmuMetrics.AverageGfxclkFrequencyPreDs

  def get_mem_freq(self, dev, metrics):
    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return self._smuq10_round(metrics.UclkFrequency)
+      case (13,0,6)|(13,0,12): return self._smuq10_round(metrics.UclkFrequency)
      case _:
        return metrics.SmuMetrics.AverageMemclkFrequencyPostDs if self.get_mem_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
               metrics.SmuMetrics.AverageMemclkFrequencyPreDs

  def get_fckl_freq(self, dev, metrics):
    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return self._smuq10_round(metrics.FclkFrequency)
+      case (13,0,6)|(13,0,12): return self._smuq10_round(metrics.FclkFrequency)
      case _:
        return metrics.SmuMetrics.AverageFclkFrequencyPostDs if self.get_mem_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
               metrics.SmuMetrics.AverageFclkFrequencyPreDs

  def get_fan_rpm_pwm(self, dev, metrics):
    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return None, None
+      case (13,0,6)|(13,0,12): return None, None
      case _: return metrics.SmuMetrics.AvgFanRpm, metrics.SmuMetrics.AvgFanPwm

  def get_power(self, dev, metrics):
    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.MaxSocketPowerLimit)
+      case (13,0,6)|(13,0,12): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.MaxSocketPowerLimit)
      case _: return metrics.SmuMetrics.AverageSocketPower, metrics.SmuMetrics.dGPU_W_MAX

  def get_mem_usage(self, dev):
@@ -280,7 +280,7 @@ class SMICtx:
      device_line = [f"{bold(dev.pcibus)} {trim(self.lspci[dev.pcibus[5:]], col_size - 20)}"] + [pad("", col_size)]
      activity_line = [f"GFX Activity {draw_bar(self.get_gfx_activity(dev, metrics) / 100, activity_line_width)}"] \
                    + [f"MEM Activity {draw_bar(self.get_mem_activity(dev, metrics) / 100, activity_line_width)}"] \
-                    + [f"MEM Usage    {draw_bar((mem_used / mem_total) / 100, activity_line_width, opt_text=mem_fmt)}"] \
+                    + [f"MEM Usage    {draw_bar(mem_used / mem_total, activity_line_width, opt_text=mem_fmt)}"] \

      temps_data, temps_data_compact = self.get_temps(dev, metrics), self.get_temps(dev, metrics, compact=True)
      temps_table = ["=== Temps (°C) ==="] + [f"{name:<16}: {color_temp(val)}" for name, val in temps_data.items()]
@@ -1,12 +1,18 @@
 #!/usr/bin/env python3

+import os
 from tinygrad.helpers import Context
 from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase
+from tinygrad.runtime.support.hcq import FileIOInterface
 from tinygrad.runtime.support.am.amdev import AMDev

 if __name__ == "__main__":
  gpus = System.pci_scan_bus(0x1002, [(0xffff, [0x74a1, 0x75a0])])
-  pcidevs = [PCIDevice(f"reset:{gpu}", gpu, bars=[0, 2, 5]) for gpu in gpus]
+  for gpu in gpus:
+    drv_path = f"/sys/bus/pci/devices/{gpu}/driver"
+    if FileIOInterface.exists(drv_path) and os.path.basename(os.readlink(drv_path)) == "amdgpu":
+      raise RuntimeError(f"amdgpu is bound to {gpu}. Stopping...")
+  pcidevs = [PCIDevice("AM", gpu, bars=[0, 2, 5]) for gpu in gpus]
  amdevs = []
  with Context(DEBUG=2):
    for pcidev in pcidevs:
@@ -19,8 +19,9 @@ 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, *, only_xcc0: bool = False):
+def parse_amdgpu_logs(log_content, register_names=None, register_objects=None, *, only_xcc0: bool = False):
  register_map = register_names or {}
+  register_objs = register_objects or {}

  def replace_register(match):
    reg = match.group(1)
@@ -37,6 +38,28 @@ def parse_amdgpu_logs(log_content, register_names=None, *, only_xcc0: bool = Fal
  # remove timing prefix
  processed_log = re.sub(r'^\[\s*\d+(?:\.\d+)?\]\s*', '', processed_log, flags=re.MULTILINE)

+  # decode register values into field dicts
+  def decode_value(match):
+    reg_name = match.group(1)
+    xcc_part = match.group(2)  # "xcc=0 " or ""
+    val_str = match.group(3)
+    val = int(val_str, 16)
+
+    reg_obj = register_objs.get(reg_name)
+    if reg_obj is not None and reg_obj.fields:
+      fields = reg_obj.decode(val)
+      # show raw for unaccounted bits
+      accounted = 0
+      for name, (start, end) in reg_obj.fields.items():
+        accounted |= (((1 << (end - start + 1)) - 1) << start)
+      unaccounted = val & ~accounted
+      parts = {k: v for k, v in fields.items() if v != 0}
+      if unaccounted: parts['_raw_unaccounted'] = hex(unaccounted)
+      return f"register {reg_name}, {xcc_part}with value {val_str} {parts}"
+    return match.group(0)
+
+  processed_log = re.sub(r'register (reg\w+), ((?:xcc=\d+ )?)with value (0x[0-9a-fA-F]+)', decode_value, processed_log)
+
  # keep only xcc=0 lines (but keep lines with no xcc at all)
  if only_xcc0:
    kept = []
@@ -50,16 +73,18 @@ def main():
  only_xcc0 = bool(getenv("ONLY_XCC0", 0))

  reg_names = {}
+  reg_objs = {}
  dev = PCIIface(None, 0)
  for x, y in dev.dev_impl.__dict__.items():
    if isinstance(y, AMRegister):
      for xcc, addr in y.addr.items():
        reg_names[addr] = f"{x}, xcc={xcc}"
+      reg_objs[x] = y

  with open(sys.argv[1], 'r') as f:
    log_content = f.read()

-  processed_log = parse_amdgpu_logs(log_content, reg_names, only_xcc0=only_xcc0)
+  processed_log = parse_amdgpu_logs(log_content, reg_names, reg_objs, only_xcc0=only_xcc0)

  with open(sys.argv[2], 'w') as f:
    f.write(processed_log)
@@ -1,136 +0,0 @@
-import os, sys, struct
-sys.path.append(os.getcwd())
-# PROFILE=1 to use
-#os.environ["PROFILE"] = "1"
-os.environ["SQTT"] = "1"
-os.environ["SQTT_ITRACE_SE_MASK"] = "1"
-os.environ["SQTT_LIMIT_SE"] = "1"
-import xml.etree.ElementTree as ET
-
-from tinygrad import nn, Tensor, Device
-from tinygrad.helpers import get_single_element
-from tinygrad.runtime.support.elf import elf_loader
-from tinygrad.runtime.ops_amd import ProfileSQTTEvent
-from extra.sqtt.attempt_sqtt_parse import parse_sqtt_print_packets
-
-def disassemble(text, root:ET.Element):
-  i = 0
-  while i < len(text):
-    ins = struct.unpack("I", text[i:i+4])[0]
-
-    # 1. Get the encoding
-    did_match = False
-    for enc_el in root.findall("./ISA/Encodings/Encoding"):
-      mask = enc_el.findtext("EncodingIdentifierMask")
-      assert len(mask)%32 == 0
-      bit_mask = int(mask, 2)
-      iden = [int(x.text, 2) for x in enc_el.find("EncodingIdentifiers").findall("EncodingIdentifier")]
-      for ide in iden:
-        if ins&bit_mask == ide:
-          did_match = True
-          break
-      if did_match: break
-    if not did_match: raise RuntimeError(f"unknown instruction {ins:08X}")
-    if len(mask) >= 64: ins = (struct.unpack("I", text[i+4:i+8])[0]<<32) | ins
-    if len(mask) >= 96: ins = (struct.unpack("I", text[i+8:i+12])[0]<<64) | ins
-    encoding_name = enc_el.findtext("EncodingName")
-
-    #print(ET.tostring(enc_el).decode())
-
-    # 2. Parse the Fields for this Encoding
-    field_data = {}
-    for field in enc_el.findall("MicrocodeFormat/BitMap/Field"):
-      # Fields can be split into multiple ranges (RangeCount > 1)
-      ranges = sorted(field.findall("BitLayout/Range"), key=lambda x: int(x.attrib.get('Order')))
-      val = 0
-      current_shift = 0
-      for rng in ranges:
-        width = int(rng.find("BitCount").text)
-        chunk = (ins >> int(rng.find("BitOffset").text)) & ((1 << width) - 1)
-        val |= (chunk << current_shift)
-        current_shift += width
-      field_data[field.find("FieldName").text] = val
-    # this is already used
-    del field_data["ENCODING"]
-
-    # 3. Extract the instruction
-    did_match = False
-    for ins_el in root.findall("./ISA/Instructions/Instruction"):
-      ins_name = ins_el.findtext("InstructionName")
-      for ins_enc in ins_el.findall("InstructionEncodings/InstructionEncoding"):
-        if ins_enc.findtext("EncodingName") == encoding_name:
-          opcode = int(ins_enc.findtext("Opcode"))
-          if "OP" in field_data and opcode == field_data["OP"]:
-            did_match = True
-            del field_data["OP"]
-            break
-        if did_match: break
-      if did_match: break
-
-    #print(ET.tostring(ins_enc).decode())
-    #print()
-    #print(field_data)
-    if not did_match:
-      print(f"{i:4X} : {ins:16x} -- {encoding_name}")
-    elif did_match:
-      params = []
-      #print(ET.tostring(ins_el).decode())
-
-      # 4. Extract the opcodes
-      for op_ins in ins_enc.findall("Operands/Operand"):
-        op_type = op_ins.findtext("OperandType")
-        op_size = op_ins.findtext("OperandSize")
-        op_fmt = op_ins.findtext("DataFormatName")
-        op_field_name = op_ins.findtext("FieldName")
-        if op_field_name is None: continue
-        assert op_field_name in field_data
-        # loop through operands for compare
-        for op_el in root.findall("./ISA/OperandTypes/OperandType"):
-          test_op_type = op_el.findtext("OperandTypeName")
-          val_dict = {}
-          for op_val in op_el.findall("OperandPredefinedValues/PredefinedValue"):
-            val_dict[int(op_val.findtext("Value"))] = op_val.findtext("Name")
-          if op_type == test_op_type:
-            if field_data[op_field_name] in val_dict:
-              print(op_type, op_size, op_fmt)
-              params.append(val_dict[field_data[op_field_name]])
-            else:
-              params.append(f"{op_type}({field_data[op_field_name]})")
-            del field_data[op_field_name]
-            #print(op_type, op_size, op_fmt, op_el, op_field_name,
-            #      field_data[op_field_name],
-            #      val_dict.get(field_data[op_field_name], "<UNK>"))
-            #print(ET.tostring(op_el).decode())
-
-      print(f"{i:4X} : {ins:16x} -- {ins_name.lower()} {', '.join(params)}", field_data)
-
-    # advance
-    i += len(mask) // 8
-
-  #print(ET.tostring(root).decode())
-
-if __name__ == "__main__":
-  # human readable manual at https://docs.amd.com/v/u/en-US/rdna35_instruction_set_architecture
-  fns = nn.state.zip_extract(Tensor.from_url("https://gpuopen.com/download/machine-readable-isa/latest/"))
-  xml_str = fns['amdgpu_isa_rdna3_5.xml'].to("CPU").data()
-  with open("/tmp/rdna35.xml", "wb") as f: f.write(bytes(xml_str))
-  root = ET.fromstring(xml_str)
-
-  a = Tensor.empty(16)+1
-  for ei in a.schedule():
-    ei.lower()
-    # get text
-    _, hdr, _ = elf_loader(ei.prg.lib)
-    text = get_single_element([x for x in hdr if x.name==".text"]).content
-
-    # llvm disassembler
-    Device["AMD"].compiler.disassemble(ei.prg.lib)
-
-    # run program
-    ei.run()
-
-  sqtt_events = [e for e in Device["AMD"].profile_events if isinstance(e, ProfileSQTTEvent)]
-  for e in sqtt_events[0:1]: # only the first SE
-    parse_sqtt_print_packets(e.blob)
-
-  disassemble(text[:0x40], root)
@@ -1,15 +0,0 @@
-from tinygrad import Tensor, nn
-import xml.etree.ElementTree as ET
-
-if __name__ == "__main__":
-  # human readable manual at https://docs.amd.com/v/u/en-US/rdna35_instruction_set_architecture
-  fns = nn.state.zip_extract(Tensor.from_url("https://gpuopen.com/download/machine-readable-isa/latest/"))
-  xml_str = fns['amdgpu_isa_rdna3_5.xml'].to("CPU").data()
-  root = ET.fromstring(xml_str)
-
-  for op_el in root.findall("./ISA/OperandTypes/OperandType"):
-    op_name = op_el.findtext("OperandTypeName")
-    val_dict = {}
-    for op_val in op_el.findall("OperandPredefinedValues/PredefinedValue"):
-      val_dict[int(op_val.findtext("Value"))] = op_val.findtext("Name")
-    print(op_name, val_dict)
@@ -1,4 +0,0 @@
-*.deb
-build
-src
-sniffer/sniff.so
@@ -1,20 +0,0 @@
-Built ROCT-Thunk-Interface (hsakmt)
-  hsakmt-roct-dev_5.4.4.99999-local_amd64.deb
-  note: installs to /opt/rocm
-Built ROCm-Device-Libs
-  Works with ROCM_PATH=/home/tiny/build/ROCm-Device-Libs/build/dist
-  rocm-device-libs_1.0.0.99999-local_amd64.deb
-Built ROCm-CompilerSupport (amd_comgr)
-  no deb, sudo make install to /usr/local
-Built ROCR-Runtime
-  hsa-rocr_1.8.0-local_amd64.deb
-  hsa-rocr-dev_1.8.0-local_amd64.deb
-Built ROCm-OpenCL-Runtime
-  rocm-ocl-icd_2.0.0-local_amd64.deb
-  ISSUE: these depend on "comgr"
-  rocm-opencl_2.0.0-local_amd64.deb
-  rocm-opencl-dev_2.0.0-local_amd64.deb
-  Did sudo make install
-
-
-
@@ -1,41 +0,0 @@
-# run two "rocm-bandwidth-test" in a loop
-# amdgpu-6.0.5-1581431.20.04
-# fixed in kernel 6.2.14
-
-[   72.153646] RIP: 0010:pm_send_runlist+0x4a/0x630 [amdgpu]
-[   72.153815] Code: 30 65 48 8b 04 25 28 00 00 00 48 89 45 d0 31 c0 80 fb 01 0f 87 aa 9d 49 00 83 e3 01 0f 85 1c 05 00 00 49 8b 3f b8 01 00 00 00 <48> 8b 97 30 01 00 00 44 8b b7 6c 01 00 00 8b 9f 70 01 00 00 8b 8a
-[   72.153900] RSP: 0018:ffffb48445c03c30 EFLAGS: 00010246
-[   72.153928] RAX: 0000000000000001 RBX: 0000000000000000 RCX: 0000000000000000
-[   72.153962] RDX: 000000000000007b RSI: ffff9395e1562558 RDI: 0000000000000000
-[   72.153996] RBP: ffffb48445c03cb8 R08: 0000000000000000 R09: 0000000000000001
-[   72.154030] R10: ffff9395c900d840 R11: 0000000000000000 R12: 0000000000000000
-[   72.154065] R13: ffff9395c9e00400 R14: 0000000000000001 R15: ffff9395e15624e0
-[   72.154099] FS:  00007f345c6463c0(0000) GS:ffff93a4aee80000(0000) knlGS:0000000000000000
-[   72.154137] CS:  0010 DS: 0000 ES: 0000 CR0: 0000000080050033
-[   72.154165] CR2: 0000000000000130 CR3: 0000000112840000 CR4: 0000000000750ee0
-[   72.154201] PKRU: 55555554
-[   72.154215] Call Trace:
-[   72.154230]  <TASK>
-[   72.154244]  map_queues_cpsch+0x75/0xc0 [amdgpu]
-[   72.154365]  debug_map_and_unlock+0x51/0x90 [amdgpu]
-[   72.154480]  debug_refresh_runlist+0x1f/0x30 [amdgpu]
-[   72.154591]  kfd_dbg_runtime_disable+0x13c/0x240 [amdgpu]
-[   72.154705]  kfd_ioctl_dbg_set_debug_trap+0x69d/0x8b0 [amdgpu]
-[   72.154820]  kfd_ioctl+0x24a/0x5b0 [amdgpu]
-[   72.154925]  ? kfd_ioctl_create_queue+0x770/0x770 [amdgpu]
-[   72.155035]  ? syscall_exit_to_user_mode+0x27/0x50
-[   72.155061]  ? exit_to_user_mode_prepare+0x3d/0x1c0
-[   72.155088]  __x64_sys_ioctl+0x95/0xd0
-[   72.155109]  do_syscall_64+0x5c/0xc0
-[   72.155128]  ? syscall_exit_to_user_mode+0x27/0x50
-[   72.155151]  ? do_syscall_64+0x69/0xc0
-[   72.155172]  entry_SYSCALL_64_after_hwframe+0x61/0xcb
-[   72.155198] RIP: 0033:0x7f345c7f63ab
-[   72.155218] Code: 0f 1e fa 48 8b 05 e5 7a 0d 00 64 c7 00 26 00 00 00 48 c7 c0 ff ff ff ff c3 66 0f 1f 44 00 00 f3 0f 1e fa b8 10 00 00 00 0f 05 <48> 3d 01 f0 ff ff 73 01 c3 48 8b 0d b5 7a 0d 00 f7 d8 64 89 01 48
-[   72.155301] RSP: 002b:00007ffc97cc89f8 EFLAGS: 00000246 ORIG_RAX: 0000000000000010
-[   72.155339] RAX: ffffffffffffffda RBX: 00007ffc97cc8a30 RCX: 00007f345c7f63ab
-[   72.155375] RDX: 00007ffc97cc8a30 RSI: 00000000c0284b82 RDI: 0000000000000003
-[   72.155411] RBP: 00000000c0284b82 R08: 0000000000000000 R09: 0000000000000000
-[   72.155447] R10: 00007f345cd4ddb0 R11: 0000000000000246 R12: 00007ffc97cc8a30
-[   72.155481] R13: 0000000000000003 R14: 00007ffc97cc8d20 R15: 0000000000000000
-[   72.155517]  </TASK>
@@ -1,41 +0,0 @@
-# run two tinygrad matrix example in a loop
-# amdgpu-6.0.5-1581431.20.04
-# NOT fixed in kernel 6.2.14
-
-[  553.016624] gmc_v11_0_process_interrupt: 30 callbacks suppressed
-[  553.016631] amdgpu 0000:0b:00.0: amdgpu: [gfxhub] page fault (src_id:0 ring:24 vmid:9 pasid:32770, for process python3 pid 10001 thread python3 pid 10001)
-[  553.016790] amdgpu 0000:0b:00.0: amdgpu:   in page starting at address 0x00007f0000000000 from client 10
-[  553.016892] amdgpu 0000:0b:00.0: amdgpu: GCVM_L2_PROTECTION_FAULT_STATUS:0x00901A30
-[  553.016974] amdgpu 0000:0b:00.0: amdgpu:      Faulty UTCL2 client ID: SDMA0 (0xd)
-[  553.017051] amdgpu 0000:0b:00.0: amdgpu:      MORE_FAULTS: 0x0
-[  553.017111] amdgpu 0000:0b:00.0: amdgpu:      WALKER_ERROR: 0x0
-[  553.017173] amdgpu 0000:0b:00.0: amdgpu:      PERMISSION_FAULTS: 0x3
-[  553.017238] amdgpu 0000:0b:00.0: amdgpu:      MAPPING_ERROR: 0x0
-[  553.017300] amdgpu 0000:0b:00.0: amdgpu:      RW: 0x0
-[  553.123921] [drm:mes_v11_0_submit_pkt_and_poll_completion.constprop.0 [amdgpu]] *ERROR* MES failed to response msg=2
-[  553.124153] amdgpu: failed to add hardware queue to MES, doorbell=0x1a16
-[  553.124195] amdgpu: MES might be in unrecoverable state, issue a GPU reset
-[  553.124237] amdgpu: Failed to restore queue 2
-[  553.124266] amdgpu: Failed to restore process queues
-[  553.124270] amdgpu: Failed to evict queue 3
-[  553.124297] amdgpu: amdgpu_amdkfd_restore_userptr_worker: Failed to resume KFD
-
-# alternative crash in kernel 6.2.14
-
-[  151.097948] gmc_v11_0_process_interrupt: 30 callbacks suppressed
-[  151.097953] amdgpu 0000:0b:00.0: amdgpu: [gfxhub] page fault (src_id:0 ring:24 vmid:8 pasid:32771, for process python3 pid 7525 thread python3 pid 7525)
-[  151.097993] amdgpu 0000:0b:00.0: amdgpu:   in page starting at address 0x00007f0000000000 from client 10
-[  151.098008] amdgpu 0000:0b:00.0: amdgpu: GCVM_L2_PROTECTION_FAULT_STATUS:0x00801A30
-[  151.098020] amdgpu 0000:0b:00.0: amdgpu:      Faulty UTCL2 client ID: SDMA0 (0xd)
-[  151.098032] amdgpu 0000:0b:00.0: amdgpu:      MORE_FAULTS: 0x0
-[  151.098042] amdgpu 0000:0b:00.0: amdgpu:      WALKER_ERROR: 0x0
-[  151.098052] amdgpu 0000:0b:00.0: amdgpu:      PERMISSION_FAULTS: 0x3
-[  151.098062] amdgpu 0000:0b:00.0: amdgpu:      MAPPING_ERROR: 0x0
-[  151.098071] amdgpu 0000:0b:00.0: amdgpu:      RW: 0x0
-[  151.209517] [drm:mes_v11_0_submit_pkt_and_poll_completion.constprop.0 [amdgpu]] *ERROR* MES failed to response msg=2
-[  151.209724] amdgpu: failed to add hardware queue to MES, doorbell=0x1002
-[  151.209734] amdgpu: MES might be in unrecoverable state, issue a GPU reset
-[  151.209743] amdgpu: Failed to restore queue 1
-[  151.209751] amdgpu: Failed to restore process queues
-[  151.209759] amdgpu: amdgpu_amdkfd_restore_userptr_worker: Failed to resume KFD
-[  151.209858] amdgpu 0000:0b:00.0: amdgpu: GPU reset begin!
@@ -1,20 +0,0 @@
-# two tinygrad + two bandwidth test
-# RDNA2, driver 6.0.5
-# recovered from this!
-
-[  136.971209] gmc_v10_0_process_interrupt: 39 callbacks suppressed
-[  136.971218] amdgpu 0000:0b:00.0: amdgpu: [gfxhub] page fault (src_id:0 ring:24 vmid:11 pasid:32773, for process rocm-bandwidth- pid 20281 thread rocm-bandwidth- pid 20281)
-[  136.971228] amdgpu 0000:0b:00.0: amdgpu:   in page starting at address 0x00007f5c2b800000 from client 0x1b (UTCL2)
-[  136.971232] amdgpu 0000:0b:00.0: amdgpu: GCVM_L2_PROTECTION_FAULT_STATUS:0x00B01A31
-[  136.971233] amdgpu 0000:0b:00.0: amdgpu:      Faulty UTCL2 client ID: SDMA0 (0xd)
-[  136.971235] amdgpu 0000:0b:00.0: amdgpu:      MORE_FAULTS: 0x1
-[  136.971236] amdgpu 0000:0b:00.0: amdgpu:      WALKER_ERROR: 0x0
-[  136.971236] amdgpu 0000:0b:00.0: amdgpu:      PERMISSION_FAULTS: 0x3
-[  136.971237] amdgpu 0000:0b:00.0: amdgpu:      MAPPING_ERROR: 0x0
-[  136.971238] amdgpu 0000:0b:00.0: amdgpu:      RW: 0x0
-...
-[  136.993979] amdgpu 0000:0b:00.0: amdgpu: IH ring buffer overflow (0x000BE5A0, 0x0003C480, 0x0003E5C0)
-[  138.209072] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00004000 flags=0x0000]
-[  138.209078] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00004d80 flags=0x0000]
-[  138.209081] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00005000 flags=0x0000]
-[  138.209084] amdgpu 0000:0b:00.0: AMD-Vi: Event logged [IO_PAGE_FAULT domain=0x001a address=0x7c00005d80 flags=0x0000]
@@ -1,33 +0,0 @@
-# ROCK-Kernel-Driver 0b579de9622f5c93021dcb7927d13926313740a2
-# non fatal "crash"
-
-[  127.418045] ------------[ cut here ]------------
-[  127.418046] User pages unexpectedly invalid
-[  127.418056] WARNING: CPU: 16 PID: 260 at drivers/gpu/drm/amd/amdgpu/amdgpu_amdkfd_gpuvm.c:3000 amdgpu_amdkfd_restore_userptr_worker+0x4d9/0x500 [amdgpu]
-[  127.418235] Modules linked in: rfcomm cmac algif_hash algif_skcipher af_alg bnep nls_iso8859_1 iwlmvm mac80211 intel_rapl_msr intel_rapl_common edac_mce_amd snd_hda_codec_realtek snd_hda_codec_generic snd_hda_codec_hdmi kvm_amd binfmt_misc snd_hda_intel snd_intel_dspcfg kvm libarc4 snd_intel_sdw_acpi snd_hda_codec btusb iwlwifi btrtl snd_hda_core btbcm btintel irqbypass btmtk snd_hwdep crct10dif_pclmul snd_pcm polyval_clmulni bluetooth snd_seq_midi snd_seq_midi_event snd_rawmidi snd_seq polyval_generic cfg80211 ghash_clmulni_intel eeepc_wmi snd_seq_device snd_timer aesni_intel asus_wmi ecdh_generic snd platform_profile crypto_simd ledtrig_audio cryptd ecc ccp soundcore sparse_keymap rapl k10temp wmi_bmof mac_hid sch_fq_codel msr parport_pc ppdev lp parport ramoops pstore_blk efi_pstore reed_solomon pstore_zone ip_tables x_tables autofs4 amdgpu hid_generic usbhid hid i2c_algo_bit drm_ttm_helper ttm video iommu_v2 drm_buddy gpu_sched drm_display_helper drm_kms_helper syscopyarea
-[  127.418276]  sysfillrect sysimgblt fb_sys_fops drm nvme nvme_core cec r8169 ahci crc32_pclmul rc_core i2c_piix4 xhci_pci libahci nvme_common xhci_pci_renesas realtek wmi
-[  127.418284] CPU: 16 PID: 260 Comm: kworker/16:1 Tainted: G        W          6.0.0 #4
-[  127.418286] Hardware name: System manufacturer System Product Name/TUF GAMING X570-PLUS (WI-FI), BIOS 3603 03/20/2021
-[  127.418287] Workqueue: events amdgpu_amdkfd_restore_userptr_worker [amdgpu]
-[  127.418455] RIP: 0010:amdgpu_amdkfd_restore_userptr_worker+0x4d9/0x500 [amdgpu]
-[  127.418601] Code: ff e8 2b 8a 96 d1 e9 66 fe ff ff 48 c7 c7 40 4f f5 c0 e8 56 7b 8a d1 0f 0b e9 2e ff ff ff 48 c7 c7 d8 d0 ed c0 e8 43 7b 8a d1 <0f> 0b e9 0a fe ff ff 4c 89 ef e8 f8 89 96 d1 e9 cb fd ff ff e8 ce
-[  127.418603] RSP: 0018:ffffb36740a83dc8 EFLAGS: 00010282
-[  127.418604] RAX: 0000000000000000 RBX: ffff9d159ee9df30 RCX: 0000000000000027
-[  127.418605] RDX: 0000000000000027 RSI: ffffb36740a83c88 RDI: ffff9d242a220568
-[  127.418606] RBP: ffffb36740a83e58 R08: ffff9d242a220560 R09: 0000000000000001
-[  127.418607] R10: 0000000000000001 R11: 0000000000000020 R12: ffff9d159ee9df98
-[  127.418607] R13: ffff9d159ee9df70 R14: ffff9d159ee9dee0 R15: ffff9d159ee9dee0
-[  127.418608] FS:  0000000000000000(0000) GS:ffff9d242a200000(0000) knlGS:0000000000000000
-[  127.418609] CS:  0010 DS: 0000 ES: 0000 CR0: 0000000080050033
-[  127.418610] CR2: 00007fd5d4715000 CR3: 0000000120ffe000 CR4: 0000000000750ee0
-[  127.418611] PKRU: 55555554
-[  127.418611] Call Trace:
-[  127.418612]  <TASK>
-[  127.418613]  process_one_work+0x21f/0x3f0
-[  127.418615]  worker_thread+0x4a/0x3c0
-[  127.418617]  ? process_one_work+0x3f0/0x3f0
-[  127.418618]  kthread+0xf0/0x120
-[  127.418619]  ? kthread_complete_and_exit+0x20/0x20
-[  127.418620]  ret_from_fork+0x22/0x30
-[  127.418622]  </TASK>
-[  127.418623] ---[ end trace 0000000000000000 ]---
@@ -1,80 +0,0 @@
-import numpy as np
-import pathlib
-from hexdump import hexdump
-from tinygrad.helpers import colored
-from extra.helpers import enable_early_exec
-early_exec = enable_early_exec()
-
-from tinygrad.runtime.ops_cl import CLProgram, CLBuffer, ROCM_LLVM_PATH
-
-ENABLE_NON_ASM = False
-
-WMMA = True
-DUAL_ALU = True
-F32 = True
-
-if ENABLE_NON_ASM:
-  buf = CLBuffer.fromCPU(np.zeros(10, np.float32))
-  prg_empty = CLProgram("code", "__kernel void code(__global float *a) { a[0] = 1; }")
-  asm_real = prg_empty.binary()
-  with open("/tmp/cc.elf", "wb") as f:
-    f.write(asm_real)
-  prg_empty([1], [1], buf, wait=True)
-  print(buf.toCPU())
-
-print(colored("creating CLBuffer", "green"))
-buf = CLBuffer.fromCPU(np.zeros(10, np.float32))
-code = open(pathlib.Path(__file__).parent / "prog.s", "r").read()
-
-gen = []
-FLOPS = 0
-MAX_REG = 251
-for j in range(1):
-  if WMMA:
-    KY, KX = 4, 4
-    for y in range(KY):
-      for x in range(KX):
-        c = (y*KX+x)*8
-        a = (KY*KX*8) + y*8
-        b = (KY*KX*8) + (KY*8) + x*8
-        gen.append(f"v_wmma_f32_16x16x16_f16 v[{c}:{c+7}], v[{a}:{a+7}], v[{b}:{b+7}], v[{c}:{c+7}]")
-        FLOPS += 16*8*2
-  else:
-    for i in range(0, MAX_REG, 6):
-      if DUAL_ALU:
-        if F32:
-          gen.append(f"v_dual_fmac_f32 v{i+0}, v{i+1}, v{i+2} :: v_dual_fmac_f32 v{i+3}, v{i+4}, v{i+5}")
-          FLOPS += 4
-        else:
-          gen.append(f"v_dual_dot2acc_f32_f16 v{i+0}, v{i+1}, v{i+2} :: v_dual_dot2acc_f32_f16 v{i+3}, v{i+4}, v{i+5}")
-          FLOPS += 8
-      else:
-        assert F32
-        gen.append(f"v_fmac_f32 v{i+0}, v{i+1}, v{i+2}")
-        gen.append(f"v_fmac_f32 v{i+3}, v{i+4}, v{i+5}")
-code = code.replace("// FLOPS", '\n'.join(gen))
-print(code)
-
-
-# fix: COMGR failed to get code object ISA name. set triple to 'amdgcn-amd-amdhsa'
-
-object = early_exec(([ROCM_LLVM_PATH / "llvm-mc", '--arch=amdgcn', '--mcpu=gfx1100', '--triple=amdgcn-amd-amdhsa', '--filetype=obj', '-'], code.encode("utf-8")))
-asm = early_exec(([ROCM_LLVM_PATH / "ld.lld", "/dev/stdin", "-o", "/dev/stdout", "--pie"], object))
-
-with open("/tmp/cc2.o", "wb") as f:
-  f.write(object)
-with open("/tmp/cc2.elf", "wb") as f:
-  f.write(asm)
-
-print(colored("creating CLProgram", "green"))
-prg = CLProgram("code", asm)
-
-print(colored("running program", "green"))
-G = 512
-FLOPS *= 100000*G*G  # loop * global_size
-for i in range(3):
-  tm = prg(buf, global_size=[G//256, G, 1], local_size=[256, 1, 1], wait=True)
-  print(f"ran in {tm*1e3:.2f} ms, {FLOPS/(tm*1e9):.2f} GFLOPS")
-
-print(colored("transferring buffer", "green"))
-print(buf.toCPU())
@@ -1,80 +0,0 @@
-.global _start
-_start:
-.rodata
-.align 0x10
-.global code.kd
-.type code.kd,STT_OBJECT
-# amd_kernel_code_t (must be at 0x440 for kernel_code_entry_byte_offset to be right)
-code.kd:
-# amd_kernel_..., amd_machine_...
-.long 0,0,0,0
-# kernel_code_entry_byte_offset, kernel_code_prefetch_byte_offset
-.long 0x00000bc0,0x00000000,0x00000000,0x00000000
-# kernel_code_prefetch_byte_size, max_scratch_backing_memory_byte_size
-.long 0,0,0,0
-# compute_pgm_rsrc1, compute_pgm_rsrc2, kernel_code_properties, workitem_private_segment_byte_size
-.long 0x60af0000,0x0000009e,0x00000408,0x00000000
-# compute_pgm_rsrc1 |= AMD_COMPUTE_PGM_RSRC_ONE_FLOAT_DENORM_MODE_32 | AMD_COMPUTE_PGM_RSRC_ONE_FLOAT_DENORM_MODE_16_64
-# compute_pgm_rsrc1 |= AMD_COMPUTE_PGM_RSRC_ONE_ENABLE_DX10_CLAMP | AMD_COMPUTE_PGM_RSRC_ONE_ENABLE_IEEE_MODE
-# compute_pgm_rsrc2 |= AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT = 0xF
-# compute_pgm_rsrc2 |= AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_X
-# kernel_code_properties |= AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_KERNARG_SEGMENT_PTR = 1
-# kernel_code_properties |= AMD_KERNEL_CODE_PROPERTIES_RESERVED1 = 1
-.text
-.global code
-.type code,STT_FUNC
-code:
-# https://llvm.org/docs/AMDGPUUsage.html#initial-kernel-execution-state
-# s[0:1] contains the kernarg_address
-# TODO: can we use s[2:3] if this was really a wave since we only alloced 2 SGPRs?
-s_load_b64 s[2:3], s[0:1], null
-
-s_mov_b32 s8, 0
-loop:
-s_addk_i32 s8, 1
-s_cmp_eq_u32 s8, 100000
-// FLOPS
-s_cbranch_scc0 loop
-
-# wait for the s_load_b64
-s_waitcnt lgkmcnt(0)
-
-v_dual_mov_b32 v0, 4 :: v_dual_mov_b32 v1, 2.0
-global_store_b32 v0, v1, s[2:3]
-
-# Deallocate all VGPRs for this wave. Use only when next instruction is S_ENDPGM.
-s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)
-s_endpgm
-s_code_end
-
-.amdgpu_metadata
-amdhsa.kernels:
-  - .args:
-      - .address_space:  global
-        .name:           a
-        .offset:         0
-        .size:           8
-        .type_name:      'float*'
-        .value_kind:     global_buffer
-    .group_segment_fixed_size: 0
-    .kernarg_segment_align: 8
-    .kernarg_segment_size: 8
-    .language:       OpenCL C
-    .language_version:
-      - 1
-      - 2
-    .max_flat_workgroup_size: 256
-    .name:           code
-    .private_segment_fixed_size: 0
-    .sgpr_count:     2
-    .sgpr_spill_count: 0
-    .symbol:         code.kd
-    .uses_dynamic_stack: false
-    .vgpr_count:     256
-    .vgpr_spill_count: 0
-    .wavefront_size: 32
-amdhsa.target:   amdgcn-amd-amdhsa--gfx1100
-amdhsa.version:
-  - 1
-  - 2
-.end_amdgpu_metadata
@@ -1,11 +0,0 @@
-#!/bin/bash
-mkdir -p src
-cd src
-git clone https://github.com/RadeonOpenCompute/ROCT-Thunk-Interface.git -b rocm-5.5.0
-git clone https://github.com/RadeonOpenCompute/ROCm-Device-Libs.git -b rocm-5.5.0
-git clone https://github.com/RadeonOpenCompute/llvm-project.git -b rocm-5.5.0 --depth 1
-git clone https://github.com/RadeonOpenCompute/ROCR-Runtime.git -b rocm-5.5.0
-git clone https://github.com/ROCm-Developer-Tools/ROCclr.git -b rocm-5.5.0
-git clone https://github.com/RadeonOpenCompute/ROCm-CompilerSupport.git -b rocm-5.5.0
-git clone https://github.com/RadeonOpenCompute/ROCm-OpenCL-Runtime.git -b rocm-5.5.0
-cd ../
@@ -1,69 +0,0 @@
-#!/bin/bash
-mkdir -p build/debs
-cd build
-
-# ROCT-Thunk-Interface (hsakmt)
-if [ ! -f debs/hsakmt-roct-dev_5.5.0.99999-local_amd64.deb ]
-then
-  mkdir -p ROCT-Thunk-Interface
-  cd ROCT-Thunk-Interface
-  cmake ../../src/ROCT-Thunk-Interface
-  make -j32 package
-  cp hsakmt-roct-dev_5.5.0.99999-local_amd64.deb ../debs
-  cd ../
-fi
-
-
-# build custom LLVM
-if [ ! -f llvm-project/bin/clang ]
-then
-  mkdir -p llvm-project
-  cd llvm-project
-  cmake -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_PROJECTS="llvm;clang;lld" -DLLVM_TARGETS_TO_BUILD="AMDGPU;X86" ../../src/llvm-project/llvm
-  make -j32
-  cd ..
-fi
-
-# use custom LLVM
-export PATH="$PWD/llvm-project/bin:$PATH"
-
-# ROCm-Device-Libs
-if [ ! -f debs/rocm-device-libs_1.0.0.99999-local_amd64.deb ]
-then
-  mkdir -p ROCm-Device-Libs
-  cd ROCm-Device-Libs
-  cmake ../../src/ROCm-Device-Libs
-  make -j32 package
-  cp rocm-device-libs_1.0.0.99999-local_amd64.deb ../debs
-  cd ../
-fi
-
-# ROCR-Runtime
-if [ ! -f debs/hsa-rocr_1.8.0-local_amd64.deb ]
-then
-  mkdir -p ROCR-Runtime
-  cd ROCR-Runtime
-  cmake ../../src/ROCR-Runtime/src
-  make -j32 package
-  cp hsa-rocr_1.8.0-local_amd64.deb ../debs
-  cp hsa-rocr-dev_1.8.0-local_amd64.deb ../debs
-  cd ../
-fi
-
-# ROCm-OpenCL-Runtime (needs ROCclr)
-if [ ! -f debs/rocm-opencl_2.0.0-local_amd64.deb ]
-then
-  mkdir -p ROCm-OpenCL-Runtime
-  cd ROCm-OpenCL-Runtime
-  cmake ../../src/ROCm-OpenCL-Runtime
-  make -j32 package
-  cp rocm-opencl_2.0.0-local_amd64.deb ../debs
-  cp rocm-opencl-dev_2.0.0-local_amd64.deb ../debs
-  cp rocm-ocl-icd_2.0.0-local_amd64.deb ../debs
-fi
-
-# ROCm-CompilerSupport (broken)
-#mkdir -p ROCm-CompilerSupport
-#cd ROCm-CompilerSupport
-#cmake ../../src/ROCm-CompilerSupport/lib/comgr
-#make -j32
@@ -1,14 +0,0 @@
-#!/bin/bash
-rm amdgpu-install_5.5.50500-1_all.deb
-wget https://repo.radeon.com/amdgpu-install/5.5/ubuntu/$(lsb_release -cs)/amdgpu-install_5.5.50500-1_all.deb
-sudo dpkg -i amdgpu-install_5.5.50500-1_all.deb
-sudo apt-get update
-
-# kernel driver
-sudo apt-get install amdgpu-dkms
-
-# for opencl
-sudo apt-get install rocm-opencl-runtime
-
-# for HIP
-sudo apt-get install hip-runtime-amd rocm-device-libs hip-dev
@@ -1,11 +0,0 @@
-#!/bin/bash -e
-clang sniff.cc -Werror -shared -fPIC -I../src/ -I../src/ROCT-Thunk-Interface/include -I../src/ROCm-Device-Libs/ockl/inc -o sniff.so -lstdc++
-#AMD_LOG_LEVEL=4 HSAKMT_DEBUG_LEVEL=7 LD_PRELOAD=$PWD/sniff.so /home/tiny/build/HIP-Examples/HIP-Examples-Applications/HelloWorld/HelloWorld
-#AMD_LOG_LEVEL=4 LD_PRELOAD=$PWD/sniff.so $HOME/build/HIP-Examples/HIP-Examples-Applications/HelloWorld/HelloWorld
-#AMD_LOG_LEVEL=5 LD_PRELOAD=$PWD/sniff.so python3 ../rdna3/asm.py
-DEBUG=5 LD_PRELOAD=$PWD/sniff.so python3 ../rdna3/asm.py
-#AMD_LOG_LEVEL=5 HSAKMT_DEBUG_LEVEL=7 DEBUG=5 LD_PRELOAD=$PWD/sniff.so strace -F python3 ../rdna3/asm.py
-#LD_PRELOAD=$PWD/sniff.so python3 ../rdna3/asm.py
-#AMD_LOG_LEVEL=4 LD_PRELOAD=$PWD/sniff.so FORWARD_ONLY=1 DEBUG=2 python3 ../../../test/test_ops.py TestOps.test_add
-#AMD_LOG_LEVEL=4 HSAKMT_DEBUG_LEVEL=7 LD_PRELOAD=$PWD/sniff.so rocm-bandwidth-test -s 0 -d 1 -m 1
-#AMD_LOG_LEVEL=4 HSAKMT_DEBUG_LEVEL=7 LD_PRELOAD=$PWD/sniff.so rocm-bandwidth-test -s 1 -d 2 -m 1
@@ -1,282 +0,0 @@
-// template copied from https://github.com/geohot/cuda_ioctl_sniffer/blob/master/sniff.cc
-
-#include <stdio.h>
-#include <stdlib.h>
-#include <string.h>
-#include <dlfcn.h>
-#include <signal.h>
-#include <ucontext.h>
-
-#include <sys/mman.h>
-
-// includes from the ROCm sources
-#include <linux/kfd_ioctl.h>
-#include <hsa.h>
-#include <amd_hsa_kernel_code.h>
-#include <ROCR-Runtime/src/core/inc/sdma_registers.h>
-using namespace rocr::AMD;
-
-#include <string>
-#include <map>
-std::map<int, std::string> files;
-std::map<uint64_t, uint64_t> ring_base_addresses;
-
-#define D(args...) fprintf(stderr, args)
-
-uint64_t doorbell_offset = -1;
-std::map<uint64_t, int> queue_types;
-
-void hexdump(void *d, int l) {
-  for (int i = 0; i < l; i++) {
-    if (i%0x10 == 0 && i != 0) printf("\n");
-    if (i%0x10 == 8) printf(" ");
-    if (i%0x10 == 0) printf("%8X: ", i);
-    printf("%2.2X ", ((uint8_t*)d)[i]);
-  }
-  printf("\n");
-}
-
-extern "C" {
-
-// https://defuse.ca/online-x86-assembler.htm#disassembly2
-static void handler(int sig, siginfo_t *si, void *unused) {
-  ucontext_t *u = (ucontext_t *)unused;
-  uint8_t *rip = (uint8_t*)u->uc_mcontext.gregs[REG_RIP];
-
-  int store_size = 0;
-  uint64_t value;
-  if (rip[0] == 0x48 && rip[1] == 0x89 && rip[2] == 0x30) {
-    // 0:  48 89 30                mov    QWORD PTR [rax],rsi
-    store_size = 8;
-    value = u->uc_mcontext.gregs[REG_RSI];
-    u->uc_mcontext.gregs[REG_RIP] += 3;
-  } else if (rip[0] == 0x4c && rip[1] == 0x89 && rip[2] == 0x28) {
-    // 0:  4c 89 28                mov    QWORD PTR [rax],r13
-    store_size = 8;
-    value = u->uc_mcontext.gregs[REG_R13];
-    u->uc_mcontext.gregs[REG_RIP] += 3;
-  } else {
-    D("segfault %02X %02X %02X %02X %02X %02X %02X %02X rip: %p addr: %p\n", rip[0], rip[1], rip[2], rip[3], rip[4], rip[5], rip[6], rip[7], rip, si->si_addr);
-    D("rax: %llx rcx: %llx rdx: %llx rsi: %llx rbx: %llx\n", u->uc_mcontext.gregs[REG_RAX], u->uc_mcontext.gregs[REG_RCX], u->uc_mcontext.gregs[REG_RDX], u->uc_mcontext.gregs[REG_RSI], u->uc_mcontext.gregs[REG_RBX]);
-    exit(-1);
-  }
-
-  uint64_t ring_base_address = ring_base_addresses[((uint64_t)si->si_addr)&0xFFF];
-  int queue_type = queue_types[((uint64_t)si->si_addr)&0xFFF];
-  D("%16p: \u001b[31mDING DONG\u001b[0m (queue_type %d) store(%d): 0x%8lx -> %p ring_base_address:0x%lx\n", rip, queue_type, store_size, value, si->si_addr, ring_base_address);
-
-  if (queue_type == KFD_IOC_QUEUE_TYPE_SDMA) {
-    uint8_t *sdma_ptr = (uint8_t*)(ring_base_address);
-    while (sdma_ptr < ((uint8_t*)(ring_base_address)+value)) {
-      D("0x%3lx: ", sdma_ptr-(uint8_t*)(ring_base_address));
-      if (sdma_ptr[0] == SDMA_OP_TIMESTAMP) {
-        D("SDMA_PKT_TIMESTAMP\n");
-        sdma_ptr += sizeof(SDMA_PKT_TIMESTAMP);
-      } else if (sdma_ptr[0] == SDMA_OP_GCR) {
-        D("SDMA_PKT_GCR\n");
-        sdma_ptr += sizeof(SDMA_PKT_GCR);
-      } else if (sdma_ptr[0] == SDMA_OP_ATOMIC) {
-        D("SDMA_PKT_ATOMIC\n");
-        sdma_ptr += sizeof(SDMA_PKT_ATOMIC);
-      } else if (sdma_ptr[0] == SDMA_OP_FENCE) {
-        D("SDMA_PKT_FENCE\n");
-        sdma_ptr += sizeof(SDMA_PKT_FENCE);
-      } else if (sdma_ptr[0] == SDMA_OP_TRAP) {
-        D("SDMA_PKT_TRAP\n");
-        sdma_ptr += sizeof(SDMA_PKT_TRAP);
-      } else if (sdma_ptr[0] == SDMA_OP_COPY && sdma_ptr[1] == SDMA_SUBOP_COPY_LINEAR) {
-        SDMA_PKT_COPY_LINEAR *pkt = (SDMA_PKT_COPY_LINEAR *)sdma_ptr;
-        D("SDMA_PKT_COPY_LINEAR: count:0x%x src:0x%lx dst:0x%lx\n", pkt->COUNT_UNION.count+1,
-          (uint64_t)pkt->SRC_ADDR_LO_UNION.src_addr_31_0 | ((uint64_t)pkt->SRC_ADDR_HI_UNION.src_addr_63_32 << 32),
-          (uint64_t)pkt->DST_ADDR_LO_UNION.dst_addr_31_0 | ((uint64_t)pkt->DST_ADDR_HI_UNION.dst_addr_63_32 << 32)
-        );
-        sdma_ptr += sizeof(SDMA_PKT_COPY_LINEAR);
-      } else {
-        D("unhandled packet type %d %d, exiting\n", sdma_ptr[0], sdma_ptr[1]);
-        break;
-      }
-    }
-
-    //hexdump((void*)(ring_base_address), 0x100);
-  } else if (queue_type == KFD_IOC_QUEUE_TYPE_COMPUTE_AQL) {
-    hsa_kernel_dispatch_packet_t *pkt = (hsa_kernel_dispatch_packet_t *)(ring_base_address+value*0x40);
-    if ((pkt->header&0xFF) == HSA_PACKET_TYPE_KERNEL_DISPATCH) {
-      D("HSA_PACKET_TYPE_KERNEL_DISPATCH -- setup:%d workgroup[%d, %d, %d] grid[%d, %d, %d] kernel_object:0x%lx kernarg_address:%p\n", pkt->setup, pkt->workgroup_size_x, pkt->workgroup_size_y, pkt->workgroup_size_z, pkt->grid_size_x, pkt->grid_size_y, pkt->grid_size_z, pkt->kernel_object, pkt->kernarg_address);
-      amd_kernel_code_t *code = (amd_kernel_code_t *)pkt->kernel_object;
-      D("kernel_code_entry_byte_offset:%lx\n", code->kernel_code_entry_byte_offset);
-      uint32_t *kernel_code = (uint32_t*)(pkt->kernel_object + code->kernel_code_entry_byte_offset);
-      int code_len = 0;
-      while (kernel_code[code_len] != 0xbf9f0000 && kernel_code[code_len] != 0) code_len++;
-      hexdump(kernel_code, code_len*4);
-      /*FILE *f = fopen("/tmp/kernel_code", "wb");
-      fwrite(kernel_code, 4, code_len, f);
-      fclose(f);
-      system("python -c 'print(\" \".join([(\"0x%02X\"%x) for x in open(\"/tmp/kernel_code\", \"rb\").read()]))' | ../build/llvm-project/bin/llvm-mc --disassemble --arch=amdgcn --mcpu=gfx1100 --show-encoding");*/
-      D("kernargs (kernarg_segment_byte_size:0x%lx)\n", code->kernarg_segment_byte_size);
-      // get length
-      int i;
-      for (i = 0; i < 0x400; i+=0x10) {
-        if (memcmp((void*)((uint64_t)pkt->kernarg_address+i), "\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00", 0x10) == 0) break;
-      }
-      hexdump((void*)pkt->kernarg_address, i+0x10);
-    } else if ((pkt->header&0xFF) == HSA_PACKET_TYPE_BARRIER_AND) {
-      hsa_barrier_and_packet_t *pkt_and = (hsa_barrier_and_packet_t *)(ring_base_address+value*0x40);
-      D("HSA_PACKET_TYPE_BARRIER_AND completion_signal:0x%lx\n", pkt_and->completion_signal.handle);
-      //hexdump((void*)(ring_base_address+value*0x40), 0x40);
-    } else if ((pkt->header&0xFF) == HSA_PACKET_TYPE_VENDOR_SPECIFIC) {
-      D("HSA_PACKET_TYPE_VENDOR_SPECIFIC\n");
-      hexdump((void*)(ring_base_address+value*0x40), 0x40);
-    } else {
-      hexdump((void*)(ring_base_address+value*0x40), 0x40);
-    }
-  }
-
-  mprotect((void *)((uint64_t)si->si_addr & ~0xFFF), 0x2000, PROT_READ | PROT_WRITE);
-  if (store_size == 8) {
-    *(volatile uint64_t*)(si->si_addr) = value;
-  } else if (store_size == 4) {
-    *(volatile uint32_t*)(si->si_addr) = value;
-  } else if (store_size == 2) {
-    *(volatile uint16_t*)(si->si_addr) = value;
-  } else {
-    D("store size not supported\n");
-    exit(-1);
-  }
-  mprotect((void *)((uint64_t)si->si_addr & ~0xFFF), 0x2000, PROT_NONE);
-}
-
-void register_sigsegv_handler() {
-  struct sigaction sa = {0};
-  sa.sa_flags = SA_SIGINFO;
-  sigemptyset(&sa.sa_mask);
-  sa.sa_sigaction = handler;
-  if (sigaction(SIGSEGV, &sa, NULL) == -1) {
-    D("ERROR: failed to register sigsegv handler");
-    exit(-1);
-  }
-  // NOTE: python (or ocl runtime?) blocks the SIGSEGV signal
-  sigset_t x;
-  sigemptyset(&x);
-  sigaddset(&x, SIGSEGV);
-  sigprocmask(SIG_UNBLOCK, &x, NULL);
-}
-
-int (*my_open)(const char *pathname, int flags, mode_t mode);
-#undef open
-int open(const char *pathname, int flags, mode_t mode) {
-  if (my_open == NULL) my_open = reinterpret_cast<decltype(my_open)>(dlsym(RTLD_NEXT, "open"));
-  int ret = my_open(pathname, flags, mode);
-  //D("open %s (0o%o) = %d\n", pathname, flags, ret);
-  files[ret] = pathname;
-  return ret;
-}
-
-
-int (*my_open64)(const char *pathname, int flags, mode_t mode);
-#undef open
-int open64(const char *pathname, int flags, mode_t mode) {
-  if (my_open64 == NULL) my_open64 = reinterpret_cast<decltype(my_open64)>(dlsym(RTLD_NEXT, "open64"));
-  int ret = my_open64(pathname, flags, mode);
-  //D("open %s (0o%o) = %d\n", pathname, flags, ret);
-  files[ret] = pathname;
-  return ret;
-}
-
-void *(*my_mmap)(void *addr, size_t length, int prot, int flags, int fd, off_t offset);
-#undef mmap
-void *mmap(void *addr, size_t length, int prot, int flags, int fd, off_t offset) {
-  if (my_mmap == NULL) my_mmap = reinterpret_cast<decltype(my_mmap)>(dlsym(RTLD_NEXT, "mmap"));
-  void *ret = my_mmap(addr, length, prot, flags, fd, offset);
-
-  if (doorbell_offset != -1 && offset == doorbell_offset) {
-    D("HIDDEN DOORBELL %p, handled by %p\n", addr, handler);
-    register_sigsegv_handler();
-    mprotect(addr, length, PROT_NONE);
-  }
-
-  if (fd != -1) D("mmapped %p (target %p) with flags 0x%x length 0x%zx fd %d %s offset 0x%lx\n", ret, addr, flags, length, fd, files[fd].c_str(), offset);
-  return ret;
-}
-
-void *(*my_mmap64)(void *addr, size_t length, int prot, int flags, int fd, off_t offset);
-#undef mmap64
-void *mmap64(void *addr, size_t length, int prot, int flags, int fd, off_t offset) { return mmap(addr, length, prot, flags, fd, offset); }
-
-int ioctl_num = 1;
-int (*my_ioctl)(int filedes, unsigned long request, void *argp) = NULL;
-#undef ioctl
-int ioctl(int filedes, unsigned long request, void *argp) {
-  if (my_ioctl == NULL) my_ioctl = reinterpret_cast<decltype(my_ioctl)>(dlsym(RTLD_NEXT, "ioctl"));
-  int ret = 0;
-  ret = my_ioctl(filedes, request, argp);
-  if (!files.count(filedes)) return ret;
-
-  uint8_t type = (request >> 8) & 0xFF;
-  uint8_t nr = (request >> 0) & 0xFF;
-  uint16_t size = (request >> 16) & 0xFFF;
-
-  D("%3d: %d = %3d(%20s) 0x%3x ", ioctl_num, ret, filedes, files[filedes].c_str(), size);
-
-  if (request == AMDKFD_IOC_SET_EVENT) {
-    kfd_ioctl_set_event_args *args = (kfd_ioctl_set_event_args *)argp;
-    D("AMDKFD_IOC_SET_EVENT event_id:%d", args->event_id);
-  } else if (request == AMDKFD_IOC_ALLOC_MEMORY_OF_GPU) {
-    kfd_ioctl_alloc_memory_of_gpu_args *args = (kfd_ioctl_alloc_memory_of_gpu_args *)argp;
-    D("AMDKFD_IOC_ALLOC_MEMORY_OF_GPU va_addr:0x%llx size:0x%llx handle:%llX gpu_id:0x%x", args->va_addr, args->size, args->handle, args->gpu_id);
-  } else if (request == AMDKFD_IOC_MAP_MEMORY_TO_GPU) {
-    kfd_ioctl_map_memory_to_gpu_args *args = (kfd_ioctl_map_memory_to_gpu_args *)argp;
-    D("AMDKFD_IOC_MAP_MEMORY_TO_GPU handle:%llX", args->handle);
-  } else if (request == AMDKFD_IOC_CREATE_EVENT) {
-    kfd_ioctl_create_event_args *args = (kfd_ioctl_create_event_args *)argp;
-    D("AMDKFD_IOC_CREATE_EVENT event_page_offset:0x%llx event_type:%d event_id:%d", args->event_page_offset, args->event_type, args->event_id);
-  } else if (request == AMDKFD_IOC_WAIT_EVENTS) {
-    D("AMDKFD_IOC_WAIT_EVENTS");
-  } else if (request == AMDKFD_IOC_SET_XNACK_MODE) {
-    D("AMDKFD_IOC_SET_XNACK_MODE");
-  } else if (request == AMDKFD_IOC_SVM || (type == 0x4b && nr == 0x20)) {
-    // NOTE: this one is variable length
-    kfd_ioctl_svm_args *args = (kfd_ioctl_svm_args *)argp;
-    D("AMDKFD_IOC_SVM start_addr:0x%llx size:0x%llx op:%d", args->start_addr, args->size, args->op);
-  } else if (request == AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU) {
-    kfd_ioctl_unmap_memory_from_gpu_args *args = (kfd_ioctl_unmap_memory_from_gpu_args *)argp;
-    D("AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU handle:%llX", args->handle);
-  } else if (request == AMDKFD_IOC_FREE_MEMORY_OF_GPU) {
-    D("AMDKFD_IOC_FREE_MEMORY_OF_GPU");
-  } else if (request == AMDKFD_IOC_SET_SCRATCH_BACKING_VA) {
-    D("AMDKFD_IOC_SET_SCRATCH_BACKING_VA");
-  } else if (request == AMDKFD_IOC_GET_TILE_CONFIG) {
-    D("AMDKFD_IOC_GET_TILE_CONFIG");
-  } else if (request == AMDKFD_IOC_SET_TRAP_HANDLER) {
-    D("AMDKFD_IOC_SET_TRAP_HANDLER");
-  } else if (request == AMDKFD_IOC_GET_VERSION) {
-    kfd_ioctl_get_version_args *args = (kfd_ioctl_get_version_args *)argp;
-    D("AMDKFD_IOC_GET_VERSION major_version:%d minor_version:%d", args->major_version, args->minor_version);
-  } else if (request == AMDKFD_IOC_GET_PROCESS_APERTURES_NEW) {
-    D("AMDKFD_IOC_GET_PROCESS_APERTURES_NEW");
-  } else if (request == AMDKFD_IOC_ACQUIRE_VM) {
-    D("AMDKFD_IOC_ACQUIRE_VM");
-  } else if (request == AMDKFD_IOC_SET_MEMORY_POLICY) {
-    D("AMDKFD_IOC_SET_MEMORY_POLICY");
-  } else if (request == AMDKFD_IOC_GET_CLOCK_COUNTERS) {
-    D("AMDKFD_IOC_GET_CLOCK_COUNTERS");
-  } else if (request == AMDKFD_IOC_CREATE_QUEUE) {
-    kfd_ioctl_create_queue_args *args = (kfd_ioctl_create_queue_args *)argp;
-    D("AMDKFD_IOC_CREATE_QUEUE\n");
-    D("queue_type:%d ring_base_address:0x%llx\n", args->queue_type, args->ring_base_address);
-    D("eop_buffer_address:0x%llx ctx_save_restore_address:0x%llx\n", args->eop_buffer_address, args->ctx_save_restore_address);
-    D("ring_size:0x%x queue_priority:%d\n", args->ring_size, args->queue_priority);
-    D("RETURNS write_pointer_address:0x%llx read_pointer_address:0x%llx doorbell_offset:0x%llx queue_id:%d\n", args->write_pointer_address, args->read_pointer_address, args->doorbell_offset, args->queue_id);
-    //D("RETURNS *write_pointer_address:0x%llx *read_pointer_address:0x%llx\n", *(uint64_t*)args->write_pointer_address, *(uint64_t*)args->read_pointer_address);
-    ring_base_addresses[args->doorbell_offset&0xFFF] = args->ring_base_address;
-    queue_types[args->doorbell_offset&0xFFF] = args->queue_type;
-    doorbell_offset = args->doorbell_offset&~0xFFF;
-  } else {
-    D("type:0x%x nr:0x%x size:0x%x", type, nr, size);
-  }
-
-  D("\n");
-  ioctl_num++;
-  return ret;
-}
-
-}
@@ -67,12 +67,11 @@ def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,in
  forward_args = ",".join(f"{dtype}{'*' if name not in symbolic_vars.values() else ''} {name}" for name,dtype,_ in (outputs+inputs if wasm else inputs+outputs))

  if not wasm:
-    thread_id = 0 # NOTE: export does not support threading, thread_id is always 0
    for name,cl in bufs_to_save.items():
      weight = ''.join(["\\x%02X"%x for x in bytes(to_mv(cl._buf.va_addr, cl._buf.size))])
      cprog.append(f"unsigned char {name}_data[] = \"{weight}\";")
    cprog += [f"{dtype_map[dtype]} {name}[{len}];" if name not in bufs_to_save else f"{dtype_map[dtype]} *{name} = ({dtype_map[dtype]} *){name}_data;" for name,(len,dtype,_key) in bufs.items() if name not in input_names+output_names]
-    cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)}, {thread_id});" for (name, args, _global_size, _local_size) in statements] + ["}"]
+    cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)});" for (name, args, _global_size, _local_size) in statements] + ["}"]
    return '\n'.join(headers + cprog)
  else:
    if bufs_to_save:
@@ -0,0 +1,101 @@
+from typing import Callable, Any
+from tinygrad import Tensor, dtypes, nn, UOp
+from tinygrad.uop.ops import KernelInfo, AxisType, Ops
+
+def quantize_to_fp8(x: Tensor, dtype=dtypes.fp8e4m3):
+  fp8_min = -448.0 if dtype == dtypes.fp8e4m3 else -57344.0
+  fp8_max = 448.0 if dtype == dtypes.fp8e4m3 else 57344.0
+  x_abs_max = x.abs().max().detach()
+  scale = fp8_max / (x_abs_max + 1e-8)
+  x_scaled = x * scale
+  x_det = x_scaled.detach()
+  x_clamped = x_det.clamp(fp8_min, fp8_max)
+  x_clamped_ste = x_scaled + (x_clamped - x_det)
+  res = x_clamped_ste.cast(dtype)
+  return res, scale.float().reciprocal()
+
+def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
+  SEQ = inp.shape[1]
+  OUT = weight.shape[0]
+  IN = weight.shape[-1]
+  seq_idx = UOp.range(SEQ, 2, AxisType.LOOP)
+  out_idx = UOp.range(OUT, 3, AxisType.LOOP)
+  batch_idx = UOp.range(output.size//SEQ//OUT, 1, AxisType.LOOP)
+  reduce_idx = UOp.range(IN, 0, AxisType.REDUCE)
+  product = (inp.index((seq_idx*IN+reduce_idx+batch_idx*IN*SEQ)) * weight.index((out_idx*IN+reduce_idx))).cast(dtypes.float)
+  reduced = product.reduce(reduce_idx, arg=Ops.ADD)
+  store_op = output.index((seq_idx*OUT+out_idx+batch_idx*OUT*SEQ), ptr=True).store(reduced).end(batch_idx, seq_idx, out_idx)
+  return store_op.sink(arg=KernelInfo(name=f"fp8_matmul_{inp.shape}x{weight.shape}"))
+
+def custom_matmul_backward(gradient: UOp, kernel: UOp) -> tuple[UOp, UOp]:
+  _, input_uop, weight_uop = kernel.src[1:]
+  input_tensor = Tensor(input_uop, device=input_uop.device)
+  grad_tensor = Tensor(gradient, device=gradient.device)
+  weight_tensor = Tensor(weight_uop, device=weight_uop.device)
+  grad_quantized, scale = quantize_to_fp8(grad_tensor)
+  scale_scalar = scale.reshape(())
+  grad_weight = Tensor.einsum("bso,bsi->oi", grad_quantized, input_tensor, dtype=dtypes.float)
+  grad_weight = grad_weight * scale_scalar
+  grad_2d = grad_quantized.reshape(grad_tensor.shape[0] * grad_tensor.shape[1], grad_tensor.shape[-1])
+  grad_input = (grad_2d.dot(weight_tensor, dtype=dtypes.float)).contiguous().reshape(input_tensor.shape) * scale
+  return (None, grad_input.uop, grad_weight.uop)
+
+class FP8Linear:
+  def __init__(self, in_features:int, out_features:int, bias:bool=True):
+    self.weight = Tensor.empty(out_features, in_features, dtype=dtypes.float32)
+    self.bias = Tensor.empty(out_features, dtype=dtypes.float32) if bias else None
+
+  def __call__(self, x: Tensor) -> Tensor:
+    original_ndim = len(x.shape)
+    if original_ndim == 2: x = x.reshape(x.shape[0], 1, x.shape[1])
+    batch, seq, _ = x.shape
+    w_fp8, w_scale = quantize_to_fp8(self.weight)
+    x_fp8, x_scale = quantize_to_fp8(x)
+    GPUS = self.weight.device
+    if isinstance(GPUS, tuple) and len(GPUS) > 1:
+      y = Tensor(Tensor.empty((batch//len(GPUS), seq, self.weight.shape[0]), dtype=dtypes.float, device=GPUS).uop.multi(0), device=GPUS)
+    else:
+      y = Tensor.empty((batch, seq, self.weight.shape[0]), dtype=dtypes.float)
+    y = Tensor.custom_kernel(y, x_fp8, w_fp8, fxn=custom_matmul, grad_fxn=custom_matmul_backward)[0]
+    y = y * w_scale * x_scale
+    if self.bias is not None: y = y + self.bias
+    if original_ndim == 2: y = y.reshape(batch, self.weight.shape[0])
+    return y.cast(x.dtype)
+
+def _replace_linear(layer: nn.Linear):
+  fp8_linear = FP8Linear(layer.weight.shape[1], layer.weight.shape[0], layer.bias is not None)
+  fp8_linear.weight = layer.weight
+  if layer.bias is not None: fp8_linear.bias = layer.bias
+  return fp8_linear
+
+def _swap_linear_with_fp8(model, module_filter_fn:Callable[[Any, str],bool]|None=None, fqn:str="", parent:Any|None=None,
+                          attr_name:str="", visited:set|None=None):
+  if visited is None: visited = set()
+  if id(model) in visited: return
+  visited.add(id(model))
+  if isinstance(model, (str, int, float, bool, type(None), Tensor, UOp)): return
+  elif isinstance(model, nn.Linear):
+    if module_filter_fn is not None and not module_filter_fn(model, fqn): return
+    fp8_linear = _replace_linear(model)
+    if parent is not None and attr_name:
+      setattr(parent, attr_name, fp8_linear)
+  elif isinstance(model, list):
+    for i, item in enumerate(model):
+      child_fqn = f"{fqn}.{i}" if fqn else str(i)
+      if isinstance(item, nn.Linear) and (module_filter_fn is None or module_filter_fn(item, child_fqn)): model[i] = _replace_linear(item)
+      else: _swap_linear_with_fp8(item, module_filter_fn, child_fqn, None, "", visited)
+  elif isinstance(model, dict):
+    for key, item in list(model.items()):
+      child_fqn = f"{fqn}.{key}" if fqn else str(key)
+      if isinstance(item, nn.Linear) and (module_filter_fn is None or module_filter_fn(item, child_fqn)): model[key] = _replace_linear(item)
+      else: _swap_linear_with_fp8(item, module_filter_fn, child_fqn, None, "", visited)
+  elif hasattr(model, "__dict__"):
+    for attr_key in list(vars(model).keys()):
+      try: attr = getattr(model, attr_key)
+      except Exception: continue
+      child_fqn = f"{fqn}.{attr_key}" if fqn else attr_key
+      _swap_linear_with_fp8(attr, module_filter_fn, child_fqn, model, attr_key, visited)
+
+def convert_to_float8_training(model, module_filter_fn:Callable[[Any,str],bool]|None=None):
+  _swap_linear_with_fp8(model, module_filter_fn, "", None, "")
+  return model
@@ -1,3 +1,2 @@
-*.s
 *.ll
 fp32_sgemm_amd
@@ -0,0 +1,504 @@
+# RDNA3 128x128 tiled GEMM kernel - DSL version
+# Computes C = A @ B for 4096x4096 float32 matrices using 128x128 tiles
+#
+# Architecture: RDNA3 (gfx1100)
+# Tile size: 128x128 (each workgroup computes one tile of C)
+# Workgroup: 128 threads (arranged as 32x4 for coalesced memory access)
+# Inner loop: 8 iterations per K-block, processing 8 columns of A and 8 rows of B
+#
+# Accumulators: 128 vgprs (v[2-129])
+
+import numpy as np
+from tinygrad import Tensor, Device, Context, GlobalCounters
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.helpers import getenv, colored
+from tinygrad.dtype import dtypes, AddrSpace
+from tinygrad.engine.realize import Estimates
+from tinygrad.renderer.amd.dsl import s, v, VCC_LO, NULL
+from tinygrad.runtime.autogen.amd.rdna3.ins import *
+
+# =============================================================================
+# Kernel constants
+# =============================================================================
+LDS_SIZE = 8320       # Local data share size in bytes
+MATRIX_DIM = 4096     # Matrix dimension N (assumes square NxN matrices)
+LDS_A_STRIDE = 0x210  # LDS stride for A tile (528 bytes)
+LDS_B_STRIDE = 0x200  # LDS stride for B tile (512 bytes)
+LDS_BASE_OFFSET = 0x1080  # Base LDS offset for tiles
+ADDR_MASK = 0x3fffff80    # Address alignment mask
+
+# =============================================================================
+# Named register assignments (VGPRs)
+# =============================================================================
+V_LANE_ID = 0             # lane_id set on startup
+# Use tile gaps (v146-159) for named regs to minimize max VGPR
+V_LANE_ID_MOD8 = 146      # lane_id & 7
+V_LANE_MOD8_X4 = 147      # (lane_id & 7) << 2
+V_LANE_DIV8_X4 = 150      # ((lane_id >> 3) & 3) << 2
+V_LDS_B_BASE = 151        # LDS B-tile base address for inner loop
+V_LDS_A_BASE = 154        # LDS A-tile base address for inner loop
+V_GLOBAL_A_ADDR = 155     # global memory A prefetch address
+V_GLOBAL_B_ADDR = 158     # global memory B prefetch address
+V_LDS_A_ADDR = 159        # single base register for A stores
+V_LDS_B_ADDR = 162        # single base register for B stores
+
+# LDS tile register destinations - SEPARATE from DATA to avoid overlap
+# A on banks 2-3, B on banks 0-1 to avoid bank conflicts in VOPD
+V_A_TILE_REGS = [130, 134, 138, 142]  # A tile: banks 2,2,2,2 (130%4=2, etc.)
+V_B_TILE_REGS = [132, 136, 140, 144, 148, 152, 156, 160]  # B tile: banks 0,0,0,0,0,0,0,0
+
+# =============================================================================
+# Named register assignments (SGPRs)
+# =============================================================================
+S_OUT_PTR = (0, 1)        # output C matrix base pointer
+S_WORKGROUP_X = 2         # workgroup_id_x (system SGPR, follows user SGPRs)
+S_WORKGROUP_Y = 3         # workgroup_id_y (system SGPR)
+S_DIM_N = 4               # matrix dimension N
+S_LOOP_BOUND = 7          # K-8 (loop termination bound)
+S_LOOP_CTR = 12           # loop counter (increments by 8)
+S_PREFETCH_FLAG = 13      # prefetch condition flag / row stride in epilogue
+S_TILE_X = 14             # workgroup_x << 7
+S_TILE_Y = 15             # workgroup_y << 7
+# Kernarg load destinations
+S_KERNARG_A = (20, 21)    # A pointer from kernarg
+S_KERNARG_B = (22, 23)    # B pointer from kernarg
+# Prefetch base pointers (8 pairs each, 16KB/256KB apart)
+S_PREFETCH_B = 24         # s[24:39] - 8 B tile pointers
+S_PREFETCH_A = 40         # s[40:55] - 8 A tile pointers
+
+# =============================================================================
+# Data tables
+# =============================================================================
+
+# Accumulator grid: ACC_GRID[a_idx][b_idx] = vgpr for C[a,b]
+# a_idx: which A value (0-7), b_idx: which B value (0-15)
+# Scattered due to VOPD bank constraints (vdst_x % 4 != vdst_y % 4)
+# Range is from v2 - v129
+ACC_GRID = [
+  [  5,  3,  9,  8,   37, 35, 41, 40,   69, 67, 73, 72,  101, 99,105,104],  # a0
+  [  4,  2,  7,  6,   36, 34, 39, 38,   68, 66, 71, 70,  100, 98,103,102],  # a1
+  [ 17, 16, 13, 11,   49, 48, 45, 43,   81, 80, 77, 75,  113,112,109,107],  # a2
+  [ 15, 14, 12, 10,   47, 46, 44, 42,   79, 78, 76, 74,  111,110,108,106],  # a3
+  [ 21, 19, 25, 24,   53, 51, 57, 56,   85, 83, 89, 88,  117,115,121,120],  # a4
+  [ 20, 18, 23, 22,   52, 50, 55, 54,   84, 82, 87, 86,  116,114,123,122],  # a5
+  [125,128, 29, 27,   33, 32, 61, 59,   65, 64, 93, 91,   97, 96,129,127],  # a6
+  [119,118, 28, 26,   31, 30, 60, 58,   63, 62, 92, 90,   95, 94,124,126],  # a7
+]
+
+# Optimized (a_pair, b_pair) iteration order for better GPU scheduling
+# Interleaves A and B pairs to maximize instruction-level parallelism
+FMAC_PAIR_ORDER = [
+  (0,0),(0,1),(1,1),(1,0), (2,0),(2,1),(3,1),(3,2), (0,2),(0,3),(1,3),(1,2), (2,2),(2,3),(3,3),(3,4),
+  (0,4),(0,5),(1,5),(1,4), (2,4),(2,5),(3,5),(3,6), (0,6),(0,7),(1,7),(1,6), (2,6),(2,7),(3,7),(3,0),
+]
+
+def derive_fmac_pattern(acc_grid, a_tile_regs=None, b_tile_regs=None):
+  """Generate 64 dual FMAC ops from accumulator grid with optimized iteration order."""
+  pattern = []
+  for idx, (a_pair, b_pair) in enumerate(FMAC_PAIR_ORDER):
+    a_even, a_odd = a_pair * 2, a_pair * 2 + 1
+    b_even, b_odd = b_pair * 2, b_pair * 2 + 1
+    a_base, b_base = a_tile_regs[a_pair], b_tile_regs[b_pair]
+    # Op 1: normal order -> C[a_even, b_even] + C[a_odd, b_odd]
+    pattern.append((acc_grid[a_even][b_even], acc_grid[a_odd][b_odd],
+                   a_base, b_base, a_base+1, b_base+1))
+    # Op 2: alternate swapping A vs B to vary register banks
+    if idx % 2 == 0:  # swap B
+      pattern.append((acc_grid[a_even][b_odd], acc_grid[a_odd][b_even],
+                     a_base, b_base+1, a_base+1, b_base))
+    else:  # swap A
+      pattern.append((acc_grid[a_odd][b_even], acc_grid[a_even][b_odd],
+                     a_base+1, b_base, a_base, b_base+1))
+  return pattern
+
+# Derived: 64 dual FMAC operations
+FMAC_PATTERN = derive_fmac_pattern(ACC_GRID, V_A_TILE_REGS, V_B_TILE_REGS)
+
+def derive_permute_swaps(acc_grid, out_regs):
+  """Derive swap sequence to permute accumulators from FMAC layout to output order.
+
+  After FMAC loop: acc_grid[a][b] holds C[a,b]
+  Output order: for row_half in 0,1; col_group in 0-3; row_in_group in 0-3; b_off in 0-3
+    -> need C[row_half*4 + row_in_group, col_group*4 + b_off] in specified reg order
+  """
+  def target_ab(i):
+    row_half, col_group = i // 64, (i // 16) % 4
+    row_in_group, b_off = (i // 4) % 4, i % 4
+    return (row_half * 4 + row_in_group, col_group * 4 + b_off)
+
+  reg_contents = {acc_grid[a][b]: (a, b) for a in range(8) for b in range(16)}
+  ab_location = {ab: r for r, ab in reg_contents.items()}
+
+  swaps = []
+  for i in range(128):
+    target_reg, needed_ab = out_regs[i], target_ab(i)
+    current_reg = ab_location[needed_ab]
+    if current_reg != target_reg:
+      swaps.append((current_reg, target_reg))
+      ab_at_target = reg_contents.get(target_reg)
+      reg_contents[target_reg], ab_location[needed_ab] = needed_ab, target_reg
+      if ab_at_target is not None:
+        reg_contents[current_reg], ab_location[ab_at_target] = ab_at_target, current_reg
+  return swaps
+
+# Derived: swap sequence to arrange accumulators for output
+# Each group of 4 registers is ascending for direct global_store_b128
+OUT_REGS = [r for i in range(32) for r in range(126 - i*4, 130 - i*4)]
+PERMUTE_SWAPS = derive_permute_swaps(ACC_GRID, OUT_REGS)
+
+# =============================================================================
+# LDS tile staging registers
+# =============================================================================
+# DATA regs receive contiguous global prefetch, then write to LDS
+# TILE regs receive scattered LDS loads (ds_load_b64 pairs), then feed FMACs
+# Contiguous layout with mod4=[3,0,1,2,3,0,1,2] for bank conflict avoidance
+V_LDS_A_DATA = [163, 164, 165, 166, 167, 168, 169, 170]
+V_LDS_B_DATA = [171, 172, 173, 174, 175, 176, 177, 178]
+
+# Initial tile prefetch: (vdst, saddr_lo) - load into A data regs using B prefetch pointers (s[24:31])
+INIT_PREFETCH = [(V_LDS_A_DATA[i], S_PREFETCH_B+2*i) for i in range(4)]
+
+# Global memory prefetch schedule: (vdst1, vdst2, addr_vreg, saddr_lo1, saddr_lo2)
+# First 2 pairs from B prefetch pointers (s[32:39]), next 4 pairs from A prefetch pointers (s[40:55])
+PREFETCH_LOADS = [(V_LDS_A_DATA[4+2*i], V_LDS_A_DATA[4+2*i+1], V_GLOBAL_B_ADDR, S_PREFETCH_B+8+4*i, S_PREFETCH_B+10+4*i) for i in range(2)] + \
+                 [(V_LDS_B_DATA[2*(i-2)], V_LDS_B_DATA[2*(i-2)+1], V_GLOBAL_A_ADDR, S_PREFETCH_A+4*(i-2), S_PREFETCH_A+2+4*(i-2)) for i in range(2, 6)]
+
+# =============================================================================
+# Kernel class
+# =============================================================================
+
+class Kernel:
+  def __init__(self, arch='gfx1100'): self.instructions, self.labels, self.pos, self.arch = [], {}, 0, arch
+  def label(self, name): self.labels[name] = self.pos
+
+  def emit(self, inst, target=None):
+    self.instructions.append(inst)
+    inst._target, inst._pos = target, self.pos
+    self.pos += inst.size()
+    return inst
+
+  def waitcnt(self, lgkm=None, vm=None):
+    """Wait for memory operations. lgkm=N waits until N lgkm ops remain, vm=N waits until N vmem ops remain."""
+    vmcnt, lgkmcnt, expcnt = vm if vm is not None else 63, lgkm if lgkm is not None else 63, 7
+    waitcnt = (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
+    self.emit(s_waitcnt(simm16=waitcnt))
+
+  def finalize(self):
+    """Patch branch offsets and return the finalized instruction list."""
+    for inst in self.instructions:
+      if inst._target is None: continue
+      offset_dwords = (self.labels[inst._target] - inst._pos - inst.size()) // 4
+      if not -32768 <= offset_dwords <= 32767: raise ValueError(f"branch to '{inst._target}' offset {offset_dwords} exceeds simm16 range")
+      inst.simm16 = offset_dwords
+    return self.instructions
+
+
+# =============================================================================
+# Kernel builder
+# =============================================================================
+
+def build_kernel(arch='gfx1100'):
+  k = Kernel(arch)
+
+  # ===========================================================================
+  # PROLOGUE: Load kernel arguments, compute tile coordinates and addresses
+  # ===========================================================================
+  k.emit(s_load_b128(sdata=s[S_KERNARG_A[0]:S_KERNARG_B[1]], sbase=s[0:1], offset=0x0, soffset=NULL))
+  k.emit(s_load_b64(sdata=s[S_OUT_PTR[0]:S_OUT_PTR[1]], sbase=s[0:1], offset=0x10, soffset=NULL))
+  k.emit(s_mov_b32(s[S_DIM_N], MATRIX_DIM))
+  k.emit(s_mov_b32(s[S_LOOP_CTR], 0))  # used by LDS swizzle, always 0 for valid workgroups
+  k.emit(s_lshl_b32(s[S_TILE_X], s[S_WORKGROUP_X], 7))
+  k.emit(s_lshl_b32(s[S_TILE_Y], s[S_WORKGROUP_Y], 7))
+
+  # Lane-derived values
+  k.emit(v_and_b32_e32(v[V_LANE_ID_MOD8], 7, v[V_LANE_ID]))
+  k.emit(v_lshrrev_b32_e32(v[4], 3, v[V_LANE_ID]))
+  k.emit(v_or_b32_e32(v[1], s[S_TILE_X], v[V_LANE_ID]))
+  k.emit(v_or_b32_e32(v[22], s[S_TILE_Y], v[4]))
+  k.emit(v_lshlrev_b32_e32(v[V_LANE_MOD8_X4], 2, v[V_LANE_ID_MOD8]))
+  k.waitcnt(lgkm=0)
+
+  # Compute 8 A and B matrix tile base pointers for prefetch
+  k.emit(s_mov_b64(s[S_PREFETCH_B:S_PREFETCH_B+1], s[S_KERNARG_B[0]:S_KERNARG_B[1]]))  # B[0]: no offset
+  for i in range(1, 8):  # B: 16KB apart
+    k.emit(s_add_u32(s[S_PREFETCH_B+i*2], s[S_KERNARG_B[0]], i * 0x4000))
+    k.emit(s_addc_u32(s[S_PREFETCH_B+i*2+1], s[S_KERNARG_B[1]], 0))
+  k.emit(s_mov_b64(s[S_PREFETCH_A:S_PREFETCH_A+1], s[S_KERNARG_A[0]:S_KERNARG_A[1]]))  # A[0]: no offset
+  for i in range(1, 8):  # A: 256KB apart
+    k.emit(s_add_u32(s[S_PREFETCH_A+i*2], s[S_KERNARG_A[0]], i * 0x40000))
+    k.emit(s_addc_u32(s[S_PREFETCH_A+i*2+1], s[S_KERNARG_A[1]], 0))
+
+  # Global prefetch addresses: B = (tile_x + lane_id) * 4, A = ((tile_y << 12) + (lane_id/8)*4K + lane_id%8) * 4
+  k.emit(v_add_nc_u32_e32(v[V_GLOBAL_B_ADDR], s[S_TILE_X], v[V_LANE_ID]))
+  k.emit(v_lshlrev_b32_e32(v[V_GLOBAL_B_ADDR], 2, v[V_GLOBAL_B_ADDR]))
+  k.emit(s_lshl_b32(s[19], s[S_TILE_Y], 12))
+  k.emit(v_lshl_add_u32(v[V_GLOBAL_A_ADDR], v[4], 12, v[V_LANE_ID_MOD8]))  # (lane_id/8)*4K + lane_id%8
+  k.emit(v_add_nc_u32_e32(v[V_GLOBAL_A_ADDR], s[19], v[V_GLOBAL_A_ADDR]))
+  k.emit(v_lshlrev_b32_e32(v[V_GLOBAL_A_ADDR], 2, v[V_GLOBAL_A_ADDR]))
+
+  # Do initial loads
+  for vdst, saddr_lo in INIT_PREFETCH:
+    k.emit(global_load_b32(vdst=v[vdst], addr=v[V_GLOBAL_B_ADDR], saddr=s[saddr_lo:saddr_lo+1]))
+  for iter in range(6):
+    vdst1, vdst2, addr, slo1, slo2 = PREFETCH_LOADS[iter]
+    k.emit(global_load_b32(vdst=v[vdst1], addr=v[addr], saddr=s[slo1:slo1+1]))
+    k.emit(global_load_b32(vdst=v[vdst2], addr=v[addr], saddr=s[slo2:slo2+1]))
+
+  # ===========================================================================
+  # LDS store address computation (bank-conflict-avoiding swizzle)
+  # ===========================================================================
+  # This section computes LDS store addresses with a swizzle pattern to avoid bank conflicts.
+  # The swizzle ensures that threads in the same wavefront write to different LDS banks.
+  # Formula: swizzled_addr = base + (lane_id & 7) * LDS_A_STRIDE + swizzle_offset
+  # where swizzle_offset depends on (lane_id >> 3) to distribute across banks.
+  k.emit(v_add_nc_u32_e32(v[9], s[S_LOOP_CTR], v[22]))  # row 0 base
+  k.emit(v_and_b32_e32(v[9], ADDR_MASK, v[9]))
+  k.emit(v_sub_nc_u32_e32(v[9], v[22], v[9]))  # row 0 swizzle offset
+  k.emit(v_lshlrev_b32_e32(v[9], 2, v[9]))  # * 4
+  k.emit(v_mad_u32_u24(v[V_LDS_B_ADDR], LDS_A_STRIDE, v[V_LANE_ID_MOD8], v[9]))
+
+  # For V_LDS_A_BASE and epilogue
+  k.emit(v_bfe_u32(v[2], v[V_LANE_ID], 3, 2))  # v[2] = (lane_id >> 3) & 3
+  k.emit(v_lshlrev_b32_e32(v[V_LANE_DIV8_X4], 2, v[2]))
+
+  # Compute LDS load/store base addresses for inner loop
+  k.emit(v_lshlrev_b32_e32(v[2], 4, v[2]))
+  k.emit(v_and_b32_e32(v[3], 0x7F, v[1]))  # simplified from 3 lines
+  k.emit(v_lshl_or_b32(v[V_LDS_B_BASE], v[V_LANE_ID_MOD8], 4, LDS_BASE_OFFSET))
+  k.emit(v_lshl_add_u32(v[V_LDS_A_ADDR], v[3], 2, LDS_BASE_OFFSET))
+  k.emit(v_lshlrev_b32_e32(v[3], 2, v[V_LANE_ID]))
+  k.emit(v_and_or_b32(v[V_LDS_A_BASE], 0x180, v[3], v[2]))
+
+  # Do initial stores
+  k.waitcnt(vm=0)
+  for i in range(4):  # A tile: 8 values via 4 stride64 stores
+    k.emit(ds_store_2addr_stride64_b32(addr=v[V_LDS_A_ADDR], data0=v[V_LDS_A_DATA[i*2]], data1=v[V_LDS_A_DATA[i*2+1]], offset0=i*4, offset1=i*4+2))
+  for i in range(8):  # B tile: 8 values via 8 scalar stores with 64-byte spacing
+    offset = i * 64
+    k.emit(ds_store_b32(addr=v[V_LDS_B_ADDR], data0=v[V_LDS_B_DATA[i]], offset0=offset & 0xFF, offset1=offset >> 8))
+
+  # Zero all 128 accumulators using VOPD dual moves (64 instructions instead of 128)
+  for i in range(0, len(OUT_REGS), 2):
+    k.emit(VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[OUT_REGS[i]], vdsty=v[OUT_REGS[i+1]], srcx0=0, srcy0=0))
+  k.emit(s_add_i32(s[S_LOOP_BOUND], s[S_DIM_N], -8))
+
+  # S_LOOP_CTR is already 0 from prologue initialization
+  k.emit(s_branch(), target='LOOP_ENTRY')
+
+  # ===========================================================================
+  # MAIN GEMM LOOP
+  # ===========================================================================
+
+  NO_DS, NO_GLOBAL = getenv("NO_DS", 0), getenv("NO_GLOBAL", 0)
+
+  k.label('LOOP_INC')
+  k.emit(s_add_i32(s[S_LOOP_CTR], s[S_LOOP_CTR], 8))
+  k.emit(s_cmp_ge_i32(s[S_LOOP_CTR], s[S_DIM_N]))
+  k.emit(s_cbranch_scc1(), target='EPILOGUE')
+
+  k.label('LOOP_ENTRY')
+  k.emit(s_cmp_lt_i32(s[S_LOOP_CTR], s[S_LOOP_BOUND]))
+  k.emit(s_cselect_b32(s[S_PREFETCH_FLAG], -1, 0))  # s_cselect doesn't modify SCC
+  k.emit(s_cbranch_scc0(), target='SKIP_PREFETCH')  # branch if loop_ctr >= loop_bound
+
+  if not NO_GLOBAL:
+    # Advance prefetch pointers (VGPR)
+    #k.emit(v_add_nc_u32_e32(v[V_GLOBAL_B_ADDR], 0x20000, v[V_GLOBAL_B_ADDR]))
+    #k.emit(v_add_nc_u32_e32(v[V_GLOBAL_A_ADDR], 0x20, v[V_GLOBAL_A_ADDR]))
+
+    # Advance prefetch pointers (64-bit adds)
+    k.emit(s_clause(simm16=31))
+    for i in range(8):
+      k.emit(s_add_u32(s[S_PREFETCH_B+i*2], s[S_PREFETCH_B+i*2], 0x20000))
+      k.emit(s_addc_u32(s[S_PREFETCH_B+i*2+1], s[S_PREFETCH_B+i*2+1], 0))
+    for i in range(8):
+      k.emit(s_add_u32(s[S_PREFETCH_A+i*2], s[S_PREFETCH_A+i*2], 0x20))
+      k.emit(s_addc_u32(s[S_PREFETCH_A+i*2+1], s[S_PREFETCH_A+i*2+1], 0))
+
+    # do the fetch
+    for vdst, saddr_lo in INIT_PREFETCH:
+      k.emit(global_load_b32(vdst=v[vdst], addr=v[V_GLOBAL_B_ADDR], saddr=s[saddr_lo:saddr_lo+1]))
+
+  k.label('SKIP_PREFETCH')
+
+  # wait for local stores to finish (either initial or loop)
+  # then sync the warp so it's safe to load local
+  k.waitcnt(lgkm=0)
+  k.emit(s_barrier())
+
+  # 8 inner loop iterations
+  for iter in range(8):
+    # Load A tile (4 pairs) and B tile (8 pairs) from LDS
+    if not NO_DS:
+      k.emit(s_clause(simm16=len(V_A_TILE_REGS) + len(V_B_TILE_REGS) - 1))  # 12 loads total: 4 A + 8 B
+      # A tile: 4 ds_load_b64
+      for i, vdst in enumerate(V_A_TILE_REGS):
+        a_off = (i & 1) * 8 + (i >> 1) * 64 + iter * LDS_A_STRIDE
+        k.emit(ds_load_b64(vdst=v[vdst:vdst+1], addr=v[V_LDS_A_BASE], offset0=a_off & 0xFF, offset1=a_off >> 8))
+      # B tile: 8 ds_load_b64
+      for i, vdst in enumerate(V_B_TILE_REGS):
+        b_off = (i & 1) * 8 + (i & 2) * 64 + (i >> 2) * 256 + iter * LDS_B_STRIDE
+        k.emit(ds_load_b64(vdst=v[vdst:vdst+1], addr=v[V_LDS_B_BASE], offset0=b_off & 0xFF, offset1=b_off >> 8))
+
+    # Issue global prefetch (first 6 iterations only)
+    if iter < 6 and not NO_GLOBAL:
+      vdst1, vdst2, addr, slo1, slo2 = PREFETCH_LOADS[iter]
+      k.emit(global_load_b32(vdst=v[vdst1], addr=v[addr], saddr=s[slo1:slo1+1]))
+      k.emit(global_load_b32(vdst=v[vdst2], addr=v[addr], saddr=s[slo2:slo2+1]))
+
+    # 64 dual FMACs
+    k.waitcnt(lgkm=0)
+    k.emit(s_clause(simm16=len(FMAC_PATTERN)-1))
+    for i, (vdst_x, vdst_y, ax, bx, ay, by) in enumerate(FMAC_PATTERN):
+      k.emit(VOPD(VOPDOp.V_DUAL_FMAC_F32, VOPDOp.V_DUAL_FMAC_F32,
+                  vdstx=v[vdst_x], vdsty=v[vdst_y], srcx0=v[ax], vsrcx1=v[bx], srcy0=v[ay], vsrcy1=v[by]))
+
+  # wait for all global loads to finish
+  # then sync the warp so it's safe to store local
+  k.waitcnt(vm=0)
+  k.emit(s_barrier())
+
+  # Store prefetched data to LDS
+  # NOTE: Register naming reflects LDS tile organization, not source matrix:
+  #   V_LDS_A_DATA (v155-162) holds data that goes to LDS A-tile region
+  #   V_LDS_B_DATA (v163-170) holds data that goes to LDS B-tile region
+  # The data sources are swapped: A-tile receives B matrix rows, B-tile receives A matrix columns
+  if not NO_DS:
+    for i in range(4):  # A tile: 8 values via 4 stride64 stores
+      k.emit(ds_store_2addr_stride64_b32(addr=v[V_LDS_A_ADDR], data0=v[V_LDS_A_DATA[i*2]], data1=v[V_LDS_A_DATA[i*2+1]], offset0=i*4, offset1=i*4+2))
+    for i in range(8):  # B tile: 8 values via 8 scalar stores with 64-byte spacing
+      offset = i * 64
+      k.emit(ds_store_b32(addr=v[V_LDS_B_ADDR], data0=v[V_LDS_B_DATA[i]], offset0=offset & 0xFF, offset1=offset >> 8))
+
+  k.emit(s_branch(), target='LOOP_INC')
+
+  # ===========================================================================
+  # EPILOGUE: Permute and store results
+  # ===========================================================================
+  k.label('EPILOGUE')
+
+  # Rearrange accumulators from FMAC layout to contiguous output order
+  for a, b in PERMUTE_SWAPS:
+    k.emit(v_swap_b32_e32(v[a], v[b]))
+
+  # Compute output base coordinates
+  # v[130] = col_base = tile_x + (lane_id & 7) * 4
+  # v[131] = row_base = tile_y + (lane_id & 0x60) + ((lane_id >> 3) & 3) * 4
+  # v[132] = 0 (for 64-bit address high part)
+  k.emit(v_add_nc_u32_e32(v[130], s[S_TILE_X], v[V_LANE_MOD8_X4]))
+  k.emit(v_and_b32_e32(v[131], 0x60, v[V_LANE_ID]))
+  k.emit(v_add_nc_u32_e32(v[131], s[S_TILE_Y], v[131]))
+  k.emit(v_add_nc_u32_e32(v[131], v[V_LANE_DIV8_X4], v[131]))
+  k.emit(v_mov_b32_e32(v[132], 0))
+
+  # Precompute row offsets: v[133-136] for rows 0-3, v[137-140] for rows 16-19
+  for base, row_off in [(133, 0), (137, 16)]:
+    if row_off: k.emit(v_add_nc_u32_e32(v[141], row_off, v[131]))
+    k.emit(v_mul_lo_u32(v[base], v[141] if row_off else v[131], s[S_DIM_N]))
+    for j in range(3): k.emit(v_add_nc_u32_e32(v[base + 1 + j], s[S_DIM_N], v[base + j]))
+
+  # s[S_PREFETCH_FLAG] = row stride in bytes (N * 4)
+  k.emit(s_lshl_b32(s[S_PREFETCH_FLAG], s[S_DIM_N], 2))
+
+  # Store 128 output values as 32 groups of 4 (128-bit stores)
+  # Layout: 2 row halves (0-3, 16-19) x 4 col groups x 4 rows = 32 stores of 4 floats
+  for i, (row_half, col_off, row_in_group) in enumerate([(rh, co, ri)
+      for rh in range(2) for co in [0, 32, 64, 96] for ri in range(4)]):
+    row = row_half * 16 + row_in_group
+    src = OUT_REGS[i*4]  # first reg of ascending group of 4
+
+    if row_in_group == 0:
+      # First row of group: compute full address
+      if col_off == 0: k.emit(v_mov_b32_e32(v[141], v[130]))
+      else: k.emit(v_add_nc_u32_e32(v[141], col_off, v[130]))
+      row_base = 133 + row if row < 4 else 137 + row - 16
+      k.emit(v_add_nc_u32_e32(v[141], v[row_base], v[141]))
+      k.emit(v_lshlrev_b32_e32(v[141], 2, v[141]))
+      k.emit(v_add_co_u32(v[141], VCC_LO, s[S_OUT_PTR[0]], v[141]))
+      k.emit(v_add_co_ci_u32_e32(v[142], s[S_OUT_PTR[1]], v[132]))
+    else:
+      # Subsequent rows: add stride
+      k.emit(v_add_co_u32(v[141], VCC_LO, s[S_PREFETCH_FLAG], v[141]))
+      k.emit(v_add_co_ci_u32_e32(v[142], v[142], v[132]))
+
+    k.emit(global_store_b128(addr=v[141:142], data=v[src:src+3], saddr=NULL))
+
+  k.emit(s_sendmsg(simm16=3))  # DEALLOC_VGPRS
+  k.emit(s_endpgm())
+
+  return k.finalize()
+
+# =============================================================================
+# Test harness
+# =============================================================================
+
+N = getenv("N", 4096)
+BLOCK_M, BLOCK_N = 128, 128
+THREADS = 128
+
+def test_matmul():
+  dev = Device[Device.DEFAULT]
+  print(f"Device arch: {dev.renderer.arch}")
+
+  insts = build_kernel(dev.renderer.arch)
+
+  rng = np.random.default_rng(42)
+  a = Tensor(rng.random((N, N), dtype=np.float32) - 0.5)
+  b = Tensor(rng.random((N, N), dtype=np.float32) - 0.5)
+  c = Tensor.empty(N, N)
+  Tensor.realize(a, b, c)
+
+  grid, local = (N // BLOCK_N, N // BLOCK_M, 1), (THREADS, 1, 1)
+  print(f"Grid: {grid}, Local: {local}")
+
+  dname:str = Device.DEFAULT
+  def asm_kernel(A:UOp, B:UOp, C:UOp) -> UOp:
+    gidxs = [UOp.special(n, f"gidx{i}") for i,n in enumerate(grid)]
+    lidxs = [UOp.special(n, f"lidx{i}") for i,n in enumerate(local)]
+    lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=max(LDS_SIZE, 65536//getenv("LIMIT_OCC", 65536)), addrspace=AddrSpace.LOCAL), (), 'lds')
+    sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs, arg=KernelInfo(name=colored("kernel", "cyan"),
+                                                                                  estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+  c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
+  ei = c.schedule()[0].lower()
+
+  ets = []
+  with Context(DEBUG=2):
+    for _ in range(getenv("CNT", 5)): ets.append(ei.run(wait=True))
+  print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")
+
+  if getenv("VERIFY", 1):
+    GlobalCounters.reset()
+    with Context(DEBUG=2): tc = (a @ b).realize()
+    with Context(DEBUG=0): err = (c - tc).square().mean().item()
+    print(f"mean squared error {err}")
+    if err != err or err > 1e-06: raise RuntimeError("matmul is wrong!")
+
+def run_sqtt():
+  """Run with SQTT profiling and write trace files."""
+  import subprocess, os
+
+  # Run test_matmul in a subprocess with SQTT enabled from the start (no verify)
+  env = {**os.environ, "AMD": "1", "SQTT": "1", "CNT": "1", "PROFILE": "1", "PYTHONPATH": ".", "VERIFY": "0"}
+  result = subprocess.run(
+    ["python", "-c", "from extra.gemm.amd_asm_matmul import test_matmul; test_matmul()"],
+    capture_output=True, text=True, env=env, timeout=120
+  )
+  print(result.stdout)
+
+  # Run roc.py to extract trace data
+  result = subprocess.run(
+    ["python", "extra/sqtt/roc.py", "--profile", "/tmp/profile.pkl.tiny", "--kernel", "kernel"],
+    capture_output=True, text=True, env={**os.environ, "DEBUG": "5"}, timeout=60
+  )
+  output = result.stdout + result.stderr
+
+  # Write full output to trace file
+  with open("/tmp/sqtt_trace.txt", "w") as f:
+    f.write(output)
+  print(f"Wrote {len(output)} bytes to /tmp/sqtt_trace.txt")
+
+if __name__ == "__main__":
+  if getenv("ASM", 0): print("\n".join(str(inst) for inst in build_kernel(Device[Device.DEFAULT].renderer.arch)))
+  elif getenv("SQTT", 0): run_sqtt()
+  else: test_matmul()
@@ -140,11 +140,11 @@ def hand_spec_kernel3():

  return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()

-def test_matmul(sink:UOp, N=N):
+def test_matmul(sink:UOp, dtype=dtypes.float32, N=N):
  rng = np.random.default_rng()
-  a = Tensor(rng.random((N, N), dtype=np.float32)-0.5)
-  b = Tensor(rng.random((N, N), dtype=np.float32)-0.5)
-  hc = Tensor.empty(N, N)
+  a = Tensor(rng.random((N, N), dtype=np.float32)-0.5, dtype=dtype)
+  b = Tensor(rng.random((N, N), dtype=np.float32)-0.5, dtype=dtype)
+  hc = Tensor.empty(N, N, dtype=dtype)
  Tensor.realize(a, b, hc)

  ei = ExecItem(sink, [t.uop.buffer for t in [hc, a, b]], prg=get_runner(Device.DEFAULT, sink))
@@ -0,0 +1,102 @@
+import atexit, functools
+from tinygrad.runtime.support.compiler_amd import HIPCompiler
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.renderer import Estimates
+from tinygrad.helpers import getenv, all_same, dedup
+from extra.gemm.asm.cdna.asm import build_kernel, GEMM_ARGS
+
+# ** CDNA4 assembly gemm
+
+WORKGROUP_SIZE = 256
+
+@functools.cache
+def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str, arch:str, wg:int) -> UOp:
+  batch, M, K = A.shape
+  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
+  assert K == K2
+  lidx = UOp.special(WORKGROUP_SIZE, "lidx0")
+  gidx = UOp.special(wg, "gidx0")
+  k = build_kernel(batch, M, N, K, A.dtype.base)
+  sink = UOp.sink(C.base, A.base, B.base, lidx, gidx,
+                  arg=KernelInfo(name=k.name, estimates=Estimates(ops=2*batch*M*N*K, mem=(batch*M*K + K*N + batch*M*N)*2)))
+  # TODO: you shouldn't have to call the compiler here, BINARY should be auto-added
+  binary = HIPCompiler(arch).compile(k.to_asm())
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=k.to_text()), UOp(Ops.BINARY, arg=binary)))
+
+counters = {"used":0, "todos":[]}
+def todo(msg:str) -> bool: counters["todos"].append(msg); return False
+atexit.register(lambda: print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used'))
+
+def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
+  if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}")
+  if a.dtype not in {dtypes.bfloat16, dtypes.float16}: return todo(f"only bfloat16/float16, got {a.dtype}")
+  batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
+  N = b.shape[1]
+  # only sharding on the batch or K is tested, others might work too
+  if isinstance(a.device, tuple):
+    if a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
+    elif a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None: batch //= len(a.device)
+    else: return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
+    dname = a.device[0]
+  else: dname = a.device
+  arch = getattr(Device[dname].renderer, "arch", "")
+  if batch not in {1, 2}: return todo(f"GEMM batch size {batch}")
+  if (key:=(M, N, K)) not in GEMM_ARGS and arch == "gfx950": return todo(f"GEMM shape not supported {key} on {arch}")
+  return True
+
+# ** UOp gemm to test Tensor.custom_kernel multi and backward correctness on non cdna4
+# note: this can be removed after we have GEMM on mixins
+
+def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
+  M, K = A.shape[0]*A.shape[1], A.shape[2]
+  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
+  assert K == K2
+  m = UOp.range(M, 1, AxisType.LOOP)
+  n = UOp.range(N, 2, AxisType.LOOP)
+  k = UOp.range(K, 0, AxisType.REDUCE)
+  mul = (A.index((m*UOp.const(dtypes.index, K)+k))*B.index((k*UOp.const(dtypes.index, N)+n))).cast(dtypes.float32)
+  red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype.base)
+  store = C.index((m*UOp.const(dtypes.index, N)+n), ptr=True).store(red).end(m, n)
+  return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
+
+# ** backward gemm, might use the asm gemm
+
+def custom_gemm_bw(gradient:UOp, kernel:UOp):
+  out, a, b = kernel.src[1:]
+  assert all_same([gradient.device, a.device, b.device, out.device])
+  a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
+  # TODO: this needs to be cleaned up and done properly, the batch dim of grad and a multi need to align
+  g_t = g_t[:a.shape[0]]
+  grad_a = (g_t @ b_t.T).uop
+  grad_b = (a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1) @ g_t.reshape(-1, g_t.shape[-1])).uop
+  return (None, grad_a, grad_b)
+
+# ** main gemm function
+
+def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
+  assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
+  counters["used"] += 1
+  squeeze = a.ndim == 2
+  if squeeze: a = a.unsqueeze(0)
+
+  batch, M, K = a.shape
+  N = b.shape[1]
+  is_multi = isinstance(a.device, tuple)
+  if (k_sharded:=is_multi and a.uop.axis == 2): K //= len(a.device)
+
+  if is_multi:
+    out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
+  else:
+    out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)
+
+  dname = a.device[0] if is_multi else a.device
+  arch = getattr(Device[dname].renderer, "arch", "")
+  if arch.startswith("gfx950") and getenv("USE_ASM", 1):
+    numWG = GEMM_ARGS[(M, N, K)][0]
+    out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname, wg=numWG, arch=arch), grad_fxn=custom_gemm_bw)[0]
+  else:
+    out = Tensor.custom_kernel(out, a, b, fxn=custom_uop_gemm, grad_fxn=custom_gemm_bw)[0]
+  if k_sharded: out = out.sum(0)
+  return out.squeeze(0) if squeeze else out
@@ -0,0 +1,76 @@
+.text
+.section	.text.
+.global	gemm
+.p2align	8
+.type	gemm,@function
+
+gemm:
+INSTRUCTIONS
+
+.section .rodata,"a",@progbits
+.p2align 6, 0x0
+.amdhsa_kernel gemm
+  # basic memory requirements
+  .amdhsa_group_segment_fixed_size 30336
+  .amdhsa_private_segment_fixed_size 0
+  .amdhsa_kernarg_size 32
+  # register usage (RSRC1)
+  .amdhsa_next_free_vgpr 256
+  .amdhsa_next_free_sgpr 100
+  # workgroup / workitem IDs (RSRC2)
+  .amdhsa_system_sgpr_workgroup_id_x 1
+  .amdhsa_system_sgpr_workgroup_id_y 1
+  .amdhsa_system_sgpr_workgroup_id_z 1
+  # user SGPRs: kernarg ptr in s[0:1]
+  .amdhsa_user_sgpr_kernarg_segment_ptr 1
+  .amdhsa_user_sgpr_count 2
+  # gfx10+ / gfx11 specifics (RSRC1[29..31])
+  .amdhsa_wavefront_size32 1
+  .amdhsa_workgroup_processor_mode 1
+  .amdhsa_memory_ordered 1
+  .amdhsa_forward_progress 1
+  # misc for gfx11
+  .amdhsa_dx10_clamp 1
+  .amdhsa_ieee_mode 1
+  .amdhsa_uses_dynamic_stack 0
+.end_amdhsa_kernel
+
+.amdgpu_metadata
+---
+amdhsa.kernels:
+  - .args:
+      - .address_space:  generic
+        .name:           C
+        .offset:         0
+        .size:           8
+        .value_kind:     global_buffer
+        .value_type:     f16
+      - .address_space:  generic
+        .name:           A
+        .offset:         8
+        .size:           8
+        .value_kind:     global_buffer
+        .value_type:     f16
+      - .address_space:  generic
+        .name:           B
+        .offset:         16
+        .size:           8
+        .value_kind:     global_buffer
+        .value_type:     f16
+    .group_segment_fixed_size: 30336
+    .kernarg_segment_align: 8
+    .kernarg_segment_size: 32
+    .max_flat_workgroup_size: 128
+    .name:           gemm
+    .private_segment_fixed_size: 0
+    .sgpr_count:     70
+    .sgpr_spill_count: 0
+    .symbol:         gemm.kd
+    .vgpr_count:     256
+    .vgpr_spill_count: 0
+    .wavefront_size: 32
+amdhsa.version:
+  - 1
+  - 1
+...
+.end_amdgpu_metadata
@@ -0,0 +1,30 @@
+import math, pathlib
+
+from tinygrad import Device, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+
+from extra.gemm.amd_uop_matmul import test_matmul
+
+N = 4096
+TN = 96
+THREADS_PER_WG = 128
+NUM_WG = math.ceil(N / TN) * math.ceil(N / TN)
+
+dname:str = Device.DEFAULT
+template:str = (pathlib.Path(__file__).parent/"template.s").read_text()
+
+def asm_kernel() -> UOp:
+  lidx = UOp.special(THREADS_PER_WG, "lidx0")
+  gidx = UOp.special(NUM_WG, "gidx0")
+
+  a = UOp.placeholder((N*N,), dtypes.half, slot=1)
+  b = UOp.placeholder((N*N,), dtypes.half, slot=2)
+  c = UOp.placeholder((N*N,), dtypes.half, slot=0)
+
+  src = template.replace("INSTRUCTIONS", (pathlib.Path(__file__).parent/"gemm.s").read_text())
+
+  sink = UOp.sink(a, b, c, lidx, gidx, arg=KernelInfo(name="gemm"))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src)))
+
+if __name__ == "__main__":
+  test_matmul(asm_kernel(), dtype=dtypes.half, N=N)
@@ -0,0 +1,179 @@
+# unpack the complete kernel descriptor of an amdgpu ELF
+# https://rocm.docs.amd.com/projects/llvm-project/en/latest/LLVM/llvm/html/AMDGPUUsage.html#code-object-v3-kernel-descriptor
+import struct, pathlib, sys
+from tinygrad.runtime.support.elf import elf_loader
+
+def bits(x, lo, hi): return (x >> lo) & ((1 << (hi - lo + 1)) - 1)
+def assert_zero(x, lo, hi): assert bits(x, lo, hi) == 0
+
+with open(sys.argv[1], "rb") as f:
+  lib = f.read()
+
+image, sections, relocs = elf_loader(lib)
+rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"))
+
+# rodata is exactly 64 bytes
+kd = image[rodata_entry:rodata_entry+64]
+desc = int.from_bytes(kd, byteorder="little")
+
+group_segment_fixed_size = bits(desc, 0, 31)
+private_segment_fixed_size = bits(desc, 32, 63)
+kernarg_size = bits(desc, 64, 95)
+reserved_127_96 = bits(desc, 96, 127)
+assert reserved_127_96 == 0
+
+print("GROUP_SEGMENT_FIXED_SIZE:", group_segment_fixed_size)
+print("PRIVATE_SEGMENT_FIXED_SIZE:", private_segment_fixed_size)
+print("KERNARG_SIZE:", kernarg_size)
+print("RESERVED 127:96:", reserved_127_96)
+
+entry_off = bits(desc, 128, 191)
+
+# sign-extend manually if needed
+if entry_off & (1 << 63):
+  entry_off -= 1 << 64
+
+print("KERNEL_CODE_ENTRY_BYTE_OFFSET:", entry_off)
+
+kd_addr = 0x1840
+entry_addr = kd_addr + entry_off
+
+print("Computed entry address: 0x%016x" % entry_addr)
+print("256B aligned:", entry_addr % 256 == 0)
+
+pgm_rsrc3 = bits(desc, 352, 383)
+pgm_rsrc1 = bits(desc, 384, 415)
+pgm_rsrc2 = bits(desc, 416, 447)
+
+print("COMPUTE_PGM_RSRC3: 0x%08x" % pgm_rsrc3)
+print("COMPUTE_PGM_RSRC1: 0x%08x" % pgm_rsrc1)
+print("COMPUTE_PGM_RSRC2: 0x%08x" % pgm_rsrc2)
+
+# rsrc 3 (gfx950)
+
+accum_offset_raw = bits(pgm_rsrc3, 0, 5)
+assert_zero(pgm_rsrc3, 6, 15)
+tg_split = bits(pgm_rsrc3, 16, 16)
+accum_offset_vgprs = (accum_offset_raw + 1) * 4
+print("RSRC3.ACCUM_OFFSET (AccVGPR index):", accum_offset_vgprs)
+print("RSRC3.TG_SPLIT:", tg_split)
+
+# rsrc 1
+
+vgpr_gran = bits(pgm_rsrc1, 0, 5)
+sgpr_gran = bits(pgm_rsrc1, 6, 9)
+assert_zero(pgm_rsrc1, 27, 28)
+
+# NOTE: this is vgprs + agprs
+vgprs_used = (vgpr_gran + 1) * 8
+assert 0 <= vgprs_used <= 512
+
+k = sgpr_gran // 2
+sgprs_used = (k + 1) * 16
+
+print("RSRC1.VGPRS:", vgprs_used)
+print("RSRC1.SGPRS:", sgprs_used)
+
+assert_zero(pgm_rsrc1, 10, 11)
+
+float_round_mode_32 = bits(pgm_rsrc1, 12, 13)
+float_round_mode_16_64 = bits(pgm_rsrc1, 15, 14)
+float_denorm_mode_32 = bits(pgm_rsrc1, 16, 17)
+float_denorm_mode_16_64 = bits(pgm_rsrc1, 18, 19)
+
+priv = bits(pgm_rsrc1, 20, 20)
+assert priv == 0
+enable_dx10_clamp_wg_rr_en = bits(pgm_rsrc1, 21, 21)
+debug_mode = bits(pgm_rsrc1, 22, 22)
+enable_ieee_mode = bits(pgm_rsrc1, 23, 23)
+bulky = bits(pgm_rsrc1, 24, 24)
+assert bulky == 0
+cdbg_user = bits(pgm_rsrc1, 25, 25)
+assert cdbg_user == 0
+fp16_ovfl = bits(pgm_rsrc1, 26, 26)
+assert_zero(pgm_rsrc1, 27, 28)  # reserved
+assert_zero(pgm_rsrc1, 29, 29)  # WGP_MODE (reserved on gfx9)
+assert_zero(pgm_rsrc1, 30, 30)  # MEM_ORDERED (reserved on gfx9)
+assert_zero(pgm_rsrc1, 31, 31)  # FWD_PROGRESS (reserved on gfx9)
+
+# rsrc 2
+
+enable_private_segment = bits(pgm_rsrc2, 0, 0)  # SCRATCH_EN
+user_sgpr_count = bits(pgm_rsrc2, 1, 5)         # USER_SGPR
+enable_trap_handler = bits(pgm_rsrc2, 6, 6)     # TRAP_PRESENT (must be 0 here)
+assert enable_trap_handler == 0
+
+enable_sgpr_workgroup_id_x = bits(pgm_rsrc2, 7, 7)
+enable_sgpr_workgroup_id_y = bits(pgm_rsrc2, 8, 8)
+enable_sgpr_workgroup_id_z = bits(pgm_rsrc2, 9, 9)
+enable_sgpr_workgroup_info = bits(pgm_rsrc2, 10, 10)
+
+enable_vgpr_workitem_id = bits(pgm_rsrc2, 11, 12)  # TIDIG_CMP_CNT enum (0..3)
+
+enable_exception_address_watch = bits(pgm_rsrc2, 13, 13)
+assert enable_exception_address_watch == 0
+enable_exception_memory = bits(pgm_rsrc2, 14, 14)
+assert enable_exception_memory == 0
+
+granulated_lds_size = bits(pgm_rsrc2, 15, 23)
+assert granulated_lds_size == 0  # spec: must be 0; CP uses dispatch packet rounding
+
+enable_exception_fp_invalid = bits(pgm_rsrc2, 24, 24)
+enable_exception_fp_denorm_src = bits(pgm_rsrc2, 25, 25)
+enable_exception_fp_div0 = bits(pgm_rsrc2, 26, 26)
+enable_exception_fp_overflow = bits(pgm_rsrc2, 27, 27)
+enable_exception_fp_underflow = bits(pgm_rsrc2, 28, 28)
+enable_exception_fp_inexact = bits(pgm_rsrc2, 29, 29)
+enable_exception_int_div0 = bits(pgm_rsrc2, 30, 30)
+
+assert_zero(pgm_rsrc2, 31, 31)
+
+print("RSRC2.ENABLE_PRIVATE_SEGMENT:", enable_private_segment)
+print("RSRC2.USER_SGPR_COUNT:", user_sgpr_count)
+print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_X:", enable_sgpr_workgroup_id_x)
+print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_Y:", enable_sgpr_workgroup_id_y)
+print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_Z:", enable_sgpr_workgroup_id_z)
+print("RSRC2.ENABLE_SGPR_WORKGROUP_INFO:", enable_sgpr_workgroup_info)
+print("RSRC2.ENABLE_VGPR_WORKITEM_ID (enum):", enable_vgpr_workitem_id)
+
+print("RSRC2.EXC_FP_INVALID:", enable_exception_fp_invalid)
+print("RSRC2.EXC_FP_DENORM_SRC:", enable_exception_fp_denorm_src)
+print("RSRC2.EXC_FP_DIV0:", enable_exception_fp_div0)
+print("RSRC2.EXC_FP_OVERFLOW:", enable_exception_fp_overflow)
+print("RSRC2.EXC_FP_UNDERFLOW:", enable_exception_fp_underflow)
+print("RSRC2.EXC_FP_INEXACT:", enable_exception_fp_inexact)
+print("RSRC2.EXC_INT_DIV0:", enable_exception_int_div0)
+
+# user sgprs
+
+enable_sgpr_private_segment_buffer = bits(desc, 448, 448)
+enable_sgpr_dispatch_ptr = bits(desc, 449, 449)
+enable_sgpr_queue_ptr = bits(desc, 450, 450)
+enable_sgpr_kernarg_segment_ptr = bits(desc, 451, 451)
+enable_sgpr_dispatch_id = bits(desc, 452, 452)
+enable_sgpr_flat_scratch_init = bits(desc, 453, 453)
+enable_sgpr_private_segment_size = bits(desc, 454, 454)
+
+assert_zero(desc, 455, 457)
+
+print("DESC.ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER:", enable_sgpr_private_segment_buffer)
+print("DESC.ENABLE_SGPR_DISPATCH_PTR:", enable_sgpr_dispatch_ptr)
+print("DESC.ENABLE_SGPR_QUEUE_PTR:", enable_sgpr_queue_ptr)
+print("DESC.ENABLE_SGPR_KERNARG_SEGMENT_PTR:", enable_sgpr_kernarg_segment_ptr)
+print("DESC.ENABLE_SGPR_DISPATCH_ID:", enable_sgpr_dispatch_id)
+print("DESC.ENABLE_SGPR_FLAT_SCRATCH_INIT:", enable_sgpr_flat_scratch_init)
+print("DESC.ENABLE_SGPR_PRIVATE_SEGMENT_SIZE:", enable_sgpr_private_segment_size)
+
+assert_zero(desc, 458, 459)
+
+uses_dynamic_stack = bits(desc, 459, 460)
+print("DESC.USES_DYNAMIC_STACK:", uses_dynamic_stack)
+
+# gfx950 only
+assert_zero(desc, 460, 463)
+kernarg_preload_spec_length = bits(desc, 464, 470)
+print("DESC.KERNARG_PRELOAD_SPEC_LENGTH:", kernarg_preload_spec_length)
+kernarg_preload_spec_offset = bits(desc, 471, 479)
+print("DESC.KERNARG_PRELOAD_SPEC_OFFSET:", kernarg_preload_spec_offset)
+
+assert_zero(desc, 480, 511)
@@ -37,7 +37,7 @@ b.copyin(row.data)
 c.copyin(mat.data)
 ret = prog(a._buf, b._buf, c._buf, global_size=[1,1,1], local_size=[8,1,1], wait=True)
 print(ret)
-out = np.frombuffer(a.as_buffer(), np.float32)
+out = np.frombuffer(a.as_memoryview(), np.float32)
 real = row.astype(np.float32)@mat.T.astype(np.float32)
 print("out:", out)
 print("real", real)
@@ -98,10 +98,10 @@ if __name__ == "__main__":
  # check correctness
  if getenv("VERIFY"):
    from tinygrad.engine.realize import run_schedule
-    triton_buf = np.frombuffer(si.bufs[0].as_buffer(), np.float16).reshape(M,N)
+    triton_buf = np.frombuffer(si.bufs[0].as_memoryview(), np.float16).reshape(M,N)
    print(triton_buf)
    run_schedule(sched)
-    tinygrad_buf = np.frombuffer(si.bufs[0].as_buffer(), np.float16).reshape(M,N)
+    tinygrad_buf = np.frombuffer(si.bufs[0].as_memoryview(), np.float16).reshape(M,N)
    print(tinygrad_buf)
    np.testing.assert_allclose(triton_buf, tinygrad_buf)
    print("correct!")
@@ -1,14 +1,15 @@
 #!/usr/bin/env python3

 import argparse, glob, os, time, subprocess, sys
+from tinygrad.helpers import temp

 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)
+  for dev in glob.glob(temp(f'{prefix}_*.lock')):
+    dev_id = dev.split('/')[-1][len(prefix)+1:-5]
+    if dev_id.startswith(target_dev): devs.append(dev_id)
  return devs

 def _do_reset_device(pci_bus): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{pci_bus}/reset'")
@@ -53,16 +54,7 @@ def cmd_show_pids(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]
+      pid = subprocess.check_output(['sudo', 'lsof', temp(f'{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")

@@ -74,7 +66,7 @@ def cmd_kill_pids(args):
      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]
+        try: pid = subprocess.check_output(['sudo', 'lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
        except subprocess.CalledProcessError: break

        print(f"Killing process {pid} (which uses {dev})")
@@ -1,8 +1,42 @@
-import argparse, os, hashlib
-from tinygrad.helpers import getenv, DEBUG, round_up, Timing, tqdm, fetch
+import argparse, os, hashlib, functools
+from typing import Iterator, Callable
+from tinygrad.helpers import getenv, DEBUG, round_up, Timing, tqdm, fetch, ceildiv
 from extra.hevc.hevc import parse_hevc_file_headers, untile_nv12, to_bgr, nv_gpu
 from tinygrad import Tensor, dtypes, Device, Variable, TinyJit

+# rounds up hevc input data to 32 bytes, so more optimal kernels can be generated
+HEVC_ROUNDUP = getenv("DATA_ROUNDUP", 32)
+
+@functools.cache
+def _hevc_jitted_decoder(out_image_size:tuple[int, int], max_hist:int, inplace:bool):
+  def hevc_decode_frame(pos:Variable, hevc_tensor:Tensor, offset:Variable, sz:Variable, opaque:Tensor, i:Variable, *hist:Tensor, outbuf:Tensor|None=None):
+    x = hevc_tensor[offset:offset+sz*HEVC_ROUNDUP].decode_hevc_frame(pos, out_image_size, opaque[i], hist)
+    if outbuf is not None: outbuf.assign(x).realize()
+    return x.realize()
+  return TinyJit(hevc_decode_frame)
+
+def hevc_decode(hevc_tensor:Tensor, opaque:Tensor, frame_info:list, luma_h:int, luma_w:int,
+                history:list[Tensor]|None=None, preallocated_outputs:list[Tensor]|None=None, warmup=False) -> Iterator[Tensor]:
+  out_image_size = luma_h + (luma_h + 1) // 2, round_up(luma_w, 64)
+  max_hist = max((hs for _, _, _, hs, _ in frame_info), default=0)
+
+  v_pos = Variable("pos", 0, max_hist + 1)
+  v_offset = Variable("offset", 0, hevc_tensor.numel()-1)
+  v_sz = Variable("sz", 1, ceildiv(hevc_tensor.numel(), HEVC_ROUNDUP))
+  v_i = Variable("i", 0, len(frame_info)-1)
+
+  decode_jit = _hevc_jitted_decoder(out_image_size, max_hist, preallocated_outputs is not None)
+  history = history or [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV").contiguous().realize() for _ in range(max_hist)]
+  assert len(history) == max_hist, f"history length {len(history)} does not match max_hist {max_hist}"
+
+  for i, (offset, sz, frame_pos, _, is_hist) in enumerate(frame_info):
+    history = history[-max_hist:] if max_hist > 0 else []
+    img = decode_jit(v_pos.bind(frame_pos), hevc_tensor, v_offset.bind(offset), v_sz.bind(ceildiv(sz, HEVC_ROUNDUP)),
+                     opaque, v_i.bind(i), *history, outbuf=preallocated_outputs[i] if preallocated_outputs else None)
+    res = preallocated_outputs[i] if preallocated_outputs else img.clone().realize()
+    if is_hist: history.append(res)
+    yield res
+
 if __name__ == "__main__":
  parser = argparse.ArgumentParser()
  parser.add_argument("--input_file", type=str, default="")
@@ -19,49 +53,32 @@ if __name__ == "__main__":
  dat_hash = hashlib.md5(dat).hexdigest()

  with Timing("prep infos: "):
-    dat_nv = hevc_tensor.to("NV")
    opaque, frame_info, w, h, luma_w, luma_h, chroma_off = parse_hevc_file_headers(dat)

  frame_info = frame_info[:getenv("MAX_FRAMES", len(frame_info))]

  # move all needed data to gpu
-  #all_slices = []
  with Timing("copy to gpu: "):
    opaque_nv = opaque.to("NV").contiguous().realize()
    hevc_tensor = hevc_tensor.to("NV")

  out_image_size = luma_h + (luma_h + 1) // 2, round_up(luma_w, 64)
-  max_hist = max(history_sz for _, _, _, history_sz, _ in frame_info)

-  # define variables
-  v_pos = Variable("pos", 0, max_hist + 1)
-  v_offset = Variable("offset", 0, hevc_tensor.numel()-1)
-  v_sz = Variable("sz", 0, hevc_tensor.numel())
-  v_i = Variable("i", 0, len(frame_info)-1)
+  # preallocate output/hist buffers
+  max_hist = max((hs for _, _, _, hs, _ in frame_info), default=0)
+  hist = [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV").contiguous().realize() for _ in range(max_hist)]
+  out_images = [Tensor.zeros(*out_image_size, dtype=dtypes.uint8, device="NV").contiguous().realize() for _ in range(len(frame_info))]

-  @TinyJit
-  def decode_jit(pos:Variable, src:Tensor, data:Tensor, *hist:Tensor):
-    return src.decode_hevc_frame(pos, out_image_size, data, hist).realize()
+  # warmup decode
+  _ = list(hevc_decode(hevc_tensor, opaque_nv, frame_info[:3], luma_h, luma_w, history=hist, preallocated_outputs=out_images))
+  Device.default.synchronize()

-  # warm up
-  history = [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV") for _ in range(max_hist)]
-  for i in range(3):
-    hevc_frame = hevc_tensor.shrink((((bound_offset:=v_offset.bind(frame_info[0][0])), bound_offset+v_sz.bind(frame_info[0][1])),))
-    decode_jit(v_pos.bind(0), hevc_frame, opaque_nv[v_i.bind(0)], *history)
-
-  out_images = []
+  # decode all frames using the iterator
  with Timing("decoding whole file: ", on_exit=(lambda et: f", {len(frame_info)} frames, {len(frame_info)/(et/1e9):.2f} fps")):
-    for i, (offset, sz, frame_pos, history_sz, is_hist) in enumerate(frame_info):
-      history = history[-max_hist:] if max_hist > 0 else []
-      # TODO: this shrink should work as a slice
-      hevc_frame = hevc_tensor.shrink((((bound_offset:=v_offset.bind(offset)), bound_offset+v_sz.bind(sz)),))
-
-      outimg = decode_jit(v_pos.bind(frame_pos), hevc_frame, opaque_nv[v_i.bind(i)], *history).clone()
-      out_images.append(outimg)
-      if is_hist: history.append(outimg)
-
+    images = list(hevc_decode(hevc_tensor, opaque_nv, frame_info, luma_h, luma_w, history=hist, preallocated_outputs=out_images))
    Device.default.synchronize()

+  # validation
  if getenv("VALIDATE", 0):
    import pickle
    if dat_hash == "b813bfdbec194fd17fdf0e3ceb8cea1c":
@@ -70,7 +87,7 @@ if __name__ == "__main__":
    else: decoded_frames = pickle.load(open(f"extra/hevc/decoded_frames_{dat_hash}.pkl", "rb"))
  else: import cv2

-  for i, img in tqdm(enumerate(out_images)):
+  for i, img in tqdm(enumerate(images)):
    if getenv("VALIDATE", 0):
      if i < len(decoded_frames) and len(decoded_frames[i]) > 0:
        img = untile_nv12(img, h, w, luma_w, chroma_off).realize()
@@ -13,7 +13,7 @@ def get_struct(argp, stype):

 def format_struct(s):
  sdats = []
-  for field_name, field_type in s._fields_:
+  for field_name, *_ in s._real_fields_:
    dat = getattr(s, field_name)
    if isinstance(dat, int): sdats.append(f"{field_name}:0x{dat:X}")
    else: sdats.append(f"{field_name}:{dat}")
@@ -46,6 +46,22 @@ def install_hook(c_function, python_function):
 # *** ioctl lib end ***

 import tinygrad.runtime.autogen.kfd as kfd_ioctl
+import tinygrad.runtime.autogen.hsa as hsa
+
+def print_aql_queue(read_pointer_address):
+  rptr_offset = getattr(hsa.amd_queue_v2_t, 'read_dispatch_id').offset
+  queue_base = read_pointer_address - rptr_offset
+  queue = hsa.amd_queue_v2_t.from_address(queue_base)
+  print(f"  AQL Queue @ 0x{queue_base:X}:")
+  for field_name, *_ in hsa.amd_queue_v2_t._real_fields_:
+    val = getattr(queue, field_name)
+    if isinstance(val, int): print(f"    {field_name}: 0x{val:X}")
+    elif hasattr(val, '_length_'):
+      arr_vals = [f"{format_struct(v)}" if hasattr(v, '_real_fields_') else f"{v:#X}" for v in val]
+      print(f"    {field_name}: [{', '.join(arr_vals)}]")
+    elif hasattr(val, '_real_fields_'): print(f"    {field_name}: {format_struct(val)}")
+    else: print(f"    {field_name}: {val}")
+
 def ioctls_from_header():
  hdr = (pathlib.Path(__file__).parent / "kfd_ioctl.h").read_text().replace("\\\n", "")
  pattern = r'#define\s+(AMDKFD_IOC_[A-Z0-9_]+)\s+AMDKFD_IOW?R?\((0x[0-9a-fA-F]+),\s+struct\s([A-Za-z0-9_]+)\)'
@@ -67,6 +83,7 @@ def ioctl(fd, request, argp):
    if name == "AMDKFD_IOC_SVM":
      out = ctypes.cast(s.attrs, ctypes.POINTER(kfd_ioctl.struct_kfd_ioctl_svm_attribute))
      for i in range(s.nattr): print(f"{i}: {kfd_ioctl.enum_kfd_ioctl_svm_attr_type.get(out[i].type):40s}: {out[i].value:#x}")
+    if name == "AMDKFD_IOC_CREATE_QUEUE" and s.queue_type == kfd_ioctl.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL: print_aql_queue(s.read_pointer_address)
  else:
    print(f"{(st-start)*1000:7.2f} ms +{et*1000.:7.2f} ms : ioctl",
          f"{idir=} {size=} {itype=} {nr=} {fd=} {ret=}", os.readlink(f"/proc/self/fd/{fd}") if fd >= 0 else "")
@@ -18,7 +18,7 @@ prg = dev.runtime("write_ones", mbin)
 prg(buf0._buf, global_size=(1,65537,1), local_size=(1,1,1), wait=True)

 import numpy as np
-def to_np(buf): return np.frombuffer(buf.as_buffer().cast(buf.dtype.base.fmt), dtype=_to_np_dtype(buf.dtype.base))
+def to_np(buf): return np.frombuffer(buf.as_memoryview().cast(buf.dtype.base.fmt), dtype=_to_np_dtype(buf.dtype.base))

 big = to_np(buf0)
 print(big)
@@ -8,14 +8,14 @@ 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,
+def snapshot_download_with_retry(*, repo_id: str, allow_patterns: list[str]|tuple[str, ...]|None=None, local_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,
+        local_dir=str(local_dir) if local_dir is not None else None,
        **kwargs
      ))
    except Exception as e:
@@ -144,14 +144,14 @@ class HuggingFaceONNXManager:
        root_path = snapshot_download_with_retry(
          repo_id=model_id,
          allow_patterns=allow_patterns,
-          cache_dir=str(self.models_dir)
+          local_dir=str(self.models_dir / model_id)
        )

        # Download config files (usually small)
        snapshot_download_with_retry(
          repo_id=model_id,
          allow_patterns=["*config.json"],
-          cache_dir=str(self.models_dir)
+          local_dir=str(self.models_dir / model_id)
        )

        model_data["download_path"] = str(root_path)
@@ -88,8 +88,8 @@ if __name__ == "__main__":
      # repo id
      # validates all onnx models inside repo
      repo_id = "/".join(path)
-      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)
+      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*.onnx", "*.onnx_data"], local_dir=DOWNLOADS_DIR / repo_id)
+      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], local_dir=DOWNLOADS_DIR / repo_id)
      config = get_config(root_path)
      for onnx_model in root_path.rglob("*.onnx"):
        rtol, atol = get_tolerances(onnx_model.name)
@@ -101,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 = 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)
+      root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=[relative_path], local_dir=DOWNLOADS_DIR / repo_id)
+      snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], local_dir=DOWNLOADS_DIR / repo_id)
      config = get_config(root_path)
      rtol, atol = get_tolerances(onnx_model)
      print(f"validating {relative_path} with truncate={args.truncate}, {rtol=}, {atol=}")
@@ -1,99 +1,105 @@
-import os, pathlib
+import os

 # TODO: there is a timing bug without this
 os.environ["AMD_AQL"] = "1"

-from tinygrad.device import Device
-from tinygrad.runtime.ops_amd import AMDProgram, HIPCompiler
+from tinygrad import Tensor, Device
+from tinygrad.helpers import getenv
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from tinygrad.renderer.amd.dsl import Reg, Inst, s, v

 NUM_WORKGROUPS = 96
 WAVE_SIZE = 32
-NUM_WAVES = 2
+NUM_WAVES = 4
 FLOPS_PER_MATMUL = 16*16*16*2
-INTERNAL_LOOP = 1_000_00
+INTERNAL_LOOP = getenv("LOOP", 10_000)
 INSTRUCTIONS_PER_LOOP = 200
-DIRECTIVE = ".amdhsa_wavefront_size32 1"

-assemblyTemplate = (pathlib.Path(__file__).parent / "template.s").read_text()
+def repeat(insts:list[Inst], n:int, counter_sreg:Reg) -> list[Inst]:
+  insts_bytes = b"".join([inst.to_bytes() for inst in insts])
+  sub_inst, cmp_inst = s_sub_u32(counter_sreg, counter_sreg, 1), s_cmp_lg_i32(counter_sreg, 0)
+  loop_sz = len(insts_bytes) + sub_inst.size() + cmp_inst.size()
+  branch_inst = s_cbranch_scc1(simm16=-((loop_sz // 4) + 1) & 0xFFFF)
+  return [s_mov_b32(counter_sreg, n)] + insts + [sub_inst, cmp_inst, branch_inst, s_endpgm()]

-def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, extra=""):
+def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, **kwargs):
  if accum:
-    instructions = "{} a[0:{}], v[{}:{}], v[{}:{}], 1{}\n".format(instruction, vgprIndices[0],
-                                                                vgprIndices[1], vgprIndices[2],
-                                                                vgprIndices[1], vgprIndices[2], extra)
+    inst = instruction(v[0:vgprIndices[0]], v[vgprIndices[1]:vgprIndices[2]], v[vgprIndices[1]:vgprIndices[2]], 1, acc_cd=1, **kwargs)
  elif dense:
-    instructions = "{} v[0:{}], v[{}:{}], v[{}:{}], 1\n".format(instruction, vgprIndices[0],
-                                                                vgprIndices[1], vgprIndices[2],
-                                                                vgprIndices[1], vgprIndices[2])
+    inst = instruction(v[0:vgprIndices[0]], v[vgprIndices[1]:vgprIndices[2]], v[vgprIndices[1]:vgprIndices[2]], 1)
  else:
-    instructions = "{} v[0:{}], v[{}:{}], v[{}:{}], v{}\n".format(instruction, vgprIndices[0],
-                                                                  vgprIndices[1], vgprIndices[2],
-                                                                  vgprIndices[3], vgprIndices[4],
-                                                                  vgprIndices[5])
-  src = assemblyTemplate.replace("INTERNAL_LOOP", str(INTERNAL_LOOP)).replace("INSTRUCTION", instructions*INSTRUCTIONS_PER_LOOP)
-  src = src.replace("DIRECTIVE", DIRECTIVE)
-  lib = COMPILER.compile(src)
-  fxn = AMDProgram(DEV, "matmul", lib)
-  elapsed = min([fxn(global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True) for _ in range(2)])
+    inst = instruction(v[0:vgprIndices[0]], v[vgprIndices[1]:vgprIndices[2]], v[vgprIndices[3]:vgprIndices[4]], v[vgprIndices[5]])
+  insts = repeat([inst for _ in range(INSTRUCTIONS_PER_LOOP)], n=INTERNAL_LOOP, counter_sreg=s[1])
+  def fxn(A:UOp) -> UOp:
+    threads = UOp.special(WAVE_SIZE * NUM_WAVES, "lidx0")
+    gidx = UOp.special(NUM_WORKGROUPS, "gidx0")
+    FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
+    sink = UOp.sink(A.base, threads, gidx, arg=KernelInfo(inst.op.name.lower(), estimates=Estimates(ops=FLOPs, mem=0)))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+  dummy = Tensor.zeros(1).contiguous().realize()
+  out = Tensor.custom_kernel(dummy, fxn=fxn)[0]
+  ei = out.schedule()[-1].lower()
+  elapsed = min([ei.run(wait=True) for _ in range(2)])
  FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
-  print(f"{instruction:<29} : {FLOPs/elapsed/10**12:.2f} T(FL)OPS")
+  print(f"{inst.op_name.lower():<29} : {FLOPs/elapsed/10**12:.2f} T(FL)OPS")

 if __name__=="__main__":
-  DEVICENUM = os.getenv("DEVICENUM", "0")
-  try:
-    DEV = Device['AMD:' + DEVICENUM]
-  except:
-    raise RuntimeError("Error while initiating AMD device")
+  DEV = Device[Device.DEFAULT]
+  arch = DEV.renderer.arch

-  COMPILER = HIPCompiler(DEV.arch)
-  if DEV.arch in {'gfx1100', 'gfx1103', 'gfx1151'}:
-    if DEV.arch == 'gfx1103': NUM_WORKGROUPS = 8
-    if DEV.arch == 'gfx1151': NUM_WORKGROUPS = 32
-    launchBenchmark("v_wmma_bf16_16x16x16_bf16", (7,8,15))
-    launchBenchmark("v_wmma_f16_16x16x16_f16", (7,8,15))
-    launchBenchmark("v_wmma_f32_16x16x16_bf16", (7,8,15))
-    launchBenchmark("v_wmma_f32_16x16x16_f16", (7,8,15))
-    launchBenchmark("v_wmma_i32_16x16x16_iu4", (7,8,9))
-    launchBenchmark("v_wmma_i32_16x16x16_iu8", (7,8,11))
-  elif DEV.arch == 'gfx1201':
+  if arch in {'gfx1100', 'gfx1103', 'gfx1151'}:
+    from tinygrad.runtime.autogen.amd.rdna3.ins import *
+    if arch == 'gfx1103': NUM_WORKGROUPS = 8
+    if arch == 'gfx1151': NUM_WORKGROUPS = 32
+    launchBenchmark(v_wmma_bf16_16x16x16_bf16, (7,8,15))
+    launchBenchmark(v_wmma_f16_16x16x16_f16, (7,8,15))
+    launchBenchmark(v_wmma_f32_16x16x16_bf16, (7,8,15))
+    launchBenchmark(v_wmma_f32_16x16x16_f16, (7,8,15))
+    launchBenchmark(v_wmma_i32_16x16x16_iu4, (7,8,9))
+    launchBenchmark(v_wmma_i32_16x16x16_iu8, (7,8,11))
+  elif arch in {'gfx1200', 'gfx1201'}:
+    from tinygrad.runtime.autogen.amd.rdna4.ins import *
+    # this instruction does not exist in the rdna4 isa, use the co version
+    s_sub_u32 = s_sub_co_u32
    NUM_WORKGROUPS = 64
-    launchBenchmark("v_wmma_bf16_16x16x16_bf16", (3,4,7))
-    launchBenchmark("v_wmma_f16_16x16x16_f16", (3,4,7))
-    launchBenchmark("v_wmma_f32_16x16x16_bf16", (7,8,11))
-    launchBenchmark("v_wmma_f32_16x16x16_f16", (7,8,11))
-    launchBenchmark("v_wmma_i32_16x16x16_iu4", (7,8,8))
-    launchBenchmark("v_wmma_i32_16x16x16_iu8", (7,8,9))
-    launchBenchmark("v_wmma_f32_16x16x16_fp8_fp8", (7,8,9))
-    launchBenchmark("v_wmma_f32_16x16x16_fp8_bf8", (7,8,9))
-    launchBenchmark("v_wmma_f32_16x16x16_bf8_fp8", (7,8,9))
-    launchBenchmark("v_wmma_f32_16x16x16_bf8_bf8", (7,8,9))
+    launchBenchmark(v_wmma_bf16_16x16x16_bf16, (3,4,7))
+    launchBenchmark(v_wmma_f16_16x16x16_f16, (3,4,7))
+    launchBenchmark(v_wmma_f32_16x16x16_bf16, (7,8,11))
+    launchBenchmark(v_wmma_f32_16x16x16_f16, (7,8,11))
+    launchBenchmark(v_wmma_i32_16x16x16_iu4, (7,8,8))
+    launchBenchmark(v_wmma_i32_16x16x16_iu8, (7,8,9))
+    launchBenchmark(v_wmma_f32_16x16x16_fp8_fp8, (7,8,9))
+    launchBenchmark(v_wmma_f32_16x16x16_fp8_bf8, (7,8,9))
+    launchBenchmark(v_wmma_f32_16x16x16_bf8_fp8, (7,8,9))
+    launchBenchmark(v_wmma_f32_16x16x16_bf8_bf8, (7,8,9))
    FLOPS_PER_MATMUL = 16*16*32*2
-    launchBenchmark("v_wmma_i32_16X16X32_iu4", (7,8,9))
-    launchBenchmark("v_swmmac_f32_16x16x32_f16", (7,8,11,12,19,20), False)
-    launchBenchmark("v_swmmac_f32_16x16x32_bf16", (7,8,11,12,19,20), False)
-    launchBenchmark("v_swmmac_f16_16x16x32_f16", (3,4,7,8,15,16), False)
-    launchBenchmark("v_swmmac_bf16_16x16x32_bf16", (3,4,7,8,15,16), False)
-    launchBenchmark("v_swmmac_i32_16x16x32_iu8", (7,8,9,10,13,14), False)
-    launchBenchmark("v_swmmac_i32_16x16x32_iu4", (7,8,8,9,10,11), False)
-    launchBenchmark("v_swmmac_f32_16x16x32_fp8_fp8", (7,8,9,10,13,14), False)
-    launchBenchmark("v_swmmac_f32_16x16x32_fp8_bf8", (7,8,9,10,13,14), False)
-    launchBenchmark("v_swmmac_f32_16x16x32_bf8_fp8", (7,8,9,10,13,14), False)
-    launchBenchmark("v_swmmac_f32_16x16x32_bf8_bf8", (7,8,9,10,13,14), False)
+    launchBenchmark(v_wmma_i32_16x16x32_iu4, (7,8,9))
+    launchBenchmark(v_swmmac_f32_16x16x32_f16, (7,8,11,12,19,20), False)
+    launchBenchmark(v_swmmac_f32_16x16x32_bf16, (7,8,11,12,19,20), False)
+    launchBenchmark(v_swmmac_f16_16x16x32_f16, (3,4,7,8,15,16), False)
+    launchBenchmark(v_swmmac_bf16_16x16x32_bf16, (3,4,7,8,15,16), False)
+    launchBenchmark(v_swmmac_i32_16x16x32_iu8, (7,8,9,10,13,14), False)
+    launchBenchmark(v_swmmac_i32_16x16x32_iu4, (7,8,8,9,10,11), False)
+    launchBenchmark(v_swmmac_f32_16x16x32_fp8_fp8, (7,8,9,10,13,14), False)
+    launchBenchmark(v_swmmac_f32_16x16x32_fp8_bf8, (7,8,9,10,13,14), False)
+    launchBenchmark(v_swmmac_f32_16x16x32_bf8_fp8, (7,8,9,10,13,14), False)
+    launchBenchmark(v_swmmac_f32_16x16x32_bf8_bf8, (7,8,9,10,13,14), False)
    FLOPS_PER_MATMUL = 16*16*64*2
-    launchBenchmark("v_swmmac_i32_16x16x64_iu4", (7,8,9,10,13,14), False)
-  elif DEV.arch == 'gfx950':
-    DIRECTIVE = ".amdhsa_accum_offset 4"
+    launchBenchmark(v_swmmac_i32_16x16x64_iu4, (7,8,9,10,13,14), False)
+  elif arch == 'gfx950':
+    from tinygrad.runtime.autogen.amd.cdna.ins import *
    NUM_WORKGROUPS = 256
    WAVE_SIZE = 64
    NUM_WAVES = 4
-    launchBenchmark("v_mfma_f32_16x16x16_f16", (3,0,1), accum=True)
-    launchBenchmark("v_mfma_f32_16x16x16_bf16", (3,0,1), accum=True)
+    launchBenchmark(v_mfma_f32_16x16x16_f16, (3,0,1), accum=True)
+    launchBenchmark(v_mfma_f32_16x16x16_bf16, (3,0,1), accum=True)
    FLOPS_PER_MATMUL = 16*16*32*2
-    launchBenchmark("v_mfma_f32_16x16x32_f16", (3,0,3), accum=True)
-    launchBenchmark("v_mfma_f32_16x16x32_bf16", (3,0,3), accum=True)
+    launchBenchmark(v_mfma_f32_16x16x32_f16, (3,0,3), accum=True)
+    launchBenchmark(v_mfma_f32_16x16x32_bf16, (3,0,3), accum=True)
    FLOPS_PER_MATMUL = 16*16*128*2
-    launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,7), accum=True) # fp8
-    launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,5), accum=True, extra=", cbsz:2 blgp:2") # fp6
-    launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,3), accum=True, extra=", cbsz:4 blgp:4") # fp4
+    launchBenchmark(v_mfma_f32_16x16x128_f8f6f4, (3,0,7), accum=True) # fp8
+    launchBenchmark(v_mfma_f32_16x16x128_f8f6f4, (3,0,5), accum=True, cbsz=2, blgp=2) # fp6
+    launchBenchmark(v_mfma_f32_16x16x128_f8f6f4, (3,0,3), accum=True, cbsz=4, blgp=4) # fp4
  else:
-    raise RuntimeError(f"arch {DEV.arch} not supported.")
+    raise RuntimeError(f"arch {arch} not supported.")
@@ -1,40 +0,0 @@
-    .text
-    .globl matmul
-    .p2align 8
-    .type matmul,@function
-matmul:
-  s_mov_b32 s1, INTERNAL_LOOP
-  s_mov_b32 s2, 0
-  inner_loop:
-    INSTRUCTION
-    s_sub_u32 s1, s1, 1
-    s_cmp_lg_i32 s1, s2
-    s_cbranch_scc1 inner_loop
-  s_endpgm
-
-.rodata
-.p2align 6
-.amdhsa_kernel matmul
-  .amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
-  .amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
-  DIRECTIVE
-.end_amdhsa_kernel
-
-.amdgpu_metadata
---
-amdhsa.version:
-  - 1
-  - 0
-amdhsa.kernels:
-  - .name: matmul
-    .symbol: matmul.kd
-    .kernarg_segment_size: 0
-    .group_segment_fixed_size: 0
-    .private_segment_fixed_size: 0
-    .kernarg_segment_align: 4
-    .wavefront_size: 32
-    .sgpr_count: 8
-    .vgpr_count: 32
-    .max_flat_workgroup_size: 1024
-...
-.end_amdgpu_metadata
@@ -5,7 +5,6 @@ from tinygrad.nn import Linear, LayerNorm, Embedding, Conv2d
 from typing import List, Optional, Union, Tuple, Dict
 from abc import ABC, abstractmethod
 from functools import lru_cache
-from PIL import Image
 import numpy as np
 import re, gzip

@@ -444,7 +443,8 @@ class OpenClipEncoder:
  # TODO:
  # Should be doable in pure tinygrad, would just require some work and verification.
  # This is very desirable since it would allow for full generation->evaluation in a single JIT call.
-  def prepare_image(self, image:Image.Image) -> Tensor:
+  def prepare_image(self, image) -> Tensor:
+    from PIL import Image
    SIZE = 224
    w, h = image.size
    scale = min(SIZE / h, SIZE / w)
@@ -55,12 +55,15 @@ class Attention:
      xqkv = x @ self.wqkv.T
      xq, xk, xv = xqkv.split([self.wq.weight.shape[0], self.wk.weight.shape[0], self.wv.weight.shape[0]], dim=2)
    else:
-      xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
+      xq, xk, xv = self.wq(x), self.wk(x.contiguous_backward()), self.wv(x)

    if self.q_norm is not None and self.k_norm is not None:
      xq = self.q_norm(xq)
      xk = self.k_norm(xk)

+    # cast_float_to_bf16 is expensive in reduction loops, break it out
+    if x.dtype == dtypes.bfloat16: xq, xk = xq.contiguous_backward(), xk.contiguous_backward()
+
    xq = xq.reshape(xq.shape[0], xq.shape[1], self.n_heads, self.head_dim)
    xk = xk.reshape(xk.shape[0], xk.shape[1], self.n_kv_heads, self.head_dim)
    xv = xv.reshape(xv.shape[0], xv.shape[1], self.n_kv_heads, self.head_dim)
@@ -86,20 +89,23 @@ class Attention:
      assert start_pos == 0
      keys, values = xk, xv

-    keys, values = repeat_kv(keys, self.n_rep), repeat_kv(values, self.n_rep)
-    xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
-    attn = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2)
+    if self.max_context:
+      keys, values = repeat_kv(keys, self.n_rep), repeat_kv(values, self.n_rep)
+      xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
+      attn = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2)
+    else:
+      xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
+      attn = xq.scaled_dot_product_attention(keys, values, is_causal=True, enable_gqa=True).transpose(1, 2)
    if getenv("STUB_ATTENTION"):
-      # TODO: do we need mask?
      from tinygrad.uop.ops import UOp, KernelInfo
      def fa_custom_forward(attn:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
        return UOp.sink(arg=KernelInfo(name="fa_custom_forward"))
      def fa_custom_backward(out_q:UOp, out_k:UOp, out_v:UOp, grad:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
        return UOp.sink(arg=KernelInfo(name="fa_custom_backward"))
      def fa_backward(grad:UOp, kernel:UOp) -> tuple[None, UOp, UOp, UOp]:
-        grad_q = Tensor.empty_like(q:=Tensor(kernel.src[1]))
-        grad_k = Tensor.empty_like(k:=Tensor(kernel.src[2]))
-        grad_v = Tensor.empty_like(v:=Tensor(kernel.src[3]))
+        grad_q = Tensor.empty_like(q:=Tensor(kernel.src[2]))
+        grad_k = Tensor.empty_like(k:=Tensor(kernel.src[3]))
+        grad_v = Tensor.empty_like(v:=Tensor(kernel.src[4]))
        ck = Tensor.custom_kernel(grad_q, grad_k, grad_v, Tensor(grad), q, k, v, fxn=fa_custom_backward)[:3]
        return (None, ck[0].uop, ck[1].uop, ck[2].uop)
      attn = Tensor.empty_like(attn).custom_kernel(xq, keys, values, fxn=fa_custom_forward, grad_fxn=fa_backward)[0]
@@ -197,7 +203,9 @@ class Transformer:
    h = self.tok_embeddings(tokens)
    freqs_cis = self.freqs_cis.cast(h.dtype)[:, 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) if seqlen > 1 else None
+    if self.max_context != 0 and seqlen > 1:
+      mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1)
+    else: mask = None
    for layer in self.layers: h = layer(h, start_pos, freqs_cis, mask)
    logits = self.output(self.norm(h))
    if math.isnan(temperature): return logits
@@ -150,7 +150,7 @@ class ResNet:
        continue # Skip FC if transfer learning

      if 'bn' not in k and 'downsample' not in k: assert obj.shape == dat.shape, (k, obj.shape, dat.shape)
-      obj.assign(dat.to(obj.device).reshape(obj.shape))
+      obj.assign(dat.to(obj.device).cast(obj.dtype).reshape(obj.shape))

 ResNet18 = lambda num_classes=1000: ResNet(18, num_classes=num_classes)
 ResNet34 = lambda num_classes=1000: ResNet(34, num_classes=num_classes)
@@ -129,7 +129,7 @@ class LSTM:
      return self.do_step(x_, hc_)

    if hc is None:
-      hc = Tensor.zeros(self.layers, 2 * x.shape[1], self.hidden_size, requires_grad=False)
+      hc = Tensor.zeros(self.layers, 2 * x.shape[1], self.hidden_size, requires_grad=False).contiguous().realize()

    output = None
    for t in range(x.shape[0]):
@@ -1,6 +1,7 @@
 # type: ignore
 import ctypes, ctypes.util, struct, platform, pathlib, re, time, os, signal
-from tinygrad.helpers import from_mv, to_mv, getenv, init_c_struct_t
+from tinygrad.helpers import from_mv, to_mv, getenv
+from tinygrad.runtime.support.c import init_c_struct_t
 from hexdump import hexdump
 start = time.perf_counter()

@@ -10,18 +11,21 @@ processor = platform.processor()
 IOCTL_SYSCALL = {"aarch64": 0x1d, "x86_64":16}[processor]
 MMAP_SYSCALL = {"aarch64": 0xde, "x86_64":0x09}[processor]

+IOCTL_PRINT = getenv("IOCTL_PRINT", getenv("IOCTL", 0))
+GRAB_PMA = getenv("GRAB_PMA", 0)
+
 def get_struct(argp, stype):
  return ctypes.cast(ctypes.c_void_p(argp), ctypes.POINTER(stype)).contents

 def dump_struct(st):
-  if getenv("IOCTL", 0) == 0: return
+  if IOCTL_PRINT == 0: return
  print("\t", st.__class__.__name__, end=" { ")
-  for v in type(st)._fields_: print(f"{v[0]}={getattr(st, v[0])}", end=" ")
+  for v in type(st)._real_fields_: print(f"{v[0]}={getattr(st, v[0])}", end=" ")
  print("}")

 def format_struct(s):
  sdats = []
-  for field in s._fields_:
+  for field in s._real_fields_:
    dat = getattr(s, field[0])
    if isinstance(dat, int): sdats.append(f"{field[0]}:0x{dat:X}")
    else: sdats.append(f"{field[0]}:{dat}")
@@ -58,6 +62,29 @@ def install_hook(c_function, python_function):
  return orig_func

 # *** ioctl lib end ***
+
+# PMA buffer tracking for raw PC sampling data (only when GRAB_PMA is enabled)
+pma_mem_handle = 0  # hMemPmaBuffer from ALLOC_PMA_STREAM
+pma_buffer_size = 0
+pma_buffer_va = 0   # actual mapped VA (found via /proc/self/maps)
+pma_get_offset = 0  # current read offset in ring buffer
+pma_pending_map = False  # flag to check for new mapping on next ioctl
+pma_maps_before = set()  # mappings before MAP_MEMORY
+pma_raw_dumps: list[bytes] = []  # raw PMA buffer dumps
+
+def get_pma_raw_dumps() -> list[bytes]: return pma_raw_dumps
+def clear_pma_raw_dumps(): pma_raw_dumps.clear()
+
+def get_proc_maps():
+  """Read current process memory mappings as set of (start, end) tuples."""
+  result = set()
+  with open("/proc/self/maps", "r") as f:
+    for line in f:
+      addr_range = line.split()[0]
+      start, end = addr_range.split("-")
+      result.add((int(start, 16), int(end, 16)))
+  return result
+
 from tinygrad.runtime.autogen import nv_570 as nv_gpu
 nvescs = {getattr(nv_gpu, x):x for x in dir(nv_gpu) if x.startswith("NV_ESC")}
 nvcmds = {getattr(nv_gpu, x):(x, getattr(nv_gpu, "struct_"+x+"_PARAMS", getattr(nv_gpu, "struct_"+x.replace("_CMD_", "_")+"_PARAMS", None))) for x in dir(nv_gpu) if \
@@ -69,6 +96,7 @@ def get_classes():
                   "NV20_SUBDEVICE_0"}
  for nm,val in nv_gpu.__dict__.items():
    if not isinstance(val, int): continue
+    if nm.endswith("PARAMETERS_MESSAGE_ID"): continue
    if 0x3000 < val < 0xffff: res[val] = nm
    if nm in known_classes: res[val] = nm
  return res
@@ -80,37 +108,92 @@ global_ioctl_id = 0
 gpus_user_modes = []
 gpus_mmio = []
 gpus_fifo = []
+offset_load = 0

@ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_int, ctypes.c_ulong, ctypes.c_void_p)
 def ioctl(fd, request, argp):
  global global_ioctl_id, gpus_user_modes, gpus_mmio
+  global pma_mem_handle, pma_buffer_size, pma_buffer_va, pma_get_offset, pma_pending_map, pma_maps_before
  global_ioctl_id += 1
+
+  # Check for new PMA buffer mapping from previous MAP_MEMORY call (only when GRAB_PMA is enabled)
+  if GRAB_PMA and pma_pending_map:
+    pma_pending_map = False
+    new_maps = get_proc_maps()
+    for start, end in new_maps - pma_maps_before:
+      if end - start == pma_buffer_size:
+        pma_buffer_va = start
+        if IOCTL_PRINT >= 1: print(f"\t  PMA buffer mapped at CPU VA=0x{pma_buffer_va:x}")
+        break
+
  st = time.perf_counter()
  ret = libc.syscall(IOCTL_SYSCALL, ctypes.c_int(fd), ctypes.c_ulong(request), ctypes.c_void_p(argp))
  et = time.perf_counter()-st
  fn = os.readlink(f"/proc/self/fd/{fd}")
-  #print(f"ioctl {request:8x} {fn:20s}")
+
  idir, size, itype, nr = (request>>30), (request>>16)&0x3FFF, (request>>8)&0xFF, request&0xFF
-  if getenv("IOCTL", 0) >= 1: print(f"#{global_ioctl_id}: ", end="")
+  if IOCTL_PRINT >= 1: print(f"#{global_ioctl_id}: ", end="")
  if itype == ord(nv_gpu.NV_IOCTL_MAGIC):
    if nr == nv_gpu.NV_ESC_RM_CONTROL:
      s = get_struct(argp, nv_gpu.NVOS54_PARAMETERS)
      if s.cmd in nvcmds:
        name, struc = nvcmds[s.cmd]
-        if getenv("IOCTL", 0) >= 1:
+        if IOCTL_PRINT >= 1:
          print(f"NV_ESC_RM_CONTROL    cmd={name:30s} hClient={s.hClient}, hObject={s.hObject}, flags={s.flags}, params={s.params}, paramsSize={s.paramsSize}, status={s.status}")

        if struc is not None: dump_struct(get_struct(s.params, struc))
        elif hasattr(nv_gpu, name+"_PARAMS"): dump_struct(get_struct(argp, getattr(nv_gpu, name+"_PARAMS")))
        elif name == "NVA06C_CTRL_CMD_GPFIFO_SCHEDULE": dump_struct(get_struct(argp, nv_gpu.NVA06C_CTRL_GPFIFO_SCHEDULE_PARAMS))
        elif name == "NV83DE_CTRL_CMD_GET_MAPPINGS": dump_struct(get_struct(s.params, nv_gpu.NV83DE_CTRL_DEBUG_GET_MAPPINGS_PARAMETERS))
+        elif name == "NVB0CC_CTRL_CMD_SET_HS_CREDITS":
+          hs_params = get_struct(s.params, nv_gpu.NVB0CC_CTRL_SET_HS_CREDITS_PARAMS)
+          dump_struct(hs_params)
+          if IOCTL_PRINT >= 2:
+            for i in range(hs_params.numEntries):
+              print(f"\t\t", end="")
+              dump_struct(hs_params.creditInfo[i])
+
+        # PMA buffer tracking (only when GRAB_PMA is enabled)
+        if GRAB_PMA and name == "NVB0CC_CTRL_CMD_ALLOC_PMA_STREAM":
+          pma_params = get_struct(s.params, nv_gpu.struct_NVB0CC_CTRL_ALLOC_PMA_STREAM_PARAMS)
+          pma_mem_handle = pma_params.hMemPmaBuffer
+          pma_buffer_size = pma_params.pmaBufferSize
+          pma_get_offset = 0  # Reset read offset for new stream
+          if IOCTL_PRINT >= 1: print(f"\t  PMA buffer: hMem=0x{pma_mem_handle:x} size={pma_buffer_size}")
+        if GRAB_PMA and name == "NVB0CC_CTRL_CMD_PMA_STREAM_UPDATE_GET_PUT":
+          pma_update = get_struct(s.params, nv_gpu.struct_NVB0CC_CTRL_PMA_STREAM_UPDATE_GET_PUT_PARAMS)
+          if pma_update.bytesAvailable > 0 and pma_buffer_va and pma_buffer_size > 0:
+            avail = pma_update.bytesAvailable
+            read_offset = pma_get_offset
+            # Handle ring buffer wrap-around
+            if pma_get_offset + avail <= pma_buffer_size:
+              pma_data = bytes(to_mv(pma_buffer_va + pma_get_offset, avail))
+            else:
+              # Wrap around: read to end, then from start
+              first_part = pma_buffer_size - pma_get_offset
+              second_part = avail - first_part
+              pma_data = bytes(to_mv(pma_buffer_va + pma_get_offset, first_part)) + bytes(to_mv(pma_buffer_va, second_part))
+            pma_raw_dumps.append(pma_data)
+            pma_get_offset = (pma_get_offset + avail) % pma_buffer_size
+            if IOCTL_PRINT >= 2:
+              print(f"\t  PMA data: {avail} bytes from offset=0x{read_offset:x}, new offset=0x{pma_get_offset:x}")
+              hexdump(pma_data)
+
+        # Dump regOps for EXEC_REG_OPS when IOCTL >= 3
+        if name == "NVB0CC_CTRL_CMD_EXEC_REG_OPS" and struc is not None and IOCTL_PRINT >= 3:
+          reg_params = get_struct(s.params, struc)
+          for i in range(reg_params.regOpCount):
+            print(f"\t\t", end="")
+            dump_struct(reg_params.regOps[i])
+            # val = (op.regValueHi << 32) | op.regValueLo
+            # print(f"\t  regOps[{i:3d}]: op={op.regOp} type={op.regType} status={op.regStatus} offset=0x{op.regOffset:08x} value=0x{val:016x}")
      else:
-        if getenv("IOCTL", 0) >= 1: print("unhandled cmd", hex(s.cmd))
+        if IOCTL_PRINT >= 1: print("unhandled cmd", hex(s.cmd))
      # format_struct(s)
      # print(f"{(st-start)*1000:7.2f} ms +{et*1000.:7.2f} ms : {ret:2d} = {name:40s}", ' '.join(format_struct(s)))
    elif nr == nv_gpu.NV_ESC_RM_ALLOC:
      s = get_struct(argp, nv_gpu.NVOS21_PARAMETERS)
-      if getenv("IOCTL", 0) >= 1: print(f"NV_ESC_RM_ALLOC    hClass={nvclasses.get(s.hClass, f'unk=0x{s.hClass:X}'):30s}, hRoot={s.hRoot}, hObjectParent={s.hObjectParent}, pAllocParms={s.pAllocParms}, hObjectNew={s.hObjectNew} status={s.status}")
+      if IOCTL_PRINT >= 1: print(f"NV_ESC_RM_ALLOC    hClass={nvclasses.get(s.hClass, f'unk=0x{s.hClass:X}'):30s}, hRoot={s.hRoot}, hObjectParent={s.hObjectParent}, pAllocParms={s.pAllocParms}, hObjectNew={s.hObjectNew} status={s.status}")
      if s.pAllocParms is not None:
        if s.hClass == nv_gpu.NV01_DEVICE_0: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV0080_ALLOC_PARAMETERS))
        if s.hClass == nv_gpu.FERMI_VASPACE_A: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV_VASPACE_ALLOCATION_PARAMETERS))
@@ -118,7 +201,8 @@ def ioctl(fd, request, argp):
        if s.hClass == nv_gpu.NV1_MEMORY_USER: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV_MEMORY_ALLOCATION_PARAMS))
        if s.hClass == nv_gpu.NV1_MEMORY_SYSTEM: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV_MEMORY_ALLOCATION_PARAMS))
        if s.hClass == nv_gpu.GT200_DEBUGGER: dump_struct(get_struct(s.pAllocParms, nv_gpu.NV83DE_ALLOC_PARAMETERS))
-        if s.hClass == nv_gpu.AMPERE_CHANNEL_GPFIFO_A:
+        if s.hClass == nv_gpu.MAXWELL_PROFILER_DEVICE: dump_struct(get_struct(s.pAllocParms, nv_gpu.NVB2CC_ALLOC_PARAMETERS))
+        if s.hClass in {nv_gpu.AMPERE_CHANNEL_GPFIFO_A, nv_gpu.BLACKWELL_CHANNEL_GPFIFO_A}:
          sx = get_struct(s.pAllocParms, nv_gpu.NV_CHANNELGPFIFO_ALLOCATION_PARAMETERS)
          dump_struct(sx)
          gpus_fifo.append((sx.gpFifoOffset, sx.gpFifoEntries))
@@ -126,31 +210,35 @@ def ioctl(fd, request, argp):
      if s.hClass == nv_gpu.TURING_USERMODE_A: gpus_user_modes.append(s.hObjectNew)
    elif nr == nv_gpu.NV_ESC_RM_MAP_MEMORY:
      # nv_ioctl_nvos33_parameters_with_fd
-      if getenv("IOCTL", 0) >= 1:
-        s = get_struct(argp, nv_gpu.NVOS33_PARAMETERS)
+      s = get_struct(argp, nv_gpu.NVOS33_PARAMETERS)
+      if IOCTL_PRINT >= 1:
        print(f"NV_ESC_RM_MAP_MEMORY   hClient={s.hClient}, hDevice={s.hDevice}, hMemory={s.hMemory}, length={s.length} flags={s.flags} pLinearAddress={s.pLinearAddress}")
+      # Track PMA buffer mapping - save maps now, check for new mapping on next ioctl (after mmap happens)
+      if GRAB_PMA and pma_mem_handle and s.hMemory == pma_mem_handle:
+        pma_maps_before = get_proc_maps()
+        pma_pending_map = True
    elif nr == nv_gpu.NV_ESC_RM_UPDATE_DEVICE_MAPPING_INFO:
-      if getenv("IOCTL", 0) >= 1:
+      if IOCTL_PRINT >= 1:
        s = get_struct(argp, nv_gpu.NVOS56_PARAMETERS)
        print(f"NV_ESC_RM_UPDATE_DEVICE_MAPPING_INFO   hClient={s.hClient}, hDevice={s.hDevice}, hMemory={s.hMemory}, pOldCpuAddress={s.pOldCpuAddress} pNewCpuAddress={s.pNewCpuAddress} status={s.status}")
    elif nr == nv_gpu.NV_ESC_RM_ALLOC_MEMORY:
-      if getenv("IOCTL", 0) >= 1:
+      if IOCTL_PRINT >= 1:
        s = get_struct(argp, nv_gpu.nv_ioctl_nvos02_parameters_with_fd)
        print(f"NV_ESC_RM_ALLOC_MEMORY  fd={s.fd}, hRoot={s.params.hRoot}, hObjectParent={s.params.hObjectParent}, hObjectNew={s.params.hObjectNew}, hClass={s.params.hClass}, flags={s.params.flags}, pMemory={s.params.pMemory}, limit={s.params.limit}, status={s.params.status}")
    elif nr == nv_gpu.NV_ESC_ALLOC_OS_EVENT:
-      if getenv("IOCTL", 0) >= 1:
+      if IOCTL_PRINT >= 1:
        s = get_struct(argp, nv_gpu.nv_ioctl_alloc_os_event_t)
        print(f"NV_ESC_ALLOC_OS_EVENT  hClient={s.hClient} hDevice={s.hDevice} fd={s.fd} Status={s.Status}")
    elif nr == nv_gpu.NV_ESC_REGISTER_FD:
-      if getenv("IOCTL", 0) >= 1:
+      if IOCTL_PRINT >= 1:
        s = get_struct(argp, nv_gpu.nv_ioctl_register_fd_t)
        print(f"NV_ESC_REGISTER_FD  fd={s.ctl_fd}")
    elif nr in nvescs:
-      if getenv("IOCTL", 0) >= 1: print(nvescs[nr])
+      if IOCTL_PRINT >= 1: print(nvescs[nr])
    else:
-      if getenv("IOCTL", 0) >= 1: print("unhandled NR", nr)
+      if IOCTL_PRINT >= 1: print("unhandled NR", nr)
  elif fn.endswith("nvidia-uvm"):
-    if getenv("IOCTL", 0) >= 1:
+    if IOCTL_PRINT >= 1:
      print(f"{nvuvms.get(request, f'UVM UNKNOWN {request=}')}")
      if nvuvms.get(request) is not None: dump_struct(get_struct(argp, getattr(nv_gpu, nvuvms.get(request)+"_PARAMS")))
      if nvuvms.get(request) == "UVM_MAP_EXTERNAL_ALLOCATION":
@@ -159,7 +247,7 @@ def ioctl(fd, request, argp):
          print("perGpuAttributes[{i}] = ", end="")
          dump_struct(st.perGpuAttributes[i])

-  if getenv("IOCTL") >= 2: print("ioctl", f"{idir=} {size=} {itype=} {nr=} {fd=} {ret=}", fn)
+  if IOCTL_PRINT >= 2: print("ioctl", f"{idir=} {size=} {itype=} {nr=} {fd=} {ret=}", fn)
  return ret

@ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_long)
@@ -172,14 +260,14 @@ def _mmap(addr, length, prot, flags, fd, offset):
  return ret

 install_hook(libc.ioctl, ioctl)
-if getenv("IOCTL") >= 3: orig_mmap_mv = install_hook(libc.mmap, _mmap)
+if getenv("IOCTL") >= 4: orig_mmap_mv = install_hook(libc.mmap, _mmap)

 import collections
 old_gpputs = collections.defaultdict(int)
 def _dump_gpfifo(mark):
  launches = []

-  # print("_dump_gpfifo:", mark)
+  print("_dump_gpfifo:", mark)
  for start, size in gpus_fifo:
    gpfifo_controls = nv_gpu.AmpereAControlGPFifo.from_address(start+size*8)
    gpfifo = to_mv(start, size * 8).cast("Q")
@@ -205,7 +293,7 @@ def make_qmd_struct_type():
    fields.append((name.replace("NVC6C0_QMDV03_00_", "").lower(), ctypes.c_uint32, data[0]-data[1]+1))
    if len(fields) >= 2 and fields[-2][0].endswith('_lower') and fields[-1][0].endswith('_upper') and fields[-1][0][:-6] == fields[-2][0][:-6]:
      fields = fields[:-2] + [(fields[-1][0][:-6], ctypes.c_uint64, fields[-1][2] + fields[-2][2])]
-  return init_c_struct_t(tuple(fields))
+  return init_c_struct_t(0x40 * 4, tuple(fields))
 qmd_struct_t = make_qmd_struct_type()
 assert ctypes.sizeof(qmd_struct_t) == 0x40 * 4

@@ -222,7 +310,7 @@ def _dump_qmd(address, packets):
    subc = (dat>>13) & 7
    mthd = (dat<<2) & 0x7FFF
    method_name = nvqcmds.get(mthd, f"unknown method #{mthd}")
-    if getenv("IOCTL", 0) >= 1:
+    if IOCTL_PRINT >= 1:
      print(f"\t\t{method_name}, {typ=} {size=} {subc=} {mthd=}")
      for j in range(size): print(f"\t\t\t{j}: {gpfifo[i+j+1]} | 0x{gpfifo[i+j+1]:x}")
    if mthd == 792:
@@ -0,0 +1 @@
+examples/
@@ -0,0 +1,135 @@
+import pickle, os, sys, functools, numpy as np
+from pathlib import Path
+
+os.environ["DEV"] = "CUDA"
+os.environ["PROFILE"] = os.environ.get("PROFILE", "2")
+from extra.nv_pma.cupti import cu_prof_ext
+cu_prof_ext.enable_auto()
+
+from tinygrad import Tensor, Device
+
+if not os.environ.get("IOCTL") or not os.environ.get("GRAB_PMA"):
+  print("Usage: GRAB_PMA=1 IOCTL=1 IOCTL_PRINT=0 python3 extra/nv_pma/collect.py")
+  sys.exit(1)
+
+assert Device.DEFAULT == "CUDA", "only works with CUDA"
+
+EXAMPLES_DIR = Path(__file__).parent / "examples"
+_collectors: list[tuple[str, callable]] = []
+
+def pcsampling_test(name: str):
+  def decorator(fn):
+    @functools.wraps(fn)
+    def wrapper():
+      cu_prof_ext.clear_pma_raw_dumps()
+      cu_prof_ext.clear_cupti_pc_samples()
+
+      fn()
+      Device["CUDA"].synchronize()
+
+      dumps = cu_prof_ext.get_pma_raw_dumps()
+      # from hexdump import hexdump
+      # hexdump(dumps[0][:0x40])
+
+      return {"test_name": name, "pma_raw_dumps": list(cu_prof_ext.get_pma_raw_dumps()), "cupti_pc_samples": list(cu_prof_ext.get_cupti_pc_samples())}
+    _collectors.append((name, wrapper))
+    return wrapper
+  return decorator
+
+# Refs
+
+@pcsampling_test("test_plus")
+def test_plus():
+  a = Tensor([1, 2, 3, 4])
+  b = Tensor([5, 6, 7, 8])
+  (a + b).realize()
+
+@pcsampling_test("test_matmul")
+def test_matmul():
+  a = Tensor(np.random.rand(12, 12).astype(np.float32))
+  b = Tensor(np.random.rand(12, 12).astype(np.float32))
+  (a @ b).realize()
+
+@pcsampling_test("test_reduce_sum")
+def test_reduce_sum():
+  a = Tensor(np.random.rand(1024).astype(np.float32))
+  a.sum().realize()
+
+@pcsampling_test("test_reduce_max")
+def test_reduce_max():
+  a = Tensor(np.random.rand(1024).astype(np.float32))
+  a.max().realize()
+
+@pcsampling_test("test_exp")
+def test_exp():
+  a = Tensor(np.random.rand(256).astype(np.float32))
+  a.exp().realize()
+
+@pcsampling_test("test_softmax")
+def test_softmax():
+  a = Tensor(np.random.rand(64, 64).astype(np.float32))
+  a.softmax().realize()
+
+@pcsampling_test("test_conv2d")
+def test_conv2d():
+  x = Tensor(np.random.rand(1, 3, 32, 32).astype(np.float32))
+  w = Tensor(np.random.rand(8, 3, 3, 3).astype(np.float32))
+  x.conv2d(w).realize()
+
+@pcsampling_test("test_large_matmul")
+def test_large_matmul():
+  a = Tensor(np.random.rand(128, 128).astype(np.float32))
+  b = Tensor(np.random.rand(128, 128).astype(np.float32))
+  (a @ b).realize()
+
+@pcsampling_test("test_elementwise_chain")
+def test_elementwise_chain():
+  a = Tensor(np.random.rand(512).astype(np.float32))
+  ((a + 1) * 2 - 0.5).relu().realize()
+
+@pcsampling_test("test_broadcast")
+def test_broadcast():
+  a = Tensor(np.random.rand(64, 1).astype(np.float32))
+  b = Tensor(np.random.rand(1, 64).astype(np.float32))
+  (a + b).realize()
+
+@pcsampling_test("test_plus_big")
+def test_plus_big():
+  a = Tensor(np.random.rand(64, 32).astype(np.float32))
+  b = Tensor(np.random.rand(64, 32).astype(np.float32))
+  (a + b).realize()
+
+def save_example(name: str, data: dict):
+  pma_bytes = sum(len(d) for d in data['pma_raw_dumps'])
+  cupti_samples = sum(r['samples'] for r in data['cupti_pc_samples'])
+  print(f"  PMA: {len(data['pma_raw_dumps'])} buffers, {pma_bytes} bytes")
+  print(f"  CUPTI: {len(data['cupti_pc_samples'])} records, {cupti_samples} samples")
+
+  outfile = EXAMPLES_DIR / f"{name}.pkl"
+  with open(outfile, "wb") as f:
+    pickle.dump(data, f)
+  print(f"  Saved to {outfile}")
+
+if __name__ == "__main__":
+  EXAMPLES_DIR.mkdir(exist_ok=True)
+
+  # Run specific tests if provided as arguments, otherwise run all
+  if len(sys.argv) > 1:
+    test_names = sys.argv[1:]
+    collectors = [(name, fn) for name, fn in _collectors if name in test_names]
+    if not collectors:
+      print(f"Unknown tests: {test_names}")
+      print(f"Available: {[name for name, _ in _collectors]}")
+      sys.exit(1)
+  else:
+    collectors = _collectors
+
+  for name, collect_fn in collectors:
+    print(f"\nCollecting {name}...")
+    try:
+      data = collect_fn()
+      save_example(name, data)
+    except Exception as e:
+      print(f"  ERROR: {e}")
+      import traceback
+      traceback.print_exc()
@@ -0,0 +1,26 @@
+# CUPTI autogen loader for nv_pma
+# To regenerate: REGEN=1 python -c "import extra.nv_pma.cupti"
+import importlib, pathlib
+from tinygrad.helpers import getenv
+
+root = pathlib.Path(__file__).parents[3]
+here = pathlib.Path(__file__).parent
+
+def load(name, dll, files, **kwargs):
+  if not (f:=here/f"{name}.py").exists() or getenv('REGEN'):
+    kwargs['args'] = kwargs.get('args', [])
+    f.write_text(importlib.import_module("tinygrad.runtime.support.autogen").gen(name, dll, files, **kwargs))
+  return importlib.import_module(f"extra.nv_pma.cupti.{name}")
+
+def __getattr__(nm):
+  match nm:
+    case "cupti":
+      return load("cupti", "'/usr/local/cuda/targets/x86_64-linux/lib/libcupti.so'", [
+        "/usr/local/cuda/include/cupti_result.h", "/usr/local/cuda/include/cupti_activity.h",
+        "/usr/local/cuda/include/cupti_callbacks.h", "/usr/local/cuda/include/cupti_events.h",
+        "/usr/local/cuda/include/cupti_metrics.h", "/usr/local/cuda/include/cupti_driver_cbid.h",
+        "/usr/local/cuda/include/cupti_runtime_cbid.h", "/usr/local/cuda/include/cupti_profiler_target.h",
+        "/usr/local/cuda/include/cupti_profiler_host.h", "/usr/local/cuda/include/cupti_pmsampling.h",
+        "/usr/local/cuda/include/generated_cuda_meta.h", "/usr/local/cuda/include/generated_cuda_runtime_api_meta.h"
+      ], args=["-D__CUDA_API_VERSION_INTERNAL", "-I/usr/local/cuda/include"], parse_macros=False)
+    case _: raise AttributeError(f"no such autogen: {nm}")
@@ -0,0 +1,164 @@
+from __future__ import annotations
+import ctypes
+from tinygrad.helpers import DEBUG, getenv
+from extra.nv_pma.cupti import cupti
+
+def stall_reason_name(reason: int) -> str:
+  name = cupti.CUpti_ActivityPCSamplingStallReason.get(reason)
+  return name.replace("CUPTI_ACTIVITY_PC_SAMPLING_STALL_", "").lower() if name else str(reason)
+
+class CUPTIProfiler:
+  def __init__(self):
+    self.initialized = False
+    self.pc_sampling_enabled = False
+    self.buffers: list[ctypes.Array] = []
+    self.kernel_stalls: dict[int, dict[int, int]] = {}
+    self.raw_buffers: list[bytes] = []
+    self.pc_samples: list[dict] = []
+
+  def _check_cupti(self, status, soft=False):
+    if status != cupti.CUPTI_SUCCESS:
+      if soft: return False
+      raise RuntimeError(f"CUPTI Error {status}")
+    return True
+
+  def init(self, ctx, device_id: int = 0, profile_level: int = 2):
+    if self.initialized: return
+
+    # Initialize profiler API
+    init_params = cupti.CUpti_Profiler_Initialize_Params()
+    init_params.structSize = 16
+    cupti.cuptiProfilerInitialize(ctypes.byref(init_params))
+
+    # Register buffer callbacks for Activity API
+    self._buf_req_cb = cupti.CUpti_BuffersCallbackRequestFunc(self._buffer_requested)
+    self._buf_comp_cb = cupti.CUpti_BuffersCallbackCompleteFunc(self._buffer_completed)
+    self._check_cupti(cupti.cuptiActivityRegisterCallbacks(self._buf_req_cb, self._buf_comp_cb))
+
+    # PROFILE=1: kernel timing, PROFILE=2: PC sampling with stall reasons
+    if profile_level >= 2:
+      # PC sampling for stall analysis (requires elevated privileges)
+      if DEBUG >= 1: print("  CUPTI: PC sampling mode (before)")
+      pc_status = cupti.cuptiActivityEnable(cupti.CUPTI_ACTIVITY_KIND_PC_SAMPLING)
+      if pc_status == cupti.CUPTI_SUCCESS:
+        config = cupti.CUpti_ActivityPCSamplingConfig()
+        config.size, config.samplingPeriod = 16, cupti.CUPTI_ACTIVITY_PC_SAMPLING_PERIOD_MIN
+        cfg_status = cupti.dll.cuptiActivityConfigurePCSampling(ctx, ctypes.byref(config))
+        if cfg_status == cupti.CUPTI_SUCCESS:
+          if DEBUG >= 1: print("  CUPTI: PC sampling mode (before stall analysis)")
+          cupti.cuptiActivityEnable(cupti.CUPTI_ACTIVITY_KIND_PC_SAMPLING_RECORD_INFO)
+          self.pc_sampling_enabled = True
+          if DEBUG >= 1: print("  CUPTI: PC sampling mode (stall analysis)")
+        elif cfg_status == 35:
+          if DEBUG >= 1: print("  CUPTI: PC sampling needs: echo 'options nvidia NVreg_RestrictProfilingToAdminUsers=0'|sudo tee /etc/modprobe.d/nvidia.conf && sudo reboot")
+      # Fall back to kernel timing if PC sampling setup failed
+      if not self.pc_sampling_enabled:
+        self._check_cupti(cupti.cuptiActivityEnable(cupti.CUPTI_ACTIVITY_KIND_KERNEL))
+    else:
+      # Kernel activity tracing for timing
+      self._check_cupti(cupti.cuptiActivityEnable(cupti.CUPTI_ACTIVITY_KIND_KERNEL))
+
+    self.initialized = True
+
+  def _buffer_requested(self, buffer, size, max_num_records):
+    buf = (ctypes.c_uint8 * 1024 * 1024)()  # 1MB buffer
+    self.buffers.append(buf)
+    buffer[0] = ctypes.cast(buf, ctypes.POINTER(ctypes.c_uint8))
+    size[0] = ctypes.sizeof(buf)
+    max_num_records[0] = 0
+
+  def _buffer_completed(self, ctx, stream_id, buffer, size, valid_size):
+    if valid_size > 0:
+      record = ctypes.POINTER(cupti.CUpti_Activity)()
+      while cupti.cuptiActivityGetNextRecord(buffer, valid_size, ctypes.byref(record)) == cupti.CUPTI_SUCCESS:
+        kind = record.contents.kind
+        if kind == cupti.CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL:
+          kernel = ctypes.cast(record, ctypes.POINTER(cupti.CUpti_ActivityKernel9)).contents
+          name = ctypes.string_at(kernel.name).decode() if kernel.name else "unknown"
+          duration_us = (kernel.end - kernel.start) / 1000.0
+          grid, block = (kernel.gridX, kernel.gridY, kernel.gridZ), (kernel.blockX, kernel.blockY, kernel.blockZ)
+          print(f"  CUPTI: {name[:40]:40s} | {duration_us:10.2f} us | grid={grid} block={block} | regs={kernel.registersPerThread:3d} smem={kernel.staticSharedMemory + kernel.dynamicSharedMemory:6d}B")
+        elif kind == cupti.CUPTI_ACTIVITY_KIND_PC_SAMPLING:
+          pc = ctypes.cast(record, ctypes.POINTER(cupti.CUpti_ActivityPCSampling3)).contents
+          cid = pc.correlationId
+          if cid not in self.kernel_stalls: self.kernel_stalls[cid] = {}
+          self.kernel_stalls[cid][pc.stallReason] = self.kernel_stalls[cid].get(pc.stallReason, 0) + pc.samples
+          self.pc_samples.append({
+            'correlationId': pc.correlationId, 'pcOffset': pc.pcOffset, 'stallReason': pc.stallReason,
+            'samples': pc.samples, 'latencySamples': pc.latencySamples, 'functionId': pc.functionId, 'sourceLocatorId': pc.sourceLocatorId
+          })
+          if DEBUG >= 3:
+            print(f"    PC {pc.pcOffset:#x} stall={stall_reason_name(pc.stallReason)} samples={pc.samples} latency={pc.latencySamples} func={pc.functionId} src={pc.sourceLocatorId}")
+        elif kind == cupti.CUPTI_ACTIVITY_KIND_PC_SAMPLING_RECORD_INFO:
+          info = ctypes.cast(record, ctypes.POINTER(cupti.CUpti_ActivityPCSamplingRecordInfo)).contents
+          cid = info.correlationId
+          if cid in self.kernel_stalls:
+            stalls = self.kernel_stalls[cid]
+            total = sum(stalls.values())
+            if total > 0:
+              top = sorted(stalls.items(), key=lambda x: -x[1])[:5]
+              stall_str = " ".join(f"{stall_reason_name(r)}:{100*c//total}%" for r,c in top if c > 0)
+              print(f"  CUPTI stalls (corr={cid}): {total} samples | {stall_str}")
+            del self.kernel_stalls[cid]
+        else: print(f"  CUPTI: Unhandled activity kind {kind}")
+
+  def flush(self):
+    if not self.initialized: return
+    self._check_cupti(cupti.cuptiActivityFlushAll(0))
+
+# Module-level profiler instance
+_profiler: CUPTIProfiler | None = None
+
+def get_profiler() -> CUPTIProfiler | None:
+  return _profiler
+
+def get_cupti_raw_buffers() -> list[bytes]:
+  return _profiler.raw_buffers if _profiler else []
+
+def clear_cupti_raw_buffers():
+  if _profiler: _profiler.raw_buffers.clear()
+
+def get_cupti_pc_samples() -> list[dict]:
+  return _profiler.pc_samples if _profiler else []
+
+def clear_cupti_pc_samples():
+  if _profiler: _profiler.pc_samples.clear()
+
+# Raw PMA buffer access (from ioctl interception)
+def get_pma_raw_dumps() -> list[bytes]:
+  try:
+    from extra.nv_gpu_driver.nv_ioctl import get_pma_raw_dumps as _get
+    return _get()
+  except ImportError: return []
+
+def clear_pma_raw_dumps():
+  try:
+    from extra.nv_gpu_driver.nv_ioctl import clear_pma_raw_dumps as _clear
+    _clear()
+  except ImportError: pass
+
+def enable(profile_level:int=2):
+  global _profiler
+  if _profiler is not None: return
+
+  _profiler = CUPTIProfiler()
+
+  # Patch CUDADevice to initialize CUPTI profiler
+  from tinygrad.runtime.ops_cuda import CUDADevice
+  _orig_init = CUDADevice.__init__
+  _orig_sync = CUDADevice.synchronize
+
+  def _patched_init(self, device: str):
+    _orig_init(self, device)
+    device_id = int(device.split(":")[1]) if ":" in device else 0
+    _profiler.init(self.context, device_id, profile_level)
+
+  def _patched_sync(self):
+    _orig_sync(self)
+    if _profiler: _profiler.flush()
+
+  CUDADevice.__init__ = _patched_init
+  CUDADevice.synchronize = _patched_sync
+
+def enable_auto():
+  if (profile_level:=getenv("PROFILE", 0)) > 0: enable(profile_level)
@@ -0,0 +1,193 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+import enum, collections
+from typing import Iterator
+from tinygrad.helpers import colored
+from tinygrad.renderer.amd.sqtt import PacketType, bits
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# STALL REASONS
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class StallReason(enum.IntEnum):
+  # Based on CUpti_ActivityPCSamplingStallReason
+  INVALID = 0
+  NONE = 1              # selected, selected_not_issued
+  INST_FETCH = 2        # branch_resolving, no_instructions
+  EXEC_DEPENDENCY = 3   # short_scoreboard, wait
+  MEMORY_DEPENDENCY = 4 # long_scoreboard
+  TEXTURE = 5           # tex_throttle
+  SYNC = 6              # barrier, membar
+  CONSTANT_MEMORY = 7   # imc_miss
+  PIPE_BUSY = 8         # mio_throttle, math_pipe_throttle
+  MEMORY_THROTTLE = 9   # drain, lg_throttle
+  NOT_SELECTED = 10     # not_selected
+  OTHER = 11            # misc, dispatch_stall
+  SLEEPING = 12         # sleeping
+
+STALL_KEY_MAP_AMPERE: dict[int, StallReason] = {
+  1: StallReason.MEMORY_THROTTLE, 15: StallReason.MEMORY_THROTTLE,
+  2: StallReason.CONSTANT_MEMORY,
+  3: StallReason.SYNC,
+  6: StallReason.INST_FETCH, 11: StallReason.INST_FETCH,
+  7: StallReason.EXEC_DEPENDENCY, 10: StallReason.EXEC_DEPENDENCY,
+  9: StallReason.MEMORY_DEPENDENCY,
+  12: StallReason.PIPE_BUSY,
+  17: StallReason.OTHER, 20: StallReason.OTHER,
+  18: StallReason.NONE,
+}
+
+STALL_KEY_MAP_BLACKWELL: dict[int, StallReason] = {
+  0x01: StallReason.MEMORY_THROTTLE, 0x0e: StallReason.MEMORY_THROTTLE,
+  0x02: StallReason.SYNC,
+  0x05: StallReason.INST_FETCH, 0x0a: StallReason.INST_FETCH,
+  0x06: StallReason.EXEC_DEPENDENCY, 0x09: StallReason.EXEC_DEPENDENCY,
+  0x08: StallReason.MEMORY_DEPENDENCY,
+  0x0b: StallReason.PIPE_BUSY, 0x0f: StallReason.PIPE_BUSY,
+  0x10: StallReason.OTHER, 0x13: StallReason.OTHER,
+  0x11: StallReason.NONE,
+}
+
+# Lookup table for extracting sample bytes from 32-byte packet (bytes 0-3, 8-31, skipping header at 4-7)
+LOOKUP_28B = [0, 1, 2, 3, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31]
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# PACKET HEADER
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class PMAHeader(PacketType):
+  num_bytes  = bits[4:0]    # number of sample bytes in this packet
+  tpc_id_lo  = bits[15:8]   # TPC identifier low 8 bits
+  tpc_id_hi  = bits[27:25]  # TPC identifier high 3 bits
+  dropped    = bits[28:28]  # dropped flag (resets byte accumulator)
+  @property
+  def tpc_id(self) -> int: return self.tpc_id_lo | (self.tpc_id_hi << 8)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# 8-BYTE SAMPLE FORMAT (Ampere/Ada/Hopper)
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class PMASampleAmpere8B(PacketType):
+  pc_raw     = bits[44:0]   # raw PC value (pc_offset = pc_raw << 4)
+  stall_key  = bits[49:45]  # stall reason key
+  wave_id    = bits[55:50]  # warp/wave identifier
+  active     = bits[62:62]  # 1 if warp was executing, 0 if scheduled but not issued
+  @property
+  def pc_offset(self) -> int: return self.pc_raw << 4
+  @property
+  def stall_reason(self) -> StallReason: return STALL_KEY_MAP_AMPERE.get(self.stall_key, StallReason.OTHER)
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# 9-BYTE SAMPLE FORMAT (Blackwell+)
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class PMASampleBlackwell9B(PacketType):
+  stall_key  = bits[5:0]    # stall reason key
+  pc_raw     = bits[60:8]   # raw PC value (pc_offset = pc_raw << 4)
+  wave_hi    = bits[7:6]    # wave_id high 2 bits
+  wave_lo    = bits[71:68]  # wave_id low 4 bits
+  active     = bits[67:67]  # 1 if warp was executing, 0 if scheduled but not issued
+  @property
+  def pc_offset(self) -> int: return self.pc_raw << 4
+  @property
+  def stall_reason(self) -> StallReason: return STALL_KEY_MAP_BLACKWELL.get(self.stall_key, StallReason.OTHER)
+  @property
+  def wave_id(self) -> int: return (self.wave_hi << 4) | self.wave_lo
+
+PMASample = PMASampleAmpere8B|PMASampleBlackwell9B
+
+def decode(data: bytes, sm_version: int = 0x800) -> Iterator[tuple[PMASample, int]]:
+  use_9byte = sm_version >= 0xa04
+  record_size = 9 if use_9byte else 8
+  sample_cls = PMASampleBlackwell9B if use_9byte else PMASampleAmpere8B
+
+  tpc_state: dict[int, list[int]] = collections.defaultdict(list)
+  for pkt_idx in range(len(data) // 32):
+    pkt = data[pkt_idx * 32:(pkt_idx + 1) * 32]
+    hdr = PMAHeader.from_raw(int.from_bytes(pkt[4:8], 'little'))
+
+    if hdr.dropped: tpc_state[hdr.tpc_id].clear()
+
+    for i in range(hdr.num_bytes):
+      tpc_state[hdr.tpc_id].append(pkt[LOOKUP_28B[i]])
+
+    while len(tpc_state[hdr.tpc_id]) >= record_size:
+      yield sample_cls.from_raw(int.from_bytes(bytes(tpc_state[hdr.tpc_id][:record_size]), 'little')), hdr.tpc_id
+      del tpc_state[hdr.tpc_id][:record_size]
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# CLI
+# ═══════════════════════════════════════════════════════════════════════════════
+
+STALL_COLORS = {
+  StallReason.NONE: "green", StallReason.INST_FETCH: "yellow", StallReason.EXEC_DEPENDENCY: "cyan",
+  StallReason.MEMORY_DEPENDENCY: "red", StallReason.SYNC: "magenta", StallReason.CONSTANT_MEMORY: "blue",
+  StallReason.PIPE_BUSY: "yellow", StallReason.MEMORY_THROTTLE: "RED", StallReason.OTHER: "white",
+}
+
+def decode_tpc_id(tpc_id:int) -> tuple[int, int, int]:
+  # NOTE: valid only for ops_nv, cuda encoding is different
+  return (tpc_id >> 5, (tpc_id >> 1) & 0xf, tpc_id & 1)
+
+def print_samples(samples:list[tuple[PMASample, int]]) -> None:
+  if not samples: return
+  base_pc = min(s.pc_offset for s, _ in samples)
+  for s, tpc_id in samples:
+    gpc, tpc, sm = decode_tpc_id(tpc_id)
+    stall_str = colored(f"{s.stall_reason.name:17}", STALL_COLORS.get(s.stall_reason, "white"))
+    print(f"pc=0x{s.pc_offset - base_pc:06x} {stall_str} ev={s.stall_key:2d} active={s.active} wave={s.wave_id:2d} gpc={gpc} tpc={tpc} sm={sm}")
+
+def print_packets(data:bytes, sm_version:int=0x800) -> None:
+  record_size = 9 if sm_version >= 0x890 else 8
+  tpc_state: dict[int, list[int]] = collections.defaultdict(list)
+  for i in range(len(data) // 32):
+    pkt = data[i * 32:(i + 1) * 32]
+    hdr = PMAHeader.from_raw(int.from_bytes(pkt[4:8], 'little'))
+    if hdr.dropped: tpc_state[hdr.tpc_id].clear()
+    for j in range(hdr.num_bytes): tpc_state[hdr.tpc_id].append(pkt[LOOKUP_28B[j]])
+    # Show complete records extracted from this packet
+    records = []
+    while len(tpc_state[hdr.tpc_id]) >= record_size:
+      records.append(bytes(tpc_state[hdr.tpc_id][:record_size]).hex())
+      del tpc_state[hdr.tpc_id][:record_size]
+    leftover = len(tpc_state[hdr.tpc_id])
+    print(f"Pkt {i:3d}: tpc={hdr.tpc_id:4d} n={hdr.num_bytes:2d} drop={hdr.dropped} left={leftover} | {' '.join(records)}")
+
+def print_aggregated(samples:list[tuple[PMASample, int]]) -> None:
+  if not samples: return
+  base_pc = min(s.pc_offset for s, _ in samples)
+  counter: collections.Counter[tuple[int, StallReason]] = collections.Counter((s.pc_offset, s.stall_reason) for s, _ in samples)
+  print(f"\nAggregated samples (base_pc=0x{base_pc:x}):")
+  for (pc, reason), cnt in sorted(counter.items()):
+    stall_str = colored(f"{reason.name:17}", STALL_COLORS.get(reason, "white"))
+    print(f"  pc=0x{pc - base_pc:06x} {stall_str} samples={cnt:4d}")
+
+if __name__ == "__main__":
+  import sys, pickle
+
+  if len(sys.argv) < 2:
+    print("Usage: python decode.py <pkl_file> [--raw] [--sm=0xNNN]")
+    sys.exit(1)
+
+  with open(sys.argv[1], "rb") as f:
+    data = pickle.load(f)
+
+  if isinstance(data, dict):
+    sm_version = 0x800  # default to Ampere
+    for arg in sys.argv:
+      if arg.startswith("--sm="): sm_version = int(arg[5:], 0)
+    dumps = [(i, x, sm_version) for i, x in enumerate(data["pma_raw_dumps"])]
+  else:
+    devs = {e.device: e for e in data if type(e).__name__ == "ProfileDeviceEvent"}
+    dumps = []
+    for i, e in enumerate(e for e in data if type(e).__name__ == "ProfilePMAEvent"):
+      dumps.append((i, e.blob, devs[e.device].props.get('sm_version', 0x800)))
+
+  for dump_idx, raw, sm_ver in dumps:
+    print(f"\n{'='*60}\nDump {dump_idx} ({len(raw)} bytes, {len(raw)//32} packets)\n{'='*60}")
+    if "--raw" in sys.argv: print_packets(raw, sm_ver)
+    else:
+      samples = list(decode(raw, sm_ver))
+      print(f"\nDecoded {len(samples)} samples:")
+      print_samples(samples)
+      print_aggregated(samples)
@@ -0,0 +1,76 @@
+import pickle, unittest
+from collections import Counter
+from pathlib import Path
+
+from extra.nv_pma.decode import decode
+from tinygrad.helpers import DEBUG
+
+EXAMPLES_DIR = Path(__file__).parent.parent / "examples"
+EXAMPLES_5090_DIR = Path(__file__).parent.parent / "examples_5090"
+
+def decode_and_aggregate(raw_dumps: list[bytes], sm_version: int = 0x800) -> Counter[tuple[int, int]]:
+  """Decode all PMA buffers and aggregate by (relative_pc, stall_reason). Each dump is normalized separately."""
+  result: Counter[tuple[int, int]] = Counter()
+  for raw in raw_dumps:
+    samples = [s for s, _ in decode(raw, sm_version)]
+    if not samples: continue
+    base_pc = min(s.pc_offset for s in samples)
+    result += Counter((s.pc_offset - base_pc, int(s.stall_reason)) for s in samples)
+  return result
+
+def cupti_to_counter(cupti_records: list[dict]) -> Counter[tuple[int, int]]:
+  """Convert CUPTI records to Counter[(pcOffset, stallReason)]."""
+  counter: Counter[tuple[int, int]] = Counter()
+  for r in cupti_records:
+    counter[(r['pcOffset'], r['stallReason'])] += r['samples']
+  return counter
+
+class TestNVProf(unittest.TestCase):
+  def _test_example(self, name: str, sm_version: int = 0x800, examples_dir: Path = EXAMPLES_DIR):
+    pkl_file = examples_dir / f"{name}.pkl"
+    if not pkl_file.exists():
+      self.skipTest(f"Example data not found: {pkl_file}. Run collect.py first.")
+
+    with open(pkl_file, "rb") as f:
+      data = pickle.load(f)
+
+    self.assertEqual(data["test_name"], name)
+    pma_agg = decode_and_aggregate(data["pma_raw_dumps"], sm_version)
+    cupti_agg = cupti_to_counter(data["cupti_pc_samples"])
+
+    if DEBUG >= 2:
+      total = sum(cupti_agg.values())
+      mismatched = sum(abs(pma_agg.get(k, 0) - v) for k, v in cupti_agg.items())
+      mismatched += sum(v for k, v in pma_agg.items() if k not in cupti_agg)
+      mismatched //= 2
+
+      print(f"\n=== Test: {name} ===")
+      print(f"Total samples: {total}, Mismatched: {mismatched} ({mismatched/total*100 if total else 0:.1f}%)")
+
+    self.assertEqual(pma_agg, cupti_agg, f"PMA: {dict(pma_agg)}\nCUPTI: {dict(cupti_agg)}")
+
+  # Ampere tests (8-byte format)
+  def test_decode_test_plus(self): self._test_example("test_plus")
+  def test_decode_test_reduce_sum(self): self._test_example("test_reduce_sum")
+  def test_decode_test_broadcast(self): self._test_example("test_broadcast")
+  def test_decode_test_matmul(self): self._test_example("test_matmul")
+  def test_decode_test_plus_big(self): self._test_example("test_plus_big")
+  def test_decode_test_elementwise_chain(self): self._test_example("test_elementwise_chain")
+  def test_decode_test_conv2d(self): self._test_example("test_conv2d")
+  def test_decode_test_large_matmul(self): self._test_example("test_large_matmul")
+
+  # Blackwell/5090 tests (9-byte format)
+  def test_5090_test_plus(self): self._test_example("test_plus", 0xa04, EXAMPLES_5090_DIR)
+  def test_5090_test_plus_big(self): self._test_example("test_plus_big", 0xa04, EXAMPLES_5090_DIR)
+  def test_5090_test_broadcast(self): self._test_example("test_broadcast", 0xa04, EXAMPLES_5090_DIR)
+  def test_5090_test_matmul(self): self._test_example("test_matmul", 0xa04, EXAMPLES_5090_DIR)
+  def test_5090_test_large_matmul(self): self._test_example("test_large_matmul", 0xa04, EXAMPLES_5090_DIR)
+  def test_5090_test_reduce_sum(self): self._test_example("test_reduce_sum", 0xa04, EXAMPLES_5090_DIR)
+  def test_5090_test_reduce_max(self): self._test_example("test_reduce_max", 0xa04, EXAMPLES_5090_DIR)
+  def test_5090_test_elementwise_chain(self): self._test_example("test_elementwise_chain", 0xa04, EXAMPLES_5090_DIR)
+  def test_5090_test_conv2d(self): self._test_example("test_conv2d", 0xa04, EXAMPLES_5090_DIR)
+  def test_5090_test_exp(self): self._test_example("test_exp", 0xa04, EXAMPLES_5090_DIR)
+  def test_5090_test_softmax(self): self._test_example("test_softmax", 0xa04, EXAMPLES_5090_DIR)
+
+if __name__ == "__main__":
+  unittest.main()
@@ -1,12 +1,13 @@
 #!/bin/bash
+set -e
 export PAGE_SIZE=1
 export PYTHONPATH=.
 export LOGOPS=/tmp/ops
 export CAPTURE_PROCESS_REPLAY=1
-rm $LOGOPS
+rm "$LOGOPS" 2>/dev/null || true
 test/external/process_replay/reset.py

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

 # extract, sort and uniq
@@ -1,31 +0,0 @@
-import random
-from extra.optimization.helpers import load_worlds, ast_str_to_lin
-from tinygrad.codegen.opt.search import actions
-from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
-from tinygrad.helpers import tqdm
-
-tactions = set()
-def test_rebuild(lin):
-  linr = Kernel(lin.ast)
-  for o in lin.applied_opts:
-    assert o in actions, f"{o} is not in actions"
-    tactions.add(o)
-    linr.apply_opt(o)
-
-  assert len(lin.sts) == len(linr.sts)
-  for st1,st2 in zip(lin.sts, linr.sts):
-    assert st1 == st2, f"{st1} != {st2}"
-
-if __name__ == "__main__":
-  ast_strs = load_worlds(False, False, False)
-  random.shuffle(ast_strs)
-  ast_strs = ast_strs[:2000]
-  for ast_str in tqdm(ast_strs):
-    lin = ast_str_to_lin(ast_str)
-    #if not lin.apply_tensor_cores():
-    lin.apply_opts(hand_coded_optimizations(lin))
-    test_rebuild(lin)
-
-  print(len(tactions), len(actions))
-  print(sorted(list(tactions)))
@@ -1,76 +0,0 @@
-import os
-import numpy as np
-import math, random
-from tinygrad.tensor import Tensor
-from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
-from tinygrad.codegen.opt.search import actions, bufs_from_lin, get_kernel_actions
-from tinygrad.nn.optim import Adam
-from extra.optimization.extract_policynet import PolicyNet
-from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, time_linearizer
-
-if __name__ == "__main__":
-  net = PolicyNet()
-  if os.path.isfile("/tmp/policynet.safetensors"): load_state_dict(net, safe_load("/tmp/policynet.safetensors"))
-  optim = Adam(get_parameters(net))
-
-  ast_strs = load_worlds()
-
-  # select a world
-  all_feats, all_acts, all_rews = [], [], []
-  while 1:
-    Tensor.training = False
-    lin = ast_str_to_lin(random.choice(ast_strs))
-    rawbufs = bufs_from_lin(lin)
-    tm = last_tm = base_tm = time_linearizer(lin, rawbufs)
-
-    # take actions
-    feats, acts, rews = [], [], []
-    while 1:
-      feat = lin_to_feats(lin)
-      feats.append(feat)
-      probs = net(Tensor([feat])).exp()[0].numpy()
-
-      # mask valid actions
-      valid_action_mask = np.zeros((len(actions)+1), dtype=np.float32)
-      for x in get_kernel_actions(lin): valid_action_mask[x] = 1
-      probs *= valid_action_mask
-      probs /= sum(probs)
-
-      act = np.random.choice(len(probs), p=probs)
-      acts.append(act)
-      if act == 0:
-        rews.append(0)
-        break
-      try:
-        lin.apply_opt(actions[act-1])
-        tm = time_linearizer(lin, rawbufs)
-        if math.isinf(tm): raise Exception("failed")
-        rews.append(((last_tm-tm)/base_tm))
-        last_tm = tm
-      except Exception:
-        rews.append(-0.5)
-        break
-      #print(f"{tm*1e6:10.2f}", lin.colored_shape())
-
-    assert len(feats) == len(acts) and len(acts) == len(rews)
-    #print(rews)
-    print(f"***** EPISODE {len(rews)} steps, {sum(rews):5.2f} reward, {base_tm*1e6:12.2f} -> {tm*1e6:12.2f} : {lin.colored_shape()}")
-    all_feats += feats
-    all_acts += acts
-    # rewards to go
-    for i in range(len(rews)-2, -1, -1): rews[i] += rews[i+1]
-    all_rews += rews
-
-    BS = 32
-    if len(all_feats) >= BS:
-      Tensor.training = True
-      x = Tensor(all_feats[:BS])
-      mask = np.zeros((BS, len(actions)+1), dtype=np.float32)
-      mask[range(BS), all_acts[:BS]] = all_rews[:BS]
-      loss = -(net(x) * Tensor(mask)).mean()
-      optim.zero_grad()
-      loss.backward()
-      optim.step()
-      all_feats = all_feats[BS:]
-      all_acts = all_acts[BS:]
-      all_rews = all_rews[BS:]
@@ -1,32 +0,0 @@
-from typing import List, Tuple
-from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.codegen.opt.search import get_kernel_actions, actions
-
-_net = None
-def beam_q_estimate(beam:List[Tuple[Kernel, float]]) -> List[Tuple[Kernel, float]]:
-  global _net
-  if _net is None:
-    from tinygrad.nn.state import load_state_dict, safe_load
-    from extra.optimization.pretrain_valuenet import ValueNet
-    _net = ValueNet(1021+len(actions), 2)
-    load_state_dict(_net, safe_load("/tmp/qnet.safetensors"), verbose=False)
-  from tinygrad.tensor import Tensor
-  from tinygrad.helpers import Context
-  from extra.optimization.helpers import lin_to_feats
-  import numpy as np
-  feats = []
-  lins = []
-  base_tms = []
-  for lin,tm in beam:
-    lin_feats = lin_to_feats(lin)
-    for a,v in get_kernel_actions(lin, include_0=False).items():
-      acts = np.zeros(len(actions))
-      acts[a-1] = 1.0
-      feats.append(np.concatenate([lin_feats, acts]))
-      lins.append(v)
-      base_tms.append(tm)
-  with Context(BEAM=0):
-    with Tensor.train(False):
-      preds = _net(Tensor(feats)).numpy()
-  pred_time = np.array(base_tms) / np.exp(preds[:, 0])
-  return sorted(zip(lins, pred_time), key=lambda x: x[1])
@@ -1,34 +0,0 @@
-import argparse
-from extra.optimization.helpers import ast_str_to_lin, time_linearizer
-
-from tinygrad import dtypes
-from tinygrad.helpers import BEAM, getenv
-from tinygrad.device import Device, Compiled
-from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
-
-
-if __name__ == '__main__':
-  parser = argparse.ArgumentParser(description="Run a search for the optimal opts for a kernel", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-  parser.add_argument("--ast", type=str, default=None, help="the ast for the kernel to be optimized")
-  parser.add_argument("--file", type=str, default=None, help="a file containing asts to be optimized, one per line")
-  args = parser.parse_args()
-
-  device: Compiled = Device[Device.DEFAULT]
-  print(f"optimizing for {Device.DEFAULT}")
-
-  if args.ast is not None:
-    ast_strs = [args.ast]
-  elif args.file is not None:
-    with open(args.file, 'r') as file:
-      ast_strs = file.readlines()
-
-  for i, ast_str in enumerate(ast_strs):
-    print(f"optimizing {i}/{len(ast_strs)}\nast={ast_str}")
-    lin = ast_str_to_lin(ast_str, opts=device.renderer)
-    rawbufs = bufs_from_lin(lin)
-    lin = beam_search(lin, rawbufs, getenv("BEAM", 8), bool(getenv("BEAM_ESTIMATE", 1)))
-
-    tm = time_linearizer(lin, rawbufs, allow_test_size=False, cnt=10)
-    print(f"final time {tm*1e6:9.0f} us: {lin.colored_shape()}")
-    print(lin.applied_opts)
--- a/Show More
+++ b/Show More