add ninja to setup-tinygrad

ci: cleanup deps
2026-08-18 15:53:24 -07:00 · 2026-08-18 15:48:48 -07:00
231 changed files with 4833 additions and 11662 deletions
@@ -4,7 +4,7 @@ inputs:
  python-version:
    description: 'Python version to use'
    required: false
-    default: '3.14'
+    default: '' # if you don't set a version, the native python version will be used
  key:
    description: 'Key for the python cache'
    required: false
@@ -59,18 +59,18 @@ runs:
        echo "OMP_NUM_THREADS=1" >> "$GITHUB_ENV"
        # no buffers should be over 300MB in CI
        echo "MAX_BUFFER_SIZE=300000000" >> "$GITHUB_ENV"
-        if [[ "$RUNNER_OS" == "Linux" ]]; then
-          echo "VIRTUAL_ENV=/opt/venv/${{ inputs.python-version }}" >> "$GITHUB_ENV"
-          echo "UV_PYTHON_INSTALL_DIR=/opt/python" >> "$GITHUB_ENV"
-        else
-          echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
-        fi

    - name: Set up uv
      uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b
      with:
        enable-cache: 'false' # see below for manual caching

+    - name: Set up Python ${{ inputs.python-version }}
+      uses: actions/setup-python@v6
+      if: inputs.python-version != ''
+      with:
+        python-version: ${{ inputs.python-version }}
+
    # **** Caching packages ****

    - name: Cache Python packages (PR)
@@ -109,15 +109,15 @@ runs:
      if: inputs.deps != ''
      shell: bash
      run: |
-        uv venv --allow-existing --python ${{ inputs.python-version }} "$VIRTUAL_ENV"
+        uv venv .venv
        DEPS="${{ inputs.deps }}"
-        uv pip install --python "$VIRTUAL_ENV" -e ".[${DEPS// /,}]" ${{ inputs.pydeps }} --torch-backend cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
+        uv pip install --python .venv -e ".[${DEPS// /,}]" ${{ inputs.pydeps }} --torch-backend 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 == ''
      shell: bash
      run: |
-        uv venv --allow-existing --python ${{ inputs.python-version }} "$VIRTUAL_ENV"
-        uv pip install --python "$VIRTUAL_ENV" -e . ${{ inputs.pydeps }}
+        uv venv .venv
+        uv pip install --python .venv -e . ${{ inputs.pydeps }}
    - name: Prune uv cache
      if: github.event_name != 'pull_request'
      shell: bash
@@ -125,10 +125,11 @@ runs:
    - name: Configure venv
      shell: bash
      run: |
+        echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
        if [[ "$RUNNER_OS" == "Windows" ]]; then
-          echo "$VIRTUAL_ENV/Scripts" >> "$GITHUB_PATH"
+          echo "${{ github.workspace }}/.venv/Scripts" >> "$GITHUB_PATH"
        else
-          echo "$VIRTUAL_ENV/bin" >> "$GITHUB_PATH"
+          echo "${{ github.workspace }}/.venv/bin" >> "$GITHUB_PATH"
        fi

    # ******************* apt *******************
@@ -194,29 +195,21 @@ runs:
        echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
        echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"

-        installed=true
-        for pkg in $pkgs; do
-          info=$(dpkg-query -W -f='${db:Status-Abbrev} ${Version}' "$pkg" 2> /dev/null || true)
-          echo "${pkg}: ${info:-not in dpkg database}"
-          [[ "$info" == ii* ]] || installed=false
-        done
-        echo "installed=$installed" >> "$GITHUB_OUTPUT"
-
    - name: Cache apt (PR)
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name == 'pull_request' && steps.apt-pkgs.outputs.installed == 'false'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name == 'pull_request'
      uses: actions/cache/restore@v5
      with:
        path: /var/cache/apt/archives/
        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
    - name: Cache apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name != 'pull_request' && steps.apt-pkgs.outputs.installed == 'false'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name != 'pull_request'
      uses: actions/cache@v5
      with:
        path: /var/cache/apt/archives/
        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}

    - name: Run apt Update + Install
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && steps.apt-pkgs.outputs.installed == 'false'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
      shell: bash
      run: |
        sudo apt -qq update || true
@@ -248,10 +241,10 @@ runs:
      if: inputs.amd == 'true' && runner.os == 'macOS'
      shell: bash
      run: |
-        sudo "$VIRTUAL_ENV/bin/python" -c "
-          from tinygrad.helpers import fetch
-          fetch('https://github.com/tinygrad/amdcomgr_dylib/releases/download/v7.2.0/libamd_comgr.dylib', name='/usr/local/lib/libamd_comgr.dylib',
-                sha256='7712fbe4fcb9fcdea49aeac989876448df975ce0a8ce7c9b15b55c15e7a05935').chmod(0o644)"
+        sudo mkdir -p /usr/local/lib
+        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/tinygrad/amdcomgr_dylib/releases/latest | \
+          jq -r '.assets[] | select(.name == "libamd_comgr.dylib").browser_download_url' | \
+          sudo xargs curl -fL -o /usr/local/lib/libamd_comgr.dylib

    # **** CUDA ****
    - name: Install CUDA
@@ -269,11 +262,8 @@ runs:
      if: inputs.ocelot == 'true'
      shell: bash
      run: |
-        sudo "$VIRTUAL_ENV/bin/python" -c "
-          from tinygrad.helpers import fetch
-          fetch('https://github.com/tinygrad/gpuocelot/releases/download/v0.1.0/libgpuocelot.${{ runner.os == 'Linux' && 'so' || 'dylib' }}',
-                name='/usr/local/lib/libgpuocelot.${{ runner.os == 'Linux' && 'so' || 'dylib' }}',
-                sha256='${{ runner.os == 'Linux' && 'a24705276a9a187111371465987b3258f8836ef512a34266e3075bc4714e125a' || '5106c998c795a36dec79eb7b2aae324a93d1338236d36eeaae232649ec457663' }}').chmod(0o644)"
+        sudo mkdir -p /usr/local/lib
+        sudo curl --output-dir /usr/local/lib -fLO https://github.com/tinygrad/gpuocelot/releases/download/v0.1.0/libgpuocelot.${{ runner.os == 'Linux' && 'so' || 'dylib' }}

    # **** WebGPU ****

@@ -281,11 +271,8 @@ runs:
      if: inputs.webgpu == 'true'
      shell: bash
      run: |
-        sudo "$VIRTUAL_ENV/bin/python" -c "
-          from tinygrad.helpers import fetch
-          fetch('https://github.com/wpmed92/pydawn/releases/download/v0.1.6/libwebgpu_dawn.${{ runner.os == 'Linux' && 'so' || 'dylib' }}',
-                name='/usr/local/lib/libwebgpu_dawn.${{ runner.os == 'Linux' && 'so' || 'dylib' }}',
-                sha256='${{ runner.os == 'Linux' && 'cf36091d266a32c9d5080f14662de44cece241987939713282ea0ff558db81c6' || '7e87c7acefda8b6af1a1c5debfedcf62958311284b8fd8d9bcf93e312e6636e3' }}').chmod(0o644)"
+        sudo mkdir -p /usr/local/lib
+        sudo curl --output-dir /usr/local/lib -fLO https://github.com/wpmed92/pydawn/releases/download/v0.1.6/libwebgpu_dawn.${{ runner.os == 'Linux' && 'so' || 'dylib' }}

    # **** LLVM ****

@@ -299,10 +286,7 @@ runs:
      if: inputs.opencl == 'true'
      shell: bash
      run: |
-        sudo "$VIRTUAL_ENV/bin/python" -c "
-          from tinygrad.helpers import fetch
-          fetch('https://github.com/sirhcm/tinymesa/releases/download/rusticl-v1/libRusticlOpenCL.so.1.0.0', name='/usr/lib/libRusticlOpenCL.so',
-                sha256='d4f48566d8fd33f6cdd8ef6de35a71966e8a8517e6f68ff3c52dbb43765a2513').chmod(0o644)"
+        sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/rusticl-v1/libRusticlOpenCL.so.1.0.0 -o /usr/lib/libRusticlOpenCL.so
        sudo mkdir -p /etc/OpenCL/vendors
        echo "/usr/lib/libRusticlOpenCL.so" | sudo tee /etc/OpenCL/vendors/rusticl.icd
        echo "RUSTICL_ENABLE=llvmpipe" >> "$GITHUB_ENV"
@@ -40,10 +40,10 @@ jobs:
      run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev liburing-dev
    - name: Regenerate autogen files
      run: |
-        find tinygrad/runtime/autogen -type f -name "*.py" -not -path "*/amd/*" -not -name "__init__.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete
+        find tinygrad/runtime/autogen -type f -name "*.py" -not -path "*/amd/*" -not -name "__init__.py" -not -name "comgr.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"
        python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv_610, nv"
-        python3 -c "from tinygrad.runtime.autogen import comgr, comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
+        python3 -c "from tinygrad.runtime.autogen import comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
        python3 -c "from tinygrad.runtime.autogen.am import *"
        python3 -c "from tinygrad.runtime.autogen.nv_regs import *"
        python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, pci, vfio"
@@ -54,7 +54,7 @@ jobs:
        python3 -c "from tinygrad.runtime.autogen import mesa"
        python3 -c "from tinygrad.runtime.autogen import avcodec"
        python3 -c "from tinygrad.runtime.autogen import llvm_qcom"
-        python3 -c "from tinygrad.runtime.autogen import mlx5, bnxt"
+        python3 -c "from tinygrad.runtime.autogen import mlx5"
        python3 -c "from tinygrad.runtime.autogen import ggml_common"
        REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
    - name: Check for differences
@@ -102,3 +102,42 @@ jobs:
      with:
        name: autogen-macos-patch
        path: autogen-macos.patch
+
+  autogen-comgr-2:
+    name: In-tree Autogen (comgr 2)
+    runs-on: ubuntu-24.04
+    timeout-minutes: 15
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: 'autogen-comgr'
+    - name: Install autogen support packages
+      run: |
+        wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
+        sudo tee /etc/apt/sources.list.d/rocm.list <<EOF
+        deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.2 $(lsb_release -cs) main
+        EOF
+        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: Regenerate autogen files
+      run: |
+        rm tinygrad/runtime/autogen/comgr.py
+        python3 -c "from tinygrad.runtime.autogen import comgr"
+    - name: Check for differences
+      run: |
+        if ! git diff --quiet; then
+          git diff
+          git diff > autogen-comgr2.patch
+          echo "Autogen mismatch detected. Patch available at: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
+          exit 1
+        fi
+    - name: Upload patch artifact
+      if: failure()
+      uses: actions/upload-artifact@v7
+      with:
+        name: autogen-comgr2-patch
+        path: autogen-comgr2.patch
@@ -88,13 +88,13 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 30
+    timeout-minutes: 60
    defaults:
      run:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -102,11 +102,16 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py amd rmmod
+        ./extra/hcq/hcq_smi.py amd kill_pids
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -116,14 +121,18 @@ jobs:
      run: python3 test/external/process_replay/reset.py
    - name: Run llama3.2
      run: BENCHMARK_LOG=llama32_3b-f16 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m llama3.2:3b-f16 --benchmark --warmup
-    - name: Run qwen3.8
-      # qwen3.8:27b doesn't fit on mac
+    - name: Run qwen3.6
+      # qwen3.6:35b-a3b doesn't fit on mac
      if: ${{ matrix.dev != 'METAL' }}
-      run: BENCHMARK_LOG=qwen38_27b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.8:27b --benchmark --warmup
+      run: BENCHMARK_LOG=qwen36_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.6:35b-a3b --benchmark --warmup
    - name: Run olmoe
      # just metal for now
      if: ${{ matrix.dev == 'METAL' }}
      run: BENCHMARK_LOG=olmoe JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m olmoe --benchmark --warmup
+    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
+      # only run on machines with multiple gpus
+      if: ${{ matrix.dev != 'METAL' }}
+      run: BENCHMARK_LOG=llama3_beam_4gpu 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 process replay tests
      uses: ./.github/actions/process-replay

@@ -134,13 +143,13 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 10
+    timeout-minutes: 60
    defaults:
      run:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -148,11 +157,12 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py amd rmmod
+        ./extra/hcq/hcq_smi.py amd kill_pids
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -172,6 +182,10 @@ jobs:
      # slow on metal
      if: ${{ matrix.dev != 'METAL' }}
      run: time BENCHMARK_LOG=cifar 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
+      # only run on machines with multiple gpus
+      if: ${{ matrix.dev != 'METAL' }}
+      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -182,13 +196,13 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['AMD', 'NV']
-    timeout-minutes: 5
+    timeout-minutes: 60
    defaults:
      run:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -196,11 +210,12 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py amd rmmod
+        ./extra/hcq/hcq_smi.py amd kill_pids
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
    - name: Symlink models and datasets
      run: |
        mkdir -p extra/datasets
@@ -212,8 +227,15 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: test/external/process_replay/reset.py
+    - name: Run MLPerf resnet eval on training data
+      run: time BENCHMARK_LOG=resnet_eval MODEL=resnet python3 examples/mlperf/model_eval.py
    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
      run: BENCHMARK_LOG=resnet_10steps 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 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 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
      uses: ./.github/actions/process-replay

@@ -224,13 +246,13 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 15
+    timeout-minutes: 60
    defaults:
      run:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -238,11 +260,12 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py amd rmmod
+        ./extra/hcq/hcq_smi.py amd kill_pids
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -262,59 +285,6 @@ jobs:
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

-  multigpubenchmark:
-    name: Multi-GPU Benchmarks (DEV=${{ matrix.dev }})
-    runs-on: [self-hosted, "${{ matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
-    strategy:
-      fail-fast: false
-      matrix:
-        dev: ['AMD', 'NV']
-    timeout-minutes: 20
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    env:
-      DEV: ${{ matrix.dev }}
-      HCQ2: '0'
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup (AMD)
-      if: ${{ matrix.dev == 'AMD' }}
-      run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
-    - name: Setup (NV)
-      if: ${{ matrix.dev == 'NV' }}
-      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
-    - name: Symlink models and datasets
-      run: |
-        mkdir -p weights
-        mkdir -p extra/datasets
-        ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
-        ln -s /raid/datasets/imagenet extra/datasets/imagenet
-    - name: setup staging db
-      if: github.ref == 'refs/heads/update_benchmark_staging'
-      run: |
-        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
-        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: reset process replay
-      run: python3 test/external/process_replay/reset.py
-    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
-      run: BENCHMARK_LOG=llama3_beam_4gpu 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 full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 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 MODEL=resnet python3 examples/mlperf/model_eval.py
-    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: BENCHMARK_LOG=resnet_10steps_6gpu 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 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
-      uses: ./.github/actions/process-replay
-
  tests:
    name: Tests (DEV=${{ matrix.dev }})
    runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
@@ -322,7 +292,7 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 10
+    timeout-minutes: 60
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -335,11 +305,12 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py amd rmmod
+        ./extra/hcq/hcq_smi.py amd kill_pids
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -416,7 +387,7 @@ jobs:
  testusbgpu:
    name: UsbGPU Benchmark
    runs-on: [self-hosted, macOS]
-    timeout-minutes: 3
+    timeout-minutes: 10
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -431,70 +402,32 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: Kill stale pids
      run: |
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
-        ./extra/hcq/hcq_smi.py nv kill_pids --sudoless
+        PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
+        PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
+    # since sudo is required for usbgpu on macos, do not write bytecode, as some of the files are owned by root
    - name: UsbGPU boot time
-      run: GMMU=0 DEBUG=2 AM_RESET=1 DEV=USB+AMD time python3.11 test/test_tiny.py TestTiny.test_plus
+      run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEBUG=2 AM_RESET=1 DEV=USB+AMD time python3.11 test/test_tiny.py TestTiny.test_plus
    - name: UsbGPU tiny tests
-      run: GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
+      run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
    - name: UsbGPU copy speeds
-      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+      run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+    #- name: UsbGPU openpilot test
+    #  run: sudo -E PYTHONPATH=. GMMU=0 DEV=USB+AMD GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
    - name: UsbGPU (USB4/TB) install script
-      run: sh extra/setup_tinygpu_osx.sh
+      run: PYTHONPATH=. sh extra/setup_tinygpu_osx.sh
    - name: UsbGPU (USB4/TB) boot time
-      run: DEBUG=3 DEV=PCI+NV:NAK time python3.11 test/test_tiny.py TestTiny.test_plus
+      run: PYTHONPATH=. DEBUG=3 DEV=PCI+NV:NAK time python3.11 test/test_tiny.py TestTiny.test_plus
    - name: UsbGPU (USB4/TB) tiny tests
-      run: DEV=PCI+NV:NAK python3.11 test/test_tiny.py
+      run: PYTHONPATH=. DEV=PCI+NV:NAK python3.11 test/test_tiny.py

-  testcomma:
-    strategy:
-      matrix:
-        dev: ['QCOM', 'QCOM:IR3']
-        version: ['0.11.0', '0.11.2']
-        model: ['vision', 'policy', 'supercombo', 'dmonitoring']
-        # exclude non-existent models
-        exclude: [{ version: '0.11.0', model: supercombo }, { version: '0.11.2', model: vision }, { version: '0.11.2', model: policy }]
-        include:
-          - version: '0.11.0'
-            model: vision
-            url: https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
-            timing: 17
-          - version: '0.11.0'
-            model: policy
-            url: https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
-            timing: 3.2
-          - version: '0.11.0'
-            model: dmonitoring
-            url: https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
-            timing: 11
-          - version: '0.11.2'
-            model: supercombo
-            url: https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
-            timing: 26
-          - dev: QCOM:IR3
-            version: '0.11.2'
-            model: supercombo
-            timing: 41
-          - version: '0.11.2'
-            model: dmonitoring
-            url: https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/3e7b31dfbc0a5234f1baf196513b77fc6af12204b8a8ffe8ee0417e48352f316
-            timing: 11
-          # IR3 dmonitoring is slightly slower
-          - dev: QCOM:IR3
-            model: dmonitoring
-            timing: 12
-      fail-fast: false
-    name: openpilot ${{ matrix.version }} compile3 ${{ matrix.model }} (DEV=${{ matrix.dev }})
+  testcommalatest:
+    name: comma Benchmark (0.11.2)
    runs-on: [self-hosted, Linux, comma]
-    timeout-minutes: 5
+    timeout-minutes: 12
    defaults:
      run:
        shell: bash -e -o pipefail {0}
    if: github.repository_owner == 'tinygrad'
-    env:
-      DEV: ${{ matrix.dev }}
-      ASSERT_MIN_STEP_TIME: ${{ matrix.timing }}
-      BENCHMARK_LOG: ${{ matrix.dev == 'QCOM:IR3' && 'ir3_' || '' }}openpilot_${{ matrix.version }}_${{ matrix.model }}
    steps:
    - name: Checkout Code
      uses: actions/checkout@v6
@@ -505,10 +438,45 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: test/external/process_replay/reset.py
-    - name: compile
-      run: FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py ${{ matrix.url }}
-    - name: run pickle
-      run: BENCHMARK_LOG="${BENCHMARK_LOG}_run_pickle" RUN_PICKLE=1 taskset -c 4-7 python3 examples/openpilot/compile3.py
+    - name: openpilot compile3 0.11.2 supercombo
+      run: BENCHMARK_LOG=openpilot_0_11_2_supercombo PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
+    - name: openpilot compile3 0.11.2 supercombo (from pickle)
+      run: BENCHMARK_LOG=openpilot_0_11_2_supercombo_run_pickle RUN_PICKLE=1 PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 DEV=QCOM taskset -c 4-7 python3 examples/openpilot/compile3.py
+    - name: IR3 openpilot compile3 0.11.2 supercombo
+      run: BENCHMARK_LOG=ir3_openpilot_0_11_2_supercombo PYTHONPATH="." ASSERT_MIN_STEP_TIME=41 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
+    - name: openpilot compile3 0.11.2 dmonitoring
+      run: BENCHMARK_LOG=openpilot_0_11_2_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/3e7b31dfbc0a5234f1baf196513b77fc6af12204b8a8ffe8ee0417e48352f316
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
+  testcommaold:
+    name: comma Benchmark (0.11.0)
+    runs-on: [self-hosted, Linux, comma]
+    timeout-minutes: 12
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: reset process replay
+      run: test/external/process_replay/reset.py
+    - name: openpilot compile3 0.11.0 driving_vision
+      run: BENCHMARK_LOG=openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot compile3 0.11.0 driving_vision (from pickle)
+      run: BENCHMARK_LOG=openpilot_0_11_0_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM taskset -c 4-7 python3 examples/openpilot/compile3.py
+    - name: IR3 openpilot compile3 0.11.0 driving_vision
+      run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=18 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot compile3 0.11.0 driving_policy
+      run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3.2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
+    - name: openpilot compile3 0.11.0 dmonitoring
+      run: BENCHMARK_LOG=openpilot_0_11_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -521,6 +489,15 @@ jobs:
        shell: bash -e -o pipefail {0}
    if: github.repository_owner == 'tinygrad'
    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: reset process replay
+      run: test/external/process_replay/reset.py
    - name: Checkout Code
      uses: actions/checkout@v6
    - name: setup staging db
@@ -544,7 +521,7 @@ jobs:
  testcommausbgpubenchmark:
    name: UsbGPU Benchmark (comma)
    runs-on: [self-hosted, Linux, comma4]
-    timeout-minutes: 10
+    timeout-minutes: 20
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -564,7 +541,7 @@ jobs:
    - name: openpilot run_pickle big_driving_supercombo
      run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_run_pickle RUN_PICKLE=1 PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py - openpilot.pkl
    - name: Test copy speeds
-      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+      run: SIZE=64e6 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds

  driverbenchmarks:
    name: PCI Driver Benchmark (DEV=${{ matrix.dev }})
@@ -573,7 +550,7 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['AMD', 'NV']
-    timeout-minutes: 5
+    timeout-minutes: 20
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -585,8 +562,9 @@ jobs:
      uses: actions/checkout@v6
    - name: Setup
      run: |
-        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} rmmod --expect
-        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} kill_pids --sudoless
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py ${{ matrix.dev == 'AMD' && 'amd' || 'nv' }} rmmod
+        ./extra/hcq/hcq_smi.py ${{ matrix.dev == 'AMD' && 'amd' || 'nv' }} kill_pids
        mkdir -p extra/datasets
        ln -s /raid/datasets/imagenet extra/datasets/imagenet
    - name: setup staging db
@@ -621,9 +599,6 @@ jobs:
      run: |
        GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
        GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
-    - name: HEVC Decode Benchmark
-      if: ${{ matrix.dev == 'NV' }}
-      run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
      if: ${{ matrix.dev == 'NV' }}
      run: BENCHMARK_LOG=resnet_10steps MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
@@ -646,7 +621,7 @@ jobs:
  llvmspeed:
    name: LLVM Speed
    runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 5
+    timeout-minutes: 20
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -8,7 +8,7 @@ permissions:
  contents: write
 jobs:
  deploy:
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v6
      - name: Configure Git Credentials
@@ -166,7 +166,7 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: windows-${{ matrix.dev }}-minimal
-          deps: testing_minimal
+          deps: testing_unit
          pydeps: ${{ matrix.dev == 'WEBGPU' && 'dawn-python' || '' }}
      - name: Set env
        shell: bash
@@ -10,7 +10,7 @@ on:

 jobs:
  deploy:
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-latest
    steps:
    - uses: actions/checkout@v6
    - name: Set up Python
@@ -10,7 +10,7 @@ concurrency:
 jobs:
  checkbranch:
    name: Check PR Branch status
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-latest
    outputs:
      branchstat: ${{ steps.brstat.outputs.stat}}
    steps:
@@ -44,7 +44,7 @@ jobs:
    permissions:
      contents: read
      pull-requests: write
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-latest
    needs: checkbranch
    if: needs.checkbranch.outputs.branchstat == 'false'
    steps:
@@ -87,7 +87,7 @@ jobs:
    name: Core Library Line Difference
    permissions:
      pull-requests: write
-    runs-on: ubuntu-24.04
+    runs-on: ubuntu-latest
    needs: checkbranch
    if: needs.checkbranch.outputs.branchstat == 'true'
    steps:
@@ -25,7 +25,6 @@ jobs:
    timeout-minutes: 10
    env:
      CHECK_OOB: 0
-      DEV: CPU
    steps:
    - name: Checkout Code
      uses: actions/checkout@v6
@@ -55,7 +54,7 @@ jobs:
              'python docs/abstractions3.py' \
              $'awk \'/```python/{flag=1;next}/```/{flag=0}flag\' README.md | python' \
              $'awk \'/```python/{flag=1;next}/```/{flag=0}flag\' docs/quickstart.md | python' \
-              'python examples/compile_efficientnet.py > recognize.c && clang -O2 recognize.c -lm -o recognize && cat test/models/efficientnet/Chicken.jpg | ./recognize | grep cock'
+              'DEV=CPU python examples/compile_efficientnet.py > recognize.c && clang -O2 recognize.c -lm -o recognize && cat test/models/efficientnet/Chicken.jpg | ./recognize | grep cock'
    - name: Test DEBUG
      run: DEBUG=100 python3 -c "from tinygrad import Tensor; N = 1024; a, b = Tensor.rand(N, N), Tensor.rand(N, N); c = (a.reshape(N, 1, N) * b.T.reshape(1, N, N)).sum(axis=2); print((c.numpy() - (a.numpy() @ b.numpy())).mean())"

@@ -75,9 +74,9 @@ jobs:
        llvm: 'true'
        ninja: 'true'
    - name: Test ResNet-18
-      run: DEV=CPU DEBUG=2 python3 extra/torch_backend/example.py
+      run: DEBUG=2 python3 extra/torch_backend/example.py
    - name: Test one op in torch tests
-      run: DEV=CPU DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
+      run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
    - name: Test Ops with TINY_BACKEND
      run: DEV=CPU:LLVM LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/backend/test_ops.py --durations=20
    - name: Custom tests
@@ -207,8 +206,6 @@ jobs:
    name: Unit Tests
    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
-    env:
-      DEV: CPU

    steps:
    - name: Checkout Code
@@ -227,8 +224,8 @@ jobs:
      run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
    - name: Run unit tests
      run: |
-        python test/null/test_device.py TestRunAsModule.test_module_runs
-        python -m pytest -n=auto test/unit/ --durations=20
+        DEV=CPU python test/null/test_device.py TestRunAsModule.test_module_runs
+        DEV=CPU python -m pytest -n=auto test/unit/ --durations=20
    - name: Run GC tests
      run: python test/external/external_uop_gc.py
    - name: External Benchmark Schedule
@@ -236,7 +233,7 @@ jobs:
    - name: Run process replay tests
      uses: ./.github/actions/process-replay
    - name: Repo line count <= 26000 lines
-      run: MAX_LINE_COUNT=26500 python sz.py
+      run: MAX_LINE_COUNT=26000 python sz.py

  spec:
    strategy:
@@ -256,14 +253,12 @@ jobs:
        deps: testing_unit
        llvm: 'true'
    - name: Test SPEC=2
-      run: SPEC=2 DEV=CPU 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 -k "not test_setitem_big" -k "not test_conv2d_ceildiv_edge_case" --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 -k "not test_setitem_big" -k "not test_conv2d_ceildiv_edge_case" --splits 2 --group ${{ matrix.group }}

  fuzzing:
    name: Fuzzing
    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 10
