Merge branch 'master' into simpler_image

ish
fix half
2026-06-24 16:53:31 -07:00 · 2026-06-24 16:14:15 -07:00 · 2026-06-24 15:43:05 -07:00 · 2026-06-24 15:16:00 -07:00 · 2026-06-24 14:34:44 -07:00 · 2026-06-24 14:27:34 -07:00
176 changed files with 6456 additions and 3783 deletions
@@ -10,7 +10,7 @@ inputs:
    required: false
    default: ''  # if you don't set a key, it doesn't cache
  deps:
-    description: 'Extra dependency groups (space separated)'
+    description: 'Extra dependency groups (comma separated)'
    required: false
    default: ''
  pydeps:
@@ -41,6 +41,10 @@ inputs:
    description: "Install LLVM?"
    required: false
    default: 'false'
+  mesa:
+    description: "Install mesa (true, false, cpu)"
+    required: false
+    default: 'false'
  tinydreno:
    description: "Install tinydreno"
    required: false
@@ -110,8 +114,7 @@ runs:
      shell: bash
      run: |
        uv venv .venv
-        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 .venv -e ".[${{ inputs.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
@@ -143,6 +146,11 @@ runs:
        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
        echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' | sudo tee -a /etc/apt/apt.conf.d/99keep-debs

+    - name: Add OpenCL Repo
+      if: inputs.opencl == 'true' && runner.os == 'Linux'
+      shell: bash
+      run: echo "deb [ allow-insecure=yes ] https://apt.repos.intel.com/oneapi all main" | sudo tee /etc/apt/sources.list.d/oneAPI.list
+
    - name: Add AMD Repo (Linux)
      if: inputs.amd == 'true' && runner.os == 'Linux'
      shell: bash
@@ -168,7 +176,10 @@ runs:
        pkgs=""
        # **** OpenCL ****
        if [[ "${{ inputs.opencl }}" == "true" ]]; then
-          pkgs+=" ocl-icd-opencl-dev"
+          pkgs+=" opencl-headers \
+            intel-oneapi-runtime-openmp=2023.2.1-16 intel-oneapi-runtime-compilers-common=2023.2.1-16 intel-oneapi-runtime-compilers=2023.2.1-16 \
+            intel-oneapi-runtime-dpcpp-sycl-opencl-cpu=2023.2.1-16 intel-oneapi-runtime-tbb-common=2021.10.0-49541 \
+            intel-oneapi-runtime-tbb=2021.10.0-49541 intel-oneapi-runtime-opencl=2023.2.1-16"
        fi
        # **** AMD ****
        if [[ "${{ inputs.amd }}" == "true" ]]; then
@@ -275,18 +286,18 @@ runs:
      shell: bash
      run: brew install llvm@20

+    # **** mesa ****
+    - name: Install mesa (linux)
+      if: inputs.mesa != 'false' && runner.os == 'Linux'
+      shell: bash
+      run: sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/v1/libtinymesa${{ inputs.mesa == 'cpu' && '_cpu' || '' }}-mesa-25.2.7-linux-amd64.so -o /usr/lib/libtinymesa${{ inputs.mesa == 'cpu' && '_cpu' || '' }}.so
+    - name: Install mesa (macOS)
+      if: inputs.mesa != 'false' && runner.os == 'macOS'
+      shell: bash
+      run: brew install sirhcm/tinymesa/tinymesa${{ inputs.mesa == 'cpu' && '_cpu' || '' }}
+
    # *** 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'
-      shell: bash
-      run: |
-        sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/rusticl-v1/libRusticlOpenCL.so.1.0.0 -o /usr/lib/libRusticlOpenCL.so
-        sudo mkdir -p /etc/OpenCL/vendors
-        echo "/usr/lib/libRusticlOpenCL.so" | sudo tee /etc/OpenCL/vendors/rusticl.icd
-        echo "RUSTICL_ENABLE=llvmpipe" >> "$GITHUB_ENV"
@@ -81,285 +81,80 @@ jobs:
  #      source /tmp/tinygrad_pytest_ci/bin/activate
  #      pytest -nauto --durations=20

-  llmbenchmark:
-    name: LLM (DEV=${{ matrix.dev }})
-    runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
-    strategy:
-      fail-fast: false
-      matrix:
-        dev: ['METAL', 'AMD', 'NV']
+  testmacbenchmark:
+    name: Mac Benchmark
+    runs-on: [self-hosted, macOS]
    timeout-minutes: 60
    defaults:
      run:
        shell: bash -e -o pipefail {0}
-    env:
-      DEV: ${{ matrix.dev }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
      uses: actions/checkout@v6
-    - 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
+        mkdir -p extra/disassemblers
+        ln -s ~/tinygrad/extra/disassemblers/applegpu extra/disassemblers/applegpu
+        ln -s ~/tinygrad/weights/sd-v1-4.ckpt weights/sd-v1-4.ckpt
+        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
+        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
    - 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 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
-      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
-    - 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
-
-  cifarbenchmark:
-    name: HLB-CIFAR10 (DEV=${{ matrix.dev }})
-    runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
-    strategy:
-      fail-fast: false
-      matrix:
-        dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 60
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    env:
-      DEV: ${{ matrix.dev }}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - 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: 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 10 CIFAR training steps
-      env:
-        ASSERT_MIN_STEP_TIME: ${{ matrix.dev == 'NV' && '130' || matrix.dev == 'AMD' && '200' || '3000' }}
-      run: BENCHMARK_LOG=cifar_10steps STEPS=10 python3 examples/hlb_cifar10.py
-    - name: Run 10 CIFAR training steps w HALF
-      env:
-        ASSERT_MIN_STEP_TIME: ${{ matrix.dev == 'NV' && '120' || matrix.dev == 'AMD' && '235' || '3000' }}
-      run: BENCHMARK_LOG=cifar_10steps_half STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
-    - name: Run full CIFAR training w 1 GPU
-      # 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
-
-  mlperfbenchmark:
-    name: MLPerf (${{ matrix.dev }})
-    runs-on: [self-hosted, Linux, "${{ matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
-    strategy:
-      fail-fast: false
-      matrix:
-        dev: ['AMD', 'NV']
-    timeout-minutes: 60
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    env:
-      DEV: ${{ matrix.dev }}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - 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 extra/datasets
-        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: 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
-
-  sdbenchmark:
-    name: Stable Diffusion (DEV=${{ matrix.dev }})
-    runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
-    strategy:
-      fail-fast: false
-      matrix:
-        dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 60
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    env:
-      DEV: ${{ matrix.dev }}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - 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: 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
+      run: python3.11 test/external/process_replay/reset.py
+    - name: Print macOS version
+      run: sw_vers
    - name: Run Stable Diffusion
-      env:
-        ASSERT_MIN_STEP_TIME: ${{ matrix.dev == 'METAL' && '720' || matrix.dev == 'AMD' && '550' || '0' }}
-      run: BENCHMARK_LOG=stable_diffusion python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
+      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
+    - name: Run Stable Diffusion without fp16
+      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing
+    - name: Run Stable Diffusion v2
+      # TODO: very slow step time
+      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing
+    # process replay can't capture this, the graph is too large
    - name: Run SDXL
-      if: ${{ matrix.dev != 'NV' }}
-      env:
-        ASSERT_MIN_STEP_TIME: ${{ matrix.dev == 'METAL' && '5000' || matrix.dev == 'AMD' && '3200' || '2000' }}
-      run: BENCHMARK_LOG=stable_diffusion_xl CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing
-    - 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' }}"]
-    strategy:
-      fail-fast: false
-      matrix:
-        dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 60
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    env:
-      DEV: ${{ matrix.dev }}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - 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: 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: Test tiny
-      run: |
-        DEBUG=2 python -m pytest -rA test/test_tiny.py
-        if [[ "${{ matrix.dev }}" == "NV" ]]; then
-          DEBUG=2 DEV=CUDA python -m pytest -rA test/test_tiny.py
-        fi
-    - name: Test tensor cores
-      run: |
-        if [[ "${{ matrix.dev }}" == "METAL" ]]; then
-          python3 test/opt/test_tensor_cores.py
-          DEBUG=2 SHOULD_USE_TC=1 python3 extra/gemm/simple_matmul.py
-          DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3 extra/gemm/simple_matmul.py
-          DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3 extra/gemm/simple_matmul.py
-          M_START=6 M_STOP=10 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=6 K_STOP=24 K_STEP=1 TC_OPT=2 DEBUG=2 python3 ./extra/gemm/fuzz_matmul.py
-        elif [[ "${{ matrix.dev }}" == "NV" ]]; then
-          ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
-          DEV=NV:PTX ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
-          DEV=CUDA SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-          DEV=CUDA SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-          DEV=CUDA SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
-          DEV=CUDA SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-          DEV=NV:PTX SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-          SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-          # TODO: too slow
-          # M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
-          # DEV=NV:PTX M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
-        else
-          python3 test/opt/test_tensor_cores.py
-          # TODO: this is flaky
-          # DEV=AMD:LLVM python3 test/opt/test_tensor_cores.py
-          SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-          SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
-          # TODO: AMD compiler bug causes this to fail
-          # HSA=1 M_START=12 M_STOP=20 M_STEP=1 N_START=12 N_STOP=20 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 DEBUG=2 python3 ./extra/gemm/fuzz_matmul.py
-        fi
+      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing
    - name: Run model inference benchmark
-      # TODO: unstable on AMD
-      if: ${{ matrix.dev != 'AMD' }}
-      run: CAPTURE_PROCESS_REPLAY=0 NOCLANG=1 python3 test/external/external_model_benchmark.py
+      run: DEV=METAL NOCLANG=1 python3.11 test/external/external_model_benchmark.py
    - name: Test speed vs torch
-      # TODO: unstable on AMD
-      if: ${{ matrix.dev != 'AMD' }}
-      env:
-        HALF: ${{ matrix.dev == 'NV' && '1' || '0' }}
-      run: CAPTURE_PROCESS_REPLAY=0 BIG=2 ${{ matrix.dev == 'METAL' && 'MPS=1' || 'TORCHCUDA=1' }} python3 test/speed/external_test_speed_v_torch.py
-    - name: Test speed vs theoretical
-      # no targets for METAL
-      if: ${{ matrix.dev != 'METAL' }}
-      run: IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+      run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py
+    - name: Test tensor cores
+      run: DEV=METAL python3.11 test/opt/test_tensor_cores.py
+    - name: Run Tensor Core GEMM (float)
+      run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py
+    - name: Run Tensor Core GEMM (half)
+      run: DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3.11 extra/gemm/simple_matmul.py
+    - name: Run Tensor Core GEMM (bfloat16)
+      run: DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py
+    - name: Fuzz Padded Tensor Core GEMM
+      run: DEV=METAL M_START=6 M_STOP=10 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=6 K_STOP=24 K_STEP=1 TC_OPT=2 DEBUG=2 python3.11 ./extra/gemm/fuzz_matmul.py
+    - name: Run llama3.2
+      run: BENCHMARK_LOG=llama32_3b-f16 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 -m tinygrad.llm -m llama3.2:3b-f16 --benchmark --warmup
+    - name: Run olmoe
+      run: BENCHMARK_LOG=olmoe JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 -m tinygrad.llm -m olmoe --benchmark --warmup
    - name: Train MNIST
-      run: time TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
-    - name: Test benchmark allreduce
-      if: ${{ matrix.dev == 'NV' }}
-      run: python test/external/external_benchmark_multitensor_allreduce.py
-    - 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
+      run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py
+
+    # NOTE: this is failing in CI. it is not failing on my machine and I don't really have a way to debug it
+    # the error is "RuntimeError: Internal Error (0000000e:Internal Error)"
+    #- name: Run 10 CIFAR training steps
+    #  run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=3000 STEPS=10 python3.11 examples/hlb_cifar10.py
+    #- name: Run 10 CIFAR training steps w HALF
+    #  run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=3000 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py
+
+    #- name: Run 10 CIFAR training steps w BF16
+    #  run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py
+    # TODO: too slow
+    # - name: Run 10 CIFAR training steps w winograd
+    #   run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py
    - uses: actions/upload-artifact@v7
-      if: ${{ matrix.dev != 'AMD' }}
      with:
-        name: Speed (${{ matrix.dev }})
+        name: Speed (Mac)
        path: |
          onnx_inference_speed.csv
    - name: Run process replay tests
@@ -401,6 +196,337 @@ jobs:
    - name: UsbGPU (USB4/TB) tiny tests
      run: PYTHONPATH=. DEV=PCI+NV:NAK python3.11 test/test_tiny.py

+  testnvidiabenchmark:
+    name: tinybox green Benchmark
+    runs-on: [self-hosted, Linux, tinyboxgreen]
+    timeout-minutes: 60
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Print nvidia-smi
+      run: nvidia-smi
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
+        mkdir -p extra/datasets
+        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: test/external/process_replay/reset.py
+    - name: Run model inference benchmark
+      run: DEV=NV CAPTURE_PROCESS_REPLAY=0 NOCLANG=1 python3 test/external/external_model_benchmark.py
+    - name: Test speed vs torch
+      run: DEV=NV CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py
+    - name: Test speed vs theoretical
+      run: DEV=NV IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+    - name: Test benchmark allreduce
+      run: DEV=NV python test/external/external_benchmark_multitensor_allreduce.py
+    - name: Test tensor cores
+      run: |
+        DEV=NV ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
+        DEV=NV:PTX ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
+    - name: Run Tensor Core GEMM (CUDA)
+      run: |
+        DEV=CUDA SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+        DEV=CUDA SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+        DEV=CUDA SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
+        DEV=CUDA SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+    - name: Run Tensor Core GEMM (PTX)
+      run: DEV=NV:PTX SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+    - name: Run Tensor Core GEMM (NV)
+      run: DEV=NV SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+    - name: Test DEV=NV
+      run: DEBUG=2 DEV=NV python -m pytest -rA test/test_tiny.py
+    - name: Test DEV=CUDA
+      run: DEBUG=2 DEV=CUDA python -m pytest -rA test/test_tiny.py
+    - name: Run Stable Diffusion
+      run: BENCHMARK_LOG=stable_diffusion DEV=NV python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
+    # TODO: too slow
+    # - name: Run SDXL
+    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 DEV=NV CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing
+    - name: Run llama3.2
+      run: DEV=NV 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
+      run: DEV=NV BENCHMARK_LOG=qwen35_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 -m tinygrad.llm -m qwen3.5:35b-a3b --benchmark --warmup
+    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
+      run: BENCHMARK_LOG=llama3_beam_4gpu DEV=NV JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
+    # - name: Run LLaMA-3 8B on 6 GPUs
+    #   run: DEV=NV CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
+    # - name: Run LLaMA-2 70B
+    #   run: DEV=NV CAPTURE_PROCESS_REPLAY=0 MAX_CONTEXT=256 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
+    - uses: actions/upload-artifact@v7
+      with:
+        name: Speed (NVIDIA)
+        path: |
+          onnx_inference_speed.csv
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
+  testmorenvidiabenchmark:
+    name: tinybox green Training Benchmark
+    runs-on: [self-hosted, Linux, tinyboxgreen]
+    timeout-minutes: 60
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
+        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
+        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
+        ln -s /raid/weights/mixtral-8x7b-32kseqlen weights/mixtral-8x7b-32kseqlen
+        ln -s /raid/weights/LLaMA-2 weights/LLaMA-2
+        mkdir -p extra/datasets
+        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: test/external/process_replay/reset.py
+    # TODO: too slow
+    # - name: Fuzz Padded Tensor Core GEMM (NV)
+    #   run: DEV=NV M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
+    # TODO: too slow
+    # - name: Fuzz Padded Tensor Core GEMM (PTX)
+    #   run: DEV=NV:PTX M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
+    - name: HEVC Decode Benchmark
+      run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 DEV=NV PYTHONPATH=. python3 extra/hevc/decode.py
+    - name: Train MNIST
+      run: time PYTHONPATH=. DEV=NV TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
+    - name: Run 10 CIFAR training steps
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=130 DEV=NV STEPS=10 python3 examples/hlb_cifar10.py
+    - name: Run 10 CIFAR training steps w HALF
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=120 DEV=NV STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
+    - name: Run 10 CIFAR training steps w BF16
+      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=120 DEV=NV STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
+    # - name: Run 10 CIFAR training steps w winograd
+    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 DEV=NV WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
+    - name: Run full CIFAR training w 1 GPU
+      run: time BENCHMARK_LOG=cifar DEV=NV DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    - name: Run full CIFAR training steps w 6 GPUS
+      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 DEV=NV 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 DEV=NV MODEL=resnet python3 examples/mlperf/model_eval.py
+    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+      run: BENCHMARK_LOG=resnet_10steps DEV=NV 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 DEV=NV 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 DEV=NV 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
+
+  testamdbenchmark:
+    name: tinybox red Benchmark
+    runs-on: [self-hosted, Linux, tinybox]
+    timeout-minutes: 60
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setcap to python
+      run: ./extra/amdpci/setup_python_cap.sh
+    - name: Remove amd modules
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
+    - name: Kill stale pids
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
+    #- name: Insert amdgpu
+    #  run: sudo modprobe amdgpu
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
+        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
+        ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
+        mkdir -p extra/datasets
+        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: test/external/process_replay/reset.py
+    #- name: setup perflevel
+    #  run: |
+    #    examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/setup.sh
+    #    rocm-smi
+    #- name: Show off tinybox
+    #  run: /opt/rocm/bin/rocm-bandwidth-test
+    # TODO: unstable on AMD
+    #- name: Run model inference benchmark
+    #  run: LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 NOCLANG=1 python3 test/external/external_model_benchmark.py
+    # TODO: unstable on AMD
+    #- name: Test speed vs torch
+    #  run: |
+    #    python3 -c "import torch; print(torch.__version__)"
+    #    LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py
+    - name: Test speed vs theoretical
+      run: DEV=AMD IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+    - name: Test tensor cores (no LLVM)
+      run: DEV=AMD python3 test/opt/test_tensor_cores.py
+    # TODO: this is flaky
+    # - name: Test tensor cores AMD:LLVM
+    #   run: DEV=AMD:LLVM python3 test/opt/test_tensor_cores.py
+    - name: Run Tensor Core GEMM (AMD)
+      run: |
+        DEV=AMD SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+        DEV=AMD SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
+    - name: Test DEV=AMD
+      run: DEBUG=2 DEV=AMD python -m pytest -rA test/test_tiny.py
+    #- name: Test HIP=1
+    #  run: DEBUG=2 HIP=1 python -m pytest -rA test/test_tiny.py
+    # TODO: AMD compiler bug causes this to fail
+    #- name: Fuzz Padded Tensor Core GEMM
+    #  run: HSA=1 M_START=12 M_STOP=20 M_STEP=1 N_START=12 N_STOP=20 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 DEBUG=2 python3 ./extra/gemm/fuzz_matmul.py
+    #- name: Remove amdgpu
+    #  run: sleep 10 && sudo rmmod amdgpu # sleep a bit to let the driver unload the prev pid.
+    - name: Test AM cold start time
+      run: time DEV=AMD AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
+    - name: Test AM warm start time
+      run: time DEV=AMD python3 test/test_tiny.py TestTiny.test_plus
+    - name: Run Stable Diffusion
+      run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 DEV=AMD python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
+    - name: Run SDXL
+      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 DEV=AMD python3 examples/sdxl.py --seed 0 --noshow --timing
+    - name: Run llama3.2
+      run: DEV=AMD 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
+      run: DEV=AMD BENCHMARK_LOG=qwen35_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 -m tinygrad.llm -m qwen3.5:35b-a3b --benchmark --warmup
+    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
+      run: BENCHMARK_LOG=llama3_beam_4gpu DEV=AMD JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
+    # - name: Run LLaMA-3 8B on 6 GPUs
+    #   run: DEV=AMD CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
+    #- name: Restore amdgpu
+    #  run: sudo modprobe amdgpu
+    # - name: Run LLaMA-2 70B
+    #   run: DEV=AMD CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
+  testmoreamdbenchmark:
+    name: tinybox red Training Benchmark
+    runs-on: [self-hosted, Linux, tinybox]
+    timeout-minutes: 60
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setcap to python
+      run: ./extra/amdpci/setup_python_cap.sh
+    - name: Remove amd modules
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
+    - name: Kill stale pids
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
+        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
+        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
+        ln -s /raid/weights/mixtral-8x7b-32kseqlen weights/mixtral-8x7b-32kseqlen
+        ln -s /raid/weights/LLaMA-2 weights/LLaMA-2
+        mkdir -p extra/datasets
+        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: test/external/process_replay/reset.py
+    - name: Test GPU crash recovery
+      run: DEV=AMD python3 -m pytest -rA test/external/external_test_gpu_crash.py
+    - name: Train MNIST
+      run: time PYTHONPATH=. DEV=AMD TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
+    - name: Run 10 CIFAR training steps
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=200 DEV=AMD STEPS=10 python3 examples/hlb_cifar10.py
+    - name: Run 10 CIFAR training steps w HALF
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=230 DEV=AMD STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
+    # - name: Run 10 CIFAR training steps w BF16
+    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 DEV=AMD STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
+    # TODO: too slow
+    # - name: Run 10 CIFAR training steps w winograd
+    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 DEV=AMD WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
+    - name: Run full CIFAR training w 1 GPU
+      run: time BENCHMARK_LOG=cifar DEV=AMD DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    - name: Run full CIFAR training steps w 6 GPUS
+      run: time BENCHMARK_LOG=cifar_6gpu DEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    # TODO: broken on some of the machines
+    #- name: Test full tinyfs load
+    #  run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
+  testmlperfamdbenchmark:
+    name: tinybox red MLPerf Benchmark
+    runs-on: [self-hosted, Linux, tinybox]
+    timeout-minutes: 60
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setcap to python
+      run: ./extra/amdpci/setup_python_cap.sh
+    - name: Remove amd modules
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
+    - name: Kill stale pids
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
+        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
+        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
+        ln -s /raid/weights/mixtral-8x7b-32kseqlen weights/mixtral-8x7b-32kseqlen
+        ln -s /raid/weights/LLaMA-2 weights/LLaMA-2
+        mkdir -p extra/datasets
+        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: test/external/process_replay/reset.py
+    - name: Run MLPerf resnet eval
+      run: time BENCHMARK_LOG=resnet_eval DEV=AMD MODEL=resnet python3 examples/mlperf/model_eval.py
+    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+      run: BENCHMARK_LOG=resnet_10steps DEV=AMD 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 DEV=AMD 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 DEV=AMD 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
+
  testcommalatest:
    name: comma Benchmark (0.11.0)
    runs-on: [self-hosted, Linux, comma]
@@ -424,7 +550,7 @@ jobs:
    - name: openpilot compile3 0.11.0 driving_vision (from pickle)
      run: BENCHMARK_LOG=openpilot_0_11_0_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM taskset -c 4-7 python3 examples/openpilot/compile3.py
    - name: IR3 openpilot compile3 0.11.0 driving_vision
-      run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=18 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
+      run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.11.0 driving_policy
      run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3.2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
    - name: openpilot compile3 0.11.0 dmonitoring
@@ -522,28 +648,31 @@ jobs:
    - 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

-  driverbenchmarks:
-    name: PCI Driver Benchmark (DEV=${{ matrix.dev }})
+  testreddriverbenchmark:
+    name: AM Benchmark
    runs-on: [self-hosted, Linux, tinyboxrandom]
-    strategy:
-      fail-fast: false
-      matrix:
-        dev: ['AMD', 'NV']
    timeout-minutes: 20
    defaults:
      run:
        shell: bash -e -o pipefail {0}
-    env:
-      DEV: ${{ matrix.dev }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
      uses: actions/checkout@v6
-    - name: Setup
+    - name: Setcap to python
+      run: ./extra/amdpci/setup_python_cap.sh
+    - name: Remove amd modules
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd rmmod
+    - name: Kill stale pids
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
+    - name: Symlink models and datasets
      run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py ${{ matrix.dev == 'AMD' && 'amd' || 'nv' }} rmmod
-        ./extra/hcq/hcq_smi.py ${{ matrix.dev == 'AMD' && 'amd' || 'nv' }} kill_pids
+        mkdir -p weights
+        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
+        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
+        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
+        ln -s /raid/weights/mixtral-8x7b-32kseqlen weights/mixtral-8x7b-32kseqlen
+        ln -s /raid/weights/LLaMA-2 weights/LLaMA-2
        mkdir -p extra/datasets
        ln -s /raid/datasets/imagenet extra/datasets/imagenet
    - name: setup staging db
@@ -554,45 +683,104 @@ jobs:
    - name: reset process replay
      run: test/external/process_replay/reset.py
    - name: Test driver cold start time
-      run: time DEBUG=3 AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
+      run: time DEBUG=3 DEV=AMD AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Test driver warm start time
-      if: ${{ matrix.dev == 'AMD' }}
-      run: time DEBUG=3 python3 test/test_tiny.py TestTiny.test_plus
+      run: time DEBUG=3 DEV=AMD python3 test/test_tiny.py TestTiny.test_plus
    - name: Test GPU crash recovery
-      if: ${{ matrix.dev == 'AMD' }}
-      run: python3 -m pytest -rA test/external/external_test_gpu_crash.py
-    - name: Test tensor cores
-      run: |
-        if [[ "${{ matrix.dev }}" == "AMD" ]]; then
-          # Fails on 9070
-          # python3 test/test_linearizer.py test/opt/test_tensor_cores.py
-          # DEV=AMD:LLVM python3 test/test_linearizer.py test/opt/test_tensor_cores.py
-          # SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-          SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
-        else
-          ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
-        fi
+      run: DEV=AMD python3 -m pytest -rA test/external/external_test_gpu_crash.py
+    # Fails on 9070
+    # - name: Test tensor cores
+    #   run: |
+    #     DEV=AMD python3 test/test_linearizer.py test/opt/test_tensor_cores.py
+    #     DEV=AMD:LLVM python3 test/test_linearizer.py test/opt/test_tensor_cores.py
+    #     DEV=AMD SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+    - name: Run Tensor Core GEMM (AMD)
+      run: DEV=AMD SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
+    - name: Test DEV=AMD
+      run: DEBUG=2 DEV=AMD python -m pytest -rA test/test_tiny.py
    - name: Test DISK copy time
-      run: TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
+      run: DEV=AMD TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
    - name: Test CPU copy time
      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: 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
+        DEV=AMD GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
+        DEV=AMD GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
+    - name: Run full CIFAR training w 1 GPU
+      run: time BENCHMARK_LOG=cifar DEV=AMD DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    # - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+    #   run: BENCHMARK_LOG=resnet_10steps DEV=AMD MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
    - name: Run 10 MLPerf Bert training steps (1 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps 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: BENCHMARK_LOG=bert_10steps DEV=AMD CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    - name: Remote
      run: |
        pkill -f 'extra/remote/serve.py' || true
        PYTHONPATH=. python3 extra/remote/serve.py 6482 &
        sleep 1
-        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 python3 test/test_tiny.py
-        if [[ "${{ matrix.dev }}" == "AMD" ]]; then
-          DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 AMD_AQL=1 python3 test/test_tiny.py
-        fi
+        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 DEV=PCI+AMD python3 test/test_tiny.py
+        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 DEV=PCI+AMD AMD_AQL=1 python3 test/test_tiny.py
+        pkill -f 'extra/remote/serve.py' || true
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
+  testgreendriverbenchmark:
+    name: NV Benchmark
+    runs-on: [self-hosted, Linux, tinyboxrandom]
+    timeout-minutes: 20
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setcap to python
+      run: ./extra/amdpci/setup_python_cap.sh
+    - name: Remove nv modules
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py nv rmmod
+    - name: Kill stale pids
+      run: PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
+        ln -s ~/tinygrad/weights/LLaMA weights/LLaMA
+        ln -s ~/tinygrad/extra/datasets/cifar-10-python.tar.gz extra/datasets/cifar-10-python.tar.gz
+        ln -s /raid/weights/mixtral-8x7b-32kseqlen weights/mixtral-8x7b-32kseqlen
+        ln -s /raid/weights/LLaMA-2 weights/LLaMA-2
+        mkdir -p extra/datasets
+        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: test/external/process_replay/reset.py
+    - name: Test driver start time
+      run: time DEBUG=3 DEV=NV python3 test/test_tiny.py TestTiny.test_plus
+    - name: Test tensor cores
+      run: DEV=NV ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
+    - name: Test DISK copy time
+      run: DEV=NV TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
+    - name: Test CPU copy time
+      run: |
+        DEV=NV GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
+        DEV=NV GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
+    - name: Test LLAMA-3
+      run: BENCHMARK_LOG=llama3_beam DEV=NV JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0
+    - name: Run full CIFAR training w 1 GPU
+      run: time BENCHMARK_LOG=cifar DEV=NV DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+      run: BENCHMARK_LOG=resnet_10steps DEV=NV MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
+    - name: Run 10 MLPerf Bert training steps (1 gpu)
+      # TODO: remove BERT_LAYERS once scheduler is fast
+      run: BENCHMARK_LOG=bert_10steps DEV=NV CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+    - name: Remote
+      run: |
+        pkill -f 'extra/remote/serve.py' || true
+        PYTHONPATH=. python3 extra/remote/serve.py 6483 &
+        sleep 1
+        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6483 DEV=NV python3 test/test_tiny.py
        pkill -f 'extra/remote/serve.py' || true
    - name: Run process replay tests
      uses: ./.github/actions/process-replay
@@ -1,8 +1,8 @@
 name: Run MLPerf Training

 on:
-  #schedule:
-  #  - cron: '5 8 * * *'  # Runs at 08:05 UTC (12:05 AM Pacific Time)
+  schedule:
+    - cron: '5 8 * * *'  # Runs at 08:05 UTC (12:05 AM Pacific Time)
  push:
    branches:
      - update_mlperf
@@ -77,14 +77,45 @@ jobs:
      run: |
        sudo apt update || true
        sudo apt install -y --no-install-recommends ninja-build
+    - name: Test one op
+      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Test ResNet-18
      run: DEBUG=2 python3 extra/torch_backend/example.py
+    - name: custom tests
+      run: python3 -m pytest -n auto extra/torch_backend/test.py --durations=20
    - name: Test one op in torch tests
      run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
    - name: Test Ops with TINY_BACKEND
      run: DEV=CPU:LLVM LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/backend/test_ops.py --durations=20
-    - 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
+    - name: Test in-place operations on views
+      run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
+    - name: Test multi-gpu
+      run: DEV=CPU:LLVM GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
+    - name: Test kernel fusion
+      run: python3 extra/torch_backend/test_kernel_fusion.py
+
+
+  torchbackendmore:
+    name: Torch Backend Tests More
+    runs-on: *linux
+    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'
+    - name: Install ninja
+      run: |
+        sudo apt update || true
+        sudo apt install -y --no-install-recommends ninja-build
+    - 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
+    - name: Test some torch tests (expect failure)
+      run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true

  bepython:
    name: Python Backend
@@ -251,9 +282,14 @@ jobs:
      with:
        key: fuzzing-unit
        deps: testing_unit
-    - name: Fuzz Tests
-      run: |
-        parallel --tagstring '[{}]' 'python test/external/fuzz_{}.py' ::: symbolic symbolic_div fast_idiv shape_ops
+    - name: Fuzz Test symbolic
+      run: python test/external/fuzz_symbolic.py
+    - name: Fuzz Test symbolic (symbolic divisors)
+      run: python test/external/fuzz_symbolic_symbolic_div.py
+    - name: Fuzz Test fast idiv
+      run: python test/external/fuzz_fast_idiv.py
+    - name: Fuzz Test shape ops
+      run: python test/external/fuzz_shape_ops.py

  testopenclimage:
    name: CL IMAGE Tests
@@ -291,8 +327,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=55 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
+          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1361 ALLOWED_GATED_READ_IMAGE=55 FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
      - name: Test openpilot CL compile fp32 (test correctness)
        run: |
          DEV=CL IMAGE=1 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
@@ -439,8 +474,9 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: linux-${{ matrix.dev }}
-        deps: "testing_unit${{ contains(matrix.dev, 'LVP') && ' mesa' || '' }}"
+        deps: testing_unit
        llvm: ${{ contains(matrix.dev, 'LLVM') || contains(matrix.dev, 'LVP') || contains(matrix.dev, 'CLANG') }}
+        mesa: ${{ contains(matrix.dev, 'LVP') && 'cpu' || 'false' }}
        webgpu: ${{ matrix.dev == 'WEBGPU' }}
        opencl: ${{ matrix.dev == 'CL' }}
    - name: Set env
@@ -492,7 +528,7 @@ jobs:
          AMD: 0
        run: |
          PYTHONPATH=. DEV=NULL:HIP:gfx1100 python extra/mmapeak/mmapeak.py
-          PYTHONPATH=. DEV=NULL:HIP:gfx950 python3 -m pytest -n=auto test/testextra/test_tk.py
+          PYTHONPATH=. DEV=NULL:HIP:gfx950 python3 -m pytest -n=auto test/testextra/test_tk.py test/backend/test_asm_gemm.py
      - name: Run matmul on MOCKKFD
        run: |
          PYTHONPATH="." DEV=MOCKKFD+AMD N=256 python3 extra/gemm/amd_asm_matmul.py
@@ -648,33 +684,6 @@ jobs:
    - 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
@@ -683,6 +692,7 @@ jobs:
          - 'CPU:CLANG'
          - 'CPU:LLVM'
          - 'CPU:LVP'
+          - 'METAL'
          - 'WEBGPU'

    name: MacOS (DEV=${{ matrix.dev }})
@@ -695,8 +705,9 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: macos-${{ matrix.dev }}
-        deps: "testing_unit${{ contains(matrix.dev, 'LVP') && ' mesa' || '' }}"
+        deps: testing_unit
        llvm: ${{ contains(matrix.dev, 'LLVM') || contains(matrix.dev, 'LVP') }}
+        mesa: ${{ contains(matrix.dev, 'LVP') && 'cpu' || 'false' }}
        webgpu: ${{ matrix.dev == 'WEBGPU' }}
    - name: Set env
      run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
@@ -704,8 +715,8 @@ jobs:
      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 backend tests
+      run: python -m pytest -n=auto test/backend --durations=20
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -762,7 +773,8 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: compile-${{ matrix.backend }}
-          deps: "testing_unit mesa"
+          deps: testing_unit
+          mesa: ${{ (matrix.backend == 'ir3' || matrix.backend == 'nak') && 'true' }}
      - name: Set env
        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
@@ -11,7 +11,7 @@ X_train -= X_train.mean()
 # *****
 # 1. Define an MNIST model.

-from tinygrad import Tensor, Context
+from tinygrad import Tensor

 l1 = Tensor.kaiming_uniform(128, 784)
 l2 = Tensor.kaiming_uniform(10, 128)
@@ -24,11 +24,11 @@ l1n, l2n = l1.numpy(), l2.numpy()
 from tinygrad.nn.optim import SGD
 optim = SGD([l1, l2])

-with Context(TRAINING=1):
-  X, Y = X_train[(samples:=Tensor.randint(128, high=X_train.shape[0]))], Y_train[samples]
-  optim.zero_grad()
-  model(X).sparse_categorical_crossentropy(Y).backward()
-  optim.schedule_step()   # this will step the optimizer without running realize
+Tensor.training = True
+X, Y = X_train[(samples:=Tensor.randint(128, high=X_train.shape[0]))], Y_train[samples]
+optim.zero_grad()
+model(X).sparse_categorical_crossentropy(Y).backward()
+optim.schedule_step()   # this will step the optimizer without running realize

 # *****
 # 3. Create a schedule (linear uop).
@@ -67,7 +67,8 @@ def example_2_hip(a:Tensor, correct):
    # the sink specifies the GLOBAL and LOCAL sizes, along with the input buffers and name
    sink = UOp.sink(UOp.special(GLOBALS, 'gidx0'), UOp.special(THREADS, 'lidx0'), out, buf,
                    arg=KernelInfo(name="hip_reduce_sum_kernel"))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT),
+                UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
  eval_harness("HIP kernel", a, lambda x: Tensor.empty(GLOBALS).custom_kernel(x, fxn=hip_reduce_sum)[0].sum(), check=correct)

 def example_3_custom_uop(a:Tensor, correct):
@@ -122,7 +123,8 @@ def example_5_custom_assembly(a:Tensor, correct):
        offset_dwords = (self.labels[inst._target] - inst._pos - inst.size()) // 4
        if not -32768 <= offset_dwords <= 32767: raise ValueError(f"branch to '{inst._target}' offset {offset_dwords} exceeds simm16 range")
        inst.simm16 = offset_dwords
-      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in self.instructions]))))
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT),
+                                   UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in self.instructions]))))

  CU_COUNT = 32
  LANES = 64
