bump amd firmware

* fused_qkv_rope in UOp try 2

* dont need that

* less

.github/actions/setup-tinygrad/action.yml

@@ -41,12 +41,12 @@ inputs:
     description: "Install LLVM?"
     required: false
     default: 'false'
-  tinydreno:
-    description: "Install tinydreno"
+  qemu:
+    description: "Install qemu?"
     required: false
     default: 'false'
-  qemu:
-    description: "Install qemu"
+  ninja:
+    description: "Install ninja?"
     required: false
     default: 'false'
 runs:
@@ -134,7 +134,7 @@ runs:
 
     # ******************* apt *******************
     - name: Setup apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
       shell: bash
       run: |
         sudo mkdir -p /var/cache/apt/archives
@@ -162,7 +162,7 @@ runs:
         echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-20 main" | sudo tee /etc/apt/sources.list.d/llvm.list
 
     - name: Compute Package List + Hash
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
       id: apt-pkgs
       shell: bash
       run: |
@@ -187,25 +187,29 @@ runs:
         if [[ "${{ inputs.qemu }}" == "true" ]]; then
           pkgs+=" qemu-user-static"
         fi
+        # **** ninja ****
+        if [[ "${{ inputs.ninja }}" == "true" ]]; then
+          pkgs+=" ninja-build"
+        fi
 
         echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
         echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"
 
     - name: Cache apt (PR)
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true') && github.event_name == 'pull_request'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name == 'pull_request'
       uses: actions/cache/restore@v5
       with:
         path: /var/cache/apt/archives/
         key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
     - name: Cache apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true') && github.event_name != 'pull_request'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name != 'pull_request'
       uses: actions/cache@v5
       with:
         path: /var/cache/apt/archives/
         key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
 
     - name: Run apt Update + Install
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
       shell: bash
       run: |
         sudo apt -qq update || true
@@ -277,12 +281,6 @@ runs:
       shell: bash
       run: brew install llvm@20
 
-    # *** tinydreno ***
-    - name: Install tinydreno (linux)
-      if: inputs.tinydreno == 'true' && runner.os == 'Linux'
-      shell: bash
-      run: sudo curl -fL https://github.com/sirhcm/tinydreno/raw/refs/heads/master/libllvm-qcom.so -o /usr/lib/libllvm-qcom.so
-
     # *** OpenCL ***
     - name: Install rusticl
       if: inputs.opencl == 'true'

.github/workflows/autogen.yml

@@ -35,7 +35,7 @@ jobs:
         key: 'autogen'
         amd: 'true'
         llvm: 'true'
-        pydeps: 'pyyaml mako'
+        deps: 'autogen'
     - name: Install autogen support packages
       run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev liburing-dev
     - name: Regenerate autogen files

.github/workflows/benchmark.yml

@@ -108,10 +108,6 @@ jobs:
     - name: Setup (NV)
       if: ${{ matrix.dev == 'NV' }}
       run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
-    - name: Symlink models and datasets
-      run: |
-        mkdir -p weights
-        ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -129,10 +125,6 @@ jobs:
       # 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
 
@@ -182,10 +174,6 @@ jobs:
       # slow on metal
       if: ${{ matrix.dev != 'METAL' }}
       run: time BENCHMARK_LOG=cifar DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
-    - name: Run full CIFAR training steps w 6 GPUS
-      # only run on machines with multiple gpus
-      if: ${{ matrix.dev != 'METAL' }}
-      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
@@ -231,11 +219,6 @@ jobs:
       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
 
@@ -285,6 +268,58 @@ jobs:
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
+  multigpubenchmark:
+    name: Multi-GPU Benchmarks (DEV=${{ matrix.dev }})
+    runs-on: [self-hosted, "${{ matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
+    strategy:
+      fail-fast: false
+      matrix:
+        dev: ['AMD', 'NV']
+    timeout-minutes: 60
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    env:
+      DEV: ${{ matrix.dev }}
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup (AMD)
+      if: ${{ matrix.dev == 'AMD' }}
+      run: |
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py amd rmmod
+        ./extra/hcq/hcq_smi.py amd kill_pids
+    - name: Setup (NV)
+      if: ${{ matrix.dev == 'NV' }}
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        mkdir -p extra/datasets
+        ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
+        ln -s /raid/datasets/imagenet extra/datasets/imagenet
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: reset process replay
+      run: python3 test/external/process_replay/reset.py
+    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
+      run: BENCHMARK_LOG=llama3_beam_4gpu JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
+    - name: Run full CIFAR training steps w 6 GPUS
+      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
+      run: BENCHMARK_LOG=resnet_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
+    - name: Run 10 MLPerf Bert training steps (6 gpu)
+      # TODO: remove BERT_LAYERS once scheduler is fast
+      run: BENCHMARK_LOG=bert_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
   tests:
     name: Tests (DEV=${{ matrix.dev }})
     runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
@@ -541,7 +576,7 @@ jobs:
     - name: openpilot run_pickle big_driving_supercombo
       run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_run_pickle RUN_PICKLE=1 PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py - openpilot.pkl
     - name: Test copy speeds
-      run: SIZE=64e6 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
 
   driverbenchmarks:
     name: PCI Driver Benchmark (DEV=${{ matrix.dev }})

.github/workflows/docs.yml

@@ -8,7 +8,7 @@ permissions:
   contents: write
 jobs:
   deploy:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     steps:
       - uses: actions/checkout@v6
       - name: Configure Git Credentials

.github/workflows/platform.yml

@@ -179,35 +179,3 @@ jobs:
       - name: Run test_tiny
         shell: bash
         run: python -m pytest -n=auto test/test_tiny.py --durations=20
-
-
-  qcomclcompiletests:
-    name: Compile-only (QCOM CL)
-    runs-on: ubuntu-24.04-arm
-    timeout-minutes: 15
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v6
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: compile-qcomcl
-          deps: testing_unit
-          tinydreno: 'true'
-      - name: Set env
-        shell: bash
-        run: printf "DEV=NULL:QCOMCL:a630\nNULL_ALLOW_COPYOUT=1" >> $GITHUB_ENV
-      - name: Run test_ops
-        shell: bash
-        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
-          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
-          python -m pytest -n=auto test/backend/test_ops.py --durations=20
-      - name: Run test_ops (IMAGE)
-        shell: bash
-        env:
-          IMAGE: 1
-          DEV: "NULL:QCOMCL:a630,IMAGE_PITCH_ALIGNMENT=64"
-        run: |
-          DEBUG=4 python test/backend/test_ops.py TestOps.test_gemm | grep read_imagef
-          python -m pytest -n=auto test/backend/test_ops.py --durations=20

.github/workflows/python-publish.yml

@@ -10,7 +10,7 @@ on:
 
 jobs:
   deploy:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     steps:
     - uses: actions/checkout@v6
     - name: Set up Python

.github/workflows/szdiff.yml

@@ -10,7 +10,7 @@ concurrency:
 jobs:
   checkbranch:
     name: Check PR Branch status
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     outputs:
       branchstat: ${{ steps.brstat.outputs.stat}}
     steps:
@@ -44,7 +44,7 @@ jobs:
     permissions:
       contents: read
       pull-requests: write
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     needs: checkbranch
     if: needs.checkbranch.outputs.branchstat == 'false'
     steps:
@@ -87,7 +87,7 @@ jobs:
     name: Core Library Line Difference
     permissions:
       pull-requests: write
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-24.04
     needs: checkbranch
     if: needs.checkbranch.outputs.branchstat == 'true'
     steps:

.github/workflows/test.yml

@@ -21,7 +21,7 @@ concurrency:
 jobs:
   docs:
     name: Docs
-    runs-on: &linux ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 10
     env:
       CHECK_OOB: 0
@@ -31,8 +31,7 @@ jobs:
     - name: Setup Environment
       uses: ./.github/actions/setup-tinygrad
       with:
-        deps: docs
-        pydeps: "capstone torch"
+        deps: "docs testing_minimal"
     - name: Build wheel and show size
       run: |
         uv build --wheel
@@ -61,7 +60,7 @@ jobs:
 
   torchbackend:
     name: Torch Backend Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     steps:
     - name: Checkout Code
@@ -73,10 +72,7 @@ jobs:
         deps: testing_unit
         pydeps: "pillow torchvision expecttest"
         llvm: 'true'
-    - name: Install ninja
-      run: |
-        sudo apt update || true
-        sudo apt install -y --no-install-recommends ninja-build
+        ninja: 'true'
     - name: Test ResNet-18
       run: DEBUG=2 python3 extra/torch_backend/example.py
     - name: Test one op in torch tests
@@ -86,9 +82,26 @@ jobs:
     - name: Custom tests
       run: DEV=CPU:LLVM GPUS=4 TINY_BACKEND=1 python3 -m pytest -nauto extra/torch_backend/test.py extra/torch_backend/test_inplace.py extra/torch_backend/test_multigpu.py extra/torch_backend/test_kernel_fusion.py --durations=20
 
+  torchbackendtrain:
+    name: Torch Backend Training
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    timeout-minutes: 15
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: torch-backend-pillow-torchvision-et-pt
+        deps: testing_unit
+        llvm: 'true'
+        ninja: 'true'
+    - name: Test beautiful_mnist in torch with TINY_BACKEND
+      run: STEPS=20 DEV=CPU TARGET_EVAL_ACC_PCT=90.0 MAX_BUFFER_SIZE=0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+
   bepython:
     name: Python Backend
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     steps:
     - name: Checkout Code
@@ -126,7 +139,7 @@ jobs:
 
   linter:
     name: Linters
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 10
 
     steps:
@@ -157,7 +170,7 @@ jobs:
 
   nulltest:
     name: Null Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
 
     steps:
@@ -191,7 +204,7 @@ jobs:
 
   unittest:
     name: Unit Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
 
     steps:
@@ -228,7 +241,7 @@ jobs:
       matrix:
         group: [1, 2]
     name: SPEC=2 (${{ matrix.group }})
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     steps:
     - name: Checkout Code
@@ -244,7 +257,7 @@ jobs:
 
   fuzzing:
     name: Fuzzing
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 10
     steps:
     - name: Checkout Code
@@ -260,7 +273,7 @@ jobs:
 
   testopenclimage:
     name: CL IMAGE Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     steps:
       - name: Checkout Code
@@ -280,7 +293,7 @@ jobs:
 
   testopenpilot:
     name: openpilot Compile Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     steps:
       - name: Checkout Code
@@ -309,7 +322,7 @@ jobs:
 
   testonnxcpu:
     name: ONNX (CPU) Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 20
 
     steps:
@@ -328,7 +341,7 @@ jobs:
 
   testoptim:
     name: Optimization Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 20
     steps:
       - name: Checkout Code
@@ -360,7 +373,7 @@ jobs:
 
   testllm:
     name: Test LLM
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     env:
       CHECK_OOB: 0
@@ -387,7 +400,7 @@ jobs:
 
   testmodels:
     name: Models
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     steps:
       - name: Checkout Code
@@ -407,7 +420,7 @@ jobs:
 
   testdsp:
     name: Linux (DSP)
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     steps:
     - name: Checkout Code
@@ -435,7 +448,7 @@ jobs:
           - 'WEBGPU'
 
     name: Linux (DEV=${{ matrix.dev }})
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 20
     steps:
     - name: Checkout Code
@@ -461,7 +474,7 @@ jobs:
 
   testamdasm:
     name: AMD ASM IDE
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 20
     env:
       DEV: MOCKKFD+AMD
@@ -507,7 +520,7 @@ jobs:
 
   hcq2:
     name: hcq2
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 5
     steps:
       - name: Checkout Code
@@ -521,16 +534,15 @@ jobs:
       - name: Run HCQ2 tests
         run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/test_tiny.py
       - name: Run HCQ2 multi-device tests
-        run: |
-          HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_multitensor.py \
-            TestMultiTensor.test_simple_add TestMultiTensor.test_shard_reduce \
-            TestMultiTensor.test_backward_sum TestMultiTensor.test_matmul_shard_0_0
+        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest -n=auto test/backend/test_multitensor.py
       - name: Run HCQ2 JIT tests
         run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_jit.py
+      - name: Run HCQ2 unit tests
+        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest test/device/test_hcq2.py
 
   testmockam:
     name: Linux (am)
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     env:
       DEV: MOCKPCI+AMD
@@ -566,7 +578,7 @@ jobs:
         arch: [gfx1100, gfx1201, gfx950]
 
     name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     env:
       DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
@@ -589,7 +601,7 @@ jobs:
         if: ${{ matrix.backend == 'amd' && matrix.arch == 'gfx950' }}
         run: PYTHONPATH=. DEV=NULL:HIP:gfx950 MXFP4=1 LLAMA_LAYERS=2 BENCHMARK=3 NULL_ALLOW_COPYOUT=1 NO_HIPCC=1 ROCM_PATH=/opt/rocm JITBEAM=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh
       - name: Run pytest (amd)
-        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM test/opt/test_tensor_cores.py --durations=20
+        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM --durations=20
       - name: Run disk copy tests
         run: python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk
       - name: Run TRANSCENDENTAL math
@@ -604,7 +616,7 @@ jobs:
         backend: [ptx, nv]
 
     name: Linux (${{ matrix.backend }})
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 20
     env:
       FORWARD_ONLY: 1
@@ -638,10 +650,17 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        backend: [ir3, nak]
-    name: Compile-only (${{ matrix.backend }})
-    runs-on: *linux
+        dev:
+          - 'NULL:IR3:a630'
+          - 'NULL:QCOMCL:a630'
+          - 'NULL:NAK:sm_120'
+    name: Compile-only (DEV=${{ matrix.dev }})
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
+    env:
+      NULL_ALLOW_COPYOUT: 1
+      DEV: ${{ matrix.dev }}${{ contains(matrix.dev, 'a630') && ',IMAGE_PITCH_ALIGNMENT=64' || '' }}
+      IMAGE: ${{ contains(matrix.dev, 'a630') && '1' || '0' }}
     steps:
       - name: Checkout Code
         uses: actions/checkout@v6
@@ -650,21 +669,14 @@ jobs:
         with:
           key: compile-${{ matrix.backend }}
           deps: "testing_unit mesa"
-      - name: Set env
+          qemu: ${{ contains(matrix.dev, 'QCOMCL') }}
+      - name: Test IMAGE
         shell: bash
-        run: printf "NULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'DEV=NULL:IR3:a630' || matrix.backend == 'nak' && 'DEV=NULL:NAK:sm_120' }}" >> $GITHUB_ENV
+        if: contains(matrix.dev, 'a630')
+        run: DEBUG=7 python3 test/backend/test_ops.py TestOps.test_gemm | grep isam
       - name: Run test_ops
         shell: bash
         run: |
           python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
           DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
           python -m pytest -n=auto test/backend/test_ops.py --durations=20
-      - name: Run test_ops (IMAGE)
-        if: matrix.backend == 'ir3'
-        shell: bash
-        env:
-          IMAGE: 1
-          DEV: "NULL:IR3:a630,IMAGE_PITCH_ALIGNMENT=64"
-        run: |
-          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_gemm | grep image_load
-          python -m pytest -n=auto test/backend/test_ops.py --durations=20

README.md

@@ -140,7 +140,7 @@ Documentation along with a quick start guide can be found on the [docs website](
 ```python
 from tinygrad import Tensor
 
-x = Tensor.eye(3)
+x = Tensor.eye(3).clone()  # clone to make it a buffer
 y = Tensor([[2.0,0,-2.0]])
 z = y.matmul(x).sum()
 z.backward()

docs/kimi_k3_mi350.md

@@ -1,224 +0,0 @@
-# Kimi K3 on 8× MI350X
-
-This branch targets text generation directly from the official `moonshotai/Kimi-K3` checkpoint at `/raid/weights/kimi-k3`. It intentionally ignores the vision tower and multimodal projector. The checkpoint remains in its official 96-shard format; the loader never converts, rewrites, or creates a second 1.56 TB copy.
-
-The checked TP8 layout consumes 196.78 GB (183.27 GiB) of text weights per GPU. The compressed MLA cache adds 28.99 GB (27 GiB) per GPU at the full 1,048,576-token context, leaving approximately 62.23 GB of each nominal 288 GB MI350X for execution buffers and allocator overhead. Start much smaller.
-
-## Resume the current optimization session
-
-Work on branch `kimi_slop`. It was cleanly rebased onto `origin/kimi_slop` commit `553bdf68e` on 2026-08-10. The retained K3 commits after that base are `1b3732a6e`, `c6ac4961d`, `1d8620471`, `224bac031`, `f53f0e7e7`, and `2b1b8c22a`; verify the current hashes with `git log` because a later rebase may rewrite them. Before starting any benchmark, check that the worktree is clean and that no model process remains:
-
-```sh
-git status --short --branch
-git log --oneline --decorate -10
-pgrep -af 'tinygrad.llm.cli|benchmark_kimi_k3' || true
-```
-
-The active acceptance target is **more than 100 tok/s decode, more than 200 tok/s prefill, and less than 180 seconds cold startup** on TP8/gfx950. None is currently met. The authoritative official-checkpoint baseline is 389.84 seconds startup, 38.65 tok/s prefill, and 6.25 tok/s decode. The 1.56 TB checkpoint has a measured 6.9 GB/s single-XFS-NVMe read ceiling, giving a roughly 227-second physical cold-read floor; meeting the startup target therefore also requires a faster storage path, not only loader code.
-
-Use the fake-weight, one-layer loop for development. Do not repeatedly load the official checkpoint while optimizing:
-
-```sh
-DEV=AMD python extra/benchmark_kimi_k3_fake.py --mode attention --iterations 30
-DEV=AMD python extra/benchmark_kimi_k3_fake.py --mode block --iterations 30
-PROFILE=1 DEV=AMD python extra/benchmark_kimi_k3_fake.py --mode block --iterations 5
-```
-
-The clean retained baseline is about 0.630 ms per attention layer and 1.37 ms per complete block, with fake initialization taking about 0.9/2 seconds respectively after the rebase. Since K3 has 93 sequential blocks, a 100 tok/s projection requires at most approximately 0.108 ms per complete block. Only run another 96-shard official validation after a candidate produces a large whole-block gain, remains finite and deterministic, and passes a direct numerical comparison. Test one candidate at a time and remove failed experiments before moving on.
-
-The immediate bottleneck is launch and synchronization granularity: an official four-token decode profile contained 6,304 kernel events, while packed expert work was only a small fraction of total GPU time. Continue with whole-component or whole-block fusion/replay work, not isolated expert microkernels. The latest fake-loop A/B retested the previously rejected dual gate/up and weighted-down MFMA prototypes: 1.374 ms baseline versus 1.375 ms fused, so they were removed again. A fused whole-core KDA recurrence was also slower in the exact fake attention gate (0.665 versus 0.633 ms) and must not be restored unchanged.
-
-Preserve these invariants when official validation resumes: use `/raid/weights/kimi-k3` directly, keep all 96 shards byte-for-byte untouched, run only one model process, begin at context 128, verify all eight devices are `gfx950`, and preserve the first failure instead of retrying over it. The most recent preserved official failure from a rejected KDA experiment was the invalid sequence `[198, 163840, 163840, 163840]`; token 163840 is outside the valid vocabulary. The retained path before that experiment produced deterministic in-range replay.
-
-After a synthetic candidate passes, run correctness and performance in this order: NULL gfx950 compile coverage, focused tests with `-n12` where supported, TP8 fake numerical comparison, official context-128 deterministic tokens, load/prefill/decode timing, and then context admission at 4K, 32K, 131K, and 262K. Run `python -m mypy tinygrad/` and `python -m ruff check .` when those tools are installed. Read `tinygrad/viz/README.md` before inspecting rewrite or device profiles.
-
-## Before renting the machine
-
-- Keep the existing 96 shards in `/raid/weights/kimi-k3`; no additional model-sized free space is required. Leave ordinary headroom for logs and temporary files.
-- The host should have roughly 3 TB RAM, in line with AMD's MI350X platform guidance. The loader itself is streaming and must not need checkpoint-sized RAM.
-- Use a recent kernel/ROCm stack supported by the host vendor, although tinygrad uses its own AMD userspace driver when `DEV=AMD`.
-- Clone this exact commit/branch and keep the official checkpoint directory separate from the repository.
-
-Validate the existing directory without modifying it:
-
-```sh
-python examples/kimi_k3_prepare.py /raid/weights/kimi-k3 --context 4096
-```
-
-For a metadata-only preflight, place the official `config.json` and `model.safetensors.index.json` in a directory and run:
-
-```sh
-python examples/kimi_k3_prepare.py /raid/weights/kimi-k3 --metadata-only
-```
-
-## Hardware admission checks
-
-Do these before loading weights. Stop if any device is missing or reports a different architecture.
-
-```sh
-lspci -d 1002:75a0
-amd-smi list
-DEV=AMD DEBUG=2 python - <<'PY'
-from tinygrad import Device
-for i in range(8):
-  dev = Device[f"AMD:{i}"]
-  print(i, dev.arch)
-PY
-```
-
-Expected architecture: `gfx950` on all eight devices. Then run the small TP8 graph tests:
-
-```sh
-python -m pytest test/unit/test_llm_k3.py test/null/test_kimi_k3.py -q -n12
-DEV=NULL:HIP:gfx950 NULL_ALLOW_COPYOUT=1 python -m pytest \
-  test/unit/test_llm_k3.py::TestKimiK3::test_chunked_recurrent_generate -q -n1
-DEV=AMD python examples/kimi_k3_smoke.py --devices 8
-```
-
-The last two commands are deliberately small. They compile CDNA4 kernels and then exercise the complete TP8 topology without loading the checkpoint.
-
-For performance iteration, use the exact-width fake-weight harness before another official load:
-
-```sh
-DEV=AMD python extra/benchmark_kimi_k3_fake.py --mode attention --iterations 20
-DEV=AMD python extra/benchmark_kimi_k3_fake.py --mode block --iterations 20
-```
-
-It retains K3's 7,168-wide residual stream, 12,288-wide KDA state, 96 heads, 128×128 recurrent matrices, TP8 layouts, top-k 16 routing, packed MXFP4 expert shapes, collectives, and decode JIT, but uses one layer and 16 fake experts. Fake attention weights initialize in about 0.9 seconds and the full block in about 3 seconds. The retained path measured 0.630 ms per fake attention layer and 1.367 ms per complete fake block, projecting about 7.87 tok/s across 93 identical blocks versus 6.25 tok/s for the official heterogeneous model. Treat this as a candidate admission benchmark, not a correctness substitute for official weights.
-
-## First official load
-
-Start at a short context so cache allocation and compilation are bounded. The loader reads disk-backed safetensors, TP-shards every destination before realizing it, and drops each source shard/projection immediately afterward.
-
-```sh
-/usr/bin/time -v env DEV=AMD DEBUG=1 python -m tinygrad.llm.cli \
-  --model /raid/weights/kimi-k3 --devices 8 --max_context 128 </dev/null 2>&1 | tee kimi-k3-load.log
-```
-
-Watch host RAM, swap, HBM, temperatures, and XGMI traffic from a second terminal. Do not start with a one-million-token cache. If loading fails, preserve the first exception and the last loader progress line; do not retry with a larger host-side cache.
-
-## Correctness and performance sequence
-
-1. Load with context 128 and generate one token.
-2. Repeat a fixed prompt twice and confirm token-for-token deterministic greedy output.
-3. Compare the first several greedy tokens against the official Transformers implementation at temperature zero.
-4. Benchmark decode only after two warm-up tokens.
-5. Benchmark prefill at 128, 512, 2K, and 8K tokens. Increase context only while HBM and compile time remain healthy.
-6. Use `VIZ=1` plus `python -m tinygrad.viz.cli` to inspect kernels; use `VIZ=2` only for short SQTT captures because it adds overhead.
-
-Example decode benchmark:
-
-```sh
-DEV=AMD DEBUG=1 python -m tinygrad.llm.cli --model /raid/weights/kimi-k3 \
-  --devices 8 --max_context 4096 --warmup --benchmark 20
-```
-
-## MI350X validation results (2026-08-10)
-
-The official directory was audited in place: 96 shards, 497,220 indexed tensors, 497,052 language tensors, and 1,560,860,324,864 total bytes. All eight devices reported `gfx950`. No checkpoint file was converted, copied, or modified, and every model run used a single process. The actual text tower is 1,559,965,606,912 bytes; its checked TP8 layout is 196,784,397,312 bytes per GPU.
-
-The preserved first full-checkpoint error was an `A_log` shape mismatch, `(128,) -> (96, 1)`. K3 stores one decay value per 128-wide KDA channel, not one per head. The loader now keeps this field replicated and applies the official channel-wise broadcast. A numerical unit test covers the distinction from the older head-wise Kimi Linear behavior.
-
-Load speed was fixed before generation. The original loader opened thousands of individual expert tensors and independently realized eight strided TP slices. The MI350 path now does the following without changing the checkpoint:
-
-- parses safetensor headers selectively, constructing disk-backed tensors only for the 2,460 non-expert entries consumed by that pass instead of materializing metadata objects for every expert entry twice;
-- copies contiguous axis-zero shards and replicas directly into their final device buffers;
-- reads a replicated tensor once and fans it out over XGMI instead of issuing eight identical direct reads (14.31 GB less RAID traffic);
-- stages an inner-axis tensor once and schedules all eight TP slices together;
-- reads each layer's contiguous 15.72 GB expert region once, reorders its lexicographically stored expert records on GPU 0, and realizes all six packed/scale destinations together;
-- retains only final MultiBuffer identities, drops the reorder graph, and flushes the 15.72 GB staging allocation before the next layer.
-
-One real expert layer leaves exactly 1,965,293,568 bytes resident on each GPU and zero bytes in the GPU-0 allocator cache. Complete context-128 loads measured 527.20 seconds before the final staging cleanup and 490.05/489.59 seconds afterward. Peak host RSS for the unprofiled correctness run was 2.11 GiB with zero swap. RAID variability produced later loads from 489.06 to 532.85 seconds.
-
-The selective-metadata and bounded-GC pass reduced non-expert loading from 125.77 to 57.77 seconds. A subsequent full official context-128 load completed in 411.49 seconds, 78.10 seconds (16.0%) faster than the 489.59-second baseline. It read the 96 shards in place with 1,049,688 KiB peak host RSS and zero swap; no weight payload was converted, copied, or modified. Direct-I/O probes measured approximately 6.9 GB/s aggregate for both one and eight concurrent 1 GiB reads. At that rate the 1.56 TB checkpoint has a roughly 227-second cold-read lower bound, so this RAID cannot meet a true cold sub-three-minute startup regardless of loader overhead.
-
-Expert staging graphs are acyclic and are released by reference counting after each layer, so the loader now suppresses unnecessary cyclic-collector scans only around that loop and restores its prior state on every exit. A quiet context-128 load then completed in 391.54 seconds, 30.14 seconds (7.1%) faster than the immediately preceding 421.68-second run, with 1.04 GiB peak RSS and zero swap, although storage variability contributes to run-to-run timing. The host used for these measurements actually mounts `/raid` from one 3.5 TB XFS NVMe, not a multi-drive RAID; shard 28 has 218 extents and live reads fell to roughly 160 MB/s there. This storage layout, plus the physical checkpoint size, remains the limiting cold-start constraint. The weights were not defragmented, copied, or modified.
-
-The fixed XTML prompt `Reply with exactly: OK` encodes to 93 tokens. After excluding the cold JIT capture from replay comparison, two greedy runs produced the identical eight-token sequence:
-
-```text
-[9545, 59991, 10580, 14404, 9545, 59991, 9545, 59991]
-```
-
-At context 128, steady prefill was 14.32 seconds (6.49 tok/s) and eight-token decode was 2.27 seconds (3.53 tok/s, 283.3 ms/token). The same first tokens remained stable at every admitted context. These rates are much lower than the planning estimates below and should be treated as the current measured baseline.
-
-The retained gfx950 serving pass enables the validated wave64 recurrent prefill kernel with 128-token chunks, uses exact BF16 decode projections, combines the routed/shared final TP partials into one collective, and tiles four adjacent packed-expert outputs during multi-token execution. On the same 93-token prompt, two replay trials produced the identical sequence `[198, 92652, 220, 80225]`. Prefill replay measured 2.418--2.482 seconds (37.47--38.46 tok/s), and eight-token decode measured 1.294 seconds (6.18 tok/s, 161.81 ms/token). Peak RSS was 2.77 GiB with zero swap. The packed prefill tile changes floating-point reduction order: direct official-layer comparison against the original kernel had maximum differences of 0.015625 for gate and 0.0078125 for down, and the end-to-end greedy sequence was stable across replay.
-
-A subsequent gfx950 decode pass split the 7,168-wide replicated BF16 projections across eight waves per 16 output channels and used CDNA4 BF16 MFMA, with one FP32 LDS reduction at the end. It is enabled only for batch-one/token-one replicated projections whose dimensions satisfy the hardware tile; prefill, the FP32 router, and the output-sharded 12,288-wide KDA gate remain unchanged. The official retained path uses it for MLA q-a/kv-a and KDA f-a. Isolated TP8 measurements improved replicated 128/576-output projections by about 16--18%; applying it to the already output-sharded KDA gate was slower and was rejected. Random-shape comparison against the generic graph had maximum/mean absolute BF16 differences of 2.0/0.1114 because the split changes reduction order. Against a serial FP32 accumulation rounded once to BF16, the 7,168-to-1,536 kernel was bit-exact in the tested sample.
-
-The final official context-128 validation loaded in 389.84 seconds with 2.71 GiB peak RSS and zero swap. Two replay trials produced the identical four-token sequence `[198, 59675, 9817, 12519]`; prefill remained 2.406 seconds (38.65 tok/s), while eight-token decode improved to 1.280 seconds (6.25 tok/s, 160.00 ms/token). A one-wave MFMA variant and a full-wave fused decode recurrence were both rejected: the former delivered 6.02 tok/s, and the latter 6.179 tok/s, while both changed the greedy sequence without a useful speed gain.
-
-A final load-first experiment increased the disk-to-HBM io_uring queue depth from one to the 32 existing bounded 2 MiB staging buffers. On a direct 1 GiB read from fragmented shard 28 it measured 6.834 GB/s versus 6.832 GB/s for the original path, so the change was rejected. The subsequent unmodified official 96-shard load completed in 389.48 seconds, confirming both the prior result and the single-NVMe lower bound. Peak RSS was 2.75 GiB with zero swap.
-
-Two direct packed-expert MFMA prototypes were also rejected after that load. A fused gate/up kernel was about 29% faster in isolation at the TP8-local shape, and a routed-down kernel which combined projection, probability weighting, and route reduction measured 1.45 ms versus 2.42 ms in isolation. End-to-end, however, stable replay produced `[198, 2338, 2127, 148297]`, prefill measured 38.87 tok/s, and decode measured 6.263 tok/s. That is indistinguishable from the retained 38.65/6.25 tok/s path while changing floating-point reduction order, so neither kernel was retained.
-
-A whole-core KDA decode experiment fused convolution, Q/K normalization, channel decay, recurrence, RMS normalization, output gating, and four persistent state updates. Its raw kernel replayed in about 109 microseconds per local KDA layer and matched a one-step synthetic reference within `9.77e-4` output and `8.13e-4` state maximum error. The exact-width fake-layer gate caught that it was slower than the retained attention path (0.665 versus 0.633 ms/layer). The already-running official validation was stopped after its first invalid greedy sequence, `[198, 163840, 163840, 163840]`, where 163840 is outside the checkpoint's vocabulary. The kernel was rejected and removed.
-
-| Maximum context | Load | Short-prompt replay | Result |
-|---:|---:|---:|---|
-| 128 | 489.59s | 14.32s | stable 8-token replay |
-| 4,096 | 489.06s | 14.32s | stable replay, zero swap |
-| 32,768 | 532.85s | 14.33s | stable replay, zero swap |
-| 131,072 | 520.91s | 14.37s | stable first token, zero swap |
-| 262,144 | 497.34s | 14.41s | stable first token, zero swap |
-
-These are maximum-context/cache admission tests with the same 93-token prompt, not full-length 32K/131K/262K prefills. The full cache allocation path was exercised, but filling those contexts remains a separate long-running throughput test.
-
-Runtime profiling bracketed four steady decode tokens. It recorded 6,304 kernel events and about 474--478 ms of summed GPU work across the eight devices inside a roughly 1.5-second profiled wall interval. The packed `mxfp4_expert_linear_wave64` kernels accounted for only about 22.5 ms summed; the largest families were small 1,792-wide reductions. This identifies launch/synchronization granularity as the immediate MI350 bottleneck rather than packed-weight bandwidth. `JIT_BATCH_SIZE=64` produced the same original 3.53 tok/s as 32. A gfx950 fused MXFP8 QDQ experiment was bit-exact but slower on the real device (about 95 microseconds versus 57--64 microseconds), so it was rejected. Combining the routed and shared final TP partials removed one collective per routed decode layer and helped raise unprofiled decode to 6.18 tok/s, but the remaining sequential launch boundaries still dominate.
-
-The checkpoint's bundled Transformers code was used as the architectural reference for channel decay and tensor mapping. A full independent Transformers/vLLM token comparison was not run on this host because the required `compressed_tensors`/serving backend is not installed; deterministic tinygrad replay and the numerical KDA, loader-layout, NULL gfx950 compile, and real TP8 smoke tests are the completed correctness gates.
-
-## Known hardware-only gate
-
-The correctness path now consumes packed MXFP4 expert weights directly on gfx950 with a wave64 software-decode kernel, so it does not create selected-expert BF16 weight expansions. MXFP8 activation quantization is still emulated. tinygrad has gfx950/CDNA4 BF16 and FP8 matrix-core support, but this branch does not yet have a hardware-validated native MXFP4×MXFP8 expert GEMM. Expect the first run to be a correctness bring-up, not production throughput. Capture profiles on MI350X before changing the representation: native FP4 work cannot be validated faithfully on the available gfx1100 cards.
-
-Recurrent prefill is fused. The gfx950 wave-parallel kernel was compared directly with the portable graph at the official per-GPU shape through 128 tokens: maximum core/state differences remained below `8e-6`/`1e-6`, outputs were finite, and replay was about 2.7 ms versus about 8 ms for the portable kernel in the isolated test. Full K3 therefore uses 128-token recurrent chunks on gfx950. Chunk size remains part of the numerical configuration because different reduction orders can select different final greedy tokens.
-
-The following serving changes apply to the official K3 path: recurrent-state reset graph capture, direct AMD scalar readback without rebuilding a scheduler graph, materialized gate/up boundaries, separate greedy decode JITs, K3's uncorrected routed probability semantics, gfx950 KDA Q/K/V and exact BF16 partial projections, one combined routed/shared final collective, a gfx950 greedy output-head kernel, the wave64 packed-expert path, and the multi-token four-output packed tile. Software MXFP8 remains in use.
-
-After hardware admission on MI350X, profile before porting those kernels. The likely implementation order is:
-
-1. A native packed MXFP4×MXFP8 grouped expert GEMM using CDNA4 matrix instructions.
-2. A wave64/MFMA KDA Q/K/V decode projection.
-3. Combined routed/shared down-projection TP partials so each layer performs one XGMI all-reduce.
-4. A CDNA4 output-head matvec and router matvec if they remain visible in the profile.
-
-Every port needs a direct numerical comparison with the generic graph and an end-to-end greedy-token comparison before performance measurements. The wave64 packed-expert kernel has compile coverage through `NULL:HIP:gfx950`; numerical and performance validation still require real MI350X hardware. None of the remaining gfx11-only kernels should be enabled on gfx950 by changing only the architecture guard.
-
-## MI350X performance expectation
-
-Treat the first rental as bring-up, not a guaranteed throughput run. The loader reads every official expert tensor once into a transient GPU-0 staging buffer (at most one packed projection), then redistributes TP8 slices over the GPU fabric; it does not generate files or require checkpoint-sized host RAM. A reasonable planning range for the full text model on eight MI350X cards is 3–8 minutes to stream and TP-shard the 1.56 TB checkpoint, 150–400 tok/s for initial short/medium prefill, and 25–60 tok/s decode with the software packed-expert path. After a native CDNA4 MXFP4×MXFP8 grouped expert kernel, wave64/MFMA recurrent projections, and XGMI collective tuning, 500+ tok/s prefill and roughly 80–150 tok/s decode are plausible targets. These ranges are engineering estimates, not measurements.
-
-The nominal HBM bandwidth is not the main uncertainty: eight MI350X devices have enough aggregate bandwidth for K3's active weights. Utilization is limited by 93 sequential layers, small routed projections, and synchronization after TP input-sharded projections. Record actual HBM and XGMI counters before deciding whether the next port should target matrix instructions or collective count.
-
-The official checkpoint also contains MoonViT-V2 and multimodal projector weights. They are skipped by the text loader. Image input remains a separate implementation and validation task.
-
-## Local TP4 performance baseline
-
-The pre-rental benchmark uses the converted `Kimi-Linear-48B-A3B-Instruct-MXFP4-v2` checkpoint on four gfx1100 GPUs. It is a useful regression test for the KDA/MLA/MoE text path, not a projection of K3 throughput on MI350X.
-
-```sh
-DEV=AMD JIT_BATCH_SIZE=64 python extra/benchmark_kimi.py \
-  /raid/models/Kimi-Linear-48B-A3B-Instruct-MXFP4-v2 \
-  --devices 4 --max-context 128 --prompt-tokens 32 --decode-tokens 32 --chunk-size 32
-```
-
-Results from 2026-08-10:
-
-- load from RAID: 44.28s for the 29.27 GB checkpoint
-- first 32-token prefill includes roughly 10s of compilation/capture
-- steady fresh-prompt prefill replay: 0.118s, 270.20 tok/s
-- steady context-32 decode replay: 101.82 tok/s, 9.82 ms/token
-- peak host RSS: 729.9 MiB; swap was not used
-
-The load, prefill, and decode targets are all met in the bounded prompt-32 run. Decode improved from 23.03 tok/s to 101.82 tok/s. The retained greedy output was checked across 32 decode steps; rejected half-wave and unrounded recurrent reductions were faster but diverged and eventually collapsed to a repeated token.
-
-Fully warmed HTTP serving was also measured with `--max_context 4096`. Startup, including weight load, capture, and replay of both serving shapes, took 113.73s. After a two-turn cache test, the first aligned 64-token request reported 271 tok/s prefill and 101 tok/s decode over 64 generated tokens. A 99-token prompt reported 254 tok/s prefill and 99 tok/s decode; decode falls slightly as MLA context grows.
-
-Recurrent serving uses only the captured 32-token prefill graph and captured single-token graph. Warmup uses two consecutive chunks so both initial and nonzero-position prefill execution are ready before the socket opens. A prompt tail shorter than 32 tokens runs through the single-token graph instead of compiling a new static shape, so no request-time JIT capture is required. Exact extensions reuse recurrent and KV state—the live second turn logged `in: 18 + 15`—while divergent prompts reset both safely. Very short prompts can report less than 200 aggregate prefill tok/s because fixed reset and single-token costs dominate; aligned and medium/long prompts exercise the 200+ tok/s prefill path.
-
-Four 7900 XTX cards provide 96 GB aggregate VRAM and about 3.84 TB/s aggregate physical memory bandwidth. Their nominal aggregate vector FP16 rate is about 245.6 TFLOP/s, or about 492 TFLOP/s through matrix instructions. Kimi Linear activates roughly 3.107B parameters per token; a simple active-weight accounting gives approximately 4.05 GB/token and an optimistic bandwidth-only ceiling near 948 tok/s. The measured decode rate is much lower because this MoE decode workload is a collection of small matrix-vector operations plus PCIe collectives, not one ideal streaming kernel.
-
-The generic loader currently rereads logical TP shards and accounts for roughly 227 GB of disk traffic for a TP4 load. RAID bandwidth hides that inefficiency locally, but a direct one-pass shard loader remains worthwhile before slow remote storage is used. It was not retained here because the attempted direct-shard graph exposed an unresolved scheduler/renderer edge; correctness and bounded memory take priority over avoiding the redundant reads.
-
-Different chunk sizes can choose a different final token because their matrix kernels use different floating-point reduction orders. Each measured shape was repeatable between cold and captured execution. For official K3 validation, compare logits/tokens against the reference at one fixed chunk size and greedy settings rather than requiring bitwise agreement between performance shapes.

examples/convert_kimi_linear.py

@@ -1,9 +0,0 @@
-import argparse
-from tinygrad.llm.kimi import convert_kimi
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser(description="Convert official Kimi-Linear-48B-A3B BF16 weights to tinygrad MXFP4/BF16")
-  parser.add_argument("source", help="downloaded moonshotai/Kimi-Linear-48B-A3B-Instruct directory")
-  parser.add_argument("output", help="output directory")
-  args = parser.parse_args()
-  convert_kimi(args.source, args.output)

examples/kimi_k3_prepare.py

@@ -1,27 +0,0 @@
-#!/usr/bin/env python3
-"""Cheap preflight for an official moonshotai/Kimi-K3 checkout. Does not load model weights."""
-import argparse, json, pathlib, shutil
-from tinygrad.llm.kimi_k3 import KIMI_K3_TP8_BYTES_PER_GPU, audit_kimi_k3_checkpoint
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("model_dir", type=pathlib.Path)
-  parser.add_argument("--metadata-only", action="store_true", help="permit absent weight shards")
-  parser.add_argument("--context", type=int, default=4096, help="context length used for the memory estimate")
-  args = parser.parse_args()
-  stats = audit_kimi_k3_checkpoint(args.model_dir, require_shards=not args.metadata_only)
-  if not 1 <= args.context <= 1_048_576: raise ValueError("--context must be between 1 and 1048576")
-
-  # K3 has 24 MLA layers. Each token stores the 512-value compressed latent plus 64 RoPE values in BF16.
-  per_gpu_weights = KIMI_K3_TP8_BYTES_PER_GPU
-  mla_cache = 24 * args.context * (512 + 64) * 2
-  hbm = 288_000_000_000
-  print(json.dumps(stats, indent=2))
-  print(f"exact text weights/GPU under this TP8 layout: {per_gpu_weights/1e9:.2f} GB ({per_gpu_weights/2**30:.2f} GiB)")
-  print(f"replicated MLA cache/GPU at {args.context:,} tokens: {mla_cache/1e9:.2f} GB ({mla_cache/2**30:.2f} GiB)")
-  print(f"nominal MI350X headroom before runtime buffers: {(hbm-per_gpu_weights-mla_cache)/1e9:.2f} GB")
-  if not args.metadata_only:
-    usage = shutil.disk_usage(args.model_dir)
-    print(f"filesystem free space: {usage.free/1e9:.2f} GB")
-
-if __name__ == "__main__": main()

examples/kimi_k3_smoke.py

@@ -1,27 +0,0 @@
-#!/usr/bin/env python3
-"""Run a reduced, architecture-complete K3 prefill/decode on tensor-parallel devices."""
-import argparse, time
-from tinygrad import Tensor, Device, dtypes, nn
-from tinygrad.llm.kimi_k3 import _shard_kimi_k3, kimi_k3_smoke_config
-from tinygrad.llm.model import Transformer
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("--devices", type=int, default=8)
-  args = parser.parse_args()
-  if args.devices not in (1, 2, 4, 8): raise ValueError("the K3 admission smoke test supports 1, 2, 4, or 8 devices")
-  devices = tuple(f"AMD:{i}" for i in range(args.devices))
-  model = Transformer(kimi_k3_smoke_config())
-  for name,value in nn.state.get_state_dict(model).items():
-    fill = 127 if name.endswith("weight_scale") else 0
-    dtype = value.dtype if value.dtype is dtypes.uint8 else dtypes.bfloat16
-    value.replace(Tensor.full(value.shape, fill, dtype=dtype, device="CPU"))
-  _shard_kimi_k3(model, devices)
-  temperature = Tensor([0.0], device=devices)
-  for label,tokens,start in (("prefill", [[1, 2]], 0), ("decode", [[3]], 2), ("decode replay", [[4]], 3)):
-    begin = time.perf_counter()
-    out = model(Tensor(tokens, dtype=dtypes.int32, device=devices), start, temperature).realize()
-    for device in devices: Device[device].synchronize()
-    print(f"{label}: shape={out.shape}, {time.perf_counter()-begin:.3f}s")
-
-if __name__ == "__main__": main()

examples/mlperf/model_train.py

@@ -1742,8 +1742,8 @@ def train_gptoss():
   )
 
   for p in optim.params:
-    grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
-    p.grad = p.zeros_like(dtype=grad_dtype).contiguous()
+    p.grad = p.zeros_like(dtype=dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype).contiguous()
+    if getattr(p, "_zero2", False): p.grad = optim.optimizers[0]._zero_shard(p.grad)
   grads = [p.grad for p in optim.params]
 
   from extra.gemm.cdna_asm_gemm import _mx_block_scale

examples/mlperf/models/gpt_oss.py

@@ -12,7 +12,7 @@ from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb
 from extra.llama_kernels.rmsnorm import rmsnorm
-from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8
+from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8, asm_gemm, can_use_asm_gemm
 from extra.gemm.moe_gemm import grouped_mx_gemm
 from extra.gemm.moe_routing import route, dispatch, combine
 
@@ -146,6 +146,7 @@ class GPTOSS:
       return w_q, w_e8.is_param_(False)
     if moe:
       qs = [_one(*shape[1:]) for _ in range(shape[0])]
+      for q in qs: q[0]._zero2 = True  # grad arrives sharded on the expert axis under ZeRO-2 (moe_gemm)
       return [q[0] for q in qs], [q[1] for q in qs]
     return _one(*shape)
 
@@ -182,12 +183,14 @@ class GPTOSS:
     xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
     xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)  # (B,N,H,D)/(B,N,KV,D)
 
-    if sliding:
-      attn = self._sliding_attention(xq, xk, xv, sinks)
-    elif getenv("HK_FLASH_ATTENTION"):
+    fa_saves = []
+    if getenv("HK_FLASH_ATTENTION"):
       from extra.thunder.amd.fa import flash_attention
-      attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks)
+      attn, _, l_vec = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks, window=self.sliding_window if sliding else 0)
       attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
+      fa_saves = [xq, xk, xv, l_vec]
+    elif sliding:
+      attn = self._sliding_attention(xq, xk, xv, sinks)
     else:
       xqm = xq.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep, self.head_dim).permute(0, 2, 3, 1, 4)
       xkm, xvm = xk.permute(0, 2, 1, 3).unsqueeze(2), xv.permute(0, 2, 1, 3).unsqueeze(2)
@@ -199,7 +202,7 @@ class GPTOSS:
       attn = (w @ xvm).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, self.n_heads * self.head_dim)
 
     out = matmul_mx(attn, wo, wo_scale) + wo_bias
-    return out, [x_normed, rrms, attn]
+    return out, [x_normed, rrms, attn] + fa_saves
 
   def feed_forward(self, x:Tensor, *, ffn_norm:Tensor, gate:Tensor, gate_bias:Tensor,
                    w_gate_up:Tensor, w_gate_up_scale:Tensor, w_gate_up_bias:Tensor,
@@ -220,6 +223,7 @@ class GPTOSS:
       z = grouped_mx_gemm(_pad_cols(y.cast(dtypes.bfloat16)), (w_down, w_down_scale), r.off)[:, :dim] \
           + (onehot @ w_down_bias.float()).cast(dtypes.bfloat16)
       out = combine(z, r, inp.shape[0], self.experts_per_tok).reshape(bsz, seqlen, dim)
+      return out, [x_normed, rrms, xg, h, y, z, r.weights, r.dest_row, r.off]
     else:
       thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
       weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)
@@ -263,7 +267,11 @@ class GPTOSS:
                         w_down=self.w_down[i], w_down_scale=self.w_down_scale[i], w_down_bias=self.w_down_bias[i])
       h, *_ = self.run_layer(h, freqs_cis, mask_full, i % 2 == 0, attn_kwargs, ffn_kwargs, save=save)
 
-    logits = self.norm(h) @ self.output.T
+    h_normed = self.norm(h)
+    pad = (-self.dim) % 256
+    h_padded, w_padded = h_normed.pad((None, None, (0, pad))), self.output.pad(((0, 0), (0, pad)))
+    if ASM_GEMM and can_use_asm_gemm(h_padded, w_padded.T): logits = asm_gemm(h_padded, w_padded.T)
+    else: logits = h_normed @ self.output.T
     return logits
 
 def _get_pads(uop:UOp) -> list[UOp]:

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh

@@ -26,7 +26,7 @@ export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
 export SPLIT_W13=${SPLIT_W13:-0}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
 
-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="float32"
 export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
@@ -46,7 +46,7 @@ export DATA_SEED=${DATA_SEED:-5760}
 export JITBEAM=${JITBEAM:-3}
 export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
 
-export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+export FAKEDATA=${FAKEDATA:-$([[ "$DEV" == NULL:* ]] && echo 1 || echo 0)} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
   export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
 fi

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh

@@ -1,8 +1,8 @@
 #!/usr/bin/env bash
 
 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
+export ROCM_PATH=${ROCM_PATH:-/opt/rocm-7.1.1}
+export PATH="$ROCM_PATH/bin:$PATH"
 export DEV=${DEV:-AMD}
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -26,7 +26,7 @@ export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
 export SPLIT_W13=${SPLIT_W13:-0}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
 
-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="float32"
 export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/gpt_oss/implementations/tinybox_8xMI350X/dev_beam.sh

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

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/gpt_oss/implementations/tinybox_8xMI350X/dev_run.sh

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

extra/benchmark_kimi.py

@@ -1,79 +0,0 @@
-#!/usr/bin/env python3
-"""Benchmark Kimi-Linear load, prefill, and decode on its TP4 checkpoint."""
-import argparse, resource, time
-from tinygrad import Device, TinyJit
-from tinygrad.helpers import profile_marker
-from tinygrad.llm.kimi import load_kimi
-
-def sync(devices:int) -> None:
-  for i in range(devices): Device[f"AMD:{i}"].synchronize()
-
-def timed_next(gen, devices:int) -> tuple[int, float]:
-  begin = time.perf_counter()
-  token = next(gen)
-  sync(devices)
-  return token, time.perf_counter()-begin
-
-def fresh_generate(model, prompt:list[int], chunk_size:int):
-  # Force recurrent/KV state reset so repeated runs and chunk sweeps measure the entire prompt,
-  # rather than silently reusing the prefix cached by the previous measurement.
-  model._cached_tokens = [-1] * len(prompt)
-  return model.generate(prompt.copy(), chunk_size=chunk_size)
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("model", help="converted Kimi-Linear-48B-A3B MXFP4-v2 directory")
-  parser.add_argument("--devices", type=int, default=4)
-  parser.add_argument("--max-context", type=int, default=128)
-  parser.add_argument("--prompt-tokens", type=int, default=32)
-  parser.add_argument("--decode-tokens", type=int, default=8)
-  parser.add_argument("--chunk-size", type=int, default=32)
-  parser.add_argument("--sweep-chunks", help="comma-separated prefill chunk sizes; uses the fastest for decode")
-  args = parser.parse_args()
-  if args.prompt_tokens < 1 or args.prompt_tokens + args.decode_tokens + 1 > args.max_context:
-    raise ValueError("prompt and decode tokens must fit within --max-context")
-
-  begin = time.perf_counter()
-  model = load_kimi(args.model, max_context=args.max_context, devices=args.devices)
-  sync(args.devices)
-  print(f"load: {time.perf_counter()-begin:.3f}s", flush=True)
-
-  prompt = [1] + [1000+i%1000 for i in range(args.prompt_tokens-1)]
-  chunks = [int(x) for x in args.sweep_chunks.split(",")] if args.sweep_chunks else [args.chunk_size]
-  if any(x < 1 or x > args.prompt_tokens for x in chunks): raise ValueError("prefill chunks must be between 1 and --prompt-tokens")
-  timings:list[tuple[float, int]] = []
-  prefill_jits:dict[int, TinyJit] = {}
-  for chunk in chunks:
-    # Recurrent prefill has a static token dimension. Give each swept shape its own capture;
-    # the rollout JIT remains shared and independently benchmarks chunk 1/decode.
-    if chunk != 1: model.prefill_jit = TinyJit(model.forward)
-    cold = fresh_generate(model, prompt, chunk)
-    first, cold_prefill = timed_next(cold, args.devices)
-    print(f"chunk {chunk}: cold prefill {cold_prefill:.3f}s, token={first}", flush=True)
-    warm = fresh_generate(model, prompt, chunk)
-    warm_first, prefill = timed_next(warm, args.devices)
-    if first != warm_first: raise RuntimeError(f"chunk {chunk} is not repeatable: cold={first}, warm={warm_first}")
-    timings.append((prefill, chunk))
-    if chunk != 1: prefill_jits[chunk] = model.prefill_jit
-    print(f"chunk {chunk}: prefill {prefill:.3f}s ({args.prompt_tokens/prefill:.3f} tok/s), token={first}", flush=True)
-
-  prefill, best_chunk = min(timings)
-  if best_chunk != 1: model.prefill_jit = prefill_jits[best_chunk]
-  warm = fresh_generate(model, prompt, best_chunk)
-  first, replay_prefill = timed_next(warm, args.devices)
-  _, cold_decode = timed_next(warm, args.devices)
-  _, capture_decode = timed_next(warm, args.devices)
-  print(f"selected chunk: {best_chunk}; prefill replay {replay_prefill:.3f}s "
-        f"({args.prompt_tokens/replay_prefill:.3f} tok/s), token={first}", flush=True)
-  print(f"cold decode: {cold_decode:.3f}s", flush=True)
-  print(f"capture decode: {capture_decode:.3f}s", flush=True)
-  profile_marker("kimi decode steady start")
-  begin = time.perf_counter()
-  output = [next(warm) for _ in range(args.decode_tokens)]
-  sync(args.devices)
-  decode = time.perf_counter()-begin
-  profile_marker("kimi decode steady end")
-  print(f"decode: {decode:.3f}s ({args.decode_tokens/decode:.3f} tok/s, {decode/args.decode_tokens*1e3:.3f} ms/tok), output={output}", flush=True)
-  print(f"peak RSS: {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024:.1f} MiB", flush=True)
-
-if __name__ == "__main__": main()

extra/benchmark_kimi_k3.py

@@ -1,83 +0,0 @@
-#!/usr/bin/env python3
-"""Bounded correctness and load/prefill/decode benchmark for the official TP8 Kimi K3 checkpoint."""
-import argparse, resource, time
-
-from tinygrad import Device
-from tinygrad.helpers import profile_marker
-from tinygrad.llm.cli import KimiK3Template, SimpleTokenizer
-from tinygrad.llm.kimi_k3 import load_kimi_k3, load_kimi_tokenizer_data
-
-def sync(devices:int) -> None:
-  for i in range(devices): Device[f"AMD:{i}"].synchronize()
-
-def fresh_generate(model, prompt:list[int], chunk_size:int):
-  # Never reuse a prefix or recurrent state across correctness/benchmark trials.
-  model._cached_tokens = [-1] * len(prompt)
-  return model.generate(prompt.copy(), chunk_size=chunk_size, temperature=0.0)
-
-def timed_next(gen, devices:int) -> tuple[int, float]:
-  begin = time.perf_counter()
-  token = next(gen)
-  sync(devices)
-  return token, time.perf_counter()-begin
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("model", help="official unmodified Kimi K3 checkpoint directory")
-  parser.add_argument("--devices", type=int, default=8)
-  parser.add_argument("--max-context", type=int, default=128)
-  parser.add_argument("--prompt", default="Reply with exactly: OK")
-  parser.add_argument("--stable-tokens", type=int, default=8)
-  parser.add_argument("--decode-tokens", type=int, default=8)
-  parser.add_argument("--chunk-size", type=int, default=128)
-  args = parser.parse_args()
-
-  begin = time.perf_counter()
-  model = load_kimi_k3(args.model, max_context=args.max_context, devices=args.devices)
-  sync(args.devices)
-  load_time = time.perf_counter()-begin
-  print(f"load: {load_time:.3f}s", flush=True)
-
-  normal, special, bos, eos = load_kimi_tokenizer_data(args.model)
-  tok = SimpleTokenizer(normal, special, "kimi-k2", bos_id=bos, eos_id=eos, eot_id=eos)
-  rendered = KimiK3Template().render(messages=[{"role":"user", "content":args.prompt}], add_generation_prompt=True)
-  prompt = tok.encode(rendered)
-  needed = len(prompt) + max(args.stable_tokens, args.decode_tokens+3)
-  if needed > args.max_context: raise ValueError(f"prompt and output need {needed} tokens but max context is {args.max_context}")
-  print(f"prompt: {len(prompt)} tokens, chunk={args.chunk_size}", flush=True)
-
-  sequences:list[list[int]] = []
-  # TinyJit executes uncaptured once, captures the second call, and replays from the third call.
-  # Compare two replay paths rather than capture numerics/timing against replay.
-  for trial in range(4):
-    gen = fresh_generate(model, prompt, args.chunk_size)
-    sequence:list[int] = []
-    prefill = 0.0
-    for step in range(args.stable_tokens):
-      token, elapsed = timed_next(gen, args.devices)
-      sequence.append(token)
-      if step == 0: prefill = elapsed
-    if trial >= 2: sequences.append(sequence)
-    label = ("uncaptured warmup", "capture warmup", "stable trial 1", "stable trial 2")[trial]
-    print(f"{label}: prefill={prefill:.3f}s "
-          f"({len(prompt)/prefill:.3f} tok/s), tokens={sequence}", flush=True)
-  if sequences[0] != sequences[1]: raise RuntimeError(f"greedy output is not repeatable: {sequences}")
-  print(f"stable text: {tok.decode(sequences[0])!r}", flush=True)
-
-  gen = fresh_generate(model, prompt, args.chunk_size)
-  profile_marker("kimi k3 steady prefill start")
-  first, prefill = timed_next(gen, args.devices)
-  profile_marker("kimi k3 steady prefill end")
-  warmup = [timed_next(gen, args.devices)[0] for _ in range(2)]
-  profile_marker("kimi k3 steady decode start")
-  begin = time.perf_counter()
-  output = [next(gen) for _ in range(args.decode_tokens)]
-  sync(args.devices)
-  decode = time.perf_counter()-begin
-  profile_marker("kimi k3 steady decode end")
-  print(f"prefill replay: {prefill:.3f}s ({len(prompt)/prefill:.3f} tok/s), token={first}", flush=True)
-  print(f"decode after warmup {warmup}: {decode:.3f}s ({args.decode_tokens/decode:.3f} tok/s, "
-        f"{decode/args.decode_tokens*1e3:.3f} ms/tok), output={output}", flush=True)
-  print(f"peak RSS: {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024:.1f} MiB", flush=True)
-
-if __name__ == "__main__": main()

extra/benchmark_kimi_k3_fake.py

@@ -1,84 +0,0 @@
-#!/usr/bin/env python3
-"""Fast exact-shape K3 KDA/layer benchmark using bounded fake weights instead of the 1.56 TB checkpoint."""
-from __future__ import annotations
-import argparse, statistics, time
-from dataclasses import replace
-
-from tinygrad import Device, Tensor, TinyJit, dtypes, nn
-from tinygrad.helpers import profile_marker
-from tinygrad.llm.kimi_k3 import kimi_k3_config
-from tinygrad.llm.model import GatedDeltaNetBlock
-
-def tp_axis(name:str) -> int|None:
-  if "ffn_gate_exps.weight" in name or "ffn_up_exps.weight" in name: return 1
-  if "ffn_gate_exps.weight_scale" in name or "ffn_up_exps.weight_scale" in name: return 1
-  if "ffn_down_exps.weight" in name or "ffn_down_exps.weight_scale" in name: return 2
-  if name.endswith(("ffn_gate_shexp.weight", "ffn_up_shexp.weight")): return 0
-  if name.endswith(("ffn_down_shexp.weight", "ffn_routed_down.weight", "ffn_routed_up.weight", "ssm_out.weight")): return 1
-  if name.endswith(("attn_q.weight", "attn_k.weight", "attn_v.weight", "ssm_g_full.weight", "ssm_f_b.weight", "ssm_beta.weight")): return 0
-  if name.endswith(("ssm_q_conv1d.weight", "ssm_k_conv1d.weight", "ssm_v_conv1d.weight", "ssm_dt.bias")): return 0
-  return None
-
-def fake_value(name:str) -> tuple[int|float, object]:
-  if name.endswith("weight_scale"): return 120, dtypes.uint8
-  if name.endswith("_exps.weight"): return 0x11, dtypes.uint8
-  if name.endswith("ssm_a"): return -0.1, dtypes.float32
-  if name.endswith("ssm_dt.bias"): return 0.1, dtypes.float32
-  if "conv1d.weight" in name: return 0.1, dtypes.float32
-  if name.endswith("exp_probs_b.bias"): return 0.0, dtypes.float32
-  if name.endswith("norm.weight"): return 1.0, dtypes.bfloat16
-  return 0.001, dtypes.bfloat16
-
-def fake_tp_tensor(shape:tuple[int, ...], value:int|float, dtype, devices:tuple[str, ...], axis:int|None) -> Tensor:
-  if axis is not None and shape[axis] % len(devices): raise ValueError(f"shape {shape} is not TP{len(devices)} divisible on axis {axis}")
-  source = Tensor.full(shape, value, dtype=dtype, device=devices[0]).clone().realize()
-  return source.shard(devices, axis=axis).realize()
-
-def sync(devices:tuple[str, ...]) -> None:
-  for device in devices: Device[device].synchronize()
-
-def main() -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument("--devices", type=int, default=8)
-  parser.add_argument("--mode", choices=("attention", "block"), default="attention")
-  parser.add_argument("--iterations", type=int, default=20)
-  args = parser.parse_args()
-  devices = tuple(f"AMD:{i}" for i in range(args.devices))
-  # One exact-width KDA layer, but only 16 fake routed experts. This retains top-k 16 and every
-  # official per-GPU matrix/state shape while keeping fake expert storage below 300 MB per layer.
-  config = replace(kimi_k3_config(4), num_blocks=1, num_experts=16, num_experts_per_tok=16, ssm_layers=(True,),
-                   attn_res_block_size=0)
-  block = GatedDeltaNetBlock(config, config.ssm)
-  begin = time.perf_counter()
-  for name,tensor in nn.state.get_state_dict(block).items():
-    if args.mode == "attention" and name.startswith(("ffn_", "exp_probs_")): continue
-    value, dtype = fake_value(name)
-    tensor.replace(fake_tp_tensor(tuple(int(x) for x in tensor.shape), value, dtype, devices, tp_axis(name)))
-  sync(devices)
-  print(f"fake weights: {time.perf_counter()-begin:.3f}s", flush=True)
-  x_source = (((Tensor.arange(config.dim, dtype=dtypes.float32).reshape(1, 1, config.dim) % 31) / 31) \
-    .cast(dtypes.bfloat16).to(devices[0])).clone().realize()
-  x = x_source.shard(devices, axis=None).realize()
-  block._init_state(x)
-  # Use direct buffer-backed state shards. The production path reaches this form after prefill;
-  # the fake harness begins immediately at decode and must not feed lazy clone graphs to TinyJit.
-  for state,axis in ((block.conv_state_q, 2), (block.conv_state_k, 2), (block.conv_state_v, 2), (block.recurrent_state, 1)):
-    state.replace(Tensor.zeros(*state.shape, dtype=state.dtype, device=devices[0]).shard(devices, axis=axis).realize())
-
-  @TinyJit
-  def run(inp:Tensor) -> Tensor:
-    if args.mode == "attention": return block._attention(block.attn_norm(inp), 0).realize()
-    return block(inp, 0).realize()
-
-  # uncaptured, capture, then replay only
-  run(x); sync(devices)
-  run(x); sync(devices)
-  samples:list[float] = []
-  profile_marker(f"fake K3 {args.mode} start")
-  for _ in range(args.iterations):
-    begin = time.perf_counter(); out = run(x); sync(devices); samples.append((time.perf_counter()-begin)*1e3)
-  profile_marker(f"fake K3 {args.mode} end")
-  print(f"{args.mode}: median={statistics.median(samples):.3f} ms/layer, min={min(samples):.3f} ms/layer, "
-        f"projected_93_layer_rate={1000/(statistics.median(samples)*93):.3f} tok/s, finite={out.float().isfinite().all().item()}")
-
-if __name__ == "__main__": main()

extra/export_model.py

@@ -35,7 +35,7 @@ def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], Li
     return name
 
   for call in iter_kernel_calls(linear):
-    arg_uops = [b for b in call.src[1:] if b.op is not Ops.BIND]
+    arg_uops = [b for b in call.src[1:] if not b.is_bound_var]
     prg = to_program(call.src[0], Device[arg_uops[0].device].renderer)
     info = prg.arg
     functions[info.function_name] = prg.src[2].arg

extra/gemm/cdna_asm_gemm.py

@@ -122,7 +122,8 @@ def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp,
   groups_x, groups_y = UOp.special(ceildiv(N, tile_n), "gidx0"), UOp.special(ceildiv(M, tile_m), "gidx1")
   lds = UOp.placeholder((163840,), dtypes.uint8, 0, AddrSpace.LOCAL)
   sink = UOp.sink(C.base, A.base, B.base, scale_a.base, scale_b.base, *(x.base for x in extra), lds, threads, groups_x, groups_y,
-                  arg=KernelInfo(f"custom_mxfp4_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K)))
+                  arg=KernelInfo(f"mxfp4_gemm_{M}_{N}_{K}",
+                                 estimates=Estimates(ops=2*M*N*K, mem=(M*half_k+N*half_k)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
   insts = build_kernel(M, N, K, tile_m, tile_n)
   return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in insts))))
 

extra/gemm/moe_gemm.py

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

extra/gemm/triton_nv_matmul.py

@@ -79,7 +79,7 @@ if __name__ == "__main__":
   linear, var_vals = C.linear_with_vars()
   last_call = linear.src[-1]
   ast = last_call.src[0]
-  bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]
+  bufs = [s.buffer for s in last_call.src[1:] if not s.is_bound_var]
 
   src = compiled.asm["ptx"]
   # specify the shared memory here so we don't need to do it dynamically

extra/hcq2/ops_amd2.py

@@ -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 tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, encode_kernargs_clike, make_cmdbuf
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, encode_kernargs_clike, make_cmdbuf
 from tinygrad.runtime.support.hcq2 import make_binary_patch
 from tinygrad.uop.ops import sint, UOp
 from tinygrad.device import Compiled, BufferSpec, Buffer, Device
@@ -179,11 +179,10 @@ class SDMAOps(FastEnum): COPY = auto(); POLL_REGMEM = auto(); FENCE = auto(); TR
 
 def sdma_copy(ctx, call):
   sz = call.src[2].max_numel() * call.src[2].dtype.itemsize
-  src_addr, dst_addr = call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs)
-  return call.ins(SDMAOps.COPY, src=tuple(UOp.const(x, dtypes.uint32) for off in range(0, sz, ctx.max_copy_size) for x in (
-    ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR),
-    ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0,
-    *data64_le(src_addr+UOp.const(off, dtypes.uint64)), *data64_le(dst_addr+UOp.const(off, dtypes.uint64)))))
+  hdr = ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR)
+  return call.ins(SDMAOps.COPY, src=tuple(x for off in range(0, sz, ctx.max_copy_size) for x in (
+    *(UOp.const(v, dtypes.uint32) for v in (hdr, ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0)),
+    *(a + UOp.const(off, dtypes.uint64) if off else a for a in (call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs))))))
 
 def sdma_wait(ctx, ins, dst, val):
   op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \
@@ -288,14 +287,16 @@ def amd_build_program(prg:UOp) -> UOp:
 
 class AMDAllocator(HCQAllocator['AMDDevice']):
   def __init__(self, dev:AMDDevice):
-    super().__init__(dev, supports_copy_from_disk=dev.has_sdma_queue, supports_transfer=dev.has_sdma_queue and not dev.is_usb())
+    super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb())
 
-  def _alloc(self, size:int, options:BufferSpec) -> HCQ2Buffer:
-    return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_sdma_queue)
+  def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
+    return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_copy_queue)
 
   def _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque)
 
-  def _do_map(self, buf:HCQ2Buffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+  def _do_map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+
+  def _do_unmap(self, buf:HCQBuffer): self.dev.iface.unmap(buf)
 
 @dataclass
 class AMDQueueDesc:
@@ -388,15 +389,24 @@ class KFDIface:
     return hcqbuf
 
   def free(self, mem):
+    self._unmap(mem)
+    if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
+    kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
+
+  def unmap(self, mem):
+    self._unmap(mem)
+    if getattr(mem, '_owns_kfd_handle', False): kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
+
+  def _unmap(self, mem):
     gpus = (ctypes.c_int32 * 1)(self.gpu_id)
     stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1)
     assert stm.n_success == 1
-    if mem.owner == self.dev:
-      if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
-      kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
 
   def map(self, mem):
-    if mem.owner is not None and mem.owner._is_cpu(): return self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
+    if mem.owner is not None and mem.owner._is_cpu():
+      mapped = self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
+      mapped._owns_kfd_handle = True
+      return mapped
 
     c_gpus = (ctypes.c_int32 * 1)(self.gpu_id)
     stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(c_gpus), n_devices=1)
@@ -468,6 +478,7 @@ class PCIIface(PCIIfaceBase):
 
   def require_profile_mode(self): return True
   def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return True # TODO: account for WGP disablement on some asics.
+  def unmap(self, mem): self.free(mem)
 
   def _compute_props(self):
     self.ip_versions = self.dev_impl.ip_ver
@@ -549,9 +560,7 @@ class AMDDevice(HCQ2Compiled):
   def is_usb(self) -> bool: return False
 
   def __init__(self, device:str=""):
-    self.device_id = int(device.split(":")[1]) if ":" in device else 0
-
-    self.iface = self._select_iface()
+    self.iface = self._select_iface(device)
 
     self.target:tuple[int, ...] = ((trgt:=self.iface.props['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
     self.arch = "gfx%d%x%x" % self.target
@@ -581,7 +590,7 @@ class AMDDevice(HCQ2Compiled):
 
     self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
     self.sdma_queues:dict = {}
-    self.has_sdma_queue = True # self.sdma_queue(0) is not None, TODO: think of this
+    self.has_copy_queue = not getenv("AMD_DISABLE_SDMA")
 
     super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, can_recover=self.is_am(), arch=self.arch)
 

extra/runbook_digitalocean_mi350x.md

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

extra/thunder/amd/fa.py

@@ -19,16 +19,33 @@ def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None
 @functools.cache
 def custom_fused_qkv_rope_forward(q:UOp, k:UOp, v:UOp, xqkv:UOp, freqs_cis:UOp,
                                   device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
-  code = (pathlib.Path(__file__).parent / "fused_qkv_rope.cpp").read_text()
-  threads = 256
-  thread_idx = UOp.special(threads, "lidx0")
-  block_idx_x, block_idx_y = UOp.special(B, "gidx0"), UOp.special(N, "gidx1")
-  sink = UOp.sink(q.base, k.base, v.base, xqkv.base, freqs_cis.base, thread_idx, block_idx_x, block_idx_y,
-                  arg=KernelInfo(name="fused_qkv_rope_forward"))
-  compile_args = ["-std=c++20", "-ffast-math", f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}",
-                  f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DTHREADS_PER_BLOCK={threads}"]
-  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
+  group_size = H // H_KV
+  q, k, v = q.reshape(B, N, H, D), k.reshape(B, N, H_KV, D), v.reshape(B, N, H_KV, D)
+  xqkv = xqkv.reshape(B, N, H_KV, group_size + 2, D)
+  b, n = UOp.range(B, 0), UOp.range(N, 1)
+  pair = UOp.range(D // 2, 2)
+  even = pair * 2
+  c = freqs_cis[0, n, 0, pair, 0].cast(dtypes.float)
+  s = freqs_cis[0, n, 0, pair, 1].cast(dtypes.float)
+  ordered:UOp|None = None
+  for kvh in range(H_KV):
+    q_out, k_out, v_out = (x.after(ordered) if ordered is not None else x for x in (q, k, v))
+    x_in = xqkv.after(ordered) if ordered is not None else xqkv
+    stores:list[UOp] = []
+    for rep in range(group_size):
+      a = x_in[b, n, kvh, rep, even].cast(dtypes.float)
+      bb = x_in[b, n, kvh, rep, even + 1].cast(dtypes.float)
+      h = kvh * group_size + rep
+      stores += [q_out[b, n, h, even].store((a * c - bb * s).cast(q.dtype)), q_out[b, n, h, even + 1].store((a * s + bb * c).cast(q.dtype))]
+    a = x_in[b, n, kvh, group_size, even].cast(dtypes.float)
+    bb = x_in[b, n, kvh, group_size, even + 1].cast(dtypes.float)
+    stores += [k_out[b, n, kvh, even].store((a * c - bb * s).cast(k.dtype)),
+               k_out[b, n, kvh, even + 1].store((a * s + bb * c).cast(k.dtype)),
+               v_out[b, n, kvh, even].store(x_in[b, n, kvh, group_size + 1, even]),
+               v_out[b, n, kvh, even + 1].store(x_in[b, n, kvh, group_size + 1, even + 1])]
+    ordered = UOp.group(*stores)
+  assert ordered is not None
+  return ordered.end(pair, n, b).sink(arg=KernelInfo(name="fused_qkv_rope_forward"))
 
 @functools.cache
 def custom_fused_qkv_rope_backward(dxqkv:UOp, dq:UOp, dk:UOp, dv:UOp, freqs_cis:UOp,
@@ -110,7 +127,49 @@ def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
   return _sharded_empty(ref.shape, ref, axis)
 
 @functools.cache
-def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink):
+def _windowed_lse(xq:Tensor, xk:Tensor, sinks, W:int) -> Tensor:
+  B, N, H, hd = xq.shape
+  H_KV = xk.shape[2]; R = H // H_KV; nb = N // W; sm = hd ** -0.5
+  q = xq.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
+  k = xk.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
+  k_prev = k.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
+  sc_d = (q @ k.transpose(-1, -2)) * sm
+  sc_p = (q @ k_prev.transpose(-1, -2)) * sm
+  li, lj = Tensor.arange(W).reshape(W, 1), Tensor.arange(W).reshape(1, W)
+  pv = (Tensor.arange(nb).reshape(nb, 1, 1) >= 1)
+  sc_d = (lj <= li).where(sc_d, -float("inf"))
+  sc_p = ((li < lj) & pv).where(sc_p, -float("inf"))
+  m = sc_d.max(-1, keepdim=True).maximum(sc_p.max(-1, keepdim=True))
+  if sinks is not None: m = m.maximum(sinks.reshape(1, H_KV, R, 1, 1, 1).float())
+  denom = (sc_d - m).exp().sum(-1, keepdim=True) + (sc_p - m).exp().sum(-1, keepdim=True)
+  if sinks is not None: denom = denom + (sinks.reshape(1, H_KV, R, 1, 1, 1).float() - m).exp()
+  return (m + denom.log()).reshape(B, H, N).unsqueeze(2)  # (B, H, 1, N), matches saved l_vec
+
+def _windowed_delta(xq:Tensor, xk:Tensor, xv:Tensor, do:Tensor, sinks, W:int) -> Tensor:
+  B, N, H, hd = xq.shape
+  H_KV = xk.shape[2]; R = H // H_KV; nb = N // W; sm = hd ** -0.5
+  q = xq.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
+  k = xk.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
+  v = xv.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
+  dob = do.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
+  k_prev = k.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
+  v_prev = v.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
+  sc_d = (q @ k.transpose(-1, -2)) * sm
+  sc_p = (q @ k_prev.transpose(-1, -2)) * sm
+  li, lj = Tensor.arange(W).reshape(W, 1), Tensor.arange(W).reshape(1, W)
+  pv = (Tensor.arange(nb).reshape(nb, 1, 1) >= 1)
+  sc_d = (lj <= li).where(sc_d, -float("inf"))
+  sc_p = ((li < lj) & pv).where(sc_p, -float("inf"))
+  m = sc_d.max(-1, keepdim=True).maximum(sc_p.max(-1, keepdim=True))
+  if sinks is not None: m = m.maximum(sinks.reshape(1, H_KV, R, 1, 1, 1).float())
+  e_d, e_p = (sc_d - m).exp(), (sc_p - m).exp()
+  denom = e_d.sum(-1, keepdim=True) + e_p.sum(-1, keepdim=True)
+  if sinks is not None: denom = denom + (sinks.reshape(1, H_KV, R, 1, 1, 1).float() - m).exp()
+  o = ((e_d / denom) @ v) + ((e_p / denom) @ v_prev)
+  delta = (dob * o).sum(-1)
+  return delta.reshape(B, H, N).unsqueeze(2)
+
+def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink, window=0):
   def grad(dou:UOp, ker:UOp) -> tuple:
     do = Tensor(dou, device=dou.device)
     attn = Tensor(ker.src[1].after(ker), device=ker.src[1].device)
@@ -118,6 +177,8 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
     xq = Tensor(ker.src[3], device=ker.src[3].device)
     xk = Tensor(ker.src[4], device=ker.src[4].device)
     xv = Tensor(ker.src[5], device=ker.src[5].device)
+    if window:
+      l_vec = _windowed_lse(xq, xk, Tensor(ker.src[6], device=ker.src[6].device) if has_sink else None, window)
 
     dq = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t)
     GROUP_SIZE = H_local // H_KV_local
@@ -128,8 +189,10 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
     # delta_vec = (do * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
     delta_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
     delta_vec, dq = Tensor.custom_kernel(delta_vec, dq, attn, do, fxn=functools.partial(custom_fa_backward_pre, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:2]
+    if window:
+      delta_vec = _windowed_delta(xq, xk, xv, do, Tensor(ker.src[6], device=ker.src[6].device) if has_sink else None, window)
 
-    dq, dk_partial, dv_partial = Tensor.custom_kernel(dq, dk_partial, dv_partial, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:3]
+    dq, dk_partial, dv_partial = Tensor.custom_kernel(dq, dk_partial, dv_partial, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, window=window))[:3]
 
     if D == 64:
       dq = dq.reshape(B, H, N//16, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2).permute(0, 1, 2, 8, 9, 10, 11, 3, 4, 6, 7, 5, 12).reshape(B, H, N, D).transpose(1, 2)
@@ -149,7 +212,7 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
   return grad
 
 # TODO: remove write_flat once scheduler can remove reshapes between custom_kernel. TestCustomKernel.test_simple_reshape
-def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False, write_flat:bool=False, sinks:Tensor|None=None):
+def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False, write_flat:bool=False, sinks:Tensor|None=None, window:int=0):
   assert attn_mask is None, "attn_mask not supported"
   assert is_causal, "only causal attention supported"
 
@@ -176,18 +239,18 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
   attn = _sharded_empty((B, N, H * D), xq, axis=shard_axis) if write_flat else _sharded_empty_like(xq, axis=shard_axis)
   l_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
 
-  grad = _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink)
+  grad = _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink, window=window)
 
   fwd_inputs = (attn, l_vec, xq, xk, xv) + ((sinks,) if has_sink else ())
-  attn, l_vec = Tensor.custom_kernel(*fwd_inputs, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, has_sink=has_sink), grad_fxn=grad)[:2]
+  attn, l_vec = Tensor.custom_kernel(*fwd_inputs, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, has_sink=has_sink, window=window), grad_fxn=grad)[:2]
 
   return attn, attn, l_vec
 
 @functools.cache
-def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None, *, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, has_sink:bool=True):
+def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None, *, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, has_sink:bool=True, window:int=0):
   code = (pathlib.Path(__file__).parent / "fa_fwd_causal.cpp").read_text()
   compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math",
-                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DATTN_SINK={int(has_sink)}"]
+                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DATTN_SINK={int(has_sink)}", f"-DWINDOW={window}"]
 
   Q_BLOCK_SIZE = 32
   NUM_WARPS = 8
@@ -247,10 +310,10 @@ def custom_fa_backward_pre(delta_vec:UOp, dq:UOp, o:UOp, do:UOp, device:str, arc
              src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
 
 @functools.cache
-def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_vec:UOp, delta_vec:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
+def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_vec:UOp, delta_vec:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, window:int=0):
   code = (pathlib.Path(__file__).parent / "fa_bwd_causal.cpp").read_text()
   compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math",
-                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}"]
+                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DWINDOW={window}"]
 
   BLOCK_SIZE_KV = 256
   GROUP_SIZE = H // H_KV

extra/thunder/amd/fused_qkv_rope.cpp

@@ -1,69 +0,0 @@
-#include <hip/hip_runtime.h>
-#include <hip/hip_bf16.h>
-
-#ifndef ATTN_B
-#define ATTN_B 2
-#endif
-#ifndef ATTN_N
-#define ATTN_N 8192
-#endif
-#ifndef ATTN_H
-#define ATTN_H 32
-#endif
-#ifndef ATTN_H_KV
-#define ATTN_H_KV 8
-#endif
-#ifndef ATTN_D
-#define ATTN_D 128
-#endif
-#ifndef THREADS_PER_BLOCK
-#define THREADS_PER_BLOCK 256
-#endif
-
-constexpr int GROUP_SIZE = ATTN_H / ATTN_H_KV;
-constexpr int HALF_D = ATTN_D / 2;
-constexpr int PACKED_D = (GROUP_SIZE + 2) * ATTN_D;
-
-extern "C" __global__ __launch_bounds__(THREADS_PER_BLOCK) void
-fused_qkv_rope_forward(
-    __hip_bfloat16*       __restrict__ q,
-    __hip_bfloat16*       __restrict__ k,
-    __hip_bfloat16*       __restrict__ v,
-    const __hip_bfloat16* __restrict__ xqkv,
-    const __hip_bfloat16* __restrict__ freqs_cis) {
-  const int b = blockIdx.x;
-  const int n = blockIdx.y;
-  const int bn = b * ATTN_N + n;
-  const int packed_bn = bn * ATTN_H_KV * PACKED_D;
-  const int q_bn = bn * ATTN_H * ATTN_D;
-  const int kv_bn = bn * ATTN_H_KV * ATTN_D;
-
-  if (threadIdx.x < HALF_D) {
-    const int pair = threadIdx.x;
-    const int even = pair << 1;
-    const float c = static_cast<float>(freqs_cis[((n * HALF_D + pair) * 2) + 0]);
-    const float s = static_cast<float>(freqs_cis[((n * HALF_D + pair) * 2) + 1]);
-
-    for (int kvh = 0; kvh < ATTN_H_KV; kvh++) {
-      const int base = packed_bn + kvh * PACKED_D;
-
-      for (int rep = 0; rep < GROUP_SIZE; rep++) {
-        const int qbase = base + rep * ATTN_D;
-        const int h = kvh * GROUP_SIZE + rep;
-        const float a = static_cast<float>(xqkv[qbase + even]);
-        const float bb = static_cast<float>(xqkv[qbase + even + 1]);
-        const int out = q_bn + h * ATTN_D + even;
-        q[out] = static_cast<__hip_bfloat16>(a * c - bb * s);
-        q[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c);
-      }
-
-      const float a = static_cast<float>(xqkv[base + GROUP_SIZE * ATTN_D + even]);
-      const float bb = static_cast<float>(xqkv[base + GROUP_SIZE * ATTN_D + even + 1]);
-      const int out = kv_bn + kvh * ATTN_D + even;
-      k[out] = static_cast<__hip_bfloat16>(a * c - bb * s);
-      k[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c);
-      v[out] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even];
-      v[out + 1] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even + 1];
-    }
-  }
-}

extra/thunder/amd/matvec_bf16_splitk.cpp

@@ -1,74 +0,0 @@
-#include "kittens.cuh"
-
-using namespace kittens;
-
-#ifndef MATVEC_N
-#define MATVEC_N 1536
-#endif
-#ifndef MATVEC_K
-#define MATVEC_K 7168
-#endif
-
-constexpr int SPLIT_WAVES = 8;
-
-template<int W>
-__device__ __forceinline__ float run_split(const bf16 *A_ptr, const bf16 *B_ptr, int out_base,
-                                           st_bf<16, 32, st_16x32_s> &As,
-                                           st_bf<16, 32, st_16x32_s> &Bs) {
-  constexpr int K = MATVEC_K;
-  rt_bf<16, 32, row_l, rt_16x32_s> A;
-  rt_bf<16, 32, row_l, rt_16x32_s> B;
-  rt_fl<16, 16, col_l, rt_16x16_s> C;
-  zero(C);
-  const int lane = laneid();
-  constexpr int k_begin = W * (K / SPLIT_WAVES), k_end = k_begin + K / SPLIT_WAVES;
-  #pragma unroll 1
-  for (int k = k_begin; k < k_end; k += 32) {
-    #pragma unroll
-    for (int idx = lane; idx < 16 * 32; idx += 64) {
-      const int row = idx / 32, col = idx % 32;
-      *reinterpret_cast<bf16 *>(reinterpret_cast<char *>(&As.data[0]) + As.swizzle({row, col})) = A_ptr[k + col];
-      *reinterpret_cast<bf16 *>(reinterpret_cast<char *>(&Bs.data[0]) + Bs.swizzle({row, col})) =
-        B_ptr[(out_base + row) * K + k + col];
-    }
-    asm volatile("s_waitcnt lgkmcnt(0)");
-    load(A, As);
-    load(B, Bs);
-    asm volatile("s_waitcnt lgkmcnt(0)");
-    mma_ABt(C, A, B, C);
-  }
-  return C.tiles[0][0].data[0].x;
-}
-
-// Eight waves split K for one 16-channel output tile.  Each wave uses MFMA on
-// a repeated activation row, then wave zero reduces the eight FP32 partials.
-__global__ __launch_bounds__(64 * SPLIT_WAVES, 1)
-void hk_bf16_matvec_splitk(bf16 *C_ptr, const bf16 *A_ptr, const bf16 *B_ptr, bf16 *unused) {
-  constexpr int N = MATVEC_N, K = MATVEC_K;
-  static_assert(N % 16 == 0 && K % (32 * SPLIT_WAVES) == 0);
-  __shared__ st_bf<16, 32, st_16x32_s> As[SPLIT_WAVES];
-  __shared__ st_bf<16, 32, st_16x32_s> Bs[SPLIT_WAVES];
-  __shared__ float partial[SPLIT_WAVES][16];
-  const int tid = threadIdx.x, wave = tid / 64, lane = tid & 63;
-  const int out_base = blockIdx.x * 16;
-  float result = 0.0f;
-  switch (wave) {
-    case 0: result = run_split<0>(A_ptr, B_ptr, out_base, As[0], Bs[0]); break;
-    case 1: result = run_split<1>(A_ptr, B_ptr, out_base, As[1], Bs[1]); break;
-    case 2: result = run_split<2>(A_ptr, B_ptr, out_base, As[2], Bs[2]); break;
-    case 3: result = run_split<3>(A_ptr, B_ptr, out_base, As[3], Bs[3]); break;
-    case 4: result = run_split<4>(A_ptr, B_ptr, out_base, As[4], Bs[4]); break;
-    case 5: result = run_split<5>(A_ptr, B_ptr, out_base, As[5], Bs[5]); break;
-    case 6: result = run_split<6>(A_ptr, B_ptr, out_base, As[6], Bs[6]); break;
-    case 7: result = run_split<7>(A_ptr, B_ptr, out_base, As[7], Bs[7]); break;
-  }
-  if (lane < 16) partial[wave][lane] = result;
-  asm volatile("s_waitcnt lgkmcnt(0)");
-  __builtin_amdgcn_s_barrier();
-  if (wave == 0 && lane < 16) {
-    float total = 0.0f;
-    #pragma unroll
-    for (int i = 0; i < SPLIT_WAVES; i++) total += partial[i][lane];
-    C_ptr[out_base + lane] = static_cast<bf16>(total);
-  }
-}

extra/thunder/tiny/tk/tiles.py

@@ -209,7 +209,7 @@ class ST:
     return cls(uop, rows, cols, layout, base_shape, ker)
 
   def swizzle(self, row, col):
-    swizzled_offset = self.base_shape.swizzle(row, col, self._uop.dtype.scalar())
+    swizzled_offset = self.base_shape.swizzle(row, col, self._uop.dtype)
 
     row = swizzled_offset // self.base_shape.cols
     col = swizzled_offset % self.base_shape.cols

extra/torch_backend/backend.py

@@ -4,7 +4,7 @@
 # A006 Lambda argument `input` is shadowing a Python builtin
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.uop.ops import Ops, GroupOp
-from tinygrad.helpers import getenv, prod, strides_for_shape, argfix
+from tinygrad.helpers import getenv, prod, strides_for_shape
 import torch.lib
 TORCH_DEBUG = getenv("TORCH_DEBUG")
 import torch, pathlib, operator, functools, weakref
@@ -73,6 +73,12 @@ def wrap_view_op(fn):
     return wrap(ret)
   return _wrap
 
+# NOTE: list assignment raises IndexError on an out of range dim, and the index must be a tuple: a list of all ints is one advanced index
+def _index_dim(self, dim, idx):
+  idxs = [slice(None)] * self.ndim
+  idxs[dim] = idx
+  return self[tuple(idxs)]
+
 view_ops = {
   "aten.view": Tensor.reshape,
   "aten._unsafe_view": Tensor.reshape,  # when are views unsafe, and do we care?
@@ -82,15 +88,13 @@ view_ops = {
   "aten.transpose.int": Tensor.transpose,
   "aten.squeeze.dim": Tensor.squeeze,
   "aten.unsqueeze": Tensor.unsqueeze,
-  "aten.select.int": lambda self, dim, idx: self[(slice(None),) * (dim%self.ndim) + (idx,)],
+  "aten.select.int": _index_dim,
   "aten.permute": Tensor.permute,
   "aten.alias": lambda self: self,
   "aten.diagonal": Tensor.diagonal,
+  "aten.slice.Tensor": lambda self, dim=0, start=None, end=None, step=1: _index_dim(self, dim, slice(start, end, step)),
   }
 
-# torch 2.10 handles this natively
-if tuple(map(int, torch.__version__.split('.')[:2])) < (2, 10): view_ops.update({"aten.detach": Tensor.detach})
-
 for k,v in view_ops.items(): torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(wrap_view_op(v))
 
 def _get_view_ops(view): return getattr(view, "_view_ops", [])
@@ -99,46 +103,21 @@ def _apply_view_ops(target, ops):
   for fn, args, kwargs in ops: target = fn(target, *args, **kwargs)
   return target
 
-# similar to https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/InferSize.h
-def _reshape_target_shape(shape:tuple[int, ...], args) -> tuple[int, ...]|None:
-  if not (req := argfix(*args)): return None
-  new_shape, infer_idx = [], -1
-  for i, s in enumerate(req):
-    if s is None: s = shape[i] if i < len(shape) else None
-    if not isinstance(s, int): return None
-    if s == -1:
-      if infer_idx != -1: return None
-      infer_idx = len(new_shape)
-    new_shape.append(s)
-  total = prod(shape)
-  if infer_idx != -1:
-    known = prod(x for x in new_shape if x != -1)
-    if known == 0:
-      if total != 0: return None
-      new_shape[infer_idx] = 0
-    else: new_shape[infer_idx] = total // known
-  return tuple(new_shape) if prod(new_shape) == total else None
-
-# TODO: can we get rid of this? only for test_flatten_reshape_add
+# a chain of reshapes is undone by reshaping the value back to the base
 def _try_simple_reshape_view_write(base: Tensor, view: Tensor, val: Tensor) -> bool:
   if not (ops := _get_view_ops(view)): return False
-  shapes = [base.shape]
-  for fn, args, _ in ops:
-    if fn is Tensor.reshape:
-      if not (next_shape := _reshape_target_shape(shapes[-1], args)): return False
-      shapes.append(next_shape)
-  if shapes[-1] != view.shape: return False
-  for s in reversed(shapes[:-1]): val = val.reshape(s)
-  base.assign(val)
+  if any(fn is not Tensor.reshape for fn, _, _ in ops): return False
+  base.assign(val.reshape(base.shape))
   return True
 
 def _view_write(base: Tensor, view: Tensor, value: Tensor) -> None:
   val = value if value.dtype == base.dtype else value.cast(base.dtype)
-  if view.shape == base.shape: return base.assign(val)
   if _try_simple_reshape_view_write(base, view, val): return
   idx_base = Tensor.arange(base.numel(), dtype=dtypes.int32).reshape(base.shape)
   idx_view = _apply_view_ops(idx_base, _get_view_ops(view)).reshape(-1)
-  flat_base = base.reshape(base.numel()).contiguous()
+  # clone, not contiguous: contiguous() on a base that already owns its buffer returns the base itself, and scattering
+  # into that is an in-place write to a buffer other tensors still hold, which setitem refuses
+  flat_base = base.reshape(base.numel()).clone()
   flat_base[idx_view] = val.reshape(-1)
   base.assign(flat_base.reshape(base.shape))
 
@@ -166,11 +145,6 @@ def _index_put_impl_(self, indices, values, accumulate=False, unsafe=False):
 def index_put(self, indices, values, accumulate=False):
   return aten.index_put(self.cpu(), [z.cpu() if isinstance(z, torch.Tensor) else None for z in indices], values.clone().cpu(), accumulate).tiny()
 
-@torch.library.impl("aten::isin.Tensor_Tensor_out", "privateuseone")
-def isin_tensor_tensor_out(x, y, *, assume_unique=False, invert=False, out=None):
-  result = (unwrap(x).unsqueeze(-1) == unwrap(y).flatten()).any(-1)
-  return out.copy_(wrap(~result if invert else result))
-
 @torch.library.impl("aten::randperm.generator_out", "privateuseone")
 def randperm_generator(n, generator=None, out=None):
   if generator is not None: raise NotImplementedError("tinygrad torch backend does not support torch.Generator for randperm")
@@ -231,49 +205,6 @@ def as_strided(tensor:torch.Tensor, size, stride, storage_offset=None):
 def _reshape_alias(tensor:torch.Tensor, size, stride):
   return _as_strided(tensor, size, stride)
 
-@torch.library.impl("aten::empty_strided", "privateuseone")
-def empty_strided(size, stride, dtype=None, layout=None, device=None, pin_memory=False):
-  if TORCH_DEBUG: print(f"empty_strided {size=} {stride=} {dtype=} {layout=} {device=} {pin_memory=}")
-  ret = Tensor.empty(*size, dtype=_from_torch_dtype(dtype or torch.get_default_dtype()), device=_from_torch_device(device))
-  # TODO: should return with requested strides
-  return wrap(ret)
-
-@torch.library.impl("aten::empty.memory_format", "privateuseone")
-def empty_memory_format(size, dtype=None, layout=None, device=None, pin_memory=False, memory_format=None):
-  if TORCH_DEBUG: print(f"empty.memory_format {size=} {dtype=} {layout=} {device=} {pin_memory=} {memory_format=}")
-  ret = Tensor.empty(*size, dtype=_from_torch_dtype(dtype or torch.get_default_dtype()), device=_from_torch_device(device))
-  return wrap(ret)
-
-@torch.library.impl("aten::max_pool2d_with_indices", "privateuseone")
-def max_pool2d_with_indices(self:torch.Tensor, kernel_size:tuple[int, ...], stride=None, padding=0, dilation=1, ceil_mode=False):
-  # TODO: supprt stride [] in tinygrad?
-  if stride is not None and len(stride) == 0: stride = None
-  ret, idx = unwrap(self).max_pool2d(kernel_size, stride, dilation, padding, ceil_mode, return_indices=True)
-  return (wrap(ret), wrap(idx.cast(dtypes.int64)))
-
-@torch.library.impl("aten::max_pool2d_with_indices_backward", "privateuseone")
-def max_pool2d_with_indices_backward(grad_out:torch.Tensor, self:torch.Tensor, kernel_size:tuple[int, ...], stride=None, padding=0, dilation=1, ceil_mode=False, indices=None):
-  return wrap(Tensor.max_unpool2d(unwrap(grad_out), unwrap(indices), output_size=unwrap(self).shape))
-
-@torch.library.impl("aten::max_unpool2d", "privateuseone")
-def max_unpool2d(self:torch.Tensor, indices:torch.Tensor, output_size):
-  return wrap(unwrap(self).max_unpool2d(unwrap(indices), output_size=output_size))
-
-@torch.library.impl("aten::arange", "privateuseone")
-def arange(end, dtype=None, device=None, pin_memory=None):
-  has_float = isinstance(end, float)
-  return wrap(Tensor.arange(0, end, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))
-
-@torch.library.impl("aten::arange.start", "privateuseone")
-def arange_start(start, end, dtype=None, device=None, pin_memory=None):
-  has_float = any(isinstance(x, float) for x in (start, end))
-  return wrap(Tensor.arange(start, end, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))
-
-@torch.library.impl("aten::arange.start_step", "privateuseone")
-def arange_start_step(start, end, step, dtype=None, device=None, pin_memory=None):
-  has_float = any(isinstance(x, float) for x in (start, end, step))
-  return wrap(Tensor.arange(start, end, step, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))
-
 @torch.library.impl("aten::convolution_overrideable", "privateuseone")
 def convolution_overrideable(input, weight, bias, stride, padding, dilation, transposed, output_padding, groups):
   if TORCH_DEBUG >= 1:
@@ -294,12 +225,27 @@ def convolution_backward_overrideable(grad_out, input, weight, stride, padding,
   grads = out.gradient(*[t for t,m in zip([input, weight, bias], output_mask) if m], gradient=grad_out)
   return tuple([wrap(grads.pop(0)) if m else None for m in output_mask])
 
-@torch.library.impl("aten::slice.Tensor", "privateuseone")
-@wrap_view_op
-def slice_tensor(self, dim=0, start=None, end=None, step=1):
-  slices = [slice(None)] * self.ndim
-  slices[dim] = slice(start, end, step)
-  return self[slices]
+# the functional scatters. without an impl aten falls back to a path that assumes a real storage: "self.has_storage() INTERNAL ASSERT FAILED"
+def _scatter_into(self, src, dim, index):
+  out = unwrap(self).clone()
+  slices = [slice(None)] * out.ndim
+  slices[dim] = index
+  out[slices] = unwrap(src).cast(out.dtype)  # torch casts src to self's dtype, tinygrad setitem demands they already match
+  return wrap(out)
+
+@torch.library.impl("aten::slice_scatter", "privateuseone")
+def slice_scatter(self, src, dim=0, start=None, end=None, step=1): return _scatter_into(self, src, dim, slice(start, end, step))
+
+@torch.library.impl("aten::select_scatter", "privateuseone")
+def select_scatter(self, src, dim, index): return _scatter_into(self, src, dim, index)
+
+@torch.library.impl("aten::diagonal_scatter", "privateuseone")
+def diagonal_scatter(self, src, offset=0, dim1=0, dim2=1):
+  # a diagonal is not one axis, so scatter through the flat indices it picks out
+  base, out = unwrap(self), unwrap(self).clone().reshape(-1)
+  idx = Tensor.arange(base.numel(), dtype=dtypes.int32).reshape(base.shape).diagonal(offset, dim1, dim2).reshape(-1)
+  out[idx] = unwrap(src).cast(base.dtype).reshape(-1)
+  return wrap(out.reshape(base.shape))
 
 @torch.library.impl("aten::slice_backward", "privateuseone")
 def slice_backward(grad_out, input_sizes, dim, start, end, step):
@@ -341,19 +287,14 @@ for dim in [1, 2, 3]:
     torch.library.impl(f"aten::{pad_type}_pad{dim}d", "privateuseone")(functools.partial(pad_forward, mode=mode))
     torch.library.impl(f"aten::{pad_type}_pad{dim}d_backward", "privateuseone")(functools.partial(pad_backward, mode=mode))
 
-def upsample(self, size, align_corners=False, mode=None): return wrap(Tensor.interpolate(unwrap(self), size, mode=mode, align_corners=align_corners))
+# the schemas are all positional: (self, output_size, align_corners, *scales) for linear, (self, output_size, *scales) for nearest.
+def upsample(self, size, *args, mode=None):
+  return wrap(Tensor.interpolate(unwrap(self), size, mode=mode, align_corners=args[0] if mode == "linear" else False))
 for i,pre in enumerate(["", "bi", "tri"]):
   torch.library.impl(f"aten::upsample_{pre}linear{i+1}d", "privateuseone")(functools.partial(upsample, mode="linear"))
   torch.library.impl(f"aten::upsample_nearest{i+1}d", "privateuseone")(functools.partial(upsample, mode="nearest"))
   torch.library.impl(f"aten::_upsample_nearest_exact{i+1}d", "privateuseone")(functools.partial(upsample, mode="nearest-exact"))
 
-@torch.library.impl("aten::scatter_add.out", "privateuseone")
-def scatter_add(self, dim, index, src, out):
-  self, index, src, out_unwrapped = unwrap(self), unwrap(index), unwrap(src), unwrap(out)
-  if self.shape == (): _apply_inplace(out_unwrapped, src)
-  else: _apply_inplace(out_unwrapped, Tensor.scatter_reduce(self, dim, index, src, reduce='sum'))
-  return out
-
 def _copy_between_devices(src, dest, cast_dtype, to_device, non_blocking=False):
   if src.is_tiny and dest.is_tiny:
     src_t, dest_t = unwrap(src), unwrap(dest)
@@ -404,15 +345,11 @@ def sort_values(input, dim=-1, descending=False, stable=True, values=None, indic
   _apply_inplace(unwrap(indices), out_indices.cast(dtypes.int64))
   return values, indices
 
-@torch.library.impl("aten::_linalg_svd", "privateuseone")
-def _linalg_svd(self, full_matrices=False):
-  U, S, Vh = unwrap(self).svd(full_matrices)
-  return wrap(U), wrap(S), wrap(Vh)
-
 # register some decompositions
 from torch._decomp import get_decompositions
 decomps = [
   aten.native_layer_norm_backward,
+  aten.native_group_norm_backward,
   aten.linalg_cross,
   aten.addmm,
   aten.addcmul,
@@ -447,12 +384,20 @@ decomps = [
   aten._softmax_backward_data, aten.embedding_dense_backward,
   aten.linalg_vector_norm,
   aten.binary_cross_entropy, aten.binary_cross_entropy_backward,
+  # the C++ mse/smooth_l1 kernels resize their out tensor, and a tiny tensor has no storage to resize
+  aten.mse_loss, aten.mse_loss_backward,
+  aten.smooth_l1_loss, aten.smooth_l1_loss_backward,
   aten.upsample_nearest2d.out,
+  # NOTE: only the "out" overload, the "vec" one is CompositeImplicitAutograd and overriding it loses the autograd kernel
+  aten.upsample_bicubic2d.out,
+  aten._adaptive_avg_pool2d,
   # activations
   aten.hardswish, aten.hardswish_backward,
   aten.hardtanh, aten.hardtanh_backward,
   aten.gelu, aten.gelu_backward,
-  aten.logical_and,
+  # NOTE: no aten.logical_or here, its decomposition reaches aten.bitwise_or through a path that checks aliasing by
+  # reading storage, which a tiny tensor has none of. it gets a direct impl below instead
+  aten.logical_and, aten.logical_xor,
   aten.randint,
   aten.eye,
   aten.hardsigmoid_backward,
@@ -495,7 +440,7 @@ simple_tensor_methods = [
   # reduce
   "all", "any", "argmax", "argmin", "cumsum", "cumprod",
   # complex
-  "avg_pool2d", "linspace"]
+  "linspace"]
 
 tiny_backend_out = {**{f"aten.{x}.out":getattr(Tensor,x) for x in simple_tensor_methods}, **{
   "aten.add.out": lambda input,other,alpha=1: input+alpha*other,
@@ -540,6 +485,8 @@ tiny_backend_out = {**{f"aten.{x}.out":getattr(Tensor,x) for x in simple_tensor_
   "aten.where.self_out": Tensor.where,
   "aten.prod.int_out": Tensor.prod,
   "aten.scatter.src_out": Tensor.scatter,
+  "aten.scatter_add.out": lambda self,dim,index,src: src if self.shape == () else Tensor.scatter_reduce(self, dim, index, src, reduce="sum"),
+  "aten.isin.Tensor_Tensor_out": lambda x,y,assume_unique=False,invert=False: (x.unsqueeze(-1)==y.flatten()).any(-1) != invert,
   # NOTE: axis=[] in torch means all, change tinygrad?
   "aten.sum.IntList_out": lambda self,axis,keepdim=False,dtype=None:
     self.sum(axis if axis is None or len(axis) else None, keepdim,
@@ -555,10 +502,9 @@ def wrap_out(f):
     assert out.shape == assigned.shape, f"shape mismatch: {assigned.shape} -> {out.shape}"
     assert out.device == assigned.device or out.device is None or assigned.device is None, f"device mismatch: {assigned.device} -> {out.device}"
     assert out.dtype == assigned.dtype, f"dtype mismatch: {assigned.dtype} -> {out.dtype}"
-    # an out= that is a view has to be written through its base, and _apply_inplace gives a deviceless base its buffer first
-    if canonical_base(out) is not out: return _apply_inplace(out, assigned) or out
-    if out.device is None and assigned.device is not None: out.replace(out.empty_like(device=assigned.device))
-    return out.assign(assigned)
+    # writing out= is an in-place write like any other: through the base if it is a view, refreshing any derived views
+    _apply_inplace(out, assigned)
+    return out
   return _wrap_out
 
 def _inplace_op(t, new_value):
@@ -566,7 +512,14 @@ def _inplace_op(t, new_value):
   else: _apply_inplace(t, new_value)
   return t
 
-tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
+# the three arange overloads are one function at different arity, and dtype/layout/device/pin_memory are keyword only in all of them
+def _arange(*args, dtype=None, **_):
+  return Tensor.arange(*args, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if any(isinstance(x, float) for x in args) else torch.int64)))
+
+def _empty(size, dtype=None, device=None, **_):
+  return Tensor.empty(*size, dtype=_from_torch_dtype(dtype or torch.get_default_dtype()), device=_from_torch_device(device))
+
+tiny_backend = {**tiny_backend_out, **{
   "aten.remainder.Scalar_Tensor": lambda x,y: x%y,
   "aten.floor_divide": lambda x,y: x//y,
   "aten.floor_divide_.Tensor": lambda x,y: x//y,
@@ -579,8 +532,8 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   # inplace ops using replace for fusion
   "aten.zero_": lambda x: x.const_like(0),
   "aten.fill_.Scalar": lambda x, y: x.const_like(y),
-  "aten.add_.Tensor": lambda self, other, alpha=1.0: self + other * alpha,
-  "aten.add_.Scalar": lambda self, other, alpha=1.0: self + other * alpha,
+  "aten.add_.Tensor": lambda self, other, alpha=1: self + other * alpha,
+  "aten.add_.Scalar": lambda self, other, alpha=1: self + other * alpha,
   "aten.mul_.Tensor": lambda self, other: self * other,
   "aten.mul_.Scalar": lambda self, other: self * other,
   # relu doesn't have an out form?
@@ -613,7 +566,9 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   # these don't work in out form, they have size 0
   "aten.abs": Tensor.abs,
   "aten.logical_not": Tensor.logical_not,
-  "aten.logical_or_": lambda x, y: x | y,
+  # compare against zero first: logical_* is bool-valued for any input dtype, while | is bitwise
+  "aten.logical_or": lambda x, y: (x != 0) | (y != 0),
+  "aten.logical_or_": lambda x, y: (x != 0) | (y != 0),
   "aten.multinomial": Tensor.multinomial,
   "aten.masked_fill_.Scalar": lambda self, mask, value: self.masked_fill(mask, value),
   "aten.masked_fill_.Tensor": lambda self, mask, value: self.masked_fill(mask, value),
@@ -622,14 +577,7 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   "aten.masked_select": Tensor.masked_select,
   "aten.all": Tensor.all,
   "aten.sgn": Tensor.sign,
-  "aten.acos": Tensor.acos,
   "aten.any": Tensor.any,
-  "aten.bitwise_not": Tensor.bitwise_not,
-  "aten.argmax": Tensor.argmax,
-  "aten.argmin": Tensor.argmin,
-  "aten.asinh": Tensor.asinh,
-  "aten.mul": Tensor.mul,
-  "aten.atanh": Tensor.atanh,
   "aten.fill_.Tensor": lambda self, value: self.const_like(value.reshape(()).item()),
   "aten.flip": Tensor.flip,
   "aten.scatter_reduce.two": Tensor.scatter_reduce,
@@ -640,10 +588,22 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
   "aten.add.Tensor": lambda input,other,alpha=1: input+alpha*other,
   "aten.linspace": lambda start, stop, steps, dtype=None, **kwargs:
     Tensor.linspace(start, stop, steps, **({"dtype": _from_torch_dtype(dtype)} if dtype is not None else {})),
+  # the functional copy_. without an impl the fallback segfaults on a tensor with no storage
+  "aten.copy": lambda self,src,non_blocking=False: src.cast(self.dtype).to(self.device).expand(self.shape),
+  "aten.arange": lambda end, **kwargs: _arange(0, end, **kwargs),
+  "aten.arange.start": _arange,
+  "aten.arange.start_step": _arange,
+  # empty_strided takes the strides and drops them: we always allocate contiguous
+  "aten.empty_strided": lambda size, stride, **kwargs: _empty(size, **kwargs),
+  "aten.empty.memory_format": _empty,
+  # TODO: supprt stride [] in tinygrad?
+  "aten.max_pool2d_with_indices": lambda self,kernel_size,stride=None,padding=0,dilation=1,ceil_mode=False: ((r:=Tensor.max_pool2d(self, kernel_size, stride or None, dilation, padding, ceil_mode, return_indices=True))[0], r[1].cast(dtypes.int64)),
+  "aten.max_pool2d_with_indices_backward": lambda grad_out,self,kernel_size,stride=None,padding=0,dilation=1,ceil_mode=False,indices=None: Tensor.max_unpool2d(grad_out, indices, output_size=self.shape),
+  "aten.max_unpool2d": lambda self,indices,output_size: Tensor.max_unpool2d(self, indices, output_size=output_size),
+  "aten._linalg_svd": lambda self,full_matrices=False: Tensor.svd(self, full_matrices),
   "aten.topk": Tensor.topk,
   "aten.constant_pad_nd": lambda self, padding, value=0.0: self.pad(padding, mode="constant", value=value).contiguous(),
-  # TODO: input contiguous is needed to prevent CFGContext circular dependency assertion for shapes >512 (see test_cumsum_arange_large)
-  "aten.cumsum": lambda self, dim: self.contiguous().cumsum(dim),
+  "aten.cumsum": lambda self, dim: self.cumsum(dim),
   "aten.logsumexp": lambda self, axis, keepdim=False: self.logsumexp(axis[0], keepdim=keepdim),
   "aten.roll": Tensor.roll,
   "aten.logcumsumexp": Tensor.logcumsumexp,
@@ -652,6 +612,7 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
     self.ones_like(**{k: v for k, v in {"dtype": _from_torch_dtype(dtype) if dtype else None,
                                         "device": _from_torch_device(device) if device else None}.items() if v is not None}),
   "aten.max.dim": lambda self, dim, keepdim=False: (self.max(dim, keepdim), self.argmax(dim, keepdim).cast(dtype=dtypes.int64)),
+  "aten.min.dim": lambda self, dim, keepdim=False: (self.min(dim, keepdim), self.argmin(dim, keepdim).cast(dtype=dtypes.int64)),
   "aten.cummax": lambda self, dim: ((r := self.cummax(dim))[0], r[1].cast(dtypes.int64)),
   "aten.cummin": lambda self, dim: ((r := self.cummin(dim))[0], r[1].cast(dtypes.int64)),
   "aten.nonzero": Tensor.nonzero,
@@ -713,15 +674,16 @@ def wrap_inplace_view_op(f):
   return nf
 
 # the aten schema says how an op is called: an inplace view retargets the view, a writable first arg is inplace,
-# and a writable out arg must have come from tiny_backend_out so that wrap_out was applied
+# and a writable out arg gets wrap_out's dtype cast, shape assert, and view write-through
 for k,v in tiny_backend.items():
   name, _, overload = k.removeprefix("aten.").partition(".")
   op = getattr(getattr(aten, name), overload or "default")
   writes = [a.name for a in op._schema.arguments if a.alias_info is not None and a.alias_info.is_write]
   if torch.Tag.inplace_view in op.tags: fxn = wrap_inplace_view_op(v)
   elif writes == [op._schema.arguments[0].name] and op._schema.returns: fxn = wrap_inplace(v)
-  elif not writes or (writes == ["out"] and k in tiny_backend_out): fxn = wrap_fxn(k, v)
-  else: raise RuntimeError(f"{k} writes {writes}: expected an inplace first arg, or an out arg with {k} in tiny_backend_out")
+  elif not writes: fxn = wrap_fxn(k, v)
+  elif writes == ["out"]: fxn = wrap_fxn(k, wrap_out(v))
+  else: raise RuntimeError(f"{k} writes {writes}: unhandled writable arg in schema")
   torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(fxn)
 
 @torch.library.impl("aten::equal", "privateuseone")

extra/torch_backend/test.py

@@ -83,6 +83,12 @@ class TestTorchBackend(unittest.TestCase):
     torch.add(torch.ones(5, device=device), torch.ones(5, device=device), out=a)
     self.assertEqual(a.detach().storage_offset(), 3)
 
+  def test_out_refreshes_views_of_base(self):
+    a = torch.zeros(4, device=device)
+    v = a[2:]
+    torch.add(torch.ones(4, device=device), torch.ones(4, device=device), out=a)
+    np.testing.assert_equal(v.cpu().numpy(), [2., 2.])
+
   @unittest.expectedFailure  # TODO: storage offset assumes a contiguous source, use UOp.contiguous_view_offset
   def test_storage_offset_non_contiguous_source(self):
     a = torch.arange(12., device=device).reshape(3,4)
@@ -166,6 +172,15 @@ class TestTorchBackend(unittest.TestCase):
     expected = np.array([[1.5, 5.2, 9.0], [13.2, 17.1, 18.4]], dtype=np.float32)
     np.testing.assert_equal(y3.cpu().numpy(), expected)
 
+  def test_argmax_argmin(self):
+    a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
+    c = a.cpu()
+    for got, want in [(a.argmax(), c.argmax()), (a.argmin(0), c.argmin(0)), (a.argmax(1, keepdim=True), c.argmax(1, keepdim=True)),
+                      (torch.min(a, 1).indices, torch.min(c, 1).indices), (torch.max(a, 1).indices, torch.max(c, 1).indices),
+                      (torch.min(a, 1).values, torch.min(c, 1).values), (torch.min(a, 1, keepdim=True).indices, torch.min(c, 1, keepdim=True).indices)]:
+      self.assertEqual(got.dtype, want.dtype)  # torch's arg reduces are int64, tinygrad's are int32
+      np.testing.assert_equal(got.cpu().numpy(), want.numpy())
+
   def test_isfinite(self):
     a = torch.ones(4, device=device)
     np.testing.assert_equal(torch.isfinite(a).cpu().numpy(), [True, True, True, True])
@@ -373,6 +388,22 @@ class TestTorchBackend(unittest.TestCase):
     for bwd_eps in [1e-5, 0.3]:
       for got, want in zip(run(device, bwd_eps), run("cpu", bwd_eps)): np.testing.assert_allclose(got, want, atol=1e-4, rtol=1e-3)
 
+  def test_groupnorm_backward(self):
+    def run(dev):
+      x = torch.arange(24., device=dev).reshape(2, 4, 3).requires_grad_()
+      w = torch.linspace(0.5, 2.0, 4).to(dev).requires_grad_()
+      torch.nn.functional.group_norm(x, 2, w, torch.zeros(4, device=dev)).square().sum().backward()
+      return x.grad.cpu().numpy(), w.grad.cpu().numpy()
+    for got, want in zip(run(device), run("cpu")): np.testing.assert_allclose(got, want, atol=1e-4, rtol=1e-3)
+
+  def test_mse_smooth_l1_loss_backward(self):
+    def run(dev, loss):
+      x = torch.arange(4., device=dev).requires_grad_()
+      loss(x, torch.ones(4, device=dev)).backward()
+      return x.grad.cpu().numpy()
+    for loss in [torch.nn.functional.mse_loss, torch.nn.functional.smooth_l1_loss]:
+      np.testing.assert_allclose(run(device, loss), run("cpu", loss), atol=1e-6)
+
   def test_batchnorm_unsqueeze(self):
     bn = torch.nn.BatchNorm2d(4).to(device)
     x = torch.randn(8, 4, 3, 3, device=device)
@@ -516,6 +547,15 @@ class TestTorchBackend(unittest.TestCase):
     cpu_res = torch.arange(20, dtype=torch.float32)[::2][1:4].numpy()
     np.testing.assert_equal(torch_res, cpu_res)
 
+  def test_select_out_of_range_dim(self):
+    a = torch.arange(12, dtype=torch.int32, device=device).reshape(3, 4)
+    with self.assertRaises(IndexError): a.select(5, 0)
+
+  def test_select_collapses_the_only_dim(self):
+    a = torch.arange(3, dtype=torch.int32, device=device)
+    self.assertEqual(a.select(0, 1).shape, ())
+    np.testing.assert_equal(a.select(0, 1).cpu().numpy(), 1)
+
   def test_slice_negative_dim(self):
     a = torch.arange(13, dtype=torch.int32, device=device).repeat(8, 1)
     torch_chunks = a.chunk(3, -1)
@@ -796,6 +836,86 @@ class TestTorchBackend(unittest.TestCase):
     np.testing.assert_allclose(w_tiny.grad.cpu().numpy(), w_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)
     np.testing.assert_allclose(b_tiny.grad.cpu().numpy(), b_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)
 
+  def test_write_through_detach_of_unrealized(self):
+    a = torch.empty(4, device=device)
+    a.detach().fill_(3)
+    np.testing.assert_equal(a.cpu().numpy(), [3, 3, 3, 3])
+
+  def test_square_transpose_inplace(self):
+    # a same-shape transpose is not a reshape: writing the transposed values straight back would scramble the base
+    a = torch.tensor([[0., 1., 2.], [3., 4., 5.], [6., 7., 8.]], device=device)
+    a.transpose(0, 1).add_(100)
+    np.testing.assert_equal(a.cpu().numpy(), [[100., 101., 102.], [103., 104., 105.], [106., 107., 108.]])
+
+  def test_interpolate(self):
+    a = torch.arange(4, dtype=torch.float32, device=device).reshape(1, 1, 2, 2)
+    nearest = torch.nn.functional.interpolate(a, scale_factor=2.0)
+    np.testing.assert_equal(nearest.cpu().numpy()[0, 0], [[0, 0, 1, 1], [0, 0, 1, 1], [2, 2, 3, 3], [2, 2, 3, 3]])
+    linear = torch.nn.functional.interpolate(a, size=(4, 4), mode="bilinear", align_corners=False)
+    ref = torch.nn.functional.interpolate(a.cpu(), size=(4, 4), mode="bilinear", align_corners=False)
+    np.testing.assert_allclose(linear.cpu().numpy(), ref.numpy(), rtol=1e-5)
+
+  def test_interpolate_bicubic_area(self):
+    a = torch.arange(32, dtype=torch.float32, device=device).reshape(1, 2, 4, 4)
+    for mode, scale in [("bicubic", 2.0), ("area", 0.5)]:
+      ref = torch.nn.functional.interpolate(a.cpu(), scale_factor=scale, mode=mode)
+      np.testing.assert_allclose(torch.nn.functional.interpolate(a, scale_factor=scale, mode=mode).cpu().numpy(), ref.numpy(), atol=1e-4)
+
+  @unittest.expectedFailure
+  def test_interpolate_bicubic_backward(self):
+    # the forward comes from a decomposition, but aten::upsample_bicubic2d_backward has none (nor does
+    # aten::_adaptive_avg_pool2d_backward, for area), so training through these modes needs a real kernel
+    x = torch.arange(32., dtype=torch.float32, device=device).reshape(1, 2, 4, 4).requires_grad_()
+    torch.nn.functional.interpolate(x, scale_factor=2.0, mode="bicubic").sum().backward()
+
+  @unittest.expectedFailure
+  def test_interpolate_inexact_scale(self):
+    # torch forwards the raw scale_factor, Tensor.interpolate recomputes it from output_size, and they disagree here
+    a = torch.arange(6, dtype=torch.float32, device=device).reshape(1, 1, 2, 3)
+    tiny = torch.nn.functional.interpolate(a, scale_factor=2.5, mode="bilinear")
+    ref = torch.nn.functional.interpolate(a.cpu(), scale_factor=2.5, mode="bilinear")
+    np.testing.assert_allclose(tiny.cpu().numpy(), ref.numpy(), rtol=1e-5)
+
+  def test_logical_or_xor(self):
+    a = torch.tensor([True, True, False, False], device=device)
+    b = torch.tensor([True, False, True, False], device=device)
+    np.testing.assert_equal(torch.logical_or(a, b).cpu().numpy(), [True, True, True, False])
+    np.testing.assert_equal(torch.logical_xor(a, b).cpu().numpy(), [False, True, True, False])
+    # bool-valued whatever the input dtype, so this is not | and ^
+    i, j = torch.tensor([2, 0, 5, 0], device=device), torch.tensor([0, 0, 1, 1], device=device)
+    np.testing.assert_equal(torch.logical_or(i, j).cpu().numpy(), [True, False, True, True])
+    np.testing.assert_equal(torch.logical_xor(i, j).cpu().numpy(), [True, False, False, True])
+
+  def test_slice_scatter(self):
+    # the scatters are functional: they return a new tensor and must leave the one they were given alone
+    a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
+    out = torch.slice_scatter(a, torch.ones(1, 4, device=device), 0, 0, 1)
+    np.testing.assert_equal(out.cpu().numpy(), [[1, 1, 1, 1], [4, 5, 6, 7], [8, 9, 10, 11]])
+    np.testing.assert_equal(a.cpu().numpy(), np.arange(12, dtype=np.float32).reshape(3, 4))
+
+  def test_slice_scatter_casts_src(self):
+    a = torch.zeros(3, 4, device=device)
+    out = torch.slice_scatter(a, torch.ones(1, 4, dtype=torch.int32, device=device), 0, 0, 1)
+    self.assertEqual(out.dtype, torch.float32)
+    np.testing.assert_equal(out.cpu().numpy()[0], np.ones(4, dtype=np.float32))
+
+  def test_select_scatter(self):
+    a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
+    out = torch.select_scatter(a, torch.ones(4, device=device), 0, 1)
+    np.testing.assert_equal(out.cpu().numpy(), [[0, 1, 2, 3], [1, 1, 1, 1], [8, 9, 10, 11]])
+
+  def test_diagonal_scatter(self):
+    a = torch.zeros(3, 3, device=device)
+    out = torch.diagonal_scatter(a, torch.arange(3, dtype=torch.float32, device=device))
+    np.testing.assert_equal(out.cpu().numpy(), np.diag([0., 1., 2.]))
+    np.testing.assert_equal(a.cpu().numpy(), np.zeros((3, 3), dtype=np.float32))
+
+  def test_copy_functional(self):
+    # without an impl this segfaults rather than fails: a regression here takes the whole run down
+    a = torch.arange(4, dtype=torch.float32, device=device)
+    out = torch.ops.aten.copy(a, torch.zeros(4, device=device))
+    np.testing.assert_equal(out.cpu().numpy(), [0., 0., 0., 0.])
+    np.testing.assert_equal(a.cpu().numpy(), [0., 1., 2., 3.])
 
 from tinygrad import Tensor
 class TestBackendHelpers(unittest.TestCase):

pyproject.toml

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

test/amd/test_emu2_pcode.py

@@ -67,32 +67,26 @@ class TestParseExpr(unittest.TestCase):
 
   def test_integer_literals(self):
     """Test parsing integer literals."""
-    self.assertEqual(parse_expr('0', {}).val, 0)
-    self.assertEqual(parse_expr('42', {}).val, 42)
-    self.assertEqual(parse_expr('42U', {}).val, 42)
+    self.assertIs(parse_expr('0', {}), UOp.const(0, dtypes.uint32))
+    self.assertIs(parse_expr('42', {}), UOp.const(42, dtypes.uint32))
+    self.assertIs(parse_expr('42U', {}), UOp.const(42, dtypes.uint32))
 
   def test_negative_integers(self):
     """Test parsing negative integer literals."""
-    result = parse_expr('-1', {})
-    self.assertEqual(result.val, -1)
-    self.assertEqual(result.dtype, dtypes.int)
+    self.assertIs(parse_expr('-1', {}), UOp.const(-1, dtypes.int))
 
   def test_float_literals(self):
     """Test parsing float literals."""
-    result = parse_expr('1.0F', {})
-    self.assertEqual(result.val, 1.0)
-    self.assertEqual(result.dtype, dtypes.float32)
+    self.assertIs(parse_expr('1.0F', {}), UOp.const(1.0, dtypes.float32))
 
   def test_hex_literals(self):
     """Test parsing hex literals."""
-    result = parse_expr('0xFF', {})
-    self.assertEqual(result.val, 255)
+    self.assertIs(parse_expr('0xFF', {}), UOp.const(255, dtypes.uint32))
 
   def test_variable_lookup(self):
     """Test variable lookup in parse_expr."""
     vrs = {'x': UOp.const(42, dtypes.uint32)}
-    result = parse_expr('x', vrs)
-    self.assertEqual(result.val, 42)
+    self.assertIs(parse_expr('x', vrs), vrs['x'])
 
   def test_binary_ops(self):
     """Test parsing binary operations."""
@@ -103,9 +97,7 @@ class TestParseExpr(unittest.TestCase):
     self.assertEqual(result.op, Ops.ADD)
 
     # Subtraction with constant folding
-    result = parse_expr('10 - 5', {})
-    self.assertEqual(result.op, Ops.CONST)
-    self.assertEqual(result.val, 5)
+    self.assertIs(parse_expr('10 - 5', {}), UOp.const(5, dtypes.uint32))
 
   def test_ternary(self):
     """Test parsing ternary expressions."""
@@ -142,15 +134,8 @@ class TestForLoopParsing(unittest.TestCase):
     S0 = UOp.const(0, dtypes.uint32)
     _vrs, assigns = parse_pcode(pcode, {'S0': S0})
 
-    # Check that the innermost value (default) is -1 (may be wrapped in CAST)
-    val = assigns[0][1]
-    # Traverse to innermost WHERE
-    while val.op == Ops.WHERE:
-      val = val.src[2]  # false branch
-    # Unwrap CAST if present
-    while val.op == Ops.CAST:
-      val = val.src[0]
-    self.assertEqual(val.val, -1)
+    # every cond folds (S0 is a const), leaving the default branch: -1 in the destination dtype
+    self.assertIs(assigns[0][1].simplify(), UOp.const(-1, dtypes.uint32))
 
   def test_ctz_parsing(self):
     """Test CTZ pcode parsing."""
@@ -262,8 +247,8 @@ class TestDSPcodePatterns(unittest.TestCase):
     _, assigns = parse_pcode(pcode, srcs)
     # Check addresses: 100 + 2*4 = 108, 100 + 5*4 = 120
     # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp
-    self.assertEqual(assigns[0][1][0].simplify().val, 108)  # type: ignore[index]
-    self.assertEqual(assigns[1][1][0].simplify().val, 120)  # type: ignore[index]
+    self.assertIs(assigns[0][1][0].simplify(), UOp.const(108, dtypes.uint32))  # type: ignore[index]
+    self.assertIs(assigns[1][1][0].simplify(), UOp.const(120, dtypes.uint32))  # type: ignore[index]
 
   def test_ds_store_data_values(self):
     """Test DS_STORE_2ADDR_B32 uses correct data values."""
@@ -280,8 +265,8 @@ class TestDSPcodePatterns(unittest.TestCase):
     _, assigns = parse_pcode(pcode, srcs)
     # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp
     # DATA[31:0] should preserve the value
-    self.assertEqual(assigns[0][1][1].simplify().val, 0xAAAAAAAA)  # type: ignore[index]
-    self.assertEqual(assigns[1][1][1].simplify().val, 0xBBBBBBBB)  # type: ignore[index]
+    self.assertIs(assigns[0][1][1].simplify(), UOp.const(0xAAAAAAAA, dtypes.uint32))  # type: ignore[index]
+    self.assertIs(assigns[1][1][1].simplify(), UOp.const(0xBBBBBBBB, dtypes.uint32))  # type: ignore[index]
 
 class TestConditionalParsing(unittest.TestCase):
   """Test conditional (if/elsif/else) pcode parsing."""
@@ -306,12 +291,12 @@ class TestConcatWidthParsing(unittest.TestCase):
   def test_permlanex16_altrow_concat(self):
     for row, expected in [(0, 1), (1, 0), (2, 3), (3, 2)]:
       parsed = parse_expr('{ row[1], ~row[0] }', {'row': UOp.const(row, dtypes.uint32)})
-      self.assertEqual(parsed.simplify().val, expected)
+      self.assertIs(parsed.simplify(), UOp.const(expected, dtypes.uint32))
 
   def test_permlane64_altlane_concat(self):
     for lane, expected in [(0, 32), (1, 33), (31, 63), (32, 0), (63, 31)]:
       parsed = parse_expr('{ ~lane[5], lane[4:0] }', {'lane': UOp.const(lane, dtypes.uint32)})
-      self.assertEqual(parsed.simplify().val, expected)
+      self.assertIs(parsed.simplify(), UOp.const(expected, dtypes.uint32))
 
   def test_permlane64_wave64_pcode_indices(self):
     vgpr = UOp.param(0, dtypes.uint32, (256,))
@@ -327,19 +312,17 @@ class TestConcatWidthParsing(unittest.TestCase):
       'S2': UOp.const(0, dtypes.uint32),
     }
 
-    def load_idx(v: UOp) -> int:
+    def check_load_idx(v: UOp, expected: int):
       simp = v.simplify()
       self.assertEqual(simp.op, Ops.LOAD)
       self.assertEqual(simp.src[0].op, Ops.INDEX)
-      idx = simp.src[0].src[1].simplify()
-      self.assertEqual(idx.op, Ops.CONST)
-      return idx.val
+      self.assertIs(simp.src[0].src[1].simplify(), UOp.const(expected, dtypes.uint32))
 
     _, assigns = parse_pcode(PCODE[VOP1Op.V_PERMLANE64_B32_E32], srcs)
     self.assertEqual(len(assigns), 64)
     for lane, (dst_idx, src_idx) in {0: (64, 32), 31: (95, 63), 32: (96, 0), 63: (127, 31)}.items():
-      self.assertEqual(assigns[lane][1][0].simplify().val, dst_idx)  # type: ignore[index]
-      self.assertEqual(load_idx(assigns[lane][1][1]), src_idx)  # type: ignore[index]
+      self.assertIs(assigns[lane][1][0].simplify(), UOp.const(dst_idx, dtypes.uint32))  # type: ignore[index]
+      check_load_idx(assigns[lane][1][1], src_idx)  # type: ignore[index]
 
 class TestAllPcode(unittest.TestCase):
   """Test that all pcode from all architectures can be parsed."""

test/backend/test_asm_gemm.py

@@ -1,4 +1,5 @@
 import unittest
+import functools
 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.helpers import getenv, system, DEV
 from extra.gemm.cdna_asm_gemm import asm_gemm, hk_bf16_atb_gemm
@@ -9,6 +10,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")
 
+@functools.cache
 def has_hipcc():
   try: system("hipcc --version")
   except Exception: return False

test/backend/test_const_folding.py

@@ -1,7 +1,7 @@
 import unittest, math
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.dtype import DTYPES_DICT
-from tinygrad.uop.ops import Ops, UOp
+from tinygrad.uop.ops import Ops, UOp, GroupOp
 from tinygrad.codegen.decomp.op import threefry2x32
 import numpy as np
 from test.helpers import not_support_multi_device
@@ -17,7 +17,7 @@ def _check_ast_count(desired_count:int, t:Tensor):
 class TestMovedConstFolding(unittest.TestCase):
   def test_contiguous_deviceless_const(self):
     t = Tensor(UOp.const(2.0, dtypes.float)).contiguous()
-    self.assertIs(t.uop.op, Ops.CONST)
+    self.assertIs(t.uop, UOp.const(2.0, dtypes.float))
     self.assertIsNone(t.uop.device)
 
   def test_add_shrunk_zero(self):
@@ -169,8 +169,8 @@ class TestMultiConstFolding(unittest.TestCase):
 class TestThreefryConstFolding(unittest.TestCase):
   def test_threefry(self):
     # THREEFRY(const,const) folds to a const once decomposed
-    x = threefry2x32(UOp.const(5, dtypes.uint64), UOp.const(10, dtypes.uint64))
-    self.assertIs(x.simplify().op, Ops.CONST)
+    x = threefry2x32(UOp.const(5, dtypes.uint64), UOp.const(10, dtypes.uint64)).simplify()
+    self.assertEqual([u.op for u in x.toposort() if u.op in GroupOp.ALU], [])
 
 class TestTautologicalCompare(unittest.TestCase):
   # without const folding, these would have triggered -Wtautological-compare in clang
@@ -188,7 +188,6 @@ class TestTautologicalCompare(unittest.TestCase):
     np.testing.assert_equal((Tensor(True) < Tensor(False)).numpy(), False)
     np.testing.assert_equal((Tensor(True) < Tensor(True)).numpy(), False)
 
-  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support NaN comparison correctly")
   def test_a_eq_a(self):
     # self eq is always true for int or bool
     a = Tensor([1, 2, 3])

test/backend/test_custom_kernel.py

@@ -4,7 +4,7 @@ import numpy as np
 from tinygrad.dtype import AddrSpace, dtypes, Invalid
 from tinygrad.uop.ops import KernelInfo, AxisType, Ops
 from tinygrad.renderer.ptx import PTXRenderer
-from test.helpers import assert_kernel_count
+from test.helpers import assert_kernel_count, KernelCountException
 
 # **** kernels ****
 
@@ -422,9 +422,8 @@ class TestCustomKernel(unittest.TestCase):
       return Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
     GlobalCounters.reset()
     y = run(x[0]).realize()
-    # it's copying the input and the output
-    # TODO: subbuffer usage has runtime specific behavior, this will be fixed after the removal of SLICE.
-    assert_kernel_count(2 if y.device in ("CL", "WEBGPU") else 1)
+    # backends that support contiguous views don't launch extra kernels
+    assert_kernel_count(2 if x[0].uop.contiguous_view() is None else 1)
     self.assertEqual(y.tolist(), [1, 2, 3, 4])
 
   @Context(DEV="CPU")
@@ -475,7 +474,7 @@ class TestCustomKernelInput(unittest.TestCase):
     y.realize()
     kernel_count = GlobalCounters.kernel_count
     self.assertEqual(y.tolist(), x.add(1).tolist())
-    self.assertLessEqual(kernel_count, max_kernels)
+    if kernel_count > max_kernels: raise KernelCountException(max_kernels, kernel_count)
     # same test with @function, input is PARAM
     from tinygrad import function
     x0 = Tensor.arange(32).clone("CPU").realize()
@@ -488,7 +487,7 @@ class TestCustomKernelInput(unittest.TestCase):
     y = run(x0).realize()
     kernel_count = GlobalCounters.kernel_count
     self.assertEqual(y.tolist(), mop_fxn(x0).add(1).tolist())
-    self.assertLessEqual(kernel_count, max_kernels)
+    if kernel_count > max_kernels: raise KernelCountException(max_kernels, kernel_count)
 
   def test_reshape(self): self._test_mop(lambda x: x.reshape(16, 2), max_kernels=2)
   def test_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T, max_kernels=3)

test/backend/test_dtype.py

@@ -340,6 +340,9 @@ class TestUint64DType(TestDType):
   DTYPE = dtypes.uint64
   def test_uint64_load(self):
     assert Tensor(2**64 - 1, dtype=dtypes.uint64).numpy() == 2**64 - 1
+  @unittest.skipIf(dtypes.double not in supported_dtypes, "needs float64")
+  def test_uint64_cast_double(self):
+    assert Tensor([2**32 + 1], dtype=dtypes.uint64).cast(dtypes.double).numpy() == 2**32 + 1
 
 @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX does indexing math with longs")
 class TestEmulatedUInt64DType(TestUint64DType):

test/backend/test_isel.py

@@ -7,7 +7,7 @@ from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops
 from tinygrad.renderer.isa import IselContext
 
 # INDEX on a register value with a constant index extracts a single element (the old GEP)
-def lane(y:UOp, i:int) -> UOp: return y.index(UOp.const(i, dtypes.int), dtype=y.dtype.scalar())
+def lane(y:UOp, i:int) -> UOp: return y.index(UOp.cconst(i, dtypes.int), dtype=y.dtype)
 
 @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only x86")
 class TestIselX86(unittest.TestCase):
@@ -46,10 +46,10 @@ class TestIselX86(unittest.TestCase):
   # complex address is [base + index*scale + displacement]
   def test_complex_address(self):
     a = UOp.variable("a", 0, 0, dtypes.int32)
-    load = UOp.param(0, dtypes.int32, (16,)).index(a + 1).load()
+    load = UOp.param(0, dtypes.int32, (16,)).index(a + UOp.cconst(1, dtypes.int32)).load()
     n = self.isel_rewrite(load)
     # displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32
-    self.assertTrue(n.src[2].op is Ops.CONST and n.src[2].dtype is dtypes.int8 and n.src[2].val == 4)
+    self.assertTrue(n.src[2].dtype is dtypes.int8 and n.src[2].src[0].op is Ops.CONST and n.src[2].src[0].val == 4)
 
 if __name__ == "__main__":
   unittest.main()

test/backend/test_jit.py

@@ -360,7 +360,7 @@ class TestJitGraphSplit(unittest.TestCase):
     self.expect(f, inp, inp_cpu,
       graph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()],
       multigraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()],
-      hcqgraph=[self.ji_graph(4)])
+      hcqgraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()]) # cpu is hcq2 now, it does not join hcq graphs
 
   def test_jit_cpu_several(self):
     if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
@@ -377,9 +377,9 @@ class TestJitGraphSplit(unittest.TestCase):
     inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
     inp_cpu = Tensor.randn(10, 10, device="CPU").realize()
     self.expect(f, inp, inp_cpu,
-      graph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
-      multigraph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
-      hcqgraph=[self.ji_graph(5)])
+      graph=[self.ji_graph(2), self.ji_comp(), self.ji_comp(), self.ji_comp()],
+      multigraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp(), self.ji_comp()],
+      hcqgraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp(), self.ji_comp()])
 
   def test_jit_multidev(self):
     if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")

test/backend/test_linearizer.py

@@ -16,8 +16,6 @@ from test.helpers import replace_opts, check_schedule
 from test.backend.test_softmax_fusion import single_kernel_softmax
 MOCKGPU = DEV.interface.startswith("MOCK")
 
-from tinygrad.uop.render import print_uops # noqa: F401 # pylint: disable=unused-import
-
 @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, ISARenderer), "isa backends don't preserve the op spec when lowering")
 class TestLinearizer(unittest.TestCase):
   def test_arg_dedup(self):
@@ -30,7 +28,7 @@ class TestLinearizer(unittest.TestCase):
     c = ((a.shrink(((0, 2),)) - a.shrink(((2, 4),))) - (b.shrink(((0, 2),)) - b.shrink(((2, 4),))))
     linear = c.schedule_linear()
     run_linear(linear)
-    rawbufs = [s.buffer for s in linear.src[-1].src[1:] if s.op is not Ops.BIND]
+    rawbufs = [s.buffer for s in linear.src[-1].src[1:] if not s.is_bound_var]
     assert len(rawbufs) == 3 and set(rawbufs[1:]) == {a.uop.base.realized, b.uop.base.realized}
     np_c = (np_a[:2] - np_a[2:]) - (np_b[:2] - np_b[2:])
     np.testing.assert_allclose(np_c, c.numpy(), atol=1e-4, rtol=1e-4)
@@ -248,11 +246,10 @@ class TestLinearizer(unittest.TestCase):
     uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[1].src)
     begin_range = [i for i, x in enumerate(uops) if x.op is Ops.RANGE][-1]
     end_range = [i for i, x in enumerate(uops) if x.op is Ops.END][0]
-    for i,u in enumerate(uops): print(i, u.op, [uops.index(s) for s in u.src], u.arg, u.dtype)
     for u in uops:
       if u.op is Ops.STORE and u.src[0].addrspace is AddrSpace.REG:
         if uops.index(u) < begin_range:
-          assert u.src[1].op is Ops.CONST
+          assert u.src[1].op not in GroupOp.ALU
         else:
           assert u.src[1].op in GroupOp.ALU
           assert begin_range < uops.index(u) < end_range
@@ -268,9 +265,9 @@ class TestLinearizer(unittest.TestCase):
     uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
     idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
     idxs = sorted(idxs, key=lambda uop: uop.arg)
-    assert (idxs[0].arg, idxs[0].src[0].val) == ('gidx0', 6), idxs[0]
-    assert (idxs[1].arg, idxs[1].src[0].val) == ('gidx1', 5), idxs[1].arg
-    assert (idxs[2].arg, idxs[2].src[0].val) == ('gidx2', 4), idxs[2].arg
+    assert (idxs[0].arg, idxs[0].src[0].src[0].val) == ('gidx0', 6), idxs[0]
+    assert (idxs[1].arg, idxs[1].src[0].src[0].val) == ('gidx1', 5), idxs[1].arg
+    assert (idxs[2].arg, idxs[2].src[0].src[0].val) == ('gidx2', 4), idxs[2].arg
 
   def test_sum_collapse(self):
     t = Tensor([2]).reshape(1, 1).expand(256, 256).sum()
@@ -411,7 +408,7 @@ def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
   last_call = linear.src[-1]
   ast = last_call.src[0]
   assert ast.op is Ops.SINK, f"helper_realized_ast expects a SINK {last_call}"
-  last_bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]
+  last_bufs = [s.buffer for s in last_call.src[1:] if not s.is_bound_var]
   # now all input buffers in last_call should be realized
   # create fresh buffers for the outputs
   bufs = [Buffer(x.device, x.size, x.dtype).allocate() if i < len(ast.src) else x for i,x in enumerate(last_bufs)]
@@ -437,7 +434,7 @@ def reset_bufs(bufs:list[Buffer]):
   for buf in bufs: buf.copy_from(Buffer("PYTHON", buf.size, buf.dtype, opaque=memoryview(bytearray(buf.nbytes))))
 
 def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[],
-                               apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[]):
+                               apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[], check_default_opt=True):
   outbufs = real_bufs[:len(realized_ast.src)]
   wanna_output = [np.array(x).flatten() for x in wanna_output]
   buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in real_bufs]
@@ -459,9 +456,7 @@ def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[]
     for buf,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(buf, want, atol=atol, rtol=rtol)
 
   # Check correctness of handcoded optimiztions.
-  reset_bufs(outbufs)
-  run_prg(opts=None)
-  for buf,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(buf, want, atol=atol, rtol=rtol)
+  if check_default_opt: check_opt(None)
   for x in opts: # Check custom transformations if any.
     check_opt(([Opt(OptOps.TC, 0, (TC_SELECT.value, TC_OPT.value, 1))] if apply_tc else [])+x)
 

test/backend/test_llama_kernels.py

@@ -99,22 +99,20 @@ class TestLocalAmax(unittest.TestCase):
     assert_kernel_count(2)
     self.assertEqual(out.tolist(), [[0., 7., 14., 21.], [28., 35., 42., 49.], [120., 135., 150., 165.], [180., 195., 210., 225.]])
 
-@unittest.skipUnless(has_hipcc() and Device.DEFAULT == "AMD", "requires hipcc to compile and amd device to run")
 class TestFusedQKVRoPE(unittest.TestCase):
   SHAPE = (2, 8192, 32, 8, 128)
 
+  def setUp(self):
+    if dtypes.bfloat16 not in Device[Device.DEFAULT].renderer.supported_dtypes(): self.skipTest("test uses bf16 inputs")
+
   def rand_bf16(self, *shape:int) -> Tensor:
     return (Tensor.randn(*shape) * 0.1).cast(dtypes.bfloat16).contiguous().realize()
 
-  def freqs_cis(self) -> Tensor:
-    _, N, _, _, D = self.SHAPE
-    return precompute_freqs_cis(D, N * 2).cast(dtypes.bfloat16).clone().realize()
-
-  def test_llama31_8b_forward(self):
+  def test_forward(self):
     Tensor.manual_seed(0)
-    B, N, H, H_KV, D = self.SHAPE
+    B, N, H, H_KV, D = 1, 32, 8, 2, 16
     GROUP = H // H_KV
-    freqs_cis = self.freqs_cis()
+    freqs_cis = (Tensor.randn(1, N * 2, 1, D // 2, 2) * 0.1).cast(dtypes.bfloat16).contiguous().realize()
 
     x = self.rand_bf16(B, N, H_KV * (GROUP + 2) * D)
     q, k, v = fused_qkv_rope(x, freqs_cis, H, H_KV, D)
@@ -131,12 +129,13 @@ class TestFusedQKVRoPE(unittest.TestCase):
       self.assertTrue(k.allclose(k_ref, atol=2e-2, rtol=0).item(), "K forward mismatch")
       self.assertTrue(v.allclose(v_ref, atol=0, rtol=0).item(), "V forward mismatch")
 
-  def test_llama31_8b_backward(self):
+  @unittest.skipUnless(has_hipcc(), "backward kernel requires hipcc to compile")
+  def test_llama31_8b(self):
     Tensor.manual_seed(1)
     B, N, H, H_KV, D = self.SHAPE
     PARTIALS = 2
     GROUP = H // H_KV
-    freqs_cis = self.freqs_cis()
+    freqs_cis = precompute_freqs_cis(D, N * 2).cast(dtypes.bfloat16).clone().realize()
     dq = self.rand_bf16(B, N, H, D)
     dk_partial = self.rand_bf16(B * PARTIALS, N, H_KV, D)
     dv_partial = self.rand_bf16(B * PARTIALS, N, H_KV, D)

test/backend/test_multitensor.py

@@ -1,12 +1,12 @@
 import unittest, random
 from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variable
-from tinygrad.uop.ops import Ops, UOp, AxisType
+from tinygrad.uop.ops import Ops, UOp, AxisType, graph_rewrite
 from tinygrad.helpers import getenv, prod, Context
 from tinygrad.nn.state import get_parameters
-from tinygrad.engine.realize import run_linear, compile_linear
+from tinygrad.engine.realize import run_linear, compile_linear, pm_beam, pm_compile
 import numpy as np
 from hypothesis import given, strategies as strat, settings
-from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule, assert_kernel_count
+from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule, assert_kernel_count, KernelCountException
 
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
@@ -79,9 +79,9 @@ class TestMultiTensor(unittest.TestCase):
   def test_shard_beam(self):
     cpu_2 = ("CPU:1", "CPU:2")
     src = Tensor.ones(16).shard(cpu_2, 0).realize()
-    pad = src.to(cpu_2[::-1]).schedule_linear().src[0]
-    with Context(BEAM=1, IGNORE_BEAM_CACHE=1): prg = compile_linear(UOp(Ops.LINEAR, src=(pad,))).src[0].src[0]
-    self.assertNotEqual(prg.src[0].arg.applied_opts, ())
+    lin = UOp(Ops.LINEAR, src=(src.to(cpu_2[::-1]).schedule_linear().src[0],))
+    with Context(BEAM=1, IGNORE_BEAM_CACHE=1): call = graph_rewrite(graph_rewrite(lin, pm_beam, ctx=1, walk=True), pm_compile, walk=True).src[0]
+    self.assertNotEqual(call.src[0].src[0].arg.applied_opts, ())
 
   def test_shard_same_device(self):
     X = Tensor.ones(256).contiguous().realize()
@@ -395,7 +395,7 @@ class TestMultiBufferView(unittest.TestCase):
     linear, var_vals = b_multi.linear_with_vars()
     if all(not d.startswith(("WEBGPU", "CL")) for d in b_multi.device):
       compiled = [call for call in linear.src if call.src[0].op is Ops.SINK]
-      self.assertEqual(len(compiled), 0, f"expected zero compiled kernels, got {len(compiled)}")
+      if len(compiled) != 0: raise KernelCountException(0, len(compiled))
     run_linear(linear, var_vals)
     np.testing.assert_equal(b_multi.numpy(), b_ref.numpy())
 

test/backend/test_ops.py

@@ -720,10 +720,11 @@ class TestOps(unittest.TestCase):
       return torch.autograd.grad(t ** c, t)[0].item()
     for x in [-math.inf, 0, 1, math.inf]:
       for c in [-1, 0, 0.3, 1, 2]:
-        tiny_out = get_tiny_gradient(x, c)
         torch_out = get_torch_gradient(x, c)
+        # the pow backward routes through exp2/log2, whose 0/inf behavior is undefined on WEBGPU
+        if Device.DEFAULT == "WEBGPU" and not math.isfinite(torch_out): continue
+        tiny_out = get_tiny_gradient(x, c)
         if math.isnan(tiny_out):
-          if Device.DEFAULT == "WEBGPU": continue # TODO: WEBGPU issue with nan
           assert math.isnan(torch_out)
         else:
           self.assertAlmostEqual(tiny_out, torch_out, msg=f"{x}, {c}")
@@ -749,6 +750,7 @@ class TestOps(unittest.TestCase):
   def test_exp2_log2_zero_times_negative(self):
     # gallivm's exp2/log2 have "undefined behavior with infs, 0s and nans", so exp2(log2(0)*y) returns 0 instead of inf
     helper_test_op(None, lambda x,y: (x.log2()*y).exp2(), lambda x,y: (x.log2()*y).exp2(), vals=[[0.0], [-0.7]], forward_only=True)
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "pow at 0 routes through exp2/log2, whose 0/inf behavior is undefined on WEBGPU")
   def test_pow_zero_const(self):
     helper_test_op(None, lambda x: x**0.3, vals=[[0.0]])
     helper_test_op(None, lambda x: x**0.0, vals=[[0.0]])
@@ -2162,6 +2164,10 @@ class TestOps(unittest.TestCase):
   def test_roll(self):
     helper_test_op([(2, 4)], lambda x: x.roll(1))
     helper_test_op([(2, 4)], lambda x: x.roll((1,)))
+    helper_test_op([(0,)], lambda x: x.roll(1, 0))
+    helper_test_op([(2, 0, 3)], lambda x: x.roll(1, 0))
+    helper_test_op([(2, 0, 3)], lambda x: x.roll(1, 1))
+    helper_test_op([(2, 0, 3)], lambda x: x.roll(1))
     self.helper_test_exception([(2, 4)], lambda x: x.roll((1, 2)), expected=RuntimeError)
     helper_test_op([(2, 4)], lambda x: x.roll(1, 0))
     helper_test_op([(2, 4)], lambda x: x.roll(-1, 0))

test/backend/test_pickle.py

@@ -2,7 +2,8 @@ import unittest, pickle, types, tracemalloc
 import numpy as np
 from tinygrad import Tensor, Device, TinyJit, Variable, dtypes
 from tinygrad.helpers import GlobalCounters, ContextVar, Context, DEV
-from tinygrad.uop.ops import PatternMatcher, UPat, UOp
+from tinygrad.uop.ops import PatternMatcher, UPat, UOp, deconstruct_function
+from test.helpers import KernelCountException
 
 class TestPickle(unittest.TestCase):
   def test_pickle_code_object(self):
@@ -11,6 +12,11 @@ class TestPickle(unittest.TestCase):
     fxn = types.FunctionType(pickle.loads(code_str), globals())
     self.assertEqual(fxn(2), 4)
 
+  def test_deconstruct_function_nested_comprehension(self):
+    # pre PEP 709, each comprehension is its own code object, so dtypes here is referenced two code objects deep
+    def fxn(): return [[dtypes.int for _ in range(2)] for _ in range(2)]
+    self.assertEqual(types.FunctionType(*deconstruct_function(fxn))(), fxn())
+
   def test_pickle_pattern_matcher(self):
     pm = PatternMatcher([(UPat.cvar('x'), lambda x: x*2)])
     sink = UOp.const(2)
@@ -36,7 +42,7 @@ class TestPickle(unittest.TestCase):
     t2:Tensor = pickle.loads(st)
     np.testing.assert_equal(t_values, t2.numpy())
     # expect at most one COPY kernel
-    self.assertLessEqual(GlobalCounters.kernel_count, 1)
+    if GlobalCounters.kernel_count > 1: raise KernelCountException(1, GlobalCounters.kernel_count)
 
   def test_pickle_realized_tensor_alt(self):
     print("** init")

test/backend/test_quantize_onnx.py

@@ -82,7 +82,7 @@ class TestQuantizeOnnxCPU(unittest.TestCase):
       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]
-      assert all(u.dtype.scalar() is dtypes.int for u in daccs)
+      assert all(u.dtype is dtypes.int for u in daccs)
 
 @unittest.skipIf(Device.DEFAULT != "DSP", "only tests for DSP")
 class TestQuantizeOnnx(unittest.TestCase):

test/backend/test_tensor.py

@@ -653,9 +653,19 @@ class TestZeroShapeTensor(unittest.TestCase):
     np.testing.assert_equal(Tensor([[1, 2]]).pad_to(2, 3).numpy(), [[1, 2, 0], [0, 0, 0]])
     np.testing.assert_equal(Tensor([[1, 2]]).pad_to(1, 3).numpy(), [[1, 2, 0]])
     np.testing.assert_equal(Tensor([[1, 2]]).pad_to(None, 3).numpy(), [[1, 2, 0]])
+    np.testing.assert_equal(Tensor([1, 2]).pad_to(4, value=2).numpy(), [1, 2, 2, 2])
+    np.testing.assert_equal(Tensor([[1, 2]]).pad_to(2, 3, value=-1).numpy(), [[1, 2, -1], [-1, -1, -1]])
+    np.testing.assert_equal(Tensor([1, 2]).pad_to(None, value=5).numpy(), [1, 2])  # no-op pad ignores the fill
     with self.assertRaises(ValueError): Tensor([1, 2]).pad_to(2, 3)
     with self.assertRaises(ValueError): Tensor([[1, 2]]).pad_to(3)
 
+  def test_max_shape(self):
+    from tinygrad import UOp
+    t = Tensor.empty(2, UOp.variable('v', 1, 32), 4)
+    self.assertEqual(t.max_shape, (2, 32, 4))
+    self.assertEqual(t.max_numel(), 2*32*4)
+    self.assertEqual(Tensor.empty(2, 3).max_shape, (2, 3))
+
   def test_shrink_into_zero(self):
     t = Tensor.rand(3, 4).realize()
     assert t.shrink((None, (2, 2))).realize().shape == (3, 0)

test/device/test_hcq2.py

@@ -0,0 +1,22 @@
+import unittest, numpy as np
+from unittest.mock import patch
+from tinygrad import Device, Tensor
+from tinygrad.helpers import getenv
+from tinygrad.runtime.support.hcq2 import HCQ_DEVS, all_devices_in
+
+@unittest.skipUnless(getenv("HCQ2") and all_devices_in(Device.DEFAULT, HCQ_DEVS), "hcq2 device required")
+class TestHCQ2(unittest.TestCase):
+  def test_copy_without_copy_queue(self):
+    with patch.object(Device[Device.DEFAULT], "has_copy_queue", False):
+      np.testing.assert_equal(Tensor(np.arange(61, dtype=np.float32)).to(Device.DEFAULT).contiguous().realize().numpy(), np.arange(61))
+
+  def test_overlapping_device_tuples(self):
+    # an op on a wide device tuple followed by an op on an overlapping smaller tuple used to MMU-fault the smaller one
+    d4, d2 = tuple(f"{Device.DEFAULT}:{i}" for i in range(4)), tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
+    ref = Tensor.arange(16).contiguous().realize()
+    Tensor(ref.uop.copy_to_device(d4)).realize()
+    out = Tensor.ones(8).shard(d2, axis=0).contiguous().realize()
+    np.testing.assert_equal(out.numpy(), np.ones(8))
+
+if __name__ == "__main__":
+  unittest.main()

test/external/external_benchmark_all2all.py

@@ -0,0 +1,17 @@
+from tinygrad import Device, Tensor, TinyJit, dtypes
+from tinygrad.helpers import Timing, Context
+
+GPUS, DEPTH, SZ = 8, 4, 128 * 2**20
+WARMUP, ITERS = 3, 5
+devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(GPUS))
+bufs = tuple(Tensor.empty(SZ, dtype=dtypes.uint8, device=dev).contiguous().realize() for _ in range(DEPTH) for dev in devs)
+
+@TinyJit
+def all_to_all(*srcs:Tensor): return Tensor.realize(*(src.to(dst) for i,src in enumerate(srcs) for j,dst in enumerate(devs) if i % GPUS != j))
+
+if __name__ == "__main__":
+  with Context(ALL2ALL=1, JIT_BATCH_SIZE=0):
+    for i in range(-WARMUP, ITERS):
+      with Timing("ALL2ALL ", lambda ns: f" {SZ*GPUS*(GPUS-1)*DEPTH/ns:.2f} GB/s", enabled=i>=0):
+        all_to_all(*bufs)
+        for dev in devs: Device[dev].synchronize()

test/external/external_test_schedule_scaling.py

@@ -1,5 +1,5 @@
-import unittest, time
-from tinygrad import Tensor
+import unittest, time, itertools
+from tinygrad import Tensor, Context
 
 class TestScheduleScaling(unittest.TestCase):
   """Test that .schedule() scales linearly with graph size (no O(n^2) behavior)."""
@@ -130,5 +130,18 @@ class TestScheduleScaling(unittest.TestCase):
       return parts[0].cat(*parts[1:])
     self._assert_linear(concat_chain)
 
+  @Context(DEV="NULL:HIP:gfx1100")
+  def test_custom_kernel_assign_scaling(self):
+    from tinygrad.uop.ops import UOp, Ops, KernelInfo
+    from tinygrad.runtime.autogen.amd.rdna3.ins import s_nop
+    count = itertools.count(0)
+    def custom_kernel_assign(n):
+      def custom_asm(out):
+        return UOp(Ops.PROGRAM, src=(UOp.sink(out, arg=KernelInfo(f"fxn_{next(count)}")),
+                                     UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=s_nop(i)) for i in range(n*8)))))
+      call = Tensor.custom_kernel(Tensor.empty(1), fxn=custom_asm)[0]
+      return Tensor.cat(*[Tensor.empty(1).assign(call+i) for i in range(n)])
+    self._assert_linear(custom_kernel_assign, n_small=50, n_large=500)
+
 if __name__ == '__main__':
   unittest.main(verbosity=2)

test/external/external_test_usb_asm24.py

@@ -6,7 +6,7 @@ import numpy as np
 class TestDevCopySpeeds(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
-    cls.sz = getenv("SIZE", 2e6)
+    cls.sz = getenv("SIZE", 2000000)
     cls.dev = Device["AMD"]
     if not cls.dev.is_usb(): raise unittest.SkipTest("only test this on USB devices")
 

test/external/external_uop_gc.py

@@ -44,7 +44,7 @@ def realized_matmul():
   z = y.matmul(x)
   Tensor.realize(z)
 def realized_gradient():
-  x = Tensor.eye(3)
+  x = Tensor.eye(3).clone()
   y = Tensor([[2.0,0,-2.0]])
   z = y.matmul(x).sum()
   z.backward()

test/helpers.py

@@ -48,6 +48,8 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
   else:
     assert isinstance(t, UOp), f"can't schedule {t}"
     linear, var_vals = Tensor(t).linear_with_vars()
+  # test compiling the linear
+  compile_linear(linear)
   kernel_cnt = sum((len(call.device) if isinstance(call.device, tuple) else 1)
                    for call in linear.src if call.src[0].op is Ops.SINK or not filter_sink)
   if kernel_cnt != allowed:
@@ -57,8 +59,6 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
         print("kernel", i+1)
         print(call.src[0])
     raise KernelCountException(allowed, kernel_cnt)
-  # test compiling the linear
-  compile_linear(linear)
   return linear, var_vals
 
 def assert_kernel_count(expected:int):
@@ -86,7 +86,9 @@ def assert_jit_cache_len(fxn, expected_len):
   if linear is None or not linear.src:
     if expected_len != 0: raise KernelCountException(expected_len, 0)
     return
-  if expected_len and all(call_is_hcq(call) for call in linear.src): expected_len = 3 # HCQ2: merged same-queue calls + finalizer + bumps
+  if expected_len and all(call_is_hcq(call) for call in linear.src): # HCQ2: one batch submitter, or fence + reset + merged calls + finalizer
+    from tinygrad.runtime.support.hcq2 import HCQ_RUNTIME_DEV
+    expected_len = 1 if HCQ_RUNTIME_DEV.value == "CPU" else 4
   if call_is_graph(linear.src[0]):
     if len(linear.src) != 1: raise KernelCountException(1, len(linear.src))
     inner = linear.src[0].src[0].src[0]  # LINEAR UOp inside CUSTOM_FUNCTION

test/mockgpu/amd/emu.py

@@ -260,19 +260,6 @@ def _cond(cond, if_true, if_false):
 def _cond_hi16(cond, val: UOp) -> UOp: return _cond(cond, _hi16(val), val)
 def _apply_opsel(val: UOp, sel_bit: int, opsel: int) -> UOp: return _hi16(val) if opsel & (1 << sel_bit) else val
 
-def _set_lane_bit(old: UOp, lane: UOp, val: UOp, exec_mask: UOp) -> UOp:
-  """Set/clear a single bit in a mask based on lane index, respecting exec mask."""
-  if old.dtype in (dtypes.uint64, dtypes.int64):
-    dt = dtypes.uint64
-    mask = UOp.const(1, dt) << lane.cast(dt)
-    new_bit = _to_u32(val).cast(dt) << lane.cast(dt)
-    cleared = old.cast(dt) & (mask ^ UOp.const(0xFFFFFFFFFFFFFFFF, dt))
-    return _lane_active(exec_mask, lane).where(cleared | new_bit, old.cast(dt))
-  mask = _c(1) << lane.cast(dtypes.uint32)
-  new_bit = _to_u32(val) << lane.cast(dtypes.uint32)
-  cleared = old & (mask ^ _c(MASK32))
-  return _lane_active(exec_mask, lane).where(cleared | new_bit, old)
-
 def _val_to_u32(val: UOp) -> UOp:
   """Convert any value to uint32 for storage (bitcast floats, cast ints)."""
   if val.dtype == dtypes.uint32: return val
@@ -532,6 +519,19 @@ class _Ctx:
       return [self.wsgpr_dyn(reg, lo), self.wsgpr_dyn(reg + _c(1), hi)]
     return [self.wsgpr_dyn(reg, val)]
 
+  def wmask_lane_bit(self, reg: UOp, lane: UOp, val: UOp, exec_mask: UOp) -> list[UOp]:
+    """Set/clear bit `lane` of the mask at `reg` from val for exec-active lanes, preserving memory for inactive lanes"""
+    active, bit = _lane_active(exec_mask, lane), _to_u32(val)
+    if self.wave_size <= 32:
+      old = self.rsgpr_dyn(reg)
+      mask = _c(1) << lane.cast(dtypes.uint32)
+      return [self.wsgpr_dyn(reg, active.where((old & (mask ^ _c(MASK32))) | (bit << lane.cast(dtypes.uint32)), old))]
+    off = (lane & _c(31, dtypes.int)).cast(dtypes.uint32)
+    mask = _c(1) << off
+    def half(old: UOp, sel: UOp) -> UOp: return sel.where(active.where((old & (mask ^ _c(MASK32))) | (bit << off), old), old)
+    return [self.wsgpr_dyn(reg, half(self.rsgpr_dyn(reg), lane < _c(32, dtypes.int))),
+            self.wsgpr_dyn(reg + _c(1), half(self.rsgpr_dyn(reg + _c(1)), _c(32, dtypes.int) <= lane))]
+
   def rmask(self, reg: UOp) -> UOp:
     """Read a lane mask (VCC/EXEC). Combines lo/hi for wave64."""
     if self.wave_size > 32: return _u64(self.rsgpr_dyn(reg), self.rsgpr_dyn(reg + _c(1)))
@@ -718,9 +718,7 @@ class _Ctx:
         raw_stores.append(('vgpr_direct', self.vgpr.index(val[0].valid(active)).store(new_val)))
         continue
       if 'D0' in dest and '[laneId]' in dest:
-        old_vcc = self.rmask(_c(VCC_LO.offset))
-        new_vcc = _set_lane_bit(old_vcc, lane, val, exec_mask)
-        raw_stores.extend([('vcc', s) for s in self.wmask(_c(VCC_LO.offset), new_vcc)])
+        raw_stores.extend([('vcc', s) for s in self.wmask_lane_bit(_c(VCC_LO.offset), lane, val, exec_mask)])
       elif dest.startswith('D0'):
         dest_suffix = re.match(r'D0\.(\w+)', dest)
         if dest_suffix is not None:
@@ -1039,13 +1037,11 @@ def _compile_sdwa(inst: irc.VOP1_SDWA | irc.VOP2_SDWA | irc.VOP2_SDWA_SDST | irc
         result = _sdwa_write(old, result, dst_sel, dst_unused)
       stores.append(ctx.wvgpr_dyn(vdst_reg, lane, result, exec_mask))
     elif dest.startswith('VCC'):
-      old_vcc = ctx.rmask(_c(VCC_LO.offset))
-      stores.extend(ctx.wmask(_c(VCC_LO.offset), _set_lane_bit(old_vcc, lane, val, exec_mask)))
+      stores.extend(ctx.wmask_lane_bit(_c(VCC_LO.offset), lane, val, exec_mask))
   if vcc_val is not None:
     # Initialize sdst to 0 before lane loop (old value may be unrelated data), then set lane bits in loop
     init_stores = [ctx.wsgpr_dyn(sdst_off, _c(0)), ctx.wsgpr_dyn(sdst_off + _c(1), _c(0))]
-    old_sdst = ctx.rmask(sdst_off)
-    stores.extend(ctx.wmask(sdst_off, _set_lane_bit(old_sdst, lane, vcc_val, exec_mask)))
+    stores.extend(ctx.wmask_lane_bit(sdst_off, lane, vcc_val, exec_mask))
     if stores:
       return UOp.sink(*init_stores, UOp.sink(*stores).end(lane), *ctx.inc_pc())
     return UOp.sink(*init_stores, *ctx.inc_pc())

test/models/test_whisper.py

@@ -74,6 +74,7 @@ class TestWhisper(unittest.TestCase):
         err
       )
 
+  @slow
   def test_transcribe_file1(self):
     self.assertEqual(transcribe_file(self.model, self.enc, TEST_FILE_1),  TRANSCRIPTION_1)
 
@@ -89,6 +90,7 @@ class TestWhisper(unittest.TestCase):
     self.assertEqual(TRANSCRIPTION_1,  transcriptions[0])
     self.assertEqual(TRANSCRIPTION_2,  transcriptions[1])
 
+  @slow
   def test_transcribe_batch21(self):
     waveforms = [load_file_waveform(TEST_FILE_2), load_file_waveform(TEST_FILE_1)]
     transcriptions = transcribe_waveform(self.model, self.enc, waveforms)

test/null/test_const_folding.py

@@ -1,39 +1,10 @@
 import unittest, itertools, math
-from tinygrad import Tensor, dtypes, Context
-from tinygrad.dtype import DType, ConstType, truncate
+from tinygrad import dtypes, Context
+from tinygrad.dtype import DType, ConstType
 from tinygrad.uop.ops import Ops, UOp
 from test.helpers import full_rewrite
 import numpy as np
 
-def _check_ast_count(desired_count:int, t:Tensor):
-  # NOTE: this has side effect because everything can be scheduled only once
-  linear = t.schedule_linear()
-  asts = [s for s in linear.src if s.src[0].op is Ops.SINK]
-  len(asts)
-  # NOT SUPPORTED ANYMORE
-  #assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"
-
-class TestUnaryOpsConstFolding(unittest.TestCase):
-  def test_all_consts_ops(self):
-    _check_ast_count(0, Tensor.ones(4).exp())
-    _check_ast_count(0, Tensor.ones(4).sqrt())
-    _check_ast_count(0, Tensor.ones(4) + Tensor.ones(4))
-    _check_ast_count(0, Tensor.ones(4) / Tensor.ones(4))
-
-  def test_cast(self):
-    _check_ast_count(0, Tensor.ones(4).cast(dtypes.int16))
-    _check_ast_count(0, Tensor.full(4, fill_value=-1).cast(dtypes.uint16))
-
-  def test_neg_folding(self):
-    _check_ast_count(0, Tensor([1, 2, 3]).mul(-1).neg())
-    _check_ast_count(0, Tensor([1, 2, 3]).neg().mul(-1))
-    _check_ast_count(0, Tensor([1, 2, 3]).neg().neg())
-
-  def test_neg_realized_no_fold(self):
-    x = Tensor.randn(32, 32)
-    x = x.clip(0, 1).realize()
-    _check_ast_count(1, x.neg())
-
 class TestWeakConstFolding(unittest.TestCase):
   def test_weakint_math(self):
     out = (UOp.const(2**40) + UOp.const(2**40)).simplify()
@@ -51,87 +22,18 @@ class TestWeakConstFolding(unittest.TestCase):
   def test_invalid_poison(self):
     self.assertTrue(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().is_invalid)
 
-  def test_cast_commits_to_dtype_grid(self):
-    # committing a weak const to a stated width puts the value on that width's grid, same as storage packing and native compilers
-    v = 1/123008  # not representable in float16
-    out = UOp.const(v).cast(dtypes.half).simplify()
-    self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.half, truncate[dtypes.half](v)))
-    self.assertNotEqual(out.val, v)
-    # the grid commit preserves the sign of zero
-    self.assertEqual(math.copysign(1, UOp.const(-0.0).cast(dtypes.half).simplify().val), -1)
-    # observable at tensor level: the const-folded comparison agrees with the committed value
-    self.assertTrue((Tensor(-3.2).cast(dtypes.float32) <= truncate[dtypes.float32](-3.2)).item())
-
-class TestBinaryOpsConstFolding(unittest.TestCase):
-  def test_add_literal_zero(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + 0)
-  def test_add_tensor_zero(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(4))
-  def test_literal_zero_add(self):
-    _check_ast_count(0, 0 + Tensor([1.0, 2, 3, 4]))
-  def test_tensor_zero_add(self):
-    _check_ast_count(0, Tensor.zeros(4) + Tensor([1.0, 2, 3, 4]))
-
-  def test_sub_literal_zero(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) - 0)
-  def test_sub_tensor_zero(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) - Tensor.zeros(4))
-
-  def test_mul_literal_zero(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * 0)
-  def test_mul_tensor_zero(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.zeros(4))
-  def test_literal_zero_mul(self):
-    _check_ast_count(0, 0 * Tensor([1.0, 2, 3, 4]) * 0)
-  def test_tensor_zero_mul(self):
-    _check_ast_count(0, Tensor.zeros(4) * Tensor([1.0, 2, 3, 4]))
-
-  def test_mul_literal_one(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * 1)
-  def test_mul_tensor_one(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) * Tensor.ones(4))
-  def test_literal_one_mul(self):
-    _check_ast_count(0, 1 * Tensor([1.0, 2, 3, 4]))
-  def test_tensor_one_mul(self):
-    _check_ast_count(0, Tensor.ones(4) * Tensor([1.0, 2, 3, 4]))
-
-  def test_bool_tensor_mul_bool(self):
-    _check_ast_count(0, Tensor([True, False]) * True)
-    _check_ast_count(0, Tensor([True, False]) * False)
-  def test_bool_mul_bool_tensor(self):
-    _check_ast_count(0, True * Tensor([True, False]))
-    _check_ast_count(0, False * Tensor([True, False]))
-
-  def test_div_literal_one(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) / 1)
-  def test_div_tensor_one(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) / Tensor.ones(4))
-
-  def test_floordiv_literal_one(self):
-    _check_ast_count(0, Tensor([1, 2, 3, 4]) // 1)
-  def test_floordiv_tensor_one(self):
-    _check_ast_count(0, Tensor([1, 2, 3, 4]) // Tensor.ones(4, dtype=dtypes.int32))
-
-  def test_pow_literal_zero(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** 0)
-  def test_pow_tensor_zero(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** Tensor.zeros(4))
-
-  def test_pow_literal_one(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** 1)
-  def test_pow_tensor_one(self):
-    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** Tensor.ones(4))
-  def test_literal_one_pow(self):
-    _check_ast_count(0, 1 ** Tensor([1.0, 2, 3, 4]))
-  def test_tensor_one_pow(self):
-    _check_ast_count(0, Tensor.ones(4) ** Tensor([1.0, 2, 3, 4]))
-
 class TestBitcastConstFolding(unittest.TestCase):
+  def test_out_of_range_source_value(self):
+    for val, src_dt, dst_dt, bits in ((3000000000, dtypes.int32, dtypes.uint32, 3000000000),
+                                      (70000, dtypes.int16, dtypes.uint16, 4464),
+                                      (-5, dtypes.uint32, dtypes.int32, -5)):
+      self.assertEqual(UOp.const(val, src_dt).bitcast(dst_dt).simplify().val, bits)
+
   def test_scalar_bitcast(self):
     def t(cases: dict[DType, ConstType]):
       for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()):
         if not math.isnan(from_v):
-          r = full_rewrite(UOp.const(from_v, from_dt).bitcast(to_dt).sink()).src[0]
+          r = UOp.const(from_v, from_dt).bitcast(to_dt).simplify()
           self.assertEqual(r.op, Ops.CONST, msg:=f"{from_dt} -> {to_dt} ({from_v} -> {to_v})")
           self.assertEqual(r.dtype, to_dt, msg)
           np.testing.assert_equal(r.val, to_v, msg)
@@ -155,24 +57,9 @@ class TestBitcastConstFolding(unittest.TestCase):
 
   def test_vec_bitcast(self):
     with Context(SPEC=0):
-      srcs = full_rewrite(UOp.const((-1, -2**31, 75), dtypes.int32).bitcast(dtypes.uint32).sink()).src
-    self.assertTrue(all(r.op is Ops.CONST and r.dtype == dtypes.uint32 for r in srcs))
-    self.assertEqual(tuple(x.val for x in srcs), (2**32-1, 2**31, 75))
-
-# folds advance indexing into basic indexing
-class TestIndexingConstFolding(unittest.TestCase):
-  def test_scalar_index(self):
-    t = Tensor.arange(16).float().reshape(1,1,4,4).clone().realize()
-    _check_ast_count(1, t[:,:,Tensor(1),:])
-    _check_ast_count(1, t[:,:,Tensor(1)+2,:])
-    _check_ast_count(1, t[:,:,Tensor(1),Tensor(0)])
-
-  def test_const_tensor_index(self):
-    # TODO: these can be 0, implement const tensor folded indexing
-    t = Tensor.arange(16).float().reshape(1,1,4,4).clone().realize()
-    _check_ast_count(1, t[:,:,Tensor.ones(2,1,dtype=dtypes.int),:])
-    _check_ast_count(1, t[:,:,Tensor.ones(1,2,dtype=dtypes.int)+2,:])
-    _check_ast_count(1, t[:,:,Tensor.ones(1,1,dtype=dtypes.int),Tensor.zeros(2,1,2,dtype=dtypes.int)])
+      result = full_rewrite(UOp.const((-1, -2**31, 75), dtypes.int32).bitcast(dtypes.uint32).sink())
+      expected = full_rewrite(UOp.const((2**32-1, 2**31, 75), dtypes.uint32).sink())
+    self.assertEqual(result.src, expected.src)
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_dtype_spec.py

@@ -51,10 +51,6 @@ class TestHelpers(unittest.TestCase):
     assert dtypes.is_float(dtypes.fp8e4m3)
     assert dtypes.is_float(dtypes.fp8e5m2)
 
-  @given(strat.sampled_from([d for d in DTYPES_DICT.values() if dtypes.is_float(d) or dtypes.is_int(d)]))
-  def test_scalar(self, dtype):
-    assert dtype.scalar() == dtype
-
   def test_from_py(self):
     assert dtypes.from_py(True) == dtypes.bool
     assert dtypes.from_py(Invalid) == dtypes.bool

test/null/test_elf.py

@@ -1,6 +1,7 @@
 import unittest, subprocess, platform
 from tinygrad.runtime.support.compiler_cpu import ClangCompiler
 from tinygrad.runtime.support.elf import elf_loader
+from tinygrad.runtime.support.c import DLL
 
 class TestElfLoader(unittest.TestCase):
   def test_load_clang_jit_strtab(self):
@@ -23,7 +24,7 @@ class TestElfLoader(unittest.TestCase):
       }
     '''
     with self.assertRaisesRegex(RuntimeError, 'evil_external_function'):
-      ClangCompiler([{'AMD64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine(), m), "native"]).compile(src)
+      elf_loader(ClangCompiler([{'AMD64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine(), m), "native"]).compile(src))
   def test_link(self):
     src = '''
       float powf(float, float); // from libm
@@ -32,7 +33,7 @@ class TestElfLoader(unittest.TestCase):
     args = ('-x', 'c', '-c', '-target', f'{platform.machine()}-none-unknown-elf', '-march=native', '-fPIC', '-O2', '-ffreestanding', '-nostdlib')
     obj = subprocess.check_output(('clang',) + args + ('-', '-o', '-'), input=src.encode())
     with self.assertRaisesRegex(RuntimeError, 'powf'): elf_loader(obj)
-    elf_loader(obj, link_libs=['m'])
+    elf_loader(obj, link_libs=[DLL('m', 'm')])
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_graph_rewrite.py

@@ -1,8 +1,7 @@
 import unittest, math
 from tinygrad import dtypes
-from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import all_same, Context
-from tinygrad.uop.ops import GroupOp, UOp, Ops, exec_alu, PatternMatcher, TrackedPatternMatcher, UPat
+from tinygrad.uop.ops import GroupOp, UOp, Ops, PatternMatcher, TrackedPatternMatcher, UPat
 from test.helpers import full_rewrite
 from hypothesis import given, strategies as strat
 
@@ -11,125 +10,14 @@ from hypothesis import given, strategies as strat
 def apply_rewrite(expr):
   return full_rewrite(expr.sink()).src[0]
 
-@Context(SPEC=0)
-def apply_rewrite_values(expr):
-  srcs = full_rewrite(expr.sink()).src
-  if len(srcs) == 1:
-    if srcs[0].op is Ops.CONST: return (srcs[0].val,)
-    if srcs[0].op is Ops.STACK: return tuple(s.val for s in srcs[0].src)
-  return tuple(s.val for s in srcs)
-
-def evaluate_uop(uop, variables):
-  if uop.op == Ops.CONST:
-    return uop.val
-  elif uop.op == Ops.PARAM and uop.arg.addrspace is AddrSpace.ALU:
-    return variables[uop.expr]
-  elif uop.op in GroupOp.ALU:
-    src_values = [evaluate_uop(src, variables) for src in uop.src]
-    return exec_alu(uop.op, uop.dtype, src_values)
-  else:
-    raise NotImplementedError(f"Unsupported UOp {uop.op}")
-
-class TestArithmeticSimplifications(unittest.TestCase):
-  def test_full_graph_rewrite_division_by_zero(self):
-    optimized_div_uop = apply_rewrite(UOp.const(10.0) / UOp.const(0.0))
-    self.assertEqual(optimized_div_uop.op, Ops.CONST)
-    self.assertTrue(math.isinf(optimized_div_uop.val) or math.isnan(optimized_div_uop.val))
-
-  def test_full_graph_rewrite_redundant_operations(self):
-    optimized_uop = apply_rewrite((UOp.const(10.0) + UOp.const(0.0)) * UOp.const(1.0))
-    self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.val, 10.0)
-
-  def test_full_graph_rewrite_large_graph(self):
-    prev_uop = UOp.const(0)
-    for i in range(1, 101):
-      prev_uop += UOp.const(i)
-    optimized_uop = apply_rewrite(prev_uop)
-    self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.val, sum(range(1, 101)))
-
-  def test_full_graph_rewrite_division_by_one(self):
-    optimized_uop = apply_rewrite(UOp.const(42.0) / UOp.const(1.0))
-    self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.val, 42.0)
-
-  def test_full_graph_rewrite_modulo_by_one(self):
-    optimized_uop = apply_rewrite(UOp.const(42) % UOp.const(1))
-    self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.val, 0)
-
-
-class TestFoldingAndReduction(unittest.TestCase):
-  @unittest.skip("reduce is removed now")
-  def test_full_graph_rewrite_constant_reduction_folding(self):
-    const1 = UOp.const(5)
-    const2 = UOp.const(10)
-    const3 = UOp.const(20)
-    optimized_sink = apply_rewrite((const1 + const2 + const3).reduce(Ops.ADD))
-    expected_sum = 5 + 10 + 20
-    self.assertEqual(optimized_sink.val, expected_sum)
-
-  @unittest.skip("reduce is removed now")
-  def test_full_graph_rewrite_reduction_with_unused_range(self):
-    const1 = UOp.const(15)
-    const2 = UOp.const(25)
-    rng = UOp.range(10, idx=0)
-    optimized_sink = apply_rewrite((const1 + const2).reduce(Ops.ADD, rng))
-    expected_sum = 10 * (15 + 25)
-    self.assertEqual(optimized_sink.val, expected_sum)
-
-  @unittest.skip("currently failing")
-  def test_full_graph_rewrite_range_reduction(self):
-    simple_range = UOp.range(5, idx=0)
-    optimized_sink = apply_rewrite(simple_range.reduce(Ops.ADD, simple_range))
-    expected_sum = sum(range(5))
-    self.assertEqual(optimized_sink.val, expected_sum)
-
-  @unittest.skip("currently failing")
-  def test_full_graph_rewrite_simple_reduction_folding(self):
-    simple_range = UOp.range(4, idx=0)
-    add_uop = simple_range + UOp.const(1)
-    optimized_sink = apply_rewrite(add_uop.reduce(Ops.ADD, simple_range))
-    expected_sum = sum(i + 1 for i in range(4))
-    self.assertEqual(optimized_sink.val, expected_sum)
-
-  @unittest.skip("currently failing")
-  def test_full_graph_rewrite_nested_loop_collapse(self):
-    outer_range = UOp.range(8, 0)
-    inner_range = UOp.range(4, 1)
-    expr = (outer_range * 10) + inner_range
-    optimized_reduce_uop = apply_rewrite(expr.reduce(Ops.ADD, outer_range, inner_range))
-    self.assertEqual(optimized_reduce_uop.op, Ops.CONST)
-    self.assertEqual(optimized_reduce_uop.val, sum((i * 10) + j for i in range(8) for j in range(4)))
-
+def const_value(uop:UOp):
+  if uop.op is Ops.CAST: uop = uop.src[0]
+  assert uop.op is Ops.CONST
+  return uop.val
 
 class TestModuloAndDivisionFolding(unittest.TestCase):
-  def test_full_graph_rewrite_modulo_folding_with_define_var(self):
-    # index dtype because div-mod rules only work on index
-    x_var_uop = UOp.variable('x', 0, 100).cast(dtypes.weakint)
-    optimized_mod_uop = apply_rewrite(((x_var_uop * 4) + 2) % 4)
-    self.assertEqual(optimized_mod_uop.op, Ops.CONST)
-    self.assertEqual(optimized_mod_uop.val, 2)
-
-  def test_full_graph_rewrite_division_folding_with_define_var(self):
-    # index dtype because div-mod rules only work on index
-    n_var_uop = UOp.variable('n', 1, 1000).cast(dtypes.weakint)
-    optimized_div_uop = apply_rewrite((n_var_uop * 6) // 3)
-    self.assertEqual(optimized_div_uop.op, Ops.MUL)
-    self.assertEqual(optimized_div_uop.src[1].val, 2)
-
-  def test_full_graph_rewrite_complex_mod_div_folding(self):
-    # index dtype because div-mod rules only work on index
-    k_var_uop = UOp.variable('k', 0, 50).cast(dtypes.weakint)
-    optimized_div_uop = apply_rewrite(((k_var_uop * 12 + 8) % 6) // 2)
-    self.assertEqual(optimized_div_uop.op, Ops.CONST)
-    self.assertEqual(optimized_div_uop.val, 1)
-
   def test_graph_rewrite_div_folding_bug(self):
-    lhs = UOp(Ops.ADD, src=(
-      UOp(Ops.STACK, arg=None, src=(UOp(Ops.SPECIAL, src=(UOp.const(32),), arg='lidx0'),)*4),
-      UOp.const((0, 256, 512, 768))))
+    lhs = UOp.stack(*(UOp.special(32, 'lidx0'),)*4) + UOp.const((0, 256, 512, 768))
     rhs = UOp.const((2,)*4)
     unopt = lhs<rhs
     opt = apply_rewrite(unopt)
@@ -137,74 +25,31 @@ class TestModuloAndDivisionFolding(unittest.TestCase):
     print(opt)
     if opt.op is Ops.STACK: self.assertFalse(all_same(opt.src))
 
-  def test_full_graph_rewrite_modulo_large_divisor(self):
-    # index dtype because div-mod rules only work on index
-    x_var_uop = UOp.variable('x', 1, 5)
-    self.assertIs(apply_rewrite(x_var_uop.cast(dtypes.weakint) % 10).render(simplify=False), x_var_uop.render(simplify=False))
-
-  def test_full_graph_rewrite_division_with_remainder(self):
-    x_var_uop = UOp.variable('x', 7, 9)
-    optimized_sink = apply_rewrite(x_var_uop // 2)
-    for x_value in range(7, 10):
-      self.assertEqual(x_value // 2, evaluate_uop(optimized_sink, {'x': x_value}))
-
-  def test_full_graph_rewrite_complex_mod_div_expression(self):
-    x_var_uop = UOp.variable('x', 1, 10)
-    optimized_sink = apply_rewrite(((x_var_uop * 5) % 3) // 2)
-    for x_value in range(1, 11):
-      original_result = ((x_value * 5) % 3) // 2
-      optimized_result = evaluate_uop(optimized_sink, {'x': x_value})
-      self.assertEqual(original_result, optimized_result)
-
-
 class TestEdgeCasesAndSpecialOperations(unittest.TestCase):
   def test_full_graph_rewrite_transcendental_edge_cases(self):
     optimized_sink = full_rewrite(UOp.const(-1.0).log2().sink(UOp.const(0.0).reciprocal()))
     optimized_log2_neg, optimized_recip_zero = optimized_sink.src
-    self.assertTrue(math.isnan(optimized_log2_neg.val), f"Expected NaN for log2(-1.0), got {optimized_log2_neg.val}")
-    self.assertTrue(math.isinf(optimized_recip_zero.val) and optimized_recip_zero.val > 0,
-                    f"Expected +inf for reciprocal(0.0), got {optimized_recip_zero.val}")
-
-  @unittest.skip("broken")
-  def test_full_graph_rewrite_modulo_negative_dividend(self):
-    x_var_uop = UOp.variable('x', -5, -1)
-    optimized_sink = full_rewrite((x_var_uop % 3).sink())
-    for x_value in range(-5, 0):
-      self.assertEqual(x_value % 3, evaluate_uop(optimized_sink.src[0], {'x': x_value}))
-
-  @unittest.skip("broken")
-  def test_full_graph_rewrite_division_negative_divisor(self):
-    x_var_uop = UOp.variable('x', 1, 5)
-    optimized_sink = full_rewrite((x_var_uop // -2).sink())
-    for x_value in range(1, 6):
-      self.assertEqual(x_value // -2, evaluate_uop(optimized_sink.src[0], {'x': x_value}))
+    log2_neg, recip_zero = const_value(optimized_log2_neg), const_value(optimized_recip_zero)
+    self.assertTrue(math.isnan(log2_neg), f"Expected NaN for log2(-1.0), got {log2_neg}")
+    self.assertTrue(math.isinf(recip_zero) and recip_zero > 0, f"Expected +inf for reciprocal(0.0), got {recip_zero}")
 
 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((1.0, 2.0, 3.0, 4.0))
-    self.assertEqual(apply_rewrite(base_vector.index(2)).val, 3.0)
+    self.assertIs(apply_rewrite(base_vector.index(2)), apply_rewrite(base_vector.src[2]))
 
   def test_gep_tuple_extraction(self):
     # GEP on a vector dtype to extract multiple elements as a vector
     base_vector = UOp.const((1.0, 2.0, 3.0, 4.0))
-    self.assertEqual(list(apply_rewrite_values(UOp.stack(*[base_vector.index(i) for i in (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((1.0, 2.0, 3.0, 4.0))
-    self.assertEqual(apply_rewrite(const_stack.index(2)).val, 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((7.0, 8.0, 9.0, 10.0))
-    self.assertEqual(list(apply_rewrite_values(UOp.stack(*[const_stack.index(i) for i in (1, 3)]))), [8.0, 10.0])
+    self.assertIs(apply_rewrite(UOp.stack(*[base_vector.index(i) for i in (2, 3)])),
+                  apply_rewrite(UOp.stack(base_vector.src[2], base_vector.src[3])))
 
   def test_vectorize_multiple_elements(self):
     # Vectorizing multiple elements using GEP
     base_vector = UOp.const((5.0, 10.0, 15.0, 20.0))
-    vectorized_uop = UOp(Ops.STACK, src=tuple(base_vector.index(i) for i in range(4)))
-    self.assertEqual(list(apply_rewrite_values(vectorized_uop)), [5.0, 10.0, 15.0, 20.0])
+    vectorized_uop = UOp.stack(*(base_vector.index(i) for i in range(4)))
+    self.assertIs(apply_rewrite(vectorized_uop), apply_rewrite(base_vector))
 
 
 import inspect
@@ -256,16 +101,6 @@ class TestSubstitute(unittest.TestCase):
     ret = substitute(ret, {a.sin():b})
     self.assertIs(ret, b.sin())
 
-  # broken due to infinite recursion
-  # NOTE: VIZ hangs and doesn't recover if you click this one
-  @unittest.skip("recursion error no longer raised")
-  def test_assert_inf_recurse(self):
-    a = UOp.variable('a', 0, 10)
-    n1 = a.sin()
-    ret = n1
-    with self.assertRaises(RecursionError):
-      ret = substitute(ret, {n1:n1.sqrt()})
-
   def test_sin_to_sqrt(self):
     a = UOp.variable('a', 0, 10, dtype=dtypes.float)
     n1 = a.sin()

test/null/test_kimi.py

@@ -1,30 +0,0 @@
-import unittest
-from tinygrad import Tensor, dtypes, nn
-from tinygrad.llm.kimi import _shard_kimi
-from tinygrad.llm.model import SSMConfig, Transformer, TransformerConfig
-
-class TestKimiTP4(unittest.TestCase):
-  def test_prefill_and_decode_graph(self):
-    devices = ("NULL:0", "NULL:1", "NULL:2", "NULL:3")
-    config = TransformerConfig(num_blocks=2, dim=32, hidden_dim=128, n_heads=4, n_kv_heads=1, norm_eps=1e-5,
-      vocab_size=64, head_dim=12, rope_theta=10000, rope_dim=4, v_head_dim=8, max_context=4, kv_lora_rank=16,
-      num_experts=8, num_experts_per_tok=2, norm_topk_prob=True, shared_expert_dim=32, ssm_layers=(True, False),
-      ssm=SSMConfig(4, 8, 4, 4, 32, True), shared_expert_gate=False, leading_dense_blocks=1, dense_hidden_dim=64,
-      routed_scaling_factor=2.446, expert_bias=True, expert_mxfp4=True, bf16_activations=True, kda_split_qkv=True)
-    model = Transformer(config)
-    for name, value in nn.state.get_state_dict(model).items():
-      fill = 127 if name.endswith("weight_scale") else 0
-      value.replace(Tensor.full(value.shape, fill, dtype=value.dtype if value.dtype is dtypes.uint8 else dtypes.bfloat16, device="NULL"))
-    _shard_kimi(model, devices)
-
-    temperature = Tensor([0.0], device=devices)
-    prefill = model(Tensor([[1, 2]], dtype=dtypes.int32, device=devices), 0, temperature).realize()
-    model(Tensor([[1, 2]], dtype=dtypes.int32, device=devices), 0, temperature).realize()  # replay prefill JIT
-    decode = model(Tensor([[3]], dtype=dtypes.int32, device=devices), 2, temperature).realize()
-    model(Tensor([[4]], dtype=dtypes.int32, device=devices), 3, temperature).realize()  # replay decode JIT
-    self.assertEqual(prefill.shape, (1, 1))
-    self.assertEqual(decode.shape, (1, 1))
-    self.assertEqual(model.blk[0].recurrent_state.uop.axis, 1)
-    self.assertEqual(model.blk[1].cache_k.dtype, dtypes.bfloat16)
-
-if __name__ == "__main__": unittest.main()

test/null/test_kimi_k3.py

@@ -1,29 +0,0 @@
-import unittest
-from tinygrad import Tensor, dtypes, nn
-from tinygrad.llm.kimi_k3 import _shard_kimi_k3
-from test.unit.test_llm_k3 import small_k3_config
-from tinygrad.llm.model import Transformer
-
-class TestKimiK3TP8(unittest.TestCase):
-  @staticmethod
-  def _model():
-    model = Transformer(small_k3_config())
-    for name,value in nn.state.get_state_dict(model).items():
-      fill = 127 if name.endswith("weight_scale") else 0
-      dtype = value.dtype if value.dtype is dtypes.uint8 else dtypes.bfloat16
-      value.replace(Tensor.full(value.shape, fill, dtype=dtype, device="NULL"))
-    _shard_kimi_k3(model, tuple(f"NULL:{i}" for i in range(8)))
-    return model
-
-  def test_prefill_decode_and_jit_replay(self):
-    devices = tuple(f"NULL:{i}" for i in range(8))
-    model = self._model()
-    temperature = Tensor([0.0], device=devices)
-    self.assertEqual(model(Tensor([[1, 2]], dtype=dtypes.int32, device=devices), 0, temperature).realize().shape, (1, 1))
-    model(Tensor([[1, 2]], dtype=dtypes.int32, device=devices), 0, temperature).realize()
-    self.assertEqual(model(Tensor([[3]], dtype=dtypes.int32, device=devices), 2, temperature).realize().shape, (1, 1))
-    model(Tensor([[4]], dtype=dtypes.int32, device=devices), 3, temperature).realize()
-    self.assertEqual(model.blk[0].recurrent_state.uop.axis, 1)
-    self.assertEqual(model.blk[1].cache_k.dtype, dtypes.bfloat16)
-
-if __name__ == "__main__": unittest.main()

test/null/test_memory_planner.py

@@ -27,10 +27,15 @@ def _make_linear(buffer_lists, copies=None):
     calls.append(UOp(Ops.CALL, src=(src0, *bufs)))
   return UOp(Ops.LINEAR, src=tuple(calls))
 
+def _get_planned_view(buf:UOp) -> tuple[UOp, int, int]|None:
+  view = buf.src[0] if buf.op is Ops.BITCAST else buf
+  if view.op is not Ops.SHRINK or view.src[0].op is not Ops.BUFFER: return None
+  return (arena:=view.src[0]), view.src[1].val * arena.dtype.itemsize, view.src[2].val * arena.dtype.itemsize
+
 def _get_arena(buf, linear, result):
   for orig_si, new_si in zip(linear.src, result.src):
     for orig, new in zip(orig_si.src[1:], new_si.src[1:]):
-      if orig is buf and new.op is Ops.SLICE: return new.src[0]
+      if orig is buf and (planned:=_get_planned_view(new)) is not None: return planned[0]
   return None
 
 def check_assign(buffer_lists, copies=None):
@@ -41,8 +46,8 @@ def check_assign(buffer_lists, copies=None):
   replace_map: dict[int, tuple[UOp, int, int]] = {}
   for orig_si, new_si in zip(linear.src, result.src):
     for orig, new in zip(orig_si.src[1:], new_si.src[1:]):
-      if new.op is Ops.SLICE and id(orig) not in replace_map:
-        replace_map[id(orig)] = (new.src[0], new.src[1].val * new.src[0].dtype.itemsize, new.arg * new.dtype.itemsize)
+      if (planned:=_get_planned_view(new)) is not None and id(orig) not in replace_map:
+        replace_map[id(orig)] = planned
 
   # verify pinned buffers are not planned
   for buf in held_bufs:

test/null/test_mnist_dataset.py

@@ -1,6 +1,7 @@
 import unittest
 from tinygrad.helpers import GlobalCounters
 from tinygrad.nn.datasets import mnist
+from test.helpers import KernelCountException
 
 class TestDataset(unittest.TestCase):
   def test_dataset_is_realized(self):
@@ -8,7 +9,7 @@ class TestDataset(unittest.TestCase):
     X_train[0].contiguous().realize()
     GlobalCounters.reset()
     X_train[0].contiguous().realize()
-    self.assertLessEqual(GlobalCounters.kernel_count, 1)  # 0 if SLICE (zero-copy), 1 otherwise
+    if GlobalCounters.kernel_count > 1: raise KernelCountException(1, GlobalCounters.kernel_count)  # 0 if SLICE (zero-copy), 1 otherwise
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_simplify_valid_idx.py

@@ -15,17 +15,13 @@ def simplify_valid_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move
 def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load+indexing_simplify, name="simplify_image_idx")
 
 def get_gated_load_uop(valid:UOp, idx:UOp):
-  return UOp(Ops.LOAD, src=(
-    UOp.param(0, dtypes.float, (1024,)).index(idx.valid(valid)),
-  ))
+  return UOp.param(0, dtypes.float, (1024,)).index(idx.valid(valid)).load()
 
 def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]):
-  return UOp(Ops.LOAD, src=(
-    UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)),
-  ))
+  return UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)).load()
 
-def Special(expr, nmax): return UOp(Ops.SPECIAL, src=(UOp.const(nmax),), arg=expr)
-def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax)
+def Special(expr, nmax): return UOp.special(nmax, expr)
+def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax, param=True)
 def Range(n, nmax): return UOp.range(nmax, n)
 
 class TestValidIdxSimplification(unittest.TestCase):
@@ -512,7 +508,7 @@ class TestDropTrueGate(unittest.TestCase):
     buf = UOp.param(0, dtypes.int, (1,))
     idx = UOp.const(0)
     true_gate = UOp.const(True)
-    index_with_gate = UOp(Ops.INDEX, src=(buf, idx.valid(true_gate)))
+    index_with_gate = buf.index(idx.valid(true_gate))
     # apply the optimization
     result = graph_rewrite(index_with_gate, sym+indexing_simplify)
     # the True valid should be dropped (INDEX should only have 2 sources)
@@ -524,13 +520,17 @@ class TestRangeShrink(unittest.TestCase):
       result = full_rewrite(sink)
     return [u for u in result.toposort() if u.op is Ops.RANGE]
 
+  def assert_range_end(self, ranges:list[UOp], end:int):
+    self.assertEqual(len(ranges), 1)
+    with Context(NOOPT=1, SPEC=0): expected = full_rewrite(UOp.const(end, dtypes.int).sink()).src[0]
+    self.assertIs(ranges[0].src[0], expected)
+
   def test_range_shrink_single_guard(self):
     # range 0..203 guarded by r < 4 everywhere -> shrink to 0..3
     r = Range(0, 204)
     load = get_gated_load_uop(r < UOp.const(4), r)
     ranges = self.get_ranges(load.sink())
-    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 4)
+    self.assert_range_end(ranges, 4)
 
   def test_range_shrink_picks_max_guard(self):
     # two loads guard the same range with r < 4 and r < 8 -> shrink to max(4, 8) = 8
@@ -538,25 +538,22 @@ class TestRangeShrink(unittest.TestCase):
     load1 = get_gated_load_uop(r < UOp.const(4), r)
     load2 = get_gated_load_uop(r < UOp.const(8), r)
     ranges = self.get_ranges(UOp.sink(load1, load2))
-    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 8)
+    self.assert_range_end(ranges, 8)
 
   def test_range_no_shrink_guard_ge_max(self):
     # guard r < 300 with range max 204 -> no shrink (guard doesn't constrain)
     r = Range(0, 204)
     load = get_gated_load_uop(r < UOp.const(300), r)
     ranges = self.get_ranges(load.sink())
-    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 204)
+    self.assert_range_end(ranges, 204)
 
   def test_range_no_shrink_when_unguarded_elsewhere(self):
     # one load guards r < 4, but another load uses r without a gate -> no shrink
     r = Range(0, 204)
     load1 = get_gated_load_uop(r < UOp.const(4), r)
-    load2 = UOp(Ops.LOAD, src=(UOp.param(1, dtypes.float, (204,)).index(r),))
+    load2 = UOp.param(1, dtypes.float, (204,)).index(r).load()
     ranges = self.get_ranges(UOp.sink(load1, load2))
-    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 204)
+    self.assert_range_end(ranges, 204)
 
   def test_range_no_shrink_when_used_in_reduce(self):
     # range used in both a gated load AND directly in the reduce expression -> no shrink
@@ -564,8 +561,7 @@ class TestRangeShrink(unittest.TestCase):
     gated_load = get_gated_load_uop(r < UOp.const(4), r)
     red = (r.cast(dtypes.float) + gated_load).reduce(r, arg=Ops.ADD)
     ranges = self.get_ranges(red.sink())
-    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 204)
+    self.assert_range_end(ranges, 204)
 
   def test_range_shrink_to_single_iteration(self):
     # guard r < 1 shrinks range to 1 -> single iteration, range eliminated entirely
@@ -580,8 +576,7 @@ class TestRangeShrink(unittest.TestCase):
     r = Range(0, 204)
     x = (r < 4).where(UOp.const(1.0), Invalid)
     ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, Invalid)).sink())
-    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 4)
+    self.assert_range_end(ranges, 4)
 
   def test_range_shrink_store_where_invalid_flipped(self):
     # above, but flipped
@@ -589,8 +584,7 @@ class TestRangeShrink(unittest.TestCase):
     r = Range(0, 204)
     x = (r < 4).where(UOp.const(1.0), Invalid)
     ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r >= 4).where(Invalid, x)).sink())
-    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 4)
+    self.assert_range_end(ranges, 4)
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_tensor.py

@@ -69,7 +69,7 @@ class TestIdxUpcast(unittest.TestCase):
     if not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)):
       assert idx.op is Ops.INDEX
       idx_val = idx.src[1]
-      self.assertFalse(idx_val.overflows(idx_val.dtype.scalar()))
+      self.assertFalse(idx_val.overflows(idx_val.dtype))
 
   # use expand to generate kernel that uses large idx
   def do_op_then_assert(self, dtype: DType, dim1, dim2, dim3):
@@ -171,6 +171,11 @@ class TestTensorConstLike(unittest.TestCase):
     t = Tensor.ones(8, 4).shard(("NULL:0", "NULL:1"), axis=0)
     with self.assertRaises(RuntimeError): t.full_like(5, device="NULL")
 
+class TestTensorShape(unittest.TestCase):
+  def test_float_shape_raises(self):
+    for dim in (2.0, 2.5):
+      with self.subTest(dim=dim), self.assertRaisesRegex(RuntimeError, "shape must be int"): Tensor.ones(dim)
+
 class TestTensorDevice(unittest.TestCase):
   def test_create_from_single_device_tuple(self):
     (Tensor([1.0], device=(Device.DEFAULT,)) + Tensor([2.0])).realize()

test/null/test_uop_graph.py

@@ -1,10 +1,9 @@
 import unittest, pytest
 from tinygrad import dtypes, Variable, Device
 from tinygrad.dtype import AddrSpace
-from tinygrad.helpers import DEBUG, Context
 from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, GroupOp, AxisType, broadcast_axes, KernelInfo
 from tinygrad.uop.symbolic import sym
-from test.helpers import to_uops_list
+from test.helpers import full_rewrite, to_uops_list
 from tinygrad.codegen import full_rewrite_to_sink
 
 simple_pm = PatternMatcher([
@@ -14,43 +13,27 @@ simple_pm = PatternMatcher([
   ((UPat.var('x') + UPat.cvar('c1')) + UPat.cvar('c2'), lambda x,c1,c2: x + (c1.val+c2.val)),
 ])
 
-def const_values(u:UOp):
-  if u.op is Ops.CONST: return (u.val,)
-  if u.op is Ops.STACK: return tuple(x.val for x in u.src)
-  raise AssertionError(f"expected const-like UOp, got {u.op}")
-
 class TestGraphRewriteConst(unittest.TestCase):
   def test_gep_const(self):
     v1 = UOp.const((0,1,2), dtypes.int)
     v2 = v1.index(1)
     ret = graph_rewrite(v2, sym)
-    self.assertEqual(ret.dtype, dtypes.int)
-    self.assertEqual(ret.val, 1)
+    self.assertIs(ret, UOp.const(1, dtypes.int))
 
   def test_add_const(self):
     v1 = UOp.const((0,1,2))
     v2 = UOp.const((5,6,7))
-    ret = graph_rewrite(v1+v2, sym)
-    self.assertEqual(ret.op, Ops.STACK)
-    self.assertEqual(const_values(ret), (5,7,9))
-
-  def test_add_const_lose_v(self):
-    v1 = UOp.const((0,1,2))
-    v2 = UOp.const((2,1,0))
-    ret = graph_rewrite(v1+v2, sym)
-    self.assertEqual(ret.op, Ops.STACK)
-    self.assertEqual(const_values(ret), (2,2,2))
+    self.assertIs(graph_rewrite(v1+v2, sym), UOp.const((5,7,9)))
 
 def xfail_broken_const_wraparound(fn):
   fn = pytest.mark.xfail(reason="const folding does not properly implement modular arithmetic")(fn)
   return unittest.expectedFailure(fn)
 class TestModularWraparound(unittest.TestCase):
   def _test(self, uop:UOp, expected:int):
-    results = to_uops_list([uop])
-    self.assertEqual(len(results), 2)  # +1 for SINK
-    self.assertEqual(results[0].op, Ops.CONST)
-    self.assertEqual(results[0].dtype, uop.dtype)
-    self.assertEqual(results[0].val, expected)
+    result = uop.simplify()
+    self.assertEqual(result.op, Ops.CONST)
+    self.assertEqual(result.dtype, uop.dtype)
+    self.assertEqual(result.val, expected)
 
   @xfail_broken_const_wraparound
   def test_cast(self):
@@ -157,7 +140,7 @@ class TestGraphRewrite(unittest.TestCase):
     self.assertEqual(nout.val, 3.0)
 
   def test_depth_2_fold(self):
-    v = UOp.variable("v", 0, 1, dtypes.float)
+    v = UOp.variable("v", 0, 1, dtypes.float, param=True)
     c1 = UOp.const(1.0)
     c2 = UOp.const(2.0)
     nout = graph_rewrite(v+c1+c2, simple_pm)
@@ -191,63 +174,25 @@ class TestGraphRewrite(unittest.TestCase):
       self.assertEqual(len([x for x in sink.toposort() if x.op is Ops.CONST]), 1)
 
 class TestUOpGraph(unittest.TestCase):
-  def test_add_constant_fold(self):
-    c1 = UOp.const(1.0, dtypes.float)
-    c2 = UOp.const(2.0, dtypes.float)
-    out = c1+c2
-    uops = to_uops_list([out])
-    self.assertEqual(len(uops), 2)  # +1 for SINK
-    out = uops[-2]
-    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 3.0)
-
   def test_where_same_fold(self):
     v = UOp.variable('tmp', 0, 1)
     c0 = UOp.const(0)
     vc = v != c0
     c1 = UOp.const(1.0, dtypes.float)
     out = vc.where(c1, c1)
-    uops = to_uops_list([out])
-    self.assertEqual(len(uops), 2)  # +1 for SINK
-    out = uops[-2]
-    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 1.0)
+    self.assertIs(out.simplify(), c1)
 
   def test_where_const_fold(self):
     bf = UOp.const(False)
     c1 = UOp.const(1.0, dtypes.float)
     c2 = UOp.const(2.0, dtypes.float)
     out = bf.where(c1, c2)
-    uops = to_uops_list([out])
-    self.assertEqual(len(uops), 2)  # +1 for SINK
-    out = uops[-2]
-    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 2.0)
+    self.assertIs(out.simplify(), c2)
 
   def test_const_cast(self):
     bf = UOp.const(False)
     out = bf.cast(dtypes.int)
-    uops = to_uops_list([out])
-    self.assertEqual(len(uops), 2)  # +1 for SINK
-    out = uops[-2]
-    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 0)
-
-  def test_const_bitcast(self):
-    bf = UOp.const(1.0, dtypes.float)
-    out = bf.bitcast(dtypes.uint32)
-    uops = to_uops_list([out])
-    self.assertEqual(len(uops), 2)  # +1 for SINK
-    out = uops[-2]
-    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 0x3F800000)
-
-  @unittest.expectedFailure
-  def test_const_shape_change_bitcast(self):
-    bf = UOp.const(0x3F).cast(dtypes.uint8)
-    out = bf.bitcast(dtypes.half)
-    uops = to_uops_list([out])
-    self.assertEqual(len(uops), 2)  # +1 for SINK
+    self.assertIs(full_rewrite(out.sink()).src[0], full_rewrite(UOp.const(0, dtypes.int).sink()).src[0])
 
   def test_devectorize_derives_lane_dtype(self):
     from tinygrad.codegen import do_devectorize
@@ -257,66 +202,11 @@ class TestUOpGraph(unittest.TestCase):
     invalid_lane_mul = next(u for u in out.src[0].toposort() if u.op is Ops.MUL)
     self.assertIs(invalid_lane_mul.dtype, dtypes.bool)
 
-  @unittest.skip("this test isn't valid uops")
-  def test_noop_vectorize_fold(self):
-    d0 = UOp.param(0, dtypes.float, (1,))
-    idx = UOp.const(0)
-    ld = d0.load(idx, dtype=dtypes.float)
-    vec = UOp(Ops.STACK, dtypes.float, (ld,))
-    x = vec.index(0)
-    alu = UOp(Ops.SQRT, src=(x, ))
-    out = UOp(Ops.STORE, src=(d0, idx, alu))
-    uops = to_uops_list([out])
-    self.assertEqual(len([x for x in uops if x.op is Ops.STACK]), 0)
-
-  @unittest.skip("this test isn't valid uops")
-  def test_gep_vec_fold(self):
-    d0 = UOp.param(0, dtypes.float, (1,))
-    d1 = UOp.param(1, dtypes.float, (1,))
-    d2 = UOp.param(2, dtypes.float, (1,))
-    idx = UOp.const(0)
-    def _test_vec(geps, count=4):
-      vec = UOp(Ops.STACK, dtypes.float, geps)
-      out = d0.index(idx).store(vec)
-      uops = to_uops_list([out])
-      if DEBUG >= 4:
-        from tinygrad import Device
-        print(Device[Device.DEFAULT].renderer.render(uops))
-      return uops[-2].src[-1]  # -2 to skip SINK
-
-    # possible
-    val = d1.index(idx).load(dtype=dtypes.float)
-    xyzw = tuple(val.index(i) for i in range(4))
-    self.assertIs(_test_vec(xyzw).op, Ops.LOAD)
-
-    # unaligned
-    val = d1.index(idx).load(dtype=dtypes.float)
-    wzyx = tuple(val.index(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)
-    xy = tuple(val.index(i) for i in range(2))
-    self.assertIs(_test_vec(xy+xy).op, Ops.STACK)
-    val = d1.index(idx).load(dtype=dtypes.float)
-    xy = tuple(val.index(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)
-    val2 = d2.index(idx).load(dtype=dtypes.float)
-    xy1 = tuple(val1.index(i) for i in range(2))
-    xy2 = tuple(val2.index(i) for i in range(2))
-    self.assertIs(_test_vec(xy1+xy2).op, Ops.STACK)
-
   def test_gep_vec_const_fold(self):
     for vec_size in [2, 4, 8]:
       consts = [UOp.const(float(i), dtypes.float) for i in range(vec_size)]
-      vec = UOp(Ops.STACK, src=tuple(consts))
-      with Context(SPEC=0):
-        uops = to_uops_list([vec.index(i) for i in range(vec_size)])
-        for uop, const in zip(uops, consts):
-          self.assertEqual(uop, const)
+      vec = UOp.stack(*consts)
+      for i, const in enumerate(consts): self.assertIs(vec.index(i), const)
 
   def test_cast_alu_fold(self):
     d0 = UOp.param(0, dtypes.bool, (1,))
@@ -326,7 +216,7 @@ class TestUOpGraph(unittest.TestCase):
     alu = (ld<1).cast(dtypes.bool)
     out = d0.index(idx).store(alu)
     uops = to_uops_list([out])
-    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)
+    self.assertEqual(len([x for x in uops if x.op is Ops.CAST and x.src[0].op is not Ops.CONST]), 0)
 
   def test_double_cast_fold(self):
     d0 = UOp.param(0, dtypes.float, (1,))
@@ -336,20 +226,15 @@ class TestUOpGraph(unittest.TestCase):
     alu = ld.cast(dtypes.float).cast(dtypes.float)
     out = d0.index(idx).store(alu)
     uops = to_uops_list([out])
-    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)
+    self.assertEqual(len([x for x in uops if x.op is Ops.CAST and x.src[0].op is not Ops.CONST]), 1)
 
   def test_depth_2_const_fold(self):
-    v = UOp.variable("tmp", 0, 1, dtypes.int)
+    v = UOp.variable("tmp", 0, 1, dtypes.int, param=True)
     c2 = UOp.const(2, dtypes.int)
     c4 = UOp.const(4, dtypes.int)
     vc = v+c2
     out = vc+c4
-    uops = to_uops_list([out])
-    self.assertEqual(len(uops), 5)  # +1 for SINK, +1 for the PARAM shape STACK
-    out = uops[-2]  # -2 to skip SINK
-    self.assertEqual(out.op, Ops.ADD)
-    self.assertEqual(out.src[1].op, Ops.CONST)
-    self.assertEqual(out.src[1].val, 6)
+    self.assertIs(out.simplify(), (v+UOp.const(6, dtypes.int)).simplify())
 
   def test_bitcast_to_same_dtype_fold(self):
     for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
@@ -360,9 +245,8 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_sub_with_cast_folds(self):
     a = Variable("a", 0, 5)
-    uops = to_uops_list([a.cast(dtypes.int)+(-a).cast(dtypes.int)])
-    assert uops[0] == UOp.const(0, dtypes.int)
-    assert uops[-1].op == Ops.SINK
+    out = a.cast(dtypes.int)+(-a).cast(dtypes.int)
+    self.assertIs(full_rewrite(out.sink()).src[0], full_rewrite(UOp.const(0, dtypes.int).sink()).src[0])
 
   def test_where_on_gated_load_fold(self):
     ridx0 = UOp.range(100, 0)
@@ -371,9 +255,10 @@ class TestUOpGraph(unittest.TestCase):
     w = (ridx0<50).where(ld, 5)
     out = UOp.param(1, dtypes.long, (100,))
     uops = to_uops_list([out.index(ridx0).store(w)])
+    expected = full_rewrite(UOp.const(5, dtypes.long).sink()).src[0]
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].val==5
+      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: self.assertIs(u.src[1], expected)
 
   def test_where_on_gated_load_folds_swapped_branches(self):
     ridx0 = UOp.range(100, 0)
@@ -381,9 +266,10 @@ class TestUOpGraph(unittest.TestCase):
     ld = d0.index(ridx0.valid((ridx0<50).logical_not()))
     w = (ridx0<50).where(5, ld)
     uops = to_uops_list([w])
+    expected = full_rewrite(UOp.const(5, dtypes.long).sink()).src[0]
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD: assert u.src[1].val==5
+      if u.op is Ops.LOAD: self.assertIs(u.src[1], expected)
 
   def test_where_on_gated_load_with_cast(self):
     ridx0 = UOp.range(100, 0)
@@ -393,9 +279,10 @@ class TestUOpGraph(unittest.TestCase):
     w = (ridx0<50).where(ld, 5.0)
     out = UOp.param(1, dtypes.float, (100,))
     uops = to_uops_list([out.index(ridx0).store(w)])
+    expected = full_rewrite(UOp.const(5, dtypes.int).sink()).src[0]
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].val == 5
+      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: self.assertIs(u.src[1], expected)
 
   def test_where_on_casted_gated_load_extra_cond(self):
     ridx0 = UOp.range(100, 0)
@@ -425,9 +312,10 @@ class TestUOpGraph(unittest.TestCase):
     val = (ridx0<50).where(5, ld)
     st = idx.store(val).end(ridx0)
     uops = to_uops_list([st])
+    expected = full_rewrite(UOp.const(5, dtypes.long).sink()).src[0]
     for u in uops:
       assert u.op is not Ops.WHERE
-      if u.op is Ops.STORE: assert u.src[1].val==5
+      if u.op is Ops.STORE: self.assertIs(u.src[1], expected)
 
   def test_load_idx_becomes_int(self):
     # mnist indexing with split reduceop
@@ -501,13 +389,6 @@ class TestUOpGraph(unittest.TestCase):
     # only the second store happens
     self.assertEqual(len([u for u in uops if u.op is Ops.STORE]), 1)
 
-  @unittest.skip("this is a uop type error")
-  def test_asserts_bad_gate(self):
-    glbl0 = UOp.param(0, dtypes.int, (1,))
-    idx = UOp.const(0)
-    bad_gate = UOp.const(1)
-    with self.assertRaises(AssertionError): to_uops_list([UOp(Ops.STORE, src=(glbl0, idx, UOp.const(42), bad_gate))])
-
   def test_after_end(self):
     r = UOp.range(10, 0)
 
@@ -575,7 +456,7 @@ class TestConstBufferize(unittest.TestCase):
     from tinygrad.schedule.rangeify import pm_const_buffer_folding, BufferizeOpts
     c = UOp.const(42.0)
     r1 = UOp.range(3, 0)
-    bufferize_with_range = UOp(Ops.STAGE, src=(c, r1), arg=BufferizeOpts(device="CPU"))
+    bufferize_with_range = c.bufferize(r1, arg=BufferizeOpts(device="CPU"))
     self.assertEqual(len(bufferize_with_range.src), 2)  # const + 1 range
 
     result = graph_rewrite(bufferize_with_range, pm_const_buffer_folding, name='test')
@@ -590,7 +471,7 @@ class TestConstBufferize(unittest.TestCase):
     c = UOp.const(3.14)
     r1 = UOp.range(3, 0)
     r2 = UOp.range(4, 1)
-    bufferize_with_ranges = UOp(Ops.STAGE, src=(c, r1, r2), arg=BufferizeOpts(device="CPU"))
+    bufferize_with_ranges = c.bufferize(r1, r2, arg=BufferizeOpts(device="CPU"))
     self.assertEqual(len(bufferize_with_ranges.src), 3)  # const + 2 ranges
 
     result = graph_rewrite(bufferize_with_ranges, pm_const_buffer_folding, name='test')

test/null/test_uop_symbolic.py

@@ -3,10 +3,10 @@ import unittest, pickle, functools, math
 import z3
 
 from tinygrad.dtype import dtypes, ConstType, DType, Invalid
-from test.helpers import get_uops
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad.uop.spec import spec_shared, type_verify
-from tinygrad.uop.symbolic import sym, pm_fold_cast_const, commutative, pm_simplify_valid, pm_move_where_on_load
+from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load
+from tinygrad.uop.weak import pm_cast_weak
 from tinygrad.uop.validate import uops_to_z3
 
 def check_uop_against_string(self, v:UOp, s:str):
@@ -16,7 +16,8 @@ def check_uop_against_string(self, v:UOp, s:str):
   s_eval = graph_rewrite(s_eval, commutative, name="cannonicalize eval")
   self.assertIs(s_eval, v, f"eval did not match simplified: {s_eval} != {v.render()} for {s}")
 
-def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.weakint): return UOp.variable(name,min_val,max_val,dtype)
+def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.weakint):
+  return UOp.variable(name, min_val, max_val, dtype, param=True)
 def uconst(val): return UOp.const(val)
 def usum(ops): return functools.reduce(lambda x,y: x+y, ops)
 def uand(ops): return functools.reduce(lambda x,y: x*y, ops)
@@ -35,7 +36,7 @@ class TestSymbolic(unittest.TestCase):
     self.assertEqual(solver.check(expr1 != expr2), z3.unsat, "simplified expression not equal to original")
 
   def helper_test_variable(self, v, n, m, s, test_z3:bool=True):
-    v_simplified = graph_rewrite(v, sym+pm_fold_cast_const, name="simplify symbolic uop")
+    v_simplified = graph_rewrite(v, sym+pm_cast_weak, name="simplify symbolic uop")
     if test_z3: self.check_equal_z3(v, v_simplified)
     nmin, nmax = v_simplified.vmin, v_simplified.vmax
     check_uop_against_string(self, v_simplified, s)
@@ -442,7 +443,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(uand([uconst(1), Variable("a", 0, 1)]), 0, 1, "a")
 
   def test_masked_shr_fold(self):
-    x = UOp.variable('x', 0, 255, dtype=dtypes.uint32)
+    x = UOp.variable('x', 0, 255, dtype=dtypes.uint32, param=True)
     self.helper_test_variable((x & -4) >> 2, 0, 63, "(x>>2)")
 
   def test_bool_or_not_tautology(self):
@@ -483,12 +484,12 @@ class TestSymbolic(unittest.TestCase):
 
   def test_div_drop_small_terms(self):
     # from openpilot, shouldnt simplify
-    gidx0 = UOp.variable("gidx0", 0, 10)
-    gidx1 = UOp.variable("gidx1", 0, 10)
-    lidx0 = UOp.variable("lidx0", 0, 1)
-    lidx1 = UOp.variable("lidx1", 0, 1)
-    ridx1005 = UOp.variable("ridx1005", 0, 2)
-    ridx1006 = UOp.variable("ridx1006", 0, 2)
+    gidx0 = UOp.variable("gidx0", 0, 10, param=True)
+    gidx1 = UOp.variable("gidx1", 0, 10, param=True)
+    lidx0 = UOp.variable("lidx0", 0, 1, param=True)
+    lidx1 = UOp.variable("lidx1", 0, 1, param=True)
+    ridx1005 = UOp.variable("ridx1005", 0, 2, param=True)
+    ridx1006 = UOp.variable("ridx1006", 0, 2, param=True)
     self.helper_test_variable((lidx1+((gidx1*18)+(ridx1005*18)+(lidx0*162))+(gidx0*2)+(ridx1006*2)+-40)//18, -3, 20,
       "(gidx1+ridx1005+lidx0*9+(gidx0+ridx1006+7)//9+-3)")
 
@@ -948,6 +949,11 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(cond.where(u0, u1), 0, 1, "((a<2)!=True)")
     self.helper_test_variable(cond.where(u0, u1).where(u0, u1), 0, 1, "(a<2)")
 
+  def test_equivalent_const_max(self):
+    x = Variable("x", -10, 10)
+    self.helper_test_variable((x < 0).where(0, x), 0, 10, "x.maximum(0)")
+    self.helper_test_variable((0 < x).where(x, 0), 0, 10, "x.maximum(0)")
+
   def test_where_combine(self):
     cond = Variable("x", 0, 3) < 2
     a = Variable("a", 0, 3)
@@ -992,7 +998,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(cond.ne(False), 0, 1, "(x<2)")
 
   def test_bitcast_chain(self):
-    a = UOp.variable("a", 0, 3, dtype=dtypes.int32)
+    a = UOp.variable("a", 0, 3, dtype=dtypes.int32, param=True)
     self.assertIs(graph_rewrite(a.bitcast(dtypes.float32).bitcast(a.dtype), sym), a)
 
   def test_negation_in_where(self):
@@ -1008,20 +1014,11 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(-a<-b, False, True, "(b<a)")
 
   def test_where_cast(self):
-    s = Variable("s", 0, 3, dtypes.int)
-    cond = s < 2
+    cond = Variable("s", 0, 3, dtypes.int) < 2
     a = Variable("a", 0, 3, dtypes.int)
-    b = Variable("b", 0, 3, dtypes.int)
-    expr = cond.where(a, b).cast(dtypes.half)
-
-    # TODO: copied from render, render does not support cast
-    glbl = UOp.param(0, dtypes.int, (1,))
-    uops = get_uops(UOp(Ops.STORE, src=(glbl.index(UOp.const(0, dtypes.int)), expr)).sink())
-    rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[1]
-
-    # the vars are now scalar PARAMs
-    pvar = {u.expr: u for u in rewritten_uop.toposort() if u.op is Ops.PARAM}
-    self.assertEqual(rewritten_uop, (pvar['s']<UOp.const(2, dtypes.int)).where(pvar['a'].cast(dtypes.half), pvar['b'].cast(dtypes.half)))
+    self.assertIs(graph_rewrite(cond.where(a, a+1).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), (a+1).cast(dtypes.half)))
+    self.assertIs(graph_rewrite(cond.where(a, uconst(2)).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), UOp.const(2, dtypes.half)))
+    self.assertIs(graph_rewrite(cond.where(a, UOp.invalid()).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), UOp.invalid()))
 
   def test_where_merge_branches(self):
     cond1 = Variable("s", 0, 10) < 6
@@ -1175,7 +1172,7 @@ class TestSymbolicVariables(unittest.TestCase):
     assert (a//4 + a//6).variables() == [a]
 
   def test_variable_min_eq_max_bind_folds(self):
-    b = Variable("x", 1, 1).bind(1)
+    b = UOp.variable("x", 1, 1).bind(1)
     s = b.simplify()
     self.assertEqual(s.op, Ops.CONST)
     self.assertEqual(s.val, 1)
@@ -1369,6 +1366,16 @@ class TestInvalidIndex(unittest.TestCase):
     c2 = UOp.const((1, Invalid, 1, 1))
     self.assertIs((c1+c2).simplify(), UOp.const((2, Invalid, Invalid, Invalid)))
 
+  def test_gated_load_keeps_index_valid(self):
+    # the load executes even on gated-off iterations: gated_given_valid must not erase its mask (PADTO OOB shape)
+    buf = UOp.param(0, dtypes.bool, (17,))
+    ridx = Variable("ridx", 0, 31)
+    cond = ridx < 17
+    load = buf.index(ridx.valid(cond))
+    out = graph_rewrite(cond.where(load.where(uconst(2), uconst(0)), UOp.invalid()), sym)
+    idx = next(u for u in out.toposort() if u.op is Ops.INDEX)
+    self.assertIs(idx.src[1].get_valid(), cond.simplify())
+
 class TestStoreLoadFolding(unittest.TestCase):
   """Tests for store(index, load(index)) -> NOOP rule. This rule matches patterns that EMERGE during simplification."""
   def test_store_load_folding(self):

test/null/test_uops.py

@@ -41,11 +41,6 @@ class TestDTypeFromUOp(unittest.TestCase):
     # an explicit (strong) const dtype is legal until the field is removed
     self.assertEqual(UOp.const(3, dtypes.int32).dtype, dtypes.int32)
 
-  def test_weak_dtype_rejected_by_program_spec(self):
-    for weak, concrete, value in ((dtypes.weakint, dtypes.int32, 1), (dtypes.weakfloat, dtypes.float32, 1.0)):
-      with self.assertRaises(RuntimeError): type_verify(UOp.const(value, weak).sink(), spec_program)
-      type_verify(UOp.const(value, concrete).sink(), spec_program)
-
   def test_invalid_stated_dtype(self):
     # UOp.const normalizes a stated dtype away (const_like/full pass their position's); the core constructor does not,
     # and the spec is what rejects a non-bool Invalid
@@ -134,7 +129,7 @@ class TestConstFloatEq(unittest.TestCase):
     self.assertFalse(Invalid != HoldsInvalid())
 
   def test_matchers_agree_on_nan(self):
-    n = UOp.const(math.nan, dtypes.float32)
+    n = UOp.const(math.nan)
     for compiled in (False, True):
       pm = PatternMatcher([(UPat(Ops.CONST, arg=math.nan), lambda: True)], compiled=compiled)
       self.assertTrue(pm.rewrite(n), f"{compiled=}")
@@ -348,10 +343,9 @@ class TestFastIdiv(unittest.TestCase):
   def test_fast_idiv_remove_powers_of_two(self):
     ridx = UOp.range(2**20, 0)
     uops = to_uops_list([ridx//(7*64)], ren=Device[Device.DEFAULT].renderer)
-    ops = [x.op for x in uops]
     # this requires shifting out the powers of two before doing fast_idiv
     # (((ridx0>>6)*18725)>>17) instead of (int)((((long)(ridx0)*1198373)>>29))
-    self.assertNotIn(Ops.CAST, ops)
+    self.assertNotIn(dtypes.long, [x.dtype for x in uops])
 
   @unittest.expectedFailure
   def test_fast_idiv_overflow(self):

test/null/test_validate_oob.py

@@ -1,10 +1,12 @@
 import unittest
-from tinygrad import dtypes, Variable
+from tinygrad import dtypes
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import Ops, UOp, AxisType
 from test.helpers import to_uops_list
 
+def Variable(name, nmin, nmax): return UOp.variable(name, nmin, nmax, param=True)
+
 class TestValidateOOB(unittest.TestCase):
   """Test z3 validation of index bounds for different ALU ops and patterns."""
 

test/null/test_viz.py

@@ -236,8 +236,8 @@ class TestViz(unittest.TestCase):
   def test_const_node_visibility(self):
     with save_viz() as viz:
       a = UOp.variable("a", 0, 10, dtype=dtypes.int)
-      z = UOp.const(0, a.dtype)
-      y = UOp.const(math.pi, dtypes.float)
+      z = UOp.const(0)
+      y = UOp.const(math.pi)
       alu = a*z
       ret = exec_rewrite(sink:=UOp.sink(alu, y), [sym])
     lst = viz.list_items()
@@ -249,7 +249,7 @@ class TestViz(unittest.TestCase):
     self.assertTrue(graphs[0][id(y)]["exclude"])
     self.assertFalse(graphs[0][id(alu)]["exclude"])
     self.assertEqual(graphs[0][id(y)]["label"].split("\n")[:2], ["CONST", "3.14159"])
-    self.assertEqual(list(graphs[1]), [id(z), id(y), id(ret)])
+    self.assertEqual(list(graphs[1]), [id(u) for u in ret.toposort()])  # rewrite graph keys follow the rewritten sink's toposort
 
   def test_const_reshape_expand_folded(self):
     # CONST->EXPAND should be folded into the ALU node, not shown as separate EXPAND nodes
@@ -305,10 +305,10 @@ class TestVizTree(unittest.TestCase):
 
   def test_tree_view(self):
     with save_viz() as viz:
-      a = UOp.variable("a",0,10)
-      b = UOp.variable("b",0,10)
-      c = UOp.variable("c",0,10)
-      d = UOp.variable("d",0,10)
+      a = UOp.variable("a",0,10,param=True)
+      b = UOp.variable("b",0,10,param=True)
+      c = UOp.variable("c",0,10,param=True)
+      d = UOp.variable("d",0,10,param=True)
       sink = UOp.sink(a+b, c+d)
       def tree_rewrite(): return graph_rewrite(sink, root, name="root")
       tree_rewrite()

test/opt/test_gen_float4.py

@@ -10,12 +10,12 @@ from test.helpers import replace_opts
 class TestFloat4(unittest.TestCase):
   @staticmethod
   def count_float4(uops: list[UOp], n=4):
-    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype.scalar() == dtypes.float and uop.shape == (4,)]),
-            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype.scalar() == dtypes.float and uop.shape == (4,)]))
+    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype == dtypes.float and uop.shape == (4,)]),
+            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.float and uop.shape == (4,)]))
   @staticmethod
   def count_half4(uops: list[UOp]):
-    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype.scalar() == dtypes.half and uop.shape == (4,)]),
-            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype.scalar() == dtypes.half and uop.shape == (4,)]))
+    return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype == dtypes.half and uop.shape == (4,)]),
+            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.half and uop.shape == (4,)]))
 
   def test_float4_basic(self):
     a = Tensor.empty(2, 8).realize()

test/opt/test_kernel_opts.py

@@ -239,6 +239,15 @@ class TestKernelOpts(unittest.TestCase):
     helper_linearizer_opt(a.sum().exp(), [[Opt(OptOps.PADTO, 0, 32)],])
     helper_linearizer_opt(a.sum(0).exp(), [[Opt(OptOps.PADTO, 1, 32)],])
 
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
+  @unittest.expectedFailure
+  def test_padto_group_full_unroll_sum(self):
+    a = Tensor.ones(2, 28, 4096, dtype=dtypes.bfloat16).realize()
+    out = ((a * 0.5).float().square()).sum(axis=(0, 2))
+    opts_to_apply = [Opt(OptOps.GROUPTOP, 1, 256), Opt(OptOps.PADTO, 3, 32), Opt(OptOps.UNROLL, 2, 0), Opt(OptOps.UPCAST, 0, 7)]
+    helper_linearizer_opt(out, [opts_to_apply], check_default_opt=False)
+
   def test_padto_sum(self):
     N = 18
     # NOTE: this setup prevents 17 * 17 contiguous merged into one dimension

test/opt/test_tensor_cores.py

@@ -79,7 +79,8 @@ class TestTensorCores(unittest.TestCase):
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
   def test_tensor_cores(self):
     for tc in Device[Device.DEFAULT].renderer.tensor_cores:
-      helper_tc_allclose(tc.dims[0], tc.dims[1], tc.dims[2], tc.dtype_in, tc.dtype_out, axis=0, tc_opt=0)
+      with self.subTest(tc=tc):
+        helper_tc_allclose(tc.dims[0], tc.dims[1], tc.dims[2], tc.dtype_in, tc.dtype_out, axis=0, tc_opt=0)
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
   def test_tensor_cores_nested_reduce(self):
@@ -185,10 +186,10 @@ class TestTensorCores(unittest.TestCase):
     # skip fp8 tcs: the unoptimized ALU baseline quantizes products to fp8 (JAX promotion), which legitimately
     # differs from the MFMA path (f32 accumulation), so the baseline-vs-TC numerical gate can't hold for fp8.
     tc = next(tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in not in dtypes.fp8s)
-    x, y = Tensor.rand(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in)
+    x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in)
     r = x.matmul(y, dtype=tc.dtype_out)
     opts = [Opt(OptOps.UNROLL, 0, 2)]
-    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
+    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False)
     for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src):
       if u.op is Ops.WMMA:
         assert u.src[-1].src[0].op != Ops.STORE
@@ -199,10 +200,10 @@ class TestTensorCores(unittest.TestCase):
   @unittest.skipIf(Device.DEFAULT in {"CPU"}, "CPU does not support using a different type for accumulation")
   def test_tensor_cores_unroll_casted_phi(self):
     tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out and tc.dtype_in not in dtypes.fp8s][0]
-    x, y = Tensor.rand(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in)
+    x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in)
     r = x.matmul(y, dtype=tc.dtype_out)
     opts = [Opt(OptOps.UNROLL, 0, 2)]
-    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
+    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False)
     for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src):
       if u.op is Ops.WMMA:
         #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
@@ -215,10 +216,10 @@ class TestTensorCores(unittest.TestCase):
   def test_tensor_cores_unroll_casted_phi_with_children(self):
     # all STORE children are outside the loop
     tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out and tc.dtype_in not in dtypes.fp8s][0]
-    x, y = Tensor.rand(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in)
+    x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in)
     r = x.matmul(y, dtype=tc.dtype_out).relu()
     opts = [Opt(OptOps.UNROLL, 0, 2)]
-    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
+    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False)
     for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src):
       if u.op is Ops.WMMA:
         #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))

test/unit/test_allreduce.py

@@ -2,6 +2,7 @@ import unittest
 from tinygrad import Tensor, UOp, dtypes
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import Ops
+from test.helpers import KernelCountException
 
 class TestRingAllReduce(unittest.TestCase):
   def test_schedule_ring(self):
@@ -13,7 +14,7 @@ class TestRingAllReduce(unittest.TestCase):
       copies = [si for si in linear.src if si.src[0].op is Ops.COPY]
       pairs = [(c.src[1].buffer.device, c.src[2].buffer.device) for c in copies]
       # N*(N-1) scatter reduce, and N*(N-1) allgather
-      self.assertEqual(len(pairs), N*(N-1)*2)
+      if len(pairs) != N*(N-1)*2: raise KernelCountException(N*(N-1)*2, len(pairs))
       # copy topology forms a ring
       self.assertEqual(len(set(pairs)), N)
 
@@ -25,8 +26,8 @@ class TestRingAllReduce(unittest.TestCase):
       linear = t.sum(0).mul(2.0).contiguous().linear_with_vars()[0]
       copies = [si for si in linear.src if si.src[0].op is Ops.COPY]
       sinks = [si for si in linear.src if si.src[0].op is Ops.SINK]
-      self.assertEqual(len(copies), 24)
-      self.assertEqual(len(sinks), 26)
+      if len(copies) != 24: raise KernelCountException(24, len(copies))
+      if len(sinks) != 26: raise KernelCountException(26, len(sinks))
 
   @Context(RING=0, ALL2ALL=0)
   def test_schedule_naive(self):
@@ -39,8 +40,8 @@ class TestRingAllReduce(unittest.TestCase):
     sinks = [si for si in linear.src if si.src[0].op is Ops.SINK]
     pairs = [(c.src[1].buffer.device, c.src[2].buffer.device) for c in copies]
 
-    self.assertEqual(len(pairs), N*(N-1))
-    self.assertEqual(len(sinks), 2)
+    if len(pairs) != N*(N-1): raise KernelCountException(N*(N-1), len(pairs))
+    if len(sinks) != 2: raise KernelCountException(2, len(sinks))
     self.assertTrue(all(dst != src for dst, src in pairs))
 
   def test_symbolic_shape(self):
@@ -64,7 +65,7 @@ class TestAllreduceCast(unittest.TestCase):
     with Context(ALLREDUCE_CAST=allreduce_cast, RING=0, SCACHE=0):
       t = Tensor.empty(4, 4, dtype=dtype).shard(ds, axis=0)
       linear = t.sum(0).linear_with_vars()[0]
-      return {si.src[1].buffer.dtype.scalar() for si in linear.src if si.src[0].op is Ops.COPY}
+      return {si.src[1].buffer.dtype for si in linear.src if si.src[0].op is Ops.COPY}
 
   def test_allreduce_cast_bf16(self):
     # with ALLREDUCE_CAST, allreduce copies stay in bfloat16 instead of promoting to float32

test/unit/test_assign.py

@@ -540,6 +540,10 @@ class TestAssign(unittest.TestCase):
     c = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
     c[0:2].bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000], dtype=dtypes.uint32)).realize()
     np.testing.assert_allclose(c.numpy(), [4.0, 3.0, 3.0, 4.0])
+    # without .realize()
+    a = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
+    a.bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.uint32))
+    np.testing.assert_allclose(a.numpy(), [4.0, 3.0, 2.0, 1.0])
 
   def test_assign_bitcast_different_size(self):
     # assign to a shape-changing bitcast view (only works on DISK currently)

test/unit/test_attention.py

@@ -1,11 +1,9 @@
 import unittest
-from types import SimpleNamespace
 import numpy as np
 from tinygrad import Tensor, dtypes, nn
-from tinygrad.llm.kimi import _shard_kimi
 from tinygrad.llm.model import (
   GatedDeltaNetBlock, SSMConfig, TransformerBlock, TransformerConfig,
-  apply_rope as apply_rope_new, iterative_topk, l2norm, precompute_freqs_cis, pairwise_topk,
+  apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk,
 )
 
 def apply_rope(x:Tensor, start_pos:int):
@@ -43,19 +41,14 @@ class TestAttention(unittest.TestCase):
     np.testing.assert_allclose(block.cache_kv[0, :, :, :seqlen, :].numpy(), expected.numpy(), rtol=1e-5, atol=1e-5)
 
 class TestGatedDeltaNetBlock(unittest.TestCase):
-  def test_kda_l2norm_matches_fla(self):
-    x = np.array([[1e-4, -2e-4, 3e-4], [1.0, 2.0, -3.0]], dtype=np.float32)
-    expected = x / np.sqrt((x*x).sum(axis=-1, keepdims=True) + 1e-6)
-    np.testing.assert_allclose(l2norm(Tensor(x)).numpy(), expected, rtol=1e-6, atol=1e-6)
-
   def _tensor_linspace(self, start:float, stop:float, shape:tuple[int, ...]) -> Tensor:
     return Tensor.linspace(start, stop, int(np.prod(shape)), dtype=dtypes.float32).reshape(*shape)
 
   def _make_config(self, **kwargs):
-    return TransformerConfig(**({"num_blocks":1, "dim":4, "hidden_dim":8, "n_heads":1, "n_kv_heads":1,
-                                 "norm_eps":1e-5, "vocab_size":32, "head_dim":4, "rope_theta":10000.0,
-                                 "rope_dim":4, "v_head_dim":4, "max_context":4, "ssm_layers":(True,),
-                                 "ssm":SSMConfig(conv_kernel=2, state_size=2, group_count=1, time_step_rank=1, inner_size=2)} | kwargs))
+    return TransformerConfig(**({"num_blocks":1, "dim":32, "hidden_dim":64, "n_heads":1, "n_kv_heads":1,
+                                 "norm_eps":1e-5, "vocab_size":32, "head_dim":32, "rope_theta":10000.0,
+                                 "rope_dim":32, "v_head_dim":32, "max_context":4, "ssm_layers":(True,),
+                                 "ssm":SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32)} | kwargs))
 
   def _make_block(self, config:TransformerConfig) -> GatedDeltaNetBlock:
     block = GatedDeltaNetBlock(config, config.ssm)
@@ -86,6 +79,10 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
       recurrent_state = cache[:, conv_flat:].reshape(cache.shape[0], block.num_v_heads, block.head_v_dim, block.head_v_dim)
       return conv_state, recurrent_state
 
+  def _reset_state(self, block:GatedDeltaNetBlock):
+    Tensor.realize(block.conv_state.assign(block.conv_state.const_like(0)),
+                   block.recurrent_state.assign(block.recurrent_state.const_like(0)))
+
   def _linear_np(self, x:np.ndarray, weight:np.ndarray) -> np.ndarray:
     return x.astype(np.float32) @ weight.T.astype(np.float32)
 
@@ -93,7 +90,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
     x_float = x.astype(np.float32)
     return (x_float / np.sqrt((x_float * x_float).mean(axis=-1, keepdims=True) + eps)) * weight.astype(np.float32)
 
-  def _normalize_np(self, x:np.ndarray, eps:float=1e-12) -> np.ndarray:
+  def _normalize_np(self, x:np.ndarray, eps:float=1e-6) -> np.ndarray:
     return x / np.maximum(np.sqrt((x * x).sum(axis=-1, keepdims=True)), eps)
 
   def _softplus_np(self, x:np.ndarray) -> np.ndarray:
@@ -155,6 +152,12 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
     x = Tensor.linspace(-1.0, 1.0, 3 * config.dim, dtype=dtypes.float32).reshape(1, 3, config.dim)
 
     expected_outs, expected_conv, expected_recurrent = self._naive_attention(block, x)
+    out = self._run_attention(block, x, 0)
+    conv_state, recurrent_state = self._cache_views(block)
+    np.testing.assert_allclose(out, np.concatenate(expected_outs, axis=1), rtol=1e-3, atol=1e-3)
+    np.testing.assert_allclose(conv_state, expected_conv[-1], rtol=1e-3, atol=1e-3)
+    np.testing.assert_allclose(recurrent_state, expected_recurrent[-1], rtol=1e-3, atol=1e-3)
+    self._reset_state(block)
 
     for step in range(x.shape[1]):
       out = self._run_attention(block, x[:, step:step+1], step)
@@ -170,7 +173,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
     prompt = Tensor.linspace(0.75, -0.75, 2 * config.dim, dtype=dtypes.float32).reshape(1, 2, config.dim)
 
     for i in range(warmup.shape[1]): self._run_attention(block, warmup[:, i:i+1], i)
-    Tensor.realize(*block._state_reset_ops())
+    self._reset_state(block)
     expected_outs, expected_conv, expected_recurrent = self._naive_attention(block, prompt)
 
     for step in range(prompt.shape[1]):
@@ -184,99 +187,64 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
                                  err_msg=f"GatedDeltaNet reset recurrent cache mismatch at step {step}")
 
   def test_kda_channel_decay(self):
-    config = self._make_config(n_heads=2, ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=2, time_step_rank=2, inner_size=4, kda=True))
-    block, x = GatedDeltaNetBlock(config, config.ssm), Tensor([[[1., 2., 0., 0.]]])
-    # f_b(f_a(x)) = [1, 2, 3, 4]
+    config = self._make_config(dim=4, hidden_dim=8, n_heads=2, head_dim=4, rope_dim=4, v_head_dim=4,
+      ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=2, time_step_rank=2, inner_size=4, kda=True))
+    block, x = GatedDeltaNetBlock(config, config.ssm), Tensor([[[1., 2., 0., 0.], [2., 1., 0., 0.]]])
     block.ssm_f_a.weight = Tensor([[1., 0., 0., 0.], [0., 1., 0., 0.]])
     block.ssm_f_b.weight = Tensor([[1., 0.], [0., 1.], [1., 1.], [2., 1.]])
     block._init_state(x)
     initial_state = Tensor.arange(8, dtype=dtypes.float32).reshape(1, 2, 2, 2)
     block.recurrent_state.assign(initial_state).realize()
     block.ssm_a = Tensor([[-1.], [-1.]])
-    block._attention(x, 0).realize()
-    alpha = np.exp(-self._softplus_np(np.arange(1, 5)).reshape(1, 2, 1, 2))
-    np.testing.assert_allclose(block.recurrent_state.numpy(), initial_state.numpy() * alpha, rtol=1e-5, atol=1e-5)
+    block._attention(x, x.shape[1]).realize()
+    alpha = np.exp(-self._softplus_np(np.array([[1, 2, 3, 4], [2, 1, 3, 5]])).reshape(2, 2, 2)).prod(0)
+    np.testing.assert_allclose(block.recurrent_state.numpy(), initial_state.numpy() * alpha[..., None], rtol=1e-5, atol=1e-5)
 
-  def test_kda_safe_gate_decay(self):
-    config = self._make_config(n_heads=2, kda_full_rank_gate=True, kda_gate_lower_bound=-5.0,
-      ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=2, time_step_rank=2, inner_size=4, kda=True))
-    block, x = GatedDeltaNetBlock(config, config.ssm), Tensor([[[1., 2., 0., 0.]]])
-    block.ssm_f_a.weight = Tensor([[1., 0., 0., 0.], [0., 1., 0., 0.]])
-    block.ssm_f_b.weight = Tensor([[1., 0.], [0., 1.], [1., 1.], [2., 1.]])
-    block.ssm_dt["bias"] = Tensor.zeros(4)
-    block.ssm_a = Tensor([[-2.], [-3.]])  # stores -exp(A_log)
-    block._init_state(x)
-    initial_state = Tensor.arange(8, dtype=dtypes.float32).reshape(1, 2, 2, 2)
-    block.recurrent_state.assign(initial_state).realize()
-    block._attention(x, 0).realize()
-    gate_logits = np.arange(1, 5, dtype=np.float32).reshape(1, 2, 2)
-    exp_a = np.array([2., 3.], dtype=np.float32).reshape(1, 2, 1)
-    alpha = np.exp(-5.0 / (1.0 + np.exp(-(exp_a * gate_logits)))).reshape(1, 2, 1, 2)
-    np.testing.assert_allclose(block.recurrent_state.numpy(), initial_state.numpy() * alpha, rtol=2e-5, atol=2e-5)
+  def test_kda_prefill_matches_decode(self):
+    config = self._make_config(ssm=SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32, kda=True))
+    block = GatedDeltaNetBlock(config, config.ssm)
+    for p in nn.state.get_parameters(block):
+      p.replace(self._tensor_linspace(-0.05, 0.05, p.shape) if len(p.shape) > 1 else self._tensor_linspace(0.05, 0.1, p.shape))
+    x = self._tensor_linspace(-0.5, 0.5, (1, 3, config.dim))
+    prefill = self._run_attention(block, x, 0)
+    prefill_conv, prefill_recurrent = self._cache_views(block)
+    self._reset_state(block)
+    decode = np.concatenate([self._run_attention(block, x[:, i:i+1], i) for i in range(3)], axis=1)
+    decode_conv, decode_recurrent = self._cache_views(block)
+    np.testing.assert_allclose(prefill, decode, rtol=1e-3, atol=1e-3)
+    np.testing.assert_allclose(prefill_conv, decode_conv, rtol=1e-3, atol=1e-3)
+    np.testing.assert_allclose(prefill_recurrent, decode_recurrent, rtol=1e-3, atol=1e-3)
 
-  def test_kda_per_channel_a(self):
-    config = self._make_config(n_heads=2, kda_full_rank_gate=True, kda_gate_lower_bound=-5.0,
-      ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=2, time_step_rank=2, inner_size=4, kda=True, channel_decay=True))
-    block, x = GatedDeltaNetBlock(config, config.ssm), Tensor([[[1., 2., 0., 0.]]])
-    block.ssm_f_a.weight = Tensor([[1., 0., 0., 0.], [0., 1., 0., 0.]])
-    block.ssm_f_b.weight = Tensor([[1., 0.], [0., 1.], [1., 1.], [2., 1.]])
-    block.ssm_dt["bias"] = Tensor.zeros(4)
-    block.ssm_a = Tensor([[-2.], [-3.]])
-    block._init_state(x)
-    initial_state = Tensor.arange(8, dtype=dtypes.float32).reshape(1, 2, 2, 2)
-    block.recurrent_state.assign(initial_state).realize()
-    block._attention(x, 0).realize()
-    gate_logits = np.arange(1, 5, dtype=np.float32).reshape(1, 2, 2)
-    exp_a = np.array([2., 3.], dtype=np.float32).reshape(1, 1, 2)
-    alpha = np.exp(-5.0 / (1.0 + np.exp(-(exp_a * gate_logits)))).reshape(1, 2, 1, 2)
-    np.testing.assert_allclose(block.recurrent_state.numpy(), initial_state.numpy() * alpha, rtol=2e-5, atol=2e-5)
+  def test_varied_chunk_sizes_match_decode(self):
+    for kda in (False, True):
+      ssm = SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32, kda=kda)
+      config = self._make_config(ssm=ssm)
+      if kda:
+        block = GatedDeltaNetBlock(config, config.ssm)
+        for p in nn.state.get_parameters(block):
+          p.replace(self._tensor_linspace(-0.05, 0.05, p.shape) if len(p.shape) > 1 else self._tensor_linspace(0.05, 0.1, p.shape))
+      else: block = self._make_block(config)
+      x = self._tensor_linspace(-0.5, 0.5, (1, 4, config.dim))
+      decode = np.concatenate([self._run_attention(block, x[:, i:i+1], i) for i in range(4)], axis=1)
+      decode_conv, decode_recurrent = self._cache_views(block)
+      for chunking in ([4], [2, 2], [1, 3], [3, 1], [2, 1, 1]):
+        self._reset_state(block)
+        outs, start = [], 0
+        for size in chunking:
+          outs.append(self._run_attention(block, x[:, start:start+size], start))
+          start += size
+        chunked_conv, chunked_recurrent = self._cache_views(block)
+        np.testing.assert_allclose(np.concatenate(outs, axis=1), decode, rtol=1e-3, atol=1e-3, err_msg=f"{kda=} {chunking=}")
+        np.testing.assert_allclose(chunked_conv, decode_conv, rtol=1e-3, atol=1e-3, err_msg=f"{kda=} {chunking=}")
+        np.testing.assert_allclose(chunked_recurrent, decode_recurrent, rtol=1e-3, atol=1e-3, err_msg=f"{kda=} {chunking=}")
 
-  def test_kda_chunked_prefill_matches_decode(self):
-    config = self._make_config(max_context=4, n_heads=2,
-      ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=2, time_step_rank=2, inner_size=4, kda=True), kda_split_qkv=True)
-    x = Tensor.linspace(-1, 1, 4*config.dim, dtype=dtypes.float32).reshape(1, 4, config.dim).cast(dtypes.bfloat16)
-
-    chunked = GatedDeltaNetBlock(config, config.ssm)
-    sequential = GatedDeltaNetBlock(config, config.ssm)
-    for value in nn.state.get_state_dict(chunked).values(): value.replace(value.cast(dtypes.bfloat16).realize())
-    sequential_state = nn.state.get_state_dict(sequential)
-    for name, value in nn.state.get_state_dict(chunked).items(): sequential_state[name].replace(value)
-
-    chunked._init_state(x)
-    chunk_out = chunked._attention(x, 0).realize()
-    sequential._init_state(x)
-    seq_out = Tensor.cat(*[sequential._attention(x[:, t:t+1], t).realize() for t in range(x.shape[1])], dim=1).realize()
-
-    np.testing.assert_allclose(chunk_out.numpy(), seq_out.numpy(), rtol=1e-5, atol=1e-5)
-    for name in ("conv_state_q", "conv_state_k", "conv_state_v"):
-      np.testing.assert_allclose(getattr(chunked, name).numpy(), getattr(sequential, name).numpy(), rtol=2e-2, atol=4e-3)
-    np.testing.assert_allclose(chunked.recurrent_state.numpy(), sequential.recurrent_state.numpy(), rtol=2e-3, atol=2e-3)
-
-  def test_kda_tp_final_token_matches_unsharded(self):
-    config = self._make_config(dim=8, hidden_dim=16, n_heads=4, n_kv_heads=4, head_dim=2, rope_dim=2, v_head_dim=2,
-      ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=4, time_step_rank=4, inner_size=8, kda=True), kda_split_qkv=True)
-    single, tp = GatedDeltaNetBlock(config, config.ssm), GatedDeltaNetBlock(config, config.ssm)
-    for name, value in nn.state.get_state_dict(single).items():
-      data = np.full(value.shape, 1.0, np.float32) if "norm.weight" in name else \
-             np.linspace(-0.2, 0.2, value.numel(), dtype=np.float32).reshape(value.shape)
-      if name == "ssm_a": data.fill(-0.1)
-      value.replace(Tensor(data, device="CPU", dtype=dtypes.bfloat16).realize())
-    tp_state = nn.state.get_state_dict(tp)
-    for name, value in nn.state.get_state_dict(single).items(): tp_state[name].replace(value)
-    devices = ("CPU", "CPU:1")
-    _shard_kimi(SimpleNamespace(blk=[tp]), devices)
-
-    x = Tensor(np.linspace(-1, 1, 32, dtype=np.float32).reshape(1, 4, 8), device="CPU", dtype=dtypes.bfloat16)
-    single._init_state(x)
-    expected = single._attention(x, 0).realize()
-    x_tp = x.shard(devices, axis=None)
-    tp._init_state(x_tp)
-    actual = tp._attention(x_tp, 0).realize()
-
-    np.testing.assert_equal(actual.numpy(), expected.numpy())
-    np.testing.assert_equal(tp.recurrent_state.numpy(), single.recurrent_state.numpy())
-    for name in ("conv_state_q", "conv_state_k", "conv_state_v"):
-      np.testing.assert_equal(getattr(tp, name).numpy(), getattr(single, name).numpy())
+  def test_start_zero_resets_realized_state(self):
+    config, x = self._make_config(max_context=3), self._tensor_linspace(-1, 1, (1, 3, 32))
+    block = self._make_block(config)
+    self._run_attention(block, x, 0)
+    restarted = self._run_attention(block, x[:, :2], 0)
+    fresh = self._run_attention(self._make_block(config), x[:, :2], 0)
+    np.testing.assert_allclose(restarted, fresh, rtol=1e-3, atol=1e-3)
 
 class TestPairwiseTopk(unittest.TestCase):
   def test_basic_topk(self):
@@ -301,13 +269,5 @@ class TestPairwiseTopk(unittest.TestCase):
         self.assertEqual(set(sel.numpy()[b, t].tolist()), expected)
         np.testing.assert_allclose(vals.numpy()[b, t], data[b, t][sel.numpy()[b, t]])
 
-  def test_iterative_matches_numpy(self):
-    rng = np.random.default_rng(42)
-    data = rng.standard_normal((2, 3, 896), dtype=np.float32)
-    vals, sel = iterative_topk(Tensor(data), 16)
-    expected = np.argsort(-data, axis=-1, stable=True)[..., :16]
-    np.testing.assert_equal(sel.numpy(), expected)
-    np.testing.assert_allclose(vals.numpy(), np.take_along_axis(data, expected, axis=-1))
-
 if __name__ == '__main__':
   unittest.main()

test/unit/test_dtype_weak.py

@@ -3,11 +3,12 @@ import tempfile, unittest, math
 from tinygrad import Tensor, dtypes, TinyJit
 from tinygrad.helpers import Context
 from tinygrad.dtype import least_upper_float
-from tinygrad.uop.ops import UOp, Ops, dtype_from_uop, graph_rewrite
+from tinygrad.uop.ops import UOp, Ops, GroupOp, dtype_from_uop, graph_rewrite
 from tinygrad.uop.weak import pm_lower_index_dtype, pm_commit_weak
 from tinygrad.uop.symbolic import symbolic_simple
 from tinygrad.uop.spec import spec_shared, type_verify
 from tinygrad.engine.jit import JitError
+from test.helpers import full_rewrite
 
 
 class TestWeakPromotion(unittest.TestCase):
@@ -76,6 +77,11 @@ class TestWeakPromotion(unittest.TestCase):
       committed = graph_rewrite((UOp.const(1).cast(dtypes.int32) + UOp.const(1.0)).cast(dtypes.float32), pm_lower_index_dtype, ctx={})
     self.assertEqual([u.dtype for u in committed.toposort() if u.op is Ops.ADD], [dtypes.float32])
 
+  def test_div_sub_operand_kept_weak(self):
+    a = Tensor.empty(4, dtype=dtypes.float32)
+    for t in (a / 1, a - 0):
+      self.assertEqual(t.uop.src[1].dtype, dtypes.weakfloat)
+
   def test_cast_weak_expression_commits_at_cast_floor(self):
     # the floor never narrows: a cast BELOW the default does not pull the compute width down with it
     with Context(DEFAULT_FLOAT=dtypes.float32):
@@ -88,6 +94,13 @@ class TestWeakPromotion(unittest.TestCase):
       out = Tensor(1.0, dtype=dtypes.float32, device="CPU") / denom
       self.assertAlmostEqual(out.item(), 1 / (70000 + 1e-5), places=10)
 
+  def test_stacked_weak_casts_convert_each_kind(self):
+    # each weak cast is a kind conversion: weakint truncates before weakfloat re-lifts (neither is only a marker)
+    x = Tensor([2.5, -3.7], dtype=dtypes.float32, device="CPU")
+    stacked = x.cast(dtypes.weakint).cast(dtypes.weakfloat)
+    self.assertIs(stacked.dtype, dtypes.weakfloat)
+    self.assertEqual(stacked.tolist(), [2.0, -3.0])
+
   def test_uop_scalar_const_lifts_kind(self):
     for dtype, value, out_dtype, const_dtype in ((dtypes.weakint, 1, dtypes.weakint, dtypes.weakint),
                                                  (dtypes.int32, 1, dtypes.int32, dtypes.weakint),
@@ -276,5 +289,18 @@ class TestSignedUint64Weakfloat(unittest.TestCase):
     self.assertAlmostEqual((i64 + u64).sin().item(), math.sin(2), places=5)  # Unary lowers before transcendental
 
 
+class TestNoRedundantWide(unittest.TestCase):
+  def wide_alu(self, t:Tensor) -> int:
+    return sum(sum(1 for u in full_rewrite(call.src[0]).toposort() if u.op in GroupOp.ALU and u.dtype in {dtypes.long, dtypes.ulong})
+               for call in t.schedule_linear().src if call.src[0].op is Ops.SINK)
+
+  def test_unbounded_long_stays_long(self):
+    self.assertGreater(self.wide_alu(Tensor.empty(16, dtype=dtypes.long)*3 + 1), 0)
+
+  def test_fancy_index_has_no_wide_alu(self):
+    j, o = Tensor([0, 1, 2]).reshape(3, 1), Tensor([0, 1]).reshape(1, 2)
+    self.assertEqual(self.wide_alu(Tensor.empty(8, 9, 10, 11, 12)[1, j, 2, o, 2]), 0)
+
+
 if __name__ == "__main__":
   unittest.main()

test/unit/test_function.py

@@ -4,7 +4,7 @@ from tinygrad.function import function
 from tinygrad import Tensor, GlobalCounters, Device
 from tinygrad.dtype import Invalid
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, ProgramInfo
-from test.helpers import assert_kernel_count
+from test.helpers import assert_kernel_count, KernelCountException
 
 class TestFunction(unittest.TestCase):
   def test_simple(self):
@@ -516,7 +516,7 @@ class TestFunctionTuple(unittest.TestCase):
     Tensor.realize(a)
     c = f(a)
 
-    self.assertEqual(count_kernels(c), 1)
+    if count_kernels(c) != 1: raise KernelCountException(1, count_kernels(c))
 
     c.sum().backward()
     Tensor.realize(a.grad)

test/unit/test_gradient.py

@@ -51,6 +51,10 @@ class TestTensorGradient(unittest.TestCase):
     with self.assertRaises(RuntimeError): x.sum().gradient(x)
     with self.assertRaises(RuntimeError): x.float().sum().gradient(x)
 
+  def test_const_target_raise(self):
+    t = Tensor(2.0)
+    with self.assertRaises(RuntimeError): (t * 2.0).gradient(t)
+
   def test_copy_to_device_gradient(self):
     t = Tensor([1.0, 2, 3]).realize()
     t.to("CPU:1").square().sum().backward()
@@ -100,7 +104,7 @@ class TestTensorGradient(unittest.TestCase):
 
   def test_implicit_broadcast_where_gradient(self):
     # WHERE with a bare ()-shape branch: the scalar's gradient counts the positions where it is selected
-    cond, x, w = Tensor([True, False, True]), Tensor([1.0, 2.0, 3.0]), Tensor(4.0)
+    cond, x, w = Tensor([True, False, True]), Tensor([1.0, 2.0, 3.0]), Tensor(4.0, dtype=dtypes.float32)
     dw = Tensor(cond.uop.alu(Ops.WHERE, x.uop, w.uop)).sum().gradient(w)[0]
     self.assertEqual(dw.shape, ())
     self.assertEqual(dw.item(), 1.0)
@@ -109,7 +113,7 @@ class TestTensorGradient(unittest.TestCase):
 
   def test_implicit_broadcast_alu_gradient(self):
     # MUL with a bare ()-shape src, no EXPAND in the graph
-    x, w = Tensor([1.0, 2.0, 3.0]), Tensor(2.0)
+    x, w = Tensor([1.0, 2.0, 3.0]), Tensor(2.0, dtype=dtypes.float32)
     m = x.uop.alu(Ops.MUL, w.uop)
     self.assertIs(m.src[1], w.uop)
     dw = Tensor(m).sum().gradient(w)[0]
@@ -118,7 +122,7 @@ class TestTensorGradient(unittest.TestCase):
 
   def test_implicit_broadcast_intermediate_accumulation(self):
     # s is used directly and through an implicit broadcast edge, each edge's gradient reduces to s's shape before they sum
-    x, p = Tensor([1.0, 2.0, 3.0]), Tensor(0.5)
+    x, p = Tensor([1.0, 2.0, 3.0]), Tensor(0.5, dtype=dtypes.float32)
     s = p.sin()
     z = Tensor(x.uop.alu(Ops.MUL, s.uop)).sum() + s
     dp = z.gradient(p)[0]

test/unit/test_hcq_graph.py

@@ -25,10 +25,10 @@ class TestHCQUnit(unittest.TestCase):
     cpu_call = UOp(Ops.PROGRAM, src=(UOp.sink(),)).call(UOp.new_buffer("CPU", 1, dtypes.float))
     gpu_devs = [d0]
 
-    # local MMIO: GPU works alone and with CPU in batch (cpu_support=True)
+    # CPU uses HCQ2 and is no longer batched into legacy HCQ graphs.
     assert HCQGraph.supports_uop(gpu_devs, gpu_call) is True
-    assert HCQGraph.supports_uop(gpu_devs, cpu_call) is True
-    assert HCQGraph.supports_uop(gpu_devs + [cpu_dev], gpu_call) is True
+    assert HCQGraph.supports_uop(gpu_devs, cpu_call) is False
+    assert HCQGraph.supports_uop(gpu_devs + [cpu_dev], gpu_call) is False
 
     # USB MMIO: GPU-only still works, but CPU batching must be rejected (cpu_support=False)
     orig_view = d0.timeline_signal.base_buf.view

test/unit/test_llm_cli_k3.py

@@ -1,30 +0,0 @@
-import unittest
-from tinygrad.llm.cli import KimiK3Template
-from tinygrad.llm.serve import StreamRouter
-
-class TestKimiK3Template(unittest.TestCase):
-  def test_simple_text_chat(self):
-    template = KimiK3Template()
-    got = template.render([{"role":"system", "content":"Be concise."}, {"role":"user", "content":"Hello"}])
-    self.assertTrue(got.startswith('<|open|>message role="system" type="thinking-effort"<|sep|>'))
-    self.assertIn('<|open|>message role="user"<|sep|>Hello<|close|>message<|sep|><|end_of_msg|>', got)
-    self.assertTrue(got.endswith('<|open|>message role="assistant"<|sep|><|open|>think<|sep|>'))
-
-  def test_preserves_assistant_thinking(self):
-    got = KimiK3Template().render([{"role":"assistant", "reasoning_content":"why", "content":"answer"}], add_generation_prompt=False)
-    self.assertIn('<|open|>think<|sep|>why<|close|>think<|sep|>', got)
-    self.assertIn('<|open|>response<|sep|>answer<|close|>response<|sep|>', got)
-
-  def test_rejects_unimplemented_modalities(self):
-    with self.assertRaisesRegex(ValueError, "text-only"):
-      KimiK3Template().render([{"role":"user", "content":[{"type":"image", "url":"x"}]}])
-    with self.assertRaisesRegex(ValueError, "tool rendering"):
-      KimiK3Template().render([{"role":"user", "content":"x"}], tools=[{"type":"function"}])
-
-  def test_xtml_stream_router(self):
-    router, routed = StreamRouter(reasoning=True, xtml=True), []
-    for piece in ("rea", "son<|close|>thi", "nk<|sep|><|open|>response<|sep|>ans", "wer<|close|>response<|sep|>"):
-      routed.extend(router.route(piece))
-    self.assertEqual(routed, [("reasoning_content", "rea"), ("reasoning_content", "son"), ("content", "ans"), ("content", "wer")])
-
-if __name__ == "__main__": unittest.main()

test/unit/test_llm_k3.py

@@ -1,139 +0,0 @@
-import tempfile, unittest
-from pathlib import Path
-from dataclasses import replace
-import numpy as np
-from tinygrad import Tensor, dtypes, nn
-from tinygrad.helpers import getenv
-from tinygrad.llm.kernels import bf16_mfma_splitk
-from tinygrad.llm.kimi_k3 import KIMI_K3_FULL_ATTN_LAYERS, KIMI_K3_SSM_LAYERS, KIMI_K3_TEXT_SIZE, KIMI_K3_TP8_BYTES_PER_GPU, \
-  _layer_sources, _load_stacked_experts, _replace, _safe_load_selected, _shard_kimi_k3, _validate_config, kimi_k3_config, kimi_k3_smoke_config
-from tinygrad.llm.model import FFNBlock, Transformer
-
-def small_k3_config(max_context:int=4): return replace(kimi_k3_smoke_config(max_context), num_experts=8)
-
-class TestKimiK3(unittest.TestCase):
-  def test_selective_safetensor_load(self):
-    with tempfile.TemporaryDirectory() as tmp:
-      path = Path(tmp) / "weights.safetensors"
-      nn.state.safe_save({"keep":Tensor.arange(8), "skip":Tensor.arange(16)}, str(path))
-      selected = _safe_load_selected(path, ["keep"])
-      self.assertEqual(list(selected), ["keep"])
-      np.testing.assert_equal(selected["keep"].numpy(), np.arange(8))
-      with self.assertRaisesRegex(ValueError, "missing tensor absent"): _safe_load_selected(path, ["absent"])
-
-  def test_smoke_config_preserves_gfx950_expert_alignment(self):
-    c = kimi_k3_smoke_config()
-    self.assertEqual(c.routed_expert_dim % 64, 0)
-    self.assertEqual((c.hidden_dim // 8) % 64, 0)
-
-  @unittest.skipUnless(getenv("DEV", "") == "NULL:HIP:gfx950", "gfx950 compile coverage")
-  def test_gfx950_mfma_splitk_compile(self):
-    x = Tensor.zeros(1, 1, 256, dtype=dtypes.bfloat16, device="NULL:HIP:gfx950")
-    weight = Tensor.zeros(16, 256, dtype=dtypes.bfloat16, device="NULL:HIP:gfx950")
-    self.assertEqual(bf16_mfma_splitk(x, weight).realize().shape, (1, 1, 16))
-
-  def test_official_config(self):
-    c = kimi_k3_config(1_048_576)
-    self.assertEqual((c.num_blocks, c.dim, c.n_heads, c.num_experts, c.num_experts_per_tok), (93, 7168, 96, 896, 16))
-    self.assertEqual((sum(KIMI_K3_SSM_LAYERS), len(KIMI_K3_FULL_ATTN_LAYERS)), (69, 24))
-    self.assertEqual(KIMI_K3_FULL_ATTN_LAYERS, (*range(3, 93, 4), 92))
-    self.assertEqual((c.routed_expert_dim, c.hidden_dim, c.shared_expert_dim), (3584, 3072, 6144))
-    self.assertTrue(c.route_weights_uncorrected and c.kda_full_rank_gate and c.attn_output_gate)
-    self.assertTrue(c.ssm is not None and c.ssm.channel_decay)
-    self.assertEqual((c.activation_situ_beta, c.activation_situ_linear_beta, c.kda_gate_lower_bound), (4.0, 25.0, -5.0))
-
-  def test_config_rejects_wrong_checkpoint(self):
-    with self.assertRaisesRegex(ValueError, "not the supported official"):
-      _validate_config({"model_type":"kimi_linear", "hidden_size":2304})
-
-  def test_official_mapping_covers_model(self):
-    model = Transformer(kimi_k3_config(1))
-    state = nn.state.get_state_dict(model)
-    targets = {"token_embd.weight", "output_norm.weight", "output.weight", "output_attn_res_norm.weight", "output_attn_res_proj.weight"}
-    for i,is_kda in enumerate(KIMI_K3_SSM_LAYERS):
-      for target in _layer_sources(i, is_kda).values(): targets.update(target.split("|"))
-      if i:
-        for name in ("ffn_gate_exps.weight", "ffn_gate_exps.weight_scale", "ffn_up_exps.weight", "ffn_up_exps.weight_scale",
-                     "ffn_down_exps.weight", "ffn_down_exps.weight_scale"): targets.add(f"blk.{i}.{name}")
-    self.assertEqual(targets, set(state))
-    self.assertEqual(state["blk.1.ffn_gate_exps.weight"].shape, (896, 3072, 1792))
-    self.assertEqual(state["blk.1.ffn_gate_exps.weight_scale"].shape, (896, 3072, 112))
-    self.assertEqual(state["blk.0.ssm_a"].shape, (128, 1))
-    _shard_kimi_k3(model, tuple(f"NULL:{i}" for i in range(8)))
-    total, per_gpu = 0, 0
-    for name,value in state.items():
-      dtype = dtypes.uint8 if name.endswith(("weight_scale", "_exps.weight")) else dtypes.float32 if name.endswith(
-        ("exp_probs_b.bias", "ssm_q_conv1d.weight", "ssm_k_conv1d.weight", "ssm_v_conv1d.weight", "ssm_norm.weight", "ssm_a", "ssm_dt.bias")) \
-        else dtypes.bfloat16
-      size = value.numel() * dtype.itemsize
-      total += size
-      per_gpu += size if value.uop.axis is None else size//8
-    self.assertEqual((total, per_gpu), (KIMI_K3_TEXT_SIZE, KIMI_K3_TP8_BYTES_PER_GPU))
-
-  def test_situ_matches_reference(self):
-    block = FFNBlock(small_k3_config())
-    gate, up = Tensor([[-8., -1., 0., 3.]]), Tensor([[-30., -2., 5., 40.]])
-    got = block._activation(gate, up).numpy()
-    g, u = gate.numpy().astype(np.float32), up.numpy().astype(np.float32)
-    expected = (4*np.tanh(g/4)/(1+np.exp(-g))) * (25*np.tanh(u/25))
-    np.testing.assert_allclose(got, expected, rtol=1e-5, atol=1e-5)
-
-  def test_attention_residual_matches_reference(self):
-    block = FFNBlock(small_k3_config())
-    block.attn_res_norm.weight.assign([1.0+i/16 for i in range(32)])
-    block.attn_res_proj.weight.assign([[(-1.0)**i/8 for i in range(32)]])
-    prefix, residual = Tensor.arange(64).reshape(2, 32).float()/16, Tensor.arange(128).reshape(2, 2, 32).float()/32
-    got = block._apply_attn_res(prefix, residual, block.attn_res_proj, block.attn_res_norm).numpy()
-    v = np.concatenate((residual.numpy(), prefix.numpy()[:, None]), axis=1).astype(np.float32)
-    k = v / np.sqrt(np.mean(v*v, axis=-1, keepdims=True) + 1e-5)
-    scores = np.sum(k * block.attn_res_norm.weight.numpy() * block.attn_res_proj.weight.numpy()[0], axis=-1)
-    probs = np.exp(scores-scores.max(axis=-1, keepdims=True))
-    probs /= probs.sum(axis=-1, keepdims=True)
-    expected = np.matmul(probs[:, None], v).squeeze(1)
-    np.testing.assert_allclose(got, expected, rtol=1e-5, atol=1e-5)
-
-  def test_tp8_schema(self):
-    model = Transformer(small_k3_config())
-    _shard_kimi_k3(model, tuple(f"NULL:{i}" for i in range(8)))
-    state = nn.state.get_state_dict(model)
-    for name,axis in (("token_embd.weight",0), ("blk.1.ffn_gate_exps.weight",1), ("blk.1.ffn_down_exps.weight_scale",2),
-                      ("blk.1.ffn_routed_down.weight",1), ("blk.0.ssm_g_full.weight",0), ("blk.1.attn_q_b.weight",0)):
-      self.assertEqual(state[name].uop.axis, axis, name)
-    self.assertIsNone(state["blk.1.attn_res_norm.weight"].uop.axis)
-    self.assertIsNone(state["blk.1.ffn_routed_norm.weight"].uop.axis)
-    self.assertIsNone(state["blk.0.ssm_a"].uop.axis)
-
-  def test_direct_expert_staging(self):
-    devices = tuple(f"PYTHON:{i}" for i in range(4))
-    sources = [Tensor([[(e*40+r*4+c)&255 for c in range(4)] for r in range(8)], dtype=dtypes.uint8,
-                      device=devices[0]).realize() for e in range(8)]
-    expected = Tensor.stack(*sources).numpy()
-    for axis in (1, 2):
-      dst = Tensor.zeros(8, 8, 4, dtype=dtypes.uint8, device=devices[0]).shard(devices, axis=axis)
-      _load_stacked_experts(dst, sources)
-      np.testing.assert_equal(dst.numpy(), expected)
-
-  def test_direct_tp_replacement(self):
-    devices = tuple(f"PYTHON:{i}" for i in range(4))
-    source = Tensor.arange(64, dtype=dtypes.float32).reshape(8, 8).realize()
-    expected = source.numpy()
-    for axis in (None, 0, 1):
-      dst = Tensor.zeros(8, 8, device="PYTHON").shard(devices, axis=axis)
-      _replace(dst, source)
-      np.testing.assert_equal(dst.numpy(), expected)
-
-  def test_chunked_recurrent_generate(self):
-    model = Transformer(small_k3_config(max_context=8))
-    for name,value in nn.state.get_state_dict(model).items():
-      fill = 127 if name.endswith("weight_scale") else 0
-      value.replace(Tensor.full(value.shape, fill, dtype=value.dtype if value.dtype is dtypes.uint8 else dtypes.bfloat16, device="PYTHON"))
-    self.assertIsInstance(next(model.generate([1], chunk_size=2)), int)
-    prompt = [1, 2, 3, 4]
-    for _ in range(3): self.assertIsInstance(next(model.generate(prompt.copy(), chunk_size=2)), int)
-    self.assertEqual(model.get_start_pos(model._cached_tokens + [42]), len(prompt))
-    self.assertEqual(model.get_start_pos([9, 2, 3, 4, 42]), 0)
-    self.assertIsInstance(next(model.generate([1, 2, 3, 4, 5], chunk_size=3)), int)
-    self.assertEqual(set(model.recurrent_greedy_prefill_jits), {2})
-    self.assertEqual(model._cached_tokens[:4], [1, 2, 3, 4])
-
-if __name__ == "__main__": unittest.main()

test/unit/test_llm_kimi.py

@@ -1,37 +0,0 @@
-import unittest
-from tinygrad import dtypes, nn
-from tinygrad.llm.kimi import KIMI_LOGICAL_BYTES, KIMI_SSM_LAYERS, KIMI_TENSOR_COUNT, _shard_kimi, _validate_kimi_state, kimi_config
-from tinygrad.llm.model import Transformer
-
-class TestKimiLinear(unittest.TestCase):
-  def test_architecture_config(self):
-    config = kimi_config(4096)
-    self.assertEqual((config.num_blocks, config.dim, config.n_heads, config.vocab_size), (27, 2304, 32, 163840))
-    self.assertEqual(tuple(i for i, is_kda in enumerate(KIMI_SSM_LAYERS) if not is_kda), (3, 7, 11, 15, 19, 23, 26))
-    self.assertEqual((config.num_experts, config.num_experts_per_tok, config.shared_expert_dim), (256, 8, 1024))
-    self.assertTrue(config.expert_mxfp4 and config.bf16_activations and config.kda_split_qkv)
-    self.assertFalse(config.shared_expert_gate)
-
-  def test_tp4_schema_and_axes(self):
-    model = Transformer(kimi_config(32))
-    state = nn.state.get_state_dict(model)
-    self.assertEqual(len(state), KIMI_TENSOR_COUNT)
-    self.assertNotIn("blk.1.ffn_gate_inp_shexp.weight", state)
-    self.assertEqual(state["blk.1.ffn_gate_exps.weight"].dtype, dtypes.uint8)
-    self.assertEqual(state["blk.1.ffn_gate_exps.weight_scale"].dtype, dtypes.uint8)
-
-    _shard_kimi(model, ("NULL:0", "NULL:1", "NULL:2", "NULL:3"))
-    state = nn.state.get_state_dict(model)
-    for name, axis in (("token_embd.weight", 0), ("blk.1.ffn_gate_exps.weight", 1),
-                       ("blk.1.ffn_down_exps.weight_scale", 2), ("blk.3.attn_k_b.weight", 0)):
-      self.assertEqual(state[name].uop.axis, axis, name)
-    self.assertIsNone(state["blk.1.attn_norm.weight"].uop.axis)
-
-  def test_converted_schema_validation(self):
-    model = Transformer(kimi_config(1))
-    state = {name:value if value.dtype is dtypes.uint8 else value.cast(dtypes.bfloat16)
-             for name,value in nn.state.get_state_dict(model).items()}
-    _validate_kimi_state(model, state)
-    self.assertEqual(sum(value.nbytes() for value in state.values()), KIMI_LOGICAL_BYTES)
-
-if __name__ == "__main__": unittest.main()

test/unit/test_llm_moe.py

@@ -2,7 +2,7 @@ import unittest
 import numpy as np
 from dataclasses import replace
 from tinygrad import Tensor
-from tinygrad.llm.model import TransformerBlock, TransformerConfig
+from tinygrad.llm.model import ExpertGating, TransformerBlock, TransformerConfig
 
 def _moe_config(dim=8, hidden=16, n_heads=2, num_experts=4, num_experts_per_tok=2):
   return TransformerConfig(
@@ -72,20 +72,6 @@ class TestMoEFeedForward(unittest.TestCase):
     expected = (Tensor([1.0]).silu().item() + Tensor([3.0]).silu().item()) / 2
     np.testing.assert_allclose(out.numpy()[0, 0, 0], expected, rtol=1e-2)
 
-  def test_kimi_correction_bias_affects_route_weights(self):
-    dim, hidden, n_heads, num_experts, k = 8, 16, 2, 4, 2
-    config = replace(_moe_config(dim, hidden, n_heads, num_experts, k), norm_topk_prob=True, expert_bias=True)
-    block = TransformerBlock(config)
-    block.ffn_gate_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) * (i + 1) for i in range(num_experts)])
-    block.ffn_up_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) for _ in range(num_experts)])
-    block.ffn_down_exps.weight = Tensor.stack(*[Tensor.eye(dim, hidden) for _ in range(num_experts)])
-    block.ffn_gate_inp.weight = Tensor.zeros(num_experts, dim)
-    block.exp_probs_b["bias"] = Tensor([0.2, 0.1, 0.0, -0.1])
-
-    out = block._feed_forward(Tensor.ones(1, 1, dim))
-    expected = (Tensor([1.0]).silu().item() * 0.7 + Tensor([2.0]).silu().item() * 0.6) / 1.3
-    np.testing.assert_allclose(out.numpy()[0, 0, 0], expected, rtol=1e-2)
-
   def test_moe_feed_forward_shared_expert(self):
     dim, hidden, n_heads = 8, 16, 2
     num_experts, k = 4, 2
@@ -110,5 +96,32 @@ class TestMoEFeedForward(unittest.TestCase):
     expected = moe_expected + shared_expected
     np.testing.assert_allclose(out.numpy(), expected, rtol=1e-2)
 
+  def test_moe_feed_forward_gating_funcs(self):
+    dim, hidden, n_heads = 8, 16, 2
+    num_experts, k = 4, 2
+    logits = np.array([4.0, 3.0, 0.0, -1.0], dtype=np.float32)
+    def softmax(x):
+      probs = np.exp(x - x.max())
+      return probs / probs.sum()
+    for gating_func in ExpertGating:
+      for norm_topk_prob in (False, True):
+        block = TransformerBlock(replace(_moe_config(dim, hidden, n_heads, num_experts, k),
+                                         expert_gating_func=gating_func, norm_topk_prob=norm_topk_prob))
+        block.ffn_gate_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) for _ in range(num_experts)])
+        block.ffn_up_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) * (i + 1) for i in range(num_experts)])
+        block.ffn_down_exps.weight = Tensor.stack(*[Tensor.eye(dim, hidden) for _ in range(num_experts)])
+        block.ffn_gate_inp.weight = Tensor((logits / dim)[None, :].repeat(dim, 0).T)
+        out = block._feed_forward(Tensor.ones(1, 1, dim)).numpy()[0, 0, 0]
+
+        if gating_func == ExpertGating.SOFTMAX: selection_scores = softmax(logits)
+        elif gating_func == ExpertGating.SIGMOID: selection_scores = 1 / (1 + np.exp(-logits))
+        elif gating_func == ExpertGating.SOFTMAX_WEIGHT: selection_scores = logits
+        else: selection_scores = np.sqrt(np.logaddexp(0, logits))
+        sel = np.argsort(selection_scores)[-k:]
+        weights = softmax(logits[sel]) if gating_func == ExpertGating.SOFTMAX_WEIGHT else selection_scores[sel]
+        if norm_topk_prob: weights /= weights.sum()
+        expected = (weights * (sel + 1)).sum() / (1 + np.exp(-1))
+        np.testing.assert_allclose(out, expected, rtol=1e-3)
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_llm_quant.py

@@ -1,68 +0,0 @@
-import unittest
-import numpy as np
-from tinygrad import Tensor, dtypes
-from tinygrad.llm.quant import MXFP4_VALUES, dequantize_mxfp4, quantize_dequantize_mxfp8, quantize_mxfp4, quantize_mxfp4_cpu
-
-class TestMXFormats(unittest.TestCase):
-  def test_mxfp4_known_codes_and_scale(self):
-    values = np.array(MXFP4_VALUES * 2, dtype=np.float32)
-    packed, scale = quantize_mxfp4(Tensor(values))
-    # Positive and negative zero are numerically identical, so nearest-value encoding canonicalizes to +0.
-    np.testing.assert_array_equal(packed.numpy(), np.array([0x10, 0x32, 0x54, 0x76, 0x90, 0xba, 0xdc, 0xfe] * 2, dtype=np.uint8))
-    np.testing.assert_array_equal(scale.numpy(), np.array([127], dtype=np.uint8))
-    np.testing.assert_array_equal(dequantize_mxfp4(packed, scale, dtypes.float32).numpy(), values)
-
-  def test_mxfp4_block_scales_and_zero(self):
-    x = Tensor(np.array([0.0]*32 + [12.0, -12.0] + [0.0]*30, dtype=np.float32))
-    packed, scale = quantize_mxfp4(x)
-    np.testing.assert_array_equal(scale.numpy(), np.array([127, 128], dtype=np.uint8))
-    np.testing.assert_allclose(dequantize_mxfp4(packed, scale, dtypes.float32).numpy(), x.numpy())
-
-  def test_mxfp4_scale_rounds_amax_over_format_max(self):
-    # OCP E8M0 scale selection rounds log2(amax / 6), rather than flooring the
-    # input exponent. At this boundary the two rules differ by a factor of two.
-    x = Tensor(np.array([8.0] + [0.0]*31, dtype=np.float32))
-    packed, scale = quantize_mxfp4(x)
-    np.testing.assert_array_equal(scale.numpy(), np.array([127], dtype=np.uint8))
-    self.assertEqual(dequantize_mxfp4(packed, scale, dtypes.float32).numpy()[0], 6.0)
-
-  def test_mxfp4_cpu_converter_matches_tensor_path(self):
-    x = Tensor(np.linspace(-13, 13, 64*32, dtype=np.float32).reshape(64, 32))
-    packed, scale = quantize_mxfp4(x)
-    cpu_packed, cpu_scale = quantize_mxfp4_cpu(x)
-    np.testing.assert_array_equal(cpu_packed.numpy(), packed.numpy())
-    np.testing.assert_array_equal(cpu_scale.numpy(), scale.numpy())
-
-  def test_mxfp4_midpoints_round_to_even(self):
-    midpoints = np.array([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0], dtype=np.float32)
-    x = Tensor(np.pad(np.concatenate((midpoints, -midpoints)), (0, 18)))
-    packed, scale = quantize_mxfp4(x)
-    expected = np.pad(np.array([0, 1, 1, 2, 2, 4, 4, 0, -1, -1, -2, -2, -4, -4], dtype=np.float32), (0, 18))
-    np.testing.assert_array_equal(dequantize_mxfp4(packed, scale, dtypes.float32).numpy(), expected)
-
-  def test_mxfp8_roundtrip_and_dtype(self):
-    # All E4M3-exact values remain exact after extracting a shared exponent.
-    x = Tensor(np.array(([0.0, 0.5, 1.0, 1.5, 2.0, -3.0, 4.0, -6.0] * 4), dtype=np.float32))
-    out = quantize_dequantize_mxfp8(x)
-    self.assertEqual(out.dtype, dtypes.bfloat16)
-    np.testing.assert_array_equal(out.float().numpy(), x.numpy())
-
-  def test_mxfp8_subnormal_and_rounding(self):
-    x = np.zeros(32, dtype=np.float32)
-    x[:5] = [1.0, 1.0625, 1.07, 2**-9, 2**-10]
-    out = quantize_dequantize_mxfp8(Tensor(x), dtype=dtypes.float32).numpy()
-    # amax / 448 rounds to an E8M0 scale of 2**-9, saturating the largest
-    # values while retaining the E4M3 subnormal quantum for this block.
-    np.testing.assert_array_equal(out[:5], [0.875, 0.875, 0.875, 2**-9, 2**-10])
-
-  def test_mxfp8_uses_full_e4m3_range(self):
-    x = np.zeros(32, dtype=np.float32)
-    x[:4] = [448.0, 416.0, 400.0, -448.0]
-    np.testing.assert_array_equal(quantize_dequantize_mxfp8(Tensor(x), dtype=dtypes.float32).numpy()[:4], [448.0, 416.0, 384.0, -448.0])
-
-  def test_mxfp8_scale_rounds_amax_over_format_max(self):
-    x = np.zeros(32, dtype=np.float32)
-    x[0] = 512.0
-    self.assertEqual(quantize_dequantize_mxfp8(Tensor(x), dtype=dtypes.float32).numpy()[0], 448.0)
-
-if __name__ == "__main__": unittest.main()

test/unit/test_llm_server.py

@@ -1,10 +1,9 @@
 import unittest
-from dataclasses import replace
 from unittest.mock import patch
 from tinygrad import Tensor, UOp
 from tinygrad.schedule import schedule_cache
 from tinygrad.llm.model import Transformer, TransformerConfig
-from tinygrad.llm.serve import StreamRouter, parse_kimi_tool_call
+from tinygrad.llm.serve import StreamRouter
 
 TEST_CONFIG = TransformerConfig(num_blocks=1, dim=64, hidden_dim=128, n_heads=2, n_kv_heads=2,
                            norm_eps=1e-5, vocab_size=100, head_dim=32, rope_theta=10000.0, rope_dim=32, v_head_dim=32, max_context=32)
@@ -14,29 +13,17 @@ V_TOKS = UOp.variable("toks", 1, 32)  # 32 is the default chunk_size in generate
 class TestTransformerGenerate(unittest.TestCase):
   def test_warmup(self):
     model, calls = Transformer(TEST_CONFIG), []
-    def generate(tokens, temperature):
-      calls.append((tokens, temperature))
+    def generate(tokens, **kwargs):
+      calls.append(tokens)
       yield from (1, 2)
     with patch.object(model, "generate", generate): model.warmup()
-    self.assertEqual(calls, [([0], 0.0), ([0], 0.0)])
+    self.assertEqual(calls, [[0], [0]])
 
-  def test_recurrent_warmup_captures_reset_replay(self):
-    model, calls = Transformer(TEST_CONFIG), []
-    model.has_recurrent_block = True
-    state = Tensor.ones(4).realize()
-    model.blk[0]._state_reset_ops = lambda: [state.assign(state.const_like(0))]
-    def generate(tokens, temperature):
-      if calls: model.reset_jit()
-      calls.append((tokens.copy(), temperature))
-      tokens.append(42)
-      yield from (1, 2)
-    with patch.object(model, "generate", generate): model.warmup()
-    prompt = [0] * (TEST_CONFIG.max_context-2)
-    self.assertEqual(calls, [(prompt, 0.0)] * 3 + [(prompt, 1.0)] * 3 + [(prompt + list(range(1, i+1)), 0.0) for i in range(1, 4)])
-    self.assertEqual(model.reset_jit.cnt, 8)
-    cache_size = len(schedule_cache)
-    model.reset_jit()
-    self.assertEqual(len(schedule_cache), cache_size)
+  def test_warmup_then_generate_with_default_chunk(self):
+    # warmup must not capture JIT graphs that generate()'s default chunk_size then rejects
+    model = Transformer(TEST_CONFIG)
+    model.warmup()
+    self.assertIsInstance(next(model.generate([5, 6, 7, 8])), int)
 
   def test_first_recurrent_generate_before_state_init(self):
     model = Transformer(TEST_CONFIG)
@@ -44,17 +31,6 @@ class TestTransformerGenerate(unittest.TestCase):
     with patch.object(Transformer, '__call__', return_value=Tensor([[42]])):
       self.assertEqual(next(model.generate([0])), 42)
 
-  def test_recurrent_prefill_tail_uses_rollout_shape(self):
-    model = Transformer(TEST_CONFIG)
-    model.has_recurrent_block = True
-    model.config = replace(model.config, recurrent_prefill_chunked=True)
-    calls = []
-    def mock_call(self, tokens, start_pos, temperature, **kwargs):
-      calls.append(tokens.shape)
-      return Tensor([[42]])
-    with patch.object(Transformer, '__call__', mock_call): next(model.generate([1, 2, 3, 4, 5, 6], chunk_size=4))
-    self.assertEqual(calls, [(1, 4), (1, 1), (1, 1)])
-
   def test_recurrent_live_state_reuse(self):
     model = Transformer(TEST_CONFIG)
     model.has_recurrent_block = True
@@ -68,37 +44,20 @@ class TestTransformerGenerate(unittest.TestCase):
       next(model.generate([1, 2, 3, 4, 5, 42, 10]))
     self.assertEqual(calls, [((1, 1), V_START_POS.bind(5)), ((1, 1), V_START_POS.bind(6))])
 
-  def test_recurrent_prompt_snapshot_reuse(self):
-    model = Transformer(TEST_CONFIG)
-    model.has_recurrent_block = True
-    state, calls = Tensor.ones(4).realize(), []
-    def mock_call(self, tokens, start_pos, temperature, **kwargs):
+  def test_recurrent_divergent_prompt_restarts(self):
+    model, calls = Transformer(TEST_CONFIG), []
+    model.has_recurrent_block, model._cached_tokens = True, [1, 2, 9]
+    def mock_call(self, tokens, start_pos, temperature):
       calls.append(start_pos)
       return Tensor([[42]])
-    with patch.object(model, "_state_tensors", return_value=[state]), patch.object(model.blk[0], "_reusable_prefix_len", return_value=0), \
-         patch.object(Transformer, '__call__', mock_call):
-      next(model.generate([1, 2, 3]))
-      state.assign(state.const_like(5)).realize()
-      model._cached_tokens = [1, 2, 3, 9, 9]
-      calls.clear()
-      self.assertEqual(model.get_start_pos([1, 2, 3, 7, 8]), 3)
-      next(model.generate([1, 2, 3, 7, 8]))
-    self.assertEqual(calls, [V_START_POS.bind(3), V_START_POS.bind(4)])
-    self.assertEqual(state.tolist(), [1.0] * 4)
+    with patch.object(Transformer, '__call__', mock_call): next(model.generate([1, 2, 10, 11]))
+    self.assertEqual(calls[0], V_START_POS.bind(0))
 
   def test_template_starts_reasoning(self):
     router = StreamRouter(reasoning=True)
     self.assertEqual(list(router.route("reasoning</think>answer")),
                      [("reasoning_content", "reasoning"), ("content", "answer")])
 
-  def test_kimi_tool_call_stream(self):
-    router = StreamRouter()
-    self.assertEqual(list(router.route("before<|tool_calls_section_beg")), [("content", "before")])
-    self.assertEqual(list(router.route("in|><|tool_call_begin|>functions.read:0<|tool_call_argument_begin|>"
-                                      '{"path":"/tmp/x"}<|tool_call_end|><|tool_calls_section_end|>')), [])
-    self.assertEqual(parse_kimi_tool_call("functions.read:0<|tool_call_argument_begin|>{\"path\":\"/tmp/x\"}"),
-                     ("read", {"path":"/tmp/x"}))
-
   def test_kv_cache_reuse(self):
     """Test that generate reuses the KV cache when tokens extend the cached prefix."""
     model = Transformer(TEST_CONFIG)
@@ -234,6 +193,12 @@ class TestTransformerGenerate(unittest.TestCase):
     # with temperature=2.0, we should see at least 2 distinct outputs across 5 runs
     self.assertGreater(len(runs), 1, "high temperature should produce varied outputs")
 
+  def test_recurrent_temperature_high_produces_variety(self):
+    model = Transformer(TEST_CONFIG)
+    model.has_recurrent_block = True
+    outputs = {model.forward(Tensor([[1]]), 0, Tensor([2.0])).item() for _ in range(5)}
+    self.assertGreater(len(outputs), 1)
+
   def test_temperature_passed_to_forward(self):
     """Temperature from generate should be passed through to __call__."""
     model = Transformer(TEST_CONFIG)

test/unit/test_metal_graph.py

@@ -1,5 +1,4 @@
 import unittest
-from unittest.mock import MagicMock
 from tinygrad import Device
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.dtype import dtypes
@@ -11,36 +10,27 @@ class TestMetalGraph(unittest.TestCase):
     self.MetalGraph = MetalGraph
     self.dev = Device[Device.DEFAULT]
 
-  def metal_buf(self, offset):
-    buf = MagicMock()
-    if offset > 0:
-      buf.op = Ops.SLICE
-      src = MagicMock()
-      src.dtype = dtypes.uint8
-      buf.src = (src, UOp.const(offset))
-      buf.dtype = dtypes.uint8
-    else:
-      buf.op = Ops.BUFFER
-    buf.device = Device.DEFAULT
-    return buf
+  def metal_buf(self, offset, bitcast=False):
+    size = 4 if bitcast else 1
+    buf = UOp.new_buffer(Device.DEFAULT, offset+size, dtypes.uint8)
+    if offset: buf = buf[offset:offset+size]
+    return buf.bitcast(dtypes.float32) if bitcast else buf
 
-  def call(self, *bufs):
-    c = MagicMock()
-    c.src = (MagicMock(op=Ops.PROGRAM),) + tuple(bufs)
-    return c
+  def supports_uop(self, *bufs):
+    return self.MetalGraph.supports_uop([self.dev], UOp(Ops.PROGRAM, src=(UOp.sink(),)).call(*bufs))
 
   def test_supports_uop_normal_offset(self):
-    assert self.MetalGraph.supports_uop([self.dev], self.call(self.metal_buf(0), self.metal_buf(100), self.metal_buf(0xFFFFFFFF))) is True
+    assert self.supports_uop(self.metal_buf(0), self.metal_buf(100), self.metal_buf(0xFFFFFFFF)) is True
 
   def test_supports_uop_overflow_offset(self):
-    assert self.MetalGraph.supports_uop([self.dev], self.call(self.metal_buf(0), self.metal_buf(0x100000000))) is False
+    assert self.supports_uop(self.metal_buf(0), self.metal_buf(0x100000000)) is False
 
-  def test_supports_uop_nonmetal_buf(self):
-    # non-SLICE ops should not be checked for offset
-    buf = MagicMock()
-    buf.op = Ops.BUFFER
-    buf.device = Device.DEFAULT
-    self.MetalGraph.supports_uop([self.dev], self.call(buf))
+  def test_supports_uop_non_view_buf(self):
+    assert self.supports_uop(self.metal_buf(0)) is True
+
+  def test_supports_uop_bitcast(self):
+    assert self.supports_uop(self.metal_buf(0xFFFFFFFF, bitcast=True)) is True
+    assert self.supports_uop(self.metal_buf(0x100000000, bitcast=True)) is False
 
 if __name__ == "__main__":
   unittest.main()

test/unit/test_multitensor.py

@@ -390,6 +390,12 @@ class TestMultiTensor(unittest.TestCase):
     self.assertEqual(out.shape, (rows, 8))
     np.testing.assert_equal(out[:3].to(Device.DEFAULT).numpy(), np.ones((3, 8)))
 
+  def test_symbolic_broadcast_consumed(self):
+    rows = Variable("rows", 1, 4).bind(3)
+    out = (Tensor.ones(rows).to(devices_2) + 1).realize()
+    self.assertEqual(out.shape, (rows,))
+    np.testing.assert_equal(out[:3].to(Device.DEFAULT).numpy(), np.full(3, 2))
+
   def test_multitensor_jit_in_list(self):
     # test MULTI tensor inside a list container - exercises the container unpacking + MULTI unpacking
     @TinyJit

test/unit/test_schedule_cache.py

@@ -1,11 +1,16 @@
 import unittest
 import functools
-from tinygrad import Tensor, Variable, UOp
+from tinygrad import Tensor, Variable, UOp, function
 from tinygrad.uop.ops import KernelInfo
 from tinygrad.schedule import schedule_cache
 
-def custom_set0_kernel(A:UOp, num:int) -> UOp:
-  return A[0].set(num).sink(arg=KernelInfo(f"custom_set0_{num}"))
+def custom_add_kernel(A:UOp, B:UOp, num:int=0) -> UOp:
+  return A[0].set(B[0] + num).sink(arg=KernelInfo(f"custom_add_{num}"))
+
+def custom_add_backward(grad_output:UOp, _) -> tuple[None, UOp]:
+  grad = Tensor.invalids(*grad_output.shape, dtype=grad_output.dtype, device=grad_output.device)
+  grad = Tensor.custom_kernel(grad, Tensor(grad_output, device=grad_output.device), fxn=functools.partial(custom_add_kernel, num=0))[0]
+  return None, grad.uop
 
 class TestScheduleCache(unittest.TestCase):
   def test_bound_variable_reuses_cache(self):
@@ -25,27 +30,27 @@ class TestScheduleCache(unittest.TestCase):
 
   def test_custom_kernel(self):
     for i in range(4):
-      a = Tensor.empty(1)
-      a = Tensor.custom_kernel(a, fxn=functools.partial(custom_set0_kernel, num=i))[0]
+      a, b = Tensor.empty(1), Tensor.ones(1)
+      a = Tensor.custom_kernel(a, b, fxn=functools.partial(custom_add_kernel, num=i))[0]
       a.realize()
-      self.assertEqual(a.item(), i)
+      self.assertEqual(a.item(), i+1)
 
   def test_same_custom_function_reuses_cache(self):
     schedule_cache.clear()
-    fxn = functools.partial(custom_set0_kernel, num=10)
+    fxn = functools.partial(custom_add_kernel, num=10)
 
     # first run
-    a = Tensor.empty(1)
-    a = Tensor.custom_kernel(a, fxn=fxn)[0]
+    a, x = Tensor.empty(1), Tensor.ones(1)
+    a = Tensor.custom_kernel(a, x, fxn=fxn)[0]
     a.realize()
-    self.assertEqual(a.item(), 10)
+    self.assertEqual(a.item(), 11)
     cache_size_after_first = len(schedule_cache)
 
     # second run with same function should reuse cache
-    b = Tensor.empty(1)
-    b = Tensor.custom_kernel(b, fxn=fxn)[0]
+    b, x = Tensor.empty(1), Tensor.ones(1)
+    b = Tensor.custom_kernel(b, x, fxn=fxn)[0]
     b.realize()
-    self.assertEqual(b.item(), 10)
+    self.assertEqual(b.item(), 11)
     self.assertEqual(len(schedule_cache), cache_size_after_first)
 
   def test_simple(self):
@@ -65,5 +70,29 @@ class TestScheduleCache(unittest.TestCase):
       print(num)
     self.assertEqual(len(schedule_cache), start_len_schedule_cache)
 
+  def test_simple_precompile(self):
+    @function(precompile=True, precompile_backward=True)
+    def f(x:Tensor) -> Tensor:
+      out = Tensor.invalids(*x.shape, dtype=x.dtype, device=x.device)
+      out = Tensor.custom_kernel(out, x, fxn=functools.partial(custom_add_kernel, num=10), grad_fxn=custom_add_backward)[0]
+      return out + x
+
+    # warmup
+    x = Tensor.ones(1).realize()
+    out = f(x)
+    out.backward(x)
+    self.assertEqual(out.item(), 12)
+    self.assertEqual(x.grad.item(), 2)
+
+    # use the cache next time function is called
+    start_len_schedule_cache = len(schedule_cache)
+    for _ in range(3):
+      x = Tensor.ones(1).realize()
+      out = f(x)
+      out.backward(x)
+      self.assertEqual(out.item(), 12)
+      self.assertEqual(x.grad.item(), 2)
+    self.assertEqual(len(schedule_cache), start_len_schedule_cache)
+
 if __name__ == "__main__":
   unittest.main()

tinygrad/codegen/__init__.py

@@ -12,7 +12,7 @@ from tinygrad.dtype import dtypes, AddrSpace
 
 # import all pattern matchers here
 from tinygrad.codegen.gpudims import pm_add_gpudims
-from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_fold_cast_const, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid
+from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid
 from tinygrad.uop.movement import mop_cleanup
 from tinygrad.codegen.decomp.dtype import pm_dtype_decomps
 from tinygrad.codegen.decomp.op import get_late_rewrite_patterns, get_simplifying_rewrite_patterns
@@ -153,8 +153,8 @@ devectorizer2 = mop_cleanup+pm_mops+PatternMatcher([
   # unpack WMMA
   (UPat(Ops.WMMA, name="u"), do_stack_wmma),
   # stacked INDEX is many INDEX
-  (UPat(Ops.INDEX, src=(UPat((Ops.PARAM, Ops.BUFFER), name="b"), UPat(Ops.STACK, name="s"))),
-   lambda b,s: UOp.stack(*[b.index(u) for u in s.src])),
+  (UPat(Ops.INDEX, src=(UPat((Ops.PARAM, Ops.BUFFER), name="b"), UPat(Ops.STACK, name="s")), name="x"),
+   lambda b,s,x: UOp.stack(*[x.replace(src=(b,u)) for u in s.src])),
   # INDEX into RESHAPE moves the RESHAPE
   (UPat(Ops.INDEX, src=(UPat((Ops.PARAM, Ops.BUFFER), name="b"), UPat(Ops.RESHAPE, name="s"))),
    lambda b,s: b.index(s.src[0]).reshape(s.shape)),
@@ -281,6 +281,10 @@ pm_implicit_barriers = PatternMatcher([
   (UPat(Ops.END, name="end"), add_war_barrier),
 ])
 
+pm_casted_consts = PatternMatcher([
+  (UPat(Ops.CONST, dtypes.all, name="c"), lambda c: UOp.cconst(c.val, c.dtype)),
+])
+
 def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
   if DEBUG >= 5: print(pyrender(ast))
@@ -301,7 +305,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
     sink = graph_rewrite(sink, pm_split_ranges+pm_flatten_range, ctx={}, name="split ranges")
 
     # symbolic (NOTE: this is a requirement for pm_simplify_ranges to be correct)
-    sink = graph_rewrite(sink, sym+pm_fold_cast_const+pm_flatten_range, name="initial symbolic")
+    sink = graph_rewrite(sink, sym+pm_flatten_range, name="initial symbolic")
 
     # optimize (schedule) the AST
     sink = graph_rewrite(sink, pm_flatten_range+pm_simplify_ranges, ctx={}, name="simplify ranges")
@@ -346,7 +350,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
 
   # lower index dtype
   # NOTE: we need indexing_simplify to remove the cast to long using the Invalid
-  sink = graph_rewrite(sink, symbolic_simple+pm_fold_cast_const+pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes")
+  sink = graph_rewrite(sink, symbolic_simple+pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes")
 
   # final symbolic before decomp
   sink = graph_rewrite(sink, symbolic, name="final symbolic")
@@ -357,7 +361,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
 
   # floordiv+mod / dtype decomp (early)
   supported_ops = tuple(ren.code_for_op.keys())
-  pm_decomp = symbolic_simple+pm_fold_cast_const+get_simplifying_rewrite_patterns(supported_ops)
+  pm_decomp = symbolic_simple+get_simplifying_rewrite_patterns(supported_ops)
   sink = graph_rewrite(sink, pm_decomp, name="early decompositions")
 
   # late decomps + move gates from unrenderable INVALID where
@@ -383,6 +387,10 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   num_params = len([x for x in sink.toposort() if x.op is Ops.PARAM and x.arg.slot != -1])
   sink = graph_rewrite(sink, pm_number_params, ctx=[num_params], name="number params with -1", walk=True)
 
+  # spell every literal as a casted const CAST(dt, CONST(value))
+  # TODO: remove once consts are always weak
+  sink = graph_rewrite(sink, pm_casted_consts, name="casted consts", walk=True)
+
   if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Output AST")
   if SPEC: type_verify(sink, spec_program)
 

tinygrad/codegen/decomp/dtype.py

@@ -33,7 +33,8 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
       return (lo:=uops[0].cast(l2i_dt[dt])), (uops[0] / 2**32).cast(l2i_dt[dt]) - ((uops[0] < 0) & lo.ne(0))
     case Ops.CAST if dt in dtypes.floats:
       small = (a1.eq(0) & (a0 >= 0)) | (a1.eq(-1) & (a0 < 0))
-      return small.where(a0.cast(dt), ((a1.cast(dtypes.float32) * (2**32)) + a0.bitcast(dtypes.uint).cast(dtypes.float32)).cast(dt))
+      cdt = dt if dt == dtypes.float64 else dtypes.float32
+      return small.where(a0.cast(dt), ((a1.cast(cdt) * (2**32)) + a0.bitcast(dtypes.uint).cast(cdt)).cast(dt))
     case Ops.CAST: return a0.bitcast(dtypes.uint).cast(dt)
     case Ops.BITCAST: return a0.bitcast(dt), a1.bitcast(dt)
     case Ops.SHL:

tinygrad/codegen/gpudims.py

@@ -57,7 +57,7 @@ def add_gpudims(ctx:Renderer, s:UOp):
 
   # get the idxs
   ki: KernelInfo = s.arg
-  if ctx.has_threads: idxs = [UOp.variable("core_id", 0, int(global_shape[0])-1, dtypes.int).cast(dtypes.weakint)]
+  if ctx.has_threads: idxs = [UOp.variable("core_id", 0, int(global_shape[0])-1, dtypes.int, param=True).cast(dtypes.weakint)]
   elif ki.dont_use_locals:
     assert not local_dims, "can't use locals if there's no local dims"
     idxs = get_grouped_dims("idx", global_shape, ctx.global_max, reverse=True)
@@ -89,7 +89,7 @@ def add_gpudims(ctx:Renderer, s:UOp):
 
 pm_device_to_var = PatternMatcher([
   # the DEVICE axis is not a program axis, it's bound per device at launch. lower it to the _device_num variable (like SPECIAL for devices)
-  (UPat(Ops.RANGE, name="r"), lambda r: UOp.variable("_device_num", 0, r.vmax, dtype=r.dtype) if r.arg[-1] is AxisType.DEVICE else None),
+  (UPat(Ops.RANGE, name="r"), lambda r: UOp.variable("_device_num", 0, r.vmax, dtype=r.dtype, param=True) if r.arg[-1] is AxisType.DEVICE else None),
   # ENDs that closed a DEVICE range no longer close it
   (UPat(Ops.END, name="e"), lambda e: e.replace(src=(e.src[0],)+tuple(s for s in e.src[1:] if s.op is not Ops.PARAM))
    if any(s.op is Ops.PARAM and s.arg.name == '_device_num' for s in e.src[1:]) else None),

tinygrad/codegen/late/coalesce.py

@@ -26,7 +26,7 @@ def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]:
     # check if idx is out of bound when X is on the wrong side of the bound: X in [c+1, vmax] or [vmin, c-1]
     lo, hi = (c + 1, X.vmax) if is_upper_bound else (X.vmin, c - 1)
     if lo <= hi:
-      fake = UOp.variable(f"fake{i}", lo, hi, X.dtype)
+      fake = UOp.variable(f"fake{i}", lo, hi, X.dtype, param=True)
       subs = [{X: fake}]
       # idx may not have X itself, so also substitute a term of X: v -> fake - (X - v)
       terms = list(X.split_uop(Ops.ADD))
@@ -149,7 +149,7 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
         grp = full_grp[:length]
         # NOTE: we apply the valid again after we determine the length
         offset = offset.valid(valid) if valid is not None else offset
-        idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(len(grp)))) if len(grp) > 1 else buf.index(offset)
+        idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(len(grp)))) if len(grp) > 1 else buf.index(offset, dtype=offsets[grp[0]][0].src[0].dtype)
         if op == Ops.STORE:
           datas = []
           for i,g in enumerate(grp):

tinygrad/codegen/late/gater.py

@@ -3,7 +3,7 @@ from tinygrad.uop.ops import PatternMatcher, UPat, Ops
 from tinygrad.dtype import Invalid, dtypes
 
 def move_where_load(gate, l, a, w):
-  return l.replace(src=(l.src[0], l.vconst_like(0) if a.is_invalid else
+  return l.replace(src=(l.src[0], l.vconst_like(0) if a.is_invalid else l.const_like(a.val) if a.op is Ops.CONST else
     a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype), l.src[2])).cast(w.dtype)
 
 pm_move_gates_from_index = PatternMatcher([

tinygrad/codegen/late/regalloc.py

@@ -4,7 +4,7 @@ from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
 from tinygrad.renderer.isa import ISARenderer, Register, greg
 from tinygrad.dtype import dtypes
 
-PSEUDO_OPS = {Ops.CONST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP, Ops.STACK}
+PSEUDO_OPS = {Ops.CONST, Ops.CAST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP, Ops.STACK}
 
 class LinearScanRegallocContext:
   # returns the uop that defines the virtual register
@@ -52,7 +52,7 @@ class LinearScanRegallocContext:
         # the value of a BUFFER is its 64bit address, XMM registers need 16 bytes
         sz = 16 if v.cons[0].size == 16 else (8 if self.vdef(v).op is Ops.BUFFER else self.vdef(v).dtype.itemsize)
         offset = self.stack_size + (sz - self.stack_size % sz) % sz
-        self.spills[v] = UOp.const(offset, dtypes.int32)
+        self.spills[v] = UOp.cconst(offset, dtypes.int32)
         self.stack_size = offset + sz
       r = alloc(cons if cons is not None else v.cons, i)
       self.insert_before.setdefault(i, []).append((v, r))
@@ -84,7 +84,7 @@ class LinearScanRegallocContext:
 
       # allocate stack array
       if u.op is Ops.BUFFER:
-        self.locals[u] = UOp.const(self.stack_size, dtypes.int32)
+        self.locals[u] = UOp.cconst(self.stack_size, dtypes.int32)
         self.stack_size += u.max_numel() * u.dtype.itemsize
 
       # loop prologue, avoid loading inside the loop
@@ -125,7 +125,7 @@ def regalloc_rewrite(ctx:LinearScanRegallocContext, x:UOp):
   # alloc/dealloc stack
   if ctx.stack_size > 0:
     sp = ctx.ren.stack_pointer()
-    offset = UOp.const(ctx.stack_size, sp.dtype)
+    offset = UOp.cconst(ctx.stack_size, sp.dtype)
     if i == 0: before = [ctx.ren.isel_matcher.rewrite(UOp(Ops.SUB, src=(sp, offset), tag=sp.tag))] + before
     elif i == len(ctx.uops) - 2: before += [ctx.ren.isel_matcher.rewrite(UOp(Ops.ADD, src=(sp, offset), tag=sp.tag))]
 

tinygrad/codegen/opt/tc.py

@@ -1,6 +1,7 @@
 import math, functools
 from dataclasses import dataclass
 from tinygrad.dtype import DType, dtypes
+from tinygrad.uop.ops import PatternMatcher, UOp, UPat, Ops
 
 @dataclass(frozen=True)
 class TensorCore: # D = A * B + C, A is (M x K), B is (K x N), C and D are (M x N)
@@ -135,6 +136,42 @@ amd_cdna4 = amd_cdna_1616128 + amd_cdna_161632 + amd_cdna_161616
 
 def get_amd(arch): return {"gfx942": amd_cdna3, "gfx950": amd_cdna4, "gfx1200": amd_rdna4, "gfx1201": amd_rdna4}.get(arch, amd_rdna3)
 
+pm_validate_wmma_rdna3 = PatternMatcher([
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.int32), lambda x: x.replace(
+    src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
+    if x.src[0].dtype == dtypes.int8 and x.src[0].max_numel() == 16 else None),
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(x.replace(
+      src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(UOp.const(j//2, dtypes.int16))
+      if j%2 == 0 else UOp.const(0.0, x.src[2].dtype)
+      for j in range(x.max_numel()*2)))),
+      arg=(*x.arg[:4], None)).index(UOp.const(i*2, dtypes.int16))
+      for i in range(x.max_numel()))) if x.max_numel() == 8 else None),
+  (UPat(Ops.WMMA, name="x"), lambda x: x.replace(
+    src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
+    if x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 16 else None),
+])
+
+pm_validate_wmma_rdna4 = PatternMatcher([
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16), lambda x: x.replace(
+    dtype=dtypes.uint16,
+    src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2].bitcast(dtypes.uint16)))
+      .bitcast(dtypes.bfloat16) if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None),
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+    lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
+    if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None)
+])
+
+pm_validate_wmma_cdna = PatternMatcher([
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+    lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
+    if x.arg[0][2] == 128 and x.src[0].dtype.itemsize <= 8 else None),
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+    lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
+    if x.max_numel() == 4 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 4 else None),
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+    lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2]))
+    if x.max_numel() == 4 and x.src[0].dtype in dtypes.fp8_ocp and x.src[0].max_numel() == 8 else None),
+])
 # ***** Apple Metal *****
 
 metal = [TensorCore(dims=(8,8,8), threads=32, elements_per_thread=(2,2,2), dtype_in=di, dtype_out=do,

tinygrad/codegen/simplify.py

@@ -1,7 +1,7 @@
 import itertools
 from typing import Callable
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start, AxisType
-from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const, invalid_gate
+from tinygrad.uop.symbolic import symbolic, invalid_gate
 from tinygrad.helpers import partition
 from tinygrad.dtype import dtypes
 
@@ -32,7 +32,7 @@ def simplify_merge_adjacent(u:UOp) -> UOp|None:
         s0, s1 = r0.src[0], r1.src[0]
         # do the merge
         new_range = r0.replace(src=(s0*s1,))
-        nidx = graph_rewrite(u, _substitute+symbolic+pm_fold_cast_const+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
+        nidx = graph_rewrite(u, _substitute+symbolic+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
                              name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}")
 
         # check if it simplifies
@@ -137,7 +137,7 @@ def reduce_collapse(red:UOp, u:UOp, pm:PatternMatcher=pm_reduce_collapse) -> UOp
     for u in included:
       for s in u.src:
         if s in included or s in replaces or s.op in {Ops.CONST, Ops.PARAM, Ops.BUFFER}: continue
-        replaces[s] = UOp.variable(f'in{len(replaces)}', s.vmin, s.vmax, s.dtype)
+        replaces[s] = UOp.variable(f'in{len(replaces)}', s.vmin, s.vmax, s.dtype, param=True)
     collapse_fxn = u.substitute(replaces).reduce(r, arg=Ops.ADD)
     sink = graph_rewrite(collapse_fxn, pm, name="reduce_collapse")
     if not no_range(sink): return None

tinygrad/device.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 from dataclasses import dataclass, replace
 from collections import defaultdict
-from typing import Any, Generic, TypeVar, Iterator, Generator, Self, TYPE_CHECKING
-import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal
+from typing import Any, Callable, Generic, TypeVar, Iterator, Generator, Self, TYPE_CHECKING
+import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal, subprocess, struct
 from tinygrad.helpers import LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, PROFILE, temp, colored
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing
 from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize, Target, unwrap, round_up
@@ -103,7 +103,7 @@ class Buffer:
   def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None, options:BufferSpec|None=None,
                initial_value:bytes|pickle.PickleBuffer|None=None, uop_refcount=0, base:Buffer|None=None, offset:int=0, preallocate=False):
     assert isinstance(dtype, DType)
-    self.device, self.size, self.dtype, self.options, self.offset, self.allocated_views = device, size, dtype, options, offset, 0
+    self.device, self.size, self.dtype, self.options, self.offset, self.allocated_views = Device.canonicalize(device), size, dtype, options, offset, 0
     self._bufs: dict[str, Any] = {}
     if base is None:
       assert offset == 0, "base buffers can't have offset"
@@ -116,7 +116,7 @@ class Buffer:
         if isinstance(initial_value, pickle.PickleBuffer): initial_value.release()
     else:
       assert base._base is None, "base can't have a base"
-      assert device == base.device, "base must have the same device"
+      assert self.device == base.device, "base must have the same device"
       self._base = base
     if preallocate: self.allocate()
   @property
@@ -133,7 +133,7 @@ class Buffer:
   # check if the underlying buffer is allocated, possibly from the base object
   def is_allocated(self) -> bool: return self.base.is_allocated() if self._base is not None else self.device in self._bufs
   def get_buf(self, device: str) -> Any:
-    if device not in self._bufs:
+    if device not in self._bufs and (device:=Device.canonicalize(device)) not in self._bufs:
       allocator = Device[device].allocator
       if device == self.device: self.ensure_allocated()
       elif self._base is not None: self._bufs[device] = allocator._offset(self._base.get_buf(device), self.nbytes, self.offset)
@@ -310,6 +310,14 @@ class Compiler:
       if self.cachekey is not None: diskcache_put(self.cachekey, src, lib)
     return lib
   def disassemble(self, lib:bytes): pass
+  def server(self, cmd:str, arch:str, *args) -> subprocess.Popen:
+    argv = f"{cmd} {pathlib.Path(__file__).parent}/runtime/support/compileserver.py {type(self).__module__}:{type(self).__name__} {arch}"
+    return subprocess.Popen(argv.split() + [str(a) for a in args], stdout=subprocess.PIPE, stdin=subprocess.PIPE, bufsize=0)
+  def compile_server(self, src:str, proc:subprocess.Popen) -> bytes:
+    unwrap(proc.stdin).write(struct.pack("I", len(src.encode())) + src.encode())
+    if (lib:=unwrap(proc.stdout).read(struct.unpack("I", unwrap(proc.stdout).read(4))[0])): return lib
+    raise CompileError("Compilation Error")
+
 
 @dataclass
 class TinyELF:
@@ -331,15 +339,18 @@ class Program(Generic[DeviceType]):
                wait=False) -> float|None: pass
 
 class Compiled:
+  ifaces:list[Callable] = []
   profile_events:list[ProfileEvent] = [ProfileDeviceEvent("CPU")] # NOTE: CPU is the default device.
 
+  has_copy_queue:bool = True
+
   pm_lower:Any = None
   pm_bufferize:Any = None
 
   def __init__(self, device:str, allocator:Allocator, renderers:list[type[Renderer]], runtime:type[Program[Self]]|None, graph=None, arch=None):
     from tinygrad.renderer import Renderer
     self.device, self.allocator, self.runtime_t, self.graph, self.renderers = device, allocator, runtime, graph, renderers or [Renderer]
-    self.arch = arch
+    self.device_id, self.arch = (int(idx) if ":" in device and (idx:=device.split(":")[1]).isdigit() else 0), arch
     self.cached_renderer:dict[Any, Renderer] = {}
 
   @property
@@ -362,11 +373,21 @@ class Compiled:
     return select_first_inited(select_by_name(self.renderers, self._renderer_name, t.renderer, f"{self.device} has no renderer {t.renderer!r}"),
                                f"No renderer for {self.device} is available", self.cached_renderer, t)
 
+  def _select_iface(self, device:str):
+    self.device_id = int(device.split(":")[1]) if ":" in device else 0
+    assert (v:=getenv(k:=f'{type(self).__name__[:-6].upper()}_IFACE', "")) == "",  \
+      f"{k}={v} is deprecated, use DEV={replace(DEV.target(type(self).__name__[:-6]), interface=v)} instead"
+    t = DEV.target(dev:=type(self).__name__[:-6])
+    filtered = select_by_name(self.ifaces, lambda i: i.__name__[:-5], t.interface, f"{dev} has no interface {t.interface!r}")
+    filtered = [i for i in filtered if t.interface.startswith("MOCK") or not i.__name__[:-5].startswith("MOCK")] # never fallback to mock ifaces
+    return select_first_inited([functools.partial(iface, self, self.device_id) for iface in filtered],
+                               f"No interface for {dev}:{self.device_id} is available")
+
   def count(self) -> int:
     """
     Returns the number of physical accelerators available to the runtime.
     """
-    return 1
+    return self.iface.count if hasattr(self, 'iface') else 1
 
   def synchronize(self):
     """
@@ -384,7 +405,7 @@ class Compiled:
     """
     Called at the end of process lifetime to allow the device to finalize.
     """
-    # override this in your device implementation
+    if hasattr(self, 'iface') and hasattr(self.iface, 'device_fini'): self.iface.device_fini()
 
 if PROFILE:
   @atexit.register
@@ -406,7 +427,7 @@ def enumerate_devices_str() -> Generator[str, None, None]:
     ren_results, iface_results = [], []
     try:
       d = Device[device]
-      for iface in [i for i in getattr(d, 'ifaces', []) if not i.__name__.startswith("MOCK")]:
+      for iface in [i for i in d.ifaces if not i.__name__.startswith("MOCK")]:
         try:
           name = iface.__name__[:-5]
           default_text, count = ("(default)", d.count()) if type(d.iface) is iface else (f"(DEV={name}+{device} to make default)", iface(d, 0).count) # type: ignore