-    env:
-      DEV: CPU
    steps:
    - name: Checkout Code
      uses: actions/checkout@v6
@@ -382,7 +377,6 @@ jobs:
    timeout-minutes: 15
    env:
      CHECK_OOB: 0
-      DEV: CPU
    steps:
      - name: Checkout Code
        uses: actions/checkout@v6
@@ -396,10 +390,10 @@ jobs:
        run: |
          parallel --link --tagstring '[{1}]' '{2}' \
            ::: llama 'llama q4' qwen3.5 qwen \
-            ::: $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model llama3.2:1b | tee /dev/stderr | grep -i rooster' \
-                $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster' \
-                $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model qwen3.5:0.8b | tee /dev/stderr | grep -i rooster' \
-                $'echo "What\'s a female chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model qwen3:0.6b | tee /dev/stderr | grep -i hen'
+            ::: $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --model llama3.2:1b | tee /dev/stderr | grep -i rooster' \
+                $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster' \
+                $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --model qwen3.5:0.8b | tee /dev/stderr | grep -i rooster' \
+                $'echo "What\'s a female chicken called? Answer with only one word." | python3 -m tinygrad.llm --model qwen3:0.6b | tee /dev/stderr | grep -i hen'
                # NOTE: qwen is dumb and only knows about female chickens

 # ****** Models Tests ******
@@ -456,8 +450,6 @@ jobs:
    name: Linux (DEV=${{ matrix.dev }})
    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 20
-    env:
-      DEV: ${{ matrix.dev }}
    steps:
    - name: Checkout Code
      uses: actions/checkout@v6
@@ -469,6 +461,8 @@ jobs:
        llvm: ${{ contains(matrix.dev, 'LLVM') || contains(matrix.dev, 'LVP') || contains(matrix.dev, 'CLANG') }}
        webgpu: ${{ matrix.dev == 'WEBGPU' }}
        opencl: ${{ matrix.dev == 'CL' }}
+    - name: Set env
+      run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
    - name: Check Device.DEFAULT and print some source
      run: |
        python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
@@ -510,7 +504,7 @@ jobs:
      - name: Run AMD renderer tests (AMD:LLVM)
        run: DEV=MOCKKFD+AMD:LLVM python -m pytest -n=auto test/amd/ --durations 20
      - name: Run SQTT profiling tests
-        run: VIZ=-2 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
+        run: PROFILE=1 SQTT=1 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
      - name: Run AMD emulated tests on NULL backend
        env:
          AMD: 0
@@ -685,5 +679,4 @@ jobs:
        run: |
          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
-          # QCOMCL compiles in qemu, too slow for parallel workers
-          ${{ contains(matrix.dev, 'QCOMCL') && 'PARALLEL=0' || '' }} python -m pytest -n=auto test/backend/test_ops.py --durations=20
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20
@@ -69,4 +69,3 @@ mutants
 dagre/
 graphlib/
 uv.lock
-pi_session_window0.jsonl
@@ -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/backend/test_ops.py test/backend/test_schedule.py test/backend/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
+        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
@@ -4,4 +4,3 @@
 - Run `python -m mypy tinygrad/` to typecheck
 - Run `python -m ruff check .` to lint
 - Read `./tinygrad/viz/README.md` for profiling and debugging rewrite rules
- Do not do amend commits. Always do a new commit if a force push to origin would be required.
@@ -2,7 +2,7 @@ import os, pytest, signal, threading

@pytest.hookimpl(wrapper=True)
 def pytest_runtest_call(item):
-  t = threading.Timer(int(os.getenv("TEST_TIMEOUT", 90)), os.kill, args=(os.getpid(), signal.SIGABRT))
+  t = threading.Timer(int(os.getenv("TEST_TIMEOUT", 300)), os.kill, args=(os.getpid(), signal.SIGABRT))
  t.start()
  try: yield
  finally:
@@ -122,7 +122,7 @@ def example_5_custom_assembly(a:Tensor, correct):
        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 UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in self.instructions]))))
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in self.instructions]))))

  CU_COUNT = 32
  LANES = 64
@@ -1667,14 +1667,15 @@ def train_llama3():
 def train_gptoss():
  from examples.mlperf.models.gpt_oss import GPTOSS, GPT_OSS_20B, apply_grad, FP8_DTYPE
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
-  from examples.mlperf.optim import GradAccClipAdamW, GradAccClipAdamWGroup, fclip_grads
+  from examples.mlperf.optim import GradAccClipAdamW, GradAccClipAdamWGroup, clip_grads

  BENCHMARK = getenv("BENCHMARK")

  config = {}
  BASEDIR            = config["BASEDIR"]                = Path(getenv("BASEDIR", "/raid/datasets/c4-8b/"))
  BS                 = config["BS"]                     = getenv("BS", 16)
-  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS
+  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
+  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
  SEED               = config["SEED"]                   = getenv("SEED", 5760)
  DATA_SEED          = config["DATA_SEED"]              = getenv("DATA_SEED", SEED)
  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
@@ -1736,8 +1737,8 @@ def train_gptoss():
  params_wd = [p for p in params if p.ndim >= 3]
  params_no_wd = [p for p in params if p.ndim < 3]
  optim = GradAccClipAdamWGroup(
-    GradAccClipAdamW(params_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=1, device=optim_device),
-    GradAccClipAdamW(params_no_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=0.0, grad_acc=1, device=optim_device),
+    GradAccClipAdamW(params_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device),
+    GradAccClipAdamW(params_no_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=0.0, grad_acc=grad_acc, device=optim_device),
  )

  for p in optim.params:
@@ -1769,32 +1770,31 @@ def train_gptoss():

  @TinyJit
  @Context(TRAINING=1)
-  def step(tokens:Tensor):
+  def minibatch(tokens:Tensor):
    if is_dp: tokens = tokens.to(None).shard(device, 0)
    if not is_sharding: tokens = tokens.to(None)
-
    logits:Tensor = model(tokens[:, :-1], save=True)
-    if getenv("FUSED_CE", 0):
-      from extra.llama_kernels.fused_ce import fused_ce_loss
-      loss = fused_ce_loss(logits.cast(dtypes.bfloat16), tokens[:, 1:], label_smoothing=0.0)
-    else:
-      loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])

    for g, new_g in zip(grads, loss.gradient(*optim.params)):
      apply_grad(g, new_g.uop)

-    Tensor.realize(loss, *grads)
+    loss_cpu = loss.flatten().float().to("CPU")
+    return loss_cpu.realize(*grads)

-    clipped_grads, grad_norm = fclip_grads(grads, 1.0)
-    optim.fstep(clipped_grads, grad_norm)
+  @TinyJit
+  def optim_step():
+    grad_norm = clip_grads(grads, grad_acc, 1.0)
+    optim.fstep(grads, grad_norm)
    scheduler.step()

-    loss_cpu = loss.flatten().float().to("CPU")
+    for g in grads: g.assign(0)
+
    lr_cpu = optim.lr.float().to("CPU")
    grad_norm_cpu = grad_norm.float().to("CPU")
-    Tensor.realize(loss_cpu, lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales)
+    Tensor.realize(lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales)

-    return loss_cpu, lr_cpu, grad_norm_cpu
+    return lr_cpu, grad_norm_cpu

  @TinyJit
  @Context(TRAINING=0)
@@ -1843,20 +1843,30 @@ def train_gptoss():
      profile_marker(f"train @ {i}")
      st = time.perf_counter()

-      ist = time.perf_counter()
+      stopped = False
+      losses, data_time, dev_time = [], 0, 0
+      for _ in range(grad_acc if i >= 2 else 1):
+        ist = time.perf_counter()
+        try: tokens = next(train_iter)
+        except StopIteration:
+          stopped = True
+          break
+        mst = time.perf_counter()
+        data_time += mst - ist
+        losses.append(minibatch(tokens).item())
+        dev_time += time.perf_counter() - mst
+      if stopped: break

-      try: tokens = next(train_iter)
-      except StopIteration: break
-      mst = time.perf_counter()
-      data_time = mst - ist
-
-      ret = step(tokens)
-      dev_time = time.perf_counter() - mst
-
-      loss, lr, grad_norm = ret[0].item(), ret[1].item(), ret[2].item()
+      gt = time.perf_counter()
+      ret = optim_step()
+      lr, grad_norm = ret[0].item(), ret[1].item()
      et = time.perf_counter()

+      loss = sum(losses) / len(losses)
+      optim_time = et - gt
+      dev_time += optim_time
      step_time = et - st
+      gbs_time = gt - st
      if BENCHMARK: step_times.append(step_time)

      i += 1
@@ -1866,7 +1876,7 @@ def train_gptoss():
      gflops = GlobalCounters.global_ops / 1e9 / dev_time
      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * 4.6e15)) * 100
      tqdm.write(
-          f"{i:5} {step_time:.3f} s step, {dev_time:.3f} s dev, {data_time:.3f} s data, {loss:.4f} loss, " \
+          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, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
      if DEBUG >= 1: tqdm.write("  mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))

@@ -1876,6 +1886,8 @@ def train_gptoss():
          "train/lr": lr,
          "train/grad_norm": grad_norm,
          "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,
@@ -183,12 +183,10 @@ class GPTOSS:
    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
    xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)  # (B,N,H,D)/(B,N,KV,D)

-    fa_saves = []
    if getenv("HK_FLASH_ATTENTION"):
      from extra.thunder.amd.fa import flash_attention
-      attn, _, l_vec = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks, window=self.sliding_window if sliding else 0)
+      attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks, window=self.sliding_window if sliding else 0)
      attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
-      fa_saves = [xq, xk, xv, l_vec]
    elif sliding:
      attn = self._sliding_attention(xq, xk, xv, sinks)
    else:
@@ -202,7 +200,7 @@ class GPTOSS:
      attn = (w @ xvm).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, self.n_heads * self.head_dim)

    out = matmul_mx(attn, wo, wo_scale) + wo_bias
-    return out, [x_normed, rrms, attn] + fa_saves
+    return out, [x_normed, rrms, attn]

  def feed_forward(self, x:Tensor, *, ffn_norm:Tensor, gate:Tensor, gate_bias:Tensor,
                   w_gate_up:Tensor, w_gate_up_scale:Tensor, w_gate_up_bias:Tensor,
@@ -223,7 +221,6 @@ class GPTOSS:
      z = grouped_mx_gemm(_pad_cols(y.cast(dtypes.bfloat16)), (w_down, w_down_scale), r.off)[:, :dim] \
          + (onehot @ w_down_bias.float()).cast(dtypes.bfloat16)
      out = combine(z, r, inp.shape[0], self.experts_per_tok).reshape(bsz, seqlen, dim)
-      return out, [x_normed, rrms, xg, h, y, z, r.weights, r.dest_row, r.off]
    else:
      thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
      weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)
@@ -282,14 +279,14 @@ def apply_grad(grad_buf:Tensor, new_grad:UOp):
  pads = _get_pads(new_grad)
  if len(pads) <= 1:
    new_grad = new_grad.cast(grad_buf.dtype)
-    grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(new_grad))
+    grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(grad_buf.uop + new_grad))
    return
  cur = grad_buf.uop
  for pad in sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0, reverse=True):
    if pad.op == Ops.PAD:
-      grad_shrink = tuple((p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg))
+      grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg)])
      buf_slice = cur.shrink(grad_shrink)
-      cur = cur.after(buf_slice.store(pad.src[0].cast(cur.dtype)))
+      cur = cur.after(buf_slice.store(buf_slice + pad.src[0].cast(cur.dtype)))
    else:
      cur = cur.after(cur.store(cur + pad.cast(cur.dtype)))
  grad_buf.uop = cur
@@ -15,7 +15,7 @@ def stochastic_round_bf16(x:Tensor) -> Tensor:
  bits = x.bitcast(dtypes.uint32)
  if isinstance(x.device, tuple):
    shape = x.uop.shard_shape if x.uop.axis is not None else x.shape
-    noise = Tensor(UOp(Ops.MSTACK, src=tuple(Tensor.rand(*shape, device=d).uop for d in x.device)))
+    noise = Tensor(UOp(Ops.MSTACK, dtypes.default_float, tuple(Tensor.rand(*shape, device=d).uop for d in x.device)))
  else:
    noise = x.rand_like()
  noise = (noise * 0xFFFF).cast(dtypes.uint32)
