low rank call in kernel graph

fix custom_kernel wrapper
call has rank <= 1
2026-08-28 04:36:07 +00:00 · 2026-08-27 10:34:54 -07:00 · 2026-08-27 10:18:37 -07:00 · 2026-08-27 10:15:35 -07:00 · 2026-08-27 10:47:21 -04:00 · 2026-08-27 17:42:37 +03:00
148 changed files with 9483 additions and 2236 deletions
@@ -61,6 +61,7 @@ runs:
        echo "MAX_BUFFER_SIZE=300000000" >> "$GITHUB_ENV"
        if [[ "$RUNNER_OS" == "Linux" ]]; then
          echo "VIRTUAL_ENV=/opt/venv/${{ inputs.python-version }}" >> "$GITHUB_ENV"
+          echo "UV_PYTHON_INSTALL_DIR=/opt/python" >> "$GITHUB_ENV"
        else
          echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
        fi
@@ -70,11 +71,6 @@ runs:
      with:
        enable-cache: 'false' # see below for manual caching

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

    - name: Cache Python packages (PR)
@@ -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"
+        python3 -c "from tinygrad.runtime.autogen import mlx5, bnxt"
        python3 -c "from tinygrad.runtime.autogen import ggml_common"
        REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
    - name: Check for differences
@@ -88,13 +88,13 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 30
    defaults:
      run:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: "0"
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -102,12 +102,11 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./extra/hcq/hcq_smi.py amd rmmod --expect
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -117,10 +116,10 @@ jobs:
      run: python3 test/external/process_replay/reset.py
    - name: Run llama3.2
      run: BENCHMARK_LOG=llama32_3b-f16 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m llama3.2:3b-f16 --benchmark --warmup
-    - name: Run qwen3.6
-      # qwen3.6:35b-a3b doesn't fit on mac
+    - name: Run qwen3.8
+      # qwen3.8:27b doesn't fit on mac
      if: ${{ matrix.dev != 'METAL' }}
-      run: BENCHMARK_LOG=qwen36_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.6:35b-a3b --benchmark --warmup
+      run: BENCHMARK_LOG=qwen38_27b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.8:27b --benchmark --warmup
    - name: Run olmoe
      # just metal for now
      if: ${{ matrix.dev == 'METAL' }}
@@ -135,13 +134,13 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 10
    defaults:
      run:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: "0"
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -149,12 +148,11 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./extra/hcq/hcq_smi.py amd rmmod --expect
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -184,13 +182,13 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 5
    defaults:
      run:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: "0"
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -198,12 +196,11 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./extra/hcq/hcq_smi.py amd rmmod --expect
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
    - name: Symlink models and datasets
      run: |
        mkdir -p extra/datasets
@@ -227,13 +224,13 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 15
    defaults:
      run:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: "0"
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -241,12 +238,11 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./extra/hcq/hcq_smi.py amd rmmod --expect
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -273,13 +269,13 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 20
    defaults:
      run:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: "0"
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -287,12 +283,11 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./extra/hcq/hcq_smi.py amd rmmod --expect
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -327,7 +322,7 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 10
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -340,12 +335,11 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./extra/hcq/hcq_smi.py amd rmmod --expect
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -422,7 +416,7 @@ jobs:
  testusbgpu:
    name: UsbGPU Benchmark
    runs-on: [self-hosted, macOS]
-    timeout-minutes: 10
+    timeout-minutes: 3
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -437,32 +431,70 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: Kill stale pids
      run: |
-        PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
-        PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
-    # since sudo is required for usbgpu on macos, do not write bytecode, as some of the files are owned by root
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+        ./extra/hcq/hcq_smi.py nv kill_pids --sudoless
    - name: UsbGPU boot time
-      run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEBUG=2 AM_RESET=1 DEV=USB+AMD time python3.11 test/test_tiny.py TestTiny.test_plus
+      run: GMMU=0 DEBUG=2 AM_RESET=1 DEV=USB+AMD time python3.11 test/test_tiny.py TestTiny.test_plus
    - name: UsbGPU tiny tests
-      run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
+      run: GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
    - name: UsbGPU copy speeds
-      run: sudo -E PYTHONDONTWRITEBYTECODE=1 SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
-    #- name: UsbGPU openpilot test
-    #  run: sudo -E PYTHONPATH=. GMMU=0 DEV=USB+AMD GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
    - name: UsbGPU (USB4/TB) install script
-      run: PYTHONPATH=. sh extra/setup_tinygpu_osx.sh
+      run: sh extra/setup_tinygpu_osx.sh
    - name: UsbGPU (USB4/TB) boot time
-      run: PYTHONPATH=. DEBUG=3 DEV=PCI+NV:NAK time python3.11 test/test_tiny.py TestTiny.test_plus
+      run: DEBUG=3 DEV=PCI+NV:NAK time python3.11 test/test_tiny.py TestTiny.test_plus
    - name: UsbGPU (USB4/TB) tiny tests
-      run: PYTHONPATH=. DEV=PCI+NV:NAK python3.11 test/test_tiny.py
+      run: DEV=PCI+NV:NAK python3.11 test/test_tiny.py

-  testcommalatest:
-    name: comma Benchmark (0.11.2)
+  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 }})
    runs-on: [self-hosted, Linux, comma]
-    timeout-minutes: 12
+    timeout-minutes: 5
    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
@@ -473,45 +505,10 @@ 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: 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: 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: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -524,15 +521,6 @@ jobs:
        shell: bash -e -o pipefail {0}
    if: github.repository_owner == 'tinygrad'
    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: setup staging db
-      if: github.ref == 'refs/heads/update_benchmark_staging'
-      run: |
-        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
-        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: reset process replay
-      run: test/external/process_replay/reset.py
    - name: Checkout Code
      uses: actions/checkout@v6
    - name: setup staging db
@@ -556,7 +544,7 @@ jobs:
  testcommausbgpubenchmark:
    name: UsbGPU Benchmark (comma)
    runs-on: [self-hosted, Linux, comma4]
-    timeout-minutes: 20
+    timeout-minutes: 10
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -585,7 +573,7 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['AMD', 'NV']
-    timeout-minutes: 20
+    timeout-minutes: 5
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -597,9 +585,8 @@ jobs:
      uses: actions/checkout@v6
    - name: Setup
      run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} rmmod
-        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} kill_pids
+        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} rmmod --expect
+        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} kill_pids --sudoless
        mkdir -p extra/datasets
        ln -s /raid/datasets/imagenet extra/datasets/imagenet
    - name: setup staging db
@@ -634,6 +621,9 @@ jobs:
      run: |
        GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
        GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
+    - name: HEVC Decode Benchmark
+      if: ${{ matrix.dev == 'NV' }}
+      run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
      if: ${{ matrix.dev == 'NV' }}
      run: BENCHMARK_LOG=resnet_10steps MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
@@ -656,7 +646,7 @@ jobs:
  llvmspeed:
    name: LLVM Speed
    runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 20
+    timeout-minutes: 5
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -233,7 +233,7 @@ jobs:
    - name: Run process replay tests
      uses: ./.github/actions/process-replay
    - name: Repo line count <= 26000 lines
-      run: MAX_LINE_COUNT=26000 python sz.py
+      run: MAX_LINE_COUNT=26500 python sz.py

  spec:
    strategy:
@@ -69,3 +69,4 @@ mutants
 dagre/
 graphlib/
 uv.lock
+pi_session_window0.jsonl
@@ -28,7 +28,7 @@ repos:
        pass_filenames: false
      - id: tests
        name: comprehensive test suite
-        entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/backend/test_ops.py test/backend/test_schedule.py test/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
+        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
        language: system
        always_run: true
        pass_filenames: false
@@ -4,3 +4,4 @@
 - Run `python -m mypy tinygrad/` to typecheck
 - Run `python -m ruff check .` to lint
 - Read `./tinygrad/viz/README.md` for profiling and debugging rewrite rules
+- Do not do amend commits. Always do a new commit if a force push to origin would be required.
@@ -1674,8 +1674,7 @@ def train_gptoss():
  config = {}
  BASEDIR            = config["BASEDIR"]                = Path(getenv("BASEDIR", "/raid/datasets/c4-8b/"))
  BS                 = config["BS"]                     = getenv("BS", 16)
-  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
-  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
+  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS
  SEED               = config["SEED"]                   = getenv("SEED", 5760)
  DATA_SEED          = config["DATA_SEED"]              = getenv("DATA_SEED", SEED)
  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
@@ -1737,8 +1736,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=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),
+    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),
  )

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

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

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

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

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

    for g in grads: g.assign(0)

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

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

  @TinyJit
  @Context(TRAINING=0)
@@ -1843,30 +1845,20 @@ def train_gptoss():
      profile_marker(f"train @ {i}")
      st = 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
+      ist = time.perf_counter()

-      gt = time.perf_counter()
-      ret = optim_step()
-      lr, grad_norm = ret[0].item(), ret[1].item()
+      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()
      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
@@ -1876,7 +1868,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, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
+          f"{i:5} {step_time:.3f} s step, {dev_time:.3f} s dev, {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())))

@@ -1886,8 +1878,6 @@ 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,
@@ -15,7 +15,7 @@ def stochastic_round_bf16(x:Tensor) -> Tensor:
  bits = x.bitcast(dtypes.uint32)
  if isinstance(x.device, tuple):
    shape = x.uop.shard_shape if x.uop.axis is not None else x.shape
-    noise = Tensor(UOp(Ops.MSTACK, dtypes.default_float, tuple(Tensor.rand(*shape, device=d).uop for d in x.device)))
+    noise = Tensor(UOp(Ops.MSTACK, src=tuple(Tensor.rand(*shape, device=d).uop for d in x.device)))
  else:
    noise = x.rand_like()
  noise = (noise * 0xFFFF).cast(dtypes.uint32)
@@ -35,7 +35,7 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=2048 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
 export SEQLEN=${SEQLEN:-8192}
@@ -44,9 +44,9 @@ export SEED=${SEED:-5760}
 export DATA_SEED=${DATA_SEED:-5760}

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

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

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

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

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

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

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

 export LOGMLPERF=1

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

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

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

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

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

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

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

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

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

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

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

    write_options = WriteOptions(write_type=WriteType.synchronous, retry_interval=5000, max_retries=5, max_retry_delay=30000, exponential_base=2)
@@ -0,0 +1,238 @@
+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
@@ -0,0 +1,118 @@
+#!/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()
@@ -0,0 +1,46 @@
+#!/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)
@@ -53,7 +53,7 @@ def _ggather_bwd(gradient:UOp, kernel:UOp) -> tuple:
    g, m, j, jo, ji = _kv_ranges(Gk, M, Dk, _blk_for(Dk))
    row = idx.index(g, m).cast(dtypes.weakint)
    val = gout.index(g, m, j).load().cast(dtypes.float32)
-    atomic = UOp(Ops.CUSTOM, dtypes.void, (gtab.index(g, row, j), val), arg=atomic_str)
+    atomic = UOp(Ops.CUSTOM, src=(gtab.index(g, row, j), val), arg=(atomic_str, dtypes.void))
    return atomic.end(g, m, jo, ji).sink(arg=KernelInfo(name=f"ggather_bwd_{M}_{Dk}", opts_to_apply=()))
  grad_table = Tensor.custom_kernel(gt, go, Tensor(idx_u, device=dev), fxn=_bwd_kernel)[0]
  return (None, grad_table.cast(table_u.dtype).uop, None)
@@ -16,9 +16,12 @@ def _do_reset_device(pci_bus): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/dev
 def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")

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

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

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

      try:
-        try: pid = subprocess.check_output(['sudo', 'lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
+        try: pid = subprocess.check_output((['sudo'] if use_sudo else []) +
+                                           ['lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
        except subprocess.CalledProcessError: break

        print(f"Killing process {pid} (which uses {dev})")
-        subprocess.run(['sudo', 'kill', '-9', pid], check=True)
+        subprocess.run((['sudo'] if use_sudo else []) + ['kill', '-9', pid], check=True)
      except subprocess.CalledProcessError as e:
        print(f"Failed to kill process for device {dev}: {e}", file=sys.stderr)

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

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

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

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

 if __name__ == "__main__":
@@ -19,7 +19,7 @@ from tinygrad.runtime.support.hcq import FileIOInterface, HCQBuffer, MMIOInterfa
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
 from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_pmc
 from tinygrad.runtime.support.system import PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
-from tinygrad.runtime.support.usb import USB3
+from tinygrad.runtime.support.usb import USB3, usb_ib, usb_push, usb_arm_bytes, pm_usb_stage, pm_usb_hostio, pm_usb_bufferize
 from tinygrad.runtime.support.memory import AddrSpace, BumpAllocator
 from tinygrad.runtime.ops_amd import SQTT, SQTT_ITRACE_SE_MASK, SQTT_LIMIT_SE, SQTT_SIMD_SEL, SQTT_TOKEN_EXCLUDE, PMC
 from tinygrad.runtime.ops_amd import EVENT_INDEX_PARTIAL_FLUSH, WAIT_REG_MEM_FUNCTION_EQ, WAIT_REG_MEM_FUNCTION_NEQ, WAIT_REG_MEM_FUNCTION_GEQ
@@ -87,7 +87,7 @@ def release_mem(ctx, address=0x0, value=0, data_sel=0, int_sel=2, ctxid=0, cache

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

 pm_pm4_opsel = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), pm4_program),
@@ -146,11 +146,14 @@ pm_pm4_opsel = PatternMatcher([
  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
 ])

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

  # the host fence at the start of the batch guarantees the ib is free to reuse
  size_dw = sum(len(ins.src) for ins in lin.src)
@@ -204,7 +207,7 @@ 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"), lambda: UOp(Ops.NOOP, dtypes.void, ())),
+  (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),
@@ -216,8 +219,7 @@ def sdma_submit(cmdbuf, devs):

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

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

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

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

@dataclass(frozen=True)
@@ -282,13 +301,14 @@ def amd_build_program(prg:UOp) -> UOp:
      wave32=bool(desc.kernel_code_properties & 0x400), private_segment_size=desc.private_segment_fixed_size, kernargs_segment_size=desc.kernarg_size,
      kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0), enable_dispatch_ptr=edp,
      enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER)
+    image = bytes(image).ljust(round_up(len(image), 4), b"\x00") # the program is uploaded as whole dwords
    buf = UOp.placeholder((len(image),), dtypes.uint8, next(UOp.unique_num), device=prg.device).rtag("program")
-    cached = _amd_program_cache[key] = prg.replace(src=(buf.after(make_binary_patch(buf, bytes(image))),), arg=(data, prg.arg))
+    cached = _amd_program_cache[key] = prg.replace(src=(buf.after(make_binary_patch(buf, image)),), arg=(data, prg.arg))
  return cached

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@functools.cache
@@ -12,7 +12,7 @@ def _custom_quantize_mxfp4(row_fp4:UOp, row_scale:UOp, col_fp4:UOp, col_scale:UO
  mem = M*N*2 + M*N + M*N//16  # read bf16, write row+col fp4 + e8m0
  outputs = (row_fp4, row_scale, col_fp4, col_scale)
  sink = UOp.sink(*(o.base for o in outputs), x.base,
-                  *(UOp(Ops.CUSTOM, dtypes.void, (o.base.index(0),), arg="") for o in outputs),
+                  *(UOp(Ops.CUSTOM, src=(o.base.index(0),), arg=("", dtypes.void)) for o in outputs),
                  UOp.special(256, "lidx0"), UOp.special(M//128, "gidx0"), UOp.special(N//64, "gidx1"),
                  arg=KernelInfo(name, estimates=Estimates(ops=12*M*N, mem=mem)))
  src = (pathlib.Path(__file__).parent/"quantize_mxfp4.cpp").read_text()
@@ -5,9 +5,9 @@ from tinygrad.helpers import getenv, DEBUG

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

 # matches meta, non hugging face weights
 # (a+i*b) * (c+i*d) = (ac-bd) + i*(ad+bc)
@@ -1,4 +1,4 @@
-import os, subprocess, sys, shlex
+import os, subprocess, sys, shlex, pickle
 from pathlib import Path
 from tinygrad.helpers import temp, getenv

@@ -23,5 +23,8 @@ if __name__ == "__main__":
      # AM_RESET=1 gets a clear trace, does not work on mi300 machines
      subprocess.run([sys.executable, *shlex.split(test)], cwd=EXAMPLES_DIR.parent.parent.parent,
                     env={**os.environ, "DEV":"AMD", "AM_RESET":"1" if not arch.startswith("gfx9") else "0", "VIZ":"-2", "PYTHONPATH":"."})
+      with open(PROFILE_PATH, "rb") as f: events = pickle.load(f)
+      with open(PROFILE_PATH, "wb") as f:
+        pickle.dump([e for e in events if type(e).__name__ in {"ProfilePMCEvent", "ProfileSQTTEvent", "ProfileProgramEvent"}], f)
      PROFILE_PATH.rename(dest:=EXAMPLES_DIR/arch/f"profile_{name}_run_{i}.pkl")
      print(f"saved SQTT trace to {dest}")
@@ -1,6 +1,6 @@
 [project]
 name = "tinygrad"
-version = "0.13.0"
+version = "0.14.0"
 description = "You like pytorch? You like micrograd? You love tinygrad! <3"
 authors = [{ name = "George Hotz" }]

@@ -84,7 +84,7 @@ testing = [
  "pillow",
  "onnx==1.19.0",
  "onnx2torch",
-  "onnxruntime",
+  "onnxruntime==1.24.1",
  "opencv-python",
  "transformers",
  "sentencepiece",
@@ -51,7 +51,7 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op
 \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
 \midrule
 \op{Param} & $(\mathbf{s})$ & slot, dtype, device?, addrspace? &
-  Placeholder with shape $\mathbf{s}$. Substituted in \op{Function}. \\[4pt]
+  Placeholder with shape $\mathbf{s}$. Substituted in \op{Call}. \\[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 &
@@ -102,9 +102,8 @@ 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{Function} & (body, $a_0$, $a_1$, \ldots) & --- & Substitute each \op{Param} $k$ in \op{Tuple} body with $a_k$. Gradient-able. \\
-\op{Call} & (body, $a_0$, $a_1$, \ldots) & --- & Opaque invocation of a compiled kernel or custom function. \\
-\op{Tuple} & $(v_0, v_1, \ldots)$ & --- & Pack values; required as \op{Function} body to return a value. \\
+\op{Call} & (body, $a_0$, $a_1$, \ldots) & --- & Substitute each \op{Param} $k$ in body with $a_k$. \\
+\op{Tuple} & $(v_0, v_1, \ldots)$ & --- & Pack values; required as \op{Call} body to return a value. \\
 \op{GetTuple} & $(T,)$ & idx & Extract element at idx from a \op{Tuple}. \\
 \bottomrule
 \end{tabular}
@@ -272,7 +271,7 @@ ALU unary & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0].\mathrm{shape}$
 Other binary & $\mathrm{src}[0].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\
 \op{CmpLt}, \op{CmpNe} & bool & broadcast & $\mathrm{src}[0].\mathrm{device}$ & from intervals \\
 \op{Where} & $\mathrm{src}[1].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & $[\min(b,c),\, \max(B,C)]$ \\[3pt]
-\op{Function}, \op{Call} & $\mathrm{src}[0].\mathrm{dtype}$ & substitute \op{Param} shapes & $\mathrm{src}[1].\mathrm{device}$ & dtype range \\
+\op{Call} & $\mathrm{src}[0].\mathrm{dtype}$ & substitute \op{Param} shapes & $\mathrm{src}[1].\mathrm{device}$ & dtype range \\
 \op{Range} & index & $()$ & \textsc{null} & $[0,\, n{-}1]$ \\
 \op{Index} & $\mathrm{src}[0].\mathrm{dtype}$ & remaining dims & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
 \op{Store} & void & $()$ & $\mathrm{src}[0].\mathrm{device}$ & --- \\
@@ -421,7 +420,7 @@ def allreduce(T):
 %% ============================================================
 \subsection*{{\color{callblue}The \texttt{@function} Decorator} \normalfont\small--- graph capture via tracing}

-The \texttt{@function} decorator transforms a Python function on Tensors into a single \op{Function} node.
+The \texttt{@function} decorator transforms a Python function on Tensors into a single \op{Call} node.

 \begin{lstlisting}
@function
@@ -436,11 +435,11 @@ When \texttt{f(x, y)} is called, the decorator:
  \item \textbf{Runs the function} lazily (no device execution), building a UOp graph from the result.
  \item \textbf{Parameterizes}: replaces each input UOp with a \op{Param}$(k)$ placeholder.
  \item \textbf{Wraps the body} in a \op{Tuple} (even for single returns) and creates\\
-        \op{Function}(\op{Tuple}(body), $x$, $y$).
+        \op{Call}(\op{Tuple}(body), $x$, $y$).
  \item \textbf{Returns} the result via \op{GetTuple}$(0)$, or one \op{GetTuple} per element for tuple returns.
 \end{enumerate}

-The result is a reusable graph fragment: the body contains only \op{Param} references, not concrete buffers. At schedule time, the \op{Function} is resolved by substituting each \op{Param}$(k)$ back with its corresponding argument $a_k$, or lowered into an opaque \op{Call} if it is to be compiled as a reusable kernel.
+The result is a reusable graph fragment: the body contains only \op{Param} references, not concrete buffers. At schedule time, the \op{Call} is resolved by substituting each \op{Param}$(k)$ back with its corresponding argument $a_k$, or compiled as a reusable kernel.

 %% ============================================================
 \subsection*{Lowering Pipeline \normalfont\small--- from Tensor graph to machine code}
@@ -1002,6 +1002,39 @@ class TestBarrier(unittest.TestCase):
    for tid in range(64):
      self.assertEqual(st.vgpr[tid][0], tid + 100 + 1000, f"tid={tid}")

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

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


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

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

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

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

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

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

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

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

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

 if __name__ == '__main__':
  unittest.main()
@@ -88,7 +88,6 @@ def run_rocprof_decoder(blobs: list[bytes], lib: bytes, base: int, target: str):
  if t.is_alive(): raise RuntimeError("rocprof decoder timeout")
  return occupancy_records, wave_insts

-@unittest.skip("TODO: fix to not require unpickling UOps.")
 class SQTTExamplesTestBase(unittest.TestCase):
  target: str
  examples: dict
@@ -188,7 +188,7 @@ class TestMXFP4(unittest.TestCase):
    M, N, K = getenv("M", 16384), getenv("N", 4096), getenv("K", 14336)
    a = Tensor.empty(M, K, dtype=dtypes.bfloat16)
    b = Tensor.empty(N, K, dtype=dtypes.bfloat16)
-    asm_gemm(a, b.T, mxfp4=True).realize()
+    for _ in range(getenv("CNT", 1)): asm_gemm(a, b.T, mxfp4=True).realize()

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

@@ -1,10 +1,10 @@
 #!/usr/bin/env python
 import unittest
 import numpy as np
-from tinygrad import dtypes, Tensor, TinyJit, GlobalCounters, Variable
+from tinygrad import Device, 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
+from test.helpers import assert_kernel_count, needs_second_gpu

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

@@ -1079,5 +1079,80 @@ 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()
@@ -7,7 +7,7 @@ from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops
 from tinygrad.renderer.isa import IselContext

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

@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only x86")
 class TestIselX86(unittest.TestCase):
@@ -58,6 +58,11 @@ class TestMultiTensor(unittest.TestCase):
    assert X.uop.ended_ranges == X.uop.src[1:]
    (X + X).realize()

+  def test_shard_invalids_contiguous(self):
+    # every store is Invalid, so none of them should become a (empty) kernel
+    t = Tensor.invalids(8).shard(devices_2, axis=0).contiguous()
+    self.assertEqual(len([c for c in t.schedule_linear().src if c.src[0].op is Ops.SINK]), 1)
+
  @unittest.expectedFailure # TODO: fix
  def test_shard_empty(self):
    GlobalCounters.reset()
@@ -6,6 +6,7 @@ from tinygrad.helpers import getenv, DEBUG, DEV, IMAGE, Context
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.renderer.nir import NIRRenderer
+from tinygrad.renderer.isa.x86 import X86Renderer

 TINY_BACKEND = getenv("TINY_BACKEND")
 if TINY_BACKEND:
@@ -808,6 +809,8 @@ class TestOps(unittest.TestCase):
    helper_test_op([], lambda: tor^0x1337, lambda: ten^0x1337, forward_only=True)
    helper_test_op([], lambda: 0x1337^tor, lambda: 0x1337^ten, forward_only=True)

+  # TODO: x86 PARAM dtype fails SPEC=2
+  @Context(SPEC=1 if isinstance(Device[Device.DEFAULT].renderer, X86Renderer) else 2)
  def test_and(self):
    data = [[1,-8,1],[32,1,6]]
    tor = torch.tensor(data, dtype=torch.int)
@@ -865,9 +868,9 @@ class TestOps(unittest.TestCase):
                   lambda: (ten << Tensor([0,2,4], dtype=dtypes.uint32)).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor.__lshift__(2), lambda: ten.__lshift__(2).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor.bitwise_left_shift(2), lambda: ten.lshift(2).cast(dtypes.int32), forward_only=True)
-    self.helper_test_exception([], lambda: torch.tensor([1.0]) << 2, lambda: Tensor([1.0]) << 2, expected=RuntimeError)
-    self.helper_test_exception([], lambda: tor << torch.tensor([1.0]), lambda: ten << Tensor([1.0]), expected=RuntimeError)
-    self.helper_test_exception([], lambda: tor << 1.0, lambda: ten << 1.0, expected=RuntimeError)
+    self.helper_test_exception([], lambda: torch.tensor([1.0]) << 2, lambda: (Tensor([1.0]) << 2).realize(), expected=RuntimeError)
+    self.helper_test_exception([], lambda: tor << torch.tensor([1.0]), lambda: (ten << Tensor([1.0])).realize(), expected=RuntimeError)
+    self.helper_test_exception([], lambda: tor << 1.0, lambda: (ten << 1.0).realize(), expected=RuntimeError)

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

  def test_lshift_signed(self):
    data = [[-1, -3, 1, 7], [0, -2147483648, 2147483647, -1]]
@@ -1807,9 +1810,10 @@ class TestOps(unittest.TestCase):
      helper_test_op([()], lambda x: torch.nn.functional.hardtanh(x, -val, val), lambda x: x.hardtanh(-val, val), grad_atol=1e-6)
  def test_asinh(self):
    helper_test_op([(45,65)], lambda x: x.asinh(), grad_atol=1e-6)
-    # TODO: this one has larger tol?
-    helper_test_op([(45,65)], lambda x: x.asinh(), atol=1e-2, rtol=2e-2, grad_rtol=2e-2, low=-300, high=-297)
+    helper_test_op([(45,65)], lambda x: x.asinh(), grad_atol=1e-6, low=-300, high=-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)
@@ -87,7 +87,8 @@ class TestOptim(unittest.TestCase):
  def test_muon(self): self._test_muon(1, {'lr': 0.001}, 1e-3, 0)
  # TODO: disabled due to big atol
  # def test_muon_high_lr(self): self._test_muon(1, {'lr': 10}, 1e-6, 3e-4)
-  def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 0.01}, 1e-3, 3e-4)
+  # NOTE: big weight_decay so a missing wd would be way over atol
+  def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 10}, 1e-3, 3e-4)
  # TODO: disabled due to big atol
  # def test_muon_high_lr_wd(self): self._test_muon(1, {'lr': 10, 'weight_decay': 0.01}, 1e-6, 5e-4)

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

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

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

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

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

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