@@ -24,7 +24,7 @@ You will see `CUDA` here on a GPU instance, or `CPU` here on a CPU instance.
 We'll use the model from [the Keras tutorial](https://keras.io/examples/vision/mnist_convnet/).

 ```python
-from tinygrad import Tensor, nn, Context
+from tinygrad import Tensor, nn

 class Model:
  def __init__(self):
@@ -74,8 +74,8 @@ We'll use the Adam optimizer. The `nn.state.get_parameters` will walk the model
 ```python
 optim = nn.optim.Adam(nn.state.get_parameters(model))
 batch_size = 128
-@Context(TRAINING=1)
 def step():
+  Tensor.training = True  # makes dropout work
  samples = Tensor.randint(batch_size, high=X_train.shape[0])
  X, Y = X_train[samples], Y_train[samples]
  optim.zero_grad()
@@ -143,6 +143,7 @@ Since we are just randomly sampling from the dataset, there's no real concept of
 for step in range(7000):
  loss = jit_step()
  if step%100 == 0:
+    Tensor.training = False
    acc = (model(X_test).argmax(axis=1) == Y_test).mean().item()
    print(f"step {step:4d}, loss {loss.item():.2f}, acc {acc*100.:.2f}%")
 ```
@@ -165,7 +165,7 @@ from extra.datasets import fetch_mnist
 Now we have everything we need to start training our neural network.
 We will be training for 1000 steps with a batch size of 64.

-We use `with Context(TRAINING=1)` to enable training mode.
+We use `with Context(TRAINING=1)` to set the internal flag `Tensor.training` to `True` during training.
 Upon exit, the flag is restored to its previous value by the context manager.

 ```python
@@ -100,7 +100,7 @@ class VLIWRenderer(Renderer):
      assert u.dtype.count in (1,8), "dtype count must be 1 or 8"

      # dumb register allocator
-      if u.op not in {Ops.STORE, Ops.SINK, Ops.INDEX}:
+      if u.op not in {Ops.STORE, Ops.SINK, Ops.GEP}:
        r[u] = reg
        reg += u.dtype.count

@@ -110,9 +110,9 @@ class VLIWRenderer(Renderer):
          inst.append({"flow": [("halt",)]})
        case Ops.CONST:
          inst.append({"load": [("const", r[u], u.arg)]})
-        case Ops.INDEX:
-          # an INDEX is just an alias to a special register in the vector
-          r[u] = r[u.src[0]] + u.src[1].arg
+        case Ops.GEP:
+          # a GEP is just an alias to a special register in the vector
+          r[u] = r[u.src[0]] + u.arg[0]
        case Ops.STACK:
          if all(s == u.src[0] for s in u.src):
            # if all sources are the same, we can broadcast
@@ -182,7 +182,7 @@ if __name__ == "__main__":
  # *** run on Machine and compare ***

  # NOTE: the scratch size needs to be reduced to 1536 when you have a register allocator
-  src = eval(prg.src[2].arg)
+  src = eval(prg.src[3].arg)
  max_regs = max(t[1] for instr in src for v in instr.values() for t in v if len(t) > 1) + 8
  print(f"{max_regs:5d} regs used" + ("" if max_regs <= 1536 else "       <-- WARNING: TOO MANY REGISTERS, MUST BE <= 1536"))
  machine = problem.Machine(mem, src, problem.DebugInfo(scratch_map={}), n_cores=1, trace=False, scratch_size=max_regs)
@@ -10,7 +10,7 @@ from extra.lr_scheduler import OneCycleLR
 from tinygrad import nn, dtypes, Tensor, Device, GlobalCounters, TinyJit, Variable
 from tinygrad.nn.state import get_state_dict
 from tinygrad.nn import optim
-from tinygrad.helpers import Context, BEAM, WINO, getenv, colored, prod, TRAINING
+from tinygrad.helpers import Context, BEAM, WINO, getenv, colored, prod
 from extra.bench_log import BenchEvent, WallTimeEvent

 cifar_mean = [0.4913997551666284, 0.48215855929893703, 0.4465309133731618]
@@ -44,7 +44,7 @@ class UnsyncedBatchNorm:
    return ret.reshape(x.shape).cast(x.dtype)

  def calc_stats(self, x:Tensor):
-    if TRAINING:
+    if Tensor.training:
      # This requires two full memory accesses to x
      # https://github.com/pytorch/pytorch/blob/c618dc13d2aa23625cb0d7ada694137532a4fa33/aten/src/ATen/native/cuda/Normalization.cuh
      # There's "online" algorithms that fix this, like https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_Online_algorithm
@@ -363,7 +363,7 @@ def train_cifar():
    st = time.monotonic()
    while i <= STEPS:
      if i % getenv("EVAL_STEPS", STEPS) == 0 and i > 1 and not getenv("DISABLE_BACKWARD"):
-        # Using Context(TRAINING=0) here actually bricks batchnorm, even with track_running_stats=True
+        # Use Tensor.training = False here actually bricks batchnorm, even with track_running_stats=True
        corrects = []
        corrects_ema = []
        losses = []
@@ -3,7 +3,7 @@ import os
 if "NOOPT" not in os.environ: os.environ["NOOPT"] = "1"
 from tinygrad import Device, nn, Tensor, dtypes
 from train_gpt2 import GPT, GPTConfig
-from tinygrad.helpers import DEV, dedup, flatten, getenv, GlobalCounters, to_function_name, Context
+from tinygrad.helpers import DEV, dedup, flatten, getenv, GlobalCounters, to_function_name
 from tinygrad.engine.realize import get_kernel
 from tinygrad.schedule.memory import memory_planner
 from tinygrad.uop.ops import Ops
@@ -23,23 +23,23 @@ if __name__ == "__main__":
  #B, T = Variable("B", 1, 128).bind(4), 64 #Variable("T", 1, 1024).bind(64)
  B, T = 4, 64

+  Tensor.training = True
  optimizer = nn.optim.Adam(nn.state.get_parameters(model), lr=1e-4)
  warmup_count = getenv("WARMUP", 3)
-  with Context(TRAINING=1):
-    for i in range(warmup_count):  # TODO: why does it take three and not two to stabilize
-      GlobalCounters.reset()
-      X = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
-      Y = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
-      _, loss = model(X, Y)
-      optimizer.zero_grad()
-      if getenv("BACKWARD", 1):
-        loss.backward()
-        tensors = optimizer.schedule_step()
-      else:
-        tensors = []
-      sched = loss.schedule(*tensors)
-      print(f"calls {i}:", len(sched))
-      #run_schedule(sched[:])
+  for i in range(warmup_count):  # TODO: why does it take three and not two to stabilize
+    GlobalCounters.reset()
+    X = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
+    Y = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
+    _, loss = model(X, Y)
+    optimizer.zero_grad()
+    if getenv("BACKWARD", 1):
+      loss.backward()
+      tensors = optimizer.schedule_step()
+    else:
+      tensors = []
+    sched = loss.schedule(*tensors)
+    print(f"calls {i}:", len(sched))
+    #run_schedule(sched[:])
  sched = memory_planner(sched)
  ast_dedup = dedup([si.ast for si in sched if si.ast.op is Ops.SINK])
  srcs = {}
@@ -2,7 +2,7 @@ import math
 from typing import Union

 from tinygrad import Tensor, nn, dtypes
-from tinygrad.helpers import prod, argfix, Context, TRAINING
+from tinygrad.helpers import prod, argfix, Context
 from tinygrad.nn.state import get_parameters
 from extra.models.unet import UNetModel

@@ -85,7 +85,7 @@ class FrozenBatchNorm2dRetinaNet(nn.BatchNorm2d):

  def __call__(self, x:Tensor) -> Tensor:
    batch_mean, batch_var = super().calc_stats(x.cast(dtypes.float32))
-    if self.track_running_stats and TRAINING:
+    if self.track_running_stats and Tensor.training:
      self.running_mean.assign((1-self.momentum) * self.running_mean + self.momentum * batch_mean.detach().cast(self.running_mean.dtype))
      self.running_var.assign((1-self.momentum) * self.running_var + self.momentum * x.numel()/(x.numel()-x.shape[1]) * batch_var.detach().cast(self.running_var.dtype))
      self.num_batches_tracked += 1
@@ -498,10 +498,11 @@ def eval_stable_diffusion():

 if __name__ == "__main__":
  # inference only
+  Tensor.training = False
+
  models = getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert").split(",")
-  with Context(TRAINING=0):
-    for m in models:
-      nm = f"eval_{m}"
-      if nm in globals():
-        print(f"eval {m}")
-        globals()[nm]()
+  for m in models:
+    nm = f"eval_{m}"
+    if nm in globals():
+      print(f"eval {m}")
+      globals()[nm]()
@@ -1,6 +1,6 @@
 # load each model here, quick benchmark
 from tinygrad import Tensor, GlobalCounters
-from tinygrad.helpers import getenv, Context
+from tinygrad.helpers import getenv
 import numpy as np

 def test_model(model, *inputs):
@@ -59,10 +59,11 @@ def spec_mrcnn():

 if __name__ == "__main__":
  # inference only for now
-  with Context(TRAINING=0):
-    for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,mrcnn").split(","):
-      nm = f"spec_{m}"
-      if nm in globals():
-        print(f"testing {m}")
-        globals()[nm]()
+  Tensor.training = False
+
+  for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,mrcnn").split(","):
+    nm = f"spec_{m}"
+    if nm in globals():
+      print(f"testing {m}")
+      globals()[nm]()

@@ -157,7 +157,6 @@ def train_resnet():
  # input_std = Tensor([0.229, 0.224, 0.225], device=GPUS, dtype=dtypes.float32).reshape(1, -1, 1, 1)
  def normalize(x): return (x.permute([0, 3, 1, 2]) - input_mean).cast(dtypes.default_float)
  @TinyJit
-  @Context(TRAINING=1)
  def train_step(X, Y):
    optimizer_group.zero_grad()
    X = normalize(X)
@@ -171,7 +170,6 @@ def train_resnet():
    return loss.realize(), top_1.realize()

  @TinyJit
-  @Context(TRAINING=0)
  def eval_step(X, Y):
    X = normalize(X)
    out = model.forward(X)
@@ -194,6 +192,7 @@ def train_resnet():
    # ** train loop **
    if MLLOGGER and RUNMLPERF:
      MLLOGGER.start(key=mllog_constants.EPOCH_START, value=e+1, metadata=dict(epoch_num=e+1))
+    Tensor.training = True
    BEAM.value = TRAIN_BEAM

    if INITMLPERF:
@@ -272,6 +271,7 @@ def train_resnet():
      eval_loss = 0.0
      eval_top_1 = 0
      eval_num_samples = 0
+      Tensor.training = False
      BEAM.value = EVAL_BEAM

      if INITMLPERF:
@@ -919,7 +919,6 @@ def train_rnnt():
  pass

@TinyJit
-@Context(TRAINING=0)
 def eval_step_bert(model, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor, masked_positions:Tensor, masked_lm_ids:Tensor,
                   masked_lm_weights:Tensor, next_sentence_labels:Tensor, GPUS):
  for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
@@ -1107,7 +1106,6 @@ def train_bert():
      MLLOGGER.start(key=mllog_constants.EPOCH_START, value=i*GBS, metadata={"epoch_num": i*GBS})

  @TinyJit
-  @Context(TRAINING=1)
  def train_step_bert(input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor,
                      masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor):
    for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
@@ -1135,6 +1133,7 @@ def train_bert():

  while train_data is not None and i < train_steps and not achieved:
    if getenv("TRAIN", 1):
+      Tensor.training = True
      BEAM.value = TRAIN_BEAM
      st = time.perf_counter()
      GlobalCounters.reset()
@@ -1187,6 +1186,7 @@ def train_bert():
      eval_lm_accs = []
      eval_clsf_accs = []
      eval_times = []
+      Tensor.training = False
      BEAM.value = EVAL_BEAM

      for j in tqdm(range(max_eval_steps), desc="Evaluating", total=max_eval_steps, disable=BENCHMARK):
@@ -1,285 +0,0 @@
-import math, os, functools
-if __name__ == "__main__":
-  os.environ["DEFAULT_FLOAT"] = "bfloat16"
-  os.environ["OPTIM_DTYPE"] = "bfloat16"
-  if "DEV" not in os.environ: os.environ["DEV"] = "NULL::gfx950"
-  # CDNA
-  os.environ["DEVICE_IN_FUNCTION_BUG"] = "1"
-  os.environ["ALL2ALL"] = "1"
-  os.environ["USE_ATOMICS"] = "1"
-from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
-from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
-from tinygrad.uop.ops import Ops, UOp
-from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
-from extra.llama_kernels.rmsnorm import rmsnorm
-from extra.gemm.cdna_asm_gemm import _mx_block_scale
-
-FP8_DTYPE = dtypes.fp8e4m3
-FP8_MAX = 448.0
-INIT_STD = 0.008
-
-def quantize_mx(x:Tensor) -> tuple[Tensor, Tensor]:
-  *batch, K = x.shape
-  scale_K = K // 32
-  amax = x.detach().float().reshape(*batch, scale_K, 32).abs().max(axis=-1)
-  e8 = (amax.maximum(1e-38).log2().floor() + 127).clamp(0, 254).cast(dtypes.uint8)
-  qscale = (127.0 - e8.cast(dtypes.float32)).exp2().reshape(*batch, scale_K, 1).expand(*batch, scale_K, 32).reshape(*batch, K)
-  x_scaled = x.float() * qscale
-  x_clamped = x_scaled + (x_scaled.detach().clamp(-FP8_MAX, FP8_MAX) - x_scaled.detach())  # STE
-  return x_clamped.cast(FP8_DTYPE), e8
-
-def _quant_dequant_fwd(x:Tensor) -> Tensor:
-  # x (2d bf16) -> bf16 value after an mxfp8 round-trip (1x32 block scaling on the last axis)
-  M, K = x.shape
-  scale_K = K // 32
-  amax = x.float().reshape(M, scale_K, 32).abs().max(axis=-1)
-  e8 = (amax.maximum(1e-38).log2().floor() + 127).clamp(0, 254).cast(dtypes.uint8)
-  qscale = (127.0 - e8.cast(dtypes.float32)).exp2().reshape(M, scale_K, 1).expand(M, scale_K, 32).reshape(M, K)
-  x_fp8 = (x.float() * qscale).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE).cast(dtypes.float32)
-  return (x_fp8 * _mx_block_scale(e8)).cast(dtypes.bfloat16)
-
-@functools.cache
-def _quant_dequant_fwd_fxn(x_p, device):
-  return _quant_dequant_fwd(Tensor(x_p, device=device))
-
-def _quant_dequant_bwd(grad:UOp, call:UOp) -> tuple:
-  return (Tensor(grad).cast(dtypes.bfloat16).uop,)
-
-def quant_dequant_mx(x:Tensor) -> Tensor:
-  fxn = _quant_dequant_fwd_fxn(x.as_param(0).uop, x.device)
-  return Tensor(UOp.maketuple(fxn.uop).call(x.uop, grad_fxn=_quant_dequant_bwd).gettuple(0))
-
-def _dequant_fwd(w_q:Tensor, w_scale:Tensor) -> Tensor:
-  return w_q.cast(dtypes.bfloat16) * _mx_block_scale(w_scale)
-
-@functools.cache
-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_block_scale(w_scale).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)
-  call = UOp.maketuple(fxn.uop).call(w_q.uop, w_scale.uop, grad_fxn=_dequant_bwd)
-  return Tensor(call.gettuple(0))
-
-def matmul_mx(x:Tensor, w_q:Tensor, w_scale:Tensor) -> Tensor:
-  l_shape = x.shape[:-1]
-  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 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)
-  x_linear = x_linear.clamp(-limit, limit)
-  return (x_glu * (alpha * x_glu).sigmoid()) * (x_linear + 1)
-
-class GPTOSS:
-  def __init__(self, dim:int, n_layers:int, n_heads:int, n_kv_heads:int, head_dim:int, n_experts:int, experts_per_tok:int,
-               intermediate_size:int, vocab_size:int, norm_eps:float=1e-5, rope_theta:int=150000, sliding_window:int=128,
-               swiglu_limit:float=7.0, max_context:int=8192):
-    self.dim, self.n_layers, self.n_heads, self.n_kv_heads, self.head_dim = dim, n_layers, n_heads, n_kv_heads, head_dim
-    self.n_rep = n_heads // n_kv_heads
-    self.n_experts, self.experts_per_tok, self.intermediate_size = n_experts, experts_per_tok, intermediate_size
-    self.vocab_size, self.norm_eps, self.sliding_window, self.swiglu_limit = vocab_size, norm_eps, sliding_window, swiglu_limit
-    self.sm_scale = 1.0 / math.sqrt(head_dim)
-
-    scaled_std = INIT_STD / math.sqrt(2 * n_layers)
-    q_dim, qkv_dim = n_heads * head_dim, head_dim * (n_heads + 2 * n_kv_heads)
-
-    # attn
-    self.wqkv, self.wqkv_scale = self._quant_weight(n_layers, qkv_dim, dim)
-    self.wqkv_bias = Tensor.zeros(n_layers, qkv_dim, dtype=dtypes.bfloat16).contiguous()
-    self.wo, self.wo_scale = self._quant_weight(n_layers, dim, q_dim, std=scaled_std)
-    self.wo_bias = Tensor.zeros(n_layers, dim, dtype=dtypes.bfloat16).contiguous()
-    self.sinks = Tensor.zeros(n_layers, n_heads, dtype=dtypes.bfloat16).contiguous()
-    self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
-
-    # moe ffn
-    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_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_bias = Tensor.zeros(n_layers, n_experts, dim, dtype=dtypes.bfloat16).contiguous()
-
-    # output
-    self.norm = nn.RMSNorm(dim, norm_eps)
-    self.tok_embeddings = nn.Embedding(vocab_size, dim)
-    self.tok_embeddings.weight = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
-    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_mx(w)
-    return w_q, w_e8.is_param_(False)
-
-  def _attn_mask(self, seqlen:int, sliding:bool, dtype) -> Tensor:
-    i, j = Tensor.arange(seqlen).reshape(seqlen, 1), Tensor.arange(seqlen).reshape(1, seqlen)
-    allowed = j <= i
-    if sliding: allowed = allowed & (i - j < self.sliding_window)
-    return allowed.where(0.0, -1e30).cast(dtype).contiguous()
-
-  def attention(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, *, attention_norm:Tensor, wqkv:Tensor, wqkv_scale:Tensor,
-                wqkv_bias:Tensor, wo:Tensor, wo_scale:Tensor, wo_bias:Tensor, sinks:Tensor):
-    bsz, seqlen, _ = x.shape
-    x_normed, rrms = rmsnorm(x, self.norm_eps)
-    qkv = matmul_mx(x_normed * attention_norm, wqkv, wqkv_scale) + wqkv_bias
-    qkv = qkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
-    xq = qkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
-    xk, xv = qkv[:, :, :, self.n_rep], qkv[:, :, :, self.n_rep + 1]
-    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
-
-    xq = xq.cast(dtypes.bfloat16).reshape(bsz, seqlen, self.n_kv_heads, self.n_rep, self.head_dim).permute(0, 2, 3, 1, 4)
-    xk = xk.cast(dtypes.bfloat16).permute(0, 2, 1, 3).unsqueeze(2)
-    xv = xv.cast(dtypes.bfloat16).permute(0, 2, 1, 3).unsqueeze(2)
-    scores = ((xq @ xk.transpose(-2, -1)).float() * self.sm_scale + mask).contiguous()
-    sink = sinks.reshape(1, self.n_kv_heads, self.n_rep, 1, 1).float()
-    m = scores.max(-1, keepdim=True).maximum(sink)
-    e = (scores - m).exp()
-    w = (e / (e.sum(-1, keepdim=True) + (sink - m).exp())).cast(dtypes.bfloat16).contiguous()
-    attn = (w @ xv).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]
-
-  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)
-
-    out = None
-    for e in range(self.n_experts):
-      gate_up = (matmul_mx(inp.contiguous_backward(), w_gate_up[e], w_gate_up_scale[e]) + w_gate_up_bias[e]).contiguous()
-      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).contiguous()
-      out = contrib if out is None else out + contrib
-    return out, [x_normed, rrms]
-
-  @function(precompile=True, precompile_backward=True)
-  def run_layer(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
-    attn, attn_saves = self.attention(x, freqs_cis, mask, **attn_kwargs)
-    h = x + attn
-    ffn, ffn_saves = self.feed_forward(h, **ffn_kwargs)
-    h = h + ffn
-    if save: return (h, *attn_saves, *ffn_saves)
-    return (h,)
-
-  def shard(self, device:tuple[str, ...], mp:bool=False):
-    assert not mp, "MP not supported"
-    from tinygrad.nn.state import get_parameters
-    for v in get_parameters(self): v.shard_(device, axis=None)
-    Tensor.realize(*get_parameters(self))
-
-  def __call__(self, tokens:Tensor, save:bool=True):
-    h = self.tok_embeddings(tokens)
-    bsz, seqlen = tokens.shape
-    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :seqlen, :, :, :]
-    mask_full = self._attn_mask(seqlen, False, dtypes.float32)
-    mask_sliding = self._attn_mask(seqlen, True, dtypes.float32)
-    for i in range(self.n_layers):
-      attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wqkv_scale=self.wqkv_scale[i],
-                         wqkv_bias=self.wqkv_bias[i], wo=self.wo[i], wo_scale=self.wo_scale[i], wo_bias=self.wo_bias[i],
-                         sinks=self.sinks[i])
-      ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], gate=self.gate[i], gate_bias=self.gate_bias[i],
-                        w_gate_up=self.w_gate_up[i], w_gate_up_scale=self.w_gate_up_scale[i], w_gate_up_bias=self.w_gate_up_bias[i],
-                        w_down=self.w_down[i], w_down_scale=self.w_down_scale[i], w_down_bias=self.w_down_bias[i])
-      mask = mask_sliding if i % 2 == 0 else mask_full
-      h, *_ = self.run_layer(h, freqs_cis, mask, attn_kwargs, ffn_kwargs, save=save)
-
-    logits = self.norm(h) @ self.output.T
-    return logits
-
-def _get_pads(uop:UOp) -> list[UOp]:
-  if uop.op == Ops.ADD: return _get_pads(uop.src[0]) + _get_pads(uop.src[1])
-  return [uop]
-
-def apply_grad(grad_buf:Tensor, new_grad:UOp):
-  pads = _get_pads(new_grad)
-  if len(pads) <= 1:
-    new_grad = new_grad.cast(grad_buf.dtype)
-    grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(grad_buf.uop + new_grad))
-    return
-  cur = grad_buf.uop
-  for pad in sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0, reverse=True):
-    if pad.op == Ops.PAD:
-      grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg)])
-      buf_slice = cur.shrink(grad_shrink)
-      cur = cur.after(buf_slice.store(buf_slice + pad.src[0].cast(cur.dtype)))
-    else:
-      cur = cur.after(cur.store(cur + pad.cast(cur.dtype)))
-  grad_buf.uop = cur
-
-GPT_OSS_20B = dict(dim=2880, n_layers=24, n_heads=64, n_kv_heads=8, head_dim=64, n_experts=32, experts_per_tok=4,
-                   intermediate_size=2880, vocab_size=128256, norm_eps=1e-5, rope_theta=150000, sliding_window=128,
-                   swiglu_limit=7.0)
-
-if __name__ == "__main__":
-  config = {}
-  BS      = config["BS"]      = getenv("BS", 16)
-  SEQLEN  = config["SEQLEN"]  = getenv("SEQLEN", 8192)
-
-  model_params = GPT_OSS_20B
-  real_vocab_size = model_params["vocab_size"]
-  if (layers := getenv("LAYERS")) != 0: model_params["n_layers"] = layers
-
-  model = GPTOSS(**model_params, max_context=SEQLEN)
-
-  state = nn.state.get_state_dict(model)
-  print("tensor count:", len(state))
-
-  from tinygrad import Device
-  is_dp = (DP := getenv("DP", 1)) > 1
-  device_count = DP
-  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))
-
-  if is_dp: model.shard(device)
-
-  # preallocate all the grad buffers and zero them out
-  grad_dtype = lambda x: dtypes.bfloat16 if x.dtype in dtypes.fp8s else x.dtype
-  grads = {x:x.zeros_like(dtype=grad_dtype(x)).contiguous() for x in state.values() if x.is_param}
-
-  # print model size
-  sz = 0
-  for k,v in state.items():
-    print(f"{colored(k, 'green' if v in grads else 'white'):30s} {str(v.shape):30s} {str(v.dtype):20s} {v.device}  {v.nbytes()/1e9:.2f} GB")
-    sz += v.nbytes()
-  print(f"total sz: {sz/1e9:.2f} GB")
-
-  with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=real_vocab_size, dtype=dtypes.int)
-  with Timing("realize weights/grads/data: "): Tensor.realize(*state.values(), *grads.values(), tokens)
-  print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
-  if is_dp: tokens = tokens.shard(device, axis=0)
-
-  @TinyJit
-  def fwd_bwd(tokens:Tensor):
-    with Timing("python forward: "):
-      logits = model(tokens[:, :-1], save=True)
-      loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-    with Timing("python backward: "):
-      for t,g in zip(grads, loss.gradient(*grads)):
-        apply_grad(grads[t], g.uop)
-    with Timing("run fwd_bwd: "): loss.realize(*grads.values())
-
-  @TinyJit
-  def optim_step():
-    for g in grads.values(): g.assign(g.zeros_like())
-    Tensor.realize(*grads.values())
-
-  for i in range(6):
-    GlobalCounters.reset()
-    profile_marker(f"step {i}")
-    with Timing(colored(f"*** step {i}: ", "red")):
-      fwd_bwd(tokens)
-      optim_step()
-  print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
@@ -14,6 +14,7 @@ export ALL2ALL=${ALL2ALL:-1}
 export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
+export USE_HK_BF16_GEMM=${USE_HK_BF16_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}
@@ -14,6 +14,7 @@ export ALL2ALL=${ALL2ALL:-1}
 export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
+export USE_HK_BF16_GEMM=${USE_HK_BF16_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}
@@ -14,6 +14,7 @@ export ALL2ALL=${ALL2ALL:-1}
 export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
+export USE_HK_BF16_GEMM=${USE_HK_BF16_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}
@@ -14,6 +14,7 @@ export ALL2ALL=${ALL2ALL:-1}
 export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
+export USE_HK_BF16_GEMM=${USE_HK_BF16_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}
@@ -3,4 +3,4 @@ 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
 SRC="AMD"; [[ $DEV == NULL* ]] && SRC="NULL"
-[ "$BENCHMARK" -le 3 ] || python -m tinygrad.viz.cli -s "$SRC" -t --interval "train @ 2" "train @ 3"
+[ $BENCHMARK -gt 3 ] && python -m tinygrad.viz.cli -s "$SRC" -t --interval "train @ 2" "train @ 3"
@@ -42,7 +42,7 @@ def compile(onnx_file):
  kernel_calls = [u for u in run_onnx_jit.captured.linear.toposort(gate=lambda x: x.op not in kernel_asts)
                  if u.op is Ops.CALL and u.src[0].op in kernel_asts]
  print(f"captured {len(kernel_calls)} kernels")
-  if getenv("TEST", 1): np.testing.assert_equal(test_val, ret, "JIT run failed")
+  np.testing.assert_equal(test_val, ret, "JIT run failed")
  print("jit run validated")

  # check gated read_image usage
@@ -50,7 +50,7 @@ def compile(onnx_file):
  read_image_count = 0
  gated_read_image_count = 0
  for call in kernel_calls:
-    _, _, source, _ = call.src[0].src
+    _, _, _, source, _ = call.src[0].src
    src = source.arg
    kernel_count += 1
    read_image_count += src.count("read_image")