@@ -27,11 +27,6 @@ def clip_grads(grads:list[Tensor], grad_acc, clip_norm) -> Tensor:
  for g in grads: g.assign((g * (clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(g.dtype))
  return total_norm

-def fclip_grads(grads:list[Tensor], clip_norm) -> Tensor:
-  total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
-  scale = (clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)
-  return [(g * scale).cast(g.dtype) for g in grads], total_norm
-
 class GradAccClipAdamW(Optimizer):
  def __init__(self, params:list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, grad_acc=1, clip_norm=1.0, device=None, fused=FUSE_OPTIM):
    super().__init__(params, lr, device, fused)
@@ -35,7 +35,7 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=2048 MAX_STEPS=1200000
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
 export SEQLEN=${SEQLEN:-8192}
@@ -44,9 +44,9 @@ export SEED=${SEED:-5760}
 export DATA_SEED=${DATA_SEED:-5760}

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

-export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+export FAKEDATA=${FAKEDATA:-$([[ "$DEV" == NULL:* ]] && echo 1 || echo 0)} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
 fi
@@ -35,7 +35,7 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=2048 MAX_STEPS=1200000
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
 export SEQLEN=${SEQLEN:-8192}
@@ -44,6 +44,6 @@ export SEED=${SEED:-$RANDOM}
 export DATA_SEED=${DATA_SEED:-5760}

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

 python3 examples/mlperf/model_train.py
@@ -17,7 +17,7 @@ export USE_ATOMICS=1
 export ASM_GEMM=1
 export WQKV=1
 export MASTER_WEIGHTS=1
-export MXFP4=1
+export FP8=1
 export ALLREDUCE_CAST=1
 export FAST_CE=1
 export FUSED_INPUT_QUANTIZE=1
@@ -26,7 +26,7 @@ export FUSED_ADD_NORM_MUL_QUANTIZE=1
 export FUSED_SILU_W13=1
 export SPLIT_W13=0

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

@@ -44,7 +44,7 @@ export SEED=$RANDOM
 export DATA_SEED=$SEED

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

 export LOGMLPERF=1

@@ -107,21 +107,14 @@ def compile(onnx_file):
  return inputs, test_val

 def test_vs_compile(run, inputs, test_val=None):
-  if (log:=bool(getenv("BENCHMARK_LOG", ""))): from extra.bench_log import WallTimeEvent, BenchEvent

  # run 20 times
  step_times = []
  for _ in range(20):
    st = time.perf_counter()
-    if log:
-      with WallTimeEvent(BenchEvent.STEP):
-        out = run(**inputs)
-        mt = time.perf_counter()
-        val = out.numpy()
-    else:
-      out = run(**inputs)
-      mt = time.perf_counter()
-      val = out.numpy()
+    out = run(**inputs)
+    mt = time.perf_counter()
+    val = out.numpy()
    et = time.perf_counter()
    step_times.append((et-st)*1e3)
    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")
@@ -167,6 +160,12 @@ def test_vs_onnx(new_inputs, test_val, onnx_file, tol):
  print("test vs onnx passed")
  return timings

+def bench(run, inputs):
+  from extra.bench_log import WallTimeEvent, BenchEvent
+  for _ in range(10):
+    with WallTimeEvent(BenchEvent.STEP):
+      run(**inputs).numpy()
+
 if __name__ == "__main__":
  if getenv("RUN_PICKLE"):
    with open(OUTPUT, "rb") as f: pickle_loaded = load_pickle(f)
@@ -182,3 +181,6 @@ if __name__ == "__main__":
    test_vs_compile(pickle_loaded, inputs, outputs)
    if getenv("SELFTEST"):
      test_vs_onnx(inputs, outputs, onnx_file, 1e-4)
+
+  if getenv("BENCHMARK_LOG", ""):
+    bench(pickle_loaded, inputs)
@@ -84,8 +84,7 @@ class AMSMI(AMDev):
    with open(f"/sys/bus/pci/devices/{self.pcibus}/power_state", "r") as f: return f.read().strip().rstrip()

 class SMICtx:
-  def __init__(self, dev_filter=None):
-    self.dev_filter = dev_filter
+  def __init__(self):
    self.devs = []
    self.opened_pcidevs = []
    self.opened_pci_resources = {}
@@ -136,7 +135,6 @@ class SMICtx:
    pattern = os.path.join('/tmp', 'am_*.lock')
    for d in [f[8:-5] for f in glob.glob(pattern)]:
      if d.startswith("usb"): continue
-      if self.dev_filter is not None and d != self.dev_filter: continue
      if d not in self.opened_pcidevs:
        self._open_am_device(d)

@@ -408,7 +406,7 @@ if __name__ == "__main__":

  try:
    if not args.list: os.system('clear')
-    smi_ctx = SMICtx(args.dev)
+    smi_ctx = SMICtx()
    while True:
      smi_ctx.rescan_devs()
      smi_ctx.draw(args.list)
@@ -35,7 +35,7 @@ class WallTimeEvent:
    return self
  def __exit__(self, *_):
    self.time = time.monotonic() - self.start
-    _events[self.event]["wall"].append((self.time, BENCHMARK_LOG.value))
+    _events[self.event]["wall"].append(self.time)
    return False

 class KernelTimeEvent:
@@ -47,19 +47,19 @@ class KernelTimeEvent:
    self.start = GlobalCounters.time_sum_s
    return self
  def __exit__(self, *_):
-    _events[self.event]["kernel"].append((GlobalCounters.time_sum_s - self.start, BENCHMARK_LOG.value))
+    _events[self.event]["kernel"].append(GlobalCounters.time_sum_s - self.start)
    return False

 def log_event_instant(event:InstantBenchEvent, value:float):
-  _events[event].append((value, BENCHMARK_LOG.value))
+  _events[event].append(value)

 if BENCHMARK_LOG:
  INFLUXDB_HOST = getenv("INFLUXDB_HOST", "")
  INFLUXDB_ORG = getenv("INFLUXDB_ORG", "tiny")
  INFLUXDB_TOKEN = getenv("INFLUXDB_TOKEN", "")

-  def _create_point(run_id, i, attempt, ref, commit, name, value, log_name, run):
-    point = Point(log_name.replace(':', '_').replace('.', '_')).tag("id", run_id).tag("index", i)
+  def _create_point(run_id, i, attempt, ref, commit, name, value, run):
+    point = Point(BENCHMARK_LOG.value).tag("id", run_id).tag("index", i)
    point = point.tag("device", Device.DEFAULT)
    point = point.tag("attempt", attempt).tag("ref", ref).tag("commit", commit)
    point = point.field(name, value).field("x", run)
@@ -91,12 +91,12 @@ if BENCHMARK_LOG:
      run_id = str(uuid.uuid4())
      if isinstance(event, BenchEvent):
        for event_type, values in _events[event].items():
-          for i, (value, log_name) in enumerate(values):
-            point = _create_point(run_id, i, attempt, ref, commit, f"{event.value}_{event_type}", value, log_name, run)
+          for i, value in enumerate(values):
+            point = _create_point(run_id, i, attempt, ref, commit, f"{event.value}_{event_type}", value, run)
            points.append(point)
      else:
-        for i, (value, log_name) in enumerate(_events[event]):
-          point = _create_point(run_id, i, attempt, ref, commit, event.value, value, log_name, run)
+        for i, value in enumerate(_events[event]):
+          point = _create_point(run_id, i, attempt, ref, commit, event.value, value, run)
          points.append(point)

    write_options = WriteOptions(write_type=WriteType.synchronous, retry_interval=5000, max_retries=5, max_retry_delay=30000, exponential_base=2)
@@ -1,238 +0,0 @@
-import ctypes, struct
-from tinygrad.helpers import ceildiv, getenv, wait_cond, DEBUG
-from tinygrad.runtime.autogen import bnxt, pci
-from tinygrad.runtime.support.system import PCIDevice, System, ipv4_to_gid
-
-BNXT_DEBUG = getenv("BNXT_DEBUG", 0)
-BNXT_ACCESS, BNXT_INIT_MASK, BNXT_RTR_MASK, BNXT_RTS_MASK = 3, 0xd, 0x41515ad, 0xae005
-BNXT_CHIMP_COMM, BNXT_CHIMP_COMM_TRIGGER = 0x0, 0x100
-BNXT_BACKING_STORE = ((0, 2), (1, 0), (2, 2), (3, 0), (4, 2), (5, 0), (6, 0), (14, 2), (15, 0))
-
-def db_value(xid, typ, index, epoch):
-  return (xid & bnxt.DBC_DBC_XID_MASK | bnxt.DBC_DBC_PATH_ROCE | typ | bnxt.BNXT_QPLIB_DBR_VALID) << 32 | \
-         index & bnxt.DBC_DBC_INDEX_MASK | epoch << bnxt.BNXT_QPLIB_DBR_EPOCH_SHIFT
-
-def _pbl(dev, paddrs, queue=False):
-  if len(paddrs) == 1: return 0, paddrs[0]
-  values = [p | bnxt.PTU_PTE_VALID for p in paddrs]
-  if queue:
-    values[-1] |= bnxt.PTU_PTE_LAST
-    if len(values) > 1: values[-2] |= bnxt.PTU_PTE_NEXT_TO_LAST
-  table, table_paddrs = dev.pci_dev.alloc_sysmem(ceildiv(len(values), 512) * 0x1000)
-  table[:len(values) * 8] = struct.pack(f"<{len(values)}Q", *values)
-  if len(table_paddrs) == 1: return 1, table_paddrs[0]
-  top, top_paddrs = dev.pci_dev.alloc_sysmem(0x1000)
-  top[:len(table_paddrs) * 8] = struct.pack(f"<{len(table_paddrs)}Q", *(p | bnxt.PTU_PTE_VALID for p in table_paddrs))
-  return 2, top_paddrs[0]
-
-def _queue(dev, stride:int=16, aux=False):
-  mem, paddrs = dev.pci_dev.alloc_sysmem(0x1000 + aux * 0x400)
-  level, base = _pbl(dev, paddrs, queue=True)
-  return {"mem":mem, "paddrs":paddrs, "stride":stride, "prod":0, "cons":0, "level":level, "base":base}
-
-def _qread(q, i):
-  off = (i & 15) * q["stride"]
-  return q["mem"][off:off + q["stride"]]
-
-def _qwrite(q, i, data, aux=False):
-  off = 0x1000 + i % 128 * 8 if aux else (i & 15) * q["stride"]
-  q["mem"][off:off + len(data)] = data
-
-class BNXTDev:
-  def __init__(self, pci_dev:PCIDevice, ip:str=getenv("BNXT_IP", "10.0.0.1")):
-    self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
-    self.bar0, self.db = pci_dev.map_bar(0, fmt='I'), pci_dev.map_bar(2, fmt='Q')
-    pci_dev.write_config(pci.PCI_COMMAND, pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
-    self.resp, self.resp_pa = pci_dev.alloc_sysmem(0x1000)
-    self.seq = 0
-
-    ver = self.hwrm("ver_get")
-    if DEBUG >= 2: print(f"bnxt {self.devfmt}: firmware {ver.hwrm_fw_maj_8b}.{ver.hwrm_fw_min_8b}.{ver.hwrm_fw_bld_8b}")
-    self.hwrm("func_reset", timeout_ms=40000)
-    caps = self.hwrm("func_qcaps", fid=0xffff)
-    self.mac, self.port_id = int.from_bytes(bytes(caps.mac_address), 'big'), caps.port_id
-    self.hwrm("func_drv_rgtr")
-    self.db_off = self.hwrm("func_qcfg", fid=0xffff).legacy_l2_db_size_kb * 1024
-
-    self.setup_backing_store()
-    self._open_rcfw()
-    self._open_l2()
-    self.local_gid = ipv4_to_gid(ip)
-    gids, mac = (ctypes.c_uint32 * 4)(*(int.from_bytes(self.local_gid[i:i + 4], 'big') for i in (12, 8, 4, 0))), self.mac.to_bytes(6, 'big')
-    smac = (ctypes.c_uint16 * 3)(*(int.from_bytes(mac[i:i + 2], 'big') for i in (0, 2, 4)))
-    self.gid_id = self.rcfw("add_gid", gid=gids, src_mac=smac).xid
-
-    if DEBUG >= 2: print(f"bnxt {self.devfmt}: booted mac={self.mac.to_bytes(6, 'big').hex(':')} gid={self.local_gid.hex()}")
-
-  def hwrm(self, name, timeout_ms=10000, **fields):
-    inp, out = getattr(bnxt, f"struct_hwrm_{name}_input"), getattr(bnxt, f"struct_hwrm_{name}_output")
-    opcode = getattr(bnxt, f"HWRM_{name.upper()}")
-    self.seq = (self.seq + 1) & 0xffff
-    data = bytes(inp(req_type=opcode, cmpl_ring=bnxt.BNXT_HWRM_NO_CMPL_RING, seq_id=self.seq, target_id=bnxt.BNXT_HWRM_TARGET,
-                     resp_addr=self.resp_pa[0], **fields))
-    self.resp[:] = bytes(len(self.resp))
-    System.memory_barrier()
-    for i, w in enumerate(memoryview(bytearray(data.ljust(bnxt.HWRM_MAX_REQ_LEN, b'\0'))).cast('I')):
-      self.bar0[BNXT_CHIMP_COMM // 4 + i] = w
-    self.bar0[BNXT_CHIMP_COMM_TRIGGER // 4] = 1
-    def hdr(): return bnxt.struct_hwrm_resp_hdr.from_buffer_copy(bytes(self.resp[:8]))
-    wait_cond(lambda: (n := hdr().resp_len) and hdr().seq_id == self.seq and self.resp[n - 1], timeout_ms=timeout_ms, msg=f"HWRM {name}")
-    ret = out.from_buffer_copy(bytes(self.resp[:ctypes.sizeof(out)]))
-    assert ret.error_code == 0, f"HWRM {name}: {ret.error_code}"
-    return ret
-
-  def setup_backing_store(self):
-    counts: dict[int, int] = {}
-    for typ, extra in BNXT_BACKING_STORE:
-      caps = self.hwrm("func_backing_store_qcaps_v2", type=typ)
-      size, splits = caps.entry_size, tuple(getattr(caps, f"split_entry_{j}") for j in range(caps.subtype_valid_cnt))
-      counts[typ] = n = counts[0] if typ == 15 else max(caps.min_num_entries, sum(splits) + extra)
-      # a zero bitmap means the type has a single instance 0
-      for instance in [i for i in range(8) if caps.instance_bit_map >> i & 1] or [0]:
-        mem, paddrs = self.pci_dev.alloc_sysmem(ceildiv(n * size, 0x1000) * 0x1000)
-        if caps.ctx_init_value:
-          for off in range(caps.ctx_init_offset, len(mem), size): mem[off] = caps.ctx_init_value
-        lvl, base = _pbl(self, paddrs)
-        self.hwrm("func_backing_store_cfg_v2", type=typ, instance=instance, entry_size=size, num_entries=n, page_dir=base,
-          page_size_pbl_level=lvl, subtype_valid_cnt=len(splits),
-          flags=bnxt.FUNC_BACKING_STORE_CFG_V2_REQ_FLAGS_BS_CFG_ALL_DONE if typ == 15 else 0,
-          **{f"split_entry_{j}": v for j, v in enumerate(splits)})
-
-  def _open_rcfw(self):
-    self.rcfw_first = True
-
-    self.creq = _queue(self)
-    self.creq_id = self.hwrm("ring_alloc", ring_type=bnxt.RING_ALLOC_REQ_RING_TYPE_NQ, page_tbl_addr=self.creq["base"],
-      page_size=12, page_tbl_depth=self.creq["level"], length=16, int_mode=bnxt.RING_ALLOC_REQ_INT_MODE_MSIX).ring_id
-
-    self.cmdq = _queue(self)
-    self.doorbell(self.creq_id, bnxt.DBC_DBC_TYPE_NQ_ARM, 0, 0)
-    init = bnxt.struct_cmdq_init(cmdq_pbl=self.cmdq["base"], creq_ring_id=self.creq_id,
-                                 cmdq_size_cmdq_lvl=16 << bnxt.CMDQ_INIT_CMDQ_SIZE_SFT)
-
-    System.memory_barrier()
-    for i, w in enumerate(memoryview(bytearray(bytes(init))).cast('I')): self.bar0[bnxt.RCFW_COMM_BASE_OFFSET // 4 + i] = w
-
-    _, p = self.pci_dev.alloc_sysmem(0x1000)
-    self.rcfw("initialize_fw", stat_ctx_id=self.hwrm("stat_ctx_alloc", stats_dma_addr=p[0], stats_dma_length=176).stat_ctx_id,
-      flags=bnxt.CMDQ_INITIALIZE_FW_FLAGS_HW_REQUESTER_RETX_SUPPORTED)
-
-    # RoCE notification ring: never armed or serviced, but CQ and L2 ring allocation require one
-    nq = _queue(self)
-    self.nq_id = self.hwrm("ring_alloc", ring_type=bnxt.RING_ALLOC_REQ_RING_TYPE_NQ, page_tbl_addr=nq["base"],
-      page_size=12, page_tbl_depth=nq["level"], length=16, logical_id=1, int_mode=bnxt.RING_ALLOC_REQ_INT_MODE_MSIX).ring_id
-
-  def rcfw(self, name, timeout_ms=20000, **fields):
-    req_t, resp_t = getattr(bnxt, f"struct_cmdq_{name}"), getattr(bnxt, f"struct_creq_{name}_resp")
-    op = getattr(bnxt, f"CMDQ_BASE_OPCODE_{name.upper()}")
-    data = bytes(req_t(opcode=op, cmd_size=(slots := ceildiv(ctypes.sizeof(req_t), 16)), **fields)).ljust(slots * 16, b'\0')
-    for i in range(slots): _qwrite(self.cmdq, self.cmdq["prod"] + i, data[i * 16:(i + 1) * 16])
-
-    self.cmdq["prod"] += slots
-    prod = self.cmdq["prod"] & 0xffff
-    if self.rcfw_first: prod, self.rcfw_first = prod | 1 << bnxt.FIRMWARE_FIRST_FLAG, False
-
-    System.memory_barrier()
-
-    self.bar0[(bnxt.RCFW_COMM_BASE_OFFSET + bnxt.RCFW_PF_VF_COMM_PROD_OFFSET) // 4] = prod
-    self.bar0[(bnxt.RCFW_COMM_BASE_OFFSET + bnxt.RCFW_COMM_TRIG_OFFSET) // 4] = bnxt.RCFW_CMDQ_TRIG_VAL
-
-    def poll():
-      h = bnxt.struct_creq_base.from_buffer_copy(bytes(_qread(self.creq, self.creq["cons"])))
-      return bool(h.v & bnxt.CREQ_BASE_V) != bool((self.creq["cons"] // 16) & 1)
-    wait_cond(poll, timeout_ms=timeout_ms, msg=f"RCFW {name}")
-
-    ret = resp_t.from_buffer_copy(bytes(_qread(self.creq, self.creq["cons"])))
-    self.creq["cons"] += 1
-
-    # NQ_ARM also publishes the CREQ consumer index, which is what frees ring space for the next command
-    self.doorbell(self.creq_id, bnxt.DBC_DBC_TYPE_NQ_ARM, self.creq["cons"] & 15, (self.creq["cons"] // 16) & 1)
-    assert ret.status == 0, f"RCFW {name}: {ret.status}"
-
-    if BNXT_DEBUG >= 1: print(f"bnxt {self.devfmt}: rcfw {name} xid={getattr(ret, 'xid', 0):#x}")
-    return ret
-
-  def doorbell(self, xid, typ, index, epoch):
-    System.memory_barrier()
-    self.db[self.db_off // 8] = db_value(xid, typ, index, epoch)
-
-  # L2 receive path, required for RoCE ingress even though no ethernet receive buffers are posted
-  def _open_l2(self):
-    cq = _queue(self)
-    ci = self.hwrm("ring_alloc", enables=bnxt.RING_ALLOC_REQ_ENABLES_NQ_RING_ID_VALID, ring_type=bnxt.RING_ALLOC_REQ_RING_TYPE_L2_CMPL,
-      page_tbl_addr=cq["base"], page_size=12, page_tbl_depth=cq["level"], length=16, nq_ring_id=self.nq_id).ring_id
-    rx = _queue(self)
-    ri = self.hwrm("ring_alloc", enables=bnxt.RING_ALLOC_REQ_ENABLES_NQ_RING_ID_VALID |
-      bnxt.RING_ALLOC_REQ_ENABLES_RX_BUF_SIZE_VALID, ring_type=bnxt.RING_ALLOC_REQ_RING_TYPE_RX, page_tbl_addr=rx["base"],
-      page_size=12, page_tbl_depth=rx["level"], length=16, rx_buf_size=640, nq_ring_id=self.nq_id).ring_id
-    vi = self.hwrm("vnic_alloc").vnic_id
-    self.hwrm("vnic_cfg", enables=bnxt.VNIC_CFG_REQ_ENABLES_MRU | bnxt.VNIC_CFG_REQ_ENABLES_DEFAULT_RX_RING_ID |
-      bnxt.VNIC_CFG_REQ_ENABLES_DEFAULT_CMPL_RING_ID, vnic_id=vi, mru=9018,
-      default_rx_ring_id=ri, default_cmpl_ring_id=ci)
-    self.hwrm("cfa_l2_filter_alloc", flags=bnxt.CFA_L2_FILTER_ALLOC_REQ_FLAGS_PATH_RX,
-      enables=bnxt.CFA_L2_FILTER_ALLOC_REQ_ENABLES_L2_ADDR | bnxt.CFA_L2_FILTER_ALLOC_REQ_ENABLES_L2_ADDR_MASK |
-      bnxt.CFA_L2_FILTER_ALLOC_REQ_ENABLES_DST_ID, l2_addr=tuple(self.mac.to_bytes(6, 'big')), l2_addr_mask=(0xff,) * 6, dst_id=vi)
-
-  def register_mem(self, paddrs:list[int], size:int, log_page_size:int=12) -> int:
-    level, base = _pbl(self, paddrs[:ceildiv(size, 1 << log_page_size)])
-    return self.rcfw("register_mr", flags=bnxt.CMDQ_REGISTER_MR_FLAGS_ALLOC_MR,
-      log2_pg_size_lvl=level << bnxt.CMDQ_REGISTER_MR_LVL_SFT | log_page_size << bnxt.CMDQ_REGISTER_MR_LOG2_PG_SIZE_SFT,
-      access=bnxt.CMDQ_REGISTER_MR_ACCESS_LOCAL_WRITE | bnxt.CMDQ_REGISTER_MR_ACCESS_REMOTE_WRITE,
-      log2_pbl_pg_size=12, pbl=base, va=paddrs[0], mr_size=size).xid
-
-class BNXTQP:
-  def __init__(self, dev:BNXTDev):
-    self.dev, self.sq_psn, self.msn = dev, 0, 0
-
-    self.cqq = _queue(dev, ctypes.sizeof(bnxt.struct_cq_base))
-    self.cq_id = dev.rcfw("create_cq", cq_size=16, pbl=self.cqq["base"],
-                          pg_size_lvl=self.cqq["level"], cq_fco_cnq_id=dev.nq_id).xid
-
-    self.sq = _queue(dev, aux=True)
-    self.qpn = dev.rcfw("create_qp", type=bnxt.CMDQ_CREATE_QP_TYPE_RC,
-      sq_size=16, sq_fwo_sq_sge=1, scq_cid=self.cq_id, rcq_cid=self.cq_id,
-      sq_pbl=self.sq["base"], sq_pg_size_sq_lvl=self.sq["level"]).xid
-    self.qp_op(1, BNXT_INIT_MASK, access=BNXT_ACCESS, pkey=0xffff)
-
-  def qp_op(self, state, mask, network_type=0, **fields):
-    self.dev.rcfw("modify_qp", qp_cid=self.qpn, modify_mask=mask,
-      network_type_en_sqd_async_notify_new_state=state | network_type, **fields)
-
-  def connect(self, qpn:int, gid:bytes, mac:int):
-    network_type = bnxt.CMDQ_MODIFY_QP_NETWORK_TYPE_ROCEV2_IPV4
-    dgid = (ctypes.c_uint32 * 4)(*(int.from_bytes(gid[i:i + 4], 'little') for i in (0, 4, 8, 12)))
-    dmac = (ctypes.c_uint16 * 3)(*(int.from_bytes(mac.to_bytes(6, 'big')[i:i + 2], 'little') for i in (0, 2, 4)))
-
-    self.qp_op(2, BNXT_RTR_MASK, network_type=network_type, qp_type=bnxt.CMDQ_MODIFY_QP_QP_TYPE_RC, access=BNXT_ACCESS,
-      pkey=0xffff, dgid=dgid, sgid_index=self.dev.gid_id, hop_limit=64, dest_mac=dmac,
-      path_mtu_pingpong_push_enable=bnxt.CMDQ_MODIFY_QP_PATH_MTU_MTU_1024, max_dest_rd_atomic=4,
-      dest_qp_id=qpn)
-    self.qp_op(3, BNXT_RTS_MASK, network_type=network_type, qp_type=bnxt.CMDQ_MODIFY_QP_QP_TYPE_RC, access=BNXT_ACCESS,
-      max_rd_atomic=1)
-
-    if BNXT_DEBUG >= 1: print(f"bnxt: QP {self.qpn:#x} connected (remote={qpn:#x})")
-
-  def _poll(self, timeout):
-    def poll():
-      base = bnxt.struct_cq_base.from_buffer_copy(bytes(_qread(self.cqq, self.cqq["cons"])))
-      return bool(base.cqe_type_toggle & bnxt.CQ_BASE_TOGGLE) == (not bool((self.cqq["cons"] // 16) & 1))
-    wait_cond(poll, timeout_ms=timeout, msg="BNXT CQ")
-    raw = bytes(_qread(self.cqq, self.cqq["cons"]))
-    self.cqq["cons"] += 1
-    self.dev.doorbell(self.cq_id, bnxt.DBC_DBC_TYPE_CQ, self.cqq["cons"] & 15, (self.cqq["cons"] // 16) & 1)
-    return raw
-
-  def rdma_write(self, rva, rkey, lva, lkey, size, timeout_ms=20000):
-    start = self.sq["prod"] & 15
-    hdr = bytes(bnxt.struct_sq_rdma_hdr(wqe_type=bnxt.SQ_RDMA_HDR_WQE_TYPE_WRITE_WQE,
-                flags=bnxt.SQ_SEND_FLAGS_SIGNAL_COMP, wqe_size=3, length=size, remote_va=rva, remote_key=rkey))
-    for i, data in enumerate((hdr[:16], hdr[16:32], bytes(bnxt.struct_sq_sge(va_or_pa=lva, l_key=lkey, size=size)))):
-      _qwrite(self.sq, start + i, data)
-    nxt = (self.sq_psn + max(1, ceildiv(size, 1024))) & 0xffffff
-    value = start << bnxt.SQ_MSN_SEARCH_START_IDX_SFT | nxt << bnxt.SQ_MSN_SEARCH_NEXT_PSN_SFT | self.sq_psn
-    _qwrite(self.sq, self.msn, struct.pack("<Q", value), aux=True)
-
-    self.msn, self.sq_psn, self.sq["prod"] = (self.msn + 1) % 128, nxt, self.sq["prod"] + 3
-    self.dev.doorbell(self.qpn, bnxt.DBC_DBC_TYPE_SQ, self.sq["prod"] & 15, (self.sq["prod"] // 16) & 1)
-    cqe = bnxt.struct_cq_req.from_buffer_copy(self._poll(timeout_ms))
-    assert cqe.status == 0
@@ -1,118 +0,0 @@
-#!/usr/bin/env python3
-"""Send and validate one RDMA WRITE between two Broadcom BNXT hosts.
-
-This follows ``extra/mlx_driver/connect.py``: sync the driver, start the remote
-endpoint over SSH, exchange QP/GID/MAC/MR metadata, move both RC QPs to RTS,
-write bytes into the remote MR, and verify the bytes on the remote host.
-
-Both PCI functions must be unbound from bnxt_en/bnxt_re first.
-"""
-import json
-import os
-import subprocess
-import sys
-from typing import Any, IO
-
-TINYGRAD = os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))
-sys.path.insert(0, TINYGRAD)
-
-from extra.bnxt_driver.bnxtdev import BNXTDev, BNXTQP
-from tinygrad.runtime.support.system import PCIDevice
-
-REMOTE_HOST = os.getenv("REMOTE_HOST", "192.168.52.213")
-REMOTE_USER = os.getenv("REMOTE_USER", "nimlgen")
-LOCAL_PCI = os.getenv("BNXT_PCI", "0000:41:00.0")
-REMOTE_PCI = os.getenv("REMOTE_PCI", "0000:41:00.0")
-LOCAL_IP = os.getenv("LOCAL_IP", "10.0.200.5")
-REMOTE_IP = os.getenv("REMOTE_IP", "10.0.200.6")
-MESSAGE = os.getenv("RDMA_MESSAGE", "Test message, rdma works!").encode()
-REMOTE = f"{REMOTE_USER}@{REMOTE_HOST}"
-SSH = ["ssh", "-o", "BatchMode=yes", "-o", "ConnectTimeout=10", "-o", "StrictHostKeyChecking=accept-new", REMOTE]
-SYNC_FILES = ("tinygrad/runtime/autogen/bnxt.py", "tinygrad/runtime/support/system.py",
-              "extra/bnxt_driver/bnxtdev.py", "extra/bnxt_driver/connect.py")
-
-def read_json(stream:IO[str], what:str) -> dict[str, Any]:
-  for line in iter(stream.readline, ""):
-    print(f"  [remote] {line}", end="")
-    try: value = json.loads(line)
-    except json.JSONDecodeError: continue
-    if isinstance(value, dict): return value
-  raise RuntimeError(f"remote exited before publishing {what}")
-
-def wait_line(stream:IO[str], text:str) -> str:
-  for line in iter(stream.readline, ""):
-    print(f"  [remote] {line}", end="")
-    if text in line: return line
-  raise RuntimeError(f"remote exited before reporting {text!r}")
-
-def send_line(stream:IO[str], value:str|dict[str, Any]):
-  stream.write((json.dumps(value) if isinstance(value, dict) else value) + "\n")
-  stream.flush()
-
-def qp_info(dev:BNXTDev, qp:BNXTQP) -> dict[str, Any]:
-  return {"qpn":qp.qpn, "mac":dev.mac.to_bytes(6, "big").hex(), "gid":dev.local_gid.hex()}
-
-def server():
-  dev = BNXTDev(PCIDevice("bnxt", os.getenv("BNXT_PCI", "0000:41:00.0")), ip=os.getenv("BNXT_IP", REMOTE_IP))
-  qp = BNXTQP(dev)
-  print(json.dumps(qp_info(dev, qp)), flush=True)
-
-  peer = json.loads(sys.stdin.readline())
-  qp.connect(peer["qpn"], bytes.fromhex(peer["gid"]), int(peer["mac"], 16))
-  print("connected", flush=True)
-
-  target, target_paddrs = dev.pci_dev.alloc_sysmem(0x1000)
-  target[:0x1000] = bytes(0x1000)
-  rkey = dev.register_mem(target_paddrs, 0x1000)
-  print(json.dumps({"target_addr":target_paddrs[0], "rkey":rkey}), flush=True)
-
-  assert sys.stdin.readline().strip() == "done"
-  received = bytes(target).rstrip(b"\0")
-  print(f"AS TEXT: {received.decode(errors='replace')!r}", flush=True)
-  print(json.dumps({"data":received.hex()}), flush=True)
-
-def sync_remote():
-  if os.getenv("SYNC", "1") == "0": return
-  print("syncing BNXT driver to remote")
-  subprocess.run(["rsync", "-azR", *SYNC_FILES, f"{REMOTE}:~/tinygrad/"], cwd=TINYGRAD, check=True)
-
-def start_remote() -> subprocess.Popen[str]:
-  print("booting remote")
-  command = (f"cd ~/tinygrad && sudo env PYTHONPATH=. PYTHONUNBUFFERED=1 BNXT_DEBUG={os.getenv('BNXT_DEBUG', '0')} "
-             f"BNXT_PCI={REMOTE_PCI} BNXT_IP={REMOTE_IP} python3 extra/bnxt_driver/connect.py --server")
-  return subprocess.Popen(SSH + [command], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=sys.stderr, text=True)
-
-def client():
-  assert 0 < len(MESSAGE) <= 0x1000
-  sync_remote()
-  remote = start_remote()
-  assert remote.stdin is not None and remote.stdout is not None
-  remote_info = read_json(remote.stdout, "QP information")
-  print("booting local")
-  dev = BNXTDev(PCIDevice("bnxt", LOCAL_PCI), ip=LOCAL_IP)
-  qp = BNXTQP(dev)
-
-  send_line(remote.stdin, qp_info(dev, qp))
-  wait_line(remote.stdout, "connected")
-  qp.connect(remote_info["qpn"], bytes.fromhex(remote_info["gid"]), int(remote_info["mac"], 16))
-  print("both QPs in RTS")
-
-  remote_target = read_json(remote.stdout, "MR information")
-  source, source_paddrs = dev.pci_dev.alloc_sysmem(0x1000)
-  source[:len(MESSAGE)] = MESSAGE
-  lkey = dev.register_mem(source_paddrs, 0x1000)
-  print(f"RDMA WRITE {len(MESSAGE)}B to remote phys 0x{remote_target['target_addr']:x}")
-  qp.rdma_write(remote_target["target_addr"], remote_target["rkey"], source_paddrs[0], lkey, len(MESSAGE))
-
-  send_line(remote.stdin, "done")
-  wait_line(remote.stdout, "AS TEXT")
-  result = read_json(remote.stdout, "RDMA result")
-  assert bytes.fromhex(result["data"]) == MESSAGE
-  print("RDMA WRITE data verified")
-
-  remote.stdin.close()
-  assert remote.wait() == 0
-  print("RDMA WRITE test complete")
-
-if __name__ == "__main__":
-  server() if "--server" in sys.argv else client()
@@ -1,46 +0,0 @@
-#!/usr/bin/env python3
-"""Local BNXT RoCEv2 RDMA WRITE loopback using the firmware's PHY loopback mode.
-
-The kernel bnxt_en/bnxt_re modules must be unloaded first.
-
-  sudo PYTHONPATH=. BNXT_PCI=0000:41:00.0 BNXT_IP=10.0.200.5 python3 extra/bnxt_driver/loopback.py
-"""
-import os
-import sys
-import time
-
-sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))
-
-from extra.bnxt_driver.bnxtdev import BNXTDev, BNXTQP
-from tinygrad.runtime.autogen import bnxt
-from tinygrad.runtime.support.system import PCIDevice
-
-BUF_SIZE = 0x1000
-BNXT_PCI = os.getenv("BNXT_PCI", "0000:41:00.0")
-BNXT_IP = os.getenv("BNXT_IP", "10.0.200.5")
-
-if __name__ == "__main__":
-  print(f"[init] BNXT at {BNXT_PCI}")
-  dev = BNXTDev(PCIDevice("bnxt", BNXT_PCI), ip=BNXT_IP)
-  tx_qp, rx_qp = BNXTQP(dev), BNXTQP(dev)
-  print(f"[init] loopback-connect TX QP 0x{tx_qp.qpn:x} <-> RX QP 0x{rx_qp.qpn:x}")
-  tx_qp.connect(rx_qp.qpn, dev.local_gid, dev.mac)
-  rx_qp.connect(tx_qp.qpn, dev.local_gid, dev.mac)
-
-  src, src_paddrs = dev.pci_dev.alloc_sysmem(BUF_SIZE)
-  dst, dst_paddrs = dev.pci_dev.alloc_sysmem(BUF_SIZE)
-  message = b"Hello from BNXT RoCE PHY loopback!"
-  src[:BUF_SIZE], dst[:BUF_SIZE] = bytes(BUF_SIZE), bytes(BUF_SIZE)
-  src[:len(message)] = message
-  lkey = dev.register_mem(src_paddrs, BUF_SIZE)
-  rkey = dev.register_mem(dst_paddrs, BUF_SIZE)
-
-  print("[loopback] enabling local PHY loopback")
-  dev.hwrm("port_phy_cfg", port_id=dev.port_id, enables=bnxt.PORT_PHY_CFG_REQ_ENABLES_LPBK, lpbk=bnxt.PORT_PHY_CFG_REQ_LPBK_LOCAL)
-  time.sleep(1)
-  tx_qp.rdma_write(dst_paddrs[0], rkey, src_paddrs[0], lkey, len(message))
-  got = bytes(dst[:len(message)])
-  print(f"[result] {got!r}")
-  assert got == message
-  print("BNXT RoCE PHY loopback RDMA WRITE passed")
-  dev.hwrm("port_phy_cfg", port_id=dev.port_id, enables=bnxt.PORT_PHY_CFG_REQ_ENABLES_LPBK, lpbk=bnxt.PORT_PHY_CFG_REQ_LPBK_NONE)
@@ -462,7 +462,7 @@ def test_matmul():
    lds = UOp.placeholder((lds_size,), dtypes.uint8, 0, AddrSpace.LOCAL)
    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.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, 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]
  linear = c.schedule_linear()

@@ -125,7 +125,7 @@ def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp,
                  arg=KernelInfo(f"mxfp4_gemm_{M}_{N}_{K}",
                                 estimates=Estimates(ops=2*M*N*K, mem=(M*half_k+N*half_k)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
  insts = build_kernel(M, N, K, tile_m, tile_n)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in insts))))

 def _mxfp4_gemm_quantized(a_q:Tensor, b_q:Tensor, scale_a:Tensor, scale_b:Tensor) -> Tensor:
  M, half_k = a_q.shape
@@ -215,7 +215,7 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
  k = UOp.range(K, 0, AxisType.REDUCE)
  mul = (A.flatten().index((m*UOp.const(K)+k))*
         B.flatten().index((k*UOp.const(N)+n))).cast(dtypes.float32)
-  red = mul.reduce(k, arg=Ops.ADD).cast(C.dtype)
+  red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype)
  store = C.flatten().index((m*UOp.const(N)+n)).store(red).end(m, n)
  return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))

@@ -53,7 +53,7 @@ def _ggather_bwd(gradient:UOp, kernel:UOp) -> tuple:
    g, m, j, jo, ji = _kv_ranges(Gk, M, Dk, _blk_for(Dk))
    row = idx.index(g, m).cast(dtypes.weakint)
    val = gout.index(g, m, j).load().cast(dtypes.float32)
-    atomic = UOp(Ops.CUSTOM, src=(gtab.index(g, row, j), val), arg=(atomic_str, dtypes.void))
+    atomic = UOp(Ops.CUSTOM, dtypes.void, (gtab.index(g, row, j), val), arg=atomic_str)
    return atomic.end(g, m, jo, ji).sink(arg=KernelInfo(name=f"ggather_bwd_{M}_{Dk}", opts_to_apply=()))
  grad_table = Tensor.custom_kernel(gt, go, Tensor(idx_u, device=dev), fxn=_bwd_kernel)[0]
  return (None, grad_table.cast(table_u.dtype).uop, None)
@@ -223,7 +223,7 @@ def test_matmul():
    lds = UOp.placeholder((lds_size,), dtypes.uint8, 0, AddrSpace.LOCAL)
    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*2*3)))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, 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]
  linear = c.schedule_linear()
@@ -16,12 +16,9 @@ def _do_reset_device(pci_bus): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/dev
 def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")

 def cmd_remove_module(args):
-  modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia"] if args.backend == "nv" else ["amdgpu"]
+  modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia", "ast"] if args.backend == "nv" else ["amdgpu"]
  to_unload = [m for m in modules if _is_module_loaded(m)]
  if not to_unload: print("Kernel modules are not loaded")
-  elif getattr(args, "expect", False):
-    print(f"Kernel modules are loaded: {to_unload}")
-    sys.exit(1)
  else:
    print("Removing kernel modules:", ", ".join(to_unload))
    try: subprocess.run(["sudo", "modprobe", "-r", *to_unload], check=True)
@@ -63,19 +60,17 @@ def cmd_show_pids(args):

 def cmd_kill_pids(args):
  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
-  use_sudo = not getattr(args, "sudoless", False)

  for dev in devs:
    for i in range(128):
      if i > 0: time.sleep(0.2)

      try:
-        try: pid = subprocess.check_output((['sudo'] if use_sudo else []) +
-                                           ['lsof', temp(f'{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})")
-        subprocess.run((['sudo'] if use_sudo else []) + ['kill', '-9', pid], check=True)
+        subprocess.run(['sudo', 'kill', '-9', pid], check=True)
      except subprocess.CalledProcessError as e:
        print(f"Failed to kill process for device {dev}: {e}", file=sys.stderr)

@@ -84,7 +79,6 @@ def add_common_commands(parent_subparsers):
  p_insmod.set_defaults(func=cmd_insert_module)

  p_rmmod = parent_subparsers.add_parser("rmmod", help="Remove a kernel module")
-  p_rmmod.add_argument("--expect", action="store_true", help="Just assert that module is already unloaded")
  p_rmmod.set_defaults(func=cmd_remove_module)

  p_reset = parent_subparsers.add_parser("reset", help="Reset a device")
@@ -97,20 +91,17 @@ def add_common_commands(parent_subparsers):

  p_reset = parent_subparsers.add_parser("kill_pids", help="Kill pids of processes using the device")
  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device")
-  p_reset.add_argument("--sudoless", action="store_true", help="Do not use sudo when detecting or killing pids")
  p_reset.set_defaults(func=cmd_kill_pids)

 if __name__ == "__main__":
  parser = argparse.ArgumentParser()
  backend_subparsers = parser.add_subparsers(dest="backend", required=True, metavar="{nv,amd}", help="Hardware backend to target")

-  nv_parser = backend_subparsers.add_parser("nv", aliases=["NV"], help="NVIDIA GPUs")
-  nv_parser.set_defaults(backend="nv")
+  nv_parser = backend_subparsers.add_parser("nv", help="NVIDIA GPUs")
  nv_commands = nv_parser.add_subparsers(dest="command", required=True)
  add_common_commands(nv_commands)

-  amd_parser = backend_subparsers.add_parser("amd", aliases=["AMD"], help="AMD GPUs")
-  amd_parser.set_defaults(backend="amd")
+  amd_parser = backend_subparsers.add_parser("amd", help="AMD GPUs")
  amd_commands = amd_parser.add_subparsers(dest="command", required=True)
  add_common_commands(amd_commands)

@@ -19,7 +19,7 @@ from tinygrad.runtime.support.hcq import FileIOInterface, HCQBuffer, MMIOInterfa
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
 from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_pmc
 from tinygrad.runtime.support.system import PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
-from tinygrad.runtime.support.usb import USB3, usb_ib, usb_push, usb_arm_bytes, pm_usb_stage, pm_usb_hostio, pm_usb_bufferize
+from tinygrad.runtime.support.usb import USB3
 from tinygrad.runtime.support.memory import AddrSpace, BumpAllocator
 from tinygrad.runtime.ops_amd import SQTT, SQTT_ITRACE_SE_MASK, SQTT_LIMIT_SE, SQTT_SIMD_SEL, SQTT_TOKEN_EXCLUDE, PMC
 from tinygrad.runtime.ops_amd import EVENT_INDEX_PARTIAL_FLUSH, WAIT_REG_MEM_FUNCTION_EQ, WAIT_REG_MEM_FUNCTION_NEQ, WAIT_REG_MEM_FUNCTION_GEQ
@@ -37,7 +37,7 @@ class PM4Ops(FastEnum):
  RELEASE_MEM = auto(); DISPATCH_DIRECT = auto(); EVENT_WRITE = auto()  # noqa: E702

 def pkt3(ctx, op:PM4Ops, *vals):
-  return UOp(Ops.INS, arg=(op, dtypes.void), src=tuple(UOp.const(x, dtypes.uint32)
+  return UOp(Ops.INS, arg=op, src=tuple(UOp.const(x, dtypes.uint32)
    for x in (ctx.pm4.PACKET3(getattr(ctx.pm4, f"PACKET3_{op.name}"), len(vals) - 1), *vals)))

 def wreg(ctx, reg:AMDReg, *args:sint, **kwargs:int):
@@ -87,7 +87,7 @@ def release_mem(ctx, address=0x0, value=0, data_sel=0, int_sel=2, ctxid=0, cache

 def memory_barrier(ctx):
  pf = '' if ctx.nbio.version[0] == 2 else '0' if ctx.nbio.version[:2] != (7, 11) else '1'
-  return UOp(Ops.LINEAR, src=(
+  return UOp(Ops.LINEAR, dtypes.void, (
    wait_reg_mem(ctx, reg=getattr(ctx.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_REQ').addr[0],
                 reg_done=getattr(ctx.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff),
    acquire_mem(ctx)))
@@ -135,25 +135,22 @@ def pm4_program(ctx, call, prg):
    wreg(ctx, ctx.gc.regCOMPUTE_START_X, 0, 0, 0, *(info.local_size or (1, 1, 1)), 0, 0),
    pkt3(ctx, PM4Ops.DISPATCH_DIRECT, *info.global_size, dispatch_init),
    pkt3(ctx, PM4Ops.EVENT_WRITE, ctx.pm4.EVENT_TYPE(ctx.soc.CS_PARTIAL_FLUSH) | ctx.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH))]
-  return UOp(Ops.LINEAR, src=tuple(ins))
+  return UOp(Ops.LINEAR, dtypes.void, tuple(ins))

 pm_pm4_opsel = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), pm4_program),

-  (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), pm4_wait),
-  (UPat(Ops.INS, arg=("barrier", dtypes.void)), pm4_barrier),
-  (UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),)), pm4_timestamp),
-  (UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
+  (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val"))), pm4_wait),
+  (UPat(Ops.INS, arg="barrier"), pm4_barrier),
+  (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),)), pm4_timestamp),
+  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
 ])

-def queue_ptrs(devs, qname:str, q:AMDQueueDesc) -> tuple[UOp, ...]:
-  return tuple(UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"{qname}_{n}")
-               for n, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
-
 def pm4_submit(ctx, lin):
  # ensure compute queues are allocated
  for d in (devs:=ctx.devs): q = Device[d].compute_queue
-  ring, wptr, doorbell, put_ptr = queue_ptrs(devs, "COMPUTE:0", q)
+  ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COMPUTE:0_{name}")
+    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))

  # the host fence at the start of the batch guarantees the ib is free to reuse
  size_dw = sum(len(ins.src) for ins in lin.src)
@@ -207,10 +204,10 @@ def sdma_timestamp(ctx, ins, dst):
 pm_sdma_opsel = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy),

-  (UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda: UOp(Ops.NOOP)),
-  (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val")), name="ins"), sdma_wait),
-  (UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),), name="ins"), sdma_timestamp),
-  (UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val")), name="ins"), sdma_store),
+  (UPat(Ops.INS, arg="barrier"), lambda: UOp(Ops.NOOP, dtypes.void, ())),
+  (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val")), name="ins"), sdma_wait),
+  (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),), name="ins"), sdma_timestamp),
+  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val")), name="ins"), sdma_store),
 ])

 def sdma_submit(cmdbuf, devs):