@unittest.skip("pm_rangeify no longer exists. test this in a different way")
@@ -77,6 +77,14 @@ class TestCStyleFailures(unittest.TestCase):

@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "tests for wgsl renderer")
 class TestWGSLFailures(unittest.TestCase):
+  def test_folded_packed_store(self):
+    b = UOp.param(0, dtypes.char, (4,))
+    idx = b.index(UOp.const(0).cast(dtypes.int))
+    store = UOp.store(idx, idx.cast(dtypes.uint32).load() & UOp.const(0xffffff00).cast(dtypes.uint32))
+    src = Device[Device.DEFAULT].renderer.render(UOp.sink(store, arg=KernelInfo()).toposort())
+    self.assertIn("atomicAnd(&data0_4[0],4294967040u);", src)
+    self.assertNotIn("atomicAdd", src)
+
  def test_multiply_infinity(self):
    # multiplying a positive constant by infinity should return infinity
    # WGSL pipelines do not handle this reliably, some of which return zero, unless infinity always comes from a read on a dynamic buffer
@@ -147,7 +147,7 @@ class TestSchedule(unittest.TestCase):
    devs = ("CPU:0", "CPU:1")
    x = Tensor.ones(2, device="CPU").shard(devs, axis=0).realize()
    out = (x.sum()*2).reshape(1).to("CPU")
-    run_linear(*check_schedule(out, 5))
+    run_linear(*check_schedule(out, 3))
    np.testing.assert_equal(out.numpy(), [4.])

 class TestLimitBufs(unittest.TestCase):
@@ -365,16 +365,6 @@ 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))
@@ -384,14 +374,6 @@ 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)
@@ -301,6 +301,14 @@ class TestSetitem(unittest.TestCase):
    self.assertListEqual(z[2:5].tolist(), [2, 2, 2])
    self.assertListEqual(z[6:7].tolist(), [3])

+class TestAssignBitcast(unittest.TestCase):
+  def test_assign_through_bitcast(self):
+    # the dest is unrealized, so callify cannot fold the BITCAST into a buffer view and the STORE keeps a
+    # BITCAST dest; the bitcast has to move to the value side or the store never reaches the buffer
+    a = Tensor.full((4,), 1.0, dtype=dtypes.float32).contiguous()
+    a.bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.uint32)).realize()
+    np.testing.assert_allclose(a.numpy(), [4.0, 3.0, 2.0, 1.0])
+
 class TestWithGrad(unittest.TestCase):
  def test_basic_setitem_works(self):
    z = Tensor.rand(8, 8)
@@ -57,6 +57,35 @@ def _test_uops_result(output_dtype, uops, res):
  run_uops([out], [buf])
  return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0]

+@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, CStyleLanguage) and
+                     dtypes.uint64 in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires C-style pointer bitcast and 64-bit ints")
+class TestBitcastBufferView(unittest.TestCase):
+  @Context(SPEC=2)
+  def test_render(self):
+    buf = UOp.param(0, dtypes.uint32, (4,))
+    uops = to_uops_list([buf.shrink(((1, 3),)).bitcast(dtypes.uint64).index(0).store(1)], ren=Device[Device.DEFAULT].renderer)
+    idx = next(u for u in uops if u.op is Ops.INDEX and u.src[0].op is Ops.BITCAST)
+    self.assertEqual(idx.src[0].src[0].op, Ops.SHRINK)
+    Device[Device.DEFAULT].renderer.render(uops)
+
+  @Context(SPEC=2)
+  def test_load(self):
+    val = 0x1122334455667788
+    src, out = UOp.param(0, dtypes.uint32, (4,)), UOp.param(1, dtypes.uint64, (1,))
+    ibuf = Buffer(Device.DEFAULT, 4, dtypes.uint32, initial_value=np.array([0, 0x55667788, 0x11223344, 0], dtype=np.uint32).tobytes())
+    obuf = Buffer(Device.DEFAULT, 1, dtypes.uint64).allocate()
+    run_uops([out.index(0).store(src.shrink(((1, 3),)).bitcast(dtypes.uint64).index(0))], [ibuf, obuf])
+    self.assertEqual(np.frombuffer(obuf.as_memoryview(), dtype=np.uint64)[0], val)
+
+  @Context(SPEC=2)
+  def test_store(self):
+    val = 0x1122334455667788
+    dst = UOp.param(0, dtypes.uint32, (6,))
+    buf = Buffer(Device.DEFAULT, 6, dtypes.uint32, initial_value=bytes(24))
+    view = dst.shrink(((1, 5),)).bitcast(dtypes.uint64)  # two stores through one view: it must inline, not get a declared vector-pointer
+    run_uops([view.index(0).store(val ^ 0xff), view.index(1).store(val)], [buf])
+    self.assertEqual(np.frombuffer(buf.as_memoryview(), dtype=np.uint64, count=2, offset=4).tolist(), [val ^ 0xff, val])
+
 class TestUOps(unittest.TestCase):
  def _equal(self, v1, v2):
    assert isinstance(v2, (float, int, bool))
@@ -23,7 +23,7 @@ class TestGPUCrash(unittest.TestCase):
    cls.is_cdna = cls.dev.target[0] < 10
    ins = importlib.import_module('tinygrad.runtime.autogen.amd.' + ('cdna' if cls.is_cdna else 'rdna3') + '.ins')
    for rdna3_name, cdna3_name in RDNA3_CDNA3_MAP.items():
-      setattr(cls, rdna3_name, getattr(ins, cdna3_name if cls.is_cdna else rdna3_name))
+      setattr(cls, rdna3_name, staticmethod(getattr(ins, cdna3_name if cls.is_cdna else rdna3_name)))

  def setUp(self):
    # Verify device works before each test
@@ -1,20 +1,15 @@
-import pathlib, re, ctypes, mmap, collections, functools, copy, os
+import 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 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()
+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)}

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

  def kfd_ioctl(self, req, argp):
    nr = req & 0xFF
-    struct = kfd_headers[nr].from_address(argp)
+    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)

-    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:
+    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):
      struct.major_version = 1
      struct.minor_version = 14