@@ -38,7 +38,7 @@ def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], Li
    arg_uops = [b for b in call.src[1:] if b.op is not Ops.BIND]
    prg = to_program(call.src[0], Device[arg_uops[0].device].renderer)
    info = prg.arg
-    functions[info.function_name] = prg.src[2].arg
+    functions[info.function_name] = prg.src[3].arg
    cargs = [name_of(bu, i == 0) for i, bu in enumerate(arg_uops)] + list(info.vars)
    statements.append((info.function_name, cargs, info.global_size, info.local_size))

@@ -462,7 +462,7 @@ def test_matmul():
    lds = UOp.placeholder((lds_size,), dtypes.uint8, 0, AddrSpace.LOCAL)
    sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs, arg=KernelInfo(name=colored("kernel", "cyan"),
                                                                                  estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
  c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
  linear = c.schedule_linear()

@@ -17,7 +17,7 @@ def make_matmul_kernel(name:str, src:str, local_size:int):
    wg_y = UOp.special(N//128, "gidx1")
    sink = UOp.sink(a.base, b.base, c.base, threads, wg_x, wg_y, arg=KernelInfo(name, estimates=Estimates(ops=2*N**3, mem=3*N*N*4)))
    lib = Device[Device.DEFAULT].compiler.compile_cached(src)
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR, src=(*sink.src, sink)),
                                 UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
  return fxn

@@ -33,7 +33,7 @@ def hand_spec_tc_cores():
  acc_load = UOp.vectorize(acc.after(gk)[0], acc.after(gk)[1])
  out = UOp(Ops.WMMA, dtypes.float.vec(2), (a_tc, b_tc, acc_load), arg=wmma_arg)

-  end_loop = UOp.group(*[acc[i].store(out.index(i)) for i in range(2)]).end(gk)
+  end_loop = UOp.group(*[acc[i].store(out.gep(i)) for i in range(2)]).end(gk)

  sink = UOp.group(*[mat_idx(c.after(end_loop), gx, gy, warp, i).store(acc[i]) for i in range(2)])
  return sink.sink(arg=KernelInfo(name="custom_metal_matmul", opts_to_apply=())).simplify()
@@ -180,7 +180,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
  # store the acc into gmem
  cp_i, cp_j = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, 10004), UOp.range(BLOCK_N//TC_N, 10005)
  c_load = lambda i: C[gx, cp_i*TC_M*WARPGROUP_SIZE + warpgroup*TC_M + (warp//16)*4+i, gy, cp_j*TC_N + warp%16]
-  store = UOp.group(*[c_load(i).store(acc[cp_j, cp_i].index(i)) for i in range(4)])
+  store = UOp.group(*[c_load(i).store(acc[cp_j, cp_i].gep(i)) for i in range(4)])
  store = store.end(cp_i, cp_j)

  return store.sink(arg=KernelInfo(name="custom_gemm", opts_to_apply=())).simplify()
@@ -197,7 +197,7 @@ wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float,
 out = UOp(Ops.WMMA, dtypes.float.vec(4), (A_in, B_in, acc_load), arg=wmma_arg)

 # store back the acc
-acc = acc.after(UOp.group(*[acc[i].store(out.index(i)) for i in range(4)]).end(K_loop))
+acc = acc.after(UOp.group(*[acc[i].store(out.gep(i)) for i in range(4)]).end(K_loop))

 # store the acc into gmem
 store = UOp.group(*[C[gx, (warp//16)*4+i, gy, warp%16].store(acc[i]) for i in range(4)])
@@ -223,7 +223,7 @@ def test_matmul():
    lds = UOp.placeholder((lds_size,), dtypes.uint8, 0, AddrSpace.LOCAL)
    sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs,
                    arg=KernelInfo(name=colored("kernel","cyan"), estimates=Estimates(ops=N*N*N*2, mem=N*N*2*3)))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

  c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
  linear = c.schedule_linear()
@@ -93,7 +93,7 @@ if __name__ == "__main__":
  info = ProgramInfo(name="matmul_kernel",
                     global_size=(M//BLOCK_SIZE_M, N//BLOCK_SIZE_N, 1), local_size=(32*compiled.metadata.num_warps, 1, 1))
  sink = UOp.sink(arg=KernelInfo(name="matmul_kernel"))
-  prg_uop = to_program(UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=src)), arg=info),
+  prg_uop = to_program(UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=src)), arg=info),
                       Device.default.renderer)
  rt = get_runtime(Device.DEFAULT, prg_uop)
  all_bufs = [x.ensure_allocated() for x in bufs]
@@ -2,12 +2,12 @@ from __future__ import annotations
 from typing import cast, Callable, TypeVar, Generic, Any
 import struct, functools, time, collections, importlib, itertools, weakref
 from dataclasses import replace, dataclass, field
-from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, DEBUG, dedup, pluralize
-from tinygrad.helpers import to_tuple, round_up, partition
+from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, DEBUG, dedup, flatten, pluralize
+from tinygrad.helpers import to_tuple, round_up
 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_simple, symbolic
-from tinygrad.dtype import dtypes, AddrSpace, truncate
+from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.runtime.support.hcq import MMIOInterface
 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
@@ -23,12 +23,11 @@ class HCQ2Compiled(Compiled):

    # default pm bufferize
    self.pm_bufferize = PatternMatcher([
-      (UPat(Ops.PARAM, tag="timeline_signal"), lambda ctx: ctx.timeline_signal()),
-      (UPat(Ops.PARAM, tag="timeline_value"), lambda ctx: ctx.timeline_value()),
-      (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx.timeline_signal("sentinel", (1 << 64) - 1)),
-      (UPat(Ops.PARAM, name="b"), lambda ctx, b:
-        Buffer(ctx.device, b.max_numel(), b.dtype.base, options=BufferSpec(host=False, uncached=True, cpu_access=True, nolru=True))
-        if b.tag is not None else None), # TODO: remove nolru
+      (UPat(Ops.BUFFER, tag="timeline_signal"), lambda ctx: ctx.timeline_signal()),
+      (UPat(Ops.BUFFER, tag="timeline_value"), lambda ctx: ctx.timeline_value()),
+      (UPat(Ops.BUFFER, tag="sentinel_signal"), lambda ctx: ctx.timeline_signal("sentinel", (1 << 64) - 1)),
+      (UPat(Ops.BUFFER, name="b"), lambda ctx, b:
+        Buffer(ctx.device, b.arg, b.dtype, options=BufferSpec(host=False, uncached=True, cpu_access=True, nolru=True))), # TODO: remove nolru
    ])

    super().__init__(device, allocator, compilers, lambda *a, **kw: None, None, arch=arch)
@@ -45,6 +44,10 @@ class HCQ2Compiled(Compiled):
    buf.as_memoryview(force_zero_copy=True).cast('Q')[0] = init_value
    return buf

+  @functools.cached_property
+  def timestamps_buf(self) -> Buffer:
+    return Buffer(self.device, 0x1000, dtypes.uint8, options=BufferSpec(cpu_access=True), preallocate=True)
+
  def synchronize(self, timeout:int|None=None):
    if not hasattr(self, 'iface'): return
    sig = self.timeline_signal()._buf.cpu_view().mv.cast('Q')
@@ -129,6 +132,35 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):

  # def _as_buffer(self, buf): return buf.cpu_view().mv

+def unwrap_after(uop):
+  while uop.op is Ops.AFTER: uop = uop.src[0]
+  return uop
+
+def make_getaddr(u, device=None):
+  if unwrap_after(u).op not in (Ops.BUFFER, Ops.SLICE, Ops.BINARY, Ops.MSTACK, Ops.MSELECT): return u
+  return UOp(Ops.GETADDR, dtypes.uint64, src=(u, UOp(Ops.DEVICE, arg=device or to_tuple(u.device)[0])))
+
+def make_ins(op, *srcs):
+  return UOp(Ops.INS, dtypes.void, tuple(UOp.const(dtypes.uint32, s) if isinstance(s, int) else s.cast(dtypes.uint32) for s in srcs), op)
+
+def make_patch(buf:UOp, off:sint, val:UOp, dtype=None) -> UOp:
+  dt = dtype or val.dtype
+  return UOp(Ops.SHRINK, buf.dtype.base, (buf, UOp.const(dtypes.int, off), UOp.const(dtypes.int, dt.itemsize))).bitcast(dt).store(val.cast(dt))
+
+def make_cmdbuf(lin, devs, tag):
+  blob, patches = b'', []
+  for s in (s for ins in lin.src for s in ins.src):
+    if s.op is not Ops.CONST: patches.append((len(blob), s))
+    blob += struct.pack(f'<{s.dtype.fmt}', s.arg if s.op is Ops.CONST else 0x0)
+  buf = UOp.new_buffer(devs, len(blob), dtypes.uint8).rtag(tag)
+  return buf.after(buf.store(UOp(Ops.BINARY, dtypes.void, src=(), arg=blob)), *[make_patch(buf, off, s) for off, s in patches])
+
+def make_mstack(uops): return uops[0] if len(uops) == 1 else UOp(Ops.MSTACK, uops[0].dtype, tuple(uops))
+
+def make_signal(devs, queue=None, sentinel=False):
+  return UOp.new_buffer(devs, 1, dtypes.uint64).rtag("sentinel_signal" if sentinel else (queue, "timeline_signal") if queue else "timeline_signal")
+def make_signal_value(devs, queue=None): return UOp.new_buffer(devs, 1, dtypes.uint64).rtag((queue, "timeline_value") if queue else "timeline_value")
+
 # *****************
 # 0. helpers

@@ -137,92 +169,91 @@ HCQ_P2P_DEVS = HCQ_DEVS | frozenset(("CPU",))

 def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c

-def unwrap_after(uop):
-  while uop.op is Ops.AFTER: uop = uop.src[0]
-  return uop
-
-def make_getaddr(u, device=None):
-  if unwrap_after(u).op not in (Ops.BUFFER, Ops.SLICE, Ops.BINARY, Ops.MSTACK, Ops.MSELECT, Ops.PARAM): return u
-  return UOp(Ops.GETADDR, dtypes.uint64, src=(u,), arg=device or to_tuple(u.device)[0])
-
-def make_ins(op, *srcs):
-  return UOp(Ops.INS, dtypes.void, tuple(UOp.const(dtypes.uint32, s) if isinstance(s, int) else s.cast(dtypes.uint32) for s in srcs), op)
-
-def make_placeholder(devs, size:int, dtype, name=None, unique=True) -> UOp:
-  return UOp.param(next(UOp.unique_num) if unique else 0, dtype.ptr(size), device=devs).rtag(name or "buf")
-
-def make_patch(buf:UOp, off:sint, val:UOp, dtype=None) -> UOp:
-  dt = dtype or val.dtype
-  return UOp(Ops.SHRINK, buf.dtype.base, (buf, UOp.const(dtypes.int, off), UOp.const(dtypes.int, dt.itemsize))).bitcast(dt).store(val.cast(dt))
-
-def make_cmdbuf(lin, devs):
-  blob, patches = b'', []
-  for s in (s for ins in lin.src for s in ins.src):
-    if s.op is not Ops.CONST: patches.append((len(blob), s))
-    blob += struct.pack(f'<{s.dtype.fmt}', s.arg if s.op is Ops.CONST else 0x0)
-  buf = make_placeholder(devs, len(blob), dtypes.uint8)
-
-  # pull patches to cmdbuf
-  afters = dedup(u for _, s in patches for u in s.toposort() if u.op is Ops.AFTER)
-  deps = tuple(d for p in afters for d in p.src[1:])
-  cmdbuf = buf.after(buf.store(UOp(Ops.BINARY, dtypes.void, src=(), arg=blob)), *[make_patch(buf, off, s) for off, s in patches], *deps)
-  return cmdbuf.substitute({p: p.src[0] for p in afters}) if afters else cmdbuf
-
-def make_mstack(uops): return uops[0] if len(uops) == 1 else UOp(Ops.MSTACK, uops[0].dtype, tuple(uops))
-
-def make_signal(devs, queue=None, sentinel=False):
-  return make_placeholder(devs, 1, dtypes.uint64, "sentinel_signal" if sentinel else (queue, "timeline_signal") if queue else "timeline_signal", unique=False)
-def make_signal_value(devs, queue=None):
-  return make_placeholder(devs, 1, dtypes.uint64, (queue, "timeline_value") if queue else "timeline_value", unique=False)
-
-def make_submit(*cmds, devs:str|tuple[str, ...], queue:str) -> UOp:
-  return UOp.custom_function("submit_cmdbuf", UOp(Ops.LINEAR, dtypes.void, 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")
-
@dataclass(frozen=True)
 class HCQInfo:
-  name:str
-  estimates:Estimates
-  device:tuple[str, ...]
-  queue:str
+  name:str = ""
+  estimates:Estimates = Estimates()
+  outs:tuple[int, ...] = ()
+  devs:tuple[str, ...] = ()

-  input_idxs:tuple[int, ...] = () # indexes into input_uops used by this call
+  params:tuple[int, ...] = ()
  inputs:int|None = None

-# *****************
-# 0.1. prep: replace buffers with params
-
-def replace_call_buffers(ctx:list[UOp], call:UOp) -> UOp|None:
-  ctx += [s for s in dedup(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)])
+  @staticmethod
+  def from_call(call:UOp) -> HCQInfo: return HCQInfo(get_call_name(call, get_call_arg_uops(call)), estimate_uop(call), get_call_outs_ins(call)[0])

 # *****************
-# 1.1. prep: staging copies
+# 1.1. prep runtimes: 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.base.itemsize, dtypes.uint8)
+  stage = UOp.new_buffer("CPU", src.buffer.nbytes, dtypes.uint8)
  return UOp(Ops.LINEAR, dtypes.void, (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.1. tag hcq calls
+# 1.2. prep runtimes: programs/kernargs

-def tag_hcq_call(ctx:itertools.count, call:UOp) -> UOp:
-  if (hcq_devs:=next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None)) is None: return call
+@functools.cache
+def get_pm_prep_program(name:str) -> PatternMatcher|None:
+  try:
+    importlib.import_module(f'tinygrad.runtime.ops_{name.lower()}') # TODO: remove that
+    return importlib.import_module(f'extra.hcq2.ops_{name.lower()}2').pm_prep_program
+  except ImportError: return None

-  queue = "COMPUTE:0" if call.src[0].op is Ops.PROGRAM else "COPY:0"
-  info = HCQInfo(get_call_name(call, get_call_arg_uops(call)), estimate_uop(call), to_tuple(hcq_devs), queue)
-  return call.replace(arg=replace(call.arg, aux=info)).rtag(next(ctx))
-pm_tag_hcq_calls = PatternMatcher([(UPat(Ops.LINEAR, name="linear"),
-  lambda ctx, linear: linear.replace(src=tuple(tag_hcq_call(ctx, s) for s in linear.src)))])
+def prep_program(call:UOp, prg:UOp) -> UOp|None:
+  dev = call.src[1].device
+  if (pm:=get_pm_prep_program(to_tuple(dev)[0].split(":")[0])) is None or (lowered:=pm.rewrite(prg)) is None: return None
+
+  data, image_bytes = lowered
+  buf = UOp.new_buffer(dev, len(image_bytes), dtypes.uint8).rtag("program")
+  blob = UOp(Ops.BINARY, dtypes.void, src=(), arg=image_bytes)
+  return prg.replace(src=(buf.after(buf.store(blob)),), arg=(data, prg.arg)).call(*call.src[1:], aux=HCQInfo.from_call(call))
+
+def prep_kernargs(call:UOp, prg:UOp) -> UOp:
+  (data, info), dev_uop = prg.arg, UOp(Ops.DEVICE, arg=call.src[1].device)
+  buf = UOp.new_buffer(dev_uop.arg, data.kernargs_alloc_size, dtypes.uint8).rtag("kernargs")
+  patches = [make_patch(buf, i*8, UOp(Ops.GETADDR, dtypes.uint64, src=(call.src[1+gi], dev_uop))) for i,gi in enumerate(info.globals)] \
+          + [make_patch(buf, len(info.globals)*8 + i*4, v, dtypes.uint32) for i,v in enumerate(info.vars)]
+  return call.replace(src=(prg.replace(src=prg.src + (buf.after(*patches),), arg=(data, info)),) + call.src[1:])
+
+pm_prep_runtime = PatternMatcher([
+  # bind generic PROGRAM device to the call's actual dev(s), then run device-specific lowering
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"),),
+    name="call", allow_any_len=True), prep_program),
+
+  # lower kernargs (PROGRAM.src[0] is now AFTER(BUFFER, COPY) — the lowered program image)
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.BUFFER).or_after(),), name="prg"),), name="call", allow_any_len=True), prep_kernargs),
+])

 # *****************
-# 2.2. deps tracking
+# 2. lowering to hcq ir
+
+def make_submit(*cmds, devs:str|tuple[str, ...], queue:str) -> UOp:
+  devs:tuple[str, ...] = to_tuple(devs)
+  return UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(UOp(Ops.LINEAR, dtypes.void, src=tuple(cmds), arg=(devs, queue)),), arg="submit")
+
+def lower_program(call:UOp, prg:UOp) -> UOp:
+  return make_submit(prg, devs=call.src[1].device, queue="COMPUTE:0").sink().call(*call.src[1:], aux=call.arg.aux).rtag("hcq")
+
+def lower_copy(call:UOp, copy:UOp) -> UOp|None:
+  dst, src = call.src[1], call.src[2]
+  if (hcq_dev:=next((b.device for b in (dst, src) if b.device.split(":")[0] in HCQ_DEVS), None)) is None: return None
+
+  cp_op = UOp(Ops.COPY, dtypes.void, src=(dst, src), arg=src.buffer.nbytes)
+  return make_submit(cp_op, devs=hcq_dev, queue="COPY:0").sink().call(*call.src[1:], aux=HCQInfo.from_call(call)).rtag("hcq")
+
+pm_lower_ops = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.BUFFER).or_after(), UPat(Ops.BUFFER).or_after()), name="prg"),),
+    name="call", allow_any_len=True), lower_program),
+  (UPat(Ops.CALL, src=(UPat(Ops.COPY, name="copy"),), name="call", allow_any_len=True), lower_copy),
+])
+
+# *****************
+# 3.1. deps tracking
 # device.timeline_signal/value are the per-device schedule epoch. Before a schedule queue accesses memory owned by device N for the first time,
 # it waits for device[N].timeline_signal >= device[N].timeline_value - 1. This orders the schedule after all prior schedules that touched device N.
 #
@@ -235,95 +266,142 @@ pm_tag_hcq_calls = PatternMatcher([(UPat(Ops.LINEAR, name="linear"),
 #
 # C programs reserve and bump timeline values, then patch command buffers with the concrete wait/signal values.

-class HCQDepsTracker(DepsTracker):
-  @staticmethod
-  def _key(buf:Any) -> tuple[Any, int, int]:
-    return (buf.arg.slot, 0, buf.max_numel() * buf.dtype.base.itemsize) if isinstance(buf, UOp) else DepsTracker._key(buf)
+@dataclass
+class DepsCtx:
+  deps:DepsTracker = field(default_factory=DepsTracker)
+  opid:itertools.count = field(default_factory=lambda: itertools.count(0))
+  last_per_queue:weakref.WeakValueDictionary[tuple[Any, str], UOp] = field(default_factory=weakref.WeakValueDictionary)
+  params:dict[tuple[int, int], Buffer] = field(default_factory=dict)

-def make_deps(u:UOp, dep_lanes:list[tuple[UOp, int, int]], nlanes:int) -> UOp:
-  deps:dict[UOp, list[int|None]] = collections.defaultdict(lambda: [None]*nlanes)
-  for dep, dlane, lane in dep_lanes: deps[dep][lane] = dlane
-  return u.after(*deps, arg=tuple(tuple(v) for v in deps.values()))
-
-def sched_sync(ctx:DepsTracker, call:UOp) -> UOp|None:
-  if not isinstance(call.arg.aux, HCQInfo): return None
-
-  refs = get_call_arg_uops(call)
-  outs, _ = get_call_outs_ins(call)
-  devices, queue = call.arg.aux.device, call.arg.aux.queue
-
-  dep_lanes:list[tuple[UOp, int, int]] = []
-  for lane, d in enumerate(devices):
-    lane_refs = [b if b.op is Ops.PARAM else mb.bufs[lane] if isinstance(mb:=b.buffer, MultiBuffer) else mb for b in refs]
-    for dep, dlane in ctx.access_resources(lane_refs, outs, (call, lane)): dep_lanes.append((dep, dlane, lane))
-
-  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.arg.aux.device[dlane], dep.arg.aux.queue) != (devices[lane], queue)]
-
-  # keep latest dep per (dep device, queue, cur lane)
-  latest = {((dep.arg.aux.device[dlane], dep.arg.aux.queue), lane): (dep, dlane) for dep, dlane, lane in sorted(dep_lanes, key=lambda x: x[0].tag)}
-  return make_deps(call, [(dep, dlane, lane) for (_, lane), (dep, dlane) in latest.items()], len(devices))
-pm_sched_sync = PatternMatcher([(UPat(Ops.CALL, name="call"), sched_sync)])
-
-# *****************
-# 2.3. merge into queues
-
-def _merged_hcq_call(calls:list[UOp]):
-  info = replace(unwrap_after(calls[0]).arg.aux, estimates=sum((unwrap_after(c).arg.aux.estimates for c in calls), start=Estimates()))
-  cmdbuf = make_submit(*calls, devs=info.device, queue=info.queue)
-  return UOp.custom_function("hcq", cmdbuf.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 calls, kept in submit order
+def get_dep_buf(ctx:DepsCtx, u:UOp, lane:int) -> Buffer:
+  # TODO: should this be a part of DepsTracker?
+  if u.op is Ops.PARAM: return ctx.params.setdefault((u.arg.slot, lane), Buffer("NULL", u.max_numel(), u.dtype.base))
+  if u.op is Ops.MSTACK: return get_dep_buf(ctx, u.src[lane], 0)
+  if u.op in (Ops.SLICE, Ops.MSELECT): return get_dep_buf(ctx, u.src[0], u.arg if u.op is Ops.MSELECT else lane)
+  return b.bufs[lane] if isinstance(b:=u.buffer, MultiBuffer) else b

+def schedule_inner_sync(ctx:DepsCtx, linear:UOp) -> UOp:
+  new_src = []
  for call in linear.src:
-    if not isinstance(unwrap_after(call).arg.aux, HCQInfo):
-      new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in list(opened_qs)] + [call]
+    if call.tag != "hcq":
+      new_src.append(call)
      continue

-    devices, queue = unwrap_after(call).arg.aux.device, unwrap_after(call).arg.aux.queue
+    new_q = ctx.last_per_queue[q.arg] = (q:=get_submit(call.src[0]).src[0]).rtag(next(ctx.opid))
+    qdevs, refs = to_tuple(new_q.arg[0]), get_call_arg_uops(call)

-    if (old:=opened_qs.pop((devices, queue), None)) is not None: 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] == queue and set(k[0]) & set(devices)]
-      new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in closing]
-      new_rec = [call]
-    opened_qs[(devices, 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)])
+    # per-lane deps, tracked per (device, queue). skip self
+    dep_lanes:list[tuple[UOp, int]] = []
+    for lane, d in enumerate(qdevs):
+      for dep in ctx.deps.access_resources([get_dep_buf(ctx, b, lane) for b in refs], call.arg.aux.outs, new_q.replace(arg=(d, new_q.arg[1]))):
+        if dep.tag != new_q.tag: dep_lanes.append((dep, lane))
+
+    # drop self-queue waits, queue self-orders
+    if qdevs[0].split(":")[0] in {"AMD", "QCOM"} or new_q.arg[1].startswith("COPY"):
+      dep_lanes = [(dep, lane) for dep, lane in dep_lanes if dep.arg != (qdevs[lane], new_q.arg[1])]
+
+    # keep latest dep per lane, group lanes
+    latest = {(dep.arg, lane): dep for dep, lane in sorted(dep_lanes, key=lambda x: x[0].tag)}
+    deps:dict[UOp, tuple[int, ...]] = collections.defaultdict(tuple)
+    for (_, lane), dep in latest.items(): deps[dep] += (lane,)
+
+    if deps: new_q = new_q.after(*deps, arg=tuple(deps.values())).rtag("deps")
+    new_src.append(call.replace(src=(call.src[0].substitute({q:new_q}),)))
+  return linear.replace(src=tuple(new_src))
+pm_schedule_inner_sync = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), schedule_inner_sync)])

 # *****************
-# 2.4. finalizer
+# 3.2. finalizer

-def add_finalizer(ctx:itertools.count, linear:UOp) -> UOp:
-  # collect by device type
+def make_finalizer(queues:list[UOp], nbump:int) -> UOp:
+  devs = tuple(dedup([d for q in queues for d in to_tuple(q.arg[0])]))
+  zero = UOp.const(dtypes.int, 0)
+  tl = make_signal_value(devs)
+
+  # queue is inc with deps
+  submit = make_submit(make_signal(devs).store(tl.index(zero)), devs=devs, queue="COMPUTE:0")
+
+  # split each (multi-device) queue into per-device deps so each finalizer lane waits on the matching device's signal
+  lane_queues = [(q.replace(arg=(d, q.arg[1])), (devs.index(d),)) for q in queues for d in to_tuple(q.arg[0])]
+  submit = submit.replace(src=(submit.src[0].after(*(q for q, _ in lane_queues), arg=tuple(l for _, l in lane_queues)).rtag("deps"),))
+
+  upd = [(tl, 1)] + [(make_signal_value(devs, queue=qn), nbump) for qn in dedup([q.arg[1] for q in queues])]
+  patches = [s.after(submit).index(zero, dtype=s.dtype.ptr()).store(s.index(zero) + inc) for s, inc in upd]
+  return UOp.barrier(*patches).sink().call(aux=HCQInfo("hcq finalizer")).rtag("hcq")
+
+def add_finalizer(ctx:DepsCtx, linear:UOp) -> UOp:
  parts:dict[str, list[UOp]] = collections.defaultdict(list)
-  for call in linear.src:
-    if (c:=unwrap_after(call)).src[0].op is not Ops.CUSTOM_FUNCTION or c.src[0].arg != "hcq": continue
-    parts[c.arg.aux.device[0].split(':')[0]].append(unwrap_after(get_submit(call).src[0].src[0]))
+  for d, q in ctx.last_per_queue.items(): parts[to_tuple(d[0])[0].split(':')[0]].append(q)

-  nbump = next(ctx)
-  finalizers = []
-  for calls in parts.values():
-    devs = tuple(dedup(d for call in calls for d in unwrap_after(call).arg.aux.device))
-    zero = UOp.const(dtypes.int, 0)
-    tl = make_signal_value(devs)
-
-    # split each (multi-device) call into per-device deps, then store the device timeline value into the device signal after them
-    dep_lanes = [(call, dlane, devs.index(d)) for call in calls for dlane, d in enumerate(unwrap_after(call).arg.aux.device)]
-    store = make_deps(make_signal(devs).store(tl.index(zero)), dep_lanes, len(devs))
-    submit = make_submit(store, devs=devs, queue="COMPUTE:0")
-
-    upd = [(tl, 1)] + [(make_signal_value(devs, queue=qn), nbump) for qn in dedup([unwrap_after(call).arg.aux.queue for call in calls])]
-    patches = [s.after(submit).index(zero, dtype=s.dtype).store(s.index(zero) + inc) for s, inc in upd]
-    finalizers.append(UOp.custom_function("hcq", UOp.barrier(*patches).sink()).call(aux=HCQInfo("hcq finalizer", Estimates(), devs, "COMPUTE:0")))
-  return linear.replace(src=linear.src + tuple(finalizers))
+  nbump = next(ctx.opid)
+  return linear.replace(src=linear.src + tuple([make_finalizer(queues, nbump) for queues in parts.values()]))
 pm_add_finalizer = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), add_finalizer)])

 # *****************
-# 2.4. global sync
+# 3.3. lower loads/stores
+
+def add_loads(ctx:set[int], deps:UOp) -> UOp:
+  cur_devs = to_tuple((cur:=deps.src[0]).arg[0])
+
+  waits = []
+  for lanes, dep in zip(deps.arg, deps.src[1:]):
+    dep_dev, queue = dep.arg # dep_dev is a single device (deps are recorded per-device)
+    ctx.add(dep.tag) # mark op to update signal.
+
+    # for lanes that need this dep, wait on the dep device's signal/value; other lanes get a passing sentinel
+    lanes = set(lanes)
+    sig = make_mstack([make_signal(dep_dev if j in lanes else d, queue=queue, sentinel=j not in lanes) for j, d in enumerate(cur_devs)])
+    val = make_mstack([make_signal_value(dep_dev if j in lanes else d, queue=queue) for j, d in enumerate(cur_devs)]).index(UOp.const(dtypes.int, 0))
+    waits.append(sig.wait(val + dep.tag))
+  return cur.replace(src=tuple(waits) + cur.src)
+pm_add_inner_loads = PatternMatcher([(UPat(Ops.AFTER, tag="deps", name="deps"), add_loads)])
+
+def add_stores(ctx:set[int], submit:UOp, q:UOp) -> UOp|None:
+  if q.tag not in ctx: return None
+  devs, queue = q.arg
+  src = q.src + (make_signal(devs, queue=queue).store(make_signal_value(devs, queue=queue).index(UOp.const(dtypes.int, 0)) + q.tag),)
+  return submit.replace(src=(q.replace(src=src, tag=None),))
+pm_add_inner_stores = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit", src=(UPat(Ops.LINEAR, name="q"),), name="submit"), add_stores)])
+
+# *****************
+# 4.1. merge queues
+
+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")
+
+def merge_sink(sinks:list[UOp]) -> UOp:
+  if len(sinks) == 1: return sinks[0]
+  submits = [get_submit(sink) for sink in sinks]
+  queues = [submit.src[0] for submit in submits]
+  anchor = submits[-1].replace(src=(queues[-1].replace(src=tuple(x for q in queues for x in q.src)),))
+  for sink, submit in zip(sinks[:-1], submits[:-1]):
+    if sink.src[0] is not submit: anchor = sink.src[0].substitute({submit: anchor}, walk=True)
+  return sinks[-1].substitute({submits[-1]: anchor}, walk=True)
+
+def merge_queues(linear:UOp) -> UOp:
+  new_src:list[UOp] = []
+  opened_qs:dict[tuple[tuple[str, ...], str], tuple[list[UOp], HCQInfo]] = {} # (devs, queue) -> (sinks, aux), kept in submit order
+
+  for call in linear.src:
+    # finalizer cannot be merged, since it bumps inner signal (this introduces race when multidevs).
+    if call.tag != "hcq" or (call.tag == "hcq" and call.arg.aux.name == "hcq finalizer"):
+      new_src += [merge_sink((sa:=opened_qs.pop(k))[0]).call(aux=sa[1]).rtag("hcq") for k in list(opened_qs)] + [call]
+      continue
+
+    devs, queue = get_submit(new_sink:=call.src[0]).src[0].arg
+    new_rec = ([new_sink], call.arg.aux)
+    if (old:=opened_qs.pop((devs, queue), None)) is not None:
+      new_rec = (old[0] + [new_sink], replace(new_rec[1], name=f"{queue.lower()} submit", estimates=old[1].estimates + new_rec[1].estimates))
+    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] == queue and set(k[0]) & set(devs)]
+      new_src += [merge_sink((sa:=opened_qs.pop(k))[0]).call(aux=sa[1]).rtag("hcq") for k in closing]
+    opened_qs[(devs, queue)] = new_rec
+
+  return linear.replace(src=tuple(new_src + [merge_sink(sinks).call(aux=aux).rtag("hcq") for sinks, aux in opened_qs.values()]))
+pm_merge_queues = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), merge_queues)])
+
+# *****************
+# 4.2. global sync

 def add_global_sync(ctx:set[tuple[str, ...]], submit:UOp, q:UOp) -> UOp|None:
  if (devs:=q.arg[0]) in ctx: return None
@@ -332,143 +410,94 @@ def add_global_sync(ctx:set[tuple[str, ...]], submit:UOp, q:UOp) -> UOp|None:
  # some devices from a command buffer might be used for the first time this schedule, so we wait for their global timeline epoch.
  wait = make_signal(devs).wait(make_signal_value(devs).index(UOp.const(dtypes.int, 0)) - 1)
  return submit.replace(src=(q.replace(src=(UOp(Ops.BARRIER, dtypes.void), wait, *q.src)),))
