move to SPEC=3

.github/actions/process-replay/action.yml

@@ -5,7 +5,6 @@ runs:
   steps:
     - name: Run process replay tests
       shell: bash
-      if: env.CAPTURE_PROCESS_REPLAY == '1'
       run: |
         export PR_TITLE=$(jq -r .pull_request.title "$GITHUB_EVENT_PATH")
         export CURRENT_SHA=${{ github.event.pull_request && github.event.pull_request.head.sha || github.sha }}

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

@@ -4,13 +4,13 @@ inputs:
   python-version:
     description: 'Python version to use'
     required: false
-    default: '' # if you don't set a version, the native python version will be used
+    default: '3.12'
   key:
     description: 'Key for the python cache'
     required: false
     default: ''  # if you don't set a key, it doesn't cache
   deps:
-    description: 'Extra dependency groups (space separated)'
+    description: 'Extra dependency groups (comma separated)'
     required: false
     default: ''
   pydeps:
@@ -41,33 +41,20 @@ inputs:
     description: "Install LLVM?"
     required: false
     default: 'false'
-  tinydreno:
-    description: "Install tinydreno"
+  mesa:
+    description: "Install mesa"
     required: false
     default: 'false'
-  qemu:
-    description: "Install qemu"
+  tinydreno:
+    description: "Install tinydreno"
     required: false
     default: 'false'
 runs:
   using: "composite"
   steps:
-    - name: Setup environment
-      shell: bash
-      run: |
-        echo "UV_CACHE_DIR=/tmp/.uv-cache" >> "$GITHUB_ENV"
-        echo "OMP_NUM_THREADS=1" >> "$GITHUB_ENV"
-        # no buffers should be over 300MB in CI
-        echo "MAX_BUFFER_SIZE=300000000" >> "$GITHUB_ENV"
-
-    - name: Set up uv
-      uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b
-      with:
-        enable-cache: 'false' # see below for manual caching
-
     - name: Set up Python ${{ inputs.python-version }}
+      id: setup-python
       uses: actions/setup-python@v6
-      if: inputs.python-version != ''
       with:
         python-version: ${{ inputs.python-version }}
 
@@ -76,23 +63,23 @@ runs:
     - name: Cache Python packages (PR)
       if: github.event_name == 'pull_request'
       id: restore-venv-pr
-      uses: actions/cache/restore@v5
+      uses: actions/cache/restore@v4
       with:
-        path: /tmp/.uv-cache
-        key: uv-${{ runner.os }}-${{ runner.arch }}-python-${{ inputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
+        path: ${{ github.workspace }}/.venv
+        key: venv-${{ runner.os }}-${{ runner.arch }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
     - name: Cache Python packages
       if: github.event_name != 'pull_request'
       id: restore-venv
       uses: actions/cache@v5
       with:
-        path: /tmp/.uv-cache
-        key: uv-${{ runner.os }}-${{ runner.arch }}-python-${{ inputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
+        path: ${{ github.workspace }}/.venv
+        key: venv-${{ runner.os }}-${{ runner.arch }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
 
     # **** Caching downloads ****
 
     - name: Cache downloads (PR)
       if: inputs.key != '' && github.event_name == 'pull_request'
-      uses: actions/cache/restore@v5
+      uses: actions/cache/restore@v4
       with:
         path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
         key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
@@ -106,26 +93,34 @@ runs:
     # **** Python deps ****
 
     - name: Install dependencies in venv (with extra)
-      if: inputs.deps != ''
+      if: inputs.deps != '' && steps.restore-venv-pr.outputs.cache-hit != 'true' && steps.restore-venv.outputs.cache-hit != 'true'
       shell: bash
       run: |
-        uv venv .venv
-        DEPS="${{ inputs.deps }}"
-        uv pip install --python .venv -e ".[${DEPS// /,}]" ${{ inputs.pydeps }} --torch-backend cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
+        python -m venv .venv
+        if [[ "$RUNNER_OS" == "Windows" ]]; then
+          source .venv/Scripts/activate
+        else
+          . .venv/bin/activate
+        fi
+        python -m pip install -e ".[${{ inputs.deps }}]" ${{ inputs.pydeps }} --extra-index-url https://download.pytorch.org/whl/cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
     - name: Install dependencies in venv (without extra)
-      if: inputs.deps == ''
+      if: inputs.deps == '' && steps.restore-venv-pr.outputs.cache-hit != 'true' && steps.restore-venv.outputs.cache-hit != 'true'
       shell: bash
       run: |
-        uv venv .venv
-        uv pip install --python .venv -e . ${{ inputs.pydeps }}
-    - name: Prune uv cache
-      if: github.event_name != 'pull_request'
-      shell: bash
-      run: uv cache prune --ci
-    - name: Configure venv
+        python -m venv .venv
+        if [[ "$RUNNER_OS" == "Windows" ]]; then
+          source .venv/Scripts/activate
+        else
+          . .venv/bin/activate
+        fi
+        python -m pip install -e . ${{ inputs.pydeps }}
+    - name: Set up venv environment
       shell: bash
       run: |
         echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
+        echo "OMP_NUM_THREADS=1" >> "$GITHUB_ENV"
+        # no buffers should be over 300MB in CI
+        echo "MAX_BUFFER_SIZE=300000000" >> "$GITHUB_ENV"
         if [[ "$RUNNER_OS" == "Windows" ]]; then
           echo "${{ github.workspace }}/.venv/Scripts" >> "$GITHUB_PATH"
         else
@@ -134,16 +129,20 @@ 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.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
       shell: bash
       run: |
-        sudo mkdir -p /var/cache/apt/archives
         sudo chown -R $USER:$USER /var/cache/apt/archives
 
         echo 'Acquire::GzipIndexes "true";' | sudo tee /etc/apt/apt.conf.d/gzip
         echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
         echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' | sudo tee -a /etc/apt/apt.conf.d/99keep-debs
 
+    - name: Add OpenCL Repo
+      if: inputs.opencl == 'true' && runner.os == 'Linux'
+      shell: bash
+      run: echo "deb [ allow-insecure=yes ] https://apt.repos.intel.com/oneapi all main" | sudo tee /etc/apt/sources.list.d/oneAPI.list
+
     - name: Add AMD Repo (Linux)
       if: inputs.amd == 'true' && runner.os == 'Linux'
       shell: bash
@@ -162,50 +161,54 @@ 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.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
       id: apt-pkgs
       shell: bash
       run: |
         pkgs=""
         # **** OpenCL ****
         if [[ "${{ inputs.opencl }}" == "true" ]]; then
-          pkgs+=" ocl-icd-opencl-dev"
+          pkgs+=" opencl-headers \
+            intel-oneapi-runtime-openmp=2023.2.1-16 intel-oneapi-runtime-compilers-common=2023.2.1-16 intel-oneapi-runtime-compilers=2023.2.1-16 \
+            intel-oneapi-runtime-dpcpp-sycl-opencl-cpu=2023.2.1-16 intel-oneapi-runtime-tbb-common=2021.10.0-49541 \
+            intel-oneapi-runtime-tbb=2021.10.0-49541 intel-oneapi-runtime-opencl=2023.2.1-16"
         fi
         # **** AMD ****
         if [[ "${{ inputs.amd }}" == "true" ]]; then
-          pkgs+=" comgr"
+          pkgs+=" hsa-rocr comgr hsa-rocr-dev liburing-dev libibverbs-dev libc6-dev"
+        fi
+        # **** CUDA ****
+        if [[ "${{ inputs.cuda }}" == "true" ]]; then
+          pkgs+=" git g++ cmake ninja-build llvm-15-dev zlib1g-dev libglew-dev \
+            flex bison libfl-dev libboost-thread-dev libboost-filesystem-dev nvidia-cuda-toolkit-gcc libzstd-dev"
         fi
         # **** WebGPU (dependencies for software-based vulkan) ****
         if [[ "${{ inputs.webgpu }}" == "true" ]]; then
-          pkgs+=" mesa-vulkan-drivers"
+          pkgs+=" libgl1 libglx-mesa0 libgl1-mesa-dri libxcb-xfixes0-dev mesa-vulkan-drivers"
         fi
         # **** LLVM ****
         if [[ "${{ inputs.llvm }}" == "true" ]]; then
           pkgs+=" libllvm20 clang-20 lld-20"
         fi
-        # **** QEMU ****
-        if [[ "${{ inputs.qemu }}" == "true" ]]; then
-          pkgs+=" qemu-user-static"
-        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'
-      uses: actions/cache/restore@v5
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true') && github.event_name == 'pull_request'
+      uses: actions/cache/restore@v4
       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.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == '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.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
       shell: bash
       run: |
         sudo apt -qq update || true
@@ -215,14 +218,8 @@ runs:
           sudo apt-get -y --allow-unauthenticated --no-install-recommends install ${{ steps.apt-pkgs.outputs.pkgs }}
         fi
 
-        sudo mkdir -p /var/cache/apt/archives
         sudo chown -R $USER:$USER /var/cache/apt/archives/
 
-    - name: Add clang to PATH (Linux)
-      if: inputs.llvm == 'true' && runner.os == 'Linux'
-      shell: bash
-      run: echo "/usr/lib/llvm-20/bin" >> "$GITHUB_PATH"
-
     # **** AMD ****
     - name: Setup AMD (Linux)
       if: inputs.amd == 'true' && runner.os == 'Linux'
@@ -242,33 +239,78 @@ runs:
           jq -r '.assets[] | select(.name == "libamd_comgr.dylib").browser_download_url' | \
           sudo xargs curl -fL -o /usr/local/lib/libamd_comgr.dylib
 
-    # **** CUDA ****
-    - name: Install CUDA
-      if: inputs.cuda == 'true'
+    # **** gpuocelot ****
+
+    - name: Install gpuocelot dependencies (MacOS)
+      if: inputs.ocelot == 'true' && runner.os == 'macOS'
       shell: bash
       run: |
-        sudo mkdir -p /usr/local/cuda/targets/x86_64-linux
-        curl -fL https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvrtc/linux-x86_64/cuda_nvrtc-linux-x86_64-11.5.119-archive.tar.xz \
-          | sudo tar -xJ -C /usr/local/cuda/targets/x86_64-linux --strip-components=1
-        echo /usr/local/cuda/targets/x86_64-linux/lib | sudo tee /etc/ld.so.conf.d/cuda-nvrtc.conf
-        sudo ldconfig
+        pkgs=(cmake ninja llvm@15 zlib glew flex bison [email protected] zstd ncurses)
+        for f in "${pkgs[@]}"; do
+          brew ls --versions "$f" >/dev/null 2>&1 || brew install --quiet "$f"
+        done
 
-    # **** gpuocelot ****
+        # Fix boost 1.85 for gpuocelot
+        ln -s /opt/homebrew/opt/[email protected] /opt/homebrew/opt/boost || true
+        ln -s /opt/homebrew/opt/boost/lib/libboost_atomic-mt.dylib /opt/homebrew/opt/boost/lib/libboost_atomic.dylib || true
+        ln -s /opt/homebrew/opt/boost/lib/libboost_thread-mt.dylib /opt/homebrew/opt/boost/lib/libboost_thread.dylib || true
+    - name: Cache gpuocelot (PR)
+      if: inputs.ocelot == 'true' && github.event_name == 'pull_request'
+      id: cache-build-pr
+      uses: actions/cache/restore@v4
+      env:
+        cache-name: cache-gpuocelot-build-1
+      with:
+        path: ${{ github.workspace }}/gpuocelot/ocelot
+        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
+    - name: Cache gpuocelot
+      if: inputs.ocelot == 'true' && github.event_name != 'pull_request'
+      id: cache-build
+      uses: actions/cache@v5
+      env:
+        cache-name: cache-gpuocelot-build-1
+      with:
+        path: ${{ github.workspace }}/gpuocelot/ocelot
+        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
+    - name: Clone/compile gpuocelot
+      if: inputs.ocelot == 'true' && steps.cache-build-pr.outputs.cache-hit != 'true' && steps.cache-build.outputs.cache-hit != 'true'
+      shell: bash
+      run: |
+        git clone --recurse-submodules https://github.com/gpuocelot/gpuocelot.git ${{ github.workspace }}/gpuocelot
+        cd ${{ github.workspace }}/gpuocelot/ocelot
+        git checkout b16039dc940dc6bc4ea0a98380495769ff35ed99
+        mkdir build
+        cd build
+
+        CMAKE_ARGS="-Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5"
+        if [[ "${{ runner.os }}" == "macOS" ]]; then
+          sudo xcode-select -s /Applications/Xcode_16.2.app/Contents/Developer
+          CMAKE_ARGS="$CMAKE_ARGS -DBoost_INCLUDE_DIR=$(brew --prefix boost)/include -DBoost_LIBRARY_DIR=$(brew --prefix boost)/lib"
+        fi
+
+        cmake .. $CMAKE_ARGS
+        ninja
     - name: Install gpuocelot
       if: inputs.ocelot == 'true'
       shell: bash
       run: |
-        sudo mkdir -p /usr/local/lib
-        sudo curl --output-dir /usr/local/lib -fLO https://github.com/tinygrad/gpuocelot/releases/download/v0.1.0/libgpuocelot.${{ runner.os == 'Linux' && 'so' || 'dylib' }}
+        cd ${{ github.workspace }}/gpuocelot/ocelot/build
+        sudo cp libgpuocelot.${{ runner.os == 'macOS' && 'dylib' || 'so' }} /usr/${{ runner.os == 'macOS' && 'local/' || '' }}lib/
 
     # **** WebGPU ****
 
-    - name: Install WebGPU dawn
-      if: inputs.webgpu == 'true'
+    - name: Install WebGPU dawn (Linux)
+      if: inputs.webgpu == 'true' && runner.os == 'Linux'
       shell: bash
       run: |
-        sudo mkdir -p /usr/local/lib
-        sudo curl --output-dir /usr/local/lib -fLO https://github.com/wpmed92/pydawn/releases/download/v0.1.6/libwebgpu_dawn.${{ runner.os == 'Linux' && 'so' || 'dylib' }}
+        sudo curl -fL https://github.com/wpmed92/pydawn/releases/download/v0.1.6/libwebgpu_dawn.so -o /usr/local/lib/libwebgpu_dawn.so
+        sudo ldconfig
+    - name: Install WebGPU dawn (macOS)
+      if: inputs.webgpu == 'true' && runner.os == 'macOS'
+      shell: bash
+      run: |
+        brew tap wpmed92/dawn
+        brew install dawn
 
     # **** LLVM ****
 
@@ -277,18 +319,18 @@ runs:
       shell: bash
       run: brew install llvm@20
 
+    # **** mesa ****
+    - name: Install mesa (linux)
+      if: inputs.mesa == 'true' && runner.os == 'Linux'
+      shell: bash
+      run: sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/v1/libtinymesa_cpu-mesa-25.2.7-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
+    - name: Install mesa (macOS)
+      if: inputs.mesa == 'true' && runner.os == 'macOS'
+      shell: bash
+      run: brew install sirhcm/tinymesa/tinymesa_cpu
+
     # *** tinydreno ***
     - name: Install tinydreno (linux)
       if: inputs.tinydreno == 'true' && runner.os == 'Linux'
       shell: bash
       run: sudo curl -fL https://github.com/sirhcm/tinydreno/raw/refs/heads/master/libllvm-qcom.so -o /usr/lib/libllvm-qcom.so
-
-    # *** OpenCL ***
-    - name: Install rusticl
-      if: inputs.opencl == 'true'
-      shell: bash
-      run: |
-        sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/rusticl-v1/libRusticlOpenCL.so.1.0.0 -o /usr/lib/libRusticlOpenCL.so
-        sudo mkdir -p /etc/OpenCL/vendors
-        echo "/usr/lib/libRusticlOpenCL.so" | sudo tee /etc/OpenCL/vendors/rusticl.icd
-        echo "RUSTICL_ENABLE=llvmpipe" >> "$GITHUB_ENV"

.github/workflows/autogen.yml

@@ -33,20 +33,23 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: 'autogen'
+        opencl: 'true'
         amd: 'true'
+        cuda: 'true'
         llvm: 'true'
+        webgpu: 'true'
+        mesa: 'true'
         pydeps: 'pyyaml mako'
     - 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
+      run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev
     - name: Regenerate autogen files
       run: |
         find tinygrad/runtime/autogen -type f -name "*.py" -not -path "*/amd/*" -not -name "__init__.py" -not -name "comgr.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete
         python3 -c "from tinygrad.runtime.autogen import opencl"
-        python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv_610, nv"
+        python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv"
         python3 -c "from tinygrad.runtime.autogen import comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
-        python3 -c "from tinygrad.runtime.autogen.am import *"
-        python3 -c "from tinygrad.runtime.autogen.nv_regs import *"
-        python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, pci, vfio"
+        python3 -c "from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v13_0_6, smu_v13_0_12, smu_v14_0_2, fw"
+        python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
         python3 -c "from tinygrad.runtime.autogen import llvm"
         python3 -c "from tinygrad.runtime.autogen import webgpu"
         python3 -c "from tinygrad.runtime.autogen import kgsl, qcom_dsp"

.github/workflows/mlperf.yml

@@ -1,8 +1,8 @@
 name: Run MLPerf Training
 
 on:
-  #schedule:
-  #  - cron: '5 8 * * *'  # Runs at 08:05 UTC (12:05 AM Pacific Time)
+  schedule:
+    - cron: '5 8 * * *'  # Runs at 08:05 UTC (12:05 AM Pacific Time)
   push:
     branches:
       - update_mlperf

.github/workflows/platform.yml

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

@@ -14,15 +14,12 @@ jobs:
     outputs:
       branchstat: ${{ steps.brstat.outputs.stat}}
     steps:
-      - name: Check code from PR branch
+      - name: Check code from PR branch 
         uses: actions/checkout@v6
         with:
           repository: ${{ github.event.pull_request.head.repo.full_name }}
           ref: ${{ github.event.pull_request.head.sha }}
           fetch-depth: 0
-          # PR code is only inspected with git rev-list, never executed
-          allow-unsafe-pr-checkout: true
-          persist-credentials: false
       - name: Check whether branch is up-to-date
         id: brstat
         run: |
@@ -54,9 +51,6 @@ jobs:
           repository: ${{ github.event.pull_request.head.repo.full_name }}
           ref: ${{ github.event.pull_request.head.sha }}
           path: pr
-          # PR code is only line-counted by master's sz.py, never executed
-          allow-unsafe-pr-checkout: true
-          persist-credentials: false
         # the base default to tinygrad master and cannot be other fork branch for security purpose
       - name: Checkout code from tinygrad master
         uses: actions/checkout@v6

AGENTS.md

@@ -1,6 +0,0 @@
-# Notes
-
-- Run tests with `-n12` for speed (e.g. `python -m pytest test/null/test_dtype.py -x -q -n12`)
-- Run `python -m mypy tinygrad/` to typecheck
-- Run `python -m ruff check .` to lint
-- Read `./tinygrad/viz/README.md` for profiling and debugging rewrite rules

README.md

@@ -72,7 +72,7 @@ As it turns out, 90% of what you need for neural networks are a decent autograd/
 Throw in an optimizer, a data loader, and some compute, and you have all you need.
 
 ```python
-from tinygrad import Tensor, nn, Context
+from tinygrad import Tensor, nn
 
 class LinearNet:
   def __init__(self):
@@ -86,7 +86,7 @@ optim = nn.optim.Adam([model.l1, model.l2], lr=0.001)
 
 x, y = Tensor.rand(4, 1, 28, 28), Tensor([2,4,3,7])  # replace with real mnist dataloader
 
-with Context(TRAINING=1):
+with Tensor.train():
   for i in range(10):
     optim.zero_grad()
     loss = model(x).sparse_categorical_crossentropy(y).backward()
@@ -140,8 +140,8 @@ Documentation along with a quick start guide can be found on the [docs website](
 ```python
 from tinygrad import Tensor
 
-x = Tensor.eye(3)
-y = Tensor([[2.0,0,-2.0]])
+x = Tensor.eye(3, requires_grad=True)
+y = Tensor([[2.0,0,-2.0]], requires_grad=True)
 z = y.matmul(x).sum()
 z.backward()
 
@@ -164,9 +164,7 @@ print(y.grad.tolist())  # dz/dy
 
 ## Contributing
 
-There has been a lot of interest in tinygrad lately. Following these guidelines will help your PR get accepted. If you do submit a PR, please include a sentence or two about why you want this merged and why you think it will improve the project.
-
-If you are a new contributor with something that looks even close to AI written, it will be closed without feedback and you may be banned from our GitHub. No human should waste time reading AI slop. And for everyone, if you used AI, disclose what you used it for.
+There has been a lot of interest in tinygrad lately. Following these guidelines will help your PR get accepted.
 
 We'll start with what will get your PR closed with a pointer to this section:
 
@@ -198,8 +196,6 @@ python3 test/backend/test_ops.py        # just the ops tests
 python3 -m pytest test/                 # whole test suite
 ```
 
-For agents, always run tests with `-n12` for speed.
-
 #### Process replay tests
 
 [Process replay](https://github.com/tinygrad/tinygrad/blob/master/test/external/process_replay/README.md) compares your PR's generated kernels against master. If your PR is a refactor or speedup without any expected behavior change, It should include [pr] in the pull request title.

conftest.py

@@ -1,10 +0,0 @@
-import os, pytest, signal, threading
-
-@pytest.hookimpl(wrapper=True)
-def pytest_runtest_call(item):
-  t = threading.Timer(int(os.getenv("TEST_TIMEOUT", 300)), os.kill, args=(os.getpid(), signal.SIGABRT))
-  t.start()
-  try: yield
-  finally:
-    t.cancel()
-    t.join()

docs/abstractions3.py

@@ -11,7 +11,7 @@ X_train -= X_train.mean()
 # *****
 # 1. Define an MNIST model.
 
-from tinygrad import Tensor, Context
+from tinygrad import Tensor
 
 l1 = Tensor.kaiming_uniform(128, 784)
 l2 = Tensor.kaiming_uniform(10, 128)
@@ -24,11 +24,11 @@ l1n, l2n = l1.numpy(), l2.numpy()
 from tinygrad.nn.optim import SGD
 optim = SGD([l1, l2])
 
-with Context(TRAINING=1):
-  X, Y = X_train[(samples:=Tensor.randint(128, high=X_train.shape[0]))], Y_train[samples]
-  optim.zero_grad()
-  model(X).sparse_categorical_crossentropy(Y).backward()
-  optim.schedule_step()   # this will step the optimizer without running realize
+Tensor.training = True
+X, Y = X_train[(samples:=Tensor.randint(128, high=X_train.shape[0]))], Y_train[samples]
+optim.zero_grad()
+model(X).sparse_categorical_crossentropy(Y).backward()
+optim.schedule_step()   # this will step the optimizer without running realize
 
 # *****
 # 3. Create a schedule (linear uop).

docs/abstractions4.py

@@ -67,7 +67,8 @@ def example_2_hip(a:Tensor, correct):
     # the sink specifies the GLOBAL and LOCAL sizes, along with the input buffers and name
     sink = UOp.sink(UOp.special(GLOBALS, 'gidx0'), UOp.special(THREADS, 'lidx0'), out, buf,
                     arg=KernelInfo(name="hip_reduce_sum_kernel"))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT),
+                UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
   eval_harness("HIP kernel", a, lambda x: Tensor.empty(GLOBALS).custom_kernel(x, fxn=hip_reduce_sum)[0].sum(), check=correct)
 
 def example_3_custom_uop(a:Tensor, correct):
@@ -88,7 +89,7 @@ def example_3_custom_uop(a:Tensor, correct):
 
     # store all the per lane accumulators to LOCAL
     local_accs = UOp.placeholder((LCLS,), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
-    local_accs = local_accs.after(local_accs[lane].store(acc[0]))
+    local_accs = local_accs.after(local_accs[lane].store(acc[0]).barrier())
 
     # accumulate LOCALs into a single per CU accumulator
     late_reduce_loop = UOp.range(LCLS, 3, AxisType.REDUCE)
@@ -122,7 +123,8 @@ def example_5_custom_assembly(a:Tensor, correct):
         offset_dwords = (self.labels[inst._target] - inst._pos - inst.size()) // 4
         if not -32768 <= offset_dwords <= 32767: raise ValueError(f"branch to '{inst._target}' offset {offset_dwords} exceeds simm16 range")
         inst.simm16 = offset_dwords
-      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in self.instructions]))))
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT),
+                                   UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in self.instructions]))))
 
   CU_COUNT = 32
   LANES = 64

docs/developer/speed.md

@@ -62,7 +62,7 @@ A lot of work can still be done here. For example, we never copy the inputs to o
 
 Many accelerators have Tensor Cores / MAC arrays / systolic arrays. The main value of these is that, since they are 2-D, they create an n^2 ratio between the compute and the input data.
 
-GPUs use Tensor Cores instead of MAC arrays to fit better in the GPU warp paradigm. This is because the output of Tensor Cores is O(n) wrt the input, while the output of MAC arrays is O(n^2)
+GPUs use Tensor Cores instead of MAC arrays to fit better in the GPU warp paradigm. This is because the output of Tensor Cores is O(n) wrt the input, while the output of MAC arrays like the AMX is O(n^2)
 
 We have a simple framework in tinygrad for adding these ALU blocks and achieving good performance from them.
 

docs/dtypes.md

@@ -1,8 +1,7 @@
 ::: tinygrad.dtype.DType
 
-::: tinygrad.dtype.DTypes
+::: tinygrad.dtype.dtypes
     options:
-        heading: dtypes
         members: true
         members_order: source
         show_labels: false

docs/mnist.md

@@ -24,7 +24,7 @@ You will see `CUDA` here on a GPU instance, or `CPU` here on a CPU instance.
 We'll use the model from [the Keras tutorial](https://keras.io/examples/vision/mnist_convnet/).
 
 ```python
-from tinygrad import Tensor, nn, Context
+from tinygrad import Tensor, nn
 
 class Model:
   def __init__(self):
@@ -74,8 +74,8 @@ We'll use the Adam optimizer. The `nn.state.get_parameters` will walk the model
 ```python
 optim = nn.optim.Adam(nn.state.get_parameters(model))
 batch_size = 128
-@Context(TRAINING=1)
 def step():
+  Tensor.training = True  # makes dropout work
   samples = Tensor.randint(batch_size, high=X_train.shape[0])
   X, Y = X_train[samples], Y_train[samples]
   optim.zero_grad()
@@ -143,6 +143,7 @@ Since we are just randomly sampling from the dataset, there's no real concept of
 for step in range(7000):
   loss = jit_step()
   if step%100 == 0:
+    Tensor.training = False
     acc = (model(X_test).argmax(axis=1) == Y_test).mean().item()
     print(f"step {step:4d}, loss {loss.item():.2f}, acc {acc*100.:.2f}%")
 ```

docs/quickstart.md

@@ -133,7 +133,7 @@ For our loss function we will be using sparse categorical cross entropy loss. Th
 ```python
 def sparse_categorical_crossentropy(self, Y, ignore_index=-1) -> Tensor:
     loss_mask = Y != ignore_index
-    y_counter = Tensor.arange(self.shape[-1], dtype=dtypes.int32).unsqueeze(0).expand(Y.numel(), self.shape[-1])
+    y_counter = Tensor.arange(self.shape[-1], dtype=dtypes.int32, requires_grad=False, device=self.device).unsqueeze(0).expand(Y.numel(), self.shape[-1])
     y = ((y_counter == Y.flatten().reshape(-1, 1)).where(-1.0, 0) * loss_mask.reshape(-1, 1)).reshape(*Y.shape, self.shape[-1])
     return self.log_softmax().mul(y).sum() / loss_mask.sum()
 ```
@@ -165,18 +165,17 @@ from extra.datasets import fetch_mnist
 Now we have everything we need to start training our neural network.
 We will be training for 1000 steps with a batch size of 64.
 
-We use `with Context(TRAINING=1)` to enable training mode.
+We use `with Tensor.train()` to set the internal flag `Tensor.training` to `True` during training.
 Upon exit, the flag is restored to its previous value by the context manager.
 
 ```python
-from tinygrad import Context
 X_train, Y_train, X_test, Y_test = fetch_mnist()
 
-with Context(TRAINING=1):
+with Tensor.train():
   for step in range(1000):
     # random sample a batch
     samp = np.random.randint(0, X_train.shape[0], size=(64))
-    batch = Tensor(X_train[samp])
+    batch = Tensor(X_train[samp], requires_grad=False)
     # get the corresponding labels
     labels = Tensor(Y_train[samp])
 
@@ -214,7 +213,7 @@ with Timing("Time: "):
   for step in range(1000):
     # random sample a batch
     samp = np.random.randint(0, X_test.shape[0], size=(64))
-    batch = Tensor(X_test[samp])
+    batch = Tensor(X_test[samp], requires_grad=False)
     # get the corresponding labels
     labels = Y_test[samp]
 
@@ -258,7 +257,7 @@ with Timing("Time: "):
   for step in range(1000):
     # random sample a batch
     samp = np.random.randint(0, X_test.shape[0], size=(64))
-    batch = Tensor(X_test[samp])
+    batch = Tensor(X_test[samp], requires_grad=False)
     # get the corresponding labels
     labels = Y_test[samp]
 

docs/runtime.md

@@ -5,7 +5,7 @@ tinygrad supports various runtimes, enabling your code to scale across a wide ra
 | Runtime | Description | Compiler Options | Requirements |
 |---------|-------------|------------------|--------------|
 | [NV](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_nv.py) | Provides acceleration for NVIDIA GPUs | nvrtc (default)<br>PTX (`DEV=NV:PTX`) | Ampere/Ada/Blackwell series GPUs.<br>You can select an interface via [the `DEV` variable](env_vars.md#dev-variable). See [NV interfaces](#nv-interfaces) for details. |
-| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | LLVM (`DEV=AMD:LLVM`)<br>HIP/COMGR (`DEV=AMD:HIP`) | CDNA3, CDNA4, RDNA3 or RDNA4 GPUs.<br>You can select an interface via [the `DEV` variable](env_vars.md#dev-variable). See [AMD interfaces](#amd-interfaces) for details. |
+| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | LLVM (`DEV=AMD:LLVM`)<br>HIP/COMGR (`DEV=AMD:HIP`) | RDNA2 or newer GPUs.<br>You can select an interface via [the `DEV` variable](env_vars.md#dev-variable). See [AMD interfaces](#amd-interfaces) for details. |
 | [QCOM](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_qcom.py) | Provides acceleration for QCOM GPUs | - | 6xx series GPUs |
 | [METAL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_metal.py) | Utilizes Metal for acceleration on Apple devices | - | M1+ Macs; Metal 3.0+ for `bfloat` support |
 | [CUDA](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cuda.py) | Utilizes CUDA for acceleration on NVIDIA GPUs | nvrtc (default)<br> PTX (`DEV=CUDA:PTX`) | NVIDIA GPU with CUDA support |
@@ -83,5 +83,9 @@ NV backend supports several interfaces for communicating with devices:
 ## CPU Arch
 The CPU renderers may be additionally configured using the arch component of [the `DEV` environment variable](env_vars.md#dev-variable).
 CPU arch should be specified as a comma-separated list of parameters, and must contain at least two values: the architecture family (ie. x86_64, arm64, or riscv64) and the cpu type (as accepted by `clang`'s `-march`).
-If native is specified as the cpu type, tinygrad (or delegate compiler) will query the host cpu type. Additional comma-separated values are interpreted as cpu feature flags. When a value is preceded by a `-` character, the corresponding feature flag will be disabled, otherwise the flag will be enabled.
+If native is specified as the cpu type, tinygrad (or delegate compiler) will query the host cpu type. Additional comma-separated values may be specified as follows:
+
+* `AMX`: emit Apple silicon AMX instructions
+
+All other additional values are interpreted as cpu feature flags. When a value is preceded by a `-` character, the corresponding feature flag will be disabled, otherwise the flag will be enabled.
 Note that enabled feature flags should not be preceded by a `+`.

docs/tensor/elementwise.md

@@ -66,8 +66,8 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.sub
 ::: tinygrad.Tensor.mul
 ::: tinygrad.Tensor.div
+::: tinygrad.Tensor.idiv
 ::: tinygrad.Tensor.mod
-::: tinygrad.Tensor.fmod
 ::: tinygrad.Tensor.bitwise_xor
 ::: tinygrad.Tensor.bitwise_and
 ::: tinygrad.Tensor.bitwise_or

docs/tinygpu.md

@@ -4,7 +4,7 @@ TinyGPU app lets you use AMD and NVIDIA GPUs on macOS over USB4/Thunderbolt with
 
 ## Requirements
 
-- macOS (13.0+)
+- macOS (12.1+)
 - USB4/Thunderbolt port
 - A supported GPU (AMD RDNA3+ or NVIDIA Ampere+)
 

examples/anthropic_challenge.py

@@ -0,0 +1,196 @@
+from tinygrad import Tensor, dtypes, Context, getenv, UOp, fetch
+from tinygrad.uop.ops import Ops, PatternMatcher, UPat
+from tinygrad.uop.symbolic import symbolic
+from tinygrad.codegen import Renderer
+from tinygrad.codegen.opt import Opt, OptOps
+
+# ************************* implementation of the problem ************************
+
+def myhash(a: Tensor) -> Tensor:
+  a = (a + 0x7ED55D16) + (a << 12)
+  a = (a ^ 0xC761C23C) ^ (a >> 19)
+  a = (a + 0x165667B1) + (a << 5)
+  a = (a + 0xD3A2646C) ^ (a << 9)
+  a = (a + 0xFD7046C5) + (a << 3)
+  a = (a ^ 0xB55A4F09) ^ (a >> 16)
+  return a
+
+def select_with_where_tree(values: Tensor, relative_idx: Tensor) -> Tensor:
+  n = values.shape[0]
+  if n == 1: return values[0].expand(relative_idx.shape)
+
+  mid = n // 2
+  left = select_with_where_tree(values[:mid], relative_idx)
+  right = select_with_where_tree(values[mid:], relative_idx - mid)
+
+  go_left = relative_idx < mid
+  return go_left.where(left, right)
+
+def tree_traversal(forest: Tensor, val: Tensor, height: int, rounds: int, where_tree_threshold=3) -> Tensor:
+  # All walkers start at idx=0
+  idx = Tensor.zeros(val.shape, device=val.device, dtype=dtypes.uint32)
+
+  for r in range(rounds):
+    level = r % (height + 1)
+    level_start = (1 << level) - 1
+    level_size = 1 << level
+
+    if level == 0:
+      # At root (level 0), all walkers are at idx=0
+      # No gather needed, just broadcast the root value
+      node_val = forest[0].expand(val.shape)
+      idx = idx * 0  # Reset to 0
+    elif level <= where_tree_threshold:
+      # Small level: use where-tree
+      level_values = forest[level_start : level_start + level_size]
+      relative_idx = (idx - level_start)
+      node_val = select_with_where_tree(level_values, relative_idx)
+    else:
+      # Large level: use gather
+      node_val = forest.gather(0, idx)
+
+    val = myhash(val ^ node_val)
+    idx = (idx << 1) + (1 + (val & 1))
+
+    # No wrap check needed! At round 10 (level becomes 0), we reset idx above.
+
+  return val.contiguous(arg=(Opt(OptOps.UPCAST, 0, 8),))
+
+# ************************* renderer for VLIW machine *************************
+
+def loop_unrolling(sink:UOp):
+  rng = [x for x in sink.toposort() if x.op is Ops.RANGE]
+  if len(rng) == 0: return None
+  print(f"unrolling loop with size {rng[0].vmax+1}")
+  unrolled_sinks = [sink.substitute({rng[0]:rng[0].const_like(i)}).src[0] for i in range(rng[0].vmax+1)]
+  return UOp.sink(*unrolled_sinks, arg=sink.arg)
+
+global_addrs = []
+vliw_prepare = PatternMatcher([
+  # loop unrolling (should be a part of tinygrad)
+  (UPat(Ops.SINK, name="sink"), loop_unrolling),
+  # cast is fake
+  (UPat(Ops.CAST, name="c"), lambda c: c.src[0]),
+  # rewrites to hardcode the addresses in memory
+  (UPat(Ops.PARAM, name="dg"), lambda dg: UOp.const(dtypes.uint, global_addrs[dg.arg])),
+  # INDEX is just plus
+  (UPat(Ops.INDEX, name="i"), lambda i: i.src[0]+i.src[1]),
+])+symbolic
+
+class VLIWRenderer(Renderer):
+  has_local = False  # TODO: this should be the default / cleaned up
+  # this says this backend supports MULACC + more. decompositions uses this
+  code_for_op: dict = {Ops.MULACC: None, Ops.ADD: "+", Ops.MUL: "*",
+                       Ops.XOR: "^", Ops.AND: "&", Ops.OR: "|",
+                       Ops.SHL: "<<", Ops.SHR: ">>", Ops.CMPLT: "<"}
+  # this matcher runs while still in graph form
+  pre_matcher = vliw_prepare
+
+  def render(self, uops:list[UOp]):
+
+    # TODO: this is a minimal renderer. for low cycle count, make it good
+    # to get speed, you need to add VLIW packing
+    # to get under 1536 regs, you need to add a register allocator
+    # we left the fun parts to you
+
+    print(f"rendering with {len(uops)} uops")
+    reg, inst = 0, []
+    r: dict[UOp, int] = {}
+    for u in uops:
+      assert u.dtype.count in (1,8), "dtype count must be 1 or 8"
+
+      # dumb register allocator
+      if u.op not in {Ops.STORE, Ops.SINK, Ops.GEP}:
+        r[u] = reg
+        reg += u.dtype.count
+
+      # render UOps to instructions
+      match u.op:
+        case Ops.SINK:
+          inst.append({"flow": [("halt",)]})
+        case Ops.CONST:
+          inst.append({"load": [("const", r[u], u.arg)]})
+        case Ops.GEP:
+          # a GEP is just an alias to a special register in the vector
+          r[u] = r[u.src[0]] + u.arg[0]
+        case Ops.STACK:
+          if all(s == u.src[0] for s in u.src):
+            # if all sources are the same, we can broadcast
+            inst.append({"valu": [("vbroadcast", r[u], r[u.src[0]])]})
+          else:
+            # this is a copy into a contiguous chunk of registers
+            inst.extend({"flow": [("add_imm", r[u]+i, r[s], 0)]} for i,s in enumerate(u.src) if r[s] != r[u]+i)
+        case Ops.LOAD:
+          op = "vload" if u.dtype.count > 1 else "load"
+          inst.append({"load": [(op, r[u], r[u.src[0]])]})
+        case Ops.STORE:
+          op = "vstore" if u.src[1].dtype.count > 1 else "store"
+          inst.append({"store": [(op, r[u.src[0]], r[u.src[1]])]})
+        case Ops.MULACC:
+          assert u.dtype.count == 8
+          inst.append({"valu": [("multiply_add", r[u], r[u.src[0]], r[u.src[1]], r[u.src[2]])]})
+        case Ops.WHERE:
+          assert u.dtype.count == 8
+          inst.append({"flow": [("vselect", r[u], r[u.src[0]], r[u.src[1]], r[u.src[2]])]})
+        case _ if u.op in self.code_for_op:
+          cat = "valu" if u.dtype.count > 1 else "alu"
+          inst.append({cat: [(self.code_for_op[u.op], r[u], r[u.src[0]], r[u.src[1]])]})
+        case _:
+          raise NotImplementedError(f"unhandled op {u.op}")
+    return repr(inst)
+
+# ************************* test and render *************************
+
+import sys, types
+PROBLEM_URL = "https://raw.githubusercontent.com/anthropics/original_performance_takehome/refs/heads/main/tests/frozen_problem.py"
+sys.modules["problem"] = problem = types.ModuleType("problem")
+exec(fetch(PROBLEM_URL).read_text(), problem.__dict__)
+
+if __name__ == "__main__":
+  batch_size = getenv("BS", 256)
+  height = 10
+  rounds = getenv("ROUNDS", 16)
+
+  # build problem
+  tree = problem.Tree.generate(height)
+  inp = problem.Input.generate(tree, batch_size, rounds)
+  mem = problem.build_mem_image(tree, inp)
+  global_addrs.extend([mem[6], mem[6], mem[4]])  # output, input, forest
+
+  # *** verify the kernel in tinygrad compared to reference ***
+
+  forest_t = Tensor(tree.values, dtype=dtypes.uint32)
+  val_t = Tensor(inp.values, dtype=dtypes.uint32)
+
+  if getenv("VERIFY", 1):
+    # verify on normal tinygrad device
+    with Context(PCONTIG=2):
+      out = tree_traversal(forest_t, val_t, height, rounds)
+      val_out = out.tolist()
+    problem.reference_kernel(tree, inp)
+    assert val_out == inp.values
+    print("verification passed")
+
+  # *** render to device ***
+
+  from tinygrad.codegen import to_program
+  with Context(PCONTIG=2, DEVECTORIZE=2, SPEC=0):
+    out = tree_traversal(forest_t, val_t, height, rounds)
+    sink = out.schedule_linear().src[-1].src[0]
+    prg = to_program(sink, VLIWRenderer())
+
+  # *** run on Machine and compare ***
+
+  # NOTE: the scratch size needs to be reduced to 1536 when you have a register allocator
+  src = eval(prg.src[3].arg)
+  max_regs = max(t[1] for instr in src for v in instr.values() for t in v if len(t) > 1) + 8
+  print(f"{max_regs:5d} regs used" + ("" if max_regs <= 1536 else "       <-- WARNING: TOO MANY REGISTERS, MUST BE <= 1536"))
+  machine = problem.Machine(mem, src, problem.DebugInfo(scratch_map={}), n_cores=1, trace=False, scratch_size=max_regs)
+  machine.run()
+  print(f"ran for {machine.cycle:5d} cycles" + ("" if machine.cycle <= 1363 else "  <-- EVEN CLAUDE GOT 1363"))
+
+  # compare to reference
+  ref_mem = mem.copy()
+  for _ in problem.reference_kernel2(ref_mem, {}): pass
+  assert machine.mem[mem[6]:mem[6]+mem[2]] == ref_mem[mem[6]:mem[6]+mem[2]]
+  print("compare passed!")

examples/audio_helpers.py

@@ -4,10 +4,10 @@ from tinygrad.dtype import DTypeLike, dtypes
 import math
 
 # rewritten from numpy
-def rfftfreq(n: int, d: float = 1.0) -> Tensor:
+def rfftfreq(n: int, d: float = 1.0, device=None) -> Tensor:
   val = 1.0 / (n * d)
   N = n // 2 + 1
-  results = Tensor.arange(N)
+  results = Tensor.arange(N, device=device)
   return results * val
 
 # just like in librosa

examples/beautiful_cartpole.py

@@ -1,6 +1,6 @@
 from typing import Tuple
 import time
-from tinygrad import Tensor, TinyJit, nn, Context
+from tinygrad import Tensor, TinyJit, nn
 import gymnasium as gym
 from tinygrad.helpers import trange
 import numpy as np  # TODO: remove numpy import
@@ -55,7 +55,7 @@ if __name__ == "__main__":
 
   @TinyJit
   def train_step(x:Tensor, selected_action:Tensor, reward:Tensor, old_log_dist:Tensor) -> Tuple[Tensor, Tensor, Tensor]:
-    with Context(TRAINING=1):
+    with Tensor.train():
       log_dist, value = model(x)
       action_mask = (selected_action.reshape(-1, 1) == Tensor.arange(log_dist.shape[1]).reshape(1, -1).expand(selected_action.shape[0], -1)).float()
 

examples/beautiful_cifar.py

@@ -9,7 +9,8 @@ from extra.lr_scheduler import OneCycleLR
 GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))]
 
 # override tinygrad defaults
-Context(DEFAULT_FLOAT=dtypes.half, FUSE_OPTIM=1).__enter__()
+dtypes.default_float = dtypes.half
+Context(FUSE_OPTIM=1).__enter__()
 
 # from https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
 batchsize = getenv("BS", 1024)
@@ -66,8 +67,8 @@ class ConvGroup:
     self.conv2 = nn.Conv2d(channels_out, channels_out, kernel_size=3, padding=1, bias=False)
     self.norm1 = nn.BatchNorm(channels_out, track_running_stats=False, eps=1e-12, momentum=hyp['net']['batch_norm_momentum'])
     self.norm2 = nn.BatchNorm(channels_out, track_running_stats=False, eps=1e-12, momentum=hyp['net']['batch_norm_momentum'])
-    cast(Tensor, self.norm1.weight).is_param_(False)
-    cast(Tensor, self.norm2.weight).is_param_(False)
+    cast(Tensor, self.norm1.weight).requires_grad = False
+    cast(Tensor, self.norm2.weight).requires_grad = False
   def __call__(self, x:Tensor) -> Tensor:
     x =    self.norm1(self.conv1(x).max_pool2d().float()).cast(dtypes.default_float).quick_gelu()
     return self.norm2(self.conv2(x).float()).cast(dtypes.default_float).quick_gelu() + x
@@ -121,7 +122,7 @@ if __name__ == "__main__":
     return ret.mul(hyp['opt']['loss_scale_scaler']*loss_batchsize_scaler).sum().div(hyp['opt']['loss_scale_scaler'])
 
   @TinyJit
-  @Context(TRAINING=1)
+  @Tensor.train()
   def train_step(idxs:Tensor) -> Tensor:
     X, Y = X_train[idxs], Y_train[idxs]
     if len(GPUS) > 1:

examples/beautiful_mnist.py

@@ -1,6 +1,6 @@
 # model based off https://medium.com/data-science/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
 from typing import Callable
-from tinygrad import Tensor, TinyJit, nn, GlobalCounters, function, Context
+from tinygrad import Tensor, TinyJit, nn, GlobalCounters, function
 from tinygrad.helpers import getenv, colored, trange
 from tinygrad.nn.datasets import mnist
 
@@ -19,7 +19,7 @@ class Model:
   def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers)
 
   @TinyJit
-  @Context(TRAINING=1)
+  @Tensor.train()
   def train_step(self, X_train:Tensor, Y_train:Tensor) -> Tensor:
     opt.zero_grad()
     samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])

examples/beautiful_mnist_multigpu.py

@@ -1,6 +1,6 @@
 # model based off https://towardsdatascience.com/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
 from typing import List, Callable
-from tinygrad import Tensor, TinyJit, nn, GlobalCounters, Device, Context
+from tinygrad import Tensor, TinyJit, nn, GlobalCounters, Device
 from tinygrad.helpers import getenv, colored, trange
 from tinygrad.nn.datasets import mnist
 
@@ -31,7 +31,7 @@ if __name__ == "__main__":
 
   @TinyJit
   def train_step() -> Tensor:
-    with Context(TRAINING=1):
+    with Tensor.train():
       opt.zero_grad()
       samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
       Xt, Yt = X_train[samples].shard_(GPUS, axis=0), Y_train[samples].shard_(GPUS, axis=0)  # we shard the data on axis 0

examples/gpt2.py

@@ -22,6 +22,10 @@ class Attention:
     self.head_dim = dim // n_heads
 
   def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]) -> Tensor:
+    if mask is not None or start_pos.val == 0:
+      # no symbolic shape qkv when consuming prompts
+      start_pos = start_pos.val
+
     if HALF: x = x.half()
     xqkv = self.c_attn(x).reshape(None, None, 3, self.n_heads, self.head_dim)
     xq, xk, xv = [xqkv[:, :, i, :, :] for i in range(3)]
@@ -34,8 +38,12 @@ class Attention:
     # update the cache
     self.cache_kv[:, :, start_pos:start_pos+seqlen, :, :].assign(Tensor.stack(xk, xv)).realize()
 
-    keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
-    values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
+    if start_pos > 0:
+      keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
+      values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
+    else:
+      keys = xk
+      values = xv
 
     xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
     return self.c_proj(xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2).reshape(bsz, seqlen, self.dim))
@@ -78,14 +86,15 @@ class Transformer:
       seqlen = tokens.shape[1]
       tok_emb = self.wte(tokens)
 
-    # start_pos is a bound Variable, so everything below it stays symbolic
-    pos_emb = self.wpe(self.allpos.shrink((None, (start_pos, start_pos+seqlen))))
+    # not symbolic when consuming the prompt
+    selected_pos = (0, seqlen) if start_pos.val == 0 else (start_pos, start_pos+1)
+    pos_emb = self.wpe(self.allpos.shrink((None, selected_pos)))
 
     h = tok_emb + pos_emb
 
     if HALF: h = h.half()
 
-    mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype).triu(start_pos+1) if seqlen > 1 else None
+    mask = Tensor.full((1, 1, seqlen, start_pos.val+seqlen), float("-inf"), dtype=h.dtype).triu(start_pos.val+1) if seqlen > 1 else None
 
     for hi in self.h: h = hi(h, start_pos, mask)
 

examples/gradaccum_mnist.py

@@ -1,6 +1,6 @@
 import itertools
 from typing import Callable
-from tinygrad import nn, Tensor, dtypes, Device, TinyJit, Context
+from tinygrad import nn, Tensor, dtypes, Device, TinyJit
 from tinygrad.helpers import getenv, trange, partition
 
 class Model:
@@ -35,21 +35,22 @@ if __name__ == "__main__":
 
   params = nn.state.get_parameters(model)
 
-  # init params
+  # init params, set requires grad on the ones we need gradients of
   for x in params:
+    if x.requires_grad is None: x.requires_grad_()
     x.replace(x.contiguous())
   Tensor.realize(*params)
 
   # split params (with grads) and buffers (without)
-  params, buffers = partition(params, lambda x: x.is_param)
+  params, buffers = partition(params, lambda x: x.requires_grad)
   print(f"params: {len(params)} buffers: {len(buffers)}")
 
   # optim params
   pos_params = list(itertools.accumulate(params, lambda x,y: x+y.numel(), initial=0))
   adam_m = Tensor.zeros(pos_params[-1], device="CPU").contiguous()
   adam_v = Tensor.zeros(pos_params[-1], device="CPU").contiguous()
-  adam_b1_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU").contiguous()
-  adam_b2_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU").contiguous()
+  adam_b1_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU", requires_grad=False).contiguous()
+  adam_b2_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU", requires_grad=False).contiguous()
   adam_params = [adam_m, adam_v, adam_b1_t, adam_b2_t]
 
   # create loss and grads. init all state so the JIT works on microbatch
@@ -59,7 +60,7 @@ if __name__ == "__main__":
   Tensor.realize(*params, *buffers, *adam_params, loss, grads)
 
   @TinyJit
-  @Context(TRAINING=1)
+  @Tensor.train()
   def microbatch():
     samples = Tensor.randint(BS // ACC_STEPS, high=X_train.shape[0])
     for t in params: t.grad = None

examples/hlb_cifar10.py

@@ -10,7 +10,7 @@ from extra.lr_scheduler import OneCycleLR
 from tinygrad import nn, dtypes, Tensor, Device, GlobalCounters, TinyJit, Variable
 from tinygrad.nn.state import get_state_dict
 from tinygrad.nn import optim
-from tinygrad.helpers import Context, BEAM, WINO, getenv, colored, prod, TRAINING
+from tinygrad.helpers import Context, BEAM, WINO, getenv, colored, prod
 from extra.bench_log import BenchEvent, WallTimeEvent
 
 cifar_mean = [0.4913997551666284, 0.48215855929893703, 0.4465309133731618]
@@ -30,9 +30,9 @@ class UnsyncedBatchNorm:
     if affine: self.weight, self.bias = Tensor.ones(sz, dtype=dtypes.float32), Tensor.zeros(sz, dtype=dtypes.float32)
     else: self.weight, self.bias = None, None
 
-    self.running_mean = Tensor.zeros(num_devices, sz, dtype=dtypes.float32).is_param_(False)
-    self.running_var = Tensor.ones(num_devices, sz, dtype=dtypes.float32).is_param_(False)
-    self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.int).is_param_(False)
+    self.running_mean = Tensor.zeros(num_devices, sz, dtype=dtypes.float32, requires_grad=False)
+    self.running_var = Tensor.ones(num_devices, sz, dtype=dtypes.float32, requires_grad=False)
+    self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.int, requires_grad=False)
 
   def __call__(self, x:Tensor):
     xr = x.reshape(self.num_devices, -1, *x.shape[1:]).cast(dtypes.float32)
@@ -44,7 +44,7 @@ class UnsyncedBatchNorm:
     return ret.reshape(x.shape).cast(x.dtype)
 
   def calc_stats(self, x:Tensor):
-    if TRAINING:
+    if Tensor.training:
       # This requires two full memory accesses to x
       # https://github.com/pytorch/pytorch/blob/c618dc13d2aa23625cb0d7ada694137532a4fa33/aten/src/ATen/native/cuda/Normalization.cuh
       # There's "online" algorithms that fix this, like https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_Online_algorithm
@@ -68,7 +68,8 @@ class UnsyncedBatchNorm:
 class BatchNorm(nn.BatchNorm2d if getenv("SYNCBN") else UnsyncedBatchNorm):
   def __init__(self, num_features):
     super().__init__(num_features, track_running_stats=False, eps=1e-12, momentum=0.85, affine=True)
-    self.weight.is_param_(False)
+    self.weight.requires_grad = False
+    self.bias.requires_grad = True
 
 class ConvGroup:
   def __init__(self, channels_in, channels_out):
@@ -152,21 +153,26 @@ def train_cifar():
 
   # ========== Model ==========
   def whitening(X, kernel_size=hyp['net']['kernel_size']):
-    def _patches(data:Tensor, patch_size=(kernel_size,kernel_size)):
+    def _cov(X):
+      return (X.T @ X) / (X.shape[0] - 1)
+
+    def _patches(data, patch_size=(kernel_size,kernel_size)):
       h, w = patch_size
-      _, c, _, _ = data.shape
-      return data._pool((h, w)).permute(1, 4, 5, 0, 3, 2).reshape(c*h*w, -1)
+      c = data.shape[1]
+      axis = (2, 3)
+      return np.lib.stride_tricks.sliding_window_view(data, window_shape=(h,w), axis=axis).transpose((0,3,2,1,4,5)).reshape((-1,c,h,w))
 
     def _eigens(patches):
-      cov = ((patches @ patches.T) / (patches.shape[1] - 1)).numpy()
-      eigvals, eigvecs = np.linalg.eigh(cov, UPLO='U')
-      return np.flip(eigvals, 0), np.flip(eigvecs.T.reshape(patches.shape[0], X.shape[1], kernel_size, kernel_size), 0)
+      n,c,h,w = patches.shape
+      Σ = _cov(patches.reshape(n, c*h*w))
+      Λ, V = np.linalg.eigh(Σ, UPLO='U')
+      return np.flip(Λ, 0), np.flip(V.T.reshape(c*h*w, c, h, w), 0)
 
     # NOTE: np.linalg.eigh only supports float32 so the whitening layer weights need to be converted to float16 manually
-    eigvals, eigvecs = _eigens(_patches(X.float()))
-    W = eigvecs/np.sqrt(eigvals+1e-2)[:,None,None,None]
+    Λ, V = _eigens(_patches(X.float().numpy()))
+    W = V/np.sqrt(Λ+1e-2)[:,None,None,None]
 
-    return Tensor(W.astype(np.float32)).cast(dtypes.default_float).is_param_(False)
+    return Tensor(W.astype(np.float32), requires_grad=False).cast(dtypes.default_float)
 
   # ========== Loss ==========
   def cross_entropy(x:Tensor, y:Tensor, reduction:str='mean', label_smoothing:float=0.0) -> Tensor:
@@ -218,7 +224,7 @@ def train_cifar():
 
   @TinyJit
   def augmentations(X:Tensor, Y:Tensor):
-    perms = Tensor.randperm(X.shape[0], device=X.device) # We reuse perms for cutmix, because they are expensive to generate
+    perms = Tensor.randperm(X.shape[0], device=X.device) # We reuse perms for cutmix, because they are expensivne to generate
     if getenv("RANDOM_CROP", 1):
       X = random_crop(X, crop_size=32)
     if getenv("RANDOM_FLIP", 1):
@@ -258,6 +264,7 @@ def train_cifar():
       # self.model_ema = copy.deepcopy(net) # won't work for opencl due to unpickeable pyopencl._cl.Buffer
       self.net_ema = SpeedyResNet(w)
       for net_ema_param, net_param in zip(get_state_dict(self.net_ema).values(), get_state_dict(net).values()):
+        net_ema_param.requires_grad = False
         net_ema_param.assign(net_param.numpy())
 
     @TinyJit
@@ -300,7 +307,7 @@ def train_cifar():
   params_bias = []
   params_non_bias = []
   for params in params_dict:
-    if params_dict[params].is_param:
+    if params_dict[params].requires_grad is not False:
       if 'bias' in params:
         params_bias.append(params_dict[params])
       else:
@@ -309,9 +316,6 @@ def train_cifar():
   opt_bias     = optim.SGD(params_bias,     lr=0.01, momentum=hyp['opt']['momentum'], nesterov=True, weight_decay=hyp['opt']['bias_decay'])
   opt_non_bias = optim.SGD(params_non_bias, lr=0.01, momentum=hyp['opt']['momentum'], nesterov=True, weight_decay=hyp['opt']['non_bias_decay'])
 
-  # realize model params and optimizer state before JIT to avoid cache misses
-  Tensor.realize(*params_dict.values(), *opt_bias.b, *opt_non_bias.b)
-
   # NOTE taken from the hlb_CIFAR repository, might need to be tuned
   initial_div_factor = hyp['opt']['initial_div_factor']
   final_lr_ratio = hyp['opt']['final_lr_ratio']
@@ -328,7 +332,9 @@ def train_cifar():
       # index 0 for bias and 1 for non-bias
       optimizer.zero_grad()
       loss.backward()
-      return loss.realize(*optimizer.schedule_step(), *lr_scheduler[0].schedule_step(), *lr_scheduler[1].schedule_step())
+      optimizer.step()
+      lr_scheduler[0].step()
+      lr_scheduler[1].step()
     return loss.realize()
 
   train_step_jitted = TinyJit(train_step)
@@ -355,11 +361,11 @@ def train_cifar():
   i = 0
   eval_acc_pct = 0.0
   batcher = fetch_batches(X_train, Y_train, BS=BS, is_train=True)
-  with Context(TRAINING=1):
+  with Tensor.train():
     st = time.monotonic()
     while i <= STEPS:
       if i % getenv("EVAL_STEPS", STEPS) == 0 and i > 1 and not getenv("DISABLE_BACKWARD"):
-        # Using Context(TRAINING=0) here actually bricks batchnorm, even with track_running_stats=True
+        # Use Tensor.training = False here actually bricks batchnorm, even with track_running_stats=True
         corrects = []
         corrects_ema = []
         losses = []

examples/llama3.py

@@ -102,7 +102,7 @@ class Int8Embedding:
     self.weight, self.scale = Tensor.ones(vocab_size, embed_size, dtype=dtypes.int8), Tensor.ones(vocab_size, dtype=dtypes.half)
 
   def __call__(self, idx:Tensor) -> Tensor:
-    if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz).unsqueeze(-1)
+    if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, requires_grad=False, device=self.weight.device).unsqueeze(-1)
     big_shp = idx.shape+(self.vocab_sz, self.embed_sz)
     arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1)).expand(big_shp), (self.weight.cast(self.scale.dtype).T*self.scale).T
     return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
@@ -123,7 +123,7 @@ def NF4Linear(block_size):
     def __call__(self, x: Tensor) -> Tensor:
       high_bits = self.weight
       low_bits = (self.weight * 2 ** 4).contiguous()
-      unpacked = Tensor.stack(high_bits, low_bits, dim=-1).div(2 ** 4, rounding_mode="trunc")
+      unpacked = Tensor.stack(high_bits, low_bits, dim=-1).idiv(2 ** 4)
       unscaled = CODE[unpacked].to(x.device).reshape(-1, block_size) * self.scale
       return x.linear(unscaled.reshape(self.out_features, self.in_features).T)
 

examples/llm.c/export.py

@@ -3,7 +3,7 @@ import os
 if "NOOPT" not in os.environ: os.environ["NOOPT"] = "1"
 from tinygrad import Device, nn, Tensor, dtypes
 from train_gpt2 import GPT, GPTConfig
-from tinygrad.helpers import DEV, dedup, flatten, getenv, GlobalCounters, to_function_name, Context
+from tinygrad.helpers import DEV, dedup, flatten, getenv, GlobalCounters, to_function_name
 from tinygrad.engine.realize import get_kernel
 from tinygrad.schedule.memory import memory_planner
 from tinygrad.uop.ops import Ops
@@ -23,23 +23,23 @@ if __name__ == "__main__":
   #B, T = Variable("B", 1, 128).bind(4), 64 #Variable("T", 1, 1024).bind(64)
   B, T = 4, 64
 
+  Tensor.training = True
   optimizer = nn.optim.Adam(nn.state.get_parameters(model), lr=1e-4)
   warmup_count = getenv("WARMUP", 3)
-  with Context(TRAINING=1):
-    for i in range(warmup_count):  # TODO: why does it take three and not two to stabilize
-      GlobalCounters.reset()
-      X = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
-      Y = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
-      _, loss = model(X, Y)
-      optimizer.zero_grad()
-      if getenv("BACKWARD", 1):
-        loss.backward()
-        tensors = optimizer.schedule_step()
-      else:
-        tensors = []
-      sched = loss.schedule(*tensors)
-      print(f"calls {i}:", len(sched))
-      #run_schedule(sched[:])
+  for i in range(warmup_count):  # TODO: why does it take three and not two to stabilize
+    GlobalCounters.reset()
+    X = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
+    Y = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
+    _, loss = model(X, Y)
+    optimizer.zero_grad()
+    if getenv("BACKWARD", 1):
+      loss.backward()
+      tensors = optimizer.schedule_step()
+    else:
+      tensors = []
+    sched = loss.schedule(*tensors)
+    print(f"calls {i}:", len(sched))
+    #run_schedule(sched[:])
   sched = memory_planner(sched)
   ast_dedup = dedup([si.ast for si in sched if si.ast.op is Ops.SINK])
   srcs = {}

examples/llm.c/train_gpt2.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 import os, math, time
 import numpy as np
-from tinygrad import Tensor, nn, fetch, Device, TinyJit, GlobalCounters, Context
+from tinygrad import Tensor, nn, fetch, Device, TinyJit, GlobalCounters
 from dataclasses import dataclass
 
 @dataclass
@@ -25,7 +25,7 @@ class CausalSelfAttention:
     self.n_embd = config.n_embd
     # not really a 'bias', more of a mask, but following the OpenAI/HF naming though
     self.bias = Tensor.ones(1, 1, config.block_size, config.block_size).tril()
-    self.bias.is_param_(False)
+    self.bias.requires_grad = False
 
   def __call__(self, x:Tensor):
     B, T, C = x.shape
@@ -99,7 +99,7 @@ class GPT:
 
   def __call__(self, idx:Tensor, targets=None):
     b, t = idx.shape
-    pos = Tensor.arange(0, t)
+    pos = Tensor.arange(0, t, device=idx.device)
 
     tok_emb = self.wte(idx) # token embeddings of shape (b, t, n_embd)
     pos_emb = self.wpe(pos) # position embeddings of shape (t, n_embd)
@@ -177,7 +177,7 @@ if __name__ == "__main__":
   if args.gpus > 1: x, y = x.shard(GPUS, axis=0), y.shard(GPUS, axis=0)
 
   @TinyJit
-  @Context(TRAINING=1)
+  @Tensor.train()
   def step(x:Tensor, y:Tensor) -> Tensor:
     _, loss = model(x, y)
     optimizer.zero_grad()
@@ -204,3 +204,4 @@ if __name__ == "__main__":
     top_k = 40
     y = model.generate(x, max_new_tokens, temperature=temperature, top_k=top_k)
     print(decode(y[0].tolist()))
+

examples/minrf.py

@@ -1,5 +1,5 @@
 # much taken from https://github.com/cloneofsimo/minRF
-from tinygrad import Tensor, nn, GlobalCounters, TinyJit, Context
+from tinygrad import Tensor, nn, GlobalCounters, TinyJit
 from tinygrad.helpers import getenv, trange
 from extra.models.llama import Attention, FeedForward, precompute_freqs_cis
 
@@ -135,7 +135,7 @@ if __name__ == "__main__":
   optimizer = nn.optim.Adam(nn.state.get_parameters(model), lr=5e-4)
 
   @TinyJit
-  @Context(TRAINING=1)
+  @Tensor.train()
   def train_step():
     if getenv("OVERFIT"): samples = Tensor.zeros(getenv("BS", 256), dtype='int')
     else: samples = Tensor.randint(getenv("BS", 256), high=X_train.shape[0])

examples/mixtral.py

@@ -1,6 +1,6 @@
 import functools, argparse, pathlib
 from tinygrad import Tensor, nn, Device, GlobalCounters, Variable
-from tinygrad.helpers import Timing, Profiling, tqdm
+from tinygrad.helpers import Timing, Profiling, CI, tqdm
 from tinygrad.nn.state import torch_load, get_state_dict
 from extra.models.llama import FeedForward, Transformer
 from extra.bench_log import BenchEvent, WallTimeEvent
@@ -36,7 +36,7 @@ if __name__ == "__main__":
     model = Transformer(n_layers=32, dim=4096, hidden_dim=14336, n_heads=32, n_kv_heads=8, norm_eps=1e-5, vocab_size=32000, feed_forward=functools.partial(MixtureFeedForward, 8), jit=False)
     model_state_dict = get_state_dict(model)
 
-    for k in (t := tqdm(state, disable=None)):
+    for k in (t := tqdm(state, disable=CI)):
       if 'feed_forward.experts.' in k:
         expert_no = int(k.split('feed_forward.experts.')[1].split('.')[0])
         device = Device.DEFAULT + ":" + str((expert_no//2)+1)
@@ -44,7 +44,7 @@ if __name__ == "__main__":
         device = Device.DEFAULT
       t.set_description(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB, loading {k} to {device}")
       model_state_dict[k].replace(state[k].to(device).half()).realize()
-  if t.disable: print(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB")
+  if CI: print(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB")
 
   from sentencepiece import SentencePieceProcessor
   spp = SentencePieceProcessor(model_file=args.weights + "/tokenizer.model")

examples/mlperf/dataloader.py

@@ -1,11 +1,11 @@
 import os, random, pickle, queue, struct, math, functools, hashlib, time
 from typing import List
 from pathlib import Path
-from multiprocessing import Queue, Process, shared_memory, connection, Lock
+from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu_count
 
 import numpy as np
 from tinygrad import dtypes, Tensor
-from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX, NUM_CPU_THREADS
+from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX
 from tinygrad.nn.state import TensorIO
 
 ### ResNet
@@ -131,7 +131,7 @@ def batch_load_resnet(batch_size=64, val=False, shuffle=True, seed=None, pad_fir
     else: X = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name}")
     Y = [None] * (batch_size*BATCH_COUNT)
 
-    for _ in range(NUM_CPU_THREADS.value):
+    for _ in range(cpu_count()):
       p = Process(target=loader_process, args=(q_in, q_out, X, seed))
       p.daemon = True
       p.start()
@@ -212,7 +212,7 @@ def batch_load_train_bert(BS:int, seed:int|None=None):
     rng.shuffle(fs)
     train_files.append(fs.pop(0))
 
-  cycle_length = min(NUM_CPU_THREADS.value, len(train_files))
+  cycle_length = min(getenv("NUM_CPU_THREADS", min(os.cpu_count(), 8)), len(train_files))
   assert cycle_length > 0, "cycle_length must be greater than 0"
 
   dataset = InterleavedDataset(train_files, cycle_length)
@@ -301,7 +301,7 @@ def batch_load_unet3d(preprocessed_dataset_dir:Path, batch_size:int=6, val:bool=
     X = Tensor.empty(*sz, dtype=dtypes.float32, device=f"disk:/dev/shm/{shm_name_x}")
     Y = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name_y}")
 
-    for _ in range(NUM_CPU_THREADS.value):
+    for _ in range(cpu_count()):
       proc = Process(target=load_unet3d_data, args=(preprocessed_dataset_dir, seed, queue_in, queue_out, X, Y))
       proc.daemon = True
       proc.start()
@@ -437,7 +437,7 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
   dataset_iter = iter(image_ids)
 
   try:
-    for _ in range(NUM_CPU_THREADS.value):
+    for _ in range(cpu_count()):
       proc = Process(
         target=load_retinanet_data,
         args=(base_dir, val, queue_in, queue_out, imgs, boxes, labels),

examples/mlperf/initializers.py

@@ -2,7 +2,7 @@ import math
 from typing import Union
 
 from tinygrad import Tensor, nn, dtypes
-from tinygrad.helpers import prod, argfix, Context, TRAINING
+from tinygrad.helpers import prod, argfix, Context
 from tinygrad.nn.state import get_parameters
 from extra.models.unet import UNetModel
 
@@ -57,7 +57,7 @@ class EmbeddingBert(nn.Embedding):
   def __call__(self, idx:Tensor) -> Tensor:
     if idx.numel() == 0: return Tensor.empty(idx.shape+(self.embed_sz,), dtype=self.weight.dtype, device=self.weight.device)
     arange_shp, weight_shp, big_shp = (1, 1, self.vocab_sz, 1), (1, 1, self.vocab_sz, self.embed_sz), idx.shape+(self.vocab_sz, self.embed_sz,)
-    if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz).reshape(arange_shp)
+    if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, requires_grad=False, device=self.weight.device).reshape(arange_shp)
     arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1,)).expand(big_shp), self.weight.cast(dtypes.default_float).reshape(weight_shp).expand(big_shp)
     return (arange == idx).where(vals, 0).sum(2, dtype=vals.dtype)
 
@@ -77,15 +77,15 @@ class FrozenBatchNorm2dRetinaNet(nn.BatchNorm2d):
   def __init__(self, sz:int, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1):
     self.eps, self.track_running_stats, self.momentum = eps, track_running_stats, momentum
 
-    self.weight = Tensor.ones(sz, dtype=dtypes.float32).is_param_(False) if affine else None
-    self.bias = Tensor.zeros(sz, dtype=dtypes.float32).is_param_(False) if affine else None
+    self.weight = Tensor.ones(sz, dtype=dtypes.float32, requires_grad=False) if affine else None
+    self.bias = Tensor.zeros(sz, dtype=dtypes.float32, requires_grad=False) if affine else None
 
-    if track_running_stats: self.running_mean, self.running_var = Tensor.zeros(sz, dtype=dtypes.float32).is_param_(False), Tensor.ones(sz, dtype=dtypes.float32).is_param_(False)
-    self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.long).is_param_(False)
+    if track_running_stats: self.running_mean, self.running_var = Tensor.zeros(sz, dtype=dtypes.float32, requires_grad=False), Tensor.ones(sz, dtype=dtypes.float32, requires_grad=False)
+    self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.long, requires_grad=False)
 
   def __call__(self, x:Tensor) -> Tensor:
     batch_mean, batch_var = super().calc_stats(x.cast(dtypes.float32))
-    if self.track_running_stats and TRAINING:
+    if self.track_running_stats and Tensor.training:
       self.running_mean.assign((1-self.momentum) * self.running_mean + self.momentum * batch_mean.detach().cast(self.running_mean.dtype))
       self.running_var.assign((1-self.momentum) * self.running_var + self.momentum * x.numel()/(x.numel()-x.shape[1]) * batch_var.detach().cast(self.running_var.dtype))
       self.num_batches_tracked += 1

examples/mlperf/model_eval.py

@@ -358,7 +358,7 @@ def eval_stable_diffusion():
       batch = batch.cat(batch[-1:].expand(bs - unpadded_bs, *batch[-1].shape))
     return batch, unpadded_bs 
 
-  @Context(TRAINING=0)
+  @Tensor.train(mode=False)
   def eval_unet(eval_inputs:list[dict], unet:UNetModel, cond_stage:FrozenOpenClipEmbedder, first_stage:AutoencoderKL,
                 inception:FidInceptionV3, clip:OpenClipEncoder) -> tuple[float, float]:
     # Eval is divided into 5 jits, one per model
@@ -498,10 +498,11 @@ def eval_stable_diffusion():
 
 if __name__ == "__main__":
   # inference only
+  Tensor.training = False
+
   models = getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert").split(",")
-  with Context(TRAINING=0):
-    for m in models:
-      nm = f"eval_{m}"
-      if nm in globals():
-        print(f"eval {m}")
-        globals()[nm]()
+  for m in models:
+    nm = f"eval_{m}"
+    if nm in globals():
+      print(f"eval {m}")
+      globals()[nm]()

examples/mlperf/model_spec.py

@@ -1,6 +1,6 @@
 # load each model here, quick benchmark
 from tinygrad import Tensor, GlobalCounters
-from tinygrad.helpers import getenv, Context
+from tinygrad.helpers import getenv
 import numpy as np
 
 def test_model(model, *inputs):
@@ -59,10 +59,11 @@ def spec_mrcnn():
 
 if __name__ == "__main__":
   # inference only for now
-  with Context(TRAINING=0):
-    for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,mrcnn").split(","):
-      nm = f"spec_{m}"
-      if nm in globals():
-        print(f"testing {m}")
-        globals()[nm]()
+  Tensor.training = False
+
+  for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,mrcnn").split(","):
+    nm = f"spec_{m}"
+    if nm in globals():
+      print(f"testing {m}")
+      globals()[nm]()
 

examples/mlperf/model_train.py

@@ -2,7 +2,7 @@ import os, time, math, functools, random, contextlib
 from pathlib import Path
 import multiprocessing
 
-from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes, Context
+from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
 from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling, profile_marker, DEBUG
 from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict, safe_load, safe_save
 from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW
@@ -157,7 +157,6 @@ def train_resnet():
   # input_std = Tensor([0.229, 0.224, 0.225], device=GPUS, dtype=dtypes.float32).reshape(1, -1, 1, 1)
   def normalize(x): return (x.permute([0, 3, 1, 2]) - input_mean).cast(dtypes.default_float)
   @TinyJit
-  @Context(TRAINING=1)
   def train_step(X, Y):
     optimizer_group.zero_grad()
     X = normalize(X)
@@ -171,7 +170,6 @@ def train_resnet():
     return loss.realize(), top_1.realize()
 
   @TinyJit
-  @Context(TRAINING=0)
   def eval_step(X, Y):
     X = normalize(X)
     out = model.forward(X)
@@ -182,11 +180,11 @@ def train_resnet():
   def fake_data_get(batch_size):
     x = Tensor.zeros(batch_size, 224, 224, 3, dtype=dtypes.uchar).contiguous()
     y = [0] * batch_size
-    return x.shard(GPUS, axis=0).realize(), Tensor(y).shard(GPUS, axis=0), y, None
+    return x.shard(GPUS, axis=0).realize(), Tensor(y, requires_grad=False).shard(GPUS, axis=0), y, None
 
   def data_get(it):
     x, y, cookie = next(it)
-    return x.shard(GPUS, axis=0).realize(), Tensor(y).shard(GPUS, axis=0), y, cookie
+    return x.shard(GPUS, axis=0).realize(), Tensor(y, requires_grad=False).shard(GPUS, axis=0), y, cookie
 
   # ** epoch loop **
   step_times = []
@@ -194,6 +192,7 @@ def train_resnet():
     # ** train loop **
     if MLLOGGER and RUNMLPERF:
       MLLOGGER.start(key=mllog_constants.EPOCH_START, value=e+1, metadata=dict(epoch_num=e+1))
+    Tensor.training = True
     BEAM.value = TRAIN_BEAM
 
     if INITMLPERF:
@@ -272,6 +271,7 @@ def train_resnet():
       eval_loss = 0.0
       eval_top_1 = 0
       eval_num_samples = 0
+      Tensor.training = False
       BEAM.value = EVAL_BEAM
 
       if INITMLPERF:
@@ -413,7 +413,7 @@ def train_retinanet():
     layers_to_train = ["layer4", "layer3", "layer2", "layer1", "conv1"][:trainable_layers]
     for k, v in get_state_dict(backbone).items():
       if all([not k.startswith(layer) for layer in layers_to_train]):
-        v.is_param_(False)
+        v.requires_grad = False
 
   def _data_get(it:Iterator[tuple[Tensor, ...]], val:bool=False):
     if val:
@@ -614,7 +614,7 @@ def train_retinanet():
 
       if getenv("RESET_STEP", 1): _train_step.reset()
 
-      with Context(TRAINING=0):
+      with Tensor.train(mode=False):
         if not RUNMLPERF:
           i, proc = 0, _fake_data_get(EVAL_BS, val=(val:=True))
         else:
@@ -784,7 +784,7 @@ def train_unet3d():
     return x.shard(GPUS, axis=0).realize(), y.shard(GPUS, axis=0), cookie
 
   @TinyJit
-  @Context(TRAINING=1)
+  @Tensor.train()
   def train_step(model, x, y):
     optim.zero_grad()
 
@@ -795,10 +795,10 @@ def train_unet3d():
     optim.step()
     return loss.realize()
 
-  @Context(TRAINING=0)
+  @Tensor.train(mode=False)
   def eval_step(model, x, y):
     y_hat, y = sliding_window_inference(model, x, y, gpus=GPUS)
-    y_hat, y = Tensor(y_hat), Tensor(y)
+    y_hat, y = Tensor(y_hat), Tensor(y, requires_grad=False)
     loss = dice_ce_loss(y_hat, y)
     score = dice_score(y_hat, y)
     return loss.realize(), score.realize()
@@ -919,7 +919,6 @@ def train_rnnt():
   pass
 
 @TinyJit
-@Context(TRAINING=0)
 def eval_step_bert(model, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor, masked_positions:Tensor, masked_lm_ids:Tensor,
                    masked_lm_weights:Tensor, next_sentence_labels:Tensor, GPUS):
   for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
@@ -1107,7 +1106,6 @@ def train_bert():
       MLLOGGER.start(key=mllog_constants.EPOCH_START, value=i*GBS, metadata={"epoch_num": i*GBS})
 
   @TinyJit
-  @Context(TRAINING=1)
   def train_step_bert(input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor,
                       masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor):
     for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
@@ -1135,6 +1133,7 @@ def train_bert():
 
   while train_data is not None and i < train_steps and not achieved:
     if getenv("TRAIN", 1):
+      Tensor.training = True
       BEAM.value = TRAIN_BEAM
       st = time.perf_counter()
       GlobalCounters.reset()
@@ -1187,6 +1186,7 @@ def train_bert():
       eval_lm_accs = []
       eval_clsf_accs = []
       eval_times = []
+      Tensor.training = False
       BEAM.value = EVAL_BEAM
 
       for j in tqdm(range(max_eval_steps), desc="Evaluating", total=max_eval_steps, disable=BENCHMARK):
@@ -1282,10 +1282,10 @@ def train_bert():
         previous_step = i
 
 def train_llama3():
-  from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8_DTYPE, MXFP8, MXFP4
+  from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8_DTYPE
   from examples.llama3 import MODEL_PARAMS
   from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
-  from examples.mlperf.optim import GradAccClipAdamW, clip_grads
+  from examples.mlperf.optim import GradAccClipAdamW
 
   INITMLPERF = getenv("INITMLPERF")
   RUNMLPERF = getenv("RUNMLPERF")
@@ -1419,7 +1419,7 @@ def train_llama3():
 
   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 = Tensor.zeros(p.shape, dtype=grad_dtype, device=p.device).contiguous()
   grads = [p.grad for p in optim.params]
 
   scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
@@ -1434,40 +1434,24 @@ def train_llama3():
     load_state_dict(scheduler, safe_load(fn), realize=False)
 
   fp8_amax = [t for ts in model._fp8_amax.values() for t in ts]
-  fp8_next_amax = [t for ts in model._fp8_next_amax.values() for t in ts]
-  fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts]
-  fp8_next_grad_amax = [t for ts in model._fp8_next_grad_amax.values() for t in ts]
-  fp8_inv_scales = list(model._fp8_inv_scale.values()) + list(model._fp8_next_inv_scale.values())
+  fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts] if hasattr(model, "_fp8_grad_amax") else []
+  fp8_inv_scales = list(model._fp8_inv_scale.values())
 
   from tinygrad.nn.state import get_state_dict
   model_state = get_state_dict(model)
-  for wname in model._fp8_inv_scale:
+  for wname in ["wqkv", "wo", "w13", "w2"]:
     w = model_state[wname]
     w._inv_scale = model._fp8_inv_scale[wname]
-    w._next_inv_scale = model._fp8_next_inv_scale[wname]
     if optim.master_params:
       idx = next(j for j, p in enumerate(optim.params) if p is w)
-      master = optim.master_params[idx]
-      inv = w._inv_scale if w._inv_scale.device == master.device else w._inv_scale.to(master.device)
-      if MXFP8:
-        from extra.gemm.cdna_asm_gemm import _mx_block_scale
-        bs = _mx_block_scale(inv.reshape(-1, inv.shape[-1])).reshape(w.shape)
-        master.assign((master * bs).contiguous())
-      else:
-        master.assign((master * inv.reshape(*inv.shape, *([1]*(w.ndim-inv.ndim)))).contiguous())
-
-  # realize everything here
-  if optim.master_params: Tensor.realize(*optim.master_params)
-  loss_acc = Tensor.zeros(1, dtype=dtypes.float32, device=device)
-  Tensor.realize(loss_acc, *optim.params, *fp8_inv_scales, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)
+      optim.master_params[idx].assign((optim.master_params[idx] * w._inv_scale.reshape(-1, *([1]*(w.ndim-1)))).contiguous())
 
   @TinyJit
   def minibatch(tokens:Tensor):
-    model.reset_amax()
     if is_dp: tokens = tokens.to(None).shard(device, 0)
     if is_mp: tokens = tokens.shard(device)
     if not is_sharding: tokens = tokens.to(None)
-    logits:Tensor = model(tokens[:, :-1], save=bool(SMALL))
+    logits:Tensor = model(tokens[:, :-1])
     if getenv("FAST_CE", 0):
       from extra.llama_kernels.fused_ce import fused_ce_loss
       loss = fused_ce_loss(logits.cast(dtypes.bfloat16), tokens[:, 1:], label_smoothing=0.0)
@@ -1477,27 +1461,24 @@ def train_llama3():
     for g, new_g in zip(grads, loss.gradient(*optim.params)):
       apply_grad(g, new_g.uop)
 
-    loss_acc.assign(loss_acc + loss.flatten().float())
-    return loss_acc.realize(*grads, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)
+    loss_cpu = loss.flatten().float().to("CPU")
+    return loss_cpu.realize(*grads, *fp8_amax, *fp8_grad_amax)
 
   @TinyJit
   def optim_step():
-    grad_norm = clip_grads(grads, grad_acc, 1.0)
-    optim.fstep(grads, grad_norm)
+    grad_norm = optim.fstep(grads)
     scheduler.step()
 
-    for g in grads: g.assign(0)
-    model.update_amax()
+    for g in grads: g.assign(g.zeros_like())
 
     lr_cpu = optim.lr.float().to("CPU")
     grad_norm_cpu = grad_norm.float().to("CPU")
-    loss_cpu = loss_acc.to("CPU")
-    Tensor.realize(lr_cpu, grad_norm_cpu, loss_cpu, loss_acc.assign(0), *grads, *fp8_inv_scales, *fp8_amax, *fp8_grad_amax)
+    Tensor.realize(lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales)
 
-    return lr_cpu, grad_norm_cpu, loss_cpu
+    return lr_cpu, grad_norm_cpu
 
   @TinyJit
-  @Context(TRAINING=0)
+  @Tensor.train(False)
   def eval_step(tokens:Tensor):
     if is_dp: tokens = tokens.to(None).shard(device, 0)
     if is_mp: tokens = tokens.shard(device)
@@ -1510,7 +1491,7 @@ def train_llama3():
   def fake_data(bs, samples):
     import numpy as np
     for _ in range(samples // bs):
-      fake_data_np = np.random.randint(0, real_vocab_size, size=(bs, SEQLEN + 1), dtype=np.int32)
+      fake_data_np = np.random.randint(0, model_params["vocab_size"], size=(bs, SEQLEN + 1), dtype=np.int32)
       yield Tensor(fake_data_np, device="NPY")
 
   def get_train_iter():
@@ -1549,8 +1530,8 @@ def train_llama3():
       st = time.perf_counter()
 
       stopped = False
-      data_time, dev_time = 0, 0
-      for _ in range(accum_steps:=grad_acc if i >= 2 else 1):
+      losses, data_time, dev_time = [], 0, 0
+      for _ in range(grad_acc if i >= 2 else 1):
         ist = time.perf_counter()
         try: tokens = next(train_iter)
         except StopIteration:
@@ -1558,15 +1539,16 @@ def train_llama3():
           break
         mst = time.perf_counter()
         data_time += mst - ist
-        minibatch(tokens)
+        losses.append(minibatch(tokens).item())
         dev_time += time.perf_counter() - mst
       if stopped: break
 
       gt = time.perf_counter()
       ret = optim_step()
-      lr, grad_norm, loss = ret[0].item(), ret[1].item(), ret[2].item() / accum_steps
+      lr, grad_norm = ret[0].item(), ret[1].item()
       et = time.perf_counter()
 
+      loss = sum(losses) / len(losses)
       optim_time = et - gt
       dev_time += optim_time
       step_time = et - st
@@ -1578,7 +1560,7 @@ def train_llama3():
 
       mem_gb = GlobalCounters.mem_used / 1e9
       gflops = GlobalCounters.global_ops / 1e9 / dev_time
-      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * (9.2e15 if MXFP4 else 4.6e15))) * 100
+      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * 4.6e15)) * 100
       tqdm.write(
           f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
           f"{lr:.12f} LR, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
@@ -1664,287 +1646,6 @@ def train_llama3():
       if MLLOGGER and RUNMLPERF:
         MLLOGGER.start(key=mllog_constants.BLOCK_START, metadata={mllog_constants.SAMPLES_COUNT: sequences_seen})
 
-def train_gptoss():
-  from examples.mlperf.models.gpt_oss import GPTOSS, GPT_OSS_20B, apply_grad, FP8_DTYPE
-  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
-  from examples.mlperf.optim import GradAccClipAdamW, GradAccClipAdamWGroup, clip_grads
-
-  BENCHMARK = getenv("BENCHMARK")
-
-  config = {}
-  BASEDIR            = config["BASEDIR"]                = Path(getenv("BASEDIR", "/raid/datasets/c4-8b/"))
-  BS                 = config["BS"]                     = getenv("BS", 16)
-  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
-  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
-  SEED               = config["SEED"]                   = getenv("SEED", 5760)
-  DATA_SEED          = config["DATA_SEED"]              = getenv("DATA_SEED", SEED)
-  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
-  TRAIN_ON_VAL       = config["TRAIN_ON_VAL"]           = getenv("TRAIN_ON_VAL", 0)
-  MAX_STEPS          = config["MAX_STEPS"]              = getenv("MAX_STEPS", 1_200_000)
-  SAMPLES            = config["SAMPLES"]                = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else MAX_STEPS * GBS)
-  EVAL_SAMPLES       = config["EVAL_SAMPLES"]           = getenv("EVAL_SAMPLES", 1024)
-  WARMUP_STEPS       = config["WARMUP_STEPS"]           = getenv("WARMUP_STEPS", 128)
-  LR                 = config["LR"]                     = getenv("LR", 4e-4 * GBS / 16)
-  END_LR             = config["END_LR"]                 = getenv("END_LR", 4e-5)
-  EVAL_FREQ          = config["EVAL_FREQ"]              = getenv("EVAL_FREQ", 12288)
-  EVAL_BS            = config["EVAL_BS"]                = getenv("EVAL_BS", 16)
-  EVAL_TARGET        = config["EVAL_TARGET"]            = getenv("EVAL_TARGET", 3.34)
-
-  opt_adamw_beta_1 = 0.9
-  opt_adamw_beta_2 = 0.95
-  opt_adamw_epsilon = 1e-5
-  opt_adamw_weight_decay = 0.1
-
-  opt_learning_rate_warmup_steps = WARMUP_STEPS
-  opt_learning_rate_decay_steps = MAX_STEPS - opt_learning_rate_warmup_steps
-  opt_base_learning_rate = LR
-  opt_end_learning_rate = END_LR
-
-  Tensor.manual_seed(SEED)  # seed for weight initialization
-
-  # ** init wandb **
-  WANDB = getenv("WANDB")
-  if WANDB:
-    import wandb
-    wandb_args = {"id": wandb_id, "resume": "must"} if (wandb_id := getenv("WANDB_RESUME", "")) else {}
-    wandb.init(config=config, **wandb_args, project="MLPerf-gpt-oss")
-
-  model_params = GPT_OSS_20B
-  model_params['vocab_size'] = getenv("VOCAB_SIZE", 128256)
-  real_vocab_size = model_params['vocab_size']
-  if (layers:=getenv("LAYERS")) != 0: model_params['n_layers'] = layers
-  if (experts:=getenv("EXPERTS")) != 0: model_params['n_experts'] = experts
-  print(f"model parameters: {model_params}")
-
-  model = GPTOSS(**model_params, max_context=SEQLEN)
-
-  params = get_parameters(model)
-
-  if getenv("EMPTYWEIGHT"):
-    for v in get_parameters(model):
-      v = v.assign(Tensor.empty(v.shape, dtype=v.dtype))
-
-  is_dp = (DP := getenv("DP", 1)) > 1
-  is_sharding = is_dp
-  device_count = DP
-  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))
-
-  model.shard(device, False)
-
-  is_offload_optim = bool(getenv("OFFLOAD_OPTIM"))
-  is_fake_offload = Device.DEFAULT == "NULL"
-  optim_device = ("CPU" if not is_fake_offload else "NULL:99") if is_offload_optim else None
-  params_wd = [p for p in params if p.ndim >= 3]
-  params_no_wd = [p for p in params if p.ndim < 3]
-  optim = GradAccClipAdamWGroup(
-    GradAccClipAdamW(params_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device),
-    GradAccClipAdamW(params_no_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=0.0, grad_acc=grad_acc, device=optim_device),
-  )
-
-  for p in optim.params:
-    grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
-    p.grad = p.zeros_like(dtype=grad_dtype).contiguous()
-  grads = [p.grad for p in optim.params]
-
-  from extra.gemm.cdna_asm_gemm import _mx_block_scale
-  model_state = get_state_dict(model)
-  def _scale_key(n):
-    if "." in n and (c:=f"{(b:=n.rsplit('.',1))[0]}_scale.{b[1]}") in model_state: return c
-    return f"{n}_scale"
-  fp8_scale_names = {n: _scale_key(n) for n, t in model_state.items() if t.dtype == FP8_DTYPE}
-  fp8_inv_scales = [model_state[sname] for sname in fp8_scale_names.values()]
-  for wname, sname in fp8_scale_names.items():
-    w, scale = model_state[wname], model_state[sname]
-    w._inv_scale = scale
-    if optim.master_params:
-      master = optim.master_params[next(j for j, p in enumerate(optim.params) if p is w)]
-      inv = scale if scale.device == master.device else scale.to(master.device)
-      bs = _mx_block_scale(inv.reshape(-1, inv.shape[-1])).reshape(w.shape)
-      master.assign((master * bs).contiguous())
-
-  scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
-
-  if optim.master_params:
-    for m in optim.master_params: m.realize()
-  Tensor.realize(*optim.params, *fp8_inv_scales)
-
-  @TinyJit
-  @Context(TRAINING=1)
-  def minibatch(tokens:Tensor):
-    if is_dp: tokens = tokens.to(None).shard(device, 0)
-    if not is_sharding: tokens = tokens.to(None)
-    logits:Tensor = model(tokens[:, :-1], save=True)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-
-    for g, new_g in zip(grads, loss.gradient(*optim.params)):
-      apply_grad(g, new_g.uop)
-
-    loss_cpu = loss.flatten().float().to("CPU")
-    return loss_cpu.realize(*grads)
-
-  @TinyJit
-  def optim_step():
-    grad_norm = clip_grads(grads, grad_acc, 1.0)
-    optim.fstep(grads, grad_norm)
-    scheduler.step()
-
-    for g in grads: g.assign(0)
-
-    lr_cpu = optim.lr.float().to("CPU")
-    grad_norm_cpu = grad_norm.float().to("CPU")
-    Tensor.realize(lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales)
-
-    return lr_cpu, grad_norm_cpu
-
-  @TinyJit
-  @Context(TRAINING=0)
-  def eval_step(tokens:Tensor):
-    if is_dp: tokens = tokens.to(None).shard(device, 0)
-    if not is_sharding: tokens = tokens.to(None)
-    logits:Tensor = model(tokens[:, :-1])
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-    return loss.flatten().float().to("CPU")
-
-  # ** data iters **
-  def fake_data(bs, samples):
-    import numpy as np
-    for _ in range(samples // bs):
-      fake_data_np = np.random.randint(0, real_vocab_size, size=(bs, SEQLEN + 1), dtype=np.int32)
-      yield Tensor(fake_data_np, device="NPY")
-
-  def get_train_iter():
-    if getenv("FAKEDATA", 0):
-      return fake_data(BS, SAMPLES)
-    else:
-      from examples.mlperf.dataloader import batch_load_llama3
-      return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=DATA_SEED, val=bool(TRAIN_ON_VAL), small=True)
-
-  if getenv("FAKEDATA", 0):
-    eval_dataset = None
-  else:
-    from examples.mlperf.dataloader import get_llama3_dataset
-    eval_dataset = get_llama3_dataset(EVAL_SAMPLES, SEQLEN, BASEDIR, val=True, small=True)
-
-  def get_eval_iter():
-    if eval_dataset is None:
-      return fake_data(EVAL_BS, EVAL_SAMPLES)
-    from examples.mlperf.dataloader import iterate_llama3_dataset
-    return iterate_llama3_dataset(eval_dataset, EVAL_BS)
-
-  num_params = sum(p.numel() for p in params) - model_params["vocab_size"]*model_params["dim"]
-  train_iter = get_train_iter()
-  i, sequences_seen = 0, 0
-  step_times = []
-
-  while i < MAX_STEPS:
-    GlobalCounters.reset()
-    actual_gbs = GBS if i >= 2 else BS
-    if getenv("TRAIN", 1):
-      profile_marker(f"train @ {i}")
-      st = time.perf_counter()
-
-      stopped = False
-      losses, data_time, dev_time = [], 0, 0
-      for _ in range(grad_acc if i >= 2 else 1):
-        ist = time.perf_counter()
-        try: tokens = next(train_iter)
-        except StopIteration:
-          stopped = True
-          break
-        mst = time.perf_counter()
-        data_time += mst - ist
-        losses.append(minibatch(tokens).item())
-        dev_time += time.perf_counter() - mst
-      if stopped: break
-
-      gt = time.perf_counter()
-      ret = optim_step()
-      lr, grad_norm = ret[0].item(), ret[1].item()
-      et = time.perf_counter()
-
-      loss = sum(losses) / len(losses)
-      optim_time = et - gt
-      dev_time += optim_time
-      step_time = et - st
-      gbs_time = gt - st
-      if BENCHMARK: step_times.append(step_time)
-
-      i += 1
-      sequences_seen += actual_gbs
-
-      mem_gb = GlobalCounters.mem_used / 1e9
-      gflops = GlobalCounters.global_ops / 1e9 / dev_time
-      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * 4.6e15)) * 100
-      tqdm.write(
-          f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
-          f"{lr:.12f} LR, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
-      if DEBUG >= 1: tqdm.write("  mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
-
-      if WANDB:
-        wandb.log({
-          "train/loss": loss,
-          "train/lr": lr,
-          "train/grad_norm": grad_norm,
-          "train/step_time": step_time,
-          "train/gbs_time": gbs_time,
-          "train/optim_time": optim_time,
-          "train/dev_time": dev_time,
-          "train/data_time": data_time,
-          "train/mem": mem_gb,
-          "train/GFLOPS": gflops,
-          "train/MFU": mfu,
-          "train/sequences_seen": sequences_seen
-        })
-
-      if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
-        tqdm.write("saving checkpoint")
-        if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
-        fn = f"{ckpt_dir}/gptoss_{i}.safe"
-        safe_save(get_state_dict(model), fn)
-
-        tqdm.write("saving optim checkpoint")
-        fn = f"{ckpt_dir}/gptoss_{i}_optim.safe"
-        safe_save(get_state_dict(scheduler), fn)
-
-      if i == BENCHMARK:
-        median_step_time = sorted(step_times)[BENCHMARK // 2]
-        estimated_steps = MAX_STEPS
-        estimated_total_minutes = int(median_step_time * estimated_steps / 60)
-        print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
-        print(f"epoch global_ops: {GlobalCounters.global_ops:_}, "
-              f"epoch global_mem: {GlobalCounters.global_mem:_}")
-
-    if (sequences_seen // EVAL_FREQ != (sequences_seen - actual_gbs) // EVAL_FREQ and (i != 1 or EVAL_FREQ == 1)) or (BENCHMARK and i == BENCHMARK):
-      if EVAL_BS == 0: return
-      tqdm.write(f"evaluating after {sequences_seen} sequences")
-      profile_marker(f"eval @ {i}")
-
-      # run eval
-      eval_losses = []
-      eval_iter = get_eval_iter()
-      tqdm.write(f"evaluating {EVAL_SAMPLES//EVAL_BS} batches of {EVAL_BS} sequences")
-
-      for j,tokens in tqdm(enumerate(eval_iter), total=EVAL_SAMPLES//EVAL_BS):
-        eval_losses += eval_step(tokens).tolist()
-
-        if BENCHMARK and (j+1) == min(BENCHMARK, EVAL_SAMPLES//EVAL_BS):
-          return
-
-      log_perplexity = sum(eval_losses) / len(eval_losses)
-
-      tqdm.write(f"eval log perplexity: {log_perplexity:.4f}")
-
-      if WANDB:
-        wandb.log({"eval/log_perplexity": log_perplexity, "eval/sequences_seen": sequences_seen})
-
-      if log_perplexity < EVAL_TARGET:
-        tqdm.write(f"target achieved after {sequences_seen} sequences")
-        if getenv("CKPT"):
-          if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
-          fn = f"{ckpt_dir}/gptoss.safe"
-          safe_save(get_state_dict(model), fn)
-        break
-
 def train_stable_diffusion():
   from extra.models.unet import UNetModel
   from examples.mlperf.dataloader import batch_load_train_stable_diffusion
@@ -2023,7 +1724,7 @@ def train_stable_diffusion():
     # move to CPU first so more GPU bufs aren't created (can trigger OOM)
     for k,v in ckpt.items(): ckpt[k] = v.detach().to("CPU")
     Tensor.realize(*[v for v in ckpt.values()])
-    for k,v in ckpt.items(): ckpt[k] = v.cast(v.dtype).contiguous()
+    for k,v in ckpt.items(): ckpt[k] = v.cast(v.dtype.base).contiguous()
     Tensor.realize(*[v for v in ckpt.values()])
     return ckpt
 
@@ -2090,7 +1791,7 @@ if __name__ == "__main__":
   elif getenv("RUNMLPERF"): bench_log_manager = WallTimeEvent(BenchEvent.MLPERF_RUN)
   else: bench_log_manager = contextlib.nullcontext()
 
-  with Context(TRAINING=1):
+  with Tensor.train():
     for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,maskrcnn,stable_diffusion").split(","):
       nm = f"train_{m}"
       if nm in globals():

examples/mlperf/models/flat_llama.py

@@ -2,8 +2,9 @@ import math, os
 if __name__ == "__main__":
   os.environ["DEFAULT_FLOAT"] = "bfloat16"
   os.environ["OPTIM_DTYPE"] = "bfloat16"
-  if "DEV" not in os.environ: os.environ["DEV"] = "NULL::gfx950"
+  if "DEV" not in os.environ: os.environ["DEV"] = "NULL"
   # CDNA
+  os.environ["EMULATE"] = "AMD_CDNA4"
   os.environ["DEVICE_IN_FUNCTION_BUG"] = "1"
   os.environ["ALL2ALL"] = "1"
   os.environ["USE_ATOMICS"] = "1"
@@ -12,7 +13,7 @@ if __name__ == "__main__":
     if "ASM_GEMM" not in os.environ:
       os.environ["ASM_GEMM"] = "1"
 from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
-from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker, round_up
+from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
 from extra.llama_kernels.rmsnorm import rmsnorm
@@ -22,10 +23,6 @@ ASM_GEMM = getenv("ASM_GEMM", 0)
 FUSED_INPUT_QUANTIZE = getenv("FUSED_INPUT_QUANTIZE", 0)
 FUSED_ADD_NORM_MUL_QUANTIZE = getenv("FUSED_ADD_NORM_MUL_QUANTIZE", 0)
 FUSED_SILU_W13 = getenv("FUSED_SILU_W13", 0)
-SPLIT_W13 = getenv("SPLIT_W13", 0)
-COLUMNWISE_WEIGHT_SCALE = getenv("COLUMNWISE_WEIGHT_SCALE", 0)
-MXFP8 = getenv("MXFP8", 0)
-MXFP4 = getenv("MXFP4", 0)
 
 FP8_DTYPE = dtypes.fp8e4m3
 FP8_GRAD_DTYPE = dtypes.fp8e5m2
@@ -38,98 +35,58 @@ def quantize_fp8(x:Tensor, amax_state:Tensor|None=None):
   return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal(), new_amax
 
 def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_scale:Tensor|None=None,
-           x_fp8:Tensor|None=None, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None, x_prequant_mx:tuple|None=None,
-           next_amax_x:Tensor|None=None) -> tuple[Tensor,...]:
+           x_fp8:Tensor|None=None, x_scale:Tensor|None=None, x_new_amax:Tensor|None=None,
+           grad_amax_state:Tensor|None=None) -> tuple[Tensor,...]:
   if not fp8:
     if ASM_GEMM:
       from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
       if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T),)
     return (x @ w.T,)
-  if MXFP4:
-    assert x is not None, "MXFP4 matmul requires an unquantized input"
-    from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm
-    if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T, mxfp4=True),)
-    return (x @ w.T,)
   assert w_inv_scale is not None, "fp8 matmul requires w_inv_scale (weights must be stored in fp8 with per-tensor scale)"
-  if MXFP8:
-    from extra.gemm.cdna_asm_gemm import asm_gemm, quantize_mxfp8, mx_pack, can_use_asm_gemm, _mx_block_scale
-    if x_prequant_mx is not None: x_q, x_e8, x_si = x_prequant_mx       # fused producer already quantized (2d)
-    else: x_q, x_e8, x_si = quantize_mxfp8(x.reshape(-1, x.shape[-1]))
-    l_shape = x.shape[:-1] if x is not None else x_q.shape[:-1]
-    if can_use_asm_gemm(x_q, w.T):
-      out = asm_gemm(x_q, w.T, mx=True, mx_scales=(x_si, x_e8, mx_pack(w_inv_scale), w_inv_scale),
-                     mx_w_stored=True).reshape(*l_shape, w.shape[0])
-    else:
-      x_phys = (x_q.cast(dtypes.bfloat16) * _mx_block_scale(x_e8)).reshape(*l_shape, x_q.shape[-1])
-      out = x_phys @ (w.cast(dtypes.bfloat16) * _mx_block_scale(w_inv_scale)).T
-    return out, x_q
   if x_fp8 is None:
-    if FUSED_INPUT_QUANTIZE:
+    if FUSED_INPUT_QUANTIZE and amax_x is not None:
       from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed
-      x_fp8, _ = quantize_fp8_delayed(x, amax_x, next_amax_x, FP8_DTYPE)
+      x_fp8, x_scale, x_new_amax, _ = quantize_fp8_delayed(x, amax_x, FP8_DTYPE)
     else:
-      x_fp8, _, new_amax_x = quantize_fp8(x, amax_state=amax_x)
-      next_amax_x.assign(new_amax_x)
+      x_fp8, x_scale, x_new_amax = quantize_fp8(x, amax_state=amax_x)
   if ASM_GEMM:
     from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
     if can_use_asm_gemm(x_fp8, w.T):
-      assert amax_x is not None
-      if COLUMNWISE_WEIGHT_SCALE:
-        out = asm_gemm(x_fp8, w.T, x_scale=amax_x, grad_amax_state=grad_amax_state,
-                       next_grad_amax_state=next_grad_amax_state, w_post_scale=w_inv_scale)
-      else:
-        out = asm_gemm(x_fp8, w.T, x_scale=amax_x, w_scale=w_inv_scale, grad_amax_state=grad_amax_state,
-                       next_grad_amax_state=next_grad_amax_state)
-      return out, x_fp8
-  return (x_fp8.dot(w.T, dtype=dtypes.float) * ((amax_x.float() + 1e-8) / FP8_MAX) * w_inv_scale).cast(dtypes.bfloat16), x_fp8
+      return asm_gemm(x_fp8, w.T, x_scale=x_scale, w_scale=w_inv_scale, grad_amax_state=grad_amax_state), x_new_amax, x_fp8, w
+  return x_fp8.dot(w.T, dtype=dtypes.float) * x_scale * w_inv_scale, x_new_amax, x_fp8, w
 
-def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor|None,
-                         next_amax_x:Tensor|None, grad_amax_state:Tensor|None, next_grad_amax_state:Tensor|None):
-  if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
+def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor, grad_amax_state:Tensor):
+  if FUSED_ADD_NORM_MUL_QUANTIZE:
     from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_rmsnorm_mul_quantize_fp8
-    x_fp8, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
-    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
-                       grad_amax_state=grad_amax_state, next_grad_amax_state=next_grad_amax_state)
+    x_fp8, x_inv_scale, new_amax, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE)
+    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, x_scale=x_inv_scale, x_new_amax=new_amax, grad_amax_state=grad_amax_state)
     return out, x_normed, rrms, ret
   x_normed, rrms = rmsnorm(x, eps)
-  out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state,
-                     next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
+  out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state)
   return out, x_normed, rrms, ret
 
-def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor|None,
-                             next_amax_x:Tensor|None, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
-  if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
+def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor):
+  if FUSED_ADD_NORM_MUL_QUANTIZE:
     from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_add_rmsnorm_mul_quantize_fp8
-    x_fp8, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
-    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
-                       grad_amax_state=grad_amax_state, next_grad_amax_state=next_grad_amax_state)
+    x_fp8, x_inv_scale, new_amax, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE)
+    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, x_scale=x_inv_scale, x_new_amax=new_amax)
     return out, h, x_normed, rrms, ret
   h = x + residual
   x_normed, rrms = rmsnorm(h, eps)
-  out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state,
-                     next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
+  out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale)
   return out, h, x_normed, rrms, ret
 
 def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor,
-                             amax_x2:Tensor|None, next_amax_x2:Tensor|None,
-                             grad_amax_xw13:Tensor|None, next_grad_amax_xw13:Tensor|None,
-                             grad_amax_xout:Tensor|None, next_grad_amax_xout:Tensor|None):
-  if FUSED_SILU_W13 and MXFP4:
-    from extra.llama_kernels.swiglu import swiglu
-    out, *ret = matmul(swiglu(x_w13), w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout,
-                       next_grad_amax_state=next_grad_amax_xout, next_amax_x=next_amax_x2)
-    return out, ret
-  if FUSED_SILU_W13 and not MXFP4:
+                             amax_x2:Tensor,
+                             grad_amax_xw13:Tensor, grad_amax_xout:Tensor):
+  if FUSED_SILU_W13:
     from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
-    x2_fp8 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13,
-                                                 next_grad_amax_state=next_grad_amax_xw13, amax_out=next_amax_x2)
-    out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, amax_x=amax_x2,
-                       grad_amax_state=grad_amax_xout, next_grad_amax_state=next_grad_amax_xout)
+    x2_fp8, x2_inv_scale, new_amax_x2 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13)
+    out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, x_scale=x2_inv_scale, x_new_amax=new_amax_x2, grad_amax_state=grad_amax_xout)
     return out, ret
   hidden = x_w13.shape[-1] // 2
   x_w1, x_w3 = x_w13[..., :hidden], x_w13[..., hidden:]
-  out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout,
-                     next_grad_amax_state=next_grad_amax_xout, next_amax_x=next_amax_x2)
+  out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout)
   return out, ret
 
 class FlatTransformer:
@@ -146,16 +103,13 @@ class FlatTransformer:
     scaled_std = 0.02 / math.sqrt(2 * n_layers)
 
     # Attention
-    self.wqkv, s_qkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
-    self.wo, s_o = self.lin_per_layer(self.n_heads * self.head_dim, dim, std=scaled_std)
+    self._init_inv_scales = []  # populated by lin_per_layer
+    self.wqkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
+    self.wo = self.lin_per_layer(self.n_heads * self.head_dim, dim, std=scaled_std)
 
     # FeedForward
-    if SPLIT_W13:
-      self.w1, s_1 = self.lin_per_layer(dim, hidden_dim)
-      self.w3, s_3 = self.lin_per_layer(dim, hidden_dim)
-    else:
-      self.w13, s_13 = self.lin_per_layer(dim, hidden_dim * 2)
-    self.w2, s_2 = self.lin_per_layer(hidden_dim, dim, std=scaled_std)
+    self.w13 = self.lin_per_layer(dim, hidden_dim * 2)
+    self.w2 = self.lin_per_layer(hidden_dim, dim, std=scaled_std)
 
     self.norm_eps = norm_eps
     self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
@@ -166,124 +120,93 @@ class FlatTransformer:
     self.tok_embeddings = nn.Embedding(vocab_size, dim)
     self.tok_embeddings.weight = Tensor.normal(vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
     self.output = Tensor.normal(1, vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
-    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).clone().is_param_(False)
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)
 
-    def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().is_param_(False)
-    n_amax = 0 if MXFP4 else n_layers
-    names = ["xqkv", "xo", "x2"]
-    names += ["x1", "x3"] if SPLIT_W13 else ["x13"]
-    self._fp8_amax = {name: [_amax() for _ in range(n_amax)] for name in names}
-    self._fp8_next_amax = {name: [_amax() for _ in range(n_amax)] for name in names}
-    grad_names = ["xqkv", "xo", "xout"]
-    grad_names += ["xw1", "xw3"] if SPLIT_W13 else ["xw13"]
-    self._fp8_grad_amax = {name: [_amax() for _ in range(n_amax)] for name in grad_names}
-    self._fp8_next_grad_amax = {name: [_amax() for _ in range(n_amax)] for name in grad_names}
-    w_scales = [("wqkv", s_qkv), ("wo", s_o), ("w2", s_2)]
-    w_scales += [("w1", s_1), ("w3", s_3)] if SPLIT_W13 else [("w13", s_13)]
-    self._fp8_inv_scale = {name: (s if MXFP8 else s.float()).contiguous().is_param_(False) for name, s in w_scales}
-    self._fp8_next_inv_scale = {name: (s if MXFP8 else s.float()).contiguous().is_param_(False) for name, s in w_scales}
+    def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().requires_grad_(False)
+    names = ["xqkv", "xo", "x13", "x2"]
+    self._fp8_amax = {name: [_amax() for _ in range(n_layers)] for name in names}
+    grad_names = ["xqkv", "xo", "xw13", "xout"]
+    self._fp8_grad_amax = {name: [_amax() for _ in range(n_layers)] for name in grad_names}
+    w_names = ["wqkv", "wo", "w13", "w2"]
+    self._fp8_inv_scale = {wname: inv_scales.float().contiguous().requires_grad_(False)
+                           for wname, inv_scales in zip(w_names, self._init_inv_scales)}
+    del self._init_inv_scales
 
-  def lin_per_layer(self, in_features:int, out_features:int, std:float=0.02, w:Tensor|None=None):
-    if w is None:
-      if getenv("ZEROS"): w = Tensor.zeros(self.n_layers, out_features, in_features)
-      else: w = Tensor.normal(self.n_layers, out_features, in_features, mean=0.0, std=std)
-    if MXFP8:
-      from extra.gemm.cdna_asm_gemm import quantize_mxfp8
-      w_q, w_e8, _ = quantize_mxfp8(w.reshape(self.n_layers * out_features, in_features))
-      return w_q.reshape(self.n_layers, out_features, in_features), w_e8.reshape(self.n_layers, out_features, in_features // 32)
-    if MXFP4:
-      # FP4 is produced dynamically so optimizer updates always start from the current BF16 weight.
-      return w.cast(dtypes.bfloat16), Tensor.ones(self.n_layers)
-    amax = (w.abs().max(axis=2) if COLUMNWISE_WEIGHT_SCALE else w.abs().flatten(1).max(1)).detach()
+  def lin_per_layer(self, in_features:int, out_features:int, std:float=0.02):
+    if getenv("ZEROS"): w = Tensor.zeros(self.n_layers, out_features, in_features)
+    else: w = Tensor.normal(self.n_layers, out_features, in_features, mean=0.0, std=std)
+    amax = w.abs().flatten(1).max(1).detach()
     scale = FP8_MAX / (amax + 1e-8)
-    inv_scale = (amax + 1e-8) / FP8_MAX
-    scale_b = scale.reshape(self.n_layers, out_features, 1) if COLUMNWISE_WEIGHT_SCALE else scale.reshape(-1, 1, 1)
-    return (w * scale_b).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE), inv_scale
+    self._init_inv_scales.append((amax + 1e-8) / FP8_MAX)
+    return (w * scale.reshape(-1, 1, 1)).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE)
 
-  def attention(self, x:Tensor, freqs_cis:Tensor, *, attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
-                amax_xqkv:Tensor|None, amax_xo:Tensor|None, s_qkv:Tensor, s_o:Tensor,
-                next_amax_xqkv:Tensor|None, next_amax_xo:Tensor|None,
-                grad_amax_xqkv:Tensor|None, grad_amax_xo:Tensor|None,
-                next_grad_amax_xqkv:Tensor|None, next_grad_amax_xo:Tensor|None):
+  def attention(self, x:Tensor, freqs_cis:Tensor, attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
+                amax_xqkv:Tensor, amax_xo:Tensor, s_qkv:Tensor, s_o:Tensor,
+                grad_amax_xqkv:Tensor, grad_amax_xo:Tensor):
     bsz, seqlen, _ = x.shape
-    saves = []
+    new_amaxs, saves = [], []
 
-    xqkv, x_normed, rrms, s = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
-                                                                  amax_x=amax_xqkv, grad_amax_state=grad_amax_xqkv,
-                                                                  next_grad_amax_state=next_grad_amax_xqkv, next_amax_x=next_amax_xqkv)
-    saves.extend([x_normed, rrms, *s, xqkv])
+    xqkv, x_normed, rrms, ret = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
+                                                     amax_x=amax_xqkv, grad_amax_state=grad_amax_xqkv)
+    saves.extend([x_normed, rrms])
+    new_amaxs.extend(ret[:1])
+    saves.extend(ret[1:] + [xqkv])
+    xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
+    xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
+    xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
+    xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
+
+    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
+    xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
+    xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
     if getenv("HK_FLASH_ATTENTION"):
-      from extra.thunder.amd.fa import flash_attention, fused_qkv_rope
-      xq, xk, xv = fused_qkv_rope(xqkv, freqs_cis, self.n_heads, self.n_kv_heads, self.head_dim)
-      attn, *save = flash_attention(xq, xk, xv, is_causal=True, write_flat=True)
+      from extra.thunder.amd.fa import flash_attention
+      attn, *save = flash_attention(xq, xk, xv, is_causal=True)
       saves.extend(save)
     else:
-      xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
-      xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
-      xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
-      xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
-      xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
-      xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
-      xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
-      attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
-    attn = attn.reshape(bsz, seqlen, -1)
+      attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True)
+    attn = attn.transpose(1, 2).reshape(bsz, seqlen, -1)
 
-    out, *s = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo,
-                               next_grad_amax_state=next_grad_amax_xo, next_amax_x=next_amax_xo)
-    saves.extend([*s, out])
-    return out, saves
+    out, *ret = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo)
+    new_amaxs.extend(ret[:1])
+    saves.extend(ret[1:] + [out])
+    return (out, *new_amaxs, *saves)
 
-  def feed_forward(self, x:Tensor, residual:Tensor, **kwargs):
-    saves = []
+  def feed_forward(self, x:Tensor, residual:Tensor, ffn_norm:Tensor, w13:Tensor, w2:Tensor,
+                   amax_x13:Tensor, amax_x2:Tensor, s_13:Tensor, s_2:Tensor,
+                   grad_amax_xw13:Tensor, grad_amax_xout:Tensor):
+    new_amaxs, saves = [], []
 
-    if SPLIT_W13:
-      h = x + residual
-      x_normed, rrms = rmsnorm(h, self.norm_eps)
-      saves.extend([x_normed, rrms])
-      inp = x_normed * kwargs["ffn_norm"]
-      x_w1, *s = matmul(inp, kwargs["w1"], amax_x=kwargs["amax_x1"], w_inv_scale=kwargs["s_1"],
-                                  grad_amax_state=kwargs["grad_amax_xw1"], next_grad_amax_state=kwargs["next_grad_amax_xw1"],
-                                  next_amax_x=kwargs["next_amax_x1"])
-      saves.extend([*s, x_w1])
-      x_w3, *s = matmul(inp, kwargs["w3"], amax_x=kwargs["amax_x3"], w_inv_scale=kwargs["s_3"],
-                                  grad_amax_state=kwargs["grad_amax_xw3"], next_grad_amax_state=kwargs["next_grad_amax_xw3"],
-                                  next_amax_x=kwargs["next_amax_x3"])
-      saves.extend([*s, x_w3])
-      if FUSED_SILU_W13 and MXFP8:
-        from extra.llama_kernels.fused_silu_mul_quantize_mxfp8 import fused_silu_mul_quantize_mxfp8
-        aq, ae8, asi = fused_silu_mul_quantize_mxfp8(x_w1.reshape(-1, x_w1.shape[-1]), x_w3.reshape(-1, x_w3.shape[-1]))
-        out, *s = matmul(None, kwargs["w2"], x_prequant_mx=(aq, ae8, asi), amax_x=kwargs["amax_x2"],
-                         w_inv_scale=kwargs["s_2"], grad_amax_state=kwargs["grad_amax_xout"],
-                         next_grad_amax_state=kwargs["next_grad_amax_xout"], next_amax_x=kwargs["next_amax_x2"])
-        out = out.reshape(*x_w1.shape[:-1], kwargs["w2"].shape[0])
-      else:
-        out, *s = matmul(x_w1.silu() * x_w3, kwargs["w2"], amax_x=kwargs["amax_x2"], w_inv_scale=kwargs["s_2"],
-                         grad_amax_state=kwargs["grad_amax_xout"], next_grad_amax_state=kwargs["next_grad_amax_xout"],
-                         next_amax_x=kwargs["next_amax_x2"])
-      saves.extend([*s, out])
-    else:
-      x_w13, h, x_normed, rrms, s = add_norm_quantize_matmul(x, residual, kwargs["ffn_norm"], kwargs["w13"], kwargs["s_13"],
-                                                                          self.norm_eps, amax_x=kwargs["amax_x13"],
-                                                                          next_amax_x=kwargs["next_amax_x13"],
-                                                                          grad_amax_state=kwargs["grad_amax_xw13"],
-                                                                          next_grad_amax_state=kwargs["next_grad_amax_xw13"])
-      saves.extend([x_normed, rrms, *s, x_w13])
-      out, s = silu_w13_quantize_matmul(x_w13, kwargs["w2"], kwargs["s_2"], amax_x2=kwargs["amax_x2"],
-                                                     next_amax_x2=kwargs["next_amax_x2"],
-                                                     grad_amax_xw13=kwargs["grad_amax_xw13"],
-                                                     next_grad_amax_xw13=kwargs["next_grad_amax_xw13"],
-                                                     grad_amax_xout=kwargs["grad_amax_xout"],
-                                                     next_grad_amax_xout=kwargs["next_grad_amax_xout"])
-      saves.extend([*s, out])
-    return out, h, saves
+    x_w13, h, x_normed, rrms, ret = add_norm_quantize_matmul(x, residual, ffn_norm, w13, s_13, self.norm_eps,
+                                                             amax_x=amax_x13)
+    saves.extend([x_normed, rrms])
+    new_amaxs.extend(ret[:1])
+    saves.extend(ret[1:] + [x_w13])
+
+    out, ret = silu_w13_quantize_matmul(x_w13, w2, s_2, amax_x2=amax_x2, grad_amax_xw13=grad_amax_xw13, grad_amax_xout=grad_amax_xout)
+    new_amaxs.extend(ret[:1])
+    saves.extend(ret[1:] + [out])
+    return (out, h, *new_amaxs, *saves)
 
   @function(precompile=True, precompile_backward=True)
-  def run_layer(self, x:Tensor, freqs_cis:Tensor, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
-    attn, attn_saves = self.attention(x, freqs_cis, **attn_kwargs)
-    ffn, h, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs)
+  def run_layer(self, x:Tensor, freqs_cis:Tensor,
+                attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
+                ffn_norm:Tensor, w13:Tensor, w2:Tensor,
+                amax_xqkv:Tensor, amax_xo:Tensor,
+                amax_x13:Tensor, amax_x2:Tensor,
+                s_qkv:Tensor, s_o:Tensor, s_13:Tensor, s_2:Tensor,
+                grad_amax_xqkv:Tensor, grad_amax_xo:Tensor,
+                grad_amax_xw13:Tensor, grad_amax_xout:Tensor):
+    attn, *attn_ret = self.attention(x, freqs_cis, attention_norm, wqkv, wo,
+                                     amax_xqkv=amax_xqkv, amax_xo=amax_xo, s_qkv=s_qkv, s_o=s_o,
+                                     grad_amax_xqkv=grad_amax_xqkv, grad_amax_xo=grad_amax_xo)
+    attn_amaxs, attn_saves = attn_ret[:2], attn_ret[2:]
+    ffn, h, *ffn_ret = self.feed_forward(x, attn, ffn_norm, w13, w2,
+                                                   amax_x13=amax_x13, amax_x2=amax_x2, s_13=s_13, s_2=s_2,
+                                                   grad_amax_xw13=grad_amax_xw13, grad_amax_xout=grad_amax_xout)
+    ffn_amaxs, ffn_saves = ffn_ret[:2], ffn_ret[2:]
     h = h + ffn
-    if save: return (h, *attn_saves, *ffn_saves)
-    else: return (h,)
+    return (h, *attn_amaxs, *ffn_amaxs, *attn_saves, *ffn_saves)
 
   def shard(self, device:tuple[str, ...], mp:bool=False):
     from tinygrad.nn.state import get_parameters
@@ -291,69 +214,39 @@ class FlatTransformer:
       for v in get_parameters(self): v.shard_(device, axis=None)
     else:
       # flat per-layer weights: axis 0 is n_layers, so shard axes are +1 vs per-layer Transformer
-      def _shard_fp8(name:str, axis:int, std:float=0.02):
-        w = getattr(self, name)
-        if MXFP8:
-          from extra.gemm.cdna_asm_gemm import quantize_mxfp8
-          w_bf16 = Tensor.empty(self.n_layers, w.shape[1], w.shape[2], dtype=dtypes.bfloat16).shard(device, axis=axis).randn_like() * std
-          w_q, w_e8, _ = quantize_mxfp8(w_bf16)
-          w.replace(w_q)
-          self._fp8_inv_scale[name].replace(w_e8.contiguous()).is_param_(False)
-          self._fp8_next_inv_scale[name].replace(w_e8.contiguous()).is_param_(False)
-        else:
-          w.shard_(device, axis=axis)
-          scale_axis = (1 if axis == 1 else None) if COLUMNWISE_WEIGHT_SCALE else None
-          self._fp8_inv_scale[name] = self._fp8_inv_scale[name].shard(device, axis=scale_axis).contiguous().is_param_(False)
-          self._fp8_next_inv_scale[name] = self._fp8_next_inv_scale[name].shard(device, axis=scale_axis).contiguous().is_param_(False)
-        Tensor.realize(w, self._fp8_inv_scale[name], self._fp8_next_inv_scale[name])
-      sstd = 0.02 / math.sqrt(2 * self.n_layers)
-      _shard_fp8("wqkv", 1)          # (n_layers, out, dim) shard out
-      _shard_fp8("wo", 2, sstd)      # (n_layers, dim, in) shard in
-      if SPLIT_W13:
-        _shard_fp8("w1", 1)
-        _shard_fp8("w3", 1)
-      else:
-        _shard_fp8("w13", 1)         # (n_layers, hidden*2, dim) shard out
-      _shard_fp8("w2", 2, sstd)      # (n_layers, dim, hidden) shard in
+      self.wqkv.shard_(device, axis=1).realize()          # (n_layers, out, dim) shard out
+      self.wo.shard_(device, axis=2).realize()            # (n_layers, dim, in) shard in
+      self.w13.shard_(device, axis=1).realize()           # (n_layers, hidden*2, dim) shard out
+      self.w2.shard_(device, axis=2).realize()            # (n_layers, dim, hidden) shard in
       self.attention_norm.shard_(device, axis=None).realize()
       self.ffn_norm.shard_(device, axis=None).realize()
       self.norm.weight.shard_(device, axis=None).realize()
       self.tok_embeddings.weight.shard_(device, axis=0).realize()
       self.output.shard_(device, axis=1).realize()
       self.freqs_cis.shard_(device, axis=None).realize()
-      for amax_dict in (self._fp8_amax, self._fp8_next_amax, self._fp8_grad_amax, self._fp8_next_grad_amax):
+      for amax_dict in (self._fp8_amax, self._fp8_grad_amax):
         for name in amax_dict:
           for i in range(len(amax_dict[name])):
-            amax_dict[name][i] = amax_dict[name][i].to(device).contiguous().is_param_(False)
+            amax_dict[name][i] = amax_dict[name][i].to(device).contiguous().requires_grad_(False)
+      for name in self._fp8_inv_scale:
+        self._fp8_inv_scale[name] = self._fp8_inv_scale[name].to(device).contiguous().requires_grad_(False)
 
-  def reset_amax(self):
-    for st in (self._fp8_next_amax, self._fp8_next_grad_amax):
-      for ts in st.values():
-        for t in ts: t.assign(0)
-
-  def update_amax(self):
-    for cur, nxt in ((self._fp8_amax, self._fp8_next_amax), (self._fp8_grad_amax, self._fp8_next_grad_amax)):
-      for name in cur:
-        for c, n in zip(cur[name], nxt[name]): c.assign(n)
-
-  def __call__(self, tokens:Tensor, save:bool=True):
+  def __call__(self, tokens:Tensor):
     h = self.tok_embeddings(tokens)
-    freqs_cis = self.freqs_cis.cast(h.dtype)
-    if not getenv("HK_FLASH_ATTENTION"): freqs_cis = freqs_cis[:, :tokens.shape[1], :, :, :]
-    a, na, ga, nga, s = self._fp8_amax, self._fp8_next_amax, self._fp8_grad_amax, self._fp8_next_grad_amax, self._fp8_inv_scale
-    def amax_kwargs(i:int, act_names:tuple[str, ...], grad_names:tuple[str, ...]) -> dict[str, Tensor|None]:
-      specs = (("amax_", a, act_names), ("next_amax_", na, act_names), ("grad_amax_", ga, grad_names), ("next_grad_amax_", nga, grad_names))
-      if MXFP4: return dict.fromkeys(f"{prefix}{name}" for prefix, _, names in specs for name in names)
-      return {f"{prefix}{name}":val[name][i] for prefix, val, names in specs for name in names}
+    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
+    a, ga, s = self._fp8_amax, self._fp8_grad_amax, self._fp8_inv_scale
     for i in range(self.n_layers):
-      attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i], s_qkv=s["wqkv"][i], s_o=s["wo"][i],
-                         **amax_kwargs(i, ("xqkv", "xo"), ("xqkv", "xo")))
-      ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], w2=self.w2[i], s_2=s["w2"][i], **amax_kwargs(i, ("x2",), ("xout",)))
-      if SPLIT_W13:
-        ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i], s_1=s["w1"][i], s_3=s["w3"][i], **amax_kwargs(i, ("x1", "x3"), ("xw1", "xw3")))
-      else:
-        ffn_kwargs.update(w13=self.w13[i], s_13=s["w13"][i], **amax_kwargs(i, ("x13",), ("xw13",)))
-      h, *_ = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)
+      h, *ret = self.run_layer(h, freqs_cis,
+                               self.attention_norm[i], self.wqkv[i], self.wo[i],
+                               self.ffn_norm[i], self.w13[i], self.w2[i],
+                               amax_xqkv=a["xqkv"][i], amax_xo=a["xo"][i],
+                               amax_x13=a["x13"][i], amax_x2=a["x2"][i],
+                               s_qkv=s["wqkv"][i], s_o=s["wo"][i],
+                               s_13=s["w13"][i], s_2=s["w2"][i],
+                               grad_amax_xqkv=ga["xqkv"][i], grad_amax_xo=ga["xo"][i],
+                               grad_amax_xw13=ga["xw13"][i], grad_amax_xout=ga["xout"][i])
+      for name, new_val in zip(["xqkv", "xo", "x13", "x2"], ret[:5]):
+        a[name][i].assign(new_val)
 
     logits = matmul(self.norm(h), self.output[0], fp8=False)[0]
     return logits
@@ -364,61 +257,42 @@ def _get_pads(uop:UOp) -> list[UOp]:
 
 def apply_grad(grad_buf:Tensor, new_grad:UOp):
   pads = _get_pads(new_grad)
+  new_grad = new_grad.cast(grad_buf.dtype)
   if len(pads) <= 1:
-    new_grad = new_grad.cast(grad_buf.dtype)
-    grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(grad_buf.uop + new_grad))
+    store = grad_buf.uop.store(grad_buf.uop + new_grad)
+    grad_buf.uop = grad_buf.uop.after(store)
     return
-  cur = grad_buf.uop
-  for pad in sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0, reverse=True):
-    if pad.op == Ops.PAD:
-      grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg)])
-      buf_slice = cur.shrink(grad_shrink)
-      cur = cur.after(buf_slice.store(buf_slice + pad.src[0].cast(cur.dtype)))
-    else:
-      cur = cur.after(cur.store(cur + pad.cast(cur.dtype)))
-  grad_buf.uop = cur
+  sorted_pads = sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0)
+  inners = [Tensor(p.src[0] if p.op == Ops.PAD else p, device=grad_buf.device).cast(grad_buf.dtype) for p in sorted_pads]
+  if getenv("FUSED_PAD_GRAD_ACCUM", 0):
+    from extra.llama_kernels.fused_pad_grad_accum import fused_pad_grad_accum, can_fused_pad_grad_accum
+    if can_fused_pad_grad_accum(grad_buf, inners):
+      grad_buf.uop = fused_pad_grad_accum(grad_buf, inners).uop
+      return
+  grad_buf.assign(grad_buf + inners[0].cat(*inners[1:], dim=0))
 
 if __name__ == "__main__":
   config = {}
   BS                 = config["BS"]                     = getenv("BS", 16)
   SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
-  SMALL              = config["SMALL"]                  = getenv("SMALL", 0)
 
   from examples.llama3 import MODEL_PARAMS
-  model_params = MODEL_PARAMS[llama_size:=getenv("LLAMA3_SIZE", "8B")]["args"]
-  # vocab_size from mixtral tokenizer
-  if not SMALL: model_params |= {"vocab_size": 32000}
-  real_vocab_size = model_params['vocab_size']
-  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params["n_layers"] = llama_layers
-
-  # pad vocab
-  if (MP := getenv("MP", 1)) > 1: model_params["vocab_size"] = round_up(model_params["vocab_size"], 256 * MP)
-  vocab_mask:Tensor = Tensor.arange(model_params["vocab_size"]).reshape(1, 1, -1) >= real_vocab_size
-
+  model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
+  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
   model = FlatTransformer(**model_params, max_context=SEQLEN)
-
   state = nn.state.get_state_dict(model)
   print("tensor count:", len(state))
 
   # shard the model
   from tinygrad import Device
-  is_dp = (DP := getenv("DP", 1)) > 1
-  is_mp = (MP := getenv("MP", 1)) > 1
-  is_sharding = is_dp or is_mp
-  device_count = max(DP, MP)
-  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))
-
-  model.shard(device, is_mp)
-
-  if is_dp: vocab_mask.shard_(device, axis=None).realize()
-  if is_mp: vocab_mask.shard_(device, axis=2).realize()
+  if (DP := getenv("DP", 1)) > 1:
+    model.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)))
+  if (MP := getenv("MP", 1)) > 1:
+    model.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)), mp=True)
 
   # preallocate all the grad buffers and zero them out
-  grad_dtype = lambda x: dtypes.bfloat16 if x.dtype in dtypes.fp8s else x.dtype
-  grads = {x:x.zeros_like(dtype=grad_dtype(x)).contiguous() for x in state.values() if x.is_param}
-
-  fp8_amax = [t for ts in model._fp8_amax.values() for t in ts]
-  fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts]
+  grads = {x:Tensor.zeros(x.shape, dtype=x.dtype, device=x.device).contiguous()
+           for x in state.values() if x.requires_grad is None}
 
   # print model size
   sz = 0
@@ -427,32 +301,23 @@ if __name__ == "__main__":
     sz += v.nbytes()
   print(f"total sz: {sz/1e9:.2f} GB")
 
-  with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=real_vocab_size, dtype=dtypes.int)
+  with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=model.vocab_size, dtype=dtypes.int)
   with Timing("realize weights/grads/data: "): Tensor.realize(*state.values(), *grads.values(), tokens)
   print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
   if DP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)), axis=0)
   if MP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)))
 
   @TinyJit
-  def fwd_bwd(tokens:Tensor):
-    with Timing("python forward: "):
-      model.reset_amax()
-      logits = model(tokens[:, :-1], save=llama_size=="8B")
-      loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
+  def jit_step(tokens:Tensor):
+    with Timing("python forward: "): loss = model(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
     with Timing("python backward: "):
       for t,g in zip(grads, loss.gradient(*grads)):
         apply_grad(grads[t], g.uop)
-    with Timing("run fwd_bwd: "): loss.realize(*grads.values(), *fp8_amax, *fp8_grad_amax)
-
-  @TinyJit
-  def optim_step():
-    for g in grads.values(): g.assign(g.zeros_like())
-    Tensor.realize(*grads.values())
+    with Timing("run step: "): loss.realize(*grads.values())
 
   for i in range(6):
     GlobalCounters.reset()
     profile_marker(f"step {i}")
     with Timing(colored(f"*** step {i}: ", "red")):
-      fwd_bwd(tokens)
-      optim_step()
+      jit_step(tokens)
   print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))

examples/mlperf/models/gpt_oss.py

@@ -1,351 +0,0 @@
-import math, os, functools
-if __name__ == "__main__":
-  os.environ["DEFAULT_FLOAT"] = "bfloat16"
-  os.environ["OPTIM_DTYPE"] = "bfloat16"
-  if "DEV" not in os.environ: os.environ["DEV"] = "NULL::gfx950"
-  # CDNA
-  os.environ["DEVICE_IN_FUNCTION_BUG"] = "1"
-  os.environ["ALL2ALL"] = "1"
-  os.environ["USE_ATOMICS"] = "1"
-from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
-from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
-from tinygrad.uop.ops import Ops, UOp
-from extra.models.llama import apply_rotary_emb
-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.moe_gemm import grouped_mx_gemm
-from extra.gemm.moe_routing import route, dispatch, combine
-
-FP8_DTYPE = dtypes.fp8e4m3
-FP8_MAX = 448.0
-INIT_STD = 0.02
-ASM_GEMM = getenv("ASM_GEMM", 0)
-
-
-def _quant_dequant_fwd(x:Tensor) -> Tensor:
-  # x (2d bf16) -> bf16 value after an mxfp8 round-trip (1x32 block scaling on the last axis)
-  M, K = x.shape
-  scale_K = K // 32
-  amax = x.float().reshape(M, scale_K, 32).abs().max(axis=-1)
-  e8 = (amax.maximum(1e-38).log2().floor() + 127).clamp(0, 254).cast(dtypes.uint8)
-  qscale = (127.0 - e8.cast(dtypes.float32)).exp2().reshape(M, scale_K, 1).expand(M, scale_K, 32).reshape(M, K)
-  x_fp8 = (x.float() * qscale).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE).cast(dtypes.float32)
-  return (x_fp8 * _mx_block_scale(e8)).cast(dtypes.bfloat16)
-
-@functools.cache
-def _quant_dequant_fwd_fxn(x_p, device):
-  return _quant_dequant_fwd(Tensor(x_p, device=device))
-
-def _quant_dequant_bwd(grad:UOp, call:UOp) -> tuple:
-  return (Tensor(grad).cast(dtypes.bfloat16).uop,)
-
-def quant_dequant_mx(x:Tensor) -> Tensor:
-  fxn = _quant_dequant_fwd_fxn(x.as_param(0).uop, x.device)
-  return Tensor(UOp.maketuple(fxn.uop).call(x.uop, grad_fxn=_quant_dequant_bwd).gettuple(0))
-
-def _mx_scale(e8:Tensor) -> Tensor:
-  return _mx_block_scale(e8) if e8.ndim == 2 else _mx_block_scale_3d(e8)
-
-def _dequant_fwd(w_q:Tensor, w_scale:Tensor) -> Tensor:
-  return w_q.cast(dtypes.bfloat16) * _mx_scale(w_scale)
-
-@functools.cache
-def _dequant_fwd_fxn(wq_p, ws_p, device):
-  return _dequant_fwd(Tensor(wq_p, device=device), Tensor(ws_p, device=device))
-
-def _dequant_bwd(grad:UOp, call:UOp) -> tuple:
-  return (Tensor(grad).cast(dtypes.bfloat16).uop, None)
-
-def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:
-  fxn = _dequant_fwd_fxn(w_q.as_param(0).uop, w_scale.as_param(1).uop, w_q.device)
-  call = UOp.maketuple(fxn.uop).call(w_q.uop, w_scale.uop, grad_fxn=_dequant_bwd)
-  return Tensor(call.gettuple(0))
-
-def matmul_mx(x:Tensor, w_q:Tensor, w_scale:Tensor) -> Tensor:
-  l_shape = x.shape[:-1]
-  if ASM_GEMM:
-    from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm, mx_pack
-    x2, K, N = x.reshape(-1, x.shape[-1]), x.shape[-1], w_q.shape[0]
-    wq, ws = w_q, w_scale
-    if (pad := (-K) % 256):
-      x2 = x2.pad(((0, 0), (0, pad)))
-      wq = wq.pad(((0, 0), (0, pad)))
-      ws = ws.pad(((0, 0), (0, pad // 32)), value=127).cast(dtypes.uint8)
-    if (npad := (-N) % 256):
-      wq = wq.pad(((0, npad), (0, 0)))
-      ws = ws.pad(((0, npad), (0, 0)), value=127).cast(dtypes.uint8)
-    x_q, x_e8, x_si = quantize_mxfp8(x2)
-    if x_si is not None and can_use_asm_gemm(x_q, wq.T):
-      out = asm_gemm(x_q, wq.T, mx=True, mx_scales=(x_si, x_e8, mx_pack(ws), ws), mx_w_stored=True)
-      return (out[:, :N] if npad else out).reshape(*l_shape, N).cast(dtypes.bfloat16)
-  x_phys = quant_dequant_mx(x.reshape(-1, x.shape[-1])).reshape(*l_shape, x.shape[-1])
-  w_phys = dequant_weight(w_q, w_scale)
-  return (x_phys @ w_phys.T).cast(dtypes.bfloat16)
-
-def _pad_to_mult(t:Tensor, axis:int, mult:int=256) -> Tensor:
-  if (r := (-t.shape[axis]) % mult) == 0: return t
-  pads = [(0, 0)] * t.ndim
-  pads[axis] = (0, r)
-  return t.pad(tuple(pads))
-
-def _pad_cols(t:Tensor) -> Tensor: return _pad_to_mult(t, -1)
-def _pad_rows(t:Tensor) -> Tensor: return _pad_to_mult(t, -2)
-
-def swiglu(x:Tensor, limit:float=7.0, alpha:float=1.702) -> Tensor:
-  x_glu, x_linear = x[..., ::2], x[..., 1::2]
-  x_glu = x_glu.clamp(max_=limit)
-  x_linear = x_linear.clamp(-limit, limit)
-  return (x_glu * (alpha * x_glu).sigmoid()) * (x_linear + 1)
-
-def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
-  freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2, dtype=dtypes.float32)[:(dim // 2)] / dim))
-  freqs = Tensor.arange(end, dtype=dtypes.float32).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
-  return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).cast(dtypes.default_float).reshape(1, end, 1, dim//2, 2)
-
-class GPTOSS:
-  def __init__(self, dim:int, n_layers:int, n_heads:int, n_kv_heads:int, head_dim:int, n_experts:int, experts_per_tok:int,
-               intermediate_size:int, vocab_size:int, norm_eps:float=1e-5, rope_theta:int=150000, sliding_window:int=128,
-               swiglu_limit:float=7.0, max_context:int=8192):
-    self.dim, self.n_layers, self.n_heads, self.n_kv_heads, self.head_dim = dim, n_layers, n_heads, n_kv_heads, head_dim
-    self.n_rep = n_heads // n_kv_heads
-    self.n_experts, self.experts_per_tok, self.intermediate_size = n_experts, experts_per_tok, intermediate_size
-    self.vocab_size, self.norm_eps, self.sliding_window, self.swiglu_limit = vocab_size, norm_eps, sliding_window, swiglu_limit
-    self.sm_scale = 1.0 / math.sqrt(head_dim)
-
-    scaled_std = INIT_STD / math.sqrt(2 * n_layers)
-    q_dim, qkv_dim = n_heads * head_dim, head_dim * (n_heads + 2 * n_kv_heads)
-
-    # attn
-    self.wqkv, self.wqkv_scale = self._quant_weight(n_layers, qkv_dim, dim)
-    self.wqkv_bias = Tensor.zeros(n_layers, qkv_dim, dtype=dtypes.bfloat16).contiguous()
-    self.wo, self.wo_scale = self._quant_weight(n_layers, dim, q_dim, std=scaled_std)
-    self.wo_bias = Tensor.zeros(n_layers, dim, dtype=dtypes.bfloat16).contiguous()
-    self.sinks = Tensor.zeros(n_layers, n_heads, dtype=dtypes.bfloat16).contiguous()
-    self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
-
-    # moe ffn
-    self.ffn_norm = Tensor.ones(n_layers, dim).contiguous()
-    self.gate = Tensor.normal(n_layers, n_experts, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
-    self.gate_bias = Tensor.zeros(n_layers, n_experts, dtype=dtypes.bfloat16).contiguous()
-    self.w_gate_up, self.w_gate_up_scale = self._quant_weight(n_layers, n_experts, intermediate_size * 2, dim, moe=True)
-    self.w_gate_up_bias = Tensor.zeros(n_layers, n_experts, intermediate_size * 2, dtype=dtypes.bfloat16).contiguous()
-    self.w_down, self.w_down_scale = self._quant_weight(n_layers, n_experts, dim, intermediate_size, std=scaled_std, moe=True)
-    self.w_down_bias = Tensor.zeros(n_layers, n_experts, dim, dtype=dtypes.bfloat16).contiguous()
-
-    # output
-    self.norm = nn.RMSNorm(dim, norm_eps)
-    self.tok_embeddings = nn.Embedding(vocab_size, dim)
-    self.tok_embeddings.weight = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
-    self.output = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
-    self.freqs_cis = precompute_freqs_cis(head_dim, max_context * 2, rope_theta).contiguous().is_param_(False)
-
-  def _quant_weight(self, *shape:int, std:float=INIT_STD, moe:bool=False):
-    def _one(*s:int):
-      w = Tensor.zeros(*s) if getenv("ZEROS") else Tensor.normal(*s, mean=0.0, std=std)
-      w_q, w_e8, _ = quantize_mxfp8(_pad_cols(_pad_rows(w)) if moe else w)
-      return w_q, w_e8.is_param_(False)
-    if moe:
-      qs = [_one(*shape[1:]) for _ in range(shape[0])]
-      return [q[0] for q in qs], [q[1] for q in qs]
-    return _one(*shape)
-
-  def _attn_mask(self, seqlen:int, dtype) -> Tensor:
-    i, j = Tensor.arange(seqlen).reshape(seqlen, 1), Tensor.arange(seqlen).reshape(1, seqlen)
-    return (j <= i).where(0.0, -1e30).cast(dtype).contiguous()
-
-  def _sliding_attention(self, xq:Tensor, xk:Tensor, xv:Tensor, sinks:Tensor) -> Tensor:
-    bsz, seqlen, H, hd = xq.shape
-    KV, R, W = self.n_kv_heads, self.n_rep, self.sliding_window
-    assert seqlen % W == 0, f"seqlen {seqlen} must be a multiple of sliding_window {W} for banded attention"
-    nb = seqlen // W
-    q = xq.reshape(bsz, seqlen, KV, R, hd).permute(0, 2, 3, 1, 4).reshape(bsz, KV, R, nb, W, hd).float()
-    k, v = (x.permute(0, 2, 1, 3).reshape(bsz, KV, 1, nb, W, hd).float() for x in (xk, xv))
-    kk, vv = (x.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb].cat(x, dim=-2) for x in (k, v))
-    sc = (q @ kk.transpose(-1, -2)) * self.sm_scale   # (B,KV,R,nb,W,2W)
-    i, j, pv = Tensor.arange(W).reshape(W, 1), Tensor.arange(2 * W).reshape(1, 2 * W), Tensor.arange(nb).reshape(nb, 1, 1) >= 1
-    sc = ((j > i) & (j <= i + W) & (pv | (j >= W))).where(sc, -float("inf"))
-    sink = sinks.reshape(1, KV, R, 1, 1, 1).float()
-    m = sc.max(-1, keepdim=True).maximum(sink)
-    e = (sc - m).exp()
-    p = (e / (e.sum(-1, keepdim=True) + (sink - m).exp())).cast(dtypes.bfloat16)
-    attn = p @ vv.cast(dtypes.bfloat16)
-    return attn.reshape(bsz, KV, R, seqlen, hd).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, H * hd)
-
-  def attention(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, sliding:bool, *, attention_norm:Tensor, wqkv:Tensor,
-                wqkv_scale:Tensor, wqkv_bias:Tensor, wo:Tensor, wo_scale:Tensor, wo_bias:Tensor, sinks:Tensor):
-    bsz, seqlen, _ = x.shape
-    x_normed, rrms = rmsnorm(x, self.norm_eps)
-    qkv = matmul_mx(x_normed * attention_norm, wqkv, wqkv_scale) + wqkv_bias
-    qkv = qkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
-    xq = qkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
-    xk, xv = qkv[:, :, :, self.n_rep], qkv[:, :, :, self.n_rep + 1]
-    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
-    xq, 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"):
-      from extra.thunder.amd.fa import flash_attention
-      attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks)
-      attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
-    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)
-      scores = (xqm @ xkm.transpose(-2, -1)).float() * self.sm_scale + mask
-      sink = sinks.reshape(1, self.n_kv_heads, self.n_rep, 1, 1).float()
-      m = scores.max(-1, keepdim=True).maximum(sink)
-      e = (scores - m).exp()
-      w = (e / (e.sum(-1, keepdim=True) + (sink - m).exp())).cast(dtypes.bfloat16)
-      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]
-
-  def feed_forward(self, x:Tensor, *, ffn_norm:Tensor, gate:Tensor, gate_bias:Tensor,
-                   w_gate_up:Tensor, w_gate_up_scale:Tensor, w_gate_up_bias:Tensor,
-                   w_down:Tensor, w_down_scale:Tensor, w_down_bias:Tensor):
-    x_normed, rrms = rmsnorm(x, self.norm_eps)
-    inp = x_normed * ffn_norm
-    logits = inp.float() @ gate.float().T + gate_bias.float()
-    dim, inter = self.dim, self.intermediate_size
-
-    if getenv("GROUPED_MOE", 0):
-      bsz, seqlen = x.shape[:2]
-      inp, logits = inp.reshape(-1, dim), logits.reshape(-1, self.n_experts)
-      r = route(logits, self.experts_per_tok, self.n_experts)
-      onehot = r.rows_e.one_hot(self.n_experts).float()
-      xg = dispatch(_pad_cols(inp.cast(dtypes.bfloat16)), r)
-      h = grouped_mx_gemm(xg, (w_gate_up, w_gate_up_scale), r.off)[:, :2*inter] + (onehot @ w_gate_up_bias.float()).cast(dtypes.bfloat16)
-      y = swiglu(h, self.swiglu_limit)
-      z = grouped_mx_gemm(_pad_cols(y.cast(dtypes.bfloat16)), (w_down, w_down_scale), r.off)[:, :dim] \
-          + (onehot @ w_down_bias.float()).cast(dtypes.bfloat16)
-      out = combine(z, r, inp.shape[0], self.experts_per_tok).reshape(bsz, seqlen, dim)
-    else:
-      thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
-      weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)
-
-      out = None
-      for e in range(self.n_experts):
-        gu_q, gu_s = w_gate_up[e][:2*inter, :dim].contiguous(), w_gate_up_scale[e][:2*inter, :dim//32].contiguous()
-        dn_q, dn_s = w_down[e][:dim, :inter].contiguous(), w_down_scale[e][:dim, :inter//32].contiguous()
-        gate_up = matmul_mx(inp, gu_q, gu_s) + w_gate_up_bias[e]
-        y = (matmul_mx(swiglu(gate_up, self.swiglu_limit), dn_q, dn_s) + w_down_bias[e]).contiguous()
-        contrib = weights[..., e:e+1].cast(y.dtype) * y
-        out = contrib if out is None else out + contrib
-    return out, [x_normed, rrms]
-
-  @function(precompile=True, precompile_backward=True)
-  def run_layer(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, sliding:bool, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
-    attn, attn_saves = self.attention(x, freqs_cis, mask, sliding, **attn_kwargs)
-    h = x + attn
-    ffn, ffn_saves = self.feed_forward(h, **ffn_kwargs)
-    h = h + ffn
-    if save: return (h, *attn_saves, *ffn_saves)
-    return (h,)
-
-  def shard(self, device:tuple[str, ...], mp:bool=False):
-    assert not mp, "MP not supported"
-    from tinygrad.nn.state import get_parameters
-    for v in get_parameters(self): v.shard_(device, axis=None)
-    Tensor.realize(*get_parameters(self))
-
-  def __call__(self, tokens:Tensor, save:bool=True):
-    h = self.tok_embeddings(tokens)
-    bsz, seqlen = tokens.shape
-    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :seqlen, :, :, :]
-    mask_full = None if getenv("HK_FLASH_ATTENTION") else self._attn_mask(seqlen, dtypes.float32)
-    for i in range(self.n_layers):
-      attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wqkv_scale=self.wqkv_scale[i],
-                         wqkv_bias=self.wqkv_bias[i], wo=self.wo[i], wo_scale=self.wo_scale[i], wo_bias=self.wo_bias[i],
-                         sinks=self.sinks[i])
-      ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], gate=self.gate[i], gate_bias=self.gate_bias[i],
-                        w_gate_up=self.w_gate_up[i], w_gate_up_scale=self.w_gate_up_scale[i], w_gate_up_bias=self.w_gate_up_bias[i],
-                        w_down=self.w_down[i], w_down_scale=self.w_down_scale[i], w_down_bias=self.w_down_bias[i])
-      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
-    return logits
-
-def _get_pads(uop:UOp) -> list[UOp]:
-  if uop.op == Ops.ADD: return _get_pads(uop.src[0]) + _get_pads(uop.src[1])
-  return [uop]
-
-def apply_grad(grad_buf:Tensor, new_grad:UOp):
-  pads = _get_pads(new_grad)
-  if len(pads) <= 1:
-    new_grad = new_grad.cast(grad_buf.dtype)
-    grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(grad_buf.uop + new_grad))
-    return
-  cur = grad_buf.uop
-  for pad in sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0, reverse=True):
-    if pad.op == Ops.PAD:
-      grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg)])
-      buf_slice = cur.shrink(grad_shrink)
-      cur = cur.after(buf_slice.store(buf_slice + pad.src[0].cast(cur.dtype)))
-    else:
-      cur = cur.after(cur.store(cur + pad.cast(cur.dtype)))
-  grad_buf.uop = cur
-
-GPT_OSS_20B = dict(dim=2880, n_layers=24, n_heads=64, n_kv_heads=8, head_dim=64, n_experts=32, experts_per_tok=4,
-                   intermediate_size=2880, vocab_size=128256, norm_eps=1e-5, rope_theta=150000, sliding_window=128,
-                   swiglu_limit=7.0)
-
-if __name__ == "__main__":
-  config = {}
-  BS      = config["BS"]      = getenv("BS", 16)
-  SEQLEN  = config["SEQLEN"]  = getenv("SEQLEN", 8192)
-
-  model_params = GPT_OSS_20B
-  real_vocab_size = model_params["vocab_size"]
-  if (layers := getenv("LAYERS")) != 0: model_params["n_layers"] = layers
-
-  model = GPTOSS(**model_params, max_context=SEQLEN)
-
-  state = nn.state.get_state_dict(model)
-  print("tensor count:", len(state))
-
-  from tinygrad import Device
-  is_dp = (DP := getenv("DP", 1)) > 1
-  device_count = DP
-  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))
-
-  if is_dp: model.shard(device)
-
-  # preallocate all the grad buffers and zero them out
-  grad_dtype = lambda x: dtypes.bfloat16 if x.dtype in dtypes.fp8s else x.dtype
-  grads = {x:x.zeros_like(dtype=grad_dtype(x)).contiguous() for x in state.values() if x.is_param}
-
-  # print model size
-  sz = 0
-  for k,v in state.items():
-    print(f"{colored(k, 'green' if v in grads else 'white'):30s} {str(v.shape):30s} {str(v.dtype):20s} {v.device}  {v.nbytes()/1e9:.2f} GB")
-    sz += v.nbytes()
-  print(f"total sz: {sz/1e9:.2f} GB")
-
-  with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=real_vocab_size, dtype=dtypes.int)
-  with Timing("realize weights/grads/data: "): Tensor.realize(*state.values(), *grads.values(), tokens)
-  print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
-  if is_dp: tokens = tokens.shard(device, axis=0)
-
-  @TinyJit
-  def fwd_bwd(tokens:Tensor):
-    with Timing("python forward: "):
-      logits = model(tokens[:, :-1], save=True)
-      loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-    with Timing("python backward: "):
-      for t,g in zip(grads, loss.gradient(*grads)):
-        apply_grad(grads[t], g.uop)
-    with Timing("run fwd_bwd: "): loss.realize(*grads.values())
-
-  @TinyJit
-  def optim_step():
-    for g in grads.values(): g.assign(g.zeros_like())
-    Tensor.realize(*grads.values())
-
-  for i in range(6):
-    GlobalCounters.reset()
-    profile_marker(f"step {i}")
-    with Timing(colored(f"*** step {i}: ", "red")):
-      fwd_bwd(tokens)
-      optim_step()
-  print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))

examples/mlperf/models/test_apply_grad.py

@@ -1,68 +0,0 @@
-import unittest
-from tinygrad import Tensor, TinyJit
-from tinygrad.nn.state import get_parameters
-from examples.mlperf.models.flat_llama import apply_grad
-
-class FlatModel:
-  def __init__(self, n_layers:int, dim:int, hidden:int):
-    self.n_layers = n_layers
-    self.w1 = Tensor.uniform(n_layers, dim, hidden, low=-0.1, high=0.1)
-    self.w2 = Tensor.uniform(n_layers, hidden, dim, low=-0.1, high=0.1)
-    self.scale = Tensor.uniform(dim, low=0.9, high=1.1)
-    self.bias = Tensor.zeros(dim).contiguous()
-
-  def __call__(self, x:Tensor) -> Tensor:
-    h = x
-    for i in range(self.n_layers):
-      h = (h @ self.w1[i]).relu() @ self.w2[i] + h
-    return (h * self.scale + self.bias).sum()
-
-class TestApplyGradE2E(unittest.TestCase):
-  def _run_with_apply_grad(self, model, xs):
-    grads = {p: Tensor.zeros(p.shape, dtype=p.dtype).contiguous().realize() for p in get_parameters(model)}
-    for x in xs:
-      loss = model(x)
-      for p, g in zip(grads, loss.gradient(*grads)):
-        apply_grad(grads[p], g.uop)
-      Tensor.realize(loss, *grads.values())
-    return [grads[p] for p in get_parameters(model)]
-
-  def _run_reference(self, model, xs):
-    for x in xs: model(x).backward()
-    return [p.grad for p in get_parameters(model)]
-
-  def _assert_close(self, got, expected, atol, rtol):
-    for g, e in zip(got, expected):
-      self.assertTrue(g.allclose(e, atol=atol, rtol=rtol).item(), f"grad mismatch (max abs diff {(g - e).abs().max().item()})")
-
-  def _assert_match(self, model, xs, atol, rtol):
-    self._assert_close(self._run_with_apply_grad(model, xs), self._run_reference(model, xs), atol, rtol)
-
-  def test_e2e_single_step(self):
-    model = FlatModel(n_layers=3, dim=8, hidden=16)
-    Tensor.realize(*get_parameters(model))
-    self._assert_match(model, [Tensor.randn(2, 8).realize()], atol=1e-4, rtol=1e-4)
-
-  def test_e2e_multi_step_accumulation(self):
-    model = FlatModel(n_layers=4, dim=8, hidden=16)
-    Tensor.realize(*get_parameters(model))
-    self._assert_match(model, [Tensor.randn(2, 8).realize() for _ in range(3)], atol=1e-4, rtol=1e-4)
-
-  def test_e2e_jit(self):
-    model = FlatModel(n_layers=3, dim=8, hidden=16)
-    Tensor.realize(*get_parameters(model))
-    grads = {p: Tensor.zeros(p.shape, dtype=p.dtype).contiguous().realize() for p in get_parameters(model)}
-
-    @TinyJit
-    def fwd_bwd(x:Tensor):
-      loss = model(x)
-      for p, g in zip(grads, loss.gradient(*grads)): apply_grad(grads[p], g.uop)
-      Tensor.realize(loss, *grads.values())
-
-    xs = [Tensor.randn(2, 8).realize() for _ in range(3)]
-    for x in xs: fwd_bwd(x)
-    self._assert_close([grads[p] for p in get_parameters(model)], self._run_reference(model, xs), atol=1e-3, rtol=1e-3)
-
-
-if __name__ == "__main__":
-  unittest.main()

examples/mlperf/models/test_flat_llama.py

@@ -3,7 +3,8 @@ os.environ["WQKV"] = "1"
 import unittest
 import numpy as np
 from tinygrad import Tensor, nn, dtypes
-from tinygrad.device import Device
+from tinygrad.nn.state import get_parameters
+from tinygrad.device import is_dtype_supported, Device
 from examples.mlperf.models.llama import Transformer
 from examples.mlperf.models.flat_llama import FlatTransformer
 
@@ -44,6 +45,8 @@ class TestFlatLlama(unittest.TestCase):
     flat = FlatTransformer(**params)
     copy_weights(flat, ref)
 
+    for p in get_parameters(ref): p.requires_grad_(True)
+    for p in get_parameters(flat): p.requires_grad_(True)
     Tensor.realize(*nn.state.get_state_dict(flat).values())
 
     tokens = Tensor([[1, 50, 100, 999, 2, 10]])
@@ -111,7 +114,7 @@ class TestFlatLlama(unittest.TestCase):
     self.assertEqual(ref_logits.shape, flat_logits.shape)
     np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
 
-  @unittest.skipUnless(dtypes.fp8e4m3 in Device[Device.DEFAULT].renderer.supported_dtypes(), "fp8 not supported on this device")
+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), "fp8 not supported on this device")
   def test_forward_fp8(self):
     import examples.mlperf.models.flat_llama as flat_llama_mod
     old_fp8 = flat_llama_mod.FP8

examples/mlperf/optim.py

@@ -1,15 +1,11 @@
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
-from tinygrad.nn.optim import Optimizer, OptimizerGroup
+from tinygrad.nn.optim import Optimizer
 from tinygrad.helpers import FUSE_OPTIM, getenv
-from tinygrad.uop.ops import UOp, Ops, AxisType
+from tinygrad.uop.ops import UOp, Ops
 
 STOCHASTIC_ROUND = getenv("STOCHASTIC_ROUND", 0)
 MASTER_WEIGHTS = getenv("MASTER_WEIGHTS", 0)
-ZERO_OPTIM = getenv("ZERO_OPTIM", 0)
-FP8_AMAX_MARGIN = getenv("FP8_AMAX_MARGIN", 1.1)
-IMMEDIATE_SCALE = getenv("IMMEDIATE_SCALE", 0)
-MXFP8 = getenv("MXFP8", 0)
 
 def stochastic_round_bf16(x:Tensor) -> Tensor:
   bits = x.bitcast(dtypes.uint32)
@@ -21,50 +17,47 @@ def stochastic_round_bf16(x:Tensor) -> Tensor:
   noise = (noise * 0xFFFF).cast(dtypes.uint32)
   return ((bits + noise) & 0xFFFF0000).bitcast(dtypes.float32).cast(dtypes.bfloat16)
 
-def clip_grads(grads:list[Tensor], grad_acc, clip_norm) -> Tensor:
-  for g in grads: g.assign(g / grad_acc)
-  total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
-  for g in grads: g.assign((g * (clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(g.dtype))
-  return total_norm
-
 class GradAccClipAdamW(Optimizer):
   def __init__(self, params:list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, grad_acc=1, clip_norm=1.0, device=None, fused=FUSE_OPTIM):
     super().__init__(params, lr, device, fused)
     self.b1, self.b2, self.eps, self.wd = b1, b2, eps, weight_decay
-    self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device) for _ in [b1, b2])
-    self.zero = bool(ZERO_OPTIM) and isinstance(self.device, tuple) and not self.fused
-    self.m = [self._zero_shard(x) for x in self._new_optim_param()]
-    self.v = [self._zero_shard(x) for x in self._new_optim_param()]
+    self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False) for _ in [b1, b2])
+    self.m = self._new_optim_param()
+    self.v = self._new_optim_param()
     self.grad_acc, self.clip_norm = grad_acc, clip_norm
-    if MASTER_WEIGHTS and self.params[0].dtype != dtypes.float32:
-      self.master_params:list[Tensor]|None = [self._zero_shard(p.to(self.device).float().contiguous()) for p in self.params]
+    self.master_params:list[Tensor]|None = [p.float().contiguous() for p in self.params] if MASTER_WEIGHTS and self.params[0].dtype != dtypes.float32 else None
+
+  def fstep(self, grads:list[Tensor]):
+    if self.fused:
+      out, extra = self._step([], grads)
+      updates = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
     else:
-      self.master_params = None
-
-  def _zero_shard(self, t:Tensor) -> Tensor:
-    if not self.zero or t.ndim < 2 or (t.shape[0] % len(self.device)) != 0: return t
-    return Tensor(t.uop._shard(0, UOp.range(len(self.device), -1, AxisType.DEVICE)).unshard(0)).clone()
-
-  def _zero_gather(self, t:Tensor) -> Tensor:
-    if not isinstance(t.device, tuple) or t.uop.axis != 0: return t
-    n, sz = len(t.device), t.shape[0] // len(t.device)
-    return Tensor.cat(*[t[p*sz:(p+1)*sz] for p in range(n)], dim=0)
-
-  def fschedule_step(self, grads:list[Tensor]) -> list[Tensor]:
-    updates, extra = self._step([], grads)
+      updates, extra = self._step([], grads)
     for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i], self.master_params[i] if self.master_params else None))
+    # collect inv_scale tensors attached to fp8 params (set by _apply_update)
     fp8_inv_scales = [tt._inv_scale for tt in self.params if hasattr(tt, '_inv_scale')]
-    fp8_next_inv_scales = [tt._next_inv_scale for tt in self.params if hasattr(tt, '_next_inv_scale')]
-    return extra + self.params + self.buffers + (self.master_params or []) + fp8_inv_scales + fp8_next_inv_scales
+    to_realize = extra+self.params+self.buffers+(self.master_params or [])+fp8_inv_scales
 
-  def fstep(self, grads:list[Tensor], grad_norm:Tensor|None=None):
-    Tensor.realize(*([grad_norm] if grad_norm is not None else []), *self.fschedule_step(grads))
+    Tensor.realize(*to_realize)
+    return extra[-1]
 
   def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
     grads = list(grads)
 
     for i in range(len(grads)):
       if grads[i].device != self.m[i].device: grads[i] = grads[i].to(self.m[i].device)
+
+    if self.fused:
+      grads[0].assign(grads[0] / self.grad_acc)
+      total_norm = grads[0].float().square().sum().sqrt()
+      grads[0].assign((grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype))
+    else:
+      for i in range(len(grads)):
+        grads[i].assign(grads[i] / self.grad_acc)
+      total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
+      for i in range(len(grads)):
+        grads[i].assign((grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype))
+
     ret = []
     self.b1_t *= self.b1
     self.b2_t *= self.b2
@@ -77,7 +70,7 @@ class GradAccClipAdamW(Optimizer):
       v_hat = v_new / (1.0 - self.b2_t)
       up = m_hat / (v_hat.sqrt() + self.eps)
       ret.append(self.lr * up)
-    return ret, [self.b1_t, self.b2_t] + self.m + self.v
+    return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]
 
   def _apply_update(self, t:Tensor, up:Tensor, master:Tensor|None=None) -> Tensor:
     w = master if master is not None else t
@@ -85,60 +78,13 @@ class GradAccClipAdamW(Optimizer):
     up = up.float().shard_like(w) + self.lr.to(w.device) * wd * w.detach()
     new_w = w.detach() - up
     if master is not None: master.assign(new_w)
-    if self.zero and not (MXFP8 and t.dtype in dtypes.fp8s): new_w = self._zero_gather(new_w)
-    # when master is offloaded to a different device than the param, results are resharded back onto the param's (sharded) device
-    offloaded = master is not None and master.device != t.device
-    if STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16:
-      out = stochastic_round_bf16(new_w)
-      return out.shard_like(t) if offloaded else out
+    if STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16: return stochastic_round_bf16(new_w)
     if t.dtype in dtypes.fp8s:
-      if MXFP8:
-        from extra.gemm.cdna_asm_gemm import quantize_mxfp8
-        w_q, w_e8, _ = quantize_mxfp8(new_w.reshape(-1, new_w.shape[-1]))
-        if self.zero: w_q, w_e8 = self._zero_gather(w_q), self._zero_gather(w_e8)
-        new_e8 = w_e8.reshape(t._inv_scale.shape)
-        t._inv_scale.assign(new_e8.shard_like(t._inv_scale) if offloaded else new_e8)
-        ret = w_q.reshape(t.shape)
-        return ret.shard_like(t) if offloaded else ret
       from examples.mlperf.models.flat_llama import FP8_MAX
-      if IMMEDIATE_SCALE:
-        amax_axis = tuple(range(t._inv_scale.ndim, new_w.ndim))
-        new_inv = ((new_w.float().abs().max(axis=amax_axis).detach() + 1e-8) / FP8_MAX).cast(t._inv_scale.dtype)
-        t._inv_scale.assign(new_inv.shard_like(t._inv_scale) if offloaded else new_inv)
-        scale = new_inv.reciprocal().reshape(*new_inv.shape, *([1]*(new_w.ndim-new_inv.ndim)))
-        ret = (new_w * scale).clamp(-FP8_MAX, FP8_MAX).cast(t.dtype)
-        return ret.shard_like(t) if offloaded else ret
-      # delayed scaling: reuse previous step's inv_scale
-      t._inv_scale.assign(t._next_inv_scale)
-      inv_scale = t._inv_scale.to(new_w.device) if offloaded else t._inv_scale
-      scale = inv_scale.reciprocal().reshape(*inv_scale.shape, *([1]*(new_w.ndim-inv_scale.ndim)))
-      scaled = (new_w * scale).clamp(-FP8_MAX, FP8_MAX)
-      ret = scaled.cast(t.dtype)
-      # update inv_scale for next step from quantized result
-      new_amax = (ret.float().abs().max(axis=tuple(range(inv_scale.ndim, ret.ndim))) * inv_scale * FP8_AMAX_MARGIN).detach()
-      new_inv = ((new_amax + 1e-8) / FP8_MAX).cast(t._inv_scale.dtype)
-      t._next_inv_scale.assign(new_inv.shard_like(t._next_inv_scale) if offloaded else new_inv)
-      return ret.shard_like(t) if offloaded else ret
-    out = new_w.cast(t.dtype)
-    return out.shard_like(t) if offloaded else out
-
-class GradAccClipAdamWGroup(OptimizerGroup):
-  def __init__(self, *optimizers:GradAccClipAdamW):
-    super().__init__(*optimizers)
-    for o in self.optimizers[1:]: o.lr = self.optimizers[0].lr
-  def fstep(self, grads:list[Tensor], grad_norm:Tensor|None=None):
-    offset = 0
-    to_realize = []
-    for o in self.optimizers:
-      n = len(o.params)
-      to_realize += o.fschedule_step(grads[offset:offset+n])
-      offset += n
-    Tensor.realize(*to_realize, *([grad_norm] if grad_norm is not None else []))
-  @property
-  def lr(self): return self.optimizers[0].lr
-  @property
-  def device(self): return self.optimizers[0].device
-  @property
-  def master_params(self):
-    mp = [mp for o in self.optimizers for mp in (o.master_params or [])]
-    return mp if mp else None
+      amax = new_w.float().abs().flatten(1).max(1).detach()  # per-layer amax for (n_layers, out, in)
+      scale = FP8_MAX / (amax + 1e-8)
+      fp8_w = (new_w * scale.reshape(-1, *([1]*(new_w.ndim-1)))).clamp(-FP8_MAX, FP8_MAX).cast(t.dtype)
+      if hasattr(t, '_inv_scale'):
+        t._inv_scale.assign(((amax + 1e-8) / FP8_MAX).cast(t._inv_scale.dtype))
+      return fp8_w
+    return new_w.cast(t.dtype)

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

@@ -1,44 +0,0 @@
-#!/usr/bin/env bash
-
-export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
-export DEV=${DEV:-AMD}
-export CHECK_OOB=0
-export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
-export DEVICE_IN_FUNCTION_BUG=1
-
-export DEBUG=${DEBUG:-2}
-export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
-export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export USE_ATOMICS=${USE_ATOMICS:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export MXFP8=${MXFP8:-1}
-export ZERO_OPTIM=${ZERO_OPTIM:-1}
-export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
-
-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
-export GBS=$((BS * GRADIENT_ACC_STEPS))
-
-export MODEL="gptoss"
-export BASEDIR="/raid/datasets/c4-8b/"
-export EVAL_TARGET=3.34 EVAL_FREQ=12288
-export END_LR="4e-5" WARMUP_STEPS=128 MAX_STEPS=1200000
-export SAMPLES=$((MAX_STEPS * GBS))
-export SEQLEN=${SEQLEN:-8192}
-
-export SEED=${SEED:-5760}
-export DATA_SEED=${DATA_SEED:-5760}
-
-export JITBEAM=${JITBEAM:-3}
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
-
-export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
-if [ -z "$FULL_LAYERS" ]; then
-  export LAYERS=${LAYERS:-2}
-fi
-
-python3 examples/mlperf/model_train.py

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

@@ -1,39 +0,0 @@
-#!/usr/bin/env bash
-
-export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
-export DEV=${DEV:-AMD}
-export CHECK_OOB=0
-export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
-export DEVICE_IN_FUNCTION_BUG=1
-
-export DEBUG=${DEBUG:-0}
-export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
-export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export USE_ATOMICS=${USE_ATOMICS:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export MXFP8=${MXFP8:-1}
-export ZERO_OPTIM=${ZERO_OPTIM:-1}
-export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
-
-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
-export GBS=$((BS * GRADIENT_ACC_STEPS))
-
-export MODEL="gptoss"
-export BASEDIR="/raid/datasets/c4-8b/"
-export EVAL_TARGET=3.34 EVAL_FREQ=12288
-export END_LR="4e-5" WARMUP_STEPS=128 MAX_STEPS=1200000
-export SAMPLES=$((MAX_STEPS * GBS))
-export SEQLEN=${SEQLEN:-8192}
-
-export SEED=${SEED:-$RANDOM}
-export DATA_SEED=${DATA_SEED:-5760}
-
-export JITBEAM=${JITBEAM:-3}
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
-
-python3 examples/mlperf/model_train.py

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

@@ -1,54 +0,0 @@
-#!/usr/bin/env bash
-
-export PYTHONPATH="."
-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
-export DEVICE_IN_FUNCTION_BUG=1
-
-export DEBUG=${DEBUG:-2}
-export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
-export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
-export USE_ATOMICS=${USE_ATOMICS:-1}
-export ASM_GEMM=${ASM_GEMM:-1}
-export WQKV=${WQKV:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export MXFP4=${MXFP4:-1}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-1}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
-export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
-export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
-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="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))
-
-export MODEL="llama3"
-export BASEDIR="/raid/datasets/c4-8b/"
-export SMALL=1
-export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
-export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
-export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
-export SAMPLES=$((MAX_STEPS * GBS))
-export SEQLEN=${SEQLEN:-8192}
-
-export SEED=${SEED:-5760}
-export DATA_SEED=${DATA_SEED:-5760}
-
-export JITBEAM=${JITBEAM:-3}
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
-
-export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
-if [ -z "$FULL_LAYERS" ]; then
-  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
-fi
-
-python3 examples/mlperf/model_train.py

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

@@ -1,49 +0,0 @@
-#!/usr/bin/env bash
-
-export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
-export DEV=${DEV:-AMD}
-export CHECK_OOB=0
-export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
-export DEVICE_IN_FUNCTION_BUG=1
-
-export DEBUG=${DEBUG:-0}
-export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
-export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
-export USE_ATOMICS=${USE_ATOMICS:-1}
-export ASM_GEMM=${ASM_GEMM:-1}
-export WQKV=${WQKV:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export MXFP4=${MXFP4:-1}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-1}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
-export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
-export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
-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="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))
-
-export MODEL="llama3"
-export BASEDIR="/raid/datasets/c4-8b/"
-export SMALL=1
-export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
-export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
-export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
-export SAMPLES=$((MAX_STEPS * GBS))
-export SEQLEN=${SEQLEN:-8192}
-
-export SEED=${SEED:-$RANDOM}
-export DATA_SEED=${DATA_SEED:-5760}
-
-export JITBEAM=${JITBEAM:-3}
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
-
-python3 examples/mlperf/model_train.py

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

@@ -1,7 +0,0 @@
-#!/bin/bash
-set -e
-
-export BENCHMARK=${BENCHMARK:-5}
-export EVAL_BS=0
-VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=${DEBUG:--0} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
-[ "$BENCHMARK" -le 3 ] || [[ $DEV == NULL* ]] || python -m tinygrad.viz.cli -s AMD -t --interval "train @ 2" "train @ 3"

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

@@ -1,58 +0,0 @@
-#!/usr/bin/env bash
-set -e  # Exit on any error
-set -o pipefail  # Make pipeline fail if any command fails
-
-export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
-export DEV=AMD
-export CHECK_OOB=0
-export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
-export DEVICE_IN_FUNCTION_BUG=1
-
-export HK_FLASH_ATTENTION=1
-export ALL2ALL=1
-export LATE_ALLREDUCE=0
-export USE_ATOMICS=1
-export ASM_GEMM=1
-export WQKV=1
-export MASTER_WEIGHTS=1
-export FP8=1
-export ALLREDUCE_CAST=1
-export FAST_CE=1
-export FUSED_INPUT_QUANTIZE=1
-export FUSED_GRAD_QUANTIZE=1
-export FUSED_ADD_NORM_MUL_QUANTIZE=1
-export FUSED_SILU_W13=1
-export SPLIT_W13=0
-
-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=8 MP=1 BS=16 EVAL_BS=8 GRADIENT_ACC_STEPS=2
-export GBS=$((BS * GRADIENT_ACC_STEPS))
-
-export MODEL="llama3"
-export BASEDIR="/raid/datasets/c4-8b/"
-export SMALL=1
-export LLAMA3_SIZE=8B
-export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
-export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
-export SAMPLES=$((MAX_STEPS * GBS))
-export SEQLEN=8192
-
-export SEED=$RANDOM
-export DATA_SEED=$SEED
-
-export JITBEAM=3
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
-
-export LOGMLPERF=1
-
-DATETIME=$(date "+%m%d%H%M")
-LOGFILE="llama31_8b_8xMI350x_${DATETIME}_${SEED}.log"
-
-# beam
-FAKEDATA=1 BENCHMARK=10 INITMLPERF=1 LLAMA_LAYERS=2 python3 examples/mlperf/model_train.py | tee "$LOGFILE"
-
-# run
-RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a "$LOGFILE"

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

@@ -1,9 +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 DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
 export DEVICE_IN_FUNCTION_BUG=1
@@ -11,24 +10,14 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
-export USE_ATOMICS=${USE_ATOMICS:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export FP8=${FP8:-1}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-0}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
-export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
-export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
-export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
-export SPLIT_W13=${SPLIT_W13:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-1} MP=${MP:-8} BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
-export GBS=$((BS * GRADIENT_ACC_STEPS))
+export DP=${DP:-1} MP=${MP:-8}
+export BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
 
 export MODEL="llama3"
 export BASEDIR="/raid/datasets/c4/"
@@ -41,9 +30,9 @@ 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=1 BENCHMARK=10
 if [ -z "$FULL_LAYERS" ]; then
-  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
+  export LLAMA_LAYERS=2
 fi
 
 python3 examples/mlperf/model_train.py

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

@@ -1,34 +1,22 @@
 #!/usr/bin/env bash
 
 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
-export DEVICE_IN_FUNCTION_BUG=1
 
 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
-export USE_ATOMICS=${USE_ATOMICS:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export FP8=${FP8:-1}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-0}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
-export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
-export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
-export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
-export SPLIT_W13=${SPLIT_W13:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-1} MP=${MP:-8} BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1152}
-export GBS=$((BS * GRADIENT_ACC_STEPS))
+export DP=${DP:-1} MP=${MP:-8}
+export BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1152}
 
 export MODEL="llama3"
 export BASEDIR="/raid/datasets/c4/"

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

@@ -1,8 +1,6 @@
 #!/usr/bin/env bash
 
 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=${DEV:-AMD}
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -11,20 +9,17 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-1}
+export FAST_CE=${FASE_CE:-1}
 export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
-export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
 export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
 export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
-export SPLIT_W13=${SPLIT_W13:-0}
-export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
+export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
@@ -48,7 +43,7 @@ export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGR
 
 export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
-  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
+  export LLAMA_LAYERS=2
 fi
 
 python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam_mp.sh

@@ -1,9 +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 DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
 export DEVICE_IN_FUNCTION_BUG=1
@@ -11,19 +10,9 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
-export USE_ATOMICS=${USE_ATOMICS:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export FP8=${FP8:-1}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-0}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
-export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
-export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
-export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
-export SPLIT_W13=${SPLIT_W13:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
@@ -46,9 +35,9 @@ 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=1 BENCHMARK=10
 if [ -z "$FULL_LAYERS" ]; then
-  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
+  export LLAMA_LAYERS=2
 fi
 
 python3 examples/mlperf/model_train.py

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

@@ -1,8 +1,6 @@
 #!/usr/bin/env bash
 
 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=${DEV:-AMD}
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -11,20 +9,17 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-1}
+export FAST_CE=${FASE_CE:-1}
 export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
-export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
 export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
 export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
-export SPLIT_W13=${SPLIT_W13:-0}
-export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
+export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run_mp.sh

@@ -1,9 +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 DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
 export DEVICE_IN_FUNCTION_BUG=1
@@ -11,19 +10,9 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
-export USE_ATOMICS=${USE_ATOMICS:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export FP8=${FP8:-1}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-0}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
-export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
-export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
-export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
-export SPLIT_W13=${SPLIT_W13:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+export BENCHMARK=5
+export EVAL_BS=0
+VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh
+SRC="AMD"; [[ $DEV == NULL* ]] && SRC="NULL"
+python -m tinygrad.viz.cli -s "$SRC" --top 20

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/run_and_time.sh

@@ -3,8 +3,6 @@ set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails
 
 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=AMD
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -12,7 +10,6 @@ export DEVICE_IN_FUNCTION_BUG=1
 
 export HK_FLASH_ATTENTION=1
 export ALL2ALL=1
-export LATE_ALLREDUCE=0
 export USE_ATOMICS=1
 export ASM_GEMM=1
 export WQKV=1
@@ -21,10 +18,9 @@ export FP8=1
 export ALLREDUCE_CAST=1
 export FAST_CE=1
 export FUSED_INPUT_QUANTIZE=1
-export FUSED_GRAD_QUANTIZE=1
 export FUSED_ADD_NORM_MUL_QUANTIZE=1
 export FUSED_SILU_W13=1
-export SPLIT_W13=0
+export FUSED_PAD_GRAD_ACCUM=1
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=8 MP=1 BS=16 EVAL_BS=8 GRADIENT_ACC_STEPS=2

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/startup_walltime.sh

@@ -4,7 +4,7 @@ export EVAL_BS=0
 export FAKEDATA=1
 export NULL_ALLOW_COPYOUT=1
 export HIP_VISIBLE_DEVICES=""
-export DEV=NULL:HIP:gfx950
+export DEV=NULL
 export JITBEAM=0
 export LLAMA_LAYERS=${LLAMA_LAYERS:-"2"}
-time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh
+time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh

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

@@ -1,45 +0,0 @@
-#!/usr/bin/env bash
-
-export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
-export DEV=${DEV:-AMD}
-export CHECK_OOB=0
-export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
-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 ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export USE_ATOMICS=${USE_ATOMICS:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export MXFP8=${MXFP8:-1}
-export ZERO_OPTIM=${ZERO_OPTIM:-1}
-export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
-
-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
-export GBS=$((BS * GRADIENT_ACC_STEPS))
-
-export MODEL="gptoss"
-export BASEDIR="/raid/datasets/c4-8b/"
-export EVAL_TARGET=3.34 EVAL_FREQ=12288
-export END_LR="4e-5" WARMUP_STEPS=128 MAX_STEPS=1200000
-export SAMPLES=$((MAX_STEPS * GBS))
-export SEQLEN=${SEQLEN:-8192}
-
-export SEED=${SEED:-5760}
-export DATA_SEED=${DATA_SEED:-5760}
-
-export JITBEAM=${JITBEAM:-3}
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
-
-export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
-if [ -z "$FULL_LAYERS" ]; then
-  export LAYERS=${LAYERS:-2}
-fi
-
-python3 examples/mlperf/model_train.py

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

@@ -1,40 +0,0 @@
-#!/usr/bin/env bash
-
-export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
-export DEV=${DEV:-AMD}
-export CHECK_OOB=0
-export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
-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 ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export USE_ATOMICS=${USE_ATOMICS:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export MXFP8=${MXFP8:-1}
-export ZERO_OPTIM=${ZERO_OPTIM:-1}
-export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
-
-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-8} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1}
-export GBS=$((BS * GRADIENT_ACC_STEPS))
-
-export MODEL="gptoss"
-export BASEDIR="/raid/datasets/c4-8b/"
-export EVAL_TARGET=3.34 EVAL_FREQ=12288
-export END_LR="4e-5" WARMUP_STEPS=128 MAX_STEPS=1200000
-export SAMPLES=$((MAX_STEPS * GBS))
-export SEQLEN=${SEQLEN:-8192}
-
-export SEED=${SEED:-$RANDOM}
-export DATA_SEED=${DATA_SEED:-5760}
-
-export JITBEAM=${JITBEAM:-3}
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
-
-python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam_mp.sh

@@ -1,54 +0,0 @@
-#!/usr/bin/env bash
-
-export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
-export DEV=${DEV:-AMD}
-export CHECK_OOB=0
-export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
-export DEVICE_IN_FUNCTION_BUG=1
-
-export DEBUG=${DEBUG:-2}
-export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
-export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
-export USE_ATOMICS=${USE_ATOMICS:-1}
-export ASM_GEMM=${ASM_GEMM:-1}
-export WQKV=${WQKV:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export FP8=${FP8:-1}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-0}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
-export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
-export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
-export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
-export SPLIT_W13=${SPLIT_W13:-1}
-export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
-
-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-1} MP=${MP:-8} BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
-export GBS=$((BS * GRADIENT_ACC_STEPS))
-
-export MODEL="llama3"
-export BASEDIR="/raid/datasets/c4-8b/"
-export SMALL=1
-export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
-export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
-export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
-export SAMPLES=$((MAX_STEPS * GBS))
-export SEQLEN=${SEQLEN:-8192}
-
-export SEED=${SEED:-5760}
-export DATA_SEED=${DATA_SEED:-5760}
-
-export JITBEAM=${JITBEAM:-3}
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
-
-export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
-if [ -z "$FULL_LAYERS" ]; then
-  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
-fi
-
-python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run_mp.sh

@@ -1,49 +0,0 @@
-#!/usr/bin/env bash
-
-export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
-export DEV=${DEV:-AMD}
-export CHECK_OOB=0
-export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
-export DEVICE_IN_FUNCTION_BUG=1
-
-export DEBUG=${DEBUG:-0}
-export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
-export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-1}
-export USE_ATOMICS=${USE_ATOMICS:-1}
-export ASM_GEMM=${ASM_GEMM:-1}
-export WQKV=${WQKV:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export FP8=${FP8:-1}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-0}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
-export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
-export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
-export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
-export SPLIT_W13=${SPLIT_W13:-1}
-export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
-
-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-export DP=${DP:-1} MP=${MP:-8} BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-32}
-export GBS=$((BS * GRADIENT_ACC_STEPS))
-
-export MODEL="llama3"
-export BASEDIR="/raid/datasets/c4-8b/"
-export SMALL=1
-export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
-export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
-export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
-export SAMPLES=$((MAX_STEPS * GBS))
-export SEQLEN=${SEQLEN:-8192}
-
-export SEED=${SEED:-$RANDOM}
-export DATA_SEED=${DATA_SEED:-5760}
-
-export JITBEAM=${JITBEAM:-3}
-export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
-
-python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh

@@ -1,5 +0,0 @@
-#!/bin/bash
-export BENCHMARK=${BENCHMARK:-5}
-export EVAL_BS=0
-VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=${DEBUG:--0} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
-[ "$BENCHMARK" -le 3 ] || [[ $DEV == NULL* ]] || python -m tinygrad.viz.cli -s AMD -t --interval "train @ 2" "train @ 3"

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/startup_walltime.sh

@@ -1,10 +0,0 @@
-#!/bin/bash
-export BENCHMARK=5
-export EVAL_BS=0
-export FAKEDATA=1
-export NULL_ALLOW_COPYOUT=1
-export HIP_VISIBLE_DEVICES=""
-export DEV=NULL:HIP:gfx950
-export JITBEAM=0
-export LLAMA_LAYERS=${LLAMA_LAYERS:-"2"}
-time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh

examples/mlperf/training_submission_v6.1/tinycorp/systems/tinybox_8xMI300X.json

@@ -1,38 +0,0 @@
-{
-    "submitter": "tinycorp",
-    "division": "closed",
-    "status": "Available on-premise",
-    "system_name": "tinybox 8xMI300X",
-    "number_of_nodes": "1",
-    "host_processors_per_node": "2",
-    "host_processor_model_name": "AMD EPYC 9354",
-    "host_processor_core_count": "32",
-    "host_processor_vcpu_count": "64",
-    "host_processor_frequency": "",
-    "host_processor_caches": "",
-    "host_processor_interconnect": "",
-    "host_memory_capacity": "2304GB",
-    "host_storage_type": "NVMe SSD",
-    "host_storage_capacity": "3x 4TB raid array",
-    "host_networking": "",
-    "host_networking_topology": "",
-    "host_memory_configuration": "24x 96GB DDR5",
-    "accelerators_per_node": "8",
-    "accelerator_model_name": "AMD Instinct MI300X 192GB HBM3",
-    "accelerator_host_interconnect": "PCIe 5.0 x16",
-    "accelerator_frequency": "",
-    "accelerator_on-chip_memories": "",
-    "accelerator_memory_configuration": "HBM3",
-    "accelerator_memory_capacity": "192GB",
-    "accelerator_interconnect": "",
-    "accelerator_interconnect_topology": "",
-    "cooling": "air",
-    "hw_notes": "",
-    "framework": "tinygrad, branch mlperf_training_v5.0",
-    "other_software_stack": {
-        "python": "3.10.16",
-        "ROCm": "3.0.0+94441cb"
-    },
-    "operating_system": "Ubuntu 24.04.1 LTS",
-    "sw_notes": ""
-  }

examples/mlperf/training_submission_v6.1/tinycorp/systems/tinybox_8xMI350X.json

@@ -1,38 +0,0 @@
-{
-    "submitter": "tinycorp",
-    "division": "closed",
-    "status": "Available on-premise",
-    "system_name": "tinybox 8xMI350X",
-    "number_of_nodes": "1",
-    "host_processors_per_node": "2",
-    "host_processor_model_name": "AMD EPYC 9575F",
-    "host_processor_core_count": "32",
-    "host_processor_vcpu_count": "64",
-    "host_processor_frequency": "",
-    "host_processor_caches": "",
-    "host_processor_interconnect": "",
-    "host_memory_capacity": "3072 GiB",
-    "host_storage_type": "NVMe SSD",
-    "host_storage_capacity": "4TB",
-    "host_networking": "",
-    "host_networking_topology": "",
-    "host_memory_configuration": "24x 128GB DDR5",
-    "accelerators_per_node": "8",
-    "accelerator_model_name": "AMD Instinct MI350X 288GB HBM3e",
-    "accelerator_host_interconnect": "PCIe 5.0 x16",
-    "accelerator_frequency": "",
-    "accelerator_on-chip_memories": "",
-    "accelerator_memory_configuration": "HBM3",
-    "accelerator_memory_capacity": "288GB",
-    "accelerator_interconnect": "",
-    "accelerator_interconnect_topology": "",
-    "cooling": "air",
-    "hw_notes": "",
-    "framework": "tinygrad, branch mlperf_training_v6.0",
-    "other_software_stack": {
-        "python": "3.12.3",
-        "ROCm": "7.1.1"
-    },
-    "operating_system": "Ubuntu 24.04.3 LTS",
-    "sw_notes": ""
-  }

examples/mlperf/training_submission_v6.1/tinycorp/systems/tinybox_green.json

@@ -1,38 +0,0 @@
-{
-  "submitter": "tinycorp",
-  "division": "closed",
-  "status": "Available on-premise",
-  "system_name": "tinybox green",
-  "number_of_nodes": "1",
-  "host_processors_per_node": "1",
-  "host_processor_model_name": "AMD EPYC 7532",
-  "host_processor_core_count": "32",
-  "host_processor_vcpu_count": "64",
-  "host_processor_frequency": "",
-  "host_processor_caches": "",
-  "host_processor_interconnect": "",
-  "host_memory_capacity": "128GB",
-  "host_storage_type": "NVMe SSD",
-  "host_storage_capacity": "4 TB raid array + 1 TB boot",
-  "host_networking": "",
-  "host_networking_topology": "",
-  "host_memory_configuration": "8x 16GB DDR4",
-  "accelerators_per_node": "6",
-  "accelerator_model_name": "NVIDIA GeForce RTX 4090",
-  "accelerator_host_interconnect": "PCIe 4.0 x16",
-  "accelerator_frequency": "",
-  "accelerator_on-chip_memories": "",
-  "accelerator_memory_configuration": "GDDR6X",
-  "accelerator_memory_capacity": "24GB",
-  "accelerator_interconnect": "",
-  "accelerator_interconnect_topology": "",
-  "cooling": "air",
-  "hw_notes": "",
-  "framework": "tinygrad, branch mlperf_training_v5.0",
-  "other_software_stack": {
-    "python": "3.10.12",
-    "CUDA": "12.4"
-  },
-  "operating_system": "Ubuntu 22.04.4",
-  "sw_notes": ""
-}

examples/mlperf/training_submission_v6.1/tinycorp/systems/tinybox_red.json

@@ -1,37 +0,0 @@
-{
-  "submitter": "tinycorp",
-  "division": "closed",
-  "status": "Available on-premise",
-  "system_name": "tinybox red",
-  "number_of_nodes": "1",
-  "host_processors_per_node": "1",
-  "host_processor_model_name": "AMD EPYC 7532",
-  "host_processor_core_count": "32",
-  "host_processor_vcpu_count": "64",
-  "host_processor_frequency": "",
-  "host_processor_caches": "",
-  "host_processor_interconnect": "",
-  "host_memory_capacity": "128GB",
-  "host_storage_type": "NVMe SSD",
-  "host_storage_capacity": "4 TB raid array + 1 TB boot",
-  "host_networking": "",
-  "host_networking_topology": "",
-  "host_memory_configuration": "8x 16GB DDR4",
-  "accelerators_per_node": "6",
-  "accelerator_model_name": "AMD Radeon RX 7900 XTX",
-  "accelerator_host_interconnect": "PCIe 4.0 x16",
-  "accelerator_frequency": "",
-  "accelerator_on-chip_memories": "",
-  "accelerator_memory_configuration": "GDDR6",
-  "accelerator_memory_capacity": "24GB",
-  "accelerator_interconnect": "",
-  "accelerator_interconnect_topology": "",
-  "cooling": "air",
-  "hw_notes": "",
-  "framework": "tinygrad, branch mlperf_training_v5.0",
-  "other_software_stack": {
-    "python": "3.10.12"
-  },
-  "operating_system": "Ubuntu 22.04.4",
-  "sw_notes": ""
-}

examples/mnist_gan.py

@@ -3,7 +3,7 @@ import torch
 from torchvision.utils import make_grid, save_image
 from tinygrad.nn.state import get_parameters
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import trange, Context
+from tinygrad.helpers import trange
 from tinygrad.nn import optim
 from tinygrad.nn.datasets import mnist
 
@@ -71,7 +71,7 @@ def train_generator(optimizer, data_fake):
 if __name__ == "__main__":
   # data for training and validation
   X_train, _, _, _ = mnist()
-  ds_noise = Tensor.randn(64, 128)
+  ds_noise = Tensor.randn(64, 128, requires_grad=False)
   # parameters
   epochs, batch_size, k = 300, 512, 1
   sample_interval = epochs // 10
@@ -86,7 +86,7 @@ if __name__ == "__main__":
   optim_g = optim.Adam(get_parameters(generator), lr=0.0002, b1=0.5)  # 0.0002 for equilibrium!
   optim_d = optim.Adam(get_parameters(discriminator), lr=0.0002, b1=0.5)
   # training loop
-  with Context(TRAINING=1):
+  with Tensor.train():
     for epoch in (t := trange(epochs)):
       loss_g, loss_d = 0.0, 0.0
       for _ in range(n_steps):

examples/openpilot/compile3.py

@@ -1,4 +1,4 @@
-import os, sys, pickle, time, re, tempfile, struct, shutil, io
+import os, sys, pickle, time, re
 import numpy as np
 if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"
 
@@ -9,39 +9,6 @@ from tinygrad.nn.onnx import OnnxRunner
 
 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
 OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
-PICKLE_OOB = getenv("PICKLE_OOB")
-
-def dump_pickle(obj, f):
-  if PICKLE_OOB:
-    # allows pickling when buffers don't fit in (CPU) RAM
-    # from openpilot/selfdrive/modeld/helpers.py
-    with tempfile.TemporaryFile(dir=".") as tmp:
-      def buffer_callback(pb: pickle.PickleBuffer):
-        m = pb.raw()
-        tmp.write(struct.pack('<q', m.nbytes))
-        tmp.write(m)
-        pb.release() # keep peak ram at ~1 buffer
-      stream = io.BytesIO()
-      pickle.Pickler(stream, protocol=5, buffer_callback=buffer_callback).dump(obj)
-      opcodes = stream.getvalue()
-      f.write(struct.pack('<q', len(opcodes)))
-      f.write(opcodes)
-      tmp.seek(0)
-      shutil.copyfileobj(tmp, f)
-  else: pickle.dump(obj, f)
-
-def load_pickle(f):
-  if PICKLE_OOB:
-    # allows unpickling when buffers don't fit in (CPU) RAM
-    # from openpilot/selfdrive/modeld/helpers.py
-    opcodes = f.read(struct.unpack('<q', f.read(8))[0])
-    def buffers():
-      while (h := f.read(8)):
-        pb = pickle.PickleBuffer(bytearray(struct.unpack('<q', h)[0]))
-        f.readinto(pb)
-        yield pb
-    return pickle.load(io.BytesIO(opcodes), buffers=buffers())
-  else: return pickle.load(f)
 
 def compile(onnx_file):
   run_onnx = OnnxRunner(onnx_file)
@@ -54,15 +21,13 @@ def compile(onnx_file):
   # TODO this seems dumb
   input_types = {k:(dtypes.float32 if v is dtypes.float16 else v) for k,v in input_types.items()}
   Tensor.manual_seed(100)
-  # replace symbolic dimensions (e.g. 'b' for dynamic batch) with 1
-  input_shapes = {k:tuple(s if isinstance(s, int) else 1 for s in shp) for k,shp in input_shapes.items()}
   inputs = {k:Tensor(Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize().numpy(), device='NPY') for k,shp in sorted(input_shapes.items())}
   if not getenv("NPY_IMG"):
     inputs = {k:Tensor(v.numpy(), device=Device.DEFAULT).realize() if 'img' in k else v for k,v in inputs.items()}
   print("created tensors")
 
-  @TinyJit(prune=True)
-  def run_onnx_jit(**kwargs): return next(iter(run_onnx({k:v.to(Device.DEFAULT) for k,v in kwargs.items()}).values())).cast('float32')
+  run_onnx_jit = TinyJit(lambda **kwargs:
+                         next(iter(run_onnx({k:v.to(Device.DEFAULT) for k,v in kwargs.items()}).values())).cast('float32'), prune=True)
   for i in range(3):
     GlobalCounters.reset()
     print(f"run {i}")
@@ -75,7 +40,7 @@ def compile(onnx_file):
   kernel_calls = [u for u in run_onnx_jit.captured.linear.toposort(gate=lambda x: x.op not in kernel_asts)
                   if u.op is Ops.CALL and u.src[0].op in kernel_asts]
   print(f"captured {len(kernel_calls)} kernels")
-  if getenv("TEST", 1): np.testing.assert_equal(test_val, ret, "JIT run failed")
+  np.testing.assert_equal(test_val, ret, "JIT run failed")
   print("jit run validated")
 
   # check gated read_image usage
@@ -83,7 +48,7 @@ def compile(onnx_file):
   read_image_count = 0
   gated_read_image_count = 0
   for call in kernel_calls:
-    _, _, source, _ = call.src[0].src
+    _, _, _, source, _ = call.src[0].src
     src = source.arg
     kernel_count += 1
     read_image_count += src.count("read_image")
@@ -98,7 +63,8 @@ def compile(onnx_file):
   if (allowed_gated_read_image:=getenv("ALLOWED_GATED_READ_IMAGE", -1)) != -1:
     assert gated_read_image_count == allowed_gated_read_image, f"different gated read_image! {gated_read_image_count=}, {allowed_gated_read_image=}"
 
-  with open(OUTPUT, "wb") as f: dump_pickle(run_onnx_jit, f)
+  with open(OUTPUT, "wb") as f:
+    pickle.dump(run_onnx_jit, f)
   mdl_sz = os.path.getsize(onnx_file)
   pkl_sz = os.path.getsize(OUTPUT)
   print(f"mdl size is {mdl_sz/1e6:.2f}M")
@@ -119,7 +85,7 @@ def test_vs_compile(run, inputs, test_val=None):
     step_times.append((et-st)*1e3)
     print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")
 
-  if (assert_time:=getenv("ASSERT_MIN_STEP_TIME", 0.0)):
+  if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
     min_time = min(step_times)
     assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
 
@@ -136,7 +102,7 @@ def test_vs_compile(run, inputs, test_val=None):
 def test_vs_onnx(new_inputs, test_val, onnx_file, tol):
   import onnx
   import onnxruntime as ort
-
+  
   onnx_inputs = {k:v.numpy() for k,v in new_inputs.items()}
   onnx_model = onnx.load(onnx_file)
 
@@ -168,15 +134,15 @@ def bench(run, inputs):
 
 if __name__ == "__main__":
   if getenv("RUN_PICKLE"):
-    with open(OUTPUT, "rb") as f: pickle_loaded = load_pickle(f)
-    inputs = {name: Tensor(Tensor.randn(*view.shape, dtype=dtype).numpy(), device=device)
+    with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
+    inputs = {name: Tensor(Tensor.randn(*[int(s) for s in view.src[1].arg], dtype=dtype).numpy(), device=device)
               for name, (view, _vars, dtype, device) in zip(pickle_loaded.captured.expected_names, pickle_loaded.captured.expected_input_info)}
     test_vs_compile(pickle_loaded, inputs)
   else:
     onnx_file = fetch(OPENPILOT_MODEL)
     inputs, outputs = compile(onnx_file)
 
-    with open(OUTPUT, "rb") as f: pickle_loaded = load_pickle(f)
+    with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
 
     test_vs_compile(pickle_loaded, inputs, outputs)
     if getenv("SELFTEST"):

examples/openpilot/load_pickle.py

@@ -1,6 +1,5 @@
-import sys
+import sys, pickle
 from extra.bench_log import WallTimeEvent, BenchEvent
-from examples.openpilot.compile3 import load_pickle
 from tinygrad.helpers import getenv
 
 PKL = sys.argv[1] if len(sys.argv) > 1 else "/tmp/openpilot.pkl"
@@ -8,7 +7,7 @@ PKL = sys.argv[1] if len(sys.argv) > 1 else "/tmp/openpilot.pkl"
 load_times = []
 
 for _ in range(10):
-  with WallTimeEvent(BenchEvent.STEP) as wte: load_pickle(open(PKL, 'rb'))
+  with WallTimeEvent(BenchEvent.STEP) as wte: pickle.load(open(PKL, 'rb'))
   load_times.append(wte.time)
   print(f"pickle load: {wte.time:6.2f} s")
 

examples/stunning_mnist.py

@@ -5,7 +5,7 @@
 #  - symbolic removal
 
 from examples.beautiful_mnist import Model
-from tinygrad import Tensor, nn, getenv, GlobalCounters, Variable, Context
+from tinygrad import Tensor, nn, getenv, GlobalCounters, Variable
 from tinygrad.nn.datasets import mnist
 from tinygrad.helpers import trange
 
@@ -26,7 +26,7 @@ if __name__ == "__main__":
   X_samp, Y_samp = X_train[samples], Y_train[samples]
   print("*** got samples")
 
-  with Context(TRAINING=1):
+  with Tensor.train():
     """
     i = UOp.range(samples.shape[0])  # TODO: fix range function on UOp
     losses = model(X_samp[i]).sparse_categorical_crossentropy(Y_samp[i]).backward().contract(i)

examples/vgg7.py

@@ -164,8 +164,8 @@ elif cmd == "train":
     x_img = image_load(samples_base + "/" + str(sample_idx) + "a.png")
     y_img = image_load(samples_base + "/" + str(sample_idx) + "b.png")
 
-    sample_x = Tensor(x_img)
-    sample_y = Tensor(y_img)
+    sample_x = Tensor(x_img, requires_grad = False)
+    sample_y = Tensor(y_img, requires_grad = False)
 
     # magic code roughly from readme example
     # An explanation, in case anyone else has to go down this path:

examples/yolov8.py

@@ -193,8 +193,8 @@ class SPPF:
     self.cv1 = Conv_Block(c1, c_, 1, 1, padding=None)
     self.cv2 = Conv_Block(c_ * 4, c2, 1, 1, padding=None)
 
-    # Pad with -inf to match PyTorch's MaxPool2d behavior.
-    self.maxpool = lambda x : x.pad((k // 2, k // 2, k // 2, k // 2), value=float('-inf')).max_pool2d(kernel_size=k, stride=1)
+    # TODO: this pads with 0s, whereas torch function pads with -infinity. This results in a < 2% difference in prediction which does not make a difference visually.
+    self.maxpool = lambda x : x.pad((k // 2, k // 2, k // 2, k // 2)).max_pool2d(kernel_size=k, stride=1)
 
   def __call__(self, x):
     x = self.cv1(x)

extra/amdpci/am_smi.py

@@ -1,14 +1,14 @@
 #!/usr/bin/env python3
 
 import time, mmap, sys, shutil, os, glob, subprocess, argparse, collections
-from tinygrad.helpers import DEBUG, NO_COLOR, colored, ansilen
+from tinygrad.helpers import DEBUG, colored, ansilen
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.autogen.am import am
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager, AMPageTableEntry
 from tinygrad.runtime.support.am.ip import AM_SOC, AM_GMC, AM_IH, AM_PSP, AM_SMU, AM_GFX, AM_SDMA
 
-def bold(s): return s if NO_COLOR else f"\033[1m{s}\033[0m"
+def bold(s): return f"\033[1m{s}\033[0m"
 
 def trim(s:str, length:int) -> str:
   if len(s) > length: return s[:length-3] + "..."
@@ -64,7 +64,7 @@ def get_bar0_size(pcibus):
 
 class AMSMI(AMDev):
   def __init__(self, pcibus, vram_bar:MMIOInterface, doorbell_bar:MMIOInterface, mmio_bar:MMIOInterface):
-    self.pcibus, self.devfmt = pcibus, pcibus
+    self.pcibus = pcibus
     self.vram, self.doorbell64, self.mmio = vram_bar, doorbell_bar, mmio_bar
     self.pci_state = self.read_pci_state()
     if self.pci_state == "D0": self._init_from_d0()
@@ -91,7 +91,6 @@ class SMICtx:
     self.prev_lines_cnt = 0
     self.prev_terminal_width = 0
     self.prev_terminal_height = 0
-    self.prev_metrics = {}
 
     remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:", "Processing accelerators:"]
     lspci = subprocess.check_output(["lspci"]).decode("utf-8").splitlines()
@@ -236,29 +235,6 @@ class SMICtx:
       case (13,0,12): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.SocketPowerLimit)
       case _: return metrics.SmuMetrics.AverageSocketPower, metrics.SmuMetrics.dGPU_W_MAX
 
-  def get_throttle_info(self, dev, metrics):
-    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6)|(13,0,12):
-        throttle_fields = [('ProchotResidencyAcc', 'Prochot'), ('PptResidencyAcc', 'PPT'),
-                           ('SocketThmResidencyAcc', 'Socket Thm'), ('VrThmResidencyAcc', 'VR Thm'), ('HbmThmResidencyAcc', 'HBM Thm')]
-        prev = self.prev_metrics.get(dev.pcibus)
-        active = []
-        if prev is not None:
-          acc_delta = metrics.AccumulationCounter - prev.AccumulationCounter
-          if acc_delta > 0:
-            for field, name in throttle_fields:
-              delta = getattr(metrics, field) - getattr(prev, field)
-              if delta > 0 and (pct := min(100, (delta * 100 + acc_delta // 2) // acc_delta)) > 0: active.append((name, pct))
-        return active
-      case _:
-        smu_mod = dev.smu.smu_mod
-        throttler_names = {getattr(smu_mod, a): a[len('THROTTLER_'):-len('_BIT')]
-                           for a in dir(smu_mod) if a.startswith('THROTTLER_') and a.endswith('_BIT')}
-        active = []
-        for i, pct in enumerate(metrics.SmuMetrics.ThrottlingPercentage):
-          if pct > 0: active.append((throttler_names.get(i, f"UNK_{i}"), int(pct)))
-        return active
-
   def get_mem_usage(self, dev):
     usage = 0
     pt_stack = [dev.mm.root_page_table]
@@ -276,7 +252,7 @@ class SMICtx:
     return usage
 
   def draw(self, once):
-    terminal_width, terminal_height = shutil.get_terminal_size(fallback=(231, 24))
+    terminal_width, terminal_height = shutil.get_terminal_size()
     if not once and (self.prev_terminal_width != terminal_width or self.prev_terminal_height != terminal_height):
       os.system('clear')
     self.prev_terminal_width, self.prev_terminal_height = terminal_width, terminal_height
@@ -305,13 +281,6 @@ class SMICtx:
                     + [f"MEM Activity {draw_bar(self.get_mem_activity(dev, metrics) / 100, activity_line_width)}"] \
                     + [f"MEM Usage    {draw_bar(mem_used / mem_total, activity_line_width, opt_text=mem_fmt)}"] \
 
-      throttle_info = self.get_throttle_info(dev, metrics)
-      if throttle_info:
-        throttle_text = colored(', '.join(f"{name} {pct}%" for name, pct in throttle_info), "red")
-      else:
-        throttle_text = colored("None", "green")
-      activity_line += [f"Throttle     {throttle_text}" + " " * (activity_line_width + 2)]
-
       temps_data, temps_data_compact = self.get_temps(dev, metrics), self.get_temps(dev, metrics, compact=True)
       temps_table = ["=== Temps (°C) ==="] + [f"{name:<16}: {color_temp(val)}" for name, val in temps_data.items()]
       temps_table_compact = ["Temps (°C):" + '/'.join([f"{color_temp(val)} {name}" for name, val in temps_data_compact.items()])]
@@ -355,8 +324,6 @@ class SMICtx:
 
       dev_content.append(device_line + activity_line + same_line([temps_table, power_table, frequency_table]))
 
-    self.prev_metrics = {dev.pcibus: m for dev, m in dev_metrics.items() if m is not None}
-
     raw_text = 'AM Monitor'.center(terminal_width) + "\n" + "=" * terminal_width + "\n\n"
     for i in range(0, len(dev_content), 2):
       if i + 1 < len(dev_content): raw_text += '\n'.join(same_line([dev_content[i], dev_content[i+1]], split=padding))

extra/benchmark_llm.py

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

extra/export_model.py

@@ -1,5 +1,5 @@
 from typing import Tuple, Dict, List, Optional
-from tinygrad.dtype import DType, dtypes, AddrSpace
+from tinygrad.dtype import DType, dtypes
 from tinygrad.tensor import Tensor
 from tinygrad.device import Device, Buffer
 from tinygrad.engine.jit import TinyJit
@@ -23,7 +23,7 @@ def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], Li
 
   def name_of(bu:UOp, is_out:bool) -> str:
     nonlocal n
-    if bu.op is Ops.PARAM: key, name, size = ("in", bu.arg.slot), f"input{bu.arg.slot}", prod(bu.shape)*bu.dtype.itemsize
+    if bu.op is Ops.PARAM: key, name, size = ("in", bu.arg), f"input{bu.arg}", prod(bu.shape)*bu.dtype.itemsize
     else:
       b = bu.buffer
       key, size = (id(b.base), b.offset, b.size, b.dtype), b.size*b.dtype.itemsize
@@ -38,8 +38,8 @@ def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], Li
     arg_uops = [b for b in call.src[1:] if b.op is not Ops.BIND]
     prg = to_program(call.src[0], Device[arg_uops[0].device].renderer)
     info = prg.arg
-    functions[info.function_name] = prg.src[2].arg
-    cargs = [name_of(bu, i == 0) for i, bu in enumerate(arg_uops)] + list(info.vars)
+    functions[info.function_name] = prg.src[3].arg
+    cargs = [name_of(bu, i == 0) for i, bu in enumerate(arg_uops)] + [v for v in info.vars if v.op is Ops.DEFINE_VAR]
     statements.append((info.function_name, cargs, info.global_size, info.local_size))
 
   return functions, statements, {name:(size, dtype, key) for name, size, dtype, key in bufs.values()}, bufs_to_save
@@ -241,8 +241,8 @@ export default {model_name};
 def export_model(model, target:str, *inputs, model_name: Optional[str] = "model", stream_weights=False):
   assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, f"only {', '.join(EXPORT_SUPPORTED_DEVICE)} are supported"
 
-  # NOTE: NUM_CPU_THREADS=1, since export does not support threading
-  with Context(JIT=2, NUM_CPU_THREADS=1): linear, output_bufs = jit_model(model, *inputs)
+  # NOTE: CPU_COUNT=1, since export does not support threading
+  with Context(JIT=2, CPU_COUNT=1): linear, output_bufs = jit_model(model, *inputs)
   functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs)
   state = get_state_dict(model)
   weight_names = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None}
@@ -253,18 +253,17 @@ def export_model(model, target:str, *inputs, model_name: Optional[str] = "model"
   symbolic_vars = OrderedDict()
   for i, (_, args, global_size, _) in enumerate(statements):
     for j, var in enumerate(args):
-      if getattr(var, "op", None) is Ops.PARAM and var.addrspace is AddrSpace.ALU and var.arg.name is not None:
+      if getattr(var, "op", None) is Ops.DEFINE_VAR and isinstance(getattr(var, "arg", None), tuple) and isinstance(var.arg[0], str):
         if var not in symbolic_vars:
-          symbolic_vars[var] = var.expr
+          symbolic_vars[var] = var.arg[0]
           bufs[symbolic_vars[var]] = (var.dtype.itemsize, var.dtype, symbolic_vars[var])
         statements[i][1][j] = symbolic_vars[var]
 
     if global_size:
       for j, dim in enumerate(global_size):
-        if getattr(dim, "op", None) is Ops.ADD and len(dim.src) == 2 and \
-           any(s.op is Ops.PARAM and s.addrspace is AddrSpace.ALU for s in dim.src) and any(s.op is Ops.CONST for s in dim.src):
+        if getattr(dim, "op", None) is Ops.ADD and len(dim.src) == 2 and {dim.src[0].op, dim.src[1].op} == {Ops.DEFINE_VAR, Ops.CONST}:
           name, val = dim.src if dim.src[1].op is Ops.CONST else reversed(dim.src)
-          global_size[j] = f"_{name.expr}[0] + {val.val}"
+          global_size[j] = f"_{name.arg[0]}[0] + {val.arg}"
 
   prg = ""
   if target == "clang":

extra/f16_decompress.py

@@ -5,10 +5,10 @@ def bit_extract(x: Tensor, e: int, s: int) -> Tensor:
   return (x >> s) & mask
 
 def u16_to_f16(x: Tensor) -> Tensor:
-  sign = bit_extract(x, 15, 15).bool()
+  sign = bit_extract(x, 15, 15).float()
   exponent = bit_extract(x, 14, 10).float()
   fraction = bit_extract(x, 9, 0).float()
-  return sign.where(-1, 1) * exponent.bool().where((exponent - 15.0).exp2() * (1 + fraction / 1024.0), 6.103515625e-5 * (fraction / 1024.0))
+  return sign.where(-1, 1) * exponent.where((exponent - 15.0).exp2() * (1 + fraction / 1024.0), 6.103515625e-5 * (fraction / 1024.0))
 
 def u32_to_f16(oo: Tensor) -> Tensor:
   f1 = u16_to_f16(oo>>16)

extra/fp8/fp8_linear.py

@@ -18,13 +18,13 @@ def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
   SEQ = inp.shape[1]
   OUT = weight.shape[0]
   IN = weight.shape[-1]
-  seq_idx = UOp.range(SEQ, 2)
-  out_idx = UOp.range(OUT, 3)
-  batch_idx = UOp.range(output.size//SEQ//OUT, 1)
+  seq_idx = UOp.range(SEQ, 2, AxisType.LOOP)
+  out_idx = UOp.range(OUT, 3, AxisType.LOOP)
+  batch_idx = UOp.range(output.size//SEQ//OUT, 1, AxisType.LOOP)
   reduce_idx = UOp.range(IN, 0, AxisType.REDUCE)
   product = (inp.index((seq_idx*IN+reduce_idx+batch_idx*IN*SEQ)) * weight.index((out_idx*IN+reduce_idx))).cast(dtypes.float)
   reduced = product.reduce(reduce_idx, arg=Ops.ADD)
-  store_op = output.index((seq_idx*OUT+out_idx+batch_idx*OUT*SEQ)).store(reduced).end(batch_idx, seq_idx, out_idx)
+  store_op = output.index((seq_idx*OUT+out_idx+batch_idx*OUT*SEQ), ptr=True).store(reduced).end(batch_idx, seq_idx, out_idx)
   return store_op.sink(arg=KernelInfo(name=f"fp8_matmul_{inp.shape}x{weight.shape}"))
 
 def custom_matmul_backward(gradient: UOp, kernel: UOp) -> tuple[UOp, UOp]:
@@ -53,7 +53,7 @@ class FP8Linear:
     x_fp8, x_scale = quantize_to_fp8(x)
     GPUS = self.weight.device
     if isinstance(GPUS, tuple) and len(GPUS) > 1:
-      y = Tensor(Tensor.empty((batch//len(GPUS), seq, self.weight.shape[0]), dtype=dtypes.float, device=GPUS).uop.unshard(0), device=GPUS)
+      y = Tensor(Tensor.empty((batch//len(GPUS), seq, self.weight.shape[0]), dtype=dtypes.float, device=GPUS).uop.multi(0), device=GPUS)
     else:
       y = Tensor.empty((batch, seq, self.weight.shape[0]), dtype=dtypes.float)
     y = Tensor.custom_kernel(y, x_fp8, w_fp8, fxn=custom_matmul, grad_fxn=custom_matmul_backward)[0]

extra/gemm/amd_asm_matmul.py

@@ -458,11 +458,10 @@ def test_matmul():
   def asm_kernel(A:UOp, B:UOp, C:UOp) -> UOp:
     gidxs = [UOp.special(n, f"gidx{i}") for i,n in enumerate(grid)]
     lidxs = [UOp.special(n, f"lidx{i}") for i,n in enumerate(local)]
-    lds_size = max(LDS_SIZE, 65536//getenv("LIMIT_OCC", 65536))
-    lds = UOp.placeholder((lds_size,), dtypes.uint8, 0, AddrSpace.LOCAL)
+    lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=max(LDS_SIZE, 65536//getenv("LIMIT_OCC", 65536)), addrspace=AddrSpace.LOCAL), (), 'lds')
     sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs, arg=KernelInfo(name=colored("kernel", "cyan"),
                                                                                   estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
   c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
   linear = c.schedule_linear()
 

extra/gemm/amd_copy_matmul.py

@@ -1,5 +1,5 @@
 from tinygrad import Device, UOp, getenv
-from tinygrad.uop.ops import AxisType, KernelInfo
+from tinygrad.uop.ops import AxisType, KernelInfo, Ops
 from tinygrad.dtype import AddrSpace, dtypes
 
 N = getenv("N", 4096)
@@ -46,8 +46,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
   # -- GLOBAL -> LOCAL --
   # wmma: spatial outer, k inner (k contiguous for vectorized WMMA tile loads)
   # gemm: k outer, spatial inner
-  A_local = UOp.placeholder((BLOCK_M, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_M), a.dtype, slot=0, addrspace=AddrSpace.LOCAL)
-  B_local = UOp.placeholder((BLOCK_N, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_N), b.dtype, slot=1, addrspace=AddrSpace.LOCAL)
+  A_local = UOp.placeholder((BLOCK_M, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_M), a.dtype.base, slot=0, addrspace=AddrSpace.LOCAL)
+  B_local = UOp.placeholder((BLOCK_N, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_N), b.dtype.base, slot=1, addrspace=AddrSpace.LOCAL)
 
   a = a.reshape(K // BLOCK_K, BLOCK_K, BLOCK_M)
   b = b.reshape(K // BLOCK_K, BLOCK_K, BLOCK_N)
@@ -58,20 +58,20 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
   B_copy = B_local.permute((1,0)) if use_wmma else B_local
   A_store = A_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(a[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
   B_store = B_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(b[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
-  # NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
-  A_local, B_local = A_local.after(A_store, B_store), B_local.after(A_store, B_store)
+  barrier = UOp.barrier(A_store, B_store)
+  A_local, B_local = A_local.after(barrier), B_local.after(barrier)
 
   # -- COMPUTE --
   lane_m, lane_n = lane // LANES_PER_WAVE_N, lane % LANES_PER_WAVE_N
 
   # accumulator (unified: both paths use (TM, TN) with scalar dtypes.float)
   acc = UOp.placeholder((TM, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
-  acc = acc.after(acc.store(acc.zeros_like(buffer=False)))
+  acc = acc.after(acc.store(acc.zeros_like()))
 
   if use_wmma:
     k = UOp.range(BLOCK_K // WMMA_K, 101, AxisType.REDUCE)
-    tile_m = UOp.range(TM // WMMA_ACC, 200)
-    tile_n = UOp.range(TN, 201)
+    tile_m = UOp.range(TM // WMMA_ACC, 200, AxisType.LOOP)
+    tile_n = UOp.range(TN, 201, AxisType.LOOP)
 
     acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0,2,1)[tile_m, tile_n]
     a_frag = A_local.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_K // WMMA_K, WMMA_K)[wave_m, tile_m, lane_n, k]
@@ -80,7 +80,7 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
       # NOTE: since this is part of K, these 2 can be anywhere in the frags and long as a and b match
       a_frag = a_frag.reshape(2, 8)[lane_m, :]
       b_frag = b_frag.reshape(2, 8)[lane_m, :]
-    wmma = UOp.wmma(a_frag, b_frag, acc_frag.after(k), (16, 16, 16), 'AMD', 32)
+    wmma = UOp(Ops.SHAPED_WMMA, dtypes.float, (a_frag, b_frag, acc_frag.after(k)), arg=((16, 16, 16), 'AMD', 32))
     acc_store = acc_frag.store(wmma).end(tile_m, tile_n)
   else:
     # registers for LOCAL -> REG
@@ -96,8 +96,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
     b_frag = b_frag.reshape(1, TN).expand(TM, TN)
     acc_store = acc.store(acc.after(k) + (a_frag * b_frag))
 
-  # store accumulator and loop (the barrier at the end of the loop is implied by the LOCAL buffers stored and loaded in the loop)
-  acc = acc.after(acc_store.end(k).end(k_tile))
+  # store accumulator and loop
+  acc = acc.after(acc_store.end(k).barrier().end(k_tile))
 
   # store accumulator to output (unified)
   c = c.reshape(WAVES_M, TM//UNROLL_M, LANES_PER_WAVE_M, UNROLL_M,

extra/gemm/amd_flash_attention.py

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

extra/gemm/amd_matmul.py

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

extra/gemm/amd_uop_matmul.py

@@ -28,10 +28,10 @@ REG_TILES_PER_WAVE_M = BLOCK_M // (WAVES_PER_BLOCK_M * LANES_PER_WAVE_M * TM)
 assert WAVES_PER_BLOCK_M*REG_TILES_PER_WAVE_M*LANES_PER_WAVE_M*TM == BLOCK_M, "M reshape is wrong"
 assert WAVES_PER_BLOCK_N*REG_TILES_PER_WAVE_N*LANES_PER_WAVE_N*TN == BLOCK_N, "N reshape is wrong"
 
-def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.WEAK): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
+def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.LOOP): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
 def copy(dest:UOp, src:UOp, rng:int, upcast=False):
   assert dest.shape == src.shape
-  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.WEAK)
+  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.LOOP)
   return dest[*rngs].store(src[*rngs]).end(*rngs)
 
 def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
@@ -66,8 +66,9 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
   B_local = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
   B_local_store = copy(B_local.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)
 
-  # NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
-  A_local, B_local = A_local.after(A_local_store, B_local_store), B_local.after(A_local_store, B_local_store)
+  # TODO: can we automate barrier?
+  barrier = UOp.barrier(A_local_store, B_local_store)
+  A_local, B_local = A_local.after(barrier), B_local.after(barrier)
 
   # open inner k range
   k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)
@@ -101,7 +102,7 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
   sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iter_m, t_m] * B_row[iter_n, t_n]).end(iter_m, iter_n, t_m, t_n)
 
   # Close k, sync, and close K tiles
-  sink = sink.end(k).end(k_tile_range)
+  sink = sink.end(k).barrier().end(k_tile_range)
 
   # ---------------------------
   # REG -> GLOBAL (epilogue)
@@ -121,7 +122,7 @@ def eval_custom_matmul(fxn, dt=dtypes.float):
   with Context(DEBUG=0): Tensor.realize(a, b)
 
   ets = []
-  with Context(DEBUG=max(2, DEBUG.value)):
+  with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2 if dt == dtypes.half else 0):
     for _ in range(NUM_RUNS):
       GlobalCounters.reset()
       tst = Tensor.custom_kernel(c, a, b, fxn=fxn)[0].realize()

extra/gemm/amx.py

@@ -0,0 +1,180 @@
+#!/usr/bin/env python3
+import numpy as np
+import time
+import sys
+np.set_printoptions(linewidth=160)
+np.set_printoptions(linewidth=1000, threshold=10000000000, suppress=False)
+from tinygrad.runtime.ops_llvm import LLVMDevice, LLVMProgram, LLVMCompiler
+from llvmlite import ir  # type: ignore
+from tinygrad.helpers import flat_mv
+from tinygrad.device import MallocAllocator
+
+# https://github.com/corsix/amx/blob/main/Instructions.md
+# 12 lines for AMX support
+from functools import partialmethod
+class AMX:
+  @staticmethod
+  def nop_op_imm5(op, imm5, builder): builder.asm(ir.FunctionType(ir.VoidType(), []), f".word (0x201000 + ({op} << 5) + {imm5}); amx op {op} imm {imm5}", "", tuple(), True)
+  @staticmethod
+  def op_gpr(op, builder, gpr): builder.asm(ir.FunctionType(ir.VoidType(), [ir.IntType(64)]), f".word (0x201000 + ({op} << 5) + 0$0 - ((0$0 >> 4) * 6)); amx op {op} reg $0", "r", (gpr,), True)
+  set, clr = partialmethod(nop_op_imm5, 17, 0), partialmethod(nop_op_imm5, 17, 1)
+  ldx, ldy, stx, sty = partialmethod(op_gpr, 0), partialmethod(op_gpr, 1), partialmethod(op_gpr, 2), partialmethod(op_gpr, 3)
+  ldz, stz, ldzi, stzi = partialmethod(op_gpr, 4), partialmethod(op_gpr, 5), partialmethod(op_gpr, 6), partialmethod(op_gpr, 7)
+  extrx, extry = partialmethod(op_gpr, 8), partialmethod(op_gpr, 9)
+  fma64, fms64, fma32, fms32 = partialmethod(op_gpr, 10), partialmethod(op_gpr, 11), partialmethod(op_gpr, 12), partialmethod(op_gpr, 13)
+  mac16, fma16, fms16 = partialmethod(op_gpr, 14), partialmethod(op_gpr, 15), partialmethod(op_gpr, 16)
+  vecint, vecfp, matint, matfp, genlut = partialmethod(op_gpr, 18), partialmethod(op_gpr, 19), partialmethod(op_gpr, 20), partialmethod(op_gpr, 21), partialmethod(op_gpr, 22)
+
+def int_const(x): return ir.Constant(ir.IntType(64), x)
+
+
+N = 4096
+# N = 1024
+# N = 64
+
+BW = N*N*4
+
+# matrix is 64M, max load bandwidth is 57 GB/s
+# cache line looks like 256 bytes (64 floats)
+
+na = np.zeros((256), dtype=np.float32)
+# na = np.zeros((N, N), dtype=np.float32)
+nb = np.random.randn(N, N).astype(np.float32)
+nc = np.random.randn(N, N).astype(np.float32)
+
+ns = nb.reshape(-1, 32).sum(axis=0)
+
+a = MallocAllocator.alloc(na.nbytes)
+b = MallocAllocator.alloc(nb.nbytes)
+c = MallocAllocator.alloc(nc.nbytes)
+
+MallocAllocator._copyin(b, flat_mv(nb.data))
+MallocAllocator._copyin(c, flat_mv(nc.data))
+
+module = ir.Module(name=__file__)
+func = ir.Function(module, ir.FunctionType(ir.IntType(64), [ir.FloatType().as_pointer()]*3), name='exec')
+
+# load all
+entry = ir.IRBuilder(func.append_basic_block(name="entry"))
+zm, xm, ym = [entry.ptrtoint(func.args[i], ir.IntType(64)) for i in range(3)]
+
+loop_1 = ir.IRBuilder(func.append_basic_block(name="loop_y"))
+loop_1_exit = ir.IRBuilder(func.append_basic_block(name="loop_y_exit"))
+exit = ir.IRBuilder(func.append_basic_block(name="exit"))
+
+y = loop_1.phi(ir.IntType(64), name="y")
+y.add_incoming(int_const(0), entry._block)
+yp = loop_1_exit.add(y, int_const(32*2))
+y.add_incoming(yp, loop_1_exit._block)
+
+prefetch_function = ir.Function(module, ir.FunctionType(ir.VoidType(), [ir.PointerType(ir.FloatType()), ir.IntType(32), ir.IntType(32), ir.IntType(32)]), name="llvm.prefetch")
+
+xptr = y
+addr = loop_1_exit.add(xm, loop_1_exit.mul(int_const(4), xptr))
+
+#prefetch_ptr = loop_1_exit.inttoptr(loop_1_exit.add(addr, int_const(128)), ir.PointerType(ir.FloatType()))
+#loop_1_exit.call(prefetch_function, [prefetch_ptr, ir.IntType(32)(0), ir.IntType(32)(2), ir.IntType(32)(1)])
+
+AMX.ldx(loop_1_exit, loop_1_exit.add(int_const(1<<62), addr))
+xptr = loop_1_exit.add(xptr, int_const(32))
+AMX.ldy(loop_1_exit, loop_1_exit.add(int_const(1<<62), loop_1_exit.add(xm, loop_1_exit.mul(int_const(4), xptr))))
+
+AMX.fma32(loop_1_exit, int_const(1 << 63 | 1 << 28))
+AMX.fma32(loop_1_exit, int_const(1 << 63 | 1 << 28 | 1 << 20 | (16*4)<<10))
+AMX.fma32(loop_1_exit, int_const(1 << 63 | 1 << 29))
+AMX.fma32(loop_1_exit, int_const(1 << 63 | 1 << 29 | 1 << 20 | (16*4)))
+
+AMX.set(entry)
+
+AMX.stz(exit, exit.add(zm, int_const(1 << 62 | (0 << 56) | 0)))
+AMX.clr(exit)
+
+entry.branch(loop_1._block)
+loop_1.branch(loop_1_exit._block)
+loop_1_exit.cbranch(loop_1_exit.icmp_unsigned("==", yp, int_const(N*N)), exit._block, loop_1._block)
+exit.ret(int_const(0))
+
+device = LLVMDevice("llvm")
+prog = LLVMProgram(device, "exec", LLVMCompiler(device).compile(str(module)))
+
+"""
+loop_1 = ir.IRBuilder(func.append_basic_block(name="loop_y"))
+loop_2 = ir.IRBuilder(func.append_basic_block(name="loop_x"))
+loop_3 = ir.IRBuilder(func.append_basic_block(name="loop_k"))
+loop_3_exit = ir.IRBuilder(func.append_basic_block(name="loop_k_exit"))
+loop_2_exit = ir.IRBuilder(func.append_basic_block(name="loop_x_exit"))
+loop_1_exit = ir.IRBuilder(func.append_basic_block(name="loop_y_exit"))
+
+y = loop_1.phi(ir.IntType(64), name="y")
+x = loop_2.phi(ir.IntType(64), name="x")
+k = loop_3.phi(ir.IntType(64), name="k")
+
+exit = ir.IRBuilder(func.append_basic_block(name="exit"))
+
+AMX.set(loop_2)
+
+# stride
+xptr = loop_3_exit.add(x, loop_3_exit.mul(k, int_const(N)))
+yptr = loop_3_exit.add(y, loop_3_exit.mul(k, int_const(N)))
+
+# if you are okay with the wrong answer, this is faster
+#xptr = loop_3_exit.add(x, loop_3_exit.mul(k, int_const(32)))
+#yptr = loop_3_exit.add(y, loop_3_exit.mul(k, int_const(32)))
+
+# double loads load 32 floats
+AMX.ldx(loop_3_exit, loop_3_exit.add(int_const(1<<62), loop_3_exit.add(xm, loop_3_exit.mul(int_const(4), xptr))))
+AMX.ldy(loop_3_exit, loop_3_exit.add(int_const(1<<62), loop_3_exit.add(ym, loop_3_exit.mul(int_const(4), yptr))))
+
+# <Z row> <X offset> <Y offset>
+AMX.fma32(loop_3_exit, int_const(0<<20 | (0*16*4)<<10 | (0*16*4)))
+AMX.fma32(loop_3_exit, int_const(1<<20 | (1*16*4)<<10 | (0*16*4)))
+AMX.fma32(loop_3_exit, int_const(2<<20 | (0*16*4)<<10 | (1*16*4)))
+AMX.fma32(loop_3_exit, int_const(3<<20 | (1*16*4)<<10 | (1*16*4)))
+
+# store
+gptr = loop_2_exit.mul(loop_2_exit.add(loop_2.mul(y, int_const(N)), x), int_const(4))
+zmp = loop_2_exit.add(zm, gptr)
+for j in range(2):
+  for r in range(16):
+    z_row = j*2
+    ptr = ((j*16)+r)*N
+    AMX.stz(loop_2_exit, loop_2_exit.add(zmp, int_const(1 << 62 | ((r*4+z_row) << 56) | ptr*4)))
+AMX.clr(loop_2_exit)
+
+yp = loop_1_exit.add(y, int_const(32))
+xp = loop_2_exit.add(x, int_const(32))
+kp = loop_3_exit.add(k, int_const(1))
+
+y.add_incoming(int_const(0), entry._block)
+x.add_incoming(int_const(0), loop_1._block)
+k.add_incoming(int_const(0), loop_2._block)
+y.add_incoming(yp, loop_1_exit._block)
+x.add_incoming(xp, loop_2_exit._block)
+k.add_incoming(kp, loop_3_exit._block)
+
+entry.branch(loop_1._block)
+loop_1.branch(loop_2._block)
+loop_2.branch(loop_3._block)
+loop_3.branch(loop_3_exit._block)
+loop_3_exit.cbranch(loop_3_exit.icmp_unsigned("==", kp, int_const(N)), loop_2_exit._block, loop_3._block)
+loop_2_exit.cbranch(loop_2_exit.icmp_unsigned("==", xp, int_const(N)), loop_1_exit._block, loop_2._block)
+loop_1_exit.cbranch(loop_1_exit.icmp_unsigned("==", yp, int_const(N)), exit._block, loop_1._block)
+exit.ret(int_const(0))
+
+device = LLVMDevice("llvm")
+prog = LLVMProgram(device, "exec", LLVMCompiler(device).compile(str(module)))
+"""
+
+def timeit(fxn):
+  st = time.perf_counter()
+  et = fxn()
+  return time.perf_counter() - st
+
+tm = min([timeit(lambda: prog(a, b, c, N**2)) for _ in range(20)])
+MallocAllocator._copyout(flat_mv(na.data), a)
+print(f"{N*N:10d} {tm*1e6:9.2f} us, {BW*1e-9/tm:.2f} GB/s")
+
+np.testing.assert_allclose(na[:ns.shape[0]], ns, atol=1e-4, rtol=1e-4)
+
+# comp = (nb.T @ nc).T
+# np.testing.assert_allclose(na, comp, atol=1e-4, rtol=1e-5)

extra/gemm/intel_xmx.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python3
+import numpy as np
+from tinygrad.runtime.ops_cl import CLProgram, CLCompiler
+from tinygrad import Device, dtypes
+from tinygrad.device import Buffer
+from hexdump import hexdump
+
+# https://github.com/intel/intel-graphics-compiler/blob/master/documentation/visa/instructions/DPAS.md
+# https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroups.html
+# https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_matrix_multiply_accumulate.html
+# https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_split_matrix_multiply_accumulate.html
+# https://hc34.hotchips.org/assets/program/conference/day1/GPU%20HPC/Intel_s%20Ponte%20Vecchio%20GPU%20-%20Architecture%20Systems%20and%20Software%20FINAL.pdf
+
+device = Device["CL"]
+
+# NOTE: only the subgroup type 8 ones work
+prog = CLProgram(device, "test", CLCompiler(device, "test").compile(f"""
+__attribute__((intel_reqd_sub_group_size(8)))
+__kernel void test(__global float* data0, const __global int* data1, const __global int8* data2) {{
+  int lidx0 = get_local_id(0);
+  int a = data1[lidx0];
+  int8 b = data2[lidx0];
+  float out = intel_sub_group_f16_f16_matrix_mad_k16(a, b, 0.0f);
+  data0[lidx0] = out;
+}}
+"""))
+#with open("/tmp/test.elf", "wb") as f: f.write(prog.lib)
+
+a = Buffer("CL", 8, dtypes.float32).allocate()
+b = Buffer("CL", 0x10, dtypes.float16).allocate()
+c = Buffer("CL", 8*0x10, dtypes.float16).allocate()
+
+row = np.array([1,2,3,4,5,6,7,8,1,2,3,4,5,6,7,8], np.float16)
+mat = np.random.random((8, 0x10)).astype(np.float16)
+
+b.copyin(row.data)
+c.copyin(mat.data)
+ret = prog(a._buf, b._buf, c._buf, global_size=[1,1,1], local_size=[8,1,1], wait=True)
+print(ret)
+out = np.frombuffer(a.as_memoryview(), np.float32)
+real = row.astype(np.float32)@mat.T.astype(np.float32)
+print("out:", out)
+print("real", real)

extra/gemm/metal_uop_matmul.py

@@ -1,5 +1,5 @@
 from tinygrad import UOp, dtypes
-from tinygrad.uop.ops import AxisType, KernelInfo, AddrSpace
+from tinygrad.uop.ops import AxisType, Ops, KernelInfo, AddrSpace
 from extra.gemm.amd_uop_matmul import test_matmul
 
 N = 2048
@@ -20,17 +20,20 @@ def hand_spec_tc_cores():
 
   gk = UOp.range(N // 8, 0, AxisType.REDUCE)
 
-  a_tc = UOp.stack(*[mat_idx(a, gx, gk, warp, i) for i in range(2)])
-  b_tc = UOp.stack(*[mat_idx(b, gk, gy, warp, i) for i in range(2)])
+  a_tc = UOp.vectorize(*[mat_idx(a, gx, gk, warp, i) for i in range(2)])
+  b_tc = UOp.vectorize(*[mat_idx(b, gk, gy, warp, i) for i in range(2)])
 
   acc = UOp.placeholder((2,), dtypes.float, slot=0, addrspace=AddrSpace.REG)
   acc = acc[0].set(0.0)
   acc = acc[1].set(0.0)
 
-  acc_load = UOp.stack(acc.after(gk)[0], acc.after(gk)[1])
-  out = UOp.wmma(a_tc, b_tc, acc_load, (8, 8, 8), 'METAL', 32)
+  # TODO: make this simple
+  wmma_arg = ('WMMA_8_8_8_float_float', (8, 8, 8), dtypes.float, dtypes.float, 'METAL', 32, (((3, 2),), ((3, 2),), ((3, 2),)), ())
 
-  end_loop = UOp.group(*[acc[i].store(out.index(i)) for i in range(2)]).end(gk)
+  acc_load = UOp.vectorize(acc.after(gk)[0], acc.after(gk)[1])
+  out = UOp(Ops.WMMA, dtypes.float.vec(2), (a_tc, b_tc, acc_load), arg=wmma_arg)
+
+  end_loop = UOp.group(*[acc[i].store(out.gep(i)) for i in range(2)]).end(gk)
 
   sink = UOp.group(*[mat_idx(c.after(end_loop), gx, gy, warp, i).store(acc[i]) for i in range(2)])
   return sink.sink(arg=KernelInfo(name="custom_metal_matmul", opts_to_apply=())).simplify()

extra/gemm/mi350x_uop_matmul.py

@@ -6,7 +6,7 @@ os.environ["AMD_LLVM"] = "0"
 from tinygrad import Tensor, Context, dtypes, UOp, GlobalCounters
 from tinygrad.helpers import DEBUG, getenv
 from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import AxisType, KernelInfo
+from tinygrad.uop.ops import AxisType, KernelInfo, Ops
 
 WARP_SIZE = 64
 
@@ -77,9 +77,9 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
   B = B.reshape((K//BLOCK_K, BLOCK_K, N//BLOCK_N, BLOCK_N))
 
   # this is the big accumulator
-  acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float, 0, AddrSpace.REG)
+  acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float.vec(4), 0, AddrSpace.REG)
   assert acc.size*WARP_SIZE*WARPGROUP_SIZE*4 == BLOCK_M*BLOCK_N
-  acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const((0.0,)*4, dtypes.float), end=init_l)
+  acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const(dtypes.float.vec(4), 0.0), end=init_l)
 
   # create locals (note A is permuted, and the stride is changed to avoid bank conflicts)
   def make_locals(slot) -> tuple[UOp, UOp]:
@@ -114,8 +114,8 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
     K_inner_loop = UOp.range(BLOCK_K//TC_K, rng, AxisType.REDUCE)
 
     # load from locals into registers
-    Ar = UOp.placeholder((BLOCK_M//TC_M//WARPGROUP_SIZE,), dtypes.half, slot=1, addrspace=AddrSpace.REG)
-    Br = UOp.placeholder((BLOCK_N//TC_N,), dtypes.half, slot=2, addrspace=AddrSpace.REG)
+    Ar = UOp.placeholder((BLOCK_M//TC_M//WARPGROUP_SIZE,), dtypes.half.vec(8), slot=1, addrspace=AddrSpace.REG)
+    Br = UOp.placeholder((BLOCK_N//TC_N,), dtypes.half.vec(8), slot=2, addrspace=AddrSpace.REG)
 
     M_load_loop = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, rng+10)
     Asl = Asl.reshape((BLOCK_K//TC_K, TC_K, BLOCK_M//TC_M//WARPGROUP_SIZE, WARPGROUP_SIZE, TC_M))
@@ -137,7 +137,8 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
     acc_load = acc_after[N_inner_loop, M_inner_loop]
 
     # do WMMA
-    out = UOp.wmma(Ar[M_inner_loop], Br[N_inner_loop], acc_load, (16, 16, 32), 'AMD', 64)
+    wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
+    out = UOp(Ops.WMMA, dtypes.float.vec(4), (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg)
 
     # store back the acc
     acc_store = acc[N_inner_loop, M_inner_loop].store(out)
@@ -179,7 +180,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
   # store the acc into gmem
   cp_i, cp_j = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, 10004), UOp.range(BLOCK_N//TC_N, 10005)
   c_load = lambda i: C[gx, cp_i*TC_M*WARPGROUP_SIZE + warpgroup*TC_M + (warp//16)*4+i, gy, cp_j*TC_N + warp%16]
-  store = UOp.group(*[c_load(i).store(acc[cp_j, cp_i].index(i)) for i in range(4)])
+  store = UOp.group(*[c_load(i).store(acc[cp_j, cp_i].gep(i)) for i in range(4)])
   store = store.end(cp_i, cp_j)
 
   return store.sink(arg=KernelInfo(name="custom_gemm", opts_to_apply=())).simplify()
@@ -191,11 +192,12 @@ acc = UOp.placeholder((4,), dtypes.float, 0, AddrSpace.REG)
 acc = acc[init_l:=UOp.range(4, 1)].set(0.0, end=init_l)
 
 # do the wmma
-acc_load = UOp.stack(*[acc.after(K_loop)[i] for i in range(4)])
-out = UOp.wmma(A_in, B_in, acc_load, (16, 16, 32), 'AMD', 64)
+acc_load = UOp.vectorize(*[acc.after(K_loop)[i] for i in range(4)])
+wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
+out = UOp(Ops.WMMA, dtypes.float.vec(4), (A_in, B_in, acc_load), arg=wmma_arg)
 
 # store back the acc
-acc = acc.after(UOp.group(*[acc[i].store(out.index(i)) for i in range(4)]).end(K_loop))
+acc = acc.after(UOp.group(*[acc[i].store(out.gep(i)) for i in range(4)]).end(K_loop))
 
 # store the acc into gmem
 store = UOp.group(*[C[gx, (warp//16)*4+i, gy, warp%16].store(acc[i]) for i in range(4)])
@@ -216,7 +218,7 @@ if __name__ == "__main__":
   ref.realize()
 
   GlobalCounters.reset()
-  with Context(DEBUG=max(2, DEBUG.value)):
+  with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2):
     tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
     tst.realize()
   print(f"{(N*M*K*2 / GlobalCounters.time_sum_s)*1e-12:.2f} REAL TFLOPS")

extra/gemm/mi350x_uop_matmul_2.py

@@ -6,7 +6,7 @@ os.environ["AMD_LLVM"] = "0"
 from tinygrad import Tensor, Context, dtypes, UOp, GlobalCounters
 from tinygrad.helpers import DEBUG, getenv
 from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import AxisType, KernelInfo
+from tinygrad.uop.ops import sint, AxisType, KernelInfo, Ops
 
 WARP_SIZE = 64
 
@@ -29,7 +29,7 @@ TID_SIZE = WARPGROUP_SIZE*WARP_SIZE
 
 def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=()):
   assert dest.shape == src.shape
-  rngs = [UOp.range(s, rng+i, AxisType.UPCAST if i in upcast else AxisType.WEAK) for i,s in enumerate(src.shape)]
+  rngs = [UOp.range(s, rng+i, AxisType.UPCAST if i in upcast else AxisType.LOOP) for i,s in enumerate(src.shape)]
   copy = dest[*rngs].store(src[*rngs]).end(*rngs)
   return dest.after(copy) if set else copy
 
@@ -37,8 +37,8 @@ def compute_on_locals(acc:UOp, Asl:UOp, Bsl:UOp, rng:int, afters:tuple[UOp, ...]
   K_inner_loop = UOp.range(BLOCK_K//TC_K, rng, AxisType.REDUCE)
 
   # load from locals into registers
-  Ar = UOp.placeholder((BLOCK_M//TC_M//WARPGROUP_SIZE,), dtypes.half, slot=1, addrspace=AddrSpace.REG)
-  Br = UOp.placeholder((BLOCK_N//TC_N,), dtypes.half, slot=2, addrspace=AddrSpace.REG)
+  Ar = UOp.placeholder((BLOCK_M//TC_M//WARPGROUP_SIZE,), dtypes.half.vec(8), slot=1, addrspace=AddrSpace.REG)
+  Br = UOp.placeholder((BLOCK_N//TC_N,), dtypes.half.vec(8), slot=2, addrspace=AddrSpace.REG)
 
   M_load_loop = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, rng+10)
   Asl = Asl.reshape(BLOCK_K//TC_K, TC_K, BLOCK_M//TC_M//WARPGROUP_SIZE, WARPGROUP_SIZE, TC_M)
@@ -60,7 +60,8 @@ def compute_on_locals(acc:UOp, Asl:UOp, Bsl:UOp, rng:int, afters:tuple[UOp, ...]
   acc_load = acc_after[N_inner_loop, M_inner_loop]
 
   # do WMMA
-  out = UOp.wmma(Ar[M_inner_loop], Br[N_inner_loop], acc_load, (16, 16, 32), 'AMD', 64)
+  wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
+  out = UOp(Ops.WMMA, dtypes.float.vec(4), (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg)
 
   # store back the acc
   acc_store = acc[N_inner_loop, M_inner_loop].store(out)
@@ -71,7 +72,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
   K_outer_loop = UOp.range(K//BLOCK_K, 0, AxisType.REDUCE)
 
   # split out the globals into blocks
-  C = C.src[0].cast(dtypes.float).reshape((M//BLOCK_M, BLOCK_M, N//BLOCK_N, BLOCK_N))
+  C = C.src[0].cast(dtypes.float.vec(4).ptr(C.ptrdtype.size)).reshape((M//BLOCK_M, BLOCK_M, N//BLOCK_N, BLOCK_N))
   A = A.reshape((M//BLOCK_M, BLOCK_M, K//BLOCK_K, BLOCK_K))[gx, :, K_outer_loop, :]
   B = B.reshape((K//BLOCK_K, BLOCK_K, N//BLOCK_N, BLOCK_N))[K_outer_loop, :, gy, :]
 
@@ -106,7 +107,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
   if getenv("COMPUTE"):
     As, Bs = As.after(barrier), Bs.after(barrier)
 
-    acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float, 0, AddrSpace.REG)
+    acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float.vec(4), 0, AddrSpace.REG)
 
     sink = compute_on_locals(acc, As, Bs, 200, afters=(barrier,), warpgroup=warpgroup, warp=warp)
     sink = sink.end(K_outer_loop)
@@ -126,7 +127,7 @@ if __name__ == "__main__":
 
 
   GlobalCounters.reset()
-  with Context(DEBUG=max(2, DEBUG.value)):
+  with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2):
     tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
     tst.realize()
   print(f"{(N*M*K*2 / GlobalCounters.time_sum_s)*1e-12:.2f} REAL TFLOPS")

extra/gemm/moe_gemm.py

@@ -1,111 +0,0 @@
-import functools, pathlib
-from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-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
-
-@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
-  E, N, K2 = B.shape
-  assert K == K2, f"{A.shape} {B.shape}"
-  assert E == n_experts, f"{E} != {n_experts}"
-  threads = UOp.special(64 * 8, "lidx0")
-  workgroups = UOp.special((M // 256) * (N // 256), "gidx0")
-  sink_inputs = (C.base, A.base, B.base, scale_A.base, scale_B.base, extra[0].base, extra[1].base, extra[2].base, threads, workgroups)
-  sink = UOp.sink(*sink_inputs,
-                  arg=KernelInfo(f"hk_grouped_mxfp8_gemm_{E}_{M}_{N}_{K}",
-                                 estimates=Estimates(ops=2*M*N*K, mem=(M*K+E*N*K)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
-  kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
-  src = (kittens_path/"grouped_mxfp8_gemm.cpp").read_text()
-  lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
-                                 "-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}",
-                                 f"-DGEMM_E={E}"]).compile_cached(src)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
-                               UOp(Ops.BINARY, arg=lib)))
-
-@functools.cache
-def custom_hk_grouped_mxfp8_wgrad(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, expert_off:UOp, *, dname:str, n_experts:int) -> UOp:
-  N, M = A.shape
-  K, M2 = B.shape
-  assert M == M2, f"{A.shape} {B.shape}"
-  E = n_experts
-  threads = UOp.special(64 * 8, "lidx0")
-  workgroups = UOp.special(E * (N // 256) * (K // 256), "gidx0")
-  sink = UOp.sink(C.base, A.base, B.base, scale_A.base, scale_B.base, expert_off.base, threads, workgroups,
-                  arg=KernelInfo(f"hk_grouped_mxfp8_wgrad_{E}_{M}_{N}_{K}",
-                                 estimates=Estimates(ops=2*M*N*K, mem=(N*M+K*M)*A.dtype.itemsize+E*N*K*C.dtype.itemsize)))
-  kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
-  src = (kittens_path/"grouped_mxfp8_wgrad.cpp").read_text()
-  lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
-                                 "-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DWGRAD_M={M}", f"-DWGRAD_N={N}", f"-DWGRAD_K={K}",
-                                 f"-DWGRAD_E={E}"]).compile_cached(src)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
-                               UOp(Ops.BINARY, arg=lib)))
-
-def grouped_mx_wgrad(g:Tensor, xg:Tensor, expert_off:Tensor, n_experts:int) -> Tensor:
-  from extra.llama_kernels.transpose_quantize_mxfp8 import transpose_quantize_mxfp8
-  M, N = g.shape
-  M2, K = xg.shape
-  assert M == M2, f"{g.shape} {xg.shape}"
-  assert M % 128 == 0 and N % 256 == 0 and K % 256 == 0, f"wgrad needs M%128,N%256,K%256, got {g.shape} {xg.shape}"
-  gT, _, g_si = transpose_quantize_mxfp8(g.contiguous())
-  xT, _, x_si = transpose_quantize_mxfp8(xg.contiguous())
-  dname = (g.device[0] if isinstance(g.device, tuple) else g.device).split(":")[0]
-  is_multi = isinstance(g.device, tuple)
-  inv = Tensor.invalids(1, n_experts * N, K, dtype=dtypes.bfloat16, device=g.device)
-  out = Tensor(inv.uop.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)
-  return out.reshape(n_experts, N, K)
-
-def mx_pack_3d(e8:Tensor) -> Tensor:
-  E, rows, scale_K = e8.shape
-  return e8.reshape(E, rows, scale_K // 4, 4).bitcast(dtypes.uint32).reshape(E, rows, scale_K // 4).permute(0, 2, 1).contiguous()
-
-@functools.cache
-def custom_grouped_mx_gemm_bw(gradient:UOp, kernel:UOp, w_stored:bool=False) -> tuple:
-  inputs = kernel.src[1:]
-  aq = Tensor(inputs[1], device=inputs[1].device)
-  bq = Tensor(inputs[2], device=inputs[2].device)
-  ae8 = Tensor(inputs[5], device=inputs[5].device)
-  be8 = Tensor(inputs[6], device=inputs[6].device)
-  E, N = bq.shape[0], bq.shape[1]
-  M, K = aq.shape
-  g = Tensor(gradient, device=aq.device).reshape(M, N).cast(dtypes.bfloat16)
-  x_phys = (aq.cast(dtypes.bfloat16) * _mx_block_scale(ae8).cast(dtypes.bfloat16))
-  w_phys = (bq.cast(dtypes.bfloat16) * _mx_block_scale_3d(be8).cast(dtypes.bfloat16))
-  expert_off = Tensor(inputs[7], device=inputs[7].device)
-  grad_x = grouped_mx_gemm(g, w_phys.transpose(1, 2), expert_off)
-  grad_w = grouped_mx_wgrad(g, x_phys, expert_off, E)
-  grad_xq = grad_x * _mx_block_scale(ae8).cast(dtypes.bfloat16)
-  grad_wq = grad_w.contiguous() if w_stored else (grad_w * _mx_block_scale_3d(be8).cast(dtypes.bfloat16)).contiguous()
-  return (None, grad_xq.uop, grad_wq.uop) + tuple(None for _ in inputs[3:])
-
-_grouped_bw_stored = functools.partial(custom_grouped_mx_gemm_bw, w_stored=True)
-
-def grouped_mx_gemm(x:Tensor, w:Tensor|tuple[Tensor, Tensor], expert_off:Tensor) -> Tensor:
-  if (pre_quantized := isinstance(w, tuple)):
-    w_q, w_e8 = w
-    E, N, K2 = w_q.shape
-  else:
-    E, N, K2 = w.shape
-  M, K = x.shape
-  assert K == K2, f"shape mismatch {x.shape} {w.shape}"
-  assert M % 256 == 0 and N % 256 == 0 and K % 128 == 0, f"grouped mxfp8 needs M%256,N%256,K%128, got {x.shape} {w.shape}"
-  dname = (x.device[0] if isinstance(x.device, tuple) else x.device).split(":")[0]
-  x_q, x_e8, x_si = quantize_mxfp8(x)
-  if not pre_quantized: w_q, w_e8, _ = quantize_mxfp8(w)
-  w_si = mx_pack_3d(w_e8)
-  xe_in, out_shape = x_e8.reshape(M, K // 32), (M, N)
-  if isinstance(x.device, tuple) and (row_axis := x.uop.axis) is not None:
-    ndev = len(x.device)
-    out = Tensor(Tensor.invalids(*(s // ndev if i == row_axis else s for i, s in enumerate(out_shape)),
-                                 dtype=dtypes.bfloat16, device=x.device).uop.unshard(row_axis), device=x.device)
-  else:
-    out = Tensor.invalids(*out_shape, dtype=dtypes.bfloat16, device=x.device)
-  return Tensor.custom_kernel(out, x_q, w_q, x_si, w_si, xe_in, w_e8, expert_off,
-                              fxn=functools.partial(custom_hk_grouped_mxfp8_gemm, dname=dname, n_experts=E),
-                              grad_fxn=(_grouped_bw_stored if pre_quantized else custom_grouped_mx_gemm_bw))[0]

extra/gemm/moe_routing.py

@@ -1,131 +0,0 @@
-from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
-
-BLOCK_ROW = 256
-
-def _sharded_invalids(shape:tuple[int, ...], dtype, device) -> Tensor:
-  if isinstance(device, tuple):
-    per = Tensor.invalids(shape[0]//len(device), *shape[1:], dtype=dtype, device=device)
-    return Tensor(per.uop.unshard(0), device=device)
-  return Tensor.invalids(*shape, dtype=dtype, device=device)
-
-def _atomic_add(device:str) -> str:
-  return "__hip_atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);" if device == "AMD" \
-    else "__atomic_fetch_add({0}, {1}, __ATOMIC_RELAXED);"
-
-def _blk_for(D:int) -> int:
-  blk = 64
-  while D % blk: blk //= 2
-  return blk
-
-def _kv_ranges(G, N, D, BLK):
-  g = UOp.range(G, 0)
-  m = UOp.range(N, 1)
-  jo = UOp.range(D // BLK, 2)
-  ji = UOp.range(BLK, 3, AxisType.LOCAL)
-  return g, m, jo * BLK + ji, jo, ji
-
-def _ggather_fwd_kernel(out:UOp, table:UOp, idx:UOp) -> UOp:
-  G, M, D = out.shape
-  g, m, j, jo, ji = _kv_ranges(G, M, D, _blk_for(D))
-  row = idx.index(g, m).cast(dtypes.weakint)
-  val = table.index(g, row, j).load()
-  return out.index(g, m, j).store(val).end(g, m, jo, ji).sink(
-    arg=KernelInfo(name=f"ggather_fwd_{M}_{D}", opts_to_apply=()))
-
-def _ggather_zero_kernel(out:UOp) -> UOp:
-  i = UOp.range(out.numel(), 0)
-  return out.flatten().index(i).store(UOp.const(0.0, out.dtype)).end(i).sink(arg=KernelInfo(name="ggather_zero"))
-
-def _sharded_zeros(shape:tuple[int, ...], dtype, device) -> Tensor:
-  return Tensor.custom_kernel(_sharded_invalids(shape, dtype, device), fxn=_ggather_zero_kernel)[0]
-
-def _ggather_bwd(gradient:UOp, kernel:UOp) -> tuple:
-  _, table_u, idx_u = kernel.src[1:4]
-  dev = table_u.device
-  device = (dev[0] if isinstance(dev, tuple) else dev).split(":")[0]
-  G, R, D = table_u.shape
-  gt = _sharded_zeros((G, R, D), dtypes.float32, dev)
-  go = Tensor(gradient, device=dev)
-  atomic_str = _atomic_add(device)
-  def _bwd_kernel(gtab:UOp, gout:UOp, idx:UOp) -> UOp:
-    Gk, M, Dk = gout.shape
-    g, m, j, jo, ji = _kv_ranges(Gk, M, Dk, _blk_for(Dk))
-    row = idx.index(g, m).cast(dtypes.weakint)
-    val = gout.index(g, m, j).load().cast(dtypes.float32)
-    atomic = UOp(Ops.CUSTOM, dtypes.void, (gtab.index(g, row, j), val), arg=atomic_str)
-    return atomic.end(g, m, jo, ji).sink(arg=KernelInfo(name=f"ggather_bwd_{M}_{Dk}", opts_to_apply=()))
-  grad_table = Tensor.custom_kernel(gt, go, Tensor(idx_u, device=dev), fxn=_bwd_kernel)[0]
-  return (None, grad_table.cast(table_u.dtype).uop, None)
-
-def grouped_gather_rows(table:Tensor, idx:Tensor, n_groups:int) -> Tensor:
-  G, R, D = table.shape
-  M = idx.shape[1]
-  out = _sharded_invalids((G, M, D), table.dtype, table.device)
-  return Tensor.custom_kernel(out, table, idx, fxn=_ggather_fwd_kernel, grad_fxn=_ggather_bwd)[0]
-
-def _gscatter_fwd_kernel(out:UOp, src:UOp, idx:UOp) -> UOp:
-  G, M, D = out.shape
-  k = idx.shape[1] // src.shape[1]
-  g, m, j, jo, ji = _kv_ranges(G, idx.shape[1], D, _blk_for(D))
-  row = idx.index(g, m).cast(dtypes.weakint)
-  val = src.index(g, (m // k).cast(dtypes.weakint), j).load()
-  return out.index(g, row, j).store(val).end(g, m, jo, ji).sink(
-    arg=KernelInfo(name=f"gscatter_fwd_{idx.shape[1]}_{D}", opts_to_apply=()))
-
-def _gscatter_bwd(gradient:UOp, kernel:UOp) -> tuple:
-  _, src_u, idx_u = kernel.src[1:4]
-  dev = src_u.device
-  G, T_l, D = src_u.shape
-  k = idx_u.shape[1] // T_l
-  sel = grouped_gather_rows(Tensor(gradient, device=dev), Tensor(idx_u, device=dev), G)
-  return (None, sel.reshape(G, T_l, k, D).sum(2).cast(src_u.dtype).uop, None)
-
-def grouped_scatter_rows(src:Tensor, idx:Tensor, m_l:int) -> Tensor:
-  G, T_l, D = src.shape
-  zero = _sharded_zeros((G, m_l, D), src.dtype, src.device)
-  return Tensor.custom_kernel(zero, src, idx, fxn=_gscatter_fwd_kernel, grad_fxn=_gscatter_bwd)[0]
-
-def m_max_for(t_local:int, experts_per_tok:int, n_experts:int) -> int:
-  return (-(-t_local * experts_per_tok // BLOCK_ROW) + n_experts) * BLOCK_ROW
-
-class Routing:
-  def __init__(self, weights:Tensor, dest_row:Tensor, off:Tensor, m_l:int, n_groups:int, t_local:int):
-    self.weights, self.dest_row = weights, dest_row
-    self.off = off
-    self.m_l, self.n_groups, self.t_local = m_l, n_groups, t_local
-
-  @property
-  def rows_e(self) -> Tensor:
-    G, E = self.off.shape[0], self.off.shape[1] - 1
-    tr = Tensor.arange(self.m_l // BLOCK_ROW, dtype=dtypes.int32).reshape(1, -1, 1) * BLOCK_ROW
-    tr = tr.shard(self.off.device) if isinstance(self.off.device, tuple) else tr.to(self.off.device)
-    tile_e = ((tr >= self.off[:, :E].reshape(G, 1, E)).sum(-1) - 1).cast(dtypes.int32)
-    return tile_e.reshape(-1, 1).expand(-1, BLOCK_ROW).reshape(-1)
-
-def n_groups_of(t:Tensor) -> int:
-  return len(t.device) if isinstance(t.device, tuple) else 1
-
-def route(logits:Tensor, experts_per_tok:int, n_experts:int) -> Routing:
-  T, E = logits.shape
-  k, G = experts_per_tok, n_groups_of(logits)
-  assert T % G == 0, f"tokens {T} must split across {G} devices"
-  T_l, m_l = T // G, m_max_for(T // G, k, n_experts)
-
-  topv, topi = logits.reshape(G, T_l, E).topk(k)
-  weights = topv.softmax(-1)
-  m = topi.reshape(G, T_l * k).cast(dtypes.int32).one_hot(E).cast(dtypes.int32)
-
-  pad = ((m.sum(1) + (BLOCK_ROW - 1)) // BLOCK_ROW) * BLOCK_ROW
-  off = pad.cumsum(1).pad(((0, 0), (1, 0)))
-  dest_row = ((m.cumsum(1) + off[:, :E].reshape(G, 1, E)) * m).sum(-1).sub(1).cast(dtypes.int32)
-  return Routing(weights, dest_row, off, m_l, G, T_l)
-
-def dispatch(x:Tensor, r:Routing) -> Tensor:
-  G, D = r.n_groups, x.shape[-1]
-  return grouped_scatter_rows(x.reshape(G, r.t_local, D), r.dest_row, r.m_l).reshape(G * r.m_l, D)
-
-def combine(y:Tensor, r:Routing, n_tokens:int, experts_per_tok:int) -> Tensor:
-  G, D, k = r.n_groups, y.shape[-1], experts_per_tok
-  sel = grouped_gather_rows(y.reshape(G, r.m_l, D), r.dest_row, G).reshape(G, r.t_local, k, D)
-  return (sel * r.weights.reshape(G, r.t_local, k, 1).cast(sel.dtype)).sum(2).reshape(n_tokens, D).cast(y.dtype)

extra/gemm/rdna4_asm_matmul.py

@@ -219,11 +219,10 @@ def test_matmul():
   def asm_kernel(A, B, C):
     gidxs = [UOp.special(n, f"gidx{i}") for i,n in enumerate(grid)]
     lidxs = [UOp.special(THREADS, "lidx0")]
-    lds_size = max(LDS_SIZE, 65536//getenv("LIMIT_OCC",2))
-    lds = UOp.placeholder((lds_size,), dtypes.uint8, 0, AddrSpace.LOCAL)
+    lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=max(LDS_SIZE, 65536//getenv("LIMIT_OCC",2)), addrspace=AddrSpace.LOCAL), (), 'lds')
     sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs,
                     arg=KernelInfo(name=colored("kernel","cyan"), estimates=Estimates(ops=N*N*N*2, mem=N*N*2*3)))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
 
   c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
   linear = c.schedule_linear()

extra/gemm/triton_nv_matmul.py

@@ -93,7 +93,7 @@ if __name__ == "__main__":
   info = ProgramInfo(name="matmul_kernel",
                      global_size=(M//BLOCK_SIZE_M, N//BLOCK_SIZE_N, 1), local_size=(32*compiled.metadata.num_warps, 1, 1))
   sink = UOp.sink(arg=KernelInfo(name="matmul_kernel"))
-  prg_uop = to_program(UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=src)), arg=info),
+  prg_uop = to_program(UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=src)), arg=info),
                        Device.default.renderer)
   rt = get_runtime(Device.DEFAULT, prg_uop)
   all_bufs = [x.ensure_allocated() for x in bufs]

extra/gemm_fragment.py

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

extra/hcq2/ops_amd2.py

@@ -1,709 +0,0 @@
-from __future__ import annotations
-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 make_binary_patch
-from tinygrad.uop.ops import sint, UOp
-from tinygrad.device import Compiled, BufferSpec, Buffer, Device
-from tinygrad.dtype import dtypes
-from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, colored, prod, ContextVar, TracingKey
-from tinygrad.helpers import VIZ, ceildiv, unwrap, pluralize, to_tuple
-from tinygrad.renderer.cstyle import HIPRenderer, HIPCCRenderer
-from tinygrad.renderer.llvmir import AMDLLVMRenderer
-from tinygrad.runtime.autogen import kfd, hsa, sqtt, amdgpu_kd, amdgpu_drm
-from tinygrad.runtime.autogen.am import am
-from tinygrad.runtime.support.elf import elf_loader
-from tinygrad.runtime.support.hcq import FileIOInterface, HCQBuffer, MMIOInterface, hcq_filter_visible_devices
-from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
-from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_pmc
-from tinygrad.runtime.support.system import PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
-from tinygrad.runtime.support.usb import USB3
-from tinygrad.runtime.support.memory import AddrSpace, BumpAllocator
-from tinygrad.runtime.ops_amd import SQTT, SQTT_ITRACE_SE_MASK, SQTT_LIMIT_SE, SQTT_SIMD_SEL, SQTT_TOKEN_EXCLUDE, PMC
-from tinygrad.runtime.ops_amd import EVENT_INDEX_PARTIAL_FLUSH, WAIT_REG_MEM_FUNCTION_EQ, WAIT_REG_MEM_FUNCTION_NEQ, WAIT_REG_MEM_FUNCTION_GEQ
-if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import
-
-from tinygrad.engine.realize import get_runtime, pm_flatten_linear
-from tinygrad.uop import FastEnum, auto
-from tinygrad.uop.ops import Ops, UPat, PatternMatcher, graph_rewrite
-
-# *****************
-# PM4
-
-class PM4Ops(FastEnum):
-  SET_SH_REG = auto(); SET_UCONFIG_REG = auto(); WAIT_REG_MEM = auto(); ACQUIRE_MEM = auto()  # noqa: E702
-  RELEASE_MEM = auto(); DISPATCH_DIRECT = auto(); EVENT_WRITE = auto()  # noqa: E702
-
-def pkt3(ctx, op:PM4Ops, *vals):
-  return UOp(Ops.INS, arg=op, src=tuple(UOp.const(x, dtypes.uint32)
-    for x in (ctx.pm4.PACKET3(getattr(ctx.pm4, f"PACKET3_{op.name}"), len(vals) - 1), *vals)))
-
-def wreg(ctx, reg:AMDReg, *args:sint, **kwargs:int):
-  if bool(args) == bool(kwargs): raise RuntimeError('One (and only one) of *args or **kwargs must be specified')
-  if ctx.pm4.PACKET3_SET_SH_REG_START <= reg.addr[0] < ctx.pm4.PACKET3_SET_SH_REG_END:
-    op, set_packet_start = PM4Ops.SET_SH_REG, ctx.pm4.PACKET3_SET_SH_REG_START
-  elif ctx.pm4.PACKET3_SET_UCONFIG_REG_START <= reg.addr[0] < ctx.pm4.PACKET3_SET_UCONFIG_REG_START + 2**16-1:
-    op, set_packet_start = PM4Ops.SET_UCONFIG_REG, ctx.pm4.PACKET3_SET_UCONFIG_REG_START
-  else: raise RuntimeError(f'Cannot set {reg.name} ({reg.addr[0]}) via pm4 packet')
-  return pkt3(ctx, op, reg.addr[0] - set_packet_start, *(args or (reg.encode(**kwargs),)))
-
-def wait_reg_mem(ctx, value, mask=0xffffffff, mem=None, reg=None, reg_done=0, op=WAIT_REG_MEM_FUNCTION_GEQ):
-  wrm_info_dw = ctx.pm4.WAIT_REG_MEM_MEM_SPACE(int(mem is not None)) | ctx.pm4.WAIT_REG_MEM_OPERATION(int(mem is None and reg_done > 0)) \
-              | ctx.pm4.WAIT_REG_MEM_FUNCTION(op) | ctx.pm4.WAIT_REG_MEM_ENGINE(0)
-  return pkt3(ctx, PM4Ops.WAIT_REG_MEM, wrm_info_dw, *(data64_le(mem) if mem is not None else (reg, reg_done)), value, mask, 4)
-
-def acquire_mem(ctx, addr=0x0, sz=(1 << 64)-1, gli=1, glm=1, glk=1, glv=1, gl1=1, gl2=1):
-  if ctx.target[0] != 9:
-    cache_flags_dw = ctx.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLI_INV(gli) \
-                   | ctx.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLM_INV(glm) | ctx.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLM_WB(glm) \
-                   | ctx.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLK_INV(glk) | ctx.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLK_WB(glk) \
-                   | ctx.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLV_INV(glv) | ctx.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL1_INV(gl1) \
-                   | ctx.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL2_INV(gl2) | ctx.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL2_WB(gl2)
-    return pkt3(ctx, PM4Ops.ACQUIRE_MEM, 0, *data64_le(sz), *data64_le(addr), 0, cache_flags_dw)
-  cp_coher_cntl = ctx.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_SH_ICACHE_ACTION_ENA(gli) | \
-                  ctx.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_SH_KCACHE_ACTION_ENA(glk) | \
-                  ctx.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TC_ACTION_ENA(gl2) | \
-                  ctx.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TCL1_ACTION_ENA(gl1) | \
-                  ctx.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TC_WB_ACTION_ENA(gl2)
-  return pkt3(ctx, PM4Ops.ACQUIRE_MEM, cp_coher_cntl, *data64_le(sz), *data64_le(addr), 0x0000000A)
-
-def release_mem(ctx, address=0x0, value=0, data_sel=0, int_sel=2, ctxid=0, cache_flush=False):
-  if ctx.target[0] != 9:
-    cache_flags_dw = 0 if not cache_flush else (ctx.pm4.PACKET3_RELEASE_MEM_GCR_GLV_INV | ctx.pm4.PACKET3_RELEASE_MEM_GCR_GL1_INV \
-                   | ctx.pm4.PACKET3_RELEASE_MEM_GCR_GL2_INV | ctx.pm4.PACKET3_RELEASE_MEM_GCR_GLM_WB \
-                   | ctx.pm4.PACKET3_RELEASE_MEM_GCR_GLM_INV | ctx.pm4.PACKET3_RELEASE_MEM_GCR_GL2_WB | ctx.pm4.PACKET3_RELEASE_MEM_GCR_SEQ)
-    event_dw = ctx.pm4.PACKET3_RELEASE_MEM_EVENT_TYPE(ctx.pm4.CACHE_FLUSH_AND_INV_TS_EVENT) \
-             | ctx.pm4.PACKET3_RELEASE_MEM_EVENT_INDEX(ctx.pm4.event_index__mec_release_mem__end_of_pipe)
-    memsel_dw = ctx.pm4.PACKET3_RELEASE_MEM_DATA_SEL(data_sel) | ctx.pm4.PACKET3_RELEASE_MEM_INT_SEL(int_sel) \
-              | ctx.pm4.PACKET3_RELEASE_MEM_DST_SEL(0)
-  else:
-    cache_flags_dw = 0 if not cache_flush else (ctx.pm4.EOP_TC_WB_ACTION_EN | ctx.pm4.EOP_TC_NC_ACTION_EN)
-    event_dw = ctx.pm4.EVENT_TYPE(ctx.pm4.CACHE_FLUSH_AND_INV_TS_EVENT) | ctx.pm4.EVENT_INDEX(ctx.pm4.event_index__mec_release_mem__end_of_pipe)
-    memsel_dw = ctx.pm4.DATA_SEL(data_sel) | ctx.pm4.INT_SEL(int_sel)
-    ctxid = 0
-  return pkt3(ctx, PM4Ops.RELEASE_MEM, event_dw | cache_flags_dw, memsel_dw, *data64_le(address), *data64_le(value), ctxid)
-
-def memory_barrier(ctx):
-  pf = '' if ctx.nbio.version[0] == 2 else '0' if ctx.nbio.version[:2] != (7, 11) else '1'
-  return UOp(Ops.LINEAR, dtypes.void, (
-    wait_reg_mem(ctx, reg=getattr(ctx.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_REQ').addr[0],
-                 reg_done=getattr(ctx.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff),
-    acquire_mem(ctx)))
-
-def pm4_wait(ctx, dst, val): return wait_reg_mem(ctx, val, mem=dst.getaddr(ctx.devs))
-
-def pm4_barrier(ctx): return memory_barrier(ctx)
-
-def pm4_store(ctx, dst, val):
-  if val.op is Ops.BINARY: return None
-  return release_mem(ctx, dst.getaddr(ctx.devs), val, ctx.pm4.data_sel__mec_release_mem__send_32_bit_low,
-                     ctx.pm4.int_sel__mec_release_mem__send_interrupt_after_write_confirm, cache_flush=True)
-
-def pm4_timestamp(ctx, dst):
-  return release_mem(ctx, dst.getaddr(ctx.devs), 0, ctx.pm4.data_sel__mec_release_mem__send_gpu_clock_counter,
-                     ctx.pm4.int_sel__mec_release_mem__none)
-
-def pm4_program(ctx, call, prg):
-  data, info = prg.arg
-  lib_gpu = prg.src[0]
-  args = encode_kernargs_clike(call, prg, ctx.devs)
-  prog_addr = lib_gpu.getaddr(ctx.devs) + data.entry_point_offset
-  scratch_addr = UOp.placeholder((data.private_segment_size,), dtypes.uint8, 0, device=ctx.devs).rtag("scratch").getaddr(ctx.devs)
-  args_addr = args.getaddr(ctx.devs)
-
-  user_regs = []
-  if data.enable_private_segment_sgpr:
-    scratch_hilo = data64_le(scratch_addr)
-    user_regs = [scratch_hilo[0], scratch_hilo[1] | 1 << 31, 0xffffffff, 0x20c14000]
-  if data.enable_dispatch_ptr: user_regs += [*data64_le(args_addr + data.kernargs_segment_size)]
-  user_regs += [*data64_le(args_addr)]
-
-  dispatch_init = ctx.gc.regCOMPUTE_DISPATCH_INITIATOR.encode(
-    **({'cs_w32_en': int(data.wave32)} if ctx.target[0] != 9 else {}), force_start_at_000=1, compute_shader_en=1)
-  ins = [acquire_mem(ctx, gli=0, gl2=0),
-    wreg(ctx, ctx.gc.regCOMPUTE_PGM_LO, *data64_le(prog_addr >> 8)),
-    wreg(ctx, ctx.gc.regCOMPUTE_PGM_RSRC1, data.rsrc1, data.rsrc2),
-    wreg(ctx, ctx.gc.regCOMPUTE_PGM_RSRC3, data.rsrc3),
-    wreg(ctx, ctx.gc.regCOMPUTE_TMPRING_SIZE, ctx.tmpring_size(data.private_segment_size))]
-  ins += [wreg(ctx, ctx.gc.regCOMPUTE_DISPATCH_SCRATCH_BASE_LO, *data64_le((scratch_addr + data.private_segment_size // ctx.xccs * xcc_id) >> 8))
-          for xcc_id in range(ctx.xccs)]
-  ins += [wreg(ctx, ctx.gc.regCOMPUTE_RESTART_X, 0, 0, 0),
-    wreg(ctx, ctx.gc.regCOMPUTE_USER_DATA_0, *user_regs),
-    wreg(ctx, ctx.gc.regCOMPUTE_RESOURCE_LIMITS, ctx.gc.regCOMPUTE_RESOURCE_LIMITS.encode(waves_per_sh=getenv("WAVES_PER_SH"))),
-    wreg(ctx, ctx.gc.regCOMPUTE_START_X, 0, 0, 0, *(info.local_size or (1, 1, 1)), 0, 0),
-    pkt3(ctx, PM4Ops.DISPATCH_DIRECT, *info.global_size, dispatch_init),
-    pkt3(ctx, PM4Ops.EVENT_WRITE, ctx.pm4.EVENT_TYPE(ctx.soc.CS_PARTIAL_FLUSH) | ctx.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH))]
-  return UOp(Ops.LINEAR, dtypes.void, tuple(ins))
-
-pm_pm4_opsel = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), pm4_program),
-
-  (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val"))), pm4_wait),
-  (UPat(Ops.INS, arg="barrier"), pm4_barrier),
-  (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),)), pm4_timestamp),
-  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
-])
-
-def pm4_submit(ctx, lin):
-  # ensure compute queues are allocated
-  for d in (devs:=ctx.devs): q = Device[d].compute_queue
-  ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COMPUTE:0_{name}")
-    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
-
-  # the host fence at the start of the batch guarantees the ib is free to reuse
-  size_dw = sum(len(ins.src) for ins in lin.src)
-  assert size_dw < (1 << 20), f"indirect buffer of {size_dw} dwords doesn't fit one packet"
-
-  ib = UOp.placeholder((size_dw,), dtypes.uint32, next(UOp.unique_num), device=devs, volatile=True).rtag("cmdbuf")
-  cmdbuf = make_cmdbuf(lin, devs, buf=ib)
-
-  # the ring itself only carries a packet pointing at the ib, wrapping the ring
-  put = put_ptr.index(zero:=UOp.const(0, dtypes.int))
-  pkt = (ctx.pm4.PACKET3(ctx.pm4.PACKET3_INDIRECT_BUFFER, 2), *data64_le(cmdbuf.getaddr(devs)), size_dw | ctx.pm4.INDIRECT_BUFFER_VALID)
-  write_pkt = UOp.barrier(*[ring.index(((put + off) % q.ring.size).cast(dtypes.int)).store(UOp.const(x, dtypes.uint32)) for off,x in enumerate(pkt)])
-
-  # advance the put/write pointers past the packet
-  bump_put_ptr = put_ptr.index(zero).store(put + len(pkt))
-  bump_wptr = wptr.index(zero).store(put + len(pkt))
-  flush = UOp.barrier(write_pkt, bump_put_ptr, bump_wptr)
-  return doorbell.after(flush).index(zero).store(put + len(pkt))
-
-pm_pm4_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"), pm4_submit)])
-
-# *****************
-# SDMA
-
-class SDMAOps(FastEnum): COPY = auto(); POLL_REGMEM = auto(); FENCE = auto(); TRAP = auto(); TIMESTAMP = auto()  # noqa: E702
-
-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)))))
-
-def sdma_wait(ctx, ins, dst, val):
-  op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \
-     | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1)
-  return ins.ins(SDMAOps.POLL_REGMEM, src=tuple(UOp.const(x, dtypes.uint32) for x in (
-    op, *data64_le(dst.getaddr(ctx.devs)), val, 0xffffffff,
-    ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))))
-
-def sdma_store(ctx, ins, dst, val):
-  op = ctx.sdma.SDMA_OP_FENCE | (ctx.sdma.SDMA_PKT_FENCE_HEADER_MTYPE(3) if ctx.target[0] != 9 else 0)
-  return UOp(Ops.LINEAR, src=(
-    ins.ins(SDMAOps.FENCE, src=tuple(UOp.const(x, dtypes.uint32) for x in (op, *data64_le(dst.getaddr(ctx.devs)), val))),
-    ins.ins(SDMAOps.TRAP, src=tuple(UOp.const(x, dtypes.uint32) for x in (ctx.sdma.SDMA_OP_TRAP, 0)))))
-
-def sdma_timestamp(ctx, ins, dst):
-  op = ctx.sdma.SDMA_OP_TIMESTAMP | ctx.sdma.SDMA_PKT_TIMESTAMP_GET_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_TIMESTAMP_GET_GLOBAL)
-  return ins.ins(SDMAOps.TIMESTAMP, src=tuple(UOp.const(x, dtypes.uint32) for x in (op, *data64_le(dst.getaddr(ctx.devs)))))
-
-pm_sdma_opsel = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy),
-
-  (UPat(Ops.INS, arg="barrier"), lambda: UOp(Ops.NOOP, dtypes.void, ())),
-  (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val")), name="ins"), sdma_wait),
-  (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),), name="ins"), sdma_timestamp),
-  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val")), name="ins"), sdma_store),
-])
-
-def sdma_submit(cmdbuf, devs):
-  # the cmdbuf to submit + the patch writes that fill it
-  size_dw, zero = cmdbuf.nbytes() // dtypes.uint32.itemsize, UOp.const(0, dtypes.int)
-
-  # the sdma queue's ring and its host-side ring/write/put pointers
-  for d in devs: q = Device[d].sdma_queue(0)
-  ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COPY:0_{name}")
-    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
-
-  # sdma needs the cmdbuf contiguous: if it won't fit before the ring end, restart at 0 and zero the tail
-  put_b = put_ptr.index(zero)
-  tail_off_dw = ((put_b % (q.ring.size * 4)) // 4).cast(dtypes.int)
-  fits = (size_dw <= q.ring.size - tail_off_dw).cast(dtypes.int)
-  start_dw = fits * tail_off_dw
-  zero_amt_dw = (1 - fits) * (q.ring.size - tail_off_dw)
-
-  # zero the wrapped tail, then copy the cmdbuf into the ring
-  zi = UOp.range(zero_amt_dw, 0, dtype=dtypes.int, src=(cmdbuf,))
-  zero_tail = ring.index(tail_off_dw + zi).store(UOp.const(0, dtypes.uint32)).end(zi)
-  i = UOp.range(UOp.const(size_dw, dtypes.int), 0, dtype=dtypes.int, src=(cmdbuf,))
-  copy_to_ring = ring.index(start_dw + i).store(cmdbuf.index(i).load()).end(i)
-
-  # advance the put/write pointers past the zeroed tail and the cmdbuf
-  next_put_b = put_b + ((zero_amt_dw + size_dw) * 4).cast(put_b.dtype)
-  bump_put_ptr = put_ptr.index(zero).store(next_put_b)
-  bump_wptr = wptr.index(zero).store(next_put_b)
-
-  # ring the doorbell once the writes have landed
-  flush = UOp.barrier(zero_tail, copy_to_ring, bump_put_ptr, bump_wptr)
-  return doorbell.after(flush).index(zero).store(next_put_b)
-
-pm_sdma_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"),
-  lambda ctx, lin: sdma_submit(make_cmdbuf(lin, ctx.devs), ctx.devs))])
-
-@dataclass(frozen=True)
-class AMDEncodeCtx:  # encode-time constants for one queue: devs (every cmdbuf address resolves into these) + gfx version + packet/ip modules
-  devs: tuple[str, ...]; target: tuple[int, ...]; pm4: Any; sdma: Any; soc: Any  # noqa: E702
-  gc: AMDIP; nbio: AMDIP; xccs: int; max_copy_size: int; tmpring_size: Callable  # noqa: E702
-
-def encode_queue(q:UOp) -> UOp|None:
-  d = Device[(devs:=to_tuple(q.arg[0]))[0]]
-  ctx = AMDEncodeCtx(devs, d.target, d.pm4, d.sdma, d.soc, d.gc, d.nbio, d.xccs, d.max_copy_size, d.tmpring_size)
-  opsel, submit = (pm_pm4_opsel, pm_pm4_submit) if q.arg[1].startswith("COMPUTE") else (pm_sdma_opsel, pm_sdma_submit)
-  return submit.rewrite(graph_rewrite(q, opsel + pm_flatten_linear, walk=True, ctx=ctx, name=f"{q.arg[1]} opsel"), ctx)
-
-@dataclass(frozen=True)
-class AMDProgramData:
-  entry_point_offset:int; rsrc1:int; rsrc2:int; rsrc3:int; wave32:bool
-  private_segment_size:int; kernargs_segment_size:int; kernargs_alloc_size:int
-  enable_dispatch_ptr:int; enable_private_segment_sgpr:int
-
-_amd_program_cache:dict[tuple[bytes,str], tuple[AMDProgramData,bytes]] = {}
-def amd_build_program(prg:UOp) -> UOp:
-  dev = Device[to_tuple(prg.device)[0]] # TODO: rm this
-  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, dev.device))) is None:
-    image, sections, relocs = elf_loader(lib)
-    rodata = next(sh.header.sh_addr for sh in sections if sh.name == ".rodata")
-    for off, sym, typ, addent in relocs:
-      assert typ == 5, f"unknown AMD reloc {typ}"  # R_AMDGPU_REL64
-      image[off:off+8] = struct.pack('<q', sym - off + addent)
-    desc = amdgpu_kd.llvm_amdhsa_kernel_descriptor_t.from_buffer_copy(bytes(image[rodata:rodata+ctypes.sizeof(amdgpu_kd.llvm_amdhsa_kernel_descriptor_t)]))
-    if (lds:=((desc.group_segment_fixed_size+511)//512)&0x1FF) > (dev.iface.props['lds_size_in_kb']*1024)//512:
-      raise RuntimeError("Too many resources requested: group_segment_size")
-    edp = desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_DISPATCH_PTR
-
-    data = AMDProgramData(entry_point_offset=rodata + desc.kernel_code_entry_byte_offset,
-      rsrc1=desc.compute_pgm_rsrc1 | ((1<<20) if dev.target[0]==11 else 0),  # priv=1 on gfx11 for cwsr
-      rsrc2=desc.compute_pgm_rsrc2 | (lds<<15), rsrc3=desc.compute_pgm_rsrc3,
-      wave32=bool(desc.kernel_code_properties & 0x400), private_segment_size=desc.private_segment_fixed_size, kernargs_segment_size=desc.kernarg_size,
-      kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0), enable_dispatch_ptr=edp,
-      enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER)
-    buf = UOp.placeholder((len(image),), dtypes.uint8, next(UOp.unique_num), device=prg.device).rtag("program")
-    cached = _amd_program_cache[key] = prg.replace(src=(buf.after(make_binary_patch(buf, bytes(image))),), arg=(data, prg.arg))
-  return cached
-
-class AMDAllocator(HCQAllocator['AMDDevice']):
-  def __init__(self, dev:AMDDevice):
-    super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb())
-
-  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_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)
-
-@dataclass
-class AMDQueueDesc:
-  ring: Buffer; read_ptr: Buffer; write_ptr: Buffer; doorbell: Buffer; put_value: Buffer  # noqa: E702
-  eop_buffer: Buffer|None = None; cwsr_buffer: Buffer|None = None; params: tuple|None = None  # noqa: E702
-
-class KFDIface:
-  kfd:FileIOInterface|None = None
-  event_page:HCQBuffer|None = None
-  gpus:list[FileIOInterface] = []
-  count:int = 0
-
-  def _is_usable_gpu(self, gpu_id):
-    with contextlib.suppress(OSError): return int(gpu_id.read()) != 0
-    return False
-
-  def __init__(self, dev, device_id):
-    self.dev = dev
-
-    kfd_topo_path = "/sys/devices/virtual/kfd/kfd/topology/nodes"
-
-    # Initialize KFD interface during first run
-    if KFDIface.kfd is None:
-      KFDIface.kfd = FileIOInterface("/dev/kfd", os.O_RDWR)
-      gpus = [g for g in FileIOInterface(kfd_topo_path).listdir() if self._is_usable_gpu(FileIOInterface(f"{kfd_topo_path}/{g}/gpu_id"))]
-      KFDIface.gpus = hcq_filter_visible_devices(sorted(gpus, key=lambda x: int(x.split('/')[-1])), "AMD")
-      KFDIface.count = len(KFDIface.gpus)
-
-    if device_id >= len(KFDIface.gpus): raise RuntimeError(f"No device found for {device_id}. Requesting more devices than the system has?")
-
-    self.gpu_id = int(FileIOInterface(f"{kfd_topo_path}/{KFDIface.gpus[device_id]}/gpu_id").read())
-    self.props = {(p:=l.split())[0]: int(p[1]) for l in FileIOInterface(f"{kfd_topo_path}/{KFDIface.gpus[device_id]}/properties").read().splitlines()}
-    self.dev_sysfs_path = f"/sys/class/drm/renderD{self.props['drm_render_minor']}/device"
-    ip_base = f"{self.dev_sysfs_path}/ip_discovery/die/0"
-    id2ip = {am.GC_HWID: am.GC_HWIP, am.SDMA0_HWID: am.SDMA0_HWIP, am.NBIF_HWID: am.NBIF_HWIP}
-    ip_hw = [(id2ip[int(hwid)], int(hwid)) for hwid in FileIOInterface(ip_base).listdir() if hwid.isnumeric() and int(hwid) in id2ip]
-    self.ip_versions = {ip:tuple(int(FileIOInterface(f'{ip_base}/{hw}/0/{part}').read()) for part in ['major','minor','revision']) for ip,hw in ip_hw}
-    self.drm_fd = FileIOInterface(f"/dev/dri/renderD{self.props['drm_render_minor']}", os.O_RDWR)
-
-    self.kfd_ver = ((ver_st:=kfd.AMDKFD_IOC_GET_VERSION(KFDIface.kfd)).major_version, ver_st.minor_version)
-    kfd.AMDKFD_IOC_ACQUIRE_VM(KFDIface.kfd, drm_fd=self.drm_fd.fd, gpu_id=self.gpu_id)
-    if self.kfd_ver >= (1,14): kfd.AMDKFD_IOC_RUNTIME_ENABLE(KFDIface.kfd, mode_mask=0)
-
-    # Set these for our device.
-    if KFDIface.event_page is None:
-      KFDIface.event_page = self.alloc(0x8000, uncached=True)
-      kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_page_offset=KFDIface.event_page.meta.handle)
-    else: self.map(KFDIface.event_page)
-
-    # Event to wait for queues completion
-    self.dev.queue_event = kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_type=kfd.KFD_IOC_EVENT_SIGNAL, auto_reset=1)
-    self.dev.queue_event_mailbox_ptr = KFDIface.event_page.va_addr + self.dev.queue_event.event_slot_index * 8
-
-    # OS events to collect memory and hardware faults
-    self.mem_fault_event = kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_type=kfd.KFD_IOC_EVENT_MEMORY)
-    self.hw_fault_event = kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_type=kfd.KFD_IOC_EVENT_HW_EXCEPTION)
-
-    self.queue_event_arr = (kfd.struct_kfd_event_data * 3)(kfd.struct_kfd_event_data(event_id=self.dev.queue_event.event_id),
-      kfd.struct_kfd_event_data(event_id=self.mem_fault_event.event_id), kfd.struct_kfd_event_data(event_id=self.hw_fault_event.event_id))
-    self.queue_event_arr_ptr = ctypes.addressof(self.queue_event_arr)
-
-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, cpu_addr=None) -> HCQBuffer:
-    flags = kfd.KFD_IOC_ALLOC_MEM_FLAGS_WRITABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_EXECUTABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_NO_SUBSTITUTE
-
-    if uncached: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_COHERENT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED | kfd.KFD_IOC_ALLOC_MEM_FLAGS_GTT
-    else: flags |= (kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR if host else kfd.KFD_IOC_ALLOC_MEM_FLAGS_VRAM)
-
-    # Make mapped cpu address to be uncachable
-    if cpu_addr is not None: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_COHERENT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED
-
-    if cpu_access or host: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_PUBLIC
-
-    if flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR:
-      buf = addr = cpu_addr or FileIOInterface.anon_mmap(0, size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | mmap.MAP_ANONYMOUS, 0)
-    else: buf, addr = 0, FileIOInterface.anon_mmap(0, size, 0, mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS | MAP_NORESERVE, 0)
-
-    try: mem = kfd.AMDKFD_IOC_ALLOC_MEMORY_OF_GPU(self.kfd, va_addr=addr, size=size, gpu_id=self.gpu_id, flags=flags, mmap_offset=buf)
-    except OSError as e:
-      if e.errno == errno.EINVAL and (flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_VRAM) and cpu_access:
-        raise MemoryError("Cannot allocate host-visible VRAM. Ensure the resizable BAR option is enabled on your system.") from e
-      if e.errno == errno.ENOMEM: raise MemoryError(f"Cannot allocate {size} bytes: no memory is available.") from e
-      raise
-
-    if not (flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR):
-      buf = self.drm_fd.mmap(mem.va_addr, mem.size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | MAP_FIXED, mem.mmap_offset)
-      assert addr == buf == mem.va_addr
-
-    view = MMIOInterface(mem.va_addr, mem.size, fmt='B') if cpu_access or host else None
-    self.map(hcqbuf:=HCQBuffer(mem.va_addr, mem.size, meta=mem, view=view, owner=self.dev))
-    return hcqbuf
-
-  def free(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)
-
-    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)
-    assert stm.n_success == 1
-    return HCQBuffer(mem.va_addr, mem.size, meta=mem.meta, owner=mem.owner)
-
-  def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
-                   xcc_id=0, idx=0):
-    queue = kfd.AMDKFD_IOC_CREATE_QUEUE(KFDIface.kfd, ring_base_address=ring._buf.va_addr, ring_size=ring._buf.size, gpu_id=self.gpu_id,
-      queue_type=queue_type, queue_percentage=kfd.KFD_MAX_QUEUE_PERCENTAGE|(xcc_id<<8), queue_priority=getenv("AMD_KFD_QUEUE_PRIORITY", 7),
-      eop_buffer_address=eop_buffer._buf.va_addr if eop_buffer else 0, eop_buffer_size=eop_buffer._buf.size if eop_buffer else 0,
-      ctl_stack_size=ctl_stack_size, ctx_save_restore_address=cwsr_buffer._buf.va_addr if cwsr_buffer else 0, ctx_save_restore_size=ctx_save_restore_size,
-      write_pointer_address=gart._buf.va_addr+wptr, read_pointer_address=gart._buf.va_addr+rptr+8*xcc_id)
-
-    if not hasattr(self, 'doorbells'):
-      self.doorbells_base = queue.doorbell_offset & (~0x1fff) # doorbell is two pages
-      self.doorbells = cast(FileIOInterface, KFDIface.kfd).mmap(0, 0x2000, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED, self.doorbells_base)
-
-    (put_value := Buffer("CPU", 1, dtypes.uint64, preallocate=True))._buf.view.view(fmt='Q')[0] = 0
-    doorbell = Buffer("CPU", 1, dtypes.uint64,
-      options=BufferSpec(external_ptr=self.doorbells + queue.doorbell_offset - self.doorbells_base), preallocate=True)
-    return AMDQueueDesc(ring=ring, doorbell=doorbell, read_ptr=gart.view(1, dtypes.uint64, rptr+8*xcc_id).ensure_allocated(),
-      write_ptr=gart.view(1, dtypes.uint64, wptr).ensure_allocated(), put_value=put_value, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer)
-
-  def sleep(self, tm:int):
-    kfd.AMDKFD_IOC_WAIT_EVENTS(KFDIface.kfd, events_ptr=self.queue_event_arr_ptr, num_events=3, wait_for_all=0, timeout=tm)
-    if self.queue_event_arr[1].memory_exception_data.gpu_id or self.queue_event_arr[2].hw_exception_data.gpu_id: self.on_device_hang()
-
-  def on_device_hang(self):
-    def _str(st): return ' '.join(f'{k[0]}={getattr(st, k[0])}' for k in st._real_fields_)
-
-    # try to collect fault info if not already set from sleep().
-    if not self.queue_event_arr[1].memory_exception_data.gpu_id and not self.queue_event_arr[2].hw_exception_data.gpu_id:
-      with contextlib.suppress(RuntimeError): self.sleep(tm=1)
-
-    report = []
-    if self.queue_event_arr[1].memory_exception_data.gpu_id:
-      report += [f"MMU fault: 0x{self.queue_event_arr[1].memory_exception_data.va:X} | {_str(self.queue_event_arr[1].memory_exception_data.failure)}"]
-    if self.queue_event_arr[2].hw_exception_data.gpu_id: report += [f"HW fault: {_str(self.queue_event_arr[2].hw_exception_data)}"]
-
-    raise RuntimeError("\n".join(report))
-
-  def require_profile_mode(self, can_set_mode=True):
-    if self.dev.target[0] == 9: return
-    fn = f'{self.dev_sysfs_path}/power_dpm_force_performance_level'
-    if (perflevel:=FileIOInterface(fn).read().strip()) != 'profile_standard':
-      if can_set_mode:
-        atexit.register(lambda: os.system(f"echo '{perflevel}' | sudo tee {fn} > /dev/null"))
-        os.system(f"echo 'profile_standard' | sudo tee {fn} > /dev/null")
-        self.require_profile_mode(can_set_mode=False)
-      else:
-        raise RuntimeError("PMC/SQTT requires stable power state: run `amd-smi set -l stable_std` for KFD iface")
-
-  @functools.cached_property
-  def drm_dev_info(self) -> amdgpu_drm.struct_drm_amdgpu_info_device:
-    amdgpu_drm.DRM_IOCTL_AMDGPU_INFO(self.drm_fd, query=amdgpu_drm.AMDGPU_INFO_DEV_INFO,
-      return_pointer=ctypes.addressof(inf:=amdgpu_drm.struct_drm_amdgpu_info_device()), return_size=ctypes.sizeof(inf))
-    return inf
-  def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return ((self.drm_dev_info.cu_bitmap[se % 4][sa + (se // 4) * 2] >> (2 * wgp)) & 0x3) == 0x3
-
-class PCIIface(PCIIfaceBase):
-  def __init__(self, dev, dev_id):
-    super().__init__(dev, dev_id, vendor=0x1002, devices=((0xffff, (0x74a1,0x744c,0x7480,0x7550,0x7551,0x7590,0x75a0)),), vram_bar=0,
-      va_start=AMMemoryManager.va_allocator.base, va_size=AMMemoryManager.va_allocator.size, dev_impl_t=AMDev)
-    self._compute_props()
-
-  def p2p_paddrs(self, paddrs:list[tuple[int,int]]) -> tuple[list[tuple[int,int]], AddrSpace]:
-    return ([(self.dev_impl.paddr2xgmi(p), sz) for p, sz in paddrs], AddrSpace.PEER) if self.dev_impl.is_hive() else super().p2p_paddrs(paddrs)
-
-  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 _compute_props(self):
-    self.ip_versions = self.dev_impl.ip_ver
-
-    gfxver = int(f"{self.dev_impl.ip_ver[am.GC_HWIP][0]:02d}{self.dev_impl.ip_ver[am.GC_HWIP][1]:02d}{self.dev_impl.ip_ver[am.GC_HWIP][2]:02d}")
-    if self.dev_impl.gc_info.header.version_major == 2:
-      cu_per_sa = self.dev_impl.gc_info.gc_num_cu_per_sh
-      max_sh_per_se = self.dev_impl.gc_info.gc_num_sh_per_se
-    else:
-      cu_per_sa = 2 * (self.dev_impl.gc_info.gc_num_wgp0_per_sa + self.dev_impl.gc_info.gc_num_wgp1_per_sa)
-      max_sh_per_se = self.dev_impl.gc_info.gc_num_sa_per_se
-
-    array_count = max_sh_per_se * self.dev_impl.gc_info.gc_num_se * self.dev_impl.gfx.xccs
-    self.props = {'cu_per_simd_array': cu_per_sa, 'simd_count': 2 * cu_per_sa * array_count, 'simd_per_cu': 2, 'array_count': array_count,
-      'max_slots_scratch_cu': self.dev_impl.gc_info.gc_max_scratch_slots_per_cu, 'max_waves_per_simd': self.dev_impl.gc_info.gc_max_waves_per_simd,
-      'simd_arrays_per_engine': max_sh_per_se, 'lds_size_in_kb': self.dev_impl.gc_info.gc_lds_size, 'num_xcc': self.dev_impl.gfx.xccs,
-      'gfx_target_version': {90403: 90402}.get(gfxver, gfxver)}
-
-  def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
-                   xcc_id=0, idx=0):
-    assert cwsr_buffer is None, "no cwsr buffer for am"
-
-    rcvr_params: tuple
-    if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA:
-      doorbell_index = self.dev_impl.sdma.setup_ring(*(rcvr_params:=(ring._buf.va_addr, ring._buf.size, gart._buf.va_addr+rptr,
-        gart._buf.va_addr+wptr, idx)))
-    else:
-      doorbell_index = self.dev_impl.gfx.setup_ring(*(rcvr_params:=(ring._buf.va_addr, ring._buf.size, gart._buf.va_addr+rptr,
-        gart._buf.va_addr+wptr, eop_buffer._buf.va_addr, eop_buffer._buf.size, is_aql:=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL), is_aql)))
-
-    (put_value := Buffer("CPU", 1, dtypes.uint64, preallocate=True))._buf.view.view(fmt='Q')[0] = 0
-    doorbell = Buffer("CPU", 1, dtypes.uint64, options=BufferSpec(external_ptr=self.dev_impl.doorbell64.addr + doorbell_index*8), preallocate=True)
-    return AMDQueueDesc(ring=ring, doorbell=doorbell, read_ptr=gart.view(1, dtypes.uint64, rptr).ensure_allocated(),
-      write_ptr=gart.view(1, dtypes.uint64, wptr).ensure_allocated(), put_value=put_value, eop_buffer=eop_buffer, params=rcvr_params)
-
-  def _collect_interrupts(self, reset=False, drain_only=False):
-    d = self.dev
-    if drain_only: d.iface.dev_impl.ih.drain()
-    else: d.iface.dev_impl.ih.interrupt_handler()
-
-    if reset and d.iface.dev_impl.recover(force=True):
-      cq = d.compute_queue
-      for b in (cq.put_value, cq.read_ptr, cq.write_ptr): b._buf.view.view(fmt='Q')[0] = 0
-      d.iface.dev_impl.gfx.setup_ring(*cq.params)
-      d.signal('timeline')._buf.cpu_view().mv.cast('Q')[0] = \
-        d.signal('value', 1).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] - 1
-
-  def sleep(self, timeout):
-    if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
-      self.pci_dev.irq_fd.read(8 * events_cnt)
-    self._collect_interrupts()
-    if self.dev_impl.is_err_state: raise RuntimeError("Device is in error state")
-
-  def on_device_hang(self):
-    self._collect_interrupts(reset=True)
-    raise RuntimeError("Device hang detected")
-
-  def device_fini(self): self.dev_impl.fini()
-
-def _mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface,), {})
-
-class AMDDevice(HCQ2Compiled):
-  pm_lower = PatternMatcher([
-    # prep program
-    (UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"), amd_build_program),
-
-    # encoding of cmdbuf
-    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),)), encode_queue),
-  ])
-
-  timestamp_divider = 100.0  # AMD GPU clock: ticks/us
-  max_scratch_psize = 0
-
-  ifaces = [KFDIface, PCIIface, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface)]
-
-  def device_props(self): return self.iface.props
-
-  def is_am(self) -> bool: return isinstance(self.iface, (PCIIface,))
-  def is_usb(self) -> bool: return False
-
-  def __init__(self, device:str=""):
-    self.device_id = int(device.split(":")[1]) if ":" in device else 0
-
-    self.iface = self._select_iface()
-
-    self.target:tuple[int, ...] = ((trgt:=self.iface.props['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
-    self.arch = "gfx%d%x%x" % self.target
-    assert (self.target in ((9,4,2),(9,5,0))) or self.target[0] in (11, 12), f"Unsupported arch: {self.arch}"
-    if DEBUG >= 1: print(f"AMDDevice: opening {self.device_id} with target {self.target} arch {self.arch}")
-
-    self.xccs = self.iface.props.get('num_xcc', 1)
-    self.se_cnt = self.iface.props['array_count'] // self.iface.props['simd_arrays_per_engine'] // self.xccs
-    self.cu_cnt = self.iface.props['simd_count'] // self.iface.props['simd_per_cu'] // self.xccs
-    self.waves_per_cu = self.iface.props['max_waves_per_simd'] * self.iface.props['simd_per_cu']
-    self.wave_cnt = (self.cu_cnt * self.waves_per_cu) if self.target[0] != 9 else min(self.cu_cnt * 40, self.se_cnt * self.xccs * 512)
-
-    self.ip_off = importlib.import_module(f"tinygrad.runtime.autogen.am.{'vega' if self.target[0] == 9 else 'navi'}_offsets")
-    self.soc = import_soc(self.target)
-    self.pm4 = importlib.import_module(f"tinygrad.runtime.autogen.am.pm4_{'soc15' if self.target[0] == 9 else 'nv'}")
-    self.sdma = import_module('sdma', min(self.iface.ip_versions[am.SDMA0_HWIP], (6, 0, 0)))
-    self.gc = AMDIP('gc', self.iface.ip_versions[am.GC_HWIP],
-                    bases={i: tuple(getattr(self.ip_off, f'GC_BASE__INST{i}_SEG{s}', 0) for s in range(6)) for i in range(6)})
-
-    self.nbio = AMDIP('nbio' if self.target[0] < 12 else 'nbif', self.iface.ip_versions[am.NBIF_HWIP],
-                      bases={i: tuple(getattr(self.ip_off, f'NBIO_BASE__INST{i}_SEG{s}', 0) for s in range(9)) for i in range(6)})
-
-    self.is_aql = getenv("AMD_AQL", int(self.xccs > 1))
-    if self.is_aql:
-      self.pm4_ibs = self.iface.alloc(0x2000 if self.is_usb() else (16 << 20), uncached=True, cpu_access=True)
-      self.pm4_ib_alloc = BumpAllocator(self.pm4_ibs.size, wrap=True)
-
-    self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
-    self.sdma_queues:dict = {}
-    self.has_copy_queue = not getenv("AMD_DISABLE_SDMA")
-
-    super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, can_recover=self.is_am(), arch=self.arch)
-
-    # Scratch setup
-    self.max_private_segment_size = 0
-    self.pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, tag="scratch", name="b"), lambda ctx, b: ctx[0].scratch_buffer(b.max_numel()))]) + self.pm_bufferize
-
-    self.pmc_enabled:bool = PROFILE > 0 and PMC > 0
-    if self.pmc_enabled:
-      self.iface.require_profile_mode()
-
-      self.pmc_sched:list[PMCSample] = []
-      self.pmc_counters = import_pmc(self.target)
-
-      # validate counters: SQ for SIMD busy/instruction counts, LDS stats, GRBM for GPU cycles, L2 cache hits/misses
-      l2, lds = ("TCC", "SQ") if self.target[0] == 9 else ("GL2C", "SQC")
-      pmc_default = f"SQ_BUSY_CYCLES,SQ_INSTS_VALU,SQ_INSTS_SALU,{lds}_LDS_IDX_ACTIVE,{lds}_LDS_BANK_CONFLICT,GRBM_GUI_ACTIVE,{l2}_HIT,{l2}_MISS"
-      for k in (PMC_COUNTERS:=getenv("PMC_COUNTERS", pmc_default).split(",")):
-        if k not in self.pmc_counters: raise RuntimeError(f"PMC counter {k} is not supported. Available: {','.join(self.pmc_counters.keys())}")
-
-      raise NotImplementedError("PMC start not migrated to hcq2 yet")
-
-    # SQTT is disabled by default because of runtime overhead and big file sizes (~200mb to Tensor.full() two 4096x4096 tensors and matmul them)
-    self.sqtt_enabled:bool = PROFILE > 0 and SQTT > 0
-    if self.sqtt_enabled:
-      self.iface.require_profile_mode()
-
-      SQTT_BUFFER_SIZE = getenv("SQTT_BUFFER_SIZE", 256) # in mb, per shader engine
-      self.sqtt_buffers = [self.allocator.alloc(SQTT_BUFFER_SIZE<<20, BufferSpec(nolru=True, uncached=True)) for _ in range(self.se_cnt * self.xccs)]
-      self.sqtt_wptrs = self.allocator.alloc(round_up(self.se_cnt * self.xccs * 4, 0x1000), BufferSpec(cpu_access=True, nolru=True))
-      self.sqtt_next_cmd_id = itertools.count(0)
-
-  def create_queue(self, queue_type, ring_size, ctx_save_restore_size=0, eop_buffer_size=0, ctl_stack_size=0, debug_memory_size=0, idx=0):
-    ring = Buffer(self.device, ring_size // 4, dtypes.uint32, options=BufferSpec(uncached=True, cpu_access=True), preallocate=True)
-    gart = Buffer(self.device, 0x100, dtypes.uint8, options=BufferSpec(uncached=True, cpu_access=True), preallocate=True)
-
-    if queue_type == kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL:
-      self.aql_gart = gart
-      self.aql_desc = hsa.amd_queue_t(queue_properties=hsa.AMD_QUEUE_PROPERTIES_IS_PTR64 | hsa.AMD_QUEUE_PROPERTIES_ENABLE_PROFILING,
-        read_dispatch_id_field_base_byte_offset=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset,
-        max_cu_id=(self.cu_cnt * self.xccs) - 1, max_wave_id=self.waves_per_cu - 1)
-      self.aql_gart._buf.cpu_view().view(fmt='B')[:ctypes.sizeof(self.aql_desc)] = bytes(self.aql_desc)
-
-    cwsr_buffer_size = round_up((ctx_save_restore_size + debug_memory_size) * self.xccs, mmap.PAGESIZE)
-    cwsr_buffer = Buffer(self.device, cwsr_buffer_size, dtypes.uint8, preallocate=True) if ctx_save_restore_size else None
-    eop_buffer = Buffer(self.device, eop_buffer_size, dtypes.uint8, preallocate=True) if eop_buffer_size else None
-
-    queue = (self.iface.create_queue(queue_type, ring, gart, rptr=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset,
-             wptr=getattr(hsa.amd_queue_t, 'write_dispatch_id').offset, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer,
-             ctx_save_restore_size=ctx_save_restore_size, ctl_stack_size=ctl_stack_size, idx=idx))
-
-    qname = f"{'COPY' if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA else 'COMPUTE'}:{idx}"
-    self.pm_bufferize = PatternMatcher([
-      (UPat(Ops.PARAM, tag=f"{qname}_{name}"), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"]
-    ]) + self.pm_bufferize
-
-    return queue
-
-  @functools.cached_property
-  def compute_queue(self) -> AMDQueueDesc:
-    # https://gitlab.freedesktop.org/agd5f/linux/-/blob/a1fc9f584c4aaf8bc1ebfa459fc57a3f26a290d8/drivers/gpu/drm/amd/amdkfd/kfd_queue.c#L391
-    sgrp_size_per_cu, hwreg_size_per_cu = 0x4000, 0x1000
-    lds_size_per_cu = self.iface.props["lds_size_in_kb"] << 10 if self.target[:2] == (9,5) else 0x10000
-    vgpr_size_per_cu = 0x60000 if self.target in {(11,0,0), (11,0,1), (11,5,1), (12,0,0), (12,0,1)} else 0x80000 if self.target[0] == 9 else 0x40000
-    wg_data_size = round_up((vgpr_size_per_cu + sgrp_size_per_cu + lds_size_per_cu + hwreg_size_per_cu) * self.cu_cnt, mmap.PAGESIZE)
-    ctl_stack_size = round_up((12 if self.target[0] != 9 else 8) * self.wave_cnt + 8 + 40, mmap.PAGESIZE)
-    return self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL if self.is_aql else kfd.KFD_IOC_QUEUE_TYPE_COMPUTE,
-      0x2000 if self.is_usb() else (16 << 20), eop_buffer_size=0x1000,
-      ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size,
-      debug_memory_size=round_up(self.wave_cnt * 32, 64))
-
-  def sdma_queue(self, idx:int):
-    if getenv("AMD_DISABLE_SDMA"): return None
-    if idx in self.sdma_queues: return self.sdma_queues[idx]
-    with contextlib.suppress(OSError):
-      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20), idx=idx)
-    return self.sdma_queues.get(idx, None)
-
-  def tmpring_size(self, private_segment_size):
-    private_segment_size = max(private_segment_size, 128)
-
-    lanes_per_wave = 64 # wave64
-    mem_alignment_size = 256 if self.target[0] != 9 else 1024
-    size_per_thread = round_up(private_segment_size, mem_alignment_size // lanes_per_wave)
-    size_per_xcc = size_per_thread * lanes_per_wave * self.iface.props['max_slots_scratch_cu'] * self.cu_cnt
-
-    # NOTE: xcc logic is correct only for GFX9.
-    max_scratch_waves = self.cu_cnt * self.iface.props['max_slots_scratch_cu'] * self.xccs
-    wave_scratch = ceildiv(lanes_per_wave * size_per_thread, mem_alignment_size)
-    num_waves = (size_per_xcc // (wave_scratch * mem_alignment_size)) // (self.se_cnt if self.target[0] != 9 else 1)
-
-    tmpring_t = getattr(hsa, f'union_COMPUTE_TMPRING_SIZE{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
-    tmpring = int.from_bytes(tmpring_t(WAVES=min(num_waves, max_scratch_waves), WAVESIZE=wave_scratch), 'little')
-
-    if hasattr(self, 'aql_desc'):
-      gfx9_rsrc = {'NUM_FORMAT':hsa.BUF_NUM_FORMAT_UINT, 'DATA_FORMAT':hsa.BUF_DATA_FORMAT_32, 'ELEMENT_SIZE':1, 'INDEX_STRIDE':3}
-      rsrc = {'DST_SEL_X':hsa.SQ_SEL_X, 'DST_SEL_Y':hsa.SQ_SEL_Y, 'DST_SEL_Z':hsa.SQ_SEL_Z, 'DST_SEL_W':hsa.SQ_SEL_W, 'ADD_TID_ENABLE':1,
-              'TYPE':hsa.SQ_RSRC_BUF, **(gfx9_rsrc if self.target[0] == 9 else {'FORMAT':hsa.BUF_FORMAT_32_UINT, 'OOB_SELECT':2})}
-      rsrc1_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD1{"_GFX11" if self.target[0] != 9 else ""}_bitfields')
-      rsrc3_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD3{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
-
-      self.aql_desc.scratch_backing_memory_location = int(self.scratch.get_buf().va_addr)
-      self.aql_desc.scratch_wave64_lane_byte_size = self.max_private_segment_size * lanes_per_wave // 64
-      self.aql_desc.scratch_resource_descriptor[:] = [lo32(self.scratch.get_buf().va_addr),
-        int.from_bytes(rsrc1_t(BASE_ADDRESS_HI=hi32(self.scratch.get_buf().va_addr), SWIZZLE_ENABLE=1), 'little'),
-        lo32(size_per_xcc), int.from_bytes(bytes(rsrc3_t(**rsrc)), 'little')]
-      self.aql_desc.compute_tmpring_size = tmpring
-      self.aql_gart._buf.cpu_view()[:ctypes.sizeof(self.aql_desc)] = bytes(self.aql_desc)
-
-    return tmpring
-
-  def scratch_buffer(self, private_segment_size):
-    AMDDevice.max_scratch_psize = private_segment_size = max(private_segment_size, 128, AMDDevice.max_scratch_psize)
-    if self.max_private_segment_size < private_segment_size:
-      lanes_per_wave = 64 # wave64
-      mem_alignment_size = 256 if self.target[0] != 9 else 1024
-      size_per_thread = round_up(private_segment_size, mem_alignment_size // lanes_per_wave)
-      size_per_xcc = size_per_thread * lanes_per_wave * self.iface.props['max_slots_scratch_cu'] * self.cu_cnt
-      self.scratch = Buffer(self.device, size_per_xcc * self.xccs, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)
-      self.max_private_segment_size = private_segment_size
-    return self.scratch
-
-  def on_device_hang(self): self.iface.on_device_hang()
-
-  def device_props(self): return self.iface.props

extra/hip_large_kernel.py

@@ -18,7 +18,7 @@ prg = dev.runtime("write_ones", mbin)
 prg(buf0._buf, global_size=(1,65537,1), local_size=(1,1,1), wait=True)
 
 import numpy as np
-def to_np(buf): return np.frombuffer(buf.as_memoryview().cast(buf.dtype.fmt), dtype=_to_np_dtype(buf.dtype))
+def to_np(buf): return np.frombuffer(buf.as_memoryview().cast(buf.dtype.base.fmt), dtype=_to_np_dtype(buf.dtype.base))
 
 big = to_np(buf0)
 print(big)