-    elif nr == kfd_ioctls.AMDKFD_IOC_ALLOC_MEMORY_OF_GPU:
+    elif nr == _ioctl_nr(kfd.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 == kfd_ioctls.AMDKFD_IOC_FREE_MEMORY_OF_GPU:
+    elif nr == _ioctl_nr(kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU):
      self.object_by_handle.pop(struct.handle)
-    elif nr == kfd_ioctls.AMDKFD_IOC_MAP_MEMORY_TO_GPU:
+    elif nr == _ioctl_nr(kfd.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 == kfd_ioctls.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU:
+    elif nr == _ioctl_nr(kfd.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 == kfd_ioctls.AMDKFD_IOC_CREATE_EVENT:
+    elif nr == _ioctl_nr(kfd.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 == kfd_ioctls.AMDKFD_IOC_CREATE_QUEUE:
+    elif nr == _ioctl_nr(kfd.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)
@@ -162,7 +159,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 == kfd_ioctls.AMDKFD_IOC_WAIT_EVENTS:
+    elif nr == _ioctl_nr(kfd.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:
@@ -170,11 +167,7 @@ class AMDDriver(VirtDriver):
          ev.memory_exception_data.va = int(os.environ["MOCKGPU_EMU_FAULTADDR"], 16)
          ev.memory_exception_data.failure.NotPresent = 1
    else:
-      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)
+      raise RuntimeError(f"unsupported kfd ioctl, {nr} {name}")
    return 0

  def _emulate_execute(self):
@@ -1,5 +1,20 @@
 # Tokenizer-based expression parser for AMD pcode
+import ast, itertools, operator, re
 from typing import Any, Callable
+
+_BINOPS = {ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul, ast.FloorDiv: operator.floordiv,
+           ast.Mod: operator.mod, ast.LShift: operator.lshift, ast.RShift: operator.rshift,
+           ast.BitAnd: operator.and_, ast.BitOr: operator.or_, ast.BitXor: operator.xor}
+def _const_int(expr: str) -> int:
+  """Evaluate a compile-time integer expression (integer literals and basic arithmetic only)."""
+  def ev(node: ast.AST) -> int:
+    if isinstance(node, ast.Expression): return ev(node.body)
+    if isinstance(node, ast.Constant) and isinstance(node.value, int): return node.value
+    if isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.USub, ast.UAdd)):
+      return (-1 if isinstance(node.op, ast.USub) else 1) * ev(node.operand)
+    if isinstance(node, ast.BinOp) and type(node.op) in _BINOPS: return _BINOPS[type(node.op)](ev(node.left), ev(node.right))
+    raise ValueError(f"not a constant integer expression: {expr!r}")
+  return ev(ast.parse(expr.strip(), mode='eval'))
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.codegen.decomp.dtype import f2f
@@ -8,6 +23,7 @@ from tinygrad.codegen.decomp.dtype import f2f
 VarVal = UOp | tuple[str, list[str], str]

 def _const(dt, v): return UOp.const(v, dt)
+def _single_value(v: UOp): return v.vmin if v.vmin == v.vmax else None
 def _u32(v): return _const(dtypes.uint32, v)
 def _u64(v): return _const(dtypes.uint64, v)
 def _to_u32(v): return v if v.dtype == dtypes.uint32 else v.bitcast(dtypes.uint32) if v.dtype.itemsize == 4 else v.cast(dtypes.uint32)
@@ -55,8 +71,8 @@ def _expr_bits(v: UOp) -> int:
  if v.op in (Ops.AND, Ops.XOR):
    widths: list[int] = []
    for src in v.src:
-      if src.op == Ops.CONST and isinstance(src.val, int) and src.val > 0 and (src.val & (src.val + 1)) == 0:
-        widths.append(src.val.bit_length())
+      if isinstance(sv:=_single_value(src), int) and sv > 0 and (sv & (sv + 1)) == 0:
+        widths.append(sv.bit_length())
    if widths: return max(widths)
  return v.dtype.bitsize

@@ -144,9 +160,9 @@ def _minmax_reduce(is_max: bool, dt, *args: UOp) -> UOp:
 def _find_two_pi_mul(x):
  if x.op != Ops.MUL or len(x.src) != 2: return None
  for i, s in enumerate(x.src):
-    if s.op == Ops.CONST and abs(s.val - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586)
+    if (sv:=_single_value(s)) is not None and abs(sv - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586)
    if s.op == Ops.MUL and len(s.src) == 2:
-      vals = [ss.val for ss in s.src if ss.op == Ops.CONST] + [ss.src[0].val for ss in s.src if ss.op == Ops.CAST and ss.src[0].op == Ops.CONST]
+      vals = [sv for ss in s.src if (sv:=_single_value(ss)) is not None]
      if len(vals) == 2 and abs(vals[0] * vals[1] - 6.283185307179586) < 1e-5: return (x.src[1-i], vals[0] * vals[1])
  return None

@@ -163,7 +179,7 @@ def _trig_reduce(x, phase=0.0):

 def _signext(val: UOp) -> UOp:
  for bits, mask, ext in [(4, 0xF, 0xFFFFFFF0), (8, 0xFF, 0xFFFFFF00), (16, 0xFFFF, 0xFFFF0000)]:
-    if (val.op == Ops.AND and len(val.src) == 2 and val.src[1].op == Ops.CONST and val.src[1].val == mask) or val.dtype.itemsize == bits // 8:
+    if (val.op == Ops.AND and len(val.src) == 2 and _single_value(val.src[1]) == mask) or val.dtype.itemsize == bits // 8:
      v32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
      sb = (v32 >> _u32(bits - 1)) & _u32(1)
      return sb.ne(_u32(0)).where(v32 | _u32(ext), v32).cast(dtypes.int)
@@ -185,7 +201,20 @@ def _abs(val: UOp) -> UOp:
 def _f_to_u(f, dt):
  clamped = (f < _const(f.dtype, 0.0)).where(_const(f.dtype, 0.0), f)
  truncated = UOp(Ops.TRUNC, src=(clamped,))
-  return (truncated >= _const(f.dtype, 2**(dt.itemsize*8))).where(_const(dt, dt.max), truncated.cast(dt))
+  res = (truncated >= _const(f.dtype, 2**(dt.itemsize*8))).where(_const(dt, dt.max), truncated.cast(dt))
+  return _isnan(f).where(_const(dt, 0), res)  # float->uint conversion of NaN is 0 on hardware
+
+def _f_to_i32(a: UOp) -> UOp:
+  """v_cvt_i32_f32/f64: truncate toward zero, saturate to [INT_MIN, INT_MAX], NaN -> 0.
+  (x86 cvttss2si returns 0x80000000 for all of these, which matches hardware only for negative overflow.)"""
+  res = (a >= _const(a.dtype, 2147483648.0)).where(_const(dtypes.int, 0x7FFFFFFF), UOp(Ops.TRUNC, src=(a,)).cast(dtypes.int))
+  return _isnan(a).where(_const(dtypes.int, 0), res)
+
+def _ftz_f32(v: UOp) -> UOp:
+  """Flush f32 denormals to signed zero (RDNA default float mode flushes denormal f32 inputs on select-style ops)."""
+  bits = v.bitcast(dtypes.uint32) if v.dtype == dtypes.float32 else v
+  return ((bits & _u32(0x7FFFFFFF)) < _u32(0x00800000)).where((bits & _u32(0x80000000)).bitcast(dtypes.float32),
+                                                              v if v.dtype == dtypes.float32 else v.bitcast(dtypes.float32))

 def _cvt_quiet(val: UOp) -> UOp:
  bits, _, _, qb, _ = _float_info(val)
@@ -230,18 +259,51 @@ def _ldexp(val: UOp, exp: UOp) -> UOp:
  if val.dtype == dtypes.uint32: val = val.bitcast(dtypes.float32)
  elif val.dtype == dtypes.uint64: val = val.bitcast(dtypes.float64)
  if exp.dtype in (dtypes.uint32, dtypes.uint64): exp = exp.cast(dtypes.int if exp.dtype == dtypes.uint32 else dtypes.int64)
-  return val * UOp(Ops.EXP2, src=(exp.cast(val.dtype),))
+  bits = val.bitcast(dtypes.uint32) if val.dtype == dtypes.float32 else val.bitcast(dtypes.uint64)
+  abs_max = _const(bits.dtype, 0x7F800000 if val.dtype == dtypes.float32 else 0x7FF0000000000000)
+  sign_mask = _const(bits.dtype, 0x80000000 if val.dtype == dtypes.float32 else 0x8000000000000000)
+  # hardware flushes denormal inputs to signed zero
+  magn_mask = _const(bits.dtype, 0x7FFFFFFF if val.dtype == dtypes.float32 else 0x7FFFFFFFFFFFFFFF)
+  is_denorm = ((bits & abs_max).eq(_const(bits.dtype, 0))) & ((bits & magn_mask).ne(_const(bits.dtype, 0)))
+  val = is_denorm.where((bits & sign_mask).bitcast(val.dtype), val)
+  # hardware propagates 0/+-inf/NaN unchanged (avoids 0*inf = NaN on the host)
+  res = val * UOp(Ops.EXP2, src=(exp.cast(val.dtype),))
+  is_special = (bits & abs_max).eq(_const(bits.dtype, 0)) | ((bits & abs_max) >= abs_max)
+  return is_special.where(val, res)

 def _frexp_mant(val: UOp) -> UOp:
  val = val.bitcast(dtypes.float32) if val.dtype == dtypes.uint32 else val.bitcast(dtypes.float64) if val.dtype == dtypes.uint64 else val
-  if val.dtype == dtypes.float32: return ((val.bitcast(dtypes.uint32) & _u32(0x807FFFFF)) | _u32(0x3f000000)).bitcast(dtypes.float32)
-  return ((val.bitcast(dtypes.uint64) & _const(dtypes.uint64, 0x800FFFFFFFFFFFFF)) |
-    _const(dtypes.uint64, 0x3fe0000000000000)).bitcast(dtypes.float64)
+  if val.dtype == dtypes.float32:
+    bits = val.bitcast(dtypes.uint32)
+    # denormal/zero inputs (exponent field == 0) return signed zero on hardware
+    return ((bits & _u32(0x7F800000)).ne(_u32(0))).where(((bits & _u32(0x807FFFFF)) | _u32(0x3F000000)).bitcast(dtypes.float32),
+                                                          (bits & _u32(0x80000000)).bitcast(dtypes.float32))
+  bits = val.bitcast(dtypes.uint64)
+  return ((bits & _const(dtypes.uint64, 0x7FF0000000000000)).ne(_const(dtypes.uint64, 0))).where(
+    ((bits & _const(dtypes.uint64, 0x800FFFFFFFFFFFFF)) | _const(dtypes.uint64, 0x3fe0000000000000)).bitcast(dtypes.float64),
+    (bits & _const(dtypes.uint64, 0x8000000000000000)).bitcast(dtypes.float64))
+
+def _msb(val: UOp, bits: int) -> UOp:
+  """Index of the highest set bit, or -1 if val == 0."""
+  dt = dtypes.uint64 if bits > 32 else dtypes.uint32
+  val = val.cast(dt) if val.dtype != dt else val
+  result = _const(dtypes.int, -1)
+  for i in range(bits - 1, -1, -1):
+    cond = ((val >> _const(dt, i)) & _const(dt, 1)).ne(_const(dt, 0)) & result.eq(_const(dtypes.int, -1))
+    result = cond.where(_const(dtypes.int, i), result)
+  return result

 def _frexp_exp(val: UOp) -> UOp:
  val = val.bitcast(dtypes.float32) if val.dtype == dtypes.uint32 else val.bitcast(dtypes.float64) if val.dtype == dtypes.uint64 else val
-  if val.dtype == dtypes.float32: return ((val.bitcast(dtypes.uint32) >> _u32(23)) & _u32(0xFF)).cast(dtypes.int) - _const(dtypes.int, 126)
-  return ((val.bitcast(dtypes.uint64) >> _const(dtypes.uint64, 52)) & _const(dtypes.uint64, 0x7FF)).cast(dtypes.int) - _const(dtypes.int, 1022)
+  if val.dtype == dtypes.float32:
+    e = (val.bitcast(dtypes.uint32) >> _u32(23)) & _u32(0xFF)
+    return e.ne(_u32(0)).where(e.cast(dtypes.int) - _const(dtypes.int, 126), _const(dtypes.int, 0))  # f32 denormals -> 0 (hardware verified)
+  bits = val.bitcast(dtypes.uint64)
+  e = (bits >> _const(dtypes.uint64, 52)) & _const(dtypes.uint64, 0x7FF)
+  mant = bits & _const(dtypes.uint64, 0xFFFFFFFFFFFFF)
+  # f64 denormals: normalized exponent = highest set mantissa bit - 1073, zero -> 0 (hardware verified)
+  denorm = mant.ne(_const(dtypes.uint64, 0)).where(_msb(mant, 52) - _const(dtypes.int, 1073), _const(dtypes.int, 0))
+  return e.ne(_const(dtypes.uint64, 0)).where(e.cast(dtypes.int) - _const(dtypes.int, 1022), denorm)

 TWO_OVER_PI = int(
  "0145f306dc9c882a53f84eafa3ea69bb81b6c52b3278872083fca2c757bd778ac36e48dc74849ba5c00c925dd413a32439fc3bd"
@@ -299,9 +361,9 @@ _FUNCS: dict[str, Callable[..., UOp]] = {
  'fma': lambda a, b, c: a * b + c,
  'i32_to_f32': lambda a: a.cast(dtypes.int).cast(dtypes.float32),
  'u32_to_f32': lambda a: a.cast(dtypes.uint32).cast(dtypes.float32),
-  'f32_to_i32': lambda a: UOp(Ops.TRUNC, src=(a.bitcast(dtypes.float32),)).cast(dtypes.int),
+  'f32_to_i32': lambda a: _f_to_i32(a.bitcast(dtypes.float32)),
  'f32_to_u32': lambda a: _f_to_u(a.bitcast(dtypes.float32), dtypes.uint32),
-  'f64_to_i32': lambda a: UOp(Ops.TRUNC, src=(a.bitcast(dtypes.float64),)).cast(dtypes.int),
+  'f64_to_i32': lambda a: _f_to_i32(a.bitcast(dtypes.float64)),
  'f64_to_u32': lambda a: _f_to_u(a.bitcast(dtypes.float64), dtypes.uint32),
  'f16_to_f32': lambda a: _f16_extract(a).cast(dtypes.float32),
  'f32_to_f16': lambda a: a.cast(dtypes.half),
@@ -360,22 +422,13 @@ _FUNCS: dict[str, Callable[..., UOp]] = {
  'fp8_to_f32': _fp8_to_f32, 'bf8_to_f32': _bf8_to_f32, 'f32_to_fp8': _f32_to_fp8, 'f32_to_bf8': _f32_to_bf8,
  'f32_to_bf16': _f32_to_bf16, 'f32_to_bf16_SR': _f32_to_bf16_sr, 'f32_to_bf16_sr': _f32_to_bf16_sr,
 }
-for is_max, name in [(False, 'min'), (True, 'max')]:
-  for dt, sfx in [(dtypes.float32, 'f32'), (dtypes.int, 'i32'), (dtypes.uint32, 'u32'), (dtypes.int16, 'i16'), (dtypes.uint16, 'u16')]:
-    _FUNCS[f'v_{name}_{sfx}'] = lambda *a, im=is_max, d=dt: _minmax_reduce(im, d, *a)
-    _FUNCS[f'v_{name}3_{sfx}'] = lambda *a, im=is_max, d=dt: _minmax_reduce(im, d, *a)
-# f16 min/max/min3/max3/med3
-for is_max, name in [(False, 'min'), (True, 'max')]:
-  _FUNCS[f'v_{name}_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
-  _FUNCS[f'v_{name}3_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
-  _FUNCS[f'v_{name}_num_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
-  _FUNCS[f'v_{name}_num_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
-  _FUNCS[f'v_{name}3_num_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
-  _FUNCS[f'v_{name}3_num_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
-  _FUNCS[f'v_{name}imum_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
-  _FUNCS[f'v_{name}imum_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
-  _FUNCS[f'v_{name}imum3_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
-  _FUNCS[f'v_{name}imum3_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
+# min/max family: min/max + 3-input (x3), IEEE num variants (f16/f32 only), and long names minimum/maximum (f16/f32 only)
+for is_max, name, full in [(False, 'min', 'minimum'), (True, 'max', 'maximum')]:
+  for dt, sfx, pre in [(dtypes.float32, 'f32', None), (dtypes.int, 'i32', None), (dtypes.uint32, 'u32', None),
+                       (dtypes.int16, 'i16', None), (dtypes.uint16, 'u16', None), (dtypes.half, 'f16', _f16_extract)]:
+    def mm(*a, im=is_max, d=dt, p=pre): return _minmax_reduce(im, d, *(a if p is None else [p(x) for x in a]))
+    extra = (f'v_{name}_num_{sfx}', f'v_{name}3_num_{sfx}', f'v_{full}_{sfx}', f'v_{full}3_{sfx}') if dt in (dtypes.float32, dtypes.half) else ()
+    for fn in (f'v_{name}_{sfx}', f'v_{name}3_{sfx}', *extra): _FUNCS[fn] = mm

 # ═══════════════════════════════════════════════════════════════════════════════
 # TOKENIZER/PARSER
@@ -497,7 +550,7 @@ class Parser:
        if not dtypes.is_int(right.dtype): right = right.cast(dtypes.uint32)
        return (left >> right) if op == '>>' else (left << right)
      case '+' | '-':
-        if op == '-' and left.op == Ops.CONST and right.op == Ops.CONST: return _const(left.dtype, left.val - right.val)
+        if op == '-' and (lv:=_single_value(left)) is not None and (rv:=_single_value(right)) is not None: return _const(left.dtype, lv - rv)
        return (left + right) if op == '+' else (left - right)
      case '*' | '/':
        # Integer promotion: promote 16-bit integers to 32-bit before multiply to avoid overflow
@@ -507,7 +560,7 @@ class Parser:
          left, right = left.cast(pdt), right.cast(pdt)
        if op == '*': return left * right
        return (left // right) if dtypes.is_int(left.dtype) else (left / right)
-      case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if left.op == Ops.CONST and left.val == 2.0 else left
+      case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if _single_value(left) == 2.0 else left

  _PREC = [('||',), ('&&',), ('|',), ('^',), ('&',), ('==', '!=', '<>'), ('>=', '<=', '>', '<'), ('>>', '<<'), ('+', '-'), ('*', '/'), ('**',)]

@@ -529,8 +582,8 @@ class Parser:
      return inner.eq(_const(inner.dtype, 0))
    if self.try_eat_val('-', 'OP'):
      inner = self.unary()
-      if inner.op == Ops.CONST:
-        return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -inner.val)
+      if (v:=_single_value(inner)) is not None:
+        return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -v)
      return inner.neg()
    if self.try_eat_val('+', 'OP'): return self.unary()
    return self.postfix()
@@ -669,15 +722,13 @@ class Parser:
      self.eat('OP')
      width = self.parse()
      self.eat('RBRACKET')
-      if width.op == Ops.CONST:
-        w = int(width.val)
+      if isinstance(w:=_single_value(width), int):
        return (base >> _to_u32(first)) & _const(base.dtype, (1 << w) - 1)
      return base
    if self.try_eat('COLON'):
      second = self.parse()
      self.eat('RBRACKET')
-      if first.op == Ops.CONST and second.op == Ops.CONST:
-        a, b = int(first.val), int(second.val)
+      if isinstance(a:=_single_value(first), int) and isinstance(b:=_single_value(second), int):
        if a < b: return _bitreverse(base, b - a + 1)
        hi, lo = a, b
        if lo >= base.dtype.itemsize * 8:
@@ -698,8 +749,7 @@ class Parser:
      dt_suffix = DTYPES.get(self.eat('IDENT').val, dtypes.uint32)
    if var_name is None:
      var_name = self._find_var_name(base)
-    if first.op == Ops.CONST:
-      idx = int(first.val)
+    if isinstance(idx:=_single_value(first), int):
      # Check for array element (var@idx)
      if var_name and f'{var_name}@{idx}' in self.vars:
        v = self.vars[f'{var_name}@{idx}']
@@ -872,7 +922,7 @@ class Parser:

  def _coerce_cmp(self, l: UOp, r: UOp) -> tuple[UOp, UOp]:
    if l.dtype != r.dtype:
-      if r.dtype == dtypes.int and r.op == Ops.CONST and r.val < 0: l = l.cast(dtypes.int)
+      if r.dtype == dtypes.int and isinstance(rv:=_single_value(r), int) and rv < 0: l = l.cast(dtypes.int)
      else: r = r.cast(l.dtype)
    return l, r

@@ -890,6 +940,8 @@ class Parser:
      return result & _isnan(l).logical_not() & _isnan(r).logical_not()
    return result

+_break_var_ids = itertools.count()  # unique names for per-loop break-tracking variables
+
 def _match_bracket(toks: list[Token], start: int) -> tuple[int, list[Token]]:
  """Match brackets from start, return (end_idx, inner_tokens)."""
  j, depth = start + 1, 1
@@ -968,9 +1020,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
          p.eat('NUM')
          p.eat('QUOTE')
        if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl'))
-        expr = p.parse().simplify()
-        assert expr.op == Ops.CONST, f"loop bound must be constant, got {expr}"
-        return int(expr.val)
+        return int(p.parse())
      start_val = parse_bound()
      p.eat('COLON')
      end_val = parse_bound()
@@ -987,7 +1037,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
        i += 1
      # Execute loop with break support
      has_break = any('break' in bl.lower() for bl in body_lines)
-      found_var = f'_found_{id(body_lines)}' if has_break else None
+      found_var = f'_found_{next(_break_var_ids)}' if has_break else None
      if found_var: env[found_var] = block_assigns[found_var] = _const(dtypes.bool, False)
      for loop_i in range(start_val, end_val + 1):
        subst_lines = [_subst_loop_var(bl, loop_var, loop_i) for bl in body_lines if not (has_break and bl.strip().lower() == 'break')]
@@ -1087,7 +1137,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
          j, slice_toks = _match_bracket(toks, j)
          slice_str = _tok_str(slice_toks)
          hi_str, lo_str = slice_str.split(':')
-          hi_val, lo_val = int(eval(hi_str.strip())), int(eval(lo_str.strip()))
+          hi_val, lo_val = _const_int(hi_str), _const_int(lo_str)
          if j < len(toks) and toks[j].type == 'DOT': j += 2  # skip .type suffix
          if j < len(toks) and toks[j].type == 'EQUALS': j += 1
          ln = parse_tokens(lane_toks, env, funcs)
@@ -1145,7 +1195,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
        hi_str = ' '.join(t.val for t in toks[bracket_start:colon_pos] if t.type != 'EOF')
        lo_str = ' '.join(t.val for t in toks[colon_pos+1:j] if t.type != 'EOF')
        try:
-          hi_val, lo_val = int(eval(hi_str)), int(eval(lo_str))
+          hi_val, lo_val = _const_int(hi_str), _const_int(lo_str)
          hi, lo = max(hi_val, lo_val), min(hi_val, lo_val)
          j += 1
          if j < len(toks) and toks[j].type == 'DOT': j += 2
@@ -1159,7 +1209,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
          block_assigns[var] = env[var] = _set_bits(old, _val_to_bits(val), hi - lo + 1, lo)
          i += 1
          continue
-        except Exception: pass
+        except (ValueError, SyntaxError): pass  # non-constant slice bounds - fall through to other statement forms
      elif toks[1].type == 'LBRACKET':  # bit index: var[expr] (only for var[...], not var.type[...])
        existing = block_assigns.get(var, env.get(var))
        if existing is not None and isinstance(existing, UOp) and \
@@ -1360,3 +1410,25 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic

 def parse_expr(expr: str, env: dict[str, VarVal], funcs: dict | None = None) -> UOp:
  return parse_tokens(tokenize(expr.strip().rstrip(';')), env, funcs)
+
+def parse_pcode(pcode: str, srcs: dict[str, UOp | int] | None = None) -> tuple[dict, list]:
+  env: dict = srcs.copy() if srcs else {}
+  assigns: list[tuple[str, UOp]] = []
+  raw_lines = [l.strip().rstrip(';') for l in pcode.split('\n') if l.strip() and not l.strip().startswith('//')]
+  # TODO: pcode.py should tokenize full pcode string instead of line-by-line, then this hack can be removed
+  lines: list[str] = []
+  for l in raw_lines:
+    if lines and re.search(r'(&&|\|\||[&|+\-*/^])\s*$', lines[-1]): lines[-1] = lines[-1] + ' ' + l
+    else: lines.append(l)
+  _, final, _ = parse_block(lines, 0, env, assigns=assigns)
+  sliced = set(d.split('[')[0] for d, _ in assigns if '[' in d)
+  for var, val in final.items():
+    if var in ['D0', 'S0', 'SCC', 'VCC', 'EXEC', 'PC', 'RETURN_DATA', 'VDATA'] and isinstance(val, UOp):
+      if var in sliced and not any(re.match(rf'{var}\.\w+\s*=', l) for l in lines): continue
+      for l in lines:
+        if (m := re.match(rf'{var}\.(\w+(?:\[\w+\])?)', l)):
+          assigns.append((f'{var}.{m.group(1)}', val))
+          break
+      else: assigns.append((var, val))
+  return env, assigns
+
@@ -0,0 +1,100 @@
+# SQTT trace encoder for the emulator (the decoder lives in tinygrad/renderer/amd/sqtt.py).
+# run_asm emits packets inline as instructions execute; finished traces end up in emu.sqtt_traces.
+from __future__ import annotations
+from tinygrad.renderer.amd.dsl import Inst
+from tinygrad.renderer.amd.sqtt import (_build_decode_tables, PACKET_TYPES_RDNA3, PacketType, InstOp,
+                                        LAYOUT_HEADER, WAVESTART, WAVEEND, INST, IMMEDIATE, VALUINST)
+
+_NIB_COUNTS = {cls: nc for _, (cls, nc, *_) in _build_decode_tables(PACKET_TYPES_RDNA3)[0].items()}
+
+def _emit_nibbles(nibbles: list[int], pkt_cls: type[PacketType], **kwargs):
+  raw = pkt_cls.encoding.default
+  for k, v in kwargs.items(): raw = pkt_cls.__dict__[k].set(raw, v)
+  nibbles.extend((raw >> (i * 4)) & 0xF for i in range(_NIB_COUNTS[pkt_cls]))
+
+def make_encoder():
+  """Build an SQTT trace encoder for the emulator. Returns (emit, finish, finalize)."""
+  from tinygrad.runtime.autogen.amd.rdna3.enum import SOPPOp as SOPPOp3
+  from tinygrad.runtime.autogen.amd.rdna4.enum import SOPPOp as SOPPOp4
+  from tinygrad.runtime.autogen.amd.rdna3 import ins as ir3
+  from tinygrad.runtime.autogen.amd.rdna4 import ins as ir4
+  from tinygrad.runtime.autogen.amd.cdna import ins as irc
+  import re
+
+  def _kinds(*names: str) -> tuple[type[Inst], ...]:
+    return tuple(getattr(m, n) for m in (ir3, ir4, irc) for n in names if hasattr(m, n))
+  _SOPP, _SMEM, _DS = _kinds('SOPP'), _kinds('SMEM'), _kinds('DS')
+  _GLOBAL, _FLAT, _SCRATCH = _kinds('GLOBAL', 'VGLOBAL'), _kinds('FLAT', 'VFLAT'), _kinds('SCRATCH', 'VSCRATCH')
+  _VALU = _kinds('VOP1', 'VOP2', 'VOP3', 'VOP3P', 'VOP3PX2', 'VOPC', 'VOPD', 'VOP3SD', 'VOP3_SDST', 'VOP1_SDST')
+
+  # SOPP classification sets
+  _SOPP_SKIP = {SOPPOp3.S_ENDPGM.value, SOPPOp3.S_ENDPGM_SAVED.value, SOPPOp3.S_ENDPGM_ORDERED_PS_DONE.value, SOPPOp3.S_DELAY_ALU.value}
+  _SOPP_IMMEDIATE = {SOPPOp3.S_NOP.value, SOPPOp3.S_CLAUSE.value, SOPPOp3.S_WAITCNT.value, SOPPOp3.S_WAITCNT_DEPCTR.value,
+                     SOPPOp3.S_WAIT_IDLE.value, SOPPOp3.S_WAIT_EVENT.value, SOPPOp3.S_SLEEP.value, SOPPOp3.S_SET_INST_PREFETCH_DISTANCE.value}
+  for _op in (SOPPOp4.S_WAIT_ALU, SOPPOp4.S_WAIT_LOADCNT, SOPPOp4.S_WAIT_STORECNT, SOPPOp4.S_WAIT_SAMPLECNT,
+              SOPPOp4.S_WAIT_BVHCNT, SOPPOp4.S_WAIT_EXPCNT, SOPPOp4.S_WAIT_DSCNT, SOPPOp4.S_WAIT_KMCNT,
+              SOPPOp4.S_WAIT_LOADCNT_DSCNT, SOPPOp4.S_WAIT_STORECNT_DSCNT):
+    _SOPP_IMMEDIATE.add(_op.value)
+  _SOPP_BARRIER = {SOPPOp3.S_BARRIER.value}
+  if hasattr(SOPPOp4, 'S_BARRIER_WAIT'): _SOPP_BARRIER.add(SOPPOp4.S_BARRIER_WAIT.value)
+  if hasattr(SOPPOp4, 'S_BARRIER_LEAVE'): _SOPP_BARRIER.add(SOPPOp4.S_BARRIER_LEAVE.value)
+  _SOPP_BRANCH = {SOPPOp3.S_BRANCH.value, SOPPOp3.S_CBRANCH_SCC0.value, SOPPOp3.S_CBRANCH_SCC1.value,
+                  SOPPOp3.S_CBRANCH_VCCZ.value, SOPPOp3.S_CBRANCH_VCCNZ.value,
+                  SOPPOp3.S_CBRANCH_EXECZ.value, SOPPOp3.S_CBRANCH_EXECNZ.value}
+
+  # VALU sub-classification patterns
+  _VALUT_4_RE = re.compile(r'V_(EXP|LOG|RCP|RSQ|SQRT|SIN|COS|CEIL|FLOOR|TRUNC|RNDNE|FRACT|FREXP)_')
+  _VALUB_2_RE = re.compile(r'V_(LSHLREV|LSHRREV|ASHRREV)_(B|I)64')
+  _VALUB_4_RE = re.compile(r'V_MAD_(U|I)64')
+  _VALUB_16_RE = re.compile(r'V_\w+_F64')
+
+  def _valu_op(op_name: str) -> InstOp|None:
+    if 'CMPX' in op_name: return InstOp.VALU1_WR_EXEC
+    if _VALUB_2_RE.search(op_name): return InstOp.VALUB_2
+    if _VALUB_4_RE.search(op_name): return InstOp.VALUB_4
+    if _VALUB_16_RE.search(op_name): return InstOp.VALUB_16
+    if _VALUT_4_RE.search(op_name): return InstOp.VALUT_4
+    return None
+
+  def _mem_op(t: type[Inst], op_name: str) -> InstOp:
+    is_store = "STORE" in op_name
+    if issubclass(t, _DS): return InstOp.LDS_WR_2 if is_store else InstOp.LDS_RD
+    if issubclass(t, _GLOBAL): return InstOp.SGMEM_WR_2 if is_store else InstOp.SGMEM_RD_1
+    if issubclass(t, _FLAT) or issubclass(t, _SCRATCH): return InstOp.FLAT_WR_3 if is_store else InstOp.FLAT_RD_2
+    return InstOp.SALU
+
+  nibbles: list[int] = []
+  started: set[int] = set()
+  _emit_nibbles(nibbles, LAYOUT_HEADER, layout=3, sel_a=6)
+
+  def emit(wave_id: int, inst: Inst, branch_taken: bool|None):
+    """Emit an SQTT packet for one executed instruction."""
+    w = wave_id & 0x1F
+    if wave_id not in started:
+      _emit_nibbles(nibbles, WAVESTART, delta=1, simd=0, wgp=0, wave=w, id7=wave_id)
+      started.add(wave_id)
+    inst_type, inst_op, op_name = type(inst), inst.op.value if hasattr(inst, 'op') else 0, inst.op.name if hasattr(inst, 'op') else ""
+    if issubclass(inst_type, _SOPP):
+      if inst_op in _SOPP_SKIP: return
+      if inst_op in _SOPP_IMMEDIATE: _emit_nibbles(nibbles, IMMEDIATE, delta=1, wave=w)
+      elif inst_op in _SOPP_BARRIER: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=InstOp.BARRIER)
+      elif inst_op in _SOPP_BRANCH: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=InstOp.JUMP if branch_taken else InstOp.JUMP_NO)
+      else: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=InstOp.SALU)
+    elif issubclass(inst_type, _VALU):
+      if (op := _valu_op(op_name)) is None: _emit_nibbles(nibbles, VALUINST, delta=1, wave=w)
+      else: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=op)
+    elif issubclass(inst_type, _SMEM): _emit_nibbles(nibbles, INST, delta=1, wave=w, op=InstOp.SMEM_RD)
+    else: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=_mem_op(inst_type, op_name))
+
+  def finish(wave_id: int):
+    """Emit WAVEEND for a completed wave."""
+    if wave_id in started: _emit_nibbles(nibbles, WAVEEND, delta=1, simd=0, wgp=0, wave=wave_id & 0x1F)
+
+  def finalize() -> bytes:
+    """Pad and return the encoded SQTT blob."""
+    while len(nibbles) % 2 != 0: nibbles.append(0)
+    nibbles.extend([0] * 32)
+    while len(nibbles) % 64 != 0: nibbles.append(0)
+    return bytes(nibbles[i] | ((nibbles[i + 1] if i + 1 < len(nibbles) else 0) << 4) for i in range(0, len(nibbles), 2))
+
+  return emit, finish, finalize
@@ -204,6 +204,13 @@ class MockUSB3:
    self.bulk_write(bytes(payload), timeout)
    return 0

+  def control_write_async(self, request:int, value:int=0, index:int=0, data:bytes=b"", timeout:int=1000) -> int:
+    self.control_write(request, value, index, data, timeout)
+    return 0
+
+  def control_read_async(self, request:int, length:int, value:int=0, index:int=0, timeout:int=1000) -> tuple[int, memoryview]:
+    return 0, self.control_read(request, length, value, index, timeout)
+
  def bulk_wait(self, tag:int): pass

  def bulk_read(self, length:int, timeout:int=1000) -> memoryview:
@@ -3,7 +3,6 @@ from tinygrad import dtypes, Context
 from tinygrad.dtype import DType, ConstType
 from tinygrad.uop.ops import Ops, UOp
 from test.helpers import full_rewrite
-import numpy as np

 class TestWeakConstFolding(unittest.TestCase):
  def test_weakint_math(self):
@@ -27,16 +26,14 @@ class TestBitcastConstFolding(unittest.TestCase):
    for val, src_dt, dst_dt, bits in ((3000000000, dtypes.int32, dtypes.uint32, 3000000000),
                                      (70000, dtypes.int16, dtypes.uint16, 4464),
                                      (-5, dtypes.uint32, dtypes.int32, -5)):
-      self.assertEqual(UOp.const(val, src_dt).bitcast(dst_dt).simplify().val, bits)
+      self.assertIs(UOp.const(val, src_dt).bitcast(dst_dt).simplify(), UOp.const(bits, dst_dt))

  def test_scalar_bitcast(self):
    def t(cases: dict[DType, ConstType]):
      for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()):
        if not math.isnan(from_v):
          r = UOp.const(from_v, from_dt).bitcast(to_dt).simplify()
-          self.assertEqual(r.op, Ops.CONST, msg:=f"{from_dt} -> {to_dt} ({from_v} -> {to_v})")
-          self.assertEqual(r.dtype, to_dt, msg)
-          np.testing.assert_equal(r.val, to_v, msg)
+          self.assertIs(r, UOp.const(to_v, to_dt), f"{from_dt} -> {to_dt} ({from_v} -> {to_v})")

    t({dtypes.int8: 0, dtypes.uint8: 0, dtypes.bool: False})
    t({dtypes.int8: 1, dtypes.uint8: 1, dtypes.bool: True})
@@ -419,7 +419,7 @@ class TestAutoCastType(unittest.TestCase):
    self.check_where_alternate_input_other(3, True, dtypes.weakint)

  def test_where_non_bool_cond_raises(self):
-    with self.assertRaises(RuntimeError): Tensor([1, 0, 2]).where(1, 0)
+    with self.assertRaises(RuntimeError): Tensor([1, 0, 2]).where(1, 0).dtype
    self.check_where_alternate_input_other(False, True, dtypes.bool)

  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
@@ -95,6 +95,26 @@ class TestLLMTokenizer(unittest.TestCase):
    self.assertEqual(template.end_turn(), "[/INST]")
    self.assertEqual(template.role("assistant"), "")

+  def test_tekken_gpt4o_split(self):
+    split = {p: SimpleTokenizer({}, {}, p)._split_to_word.findall for p in ("tekken", "gpt-4o")}
+    shared = {
+      "HelloWorld": ["Hello", "World"],
+      " ÜNICODE": [" ÜNICODE"],  # Ü: non-ascii upper joins the run
+      "é café": ["é", " café"],  # first é is e + U+0301 combining acute (NFD)
+      "เพื่อน วิ": ["เพื่อน", " วิ"],  # thai vowel marks stay in the word
+      "a/b\r\n x": ["a", "/b", "\r\n", " x"],  # punct tail eats /
+    }
+    for s, want in shared.items():
+      self.assertEqual(split["tekken"](s), want, f"tekken {s!r}")
+      self.assertEqual(split["gpt-4o"](s), want, f"gpt-4o {s!r}")
+    differ = [
+      ("12345", list("12345"), ["123", "45"]),  # digits: tekken single, o200k groups {1,3}
+      ("it's I'M don'T", ["it", "'s", " I", "'M", " don", "'T"], ["it's", " I'M", " don'T"]),  # contraction: o200k inline suffix
+    ]
+    for s, tk, go in differ:
+      self.assertEqual(split["tekken"](s), tk, f"tekken {s!r}")
+      self.assertEqual(split["gpt-4o"](s), go, f"gpt-4o {s!r}")
+
  def test_stream_decoder(self):
    """stream_decoder buffers incomplete UTF-8: token 25677 has 3/4 of emoji, token 138 completes it."""
    bs = [*range(33, 127), *range(161, 173), *range(174, 256)]
@@ -3,7 +3,7 @@ import unittest, itertools
 from tinygrad.codegen.late.coalesce import indexing_simplify
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite
-from tinygrad.uop.weak import pm_lower_index_dtype
+from tinygrad.uop.weak import pm_commit_weak
 from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load
 from tinygrad.helpers import Context
 from test.helpers import full_rewrite
@@ -496,7 +496,7 @@ class TestImageSimplification(unittest.TestCase):
    idx_y = (f + UOp.const(1.0)).cast(dtypes.int)
    load = get_load_image_uop((10, 10, 4), (UOp.const(-1) < idx_y) & (idx_y < UOp.const(10)),
                              (Special("gidx0", 10), idx_y))
-    off = graph_rewrite(load.sink(), pm_lower_index_dtype+indexing_simplify, ctx={}).src[0].src[0]
+    off = graph_rewrite(load.sink(), pm_commit_weak+indexing_simplify).src[0].src[0]
    self.assertEqual(off.src[1].get_valid(), UOp.const(True))

 class TestDropTrueGate(unittest.TestCase):
@@ -202,6 +202,11 @@ class TestUOpGraph(unittest.TestCase):
    invalid_lane_mul = next(u for u in out.src[0].toposort() if u.op is Ops.MUL)
    self.assertIs(invalid_lane_mul.dtype, dtypes.bool)

+  def test_devectorize_zero_sized_scalar_expand(self):
+    from tinygrad.codegen import devectorizer2
+    expanded = UOp.const(1.0).reshape(1, 1).expand(0, 3)
+    self.assertEqual(graph_rewrite(expanded, devectorizer2).shape, (0, 3))
+
  def test_gep_vec_const_fold(self):
    for vec_size in [2, 4, 8]:
      consts = [UOp.const(float(i), dtypes.float) for i in range(vec_size)]
@@ -230,11 +235,11 @@ class TestUOpGraph(unittest.TestCase):

  def test_depth_2_const_fold(self):
    v = UOp.variable("tmp", 0, 1, dtypes.int, param=True)
-    c2 = UOp.const(2, dtypes.int)
-    c4 = UOp.const(4, dtypes.int)
+    c2 = UOp.const(2)
+    c4 = UOp.const(4)
    vc = v+c2
    out = vc+c4
-    self.assertIs(out.simplify(), (v+UOp.const(6, dtypes.int)).simplify())
+    self.assertIs(out.simplify(), (v+UOp.const(6)).simplify())

  def test_bitcast_to_same_dtype_fold(self):
    for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
@@ -245,7 +250,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_sub_with_cast_folds(self):
    a = Variable("a", 0, 5)
-    out = a.cast(dtypes.int)+(-a).cast(dtypes.int)
+    out = a+(-a)
    self.assertIs(full_rewrite(out.sink()).src[0], full_rewrite(UOp.const(0, dtypes.int).sink()).src[0])

  def test_where_on_gated_load_fold(self):
@@ -429,7 +434,7 @@ class TestReduceCollapse(unittest.TestCase):

 class TestMovementOps(unittest.TestCase):
  def test_pm_mops_partial_reshape_index_removes_reshape(self):
-    from tinygrad.schedule.rangeify import pm_mops
+    from tinygrad.schedule.prepare import pm_mops
    src = UOp.param(0, dtypes.float, shape=(32, 4))
    r0, r1 = UOp.range(4, 0), UOp.range(8, 1)
    result = graph_rewrite(src.reshape((4, 8, 4)).index(r0, r1), pm_mops, name="test")
@@ -439,7 +444,7 @@ class TestMovementOps(unittest.TestCase):
    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.rangeify import pm_mops
+    from tinygrad.schedule.prepare import pm_mops
    src = UOp.param(0, dtypes.float, shape=(2, 6))
    result = graph_rewrite(src.reshape((2, 3, 2)).index(UOp.range(2, 0)), pm_mops, name="test")
    self.assertEqual(result.op, Ops.INDEX)
@@ -5,8 +5,7 @@ import z3
 from tinygrad.dtype import dtypes, ConstType, DType, Invalid
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad.uop.spec import spec_shared, type_verify
-from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load
-from tinygrad.uop.weak import pm_cast_weak
+from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load, symbolic_simple
 from tinygrad.uop.validate import uops_to_z3

 def check_uop_against_string(self, v:UOp, s:str):
@@ -36,7 +35,7 @@ class TestSymbolic(unittest.TestCase):
    self.assertEqual(solver.check(expr1 != expr2), z3.unsat, "simplified expression not equal to original")

  def helper_test_variable(self, v, n, m, s, test_z3:bool=True):
-    v_simplified = graph_rewrite(v, sym+pm_cast_weak, name="simplify symbolic uop")
+    v_simplified = graph_rewrite(v, sym, name="simplify symbolic uop")
    if test_z3: self.check_equal_z3(v, v_simplified)
    nmin, nmax = v_simplified.vmin, v_simplified.vmax
    check_uop_against_string(self, v_simplified, s)
@@ -149,6 +148,13 @@ class TestSymbolic(unittest.TestCase):
  def test_xor_0(self):
    self.helper_test_variable(Variable("a", 0, 8, dtypes.int) ^ 0, 0, 8, "a", test_z3=False)

+  def test_or_0(self):
+    self.helper_test_variable(Variable("a", 0, 8, dtypes.int) | 0, 0, 8, "a", test_z3=False)
+
+  def test_shift_0(self):
+    self.helper_test_variable(Variable("a", 0, 8, dtypes.int) << 0, 0, 8, "a")
+    self.helper_test_variable(Variable("a", 0, 8, dtypes.int) >> 0, 0, 8, "a")
+
  def test_xor_self_inverse(self):
    self.helper_test_variable((Variable("a", 0, 8, dtypes.int) ^ 5) ^ 5, 0, 8, "a", test_z3=False)

@@ -442,10 +448,20 @@ class TestSymbolic(unittest.TestCase):
  def test_and_remove(self):
    self.helper_test_variable(uand([uconst(1), Variable("a", 0, 1)]), 0, 1, "a")

+  def test_zero_div_zero_bottom_up(self):
+    # codegen runs symbolic_simple bottom_up, so the 0/0 is rewritten before its consts fold.
+    # without the guard the unsound x/x -> 1 below it claims this one.
+    z = UOp.const(0.0)
+    self.assertTrue(math.isnan(graph_rewrite(z/z, symbolic_simple, bottom_up=True).arg))
+
  def test_masked_shr_fold(self):
    x = UOp.variable('x', 0, 255, dtype=dtypes.uint32, param=True)
    self.helper_test_variable((x & -4) >> 2, 0, 63, "(x>>2)")

+  def test_masked_idiv_fold(self):
+    x = UOp.variable('x', 0, 255, dtype=dtypes.uint32, param=True)
+    self.helper_test_variable((x & -4) // 4, 0, 63, "(x//4)")
+
  def test_bool_or_not_tautology(self):
    a = Variable("a", 0, 10)
    c = a<10
@@ -1017,7 +1033,7 @@ class TestSymbolic(unittest.TestCase):
    cond = Variable("s", 0, 3, dtypes.int) < 2
    a = Variable("a", 0, 3, dtypes.int)
    self.assertIs(graph_rewrite(cond.where(a, a+1).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), (a+1).cast(dtypes.half)))
-    self.assertIs(graph_rewrite(cond.where(a, uconst(2)).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), UOp.const(2, dtypes.half)))
+    self.assertIs(graph_rewrite(cond.where(a, uconst(2)).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), uconst(2.0)))
    self.assertIs(graph_rewrite(cond.where(a, UOp.invalid()).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), UOp.invalid()))

  def test_where_const_gate_keeps_stated_width(self):
@@ -1449,6 +1465,11 @@ class TestGatedUopGivenValid(unittest.TestCase):
    self.assertEqual(idx, (r0 < 3).where(expected_vec, UOp.invalid()))

 class TestRangeSplitting(unittest.TestCase):
+  def test_end_preserves_constant_backedge(self):
+    loop, backedge = UOp.loop(0), UOp.const(False)
+    end = graph_rewrite(UOp(Ops.NOOP).end(loop, backedge), sym)
+    self.assertEqual(end.src, (UOp(Ops.NOOP), loop, backedge))
+
  def test_range_split_on_mod(self):
    # test that mark_range_mod splits RANGE(8) into RANGE(4)*2 + RANGE(2) when used with %2
    from tinygrad.codegen.simplify import pm_split_ranges, pm_flatten_range
@@ -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), dtype=dtypes.float)
+    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)
@@ -6,7 +6,7 @@ from tinygrad.helpers import Timing, Context, cdiv
 from tinygrad.dtype import dtypes, AddrSpace, ConstFloat, Invalid  # noqa: F401
 from tinygrad.device import Device
 from tinygrad.uop.ops import Ops, AxisType, ParamArg, PatternMatcher, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
-from tinygrad.uop.weak import pm_lower_index_dtype
+from tinygrad.uop.weak import pm_lower_weak
 from tinygrad.uop.spec import spec_program, spec_shared, type_verify
 from tinygrad.uop.symbolic import sym, pm_remove_invalid
 from test.helpers import eval_uop, to_uops_list
@@ -56,7 +56,7 @@ class TestDTypeFromUOp(unittest.TestCase):
      if u.is_invalid)), (dtypes.float32, dtypes.float32, dtypes.bool))
    invalid, value = UOp.invalid(), UOp.const(1, dtypes.float32)
    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, dtypes.float32, src=(value, invalid)), UOp(Ops.ADD, dtypes.float32, src=(value, 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)
    gate, value = UOp.param(0, dtypes.bool, ()), UOp.param(1, dtypes.float, ())
@@ -64,7 +64,7 @@ class TestDTypeFromUOp(unittest.TestCase):
    type_verify(out.sink(), spec_program)

  def test_remove_invalid_stack_lanes(self):
-    stack = UOp(Ops.STACK, dtypes.half, (UOp.const(1, dtypes.half), UOp.invalid()))
+    stack = UOp(Ops.STACK, src=(UOp.const(1, dtypes.half), UOp.invalid()))
    out = graph_rewrite(stack, pm_remove_invalid)
    self.assertEqual(out.src, (UOp.const(1, dtypes.half), UOp.const(0, dtypes.half)))
    type_verify(out.sink(), spec_program)
@@ -76,16 +76,18 @@ class TestLowerIndexDtype(unittest.TestCase):
    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_index_dtype)