@@ -219,7 +216,8 @@ def sdma_submit(cmdbuf, devs):

  # the sdma queue's ring and its host-side ring/write/put pointers
  for d in devs: q = Device[d].sdma_queue(0)
-  ring, wptr, doorbell, put_ptr = queue_ptrs(devs, "COPY:0", q)
+  ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COPY:0_{name}")
+    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))

  # sdma needs the cmdbuf contiguous: if it won't fit before the ring end, restart at 0 and zero the tail
  put_b = put_ptr.index(zero)
@@ -246,32 +244,15 @@ def sdma_submit(cmdbuf, devs):
 pm_sdma_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"),
  lambda ctx, lin: sdma_submit(make_cmdbuf(lin, ctx.devs), ctx.devs))])

-# *****************
-# USB submit
-
-def amd_usb_submit(ctx, lin):
-  for d in ctx.devs: q = Device[d].compute_queue if (comp:=ctx.qname.startswith("COMPUTE")) else Device[d].sdma_queue(0)
-
-  if nb:=usb_arm_bytes(ctx.pre, Device[ctx.devs[0]].iface.usb_sram):
-    poke = (ctx.sdma.SDMA_OP_WRITE, *data64_le(Device[ctx.devs[0]].iface.cq_buf.va_addr + 12), 0, 0)
-    lin = lin.replace(src=lin.src + (UOp(Ops.INS, arg=("poke", dtypes.void), src=tuple(UOp.const(x, dtypes.uint32) for x in poke)),))
-
-  ib_host, ib_gpu, pkt_dw = usb_ib(ctx.devs, lin, 32 if comp else 0x100, nb)
-  pkt = (ctx.pm4.PACKET3(ctx.pm4.PACKET3_INDIRECT_BUFFER,2),*data64_le(ib_gpu.getaddr(ctx.devs)),pkt_dw|ctx.pm4.INDIRECT_BUFFER_VALID) if comp else ()
-  return usb_push(ctx.devs, *queue_ptrs(ctx.devs, ctx.qname, q), ib_host, ib_gpu, pkt, 4 if comp else 1)
-
-pm_usb_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"), amd_usb_submit)])
-
@dataclass(frozen=True)
 class AMDEncodeCtx:  # encode-time constants for one queue: devs (every cmdbuf address resolves into these) + gfx version + packet/ip modules
  devs: tuple[str, ...]; target: tuple[int, ...]; pm4: Any; sdma: Any; soc: Any  # noqa: E702
-  gc: AMDIP; nbio: AMDIP; xccs: int; max_copy_size: int; tmpring_size: Callable; qname: str; pre: UOp # pre: the queue before opsel
+  gc: AMDIP; nbio: AMDIP; xccs: int; max_copy_size: int; tmpring_size: Callable  # noqa: E702

 def encode_queue(q:UOp) -> UOp|None:
  d = Device[(devs:=to_tuple(q.arg[0]))[0]]
-  ctx = AMDEncodeCtx(devs, d.target, d.pm4, d.sdma, d.soc, d.gc, d.nbio, d.xccs, d.max_copy_size, d.tmpring_size, q.arg[1], q)
-  opsel = pm_pm4_opsel if (comp:=q.arg[1].startswith("COMPUTE")) else pm_sdma_opsel
-  submit = d.pm_submit if d.pm_submit is not None else (pm_pm4_submit if comp else pm_sdma_submit)
+  ctx = AMDEncodeCtx(devs, d.target, d.pm4, d.sdma, d.soc, d.gc, d.nbio, d.xccs, d.max_copy_size, d.tmpring_size)
+  opsel, submit = (pm_pm4_opsel, pm_pm4_submit) if q.arg[1].startswith("COMPUTE") else (pm_sdma_opsel, pm_sdma_submit)
  return submit.rewrite(graph_rewrite(q, opsel + pm_flatten_linear, walk=True, ctx=ctx, name=f"{q.arg[1]} opsel"), ctx)

@dataclass(frozen=True)
@@ -280,11 +261,10 @@ class AMDProgramData:
  private_segment_size:int; kernargs_segment_size:int; kernargs_alloc_size:int
  enable_dispatch_ptr:int; enable_private_segment_sgpr:int

-_amd_program_cache:dict[tuple[bytes, tuple[str, ...]], UOp] = {}
+_amd_program_cache:dict[tuple[bytes,str], tuple[AMDProgramData,bytes]] = {}
 def amd_build_program(prg:UOp) -> UOp:
  dev = Device[to_tuple(prg.device)[0]] # TODO: rm this
-  # key on the full device tuple: the same lib can be built for different device sets, each needs its own program buffer
-  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, to_tuple(prg.device)))) is None:
+  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, dev.device))) is None:
    image, sections, relocs = elf_loader(lib)
    rodata = next(sh.header.sh_addr for sh in sections if sh.name == ".rodata")
    for off, sym, typ, addent in relocs:
@@ -301,14 +281,13 @@ def amd_build_program(prg:UOp) -> UOp:
      wave32=bool(desc.kernel_code_properties & 0x400), private_segment_size=desc.private_segment_fixed_size, kernargs_segment_size=desc.kernarg_size,
      kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0), enable_dispatch_ptr=edp,
      enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER)
-    image = bytes(image).ljust(round_up(len(image), 4), b"\x00") # the program is uploaded as whole dwords
    buf = UOp.placeholder((len(image),), dtypes.uint8, next(UOp.unique_num), device=prg.device).rtag("program")
-    cached = _amd_program_cache[key] = prg.replace(src=(buf.after(make_binary_patch(buf, image)),), arg=(data, prg.arg))
+    cached = _amd_program_cache[key] = prg.replace(src=(buf.after(make_binary_patch(buf, bytes(image))),), arg=(data, prg.arg))
  return cached

 class AMDAllocator(HCQAllocator['AMDDevice']):
  def __init__(self, dev:AMDDevice):
-    super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb)
+    super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb())

  def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
    return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_copy_queue)
@@ -544,7 +523,8 @@ class PCIIface(PCIIfaceBase):
      cq = d.compute_queue
      for b in (cq.put_value, cq.read_ptr, cq.write_ptr): b._buf.view.view(fmt='Q')[0] = 0
      d.iface.dev_impl.gfx.setup_ring(*cq.params)
-      d.signal('timeline')._buf.cpu_view().view(fmt='Q')[0] = d.signal('value', 1, device="CPU")._buf.cpu_view().view(fmt='Q')[0] - 1
+      d.signal('timeline')._buf.cpu_view().mv.cast('Q')[0] = \
+        d.signal('value', 1).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] - 1

  def sleep(self, timeout):
    if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
@@ -558,32 +538,6 @@ class PCIIface(PCIIfaceBase):

  def device_fini(self): self.dev_impl.fini()

-class USBIface(PCIIface):
-  def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called
-    if dev_id >= len(visible:=hcq_filter_visible_devices(USB3.list_devices(0xADD1, 0x0001) + USB3.list_devices(0x3801, 0x0001), "AMD")):
-      raise RuntimeError(f"AMD:{dev_id} does not exist ({pluralize('device', len(visible))} available)")
-    self.dev, self.pci_dev, self.vram_bar, self.count = dev, USBPCIDevice("AM", *visible[dev_id]), 0, len(visible)
-    self.dev_impl = AMDev(self.pci_dev)
-    self._compute_props()
-    self.sram = self._dma_region(ctrl_addr=0xf000, sys_addr=0x200000, size=0x80000)
-    self.cq_buf = self._dma_region(ctrl_addr=0xb800, sys_addr=0x822000, size=0x1000) # +12 is the dword that releases an armed read
-    self.usb_handle = unwrap(ctypes.cast(self.pci_dev.usb.usb.handle, ctypes.c_void_p).value)
-
-  def _dma_region(self, ctrl_addr, sys_addr, size):
-    region = self.dev_impl.mm.map_range(vaddr:=self.dev_impl.mm.alloc_vaddr(size=size), size, [(sys_addr, size)], aspace=AddrSpace.SYS, uncached=True)
-    return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(region, has_cpu_mapping=False), view=self.pci_dev.dma_view(ctrl_addr, size), owner=self.dev)
-
-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, **kwargs) -> HCQBuffer:
-    # everything, even host-style signals, lives in vram: gpu writes into the bridge's own memory collide with an armed 0xF2 read stream
-    return super().alloc(size, host=False, uncached=uncached, cpu_access=cpu_access or host, contiguous=contiguous, force_devmem=True, **kwargs)
-
-  def sleep(self, timeout): pass
-
-  # we don't own the sram region, so the buffer never frees it
-  @functools.cached_property
-  def usb_sram(self) -> Buffer:
-    return Buffer(self.dev.device, (b:=self.sram).size, dtypes.uint8, options=BufferSpec(external_ptr=b.va_addr, nolru=True)).allocate(opaque=b)
-
 def _mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface,), {})

 class AMDDevice(HCQ2Compiled):
@@ -594,21 +548,19 @@ class AMDDevice(HCQ2Compiled):
    # encoding of cmdbuf
    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),)), encode_queue),
  ])
-  pm_submit: PatternMatcher|None = None

  timestamp_divider = 100.0  # AMD GPU clock: ticks/us
  max_scratch_psize = 0

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

  def device_props(self): return self.iface.props

  def is_am(self) -> bool: return isinstance(self.iface, (PCIIface,))
+  def is_usb(self) -> bool: return False

  def __init__(self, device:str=""):
    self.iface = self._select_iface(device)
-    self.is_usb = isinstance(self.iface, USBIface)
-    if self.is_usb: self.rt_nbytes = 4 << 20

    self.target:tuple[int, ...] = ((trgt:=self.iface.props['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
    self.arch = "gfx%d%x%x" % self.target
@@ -633,7 +585,7 @@ class AMDDevice(HCQ2Compiled):

    self.is_aql = getenv("AMD_AQL", int(self.xccs > 1))
    if self.is_aql:
-      self.pm4_ibs = self.iface.alloc(0x2000 if self.is_usb else (16 << 20), uncached=True, cpu_access=True)
+      self.pm4_ibs = self.iface.alloc(0x2000 if self.is_usb() else (16 << 20), uncached=True, cpu_access=True)
      self.pm4_ib_alloc = BumpAllocator(self.pm4_ibs.size, wrap=True)

    self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
@@ -646,10 +598,6 @@ class AMDDevice(HCQ2Compiled):
    self.max_private_segment_size = 0
    self.pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, tag="scratch", name="b"), lambda ctx, b: ctx[0].scratch_buffer(b.max_numel()))]) + self.pm_bufferize

-    if self.is_usb:
-      self.pm_bufferize = pm_usb_bufferize + self.pm_bufferize
-      self.pm_stage_copy, self.pm_host_lower, self.pm_submit = pm_usb_stage, pm_usb_hostio, pm_usb_submit
-
    self.pmc_enabled:bool = PROFILE > 0 and PMC > 0
    if self.pmc_enabled:
      self.iface.require_profile_mode()
@@ -710,7 +658,7 @@ class AMDDevice(HCQ2Compiled):
    wg_data_size = round_up((vgpr_size_per_cu + sgrp_size_per_cu + lds_size_per_cu + hwreg_size_per_cu) * self.cu_cnt, mmap.PAGESIZE)
    ctl_stack_size = round_up((12 if self.target[0] != 9 else 8) * self.wave_cnt + 8 + 40, mmap.PAGESIZE)
    return self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL if self.is_aql else kfd.KFD_IOC_QUEUE_TYPE_COMPUTE,
-      0x2000 if self.is_usb else (16 << 20), eop_buffer_size=0x1000,
+      0x2000 if self.is_usb() else (16 << 20), eop_buffer_size=0x1000,
      ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size,
      debug_memory_size=round_up(self.wave_cnt * 32, 64))

@@ -718,7 +666,7 @@ class AMDDevice(HCQ2Compiled):
    if getenv("AMD_DISABLE_SDMA"): return None
    if idx in self.sdma_queues: return self.sdma_queues[idx]
    with contextlib.suppress(OSError):
-      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x2000 if self.is_usb else (16 << 20), idx=idx)
+      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20), idx=idx)
    return self.sdma_queues.get(idx, None)

  def tmpring_size(self, private_segment_size):
@@ -1,6 +1,8 @@
 from __future__ import annotations
 import functools, pathlib
-from tinygrad import Tensor
+from dataclasses import replace
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import shape_to_shape_arg
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler

 FP8_MAX = 448.0
@@ -10,7 +12,7 @@ NUM_WG, THREADS_PER_WG = 1024, 256
@functools.cache
 def _local_abs_max_fxn(x_p, device):
  x = Tensor(x_p, device=device)
-  inner = Tensor(x.uop.src[0]) if x.uop.axis is not None else x  # the per-shard view of the flat param
+  inner = Tensor(x.uop.replace(src=(shape_to_shape_arg(x.uop.shard_shape),), arg=replace(x.uop.arg, axis=None))) if x.uop.axis is not None else x
  return (inner.abs().max(),)

 def local_abs_max(x:Tensor) -> Tensor:
@@ -50,7 +50,7 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:
  else: raise NotImplementedError(f"no atomic max for device {device}")
  amax_idx = amax_out.reshape((1,)).index(UOp.const(0))
  max_val = lds[0].load()
-  atomic = UOp(Ops.CUSTOM, src=(amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=(atomic_arg, dtypes.void))
+  atomic = UOp(Ops.CUSTOM, dtypes.void, (amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=atomic_arg)
  return atomic.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=()))

@functools.cache
@@ -12,7 +12,7 @@ def _custom_quantize_mxfp4(row_fp4:UOp, row_scale:UOp, col_fp4:UOp, col_scale:UO
  mem = M*N*2 + M*N + M*N//16  # read bf16, write row+col fp4 + e8m0
  outputs = (row_fp4, row_scale, col_fp4, col_scale)
  sink = UOp.sink(*(o.base for o in outputs), x.base,
-                  *(UOp(Ops.CUSTOM, src=(o.base.index(0),), arg=("", dtypes.void)) for o in outputs),
+                  *(UOp(Ops.CUSTOM, dtypes.void, (o.base.index(0),), arg="") for o in outputs),
                  UOp.special(256, "lidx0"), UOp.special(M//128, "gidx0"), UOp.special(N//64, "gidx1"),
                  arg=KernelInfo(name, estimates=Estimates(ops=12*M*N, mem=mem)))
  src = (pathlib.Path(__file__).parent/"quantize_mxfp4.cpp").read_text()
@@ -3,7 +3,7 @@ import os
 # TODO: there is a timing bug without this
 os.environ["AMD_AQL"] = "1"

-from tinygrad import Tensor, Device, GlobalCounters, Context, dtypes
+from tinygrad import Tensor, Device, GlobalCounters, Context
 from tinygrad.helpers import getenv, DEV
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
@@ -37,7 +37,7 @@ def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, **kwargs)
    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.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, 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]
  linear = out.schedule_linear()
@@ -5,9 +5,9 @@ from tinygrad.helpers import getenv, DEBUG

 # https://github.com/facebookresearch/llama/blob/1076b9c51c77ad06e9d7ba8a4c6df775741732bd/llama/model.py#L47
 def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