-pm_add_global_sync = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),), name="submit"), add_global_sync)])
+pm_add_global_sync = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit", src=(UPat(Ops.LINEAR, name="q"),), name="submit"), add_global_sync)])

 # *****************
-# 3.1. lower loads/stores
+# 4.3. annotate exec devs

-def add_loads(ctx:set[int], submit:UOp, q:UOp) -> UOp|None:
-  cur_devs = q.arg[0]
-  new_src:list[UOp] = []
-  for s in q.src:
-    if s.op is Ops.AFTER:
-      for lanes, dep in zip(s.arg, s.src[1:]):
-        devs, queue = dep.arg.aux.device, dep.arg.aux.queue
-        ctx.add(dep.tag) # mark op to update signal.
-
-        sig = make_mstack([make_signal(d if dl is None else devs[dl], queue=queue, sentinel=dl is None) for dl, d in zip(lanes, cur_devs)])
-        val = make_mstack([make_signal_value(d if dl is None else devs[dl], queue=queue) for dl, d in zip(lanes, cur_devs)]).index(UOp.const(dtypes.int, 0))
-        new_src.append(sig.wait(val + dep.tag))
-      s = s.src[0]
-    new_src.append(s)
-  return submit.replace(src=(q.replace(src=tuple(new_src)),))
-pm_add_inner_loads = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),), name="submit"), add_loads)])
-
-def add_stores(ctx:set[int], submit:UOp, q:UOp) -> UOp|None:
-  devs, queue = q.arg
-  new_src:list[UOp] = []
-  for op in q.src:
-    new_src.append(op)
-    if (sigval:=unwrap_after(op).tag) in ctx:
-      new_src.append(make_signal(devs, queue=queue).store(make_signal_value(devs, queue=queue).index(UOp.const(dtypes.int, 0)) + sigval))
-  return submit.replace(src=(q.replace(src=tuple(new_src)),))
-pm_add_inner_stores = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),), name="submit"), add_stores)])
-
-# *****************
-# 2.1. hcq lowering: programs
-
-def encode_kernargs_clike(call:UOp, prg:UOp, devs:str|tuple[str, ...]) -> UOp:
-  data, info = prg.arg
-  call_args = get_call_arg_uops(call)
-  buf = make_placeholder(devs, data.kernargs_alloc_size, dtypes.uint8, name="kernargs")
-  patches = [make_patch(buf, i*8, make_getaddr(call_args[gi], devs)) for i,gi in enumerate(info.globals)] \
-          + [make_patch(buf, len(info.globals)*8 + i*4, v, dtypes.uint32) for i,v in enumerate(info.vars)]
-  return buf.after(*patches)
-
-# *****************
-# 2.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)])
-
-# *****************
-
-pm_early_simplify = PatternMatcher([
-  # getaddr(slice(base, off)) -> getaddr(base) + byte offset
-  (UPat(Ops.GETADDR, src=(UPat(Ops.SLICE, src=(UPat.var("base"), UPat.cvar("off"))),), name="g"),
-   lambda g, base, off: g.replace(src=(base,)) + UOp.const(dtypes.uint64, off.arg * base.dtype.itemsize)),
-])
+pm_annotate_devs = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"),
+  lambda call: call.replace(arg=replace(call.arg, aux=replace(call.arg.aux, devs=get_submit(call.src[0]).src[0].arg[0]))))])

 # *****************
+# 4.4. replace params with per-submit input address loads

 def replace_params(call:UOp) -> UOp|None:
-  gaddrs = [u for u in call.src[0].toposort(enter_calls=False) if u.op is Ops.GETADDR and u.src[0].op is Ops.PARAM and u.src[0].tag is None]
-  if not gaddrs: return None
+  if not (params:={u:u.arg.slot for u in call.src[0].toposort() if u.op is Ops.PARAM and u.addrspace is AddrSpace.GLOBAL}): return None

-  idxs:dict[int, int] = {}
-  for g in gaddrs: idxs.setdefault(g.src[0].arg.slot, len(idxs))
+  # fill new info
+  hcqinfo = replace(call.arg.aux, params=tuple(sorted(set(params.values()))), inputs=len(get_call_arg_uops(call)))

-  inputs = make_placeholder(call.arg.aux.device, len(idxs), dtypes.uint64, "inputs")
-  body = call.src[0].substitute({g: inputs.index(UOp.const(dtypes.int, idxs[g.src[0].arg.slot])).load() for g in gaddrs})
-  return call.replace(src=(body, *call.src[1:], inputs), arg=replace(call.arg, aux=replace(call.arg.aux, input_idxs=tuple(idxs))))
-pm_replace_params = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), replace_params)])
+  inputs = UOp.new_buffer(get_submit(call.src[0]).src[0].arg[0], len(hcqinfo.params), dtypes.uint64).rtag("inputs")
+
+  slot2idx = {s:i for i,s in enumerate(hcqinfo.params)}
+  body = call.src[0].substitute({u:inputs.index(UOp.const(dtypes.int, slot2idx[s])).load() for u,s in params.items()})
+
+  return call.replace(src=(body, *call.src[1:], inputs), arg=replace(call.arg, aux=hcqinfo))
+pm_replace_params = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), replace_params)])

 # *****************
+# 5.1. encode cmdbufs

-def changes_per_submit(u:UOp) -> bool: return u.op in (Ops.LOAD, Ops.INDEX) or (u.op is Ops.PARAM and u.tag is None)
-def is_placeholder(b:UOp) -> bool: return (b.op is Ops.PARAM and b.tag is not None) or (b.op is Ops.MSTACK and all(is_placeholder(x) for x in b.src))
-def is_link_patch(s:UOp) -> bool: return is_placeholder(s.buf_uop) and not any(changes_per_submit(u) for u in s.backward_slice)
+@functools.cache
+def get_pm_lower(name:str) -> PatternMatcher|None:
+  try:
+    importlib.import_module(f'tinygrad.runtime.ops_{name.lower()}') # TODO: remove that
+    return importlib.import_module(f'extra.hcq2.ops_{name.lower()}2').pm_lower
+  except ImportError: return None

-def trim_link_patches(ctx:list[UOp], a:UOp) -> UOp|None:
-  links, kept = partition(a.src[1:], is_link_patch)
-  ctx += links
-  return a.src[0].after(*kept) 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 encode_cmdbuf(submit:UOp, lin:UOp) -> UOp|None:
+  if (pm:=get_pm_lower(to_tuple(lin.arg[0])[0].split(":")[0])) is None: return None
+  return pm.rewrite(submit)
+pm_encode_cmdbufs = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit", src=(UPat(Ops.LINEAR, name="lin"),), name="submit"), encode_cmdbuf)])

-def split_patches(call:UOp) -> UOp|None:
-  # trim link-time patches
-  body = graph_rewrite(call.src[0], pm_trim_link_patches, ctx=(lt_patches:=[]), name=f"trim link-time patches ({call.arg.aux.name})")
+# *****************
+# 5.2. lift patches to the command buffer (root)

-  units:dict[UOp, None] = {}
-  def unit_gate(u:UOp) -> bool:
-    if (is_plc:=is_placeholder(u)): units[u] = None
-    return not is_plc
-  body.toposort(gate=unit_gate)
-
-  srcs = dedup(list(call.src[1:]) + list(units) + [s.buf_uop for s in lt_patches])
-  param_sub = {u: UOp.param(i, u.dtype, device=u.device) for i,u in enumerate(srcs)}
-  for b in dedup(s.buf_uop for s in lt_patches):
-    idx = param_sub[b].arg.slot
-    srcs[idx] = srcs[idx].after(*dedup(s for s in lt_patches if s.buf_uop is b))
-
-  param_sub |= {v: v.replace(arg=replace(v.arg, slot=-1)) for v in body.variables() if v.op is Ops.PARAM}
-
-  info = replace(call.arg.aux, inputs=next((i for i,u in enumerate(srcs) if unwrap_after(u).tag == "inputs"), None))
-  return call.replace(src=(body.substitute(param_sub), *srcs), arg=replace(call.arg, aux=info))
-pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), split_patches)])
+def lift_patches_to_cmdbuf(cmdbuf:UOp) -> UOp|None:
+  if not (patches:=dedup(u for store in cmdbuf.src[1:] for u in store.toposort() if u.op is Ops.AFTER)): return None
+  deps = tuple(d for p in patches for d in p.src[1:])
+  return cmdbuf.replace(src=cmdbuf.src + deps).substitute({p: p.src[0] for p in patches})
+pm_lift_patches_to_cmdbuf = PatternMatcher([
+  (UPat(Ops.AFTER, src=(UPat(Ops.BUFFER, tag={"compute", "copy"}),), allow_any_len=True, name="cmdbuf"), lift_patches_to_cmdbuf),
+])

 # *****************
 # 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 (maxtags:={"scratch"}) | (sumtags:={"program", "kernargs"})]
+  bufs = [b for b in call.src[0].toposort() if b.op is Ops.BUFFER and b.tag in (maxtags:={"scratch"}) | (sumtags:={"program", "kernargs"})]

  off_per_buf:dict[UOp, int] = {}
  size_per_tag:dict[str, int] = {}
  for b in bufs:
-    bsz = b.max_numel()
-    if b.tag in maxtags: size_per_tag[b.tag] = max(size_per_tag.get(b.tag, 0), bsz)
+    if b.tag in maxtags: size_per_tag[b.tag] = max(size_per_tag.get(b.tag, 0), b.arg)
    elif b.tag in sumtags:
      off_per_buf[b] = round_up(size_per_tag.get(b.tag, 0), {"program": 0x1000}.get(b.tag, 128))
-      size_per_tag[b.tag] = off_per_buf[b] + bsz
+      size_per_tag[b.tag] = off_per_buf[b] + b.arg

  count_per_tag = collections.Counter(b.tag for b in bufs)
  ref_bufs = {b.tag:b for b in bufs if count_per_tag[b.tag] > 1}
-  bases = {tag:UOp.new_buffer(b.device, size_per_tag[tag], b.dtype).rtag(tag) for tag,b in ref_bufs.items()}
-  subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(dtypes.weakint, off_per_buf.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases}
+  bases = {tag:UOp.new_buffer(b.src[1].arg, size_per_tag[tag], b.dtype).rtag(tag) for tag,b in ref_bufs.items()}
+  subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(dtypes.weakint, off_per_buf.get(b, 0))), b.arg) 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)])
+pm_pack_placeholders = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), pack_hcq_placeholders)])
+
+# *****************
+# 5.4. capture buffers reachable from each hcq call as BIND, so we don't drop their refs
+
+def hold_call_buffers(call:UOp) -> UOp|None:
+  if not (bufs:=tuple(dedup(u for u in call.src[0].toposort() if u.op is Ops.BUFFER and u not in call.src))): return None
+  return call.replace(src=call.src + (UOp(Ops.BIND, dtypes.void, src=bufs),))
+pm_hold_call_buffers = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), hold_call_buffers)])

 # *****************
 # 6. bufferize placeholders: replace placeholders with real buffers.

 def bufferize_buf(buf:UOp) -> UOp|None:
  if buf.tag is None: return None
-  uops = tuple(UOp.from_buffer((dv:=Device[dev]).pm_bufferize.rewrite(buf, ctx=dv), "CPU") for dev in to_tuple(buf.device))
+  uops = tuple(UOp.from_buffer((dv:=Device[dev]).pm_bufferize.rewrite(buf, ctx=dv), "CPU") for dev in to_tuple(buf.src[1].arg))
  return make_mstack(uops)
-pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, name="buf"), bufferize_buf)])
+pm_bufferize = PatternMatcher([(UPat(Ops.BUFFER, name="buf"), bufferize_buf)])

 # *****************
 # 7. resolve patches
@@ -482,20 +511,20 @@ def fold_blob_store(buf:UOp, blob:UOp) -> UOp:

 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)):
-    struct.pack_into(f'<{v.dtype.fmt}', b.ensure_allocated()._buf.cpu_view().mv.cast('B'), off.arg * buf.dtype.base.itemsize, truncate[v.dtype](v.arg))
+    struct.pack_into(f'<{v.dtype.fmt}', b.ensure_allocated()._buf.cpu_view().mv.cast('B'), off.arg * buf.dtype.base.itemsize, v.arg)
  return UOp(Ops.NOOP)

 def resolve_getaddr(buf:UOp, g:UOp) -> UOp:
  if buf.op not in (Ops.BUFFER, Ops.MSTACK, Ops.MSELECT): return buf
-  devs, b = to_tuple(g.arg), buf.buffer
+  devs, b = to_tuple(g.src[1].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.vec(len(addrs)), addrs)

-def resolve_getaddr_slice(bv:UOp, g:UOp) -> UOp:
+def resolve_getaddr_slice(bv:UOp, dev:UOp) -> UOp:
  itemsize = bv.src[0].dtype.itemsize if unwrap_after(bv.src[0]).op in (Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize
-  return UOp(Ops.GETADDR, dtypes.uint64, src=(bv.src[0],), arg=g.arg) + UOp.const(dtypes.uint64, bv.src[1].arg * itemsize)
+  return UOp(Ops.GETADDR, dtypes.uint64, src=(bv.src[0], dev)) + UOp.const(dtypes.uint64, bv.src[1].arg * itemsize)

 pm_resolve_patches = PatternMatcher([
  # multi
@@ -507,54 +536,62 @@ pm_resolve_patches = PatternMatcher([
    lambda shr, bv: shr.replace(src=(bv.src[0], shr.src[1] + bv.src[1].cast(shr.src[1].dtype), shr.src[2]))),

  # getaddr
-  (UPat(Ops.GETADDR, src=(UPat(Ops.SLICE, name="bv"),), name="g"), resolve_getaddr_slice), # getaddr(slice(x)) -> offset+getaddr(x)
-  (UPat(Ops.GETADDR, src=(UPat(name="buf"),), name="g"), resolve_getaddr),
+  (UPat(Ops.GETADDR, src=(UPat(Ops.SLICE, name="bv"), UPat(Ops.DEVICE, name="dev"))), resolve_getaddr_slice), # getaddr(slice(x)) -> offset+getaddr(x)
+  (UPat(Ops.GETADDR, src=(UPat(name="buf"), UPat(Ops.DEVICE)), name="g"), resolve_getaddr),

  # folders
  (UPat({Ops.BUFFER, Ops.SLICE, Ops.MSTACK}, name="buf").store(UPat(Ops.BINARY, name="blob")), fold_blob_store),
  (UPat(Ops.SHRINK, src=(UPat({Ops.BUFFER, Ops.SLICE, Ops.MSTACK}, name="buf"), UPat.cvar("off"), UPat(Ops.CONST))).bitcast()
    .store(UPat.any(UPat.cvar("val"), UPat(Ops.STACK, name="val"))), fold_const_store),
-])
+]) + symbolic_simple

 # *****************
 # 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["CPU"].renderer),)))])
+def to_param(bufs:list[UOp], ref:UOp) -> UOp:
+  if ref not in bufs: bufs.append(ref)
+  return UOp.placeholder((ref.buffer.size,), ref.dtype, bufs.index(ref))
+pm_to_param = PatternMatcher([(UPat({Ops.MSELECT, Ops.MSTACK, Ops.BUFFER}, name="r"), lambda ctx, r: to_param(ctx, r))])

-hcq_compile_cache:dict[bytes, tuple[UOp, tuple[UOp, ...]]] = {}
+def parametrize_host_buffers(call:UOp) -> UOp:
+  # preserve original order of args
+  body = graph_rewrite(call.src[0], pm_to_param, ctx=(bufs:=list(get_call_arg_uops(call))), bottom_up=True, name="parametrize host buffers")

-@track_rewrites(lambda linear,input_uops,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
-def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None) -> 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")
+  # move vars to new slots
+  var_slots = {nm:len(bufs)+i for i,nm in enumerate(sorted({v.expr for v in body.variables() if v.op is Ops.PARAM}))}
+  body = body.substitute({v:v.replace(arg=replace(v.arg, slot=var_slots[v.expr])) for v in body.variables() if v.op is Ops.PARAM})

-  if (final_linear:=(hcq_compile_cache.get(cache_key:=linear.key))) is None:
-    # schedule
-    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")
-    linear = graph_rewrite(linear, pm_tag_hcq_calls, ctx=(enumerator:=itertools.count(0)), walk=True, name="tag hcq calls")
-    linear = graph_rewrite(linear, pm_sched_sync, ctx=HCQDepsTracker(), walk=True, name="schedule sync")
-    linear = linear.substitute({s: p for p, s in back_map.items()}, walk=True)
-    linear = graph_rewrite(linear, pm_merge_queues, walk=True, name="merge queues")
-    linear = graph_rewrite(linear, pm_add_finalizer, ctx=enumerator, walk=True, name="add finalizer")
-    linear = graph_rewrite(linear, pm_add_global_sync, ctx=set(), walk=True, name="add global sync", enter_calls=True)
+  return call.replace(src=(body, *bufs) + tuple(x for x in call.src[1:] if x.op is Ops.BIND))
+pm_parametrize_host_buffers = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), parametrize_host_buffers)])

-    # lowering to hcq ir
-    linear = graph_rewrite(linear, pm_add_inner_loads, ctx=(waited:=set()), walk=True, name="add loads", enter_calls=True)
-    linear = graph_rewrite(linear, pm_add_inner_stores, ctx=waited, walk=True, name="add stores", enter_calls=True)
-    linear = graph_rewrite(linear, pm_encode_cmdbufs, walk=True, name="encode cmdbufs", enter_calls=True)
+def callify_hcq(call:UOp) -> UOp:
+  prg = to_program(call.src[0].sink(arg=KernelInfo("hcq_submit"), tag=1), Device["CPU"].renderer)
+  return UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(prg,), arg="hcq").call(*call.src[1:], aux=call.arg.aux)
+pm_callify_hcq = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), callify_hcq)])

-    # pie
-    linear = graph_rewrite(linear, pm_early_simplify + symbolic, bottom_up=False, name="early simplify patches")
-    linear = graph_rewrite(linear, pm_replace_params, walk=True, name="replace params")
-    linear = graph_rewrite(linear, pm_split_patches, walk=True, name="split patches")
+@track_rewrites(lambda _,ret: f"HCQ Schedule {pluralize('Kernel', len(ret.src))}")
+def hcq_schedule(linear:UOp) -> UOp:
+  linear = graph_rewrite(linear, pm_insert_copy_staging + pm_flatten_linear, name="insert copy staging")
+  linear = graph_rewrite(linear, pm_prep_runtime, name="prepare runtime")

-    # and compile it
-    final_linear = hcq_compile_cache[cache_key] = graph_rewrite(linear, pm_callify_hcq, name="callify hcq", enter_calls=True)
+  linear = graph_rewrite(linear, pm_lower_ops, name="lower ops into hcq ir")
+  linear = graph_rewrite(linear, pm_schedule_inner_sync, ctx=(deps_ctx:=DepsCtx()), walk=True, name="schedule inner sync")
+  linear = graph_rewrite(linear, pm_add_finalizer, ctx=deps_ctx, walk=True, name="add finalizer")
+  linear = graph_rewrite(linear, pm_add_inner_loads, ctx=(waited:=set()), walk=True, name="add loads", enter_calls=True)
+  linear = graph_rewrite(linear, pm_add_inner_stores, ctx=waited, walk=True, name="add stores", enter_calls=True)
+  linear = graph_rewrite(linear, pm_merge_queues, name="merge queues")
+  linear = graph_rewrite(linear, pm_add_global_sync, ctx=set(), walk=True, name="add global sync", enter_calls=True)
+  linear = graph_rewrite(linear, pm_annotate_devs, name="annotate devs")
+  linear = graph_rewrite(linear, pm_replace_params, name="replace params")
+  linear = graph_rewrite(linear, pm_encode_cmdbufs, walk=True, name="encode cmdbufs", enter_calls=True)
+  linear = graph_rewrite(linear, pm_lift_patches_to_cmdbuf, name="lift patches to cmdbuf", enter_calls=True)
+  linear = graph_rewrite(linear, pm_pack_placeholders, walk=True, name="pack placeholders")
+  linear = graph_rewrite(linear, pm_hold_call_buffers, walk=True, name="hold call buffers")

-  return final_linear
+  # realize starts from here
+  linear = graph_rewrite(linear, pm_bufferize, bottom_up=True, walk=True, name="bufferize placeholders", enter_calls=True)
+  linear = graph_rewrite(linear, pm_resolve_patches, bottom_up=False, name="simplify patches", enter_calls=True)
+  linear = graph_rewrite(linear, pm_parametrize_host_buffers, walk=True, name="parametrize host buffers")
+  linear = graph_rewrite(linear, pm_callify_hcq, name="callify hcq")

-@track_rewrites(lambda _,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}")
-def hcq_link(linear:UOp) -> UOp:
-  linear = graph_rewrite(linear, pm_bufferize, bottom_up=True, walk=True, name="bufferize placeholders")
-  return graph_rewrite(linear, pm_resolve_patches + symbolic, bottom_up=False, name="simplify patches")
+  return linear
@@ -3,7 +3,7 @@ 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_getaddr, make_ins, make_cmdbuf, make_placeholder
+from extra.hcq2.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, make_getaddr, make_ins, make_cmdbuf
 from tinygrad.uop.ops import sint, UOp
 from tinygrad.device import Compiled, BufferSpec, Buffer, Device
 from tinygrad.dtype import dtypes
@@ -102,12 +102,11 @@ def pm4_timestamp(ctx, dst):
  return release_mem(ctx, make_getaddr(dst, ctx.devs), 0, ctx.pm4.data_sel__mec_release_mem__send_gpu_clock_counter,
                     ctx.pm4.int_sel__mec_release_mem__none)

-def pm4_program(ctx, call, prg):
+def pm4_program(ctx, prg):
  data, info = prg.arg
-  lib_gpu = prg.src[0]
-  args = encode_kernargs_clike(call, prg, ctx.devs)
+  lib_gpu, args = prg.src
  prog_addr = make_getaddr(lib_gpu, ctx.devs) + data.entry_point_offset
-  scratch_addr = make_getaddr(make_placeholder(ctx.devs, data.private_segment_size, dtypes.uint8, "scratch", unique=False), ctx.devs)
+  scratch_addr = make_getaddr(UOp.new_buffer(lib_gpu.device, data.private_segment_size, dtypes.uint8).rtag("scratch"), ctx.devs)
  args_addr = make_getaddr(args, ctx.devs)

  user_regs = []
@@ -135,20 +134,19 @@ def pm4_program(ctx, call, prg):
  return UOp(Ops.LINEAR, dtypes.void, tuple(ins))

 pm_pm4_opsel = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), pm4_program),
-
  (UPat(Ops.WAIT, src=(UPat(name="dst"), UPat(name="val"))), pm4_wait),
  (UPat(Ops.BARRIER), pm4_barrier),
+  (UPat(Ops.PROGRAM, name="prg"), pm4_program),
  (UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", src=(UPat(name="dst"),)), pm4_timestamp),
  (UPat(Ops.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)
+  size, zero = UOp.const(dtypes.uint32, cmdbuf.src[0].arg // dtypes.uint32.itemsize), UOp.const(dtypes.int, 0)

  # 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 = (make_placeholder(devs, b.size, b.dtype, ("COMPUTE:0", name), unique=False)
+  ring, wptr, doorbell, put_ptr = (UOp.new_buffer(devs, b.size, b.dtype).rtag(("COMPUTE:0", name))
    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))

  # place the cmdbuf at the ring's write offset, wrapping the ring
@@ -158,31 +156,29 @@ def pm4_submit(cmdbuf, devs):
  ring_idx = ((put + i.cast(put.dtype)) % q.ring.size).cast(dtypes.int)

  # copy the cmdbuf into the ring and advance the put/write pointers
-  copy_to_ring = ring.index(ring_idx, ptr=True).store(
-    cmdbuf.index(i*4, ptr=True).cast(dtypes.uint32.ptr()).load()).end(i)
-  bump_put_ptr = put_ptr.index(zero, ptr=True).store(next_put)
-  bump_wptr = wptr.index(zero, ptr=True).store(next_put)
+  copy_to_ring = ring.index(ring_idx, dtype=ring.dtype.ptr()).store(
+    cmdbuf.index(i*4, dtype=cmdbuf.dtype.ptr()).cast(dtypes.uint32.ptr()).load()).end(i)
+  bump_put_ptr = put_ptr.index(zero, dtype=put_ptr.dtype.ptr()).store(next_put)
+  bump_wptr = wptr.index(zero, dtype=wptr.dtype.ptr()).store(next_put)

  # 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, ptr=True).store(next_put)
+  return doorbell.after(flush).index(zero, dtype=doorbell.dtype.ptr()).store(next_put)

 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])))])
+  lambda lin: pm4_submit(make_cmdbuf(lin, to_tuple(lin.arg[0]), "compute"), to_tuple(lin.arg[0])))])

 # *****************
 # SDMA

 class SDMAOps(FastEnum): COPY = auto(); POLL_REGMEM = auto(); FENCE = auto(); TRAP = auto(); TIMESTAMP = auto()  # noqa: E702

-def sdma_copy(ctx, call):
-  dst, src = call.src[1], call.src[2]
-  sz = src.max_numel() * src.dtype.base.itemsize
+def sdma_copy(ctx, dst, src, copy):
  src_addr, dst_addr = make_getaddr(src, ctx.devs), make_getaddr(dst, ctx.devs)
  return UOp(Ops.LINEAR, dtypes.void, tuple([make_ins(SDMAOps.COPY,
     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)) for off in range(0, sz, ctx.max_copy_size)]))
+     ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(copy.arg - off, ctx.max_copy_size) - 1), 0,
+     *data64_le(src_addr + off), *data64_le(dst_addr + off)) for off in range(0, copy.arg, ctx.max_copy_size)]))

 def sdma_wait(ctx, dst, val):
  op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \
@@ -200,21 +196,20 @@ def sdma_timestamp(ctx, dst):
  return make_ins(SDMAOps.TIMESTAMP, op, *data64_le(make_getaddr(dst, ctx.devs)))

 pm_sdma_opsel = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy),
-
  (UPat(Ops.BARRIER), lambda: UOp(Ops.NOOP, dtypes.void, ())),
  (UPat(Ops.WAIT, src=(UPat(name="dst"), UPat(name="val"))), sdma_wait),
+  (UPat(Ops.COPY, src=(UPat(name="dst"), UPat(name="src")), name="copy"), sdma_copy),
  (UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", src=(UPat(name="dst"),)), sdma_timestamp),
  (UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), sdma_store),
 ])

 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.src[0].arg // dtypes.uint32.itemsize, UOp.const(dtypes.int, 0)

  # 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 = (make_placeholder(devs, b.size, b.dtype, ("COPY:0", name), unique=False)
+  ring, wptr, doorbell, put_ptr = (UOp.new_buffer(devs, b.size, b.dtype).rtag(("COPY:0", name))
    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))

  # sdma needs the cmdbuf contiguous: if it won't fit before the ring end, restart at 0 and zero the tail
@@ -226,33 +221,22 @@ 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, ptr=True).store(UOp.const(dtypes.uint32, 0)).end(zi)
+  zero_tail = ring.index(tail_off_dw + zi, dtype=ring.dtype.ptr()).store(UOp.const(dtypes.uint32, 0)).end(zi)
  i = UOp.range(UOp.const(dtypes.int, size_dw), 0, dtype=dtypes.int, src=(cmdbuf,))
-  copy_to_ring = ring.index(start_dw + i, ptr=True).store(
-    cmdbuf.index(i*4, ptr=True).cast(dtypes.uint32.ptr()).load()).end(i)
+  copy_to_ring = ring.index(start_dw + i, dtype=ring.dtype.ptr()).store(
+    cmdbuf.index(i*4, dtype=cmdbuf.dtype.ptr()).cast(dtypes.uint32.ptr()).load()).end(i)

  # advance the put/write pointers past the zeroed tail and the cmdbuf
  next_put_b = put_b + ((zero_amt_dw + size_dw) * 4).cast(put_b.dtype)
-  bump_put_ptr = put_ptr.index(zero, ptr=True).store(next_put_b)
-  bump_wptr = wptr.index(zero, ptr=True).store(next_put_b)
+  bump_put_ptr = put_ptr.index(zero, dtype=put_ptr.dtype.ptr()).store(next_put_b)
+  bump_wptr = wptr.index(zero, dtype=wptr.dtype.ptr()).store(next_put_b)

  # ring the doorbell once the writes have landed
  flush = UOp.barrier(zero_tail, copy_to_ring, bump_put_ptr, bump_wptr)
-  return doorbell.after(flush).index(zero, ptr=True).store(next_put_b)
+  return doorbell.after(flush).index(zero, dtype=doorbell.dtype.ptr()).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])))])
-
-@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
-
-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"))
+  lambda lin: sdma_submit(make_cmdbuf(lin, to_tuple(lin.arg[0]), "copy"), to_tuple(lin.arg[0])))])

@dataclass(frozen=True)
 class AMDProgramData:
@@ -261,9 +245,10 @@ class AMDProgramData:
  enable_dispatch_ptr:int; enable_private_segment_sgpr:int

 _amd_program_cache:dict[tuple[bytes,str], tuple[AMDProgramData,bytes]] = {}
+
 def amd_build_program(prg:UOp) -> UOp:
-  dev = Device[to_tuple(prg.device)[0]] # TODO: rm this
-  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, dev.device))) is None:
+  dev = Device[prg.src[1].arg] # TODO: rm this
+  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[4].arg, dev.device))) is None:
    image, sections, relocs = elf_loader(lib)
    rodata = next(sh.header.sh_addr for sh in sections if sh.name == ".rodata")
    for off, sym, typ, addent in relocs:
@@ -273,17 +258,22 @@ def amd_build_program(prg:UOp) -> UOp:
    if (lds:=((desc.group_segment_fixed_size+511)//512)&0x1FF) > (dev.iface.props['lds_size_in_kb']*1024)//512:
      raise RuntimeError("Too many resources requested: group_segment_size")
    edp = desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_DISPATCH_PTR
-
-    data = AMDProgramData(entry_point_offset=rodata + desc.kernel_code_entry_byte_offset,
+    cached = _amd_program_cache[key] = (AMDProgramData(
+      entry_point_offset=rodata + desc.kernel_code_entry_byte_offset,
      rsrc1=desc.compute_pgm_rsrc1 | ((1<<20) if dev.target[0]==11 else 0),  # priv=1 on gfx11 for cwsr
      rsrc2=desc.compute_pgm_rsrc2 | (lds<<15), rsrc3=desc.compute_pgm_rsrc3,
-      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)
-    buf = make_placeholder(prg.device, len(image), dtypes.uint8, "program")
-    cached = _amd_program_cache[key] = prg.replace(src=(buf.after(buf.store(UOp(Ops.BINARY, dtypes.void, src=(), arg=bytes(image)))),), arg=(data, prg.arg))
+      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), bytes(image))
  return cached

+pm_prep_program = PatternMatcher([
+  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE, arg="AMD"), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"), amd_build_program),
+])
+
 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())
@@ -526,15 +516,23 @@ class PCIIface(PCIIfaceBase):

 def _mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface,), {})

+@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
+
+def encode_queue(q:UOp) -> UOp|None:
+  if not (isinstance(q.arg, tuple) and len(q.arg) == 2 and isinstance(q.arg[1], str) and q.arg[1].startswith(("COMPUTE", "COPY"))): return 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"))
+
+pm_lower = PatternMatcher([
+  (UPat(Ops.CUSTOM_FUNCTION, arg="submit", src=(UPat(Ops.LINEAR, name="q"),)), encode_queue),
+])
+
 class AMDDevice(HCQ2Compiled):
-  pm_lower = PatternMatcher([
-    # prep program
-    (UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"), amd_build_program),
-
-    # encoding of cmdbuf
-    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),)), encode_queue),
-  ])
-
  timestamp_divider = 100.0  # AMD GPU clock: ticks/us

  ifaces = [KFDIface, PCIIface]
@@ -581,7 +579,7 @@ class AMDDevice(HCQ2Compiled):

    # Scratch setup
    self.max_private_segment_size = 0