-    self.assertTrue(all(u.dtype != dtypes.weakint for u in lowered.backward_slice_with_self), "lowering must resolve all weakint")
+    lowered = graph_rewrite(shrink.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")
    sh = next(u for u in lowered.backward_slice_with_self if u.op is Ops.SHRINK)
    self.assertEqual(sh.src[1].dtype, dtypes.long)

  def test_reg_buffer_size_lowers(self):
    reg = UOp.placeholder((4,), dtypes.float, 0, addrspace=AddrSpace.REG)
    self.assertEqual(reg.src[0].dtype, dtypes.weakint)
-    lowered = graph_rewrite(reg.sink(), pm_lower_index_dtype)
-    self.assertTrue(all(u.dtype != dtypes.weakint for u in lowered.backward_slice_with_self), "lowering must resolve all 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):
@@ -280,9 +282,9 @@ class TestFastIdiv(unittest.TestCase):
  def test_division_power_of_two(self):
    for dt in (dtypes.int32, dtypes.uint32):
      g = UOp.param(0, dt, (3,))
-      c = UOp.const(2).cast(dt)
+      c = UOp.const(2)
      l = g.index(c)
-      a = UOp(Ops.CDIV, dt, (l, c))
+      a = UOp(Ops.CDIV, src=(l, c))
      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
      Device[Device.DEFAULT].renderer.render(uops)
      ops = [x.op for x in uops]
@@ -293,31 +295,32 @@ class TestFastIdiv(unittest.TestCase):
    # 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,))
-      c = UOp.const(8).cast(dt)
-      a = UOp(Ops.FLOORMOD, dt, (g.index(c), c))
+      c = UOp.const(8)
+      a = UOp(Ops.FLOORMOD, src=(g.index(c), c))
      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
      ops = [x.op for x in uops]
      self.assertIn(Ops.AND, ops, f"For dtype={dt} FLOORMOD by pow2 did not simplify to AND")
      self.assertNotIn(Ops.CMOD, ops, f"For dtype={dt} FLOORMOD by pow2 left a MOD")
      self.assertNotIn(Ops.FLOORMOD, ops, f"For dtype={dt} FLOORMOD survived past late rewrite")

-  def test_floordiv_power_of_two_uint(self):
-    # uint FLOORDIV by a power of two lowers to a shift, leaving no IDIV/FLOORDIV in the kernel
-    for dt in (dtypes.uint32, dtypes.uint64):
+  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,))
-      c = UOp.const(2).cast(dt)
-      a = UOp(Ops.FLOORDIV, dt, (g.index(c), c))
+      c = UOp.const(2)
+      a = UOp(Ops.FLOORDIV, src=(g.index(c), c))
      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
      ops = [x.op for x in uops]
      self.assertIn(Ops.SHR, ops, f"For dtype={dt} FLOORDIV by power of two did not simplify to shift")
      self.assertNotIn(Ops.CDIV, ops, f"For dtype={dt} FLOORDIV by power of two did not simplify to shift")
