kill_pids --sudoless

hcq_smi rmmod --expect and skip setup_python_cap
benchmarks: remove some usage of sudo
2026-08-25 10:43:05 -07:00 · 2026-08-25 10:36:24 -07:00 · 2026-08-25 10:05:30 -07:00 · 2026-08-25 19:03:19 +03:00 · 2026-08-25 11:46:31 -04:00 · 2026-08-25 18:23:22 +03:00
325 changed files with 14737 additions and 7045 deletions
@@ -4,7 +4,7 @@ inputs:
  python-version:
    description: 'Python version to use'
    required: false
-    default: '' # if you don't set a version, the native python version will be used
+    default: '3.14'
  key:
    description: 'Key for the python cache'
    required: false
@@ -41,12 +41,12 @@ inputs:
    description: "Install LLVM?"
    required: false
    default: 'false'
-  tinydreno:
-    description: "Install tinydreno"
+  qemu:
+    description: "Install qemu?"
    required: false
    default: 'false'
-  qemu:
-    description: "Install qemu"
+  ninja:
+    description: "Install ninja?"
    required: false
    default: 'false'
 runs:
@@ -59,18 +59,18 @@ runs:
        echo "OMP_NUM_THREADS=1" >> "$GITHUB_ENV"
        # no buffers should be over 300MB in CI
        echo "MAX_BUFFER_SIZE=300000000" >> "$GITHUB_ENV"
+        if [[ "$RUNNER_OS" == "Linux" ]]; then
+          echo "VIRTUAL_ENV=/opt/venv/${{ inputs.python-version }}" >> "$GITHUB_ENV"
+          echo "UV_PYTHON_INSTALL_DIR=/opt/python" >> "$GITHUB_ENV"
+        else
+          echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
+        fi

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

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

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

    # ******************* apt *******************
    - name: Setup apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
      shell: bash
      run: |
+        sudo mkdir -p /var/cache/apt/archives
        sudo chown -R $USER:$USER /var/cache/apt/archives

        echo 'Acquire::GzipIndexes "true";' | sudo tee /etc/apt/apt.conf.d/gzip
@@ -161,7 +161,7 @@ runs:
        echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-20 main" | sudo tee /etc/apt/sources.list.d/llvm.list

    - name: Compute Package List + Hash
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
      id: apt-pkgs
      shell: bash
      run: |
@@ -186,25 +186,29 @@ runs:
        if [[ "${{ inputs.qemu }}" == "true" ]]; then
          pkgs+=" qemu-user-static"
        fi
+        # **** ninja ****
+        if [[ "${{ inputs.ninja }}" == "true" ]]; then
+          pkgs+=" ninja-build"
+        fi

        echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
        echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"

    - name: Cache apt (PR)
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true') && github.event_name == 'pull_request'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name == 'pull_request'
      uses: actions/cache/restore@v5
      with:
        path: /var/cache/apt/archives/
        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
    - name: Cache apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true') && github.event_name != 'pull_request'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name != 'pull_request'
      uses: actions/cache@v5
      with:
        path: /var/cache/apt/archives/
        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}

    - name: Run apt Update + Install
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
      shell: bash
      run: |
        sudo apt -qq update || true
@@ -214,6 +218,7 @@ runs:
          sudo apt-get -y --allow-unauthenticated --no-install-recommends install ${{ steps.apt-pkgs.outputs.pkgs }}
        fi

+        sudo mkdir -p /var/cache/apt/archives
        sudo chown -R $USER:$USER /var/cache/apt/archives/

    - name: Add clang to PATH (Linux)
@@ -275,12 +280,6 @@ runs:
      shell: bash
      run: brew install llvm@20

-    # *** tinydreno ***
-    - name: Install tinydreno (linux)
-      if: inputs.tinydreno == 'true' && runner.os == 'Linux'
-      shell: bash
-      run: sudo curl -fL https://github.com/sirhcm/tinydreno/raw/refs/heads/master/libllvm-qcom.so -o /usr/lib/libllvm-qcom.so
-
    # *** OpenCL ***
    - name: Install rusticl
      if: inputs.opencl == 'true'
@@ -35,7 +35,7 @@ jobs:
        key: 'autogen'
        amd: 'true'
        llvm: 'true'
-        pydeps: 'pyyaml mako'
+        deps: 'autogen'
    - name: Install autogen support packages
      run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev liburing-dev
    - name: Regenerate autogen files
@@ -88,12 +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: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -101,13 +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
-    - name: Symlink models and datasets
-      run: |
-        mkdir -p weights
-        ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
+        ./extra/hcq/hcq_smi.py amd rmmod --expect
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+    - name: Setup (NV)
+      if: ${{ matrix.dev == 'NV' }}
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -117,18 +116,14 @@ 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.5
-      # qwen3.5: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=qwen35_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.5: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' }}
      run: BENCHMARK_LOG=olmoe JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m olmoe --benchmark --warmup
-    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
-      # only run on machines with multiple gpus
-      if: ${{ matrix.dev != 'METAL' }}
-      run: BENCHMARK_LOG=llama3_beam_4gpu JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -139,12 +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: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -152,9 +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: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -174,10 +172,6 @@ jobs:
      # slow on metal
      if: ${{ matrix.dev != 'METAL' }}
      run: time BENCHMARK_LOG=cifar DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
-    - name: Run full CIFAR training steps w 6 GPUS
-      # only run on machines with multiple gpus
-      if: ${{ matrix.dev != 'METAL' }}
-      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -188,12 +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: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -201,9 +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: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
    - name: Symlink models and datasets
      run: |
        mkdir -p extra/datasets
@@ -215,15 +212,8 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: test/external/process_replay/reset.py
-    - name: Run MLPerf resnet eval on training data
-      run: time BENCHMARK_LOG=resnet_eval MODEL=resnet python3 examples/mlperf/model_eval.py
    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
      run: BENCHMARK_LOG=resnet_10steps DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
-    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: BENCHMARK_LOG=resnet_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
-    - name: Run 10 MLPerf Bert training steps (6 gpu)
-      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -234,12 +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: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -247,9 +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: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -269,6 +262,59 @@ jobs:
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

+  multigpubenchmark:
+    name: Multi-GPU Benchmarks (DEV=${{ matrix.dev }})
+    runs-on: [self-hosted, "${{ matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
+    strategy:
+      fail-fast: false
+      matrix:
+        dev: ['AMD', 'NV']
+    timeout-minutes: 20
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    env:
+      DEV: ${{ matrix.dev }}
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup (AMD)
+      if: ${{ matrix.dev == 'AMD' }}
+      run: |
+        ./extra/hcq/hcq_smi.py amd rmmod --expect
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
+    - name: Setup (NV)
+      if: ${{ matrix.dev == 'NV' }}
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        mkdir -p extra/datasets
+        ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
+        ln -s /raid/datasets/imagenet extra/datasets/imagenet
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: reset process replay
+      run: python3 test/external/process_replay/reset.py
+    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
+      run: BENCHMARK_LOG=llama3_beam_4gpu JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
+    - name: Run full CIFAR training steps w 6 GPUS
+      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    - name: Run MLPerf resnet eval on training data
+      run: time BENCHMARK_LOG=resnet_eval MODEL=resnet python3 examples/mlperf/model_eval.py
+    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
+      run: BENCHMARK_LOG=resnet_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
+    - name: Run 10 MLPerf Bert training steps (6 gpu)
+      # TODO: remove BERT_LAYERS once scheduler is fast
+      run: BENCHMARK_LOG=bert_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
  tests:
    name: Tests (DEV=${{ matrix.dev }})
    runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
@@ -276,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}
@@ -289,9 +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: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -368,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}
@@ -391,7 +439,7 @@ jobs:
    - name: UsbGPU tiny tests
      run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
    - name: UsbGPU copy speeds
-      run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+      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
    - name: UsbGPU (USB4/TB) install script
@@ -402,6 +450,35 @@ jobs:
      run: PYTHONPATH=. DEV=PCI+NV:NAK python3.11 test/test_tiny.py

  testcommalatest:
+    name: comma Benchmark (0.11.2)
+    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.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
@@ -432,35 +509,6 @@ jobs:
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

-  testcommaold:
-    name: comma Benchmark (0.10.1)
-    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: DEBUG=2 openpilot compile3 0.10.1 driving_vision
-      run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot compile3 0.10.1 driving_vision
-      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot compile3 0.10.1 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_1_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/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
-    - name: openpilot compile3 0.10.1 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
  testqualcommdsp:
    name: DSP Benchmark
    runs-on: [self-hosted, Linux, comma4]
@@ -470,15 +518,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
@@ -502,7 +541,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}
@@ -515,14 +554,14 @@ jobs:
      run: |
        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: openpilot compile3 0.10.1 driving_vision
-      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." GMMU=0 DEV=USB+AMD:LLVM ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot load_pickle 0.10.1 driving_vision
-      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
-    - name: openpilot run_pickle 0.10.1 driving_vision
-      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py
+    - name: openpilot compile3 big_driving_supercombo
+      run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo PICKLE_OOB=1 PYTHONPATH="." TC_OPT=2 GMMU=0 DEV=USB+AMD:LLVM ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/10926f2c0911821ca0e72439c1c3bf3ec11f0a08789aa14b7ee8f25379b2afa4 openpilot.pkl
+    - name: openpilot load_pickle big_driving_supercombo
+      run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_load_pickle PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_LOAD_TIME=25 python3 examples/openpilot/load_pickle.py openpilot.pkl
+    - name: openpilot run_pickle big_driving_supercombo
+      run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_run_pickle RUN_PICKLE=1 PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py - openpilot.pkl
    - name: Test copy speeds
-      run: SIZE=64e6 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds

  driverbenchmarks:
    name: PCI Driver Benchmark (DEV=${{ matrix.dev }})
@@ -531,7 +570,7 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['AMD', 'NV']
-    timeout-minutes: 20
+    timeout-minutes: 5
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -543,9 +582,8 @@ jobs:
      uses: actions/checkout@v6
    - name: Setup
      run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py ${{ matrix.dev == 'AMD' && 'amd' || 'nv' }} rmmod
-        ./extra/hcq/hcq_smi.py ${{ matrix.dev == 'AMD' && 'amd' || 'nv' }} kill_pids
+        ./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
@@ -580,6 +618,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
@@ -602,7 +643,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
@@ -8,7 +8,7 @@ permissions:
  contents: write
 jobs:
  deploy:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
    steps:
      - uses: actions/checkout@v6
      - name: Configure Git Credentials
@@ -0,0 +1,181 @@
+name: Platform Tests
+env:
+  # increment this when downloads substantially change to avoid the internet
+  CACHE_VERSION: '19'
+  CAPTURE_PROCESS_REPLAY: ${{ github.event_name == 'pull_request' && contains(github.event.pull_request.title, '[pr]') && '1' || '0' }}
+  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+  PYTHONPATH: ${{ github.workspace }}
+  CHECK_OOB: 1
+
+on:
+  push:
+    branches:
+      - master
+  pull_request:
+  workflow_dispatch:
+
+concurrency:
+  group: platform-${{ github.event_name }}-${{ github.event_name == 'pull_request' && github.event.pull_request.number || github.run_id }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
+jobs:
+
+# ****** OSX Tests ******
+
+  unittestmacos:
+    name: MacOS (unit)
+    runs-on: macos-26
+    timeout-minutes: 20
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: unittest-macos
+        deps: testing_unit
+    - name: Run unit tests
+      run: DEV=METAL python -m pytest -n=auto test/unit/ --durations=20
+    - name: Test tensor core ops (fake)
+      run: DEV=METAL DEBUG=3 TC=2 python test/backend/test_ops.py TestOps.test_gemm
+    - name: Test tensor core ops (real)
+      run: DEV=METAL DEBUG=3 python test/backend/test_ops.py TestOps.test_big_gemm
+    - name: Test Beam Search
+      run: DEV=METAL IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
+    - name: Test Device Specific
+      run: DEV=METAL python3 -m pytest test/device/test_metal.py
+    #- name: Fuzz Test linearizer
+    #  run: DEV=METAL DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
+  unittestmacosmock:
+    name: MacOS (unit, mock)
+    runs-on: macos-26
+    timeout-minutes: 20
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: unittest-macos-mock
+        deps: testing_unit
+        amd: 'true'
+        ocelot: 'true'
+    - name: Run NULL backend tests
+      run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --durations=20
+    - name: Run pytest (amd)
+      env:
+        DEV: MOCKKFD+AMD
+        FORWARD_ONLY: 1
+      run: |
+        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
+    - name: Run pytest (ptx)
+      env:
+        DEV: "MOCK+NV:PTX"
+        FORWARD_ONLY: 1
+        # TODO: failing due to library loading error
+        CAPTURE_PROCESS_REPLAY: 0
+      run: |
+        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py \
+          test/testextra/test_hevc.py::TestHevc::test_hevc_decode_compile --durations=20
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
+  testmetal:
+    strategy:
+      fail-fast: false
+      matrix:
+        group: [1, 2]
+    name: MacOS (DEV=METAL) (${{ matrix.group }})
+    runs-on: macos-26
+    timeout-minutes: 20
+    env:
+      DEV: METAL
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: macos-metal
+        deps: testing_unit
+    - name: Check Device.DEFAULT and print some source
+      run: |
+        python -c "from tinygrad import Device; assert Device.DEFAULT == 'METAL'"
+        DEBUG=4 python test/test_tiny.py TestTiny.test_plus
+    - name: Run backend tests
+      run: python -m pytest -n=auto test/backend --durations=20 --splits 2 --group ${{ matrix.group }}
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
+  testmacos:
+    strategy:
+      fail-fast: false
+      matrix:
+        dev:
+          - 'CPU:CLANG'
+          - 'CPU:LLVM'
+          - 'CPU:LVP'
+          - 'WEBGPU'
+
+    name: MacOS (DEV=${{ matrix.dev }})
+    runs-on: macos-26
+    timeout-minutes: 20
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: macos-${{ matrix.dev }}
+        deps: "testing_unit${{ contains(matrix.dev, 'LVP') && ' mesa' || '' }}"
+        llvm: ${{ contains(matrix.dev, 'LLVM') || contains(matrix.dev, 'LVP') }}
+        webgpu: ${{ matrix.dev == 'WEBGPU' }}
+    - name: Set env
+      run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
+    - name: Check Device.DEFAULT and print some source
+      run: |
+        python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
+        DEBUG=4 python test/test_tiny.py TestTiny.test_plus
+    - name: Run test_tiny
+      run: python -m pytest -n=auto test/test_tiny.py --durations=20
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
+# ****** Windows Tests ******
+
+  testwindows:
+    strategy:
+      fail-fast: false
+      matrix:
+        dev:
+          - 'CPU:CLANG'
+          - 'CPU:LLVM'
+          - 'CPU:X86'
+          - 'WEBGPU'
+
+    name: Windows (DEV=${{ matrix.dev }})
+    runs-on: windows-2025
+    timeout-minutes: 15
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v6
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: windows-${{ matrix.dev }}-minimal
+          deps: testing_minimal
+          pydeps: ${{ matrix.dev == 'WEBGPU' && 'dawn-python' || '' }}
+      - name: Set env
+        shell: bash
+        run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
+      - name: Check Device.DEFAULT and print some source
+        shell: bash
+        run: |
+          python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
+          DEBUG=4 python test/test_tiny.py TestTiny.test_plus
+      - name: Run test_tiny
+        shell: bash
+        run: python -m pytest -n=auto test/test_tiny.py --durations=20
@@ -10,7 +10,7 @@ on:

 jobs:
  deploy:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
    steps:
    - uses: actions/checkout@v6
    - name: Set up Python
@@ -10,7 +10,7 @@ concurrency:
 jobs:
  checkbranch:
    name: Check PR Branch status
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
    outputs:
      branchstat: ${{ steps.brstat.outputs.stat}}
    steps:
@@ -44,7 +44,7 @@ jobs:
    permissions:
      contents: read
      pull-requests: write
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
    needs: checkbranch
    if: needs.checkbranch.outputs.branchstat == 'false'
    steps:
@@ -87,7 +87,7 @@ jobs:
    name: Core Library Line Difference
    permissions:
      pull-requests: write
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
    needs: checkbranch
    if: needs.checkbranch.outputs.branchstat == 'true'
    steps:
@@ -21,7 +21,7 @@ concurrency:
 jobs:
  docs:
    name: Docs
-    runs-on: &linux ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 10
    env:
      CHECK_OOB: 0
@@ -31,8 +31,7 @@ jobs:
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
-        deps: docs
-        pydeps: "capstone torch"
+        deps: "docs testing_minimal"
    - name: Build wheel and show size
      run: |
        uv build --wheel
@@ -61,7 +60,7 @@ jobs:

  torchbackend:
    name: Torch Backend Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -73,10 +72,7 @@ jobs:
        deps: testing_unit
        pydeps: "pillow torchvision expecttest"
        llvm: 'true'
-    - name: Install ninja
-      run: |
-        sudo apt update || true
-        sudo apt install -y --no-install-recommends ninja-build
+        ninja: 'true'
    - name: Test ResNet-18
      run: DEBUG=2 python3 extra/torch_backend/example.py
    - name: Test one op in torch tests
@@ -86,9 +82,26 @@ jobs:
    - name: Custom tests
      run: DEV=CPU:LLVM GPUS=4 TINY_BACKEND=1 python3 -m pytest -nauto extra/torch_backend/test.py extra/torch_backend/test_inplace.py extra/torch_backend/test_multigpu.py extra/torch_backend/test_kernel_fusion.py --durations=20

+  torchbackendtrain:
+    name: Torch Backend Training
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    timeout-minutes: 15
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: torch-backend-pillow-torchvision-et-pt
+        deps: testing_unit
+        llvm: 'true'
+        ninja: 'true'
+    - name: Test beautiful_mnist in torch with TINY_BACKEND
+      run: STEPS=20 DEV=CPU TARGET_EVAL_ACC_PCT=90.0 MAX_BUFFER_SIZE=0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+
  bepython:
    name: Python Backend
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -126,7 +139,7 @@ jobs:

  linter:
    name: Linters
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 10

    steps:
@@ -157,7 +170,7 @@ jobs:

  nulltest:
    name: Null Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15

    steps:
@@ -167,6 +180,7 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: unittest-13
+        python-version: '3.11'
        deps: testing_unit
        llvm: 'true'
        amd: 'true'
@@ -176,20 +190,21 @@ jobs:
      run: |
        DEV=NULL python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
        DEV=NULL VIZ=1 python3 -m pytest -n=auto test/null/test_viz.py
+        DEBUG=7 python -m tinygrad.viz.cli --json | jq empty
    # TODO: too slow
    # - name: Run SDXL on NULL backend
    #   run: DEV=NULL DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
    - name: Run Clip tests for SD MLPerf on NULL backend
      run: DEV=NULL python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
    - name: Run AMD emulated BERT training on NULL backend
-      run: DEV=NULL::gfx1201 NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+      run: DEV=NULL::gfx1201 NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    # TODO: support fake weights
    #- name: Run LLaMA 7B on 4 fake devices
    #  run: DEV=NULL python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing

  unittest:
    name: Unit Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15

    steps:
@@ -199,6 +214,7 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: unittest-13
+        python-version: '3.11'
        pydeps: "pre-commit"
        deps: testing_unit
        llvm: 'true'
@@ -216,8 +232,8 @@ jobs:
      run: python3 test/external/external_benchmark_schedule.py
    - name: Run process replay tests
      uses: ./.github/actions/process-replay
-    - name: Repo line count < 25000 lines
-      run: MAX_LINE_COUNT=25000 python sz.py
+    - name: Repo line count <= 26000 lines
+      run: MAX_LINE_COUNT=26500 python sz.py

  spec:
    strategy:
@@ -225,7 +241,7 @@ jobs:
      matrix:
        group: [1, 2]
    name: SPEC=2 (${{ matrix.group }})
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -241,7 +257,7 @@ jobs:

  fuzzing:
    name: Fuzzing
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 10
    steps:
    - name: Checkout Code
@@ -257,7 +273,7 @@ jobs:

  testopenclimage:
    name: CL IMAGE Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -277,7 +293,7 @@ jobs:

  testopenpilot:
    name: openpilot Compile Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -291,7 +307,7 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1361 ALLOWED_GATED_READ_IMAGE=54 FLOAT16=1 DEV="CL::IMAGE_PITCH_ALIGNMENT=64" IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1361 ALLOWED_GATED_READ_IMAGE=38 FLOAT16=1 DEV="CL::IMAGE_PITCH_ALIGNMENT=64" IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
          # IMAGE_PITCH_ALIGNMENT=64 matches adreno 630
      - name: Test openpilot CL compile fp32 (test correctness)
        run: |
@@ -306,7 +322,7 @@ jobs:

  testonnxcpu:
    name: ONNX (CPU) Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 20

    steps:
@@ -325,7 +341,7 @@ jobs:

  testoptim:
    name: Optimization Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 20
    steps:
      - name: Checkout Code
@@ -350,12 +366,14 @@ jobs:
        run: DEV=NULL NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
      - name: Test llama 3 training
        run: DEV=NULL NULL_ALLOW_COPYOUT=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
+      - name: Test gpt-oss training
+        run: DEV=NULL NULL_ALLOW_COPYOUT=1 SAMPLES=32 BS=2 SEQLEN=128 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 MXFP8=1 VOCAB_SIZE=32000 LAYERS=2 EXPERTS=4 MODEL=gptoss PYTHONPATH=. python3 examples/mlperf/model_train.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

  testllm:
    name: Test LLM
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
    env:
      CHECK_OOB: 0
@@ -382,7 +400,7 @@ jobs:

  testmodels:
    name: Models
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -402,7 +420,7 @@ jobs:

  testdsp:
    name: Linux (DSP)
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -430,7 +448,7 @@ jobs:
          - 'WEBGPU'

    name: Linux (DEV=${{ matrix.dev }})
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 20
    steps:
    - name: Checkout Code
@@ -456,7 +474,7 @@ jobs:

  testamdasm:
    name: AMD ASM IDE
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 20
    env:
      DEV: MOCKKFD+AMD
@@ -478,7 +496,7 @@ jobs:
          wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
          echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-21 main" | sudo tee /etc/apt/sources.list.d/llvm.list
          sudo apt-get update
-          sudo apt-get install llvm-21 llvm-21-tools cloc
+          sudo apt-get install -y llvm-21 llvm-21-tools cloc
      - name: Install rocprof-trace-decoder
        run: sudo PYTHONPATH="." ./extra/sqtt/install_rocprof_decoder.py
      - name: Run AMD renderer tests
@@ -486,7 +504,7 @@ jobs:
      - name: Run AMD renderer tests (AMD:LLVM)
        run: DEV=MOCKKFD+AMD:LLVM python -m pytest -n=auto test/amd/ --durations 20
      - name: Run SQTT profiling tests
-        run: PROFILE=1 SQTT=1 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
+        run: VIZ=-2 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
      - name: Run AMD emulated tests on NULL backend
        env:
          AMD: 0
@@ -502,7 +520,7 @@ jobs:

  hcq2:
    name: hcq2
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 5
    steps:
      - name: Checkout Code
@@ -516,16 +534,15 @@ jobs:
      - name: Run HCQ2 tests
        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/test_tiny.py
      - name: Run HCQ2 multi-device tests
-        run: |
-          HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_multitensor.py \
-            TestMultiTensor.test_simple_add TestMultiTensor.test_shard_reduce \
-            TestMultiTensor.test_backward_sum TestMultiTensor.test_matmul_shard_0_0
+        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest -n=auto test/backend/test_multitensor.py
      - name: Run HCQ2 JIT tests
        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_jit.py
+      - name: Run HCQ2 unit tests
+        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest test/device/test_hcq2.py

  testmockam:
    name: Linux (am)
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
    env:
      DEV: MOCKPCI+AMD
@@ -561,7 +578,7 @@ jobs:
        arch: [gfx1100, gfx1201, gfx950]

    name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
    env:
      DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
@@ -580,6 +597,9 @@ jobs:
        run: |
          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['AMD'], Device.DEFAULT"
          DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
+      - name: Run MXFP4 Llama training on NULL backend
+        if: ${{ matrix.backend == 'amd' && matrix.arch == 'gfx950' }}
+        run: PYTHONPATH=. DEV=NULL:HIP:gfx950 MXFP4=1 LLAMA_LAYERS=2 BENCHMARK=3 NULL_ALLOW_COPYOUT=1 NO_HIPCC=1 ROCM_PATH=/opt/rocm JITBEAM=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh
      - name: Run pytest (amd)
        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM --durations=20
      - name: Run disk copy tests
@@ -596,7 +616,7 @@ jobs:
        backend: [ptx, nv]

    name: Linux (${{ matrix.backend }})
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 20
    env:
      FORWARD_ONLY: 1
@@ -624,160 +644,23 @@ jobs:
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

-# ****** OSX Tests ******
-
-  unittestmacos:
-    name: MacOS (unit)
-    runs-on: &macos macos-26
-    timeout-minutes: 20
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: unittest-macos
-        deps: testing_unit
-        amd: 'true'
-        ocelot: 'true'
-    - name: Run unit tests
-      run: DEV=METAL python -m pytest -n=auto test/unit/ --durations=20
-    - name: Run NULL backend tests
-      run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --durations=20
-    - name: Test tensor core ops (fake)
-      run: DEV=METAL DEBUG=3 TC=2 python test/backend/test_ops.py TestOps.test_gemm
-    - name: Test tensor core ops (real)
-      run: DEV=METAL DEBUG=3 python test/backend/test_ops.py TestOps.test_big_gemm
-    - name: Test Beam Search
-      run: DEV=METAL IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
-    - name: Test Device Specific
-      run: DEV=METAL python3 -m pytest test/device/test_metal.py
-    #- name: Fuzz Test linearizer
-    #  run: DEV=METAL DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
-    - name: Run pytest (amd)
-      env:
-        DEV: MOCKKFD+AMD
-        FORWARD_ONLY: 1
-      run: |
-        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
-    - name: Run pytest (ptx)
-      env:
-        DEV: "MOCK+NV:PTX"
-        FORWARD_ONLY: 1
-        # TODO: failing due to library loading error
-        CAPTURE_PROCESS_REPLAY: 0
-      run: |
-        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
-  testmetal:
-    strategy:
-      fail-fast: false
-      matrix:
-        group: [1, 2]
-    name: MacOS (DEV=METAL) (${{ matrix.group }})
-    runs-on: *macos
-    timeout-minutes: 20
-    env:
-      DEV: METAL
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: macos-metal
-        deps: testing_unit
-    - name: Check Device.DEFAULT and print some source
-      run: |
-        python -c "from tinygrad import Device; assert Device.DEFAULT == 'METAL'"
-        DEBUG=4 python test/test_tiny.py TestTiny.test_plus
-    - name: Run backend tests
-      run: python -m pytest -n=auto test/backend --durations=20 --splits 2 --group ${{ matrix.group }}
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
-  testmacos:
-    strategy:
-      fail-fast: false
-      matrix:
-        dev:
-          - 'CPU:CLANG'
-          - 'CPU:LLVM'
-          - 'CPU:LVP'
-          - 'WEBGPU'
-
-    name: MacOS (DEV=${{ matrix.dev }})
-    runs-on: *macos
-    timeout-minutes: 20
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: macos-${{ matrix.dev }}
-        deps: "testing_unit${{ contains(matrix.dev, 'LVP') && ' mesa' || '' }}"
-        llvm: ${{ contains(matrix.dev, 'LLVM') || contains(matrix.dev, 'LVP') }}
-        webgpu: ${{ matrix.dev == 'WEBGPU' }}
-    - name: Set env
-      run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
-    - name: Check Device.DEFAULT and print some source
-      run: |
-        python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
-        DEBUG=4 python test/test_tiny.py TestTiny.test_plus
-    - name: Run test_tiny
-      run: python -m pytest -n=auto test/test_tiny.py --durations=20
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
-# ****** Windows Tests ******
-
-  testwindows:
-    strategy:
-      fail-fast: false
-      matrix:
-        dev:
-          - 'CPU:CLANG'
-          - 'CPU:LLVM'
-          - 'CPU:X86'
-          - 'WEBGPU'
-
-    name: Windows (DEV=${{ matrix.dev }})
-    runs-on: windows-2025
-    timeout-minutes: 15
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v6
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: windows-${{ matrix.dev }}-minimal
-          deps: testing_unit
-          pydeps: ${{ matrix.dev == 'WEBGPU' && 'dawn-python' || '' }}
-      - name: Set env
-        shell: bash
-        run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
-      - name: Check Device.DEFAULT and print some source
-        shell: bash
-        run: |
-          python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
-          DEBUG=4 python test/test_tiny.py TestTiny.test_plus
-      - name: Run test_tiny
-        shell: bash
-        run: python -m pytest -n=auto test/test_tiny.py --durations=20
-
 # ****** Compile-only Tests ******

  compiletests:
    strategy:
      fail-fast: false
      matrix:
-        backend: [ir3, nak]
-    name: Compile-only (${{ matrix.backend }})
-    runs-on: *linux
+        dev:
+          - 'NULL:IR3:a630'
+          - 'NULL:QCOMCL:a630'
+          - 'NULL:NAK:sm_120'
+    name: Compile-only (DEV=${{ matrix.dev }})
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
+    env:
+      NULL_ALLOW_COPYOUT: 1
+      DEV: ${{ matrix.dev }}${{ contains(matrix.dev, 'a630') && ',IMAGE_PITCH_ALIGNMENT=64' || '' }}
+      IMAGE: ${{ contains(matrix.dev, 'a630') && '1' || '0' }}
    steps:
      - name: Checkout Code
        uses: actions/checkout@v6
@@ -786,51 +669,15 @@ jobs:
        with:
          key: compile-${{ matrix.backend }}
          deps: "testing_unit mesa"
-      - name: Set env
+          qemu: ${{ contains(matrix.dev, 'QCOMCL') }}
+      - name: Test IMAGE
        shell: bash
-        run: printf "NULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'DEV=NULL:IR3:a630' || matrix.backend == 'nak' && 'DEV=NULL:NAK:sm_120' }}" >> $GITHUB_ENV
+        if: contains(matrix.dev, 'a630')
+        run: DEBUG=7 python3 test/backend/test_ops.py TestOps.test_gemm | grep isam
      - name: Run test_ops
        shell: bash
        run: |
          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
-          python -m pytest -n=auto test/backend/test_ops.py --durations=20
-      - name: Run test_ops (IMAGE)
-        if: matrix.backend == 'ir3'
-        shell: bash
-        env:
-          IMAGE: 1
-          DEV: "NULL:IR3:a630,IMAGE_PITCH_ALIGNMENT=64"
-        run: |
-          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_gemm | grep image_load
-          python -m pytest -n=auto test/backend/test_ops.py --durations=20
-  qcomclcompiletests:
-    name: Compile-only (QCOM CL)
-    runs-on: ubuntu-24.04-arm
-    timeout-minutes: 15
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v6
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: compile-qcomcl
-          deps: testing_unit
-          tinydreno: 'true'
-      - name: Set env
-        shell: bash
-        run: printf "DEV=NULL:QCOMCL:a630\nNULL_ALLOW_COPYOUT=1" >> $GITHUB_ENV
-      - name: Run test_ops
-        shell: bash
-        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
-          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
-          python -m pytest -n=auto test/backend/test_ops.py --durations=20
-      - name: Run test_ops (IMAGE)
-        shell: bash
-        env:
-          IMAGE: 1
-          DEV: "NULL:QCOMCL:a630,IMAGE_PITCH_ALIGNMENT=64"
-        run: |
-          DEBUG=4 python test/backend/test_ops.py TestOps.test_gemm | grep read_imagef
-          python -m pytest -n=auto test/backend/test_ops.py --durations=20
+          # QCOMCL compiles in qemu, too slow for parallel workers
+          ${{ contains(matrix.dev, 'QCOMCL') && 'PARALLEL=0' || '' }} python -m pytest -n=auto test/backend/test_ops.py --durations=20
@@ -69,3 +69,4 @@ mutants
 dagre/
 graphlib/
 uv.lock
+pi_session_window0.jsonl
@@ -3,3 +3,5 @@
 - Run tests with `-n12` for speed (e.g. `python -m pytest test/null/test_dtype.py -x -q -n12`)
 - 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.
@@ -140,7 +140,7 @@ Documentation along with a quick start guide can be found on the [docs website](
 ```python
 from tinygrad import Tensor

-x = Tensor.eye(3)
+x = Tensor.eye(3).clone()  # clone to make it a buffer
 y = Tensor([[2.0,0,-2.0]])
 z = y.matmul(x).sum()
 z.backward()
@@ -88,7 +88,7 @@ def example_3_custom_uop(a:Tensor, correct):

    # store all the per lane accumulators to LOCAL
    local_accs = UOp.placeholder((LCLS,), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
-    local_accs = local_accs.after(local_accs[lane].store(acc[0]).barrier())
+    local_accs = local_accs.after(local_accs[lane].store(acc[0]))

    # accumulate LOCALs into a single per CU accumulator
    late_reduce_loop = UOp.range(LCLS, 3, AxisType.REDUCE)
@@ -1,7 +1,8 @@
 ::: tinygrad.dtype.DType

-::: tinygrad.dtype.dtypes
+::: tinygrad.dtype.DTypes
    options:
+        heading: dtypes
        members: true
        members_order: source
        show_labels: false
@@ -9,8 +9,7 @@ from extra.lr_scheduler import OneCycleLR
 GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))]

 # override tinygrad defaults
-dtypes.default_float = dtypes.half
-Context(FUSE_OPTIM=1).__enter__()
+Context(DEFAULT_FLOAT=dtypes.half, FUSE_OPTIM=1).__enter__()

 # from https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
 batchsize = getenv("BS", 1024)
@@ -22,10 +22,6 @@ class Attention:
    self.head_dim = dim // n_heads

  def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]) -> Tensor:
-    if mask is not None or start_pos.val == 0:
-      # no symbolic shape qkv when consuming prompts
-      start_pos = start_pos.val
-
    if HALF: x = x.half()
    xqkv = self.c_attn(x).reshape(None, None, 3, self.n_heads, self.head_dim)
    xq, xk, xv = [xqkv[:, :, i, :, :] for i in range(3)]
@@ -38,12 +34,8 @@ class Attention:
    # update the cache
    self.cache_kv[:, :, start_pos:start_pos+seqlen, :, :].assign(Tensor.stack(xk, xv)).realize()

-    if start_pos > 0:
-      keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
-      values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
-    else:
-      keys = xk
-      values = xv
+    keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
+    values = self.cache_kv[1][:, :start_pos+seqlen, :, :]

    xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
    return self.c_proj(xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2).reshape(bsz, seqlen, self.dim))
@@ -86,15 +78,14 @@ class Transformer:
      seqlen = tokens.shape[1]
      tok_emb = self.wte(tokens)

-    # not symbolic when consuming the prompt
-    selected_pos = (0, seqlen) if start_pos.val == 0 else (start_pos, start_pos+1)
-    pos_emb = self.wpe(self.allpos.shrink((None, selected_pos)))
+    # start_pos is a bound Variable, so everything below it stays symbolic
+    pos_emb = self.wpe(self.allpos.shrink((None, (start_pos, start_pos+seqlen))))

    h = tok_emb + pos_emb

    if HALF: h = h.half()

-    mask = Tensor.full((1, 1, seqlen, start_pos.val+seqlen), float("-inf"), dtype=h.dtype).triu(start_pos.val+1) if seqlen > 1 else None
+    mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype).triu(start_pos+1) if seqlen > 1 else None

    for hi in self.h: h = hi(h, start_pos, mask)

@@ -1,11 +1,11 @@
 import os, random, pickle, queue, struct, math, functools, hashlib, time
 from typing import List
 from pathlib import Path
-from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu_count
+from multiprocessing import Queue, Process, shared_memory, connection, Lock

 import numpy as np
 from tinygrad import dtypes, Tensor
-from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX
+from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX, NUM_CPU_THREADS
 from tinygrad.nn.state import TensorIO

 ### ResNet
@@ -131,7 +131,7 @@ def batch_load_resnet(batch_size=64, val=False, shuffle=True, seed=None, pad_fir
    else: X = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name}")
    Y = [None] * (batch_size*BATCH_COUNT)

-    for _ in range(cpu_count()):
+    for _ in range(NUM_CPU_THREADS.value):
      p = Process(target=loader_process, args=(q_in, q_out, X, seed))
      p.daemon = True
      p.start()
@@ -212,7 +212,7 @@ def batch_load_train_bert(BS:int, seed:int|None=None):
    rng.shuffle(fs)
    train_files.append(fs.pop(0))

-  cycle_length = min(getenv("NUM_CPU_THREADS", min(os.cpu_count(), 8)), len(train_files))
+  cycle_length = min(NUM_CPU_THREADS.value, len(train_files))
  assert cycle_length > 0, "cycle_length must be greater than 0"

  dataset = InterleavedDataset(train_files, cycle_length)
@@ -301,7 +301,7 @@ def batch_load_unet3d(preprocessed_dataset_dir:Path, batch_size:int=6, val:bool=
    X = Tensor.empty(*sz, dtype=dtypes.float32, device=f"disk:/dev/shm/{shm_name_x}")
    Y = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name_y}")

-    for _ in range(cpu_count()):
+    for _ in range(NUM_CPU_THREADS.value):
      proc = Process(target=load_unet3d_data, args=(preprocessed_dataset_dir, seed, queue_in, queue_out, X, Y))
      proc.daemon = True
      proc.start()
@@ -437,7 +437,7 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
  dataset_iter = iter(image_ids)

  try:
-    for _ in range(cpu_count()):
+    for _ in range(NUM_CPU_THREADS.value):
      proc = Process(
        target=load_retinanet_data,
        args=(base_dir, val, queue_in, queue_out, imgs, boxes, labels),
@@ -1282,10 +1282,10 @@ def train_bert():
        previous_step = i

 def train_llama3():
-  from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8_DTYPE, MXFP8
+  from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8_DTYPE, MXFP8, MXFP4
  from examples.llama3 import MODEL_PARAMS
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
-  from examples.mlperf.optim import GradAccClipAdamW
+  from examples.mlperf.optim import GradAccClipAdamW, clip_grads

  INITMLPERF = getenv("INITMLPERF")
  RUNMLPERF = getenv("RUNMLPERF")
@@ -1434,9 +1434,9 @@ def train_llama3():
    load_state_dict(scheduler, safe_load(fn), realize=False)

  fp8_amax = [t for ts in model._fp8_amax.values() for t in ts]
-  fp8_next_amax = [t for ts in model._fp8_next_amax.values() for t in ts] if hasattr(model, "_fp8_next_amax") else []
-  fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts] if hasattr(model, "_fp8_grad_amax") else []
-  fp8_next_grad_amax = [t for ts in model._fp8_next_grad_amax.values() for t in ts] if hasattr(model, "_fp8_next_grad_amax") else []
+  fp8_next_amax = [t for ts in model._fp8_next_amax.values() for t in ts]
+  fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts]
+  fp8_next_grad_amax = [t for ts in model._fp8_next_grad_amax.values() for t in ts]
  fp8_inv_scales = list(model._fp8_inv_scale.values()) + list(model._fp8_next_inv_scale.values())

  from tinygrad.nn.state import get_state_dict
@@ -1458,12 +1458,12 @@ def train_llama3():

  # realize everything here
  if optim.master_params: Tensor.realize(*optim.master_params)
-  Tensor.realize(*optim.params, *fp8_inv_scales, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)
+  loss_acc = Tensor.zeros(1, dtype=dtypes.float32, device=device)
+  Tensor.realize(loss_acc, *optim.params, *fp8_inv_scales, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)

  @TinyJit
  def minibatch(tokens:Tensor):
-    for nxt in fp8_next_amax: nxt.assign(0)
-    for nxt in fp8_next_grad_amax: nxt.assign(0)
+    model.reset_amax()
    if is_dp: tokens = tokens.to(None).shard(device, 0)
    if is_mp: tokens = tokens.shard(device)
    if not is_sharding: tokens = tokens.to(None)
@@ -1477,23 +1477,24 @@ def train_llama3():
    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, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)
+    loss_acc.assign(loss_acc + loss.flatten().float())
+    return loss_acc.realize(*grads, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)

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

    for g in grads: g.assign(0)
-    for cur, nxt in zip(fp8_amax, fp8_next_amax): cur.assign(nxt)
-    for cur, nxt in zip(fp8_grad_amax, fp8_next_grad_amax): cur.assign(nxt)
+    model.update_amax()

    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, *fp8_amax, *fp8_grad_amax)
+    loss_cpu = loss_acc.to("CPU")
+    Tensor.realize(lr_cpu, grad_norm_cpu, loss_cpu, loss_acc.assign(0), *grads, *fp8_inv_scales, *fp8_amax, *fp8_grad_amax)

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

  @TinyJit
  @Context(TRAINING=0)
@@ -1548,8 +1549,8 @@ def train_llama3():
      st = time.perf_counter()

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

      gt = time.perf_counter()
      ret = optim_step()
-      lr, grad_norm = ret[0].item(), ret[1].item()
+      lr, grad_norm, loss = ret[0].item(), ret[1].item(), ret[2].item() / accum_steps
      et = time.perf_counter()

-      loss = sum(losses) / len(losses)
      optim_time = et - gt
      dev_time += optim_time
      step_time = et - st
@@ -1578,7 +1578,7 @@ def train_llama3():

      mem_gb = GlobalCounters.mem_used / 1e9
      gflops = GlobalCounters.global_ops / 1e9 / dev_time
-      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * 4.6e15)) * 100
+      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * (9.2e15 if MXFP4 else 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"{lr:.12f} LR, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
@@ -1667,7 +1667,7 @@ def train_llama3():
 def train_gptoss():
  from examples.mlperf.models.gpt_oss import GPTOSS, GPT_OSS_20B, apply_grad, FP8_DTYPE
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
-  from examples.mlperf.optim import GradAccClipAdamW
+  from examples.mlperf.optim import GradAccClipAdamW, GradAccClipAdamWGroup, clip_grads

  BENCHMARK = getenv("BENCHMARK")

@@ -1710,9 +1710,10 @@ def train_gptoss():
    wandb.init(config=config, **wandb_args, project="MLPerf-gpt-oss")

  model_params = GPT_OSS_20B
-  model_params['vocab_size'] = 128256
+  model_params['vocab_size'] = getenv("VOCAB_SIZE", 128256)
  real_vocab_size = model_params['vocab_size']
  if (layers:=getenv("LAYERS")) != 0: model_params['n_layers'] = layers
+  if (experts:=getenv("EXPERTS")) != 0: model_params['n_experts'] = experts
  print(f"model parameters: {model_params}")

  model = GPTOSS(**model_params, max_context=SEQLEN)
@@ -1733,16 +1734,24 @@ def train_gptoss():
  is_offload_optim = bool(getenv("OFFLOAD_OPTIM"))
  is_fake_offload = Device.DEFAULT == "NULL"
  optim_device = ("CPU" if not is_fake_offload else "NULL:99") if is_offload_optim else None
-  optim = GradAccClipAdamW(params, 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)
+  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),
+  )

  for p in optim.params:
-    grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
-    p.grad = p.zeros_like(dtype=grad_dtype).contiguous()
+    p.grad = p.zeros_like(dtype=dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype).contiguous()
+    if getattr(p, "_zero2", False): p.grad = optim.optimizers[0]._zero_shard(p.grad)
  grads = [p.grad for p in optim.params]

  from extra.gemm.cdna_asm_gemm import _mx_block_scale
  model_state = get_state_dict(model)
-  fp8_scale_names = {n: f"{n}_scale" for n, t in model_state.items() if t.dtype == FP8_DTYPE}
+  def _scale_key(n):
+    if "." in n and (c:=f"{(b:=n.rsplit('.',1))[0]}_scale.{b[1]}") in model_state: return c
+    return f"{n}_scale"
+  fp8_scale_names = {n: _scale_key(n) for n, t in model_state.items() if t.dtype == FP8_DTYPE}
  fp8_inv_scales = [model_state[sname] for sname in fp8_scale_names.values()]
  for wname, sname in fp8_scale_names.items():
    w, scale = model_state[wname], model_state[sname]
@@ -1764,8 +1773,13 @@ def train_gptoss():
  def minibatch(tokens:Tensor):
    if is_dp: tokens = tokens.to(None).shard(device, 0)
    if not is_sharding: tokens = tokens.to(None)
+
    logits:Tensor = model(tokens[:, :-1], save=True)
-    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)
@@ -1775,7 +1789,8 @@ def train_gptoss():

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

    for g in grads: g.assign(0)
@@ -25,6 +25,7 @@ FUSED_SILU_W13 = getenv("FUSED_SILU_W13", 0)
 SPLIT_W13 = getenv("SPLIT_W13", 0)
 COLUMNWISE_WEIGHT_SCALE = getenv("COLUMNWISE_WEIGHT_SCALE", 0)
 MXFP8 = getenv("MXFP8", 0)
+MXFP4 = getenv("MXFP4", 0)

 FP8_DTYPE = dtypes.fp8e4m3
 FP8_GRAD_DTYPE = dtypes.fp8e5m2
@@ -44,6 +45,11 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca
      from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
      if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T),)
    return (x @ w.T,)
+  if MXFP4:
+    assert x is not None, "MXFP4 matmul requires an unquantized input"
+    from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm
+    if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T, mxfp4=True),)
+    return (x @ w.T,)
  assert w_inv_scale is not None, "fp8 matmul requires w_inv_scale (weights must be stored in fp8 with per-tensor scale)"
  if MXFP8:
    from extra.gemm.cdna_asm_gemm import asm_gemm, quantize_mxfp8, mx_pack, can_use_asm_gemm, _mx_block_scale
@@ -77,9 +83,9 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca
      return out, x_fp8
  return (x_fp8.dot(w.T, dtype=dtypes.float) * ((amax_x.float() + 1e-8) / FP8_MAX) * w_inv_scale).cast(dtypes.bfloat16), x_fp8

-def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
-                         next_amax_x:Tensor, grad_amax_state:Tensor, next_grad_amax_state:Tensor):
-  if FUSED_ADD_NORM_MUL_QUANTIZE:
+def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor|None,
+                         next_amax_x:Tensor|None, grad_amax_state:Tensor|None, next_grad_amax_state:Tensor|None):
+  if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
    from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_rmsnorm_mul_quantize_fp8
    x_fp8, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
@@ -90,9 +96,9 @@ def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, ep
                     next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
  return out, x_normed, rrms, ret

-def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
-                             next_amax_x:Tensor, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
-  if FUSED_ADD_NORM_MUL_QUANTIZE:
+def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor|None,
+                             next_amax_x:Tensor|None, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
+  if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
    from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_add_rmsnorm_mul_quantize_fp8
    x_fp8, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
@@ -105,10 +111,15 @@ def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w
  return out, h, x_normed, rrms, ret

 def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor,
-                             amax_x2:Tensor, next_amax_x2:Tensor,
-                             grad_amax_xw13:Tensor, next_grad_amax_xw13:Tensor,
-                             grad_amax_xout:Tensor, next_grad_amax_xout:Tensor):
-  if FUSED_SILU_W13:
+                             amax_x2:Tensor|None, next_amax_x2:Tensor|None,
+                             grad_amax_xw13:Tensor|None, next_grad_amax_xw13:Tensor|None,
+                             grad_amax_xout:Tensor|None, next_grad_amax_xout:Tensor|None):
+  if FUSED_SILU_W13 and MXFP4:
+    from extra.llama_kernels.swiglu import swiglu
+    out, *ret = matmul(swiglu(x_w13), w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout,
+                       next_grad_amax_state=next_grad_amax_xout, next_amax_x=next_amax_x2)
+    return out, ret
+  if FUSED_SILU_W13 and not MXFP4:
    from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
    x2_fp8 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13,
                                                 next_grad_amax_state=next_grad_amax_xw13, amax_out=next_amax_x2)
@@ -158,14 +169,15 @@ class FlatTransformer:
    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).clone().is_param_(False)

    def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().is_param_(False)
+    n_amax = 0 if MXFP4 else n_layers
    names = ["xqkv", "xo", "x2"]
    names += ["x1", "x3"] if SPLIT_W13 else ["x13"]
-    self._fp8_amax = {name: [_amax() for _ in range(n_layers)] for name in names}
-    self._fp8_next_amax = {name: [_amax() for _ in range(n_layers)] for name in names}
+    self._fp8_amax = {name: [_amax() for _ in range(n_amax)] for name in names}
+    self._fp8_next_amax = {name: [_amax() for _ in range(n_amax)] for name in names}
    grad_names = ["xqkv", "xo", "xout"]
    grad_names += ["xw1", "xw3"] if SPLIT_W13 else ["xw13"]
-    self._fp8_grad_amax = {name: [_amax() for _ in range(n_layers)] for name in grad_names}
-    self._fp8_next_grad_amax = {name: [_amax() for _ in range(n_layers)] for name in grad_names}
+    self._fp8_grad_amax = {name: [_amax() for _ in range(n_amax)] for name in grad_names}
+    self._fp8_next_grad_amax = {name: [_amax() for _ in range(n_amax)] for name in grad_names}
    w_scales = [("wqkv", s_qkv), ("wo", s_o), ("w2", s_2)]
    w_scales += [("w1", s_1), ("w3", s_3)] if SPLIT_W13 else [("w13", s_13)]
    self._fp8_inv_scale = {name: (s if MXFP8 else s.float()).contiguous().is_param_(False) for name, s in w_scales}
@@ -179,6 +191,9 @@ class FlatTransformer:
      from extra.gemm.cdna_asm_gemm import quantize_mxfp8
      w_q, w_e8, _ = quantize_mxfp8(w.reshape(self.n_layers * out_features, in_features))
      return w_q.reshape(self.n_layers, out_features, in_features), w_e8.reshape(self.n_layers, out_features, in_features // 32)
+    if MXFP4:
+      # FP4 is produced dynamically so optimizer updates always start from the current BF16 weight.
+      return w.cast(dtypes.bfloat16), Tensor.ones(self.n_layers)
    amax = (w.abs().max(axis=2) if COLUMNWISE_WEIGHT_SCALE else w.abs().flatten(1).max(1)).detach()
    scale = FP8_MAX / (amax + 1e-8)
    inv_scale = (amax + 1e-8) / FP8_MAX
@@ -186,9 +201,10 @@ class FlatTransformer:
    return (w * scale_b).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE), inv_scale

  def attention(self, x:Tensor, freqs_cis:Tensor, *, attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
-                amax_xqkv:Tensor, amax_xo:Tensor, s_qkv:Tensor, s_o:Tensor,
-                next_amax_xqkv:Tensor, next_amax_xo:Tensor,
-                grad_amax_xqkv:Tensor, grad_amax_xo:Tensor, next_grad_amax_xqkv:Tensor, next_grad_amax_xo:Tensor):
+                amax_xqkv:Tensor|None, amax_xo:Tensor|None, s_qkv:Tensor, s_o:Tensor,
+                next_amax_xqkv:Tensor|None, next_amax_xo:Tensor|None,
+                grad_amax_xqkv:Tensor|None, grad_amax_xo:Tensor|None,
+                next_grad_amax_xqkv:Tensor|None, next_grad_amax_xo:Tensor|None):
    bsz, seqlen, _ = x.shape
    saves = []

@@ -310,28 +326,33 @@ class FlatTransformer:
          for i in range(len(amax_dict[name])):
            amax_dict[name][i] = amax_dict[name][i].to(device).contiguous().is_param_(False)

+  def reset_amax(self):
+    for st in (self._fp8_next_amax, self._fp8_next_grad_amax):
+      for ts in st.values():
+        for t in ts: t.assign(0)
+
+  def update_amax(self):
+    for cur, nxt in ((self._fp8_amax, self._fp8_next_amax), (self._fp8_grad_amax, self._fp8_next_grad_amax)):
+      for name in cur:
+        for c, n in zip(cur[name], nxt[name]): c.assign(n)
+
  def __call__(self, tokens:Tensor, save:bool=True):
    h = self.tok_embeddings(tokens)
    freqs_cis = self.freqs_cis.cast(h.dtype)
    if not getenv("HK_FLASH_ATTENTION"): freqs_cis = freqs_cis[:, :tokens.shape[1], :, :, :]
    a, na, ga, nga, s = self._fp8_amax, self._fp8_next_amax, self._fp8_grad_amax, self._fp8_next_grad_amax, self._fp8_inv_scale
+    def amax_kwargs(i:int, act_names:tuple[str, ...], grad_names:tuple[str, ...]) -> dict[str, Tensor|None]:
+      specs = (("amax_", a, act_names), ("next_amax_", na, act_names), ("grad_amax_", ga, grad_names), ("next_grad_amax_", nga, grad_names))
+      if MXFP4: return dict.fromkeys(f"{prefix}{name}" for prefix, _, names in specs for name in names)
+      return {f"{prefix}{name}":val[name][i] for prefix, val, names in specs for name in names}
    for i in range(self.n_layers):
-      attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i],
-                         amax_xqkv=a["xqkv"][i], amax_xo=a["xo"][i], s_qkv=s["wqkv"][i], s_o=s["wo"][i],
-                         next_amax_xqkv=na["xqkv"][i], next_amax_xo=na["xo"][i],
-                         grad_amax_xqkv=ga["xqkv"][i], grad_amax_xo=ga["xo"][i],
-                         next_grad_amax_xqkv=nga["xqkv"][i], next_grad_amax_xo=nga["xo"][i])
-      ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], w2=self.w2[i],
-                        amax_x2=a["x2"][i], s_2=s["w2"][i], grad_amax_xout=ga["xout"][i], next_grad_amax_xout=nga["xout"][i],
-                        next_amax_x2=na["x2"][i])
+      attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i], s_qkv=s["wqkv"][i], s_o=s["wo"][i],
+                         **amax_kwargs(i, ("xqkv", "xo"), ("xqkv", "xo")))
+      ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], w2=self.w2[i], s_2=s["w2"][i], **amax_kwargs(i, ("x2",), ("xout",)))
      if SPLIT_W13:
-        ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i], amax_x1=a["x1"][i], amax_x3=a["x3"][i],
-                          next_amax_x1=na["x1"][i], next_amax_x3=na["x3"][i],
-                          s_1=s["w1"][i], s_3=s["w3"][i], grad_amax_xw1=ga["xw1"][i], grad_amax_xw3=ga["xw3"][i],
-                          next_grad_amax_xw1=nga["xw1"][i], next_grad_amax_xw3=nga["xw3"][i])
+        ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i], s_1=s["w1"][i], s_3=s["w3"][i], **amax_kwargs(i, ("x1", "x3"), ("xw1", "xw3")))
      else:
-        ffn_kwargs.update(w13=self.w13[i], amax_x13=a["x13"][i], s_13=s["w13"][i], grad_amax_xw13=ga["xw13"][i],
-                          next_grad_amax_xw13=nga["xw13"][i], next_amax_x13=na["x13"][i])
+        ffn_kwargs.update(w13=self.w13[i], s_13=s["w13"][i], **amax_kwargs(i, ("x13",), ("xw13",)))
      h, *_ = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)

    logits = matmul(self.norm(h), self.output[0], fp8=False)[0]
@@ -415,9 +436,7 @@ if __name__ == "__main__":
  @TinyJit
  def fwd_bwd(tokens:Tensor):
    with Timing("python forward: "):
-      for amax_dict in (model._fp8_next_amax, model._fp8_next_grad_amax):
-        for ts in amax_dict.values():
-          for nxt in ts: nxt.assign(0)
+      model.reset_amax()
      logits = model(tokens[:, :-1], save=llama_size=="8B")
      loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
    with Timing("python backward: "):
@@ -12,11 +12,15 @@ from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb
 from extra.llama_kernels.rmsnorm import rmsnorm
-from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8
+from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8, asm_gemm, can_use_asm_gemm
+from extra.gemm.moe_gemm import grouped_mx_gemm
+from extra.gemm.moe_routing import route, dispatch, combine

 FP8_DTYPE = dtypes.fp8e4m3
 FP8_MAX = 448.0
-INIT_STD = 0.008
+INIT_STD = 0.02
+ASM_GEMM = getenv("ASM_GEMM", 0)
+

 def _quant_dequant_fwd(x:Tensor) -> Tensor:
  # x (2d bf16) -> bf16 value after an mxfp8 round-trip (1x32 block scaling on the last axis)
@@ -50,8 +54,7 @@ def _dequant_fwd_fxn(wq_p, ws_p, device):
  return _dequant_fwd(Tensor(wq_p, device=device), Tensor(ws_p, device=device))

 def _dequant_bwd(grad:UOp, call:UOp) -> tuple:
-  w_scale = Tensor(call.src[2])
-  return ((Tensor(grad).cast(dtypes.bfloat16) * _mx_scale(w_scale).cast(dtypes.bfloat16)).uop, None)
+  return (Tensor(grad).cast(dtypes.bfloat16).uop, None)

 def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:
  fxn = _dequant_fwd_fxn(w_q.as_param(0).uop, w_scale.as_param(1).uop, w_q.device)
@@ -60,10 +63,34 @@ def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:

 def matmul_mx(x:Tensor, w_q:Tensor, w_scale:Tensor) -> Tensor:
  l_shape = x.shape[:-1]
+  if ASM_GEMM:
+    from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm, mx_pack
+    x2, K, N = x.reshape(-1, x.shape[-1]), x.shape[-1], w_q.shape[0]
+    wq, ws = w_q, w_scale
+    if (pad := (-K) % 256):
+      x2 = x2.pad(((0, 0), (0, pad)))
+      wq = wq.pad(((0, 0), (0, pad)))
+      ws = ws.pad(((0, 0), (0, pad // 32)), value=127).cast(dtypes.uint8)
+    if (npad := (-N) % 256):
+      wq = wq.pad(((0, npad), (0, 0)))
+      ws = ws.pad(((0, npad), (0, 0)), value=127).cast(dtypes.uint8)
+    x_q, x_e8, x_si = quantize_mxfp8(x2)
+    if x_si is not None and can_use_asm_gemm(x_q, wq.T):
+      out = asm_gemm(x_q, wq.T, mx=True, mx_scales=(x_si, x_e8, mx_pack(ws), ws), mx_w_stored=True)
+      return (out[:, :N] if npad else out).reshape(*l_shape, N).cast(dtypes.bfloat16)
  x_phys = quant_dequant_mx(x.reshape(-1, x.shape[-1])).reshape(*l_shape, x.shape[-1])
  w_phys = dequant_weight(w_q, w_scale)
  return (x_phys @ w_phys.T).cast(dtypes.bfloat16)

+def _pad_to_mult(t:Tensor, axis:int, mult:int=256) -> Tensor:
+  if (r := (-t.shape[axis]) % mult) == 0: return t
+  pads = [(0, 0)] * t.ndim
+  pads[axis] = (0, r)
+  return t.pad(tuple(pads))
+
+def _pad_cols(t:Tensor) -> Tensor: return _pad_to_mult(t, -1)
+def _pad_rows(t:Tensor) -> Tensor: return _pad_to_mult(t, -2)
+
 def swiglu(x:Tensor, limit:float=7.0, alpha:float=1.702) -> Tensor:
  x_glu, x_linear = x[..., ::2], x[..., 1::2]
  x_glu = x_glu.clamp(max_=limit)
@@ -100,9 +127,9 @@ class GPTOSS:
    self.ffn_norm = Tensor.ones(n_layers, dim).contiguous()
    self.gate = Tensor.normal(n_layers, n_experts, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
    self.gate_bias = Tensor.zeros(n_layers, n_experts, dtype=dtypes.bfloat16).contiguous()
-    self.w_gate_up, self.w_gate_up_scale = self._quant_weight(n_layers, n_experts, intermediate_size * 2, dim)
+    self.w_gate_up, self.w_gate_up_scale = self._quant_weight(n_layers, n_experts, intermediate_size * 2, dim, moe=True)
    self.w_gate_up_bias = Tensor.zeros(n_layers, n_experts, intermediate_size * 2, dtype=dtypes.bfloat16).contiguous()
-    self.w_down, self.w_down_scale = self._quant_weight(n_layers, n_experts, dim, intermediate_size, std=scaled_std)
+    self.w_down, self.w_down_scale = self._quant_weight(n_layers, n_experts, dim, intermediate_size, std=scaled_std, moe=True)
    self.w_down_bias = Tensor.zeros(n_layers, n_experts, dim, dtype=dtypes.bfloat16).contiguous()

    # output
@@ -112,10 +139,16 @@ class GPTOSS:
    self.output = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
    self.freqs_cis = precompute_freqs_cis(head_dim, max_context * 2, rope_theta).contiguous().is_param_(False)

-  def _quant_weight(self, *shape:int, std:float=INIT_STD):
-    w = Tensor.zeros(*shape) if getenv("ZEROS") else Tensor.normal(*shape, mean=0.0, std=std)
-    w_q, w_e8, _ = quantize_mxfp8(w)
-    return w_q, w_e8.is_param_(False)
+  def _quant_weight(self, *shape:int, std:float=INIT_STD, moe:bool=False):
+    def _one(*s:int):
+      w = Tensor.zeros(*s) if getenv("ZEROS") else Tensor.normal(*s, mean=0.0, std=std)
+      w_q, w_e8, _ = quantize_mxfp8(_pad_cols(_pad_rows(w)) if moe else w)
+      return w_q, w_e8.is_param_(False)
+    if moe:
+      qs = [_one(*shape[1:]) for _ in range(shape[0])]
+      for q in qs: q[0]._zero2 = True  # grad arrives sharded on the expert axis under ZeRO-2 (moe_gemm)
+      return [q[0] for q in qs], [q[1] for q in qs]
+    return _one(*shape)

  def _attn_mask(self, seqlen:int, dtype) -> Tensor:
    i, j = Tensor.arange(seqlen).reshape(seqlen, 1), Tensor.arange(seqlen).reshape(1, seqlen)
@@ -150,12 +183,14 @@ class GPTOSS:
    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
    xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)  # (B,N,H,D)/(B,N,KV,D)

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

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

  def feed_forward(self, x:Tensor, *, ffn_norm:Tensor, gate:Tensor, gate_bias:Tensor,
                   w_gate_up:Tensor, w_gate_up_scale:Tensor, w_gate_up_bias:Tensor,
                   w_down:Tensor, w_down_scale:Tensor, w_down_bias:Tensor):
    x_normed, rrms = rmsnorm(x, self.norm_eps)
    inp = x_normed * ffn_norm
-
    logits = inp.float() @ gate.float().T + gate_bias.float()
-    thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
-    weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)
+    dim, inter = self.dim, self.intermediate_size

-    out = None
-    for e in range(self.n_experts):
-      gate_up = matmul_mx(inp, w_gate_up[e], w_gate_up_scale[e]) + w_gate_up_bias[e]
-      y = (matmul_mx(swiglu(gate_up, self.swiglu_limit), w_down[e], w_down_scale[e]) + w_down_bias[e]).contiguous()
-      contrib = weights[..., e:e+1].cast(y.dtype) * y
-      out = contrib if out is None else out + contrib
+    if getenv("GROUPED_MOE", 0):
+      bsz, seqlen = x.shape[:2]
+      inp, logits = inp.reshape(-1, dim), logits.reshape(-1, self.n_experts)
+      r = route(logits, self.experts_per_tok, self.n_experts)
+      onehot = r.rows_e.one_hot(self.n_experts).float()
+      xg = dispatch(_pad_cols(inp.cast(dtypes.bfloat16)), r)
+      h = grouped_mx_gemm(xg, (w_gate_up, w_gate_up_scale), r.off)[:, :2*inter] + (onehot @ w_gate_up_bias.float()).cast(dtypes.bfloat16)
+      y = swiglu(h, self.swiglu_limit)
+      z = grouped_mx_gemm(_pad_cols(y.cast(dtypes.bfloat16)), (w_down, w_down_scale), r.off)[:, :dim] \
+          + (onehot @ w_down_bias.float()).cast(dtypes.bfloat16)
+      out = combine(z, r, inp.shape[0], self.experts_per_tok).reshape(bsz, seqlen, dim)
+      return out, [x_normed, rrms, xg, h, y, z, r.weights, r.dest_row, r.off]
+    else:
+      thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
+      weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)
+
+      out = None
+      for e in range(self.n_experts):
+        gu_q, gu_s = w_gate_up[e][:2*inter, :dim].contiguous(), w_gate_up_scale[e][:2*inter, :dim//32].contiguous()
+        dn_q, dn_s = w_down[e][:dim, :inter].contiguous(), w_down_scale[e][:dim, :inter//32].contiguous()
+        gate_up = matmul_mx(inp, gu_q, gu_s) + w_gate_up_bias[e]
+        y = (matmul_mx(swiglu(gate_up, self.swiglu_limit), dn_q, dn_s) + w_down_bias[e]).contiguous()
+        contrib = weights[..., e:e+1].cast(y.dtype) * y
+        out = contrib if out is None else out + contrib
    return out, [x_normed, rrms]

  @function(precompile=True, precompile_backward=True)
@@ -216,7 +267,11 @@ class GPTOSS:
                        w_down=self.w_down[i], w_down_scale=self.w_down_scale[i], w_down_bias=self.w_down_bias[i])
      h, *_ = self.run_layer(h, freqs_cis, mask_full, i % 2 == 0, attn_kwargs, ffn_kwargs, save=save)

-    logits = self.norm(h) @ self.output.T
+    h_normed = self.norm(h)
+    pad = (-self.dim) % 256
+    h_padded, w_padded = h_normed.pad((None, None, (0, pad))), self.output.pad(((0, 0), (0, pad)))
+    if ASM_GEMM and can_use_asm_gemm(h_padded, w_padded.T): logits = asm_gemm(h_padded, w_padded.T)
+    else: logits = h_normed @ self.output.T
    return logits

 def _get_pads(uop:UOp) -> list[UOp]:
@@ -1,8 +1,8 @@
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
-from tinygrad.nn.optim import Optimizer
+from tinygrad.nn.optim import Optimizer, OptimizerGroup
 from tinygrad.helpers import FUSE_OPTIM, getenv
-from tinygrad.uop.ops import UOp, Ops
+from tinygrad.uop.ops import UOp, Ops, AxisType

 STOCHASTIC_ROUND = getenv("STOCHASTIC_ROUND", 0)
 MASTER_WEIGHTS = getenv("MASTER_WEIGHTS", 0)
@@ -15,12 +15,18 @@ 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)
  return ((bits + noise) & 0xFFFF0000).bitcast(dtypes.float32).cast(dtypes.bfloat16)

+def clip_grads(grads:list[Tensor], grad_acc, clip_norm) -> Tensor:
+  for g in grads: g.assign(g / grad_acc)
+  total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
+  for g in grads: g.assign((g * (clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(g.dtype))
+  return total_norm
+
 class GradAccClipAdamW(Optimizer):
  def __init__(self, params:list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, grad_acc=1, clip_norm=1.0, device=None, fused=FUSE_OPTIM):
    super().__init__(params, lr, device, fused)
@@ -36,46 +42,29 @@ class GradAccClipAdamW(Optimizer):
      self.master_params = None

  def _zero_shard(self, t:Tensor) -> Tensor:
-    if not self.zero or (t.shape[0] % len(self.device)) != 0: return t
-    return Tensor(t.uop._shard(0, len(self.device)).multi(0)).clone()
+    if not self.zero or t.ndim < 2 or (t.shape[0] % len(self.device)) != 0: return t
+    return Tensor(t.uop._shard(0, UOp.range(len(self.device), -1, AxisType.DEVICE)).unshard(0)).clone()

  def _zero_gather(self, t:Tensor) -> Tensor:
    if not isinstance(t.device, tuple) or t.uop.axis != 0: return t
    n, sz = len(t.device), t.shape[0] // len(t.device)
    return Tensor.cat(*[t[p*sz:(p+1)*sz] for p in range(n)], dim=0)

-  def fstep(self, grads:list[Tensor]):
-    if self.fused:
-      out, extra = self._step([], grads)
-      updates = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
-    else:
-      updates, extra = self._step([], grads)
+  def fschedule_step(self, grads:list[Tensor]) -> list[Tensor]:
+    updates, extra = self._step([], grads)
    for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i], self.master_params[i] if self.master_params else None))
-    # collect inv_scale tensors attached to fp8 params (set by _apply_update)
    fp8_inv_scales = [tt._inv_scale for tt in self.params if hasattr(tt, '_inv_scale')]
    fp8_next_inv_scales = [tt._next_inv_scale for tt in self.params if hasattr(tt, '_next_inv_scale')]
-    to_realize = extra+self.params+self.buffers+(self.master_params or [])+fp8_inv_scales+fp8_next_inv_scales
+    return extra + self.params + self.buffers + (self.master_params or []) + fp8_inv_scales + fp8_next_inv_scales

-    Tensor.realize(*to_realize)
-    return extra[-1]
+  def fstep(self, grads:list[Tensor], grad_norm:Tensor|None=None):
+    Tensor.realize(*([grad_norm] if grad_norm is not None else []), *self.fschedule_step(grads))

  def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
    grads = list(grads)

    for i in range(len(grads)):
      if grads[i].device != self.m[i].device: grads[i] = grads[i].to(self.m[i].device)
-
-    if self.fused:
-      grads[0].assign(grads[0] / self.grad_acc)
-      total_norm = grads[0].float().square().sum().sqrt()
-      grads[0].assign((grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype))
-    else:
-      for i in range(len(grads)):
-        grads[i].assign(grads[i] / self.grad_acc)
-      total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
-      for i in range(len(grads)):
-        grads[i].assign((grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype))
-
    ret = []
    self.b1_t *= self.b1
    self.b2_t *= self.b2
@@ -88,7 +77,7 @@ class GradAccClipAdamW(Optimizer):
      v_hat = v_new / (1.0 - self.b2_t)
      up = m_hat / (v_hat.sqrt() + self.eps)
      ret.append(self.lr * up)
-    return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]
+    return ret, [self.b1_t, self.b2_t] + self.m + self.v

  def _apply_update(self, t:Tensor, up:Tensor, master:Tensor|None=None) -> Tensor:
    w = master if master is not None else t
@@ -109,7 +98,7 @@ class GradAccClipAdamW(Optimizer):
        if self.zero: w_q, w_e8 = self._zero_gather(w_q), self._zero_gather(w_e8)
        new_e8 = w_e8.reshape(t._inv_scale.shape)
        t._inv_scale.assign(new_e8.shard_like(t._inv_scale) if offloaded else new_e8)
-        ret = w_q.reshape(new_w.shape)
+        ret = w_q.reshape(t.shape)
        return ret.shard_like(t) if offloaded else ret
      from examples.mlperf.models.flat_llama import FP8_MAX
      if IMMEDIATE_SCALE:
@@ -132,3 +121,24 @@ class GradAccClipAdamW(Optimizer):
      return ret.shard_like(t) if offloaded else ret
    out = new_w.cast(t.dtype)
    return out.shard_like(t) if offloaded else out
+
+class GradAccClipAdamWGroup(OptimizerGroup):
+  def __init__(self, *optimizers:GradAccClipAdamW):
+    super().__init__(*optimizers)
+    for o in self.optimizers[1:]: o.lr = self.optimizers[0].lr
+  def fstep(self, grads:list[Tensor], grad_norm:Tensor|None=None):
+    offset = 0
+    to_realize = []
+    for o in self.optimizers:
+      n = len(o.params)
+      to_realize += o.fschedule_step(grads[offset:offset+n])
+      offset += n
+    Tensor.realize(*to_realize, *([grad_norm] if grad_norm is not None else []))
+  @property
+  def lr(self): return self.optimizers[0].lr
+  @property
+  def device(self): return self.optimizers[0].device
+  @property
+  def master_params(self):
+    mp = [mp for o in self.optimizers for mp in (o.master_params or [])]
+    return mp if mp else None
@@ -1,8 +1,8 @@
 #!/usr/bin/env bash

 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
+export ROCM_PATH=${ROCM_PATH:-/opt/rocm-7.1.1}
+export PATH="$ROCM_PATH/bin:$PATH"
 export DEV=${DEV:-AMD}
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -16,7 +16,7 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export FP8=${FP8:-1}
+export MXFP4=${MXFP4:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
 export FAST_CE=${FAST_CE:-1}
 export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
@@ -26,7 +26,7 @@ export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
 export SPLIT_W13=${SPLIT_W13:-0}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}

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

@@ -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}
@@ -1,8 +1,8 @@
 #!/usr/bin/env bash

 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
+export ROCM_PATH=${ROCM_PATH:-/opt/rocm-7.1.1}
+export PATH="$ROCM_PATH/bin:$PATH"
 export DEV=${DEV:-AMD}
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -16,7 +16,7 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export FP8=${FP8:-1}
+export MXFP4=${MXFP4:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
 export FAST_CE=${FAST_CE:-1}
 export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
@@ -26,7 +26,7 @@ export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
 export SPLIT_W13=${SPLIT_W13:-0}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}

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

@@ -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}
@@ -1,4 +1,6 @@
 #!/bin/bash
+set -e
+
 export BENCHMARK=${BENCHMARK:-5}
 export EVAL_BS=0
 VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=${DEBUG:--0} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
@@ -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))

@@ -10,6 +10,8 @@ export DEVICE_IN_FUNCTION_BUG=1

 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
+export GROUPED_MOE=${GROUPED_MOE:-1}
 export ALL2ALL=${ALL2ALL:-1}
 export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
@@ -10,6 +10,8 @@ export DEVICE_IN_FUNCTION_BUG=1

 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ASM_GEMM=${ASM_GEMM:-1}
+export GROUPED_MOE=${GROUPED_MOE:-1}
 export ALL2ALL=${ALL2ALL:-1}
 export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
@@ -1,4 +1,4 @@
-import os, sys, pickle, time, re
+import os, sys, pickle, time, re, tempfile, struct, shutil, io
 import numpy as np
 if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"

@@ -9,6 +9,39 @@ from tinygrad.nn.onnx import OnnxRunner

 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
 OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
+PICKLE_OOB = getenv("PICKLE_OOB")
+
+def dump_pickle(obj, f):
+  if PICKLE_OOB:
+    # allows pickling when buffers don't fit in (CPU) RAM
+    # from openpilot/selfdrive/modeld/helpers.py
+    with tempfile.TemporaryFile(dir=".") as tmp:
+      def buffer_callback(pb: pickle.PickleBuffer):
+        m = pb.raw()
+        tmp.write(struct.pack('<q', m.nbytes))
+        tmp.write(m)
+        pb.release() # keep peak ram at ~1 buffer
+      stream = io.BytesIO()
+      pickle.Pickler(stream, protocol=5, buffer_callback=buffer_callback).dump(obj)
+      opcodes = stream.getvalue()
+      f.write(struct.pack('<q', len(opcodes)))
+      f.write(opcodes)
+      tmp.seek(0)
+      shutil.copyfileobj(tmp, f)
+  else: pickle.dump(obj, f)
+
+def load_pickle(f):
+  if PICKLE_OOB:
+    # allows unpickling when buffers don't fit in (CPU) RAM
+    # from openpilot/selfdrive/modeld/helpers.py
+    opcodes = f.read(struct.unpack('<q', f.read(8))[0])
+    def buffers():
+      while (h := f.read(8)):
+        pb = pickle.PickleBuffer(bytearray(struct.unpack('<q', h)[0]))
+        f.readinto(pb)
+        yield pb
+    return pickle.load(io.BytesIO(opcodes), buffers=buffers())
+  else: return pickle.load(f)

 def compile(onnx_file):
  run_onnx = OnnxRunner(onnx_file)
@@ -28,8 +61,8 @@ def compile(onnx_file):
    inputs = {k:Tensor(v.numpy(), device=Device.DEFAULT).realize() if 'img' in k else v for k,v in inputs.items()}
  print("created tensors")

-  run_onnx_jit = TinyJit(lambda **kwargs:
-                         next(iter(run_onnx({k:v.to(Device.DEFAULT) for k,v in kwargs.items()}).values())).cast('float32'), prune=True)
+  @TinyJit(prune=True)
+  def run_onnx_jit(**kwargs): return next(iter(run_onnx({k:v.to(Device.DEFAULT) for k,v in kwargs.items()}).values())).cast('float32')
  for i in range(3):
    GlobalCounters.reset()
    print(f"run {i}")
@@ -65,8 +98,7 @@ def compile(onnx_file):
  if (allowed_gated_read_image:=getenv("ALLOWED_GATED_READ_IMAGE", -1)) != -1:
    assert gated_read_image_count == allowed_gated_read_image, f"different gated read_image! {gated_read_image_count=}, {allowed_gated_read_image=}"

-  with open(OUTPUT, "wb") as f:
-    pickle.dump(run_onnx_jit, f)
+  with open(OUTPUT, "wb") as f: dump_pickle(run_onnx_jit, f)
  mdl_sz = os.path.getsize(onnx_file)
  pkl_sz = os.path.getsize(OUTPUT)
  print(f"mdl size is {mdl_sz/1e6:.2f}M")
@@ -75,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")
@@ -128,15 +167,9 @@ 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 = pickle.load(f)
+    with open(OUTPUT, "rb") as f: pickle_loaded = load_pickle(f)
    inputs = {name: Tensor(Tensor.randn(*view.shape, dtype=dtype).numpy(), device=device)
              for name, (view, _vars, dtype, device) in zip(pickle_loaded.captured.expected_names, pickle_loaded.captured.expected_input_info)}
    test_vs_compile(pickle_loaded, inputs)
@@ -144,11 +177,8 @@ if __name__ == "__main__":
    onnx_file = fetch(OPENPILOT_MODEL)
    inputs, outputs = compile(onnx_file)

-    with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
+    with open(OUTPUT, "rb") as f: pickle_loaded = load_pickle(f)

    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)
@@ -1,5 +1,6 @@
-import sys, pickle
+import sys
 from extra.bench_log import WallTimeEvent, BenchEvent
+from examples.openpilot.compile3 import load_pickle
 from tinygrad.helpers import getenv

 PKL = sys.argv[1] if len(sys.argv) > 1 else "/tmp/openpilot.pkl"
@@ -7,7 +8,7 @@ PKL = sys.argv[1] if len(sys.argv) > 1 else "/tmp/openpilot.pkl"
 load_times = []

 for _ in range(10):
-  with WallTimeEvent(BenchEvent.STEP) as wte: pickle.load(open(PKL, 'rb'))
+  with WallTimeEvent(BenchEvent.STEP) as wte: load_pickle(open(PKL, 'rb'))
  load_times.append(wte.time)
  print(f"pickle load: {wte.time:6.2f} s")

@@ -0,0 +1,31 @@
+import argparse, time
+from tinygrad.llm.model import Transformer
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--model", required=True, help="path to gguf model")
+  parser.add_argument("--max-context", type=int, default=8192, help="max context length (default: %(default)s)")
+  parser.add_argument("--prompt-tokens", type=int, default=1024, help="number of prompt tokens (default: %(default)s)")
+  parser.add_argument("--decode-tokens", type=int, default=16, help="number of tokens to decode (default: %(default)s)")
+  parser.add_argument("--chunk-size", type=int, default=32, help="chunk size for prefill (default: %(default)s)")
+  args = parser.parse_args()
+
+  st = time.perf_counter()
+  model, _ = Transformer.from_gguf(args.model, args.max_context)
+  print(f"load {time.perf_counter()-st:.3f}s", flush=True)
+
+  st = time.perf_counter()
+  model.warmup()
+  print(f"warm {time.perf_counter()-st:.3f}s", flush=True)
+
+  prompt = [257] + [1000+i%1000 for i in range(args.prompt_tokens-1)]
+  gen = model.generate(prompt, chunk_size=args.chunk_size)
+  st = time.perf_counter()
+  # first token is time-to-first-token; counted as part of prefill
+  output = [next(gen)]
+  pt = time.perf_counter()
+  print(f"prefill {args.prompt_tokens/(pt-st):.3f} tok/s", flush=True)
+
+  for _ in range(args.decode_tokens): output.append(next(gen))
+  et = time.perf_counter()
+  print(f"decode {args.decode_tokens/(et-pt):.3f} tok/s output {output}", flush=True)
@@ -35,7 +35,7 @@ def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], Li
    return name

  for call in iter_kernel_calls(linear):
-    arg_uops = [b for b in call.src[1:] if b.op is not Ops.BIND]
+    arg_uops = [b for b in call.src[1:] if not b.is_bound_var]
    prg = to_program(call.src[0], Device[arg_uops[0].device].renderer)
    info = prg.arg
    functions[info.function_name] = prg.src[2].arg
@@ -241,8 +241,8 @@ export default {model_name};
 def export_model(model, target:str, *inputs, model_name: Optional[str] = "model", stream_weights=False):
  assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, f"only {', '.join(EXPORT_SUPPORTED_DEVICE)} are supported"

-  # NOTE: CPU_COUNT=1, since export does not support threading
-  with Context(JIT=2, CPU_COUNT=1): linear, output_bufs = jit_model(model, *inputs)
+  # NOTE: NUM_CPU_THREADS=1, since export does not support threading
+  with Context(JIT=2, NUM_CPU_THREADS=1): linear, output_bufs = jit_model(model, *inputs)
  functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs)
  state = get_state_dict(model)
  weight_names = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None}
@@ -264,7 +264,7 @@ def export_model(model, target:str, *inputs, model_name: Optional[str] = "model"
        if getattr(dim, "op", None) is Ops.ADD and len(dim.src) == 2 and \
           any(s.op is Ops.PARAM and s.addrspace is AddrSpace.ALU for s in dim.src) and any(s.op is Ops.CONST for s in dim.src):
          name, val = dim.src if dim.src[1].op is Ops.CONST else reversed(dim.src)
-          global_size[j] = f"_{name.expr}[0] + {val.arg}"
+          global_size[j] = f"_{name.expr}[0] + {val.val}"

  prg = ""
  if target == "clang":
@@ -18,9 +18,9 @@ def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
  SEQ = inp.shape[1]
  OUT = weight.shape[0]
  IN = weight.shape[-1]
-  seq_idx = UOp.range(SEQ, 2, AxisType.LOOP)
-  out_idx = UOp.range(OUT, 3, AxisType.LOOP)
-  batch_idx = UOp.range(output.size//SEQ//OUT, 1, AxisType.LOOP)
+  seq_idx = UOp.range(SEQ, 2)
+  out_idx = UOp.range(OUT, 3)
+  batch_idx = UOp.range(output.size//SEQ//OUT, 1)
  reduce_idx = UOp.range(IN, 0, AxisType.REDUCE)
  product = (inp.index((seq_idx*IN+reduce_idx+batch_idx*IN*SEQ)) * weight.index((out_idx*IN+reduce_idx))).cast(dtypes.float)
  reduced = product.reduce(reduce_idx, arg=Ops.ADD)
@@ -53,7 +53,7 @@ class FP8Linear:
    x_fp8, x_scale = quantize_to_fp8(x)
    GPUS = self.weight.device
    if isinstance(GPUS, tuple) and len(GPUS) > 1:
-      y = Tensor(Tensor.empty((batch//len(GPUS), seq, self.weight.shape[0]), dtype=dtypes.float, device=GPUS).uop.multi(0), device=GPUS)
+      y = Tensor(Tensor.empty((batch//len(GPUS), seq, self.weight.shape[0]), dtype=dtypes.float, device=GPUS).uop.unshard(0), device=GPUS)
    else:
      y = Tensor.empty((batch, seq, self.weight.shape[0]), dtype=dtypes.float)
    y = Tensor.custom_kernel(y, x_fp8, w_fp8, fxn=custom_matmul, grad_fxn=custom_matmul_backward)[0]
@@ -58,8 +58,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
  B_copy = B_local.permute((1,0)) if use_wmma else B_local
  A_store = A_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(a[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
  B_store = B_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(b[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
-  barrier = UOp.barrier(A_store, B_store)
-  A_local, B_local = A_local.after(barrier), B_local.after(barrier)
+  # NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
+  A_local, B_local = A_local.after(A_store, B_store), B_local.after(A_store, B_store)

  # -- COMPUTE --
  lane_m, lane_n = lane // LANES_PER_WAVE_N, lane % LANES_PER_WAVE_N
@@ -70,8 +70,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:

  if use_wmma:
    k = UOp.range(BLOCK_K // WMMA_K, 101, AxisType.REDUCE)
-    tile_m = UOp.range(TM // WMMA_ACC, 200, AxisType.LOOP)
-    tile_n = UOp.range(TN, 201, AxisType.LOOP)
+    tile_m = UOp.range(TM // WMMA_ACC, 200)
+    tile_n = UOp.range(TN, 201)

    acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0,2,1)[tile_m, tile_n]
    a_frag = A_local.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_K // WMMA_K, WMMA_K)[wave_m, tile_m, lane_n, k]
@@ -96,8 +96,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
    b_frag = b_frag.reshape(1, TN).expand(TM, TN)
    acc_store = acc.store(acc.after(k) + (a_frag * b_frag))

-  # store accumulator and loop
-  acc = acc.after(acc_store.end(k).barrier().end(k_tile))
+  # store accumulator and loop (the barrier at the end of the loop is implied by the LOCAL buffers stored and loaded in the loop)
+  acc = acc.after(acc_store.end(k).end(k_tile))

  # store accumulator to output (unified)
  c = c.reshape(WAVES_M, TM//UNROLL_M, LANES_PER_WAVE_M, UNROLL_M,
@@ -1,203 +0,0 @@
-from tinygrad import Tensor, UOp, getenv
-from tinygrad.uop.ops import AxisType, KernelInfo, Ops
-from tinygrad.dtype import AddrSpace, dtypes
-from tinygrad.helpers import DEBUG, GlobalCounters, Context
-import math
-
-BLOCK_M, BLOCK_N = 64, 64
-WARP_SIZE = 32
-WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
-WAVES_M, WAVES_N = 4, 1
-LANES_PER_WAVE_M, LANES_PER_WAVE_N = 2, 16
-WMMA_ACC = WMMA_M // LANES_PER_WAVE_M
-THREADS_PER_BLOCK = WARP_SIZE * WAVES_M * WAVES_N
-LDS_PAD = 4  # pad LDS rows to reduce bank conflicts
-
-WMMA_ARG = (WMMA_M, WMMA_N, WMMA_K), 'AMD', 32
-LOG2E = math.log2(math.e)
-
-def warp_shfl_xor(val, offset, lane):
-  """Read val from lane ^ offset using ds_bpermute."""
-  idx = ((lane ^ offset) * 4).cast(dtypes.int)
-  if val.op is Ops.INDEX and val.addrspace == AddrSpace.REG: val = val.load()
-  return UOp(Ops.CUSTOM, dtypes.float, (idx, val),
-             arg="__builtin_bit_cast(float, __builtin_amdgcn_ds_bpermute({0}, __builtin_bit_cast(int, {1})))")
-
-def warp_reduce_max(val, lane):
-  """Tree reduce MAX across LANES_PER_WAVE_N=16 lanes."""
-  for offset in [8, 4, 2, 1]:
-    val = UOp(Ops.MAX, dtypes.float, (val, warp_shfl_xor(val, offset, lane)))
-  return val
-
-def warp_reduce_sum(val, lane):
-  """Tree reduce SUM across LANES_PER_WAVE_N=16 lanes."""
-  for offset in [8, 4, 2, 1]:
-    val = val + warp_shfl_xor(val, offset, lane)
-  return val
-
-def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
-  # inputs are (B*H, N, D)
-  BH, N, D = q.shape
-  assert N % BLOCK_M == 0 and N % BLOCK_N == 0, f"N={N} must be divisible by BLOCK_M={BLOCK_M} and BLOCK_N={BLOCK_N}"
-  assert D % WMMA_K == 0 and D % LANES_PER_WAVE_N == 0, f"D={D} must be divisible by WMMA_K={WMMA_K} and LANES_PER_WAVE_N={LANES_PER_WAVE_N}"
-  assert BLOCK_M % (WAVES_M * WMMA_M) == 0 and BLOCK_N % LANES_PER_WAVE_N == 0
-  TM = BLOCK_M // (WAVES_M * LANES_PER_WAVE_M)
-  TN = BLOCK_N // (WAVES_N * LANES_PER_WAVE_N)
-  TD = D // (WAVES_N * LANES_PER_WAVE_N)
-  SCALE = 1.0 / math.sqrt(D)
-
-  block_bh = UOp.range(BH, 0, AxisType.GLOBAL)
-  block_m = UOp.range(N // BLOCK_M, 1, AxisType.GLOBAL)
-
-  q = q.reshape(BH, N//BLOCK_M, BLOCK_M, D)[block_bh, block_m]
-  k = k.reshape(BH, N//BLOCK_N, BLOCK_N, D)[block_bh]
-  v = v.reshape(BH, N//BLOCK_N, BLOCK_N, D)[block_bh]
-  o = o.reshape(BH, N//BLOCK_M, BLOCK_M, D)[block_bh, block_m]
-
-  wave_m = UOp.range(WAVES_M, 2, AxisType.LOCAL)
-  wave_n = UOp.range(WAVES_N, 3, AxisType.LOCAL)
-  lane = UOp.range(WARP_SIZE, -1, AxisType.WARP)
-  tid = (wave_m * WAVES_N + wave_n) * WARP_SIZE + lane
-  lane_m = lane // LANES_PER_WAVE_N
-  lane_n = lane % LANES_PER_WAVE_N
-
-  # LDS allocation: slot 0 = Q then P (shared), slot 1 = K then V
-  # TODO: the memory planner should be able to find this reuse
-  ELEMS_PER_THREAD = BLOCK_M * D // THREADS_PER_BLOCK
-  QP_lds = UOp.placeholder((BLOCK_M, D + LDS_PAD), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL)
-  KV_lds = UOp.placeholder((BLOCK_N, D + LDS_PAD), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL)[:, :D]
-
-  # register state
-  acc = UOp.placeholder((TM, TD), dtypes.float, slot=2, addrspace=AddrSpace.REG)
-  m_i = UOp.placeholder((TM,), dtypes.float, slot=3, addrspace=AddrSpace.REG)
-  l_i = UOp.placeholder((TM,), dtypes.float, slot=4, addrspace=AddrSpace.REG)
-  acc = acc.after(acc.store(acc.const_like(0)))
-  m_i = m_i.after(m_i.store(m_i.const_like(-math.inf)))
-  l_i = l_i.after(l_i.store(l_i.const_like(0)))
-
-  # ====== KV tile loop ======
-  n_tile = UOp.range(N // BLOCK_N, 100, AxisType.REDUCE)
-
-  # load Q + K into LDS (Q reloaded each iteration since P overwrites slot 0)
-  Q_lds = QP_lds[:, :D]
-  Q_store = Q_lds.after(n_tile).reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
-    q.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
-  K_store = KV_lds.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
-    k[n_tile].reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
-  qk_load_barrier = UOp.barrier(UOp.group(Q_store, K_store))
-  Q_lds = Q_lds.after(qk_load_barrier)
-  KV_lds_k = KV_lds.after(qk_load_barrier)
-
-  # -- S = Q @ K^T via WMMA (re-init each n_tile) --
-  S_reg = UOp.placeholder((TM, TN), dtypes.float, slot=6, addrspace=AddrSpace.REG)
-  S_reg = S_reg.after(S_reg.after(n_tile).store(S_reg.const_like(0)))
-  k_qk = UOp.range(D // WMMA_K, 101, AxisType.REDUCE)
-  tm1 = UOp.range(TM // WMMA_ACC, 200, AxisType.LOOP)
-  tn1 = UOp.range(TN, 201, AxisType.LOOP)
-  S_frag = S_reg.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0, 2, 1)[tm1, tn1]
-  q_frag = Q_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, D // WMMA_K, WMMA_K)[wave_m, tm1, lane_n, k_qk]
-  k_frag = KV_lds_k.reshape(WAVES_N, TN, WMMA_N, D // WMMA_K, WMMA_K)[wave_n, tn1, lane_n, k_qk]
-  qk = UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), *WMMA_ARG)
-  qk_done = S_frag.store(qk).end(tm1, tn1).end(k_qk)
-  S_reg = S_reg.after(qk_done)
-
-  # -- softmax in registers with warp shuffles --
-  S_reg = S_reg.after(S_reg.store(S_reg * SCALE))
-
-  # per-thread local row max over TN=4 elements, then warp reduce across 16 lanes
-  m_ij = UOp.placeholder((TM,), dtypes.float, slot=7, addrspace=AddrSpace.REG)
-  m_ij = m_ij.after(m_ij.after(n_tile).store(m_ij.const_like(-math.inf)))
-  rm2 = UOp.range(TN, 261, AxisType.REDUCE)
-  m_ij = m_ij.after(m_ij.store(m_ij.after(rm2).maximum(S_reg[:, rm2])).end(rm2))
-  # warp reduce max (in-place)
-  ri_w = UOp.range(TM, 270, AxisType.LOOP)
-  m_ij = m_ij.after(m_ij[ri_w].store(warp_reduce_max(m_ij[ri_w], lane)).end(ri_w))
-
-  # compute P = exp(S - m_ij) in S_reg
-  S_reg = S_reg.after(S_reg.store(((S_reg - m_ij.reshape(TM, 1).expand(TM, TN)) * LOG2E).exp2()))
-
-  p_local = UOp.placeholder((TM,), dtypes.float, slot=8, addrspace=AddrSpace.REG)
-  p_local = p_local.after(p_local.after(n_tile).store(p_local.const_like(0)))
-  rp2 = UOp.range(TN, 291, AxisType.REDUCE)
-  p_local = p_local.after(p_local.store(p_local.after(rp2) + S_reg[:, rp2]).end(rp2))
-  ri_ws = UOp.range(TM, 295, AxisType.LOOP)
-  p_sum = p_local.after(p_local[ri_ws].store(warp_reduce_sum(p_local[ri_ws], lane)).end(ri_ws))
-
-  # write P = exp(S - m_ij) to P_lds (reuses slot 0, Q no longer needed)
-  P_lds = QP_lds[:, :BLOCK_N]
-  P_write = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_ACC, LANES_PER_WAVE_M, WAVES_N, TN, LANES_PER_WAVE_N)
-  P_write = P_write.permute((0, 4, 3, 6, 1, 2, 5)).reshape(THREADS_PER_BLOCK, TM, TN)
-  P_store = P_write[tid].store(S_reg.cast(dtypes.half))
-
-  # -- online softmax correction --
-  ri4 = UOp.range(TM, 330, AxisType.LOOP)
-  m_new_val = m_i[ri4].maximum(m_ij[ri4])
-  alpha_val = ((m_i[ri4] - m_new_val) * LOG2E).exp2()
-  beta_val = ((m_ij[ri4] - m_new_val) * LOG2E).exp2()
-  rj4 = UOp.range(TD, 331, AxisType.LOOP)
-  correction = UOp.group(
-    acc[ri4, rj4].store(alpha_val * acc[ri4, rj4]).end(rj4),
-    l_i[ri4].store(alpha_val * l_i[ri4] + beta_val * p_sum[ri4]),
-    m_i[ri4].store(m_new_val),
-  ).end(ri4)
-  acc = acc.after(correction)
-  l_i = l_i.after(correction)
-  m_i = m_i.after(correction)
-
-  # load V into KV_lds (must wait for QK WMMA to finish reading K from KV_lds)
-  V_store = KV_lds.after(qk_done).reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
-    v[n_tile].reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
-  pv_barrier = UOp.barrier(UOp.group(P_store, V_store))
-  P_lds = P_lds.after(pv_barrier)
-  KV_lds_v = KV_lds.after(pv_barrier)
-
-  # -- acc += P @ V via WMMA --
-  k_pv = UOp.range(BLOCK_N // WMMA_K, 400, AxisType.REDUCE)
-  tm2 = UOp.range(TM // WMMA_ACC, 401, AxisType.LOOP)
-  tn2 = UOp.range(TD, 402, AxisType.LOOP)
-  acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TD).permute(0, 2, 1)[tm2, tn2]
-  p_frag = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_N // WMMA_K, WMMA_K)[wave_m, tm2, lane_n, k_pv]
-  v_frag = KV_lds_v.reshape(WAVES_N, TD, WMMA_N, BLOCK_N // WMMA_K, WMMA_K)[wave_n, tn2, lane_n, k_pv]
-  pv = UOp.wmma(p_frag, v_frag, acc_frag.after(k_pv), *WMMA_ARG)
-
-  # end KV tile loop
-  n_tile_end = acc_frag.store(pv).end(tm2, tn2).end(k_pv).barrier().end(n_tile)
-  acc = acc.after(n_tile_end)
-  l_i = l_i.after(n_tile_end)
-  m_i = m_i.after(n_tile_end)
-
-  # normalize: acc /= l_i
-  acc = acc.after(acc.store(acc * (1 / l_i).reshape(TM, 1).expand(TM, TD)))
-
-  # store output
-  o = o.reshape(WAVES_M, TM // WMMA_ACC, WMMA_ACC, LANES_PER_WAVE_M, WAVES_N, TD, LANES_PER_WAVE_N)
-  o = o.permute((0, 4, 3, 6, 1, 2, 5)).reshape(THREADS_PER_BLOCK, TM, TD)
-  return o[tid].store(acc).end(wave_m, wave_n, lane).end(block_m, block_bh).sink(arg=KernelInfo(opts_to_apply=()))
-
-if __name__ == "__main__":
-  B, H, N, D = getenv("B", 1), getenv("H", 32), getenv("N", 1024), getenv("D", 64)
-  q = Tensor.rand(B, H, N, D).cast(dtypes.half)
-  k = Tensor.rand(B, H, N, D).cast(dtypes.half)
-  v = Tensor.rand(B, H, N, D).cast(dtypes.half)
-  o = Tensor.empty(B, H, N, D, dtype=dtypes.float)
-  with Context(DEBUG=0): Tensor.realize(q, k, v)
-
-  q_flat, k_flat, v_flat, o_flat = q.reshape(B*H, N, D), k.reshape(B*H, N, D), v.reshape(B*H, N, D), o.reshape(B*H, N, D)
-  NUM_RUNS = getenv("CNT", 5)
-  ets = []
-  with Context(DEBUG=2):
-    for _ in range(NUM_RUNS):
-      GlobalCounters.reset()
-      tst = Tensor.custom_kernel(o_flat, q_flat, k_flat, v_flat, fxn=amd_flash_attention)[0].realize()
-      ets.append(GlobalCounters.time_sum_s)
-  print(f"best time: {min(ets)*1e3:.2f}ms")
-
-  if getenv("VERIFY", 1):
-    with Context(DEBUG=0):
-      ref = q.float().scaled_dot_product_attention(k.float(), v.float()).reshape(B*H, N, D).realize()
-      err = (ref - tst).square().mean().item()
-    print(f"mean squared error {err}")
-    if err > 1e-2:
-      raise RuntimeError("flash attention is wrong!")
-    else:
-      print("flash attention is correct!")
@@ -28,10 +28,10 @@ REG_TILES_PER_WAVE_M = BLOCK_M // (WAVES_PER_BLOCK_M * LANES_PER_WAVE_M * TM)
 assert WAVES_PER_BLOCK_M*REG_TILES_PER_WAVE_M*LANES_PER_WAVE_M*TM == BLOCK_M, "M reshape is wrong"
 assert WAVES_PER_BLOCK_N*REG_TILES_PER_WAVE_N*LANES_PER_WAVE_N*TN == BLOCK_N, "N reshape is wrong"

-def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.LOOP): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
+def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.WEAK): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
 def copy(dest:UOp, src:UOp, rng:int, upcast=False):
  assert dest.shape == src.shape
-  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.LOOP)
+  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.WEAK)
  return dest[*rngs].store(src[*rngs]).end(*rngs)

 def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
@@ -66,9 +66,8 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
  B_local = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
  B_local_store = copy(B_local.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)

-  # TODO: can we automate barrier?
-  barrier = UOp.barrier(A_local_store, B_local_store)
-  A_local, B_local = A_local.after(barrier), B_local.after(barrier)
+  # NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
+  A_local, B_local = A_local.after(A_local_store, B_local_store), B_local.after(A_local_store, B_local_store)

  # open inner k range
  k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)
@@ -102,7 +101,7 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
  sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iter_m, t_m] * B_row[iter_n, t_n]).end(iter_m, iter_n, t_m, t_n)

  # Close k, sync, and close K tiles
-  sink = sink.end(k).barrier().end(k_tile_range)
+  sink = sink.end(k).end(k_tile_range)

  # ---------------------------
  # REG -> GLOBAL (epilogue)
@@ -1,10 +1,12 @@
-import atexit, functools, pathlib
+import atexit, functools, math, pathlib
 from tinygrad import Tensor, Device, dtypes
+from tinygrad.dtype import AddrSpace
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
 from tinygrad.renderer import Estimates
-from tinygrad.helpers import getenv, all_same, DEBUG
+from tinygrad.helpers import getenv, all_same, DEBUG, ceildiv
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8
+from extra.llama_kernels.quantize_mxfp4 import quantize_mxfp4

 TILE_M, TILE_N, TILE_K = 256, 256, 64

@@ -72,7 +74,7 @@ def hk_fp8_atb_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, g_amax:Tensor|
    elif b.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M, N // ndev, dtype=dtypes.bfloat16, device=a.device), 2
    elif a.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M // ndev, N, dtype=dtypes.bfloat16, device=a.device), 1
    else: inv, out_axis, reduce_out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a.device), 0, True
-    out = Tensor(inv.uop.multi(out_axis), device=a.device)
+    out = Tensor(inv.uop.unshard(out_axis), device=a.device)
    dname = a.device[0]
  else:
    out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a.device)
@@ -107,6 +109,43 @@ def custom_hk_mxfp8_gemm(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, *extra:U
  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
                               UOp(Ops.BINARY, arg=lib)))

+# ** MXFP4 GEMM custom kernel
+
+@functools.cache
+def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp, tile_m:int, tile_n:int) -> UOp:
+  from extra.gemm.gemm_mxfp4 import build_kernel
+  M, half_k = math.prod(A.shape[:-1]), A.shape[-1]
+  N, half_k_b = math.prod(B.shape[:-1]), B.shape[-1]
+  K = half_k * 2
+  assert half_k == half_k_b and math.prod(C.shape[:-1]) == M and C.shape[-1] == N
+  threads = UOp.special(256, "lidx0")
+  groups_x, groups_y = UOp.special(ceildiv(N, tile_n), "gidx0"), UOp.special(ceildiv(M, tile_m), "gidx1")
+  lds = UOp.placeholder((163840,), dtypes.uint8, 0, AddrSpace.LOCAL)
+  sink = UOp.sink(C.base, A.base, B.base, scale_a.base, scale_b.base, *(x.base for x in extra), lds, threads, groups_x, groups_y,
+                  arg=KernelInfo(f"mxfp4_gemm_{M}_{N}_{K}",
+                                 estimates=Estimates(ops=2*M*N*K, mem=(M*half_k+N*half_k)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
+  insts = build_kernel(M, N, K, tile_m, tile_n)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in insts))))
+
+def _mxfp4_gemm_quantized(a_q:Tensor, b_q:Tensor, scale_a:Tensor, scale_b:Tensor) -> Tensor:
+  M, half_k = a_q.shape
+  N, half_k_b = b_q.shape
+  assert half_k == half_k_b
+  is_multi = isinstance(a_q.device, tuple)
+  reduce_out = is_multi and (a_q.uop.axis == 1 or b_q.uop.axis == 1)
+  if not is_multi: out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a_q.device)
+  elif reduce_out: out = Tensor(Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a_q.device).uop.unshard(0), device=a_q.device)
+  elif a_q.uop.axis == 0:
+    out = Tensor(Tensor.invalids(1, M//len(a_q.device), N, dtype=dtypes.bfloat16, device=a_q.device).uop.unshard(1), device=a_q.device)
+  elif b_q.uop.axis == 0:
+    out = Tensor(Tensor.invalids(1, M, N//len(a_q.device), dtype=dtypes.bfloat16, device=a_q.device).uop.unshard(2), device=a_q.device)
+  else: out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a_q.device)
+  tile_m, tile_n = next((tm, tn) for tm, tn in ((256, 256), (192, 256), (128, 512)) if M % tm == N % tn == 0)
+  out = Tensor.custom_kernel(out, a_q, b_q, scale_a, scale_b,
+                             fxn=functools.partial(custom_mxfp4_gemm, tile_m=tile_m, tile_n=tile_n))[0]
+  if reduce_out: out = out.sum(0)
+  return out.squeeze(0)
+
 def quantize_mxfp8(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
  # 1x32 block scaling along the last axis
  *batch, K = x.shape
@@ -171,13 +210,13 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
  M, K = A.shape[0]*A.shape[1], A.shape[2]
  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
  assert K == K2
-  m = UOp.range(M, 1, AxisType.LOOP)
-  n = UOp.range(N, 2, AxisType.LOOP)
+  m = UOp.range(M, 1)
+  n = UOp.range(N, 2)
  k = UOp.range(K, 0, AxisType.REDUCE)
-  mul = (A.flatten().index((m*UOp.const(dtypes.weakint, K)+k))*
-         B.flatten().index((k*UOp.const(dtypes.weakint, N)+n))).cast(dtypes.float32)
+  mul = (A.flatten().index((m*UOp.const(K)+k))*
+         B.flatten().index((k*UOp.const(N)+n))).cast(dtypes.float32)
  red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype)
-  store = C.flatten().index((m*UOp.const(dtypes.weakint, N)+n)).store(red).end(m, n)
+  store = C.flatten().index((m*UOp.const(N)+n)).store(red).end(m, n)
  return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))

 # ** bf16 A @ B.T kernel in C
@@ -234,7 +273,7 @@ def hk_bf16_atb_gemm(a:Tensor, b:Tensor) -> Tensor:
    elif b.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M, N // ndev, dtype=a.dtype, device=a.device), 2
    elif a.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M // ndev, N, dtype=a.dtype, device=a.device), 1
    else: inv, out_axis, reduce_out = Tensor.invalids(1, M, N, dtype=a.dtype, device=a.device), 0, True
-    out = Tensor(inv.uop.multi(out_axis), device=a.device)
+    out = Tensor(inv.uop.unshard(out_axis), device=a.device)
    dname = a.device[0]
  else:
    out = Tensor.invalids(1, M, N, dtype=a.dtype, device=a.device)
@@ -244,7 +283,6 @@ def hk_bf16_atb_gemm(a:Tensor, b:Tensor) -> Tensor:
  if reduce_out: out = out.sum(0)
  return out.squeeze(0) if out.ndim == 3 else out

-
 # ** backward gemm, might use the asm gemm

 def custom_gemm_bw(gradient:UOp, kernel:UOp, n_scales:int=2, has_grad_amax:bool=False, has_w_post:bool=False):
@@ -341,13 +379,30 @@ def custom_mx_gemm_bw(gradient:UOp, kernel:UOp, has_w_post:bool, w_stored:bool=F
  if wp is not None: grad_b = grad_b / wp.reshape(-1, 1)
  return (None, grad_a.uop, grad_b.uop) + tuple(None for _ in inputs[3:])

+# ** mxfp4 gemm backward
+
+def custom_mxfp4_gemm_bw(gradient:UOp, kernel:UOp):
+  inputs = kernel.src[1:]  # out, row operands/scales, BF16 operands, column operands/scales
+  assert len(inputs) == 11
+  a, w = Tensor(inputs[5], device=inputs[5].device), Tensor(inputs[6], device=inputs[6].device)
+  a_col, scale_a_col = Tensor(inputs[7], device=a.device), Tensor(inputs[8], device=a.device)
+  w_col, scale_w_col = Tensor(inputs[9], device=a.device), Tensor(inputs[10], device=a.device)
+  g = Tensor(gradient, device=a.device)[:a.shape[0]].cast(dtypes.bfloat16)
+  g_row, scale_g_row, g_col, scale_g_col = quantize_mxfp4(g, flatten_row=True)
+  grad_a = _mxfp4_gemm_quantized(g_row, w_col, scale_g_row, scale_w_col).reshape(*a.shape[:-1], w.shape[-1])
+  grad_w = _mxfp4_gemm_quantized(g_col, a_col, scale_g_col, scale_a_col).reshape(w.shape)
+  return (None, None, None, None, None, grad_a.uop, grad_w.uop, None, None, None, None)
+
 # ** main gemm function

 def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=None, grad_amax_state:Tensor|None=None,
             next_grad_amax_state:Tensor|None=None,
             w_post_scale:Tensor|None=None, mx:bool=False, mx_scales:tuple|None=None, mx_w_stored:bool=False, g_amax:Tensor|None=None,
-             a_pretranspose:Tensor|None=None) -> Tensor:
+             a_pretranspose:Tensor|None=None, mxfp4:bool=False) -> Tensor:
  assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
+  if mxfp4:
+    assert not mx and mx_scales is None, "mxfp4 owns quantization; mx/mx_scales are for mxfp8"
+    assert a.dtype == dtypes.bfloat16, f"cannot quantize {a.dtype} to mxfp4"
  counters["used"] += 1
  unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0
  if unfold_batch:
@@ -355,7 +410,7 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N
    a = a.reshape(a.shape[0]*a.shape[1], a.shape[2])
  squeeze = a.ndim == 2
  if squeeze: a = a.unsqueeze(0)
-  out_dtype = dtypes.bfloat16 if a.dtype == FP8_DTYPE else a.dtype
+  out_dtype = dtypes.bfloat16 if a.dtype == FP8_DTYPE or mxfp4 else a.dtype

  batch, M, K = a.shape
  N = b.shape[1]
@@ -366,11 +421,11 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N

  if is_multi:
    if n_sharded:
-      out = Tensor(Tensor.invalids(batch, M, N//len(a.device), dtype=out_dtype, device=a.device).uop.multi(2), device=a.device)
+      out = Tensor(Tensor.invalids(batch, M, N//len(a.device), dtype=out_dtype, device=a.device).uop.unshard(2), device=a.device)
    elif m_sharded:
-      out = Tensor(Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device).uop.multi(1), device=a.device)
+      out = Tensor(Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device).uop.unshard(1), device=a.device)
    else:
-      out = Tensor(Tensor.invalids(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=out_dtype, device=a.device).uop.multi(0),
+      out = Tensor(Tensor.invalids(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=out_dtype, device=a.device).uop.unshard(0),
                   device=a.device)
  else:
    out = Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device)
@@ -378,7 +433,15 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N
  renderer = Device[dname:=(a.device[0] if is_multi else a.device)].renderer
  dname, arch = dname.split(":")[0], renderer.target.arch
  if arch.startswith("gfx950") and getenv("USE_ASM", 1):
-    if mx:
+    if mxfp4:
+      tile_m, tile_n = next((tm, tn) for tm, tn in ((256, 256), (192, 256), (128, 512)) if (batch*M) % tm == N % tn == 0)
+      fxn = functools.partial(custom_mxfp4_gemm, tile_m=tile_m, tile_n=tile_n)
+      w = b.T
+      a_q, scale_a, a_col, scale_a_col = quantize_mxfp4(a, shuffle_col=True)
+      b_q, scale_b, b_col, scale_b_col = quantize_mxfp4(w, shuffle_row=True, shuffle_col=True)
+      out = Tensor.custom_kernel(out, a_q, b_q, scale_a, scale_b, a, w,
+                                 a_col, scale_a_col, b_col, scale_b_col, fxn=fxn, grad_fxn=custom_mxfp4_gemm_bw)[0]
+    elif mx:
      # mxfp8 1x32 block scaling
      if mx_scales is not None:
        a_si, a_e8, b_si, b_e8 = mx_scales
@@ -79,7 +79,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
  # this is the big accumulator
  acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float, 0, AddrSpace.REG)
  assert acc.size*WARP_SIZE*WARPGROUP_SIZE*4 == BLOCK_M*BLOCK_N
-  acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const(dtypes.float, (0.0,)*4), end=init_l)
+  acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const((0.0,)*4, dtypes.float), end=init_l)

  # create locals (note A is permuted, and the stride is changed to avoid bank conflicts)
  def make_locals(slot) -> tuple[UOp, UOp]:
@@ -29,7 +29,7 @@ TID_SIZE = WARPGROUP_SIZE*WARP_SIZE

 def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=()):
  assert dest.shape == src.shape
-  rngs = [UOp.range(s, rng+i, AxisType.UPCAST if i in upcast else AxisType.LOOP) for i,s in enumerate(src.shape)]
+  rngs = [UOp.range(s, rng+i, AxisType.UPCAST if i in upcast else AxisType.WEAK) for i,s in enumerate(src.shape)]
  copy = dest[*rngs].store(src[*rngs]).end(*rngs)
  return dest.after(copy) if set else copy

@@ -1,10 +1,32 @@
 import functools, pathlib
 from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.helpers import getenv
 from tinygrad.renderer import Estimates
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 from extra.gemm.cdna_asm_gemm import quantize_mxfp8, _mx_block_scale, _mx_block_scale_3d

+ZERO_OPTIM = getenv("ZERO_OPTIM", 0)
+
+def reduce_scatter_devaxis(out:Tensor, shard_axis:int=0) -> Tensor:
+  # out: sharded on the device axis, shape (ndev, *rest); return the device-axis sum left sharded on shard_axis.
+  u = out.uop
+  devs, rest = u.device, u.shape[1:]
+  assert rest[shard_axis] % len(devs) == 0, f"reduce_scatter needs even shards: {rest[shard_axis]} % {len(devs)}"
+  # reach the raw per-device buffer below the UNSHARD, keeping the AFTERs so reads stay ordered after the kernel writes
+  node, barriers = u, []
+  while node.op is not Ops.UNSHARD:
+    if node.op is Ops.AFTER: barriers += node.src[1:]
+    node = node.src[0]
+  mbuf = node.src[0].after(*barriers) if barriers else node.src[0]
+  sz = rest[shard_axis] // len(devs)
+  shards = []
+  for i in range(len(devs)):
+    bounds = tuple((0,s) if a != shard_axis else (i*sz,(i+1)*sz) for a,s in enumerate(rest))
+    contribs = [mbuf.mselect(j).reshape(rest).shrink(bounds).copy_to_device(devs[i]) for j in range(len(devs))]
+    shards.append(functools.reduce(lambda a,b: a.alu(Ops.ADD, b), contribs))
+  return Tensor(UOp.mstack(*shards).unshard(shard_axis, UOp.range(len(devs), -1, AxisType.DEVICE)), device=devs)
+
@functools.cache
 def custom_hk_grouped_mxfp8_gemm(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, *extra:UOp, dname:str, n_experts:int) -> UOp:
  M, K = A.shape
@@ -55,10 +77,11 @@ def grouped_mx_wgrad(g:Tensor, xg:Tensor, expert_off:Tensor, n_experts:int) -> T
  dname = (g.device[0] if isinstance(g.device, tuple) else g.device).split(":")[0]
  is_multi = isinstance(g.device, tuple)
  inv = Tensor.invalids(1, n_experts * N, K, dtype=dtypes.bfloat16, device=g.device)
-  out = Tensor(inv.uop.multi(0), device=g.device) if is_multi else inv
+  out = Tensor(inv.uop.unshard(0), device=g.device) if is_multi else inv
  out = Tensor.custom_kernel(out, gT, xT, g_si, x_si, expert_off,
                             fxn=functools.partial(custom_hk_grouped_mxfp8_wgrad, dname=dname, n_experts=n_experts))[0]
-  out = out.sum(0) if is_multi else out.squeeze(0)
+  if is_multi and ZERO_OPTIM: out = reduce_scatter_devaxis(out, 0)
+  else: out = out.sum(0) if is_multi else out.squeeze(0)
  return out.reshape(n_experts, N, K)

 def mx_pack_3d(e8:Tensor) -> Tensor:
@@ -103,7 +126,7 @@ def grouped_mx_gemm(x:Tensor, w:Tensor|tuple[Tensor, Tensor], expert_off:Tensor)
  if isinstance(x.device, tuple) and (row_axis := x.uop.axis) is not None:
    ndev = len(x.device)
    out = Tensor(Tensor.invalids(*(s // ndev if i == row_axis else s for i, s in enumerate(out_shape)),
-                                 dtype=dtypes.bfloat16, device=x.device).uop.multi(row_axis), device=x.device)
+                                 dtype=dtypes.bfloat16, device=x.device).uop.unshard(row_axis), device=x.device)
  else:
    out = Tensor.invalids(*out_shape, dtype=dtypes.bfloat16, device=x.device)
  return Tensor.custom_kernel(out, x_q, w_q, x_si, w_si, xe_in, w_e8, expert_off,
@@ -0,0 +1,131 @@
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+
+BLOCK_ROW = 256
+
+def _sharded_invalids(shape:tuple[int, ...], dtype, device) -> Tensor:
+  if isinstance(device, tuple):
+    per = Tensor.invalids(shape[0]//len(device), *shape[1:], dtype=dtype, device=device)
+    return Tensor(per.uop.unshard(0), device=device)
+  return Tensor.invalids(*shape, dtype=dtype, device=device)
+
+def _atomic_add(device:str) -> str:
+  return "__hip_atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);" if device == "AMD" \
+    else "__atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED);"
+
+def _blk_for(D:int) -> int:
+  blk = 64
+  while D % blk: blk //= 2
+  return blk
+
+def _kv_ranges(G, N, D, BLK):
+  g = UOp.range(G, 0)
+  m = UOp.range(N, 1)
+  jo = UOp.range(D // BLK, 2)
+  ji = UOp.range(BLK, 3, AxisType.LOCAL)
+  return g, m, jo * BLK + ji, jo, ji
+
+def _ggather_fwd_kernel(out:UOp, table:UOp, idx:UOp) -> UOp:
+  G, M, D = out.shape
+  g, m, j, jo, ji = _kv_ranges(G, M, D, _blk_for(D))
+  row = idx.index(g, m).cast(dtypes.weakint)
+  val = table.index(g, row, j).load()
+  return out.index(g, m, j).store(val).end(g, m, jo, ji).sink(
+    arg=KernelInfo(name=f"ggather_fwd_{M}_{D}", opts_to_apply=()))
+
+def _ggather_zero_kernel(out:UOp) -> UOp:
+  i = UOp.range(out.numel(), 0)
+  return out.flatten().index(i).store(UOp.const(0.0, out.dtype)).end(i).sink(arg=KernelInfo(name="ggather_zero"))
+
+def _sharded_zeros(shape:tuple[int, ...], dtype, device) -> Tensor:
+  return Tensor.custom_kernel(_sharded_invalids(shape, dtype, device), fxn=_ggather_zero_kernel)[0]
+
+def _ggather_bwd(gradient:UOp, kernel:UOp) -> tuple:
+  _, table_u, idx_u = kernel.src[1:4]
+  dev = table_u.device
+  device = (dev[0] if isinstance(dev, tuple) else dev).split(":")[0]
+  G, R, D = table_u.shape
+  gt = _sharded_zeros((G, R, D), dtypes.float32, dev)
+  go = Tensor(gradient, device=dev)
+  atomic_str = _atomic_add(device)
+  def _bwd_kernel(gtab:UOp, gout:UOp, idx:UOp) -> UOp:
+    Gk, M, Dk = gout.shape
+    g, m, j, jo, ji = _kv_ranges(Gk, M, Dk, _blk_for(Dk))
+    row = idx.index(g, m).cast(dtypes.weakint)
+    val = gout.index(g, m, j).load().cast(dtypes.float32)
+    atomic = UOp(Ops.CUSTOM, src=(gtab.index(g, row, j), val), arg=atomic_str)
+    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)
+
+def grouped_gather_rows(table:Tensor, idx:Tensor, n_groups:int) -> Tensor:
+  G, R, D = table.shape
+  M = idx.shape[1]
+  out = _sharded_invalids((G, M, D), table.dtype, table.device)
+  return Tensor.custom_kernel(out, table, idx, fxn=_ggather_fwd_kernel, grad_fxn=_ggather_bwd)[0]
+
+def _gscatter_fwd_kernel(out:UOp, src:UOp, idx:UOp) -> UOp:
+  G, M, D = out.shape
+  k = idx.shape[1] // src.shape[1]
+  g, m, j, jo, ji = _kv_ranges(G, idx.shape[1], D, _blk_for(D))
+  row = idx.index(g, m).cast(dtypes.weakint)
+  val = src.index(g, (m // k).cast(dtypes.weakint), j).load()
+  return out.index(g, row, j).store(val).end(g, m, jo, ji).sink(
+    arg=KernelInfo(name=f"gscatter_fwd_{idx.shape[1]}_{D}", opts_to_apply=()))
+
+def _gscatter_bwd(gradient:UOp, kernel:UOp) -> tuple:
+  _, src_u, idx_u = kernel.src[1:4]
+  dev = src_u.device
+  G, T_l, D = src_u.shape
+  k = idx_u.shape[1] // T_l
+  sel = grouped_gather_rows(Tensor(gradient, device=dev), Tensor(idx_u, device=dev), G)
+  return (None, sel.reshape(G, T_l, k, D).sum(2).cast(src_u.dtype).uop, None)
+
+def grouped_scatter_rows(src:Tensor, idx:Tensor, m_l:int) -> Tensor:
+  G, T_l, D = src.shape
+  zero = _sharded_zeros((G, m_l, D), src.dtype, src.device)
+  return Tensor.custom_kernel(zero, src, idx, fxn=_gscatter_fwd_kernel, grad_fxn=_gscatter_bwd)[0]
+
+def m_max_for(t_local:int, experts_per_tok:int, n_experts:int) -> int:
+  return (-(-t_local * experts_per_tok // BLOCK_ROW) + n_experts) * BLOCK_ROW
+
+class Routing:
+  def __init__(self, weights:Tensor, dest_row:Tensor, off:Tensor, m_l:int, n_groups:int, t_local:int):
+    self.weights, self.dest_row = weights, dest_row
+    self.off = off
+    self.m_l, self.n_groups, self.t_local = m_l, n_groups, t_local
+
+  @property
+  def rows_e(self) -> Tensor:
+    G, E = self.off.shape[0], self.off.shape[1] - 1
+    tr = Tensor.arange(self.m_l // BLOCK_ROW, dtype=dtypes.int32).reshape(1, -1, 1) * BLOCK_ROW
+    tr = tr.shard(self.off.device) if isinstance(self.off.device, tuple) else tr.to(self.off.device)
+    tile_e = ((tr >= self.off[:, :E].reshape(G, 1, E)).sum(-1) - 1).cast(dtypes.int32)
+    return tile_e.reshape(-1, 1).expand(-1, BLOCK_ROW).reshape(-1)
+
+def n_groups_of(t:Tensor) -> int:
+  return len(t.device) if isinstance(t.device, tuple) else 1
+
+def route(logits:Tensor, experts_per_tok:int, n_experts:int) -> Routing:
+  T, E = logits.shape
+  k, G = experts_per_tok, n_groups_of(logits)
+  assert T % G == 0, f"tokens {T} must split across {G} devices"
+  T_l, m_l = T // G, m_max_for(T // G, k, n_experts)
+
+  topv, topi = logits.reshape(G, T_l, E).topk(k)
+  weights = topv.softmax(-1)
+  m = topi.reshape(G, T_l * k).cast(dtypes.int32).one_hot(E).cast(dtypes.int32)
+
+  pad = ((m.sum(1) + (BLOCK_ROW - 1)) // BLOCK_ROW) * BLOCK_ROW
+  off = pad.cumsum(1).pad(((0, 0), (1, 0)))
+  dest_row = ((m.cumsum(1) + off[:, :E].reshape(G, 1, E)) * m).sum(-1).sub(1).cast(dtypes.int32)
+  return Routing(weights, dest_row, off, m_l, G, T_l)
+
+def dispatch(x:Tensor, r:Routing) -> Tensor:
+  G, D = r.n_groups, x.shape[-1]
+  return grouped_scatter_rows(x.reshape(G, r.t_local, D), r.dest_row, r.m_l).reshape(G * r.m_l, D)
+
+def combine(y:Tensor, r:Routing, n_tokens:int, experts_per_tok:int) -> Tensor:
+  G, D, k = r.n_groups, y.shape[-1], experts_per_tok
+  sel = grouped_gather_rows(y.reshape(G, r.m_l, D), r.dest_row, G).reshape(G, r.t_local, k, D)
+  return (sel * r.weights.reshape(G, r.t_local, k, 1).cast(sel.dtype)).sum(2).reshape(n_tokens, D).cast(y.dtype)
@@ -79,7 +79,7 @@ if __name__ == "__main__":
  linear, var_vals = C.linear_with_vars()
  last_call = linear.src[-1]
  ast = last_call.src[0]
-  bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]
+  bufs = [s.buffer for s in last_call.src[1:] if not s.is_bound_var]

  src = compiled.asm["ptx"]
  # specify the shared memory here so we don't need to do it dynamically
@@ -0,0 +1,139 @@
+"""
+tilelang-style matmul_relu written with tinygrad UOp APIs.
+
+Reference tilelang kernel:
+  @tilelang.jit
+  def matmul_relu(A, B, block_M=64, block_N=64, block_K=64,
+                  dtype=T.float16, accum_dtype=T.float32):
+    M, N, K = T.const('M, N, K')
+    C = T.empty([M, N], dtype)
+    with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
+      A_shared = T.alloc_shared((block_M, block_K), dtype)
+      B_shared = T.alloc_shared((block_K, block_N), dtype)
+      C_local  = T.alloc_fragment((block_M, block_N), accum_dtype)
+      T.clear(C_local)
+      for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
+        T.copy(A[by * block_M, ko * block_K], A_shared)
+        T.copy(B[ko * block_K, bx * block_N], B_shared)
+        T.gemm(A_shared, B_shared, C_local)
+      for i, j in T.Parallel(block_M, block_N):
+        C_local[i, j] = T.max(C_local[i, j], 0)
+      T.copy(C_local, C[by * block_M, bx * block_N])
+    return C
+"""
+
+from tinygrad.dtype import dtypes, AddrSpace, DType
+from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
+from tinygrad.helpers import cdiv, getenv
+from tinygrad.tensor import Tensor
+
+# ---------------------------------------------------------------------------
+# tilelang builtins, expressed with tinygrad UOp APIs
+# ---------------------------------------------------------------------------
+
+def alloc_shared(shape:tuple[int, ...], dtype:DType, slot:int) -> UOp:
+  """T.alloc_shared: one LOCAL buffer shared by all threads in the block."""
+  return UOp.placeholder(tuple(shape), dtype, slot, AddrSpace.LOCAL)
+
+def alloc_fragment(shape:tuple[int, ...], dtype:DType, slot:int, axes:tuple[int, ...], rngs:tuple[UOp, ...]) -> UOp:
+  """T.alloc_fragment: per-thread REG fragment + UNSHARD over the LOCAL thread grid."""
+  assert len(axes) == len(rngs)
+  assert all(tnum.op is Ops.RANGE and tnum.arg[-1] is AxisType.LOCAL for tnum in rngs), "fragments shard over LOCAL ranges"
+  by_axis = dict(zip(axes, rngs))
+  shard_shape = tuple(s // (int(by_axis[i].vmax)+1) if i in by_axis else s for i, s in enumerate(shape))
+  fragment = UOp.placeholder(shard_shape, dtype, slot, AddrSpace.REG)
+  return fragment.unshard(axes, rngs)
+
+# ---------------------------------------------------------------------------
+# GEMM kernel: C = relu(A @ B), float inputs (fp16 or fp32), fp32 fragment accumulator, no WMMA
+# ---------------------------------------------------------------------------
+
+# 64x64 output tile per block, 128 threads as an 8x16 grid; each thread owns an 8x4 fragment sub-tile
+# (the 2-D per-thread layout tilelang infers for this GEMM). The 4 contiguous columns (TN=4) are what
+# let codegen vectorize loads/stores to float4, matching tilelang's lowering exactly.
+BLOCK_M = BLOCK_N = BLOCK_K = 64
+TY = 8
+TX = 16
+THREADS = TY * TX
+TM = BLOCK_M // TY   # fragment rows per thread  (8)
+TN = BLOCK_N // TX   # fragment columns per thread (4)
+
+def matmul_relu_kernel(c:UOp, a:UOp, b:UOp) -> UOp:
+  """C[M, N] = relu(A[M, K] @ B[K, N]) -- one 64x64 tile per block, locals + a 2-D fragment."""
+  M, K = a.shape
+  K2, N = b.shape
+  assert K == K2 and a.dtype == b.dtype == c.dtype and not dtypes.is_int(a.dtype)
+  assert not (K % BLOCK_K or M % BLOCK_M or N % BLOCK_N), "test sizes must be multiples of the block sizes"
+
+  # with T.Kernel(T.ceildiv(N, BLOCK_N), T.ceildiv(M, BLOCK_M), threads=128) as (bx, by):
+  bx = UOp.range(cdiv(N, BLOCK_N), 0, AxisType.GLOBAL)
+  by = UOp.range(cdiv(M, BLOCK_M), 1, AxisType.GLOBAL)
+
+  # 16*8 threads = 128 threads
+  tx = UOp.range(TX, 2, AxisType.LOCAL)
+  ty = UOp.range(TY, 3, AxisType.LOCAL)
+
+  # shared + fragment (regs)
+  A_shared = alloc_shared((BLOCK_M, BLOCK_K), a.dtype, 0)
+  B_shared = alloc_shared((BLOCK_K, BLOCK_N), b.dtype, 1)
+  C_local = alloc_fragment((TM, TY, TX, TN), dtypes.float32, 0, (1, 2), (ty, tx))
+
+  # zero out the regs to start. this is expanded by the devectorizer
+  C_local = C_local.after(C_local.store(0.0))
+
+  # for ko in T.Pipelined(T.ceildiv(K, BLOCK_K), num_stages=3):
+  ko = UOp.range(cdiv(K, BLOCK_K), 6, AxisType.LOOP)
+
+  # index the outer matrices
+  a = a.rearrange("(m bm) (k bk) -> m k bm bk", bm=BLOCK_M, bk=BLOCK_K)[by, ko]
+  b = b.rearrange("(k bk) (n bn) -> k n bk bn", bk=BLOCK_K, bn=BLOCK_N)[ko, bx]
+  c = c.rearrange("(m bm) (n bn) -> m n bm bn", bm=BLOCK_M, bn=BLOCK_N)[by, bx]
+
+  # T.copy: A_shared <- a, B_shared <- b
+  def with_threads(x:UOp): return x.rearrange("(tm ty) (tx tn) -> ty tx tm tn", tm=TM, tn=TN)[ty, tx]
+  A_shared = A_shared.after(with_threads(A_shared).store(with_threads(a)))
+  B_shared = B_shared.after(with_threads(B_shared).store(with_threads(b)))
+
+  # T.gemm(A_shared, B_shared, C_local), no WMMA
+  kk = UOp.range(BLOCK_K, 11, AxisType.LOOP)
+  ir = UOp.range(TM, 12, AxisType.LOOP)
+  jj = UOp.range(TN, 13, AxisType.UPCAST)
+  acc = C_local.after(kk)[ir, ty, tx, jj] + A_shared[ir*TM + ty, kk].cast(dtypes.float32) * B_shared[kk, tx*TN + jj].cast(dtypes.float32)
+  # closing the ko loop here too; codegen adds the barrier so no thread overwrites the tiles while others still read them
+  C_local = C_local[ir, ty, tx, jj].set(acc, end=(kk, ir, jj, ko))
+
+  # c <- C_local (with relu and cast): every thread stores its shard's sub-view of the output tile
+  c_st = c.reshape(C_local.shape).store(C_local.relu().cast(c.dtype))
+
+  # close the locals and globals
+  return c_st.end(tx, ty, bx, by).sink(arg=KernelInfo(name="matmul_relu", opts_to_apply=()))
+
+# ---------------------------------------------------------------------------
+# python wrapper: same signature as the tilelang function
+# ---------------------------------------------------------------------------
+
+def matmul_relu(a:Tensor, b:Tensor) -> Tensor:
+  """C = relu(A @ B), fp16 in/out with an fp32 fragment accumulator."""
+  c = Tensor.empty(a.shape[0], b.shape[1], dtype=a.dtype, device=a.device)
+  return c.custom_kernel(a, b, fxn=matmul_relu_kernel)[0]
+
+# ---------------------------------------------------------------------------
+# test
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+  from tinygrad import Device
+  assert Device[Device.DEFAULT].renderer.has_local, "this GPU-style kernel needs a backend with local memory (LOCAL ranges + barriers)"
+  M = K = N = getenv("N", 256)  # 4x4 grid of 64x64 tiles, 4 K chunks
+  dtype_in = dtypes.half if getenv("HALF") else dtypes.float
+
+  a = Tensor.randn(M, K, dtype=dtype_in).contiguous()
+  b = Tensor.randn(K, N, dtype=dtype_in).contiguous()
+  ref = (a @ b).relu().realize()
+
+  for _ in range(10):
+    out = matmul_relu(a, b).realize()
+
+  import numpy as np
+  np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=1e-1, rtol=1e-2)
+  print("matmul_relu passed!")
@@ -19,6 +19,9 @@ def cmd_remove_module(args):
  modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia", "ast"] if args.backend == "nv" else ["amdgpu"]
  to_unload = [m for m in modules if _is_module_loaded(m)]
  if not to_unload: print("Kernel modules are not loaded")
+  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,17 +97,20 @@ def add_common_commands(parent_subparsers):

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

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

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

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

@@ -1,544 +0,0 @@
-from __future__ import annotations
-from typing import cast, Callable, TypeVar, Generic, Any
-import struct, functools, time, collections, itertools
-from dataclasses import replace, dataclass
-from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, dedup, pluralize, JIT_BATCH_SIZE, unwrap
-from tinygrad.helpers import to_tuple, round_up, partition, data64_le, panic, ContextVar
-from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer
-from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, track_rewrites, GroupOp
-from tinygrad.uop.symbolic import symbolic
-from tinygrad.dtype import dtypes, truncate
-from tinygrad.runtime.support.hcq import MMIOInterface
-from tinygrad.runtime.support.memory import BumpAllocator
-from tinygrad.renderer import Renderer, Estimates
-from tinygrad.engine.realize import to_program, get_call_arg_uops, get_call_name, get_call_outs_ins, estimate_uop, pm_flatten_linear
-from tinygrad.engine.jit import DepsTracker
-
-# *****************
-# 0. helpers
-
-HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQ2Compiled')
-
-HCQ_RUNTIME_DEV = ContextVar("HCQ_RUNTIME_DEV", "CPU")
-
-HCQ_DEVS = frozenset(("AMD",))
-HCQ_P2P_DEVS = HCQ_DEVS | frozenset(("CPU",))
-HCQ_CACHE_TAGS = frozenset(("program", "systems", "template"))
-
-@dataclass(frozen=True)
-class HCQInfo:
-  name:str
-  estimates:Estimates
-  device:tuple[str, ...]
-  queue:str
-
-  input_idxs:tuple[int, ...] = () # indexes into input_uops used by this call
-  inputs:int|None = None
-
-def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c
-
-def unwrap_mstack(u):
-  return tuple(x for s in u.src for x in unwrap_mstack(s)) if u.op is Ops.MSTACK else (unwrap_mstack(u.src[0]) if u.op in {Ops.MSELECT, Ops.SLICE} else (u,))
-
-def make_patch(buf:UOp, off:sint, val:UOp) -> UOp:
-  return buf.index(UOp.const(dtypes.int, off // buf.dtype.itemsize)).store(val.simplify().cast(buf.dtype))
-
-def make_binary_patch(buf:UOp, blob:bytes) -> UOp:
-  data = UOp(Ops.BINARY, src=(), arg=blob).bitcast(buf.dtype)
-  r = UOp.range(len(blob) // buf.dtype.itemsize, 0, dtype=dtypes.int, src=(buf, data))
-  return buf.index(r).store(data.index(r).load()).end(r)
-
-def make_cmdbuf(lin, devs):
-  blob, patches = b'', []
-  for s in (s for ins in lin.src for s in ins.src):
-    if (ssimp:=s.simplify()).op is not Ops.CONST: patches.append((len(blob), ssimp))
-    blob += struct.pack(f'<{ssimp.dtype.fmt}', ssimp.arg if ssimp.op is Ops.CONST else 0x0)
-  cmdbuf = UOp.placeholder((len(blob) // 4,), dtypes.uint32, next(UOp.unique_num), device=devs).rtag("cmdbuf")
-  return cmdbuf.after(make_binary_patch(cmdbuf, blob), *[make_patch(cmdbuf, off, s) for off, s in patches])
-
-def make_signal(devs, queue="COMPUTE:0", sentinel=False):
-  return UOp.placeholder((1,), dtypes.uint64, 0, device=devs).rtag("sentinel_signal" if sentinel else f"{queue}_timeline_signal")
-def make_signal_value(devs, queue="COMPUTE:0"):
-  return UOp.placeholder((1,), dtypes.uint64, 0, device=devs).rtag(f"{queue}_timeline_value")
-
-def make_submit(*cmds, devs:str|tuple[str, ...], queue:str) -> UOp:
-  return UOp.custom_function("submit_cmdbuf", UOp(Ops.LINEAR, src=tuple(cmds), arg=(to_tuple(devs), queue)))
-def get_submit(ast:UOp) -> UOp: return next(u for u in ast.toposort() if u.op is Ops.CUSTOM_FUNCTION and u.arg == "submit_cmdbuf")
-
-def encode_kernargs_clike(call:UOp, prg:UOp, devs:str|tuple[str, ...]) -> UOp:
-  data, info = prg.arg
-  buf = UOp.placeholder((data.kernargs_alloc_size // 4,), dtypes.uint32, next(UOp.unique_num), device=devs).rtag("kernargs")
-  words = [w for gi in info.globals for w in data64_le(get_call_arg_uops(call)[gi].getaddr(devs))] + list(info.vars)
-  return buf.after(*[make_patch(buf, i * 4, w) for i, w in enumerate(words)])
-
-# *****************
-# 0.1. prep: replace buffers with params
-
-def replace_call_buffers(ctx:list[UOp], call:UOp) -> UOp|None:
-  ctx += [s for s in call.src[1:] if s not in ctx and s.op not in (Ops.PARAM, Ops.BIND)]
-  return call.replace(src=call.src[:1] + tuple(s if s.op in (Ops.PARAM, Ops.BIND) else s.param_like(ctx.index(s)) for s in call.src[1:]))
-pm_replace_buffers = PatternMatcher([(UPat(Ops.CALL, name="call"), replace_call_buffers)])
-
-# *****************
-# 1.1. prep: staging copies
-
-def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS) and not all_devices_in(b.device, HCQ_P2P_DEVS)
-
-def stage_copy(dst:UOp, src:UOp) -> UOp|None:
-  if not (_need_staging(src, dst) or _need_staging(dst, src)): return None
-
-  stage = UOp.new_buffer("CPU", src.max_numel() * src.dtype.itemsize, dtypes.uint8)
-  return UOp(Ops.LINEAR, src=(src.copy_to_device("CPU").call(stage, src), stage.copy_to_device(dst.device).call(dst, stage)))
-pm_insert_copy_staging = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy)])
-
-# *****************
-# 2. deps
-
-class HCQDepsTracker(DepsTracker):
-  @staticmethod
-  def _key(buf:Any) -> tuple[Any, int, int]:
-    return (buf.arg.slot, 0, buf.max_numel() * buf.dtype.itemsize) if isinstance(buf, UOp) else DepsTracker._key(buf)
-
-def _get_call_bufs_by_lane(call:UOp, devices:tuple[str, ...]) -> list[list[Any]]:
-  refs = get_call_arg_uops(call)
-  return [[b if b.op is Ops.PARAM else mb.bufs[lane] if isinstance(mb:=b.buffer, MultiBuffer) else mb for b in refs] for lane in range(len(devices))]
-
-def _get_deps(ctx:DepsTracker, bufs_by_lane:list[list[Any]], write, key:tuple[tuple[str, ...], str, int]) -> list[tuple[tuple, int, int]]:
-  dep_lanes:list[tuple[tuple, int, int]] = []
-  for lane, bufs in enumerate(bufs_by_lane):
-    dep_lanes += [(dep, dlane, lane) for dep, dlane in ctx.access_resources(bufs, write if write is not None else range(len(bufs)), (key, lane))]
-  return dep_lanes
-
-def _build_wait_cmds(dep_lanes:list[tuple[tuple, int, int]], devices:tuple[str, ...], queue:str) -> tuple[list[UOp], set[int]]:
-  # opt1: same-queue ops are fifo-ordered
-  if devices[0].split(":")[0] in {"AMD", "QCOM"} or queue.startswith("COPY"):
-    dep_lanes = [(dep, dlane, lane) for dep, dlane, lane in dep_lanes if (dep[0][dlane], dep[1]) != (devices[lane], queue)]
-
-  # opt2: keep latest dep per (dep device, queue, cur lane)
-  latest = {((dep[0][dlane], dep[1]), lane): (dep, dlane) for dep, dlane, lane in sorted(dep_lanes, key=lambda x: x[0][2])}
-  deps:dict[tuple, list[int|None]] = collections.defaultdict(lambda: [None]*len(devices))
-  for (_, lane), (dep, dlane) in latest.items(): deps[dep][lane] = dlane
-
-  waits = []
-  for (ddevs, dqueue, dtag), lanes in deps.items():
-    sig = UOp.mstack(*[make_signal(d if dl is None else ddevs[dl], queue=dqueue, sentinel=dl is None) for dl, d in zip(lanes, devices)])
-    val = UOp.mstack(*[make_signal_value(d if dl is None else ddevs[dl], queue=dqueue) for dl, d in zip(lanes, devices)])
-    waits.append(UOp(Ops.INS, arg="wait", src=(sig, val.index(UOp.const(dtypes.int, 0)) + dtag)))
-  return waits, {dtag for _, _, dtag in deps}
-
-def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[tuple[tuple[str, ...], str]],
-                      tracker:HCQDepsTracker) -> tuple[list[UOp], set[int]]:
-  # collect all buffers which belong to devices
-  dev_bufs:dict[str, dict[int, Any]] = collections.defaultdict(dict)
-  for call, devices in batch:
-    for b in itertools.chain.from_iterable(_get_call_bufs_by_lane(call, devices)):
-      for bd in to_tuple(b.device): dev_bufs[bd][id(b)] = b
-
-  zero, n, finalizers, waited = UOp.const(dtypes.int, 0), len(batch_info), [], set()
-  for _, devgroup in itertools.groupby(sorted(dedup([d for devs, _ in batch_info for d in devs])), key=lambda d: d.split(":")[0]):
-    devs = tuple(devgroup)
-
-    # to finalize the batch, sync all accesses from other devices to buffers that belong to this device
-    fin_deps = [dl for dl in _get_deps(tracker, [list(dev_bufs[d].values()) for d in devs], None, key=(devs, "COMPUTE:0", n)) if dl[0][2] < n]
-    waits, cur_waited = _build_wait_cmds(fin_deps, devs, "COMPUTE:0")
-    waited |= cur_waited
-
-    # wait the syncs, store the device epoch; value bumps are a separate call: no lane may bump until every lane has patched its waits
-    store = UOp(Ops.INS, arg="store", src=(make_signal(devs), (tl:=make_signal_value(devs)).index(zero) + n))
-    submit = make_submit(*waits, store, devs=devs, queue="COMPUTE:0")
-    upd = [(tl, n + 1)] + [(make_signal_value(devs, queue=qn), n)
-                           for qn in dedup([qn for bdevs, qn in batch_info if set(bdevs) & set(devs)]) if qn != "COMPUTE:0"]
-    bump = UOp.barrier(*[s.index(zero, dtype=s.dtype).store(s.index(zero) + inc) for s, inc in upd])
-    finalizers += [UOp.custom_function("hcq", b.sink()).call(aux=HCQInfo("hcq_finalizer", Estimates(), devs, "COMPUTE:0")) for b in (submit, bump)]
-  return finalizers, waited
-
-def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]]) -> list[UOp]:
-  batch_info = [(devices, "COMPUTE:0" if call.src[0].op is Ops.PROGRAM else "COPY:0") for call, devices in batch]
-
-  # schedule deps
-  waited:set[int] = set()
-  deps_tracker = HCQDepsTracker()
-  call_waits:list[list[UOp]] = []
-  for tag, ((call, _), (devices, queue)) in enumerate(zip(batch, batch_info)):
-    deps = _get_deps(deps_tracker, _get_call_bufs_by_lane(call, devices), get_call_outs_ins(call)[0], key=(devices, queue, tag))
-    cmds, cur_waited = _build_wait_cmds(deps, devices, queue)
-    call_waits.append(cmds)
-    waited |= cur_waited
-
-  # build finalizers
-  finalizers, finalizer_waited = _build_finalizers(batch, batch_info, deps_tracker)
-  waited |= finalizer_waited
-
-  src = []
-  for tag, ((call, _), (devices, queue), q) in enumerate(zip(batch, batch_info, call_waits)):
-    # first queue use, sync prior device work with main signal
-    if batch_info.index((devices, queue)) == tag:
-      q = [UOp(Ops.INS, arg="barrier", src=()), UOp(Ops.INS, arg="wait", src=(make_signal(devices), make_signal_value(devices).index(0) - 1))] + q
-
-    # and make hcq call
-    info = HCQInfo(get_call_name(call, get_call_arg_uops(call)), estimate_uop(call), devices, queue)
-    q += [call.replace(arg=replace(call.arg, aux=info))]
-
-    # signal queue timeline if someone waits for us
-    if tag in waited: q += [UOp(Ops.INS, arg="store", src=(make_signal(devices, queue), make_signal_value(devices, queue).index(0) + tag))]
-    src.append(UOp.custom_function("hcq", make_submit(*q, devs=devices, queue=queue).sink()).call(name="hcq", aux=info))
-  return src + finalizers
-
-def sched_hcq_batches(l:UOp) -> UOp:
-  srcs:list[UOp] = []
-  batch:list[tuple[UOp, tuple[str, ...]]] = []
-  for call in l.src:
-    if (devs:=next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None)) is not None: batch.append((call, to_tuple(devs)))
-    else: srcs, batch = srcs + _finalize_batch(batch) + [call], []
-  return l.replace(src=tuple(srcs + _finalize_batch(batch)))
-pm_sched_hcq_batches = PatternMatcher([(UPat(Ops.LINEAR, name="l"), sched_hcq_batches)])
-
-# *****************
-# 3. merge into queues
-
-def _merged_hcq_call(calls:list[UOp]) -> UOp: # TODO: simplify?
-  if len(calls) == 1: return calls[0]
-  info = replace(calls[0].arg.aux, name=f"submit {calls[0].arg.aux.queue} ({len(calls)})",
-                 estimates=sum((c.arg.aux.estimates for c in calls), start=Estimates()))
-  cmds = [cmd for c in calls for cmd in get_submit(c).src[0].src]
-  return UOp.custom_function("hcq", make_submit(*cmds, devs=info.device, queue=info.queue).sink()).call(name="hcq", aux=info)
-
-def merge_queues(linear:UOp) -> UOp:
-  new_src:list[UOp] = []
-  opened_qs:dict[tuple[tuple[str, ...], str], list[UOp]] = {} # (devs, queue) -> list of hcq calls, kept in submit order
-  limits = collections.defaultdict(lambda: JIT_BATCH_SIZE.value)
-
-  for call in linear.src:
-    if not isinstance(info:=call.arg.aux, HCQInfo) or info.name == "hcq_finalizer": # non-hcq call or finalizer: close all open queues
-      new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in list(opened_qs)] + [call]
-      continue
-
-    if (old:=opened_qs.pop(key:=(info.device, info.queue), None)) is not None:
-      if limits[key] and len(old) >= limits[key]: new_src, old, limits[key] = new_src + [_merged_hcq_call(old)], [], limits[key] * 2
-      new_rec = old + [call]
-    else:
-      # no such queue opened: close every open submit on this queue that shares a device, so submit order is kept
-      closing = [k for k in opened_qs if k[1] == info.queue and set(k[0]) & set(info.device)]
-      new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in closing]
-      new_rec = [call]
-    opened_qs[(info.device, info.queue)] = new_rec
-  return linear.replace(src=tuple(new_src + [_merged_hcq_call(c) for c in opened_qs.values()]))
-pm_merge_queues = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), merge_queues)])
-
-# *****************
-# 4.2. hcq lowering: ops to ir
-
-def encode_cmdbuf(submit:UOp, lin:UOp) -> UOp|None:
-  if (pm:=Device.get_class(lin.arg[0][0]).pm_lower) is None: return None
-  return graph_rewrite(submit, pm, name=f"encode {lin.arg[0]}", enter_calls=True)
-pm_encode_cmdbufs = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="lin"),), name="submit"), encode_cmdbuf)])
-
-# *****************
-
-def is_value_known_at_link(val:UOp) -> bool:
-  runtime_reads = [u for u in val.toposort() if u.op in (Ops.LOAD, Ops.INDEX)]
-  addressed_bufs = [b for g in val.toposort() if g.op is Ops.GETADDR for b in unwrap_mstack(g.buf_uop)]
-
-  # addr of input params is not known at link time
-  return not val.variables() and not runtime_reads and all(b.op is not Ops.PARAM or b.tag is not None for b in addressed_bufs)
-
-def is_link_patch(p:UOp, jit:bool) -> bool:
-  store = p.src[0] if (is_binary_patch:=p.op is Ops.END) else p
-  if not jit: return store.buf_uop.tag == "program"
-  return is_binary_patch or (store.op is Ops.STORE and is_value_known_at_link(store.src[1]))
-
-def trim_link_patches(ctx:tuple[bool, list[UOp]], a:UOp) -> UOp|None:
-  links, kept = partition(a.src[1:], lambda p: is_link_patch(p, ctx[0]))
-
-  # keep all patches from the link-time patches' subtrees in the C code
-  afters = [u for u in UOp.sink(*links).toposort() if u.op is Ops.AFTER]
-  ctx[1].extend(UOp.sink(*links).substitute({p: p.src[0] for p in afters}).src)
-  return a.src[0].after(*kept, *[d for p in afters for d in p.src[1:]]) if links else None
-pm_trim_link_patches = PatternMatcher([(UPat(Ops.AFTER, src=(UPat((Ops.PARAM, Ops.MSTACK)),), allow_any_len=True, name="a"), trim_link_patches)])
-
-def split_patches(ctx:bool, call:UOp) -> UOp|None:
-  body = graph_rewrite(call.src[0], pm_trim_link_patches, ctx=(ctx, lt_patches:=[]), name=f"trim link-time patches ({call.arg.aux.name})")
-
-  lt_srcs = collections.defaultdict(list)
-  for p in lt_patches: lt_srcs[p.buf_uop].append(p)
-  return call.replace(src=(body, *call.src[1:], *[b.after(*ps) for b,ps in lt_srcs.items()]))
-pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), split_patches)])
-
-# *****************
-
-def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[dict[UOp, UOp], tuple[UOp, ...]]:
-  bare = {g: g.replace(src=(g.src[0].without_after,)) for g in gaddrs}
-
-  order = sorted(dedup(bare.values()), key=lambda g: ((b:=unwrap_mstack(g.buf_uop)[0]).arg.slot, repr(b.tag)))
-  slots = {g:i for i,g in enumerate(order)}
-  table = UOp.placeholder((len(order),), dtypes.uint64, next(UOp.unique_num), device=call.arg.aux.device).rtag(name)
-
-  reads = {g: table.after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ()).index(UOp.const(dtypes.int, slots[bare[g]])).load() for g in gaddrs}
-  return reads, (table.after(*[make_patch(table, i * table.dtype.itemsize, addr) for addr, i in slots.items()]),) if slots else ()
-
-def make_blob_bufs(call:UOp, blobs:list[UOp]) -> tuple[dict[UOp, UOp], tuple[UOp, ...]]:
-  bufs = {b: UOp.placeholder((b.max_numel(),), b.dtype, next(UOp.unique_num), device=call.arg.aux.device).rtag("template") for b in blobs}
-  return bufs, tuple(buf.after(make_binary_patch(buf, b.src[0].arg)) for b,buf in bufs.items())
-
-def rm_rt_uops(call:UOp) -> UOp|None:
-  if not (rt_uops:=[u for u in call.src[0].toposort() if u.op is Ops.GETADDR or (u.op is Ops.BITCAST and u.src[0].op is Ops.BINARY)]): return None
-  gaddrs, blobs = partition(rt_uops, lambda u: u.op is Ops.GETADDR)
-  inputs, internals = partition(gaddrs, lambda g: all(x.op is Ops.PARAM and x.tag is None for x in unwrap_mstack(g.buf_uop)))
-  runtimes, systems = partition(internals, lambda g: any(x.tag in {"program", "kernargs", "cmdbuf"} for x in unwrap_mstack(g.buf_uop)))
-
-  # exec fills the inputs table with the input addresses every run, so it has no fill patches
-  (reads, _), *tables = [make_addr_table(call, gs, n) for gs,n in ((inputs, "inputs"), (runtimes, "runtime"), (systems, "systems"))] + \
-                        [make_blob_bufs(call, blobs)]
-  reads, fills = reads | {k:v for r,_ in tables for k,v in r.items()}, [f for _,fs in tables for f in fs]
-  return call.replace(src=(call.src[0].substitute(reads), *call.src[1:], *fills),
-                      arg=replace(call.arg, aux=replace(call.arg.aux, input_idxs=tuple(sorted(dedup(g.buf_uop.arg.slot for g in inputs))))))
-pm_rm_rt_uops = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), rm_rt_uops)])
-
-# *****************
-
-def replace_params(call:UOp) -> UOp|None:
-  body, variables, param_ops = call.src[0], call.src[0].variables(), {Ops.PARAM, Ops.MSTACK}
-  args = dedup([s for u in body.toposort(gate=lambda u: u.op not in param_ops) for s in u.src if s.op in param_ops and s not in variables])
-
-  patched, refhold = partition(call.src[1:], lambda x: x.src[0] in args)
-  by_root = {p.src[0]: p for p in patched}
-  c_args = [by_root.get(a, a) for a in args]
-
-  sub = {(b:=u.without_after): UOp.param(i, u.dtype, shape=b.shape, device=u.device) for i,u in enumerate(c_args)} | \
-        {v: v.replace(arg=replace(v.arg, slot=-1)) for v in variables if v.op is Ops.PARAM}
-  info = replace(call.arg.aux, inputs=next((i for i,u in enumerate(c_args) if u.tag == "inputs"), None))
-  return call.replace(src=(body.substitute(sub), *c_args, *refhold), arg=replace(call.arg, aux=info)) # TODO: call.after(*refhold)?
-pm_replace_params = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), replace_params)])
-
-# *****************
-
-def resolve_getaddr_slice(bv:UOp, g:UOp) -> UOp:
-  base = bv.src[0].after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ())
-  itemsize = bv.src[0].dtype.itemsize if bv.src[0].without_after.op in (Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize
-  return UOp(Ops.GETADDR, dtypes.uint64, src=(base,), arg=g.arg) + UOp.const(dtypes.uint64, bv.src[1].arg * itemsize)
-
-pm_early_simplify = PatternMatcher([
-  (UPat(Ops.GETADDR, src=(UPat.any(sl:=UPat(Ops.SLICE, name="bv"), sl.after(allow_any_len=True)),), name="g"), resolve_getaddr_slice),
-  (UPat(Ops.INDEX, src=(UPat(Ops.SLICE, name="bv"),), allow_any_len=True, name="x"),
-   lambda bv,x: x.replace(src=(bv.src[0], x.src[1] + bv.src[1].cast(x.src[1].dtype), *x.src[2:]))),
-])
-
-# *****************
-# 5.3. pack placeholders buffers
-
-def pack_hcq_placeholders(call:UOp) -> UOp|None:
-  bufs = [b for b in call.src[0].toposort() if b.op is Ops.PARAM and b.tag in {"scratch", "kernargs"}]
-  offs, sizes = {}, {}
-  for b in bufs:
-    if b.tag == "scratch": sizes[b.tag] = max(sizes.get(b.tag, 0), b.max_numel())
-    else:
-      offs[b] = round_up(sizes.get(b.tag, 0), 128 // b.dtype.itemsize)
-      sizes[b.tag] = offs[b] + b.max_numel()
-  counts = collections.Counter(b.tag for b in bufs)
-  bases = {b.tag:UOp.placeholder((sizes[b.tag],), b.dtype, next(UOp.unique_num), device=b.device).rtag(b.tag) for b in bufs if counts[b.tag] > 1}
-  subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(dtypes.weakint, offs.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases}
-  return call.replace(src=(call.src[0].substitute(subs, walk=True), *call.src[1:])) if subs else None
-pm_pack_placeholders = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), pack_hcq_placeholders)])
-
-# *****************
-# 8. callify hcq programs
-
-pm_callify_hcq = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="hcq", src=(UPat(Ops.SINK),), name="cf"),
-  lambda cf: cf.replace(src=(to_program(cf.src[0].replace(arg=KernelInfo("hcq_submit"), tag=1), Device[HCQ_RUNTIME_DEV.value].renderer),)))])
-
-hcq_compile_cache:dict[tuple[bytes, bool], UOp] = {}
-
-@track_rewrites(lambda linear,input_uops,jit,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
-def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None, jit=False) -> UOp:
-  if input_uops is not None: linear = graph_rewrite(linear, pm_replace_buffers, ctx=input_uops, walk=True, enter_calls=True, name="replace buffer")
-
-  if (final_linear:=(hcq_compile_cache.get(cache_key:=(linear.key, jit)))) is None:
-    # prep
-    linear = linear.substitute(back_map:={s.param_like(i): s for i,s in enumerate(input_uops)} if input_uops is not None else {}, walk=True)
-    linear = graph_rewrite(linear, pm_insert_copy_staging + pm_flatten_linear, name="insert copy staging")
-
-    # schedule
-    linear = graph_rewrite(linear, pm_sched_hcq_batches, walk=True, name="schedule hcq batches")
-    linear = linear.substitute({s: p for p, s in back_map.items()}, walk=True, enter_calls=True)
-    linear = graph_rewrite(linear, pm_merge_queues, walk=True, name="merge queues")
-
-    # lowering to hcq ir
-    linear = graph_rewrite(linear, pm_encode_cmdbufs, walk=True, name="encode cmdbufs", enter_calls=True)
-    linear = graph_rewrite(linear, pm_pack_placeholders, walk=True, name="pack placeholders")
-
-    # pie
-    linear = graph_rewrite(linear, pm_split_patches, ctx=jit, walk=True, name="split rt/lt patches")
-    linear = graph_rewrite(linear, pm_early_simplify + symbolic, bottom_up=False, name="simplify packed placeholders", enter_calls=True)
-    linear = graph_rewrite(linear, pm_rm_rt_uops, walk=True, name="replace rt uops")
-    linear = graph_rewrite(linear, pm_replace_params, walk=True, name="replace with args")
-
-    # and compile it
-    final_linear = hcq_compile_cache[cache_key] = graph_rewrite(linear, pm_callify_hcq, name="callify hcq", enter_calls=True)
-
-  return final_linear
-
-# *****************
-# 6. bufferize placeholders: replace placeholders with real buffers.
-
-def bufferize_buf(ctx:bool, buf:UOp) -> UOp|None:
-  if buf.tag is None: return None
-  return UOp.mstack(*(UOp.from_buffer((dv:=Device[dev]).pm_bufferize.rewrite(buf, ctx=(dv, ctx)), HCQ_RUNTIME_DEV.value)
-                      for dev in to_tuple(buf.device)))
-pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, name="buf"), bufferize_buf)])
-
-# *****************
-# 7. resolve patches
-
-def push_stack(op, s): return UOp(Ops.STACK, op.dtype.scalar(),
-  tuple(op.replace(dtype=op.dtype.scalar(), src=tuple(x if y is s else y for y in op.src)) for x in s.src))
-
-def fold_binary(buf:UOp, blob:UOp) -> UOp:
-  for b in (m.bufs if isinstance(m:=buf.buffer, MultiBuffer) else (m,)):
-    b.ensure_allocated().as_memoryview(force_zero_copy=True, no_sync=True).cast('B')[:len(blob.arg)] = blob.arg
-  return UOp(Ops.NOOP)
-
-def fold_const_store(buf:UOp, off:UOp, val:UOp) -> UOp:
-  for b, v in zip((bs:=mb.bufs if isinstance((mb:=buf.buffer), MultiBuffer) else (mb,)), val.src if val.op is Ops.STACK else (val,)*len(bs)):
-    data = struct.pack(f'<{v.dtype.fmt}', truncate[v.dtype](v.arg))
-    b.ensure_allocated().as_memoryview(force_zero_copy=True, no_sync=True).cast('B')[(byte_off:=off.arg*buf.dtype.itemsize):byte_off+len(data)] = data
-  return UOp(Ops.NOOP)
-
-def resolve_getaddr(buf:UOp, g:UOp) -> UOp:
-  assert buf.op in (Ops.BUFFER, Ops.MSTACK, Ops.MSELECT), f"{buf.op}"
-
-  devs, b = g.arg, buf.buffer
-  bufs = tuple(cast(Buffer, x.buffer) for x in buf.src) if buf.op is Ops.MSTACK else tuple(b.bufs if isinstance(b, MultiBuffer) else (b,)*len(devs))
-  assert len(bufs) == len(devs), f"can't resolve {len(bufs)} buffers on {len(devs)} devices"
-  addrs = tuple(UOp.const(dtypes.uint64, x.get_buf(d).va_addr) for x, d in zip(bufs, devs))
-  return addrs[0] if len(addrs) == 1 else UOp(Ops.STACK, dtypes.uint64, addrs)
-
-pm_resolve_patches = PatternMatcher([
-  # multi
-  (UPat(GroupOp.ALU, src=[UPat(Ops.STACK, name="s"), UPat(Ops.CONST)], name="op"), push_stack),
-  (UPat(Ops.CAST, src=(UPat(Ops.STACK, name="s"),), name="op"), push_stack),
-
-  # getaddr
-  (UPat(Ops.GETADDR, src=(UPat(name="buf"),), name="g"), resolve_getaddr),
-
-  # folders
-  (UPat(name="buf").index(UPat(Ops.RANGE), allow_any_len=True)
-    .store(UPat.any(UPat(Ops.BINARY, name="blob"), UPat(Ops.BINARY, name="blob").bitcast()).index(UPat(Ops.RANGE), allow_any_len=True).load())
-    .end(UPat(Ops.RANGE)), fold_binary),
-  (UPat({Ops.BUFFER, Ops.SLICE, Ops.MSTACK}, name="buf").index(UPat.cvar("off"))
-    .store(UPat.any(UPat.cvar("val"), UPat(Ops.STACK, name="val"))), fold_const_store),
-])
-
-pm_assert_no_afters = PatternMatcher([(UPat(Ops.AFTER, name="a"), lambda a: panic(RuntimeError, f"AFTER left at hcq_link: {a.src[0].op}"))])
-
-hcq_link_cache:dict[tuple[bytes, tuple[str, ...]], UOp] = {}
-
-def link_cache_key(a:UOp): return a.key, to_tuple(a.device)
-pm_link_cache = PatternMatcher([(UPat(Ops.AFTER, name="a"), lambda a: hcq_link_cache.get(link_cache_key(a)))])
-
-@track_rewrites(lambda _,jit,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}")
-def hcq_link(linear:UOp, jit=False) -> UOp:
-  cacheable = {(j,i):a for j,c in enumerate(linear.src) for i,a in enumerate(c.src[1:], 1)
-               if a.op is Ops.AFTER and unwrap_mstack(a.src[0])[0].tag in HCQ_CACHE_TAGS}
-  hits = {a.src[0]:hcq_link_cache[key] for a in cacheable.values() if (key:=link_cache_key(a)) in hcq_link_cache}
-  linear = graph_rewrite(linear, pm_link_cache, name="apply link cache").substitute(hits, walk=True)
-  linear = graph_rewrite(linear, pm_bufferize, ctx=jit, bottom_up=True, walk=True, name="bufferize placeholders")
-  linear = graph_rewrite(linear, pm_resolve_patches + symbolic, bottom_up=False, name="simplify patches")
-  linear = graph_rewrite(linear, pm_assert_no_afters, name="assert no afters")
-  for (j,i),a in cacheable.items(): hcq_link_cache.setdefault(link_cache_key(a), linear.src[j].src[i])
-  return linear
-
-# *****************
-# Device classes
-
-class HCQ2Compiled(Compiled):
-  timestamp_divider: float = 1000.0
-
-  def __init__(self, device:str, allocator:HCQAllocator, compilers:list[type[Renderer]], runtime, can_recover:bool=False, arch=None):
-    self.device_id:int = int(device.split(":")[1]) if ":" in device else 0
-
-    self.pm_bufferize = PatternMatcher([
-      (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx[0].timeline_signal("sentinel", (1 << 64) - 1)),
-      (UPat(Ops.PARAM, name="b"), lambda ctx, b: None if b.tag is None else ctx[0].new_buffer(b, jit=ctx[1]))
-    ])
-
-    super().__init__(device, allocator, compilers, lambda *a, **kw: None, None, arch=arch)
-
-    self.rt_buffer = Buffer(self.device, 64 << 20, dtypes.uint8, options=BufferSpec(uncached=True, cpu_access=True))
-    self.rt_allocator = BumpAllocator(64 << 20, wrap=False)
-
-  def new_buffer(self, b:UOp, jit:bool) -> Buffer:
-    if jit or b.tag in HCQ_CACHE_TAGS:
-      return Buffer(self.device, b.max_numel(), b.dtype, options=BufferSpec(uncached=True, cpu_access=True, nolru=True))
-    return self.rt_buffer.view(b.max_numel(), b.dtype, self.rt_allocator.alloc(b.max_numel() * b.dtype.itemsize, alignment=128))
-
-  @functools.cache
-  def timeline_signal(self, queue:str="COMPUTE:0", init_value:int=0) -> Buffer:
-    buf = Buffer(self.device, 1, dtypes.uint64, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
-    buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = init_value
-    return buf
-
-  @functools.cache
-  def timeline_value(self, queue:str="COMPUTE:0", init_value:int=1) -> Buffer:
-    buf = Buffer("CPU", 1, dtypes.uint64, preallocate=True)
-    buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = init_value
-    return buf
-
-  def synchronize(self, timeout:int|None=None):
-    if not hasattr(self, 'iface'): return
-    sig = self.timeline_signal().as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
-    tl = self.timeline_value().as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
-    st = time.perf_counter()
-    while sig[0] < tl[0] - 1:
-      if time.perf_counter() - st > (timeout or 3000) / 1000: self.on_device_hang()
-
-  def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent.
-
-  def count(self) -> int: return self.iface.count if hasattr(self, 'iface') else 1
-
-  def _select_iface(self):
-    assert (v:=getenv(k:=f'{type(self).__name__[:-6].upper()}_IFACE', "")) == "",  \
-      f"{k}={v} is deprecated, use DEV={replace(DEV.target(type(self).__name__[:-6]), interface=v)} instead"
-    assert hasattr(self, "ifaces"), "must have ifaces to select an iface"
-    t = DEV.target(dev:=type(self).__name__[:-6])
-    filtered = select_by_name(self.ifaces, lambda i: i.__name__[:-5], t.interface, f"{dev} has no interface {t.interface!r}")
-    filtered = [i for i in filtered if t.interface.startswith("MOCK") or not i.__name__[:-5].startswith("MOCK")] # never fall back to mock ifaces
-    return select_first_inited([functools.partial(cast(Callable, iface), self, self.device_id) for iface in filtered],
-                               f"No interface for {dev}:{self.device_id} is available")
-
-  def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"
-
-  def finalize(self):
-    try: self.synchronize() # try to finalize the device in any case
-    except RuntimeError as e: print(f"{self.device} synchronization failed before finalizing: {e}")
-
-    # if the device has an interface, call device_fini to clean up resources
-    if hasattr(self, 'iface') and hasattr(self.iface, 'device_fini'): self.iface.device_fini()
-
-@dataclass
-class HCQ2Buffer:
-  va_addr:sint
-  meta:Any=None
-  view:MMIOInterface|None=None
-
-  def offset(self, offset:int, size:int) -> HCQ2Buffer:
-    return HCQ2Buffer(self.va_addr+offset, meta=self.meta, view=(self.view.view(offset=offset, size=size) if self.view is not None else None))
-
-class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
-  def _as_buffer(self, buf:HCQ2Buffer) -> memoryview:
-    return unwrap(buf.view).mv
-
-  def _map(self, buf:HCQ2Buffer) -> HCQ2Buffer:
-    if not hasattr(self, '_do_map'): raise NotImplementedError("map failed: no method implemented")
-    return self._do_map(buf)
-
-  @suppress_finalizing
-  def _free(self, buf:HCQ2Buffer, options:BufferSpec|None=None):
-    if options is not None and options.external_ptr is not None: return
-    self.dev.synchronize()
-    if hasattr(self, '_do_free'): self._do_free(buf, options)
-
-  def _unmap(self, mb):
-    self.dev.synchronize()
-    self.dev.iface.free(mb)
-
-  def _offset(self, buf, size:int, offset:int) -> HCQ2Buffer: return buf.offset(offset=offset, size=size)
@@ -3,8 +3,8 @@ from typing import cast, Any, Callable
 import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools, collections, atexit
 assert sys.platform != 'win32'
 from dataclasses import dataclass
-from extra.hcq2.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, encode_kernargs_clike, make_cmdbuf
-from extra.hcq2.hcq2 import make_binary_patch
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, encode_kernargs_clike, make_cmdbuf
+from tinygrad.runtime.support.hcq2 import make_binary_patch
 from tinygrad.uop.ops import sint, UOp
 from tinygrad.device import Compiled, BufferSpec, Buffer, Device
 from tinygrad.dtype import dtypes
@@ -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
@@ -37,7 +37,7 @@ class PM4Ops(FastEnum):
  RELEASE_MEM = auto(); DISPATCH_DIRECT = auto(); EVENT_WRITE = auto()  # noqa: E702

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

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

 def memory_barrier(ctx):
  pf = '' if ctx.nbio.version[0] == 2 else '0' if ctx.nbio.version[:2] != (7, 11) else '1'
-  return UOp(Ops.LINEAR, 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,31 +146,34 @@ pm_pm4_opsel = PatternMatcher([
  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
 ])

-def pm4_submit(cmdbuf, devs):
-  size, zero = UOp.const(dtypes.uint32, cmdbuf.nbytes() // dtypes.uint32.itemsize), UOp.const(dtypes.int, 0)
+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)))

-  # the compute queue's ring and its host-side ring/write/put pointers (placeholders, resolved in pm_bufferize)
-  for d in 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)))
+def pm4_submit(ctx, lin):
+  # ensure compute queues are allocated
+  for d in (devs:=ctx.devs): q = Device[d].compute_queue
+  ring, wptr, doorbell, put_ptr = queue_ptrs(devs, "COMPUTE:0", q)

-  # place the cmdbuf at the ring's write offset, wrapping the ring
-  put = put_ptr.index(zero)
-  next_put = put + size.cast(put.dtype)
-  i = UOp.range(size, 0, dtype=dtypes.int, src=(cmdbuf,))
-  ring_idx = ((put + i.cast(put.dtype)) % q.ring.size).cast(dtypes.int)
+  # 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)
+  assert size_dw < (1 << 20), f"indirect buffer of {size_dw} dwords doesn't fit one packet"

-  # copy the cmdbuf into the ring and advance the put/write pointers
-  copy_to_ring = ring.index(ring_idx).store(cmdbuf.index(i).load()).end(i)
-  bump_put_ptr = put_ptr.index(zero).store(next_put)
-  bump_wptr = wptr.index(zero).store(next_put)
+  ib = UOp.placeholder((size_dw,), dtypes.uint32, next(UOp.unique_num), device=devs, volatile=True).rtag("cmdbuf")
+  cmdbuf = make_cmdbuf(lin, devs, buf=ib)

-  # ring the doorbell once the copy and pointer bumps have landed
-  flush = UOp.barrier(copy_to_ring, bump_put_ptr, bump_wptr)
-  return doorbell.after(flush).index(zero).store(next_put)
+  # the ring itself only carries a packet pointing at the ib, wrapping the ring
+  put = put_ptr.index(zero:=UOp.const(0, dtypes.int))
+  pkt = (ctx.pm4.PACKET3(ctx.pm4.PACKET3_INDIRECT_BUFFER, 2), *data64_le(cmdbuf.getaddr(devs)), size_dw | ctx.pm4.INDIRECT_BUFFER_VALID)
+  write_pkt = UOp.barrier(*[ring.index(((put + off) % q.ring.size).cast(dtypes.int)).store(UOp.const(x, dtypes.uint32)) for off,x in enumerate(pkt)])

-pm_pm4_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"),
-  lambda lin: pm4_submit(make_cmdbuf(lin, to_tuple(lin.arg[0])), to_tuple(lin.arg[0])))])
+  # advance the put/write pointers past the packet
+  bump_put_ptr = put_ptr.index(zero).store(put + len(pkt))
+  bump_wptr = wptr.index(zero).store(put + len(pkt))
+  flush = UOp.barrier(write_pkt, bump_put_ptr, bump_wptr)
+  return doorbell.after(flush).index(zero).store(put + len(pkt))
+
+pm_pm4_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"), pm4_submit)])

 # *****************
 # SDMA
@@ -179,32 +182,32 @@ class SDMAOps(FastEnum): COPY = auto(); POLL_REGMEM = auto(); FENCE = auto(); TR

 def sdma_copy(ctx, call):
  sz = call.src[2].max_numel() * call.src[2].dtype.itemsize
-  src_addr, dst_addr = call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs)
-  return call.ins(SDMAOps.COPY, src=tuple(UOp.const(dtypes.uint32, x) for off in range(0, sz, ctx.max_copy_size) for x in (
-    ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR),
-    ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0, *data64_le(src_addr+off), *data64_le(dst_addr+off))))
+  hdr = ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR)
+  return call.ins(SDMAOps.COPY, src=tuple(x for off in range(0, sz, ctx.max_copy_size) for x in (
+    *(UOp.const(v, dtypes.uint32) for v in (hdr, ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0)),
+    *(a + UOp.const(off, dtypes.uint64) if off else a for a in (call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs))))))

 def sdma_wait(ctx, ins, dst, val):
  op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \
     | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1)
-  return ins.ins(SDMAOps.POLL_REGMEM, src=tuple(UOp.const(dtypes.uint32, x) for x in (
+  return ins.ins(SDMAOps.POLL_REGMEM, src=tuple(UOp.const(x, dtypes.uint32) for x in (
    op, *data64_le(dst.getaddr(ctx.devs)), val, 0xffffffff,
    ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))))

 def sdma_store(ctx, ins, dst, val):
  op = ctx.sdma.SDMA_OP_FENCE | (ctx.sdma.SDMA_PKT_FENCE_HEADER_MTYPE(3) if ctx.target[0] != 9 else 0)
  return UOp(Ops.LINEAR, src=(
-    ins.ins(SDMAOps.FENCE, src=tuple(UOp.const(dtypes.uint32, x) for x in (op, *data64_le(dst.getaddr(ctx.devs)), val))),
-    ins.ins(SDMAOps.TRAP, src=tuple(UOp.const(dtypes.uint32, x) for x in (ctx.sdma.SDMA_OP_TRAP, 0)))))
+    ins.ins(SDMAOps.FENCE, src=tuple(UOp.const(x, dtypes.uint32) for x in (op, *data64_le(dst.getaddr(ctx.devs)), val))),
+    ins.ins(SDMAOps.TRAP, src=tuple(UOp.const(x, dtypes.uint32) for x in (ctx.sdma.SDMA_OP_TRAP, 0)))))

 def sdma_timestamp(ctx, ins, dst):
  op = ctx.sdma.SDMA_OP_TIMESTAMP | ctx.sdma.SDMA_PKT_TIMESTAMP_GET_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_TIMESTAMP_GET_GLOBAL)
-  return ins.ins(SDMAOps.TIMESTAMP, src=tuple(UOp.const(dtypes.uint32, x) for x in (op, *data64_le(dst.getaddr(ctx.devs)))))
+  return ins.ins(SDMAOps.TIMESTAMP, src=tuple(UOp.const(x, dtypes.uint32) for x in (op, *data64_le(dst.getaddr(ctx.devs)))))

 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),
@@ -212,12 +215,11 @@ pm_sdma_opsel = PatternMatcher([

 def sdma_submit(cmdbuf, devs):
  # the cmdbuf to submit + the patch writes that fill it
-  size_dw, zero = cmdbuf.nbytes() // dtypes.uint32.itemsize, UOp.const(dtypes.int, 0)
+  size_dw, zero = cmdbuf.nbytes() // dtypes.uint32.itemsize, UOp.const(0, dtypes.int)

  # 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)
@@ -228,8 +230,8 @@ def sdma_submit(cmdbuf, devs):

  # zero the wrapped tail, then copy the cmdbuf into the ring
  zi = UOp.range(zero_amt_dw, 0, dtype=dtypes.int, src=(cmdbuf,))
-  zero_tail = ring.index(tail_off_dw + zi).store(UOp.const(dtypes.uint32, 0)).end(zi)
-  i = UOp.range(UOp.const(dtypes.int, size_dw), 0, dtype=dtypes.int, src=(cmdbuf,))
+  zero_tail = ring.index(tail_off_dw + zi).store(UOp.const(0, dtypes.uint32)).end(zi)
+  i = UOp.range(UOp.const(size_dw, dtypes.int), 0, dtype=dtypes.int, src=(cmdbuf,))
  copy_to_ring = ring.index(start_dw + i).store(cmdbuf.index(i).load()).end(i)

  # advance the put/write pointers past the zeroed tail and the cmdbuf
@@ -242,18 +244,35 @@ def sdma_submit(cmdbuf, devs):
  return doorbell.after(flush).index(zero).store(next_put_b)

 pm_sdma_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"),
-  lambda lin: sdma_submit(make_cmdbuf(lin, to_tuple(lin.arg[0])), to_tuple(lin.arg[0])))])
+  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)
-  return submit.rewrite(graph_rewrite(q, opsel + pm_flatten_linear, walk=True, ctx=ctx, name=f"{q.arg[1]} opsel"))
+  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)
 class AMDProgramData:
@@ -261,10 +280,11 @@ class AMDProgramData:
  private_segment_size:int; kernargs_segment_size:int; kernargs_alloc_size:int
  enable_dispatch_ptr:int; enable_private_segment_sgpr:int

-_amd_program_cache:dict[tuple[bytes,str], tuple[AMDProgramData,bytes]] = {}
+_amd_program_cache:dict[tuple[bytes, tuple[str, ...]], UOp] = {}
 def amd_build_program(prg:UOp) -> UOp:
  dev = Device[to_tuple(prg.device)[0]] # TODO: rm this
-  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, dev.device))) is None:
+  # key on the full device tuple: the same lib can be built for different device sets, each needs its own program buffer
+  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, to_tuple(prg.device)))) is None:
    image, sections, relocs = elf_loader(lib)
    rodata = next(sh.header.sh_addr for sh in sections if sh.name == ".rodata")
    for off, sym, typ, addent in relocs:
@@ -281,20 +301,23 @@ 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_sdma_queue, supports_transfer=dev.has_sdma_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) -> HCQ2Buffer:
-    return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_sdma_queue)
+  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)

  def _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque)

-  def _do_map(self, buf:HCQ2Buffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+  def _do_map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+
+  def _do_unmap(self, buf:HCQBuffer): self.dev.iface.unmap(buf)

@dataclass
 class AMDQueueDesc:
@@ -387,15 +410,24 @@ class KFDIface:
    return hcqbuf

  def free(self, mem):
+    self._unmap(mem)
+    if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
+    kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
+
+  def unmap(self, mem):
+    self._unmap(mem)
+    if getattr(mem, '_owns_kfd_handle', False): kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
+
+  def _unmap(self, mem):
    gpus = (ctypes.c_int32 * 1)(self.gpu_id)
    stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1)
    assert stm.n_success == 1
-    if mem.owner == self.dev:
-      if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
-      kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)

  def map(self, mem):
-    if mem.owner is not None and mem.owner._is_cpu(): return self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
+    if mem.owner is not None and mem.owner._is_cpu():
+      mapped = self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
+      mapped._owns_kfd_handle = True
+      return mapped

    c_gpus = (ctypes.c_int32 * 1)(self.gpu_id)
    stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(c_gpus), n_devices=1)
@@ -467,6 +499,7 @@ class PCIIface(PCIIfaceBase):

  def require_profile_mode(self): return True
  def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return True # TODO: account for WGP disablement on some asics.
+  def unmap(self, mem): self.free(mem)

  def _compute_props(self):
    self.ip_versions = self.dev_impl.ip_ver
@@ -507,11 +540,11 @@ class PCIIface(PCIIfaceBase):
    if drain_only: d.iface.dev_impl.ih.drain()
    else: d.iface.dev_impl.ih.interrupt_handler()

-    if reset and d.iface.dev_impl.recover():
+    if reset and d.iface.dev_impl.recover(force=True):
      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.timeline_signal()._buf.cpu_view().mv.cast('Q')[0] = d.timeline_value().as_memoryview(force_zero_copy=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))):
@@ -525,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):
@@ -535,18 +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.device_id = int(device.split(":")[1]) if ":" in device else 0
-
-    self.iface = self._select_iface()
+    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
@@ -571,12 +633,12 @@ 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
    self.sdma_queues:dict = {}
-    self.has_sdma_queue = True # self.sdma_queue(0) is not None, TODO: think of this
+    self.has_copy_queue = not getenv("AMD_DISABLE_SDMA")

    super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, can_recover=self.is_am(), arch=self.arch)

@@ -584,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()
@@ -631,9 +697,6 @@ class AMDDevice(HCQ2Compiled):
    qname = f"{'COPY' if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA else 'COMPUTE'}:{idx}"
    self.pm_bufferize = PatternMatcher([
      (UPat(Ops.PARAM, tag=f"{qname}_{name}"), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"]
-    ] + [
-      (UPat(Ops.PARAM, tag=f"{qname}_timeline_signal"), lambda ctx, q=qname: ctx[0].timeline_signal(q)),
-      (UPat(Ops.PARAM, tag=f"{qname}_timeline_value"), lambda ctx, q=qname: ctx[0].timeline_value(q)),
    ]) + self.pm_bufferize

    return queue
@@ -647,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))

@@ -655,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):
@@ -692,7 +755,7 @@ class AMDDevice(HCQ2Compiled):
    return tmpring

  def scratch_buffer(self, private_segment_size):
-    private_segment_size = max(private_segment_size, 128)
+    AMDDevice.max_scratch_psize = private_segment_size = max(private_segment_size, 128, AMDDevice.max_scratch_psize)
    if self.max_private_segment_size < private_segment_size:
      lanes_per_wave = 64 # wave64
      mem_alignment_size = 256 if self.target[0] != 9 else 1024
@@ -31,12 +31,12 @@ def dname_of(device) -> str:

 def alloc_like(shape, dtype, device, axis=None) -> Tensor:
  if isinstance(device, tuple) and axis is not None:
-    return Tensor(Tensor.invalids(*shard_shape(shape, axis, len(device)), dtype=dtype, device=device).uop.multi(axis), device=device)
+    return Tensor(Tensor.invalids(*shard_shape(shape, axis, len(device)), dtype=dtype, device=device).uop.unshard(axis), device=device)
  return Tensor.invalids(*shape, dtype=dtype, device=device)

 def alloc_local(shape, dtype, device, axis=None) -> Tensor:
  if isinstance(device, tuple) and axis is not None:
-    return Tensor(Tensor.invalids(*shape, dtype=dtype, device=device).uop.multi(0), device=device)
+    return Tensor(Tensor.invalids(*shape, dtype=dtype, device=device).uop.unshard(0), device=device)
  return Tensor.invalids(*shape, dtype=dtype, device=device)

 def compile_hip(src:str, defines:list[str]):
@@ -48,7 +48,7 @@ def _fused_ce_loss_bwd(gradient:UOp, kernel:UOp, label_smoothing:float):
    axis = logits_u.axis
    ndev = len(device)
    local_shape = tuple(s//ndev if i == axis else s for i,s in enumerate((MBS, SEQ, VOCAB)))
-    d_logits = Tensor(Tensor.invalids(*local_shape, dtype=dtypes.bfloat16, device=device).uop.multi(axis), device=device)
+    d_logits = Tensor(Tensor.invalids(*local_shape, dtype=dtypes.bfloat16, device=device).uop.unshard(axis), device=device)
    rows_per_dev = local_shape[0] * local_shape[1]
    seq_per_dev = local_shape[1]
  else:
@@ -74,11 +74,11 @@ def fused_ce_loss(logits:Tensor, targets:Tensor, label_smoothing:float=0.1) -> T
    axis = logits.uop.axis
    assert axis in (0, 1), f"unsupported sharding axis={axis} for CE loss"
    ndev = len(logits.device)
-    loss_out = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0),
+    loss_out = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.unshard(0),
                      device=logits.device)
-    max_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0),
+    max_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.unshard(0),
                      device=logits.device)
-    lse_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0),
+    lse_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.unshard(0),
                      device=logits.device)
    local_shape = tuple(s//ndev if i == axis else s for i,s in enumerate(logits.shape))
    rows_per_dev = local_shape[0] * local_shape[1]
@@ -16,7 +16,7 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:

  wg = UOp.range(NUM_WG, 0, AxisType.GLOBAL)
  tid = UOp.range(THREADS_PER_WG, 1, AxisType.LOCAL)
-  it = UOp.range((n_elems // VEC) // (NUM_WG * THREADS_PER_WG), 2, AxisType.LOOP)
+  it = UOp.range((n_elems // VEC) // (NUM_WG * THREADS_PER_WG), 2, AxisType.WEAK)
  lane = UOp.range(VEC, 3, AxisType.UNROLL)

  idx = (((it * NUM_WG + wg) * THREADS_PER_WG + tid) * VEC) + lane
@@ -36,21 +36,21 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:
  lmax_val = lmax.after(lmax_store.end(it))[0]

  lds = UOp.placeholder((THREADS_PER_WG,), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
-  lds = lds.after(lds[tid].store(lmax_val).barrier())
+  lds = lds.after(lds[tid].store(lmax_val))

  step = THREADS_PER_WG // 2
  while step:
    active = tid < step
    other = lds[(tid + step).valid(active)].load()
-    lds = lds.after(lds[tid.valid(active)].store(lds[tid].maximum(other)).barrier())
+    lds = lds.after(lds[tid.valid(active)].store(lds[tid].maximum(other)))
    step //= 2

  device = device[0].split(":")[0] if isinstance(device, tuple) else device.split(":")[0]
  if device in {"AMD", "NULL"}: atomic_arg = "if ({2} > {3}) __hip_atomic_fetch_max((int*){0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);"
  else: raise NotImplementedError(f"no atomic max for device {device}")
-  amax_idx = amax_out.reshape((1,)).index(UOp.const(dtypes.weakint, 0))
+  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)
  return atomic.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=()))

@functools.cache
@@ -0,0 +1,36 @@
+import functools, math, pathlib
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from extra.llama_kernels import alloc_like, compile_hip
+
+@functools.cache
+def _custom_quantize_mxfp4(row_fp4:UOp, row_scale:UOp, col_fp4:UOp, col_scale:UOp, x:UOp, *, shuffle_row:bool, shuffle_col:bool) -> UOp:
+  M, N = math.prod(x.shape[:-1]), x.shape[-1]
+  assert M % 256 == 0 and N % 256 == 0, f"MXFP4 quantization requires multiples of 256, got {x.shape}"
+  name = f"quantize_mxfp4_{int(shuffle_row)}_{int(shuffle_col)}_{M}_{N}"
+  mem = M*N*2 + M*N + M*N//16  # read bf16, write row+col fp4 + e8m0
+  outputs = (row_fp4, row_scale, col_fp4, col_scale)
+  sink = UOp.sink(*(o.base for o in outputs), x.base,
+                  *(UOp(Ops.CUSTOM, src=(o.base.index(0),), arg="") 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()
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
+    UOp(Ops.BINARY, arg=compile_hip(src, [f"-DKERNEL_NAME={name}", f"-DM_DIM={M}", f"-DN_DIM={N}",
+                                             f"-DSHUFFLE_ROWWISE_FP4_VALUE={int(shuffle_row)}",
+                                             f"-DSHUFFLE_COLWISE_FP4_VALUE={int(shuffle_col)}"]))))
+
+def quantize_mxfp4(x:Tensor, *, shuffle_row:bool=False, shuffle_col:bool=False, flatten_row:bool=False) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+  assert x.dtype == dtypes.bfloat16 and x.ndim >= 2, f"expected BF16 matrix, got {x.dtype} {x.shape}"
+  M, N = math.prod(x.shape[:-1]), x.shape[-1]
+  assert M % 256 == 0 and N % 256 == 0, f"MXFP4 quantization requires multiples of 256, got {x.shape}"
+  axis = x.uop.axis if isinstance(x.device, tuple) else None
+  row_axis = 0 if flatten_row and axis is not None else axis
+  col_axis = None if axis is None else (0 if axis == x.ndim-1 else 1)
+  outputs = (alloc_like((M, N//2) if flatten_row else (*x.shape[:-1], N//2), dtypes.uint8, x.device, row_axis),
+             alloc_like((M, N//32) if flatten_row else (*x.shape[:-1], N//32), dtypes.uint8, x.device, row_axis),
+             alloc_like((N, M//2), dtypes.uint8, x.device, col_axis),
+             alloc_like((N, M//32), dtypes.uint8, x.device, col_axis))
+  fxn = functools.partial(_custom_quantize_mxfp4, shuffle_row=shuffle_row, shuffle_col=shuffle_col)
+  return tuple(Tensor.custom_kernel(*outputs, x, fxn=fxn)[:4])
@@ -0,0 +1,226 @@
+// Copyright (c) 2025-2026, Advanced Micro Devices, Inc. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#include <hip/hip_runtime.h>
+#include <cstdint>
+
+#if !defined(KERNEL_NAME) || !defined(M_DIM) || !defined(N_DIM) || !defined(SHUFFLE_ROWWISE_FP4_VALUE) || \
+    !defined(SHUFFLE_COLWISE_FP4_VALUE)
+#error kernel dimensions and layouts must be defined
+#endif
+
+namespace {
+
+constexpr int BLOCK = 32;
+constexpr int TILE_M = 128;
+constexpr int TILE_N = 64;
+constexpr int THREADS = 256;
+constexpr int THREADS_PER_ROW = 8;
+constexpr int VALUES_PER_THREAD = 4;
+constexpr int SMEM_STRIDE = BLOCK + 2;
+constexpr int M = M_DIM;
+constexpr int N = N_DIM;
+constexpr int M_PACKED = M / 2;
+constexpr int N_PACKED = N / 2;
+constexpr int M_SCALES = M / BLOCK;
+constexpr int N_SCALES = N / BLOCK;
+constexpr bool SHUFFLE_ROWWISE_FP4 = SHUFFLE_ROWWISE_FP4_VALUE;
+constexpr bool SHUFFLE_COLWISE_FP4 = SHUFFLE_COLWISE_FP4_VALUE;
+
+static_assert(M % 256 == 0 && N % 256 == 0);
+
+struct Quantized4 {
+  uint16_t fp4;
+  uint8_t scale;
+};
+
+__device__ __forceinline__ float swizzle_xor1(float value) {
+  float result;
+  asm volatile("ds_swizzle_b32 %0, %1 offset:0x041f\n\ts_waitcnt lgkmcnt(0)" : "=v"(result) : "v"(value));
+  return result;
+}
+
+__device__ __forceinline__ float swizzle_xor2(float value) {
+  float result;
+  asm volatile("ds_swizzle_b32 %0, %1 offset:0x081f\n\ts_waitcnt lgkmcnt(0)" : "=v"(result) : "v"(value));
+  return result;
+}
+
+__device__ __forceinline__ float swizzle_xor4(float value) {
+  float result;
+  asm volatile("ds_swizzle_b32 %0, %1 offset:0x101f\n\ts_waitcnt lgkmcnt(0)" : "=v"(result) : "v"(value));
+  return result;
+}
+
+__device__ __forceinline__ float max8(float value) {
+  value = fmaxf(value, swizzle_xor4(value));
+  value = fmaxf(value, swizzle_xor2(value));
+  return fmaxf(value, swizzle_xor1(value));
+}
+
+__device__ __forceinline__ float4 load_bf16x4(const uint16_t* values) {
+  const uint32_t lo = *reinterpret_cast<const uint32_t*>(values);
+  const uint32_t hi = *reinterpret_cast<const uint32_t*>(values + 2);
+  return make_float4(__uint_as_float(lo << 16), __uint_as_float(lo & 0xffff0000u),
+                     __uint_as_float(hi << 16), __uint_as_float(hi & 0xffff0000u));
+}
+
+__device__ __forceinline__ void hadamard16(float4& value, int lane) {
+  const float a0 = value.x + value.y, a1 = value.x - value.y;
+  const float a2 = value.z + value.w, a3 = value.z - value.w;
+  value = make_float4(a0 + a2, a1 + a3, a0 - a2, a1 - a3);
+
+  const float4 xor1 = make_float4(swizzle_xor1(value.x), swizzle_xor1(value.y), swizzle_xor1(value.z), swizzle_xor1(value.w));
+  value = lane & 1 ? make_float4(xor1.x - value.x, xor1.y - value.y, xor1.z - value.z, xor1.w - value.w)
+                   : make_float4(xor1.x + value.x, xor1.y + value.y, xor1.z + value.z, xor1.w + value.w);
+
+  const float4 xor2 = make_float4(swizzle_xor2(value.x), swizzle_xor2(value.y), swizzle_xor2(value.z), swizzle_xor2(value.w));
+  value = lane & 2 ? make_float4(xor2.x - value.x, xor2.y - value.y, xor2.z - value.z, xor2.w - value.w)
+                   : make_float4(xor2.x + value.x, xor2.y + value.y, xor2.z + value.z, xor2.w + value.w);
+  value.x *= 0.25f;
+  value.y *= 0.25f;
+  value.z *= 0.25f;
+  value.w *= 0.25f;
+}
+
+__device__ __forceinline__ uint8_t e8m0_scale(float amax, float& scale) {
+  if (amax == 0.0f) {
+    scale = 1.0f;
+    return 127;
+  }
+
+  const uint32_t rounded = (__float_as_uint(amax) + 0x200000u) & 0xff800000u;
+  int exponent = static_cast<int>((rounded >> 23) & 0xff) - 129;
+  exponent = exponent < -127 ? -127 : exponent > 127 ? 127 : exponent;
+  scale = exponent == -127 ? __uint_as_float(0x00400000u) : __uint_as_float(static_cast<uint32_t>(exponent + 127) << 23);
+  return static_cast<uint8_t>(exponent + 127);
+}
+
+__device__ __forceinline__ uint16_t pack_fp4(float4 value, float scale) {
+  uint32_t lo = 0, hi = 0;
+  asm volatile("v_cvt_scalef32_pk_fp4_f32 %0, %1, %2, %3" : "+v"(lo) : "v"(value.x), "v"(value.y), "v"(scale));
+  asm volatile("v_cvt_scalef32_pk_fp4_f32 %0, %1, %2, %3" : "+v"(hi) : "v"(value.z), "v"(value.w), "v"(scale));
+  return static_cast<uint16_t>(lo | (hi << 8));
+}
+
+__device__ __forceinline__ Quantized4 quantize(float4 value, int lane) {
+  hadamard16(value, lane);
+  const float local_max = fmaxf(fmaxf(fabsf(value.x), fabsf(value.y)), fmaxf(fabsf(value.z), fabsf(value.w)));
+  float scale;
+  const uint8_t e8m0 = e8m0_scale(max8(local_max), scale);
+  return {pack_fp4(value, scale), e8m0};
+}
+
+__device__ __forceinline__ void store_scale(uint8_t* output, int row, int col, int cols, uint8_t value) {
+  const int tile = ((row >> 5) * (cols >> 3) + (col >> 3)) << 8;
+  const int offset = ((col & 3) << 6) + ((row & 15) << 2) + (((col >> 2) & 1) << 1) + ((row >> 4) & 1);
+  output[tile + offset] = value;
+}
+
+template<bool Shuffled>
+__device__ __forceinline__ void store_fp4(uint8_t* output, int row, int col, int packed_cols, uint16_t value) {
+  int index = row * packed_cols + col;
+  if constexpr (Shuffled) {
+    const int tile = (row >> 4) * (packed_cols << 4) + (col >> 5) * 512;
+    const int offset = ((col >> 4) & 1) * 256 + (row & 15) * 16 + (col & 15);
+    index = tile + offset;
+  }
+  *reinterpret_cast<uint16_t*>(output + index) = value;
+}
+
+__device__ __forceinline__ void load_tile(uint16_t* tile, const uint16_t* input, int tile_m, int tile_n) {
+  const int row = threadIdx.x / THREADS_PER_ROW;
+  const int col = threadIdx.x % THREADS_PER_ROW * VALUES_PER_THREAD;
+  const uint64_t packed = *reinterpret_cast<const uint64_t*>(input + (tile_m + row) * N + tile_n + col);
+  *reinterpret_cast<uint32_t*>(tile + row * SMEM_STRIDE + col) = static_cast<uint32_t>(packed);
+  *reinterpret_cast<uint32_t*>(tile + row * SMEM_STRIDE + col + 2) = static_cast<uint32_t>(packed >> 32);
+}
+
+__device__ __forceinline__ void quantize_row(uint16_t* tile, uint8_t* fp4_output, uint8_t* scale_output,
+                                             int tile_m, int tile_n, int local_row, int lane) {
+  const int row = tile_m + local_row;
+  const int col = lane * VALUES_PER_THREAD;
+  const Quantized4 result = quantize(load_bf16x4(tile + local_row * SMEM_STRIDE + col), lane);
+  store_fp4<SHUFFLE_ROWWISE_FP4>(fp4_output, row, (tile_n + col) / 2, N_PACKED, result.fp4);
+  if (lane == 0) store_scale(scale_output, row, tile_n / BLOCK, N_SCALES, result.scale);
+}
+
+__device__ __forceinline__ Quantized4 quantize_col(uint16_t* tile, int col, int lane) {
+  const int row = lane * VALUES_PER_THREAD;
+  return quantize(make_float4(
+    __uint_as_float(static_cast<uint32_t>(tile[(row + 0) * SMEM_STRIDE + col]) << 16),
+    __uint_as_float(static_cast<uint32_t>(tile[(row + 1) * SMEM_STRIDE + col]) << 16),
+    __uint_as_float(static_cast<uint32_t>(tile[(row + 2) * SMEM_STRIDE + col]) << 16),
+    __uint_as_float(static_cast<uint32_t>(tile[(row + 3) * SMEM_STRIDE + col]) << 16)), lane);
+}
+
+} // namespace
+
+extern "C" __global__ __launch_bounds__(THREADS, 8)
+void KERNEL_NAME(uint8_t* __restrict__ rowwise_fp4, uint8_t* __restrict__ rowwise_scale,
+                 uint8_t* __restrict__ colwise_fp4, uint8_t* __restrict__ colwise_scale,
+                 const uint16_t* __restrict__ input) {
+  __shared__ uint16_t tile[BLOCK * SMEM_STRIDE];
+  const int tid = threadIdx.x;
+  const int line = tid / THREADS_PER_ROW;
+  const int lane = tid % THREADS_PER_ROW;
+  const int block_m = blockIdx.x * TILE_M;
+  const int block_n = blockIdx.y * TILE_N;
+
+  if constexpr (!SHUFFLE_COLWISE_FP4) {
+    uint16_t col_fp4[TILE_N / BLOCK][TILE_M / BLOCK];
+    uint8_t col_scale[TILE_N / BLOCK][TILE_M / BLOCK];
+
+    for (int chunk_m = 0; chunk_m < TILE_M / BLOCK; chunk_m++) {
+      for (int chunk_n = 0; chunk_n < TILE_N / BLOCK; chunk_n++) {
+        const int tile_m = block_m + chunk_m * BLOCK;
+        const int tile_n = block_n + chunk_n * BLOCK;
+        load_tile(tile, input, tile_m, tile_n);
+        __syncthreads();
+
+        quantize_row(tile, rowwise_fp4, rowwise_scale, tile_m, tile_n, line, lane);
+        const Quantized4 result = quantize_col(tile, line, lane);
+        col_fp4[chunk_n][chunk_m] = result.fp4;
+        col_scale[chunk_n][chunk_m] = result.scale;
+        __syncthreads();
+      }
+    }
+
+    for (int chunk_n = 0; chunk_n < TILE_N / BLOCK; chunk_n++) {
+      for (int chunk_m = 0; chunk_m < TILE_M / BLOCK; chunk_m++)
+        tile[line * BLOCK + chunk_m * THREADS_PER_ROW + lane] = col_fp4[chunk_n][chunk_m];
+      __syncthreads();
+
+      for (int round = 0; round < BLOCK / THREADS_PER_ROW; round++) {
+        const int col = round * THREADS_PER_ROW + tid / BLOCK;
+        const int row_pair = tid % BLOCK;
+        *reinterpret_cast<uint16_t*>(colwise_fp4 + (block_n + chunk_n * BLOCK + col) * M_PACKED + block_m / 2 + row_pair * 2) =
+          tile[col * BLOCK + row_pair];
+      }
+
+      if (lane == 0) {
+        const int col = block_n + chunk_n * BLOCK + line;
+        for (int chunk_m = 0; chunk_m < TILE_M / BLOCK; chunk_m++)
+          store_scale(colwise_scale, col, block_m / BLOCK + chunk_m, M_SCALES, col_scale[chunk_n][chunk_m]);
+      }
+      __syncthreads();
+    }
+  } else {
+    for (int chunk_m = 0; chunk_m < TILE_M / BLOCK; chunk_m++) {
+      for (int chunk_n = 0; chunk_n < TILE_N / BLOCK; chunk_n++) {
+        const int tile_m = block_m + chunk_m * BLOCK;
+        const int tile_n = block_n + chunk_n * BLOCK;
+        load_tile(tile, input, tile_m, tile_n);
+        __syncthreads();
+
+        quantize_row(tile, rowwise_fp4, rowwise_scale, tile_m, tile_n, line, lane);
+        const int row = lane * VALUES_PER_THREAD;
+        const int col = tile_n + line;
+        const Quantized4 result = quantize_col(tile, line, lane);
+        store_fp4<true>(colwise_fp4, col, (tile_m + row) / 2, M_PACKED, result.fp4);
+        if (lane == 0) store_scale(colwise_scale, col, tile_m / BLOCK, M_SCALES, result.scale);
+        __syncthreads();
+      }
+    }
+  }
+}
@@ -0,0 +1,49 @@
+import functools, math
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, KernelInfo
+from tinygrad.renderer import Estimates
+from extra.llama_kernels import alloc_like
+
+LOG2E = 1.4426950408889634
+
+@functools.cache
+def _custom_swiglu(out:UOp, x_w13:UOp) -> UOp:
+  rows, hidden = math.prod(x_w13.shape[:-1]), x_w13.shape[-1]//2
+  n_elems = rows * hidden
+  out, x_w13 = out.reshape(n_elems), x_w13.reshape(rows, 2*hidden)
+  i = UOp.range(n_elems, 0)
+  row, col = i // hidden, i % hidden
+  act, gate = x_w13[row, col].cast(dtypes.float), x_w13[row, hidden+col].cast(dtypes.float)
+  sigmoid = (1.0 + (-LOG2E * act).exp2()).reciprocal()
+  store = out[i].store((act * sigmoid * gate).cast(out.dtype))
+  return store.end(i).sink(arg=KernelInfo(f"swiglu_fwd_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=6*n_elems)))
+
+@functools.cache
+def _custom_swiglu_bwd(grad_out:UOp, x_w13:UOp, grad_act:UOp) -> UOp:
+  rows, hidden = math.prod(x_w13.shape[:-1]), x_w13.shape[-1]//2
+  n_elems = rows * hidden
+  grad_out, x_w13, grad_act = grad_out.reshape(rows, 2*hidden), x_w13.reshape(rows, 2*hidden), grad_act.reshape(n_elems)
+  i = UOp.range(n_elems, 0)
+  row, col = i // hidden, i % hidden
+  act, gate = x_w13[row, col].cast(dtypes.float), x_w13[row, hidden+col].cast(dtypes.float)
+  grad = grad_act[i].cast(dtypes.float)
+  sigmoid = (1.0 + (-LOG2E * act).exp2()).reciprocal()
+  silu = act * sigmoid
+  dact = grad_out[row, col].store((grad * (sigmoid + silu * (1.0 - sigmoid)) * gate).cast(grad_out.dtype))
+  dgate = grad_out.after(dact)[row, hidden+col].store((grad * silu).cast(grad_out.dtype))
+  return dgate.end(i).sink(arg=KernelInfo(f"swiglu_bwd_{n_elems}", estimates=Estimates(ops=10*n_elems, mem=10*n_elems)))
+
+def _swiglu_bwd(gradient:UOp, kernel:UOp):
+  _, x_w13 = kernel.src[1:]
+  axis = x_w13.axis if isinstance(x_w13.device, tuple) else None
+  grad_out = alloc_like(x_w13.shape, dtypes.bfloat16, x_w13.device, axis)
+  grad_out, *_ = Tensor.custom_kernel(grad_out, Tensor(x_w13, device=x_w13.device), Tensor(gradient, device=x_w13.device),
+                                      fxn=_custom_swiglu_bwd)
+  return (None, grad_out.uop)
+
+def swiglu(x_w13:Tensor) -> Tensor:
+  assert x_w13.dtype == dtypes.bfloat16 and x_w13.ndim >= 2 and x_w13.shape[-1] % 32 == 0
+  *prefix, two_k = x_w13.shape
+  axis = x_w13.uop.axis if isinstance(x_w13.device, tuple) else None
+  out = alloc_like((*prefix, two_k//2), dtypes.bfloat16, x_w13.device, axis)
+  return Tensor.custom_kernel(out, x_w13, fxn=_custom_swiglu, grad_fxn=_swiglu_bwd)[0]
@@ -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)
@@ -4,7 +4,7 @@ from hexdump import hexdump
 from copy import deepcopy
 import pathlib, sys
 from tinygrad.helpers import to_mv, getenv
-from tinygrad.runtime.autogen import adreno
+from tinygrad.runtime.autogen import mesa
 sys.path.append(pathlib.Path(__file__).parent.parent.parent.as_posix())

 IOCTL = getenv("IOCTL", 0)
@@ -23,7 +23,7 @@ for child in xml.getroot():
 CAPTURED_STATE = {}

 REGS = {}
-for k, v in adreno.__dict__.items():
+for k, v in mesa.__dict__.items():
  if k.startswith("REG_") and isinstance(v, int) and v > 1024: REGS[v] = k

 from extra.qcom_gpu_driver import msm_kgsl
@@ -42,7 +42,7 @@ def get_struct(argp, stype):

 def format_struct(s):
  sdats = []
-  for field_name, *_ in s._real_fields_:
+  for field_name, *_ in s._fields_:
    if field_name in {"__pad", "PADDING_0"}: continue
    dat = getattr(s, field_name)
    if isinstance(dat, int): sdats.append(f"{field_name}:0x{dat:X}")
@@ -96,9 +96,9 @@ def parse_cmd_buf(dat):
        CAPTURED_STATE['LOAD_FRAGS'].append((state_block, state_type, num_unit, dst_off))

        if state_block == SB6_CS_SHADER:
-          from extra.disassemblers.adreno import disasm_raw
+          from tinygrad.runtime.support.compiler_mesa import disas_adreno
          if state_type == ST6_SHADER and IOCTL > 3:
-            disasm_raw(get_mem(((vals[2] << 32) | vals[1]), num_unit * 128))
+            disas_adreno(get_mem(((vals[2] << 32) | vals[1]), num_unit * 128))
          if state_type == ST6_CONSTANTS:
            x = get_mem(((vals[2] << 32) | vals[1]), num_unit*4)
            CAPTURED_STATE['constants'] = x[:]
@@ -142,7 +142,7 @@ def parse_cmd_buf(dat):
      vals = struct.unpack("I"*size, dat[ptr+4:ptr+4+4*size])
      if IOCTL > 0: print(f"{ptr:3X} -- typ 4: {size=:3d}, {reg_name}", hprint(vals))
      for vi,v in enumerate(vals): CAPTURED_STATE[offset+vi] = v
-      if offset == adreno.REG_A6XX_SP_CS_CONFIG:
+      if offset == mesa.REG_A6XX_SP_CS_CONFIG:
        val = vals[0]
        if IOCTL > 0:
          print(f"\tBINDLESS_TEX={(val >> 0) & 0b1}")
@@ -215,79 +215,3 @@ def install_hook(c_function, python_function):

 libc = ctypes.CDLL(ctypes.util.find_library("libc"))
 install_hook(libc.ioctl, ioctl)
-
-def before_launch():
-  global CAPTURED_STATE
-  CAPTURED_STATE.clear()
-def collect_last_launch_state():
-  global CAPTURED_STATE
-  return deepcopy(CAPTURED_STATE)
-def compare_launch_state(state, good_state):
-  cmp = [
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_NTEX__MASK),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_NSAMP__MASK),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_NIBO__MASK),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_ENABLED),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_BINDLESS_TEX),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_BINDLESS_SAMP),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_BINDLESS_IBO),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_BINDLESS_UBO),
-
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_HALFREGFOOTPRINT__MASK),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_FULLREGFOOTPRINT__MASK),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_BRANCHSTACK__MASK),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_FULLREGFOOTPRINT__MASK),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_THREADMODE__MASK),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_EARLYPREAMBLE),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_MERGEDREGS),
-
-    (adreno.REG_A6XX_SP_CS_PVT_MEM_PARAM, adreno.A6XX_SP_CS_PVT_MEM_PARAM_MEMSIZEPERITEM__MASK),
-    (adreno.REG_A6XX_SP_CS_PVT_MEM_PARAM, adreno.A6XX_SP_CS_PVT_MEM_PARAM_HWSTACKSIZEPERTHREAD__MASK),
-
-    (adreno.REG_A6XX_SP_CS_UNKNOWN_A9B1, adreno.A6XX_SP_CS_UNKNOWN_A9B1_UNK5),
-    (adreno.REG_A6XX_SP_CS_UNKNOWN_A9B1, adreno.A6XX_SP_CS_UNKNOWN_A9B1_UNK6),
-
-    (adreno.REG_A6XX_SP_CS_BRANCH_COND, 0xffffffff),
-
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_0, adreno.A6XX_HLSQ_CS_NDRANGE_0_KERNELDIM__MASK),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_0, adreno.A6XX_HLSQ_CS_NDRANGE_0_LOCALSIZEX__MASK),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_0, adreno.A6XX_HLSQ_CS_NDRANGE_0_LOCALSIZEY__MASK),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_0, adreno.A6XX_HLSQ_CS_NDRANGE_0_LOCALSIZEZ__MASK),
-
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_1, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_2, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_3, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_4, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_5, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_6, 0xffffffff),
-
-    (adreno.REG_A6XX_HLSQ_CS_CNTL_0, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_CNTL_1, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_KERNEL_GROUP_X, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_KERNEL_GROUP_Y, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_KERNEL_GROUP_Z, 0xffffffff),
-  ]
-
-  for x,m in cmp:
-    print(f"Field {REGS[x]}, mask: 0x{m:X} cmp: {state.get(x, 0) & m} vs {good_state.get(x, 0) & m}")
-    if state.get(x, 0) & m != good_state.get(x, 0) & m:
-      return False, f"Field {REGS[x]}, mask: 0x{m:X} mismatch: {state.get(x, 0) & m} vs {good_state.get(x, 0) & m}"
-
-  for n in ['descriptors', 'ibos']:
-    if n not in good_state: continue
-    mv1, mv2 = state.get(n), good_state.get(n)
-
-    if len(mv1) != len(mv2): return False, f"{n}: len mismatch {len(mv1)} != {len(mv2)}"
-    mv1 = memoryview(bytearray(mv1)).cast('I')
-    mv2 = memoryview(bytearray(mv2)).cast('I')
-    for i in range(len(mv2)):
-      if i % 8 == 5 or i % 8 == 4: continue # addresses
-      if mv1[i]!=mv2[i]: return False, f"{n}: content mismatch {i} {mv1[i]} {mv2[i]}"
-
-  for n in ['samplers']:
-    if n not in good_state: continue
-    mv1, mv2 = state.get(n), good_state.get(n)
-    if len(mv1) != len(mv2): return False, f"{n}: len mismatch {len(mv1)} != {len(mv2)}"
-    if any(mv1[i]!=mv2[i] for i in range(len(mv1))): return False, f"{n}: content mismatch"
-
-  return True, "PASS"
@@ -0,0 +1,217 @@
+# Runbook: Llama 3 8B Training on DigitalOcean MI350X
+
+## Machine Specs
+- 8x MI350X GPUs (gfx950, device ID 75b0), 288GB VRAM each
+- 2TB RAM, 192 CPUs, 2TB disk
+- ROCm 7.14 at `/opt/rocm` (NOT `/opt/rocm-7.1.1` like the submission scripts assume)
+- Python 3.12
+
+## Phase 1: System Setup
+
+### 1.1 Install packages
+```bash
+apt-get update
+apt-get install -y python3-pip python3-venv git tmux rclone clang
+```
+
+### 1.2 Install Python deps
+```bash
+python3 -m pip install --break-system-packages --ignore-installed typing-extensions numpy tqdm wandb tiktoken sentencepiece
+```
+Note: `--ignore-installed typing-extensions` is needed because the base image ships typing-extensions 4.10.0 without a RECORD file, so pip cannot uninstall it.
+
+### 1.3 Install ROCm dev headers
+The base image has ROCm runtime but NOT the HIP dev headers. Need:
+```bash
+apt-get install -y amdrocm-core-dev
+```
+This installs `hip/hip_runtime.h` at `/opt/rocm/core-7.14/include/hip/hip_runtime.h`.
+The symlink `/opt/rocm/include` → `/opt/rocm/core-7.14/include` makes it available at `/opt/rocm/include/hip/hip_runtime.h`.
+
+### 1.4 Configure ROCm comgr
+ROCm 7.14 ships comgr 3.3 at `/opt/rocm/lib/libamd_comgr.so`. tinygrad's DLL loader needs explicit env vars to find it (it searches for `libcomgr.so*` by default, not `libamd_comgr.so*`). Set these in the run command:
+```bash
+export COMGR_PATH=/opt/rocm/lib/libamd_comgr.so
+export COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so
+```
+Also add ROCm libs to ldconfig so comgr's shared library dependencies resolve:
+```bash
+cat > /etc/ld.so.conf.d/rocm.conf << 'EOF'
+/opt/rocm/lib
+/opt/rocm/lib/llvm/lib
+/opt/rocm/lib/rocm_sysdeps/lib
+EOF
+ldconfig
+```
+
+### 1.5 Install geohot tmux config
+```bash
+curl -sL https://raw.githubusercontent.com/geohot/configuration/master/.tmux.conf -o ~/.tmux.conf
+```
+
+### 1.6 Reload amdgpu driver
+tinygrad's HCQ backend needs `/dev/kfd` which is created by the amdgpu kernel driver.
+If the driver was unloaded, reload it:
+```bash
+modprobe amdgpu
+ls /dev/kfd  # should exist
+```
+
+## Phase 2: Clone tinygrad
+```bash
+cd /root
+git clone https://github.com/tinygrad/tinygrad.git
+cd tinygrad
+python3 -m pip install --break-system-packages -e .
+```
+
+## Phase 3: Download C4 Dataset
+
+The C4 data is on the MLCommons Cloudflare R2 bucket in Megatron-LM indexed format.
+
+```bash
+rclone config create mlc-training s3 provider=Cloudflare \
+  access_key_id=76ea42eadb867e854061a1806220ee1e \
+  secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 \
+  endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
+
+mkdir -p /raid/datasets/c4-8b
+rclone copy mlc-training:mlcommons-training-wg-public/llama3_1/datasets/c4/llama3_1_8b/ /raid/datasets/c4-8b/ -P
+```
+
+Files downloaded (~85GB total, ~6 minutes):
+- `c4-train.en_6_text_document.bin` (79 GB)
+- `c4-train.en_6_text_document.idx` (870 MB)
+- `c4-validation-91205-samples.en_text_document.bin` (159 MB)
+- `c4-validation-91205-samples.en_text_document.idx` (1.8 MB)
+- `LICENSE.txt`, `NOTICE.txt`
+
+**Wait for rclone to fully complete before starting training.** Starting training while the dataset is still downloading will read a truncated .bin file, causing `ValueError: all input arrays must have the same shape` in the dataloader. The stale `.index_cache` and `.blend_cache` files must also be deleted if this happens:
+```bash
+rm -f /raid/datasets/c4-8b/*.index_cache /raid/datasets/c4-8b/*.blend_cache
+```
+
+## Phase 4: wandb Login
+```bash
+wandb login
+```
+Enter API key from https://wandb.ai/authorize
+
+Alternatively, pass the key directly:
+```bash
+wandb login <API_KEY>
+```
+
+## Phase 5: Run Training
+
+Run training in tmux so it survives SSH disconnects:
+```bash
+tmux new-session -d -s train 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=AMD:HIP ROCM_PATH=/opt/rocm WANDB=1 bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh 2>&1 | tee /root/train.log'
+```
+Attach with `tmux attach -t train`.
+
+### 5.1 Smoke test (beam search, 2 layers, real data)
+Always run beam first to validate the pipeline:
+```bash
+tmux new-session -d -s beam 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=AMD:HIP ROCM_PATH=/opt/rocm bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh 2>&1 | tee /root/beam.log'
+```
+
+The beam test runs 10 training steps with 2 layers. Expected results:
+- ~0.29s per step after warmup
+- ~700K GFLOPS, ~7% MFU (low because only 2 layers)
+- ~380 GB VRAM used
+- Loss stable at ~12.55 with random init
+
+### 5.2 Full training run
+```bash
+tmux new-session -d -s train 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=AMD:HIP ROCM_PATH=/opt/rocm WANDB=1 bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh 2>&1 | tee /root/train.log'
+```
+
+## Environment Variable Reference
+
+| Variable | Value | Why |
+|---|---|---|
+| `COMGR_PATH` | `/opt/rocm/lib/libamd_comgr.so` | tinygrad's DLL loader needs explicit path to find comgr 3.3 |
+| `COMGR_3_PATH` | `/opt/rocm/lib/libamd_comgr.so` | comgr 3.x uses a separate `comgr_3` module with its own path var |
+| `CC` | `/opt/rocm/core-7.14/lib/llvm/bin/clang` | System clang doesn't know gfx950; must use ROCm's bundled clang |
+| `DEV` | `AMD:HIP` | Force HIPRenderer (comgr-based) over HIPCCRenderer (hipcc subprocess) |
+| `ROCM_PATH` | `/opt/rocm` | Script defaults to `/opt/rocm-7.1.1` which doesn't exist |
+| `WANDB` | `1` | Enable wandb logging (off by default) |
+
+## Architecture
+
+| Component | Source file |
+|---|---|
+| Model | `examples/mlperf/models/flat_llama.py` — FlatTransformer, FP8 MXFP4 weights, fused QKV, flash attention |
+| Trainer | `examples/mlperf/model_train.py` → `train_llama3()` |
+| Optimizer | `examples/mlperf/optim.py` — GradAccClipAdamW, master weights, FP8 re-quant |
+| LR schedule | `examples/mlperf/lr_schedulers.py` — CosineAnnealingLRWithWarmup |
+| Dataloader | `examples/mlperf/dataloader.py` — Megatron-LM indexed bin format |
+| ASM GEMM | `extra/gemm/cdna_asm_gemm.py` — gfx950 MFMA assembly, MXFP4 |
+| Flash attention | `extra/thunder/amd/fa.py` |
+| Fused kernels | `extra/llama_kernels/` — rmsnorm, silu, quantize, fused_ce |
+| GPU driver | `tinygrad/runtime/ops_amd.py` — HCQ, direct KFD ioctl |
+| Renderer | `tinygrad/renderer/cstyle.py` — HIPRenderer for gfx950 |
+| comgr compiler | `tinygrad/runtime/support/compiler_amd.py` — HIPCompiler using comgr 3.3 |
+
+## Troubleshooting
+
+### `'hip/hip_runtime.h' file not found`
+Install `amdrocm-core-dev`:
+```bash
+apt-get install -y amdrocm-core-dev
+```
+
+### `'gfx950' is not a recognized processor` + LLVM crash
+System clang doesn't know gfx950. Set `CC=/opt/rocm/core-7.14/lib/llvm/bin/clang`.
+
+### `comgr not available: try setting COMGR_PATH?`
+Add ROCm libs to ldconfig and set `COMGR_PATH` and `COMGR_3_PATH`:
+```bash
+# /etc/ld.so.conf.d/rocm.conf should contain /opt/rocm/lib paths
+ldconfig
+```
+
+### `comgr not available: try setting COMGR_3_PATH?`
+comgr 3.x uses a separate module. Set `COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so` too.
+
+### `No such file or directory: 'clang'`
+Install clang: `apt-get install -y clang` (for CPU compilation).
+For gfx950 HIP compilation, comgr (not clang) is used — ensure the ROCm 7.14 comgr 3.3 is properly loaded via `COMGR_PATH` and `COMGR_3_PATH`.
+
+## Appendix: KVM Virtualization Observations
+
+### Virtualization detection
+```
+$ systemd-detect-virt
+kvm
+$ lspci -nn | grep AMD
+83:00.0 ... Device [1002:75b0]
+```
+CPU flags include `hypervisor`. `dmesg` shows `Hypervisor detected: KVM`.
+
+### Working path: amdgpu driver (KFDIface)
+The amdgpu driver loads on boot and binds to all 8 GPUs, creating `/dev/kfd` and 64 renderD nodes (`/dev/dri/renderD128` through `/dev/dri/renderD191`). tinygrad's `KFDIface` enumerates GPUs through `/sys/devices/virtual/kfd/kfd/topology/nodes` and uses `/dev/kfd` for ioctl. No PCI device ID patching is needed — the KFD path does not use `PCIIface` or `AMDev._run_discovery()`.
+
+This is the working configuration. No code changes to tinygrad are required.
+
+### PCIIface path (does not work on this VM)
+For reference, the `PCIIface` path was also explored but does not work in this KVM guest:
+
+- `PCIIface` in `ops_amd.py` does not list device ID `0x75b0`. Adding it allows PCI detection but `AMDev._run_discovery()` fails because the VRAM BAR reads all `0xFF`.
+- This was observed with the GPU unbound from any driver, after PCI reset, and with VFIO bound.
+- VFIO binding (`vfio-pci` with `enable_unsafe_noiommu_mode=1`) succeeded but VRAM BAR still reads all `0xFF`.
+- No IOMMU in guest — `dmesg` has no `AMD-Vi` entries, PCI devices have no `iommu_group` symlink.
+
+### amdgpu driver behavior
+On first boot, amdgpu loaded and bound to all 8 GPUs. On one boot it failed to initialize:
+```
+[  799.780369] amdgpu 0000:83:00.0: Failed to alloc msi vectors
+[  799.781476] amdgpu 0000:83:00.0: sw_init of IP block <vega20_ih> failed -22
+[  799.782724] amdgpu 0000:83:00.0: amdgpu_device_ip_init failed
+[  799.793885] amdgpu 0000:83:00.0: Fatal error during GPU init
+```
+On a subsequent boot, amdgpu initialized successfully (SMU initialized, VRAM ready). After unbinding all 8 GPUs from amdgpu, `rmmod amdgpu` wedged the module (stuck in "Unloading" state in `/proc/modules`), requiring a full VM reboot.
+
+### No fan control
+No `fan*` or `pwm*` hwmon entries exist. Only `temp*`, `power*`, `freq*` are exposed. GPU temps read 56-63°C, power ~265W per GPU.
@@ -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}")
@@ -192,7 +192,7 @@ def unpack_insts(viz_data, i:int, j:int, data:dict) -> dict:
    prev_instr = max(prev_instr, e.time + e.dur)
  summary = [{"label":"Total Cycles", "value":w.end_time-w.begin_time}, {"label":"SE", "value":w.se}, {"label":"CU", "value":w.cu},
             {"label":"SIMD", "value":w.simd}, {"label":"Wave ID", "value":w.wave_id}, {"label":"Run number", "value":data["run_number"]}]
-  return {"rows":[tuple(v.values()) for v in rows.values()], "cols":columns, "metadata":[summary], "ref":viz_data.ref_map.get(data["prg"].name)}
+  return {"rows":[tuple(v.values()) for v in rows.values()], "cols":columns, "metadata":[summary],"ref":viz_data.ref_map.get(data["prg"].profile_key)}

 def print_data(data:dict) -> None:
  from tabulate import tabulate
@@ -2,7 +2,7 @@ import math, pathlib, functools, struct

 from tinygrad import Device, Tensor
 from tinygrad.dtype import DTypeLike, dtypes
-from tinygrad.helpers import DEBUG
+from tinygrad.helpers import DEBUG, getenv
 from tinygrad.renderer import Estimates
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 from tinygrad.runtime.support.elf import elf_loader
@@ -14,21 +14,38 @@ def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None
  shard_axis = ref.uop.axis if axis is None else axis
  shape = tuple(s // len(ref.device) if i == shard_axis else s for i, s in enumerate(shape))
  axis = ref.uop.axis if axis is None else axis
-  return Tensor(Tensor.invalids(*shape, dtype=dtype, device=ref.device).uop.multi(axis), dtype=dtype, device=ref.device)
+  return Tensor(Tensor.invalids(*shape, dtype=dtype, device=ref.device).uop.unshard(axis), dtype=dtype, device=ref.device)

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

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

-@functools.cache
-def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink):
+def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink, window=0):
  def grad(dou:UOp, ker:UOp) -> tuple:
    do = Tensor(dou, device=dou.device)
    attn = Tensor(ker.src[1].after(ker), device=ker.src[1].device)
@@ -129,7 +145,7 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
    delta_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
    delta_vec, dq = Tensor.custom_kernel(delta_vec, dq, attn, do, fxn=functools.partial(custom_fa_backward_pre, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:2]

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

    if D == 64:
      dq = dq.reshape(B, H, N//16, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2).permute(0, 1, 2, 8, 9, 10, 11, 3, 4, 6, 7, 5, 12).reshape(B, H, N, D).transpose(1, 2)
@@ -149,7 +165,7 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
  return grad

 # TODO: remove write_flat once scheduler can remove reshapes between custom_kernel. TestCustomKernel.test_simple_reshape
-def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False, write_flat:bool=False, sinks:Tensor|None=None):
+def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False, write_flat:bool=False, sinks:Tensor|None=None, window:int=0):
  assert attn_mask is None, "attn_mask not supported"
  assert is_causal, "only causal attention supported"

@@ -176,18 +192,18 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
  attn = _sharded_empty((B, N, H * D), xq, axis=shard_axis) if write_flat else _sharded_empty_like(xq, axis=shard_axis)
  l_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)

-  grad = _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink)
+  grad = _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink, window=window)

  fwd_inputs = (attn, l_vec, xq, xk, xv) + ((sinks,) if has_sink else ())
-  attn, l_vec = Tensor.custom_kernel(*fwd_inputs, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, has_sink=has_sink), grad_fxn=grad)[:2]
+  attn, l_vec = Tensor.custom_kernel(*fwd_inputs, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, has_sink=has_sink, window=window), grad_fxn=grad)[:2]

  return attn, attn, l_vec

@functools.cache
-def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None, *, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, has_sink:bool=True):
+def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None, *, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, has_sink:bool=True, window:int=0):
  code = (pathlib.Path(__file__).parent / "fa_fwd_causal.cpp").read_text()
  compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math",
-                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DATTN_SINK={int(has_sink)}"]
+                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DATTN_SINK={int(has_sink)}", f"-DWINDOW={window}"]

  Q_BLOCK_SIZE = 32
  NUM_WARPS = 8
@@ -206,10 +222,11 @@ def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None
                  arg=KernelInfo(name="custom_fa_forward", estimates=estimates))

  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-  lib = bytearray(lib)
-  rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
-  struct.pack_into('<I', lib, rodata_off, 160000)
-  lib = bytes(lib)
+  if not getenv("NO_HIPCC"):
+    lib = bytearray(lib)
+    rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
+    struct.pack_into('<I', lib, rodata_off, 160000)
+    lib = bytes(lib)

  return UOp(Ops.PROGRAM,
             src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
@@ -236,19 +253,20 @@ def custom_fa_backward_pre(delta_vec:UOp, dq:UOp, o:UOp, do:UOp, device:str, arc
                  arg=KernelInfo(name="custom_fa_backward_pre", estimates=estimates))

  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-  lib = bytearray(lib)
-  rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
-  struct.pack_into('<I', lib, rodata_off, 160000)
-  lib = bytes(lib)
+  if not getenv("NO_HIPCC"):
+    lib = bytearray(lib)
+    rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
+    struct.pack_into('<I', lib, rodata_off, 160000)
+    lib = bytes(lib)

  return UOp(Ops.PROGRAM,
             src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))

@functools.cache
-def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_vec:UOp, delta_vec:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
+def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_vec:UOp, delta_vec:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, window:int=0):
  code = (pathlib.Path(__file__).parent / "fa_bwd_causal.cpp").read_text()
  compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math",
-                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}"]
+                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DWINDOW={window}"]

  BLOCK_SIZE_KV = 256
  GROUP_SIZE = H // H_KV
@@ -268,10 +286,11 @@ def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_ve
                  arg=KernelInfo(name="custom_fa_backward", estimates=estimates))

  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-  lib = bytearray(lib)
-  rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
-  struct.pack_into('<I', lib, rodata_off, 160000)
-  lib = bytes(lib)
+  if not getenv("NO_HIPCC"):
+    lib = bytearray(lib)
+    rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
+    struct.pack_into('<I', lib, rodata_off, 160000)
+    lib = bytes(lib)

  return UOp(Ops.PROGRAM,
             src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
@@ -298,10 +317,11 @@ def custom_fa_backward_post(dq_out:UOp, dq_in:UOp, device:str, arch:str, B:int,
                  arg=KernelInfo(name="custom_fa_backward_post", estimates=estimates))

  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-  lib = bytearray(lib)
-  rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
-  struct.pack_into('<I', lib, rodata_off, 160000)
-  lib = bytes(lib)
+  if not getenv("NO_HIPCC"):
+    lib = bytearray(lib)
+    rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
+    struct.pack_into('<I', lib, rodata_off, 160000)
+    lib = bytes(lib)

  return UOp(Ops.PROGRAM,
             src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
@@ -43,6 +43,10 @@ constexpr int SLICE_QO = 32;
 constexpr int DOT_SLICE_QO = 16;
 constexpr int WARP_SIZE_KV = 64; // warp size for KV
 constexpr bool causal = true;
+// WINDOW>0: sliding-window backward (query i sees keys in [i-WINDOW+1, i])
+#ifndef WINDOW
+#define WINDOW 0
+#endif

 #define NUM_WARPS 4
 #define NUM_THREADS (kittens::WARP_THREADS * NUM_WARPS)
@@ -88,7 +92,12 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
  const int k_start_min = j_min * WARP_SIZE_KV;
  // first Q step that can overlap this K_span:
  const int first_step = max(0, k_start_min / STEP_QO);
+#if WINDOW
+  // cap the Q loop, padded by 2 masked steps: the epilogue's deferred dq path miscomputes in-window tail queries
+  const int num_steps_per_head = min(total_steps_per_head - first_step, (BLOCK_SIZE_KV + WINDOW) / STEP_QO + 2);
+#else
  const int num_steps_per_head = total_steps_per_head - first_step;
+#endif
  const int num_steps = num_steps_per_head * HEADS_PER_WG;
  const int k_pos = j * WARP_SIZE_KV;

@@ -380,6 +389,13 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
            mov<0, 1, neg_inf_v>(P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
+#if WINDOW
+          // window lower boundary, mirror of the causal edge
+          } else if (q_pos - k_pos == WINDOW) {
+            make_window<0, 0, neg_inf_v>(P_ij, P_ij);
+          } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -533,6 +549,8 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
+        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
+        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -638,6 +656,13 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
            make_causal<0, 1, neg_inf_v>(P_ij, P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
+#if WINDOW
+          } else if (q_pos - k_pos == WINDOW) {
+            mov<0, 0, neg_inf_v>(P_ij);
+            make_window<0, 1, neg_inf_v>(P_ij, P_ij);
+          } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -791,6 +816,8 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
+        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
+        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -895,6 +922,14 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
            // Apply the causal mask to [0, 2] and set [0, 3:4] to -inf
            make_causal<0, 2, neg_inf_v>(P_ij, P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
+#if WINDOW
+          } else if (q_pos - k_pos == WINDOW) {
+            mov<0, 0, neg_inf_v>(P_ij);
+            mov<0, 1, neg_inf_v>(P_ij);
+            make_window<0, 2, neg_inf_v>(P_ij, P_ij);
+          } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -1048,6 +1083,8 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
+        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
+        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -1151,6 +1188,15 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
          } else if (q_pos == k_pos) {
            // Apply the causal mask to [0, 3]
            make_causal<0, 3, neg_inf_v>(P_ij, P_ij);
+#if WINDOW
+          } else if (q_pos - k_pos == WINDOW) {
+            mov<0, 0, neg_inf_v>(P_ij);
+            mov<0, 1, neg_inf_v>(P_ij);
+            mov<0, 2, neg_inf_v>(P_ij);
+            make_window<0, 3, neg_inf_v>(P_ij, P_ij);
+          } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -1303,6 +1349,8 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
+        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
+        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -1428,6 +1476,13 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
            mov<0, 1, neg_inf_v>(P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
+#if WINDOW
+          // window lower boundary, mirror of the causal edge
+          } else if (q_pos - k_pos == WINDOW) {
+            make_window<0, 0, neg_inf_v>(P_ij, P_ij);
+          } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -1582,6 +1637,8 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
+        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
+        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -1689,6 +1746,13 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
            make_causal<0, 1, neg_inf_v>(P_ij, P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
+#if WINDOW
+          } else if (q_pos - k_pos == WINDOW) {
+            mov<0, 0, neg_inf_v>(P_ij);
+            make_window<0, 1, neg_inf_v>(P_ij, P_ij);
+          } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -1842,6 +1906,8 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
+        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
+        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -1946,6 +2012,14 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
            // Apply the causal mask to [0, 2] and set [0, 3:4] to -inf
            make_causal<0, 2, neg_inf_v>(P_ij, P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
+#if WINDOW
+          } else if (q_pos - k_pos == WINDOW) {
+            mov<0, 0, neg_inf_v>(P_ij);
+            mov<0, 1, neg_inf_v>(P_ij);
+            make_window<0, 2, neg_inf_v>(P_ij, P_ij);
+          } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -2099,6 +2173,8 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
+        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
+        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -2202,6 +2278,15 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
          } else if (q_pos == k_pos) {
            // Apply the causal mask to [0, 3]
            make_causal<0, 3, neg_inf_v>(P_ij, P_ij);
+#if WINDOW
+          } else if (q_pos - k_pos == WINDOW) {
+            mov<0, 0, neg_inf_v>(P_ij);
+            mov<0, 1, neg_inf_v>(P_ij);
+            mov<0, 2, neg_inf_v>(P_ij);
+            make_window<0, 3, neg_inf_v>(P_ij, P_ij);
+          } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -2354,6 +2439,8 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
+        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
+        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -2471,6 +2558,12 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
            mov<0, 1, neg_inf_v>(P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
+#if WINDOW
+        } else if (q_pos - k_pos == WINDOW) {
+            make_window<0, 0, neg_inf_v>(P_ij, P_ij);
+        } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
        }
      }
      mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -2732,6 +2825,13 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
            make_causal<0, 1, neg_inf_v>(P_ij, P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
+#if WINDOW
+        } else if (q_pos - k_pos == WINDOW) {
+            mov<0, 0, neg_inf_v>(P_ij);
+            make_window<0, 1, neg_inf_v>(P_ij, P_ij);
+        } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
        }
      }
      mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -2988,6 +3088,14 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
            // Apply the causal mask to [0, 2] and set [0, 3:4] to -inf
            make_causal<0, 2, neg_inf_v>(P_ij, P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
+#if WINDOW
+        } else if (q_pos - k_pos == WINDOW) {
+            mov<0, 0, neg_inf_v>(P_ij);
+            mov<0, 1, neg_inf_v>(P_ij);
+            make_window<0, 2, neg_inf_v>(P_ij, P_ij);
+        } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
        }
      }
      mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -3244,6 +3352,15 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
        } else if (q_pos == k_pos) {
            // Apply the causal mask to [0, 3]
            make_causal<0, 3, neg_inf_v>(P_ij, P_ij);
+#if WINDOW
+        } else if (q_pos - k_pos == WINDOW) {
+            mov<0, 0, neg_inf_v>(P_ij);
+            mov<0, 1, neg_inf_v>(P_ij);
+            mov<0, 2, neg_inf_v>(P_ij);
+            make_window<0, 3, neg_inf_v>(P_ij, P_ij);
+        } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
        }
      }
      mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -34,6 +34,10 @@ constexpr int ATTN_D = 128; // dimension
 constexpr int Q_BLOCK_SIZE = 32; // q block size
 constexpr int KV_BLOCK_SIZE = 64; // kv block size
 constexpr bool causal = true;
+// WINDOW>0: sliding-window attention, query i attends keys in [i-WINDOW+1, i]
+#ifndef WINDOW
+#define WINDOW 0
+#endif

 #define NUM_WARPS 8
 #define NUM_THREADS (kittens::WARP_THREADS * NUM_WARPS)
@@ -82,11 +86,26 @@ template<typename T=float, typename L=col_l, typename S=rt_16x32_4_s> using attn

 /**********************************************************/
 template<int THR_X, int THR_Y>
-__device__ inline void mask_vec2_imm(uint32_t rel_vgpr, uint32_t neg_inf_vgpr,
+__device__ inline void mask_vec2_imm(uint32_t rel_vgpr, uint32_t rel_hi_vgpr, uint32_t neg_inf_vgpr,
                                     uint32_t& x_ref, uint32_t& y_ref) {

    uint64_t x_mask, y_mask;
-    // uint32_t ox, oy;
+#if WINDOW
+    // causal+window in one asm block to not disturb register allocation
+    asm volatile(
+        "v_cmp_lt_i32_e64 %0, %4, %5\n\t"
+        "v_cmp_lt_i32_e64 %1, %4, %7\n\t"
+        "v_cndmask_b32_e64 %2, %2, %6, %0\n\t"
+        "v_cndmask_b32_e64 %3, %3, %6, %1\n\t"
+        "v_cmp_ge_i32_e64 %0, %8, %5\n\t"
+        "v_cmp_ge_i32_e64 %1, %8, %7\n\t"
+        "v_cndmask_b32_e64 %2, %2, %6, %0\n\t"
+        "v_cndmask_b32_e64 %3, %3, %6, %1\n\t"
+        : "=s"(x_mask), "=s"(y_mask), "+v"(x_ref), "+v"(y_ref)
+        : "v"(rel_vgpr), "n"(THR_X), "v"(neg_inf_vgpr), "n"(THR_Y), "v"(rel_hi_vgpr)
+        : "vcc"
+    );
+#else
    asm volatile(
        // x: rel < THR_X ?
        "v_cmp_lt_i32_e64 %0, %6, %7\n\t"
@@ -99,7 +118,7 @@ __device__ inline void mask_vec2_imm(uint32_t rel_vgpr, uint32_t neg_inf_vgpr,
          "n"(THR_X), "v"(neg_inf_vgpr), "n"(THR_Y)
        : "vcc"
    );
-    // x_ref = ox; y_ref = oy;
+#endif
 }

 template<ducks::rt::col_layout RT>
@@ -122,6 +141,8 @@ __device__ inline void mask_kv_tile(RT &dst, int q_abs, int k_abs, uint32_t neg_
        // (smaller rel ⇒ more "future" keys that must be -inf)
        const int rel0 = q_pos - (k_base + row_base);
        const uint32_t rel = static_cast<uint32_t>(rel0);
+        // rel-WINDOW keeps THR within the inline-constant range
+        const uint32_t rel_hi = static_cast<uint32_t>(rel0 - WINDOW);

        #pragma unroll
        for (int j = 0; j < dst.width; ++j) {
@@ -145,14 +166,14 @@ __device__ inline void mask_kv_tile(RT &dst, int q_abs, int k_abs, uint32_t neg_
            //  - reuse a single neg_inf register
            //  - keep VCC live across the pair
            //  - avoid reloading -inf or recomputing rel
-            mask_vec2_imm< 0, 1 >(rel, neg_inf_v, d0x, d0y);
-            mask_vec2_imm< 2, 3 >(rel, neg_inf_v, d1x, d1y);
-            mask_vec2_imm< 8, 9 >(rel, neg_inf_v, d2x, d2y);
-            mask_vec2_imm<10,11 >(rel, neg_inf_v, d3x, d3y);
-            mask_vec2_imm<16,17 >(rel, neg_inf_v, d4x, d4y);
-            mask_vec2_imm<18,19 >(rel, neg_inf_v, d5x, d5y);
-            mask_vec2_imm<24,25 >(rel, neg_inf_v, d6x, d6y);
-            mask_vec2_imm<26,27 >(rel, neg_inf_v, d7x, d7y);
+            mask_vec2_imm< 0, 1 >(rel, rel_hi, neg_inf_v, d0x, d0y);
+            mask_vec2_imm< 2, 3 >(rel, rel_hi, neg_inf_v, d1x, d1y);
+            mask_vec2_imm< 8, 9 >(rel, rel_hi, neg_inf_v, d2x, d2y);
+            mask_vec2_imm<10,11 >(rel, rel_hi, neg_inf_v, d3x, d3y);
+            mask_vec2_imm<16,17 >(rel, rel_hi, neg_inf_v, d4x, d4y);
+            mask_vec2_imm<18,19 >(rel, rel_hi, neg_inf_v, d5x, d5y);
+            mask_vec2_imm<24,25 >(rel, rel_hi, neg_inf_v, d6x, d6y);
+            mask_vec2_imm<26,27 >(rel, rel_hi, neg_inf_v, d7x, d7y);
        }
    }
 }
@@ -201,6 +222,16 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    else max_num_tiles = num_tiles;
    const int q_start_pos = tile_idx * Q_BLOCK_SIZE;

+#if WINDOW
+    // start at the first in-window tile; clamp keeps >=4 tiles for the pipeline unroll
+    const int block_min_q = block_tile_idx * NUM_WARPS * Q_BLOCK_SIZE;
+    int min_tile = (block_min_q - WINDOW + 1) / KV_BLOCK_SIZE;
+    if (min_tile < 0) min_tile = 0;
+    if (min_tile > max_num_tiles - 4) min_tile = max(0, max_num_tiles - 4);
+#else
+    constexpr int min_tile = 0;
+#endif
+
    constexpr float TEMPERATURE_SCALE = (D == 128) ? 0.08838834764f*1.44269504089f : 0.125f*1.44269504089f;
    uint32_t neg_inf_v = 0xff800000;

@@ -231,21 +262,23 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    G::prefill_swizzled_offsets<1, false>(k_smem[0], g.Kg, swizzled_offsets_K);
    G::prefill_swizzled_offsets<1, false>(v_smem[0], g.Vg, swizzled_offsets_V);

-    G::load<1, false>(k_smem[0], g.Kg, {batch_idx, 0, head_idx_kv, 0}, swizzled_offsets_K);
+    G::load<1, false>(k_smem[0], g.Kg, {batch_idx, min_tile, head_idx_kv, 0}, swizzled_offsets_K);
    __builtin_amdgcn_s_waitcnt(0);
    __builtin_amdgcn_sched_barrier(0);
    __builtin_amdgcn_s_barrier();

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

    // All warps then collaboratively load in the first slice of V (V0) and the second slice of K (K1) into shared memory
-    G::load<1, false>(k_smem[1], g.Kg, {batch_idx, 1, head_idx_kv, 0}, swizzled_offsets_K);
+    G::load<1, false>(k_smem[1], g.Kg, {batch_idx, min_tile + 1, head_idx_kv, 0}, swizzled_offsets_K);
    // All warps then load in the first slice of K (K0)
-    G::load<1, false>(v_smem[0], g.Vg, {batch_idx, 0, head_idx_kv, 0}, swizzled_offsets_V);
+    G::load<1, false>(v_smem[0], g.Vg, {batch_idx, min_tile, head_idx_kv, 0}, swizzled_offsets_V);
    load(k_reg, k_smem[0]);
    __builtin_amdgcn_sched_barrier(0);
    asm volatile("s_waitcnt lgkmcnt(0)");
@@ -257,15 +290,25 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    zero(att_block[0]);
    transpose(k_reg_transposed, k_reg);
    mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
+    #if WINDOW
+    mul(att_block[0], att_block[0], TEMPERATURE_SCALE);
+    #endif
    __builtin_amdgcn_sched_barrier(0);
    if constexpr (causal) {
-        const int kv_end_pos = (1) * KV_BLOCK_SIZE;
-        if (__builtin_expect(q_start_pos < kv_end_pos, 0)) {  // Only mask if needed
-            mask_kv_tile(att_block[0], tile_idx, 0, neg_inf_v, lane);
+        const int kv_end_pos = (min_tile + 1) * KV_BLOCK_SIZE;
+        if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, WINDOW ? 1 : 0)) {
+            mask_kv_tile(att_block[0], tile_idx, min_tile, neg_inf_v, lane);
        }
    }
    // Each warp performs a partial softmax of QK0 (i.e. some of the online softmax up until but not including the second exponential scaling of the attention block likely)
+#if WINDOW
+    // floor the max: min_tile can be fully masked, which would NaN via exp2(-inf - -inf)
+    zero(max_vec_prev);
+    add(max_vec_prev, max_vec_prev, -1e4f);
+    col_max(max_vec, att_block[0], max_vec_prev);
+#else
    col_max(max_vec, att_block[0]);
+#endif

    copy(max_vec_prev, max_vec);
    exp2(scale_vec, scale_vec);
@@ -284,21 +327,28 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    // All warps then load in the second slice of K (K1)
    load(k_reg, k_smem[1]);
    // All warps then collaboratively load in the third slice of K (K2) into shared memory
-    G::load<1, false>(k_smem[0], g.Kg, {batch_idx, 2, head_idx_kv, 0}, swizzled_offsets_K);
+    G::load<1, false>(k_smem[0], g.Kg, {batch_idx, min_tile + 2, head_idx_kv, 0}, swizzled_offsets_K);
    // All warps then collaboratively load in the second slice of V (V1) into shared memory
-    G::load<1, false>(v_smem[1], g.Vg, {batch_idx, 1, head_idx_kv, 0}, swizzled_offsets_V);
+    G::load<1, false>(v_smem[1], g.Vg, {batch_idx, min_tile + 1, head_idx_kv, 0}, swizzled_offsets_V);
    asm volatile("s_waitcnt lgkmcnt(0)");
    asm volatile("s_waitcnt vmcnt(" FA_VM4 ")");
    __builtin_amdgcn_sched_barrier(0);
    __builtin_amdgcn_s_barrier();

    // hot loop
-    for (int j = 3; j < max_num_tiles - 1; j += 2) {
+    for (int j = min_tile + 3; j < max_num_tiles - 1; j += 2) {
        // Cluster 0:
        //      QK1
        zero(att_block[1]);
        transpose(k_reg_transposed, k_reg);
        mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]);
+        #if WINDOW
+        mul(att_block[1], att_block[1], TEMPERATURE_SCALE);
+        #endif
+#if WINDOW
+        // window masks interior tiles that causal skips
+        mask_kv_tile(att_block[1], tile_idx, j - 2, neg_inf_v, lane);
+#endif
        //      Finish softmax for QK0
        exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]);
        mul(norm_vec, norm_vec, scale_vec);
@@ -359,6 +409,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
        zero(att_block[0]);
        transpose(k_reg_transposed, k_reg);
        mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
+        #if WINDOW
+        mul(att_block[0], att_block[0], TEMPERATURE_SCALE);
+        #endif
        //      Finish softmax for QK1
        exp2(att_block[1].tiles[1][0], att_block[1].tiles[1][0]);
        mul(norm_vec, norm_vec, scale_vec);
@@ -379,7 +432,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
        load(v_reg, v_smem[1]);
        if constexpr (causal) {
            const int kv_end_pos = (j) * KV_BLOCK_SIZE;
-            if (q_start_pos < kv_end_pos) {  // Only mask if needed
+            if (WINDOW || q_start_pos < kv_end_pos) {
                mask_kv_tile(att_block[0], tile_idx, j - 1, neg_inf_v, lane);
            }
        }
@@ -427,6 +480,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    zero(att_block[1]);
    transpose(k_reg_transposed, k_reg);
    mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]);
+    #if WINDOW
+    mul(att_block[1], att_block[1], TEMPERATURE_SCALE);
+    #endif
    //      Finish softmax for QK2
    exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]);
    mul(norm_vec, norm_vec, scale_vec);
@@ -447,7 +503,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    load(v_reg, v_smem[0]);
    if constexpr (causal) {
        const int kv_end_pos = (max_num_tiles - 2) * KV_BLOCK_SIZE;
-        if (__builtin_expect(q_start_pos < kv_end_pos, 0)) {  // Only mask if needed
+        if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, WINDOW ? 1 : 0)) {
            mask_kv_tile(att_block[1], tile_idx, max_num_tiles - 3, neg_inf_v, lane);
        }
    }
@@ -493,6 +549,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    zero(att_block[0]);
    transpose(k_reg_transposed, k_reg);
    mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
+    #if WINDOW
+    mul(att_block[0], att_block[0], TEMPERATURE_SCALE);
+    #endif
    //      Finish softmax for QK3
    exp2(att_block[1].tiles[1][0], att_block[1].tiles[1][0]);
    mul(norm_vec, norm_vec, scale_vec);
@@ -510,7 +569,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    load(v_reg, v_smem[1]);
    if constexpr (causal) {
        const int kv_end_pos = (max_num_tiles - 1) * KV_BLOCK_SIZE;
-        if (__builtin_expect(q_start_pos < kv_end_pos, 1)) {  // Only mask if needed
+        if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, 1)) {
            mask_kv_tile(att_block[0], tile_idx, max_num_tiles - 2, neg_inf_v, lane);
        }
    }
@@ -555,6 +614,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    zero(att_block[1]);
    transpose(k_reg_transposed, k_reg);
    mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]);
+    #if WINDOW
+    mul(att_block[1], att_block[1], TEMPERATURE_SCALE);
+    #endif
    //      Finish softmax for QK4
    exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]);
    mul(norm_vec, norm_vec, scale_vec);
@@ -572,7 +634,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    load(v_reg, v_smem[0]);
    if constexpr (causal) {
        const int kv_end_pos = (max_num_tiles) * KV_BLOCK_SIZE;
-        if (__builtin_expect(q_start_pos < kv_end_pos, 1)) {  // Only mask if needed
+        if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, 1)) {
            mask_kv_tile(att_block[1], tile_idx, max_num_tiles - 1, neg_inf_v, lane);
        }
    }
@@ -1,69 +0,0 @@
-#include <hip/hip_runtime.h>
-#include <hip/hip_bf16.h>
-
-#ifndef ATTN_B
-#define ATTN_B 2
-#endif
-#ifndef ATTN_N
-#define ATTN_N 8192
-#endif
-#ifndef ATTN_H
-#define ATTN_H 32
-#endif
-#ifndef ATTN_H_KV
-#define ATTN_H_KV 8
-#endif
-#ifndef ATTN_D
-#define ATTN_D 128
-#endif
-#ifndef THREADS_PER_BLOCK
-#define THREADS_PER_BLOCK 256
-#endif
-
-constexpr int GROUP_SIZE = ATTN_H / ATTN_H_KV;
-constexpr int HALF_D = ATTN_D / 2;
-constexpr int PACKED_D = (GROUP_SIZE + 2) * ATTN_D;
-
-extern "C" __global__ __launch_bounds__(THREADS_PER_BLOCK) void
-fused_qkv_rope_forward(
-    __hip_bfloat16*       __restrict__ q,
-    __hip_bfloat16*       __restrict__ k,
-    __hip_bfloat16*       __restrict__ v,
-    const __hip_bfloat16* __restrict__ xqkv,
-    const __hip_bfloat16* __restrict__ freqs_cis) {
-  const int b = blockIdx.x;
-  const int n = blockIdx.y;
-  const int bn = b * ATTN_N + n;
-  const int packed_bn = bn * ATTN_H_KV * PACKED_D;
-  const int q_bn = bn * ATTN_H * ATTN_D;
-  const int kv_bn = bn * ATTN_H_KV * ATTN_D;
-
-  if (threadIdx.x < HALF_D) {
-    const int pair = threadIdx.x;
-    const int even = pair << 1;
-    const float c = static_cast<float>(freqs_cis[((n * HALF_D + pair) * 2) + 0]);
-    const float s = static_cast<float>(freqs_cis[((n * HALF_D + pair) * 2) + 1]);
-
-    for (int kvh = 0; kvh < ATTN_H_KV; kvh++) {
-      const int base = packed_bn + kvh * PACKED_D;
-
-      for (int rep = 0; rep < GROUP_SIZE; rep++) {
-        const int qbase = base + rep * ATTN_D;
-        const int h = kvh * GROUP_SIZE + rep;
-        const float a = static_cast<float>(xqkv[qbase + even]);
-        const float bb = static_cast<float>(xqkv[qbase + even + 1]);
-        const int out = q_bn + h * ATTN_D + even;
-        q[out] = static_cast<__hip_bfloat16>(a * c - bb * s);
-        q[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c);
-      }
-
-      const float a = static_cast<float>(xqkv[base + GROUP_SIZE * ATTN_D + even]);
-      const float bb = static_cast<float>(xqkv[base + GROUP_SIZE * ATTN_D + even + 1]);
-      const int out = kv_bn + kvh * ATTN_D + even;
-      k[out] = static_cast<__hip_bfloat16>(a * c - bb * s);
-      k[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c);
-      v[out] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even];
-      v[out + 1] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even + 1];
-    }
-  }
-}
@@ -97,4 +97,33 @@ __device__ inline static void atomic_pk_add_bf16_with_warpid(const GL &dst, cons
            }(std::make_index_sequence<RT::width>{});
        }.template operator()<Ns>(), ...);
    }(std::make_index_sequence<RT::height>{});
-}
+}
+// make_window: complement of make_causal for the window lower boundary (q_pos-k_pos == WINDOW). masks = ~(causal masks)
+template<int N, int M, int GPR, ducks::art::all T0, ducks::art::all T1>
+__device__ static inline void make_window(T0 &dst, const T1 &src) {
+    static_assert(std::is_same_v<typename T0::T, float> && std::is_same_v<typename T1::T, float>, "Only float to float window mask is supported");
+    static_assert(std::is_same_v<typename T0::layout, typename T1::layout>, "Only same layout is supported");
+    static_assert(std::is_same_v<typename T0::shape, typename T1::shape>, "Only same shape is supported");
+
+    if constexpr (std::is_same_v<typename T0::layout, typename ducks::rt_layout::col> && std::is_same_v<typename T0::shape, typename ducks::rt_shape::rt_16x16>) {
+        using range_type_T0 = ducks::art::get_nth_range_t<typename T0::register_ranges, N * T0::width + M>;
+        using registers_T0 = ducks::art::split_many_t<ducks::art::type_list<range_type_T0>, 1>;
+        using range_type_T1 = ducks::art::get_nth_range_t<typename T1::register_ranges, N * T1::width + M>;
+        using registers_T1 = ducks::art::split_many_t<ducks::art::type_list<range_type_T1>, 1>;
+        static_assert(registers_T0::size == registers_T1::size);
+
+        uint64_t window_mask = 0x1FFF01FF001F0001;
+        macros::v_cndmask_b32_e64<ducks::art::get_nth_range_t<registers_T0, 0>::lo, ducks::art::get_nth_range_t<registers_T1, 0>::lo, GPR>(window_mask);
+
+        window_mask = 0x3FFF03FF003F0003;
+        macros::v_cndmask_b32_e64<ducks::art::get_nth_range_t<registers_T0, 1>::lo, ducks::art::get_nth_range_t<registers_T1, 1>::lo, GPR>(window_mask);
+
+        window_mask = 0x7FFF07FF007F0007;
+        macros::v_cndmask_b32_e64<ducks::art::get_nth_range_t<registers_T0, 2>::lo, ducks::art::get_nth_range_t<registers_T1, 2>::lo, GPR>(window_mask);
+
+        window_mask = 0xFFFF0FFF00FF000F;
+        macros::v_cndmask_b32_e64<ducks::art::get_nth_range_t<registers_T0, 3>::lo, ducks::art::get_nth_range_t<registers_T1, 3>::lo, GPR>(window_mask);
+    } else {
+        static_assert(false, "Unsupported window mask");
+    }
+}
@@ -16,7 +16,7 @@ def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None) -> Tensor:
  if not isinstance(ref.device, tuple): return Tensor.empty(*shape, dtype=ref.dtype, device=ref.device)
  shape = tuple(s // len(ref.device) if i == ref.uop.axis else s for i, s in enumerate(shape))
  axis = ref.uop.axis if axis is None else axis
-  return Tensor(Tensor.empty(*shape, dtype=ref.dtype, device=ref.device).uop.multi(axis), dtype=ref.dtype, device=ref.device)
+  return Tensor(Tensor.empty(*shape, dtype=ref.dtype, device=ref.device).uop.unshard(axis), dtype=ref.dtype, device=ref.device)

 def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
  return _sharded_empty(ref.shape, ref, axis)
@@ -224,7 +224,7 @@ class Group:

      # store to shared memory
      red_local_store = red_local[self.laneid].store(red_reg[0])
-      red_local = red_local.after(red_local_store.barrier()).reshape(red_local.shape)
+      red_local = red_local.after(red_local_store).reshape(red_local.shape)

      # reduce from shared memory
      for inner in self.ker.range(3, axis_type=AxisType.REDUCE, track=False):
@@ -258,7 +258,7 @@ class Group:

      # store to shared memory
      red_local_store = red_local[self.laneid].store(red_reg[0])
-      red_local = red_local.after(red_local_store.barrier()).reshape(red_local.shape)
+      red_local = red_local.after(red_local_store).reshape(red_local.shape)

      # reduce from shared memory
      for inner in self.ker.range(3, axis_type=AxisType.REDUCE, track=False):
@@ -342,7 +342,7 @@ class Group:
          if src.dtype != dst.dtype:
            src_load = src_load.cast(dst.dtype)
          dst_store = dst[*dst_idxs, height, width, srow, scol].store(src_load)
-          dst_store = dst_store.end(height, width, outer, inner).barrier()
+          dst_store = dst_store.end(height, width, outer, inner)
    elif dst.addrspace == AddrSpace.REG and src.addrspace == AddrSpace.GLOBAL and isinstance(dst, RT):
      srcf = src.flatten()
      row_stride = prod(src.shape[axis+1:])
@@ -48,14 +48,14 @@ class Kernel(AbstractContextManager):
  @property
  def warpgroup(self): return self.group(4)

-  def range(self, start:int, end:int=0, step:int=1, axis_type:AxisType=AxisType.LOOP, track:bool=True):
+  def range(self, start:int, end:int=0, step:int=1, axis_type:AxisType=AxisType.WEAK, track:bool=True):
    if end == 0: start, end = 0, start
    rng = _tk_range(start, end, step, axis_type, self.range_id)
    self.range_id += 1
    if track: self.range_stack.append(rng)
    return rng

-  def raw_range(self, end:int=0, axis_type:AxisType=AxisType.LOOP):
+  def raw_range(self, end:int=0, axis_type:AxisType=AxisType.WEAK):
    rng = UOp.range(end, self.range_id, axis_type=axis_type)
    self.range_id += 1
    return rng
@@ -209,7 +209,7 @@ class ST:
    return cls(uop, rows, cols, layout, base_shape, ker)

  def swizzle(self, row, col):
-    swizzled_offset = self.base_shape.swizzle(row, col, self._uop.dtype.scalar())
+    swizzled_offset = self.base_shape.swizzle(row, col, self._uop.dtype)

    row = swizzled_offset // self.base_shape.cols
    col = swizzled_offset % self.base_shape.cols
@@ -3,8 +3,8 @@
 # A002 Function argument `input` is shadowing a Python builtin
 # A006 Lambda argument `input` is shadowing a Python builtin
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.uop.ops import Ops
-from tinygrad.helpers import getenv, prod, strides_for_shape, argfix
+from tinygrad.uop.ops import Ops, GroupOp
+from tinygrad.helpers import getenv, prod, strides_for_shape
 import torch.lib
 TORCH_DEBUG = getenv("TORCH_DEBUG")
 import torch, pathlib, operator, functools, weakref
@@ -18,12 +18,16 @@ def _to_torch_device(device: str): return torch.device("tiny", int(device.partit

 import torch.utils.cpp_extension
 mod = torch.utils.cpp_extension.load(name="custom_device_extension", sources=[str(pathlib.Path(__file__).parent / "wrapped_tensor.cpp")])
+# TODO: this assumes a contiguous source, so PERMUTE/EXPAND/PAD/FLIP are wrong. UOp.contiguous_view_offset does it
+# properly, but it needs a device (these are deviceless)
+alias_ops = GroupOp.Movement | {Ops.BITCAST, Ops.DETACH, Ops.AFTER}
 def calculate_storage_offset(x: Tensor) -> int:
-  offset = 0
-  for u in x.uop.toposort():
-    if u.op == Ops.SHRINK:
+  offset, u = 0, x.uop
+  while u.op in alias_ops:
+    if u.op is Ops.SHRINK:
      u_strides = strides_for_shape(u.src[0].shape)
      for i, (start, _) in enumerate(u.marg): offset += start * u_strides[i]
+    u = u.src[0]
  return offset
 def wrap(x: Tensor, dev: torch.device|None=None) -> torch.Tensor:
  x._strides = strides_for_shape(x.shape) # always recalculate
@@ -69,6 +73,12 @@ def wrap_view_op(fn):
    return wrap(ret)
  return _wrap

+# NOTE: list assignment raises IndexError on an out of range dim, and the index must be a tuple: a list of all ints is one advanced index
+def _index_dim(self, dim, idx):
+  idxs = [slice(None)] * self.ndim
+  idxs[dim] = idx
+  return self[tuple(idxs)]
+
 view_ops = {
  "aten.view": Tensor.reshape,
  "aten._unsafe_view": Tensor.reshape,  # when are views unsafe, and do we care?
@@ -78,15 +88,13 @@ view_ops = {
  "aten.transpose.int": Tensor.transpose,
  "aten.squeeze.dim": Tensor.squeeze,
  "aten.unsqueeze": Tensor.unsqueeze,
-  "aten.select.int": lambda self, dim, idx: self[(slice(None),) * (dim%self.ndim) + (idx,)],
+  "aten.select.int": _index_dim,
  "aten.permute": Tensor.permute,
  "aten.alias": lambda self: self,
  "aten.diagonal": Tensor.diagonal,
+  "aten.slice.Tensor": lambda self, dim=0, start=None, end=None, step=1: _index_dim(self, dim, slice(start, end, step)),
  }

-# torch 2.10 handles this natively
-if tuple(map(int, torch.__version__.split('.')[:2])) < (2, 10): view_ops.update({"aten.detach": Tensor.detach})
-
 for k,v in view_ops.items(): torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(wrap_view_op(v))

 def _get_view_ops(view): return getattr(view, "_view_ops", [])
@@ -95,46 +103,21 @@ def _apply_view_ops(target, ops):
  for fn, args, kwargs in ops: target = fn(target, *args, **kwargs)
  return target

-# similar to https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/InferSize.h
-def _reshape_target_shape(shape:tuple[int, ...], args) -> tuple[int, ...]|None:
-  if not (req := argfix(*args)): return None
-  new_shape, infer_idx = [], -1
-  for i, s in enumerate(req):
-    if s is None: s = shape[i] if i < len(shape) else None
-    if not isinstance(s, int): return None
-    if s == -1:
-      if infer_idx != -1: return None
-      infer_idx = len(new_shape)
-    new_shape.append(s)
-  total = prod(shape)
-  if infer_idx != -1:
-    known = prod(x for x in new_shape if x != -1)
-    if known == 0:
-      if total != 0: return None
-      new_shape[infer_idx] = 0
-    else: new_shape[infer_idx] = total // known
-  return tuple(new_shape) if prod(new_shape) == total else None
-
-# TODO: can we get rid of this? only for test_flatten_reshape_add
+# a chain of reshapes is undone by reshaping the value back to the base
 def _try_simple_reshape_view_write(base: Tensor, view: Tensor, val: Tensor) -> bool:
  if not (ops := _get_view_ops(view)): return False
-  shapes = [base.shape]
-  for fn, args, _ in ops:
-    if fn is Tensor.reshape:
-      if not (next_shape := _reshape_target_shape(shapes[-1], args)): return False
-      shapes.append(next_shape)
-  if shapes[-1] != view.shape: return False
-  for s in reversed(shapes[:-1]): val = val.reshape(s)
-  base.assign(val)
+  if any(fn is not Tensor.reshape for fn, _, _ in ops): return False
+  base.assign(val.reshape(base.shape))
  return True

 def _view_write(base: Tensor, view: Tensor, value: Tensor) -> None:
  val = value if value.dtype == base.dtype else value.cast(base.dtype)
-  if view.shape == base.shape: return base.assign(val)
  if _try_simple_reshape_view_write(base, view, val): return
  idx_base = Tensor.arange(base.numel(), dtype=dtypes.int32).reshape(base.shape)
  idx_view = _apply_view_ops(idx_base, _get_view_ops(view)).reshape(-1)
-  flat_base = base.reshape(base.numel()).contiguous()
+  # clone, not contiguous: contiguous() on a base that already owns its buffer returns the base itself, and scattering
+  # into that is an in-place write to a buffer other tensors still hold, which setitem refuses
+  flat_base = base.reshape(base.numel()).clone()
  flat_base[idx_view] = val.reshape(-1)
  base.assign(flat_base.reshape(base.shape))

@@ -162,11 +145,6 @@ def _index_put_impl_(self, indices, values, accumulate=False, unsafe=False):
 def index_put(self, indices, values, accumulate=False):
  return aten.index_put(self.cpu(), [z.cpu() if isinstance(z, torch.Tensor) else None for z in indices], values.clone().cpu(), accumulate).tiny()

-@torch.library.impl("aten::isin.Tensor_Tensor_out", "privateuseone")
-def isin_tensor_tensor_out(x, y, *, assume_unique=False, invert=False, out=None):
-  result = (unwrap(x).unsqueeze(-1) == unwrap(y).flatten()).any(-1)
-  return out.copy_(wrap(~result if invert else result))
-
@torch.library.impl("aten::randperm.generator_out", "privateuseone")
 def randperm_generator(n, generator=None, out=None):
  if generator is not None: raise NotImplementedError("tinygrad torch backend does not support torch.Generator for randperm")
@@ -220,56 +198,13 @@ def _as_strided(tensor:Tensor, size, stride, storage_offset=0):

@torch.library.impl("aten::as_strided", "privateuseone")
 def as_strided(tensor:torch.Tensor, size, stride, storage_offset=None):
-  storage_offset = storage_offset or tensor.storage_offset()
+  if storage_offset is None: storage_offset = tensor.storage_offset()
  return _as_strided(tensor, size, stride, storage_offset)

@torch.library.impl("aten::_reshape_alias", "privateuseone")
 def _reshape_alias(tensor:torch.Tensor, size, stride):
  return _as_strided(tensor, size, stride)

-@torch.library.impl("aten::empty_strided", "privateuseone")
-def empty_strided(size, stride, dtype, layout=None, device=None, pin_memory=False):
-  if TORCH_DEBUG: print(f"empty_strided {size=} {stride=} {dtype=} {layout=} {device=} {pin_memory=}")
-  ret = Tensor.empty(*size, dtype=_from_torch_dtype(dtype), device=_from_torch_device(device)).contiguous()
-  # TODO: should return with requested strides
-  return wrap(ret)
-
-@torch.library.impl("aten::empty.memory_format", "privateuseone")
-def empty_memory_format(size, dtype=None, layout=None, device=None, pin_memory=False, memory_format=None):
-  if TORCH_DEBUG: print(f"empty.memory_format {size=} {dtype=} {layout=} {device=} {pin_memory=} {memory_format=}")
-  ret = Tensor.empty(*size, dtype=_from_torch_dtype(dtype or torch.get_default_dtype()), device=_from_torch_device(device)).contiguous()
-  return wrap(ret)
-
-@torch.library.impl("aten::max_pool2d_with_indices", "privateuseone")
-def max_pool2d_with_indices(self:torch.Tensor, kernel_size:tuple[int, ...], stride=None, padding=0, dilation=1, ceil_mode=False):
-  # TODO: supprt stride [] in tinygrad?
-  if stride is not None and len(stride) == 0: stride = None
-  ret, idx = unwrap(self).max_pool2d(kernel_size, stride, dilation, padding, ceil_mode, return_indices=True)
-  return (wrap(ret), wrap(idx.cast(dtypes.int64)))
-
-@torch.library.impl("aten::max_pool2d_with_indices_backward", "privateuseone")
-def max_pool2d_with_indices_backward(grad_out:torch.Tensor, self:torch.Tensor, kernel_size:tuple[int, ...], stride=None, padding=0, dilation=1, ceil_mode=False, indices=None):
-  return wrap(Tensor.max_unpool2d(unwrap(grad_out), unwrap(indices), output_size=unwrap(self).shape))
-
-@torch.library.impl("aten::max_unpool2d", "privateuseone")
-def max_unpool2d(self:torch.Tensor, indices:torch.Tensor, output_size):
-  return wrap(unwrap(self).max_unpool2d(unwrap(indices), output_size=output_size))
-
-@torch.library.impl("aten::arange", "privateuseone")
-def arange(end, dtype=None, device=None, pin_memory=None):
-  has_float = isinstance(end, float)
-  return wrap(Tensor.arange(0, end, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))
-
-@torch.library.impl("aten::arange.start", "privateuseone")
-def arange_start(start, end, dtype=None, device=None, pin_memory=None):
-  has_float = any(isinstance(x, float) for x in (start, end))
-  return wrap(Tensor.arange(start, end, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))
-
-@torch.library.impl("aten::arange.start_step", "privateuseone")
-def arange_start_step(start, end, step, dtype=None, device=None, pin_memory=None):
-  has_float = any(isinstance(x, float) for x in (start, end, step))
-  return wrap(Tensor.arange(start, end, step, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))
-
@torch.library.impl("aten::convolution_overrideable", "privateuseone")
 def convolution_overrideable(input, weight, bias, stride, padding, dilation, transposed, output_padding, groups):
  if TORCH_DEBUG >= 1:
@@ -290,12 +225,27 @@ def convolution_backward_overrideable(grad_out, input, weight, stride, padding,
  grads = out.gradient(*[t for t,m in zip([input, weight, bias], output_mask) if m], gradient=grad_out)
  return tuple([wrap(grads.pop(0)) if m else None for m in output_mask])

-@torch.library.impl("aten::slice.Tensor", "privateuseone")
-@wrap_view_op
-def slice_tensor(self, dim=0, start=None, end=None, step=1):
-  slices = [slice(None)] * self.ndim
-  slices[dim] = slice(start, end, step)
-  return self[slices]
+# the functional scatters. without an impl aten falls back to a path that assumes a real storage: "self.has_storage() INTERNAL ASSERT FAILED"
+def _scatter_into(self, src, dim, index):
+  out = unwrap(self).clone()
+  slices = [slice(None)] * out.ndim
+  slices[dim] = index
+  out[slices] = unwrap(src).cast(out.dtype)  # torch casts src to self's dtype, tinygrad setitem demands they already match
+  return wrap(out)
+
+@torch.library.impl("aten::slice_scatter", "privateuseone")
+def slice_scatter(self, src, dim=0, start=None, end=None, step=1): return _scatter_into(self, src, dim, slice(start, end, step))
+
+@torch.library.impl("aten::select_scatter", "privateuseone")
+def select_scatter(self, src, dim, index): return _scatter_into(self, src, dim, index)
+
+@torch.library.impl("aten::diagonal_scatter", "privateuseone")
+def diagonal_scatter(self, src, offset=0, dim1=0, dim2=1):
+  # a diagonal is not one axis, so scatter through the flat indices it picks out
+  base, out = unwrap(self), unwrap(self).clone().reshape(-1)
+  idx = Tensor.arange(base.numel(), dtype=dtypes.int32).reshape(base.shape).diagonal(offset, dim1, dim2).reshape(-1)
+  out[idx] = unwrap(src).cast(base.dtype).reshape(-1)
+  return wrap(out.reshape(base.shape))

@torch.library.impl("aten::slice_backward", "privateuseone")
 def slice_backward(grad_out, input_sizes, dim, start, end, step):
@@ -337,19 +287,14 @@ for dim in [1, 2, 3]:
    torch.library.impl(f"aten::{pad_type}_pad{dim}d", "privateuseone")(functools.partial(pad_forward, mode=mode))
    torch.library.impl(f"aten::{pad_type}_pad{dim}d_backward", "privateuseone")(functools.partial(pad_backward, mode=mode))

-def upsample(self, size, align_corners=False, mode=None): return wrap(Tensor.interpolate(unwrap(self), size, mode=mode, align_corners=align_corners))
+# the schemas are all positional: (self, output_size, align_corners, *scales) for linear, (self, output_size, *scales) for nearest.
+def upsample(self, size, *args, mode=None):
+  return wrap(Tensor.interpolate(unwrap(self), size, mode=mode, align_corners=args[0] if mode == "linear" else False))
 for i,pre in enumerate(["", "bi", "tri"]):
  torch.library.impl(f"aten::upsample_{pre}linear{i+1}d", "privateuseone")(functools.partial(upsample, mode="linear"))
  torch.library.impl(f"aten::upsample_nearest{i+1}d", "privateuseone")(functools.partial(upsample, mode="nearest"))
  torch.library.impl(f"aten::_upsample_nearest_exact{i+1}d", "privateuseone")(functools.partial(upsample, mode="nearest-exact"))

-@torch.library.impl("aten::scatter_add.out", "privateuseone")
-def scatter_add(self, dim, index, src, out):
-  self, index, src, out_unwrapped = unwrap(self), unwrap(index), unwrap(src), unwrap(out)
-  if self.shape == (): _apply_inplace(out_unwrapped, src)
-  else: _apply_inplace(out_unwrapped, Tensor.scatter_reduce(self, dim, index, src, reduce='sum'))
-  return out
-
 def _copy_between_devices(src, dest, cast_dtype, to_device, non_blocking=False):
  if src.is_tiny and dest.is_tiny:
    src_t, dest_t = unwrap(src), unwrap(dest)
@@ -400,15 +345,11 @@ def sort_values(input, dim=-1, descending=False, stable=True, values=None, indic
  _apply_inplace(unwrap(indices), out_indices.cast(dtypes.int64))
  return values, indices

-@torch.library.impl("aten::_linalg_svd", "privateuseone")
-def _linalg_svd(self, full_matrices=False):
-  U, S, Vh = unwrap(self).svd(full_matrices)
-  return wrap(U), wrap(S), wrap(Vh)
-
 # register some decompositions
 from torch._decomp import get_decompositions
 decomps = [
  aten.native_layer_norm_backward,
+  aten.native_group_norm_backward,
  aten.linalg_cross,
  aten.addmm,
  aten.addcmul,
@@ -443,12 +384,20 @@ decomps = [
  aten._softmax_backward_data, aten.embedding_dense_backward,
  aten.linalg_vector_norm,
  aten.binary_cross_entropy, aten.binary_cross_entropy_backward,
+  # the C++ mse/smooth_l1 kernels resize their out tensor, and a tiny tensor has no storage to resize
+  aten.mse_loss, aten.mse_loss_backward,
+  aten.smooth_l1_loss, aten.smooth_l1_loss_backward,
  aten.upsample_nearest2d.out,
+  # NOTE: only the "out" overload, the "vec" one is CompositeImplicitAutograd and overriding it loses the autograd kernel
+  aten.upsample_bicubic2d.out,
+  aten._adaptive_avg_pool2d,
  # activations
  aten.hardswish, aten.hardswish_backward,
  aten.hardtanh, aten.hardtanh_backward,
  aten.gelu, aten.gelu_backward,
-  aten.logical_and,
+  # NOTE: no aten.logical_or here, its decomposition reaches aten.bitwise_or through a path that checks aliasing by
+  # reading storage, which a tiny tensor has none of. it gets a direct impl below instead
+  aten.logical_and, aten.logical_xor,
  aten.randint,
  aten.eye,
  aten.hardsigmoid_backward,
@@ -491,7 +440,7 @@ simple_tensor_methods = [
  # reduce
  "all", "any", "argmax", "argmin", "cumsum", "cumprod",
  # complex
-  "avg_pool2d", "linspace"]
+  "linspace"]

 tiny_backend_out = {**{f"aten.{x}.out":getattr(Tensor,x) for x in simple_tensor_methods}, **{
  "aten.add.out": lambda input,other,alpha=1: input+alpha*other,
@@ -536,6 +485,8 @@ tiny_backend_out = {**{f"aten.{x}.out":getattr(Tensor,x) for x in simple_tensor_
  "aten.where.self_out": Tensor.where,
  "aten.prod.int_out": Tensor.prod,
  "aten.scatter.src_out": Tensor.scatter,
+  "aten.scatter_add.out": lambda self,dim,index,src: src if self.shape == () else Tensor.scatter_reduce(self, dim, index, src, reduce="sum"),
+  "aten.isin.Tensor_Tensor_out": lambda x,y,assume_unique=False,invert=False: (x.unsqueeze(-1)==y.flatten()).any(-1) != invert,
  # NOTE: axis=[] in torch means all, change tinygrad?
  "aten.sum.IntList_out": lambda self,axis,keepdim=False,dtype=None:
    self.sum(axis if axis is None or len(axis) else None, keepdim,
@@ -551,8 +502,9 @@ def wrap_out(f):
    assert out.shape == assigned.shape, f"shape mismatch: {assigned.shape} -> {out.shape}"
    assert out.device == assigned.device or out.device is None or assigned.device is None, f"device mismatch: {assigned.device} -> {out.device}"
    assert out.dtype == assigned.dtype, f"dtype mismatch: {assigned.dtype} -> {out.dtype}"
-    if out.device is None and assigned.device is not None: out.replace(out.empty_like(device=assigned.device))
-    return out.assign(assigned)
+    # writing out= is an in-place write like any other: through the base if it is a view, refreshing any derived views
+    _apply_inplace(out, assigned)
+    return out
  return _wrap_out

 def _inplace_op(t, new_value):
@@ -560,7 +512,14 @@ def _inplace_op(t, new_value):
  else: _apply_inplace(t, new_value)
  return t

-tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
+# the three arange overloads are one function at different arity, and dtype/layout/device/pin_memory are keyword only in all of them
+def _arange(*args, dtype=None, **_):
+  return Tensor.arange(*args, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if any(isinstance(x, float) for x in args) else torch.int64)))
+
+def _empty(size, dtype=None, device=None, **_):
+  return Tensor.empty(*size, dtype=_from_torch_dtype(dtype or torch.get_default_dtype()), device=_from_torch_device(device))
+
+tiny_backend = {**tiny_backend_out, **{
  "aten.remainder.Scalar_Tensor": lambda x,y: x%y,
  "aten.floor_divide": lambda x,y: x//y,
  "aten.floor_divide_.Tensor": lambda x,y: x//y,
@@ -573,8 +532,8 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
  # inplace ops using replace for fusion
  "aten.zero_": lambda x: x.const_like(0),
  "aten.fill_.Scalar": lambda x, y: x.const_like(y),
-  "aten.add_.Tensor": lambda self, other, alpha=1.0: self + other * alpha,
-  "aten.add_.Scalar": lambda self, other, alpha=1.0: self + other * alpha,
+  "aten.add_.Tensor": lambda self, other, alpha=1: self + other * alpha,
+  "aten.add_.Scalar": lambda self, other, alpha=1: self + other * alpha,
  "aten.mul_.Tensor": lambda self, other: self * other,
  "aten.mul_.Scalar": lambda self, other: self * other,
  # relu doesn't have an out form?
@@ -607,7 +566,9 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
  # these don't work in out form, they have size 0
  "aten.abs": Tensor.abs,
  "aten.logical_not": Tensor.logical_not,
-  "aten.logical_or_": lambda x, y: x | y,
+  # compare against zero first: logical_* is bool-valued for any input dtype, while | is bitwise
+  "aten.logical_or": lambda x, y: (x != 0) | (y != 0),
+  "aten.logical_or_": lambda x, y: (x != 0) | (y != 0),
  "aten.multinomial": Tensor.multinomial,
  "aten.masked_fill_.Scalar": lambda self, mask, value: self.masked_fill(mask, value),
  "aten.masked_fill_.Tensor": lambda self, mask, value: self.masked_fill(mask, value),
@@ -616,14 +577,7 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
  "aten.masked_select": Tensor.masked_select,
  "aten.all": Tensor.all,
  "aten.sgn": Tensor.sign,
-  "aten.acos": Tensor.acos,
  "aten.any": Tensor.any,
-  "aten.bitwise_not": Tensor.bitwise_not,
-  "aten.argmax": Tensor.argmax,
-  "aten.argmin": Tensor.argmin,
-  "aten.asinh": Tensor.asinh,
-  "aten.mul": Tensor.mul,
-  "aten.atanh": Tensor.atanh,
  "aten.fill_.Tensor": lambda self, value: self.const_like(value.reshape(()).item()),
  "aten.flip": Tensor.flip,
  "aten.scatter_reduce.two": Tensor.scatter_reduce,
@@ -634,10 +588,22 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
  "aten.add.Tensor": lambda input,other,alpha=1: input+alpha*other,
  "aten.linspace": lambda start, stop, steps, dtype=None, **kwargs:
    Tensor.linspace(start, stop, steps, **({"dtype": _from_torch_dtype(dtype)} if dtype is not None else {})),
+  # the functional copy_. without an impl the fallback segfaults on a tensor with no storage
+  "aten.copy": lambda self,src,non_blocking=False: src.cast(self.dtype).to(self.device).expand(self.shape),
+  "aten.arange": lambda end, **kwargs: _arange(0, end, **kwargs),
+  "aten.arange.start": _arange,
+  "aten.arange.start_step": _arange,
+  # empty_strided takes the strides and drops them: we always allocate contiguous
+  "aten.empty_strided": lambda size, stride, **kwargs: _empty(size, **kwargs),
+  "aten.empty.memory_format": _empty,
+  # TODO: supprt stride [] in tinygrad?
+  "aten.max_pool2d_with_indices": lambda self,kernel_size,stride=None,padding=0,dilation=1,ceil_mode=False: ((r:=Tensor.max_pool2d(self, kernel_size, stride or None, dilation, padding, ceil_mode, return_indices=True))[0], r[1].cast(dtypes.int64)),
+  "aten.max_pool2d_with_indices_backward": lambda grad_out,self,kernel_size,stride=None,padding=0,dilation=1,ceil_mode=False,indices=None: Tensor.max_unpool2d(grad_out, indices, output_size=self.shape),
+  "aten.max_unpool2d": lambda self,indices,output_size: Tensor.max_unpool2d(self, indices, output_size=output_size),
+  "aten._linalg_svd": lambda self,full_matrices=False: Tensor.svd(self, full_matrices),
  "aten.topk": Tensor.topk,
  "aten.constant_pad_nd": lambda self, padding, value=0.0: self.pad(padding, mode="constant", value=value).contiguous(),
-  # TODO: input contiguous is needed to prevent CFGContext circular dependency assertion for shapes >512 (see test_cumsum_arange_large)
-  "aten.cumsum": lambda self, dim: self.contiguous().cumsum(dim),
+  "aten.cumsum": lambda self, dim: self.cumsum(dim),
  "aten.logsumexp": lambda self, axis, keepdim=False: self.logsumexp(axis[0], keepdim=keepdim),
  "aten.roll": Tensor.roll,
  "aten.logcumsumexp": Tensor.logcumsumexp,
@@ -646,41 +612,13 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
    self.ones_like(**{k: v for k, v in {"dtype": _from_torch_dtype(dtype) if dtype else None,
                                        "device": _from_torch_device(device) if device else None}.items() if v is not None}),
  "aten.max.dim": lambda self, dim, keepdim=False: (self.max(dim, keepdim), self.argmax(dim, keepdim).cast(dtype=dtypes.int64)),
+  "aten.min.dim": lambda self, dim, keepdim=False: (self.min(dim, keepdim), self.argmin(dim, keepdim).cast(dtype=dtypes.int64)),
  "aten.cummax": lambda self, dim: ((r := self.cummax(dim))[0], r[1].cast(dtypes.int64)),
  "aten.cummin": lambda self, dim: ((r := self.cummin(dim))[0], r[1].cast(dtypes.int64)),
  "aten.nonzero": Tensor.nonzero,
  "aten.unfold": Tensor.unfold,
 }}

-# operations that need inplace treatment (use _inplace_op instead of wrap_fxn) AKA return original tensor
-inplace_ops = {
-  "aten.zero_",
-  "aten.fill_.Scalar",
-  "aten.fill_.Tensor",
-  "aten.add_.Tensor",
-  "aten.add_.Scalar",
-  "aten.mul_.Tensor",
-  "aten.mul_.Scalar",
-  "aten.floor_divide_.Tensor",
-  "aten.__ilshift__.Scalar",
-  "aten.__irshift__.Scalar",
-  "aten.relu_",
-  "aten.random_",
-  "aten.random_.from",
-  "aten.uniform_",
-  "aten.normal_",
-  "aten.logical_or_",
-  "aten.masked_fill_.Scalar",
-  "aten.masked_fill_.Tensor",
-}
-
-inplace_view_ops = {
-  "aten.squeeze_.dim",
-  "aten.unsqueeze_",
-  "aten.transpose_",
-  "aten.t_",
-}
-
 def wrap_fxn(k,f):
  def nf(*args, **kwargs):
    if TORCH_DEBUG:
@@ -694,7 +632,7 @@ def wrap_fxn(k,f):
    else: raise RuntimeError(f"unknown output type {type(out)}")
  return nf

-def wrap_inplace(k,f):
+def wrap_inplace(f):
  def nf(*args, **kwargs):
    orig = args[0]
    args, kwargs = unwrap_args(args, kwargs)
@@ -702,7 +640,7 @@ def wrap_inplace(k,f):
    return orig
  return nf

-def wrap_inplace_view_op(k,f):
+def wrap_inplace_view_op(f):
  def nf(*args, **kwargs):
    orig = args[0]
    args, kwargs = unwrap_args(args, kwargs)
@@ -735,11 +673,18 @@ def wrap_inplace_view_op(k,f):
    return orig
  return nf

+# the aten schema says how an op is called: an inplace view retargets the view, a writable first arg is inplace,
+# and a writable out arg gets wrap_out's dtype cast, shape assert, and view write-through
 for k,v in tiny_backend.items():
-  if k in inplace_view_ops: wrapper = wrap_inplace_view_op
-  elif k in inplace_ops: wrapper = wrap_inplace
-  else: wrapper = wrap_fxn
-  torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(wrapper(k,v))
+  name, _, overload = k.removeprefix("aten.").partition(".")
+  op = getattr(getattr(aten, name), overload or "default")
+  writes = [a.name for a in op._schema.arguments if a.alias_info is not None and a.alias_info.is_write]
+  if torch.Tag.inplace_view in op.tags: fxn = wrap_inplace_view_op(v)
+  elif writes == [op._schema.arguments[0].name] and op._schema.returns: fxn = wrap_inplace(v)
+  elif not writes: fxn = wrap_fxn(k, v)
+  elif writes == ["out"]: fxn = wrap_fxn(k, wrap_out(v))
+  else: raise RuntimeError(f"{k} writes {writes}: unhandled writable arg in schema")
+  torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(fxn)

@torch.library.impl("aten::equal", "privateuseone")
 def equal(x: torch.Tensor, y: torch.Tensor): return (x==y).all().item()
@@ -775,21 +720,17 @@ def native_batch_norm(input, weight, bias, running_mean, running_var, training,
@torch.library.impl("aten::native_batch_norm_backward", "privateuseone")
 def native_batch_norm_backward(grad_out, input, weight, running_mean, running_var, save_mean, save_invstd, train, eps, output_mask):
  grad_out_t, input_t = unwrap(grad_out), unwrap(input)
-  weight_t = unwrap(weight) if weight is not None else None
-  save_mean_t = unwrap(save_mean)
-  save_invstd_t = unwrap(save_invstd)
-  out = input_t.batchnorm(weight_t, None, save_mean_t, save_invstd_t)
-  targets = [t for t, m in zip([input_t, weight_t], output_mask[:2]) if t is not None and m]
-  if targets:
-    grads = out.gradient(*targets, gradient=grad_out_t)
-    grad_input = grads.pop(0) if output_mask[0] else None
-    grad_weight = grads.pop(0) if output_mask[1] and weight_t is not None else None
-  else:
-    grad_input, grad_weight = None, None
-  grad_bias = grad_out_t.sum(axis=tuple(x for x in range(grad_out_t.ndim) if x != 1)) if output_mask[2] else None
-  return (wrap(grad_input) if grad_input is not None else None,
-          wrap(grad_weight) if grad_weight is not None else None,
-          wrap(grad_bias) if grad_bias is not None else None)
+  dims, shape = tuple(x for x in range(input_t.ndim) if x != 1), (1, -1) + (1,)*(input_t.ndim-2)
+  # training differentiates the batch stats it was given, eval treats the running stats as constants
+  if train: mean, invstd = unwrap(save_mean), unwrap(save_invstd)
+  else: mean, invstd = unwrap(running_mean), unwrap(running_var).add(eps).rsqrt()
+  xhat = (input_t - mean.reshape(shape)) * invstd.reshape(shape)
+  grad_bias, grad_weight = grad_out_t.sum(axis=dims), (grad_out_t * xhat).sum(axis=dims)
+  grad_input = grad_out_t if not train else \
+    grad_out_t - (grad_bias.reshape(shape) + xhat * grad_weight.reshape(shape)) / (input_t.numel() // input_t.shape[1])
+  grad_input = grad_input * invstd.reshape(shape) * (unwrap(weight).reshape(shape) if weight is not None else 1)
+  return (wrap(grad_input) if output_mask[0] else None, wrap(grad_weight) if output_mask[1] else None,
+          wrap(grad_bias) if output_mask[2] else None)

 # _pad_circular is not CompositeImplicitAutograd (unlike reflect/replicate pad)
 # we need torch.autograd.Function with explicit AutogradPrivateUse1 registration
@@ -71,6 +71,54 @@ class TestTorchBackend(unittest.TestCase):
    a = a.as_strided((1,1,5,5), (50,50,7,1), storage_offset=21)
    np.testing.assert_equal(a.cpu().numpy().sum(-1), [[[115,150,185,220,255]]])

+  def test_storage_offset_of_computed_tensor(self):
+    # a computed result owns its storage, so a slice anywhere in its history must not shift the offset
+    a = torch.arange(8., device=device)
+    self.assertEqual((a[3:]+1).storage_offset(), 0)
+
+  def test_storage_offset_through_aliases(self):
+    a = torch.arange(8., device=device)[3:]
+    self.assertEqual(a.detach().storage_offset(), 3)
+    self.assertEqual(a.view(torch.int32).storage_offset(), 3)
+    torch.add(torch.ones(5, device=device), torch.ones(5, device=device), out=a)
+    self.assertEqual(a.detach().storage_offset(), 3)
+
+  def test_out_refreshes_views_of_base(self):
+    a = torch.zeros(4, device=device)
+    v = a[2:]
+    torch.add(torch.ones(4, device=device), torch.ones(4, device=device), out=a)
+    np.testing.assert_equal(v.cpu().numpy(), [2., 2.])
+
+  @unittest.expectedFailure  # TODO: storage offset assumes a contiguous source, use UOp.contiguous_view_offset
+  def test_storage_offset_non_contiguous_source(self):
+    a = torch.arange(12., device=device).reshape(3,4)
+    self.assertEqual(a.permute(1,0)[1:].storage_offset(), 1)
+    self.assertEqual(a.flatten()[3:].flip(0).storage_offset(), 0)
+
+  def test_as_strided_explicit_zero_offset(self):
+    # storage_offset=0 is a real offset, not "unspecified": it must not fall back to the input's own offset
+    a = torch.arange(6., device=device)
+    np.testing.assert_equal(a[3:].as_strided((2,), (1,), 0).cpu().numpy(), [0,1])
+    np.testing.assert_equal(a[3:].as_strided((2,), (1,)).cpu().numpy(), [3,4])
+
+  def test_empty_strided_default_dtype(self):
+    self.assertEqual(torch.empty_strided((2,3), (1,2), device=device).dtype, torch.get_default_dtype())
+
+  @unittest.expectedFailure  # TODO: empty_strided ignores the requested strides, the backend treats everything as contiguous
+  def test_empty_strided_honors_strides(self):
+    self.assertEqual(tuple(torch.empty_strided((2,3), (1,2), device=device).stride()), (1,2))
+
+  @unittest.expectedFailure  # TODO: torch refuses an out= that overlaps an input, we compute silently
+  def test_out_overlapping_input_is_rejected(self):
+    x = torch.arange(6., device=device)
+    with self.assertRaises(RuntimeError): torch.add(x[:-1], 10, out=x[1:])
+
+  def test_out_disjoint_input_is_allowed(self):
+    # torch permits an out= that shares a base with an input as long as they do not overlap
+    x, xc = torch.arange(6., device=device), torch.arange(6.)
+    torch.add(x[:3], 10, out=x[3:]); torch.add(xc[:3], 10, out=xc[3:])
+    np.testing.assert_equal(x.cpu().numpy(), xc.numpy())
+
  def test_plus_inplace(self):
    a = torch.ones(4, device=device)
    b = torch.ones(4, device=device)
@@ -124,6 +172,15 @@ class TestTorchBackend(unittest.TestCase):
    expected = np.array([[1.5, 5.2, 9.0], [13.2, 17.1, 18.4]], dtype=np.float32)
    np.testing.assert_equal(y3.cpu().numpy(), expected)

+  def test_argmax_argmin(self):
+    a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
+    c = a.cpu()
+    for got, want in [(a.argmax(), c.argmax()), (a.argmin(0), c.argmin(0)), (a.argmax(1, keepdim=True), c.argmax(1, keepdim=True)),
+                      (torch.min(a, 1).indices, torch.min(c, 1).indices), (torch.max(a, 1).indices, torch.max(c, 1).indices),
+                      (torch.min(a, 1).values, torch.min(c, 1).values), (torch.min(a, 1, keepdim=True).indices, torch.min(c, 1, keepdim=True).indices)]:
+      self.assertEqual(got.dtype, want.dtype)  # torch's arg reduces are int64, tinygrad's are int32
+      np.testing.assert_equal(got.cpu().numpy(), want.numpy())
+
  def test_isfinite(self):
    a = torch.ones(4, device=device)
    np.testing.assert_equal(torch.isfinite(a).cpu().numpy(), [True, True, True, True])
@@ -316,6 +373,37 @@ class TestTorchBackend(unittest.TestCase):
    assert b.shape == (4, 2, 3)
    np.testing.assert_equal(b.cpu().numpy(), a.cpu().numpy().transpose(2, 0, 1))

+  def test_batchnorm_backward_realized_stats(self):
+    # the saved stats are a function of input in training, so grad_input must flow through them even when handed in realized.
+    # the backward eps is unused in training: torch differentiates the save_invstd it was given
+    x0, g0 = torch.randn(8, 4, 3, 3), torch.randn(8, 4, 3, 3)
+    def run(dev, bwd_eps):
+      x, go = x0.to(dev), g0.to(dev)
+      w, b = torch.linspace(0.5, 2.0, 4).to(dev), torch.zeros(4, device=dev)
+      rm, rv = torch.zeros(4, device=dev), torch.ones(4, device=dev)
+      out, sm, si = torch.ops.aten.native_batch_norm(x, w, b, rm, rv, True, 0.1, 1e-5)
+      grads = torch.ops.aten.native_batch_norm_backward(go, x, w, rm, rv, sm.clone().detach(), si.clone().detach(),
+                                                        True, bwd_eps, [True,True,True])
+      return [t.cpu().numpy() for t in grads]
+    for bwd_eps in [1e-5, 0.3]:
+      for got, want in zip(run(device, bwd_eps), run("cpu", bwd_eps)): np.testing.assert_allclose(got, want, atol=1e-4, rtol=1e-3)
+
+  def test_groupnorm_backward(self):
+    def run(dev):
+      x = torch.arange(24., device=dev).reshape(2, 4, 3).requires_grad_()
+      w = torch.linspace(0.5, 2.0, 4).to(dev).requires_grad_()
+      torch.nn.functional.group_norm(x, 2, w, torch.zeros(4, device=dev)).square().sum().backward()
+      return x.grad.cpu().numpy(), w.grad.cpu().numpy()
+    for got, want in zip(run(device), run("cpu")): np.testing.assert_allclose(got, want, atol=1e-4, rtol=1e-3)
+
+  def test_mse_smooth_l1_loss_backward(self):
+    def run(dev, loss):
+      x = torch.arange(4., device=dev).requires_grad_()
+      loss(x, torch.ones(4, device=dev)).backward()
+      return x.grad.cpu().numpy()
+    for loss in [torch.nn.functional.mse_loss, torch.nn.functional.smooth_l1_loss]:
+      np.testing.assert_allclose(run(device, loss), run("cpu", loss), atol=1e-6)
+
  def test_batchnorm_unsqueeze(self):
    bn = torch.nn.BatchNorm2d(4).to(device)
    x = torch.randn(8, 4, 3, 3, device=device)
@@ -459,6 +547,15 @@ class TestTorchBackend(unittest.TestCase):
    cpu_res = torch.arange(20, dtype=torch.float32)[::2][1:4].numpy()
    np.testing.assert_equal(torch_res, cpu_res)

+  def test_select_out_of_range_dim(self):
+    a = torch.arange(12, dtype=torch.int32, device=device).reshape(3, 4)
+    with self.assertRaises(IndexError): a.select(5, 0)
+
+  def test_select_collapses_the_only_dim(self):
+    a = torch.arange(3, dtype=torch.int32, device=device)
+    self.assertEqual(a.select(0, 1).shape, ())
+    np.testing.assert_equal(a.select(0, 1).cpu().numpy(), 1)
+
  def test_slice_negative_dim(self):
    a = torch.arange(13, dtype=torch.int32, device=device).repeat(8, 1)
    torch_chunks = a.chunk(3, -1)
@@ -739,9 +836,103 @@ class TestTorchBackend(unittest.TestCase):
    np.testing.assert_allclose(w_tiny.grad.cpu().numpy(), w_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)
    np.testing.assert_allclose(b_tiny.grad.cpu().numpy(), b_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)

+  def test_write_through_detach_of_unrealized(self):
+    a = torch.empty(4, device=device)
+    a.detach().fill_(3)
+    np.testing.assert_equal(a.cpu().numpy(), [3, 3, 3, 3])
+
+  def test_square_transpose_inplace(self):
+    # a same-shape transpose is not a reshape: writing the transposed values straight back would scramble the base
+    a = torch.tensor([[0., 1., 2.], [3., 4., 5.], [6., 7., 8.]], device=device)
+    a.transpose(0, 1).add_(100)
+    np.testing.assert_equal(a.cpu().numpy(), [[100., 101., 102.], [103., 104., 105.], [106., 107., 108.]])
+
+  def test_interpolate(self):
+    a = torch.arange(4, dtype=torch.float32, device=device).reshape(1, 1, 2, 2)
+    nearest = torch.nn.functional.interpolate(a, scale_factor=2.0)
+    np.testing.assert_equal(nearest.cpu().numpy()[0, 0], [[0, 0, 1, 1], [0, 0, 1, 1], [2, 2, 3, 3], [2, 2, 3, 3]])
+    linear = torch.nn.functional.interpolate(a, size=(4, 4), mode="bilinear", align_corners=False)
+    ref = torch.nn.functional.interpolate(a.cpu(), size=(4, 4), mode="bilinear", align_corners=False)
+    np.testing.assert_allclose(linear.cpu().numpy(), ref.numpy(), rtol=1e-5)
+
+  def test_interpolate_bicubic_area(self):
+    a = torch.arange(32, dtype=torch.float32, device=device).reshape(1, 2, 4, 4)
+    for mode, scale in [("bicubic", 2.0), ("area", 0.5)]:
+      ref = torch.nn.functional.interpolate(a.cpu(), scale_factor=scale, mode=mode)
+      np.testing.assert_allclose(torch.nn.functional.interpolate(a, scale_factor=scale, mode=mode).cpu().numpy(), ref.numpy(), atol=1e-4)
+
+  @unittest.expectedFailure
+  def test_interpolate_bicubic_backward(self):
+    # the forward comes from a decomposition, but aten::upsample_bicubic2d_backward has none (nor does
+    # aten::_adaptive_avg_pool2d_backward, for area), so training through these modes needs a real kernel
+    x = torch.arange(32., dtype=torch.float32, device=device).reshape(1, 2, 4, 4).requires_grad_()
+    torch.nn.functional.interpolate(x, scale_factor=2.0, mode="bicubic").sum().backward()
+
+  @unittest.expectedFailure
+  def test_interpolate_inexact_scale(self):
+    # torch forwards the raw scale_factor, Tensor.interpolate recomputes it from output_size, and they disagree here
+    a = torch.arange(6, dtype=torch.float32, device=device).reshape(1, 1, 2, 3)
+    tiny = torch.nn.functional.interpolate(a, scale_factor=2.5, mode="bilinear")
+    ref = torch.nn.functional.interpolate(a.cpu(), scale_factor=2.5, mode="bilinear")
+    np.testing.assert_allclose(tiny.cpu().numpy(), ref.numpy(), rtol=1e-5)
+
+  def test_logical_or_xor(self):
+    a = torch.tensor([True, True, False, False], device=device)
+    b = torch.tensor([True, False, True, False], device=device)
+    np.testing.assert_equal(torch.logical_or(a, b).cpu().numpy(), [True, True, True, False])
+    np.testing.assert_equal(torch.logical_xor(a, b).cpu().numpy(), [False, True, True, False])
+    # bool-valued whatever the input dtype, so this is not | and ^
+    i, j = torch.tensor([2, 0, 5, 0], device=device), torch.tensor([0, 0, 1, 1], device=device)
+    np.testing.assert_equal(torch.logical_or(i, j).cpu().numpy(), [True, False, True, True])
+    np.testing.assert_equal(torch.logical_xor(i, j).cpu().numpy(), [True, False, False, True])
+
+  def test_slice_scatter(self):
+    # the scatters are functional: they return a new tensor and must leave the one they were given alone
+    a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
+    out = torch.slice_scatter(a, torch.ones(1, 4, device=device), 0, 0, 1)
+    np.testing.assert_equal(out.cpu().numpy(), [[1, 1, 1, 1], [4, 5, 6, 7], [8, 9, 10, 11]])
+    np.testing.assert_equal(a.cpu().numpy(), np.arange(12, dtype=np.float32).reshape(3, 4))
+
+  def test_slice_scatter_casts_src(self):
+    a = torch.zeros(3, 4, device=device)
+    out = torch.slice_scatter(a, torch.ones(1, 4, dtype=torch.int32, device=device), 0, 0, 1)
+    self.assertEqual(out.dtype, torch.float32)
+    np.testing.assert_equal(out.cpu().numpy()[0], np.ones(4, dtype=np.float32))
+
+  def test_select_scatter(self):
+    a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
+    out = torch.select_scatter(a, torch.ones(4, device=device), 0, 1)
+    np.testing.assert_equal(out.cpu().numpy(), [[0, 1, 2, 3], [1, 1, 1, 1], [8, 9, 10, 11]])
+
+  def test_diagonal_scatter(self):
+    a = torch.zeros(3, 3, device=device)
+    out = torch.diagonal_scatter(a, torch.arange(3, dtype=torch.float32, device=device))
+    np.testing.assert_equal(out.cpu().numpy(), np.diag([0., 1., 2.]))
+    np.testing.assert_equal(a.cpu().numpy(), np.zeros((3, 3), dtype=np.float32))
+
+  def test_copy_functional(self):
+    # without an impl this segfaults rather than fails: a regression here takes the whole run down
+    a = torch.arange(4, dtype=torch.float32, device=device)
+    out = torch.ops.aten.copy(a, torch.zeros(4, device=device))
+    np.testing.assert_equal(out.cpu().numpy(), [0., 0., 0., 0.])
+    np.testing.assert_equal(a.cpu().numpy(), [0., 1., 2., 3.])

 from tinygrad import Tensor
 class TestBackendHelpers(unittest.TestCase):
+  def test_unwrap_rejects_foreign_tensor(self):
+    # unwrap casts to the tiny impl, so a tensor from another backend must be refused rather than reinterpreted
+    with self.assertRaises(RuntimeError): extra.torch_backend.backend.unwrap(torch.ones(4))
+
+  def test_update_metadata_rejects_foreign_tensor(self):
+    # resizing a tensor we don't own would expose memory past its allocation
+    t = torch.ones(4)
+    with self.assertRaises(RuntimeError): extra.torch_backend.backend.mod.update_metadata(t, [8], [1], 0)
+    self.assertEqual(t.shape, (4,))
+
+  def test_unwrap_parameter_and_detached(self):
+    # nn.Parameter and detach rebuild the base OpaqueTensorImpl, which unwrap still has to accept
+    extra.torch_backend.backend.unwrap(torch.nn.Parameter(torch.ones(4, device="tiny")))
+    extra.torch_backend.backend.unwrap(torch.ones(4, device="tiny").detach())

  def test_calculate_storage_offset_no_shrink(self):
    t = Tensor.ones(3, 4)
@@ -124,16 +124,21 @@ at::Tensor wrap_tensor(py::object &py_obj, c10::ScalarType dtype, c10::DeviceInd
    sizes, strides, storage_offset);
 }

+// shallow_copy_and_detach (nn.Parameter, aten.detach) rebuilds the base OpaqueTensorImpl, so that is the type every tiny tensor has
+at::OpaqueTensorImpl<std::shared_ptr<c10::SafePyObject>> *tiny_impl(const at::Tensor &tensor) {
+  auto* impl = dynamic_cast<at::OpaqueTensorImpl<std::shared_ptr<c10::SafePyObject>>*>(tensor.unsafeGetTensorImpl());
+  TORCH_CHECK(impl != nullptr, "expected a tiny tensor, got a ", tensor.device().str(), " one. move it with .to(\"tiny\") first");
+  return impl;
+}
+
 py::object unwrap_tensor(const at::Tensor &tensor) {
-  auto* impl = tensor.unsafeGetTensorImpl();
-  auto* opaque_impl = static_cast<at::TinyOpaqueTensorImpl<std::shared_ptr<c10::SafePyObject>>*>(impl);
-  std::shared_ptr<c10::SafePyObject> tiny = opaque_impl->opaque_handle();
+  std::shared_ptr<c10::SafePyObject> tiny = tiny_impl(tensor)->opaque_handle();
  return py::reinterpret_borrow<py::object>(tiny->ptr(getPyInterpreter()));
 }

 void update_metadata(const at::Tensor &tensor, const std::vector<int64_t> &sizes,
                     const std::vector<int64_t> &strides, int64_t storage_offset) {
-  auto* impl = tensor.unsafeGetTensorImpl();
+  auto* impl = tiny_impl(tensor);
  impl->set_allow_tensor_metadata_change(true);
  impl->set_sizes_and_strides(sizes, strides, storage_offset);
 }
@@ -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" }]

@@ -111,6 +111,10 @@ docs = [
  "numpy",
 ]
 mesa = ["tinymesa==25.2.7.2"]
+autogen = [
+  "pyyaml",
+  "mako",
+]


 [tool.mutmut]
@@ -29,6 +29,7 @@
 \definecolor{axbrred}{HTML}{E53935}      % GROUP_REDUCE
 \definecolor{axyellow}{HTML}{F9A825}     % UPCAST
 \definecolor{axmagenta}{HTML}{7B1FA2}    % UNROLL
+\definecolor{axgreen}{HTML}{2E7D32}      % DEVICE

 \title{tinygrad: a single dialect from Tensor programs to Command Buffers}
 \author{tinygrad, Corp. \\ \texttt{[email protected]}}
@@ -79,6 +80,7 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op
 \op{Index} & $(T, i_0, i_1, \ldots)$ & --- & Index from left. $()$-shaped $i$ removes dim; $(k,)$-shaped makes it $k$. \\
 \op{Stack} & $(T_0, T_1, \ldots)$ & --- & Join along a newly created leading axis. All shapes must match. \\
 \op{Bitcast} & $(T,)$ & dtype & Reinterpret storage as target dtype; preserve total bytes. \\
+\op{Unshard} & $(T, R_0, R_1, \ldots)$ & axes $(a_0, a_1, \ldots)$ & Concatenate shards of \op{Range} $R_k$ along axis $a_k$; $R_k$ is outer. \\
 \bottomrule
 \end{tabular}

@@ -258,6 +260,7 @@ Every UOp has a \textbf{dtype}, \textbf{shape}, \textbf{device}, \textbf{addrspa
 \op{Const} & from arg & $()$ & \textsc{null} & $[v, v]$ \\
 \op{Param} & from arg & from $\mathrm{src}[0]$ & from arg & from src or dtype range \\[3pt]
 Movement ops & $\mathrm{src}[0].\mathrm{dtype}$ & (see op) & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
+\op{Unshard} & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0]$, each $a_k \times n_k$ & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
 \op{Reduce} & $\mathrm{src}[0].\mathrm{dtype}$ & remove first $n$ axes & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\[3pt]
 \op{Cast} & from arg & $\mathrm{src}[0].\mathrm{shape}$ & $\mathrm{src}[0].\mathrm{device}$ & clamped to dtype \\
 \op{Bitcast} & from arg & $\mathrm{src}[0].\mathrm{shape}$ & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\
@@ -283,9 +286,9 @@ $[a,A]$, $[b,B]$, $[c,C]$ denote min\_max of $\mathrm{src}[0]$, $\mathrm{src}[1]
 Default \emph{dtype range}: $[\mathrm{dtype\_min},\, \mathrm{dtype\_max}]$.

 \medskip
-\textbf{axis} tracks the multi-device sharding dimension. \op{Buffer} with $n$-tuple device: axis $= 0$ (device dim).
-\op{Reshape} remaps axis to preserve the shard boundary. \op{Permute} follows the permutation. \op{Expand} shifts axis right by $|\mathbf{n}|$.
-\op{Reduce} on the shard axis $\to$ \textsc{null} (shard axis is among the first $n$ axes). \op{Replicated} on the shard axis $\to$ \textsc{null}. \op{Copy} $\to$ \textsc{null}. ALU ops inherit from sources. Default: \textsc{null}.
+\textbf{sharding} tracks multi-device sharding as a set of (axis, \op{Range}) pairs. \op{Unshard} defines it: arg is the tuple of sharded axes, one \op{Range} in src per axis (positional: the $k$-th \op{Range} shards the $k$-th axis). \op{Buffer} with $n$-tuple device: sharded on axis $0$ (device dim). The single-axis convenience \textbf{axis} is \textsc{null} unless exactly one axis is sharded.
+\op{Reshape} remaps each sharded axis to preserve its shard boundary. \op{Permute} follows the permutation. \op{Expand} shifts all sharded axes right by $|\mathbf{n}|$.
+\op{Reduce} on a sharded axis drops it. \op{Replicated} on the shard axis $\to$ \textsc{null}. \op{Copy} $\to$ \textsc{null}. ALU ops inherit from sources. Default: \textsc{null}.

 %% ============================================================
 \subsection*{Kernel Optimizations (OptOps) \normalfont\small--- schedule-level transforms on kernel ranges}
@@ -297,6 +300,7 @@ Each kernel's iteration space is a set of \op{Range} axes. Every range has an \t
 \toprule
 \textbf{AxisType} & \textbf{Letter} & \textbf{Split from} & \textbf{Direction} & \textbf{Semantics} \\
 \midrule
+{\color{axgreen}\texttt{DEVICE}} & \texttt{d} & --- & --- & Multi-device sharding dimension. \\
 {\color{axblue}\texttt{GLOBAL}} & \texttt{g} & --- & --- & GPU global workgroup dimension. \\
 {\color{axcyan}\texttt{LOCAL}} & \texttt{l} & g, L & inner & Workgroup local dimension (shared memory). \\
 {\color{axbrcyan}\texttt{WARP}} & \texttt{w} & \multicolumn{2}{l}{(created by \op{TC})} & Warp-level lanes for tensor cores. \\
@@ -378,6 +382,8 @@ def scatter_add(T, idx, val):
 Let $D = (d_0, \ldots, d_{n-1})$ be an $n$-tuple device.
 \op{Copy} to an $n$-tuple device reshards with axis $= 0$. \op{Copy} never changes shape.

+\textbf{Sharding} splits a tensor along an axis across $n$ devices. It opens a \op{Range} of type \texttt{DEVICE} (a symbolic per-device index $d$), shrinks each device's view to its piece, then closes the range with \op{Unshard}$(T, R, a)$. The result is a logical tensor whose shape along axis $a$ is the full size; each device holds $1/n$ of it. \op{Unshard} is the inverse of sharding --- it marks the boundary between per-device computation and the logical multi-device tensor. The range need not be \texttt{DEVICE}; e.g.\ a \texttt{WARP} range closes the same way, concatenating per-lane shards along $a$ with the range as the outer factor. A tensor may be sharded along several axes at once: \op{Unshard}$(T, R_0, R_1, \ldots;\; a_0, a_1, \ldots)$ carries one \op{Range} per sharded axis, and every movement op maps each sharded axis independently.
+
 \begin{lstlisting}
 # T has shape (s,) on a single device.

@@ -169,11 +169,10 @@ def run_program_emu(instructions: list, n_lanes: int = 1) -> WaveState:
  return parse_output(bytes(out_buf), n_lanes)

 def run_program_hw(instructions: list, n_lanes: int = 1) -> WaveState:
-  """Run instructions on real AMD hardware via HIPCompiler and AMDProgram."""
-  from tinygrad.device import Device
-  from tinygrad.runtime.ops_amd import AMDProgram
+  """Run instructions on real AMD hardware via HIPCompiler and the AMD runtime."""
+  from tinygrad.device import Device, TinyELF
  from tinygrad.runtime.support.compiler_amd import HIPCompiler
-  from tinygrad.helpers import flat_mv
+  from tinygrad.helpers import Target, flat_mv

  dev = Device["AMD"]
  compiler = HIPCompiler(dev.arch)  # type: ignore[attr-defined]
@@ -223,7 +222,7 @@ amdhsa.kernels:
 """

  lib = compiler.compile(asm_src)
-  prg = AMDProgram(dev, "test", lib)  # type: ignore[arg-type]
+  prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))

  buf_sz = _out_bytes(n_lanes)
  out_gpu = dev.allocator.alloc(buf_sz)
@@ -5,7 +5,7 @@ gfx950 hardware when USE_HW=1.
 """
 import ctypes, struct, unittest
 import tinygrad.runtime.autogen.amd.cdna.ins as cdna
-from tinygrad.helpers import flat_mv
+from tinygrad.helpers import Target, flat_mv
 from tinygrad.renderer.amd.dsl import NULL
 from test.amd.hw.helpers import USE_HW, assemble
 from test.mockgpu.amd.emu import run_asm
@@ -42,8 +42,7 @@ def _run_emu(instructions: list, out_reg: int = 2) -> int:
  return out_buf[0]

 def _run_hw(instructions: list, out_reg: int = 2) -> int:
-  from tinygrad.device import Device
-  from tinygrad.runtime.ops_amd import AMDProgram
+  from tinygrad.device import Device, TinyELF
  from tinygrad.runtime.support.compiler_amd import HIPCompiler

  dev = Device["AMD"]
@@ -86,7 +85,7 @@ amdhsa.kernels:
 ...
 .end_amdgpu_metadata
 """
-  prg = AMDProgram(dev, "test", HIPCompiler(dev.arch).compile(asm_src))
+  prg = dev.runtime(TinyELF(HIPCompiler(dev.arch).compile(asm_src), "test", Target("AMD", arch=dev.arch), ()))
  prg(global_size=(1, 1, 1), local_size=(LANES, 1, 1), wait=True)
  out = bytearray(LANES * 4)
  dev.allocator._copyout(flat_mv(memoryview(out)), out_gpu)
@@ -6,7 +6,7 @@ when USE_HW=1.
 """
 import ctypes, unittest
 from tinygrad.runtime.autogen.amd.rdna3.ins import *
-from tinygrad.helpers import flat_mv
+from tinygrad.helpers import Target, flat_mv
 from test.amd.hw.helpers import USE_HW, assemble
 from test.mockgpu.amd.emu import run_asm

@@ -37,8 +37,7 @@ def _run_wave64_emu(instructions: list, out_reg: int = 1) -> list[int]:
  return list(out_buf)

 def _run_wave64_hw(instructions: list, out_reg: int = 1) -> list[int]:
-  from tinygrad.device import Device
-  from tinygrad.runtime.ops_amd import AMDProgram
+  from tinygrad.device import Device, TinyELF
  from tinygrad.runtime.support.compiler_amd import HIPCompiler

  dev = Device["AMD"]
@@ -84,7 +83,7 @@ amdhsa.kernels:
 .end_amdgpu_metadata
 """
  lib = compiler.compile(asm_src)
-  prg = AMDProgram(dev, "test", lib)  # type: ignore[arg-type]
+  prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
  out_gpu = dev.allocator.alloc(WAVE64 * 4)
  prg(out_gpu, global_size=(1, 1, 1), local_size=(WAVE64, 1, 1), wait=True)
  out = bytearray(WAVE64 * 4)
@@ -5,7 +5,7 @@ real RDNA4 hardware when USE_HW=1.
 """
 import ctypes, unittest
 import tinygrad.runtime.autogen.amd.rdna4.ins as r4
-from tinygrad.helpers import flat_mv
+from tinygrad.helpers import Target, flat_mv
 from tinygrad.renderer.amd.dsl import NULL
 from test.amd.hw.helpers import USE_HW, assemble
 from test.mockgpu.amd.emu import run_asm
@@ -36,8 +36,7 @@ def _run_emu(instructions: list, out_reg: int = 2) -> list[int]:
  return list(out_buf)

 def _run_hw(instructions: list, out_reg: int = 2) -> list[int]:
-  from tinygrad.device import Device
-  from tinygrad.runtime.ops_amd import AMDProgram
+  from tinygrad.device import Device, TinyELF
  from tinygrad.runtime.support.compiler_amd import HIPCompiler

  dev = Device['AMD']
@@ -85,7 +84,7 @@ amdhsa.kernels:
 .end_amdgpu_metadata
 """
  lib = compiler.compile(asm_src)
-  prg = AMDProgram(dev, 'test', lib)
+  prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
  out_gpu = dev.allocator.alloc(LANES * 4)
  prg(out_gpu, global_size=(1, 1, 1), local_size=(LANES, 1, 1), wait=True)
  out = bytearray(LANES * 4)
@@ -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()
--- a/Show More
+++ b/Show More