-    self.pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, tag="scratch", name="b"), lambda ctx, b: ctx.scratch_buffer(b.max_numel()))]) + self.pm_bufferize
+    self.pm_bufferize = PatternMatcher([(UPat(Ops.BUFFER, tag="scratch", name="b"), lambda ctx, b: ctx.scratch_buffer(b.arg))]) + self.pm_bufferize

    self.pmc_enabled:bool = PROFILE > 0 and PMC > 0
    if self.pmc_enabled:
@@ -629,10 +627,10 @@ 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={(qname, name)}), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"]
+      (UPat(Ops.BUFFER, tag={(qname, name)}), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"]
    ] + [
-      (UPat(Ops.PARAM, tag={(qname, "timeline_signal")}), lambda ctx, q=qname: ctx.timeline_signal(q)),
-      (UPat(Ops.PARAM, tag={(qname, "timeline_value")}), lambda ctx, q=qname: ctx.timeline_value(q)),
+      (UPat(Ops.BUFFER, tag={(qname, "timeline_signal")}), lambda ctx, q=qname: ctx.timeline_signal(q)),
+      (UPat(Ops.BUFFER, tag={(qname, "timeline_value")}), lambda ctx, q=qname: ctx.timeline_value(q)),
    ]) + self.pm_bufferize

    return queue
@@ -21,7 +21,7 @@ def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_buf:UOp,
                  xw13.base, grad_x2.base, amax_state.base, grad_amax_state.base, threads, workgroups,
                  arg=KernelInfo(f"fused_silu_mul_bwd_w13_{n_elems}", estimates=Estimates(ops=10*n_elems, mem=mem)))
  src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_bwd_w13.cpp", n_elems, hidden)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))

@functools.cache
@@ -34,7 +34,7 @@ def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_buf:UOp, xw13:UOp, amax_state:
  sink = UOp.sink(fp8_out.base, amax_buf.base, xw13.base, amax_state.base, threads, workgroups,
                  arg=KernelInfo(f"fused_silu_mul_cast_amax_w13_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=mem)))
  src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_fwd_w13.cpp", n_elems, hidden)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))

 def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
@@ -0,0 +1,41 @@
+from __future__ import annotations
+import functools, pathlib
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from extra.llama_kernels import THREADS_PER_WG, alloc_like, dname_of, compile_hip
+
+TILE = 64
+
+@functools.cache
+def _custom_fp8_transpose(out:UOp, inp:UOp, dname:str) -> UOp:
+  M, N = inp.shape
+  num_wg = (M // TILE) * (N // TILE)
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(num_wg, "gidx0")
+  mem = M * N * 2  # one byte read + one byte write per element
+  sink = UOp.sink(out.base, inp.base, threads, workgroups,
+                  arg=KernelInfo(f"fp8_transpose_{M}_{N}",
+                                 estimates=Estimates(ops=M*N, mem=mem)))
+  src = (pathlib.Path(__file__).parent/"fp8_transpose.cpp").read_text()
+  defines = [f"-DM_DIM={M}", f"-DN_DIM={N}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
+
+def fast_fp8_transpose(t:Tensor) -> Tensor:
+  assert t.ndim == 2, f"fast_fp8_transpose needs 2D input, got shape {t.shape}"
+  assert t.dtype in dtypes.fp8s, f"fast_fp8_transpose needs fp8 dtype, got {t.dtype}"
+  M, N = t.shape
+  assert M % TILE == 0 and N % TILE == 0, f"M={M}, N={N} must be multiples of {TILE}"
+
+  device = t.device
+  axis = t.uop.axis if isinstance(device, tuple) else None
+  out_axis = None
+  if axis == 0: out_axis = 1
+  elif axis == 1: out_axis = 0
+  elif axis is not None:
+    raise ValueError(f"fast_fp8_transpose: unsupported axis {axis}")
+
+  out = alloc_like((N, M), t.dtype, device, out_axis)
+  fxn = functools.partial(_custom_fp8_transpose, dname=dname_of(device))
+  out, _ = Tensor.custom_kernel(out, t, fxn=fxn)
+  return out
@@ -0,0 +1,74 @@
+#include <hip/hip_runtime.h>
+
+// LDS-staged 64x64 fp8 transpose.
+//   in : (M_DIM, N_DIM) fp8 contiguous
+//   out: (N_DIM, M_DIM) fp8 contiguous, out[c][r] = in[r][c]
+//
+// One WG processes one 64x64 output tile.  Each thread reads one uint4 (16 fp8) coalesced
+// from input rows, stages into LDS, then writes one uint4 coalesced to the output (whose
+// 16 fp8 come from 16 different input rows via in-LDS gather).
+//
+// LDS layout: lds[64][LDS_STRIDE] with LDS_STRIDE=65 (1 byte pad) to mitigate bank conflicts
+// during the column-direction read of the write phase.
+
+#ifndef M_DIM
+#define M_DIM 16384
+#endif
+#ifndef N_DIM
+#define N_DIM 28672
+#endif
+#ifndef THREADS_PER_WG
+#define THREADS_PER_WG 256
+#endif
+
+constexpr int TILE = 64;
+constexpr int VEC = 16;                    // fp8 per uint4 (128-bit) load/store
+constexpr int LDS_PAD = 1;
+constexpr int LDS_STRIDE = TILE + LDS_PAD; // 65 fp8 per row
+
+static_assert(THREADS_PER_WG * VEC == TILE * TILE, "256 threads * 16 fp8 = 64*64");
+static_assert(M_DIM % TILE == 0, "M_DIM must be a multiple of 64");
+static_assert(N_DIM % TILE == 0, "N_DIM must be a multiple of 64");
+
+constexpr int N_TILES_N = N_DIM / TILE;
+
+struct alignas(16) fp8x16 { uint8_t v[16]; };
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
+fp8_transpose(uint8_t* __restrict__ out,         // (N_DIM, M_DIM)
+              const uint8_t* __restrict__ in)    // (M_DIM, N_DIM)
+{
+  __shared__ uint8_t lds[TILE * LDS_STRIDE];
+
+  const int tid     = threadIdx.x;
+  const int wg_id   = blockIdx.x;
+  const int tile_r  = wg_id / N_TILES_N;       // tile index along M dim of input
+  const int tile_c  = wg_id % N_TILES_N;       // tile index along N dim of input
+
+  const int a = tid / (TILE / VEC);            // 0..63 (row within tile during read; col within tile during write)
+  const int b = tid % (TILE / VEC);            // 0..3
+  const int b16 = b * VEC;                     // 0,16,32,48
+
+  // ---- Read phase: input rows -> LDS rows
+  {
+    const long long src = (long long)(tile_r * TILE + a) * (long long)N_DIM
+                        + (long long)(tile_c * TILE + b16);
+    fp8x16 v = *reinterpret_cast<const fp8x16*>(&in[src]);
+    *reinterpret_cast<fp8x16*>(&lds[a * LDS_STRIDE + b16]) = v;
+  }
+  __syncthreads();
+
+  // ---- Write phase: LDS columns (gathered) -> output rows
+  // out[(tile_c*TILE + a)][(tile_r*TILE + b16 + i)] = in[(tile_r*TILE + b16 + i)][(tile_c*TILE + a)]
+  //                                                  = lds[b16 + i][a]
+  {
+    fp8x16 v;
+    #pragma unroll
+    for (int i = 0; i < VEC; ++i) {
+      v.v[i] = lds[(b16 + i) * LDS_STRIDE + a];
+    }
+    const long long dst = (long long)(tile_c * TILE + a) * (long long)M_DIM
+                        + (long long)(tile_r * TILE + b16);
+    *reinterpret_cast<fp8x16*>(&out[dst]) = v;
+  }
+}
@@ -22,7 +22,7 @@ def _custom_fwd(fp8_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_buf:UOp,
  defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={HIDDEN}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}",
             f"-DEPS_LITERAL={eps_val}f"]
  src = _src()
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))

@functools.cache
@@ -39,7 +39,7 @@ def _custom_fwd_add(fp8_out:UOp, h_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax
  defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={HIDDEN}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}",
             f"-DEPS_LITERAL={eps_val}f", f"-DHAS_RESIDUAL=1"]
  src = _src()
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))

@functools.cache
@@ -55,7 +55,7 @@ def _custom_bwd(grad_x:UOp, grad_weight_partial:UOp,
                                 estimates=Estimates(ops=8*n_elems, mem=mem)))
  defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={HIDDEN}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
  src = _src_bwd()
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))

 def _bwd_common(fp8_grad_u, h_grad_u, x_u, x_normed_u, rrms_u, weight_u, amax_state_u, kernel:UOp):
@@ -1,37 +0,0 @@
-from __future__ import annotations
-import functools, pathlib
-from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.renderer import Estimates
-from extra.llama_kernels import THREADS_PER_WG, alloc_like, dname_of, compile_hip
-
-TILE_N = THREADS_PER_WG   # 256
-BLK = 32
-
-@functools.cache
-def _custom_transpose_quantize_mxfp8(q:UOp, e8:UOp, g:UOp, dname:str) -> UOp:
-  M, N = g.shape
-  num_wg = (M // BLK) * (N // TILE_N)
-  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(num_wg, "gidx0")
-  mem = M * N * 2 + M * N + (M // BLK) * N   # read bf16, write fp8 + e8
-  sink = UOp.sink(q.base, e8.base, g.base, threads, workgroups,
-                  arg=KernelInfo(f"transpose_quantize_mxfp8_{M}_{N}", estimates=Estimates(ops=M*N, mem=mem)))
-  src = (pathlib.Path(__file__).parent/"transpose_quantize_mxfp8.cpp").read_text()
-  defines = [f"-DM_DIM={M}", f"-DN_DIM={N}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
-  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, defines))))
-
-def transpose_quantize_mxfp8(g:Tensor) -> tuple[Tensor, Tensor, Tensor]:
-  # fused g.T quantize: returns (q, e8, si) == quantize_mxfp8(g.T) — q (N,M) fp8, e8 (N, M/32), si packed (M/128, N)
-  assert g.ndim == 2 and g.dtype == dtypes.bfloat16, f"{g.shape} {g.dtype}"
-  from extra.gemm.cdna_asm_gemm import FP8_DTYPE, mx_pack
-  M, N = g.shape
-  assert M % BLK == 0 and N % TILE_N == 0, f"M={M} must%{BLK}, N={N} must%{TILE_N}"
-  device = g.device
-  axis = g.uop.axis if isinstance(device, tuple) else None
-  out_axis = None if axis is None else (1 if axis == 0 else 0)
-  q = alloc_like((N, M), FP8_DTYPE, device, out_axis)
-  e8 = alloc_like((N, M // BLK), dtypes.uint8, device, out_axis)
-  fxn = functools.partial(_custom_transpose_quantize_mxfp8, dname=dname_of(device))
-  q, e8, *_ = Tensor.custom_kernel(q, e8, g, fxn=fxn)
-  return q, e8, mx_pack(e8)
@@ -1,62 +0,0 @@
-#include <hip/hip_runtime.h>
-#include <hip/hip_fp8.h>
-#include <hip/hip_bf16.h>
-
-#ifndef M_DIM
-#define M_DIM 8192
-#endif
-#ifndef N_DIM
-#define N_DIM 14336
-#endif
-#ifndef THREADS_PER_WG
-#define THREADS_PER_WG 256
-#endif
-
-constexpr int BLK = 32;
-constexpr int TILE_M = BLK;                 // one mxfp8 block along M per tile
-constexpr int TILE_N = THREADS_PER_WG;      // 256, one output column per thread
-constexpr int LDS_STRIDE = TILE_N + 1;      // +1 pad: stride 257 ≡ 1 (mod 32) -> conflict-free column reads
-constexpr int N_TILES_N = N_DIM / TILE_N;
-constexpr float FP8_MAX = 448.0f;
-
-static_assert(M_DIM % TILE_M == 0, "M_DIM must be a multiple of 32");
-static_assert(N_DIM % TILE_N == 0, "N_DIM must be a multiple of 256");
-
-extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
-transpose_quantize_mxfp8(__hip_fp8_storage_t* __restrict__ q,        // (N_DIM, M_DIM)
-                         uint8_t* __restrict__ e8_out,               // (N_DIM, M_DIM/32)
-                         const __hip_bfloat16* __restrict__ g)       // (M_DIM, N_DIM)
-{
-  __shared__ __hip_bfloat16 lds[TILE_M * LDS_STRIDE];
-  const int tid    = threadIdx.x;
-  const int tile_m = blockIdx.x / N_TILES_N;     // which 32-block along M
-  const int tile_n = blockIdx.x % N_TILES_N;
-
-  #pragma unroll
-  for (int mm = 0; mm < TILE_M; mm++)
-    lds[mm * LDS_STRIDE + tid] = g[(long long)(tile_m * TILE_M + mm) * N_DIM + (tile_n * TILE_N + tid)];
-  __syncthreads();
-
-  float vals[TILE_M];
-  float amax = 0.0f;
-  #pragma unroll
-  for (int mm = 0; mm < TILE_M; mm++) {
-    float v = (float)lds[mm * LDS_STRIDE + tid];
-    vals[mm] = v;
-    amax = fmaxf(amax, fabsf(v));
-  }
-  int e8 = (int)floorf(log2f(fmaxf(amax, 1e-38f))) + 127;
-  e8 = max(0, min(254, e8));
-  float qscale = exp2f((float)(127 - e8));
-
-  const long long n = tile_n * TILE_N + tid;
-  __hip_fp8_storage_t out[TILE_M];
-  #pragma unroll
-  for (int mm = 0; mm < TILE_M; mm++)
-    out[mm] = __hip_cvt_float_to_fp8(fmaxf(-FP8_MAX, fminf(FP8_MAX, vals[mm] * qscale)), __HIP_SATFINITE, __HIP_E4M3);
-  // 32 contiguous fp8 along M -> two 16-byte vector stores
-  long long obase = n * M_DIM + (long long)(tile_m * TILE_M);
-  *reinterpret_cast<uint4*>(&q[obase])      = *reinterpret_cast<uint4*>(&out[0]);
-  *reinterpret_cast<uint4*>(&q[obase + 16]) = *reinterpret_cast<uint4*>(&out[16]);
-  e8_out[n * (M_DIM / BLK) + tile_m] = (uint8_t)e8;
-}
@@ -37,7 +37,7 @@ def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, **kwargs)
    gidx = UOp.special(NUM_WORKGROUPS, "gidx0")
    FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
    sink = UOp.sink(A.base, threads, gidx, arg=KernelInfo(inst.op.name.lower(), estimates=Estimates(ops=FLOPs, mem=0)))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
  dummy = Tensor.zeros(1).contiguous().realize()
  out = Tensor.custom_kernel(dummy, fxn=fxn)[0]
  linear = out.schedule_linear()
@@ -1,6 +1,6 @@
 from tinygrad.tensor import Tensor
 from tinygrad.nn import Conv2d, LayerNorm, LayerNorm2d, Linear
-from tinygrad.helpers import fetch, get_child, Context
+from tinygrad.helpers import fetch, get_child

 class Block:
  def __init__(self, dim):
@@ -58,6 +58,7 @@ if __name__ == "__main__":
  from test.models.test_efficientnet import chicken_img, preprocess, _LABELS
  img = Tensor(preprocess(chicken_img))

-  with Context(TRAINING=0):
-    out = model(img).numpy()
+  Tensor.training = False
+
+  out = model(img).numpy()
  print(_LABELS[out.argmax()])
@@ -5,7 +5,7 @@ import numpy as np
 from pathlib import Path
 from tinygrad import nn, Tensor, dtypes
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.helpers import get_child, fetch, TRAINING
+from tinygrad.helpers import get_child, fetch
 from tinygrad.nn.state import torch_load
 from examples.mlperf.helpers import BoxCoder
 from extra.models.resnet import ResNet
@@ -1069,7 +1069,7 @@ class RoIBoxHead:
  def __call__(self, features, proposals, targets=None):
    x = self.feature_extractor(features, proposals)
    class_logits, box_regression = self.predictor(x)
-    if not TRAINING:
+    if not Tensor.training:
      result = self.post_processor((class_logits, box_regression), proposals)
      return x, result, {}

@@ -1111,7 +1111,7 @@ class Mask:
    x = self.feature_extractor(features, proposals)
    if x is not None:
      mask_logits = self.predictor(x)
-      if not TRAINING:
+      if not Tensor.training:
        result = self.post_processor(mask_logits, proposals)
        return x, result, {}
    return x, [], {}
@@ -1,6 +1,6 @@
 import math
 from tinygrad import Tensor, dtypes
-from tinygrad.helpers import flatten, get_child, TRAINING
+from tinygrad.helpers import flatten, get_child
 from examples.mlperf.helpers import generate_anchors, BoxCoder
 from examples.mlperf.losses import sigmoid_focal_loss, l1_loss
 from extra.models.resnet import ResNet
@@ -141,7 +141,7 @@ class ClassificationHead:
    out = [self.cls_logits(feat.sequential(self.conv)).permute(0, 2, 3, 1).reshape(feat.shape[0], -1, self.num_classes) for feat in x]
    out = out[0].cat(*out[1:], dim=1)

-    if TRAINING:
+    if Tensor.training:
      assert labels is not None and matches is not None, "labels and matches should be passed in when training"
      return self._compute_loss(out.cast(dtypes.float32), labels, matches)

@@ -167,7 +167,7 @@ class RegressionHead:
    out = [self.bbox_reg(feat.sequential(self.conv)).permute(0, 2, 3, 1).reshape(feat.shape[0], -1, 4) for feat in x]
    out = out[0].cat(*out[1:], dim=1)

-    if TRAINING:
+    if Tensor.training:
      assert bboxes is not None and matches is not None and anchors is not None, "bboxes, matches, and anchors should be passed in when training"
      return self._compute_loss(out, bboxes, matches, anchors)

@@ -187,7 +187,7 @@ class RetinaHead:
    self.regression_head = RegressionHead(in_channels, num_anchors)

  def __call__(self, x:Tensor, **kwargs) -> Tensor|dict[str, Tensor]:
-    if TRAINING:
+    if Tensor.training:
      return {
        "classification_loss": self.classification_head(x, labels=kwargs["labels"], matches=kwargs["matches"]),
        "regression_loss": self.regression_head(x, bboxes=kwargs["bboxes"], matches=kwargs["matches"], anchors=kwargs["anchors"])
@@ -111,7 +111,7 @@ def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, device:str, arch:st
  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)))
+             src=(sink, UOp(Ops.DEVICE, arg=device), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))

@functools.cache
 def custom_fa_backward_pre(delta_vec:UOp, dq:UOp, o:UOp, do:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
@@ -141,7 +141,7 @@ def custom_fa_backward_pre(delta_vec:UOp, dq:UOp, o:UOp, do:UOp, device:str, arc
  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)))
+             src=(sink, UOp(Ops.DEVICE, arg=device), 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):
@@ -171,7 +171,7 @@ def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_ve
  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)))
+             src=(sink, UOp(Ops.DEVICE, arg=device), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))

@functools.cache
 def custom_fa_backward_post(dq_out:UOp, dq_in:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
@@ -201,4 +201,4 @@ def custom_fa_backward_post(dq_out:UOp, dq_in:UOp, device:str, arch:str, B:int,
  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)))
+             src=(sink, UOp(Ops.DEVICE, arg=device), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
@@ -93,7 +93,7 @@ class Group:
          d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])

          out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
-          c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
+          c_i = [c[height, width, i].store(out.gep(i)) for i in range(4)]
          c_store = UOp.group(*c_i).end(height, width, inner)

    self.ker.push_store(c_store, c)
@@ -123,7 +123,7 @@ class Group:
          d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])

          out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
-          c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
+          c_i = [c[height, width, i].store(out.gep(i)) for i in range(4)]
          c_store = UOp.group(*c_i).end(height, width, inner)

    self.ker.push_store(c_store, c)
@@ -153,7 +153,7 @@ class Group:
          d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])

          out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
-          c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
+          c_i = [c[height, width, i].store(out.gep(i)) for i in range(4)]
          c_store = UOp.group(*c_i).end(height, width, inner)

    self.ker.push_store(c_store, c)
@@ -183,7 +183,7 @@ class Group:
          d_in = UOp.vectorize(*[c[height, width, i] for i in range(4)])

          out = UOp(Ops.WMMA, dtypes.float32.vec(4), (a_in, b_in, d_in), arg=wmma_arg)
-          c_i = [c[height, width, i].store(out.index(i)) for i in range(4)]
+          c_i = [c[height, width, i].store(out.gep(i)) for i in range(4)]
          c_store = UOp.group(*c_i).end(height, width, inner)

    self.ker.push_store(c_store, c)
@@ -1,6 +1,4 @@
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import CHUNK_SIZE
-from tinygrad.nn.state import fs_load
 import argparse, math, hashlib

 def _python_hash_1mb(data:bytes|bytearray):
@@ -9,15 +7,15 @@ def _python_hash_1mb(data:bytes|bytearray):
  return hashlib.shake_128(b''.join(chunk_hashes)).digest(16)

 def hash_file(data: bytes|bytearray):
-  if len(data) % CHUNK_SIZE != 0: data += bytes(CHUNK_SIZE - len(data) % CHUNK_SIZE)
-  base_chunks = math.ceil(len(data) / CHUNK_SIZE)
-  tree_depth = math.ceil(math.log(base_chunks, CHUNK_SIZE // 16))
+  if len(data) % Tensor.CHUNK_SIZE != 0: data += bytes(Tensor.CHUNK_SIZE - len(data) % Tensor.CHUNK_SIZE)
+  base_chunks = math.ceil(len(data) / Tensor.CHUNK_SIZE)
+  tree_depth = math.ceil(math.log(base_chunks, Tensor.CHUNK_SIZE // 16))

  for _ in range(tree_depth + 1):
-    data_chunks = [data[i:i+CHUNK_SIZE] for i in range(0, len(data), CHUNK_SIZE)]
+    data_chunks = [data[i:i+Tensor.CHUNK_SIZE] for i in range(0, len(data), Tensor.CHUNK_SIZE)]
    data_chunk_hashes = [_python_hash_1mb(chunk) for chunk in data_chunks]
    data = b''.join(data_chunk_hashes)
-    if len(data) % CHUNK_SIZE != 0: data += bytes(CHUNK_SIZE - len(data) % CHUNK_SIZE)
+    if len(data) % Tensor.CHUNK_SIZE != 0: data += bytes(Tensor.CHUNK_SIZE - len(data) % Tensor.CHUNK_SIZE)

  return data[:16]

@@ -29,7 +27,7 @@ if __name__ == "__main__":
  parser.add_argument("--check", action="store_true", help="verify the file hash after fetching")
  args = parser.parse_args()

-  fs_load(Tensor(bytes.fromhex(args.hash), device="CPU"), args.len).to(f"disk:{args.dest}").realize()
+  Tensor(bytes.fromhex(args.hash), device="CPU").fs_load(args.len).to(f"disk:{args.dest}").realize()

  if args.check:
    with open(args.dest, "rb") as f:
@@ -3,7 +3,6 @@ from pathlib import Path

 from tinygrad.tensor import Tensor
 from tinygrad.helpers import tqdm, getenv
-from tinygrad.nn.state import fs_load

 raid_root = Path(getenv("RAID_ROOT", "/raid"))

@@ -15,7 +14,7 @@ def fetch_file(item):
  path.parent.mkdir(parents=True, exist_ok=True)

  try:
-    pt = fs_load(Tensor(bytes.fromhex(h), device="CPU"), size).to(f"disk:{path.as_posix()}").realize()
+    pt = Tensor(bytes.fromhex(h), device="CPU").fs_load(size).to(f"disk:{path.as_posix()}").realize()
  except Exception as e:
    print(f"error fetching {path}, {h}, {size}: {e}")
    raise
@@ -23,7 +22,7 @@ def fetch_file(item):
  pt.uop.buffer.deallocate()

 def fetch_mapping(h, l):
-  mapping_tensor = fs_load(Tensor(bytes.fromhex(h)), l).realize()
+  mapping_tensor = Tensor(bytes.fromhex(h)).fs_load(l).realize()
  mapping = mapping_tensor.data().tobytes().decode()
  mapping = json.loads(mapping)
  mapped_files = mapping.items()
@@ -3,13 +3,12 @@ import multiprocessing, json

 from tinygrad.tensor import Tensor
 from tinygrad.helpers import tqdm
-from tinygrad.nn.state import fs_store

 raid_root = Path("/raid")

 def upload_file(path: Path):
  pt = Tensor(path).realize()
-  h = fs_store(pt).realize()
+  h = pt.fs_store().realize()
  pt.uop.realized.deallocate()
  return h.data().hex(), path, pt.nbytes()

@@ -27,6 +26,6 @@ if __name__ == "__main__":

  mapping = json.dumps(mapping).encode()
  mapping_tensor = Tensor(mapping, device="CPU")
-  h = fs_store(mapping_tensor).realize()
+  h = mapping_tensor.fs_store().realize()

  print(f"final hash: {h.data().hex()}, size: {len(mapping)}")
@@ -1,14 +1,10 @@
 import unittest
 import torch
 import tinygrad.nn.torch
+torch.set_default_device("tiny")
 import numpy as np

 class TestTorchBackendInplace(unittest.TestCase):
-  def setUp(self):
-    self._prev_device = torch.get_default_device()
-    torch.set_default_device("tiny")
-  def tearDown(self): torch.set_default_device(self._prev_device)
-
  def test_zero(self):
    a = torch.ones(4)
    a.zero_()
@@ -2,15 +2,11 @@ import unittest
 from tinygrad.helpers import getenv
 import torch
 import tinygrad.nn.torch
+torch.set_default_device("tiny")
 import numpy as np

@unittest.skipIf(getenv("GPUS",1)<=1, "only single GPU")
 class TestTorchBackendMultiGPU(unittest.TestCase):
-  def setUp(self):
-    self._prev_device = torch.get_default_device()
-    torch.set_default_device("tiny")
-  def tearDown(self): torch.set_default_device(self._prev_device)
-
  def test_transfer(self):
    a = torch.Tensor([[1,2],[3,4]]).to("tiny:0")
    b = torch.Tensor([[3,2],[1,0]]).to("tiny:1")
@@ -40,7 +40,7 @@ def train(model, X_train, Y_train, optim, steps, BS=128, lossfn=lambda out,y: ou

 def evaluate(model, X_test, Y_test, num_classes=None, BS=128, return_predict=False, transform=lambda x: x,
             target_transform=lambda y: y):
-  @Context(TRAINING=0)
+  Tensor.training = False
  def numpy_eval(Y_test, num_classes):
    Y_test_preds_out = np.zeros(list(Y_test.shape)+[num_classes])
    for i in trange((len(Y_test)-1)//BS+1, disable=None):
@@ -20,7 +20,6 @@ include-package-data = true
 packages = [
  'tinygrad',
  'tinygrad.codegen',
-  'tinygrad.codegen.decomp',
  'tinygrad.codegen.opt',
  'tinygrad.codegen.late',
  'tinygrad.engine',
@@ -111,7 +110,6 @@ docs = [
  "black",
  "numpy",
 ]
-mesa = ["tinymesa==25.2.7.2"]


 [tool.mutmut]
@@ -227,7 +227,7 @@ Ternary & $(P, A, B)$
 \midrule
 \op{Barrier} & (deps\ldots) & --- & Synchronize threads within a workgroup. \\
 \op{Ins} & \ldots & \ldots & A single machine instruction (e.g.\ AMD ISA). \\
-\op{GetAddr} & (buf,) & dev & Lower buf to its address on device dev. \\
+\op{GetAddr} & (buf, dev) & --- & Lower buf to its address on device dev. \\
 \op{Special} & (bound,) & name & GPU thread/workgroup index (e.g.\ \texttt{gidx0}, \texttt{lidx1}). \\
 \op{If} & (gate,) & --- & Begin conditional execution block. \\
 \op{Endif} & (if,) & --- & End conditional execution block. \\
@@ -30,7 +30,7 @@ def custom_add_one(A:UOp) -> UOp:
    s_endpgm(),
  ]
  sink = UOp.sink(A.base, threads, arg=KernelInfo(f"custom_add_one_{A.numel()}", estimates=Estimates(ops=A.numel(), mem=A.numel()*4*2)))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

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

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

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

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

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

@unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
 class TestAsmKernel(unittest.TestCase):
@@ -89,7 +89,7 @@ class TestTinygradIntegration(unittest.TestCase):
    assert len(sink_items) > 0, "No SINK in schedule"
    renderer = AMDLLVMRenderer(Target("AMD", arch='gfx1100'))
    prg = to_program(sink_items[0].src[0], renderer)
-    lib = renderer.compiler.compile(prg.src[2].arg)
+    lib = renderer.compiler.compile(prg.src[3].arg)
    return next(s.content for s in elf_loader(lib)[1] if s.name == ".text")

  def test_simple_add_kernel(self):
@@ -15,7 +15,7 @@ dev = Device["AMD"]
 a = Tensor([1.0]).realize()
 b = a + 1
 linear = compile_linear(b.schedule_linear())
-lib = bytearray(linear.src[-1].src[0].src[3].arg)
+lib = bytearray(linear.src[-1].src[0].src[4].arg)

 # Find s_endpgm (0xBFB00000) and replace with V_MOVRELD_B32 (op=66) which has no pcode
 # VOP1 encoding: bits[31:25]=0x7E, op=bits[16:9], so op=66 -> 66<<9 = 0x8400
@@ -83,8 +83,8 @@ def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelSnapshot], dict[int, in
            buf_data[dst_id] = src_data
      elif ast.op is Ops.PROGRAM:
        info = ast.arg
-        if len(ast.src) > 3 and ast.src[3].op is Ops.BINARY:
-          lib = bytes(ast.src[3].arg)
+        if len(ast.src) > 4 and ast.src[4].op is Ops.BINARY:
+          lib = bytes(ast.src[4].arg)
          _, sections, _ = elf_loader(lib)
          for sec in sections:
            if sec.name == '.text':
@@ -98,7 +98,7 @@ def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelSnapshot], dict[int, in
                buf_sizes.append(b.nbytes)
              kernels.append(KernelSnapshot(
                code=bytes(sec.content),
-                src=ast.src[2].arg,
+                src=ast.src[3].arg,
                global_size=tuple(info.global_size),
                local_size=tuple(info.local_size),
                buf_idxs=buf_idxs,
@@ -187,26 +187,6 @@ class TestIndexing(unittest.TestCase):
    for i in idx.flatten().numpy(): expected_grad[i] += 2
    np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)

-  @unittest.skipIf(Device.DEFAULT not in ("CPU", "AMD"), "atomics only on AMD/CPU")
-  @Context(USE_ATOMICS=1, SPEC=1)
-  def test_embedding_backward_padded_embed(self):
-    from tinygrad.renderer.cstyle import CStyleLanguage
-    if Device.DEFAULT == "CPU" and not isinstance(Device["CPU"].renderer, CStyleLanguage): self.skipTest("CPU needs Clang renderer")
-    vocab_size, embed_size = 1000, 300
-    bs, seqlen = 4, 256
-    idx = Tensor.randint(bs, seqlen, high=vocab_size)
-    emb = nn.Embedding(vocab_size, embed_size)
-    emb.weight = Tensor.ones(vocab_size, embed_size)
-    gt = Tensor.zeros(bs, seqlen, embed_size)
-    Tensor.realize(idx, emb.weight, gt)
-    loss = (emb(idx)-gt).square().sum()
-    loss.backward()
-    emb.weight.grad.realize()
-    # correctness check
-    expected_grad = np.zeros((vocab_size, embed_size), dtype=np.float32)
-    for i in idx.flatten().numpy(): expected_grad[i] += 2
-    np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)
-
  @needs_second_gpu
  @unittest.skipIf(Device.DEFAULT not in ("CPU", "AMD"), "atomics only on AMD/CPU")
  @Context(USE_ATOMICS=1, SPEC=1)
@@ -9,7 +9,7 @@ from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8, FP8_MAX
 # Use DEV=NULL:HIP:gfx950 to also test the assembly
 def is_cdna4(): return Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950")

-def run_asm_gemm(a_shape, b_shape, dtype=dtypes.bfloat16, a_shard=None, b_shard=None, gpus:int=1) -> None:
+def run_asm_gemm(a_shape, b_shape, dtype=dtypes.float16, a_shard=None, b_shard=None, gpus:int=1) -> None:
  Tensor.manual_seed(0)
  input_dtype = dtypes.bfloat16 if dtype == FP8_DTYPE else dtype
  a_rand = Tensor.randn(a_shape, dtype=dtypes.float).sub(0.5).cast(input_dtype)
@@ -64,31 +64,31 @@ def run_asm_gemm(a_shape, b_shape, dtype=dtypes.bfloat16, a_shard=None, b_shard=
    assert a.grad.allclose(a_ref.grad, atol=grad_atol, rtol=grad_rtol).item(), "grad_a mismatch"
    assert b.grad.allclose(b_ref.grad, atol=grad_atol, rtol=grad_rtol).item(), "grad_b mismatch"

-def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.bfloat16, gpus:int=1) -> None:
+def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=1) -> None:
  run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=0, b_shard=None, gpus=gpus)

-def verify_asm_gemm_k_sharded(M:int, N:int, K:int, dtype=dtypes.bfloat16, gpus:int=8) -> None:
+def verify_asm_gemm_k_sharded(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=8) -> None:
  run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=1, b_shard=0, gpus=gpus)

-def verify_asm_gemm_n_sharded(batch:int, M:int, N:int, K:int, dtype=dtypes.bfloat16, gpus:int=2) -> None:
+def verify_asm_gemm_n_sharded(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
  run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=None, b_shard=1, gpus=gpus)

-def verify_asm_gemm_m_sharded(M:int, N:int, K:int, dtype=dtypes.bfloat16, gpus:int=2) -> None:
+def verify_asm_gemm_m_sharded(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
  run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=0, b_shard=None, gpus=gpus)

-def verify_asm_gemm_n_sharded_2d(M:int, N:int, K:int, dtype=dtypes.bfloat16, gpus:int=2) -> None:
+def verify_asm_gemm_n_sharded_2d(M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
  run_asm_gemm((M, K), (K, N), dtype=dtype, a_shard=None, b_shard=1, gpus=gpus)

-def verify_asm_gemm_k_sharded_3d(batch:int, M:int, N:int, K:int, dtype=dtypes.bfloat16, gpus:int=2) -> None:
+def verify_asm_gemm_k_sharded_3d(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=2) -> None:
  run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=2, b_shard=0, gpus=gpus)

 # 128x smaller than usual
 # uses the UOp GEMM, runs on non CDNA4 and CI
-@unittest.skipUnless(dtypes.bfloat16 in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half")
+@unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half")
 class TestGemm(unittest.TestCase):
  def setUp(self):
    if is_cdna4(): self.skipTest("shapes are too small for the assembly GEMM")
-  def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 32), N, N, dtype=dtypes.bfloat16)
+  def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 32), N, N, dtype=dtypes.half)
  def test_gemm(self): verify_asm_gemm(1, 64, 32, 112)
  def test_gemm_batched(self): verify_asm_gemm(2, 64, 32, 32)
  @needs_second_gpu