+      self.assertNotIn(Ops.CMOD, ops, f"For dtype={dt} FLOORDIV by pow2 kept the round toward zero correction")
      self.assertNotIn(Ops.FLOORDIV, ops, f"For dtype={dt} FLOORDIV survived past late rewrite")

  @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,))
-    c = UOp.const(3).cast(dtypes.uint)
+    c = UOp.const(3)
    l = g.index(c)
    a = UOp(Ops.CDIV, src=(l, c))
    uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -337,7 +340,7 @@ class TestFastIdiv(unittest.TestCase):
  def test_fast_idiv_bounded_numerator_zero(self):
    x = UOp.variable("x", 0, 1, dtype=dtypes.int32)
    for val in range(2):
-      self.assertEqual(eval_uop(x.alu(Ops.CDIV, UOp.const(3).cast(x.dtype)), vals=(val,)), cdiv(val, 3))
+      self.assertEqual(eval_uop(x.alu(Ops.CDIV, UOp.const(3)), vals=(val,)), cdiv(val, 3))

  @Context(DISABLE_FAST_IDIV=0)
  def test_fast_idiv_remove_powers_of_two(self):
@@ -362,7 +365,7 @@ class TestFastIdiv(unittest.TestCase):

  def test_disable_fast_idiv(self):
    g = UOp.param(0, dtypes.uint32, (4,))
-    c = UOp.const(3).cast(dtypes.uint)
+    c = UOp.const(3)
    l = g.index(c)
    a = UOp(Ops.CDIV, src=(l, c))
    with Context(DISABLE_FAST_IDIV=1):
@@ -466,10 +469,10 @@ class TestUOpRender(unittest.TestCase):
    self.assertEqual(UOp.range(1, 0, src=(shrink,), dtype=dtypes.int).render(simplify=False), "r0")

  def test_render_vectorize_empty(self):
-    u = UOp(Ops.STACK, dtype=dtypes.void, src=())
+    u = UOp(Ops.STACK, src=())
    self.assertEqual(u.render(simplify=False), "{}")
  def test_render_vectorize_empty_simplified(self):
-    u = UOp(Ops.STACK, dtype=dtypes.void, src=())
+    u = UOp(Ops.STACK, src=())
    self.assertEqual(u.render(), "{}")
  def test_render_vectorize_same(self):
    u = UOp(Ops.STACK, src=(UOp.const(0),)*3)
@@ -14,37 +14,37 @@ class TestValidateOOB(unittest.TestCase):
  def test_const_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (16,))
-      to_uops_list([buf.index(UOp.const(0)).load(dtype=dtypes.int)])  # valid
-      to_uops_list([buf.index(UOp.const(15)).load(dtype=dtypes.int)])  # valid (last element)
+      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):
-        to_uops_list([buf.index(UOp.const(16)).load(dtype=dtypes.int)])  # off by one
+        to_uops_list([buf.index(UOp.const(16)).load()])  # off by one
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(UOp.const(42)).load(dtype=dtypes.int)])  # way out
+        to_uops_list([buf.index(UOp.const(42)).load()])  # way out

  def test_variable_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (16,))
-      to_uops_list([buf.index(Variable("i", 0, 15)).load(dtype=dtypes.int)])  # valid
+      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(dtype=dtypes.int)])  # oob
+        to_uops_list([buf.index(Variable("i", 0, 20)).load()])  # oob
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(Variable("i", -5, 10)).load(dtype=dtypes.int)])  # negative
+        to_uops_list([buf.index(Variable("i", -5, 10)).load()])  # negative

  def test_range_with_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (16,))
      r = UOp.range(42, 0, AxisType.GLOBAL)
-      to_uops_list([buf.index(r.valid(r < 16)).load(dtype=dtypes.int)])  # valid
+      to_uops_list([buf.index(r.valid(r < 16)).load()])  # valid
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(r.valid(r < 17)).load(dtype=dtypes.int)])  # oob
+        to_uops_list([buf.index(r.valid(r < 17)).load()])  # oob

  def test_variable_with_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (16,))
      v = Variable("v", -5, 80)
-      to_uops_list([buf.index(v.valid((v >= 0) & (v < 16))).load(dtype=dtypes.int)])  # valid
+      to_uops_list([buf.index(v.valid((v >= 0) & (v < 16))).load()])  # valid
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(v.valid(v < 20)).load(dtype=dtypes.int)])  # negative not masked
+        to_uops_list([buf.index(v.valid(v < 20)).load()])  # negative not masked

  def test_gated_store(self):
    with Context(CHECK_OOB=1, SPEC=2):
@@ -58,62 +58,62 @@ class TestValidateOOB(unittest.TestCase):
  def test_floordiv(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (16,))
-      to_uops_list([buf.index(UOp.range(32, 0, AxisType.GLOBAL) // 2).load(dtype=dtypes.int)])  # 0..15 valid
+      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(dtype=dtypes.int)])  # 0..16 oob
+        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,))
      r = UOp.range(100, 0, AxisType.GLOBAL)
-      to_uops_list([buf.index(r % 16).load(dtype=dtypes.int)])  # 0..15 valid
+      to_uops_list([buf.index(r % 16).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(r % 20).load(dtype=dtypes.int)])  # 0..19 oob
+        to_uops_list([buf.index(r % 20).load()])  # 0..19 oob

  def test_shr(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (16,))
-      to_uops_list([buf.index(UOp.range(64, 0, AxisType.GLOBAL) >> 2).load(dtype=dtypes.int)])  # 0..15 valid
+      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(dtype=dtypes.int)])  # 0..31 oob
+        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,))
      r = UOp.range(8, 0, AxisType.GLOBAL)
-      to_uops_list([buf.index(r << 2).load(dtype=dtypes.int)])  # 0..28 valid
+      to_uops_list([buf.index(r << 2).load()])  # 0..28 valid
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(r << 4).load(dtype=dtypes.int)])  # 0..112 oob
+        to_uops_list([buf.index(r << 4).load()])  # 0..112 oob

  def test_and(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (16,))
      r = UOp.range(100, 0, AxisType.GLOBAL)
-      to_uops_list([buf.index(r & 15).load(dtype=dtypes.int)])  # 0..15 valid
+      to_uops_list([buf.index(r & 15).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(r & 31).load(dtype=dtypes.int)])  # 0..31 oob
+        to_uops_list([buf.index(r & 31).load()])  # 0..31 oob
      # align masks round down to a multiple of 2^k
-      to_uops_list([buf.index((r & -4).valid(r < 16)).load(dtype=dtypes.int)])  # 0..12 valid
+      to_uops_list([buf.index((r & -4).valid(r < 16)).load()])  # 0..12 valid
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(r & -2).load(dtype=dtypes.int)])  # 0..100 oob
+        to_uops_list([buf.index(r & -2).load()])  # 0..100 oob
      # other masks can't be modeled as mod
      with self.assertRaisesRegex(RuntimeError, "z3 int AND only supports"):
-        to_uops_list([buf.index(r & 21).load(dtype=dtypes.int)])
+        to_uops_list([buf.index(r & 21).load()])

  def test_max(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (16,))
-      to_uops_list([buf.index(Variable("v", -10, 15).maximum(0)).load(dtype=dtypes.int)])  # 0..15 valid
+      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(dtype=dtypes.int)])  # 0..20 oob
+        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,))
      r = UOp.range(32, 0, AxisType.GLOBAL)
-      to_uops_list([buf.index(r.valid((r < 8) ^ ((r >= 8) & (r < 16)))).load(dtype=dtypes.int)])  # 0..15 valid
+      to_uops_list([buf.index(r.valid((r < 8) ^ ((r >= 8) & (r < 16)))).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(r.valid((r < 10) ^ (r >= 20))).load(dtype=dtypes.int)])  # 0..9,20..31 oob
+        to_uops_list([buf.index(r.valid((r < 10) ^ (r >= 20))).load()])  # 0..9,20..31 oob

  # cast patterns
  def test_float_cast_in_index(self):
@@ -121,13 +121,13 @@ class TestValidateOOB(unittest.TestCase):
      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(dtype=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,))
      r = UOp.range(20, 0)
-      to_uops_list([buf.index(r.valid(r.cast(dtypes.bool).logical_not())).load(dtype=dtypes.int)])  # only r=0 valid
+      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):
@@ -135,18 +135,18 @@ class TestValidateOOB(unittest.TestCase):
      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(dtype=dtypes.int).cast(dtypes.weakint)
-      to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32))).load(dtype=dtypes.int)])  # valid
+      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
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 64))).load(dtype=dtypes.int)])  # oob
+        to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 64))).load()])  # oob

  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,))
      shrink = UOp(Ops.SHRINK, src=(buf0, UOp.const(0, dtypes.int), UOp.const(4)))
-      ld0 = shrink.load(dtype=dtypes.int).index(0)
-      to_uops_list([buf1.index(ld0.valid((ld0 >= 0) & (ld0 < 64))).load(dtype=dtypes.int)])
+      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):
@@ -185,18 +185,18 @@ class TestViz(unittest.TestCase):
    @dataclass(frozen=True)
    class TestStruct:
      colored_field: str
-    a = UOp(Ops.CUSTOM, arg=TestStruct(colored("xyz", "magenta")+colored("12345", "blue")))
+    a = UOp(Ops.PYLITERAL, arg=TestStruct(colored("xyz", "magenta")+colored("12345", "blue")))
    a2 = uop_to_json(VizData(), a)[id(a)]
-    self.assertEqual(ansistrip(a2["label"]), f"CUSTOM\n{TestStruct.__qualname__}(colored_field='xyz12345')")
+    self.assertEqual(ansistrip(a2["label"]), f"PYLITERAL\n{TestStruct.__qualname__}(colored_field='xyz12345')")

  def test_colored_label_multiline(self):
    with save_viz() as viz:
      arg = colored("x", "green")+"\n"+colored("y", "red")+colored("z", "yellow")+colored("ww\nw", "magenta")
      src = [Tensor.empty(1).uop for _ in range(10)]
-      a = UOp(Ops.CUSTOM, src=tuple(src), arg=arg)
+      a = UOp(Ops.PYLITERAL, src=tuple(src), arg=arg)
      exec_rewrite(a, [PatternMatcher([])])
    a2 = next(viz.get_details(0, 0))["graph"][id(a)]
-    self.assertEqual(ansistrip(a2["label"]), "CUSTOM\nx\nyzww\nw")
+    self.assertEqual(ansistrip(a2["label"]), "PYLITERAL\nx\nyzww\nw")

  def test_inf_loop(self):
    a = UOp.const(3)
@@ -347,7 +347,7 @@ class TestVizGC(unittest.TestCase):
      init = bufs_allocated()
      a = UOp.new_buffer("NULL", 10, dtypes.char)
      a.buffer.allocate()
-      exec_rewrite(UOp(Ops.CUSTOM, src=(a,), arg=a), [PatternMatcher([])])
+      exec_rewrite(UOp(Ops.PYLITERAL, src=(a,), arg=a), [PatternMatcher([])])
      del a
      self.assertEqual(bufs_allocated()-init, 0)
    lst = viz.list_items()
@@ -474,7 +474,7 @@ class TestVizIntegration(unittest.TestCase):
    def custom_fn(X:UOp):
      X = X.flatten()
      i = UOp.range(X.numel(), 0)
-      custom_op = UOp(Ops.CUSTOMI, src=(X[i],), arg="{} + undeclared_name")
+      custom_op = UOp(Ops.CUSTOMI, src=(X[i],), arg=("{} + undeclared_name", X.dtype))
      return X[i].store(custom_op).end(i).sink(arg=KernelInfo(name=f"custom_fn_{X.numel()}"))
    x = Tensor.custom_kernel(Tensor.empty(1, device="CPU"), fxn=custom_fn)[0]
    with save_viz() as viz:
@@ -21,7 +21,7 @@ class TestBenchLog(unittest.TestCase):
    # check event list
    for event in BenchEvent:
      self.assertEqual(len(_events[event]["wall"]), 1)
-      self.assertGreater(_events[event]["wall"][0], 0)
+      self.assertGreater(_events[event]["wall"][0][0], 0)

  def test_log_double_wall_time(self):
    for event in BenchEvent:
@@ -35,8 +35,8 @@ class TestBenchLog(unittest.TestCase):
    # check event list
    for event in BenchEvent:
      self.assertEqual(len(_events[event]["wall"]), 2)
-      self.assertGreater(_events[event]["wall"][0], 0)
-      self.assertGreater(_events[event]["wall"][1], 0)
+      self.assertGreater(_events[event]["wall"][0][0], 0)
+      self.assertGreater(_events[event]["wall"][1][0], 0)

  @skipIf(_SKIP_KERNEL_TIMING, "ci timing is not accurate")
  def test_log_single_kernel_time(self):
@@ -52,8 +52,8 @@ class TestBenchLog(unittest.TestCase):
    # check event list
    for event in BenchEvent:
      self.assertEqual(len(_events[event]["kernel"]), 1)
-      self.assertLess(_events[event]["kernel"][0], wall_times[0])
-      self.assertGreater(_events[event]["kernel"][0], 0)
+      self.assertLess(_events[event]["kernel"][0][0], wall_times[0])
+      self.assertGreater(_events[event]["kernel"][0][0], 0)

  @skipIf(_SKIP_KERNEL_TIMING, "ci cuda timing is not accurate")
  def test_interleaved_wall_kernel_time(self):
@@ -74,8 +74,8 @@ class TestBenchLog(unittest.TestCase):
    for event in BenchEvent:
      self.assertEqual(len(_events[event]["wall"]), 1)
      self.assertEqual(len(_events[event]["kernel"]), 1)
-      self.assertLess(_events[event]["kernel"][0], wall_times[0])
-      self.assertGreater(_events[event]["kernel"][0], 0)
+      self.assertLess(_events[event]["kernel"][0][0], wall_times[0])
+      self.assertGreater(_events[event]["kernel"][0][0], 0)

  @skipIf(_SKIP_KERNEL_TIMING, "ci cuda timing is not accurate")
  def test_stacked_wall_kernel_time(self):
@@ -93,10 +93,10 @@ class TestBenchLog(unittest.TestCase):
    for event in BenchEvent:
      self.assertEqual(len(_events[event]["wall"]), 2)
      self.assertEqual(len(_events[event]["kernel"]), 2)
-      self.assertLess(_events[event]["kernel"][0], _events[event]["wall"][0])
-      self.assertGreater(_events[event]["kernel"][0], 0)
-      self.assertLess(_events[event]["kernel"][1], _events[event]["wall"][1])
-      self.assertGreater(_events[event]["kernel"][1], 0)
+      self.assertLess(_events[event]["kernel"][0][0], _events[event]["wall"][0][0])
+      self.assertGreater(_events[event]["kernel"][0][0], 0)
+      self.assertLess(_events[event]["kernel"][1][0], _events[event]["wall"][1][0])
+      self.assertGreater(_events[event]["kernel"][1][0], 0)

  def test_log_instant_event(self):
    for event in InstantBenchEvent:
@@ -105,7 +105,7 @@ class TestBenchLog(unittest.TestCase):
    # check event list
    for event in InstantBenchEvent:
      self.assertEqual(len(_events[event]), 1)
-      self.assertEqual(_events[event][0], 1000)
+      self.assertEqual(_events[event][0][0], 1000)

 if __name__ == '__main__':
  unittest.main()
@@ -3,6 +3,7 @@ from tinygrad import Tensor, UOp, dtypes
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import Ops
 from test.helpers import KernelCountException
+from tinygrad.engine.realize import run_linear

 class TestRingAllReduce(unittest.TestCase):
  def test_schedule_ring(self):
@@ -21,13 +22,26 @@ class TestRingAllReduce(unittest.TestCase):
  def test_schedule_all2all(self):
    with Context(ALL2ALL=2):
      N = 4
+      M = N*100
      ds = tuple(f"CPU:{i}" for i in range(N))
-      t = Tensor.empty(N, N*100).shard(ds, axis=0).realize()
-      linear = t.sum(0).mul(2.0).contiguous().linear_with_vars()[0]
+      x = Tensor.arange(N*M, dtype=dtypes.float).reshape(N, M)
+      t = (x*x).clone().shard(ds, axis=0).realize()
+      out = t.sum(0).mul(2.).contiguous()
+      linear, var_vals = out.linear_with_vars()
      copies = [si for si in linear.src if si.src[0].op is Ops.COPY]
      sinks = [si for si in linear.src if si.src[0].op is Ops.SINK]
-      if len(copies) != 24: raise KernelCountException(24, len(copies))
-      if len(sinks) != 26: raise KernelCountException(26, len(sinks))
+      # 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)
+      if len(sinks) != sink_count: raise KernelCountException(sink_count, len(sinks))
+      # correctness
+      run_linear(linear, var_vals)
+      expected = [2*sum((d*M+i)**2 for d in range(N)) for i in range(M)]
+      dev_nums = Tensor.arange(1, N+1, dtype=dtypes.float).reshape(N, 1).expand(N, M).shard(ds, axis=0)
+      shards = out.reshape(1, M).expand(N, M)+dev_nums
+      self.assertListEqual(shards.tolist(), [[x+d+1 for x in expected] for d in range(N)])

  @Context(RING=0, ALL2ALL=0)
  def test_schedule_naive(self):
