also skip casted consts in expand_broadcast

don't expand consts in expand_broadcast, fix SPEC=2
remove broadcast handling from rangeify, prepare handles it now
2026-08-26 15:22:07 -07:00 · 2026-08-26 15:10:01 -07:00 · 2026-08-26 13:18:24 -07:00 · 2026-08-26 13:08:41 -07:00 · 2026-08-26 13:06:04 -07:00 · 2026-08-26 11:13:23 -07:00
124 changed files with 1327 additions and 7797 deletions
@@ -194,29 +194,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 +240,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 +261,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 +270,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 +285,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
@@ -94,7 +94,7 @@ jobs:
        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
@@ -140,7 +140,7 @@ jobs:
        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
@@ -188,7 +188,7 @@ jobs:
        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
@@ -230,7 +230,7 @@ jobs:
        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
@@ -275,7 +275,7 @@ jobs:
        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
@@ -446,55 +446,14 @@ jobs:
    - name: UsbGPU (USB4/TB) tiny tests
      run: 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 +464,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

@@ -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
@@ -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"
@@ -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
@@ -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,7 +1770,7 @@ 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)

@@ -1783,18 +1784,22 @@ def train_gptoss():
    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 +1848,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 +1881,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 +1891,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,
@@ -282,14 +282,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
@@ -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)
@@ -44,7 +44,7 @@ export SEED=${SEED:-5760}
 export DATA_SEED=${DATA_SEED:-5760}

 export JITBEAM=${JITBEAM:-3}
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=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}
 if [ -z "$FULL_LAYERS" ]; then
@@ -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
@@ -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

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

@@ -20,39 +20,34 @@ def v_mfma_fp4(dst, a, b, opsel, opsel_hi, scale_a, scale_b):
 def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
  k = Kernel()
  scale_k = K // 32
-  k.emit(s_and_b32(s[1], s[1], LIT, 65535))
  if (tile_m, tile_n) == (128, 512):
+    k.emit(s_and_b32(s[1], s[1], LIT, 65535))
    k.emit(s_mov_b32(s[47], s[2]))
    k.emit(s_mov_b32(s[48], s[3]))
-  k.emit(s_load_dwordx2(s[4:5], s[0:1], s[0], 0, 0, 0, 0, 1))
-  k.emit(s_mov_b32(s[8], 0))
-  k.emit(s_mov_b32(s[9], 0))
-  k.emit(s_load_dwordx2(s[12:13], s[0:1], s[0], 8, 0, 0, 0, 1))
-  k.emit(s_load_dwordx2(s[16:17], s[0:1], s[0], 16, 0, 0, 0, 1))
-  k.emit(s_mov_b32(s[36], N))
-  k.emit(s_mov_b32(s[37], K))
-  k.emit(s_mov_b32(s[38], K))
-  k.emit(s_mov_b32(s[43], M))
-  k.emit(s_mov_b32(s[44], N))
-  k.emit(s_mov_b32(s[45], K))
-  k.emit(s_load_dwordx2(s[20:21], s[0:1], s[0], 24, 0, 0, 0, 1))
-  k.emit(s_load_dwordx2(s[24:25], s[0:1], s[0], 32, 0, 0, 0, 1))
-  k.emit(s_mov_b32(s[39], scale_k))
-  k.emit(s_mov_b32(s[40], scale_k))
-  k.emit(v_lshrrev_b32_e32(v[1], 10))
-  k.emit(v_lshrrev_b32_e32(v[2], 10, v[1]))
-  k.emit(v_and_b32_e32(v[2], LIT, v[2], 1023))
-  k.emit(v_and_b32_e32(v[1], LIT, v[1], 1023))
-  k.emit(v_and_b32_e32(v[0], LIT, v[0], 1023))
-  k.emit(v_lshrrev_b32_e32(v[3], 6))
-  k.emit(v_and_b32_e32(v[0], 63))
-  if (tile_m, tile_n) == (256, 256):
-    k.emit(s_mov_b32(s[49], s[2]))
-    k.emit(s_mov_b32(s[47], s[3]))
-  k.emit(v_readfirstlane_b32_e32(v[46], v[3]))
-  k.emit(s_waitcnt(49279))
-
-  if (tile_m, tile_n) == (128, 512):
+    k.emit(s_load_dwordx2(s[4:5], s[0:1], s[0], 0, 0, 0, 0, 1))
+    k.emit(s_mov_b32(s[8], 0))
+    k.emit(s_mov_b32(s[9], 0))
+    k.emit(s_load_dwordx2(s[12:13], s[0:1], s[0], 8, 0, 0, 0, 1))
+    k.emit(s_load_dwordx2(s[16:17], s[0:1], s[0], 16, 0, 0, 0, 1))
+    k.emit(s_mov_b32(s[36], N))
+    k.emit(s_mov_b32(s[37], K))
+    k.emit(s_mov_b32(s[38], K))
+    k.emit(s_mov_b32(s[43], M))
+    k.emit(s_mov_b32(s[44], N))
+    k.emit(s_mov_b32(s[45], K))
+    k.emit(s_load_dwordx2(s[20:21], s[0:1], s[0], 24, 0, 0, 0, 1))
+    k.emit(s_load_dwordx2(s[24:25], s[0:1], s[0], 32, 0, 0, 0, 1))
+    k.emit(s_mov_b32(s[39], scale_k))
+    k.emit(s_mov_b32(s[40], scale_k))
+    k.emit(v_lshrrev_b32_e32(v[1], 10))
+    k.emit(v_lshrrev_b32_e32(v[2], 10, v[1]))
+    k.emit(v_and_b32_e32(v[2], LIT, v[2], 1023))
+    k.emit(v_and_b32_e32(v[1], LIT, v[1], 1023))
+    k.emit(v_and_b32_e32(v[0], LIT, v[0], 1023))
+    k.emit(v_lshrrev_b32_e32(v[3], 6))
+    k.emit(v_and_b32_e32(v[0], 63))
+    k.emit(v_readfirstlane_b32_e32(v[46], v[3]))
+    k.emit(s_waitcnt(49279))
    for i in range(2):
      k.emit(s_mov_b32(s[6 + i * 8], -16))
      k.emit(s_mov_b32(s[10 + i * 12], -16))
@@ -1218,6 +1213,31 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
    k.emit(s_waitcnt())
    k.emit(s_endpgm())
  elif (tile_m, tile_n) == (192, 256):
+    k.emit(s_and_b32(s[1], s[1], LIT, 65535))
+    k.emit(s_load_dwordx2(s[4:5], s[0:1], s[0], 0, 0, 0, 0, 1))
+    k.emit(s_mov_b32(s[8], 0))
+    k.emit(s_mov_b32(s[9], 0))
+    k.emit(s_load_dwordx2(s[12:13], s[0:1], s[0], 8, 0, 0, 0, 1))
+    k.emit(s_load_dwordx2(s[16:17], s[0:1], s[0], 16, 0, 0, 0, 1))
+    k.emit(s_mov_b32(s[36], N))
+    k.emit(s_mov_b32(s[37], K))
+    k.emit(s_mov_b32(s[38], K))
+    k.emit(s_mov_b32(s[43], M))
+    k.emit(s_mov_b32(s[44], N))
+    k.emit(s_mov_b32(s[45], K))
+    k.emit(s_load_dwordx2(s[20:21], s[0:1], s[0], 24, 0, 0, 0, 1))
+    k.emit(s_load_dwordx2(s[24:25], s[0:1], s[0], 32, 0, 0, 0, 1))
+    k.emit(s_mov_b32(s[39], scale_k))
+    k.emit(s_mov_b32(s[40], scale_k))
+    k.emit(v_lshrrev_b32_e32(v[1], 10))
+    k.emit(v_lshrrev_b32_e32(v[2], 10, v[1]))
+    k.emit(v_and_b32_e32(v[2], LIT, v[2], 1023))
+    k.emit(v_and_b32_e32(v[1], LIT, v[1], 1023))
+    k.emit(v_and_b32_e32(v[0], LIT, v[0], 1023))
+    k.emit(v_lshrrev_b32_e32(v[3], 6))
+    k.emit(v_and_b32_e32(v[0], 63))
+    k.emit(v_readfirstlane_b32_e32(v[46], v[3]))
+    k.emit(s_waitcnt(49279))
    k.emit(s_mul_i32(s[63], LIT, 8, 192))
    k.emit(v_cvt_f32_u32_e32(v[4], s[63]))
    k.emit(s_sub_i32(s[62], 0, s[63]))
@@ -2214,22 +2234,49 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
    k.emit(s_waitcnt())
    k.emit(s_endpgm())
  elif (tile_m, tile_n) == (256, 256):
+    k.emit(s_and_b32(s[1], s[1], LIT, 65535))
+    k.emit(s_load_dwordx2(s[4:5], s[0:1], s[0], 0, 0, 0, 0, 1))
+    k.emit(s_mov_b32(s[8], 0))
+    k.emit(s_mov_b32(s[9], 0))
+    k.emit(s_load_dwordx2(s[12:13], s[0:1], s[0], 8, 0, 0, 0, 1))
+    k.emit(s_load_dwordx2(s[16:17], s[0:1], s[0], 16, 0, 0, 0, 1))
+    k.emit(s_mov_b32(s[40], N))
+    k.emit(s_mov_b32(s[41], K))
+    k.emit(s_mov_b32(s[42], K))
+    k.emit(s_mov_b32(s[43], M))
+    k.emit(s_mov_b32(s[44], N))
+    k.emit(s_mov_b32(s[45], K))
+    k.emit(s_load_dwordx2(s[20:21], s[0:1], s[0], 24, 0, 0, 0, 1))
+    k.emit(s_load_dwordx2(s[24:25], s[0:1], s[0], 32, 0, 0, 0, 1))
+    k.emit(s_mov_b32(s[36], scale_k))
+    k.emit(s_mov_b32(s[37], scale_k))
+    k.emit(v_lshrrev_b32_e32(v[1], 10))
+    k.emit(v_lshrrev_b32_e32(v[2], 10, v[1]))
+    k.emit(v_and_b32_e32(v[2], LIT, v[2], 1023))
+    k.emit(v_and_b32_e32(v[1], LIT, v[1], 1023))
+    k.emit(v_and_b32_e32(v[0], LIT, v[0], 1023))
+    k.emit(v_lshrrev_b32_e32(v[3], 6))
+    k.emit(v_and_b32_e32(v[0], 63))
+    k.emit(s_mov_b32(s[46], s[2]))
+    k.emit(s_mov_b32(s[47], s[3]))
+    k.emit(v_readfirstlane_b32_e32(v[49], v[3]))
+    k.emit(s_waitcnt(49279))
    k.emit(s_add_u32(s[55], s[44], LIT, 255))
    k.emit(s_lshr_b32(s[54], s[55], 8))
    k.emit(s_mul_i32(s[48], s[54], s[47]))
-    k.emit(s_add_i32(s[48], s[48], s[49]))
+    k.emit(s_add_i32(s[48], s[48], s[46]))
    k.emit(s_add_u32(s[55], s[43], LIT, 255))
    k.emit(s_lshr_b32(s[52], s[55], 8))
    k.emit(s_lshl_b32(s[52], s[52], 5))
-    k.emit(s_mov_b32(s[49], 0))
+    k.emit(s_mov_b32(s[46], 0))
    k.label('L2_00E8')
    k.emit(s_cmp_lt_i32(s[48], s[52]))
    k.emit(s_cbranch_scc1(3), target='L2_00FC')
    k.emit(s_sub_i32(s[48], s[48], s[52]))
-    k.emit(s_add_i32(s[49], s[49], 32))
+    k.emit(s_add_i32(s[46], s[46], 32))
    k.emit(s_branch(65531), target='L2_00E8')
    k.label('L2_00FC')
-    k.emit(s_sub_i32(s[54], s[54], s[49]))
+    k.emit(s_sub_i32(s[54], s[54], s[46]))
    k.emit(s_cmp_lt_i32(s[54], 32))
    k.emit(s_cbranch_scc1(3), target='L2_0114')
    k.emit(s_lshr_b32(s[47], s[48], 5))
@@ -2264,7 +2311,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
    k.emit(s_mul_i32(s[52], s[54], s[47]))
    k.emit(s_sub_i32(s[52], s[48], s[52]))
    k.label('L2_0194')
-    k.emit(s_add_i32(s[49], s[52], s[49]))
+    k.emit(s_add_i32(s[46], s[52], s[46]))
    k.emit(s_mov_b32(s[6], -16))
    k.emit(s_mov_b32(s[10], -16))
    k.emit(s_mov_b32(s[18], -16))
@@ -2281,18 +2328,18 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
    k.emit(s_or_b32(s[9], s[9], LIT, 262144))
    k.emit(s_or_b32(s[17], s[17], LIT, 262144))
    k.emit(s_or_b32(s[13], s[13], LIT, 262144))
-    k.emit(s_lshr_b32(s[37], s[37], 1))
-    k.emit(s_mul_i32(s[52], s[37], s[43]))
+    k.emit(s_lshr_b32(s[41], s[41], 1))
+    k.emit(s_mul_i32(s[52], s[41], s[43]))
    k.emit(s_mov_b32(s[14], s[52]))
-    k.emit(s_lshr_b32(s[38], s[38], 1))
-    k.emit(s_mul_i32(s[52], s[38], s[44]))
+    k.emit(s_lshr_b32(s[42], s[42], 1))
+    k.emit(s_mul_i32(s[52], s[42], s[44]))
    k.emit(s_mov_b32(s[18], s[52]))
    k.emit(s_add_u32(s[52], s[43], 31))
    k.emit(s_lshr_b32(s[52], s[52], 5))
    k.emit(s_lshl_b32(s[52], s[52], 5))
-    k.emit(s_mul_i32(s[53], s[52], s[39]))
+    k.emit(s_mul_i32(s[53], s[52], s[36]))
    k.emit(s_mov_b32(s[22], s[53]))
-    k.emit(s_mul_i32(s[53], s[44], s[40]))
+    k.emit(s_mul_i32(s[53], s[44], s[37]))
    k.emit(s_mov_b32(s[26], s[53]))
    k.emit(s_mov_b32(s[23], LIT, 131072))
    k.emit(s_mov_b32(s[27], LIT, 131072))
@@ -2309,23 +2356,23 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
    k.emit(v_add_u32_e32(v[5], v[5], v[6]))
    k.emit(v_and_b32_e32(v[4], 1, v[4]))
    k.emit(v_add_u32_e32(v[5], v[5], v[4]))
-    k.emit(v_mul_lo_u32(v[212], s[37], v[5]))
+    k.emit(v_mul_lo_u32(v[212], s[41], v[5]))
    k.emit(v_and_b32_e32(v[4], 7))
    k.emit(v_lshlrev_b32_e32(v[4], 4, v[4]))
    k.emit(v_add_u32_e32(v[212], v[212], v[4]))
-    k.emit(s_lshr_b32(s[52], s[46], 1))
+    k.emit(s_lshr_b32(s[52], s[49], 1))
    k.emit(s_mul_i32(s[52], s[52], 8))
-    k.emit(s_and_b32(s[53], s[46], 1))
+    k.emit(s_and_b32(s[53], s[49], 1))
    k.emit(s_mul_i32(s[53], s[53], 2))
    k.emit(s_add_u32(s[52], s[52], s[53]))
    k.emit(s_mul_i32(s[53], s[47], LIT, 256))
    k.emit(s_add_u32(s[52], s[52], s[53]))
-    k.emit(s_mul_i32(s[52], s[37], s[52]))
+    k.emit(s_mul_i32(s[52], s[41], s[52]))
    k.emit(v_add_u32_e32(v[212], s[52], v[212]))
-    k.emit(s_mul_i32(s[52], s[37], 32))
+    k.emit(s_mul_i32(s[52], s[41], 32))
    for i in range(7):
      k.emit(v_add_u32_e32(v[213 + i * 1], s[52], v[212 + i * 1]))
-    k.emit(s_mul_i32(s[59], LIT, s[46], 1056))
+    k.emit(s_mul_i32(s[59], LIT, s[49], 1056))
    k.emit(s_add_u32(s[59], LIT, s[59], 4096))
    k.emit(v_and_b32_e32(v[4], 15))
    k.emit(v_lshrrev_b32_e32(v[5], 3, v[4]))
@@ -2349,35 +2396,35 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
    k.emit(v_add_u32_e32(v[221], LIT, v[220], 33792))
    k.emit(v_lshlrev_b32_e32(v[222], 2))
    k.emit(s_mul_i32(s[52], s[47], LIT, 256))
-    k.emit(s_mul_i32(s[53], s[46], 32))
+    k.emit(s_mul_i32(s[53], s[49], 32))
    k.emit(s_add_i32(s[52], s[53], s[52]))
-    k.emit(s_mul_i32(s[53], s[52], s[39]))
+    k.emit(s_mul_i32(s[53], s[52], s[36]))
    k.emit(v_add_u32_e32(v[222], s[53], v[222]))
-    k.emit(s_mul_i32(s[53], LIT, s[39], 128))
+    k.emit(s_mul_i32(s[53], LIT, s[36], 128))
    k.emit(v_add_u32_e32(v[223], s[53], v[222]))
-    k.emit(s_mul_i32(s[60], s[46], LIT, 256))
+    k.emit(s_mul_i32(s[60], s[49], LIT, 256))
    k.emit(s_add_i32(s[60], s[60], 0))
    k.emit(v_lshlrev_b32_e32(v[224], 2))
    k.emit(v_add_u32_e32(v[224], 0, v[224]))
    k.emit(v_lshlrev_b32_e32(v[225], 4))
-    k.emit(s_mul_i32(s[52], s[49], LIT, 256))
-    k.emit(s_mul_i32(s[53], s[46], 64))
+    k.emit(s_mul_i32(s[52], s[46], LIT, 256))
+    k.emit(s_mul_i32(s[53], s[49], 64))
    k.emit(s_add_u32(s[52], s[52], s[53]))
-    k.emit(s_mul_i32(s[52], s[52], s[38]))
+    k.emit(s_mul_i32(s[52], s[52], s[42]))
    k.emit(v_add_u32_e32(v[225], s[52], v[225]))
-    k.emit(s_mul_i32(s[52], 16, s[38]))
+    k.emit(s_mul_i32(s[52], 16, s[42]))
    k.emit(v_add_u32_e32(v[226], s[52], v[225]))
    k.emit(v_add_u32_e32(v[227], s[52], v[226]))
    k.emit(v_add_u32_e32(v[228], s[52], v[227]))
    for i in range(4):
      k.emit(v_add_u32_e32(v[229 + i * 1], LIT, v[225 + i * 1], 1024))
    k.emit(v_lshlrev_b32_e32(v[233], 2))
-    k.emit(s_mul_i32(s[52], s[49], LIT, 256))
-    k.emit(s_mul_i32(s[53], s[46], 64))
+    k.emit(s_mul_i32(s[52], s[46], LIT, 256))
+    k.emit(s_mul_i32(s[53], s[49], 64))
    k.emit(s_add_i32(s[52], s[53], s[52]))
-    k.emit(s_mul_i32(s[53], s[52], s[40]))
+    k.emit(s_mul_i32(s[53], s[52], s[37]))
    k.emit(v_add_u32_e32(v[233], s[53], v[233]))
-    k.emit(s_mul_i32(s[52], 32, s[40]))
+    k.emit(s_mul_i32(s[52], 32, s[37]))
    k.emit(v_add_u32_e32(v[234], s[52], v[233]))
    k.emit(s_mov_b32(s[61], LIT, 128))
    k.emit(s_mov_b32(s[62], LIT, 2048))
@@ -2463,18 +2510,18 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
    k.emit(ds_read_b32(v[201], v[224], v[0], v[0], 0, 0, 1))
    k.emit(ds_read_b32(v[202], v[224], v[0], v[0], 0, 0, 2))
    k.emit(ds_read_b32(v[203], v[224], v[0], v[0], 0, 0, 3))
-    k.emit(s_lshl_b32(s[36], s[36], 1))
+    k.emit(s_lshl_b32(s[40], s[40], 1))
    k.emit(s_mul_i32(s[52], s[47], LIT, 256))
-    k.emit(s_mul_hi_u32(s[53], s[52], s[36]))
+    k.emit(s_mul_hi_u32(s[53], s[52], s[40]))
    k.emit(s_add_u32(s[5], s[5], s[53]))
-    k.emit(s_mul_i32(s[53], s[52], s[36]))
+    k.emit(s_mul_i32(s[53], s[52], s[40]))
    k.emit(s_add_u32(s[4], s[4], s[53]))
    k.emit(s_addc_u32(s[5], 0, s[5]))
    k.emit(s_sub_i32(s[52], s[43], s[52]))
-    k.emit(s_mul_i32(s[52], s[52], s[36]))
+    k.emit(s_mul_i32(s[52], s[52], s[40]))
    k.emit(s_mov_b32(s[6], s[52]))
    k.emit(v_and_b32_e64(v[235], v[0], 15))
-    k.emit(v_mul_lo_u32(v[235], v[235], s[36]))
+    k.emit(v_mul_lo_u32(v[235], v[235], s[40]))
    k.emit(v_lshrrev_b32_e32(v[4], 5))
    k.emit(v_mul_i32_i24_e32(v[4], 16, v[4]))
    k.emit(v_add_u32_e32(v[235], v[4], v[235]))
@@ -2482,12 +2529,12 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
    k.emit(v_and_b32_e32(v[4], 1, v[4]))
    k.emit(v_mul_i32_i24_e32(v[4], 32, v[4]))
    k.emit(v_add_u32_e32(v[235], v[4], v[235]))
-    k.emit(s_mul_i32(s[52], s[49], LIT, 256))
-    k.emit(s_mul_i32(s[53], s[46], 64))
+    k.emit(s_mul_i32(s[52], s[46], LIT, 256))
+    k.emit(s_mul_i32(s[53], s[49], 64))
    k.emit(s_add_i32(s[52], s[52], s[53]))
    k.emit(s_lshl_b32(s[52], s[52], 1))
    k.emit(v_add_u32_e32(v[235], s[52], v[235]))
-    k.emit(s_mul_i32(s[53], s[36], 16))
+    k.emit(s_mul_i32(s[53], s[40], 16))
    for i in range(15):
      k.emit(v_add_u32_e64(v[236 + i * 1], v[235 + i * 1], s[53]))
    k.emit(s_mov_b32(s[50], 0))
@@ -2496,7 +2543,7 @@ def build_kernel(M: int, N: int, K: int, tile_m: int, tile_n: int):
      k.emit(s_cmp_lt_u32(LIT, s[51], 512 + i * -256))
      k.emit(s_cselect_b32(s[61 + i * 1], s[61 + i * 1], 0))
      k.emit(s_cselect_b32(s[63 + i * 1], s[63 + i * 1], 0))
-    k.emit(s_cmp_lt_i32(s[46], 2))
+    k.emit(s_cmp_lt_i32(s[49], 2))
    k.emit(s_cbranch_scc0(1367), target='L2_25B8')
    k.label('L2_105C')
    k.emit(s_waitcnt(122))
@@ -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()
@@ -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):
@@ -140,10 +140,10 @@ def pm4_program(ctx, call, prg):
 pm_pm4_opsel = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), pm4_program),

-  (UPat(Ops.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, ...]:
@@ -207,10 +207,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)),
+  (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):
@@ -254,7 +254,7 @@ def amd_usb_submit(ctx, lin):

  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)),))
+    lin = lin.replace(src=lin.src + (UOp(Ops.INS, arg="poke", 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 ()
@@ -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:
@@ -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()
@@ -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={}):
  """
@@ -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]
@@ -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):
@@ -152,7 +152,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 +183,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 +218,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),
@@ -299,7 +299,7 @@ class TestConcatWidthParsing(unittest.TestCase):
      self.assertIs(parsed.simplify(), UOp.const(expected, dtypes.uint32))

  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),
@@ -330,7 +330,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(),
@@ -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):
@@ -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"))

@@ -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,12 +41,12 @@ 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 + UOp.cconst(1, dtypes.int32)).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)
@@ -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
@@ -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)
@@ -187,13 +187,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()
@@ -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')
@@ -359,13 +359,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]])
@@ -1820,7 +1813,6 @@ class TestOps(unittest.TestCase):
    helper_test_op([(45,65)], lambda x: x.asinh(), grad_atol=1e-6, 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)
@@ -2817,7 +2809,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 +2962,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()
@@ -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)))
@@ -78,7 +78,7 @@ 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)
+    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())
@@ -93,9 +93,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 +110,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,))
@@ -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):
@@ -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):
@@ -62,7 +62,7 @@ def _test_uops_result(output_dtype, uops, res):
 class TestBitcastBufferView(unittest.TestCase):
  @Context(SPEC=2)
  def test_render(self):
-    buf = UOp.param(0, dtypes.uint32, 4)
+    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)
@@ -71,7 +71,7 @@ class TestBitcastBufferView(unittest.TestCase):
  @Context(SPEC=2)
  def test_load(self):
    val = 0x1122334455667788
-    src, out = UOp.param(0, dtypes.uint32, 4), UOp.param(1, dtypes.uint64, 1)
+    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])
@@ -80,7 +80,7 @@ class TestBitcastBufferView(unittest.TestCase):
  @Context(SPEC=2)
  def test_store(self):
    val = 0x1122334455667788
-    dst = UOp.param(0, dtypes.uint32, 6)
+    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])
@@ -249,8 +249,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 +278,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 +287,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)
@@ -43,7 +43,7 @@ def nested_loop_kernel(C:UOp) -> UOp:
  return C[0].store(i[0].load()).sink(arg=KernelInfo(name="nested_loop", opts_to_apply=()))

 def wait_ext_kernel() -> UOp:
-  sig = UOp.param(0, dtypes.int, 1, volatile=True)
+  sig = UOp.param(0, dtypes.int, (1,), volatile=True)
  l = UOp.loop(0)
  v = sig.after(l)[0].load()
  e = v.end(l, v < 1)
@@ -3,10 +3,10 @@ from unittest.mock import patch
 from tinygrad import Device, Tensor
 from tinygrad.device import Buffer
 from tinygrad.dtype import dtypes
-from tinygrad.helpers import HCQ2
+from tinygrad.helpers import getenv
 from tinygrad.runtime.support.hcq2 import HCQ_DEVS, all_devices_in

-@unittest.skipUnless(HCQ2 and all_devices_in(Device.DEFAULT, HCQ_DEVS), "hcq2 device required")
+@unittest.skipUnless(getenv("HCQ2") and all_devices_in(Device.DEFAULT, HCQ_DEVS), "hcq2 device required")
 class TestHCQ2(unittest.TestCase):
  def test_copy_without_copy_queue(self):
    with patch.object(Device[Device.DEFAULT], "has_copy_queue", False):
@@ -28,7 +28,7 @@ def vision_conv_143():
  c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float))
  c49 = UOp.param(2, dtypes.half, shape=(64, 49, 4))
  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
-  c63 = UOp.param(3, dtypes.float, 128)
+  c63 = UOp.param(3, dtypes.float, (128,))
  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
  c67 = c0.index((c2*128+c5+c8*4096)).store(c65).end(c8, c2, c5)

@@ -54,7 +54,7 @@ def vision_conv_153():
  c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float))
  c49 = UOp.param(2, dtypes.half, shape=(128, 49, 4))
  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
-  c63 = UOp.param(3, dtypes.float, 256)
+  c63 = UOp.param(3, dtypes.float, (256,))
  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
  c67 = c0.index((c2*256+c5+c8*4096)).store(c65).end(c8, c2, c5)

@@ -73,7 +73,7 @@ def dm_conv_172():
  c18 = UOp.range(8, 2, AxisType.REDUCE)
  c23 = UOp.param(2, dtypes.half, shape=(240, 128, 4))
  c35 = c5.index((c7*4+c10+c13*128+c18*1536))*c23.index((c10*4+c2%4+c7*16+c2//4*512))
-  c37 = UOp.param(3, dtypes.float, 960)
+  c37 = UOp.param(3, dtypes.float, (960,))
  c39 = c35.reduce(c7, c10, arg=Ops.ADD)+c37.index(c2)
  c50 = (1.0+((c39+0.044708251953125*(c39*(c39*c39)))*-2.3021129851685216).exp2()).reciprocal()*c39
  c53 = c50.reduce(c18, c13, arg=Ops.ADD)*0.010416666666666666
@@ -7,8 +7,8 @@ BENCHMARK_OPS = {Ops.INDEX, Ops.STAGE}

@functools.cache
 def create_uop(a:int) -> UOp:
-  op, src, arg, *rest = trace.uop_fields[a]
-  return UOp(op, tuple(create_uop(s) for s in src), arg, *rest)
+  op, dtype, src, arg, *rest = trace.uop_fields[a]
+  return UOp(op, dtype, tuple(create_uop(s) for s in src), arg, *rest)

 if __name__ == "__main__":
  # load rewrite trace
@@ -36,8 +36,8 @@ class TestGPUCrash(unittest.TestCase):

  def _run_insts(self, insts: list[Inst]):
    buf = UOp.new_buffer("AMD", 64, dtypes.uint8)
-    sink = UOp.sink(UOp.param(0, dtypes.uint8, 64, device="AMD"), UOp.special(1, "lidx0"), arg=KernelInfo("test"))
-    prg = UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=(i, dtypes.void)) for i in insts))))
+    sink = UOp.sink(UOp.param(0, dtypes.uint8, (64,), device="AMD"), UOp.special(1, "lidx0"), arg=KernelInfo("test"))
+    prg = UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=i) for i in insts))))
    run_linear(UOp(Ops.LINEAR, src=(prg.call(buf),)), wait=True)

  def _assert_gpu_fault(self, func):
@@ -54,12 +54,6 @@ class TestMainOnnxOps(TestOnnxOps):
    outputs = ["squeezed"]
    self.helper_test_single_op("Squeeze", inputs, attributes, outputs)

-  def test_mean_variance_normalization_axes(self):
-    inputs = {"x": np.random.randn(2, 3, 4, 5).astype(np.float32)}
-    attributes = {"axes": [2, 3]}
-    outputs = ["out"]
-    self.helper_test_single_op("MeanVarianceNormalization", inputs, attributes, outputs)
-
  def test_conv(self):
    # test VALID auto_pad
    inputs = {
@@ -241,15 +235,6 @@ class TestMainOnnxOps(TestOnnxOps):
    outputs = ["y"]
    self.helper_test_single_op("MaxUnpool", inputs, attributes, outputs)

-  def test_maxunpool_pads(self):
-    # per-axis pads shrink the output: spatial dim is (i-1)*stride + kernel - pad_begin - pad_end -> (2, 4), and indices index into that output
-    # NOTE: indices must be in bounds of that output; ORT aborts the process on out-of-bounds indices
-    xT = np.array([[[[5, 6], [7, 8]]]], dtype=np.float32)
-    xI = np.array([[[[0, 3], [4, 7]]]], dtype=np.int64)
-    inputs = {"x": xT, "indices": xI}
-    attributes = {"kernel_shape": [2, 2], "strides": [2, 2], "pads": [1, 0, 1, 0]}
-    self.helper_test_single_op("MaxUnpool", inputs, attributes, ["y"])
-
  def test_averagepool_3d_dilations_large_count_include_pad_is_1_ceil_mode_is_True(self):
    # https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-13
    inputs = {"x": np.random.randn(1, 1, 32, 32, 32).astype(np.float32)}
@@ -1,5 +1,5 @@
 import unittest, time, itertools
-from tinygrad import Tensor, Context, dtypes
+from tinygrad import Tensor, Context

 class TestScheduleScaling(unittest.TestCase):
  """Test that .schedule() scales linearly with graph size (no O(n^2) behavior)."""
@@ -138,7 +138,7 @@ class TestScheduleScaling(unittest.TestCase):
    def custom_kernel_assign(n):
      def custom_asm(out):
        return UOp(Ops.PROGRAM, src=(UOp.sink(out, arg=KernelInfo(f"fxn_{next(count)}")),
-                                     UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=(s_nop(i), dtypes.void)) for i in range(n*8)))))
+                                     UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=s_nop(i)) for i in range(n*8)))))
      call = Tensor.custom_kernel(Tensor.empty(1), fxn=custom_asm)[0]
      return Tensor.cat(*[Tensor.empty(1).assign(call+i) for i in range(n)])
    self._assert_linear(custom_kernel_assign, n_small=50, n_large=500)
@@ -122,7 +122,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple
  for buf_dt, data in inputs or []:
    bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
    allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
-  g = UOp.param(0, uop.dtype, 1)
+  g = UOp.param(0, uop.dtype, (1,))
  prg = to_program(UOp.store(g.index(UOp.const(0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
  prog = dev.runtime(prg.to_elf())
  prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs, vals=vals)
@@ -1,15 +1,20 @@
-import ctypes, mmap, collections, functools, copy, os
+import pathlib, re, ctypes, mmap, collections, functools, copy, os
 from tinygrad.runtime.autogen import kfd, amdgpu_drm, libc
 import tinygrad.runtime.autogen.am.am as am
 from tinygrad.helpers import from_mv
 from test.mockgpu.driver import VirtDriver, VirtFileDesc, TextFileDesc, DirFileDesc, VirtFile
 from test.mockgpu.amd.amdgpu import AMDGPU, gpu_props, GFX_TARGET_VERSION, MOCKGPU_ARCH

-def _ioctl_nr(ioctl: functools.partial) -> int: return ioctl.args[2]
-
-kfd_ioctl_info = {
-  _ioctl_nr(ioctl): (name, ioctl.args[3]) for name, ioctl in vars(kfd).items()
-  if name.startswith("AMDKFD_IOC_") and isinstance(ioctl, functools.partial)}
+def ioctls_from_header():
+  # hdrpy = (pathlib.Path(__file__).parent.parent.parent.parent / "tinygrad" / "runtime" / "autogen" / "kfd.py").read_text()
+  # pattern = r'# (AMDKFD_IOC_[A-Z0-9_]+)\s=\s_(IOW?R?).*\(( 0x[0-9a-fA-F]+) ,\s+struct\s([A-Za-z0-9_]+)\s+\)'
+  # matches = re.findall(pattern, hdrpy, re.MULTILINE)
+  hdr = (pathlib.Path(__file__).parent.parent.parent.parent / "extra" / "hip_gpu_driver" / "kfd_ioctl.h").read_text().replace("\\\n", "")
+  pattern = r'#define\s+(AMDKFD_IOC_[A-Z0-9_]+)\s+AMDKFD_(IOW?R?)\((0x[0-9a-fA-F]+),\s+struct\s([A-Za-z0-9_]+)\)'
+  matches = re.findall(pattern, hdr, re.MULTILINE)
+  return type("KFD_IOCTLS", (object, ), {name: int(nr, 0x10) for name, _, nr, _ in matches}), \
+         {int(nr, 0x10): getattr(kfd, "struct_"+sname, None) for name, idir, nr, sname in matches}
+kfd_ioctls, kfd_headers = ioctls_from_header()

 class KFDFileDesc(VirtFileDesc):
  def __init__(self, fd, driver):
@@ -111,44 +116,42 @@ class AMDDriver(VirtDriver):

  def kfd_ioctl(self, req, argp):
    nr = req & 0xFF
-    if nr not in kfd_ioctl_info: raise RuntimeError(f"unknown kfd ioctl, {nr} unknown")
-    name, struct_type = kfd_ioctl_info[nr]
-    struct = struct_type.from_address(argp)
+    struct = kfd_headers[nr].from_address(argp)

-    if nr == _ioctl_nr(kfd.AMDKFD_IOC_ACQUIRE_VM): pass
-    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_RUNTIME_ENABLE): pass
-    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_GET_VERSION):
+    if nr == kfd_ioctls.AMDKFD_IOC_ACQUIRE_VM: pass
+    elif nr == kfd_ioctls.AMDKFD_IOC_RUNTIME_ENABLE: pass
+    elif nr == kfd_ioctls.AMDKFD_IOC_GET_VERSION:
      struct.major_version = 1
      struct.minor_version = 14
-    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_ALLOC_MEMORY_OF_GPU):
+    elif nr == kfd_ioctls.AMDKFD_IOC_ALLOC_MEMORY_OF_GPU:
      if struct.gpu_id not in self.gpus: return -1
      struct.handle = self._alloc_handle()
      self.object_by_handle[struct.handle] = copy.deepcopy(struct) # save memory struct to know what mem it is
      # Track signal memory (uncached + coherent) - progress queues when written to
      if struct.flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED:
        self.track_address(struct.va_addr, struct.va_addr + struct.size, lambda mv,off: None, lambda mv, off: self._emulate_execute())
-    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU):
+    elif nr == kfd_ioctls.AMDKFD_IOC_FREE_MEMORY_OF_GPU:
      self.object_by_handle.pop(struct.handle)
-    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU):
+    elif nr == kfd_ioctls.AMDKFD_IOC_MAP_MEMORY_TO_GPU:
      dev_ids = (ctypes.c_int32 * struct.n_devices).from_address(struct.device_ids_array_ptr)
      for i in range(struct.n_devices):
        gpu = self.gpus[dev_ids[i]]
        mem_obj = self.object_by_handle[struct.handle]
        gpu.map_range(mem_obj.va_addr, mem_obj.size)
        struct.n_success = i + 1
-    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU):
+    elif nr == kfd_ioctls.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU:
      dev_ids = (ctypes.c_int32 * struct.n_devices).from_address(struct.device_ids_array_ptr)
      for i in range(struct.n_devices):
        gpu = self.gpus[dev_ids[i]]
        mem_obj = self.object_by_handle[struct.handle]
        gpu.unmap_range(mem_obj.va_addr, mem_obj.size)
        struct.n_success = i + 1
-    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_CREATE_EVENT):
+    elif nr == kfd_ioctls.AMDKFD_IOC_CREATE_EVENT:
      struct.event_slot_index = self._alloc_next_event_slot()
      struct.event_id = struct.event_slot_index

      if struct.event_type == kfd.KFD_IOC_EVENT_MEMORY: self.mmu_event_ids.append(struct.event_id)
-    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_CREATE_QUEUE):
+    elif nr == kfd_ioctls.AMDKFD_IOC_CREATE_QUEUE:
      gpu = self.gpus[struct.gpu_id]
      if struct.queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA:
        gpu.add_sdma_queue(struct.ring_base_address, struct.ring_size, struct.read_pointer_address, struct.write_pointer_address)
@@ -159,7 +162,7 @@ class AMDDriver(VirtDriver):
      # Track writes to doorbell, calling callback
      struct.doorbell_offset = self._alloc_doorbell(struct.gpu_id)
      self.track_address(struct.doorbell_offset, struct.doorbell_offset + 8, lambda mv,off: None, lambda mv, off: self._emulate_execute())
-    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_WAIT_EVENTS):
+    elif nr == kfd_ioctls.AMDKFD_IOC_WAIT_EVENTS:
      evs = (kfd.struct_kfd_event_data * struct.num_events).from_address(struct.events_ptr)
      for ev in evs:
        if ev.event_id in self.mmu_event_ids and "MOCKGPU_EMU_FAULTADDR" in os.environ:
@@ -167,7 +170,11 @@ class AMDDriver(VirtDriver):
          ev.memory_exception_data.va = int(os.environ["MOCKGPU_EMU_FAULTADDR"], 16)
          ev.memory_exception_data.failure.NotPresent = 1
    else:
-      raise RuntimeError(f"unsupported kfd ioctl, {nr} {name}")
+      name = "unknown"
+      for k,v in kfd_ioctls.__dict__.items():
+        if nr == v: name = k
+      assert False, f"unknown kfd ioctl, {nr} {name}"
+      exit(1)
    return 0

  def _emulate_execute(self):
@@ -321,10 +321,10 @@ def _int_clamp(op_name: str, srcs: dict) -> UOp | None:
 class _Ctx:
  """Context for instruction compilation - holds buffers and helpers."""
  __slots__ = ('inst_size', 'dyn_fields', '_axis_id', 'wave_size', 'vgpr', 'accvgpr')
-  sgpr = UOp.param(0, dtypes.uint32, SGPR_COUNT)
-  vmem = UOp.param(2, dtypes.uint32, 1 << 46)
-  lds = UOp.param(3, dtypes.uint32, 16384)
-  scratch = UOp.param(4, dtypes.uint8, 1 << 30)
+  sgpr = UOp.param(0, dtypes.uint32, (SGPR_COUNT,))
+  vmem = UOp.param(2, dtypes.uint32, (1 << 46,))
+  lds = UOp.param(3, dtypes.uint32, (16384,))
+  scratch = UOp.param(4, dtypes.uint8, (1 << 30,))
  # Cache PARAM UOps by wave_size so all _Ctx instances with same wave_size share identical UOp references
  _vgpr_cache: dict[int, UOp] = {}
  _accvgpr_cache: dict[int, UOp] = {}
@@ -332,10 +332,10 @@ class _Ctx:
  def __init__(self, inst_size: int, wave_size: int = 32):
    self.inst_size, self._axis_id, self.wave_size = inst_size, 0, wave_size
    self.dyn_fields: list[tuple[int, int]] = []  # (lo, hi) of fields read dynamically
-    if wave_size not in _Ctx._vgpr_cache: _Ctx._vgpr_cache[wave_size] = UOp.param(1, dtypes.uint32, 256 * wave_size)
+    if wave_size not in _Ctx._vgpr_cache: _Ctx._vgpr_cache[wave_size] = UOp.param(1, dtypes.uint32, (256 * wave_size,))
    self.vgpr = _Ctx._vgpr_cache[wave_size]
    if wave_size == 64:
-      if wave_size not in _Ctx._accvgpr_cache: _Ctx._accvgpr_cache[wave_size] = UOp.param(5, dtypes.uint32, 256 * wave_size)
+      if wave_size not in _Ctx._accvgpr_cache: _Ctx._accvgpr_cache[wave_size] = UOp.param(5, dtypes.uint32, (256 * wave_size,))
      self.accvgpr = _Ctx._accvgpr_cache[wave_size]
    else:
      self.accvgpr = self.vgpr
@@ -107,7 +107,7 @@ class TestGroupedDims(unittest.TestCase):

  def test_global_prod_max(self):
    g, l = UOp.range(256, 0, AxisType.GLOBAL), UOp.range(256, 1, AxisType.LOCAL)
-    sink = UOp.param(0, dtypes.float, 512).index(g + l).store(UOp.const(1.0)).end(g, l).sink(arg=KernelInfo())
+    sink = UOp.param(0, dtypes.float, (512,)).index(g + l).store(UOp.const(1.0)).end(g, l).sink(arg=KernelInfo())
    class R(Renderer): global_max, local_max, global_prod_max = (256, 256, 256), (128, 128, 128), (128, 128, 128)
    specials = [u for u in add_gpudims(R(Target()), sink).toposort() if u.op is Ops.SPECIAL]
    self.assertGreater(len([s for s in specials if "lidx" in s.arg]), 1)
@@ -7,14 +7,14 @@ from tinygrad.codegen import to_program

 class TestLinearizerFailures(unittest.TestCase):
  def test_fail_1(self):
-    c0 = UOp.param(0, dtypes.float, 64)
+    c0 = UOp.param(0, dtypes.float, (64,))
    c1 = UOp.range(UOp.const(2), 1, AxisType.WEAK)
    c2 = UOp.range(UOp.const(32), 2, AxisType.WEAK)
    c3 = ((c1*UOp.const(32))+c2)
-    c4 = UOp.param(1, dtypes.float, 163840)
+    c4 = UOp.param(1, dtypes.float, (163840,))
    c5 = UOp.range(UOp.const(2560), 0, AxisType.REDUCE)
    c6 = c4.index(((((((c5//UOp.const(8))%UOp.const(8))*UOp.const(8))+(c5%UOp.const(8)))+(((c2*UOp.const(40))+(c5//UOp.const(64)))*UOp.const(64)))+(c1*UOp.const(81920))))
-    c7 = UOp.param(2, dtypes.float, 64)
+    c7 = UOp.param(2, dtypes.float, (64,))
    c8 = c7.index(c3)
    c9 = ((((c6+(c8*UOp.const(-1.0)))*(c6+(c8*UOp.const(-1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(0.000390625))+UOp.const(1e-05)).sqrt().reciprocal()
    c10 = c0.index(c3).store(c9).end(c1, c2)
@@ -40,7 +40,7 @@ class TestPatternMatcher(unittest.TestCase):
      assert len(x.src) == 0
      return UOp(Ops.CONST, src=(UOp(Ops.CONST),))
    matcher = PatternMatcher([(UPat(Ops.CONST, src=(), name="x"), fxn)])
-    c1 = UOp(Ops.CONST, arg=1.0)
+    c1 = UOp(Ops.CONST, dtypes.float, arg=1.0)
    # second rewrite shouldn't match anything
    c1 = matcher.rewrite(c1)
    c1 = matcher.rewrite(c1)
@@ -15,7 +15,7 @@ def simplify_valid_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move
 def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load+indexing_simplify, name="simplify_image_idx")

 def get_gated_load_uop(valid:UOp, idx:UOp):
-  return UOp.param(0, dtypes.float, 1024).index(idx.valid(valid)).load()
+  return UOp.param(0, dtypes.float, (1024,)).index(idx.valid(valid)).load()

 def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]):
  return UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)).load()
@@ -505,7 +505,7 @@ class TestDropTrueGate(unittest.TestCase):
    from tinygrad.codegen.late.coalesce import indexing_simplify
    from tinygrad.uop.ops import graph_rewrite
    from tinygrad.uop.symbolic import sym
-    buf = UOp.param(0, dtypes.int, 1)
+    buf = UOp.param(0, dtypes.int, (1,))
    idx = UOp.const(0)
    true_gate = UOp.const(True)
    index_with_gate = buf.index(idx.valid(true_gate))
@@ -551,7 +551,7 @@ class TestRangeShrink(unittest.TestCase):
    # one load guards r < 4, but another load uses r without a gate -> no shrink
    r = Range(0, 204)
    load1 = get_gated_load_uop(r < UOp.const(4), r)
-    load2 = UOp.param(1, dtypes.float, 204).index(r).load()
+    load2 = UOp.param(1, dtypes.float, (204,)).index(r).load()
    ranges = self.get_ranges(UOp.sink(load1, load2))
    self.assert_range_end(ranges, 204)

@@ -575,7 +575,7 @@ class TestRangeShrink(unittest.TestCase):
    from tinygrad.dtype import Invalid
    r = Range(0, 204)
    x = (r < 4).where(UOp.const(1.0), Invalid)
-    ranges = self.get_ranges(UOp.param(0, dtypes.float, 204).index(r).store((r < 4).where(x, Invalid)).sink())
+    ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, Invalid)).sink())
    self.assert_range_end(ranges, 4)

  def test_range_shrink_store_where_invalid_flipped(self):
@@ -583,7 +583,7 @@ class TestRangeShrink(unittest.TestCase):
    from tinygrad.dtype import Invalid
    r = Range(0, 204)
    x = (r < 4).where(UOp.const(1.0), Invalid)
-    ranges = self.get_ranges(UOp.param(0, dtypes.float, 204).index(r).store((r >= 4).where(Invalid, x)).sink())
+    ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r >= 4).where(Invalid, x)).sink())
    self.assert_range_end(ranges, 4)

 if __name__ == '__main__':
@@ -9,7 +9,7 @@ class TestTranscendentalFunctions(unittest.TestCase):
  def test_payne_hanek_reduction(self):
    # TODO: Test constant input when constant folding is fixed (or maybe test both variants)
    # Load input value from a buffer to prevent constant folding
-    input_buf = UOp.param(1, dtypes.double, 1)
+    input_buf = UOp.param(1, dtypes.double, (1,))
    loaded_value = input_buf.index(UOp.const(0)).load()
    def eval_payne_hanek_reduction(v:float) -> tuple[float, int]:
      return tuple(eval_uop(u, [(dtypes.float64, [v])]) for u in payne_hanek_reduction(loaded_value))
@@ -214,8 +214,8 @@ class TestUOpGraph(unittest.TestCase):
      for i, const in enumerate(consts): self.assertIs(vec.index(i), const)

  def test_cast_alu_fold(self):
-    d0 = UOp.param(0, dtypes.bool, 1)
-    d1 = UOp.param(1, dtypes.int, 1)
+    d0 = UOp.param(0, dtypes.bool, (1,))
+    d1 = UOp.param(1, dtypes.int, (1,))
    idx = UOp.const(0)
    ld = d1.index(idx)
    alu = (ld<1).cast(dtypes.bool)
@@ -224,8 +224,8 @@ class TestUOpGraph(unittest.TestCase):
    self.assertEqual(len([x for x in uops if x.op is Ops.CAST and x.src[0].op is not Ops.CONST]), 0)

  def test_double_cast_fold(self):
-    d0 = UOp.param(0, dtypes.float, 1)
-    d1 = UOp.param(1, dtypes.int, 1)
+    d0 = UOp.param(0, dtypes.float, (1,))
+    d1 = UOp.param(1, dtypes.int, (1,))
    idx = UOp.const(0, dtypes.int)
    ld = d1.index(idx)
    alu = ld.cast(dtypes.float).cast(dtypes.float)
@@ -243,7 +243,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_bitcast_to_same_dtype_fold(self):
    for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
-      d0 = UOp.param(0, dt, 1)
+      d0 = UOp.param(0, dt, (1,))
      v = d0.index(UOp.const(0))
      uops = to_uops_list([v.bitcast(dt)])
      self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST and x.dtype is dt]), 0, f"dtype = {dt}")
@@ -255,10 +255,10 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_on_gated_load_fold(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.long, 100)
+    d0 = UOp.param(0, dtypes.long, (100,))
    ld = d0.index(ridx0.valid(ridx0<50))
    w = (ridx0<50).where(ld, 5)
-    out = UOp.param(1, dtypes.long, 100)
+    out = UOp.param(1, dtypes.long, (100,))
    uops = to_uops_list([out.index(ridx0).store(w)])
    expected = full_rewrite(UOp.const(5, dtypes.long).sink()).src[0]
    for u in uops:
@@ -267,7 +267,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_on_gated_load_folds_swapped_branches(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.long, 100)
+    d0 = UOp.param(0, dtypes.long, (100,))
    ld = d0.index(ridx0.valid((ridx0<50).logical_not()))
    w = (ridx0<50).where(5, ld)
    uops = to_uops_list([w])
@@ -278,11 +278,11 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_on_gated_load_with_cast(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.int, 100)
+    d0 = UOp.param(0, dtypes.int, (100,))
    gate_idx = ridx0.valid((ridx0<50))
    ld = d0.index(gate_idx).cast(dtypes.float)
    w = (ridx0<50).where(ld, 5.0)
-    out = UOp.param(1, dtypes.float, 100)
+    out = UOp.param(1, dtypes.float, (100,))
    uops = to_uops_list([out.index(ridx0).store(w)])
    expected = full_rewrite(UOp.const(5, dtypes.int).sink()).src[0]
    for u in uops:
@@ -291,27 +291,27 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_on_casted_gated_load_extra_cond(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.float, 100)
+    d0 = UOp.param(0, dtypes.float, (100,))
    ld = d0.index(ridx0.valid(ridx0<50))
    w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(0.0)).cast(dtypes.half)
-    out = UOp.param(1, dtypes.half, 100)
+    out = UOp.param(1, dtypes.half, (100,))
    uops = to_uops_list([out.index(ridx0).store(w)])
    for u in uops:
      assert u.op is not Ops.WHERE

  def test_where_on_casted_gated_load_extra_cond_swapped(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.float, 100)
+    d0 = UOp.param(0, dtypes.float, (100,))
    ld = d0.index(ridx0.valid(ridx0<50))
    w = ((ridx0<50) & (ridx0>30)).where(UOp.const(0.0), ld).cast(dtypes.half)
-    out = UOp.param(1, dtypes.half, 100)
+    out = UOp.param(1, dtypes.half, (100,))
    uops = to_uops_list([out.index(ridx0).store(w)])
    for u in uops:
      assert u.op is not Ops.WHERE

  def test_where_in_store_becomes_gate(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.long, 100)
+    d0 = UOp.param(0, dtypes.long, (100,))
    idx = d0.index(ridx0)
    ld = idx.load()
    val = (ridx0<50).where(5, ld)
@@ -325,14 +325,14 @@ class TestUOpGraph(unittest.TestCase):
  def test_load_idx_becomes_int(self):
    # mnist indexing with split reduceop
    # Make sure we are not doign math on the loaded index, which would promote it to long
-    c0 = UOp.param(0, dtypes.uchar, 128000)
+    c0 = UOp.param(0, dtypes.uchar, (128000,))
    c1 = UOp.range(UOp.const(512), 1, AxisType.WEAK)
    c2 = UOp.range(UOp.const(250), 2, AxisType.WEAK)
-    c3 = UOp.param(1, dtypes.int, 512)
+    c3 = UOp.param(1, dtypes.int, (512,))
    c4 = c3.index(c1)
    c5 = UOp.range(UOp.const(240), 0, AxisType.REDUCE)
    c6 = ((c2*UOp.const(240))+c5)
-    c7 = UOp.param(2, dtypes.uchar, 60000)
+    c7 = UOp.param(2, dtypes.uchar, (60000,))
    c8 = c7.index(c6)
    c9 = ((c4<0).where((c4+60000), c4)!=c6.cast(dtypes.int)).where(0, c8.cast(dtypes.uint).cast(dtypes.uchar)).reduce(c5, arg=Ops.ADD)
    c10 = c0.index(((c1*UOp.const(250))+c2)).store(c9).end(c1, c2)
@@ -342,14 +342,14 @@ class TestUOpGraph(unittest.TestCase):

  def test_load_idx_no_math_on_loaded(self):
    # test the (x+y)<c pattern where x has loads - we shouldn't do math on loaded indices
-    c0 = UOp.param(0, dtypes.uchar, 128000)
+    c0 = UOp.param(0, dtypes.uchar, (128000,))
    c1 = UOp.range(UOp.const(512), 1, AxisType.WEAK)
    c2 = UOp.range(UOp.const(250), 2, AxisType.WEAK)
-    c3 = UOp.param(1, dtypes.int, 512)
+    c3 = UOp.param(1, dtypes.int, (512,))
    c4 = c3.index(c1)  # c4 is a load
    c5 = UOp.range(UOp.const(240), 0, AxisType.REDUCE)
    c6 = ((c2*UOp.const(240))+c5)
-    c7 = UOp.param(2, dtypes.uchar, 60000)
+    c7 = UOp.param(2, dtypes.uchar, (60000,))
    c8 = c7.index(c6)
    # (loaded + range) < const pattern - loaded value shouldn't be promoted to long
    loaded_idx = c4.cast(dtypes.weakint)
@@ -361,9 +361,9 @@ class TestUOpGraph(unittest.TestCase):
      self.assertNotEqual(u.dtype, dtypes.long)

  def test_fold_gated_load(self):
-    glbl0 = UOp.param(0, dtypes.int, 1)
-    glbl1 = UOp.param(1, dtypes.int, 1)
-    glbl2 = UOp.param(2, dtypes.int, 1)
+    glbl0 = UOp.param(0, dtypes.int, (1,))
+    glbl1 = UOp.param(1, dtypes.int, (1,))
+    glbl2 = UOp.param(2, dtypes.int, (1,))
    idx = UOp.const(0)
    ld0 = glbl1.index(UOp.invalid())
    ld1 = glbl2.index(idx.valid(UOp.const(True)))
@@ -372,7 +372,7 @@ class TestUOpGraph(unittest.TestCase):
    self.assertEqual(len([u for u in uops if u.op is Ops.LOAD]), 1)

  def test_fold_gated_load_local(self):
-    glbl0 = UOp.param(0, dtypes.int, 16)
+    glbl0 = UOp.param(0, dtypes.int, (16,))
    smem = UOp.placeholder((18,), dtypes.int, slot=0, addrspace=AddrSpace.LOCAL)
    lidx = UOp.special(16, "lidx0")
    st = smem.index(lidx).store(glbl0.index(lidx).load())
@@ -385,7 +385,7 @@ class TestUOpGraph(unittest.TestCase):
    self.assertEqual(len([u for u in uops if u.op is Ops.LOAD]), 2)

  def test_fold_gated_store(self):
-    glbl = UOp.param(0, dtypes.int, 1)
+    glbl = UOp.param(0, dtypes.int, (1,))
    idx0 = UOp.const(0)
    val = UOp.const(42)
    st0 = glbl.index(UOp.invalid()).store(val)
@@ -425,7 +425,7 @@ class TestReduceCollapse(unittest.TestCase):

  def test_reduce_shapeless_const_unroll(self):
    """a REDUCE over a shapeless CONST (e.g. x*0 folded late in codegen) must collapse before the expander"""
-    out = UOp.param(0, dtypes.float, 1)
+    out = UOp.param(0, dtypes.float, (1,))
    red = UOp.const(3.0).cast(dtypes.float).reduce(UOp.range(4, 0, AxisType.UNROLL), arg=(Ops.ADD, 0))
    ast = UOp.sink(out.index(UOp.const(0)).store(red)).replace(arg=KernelInfo())
    uops = full_rewrite_to_sink(ast, Device["CPU"].renderer, optimize=False).toposort()
@@ -441,8 +441,7 @@ class TestMovementOps(unittest.TestCase):
    self.assertEqual(result.op, Ops.INDEX)
    self.assertIs(result.src[0], src)
    self.assertEqual(result.shape, (4,))
-    # the only RESHAPE is src itself: the view of the flat param, the extra reshape was folded into the INDEX
-    self.assertEqual([u for u in result.toposort() if u.op is Ops.RESHAPE], [src])
+    self.assertNotIn(Ops.RESHAPE, [u.op for u in result.toposort()])

  def test_pm_mops_partial_reshape_index_suffix_mismatch_does_nothing(self):
    from tinygrad.schedule.prepare import pm_mops
@@ -4,21 +4,21 @@ from tinygrad import UOp
 class TestUOpRepr(unittest.TestCase):
  def test_simple_const(self):
    a = UOp.const(42)
-    self.assertEqual(repr(a), "UOp(Ops.CONST, arg=42, src=())")
+    self.assertEqual(repr(a), "UOp(Ops.CONST, dtypes.weakint, arg=42, src=())")
  def test_different_consts(self):
    a, b = UOp.const(42), UOp.const(3)
    expected = (
-      "UOp(Ops.ADD, arg=None, src=(\n" +
-      "  UOp(Ops.CONST, arg=42, src=()),\n" +
-      "  UOp(Ops.CONST, arg=3, src=()),))"
+      "UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
+      "  UOp(Ops.CONST, dtypes.weakint, arg=42, src=()),\n" +
+      "  UOp(Ops.CONST, dtypes.weakint, arg=3, src=()),))"
    )
    self.assertEqual(repr(a+b), expected)
  def test_walrus_operator_indentation(self):
    # The reference should have the same indentation as the definition
    a = UOp.const(42)
    expected = (
-      "UOp(Ops.ADD, arg=None, src=(\n" +
-      "  x0:=UOp(Ops.CONST, arg=42, src=()),\n" +
+      "UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
+      "  x0:=UOp(Ops.CONST, dtypes.weakint, arg=42, src=()),\n" +
      "  x0,))"
    )
    self.assertEqual(repr(a+a), expected)
@@ -26,9 +26,9 @@ class TestUOpRepr(unittest.TestCase):
    # Ensure indentation is consistent at multiple levels
    b = (a:=UOp.const(1)) + a
    expected = (
-      "UOp(Ops.MUL, arg=None, src=(\n" +
-      "  x0:=UOp(Ops.ADD, arg=None, src=(\n" +
-      "    x1:=UOp(Ops.CONST, arg=1, src=()),\n" +
+      "UOp(Ops.MUL, dtypes.weakint, arg=None, src=(\n" +
+      "  x0:=UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
+      "    x1:=UOp(Ops.CONST, dtypes.weakint, arg=1, src=()),\n" +
      "    x1,)),\n" +
      "  x0,))"
    )
@@ -1392,7 +1392,7 @@ class TestInvalidIndex(unittest.TestCase):

  def test_gated_load_keeps_index_valid(self):
    # the load executes even on gated-off iterations: gated_given_valid must not erase its mask (PADTO OOB shape)
-    buf = UOp.param(0, dtypes.bool, 17)
+    buf = UOp.param(0, dtypes.bool, (17,))
    ridx = Variable("ridx", 0, 31)
    cond = ridx < 17
    load = buf.index(ridx.valid(cond))
@@ -1404,7 +1404,7 @@ class TestStoreLoadFolding(unittest.TestCase):
  """Tests for store(index, load(index)) -> NOOP rule. This rule matches patterns that EMERGE during simplification."""
  def test_store_load_folding(self):
    # store(idx, load(idx)) -> NOOP, including emergent patterns like store(idx, load(idx) + 0)
-    buf = UOp.param(0, dtypes.int, 1)
+    buf = UOp.param(0, dtypes.int, (1,))
    index = buf.index(UOp.const(0))
    # Direct: store(idx, load(idx)) -> NOOP
    self.assertEqual(graph_rewrite(index.store(index.load()), sym).op, Ops.NOOP)
@@ -1417,7 +1417,7 @@ class TestStoreLoadFolding(unittest.TestCase):

 class TestMoveWhereOnLoad(unittest.TestCase):
  def test_bool_index_preserves_dtype(self):
-    buf = UOp.param(0, dtypes.bool, 8)
+    buf = UOp.param(0, dtypes.bool, (8,))
    a = Variable("a", 0, 7)
    r = UOp.range(8, 0)
    # cond has a range that the rewrite can move into the valid: gate (a<4) goes into load valid
@@ -1475,7 +1475,7 @@ class TestRangeSplitting(unittest.TestCase):
    from tinygrad.codegen.simplify import pm_split_ranges, pm_flatten_range
    r0 = UOp.range(uconst(8), 0)
    # create a simple expression using the range with mod: store range%2 to a buffer
-    buf = UOp.param(0, dtypes.int, 1)
+    buf = UOp.param(0, dtypes.int, (1,))
    val = (r0 % uconst(2)).cast(dtypes.int)
    store = UOp(Ops.STORE, src=(buf.index(uconst(0)), val))
    sink = UOp(Ops.SINK, src=(UOp(Ops.END, src=(store, r0)),))
@@ -82,7 +82,7 @@ class TestVminVmaxProperties(unittest.TestCase):
  def test_vmin_vmax_multiplication_0_inf(self):
    # vmin and vmax for multiplication with a variable
    x = UOp.const(0.0)
-    y = UOp.load(UOp.param(0, dtypes.float, 1), UOp.const(0))
+    y = UOp.load(UOp.param(0, dtypes.float, (1,)), UOp.const(0))
    uop = x * y
    # TODO: these should be 0, but definitely should not be nan
    self.assertEqual(uop.vmin, -math.inf)
@@ -332,7 +332,7 @@ class TestVminVmaxVConst(unittest.TestCase):

  def test_vmin_vmax_vector_with_gep(self):
    # vmin and vmax for a vector constant of bool values
-    d1 = UOp.param(1, dtypes.int, 1)
+    d1 = UOp.param(1, dtypes.int, (1,))
    idx = UOp.const(0)
    val = UOp(Ops.LOAD, src=(d1.index(idx),))
    uop = (val // 32)
@@ -38,9 +38,15 @@ class TestDTypeFromUOp(unittest.TestCase):
    self.assertEqual(UOp(Ops.CONST, arg=ConstFloat(3.0)).dtype, dtypes.weakfloat)
    self.assertEqual(UOp(Ops.CONST, arg=True).dtype, dtypes.bool)
    self.assertEqual(UOp(Ops.CONST, arg=Invalid).dtype, dtypes.bool)
-    # UOp.const at a strong dtype builds the CAST that carries it
+    # an explicit (strong) const dtype is legal until the field is removed
    self.assertEqual(UOp.const(3, dtypes.int32).dtype, dtypes.int32)

+  def test_invalid_stated_dtype(self):
+    # UOp.const normalizes a stated dtype away (const_like/full pass their position's); the core constructor does not,
+    # and the spec is what rejects a non-bool Invalid
+    self.assertIs(UOp.const(Invalid, dtypes.float32), UOp.invalid())
+    with self.assertRaises(RuntimeError): type_verify(UOp(Ops.CONST, dtypes.float32, arg=Invalid), spec_shared)
+
  def test_invalid_dtype_and_consumers(self):
    invalid = UOp.invalid()
    self.assertIs(invalid.dtype, dtypes.bool)
@@ -52,7 +58,7 @@ class TestDTypeFromUOp(unittest.TestCase):
    for u in (UOp.param(0, dtypes.bool, ()).where(value, invalid), value+invalid, UOp.stack(value, invalid)): self.assertIs(u.src[-1], invalid)
    for u in (UOp(Ops.STACK, src=(value, invalid)), UOp(Ops.ADD, src=(value, invalid)),
              UOp.const(True).where(value, invalid), UOp(Ops.CMPLT, src=(invalid, value)), UOp(Ops.CMPLT, src=(value, invalid)),
-              UOp.param(0, dtypes.float32, 4).index(invalid)): type_verify(u, spec_shared)
+              UOp.param(0, dtypes.float32, (4,)).index(invalid)): type_verify(u, spec_shared)
    gate, value = UOp.param(0, dtypes.bool, ()), UOp.param(1, dtypes.float, ())
    self.assertIs((out:=graph_rewrite(gate.where(value, UOp.invalid()), pm_remove_invalid)).src[2], UOp.const(0, dtypes.float))
    type_verify(out.sink(), spec_program)
@@ -67,7 +73,7 @@ class TestLowerIndexDtype(unittest.TestCase):
  def test_gated_shrink_lowers_to_selected_width(self):
    # coalesce builds gated SHRINKs for masked vectorized loads; lowering must resolve them at the
    # width the offset bounds select (this one needs long)
-    buf = UOp.param(0, dtypes.float, 2**31+64)
+    buf = UOp.param(0, dtypes.float, (2**31+64,))
    i = UOp.variable("i", 0, 2**28)
    shrink = UOp(Ops.SHRINK, src=(buf, (i*24).valid(i < 2**28), UOp.const(4)))
    lowered = graph_rewrite(shrink.sink(), pm_lower_weak)
@@ -78,10 +84,11 @@ class TestLowerIndexDtype(unittest.TestCase):

  def test_reg_buffer_size_lowers(self):
    reg = UOp.placeholder((4,), dtypes.float, 0, addrspace=AddrSpace.REG)
-    self.assertEqual(reg.arg.size, 4)
+    self.assertEqual(reg.src[0].dtype, dtypes.weakint)
    lowered = graph_rewrite(reg.sink(), pm_lower_weak)
    self.assertTrue(all(u.op is Ops.CONST for u in lowered.backward_slice_with_self if u.dtype in dtypes.weaks),
                    "lowering must resolve every weak width, except a typed literal's value half")
+    self.assertEqual(next(u for u in lowered.backward_slice_with_self if u.op is Ops.BUFFER).src[0].dtype, dtypes.int)

 class TestSafeCast(unittest.TestCase):
  def test_cast_folds(self):
@@ -214,7 +221,7 @@ class TestExecALU(unittest.TestCase):

 class TestGatedStoreRewrite(unittest.TestCase):
  def test_tiny_gate_store(self):
-    gmem = UOp.param(0, dtypes.float, 8)
+    gmem = UOp.param(0, dtypes.float, (8,))
    gidx0 = UOp.special(4, 'gidx0')
    gate = gidx0<UOp.const(1)
    idx = UOp(Ops.INDEX, src=(gmem, (gidx0 * UOp.const(2)).valid(gate)))
@@ -230,8 +237,8 @@ class TestGatedStoreRewrite(unittest.TestCase):
    self.assertEqual(len(gated_uops[-1].src), 2)

  def test_gate_some_stores(self):
-    gmem0 = UOp.param(0, dtypes.float, 8)
-    gmem1 = UOp.param(1, dtypes.float, 8)
+    gmem0 = UOp.param(0, dtypes.float, (8,))
+    gmem1 = UOp.param(1, dtypes.float, (8,))
    gidx0 = UOp.special(4, 'gidx0')
    idx = gidx0 * UOp.const(2)
    idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gidx0<UOp.const(1))))
@@ -250,8 +257,8 @@ class TestGatedStoreRewrite(unittest.TestCase):
  # scaled down version of TestLinearizerDumb.test_unmerged_ifs
  @unittest.skip("we don't merge ifs anymore")
  def test_merge_ifs_alt(self):
-    gmem0 = UOp.param(0, dtypes.float, 8)
-    gmem1 = UOp.param(1, dtypes.float, 8)
+    gmem0 = UOp.param(0, dtypes.float, (8,))
+    gmem1 = UOp.param(1, dtypes.float, (8,))
    gidx0 = UOp.special(4, 'gidx0')
    idx = gidx0*UOp.const(2)
    gate = gidx0<UOp.const(1)
@@ -274,7 +281,7 @@ class TestGatedStoreRewrite(unittest.TestCase):
 class TestFastIdiv(unittest.TestCase):
  def test_division_power_of_two(self):
    for dt in (dtypes.int32, dtypes.uint32):
-      g = UOp.param(0, dt, 3)
+      g = UOp.param(0, dt, (3,))
      c = UOp.const(2)
      l = g.index(c)
      a = UOp(Ops.CDIV, src=(l, c))
@@ -287,7 +294,7 @@ class TestFastIdiv(unittest.TestCase):
  def test_floormod_power_of_two(self):
    # FLOORMOD by a power of two lowers to AND (correct floor mod for any sign in two's complement)
    for dt in (dtypes.int32, dtypes.uint32):
-      g = UOp.param(0, dt, 9)
+      g = UOp.param(0, dt, (9,))
      c = UOp.const(8)
      a = UOp(Ops.FLOORMOD, src=(g.index(c), c))
      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -299,7 +306,7 @@ class TestFastIdiv(unittest.TestCase):
  def test_floordiv_power_of_two(self):
    # FLOORDIV by a power of two lowers to a shift, with no round toward zero correction (a shift is exactly floor division)
    for dt in (dtypes.int32, dtypes.uint32, dtypes.int64, dtypes.uint64):
-      g = UOp.param(0, dt, 3)
+      g = UOp.param(0, dt, (3,))
      c = UOp.const(2)
      a = UOp(Ops.FLOORDIV, src=(g.index(c), c))
      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -312,7 +319,7 @@ class TestFastIdiv(unittest.TestCase):
  @Context(DISABLE_FAST_IDIV=0)
  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
  def test_fast_idiv_and_mod(self):
-    g = UOp.param(0, dtypes.uint32, 4)
+    g = UOp.param(0, dtypes.uint32, (4,))
    c = UOp.const(3)
    l = g.index(c)
    a = UOp(Ops.CDIV, src=(l, c))
@@ -346,7 +353,7 @@ class TestFastIdiv(unittest.TestCase):
  @unittest.expectedFailure
  def test_fast_idiv_overflow(self):
    # This will be possible with a slightly different method for fast_idiv
-    g = UOp.param(0, dtypes.uint32, 8)
+    g = UOp.param(0, dtypes.uint32, (8,))
    c = UOp.const(7).cast(dtypes.uint)
    l = UOp(Ops.LOAD, src=(g.index(c),))
    a = UOp(Ops.CDIV, src=(l, c))
@@ -357,7 +364,7 @@ class TestFastIdiv(unittest.TestCase):
    self.assertNotIn(Ops.CDIV, ops)

  def test_disable_fast_idiv(self):
-    g = UOp.param(0, dtypes.uint32, 4)
+    g = UOp.param(0, dtypes.uint32, (4,))
    c = UOp.const(3)
    l = g.index(c)
    a = UOp(Ops.CDIV, src=(l, c))
@@ -400,8 +407,8 @@ class TestUOpMethod(unittest.TestCase):
      self.assertIs((x != x).simplify(), x.const_like(False, dtypes.bool))

  def test_replace(self):
-    x = UOp.param(0, dtypes.int, 1)
-    self.assertEqual(x.replace(arg=UOp.param(1, dtypes.int, 1).arg).arg.slot, 1)
+    x = UOp.param(0, dtypes.int, (1,))
+    self.assertEqual(x.replace(arg=UOp.param(1, dtypes.int, (1,)).arg).arg.slot, 1)
    with self.assertRaises(AssertionError): x.replace(field="a")

  def test_const_zero_neg_zero_different(self):
@@ -456,7 +463,7 @@ class TestUOpRender(unittest.TestCase):
  def test_render_ssimplified_marg_outside_toposort(self):
    r = UOp.range(UOp.const(16, dtypes.int), 2, AxisType.WEAK, dtype=dtypes.int)
    offset = (r * 2) + (r * 2)
-    shrink = UOp(Ops.SHRINK, src=(UOp.param(0, dtypes.uint, 32), offset, UOp.const(2, dtypes.int)))
+    shrink = UOp(Ops.SHRINK, src=(UOp.param(0, dtypes.uint, (32,)), offset, UOp.const(2, dtypes.int)))
    self.assertIsNot(shrink.src[1], shrink.marg[0][0])
    self.assertEqual(shrink.render(simplify=False), "p0.shrink((((r2*4), 2),))")
    self.assertEqual(UOp.range(1, 0, src=(shrink,), dtype=dtypes.int).render(simplify=False), "r0")
@@ -146,7 +146,7 @@ class TestUOpsStats(unittest.TestCase):

  #MULACC should have the same stats as MUL + ADD
  def test_mulacc(self):
-    globl = UOp.param(0, dtypes.int, 3)
+    globl = UOp.param(0, dtypes.int, (3,))
    o1 = UOp.const(1, dtypes.int)
    o2 = UOp.const(2, dtypes.int)
    u1 = globl.index(o1)
@@ -156,7 +156,7 @@ class TestUOpsStats(unittest.TestCase):
    u5 = UOp(Ops.ADD, src=(u4,u3))
    uops = tuple(u5.toposort())

-    globl = UOp.param(0, dtypes.int, 3)
+    globl = UOp.param(0, dtypes.int, (3,))
    o1 = UOp.const(1, dtypes.int)
    o2 = UOp.const(2, dtypes.int)
    u1 = globl.index(o1)
@@ -13,7 +13,7 @@ class TestValidateOOB(unittest.TestCase):
  # basic index patterns
  def test_const_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      to_uops_list([buf.index(UOp.const(0)).load()])  # valid
      to_uops_list([buf.index(UOp.const(15)).load()])  # valid (last element)
      with self.assertRaises(RuntimeError):
@@ -23,7 +23,7 @@ class TestValidateOOB(unittest.TestCase):

  def test_variable_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      to_uops_list([buf.index(Variable("i", 0, 15)).load()])  # valid
      with self.assertRaises(RuntimeError):
        to_uops_list([buf.index(Variable("i", 0, 20)).load()])  # oob
@@ -32,7 +32,7 @@ class TestValidateOOB(unittest.TestCase):

  def test_range_with_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      r = UOp.range(42, 0, AxisType.GLOBAL)
      to_uops_list([buf.index(r.valid(r < 16)).load()])  # valid
      with self.assertRaises(RuntimeError):
@@ -40,7 +40,7 @@ class TestValidateOOB(unittest.TestCase):

  def test_variable_with_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      v = Variable("v", -5, 80)
      to_uops_list([buf.index(v.valid((v >= 0) & (v < 16))).load()])  # valid
      with self.assertRaises(RuntimeError):
@@ -48,7 +48,7 @@ class TestValidateOOB(unittest.TestCase):

  def test_gated_store(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      v = Variable("v", 0, 20)
      to_uops_list([buf.index(v.valid(v < 16)).store(0)])  # valid
      with self.assertRaises(RuntimeError):
@@ -57,14 +57,14 @@ class TestValidateOOB(unittest.TestCase):
  # ALU ops in index
  def test_floordiv(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      to_uops_list([buf.index(UOp.range(32, 0, AxisType.GLOBAL) // 2).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
        to_uops_list([buf.index(UOp.range(34, 0, AxisType.GLOBAL) // 2).load()])  # 0..16 oob

  def test_mod(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      r = UOp.range(100, 0, AxisType.GLOBAL)
      to_uops_list([buf.index(r % 16).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
@@ -72,14 +72,14 @@ class TestValidateOOB(unittest.TestCase):

  def test_shr(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      to_uops_list([buf.index(UOp.range(64, 0, AxisType.GLOBAL) >> 2).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
        to_uops_list([buf.index(UOp.range(128, 0, AxisType.GLOBAL) >> 2).load()])  # 0..31 oob

  def test_shl(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 64)
+      buf = UOp.param(0, dtypes.int, (64,))
      r = UOp.range(8, 0, AxisType.GLOBAL)
      to_uops_list([buf.index(r << 2).load()])  # 0..28 valid
      with self.assertRaises(RuntimeError):
@@ -87,7 +87,7 @@ class TestValidateOOB(unittest.TestCase):

  def test_and(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      r = UOp.range(100, 0, AxisType.GLOBAL)
      to_uops_list([buf.index(r & 15).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
@@ -102,14 +102,14 @@ class TestValidateOOB(unittest.TestCase):

  def test_max(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      to_uops_list([buf.index(Variable("v", -10, 15).maximum(0)).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
        to_uops_list([buf.index(Variable("v2", -10, 20).maximum(0)).load()])  # 0..20 oob

  def test_xor_in_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      r = UOp.range(32, 0, AxisType.GLOBAL)
      to_uops_list([buf.index(r.valid((r < 8) ^ ((r >= 8) & (r < 16)))).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
@@ -118,22 +118,22 @@ class TestValidateOOB(unittest.TestCase):
  # cast patterns
  def test_float_cast_in_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      r = UOp.range(20, 0)
      i = (r.cast(dtypes.float) * 0.68).trunc().cast(dtypes.int)
      to_uops_list([buf.index(i.valid((i >= 0) & (i < 16))).load()])

  def test_bool_cast_in_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 1)
+      buf = UOp.param(0, dtypes.int, (1,))
      r = UOp.range(20, 0)
      to_uops_list([buf.index(r.valid(r.cast(dtypes.bool).logical_not())).load()])  # only r=0 valid

  # load result as index/mask
  def test_load_as_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf0 = UOp.param(0, dtypes.int, 16)
-      buf1 = UOp.param(1, dtypes.int, 64)
+      buf0 = UOp.param(0, dtypes.int, (16,))
+      buf1 = UOp.param(1, dtypes.int, (64,))
      r = UOp.range(42, 0, AxisType.GLOBAL)
      ld0 = buf0.index(r.valid(r < 8)).load().cast(dtypes.weakint)
      to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32))).load()])  # valid
@@ -142,16 +142,16 @@ class TestValidateOOB(unittest.TestCase):

  def test_load_from_shrink_as_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf0 = UOp.param(0, dtypes.int, 16)
-      buf1 = UOp.param(1, dtypes.int, 64)
+      buf0 = UOp.param(0, dtypes.int, (16,))
+      buf1 = UOp.param(1, dtypes.int, (64,))
      shrink = UOp(Ops.SHRINK, src=(buf0, UOp.const(0, dtypes.int), UOp.const(4)))
      ld0 = shrink.load().index(0)
      to_uops_list([buf1.index(ld0.valid((ld0 >= 0) & (ld0 < 64))).load()])

  def test_load_bool_as_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf_bool = UOp.param(0, dtypes.bool, 16)
-      buf_int = UOp.param(1, dtypes.int, 8)
+      buf_bool = UOp.param(0, dtypes.bool, (16,))
+      buf_int = UOp.param(1, dtypes.int, (8,))
      gidx = UOp(Ops.SPECIAL, src=(UOp.const(16),), arg="gidx0")
      ld_bool = buf_bool.index(gidx).load()
      with self.assertRaises(RuntimeError):
@@ -162,7 +162,7 @@ class TestValidateOOB(unittest.TestCase):
  def test_in_bounds_access_gated_local(self):
    with Context(CHECK_OOB=1):
      # Define buffers
-      gbuf = UOp.param(0, dtypes.uint, 400)
+      gbuf = UOp.param(0, dtypes.uint, (400,))
      sbuf = UOp.placeholder((8,), dtypes.uint, slot=0, addrspace=AddrSpace.LOCAL)

      # Define indices, valids and barrier
@@ -186,8 +186,8 @@ class TestValidateOOB(unittest.TestCase):
  @unittest.skip("Bool load is not supported yet")
  def test_load_mask(self):
    with Context(CHECK_OOB=1):
-      glbl0 = UOp.param(0, dtypes.int, 16)
-      mask = UOp.param(0, dtypes.bool, 16)
+      glbl0 = UOp.param(0, dtypes.int, (16,))
+      mask = UOp.param(0, dtypes.bool, (16,))
      ridx = UOp.range(20, 0)
      ld0 = UOp(Ops.LOAD, src=(glbl0.index(UOp.const(ridx<16&mask, ridx))))
      to_uops_list([ld0])
@@ -244,7 +244,7 @@ class TestViz(unittest.TestCase):
    self.assertEqual(len(lst), 1)
    graphs = [x["graph"] for x in viz.get_details(0, 0)]
    # const is always in the graph, client side hides exclude=True nodes by default
-    self.assertEqual(list(graphs[0]), [id(a), id(z), id(alu), id(y), id(sink)])
+    self.assertEqual(list(graphs[0]), [id(a.src[0]), id(a), id(z), id(alu), id(y), id(sink)])
    self.assertTrue(graphs[0][id(z)]["exclude"])
    self.assertTrue(graphs[0][id(y)]["exclude"])
    self.assertFalse(graphs[0][id(alu)]["exclude"])
@@ -841,7 +841,7 @@ class TestCfg(unittest.TestCase):
      lidx = UOp.special(1, "lidx0")
      gidx = UOp.special(1, "gidx0")
      sink = UOp.sink(out.base, lidx, gidx, arg=KernelInfo(name=name))
-      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]))))
    with save_viz() as viz:
      with Context(DEV="NULL::gfx1100"):
        out = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
@@ -33,8 +33,8 @@ class TestRingAllReduce(unittest.TestCase):
      # N*(N-1) copies for input and output
      copy_count = N*(N-1)*2
      if len(copies) != copy_count: raise KernelCountException(copy_count, len(copies))
-      # N*(N-1) shrinks from other devices becoming contigs, N ALU, N extra contig, reassembly (cat), and mul
-      sink_count = (N*(N-1))+(N)+(N)+(1)+(1)
+      # N*N shrinks becoming contigs, N ALU, N extra contig, reassembly (cat), and mul
+      sink_count = (N*N)+(N)+(N)+(1)+(1)
      if len(sinks) != sink_count: raise KernelCountException(sink_count, len(sinks))
      # correctness
      run_linear(linear, var_vals)
@@ -1,10 +1,10 @@
 #!/usr/bin/env python
 import unittest
 import numpy as np
-from tinygrad import Device, dtypes, Tensor, TinyJit, GlobalCounters, Variable
+from tinygrad import dtypes, Tensor, TinyJit, GlobalCounters, Variable
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.helpers import temp, DEV, Context
-from test.helpers import assert_kernel_count, needs_second_gpu
+from test.helpers import assert_kernel_count

 N = 200  # has to be bigger than the cache to fail

@@ -1079,80 +1079,5 @@ class TestBatchNormRunningStats(unittest.TestCase):
    with Context(TRAINING=1): bn(x).realize()
    self.assertTrue(bn.running_mean.uop.base.is_realized)

-class TestMultiAssign(unittest.TestCase):
-  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
-
-  @needs_second_gpu
-  def setUp(self): pass
-
-  def test_multi_assign_realized(self):
-    out = Tensor.zeros(4).shard(self.device, 0).contiguous().realize()
-    ones = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
-    out.assign(ones).realize()
-    self.assertListEqual(out.tolist(), [1,1,1,1])
-
-  def test_multi_assign_unrealized(self):
-    out = Tensor.zeros(4).contiguous().realize().shard(self.device, 0)
-    ones = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
-    out.assign(ones).realize()
-    self.assertListEqual(out.tolist(), [1,1,1,1])
-
-  def test_multi_assign_both_unrealized(self):
-    out = Tensor.zeros(4).contiguous().realize().shard(self.device, 0)
-    ones = Tensor.ones(4).contiguous().realize().shard(self.device, 0)
-    out.assign(ones).realize()
-    self.assertListEqual(out.tolist(), [1,1,1,1])
-
-  def test_multi_assign_scalar(self):
-    out = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
-    out.assign(0).realize()
-    self.assertListEqual(out.tolist(), [0,0,0,0])
-
-  def test_multi_assign_const_like(self):
-    out = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
-    out.assign(out.const_like(7)).realize()
-    self.assertListEqual(out.tolist(), [7,7,7,7])
-
-  def test_multi_assign_piece(self):
-    out = Tensor.zeros(4,4).shard(self.device, 0).contiguous().realize()
-    ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
-    out[:, 2:3].assign(ones).realize()
-    self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
-
-  def test_multi_assign_piece_noncontig(self):
-    out = Tensor.zeros(4,4).contiguous().realize().shard(self.device, 0).realize()
-    ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
-    out[:, 2:3].assign(ones).realize()
-    self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
-
-  @unittest.expectedFailure
-  def test_multi_assign_piece_unrealized(self):
-    out = Tensor.zeros(4,4).contiguous().realize().shard(self.device, 0)
-    ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
-    out[:, 2:3].assign(ones).realize()
-    self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
-
-  def test_multi_assign_var_offset(self):
-    out = Tensor.zeros(4,4).contiguous().realize().shard(self.device, 0).realize()
-    ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
-    vi = Variable("i", 0, 3).bind(2)
-    out[:, vi:vi+1].assign(ones).realize()
-    self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
-
-  def test_multi_assign_var_offset_jit_none(self): self.test_multi_assign_var_offset_jit(None)
-  def test_multi_assign_var_offset_jit(self, shard_axis=0):
-    out = Tensor.zeros(4,6).contiguous().realize().shard(self.device, shard_axis).realize()
-    ones = Tensor.ones(4,1).shard(self.device, shard_axis).contiguous().realize()
-
-    @TinyJit
-    def f(out:Tensor, vi):
-      out[:, vi:vi+1].assign(ones).realize()
-      ones.assign(ones+1).realize()
-
-    vi = Variable("i", 0, 5)
-    for i in range(1,5):
-      GlobalCounters.reset()
-      f(out, vi.bind(i))
-    self.assertListEqual(out.tolist(), [[0,1,2,3,4,0]]*4)
 if __name__ == "__main__":
  unittest.main()
@@ -16,7 +16,7 @@ def apply_rope(x:Tensor, start_pos:int):

 class TestLinear(unittest.TestCase):
  def test_recovers_packed_ggml_weight(self):
-    for ggml_type,packed_size,words in ((13, 176, 44), (14, 210, 53), (23, 136, 34)):
+    for ggml_type,packed_size,words in ((13, 176, 44), (14, 210, 210), (23, 136, 34)):
      packed = Tensor.empty(packed_size+4, dtype=dtypes.uint8, device="CPU")[4:]
      decoded = ggml_data_to_tensor(packed, 256, ggml_type).reshape(1, 256)
      linear = Linear(256, 1, bias=False)
@@ -73,10 +73,10 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
    return Tensor.linspace(start, stop, int(np.prod(shape)), dtype=dtypes.float32).reshape(*shape)

  def _make_config(self, **kwargs):
-    return TransformerConfig(**({"num_blocks":1, "dim":8, "hidden_dim":16, "n_heads":1, "n_kv_heads":1,
-                                 "norm_eps":1e-5, "vocab_size":32, "head_dim":8, "rope_theta":10000.0,
-                                 "rope_dim":8, "v_head_dim":8, "max_context":4, "ssm_layers":(True,),
-                                 "ssm":SSMConfig(conv_kernel=2, state_size=4, group_count=1, time_step_rank=1, inner_size=4)} | kwargs))
+    return TransformerConfig(**({"num_blocks":1, "dim":32, "hidden_dim":64, "n_heads":1, "n_kv_heads":1,
+                                 "norm_eps":1e-5, "vocab_size":32, "head_dim":32, "rope_theta":10000.0,
+                                 "rope_dim":32, "v_head_dim":32, "max_context":4, "ssm_layers":(True,),
+                                 "ssm":SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32)} | kwargs))

  def _make_block(self, config:TransformerConfig) -> GatedDeltaNetBlock:
    block = GatedDeltaNetBlock(config, config.ssm)
@@ -229,7 +229,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
    np.testing.assert_allclose(block.recurrent_state.numpy(), initial_state.numpy() * alpha[..., None], rtol=1e-5, atol=1e-5)

  def test_kda_prefill_matches_decode(self):
-    config = self._make_config(ssm=SSMConfig(conv_kernel=2, state_size=4, group_count=1, time_step_rank=1, inner_size=4, kda=True))
+    config = self._make_config(ssm=SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32, kda=True))
    block = GatedDeltaNetBlock(config, config.ssm)
    for p in nn.state.get_parameters(block):
      p.replace(self._tensor_linspace(-0.05, 0.05, p.shape) if len(p.shape) > 1 else self._tensor_linspace(0.05, 0.1, p.shape))
@@ -245,7 +245,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase):

  def test_varied_chunk_sizes_match_decode(self):
    for kda in (False, True):
-      ssm = SSMConfig(conv_kernel=2, state_size=4, group_count=1, time_step_rank=1, inner_size=4, kda=kda)
+      ssm = SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32, kda=kda)
      config = self._make_config(ssm=ssm)
      if kda:
        block = GatedDeltaNetBlock(config, config.ssm)
@@ -267,8 +267,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
        np.testing.assert_allclose(chunked_recurrent, decode_recurrent, rtol=1e-3, atol=1e-3, err_msg=f"{kda=} {chunking=}")

  def test_start_zero_resets_realized_state(self):
-    config = self._make_config(max_context=3)
-    x = self._tensor_linspace(-1, 1, (1, 3, config.dim))
+    config, x = self._make_config(max_context=3), self._tensor_linspace(-1, 1, (1, 3, 32))
    block = self._make_block(config)
    self._run_attention(block, x, 0)
    restarted = self._run_attention(block, x[:, :2], 0)
@@ -1,115 +0,0 @@
-import struct, unittest
-from types import SimpleNamespace
-from unittest.mock import patch
-
-from tinygrad.runtime.autogen import bnxt
-from extra.bnxt_driver.bnxtdev import BNXT_BACKING_STORE, BNXTDev, BNXTQP, _queue, _qwrite, ipv4_to_gid
-
-class FakePCI:
-  def __init__(self): self.next_addr, self.allocations = 0x100000, []
-  def alloc_sysmem(self, size, contiguous=False):
-    pages = [self.next_addr+i*0x1000 for i in range((size+0xfff)//0x1000)]
-    self.next_addr += len(pages)*0x1000
-    self.allocations.append(mem := bytearray(size))
-    return mem, pages
-
-class FakeDev:
-  def __init__(self): self.pci_dev, self.calls = FakePCI(), []
-  def hwrm(self, name, **fields):
-    self.calls.append((name, fields))
-    typ = fields.get("type", 0)
-    return SimpleNamespace(ctx_init_value=0x5a, ctx_init_offset=4, entry_size=16 if typ == 0 else 4,
-                           subtype_valid_cnt=typ == 0, split_entry_0=2, instance_bit_map=5 if typ == 0 else 1, min_num_entries=0)
-
-class FakeRCFW:
-  def __init__(self): self.calls, self.doorbells = [], []
-  def exec(self, name, **fields):
-    self.calls.append((name, fields))
-    return SimpleNamespace(xid={"create_cq":77, "create_qp":88, "register_mr":0x5678}.get(name, 0))
-  def doorbell(self, *args, **kwargs): self.doorbells.append((args, kwargs))
-
-class FakeQPDev:
-  def __init__(self): self.pci_dev, self.fw, self.gid_id, self.nq_id = FakePCI(), FakeRCFW(), 9, 41
-  def rcfw(self, *args, **kwargs): return self.fw.exec(*args, **kwargs)
-  def doorbell(self, *args, **kwargs): self.fw.doorbell(*args, **kwargs)
-
-class TestMemory(unittest.TestCase):
-  def test_cmdq_and_sq_aux(self):
-    dev = FakeDev()
-    cmdq, sq = _queue(dev), _queue(dev, aux=True)
-    self.assertEqual((cmdq["level"], cmdq["base"]), (0, 0x100000))
-    _qwrite(sq, 3, b"ABCDEFGH", aux=True)
-    self.assertEqual(bytes(sq["mem"][0x1018:0x1020]), b"ABCDEFGH")
-
-  def test_f320_backing_layout_and_final_marker(self):
-    self.assertEqual(len(BNXT_BACKING_STORE), 9)
-    dev = FakeDev()
-    small = ((0, 6), (15, 0))
-    with patch("extra.bnxt_driver.bnxtdev.BNXT_BACKING_STORE", small): BNXTDev.setup_backing_store(dev)
-    cfg = [fields for name, fields in dev.calls if name == "func_backing_store_cfg_v2"]
-    self.assertEqual([(x["type"], x["instance"]) for x in cfg], [(0, 0), (0, 2), (15, 0)])
-    self.assertTrue(all(not x["flags"] for x in cfg[:-1]))
-    self.assertEqual(cfg[-1]["flags"], bnxt.FUNC_BACKING_STORE_CFG_V2_REQ_FLAGS_BS_CFG_ALL_DONE)
-    self.assertEqual((dev.pci_dev.allocations[0][4], dev.pci_dev.allocations[0][20]), (0x5a, 0x5a))
-
-class TestRCFW(unittest.TestCase):
-  def setUp(self):
-    patch("extra.bnxt_driver.bnxtdev.System.memory_barrier").start()
-    self.addCleanup(patch.stopall)
-
-  def test_doorbell_encodes_xid_type_and_index(self):
-    dev = BNXTDev.__new__(BNXTDev)
-    dev.db, dev.db_off = [0]*1024, 0x1000
-    dev.doorbell(0x123456, bnxt.DBC_DBC_TYPE_CQ_ARMALL, 0x456, epoch=1)
-    key = dev.db[0x1000//8]
-    self.assertEqual(key >> 32,
-      0x123456 & bnxt.DBC_DBC_XID_MASK | bnxt.DBC_DBC_PATH_ROCE | bnxt.DBC_DBC_TYPE_CQ_ARMALL | bnxt.BNXT_QPLIB_DBR_VALID)
-    self.assertEqual(key & 0xffffffff, 0x456 | 1<<bnxt.BNXT_QPLIB_DBR_EPOCH_SHIFT)
-
-  def test_command_uses_first_flag(self):
-    dev = BNXTDev.__new__(BNXTDev)
-    dev.bar0, dev.cmdq, dev.creq = [0]*1024, _queue(FakeDev()), _queue(FakeDev())
-    dev.rcfw_first, dev.creq_id = True, 23
-    dev.doorbell = lambda *args: None
-    _qwrite(dev.creq, 0, bytes(bnxt.struct_creq_query_version_resp(type=bnxt.CREQ_BASE_TYPE_QP_EVENT, cookie=0, v=1)))
-    ret = dev.rcfw("query_version")
-    req = bnxt.struct_cmdq_query_version.from_buffer_copy(bytes(dev.cmdq["mem"][:16]))
-    prod = dev.bar0[(bnxt.RCFW_COMM_BASE_OFFSET+bnxt.RCFW_PF_VF_COMM_PROD_OFFSET)//4]
-    self.assertEqual((req.cookie, ret.cookie, prod), (0, 0, 1 | 1<<bnxt.FIRMWARE_FIRST_FLAG))
-
-class TestFastPath(unittest.TestCase):
-  def test_unified_mr(self):
-    dev = BNXTDev.__new__(BNXTDev)
-    fw = FakeRCFW()
-    dev.pci_dev, dev.rcfw = FakePCI(), fw.exec
-    self.assertEqual(dev.register_mem([0x800000, 0x900000], 0x2000), 0x5678)
-    mr = fw.calls[-1][1]
-    self.assertEqual((mr["flags"], mr["va"], mr["mr_size"], mr["log2_pg_size_lvl"]),
-      (bnxt.CMDQ_REGISTER_MR_FLAGS_ALLOC_MR, 0x800000, 0x2000,
-       1<<bnxt.CMDQ_REGISTER_MR_LVL_SFT | 12<<bnxt.CMDQ_REGISTER_MR_LOG2_PG_SIZE_SFT))
-
-  def test_qp_creation_and_connect_use_f320_layout(self):
-    dev = FakeQPDev()
-    qp = BNXTQP(dev)
-    create = next(fields for name, fields in dev.fw.calls if name == "create_qp")
-    self.assertEqual((create["sq_size"], "rq_size" in create, qp.qpn), (16, False, 88))
-    qp.connect(0x123, ipv4_to_gid("10.0.0.2"), 0x001122334455)
-    rtr, rts = dev.fw.calls[-2][1], dev.fw.calls[-1][1]
-    self.assertEqual((bytes(rtr["dgid"]), bytes(rtr["dest_mac"])),
-      (ipv4_to_gid("10.0.0.2"), bytes.fromhex("001122334455")))
-    self.assertEqual((rtr["modify_mask"], rts["modify_mask"]), (0x41515ad, 0xae005))
-
-  def test_rdma_write_builds_three_slots_and_host_msn(self):
-    qp = BNXTQP.__new__(BNXTQP)
-    qp.dev, qp.qpn = FakeQPDev(), 88
-    qp.sq, qp.sq_psn, qp.msn = _queue(FakeDev(), aux=True), 5, 0
-    qp._poll = lambda timeout: bytes(bnxt.struct_cq_req())
-    qp.rdma_write(0x1122334455667788, 0x99aa, 0x12345000, 0x55aa, 100)
-    hdr = bnxt.struct_sq_rdma_hdr.from_buffer_copy(bytes(qp.sq["mem"][:32]))
-    sge = bnxt.struct_sq_sge.from_buffer_copy(bytes(qp.sq["mem"][32:48]))
-    self.assertEqual((hdr.remote_va, hdr.remote_key, hdr.length, sge.va_or_pa, sge.l_key, sge.size),
-      (0x1122334455667788, 0x99aa, 100, 0x12345000, 0x55aa, 100))
-    self.assertEqual(struct.unpack_from("<Q", qp.sq["mem"], 0x1000)[0], 6<<24 | 5)
-    self.assertEqual(qp.dev.fw.doorbells, [((88, bnxt.DBC_DBC_TYPE_SQ, 3, 0), {})])
-
-if __name__ == "__main__": unittest.main()
@@ -1,40 +0,0 @@
-import ctypes, unittest
-from unittest.mock import patch
-
-from tinygrad.runtime.autogen import bnxt
-from extra.bnxt_driver.bnxtdev import BNXT_CHIMP_COMM, BNXT_CHIMP_COMM_TRIGGER, BNXTDev
-
-class Mailbox:
-  def __init__(self, trigger): self.words, self.trigger = {}, trigger
-  def __setitem__(self, idx, val):
-    self.words[idx] = val
-    if idx == BNXT_CHIMP_COMM_TRIGGER//4: self.trigger()
-  def request(self):
-    base = BNXT_CHIMP_COMM//4
-    return b"".join(self.words.get(base+i, 0).to_bytes(4, "little") for i in range(bnxt.HWRM_MAX_REQ_LEN//4))
-
-def fake_dev():
-  dev = BNXTDev.__new__(BNXTDev)
-  dev.resp, dev.resp_pa, dev.seq = bytearray(0x1000), [0x6789a000], 0
-  return dev
-
-def reply(dev, out_type):
-  req = bnxt.struct_hwrm_cmd_hdr.from_buffer_copy(dev.bar0.request())
-  out = out_type(req_type=req.req_type, seq_id=req.seq_id, resp_len=ctypes.sizeof(out_type), valid=1)
-  dev.resp[:ctypes.sizeof(out_type)] = bytes(out)
-
-class TestHWRM(unittest.TestCase):
-  def setUp(self):
-    self.barrier = patch("extra.bnxt_driver.bnxtdev.System.memory_barrier").start()
-    self.addCleanup(patch.stopall)
-
-  def test_request(self):
-    dev = fake_dev()
-    dev.bar0 = Mailbox(lambda: reply(dev, bnxt.struct_hwrm_func_qcaps_output))
-    dev.hwrm("func_qcaps", fid=0xffff)
-    req = bnxt.struct_hwrm_func_qcaps_input.from_buffer_copy(dev.bar0.request())
-    self.assertEqual((req.req_type, req.seq_id, req.resp_addr, req.fid),
-      (bnxt.HWRM_FUNC_QCAPS, 1, dev.resp_pa[0], 0xffff))
-    self.barrier.assert_called_once_with()
-
-if __name__ == "__main__": unittest.main()
@@ -52,11 +52,6 @@ class TestCall(unittest.TestCase):
    np.testing.assert_allclose(a.grad.numpy(), gt_a_grad, rtol=1e-5)
    np.testing.assert_allclose(b.grad.numpy(), gt_b_grad, rtol=1e-5)

-  def test_call_scalar_param_shape_mismatch(self):
-    scalar_fxn = UOp.param(0, dtypes.float, ()) * 2
-    with self.assertRaisesRegex(TypeError, "shape mismatch: expected scalar"):
-      Tensor.call(Tensor.ones(2), fxn=scalar_fxn).realize()
-
  def test_call_gemm(self):
    M, K, N = 4, 8, 4
    a = Tensor.randn(M, K)
@@ -116,7 +116,7 @@ class TestWeakPromotion(unittest.TestCase):

  def test_store_weak_value_uses_destination_dtype(self):
    with Context(DEFAULT_FLOAT=dtypes.float16):
-      dst = UOp.param(0, dtypes.bfloat16, 1).index(UOp.const(0).cast(dtypes.int32))
+      dst = UOp.param(0, dtypes.bfloat16, (1,)).index(UOp.const(0).cast(dtypes.int32))
      gate = UOp.const(True)
      out = graph_rewrite(dst.store(UOp.const(5.0), gate), pm_commit_weak)
    # a bare weak CONST commits directly: the pass runs without symbolic, so a CAST here would survive it
@@ -132,7 +132,7 @@ class TestWeakPromotion(unittest.TestCase):

  def test_derivable_const_rounds_at_the_derived_width(self):
    # re-rounds a derivable const in place (still bare) so value-keyed folds (x*1 -> x, x*-1 -> NEG) still fire
-    x = UOp.param(0, dtypes.float32, 1).index(UOp.const(0).cast(dtypes.int32)).load()
+    x = UOp.param(0, dtypes.float32, (1,)).index(UOp.const(0).cast(dtypes.int32)).load()
    mul = graph_rewrite(x * UOp.const(-0.9999999893980771), symbolic_simple+pm_commit_weak)
    self.assertIs(mul.src[1], UOp.const(-1.0))
    self.assertIs(graph_rewrite(x * UOp.const(1.0000000106), symbolic_simple+pm_commit_weak), x)
@@ -191,7 +191,7 @@ class TestWeakPromotion(unittest.TestCase):
    # float bitwise builds, the spec rejects it
    with Context(SPEC=1):
      f32, wf = UOp.const(1.0, dtypes.float32), UOp.const(1.0)
-      for bad in (f32.alu(Ops.AND, f32), UOp(Ops.AND, (f32, f32)), UOp(Ops.AND, (wf, wf))):
+      for bad in (f32.alu(Ops.AND, f32), UOp(Ops.AND, dtypes.float32, (f32, f32)), UOp(Ops.AND, dtypes.int32, (wf, wf))):
        with self.assertRaises(RuntimeError): type_verify([bad], spec_shared)

  def test_integer_values(self):
@@ -5,33 +5,17 @@ from tinygrad.llm.gguf import _ggml_iq_grid, ggml_data_to_tensor, gguf_load
 from tinygrad.runtime.autogen import ggml_common as _ggml
 import numpy as np
 from gguf import GGUFReader, GGUFValueType, GGMLQuantizationType, GGML_QUANT_SIZES, dequantize, quantize
-from gguf.quants import IQ1_S, IQ2_S, IQ2_XS, IQ2_XXS, IQ3_S, IQ3_XXS
+from gguf.quants import IQ2_S, IQ3_S, IQ3_XXS

 ggml_test_block_count = 4
 supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()

 class TestGGUFTables(unittest.TestCase):
-  def test_iq2_xxs_grid_matches_gguf_py(self):
-    IQ2_XXS.init_grid()
-    grid = _ggml_iq_grid(Device.DEFAULT, _ggml.iq2xxs_grid, (256, 8)).numpy()
-    np.testing.assert_equal(grid, IQ2_XXS.grid.reshape(256, 8))
-
-  def test_iq2_xs_grid_matches_gguf_py(self):
-    IQ2_XS.init_grid()
-    grid = _ggml_iq_grid(Device.DEFAULT, _ggml.iq2xs_grid, (512, 8)).numpy()
-    np.testing.assert_equal(grid, IQ2_XS.grid.reshape(512, 8))
-
  def test_iq2_s_grid_matches_gguf_py(self):
    IQ2_S.init_grid()
    grid = _ggml_iq_grid(Device.DEFAULT, _ggml.iq2s_grid, (1024, 8)).numpy()
    np.testing.assert_equal(grid, IQ2_S.grid.reshape(1024, 8))

-  def test_iq1_s_grid_matches_gguf_py(self):
-    IQ1_S.init_grid()
-    grid = _ggml_iq_grid(Device.DEFAULT, _ggml.iq1s_grid, (2048, 8)).numpy()
-    grid = np.where(grid > 127, grid - 256, grid)
-    np.testing.assert_equal(grid, IQ1_S.grid.reshape(2048, 8))
-
  def test_iq3_xxs_grid_matches_gguf_py(self):
    IQ3_XXS.init_grid()
    grid = _ggml_iq_grid(Device.DEFAULT, _ggml.iq3xxs_grid, (256, 4)).numpy()
@@ -55,56 +39,6 @@ class TestGGUF(unittest.TestCase):
    expected = np.arange(1, 33, dtype=np.float32) * 2.0
    np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 32, GGMLQuantizationType.Q8_0.value).numpy().flatten(), expected)

-  def test_dequantization_q2_k_hardcoded(self):
-    # Q2_K: scales[16] + qs[64] + d(fp16) + dmin(fp16). 16 sub-blocks of 16, x = d*(scale&0xF)*q - dmin*(scale>>4)
-    scales, qs = bytes([0x11]*16), bytes([0x55]*64)  # scale=1, min=1; qs=0x55 -> 2-bit quants of 1
-    d, dmin = np.float16(1.0).tobytes(), np.float16(0.0).tobytes()
-    block = np.frombuffer(scales + qs + d + dmin, dtype=np.uint8).copy()
-    np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 256, 10).numpy().flatten(), np.ones(256, dtype=np.float32))
-
-  def test_dequantization_q3_k_hardcoded(self):
-    # Q3_K: hmask[32] + qs[64] + scales[12] + d(fp16). 16 sub-blocks of 16, x = d * (scale-32) * (q - (hbit?0:4))
-    # 6-bit scales 32..47 so (scale-32) = 0..15; qs=0x55 -> 2-bit quants of 1; d=1.0
-    scales = bytes([0x80, 0x91, 0xA2, 0xB3, 0xC4, 0xD5, 0xE6, 0xF7, 0xAA, 0xAA, 0xAA, 0xAA])
-    d = np.float16(1.0).tobytes()
-    qs, ones = bytes([0x55]*64), np.ones(16, dtype=np.float32)
-    # hmask all-ones: high bit set, q=1; hmask zeros: subtract 4, q=-3
-    for hmask, q in ((bytes([0xFF]*32), 1.0), (bytes([0x00]*32), -3.0)):
-      block = np.frombuffer(hmask + qs + scales + d, dtype=np.uint8).copy()
-      expected = np.concatenate([q * s * ones for s in range(16)])
-      np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 256, 11).numpy().flatten(), expected)
-
-  def test_dequantization_iq2_xxs_hardcoded(self):
-    # IQ2_XXS: d + 8 groups of (4 grid bytes + uint32 signs/scale). grid[0]=all 0x08, scale=0, signs=0
-    # db = 1.0 * (0.5 + 0) * 0.25 = 0.125; 0.125 * 8 = 1.0
-    block = np.frombuffer(np.float16(1.0).tobytes() + bytes(64), dtype=np.uint8).copy()
-    np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 256, 16).numpy().flatten(), np.ones(256, dtype=np.float32))
-
-  def test_dequantization_iq2_xs_hardcoded(self):
-    # IQ2_XS: d + 32 uint16 qs + 8 scale bytes. qs=0 -> grid[0]=all 0x08, signs=0; scales=0
-    block = np.frombuffer(np.float16(1.0).tobytes() + bytes(64) + bytes(8), dtype=np.uint8).copy()
-    np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 256, 17).numpy().flatten(), np.ones(256, dtype=np.float32))
-
-  def test_dequantization_iq1_s_hardcoded(self):
-    # IQ1_S: d + qs[32] + qh[16]. qs=qh=0 -> grid[0]=all -1, scale=1, delta=+0.125 -> -0.875
-    block = np.frombuffer(np.float16(1.0).tobytes() + bytes(48), dtype=np.uint8).copy()
-    expected = np.full(256, -0.875, dtype=np.float32)
-    np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 256, 19).numpy().flatten(), expected)
-
-  def test_dequantization_iq1_m_hardcoded(self):
-    # IQ1_M: qs[32] + qh[16] + scales[8]. f16 1.0=0x3C00 packed in high nibbles; qs=qh=0 -> -0.875
-    scales = bytes([0x00, 0x00, 0x00, 0x00, 0x00, 0xC0, 0x00, 0x30])
-    block = np.frombuffer(bytes(48) + scales, dtype=np.uint8).copy()
-    expected = np.full(256, -0.875, dtype=np.float32)
-    np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 256, 29).numpy().flatten(), expected)
-
-  def test_dequantization_iq4_nl_hardcoded(self):
-    # IQ4_NL: 2-byte fp16 scale + 16 packed bytes. low nibbles first, then high
-    lut = list(_ggml.kvalues_iq4nl)
-    block = np.frombuffer(np.float16(1.0).tobytes() + bytes(range(16)), dtype=np.uint8).copy()
-    expected = np.array(lut + [lut[0]]*16, dtype=np.float32)
-    np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 32, 20).numpy().flatten(), expected)
-
  def test_dequantization_mxfp4_hardcoded(self):
    # MXFP4: 1 byte shared exponent E + 16 packed bytes (32 x 4-bit values)
    # nibble: bit3=sign, bit2:1=exp, bit0=mant; E=128 gives scale=1.0
@@ -118,20 +52,13 @@ class TestGGUF(unittest.TestCase):
  def test_dequantization_q5_0(self): self._test_dequantization(GGMLQuantizationType.Q5_0)
  def test_dequantization_q5_1(self): self._test_dequantization(GGMLQuantizationType.Q5_1)
  def test_dequantization_q8_0(self): self._test_dequantization(GGMLQuantizationType.Q8_0)
-  def test_dequantization_q2_k(self): self._test_dequantization(GGMLQuantizationType.Q2_K)
-  def test_dequantization_q3_k(self): self._test_dequantization(GGMLQuantizationType.Q3_K)
  def test_dequantization_q4_k(self): self._test_dequantization(GGMLQuantizationType.Q4_K)
  def test_dequantization_q5_k(self): self._test_dequantization(GGMLQuantizationType.Q5_K)
  def test_dequantization_q6_k(self): self._test_dequantization(GGMLQuantizationType.Q6_K)
-  def test_dequantization_iq2_xxs(self): self._test_dequantization(GGMLQuantizationType.IQ2_XXS)
-  def test_dequantization_iq2_xs(self): self._test_dequantization(GGMLQuantizationType.IQ2_XS)
  def test_dequantization_iq3_xxs(self): self._test_dequantization(GGMLQuantizationType.IQ3_XXS)
-  def test_dequantization_iq1_s(self): self._test_dequantization(GGMLQuantizationType.IQ1_S)
-  def test_dequantization_iq4_nl(self): self._test_dequantization(GGMLQuantizationType.IQ4_NL)
  def test_dequantization_iq3_s(self): self._test_dequantization(GGMLQuantizationType.IQ3_S)
  def test_dequantization_iq2_s(self): self._test_dequantization(GGMLQuantizationType.IQ2_S)
  def test_dequantization_iq4_xs(self): self._test_dequantization(GGMLQuantizationType.IQ4_XS)
-  def test_dequantization_iq1_m(self): self._test_dequantization(GGMLQuantizationType.IQ1_M)
  def test_dequantization_mxfp4(self): self._test_dequantization(GGMLQuantizationType.MXFP4)
  @unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, "Backend must support bfloat16")
  def test_dequantization_bf16(self): self._test_dequantization(GGMLQuantizationType.BF16)
@@ -276,23 +203,12 @@ class TestGGUFGEMV(unittest.TestCase):
      q_data = rng.integers(0, 256, size=n_blocks * type_size, dtype=np.uint8).reshape(n_blocks, type_size)
      scales = np.float16(rng.standard_normal(n_blocks * 4)).view(np.uint8).reshape(n_blocks, -1)
      if qtype in (GGMLQuantizationType.Q5_0, GGMLQuantizationType.Q8_0,
-                   GGMLQuantizationType.IQ2_XXS, GGMLQuantizationType.IQ2_XS,
-                   GGMLQuantizationType.IQ3_XXS, GGMLQuantizationType.IQ4_NL,
-                   GGMLQuantizationType.IQ1_S, GGMLQuantizationType.IQ2_S,
+                   GGMLQuantizationType.IQ3_XXS,
+                   GGMLQuantizationType.IQ2_S,
                   GGMLQuantizationType.IQ3_S, GGMLQuantizationType.IQ4_XS): q_data[:, :2] = scales[:, :2]  # d at offset 0
      elif qtype in (GGMLQuantizationType.Q5_1, GGMLQuantizationType.Q4_K, GGMLQuantizationType.Q5_K):
        q_data[:, :4] = scales[:, :4]  # d, m/dmin at offset 0
-      elif qtype == GGMLQuantizationType.Q2_K: q_data[:, -4:] = scales[:, :4]  # d, dmin at end
-      elif qtype in (GGMLQuantizationType.Q6_K, GGMLQuantizationType.Q3_K): q_data[:, -2:] = scales[:, :2]  # d at end
-      elif qtype == GGMLQuantizationType.IQ1_M:
-        s = np.float16(rng.standard_normal(n_blocks)).view(np.uint16)
-        sc = q_data[:, -8:].copy().view(np.uint16).reshape(n_blocks, 4)
-        sc &= np.uint16(0x0FFF)
-        sc[:, 0] |= (s & np.uint16(0x000F)) << 12
-        sc[:, 1] |= (s & np.uint16(0x00F0)) << 8
-        sc[:, 2] |= (s & np.uint16(0x0F00)) << 4
-        sc[:, 3] |= (s & np.uint16(0xF000))
-        q_data[:, -8:] = sc.reshape(n_blocks, -1).view(np.uint8)
+      elif qtype == GGMLQuantizationType.Q6_K: q_data[:, -2:] = scales[:, :2]               # d at end
      elif qtype == GGMLQuantizationType.MXFP4: q_data[:, 0] = rng.integers(120, 136, size=n_blocks, dtype=np.uint8) # constrain byte0
      q_data = q_data.flatten()
    ref = dequantize(q_data, qtype).reshape(rows, cols)
@@ -319,20 +235,13 @@ class TestGGUFGEMV(unittest.TestCase):
  def test_gguf_gemv_q8_0(self): self._test_gguf_gemv(GGMLQuantizationType.Q8_0)
  def test_gguf_gemv_q5_0(self): self._test_gguf_gemv(GGMLQuantizationType.Q5_0)
  def test_gguf_gemv_q5_1(self): self._test_gguf_gemv(GGMLQuantizationType.Q5_1)
-  def test_gguf_gemv_q2_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q2_K)
-  def test_gguf_gemv_q3_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q3_K)
  def test_gguf_gemv_q4_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q4_K)
  def test_gguf_gemv_q5_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q5_K)
  def test_gguf_gemv_q6_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q6_K)
-  def test_gguf_gemv_iq2_xxs(self): self._test_gguf_gemv(GGMLQuantizationType.IQ2_XXS)
-  def test_gguf_gemv_iq2_xs(self): self._test_gguf_gemv(GGMLQuantizationType.IQ2_XS)
  def test_gguf_gemv_iq3_xxs(self): self._test_gguf_gemv(GGMLQuantizationType.IQ3_XXS)
-  def test_gguf_gemv_iq1_s(self): self._test_gguf_gemv(GGMLQuantizationType.IQ1_S)
-  def test_gguf_gemv_iq4_nl(self): self._test_gguf_gemv(GGMLQuantizationType.IQ4_NL)
  def test_gguf_gemv_iq3_s(self): self._test_gguf_gemv(GGMLQuantizationType.IQ3_S)
  def test_gguf_gemv_iq2_s(self): self._test_gguf_gemv(GGMLQuantizationType.IQ2_S)
  def test_gguf_gemv_iq4_xs(self): self._test_gguf_gemv(GGMLQuantizationType.IQ4_XS)
-  def test_gguf_gemv_iq1_m(self): self._test_gguf_gemv(GGMLQuantizationType.IQ1_M)
  def test_gguf_gemv_mxfp4(self): self._test_gguf_gemv(GGMLQuantizationType.MXFP4)
  @unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, "Backend must support bfloat16")
  def test_gguf_gemv_bf16(self): self._test_gguf_gemv(GGMLQuantizationType.BF16)
@@ -20,13 +20,11 @@ class TestQ8Quantize(unittest.TestCase):
  def test_values_and_scales(self):
    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
    x = np.linspace(-3.1, 2.7, 64, dtype=np.float32).reshape(2, 32)
-    quant, scale, gsum = q8_quantize(Tensor(x), 2, 32)
+    quant, scale = q8_quantize(Tensor(x), 2, 32)
    scale_np = np.maximum(np.max(np.abs(x), axis=-1, keepdims=True) / 127, 1e-8)
    expected = np.clip(np.rint(x / scale_np), -127, 127).astype(np.int8)
    np.testing.assert_array_equal(quant.bitcast(dtypes.int8).reshape(2, 32).numpy(), expected)
    np.testing.assert_allclose(scale.numpy(), scale_np, rtol=1e-6)
-    # xsum holds the two per-16 sums per 32-wide group
-    np.testing.assert_array_equal(gsum.numpy().reshape(2, 2), expected.reshape(2, 2, 16).sum(-1).astype(np.float32))

  def test_q6_linear_compiles(self):
    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
@@ -38,9 +36,7 @@ class TestQ8Quantize(unittest.TestCase):
    linear = Linear(256, 1, bias=False)
    nn.state.load_state_dict(linear, {"weight":decoded}, verbose=False, realize=False)
    self.assertTrue(np.isfinite(linear(Tensor.randn(1, 256)).realize().item()))
-    # the Q6 weight is repacked: 210-byte blocks padded to 212 (one block = 53 words)
-    self.assertEqual(linear.weight.uop.buf_uop.buffer.nbytes, 53*4)
-    self.assertEqual(linear.weight.dtype, dtypes.uint32)
+    self.assertEqual(linear.weight.uop.buf_uop.buffer.offset, 4)

  def test_q4_k_linear(self):
    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
@@ -852,6 +852,82 @@ class TestMultiFromUnrenderable(unittest.TestCase):
    np.testing.assert_equal(ll.numpy(), np.arange(100)+1)

@unittest.skipIf(not_support_multi_device(), "need multi")
+class TestMultiAssign(unittest.TestCase):
+  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
+
+  @needs_second_gpu
+  def setUp(self): pass
+
+  def test_multi_assign_realized(self):
+    out = Tensor.zeros(4).shard(self.device, 0).contiguous().realize()
+    ones = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
+    out.assign(ones).realize()
+    self.assertListEqual(out.tolist(), [1,1,1,1])
+
+  def test_multi_assign_unrealized(self):
+    out = Tensor.zeros(4).contiguous().realize().shard(self.device, 0)
+    ones = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
+    out.assign(ones).realize()
+    self.assertListEqual(out.tolist(), [1,1,1,1])
+
+  def test_multi_assign_both_unrealized(self):
+    out = Tensor.zeros(4).contiguous().realize().shard(self.device, 0)
+    ones = Tensor.ones(4).contiguous().realize().shard(self.device, 0)
+    out.assign(ones).realize()
+    self.assertListEqual(out.tolist(), [1,1,1,1])
+
+  def test_multi_assign_scalar(self):
+    out = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
+    out.assign(0).realize()
+    self.assertListEqual(out.tolist(), [0,0,0,0])
+
+  def test_multi_assign_const_like(self):
+    out = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
+    out.assign(out.const_like(7)).realize()
+    self.assertListEqual(out.tolist(), [7,7,7,7])
+
+  def test_multi_assign_piece(self):
+    out = Tensor.zeros(4,4).shard(self.device, 0).contiguous().realize()
+    ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
+    out[:, 2:3].assign(ones).realize()
+    self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
+
+  def test_multi_assign_piece_noncontig(self):
+    out = Tensor.zeros(4,4).contiguous().realize().shard(self.device, 0).realize()
+    ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
+    out[:, 2:3].assign(ones).realize()
+    self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
+
+  @unittest.expectedFailure
+  def test_multi_assign_piece_unrealized(self):
+    out = Tensor.zeros(4,4).contiguous().realize().shard(self.device, 0)
+    ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
+    out[:, 2:3].assign(ones).realize()
+    self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
+
+  def test_multi_assign_var_offset(self):
+    out = Tensor.zeros(4,4).contiguous().realize().shard(self.device, 0).realize()
+    ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
+    vi = Variable("i", 0, 3).bind(2)
+    out[:, vi:vi+1].assign(ones).realize()
+    self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
+
+  def test_multi_assign_var_offset_jit_none(self): self.test_multi_assign_var_offset_jit(None)
+  def test_multi_assign_var_offset_jit(self, shard_axis=0):
+    out = Tensor.zeros(4,6).contiguous().realize().shard(self.device, shard_axis).realize()
+    ones = Tensor.ones(4,1).shard(self.device, shard_axis).contiguous().realize()
+
+    @TinyJit
+    def f(out:Tensor, vi):
+      out[:, vi:vi+1].assign(ones).realize()
+      ones.assign(ones+1).realize()
+
+    vi = Variable("i", 0, 5)
+    for i in range(1,5):
+      GlobalCounters.reset()
+      f(out, vi.bind(i))
+    self.assertListEqual(out.tolist(), [[0,1,2,3,4,0]]*4)
+
@unittest.skipIf(not_support_multi_device(), "need multi")
 class TestMultiSetitem(unittest.TestCase):
  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
@@ -126,7 +126,7 @@ def do_devectorize(b:UOp):
  if not all(x.shape == b.shape or x.base.is_invalid for x in b.src): return None
  src = []
  for idx_c in itertools.product(*[[UOp.const(i) for i in range(x)] for x in b.shape]):
-    src.append(b.replace(src=tuple(x.base if x.base.is_invalid else x.index(*idx_c) for x in b.src)))
+    src.append(b.replace(dtype=None, src=tuple(x.base if x.base.is_invalid else x.index(*idx_c) for x in b.src)))
  return UOp.stack(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src)

 def do_stack_wmma(u:UOp):
@@ -391,15 +391,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
  sink = graph_rewrite(sink, pm_number_params, ctx=[num_params], name="number params with -1", walk=True)

  if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Output AST")
-  if SPEC:
-    import os
-    if os.environ.get("DBGTV"):
-      try: type_verify(sink, spec_program)
-      except RuntimeError:
-        from tinygrad.uop.render import print_uops
-        print_uops(list(sink.toposort()))
-        raise
-    else: type_verify(sink, spec_program)
+  if SPEC: type_verify(sink, spec_program)

  # return the rewritten sink
  return sink
@@ -443,7 +435,7 @@ def do_estimates(prg:UOp, sink:UOp, lin:UOp) -> UOp|None:
  return prg.replace(src=(sink.replace(arg=replace(sink.arg, estimates=Estimates.from_uops(lin.src, ignore_indexing=True))),)+prg.src[1:])

 def do_assemble(ctx:Renderer, prg:UOp, lin:UOp) -> UOp:
-  src = "\n".join(str(u.arg[0]) for u in lin.src)
+  src = "\n".join(str(u.arg) for u in lin.src)
  if DEBUG >= 4: print(src)
  binary = ctx.asm(prg, lin)
  return prg.replace(src=prg.src[:2]+(UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
@@ -140,8 +140,8 @@ def f2f_store(st, idx, val, fr:DType, to:DType):
 pm_long_decomp: PatternMatcher = PatternMatcher([
  # the decomp's own bottom-up rewrite can mint bare consts mid-flight: word splitting commits them at the long sibling's dtype
  (UPat(GroupOp.All, name='x'), lambda x: commit_weak_consts(x, next((s.dtype for s in x.src if s.dtype in l2i_dt), None))),
-  (UPat(GroupOp.Defines, tuple(l2i_dt.keys()), name="x"), lambda x:
-   UOp(x.op, arg=replace(x.arg, dtype=l2i_dt[x.dtype], size=None if x.arg.size is None else x.arg.size*2), tag=x.tag)),
+  (UPat(GroupOp.Defines, tuple(l2i_dt.keys()), src=(UPat.var("sz"),), name="x"), lambda x,sz:
+   UOp(x.op, src=(sz*2,), arg=replace(x.arg, dtype=l2i_dt[x.dtype]), tag=x.tag)),
  (UPat(Ops.INDEX, tuple(l2i_dt.keys()), name='x'), lambda x:
   reindex(x, x.tag[0]).replace(tag=None) if x.tag is not None else None),
  (UPat(Ops.STORE, src=(UPat.var('idx', tuple(l2i_dt.keys())), UPat.var('val')), name='st'), lambda st,idx,val:
@@ -73,7 +73,7 @@ def add_gpudims(ctx:Renderer, s:UOp):
  subs = {}
  for r in s_topo:
    # look for local INDEXes that are not used in the GLOBAL store, then add them as an INVALID
-    if r.op is Ops.STORE and len((idx := r.src[0]).src) and idx.src[0].addrspace == AddrSpace.GLOBAL:
+    if r.op is Ops.STORE and (idx := r.src[0]).src[0].addrspace == AddrSpace.GLOBAL:
      missing_locals = [all_ranges[rng] for rng in local_dims if all_ranges[rng] not in idx.ranges]
      if len(missing_locals):
        assert len(idx.src) == 2, "index has 2 sources"
@@ -1,8 +1,7 @@
 import itertools, functools
 from collections import defaultdict
-from dataclasses import replace
 from tinygrad.dtype import dtypes, AddrSpace, Invalid, DType
-from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp, graph_rewrite
+from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp, shape_to_shape_arg, graph_rewrite
 from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate, sym
 from tinygrad.helpers import getenv, IMAGE, OSX, ceildiv, is_image_shape
 from tinygrad.renderer import Renderer
@@ -86,8 +85,7 @@ def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
  if len(cands) == 0: return None
  # and tiebreak with indexing complexity (ie. number of nodes)
  h, w, cidx = cands[0] if len(cands) == 1 else min(cands, key=lambda cand: len(cand[2].index(1).simplify().backward_slice))
-  # the image dims are stored in the param's arg, the size stays the flat buffer len
-  buf = buf.replace(arg=replace(buf.arg, image=(h, w)))
+  buf = buf.replace(src=(shape_to_shape_arg((h, w, 4)),))
  shapes[buf.arg.slot] = (h, w)
  if valid.op is not Ops.CONST or valid.val is not True:
    return buf.index(cidx.src[1].valid(valid), cidx.src[0].valid(valid))
@@ -105,7 +103,7 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
  if getenv("DMC"): return sink

  # collect
-  memory: defaultdict[tuple[Ops, UOp, UOp|str, UOp, object], dict[int, list[UOp]]] = defaultdict(dict)
+  memory: defaultdict[tuple[Ops, UOp, UOp|str, UOp], dict[int, list[UOp]]] = defaultdict(dict)
  for u in sink.toposort():
    # TODO: this should handle images too, it's just memory coalescing
    if u.op in {Ops.LOAD, Ops.STORE}:
@@ -120,12 +118,11 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
      elif idx.op is Ops.CONST and idx.val is Invalid: root_src, arg = "INVALID", 0
      elif idx.op is Ops.CONST: root_src, arg = "CONST", idx.val
      else: root_src, arg = idx, 0
-      # loads/stores only coalesce with others carrying the same arg (e.g. the nontemporal flag)
-      memory[(u.op, buf, root_src, valid, u.arg)].setdefault(arg, []).append(u)
+      memory[(u.op, buf, root_src, valid)].setdefault(arg, []).append(u)

  # build replacements
  replacements = {}
-  for (op,buf,base,valid,ld_arg),offsets in memory.items():
+  for (op,buf,base,valid),offsets in memory.items():
    # allowed lengths (copied in)
    lengths = []
    must_divide = True
@@ -160,7 +157,7 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
          store = idx.store(UOp.stack(*datas) if len(datas) > 1 else datas[0])
          for i,g in enumerate(grp): replacements[offsets[g][0]] = store
        else:
-          ld = idx.load(arg=ld_arg)
+          ld = idx.load()
          for i,g in enumerate(grp):
            for oo in offsets[g]:
              replacements[oo] = ld.index(i) if len(grp) > 1 else ld
@@ -4,7 +4,7 @@ from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
 from tinygrad.renderer.isa import ISARenderer, Register, greg
 from tinygrad.dtype import dtypes

-PSEUDO_OPS = {Ops.CONST, Ops.CAST, Ops.BITCAST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP, Ops.STACK}
+PSEUDO_OPS = {Ops.CONST, Ops.CAST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP, Ops.STACK}

 class LinearScanRegallocContext:
  # returns the uop that defines the virtual register
@@ -120,6 +120,8 @@ class DTypes:
  uint32: Final[DType] = DType.new(6, 32, "unsigned int", 'I')
  int64: Final[DType] = DType.new(7, 64, "long", 'q')
  uint64: Final[DType] = DType.new(8, 64, "unsigned long", 'Q')
+  _uint128: Final[DType] = DType.new(8, 128, "uint128", None)
+  _uint256: Final[DType] = DType.new(8, 256, "uint256", None)
  weakfloat: Final[DType] = DType.new(9, 800, "weakfloat", None)
  fp8e4m3: Final[DType] = DType.new(10, 8, "float8_e4m3", None)
  fp8e5m2: Final[DType] = DType.new(11, 8, "float8_e5m2", None)
@@ -67,7 +67,7 @@ def jit_lower(linear:UOp, held_bufs:set[UOp], input_uops:list[UOp]) -> UOp:
  if VIZ: graph_rewrite(linear, PatternMatcher([]), name="View captured linear")

  # parametrize input buffers: map each input buffer UOp to a PARAM with the correct slot index
-  linear = linear.substitute({u: UOp.param(i, u.dtype, u.max_numel(), u.device) for i,u in enumerate(input_uops)}, walk=True)
+  linear = linear.substitute({u: UOp.param(i, u.dtype, u.shape, u.device) for i,u in enumerate(input_uops)}, walk=True)
  linear = memory_plan_rewrite(linear, held_bufs)
  linear = compile_linear(linear, beam=getenv("JITBEAM", BEAM.value))
  if JIT < 2: linear = graph_split_rewrite(linear, max_batch_size=JIT_BATCH_SIZE.value)
@@ -2,8 +2,8 @@ from __future__ import annotations
 from typing import cast, Iterator, Any, Sequence
 import random, itertools, math, weakref, array, decimal
 from dataclasses import dataclass, replace, field
-from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansipad, all_int, prod, flatten, Context, to_tuple, tqdm, dedup
-from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, HCQ2, PROFILE, ProfilePointEvent, cpu_events, perf_counter_us
+from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansipad, all_int, prod, flatten, Context, getenv, to_tuple, tqdm, dedup
+from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, perf_counter_us
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite, ProgramInfo
 from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
 from tinygrad.dtype import dtypes
@@ -305,17 +305,17 @@ pm_exec = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="validate", name="ast"),), name="call", allow_any_len=True), exec_validate),
 ])

-from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link, HCQ_RUNTIME_DEV # noqa: E402 # down here, hcq2 imports realize
+if getenv("HCQ2"): from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link, HCQ_RUNTIME_DEV # noqa: E402 # down here, hcq2 imports realize

 def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, profile:bool|None=None) -> UOp:
  if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
  if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
  linear = lower_and_compile(linear)
  linear = graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)
-  if HCQ2: linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile)
+  if getenv("HCQ2"): linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile)
  return linear

-def link_linear(linear:UOp, cache=True) -> UOp: return hcq_link(linear, cache=cache) if HCQ2 else linear
+def link_linear(linear:UOp, cache=True) -> UOp: return hcq_link(linear, cache=cache) if getenv("HCQ2") else linear

 def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:Sequence[UOp]=(), update_stats=True, jit=False, wait=False):
  inputs = list(input_uops)
@@ -240,7 +240,7 @@ TRANSCENDENTAL, NOLOCALS = ContextVar("TRANSCENDENTAL", 1), ContextVar("NOLOCALS
 SPLIT_REDUCEOP, NO_MEMORY_PLANNER, LRU = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("LRU", 1)
 RING, ALL2ALL, ALLREDUCE_CAST = ContextVar("RING", 1), ContextVar("ALL2ALL", 0), ContextVar("ALLREDUCE_CAST", 1)
 CACHELEVEL, IGNORE_BEAM_CACHE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0)
-VALIDATE_WITH_CPU, HCQ2 = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("HCQ2", 0)
+VALIDATE_WITH_CPU = ContextVar("VALIDATE_WITH_CPU", 0)
 # TODO: this is broken for some indexing
 DISABLE_FAST_IDIV = ContextVar("DISABLE_FAST_IDIV", 1)
 FUSE_OPTIM = ContextVar("FUSE_OPTIM", 0)
@@ -250,16 +250,16 @@ EMULATED_DTYPES = ContextVar("EMULATED_DTYPES", "")
 DEFAULT_FLOAT, DEFAULT_INT = ContextVar("DEFAULT_FLOAT", "float32"), ContextVar("DEFAULT_INT", "int32")
 CAPTURE_PROCESS_REPLAY = ContextVar("CAPTURE_PROCESS_REPLAY", 0)
 def _get_cpu_count() -> int:
-  # os.process_cpu_count is available in 3.13+, then try affinity, then fallback to cpu_count
-  count = (os.process_cpu_count() if hasattr(os, "process_cpu_count") else
-           len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count()) or 1
-  # limit with cgroup v2 (containers with --cpus=N)
+  # os.process_cpu_count (3.13+) respects cgroup limits
+  if hasattr(os, "process_cpu_count"): return max(1, os.process_cpu_count() or 1)
+  # cgroup v2 (containers with --cpus=N)
  try:
    with open("/sys/fs/cgroup/cpu.max") as f:
      quota, period = f.read().strip().split()
-      if quota != "max": count = min(count, max(1, int(quota) // int(period)))
+      if quota != "max": return max(1, int(quota) // int(period))
  except (FileNotFoundError, ValueError, ZeroDivisionError): pass
-  return count
+  # fall back to affinity (respects taskset but not cgroup quota)
+  return max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1))
 NUM_CPU_THREADS = ContextVar("NUM_CPU_THREADS", _get_cpu_count())
 NULL_ALLOW_COPYOUT = ContextVar("NULL_ALLOW_COPYOUT", 0)
 # VIZ implies PROFILE, but you can run PROFILE without VIZ
@@ -145,7 +145,6 @@ def main():
  parser.add_argument("--serve", nargs='?', type=int, const=8000, metavar="PORT", help="Run OpenAI compatible API (optional port, default 8000)")
  parser.add_argument("--warmup", action="store_true", help="warmup the JIT")
  parser.add_argument("--benchmark", nargs='?', type=int, const=20, metavar="COUNT", help="Benchmark tok/s (optional count, default 20)")
-  parser.add_argument("--no_chat_template", action="store_true", help="Don't use the model's chat template, always use the fallback template")
  args = parser.parse_args()

  # load the model
@@ -161,7 +160,7 @@ def main():

  # use the model's chat template if jinja2 is available (enables model-specific formatting)
  template: jinja2.Template|FallbackTemplate = FallbackTemplate(tok)
-  if not args.no_chat_template and (ct := kv.get('tokenizer.chat_template')) is not None:
+  if (ct := kv.get('tokenizer.chat_template')) is not None:
    try:
      import jinja2
      env = jinja2.Environment()
@@ -18,9 +18,7 @@ _GGML_NATIVE = {0: dtypes.float32, 1: dtypes.float16, 24: dtypes.int8, 25: dtype

 # quant types {ggml_type: (number of elements, number of bytes)}
 _GGML_QUANT = {2:(32,18), 3:(32,20), 6:(32,22), 7:(32,24), 8:(32,34),
-               10:(256,84), 11:(256,110), 12:(256,144), 13:(256,176), 14:(256,210),
-               16:(256,66), 17:(256,74), 18:(256,98), 19:(256,50), 20:(32,18), 21:(256,110), 22:(256,82), 23:(256,136),
-               29:(256,56), 39:(32,17), 41:(128,18)}
+               12:(256,144), 13:(256,176), 14:(256,210), 18:(256,98), 21:(256,110), 22:(256,82), 23:(256,136), 39:(32,17), 41:(128,18)}

 def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
  """
@@ -29,9 +27,8 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
  Supported native types: float32 (id: 0), float16 (id: 1), int8 (id: 24),
  int16 (id: 25), int32 (id: 26), int64 (id: 27), float64 (id: 28), bfloat16 (id: 30)
  Supported quantized types: Q4_0 (id: 2), Q4_1 (id: 3), Q5_0 (id: 6),
-  Q5_1 (id: 7), Q8_0 (id: 8), Q2_K (id: 10), Q3_K (id: 11), Q4_K (id: 12), Q5_K (id: 13),
-  Q6_K (id: 14), IQ2_XXS (id: 16), IQ2_XS (id: 17), IQ3_XXS (id: 18), IQ1_S (id: 19),
-  IQ4_NL (id: 20), IQ3_S (id: 21), IQ2_S (id: 22), IQ4_XS (id: 23), IQ1_M (id: 29), MXFP4 (id: 39), Q1_0 (id: 41)
+  Q5_1 (id: 7), Q8_0 (id: 8), Q4_K (id: 12), Q5_K (id: 13),
+  Q6_K (id: 14), IQ3_XXS (id: 18), IQ3_S (id: 21), IQ2_S (id: 22), IQ4_XS (id: 23), MXFP4 (id: 39), Q1_0 (id: 41)
  """
  # https://github.com/ggerganov/ggml/blob/323951f1bdcdfbd5b5ff3a9a7c3770e63b1a560e/include/ggml.h#L356

@@ -57,19 +54,6 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
      q = q_to_uint8(blocks[:,qh_off+4:], 4).bitcast(dtypes.int8) + qh * 16
      return q * d + (blocks[:,2:4].bitcast(dtypes.float16).cast(dtypes.float32) if ggml_type == 7 else -16 * d)
    if ggml_type == 8: return blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32) * blocks[:,2:].bitcast(dtypes.int8)
-    # Q2_K: 256 elements per 84-byte block (scales:16, qs:64, d:2, dmin:2)
-    if ggml_type == 10:
-      d, dmin = (blocks[:,i:i+2].bitcast(dtypes.float16).cast(dtypes.float32).unsqueeze(-1) for i in [80, 82])
-      sc = blocks[:, :16]
-      q = q_to_uint8(blocks[:, 16:80].reshape((-1, 2, 32)), 2).reshape((-1, 16, 16))
-      return (d * sc.bitwise_and(0xF).unsqueeze(-1) * q - dmin * sc.rshift(4).unsqueeze(-1)).flatten(-2)
-    # Q3_K: 256 elements per 110-byte block (hmask:32, qs:64, scales:12, d:2)
-    if ggml_type == 11:
-      d = blocks[:,-2:].bitcast(dtypes.float16).cast(dtypes.float32).unsqueeze(-1)
-      sc = q_to_uint8(blocks[:,96:104], 4).bitwise_or(q_to_uint8(blocks[:,104:108], 2).lshift(4)).bitcast(dtypes.int8) - 32
-      q = q_to_uint8(blocks[:,32:96].reshape((-1, 2, 32)), 2).reshape((-1, 16, 16))
-      qh = q_to_uint8(blocks[:,:32], 1).reshape((-1, 16, 16))
-      return (d * sc.unsqueeze(-1) * (q.bitcast(dtypes.int8) - qh.bitwise_xor(1).lshift(2).bitcast(dtypes.int8))).flatten(-2)
     # Q4_K: 256 elements per 144-byte block (d:2, dmin:2, scales:12, qs:128)
     # Q5_K: 256 elements per 176-byte block (d:2, dmin:2, scales:12, qh:32, qs:128)
    if ggml_type in (12, 13):
@@ -95,41 +79,6 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
      signs = (q_to_uint8(even_signs[sign_idx].reshape((-1, 32, 1)), 1) == 0).where(1.0, -1.0).reshape((-1, 8, 4, 8))
      grid = _ggml_iq_grid(t.device, _ggml.iq3xxs_grid, (256, 4))[blocks[:, 2:66]].reshape((-1, 8, 4, 8))
      return (db * grid * signs).flatten(-3)
-    # IQ2_XXS: 256 elements per 66-byte block (d:2, qs:64). 8 groups of 32: 4 grid bytes + packed signs/scale.
-    if ggml_type == 16:
-      d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1))
-      qs_u32 = blocks[:, 2:].bitcast(dtypes.uint32).reshape((-1, 8, 2))
-      db = d * (qs_u32[:, :, 1].rshift(28).cast(dtypes.float32) + 0.5).reshape((-1, 8, 1, 1)) * 0.25
-      sign_idx = qs_u32[:, :, 1].unsqueeze(-1).rshift(Tensor.const((0, 7, 14, 21), dtypes.uint32))
-      sign_idx = sign_idx.bitwise_and(0x7F).reshape((-1, 32)).cast(dtypes.int32)
-      even_signs = Tensor([i | (0x80 if i.bit_count() % 2 else 0) for i in range(128)], dtype=dtypes.uint8, device=t.device)
-      signs = (q_to_uint8(even_signs[sign_idx].reshape((-1, 32, 1)), 1) == 0).where(1.0, -1.0).reshape((-1, 8, 4, 8))
-      grid = _ggml_iq_grid(t.device, _ggml.iq2xxs_grid, (256, 8))[blocks[:, 2:].reshape((-1, 8, 8))[:, :, :4]].reshape((-1, 8, 4, 8))
-      return (db * grid * signs).flatten(-3)
-    # IQ2_XS: 256 elements per 74-byte block (d:2, qs:64 as uint16, scales:8)
-    if ggml_type == 17:
-      d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1))
-      db = d * (q_to_uint8(blocks[:, 66:74].reshape((-1, 8, 1)), 4).reshape((-1, 16)).cast(dtypes.float32) + 0.5).reshape((-1, 16, 1, 1)) * 0.25
-      qs = blocks[:, 2:66].bitcast(dtypes.uint16)
-      sign_idx = qs.rshift(9).cast(dtypes.int32)
-      even_signs = Tensor([i | (0x80 if i.bit_count() % 2 else 0) for i in range(128)], dtype=dtypes.uint8, device=t.device)
-      signs = (q_to_uint8(even_signs[sign_idx].reshape((-1, 32, 1)), 1) == 0).where(1.0, -1.0).reshape((-1, 16, 2, 8))
-      grid = _ggml_iq_grid(t.device, _ggml.iq2xs_grid, (512, 8))[qs.bitwise_and(511)].reshape((-1, 16, 2, 8))
-      return (db * grid * signs).flatten(-3)
-    # IQ1_S: 256 elements per 50-byte block (d:2, qs:32, qh:16). grid bytes are int8 {-1,0,1}.
-    if ggml_type == 19:
-      d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1))
-      qh = blocks[:, 34:50].bitcast(dtypes.uint16)
-      dl = d * (qh.rshift(12).bitwise_and(7).cast(dtypes.float32) * 2 + 1).reshape((-1, 8, 1, 1))
-      delta = (qh.bitwise_and(0x8000) == 0).where(0.125, -0.125).reshape((-1, 8, 1, 1))
-      qh_hi = qh.unsqueeze(-1).rshift(Tensor.const((0, 3, 6, 9), dtypes.uint16)).bitwise_and(7).lshift(8)
-      q = blocks[:, 2:34].cast(dtypes.uint16) + qh_hi.reshape((-1, 32))
-      grid = _ggml_iq_grid(t.device, _ggml.iq1s_grid, (2048, 8))[q].reshape((-1, 8, 4, 8))
-      grid = (grid > 127).where(grid - 256, grid)
-      return (dl * (grid + delta)).flatten(-3)
-    if ggml_type == 20:
-      d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32)
-      return d * Tensor(list(_ggml.kvalues_iq4nl), dtype=dtypes.float32, device=t.device)[q_to_uint8(blocks[:, 2:], 4)]
    if ggml_type == 21:
      d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1))
      scales = (1 + 2 * q_to_uint8(blocks[:, 106:110].reshape((-1, 4, 1)), 4).reshape((-1, 8))).cast(dtypes.float32).reshape((-1, 8, 1, 1))
@@ -153,20 +102,6 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
      scales = (scales_l.bitwise_or(scales_h.lshift(4)).bitcast(dtypes.int8) - 32).cast(dtypes.float32).reshape((-1, 8, 1))
      q = (qs:=blocks[:, 8:].reshape((-1, 8, 16))).bitwise_and(0xF).cat(qs.rshift(4), dim=2)
      return (d * scales * iq4_xs_lut[q]).flatten(-2)
-    # IQ1_M: 256 elements per 56-byte block (qs:32, qh:16, scales:8). f16 scale packed in high nibbles.
-    if ggml_type == 29:
-      sc16 = blocks[:, 48:56].bitcast(dtypes.uint16)
-      d = sc16.bitwise_and(0xF000).rshift(Tensor.const((12, 8, 4, 0), dtypes.uint16))
-      d = d[:, 0:1].bitwise_or(d[:, 1:2]).bitwise_or(d[:, 2:3]).bitwise_or(d[:, 3:4])
-      d = d.bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1, 1))
-      scales = sc16.unsqueeze(-1).rshift(Tensor.const((0, 3, 6, 9), dtypes.uint16)).bitwise_and(7)
-      dl = d * (scales.cast(dtypes.float32) * 2 + 1).reshape((-1, 8, 2, 1, 1))
-      qh_n = Tensor.stack(blocks[:, 32:48].bitwise_and(0x0F), blocks[:, 32:48].rshift(4), dim=-1).reshape((-1, 32))
-      q = blocks[:, :32].cast(dtypes.uint16) + qh_n.bitwise_and(7).cast(dtypes.uint16).lshift(8)
-      delta = (qh_n.bitwise_and(0x08) == 0).where(0.125, -0.125).reshape((-1, 8, 2, 2, 1))
-      grid = _ggml_iq_grid(t.device, _ggml.iq1s_grid, (2048, 8))[q].reshape((-1, 8, 2, 2, 8))
-      grid = (grid > 127).where(grid - 256, grid)
-      return (dl * (grid + delta)).flatten(-4)
    if ggml_type == 39:
      e = blocks[:, 0].cast(dtypes.uint32)
      small_bits = Tensor([0x00200000, 0x00400000], dtype=dtypes.uint32, device=t.device)[e.clip(0, 1).cast(dtypes.int32)] # e = 0 or e = 1 case
@@ -7,13 +7,12 @@ from tinygrad.dtype import AddrSpace, dtypes
 from tinygrad.helpers import prod
 from tinygrad.uop.ops import AxisType, KernelInfo, Ops, resolve

-BLOCK_M, BLOCK_N, WARP_SIZE = 32, 32, 32
+BLOCK_M, BLOCK_N, DECODE_HEAD_TILE, WARP_SIZE = 32, 32, 8, 32
 WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
 WAVES_M, WAVES_N, LANES_PER_WAVE_M, LANES_PER_WAVE_N = 2, 2, 2, 16
 WMMA_ACC, THREADS_PER_BLOCK = WMMA_M // LANES_PER_WAVE_M, WARP_SIZE * WAVES_M * WAVES_N
 LDS_PAD, WMMA_ARG, LOG2E = 4, ((WMMA_M, WMMA_N, WMMA_K), 'AMD', 32), math.log2(math.e)
 Q4_K, Q5_K, Q6_K, IQ4_XS, GGML_BLOCK_SIZE, Q8_GROUP_SIZE, Q4_WORDS, Q5_WORDS, Q6_BYTES, IQ4_WORDS = 12, 13, 14, 23, 256, 32, 36, 44, 210, 34
-Q6_PADDED, Q6_WORDS = 212, 53  # the 210-byte Q6 blocks are padded to 212 bytes so they are word-addressable
 QUANT_SIZES = {Q4_K: Q4_WORDS*4, Q5_K: Q5_WORDS*4, Q6_K: Q6_BYTES, IQ4_XS: IQ4_WORDS*4}  # bytes per 256-weight block

 def kernel_var(x:UOp) -> UOp:
@@ -55,43 +54,21 @@ class Linear(nn.Linear):
    self.in_features, self.out_features = in_features, out_features
  def set_quantized(self, decoded:Tensor):
    packed_sizes = {decoded.numel() // 256 * type_size:typ for typ,type_size in QUANT_SIZES.items()}
-    graph = decoded.uop.toposort()
-    raw = next((u for u in graph if u.op is Ops.SHRINK and u.dtype == dtypes.uint8 and prod(u.shape) in packed_sizes), None)
+    raw = next((u for u in decoded.uop.toposort() if u.op is Ops.SHRINK and u.dtype == dtypes.uint8 and prod(u.shape) in packed_sizes), None)
    if raw is None: return
-    ggml_type = packed_sizes[prod(raw.shape)]
-    # the packed byte rate alone can't distinguish same-rate formats (Q4_0 vs Q4_K, Q5_0 vs Q5_K, MXFP4 vs IQ4_XS).
-    # the supported formats are 256-wide superblocks: their decode views the packed bytes at the superblock width
-    # (ggml_data_to_tensor reshapes to (-1, QUANT_SIZES[type])), while same-rate 32-wide formats reshape to 17-22
-    if not any(u.op is Ops.RESHAPE and u.shape[-1:] == (QUANT_SIZES[ggml_type],) for u in graph): return
    raw_offset = raw.contiguous_view_offset()
    assert raw_offset is not None and raw_offset % 4 == 0 and raw.buf_uop.dtype == dtypes.uint8
-    self.ggml_type = ggml_type
+    self.ggml_type = packed_sizes[prod(raw.shape)]
    # store a typed buffer view: a lazy BITCAST is decomposed into byte-combining ALU before custom-kernel
    # scheduling and would copy the entire packed weight on every JIT graph
-    if self.ggml_type == Q6_K:
-      # Q6 blocks are 210 bytes, so consecutive blocks are only 2-byte aligned. pad each block to 212 bytes
-      # (a one-time copy at load) so the kernel can do all its reads as aligned u32 words
-      nbytes, nblocks = raw.max_numel(), raw.max_numel() // Q6_BYTES
-      byte_view = Tensor(UOp.from_buffer(cast(Buffer, raw.buf_uop.buffer).view(nbytes, dtypes.uint8, raw_offset)))
-      padded = byte_view.reshape((nblocks, Q6_BYTES)).pad_to((nblocks, Q6_PADDED)).contiguous().realize()
-      self.weight = Tensor(UOp.from_buffer(cast(Buffer, padded.uop.buf_uop.buffer).view(nblocks * Q6_WORDS, dtypes.uint32, 0)))
-    else:
-      self.weight = Tensor(UOp.from_buffer(cast(Buffer, raw.buf_uop.buffer)
-        .view(raw.max_numel() * raw.dtype.itemsize // dtypes.uint32.itemsize, dtypes.uint32, raw_offset)))
+    packed_dtype = dtypes.uint8 if self.ggml_type == Q6_K else dtypes.uint32
+    self.weight = Tensor(UOp.from_buffer(cast(Buffer, raw.buf_uop.buffer)
+      .view(raw.max_numel() * raw.dtype.itemsize // packed_dtype.itemsize, packed_dtype, raw_offset)))
  def __call__(self, x:Tensor) -> Tensor:
    supported = self.use_custom_quant and amd_custom_kernels_supported(self.weight.device)
    if self.ggml_type is None and supported:
      self.set_quantized(self.weight)
-      if self.ggml_type is None:
-        # tiny dense fp16 matmul (e.g. the ssm beta/alpha head rows): single fp16 gemv kernel instead of a
-        # generic matmul schedule, and realize the densely packed weight once if it is still a lazy ggml view
-        if self.weight.dtype in (dtypes.half, dtypes.float, dtypes.bfloat16) and self.out_features <= 2048 \
-          and self.in_features % (WARP_SIZE*4) == 0:
-          numel, max_shape = x.numel(), x.max_shape
-          if isinstance(numel, int) or prod(max_shape) // self.in_features <= 32:
-            out = f16_gemv(self, x if isinstance(numel, int) else x.pad_to(max_shape))
-            return out if isinstance(numel, int) else out.shrink(tuple((0, s) for s in (*x.shape[:-1], self.out_features)))
-        self.use_custom_quant = supported = False  # not a supported quant format
+      if self.ggml_type is None: self.use_custom_quant = supported = False  # not a supported quant format
    if self.ggml_type in (Q4_K, Q5_K, Q6_K, IQ4_XS) and supported:
      if isinstance(x.numel(), int): return q8_linear(self, x)
      # symbolic token count: pad to the max chunk size so the kernels see static shapes, garbage rows are sliced off
@@ -100,41 +77,32 @@ class Linear(nn.Linear):
    return super().__call__(x)

 def _amd_dp4a(a:UOp, b:UOp, c:UOp) -> UOp:
-  # int8 4-wide dot, widened to scalar multiply-adds (2% decode slower than the sudot4 builtin, but portable)
-  for i in range(4):
-    av = ((a >> (8*i)) & 255).cast(dtypes.uint8).bitcast(dtypes.int8).int()
-    bv = ((b >> (8*i)) & 255).cast(dtypes.uint8).bitcast(dtypes.int8).int()
-    c = c + av*bv
-  return c
+  return UOp(Ops.CUSTOMI, src=(a.int(), b.int(), c), arg=("__builtin_amdgcn_sudot4(true, {}, true, {}, {}, false)", dtypes.int32))

 def _amd_byte_perm(a:UOp, b:UOp, selectors:UOp) -> UOp:
  return UOp(Ops.CUSTOMI, src=tuple(x.cast(dtypes.uint32) for x in (a, b, selectors)), arg=("__builtin_amdgcn_perm({}, {}, {})", dtypes.uint32))

 def _amd_load(ptr:UOp, lanes:int|None=None) -> UOp:
  assert ptr.op is Ops.INDEX
-  # nontemporal scalar load: streamed weights must not evict the activations/KV cache from L2
-  if lanes is None: return ptr.load(arg="nontemporal")
+  if lanes is None: return UOp(Ops.CUSTOMI, src=(ptr,), arg=("__builtin_nontemporal_load({0})", ptr.dtype))
  buf, coords = ptr.src[0], ptr.src[1:]
  idx = sum((coord*math.prod(buf.shape[i+1:]) for i,coord in enumerate(coords)), UOp.const(0))
-  return UOp(Ops.SHRINK, src=(buf.flatten(), idx, UOp.const(lanes))).load()
+  return UOp(Ops.SHRINK, src=(buf.flatten(), idx, UOp.const(lanes))).load(dtype=ptr.dtype)

 def _load_byte(raw:UOp, base:UOp, offset:UOp) -> UOp: return (raw[base + offset//4] >> ((offset&3)*8).cast(dtypes.uint32)) & 255
 def _half(value:UOp) -> UOp: return value.cast(dtypes.uint16).bitcast(dtypes.float16).float()

 def _iq4_bytes(packed:UOp, shift:int) -> UOp:
-  # the non-linear iq4nl table as a byte lookup: 3 byte_perms beat any arithmetic/select-tree form (~60% decode)
  selectors = (packed >> shift) & 0x0f0f0f0f
  low = _amd_byte_perm(UOp.const(0xf6eaddcf, dtypes.uint32), UOp.const(0xbfad9881, dtypes.uint32), selectors)
  high = _amd_byte_perm(UOp.const(0x71594535, dtypes.uint32), UOp.const(0x26190d01, dtypes.uint32), selectors & 0x07070707)
  return _amd_byte_perm(high, low, 0x03020100 | ((selectors & 0x08080808) >> 1))

 def _q5_scales(raw:UOp, base:UOp, subgroup:UOp) -> tuple[UOp, UOp, UOp, UOp]:
-  # scales/mins (6-bit each) live in block bytes 4-15: three words total, same for the whole super-block's lanes
-  w1, w2, w3 = _amd_load(raw[base+1]), _amd_load(raw[base+2]), _amd_load(raw[base+3])
-  sb = (subgroup & 3) * 8  # byte within word
-  byte1, byte2, byte3 = (w1 >> sb) & 255, (w2 >> sb) & 255, (w3 >> sb) & 255
-  scale = (subgroup < 4).where(byte1 & 63, (byte3 & 15) | ((byte1 >> 6) << 4))
-  minimum = (subgroup < 4).where(byte2 & 63, (byte3 >> 4) | ((byte2 >> 6) << 4))
+  scale = (subgroup < 4).where(_load_byte(raw, base, 4 + subgroup) & 63,
+    (_load_byte(raw, base, 8 + subgroup) & 15) | ((_load_byte(raw, base, subgroup) >> 6) << 4))
+  minimum = (subgroup < 4).where(_load_byte(raw, base, 8 + subgroup) & 63,
+    (_load_byte(raw, base, 8 + subgroup) >> 4) | ((_load_byte(raw, base, 4 + subgroup) >> 6) << 4))
  d, dmin = (raw[base] & 0xffff).cast(dtypes.uint16), (raw[base] >> 16).cast(dtypes.uint16)
  return _half(d), _half(dmin), scale.float(), minimum.float()

@@ -150,66 +118,52 @@ def iq4_half_lut(device:str) -> Tensor:
                dtype=dtypes.float16, device=device).bitcast(dtypes.uint32).contiguous()

@functools.cache
-def _q8_quantize_kernel(q:UOp, scale:UOp, xsum:UOp, x:UOp, tokens:int, in_features:int) -> UOp:
+def _q8_quantize_kernel(q:UOp, scale:UOp, x:UOp, tokens:int, in_features:int) -> UOp:
  groups = in_features//Q8_GROUP_SIZE
-  token_group, lane = UOp.range(tokens*groups, 0, axis_type=AxisType.GLOBAL), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  token_group, lane = UOp.range(tokens*groups, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
  token, group = token_group//groups, token_group%groups
  x = x.reshape(tokens, groups, 32)
  group_scale = (warp_reduce(x[token, group, lane].float().abs(), maximum=True, full_wave=True) / 127).maximum(1e-8)
  word_lane = lane.minimum(7)
  xs = tuple(x[token, group, word_lane*4+i].float() for i in range(4))
-  qs = tuple((v/group_scale).round().clip(-127, 127).cast(dtypes.int8) for v in xs)
-  word = sum((v.cast(dtypes.uint8).cast(dtypes.uint32) << (i*8) for i, v in enumerate(qs)), UOp.const(0, dtypes.uint32))
-  # per-16 sums of the quantized values (lanes 0-3 / 4-7): Q4_K/Q5_K need the 32-sum, Q6_K the 16-sums
-  part = (lane < 8).where(sum((v.cast(dtypes.int32) for v in qs), UOp.const(0, dtypes.int32)), UOp.const(0, dtypes.int32))
-  gsum = [warp_reduce(((lane & 4).eq(h*4)).where(part, UOp.const(0, dtypes.int32)), full_wave=True) for h in range(2)]
-  store_half = (lane & 4) >> 2
-  stores = (q[token, group, lane.valid(lane < 8)].store(word),
-            UOp.group(scale[token, group.valid(lane.eq(0))].store(group_scale),
-                      xsum[token, group, store_half.valid(lane.eq(0) | lane.eq(4))].store(
-                        store_half.eq(0).where(gsum[0].float(), gsum[1].float()))))
+  word = sum(((v/group_scale).round().clip(-127, 127).cast(dtypes.int8).cast(dtypes.uint8).cast(dtypes.uint32) << (i*8)
+              for i,v in enumerate(xs)), UOp.const(0, dtypes.uint32))
+  stores = (q[token, group, lane.valid(lane < 8)].store(word), scale[token, group.valid(lane.eq(0))].store(group_scale))
  return UOp.group(*stores).end(token_group, lane).sink(arg=KernelInfo(name="q8_quantize", opts_to_apply=()))

-def q8_quantize(x:Tensor, tokens:int, in_features:int) -> tuple[Tensor, Tensor, Tensor]:
+def q8_quantize(x:Tensor, tokens:int, in_features:int) -> tuple[Tensor, Tensor]:
  groups = in_features//Q8_GROUP_SIZE
  q = Tensor.empty(tokens, groups, 8, dtype=dtypes.uint32, device=x.device)
  scale = Tensor.empty(tokens, groups, dtype=dtypes.float32, device=x.device)
-  xsum = Tensor.empty(tokens, groups, 2, dtype=dtypes.float32, device=x.device)
-  q, scale, xsum = Tensor.custom_kernel(q, scale, xsum, x, fxn=functools.partial(_q8_quantize_kernel, tokens=tokens, in_features=in_features))[:3]
-  return q, scale, xsum
+  q, scale = Tensor.custom_kernel(q, scale, x, fxn=functools.partial(_q8_quantize_kernel, tokens=tokens, in_features=in_features))[:2]
+  return q, scale

 def _decode_linear(out:UOp, out_features:int, group_count:int, group_dot, name:str) -> UOp:
-  chunks = out.shape[2]
-  # two-dim global grid instead of one flat grid: no div/mods needed to decompose the gid
-  token_output = UOp.range(out.shape[0]*out_features, 0, axis_type=AxisType.GLOBAL)
-  chunk, lane = UOp.range(chunks, 1, axis_type=AxisType.GLOBAL), UOp.range(32, 2, axis_type=AxisType.LOCAL)
-  token, output = token_output // out_features, token_output % out_features
-  group = (lane+chunk*32).minimum(group_count-1)
-  value = group_dot(token, output, group) if chunks*32 == group_count else \
-    (lane+chunk*32 < group_count).where(group_dot(token, output, group), UOp.const(0, dtypes.float32))
+  chunks = (group_count+31)//32
+  token_output_chunk, lane = UOp.range(out.shape[0]*out_features*chunks, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  token, output, chunk = token_output_chunk // (out_features*chunks), (token_output_chunk//chunks) % out_features, token_output_chunk % chunks
+  group = lane+chunk*32
+  value = group_dot(token, output, group) if group_count % 32 == 0 else \
+    (group < group_count).where(group_dot(token, output, group.minimum(group_count-1)), UOp.const(0, dtypes.float32))
  total = warp_reduce(value, full_wave=True)
-  return out[token, output, chunk.valid(lane.eq(0))].store(total.cast(out.dtype)).end(token_output, chunk, lane).sink(
+  return out[token, output, chunk.valid(lane.eq(0))].store(total.cast(out.dtype)).end(token_output_chunk, lane).sink(
    arg=KernelInfo(name=name, opts_to_apply=()))

@functools.cache
-def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, xs:UOp, out_features:int, in_features:int, ggml_type:int) -> UOp:
+def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, out_features:int, in_features:int, ggml_type:int) -> UOp:
  group_count = in_features // Q8_GROUP_SIZE
  def group_dot(token:UOp, output:UOp, group:UOp) -> UOp:
    block, subgroup = group // 8, group % 8
    xwords = _amd_load(xq[token, group, 0], 8)
    if ggml_type in (Q4_K, Q5_K):
      base = (output * in_features//GGML_BLOCK_SIZE + block) * (Q4_WORDS if ggml_type == Q4_K else Q5_WORDS)
-      qs_base, dot = base + (4 if ggml_type == Q4_K else 12) + (subgroup//2)*8, UOp.const(0, dtypes.int32)
-      # vectorize the 8 packed-weight words and (for Q5_K) the 32-byte high-bit bitmap
-      qs_pair = (_amd_load(raw[qs_base], 4), _amd_load(raw[qs_base+4], 4))
-      if ggml_type == Q5_K: qh_pair = (_amd_load(raw[base+4], 4), _amd_load(raw[base+8], 4))
+      qs_base, dot, qsum = base + (4 if ggml_type == Q4_K else 12) + (subgroup//2)*8, UOp.const(0, dtypes.int32), UOp.const(0, dtypes.int32)
      for word_idx in range(8):
-        word = (qs_pair[word_idx//4][word_idx%4] >> ((subgroup&1)*4).cast(dtypes.uint32)) & 0x0f0f0f0f
-        if ggml_type == Q5_K: word |= ((qh_pair[word_idx//4][word_idx%4] >> subgroup.cast(dtypes.uint32)) & 0x01010101) << 4
-        dot = _amd_dp4a(word, xwords[word_idx], dot)
+        word = (raw[qs_base+word_idx] >> ((subgroup&1)*4).cast(dtypes.uint32)) & 0x0f0f0f0f
+        if ggml_type == Q5_K: word |= ((raw[base+4+word_idx] >> subgroup.cast(dtypes.uint32)) & 0x01010101) << 4
+        dot, qsum = _amd_dp4a(word, xwords[word_idx], dot), _amd_dp4a(UOp.const(0x01010101, dtypes.uint32), xwords[word_idx], qsum)
      d, dmin, scale, minimum = _q5_scales(raw, base, subgroup)
-      gsum = xs[token, group, 0].load() + xs[token, group, 1].load()
-      return (dot.float()*d*scale - gsum*dmin*minimum) * xd[token, group]
+      return (dot.float()*d*scale - qsum.float()*dmin*minimum) * xd[token, group]
    if ggml_type == IQ4_XS:
      base = (output * in_features//GGML_BLOCK_SIZE + block) * IQ4_WORDS
      dot = UOp.const(0, dtypes.int32)
@@ -218,57 +172,45 @@ def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, xs:UOp, out_features:
        dot = _amd_dp4a(_iq4_bytes(packed, 4*(word_idx//4)), xwords[word_idx], dot)
      d, scale = _iq4_scales(raw, base, subgroup)
      return dot.float() * xd[token, group] * d * scale
-    # the packed rows were padded to 212 bytes (53 words) per 256-block in set_quantized: everything is word-aligned
-    base = (output*in_features//GGML_BLOCK_SIZE+block)*Q6_WORDS
-    # the subgroup's 8 ql words and 8 qh words are contiguous: two 16-byte vector loads each
-    lows = tuple(_amd_load(raw[base + (subgroup//4)*16 + (subgroup%2)*8 + half*4], 4) for half in range(2))
-    highs = tuple(_amd_load(raw[base + 32 + (subgroup//4)*8 + half*4], 4) for half in range(2))
+    base = (output*in_features//GGML_BLOCK_SIZE+block)*Q6_BYTES
    dots = [UOp.const(0, dtypes.int32)] * 2
    for word_idx in range(8):
-      within = (subgroup*32 + word_idx*4)%128
-      low = lows[word_idx//4][word_idx%4] >> ((within//64)*4).cast(dtypes.uint32)
-      high = highs[word_idx//4][word_idx%4] >> ((within//32)*2).cast(dtypes.uint32)
-      # 4 values per word: (low nibble) | (2 high bits << 4). values stay positive, so the int8-bitcast/-32 of the
-      # naive dequant is skipped and the -32 offset is applied later via the per-16 sums of the quantized inputs
-      word = (low & 0x0f0f0f0f) | ((high & 0x03030303) << 4)
+      pos, within = subgroup*32 + word_idx*4, (subgroup*32 + word_idx*4)%128
+      low = _amd_load(raw[base + (pos//128)*64 + within%64], 4) >> ((within//64)*4).cast(dtypes.uint8)
+      high = _amd_load(raw[base + 128 + (pos//128)*32 + within%32], 4) >> ((within//32)*2).cast(dtypes.uint8)
+      quant = ((low & 15) | ((high & 3) << 4)).bitcast(dtypes.int8) - 32
+      word = sum((quant[i].cast(dtypes.uint8).cast(dtypes.uint32) << (i*8) for i in range(4)), UOp.const(0, dtypes.uint32))
      dots[word_idx//4] = _amd_dp4a(word, xwords[word_idx], dots[word_idx//4])
-    scales = [((raw[base + 48 + (subgroup*2+i)//4] >> (((subgroup*2+i)%4)*8).cast(dtypes.uint32)) & 255)
-              .cast(dtypes.uint8).bitcast(dtypes.int8).float() for i in range(2)]
-    gsum = [xs[token, group, i].load() * 32 for i in range(2)]
-    return ((dots[0].float() - gsum[0])*scales[0] + (dots[1].float() - gsum[1])*scales[1]) * xd[token, group] * _half(raw[base+52] & 0xffff)
+    scales = [raw[base + 192 + subgroup*2+i].cast(dtypes.uint8).bitcast(dtypes.int8).float() for i in range(2)]
+    dbits = raw[base+208].cast(dtypes.uint16) | (raw[base+209].cast(dtypes.uint16) << 8)
+    return (dots[0].float()*scales[0] + dots[1].float()*scales[1]) * xd[token, group] * _half(dbits)
  names = {Q4_K: "linear_q4_k", Q5_K: "linear_q5_k", IQ4_XS: "linear_iq4_xs", Q6_K: "linear_q6"}
  return _decode_linear(out, out_features, group_count, group_dot, names[ggml_type])

 def _wmma_layout(out:UOp, out_features:int, token_tile:int, output_tiles:int):
  output_waves = 2 if out_features % (32*output_tiles) == 0 else 1
  token_block, output_block = UOp.range(out.shape[0]//token_tile, 0), UOp.range(out_features//(16*output_tiles*output_waves), 1)
-  # lane is a hardware WARP range (like the flash kernel): the fragment math stays visible without being
-  # range-split into nested loops, which would scramble the WMMA fragment layout
-  lane, wave = UOp.range(WARP_SIZE, -1, axis_type=AxisType.WARP), UOp.range(output_waves, 3, axis_type=AxisType.LOCAL)
-  col, half = lane % 16, lane // 16
+  lane, wave = UOp.range(WARP_SIZE, 2, axis_type=AxisType.LOCAL), UOp.range(output_waves, 3, axis_type=AxisType.LOCAL)
+  hw_lane = UOp(Ops.CUSTOM, src=(lane.int(),), arg=("__builtin_amdgcn_mbcnt_lo(-1, 0)", dtypes.int32)).cast(dtypes.weakint)
+  col, half = hw_lane % 16, hw_lane // 16
  outputs = tuple((output_block*output_waves+wave)*(16*output_tiles) + tile*16 + col for tile in range(output_tiles))
  inputs = tuple(token_block*token_tile + tile*16 + col for tile in range(token_tile//16))
  tokens = tuple(tuple(token_block*token_tile + tile*16 + half*8 + i for i in range(8)) for tile in range(token_tile//16))
  return output_waves, token_block, output_block, lane, wave, half, outputs, inputs, tokens

-def _wmma_stores(out, outputs, tokens, accs, update, half, lane, wave, output_waves):
-  # the accumulator fragment halves are exchanged between lane pairs (l, l^16) through LDS (a ds_swizzle without CUSTOM)
-  flat_accs = [acc for output_accs in accs for acc in output_accs]
-  lds = UOp.placeholder((output_waves, 32, len(flat_accs)*8), dtypes.float32, slot=33, addrspace=AddrSpace.LOCAL)
-  stores = [lds[wave, lane, a*8+i].store(acc.after(update)[i].load()) for a,acc in enumerate(flat_accs) for i in range(8)]
-  lds = lds.after(UOp.barrier(UOp.group(*stores)))
-  def values(ai:int) -> tuple[UOp, ...]:
-    own = tuple(lds[wave, lane, ai*8+i].load() for i in range(8))
-    peer = tuple(lds[wave, lane ^ 16, ai*8+i].load() for i in range(8))
+def _wmma_stores(out, outputs, tokens, accs, update, half):
+  def values(acc:UOp) -> tuple[UOp, ...]:
+    vals = tuple(acc.after(update)[i].load() for i in range(8))
+    swapped = tuple(UOp(Ops.CUSTOM, src=(value,),
+      arg=("__builtin_bit_cast(float, __builtin_amdgcn_ds_swizzle(__builtin_bit_cast(int, {0}), 50688))", dtypes.float32)) for value in vals)
    low = half.eq(0)
-    return tuple(low.where(own[i], peer[i+4]) if j == 0 else low.where(peer[i], own[i+4]) for i in range(4) for j in range(2))
-  tt = len(tokens)
-  return [out[token, output].store(value) for ot,(output,output_accs) in enumerate(zip(outputs, accs))
-          for tile,(tile_tokens,_acc) in enumerate(zip(tokens, output_accs)) for token,value in zip(tile_tokens, values(ot*tt+tile))]
+    return tuple(low.where(vals[i], swapped[i+4]) if j == 0 else low.where(swapped[i], vals[i+4]) for i in range(4) for j in range(2))
+  return [out[token, output].store(value) for output,output_accs in zip(outputs, accs)
+          for tile_tokens,acc in zip(tokens, output_accs) for token,value in zip(tile_tokens, values(acc))]

 def _quant_linear_wmma(out, x, out_features, in_features, type_words, layout, dequant, name):
  x = x.reshape(out.shape[0], in_features)
-  output_waves, token_block, output_block, lane, wave, physical_half, outputs, input_tokens, tokens = layout
+  _, token_block, output_block, lane, wave, physical_half, outputs, input_tokens, tokens = layout
  token_tile, output_tiles = len(tokens)*16, len(outputs)
  output_words = in_features // GGML_BLOCK_SIZE * type_words
  accs = tuple(tuple(UOp.placeholder((8,), dtypes.float32, slot=ot*(token_tile//16)+tile, addrspace=AddrSpace.REG)
@@ -287,8 +229,8 @@ def _quant_linear_wmma(out, x, out_features, in_features, type_words, layout, de
        wmma_accs[output_tile][tile] = UOp.wmma(afrag, bfrag, previous, *WMMA_ARG)
  update = UOp.group(*(acc.store(value) for output_accs,output_values in zip(accs, wmma_accs)
                       for acc,value in zip(output_accs, output_values))).end(group)
-  stores = _wmma_stores(out, outputs, tokens, accs, update, physical_half, lane, wave, output_waves)
-  return UOp.group(*stores).end(token_block, output_block, lane, wave).sink(arg=KernelInfo(name=name, opts_to_apply=()))
+  return UOp.group(*_wmma_stores(out, outputs, tokens, accs, update, physical_half)).end(token_block, output_block, lane, wave).sink(
+    arg=KernelInfo(name=name, opts_to_apply=()))

@functools.cache
 def _q5_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, out_features:int, in_features:int, ggml_type:int) -> UOp:
@@ -322,10 +264,9 @@ def _iq4_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, lut:UOp, out_features:i
      pairs = tuple(lut[((raw[base + 2 + subgroup*4 + word] >> (byte*8)) & 255).cast(dtypes.weakint)]
                    for word in range(4) for byte in range(4))
      return tuple((_half((pair >> (half*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in pairs)
-    # a subgroup-half gathers the lo (half=0) or hi (half=1) nibbles of byte pairs of each packed word
-    lut_pairs = (lut[(((raw[base+2+subgroup*4+i] >> (8*j+4*half)) & 15) |
-                      (((raw[base+2+subgroup*4+i] >> (8*j+8+4*half)) & 15) << 4)).cast(dtypes.weakint)]
-                 for i in range(4) for j in (0, 2))
+    def nibble(packed:UOp, index:int): return (packed >> (8*index+4*half)) & 15
+    lut_pairs = (lut[(nibble(packed, i) | nibble(packed, i+1)<<4).cast(dtypes.weakint)]
+                 for packed in (raw[base+2+subgroup*4+i] for i in range(4)) for i in (0, 2))
    return tuple((_half((pair >> (i*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in lut_pairs for i in range(2))
  return _quant_linear_wmma(out, x, out_features, in_features, IQ4_WORDS, layout, dequant, "linear_iq4_xs_f16_wmma")

@@ -346,163 +287,67 @@ def q8_linear(layer:Linear, x:Tensor) -> Tensor:
    fxn = _iq4_linear_f16_wmma_kernel if layer.ggml_type == IQ4_XS else functools.partial(_q5_linear_f16_wmma_kernel, ggml_type=layer.ggml_type)
    extra = (iq4_half_lut(str(x.device)).uop,) if layer.ggml_type == IQ4_XS else ()
    return run(fxn, out, raw, x.cast(dtypes.float16).contiguous().uop, *extra)
-  xq_, xd, xs = q8_quantize(x, tokens, in_features)
+  xq, xd = q8_quantize(x, tokens, in_features)
  decode = functools.partial(_quant_decode_kernel, ggml_type=layer.ggml_type)
  out = Tensor.empty(tokens, out_features, (in_features+1023)//1024, dtype=dtypes.float32, device=x.device).uop
-  return run(decode, out, raw, xq_.uop, xd.uop, xs.uop)
-
-# ******** tiny dense fp16 gemv ********
-
-@functools.cache
-def _amd_f16_gemv_kernel(out:UOp, w:UOp, x:UOp, *rest:UOp, in_features:int, out_features:int, tokens:int) -> UOp:
-  bias: UOp|None = rest[0] if rest else None
-  # one block per (token, output row), 32 lanes accumulate 4-wide chunks of the row
-  lanes, val_chunk = WARP_SIZE, 4
-  token, out_row = UOp.range(tokens, 0, AxisType.GLOBAL), UOp.range(out_features, 1, AxisType.GLOBAL)
-  lane = UOp.range(lanes, 2, axis_type=AxisType.LOCAL)
-  per = in_features // (lanes * val_chunk)
-  assert per * lanes * val_chunk == in_features
-  w = w.reshape((out_features, per, lanes*val_chunk))
-  x = x.reshape((tokens, per, lanes*val_chunk))
-  acc = UOp.const(0, dtypes.float32)
-  for i in range(per):
-    for j in range(val_chunk):
-      acc = acc + w[out_row, i, lane*val_chunk + j].load().float() * x[token, i, lane*val_chunk + j].load().float()
-  total = warp_reduce(acc, full_wave=True)
-  if bias is not None: total = total + bias[token, out_row].load().float()
-  return out[token, out_row.valid(lane.eq(0))].store(total).end(token, out_row, lane).sink(arg=KernelInfo(name="linear_f16_gemv", opts_to_apply=()))
-
-def _view_back(t:Tensor) -> Tensor:
-  """strip top-of-chain CAST(s) from a lazy weight: reading the raw file bytes in the kernel instead of
-  materializing the cast into a fresh buffer every step"""
-  uop = t.uop
-  while uop.op is Ops.CAST: uop = uop.src[0]
-  return Tensor(uop).reshape(t.shape)
-
-def f16_gemv(layer:Linear, x:Tensor) -> Tensor:
-  tokens = prod(x.shape[:-1])
-  assert isinstance(tokens, int)
-  weight = _view_back(layer.weight)
-  x = x.contiguous() if x.dtype == dtypes.half else x.cast(dtypes.half).contiguous()
-  out = Tensor.empty(tokens, layer.out_features, dtype=dtypes.float32, device=x.device)
-  fxn = functools.partial(_amd_f16_gemv_kernel, in_features=layer.in_features, out_features=layer.out_features, tokens=tokens)
-  srcs = (out, weight.reshape(-1), x.reshape(tokens, layer.in_features)) + (() if layer.bias is None else (_view_back(layer.bias),))
-  return Tensor.custom_kernel(*srcs, fxn=fxn)[0].reshape(*x.shape[:-1], layer.out_features)
+  return run(decode, out, raw, xq.uop, xd.uop)

 # ******** flash attention on the KV cache ********

-def _vec_load(ptr:UOp, lanes:int) -> tuple[UOp, ...]:
-  if lanes == 1: return (ptr.load().float(),)
-  vec = _amd_load(ptr, lanes)
-  return tuple(vec[i].float() for i in range(lanes))
-
@functools.cache
-def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, valid_kv_len, max_kv_len, block_n, waves=4):
+def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, valid_kv_len, max_kv_len, block_n):
  valid_kv_len = _unbind(valid_kv_len)
  _, B, H_KV, N, D = cast(tuple[int, int, int, int, int], cache_kv.shape)
  _, H, M, _ = cast(tuple[int, int, int, int], q.shape)
  assert M == 1 and H % H_KV == 0 and D % WARP_SIZE == 0 and max_kv_len <= N and max_kv_len % block_n == 0
-  G, CHUNK, DPL, WAVES = H // H_KV, block_n, D // WARP_SIZE, waves
-  assert CHUNK % WAVES == 0
-  SEC = CHUNK // WAVES  # keys each wave scans independently
-  live_chunks = (valid_kv_len+CHUNK-1)//CHUNK
-  live_chunks = min(live_chunks, out.shape[2]) if isinstance(live_chunks, int) else live_chunks.minimum(out.shape[2])
-  block_bhkv, block_chunk = UOp.range(B*H_KV, 0, AxisType.GLOBAL), UOp.range(live_chunks, 1, AxisType.GLOBAL)
-  lane, wave = UOp.range(WARP_SIZE, 2, axis_type=AxisType.LOCAL), UOp.range(WAVES, 3, axis_type=AxisType.LOCAL)
-  b, kv_head = block_bhkv // H_KV, block_bhkv % H_KV
-  # per-lane query fragments for every GQA head, kept packed in registers; unpacked at use
-  qf = tuple(_vec_load(q[b, kv_head*G+h, 0, lane*DPL], DPL) for h in range(G))
-  zerof = UOp.const(0, dtypes.float)
-  valids: list[UOp] = []
-  scores: list[list[UOp]] = [[zerof]*G for _ in range(SEC)]
-  vfrags: list[tuple[UOp, ...]] = [()]*SEC
-  for j in range(SEC):
-    key = block_chunk*CHUNK + wave*SEC + j
-    valid = key < valid_kv_len
-    valids.append(valid)
-    kfrag = _vec_load(cache_kv[0, b, kv_head, key, lane*DPL], DPL)
-    # V is prefetched in the score pass so both streams are in flight together
-    vfrags[j] = _vec_load(cache_kv[1, b, kv_head, key, lane*DPL], DPL)
-    for h in range(G):
-      s = warp_reduce(sum((qf[h][i]*kfrag[i] for i in range(DPL)), UOp.const(0, dtypes.float)), full_wave=True) * (1/math.sqrt(D))
-      scores[j][h] = valid.where(s, UOp.const(-math.inf, dtypes.float))
-  ninf = UOp.const(-math.inf, dtypes.float)
-  row_max = [functools.reduce(UOp.maximum, (scores[j][h] for j in range(SEC)), ninf) for h in range(G)]
-  accs:list[list[UOp]] = [[UOp.const(0, dtypes.float)] * DPL for _ in range(G)]
-  row_sums:list[UOp] = [UOp.const(0, dtypes.float) for _ in range(G)]
-  for j in range(SEC):
-    for h in range(G):
-      beta = valids[j].where(((scores[j][h]-row_max[h])*LOG2E).exp2(), UOp.const(0, dtypes.float))
-      accs[h] = [a + beta*v for a, v in zip(accs[h], vfrags[j])]
-      row_sums[h] = row_sums[h] + beta
-  # exchange across the block's waves through LDS (fp16 halves LDS so more blocks fit per CU)
-  acc_lds = UOp.placeholder((WAVES, G, D), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL)
-  ml_lds = UOp.placeholder((WAVES, G, 2), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
-  lds_acc = acc_lds.reshape(WAVES, G, WARP_SIZE, DPL)
-  stores = [lds_acc[wave, h, lane].store(UOp.stack(*accs[h]).cast(dtypes.half)) for h in range(G)]
-  # NOTE: duplicate stores of the same value from every lane are harmless here
-  stores += [ml_lds[wave, h, i].store(x) for h in range(G) for i, x in enumerate((row_max[h], row_sums[h]))]
-  barrier = UOp.barrier(UOp.group(*stores))
-  acc_lds, ml_lds = acc_lds.after(barrier), ml_lds.after(barrier)
-  tid = wave*WARP_SIZE + lane
-  final_stores:list[UOp] = []
-  for i in range(-(-G*D//(WAVES*WARP_SIZE))):
-    flat = tid + i*WAVES*WARP_SIZE
-    h, d = flat // D, flat % D
-    M = functools.reduce(UOp.maximum, (ml_lds[w, h, 0].load() for w in range(WAVES)), ninf)
-    val = sum((((ml_lds[w, h, 0].load()-M)*LOG2E).exp2() * acc_lds[w, h, d].load().float() for w in range(WAVES)), UOp.const(0, dtypes.float))
-    oidx = out[b, kv_head*G + h, block_chunk, d]
-    if G*D % (WAVES*WARP_SIZE): oidx = out[b, (kv_head*G + h).valid(flat < G*D), block_chunk, d]
-    final_stores.append(oidx.store(val))
-  hstat = tid
-  M = functools.reduce(UOp.maximum, (ml_lds[w, hstat, 0].load() for w in range(WAVES)), ninf)
-  L = sum((((ml_lds[w, hstat, 0].load()-M)*LOG2E).exp2() * ml_lds[w, hstat, 1].load() for w in range(WAVES)), UOp.const(0, dtypes.float))
-  q_head = (kv_head*G + hstat).valid(hstat < G) if WAVES*WARP_SIZE > G else kv_head*G + hstat
-  final_stores += [stats[b, q_head, block_chunk, 0].store(M), stats[b, q_head, block_chunk, 1].store(L)]
-  return UOp.group(*final_stores).end(lane, wave, block_chunk, block_bhkv).sink(arg=KernelInfo(name="flash_decode_partial", opts_to_apply=()))
-
-@functools.cache
-def _amd_flash_decode_combine(o:UOp, partial:UOp, stats:UOp, live:int|UOp) -> UOp:
-  # one wave per (batch, head, 64-dim tile): every lane redundantly weights its chunks; no cross-lane traffic
-  live = _unbind(live)
-  B, H, C, D = cast(tuple[int, int, int, int], partial.shape)
-  DT = 64 if D % 64 == 0 else WARP_SIZE  # dims per block
-  assert D % DT == 0
-  block_bh, block_dt = UOp.range(B*H, 0, AxisType.GLOBAL), UOp.range(D//DT, 1, AxisType.GLOBAL)
-  lane = UOp.range(WARP_SIZE, 2, axis_type=AxisType.LOCAL)
-  b, h = block_bh // H, block_bh % H
-  NPD = DT // WARP_SIZE  # output dims per lane
-  dims = tuple(block_dt*DT + lane*NPD + i for i in range(NPD))
-  chunk = UOp.range(live, 100, AxisType.REDUCE)
-  def iloop(ph, val): return ph.store(ph.const_like(val))
-  chunk_max = UOp.placeholder((1,), dtypes.float, slot=0, addrspace=AddrSpace.REG)
-  chunk_max_i = chunk_max.after(iloop(chunk_max, -math.inf))
-  update0 = chunk_max_i.store(chunk_max_i.after(chunk).maximum(stats[b, h, chunk, 0].load())).end(chunk)
-  chunk_max = chunk_max_i.after(update0)
-  chunk2 = UOp.range(live, 101, AxisType.REDUCE)
-  acc = UOp.placeholder((NPD,), dtypes.float, slot=1, addrspace=AddrSpace.REG)
-  weight_sum = UOp.placeholder((1,), dtypes.float, slot=2, addrspace=AddrSpace.REG)
-  acc_i, weight_sum_i = acc.after(iloop(acc, 0)), weight_sum.after(iloop(weight_sum, 0))
-  w = ((stats[b, h, chunk2, 0].load()-chunk_max)*LOG2E).exp2()
-  update1 = UOp.group(*[acc_i[i].store(acc_i.after(chunk2)[i].load() + w*partial[b, h, chunk2, d].load()) for i, d in enumerate(dims)],
-                      weight_sum_i[0].store(weight_sum_i.after(chunk2)[0].load() + w*stats[b, h, chunk2, 1].load())).end(chunk2)
-  acc, weight_sum = acc_i.after(update1), weight_sum_i.after(update1)
-  inv = 1 / weight_sum[0].load()
-  return UOp.group(*[o[b, h, 0, d].store(acc[i].load() * inv) for i, d in enumerate(dims)]) \
-    .end(lane, block_dt, block_bh).sink(arg=KernelInfo(name="flash_decode_combine", opts_to_apply=()))
+  G, CHUNK, DV, heads_per_wave = H // H_KV, block_n, D // WARP_SIZE, 2
+  head_tile = min(DECODE_HEAD_TILE, G)  # share each KV stream across two GQA heads per wave
+  assert G % head_tile == 0 and head_tile % heads_per_wave == 0
+  decode_waves, decode_group = head_tile // heads_per_wave, 4
+  block_bhkv = UOp.range(B*H_KV*(G//head_tile), 0, AxisType.GLOBAL)
+  valid_chunks = (valid_kv_len+CHUNK-1)//CHUNK
+  group_count = min(valid_chunks, out.shape[2]) if isinstance(valid_chunks, int) else valid_chunks.minimum(out.shape[2])
+  block_n, lane = UOp.range(group_count, 1, AxisType.GLOBAL), UOp.range(WARP_SIZE, 2, axis_type=AxisType.LOCAL)
+  wave = UOp.range(decode_waves, 3, axis_type=AxisType.LOCAL)
+  head_group, bhkv = block_bhkv % (G//head_tile), block_bhkv // (G//head_tile)
+  b, kv_head = bhkv // H_KV, bhkv % H_KV
+  dims = tuple(lane + i*WARP_SIZE for i in range(DV))
+  acc, row_max, row_sum = _reg((heads_per_wave, DV), 0, 0), _reg((heads_per_wave,), 1, -math.inf), _reg((heads_per_wave,), 2, 0)
+  groups_per_chunk, offset = CHUNK // decode_group, UOp.range(((valid_chunks+group_count-1)//group_count)*(CHUNK//decode_group), 100, AxisType.REDUCE)
+  chunk = block_n + (offset // groups_per_chunk) * group_count
+  keys = tuple(chunk*CHUNK + (offset % groups_per_chunk)*decode_group + i for i in range(decode_group))
+  valid = tuple(key < valid_kv_len for key in keys)
+  kvals, vvals = (tuple(tuple(is_valid.where(cache_kv[kv, b, kv_head, key, d].float(), UOp.const(0, dtypes.float)) for d in dims)
+    for key,is_valid in zip(keys, valid)) for kv in range(2))
+  q_heads = tuple(kv_head*G + head_group*head_tile + wave*heads_per_wave + head for head in range(heads_per_wave))
+  updates:list[UOp] = []
+  for head,q_head in enumerate(q_heads):
+    scores = tuple(warp_reduce(sum((q[b, q_head, 0, d].float()*k for d,k in zip(dims, key_kvals)),
+                                   UOp.const(0, dtypes.float)), full_wave=True) / math.sqrt(D) for key_kvals in kvals)
+    prev_acc, prev_max, prev_sum = acc.after(offset)[head], row_max.after(offset)[head], row_sum.after(offset)[head]
+    new_max = functools.reduce(lambda a,vs:a.maximum(vs[0].where(vs[1], UOp.const(-math.inf, dtypes.float))), zip(valid, scores), prev_max)
+    alpha = ((prev_max-new_max)*LOG2E).exp2()
+    betas = tuple(is_valid.where(((score-new_max)*LOG2E).exp2(), UOp.const(0, dtypes.float)) for is_valid,score in zip(valid, scores))
+    updates += [acc[head].store(prev_acc*alpha + sum((UOp.stack(*value)*beta for value,beta in zip(vvals, betas)), acc[head].const_like(0))),
+                row_sum[head].store(prev_sum*alpha + sum(betas, UOp.const(0, dtypes.float))), row_max[head].store(new_max)]
+  update = UOp.group(*updates).end(offset)
+  acc, row_max, row_sum = acc.after(update), row_max.after(update), row_sum.after(update)
+  stores = [out[b, q_head, block_n, d].store(acc[head, i]) for head,q_head in enumerate(q_heads) for i,d in enumerate(dims)] + \
+    [stats[b, q_head.valid(lane.eq(0)), block_n, i].store(x[head]) for head,q_head in enumerate(q_heads) for i,x in enumerate((row_max, row_sum))]
+  return UOp.group(*stores).end(lane, wave, block_n, block_bhkv).sink(arg=KernelInfo(name="flash_decode_partial", opts_to_apply=()))

 def amd_flash_attention_decode(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, max_kv_len:int) -> Tensor:
  B, H, D = cache_kv.shape[1], q.shape[1], cache_kv.shape[4]
-  chunks = min(256, max_kv_len // 64)
+  chunks = min(64, max_kv_len // 128)
  partial = Tensor.empty(B, H, chunks, D, dtype="float32", device=q.device)
  stats = Tensor.empty(B, H, chunks, 2, dtype="float32", device=q.device)
-  fxn = functools.partial(_amd_flash_attention_decode_partial, valid_kv_len=valid_kv_len, max_kv_len=max_kv_len, block_n=64, waves=16)
+  fxn = functools.partial(_amd_flash_attention_decode_partial, valid_kv_len=valid_kv_len, max_kv_len=max_kv_len, block_n=128)
  partial, stats = Tensor.custom_kernel(partial, stats, q, cache_kv, fxn=fxn)[:2]
-  live = (valid_kv_len+63)//64
+  live = (valid_kv_len+127)//128
  live = min(live, chunks) if isinstance(live, int) else live.minimum(chunks)
-  out = Tensor.empty(B, H, 1, D, dtype="float32", device=q.device)
-  fxn = functools.partial(_amd_flash_decode_combine, live=live)
-  return Tensor.custom_kernel(out, partial, stats, fxn=fxn)[0]
+  partial, stats = partial[:, :, :live], stats[:, :, :live]
+  weights = ((stats[..., 0]-stats[..., 0].max(2, keepdim=True))*LOG2E).exp2()
+  return ((partial*weights.unsqueeze(-1)).sum(2) / (stats[..., 1]*weights).sum(2, keepdim=True)).unsqueeze(2)

@functools.cache
 def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, valid_kv_len:int|UOp, q_start:int|UOp|None=None) -> UOp:
@@ -870,7 +870,7 @@ class ElementwiseMixin(CreationMixin):
    print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).asinh().numpy())
    ```
    """
-    return (sg:=(self<0).where(-1.0, 1.0)) * (self*sg + (self.square() + 1).sqrt()).log()
+    return self.sign() * (self.abs() + (self.square() + 1).sqrt()).log()

  def acosh(self) -> Self:
    """
@@ -31,10 +31,6 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
    return (None,) + k.arg.grad_fxn(on_dev(ctx, 0), k)
  assert fxn.op is Ops.TUPLE, f"expected TUPLE body for gradient, got {fxn.op}"
  params = {x.arg.slot:x for x in fxn.toposort(enter_calls=False) if x.op == Ops.PARAM}
-  # grads are collected at the flat param storage: reshape to each arg's view (max view shrunk to symbolic)
-  def shaped_grad(grad:UOp, i:int) -> UOp:
-    a = args[i]
-    return grad.view_as(a.shard_shape, a.axis) if a.axis is not None and isinstance(a.device, tuple) else grad.view_as(a._shape)
  grad_args = ctx.src
  root_grad = UOp(Ops.TUPLE, src=tuple(UOp(Ops.NOOP) if g.op is Ops.NOOP else
    g if g.device is None else g.param_like(len(args)+i) for i,g in enumerate(grad_args)))
@@ -43,7 +39,7 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
  fwd_subs = {src: src.param_like(len(args)+len(grad_args)+i) for i, src in enumerate(fxn.src)} if k.arg.precompile else {}
  fwd_outs = tuple(k.gettuple(i) for i in range(len(fxn.src))) if k.arg.precompile else ()
  # collect needed gradient bodies, compact unused params, create a single backward CALL
-  grad_bodies = [(i, shaped_grad(grads[p], i)) for i in needed if (p:=params.get(i)) is not None and p in grads]
+  grad_bodies = [(i, grads[p]) for i in needed if (p:=params.get(i)) is not None and p in grads]
  bwd_body = UOp.maketuple(*(gb for _, gb in grad_bodies)).substitute(fwd_subs, walk=True)
  bwd_body = renumber_invalid_outputs(bwd_body)
  bwd_body, compact_args = _compact_params(bwd_body, (*args, *grad_args, *fwd_outs))
@@ -1067,7 +1067,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
      reshape[i] = expand[i] = size[i]
      if mode == "linear":
        arr = type(self).arange(size[i])
-        num, den = (arr*(in_sz-1), max(size[i]-1, 1)) if align_corners else ((arr*2+1)*in_sz - size[i], size[i]*2)
+        num, den = (arr*(in_sz-1), size[i]-1) if align_corners else ((arr*2+1)*in_sz - size[i], size[i]*2)
        num = num.clip(0, (in_sz-1)*den)
        low, high, perc = [y.reshape(reshape).expand(expand) for y in (num//den, (num+den-1)//den, (num % den).cast(dtypes.float32)/den)]
        x = x.gather(i, low).lerp(x.gather(i, high), perc)
@@ -99,7 +99,7 @@ class Conv2d:
    if isinstance(padding, str):
      if padding.lower() != 'same': raise ValueError(f"Invalid padding string {padding!r}, only 'same' is supported")
      if stride != 1: raise ValueError("padding='same' is not supported for strided convolutions")
-      pad = [(d*(k-1)//2, d*(k-1) - d*(k-1)//2) for d,k in zip(make_tuple(dilation, len(self.kernel_size))[::-1], self.kernel_size[::-1])]
+      pad = [(d*(k-1)//2, d*(k-1) - d*(k-1)//2) for d,k in zip(make_tuple(dilation, len(self.kernel_size)), self.kernel_size[::-1])]
      padding = tuple(flatten(pad))
    self.stride, self.dilation, self.groups, self.padding = stride, dilation, groups, padding
    scale = 1 / math.sqrt(in_channels * prod(self.kernel_size))
@@ -556,8 +556,8 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    return tuple(t if cond else e for t,e in zip(then_out.values(), else_out.values()))

  def Identity(x:Tensor): return x
-  def Constant(sparse_value:Tensor|None=None, value:Tensor|None=None, value_float:float|None=None, value_floats:tuple[float, ...]|None=None,
-              value_int:int|None=None, value_ints:tuple[int, ...]|None=None, value_string:str|None=None, value_strings:tuple[str, ...]|None=None):
+  def Constant(sparse_value:Tensor|None=None, value:Tensor|None=None, value_float:float|None=None, value_floats:list[float]|None=None,
+              value_int:int|None=None, value_ints:list[int]|None=None, value_string:str|None=None, value_strings:list[str]|None=None):
    if value is not None: return value
    if value_float is not None: return Tensor(value_float, dtype=dtypes.float32)
    if value_floats is not None: return Tensor(list(value_floats), dtype=dtypes.float32)
@@ -594,7 +594,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT

  # ***** Unary Ops (math) *****
  def Not(x:Tensor): return x.logical_not()
-  def Clip(x: Tensor, min:Tensor|float|None=None, max:Tensor|float|None=None): return x if min is None and max is None else x.clip(min, max)  # noqa: A002 # pylint: disable=redefined-builtin
+  def Clip(x: Tensor, min:Tensor|None=None, max:Tensor|None=None): return x if min is None and max is None else x.clip(min, max)  # noqa: A002 # pylint: disable=redefined-builtin
  def IsInf(x:Tensor, detect_negative:int=1, detect_positive:int=1): return x.isinf(bool(detect_positive), bool(detect_negative))

  # ***** Unary Ops (activation) *****
@@ -643,26 +643,26 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
  def Min(*data_0:Tensor): return functools.reduce(Tensor.minimum, data_0)
  def Sum(*data_0:Tensor): return functools.reduce(Tensor.add, data_0)
  def Mean(*data_0:Tensor): return Sum(*data_0) / len(data_0)
-  def ReduceMax(data:Tensor, axes:Sequence[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+  def ReduceMax(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
    return data.max(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
-  def ReduceMin(data:Tensor, axes:Sequence[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+  def ReduceMin(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
    return data.min(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
-  def ReduceSum(data:Tensor, axes:Sequence[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+  def ReduceSum(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
    return data.sum(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
-  def ReduceMean(data:Tensor, axes:Sequence[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+  def ReduceMean(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
    return data.mean(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
-  def ReduceSumSquare(data:Tensor, axes:Sequence[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+  def ReduceSumSquare(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
    return ReduceSum(data.square(), axes, keepdims, noop_with_empty_axes)
-  def ReduceProd(data:Tensor, axes:Sequence[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+  def ReduceProd(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
    return data.prod(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
-  def ReduceL1(data:Tensor, axes:Sequence[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+  def ReduceL1(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
    return ReduceSum(data.abs(), axes, keepdims, noop_with_empty_axes)
-  def ReduceL2(data:Tensor, axes:Sequence[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+  def ReduceL2(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
    dtype = dtypes.float if data.dtype in (dtypes.float16, dtypes.bfloat16) else data.dtype
    return ReduceSum(data.cast(dtype).square(), axes, keepdims, noop_with_empty_axes).sqrt().cast(data.dtype)
-  def ReduceLogSum(data:Tensor, axes:Sequence[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+  def ReduceLogSum(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
    return ReduceSum(data, axes, keepdims, noop_with_empty_axes).log()
-  def ReduceLogSumExp(data:Tensor, axes:Sequence[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
+  def ReduceLogSumExp(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
    return ReduceSum(data.exp(), axes, keepdims, noop_with_empty_axes).log()
  def ArgMax(x:Tensor, axis:int=0, keepdims:int=1, select_last_index:int=0):
    if select_last_index: return ((int(x.shape[axis])-1) - x.flip(axis).argmax(axis, keepdim=keepdims)).cast(dtypes.int64)
@@ -671,32 +671,32 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    return ArgMax(-x, axis=axis, keepdims=keepdims, select_last_index=select_last_index)

  # ***** Movement Ops *****
-  def Reshape(data:Tensor, shape:Sequence[int], allowzero:int=0):
+  def Reshape(data:Tensor, shape:list[int], allowzero:int=0):
    return data.reshape([x if x != 0 else (0 if allowzero else data.shape[i]) for i,x in enumerate(shape)])
  def Flatten(x:Tensor, axis:int=1): return x.reshape(prod(x.shape[0:axis]), -1)
  def Expand(x:Tensor, shape:list[int]): return x.expand(_broadcast_shape(x.shape, tuple(shape)))
  def Shrink(x:Tensor, bias:float=0.0, lambd:float=0.5): return (x < -lambd)*(x+bias) + (x > lambd)*(x-bias)
-  def Transpose(x:Tensor, perm:tuple[int, ...]|None=None): return x.permute(order=perm or list(range(x.ndim)[::-1]))
+  def Transpose(x:Tensor, perm:list[int]|None=None): return x.permute(order=perm or list(range(x.ndim)[::-1]))

-  def Squeeze(data:Tensor, axes:Sequence[int]|None=None):
+  def Squeeze(data:Tensor, axes:list[int]|None=None):
    return data.squeeze() if axes is None else functools.reduce(lambda d, dim: d.squeeze(dim), sorted(axes, reverse=True), data)
-  def Unsqueeze(data:Tensor, axes:Sequence[int]): return functools.reduce(lambda d, dim: d.unsqueeze(dim), sorted(axes), data)
+  def Unsqueeze(data:Tensor, axes:list[int]): return functools.reduce(lambda d, dim: d.unsqueeze(dim), sorted(axes), data)

  def Tile(x:Tensor, repeats:list[int]): return x.repeat(repeats)
  def Concat(*xs:Tensor, axis:int): return Tensor.cat(*xs, dim=axis)
-  def Slice(data:Tensor, starts:Sequence[int], ends:Sequence[int], axes:Sequence[int]|None=None, steps:list[int]|None=None):
+  def Slice(data:Tensor, starts:list[int], ends:list[int], axes:list[int]|None=None, steps:list[int]|None=None):
    axes = axes or list(range(data.ndim))
    steps = steps or [1] * data.ndim
    slices = [slice(None)] * data.ndim
    for i, axis in enumerate(axes): slices[axis] = slice(starts[i], ends[i], steps[i])
    return data[tuple(slices)]

-  def Split(data:Tensor, split:Sequence[int]|None=None, num_outputs:int=0, axis:int=0):
+  def Split(data:Tensor, split:list[int]|None=None, num_outputs:int=0, axis:int=0):
    sz = int(data.shape[axis])
    if split is None: split = [sz // num_outputs + (1 if i < sz % num_outputs else 0) for i in range(num_outputs)]
    return data.split(split, axis)

-  def Pad(x:Tensor, pads:Sequence[int], constant_value:ConstType|None=None, axes:list[int]|None=None,
+  def Pad(x:Tensor, pads:list[int], constant_value:ConstType|None=None, axes:list[int]|None=None,
          mode:Literal["constant", "reflect", "edge", "wrap"]="constant", value=0):
    value = _resolve_const(value if constant_value is None else constant_value)
    axes = axes or list(range(x.ndim))
@@ -704,7 +704,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    for i,axis in enumerate(axes): real_pads[axis%x.ndim], real_pads[axis%x.ndim+x.ndim] = pads[i], pads[i+len(axes)]
    return x.pad(padding=_onnx_pads_to_tiny_pads(real_pads), mode={"edge":"replicate", "wrap":"circular"}.get(mode, mode), value=value)

-  def CenterCropPad(t:Tensor, shape:list[int], axes:tuple[int, ...]|None=None):
+  def CenterCropPad(t:Tensor, shape:list[int], axes:list[int]|None=None):
    shrink_arg:list[None|tuple[sint,sint]] = [None] * t.ndim
    pad_arg:list[None|tuple[sint,sint]] = [None] * t.ndim
    for s, x in zip(shape, axes or range(t.ndim)):
@@ -714,26 +714,26 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    return t.shrink(tuple(shrink_arg)).pad(tuple(pad_arg))

  # ***** Processing Ops *****
-  def AveragePool(X: Tensor, kernel_shape:tuple[int, ...], auto_pad:AUTO_PAD_OPTIONS="NOTSET", ceil_mode:int=0, count_include_pad:int=0,
-                  dilations:tuple[int, ...]|int=1, pads:tuple[int, ...]|int=0, strides:tuple[int, ...]|int=1):
+  def AveragePool(X: Tensor, kernel_shape:list[int], auto_pad:AUTO_PAD_OPTIONS="NOTSET", ceil_mode:int=0, count_include_pad:int=0,
+                  dilations:list[int]|int=1, pads:list[int]|int=0, strides:list[int]|int=1):
    pool_pads = _resolve_pool_pads(X, pads, kernel_shape, dilations, strides, auto_pad)
    return X.avg_pool2d(tuple(kernel_shape), strides, dilations, pool_pads, ceil_mode=ceil_mode, count_include_pad=count_include_pad)

-  def MaxPool(X: Tensor, kernel_shape:tuple[int, ...], auto_pad:AUTO_PAD_OPTIONS="NOTSET", ceil_mode:int=0, dilations:tuple[int, ...]|int=1,
-              pads:tuple[int, ...]|int=0, storage_order:int=0, strides:tuple[int, ...]|int=1):
+  def MaxPool(X: Tensor, kernel_shape:list[int], auto_pad:AUTO_PAD_OPTIONS="NOTSET", ceil_mode:int=0, dilations:list[int]|int=1, pads:list[int]|int=0,
+              storage_order:int=0, strides:list[int]|int=1):
    pool_pads = _resolve_pool_pads(X, pads, kernel_shape, dilations, strides, auto_pad)
    out = X.max_pool2d(tuple(kernel_shape), strides, dilations, pool_pads, ceil_mode=ceil_mode, return_indices=True)
    ret, idx = cast(tuple[Tensor, Tensor], out)
    return ret, idx.transpose(-2, -1).cast(dtypes.int64) if storage_order else idx.cast(dtypes.int64)

-  def Conv(X: Tensor, W: Tensor, B:Tensor|None=None, auto_pad:AUTO_PAD_OPTIONS="NOTSET", dilations:tuple[int, ...]|int=1, group:int=1,
-          kernel_shape:tuple[int, ...]|None=None, pads:tuple[int, ...]|int=0, strides:tuple[int, ...]|int=1):
+  def Conv(X: Tensor, W: Tensor, B:Tensor|None=None, auto_pad:AUTO_PAD_OPTIONS="NOTSET", dilations:list[int]|int=1, group:int=1,
+          kernel_shape:list[int]|None=None, pads:list[int]|int=0, strides:list[int]|int=1):
    return X.conv2d(W, B, stride=strides, groups=group, dilation=dilations,
                    padding=_resolve_pool_pads(X, pads, kernel_shape or W.shape[2:], dilations, strides, auto_pad))

-  def ConvTranspose(X: Tensor, W: Tensor, B:Tensor|None=None, auto_pad:AUTO_PAD_OPTIONS="NOTSET", dilations:tuple[int, ...]|int=1, group:int=1,
-                    kernel_shape:tuple[int, ...]|None=None, pads:Sequence[int]|None=None, output_shape:Sequence[int]|None=None,
-                    output_padding:tuple[int, ...]|int=0, strides:tuple[int, ...]|int=1):
+  def ConvTranspose(X: Tensor, W: Tensor, B:Tensor|None=None, auto_pad:AUTO_PAD_OPTIONS="NOTSET", dilations:list[int]|int=1, group:int=1,
+                    kernel_shape:list[int]|None=None, pads:list[int]|None=None, output_shape:list[int]|None=None, output_padding:list[int]|int=0,
+                    strides:list[int]|int=1):
    input_shape_, kernel_shape_ = X.shape[2:], (kernel_shape or W.shape[2:])
    strides_, dilations_, output_padding_ = (make_tuple(x, len(input_shape_)) for x in (strides, dilations, output_padding))
    if output_shape is not None: # we pad according to output_shape
@@ -747,10 +747,10 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    pads = _onnx_pads_to_tiny_pads(pads)
    return X.conv_transpose2d(W, B, group, strides_, dilations_, pads, output_padding_)

-  def MaxUnpool(xT: Tensor, xI: Tensor, outshape: list[int]|None=None, kernel_shape:Sequence[int]|None=None, pads:tuple[int, ...]|int=0,
-                strides:tuple[int, ...]|int=1):
+  def MaxUnpool(xT: Tensor, xI: Tensor, outshape: list[int]|None=None, kernel_shape:list[int]|None=None, pads:list[int]|int=0,
+                strides:list[int]|int=1):
    if kernel_shape is None: kernel_shape = []
-    pads_: int | tuple[int, ...] = pads if isinstance(pads, int) else _onnx_pads_to_tiny_pads(pads)
+    pads_: int | tuple[int, ...] = tuple(pads) if isinstance(pads, list) else pads
    return Tensor.max_unpool2d(xT, xI, tuple(kernel_shape), strides, 1, pads_, outshape if outshape is None else tuple(outshape))

  def GlobalAveragePool(X:Tensor): return X.mean(axis=tuple(range(2, X.ndim)), keepdim=True)
@@ -761,7 +761,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    if C is not None: ret = ret + beta * (C if broadcast == 0 else C.reshape([-1 if i < len(C.shape) else 1 for i in range(ret.ndim)][::-1]))
    return ret

-  def Einsum(*Inputs:Tensor, equation:str): return Tensor.einsum(equation, *Inputs)
+  def Einsum(*Inputs:list[Tensor], equation:str): return Tensor.einsum(equation, *Inputs)

  def CumSum(X:Tensor, axis:int|list[int], exclusive:int=0, reverse:int=0):
    axis = X._resolve_dim(_resolve_const(axis))
@@ -774,8 +774,8 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    k_ = _resolve_const(k)
    return x.triu(k_) if upper else x.tril(k_)

-  def Resize(X:Tensor, roi:list[float]|None=None, scales:Sequence[float]|None=None, sizes:list[int]|None=None, antialias:int=0,
-        axes:Sequence[int]|None=None, coordinate_transformation_mode:str='half_pixel', cubic_coeff_a:float=-0.75, exclude_outside:int=0,
+  def Resize(X:Tensor, roi:list[float]|None=None, scales:list[float]|None=None, sizes:list[int]|None=None, antialias:int=0,
+        axes:list[int]|None=None, coordinate_transformation_mode:str='half_pixel', cubic_coeff_a:float=-0.75, exclude_outside:int=0,
        extrapolation_value:float=0.0, keep_aspect_ratio_policy:str='stretch', mode:str='nearest', nearest_mode:str='round_prefer_floor'):
    def _apply_transformation(input_sz, output_sz, scale_dim, mode):
      index = Tensor.arange(output_sz)
@@ -876,7 +876,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
        gathered_values = [X.gather(i, idx) for idx in expanded_indices]
        X = sum(v * c for v, c in zip(gathered_values, expanded_coeffs))
    return X.permute(*argsort(perm)) if perm else X
-  def Upsample(X:Tensor, scales:Sequence[float], mode:str): return Resize(X=X, scales=scales, mode=mode)  # deprecated
+  def Upsample(X, scales, mode): return Resize(X=X, scales=scales, mode=mode)  # deprecated

  def TopK(X:Tensor, K:int|list[int], axis:int=-1, largest:int=1, sorted:int=1):  # noqa: A002 # pylint: disable=redefined-builtin
    val, idx = X.topk(_resolve_const(K), axis, bool(largest), bool(sorted))
@@ -937,8 +937,9 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    if segment_embedding is not None: embedding_sum = embedding_sum + embedding(segment_ids, segment_embedding.shape[0], segment_embedding)
    out = embedding_sum.layernorm(eps=epsilon) * gamma + beta
    return out, None, embedding_sum
-  def MeanVarianceNormalization(x:Tensor, axes:Sequence[int]=(0,2,3)):
-    return (x - x.mean(axes, keepdim=True)) / (x.std(axes, keepdim=True, correction=0) + 1e-9)
+  def MeanVarianceNormalization(x:Tensor, axis:list[int]|None=None):
+    if axis is None: axis = [0,2,3]
+    return (x - x.mean(axis, keepdim=True)) / (x.std(axis, keepdim=True, correction=0) + 1e-9)

  def LpNormalization(x:Tensor, axis:int=-1, p:int=2):
    return x / (x.abs().sum(axis, keepdim=True) if p == 1 else x.square().sum(axis, keepdim=True).sqrt())
@@ -1000,7 +1001,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT

  def attention_contrib(x:Tensor, weights:Tensor, bias:Tensor|None=None, mask_index:Tensor|None=None, past:Tensor|None=None,
                        attention_bias:Tensor|None=None, past_sequence_length:Tensor|None=None, do_rotary:int=0, mask_filter_value:float=-10000.0,
-                        num_heads:int|None=None, past_present_share_buffer:int|None=None, qkv_hidden_sizes:Sequence[int]|None=None,
+                        num_heads:int|None=None, past_present_share_buffer:int|None=None, qkv_hidden_sizes:list[int]|None=None,
                        rotary_embedding_dim:int|None=None, scale:float|None=None, unidirectional:int=0):
    assert not do_rotary and not attention_bias, "TODO"
    if qkv_hidden_sizes is None: qkv_hidden_sizes = [int(weights.shape[1] // 3)] * 3
@@ -1245,8 +1246,8 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
    ret = _qlinearop_float(GlobalAveragePool, [X], [x_zero_point], [x_scale], y_scale, y_zero_point)
    return ret.permute(0, *range(2, ret.ndim), 1) if channels_last else ret  # NCHW -> NHWC

-  def ConvInteger(x: Tensor, w: Tensor, x_zero_point:Tensor = Tensor(0), w_zero_point:Tensor = Tensor(0), **opts) -> Tensor:
-    return _op_integer(Conv, [x,w], [x_zero_point,w_zero_point], **opts)
+  def ConvInteger(x: Tensor, w: Tensor, x_zero_point:Tensor = Tensor(0), w_zero_point:Tensor = Tensor(0), B: Tensor | None = None, **opts) -> Tensor:
+    return _op_integer(Conv, [x,w], [x_zero_point,w_zero_point], **{"B":B, **opts})

  def MatMulInteger(A: Tensor, B: Tensor, a_zero_point: Tensor = Tensor(0), b_zero_point: Tensor = Tensor(0)) -> Tensor:
    return _op_integer(Tensor.matmul, [A,B], [a_zero_point,b_zero_point])
@@ -9,8 +9,8 @@ from tinygrad.device import Compiler

 # an access takes its dtype from the buffer it indexes, so accessing at another dtype restates the storage on the buffer that owns it
 def with_storage(x:UOp, dt:DType) -> UOp:
-  if x.op in {Ops.PARAM, Ops.BUFFER}: return x.replace(arg=replace(x.arg, dtype=dt))
-  return x.replace(src=(with_storage(x.src[0], dt),)+x.src[1:])
+  if x.op in {Ops.PARAM, Ops.BUFFER}: return x.replace(dtype=None, arg=replace(x.arg, dtype=dt))
+  return x.replace(dtype=None, src=(with_storage(x.src[0], dt),)+x.src[1:])

@dataclass(frozen=True)
 class Estimates:
@@ -13,7 +13,7 @@ from tinygrad.runtime.autogen.amd.cdna.ins import s_nop as s_nop_cdna

 _arch_map = {"gfx9": "cdna", "gfx10": "rdna3", "gfx11": "rdna3", "gfx12": "rdna4"}
 def assemble_linear(prg:UOp, lin:UOp, arch:str) -> bytes:
-  insts = [u.arg[0] for u in lin.src]
+  insts = [u.arg for u in lin.src]

  # ** scan for max vgpr/sgpr/accvgpr
  max_vgpr, max_sgpr, max_accvgpr = 0, 0, 0
@@ -188,12 +188,10 @@ class CStyleLanguage(Renderer):
    return prefix + self.type_map.get(dtype, dtype.name) + suffix

  def render_type(self, u:UOp): return self._render_dtype(u.dtype, u.max_numel(), u.addrspace, shape=u._shape)
-  def render_ptr(self, u:UOp):
-    # the address of an access, vector-cast if the access reads/writes more lanes than the pointer's scalar type
+  def render_access(self, u:UOp):
    if u.max_numel() > 1 or u.dtype != u.src[0].dtype:
-      return f"(({self._render_dtype(u.dtype, u.max_numel(), u.addrspace, override_ptr=True, shape=u._shape)})({self[u]}))"
-    else: return f"{self[u]}"
-  def render_access(self, u:UOp): return f"*{self.render_ptr(u)}"
+      return f"*(({self._render_dtype(u.dtype, u.max_numel(), u.addrspace, override_ptr=True, shape=u._shape)})({self[u]}))"
+    else: return f"*{self[u]}"
  def render_cast(self, u:UOp, val:str) -> str: return f"({self.render_type(u)})({val})"

  # LEGACY
@@ -247,7 +245,7 @@ class CStyleLanguage(Renderer):
        (u.op in {Ops.STACK, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA"))):
        r[u] = l
      else:
-        if u.op not in {Ops.RANGE, Ops.BUFFER} and u.dtype != dtypes.void:
+        if u.op not in {Ops.RANGE, Ops.STORE, Ops.BUFFER} and u.dtype != dtypes.void:
          l = f"{self.render_type(u)} {r[u]} = {l}" + (";" if u.op is not Ops.SPECIAL else "")
        kernel.append("\n".join("  "*depth + line for line in l.split("\n")))
        if prefix: c[prefix] += 1  # if it was used, increment
@@ -511,9 +509,6 @@ class HIPRenderer(CStyleLanguage):
        (UPat(Ops.CAST, dtypes.float, (UPat.var("y", dtypes.fp8s),), name="x",),
          lambda ctx,x,y: f"__builtin_amdgcn_cvt_f32_{('fp8', 'bf8')[fp8_index(y.dtype)]}((unsigned int){ctx[x.src[0]]}, 0)"),
      ]) + base_rewrite
-    # a LOAD flagged nontemporal renders as the cache-bypassing builtin (only used on global loads)
-    self.string_rewrite = PatternMatcher([(UPat(Ops.LOAD, arg="nontemporal", src=(UPat.var("bidx"),)),
-      lambda ctx,bidx: f"__builtin_nontemporal_load({ctx.render_ptr(bidx)})")]) + self.string_rewrite

  # https://clang.llvm.org/docs/AttributeReference.html#amdgpu-flat-work-group-size
  # NOTE: this makes hlb_cifar10 twice as fast, there may be more gains in tweaking these parameters
@@ -27,7 +27,7 @@ class IselContext:
    return Register(f"v{next(self.reg_n)}", 0, _cons=cons if isinstance(cons, tuple) else (cons,))

 def greg(u:UOp):
-  if u.op in {Ops.NOOP, Ops.AFTER, Ops.BITCAST} and u.src: return greg(u.src[0])
+  if u.op in {Ops.NOOP, Ops.AFTER} and u.src: return greg(u.src[0])
  if isinstance(u.tag, tuple): return u.tag[0]
  return u.tag

@@ -2,10 +2,9 @@
 # allow semicolons to put multiple ops on one line
 import sys, struct, functools
 from typing import cast
-from dataclasses import replace
 from tinygrad.dtype import dtypes, DType, truncate, AddrSpace
 from tinygrad.uop import FastEnum, auto, Ops, GroupOp
-from tinygrad.uop.ops import UOp, UPat, PatternMatcher, promo_dtype
+from tinygrad.uop.ops import UOp, UPat, PatternMatcher
 from tinygrad.renderer.isa import ISARenderer, IselContext, Register, PreRegAllocContext, greg
 from tinygrad.helpers import getenv, NUM_CPU_THREADS, unwrap, Target

@@ -34,11 +33,14 @@ class X86Ops(FastEnum):
  # bitcasts
  VMOVD = auto(); VMOVQ = auto(); VMOVDm = auto(); VMOVQm = auto()
  # comparisons
+  VUCOMISS = auto(); VUCOMISD = auto()
  VCMPSS = auto(); VCMPSD = auto(); VCMPPS = auto(); VCMPPD = auto()
+  VPCMPGTB = auto(); VPCMPGTW = auto(); VPCMPGTD = auto(); VPCMPGTQ = auto()
+  VPCMPEQB = auto(); VPCMPEQW = auto(); VPCMPEQD = auto(); VPCMPEQQ = auto()
  SETNE = auto(); SETE = auto(); SETL = auto(); SETB = auto()
  # where
  CMOVNE = auto(); CMOVE = auto(); CMOVL = auto(); CMOVB = auto()
-  VBLENDVPS = auto(); VBLENDVPD = auto()
+  VPBLENDVB = auto(); VBLENDVPS = auto(); VBLENDVPD = auto()
  # jumps
  JNE = auto(); JE = auto(); JL = auto(); JB = auto(); JGE = auto(); JMP = auto()
  # vectorize / gep
@@ -88,11 +90,12 @@ class X86GroupOp:
                X86Ops.VADDSS, X86Ops.VADDSD, X86Ops.VADDPS, X86Ops.VADDPD, X86Ops.VSUBSS, X86Ops.VSUBSD, X86Ops.VSUBPS, X86Ops.VSUBPD,
                X86Ops.VMULSS, X86Ops.VMULSD, X86Ops.VMULPS, X86Ops.VMULPD, X86Ops.VDIVSS, X86Ops.VDIVSD, X86Ops.VDIVPS, X86Ops.VDIVPD,
                X86Ops.VPADDB, X86Ops.VPADDW, X86Ops.VPADDD, X86Ops.VPADDQ, X86Ops.VPSUBB, X86Ops.VPSUBW, X86Ops.VPSUBD, X86Ops.VPSUBQ,
-                X86Ops.VBLENDVPS, X86Ops.VBLENDVPD, X86Ops.VCMPSS, X86Ops.VCMPSD, X86Ops.VCMPPS, X86Ops.VCMPPD,
+                X86Ops.VPCMPEQB, X86Ops.VPCMPEQW, X86Ops.VPCMPEQD, X86Ops.VPCMPEQQ, X86Ops.VPBLENDVB, X86Ops.VBLENDVPS, X86Ops.VBLENDVPD,
+                X86Ops.VPCMPGTB, X86Ops.VPCMPGTW, X86Ops.VPCMPGTD, X86Ops.VPCMPGTQ, X86Ops.VCMPSS, X86Ops.VCMPSD, X86Ops.VCMPPS, X86Ops.VCMPPD,
                X86Ops.VPMULLW, X86Ops.VPMULLD, X86Ops.VROUNDSS, X86Ops.VROUNDSD, X86Ops.VSQRTSS, X86Ops.VSQRTSD, X86Ops.VINSERTPS,
                X86Ops.VPINSRB, X86Ops.VPINSRW, X86Ops.VPINSRD, X86Ops.VPINSRQ, X86Ops.VPAND, X86Ops.VPOR, X86Ops.VPXOR, X86Ops.VPSLLVD,
                X86Ops.VPSLLVQ, X86Ops.VPSRLVD, X86Ops.VPSRLVQ, X86Ops.VPSRAVD, X86Ops.CMOVNE, X86Ops.CMOVE, X86Ops.CMOVL, X86Ops.CMOVB,
-                X86Ops.VCVTSI2SS, X86Ops.VCVTSI2SD, X86Ops.VCVTSS2SD, X86Ops.VCVTSD2SS, X86Ops.IDIV, X86Ops.DIV}
+                X86Ops.VCVTSI2SS, X86Ops.VCVTSI2SD, X86Ops.VCVTSS2SD, X86Ops.VCVTSD2SS, X86Ops.VUCOMISS, X86Ops.VUCOMISD, X86Ops.IDIV, X86Ops.DIV}

  # X86Ops that can write to memory
  WriteMem = {X86Ops.MOVm, X86Ops.MOVi, X86Ops.VMOVSSm, X86Ops.VMOVSDm, X86Ops.VMOVUPSm, X86Ops.VMOVDm, X86Ops.VMOVQm,
@@ -107,7 +110,7 @@ class X86GroupOp:
  # X86Ops that write flags or can modify flags to undefined values
  WriteFlags = {X86Ops.CMP, X86Ops.CMPi, X86Ops.ADD, X86Ops.ADDi, X86Ops.SUB, X86Ops.SUBi, X86Ops.IMUL, X86Ops.IMULi, X86Ops.IDIV, X86Ops.DIV,
                X86Ops.SHL, X86Ops.SHLi, X86Ops.SHR, X86Ops.SHRi, X86Ops.SAR, X86Ops.SARi, X86Ops.AND, X86Ops.ANDi, X86Ops.XOR, X86Ops.XORi,
-                X86Ops.OR, X86Ops.ORi}
+                X86Ops.OR, X86Ops.ORi, X86Ops.VUCOMISS, X86Ops.VUCOMISD}

  # X86Ops whose first src is the rm field
  Rm1st = ReadMem1st | (ReadMem2nd & TwoAddress) | {X86Ops.VPSRLDQ}
@@ -115,6 +118,8 @@ class X86GroupOp:
  # X86Ops whose second src is the rm field
  Rm2nd = ReadMem2nd

+  All = set(X86Ops)
+
 # ***** X86 legalization *****

 extra_matcher = PatternMatcher([
@@ -140,14 +145,14 @@ extra_matcher = PatternMatcher([
  # float16 alus are done in float32
  (UPat(GroupOp.ALU, dtypes.float16, name="x"), lambda x: UOp(x.op,
   src=tuple(s.cast(dtypes.float) if s.dtype != dtypes.bool else s for s in x.src)).cast(x.dtype)),
-  (UPat(GroupOp.Comparison, src=[UPat(dtype=dtypes.float16), UPat()], name="x"),
-   lambda x: UOp(x.op, src=tuple(s.cast(dtypes.float32) for s in x.src)).cast(x.dtype)),
+  (UPat(GroupOp.Comparison, src=(UPat.var("a", dtypes.float16), UPat.var("b")), name="x"),
+   lambda x,a,b: UOp(x.op, src=(a.cast(dtypes.float32), b.cast(dtypes.float32))).cast(x.dtype)),
  # no cmpne for packed ints, y != x => !(y==x)
  (UPat(Ops.CMPNE, src=(UPat.var("y", dtypes.ints), UPat.var("x")), name="cmp"),
   lambda y,x,cmp: UOp(Ops.CMPEQ, src=(y,x))^True if y.max_numel() > 1 else None),
-  # a float WHERE blends at the width of its value, so it needs a comparison at that width to make the mask
+  # float WHERE needs a mask unless its comparison already has a float operand
  (UPat.var("m", dtypes.bool).where(UPat.var("a", dtypes.floats+(dtypes.weakfloat,)), UPat.var("b")).named("w"),
-   lambda m,a,b,w: m.cast(w.dtype).ne(0).where(a, b) if w.dtype in dtypes.floats and promo_dtype(m.src) is not w.dtype else None),
+   lambda m,a,b,w: m.cast(w.dtype).ne(0).where(a, b) if w.dtype in dtypes.floats and not dtypes.is_float(m.src[0].dtype) else None),
  # rewrite -x -> 0 - x
  (UPat(Ops.NEG, name="x"), lambda x: UOp(Ops.SUB, src=(x.const_like(0),) + x.src)),
  # TODO: add support for mod, requires support for accessing the 2nd+ reg of a multi output instruction
@@ -161,34 +166,36 @@ def scratch_buffer(elem_dt:DType, count:int, slot:int) -> UOp:

 def gated_load(ctx, addr:UOp, alt:UOp, gate:UOp, x:UOp):
  local = scratch_buffer(addr.src[0].dtype, x.max_numel(), next(ctx))
-  local_idx = local.index(UOp.cconst(0, dtypes.int32))
-  # the AFTER orders the load after the scratch store
-  sel = gate.where(addr, local_idx)
-  return UOp(Ops.AFTER, src=(sel, (local_idx if x.max_numel() == 1 else local).store(alt))).load()
+  local_idx = local.index(UOp.cconst(0, dtypes.int32), dtype=dtypes.uint64)
+  # the selected address is a 64bit value, the AFTER orders the load after the scratch store and carries the element dtype for the encoder
+  sel = gate.where(addr.replace(dtype=dtypes.uint64), local_idx)
+  ptr = UOp(Ops.AFTER, addr.dtype, (sel, (local_idx if x.max_numel() == 1 else local).store(alt)))
+  return ptr.load(dtype=x.dtype)

 def gated_store(addr:UOp, gate:UOp, val:UOp):
  local = scratch_buffer(addr.src[0].dtype, val.max_numel(), -1)
-  sel = gate.where(addr, local.index(UOp.cconst(0, dtypes.int32)))
-  return UOp(Ops.AFTER, src=(sel,)).store(val)
-
-# a gate the flags can be picked with, or the bool compared to zero that replaces one they can't: only an integer
-# comparison sets the flags, see cmp. NOTE: the 0 is int so the bool zero-extends and compares as int (a byte compare renders
-# different kernels)
-def flag_gate(m:UOp) -> UOp|None:
-  return None if m.op in GroupOp.Comparison and m.src[0].dtype not in dtypes.floats else m.ne(UOp.cconst(0, dtypes.int))
+  sel = gate.where(addr.replace(dtype=dtypes.uint64), local.index(UOp.cconst(0, dtypes.int32), dtype=dtypes.uint64))
+  return UOp(Ops.AFTER, addr.dtype, (sel,)).store(val)

 # legalize the new style graph for isel. NOTE: this runs after the spec is verified, some of these rewrites violate it
 pre_isel_matcher = PatternMatcher([
-  # widening a scalar uint32 is free, the 32bit write that produced it already zeroed the upper half
-  (UPat.var("y", dtypes.uint32).cast(dtypes.int64s, name="x"), lambda y,x: x.replace(op=Ops.BITCAST) if y.max_numel() == 1 else None),
+  # noop casts: zero extending scalar 32bit int, same-width signed/unsigned, narrowing scalar int
+  (UPat.var("y", dtypes.uint32).cast(dtypes.int64s, name="x"), lambda y,x: x.replace(op=Ops.NOOP, arg=None) if y.max_numel() == 1 else None),
  (UPat.var("y", dtypes.ints+(dtypes.bool,)).cast(dtypes.ints, name="x"),
-   lambda y,x: x.replace(op=Ops.BITCAST) if x.dtype.itemsize == y.dtype.itemsize else None),
+   lambda y,x: x.replace(op=Ops.NOOP, arg=None) if x.dtype.itemsize == y.dtype.itemsize and y.max_numel() == 1 else None),
+  (UPat.var("y", dtypes.ints).cast(dtypes.ints, name="x"),
+   lambda y,x: x.replace(op=Ops.NOOP, arg=None) if x.dtype.itemsize < y.dtype.itemsize and y.max_numel() == 1 else None),
+  # bitcasts between scalar floats and ints are real, rest are noops
+  (UPat.var("y").bitcast().named("x"), lambda y,x: None if y.dtype in dtypes.floats and x.dtype in dtypes.ints or \
+   y.dtype in dtypes.ints and x.dtype in dtypes.floats else x.replace(op=Ops.NOOP, arg=None)),
  # gated load/store become a conditional move on the address, the load/store are unconditional
  (UPat((Ops.INDEX, Ops.SHRINK), name="addr").load(UPat.var("alt"), UPat.var("gate"), name="x"), gated_load),
  (UPat((Ops.INDEX, Ops.SHRINK), name="addr").store(UPat.var("val"), UPat.var("gate")), gated_store),
-  # a conditional backedge picks with the flags, and so does the cmove, which is legalized in isel
-  (UPat(Ops.END, src=(UPat(), UPat(), UPat.var("m", dtypes.bool)), name="x"),
-   lambda m,x: x.replace(src=x.src[:2]+(g,)) if (g:=flag_gate(m)) is not None else None),
+  # TODO: remove this once we allow all flag producing ops in cmove
+  # if gate in scalar int cmove is not a comparison need to add one to set the flag
+  # NOTE: the 0 is int so the bool gate zero-extends and compares as int (a byte compare renders different kernels)
+  (UPat.var("m", dtypes.bool).where(UPat.var("a"), UPat.var("b")),
+   lambda m,a,b: m.ne(UOp.cconst(0, dtypes.int)).where(a,b) if m.op not in GroupOp.Comparison else None),
 ])

 # ***** X86 registers *****
@@ -210,29 +217,27 @@ CALLEE_SAVED = (RBX, RBP, GPR[12], GPR[13], GPR[14], GPR[15]) + ((RSI, RDI) + XM

 reg_strs = {"rax": {4:"eax", 2:"ax", 1:"al"}, "rcx": {4:"ecx", 2:"cx", 1:"cl"}, "rdx": {4:"edx", 2:"dx", 1:"dl"}, "rbx": {4:"ebx", 2:"bx", 1:"bl"},
        "rsp": {4:"esp", 2:"sp", 1:"spl"}, "rbp": {4:"ebp", 2:"bp", 1:"bpl"}, "rsi": {4:"esi", 2:"si", 1:"sil"}, "rdi": {4:"edi", 2:"di", 1:"dil"},
-        **{f"r{i}": {4:f"r{i}d", 2:f"r{i}w", 1:f"r{i}b"} for i in range(8, 16)}}
+        **{f"r{i}": {4:f"r{i}d", 2:f"r{i}w", 1:f"r{i}b"} for i in range(8, 16)}, **{f"xmm{i}": {64:f"zmm{i}", 32:f"ymm{i}"} for i in range(16)}}

 # ***** X86 instruction selection *****
 def base(x:UOp, i:int) -> UOp: return s.src[0] if (s:=x.src[i]).op is Ops.INDEX else s
 def lane(x:UOp, i:int) -> int: return s.src[1].src[0].val if (s:=x.src[i]).op is Ops.INDEX else 0
 def to_int(dt:DType): return {dtypes.float16: dtypes.int16, dtypes.float32: dtypes.int32, dtypes.float64: dtypes.int64}[dt]
-def def_reg(dt:DType, reg:Register|None=None) -> UOp: return UOp(Ops.INS, arg=(X86Ops.DEFINE, dt), tag=None if reg is None else (reg,))
+def def_reg(dt:DType, reg:Register|None=None) -> UOp: return UOp(Ops.INS, dt, arg=X86Ops.DEFINE, tag=None if reg is None else (reg,))
 def imm(dt:DType, v:int) -> UOp: return UOp.cconst(truncate[dt](v), dt).rtag()
 def to_imm(c:UOp) -> UOp|None:
  if not (c.op is Ops.CAST and (v:=c.src[0]).op is Ops.CONST): return None
  if c.dtype in dtypes.int64s: return imm(dtypes.int32, v.val) if not v.overflows(dtypes.int32) else None
  if c.dtype in dtypes.ints+(dtypes.bool,): return imm(c.dtype, v.val)
  return None
-# the flag path, which only an integer comparison can take: an x86 float compare sets carry, zero and parity together when an
-# operand is NaN, so a NaN reads as "below" and as "equal", and it clears sign and overflow, so nothing reads as "less"
 def cmp(x:UOp) -> UOp:
-  if x.src[0].dtype in dtypes.floats: raise RuntimeError(f"no flag compare for {x.src[0].dtype}, a float gate must be a mask")
+  if x.src[0].dtype is dtypes.float32: return x.ins(X86Ops.VUCOMISS, dtype=dtypes.void)
+  if x.src[0].dtype is dtypes.float64: return x.ins(X86Ops.VUCOMISD, dtype=dtypes.void)
  return x.ins(X86Ops.CMP, dtype=dtypes.void) if (i:=to_imm(x.src[1])) is None else x.ins(X86Ops.CMPi, dtype=dtypes.void, src=(x.src[0], i))
-# comparisons that produce masks, the mask has the width of the operands
-def mask(x:UOp) -> UOp:
-  dt, v = x.src[0].dtype, imm(dtypes.uint8, {Ops.CMPLT: 1, Ops.CMPNE: 4, Ops.CMPEQ: 0}[x.op])
-  if dt is dtypes.float32: return x.ins(X86Ops.VCMPSS if x.max_numel() == 1 else X86Ops.VCMPPS, dtype=dt, src=x.src + (v,))
-  return x.ins(X86Ops.VCMPSD if x.max_numel() == 1 else X86Ops.VCMPPD, dtype=dt, src=x.src + (v,))
+def vcmp(x:UOp) -> UOp:
+  v = imm(dtypes.uint8, {Ops.CMPLT: 1, Ops.CMPNE: 4, Ops.CMPEQ: 0}[x.op])
+  if x.dtype is dtypes.float32: return x.ins(X86Ops.VCMPSS if x.max_numel() == 1 else X86Ops.VCMPPS, src=x.src + (v,))
+  return x.ins(X86Ops.VCMPSD if x.max_numel() == 1 else X86Ops.VCMPPD, src=x.src + (v,))

 # vinsertps xmm2, xmm0, xmm1, imm
 # inserts any 32 bit element in xmm1 into any position in xmm0 according to immm, result is written to xmm2
@@ -245,9 +250,9 @@ def vinsertps(x:UOp) -> UOp:

 # vpinsq xmm2, xmm0, rax, imm
 # inserts element in rax into any position in xmm0, result is written to xmm2 according to imm
-def vpins(x:UOp, srcs:tuple[UOp, ...]) -> UOp:
+def vpins(x:UOp) -> UOp:
  op = {1: X86Ops.VPINSRB, 2: X86Ops.VPINSRW, 4: X86Ops.VPINSRD, 8: X86Ops.VPINSRQ}[x.dtype.itemsize]
-  return functools.reduce(lambda ret,i: x.ins(op, src=(ret, srcs[i], imm(dtypes.uint8, i))), range(len(srcs)), def_reg(x.dtype))
+  return functools.reduce(lambda ret,i: x.ins(op, src=(ret, x.src[i], imm(dtypes.uint8, i))), range(len(x.src)), def_reg(x.dtype))

 # we don't call ctx.vreg on the srcs to avoid duplicates, a rewrite will assign the tuple of valid registers to a vreg
 def idiv(ctx:IselContext, x:UOp) -> UOp:
@@ -257,8 +262,8 @@ def idiv(ctx:IselContext, x:UOp) -> UOp:
  elif x.dtype in dtypes.uints: ext = [x.ins(X86Ops.MOVi, src=(imm(min(dtypes.uint32, x.dtype), 0),), tag=(RDX,))]
  else: ext = [x.ins(X86Ops.SARi, src=(x.src[0], imm(dtypes.uint8, x.dtype.itemsize * 8 - 1)), tag=(RDX,))]
  # for 8bit need to zero/sign extend al to ah
-  if x.dtype is dtypes.uint8: dividend = UOp(Ops.INS, arg=(X86Ops.MOVZX, dtypes.int16), src=(x.src[0],), tag=(RAX,))
-  elif x.dtype is dtypes.int8: dividend = UOp(Ops.INS, arg=(X86Ops.MOVSX, dtypes.int16), src=(x.src[0],), tag=(RAX,))
+  if x.dtype is dtypes.uint8: dividend = UOp(Ops.INS, arg=X86Ops.MOVZX, dtype=dtypes.int16, src=(x.src[0],), tag=(RAX,))
+  elif x.dtype is dtypes.int8: dividend = UOp(Ops.INS, arg=X86Ops.MOVSX, dtype=dtypes.int16, src=(x.src[0],), tag=(RAX,))
  else: dividend = x.ins(X86Ops.MOV, src=(x.src[0],), tag=(RAX,))
  # divisor can't be in rax or rdx
  divisor = x.ins(X86Ops.MOV, src=(x.src[1],), tag=tuple(r for r in WGPR if r not in (RAX, RDX)))
@@ -277,7 +282,7 @@ def shift(x:UOp, op:X86Ops) -> UOp:
 # it is materialized as an immediate so the address stays correct if the base register is ever spilled and refilled
 def fold_address(x:UOp) -> tuple[UOp, UOp, UOp, UOp]:
  def _disp(v:int) -> UOp: return imm(dtypes.int32 if abs(v) > dtypes.int8.max else dtypes.int8, v)
-  def _cast(v:UOp) -> UOp: return v.cast(dtypes.int64) if v.vmin < 0 else v.cast(dtypes.uint32) if v.dtype.itemsize < 4 else v
+  def _cast(v:UOp) -> UOp: return v.cast(dtypes.int64) if v.vmin < 0 else v
  if x.op not in {Ops.INDEX, Ops.SHRINK}: return (x, UOp(Ops.NOOP), _disp(0), imm(dtypes.uint8, x.dtype.itemsize))
  base, idx = x.src[0], x.src[1]
  # buffers are indexed by element, everything else (the stack pointer) by byte
@@ -288,19 +293,15 @@ def fold_address(x:UOp) -> tuple[UOp, UOp, UOp, UOp]:
  if idx.op is Ops.CAST and idx.src[0].op is Ops.CONST: return (base, UOp(Ops.NOOP), _disp(idx.src[0].val * scale), sz)
  return (base, _cast(idx), _disp(0), sz)

-# addresses are 64bit values
-def lea(x:UOp) -> UOp: return x.ins(X86Ops.LEA, dtype=dtypes.uint64, src=fold_address(x))
-
 def abi(ctx:IselContext, x:UOp) -> UOp|None:
  if isinstance(x.tag, tuple): return None
  i = ctx.func_args.index(x)
  # buffer params hold addresses, their value moves as a 64bit int
  dt = dtypes.uint64 if x.op is Ops.PARAM and x.arg.addrspace is AddrSpace.GLOBAL else x.dtype
-  arg = replace(x.arg, dtype=dt) if x.op is Ops.PARAM else x.arg
  # the shape srcs of a PARAM are not values, tag them so they aren't materialized into registers
-  def _reg_arg(r:Register) -> tuple[UOp, ...]: return (x.replace(arg=arg, src=tuple(s.rtag() for s in x.src), tag=(r,)),)
+  def _reg_arg(r:Register) -> tuple[UOp, ...]: return (x.replace(dtype=dt, src=tuple(s.rtag() for s in x.src), tag=(r,)),)
  def _stack_arg(disp:int):
-    return (def_reg(dtypes.uint64, RSP), UOp(Ops.NOOP), UOp(Ops.INS, arg=(X86Ops.FRAME_INDEX, dtypes.int32), tag=disp), imm(dtypes.uint8, 8))
+    return (def_reg(dtypes.uint64, RSP), UOp(Ops.NOOP), UOp(Ops.INS, arg=X86Ops.FRAME_INDEX, dtype=dtypes.int32, tag=disp), imm(dtypes.uint8, 8))
  if sys.platform == "win32": src = _reg_arg((RCX, RDX, GPR[8], GPR[9])[i]) if i < 4 else _stack_arg((i-3)*8+32)
  else: src = _reg_arg((RDI, RSI, RDX, RCX, GPR[8], GPR[9])[i]) if i < 6 else _stack_arg((i-5)*8)
  # this move "cleanses" the abi register constraint
@@ -311,7 +312,7 @@ GPR_DEST_OPS = {X86Ops.VPEXTRB, X86Ops.VPEXTRW, X86Ops.VPEXTRD, X86Ops.VPEXTRQ,
 XMM_OPS = {op for op in X86Ops if op.name.startswith('V')} - GPR_DEST_OPS

 def _is_vec_xmm(y: UOp) -> bool:
-  return (y.op is Ops.INS and y.arg[0] in XMM_OPS) or (y.op not in (Ops.BUFFER, Ops.PARAM, Ops.AFTER, Ops.INS) and y.max_numel() > 1)
+  return (y.op is Ops.INS and y.arg in XMM_OPS) or (y.op not in (Ops.BUFFER, Ops.PARAM, Ops.AFTER, Ops.INS) and y.max_numel() > 1)

 def _xmm_sz(x: UOp) -> X86Ops:
  bits = x.max_numel() * x.dtype.itemsize
@@ -327,10 +328,10 @@ def _xmm_sz_m(x: UOp) -> X86Ops:

 def alloc_vregs(ctx:IselContext, x:UOp) -> UOp|None:
  # register placeholders with real registers
-  if x.op is Ops.INS and x.arg[0] is X86Ops.DEFINE and x.tag is not None: return None
-  if x.op is Ops.INS and x.arg[0] is X86Ops.LOOP_CMP: return None
+  if x.arg is X86Ops.DEFINE and x.tag is not None: return None
+  if x.arg is X86Ops.LOOP_CMP: return None
  # this is an immediate
-  if x.op is Ops.INS and x.arg[0] is X86Ops.FRAME_INDEX: return None
+  if x.arg is X86Ops.FRAME_INDEX: return None
  # no register definition
  if x.dtype is dtypes.void: return None
  # already allocated vregs
@@ -339,16 +340,18 @@ def alloc_vregs(ctx:IselContext, x:UOp) -> UOp|None:
  defs = []
  if isinstance(x.tag, tuple): defs = [ctx.vreg(x.tag)]
  elif x.op is Ops.BUFFER: defs = [ctx.vreg(WGPR)]
-  elif x.dtype in dtypes.floats or (x.op is Ops.INS and x.arg[0] in XMM_OPS) or x.max_numel() > 1: defs = [ctx.vreg(XMM)]
+  elif x.dtype in dtypes.floats or (x.op is Ops.INS and x.arg in XMM_OPS) or x.max_numel() > 1: defs = [ctx.vreg(XMM)]
  elif x.dtype in dtypes.ints+(dtypes.bool,): defs = [ctx.vreg(WGPR)]
  # TODO: add this once the scheduler can track register pressure
-  # if x.arg[0] in X86GroupOp.WriteFlags: defs.append(ctx.vreg(RFLAGS))
+  # if x.arg in X86GroupOp.WriteFlags: defs.append(ctx.vreg(RFLAGS))
  # the size src of a BUFFER is not a value, tag it so it isn't materialized into a register
  if x.op is Ops.BUFFER: return x.replace(src=tuple(s.rtag() for s in x.src), tag=tuple(defs))
  return x.replace(tag=tuple(defs))

 isel_matcher = PatternMatcher([
  # **** Op -> Op ****
+  # cast of void is a noop
+  (UPat.var("y").cast(name="x"), lambda y,x: y if y.dtype == dtypes.void else None),
  # range is lowered to acc, cmp, jmp after regalloc
  (UPat(Ops.RANGE, src=(UPat.cvar("c").cast(),), allow_any_len=True, name="x"), lambda c,x: x.replace(src=(imm(x.dtype, c.val),) + x.src[1:])),
  (UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(tag=(ctx.vreg(WGPR),)) if not isinstance(x.tag, tuple) else None),
@@ -360,29 +363,25 @@ isel_matcher = PatternMatcher([
  # so regalloc builds the prologue/epilogue naturally
  (UPat(Ops.SINK, name="x"), lambda x:
   x.replace(src=(x.ins(X86Ops.RET, src=x.src + tuple(def_reg(dtypes.uint64 if r in GPR else dtypes.float64, r) for r in CALLEE_SAVED)),)) \
-    if not x.src or x.src[0].op is not Ops.INS or x.src[0].arg[0] is not X86Ops.RET else None),
+    if not x.src or x.src[0].arg is not X86Ops.RET else None),
  # function abi constraints
  (UPat((Ops.PARAM, Ops.SPECIAL), name="x"), abi),
-  # conditional moves between addresses, lea both srcs
-  (UPat.var("m").where(UPat((Ops.INDEX, Ops.SHRINK), name="a"), UPat((Ops.INDEX, Ops.SHRINK), name="b")), lambda m,a,b:
-   m.where(lea(a), lea(b)) if not _is_vec_xmm(a.src[0]) else None),
  # constants that can't be immediates, move them to registers
  (UPat.cvar("c").cast(dtypes.int64s, name="x"), lambda c,x: x.ins(X86Ops.MOVABS, src=(imm(x.dtype, c.val),)) if not x.tag else None),
  (UPat.cvar("c").cast(dtypes.ints+(dtypes.bool,), name="x"), lambda c,x: x.ins(X86Ops.MOVi, src=(imm(x.dtype, c.val),)) if not x.tag else None),
  (UPat.cvar("c").cast(dtypes.floats, name="x"), lambda c,x:
   UOp.cconst(struct.unpack((dt:=to_int(x.dtype)).fmt, struct.pack(x.dtype.fmt, c.val))[0], dt).bitcast(x.dtype) if not x.tag else None),
-  # conditional moves that use masks, the mask has the width of the values
-  (UPat(GroupOp.Comparison, src=(UPat(dtype=dtypes.float32), UPat()), name="m").where(UPat.var("a", dtypes.float32), UPat.var("b")), lambda m,a,b:
-   a.ins(X86Ops.VBLENDVPS, src=(b, a, mask(m)))),
-  (UPat(GroupOp.Comparison, src=(UPat(dtype=dtypes.float64), UPat()), name="m").where(UPat.var("a", dtypes.float64), UPat.var("b")), lambda m,a,b:
-   a.ins(X86Ops.VBLENDVPD, src=(b, a, mask(m)))),
+  # conditional moves that use masks NOTE: these currently assume a mask producing cmp exists
+  (UPat.var("m").where(UPat.var("a", dtypes.int8s+dtypes.int16s+dtypes.int32s+(dtypes.int64,)), UPat.var("b")), lambda m,a,b:
+   a.ins(X86Ops.VPBLENDVB, src=(b, a, m.replace(dtype=m.src[0].dtype))) if a.max_numel() > 1 else None),
+  (UPat.var("m").where(UPat.var("a", dtypes.float32), UPat.var("b")), lambda m,a,b:
+   a.ins(X86Ops.VBLENDVPS, src=(b, a, m.replace(dtype=m.src[0].dtype)))),
+  (UPat.var("m").where(UPat.var("a", dtypes.float64), UPat.var("b")), lambda m,a,b:
+   a.ins(X86Ops.VBLENDVPD, src=(b, a, m.replace(dtype=m.src[0].dtype)))),
  # in this case we have a mask producing comparison whose user expects a bool, so we convert to bool
  (UPat(GroupOp.Comparison, dtypes.bool, (UPat.var("y", (dtypes.float32, dtypes.float64)), UPat()), name="x"), lambda y,x:
-   UOp(Ops.AND, src=(mask(x).bitcast(dt:=to_int(y.dtype)), UOp.cconst(1, dt))).bitcast(dtypes.bool)),
+   UOp(Ops.AND, src=(x.replace(dtype=y.dtype).bitcast(dt:=to_int(y.dtype)), UOp.cconst(1, dt))).f(Ops.NOOP, dtype=dtypes.bool)),
  # conditional moves that use flags
-  # TODO: remove this once we allow all flag producing ops in cmove
-  # the blends took every float gate a mask can serve, so a gate that is still not an integer comparison becomes one here
-  (UPat.var("m", dtypes.bool).where(UPat.var("a"), UPat.var("b")), lambda m,a,b: g.where(a, b) if (g:=flag_gate(m)) is not None else None),
  (UPat(Ops.CMPLT, src=(UPat(dtype=dtypes.sints), UPat()), name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b:
   a.ins(X86Ops.CMOVL, src=(b, a, cmp(m)))),
  (UPat(Ops.CMPLT, name="m").where(UPat.var("a"), UPat.var("b")), lambda m,a,b: a.ins(X86Ops.CMOVB, src=(b, a, cmp(m)))),
@@ -398,6 +397,16 @@ isel_matcher = PatternMatcher([
  (UPat(Ops.CMPLT, dtypes.bool, name="x"), lambda x: x.ins(X86Ops.SETL, src=(cmp(x),))),
  (UPat(Ops.CMPEQ, dtypes.bool, name="x"), lambda x: x.ins(X86Ops.SETE, src=(cmp(x),))),
  (UPat(Ops.CMPNE, dtypes.bool, name="x"), lambda x: x.ins(X86Ops.SETNE, src=(cmp(x),))),
+  # comparisons that produce masks (these aren't bool dtype)
+  (UPat(GroupOp.Comparison, src=(UPat(dtype=(dtypes.float32, dtypes.float64)), UPat()), name="x"), vcmp),
+  (UPat(Ops.CMPEQ, src=(UPat(dtype=dtypes.int8s), UPat()), name="x"), lambda x: x.ins(X86Ops.VPCMPEQB)),
+  (UPat(Ops.CMPEQ, src=(UPat(dtype=dtypes.int16s), UPat()), name="x"), lambda x: x.ins(X86Ops.VPCMPEQW)),
+  (UPat(Ops.CMPEQ, src=(UPat(dtype=dtypes.int32s), UPat()), name="x"), lambda x: x.ins(X86Ops.VPCMPEQD)),
+  (UPat(Ops.CMPEQ, src=(UPat(dtype=dtypes.int64s), UPat()), name="x"), lambda x: x.ins(X86Ops.VPCMPEQQ)),
+  (UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int8s), UPat.var("b")), name="x"), lambda a,b,x: x.ins(X86Ops.VPCMPGTB, src=(b, a))),
+  (UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int16s), UPat.var("b")), name="x"), lambda a,b,x: x.ins(X86Ops.VPCMPGTW, src=(b, a))),
+  (UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int32s), UPat.var("b")), name="x"), lambda a,b,x: x.ins(X86Ops.VPCMPGTD, src=(b, a))),
+  (UPat(Ops.CMPLT, src=(UPat.var("a", dtypes.int64s), UPat.var("b")), name="x"), lambda a,b,x: x.ins(X86Ops.VPCMPGTQ, src=(b, a))),
  # float unary
  (UPat.var("y", dtypes.float32).sqrt().named("x"), lambda y,x: x.ins(X86Ops.VSQRTSS, src=(y, y)) if x.max_numel() == 1 else x.ins(X86Ops.VSQRTPS)),
  (UPat.var("y", dtypes.float64).sqrt().named("x"), lambda y,x: x.ins(X86Ops.VSQRTSD, src=(y, y)) if x.max_numel() == 1 else x.ins(X86Ops.VSQRTPD)),
@@ -406,9 +415,10 @@ isel_matcher = PatternMatcher([
  (UPat.var("y", dtypes.float64).trunc().named("x"), lambda y,x:
   x.ins(X86Ops.VROUNDSD, src=(y, y, imm(dtypes.uint8, 3))) if x.max_numel() == 1 else x.ins(X86Ops.VROUNDPD, src=(y, imm(dtypes.uint8, 3)))),
  # for float16 we route the srcs through gprs, this is suboptimal for values in xmms, in that case we want vpunpcklwd
-  (UPat(Ops.STACK, dtypes.float16, name="x"), lambda x: vpins(x, tuple(s.bitcast(dtypes.int16) for s in x.src))),
+  (UPat(Ops.STACK, dtypes.float16, name="x"), lambda x:
+   vpins(x.replace(src=tuple(s.bitcast(dtypes.int16) for s in x.src)))),
  (UPat(Ops.STACK, dtypes.float32, name="x"), vinsertps),
-  (UPat(Ops.STACK, dtypes.ints+(dtypes.bool,), name="x"), lambda x: vpins(x, x.src)),
+  (UPat(Ops.STACK, dtypes.ints+(dtypes.bool,), name="x"), vpins),
  # INDEX on a vector register value extracts a single element
  (UPat.var("y", dtypes.int8s+(dtypes.bool,)).index(UPat.cvar("c").cast(), name="x"),
   lambda y,c,x: x.ins(X86Ops.VPEXTRB, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
@@ -490,12 +500,11 @@ isel_matcher = PatternMatcher([
  (UPat.var("y", dtypes.float64).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSD2SS, src=(y, y))),
  (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SS, src=(def_reg(x.dtype), y))),
  (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float64, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SD, src=(def_reg(x.dtype), y))),
-  (UPat(dtype=(dtypes.uint8, dtypes.uint16, dtypes.bool)).cast(dtypes.ints, name="x"), lambda x:
+  (UPat(dtype=dtypes.uints+(dtypes.bool,)).cast(dtypes.ints, name="x"), lambda x:
   x.ins(X86Ops.MOVZX) if x.max_numel() == 1 and x.src[0].dtype.itemsize < x.dtype.itemsize else None),
  (UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.MOVSXD) if x.max_numel() == 1 else None),
  (UPat(dtype=dtypes.sints).cast(dtypes.ints, name="x"), lambda x:
   x.ins(X86Ops.MOVSX) if x.max_numel() == 1 and x.src[0].dtype.itemsize < x.dtype.itemsize else None),
-  (UPat(dtype=dtypes.ints).cast(dtypes.ints, name="x"), lambda x: x.ins(X86Ops.MOV) if x.max_numel() == 1 else None),
  (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBW)),
  (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBD)),
  (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBQ)),
@@ -510,13 +519,14 @@ isel_matcher = PatternMatcher([
  (UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXDQ)),
  # bitcasts between scalar floats and ints
  (UPat.var("y", dtypes.float16).bitcast(dtypes.int16s).named("x"), lambda y,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, 0)))),
-  (UPat(dtype=dtypes.int16s).bitcast(dtypes.float16).named("x"), lambda x: vpins(x, x.src)),
+  (UPat(dtype=dtypes.int16s).bitcast(dtypes.float16).named("x"), vpins),
  (UPat(dtype=dtypes.int32s).bitcast(dtypes.float32).named("x"), lambda x: x.ins(X86Ops.VMOVD)),
  (UPat(dtype=dtypes.int64s).bitcast(dtypes.float64).named("x"), lambda x: x.ins(X86Ops.VMOVQ)),
  (UPat(dtype=dtypes.float32).bitcast(dtypes.int32s).named("x"), lambda x: x.ins(X86Ops.VMOVDm)),
  (UPat(dtype=dtypes.float64).bitcast(dtypes.int64s).named("x"), lambda x: x.ins(X86Ops.VMOVQm)),
  # index on a buffer (or the stack pointer) computes an address, addresses are 64bit values
-  (UPat((Ops.INDEX, Ops.SHRINK), name="x"), lambda x: lea(x) if not _is_vec_xmm(x.src[0]) else None),
+  (UPat((Ops.INDEX, Ops.SHRINK), name="x"),
+   lambda x: x.ins(X86Ops.LEA, dtype=dtypes.uint64, src=fold_address(x)) if not _is_vec_xmm(x.src[0]) else None),
  # TODO: fuse stores, very few cases -- store cmp becomes setcc, store gep int becomes vpextr, store bitcast to int becomes vmovd/q
  # copy, load, store
  # NOTE: copy here violates the spec, it only happens post register allocation when a reg to reg move needs to be inserted
@@ -545,7 +555,7 @@ isel_matcher = PatternMatcher([
 # so we rematerialize. This is different from rematerialization you might want to do in regalloc because it is not optional,
 # regalloc shouldn't rematerialize if a src of the instruction is dead, but here you need to as there's no fallback load from stack
 def flag_rematerialize(ctx:PreRegAllocContext, x:UOp):
-  flag_def = x if x.op in (Ops.RANGE, Ops.END) or x.arg[0] in X86GroupOp.WriteFlags else x.src[-1] if x.arg[0] in X86GroupOp.ReadFlags else None
+  flag_def = x if x.arg in X86GroupOp.WriteFlags or x.op in (Ops.RANGE, Ops.END) else x.src[-1] if x.arg in X86GroupOp.ReadFlags else None
  if flag_def is None: return None
  if ctx.lock is not None and ctx.lock is not flag_def: ctx.clobbered.add(ctx.lock)
  ctx.lock = flag_def
@@ -561,19 +571,19 @@ pre_regalloc_matcher = PatternMatcher([
 # TODO: control flow should be overhauled so that this isn't necessary
 def lower_range(ctx, x:UOp) -> tuple[UOp, list[UOp]]:
  loop_label = "_".join(str(i) for i in x.arg[:-1])
-  label = UOp(Ops.INS, arg=(X86Ops.LABEL, dtypes.void), tag=f".LOOP_{loop_label}")
+  label = UOp(Ops.INS, arg=X86Ops.LABEL, tag=f".LOOP_{loop_label}")
  # loop, cmp on backedge all we need is a jmp tag
  if x.dtype is dtypes.void: return (label, [label])
  else:
    acc = x.ins(X86Ops.MOVi, src=(imm(x.dtype, 0),) + x.src[1:])
-    cmp = UOp(Ops.INS, arg=(X86Ops.CMPi if x.src[0].op is Ops.CAST else X86Ops.CMP, dtypes.void), src=(acc, x.src[0]))
-    jump_out = UOp(Ops.INS, arg=(X86Ops.JGE, dtypes.void), src=(cmp,), tag=f".LOOP_OUT_{loop_label}")
+    cmp = UOp(Ops.INS, arg=X86Ops.CMPi if x.src[0].op is Ops.CAST else X86Ops.CMP, src=(acc, x.src[0]))
+    jump_out = UOp(Ops.INS, arg=X86Ops.JGE, src=(cmp,), tag=f".LOOP_OUT_{loop_label}")
    ctx.loop_label[acc] = loop_label
    return (acc, [acc, label, cmp, jump_out])

 def lower_end(ctx, x:UOp) -> tuple[UOp, list[UOp]]:
-  end_label = UOp(Ops.INS, arg=(X86Ops.LABEL, dtypes.void), tag=f".LOOP_OUT_{ctx.loop_label[x.src[1]]}")
-  jmp = UOp(Ops.INS, arg=(X86Ops.JMP, dtypes.void), tag=f".LOOP_{ctx.loop_label[x.src[1]]}")
+  end_label = UOp(Ops.INS, arg=X86Ops.LABEL, tag=f".LOOP_OUT_{ctx.loop_label[x.src[1]]}")
+  jmp = UOp(Ops.INS, arg=X86Ops.JMP, tag=f".LOOP_{ctx.loop_label[x.src[1]]}")
  inc = x.src[1].ins(X86Ops.ADDi, src=(imm(x.src[1].dtype, 1),))
  return (inc, [inc, jmp, end_label])

@@ -585,16 +595,16 @@ def lower_loop(ctx, x:UOp) -> tuple[UOp, list[UOp]]:
 # final rewrite to match the isa spec
 post_regalloc_matcher = PatternMatcher([
  # rewrite FRAME_INDEX to IMM now that the stack size is known
-  (UPat(Ops.INS, name="x"), lambda ctx,x: (nx:=UOp.cconst(ctx.stack_size + x.tag, x.dtype), [nx]) if x.arg[0] is X86Ops.FRAME_INDEX else None),
+  (UPat(Ops.INS, arg=X86Ops.FRAME_INDEX, name="x"), lambda ctx,x: (nx:=UOp.cconst(ctx.stack_size + x.tag, x.dtype), [nx])),
  # expand the cmp here so we can preserve rng src edge to get label from ctx
-  (UPat(Ops.INS, name="x"), lambda ctx,x: lower_loop(ctx, x) if x.arg[0] is X86Ops.LOOP_CMP else None),
+  (UPat(Ops.INS, arg=X86Ops.LOOP_CMP, name="x"), lower_loop),
  # rewrite RANGE to ACC = 0 -> LABEL -> JUMP if ACC >= loop bound
  (UPat(Ops.RANGE, name="x"), lower_range),
  # rewrite END to ACC + 1 -> JUMP -> LABEL, also add the out of loop JUMP to the src so this becomes the jump target
  (UPat(Ops.END, name="x"), lower_end),
  # rewrite two address instructions to two address form, if reused src wasn't coalesced insert a move
  (UPat(Ops.INS, name="x"), lambda ctx,x: (nx:=x.replace(src=x.src[1:]),
-   [ctx.ren.copy(x.src[0], greg(x)), nx] if greg(x) != greg(x.src[0]) else [nx]) if x.arg[0] in X86GroupOp.TwoAddress else None),
+   [ctx.ren.copy(x.src[0], greg(x)), nx] if greg(x) != greg(x.src[0]) else [nx]) if x.arg in X86GroupOp.TwoAddress else None),
 ])

 # ***** X86 instruction encoding *****
@@ -619,15 +629,16 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
    r, _x, b = reg >> 3, idx >> 3, rm >> 3
    if sel: # VEX bytes
      vvvv = cast(Register, greg(vvvv_uop)).index if vvvv_uop is not None else 0
-      if sel == 1 and _x == b == we == 0: inst += bytes([0xC5, (~r & 0b1) << 7 | (~vvvv & 0b1111) << 3 | pp])
-      else: inst += bytes([0xC4, (~r & 0b1) << 7 | (~_x & 0b1) << 6 | (~b & 0b1) << 5 | sel, we << 7 | (~vvvv & 0b1111) << 3 | pp])
+      l = (max(reg_sz, rm_sz) > 16) & 0b1
+      if sel == 1 and _x == b == we == 0: inst += bytes([0xC5, (~r & 0b1) << 7 | (~vvvv & 0b1111) << 3 | l << 2 | pp])
+      else: inst += bytes([0xC4, (~r & 0b1) << 7 | (~_x & 0b1) << 6 | (~b & 0b1) << 5 | sel, we << 7 | (~vvvv & 0b1111) << 3 | l << 2 | pp])
    else: # optional PREFIX and REX bytes
      # PREFIX byte signaling 16 bit variant of instruction
      if sz == 2: inst += bytes([0x66])
      # bit signaling 64 bit variant of instruction
      w = sz == 8
      # legacy 8bit opcode is 1 less than 16-64bit variants
-      demote = (rm_sz == 1 or reg_sz == 1) and x.arg[0] not in X86GroupOp.ReadFlags | {X86Ops.LEA}
+      demote = (rm_sz == 1 or reg_sz == 1) and x.arg not in X86GroupOp.ReadFlags | {X86Ops.LEA}
      # REX byte is required when 64 bit or an extended reg is used (index 8 - 15) or lower 8 bits of (rsp, rbp, rsi, rdi) are accessed
      if w | r | _x | b | (reg_sz == 1 & reg >> 2) | (rm_sz == 1 & rm >> 2) | (demote and disp_uop is None and rm >= 4):
        inst += bytes([0b0100 << 4 | w << 3 | r << 2 | _x << 1 | b])
@@ -669,22 +680,22 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
  # get the encoding structure of the uop
  # when a uop writes to memory it takes the form of a store, dtype is void, no definition
  address:tuple[UOp|None, ...]
-  if x.arg[0] in X86GroupOp.WriteMem:
+  if x.arg in X86GroupOp.WriteMem:
    if len(x.src) > 4: address, rest = x.src[:4], x.src[4:]
    else: address, rest = (x, None, None, None), x.src
    imm_uop = rest[:1] if rest and rest[0].op is Ops.CAST else (None,)
    return _encode(rest[0], *address, *(None, *rest[1:])) if reg is None else _encode(None, *address, *(None, *imm_uop))

-  if x.arg[0] in X86GroupOp.Rm1st:
+  if x.arg in X86GroupOp.Rm1st:
    if len(x.src) > 3: address, rest = x.src[:4], x.src[4:]
    else: address, rest = (x.src[0], None, None, None), x.src[1:]
    imm_uop = rest[:1] if rest and rest[0].op is Ops.CAST else (None,)
    return _encode(x, *address, *(None, *imm_uop)) if reg is None else _encode(None, *address, *(x if sel else None, *imm_uop))

-  if x.arg[0] in X86GroupOp.Rm2nd:
+  if x.arg in X86GroupOp.Rm2nd:
    if len(x.src) > 4: address, rest = x.src[1:5], x.src[:1] + x.src[5:]
    else: address, rest = (x.src[1], None, None, None), x.src[:1] + x.src[2:]
-    # cmp reg, rm doesn't define a new register
+    # cmp/vucomiss reg, rm don't define a new register
    return _encode(x, *address, *rest) if x.dtype is not dtypes.void else _encode(rest[0], *address)

  return None
@@ -749,12 +760,18 @@ encodings = {
  # packed int binary
  X86Ops.VPSLLVD: lambda x: encode(x, 0x47, pp=1, sel=2), X86Ops.VPSLLVQ: lambda x: encode(x, 0x47, pp=1, sel=2, we=1),
  X86Ops.VPSRLVD: lambda x: encode(x, 0x45, pp=1, sel=2), X86Ops.VPSRLVQ: lambda x: encode(x, 0x45, pp=1, sel=2, we=1),
+  X86Ops.VPCMPGTB: lambda x: encode(x, 0x64, pp=1, sel=1), X86Ops.VPCMPGTW: lambda x: encode(x, 0x65, pp=1, sel=1),
+  X86Ops.VPCMPGTD: lambda x: encode(x, 0x66, pp=1, sel=1), X86Ops.VPCMPGTQ: lambda x: encode(x, 0x37, pp=1, sel=2),
+  X86Ops.VPCMPEQB: lambda x: encode(x, 0x74, pp=1, sel=1), X86Ops.VPCMPEQW: lambda x: encode(x, 0x75, pp=1, sel=1),
+  X86Ops.VPCMPEQD: lambda x: encode(x, 0x76, pp=1, sel=1), X86Ops.VPCMPEQQ: lambda x: encode(x, 0x29, pp=1, sel=2),
  X86Ops.VPMULLW: lambda x: encode(x, 0xD5, pp=1, sel=1), X86Ops.VPMULLD: lambda x: encode(x, 0x40, pp=1, sel=2),
  X86Ops.VPADDB: lambda x: encode(x, 0xFC, pp=1, sel=1), X86Ops.VPADDW: lambda x: encode(x, 0xFD, pp=1, sel=1),
  X86Ops.VPADDD: lambda x: encode(x, 0xFE, pp=1, sel=1), X86Ops.VPADDQ: lambda x: encode(x, 0xD4, pp=1, sel=1),
  X86Ops.VPSUBB: lambda x: encode(x, 0xF8, pp=1, sel=1), X86Ops.VPSUBW: lambda x: encode(x, 0xF9, pp=1, sel=1),
  X86Ops.VPSUBD: lambda x: encode(x, 0xFA, pp=1, sel=1), X86Ops.VPSUBQ: lambda x: encode(x, 0xFB, pp=1, sel=1),
  X86Ops.VPSRAVD: lambda x: encode(x, 0x46, pp=1, sel=2),
+  # float cmp
+  X86Ops.VUCOMISS: lambda x: encode(x, 0x2E, pp=0, sel=1), X86Ops.VUCOMISD: lambda x: encode(x, 0x2E, pp=1, sel=1),
  # scalar / packed float binary
  X86Ops.VADDSS: lambda x: encode(x, 0x58, pp=2, sel=1), X86Ops.VADDPS: lambda x: encode(x, 0x58, pp=0, sel=1),
  X86Ops.VADDSD: lambda x: encode(x, 0x58, pp=3, sel=1), X86Ops.VADDPD: lambda x: encode(x, 0x58, pp=1, sel=1),
@@ -770,6 +787,7 @@ encodings = {
  X86Ops.CMOVB: lambda x: encode(x, 0x0F42), X86Ops.CMOVL: lambda x: encode(x, 0x0F4C),
  X86Ops.CMOVE: lambda x: encode(x, 0x0F44), X86Ops.CMOVNE: lambda x: encode(x, 0x0F45),
  X86Ops.VBLENDVPS: lambda x: encode(x, 0x4A, pp=1, sel=3), X86Ops.VBLENDVPD: lambda x: encode(x, 0x4B, pp=1, sel=3),
+  X86Ops.VPBLENDVB: lambda x: encode(x, 0x4C, pp=1, sel=3),
  # shuffles
  X86Ops.VPSRLDQ: lambda x: encode(x, 0x73, reg=3, pp=1, sel=1),
  X86Ops.VPINSRB: lambda x: encode(x, 0x20, pp=1, sel=3), X86Ops.VPINSRW: lambda x: encode(x, 0xC4, pp=1, sel=1),
@@ -805,28 +823,28 @@ class X86Renderer(ISARenderer):
    super().__init__(target)
    from tinygrad.runtime.support.compiler_cpu import X86Compiler
    self.compiler = X86Compiler()
-  def is_two_address(self, x:UOp) -> bool: return x.op is Ops.INS and x.arg[0] in X86GroupOp.TwoAddress
+  def is_two_address(self, x:UOp) -> bool: return x.arg in X86GroupOp.TwoAddress
  def stack_pointer(self) -> UOp: return def_reg(dtypes.uint64, RSP)
  # the value of a BUFFER is its address, it moves through registers and the stack as a 64bit int
  def copy(self, x:UOp, reg:Register):
-    if x.op is Ops.BUFFER: x = x.replace(arg=replace(x.arg, dtype=dtypes.uint64))
-    ret = isel_matcher.rewrite(UOp(Ops.COPY, src=(x,), tag=reg))
+    dt = dtypes.uint64 if x.op is Ops.BUFFER else x.dtype
+    ret = isel_matcher.rewrite(UOp(Ops.COPY, dt, (x,), tag=reg))
    assert ret is not None, f"failed to copy {x}"
    return ret

  def spill(self, disp:UOp, x:UOp) -> UOp:
-    if x.op is Ops.BUFFER: x = x.replace(arg=replace(x.arg, dtype=dtypes.uint64))
+    if x.op is Ops.BUFFER: x = x.replace(dtype=dtypes.uint64)
    is_xmm = isinstance(x.tag, tuple) and x.tag[0].cons[0].size == 16
    op = X86Ops.VMOVUPSm if is_xmm else X86Ops.MOVm
-    return UOp(Ops.INS, src=fold_address(self.stack_pointer().index(disp)) + (x,), arg=(op, dtypes.void), tag=x.tag)
+    return UOp(Ops.INS, src=fold_address(self.stack_pointer().index(disp)) + (x,), arg=op, tag=x.tag)

  def fill(self, disp:UOp, x:UOp, reg:Register) -> UOp:
    is_xmm = reg.cons[0].size == 16
    dt = dtypes.uint64 if x.op is Ops.BUFFER else x.dtype
-    return UOp(Ops.INS, src=fold_address(self.stack_pointer().index(disp)), arg=(X86Ops.VMOVUPS if is_xmm else X86Ops.MOV, dt), tag=(reg,))
+    return UOp(Ops.INS, dt, fold_address(self.stack_pointer().index(disp)), X86Ops.VMOVUPS if is_xmm else X86Ops.MOV, (reg,))

  def asm_str(self, uops:list[UOp], function_name:str) -> str:
-    def _format_op(x:UOp) -> str: return f"    {(o[7:-1] if (o:=str(x.arg[0]))[-1] in ('i', 'm') else o[7:]).lower():7s}"
+    def _format_op(x:UOp) -> str: return f"    {(o[7:-1] if (o:=str(x.arg))[-1] in ('i', 'm') else o[7:]).lower():7s}"
    def _format_operands(x:UOp) -> str:
      def _format(src:tuple[UOp, ...]) -> list[str]:
        return [str(s.src[0].val) if s.op is Ops.CAST else reg_strs[o].get(s.dtype.itemsize, o) if \
@@ -834,17 +852,17 @@ class X86Renderer(ISARenderer):
      def _mem_adress(base:UOp, idx:UOp, disp:UOp, sz:UOp) -> list[str]:
        return [f"[{greg(base)}" + (f" + {greg(idx)}*{sz.src[0].val}" if greg(idx) else "") + (f" + {d}" if (d:=disp.src[0].val) else "") + "]"]

-      if len(x.src) > 4 and x.arg[0] in X86GroupOp.WriteMem: ret = _mem_adress(*x.src[:4]) + _format(x.src[4:])
-      elif len(x.src) > 3 and x.arg[0] in X86GroupOp.Rm1st: ret = _format((x,)) + _mem_adress(*x.src[:4]) + _format(x.src[4:])
-      elif len(x.src) > 4 and x.arg[0] in X86GroupOp.Rm2nd: ret = _format((x, x.src[0])) + _mem_adress(*x.src[1:5]) + _format(x.src[5:])
+      if len(x.src) > 4 and x.arg in X86GroupOp.WriteMem: ret = _mem_adress(*x.src[:4]) + _format(x.src[4:])
+      elif len(x.src) > 3 and x.arg in X86GroupOp.Rm1st: ret = _format((x,)) + _mem_adress(*x.src[:4]) + _format(x.src[4:])
+      elif len(x.src) > 4 and x.arg in X86GroupOp.Rm2nd: ret = _format((x, x.src[0])) + _mem_adress(*x.src[1:5]) + _format(x.src[5:])
      else: ret = _format((x,) + x.src)
      return ", ".join(ret)

    asm = [f".{function_name}:"]
    for u in uops:
-      if u.op is not Ops.INS or u.arg[0] is X86Ops.DEFINE: continue
-      if u.arg[0] is X86Ops.LABEL: asm.append(f"{str(u.tag)}:")
-      elif u.arg[0] is X86Ops.RET: asm.append(_format_op(u))
+      if u.op is not Ops.INS or u.arg is X86Ops.DEFINE: continue
+      if u.arg is X86Ops.LABEL: asm.append(f"{str(u.tag)}:")
+      elif u.arg is X86Ops.RET: asm.append(_format_op(u))
      else: asm.append(_format_op(u) + " " + _format_operands(u))
    return "\n".join(asm)

@@ -853,15 +871,15 @@ class X86Renderer(ISARenderer):
    jumps: dict[UOp, int] = {}
    binary = bytearray()
    for u in uops:
-      if u.op is not Ops.INS or u.arg[0] is X86Ops.DEFINE: continue
-      if u.arg[0] is X86Ops.LOOP_CMP: continue
-      if u.arg[0] is X86Ops.LABEL:
+      if u.op is not Ops.INS or u.arg is X86Ops.DEFINE: continue
+      if u.arg is X86Ops.LOOP_CMP: continue
+      if u.arg is X86Ops.LABEL:
        targets[u.tag] = len(binary)
        continue
-      if u.arg[0] not in encodings or (l:=encodings[u.arg[0]](u)) is None:
-        raise RuntimeError(f"failed to encode {u.arg[0]} with {u.dtype} srcs {[x.dtype for x in u.src]}")
+      if u.arg not in encodings or (l:=encodings[u.arg](u)) is None:
+        raise RuntimeError(f"failed to encode {u.arg} with {u.dtype} srcs {[x.dtype for x in u.src]}")
      binary.extend(l)
-      if u.arg[0] in (X86Ops.JL, X86Ops.JB, X86Ops.JE, X86Ops.JNE, X86Ops.JGE, X86Ops.JMP): jumps[u] = len(binary)
+      if u.arg in (X86Ops.JL, X86Ops.JB, X86Ops.JE, X86Ops.JNE, X86Ops.JGE, X86Ops.JMP): jumps[u] = len(binary)
    # fixup jump targets now that encoding size is known
    for u in uops:
      if (t:=jumps.get(u)) is not None: binary[t-4:t] = (targets[u.tag] - t).to_bytes(4, 'little', signed=True)
@@ -125,7 +125,7 @@ class NIRRenderer(Renderer):
    (UPat.var('x', dtype=dtypes.bool)<UPat.var('y'), lambda x,y: (x^True)&y),
    # a bool is one bit in NIR but a byte in memory, so every access to a bool buffer goes through a uint8 view of it
    (UPat(Ops.LOAD, dtypes.bool, name="x"),
-     lambda x: x.replace(src=(with_storage(x.src[0], dtypes.uint8),)+((x.src[1].cast(dtypes.uint8),) if len(x.src)>=2 else ())
+     lambda x: x.replace(dtype=None, src=(with_storage(x.src[0], dtypes.uint8),)+((x.src[1].cast(dtypes.uint8),) if len(x.src)>=2 else ())
      +x.src[2:]).cast(dtypes.bool)),
    (UPat(Ops.STORE, src=(UPat(name="idx"), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
     lambda x,idx: x.replace(src=(with_storage(idx, dtypes.uint8), x.src[1].cast(dtypes.uint8))+x.src[2:])),
@@ -47,7 +47,7 @@ ptx_matcher = PatternMatcher([
    lambda x: (UOp(x.op, src=tuple(vv.cast(dtypes.float32) for vv in x.src), arg=x.arg).cast(dtypes.half))),
  # a bool is a predicate register in PTX but a byte in memory, so a bool buffer is accessed through a uint8 view of it
  (UPat(Ops.LOAD, dtypes.bool, src=(UPat(name="idx"),), name="x", allow_any_len=True),
-   lambda x,idx: x.replace(src=(with_storage(idx, dtypes.uint8),) + ((x.src[1].cast(dtypes.uint8),) if len(x.src) >= 2 else ())
+   lambda x,idx: x.replace(dtype=None, src=(with_storage(idx, dtypes.uint8),) + ((x.src[1].cast(dtypes.uint8),) if len(x.src) >= 2 else ())
     + x.src[2:]).cast(dtypes.bool) if idx.addrspace != AddrSpace.REG else None),
  (UPat(Ops.STORE, src=(UPat(name="idx"), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
   lambda x,idx: x.replace(src=(with_storage(idx, dtypes.uint8), x.src[1].cast(dtypes.uint8))+x.src[2:]) if idx.addrspace != AddrSpace.REG else None),
@@ -10,14 +10,10 @@ rocr_src = "https://github.com/ROCm/rocm-systems/archive/refs/tags/rocm-7.1.1.ta
 linux_headers_deb = "https://snapshot.debian.org/archive/debian/20260207T145350Z/pool/main/l/linux/linux-libc-dev_6.18.9-1_all.deb"
 linux_headers_kern_deb = "https://snapshot.debian.org/archive/debian/20260207T145350Z/pool/main/l/linux/linux-headers-6.18.9+deb14-common_6.18.9-1_all.deb"
 liburing_src = "https://raw.githubusercontent.com/axboe/liburing/refs/tags/liburing-2.14/src/include/liburing.h"
-bnxt_src = ["https://raw.githubusercontent.com/torvalds/linux/v6.18/drivers/" + s for s in
-            ("infiniband/hw/bnxt_re/roce_hsi.h", "infiniband/hw/bnxt_re/qplib_rcfw.h", "infiniband/hw/bnxt_re/qplib_res.h",
-             "net/ethernet/broadcom/bnxt/bnxt_hwrm.h")]
 ggml_common_src = "https://raw.githubusercontent.com/ggml-org/ggml/d4fcfe88a8bcf5c9840be14be6c2fbf1f5b3b2db/src/ggml-common.h"
 cudart_src = "https://developer.download.nvidia.com/compute/cuda/redist/cuda_cudart/linux-x86_64/cuda_cudart-linux-x86_64-12.0.146-archive.tar.xz"
 nvrtc_src = "https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvrtc/linux-x86_64/cuda_nvrtc-linux-x86_64-12.0.140-archive.tar.xz"
 opencl_src = "https://github.com/KhronosGroup/OpenCL-Headers/archive/2e30669d48718fd460f085b4b35b160dad51ce9d.tar.gz"
-comgr_2_src = "https://repo.radeon.com/rocm/apt/6.2/pool/main/c/comgr/comgr_2.8.0.60200-66~24.04_amd64.deb"
 macossdk = "/var/db/xcode_select_link/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk"

 llvm_lib = (
@@ -54,8 +50,6 @@ def load(name, files, **kwargs):
    if srcs: td.cleanup()
  return importlib.import_module(f"{path}.{name.replace('/', '.')}")

-def _extract_deb(path:str): subprocess.run("ar x *.deb && tar xf data.tar.*", cwd=path, shell=True, check=True)
-
 def __getattr__(nm):
  match nm:
    case "libc":
@@ -105,11 +99,14 @@ def __getattr__(nm):
    # this defines all syscall numbers. should probably unify linux autogen?
    case "io_uring":
      return load("io_uring", ["{}/liburing.h", "{}/usr/include/linux/io_uring.h", "{}/usr/include/asm-generic/unistd.h"],
-                  args=["-I{}/usr/include"], srcs=[linux_headers_deb, liburing_src], rules=[('__NR', 'NR')], preprocess=_extract_deb)
+                  args=["-I{}/usr/include"], srcs=[linux_headers_deb, liburing_src], rules=[('__NR', 'NR')],
+                  preprocess=lambda path: subprocess.run(f"ar x {linux_headers_deb.split('/')[-1]} && tar xf data.tar.xz", cwd=path, shell=True, check=True))
    case "llvm": return load("llvm", lambda: [system("llvm-config-20 --includedir")+"/llvm-c/**/*.h"], dll=llvm_lib,
                             args=lambda: system("llvm-config-20 --cflags").split(), recsym=True, prolog=["from tinygrad.helpers import WIN, OSX"])
-    case "pci": return load("pci", ["{}/usr/include/linux/pci_regs.h"], srcs=linux_headers_deb, preprocess=_extract_deb)
-    case "vfio": return load("vfio", ["{}/usr/include/linux/vfio.h"], args=["-I{}/usr/include"], srcs=linux_headers_deb, preprocess=_extract_deb)
+    case "pci": return load("pci", ["{}/usr/include/linux/pci_regs.h"], srcs=linux_headers_deb,
+                             preprocess=lambda path: subprocess.run(f"ar x {linux_headers_deb.split('/')[-1]} && tar xf data.tar.xz", cwd=path, shell=True, check=True))
+    case "vfio": return load("vfio", ["{}/usr/include/linux/vfio.h"], args=["-I{}/usr/include"], srcs=linux_headers_deb,
+                             preprocess=lambda path: subprocess.run(f"ar x {linux_headers_deb.split('/')[-1]} && tar xf data.tar.xz", cwd=path, shell=True, check=True))
    # could add rule: WGPU_COMMA -> ','
    case "webgpu": return load("webgpu", [root/"extra/webgpu/webgpu.h"], dll=webgpu_lib,
                               prolog=["from tinygrad.helpers import WIN, OSX", "import sysconfig, os"])
@@ -119,10 +116,9 @@ def __getattr__(nm):
                            dll="os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamdhip64.so'",
                            args=["-D__HIP_PLATFORM_AMD__", "-I/opt/rocm/include", "-x", "c++"], prolog=["import os"])
    case "comgr" | "comgr_3":
-      prefix = "{}/opt/rocm-6.2.0" if nm == "comgr" else "/opt/rocm"
-      return load(nm, [f"{prefix}/include/amd_comgr/amd_comgr.h"], dll="[os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamd_comgr.so', 'amd_comgr']",
-                  args=["-D__HIP_PLATFORM_AMD__", f"-I{prefix}/include", "-x", "c++"], prolog=["import os"], srcs=comgr_2_src if nm == "comgr" else None,
-                  **({'preprocess':_extract_deb} if nm == "comgr" else {}))
+      return load("comgr_3" if nm == "comgr_3" else "comgr", ["/opt/rocm/include/amd_comgr/amd_comgr.h"],
+                  dll= "[os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamd_comgr.so', 'amd_comgr']",
+                  args=["-D__HIP_PLATFORM_AMD__", "-I/opt/rocm/include", "-x", "c++"], prolog=["import os"])
    case "hsa": return load("hsa", [*[f"{{}}/projects/rocr-runtime/runtime/hsa-runtime/core/inc/{s}.h" for s in ["registers"]],
                                    *[f"{{}}/projects/rocr-runtime/runtime/hsa-runtime/inc/{s}.h" for s in [
                                        "hsa", "hsa_ext_amd", "amd_hsa_signal", "amd_hsa_queue", "amd_hsa_kernel_code",
@@ -185,17 +181,6 @@ def __getattr__(nm):
      return load("mlx5", [root/"extra/mlx_driver/mlx5.h", f"{kh}/mlx5_ifc.h"], srcs=linux_headers_kern_deb,
                  args=["-Du8=unsigned char", "-Du16=unsigned short", "-Du32=unsigned int", "-Du64=unsigned long long",
                        "-D__be16=unsigned short", "-D__be32=unsigned int", "-D__be64=unsigned long long", f"-I{kh}"],
-                  preprocess=_extract_deb)
-    case "bnxt":
-      kh = "{}/usr/src/linux-headers-6.18.9+deb14-common/include"
-      return load("bnxt", [f"{kh}/linux/bnxt/hsi.h", *[f"{{}}/{s.split('/')[-1]}" for s in bnxt_src]],
-                  srcs=[linux_headers_kern_deb, *bnxt_src],
-                  args=["-Du8=unsigned char", "-Du32=unsigned int", "-Du64=unsigned long long", "-D__le16=unsigned short",
-                        "-D__le32=unsigned int", "-D__le64=unsigned long long", "-D__be16=unsigned short", "-D__be32=unsigned int", f"-I{kh}"],
-                  patterns=[r"hwrm_((ver_get|func_(qcaps|qcfg|reset|drv_rgtr|backing_store_(qcaps|cfg)_v2)|stat_ctx_alloc|ring_alloc"
-                            r"|vnic_(alloc|cfg)|cfa_l2_filter_alloc|port_phy_cfg)_(input|output)|(cmd|resp)_hdr)$",
-                            r"((cmdq|creq)_(base|init|add_gid|create_(cq|qp)|initialize_fw|modify_qp|query_version|register_mr)(_resp)?"
-                            r"|cq_(base|req)|sq_(rdma_hdr|sge))$",
-                            r"(BNXT|CMDQ|CREQ|CQ|SQ|DBC|PTU|RCFW|HWRM|VNIC|RING_ALLOC|STAT_CTX|CFA_L2_FILTER|PORT_PHY_CFG|FIRMWARE_FIRST"
-                            r"|FUNC_(QCAPS|QCFG|RESET|DRV_RGTR|BACKING_STORE))_"], preprocess=_extract_deb)
+                  preprocess=lambda path: subprocess.run(f"ar x {linux_headers_kern_deb.split('/')[-1]} && tar xf data.tar.xz",
+                                                         cwd=path, shell=True, check=True))
    case _: raise AttributeError(f"no such autogen: {nm}")
@@ -22,12 +22,11 @@ reg_files = {
 reg_patterns = {
  "gc": ["GCVM", "GCMC_VM", "CP_(HQD|MQD|MEC|ME_CNTL|PERFMON|RB_WPTR_POLL_CNTL|INT_CNTL|STAT|PFP_PRGRM|ME_PRGRM|COHER_START)", "COMPUTE_",
         "(SQ|GL2C|TCC)_PERFCOUNTER", "SQ_THREAD_TRACE", "SPI_(CONFIG_CNTL|COMPUTE_QUEUE_RESET)", "GRBM", "SH_MEM", "RLC", "TCP", "GB_ADDR_CONFIG",
-         "SDMA[01]_(WATCHDOG_CNTL|UTCL1_(CNTL|PAGE)|MCU_CNTL|F32_CNTL|CNTL|QUEUE0_|RLC_CGCG_CTRL)", "SCRATCH_REG[0-367]"],
+         "SDMA[01]_(WATCHDOG_CNTL|UTCL1_(CNTL|PAGE)|MCU_CNTL|F32_CNTL|CNTL|QUEUE0_|RLC_CGCG_CTRL)", "SCRATCH_REG[67]"],
  "mmhub": ["MMVM", "MMMC_VM", "MM_ATC_L2_MISC_CG"],
  "nbio": (nbio:=["BIF_BX_PF[01]_GPU_HDP_FLUSH", "BIF_BX_PF0_RSMU", "BIF_BX0_(REMAP_HDP_MEM_FLUSH_CNTL|BIF_DOORBELL_INT_CNTL|PCIE_INDEX2|PCIE_DATA2)",
                  "BIFC_(DOORBELL_ACCESS_EN_PF|GFX_INT_MONITOR_MASK)", "XCC_DOORBELL_FENCE", "DOORBELL0_CTRL_ENTRY", "GDC_S2A0_S2A_DOORBELL_ENTRY",
-                  "S2A_DOORBELL_ENTRY", "RCC_DEV0_EPF0_RCC_DOORBELL_APER_EN",
-                  "BIF_BX_DEV0_EPF0_VF0_HDP_MEM_COHERENCY_FLUSH_CNTL", "RCC_DEV0_EPF2_STRAP2"]),
+                  "S2A_DOORBELL_ENTRY", "RCC_DEV0_EPF0_RCC_DOORBELL_APER_EN", "RCC_DEV0_EPF2_STRAP2"]),
  "nbif": nbio,
  "mp": ["MP([01]|ASP)_SMN_C2PMSG"], "hdp": ["HDP_MEM_POWER_CTRL"], "oss": ["IH_"], "sdma": ["SDMA_GFX", "SDMA_CNTL"]
 }
@@ -39,7 +38,7 @@ def __getattr__(nm):
    case "am": return load("am/am", [root/f"extra/amdpci/headers/{s}.h" for s in ["v11_structs", "v12_structs", "amdgpu_vm",
      "discovery", "amdgpu_ucode", "psp_gfx_if", "amdgpu_psp", "amdgpu_irq", "amdgpu_doorbell"]] + [f"{AMD}/amdkfd/soc15_int.h"] + \
      [f"{AMDINC}/ivsrcid/{s}.h" for s in [f"gfx/irqsrcs_gfx_{x}_0" for x in ('9','11_0','12_0')] + [f"sdma0/irqsrcs_sdma0_{x}_0" for x in (4,5)]] + \
-      [f"{AMDINC}/{s}.h" for s in ["v9_structs", "soc15_ih_clientid"]] + [f"{AMD}/amdgpu/{s}.h" for s in ["mxgpu_nv", "amdgpu_virt"]], args=inc, srcs=am_src, rules=kern_rules)
+      [f"{AMDINC}/{s}.h" for s in ["v9_structs", "soc15_ih_clientid"]], args=inc, srcs=am_src, rules=kern_rules)
    case "pm4_soc15": return load("am/pm4_soc15", [f"{AMD}/amdkfd/kfd_pm4_headers_ai.h", f"{AMD}/amdgpu/soc15d.h"], srcs=am_src)
    case "pm4_nv": return load("am/pm4_nv", [f"{AMD}/amdkfd/kfd_pm4_headers_ai.h", f"{AMD}/amdgpu/nvd.h"], srcs=am_src)
    case "sdma_4_0_0": return load("am/sdma_4_0_0", [root/"extra/hip_gpu_driver/sdma_registers.h", f"{AMD}/amdgpu/vega10_sdma_pkt_open.h"],
@@ -4110,110 +4110,6 @@ class struct_v9_gfx_meta_data(c.Struct):
 struct_v9_gfx_meta_data.register_fields([('ce_payload', struct_v9_ce_ib_state, 0), ('reserved1', c.Array[uint32_t, Literal[54]], 40), ('de_payload', struct_v9_de_ib_state, 256), ('DeIbBaseAddrLo', uint32_t, 364), ('DeIbBaseAddrHi', uint32_t, 368), ('reserved2', c.Array[uint32_t, Literal[931]], 372)])
 enum_soc15_ih_clientid: dict[int, str] = {(SOC15_IH_CLIENTID_IH:=0): 'SOC15_IH_CLIENTID_IH', (SOC15_IH_CLIENTID_ACP:=1): 'SOC15_IH_CLIENTID_ACP', (SOC15_IH_CLIENTID_ATHUB:=2): 'SOC15_IH_CLIENTID_ATHUB', (SOC15_IH_CLIENTID_BIF:=3): 'SOC15_IH_CLIENTID_BIF', (SOC15_IH_CLIENTID_DCE:=4): 'SOC15_IH_CLIENTID_DCE', (SOC15_IH_CLIENTID_ISP:=5): 'SOC15_IH_CLIENTID_ISP', (SOC15_IH_CLIENTID_PCIE0:=6): 'SOC15_IH_CLIENTID_PCIE0', (SOC15_IH_CLIENTID_RLC:=7): 'SOC15_IH_CLIENTID_RLC', (SOC15_IH_CLIENTID_SDMA0:=8): 'SOC15_IH_CLIENTID_SDMA0', (SOC15_IH_CLIENTID_SDMA1:=9): 'SOC15_IH_CLIENTID_SDMA1', (SOC15_IH_CLIENTID_SE0SH:=10): 'SOC15_IH_CLIENTID_SE0SH', (SOC15_IH_CLIENTID_SE1SH:=11): 'SOC15_IH_CLIENTID_SE1SH', (SOC15_IH_CLIENTID_SE2SH:=12): 'SOC15_IH_CLIENTID_SE2SH', (SOC15_IH_CLIENTID_SE3SH:=13): 'SOC15_IH_CLIENTID_SE3SH', (SOC15_IH_CLIENTID_UVD1:=14): 'SOC15_IH_CLIENTID_UVD1', (SOC15_IH_CLIENTID_THM:=15): 'SOC15_IH_CLIENTID_THM', (SOC15_IH_CLIENTID_UVD:=16): 'SOC15_IH_CLIENTID_UVD', (SOC15_IH_CLIENTID_VCE0:=17): 'SOC15_IH_CLIENTID_VCE0', (SOC15_IH_CLIENTID_VMC:=18): 'SOC15_IH_CLIENTID_VMC', (SOC15_IH_CLIENTID_XDMA:=19): 'SOC15_IH_CLIENTID_XDMA', (SOC15_IH_CLIENTID_GRBM_CP:=20): 'SOC15_IH_CLIENTID_GRBM_CP', (SOC15_IH_CLIENTID_ATS:=21): 'SOC15_IH_CLIENTID_ATS', (SOC15_IH_CLIENTID_ROM_SMUIO:=22): 'SOC15_IH_CLIENTID_ROM_SMUIO', (SOC15_IH_CLIENTID_DF:=23): 'SOC15_IH_CLIENTID_DF', (SOC15_IH_CLIENTID_VCE1:=24): 'SOC15_IH_CLIENTID_VCE1', (SOC15_IH_CLIENTID_PWR:=25): 'SOC15_IH_CLIENTID_PWR', (SOC15_IH_CLIENTID_RESERVED:=26): 'SOC15_IH_CLIENTID_RESERVED', (SOC15_IH_CLIENTID_UTCL2:=27): 'SOC15_IH_CLIENTID_UTCL2', (SOC15_IH_CLIENTID_EA:=28): 'SOC15_IH_CLIENTID_EA', (SOC15_IH_CLIENTID_UTCL2LOG:=29): 'SOC15_IH_CLIENTID_UTCL2LOG', (SOC15_IH_CLIENTID_MP0:=30): 'SOC15_IH_CLIENTID_MP0', (SOC15_IH_CLIENTID_MP1:=31): 'SOC15_IH_CLIENTID_MP1', (SOC15_IH_CLIENTID_MAX:=32): 'SOC15_IH_CLIENTID_MAX', (SOC15_IH_CLIENTID_VCN:=16): 'SOC15_IH_CLIENTID_VCN', (SOC15_IH_CLIENTID_VCN1:=14): 'SOC15_IH_CLIENTID_VCN1', (SOC15_IH_CLIENTID_SDMA2:=1): 'SOC15_IH_CLIENTID_SDMA2', (SOC15_IH_CLIENTID_SDMA3:=4): 'SOC15_IH_CLIENTID_SDMA3', (SOC15_IH_CLIENTID_SDMA3_Sienna_Cichlid:=5): 'SOC15_IH_CLIENTID_SDMA3_Sienna_Cichlid', (SOC15_IH_CLIENTID_SDMA4:=5): 'SOC15_IH_CLIENTID_SDMA4', (SOC15_IH_CLIENTID_SDMA5:=17): 'SOC15_IH_CLIENTID_SDMA5', (SOC15_IH_CLIENTID_SDMA6:=19): 'SOC15_IH_CLIENTID_SDMA6', (SOC15_IH_CLIENTID_SDMA7:=24): 'SOC15_IH_CLIENTID_SDMA7', (SOC15_IH_CLIENTID_VMC1:=6): 'SOC15_IH_CLIENTID_VMC1'}
 enum_soc21_ih_clientid: dict[int, str] = {(SOC21_IH_CLIENTID_IH:=0): 'SOC21_IH_CLIENTID_IH', (SOC21_IH_CLIENTID_ATHUB:=2): 'SOC21_IH_CLIENTID_ATHUB', (SOC21_IH_CLIENTID_BIF:=3): 'SOC21_IH_CLIENTID_BIF', (SOC21_IH_CLIENTID_DCN:=4): 'SOC21_IH_CLIENTID_DCN', (SOC21_IH_CLIENTID_ISP:=5): 'SOC21_IH_CLIENTID_ISP', (SOC21_IH_CLIENTID_MP3:=6): 'SOC21_IH_CLIENTID_MP3', (SOC21_IH_CLIENTID_RLC:=7): 'SOC21_IH_CLIENTID_RLC', (SOC21_IH_CLIENTID_GFX:=10): 'SOC21_IH_CLIENTID_GFX', (SOC21_IH_CLIENTID_IMU:=11): 'SOC21_IH_CLIENTID_IMU', (SOC21_IH_CLIENTID_VCN1:=14): 'SOC21_IH_CLIENTID_VCN1', (SOC21_IH_CLIENTID_THM:=15): 'SOC21_IH_CLIENTID_THM', (SOC21_IH_CLIENTID_VCN:=16): 'SOC21_IH_CLIENTID_VCN', (SOC21_IH_CLIENTID_VPE1:=17): 'SOC21_IH_CLIENTID_VPE1', (SOC21_IH_CLIENTID_VMC:=18): 'SOC21_IH_CLIENTID_VMC', (SOC21_IH_CLIENTID_GRBM_CP:=20): 'SOC21_IH_CLIENTID_GRBM_CP', (SOC21_IH_CLIENTID_ROM_SMUIO:=22): 'SOC21_IH_CLIENTID_ROM_SMUIO', (SOC21_IH_CLIENTID_DF:=23): 'SOC21_IH_CLIENTID_DF', (SOC21_IH_CLIENTID_VPE:=24): 'SOC21_IH_CLIENTID_VPE', (SOC21_IH_CLIENTID_PWR:=25): 'SOC21_IH_CLIENTID_PWR', (SOC21_IH_CLIENTID_LSDMA:=26): 'SOC21_IH_CLIENTID_LSDMA', (SOC21_IH_CLIENTID_MP0:=30): 'SOC21_IH_CLIENTID_MP0', (SOC21_IH_CLIENTID_MP1:=31): 'SOC21_IH_CLIENTID_MP1', (SOC21_IH_CLIENTID_MAX:=32): 'SOC21_IH_CLIENTID_MAX'}
-enum_idh_request: dict[int, str] = {(IDH_REQ_GPU_INIT_ACCESS:=1): 'IDH_REQ_GPU_INIT_ACCESS', (IDH_REL_GPU_INIT_ACCESS:=2): 'IDH_REL_GPU_INIT_ACCESS', (IDH_REQ_GPU_FINI_ACCESS:=3): 'IDH_REQ_GPU_FINI_ACCESS', (IDH_REL_GPU_FINI_ACCESS:=4): 'IDH_REL_GPU_FINI_ACCESS', (IDH_REQ_GPU_RESET_ACCESS:=5): 'IDH_REQ_GPU_RESET_ACCESS', (IDH_REQ_GPU_INIT_DATA:=6): 'IDH_REQ_GPU_INIT_DATA', (IDH_LOG_VF_ERROR:=200): 'IDH_LOG_VF_ERROR', (IDH_READY_TO_RESET:=201): 'IDH_READY_TO_RESET', (IDH_RAS_POISON:=202): 'IDH_RAS_POISON', (IDH_REQ_RAS_ERROR_COUNT:=203): 'IDH_REQ_RAS_ERROR_COUNT', (IDH_REQ_RAS_CPER_DUMP:=204): 'IDH_REQ_RAS_CPER_DUMP', (IDH_REQ_RAS_BAD_PAGES:=205): 'IDH_REQ_RAS_BAD_PAGES', (IDH_REQ_RAS_CHK_CRITI:=206): 'IDH_REQ_RAS_CHK_CRITI'}
-enum_idh_event: dict[int, str] = {(IDH_CLR_MSG_BUF:=0): 'IDH_CLR_MSG_BUF', (IDH_READY_TO_ACCESS_GPU:=1): 'IDH_READY_TO_ACCESS_GPU', (IDH_FLR_NOTIFICATION:=2): 'IDH_FLR_NOTIFICATION', (IDH_FLR_NOTIFICATION_CMPL:=3): 'IDH_FLR_NOTIFICATION_CMPL', (IDH_SUCCESS:=4): 'IDH_SUCCESS', (IDH_FAIL:=5): 'IDH_FAIL', (IDH_QUERY_ALIVE:=6): 'IDH_QUERY_ALIVE', (IDH_REQ_GPU_INIT_DATA_READY:=7): 'IDH_REQ_GPU_INIT_DATA_READY', (IDH_RAS_POISON_READY:=8): 'IDH_RAS_POISON_READY', (IDH_PF_SOFT_FLR_NOTIFICATION:=9): 'IDH_PF_SOFT_FLR_NOTIFICATION', (IDH_RAS_ERROR_DETECTED:=10): 'IDH_RAS_ERROR_DETECTED', (IDH_RAS_ERROR_COUNT_READY:=11): 'IDH_RAS_ERROR_COUNT_READY', (IDH_RAS_CPER_DUMP_READY:=14): 'IDH_RAS_CPER_DUMP_READY', (IDH_RAS_BAD_PAGES_READY:=15): 'IDH_RAS_BAD_PAGES_READY', (IDH_RAS_BAD_PAGES_NOTIFICATION:=16): 'IDH_RAS_BAD_PAGES_NOTIFICATION', (IDH_UNRECOV_ERR_NOTIFICATION:=17): 'IDH_UNRECOV_ERR_NOTIFICATION', (IDH_REQ_RAS_CHK_CRITI_READY:=18): 'IDH_REQ_RAS_CHK_CRITI_READY', (IDH_TEXT_MESSAGE:=255): 'IDH_TEXT_MESSAGE'}
-class struct_amdgpu_virt_ops(c.Struct): pass
-enum_amdgpu_sriov_vf_mode: dict[int, str] = {(SRIOV_VF_MODE_BARE_METAL:=0): 'SRIOV_VF_MODE_BARE_METAL', (SRIOV_VF_MODE_ONE_VF:=1): 'SRIOV_VF_MODE_ONE_VF', (SRIOV_VF_MODE_MULTI_VF:=2): 'SRIOV_VF_MODE_MULTI_VF'}
-@c.record
-class struct_amdgpu_mm_table(c.Struct):
-  SIZE = 24
-  bo: c.POINTER[struct_amdgpu_bo]
-  cpu_addr: c.POINTER[ctypes.c_uint32]
-  gpu_addr: int
-class struct_amdgpu_bo(c.Struct): pass
-struct_amdgpu_mm_table.register_fields([('bo', c.POINTER[struct_amdgpu_bo], 0), ('cpu_addr', c.POINTER[uint32_t], 8), ('gpu_addr', uint64_t, 16)])
-class struct_amdgpu_vf_error_buffer(c.Struct): pass
-@c.record
-class struct_amdgpu_virt_fw_reserve(c.Struct):
-  SIZE = 32
-  p_pf2vf: c.POINTER[struct_amd_sriov_msg_pf2vf_info_header]
-  p_vf2pf: c.POINTER[struct_amd_sriov_msg_vf2pf_info_header]
-  ras_telemetry: ctypes.c_void_p
-  checksum_key: int
-@c.record
-class struct_amd_sriov_msg_pf2vf_info_header(c.Struct):
-  SIZE = 16
-  size: int
-  version: int
-  reserved: c.Array[ctypes.c_uint32, Literal[2]]
-struct_amd_sriov_msg_pf2vf_info_header.register_fields([('size', uint32_t, 0), ('version', uint32_t, 4), ('reserved', c.Array[uint32_t, Literal[2]], 8)])
-@c.record
-class struct_amd_sriov_msg_vf2pf_info_header(c.Struct):
-  SIZE = 16
-  size: int
-  version: int
-  reserved: c.Array[ctypes.c_uint32, Literal[2]]
-struct_amd_sriov_msg_vf2pf_info_header.register_fields([('size', uint32_t, 0), ('version', uint32_t, 4), ('reserved', c.Array[uint32_t, Literal[2]], 8)])
-struct_amdgpu_virt_fw_reserve.register_fields([('p_pf2vf', c.POINTER[struct_amd_sriov_msg_pf2vf_info_header], 0), ('p_vf2pf', c.POINTER[struct_amd_sriov_msg_vf2pf_info_header], 8), ('ras_telemetry', ctypes.c_void_p, 16), ('checksum_key', ctypes.c_uint32, 24)])
-enum_AMDGIM_FEATURE_FLAG: dict[int, str] = {(AMDGIM_FEATURE_ERROR_LOG_COLLECT:=1): 'AMDGIM_FEATURE_ERROR_LOG_COLLECT', (AMDGIM_FEATURE_GIM_LOAD_UCODES:=2): 'AMDGIM_FEATURE_GIM_LOAD_UCODES', (AMDGIM_FEATURE_GIM_FLR_VRAMLOST:=4): 'AMDGIM_FEATURE_GIM_FLR_VRAMLOST', (AMDGIM_FEATURE_GIM_MM_BW_MGR:=8): 'AMDGIM_FEATURE_GIM_MM_BW_MGR', (AMDGIM_FEATURE_PP_ONE_VF:=16): 'AMDGIM_FEATURE_PP_ONE_VF', (AMDGIM_FEATURE_INDIRECT_REG_ACCESS:=32): 'AMDGIM_FEATURE_INDIRECT_REG_ACCESS', (AMDGIM_FEATURE_AV1_SUPPORT:=64): 'AMDGIM_FEATURE_AV1_SUPPORT', (AMDGIM_FEATURE_VCN_RB_DECOUPLE:=128): 'AMDGIM_FEATURE_VCN_RB_DECOUPLE', (AMDGIM_FEATURE_MES_INFO_ENABLE:=256): 'AMDGIM_FEATURE_MES_INFO_ENABLE', (AMDGIM_FEATURE_RAS_CAPS:=512): 'AMDGIM_FEATURE_RAS_CAPS', (AMDGIM_FEATURE_RAS_TELEMETRY:=1024): 'AMDGIM_FEATURE_RAS_TELEMETRY', (AMDGIM_FEATURE_RAS_CPER:=2048): 'AMDGIM_FEATURE_RAS_CPER', (AMDGIM_FEATURE_XGMI_TA_EXT_PEER_LINK:=4096): 'AMDGIM_FEATURE_XGMI_TA_EXT_PEER_LINK'}
-enum_AMDGIM_REG_ACCESS_FLAG: dict[int, str] = {(AMDGIM_FEATURE_IH_REG_PSP_EN:=1): 'AMDGIM_FEATURE_IH_REG_PSP_EN', (AMDGIM_FEATURE_MMHUB_REG_RLC_EN:=2): 'AMDGIM_FEATURE_MMHUB_REG_RLC_EN', (AMDGIM_FEATURE_GC_REG_RLC_EN:=4): 'AMDGIM_FEATURE_GC_REG_RLC_EN', (AMDGIM_FEATURE_L1_TLB_CNTL_PSP_EN:=8): 'AMDGIM_FEATURE_L1_TLB_CNTL_PSP_EN', (AMDGIM_FEATURE_REG_ACCESS_SQ_CONFIG:=16): 'AMDGIM_FEATURE_REG_ACCESS_SQ_CONFIG'}
-@c.record
-class struct_amdgim_pf2vf_info_v1(c.Struct):
-  SIZE = 44
-  header: struct_amd_sriov_msg_pf2vf_info_header
-  uvd_enc_max_pixels_count: int
-  uvd_enc_max_bandwidth: int
-  vce_enc_max_pixels_count: int
-  vce_enc_max_bandwidth: int
-  mecfw_kboffset: int
-  feature_flags: int
-  checksum: int
-struct_amdgim_pf2vf_info_v1.register_fields([('header', struct_amd_sriov_msg_pf2vf_info_header, 0), ('uvd_enc_max_pixels_count', ctypes.c_uint32, 16), ('uvd_enc_max_bandwidth', ctypes.c_uint32, 20), ('vce_enc_max_pixels_count', ctypes.c_uint32, 24), ('vce_enc_max_bandwidth', ctypes.c_uint32, 28), ('mecfw_kboffset', ctypes.c_uint32, 32), ('feature_flags', ctypes.c_uint32, 36), ('checksum', ctypes.c_uint32, 40)])
-@c.record
-class struct_amdgim_vf2pf_info_v1(c.Struct):
-  SIZE = 128
-  header: struct_amd_sriov_msg_vf2pf_info_header
-  driver_version: c.Array[ctypes.c_char, Literal[64]]
-  driver_cert: int
-  os_info: int
-  fb_usage: int
-  gfx_usage: int
-  gfx_health: int
-  compute_usage: int
-  compute_health: int
-  vce_enc_usage: int
-  vce_enc_health: int
-  uvd_enc_usage: int
-  uvd_enc_health: int
-  checksum: int
-struct_amdgim_vf2pf_info_v1.register_fields([('header', struct_amd_sriov_msg_vf2pf_info_header, 0), ('driver_version', c.Array[ctypes.c_char, Literal[64]], 16), ('driver_cert', ctypes.c_uint32, 80), ('os_info', ctypes.c_uint32, 84), ('fb_usage', ctypes.c_uint32, 88), ('gfx_usage', ctypes.c_uint32, 92), ('gfx_health', ctypes.c_uint32, 96), ('compute_usage', ctypes.c_uint32, 100), ('compute_health', ctypes.c_uint32, 104), ('vce_enc_usage', ctypes.c_uint32, 108), ('vce_enc_health', ctypes.c_uint32, 112), ('uvd_enc_usage', ctypes.c_uint32, 116), ('uvd_enc_health', ctypes.c_uint32, 120), ('checksum', ctypes.c_uint32, 124)])
-@c.record
-class struct_amdgim_vf2pf_info_v2(c.Struct):
-  SIZE = 1024
-  header: struct_amd_sriov_msg_vf2pf_info_header
-  checksum: int
-  driver_version: c.Array[ctypes.c_ubyte, Literal[64]]
-  driver_cert: int
-  os_info: int
-  fb_usage: int
-  gfx_usage: int
-  gfx_health: int
-  compute_usage: int
-  compute_health: int
-  vce_enc_usage: int
-  vce_enc_health: int
-  uvd_enc_usage: int
-  uvd_enc_health: int
-  reserved: c.Array[ctypes.c_uint32, Literal[224]]
-struct_amdgim_vf2pf_info_v2.register_fields([('header', struct_amd_sriov_msg_vf2pf_info_header, 0), ('checksum', uint32_t, 16), ('driver_version', c.Array[uint8_t, Literal[64]], 20), ('driver_cert', uint32_t, 84), ('os_info', uint32_t, 88), ('fb_usage', uint32_t, 92), ('gfx_usage', uint32_t, 96), ('gfx_health', uint32_t, 100), ('compute_usage', uint32_t, 104), ('compute_health', uint32_t, 108), ('vce_enc_usage', uint32_t, 112), ('vce_enc_health', uint32_t, 116), ('uvd_enc_usage', uint32_t, 120), ('uvd_enc_health', uint32_t, 124), ('reserved', c.Array[uint32_t, Literal[224]], 128)])
-@c.record
-class struct_amdgpu_virt_ras_err_handler_data(c.Struct):
-  SIZE = 24
-  bps: c.POINTER[struct_eeprom_table_record]
-  bps_bo: c.POINTER[c.POINTER[struct_amdgpu_bo]]
-  count: int
-  last_reserved: int
-class struct_eeprom_table_record(c.Struct): pass
-struct_amdgpu_virt_ras_err_handler_data.register_fields([('bps', c.POINTER[struct_eeprom_table_record], 0), ('bps_bo', c.POINTER[c.POINTER[struct_amdgpu_bo]], 8), ('count', ctypes.c_int32, 16), ('last_reserved', ctypes.c_int32, 20)])
-class struct_amdgpu_virt_ras(c.Struct): pass
-@c.record
-class struct_amdgpu_virt_region(c.Struct):
-  SIZE = 8
-  offset: int
-  size_kb: int
-struct_amdgpu_virt_region.register_fields([('offset', uint32_t, 0), ('size_kb', uint32_t, 4)])
-class struct_amdgpu_virt(c.Struct): pass
-class struct_amdgpu_video_codec_info(c.Struct): pass
 AMDGPU_VM_MAX_UPDATE_SIZE = 0x3FFFF
 AMDGPU_PTE_VALID = (1 << 0)
 AMDGPU_PTE_SYSTEM = (1 << 1)
@@ -4569,47 +4465,4 @@ SDMA0_5_0__SRCID__SDMA_CTXEMPTY = 243
 SDMA0_5_0__SRCID__SDMA_DOORBELL_INVALID = 244
 SDMA0_5_0__SRCID__SDMA_FROZEN = 245
 SDMA0_5_0__SRCID__SDMA_POLL_TIMEOUT = 246
-SDMA0_5_0__SRCID__SDMA_SRBMWRITE = 247
-NV_MAILBOX_POLL_ACK_TIMEDOUT = 500
-NV_MAILBOX_POLL_MSG_TIMEDOUT = 15000
-NV_MAILBOX_POLL_FLR_TIMEDOUT = 10000
-NV_MAILBOX_POLL_MSG_REP_MAX = 11
-mmMAILBOX_CONTROL = 0xE5E
-NV_MAIBOX_CONTROL_TRN_OFFSET_BYTE = (mmMAILBOX_CONTROL * 4)
-NV_MAIBOX_CONTROL_RCV_OFFSET_BYTE = (NV_MAIBOX_CONTROL_TRN_OFFSET_BYTE + 1)
-mmMAILBOX_MSGBUF_TRN_DW0 = 0xE56
-mmMAILBOX_MSGBUF_TRN_DW1 = 0xE57
-mmMAILBOX_MSGBUF_TRN_DW2 = 0xE58
-mmMAILBOX_MSGBUF_TRN_DW3 = 0xE59
-mmMAILBOX_MSGBUF_RCV_DW0 = 0xE5A
-mmMAILBOX_MSGBUF_RCV_DW1 = 0xE5B
-mmMAILBOX_MSGBUF_RCV_DW2 = 0xE5C
-mmMAILBOX_MSGBUF_RCV_DW3 = 0xE5D
-mmMAILBOX_INT_CNTL = 0xE5F
-AMDGPU_SRIOV_CAPS_SRIOV_VBIOS = (1 << 0)
-AMDGPU_SRIOV_CAPS_ENABLE_IOV = (1 << 1)
-AMDGPU_SRIOV_CAPS_IS_VF = (1 << 2)
-AMDGPU_PASSTHROUGH_MODE = (1 << 3)
-AMDGPU_SRIOV_CAPS_RUNTIME = (1 << 4)
-AMDGPU_VF_MMIO_ACCESS_PROTECT = (1 << 5)
-AMDGPU_RLCG_GC_WRITE_LEGACY = (0x8 << 28)
-AMDGPU_RLCG_GC_WRITE = (0x0 << 28)
-AMDGPU_RLCG_GC_READ = (0x1 << 28)
-AMDGPU_RLCG_MMHUB_WRITE = (0x2 << 28)
-AMDGPU_RLCG_VFGATE_DISABLED = 0x4000000
-AMDGPU_RLCG_WRONG_OPERATION_TYPE = 0x2000000
-AMDGPU_RLCG_REG_NOT_IN_RANGE = 0x1000000
-AMDGPU_RLCG_SCRATCH1_ADDRESS_MASK = 0xFFFFF
-AMDGPU_RLCG_SCRATCH1_ERROR_MASK = 0xF000000
-mmRCC_IOV_FUNC_IDENTIFIER = 0xDE5
-mmBIF_IOV_FUNC_IDENTIFIER = 0x1503
-AMDGPU_VF2PF_UPDATE_MAX_RETRY_LIMIT = 2
-AMDGPU_SRIOV_CRIT_DATA_SIGNATURE = "INDA"
-AMDGPU_SRIOV_CRIT_DATA_SIG_LEN = 4
-AMDGPU_VF_ERROR_ENTRY_SIZE = 16
-AMDGIM_DATAEXCHANGE_OFFSET = (64 * 1024)
-AMDGIM_GET_STRUCTURE_RESERVED_SIZE = lambda total,u8,u16,u32,u64: (total - (((u8)+3) / 4 + ((u16)+1) / 2 + (u32) + (u64)*2)) # type: ignore
-AMDGPU_VIRT_CAPS_LIST = lambda X: X(AMDGPU_VIRT_CAP_POWER_LIMIT) # type: ignore
-amdgpu_sriov_fullaccess = lambda adev: (amdgpu_sriov_vf((adev)) and not amdgpu_sriov_runtime((adev))) # type: ignore
-amdgpu_sriov_rlcg_error_report_enabled = lambda adev: (amdgpu_sriov_reg_indirect_mmhub(adev) or amdgpu_sriov_reg_indirect_gc(adev)) # type: ignore
-amdgpu_sriov_multi_vf_mode = lambda adev: (amdgpu_sriov_vf(adev) and not amdgpu_sriov_is_pp_one_vf(adev)) # type: ignore
+SDMA0_5_0__SRCID__SDMA_SRBMWRITE = 247
@@ -510,10 +510,6 @@ gc_9_4_3 = {
  'regTCP_PERFCOUNTER_FILTER': (4793, 0, {'buffer': (0, 0), 'flat': (1, 1), 'dim': (2, 4), 'data_format': (5, 10), 'num_format': (11, 14), 'sw_mode': (15, 19), 'num_samples': (20, 21), 'opcode_type': (22, 24), 'glc': (25, 25), 'slc': (26, 26), 'compression_enable': (27, 27), 'addr_mode': (28, 30)}),
  'regTCP_PERFCOUNTER_FILTER_EN': (4794, 0, {'buffer': (0, 0), 'flat': (1, 1), 'dim': (2, 2), 'data_format': (3, 3), 'num_format': (4, 4), 'sw_mode': (5, 5), 'num_samples': (6, 6), 'opcode_type': (7, 7), 'glc': (8, 8), 'slc': (9, 9), 'compression_enable': (10, 10), 'addr_mode': (11, 11)}),
  'regCP_PERFMON_CNTX_CNTL': (216, 1, {'perfmon_enable': (31, 31)}),
-  'regSCRATCH_REG0': (8256, 1, {'scratch_reg0': (0, 31)}),
-  'regSCRATCH_REG1': (8257, 1, {'scratch_reg1': (0, 31)}),
-  'regSCRATCH_REG2': (8258, 1, {'scratch_reg2': (0, 31)}),
-  'regSCRATCH_REG3': (8259, 1, {'scratch_reg3': (0, 31)}),
  'regSCRATCH_REG6': (8262, 1, {'scratch_reg6': (0, 31)}),
  'regSCRATCH_REG7': (8263, 1, {'scratch_reg7': (0, 31)}),
  'regCP_COHER_START_DELAY': (8315, 1, {'start_delay_count': (0, 5)}),
@@ -1795,10 +1791,6 @@ gc_11_0_0 = {
  'regTCP_CNTL2': (6563, 1, {'ls_disable_clocks': (0, 7), 'tcp_fmt_mgcg_disable': (8, 8), 'tcpf_latency_bypass_disable': (9, 9), 'tcp_write_data_mgcg_disable': (10, 10), 'tcp_inner_block_mgcg_disable': (11, 11), 'tcp_adrs_img_calc_mgcg_disable': (12, 12), 'v64_combine_enable': (13, 13), 'tagram_addr_swizzle_disable': (14, 14), 'return_order_override': (15, 15), 'power_opt_disable': (16, 16), 'gcr_rsp_fgcg_disable': (17, 17), 'perf_en_override': (18, 19), 'tc_td_ram_clken_disable': (20, 20), 'tc_td_data_clken_disable': (21, 21), 'tcp_gl1_req_clken_disable': (22, 22), 'tcp_gl1r_src_clken_disable': (23, 23), 'spare_bit': (26, 26), 'tagram_xy_bias_override': (27, 28), 'tcp_req_mgcg_disable': (29, 29), 'tcp_miss_mgcg_disable': (30, 30), 'disable_mipmap_param_calc_self_gating': (31, 31)}),
  'regTCP_DEBUG_INDEX': (6565, 1, {'index': (0, 4)}),
  'regTCP_DEBUG_DATA': (6566, 1, {'data': (0, 17)}),
-  'regSCRATCH_REG0': (8256, 1, {'scratch_reg0': (0, 31)}),
-  'regSCRATCH_REG1': (8257, 1, {'scratch_reg1': (0, 31)}),
-  'regSCRATCH_REG2': (8258, 1, {'scratch_reg2': (0, 31)}),
-  'regSCRATCH_REG3': (8259, 1, {'scratch_reg3': (0, 31)}),
  'regSCRATCH_REG6': (8262, 1, {'scratch_reg6': (0, 31)}),
  'regSCRATCH_REG7': (8263, 1, {'scratch_reg7': (0, 31)}),
  'regRLC_GPM_PERF_COUNT_0': (8512, 1, {'feature_sel': (0, 3), 'se_index': (4, 7), 'sa_index': (8, 11), 'wgp_index': (12, 15), 'event_sel': (16, 17), 'unused': (18, 19), 'enable': (20, 20), 'reserved': (21, 31)}),
@@ -3374,10 +3366,6 @@ gc_11_0_3 = {
  'regTCP_CNTL': (6562, 1, {'force_hit': (0, 0), 'force_miss': (1, 1), 'flat_buf_cache_swizzle': (5, 5), 'td_data_en_override': (6, 6), 'enable_128b_dcc_comp_read_for_indep64': (7, 7), 'disable_write_combining': (9, 9), 'force_eow_total_cnt': (15, 20), 'force_eow_set_cnt': (22, 26), 'disable_z_map': (28, 28), 'force_order_between_read_write_to_same_address': (29, 29), 'astc_ve_msb_tolerant': (31, 31)}),
  'regTCP_CNTL2': (6563, 1, {'ls_disable_clocks': (0, 7), 'tcp_fmt_mgcg_disable': (8, 8), 'tcpf_latency_bypass_disable': (9, 9), 'tcp_write_data_mgcg_disable': (10, 10), 'tcp_inner_block_mgcg_disable': (11, 11), 'tcp_adrs_img_calc_mgcg_disable': (12, 12), 'v64_combine_enable': (13, 13), 'tagram_addr_swizzle_disable': (14, 14), 'return_order_override': (15, 15), 'power_opt_disable': (16, 16), 'gcr_rsp_fgcg_disable': (17, 17), 'perf_en_override': (18, 19), 'tc_td_ram_clken_disable': (20, 20), 'tc_td_data_clken_disable': (21, 21), 'tcp_gl1_req_clken_disable': (22, 22), 'tcp_gl1r_src_clken_disable': (23, 23), 'spare_bit': (26, 26), 'tagram_xy_bias_override': (27, 28), 'tcp_req_mgcg_disable': (29, 29), 'tcp_miss_mgcg_disable': (30, 30), 'disable_mipmap_param_calc_self_gating': (31, 31)}),
  'regTCP_CREDIT': (6564, 1, {'lfifo_ram_depth': (0, 9), 'gl1_req_credit': (10, 15), 'req_fifo_credit': (16, 22), 'td_ram_credit': (23, 27), 'td_data_credit': (29, 31)}),
-  'regSCRATCH_REG0': (8256, 1, {'scratch_reg0': (0, 31)}),
-  'regSCRATCH_REG1': (8257, 1, {'scratch_reg1': (0, 31)}),
-  'regSCRATCH_REG2': (8258, 1, {'scratch_reg2': (0, 31)}),
-  'regSCRATCH_REG3': (8259, 1, {'scratch_reg3': (0, 31)}),
  'regSCRATCH_REG6': (8262, 1, {'scratch_reg6': (0, 31)}),
  'regSCRATCH_REG7': (8263, 1, {'scratch_reg7': (0, 31)}),
  'regRLC_GPM_PERF_COUNT_0': (8512, 1, {'feature_sel': (0, 3), 'se_index': (4, 7), 'sa_index': (8, 11), 'wgp_index': (12, 15), 'event_sel': (16, 17), 'unused': (18, 19), 'enable': (20, 20), 'reserved': (21, 31)}),
@@ -4803,10 +4791,6 @@ gc_11_5_0 = {
  'regTCP_STATUS': (6561, 1, {'tcp_busy': (0, 0), 'input_busy': (1, 1), 'adrs_busy': (2, 2), 'tagrams_busy': (3, 3), 'cntrl_busy': (4, 4), 'lfifo_busy': (5, 5), 'read_busy': (6, 6), 'format_busy': (7, 7), 'vm_busy': (8, 8), 'memif_busy': (9, 9), 'gcr_busy': (10, 10), 'ofifo_busy': (11, 11), 'ofifo_queue_busy': (12, 13), 'xnack_prt': (15, 15)}),
  'regTCP_CNTL': (6562, 1, {'force_hit': (0, 0), 'force_miss': (1, 1), 'flat_buf_cache_swizzle': (5, 5), 'td_data_en_override': (6, 6), 'enable_128b_dcc_comp_read_for_indep64': (7, 7), 'disable_decompression_power_opt': (8, 8), 'disable_write_combining': (9, 9), 'force_eow_total_cnt': (15, 20), 'force_eow_set_cnt': (22, 26), 'disable_z_map': (28, 28), 'force_order_between_read_write_to_same_address': (29, 29), 'astc_ve_msb_tolerant': (31, 31)}),
  'regTCP_CNTL2': (6563, 1, {'ls_disable_clocks': (0, 7), 'tcp_fmt_mgcg_disable': (8, 8), 'tcpf_latency_bypass_disable': (9, 9), 'tcp_write_data_mgcg_disable': (10, 10), 'tcp_inner_block_mgcg_disable': (11, 11), 'tcp_adrs_img_calc_mgcg_disable': (12, 12), 'v64_combine_enable': (13, 13), 'tagram_addr_swizzle_disable': (14, 14), 'return_order_override': (15, 15), 'power_opt_disable': (16, 16), 'gcr_rsp_fgcg_disable': (17, 17), 'perf_en_override': (18, 19), 'tcp_gl1_req_clken_disable': (22, 22), 'tcp_gl1r_src_clken_disable': (23, 23), 'tcp_force_2x_to_load': (24, 24), 'spare_bit': (26, 26), 'tagram_xy_bias_override': (27, 28), 'tcp_req_mgcg_disable': (29, 29), 'tcp_miss_mgcg_disable': (30, 30), 'disable_mipmap_param_calc_self_gating': (31, 31)}),
-  'regSCRATCH_REG0': (8256, 1, {'scratch_reg0': (0, 31)}),
-  'regSCRATCH_REG1': (8257, 1, {'scratch_reg1': (0, 31)}),
-  'regSCRATCH_REG2': (8258, 1, {'scratch_reg2': (0, 31)}),
-  'regSCRATCH_REG3': (8259, 1, {'scratch_reg3': (0, 31)}),
  'regSCRATCH_REG6': (8262, 1, {'scratch_reg6': (0, 31)}),
  'regSCRATCH_REG7': (8263, 1, {'scratch_reg7': (0, 31)}),
  'regRLC_GPM_PERF_COUNT_0': (8512, 1, {'feature_sel': (0, 3), 'se_index': (4, 7), 'sa_index': (8, 11), 'wgp_index': (12, 15), 'event_sel': (16, 17), 'unused': (18, 19), 'enable': (20, 20), 'reserved': (21, 31)}),
@@ -6064,10 +6048,6 @@ gc_12_0_0 = {
  'regCP_ME_CNTL': (2051, 1, {'ce_invalidate_icache': (4, 4), 'pfp_invalidate_icache': (6, 6), 'me_invalidate_icache': (8, 8), 'pfp_pipe0_disable': (12, 12), 'pfp_pipe1_disable': (13, 13), 'me_pipe0_disable': (14, 14), 'me_pipe1_disable': (15, 15), 'ce_pipe0_reset': (16, 16), 'ce_pipe1_reset': (17, 17), 'pfp_pipe0_reset': (18, 18), 'pfp_pipe1_reset': (19, 19), 'me_pipe0_reset': (20, 20), 'me_pipe1_reset': (21, 21), 'ce_halt': (24, 24), 'ce_step': (25, 25), 'pfp_halt': (26, 26), 'pfp_step': (27, 27), 'me_halt': (28, 28), 'me_step': (29, 29)}),
  'regGRBM_GFX_CNTL': (2304, 1, {'pipeid': (0, 1), 'meid': (2, 3), 'vmid': (4, 7), 'queueid': (8, 10), 'ctxid': (11, 13)}),
  'regGRBM_NOWHERE': (2305, 1, {'data': (0, 31)}),
-  'regSCRATCH_REG0': (8256, 1, {'scratch_reg0': (0, 31)}),
-  'regSCRATCH_REG1': (8257, 1, {'scratch_reg1': (0, 31)}),
-  'regSCRATCH_REG2': (8258, 1, {'scratch_reg2': (0, 31)}),
-  'regSCRATCH_REG3': (8259, 1, {'scratch_reg3': (0, 31)}),
  'regSCRATCH_REG6': (8262, 1, {'scratch_reg6': (0, 31)}),
  'regSCRATCH_REG7': (8263, 1, {'scratch_reg7': (0, 31)}),
  'regRLC_GPM_PERF_COUNT_0': (8512, 1, {'feature_sel': (0, 3), 'se_index': (4, 7), 'sa_index': (8, 11), 'wgp_index': (12, 15), 'event_sel': (16, 17), 'unused': (18, 19), 'enable': (20, 20), 'reserved': (21, 31)}),
@@ -8997,7 +8977,6 @@ nbio_4_3_0 = {
  'regBIF_BX_PF0_GPU_HDP_FLUSH_REQ': (262, 2, {'cp0': (0, 0), 'cp1': (1, 1), 'cp2': (2, 2), 'cp3': (3, 3), 'cp4': (4, 4), 'cp5': (5, 5), 'cp6': (6, 6), 'cp7': (7, 7), 'cp8': (8, 8), 'cp9': (9, 9), 'sdma0': (10, 10), 'sdma1': (11, 11), 'rsvd_eng0': (12, 12), 'rsvd_eng1': (13, 13), 'rsvd_eng2': (14, 14), 'rsvd_eng3': (15, 15), 'rsvd_eng4': (16, 16), 'rsvd_eng5': (17, 17), 'rsvd_eng6': (18, 18), 'rsvd_eng7': (19, 19), 'rsvd_eng8': (20, 20), 'rsvd_eng9': (21, 21), 'rsvd_eng10': (22, 22), 'rsvd_eng11': (23, 23), 'rsvd_eng12': (24, 24), 'rsvd_eng13': (25, 25), 'rsvd_eng14': (26, 26), 'rsvd_eng15': (27, 27), 'rsvd_eng16': (28, 28), 'rsvd_eng17': (29, 29), 'rsvd_eng18': (30, 30), 'rsvd_eng19': (31, 31)}),
  'regBIF_BX_PF0_GPU_HDP_FLUSH_DONE': (263, 2, {'cp0': (0, 0), 'cp1': (1, 1), 'cp2': (2, 2), 'cp3': (3, 3), 'cp4': (4, 4), 'cp5': (5, 5), 'cp6': (6, 6), 'cp7': (7, 7), 'cp8': (8, 8), 'cp9': (9, 9), 'sdma0': (10, 10), 'sdma1': (11, 11), 'rsvd_eng0': (12, 12), 'rsvd_eng1': (13, 13), 'rsvd_eng2': (14, 14), 'rsvd_eng3': (15, 15), 'rsvd_eng4': (16, 16), 'rsvd_eng5': (17, 17), 'rsvd_eng6': (18, 18), 'rsvd_eng7': (19, 19), 'rsvd_eng8': (20, 20), 'rsvd_eng9': (21, 21), 'rsvd_eng10': (22, 22), 'rsvd_eng11': (23, 23), 'rsvd_eng12': (24, 24), 'rsvd_eng13': (25, 25), 'rsvd_eng14': (26, 26), 'rsvd_eng15': (27, 27), 'rsvd_eng16': (28, 28), 'rsvd_eng17': (29, 29), 'rsvd_eng18': (30, 30), 'rsvd_eng19': (31, 31)}),
  'regRCC_DEV0_EPF0_RCC_DOORBELL_APER_EN': (192, 2, {'bif_doorbell_aper_en': (0, 0)}),
-  'regBIF_BX_DEV0_EPF0_VF0_HDP_MEM_COHERENCY_FLUSH_CNTL': (247, 2, {'hdp_mem_flush_addr': (0, 0)}),
  'regRCC_DEV0_EPF2_STRAP2': (53506, 5, {'strap_no_soft_reset_dev0_f2': (7, 7), 'strap_resize_bar_en_dev0_f2': (8, 8), 'strap_max_pasid_width_dev0_f2': (9, 13), 'strap_msi_pervector_mask_cap_dev0_f2': (14, 14), 'strap_aer_en_dev0_f2': (16, 16), 'strap_acs_en_dev0_f2': (17, 17), 'strap_cpl_abort_err_en_dev0_f2': (20, 20), 'strap_dpa_en_dev0_f2': (21, 21), 'strap_vc_en_dev0_f2': (23, 23), 'strap_msi_multi_cap_dev0_f2': (24, 26), 'strap_pasid_en_dev0_f2': (28, 28), 'strap_pasid_exe_permission_supported_dev0_f2': (29, 29), 'strap_pasid_global_invalidate_supported_dev0_f2': (30, 30), 'strap_pasid_priv_mode_supported_dev0_f2': (31, 31)}),
  'regRCC_DEV0_EPF2_STRAP20': (53524, 5, {}),
  'regBIF_BX_PF1_GPU_HDP_FLUSH_REQ': (36390, 5, {'cp0': (0, 0), 'cp1': (1, 1), 'cp2': (2, 2), 'cp3': (3, 3), 'cp4': (4, 4), 'cp5': (5, 5), 'cp6': (6, 6), 'cp7': (7, 7), 'cp8': (8, 8), 'cp9': (9, 9), 'sdma0': (10, 10), 'sdma1': (11, 11), 'rsvd_eng0': (12, 12), 'rsvd_eng1': (13, 13), 'rsvd_eng2': (14, 14), 'rsvd_eng3': (15, 15), 'rsvd_eng4': (16, 16), 'rsvd_eng5': (17, 17), 'rsvd_eng6': (18, 18), 'rsvd_eng7': (19, 19), 'rsvd_eng8': (20, 20), 'rsvd_eng9': (21, 21), 'rsvd_eng10': (22, 22), 'rsvd_eng11': (23, 23), 'rsvd_eng12': (24, 24), 'rsvd_eng13': (25, 25), 'rsvd_eng14': (26, 26), 'rsvd_eng15': (27, 27), 'rsvd_eng16': (28, 28), 'rsvd_eng17': (29, 29), 'rsvd_eng18': (30, 30), 'rsvd_eng19': (31, 31)}),
@@ -9105,7 +9084,6 @@ nbio_7_9_0 = {
  'regS2A_DOORBELL_ENTRY_14_CTRL': (31374, 5, {'s2a_doorbell_port14_enable': (0, 0), 's2a_doorbell_port14_awid': (1, 5), 's2a_doorbell_port14_fence_enable': (6, 6), 's2a_doorbell_port14_range_offset': (7, 16), 's2a_doorbell_port14_range_size': (17, 24), 's2a_doorbell_port14_64bit_support_dis': (25, 25), 's2a_doorbell_port14_need_deduct_range_offset': (26, 26), 's2a_doorbell_port14_awaddr_31_28_value': (28, 31)}),
  'regS2A_DOORBELL_ENTRY_15_CTRL': (31375, 5, {'s2a_doorbell_port15_enable': (0, 0), 's2a_doorbell_port15_awid': (1, 5), 's2a_doorbell_port15_fence_enable': (6, 6), 's2a_doorbell_port15_range_offset': (7, 16), 's2a_doorbell_port15_range_size': (17, 24), 's2a_doorbell_port15_64bit_support_dis': (25, 25), 's2a_doorbell_port15_need_deduct_range_offset': (26, 26), 's2a_doorbell_port15_awaddr_31_28_value': (28, 31)}),
  'regXCC_DOORBELL_FENCE': (29708, 5, {'xcc_0_doorbell_fence': (0, 0), 'xcc_1_doorbell_fence': (1, 1), 'xcc_2_doorbell_fence': (2, 2), 'xcc_3_doorbell_fence': (3, 3), 'xcc_4_doorbell_fence': (4, 4), 'xcc_5_doorbell_fence': (5, 5), 'xcc_6_doorbell_fence': (6, 6), 'xcc_7_doorbell_fence': (7, 7), 'shub_slv_mode': (16, 16), 'rmote_cp_sent': (17, 17), 'cp_0_sent': (18, 18), 'cp_1_sent': (19, 19), 'cp_2_sent': (20, 20), 'cp_3_sent': (21, 21), 'cp_4_sent': (22, 22), 'cp_5_sent': (23, 23), 'cp_6_sent': (24, 24), 'cp_7_sent': (25, 25), 'remote_client_sent': (26, 26), 'remote_client_clr_pending': (27, 27)}),
-  'regBIF_BX_DEV0_EPF0_VF0_HDP_MEM_COHERENCY_FLUSH_CNTL': (247, 2, {'hdp_mem_flush_addr': (0, 0)}),
 }
 nbio_7_11_0 = {
  'regBIF_BX0_PCIE_INDEX2': (32782, 0, {'pcie_index2': (0, 31)}),
@@ -9153,7 +9131,6 @@ nbif_6_3_1 = {
  'regRCC_DEV0_EPF2_STRAP20': (53524, 5, {}),
  'regBIF_BX_PF1_GPU_HDP_FLUSH_REQ': (36390, 5, {'cp0': (0, 0), 'cp1': (1, 1), 'cp2': (2, 2), 'cp3': (3, 3), 'cp4': (4, 4), 'cp5': (5, 5), 'cp6': (6, 6), 'cp7': (7, 7), 'cp8': (8, 8), 'cp9': (9, 9), 'sdma0': (10, 10), 'sdma1': (11, 11), 'rsvd_eng0': (12, 12), 'rsvd_eng1': (13, 13), 'rsvd_eng2': (14, 14), 'rsvd_eng3': (15, 15), 'rsvd_eng4': (16, 16), 'rsvd_eng5': (17, 17), 'rsvd_eng6': (18, 18), 'rsvd_eng7': (19, 19), 'rsvd_eng8': (20, 20), 'rsvd_eng9': (21, 21), 'rsvd_eng10': (22, 22), 'rsvd_eng11': (23, 23), 'rsvd_eng12': (24, 24), 'rsvd_eng13': (25, 25), 'rsvd_eng14': (26, 26), 'rsvd_eng15': (27, 27), 'rsvd_eng16': (28, 28), 'rsvd_eng17': (29, 29), 'rsvd_eng18': (30, 30), 'rsvd_eng19': (31, 31)}),
  'regBIF_BX_PF1_GPU_HDP_FLUSH_DONE': (36391, 5, {'cp0': (0, 0), 'cp1': (1, 1), 'cp2': (2, 2), 'cp3': (3, 3), 'cp4': (4, 4), 'cp5': (5, 5), 'cp6': (6, 6), 'cp7': (7, 7), 'cp8': (8, 8), 'cp9': (9, 9), 'sdma0': (10, 10), 'sdma1': (11, 11), 'rsvd_eng0': (12, 12), 'rsvd_eng1': (13, 13), 'rsvd_eng2': (14, 14), 'rsvd_eng3': (15, 15), 'rsvd_eng4': (16, 16), 'rsvd_eng5': (17, 17), 'rsvd_eng6': (18, 18), 'rsvd_eng7': (19, 19), 'rsvd_eng8': (20, 20), 'rsvd_eng9': (21, 21), 'rsvd_eng10': (22, 22), 'rsvd_eng11': (23, 23), 'rsvd_eng12': (24, 24), 'rsvd_eng13': (25, 25), 'rsvd_eng14': (26, 26), 'rsvd_eng15': (27, 27), 'rsvd_eng16': (28, 28), 'rsvd_eng17': (29, 29), 'rsvd_eng18': (30, 30), 'rsvd_eng19': (31, 31)}),
-  'regBIF_BX_DEV0_EPF0_VF0_HDP_MEM_COHERENCY_FLUSH_CNTL': (247, 2, {'hdp_mem_flush_addr': (0, 0)}),
 }
 mp_11_0_0 = {
  'mmMP0_SMN_C2PMSG_32': (96, 0, {'content': (0, 31)}),
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	feb860a7b7	also skip casted consts in expand_broadcast	2026-08-26 15:22:07 -07:00
geohot	80cde0d70a	don't expand consts in expand_broadcast, fix SPEC=2	2026-08-26 15:10:01 -07:00
geohot	03a593c601	remove broadcast handling from rangeify, prepare handles it now	2026-08-26 13:18:24 -07:00
geohot	58134bfa59	Merge remote-tracking branch 'origin/master' into expand_in_prepare # Conflicts: # tinygrad/schedule/indexing.py	2026-08-26 13:08:41 -07:00
geohot	1e1e68a2a6	fix	2026-08-26 13:06:04 -07:00
geohot	23d5efe25d	move expand into prepare	2026-08-26 11:13:23 -07:00