@@ -107,7 +107,7 @@ class TestGemm(unittest.TestCase):
 # uses the smallest size for the cdna assembly gemm
 class TestAsmGEMM(unittest.TestCase):
  def setUp(self):
-    if not is_cdna4() or not has_hipcc():
+    if not is_cdna4():
      self.skipTest("assembly gemm is only for cdna4")

  def test_tiny(self): verify_asm_gemm(1, 256, 256, 64)
@@ -145,7 +145,7 @@ class TestGemmLlama(unittest.TestCase):
  dtype = dtypes.bfloat16

  def setUp(self):
-    if not is_cdna4() or DEV.interface.startswith("MOCK") or not has_hipcc():
+    if not is_cdna4() or DEV.interface.startswith("MOCK"):
      self.skipTest("very slow on non mi350x")

  def test_empty(self): asm_gemm(Tensor.empty(N:=getenv("N", 4096), N, dtype=self.dtype), Tensor.empty(N, N, dtype=self.dtype)).realize()
@@ -380,5 +380,23 @@ class TestHkBf16AtbGemm(unittest.TestCase):
  @needs_second_gpu
  def test_m_sharded(self): run_atb_gemm(256, 512, 256, a_shard=2, b_shard=None, gpus=2)

+class TestMagicGu(unittest.TestCase):
+  def test_magicgu_matches_old(self):
+    from extra.gemm.cdna_asm_gemm import _magicgu_mulhi, TILE_M, TILE_N, TILE_K
+    old_iters_args = {64: (67108864, 0), 128: (33554432, 0), 224: (613566757, 2147483656)}
+    old_gemm_shapes = [
+      (8192, 4096, 4096), (8192, 14336, 4096), (8192, 4096, 14336),
+      (8192, 8192, 8192), (4096, 4096, 4096), (4096, 14336, 4096),
+      (4096, 14336, 8192), (4096, 4096, 14336), (14336, 4096, 8192),
+      (4096, 8192, 14336), (4096, 4096, 8192), (4096, 8192, 4096),
+    ]
+    for M, N, K in old_gemm_shapes:
+      iters = K // TILE_K
+      total = (M // TILE_M) * (N // TILE_N) * iters
+      for batch in [1, 2]:
+        magic, shift = _magicgu_mulhi(iters, total * batch)
+        old_magic, old_shift = old_iters_args[iters]
+        self.assertEqual((magic, shift), (old_magic, old_shift), f"mismatch for ({M},{N},{K}) batch={batch} iters={iters}")
+
 if __name__ == "__main__":
  unittest.main()
@@ -2,7 +2,6 @@ import unittest, math
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.dtype import DTYPES_DICT
 from tinygrad.uop.ops import Ops, UOp
-from tinygrad.codegen.decomp.op import threefry2x32
 import numpy as np
 from test.helpers import not_support_multi_device

@@ -168,8 +167,7 @@ class TestMultiConstFolding(unittest.TestCase):

 class TestThreefryConstFolding(unittest.TestCase):
  def test_threefry(self):
-    # THREEFRY(const,const) folds to a const once decomposed
-    x = threefry2x32(UOp.const(dtypes.uint64, 5), UOp.const(dtypes.uint64, 10))
+    x = UOp.const(dtypes.uint64, 5).threefry(UOp.const(dtypes.uint64, 10))
    self.assertIs(x.simplify().op, Ops.CONST)

 class TestTautologicalCompare(unittest.TestCase):
@@ -51,7 +51,7 @@ def flip_contract_kernel(dest:UOp, src:UOp):
  i = UOp.range(dest.shape[0], 0)
  j = UOp.range(dest.shape[1], 1, AxisType.UPCAST)
  vec = src[i, j].contract(j)
-  store = UOp.group(*[dest[i, k].store(vec.index(3-k)) for k in range(4)])
+  store = UOp.group(*[dest[i, k].store(vec.gep(3-k)) for k in range(4)])
  return store.end(i, j).sink(arg=KernelInfo(name=f"flip_contract_{dest.numel()}", opts_to_apply=()))

 def slice_sum_kernel(dest:UOp, src:UOp):
@@ -430,7 +430,7 @@ class TestCustomKernel(unittest.TestCase):
    binary = Device[a.device].renderer.compiler.compile(src)
    def custom_src_kernel(A:UOp) -> UOp:
      sink = UOp.sink(A, arg=KernelInfo(name="test_src"))
-      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())),
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="CPU"), UOp(Ops.LINEAR, src=tuple(sink.toposort())),
                                   UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))

    a = Tensor.custom_kernel(a, fxn=custom_src_kernel)[0]
@@ -47,7 +47,7 @@ class TestLinearizer(unittest.TestCase):
    tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
    out = tst.neg().cast(dtypes.char).cast(dtypes.int).cast(dtypes.char) * 2
    ast = helper_linearizer_opt(out)
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)

  @unittest.expectedFailure
@@ -55,7 +55,7 @@ class TestLinearizer(unittest.TestCase):
    tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
    out = tst.neg().cast(dtypes.char).cast(dtypes.int) * 2
    ast = helper_linearizer_opt(out)
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)

  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx")
@@ -65,7 +65,7 @@ class TestLinearizer(unittest.TestCase):
    b = Tensor.empty(16)
    out = img.conv2d(w, b)
    ast = helper_linearizer_opt(out)
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
    # slice at the last loop end
    uslice = [i for i,u in enumerate(uops) if u.op == Ops.END][-1]
    # only valid test if outermost range is the reduce
@@ -86,7 +86,7 @@ class TestLinearizer(unittest.TestCase):
    a = Tensor.randn(2, ).realize()
    out = a.reshape(2, 1).expand(2, 3).sum()
    ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)).sum()])
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    assert len(ranges) == 1 # NOTE: it collapses now

@@ -94,7 +94,7 @@ class TestLinearizer(unittest.TestCase):
    a = Tensor.randn(2, ).realize()
    out = a.reshape(2, 1).expand(2, 3).expand(2, 2, 3).sum()
    ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)), (2, 2, 3)).sum()])
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    assert len(ranges) == 1 # NOTE: it collapses now

@@ -102,7 +102,7 @@ class TestLinearizer(unittest.TestCase):
    a = Tensor([2, 2]).realize()
    out = a.reshape(2, 1).pad(((1, 1), (1, 1)), value=2).sum()
    ast = helper_linearizer_opt(out, wanna_output=[24])
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    # RANGE -> ALU -> RANGE -> ALU + LOAD -> STORE
    assert any(x.op in GroupOp.ALU for x in uops[ranges[0]:ranges[1]])
@@ -115,7 +115,7 @@ class TestLinearizer(unittest.TestCase):
    b = Tensor.randn(1, 1).realize()
    out = (a + b[0]).sum() + b[0]
    ast = helper_linearizer_opt(out, wanna_output=[(a.numpy()+b.numpy()[0]).sum()+b.numpy()])
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    # LOAD -> RANGE -> LOAD -> STORE
    assert len([x for x in uops[:ranges[0]] if x.op is Ops.LOAD]) == 1
@@ -125,7 +125,7 @@ class TestLinearizer(unittest.TestCase):
    b = Tensor.randn(1, 1).realize()
    out = (a.reshape(2, 1).expand(2, 3) + b[0]).sum() + b[0]
    ast = helper_linearizer_opt(out, wanna_output=[(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)) + b.numpy()[0]).sum() + b.numpy()])
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    assert len(ranges) == 1 # NOTE: it collapses now

@@ -137,7 +137,7 @@ class TestLinearizer(unittest.TestCase):
    r = a[:-1] + a[1:]

    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
-                       renderer=Device[Device.DEFAULT].renderer).src[1].src)
+                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
    num_loads = len([uop for uop in uops if uop.op is Ops.LOAD])
    assert num_loads <= 4, "more load uops than needed"
    assert num_loads >= 1, "expected at least one load uop"
@@ -150,7 +150,7 @@ class TestLinearizer(unittest.TestCase):
    r = a.expand([2]) + b.expand([2])

    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
-                       renderer=Device[Device.DEFAULT].renderer).src[1].src)
+                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
    num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
    assert num_ops <= 1, "more alu uops than needed"

@@ -160,7 +160,7 @@ class TestLinearizer(unittest.TestCase):
    r = Tensor.conv2d(x,w,padding=1).relu()

    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0],
-      [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).src[1].src)
+      [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).src[2].src)
    accs = [u for u in uops if u.op is Ops.BUFFER and u.addrspace is AddrSpace.REG]
    stores = [u for u in uops if u.op is Ops.STORE]
    assert len(accs) == 0  # it's removed now
@@ -173,7 +173,7 @@ class TestLinearizer(unittest.TestCase):
    out = Tensor.ones(64,64).contiguous() @ Tensor.ones(64,64).contiguous()
    prg = to_program(replace_opts(out.schedule_linear().src[-1].src[0], [Opt(OptOps.LOCAL, axis=0, arg=4)]),
                      renderer=Device[Device.DEFAULT].renderer)
-    self.assertEqual(len(prg.src[2].arg.split("for")), 5)
+    self.assertEqual(len(prg.src[3].arg.split("for")), 5)

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@@ -185,7 +185,7 @@ class TestLinearizer(unittest.TestCase):
    opts_to_apply = [Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
    program = to_program(replace_opts(r.schedule_linear().src[-1].src[0], opts_to_apply), renderer=Device[Device.DEFAULT].renderer)

-    stores = [u for u in tuple(program.src[1].src) if u.op is Ops.STORE and u.src[0].addrspace != AddrSpace.REG]
+    stores = [u for u in tuple(program.src[2].src) if u.op is Ops.STORE and u.src[0].addrspace != AddrSpace.REG]

    # the first store is to lds and can be upcasted
    assert stores[0].src[1].max_numel() == 4
@@ -199,7 +199,7 @@ class TestLinearizer(unittest.TestCase):
    a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
    r = Tensor.stack(a, b)
    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
-                       renderer=Device[Device.DEFAULT].renderer).src[1].src)
+                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
    num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
    assert num_ops == 0, "more alu uops than needed"

@@ -210,14 +210,14 @@ class TestLinearizer(unittest.TestCase):
        a = Tensor([1, 2, 3], dtype=tensor_dtype).sum()
        realized_ast = a.schedule_linear().src[-1].src[0]
        program = to_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
-        local = [uop for uop in tuple(program.src[1].src) if uop.op is Ops.BUFFER and uop.addrspace in (AddrSpace.LOCAL, AddrSpace.REG)]
+        local = [uop for uop in tuple(program.src[2].src) if uop.op is Ops.BUFFER and uop.addrspace in (AddrSpace.LOCAL, AddrSpace.REG)]
        assert local[0].dtype.base == acc_dtype

  def test_arg_acc_dtype(self):
    def helper_arg_acc_dtype(c: Tensor, expected_dtype:DType):
      realized_ast = c.schedule_linear().src[-1].src[0]
      program = to_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
-      local = [uop for uop in tuple(program.src[1].src) if uop.op is Ops.BUFFER and uop.addrspace in (AddrSpace.LOCAL, AddrSpace.REG)]
+      local = [uop for uop in tuple(program.src[2].src) if uop.op is Ops.BUFFER and uop.addrspace in (AddrSpace.LOCAL, AddrSpace.REG)]
      self.assertEqual(local[0].dtype.base, expected_dtype)

    tests = (
@@ -244,7 +244,7 @@ class TestLinearizer(unittest.TestCase):
    opt = [Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4)]
    ast = helper_linearizer_opt(r, [opt])
    # the uops graph is reg BUFFER -> 4x STORE 0.0 -> RANGE -> 4x ALU -> 4x STORE -> ENDRANGE
-    uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[1].src)
+    uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[2].src)
    begin_range = [i for i, x in enumerate(uops) if x.op is Ops.RANGE][-1]
    end_range = [i for i, x in enumerate(uops) if x.op is Ops.END][0]
    for i,u in enumerate(uops): print(i, u.op, [uops.index(s) for s in u.src], u.arg, u.dtype)
@@ -264,7 +264,7 @@ class TestLinearizer(unittest.TestCase):
    # shrink so that the dims do not collapse
    t = Tensor.ones(5, 6, 7).contiguous().realize().shrink(((0, 4), (0, 5), (0, 6)))
    ast = helper_linearizer_opt(t+1)
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
+    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
    idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
    idxs = sorted(idxs, key=lambda uop: uop.arg)
    assert (idxs[0].arg, idxs[0].src[0].arg) == ('gidx0', 6), idxs[0]
@@ -297,12 +297,12 @@ class TestLinearizer(unittest.TestCase):
    run_linear(linear, var_vals)
    np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.])
    program = to_program(replace_opts(linear.src[-1].src[0], []), renderer=Device[Device.DEFAULT].renderer)
-    assert not any(u.op == Ops.WHERE for u in tuple(program.src[1].src)), "found where where where should be folded"
+    assert not any(u.op == Ops.WHERE for u in tuple(program.src[2].src)), "found where where where should be folded"

  def test_phi_simplification(self):
    def helper(t, max_ops=0):
      ast = helper_linearizer_opt(t)
-      uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[1].src)
+      uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src)
      # ignore kernel optimized IF statements for now
      if if_op:=next((u for u in uops if u.op is Ops.IF), None):
        uops = uops[:uops.index(if_op)]
@@ -335,7 +335,7 @@ class TestLinearizer(unittest.TestCase):
    out = x.matmul(y)
    with Context(TC=0):
      ast = helper_linearizer_opt(out)
-      uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[1].src)
+      uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src)
    # check that the float4 cast collapses
    store_vals = [u.src[1] for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
    for val in store_vals:
@@ -347,7 +347,7 @@ class TestLinearizer(unittest.TestCase):
    x = Tensor.randn((4,3,6,6)).realize()
    out = x.flip((0,1)).contiguous()
    ast = helper_linearizer_opt(out)
-    store_val = [u.src[1] for u in tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[1].src) if u.op is Ops.STORE][0]
+    store_val = [u.src[1] for u in tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src) if u.op is Ops.STORE][0]
    assert store_val.dtype == dtypes.float.vec(4) and store_val.op is not Ops.STACK

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@@ -360,14 +360,15 @@ class TestLinearizer(unittest.TestCase):
            Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 2)] # upcast accs in both reduces
    ast = helper_linearizer_opt(out, opts=[opt])
    def get_recursive(uop): return set.union(set(uop.src), [uop], *[get_recursive(v) for v in uop.src])
-    uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[1].src)
-    local_stores = [u for u in uops if u.op is Ops.STORE and u.addrspace == AddrSpace.LOCAL]
-    global_stores = [u for u in uops if u.op is Ops.STORE and u.addrspace == AddrSpace.GLOBAL]
+    uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    local_stores = [u for u in uops if u.op is Ops.STORE and any(
+      x.op is Ops.BUFFER and x.addrspace is AddrSpace.LOCAL for x in get_recursive(u.src[0]))]
+    global_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.PARAM for x in get_recursive(u.src[0]))]
    barrier = [u for u in uops if u.op is Ops.BARRIER]
    assert len(barrier) == 1
    # check that the float4 cast collapses for all stores
    for store in local_stores+global_stores:
-      assert store.src[1].max_numel() > 1, f"store shape {store.src[1].shape} on {store.addrspace}"
+      assert store.src[1].max_numel() > 1 # and store.src[2].op is not Ops.VECTORIZE
    # # check the children's vins
    # TODO: src ALU are not the same, should it?
    # assert barrier.src == tuple(local_stores)
@@ -382,7 +383,7 @@ class TestLinearizer(unittest.TestCase):
    x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
    r = (x@y).relu()
    ast = helper_linearizer_opt(r)
-    uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[1].src)
+    uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src)
    stores = [u for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]

    # the float4 value stores directly in lds and we skip upcast
@@ -89,7 +89,7 @@ class TestLocalAmax(unittest.TestCase):
    x = Tensor.arange(16).reshape(4, 4).cast(dtypes.float).clone(devices[0]).realize().shard(devices, axis=0).realize()
    GlobalCounters.reset()
    out = (x * local_abs_max(x)).clone().realize()
-    self.assertEqual(GlobalCounters.kernel_count, 2)
+    self.assertEqual(GlobalCounters.kernel_count, 4)
    self.assertEqual(out.tolist(), [[0., 7., 14., 21.], [28., 35., 42., 49.], [120., 135., 150., 165.], [180., 195., 210., 225.]])

 if __name__ == '__main__':
@@ -1449,7 +1449,6 @@ class TestOps(unittest.TestCase):
  def test_small_gemm_eye(self):
    helper_test_op(None, lambda x,y: x.matmul(y), lambda x,y: x@y, vals=[np.eye(8).astype(np.float32), np.eye(8).astype(np.float32)])
  @unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "not precise enough when emulating")
-  @unittest.skipIf(IMAGE>0, "image does math in float32")
  def test_gemm_fp16(self):
    helper_test_op([(64,64), (64,64)], lambda x,y: x.half().matmul(y.half()), atol=5e-3, rtol=5e-3, grad_atol=5e-3, grad_rtol=5e-3)
  def test_gemm(self):
@@ -1602,12 +1601,6 @@ class TestOps(unittest.TestCase):
        helper_test_op(None, lambda x,y: x.isclose(y), vals=[[a], [b]], forward_only=True)
        helper_test_op(None, lambda x,y: x.isclose(y, equal_nan=True), vals=[[a], [b]], forward_only=True)

-  def test_isclose_scalar(self):
-    # torch needs a tensor
-    helper_test_op([(3, 4, 5, 6)], lambda x: x.isclose(torch.tensor(1.0)), lambda x: x.isclose(1.0), forward_only=True)
-    helper_test_op(None, lambda x: x.isclose(torch.tensor(1.0)), lambda x: x.isclose(1.0),
-                   vals=[[1.0, 1.0 + 1e-7, 2.0, math.inf, -math.inf, math.nan]], forward_only=True)
-
  def test_mean(self):
    helper_test_op([(3,4,5,6)], lambda x: x.mean())
    helper_test_op([()], lambda x: x.mean())
@@ -3,7 +3,6 @@ import torch
 import unittest
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.nn.optim import Adam, SGD, AdamW, Muon, LAMB
-from tinygrad.helpers import Context
 from test.helpers import needs_second_gpu, slow

 np.random.seed(1337)
@@ -46,7 +45,11 @@ def step(tensor, optim, steps=1, teeny=False, **kwargs):

@slow
 class TestOptim(unittest.TestCase):
-  def setUp(self): self.enterContext(Context(TRAINING=1))
+  def setUp(self):
+    self.old_training = Tensor.training
+    Tensor.training = True
+  def tearDown(self):
+    Tensor.training = self.old_training

  def _test_optim(self, tinygrad_optim, torch_optim, steps, opts, atol, rtol):
    for x,y in zip(step(Tensor, tinygrad_optim, steps, **opts),
@@ -154,11 +157,13 @@ class TestOptim(unittest.TestCase):
    t = Tensor.ones((1,1))
    optimizer = Adam([t])
    optimizer.zero_grad()
+    old_state = Tensor.training
    t.sum().backward()
-    with Context(TRAINING=0):
-      self.assertRaises(RuntimeError, optimizer.step)
-    with Context(TRAINING=1):
-      optimizer.step()
+    Tensor.training = False
+    self.assertRaises(RuntimeError, optimizer.step)
+    Tensor.training = True
+    optimizer.step()
+    Tensor.training = old_state

  def test_lamb_cpu_offload(self):
    # test that LAMB works when optimizer params (m, v, b1_t, b2_t) are moved to CPU
@@ -43,10 +43,10 @@ def sexec(out:Tensor, opts:list[Opt], replace_src=None, run_count=3):
  call = linear.src[-1]
  prg = to_program(replace_opts(call.src[0], opts), renderer=Device[Device.DEFAULT].renderer)
  if replace_src is not None:
-    old_name = prg.src[2].arg.split("__attribute__((noinline)) void ")[1].split("(")[0]
-    new_src = replace_src + "/* DSP boilerplate */" + prg.src[2].arg.split("/* DSP boilerplate */")[1].replace(old_name, "fxn")
+    old_name = prg.src[3].arg.split("__attribute__((noinline)) void ")[1].split("(")[0]
+    new_src = replace_src + "/* DSP boilerplate */" + prg.src[3].arg.split("/* DSP boilerplate */")[1].replace(old_name, "fxn")
    # drop BINARY and replace SOURCE so run_linear recompiles
-    prg = prg.replace(src=prg.src[:2] + (UOp(Ops.SOURCE, arg=new_src),))
+    prg = prg.replace(src=prg.src[:3] + (UOp(Ops.SOURCE, arg=new_src),))
  linear = linear.replace(src=linear.src[:-1] + (call.replace(src=(prg, *call.src[1:])),))
  for _ in range(run_count): run_linear(linear)

@@ -81,7 +81,7 @@ class TestQuantizeOnnxCPU(unittest.TestCase):
    with Context(QUANTIZE=1):
      linear = run_onnx({"input":inp})["output"].schedule_linear()
      prg = to_program(linear.src[-2].src[0], renderer=Device[Device.DEFAULT].renderer)
-      daccs = [u for u in tuple(prg.src[1].src) if u.op is Ops.BUFFER and u.addrspace is AddrSpace.REG]
+      daccs = [u for u in tuple(prg.src[2].src) if u.op is Ops.BUFFER and u.addrspace is AddrSpace.REG]
      assert all(u.dtype.scalar() is dtypes.int for u in daccs)

@unittest.skipIf(Device.DEFAULT != "DSP", "only tests for DSP")
@@ -77,7 +77,7 @@ class TestRandomness(unittest.TestCase):
      ast = call.src[0]
      if ast.op is Ops.SINK:
        prg = to_program(ast, renderer=Device[Device.DEFAULT].renderer)
-        for u in tuple(prg.src[1].src):
+        for u in tuple(prg.src[2].src):
          self.assertNotIn(u.dtype, {dtypes.long, dtypes.ulong}, msg=f"long found in {prg.arg.name}")

  def test_threefry_against_reference_full(self):
@@ -63,7 +63,7 @@ class TestCStyleFailures(unittest.TestCase):
    for _ in range(5): ret = python_alu[op](ret, Tensor.empty(1, dtype=dtype))
    linear = ret.schedule_linear()
    assert len(linear.src) == 1
-    src = to_program(linear.src[0].src[0], Device[Device.DEFAULT].renderer).src[2].arg
+    src = to_program(linear.src[0].src[0], Device[Device.DEFAULT].renderer).src[3].arg
    self.assertEqual("("*5 not in src, should_strip_paren)

  def test_repeat_add(self): self._test_src_strip_paren(Ops.ADD)
@@ -166,13 +166,6 @@ class TestSchedule(unittest.TestCase):
    self.assertEqual(t.device, dev)
    np.testing.assert_equal(t[:64].numpy(), np.arange(64).cumsum())

-  def test_copy_multi_scalar(self):
-    devs = ("CPU:0", "CPU:1")
-    x = Tensor.ones(2, device="CPU").shard(devs, axis=0).realize()
-    out = (x.sum()*2).reshape(1).to("CPU")
-    run_linear(*check_schedule(out, 5))
-    np.testing.assert_equal(out.numpy(), [4.])
-
 class TestLimitBufs(unittest.TestCase):
  @unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT == "NV", "crashes in ocelot")
  def test_limit_bufs_with_var(self):
@@ -455,14 +448,6 @@ class TestCopyFolding(unittest.TestCase):
    b.realize()
    self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])

-  def test_permute_copy_to_device(self):
-    b = Tensor([[0, 1, 2, 3], [4, 5, 6, 7]], device="CPU").permute(1, 0).to("PYTHON")
-    self.assertListEqual(b.tolist(), [[0, 4], [1, 5], [2, 6], [3, 7]])
-
-  def test_flip_copy_to_device(self):
-    b = Tensor([0, 1, 2, 3], device="CPU").flip(0).to("PYTHON")
-    self.assertListEqual(b.tolist(), [3, 2, 1, 0])
-
 class TestUOpBecome(unittest.TestCase):
  def test_setitem_offset(self):
    a = Tensor.full((16,), 0.).contiguous().realize()
@@ -239,7 +239,6 @@ class TestAssembly(unittest.TestCase):
    self.assertIn(Ops.SHL, ops)
    self.assertIn(Ops.MUL, ops)

-  @unittest.skip("this is a questionable microoptimization i won't enforce")
  def test_mulacc_unrolled(self):
    # test that     acc = acc + a0*b0 + a1*b1 + a2*b2 + a3*b3
    # is not        acc = acc + (a0*b0 + a1*b1 + a2*b2 + a3*b3)
@@ -250,7 +249,7 @@ class TestAssembly(unittest.TestCase):
    opts_to_apply = [Opt(OptOps.UNROLL, 0, 4)]
    ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
    program = to_program(ast, Device[Device.DEFAULT].renderer)
-    uops = tuple(program.src[1].src)
+    uops = tuple(program.src[2].src)
    self.assertGreaterEqual(len([x.op for x in uops if x.op is Ops.MULACC]), 4)

  def test_mulacc_shl(self):
@@ -1,68 +0,0 @@
-from tinygrad import Tensor, dtypes, Device, TinyJit
-from tinygrad.helpers import getenv, Timing, round_up
-
-# llama 70b shapes
-DIM, HIDDEN, N_HEADS, N_KV_HEADS, HEAD_DIM = 8192, 28672, 64, 8, 128
-QKV_OUT = N_HEADS*HEAD_DIM + 2*N_KV_HEADS*HEAD_DIM
-LAYER_SHAPES = {
-  "wqkv": ((QKV_OUT, DIM), 0),
-  "wo": ((DIM, N_HEADS*HEAD_DIM), 1),
-  "w1": ((HIDDEN, DIM), 0),
-  "w3": ((HIDDEN, DIM), 0),
-  "w2": ((DIM, HIDDEN), 1),
-}
-
-def build_params(layers:int, devs:tuple[str,...], dtype):
-  ts, nbytes = [], 0
-  ndev = len(devs)
-  for _ in range(layers):
-    for (a,b), ax in LAYER_SHAPES.values():
-      shape = (round_up(a, ndev) if ax==0 else a, round_up(b, ndev) if ax==1 else b)
-      t = Tensor.empty(*shape, dtype=dtype).shard(devs, axis=ax).realize()
-      ts.append(t)
-      nbytes += t.numel()*dtype.itemsize
-  return ts, nbytes
-
-def sync():
-  Device["CPU"].synchronize()
-  Device[Device.DEFAULT].synchronize()
-
-if __name__ == "__main__":
-  ngpus = getenv("GPUS", 8)
-  layers = getenv("LAYERS", 8)
-  devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(ngpus))
-
-  # to cpu
-  grads, nbytes = build_params(layers, devs, dtypes.bfloat16)
-  cpu_dst = [Tensor.rand(*g.shape, dtype=dtypes.bfloat16, device="CPU").contiguous().realize() for g in grads]
-  sync()
-  def to_cpu():
-    for d,g in zip(cpu_dst, grads):
-      d.assign(g.to("CPU"))
-    return Tensor.realize(*cpu_dst)
-  to_cpu_jit = TinyJit(to_cpu)
-  for _ in range(3):
-    to_cpu_jit()
-    sync()
-  for _ in range(5):
-    with Timing("to cpu ", on_exit=lambda ns: f" @ {nbytes/ns:.2f} GB/s"):
-      to_cpu_jit()
-      sync()
-
-  # to gpu
-  from extra.gemm.cdna_asm_gemm import FP8_DTYPE
-  w_gpu, nbytes = build_params(layers, devs, FP8_DTYPE)
-  cpu_src = [Tensor.rand(*w.shape, dtype=FP8_DTYPE, device="CPU").contiguous().realize() for w in w_gpu]
-  sync()
-  def to_gpu():
-    for w,c in zip(w_gpu, cpu_src):
-      w.assign(c.shard_like(w))
-    return Tensor.realize(*w_gpu)
-  to_gpu_jit = TinyJit(to_gpu)
-  for _ in range(3):
-    to_gpu_jit()
-    sync()
-  for _ in range(5):
-    with Timing("to gpu ", on_exit=lambda ns: f" @ {nbytes/ns:.2f} GB/s"):
-      to_gpu_jit()
-      sync()
@@ -1,15 +1,16 @@
 from tinygrad import nn, Tensor, dtypes
-from tinygrad.helpers import DEV, Timing, Context
+from tinygrad.helpers import DEV, Timing

 from extra.models.llama import Transformer
 from examples.llama3 import MODEL_PARAMS

 if __name__ == "__main__":
  DEV.value = "NULL"
+  Tensor.training = True
  #model_size = "8B"
  model_size = "405B"

-  with Context(TRAINING=1), Timing("total "):
+  with Timing("total "):
    with Timing("***** create model in    "):
      model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=nn.Linear, embedding=nn.Embedding,
                          max_context=1024, jit=True, disable_kv_cache=True)
@@ -63,8 +63,8 @@ class BenchmarkResnetTrain(unittest.TestCase):
    with Context(TRACK_MATCH_STATS=0): Tensor.realize(*[t.assign(t.detach().contiguous()) for t in get_parameters(optim)])

    JITCNT = getenv("JITCNT", 1)
+    Tensor.training = True
    @TinyJit
-    @Context(TRAINING=1)
    def step(x):
      optim.zero_grad()
      x.grad = None
@@ -4,7 +4,7 @@ import numpy as np
 import torch

 from tinygrad import GlobalCounters, Tensor, Device
-from tinygrad.helpers import getenv, Context
+from tinygrad.helpers import getenv
 from tinygrad.nn.state import get_parameters
 from tinygrad.engine.realize import capturing, run_linear
 from tinygrad.tensor import _to_np_dtype
@@ -38,7 +38,11 @@ from extra.models.vit import ViT