@@ -5,6 +5,8 @@ from tinygrad.llm.model import (
  GatedDeltaNetBlock, SSMConfig, TransformerBlock, TransformerConfig,
  apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk,
 )
+from tinygrad.llm.kernels.amd import Linear, gated_delta_prefill, amd_custom_kernels_supported
+from tinygrad.llm.gguf import ggml_data_to_tensor

 def apply_rope(x:Tensor, start_pos:int):
  B, H, T, Hd = x.shape
@@ -12,6 +14,15 @@ def apply_rope(x:Tensor, start_pos:int):
  freqs_cis = precompute_freqs_cis(Hd, start_pos+T)[start_pos:start_pos+T]
  return apply_rope_new(x, freqs_cis)

+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)):
+      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)
+      linear.set_quantized(decoded)
+      self.assertEqual((linear.ggml_type, linear.weight.numel()), (ggml_type, words))
+
 class TestAttention(unittest.TestCase):
  def test_apply_rope(self):
    x = Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32)
@@ -41,6 +52,23 @@ class TestAttention(unittest.TestCase):
    np.testing.assert_allclose(block.cache_kv[0, :, :, :seqlen, :].numpy(), expected.numpy(), rtol=1e-5, atol=1e-5)

 class TestGatedDeltaNetBlock(unittest.TestCase):
+  def test_gated_delta_rectangular_state_and_row_decay(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    rng = np.random.default_rng(42)
+    q, k = (rng.normal(size=(1, 1, 3, 32)).astype(np.float32) for _ in range(2))
+    v, beta = rng.normal(size=(1, 1, 3, 4)).astype(np.float32), rng.uniform(size=(1, 1, 3)).astype(np.float32)
+    alpha, initial = rng.uniform(0.8, 1, size=(1, 1, 3, 4)).astype(np.float32), rng.normal(size=(1, 1, 4, 32)).astype(np.float32)
+    expected_state, expected_out = initial.copy(), np.empty_like(v)
+    for t in range(3):
+      previous, av = expected_state.copy(), alpha[:, :, t, :, None]
+      delta = (v[:, :, t] - (previous*k[:, :, t, None]).sum(-1)*alpha[:, :, t]) * beta[:, :, t, None]
+      expected_state = previous*av + delta[..., None]*k[:, :, t, None, :]
+      expected_out[:, :, t] = (previous*q[:, :, t, None]).sum(-1)*alpha[:, :, t] + delta*(q[:, :, t]*k[:, :, t]).sum(-1)
+    state = Tensor(initial).contiguous().realize()
+    out = gated_delta_prefill(Tensor(q), Tensor(k), Tensor(v), Tensor(beta), Tensor(alpha), state).realize()
+    np.testing.assert_allclose(out.numpy(), expected_out, rtol=1e-4, atol=1e-4)
+    np.testing.assert_allclose(state.numpy(), expected_state, rtol=1e-4, atol=1e-4)
+
  def _tensor_linspace(self, start:float, stop:float, shape:tuple[int, ...]) -> Tensor:
    return Tensor.linspace(start, stop, int(np.prod(shape)), dtype=dtypes.float32).reshape(*shape)

@@ -0,0 +1,115 @@
+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()
@@ -0,0 +1,40 @@
+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()
@@ -4,7 +4,7 @@ from tinygrad import Tensor, dtypes, TinyJit
 from tinygrad.helpers import Context
 from tinygrad.dtype import least_upper_float
 from tinygrad.uop.ops import UOp, Ops, GroupOp, dtype_from_uop, graph_rewrite
-from tinygrad.uop.weak import pm_lower_index_dtype, pm_commit_weak
+from tinygrad.uop.weak import pm_commit_weak
 from tinygrad.uop.symbolic import symbolic_simple
 from tinygrad.uop.spec import spec_shared, type_verify
 from tinygrad.engine.jit import JitError
@@ -74,7 +74,7 @@ class TestWeakPromotion(unittest.TestCase):
    recips = [u for u in (x / y)._uop.toposort() if u.op is Ops.RECIPROCAL]
    self.assertEqual([(u.dtype, u.src[0].dtype) for u in recips], [(dtypes.float32, dtypes.float32)])
    with Context(DEFAULT_FLOAT=dtypes.float16):
-      committed = graph_rewrite((UOp.const(1).cast(dtypes.int32) + UOp.const(1.0)).cast(dtypes.float32), pm_lower_index_dtype, ctx={})
+      committed = graph_rewrite((UOp.const(1).cast(dtypes.int32) + UOp.const(1.0)).cast(dtypes.float32), pm_commit_weak)
    self.assertEqual([u.dtype for u in committed.toposort() if u.op is Ops.ADD], [dtypes.float32])

  def test_div_sub_operand_kept_weak(self):
@@ -85,7 +85,7 @@ class TestWeakPromotion(unittest.TestCase):
  def test_cast_weak_expression_commits_at_cast_floor(self):
    # the floor never narrows: a cast BELOW the default does not pull the compute width down with it
    with Context(DEFAULT_FLOAT=dtypes.float32):
-      narrowed = graph_rewrite((UOp.const(1.0) + UOp.const(2.0)).cast(dtypes.float16), pm_lower_index_dtype, ctx={})
+      narrowed = graph_rewrite((UOp.const(1.0) + UOp.const(2.0)).cast(dtypes.float16), pm_commit_weak)
    self.assertEqual((narrowed.dtype, narrowed.src[0].dtype), (dtypes.float16, dtypes.float32))

  def test_cast_weak_expression_value_uses_cast_floor(self):
@@ -114,27 +114,35 @@ class TestWeakPromotion(unittest.TestCase):
    self.assertIsInstance((x + 2).src[1].val, float)
    self.assertIs(x + UOp.const(2), x + 2)

-  def test_index_dtype_ignores_weakness(self):
-    with Context(SPEC=2):
-      idx = UOp.const(0).cast(dtypes.int32)
-      weak = UOp.const(1.0).expand((1,))
-      self.assertEqual(UOp(Ops.INDEX, dtypes.float32, (weak, idx)).dtype, dtypes.float32)
-      with self.assertRaisesRegex(RuntimeError, "bad dtype"): UOp(Ops.INDEX, dtypes.int32, (weak, idx))
-
  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))
      gate = UOp.const(True)
-      out = graph_rewrite(dst.store(UOp.const(5.0), gate), pm_lower_index_dtype, ctx={})
+      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
    self.assertEqual((out.src[1], out.src[2]), (UOp.const(5.0, dtypes.bfloat16), gate))

  def test_weak_srcs_commit_only_at_a_concrete_lub(self):
    weak_lub = UOp(Ops.ADD, src=(UOp.const(1), UOp.const(1.0)))
-    self.assertIs(graph_rewrite(weak_lub, pm_lower_index_dtype, ctx={}), weak_lub)
+    self.assertIs(graph_rewrite(weak_lub, pm_commit_weak), weak_lub)
    concrete = UOp.const(2.0).cast(dtypes.float16)
-    where = graph_rewrite(UOp(Ops.WHERE, src=(UOp.const(True), concrete, UOp.const(1.0))), pm_lower_index_dtype, ctx={})
-    self.assertEqual(tuple(x.dtype for x in where.src), (dtypes.bool, dtypes.float16, dtypes.float16))
+    # the weak arm stays bare: its sibling states the width, so the WHERE already derives float16 for it
+    where = graph_rewrite(UOp(Ops.WHERE, src=(UOp.const(True), concrete, UOp.const(1.0))), pm_commit_weak)
+    self.assertEqual((where.dtype, tuple(x.dtype for x in where.src)), (dtypes.float16, (dtypes.bool, dtypes.float16, dtypes.weakfloat)))
+
+  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()
+    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)
+
+  def test_committed_const_conversion_folds_for_native_format(self):
+    folded = graph_rewrite(UOp.const(16256, dtypes.ushort).cast(dtypes.uint), symbolic_simple)
+    self.assertIs(folded, UOp.const(16256, dtypes.uint))
+    # fmt-less targets are lowered by renderer rewrites, where collapsing this pair would cycle with float-intermediate insertion.
+    emulated = UOp.const(1.0, dtypes.float).cast(dtypes.bfloat16)
+    self.assertIs(graph_rewrite(emulated, symbolic_simple), emulated)

  def test_weak_shift_lhs_commits_the_node(self):
    # a shift derives its lhs's dtype, so committing the lhs restates the root (WGSL's packed store writes `mask << shift_am`)
@@ -175,11 +183,11 @@ class TestWeakPromotion(unittest.TestCase):
    self.assertEqual(dtype_from_uop(Ops.SHL, (UOp.const(1, dtypes.int8), UOp.const(1, dtypes.uint32)), None), dtypes.int8)
    self.assertEqual(UOp.const(1).alu(Ops.SHL, UOp.const(1, dtypes.uint)).dtype, dtypes.weakint)
    self.assertEqual((v & 3).dtype, dtypes.weakint)
-    with self.assertRaises(RuntimeError): Tensor.const(1.0) << Tensor.const(1.0)
-    with self.assertRaises(RuntimeError): UOp.const(1, dtypes.int32).alu(Ops.SHL, UOp.const(1, dtypes.float64))
+    with self.assertRaises(RuntimeError): (Tensor.const(1.0) << Tensor.const(1.0)).dtype
+    with self.assertRaises(RuntimeError): UOp.const(1, dtypes.int32).alu(Ops.SHL, UOp.const(1, dtypes.float64)).dtype
    for op in (Ops.SHL, Ops.SHR):
      with self.assertRaises(RuntimeError):
-        UOp.const(1, dtypes.float32).alu(op, UOp.const(1, dtypes.int32))
+        UOp.const(1, dtypes.float32).alu(op, UOp.const(1, dtypes.int32)).dtype
    # float bitwise builds, the spec rejects it
    with Context(SPEC=1):
      f32, wf = UOp.const(1.0, dtypes.float32), UOp.const(1.0)