-  freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2, dtype=dtypes.float32)[:(dim // 2)] / dim))
-  freqs = Tensor.arange(end, dtype=dtypes.float32).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
-  return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).cast(dtypes.default_float).reshape(1, end, 1, dim//2, 2)
+  freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2)[:(dim // 2)] / dim))
+  freqs = Tensor.arange(end).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
+  return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).reshape(1, end, 1, dim//2, 2)

 # matches meta, non hugging face weights
 # (a+i*b) * (c+i*d) = (ac-bd) + i*(ad+bc)
@@ -1,12 +1,10 @@
 from tinygrad import Tensor
-from tinygrad.helpers import NUM_CPU_THREADS
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.nn.onnx import OnnxRunner, OnnxValue
 import numpy as np
 import onnxruntime as ort
 ort_options = ort.SessionOptions()
 ort_options.log_severity_level = 3
-ort_options.intra_op_num_threads = NUM_CPU_THREADS.value

 def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
  """
@@ -1,4 +1,4 @@
-import os, subprocess, sys, shlex, pickle
+import os, subprocess, sys, shlex
 from pathlib import Path
 from tinygrad.helpers import temp, getenv

@@ -23,8 +23,5 @@ if __name__ == "__main__":
      # AM_RESET=1 gets a clear trace, does not work on mi300 machines
      subprocess.run([sys.executable, *shlex.split(test)], cwd=EXAMPLES_DIR.parent.parent.parent,
                     env={**os.environ, "DEV":"AMD", "AM_RESET":"1" if not arch.startswith("gfx9") else "0", "VIZ":"-2", "PYTHONPATH":"."})
-      with open(PROFILE_PATH, "rb") as f: events = pickle.load(f)
-      with open(PROFILE_PATH, "wb") as f:
-        pickle.dump([e for e in events if type(e).__name__ in {"ProfilePMCEvent", "ProfileSQTTEvent", "ProfileProgramEvent"}], f)
      PROFILE_PATH.rename(dest:=EXAMPLES_DIR/arch/f"profile_{name}_run_{i}.pkl")
      print(f"saved SQTT trace to {dest}")
@@ -192,7 +192,7 @@ def unpack_insts(viz_data, i:int, j:int, data:dict) -> dict:
    prev_instr = max(prev_instr, e.time + e.dur)
  summary = [{"label":"Total Cycles", "value":w.end_time-w.begin_time}, {"label":"SE", "value":w.se}, {"label":"CU", "value":w.cu},
             {"label":"SIMD", "value":w.simd}, {"label":"Wave ID", "value":w.wave_id}, {"label":"Run number", "value":data["run_number"]}]
-  return {"rows":[tuple(v.values()) for v in rows.values()], "cols":columns, "metadata":[summary],"ref":viz_data.ref_map.get(data["prg"].profile_key)}
+  return {"rows":[tuple(v.values()) for v in rows.values()], "cols":columns, "metadata":[summary], "ref":viz_data.ref_map.get(data["prg"].name)}

 def print_data(data:dict) -> None:
  from tabulate import tabulate
@@ -19,33 +19,16 @@ def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None
@functools.cache
 def custom_fused_qkv_rope_forward(q:UOp, k:UOp, v:UOp, xqkv:UOp, freqs_cis:UOp,
                                  device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
-  group_size = H // H_KV
-  q, k, v = q.reshape(B, N, H, D), k.reshape(B, N, H_KV, D), v.reshape(B, N, H_KV, D)
-  xqkv = xqkv.reshape(B, N, H_KV, group_size + 2, D)
-  b, n = UOp.range(B, 0), UOp.range(N, 1)
-  pair = UOp.range(D // 2, 2)
-  even = pair * 2
-  c = freqs_cis[0, n, 0, pair, 0].cast(dtypes.float)
-  s = freqs_cis[0, n, 0, pair, 1].cast(dtypes.float)
-  ordered:UOp|None = None
-  for kvh in range(H_KV):
-    q_out, k_out, v_out = (x.after(ordered) if ordered is not None else x for x in (q, k, v))
-    x_in = xqkv.after(ordered) if ordered is not None else xqkv
-    stores:list[UOp] = []
-    for rep in range(group_size):
-      a = x_in[b, n, kvh, rep, even].cast(dtypes.float)
-      bb = x_in[b, n, kvh, rep, even + 1].cast(dtypes.float)
-      h = kvh * group_size + rep
-      stores += [q_out[b, n, h, even].store((a * c - bb * s).cast(q.dtype)), q_out[b, n, h, even + 1].store((a * s + bb * c).cast(q.dtype))]
-    a = x_in[b, n, kvh, group_size, even].cast(dtypes.float)
-    bb = x_in[b, n, kvh, group_size, even + 1].cast(dtypes.float)
-    stores += [k_out[b, n, kvh, even].store((a * c - bb * s).cast(k.dtype)),
-               k_out[b, n, kvh, even + 1].store((a * s + bb * c).cast(k.dtype)),
-               v_out[b, n, kvh, even].store(x_in[b, n, kvh, group_size + 1, even]),
-               v_out[b, n, kvh, even + 1].store(x_in[b, n, kvh, group_size + 1, even + 1])]
-    ordered = UOp.group(*stores)
-  assert ordered is not None
-  return ordered.end(pair, n, b).sink(arg=KernelInfo(name="fused_qkv_rope_forward"))
+  code = (pathlib.Path(__file__).parent / "fused_qkv_rope.cpp").read_text()
+  threads = 256
+  thread_idx = UOp.special(threads, "lidx0")
+  block_idx_x, block_idx_y = UOp.special(B, "gidx0"), UOp.special(N, "gidx1")
+  sink = UOp.sink(q.base, k.base, v.base, xqkv.base, freqs_cis.base, thread_idx, block_idx_x, block_idx_y,
+                  arg=KernelInfo(name="fused_qkv_rope_forward"))
+  compile_args = ["-std=c++20", "-ffast-math", f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}",
+                  f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DTHREADS_PER_BLOCK={threads}"]
+  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))

@functools.cache
 def custom_fused_qkv_rope_backward(dxqkv:UOp, dq:UOp, dk:UOp, dv:UOp, freqs_cis:UOp,
@@ -126,6 +109,49 @@ def fused_qkv_rope(xqkv:Tensor, freqs_cis:Tensor, n_heads:int, n_kv_heads:int, h
 def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
  return _sharded_empty(ref.shape, ref, axis)

+@functools.cache
+def _windowed_lse(xq:Tensor, xk:Tensor, sinks, W:int) -> Tensor:
+  B, N, H, hd = xq.shape
+  H_KV = xk.shape[2]; R = H // H_KV; nb = N // W; sm = hd ** -0.5
+  q = xq.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
+  k = xk.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
+  k_prev = k.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
+  sc_d = (q @ k.transpose(-1, -2)) * sm
+  sc_p = (q @ k_prev.transpose(-1, -2)) * sm
+  li, lj = Tensor.arange(W).reshape(W, 1), Tensor.arange(W).reshape(1, W)
+  pv = (Tensor.arange(nb).reshape(nb, 1, 1) >= 1)
+  sc_d = (lj <= li).where(sc_d, -float("inf"))
+  sc_p = ((li < lj) & pv).where(sc_p, -float("inf"))
+  m = sc_d.max(-1, keepdim=True).maximum(sc_p.max(-1, keepdim=True))
+  if sinks is not None: m = m.maximum(sinks.reshape(1, H_KV, R, 1, 1, 1).float())
+  denom = (sc_d - m).exp().sum(-1, keepdim=True) + (sc_p - m).exp().sum(-1, keepdim=True)
+  if sinks is not None: denom = denom + (sinks.reshape(1, H_KV, R, 1, 1, 1).float() - m).exp()
+  return (m + denom.log()).reshape(B, H, N).unsqueeze(2)  # (B, H, 1, N), matches saved l_vec
+
+def _windowed_delta(xq:Tensor, xk:Tensor, xv:Tensor, do:Tensor, sinks, W:int) -> Tensor:
+  B, N, H, hd = xq.shape
+  H_KV = xk.shape[2]; R = H // H_KV; nb = N // W; sm = hd ** -0.5
+  q = xq.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
+  k = xk.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
+  v = xv.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
+  dob = do.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
+  k_prev = k.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
+  v_prev = v.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
+  sc_d = (q @ k.transpose(-1, -2)) * sm
+  sc_p = (q @ k_prev.transpose(-1, -2)) * sm
+  li, lj = Tensor.arange(W).reshape(W, 1), Tensor.arange(W).reshape(1, W)
+  pv = (Tensor.arange(nb).reshape(nb, 1, 1) >= 1)
+  sc_d = (lj <= li).where(sc_d, -float("inf"))
+  sc_p = ((li < lj) & pv).where(sc_p, -float("inf"))
+  m = sc_d.max(-1, keepdim=True).maximum(sc_p.max(-1, keepdim=True))
+  if sinks is not None: m = m.maximum(sinks.reshape(1, H_KV, R, 1, 1, 1).float())
+  e_d, e_p = (sc_d - m).exp(), (sc_p - m).exp()
+  denom = e_d.sum(-1, keepdim=True) + e_p.sum(-1, keepdim=True)
+  if sinks is not None: denom = denom + (sinks.reshape(1, H_KV, R, 1, 1, 1).float() - m).exp()
+  o = ((e_d / denom) @ v) + ((e_p / denom) @ v_prev)
+  delta = (dob * o).sum(-1)
+  return delta.reshape(B, H, N).unsqueeze(2)
+
 def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink, window=0):
  def grad(dou:UOp, ker:UOp) -> tuple:
    do = Tensor(dou, device=dou.device)
@@ -134,6 +160,8 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
    xq = Tensor(ker.src[3], device=ker.src[3].device)
    xk = Tensor(ker.src[4], device=ker.src[4].device)
    xv = Tensor(ker.src[5], device=ker.src[5].device)
+    if window:
+      l_vec = _windowed_lse(xq, xk, Tensor(ker.src[6], device=ker.src[6].device) if has_sink else None, window)

    dq = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t)
    GROUP_SIZE = H_local // H_KV_local
@@ -144,6 +172,8 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
    # delta_vec = (do * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
    delta_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
    delta_vec, dq = Tensor.custom_kernel(delta_vec, dq, attn, do, fxn=functools.partial(custom_fa_backward_pre, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:2]
+    if window:
+      delta_vec = _windowed_delta(xq, xk, xv, do, Tensor(ker.src[6], device=ker.src[6].device) if has_sink else None, window)

    dq, dk_partial, dv_partial = Tensor.custom_kernel(dq, dk_partial, dv_partial, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, window=window))[:3]

@@ -269,9 +269,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p

    qo_tile<D, float> q_reg_fl;
    load<1, qo_tile<D, float>, _gl_QKVO>(q_reg_fl, g.Qg, {batch_idx, tile_idx, head_idx, 0});
-    #if !WINDOW
    mul(q_reg_fl, q_reg_fl, TEMPERATURE_SCALE);  // Use sqrtf for clarity
-    #endif
    copy(q_reg, q_reg_fl);
    transpose(q_reg_transposed, q_reg);

@@ -290,9 +288,6 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    zero(att_block[0]);
    transpose(k_reg_transposed, k_reg);
    mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
-    #if WINDOW
-    mul(att_block[0], att_block[0], TEMPERATURE_SCALE);
-    #endif
    __builtin_amdgcn_sched_barrier(0);
    if constexpr (causal) {
        const int kv_end_pos = (min_tile + 1) * KV_BLOCK_SIZE;
@@ -342,9 +337,6 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
        zero(att_block[1]);
        transpose(k_reg_transposed, k_reg);
        mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]);
-        #if WINDOW
-        mul(att_block[1], att_block[1], TEMPERATURE_SCALE);
-        #endif
 #if WINDOW
        // window masks interior tiles that causal skips
        mask_kv_tile(att_block[1], tile_idx, j - 2, neg_inf_v, lane);
@@ -409,9 +401,6 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
        zero(att_block[0]);
        transpose(k_reg_transposed, k_reg);
        mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
-        #if WINDOW
-        mul(att_block[0], att_block[0], TEMPERATURE_SCALE);
-        #endif
        //      Finish softmax for QK1
        exp2(att_block[1].tiles[1][0], att_block[1].tiles[1][0]);
        mul(norm_vec, norm_vec, scale_vec);
@@ -480,9 +469,6 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    zero(att_block[1]);
    transpose(k_reg_transposed, k_reg);
    mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]);
-    #if WINDOW
-    mul(att_block[1], att_block[1], TEMPERATURE_SCALE);
-    #endif
    //      Finish softmax for QK2
    exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]);
    mul(norm_vec, norm_vec, scale_vec);
@@ -549,9 +535,6 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    zero(att_block[0]);
    transpose(k_reg_transposed, k_reg);
    mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
-    #if WINDOW
-    mul(att_block[0], att_block[0], TEMPERATURE_SCALE);
-    #endif
    //      Finish softmax for QK3
    exp2(att_block[1].tiles[1][0], att_block[1].tiles[1][0]);
    mul(norm_vec, norm_vec, scale_vec);
@@ -614,9 +597,6 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    zero(att_block[1]);
    transpose(k_reg_transposed, k_reg);
    mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]);
-    #if WINDOW
-    mul(att_block[1], att_block[1], TEMPERATURE_SCALE);
-    #endif
    //      Finish softmax for QK4
    exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]);
    mul(norm_vec, norm_vec, scale_vec);
@@ -0,0 +1,69 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+
+#ifndef ATTN_B
+#define ATTN_B 2
+#endif
+#ifndef ATTN_N
+#define ATTN_N 8192
+#endif
+#ifndef ATTN_H
+#define ATTN_H 32
+#endif
+#ifndef ATTN_H_KV
+#define ATTN_H_KV 8
+#endif
+#ifndef ATTN_D
+#define ATTN_D 128
+#endif
+#ifndef THREADS_PER_BLOCK
+#define THREADS_PER_BLOCK 256
+#endif
+
+constexpr int GROUP_SIZE = ATTN_H / ATTN_H_KV;
+constexpr int HALF_D = ATTN_D / 2;
+constexpr int PACKED_D = (GROUP_SIZE + 2) * ATTN_D;
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_BLOCK) void
+fused_qkv_rope_forward(
+    __hip_bfloat16*       __restrict__ q,
+    __hip_bfloat16*       __restrict__ k,
+    __hip_bfloat16*       __restrict__ v,
+    const __hip_bfloat16* __restrict__ xqkv,
+    const __hip_bfloat16* __restrict__ freqs_cis) {
+  const int b = blockIdx.x;
+  const int n = blockIdx.y;
+  const int bn = b * ATTN_N + n;
+  const int packed_bn = bn * ATTN_H_KV * PACKED_D;
+  const int q_bn = bn * ATTN_H * ATTN_D;
+  const int kv_bn = bn * ATTN_H_KV * ATTN_D;
+
+  if (threadIdx.x < HALF_D) {
+    const int pair = threadIdx.x;
+    const int even = pair << 1;
+    const float c = static_cast<float>(freqs_cis[((n * HALF_D + pair) * 2) + 0]);
+    const float s = static_cast<float>(freqs_cis[((n * HALF_D + pair) * 2) + 1]);
+
+    for (int kvh = 0; kvh < ATTN_H_KV; kvh++) {
+      const int base = packed_bn + kvh * PACKED_D;
+
+      for (int rep = 0; rep < GROUP_SIZE; rep++) {
+        const int qbase = base + rep * ATTN_D;
+        const int h = kvh * GROUP_SIZE + rep;
+        const float a = static_cast<float>(xqkv[qbase + even]);
+        const float bb = static_cast<float>(xqkv[qbase + even + 1]);
+        const int out = q_bn + h * ATTN_D + even;
+        q[out] = static_cast<__hip_bfloat16>(a * c - bb * s);
+        q[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c);
+      }
+
+      const float a = static_cast<float>(xqkv[base + GROUP_SIZE * ATTN_D + even]);
+      const float bb = static_cast<float>(xqkv[base + GROUP_SIZE * ATTN_D + even + 1]);
+      const int out = kv_bn + kvh * ATTN_D + even;
+      k[out] = static_cast<__hip_bfloat16>(a * c - bb * s);
+      k[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c);
+      v[out] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even];
+      v[out + 1] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even + 1];
+    }
+  }
+}
@@ -1,6 +1,6 @@
 [project]
 name = "tinygrad"
-version = "0.14.0"
+version = "0.13.0"
 description = "You like pytorch? You like micrograd? You love tinygrad! <3"
 authors = [{ name = "George Hotz" }]

@@ -84,7 +84,7 @@ testing = [
  "pillow",
  "onnx==1.19.0",
  "onnx2torch",
-  "onnxruntime==1.24.1",
+  "onnxruntime",
  "opencv-python",
  "transformers",
  "sentencepiece",
@@ -50,10 +50,10 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op
 \toprule
 \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
 \midrule
-\op{Param} & () & slot, dtype, size?, device?, addrspace? &
-  Placeholder with flat storage of $\mathrm{size}$ elements. Substituted in \op{Call}. \\[4pt]
-\op{Buffer} & () & slot, dtype, size, device, addrspace &
-  Concrete buffer slot with flat storage of $\mathrm{size}$ elements. \\
+\op{Param} & $(\mathbf{s})$ & slot, dtype, device?, addrspace? &
+  Placeholder with shape $\mathbf{s}$. Substituted in \op{Function}. \\[4pt]
+\op{Buffer} & $(\mathbf{s})$ & slot, dtype, device, addrspace &
+  Concrete buffer slot with shape $\mathbf{s}$. If device is a tuple, it creates the fully sized buffer across multiple devices. \\
 \op{Const} & () & value, dtype &
  A scalar constant with shape $(\ )$. \\
  & & & Form vector consts with \op{Stack} \\
@@ -256,9 +256,9 @@ Every UOp has a \textbf{dtype}, \textbf{shape}, \textbf{device}, \textbf{addrspa
 \toprule
 \textbf{Op} & \textbf{dtype} & \textbf{shape} & \textbf{device} & \textbf{min\_max} \\
 \midrule
-\op{Buffer} & from arg & from arg ($\mathrm{size}$) & from arg & dtype range \\
+\op{Buffer} & from arg & from $\mathrm{src}[0]$ & from arg & dtype range \\
 \op{Const} & from arg & $()$ & \textsc{null} & $[v, v]$ \\
-\op{Param} & from arg & from arg ($\mathrm{size}$) & from arg & from src or dtype range \\[3pt]
+\op{Param} & from arg & from $\mathrm{src}[0]$ & from arg & from src or dtype range \\[3pt]
 Movement ops & $\mathrm{src}[0].\mathrm{dtype}$ & (see op) & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
 \op{Unshard} & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0]$, each $a_k \times n_k$ & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
 \op{Reduce} & $\mathrm{src}[0].\mathrm{dtype}$ & remove first $n$ axes & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\[3pt]
@@ -1002,39 +1002,6 @@ class TestBarrier(unittest.TestCase):
    for tid in range(64):
      self.assertEqual(st.vgpr[tid][0], tid + 100 + 1000, f"tid={tid}")

-class TestSMaxMinSCCRegressions(unittest.TestCase):
-  """Regression test: S_MAX sets SCC only on strict inequality (equal operands -> SCC=0)."""
-
-  def test_s_max_i32_equal_scc(self):
-    st = run_program([s_mov_b32(s[4], 64), s_mov_b32(s[5], 64), s_max_i32(s[6], s[4], s[5])], n_lanes=1)
-    self.assertEqual(st.scc, 0)
-    self.assertEqual(st.sgpr[6], 64)
-    st = run_program([s_mov_b32(s[4], 65), s_mov_b32(s[5], 64), s_max_i32(s[6], s[4], s[5])], n_lanes=1)
-    self.assertEqual(st.scc, 1)  # still set when strictly greater
-
-  def test_s_max_u32_equal_scc(self):
-    st = run_program([s_mov_b32(s[4], 64), s_mov_b32(s[5], 64), s_max_u32(s[6], s[4], s[5])], n_lanes=1)
-    self.assertEqual(st.scc, 0)
-
-class TestAbsdiffOverflowRegressions(unittest.TestCase):
-  """Regression test: S_ABSDIFF_I32 computes abs on the WRAPPED 32-bit difference (found by random difftest vs hardware)."""
-
-  def test_s_absdiff_wrapped(self):
-    # |45 - (-2147483647)| overflows int32; hardware takes abs of the wrapped 32-bit difference
-    instructions = [s_mov_b32(s[4], 45), s_mov_b32(s[5], 0x80000001), s_absdiff_i32(s[6], s[4], s[5])]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[6], 0x7FFFFFD4)
-    self.assertEqual(st.scc, 1)
-    # INT_MIN - 1 wraps to +2147483647, already positive
-    instructions = [s_mov_b32(s[4], 0x80000000), s_mov_b32(s[5], 1), s_absdiff_i32(s[6], s[4], s[5])]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[6], 0x7FFFFFFF)
-    # equality -> 0 and SCC=0
-    instructions = [s_mov_b32(s[4], 7), s_mov_b32(s[5], 7), s_absdiff_i32(s[6], s[4], s[5])]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.sgpr[6], 0)
-    self.assertEqual(st.scc, 0)
-

 if __name__ == '__main__':
  unittest.main()
@@ -1629,66 +1629,5 @@ class TestSwap(unittest.TestCase):
    self.assertEqual(st.vgpr[0][1], 0x55555555)


-class TestCvtFrexpRegressions(unittest.TestCase):
-  """Regression tests for float<->int conversion and FREXP corner cases (found by random difftest vs hardware)."""
-
-  def test_cvt_i32_f32_nan_is_zero(self):
-    """v_cvt_i32_f32 of NaN is 0, not INT_MIN (x86 cvttss2si returns INT_MIN)."""
-    for nan in (0x7FC00000, 0xFFC00000, 0x7F800001):
-      st = run_program([v_mov_b32_e32(v[0], nan), v_cvt_i32_f32_e32(v[1], v[0])], n_lanes=1)
-      self.assertEqual(st.vgpr[0][1], 0, f"nan=0x{nan:08x}")
-
-  def test_cvt_i32_f32_positive_overflow(self):
-    """v_cvt_i32_f32 saturates positive overflow/inf to INT_MAX, not INT_MIN."""
-    for bits in (0x7F800000, 0x4F000000, 0x4F800000):  # +inf, 2^31, ~2^32
-      st = run_program([v_mov_b32_e32(v[0], bits), v_cvt_i32_f32_e32(v[1], v[0])], n_lanes=1)
-      self.assertEqual(st.vgpr[0][1], 0x7FFFFFFF, f"bits=0x{bits:08x}")
-
-  def test_cvt_i32_f32_negative_overflow(self):
-    """v_cvt_i32_f32 saturates negative overflow/-inf to INT_MIN."""
-    for bits in (0xFF800000, 0xCF000001):  # -inf, below -2^31
-      st = run_program([v_mov_b32_e32(v[0], bits), v_cvt_i32_f32_e32(v[1], v[0])], n_lanes=1)
-      self.assertEqual(st.vgpr[0][1], 0x80000000, f"bits=0x{bits:08x}")
-
-  def test_cvt_u32_f32_nan_is_zero(self):
-    """v_cvt_u32_f32 of NaN is 0, not UINT_MAX."""
-    for nan in (0x7FC00000, 0xFFC00000, 0x7F800001):
-      st = run_program([v_mov_b32_e32(v[0], nan), v_cvt_u32_f32_e32(v[1], v[0])], n_lanes=1)
-      self.assertEqual(st.vgpr[0][1], 0, f"nan=0x{nan:08x}")
-
-  def test_cvt_i32_f64_nan_and_overflow(self):
-    """v_cvt_i32_f64: NaN -> 0, positive overflow/+inf -> INT_MAX."""
-    st = run_program([v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0x7FF80000), v_cvt_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0)
-    st = run_program([v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0x41F00000), v_cvt_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x7FFFFFFF)  # 2^32 -> INT_MAX
-
-  def test_frexp_f32_denormal(self):
-    """v_frexp_exp/mant_f32 of denormal/zero inputs is (0, signed zero) on hardware."""
-    for bits in (0x00000001, 0x007FFFFF, 0x00000000):
-      st = run_program([v_mov_b32_e32(v[0], bits), v_frexp_exp_i32_f32_e32(v[1], v[0]), v_frexp_mant_f32_e32(v[2], v[0])], n_lanes=1)
-      self.assertEqual(st.vgpr[0][1] & 0xFFFFFFFF, 0, f"exp bits=0x{bits:08x}")
-      self.assertEqual(st.vgpr[0][2], bits & 0x80000000, f"mant bits=0x{bits:08x}")
-    # negative denormal: mant is -0.0
-    st = run_program([v_mov_b32_e32(v[0], 0x80000001), v_frexp_mant_f32_e32(v[2], v[0])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x80000000)
-
-  def test_frexp_f64_denormal(self):
-    """v_frexp_exp_f64 of a denormal returns the normalized exponent (-1073 for min-denormal); zero -> 0."""
-    st = run_program([v_mov_b32_e32(v[0], 1), v_mov_b32_e32(v[1], 0), v_frexp_exp_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][2] & 0xFFFFFFFF, 0xFFFFFBCF)  # -1073
-    st = run_program([v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_frexp_exp_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0)
-
-  def test_frexp_exp_inf_nan(self):
-    """v_frexp_exp of +/-inf and NaN is 0 on hardware (host frexp gives 129/1024), for both f32 and f64."""
-    for bits in (0x7F800000, 0xFF800000, 0x7FC00000):
-      st = run_program([v_mov_b32_e32(v[0], bits), v_frexp_exp_i32_f32_e32(v[1], v[0])], n_lanes=1)
-      self.assertEqual(st.vgpr[0][1] & 0xFFFFFFFF, 0, f"f32 bits=0x{bits:08x}")
-    for lo, hi in ((0, 0x7FF00000), (0, 0xFFF00000), (0, 0x7FF80000), (1, 0x7FF00000)):
-      st = run_program([v_mov_b32_e32(v[0], lo), v_mov_b32_e32(v[1], hi), v_frexp_exp_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
-      self.assertEqual(st.vgpr[0][2] & 0xFFFFFFFF, 0, f"f64 bits=0x{hi:08x}{lo:08x}")
-
-
 if __name__ == '__main__':
  unittest.main()
@@ -989,53 +989,6 @@ class TestCarryOps(unittest.TestCase):
    self.assertEqual(st.vgpr[0][0], 0)  # 0xFFFFFFFF + 1 + 0 = 0 (overflow)
    self.assertEqual(st.vcc, 0xDEADBEEF)  # VCC unchanged - carry was discarded

-class TestSelectFlushRegressions(unittest.TestCase):
-  """Regression tests: f32 MIN/MAX flush denormal inputs to signed zero (select-style ops propagate inputs bitwise)."""
-
-  def test_v_min_f32_denormal_flush(self):
-    """min(denormal, 1.0) is +0, min(-denormal, -1.0) is -0."""
-    st = run_program([v_mov_b32_e32(v[0], 0x00000001), v_mov_b32_e32(v[1], 0x3F800000), v_min_f32_e32(v[2], v[0], v[1])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x00000000)
-    # flush(-denormal) = -0.0 > -1.0, so the result is -1.0 (both operand orders)
-    st = run_program([v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 0xBF800000), v_min_f32_e32(v[2], v[0], v[1])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0xBF800000)
-    st = run_program([v_mov_b32_e32(v[1], 0xBF800000), v_mov_b32_e32(v[2], 0x80000001), v_min_f32_e32(v[3], v[1], v[2])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][3], 0xBF800000)
-
-  def test_v_max_f32_denormal_flush(self):
-    """max(-denormal, -1.0) is -0; max(+denormal, -0) is +0."""
-    st = run_program([v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 0xBF800000), v_max_f32_e32(v[2], v[0], v[1])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x80000000)
-    st = run_program([v_mov_b32_e32(v[0], 0x00000001), v_mov_b32_e32(v[1], 0x80000000), v_max_f32_e32(v[2], v[0], v[1])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x00000000)
-
-
-class TestCarryExecRegressions(unittest.TestCase):
-  """Regression tests: per-lane VCC writes (carry ops) zero inactive lane bits - VCC = mask & EXEC, never preserved."""
-
-  def test_co_ci_e32_vcc_masked_by_exec(self):
-    """v_sub_co_ci_u32_e32 with EXEC=0xFFFF0000: hw clears inactive VCC bits instead of preserving them."""
-    instructions = [
-      s_mov_b32(EXEC_LO, 0xFFFF0000),
-      s_mov_b32(VCC_LO, 0xFFFFFFFF),  # preset all bits
-      v_mov_b32_e32(v[0], 0xFFFFFFFE), v_mov_b32_e32(v[1], 0x80000000),
-      v_sub_co_ci_u32_e32(v[2], v[0], v[1]),  # active lanes: no borrow
-    ]
-    st = run_program(instructions, n_lanes=32)
-    self.assertEqual(st.vcc, 0x00000000)
-
-  def test_co_ci_e32_vcc_masked_by_exec_ones(self):
-    """Same with all-ones carry: VCC = borrow_mask & EXEC."""
-    instructions = [
-      s_mov_b32(EXEC_LO, 0x0F0F0F0F),
-      s_mov_b32(VCC_LO, 0),
-      v_mov_b32_e32(v[0], 0xFFFFFFFF), v_mov_b32_e32(v[1], 1),
-      v_add_co_ci_u32_e32(v[2], v[0], v[1]),  # all lanes would carry if active
-    ]
-    st = run_program(instructions, n_lanes=32)
-    self.assertEqual(st.vcc, 0x0F0F0F0F)
-    self.assertEqual(st.vgpr[31][2], 0)  # 0xFFFFFFFF + 1 wraps to 0 in active lanes
-

 if __name__ == '__main__':
  unittest.main()
@@ -4,7 +4,6 @@ Includes: v_fma_f32, v_div_scale_f32, v_div_fmas_f32, v_div_fixup_f32,
          v_alignbit_b32, v_bfe_i32, v_mad_u64_u32, v_readlane_b32, v_writelane_b32
 """
 import unittest
-from tinygrad.helpers import OSX
 from test.amd.hw.helpers import *

 class TestFMA(unittest.TestCase):
@@ -3265,23 +3264,6 @@ class TestVOP3ClampMAD(unittest.TestCase):
    # 0xFFFF * 2 = 0x1FFFE, low 16 bits = 0xFFFE
    self.assertEqual(st.vgpr[0][3] & 0xFFFF, 0xFFFE, f"expected 0xFFFE, got 0x{st.vgpr[0][3] & 0xFFFF:04x}")

-class TestMadNarrowClampRegressions(unittest.TestCase):
-  """Regression tests: mad i16/i24 with clamp saturate to narrow output range (found by random difftest vs hardware)."""
-
-  def test_mad_i16_clamp_sat_max(self):
-    # neg/src-floggled 16-bit mul operands are sign-extended after toggling bit15; sum > INT_MAX saturates
-    instructions = [s_mov_b32(s[4], 1232348160), v_mov_b32_e32(v[3], 0x80000000),
-                    v_mov_b32_e32(v[1], 0x7F7FFFFF), v_mad_i32_i16(v[0], s[4], v[3], v[1], 0, 3, 5, 1)]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 0x7FFFFFFF)
-
-  def test_mad_i24_clamp_sat_min(self):
-    # sext24(-6344704) * sext24(+4210688) << -2^31 saturates to INT_MIN
-    instructions = [s_mov_b32(s[7], 4290772992), v_mov_b32_e32(v[1], 1077936128),
-                    v_mad_i32_i24(v[0], s[7], v[1], v[1], 1, 0, 0, 1)]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][0], 0x80000000)
-

 class TestCvtPkF16(unittest.TestCase):
  """Tests for V_CVT_PK_RTZ_F16_F32 - pack two f32 to f16 with round toward zero."""
@@ -3669,80 +3651,6 @@ class TestPermlane(unittest.TestCase):
    self.assertEqual(st.vgpr[21][1], 5)
    self.assertEqual(st.vgpr[31][1], 15)