@unittest.skipUnless(Device.DEFAULT == "CL", "Not Implemented")
 class TestInferenceMinKernels(unittest.TestCase):
-  def setUp(self): self.enterContext(Context(TRAINING=0))
+  def setUp(self):
+    self.training_old = Tensor.training
+    Tensor.training = False
+  def tearDown(self):
+    Tensor.training = self.training_old

  def test_convnext(self):
    model = ConvNeXt()
@@ -8,7 +8,6 @@ from extra.lr_scheduler import LRSchedulerGroup

 from tinygrad.tensor import Tensor
 from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup
-from tinygrad.helpers import Context

 from test.external.mlperf_resnet.lars_optimizer import LARSOptimizer

@@ -104,7 +103,11 @@ def create_tf_polylr(initial_lr, end_lr, train_steps, warmup, power=2, skip_list
                                      initial_learning_rate=initial_lr, end_learning_rate=end_lr, warmup_epochs=warmup)

 class ExternalTestOptim(unittest.TestCase):
-  def setUp(self): self.enterContext(Context(TRAINING=1))
+  def setUp(self):
+    self.old_training = Tensor.training
+    Tensor.training = True
+  def tearDown(self):
+    Tensor.training = self.old_training

  def _test_optim(self, tinygrad_optim, tensorflow_optim, steps, opts, atol, rtol, tiny_sched=None, tf_sched=None, schedopts=None, do_optim=True):
    for x,y in zip(step(tinygrad_optim, steps=steps, kwargs=opts, scheduler=tiny_sched, schedopts=schedopts, do_optim=do_optim),
@@ -3,7 +3,7 @@ import z3
 from tinygrad import dtypes, Device
 from tinygrad.uop.validate import uops_to_z3, z3_cdiv
 from tinygrad.uop.ops import UOp
-from tinygrad.codegen.decomp.op import fast_idiv
+from tinygrad.uop.decompositions import fast_idiv
 random.seed(42)

 powers_of_two = [2**i for i in range(64)]
@@ -11,7 +11,7 @@ if __name__ == "__main__":
  for i in range(10_000):
    if i % 1000 == 0:
      print(f"Progress: {i}")
-    dt = random.choice(dtypes.ints)
+    dt = random.choice(dtypes.ints + tuple(dt.vec(4) for dt in dtypes.ints))
    u = UOp.variable('x', random.randint(dt.min, 0), random.randint(1, dt.max), dtype=dt)
    d = random.randint(1, max(1, u.vmax)*2)
    if d in powers_of_two: continue
@@ -41,8 +41,8 @@ class BenchmarkBertTrain(unittest.TestCase):
    with Context(TRACK_MATCH_STATS=0): Tensor.realize(*[t.assign(t.detach().contiguous()) for t in get_parameters(optim)])

    JITCNT = getenv("JITCNT", 1)
+    Tensor.training = True
    @TinyJit
-    @Context(TRAINING=1)
    def step(inputs):
      optim.zero_grad()
      for i in inputs: i.grad = None
@@ -31,6 +31,8 @@ if __name__ == "__main__":
  assert all(os.path.exists(os.path.join(INIT_CKPT_DIR, f)) for f in required_files), \
    f"Missing checkpoint files in INIT_CKPT_DIR: {required_files}"

+  Tensor.training = False
+
  model = get_mlperf_bert_model(INIT_CKPT_DIR)

  for _, x in get_state_dict(model).items():
@@ -48,9 +48,9 @@ def replay_to_program(p:UOp, ast:UOp, renderer:Renderer) -> tuple[str, str, tupl
    if sink_arg.beam: sink_arg = replace(sink_arg, opts_to_apply=p.src[0].arg.applied_opts)
    input_ast = ast.replace(arg=replace(sink_arg, name=p.src[0].arg.name))
  p2 = to_program(input_ast, renderer=renderer)
-  device = renderer.target.device
+  device = p.src[1].arg
  def to_str(ret:UOp) -> str:
-    src = next(x.arg for x in ret.src if x.op is Ops.SOURCE)
+    src = ret.src[3].arg
    # PYTHON renderer pickles UOps, first unpickle and decode here
    if device.startswith("PYTHON"): return "\n".join([str(x) for x in pickle.loads(base64.b64decode(src))])
    return src
@@ -84,7 +84,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple
    allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
  g = UOp.param(0, uop.dtype.ptr(1))
  prg = to_program(UOp.store(g.index(UOp.const(dtypes.int, 0), ptr=True), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
-  prog = PythonProgram("run", PythonCompiler().compile(prg.src[2].arg))
+  prog = PythonProgram("run", PythonCompiler().compile(prg.src[3].arg))
  prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs, vals=vals)
  return out_buf.cast(uop.dtype.fmt or "").tolist()[0]

@@ -375,7 +375,7 @@ def _mem_store(mem: UOp, addr: UOp, val: UOp, active: UOp, addr_bits: int = 32,
  """Conditional memory store with sub-word support. Returns list of store UOps."""
  adt = dtypes.uint64 if addr_bits == 64 else dtypes.uint32
  word_addr = addr >> UOp.const(adt, 2)
-  idx = mem.index(word_addr.valid(active))
+  idx = mem.index(word_addr.cast(dtypes.int).valid(active))
  if data_bits == 32: return [idx.store(active.where(_to_u32(val), idx))]
  # Sub-word store: read-modify-write with mask
  byte_pos = addr.cast(dtypes.uint32) & _c(3)
@@ -388,7 +388,7 @@ def _mem_store(mem: UOp, addr: UOp, val: UOp, active: UOp, addr_bits: int = 32,
  is_cross = byte_pos.eq(_c(3))
  cross_word0 = (idx & _c(0x00FFFFFF)) | ((val_u32 & _c(0xFF)) << _c(24))
  store0 = idx.store(active.where(is_cross.where(cross_word0, new_word), idx))
-  next_idx = mem.index((word_addr + UOp.const(adt, 1)).valid(active & is_cross))
+  next_idx = mem.index((word_addr + UOp.const(adt, 1)).cast(dtypes.int).valid(active & is_cross))
  cross_word1 = (next_idx & _c(0xFFFFFF00)) | ((val_u32 >> _c(8)) & _c(0xFF))
  return [store0, next_idx.store((active & is_cross).where(cross_word1, next_idx))]

@@ -398,7 +398,7 @@ def _mem_store_bytes(mem: UOp, addr: UOp, val: UOp, active: UOp, data_bits: int
  val_u32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
  for i in range(data_bits // 8):
    byte_val = (val_u32 >> UOp.const(dtypes.uint32, i * 8)) & UOp.const(dtypes.uint32, 0xFF)
-    stores.append(mem.index((addr + UOp.const(dtypes.uint64, i)).valid(active)).store(byte_val.cast(dtypes.uint8)))
+    stores.append(mem.index((addr + UOp.const(dtypes.uint64, i)).cast(dtypes.int).valid(active)).store(byte_val.cast(dtypes.uint8)))
  return stores

 def _collect_data_slices(assigns: list[tuple[str, UOp]], data_prefix: str, pcode_vars: dict | None = None, op_name: str = "") -> dict[int, UOp]:
@@ -463,7 +463,7 @@ class _Ctx:
    """Read instruction dword from vmem at PC + dword_idx*4."""
    pc = self.rpc()
    addr = pc if dword_idx == 0 else pc + UOp.const(dtypes.uint64, dword_idx * 4)
-    return self.vmem.index(addr >> UOp.const(dtypes.uint64, 2), ptr=True).load()
+    return self.vmem.index((addr >> UOp.const(dtypes.uint64, 2)).cast(dtypes.int), ptr=True).load()

  def inst_field(self, field) -> UOp:
    """Extract field bits from instruction encoding. Tracks field for canonical key computation."""
@@ -516,14 +516,14 @@ class _Ctx:
  # Dynamic register access (takes UOp index instead of int)
  def rsgpr_dyn(self, reg: UOp, valid: UOp | None = None) -> UOp:
    """Read SGPR with dynamic register index."""
-    if valid is not None: return self.sgpr.index(reg.valid(valid), ptr=True).load()
-    return self.sgpr.index(reg, ptr=True).load()
+    if valid is not None: return self.sgpr.index(reg.cast(dtypes.int).valid(valid), ptr=True).load()
+    return self.sgpr.index(reg.cast(dtypes.int), ptr=True).load()

  def wsgpr_dyn(self, reg: UOp, val: UOp) -> UOp:
    """Write SGPR with dynamic register index. On RDNA, index 124 = NULL (writes discarded). On CDNA, index 124 = M0 (read/write)."""
    # RDNA: NULL (124) discards writes. CDNA: M0 (124) is writable.
    valid = None if self.wave_size == 64 else reg.ne(_c(124))
-    return self.sgpr.index(reg.valid(valid) if valid is not None else reg, ptr=True).store(val.cast(dtypes.uint32))
+    return self.sgpr.index(reg.cast(dtypes.int).valid(valid) if valid is not None else reg.cast(dtypes.int), ptr=True).store(val.cast(dtypes.uint32))

  def wmask(self, reg: UOp, val: UOp) -> list[UOp]:
    """Write a lane mask (VCC/EXEC). Splits into lo/hi for wave64."""
@@ -656,7 +656,7 @@ class _Ctx:
    stores = []
    for dest, val in assigns:
      if dest.startswith('D0'): stores.append(self.wsgpr_dyn(vdst_off, val.cast(dtypes.uint32)))
-      elif dest.startswith('VGPR['): stores.append(self.vgpr.index(val[0]).store(val[1].cast(dtypes.uint32)))
+      elif dest.startswith('VGPR['): stores.append(self.vgpr.index(val[0].cast(dtypes.int)).store(val[1].cast(dtypes.uint32)))
    return UOp.sink(*stores, *self.inc_pc())

  def compile_vop_pcode(self, op, srcs: dict[str, UOp | int], lane: UOp, vdst_reg: UOp, exec_mask: UOp,
@@ -711,11 +711,11 @@ class _Ctx:
        # VGPR bit-slice: (vgpr_idx, rhs_val, hi_bit, lo_bit) - hi/lo are UOp constants
        hi_bit, lo_bit = int(val[2].arg), int(val[3].arg)
        width = hi_bit - lo_bit + 1
-        old = self.vgpr.index(val[0], ptr=True).load()
+        old = self.vgpr.index(val[0].cast(dtypes.int), ptr=True).load()
        new_val = _set_bits(old, _val_to_bits(val[1]), width, lo_bit).cast(dtypes.uint32)
        active = _lane_active(exec_mask, lane)
        if len(val) > 4: active = active & _to_bool(val[4])
-        raw_stores.append(('vgpr_direct', self.vgpr.index(val[0].valid(active)).store(new_val)))
+        raw_stores.append(('vgpr_direct', self.vgpr.index(val[0].cast(dtypes.int).valid(active)).store(new_val)))
        continue
      if 'D0' in dest and '[laneId]' in dest:
        old_vcc = self.rmask(_c(VCC_LO.offset))
@@ -834,7 +834,7 @@ def _compile_smem(inst: ir3.SMEM | ir4.SMEM, ctx: _Ctx) -> UOp:
  nval = int(part.removeprefix('DWORD').removeprefix('X') or '1') if 'DWORD' in part else int(part[1:]) / 32 * (-1 if part[0] == 'I' else 1)
  ndwords = max(1, int(abs(nval)))
  dword_base = addr >> UOp.const(dtypes.uint64, 2)
-  vals = [ctx.vmem.index(dword_base + UOp.const(dtypes.uint64, i)) for i in range(ndwords)]
+  vals = [ctx.vmem.index((dword_base + UOp.const(dtypes.uint64, i)).cast(dtypes.int)) for i in range(ndwords)]
  if abs(nval) < 1:
    nbits = int(abs(nval) * 32)
    byte_off = (addr & UOp.const(dtypes.uint64, 3)).cast(dtypes.uint32) * UOp.const(dtypes.uint32, 8)
@@ -1819,7 +1819,7 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
    srcs = {'ADDR': addr_reg, 'DATA0': vdata_reg, 'VDST': vdst_reg, 'OFFSET': offset,
            'EXEC': exec_mask.cast(dtypes.uint64), '_vgpr': ctx.vgpr, '_wave_size': ctx.wave_size}
    _, assigns = parse_pcode(pcode, srcs)
-    stores = [ctx.vgpr.index(val[0]).store(val[1].cast(dtypes.uint32)) for dest, val in assigns if dest.startswith('VGPR[')]
+    stores = [ctx.vgpr.index(val[0].cast(dtypes.int)).store(val[1].cast(dtypes.uint32)) for dest, val in assigns if dest.startswith('VGPR[')]
    return UOp.sink(*stores, *ctx.inc_pc())

  def make_addr(lane: UOp) -> UOp:
@@ -1853,7 +1853,7 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
  def wmem(addr: UOp, val: UOp, active: UOp, data_bits: int = 32) -> UOp:
    if data_bits < 32:
      # Sub-dword LDS write: read-modify-write within the uint32 slot
-      word_addr = addr >> addr_shift
+      word_addr = (addr >> addr_shift).cast(dtypes.int)
      idx = mem.index(word_addr.valid(active))
      byte_pos = addr.cast(dtypes.uint32) & _c(3)
      byte_shift = byte_pos * _c(8)
@@ -1861,7 +1861,7 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
      mask = size_mask << byte_shift
      new_word = (idx & (mask ^ _c(0xFFFFFFFF))) | ((val.cast(dtypes.uint32) & size_mask) << byte_shift)
      return idx.store(active.where(new_word, idx))
-    idx = mem.index(addr >> addr_shift)
+    idx = mem.index((addr >> addr_shift).cast(dtypes.int))
    return idx.store(active.where(val, idx.load()))

  def make_srcs(lane: UOp) -> dict:
@@ -2011,7 +2011,7 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp:
    for i in range(n_dwords):
      word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2)
      val = in_bounds.where(mem.index(word_addr.cast(dtypes.int64), ptr=True).load(), _c(0))
-      lds_idx = (lds_addr + _c(i * 4)) >> _c(2)
+      lds_idx = ((lds_addr + _c(i * 4)) >> _c(2)).cast(dtypes.int)
      lds_slot = ctx.lds.index(lds_idx.valid(active))
      stores.append(lds_slot.store(active.where(val, lds_slot)))
  elif is_store:
@@ -2,7 +2,7 @@
 from typing import Any, Callable
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import Ops, UOp
-from tinygrad.codegen.decomp.dtype import f2f
+from tinygrad.uop.decompositions import f2f

 # Type alias for vars dict: stores UOps and tuples for lambda definitions
 VarVal = UOp | tuple[str, list[str], str]
@@ -831,7 +831,7 @@ class Parser:
    def mindex(idx:UOp, ptr=False): return mem.index(idx.valid(active) if active is not None else idx, ptr=ptr)
    byte_mem = mem.dtype.base == dtypes.uint8
    if byte_mem:
-      idx = addr
+      idx = addr.cast(dtypes.int)
      if dt in (dtypes.uint64, dtypes.int64, dtypes.float64):
        val = _u32(0).cast(dtypes.uint64)
        for i in range(8): val = val | (mindex(idx + _const(dtypes.int, i), ptr=True).load().cast(dtypes.uint64) << _u64(i * 8))
@@ -845,10 +845,10 @@ class Parser:
        val = _u32(0)
        for i in range(4): val = val | (mindex(idx + _const(dtypes.int, i), ptr=True).load().cast(dtypes.uint32) << _u32(i * 8))
    else:
-      idx = addr >> _const(addr.dtype, 2)
+      idx = (addr >> _const(addr.dtype, 2)).cast(dtypes.int)
      val = mindex(idx)
      if dt in (dtypes.uint64, dtypes.int64, dtypes.float64):
-        idx2 = (addr + _const(adt, 4)) >> _const(adt, 2)
+        idx2 = ((addr + _const(adt, 4)) >> _const(adt, 2)).cast(dtypes.int)
        val = val.cast(dtypes.uint64) | (mindex(idx2).cast(dtypes.uint64) << _u64(32))
      elif dt in (dtypes.uint8, dtypes.int8): val = (val >> ((addr & _const(adt, 3)).cast(dtypes.uint32) * _u32(8))) & _u32(0xFF)
      elif dt in (dtypes.uint16, dtypes.int16):
@@ -21,7 +21,7 @@ class TestDisassembly(unittest.TestCase):
    c = Tensor([1], dtype=dtypes.float16) + Tensor([1], dtype=dtypes.float16)
    s = c.schedule_linear().src[-1]
    p = to_program(s.src[0], Device[Device.DEFAULT].renderer)
-    lib = Device[Device.DEFAULT].compiler.compile(p.src[2].arg)
+    lib = Device[Device.DEFAULT].compiler.compile(p.src[3].arg)
    out = io.StringIO()
    with redirect_stdout(out): Device[Device.DEFAULT].compiler.disassemble(lib)
    assert "fcvt" not in out.getvalue()
@@ -94,17 +94,6 @@ class TestGroupedDims(unittest.TestCase):
    assert idxs[2].op is Ops.SPECIAL, f"expected SPECIAL for direct-mapped dim, got {idxs[2].op}"
    assert idxs[3].op is Ops.SPECIAL, f"expected SPECIAL for direct-mapped dim, got {idxs[3].op}"

-  def test_grouped_dims_high_rank(self):
-    # 4D collapsed onto 2 axes
-    self._check_grouped_dims("gidx", (4,4,4,4), (16,16), False, [16,16])
-    # 4D untouched
-    self._check_grouped_dims("gidx", (2,3,4,5), None, False, [2,3,4,5])
-    idxs = get_grouped_dims("gidx", (2,3,4,5), None, False)
-    assert all(u.op is Ops.SPECIAL for u in idxs), f"expected all-SPECIAL when untouched, got {[u.op for u in idxs]}"
-    # 5D and 6D collapsed onto 3 axes
-    self._check_grouped_dims("gidx", (2,2,2,2,2), (4,4,4), False, [4,4,2])
-    self._check_grouped_dims("gidx", (2,2,2,2,2,2), (8,8,8), False, [8,4,2])
-
  def test_global_prod_max(self):
    g, l = UOp.range(256, 0, AxisType.GLOBAL), UOp.range(256, 1, AxisType.LOCAL)
    sink = UOp.param(0, dtypes.float.ptr()).index(g + l).store(UOp.const(dtypes.float, 1.0)).end(g, l).sink(arg=KernelInfo())
@@ -183,27 +183,33 @@ class TestGEPAndVectorizeRewrite(unittest.TestCase):
  def test_gep_single_element_extraction(self):
    # GEP on a vector dtype to extract a single element
    base_vector = UOp.const(dtypes.float32.vec(4), (1.0, 2.0, 3.0, 4.0))
-    self.assertEqual(apply_rewrite(base_vector.index(2)).arg, 3.0)
+    self.assertEqual(apply_rewrite(base_vector.gep(2)).arg, 3.0)

  def test_gep_tuple_extraction(self):
    # GEP on a vector dtype to extract multiple elements as a vector
    base_vector = UOp.const(dtypes.float32.vec(4), (1.0, 2.0, 3.0, 4.0))
-    self.assertEqual(list(apply_rewrite_values(UOp.vectorize(*[base_vector.index(i) for i in (2, 3)]))), [3.0, 4.0])
+    self.assertEqual(list(apply_rewrite_values(base_vector.gep((2, 3)))), [3.0, 4.0])

  def test_gep_on_const_stack(self):
    # GEP on a const STACK to extract a single element
    const_stack = UOp.const(dtypes.float32.vec(4), (1.0, 2.0, 3.0, 4.0))
-    self.assertEqual(apply_rewrite(const_stack.index(2)).arg, 3.0)
+    self.assertEqual(apply_rewrite(const_stack.gep(2)).arg, 3.0)

  def test_gep_tuple_on_const_stack(self):
    # GEP on a const STACK using a tuple to extract multiple elements
    const_stack = UOp.const(dtypes.float32.vec(4), (7.0, 8.0, 9.0, 10.0))
-    self.assertEqual(list(apply_rewrite_values(UOp.vectorize(*[const_stack.index(i) for i in (1, 3)]))), [8.0, 10.0])
+    self.assertEqual(list(apply_rewrite_values(const_stack.gep((1, 3)))), [8.0, 10.0])
+
+  def test_gep_gep_simplification(self):
+    # Nested GEP simplification on a vector dtype
+    base_vector = UOp.const(dtypes.float32.vec(4), (10.0, 20.0, 30.0, 40.0))
+    gep_inner = base_vector.gep(1)  # Extract 2nd element (20.0)
+    self.assertEqual(apply_rewrite(gep_inner.gep(0)).arg, 20.0)

  def test_vectorize_multiple_elements(self):
    # Vectorizing multiple elements using GEP
    base_vector = UOp.const(dtypes.float32.vec(4), (5.0, 10.0, 15.0, 20.0))
-    vectorized_uop = UOp(Ops.STACK, dtypes.float32.vec(4), src=tuple(base_vector.index(i) for i in range(4)))
+    vectorized_uop = UOp(Ops.STACK, dtypes.float32.vec(4), src=(base_vector.gep(0), base_vector.gep(1), base_vector.gep(2), base_vector.gep(3)))
    self.assertEqual(list(apply_rewrite_values(vectorized_uop)), [5.0, 10.0, 15.0, 20.0])


@@ -1,6 +1,6 @@
 import ctypes, gzip, unittest, timeit, pickle
 from tinygrad import Variable
-from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, mv_address, count, all_same
+from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, mv_address, get_contraction, count, all_same
 from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
 from tinygrad.helpers import ceildiv, ansistrip, get_shape
 from tinygrad.tensor import Tensor
@@ -273,6 +273,75 @@ class TestMemoryview(unittest.TestCase):
    mva_us = timeit.timeit(lambda: mv_address(x), number=iters) * 1e6 / iters
    print(f"from_mv vs mv_address: {fmv_us:8.3f} µs vs {mva_us:8.3f} µs")

+class TestGetContraction(unittest.TestCase):
+  def test_contraction(self):
+    r = get_contraction((1,2,3,4), (2,3,4))
+    self.assertEqual(r, [[0, 1], [2], [3]])
+
+    r = get_contraction((2,1,3,4), (2,3,4))
+    self.assertEqual(r, [[0], [1, 2], [3]])
+
+    r = get_contraction((1,2,3,1,4), (1,2,3,4))
+    self.assertEqual(r, [[], [0, 1], [2], [3, 4]])
+
+    r = get_contraction((1,2,3,1,4,1,1), (2,3,4))
+    self.assertEqual(r, [[0, 1], [2], [3, 4, 5, 6]])
+
+    r = get_contraction((1,2,3,4), (1,2,3*4))
+    self.assertEqual(r, [[], [0, 1], [2, 3]])
+
+    r = get_contraction((1,2,3,4), (2,1,3,4))
+    self.assertEqual(r, [[0, 1], [], [2], [3]])
+
+    r = get_contraction((1,2,3,4), (1,1,2*3*4,1))
+    self.assertEqual(r, [[], [], [0,1,2,3], []])
+
+    r = get_contraction((2,1,3,4), (1,2,3,4))
+    self.assertEqual(r, [[], [0], [1, 2], [3]])
+
+    r = get_contraction((1,2,3,4), (2*3*4,1,1,1))
+    self.assertEqual(r, [[0, 1, 2, 3], [], [], []])
+
+    r = get_contraction((4,4,4,4), (16,1,16))
+    self.assertEqual(r, [[0, 1], [], [2, 3]])
+
+    r = get_contraction((1,2,3,4,1,1,1), (2,3,4))
+    self.assertEqual(r, [[0, 1], [2], [3, 4, 5, 6]])
+
+    r = get_contraction((1,2,3,4), (1,2,3,4,1))
+    self.assertEqual(r, [[], [0, 1], [2], [3], []])
+
+    r = get_contraction((14,1,384,14,1,1,1,1), (1,14,384,14))
+    self.assertEqual(r, [[], [0], [1,2], [3,4,5,6,7]])
+
+    r = get_contraction((14,1,384,1,14,1,1,1,1), (1,14,384,14))
+    self.assertEqual(r, [[], [0], [1,2], [3,4,5,6,7,8]])
+
+    r = get_contraction((512, 512), (1, 1, 512, 1, 1, 1, 1, 512))
+    self.assertEqual(r, [[], [], [0], [], [], [], [], [1]])
+
+    r = get_contraction((1,2,3,4), (1,2,6,2))
+    self.assertEqual(r, None)
+
+  def test_contraction_ones(self):
+    r = get_contraction((1,), (1,1,1))
+    self.assertEqual(r, [[], [], [0]])
+
+    r = get_contraction((1,1), (1,1,1))
+    self.assertEqual(r, [[], [], [0, 1]])
+
+    r = get_contraction((1,1,1,1), (1,))
+    self.assertEqual(r, [[0,1,2,3]])
+
+    r = get_contraction((1,1,1,1), (1,1))
+    self.assertEqual(r, [[], [0,1,2,3]])
+
+    r = get_contraction((1,1,1,1), (1,1,1))
+    self.assertEqual(r, [[], [], [0,1,2,3]])
+
+    r = get_contraction((1,1,1,1), (1,1,1,1))
+    self.assertEqual(r, [[], [], [], [0,1,2,3]])
+
 class TestGetShape(unittest.TestCase):
  def test_get_shape(self):
    assert get_shape(2) == ()
@@ -15,7 +15,7 @@ class TestLinearizerRewrite(unittest.TestCase):
      opts_to_apply.append(Opt(OptOps.UNROLL, 0, 4))
      ast = si.src[0].replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
      prg = to_program(ast, Device["CPU"].renderer)
-      print(prg.src[2].arg)
+      print(prg.src[3].arg)

  def test_arange(self):
    out = Tensor.arange(32).clone("NULL")
@@ -25,7 +25,7 @@ class TestLinearizerRewrite(unittest.TestCase):
      opts_to_apply.append(Opt(OptOps.UPCAST, 0, 4))
      ast = si.src[0].replace(arg=KernelInfo(opts_to_apply=tuple(opts_to_apply)))
      prg = to_program(ast, Device["CPU"].renderer)
-      print(prg.src[2].arg)
+      print(prg.src[3].arg)

  def test_kernel_info(self):
    out = Tensor.arange(4).clone("NULL")
@@ -50,7 +50,7 @@ class TestPatternMatcher(unittest.TestCase):
    def fxn(ctx, x):
      ctx.append(True)
      assert len(x.src) == 0
-      return x.replace(src=(UOp(Ops.NOOP),))
+      return x.replace(src=(UOp(Ops.DEVICE, arg="blah"),))
    matcher = PatternMatcher([(UPat(Ops.CONST, src=(), name="x"), fxn)])
    c1 = UOp.const(dtypes.float, 1.0)
    # second rewrite shouldn't match anything
@@ -266,7 +266,7 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.empty(big_enough).realize()
    with Context(SPLIT_REDUCEOP=1):
      out = (x - x.max(keepdim=True)).max()
-      check_schedule(out, 3)
+      check_schedule(out, 4)

  def test_example_matmul_contig(self):
    x = Tensor.eye(64).clone().realize()
@@ -355,7 +355,8 @@ class TestSchedule(unittest.TestCase):
    b = Tensor.empty((1, 16)).realize()
    out0 = a.sum() + 2
    out1 = a.sum() + b
-    check_schedule([out0, out1], 2)
+    # check_schedule([out0, out1], 2)
+    check_schedule([out0, out1], 3)

  def test_scaled_dot_product_attention_multireduce_fusion(self):
    q = Tensor.empty(32,8,16,8).realize()
@@ -545,7 +546,8 @@ class TestSchedule(unittest.TestCase):
    a = Tensor.empty(3, 4, 5).abs().realize()
    b = Tensor.empty(3, 4, 5).abs().realize()
    out = (a.log2().pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum()+b).abs().log2().pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum().contiguous()
-    check_schedule(out, 1)
+    # check_schedule(out, 1)
+    check_schedule(out, 2)

  def test_shrink_pad_safe(self):
    a = Tensor.ones((3, )).contiguous().realize()
@@ -1840,7 +1842,7 @@ class TestFusionOp(unittest.TestCase):
    linear = a.schedule_linear()
    prg = to_program(linear.src[-1].src[0], renderer=Device[Device.DEFAULT].renderer)
    self.assertLess(time.perf_counter()-st, 2.0)
-    assert len(prg.src[2].arg.splitlines()) < 250
+    assert len(prg.src[3].arg.splitlines()) < 250

  def test_recursive_add_cmp(self):
    st = time.perf_counter()
@@ -1,6 +1,6 @@
 import unittest, itertools

-from tinygrad.codegen.late.coalese import indexing_simplify
+from tinygrad.codegen.late.devectorizer import load_store_indexing
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite
 from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load
@@ -11,7 +11,7 @@ from test.null.test_uop_symbolic import check_uop_against_string
 # symbolic-only idx + valid simplification (no late lowering of FLOORDIV/FLOORMOD)
 def simplify_valid_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load, name="simplify_valid_idx")
 # image-aware idx + valid simplification: adds the codegen-layer matcher that drops provably in-bounds gates
-def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load+indexing_simplify, name="simplify_image_idx")
+def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load+load_store_indexing, name="simplify_image_idx")

 def get_gated_load_uop(valid:UOp, idx:UOp):
  return UOp(Ops.LOAD, dtypes.float, (
@@ -495,7 +495,7 @@ class TestImageSimplification(unittest.TestCase):
 class TestDropTrueGate(unittest.TestCase):
  def test_drop_true_gate_on_index(self):
    # test that INDEX with a constant True valid gets simplified to drop the valid
-    from tinygrad.codegen.late.coalese import indexing_simplify
+    from tinygrad.codegen.late.devectorizer import load_store_indexing
    from tinygrad.uop.ops import graph_rewrite
    from tinygrad.uop.symbolic import sym
    buf = UOp.param(0, dtypes.int.ptr())
@@ -503,7 +503,7 @@ class TestDropTrueGate(unittest.TestCase):
    true_gate = UOp.const(dtypes.bool, True)
    index_with_gate = UOp(Ops.INDEX, dtypes.int.ptr(), (buf, idx.valid(true_gate)))
    # apply the optimization
-    result = graph_rewrite(index_with_gate, sym+indexing_simplify)
+    result = graph_rewrite(index_with_gate, sym+load_store_indexing)
    # the True valid should be dropped (INDEX should only have 2 sources)
    self.assertEqual(len(result.src), 2, "True valid should be dropped from INDEX")

@@ -1,7 +1,7 @@
 # tensor tests that pass on NULL backend (no copyout needed)
 import numpy as np
 import unittest
-from tinygrad import Tensor, Device, dtypes
+from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
@@ -12,6 +12,15 @@ x_init = np.random.randn(1,3).astype(np.float32)
 W_init = np.random.randn(3,3).astype(np.float32)
 m_init = np.random.randn(1,3).astype(np.float32)

+class TestTrainMode(unittest.TestCase):
+  def test_train_mode(self):
+    assert not Tensor.training
+    @Context(TRAINING=1)
+    def f():
+      assert Tensor.training
+    f()
+    assert not Tensor.training
+
 class TestInferenceMode(unittest.TestCase):
  def test_inference(self):
    x = Tensor(x_init)
@@ -57,7 +66,7 @@ class TestIdxUpcast(unittest.TestCase):
      if ast.op is Ops.SINK:
        renderer = Device[si.src[1].buffer.device].renderer
        prg = to_program(ast, renderer)
-        return tuple(prg.src[1].src)
+        return tuple(prg.src[2].src)

  def _assert(self, dtype: DType, a: Tensor):
    uops = self._schedule_render(a)
@@ -1,10 +1,9 @@
 import math, unittest
-from dataclasses import replace
-from tinygrad import Tensor, dtypes, Context
-from tinygrad.uop.ops import ParamArg, UOp, UPat, Ops, PatternMatcher, graph_rewrite
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite

 _strip_unique_pm = PatternMatcher([
-  (UPat(Ops.BUFFER, name="b"), lambda b: b.replace(arg=replace(b.arg, slot=0)) if isinstance(b.arg, ParamArg) and b.arg.slot != 0 else None),
+  (UPat((Ops.UNIQUE, Ops.LUNIQUE), name="u"), lambda u: u.replace(arg=0) if u.arg != 0 else None),
 ])
 def _strip_unique(u: UOp) -> UOp: return graph_rewrite(u, _strip_unique_pm)

@@ -395,24 +394,6 @@ class TestTensorUOpConv2d(unittest.TestCase):
  def test_conv2d_3d(self):
    w = _t(1, 1, 2, 2, 2).float()
    _check(self, _t(1, 1, 3, 3, 3).float(), lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop))
-  def test_conv2d_winograd(self):
-    w, a = _t(2, 2, 3, 3).float(), _t(1, 2, 6, 6).float()
-    with Context(WINO=0): direct = a.conv2d(w).uop
-    with Context(WINO=1):
-      self.assertIsNot(a.conv2d(w).uop, direct)
-      _check(self, a, lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop))
-  def test_conv2d_image(self):
-    w, a = _t(4, 4, 3, 3).float(), _t(1, 4, 8, 8).float()
-    with Context(IMAGE=0): direct = a.conv2d(w).uop
-    with Context(IMAGE=1):
-      self.assertIsNot(a.conv2d(w).uop, direct)
-      _check(self, a, lambda x: x.conv2d(w if isinstance(x, Tensor) else w.uop))
-  def test_dot_image(self):
-    y, a = _t(4, 3).float(), _t(2, 4).float()
-    with Context(IMAGE=0): direct = a.dot(y).uop
-    with Context(IMAGE=1):
-      self.assertIsNot(a.dot(y).uop, direct)
-      _check(self, a, lambda x: x.dot(y if isinstance(x, Tensor) else y.uop))
  def test_conv_transpose2d_basic(self):
    w = _t(1, 1, 2, 2).float()
    _check(self, _t(1, 1, 3, 3).float(), lambda x: x.conv_transpose2d(w if isinstance(x, Tensor) else w.uop))
@@ -456,6 +437,7 @@ class TestTensorUOpSVD(unittest.TestCase):
  def test_svd_batched(self):   self._check(_t(2, 2, 2).float())
  def test_svd_nonfull(self):   self._check(_t(3, 2).float(), full_matrices=False)

+# UOp.empty / UOp.empty_like are the canonical buffer allocators; Tensor.empty / Tensor.empty_like just forward.
 class TestUOpEmpty(unittest.TestCase):
  def test_empty_dtype_string(self):
    self.assertEqual(UOp.empty((3, 4), dtype="float32").dtype, dtypes.float32)
@@ -474,12 +456,11 @@ class TestUOpEmpty(unittest.TestCase):
      self.assertTrue(u.has_buffer_identity())

  def test_empty_direct_singleton_tuple_device(self):
-    u = UOp.empty((4,), dtype=dtypes.float32, device=("NULL:0",))
+    # regression: direct UOp.empty with a singleton-tuple device + axis must not trip .multi()'s tuple assert
+    u = UOp.empty((4,), dtype=dtypes.float32, device=("NULL:0",), axis=0)
    self.assertEqual((u.shape, u.device, u.axis), ((4,), "NULL", None))

 class TestTensorUOpCreation(unittest.TestCase):
-  def test_empty(self):
-    self.assertIs(_strip_unique(Tensor.empty(2, 3).uop), _strip_unique(UOp.empty(2, 3)))
  def test_full(self):
    self.assertIs(_strip_unique(Tensor.full((2, 3), 42).uop), _strip_unique(UOp.full((2, 3), 42)))
  def test_full_kwargs(self):
@@ -1,25 +1,24 @@
 import unittest
 from tinygrad import Tensor
-from tinygrad.nn.state import fs_store, fs_load

 class TestLoadStore(unittest.TestCase):
  def test_load_shape(self):
-    t = fs_load(Tensor(bytes(16)), 1024)
+    t = Tensor(bytes(16)).fs_load(1024)
    assert t.shape == (1024,), t.shape
    t.schedule_linear()

  def test_store_shape(self):
-    t = fs_store(Tensor.zeros(1024))
+    t = Tensor.zeros(1024).fs_store()
    assert t.shape == (16,), t.shape
    t.schedule_linear()

  def test_load_large_shape(self):
-    t = fs_load(Tensor(bytes(16)), 10_000_000)
+    t = Tensor(bytes(16)).fs_load(10_000_000)
    assert t.shape == (10_000_000,), t.shape
    t.schedule_linear()

  def test_store_large_shape(self):
-    t = fs_store(Tensor.zeros(10_000_000))
+    t = Tensor.zeros(10_000_000).fs_store()
    assert t.shape == (16,), t.shape
    t.schedule_linear()

@@ -2,7 +2,8 @@ import unittest, math
 import numpy as np
 from tinygrad import dtypes
 from tinygrad.uop.ops import UOp
-from tinygrad.codegen.decomp.transcendental import payne_hanek_reduction, cody_waite_reduction, frexp, rintk, pow2if
+from tinygrad.uop.decompositions import TRANSCENDENTAL_DTYPES, payne_hanek_reduction, cody_waite_reduction
+from tinygrad.uop.decompositions import frexp, rintk, xpow, xexp2, xlog2, trig_poly, pow2if
 from test.helpers import eval_uop

 class TestTranscendentalFunctions(unittest.TestCase):
@@ -70,5 +71,39 @@ class TestTranscendentalFunctions(unittest.TestCase):
    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -10), dtypes.float)), 2**-10)
    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -63), dtypes.float)), 2**-63)