@@ -0,0 +1,111 @@
+import unittest
+import numpy as np
+from tinygrad import Tensor, UOp, dtypes, nn
+from tinygrad.llm.kernels.amd import Linear, amd_custom_kernels_supported, q8_quantize, flash_attention
+from tinygrad.llm.gguf import ggml_data_to_tensor
+
+class TestQ8Quantize(unittest.TestCase):
+  def test_word_quant_weights_use_typed_buffer_view(self):
+    for ggml_type, type_size in ((13, 176), (23, 136)):
+      with self.subTest(ggml_type=ggml_type):
+        raw = Tensor(np.zeros(type_size + 4, dtype=np.uint8), device="CPU").contiguous().realize()[4:]
+        decoded = ggml_data_to_tensor(raw, 256, ggml_type).reshape(1, 256)
+        linear = Linear(256, 1, bias=False)
+        linear.set_quantized(decoded)
+        self.assertEqual(linear.ggml_type, ggml_type)
+        self.assertEqual(linear.weight.dtype, dtypes.uint32)
+        self.assertEqual(linear.weight.nbytes(), type_size)
+        self.assertEqual(linear.weight.uop.buf_uop.buffer.offset, 4)
+
+  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)
+    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")
+    rng = np.random.default_rng(42)
+    packed = rng.integers(0, 256, 210, dtype=np.uint8)
+    packed[-2:] = np.array([0.01], dtype=np.float16).view(np.uint8)
+    raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:]
+    decoded = ggml_data_to_tensor(raw, 256, 14).reshape(1, 256)
+    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)
+
+  def test_q4_k_linear(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    rng = np.random.default_rng(42)
+    in_features, blocks = 2048, 16*2048//256
+    packed = rng.integers(0, 256, blocks*144, dtype=np.uint8)
+    for i in range(blocks): packed[i*144:i*144+4] = np.array([0.01, 0.002], dtype=np.float16).view(np.uint8)
+    raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:]
+    decoded = ggml_data_to_tensor(raw, 16*in_features, 12).reshape(16, in_features)
+    weight = decoded.numpy()
+    linear = Linear(in_features, 16, bias=False)
+    nn.state.load_state_dict(linear, {"weight":decoded}, verbose=False, realize=False)
+    x = rng.normal(size=(3, in_features)).astype(np.float32)
+    scale = np.maximum(np.abs(x).reshape(3, in_features//32, 32).max(-1, keepdims=True) / 127, 1e-8)
+    xq = np.clip(np.rint(x.reshape(3, in_features//32, 32) / scale), -127, 127) * scale
+    np.testing.assert_allclose(linear(Tensor(x)).numpy(), xq.reshape(3, in_features) @ weight.T, rtol=2e-3, atol=2e-2)
+    self.assertEqual(linear.ggml_type, 12)
+
+  def test_q6_linear_multiple_tokens(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    rng = np.random.default_rng(42)
+    in_features, blocks = 2048, 16*2048//256
+    packed = rng.integers(0, 256, blocks*210, dtype=np.uint8)
+    for i in range(blocks): packed[i*210+208:i*210+210] = np.array([0.01], dtype=np.float16).view(np.uint8)
+    raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:]
+    decoded = ggml_data_to_tensor(raw, 16*in_features, 14).reshape(16, in_features)
+    weight = decoded.numpy()
+    linear = Linear(in_features, 16, bias=False)
+    nn.state.load_state_dict(linear, {"weight":decoded}, verbose=False, realize=False)
+    x = rng.normal(size=(3, in_features)).astype(np.float32)
+    scale = np.maximum(np.abs(x).reshape(3, in_features//32, 32).max(-1, keepdims=True) / 127, 1e-8)
+    xq = np.clip(np.rint(x.reshape(3, in_features//32, 32) / scale), -127, 127) * scale
+    np.testing.assert_allclose(linear(Tensor(x)).numpy(), xq.reshape(3, in_features) @ weight.T, rtol=2e-3, atol=2e-2)
+    self.assertEqual(linear.ggml_type, 14)
+
+    # symbolic token counts take the padded kernel path and give the same results
+    generic = Linear(in_features, 16, bias=False)
+    nn.state.load_state_dict(generic, {"weight":decoded}, verbose=False, realize=False)
+    sym = Tensor(np.concatenate([x, np.zeros((1, in_features), np.float32)])).contiguous()[:UOp.variable("tokens", 1, 4).bind(3)]
+    np.testing.assert_allclose(generic(sym)[:3].numpy(), xq.reshape(3, in_features) @ weight.T, rtol=2e-3, atol=2e-2)
+    self.assertTrue(generic.use_custom_quant)
+    self.assertEqual(generic.ggml_type, 14)
+
+  def test_attention_uses_physical_cache_length(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    q, k, v = Tensor.zeros(1, 2, 1, 32), Tensor.randn(1, 1, 1, 32), Tensor.randn(1, 1, 1, 32)
+    cache = Tensor.empty(2, 1, 1, 256, 32, dtype=dtypes.half).contiguous()
+    assigned = Tensor(cache.uop.after(cache[:, :, :, 0:1, :].uop.store(Tensor.stack(k, v).cast(dtypes.half).uop)))
+    out = flash_attention(q, assigned, 1).realize()
+    np.testing.assert_allclose(out.numpy(), v.expand(1, 2, 1, 32).numpy(), rtol=2e-2, atol=2e-2)
+
+  def test_prefill_attention_unaligned_start(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    rng = np.random.default_rng(42)
+    start_pos = 1718
+    q = Tensor.zeros(1, 8, 32, 128)
+    old_kv = rng.normal(size=(2, 1, 1, start_pos, 128)).astype(np.float32)
+    new_kv = rng.normal(size=(2, 1, 1, 32, 128)).astype(np.float32)
+    cache = Tensor.zeros(2, 1, 1, 2048, 128, dtype=dtypes.half).contiguous()
+    Tensor.realize(cache[:, :, :, :start_pos].assign(Tensor(old_kv).cast(dtypes.half)))
+    sp = UOp.variable("start_pos", 0, 2047).bind(start_pos)
+    assigned = Tensor(cache.uop.after(cache[:, :, :, sp:sp+32, :].uop.store(Tensor(new_kv).cast(dtypes.half).uop)))
+    out = flash_attention(q, assigned, sp+32).realize()
+    values = np.concatenate([old_kv[1, 0, 0], new_kv[1, 0, 0]]).astype(np.float16).astype(np.float32)
+    expected = np.stack([values[:start_pos+i+1].mean(0) for i in range(32)])[None, None].repeat(8, axis=1)
+    np.testing.assert_allclose(out.numpy(), expected, rtol=2e-3, atol=2e-3)
+
+if __name__ == "__main__": unittest.main()
@@ -44,7 +44,10 @@ class TestTransformerGenerate(unittest.TestCase):
      return Tensor([[42]])
    with patch.object(Transformer, '__call__', mock_call):
      next(model.generate([1, 2, 3, 4, 5, 42, 10]))
-    self.assertEqual(calls, [((1, 1), V_START_POS.bind(5)), ((1, 1), V_START_POS.bind(6))])
+    # resumes from the reused state at position 5 and consumes the 2 new tokens (one chunk or two decode steps)
+    self.assertEqual(calls[0][1], V_START_POS.bind(5))
+    def ntok(shape): return shape[1] if isinstance(shape[1], int) else shape[1].unbind()[1]
+    self.assertEqual(sum(ntok(c[0]) for c in calls), 2)

  def test_recurrent_divergent_prompt_restarts(self):
    model, calls = Transformer(TEST_CONFIG), []
@@ -852,82 +852,6 @@ 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))
@@ -3,7 +3,7 @@ import itertools, functools
 from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC
 from tinygrad.helpers import ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT, NUM_CPU_THREADS, TC_SELECT, TC_OPT, TracingKey, Context, panic
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, Ops, UPat, rewrite_group, KernelInfo, ProgramInfo, GroupOp, AxisType
-from tinygrad.uop.weak import pm_lower_index_dtype, pm_commit_weak, pm_cast_weak
+from tinygrad.uop.weak import pm_lower_weak, pm_commit_weak, pm_cast_const
 from tinygrad.uop.render import pyrender
 from tinygrad.uop.spec import type_verify, spec_tensor, spec_program
 from tinygrad.renderer import Renderer, Estimates
@@ -22,11 +22,11 @@ from tinygrad.codegen.opt.postrange import apply_opts
 from tinygrad.codegen.late.gater import pm_move_gates_from_index
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse, pm_reduce_unparented
 from tinygrad.schedule.multi import multi_pm
-from tinygrad.schedule.rangeify import pm_mops
+from tinygrad.schedule.prepare import pm_mops
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
 from tinygrad.codegen.late.regalloc import LinearScanRegallocContext, pm_regalloc_rewrite
 from tinygrad.codegen.late.coalesce import memory_coalescing, pm_simplify_add_image
-from tinygrad.helpers import all_same, flatten, argsort, partition
+from tinygrad.helpers import all_same, all_int, flatten, argsort, partition
 from tinygrad.uop.ops import _broadcast_shape, identity_element
 from tinygrad.schedule.rangeify import BufferizeOpts

@@ -162,9 +162,10 @@ devectorizer2 = mop_cleanup+pm_mops+PatternMatcher([
  (UPat(Ops.RESHAPE, dtype=dtypes.void, name="x"), lambda x: x.src[0]),
  # reshape of a single element shaped value to scalar is an index
  (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0].index(0) if x.marg == () and x.src[0].shape == (1,) else None),
-  # EXPAND on scalar -> STACK
+  # EXPAND on scalar -> nested STACKs with the same shape
  (UPat(Ops.EXPAND, src=(UPat.var("x"), UPat()), name="out"),
-   lambda x,out: UOp.stack(*([x]*out.max_numel())) if x.shape == () and out.shape == (out.max_numel(),) else None),
+   lambda x,out: functools.reduce(lambda x,s: UOp.stack(*([x]*s)), reversed(out.shape), x)
+   if x.shape == () and all_int(out.shape) and 0 not in out.shape else None),
 ])

 def fix_group_for_reduce(x:UOp):
@@ -233,10 +234,11 @@ pm_reduce_local = pm_wmma_add+PatternMatcher([
  (UPat(Ops.SINK, name="sink"), merge_reduce_ends),
 ])+pm_clean_up_group_sink

+def is_shape_changing_bitcast(u:UOp): return u.op is Ops.BITCAST and u.shape != u.src[0].shape
 def maybe_load(u:UOp): return u.load() if u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL, AddrSpace.REG) else u
 pm_add_loads = PatternMatcher([
-  # BITCAST?
-  (UPat(GroupOp.Elementwise|{Ops.REDUCE,Ops.WMMA,Ops.STACK}, name="x"), lambda x: x.replace(src=tuple([maybe_load(u) for u in x.src]))),
+  (UPat(GroupOp.Elementwise|{Ops.REDUCE,Ops.WMMA,Ops.STACK}, name="x"),
+   lambda x: None if is_shape_changing_bitcast(x) else x.replace(src=tuple(map(maybe_load, x.src)))),
  (UPat(Ops.STORE, name="x"), lambda x: x.replace(src=(x.src[0], maybe_load(x.src[1]))+x.src[2:])),
 ])

@@ -281,10 +283,6 @@ pm_implicit_barriers = PatternMatcher([
  (UPat(Ops.END, name="end"), add_war_barrier),
 ])

-pm_casted_consts = PatternMatcher([
-  (UPat(Ops.CONST, dtypes.all, name="c"), lambda c: UOp.cconst(c.val, c.dtype)),
-])
-
 def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
  if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
  if DEBUG >= 5: print(pyrender(ast))
@@ -346,11 +344,13 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:

  # extra symbolic before decomp. crashes without this?
  # NOTE: also run indexing_simplify here, while the index is still weakint and (x+y)*c -> x*c+y*c applies
-  sink = graph_rewrite(sink, sym+indexing_simplify, name="extra symbolic")
+  # commit widths minted in this fixpoint before lowering inspects INDEX shapes
+  sink = graph_rewrite(sink, sym+indexing_simplify+pm_commit_weak, name="extra symbolic")

-  # lower index dtype
+  # the boundary: required compute dtypes settle here; derivable const edges may stay bare
  # NOTE: we need indexing_simplify to remove the cast to long using the Invalid
-  sink = graph_rewrite(sink, symbolic_simple+pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes")
+  # NOTE: symbolic must NOT be composed here -- pm_data_invalid pushes the weak result CAST into a gated WHERE, remaking the weak node, and it cycles
+  sink = graph_rewrite(sink, pm_lower_weak+indexing_simplify, name="lower all index dtypes")

  # final symbolic before decomp
  sink = graph_rewrite(sink, symbolic, name="final symbolic")
@@ -374,12 +374,11 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:

  # final rules for the renderer (without sym)
  extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([])
-  pm_final_rewrite = pm_commit_weak+pm_cast_weak+pm_decomp+extra_matcher+pm_split_ends
+  pm_final_rewrite = pm_commit_weak+pm_decomp+extra_matcher+pm_split_ends
  sink = graph_rewrite(sink, pm_final_rewrite+pm_remove_invalid, ctx=ren, name="final rewrite")

-  # spell every literal as a casted const CAST(dt, CONST(value))
-  # TODO: remove once consts are always weak
-  sink = graph_rewrite(sink, pm_casted_consts, name="casted consts", walk=True)
+  # commit every const still bare so no renderer reads one
+  sink = graph_rewrite(sink, pm_cast_const, name="cast consts")

  # add implicit barriers (stores/loads through LOCAL memory ordered by AFTER or across loop iterations need workgroup barriers)
  sink = graph_rewrite(sink, pm_implicit_barriers, name="add implicit barriers")
@@ -1,8 +1,9 @@
 from dataclasses import replace
 from tinygrad.dtype import dtypes, DType, truncate
-from tinygrad.helpers import flatten, DEBUG, EMULATED_DTYPES, Context, SPEC
+from tinygrad.helpers import flatten, DEBUG, EMULATED_DTYPES
 from tinygrad.uop import GroupOp
-from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite, ParamArg
+from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite
+from tinygrad.uop.weak import commit_weak_consts
 from tinygrad.renderer import Renderer
 from tinygrad.codegen.decomp.transcendental import exponent_bias, shl, shr

@@ -25,10 +26,10 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
  match op:
    case Ops.NEG: return l2i(Ops.SUB, dt, zero, zero, *uops)
    case Ops.CAST if dt in (dtypes.long, dtypes.ulong) and uops[0].dtype not in dtypes.floats:
-      # the high word is the sign extension; bool has no sign, test the already-cast low word instead (bool < 0 would promote to weakint)
+      # the high word is the sign extension, and unsigned and bool sources zero extend
      x, lo = uops[0], uops[0].cast(l2i_dt[dt])
-      sign = lo if x.dtype is dtypes.bool else x
-      return lo, (sign < sign.const_like(0)).where(lo.const_like(-1), lo.const_like(0))
+      if x.dtype is dtypes.bool or x.dtype in dtypes.uints: return lo, lo.const_like(0)
+      return lo, (x < x.const_like(0)).where(lo.const_like(-1), lo.const_like(0))
    case Ops.CAST if dt in (dtypes.long, dtypes.ulong):
      return (lo:=uops[0].cast(l2i_dt[dt])), (uops[0] / 2**32).cast(l2i_dt[dt]) - ((uops[0] < 0) & lo.ne(0))
    case Ops.CAST if dt in dtypes.floats:
@@ -127,19 +128,22 @@ def f2f_clamp(val:UOp, dt:DType, sat=True) -> UOp:
  return val.ne(val).where(val, (val < -mx).where(-sat, (mx < val).where(sat, val)))

 def f2f_load(x: UOp, fr:DType, to:DType) -> UOp:
-  if (n:=x.max_numel()) == 1: return f2f(x.replace(dtype=f2f_dt[fr]), fr, to)
-  return UOp(Ops.STACK, src=tuple(f2f(x.replace(dtype=f2f_dt[fr], src=(reindex(x.src[0], i, 1),)), fr, to) for i in range(n)))
+  storage_idx = graph_rewrite(x.src[0], pm_float_decomp, ctx=(fr, to), bottom_up=True)
+  if (n:=x.max_numel()) == 1: return f2f(storage_idx.load(*x.src[1:]), fr, to)
+  return UOp(Ops.STACK, src=tuple(f2f(reindex(storage_idx, i, 1).load(*x.src[1:]), fr, to) for i in range(n)))

 def f2f_store(st, idx, val, fr:DType, to:DType):
  if (n:=val.max_numel()) == 1: return st.replace(src=(idx, f2f(val.bitcast(f2f_dt[to]), to, fr)))
  return UOp.group(*(st.replace(src=(reindex(idx, i, 1), f2f(val.index(i).bitcast(f2f_dt[to]), to, fr))) for i in range(n)))

 # tag is the 32-bit word this node becomes - (0 for the low word, 1 for the high, the dtype the consumer wants)
-pm_long_decomp = PatternMatcher([
-  (UPat(GroupOp.Defines, src=(UPat.var("sz"),), name="x"), lambda x,sz:
-   x.replace(dtype=l2i_dt[x.dtype], arg=replace(x.arg, dtype=l2i_dt[x.dtype]), src=(sz*2,)) if x.dtype in l2i_dt else None),
+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()), 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(dtype=x.tag[1], tag=None) if x.tag is not None else None),
+   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:
   st.replace(src=(idx.rtag((0, dt:=l2i_dt[idx.dtype])), val.rtag((0, dt)))).group(
     st.replace(src=(idx.rtag((1, dt)), val.rtag((1, dt))))) if val.tag is None else None),
@@ -147,6 +151,9 @@ pm_long_decomp = PatternMatcher([
   split_l2i(ctx, x.op, dt:=l2i_dt[a.dtype], *flatten((s.rtag((0, dt)), s.rtag((1, dt))) for s in x.src))),
  (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda ctx,a,x:
   split_l2i(ctx, Ops.BITCAST, l2i_dt[x.dtype], a.rtag((0, dt:=l2i_dt[a.dtype])), a.rtag((1, dt)))[x.tag[0]]),
+  # a const splits by value; the general CAST arm below would drop its high word
+  (UPat(Ops.CAST, src=(UPat(Ops.CONST, name='c'),), tag={(w, dt) for w in (0, 1) for dt in l2i_dt.values()}, name='x'),
+   lambda x,c: UOp.const(truncate[x.tag[1]](c.val >> (32*x.tag[0])), x.tag[1])),
  (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a'),), name="x"), lambda ctx,a,x:
   split_l2i(ctx, x.op, x.dtype, a)[x.tag[0]] if x.tag is not None else None),
  (UPat(Ops.CAST, src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda ctx,a,x:
@@ -159,21 +166,22 @@ pm_long_decomp = PatternMatcher([
  (UPat((*(GroupOp.ALU - GroupOp.Comparison - {Ops.SHL, Ops.SHR, Ops.WHERE}), Ops.BITCAST), tuple(l2i_dt.keys()), name="x"), lambda ctx,x:
   split_l2i(ctx, x.op, l2i_dt[x.dtype], *flatten((a.rtag((0, l2i_dt[x.dtype])), a.rtag((1, l2i_dt[x.dtype]))) for a in x.src))[x.tag[0]]
   if x.tag is not None else None),
-  (UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda x,idx:
-   x.replace(dtype=l2i_dt[x.dtype], src=(reindex(idx, x.tag[0]).replace(dtype=l2i_dt[x.dtype], tag=None),), tag=None) if x.tag is not None else None),
-  (UPat(Ops.CONST, tag={(w, dt) for w in (0, 1) for dt in l2i_dt.values()}, name='x'), lambda x:
-   UOp.const(truncate[x.tag[1]]((x.val >> 32) if x.tag[0] == 1 else (x.val & 0xFFFFFFFF)), x.tag[1]))
+  (UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda ctx,x,idx:
+   reindex(graph_rewrite(idx, pm_long_decomp, ctx=ctx, bottom_up=True), x.tag[0]).replace(tag=None).load() if x.tag is not None else None)
 ])

 # float decomposition patterns - ctx is (fr, to) tuple
-pm_float_decomp = PatternMatcher([
-  (UPat((*GroupOp.Defines, Ops.INDEX, Ops.SHRINK), name="x"), lambda ctx,x:
-   x.replace(dtype=f2f_dt[ctx[0]], arg=replace(x.arg, dtype=f2f_dt[ctx[0]]) if isinstance(x.arg, ParamArg) else x.arg, tag=ctx[0])
-   if x.dtype == ctx[0] and (x.op is not Ops.INDEX or x.src[0].op not in {Ops.LOAD, Ops.STACK}) else None),
+pm_float_decomp: PatternMatcher = PatternMatcher([
+  (UPat(GroupOp.Defines, name="x"), lambda ctx,x:
+   UOp(x.op, src=x.src, arg=replace(x.arg, dtype=f2f_dt[ctx[0]]), tag=ctx[0]) if x.dtype == ctx[0] else None),
+  # INDEX into a LOAD/STACK selects a lane of an already converted value, the load rules below own those
+  (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat(GroupOp.All-{Ops.LOAD, Ops.STACK}),), allow_any_len=True, name="x"), lambda ctx,x:
+   UOp(x.op, src=(graph_rewrite(x.src[0], pm_float_decomp, ctx=ctx, bottom_up=True), *x.src[1:]), arg=x.arg, tag=ctx[0])
+   if x.dtype == ctx[0] else None),
  (UPat(Ops.LOAD, dtypes.floats, name="x"), lambda ctx,x: f2f_load(x, *ctx) if x.dtype == ctx[0] else None),
  # bitcasted load should just replace load
  (UPat(Ops.BITCAST, src=(UPat(Ops.LOAD, name="ld"),), name="bc"), lambda ctx,bc,ld:
-   ld.replace(dtype=f2f_dt[ctx[0]]).bitcast(bc.dtype) if ld.dtype == ctx[0] else None),
+   graph_rewrite(ld.src[0], pm_float_decomp, ctx=ctx, bottom_up=True).load(*ld.src[1:]).bitcast(bc.dtype) if ld.dtype == ctx[0] else None),
  # bitcast from
  (UPat(Ops.BITCAST, src=(UPat.var("x", dtypes.floats),), name="bc"), lambda ctx,bc,x:
   bc.replace(src=(f2f(x.bitcast(f2f_dt[ctx[1]]), ctx[1], ctx[0]),)) if x.dtype == ctx[1] and bc.dtype.bitsize == ctx[0].bitsize else None),
@@ -182,26 +190,21 @@ pm_float_decomp = PatternMatcher([
   f2f(x.bitcast(f2f_dt[ctx[0]]), ctx[0], ctx[1]) if bc.dtype == ctx[0] else None),
  (UPat(Ops.CAST, dtypes.floats, src=(UPat.var("val"),), name="x"), lambda ctx,x,val:
   f2f_clamp(val.cast(ctx[1]), ctx[0]) if x.dtype == ctx[0] else None),
-  # a CONST has no srcs to cast, it restates its value at the emulating dtype
-  (UPat(Ops.CONST, dtypes.floats, name="x"), lambda ctx,x: UOp.const(x.val, ctx[1]) if x.dtype == ctx[0] else None),
  (UPat(GroupOp.All-GroupOp.Defines-{Ops.CAST, Ops.BITCAST, Ops.CONST}, dtypes.floats, name="x"), lambda ctx,x:
-   x.replace(dtype=ctx[1], src=tuple(s.cast(ctx[1]) if s.dtype == ctx[0] else s for s in x.src))
-   if x.dtype == ctx[0] else None),
+   UOp(x.op, src=tuple(s.cast(ctx[1]) if s.dtype == ctx[0] else s for s in x.src), arg=x.arg, tag=x.tag) if x.dtype == ctx[0] else None),
  (UPat(Ops.STORE, src=(UPat.var("idx"), UPat(Ops.BITCAST, dtypes.floats, name="val")), name='st'), lambda ctx,st,idx,val:
-   st.replace(src=(idx, val.replace(dtype=f2f_dt[ctx[0]]))) if val.dtype == ctx[0] and idx.tag == ctx[0] else None),
-  (UPat(Ops.STORE, src=(UPat.var("idx"), UPat.var("val", dtypes.floats)), name='st'), lambda ctx,st,idx,val:
-   f2f_store(st, idx, val, *ctx) if val.dtype == ctx[1] and (idx:=idx.src[0] if idx.op == Ops.CAST else idx).tag == ctx[0] else None),
+   st.replace(src=(idx, val.src[0].bitcast(f2f_dt[ctx[0]]))) if val.dtype == ctx[0] and idx.tag == ctx[0] else None),
+  (UPat(Ops.STORE, src=(UPat.var("idx").or_casted(), UPat.var("val", dtypes.floats)), name='st'), lambda ctx,st,idx,val:
+   f2f_store(st, idx, val, *ctx) if val.dtype == ctx[1] and idx.tag == ctx[0] else None),
 ])

 def do_dtype_decomps(sink:UOp, ctx:tuple[set[DType], Renderer]) -> UOp:
  def _should_emulate(dt): return dt in EMULATED_DTYPES.tolist(dtypes) or dt not in ctx[1].supported_dtypes()
-  # NOTE: dtype decomp creates intermediate UOps that don't follow the spec (e.g. half LOAD on ushort BUFFER)
-  with Context(SPEC=min(SPEC.value, 1)):
-    for fr in sorted(filter(_should_emulate, ctx[0])):
-      to = dtypes.int if fr == dtypes.long else dtypes.half if not _should_emulate(dtypes.half) and fr in dtypes.fp8s else dtypes.float
-      if DEBUG >= 2: print(f"emulating {fr} as {to}")
-      pm = pm_float_decomp if fr in dtypes.floats else pm_long_decomp
-      sink = graph_rewrite(sink, pm, name=f"decomp {fr} -> {to}", ctx={} if pm is pm_long_decomp else (fr, to), bottom_up=True)
+  for fr in sorted(filter(_should_emulate, ctx[0])):
+    to = dtypes.int if fr == dtypes.long else dtypes.half if not _should_emulate(dtypes.half) and fr in dtypes.fp8s else dtypes.float
+    if DEBUG >= 2: print(f"emulating {fr} as {to}")
+    pm = pm_float_decomp if fr in dtypes.floats else pm_long_decomp
+    sink = graph_rewrite(sink, pm, name=f"decomp {fr} -> {to}", ctx={} if pm is pm_long_decomp else (fr, to), bottom_up=True)
  ctx[0].clear()
  return sink

@@ -75,7 +75,11 @@ powers_of_two: dict[int, int] = {2**i:i for i in range(64)}
@functools.cache
 def get_simplifying_rewrite_patterns(ops:tuple[Ops, ...]) -> PatternMatcher:
  # these are rewrites that make things simpler
-  pat: list[tuple[UPat, Callable]] = [(UPat.var("a")//UPat.var("b"), floordiv_to_idiv)]
+  pat: list[tuple[UPat, Callable]] = []
+  # FLOORDIV by 2**y -> x >> y (an arithmetic shift is exactly floor division for any sign); fires before floordiv_to_idiv
+  if Ops.SHR in ops: pat.append((UPat.var("x", dtypes.ints)//UPat.cvar("c"),
+    lambda x,c: x >> v if (v:=powers_of_two.get(c.val, 0)) else None))
+  pat.append((UPat.var("a")//UPat.var("b"), floordiv_to_idiv))
  # FLOORMOD by 2**y -> x & (2**y-1) (correct floor mod for any sign in two's complement); fires before floormod_to_mod
  if Ops.AND in ops: pat.append((UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.val-1) if c.val in powers_of_two else None))
  pat.append((UPat.var("a")%UPat.var("b"), floormod_to_mod))
@@ -90,8 +90,8 @@ def payne_hanek_reduction(d:UOp) -> tuple[UOp, UOp]:
    if count+offset < len(two_over_pi_f) - 1:
      an = i.ne(count).where(_take(an, offset, count=count+1), an.const_like(two_over_pi_f[count+offset]))
    return an
-  def _shl_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) * pow2if(y, d.dtype).cast(dtypes.uint64)).cast(dtypes.uint32)
-  def _shr_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) // pow2if(y, d.dtype).cast(dtypes.uint64)).cast(dtypes.uint32)
+  def _shl_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) << y.cast(dtypes.uint64)).cast(dtypes.uint32)
+  def _shr_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) >> y.cast(dtypes.uint64)).cast(dtypes.uint32)

  a = [_take(UOp.const(0, dtypes.uint32), i) for i in range(4)]
  #  (two_over_pi_f[Int(i) + n] << e) | (two_over_pi_f[Int(i) + n+1] >> (nbits - e))
@@ -95,7 +95,6 @@ def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
 pm_simplify_add_image = PatternMatcher([
  (UPat(Ops.SHRINK, src=(UPat(Ops.PARAM, name="buf"), UPat(name="x"), UPat(arg=4))), transform_to_image),
  # image load/store is always float
-  (UPat(Ops.INDEX, dtype=dtypes.float, name="x").load(dtype=dtypes.half), lambda x: x.load().cast(dtypes.half)),
  (UPat(Ops.INDEX, dtype=dtypes.float, name="x").store(UPat(name="d", dtype=dtypes.half)), lambda x,d: x.store(d.cast(dtypes.float))),
  (UPat.var("x", dtype=dtypes.float).cast(dtypes.half).cast(dtypes.float), lambda x: x),
 ])
@@ -104,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], dict[int, list[UOp]]] = defaultdict(dict)
+  memory: defaultdict[tuple[Ops, UOp, UOp|str, UOp, object], 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}:
@@ -119,11 +118,12 @@ 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
-      memory[(u.op, buf, root_src, valid)].setdefault(arg, []).append(u)
+      # 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)

  # build replacements
  replacements = {}
-  for (op,buf,base,valid),offsets in memory.items():
+  for (op,buf,base,valid,ld_arg),offsets in memory.items():
    # allowed lengths (copied in)
    lengths = []
    must_divide = True
@@ -158,7 +158,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()
+          ld = idx.load(arg=ld_arg)
          for i,g in enumerate(grp):
            for oo in offsets[g]:
              replacements[oo] = ld.index(i) if len(grp) > 1 else ld
@@ -193,7 +193,7 @@ class Scheduler:
      for b in self.bufs:
        if rng in (i:=b.src[1].get_idx()).backward_slice_with_self:
          nb = b.replace(src=(b.src[0], i.valid(valid&b.src[1].get_valid())))
-          replaces[b] = nb if b in store_targets else valid.where(nb, UOp.const(Invalid, b.dtype))
+          replaces[b] = nb if b in store_targets else valid.where(nb, UOp.const(Invalid))
      self.ast = self.ast.substitute(replaces, f"padto {rng.arg[:-1]} {opt.arg}")
    elif opt.op is OptOps.SWAP:
      try:
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	6234b42cf7	low rank call in kernel graph	2026-08-27 10:34:54 -07:00
geohot	ae9da86538	fix custom_kernel wrapper	2026-08-27 10:18:37 -07:00
geohot	df6df141ec	call has rank <= 1	2026-08-27 10:15:35 -07:00
chenyuandGitHub	54d69f439f	fix test_assign for no subbuffer backends (#17757 ) * test_assign to test/backend * fix	2026-08-27 10:47:21 -04:00
nimlgenandGitHub	980c42d611	hcq2 tiny speedups (#17773 )	2026-08-27 17:42:37 +03:00
nimlgenandGitHub	7eb763a3c2	bnxt to extra (#17772 ) * bnxt to extra * x * x * les	2026-08-27 17:29:17 +03:00
Teddy TennantandGitHub	5c3d044465	fix asinh gradient at zero (#17758 )	2026-08-27 09:51:17 -04:00
nimlgenandGitHub	f19d89e29a	hcq2: a bit cleaner deps (#17771 )	2026-08-27 14:24:59 +03:00
tonyfloatersuandGitHub	2d4e7afcbf	mockgpu/amd: use generated KFD ioctl metadata (#17761 ) * init commit * init commit * minimize to single file scope	2026-08-27 12:24:37 +03:00
qazalandGitHub	0d66815849	set BEAM_PADTO=0 in llama (#17769 )	2026-08-27 16:45:33 +09:00
qazalandGitHub	d16a2510ad	remove some dead scheduler tests (#17768 ) * remove tests that use .uop * that test does nothing	2026-08-27 16:01:28 +09:00
qazalandGitHub	118e812b31	simplify more same device copies (#17766 ) * simplify more same device copies * stack of contigs or copies * rename	2026-08-27 12:54:12 +09:00
sirhcmandGitHub	0820c7a424	benchmarks: openpilot matrix (#17714 )	2026-08-26 18:53:55 -04:00
wozeparrotandGitHub	22ffff3258	gptoss: single jit (#17525 )	2026-08-26 15:32:21 -07:00
George HotzandGitHub	a0a901c8e4	faster qwen 3.8 (#17720 ) * faster qwen * test fix * dead code * fix gguf issue * pretty nt loads * lil * use warp	2026-08-26 15:23:25 -07:00
sirhcmandGitHub	c015351ac5	fix _get_cpu_count for docker --cpus=N in python 3.13+ (#17760 )	2026-08-26 16:53:30 -04:00
nimlgenandGitHub	6074c002e1	hcq2: fix jit (#17747 ) * hcq2: fix reduce * Dx * inputs table * emoty commit	2026-08-26 22:36:36 +03:00
George HotzandGitHub	6042b87272	delete PCONTIG [PR] (#17756 ) * delete PCONTIG * cleanups	2026-08-26 12:07:28 -07:00
sirhcmandGitHub	cc72b9f7be	cleanup BENCHMARK_LOG (#17754 )	2026-08-26 14:55:54 -04:00
chenyuandGitHub	6a3b297548	fix PTX NIR SPEC=2 for bool [pr] (#17753 ) * fix PTX NIR SPEC=2 for bool [pr] storing bool with uint8 needs to pass SPEC * the fix	2026-08-26 14:07:29 -04:00
George HotzandGitHub	ea6c82f3be	small changes from new rangeify (#17752 )	2026-08-26 10:33:31 -07:00
chenyuandGitHub	0abcf09b74	never bufferize_to_store weak input [PR] (#17751 ) github github	2026-08-26 13:25:52 -04:00
chenyuandGitHub	4bdc865131	delete unused rewrite rules [PR] (#17748 )	2026-08-26 10:50:42 -04:00
Teddy TennantandGitHub	4c20f1d357	fix asinh precision loss on negative inputs (#17749 )	2026-08-26 10:08:04 -04:00
qazalandGitHub	ecf79e260d	better all2all schedule test (#17746 ) * better all2all schedule test * deconstruct those numbers * reorder	2026-08-26 15:46:10 +09:00
b1tgandGitHub	9860e5d285	llm tokenizer: fix tekken, add gpt4o (#17733 )	2026-08-25 23:30:35 -07:00
George HotzandGitHub	625c05df1e	fix am_smi to respect dev (#17742 )	2026-08-25 19:21:55 -07:00
sirhcmandGitHub	b49c03fb1c	benchmarks: don't use sudo on mac (#17740 )	2026-08-25 21:06:04 -04:00
George HotzandGitHub	dc04c7820e	lil fixes from new_rangeify (#17741 ) * lil fixes from new_rangeify * gpt sol review	2026-08-25 17:15:28 -07:00
chenyuandGitHub	6ece327cf3	CUSTOM arg is (str, dtype) [PR] (#17737 )	2026-08-25 20:15:06 -04:00
wozeparrotandGitHub	07268b724f	fix: external_test_gpu_crash on python 3.14 (#17739 )	2026-08-25 15:43:30 -07:00
sirhcmandGitHub	55032514ce	pin onnxruntime==1.24.1 (#17738 )	2026-08-25 17:53:15 -04:00
YassineYousfiandGitHub	2824504f90	usb amd: yield between signal polls (#17712 ) * AMD_USB_POLL_US * com * its micro	2026-08-25 14:11:53 -07:00
sirhcmandGitHub	7fef98c86e	benchmarks: remove usage of sudo from linux runners (#17735 )	2026-08-25 16:21:08 -04:00
geohot	b831ca62d9	Reapply "disk cache: thread-local db conn (#17694 )" This reverts commit `df528499ce`.	2026-08-25 13:20:15 -07:00
nimlgenandGitHub	d851aca9ae	hcq2: fix copy/call in usb (#17732 )	2026-08-25 19:03:19 +03:00
chenyuandGitHub	7dc8b666e7	wgsl cast before load for packed [pr] (#17731 ) instead of explicit dtype on load	2026-08-25 11:46:31 -04:00
nimlgenandGitHub	023bfdb380	fix viz for hcq2 (#17729 )	2026-08-25 18:23:22 +03:00
chenyuandGitHub	9f01775cf4	dtype_from_uop for THREEFRY and FDIV (#17727 )	2026-08-25 08:43:59 -04:00
chenyuandGitHub	9607787ce1	delete more explicitly set dtype to UOp [PR] (#17726 ) * delete more explicitly set dtype to UOp [PR] * not that	2026-08-25 08:32:45 -04:00
qazalandGitHub	ab79879613	bring back sqtt examples tests (#17725 )	2026-08-25 17:00:18 +09:00
George HotzandGitHub	a5678317c2	split rangeify to prepare.py (#17722 )	2026-08-24 19:12:13 -07:00
chenyuandGitHub	1d694dd700	remove more explicitly set dtype [PR] (#17721 )	2026-08-24 22:03:22 -04:00
chenyuandGitHub	65ca68567e	const cleanups [PR] (#17719 )	2026-08-24 21:37:14 -04:00
chenyuandGitHub	a7df1a1ace	const are weak 3 [pr] (#17695 )	2026-08-24 20:32:57 -04:00
RaineandGitHub	d9004cff22	use native shifts in Payne Hayek instead of pow2 mul/divs (#17717 ) * use native shifts instead of pow2 mul/divs * wow	2026-08-24 17:11:46 -07:00
wozeparrotandGitHub	6f87158d77	feat: bump version to 0.14.0 (#17716 )	2026-08-24 15:40:15 -07:00
sirhcmandGitHub	021c015eb4	benchmarks: more aggressive timeouts (#17715 )	2026-08-24 18:32:36 -04:00
sirhcmandGitHub	2e7790f16c	benchmarks: cleanup dsp benchmark (#17706 )	2026-08-24 18:09:01 -04:00
chenyuandGitHub	8df3dac0ec	remove redundant UOp dtype [PR] (#17713 )	2026-08-24 18:00:02 -04:00
George HotzandGitHub	ed110993d3	mi350p: aqua vanjaram bringup for the raw PCI driver (#17693 ) * mi350p: aqua vanjaram bringup for the raw PCI driver Minimum stable bringup, all root-caused: - disable ASPM on the PCI path: link L1 across retimers makes GPU reads oscillate to 0xffffffff (no-op on non-sysfs backends) - probe live AIDs via MMHUB FB_LOCATION; dead ones read 0xffffffff and indirect writes to them corrupt the fabric - aqua reset semantics: full boot over a live state can kill the fabric, partial boot + reset_mec is the deepest safe reset - tolerate the HQD dequeue timeout like the kernel (wedged waves can survive RESET_WAVES) * am/aspm: simplify to GPU-endpoint-only clearing via pci config abstractions * fix AID derivation for non-dense sdma instance keys + bound the aspm capability walk for 0xff reads	2026-08-24 14:32:16 -07:00
chenyuandGitHub	f4cc28824c	delete GETTUPLE on UNSHARD rule [PR] (#17711 ) not used	2026-08-24 17:25:28 -04:00
George HotzandGitHub	76bf6b7eec	llm: add DEBUG=2 for weights loading (#17710 ) * llm: add DEBUG=2 for weights loading * 00	2026-08-24 14:09:34 -07:00
geohot	08fbd25f1a	AGENTS: no amend commits	2026-08-24 13:34:42 -07:00
George HotzandGitHub	5aabbb1991	fix Muon weight decay being a no-op (#17709 ) * fix Muon weight decay being a no-op LARS._step computed the post-momentum weight decayed param but only used it for a dtype cast, so the decay was never applied. Fold the decay into the update instead, matching torch's param.mul_(1 - lrwd). test_muon_wd passed anyway since lrwd=1e-5 is far below atol, so also bump the test's weight_decay to 10 to actually exercise it. * muon: apply weight decay after lr scaling keeps the decoupled decay independent of the LARS trust ratio r, matching torch's param.mul_(1 - lr*wd). no behavior change today since r is always 1.0 on the pre_wd=False path (Muon has tcoef=0).	2026-08-24 13:33:10 -07:00
nimlgenandGitHub	7a887e84b2	am: use gc regs for xgmi (#17707 )	2026-08-24 22:42:35 +03:00
chenyuandGitHub	7b21ffac00	PYLITERAL has dtypes.void (#17705 ) also deleted Ops.WIAT, it currently does not have entry in dtype_from_uop	2026-08-24 14:37:15 -04:00
sirhcmandGitHub	9b481a4893	dsp: cleanup temp linker script (#17704 )	2026-08-24 14:05:00 -04:00
qazalandGitHub	9b0c65688d	llama script changes from the speedups branch (#17701 )	2026-08-24 21:10:02 +09:00
George HotzandGitHub	7e561fcb97	mergable fast RDNA3 Qwen 3.8 (#17512 ) * mergable fast RDNA3 Qwen 3.6 * AMD * quant 256 multiple * cleanup cast * llm kernels: adapt to Ops.BIND removal Variables are 0-d ALU BUFFERs in the tensor graph and take the ALU PARAM form inside kernels (UOp.variable(param=True)). Add kernel_var helper for the conversion, and keep start_pos in bound form at the graph level so function implicit-input collection and the schedule's binds rename-back line up. * adaptive prefill chunk sizes for recurrent models + iq4xs model entry one TinyJit per static prefill chunk size: capture 128 and 32 at warmup, generate picks the largest that fits the remaining prompt. long prompts prefill 2x faster (555 tok/s on Qwen3.6-27B IQ4_XS) without pushing short prompts through token-by-token decode. * minimize diff: early-return custom attention path, keep master state init * minimize: master _attention with gated fused-scan swap, kernels/amd only, single chunk size - GatedDeltaNetBlock._attention keeps master's symbolic-padding structure; the recurrent scan is swapped for the fused gated_delta_prefill kernel only on RDNA3 with static shapes (fast_scan), everything else uses the old path - all AMD kernel code lives in tinygrad/llm/kernels/amd.py (drop kernels/__init__.py, drop the generic fallback kernel - the old scan covers non-RDNA3) - single prefill chunk size 32; non-RDNA3 recurrent keeps master's chunk_size=1 - the conv+normalize miscompile doesn't trigger with master's window-buffer conv, so the contiguous workaround is dropped * warmup: single code path for fast and old recurrent * drop fast_scan/fast_recurrent flags, inline the RDNA3 gate (cached) * generate: chunk size is always 32, no device gating static chunks for recurrent models everywhere: the fused kernel path on RDNA3, the old scan elsewhere (which is also faster chunked than token-by-token) * warmup: drop redundant _init_state loop (lazy init in the eager step covers it) * symbolic-length prefill with the custom kernels the prefill path is fully symbolic again (master's generate, one prefill graph for every chunk size, no static-tail decode): padded steps are exact no-ops in the scan (beta=0, alpha=exp(0)=1), flash attention positions queries at start_pos instead of valid_kv_len-M, and quant linears pad to the chunk bucket prefill 401 tok: 284 -> 348 tok/s on Qwen3.8-27B IQ4_XS (tail chunks no longer decode token-by-token), decode unchanged at 45 tok/s * cli: default qwen3.6:27b to the fast IQ4_XS quant, add qwen3.8:27b Q4_K_M falls back to slow inline dequant with the custom kernels, IQ4_XS is the fast path. qwen3.8 quants use unsloth's UD (dynamic) naming * warmup: back to master's two-liner plus a cache reset with symbolic prefill, generate([0])'s 1-token chunk captures the symbolic prefill graph that serves every chunk size, and JIT batching on capture measurably doesn't matter with the fused kernels (347.7 tok/s either way) * cli: pin qwen3.8:27b to the pre-UD revision the UD-IQ4_XS replacement mixes in Q3_K tensors (ggml type 11) the loader doesn't support; the pinned revision is byte-identical to the known-good file * warmup: identical to master the leftover cache is self-consistent: get_start_pos only reuses a full strict-prefix match, everything else restarts with a state reset * model: hoist the quantized_attention import to the top level * hoist the GDN query scale out of the branch, restore master dtype.py the scale is the same op in both paths, apply it once after the transpose. the dtype.py diff was a stale pre-SPEC=2 copy, not intentional work * gated_delta_prefill: don't pass the bound start_pos as a call src device-less param buffers in call srcs crash hcq2's _get_enqueue_devs. the var already reaches the graph through the state AFTER chain (conv state store), same as the flash kernels' valid_end * llm: half KV cache with custom flash kernels, drop the int8 quantized cache matches master's new half cache default: no scales, no packing, one less buffer. the store casts to half explicitly (buffer-only half usage misses the renderer's half define). 45.5 tok/s decode, 348.7 tok/s prefill — same as int8 * llm: zero-init the KV cache the int8 path was accidentally protected from uninitialized memory by its zero-initialized scale buffer; with a plain half cache the flash prefill kernel's PV wmma computes 0NaN=NaN on masked lanes past the valid region (manifested as garbage tokens at 32k context where the allocator reuses dirty VRAM) * gate that * llm/kernels/amd: reorganize by kernel family, drop the clutter sections: shared helpers, quant linear, flash attention, gated delta prefill. no AxisType.WEAK (default), no ALLOW_DEVICE_USAGE override (unneeded), magic numbers become names (QUANT_SIZES, Q5_K/Q6_K/IQ4_XS), merged wrapper layers (flash_attention_causal_cached folded into flash_attention), one _unbind helper for the bound-var dance * test: universal recurrent reuse assertion, fix lambda lint * 1-token chunks have a static shape: they are decode steps a 1-token chunk routes to the decode graph via the existing dispatch, so warmup and decode-only workloads never build the big symbolic prefill graph: CI benchmark command 12m50s -> 5m29s (master: 6m48s), 220 -> 123 compile jobs also restores the ALLOW_DEVICE_USAGE override in amd_custom_kernels_supported: Device[] asserts inside @function contexts (ALLOW_DEVICE_USAGE=0), and the first gate call can happen there depending on test order * generate: back to plain symbolic binding, the static-1 rule wasn't worth it * custom kernels: Q4_K support (ggml type 12) Q4_K is Q5_K without the high-bit array: same d/dmin/scales layout (so _q5_scales works unchanged), 144-byte blocks, qs at word 4. both the dp4a decode kernel and the WMMA prefill kernel take a ggml_type branch now. Qwen3-8B Q4_K_M: decode 16.5 -> 114.8 tok/s, prefill 69 -> 536 tok/s * raise line count to 26500 (qwen did it) * benchmark qwen3.8 * little updates	2026-08-23 22:46:35 -07:00
YassineYousfiandGitHub	a5ea95d8b8	usb: don't fail async transfers on signal interruption (#17692 )	2026-08-23 21:30:00 -07:00
geohot	df528499ce	Revert "disk cache: thread-local db conn (#17694 )" This reverts commit `11edcc144c`.	2026-08-23 21:13:29 -07:00
qazalandGitHub	4ee114b42a	llama rope freqs in fp32 (#17698 )	2026-08-24 12:49:55 +09:00
geohot	9d0cd0ebcb	hotfix: switch benchmark to qwen3.8	2026-08-23 20:34:30 -07:00
YassineYousfiandGitHub	11edcc144c	disk cache: thread-local db conn (#17694 )	2026-08-23 19:18:47 -07:00
chenyuandGitHub	a31aca9e52	don't hardcode dtype.int const in decode_hevc_frame (#17696 )	2026-08-23 22:00:41 -04:00
chenyuandGitHub	8b164aefea	test update from weak const (#17689 )	2026-08-23 17:30:16 -04:00
George HotzandGitHub	477b573807	update llm kv cache to be half (#17690 ) * update llm kv cache to be half / chunk_size to always be 32 * just dtype	2026-08-23 14:02:54 -07:00
geohot	bb0e99acbf	hotfix: skip that nan test on mac	2026-08-23 08:40:46 -07:00
nimlgenandGitHub	93865e2c66	hcq2: sunday housekeeping (#17686 ) * hcq2: sunday housekeeping * x	2026-08-23 17:20:40 +03:00
George HotzandGitHub	5b60a09ab0	some fixes for the AMD emulator (#17684 ) * some fixes for the AMD emulator * simpler * revert * min	2026-08-22 22:48:07 -07:00
chenyuandGitHub	b0a1285330	lil decomp cleanup [PR] (#17682 )	2026-08-22 18:11:14 -04:00
nimlgenandGitHub	a2e64e16aa	hcq2: early usb (#17683 ) * hcq2: usb interface and submit * x * x * x * x * r * x	2026-08-23 00:22:39 +03:00
chenyuandGitHub	a9069c177a	make decomp pass SPEC=2 [PR] (#17681 )	2026-08-22 08:01:15 -04:00
chenyuandGitHub	8950942e75	remove explicit dtype for NOOP and decomp [PR] (#17678 )	2026-08-21 22:27:36 -04:00
chenyuandGitHub	356f665377	test update for weak const (#17675 )	2026-08-21 21:50:29 -04:00
chenyuandGitHub	7204d46786	delete dtype_from_uop INDEX exempt (#17674 )	2026-08-21 21:08:06 -04:00
George HotzandGitHub	af242819d8	refactor the AMD emulator slop (kimi) (#17673 ) * refactor the AMD emulator slop (kimi) * mypy	2026-08-21 18:00:58 -07:00
wozeparrotandGitHub	52596dbf38	gptoss: fused ce (#17672 )	2026-08-21 16:19:28 -07:00
sirhcmandGitHub	07cce78cec	compile3: log printed timings (#17671 )	2026-08-21 19:18:32 -04:00
George HotzandGitHub	f986829461	keep IndexingContext scoped in indexing (#17670 )	2026-08-21 14:12:39 -07:00
nimlgenandGitHub	4fd4eafb23	nv: hevc (#17661 ) * nv: hevc * x * nv: zero the nvdec scratch buffers	2026-08-21 23:37:09 +03:00
chenyuandGitHub	daa154aa22	FLOORDIV to SHR for powers of 2 [pr] (#17669 )	2026-08-21 16:28:04 -04:00
nimlgenandGitHub	12f889aaad	hcq2: parallel compile (#17667 ) * hcq2: parallel compile * Dx	2026-08-21 23:27:21 +03:00
sirhcmandGitHub	298748ebd3	ci: remove setup-python (#17665 )	2026-08-21 16:19:47 -04:00
George HotzandGitHub	3082956a17	usb copyin: async arm and drain, 323 MB/s on comma (#17663 ) * usb copyin: can safely arm before drain * perf counter * 294 MB/s for comma * free speed with async transfers	2026-08-21 13:12:03 -07:00
chenyuandGitHub	402bea7ddd	l2i and sign_extend cleanups [pr] (#17668 ) towards good threefry decomp	2026-08-21 15:55:45 -04:00
chenyuandGitHub	8f9cbdf0cc	few more self folding [pr] (#17657 )	2026-08-21 15:21:28 -04:00
chenyuandGitHub	8f59041ee5	more wgsl pack cleanups [PR] (#17664 )	2026-08-21 14:34:25 -04:00
nimlgenandGitHub	1cf8a2c7fe	hcq2: use shrink.bitcast (#17653 ) * hcq2: shrink.bitcast * x * x * x * s * x * Dx * Revert "hotfix: disable HCQ2" This reverts commit `a57188ea6d`. * x	2026-08-21 21:24:34 +03:00