-class TestClampLdExpRegressions(unittest.TestCase):
-  """Regression tests for f32 clamp (-0 -> +0) and ldexp input passthrough."""
-
-  def test_clamp_negative_zero(self):
-    """clmp=1 maps -0.0 to +0.0 (found by random difftest vs hardware)."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0x80000000), v_mov_b32_e32(v[1], 0x80000000),
-      v_add_f32_e64(v[2], v[0], v[1], clmp=1),  # -0 + -0 = -0, clamp -> +0
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x00000000)
-    instructions = [
-      v_mov_b32_e32(v[0], 0x3F800000), v_mov_b32_e32(v[1], 0x80000000),
-      v_min_f32_e64(v[2], v[0], v[1], clmp=1),  # min(1.0, -0) = -0, clamp -> +0
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x00000000)
-
-  def test_ldexp_special_inputs(self):
-    """v_ldexp_f32 of 0/-0/inf/NaN propagates the input instead of computing val * 2**exp (0*inf = NaN on host)."""
-    # -0.0 * 2^INT_MIN = -0.0 (src1 as integer exponent; huge negative)
-    instructions = [v_mov_b32_e32(v[0], 0x80000000), v_mov_b32_e32(v[1], 0x80000000), v_ldexp_f32(v[2], v[0], v[1])]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x80000000)
-    # inf stays inf even with negative exponent
-    instructions = [v_mov_b32_e32(v[0], 0x7F800000), v_mov_b32_e32(v[1], 0xFFFFFF80), v_ldexp_f32(v[2], v[0], v[1])]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x7F800000)
-
-  def test_ldexp_denormal_flush(self):
-    """v_ldexp_f32/f64 flush denormal inputs to signed zero (found by random difftest vs hardware)."""
-    # ldexp(+denorm, 1) = +0, ldexp(-denorm, 250) = -0
-    for src, exp_val, want in [(0x00000001, 1, 0x00000000), (0x80000001, 250, 0x80000000)]:
-      st = run_program([v_mov_b32_e32(v[0], src), v_mov_b32_e32(v[1], exp_val), v_ldexp_f32(v[2], v[0], v[1])], n_lanes=1)
-      self.assertEqual(st.vgpr[0][2], want)
-
-  def test_v_mul_neg_modifier_nan_sign(self):
-    """neg modifier is a pure sign-bit toggle on a NaN operand; result keeps that sign (found by random difftest)."""
-    # mul(normal, NEG(ABS(qNaN))): NaN payload negated in the operand stays negative qNaN
-    instructions = [v_mov_b32_e32(v[0], 0xC96CF47F), v_mov_b32_e32(v[1], 0x7FC00000),
-                    v_mul_f32_e64(v[2], v[0], v[1], s[0], 0, 7, 6)]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0xFFC00000)
-    # plain neg modifier still applies to non-NaN values: mul(-1.0, NEG(2.0)) = +2.0
-    st = run_program([v_mov_b32_e32(v[0], 0xBF800000), v_mov_b32_e32(v[1], 0x40000000),
-                      v_mul_f32_e64(v[2], v[0], v[1], s[0], 0, 2, 0)], n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0x40000000)
-
-
-class TestNaNPropagationRegressions(unittest.TestCase):
-  """Regression tests: float arithmetic propagates a NaN from the FIRST NaN operand, quieted with its own sign/payload."""
-
-  @unittest.skipIf(OSX, "broken on mac, TODO: why?")
-  def test_mul_nan_priority(self):
-    # first NaN operand wins (sign+payload), not x86's second-source propagation
-    for a, b, want in [(0x7FC00001, 0x7F800003, 0x7FC00001), (0xFFC00005, 0x7F800003, 0xFFC00005),
-                       (0x7F800001, 0xFFC00005, 0x7FC00001), (0xFF9F1800, 0x7F800001, 0xFFDF1800)]:
-      st = run_program([v_mov_b32_e32(v[0], a), v_mov_b32_e32(v[1], b),
-                        v_mul_f32_e32(v[2], v[0], v[1])], n_lanes=1)
-      self.assertEqual(st.vgpr[0][2], want, f"mul({a:#x}, {b:#x})")
-
-class TestMinMaxFlushE64Regressions(unittest.TestCase):
-  """Regression tests: f32 min/max/median flush denormal inputs to signed zero (e64 forms)."""
-
-  def test_v_min3_f32_denormal_flush(self):
-    st = run_program([v_mov_b32_e32(v[0], 0x00000001), v_mov_b32_e32(v[1], 0x3F800000), v_mov_b32_e32(v[2], 0x40000000),
-                      v_min3_f32(v[3], v[0], v[1], v[2])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][3], 0x00000000)  # min(+denorm, 1, 2) = +0
-
-  def test_v_med3_f32_denormal_flush(self):
-    st = run_program([v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 0x3F800000), v_mov_b32_e32(v[2], 0x40000000),
-                      v_med3_f32(v[3], v[0], v[1], v[2])], n_lanes=1)
-    self.assertEqual(st.vgpr[0][3], 0x3F800000)  # med(-0, 1, 2) = 1
-

 if __name__ == '__main__':
  unittest.main()
@@ -973,71 +973,6 @@ class TestCmpxPartialWavefront(unittest.TestCase):
    self.assertEqual(st.sgpr[EXEC_LO.offset] & 0xFFFFFFFF, 0x4,
                     "Only lane 2 should be active after v_cmpx_eq_u32_e64")

-class TestClassDenormalRegressions(unittest.TestCase):
-  """Regression tests: V_CMP_CLASS classifies denormals as DENORMAL (raw bits), not as zero class."""
-
-  def test_class_pos_denormal(self):
-    for bits in (0x00000001, 0x007FFFFF):
-      instructions = [v_mov_b32_e32(v[0], bits), v_mov_b32_e32(v[1], 0x80), v_cmp_class_f32_e64(VCC_LO, v[0], v[1])]
-      st = run_program(instructions, n_lanes=1)
-      self.assertEqual(st.vcc, 1, f"bits=0x{bits:08x}")  # n_lanes=1
-      # ...and it is not the zero class
-      instructions = [v_mov_b32_e32(v[0], bits), v_mov_b32_e32(v[1], 0x40), v_cmp_class_f32_e64(VCC_LO, v[0], v[1])]
-      st = run_program(instructions, n_lanes=1)
-      self.assertEqual(st.vcc, 0, f"bits=0x{bits:08x}")
-
-  def test_class_neg_denormal(self):
-    instructions = [v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 0x10), v_cmp_class_f32_e64(VCC_LO, v[0], v[1])]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc, 1)  # n_lanes=1
-    instructions = [v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 0x20), v_cmp_class_f32_e64(VCC_LO, v[0], v[1])]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc, 0)  # not the negative-zero class
-
-
-class TestIntCmpModRegressions(unittest.TestCase):
-  """Regression tests: int compares (i32/u32) honor abs/neg as bit-level sign clear/flip (not integer abs/negate)."""
-
-  def test_cmp_i32_abs_neg_bit_level(self):
-    # abs(0x80000001) = 1 -> 1 > 1 is false (integer abs would give 2147483647 > 1)
-    instructions = [v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 1), v_cmp_gt_i32_e64(VCC_LO, v[0], v[1], abs=1)]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc, 0)
-    # neg(0x80000001) flips the sign bit -> 1 > 2 is false (integer negate would give 2147483647 > 2)
-    instructions = [v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 2), v_cmp_gt_i32_e64(VCC_LO, v[0], v[1], neg=1)]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc, 0)
-
-  def test_cmp_u32_abs_bit_level(self):
-    # abs(0x80000000) = 0 -> 0 < 1 is true
-    instructions = [v_mov_b32_e32(v[0], 0x80000000), v_mov_b32_e32(v[1], 1), v_cmp_lt_u32_e64(VCC_LO, v[0], v[1], abs=1)]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vcc, 1)  # n_lanes=1
-
-
-class TestCmpxSdstRegressions(unittest.TestCase):
-  """Regression tests: V_CMPX_*_E64 writes EXEC only, never SDST (hardware verified)."""
-
-  def test_cmpx_e64_no_sdst(self):
-    instructions = [
-      s_mov_b32(VCC_LO, 0),  # preset VCC to 0
-      v_mov_b32_e32(v[0], 0x3F800000), v_mov_b32_e32(v[1], 0x40000000),
-      v_cmpx_lt_f32_e64(VCC_LO, v[0], v[1]),  # 1.0 < 2.0
-    ]
-    st = run_program(instructions, n_lanes=32)
-    self.assertEqual(st.sgpr[EXEC_LO.offset], 0xFFFFFFFF)  # EXEC updated
-    self.assertEqual(st.vcc, 0)  # but VCC untouched
-
-  def test_cmpx_e64_partial_exec(self):
-    instructions = [
-      s_mov_b32(EXEC_LO, 0x0F0F0F0F),
-      s_mov_b32(VCC_LO, 0xFFFFFFFF),
-      v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0x3F800000),
-      v_cmpx_lt_f32_e64(VCC_LO, v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=32)
-    self.assertEqual(st.sgpr[EXEC_LO.offset], 0x0F0F0F0F)  # EXEC = computed & old EXEC
-

 if __name__ == '__main__':
  unittest.main()
@@ -30,7 +30,7 @@ def custom_add_one(A:UOp) -> UOp:
    s_endpgm(),
  ]
  sink = UOp.sink(A.base, threads, arg=KernelInfo(f"custom_add_one_{A.numel()}", estimates=Estimates(ops=A.numel(), mem=A.numel()*4*2)))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

 def custom_add_var(A:UOp, B:UOp) -> UOp:
  A,B = A.flatten(), B.flatten()
@@ -49,7 +49,7 @@ def custom_add_var(A:UOp, B:UOp) -> UOp:
    s_endpgm(),
  ]
  sink = UOp.sink(A.base, B.base, var, threads, arg=KernelInfo(f"custom_add_var_{A.numel()}"))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

 def custom_wave_sync(A:UOp, arch:str) -> UOp:
  # 4 waves across 1024 WG — enough to saturate a SIMD with many concurrent WGs
@@ -63,7 +63,7 @@ def custom_wave_sync(A:UOp, arch:str) -> UOp:
    insts += [s_nop(0)]*4
  insts.append(s_endpgm())
  sink = UOp.sink(A.base, threads, wg, arg=KernelInfo("custom_wave_sync"))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

 def custom_lds_sync(A:UOp, arch:str) -> UOp:
  A = A.flatten()
@@ -97,7 +97,7 @@ def custom_lds_sync(A:UOp, arch:str) -> UOp:
    isa.s_endpgm(),
  ]
  sink = UOp.sink(A.base, lds, threads, wg, arg=KernelInfo("custom_lds_sync"))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

 def custom_handwritten(A:UOp) -> UOp:
  A = A.flatten()
@@ -143,7 +143,7 @@ def custom_handwritten(A:UOp) -> UOp:
  k.emit(r4.s_endpgm())
  insts = k.finalize()
  sink = UOp.sink(A.base, threads, wg, lds, arg=KernelInfo("custom_handwritten"))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

 def custom_data_deps(A:UOp) -> UOp:
  A = A.flatten()
@@ -159,7 +159,7 @@ def custom_data_deps(A:UOp) -> UOp:
  k.emit(s_endpgm())
  insts = k.finalize()
  sink = UOp.sink(A.base, threads, arg=KernelInfo("custom_data_deps"))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

@unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
 class TestAsmKernel(unittest.TestCase):
@@ -67,26 +67,32 @@ class TestParseExpr(unittest.TestCase):

  def test_integer_literals(self):
    """Test parsing integer literals."""
-    self.assertIs(parse_expr('0', {}), UOp.const(0, dtypes.uint32))
-    self.assertIs(parse_expr('42', {}), UOp.const(42, dtypes.uint32))
-    self.assertIs(parse_expr('42U', {}), UOp.const(42, dtypes.uint32))
+    self.assertEqual(parse_expr('0', {}).val, 0)
+    self.assertEqual(parse_expr('42', {}).val, 42)
+    self.assertEqual(parse_expr('42U', {}).val, 42)

  def test_negative_integers(self):
    """Test parsing negative integer literals."""
-    self.assertIs(parse_expr('-1', {}), UOp.const(-1, dtypes.int))
+    result = parse_expr('-1', {})
+    self.assertEqual(result.val, -1)
+    self.assertEqual(result.dtype, dtypes.int)

  def test_float_literals(self):
    """Test parsing float literals."""
-    self.assertIs(parse_expr('1.0F', {}), UOp.const(1.0, dtypes.float32))
+    result = parse_expr('1.0F', {})
+    self.assertEqual(result.val, 1.0)
+    self.assertEqual(result.dtype, dtypes.float32)

  def test_hex_literals(self):
    """Test parsing hex literals."""
-    self.assertIs(parse_expr('0xFF', {}), UOp.const(255, dtypes.uint32))
+    result = parse_expr('0xFF', {})
+    self.assertEqual(result.val, 255)

  def test_variable_lookup(self):
    """Test variable lookup in parse_expr."""
    vrs = {'x': UOp.const(42, dtypes.uint32)}
-    self.assertIs(parse_expr('x', vrs), vrs['x'])
+    result = parse_expr('x', vrs)
+    self.assertEqual(result.val, 42)

  def test_binary_ops(self):
    """Test parsing binary operations."""
@@ -97,7 +103,9 @@ class TestParseExpr(unittest.TestCase):
    self.assertEqual(result.op, Ops.ADD)

    # Subtraction with constant folding
-    self.assertIs(parse_expr('10 - 5', {}), UOp.const(5, dtypes.uint32))
+    result = parse_expr('10 - 5', {})
+    self.assertEqual(result.op, Ops.CONST)
+    self.assertEqual(result.val, 5)

  def test_ternary(self):
    """Test parsing ternary expressions."""
@@ -134,8 +142,15 @@ class TestForLoopParsing(unittest.TestCase):
    S0 = UOp.const(0, dtypes.uint32)
    _vrs, assigns = parse_pcode(pcode, {'S0': S0})

-    # every cond folds (S0 is a const), leaving the default branch: -1 in the destination dtype
-    self.assertIs(assigns[0][1].simplify(), UOp.const(-1, dtypes.uint32))
+    # Check that the innermost value (default) is -1 (may be wrapped in CAST)
+    val = assigns[0][1]
+    # Traverse to innermost WHERE
+    while val.op == Ops.WHERE:
+      val = val.src[2]  # false branch
+    # Unwrap CAST if present
+    while val.op == Ops.CAST:
+      val = val.src[0]
+    self.assertEqual(val.val, -1)

  def test_ctz_parsing(self):
    """Test CTZ pcode parsing."""
@@ -152,7 +167,7 @@ class TestDSPcodePatterns(unittest.TestCase):

  def test_global_atomic_add_f32_parsing(self):
    """Test GLOBAL_ATOMIC_ADD_F32 keeps memory values in float dtype."""
-    vmem = UOp.param(2, dtypes.uint32, 1024)
+    vmem = UOp.param(2, dtypes.uint32, (1024,))
    srcs = {
      'ADDR': UOp.const(0, dtypes.uint64),
      'DATA': UOp.const(0x3f800000, dtypes.uint32),
@@ -183,7 +198,7 @@ class TestDSPcodePatterns(unittest.TestCase):
  def test_mem_read_parsing(self):
    """Test MEM[addr].type read expression parsing."""
    # Create a mock LDS buffer
-    lds = UOp.param(3, dtypes.uint32, 16384)
+    lds = UOp.param(3, dtypes.uint32, (16384,))
    addr = UOp.const(0, dtypes.uint32)
    vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(0, dtypes.uint32)}

@@ -218,7 +233,7 @@ class TestDSPcodePatterns(unittest.TestCase):
    pcode = PCODE.get(DSOp.DS_LOAD_2ADDR_B32)
    self.assertIsNotNone(pcode)
    assert pcode is not None
-    lds = UOp.param(3, dtypes.uint32, 16384)
+    lds = UOp.param(3, dtypes.uint32, (16384,))
    srcs = {
      'ADDR': UOp.const(0, dtypes.uint32),
      'OFFSET0': UOp.const(0, dtypes.uint32),
@@ -247,8 +262,8 @@ class TestDSPcodePatterns(unittest.TestCase):
    _, assigns = parse_pcode(pcode, srcs)
    # Check addresses: 100 + 2*4 = 108, 100 + 5*4 = 120
    # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp
-    self.assertIs(assigns[0][1][0].simplify(), UOp.const(108, dtypes.uint32))  # type: ignore[index]
-    self.assertIs(assigns[1][1][0].simplify(), UOp.const(120, dtypes.uint32))  # type: ignore[index]
+    self.assertEqual(assigns[0][1][0].simplify().val, 108)  # type: ignore[index]
+    self.assertEqual(assigns[1][1][0].simplify().val, 120)  # type: ignore[index]

  def test_ds_store_data_values(self):
    """Test DS_STORE_2ADDR_B32 uses correct data values."""
@@ -265,8 +280,8 @@ class TestDSPcodePatterns(unittest.TestCase):
    _, assigns = parse_pcode(pcode, srcs)
    # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp
    # DATA[31:0] should preserve the value
-    self.assertIs(assigns[0][1][1].simplify(), UOp.const(0xAAAAAAAA, dtypes.uint32))  # type: ignore[index]
-    self.assertIs(assigns[1][1][1].simplify(), UOp.const(0xBBBBBBBB, dtypes.uint32))  # type: ignore[index]
+    self.assertEqual(assigns[0][1][1].simplify().val, 0xAAAAAAAA)  # type: ignore[index]
+    self.assertEqual(assigns[1][1][1].simplify().val, 0xBBBBBBBB)  # type: ignore[index]

 class TestConditionalParsing(unittest.TestCase):
  """Test conditional (if/elsif/else) pcode parsing."""
@@ -291,15 +306,15 @@ class TestConcatWidthParsing(unittest.TestCase):
  def test_permlanex16_altrow_concat(self):
    for row, expected in [(0, 1), (1, 0), (2, 3), (3, 2)]:
      parsed = parse_expr('{ row[1], ~row[0] }', {'row': UOp.const(row, dtypes.uint32)})
-      self.assertIs(parsed.simplify(), UOp.const(expected, dtypes.uint32))
+      self.assertEqual(parsed.simplify().val, expected)

  def test_permlane64_altlane_concat(self):
    for lane, expected in [(0, 32), (1, 33), (31, 63), (32, 0), (63, 31)]:
      parsed = parse_expr('{ ~lane[5], lane[4:0] }', {'lane': UOp.const(lane, dtypes.uint32)})
-      self.assertIs(parsed.simplify(), UOp.const(expected, dtypes.uint32))
+      self.assertEqual(parsed.simplify().val, expected)

  def test_permlane64_wave64_pcode_indices(self):
-    vgpr = UOp.param(0, dtypes.uint32, 256)
+    vgpr = UOp.param(0, dtypes.uint32, (256,))
    srcs = {
      'SRC0': UOp.const(0, dtypes.uint32),
      'VDST': UOp.const(1, dtypes.uint32),
@@ -312,17 +327,19 @@ class TestConcatWidthParsing(unittest.TestCase):
      'S2': UOp.const(0, dtypes.uint32),
    }

-    def check_load_idx(v: UOp, expected: int):
+    def load_idx(v: UOp) -> int:
      simp = v.simplify()
      self.assertEqual(simp.op, Ops.LOAD)
      self.assertEqual(simp.src[0].op, Ops.INDEX)
-      self.assertIs(simp.src[0].src[1].simplify(), UOp.const(expected, dtypes.uint32))
+      idx = simp.src[0].src[1].simplify()
+      self.assertEqual(idx.op, Ops.CONST)
+      return idx.val

    _, assigns = parse_pcode(PCODE[VOP1Op.V_PERMLANE64_B32_E32], srcs)
    self.assertEqual(len(assigns), 64)
    for lane, (dst_idx, src_idx) in {0: (64, 32), 31: (95, 63), 32: (96, 0), 63: (127, 31)}.items():
-      self.assertIs(assigns[lane][1][0].simplify(), UOp.const(dst_idx, dtypes.uint32))  # type: ignore[index]
-      check_load_idx(assigns[lane][1][1], src_idx)  # type: ignore[index]
+      self.assertEqual(assigns[lane][1][0].simplify().val, dst_idx)  # type: ignore[index]
+      self.assertEqual(load_idx(assigns[lane][1][1]), src_idx)  # type: ignore[index]

 class TestAllPcode(unittest.TestCase):
  """Test that all pcode from all architectures can be parsed."""
@@ -330,7 +347,7 @@ class TestAllPcode(unittest.TestCase):
  def _make_srcs(self):
    """Create dummy source variables for pcode parsing."""
    u32, u64 = lambda v=0: UOp.const(v, dtypes.uint32), lambda v=0: UOp.const(v, dtypes.uint64)
-    lds = UOp.param(3, dtypes.uint32, 16384)
+    lds = UOp.param(3, dtypes.uint32, (16384,))
    return {'laneId': u32(), 'laneID': u32(), 'S0': u32(), 'S1': u32(), 'S2': u32(), 'S3': u32(), 'SRC0': u32(),
            'D0': u32(), 'D1': u32(), 'DST': u32(), 'VDST': u32(), 'SDST': u32(),
            'VCC': u64(), 'VCCZ': u32(), 'EXEC': u64(), 'EXEC_LO': u32(), 'EXECZ': u32(), 'SCC': u32(),
@@ -88,6 +88,7 @@ def run_rocprof_decoder(blobs: list[bytes], lib: bytes, base: int, target: str):
  if t.is_alive(): raise RuntimeError("rocprof decoder timeout")
  return occupancy_records, wave_insts

+@unittest.skip("TODO: fix to not require unpickling UOps.")
 class SQTTExamplesTestBase(unittest.TestCase):
  target: str
  examples: dict
@@ -1,28 +1,30 @@
 import unittest, contextlib
 from tinygrad import Device, Tensor, Context, TinyJit
-from tinygrad.device import Compiled, ProfileProgramEvent
+from tinygrad.device import Compiled, ProfileProgramEvent, ProfileDeviceEvent
 from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
 from tinygrad.viz.serve import load_amd_counters, VizData

@contextlib.contextmanager
 def save_sqtt():
-  Device[Device.DEFAULT].synchronize()
-  profile_start = len(Compiled.profile_events)
  data = VizData()
  yield data.ctxs
  Device[Device.DEFAULT].synchronize()
  Device[Device.DEFAULT]._at_profile_finalize()
-  load_amd_counters(data, [e for e in Compiled.profile_events[:profile_start] if isinstance(e, ProfileProgramEvent)] +
-                          Compiled.profile_events[profile_start:])
+  load_amd_counters(data, Compiled.profile_events)
  data.ctxs[:] = [r for r in data.ctxs if r["name"].startswith("SQTT")]

@unittest.skipUnless(Device.DEFAULT == "AMD", "only runs on AMD")
 class TestSQTTProfiler(unittest.TestCase):
+  # TODO: can we enable SQTT profiling in context?
  @classmethod
  def setUpClass(cls):
    if not Device[Device.DEFAULT].sqtt_enabled: raise unittest.SkipTest("device must be in SQTT profiling mode")

+  def setUp(self):
+    Device[Device.DEFAULT].synchronize()
+    Compiled.profile_events[:] = [e for e in Compiled.profile_events if isinstance(e, (ProfileProgramEvent, ProfileDeviceEvent))]
+
  def test_simple(self):
    t = Tensor.empty(1) + 1
    with save_sqtt() as sqtt:
@@ -1,5 +1,4 @@
 import unittest
-import functools
 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.helpers import getenv, system, DEV
 from extra.gemm.cdna_asm_gemm import asm_gemm, hk_bf16_atb_gemm
@@ -10,7 +9,6 @@ from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8, FP8_MAX
 # Use DEV=NULL:HIP:gfx950 to also test the assembly
 def is_cdna4(): return Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950")

-@functools.cache
 def has_hipcc():
  try: system("hipcc --version")
  except Exception: return False
@@ -188,7 +186,7 @@ class TestMXFP4(unittest.TestCase):
    M, N, K = getenv("M", 16384), getenv("N", 4096), getenv("K", 14336)
    a = Tensor.empty(M, K, dtype=dtypes.bfloat16)
    b = Tensor.empty(N, K, dtype=dtypes.bfloat16)
-    for _ in range(getenv("CNT", 1)): asm_gemm(a, b.T, mxfp4=True).realize()
+    asm_gemm(a, b.T, mxfp4=True).realize()

 # test the Asm GEMM with Llama shapes, only run on the real machine for speed

@@ -1,7 +1,7 @@
 import unittest, math
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.dtype import DTYPES_DICT
-from tinygrad.uop.ops import Ops, UOp, GroupOp
+from tinygrad.uop.ops import Ops, UOp
 from tinygrad.codegen.decomp.op import threefry2x32
 import numpy as np
 from test.helpers import not_support_multi_device
@@ -17,7 +17,7 @@ def _check_ast_count(desired_count:int, t:Tensor):
 class TestMovedConstFolding(unittest.TestCase):
  def test_contiguous_deviceless_const(self):
    t = Tensor(UOp.const(2.0, dtypes.float)).contiguous()
-    self.assertIs(t.uop, UOp.const(2.0, dtypes.float))
+    self.assertIs(t.uop.op, Ops.CONST)
    self.assertIsNone(t.uop.device)

  def test_add_shrunk_zero(self):
@@ -169,8 +169,8 @@ class TestMultiConstFolding(unittest.TestCase):
 class TestThreefryConstFolding(unittest.TestCase):
  def test_threefry(self):
    # THREEFRY(const,const) folds to a const once decomposed
-    x = threefry2x32(UOp.const(5, dtypes.uint64), UOp.const(10, dtypes.uint64)).simplify()
-    self.assertEqual([u.op for u in x.toposort() if u.op in GroupOp.ALU], [])
+    x = threefry2x32(UOp.const(5, dtypes.uint64), UOp.const(10, dtypes.uint64))
+    self.assertIs(x.simplify().op, Ops.CONST)

 class TestTautologicalCompare(unittest.TestCase):
  # without const folding, these would have triggered -Wtautological-compare in clang
@@ -4,7 +4,7 @@ import numpy as np
 from tinygrad.dtype import AddrSpace, dtypes, Invalid
 from tinygrad.uop.ops import KernelInfo, AxisType, Ops
 from tinygrad.renderer.ptx import PTXRenderer
-from test.helpers import assert_kernel_count, KernelCountException
+from test.helpers import assert_kernel_count

 # **** kernels ****

@@ -474,7 +474,7 @@ class TestCustomKernelInput(unittest.TestCase):
    y.realize()
    kernel_count = GlobalCounters.kernel_count
    self.assertEqual(y.tolist(), x.add(1).tolist())
-    if kernel_count > max_kernels: raise KernelCountException(max_kernels, kernel_count)
+    self.assertLessEqual(kernel_count, max_kernels)
    # same test with @function, input is PARAM
    from tinygrad import function
    x0 = Tensor.arange(32).clone("CPU").realize()
@@ -487,7 +487,7 @@ class TestCustomKernelInput(unittest.TestCase):
    y = run(x0).realize()
    kernel_count = GlobalCounters.kernel_count
    self.assertEqual(y.tolist(), mop_fxn(x0).add(1).tolist())
-    if kernel_count > max_kernels: raise KernelCountException(max_kernels, kernel_count)
+    self.assertLessEqual(kernel_count, max_kernels)

  def test_reshape(self): self._test_mop(lambda x: x.reshape(16, 2), max_kernels=2)
  def test_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T, max_kernels=3)
@@ -258,11 +258,6 @@ class TestDoubleDType(TestDType):
      a = [2, 3, 4]
      np.testing.assert_allclose(func(Tensor(a, dtype=self.DTYPE)).numpy(), func(torch.tensor(a, dtype=torch.float64)), rtol=1e-12, atol=1e-12)

-  def test_float32_compare_selecting_float64(self):
-    a = Tensor([1.0, 2.0, 5.0, 9.0], dtype=dtypes.float32)
-    p, q = Tensor([10., 20., 30., 40.], dtype=self.DTYPE), Tensor([50., 60., 70., 80.], dtype=self.DTYPE)
-    _test_op(lambda: (a < 3.0).where(p, q), self.DTYPE, [10., 20., 70., 80.])
-
  def test_float64_to_float32_cast_inf(self):
    _test_op(lambda: Tensor([3.4e40, 3.4e38, 1, 0], dtype=dtypes.float64).cast(dtypes.float32),
             dtypes.float32, [float('inf'), 3.4e38, 1, 0])
@@ -328,10 +323,7 @@ class TestUint16DType(TestDType):
 class TestInt32DType(TestDType): DTYPE = dtypes.int32
 class TestUint32DType(TestDType): DTYPE = dtypes.uint32

-class TestInt64DType(TestDType):
-  DTYPE = dtypes.int64
-  def test_int64_to_uint32_to_int64(self):
-    _test_op(lambda: Tensor([0x12345678ABCDEF01], dtype=dtypes.int64).cast(dtypes.uint32).cast(dtypes.int64), dtypes.int64, [2882400001])
+class TestInt64DType(TestDType): DTYPE = dtypes.int64

@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
 class TestEmulatedInt64DType(TestInt64DType):
@@ -431,11 +423,6 @@ class TestDtypeUsage(unittest.TestCase):
        t = Tensor([[1, 2], [3, 4]], dtype=d)
        (t*t).max().item()

-  def test_where_float16_compare_to_const(self):
-    # t > 0 is CMPLT(0, t): the float16 operand is on the right
-    t = Tensor([-1.0, 1.0], dtype=dtypes.float16)
-    np.testing.assert_equal((t > 0).where(Tensor.ones(2, dtype=dtypes.float16), Tensor.zeros(2, dtype=dtypes.float16)).numpy(), [0.0, 1.0])
-
@unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, f"no bfloat16 on {Device.DEFAULT}")
 class TestOpsBFloat16(unittest.TestCase):
  def test_cast(self):
@@ -6,8 +6,6 @@ from tinygrad.tensor import _to_np_dtype
 from tinygrad.runtime.ops_python import from_storage_scalar
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
-from tinygrad.renderer.llvmir import CPULLVMRenderer
-from tinygrad.renderer.isa.x86 import X86Renderer
 from tinygrad.uop import Ops
 import numpy as np
 import pytest
@@ -66,8 +64,6 @@ ht.fp8e5m2fnuz = ht.uint8
 def universal_test(a, b, dtype, op):
  if not isinstance(op, tuple): op = (op, op)
  if op[0] == operator.mod and b == 0: return
-  # TODO: throws floating point exception
-  if isinstance(Device[Device.DEFAULT].renderer, (X86Renderer, CPULLVMRenderer)) and op[0] == operator.mod and a == dtype.min and b == -1: return
  # lt and max with nan is undefined in tinygrad
  if op[0] in (operator.lt, Tensor.maximum) and (math.isnan(a) or math.isnan(b)): return
  ta, tb = Tensor([a], dtype=dtype), Tensor([b], dtype=dtype)
@@ -4,7 +4,7 @@ from tinygrad.uop.ops import UOp, Ops
 from tinygrad.dtype import dtypes
 from tinygrad.renderer.isa.x86 import X86Ops, X86Renderer, RBP, RDI, RSP, RSI, RAX, RDX, XMM, GPR, imm, def_reg

-def ins(op, dt, src, tag=None): return UOp(Ops.INS, arg=(op, dt), src=src, tag=tag)
+def ins(op, dt, src, tag=None): return UOp(Ops.INS, arg=op, dtype=dt, src=src, tag=tag)

@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only on x86")
 class TestEncodingsX86(unittest.TestCase):
@@ -100,6 +100,13 @@ class TestEncodingsX86(unittest.TestCase):
    # vaddss xmm0, xmm0, xmm8
    self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C4 C1 7A 58 C0"))

+  # test ymm encoding
+  def test_ymm_encoding(self):
+    xmm0, xmm1 = def_reg(dtypes._uint256, XMM[0]), def_reg(dtypes._uint256, XMM[1])
+    add = ins(X86Ops.VADDPS, dtypes._uint256, (xmm0, xmm1), XMM[0])
+    # vaddps ymm0, ymm0, ymm1
+    self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C5 FC 58 C1"))
+
  # test encoding where register is in the immediate field
  def test_reg_in_imm_field(self):
    xmm0, xmm1, xmm2 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[1]), def_reg(dtypes.float32, XMM[2])
@@ -136,7 +143,7 @@ class TestEncodingsX86(unittest.TestCase):

  # cmoves have the cmp as the last src even though it is not explicitly used, the cmp doesn't define a reg and is ignored in the encoding
  def test_cmove_ignore_cmp(self):
-    cmove = ins(X86Ops.CMOVE, dtypes.int32, (def_reg(dtypes.int32, RAX), UOp(Ops.INS, arg=(X86Ops.CMP, dtypes.void))), RDX)
+    cmove = ins(X86Ops.CMOVE, dtypes.int32, (def_reg(dtypes.int32, RAX), UOp(Ops.INS, arg=X86Ops.CMP)), RDX)
    # cmove edx, eax
    self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 D0"))

@@ -7,7 +7,7 @@ from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops
 from tinygrad.renderer.isa import IselContext

 # INDEX on a register value with a constant index extracts a single element (the old GEP)
-def lane(y:UOp, i:int) -> UOp: return y.index(UOp.cconst(i, dtypes.int))
+def lane(y:UOp, i:int) -> UOp: return y.index(UOp.const(i, dtypes.int), dtype=y.dtype)

@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only x86")
 class TestIselX86(unittest.TestCase):
@@ -20,7 +20,7 @@ class TestIselX86(unittest.TestCase):
      with self.subTest(dtype=dt):
        v = [UOp.variable(str(i), 0, 0, dt) for i in range(nargs)]
        n = self.isel_rewrite(expr(*v))
-        self.assertIs(n.arg[0], op)
+        self.assertIs(n.arg, op)

  def test_cmove(self):
    a = UOp.variable("a", 0, 0, dtypes.int32)
@@ -29,9 +29,9 @@ class TestIselX86(unittest.TestCase):
    d = (a != b).where(a, b)
    f = c + d
    n = self.isel_rewrite(f)
-    self.assertTrue(n.src[0].arg[0] is X86Ops.CMOVL and n.src[1].arg[0] is X86Ops.CMOVNE)
+    self.assertTrue(n.src[0].arg is X86Ops.CMOVL and n.src[1].arg is X86Ops.CMOVNE)
    # both comparisons become the same instruction
-    self.assertTrue(n.src[0].src[2] == n.src[1].src[2] and n.src[0].src[2].arg[0] is X86Ops.CMP)
+    self.assertTrue(n.src[0].src[2] == n.src[1].src[2] and n.src[0].src[2].arg is X86Ops.CMP)

  def test_vinsertps(self):
    a = UOp.variable("a", 0, 0, dtypes.float32)
@@ -41,15 +41,15 @@ class TestIselX86(unittest.TestCase):

    valid = [UOp.stack(lane(a, 0), lane(b, 1), lane(a, 2), lane(b, 3)),
             UOp.stack(lane(a, 3), lane(b, 2), lane(c, 1), d)]
-    for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg[0], X86Ops.VINSERTPS)
+    for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VINSERTPS)

  # complex address is [base + index*scale + displacement]
  def test_complex_address(self):
    a = UOp.variable("a", 0, 0, dtypes.int32)
-    load = UOp.param(0, dtypes.int32, 16).index(a + UOp.cconst(1, dtypes.int32)).load()
+    load = UOp.param(0, dtypes.int32, (16,)).index(a + 1).load()
    n = self.isel_rewrite(load)
    # displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32
-    self.assertTrue(n.src[2].dtype is dtypes.int8 and n.src[2].src[0].op is Ops.CONST and n.src[2].src[0].val == 4)
+    self.assertTrue(n.src[2].op is Ops.CONST and n.src[2].dtype is dtypes.int8 and n.src[2].val == 4)

 if __name__ == "__main__":
  unittest.main()
@@ -6,7 +6,7 @@ from test.helpers import assert_jit_cache_len, call_is_graph, not_support_multi_
 from test.unit.test_jit import _simple_test
 from tinygrad import Tensor, Variable, TinyJit, Device, dtypes
 from tinygrad.engine.jit import graph_class
-from tinygrad.helpers import JIT, DEV, GlobalCounters, HCQ2
+from tinygrad.helpers import JIT, DEV, GlobalCounters
 from tinygrad.uop.ops import Ops
 from tinygrad.renderer.isa.x86 import X86Renderer

@@ -235,7 +235,6 @@ class TestJitPrune(unittest.TestCase):
    assert_jit_cache_len(w2_prune, 1)

 class TestJitFree(unittest.TestCase):
-  @unittest.skipIf(HCQ2, "hcq2 keeps refs to intermediate buffers")
  def test_free_intermediates(self):
    ext_tensor = Tensor([1,24,23,45,1])
    @TinyJit
@@ -16,6 +16,8 @@ from test.helpers import replace_opts, check_schedule
 from test.backend.test_softmax_fusion import single_kernel_softmax
 MOCKGPU = DEV.interface.startswith("MOCK")

+from tinygrad.uop.render import print_uops # noqa: F401 # pylint: disable=unused-import
+
@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, ISARenderer), "isa backends don't preserve the op spec when lowering")
 class TestLinearizer(unittest.TestCase):
  def test_arg_dedup(self):
@@ -246,6 +248,7 @@ class TestLinearizer(unittest.TestCase):
    uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[1].src)
    begin_range = [i for i, x in enumerate(uops) if x.op is Ops.RANGE][-1]
    end_range = [i for i, x in enumerate(uops) if x.op is Ops.END][0]
+    for i,u in enumerate(uops): print(i, u.op, [uops.index(s) for s in u.src], u.arg, u.dtype)
    for u in uops:
      if u.op is Ops.STORE and u.src[0].addrspace is AddrSpace.REG:
        if uops.index(u) < begin_range:
@@ -258,6 +261,7 @@ class TestLinearizer(unittest.TestCase):
        assert end_range < uops.index(u)

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  @unittest.skipIf(Device[Device.DEFAULT].renderer.casted_consts, "reads a literal, which is casted here. TODO: flip this")
  def test_default_global_reversed(self):
    # shrink so that the dims do not collapse
    t = Tensor.ones(5, 6, 7).contiguous().realize().shrink(((0, 4), (0, 5), (0, 6)))
@@ -265,9 +269,9 @@ class TestLinearizer(unittest.TestCase):
    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
    idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
    idxs = sorted(idxs, key=lambda uop: uop.arg)
-    assert (idxs[0].arg, idxs[0].src[0].src[0].val) == ('gidx0', 6), idxs[0]
-    assert (idxs[1].arg, idxs[1].src[0].src[0].val) == ('gidx1', 5), idxs[1].arg
-    assert (idxs[2].arg, idxs[2].src[0].src[0].val) == ('gidx2', 4), idxs[2].arg
+    assert (idxs[0].arg, idxs[0].src[0].val) == ('gidx0', 6), idxs[0]
+    assert (idxs[1].arg, idxs[1].src[0].val) == ('gidx1', 5), idxs[1].arg
+    assert (idxs[2].arg, idxs[2].src[0].val) == ('gidx2', 4), idxs[2].arg

  def test_sum_collapse(self):
    t = Tensor([2]).reshape(1, 1).expand(256, 256).sum()
@@ -11,16 +11,16 @@ from tinygrad.codegen import to_program
 class TestLinearizerFailure(unittest.TestCase):
  @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
  def test_failure_beam_mnist(self):
-    c0 = UOp.param(0, dtypes.uchar, 4014080)
+    c0 = UOp.param(0, dtypes.uchar, (4014080,))
    c1 = UOp.range(UOp.const(512), 0, AxisType.GLOBAL)
    c2 = UOp.range(UOp.const(784), 1, AxisType.GLOBAL)
    c3 = UOp.range(UOp.const(10), 3, AxisType.GLOBAL)
-    c4 = UOp.param(1, dtypes.int, 512)
+    c4 = UOp.param(1, dtypes.int, (512,))
    c5 = c4.index(c1.valid(UOp.const(True)))
    c6 = UOp.range(UOp.const(6000), 1004, AxisType.REDUCE)
    c7 = UOp.range(UOp.const(3750), 2006, AxisType.REDUCE)
    c8 = UOp.range(UOp.const(16), 2007, AxisType.GROUP_REDUCE)
-    c9 = UOp.param(2, dtypes.uchar, 47040000)
+    c9 = UOp.param(2, dtypes.uchar, (47040000,))
    c10 = c9.index((((c3*UOp.const(4704000))+c2)+(c6*UOp.const(784))).valid(UOp.const(True)))
    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(6000))+c6)+((c7*UOp.const(16))+c8)).alu(Ops.CMPLT, UOp.const(59999)).where(UOp.const(0).cast(dtypes.int), UOp.const(1).cast(dtypes.int)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(-1).cast(dtypes.int))).where(UOp.const(0).cast(dtypes.uchar), c10).reduce(c6, arg=Ops.ADD)
    c12 = c0.index((((c1*UOp.const(7840))+(c2*UOp.const(10)))+c3).valid(UOp.const(True))).store(c11).end(c1, c2, c3)
@@ -9,7 +9,7 @@ from extra.llama_kernels.swiglu import swiglu
 from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
 from extra.thunder.amd.fa import custom_fused_qkv_rope_backward, fused_qkv_rope
 from test.helpers import needs_second_gpu, assert_kernel_count
-from test.backend.test_asm_gemm import has_hipcc, is_cdna4
+from test.backend.test_asm_gemm import has_hipcc

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

+@unittest.skipUnless(has_hipcc() and Device.DEFAULT == "AMD", "requires hipcc to compile and amd device to run")
 class TestFusedQKVRoPE(unittest.TestCase):
  SHAPE = (2, 8192, 32, 8, 128)

-  def setUp(self):
-    if dtypes.bfloat16 not in Device[Device.DEFAULT].renderer.supported_dtypes(): self.skipTest("test uses bf16 inputs")
-
  def rand_bf16(self, *shape:int) -> Tensor:
    return (Tensor.randn(*shape) * 0.1).cast(dtypes.bfloat16).contiguous().realize()

-  def test_forward(self):
+  def freqs_cis(self) -> Tensor:
+    _, N, _, _, D = self.SHAPE
+    return precompute_freqs_cis(D, N * 2).cast(dtypes.bfloat16).clone().realize()
+
+  def test_llama31_8b_forward(self):
    Tensor.manual_seed(0)
-    B, N, H, H_KV, D = 1, 32, 8, 2, 16
+    B, N, H, H_KV, D = self.SHAPE
    GROUP = H // H_KV
-    freqs_cis = (Tensor.randn(1, N * 2, 1, D // 2, 2) * 0.1).cast(dtypes.bfloat16).contiguous().realize()
+    freqs_cis = self.freqs_cis()

    x = self.rand_bf16(B, N, H_KV * (GROUP + 2) * D)
    q, k, v = fused_qkv_rope(x, freqs_cis, H, H_KV, D)
@@ -129,13 +131,12 @@ class TestFusedQKVRoPE(unittest.TestCase):
      self.assertTrue(k.allclose(k_ref, atol=2e-2, rtol=0).item(), "K forward mismatch")
      self.assertTrue(v.allclose(v_ref, atol=0, rtol=0).item(), "V forward mismatch")

-  @unittest.skipUnless(has_hipcc() and is_cdna4(), "backward kernel requires hipcc to compile")
-  def test_llama31_8b(self):
+  def test_llama31_8b_backward(self):
    Tensor.manual_seed(1)
    B, N, H, H_KV, D = self.SHAPE
    PARTIALS = 2
    GROUP = H // H_KV
-    freqs_cis = precompute_freqs_cis(D, N * 2).cast(dtypes.bfloat16).clone().realize()
+    freqs_cis = self.freqs_cis()
    dq = self.rand_bf16(B, N, H, D)
    dk_partial = self.rand_bf16(B * PARTIALS, N, H_KV, D)
    dv_partial = self.rand_bf16(B * PARTIALS, N, H_KV, D)
@@ -3,10 +3,10 @@ from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variab
 from tinygrad.uop.ops import Ops, UOp, AxisType, graph_rewrite
 from tinygrad.helpers import getenv, prod, Context
 from tinygrad.nn.state import get_parameters
-from tinygrad.engine.realize import run_linear, compile_linear, lower_and_compile, pm_beam
+from tinygrad.engine.realize import run_linear, compile_linear, pm_beam, pm_compile
 import numpy as np
 from hypothesis import given, strategies as strat, settings
-from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule, assert_kernel_count, KernelCountException
+from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule, assert_kernel_count

 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
@@ -58,11 +58,6 @@ class TestMultiTensor(unittest.TestCase):
    assert X.uop.ended_ranges == X.uop.src[1:]
    (X + X).realize()

-  def test_shard_invalids_contiguous(self):
-    # every store is Invalid, so none of them should become a (empty) kernel
-    t = Tensor.invalids(8).shard(devices_2, axis=0).contiguous()
-    self.assertEqual(len([c for c in t.schedule_linear().src if c.src[0].op is Ops.SINK]), 1)
-
  @unittest.expectedFailure # TODO: fix
  def test_shard_empty(self):
    GlobalCounters.reset()
@@ -77,15 +72,15 @@ class TestMultiTensor(unittest.TestCase):
    X.shard_(devices_2, 0)
    out = (X + X)
    linear = compile_linear(out.schedule_linear())
-    uops = [call.src[0].src[0] for call in linear.src if call.src[0].op is Ops.PROGRAM]
+    names = [call.src[0].src[0].arg.name for call in linear.src if call.src[0].op is Ops.PROGRAM]
    run_linear(linear)
-    self.assertEqual(len(set(uops)), 1, "function was relinearized")
+    self.assertEqual(len(set(names)), 1, "function was relinearized")

  def test_shard_beam(self):
    cpu_2 = ("CPU:1", "CPU:2")
    src = Tensor.ones(16).shard(cpu_2, 0).realize()
    lin = UOp(Ops.LINEAR, src=(src.to(cpu_2[::-1]).schedule_linear().src[0],))
-    with Context(BEAM=1, IGNORE_BEAM_CACHE=1): call = lower_and_compile(graph_rewrite(lin, pm_beam, ctx=1, walk=True)).src[0]
+    with Context(BEAM=1, IGNORE_BEAM_CACHE=1): call = graph_rewrite(graph_rewrite(lin, pm_beam, ctx=1, walk=True), pm_compile, walk=True).src[0]
    self.assertNotEqual(call.src[0].src[0].arg.applied_opts, ())

  def test_shard_same_device(self):
@@ -187,13 +182,6 @@ class TestMultiTensor(unittest.TestCase):
        a,b = jit_allreduce(Tensor.rand(256, 256))
        np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)

-  def test_allreduce_all2all_jit(self):
-    with Context(ALL2ALL=2):
-      jit_allreduce = TinyJit(_test_allreduce)
-      for _ in range(5):
-        a,b = jit_allreduce(Tensor.rand(256, 256))
-        np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
-
  def test_multitensor_jit_input(self):
    @TinyJit
    def f(x): return (x+1).contiguous().sum()
@@ -407,7 +395,7 @@ class TestMultiBufferView(unittest.TestCase):
    linear, var_vals = b_multi.linear_with_vars()
    if all(not d.startswith(("WEBGPU", "CL")) for d in b_multi.device):
      compiled = [call for call in linear.src if call.src[0].op is Ops.SINK]
-      if len(compiled) != 0: raise KernelCountException(0, len(compiled))
+      self.assertEqual(len(compiled), 0, f"expected zero compiled kernels, got {len(compiled)}")
    run_linear(linear, var_vals)
    np.testing.assert_equal(b_multi.numpy(), b_ref.numpy())

@@ -135,7 +135,7 @@ class TestNN(unittest.TestCase):
  def test_conv2d_same_padding_large_kernel(self):
    self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=16, DIMS=[28, 33], C2=32, K=9, S=1, P='same')
  def test_conv2d_same_padding_with_dilation(self):
-    self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=3, DIMS=[28, 31], C2=32, K=(3,5), S=1, P='same', D=(2,3))
+    self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=3, DIMS=[28, 28], C2=32, K=3, S=1, P='same', D=3)

  def test_conv2d_same_padding_invalid_stride(self):
    self.assertRaises(ValueError, Conv2d, in_channels=16, out_channels=32, kernel_size=2, stride=2, padding='same')
@@ -6,7 +6,6 @@ from tinygrad.helpers import getenv, DEBUG, DEV, IMAGE, Context
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.renderer.nir import NIRRenderer
-from tinygrad.renderer.isa.x86 import X86Renderer

 TINY_BACKEND = getenv("TINY_BACKEND")
 if TINY_BACKEND:
@@ -359,13 +358,6 @@ class TestOps(unittest.TestCase):
      lambda x: torch.where(x > 0.5, 4, 2).type(torch.int32).permute((1, 0)),
      lambda x: (x > 0.5).where(4, 2).clone().permute((1, 0)), forward_only=True)

-  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "software vulkan evaluates a NaN != x as false")
-  def test_where_nan_cond(self):
-    # a NaN compares false against everything except !=.
-    for fxn in (lambda x: x<1, lambda x: x>1, lambda x: x!=1, lambda x: x==1):
-      helper_test_op(None, lambda x,a,b: torch.where(fxn(x), a, b), lambda x,a,b: fxn(x).where(a, b), forward_only=True,
-                     vals=[[math.nan, 1.0, 2.0, -1.0], [10, 20, 30, 40], [-1, -2, -3, -4]])
-
  def _test_cmp(self, fxn, reverse=True):
    # test different dtypes
    helper_test_op(None, fxn, fxn, forward_only=True, vals=[[0.,1,2], [2.,1,0]])
@@ -816,8 +808,6 @@ class TestOps(unittest.TestCase):
    helper_test_op([], lambda: tor^0x1337, lambda: ten^0x1337, forward_only=True)
    helper_test_op([], lambda: 0x1337^tor, lambda: 0x1337^ten, forward_only=True)

-  # TODO: x86 PARAM dtype fails SPEC=2
-  @Context(SPEC=1 if isinstance(Device[Device.DEFAULT].renderer, X86Renderer) else 2)
  def test_and(self):
    data = [[1,-8,1],[32,1,6]]
    tor = torch.tensor(data, dtype=torch.int)
@@ -832,10 +822,6 @@ class TestOps(unittest.TestCase):
    helper_test_op([], lambda: tor0&tor1, lambda: ten0&ten1, forward_only=True)

    helper_test_op(None, lambda x: (1 < x) & (x < 2), forward_only=True, vals=[[1.2, 1.2, 1.2, 3.2]])
-    helper_test_op([(3000,)]*10, lambda *xs: (sum(xs[1:], xs[0]) > 5) & (xs[0] < 0.9), forward_only=True)
-
-    if not COMPILE_ONLY:
-      np.testing.assert_equal((Tensor(2**64-1, dtype=dtypes.uint64) & 0xFFFFFFFF).numpy(), 0xFFFFFFFF)

  def test_or(self):
    data = [[1,-8,1],[32,1,6]]
@@ -875,9 +861,9 @@ class TestOps(unittest.TestCase):
                   lambda: (ten << Tensor([0,2,4], dtype=dtypes.uint32)).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor.__lshift__(2), lambda: ten.__lshift__(2).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor.bitwise_left_shift(2), lambda: ten.lshift(2).cast(dtypes.int32), forward_only=True)
-    self.helper_test_exception([], lambda: torch.tensor([1.0]) << 2, lambda: (Tensor([1.0]) << 2).realize(), expected=RuntimeError)
-    self.helper_test_exception([], lambda: tor << torch.tensor([1.0]), lambda: (ten << Tensor([1.0])).realize(), expected=RuntimeError)
-    self.helper_test_exception([], lambda: tor << 1.0, lambda: (ten << 1.0).realize(), expected=RuntimeError)
+    self.helper_test_exception([], lambda: torch.tensor([1.0]) << 2, lambda: Tensor([1.0]) << 2, expected=RuntimeError)
+    self.helper_test_exception([], lambda: tor << torch.tensor([1.0]), lambda: ten << Tensor([1.0]), expected=RuntimeError)
+    self.helper_test_exception([], lambda: tor << 1.0, lambda: ten << 1.0, expected=RuntimeError)

  def test_rshift(self):
    data = [[0,1,2],[1<<8,1<<16,1<<31-1]]
@@ -891,8 +877,8 @@ class TestOps(unittest.TestCase):
                   lambda: (ten >> Tensor([0,2,4], dtype=dtypes.uint32)).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor.__rshift__(2), lambda: ten.__rshift__(2).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor.bitwise_right_shift(2), lambda: ten.rshift(2).cast(dtypes.int32), forward_only=True)
-    self.helper_test_exception([], lambda: torch.tensor([4.0]) >> 1, lambda: (Tensor([4.0]) >> 1).realize(), expected=RuntimeError)
-    self.helper_test_exception([], lambda: tor >> torch.tensor([1.0]), lambda: (ten >> Tensor([1.0])).realize(), expected=RuntimeError)
+    self.helper_test_exception([], lambda: torch.tensor([4.0]) >> 1, lambda: Tensor([4.0]) >> 1, expected=RuntimeError)
+    self.helper_test_exception([], lambda: tor >> torch.tensor([1.0]), lambda: ten >> Tensor([1.0]), expected=RuntimeError)

  def test_lshift_signed(self):
    data = [[-1, -3, 1, 7], [0, -2147483648, 2147483647, -1]]
@@ -1817,10 +1803,9 @@ class TestOps(unittest.TestCase):
      helper_test_op([()], lambda x: torch.nn.functional.hardtanh(x, -val, val), lambda x: x.hardtanh(-val, val), grad_atol=1e-6)
  def test_asinh(self):
    helper_test_op([(45,65)], lambda x: x.asinh(), grad_atol=1e-6)
-    helper_test_op([(45,65)], lambda x: x.asinh(), grad_atol=1e-6, low=-300, high=-297)
+    # TODO: this one has larger tol?
+    helper_test_op([(45,65)], lambda x: x.asinh(), atol=1e-2, rtol=2e-2, grad_rtol=2e-2, low=-300, high=-297)
    helper_test_op([(45,65)], lambda x: x.asinh(), grad_atol=1e-6, low=300, high=303)
-    helper_test_op([(45,65)], lambda x: x.asinh(), grad_atol=1e-6, low=-1e10, high=-1e9)
-    helper_test_op(None, lambda x: x.asinh(), grad_atol=1e-6, vals=[[-1.0, 0.0, 1.0]])
  def test_acosh(self):
    helper_test_op([(45,65)], lambda x: x.acosh(), grad_atol=1e-6)
    helper_test_op([(45,65)], lambda x: x.acosh(), grad_atol=1e-3, grad_rtol=1e-2, low=-300, high=-297)
@@ -2179,10 +2164,6 @@ class TestOps(unittest.TestCase):
  def test_roll(self):
    helper_test_op([(2, 4)], lambda x: x.roll(1))
    helper_test_op([(2, 4)], lambda x: x.roll((1,)))
-    helper_test_op([(0,)], lambda x: x.roll(1, 0))
-    helper_test_op([(2, 0, 3)], lambda x: x.roll(1, 0))
-    helper_test_op([(2, 0, 3)], lambda x: x.roll(1, 1))
-    helper_test_op([(2, 0, 3)], lambda x: x.roll(1))
    self.helper_test_exception([(2, 4)], lambda x: x.roll((1, 2)), expected=RuntimeError)
    helper_test_op([(2, 4)], lambda x: x.roll(1, 0))
    helper_test_op([(2, 4)], lambda x: x.roll(-1, 0))
@@ -2817,7 +2798,7 @@ class TestOps(unittest.TestCase):
        lambda x: Tensor.interpolate(x, size=out_sz, mode="linear"))

  def test_interpolate_linear_corners_aligned(self):
-    for in_sz, out_sz in [((52,),(29,)), ((29,),(52,)), ((29,),(1,))]:
+    for in_sz, out_sz in [((52,),(29,)), ((29,),(52,))]:
      helper_test_op([(2,3)+in_sz],
        lambda x: torch.nn.functional.interpolate(x, size=out_sz, mode="linear", align_corners=True),
        lambda x: Tensor.interpolate(x, size=out_sz, mode="linear", align_corners=True))
@@ -2970,10 +2951,6 @@ class TestOps(unittest.TestCase):
    data = [math.inf, -math.inf, math.nan]
    helper_test_op((), lambda: torch.tensor(data)[torch.tensor([0, 1, 2])], lambda: Tensor(data)[Tensor([0, 1, 2])])

-  def test_fancy_indexing_index_dtypes(self):
-    helper_test_op((), lambda: torch.tensor([10., 20., 30., 40.])[torch.tensor([1, 2, 3, 0])],
-                       lambda: Tensor([10., 20., 30., 40.])[Tensor([1, 2, 3, 0], dtype=dtypes.uint8)])
-
  @slow_test
  def test_slice_fancy_indexing_no_dim_collapse(self):
    a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
@@ -87,8 +87,7 @@ class TestOptim(unittest.TestCase):
  def test_muon(self): self._test_muon(1, {'lr': 0.001}, 1e-3, 0)
  # TODO: disabled due to big atol
  # def test_muon_high_lr(self): self._test_muon(1, {'lr': 10}, 1e-6, 3e-4)
-  # NOTE: big weight_decay so a missing wd would be way over atol
-  def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 10}, 1e-3, 3e-4)
+  def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 0.01}, 1e-3, 3e-4)
  # TODO: disabled due to big atol
  # def test_muon_high_lr_wd(self): self._test_muon(1, {'lr': 10, 'weight_decay': 0.01}, 1e-6, 5e-4)

@@ -3,7 +3,6 @@ import numpy as np
 from tinygrad import Tensor, Device, TinyJit, Variable, dtypes
 from tinygrad.helpers import GlobalCounters, ContextVar, Context, DEV
 from tinygrad.uop.ops import PatternMatcher, UPat, UOp, deconstruct_function
-from test.helpers import KernelCountException

 class TestPickle(unittest.TestCase):
  def test_pickle_code_object(self):
@@ -42,7 +41,7 @@ class TestPickle(unittest.TestCase):
    t2:Tensor = pickle.loads(st)
    np.testing.assert_equal(t_values, t2.numpy())
    # expect at most one COPY kernel
-    if GlobalCounters.kernel_count > 1: raise KernelCountException(1, GlobalCounters.kernel_count)
+    self.assertLessEqual(GlobalCounters.kernel_count, 1)

  def test_pickle_realized_tensor_alt(self):
    print("** init")
@@ -1,6 +1,6 @@
 import unittest
-from tinygrad import Tensor, Device, dtypes, Variable
-from tinygrad.helpers import Context, GlobalCounters, getenv, DEBUG
+from tinygrad import Tensor, nn, Device, dtypes, Variable
+from tinygrad.helpers import Context, GlobalCounters, getenv, PCONTIG, DEBUG
 from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops, UOp
 from tinygrad.codegen.opt import OptOps, Opt
 from tinygrad.renderer.ptx import PTXRenderer
@@ -14,7 +14,7 @@ class TestDoubleMatmul(unittest.TestCase):
      self.ref = (self.a @ self.b @ self.c).realize()

  def _test(self, opts):
-    with Context(DEBUG=max(2, DEBUG.value)):
+    with Context(PCONTIG=2, DEBUG=max(2, DEBUG.value)):
      out = (self.a @ self.b @ self.c).contiguous(arg=opts).realize()

    with Context(DEBUG=0):
@@ -88,15 +88,16 @@ class TestRangeifyEdgeCase(unittest.TestCase):
    res = Tensor.cat(a, c, dim=0)
    self.assertEqual(res.numpy()[-1, :16].tolist(), [512] * 16)

-  def test_multi_gather(self):
+  def test_pcontig_multi_gather(self):
    # regression test: local bufferize must have device set for const_like to work
-    # NOTE: with uint type, this will become a long and fail on WEBGPU
-    forest = Tensor(list(range(8)), dtype='int')
-    idx = Tensor([0, 0], dtype='int')
-    node_val = forest.gather(0, idx)
-    idx2 = idx * 2 + 1
-    node_val2 = forest.gather(0, idx2)
-    result = (node_val + node_val2).numpy()
+    with Context(PCONTIG=2):
+      # NOTE: with uint type, this will become a long and fail on WEBGPU
+      forest = Tensor(list(range(8)), dtype='int')
+      idx = Tensor([0, 0], dtype='int')
+      node_val = forest.gather(0, idx)
+      idx2 = idx * 2 + 1
+      node_val2 = forest.gather(0, idx2)
+      result = (node_val + node_val2).numpy()
    self.assertEqual(result.tolist(), [1, 1])

 if getenv("BIG") > 2:
@@ -117,6 +118,65 @@ def fa():
  GlobalCounters.reset()
  return q.scaled_dot_product_attention(k, v)

+def fa_bw():
+  Tensor.manual_seed(1337)
+  with Context(DEBUG=0):
+    q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
+    attn_output = nn.Linear(HEADS*EMB, HEADS*EMB, bias=False)
+    attn_output.weight.realize()
+    target = Tensor.rand(BS, SEQLEN, HEADS*EMB).contiguous().realize()
+
+  GlobalCounters.reset()
+  attn = q.scaled_dot_product_attention(k, v).contiguous().contiguous_backward()
+  attn = attn.transpose(1, 2).reshape(BS, SEQLEN, -1)
+  out = attn_output(attn)
+  loss = (out - target).square().mean()
+  loss.backward()
+  #ret = [out, Tensor.stack(q.grad, k.grad, v.grad, dim=-1)]
+  #ret = [out, Tensor.stack(q.grad, k.grad, dim=-1), v.grad]
+  ret = [out, q.grad, k.grad, v.grad]
+  Tensor.realize(*ret)
+  return ret
+
+@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "broken in LVP and PTX")
+class TestPcontig(unittest.TestCase):
+  def test_flash_attention_bw(self):
+    with Context(PCONTIG=max(2, PCONTIG.value), DEBUG=2):
+      grads = fa_bw()
+      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
+
+    with Context(PCONTIG=0, DEBUG=2):
+      cmp_grads = fa_bw()
+      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
+
+    with Context(DEBUG=0):
+      mses = [((x-y)**2).sum().item() for x,y in zip(grads, cmp_grads)]
+    mse = sum(mses)
+    print(f"mse: {mse}")
+    self.assertLessEqual(mse, 1e-6)
+
+  def test_flash_attention(self, opts=None):
+    with Context(PCONTIG=2, DEBUG=max(2, DEBUG.value)):
+      ret = fa().realize() if opts is None else fa().contiguous(arg=opts).realize()
+      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
+    with Context(DEBUG=2):
+      cmp = fa().realize()
+      print(f"{GlobalCounters.global_ops/1e9:.2f} GFLOPS")
+    with Context(DEBUG=0):
+      mse = ((cmp-ret)**2).sum().item()
+    print(f"mse: {mse}")
+    self.assertLessEqual(mse, 1e-6)
+
+  def test_flash_attention_opt(self):
+    opts = ()
+    # columns in top matrix
+    opts += (Opt(OptOps.UPCAST, 0, 4),)
+    # columns in bottom matrix
+    opts += (Opt(OptOps.UPCAST, 3, 4),)
+    # rows in all the matrix
+    opts += (Opt(OptOps.UPCAST, 4, 4),)
+    self.test_flash_attention(opts)
+
 # contiguous + reduce can support ranges?

@unittest.skip("pm_rangeify no longer exists. test this in a different way")
@@ -23,8 +23,8 @@ def _test_uop_result(inputs:list[Tensor], sink:UOp, local_size=None):

 def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
  dtype = alu_src_uops[0].dtype
-  a = UOp.param(0, dtype, 1)
-  b = UOp.param(1, dtype, 1)
+  a = UOp.param(0, dtype, (1,))
+  b = UOp.param(1, dtype, (1,))
  idx = UOp.const(0)
  ld = b.index(idx).load()
  alu = ld.alu(alu_op, *alu_src_uops)
@@ -34,7 +34,7 @@ def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
 class TestRendererFailures(unittest.TestCase):
  @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
  def test_gated_store_with_alu(self):
-    a = UOp.param(0, dtypes.int, 4)
+    a = UOp.param(0, dtypes.int, (4,))
    gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
    gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0.valid(gate_alu)), UOp.const(1).cast(dtypes.int)))
    sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo())
@@ -43,7 +43,7 @@ class TestRendererFailures(unittest.TestCase):

  @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
  def test_gated_store_with_alu_2d(self):
-    a = UOp.param(0, dtypes.int, 8)
+    a = UOp.param(0, dtypes.int, (8,))
    gate_alu_0 = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
    gate_alu_1 = (lidx1:=UOp.special(2, 'lidx1')).ne(0)
    gated_alu_store = UOp(Ops.STORE, src=(a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(1).cast(dtypes.int)))
@@ -77,14 +77,6 @@ class TestCStyleFailures(unittest.TestCase):

@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "tests for wgsl renderer")
 class TestWGSLFailures(unittest.TestCase):
-  def test_folded_packed_store(self):
-    b = UOp.param(0, dtypes.char, 4)
-    idx = b.index(UOp.const(0).cast(dtypes.int))
-    store = UOp.store(idx, idx.cast(dtypes.uint32).load() & UOp.const(0xffffff00).cast(dtypes.uint32))
-    src = Device[Device.DEFAULT].renderer.render(UOp.sink(store, arg=KernelInfo()).toposort())
-    self.assertIn("atomicAnd(&data0_4[0],4294967040u);", src)
-    self.assertNotIn("atomicAdd", src)
-
  def test_multiply_infinity(self):
    # multiplying a positive constant by infinity should return infinity
    # WGSL pipelines do not handle this reliably, some of which return zero, unless infinity always comes from a read on a dynamic buffer
@@ -93,9 +85,9 @@ class TestWGSLFailures(unittest.TestCase):

  # WGSL has a specific select(alt, val, gate) ternary operator instead of gate?val:alt
  def test_gated_load(self):
-    a = UOp.param(0, dtypes.int, 4)
-    b = UOp.param(1, dtypes.int, 4)
-    c = UOp.param(2, dtypes.int, 4)
+    a = UOp.param(0, dtypes.int, (4,))
+    b = UOp.param(1, dtypes.int, (4,))
+    c = UOp.param(2, dtypes.int, (4,))
    lidx0 = UOp.special(4, "lidx0")
    gate = lidx0.ne(0)
    alt = c.index(lidx0).load()
@@ -110,7 +102,7 @@ class TestWGSLFailures(unittest.TestCase):
 class TestPTXFailures(unittest.TestCase):
  @unittest.skip("INDEX can only have a gate ALU parent, not an IF")
  def test_gated_store_with_if(self):
-    a = UOp.param(0, dtypes.int, 4)
+    a = UOp.param(0, dtypes.int, (4,))
    gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
    val = UOp.const(1).cast(dtypes.int)
    if_uop = UOp(Ops.IF, src=(gate_alu,))
@@ -147,7 +147,7 @@ class TestSchedule(unittest.TestCase):
    devs = ("CPU:0", "CPU:1")
    x = Tensor.ones(2, device="CPU").shard(devs, axis=0).realize()
    out = (x.sum()*2).reshape(1).to("CPU")
-    run_linear(*check_schedule(out, 3))
+    run_linear(*check_schedule(out, 5))
    np.testing.assert_equal(out.numpy(), [4.])

 class TestLimitBufs(unittest.TestCase):
@@ -176,7 +176,7 @@ class TestLimitBufs(unittest.TestCase):

  def test_limit_bufs_linear_scaling(self):
    def sched_time(n):
-      with Context(TRACK_MATCH_STATS=0, DEBUG=0, PARALLEL=0):
+      with Context(TRACK_MATCH_STATS=0, DEBUG=0):
        bufs = [Tensor.ones(16).contiguous().realize() for _ in range(4)]
        root = bufs[0]
        for i in range(n): root = root + bufs[i % 4]
@@ -365,6 +365,16 @@ class TestCopyFolding(unittest.TestCase):
    b = a.to("CPU")
    self.assertListEqual(b.tolist(), [2.])

+  def test_copy_to_same_device(self):
+    a = Tensor.empty(4).uop
+    b = a.copy_to_device(a.device)
+    check_schedule(b, 1, filter_sink=False) # TODO: 0?
+
+  def test_copy_to_same_device_alt(self):
+    a = Tensor.empty(4, 4).uop
+    b = a.copy_to_device(a.device)
+    check_schedule(b, 1, filter_sink=False) # TODO: 0?
+
  def test_copy_to_same_device_sched(self):
    a = Tensor.ones(4).contiguous().realize().uop.buf_uop
    t = Tensor(a.copy_to_device(a.device))
@@ -374,35 +384,40 @@ class TestCopyFolding(unittest.TestCase):
    assert t.uop.is_realized, f"didn't realize Tensor {t}"
    self.assertListEqual(t.tolist(), [1.,1.,1.,1.])

+  @unittest.skip("same-device copies are no-ops")
+  def test_self_assign_same_device_copy(self):
+    a = Tensor.ones(4, 4).contiguous().realize()
+    # use copy_to_device to bypass Tensor.to() shortcircuit and force a real same-device COPY in the graph
+    a.assign(Tensor(a.uop.copy_to_device(a.device), a.device))
+    run_linear(*check_schedule(a, 2, filter_sink=False))
+    self.assertListEqual(a.tolist(), [[1.]*4]*4)
+
  def test_clone(self):
    a = Tensor.empty(4)
    check_schedule(a.clone(), 1, filter_sink=False)

  def test_shrink_copy(self):
-    a = Tensor.arange(4).clone("CPU:1").realize()
-    b = a.to("CPU:2").shrink(((1, 3),)).to("CPU:3")
-    GlobalCounters.reset()
-    run_linear(*check_schedule(b, 3, filter_sink=False))
-    # extra E kernel, copy exactly 4 bytes
-    self.assertEqual(GlobalCounters.global_mem, 4*4 + 2*4*2 + 2*4)
-    self.assertListEqual(b.tolist(), [1, 2])
+    a = Tensor.arange(4)
+    view = a.shrink(((0, 2),))
+    b = view.clone()
+    run_linear(*check_schedule(b, 1, filter_sink=False))
+    self.assertEqual(b.uop.base.buffer.size, 2)
+    self.assertEqual(b.uop.numel(), 2)
+    self.assertListEqual(b.tolist(), [0, 1])

  def test_expanded_copy(self):
-    a = Tensor.arange(4).clone("CPU:1").realize()
-    b = a.to("CPU:2").reshape(4, 1).expand(4, 2).to("CPU:3")
-    GlobalCounters.reset()
-    run_linear(*check_schedule(b, 3, filter_sink=False))
-    # TODO: expands before copy
-    self.assertEqual(GlobalCounters.global_mem, 4*4 + (4*4 + 8*4) + 8*4)
-    self.assertListEqual(b.tolist(), [[0, 0], [1, 1], [2, 2], [3, 3]])
+    a = Tensor.arange(2)
+    view = a.reshape(2, 1).expand(2, 2)
+    b = view.clone()
+    run_linear(*check_schedule(b, 1, filter_sink=False))
+    self.assertEqual(b.uop.base.buffer.size, 4)
+    self.assertEqual(b.uop.numel(), 4)
+    self.assertListEqual(b.tolist(), [[0, 0], [1, 1]])

  def test_permuted_copy(self):
-    a = Tensor.arange(4).clone("CPU:1").realize()
-    b = a.to("CPU:2").reshape(2, 2).permute(1, 0).to("CPU:3")
-    GlobalCounters.reset()
-    run_linear(*check_schedule(b, 3, filter_sink=False))
-    # permutes before copy
-    self.assertEqual(GlobalCounters.global_mem, 4*4 + (4*4 + 4*4) + 4*4)
+    a = Tensor.arange(4)
+    b = a.reshape(2, 2).permute(1, 0)
+    b.realize()
    self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])

  def test_permute_on_disk(self):
@@ -301,14 +301,6 @@ class TestSetitem(unittest.TestCase):
    self.assertListEqual(z[2:5].tolist(), [2, 2, 2])
    self.assertListEqual(z[6:7].tolist(), [3])

-class TestAssignBitcast(unittest.TestCase):
-  def test_assign_through_bitcast(self):
-    # the dest is unrealized, so callify cannot fold the BITCAST into a buffer view and the STORE keeps a
-    # BITCAST dest; the bitcast has to move to the value side or the store never reaches the buffer
-    a = Tensor.full((4,), 1.0, dtype=dtypes.float32).contiguous()
-    a.bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.uint32)).realize()
-    np.testing.assert_allclose(a.numpy(), [4.0, 3.0, 2.0, 1.0])
-
 class TestWithGrad(unittest.TestCase):
  def test_basic_setitem_works(self):
    z = Tensor.rand(8, 8)
@@ -2,7 +2,7 @@ from typing import Optional, Any
 import unittest, math
 import numpy as np
 from tinygrad.tensor import Tensor, _to_np_dtype
-from tinygrad.helpers import Context, ceildiv
+from tinygrad.helpers import Context
 from tinygrad.dtype import dtypes, DType, AddrSpace, ConstFloat  # noqa: F401
 from tinygrad.device import Buffer, Device
 from tinygrad.uop.ops import Ops, UOp, KernelInfo, AxisType, buffers
@@ -20,8 +20,8 @@ def run_uops(uops_list:list[UOp], bufs:list[Buffer]):

 def uop(uops:list[UOp], op:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
  if op is Ops.CONST: uops.append(UOp.const(arg).cast(dtype))
-  elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, 1))
-  else: uops.append(UOp(op, tuple(src), arg))
+  elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, shape=(1,)))
+  else: uops.append(UOp(op, dtype, tuple(src), arg))
  return uops[-1]

 def _test_single_value(vals, op, dts):
@@ -57,35 +57,6 @@ def _test_uops_result(output_dtype, uops, res):
  run_uops([out], [buf])
  return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0]

-@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, CStyleLanguage) and
-                     dtypes.uint64 in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires C-style pointer bitcast and 64-bit ints")
-class TestBitcastBufferView(unittest.TestCase):
-  @Context(SPEC=2)
-  def test_render(self):
-    buf = UOp.param(0, dtypes.uint32, 4)
-    uops = to_uops_list([buf.shrink(((1, 3),)).bitcast(dtypes.uint64).index(0).store(1)], ren=Device[Device.DEFAULT].renderer)
-    idx = next(u for u in uops if u.op is Ops.INDEX and u.src[0].op is Ops.BITCAST)
-    self.assertEqual(idx.src[0].src[0].op, Ops.SHRINK)
-    Device[Device.DEFAULT].renderer.render(uops)
-
-  @Context(SPEC=2)
-  def test_load(self):
-    val = 0x1122334455667788
-    src, out = UOp.param(0, dtypes.uint32, 4), UOp.param(1, dtypes.uint64, 1)
-    ibuf = Buffer(Device.DEFAULT, 4, dtypes.uint32, initial_value=np.array([0, 0x55667788, 0x11223344, 0], dtype=np.uint32).tobytes())
-    obuf = Buffer(Device.DEFAULT, 1, dtypes.uint64).allocate()
-    run_uops([out.index(0).store(src.shrink(((1, 3),)).bitcast(dtypes.uint64).index(0))], [ibuf, obuf])
-    self.assertEqual(np.frombuffer(obuf.as_memoryview(), dtype=np.uint64)[0], val)
-
-  @Context(SPEC=2)
-  def test_store(self):
-    val = 0x1122334455667788
-    dst = UOp.param(0, dtypes.uint32, 6)
-    buf = Buffer(Device.DEFAULT, 6, dtypes.uint32, initial_value=bytes(24))
-    view = dst.shrink(((1, 5),)).bitcast(dtypes.uint64)  # two stores through one view: it must inline, not get a declared vector-pointer
-    run_uops([view.index(0).store(val ^ 0xff), view.index(1).store(val)], [buf])
-    self.assertEqual(np.frombuffer(buf.as_memoryview(), dtype=np.uint64, count=2, offset=4).tolist(), [val ^ 0xff, val])
-
 class TestUOps(unittest.TestCase):
  def _equal(self, v1, v2):
    assert isinstance(v2, (float, int, bool))
@@ -222,16 +193,15 @@ class TestLocalAccess(unittest.TestCase):
  @unittest.skipUnless(Device.DEFAULT == "WEBGPU", "Test local memory size for packed data types")
  def test_packed_smem_size(self):
    _dtypes = [dtypes.char, dtypes.uchar, dtypes.short, dtypes.ushort, dtypes.half]
-    # a partial word still needs a whole word, so sizes that don't fill one must round up
-    for size in (16, 5):
-      for dtype in _dtypes:
-        temp = UOp.placeholder((size,), dtype, slot=0, addrspace=AddrSpace.LOCAL)
-        uops = to_uops_list([temp], ren=Device[Device.DEFAULT].renderer)
-        out = Device[Device.DEFAULT].renderer.render(uops)
-        # half is supported in wgsl, so it doesn't have to be packed
-        corrected_size = ceildiv(size, 4//dtype.itemsize) if dtype != dtypes.half else size
-        # temp0: array<{Device[Device.DEFAULT].renderer.buf_map(dtype)},{corrected_size}>;
-        self.assertIn(f",{corrected_size}>;", out)
+    size = 16
+    for dtype in _dtypes:
+      temp = UOp.placeholder((size,), dtype, slot=0, addrspace=AddrSpace.LOCAL)
+      uops = to_uops_list([temp], ren=Device[Device.DEFAULT].renderer)
+      out = Device[Device.DEFAULT].renderer.render(uops)
+      # half is supported in wgsl, so it doesn't have to be packed
+      corrected_size = size//(4//dtype.itemsize) if dtype != dtypes.half else size
+      # temp0: array<{Device[Device.DEFAULT].renderer.buf_map(dtype)},{corrected_size}>;
+      self.assertIn(f",{corrected_size}>;", out)

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared memory")
  @unittest.skip("tinygrad doesn't support this behavior")
@@ -249,8 +219,8 @@ class TestLocalAccess(unittest.TestCase):
@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "This only tests assembly backends")
 class TestAssembly(unittest.TestCase):
  def test_bitshift_left(self):
-    g1 = UOp.param(0, dtypes.int32, 3)
-    out = UOp.param(1, dtypes.int32, 2)
+    g1 = UOp.param(0, dtypes.int32, shape=(3,))
+    out = UOp.param(1, dtypes.int32, shape=(2,))
    c1 = UOp.const(2)
    c2 = UOp.const(3)
    l1 = g1.index(c1)
@@ -278,7 +248,7 @@ class TestAssembly(unittest.TestCase):
    self.assertGreaterEqual(len([x.op for x in uops if x.op is Ops.MULACC]), 4)

  def test_mulacc_shl(self):
-    g1 = UOp.param(0, dtypes.int32, 2)
+    g1 = UOp.param(0, dtypes.int32, shape=(2,))
    c1 = UOp.const(0)
    c2 = UOp.const(1)
    expr = g1.index(c1) * UOp.const(4096) + g1.index(c2)
@@ -287,7 +257,7 @@ class TestAssembly(unittest.TestCase):
    self.assertIn(Ops.MULACC, [x.op for x in uops])

  def test_use_cmpeq(self):
-    g = UOp.param(0, dtypes.uint32, 8)
+    g = UOp.param(0, dtypes.uint32, shape=(8,))
    c = UOp.const(7)
    comp = g.index(c).ne(c).ne(True)
    uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer)
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
sirhcm	c0d9c79668	add ninja to setup-tinygrad Autogen / In-tree Autogen (comgr 2) (pull_request) Successful in 1m30s Details Autogen / In-tree Autogen (pull_request) Failing after 2m36s Details Check Line Counts / Check PR Branch status (pull_request_target) Successful in 15s Details Check Line Counts / Core Library Line Difference (pull_request_target) Failing after 4s Details Unit Tests / Docs (pull_request) Successful in 3m10s Details Unit Tests / Torch Backend Tests (pull_request) Successful in 6m48s Details Unit Tests / Torch Backend Training (pull_request) Successful in 6m11s Details Unit Tests / Python Backend (pull_request) Successful in 4m7s Details Unit Tests / Linters (pull_request) Successful in 4m17s Details Unit Tests / Null Tests (pull_request) Successful in 7m19s Details Unit Tests / Unit Tests (pull_request) Successful in 7m9s Details Unit Tests / SPEC=2 (1) (pull_request) Successful in 6m58s Details Unit Tests / SPEC=2 (2) (pull_request) Successful in 7m12s Details Unit Tests / Fuzzing (pull_request) Successful in 3m23s Details Unit Tests / CL IMAGE Tests (pull_request) Successful in 6m20s Details Unit Tests / openpilot Compile Tests (pull_request) Successful in 7m27s Details Unit Tests / ONNX (CPU) Tests (pull_request) Failing after 9m49s Details Unit Tests / Optimization Tests (pull_request) Successful in 8m0s Details Unit Tests / Test LLM (pull_request) Successful in 4m9s Details Unit Tests / Models (pull_request) Successful in 5m16s Details Unit Tests / Linux (DSP) (pull_request) Successful in 4m22s Details Unit Tests / Linux (DEV=CL) (pull_request) Successful in 6m53s Details Unit Tests / Linux (DEV=CPU:CLANG) (pull_request) Successful in 7m14s Details Unit Tests / Linux (DEV=CPU:LLVM) (pull_request) Successful in 6m49s Details Unit Tests / Linux (DEV=CPU:LVP) (pull_request) Successful in 6m38s Details Unit Tests / Linux (DEV=CPU:X86) (pull_request) Successful in 5m52s Details Unit Tests / Linux (DEV=WEBGPU) (pull_request) Successful in 6m40s Details Unit Tests / AMD ASM IDE (pull_request) Successful in 5m43s Details Unit Tests / hcq2 (pull_request) Failing after 4m18s Details Unit Tests / Linux (amd gfx1100) (pull_request) Successful in 7m36s Details Unit Tests / Linux (amdllvm gfx1100) (pull_request) Successful in 7m0s Details Unit Tests / Linux (amd gfx1201) (pull_request) Successful in 7m36s Details Unit Tests / Linux (amdllvm gfx1201) (pull_request) Successful in 7m11s Details Unit Tests / Linux (amd gfx950) (pull_request) Successful in 7m1s Details Unit Tests / Linux (amdllvm gfx950) (pull_request) Successful in 4m30s Details Unit Tests / Linux (nv) (pull_request) Successful in 5m12s Details Unit Tests / Linux (ptx) (pull_request) Successful in 3m50s Details Unit Tests / Compile-only (DEV=NULL:IR3:a630) (pull_request) Successful in 2m28s Details Unit Tests / Compile-only (DEV=NULL:NAK:sm_120) (pull_request) Successful in 1m53s Details Unit Tests / Compile-only (DEV=NULL:QCOMCL:a630) (pull_request) Failing after 9m25s Details Unit Tests / Linux (am) (pull_request) Successful in 3m52s Details Autogen / In-tree Autogen (macos) (pull_request) Canceled after 0s Details Platform Tests / MacOS (unit) (pull_request) Canceled after 0s Details Platform Tests / MacOS (unit, mock) (pull_request) Canceled after 0s Details Platform Tests / MacOS (DEV=METAL) (1) (pull_request) Canceled after 0s Details Platform Tests / MacOS (DEV=METAL) (2) (pull_request) Canceled after 0s Details Platform Tests / MacOS (DEV=CPU:CLANG) (pull_request) Canceled after 0s Details Platform Tests / MacOS (DEV=CPU:LLVM) (pull_request) Canceled after 0s Details Platform Tests / MacOS (DEV=CPU:LVP) (pull_request) Canceled after 0s Details Platform Tests / MacOS (DEV=WEBGPU) (pull_request) Canceled after 0s Details Platform Tests / Windows (DEV=CPU:CLANG) (pull_request) Canceled after 0s Details Platform Tests / Windows (DEV=CPU:LLVM) (pull_request) Canceled after 0s Details Platform Tests / Windows (DEV=CPU:X86) (pull_request) Canceled after 0s Details Platform Tests / Windows (DEV=WEBGPU) (pull_request) Canceled after 0s Details	2026-08-18 15:53:24 -07:00
sirhcm	19dc868750	ci: cleanup deps	2026-08-18 15:48:48 -07:00