+class TestTranscendentalVectorizedFunctions(unittest.TestCase):
+  # given a scalar and vectorized input, check that the fxn outputs have the same
+  # scalar_dtypes, args, ops, and vcount (only for vectorized input)
+
+  def _check_uop_vcount(self, u:tuple|UOp, vcount:int):
+    # check all UOps in u are vectorized with vcount
+    if isinstance(u, UOp):
+      assert u.dtype.vcount == vcount, f'expected {vcount=} but got {u.dtype.vcount=} for UOp\n{u=}'
+    [self._check_uop_vcount(x, vcount) for x in (u if isinstance(u, tuple) else u.src)]
+
+  def _check_uops_match(self, u1:tuple|UOp, u2:tuple|UOp):
+    # check all UOps in u1, u2 have the same scalar_dtype, args, ops
+    if isinstance(u1, UOp) and isinstance(u2, UOp):
+      assert u1.dtype.scalar() == u2.dtype.scalar(), f'expected {u1.dtype.scalar()=} but got {u2.dtype.scalar()=} for UOps\n{u1=}\n{u2}'
+      assert u1.arg == u2.arg or (math.isnan(u1.arg) and math.isnan(u2.arg)), f'expected {u1.arg=} but got {u2.arg=} for UOps\n{u1=}\n{u2}'
+      assert u1.op == u2.op, f'expected {u1.op=} but got {u2.op=} for UOps\n{u1=}\n{u2}'
+    [self._check_uops_match(x1, x2) for x1, x2 in zip((u1 if isinstance(u1, tuple) else u1.src), (u2 if isinstance(u2, tuple) else u2.src))]
+
+  def _test_vectorized(self, fxn, scalar_dtypes=TRANSCENDENTAL_DTYPES, vals=[-2,1.3,194], vcounts=[1,4,19]):
+    for scalar_dtype in scalar_dtypes:
+      for val in vals:
+        for vcount in vcounts:
+          in_scalar, in_vec = UOp.const(scalar_dtype, val), UOp.const(scalar_dtype.vec(vcount), val)
+          out_scalar, out_vec = fxn(in_scalar), fxn(in_vec)
+          self._check_uops_match(out_scalar, out_vec)
+          self._check_uop_vcount(out_vec, vcount)
+
+  def test_xpow(self): return self._test_vectorized(lambda x: xpow(x, x))
+  def test_xexp2(self): return self._test_vectorized(xexp2)
+  def test_xlog2(self): return self._test_vectorized(xlog2)
+  def test_payne_hanek_reduction(self): return self._test_vectorized(payne_hanek_reduction)
+  def test_cody_waite_reduction(self): return self._test_vectorized(cody_waite_reduction)
+  def test_trig_poly(self): return self._test_vectorized(lambda x: trig_poly(x, [0.0], [1.0]))
+
 if __name__ == '__main__':
  unittest.main()
@@ -2,8 +2,9 @@ import unittest, pytest
 from tinygrad import dtypes, Variable
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import DEBUG, Context
-from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, GroupOp, AxisType
+from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, track_rewrites, graph_rewrite, GroupOp, AxisType
 from tinygrad.uop.symbolic import sym
+from tinygrad.codegen.late.expander import expander
 from test.helpers import to_uops_list

 simple_pm = PatternMatcher([
@@ -21,23 +22,32 @@ def const_values(u:UOp):
 class TestGraphRewriteConst(unittest.TestCase):
  def test_gep_const(self):
    v1 = UOp.const(dtypes.int.vec(3), (0,1,2))
-    v2 = v1.index(1)
+    v2 = v1.gep(1)
    ret = graph_rewrite(v2, sym)
    self.assertEqual(ret.dtype, dtypes.int)
    self.assertEqual(ret.arg, 1)

+  def test_gep_const_single(self):
+    v1 = UOp.const(dtypes.int.vec(3), 4)
+    v2 = v1.gep(1)
+    ret = graph_rewrite(v2, sym)
+    self.assertEqual(ret.dtype, dtypes.int)
+    self.assertEqual(ret.arg, 4)
+
  def test_add_const(self):
-    v1 = UOp.const(dtypes.int, (0,1,2))
-    v2 = UOp.const(dtypes.int, (5,6,7))
+    v1 = UOp.const(dtypes.int.vec(3), (0,1,2))
+    v2 = UOp.const(dtypes.int.vec(3), (5,6,7))
    ret = graph_rewrite(v1+v2, sym)
    self.assertEqual(ret.op, Ops.STACK)
+    self.assertEqual(ret.dtype, dtypes.int.vec(3))
    self.assertEqual(const_values(ret), (5,7,9))

  def test_add_const_lose_v(self):
-    v1 = UOp.const(dtypes.int, (0,1,2))
-    v2 = UOp.const(dtypes.int, (2,1,0))
+    v1 = UOp.const(dtypes.int.vec(3), (0,1,2))
+    v2 = UOp.const(dtypes.int.vec(3), (2,1,0))
    ret = graph_rewrite(v1+v2, sym)
    self.assertEqual(ret.op, Ops.STACK)
+    self.assertEqual(ret.dtype, dtypes.int.vec(3))
    self.assertEqual(const_values(ret), (2,2,2))

 def xfail_broken_const_wraparound(fn):
@@ -254,7 +264,7 @@ class TestUOpGraph(unittest.TestCase):
    idx = UOp.const(dtypes.int, 0)
    ld = d0.load(idx, dtype=dtypes.float.vec(2))
    vec = UOp(Ops.STACK, dtypes.float.vec(2), (ld,))
-    x = vec.index(0)
+    x = UOp(Ops.GEP, dtypes.float, (vec, ), arg=0)
    alu = UOp(Ops.SQRT, dtypes.float, (x, ))
    out = UOp(Ops.STORE, dtypes.void, (d0, idx, alu))
    uops = to_uops_list([out])
@@ -277,27 +287,27 @@ class TestUOpGraph(unittest.TestCase):

    # possible
    val = d1.index(idx).load(dtype=dtypes.float.vec(4))
-    xyzw = tuple(val.index(i) for i in range(4))
+    xyzw = tuple(val.gep(i) for i in range(4))
    self.assertIs(_test_vec(xyzw).op, Ops.LOAD)

    # unaligned
    val = d1.index(idx).load(dtype=dtypes.float.vec(4))
-    wzyx = tuple(val.index(i) for i in reversed(range(4)))
+    wzyx = tuple(val.gep(i) for i in reversed(range(4)))
    self.assertIs(_test_vec(wzyx).op, Ops.STACK)

    # different_size
    val = d1.index(idx).load(dtype=dtypes.float.vec(2))
-    xy = tuple(val.index(i) for i in range(2))
+    xy = tuple(val.gep(i) for i in range(2))
    self.assertIs(_test_vec(xy+xy).op, Ops.STACK)
    val = d1.index(idx).load(dtype=dtypes.float.vec(4))
-    xy = tuple(val.index(i) for i in range(2))
+    xy = tuple(val.gep(i) for i in range(2))
    self.assertIs(_test_vec(xy, count=2).op, Ops.STACK)

    # different vals
    val1 = d1.index(idx).load(dtype=dtypes.float.vec(2))
    val2 = d2.index(idx).load(dtype=dtypes.float.vec(2))
-    xy1 = tuple(val1.index(i) for i in range(2))
-    xy2 = tuple(val2.index(i) for i in range(2))
+    xy1 = tuple(val1.gep(i) for i in range(2))
+    xy2 = tuple(val2.gep(i) for i in range(2))
    self.assertIs(_test_vec(xy1+xy2).op, Ops.STACK)

  def test_gep_vec_const_fold(self):
@@ -305,7 +315,7 @@ class TestUOpGraph(unittest.TestCase):
      consts = [UOp.const(dtypes.float, float(i)) for i in range(vec_size)]
      vec = UOp(Ops.STACK, dtypes.float.vec(vec_size), tuple(consts))
      with Context(SPEC=0):
-        uops = to_uops_list([vec.index(i) for i in range(vec_size)])
+        uops = to_uops_list([UOp(Ops.GEP, dtypes.float, (vec,), (i,)) for i in range(vec_size)])
        for uop, const in zip(uops, consts):
          self.assertEqual(uop, const)

@@ -571,6 +581,161 @@ class TestUOpGraph(unittest.TestCase):
    a = c.after(e)
    self.assertNotIn(r, a.ranges)

+@track_rewrites()
+def expander_rewrite(sink): return graph_rewrite(sink, sym + expander)
+
+class TestExpander(unittest.TestCase):
+  def test_expand_add_broadcast(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(x for x in range(4))),), ((1,4),))
+    sink = expander_rewrite(e1+3)
+    assert sink.op is Ops.UNROLL and len(const_values(sink.src[0])) == 4
+    self.assertTupleEqual(const_values(sink.src[0]), (3,4,5,6))
+
+  def test_contract_simple(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(x for x in range(4))),), ((1,4),))
+    con = UOp(Ops.CONTRACT, dtypes.int.vec(4), (e1,), ((1,4),))
+    sink = expander_rewrite(con)
+    self.assertEqual(sink.op, Ops.STACK)
+    self.assertTupleEqual(const_values(sink), (0,1,2,3))
+
+  def test_contract_axis_1(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(16), tuple(x for x in range(16))),), ((1,4),(2,4)))
+    con = UOp(Ops.CONTRACT, dtypes.int.vec(4), (e1,), ((1,4),))
+    sink = expander_rewrite(con)
+    vals = const_values(sink.src[0])
+    assert sink.op is Ops.UNROLL and len(vals) == 16 and sink.arg == ((2,4),)
+    assert sink.src[0].op is Ops.STACK
+    self.assertTupleEqual(vals[0:4], (0,4,8,12))
+    self.assertTupleEqual(vals[12:], (3,7,11,15))
+
+  def test_contract_axis_2(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(16), tuple(x for x in range(16))),), ((1,4),(2,4)))
+    con = UOp(Ops.CONTRACT, dtypes.int.vec(4), (e1,), ((2,4),))
+    sink = expander_rewrite(con)
+    vals = const_values(sink.src[0])
+    assert sink.op is Ops.UNROLL and len(vals) == 16 and sink.arg == ((1,4),)
+    assert sink.src[0].op is Ops.STACK
+    self.assertTupleEqual(vals[0:4], (0,1,2,3))
+    self.assertTupleEqual(vals[12:], (12,13,14,15))
+
+  def test_contract_axis_2_big(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(16), tuple(x for x in range(16))),), ((1,2),(2,2),(3,2),(4,2)))
+    con = UOp(Ops.CONTRACT, dtypes.int.vec(2), (e1,), ((2,2),))
+    sink = expander_rewrite(con)
+    assert sink.op is Ops.UNROLL and sink.arg == ((1, 2), (3, 2), (4, 2))
+    vals = const_values(sink.src[0])
+    self.assertTupleEqual(vals[0:2], (0,4))
+    self.assertTupleEqual(vals[12:14], (10,14))
+
+  def test_contract_multi_axis(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(16), tuple(x for x in range(16))),), ((1,2),(2,2),(3,2),(4,2)))
+    sink = expander_rewrite(UOp(Ops.CONTRACT, dtypes.int.vec(4), (e1,), ((3, 2), (2, 2))))
+    assert sink.op is Ops.UNROLL and sink.arg == ((1, 2), (4, 2))
+    self.assertTupleEqual(const_values(sink.src[0])[0:4], (0, 4, 2, 6))
+    sink = expander_rewrite(UOp(Ops.CONTRACT, dtypes.int.vec(4), (e1,), ((2, 2), (3, 2))))
+    assert sink.op is Ops.UNROLL and sink.arg == ((1, 2), (4, 2))
+    self.assertTupleEqual(const_values(sink.src[0])[0:4], (0, 2, 4, 6))
+
+  def test_contract_mid(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(8), tuple(x for x in range(8))),), ((1,2),(2,2),(3,2)))
+    con = UOp(Ops.CONTRACT, dtypes.int.vec(2), (e1,), ((2,2),))
+    sink = expander_rewrite(con)
+    assert sink.op is Ops.UNROLL and sink.arg == ((1,2),(3,2))
+    assert sink.src[0].op is Ops.STACK and len(const_values(sink.src[0])) == 8
+    self.assertTupleEqual(const_values(sink.src[0]), (0,2,1,3,4,6,5,7))
+
+  def test_contract_no_expand(self):
+    e1 = UOp.variable("i", 0, 10, dtype=dtypes.int)
+    con = UOp(Ops.CONTRACT, dtypes.int.vec(2), (e1,), ((2,2),))
+    sink = expander_rewrite(con)
+    assert sink.op is Ops.STACK and len(sink.src) == 2
+    assert sink.src[0] == sink.src[1]
+
+  def test_contract_half_expand(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(x for x in range(4))),), ((1,4),))
+    con = UOp(Ops.CONTRACT, dtypes.int.vec(8), (e1,), ((1,4), (2,2)))
+    sink = expander_rewrite(con)
+    vals = const_values(sink)
+    assert sink.op is Ops.STACK and len(vals) == 8
+    assert vals[0] == vals[1]
+    assert vals[0] != vals[2]
+    assert vals[6] == vals[7]
+
+  def test_expand_same_axis(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(x for x in range(4))),), ((1,4),))
+    e2 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(4*x for x in range(4))),), ((1,4),))
+    sink = expander_rewrite(e1+e2)
+    self.assertEqual(sink.op, Ops.UNROLL)
+    self.assertEqual(sink.src[0].op, Ops.STACK)
+    self.assertTupleEqual(const_values(sink.src[0]), (0,5,10,15))
+
+  def test_expand_different_axis(self, flip=False):
+    e1 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(4*x for x in range(4))),), ((1,4),))
+    e2 = UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(4), tuple(x for x in range(4))),), ((2,4),))
+    sink = expander_rewrite((e2+e1) if flip else (e1+e2))
+    vals = const_values(sink.src[0])
+    assert sink.op is Ops.UNROLL and len(vals) == 16
+    assert sink.arg == ((1, 4), (2, 4))
+    self.assertTupleEqual(vals, (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15))
+
+  def test_expand_different_axis_flip(self): self.test_expand_different_axis(True)
+
+  @unittest.skip("no longer supported")
+  def test_reduce_known_axis(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((1,4),))
+    sink = (3*e1).reduce(e1, arg=Ops.ADD)
+    sink = expander_rewrite(sink)
+    assert sink.op is Ops.CONST
+    self.assertEqual(sink.arg, 3*(0+1+2+3))
+
+  @unittest.skip("no longer supported")
+  def test_reduce_const(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((1,4),))
+    sink = UOp.const(dtypes.int, 3).reduce(e1, arg=Ops.ADD)
+    sink = expander_rewrite(sink)
+    assert sink.op is Ops.CONST
+    self.assertEqual(sink.arg, 3*4)
+
+  @unittest.skip("no longer supported")
+  def test_double_expand(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((2,4),))
+    e2 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, 4+x) for x in range(4)), ((2,4),))
+    e = UOp(Ops.UNROLL, dtypes.int, (e1, e2), ((1,2),))
+    sink = expander_rewrite(e)
+    assert sink.op is Ops.UNROLL and len(sink.src) == 8
+    assert sink.arg == ((1, 2), (2, 4))
+    self.assertListEqual([x.arg for x in sink.src], [0,1,2,3,4,5,6,7])
+
+  @unittest.skip("no longer supported")
+  def test_double_expand_reverse(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((1,4),))
+    e2 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, 4+x) for x in range(4)), ((1,4),))
+    e = UOp(Ops.UNROLL, dtypes.int, (e1, e2), ((2,2),))
+    sink = expander_rewrite(e)
+    assert sink.op is Ops.UNROLL and len(sink.src) == 8
+    assert sink.arg == ((1, 4), (2, 2))
+    self.assertListEqual([x.arg for x in sink.src], [0, 4, 1, 5, 2, 6, 3, 7])
+
+  @unittest.skip("no longer supported")
+  def test_double_expand_middle(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((1,2),(3,2)))
+    e2 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, 4+x) for x in range(4)), ((1,2),(3,2)))
+    e = UOp(Ops.UNROLL, dtypes.int, (e1, e2), ((2,2),))
+    sink = expander_rewrite(e)
+    assert sink.op is Ops.UNROLL and len(sink.src) == 8
+    assert sink.arg == ((1, 2), (2, 2), (3, 2))
+    self.assertListEqual([x.arg for x in sink.src], [0, 1, 4, 5, 2, 3, 6, 7])
+
+  # does this need to work?
+  @unittest.expectedFailure
+  @unittest.skip
+  def test_reduce_different_axis(self):
+    e1 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((1,4),))
+    e2 = UOp(Ops.UNROLL, dtypes.int, tuple(UOp.const(dtypes.int, x) for x in range(4)), ((2,4),))
+    sink = e1.reduce(e2, arg=Ops.ADD)
+    sink = expander_rewrite(sink)
+    print(sink)
+
 class TestReduceCollapse(unittest.TestCase):
  def test_multi_range_reduce_add(self):
    """Test that (x + y).reduce(r1, r2) distributes over multiple ranges"""
@@ -606,6 +771,45 @@ class TestMovementOps(unittest.TestCase):
    self.assertEqual(result.op, Ops.INDEX)
    self.assertEqual(result.src[0].op, Ops.RESHAPE)

+class TestLoadStoreFolding(unittest.TestCase):
+  def test_gated_load_gep_preserves_alt(self):
+    """Test that LOAD(GEP, alt) preserves alt value after rewrite"""
+    from tinygrad.codegen.late.devectorizer import load_store_folding
+    buf = UOp.param(0, dtypes.float.vec(4).ptr())
+    idx = UOp.const(dtypes.int, 0)
+    gate = UOp.const(dtypes.bool, True)
+    gated_index = buf.index(idx.valid(gate))
+    gep = gated_index.gep(0)
+    alt = UOp.const(dtypes.float, 42.0)
+    gated_load = gep.load(alt)
+    self.assertEqual(len(gated_load.src), 2)  # GEP + alt
+    result = graph_rewrite(gated_load, load_store_folding, name='test')
+    # After rewrite, should still have alt value preserved
+    self.assertEqual(result.op, Ops.GEP)
+    inner_load = result.src[0]
+    self.assertEqual(inner_load.op, Ops.LOAD)
+    self.assertEqual(len(inner_load.src), 2)  # INDEX + alt
+
+  def test_gated_load_ptrcat_preserves_alt(self):
+    """Test that LOAD(PTRCAT, alt) preserves alt value after rewrite"""
+    from tinygrad.codegen.late.devectorizer import load_store_folding
+    buf1 = UOp.param(0, dtypes.float.ptr())
+    buf2 = UOp.param(1, dtypes.float.ptr())
+    idx = UOp.const(dtypes.int, 0)
+    idx1 = buf1.index(idx)
+    idx2 = buf2.index(idx)
+    ptrcat = UOp(Ops.PTRCAT, dtypes.float.ptr().vec(2), (idx1, idx2))
+    alt = UOp.const(dtypes.float.vec(2), 42.0)
+    gated_load = ptrcat.load(alt)
+    self.assertEqual(len(gated_load.src), 2)  # PTRCAT + alt
+    result = graph_rewrite(gated_load, load_store_folding, name='test')
+    # After rewrite, should be CAT of LOADs, each preserving alt
+    self.assertEqual(result.op, Ops.VCAT)
+    for inner_load in result.src:
+      self.assertEqual(inner_load.op, Ops.LOAD)
+      self.assertEqual(len(inner_load.src), 2)  # INDEX + alt
+    self.assertEqual(inner_load.src[1].arg, 42.0)  # alt value preserved
+
 class TestConstBufferize(unittest.TestCase):
  def test_const_bufferize_with_ranges(self):
    """Test that CONST.BUFFERIZE with ranges is folded correctly.
@@ -825,26 +825,6 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable((x//10)*10 + x%10, 0, 119, "(a*10+(a+b//5)//2*10+(b+a*5)%10)")
    self.helper_test_variable((x//10)*2 + (x//5)%2, 0, 23, "(a*3+b//5)")

-  def test_div_mod_recombine_merged_quotient(self):
-    # recombine finds the quotient base//div even when stored merged as base0//(d0*div), including with an offset
-    x = Variable("x", 0, 199)
-    self.helper_test_variable(((x//3)%4)*2 + ((x//12)%5)*8, 0, 38, "x//3%20*2")  # nested-merged quotient
-    self.helper_test_variable(((x//3 + 1)%4) + ((x+3)//12)*4, 1, 67, "x//3+1")   # offset-merged quotient
-
-  def test_div_mod_recombine_negative_div(self):
-    # partial recombine only needs d>0, div can be negative: (x%div) + ((x//div)%d)*div -> x%(div*d)
-    x = Variable("x", 0, 199)
-    self.helper_test_variable(x%(-3) + ((x//(-3))%5)*(-3), -14, 0, "x%-15")
-
-  def test_div_mod_recombine_shifted_quotient(self):
-    # when vmin<0 blocks const reduction on the mod side, the quotient is stored const-shifted: (x-50)//3 -> (x+1)//3 - 17.
-    # recombine only needs a quotient of some b congruent to base mod div, so the shift folds into the result
-    x = Variable("x", 0, 100)
-    y = Variable("y", 0, 99)
-    self.helper_test_variable((x-50)%3 + ((x-50)//3)*3, -50, 50, "(x+-50)")                # shifted literal quotient
-    self.helper_test_variable((x-50)%3 + (((x-50)//3)%5)*3, 0, 14, "((x+-50)%15)")         # shift inside the partial's mod
-    self.helper_test_variable(((y-50)//5)%4 + ((y-50)//20)*4, -10, 9, "(y//5+-10)")        # merged and shifted
-
  def test_div_mod_recombine_in_additive_sum(self):
    x = Variable("x", 0, 31)
    y = Variable("y", 0, 5)
@@ -319,7 +319,7 @@ class TestVminVmaxVConst(unittest.TestCase):
    d1 = UOp.param(1, dtypes.int.ptr())
    idx = UOp.const(dtypes.int, 0)
    val = UOp(Ops.LOAD, dtypes.int.vec(2), (d1.index(idx).cast(dtypes.int.vec(2).ptr()),))
-    uop = (val // 32).index(0)
+    uop = (val // 32).gep(0)
    self.assertEqual(uop.vmin, -67108864)
    self.assertEqual(uop.vmax, 67108863)

@@ -5,7 +5,7 @@ from tinygrad.tensor import Tensor
 from tinygrad.helpers import Timing, Context, cdiv
 from tinygrad.dtype import dtypes, ConstFloat  # noqa: F401
 from tinygrad.device import Device
-from tinygrad.uop.ops import Ops, ParamArg, UOp, UPat, exec_alu  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
+from tinygrad.uop.ops import Ops, UOp, UPat, exec_alu
 from tinygrad.uop.spec import spec_shared
 from tinygrad.uop.symbolic import sym
 from test.helpers import eval_uop, to_uops_list
@@ -318,6 +318,10 @@ class TestUOpStr(unittest.TestCase):
    vec = UOp(Ops.STACK, dtypes.int.vec(4), tuple(UOp.const(dtypes.int, x) for x in range(4)))
    assert str(eval(str(vec))) == str(vec)

+  def test_device_arg(self):
+    device = UOp(Ops.DEVICE, arg="CL")
+    assert str(eval(str(device))) == str(device)
+
  def test_reduceop_arg(self):
    sum_uop = Tensor.empty(32, 32).sum().uop
    assert str(eval(str(sum_uop))) == str(sum_uop)
@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor
-from tinygrad.helpers import GlobalCounters
+from tinygrad.helpers import GlobalCounters, DEV
 from tinygrad.engine.realize import compile_linear, estimate_uop
 from tinygrad.codegen import to_program
 from tinygrad.renderer import Estimates
@@ -90,6 +90,7 @@ class TestUOpsStatsMatmulHalf(unittest.TestCase):
    expected_ops = N ** 3 * 2
    self.assertEqual(expected_ops, GlobalCounters.global_ops)

+  @unittest.skipIf(DEV.arch=="INTEL", "intel gets 524288 != 524352")
  def test_bigger_matmul_half(self): self.test_simple_matmul_half(64)

  def test_batched_matmul_half(self, N=16):
@@ -165,37 +166,26 @@ class TestStatsOptimized(unittest.TestCase):
  @classmethod
  def setUpClass(cls):
    cls.ast_gemm = (Tensor.empty(N, N) @ Tensor.empty(N, N)).schedule_linear().src[-1].src[0]
-    cls.ast_gemm_half = (Tensor.empty(N, N, dtype=dtypes.half) @ Tensor.empty(N, N, dtype=dtypes.half)).schedule_linear().src[-1].src[0]
    cls.ast_reduce = (Tensor.empty(N*N).sum()).schedule_linear().src[-1].src[0]

-  def check_gemm(self, p:UOp, extra_flops=0, half=False):
+  def check_gemm(self, p:UOp, extra_flops=0):
    est = p.src[0].arg.estimates
    print(p.arg.name, est.ops, est.mem, est.lds)
    self.assertEqual(est.ops, 2*N*N*N + extra_flops)  # N**3 mulaccs
-    self.assertEqual(est.mem, 3*N*N*(2 if half else 4)) # 3 NxN mats with floats
+    self.assertEqual(est.mem, 3*N*N*4) # 3 NxN mats with floats

  def test_gemm(self):
    p = to_program(replace_opts(self.ast_gemm, []), renderer=Device[Device.DEFAULT].renderer)
    self.check_gemm(p)
    self.assertEqual(p.src[0].arg.estimates.lds, 2*N*N*N*4 + 4*N*N)

-  @unittest.skip("fails locally on AMD")
-  def test_gemm_tc_unroll_half(self):
-    try:
-      p = to_program(replace_opts(self.ast_gemm_half, [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)]),
-                      renderer=Device[Device.DEFAULT].renderer)
-    except KernelOptError:
-      raise unittest.SkipTest("no tensor cores")
-    print(p.src[2].arg)
-    self.check_gemm(p, half=True)
-
  def test_gemm_tc_unroll(self):
    try:
      p = to_program(replace_opts(self.ast_gemm, [Opt(OptOps.TC, 0, (-1, 0, 1)), Opt(OptOps.UNROLL, 0, 2)]),
                      renderer=Device[Device.DEFAULT].renderer)
    except KernelOptError:
      raise unittest.SkipTest("no tensor cores")
-    print(p.src[2].arg)
+    print(p.src[3].arg)
    self.check_gemm(p)

  # this is a good lesson about why UPCASTing is a good idea
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
George HotzandGitHub	f34efbe9f3	Merge branch 'master' into simpler_image	2026-06-24 16:53:31 -07:00
geohot	0f40b9c3b5	ish	2026-06-24 16:14:15 -07:00
geohot	06bd987dcb	fix half	2026-06-24 15:43:05 -07:00
geohot	1cbe0ba1f2	handle demote to non image	2026-06-24 15:16:00 -07:00
geohot	548154c8f2	image index	2026-06-24 14:34:44 -07:00
geohot	2d1d91dd99	fix indexing	2026-06-24 14:27:34 -07:00
geohot	1da09eb82d	//4	2026-06-24 14:07:01 -07:00
geohot	5825ffbbcf	simpler image (no stupid ai)	2026-06-24 14:03:41 -07:00