fix assign copy

assign after copy shouldn't contig
2026-08-17 00:18:28 +00:00 · 2026-02-18 12:14:18 +08:00 · 2026-02-18 12:03:20 +08:00
367 changed files with 22483 additions and 15597 deletions
@@ -45,16 +45,12 @@ inputs:
    description: "Install mesa"
    required: false
    default: 'false'
-  tinydreno:
-    description: "Install tinydreno"
-    required: false
-    default: 'false'
 runs:
  using: "composite"
  steps:
    - name: Set up Python ${{ inputs.python-version }}
      id: setup-python
-      uses: actions/setup-python@v6
+      uses: actions/setup-python@v5
      with:
        python-version: ${{ inputs.python-version }}

@@ -66,14 +62,14 @@ runs:
      uses: actions/cache/restore@v4
      with:
        path: ${{ github.workspace }}/.venv
-        key: venv-${{ runner.os }}-${{ runner.arch }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
+        key: venv-${{ runner.os }}-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
+      uses: actions/cache@v4
      with:
        path: ${{ github.workspace }}/.venv
-        key: venv-${{ runner.os }}-${{ runner.arch }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
+        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}

    # **** Caching downloads ****

@@ -85,7 +81,7 @@ runs:
        key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
    - name: Cache downloads
      if: inputs.key != '' && github.event_name != 'pull_request'
-      uses: actions/cache@v5
+      uses: actions/cache@v4
      with:
        path: ${{ runner.os == 'Linux' && '~/.cache/tinygrad/downloads/' || '~/Library/Caches/tinygrad/downloads/' }}
        key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
@@ -199,13 +195,13 @@ runs:
      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 }}
+        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
    - name: Cache apt
      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
+      uses: actions/cache@v4
      with:
        path: /var/cache/apt/archives/
-        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-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.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
@@ -237,7 +233,7 @@ runs:
      shell: bash
      run: |
        sudo mkdir -p /usr/local/lib
-        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/tinygrad/amdcomgr_dylib/releases/latest | \
+        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/nimlgen/amdcomgr_dylib/releases/latest | \
          jq -r '.assets[] | select(.name == "libamd_comgr.dylib").browser_download_url' | \
          sudo xargs curl -fL -o /usr/local/lib/libamd_comgr.dylib
        cargo build --release --manifest-path ./extra/remu/Cargo.toml
@@ -269,7 +265,7 @@ runs:
    - name: Cache gpuocelot
      if: inputs.ocelot == 'true' && github.event_name != 'pull_request'
      id: cache-build
-      uses: actions/cache@v5
+      uses: actions/cache@v4
      env:
        cache-name: cache-gpuocelot-build-1
      with:
@@ -287,7 +283,6 @@ runs:

        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

@@ -331,9 +326,3 @@ runs:
      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
@@ -28,11 +28,10 @@ jobs:
    timeout-minutes: 15
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
-        key: 'autogen'
        opencl: 'true'
        amd: 'true'
        cuda: 'true'
@@ -56,7 +55,6 @@ jobs:
        python3 -c "from tinygrad.runtime.autogen import libusb"
        python3 -c "from tinygrad.runtime.autogen import mesa"
        python3 -c "from tinygrad.runtime.autogen import avcodec"
-        python3 -c "from tinygrad.runtime.autogen import llvm_qcom"
        REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
    - name: Check for differences
      run: |
@@ -68,7 +66,7 @@ jobs:
        fi
    - name: Upload patch artifact
      if: failure()
-      uses: actions/upload-artifact@v7
+      uses: actions/upload-artifact@v4
      with:
        name: autogen-ubuntu-patch
        path: autogen-ubuntu.patch
@@ -79,11 +77,10 @@ jobs:
    timeout-minutes: 15
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
-        key: 'autogen-mac'
        llvm: 'true'
    - name: Regenerate autogen files
      run: |
@@ -99,7 +96,7 @@ jobs:
        fi
    - name: Upload patch artifact
      if: failure()
-      uses: actions/upload-artifact@v7
+      uses: actions/upload-artifact@v4
      with:
        name: autogen-macos-patch
        path: autogen-macos.patch
@@ -110,11 +107,9 @@ jobs:
    timeout-minutes: 15
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: 'autogen-comgr'
    - name: Install autogen support packages
      run: |
        wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
@@ -138,7 +133,7 @@ jobs:
        fi
    - name: Upload patch artifact
      if: failure()
-      uses: actions/upload-artifact@v7
+      uses: actions/upload-artifact@v4
      with:
        name: autogen-comgr2-patch
        path: autogen-comgr2.patch
@@ -21,9 +21,6 @@ jobs:
  # the 3 minute timeout should not be raised
  testmacpytest:
    name: Mac pytest
-    env:
-      CI: ""
-      CAPTURE_PROCESS_REPLAY: "0"
    runs-on: [self-hosted, macOS]
    timeout-minutes: 3
    defaults:
@@ -32,7 +29,7 @@ jobs:
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    # brew install uv
    - name: setup python environment
      run: |
@@ -44,12 +41,22 @@ jobs:
      run: |
        echo "CACHEDB=/tmp/pytest-db-ci.db" >> $GITHUB_ENV
        rm -f /tmp/pytest-db-ci*
+    # TODO: remove this step once all old caches are migrated
+    - name: Migrate old huggingface cache (symlinks break onnxruntime 1.24+)
+      run: |
+        cd ~/Library/Caches/tinygrad/downloads/models 2>/dev/null || exit 0
+        for old_dir in models--*; do
+          [ -d "$old_dir" ] || continue
+          repo_id=$(echo "$old_dir" | sed 's/models--//; s/--/\//g')
+          snapshot=$(ls -1 "$old_dir/snapshots" 2>/dev/null | head -1)
+          [ -n "$snapshot" ] || continue
+          mkdir -p "$repo_id"
+          cp -RLn "$old_dir/snapshots/$snapshot/"* "$repo_id/" 2>/dev/null || true
+        done
    - name: Run pytest -nauto
      run: |
        source /tmp/tinygrad_pytest_ci/bin/activate
        pytest -nauto --durations=20
-    - name: openpilot compile3 0.10.1 driving_vision
-      run: FLOAT16=1 DEV=CL IMAGE=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx

  testmacbenchmark:
    name: Mac Benchmark
@@ -64,7 +71,7 @@ jobs:
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -94,17 +101,17 @@ jobs:
    - name: Run SDXL
      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing
    - name: Run model inference benchmark
-      run: DEV=METAL NOCLANG=1 python3.11 test/external/external_model_benchmark.py
+      run: METAL=1 NOCLANG=1 python3.11 test/external/external_model_benchmark.py
    - name: Test speed vs torch
      run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py
    - name: Test tensor cores
-      run: DEV=METAL python3.11 test/opt/test_tensor_cores.py
+      run: METAL=1 python3.11 test/opt/test_tensor_cores.py
    - name: Test AMX tensor cores
      run: |
-        DEBUG=2 DEV=CPU CPU_LLVM=0 AMX=1 python3.11 test/opt/test_tensor_cores.py
-        DEBUG=2 DEV=CPU CPU_LLVM=1 AMX=1 python3.11 test/opt/test_tensor_cores.py
-        DEBUG=2 DEV=CPU CPU_LLVM=0 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
-        DEBUG=2 DEV=CPU CPU_LLVM=1 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
+        DEBUG=2 CPU=1 CPU_LLVM=0 AMX=1 python3.11 test/opt/test_tensor_cores.py
+        DEBUG=2 CPU=1 CPU_LLVM=1 AMX=1 python3.11 test/opt/test_tensor_cores.py
+        DEBUG=2 CPU=1 CPU_LLVM=0 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
+        DEBUG=2 CPU=1 CPU_LLVM=1 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
    - name: Run Tensor Core GEMM (float)
      run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py
    - name: Run Tensor Core GEMM (half)
@@ -112,7 +119,7 @@ jobs:
    - name: Run Tensor Core GEMM (bfloat16)
      run: DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py
    - name: Fuzz Padded Tensor Core GEMM
-      run: DEV=METAL M_START=6 M_STOP=10 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=6 K_STOP=24 K_STEP=1 TC_OPT=2 DEBUG=2 python3.11 ./extra/gemm/fuzz_matmul.py
+      run: METAL=1 M_START=6 M_STOP=10 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=6 K_STOP=24 K_STEP=1 TC_OPT=2 DEBUG=2 python3.11 ./extra/gemm/fuzz_matmul.py
    - name: Run LLaMA
      run: |
        BENCHMARK_LOG=llama_nojit JIT=0 python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
@@ -154,7 +161,7 @@ jobs:
    # TODO: too slow
    # - name: Run 10 CIFAR training steps w winograd
    #   run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py
-    - uses: actions/upload-artifact@v7
+    - uses: actions/upload-artifact@v4
      with:
        name: Speed (Mac)
        path: |
@@ -174,7 +181,7 @@ jobs:
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -185,17 +192,17 @@ jobs:
        PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
        PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
    - name: UsbGPU boot time
-      run: sudo -E PYTHONPATH=. GMMU=0 DEBUG=2 AM_RESET=1 DEV=AMD AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
+      run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
    - name: UsbGPU tiny tests
-      run: sudo -E PYTHONPATH=. GMMU=0 DEV=AMD AMD_IFACE=USB python3.11 test/test_tiny.py
+      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
    - name: UsbGPU copy speeds
-      run: sudo -E PYTHONPATH=. GMMU=0 DEV=AMD AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
    #- name: UsbGPU openpilot test
-    #  run: sudo -E PYTHONPATH=. GMMU=0 DEV=AMD AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+    #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
    - name: UsbGPU (USB4/TB) boot time
-      run: PYTHONPATH=. DEBUG=3 DEV=NV NV_IFACE=PCI NV_NAK=1 time python3.11 test/test_tiny.py TestTiny.test_plus
+      run: PYTHONPATH=. DEBUG=3 NV=1 NV_IFACE=PCI NV_NAK=1 time python3.11 test/test_tiny.py TestTiny.test_plus
    - name: UsbGPU (USB4/TB) tiny tests
-      run: PYTHONPATH=. DEV=NV NV_IFACE=PCI NV_NAK=1 python3.11 test/test_tiny.py
+      run: PYTHONPATH=. NV=1 NV_IFACE=PCI NV_NAK=1 python3.11 test/test_tiny.py

  testnvidiabenchmark:
    name: tinybox green Benchmark
@@ -207,7 +214,7 @@ jobs:
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Print nvidia-smi
      run: nvidia-smi
    - name: Symlink models and datasets
@@ -227,67 +234,67 @@ jobs:
    - name: reset process replay
      run: test/external/process_replay/reset.py
    - name: Run model inference benchmark
-      run: DEV=NV CAPTURE_PROCESS_REPLAY=0 NOCLANG=1 python3 test/external/external_model_benchmark.py
+      run: NV=1 CAPTURE_PROCESS_REPLAY=0 NOCLANG=1 python3 test/external/external_model_benchmark.py
    - name: Test speed vs torch
-      run: DEV=NV CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py
+      run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py
    - name: Test speed vs theoretical
-      run: DEV=NV IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+      run: NV=1 IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
    - name: Test benchmark allreduce
-      run: DEV=NV python test/external/external_benchmark_multitensor_allreduce.py
+      run: NV=1 python test/external/external_benchmark_multitensor_allreduce.py
    - name: Test tensor cores
      run: |
-        DEV=NV ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
-        DEV=NV NV_PTX=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
+        NV=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
+        NV=1 NV_PTX=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
    - name: Run Tensor Core GEMM (CUDA)
      run: |
-        DEV=CUDA SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-        DEV=CUDA SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-        DEV=CUDA SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
-        DEV=CUDA SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+        CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+        CUDA=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+        CUDA=1 SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
+        CUDA=1 SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
    - name: Run Tensor Core GEMM (PTX)
-      run: DEV=NV NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+      run: NV=1 NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
    - name: Run Tensor Core GEMM (NV)
-      run: DEV=NV SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-    - name: Test DEV=NV
-      run: DEBUG=2 DEV=NV python -m pytest -rA test/test_tiny.py
-    - name: Test DEV=CUDA
-      run: DEBUG=2 DEV=CUDA python -m pytest -rA test/test_tiny.py
+      run: NV=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+    - name: Test NV=1
+      run: DEBUG=2 NV=1 python -m pytest -rA test/test_tiny.py
+    - name: Test CUDA=1
+      run: DEBUG=2 CUDA=1 python -m pytest -rA test/test_tiny.py
    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion DEV=NV python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
+      run: BENCHMARK_LOG=stable_diffusion NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
    # TODO: too slow
    # - name: Run SDXL
-    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 DEV=NV CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing
+    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing
    - name: Run LLaMA
      run: |
-        BENCHMARK_LOG=llama_nojit DEV=NV JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
-        BENCHMARK_LOG=llama DEV=NV JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=llama_nojit NV=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=llama NV=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
    - name: Run LLaMA with BEAM
-      run: BENCHMARK_LOG=llama_beam DEV=NV JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
+      run: BENCHMARK_LOG=llama_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
    # - name: Run LLaMA 7B on 4 GPUs
-    #   run: DEV=NV CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing
+    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing
    # - name: Run LLaMA 7B on 6 GPUs
-    #   run: DEV=NV CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
+    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
    - name: Run LLaMA-3 8B BEAM
-      run: BENCHMARK_LOG=llama3_beam DEV=NV JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
+      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
-      run: BENCHMARK_LOG=llama3_beam_4gpu DEV=NV JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
+      run: BENCHMARK_LOG=llama3_beam_4gpu NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
    - name: Run quantized LLaMA3
      run: BENCHMARK_LOG=llama3_fp8 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --temperature 0 --benchmark --quantize fp8
    # - name: Run LLaMA-3 8B on 6 GPUs
-    #   run: DEV=NV CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
+    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
    # - name: Run LLaMA-2 70B
-    #   run: DEV=NV CAPTURE_PROCESS_REPLAY=0 MAX_CONTEXT=256 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
+    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 MAX_CONTEXT=256 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
    - name: Run Mixtral 8x7B
-      run: time BENCHMARK_LOG=mixtral DEV=NV CAPTURE_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing
+      run: time BENCHMARK_LOG=mixtral NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing
    - name: Run GPT2
      run: |
-        BENCHMARK_LOG=gpt2_nojit DEV=NV JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
-        BENCHMARK_LOG=gpt2 DEV=NV JIT=1 ASSERT_MIN_STEP_TIME=4 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=gpt2_nojit NV=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=gpt2 NV=1 JIT=1 ASSERT_MIN_STEP_TIME=4 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
    - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half DEV=NV HALF=1 ASSERT_MIN_STEP_TIME=6 python3 examples/gpt2.py --count 10 --temperature 0 --timing
+      run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=6 python3 examples/gpt2.py --count 10 --temperature 0 --timing
    - name: Run GPT2 w HALF/BEAM
-      run: BENCHMARK_LOG=gpt2_half_beam DEV=NV HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing
-    - uses: actions/upload-artifact@v7
+      run: BENCHMARK_LOG=gpt2_half_beam NV=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing
+    - uses: actions/upload-artifact@v4
      with:
        name: Speed (NVIDIA)
        path: |
@@ -305,7 +312,7 @@ jobs:
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -325,35 +332,35 @@ jobs:
      run: test/external/process_replay/reset.py
    # TODO: too slow
    # - name: Fuzz Padded Tensor Core GEMM (NV)
-    #   run: DEV=NV M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
+    #   run: NV=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
    # TODO: too slow
    # - name: Fuzz Padded Tensor Core GEMM (PTX)
-    #   run: DEV=NV NV_PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
+    #   run: NV=1 NV_PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
    - name: HEVC Decode Benchmark
-      run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 DEV=NV PYTHONPATH=. python3 extra/hevc/decode.py
+      run: VALIDATE=1 MAX_FRAMES=100 JITBEAM=1 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
    - name: Train MNIST
-      run: time PYTHONPATH=. DEV=NV TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
+      run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
    - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=120 DEV=NV STEPS=10 python3 examples/hlb_cifar10.py
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 python3 examples/hlb_cifar10.py
    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=120 DEV=NV STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=110 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
    - name: Run 10 CIFAR training steps w BF16
-      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=120 DEV=NV STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
+      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
    # - name: Run 10 CIFAR training steps w winograd
-    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 DEV=NV WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
+    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar DEV=NV DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 DEV=NV DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    - name: Run MLPerf resnet eval on training data
-      run: time BENCHMARK_LOG=resnet_eval DEV=NV MODEL=resnet python3 examples/mlperf/model_eval.py
+      run: time BENCHMARK_LOG=resnet_eval NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps DEV=NV DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
+      run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: BENCHMARK_LOG=resnet_10steps_6gpu DEV=NV CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
+      run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
    - name: Run 10 MLPerf Bert training steps (6 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu DEV=NV CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+      run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

@@ -367,7 +374,7 @@ jobs:
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setcap to python
      run: ./extra/amdpci/setup_python_cap.sh
    - name: Remove amd modules
@@ -409,18 +416,18 @@ jobs:
    #    python3 -c "import torch; print(torch.__version__)"
    #    LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py
    - name: Test speed vs theoretical
-      run: DEV=AMD IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+      run: AMD=1 IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
    - name: Test tensor cores AMD_LLVM=0
-      run: DEV=AMD AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
+      run: AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
    # TODO: this is flaky
    # - name: Test tensor cores AMD_LLVM=1
-    #   run: DEV=AMD AMD_LLVM=1 python3 test/opt/test_tensor_cores.py
+    #   run: AMD=1 AMD_LLVM=1 python3 test/opt/test_tensor_cores.py
    - name: Run Tensor Core GEMM (AMD)
      run: |
-        DEV=AMD SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
-        DEV=AMD SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
-    - name: Test DEV=AMD
-      run: DEBUG=2 DEV=AMD python -m pytest -rA test/test_tiny.py
+        AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+        AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
+    - name: Test AMD=1
+      run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
    #- name: Test HIP=1
    #  run: DEBUG=2 HIP=1 python -m pytest -rA test/test_tiny.py
    # TODO: AMD compiler bug causes this to fail
@@ -429,43 +436,43 @@ jobs:
    #- name: Remove amdgpu
    #  run: sleep 10 && sudo rmmod amdgpu # sleep a bit to let the driver unload the prev pid.
    - name: Test AM cold start time
-      run: time DEV=AMD AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
+      run: time AMD=1 AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Test AM warm start time
-      run: time DEV=AMD python3 test/test_tiny.py TestTiny.test_plus
+      run: time AMD=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 DEV=AMD python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
+      run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing
    - name: Run SDXL
-      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 DEV=AMD python3 examples/sdxl.py --seed 0 --noshow --timing
+      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing
    - name: Run LLaMA 7B
      run: |
-        BENCHMARK_LOG=llama_nojit DEV=AMD JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
-        BENCHMARK_LOG=llama DEV=AMD JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=llama_nojit AMD=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=llama AMD=1 JIT=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
    - name: Run LLaMA 7B with BEAM
-      run: BENCHMARK_LOG=llama_beam DEV=AMD JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
+      run: BENCHMARK_LOG=llama_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing
    # - name: Run LLaMA 7B on 4 GPUs
-    #   run: DEV=AMD CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing
+    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing
    # - name: Run LLaMA 7B on 6 GPUs
-    #   run: DEV=AMD CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
+    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
    - name: Run LLaMA-3 8B BEAM
-      run: BENCHMARK_LOG=llama3_beam DEV=AMD JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
+      run: BENCHMARK_LOG=llama3_beam AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
-      run: BENCHMARK_LOG=llama3_beam_4gpu DEV=AMD JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
+      run: BENCHMARK_LOG=llama3_beam_4gpu AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
    # - name: Run LLaMA-3 8B on 6 GPUs
-    #   run: DEV=AMD CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
+    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
    #- name: Restore amdgpu
    #  run: sudo modprobe amdgpu
    # - name: Run LLaMA-2 70B
-    #   run: DEV=AMD CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
+    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing
    - name: Run Mixtral 8x7B
-      run: time BENCHMARK_LOG=mixtral DEV=AMD python3 examples/mixtral.py --temperature 0 --count 10 --timing
+      run: time BENCHMARK_LOG=mixtral AMD=1 python3 examples/mixtral.py --temperature 0 --count 10 --timing
    - name: Run GPT2
      run: |
-        BENCHMARK_LOG=gpt2_nojit DEV=AMD JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
-        BENCHMARK_LOG=gpt2 DEV=AMD JIT=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=gpt2_nojit AMD=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
+        BENCHMARK_LOG=gpt2 AMD=1 JIT=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing
    - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half DEV=AMD HALF=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --count 10 --temperature 0 --timing
+      run: BENCHMARK_LOG=gpt2_half AMD=1 HALF=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --count 10 --temperature 0 --timing
    - name: Run GPT2 w HALF/BEAM
-      run: BENCHMARK_LOG=gpt2_half_beam DEV=AMD HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing
+      run: BENCHMARK_LOG=gpt2_half_beam AMD=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

@@ -479,7 +486,7 @@ jobs:
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setcap to python
      run: ./extra/amdpci/setup_python_cap.sh
    - name: Remove amd modules
@@ -504,23 +511,22 @@ jobs:
    - name: reset process replay
      run: test/external/process_replay/reset.py
    - name: Train MNIST
-      run: time PYTHONPATH=. DEV=AMD TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
+      run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
    - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=200 DEV=AMD STEPS=10 python3 examples/hlb_cifar10.py
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py
    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=230 DEV=AMD STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
    # - name: Run 10 CIFAR training steps w BF16
-    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 DEV=AMD STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
+    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py
    # TODO: too slow
    # - name: Run 10 CIFAR training steps w winograd
-    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 DEV=AMD WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
+    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar DEV=AMD DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu DEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
-    # TODO: broken on some of the machines
-    #- name: Test full tinyfs load
-    #  run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
+      run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    - name: Test full tinyfs load
+      run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

@@ -534,7 +540,7 @@ jobs:
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setcap to python
      run: ./extra/amdpci/setup_python_cap.sh
    - name: Remove amd modules
@@ -559,14 +565,14 @@ jobs:
    - name: reset process replay
      run: test/external/process_replay/reset.py
    - name: Run MLPerf resnet eval
-      run: time BENCHMARK_LOG=resnet_eval DEV=AMD MODEL=resnet python3 examples/mlperf/model_eval.py
+      run: time BENCHMARK_LOG=resnet_eval AMD=1 MODEL=resnet python3 examples/mlperf/model_eval.py
    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps DEV=AMD DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
+      run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: BENCHMARK_LOG=resnet_10steps_6gpu DEV=AMD CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
+      run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
    - name: Run 10 MLPerf Bert training steps (6 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu DEV=AMD CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+      run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

@@ -580,7 +586,7 @@ jobs:
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -588,54 +594,33 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: test/external/process_replay/reset.py
-    - name: openpilot compile3 0.11.0 driving_vision
-      run: BENCHMARK_LOG=openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
-    - name: IR3 openpilot compile3 0.11.0 driving_vision
-      run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM QCOM_IR3=1 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot compile3 0.11.0 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
-    - name: openpilot compile3 0.11.0 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_11_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: openpilot compile3 0.10.0 driving_policy
+      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
+    - name: openpilot compile3 0.10.0 dmonitoring
+      run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: DEBUG=2 openpilot compile3 0.10.1 driving_vision
+      run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+    - name: DEBUG=2 IMAGE=1 openpilot compile3 0.10.1 driving_vision
      run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+    - name: IMAGE=1 openpilot compile3 0.10.1 driving_vision
+      run: BENCHMARK_LOG=image_1_openpilot_0_10_1_vision PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.10.1 driving_vision
-      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.10.1 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
    - name: openpilot compile3 0.10.1 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: benchmark MobileNetV2 on DSP
      run: |
        # generate quantized weights
        ln -s /data/home/tiny/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
        ln -s /data/home/tiny/tinygrad/testsig-*.so .
-        PYTHONPATH=. CC=clang-19 DEV=CPU CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
+        PYTHONPATH=. CC=clang-19 CPU=1 CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
        # benchmark on DSP with NOOPT=1, the devectorizer has issues
-        PYTHONPATH=. CC=clang-19 DEV=DSP NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
+        PYTHONPATH=. CC=clang-19 DSP=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

-  testcommausbgpubenchmark:
-    name: UsbGPU Benchmark (comma)
-    runs-on: [self-hosted, Linux, comma4]
-    timeout-minutes: 20
-    defaults:
-      run:
-        shell: bash -e -o pipefail {0}
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: setup staging db
-      if: github.ref == 'refs/heads/update_benchmark_staging'
-      run: |
-        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
-        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: openpilot compile3 0.10.1 driving_vision
-      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." GMMU=0 DEV=AMD AMD_LLVM=1 AMD_IFACE=USB ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot load_pickle 0.10.1 driving_vision
-      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." GMMU=0 DEV=AMD AMD_IFACE=USB ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
-
  testreddriverbenchmark:
    name: AM Benchmark
    runs-on: [self-hosted, Linux, tinyboxrandom]
@@ -646,7 +631,7 @@ jobs:
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setcap to python
      run: ./extra/amdpci/setup_python_cap.sh
    - name: Remove amd modules
@@ -671,40 +656,32 @@ jobs:
    - name: reset process replay
      run: test/external/process_replay/reset.py
    - name: Test driver cold start time
-      run: time DEBUG=3 DEV=AMD AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
+      run: time DEBUG=3 AMD=1 AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Test driver warm start time
-      run: time DEBUG=3 DEV=AMD python3 test/test_tiny.py TestTiny.test_plus
+      run: time DEBUG=3 AMD=1 python3 test/test_tiny.py TestTiny.test_plus
    # Fails on 9070
    # - name: Test tensor cores
    #   run: |
-    #     DEV=AMD AMD_LLVM=0 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
-    #     DEV=AMD AMD_LLVM=1 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
-    #     DEV=AMD SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+    #     AMD=1 AMD_LLVM=0 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
+    #     AMD=1 AMD_LLVM=1 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
+    #     AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
    - name: Run Tensor Core GEMM (AMD)
-      run: DEV=AMD SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
-    - name: Test DEV=AMD
-      run: DEBUG=2 DEV=AMD python -m pytest -rA test/test_tiny.py
+      run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py
+    - name: Test AMD=1
+      run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
    - name: Test DISK copy time
-      run: DEV=AMD TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
+      run: AMD=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
    - name: Test CPU copy time
      run: |
-        DEV=AMD GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
-        DEV=AMD GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
+        AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
+        AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar DEV=AMD DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    # - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-    #   run: BENCHMARK_LOG=resnet_10steps DEV=AMD MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
+    #   run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
    - name: Run 10 MLPerf Bert training steps (1 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps DEV=AMD CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
-    - name: Remote
-      run: |
-        pkill -f 'extra/remote/serve.py' || true
-        PYTHONPATH=. python3 extra/remote/serve.py 6482 &
-        sleep 1
-        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 DEV=AMD AMD_IFACE=PCI python3 test/test_tiny.py
-        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 DEV=AMD AMD_AQL=1 AMD_IFACE=PCI python3 test/test_tiny.py
-        pkill -f 'extra/remote/serve.py' || true
+      run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

@@ -718,7 +695,7 @@ jobs:
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setcap to python
      run: ./extra/amdpci/setup_python_cap.sh
    - name: Remove nv modules
@@ -743,30 +720,23 @@ jobs:
    - name: reset process replay
      run: test/external/process_replay/reset.py
    - name: Test driver start time
-      run: time DEBUG=3 DEV=NV python3 test/test_tiny.py TestTiny.test_plus
+      run: time DEBUG=3 NV=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Test tensor cores
-      run: DEV=NV ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
+      run: NV=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
    - name: Test DISK copy time
-      run: DEV=NV TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
+      run: NV=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
    - name: Test CPU copy time
      run: |
-        DEV=NV GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
-        DEV=NV GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
+        NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
+        NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
    - name: Test LLAMA-3
-      run: BENCHMARK_LOG=llama3_beam DEV=NV JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0
+      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar DEV=NV DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps DEV=NV MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
+      run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
    - name: Run 10 MLPerf Bert training steps (1 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps DEV=NV CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
-    - name: Remote
-      run: |
-        pkill -f 'extra/remote/serve.py' || true
-        PYTHONPATH=. python3 extra/remote/serve.py 6483 &
-        sleep 1
-        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6483 DEV=NV python3 test/test_tiny.py
-        pkill -f 'extra/remote/serve.py' || true
+      run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
@@ -14,7 +14,7 @@ jobs:

    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Remove amdgpu
      run: sudo rmmod amdgpu || true
    - name: Cleanup running AM processes
@@ -22,10 +22,10 @@ jobs:
    - name: Run SDXL with new search
      # TODO: GCVM_L2_PROTECTION_FAULT_STATUS with llvm19
      run: |
-        BENCHMARK_LOG=search_sdxl PYTHONPATH=. DEV=AMD JITBEAM=2 IGNORE_BEAM_CACHE=1 CCACHE=0 python examples/sdxl.py --noshow --timing --seed 0
+        BENCHMARK_LOG=search_sdxl PYTHONPATH=. AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CCACHE=0 python examples/sdxl.py --noshow --timing --seed 0
    - name: Run SDXL with cached search
      run: |
-        BENCHMARK_LOG=search_sdxl_cached PYTHONPATH=. DEV=AMD JITBEAM=2 python examples/sdxl.py --noshow --timing --seed 0
+        BENCHMARK_LOG=search_sdxl_cached PYTHONPATH=. AMD=1 JITBEAM=2 python examples/sdxl.py --noshow --timing --seed 0
    - name: Run winograd cifar with new search
      run: |
        BENCHMARK_LOG=search_wino_cifar WINO=1 DEFAULT_FLOAT=HALF JITBEAM=4 IGNORE_BEAM_CACHE=1 CCACHE=0 BS=1024 STEPS=500 python examples/hlb_cifar10.py
@@ -10,16 +10,16 @@ jobs:
  deploy:
    runs-on: ubuntu-latest
    steps:
-      - uses: actions/checkout@v6
+      - uses: actions/checkout@v4
      - name: Configure Git Credentials
        run: |
          git config user.name github-actions[bot]
          git config user.email 41898282+github-actions[bot]@users.noreply.github.com
-      - uses: actions/setup-python@v6
+      - uses: actions/setup-python@v5
        with:
          python-version: 3.x
      - run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV
-      - uses: actions/cache@v5
+      - uses: actions/cache@v4
        with:
          key: mkdocs-material-${{ env.cache_id }}
          path: .cache
@@ -16,7 +16,7 @@ jobs:

    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Cleanup running AM processes
      run: python extra/amdpci/am_smi.py --pids --kill
    - name: Symlink datasets
@@ -12,9 +12,9 @@ jobs:
  deploy:
    runs-on: ubuntu-latest
    steps:
-    - uses: actions/checkout@v6
+    - uses: actions/checkout@v4
    - name: Set up Python
-      uses: actions/setup-python@v6
+      uses: actions/setup-python@v2
      with:
        python-version: '3.x'
    - name: Install dependencies
@@ -15,7 +15,7 @@ jobs:
      branchstat: ${{ steps.brstat.outputs.stat}}
    steps:
      - name: Check code from PR branch 
-        uses: actions/checkout@v6
+        uses: actions/checkout@v4
        with:
          repository: ${{ github.event.pull_request.head.repo.full_name }}
          ref: ${{ github.event.pull_request.head.sha }}
@@ -46,18 +46,18 @@ jobs:
    if: needs.checkbranch.outputs.branchstat == 'false'
    steps:
      - name: Checkout code from PR branch
-        uses: actions/checkout@v6
+        uses: actions/checkout@v4
        with:
          repository: ${{ github.event.pull_request.head.repo.full_name }}
          ref: ${{ github.event.pull_request.head.sha }}
          path: pr
        # 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
+        uses: actions/checkout@v4
        with:
          path: base
      - name: Set up Python 3.12
-        uses: actions/setup-python@v6
+        uses: actions/setup-python@v5
        with:
          python-version: '3.12'
      - name: Count Line Diff
@@ -71,7 +71,7 @@ jobs:
          echo "EOF" >> "$GITHUB_ENV"
      - name: Comment Code Line Diff
        continue-on-error: false
-        uses: marocchino/sticky-pull-request-comment@v3
+        uses: marocchino/sticky-pull-request-comment@v2
        with:
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
          ignore_empty: true
@@ -89,7 +89,7 @@ jobs:
    steps:
      - name: Comment Rebase
        continue-on-error: false
-        uses: marocchino/sticky-pull-request-comment@v3
+        uses: marocchino/sticky-pull-request-comment@v2
        with:
          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
          skip_unchanged: true
@@ -1,7 +1,7 @@
 name: Unit Tests
 env:
  # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '18'
+  CACHE_VERSION: '16'
  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
  PYTHONPATH: ${{ github.workspace }}
@@ -21,7 +21,7 @@ jobs:
    timeout-minutes: 20
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
@@ -29,9 +29,9 @@ jobs:
        deps: testing_unit
        llvm: 'true'
    - name: Speed Test
-      run: DEV=CPU CPU_LLVM=1 THREADS=0 python3 test/speed/external_test_speed_v_torch.py
+      run: CPU=1 CPU_LLVM=1 THREADS=0 python3 test/speed/external_test_speed_v_torch.py
    - name: Speed Test (BEAM=2)
-      run: BEAM=2 DEV=CPU CPU_LLVM=1 THREADS=0 python3 test/speed/external_test_speed_v_torch.py
+      run: BEAM=2 CPU=1 CPU_LLVM=1 THREADS=0 python3 test/speed/external_test_speed_v_torch.py

  docs:
    name: Docs
@@ -41,7 +41,7 @@ jobs:
      CHECK_OOB: 0
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
@@ -70,7 +70,7 @@ jobs:
        source venv/bin/activate
        pip install $GITHUB_WORKSPACE
        cp $GITHUB_WORKSPACE/examples/beautiful_mnist.py .
-        BS=2 STEPS=10 MAX_BUFFER_SIZE=0 python beautiful_mnist.py
+        BS=2 STEPS=10 python beautiful_mnist.py
    - name: Test Docs Build
      run: python -m mkdocs build --strict
    - name: Test Docs
@@ -83,7 +83,7 @@ jobs:
      run: DEBUG=100 python3 -c "from tinygrad import Tensor; N = 1024; a, b = Tensor.rand(N, N), Tensor.rand(N, N); c = (a.reshape(N, 1, N) * b.T.reshape(1, N, N)).sum(axis=2); print((c.numpy() - (a.numpy() @ b.numpy())).mean())"
    - name: Compile EfficientNet to C and test it
      run: |
-        DEV=CPU CPU_LLVM=0 python examples/compile_efficientnet.py > recognize.c
+        CPU=1 CPU_LLVM=0 python examples/compile_efficientnet.py > recognize.c
        clang -O2 recognize.c -lm -o recognize
        cat test/models/efficientnet/Chicken.jpg | ./recognize | grep cock

@@ -93,7 +93,7 @@ jobs:
    timeout-minutes: 15
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
@@ -114,11 +114,11 @@ jobs:
    - name: Test one op in torch tests
      run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
    - name: Test Ops with TINY_BACKEND
-      run: DEV=CPU CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/backend/test_ops.py --durations=20
+      run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/backend/test_ops.py --durations=20
    - name: Test in-place operations on views
      run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
    - name: Test multi-gpu
-      run: DEV=CPU CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
+      run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
    - name: Test kernel fusion
      run: python3 extra/torch_backend/test_kernel_fusion.py

@@ -129,7 +129,7 @@ jobs:
    timeout-minutes: 15
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
@@ -141,7 +141,7 @@ jobs:
        sudo apt update || true
        sudo apt install -y --no-install-recommends ninja-build
    - name: Test beautiful_mnist in torch with TINY_BACKEND
-      run: STEPS=20 DEV=CPU TARGET_EVAL_ACC_PCT=90.0 MAX_BUFFER_SIZE=0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+      run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
    - name: Test some torch tests (expect failure)
      run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true

@@ -151,67 +151,67 @@ jobs:
    timeout-minutes: 15
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
        key: be-minimal
        deps: testing_unit
    - name: Test dtype with Python emulator
-      run: DEBUG=1 DEV=PYTHON python3 -m pytest -n=auto test/backend/test_dtype.py test/backend/test_dtype_alu.py
+      run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/backend/test_dtype.py test/backend/test_dtype_alu.py
    - name: Test ops with Python emulator
-      run: DEBUG=2 SKIP_SLOW_TEST=1 DEV=PYTHON python3 -m pytest -n=auto test/backend/test_ops.py --durations=20
+      run: DEBUG=2 SKIP_SLOW_TEST=1 PYTHON=1 python3 -m pytest -n=auto test/backend/test_ops.py --durations=20
    - name: Test uops with Python emulator
-      run: DEV=PYTHON python3 -m pytest test/backend/test_uops.py --durations=20
+      run: PYTHON=1 python3 -m pytest test/backend/test_uops.py --durations=20
    - name: Test symbolic with Python emulator
-      run: DEV=PYTHON python3 test/backend/test_symbolic_ops.py
+      run: PYTHON=1 python3 test/backend/test_symbolic_ops.py
    - name: test_renderer_failures with Python emulator
-      run: DEV=PYTHON python3 -m pytest -rA test/backend/test_renderer_failures.py::TestRendererFailures
-    - name: Test IMAGE support
+      run: PYTHON=1 python3 -m pytest -rA test/backend/test_renderer_failures.py::TestRendererFailures
+    - name: Test IMAGE=2 support
      run: |
-        IMAGE=1 DEV=PYTHON python3 test/backend/test_ops.py TestOps.test_gemm
-        IMAGE=1 DEV=PYTHON python3 test/backend/test_ops.py TestOps.test_simple_conv2d
+        IMAGE=2 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
+        IMAGE=2 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_simple_conv2d
    - name: Test emulated METAL tensor cores
      run: |
-        DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 DEV=PYTHON python3 test/backend/test_ops.py TestOps.test_big_gemm
-        DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 DEV=PYTHON python3 test/opt/test_tensor_cores.py
+        DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_big_gemm
+        DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
    - name: Test emulated AMX tensor cores
-      run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 DEV=PYTHON python3 test/backend/test_ops.py TestOps.test_gemm
+      run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
    - name: Test emulated AMD tensor cores
      run: |
-        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 DEV=PYTHON N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 DEV=PYTHON N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 DEV=PYTHON N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 DEV=PYTHON N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 DEV=PYTHON python3 test/opt/test_tensor_cores.py
+        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
    - name: Test emulated AMD MFMA tensor cores
      run: |
-        DEBUG=2 EMULATE=AMD_MFMA FORWARD_ONLY=1 DEV=PYTHON N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 EMULATE=AMD_MFMA FORWARD_ONLY=1 DEV=PYTHON python3 test/opt/test_tensor_cores.py
+        DEBUG=2 EMULATE=AMD_MFMA FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD_MFMA FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
    - name: Test emulated AMD RDNA4 tensor cores
      run: |
-        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 DEV=PYTHON N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 DEV=PYTHON N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 DEV=PYTHON N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 DEV=PYTHON N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 DEV=PYTHON python3 test/opt/test_tensor_cores.py
+        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
    - name: Test emulated CUDA tensor cores
      run: |
-        DEBUG=2 EMULATE=CUDA FORWARD_ONLY=1 DEV=PYTHON python3 test/backend/test_ops.py TestOps.test_gemm_fp16
-        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 DEV=PYTHON python3 test/backend/test_ops.py TestOps.test_gemm
-        DEBUG=2 EMULATE=CUDA_SM75 FORWARD_ONLY=1 DEV=PYTHON python3 test/backend/test_ops.py TestOps.test_gemm_fp16
-        DEBUG=2 EMULATE=CUDA_SM89 ALLOW_TF32=1 FORWARD_ONLY=1 DEV=PYTHON python3 test/opt/test_tensor_cores.py
+        DEBUG=2 EMULATE=CUDA FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm_fp16
+        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
+        DEBUG=2 EMULATE=CUDA_SM75 FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm_fp16
+        DEBUG=2 EMULATE=CUDA_SM89 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
    - name: Test emulated INTEL OpenCL tensor cores
-      run: DEBUG=2 EMULATE=INTEL FORWARD_ONLY=1 DEV=PYTHON HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
+      run: DEBUG=2 EMULATE=INTEL FORWARD_ONLY=1 PYTHON=1 HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
    - name: Test emulated AMX tensor cores
-      run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 DEV=PYTHON python3 test/opt/test_tensor_cores.py
+      run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
    - name: Test device flop counts
      run: |
-        DEBUG=2 EMULATE=METAL DEV=PYTHON python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 EMULATE=AMD DEV=PYTHON python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 EMULATE=CUDA DEV=PYTHON python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 EMULATE=INTEL DEV=PYTHON python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
-        DEBUG=2 AMX=1 EMULATE=AMX DEV=PYTHON python3 ./test/null/test_uops_stats.py TestUOpsStats.test_simple_matmul
+        DEBUG=2 EMULATE=METAL PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=AMD PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=CUDA PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=INTEL PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 AMX=1 EMULATE=AMX PYTHON=1 python3 ./test/null/test_uops_stats.py TestUOpsStats.test_simple_matmul

  linter:
    name: Linters
@@ -220,7 +220,7 @@ jobs:

    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
@@ -244,37 +244,6 @@ jobs:
    - name: Run TYPED=1
      run: CHECK_OOB=0 DEV=CPU TYPED=1 python test/test_tiny.py

-  nulltest:
-    name: Null Tests
-    runs-on: ubuntu-latest
-    timeout-minutes: 15
-
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: unittest-13
-        pydeps: "pillow ftfy regex pre-commit"
-        deps: testing_unit
-        llvm: 'true'
-        amd: 'true'
-    - name: Run NULL backend tests
-      run: DEV=NULL python -m pytest -n=auto test/null/ --durations=20
-    - name: Run targeted tests on NULL backend
-      run: DEV=NULL python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
-    # TODO: too slow
-    # - name: Run SDXL on NULL backend
-    #   run: DEV=NULL DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
-    - name: Run Clip tests for SD MLPerf on NULL backend
-      run: DEV=NULL python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
-    - name: Run AMD emulated BERT training on NULL backend
-      run: EMULATE=AMD_RDNA4 DEV=NULL NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
-    # TODO: support fake weights
-    #- name: Run LLaMA 7B on 4 fake devices
-    #  run: DEV=NULL python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
-
  unittest:
    name: Unit Tests
    runs-on: ubuntu-latest
@@ -282,7 +251,7 @@ jobs:

    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
@@ -297,8 +266,22 @@ jobs:
      run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
    - name: Run unit tests
      run: |
-        DEV=CPU python test/null/test_device.py TestRunAsModule.test_module_runs
-        DEV=CPU python -m pytest -n=auto test/unit/ --durations=20
+        CPU=1 python test/null/test_device.py TestRunAsModule.test_module_runs
+        CPU=1 python -m pytest -n=auto test/unit/ --durations=20
+    - name: Run NULL backend tests
+      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
+    - name: Run targetted tests on NULL backend
+      run: NULL=1 python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
+    # TODO: too slow
+    # - name: Run SDXL on NULL backend
+    #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
+    - name: Run Clip tests for SD MLPerf on NULL backend
+      run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
+    - name: Run AMD emulated BERT training on NULL backend
+      run: EMULATE=AMD_RDNA4 NULL=1 NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+    # TODO: support fake weights
+    #- name: Run LLaMA 7B on 4 fake devices
+    #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
    - name: Run GC tests
      run: python test/external/external_uop_gc.py
    - name: External Benchmark Schedule
@@ -325,7 +308,7 @@ jobs:
    timeout-minutes: 15
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
@@ -341,7 +324,7 @@ jobs:
    timeout-minutes: 10
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
@@ -362,17 +345,18 @@ jobs:
    timeout-minutes: 15
    steps:
      - name: Checkout Code
-        uses: actions/checkout@v6
+        uses: actions/checkout@v4
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
          key: gpu-image
          deps: testing_unit
          opencl: 'true'
-      - name: Test CL IMAGE=1 ops
+      - name: Test CL IMAGE=2 ops
        run: |
-          DEV=CL IMAGE=1 python -m pytest -n=auto test/backend/test_ops.py --durations=20
-          DEV=CL IMAGE=1 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
+          CL=1 IMAGE=2 python -m pytest -n=auto test/backend/test_ops.py --durations=20
+        # TODO: training is broken
+        # CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -382,7 +366,7 @@ jobs:
    timeout-minutes: 10
    steps:
      - name: Checkout Code
-        uses: actions/checkout@v6
+        uses: actions/checkout@v4
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
@@ -390,13 +374,13 @@ jobs:
          deps: testing
          opencl: 'true'
      - name: Generate Dataset
-        run: DEV=CL extra/optimization/generate_dataset.sh
+        run: CL=1 extra/optimization/generate_dataset.sh
      - name: Run Kernel Count Test
-        run: DEV=CL python -m pytest -n=auto test/external/external_test_opt.py
+        run: CL=1 python -m pytest -n=auto test/external/external_test_opt.py
      - name: Run fused optimizer tests
-        run: DEV=CL FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py test/backend/test_optim.py -k "not muon"
+        run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py test/backend/test_optim.py -k "not muon"
      - name: Upload artifact
-        uses: actions/upload-artifact@v7
+        uses: actions/upload-artifact@v4
        with:
          name: sops.gz
          path: /tmp/sops.gz
@@ -407,7 +391,7 @@ jobs:
    timeout-minutes: 15
    steps:
      - name: Checkout Code
-        uses: actions/checkout@v6
+        uses: actions/checkout@v4
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
@@ -417,13 +401,13 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1486 ALLOWED_GATED_READ_IMAGE=17 FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1397 ALLOWED_GATED_READ_IMAGE=94 FLOAT16=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
      - name: Test openpilot CL compile fp16
-        run: FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+        run: FLOAT16=1 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
      - name: Test openpilot CL compile fp32 (test correctness)
-        run: DEV=CL IMAGE=1 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
+        run: DEBUGCL=1 CL=1 IMAGE=2 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
      - name: Test openpilot LLVM compile fp16
-        run: IMAGE=1 FLOAT16=1 DEV=CPU CPU_LLVM=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+        run: FLOAT16=1 CPU=1 CPU_LLVM=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -436,7 +420,7 @@ jobs:

    steps:
      - name: Checkout Code
-        uses: actions/checkout@v6
+        uses: actions/checkout@v4
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
@@ -445,15 +429,15 @@ jobs:
          python-version: '3.12'
          llvm: 'true'
      - name: Test ONNX (CPU)
-        run: DEV=CPU CPU_LLVM=0 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+        run: CPU=1 CPU_LLVM=0 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
      - name: Test ONNX (LLVM)
-        run: DEV=CPU CPU_LLVM=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+        run: CPU=1 CPU_LLVM=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
      - name: Test ONNX Runner (CPU)
-        run: DEV=CPU CPU_LLVM=0 python3 test/external/external_test_onnx_runner.py
+        run: CPU=1 CPU_LLVM=0 python3 test/external/external_test_onnx_runner.py
      - name: Test Additional ONNX Ops (CPU)
-        run: DEV=CPU CPU_LLVM=0 python3 test/external/external_test_onnx_ops.py
+        run: CPU=1 CPU_LLVM=0 python3 test/external/external_test_onnx_ops.py
      - name: Test Quantize ONNX
-        run: DEV=CPU CPU_LLVM=0 python3 test/backend/test_quantize_onnx.py
+        run: CPU=1 CPU_LLVM=0 python3 test/backend/test_quantize_onnx.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -463,7 +447,7 @@ jobs:
    timeout-minutes: 20
    steps:
      - name: Checkout Code
-        uses: actions/checkout@v6
+        uses: actions/checkout@v4
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
@@ -473,21 +457,21 @@ jobs:
          python-version: '3.12'
          opencl: 'true'
      - name: Test ONNX (CL)
-        run: DEV=CL python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+        run: CL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
      #- name: Test Optimization Helpers
      #  run: DEBUG=1 python3 extra/optimization/test_helpers.py
      #- name: Test Action Space
-      #  run: DEBUG=1 DEV=CL python3 extra/optimization/get_action_space.py
+      #  run: DEBUG=1 CL=1 python3 extra/optimization/get_action_space.py
      - name: Test Beam Search
-        run: DEV=CL IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
+        run: CL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
      - name: Test MLPerf stuff
-        run: DEV=CL python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
-      - name: DEV=NULL beautiful_mnist_multigpu
-        run: DEV=NULL NULL_ALLOW_COPYOUT=1 python examples/beautiful_mnist_multigpu.py
+        run: CL=1 python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
+      - name: NULL=1 beautiful_mnist_multigpu
+        run: NULL=1 NULL_ALLOW_COPYOUT=1 python examples/beautiful_mnist_multigpu.py
      - name: Test Bert training
-        run: DEV=NULL NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+        run: NULL=1 NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
      - name: Test llama 3 training
-        run: DEV=NULL NULL_ALLOW_COPYOUT=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
+        run: NULL=1 NULL_ALLOW_COPYOUT=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -499,18 +483,13 @@ jobs:
      CHECK_OOB: 0
    steps:
      - name: Checkout Code
-        uses: actions/checkout@v6
+        uses: actions/checkout@v4
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
          key: apps_llm
-      - name: Test 1B LLM (llama)
-        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model llama3.2:1b | tee /dev/stderr | grep -i rooster
-      - name: Test 1B LLM (llama q4)
-        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster
-      - name: Test 1B LLM (qwen)
-        # NOTE: qwen is dumb and only knows about female chickens
-        run: echo "What's a female chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model qwen3:0.6b | tee /dev/stderr | grep -i hen
+      - name: Test 1B LLM
+        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm | grep -i rooster

 # ****** Models Tests ******

@@ -520,7 +499,7 @@ jobs:
    timeout-minutes: 15
    steps:
      - name: Checkout Code
-        uses: actions/checkout@v6
+        uses: actions/checkout@v4
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
@@ -529,11 +508,11 @@ jobs:
          opencl: 'true'
          llvm: 'true'
      - name: Test models (llvm)
-        run: DEV=CPU CPU_LLVM=1 python -m pytest -n=auto test/models --durations=20
+        run: CPU=1 CPU_LLVM=1 python -m pytest -n=auto test/models --durations=20
      - name: Test models (opencl)
-        run: DEV=CL python -m pytest -n=auto test/models --durations=20
+        run: CL=1 python -m pytest -n=auto test/models --durations=20
      - name: Test models (cpu)
-        run: DEV=CPU CPU_LLVM=0 python -m pytest -n=auto test/models --durations=20
+        run: CPU=1 CPU_LLVM=0 python -m pytest -n=auto test/models --durations=20
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -543,7 +522,7 @@ jobs:
    timeout-minutes: 20
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
@@ -551,9 +530,9 @@ jobs:
        deps: testing
        python-version: '3.12'
    - name: Test models (Metal)
-      run: DEV=METAL python -m pytest -n=auto test/models --durations=20
+      run: METAL=1 python -m pytest -n=auto test/models --durations=20
    - name: Test LLaMA compile speed
-      run: DEV=METAL python test/external/external_test_speed_llama.py
+      run: METAL=1 python test/external/external_test_speed_llama.py

 # ****** Feature Tests ******

@@ -563,7 +542,7 @@ jobs:
    timeout-minutes: 15
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
@@ -572,11 +551,11 @@ jobs:
        pydeps: "pillow"
        llvm: "true"
    - name: Test LLVM=1 DEVECTORIZE=0
-      run: DEV=CPU CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/backend/test_ops.py
+      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/backend/test_ops.py
    - name: Test LLVM=1 DEVECTORIZE=0 for model
-      run: DEV=CPU CPU_LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
-    - name: Test DEV=CPU DEVECTORIZE=0
-      run: DEV=CPU CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/backend/test_ops.py
+      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
+    - name: Test CPU=1 DEVECTORIZE=0
+      run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/backend/test_ops.py

  testdsp:
    name: Linux (DSP)
@@ -584,7 +563,7 @@ jobs:
    timeout-minutes: 15
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
@@ -593,9 +572,9 @@ jobs:
        pydeps: "onnx==1.18.0 onnxruntime ml_dtypes"
        llvm: "true"
    - name: Set up Docker Buildx
-      uses: docker/setup-buildx-action@v4
+      uses: docker/setup-buildx-action@v3
    - name: Build QEMU Docker with cache
-      uses: docker/build-push-action@v7
+      uses: docker/build-push-action@v4
      with:
        file: extra/dsp/Dockerfile
        push: false
@@ -606,11 +585,11 @@ jobs:
    - name: Set MOCKDSP env
      run: printf "MOCKDSP=1" >> $GITHUB_ENV
    - name: Run test_tiny on DSP
-      run: DEBUG=2 DEV=DSP python test/test_tiny.py
+      run: DEBUG=2 DSP=1 python test/test_tiny.py
    - name: Test transcendentals
-      run: CC=clang-20 DEBUG=2 DEV=DSP python test/backend/test_transcendental.py TestTranscendentalVectorized
+      run: CC=clang-20 DEBUG=2 DSP=1 python test/backend/test_transcendental.py TestTranscendentalVectorized
    - name: Test quantize onnx
-      run: DEBUG=2 DEV=DSP python3 test/backend/test_quantize_onnx.py
+      run: DEBUG=2 DSP=1 python3 test/backend/test_quantize_onnx.py

  testwebgpu:
    name: Linux (WebGPU)
@@ -618,7 +597,7 @@ jobs:
    timeout-minutes: 20
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
@@ -628,11 +607,11 @@ jobs:
        webgpu: 'true'
    - name: Check Device.DEFAULT (WEBGPU) and print some source
      run: |
-        DEV=WEBGPU python -c "from tinygrad import Device; assert Device.DEFAULT == 'WEBGPU', Device.DEFAULT"
-        DEV=WEBGPU DEBUG=4 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
+        WEBGPU=1 python -c "from tinygrad import Device; assert Device.DEFAULT == 'WEBGPU', Device.DEFAULT"
+        WEBGPU=1 DEBUG=4 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Run selected webgpu tests
      run: |
-          DEV=WEBGPU WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/backend --durations=20
+          WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/backend --durations=20
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -641,12 +620,12 @@ jobs:
    runs-on: ubuntu-24.04
    timeout-minutes: 20
    env:
-      DEV: AMD
+      AMD: 1
      PYTHON_REMU: 1
      MOCKGPU: 1
    steps:
      - name: Checkout Code
-        uses: actions/checkout@v6
+        uses: actions/checkout@v4
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
@@ -665,75 +644,46 @@ jobs:
          sudo apt-get update
          sudo apt-get install llvm-21 llvm-21-tools cloc
      - name: Install rocprof-trace-decoder
-        run: sudo PYTHONPATH="." ./extra/sqtt/install_rocprof_decoder.py
+        run: sudo PYTHONPATH="." ./extra/sqtt/install_sqtt_decoder.py
      - name: Run AMD renderer tests
        run: AMD_LLVM=0 python -m pytest -n=auto test/amd/ --durations 20
      - name: Run AMD renderer tests (AMD_LLVM=1)
        run: AMD_LLVM=1 python -m pytest -n=auto test/amd/ --durations 20
-      - name: Run SQTT profiling tests
-        run: PROFILE=1 SQTT=1 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
+      - name: Run TestOps.test_add with SQTT
+        run: |
+          VIZ=-2 DEBUG=5 python3 test/backend/test_ops.py TestOps.test_add
+          extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
      - name: Run AMD emulated tests on NULL backend
        env:
          AMD: 0
        run: |
-          PYTHONPATH=. DEV=NULL EMULATE=AMD python extra/mmapeak/mmapeak.py
-          PYTHONPATH=. DEV=NULL EMULATE=AMD_CDNA4 python3 -m pytest -n=auto test/testextra/test_tk.py test/backend/test_asm_gemm.py
+          PYTHONPATH=. NULL=1 EMULATE=AMD python extra/mmapeak/mmapeak.py
+          PYTHONPATH=. NULL=1 EMULATE=AMD_CDNA4 python3 -m pytest -n=auto test/testextra/test_tk.py test/backend/test_asm_gemm.py
      - name: Run ASM matmul on MOCKGPU
-        run: PYTHONPATH="." DEV=AMD MOCKGPU=1 N=256 python3 extra/gemm/amd_asm_matmul.py
+        run: PYTHONPATH="." AMD=1 MOCKGPU=1 N=256 python3 extra/gemm/amd_asm_matmul.py
      - name: Run LLVM test
        run: AMD_LLVM=1 python test/device/test_amd_llvm.py

-  testmockam:
-    name: Linux (am)
-    runs-on: ubuntu-24.04
-    timeout-minutes: 15
-    env:
-      DEV: AMD
-      MOCKGPU: 1
-      AMD_IFACE: PCI
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v6
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: mockam
-          deps: testing_unit
-          amd: 'true'
-      - name: Run test_tiny on MOCKAM
-        run: python test/test_tiny.py
-      - name: Run test_tiny on MOCKAM USB
-        run: GMMU=0 AMD_IFACE=USB python test/test_tiny.py
-      - name: Run test_hcq on MOCKAM
-        run: python -m pytest test/device/test_hcq.py
-      - name: Run disk copy tests on MOCKAM
-        run: python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk
-      - name: Run test_tiny on MOCKAM Remote
-        run: |
-          python extra/remote/serve.py 6667 &
-          sleep 2
-          REMOTE=127.0.0.1:6667 python test/test_tiny.py
-          REMOTE=127.0.0.1:6667 python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk; kill %1
-
  testamd:
    strategy:
      fail-fast: false
      matrix:
        backend: [amd, amdllvm]
-        arch: [rdna3, rdna4, cdna4]
+        arch: [rdna3, rdna4]
+        #arch: [rdna3, rdna4, cdna4]

    name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
    runs-on: ubuntu-22.04
    timeout-minutes: 15
    env:
-      DEV: AMD
+      AMD: 1
      MOCKGPU: 1
      MOCKGPU_ARCH: ${{ matrix.arch }}
      SKIP_SLOW_TEST: 1
      AMD_LLVM: ${{ matrix.backend == 'amdllvm' && '1' || matrix.backend != 'amdllvm' && '0' }}
    steps:
      - name: Checkout Code
-        uses: actions/checkout@v6
+        uses: actions/checkout@v4
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
@@ -746,9 +696,7 @@ jobs:
          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['AMD'], Device.DEFAULT"
          DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
      - name: Run pytest (amd)
-        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM --durations=20
-      - name: Run disk copy tests
-        run: python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk
+        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/testextra/test_cfg_viz.py test/external/external_test_am.py --durations=20
      - name: Run TRANSCENDENTAL math
        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
      - name: Run process replay tests
@@ -768,7 +716,7 @@ jobs:
      FORWARD_ONLY: 1
    steps:
      - name: Checkout Code
-        uses: actions/checkout@v6
+        uses: actions/checkout@v4
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
@@ -777,7 +725,7 @@ jobs:
          cuda: 'true'
          ocelot: 'true'
      - name: Set env
-        run: printf "${{ matrix.backend == 'ptx' && 'DEV=CUDA\nCUDA_PTX=1' || matrix.backend == 'nv' && 'DEV=NV\nSKIP_SLOW_TEST=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'PTX' && 'CUDA=1\nCUDA_PTX=1' || matrix.backend == 'nv' && 'NV=1\nSKIP_SLOW_TEST=1' }}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CUDA','NV'], Device.DEFAULT"
@@ -801,7 +749,7 @@ jobs:
    timeout-minutes: 20
    steps:
      - name: Checkout Code
-        uses: actions/checkout@v6
+        uses: actions/checkout@v4
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
@@ -811,7 +759,7 @@ jobs:
          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
          mesa: ${{ matrix.backend == 'lvp' && 'true' }}
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'DEV=CPU\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'DEV=CPU\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'DEV=CL' || matrix.backend == 'lvp' && 'DEV=CPU\nCPU_LVP=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'CL=1' || matrix.backend == 'lvp' && 'CPU=1\nCPU_LVP=1' }}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT"
@@ -831,7 +779,7 @@ jobs:
    timeout-minutes: 20
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
@@ -843,27 +791,25 @@ jobs:
        ocelot: 'true'
        llvm: 'true'
    - name: Run unit tests
-      run: DEV=METAL python -m pytest -n=auto test/unit/ --durations=20
+      run: METAL=1 python -m pytest -n=auto test/unit/ --durations=20
    - name: Run NULL backend tests
-      run: DEV=NULL python -m pytest -n=auto test/null/ --durations=20
+      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
    - name: Run ONNX
-      run: DEV=METAL python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+      run: METAL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
    - name: Test tensor core ops (fake)
-      run: DEV=METAL DEBUG=3 TC=2 python test/backend/test_ops.py TestOps.test_gemm
+      run: METAL=1 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
+      run: METAL=1 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
+      run: METAL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
    #- name: Fuzz Test linearizer
-    #  run: DEV=METAL DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
+    #  run: METAL=1 DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
    - name: Run TRANSCENDENTAL math
-      run: DEV=METAL TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
+      run: METAL=1 TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
    - name: Run pytest (amd)
      env:
        MOCKGPU: 1
-        DEV: AMD
+        AMD: 1
        AMD_LLVM: 0
        FORWARD_ONLY: 1
      run: |
@@ -871,7 +817,7 @@ jobs:
    - name: Run pytest (amd with llvm backend)
      env:
        MOCKGPU: 1
-        DEV: AMD
+        AMD: 1
        AMD_LLVM: 1
        FORWARD_ONLY: 1
      run: |
@@ -880,7 +826,7 @@ jobs:
      env:
        MOCKGPU: 1
        NV_PTX: 1
-        DEV: NV
+        NV: 1
        FORWARD_ONLY: 1
        # TODO: failing due to library loading error
        CAPTURE_PROCESS_REPLAY: 0
@@ -895,7 +841,7 @@ jobs:
    timeout-minutes: 10
    steps:
    - name: Checkout Code
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
@@ -903,9 +849,9 @@ jobs:
        deps: testing
        webgpu: 'true'
    - name: Build WEBGPU Efficientnet
-      run: DEV=WEBGPU WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m examples.compile_efficientnet
+      run: WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m examples.compile_efficientnet
    - name: Run selected webgpu tests
-      run: DEV=WEBGPU WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m pytest -n=auto test/backend --durations=20
+      run: WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m pytest -n=auto test/backend --durations=20
    #- name: Clean npm cache
    #  run: npm cache clean --force
    #- name: Install Puppeteer
@@ -924,7 +870,7 @@ jobs:
    #    cp $GITHUB_WORKSPACE/test/web/test_viz.js .
    #    node test_viz.js
    - name: Test ONNX Runner (WEBGPU)
-      run: DEV=WEBGPU python3 test/external/external_test_onnx_runner.py
+      run: WEBGPU=1 python3 test/external/external_test_onnx_runner.py

  osxtests:
    strategy:
@@ -936,7 +882,7 @@ jobs:
    timeout-minutes: 20
    steps:
      - name: Checkout Code
-        uses: actions/checkout@v6
+        uses: actions/checkout@v4
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
@@ -945,7 +891,7 @@ jobs:
          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
          mesa: ${{ matrix.backend == 'lvp' && 'true' }}
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'DEV=CPU\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'DEV=CPU\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'metal' && 'DEV=METAL' || matrix.backend == 'lvp' && 'DEV=CPU\nCPU_LVP=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'metal' && 'METAL=1' || matrix.backend == 'lvp' && 'CPU=1\nCPU_LVP=1' }}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU','LVP':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
@@ -971,7 +917,7 @@ jobs:
    timeout-minutes: 15
    steps:
      - name: Checkout Code
-        uses: actions/checkout@v6
+        uses: actions/checkout@v4
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
@@ -980,7 +926,7 @@ jobs:
          pydeps: ${{ matrix.backend == 'webgpu' && 'dawn-python' || '' }}
      - name: Set env
        shell: bash
-        run:  printf "${{ matrix.backend == 'llvm' && 'DEV=CPU\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'DEV=CPU\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'webgpu' && 'DEV=WEBGPU'}}" >> $GITHUB_ENV
+        run:  printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'webgpu' && 'WEBGPU=1'}}" >> $GITHUB_ENV
      - name: Run unit tests
        if: matrix.backend=='llvm'
        # test_newton_schulz hits RecursionError
@@ -988,7 +934,7 @@ jobs:
      - name: Run NULL backend tests
        if: matrix.backend=='llvm'
        shell: bash
-        run: CPU=0 CPU_LLVM=0 DEV=NULL python -m pytest -n=auto test/null/ --ignore=test/null/test_elf.py --durations=20
+        run: CPU=0 CPU_LLVM=0 NULL=1 python -m pytest -n=auto test/null/ --ignore=test/null/test_elf.py --durations=20
      - name: Run pytest (${{ matrix.backend }})
        shell: bash
        run: |
@@ -1007,7 +953,7 @@ jobs:
    timeout-minutes: 15
    steps:
      - name: Checkout Code
-        uses: actions/checkout@v6
+        uses: actions/checkout@v4
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
@@ -1017,30 +963,7 @@ jobs:
          python-version: '3.12'
      - name: Set env
        shell: bash
-        run: printf "DEV=NULL\nNULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=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
-  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'
-          python-version: '3.12'
-      - name: Set env
-        shell: bash
-        run: printf "DEV=NULL\nNULL_ALLOW_COPYOUT=1\nNULL_QCOMCL=1" >> $GITHUB_ENV
+        run: printf "NULL=1\nNULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'NULL_IR3=1' || matrix.backend == 'nak' && 'NULL_NAK=1' }}" >> $GITHUB_ENV
      - name: Run test_ops
        shell: bash
        run: |
@@ -66,5 +66,3 @@ target
 .mypy_cache
 mutants
 .mutmut-cache
-dagre/
-graphlib/
@@ -3,7 +3,7 @@
 AM driver is a userspace driver targeting AMD's RDNA3/RDNA4. You only need tinygrad to send compute tasks to your GPU!

 ## How to run?
-Make sure that amdgpu module is unloaded and just run tinygrad with `DEV=AMD`!
+Make sure that amdgpu module is unloaded and just run tinygrad with `AMD=1`!

 Optional requirements:

@@ -10,7 +10,7 @@ Directories are listed in order of how they are processed.

 Group UOps into kernels.

-::: tinygrad.schedule.rangeify.get_kernel_graph
+::: tinygrad.schedule.rangeify.get_rangeify_map
    options:
        members: false
        show_labels: false
@@ -3,7 +3,7 @@
 This is a list of environment variable that control the runtime behavior of tinygrad and its examples.
 Most of these are self-explanatory, and are usually used to set an option at runtime.

-Example: `DEV=CL DEBUG=4 python3 -m pytest`
+Example: `CL=1 DEBUG=4 python3 -m pytest`

 However you can also decorate a function to set a value only inside that function.

@@ -31,7 +31,12 @@ These control the behavior of core tinygrad even when used as a library.
 Variable | Possible Value(s) | Description
 ---|---|---
 DEBUG               | [1-7]      | enable debugging output (operations, timings, speed, generated code and more)
-DEV                 | [AMD, NV, ...] | enable a specific backend
+CL                  | [1]        | enable OpenCL backend
+CUDA                | [1]        | enable CUDA backend
+AMD                 | [1]        | enable AMD backend
+NV                  | [1]        | enable NV backend
+METAL               | [1]        | enable Metal backend (for Mac M1 and after)
+CPU                 | [1]        | enable CPU backend
 BEAM                | [#]        | number of beams in kernel beam search
 DEFAULT_FLOAT       | [HALF, ...]| specify the default float dtype (FLOAT32, HALF, BFLOAT16, FLOAT64, ...), default to FLOAT32
 IMAGE               | [1-2]      | enable 2d specific optimizations
@@ -1,6 +1,6 @@
 # Runtimes

-tinygrad supports various runtimes, enabling your code to scale across a wide range of devices. The default runtime can be automatically selected based on the available hardware, or you can force a specific runtime to be default using environment variables (e.g., `DEV=CPU`).
+tinygrad supports various runtimes, enabling your code to scale across a wide range of devices. The default runtime can be automatically selected based on the available hardware, or you can force a specific runtime to be default using environment variables (e.g., `CPU=1`).

 | Runtime | Description | Compiler Options | Requirements |
 |---------|-------------|------------------|--------------|
@@ -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
+from tinygrad import Tensor, TinyJit, nn, GlobalCounters
 from tinygrad.helpers import getenv, colored, trange
 from tinygrad.nn.datasets import mnist

@@ -15,31 +15,30 @@ class Model:
      nn.BatchNorm(64), Tensor.max_pool2d,
      lambda x: x.flatten(1), nn.Linear(576, 10)]

-  @function
  def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers)

-  @TinyJit
-  @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])
-    loss = self(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
-    return loss.realize(*opt.schedule_step())
-
-  @TinyJit
-  def get_test_acc(self, X_test:Tensor, Y_test:Tensor) -> Tensor: return (self(X_test).argmax(axis=1) == Y_test).mean()*100
-
 if __name__ == "__main__":
  X_train, Y_train, X_test, Y_test = mnist(fashion=getenv("FASHION"))

  model = Model()
  opt = (nn.optim.Muon if getenv("MUON") else nn.optim.SGD if getenv("SGD") else nn.optim.Adam)(nn.state.get_parameters(model))

+  @TinyJit
+  @Tensor.train()
+  def train_step() -> Tensor:
+    opt.zero_grad()
+    samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
+    loss = model(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
+    return loss.realize(*opt.schedule_step())
+
+  @TinyJit
+  def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100
+
  test_acc = float('nan')
  for i in (t:=trange(getenv("STEPS", 70))):
    GlobalCounters.reset()   # NOTE: this makes it nice for DEBUG=2 timing
-    loss = model.train_step(X_train, Y_train)
-    if i%10 == 9: test_acc = model.get_test_acc(X_test, Y_test).item()
+    loss = train_step()
+    if i%10 == 9: test_acc = get_test_acc().item()
    t.set_description(f"loss: {loss.item():6.2f} test_accuracy: {test_acc:5.2f}%")

  # verify eval acc
@@ -5,7 +5,7 @@ from extra.onnx_helpers import get_example_inputs, validate

 def load_onnx_model(onnx_file):
  run_onnx = OnnxRunner(onnx_file)
-  run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(None) for k,v in kwargs.items()}).values())), prune=True)
+  run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(None) for k,v in kwargs.items()}).values())), prune=True, optimize=True)
  return run_onnx_jit, run_onnx.graph_inputs

 if __name__ == "__main__":
@@ -1,10 +1,9 @@
 from pathlib import Path
 from extra.models.efficientnet import EfficientNet
 from tinygrad.tensor import Tensor
-from tinygrad.device import Device
 from tinygrad.nn.state import get_state_dict, safe_save, safe_load, load_state_dict
 from extra.export_model import export_model
-from tinygrad.helpers import fetch
+from tinygrad.helpers import getenv, fetch
 import ast

 if __name__ == "__main__":
@@ -13,13 +12,13 @@ if __name__ == "__main__":
  dirname = Path(__file__).parent
  # exporting a model that's loaded from safetensors doesn't work without loading in from safetensors first
  # loading the state dict from a safetensor file changes the generated kernels
-  if Device.DEFAULT == "WEBGPU":
+  if getenv("WEBGPU"):
    safe_save(get_state_dict(model), (dirname / "net.safetensors").as_posix())
    load_state_dict(model, safe_load(str(dirname / "net.safetensors")))
-  mode = "clang" if Device.DEFAULT == "CPU" else "webgpu" if Device.DEFAULT == "WEBGPU" else ""
+  mode = "clang" if getenv("CPU", "") != "" else "webgpu" if getenv("WEBGPU", "") != "" else ""
  prg, inp_sizes, out_sizes, state = export_model(model, mode, Tensor.randn(1,3,224,224))
-  if Device.DEFAULT != "CPU":
-    ext = "js" if Device.DEFAULT == "WEBGPU" else "json"
+  if getenv("CPU", "") == "":
+    ext = "js" if getenv("WEBGPU", "") != "" else "json"
    with open(dirname / f"net.{ext}", "w") as text_file:
      text_file.write(prg)
  else:
@@ -69,6 +68,6 @@ if __name__ == "__main__":
    else printf("%s\\n", lbls[best_idx]);
  }""")

-    # DEV=CPU python3 examples/compile_efficientnet.py | clang -O2 -lm -x c - -o recognize && DEBUG=1 time ./recognize docs/showcase/stable_diffusion_by_tinygrad.jpg
+    # CPU=1 python3 examples/compile_efficientnet.py | clang -O2 -lm -x c - -o recognize && DEBUG=1 time ./recognize docs/showcase/stable_diffusion_by_tinygrad.jpg
    # category : 281 (tabby, tabby cat) with 9.452788
    print('\n'.join(cprog))
@@ -19,8 +19,8 @@ cifar_std = [0.24703225141799082, 0.24348516474564, 0.26158783926049628]
 BS, STEPS = getenv("BS", 512), getenv("STEPS", 1000)
 EVAL_BS = getenv("EVAL_BS", BS)
 GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))]
-assert BS % len(GPUS) == 0, f"{BS=} is not a multiple of {len(GPUS)=}"
-assert EVAL_BS % len(GPUS) == 0, f"{EVAL_BS=} is not a multiple of {len(GPUS)=}"
+assert BS % len(GPUS) == 0, f"{BS=} is not a multiple of {len(GPUS)=}, uneven multi GPU is slow"
+assert EVAL_BS % len(GPUS) == 0, f"{EVAL_BS=} is not a multiple of {len(GPUS)=}, uneven multi GPU is slow"

 class UnsyncedBatchNorm:
  def __init__(self, sz:int, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1, num_devices=len(GPUS)):
@@ -65,7 +65,17 @@ def loader_process(q_in, q_out, X:Tensor, seed):
      else:
        # pad data with training mean
        img = np.tile(np.array([[[123.68, 116.78, 103.94]]], dtype=np.uint8), (224, 224, 1))
-      X[idx].flatten().assign(img.tobytes())
+
+      # broken out
+      #img_tensor = Tensor(img.tobytes(), device='CPU')
+      #storage_tensor = X[idx].contiguous().realize().lazydata.base.realized
+      #storage_tensor._copyin(img_tensor.numpy())
+
+      # faster
+      X[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = img.tobytes()
+
+      # ideal
+      #X[idx].assign(img.tobytes())   # NOTE: this is slow!
      q_out.put(idx)
    q_out.put(None)

@@ -254,8 +264,8 @@ def load_unet3d_data(preprocessed_dataset_dir, seed, queue_in, queue_out, X:Tens
      x = random_brightness_augmentation(x)
      x = gaussian_noise(x)

-    X[idx].flatten().assign(x.tobytes())
-    Y[idx].flatten().assign(y.tobytes())
+    X[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = x.tobytes()
+    Y[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = y.tobytes()

    queue_out.put(idx)
  queue_out.put(None)
@@ -369,12 +379,12 @@ def load_retinanet_data(base_dir:Path, val:bool, queue_in:Queue, queue_out:Queue
      clipped_match_idxs = np.clip(match_idxs, 0, None)
      clipped_boxes, clipped_labels = tgt["boxes"][clipped_match_idxs], tgt["labels"][clipped_match_idxs]

-      boxes[idx].flatten().assign(clipped_boxes.tobytes())
-      labels[idx].flatten().assign(clipped_labels.tobytes())
-      matches[idx].flatten().assign(match_idxs.tobytes())
-      anchors[idx].flatten().assign(anchor.tobytes())
+      boxes[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = clipped_boxes.tobytes()
+      labels[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = clipped_labels.tobytes()
+      matches[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = match_idxs.tobytes()
+      anchors[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = anchor.tobytes()

-    imgs[idx].flatten().assign(img.tobytes())
+    imgs[idx].contiguous().realize().uop.base.realized.as_memoryview(force_zero_copy=True)[:] = img.tobytes()

    queue_out.put(idx)
  queue_out.put(None)
@@ -396,7 +406,6 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
      queue_in.put((idx, img, tgt))

  def _setup_shared_mem(shm_name:str, size:tuple[int, ...], dtype:dtypes) -> tuple[shared_memory.SharedMemory, Tensor]:
-    shm_name = f"{shm_name}_{os.getpid()}"
    if os.path.exists(f"/dev/shm/{shm_name}"): os.unlink(f"/dev/shm/{shm_name}")
    shm = shared_memory.SharedMemory(name=shm_name, create=True, size=prod(size))
    shm_tensor = Tensor.empty(*size, dtype=dtype, device=f"disk:/dev/shm/{shm_name}")
@@ -3,7 +3,7 @@ from pathlib import Path
 import multiprocessing

 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.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling, profile_marker
 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

@@ -1282,7 +1282,7 @@ def train_bert():
        previous_step = i

 def train_llama3():
-  from examples.mlperf.models.flat_llama import FlatTransformer
+  from extra.models.llama import Transformer
  from examples.llama3 import MODEL_PARAMS
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
  from examples.mlperf.optim import GradAccClipAdamW
@@ -1309,7 +1309,7 @@ def train_llama3():
  EVAL_BS            = config["EVAL_BS"]                = getenv("EVAL_BS", 16)
  EVAL_TARGET        = config["EVAL_TARGET"]            = getenv("EVAL_TARGET", 5.6)

-  # LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. DEV=AMD AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
+  # LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
  # trains to 7

  opt_adamw_beta_1 = 0.9
@@ -1335,43 +1335,49 @@ def train_llama3():
  model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
  # vocab_size from the 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
  print(f"model parameters: {model_params}")

-  # 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 = FlatTransformer(**model_params, max_context=SEQLEN)
-
+  model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
  params = get_parameters(model)
+  # weights are all bfloat16 for now
+  assert params and all(p.dtype == dtypes.bfloat16 for p in params)

  if getenv("FAKEDATA"):
    for v in get_parameters(model):
      v = v.assign(Tensor.empty(v.shape))

-  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))
+  if (DP := getenv("DP", 1)) > 1:
+    device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
+    for v in get_parameters(model):
+      v.shard_(device, axis=None)

-  model.shard(device, is_mp)
+  if (MP := getenv("MP", 1)) > 1:
+    device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
+    for k,v in get_state_dict(model).items():
+      if 'scale' in k: v.shard_(device, axis=None)  # from quantized
+      elif '.attention.wq' in k: v.shard_(device, axis=0)
+      elif '.attention.wk' in k: v.shard_(device, axis=0)
+      elif '.attention.wv' in k: v.shard_(device, axis=0)
+      elif '.attention.wo' in k: v.shard_(device, axis=1)
+      elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
+      elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
+      elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
+      elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
+      elif 'output.weight' in k: v.shard_(device, axis=0)
+      else:
+        # attention_norm, ffn_norm, norm
+        v.shard_(device, axis=None)
+      # prevents memory spike on device 0
+      v.realize()

-  if is_dp: vocab_mask.shard_(device, axis=None).realize()
-  if is_mp: vocab_mask.shard_(device, axis=2).realize()
-
-  is_offload_optim = bool(getenv("OFFLOAD_OPTIM"))
-  is_fake_offload = Device.DEFAULT == "NULL"
-  optim_device = ("CPU" if not is_fake_offload else "NULL:99") if is_offload_optim else None
-  optim = GradAccClipAdamW(params, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2,
-                           eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)
+  optim = GradAccClipAdamW(get_parameters(model), 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)

  # init grads
  for p in optim.params:
-    p.grad = Tensor.zeros_like(p).contiguous()
-  grads = [p.grad for p in optim.params]
+    p.grad = p.zeros_like().contiguous().realize()
+  grads: list[Tensor] = [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)

@@ -1386,47 +1392,51 @@ def train_llama3():

  @TinyJit
  def minibatch(tokens:Tensor):
-    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])
-    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
-
+    tokens = tokens.to(None)
+    if (DP := getenv("DP", 1)) > 1:
+      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
+      tokens = tokens.shard(device, 0)
+    if (MP := getenv("MP", 1)) > 1:
+      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
+      tokens = tokens.shard(device)
+    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
+    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
    loss.backward()
    assert all(p.grad is g for p,g in zip(optim.params, grads))
-
-    loss_cpu = loss.flatten().float().to("CPU")
-    return loss_cpu.realize(*grads)
+    Tensor.realize(loss, *grads)
+    return loss.flatten().float().to("CPU")

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

-    for g in grads: g.assign(g.zeros_like())
+    for g in grads:
+      g.assign(g.zeros_like())

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

-    return lr_cpu, grad_norm_cpu
+    return lr.float().to("CPU")

  @TinyJit
  @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)
-    if not is_sharding: tokens = tokens.to(None)
-    logits:Tensor = model(tokens[:, :-1])
-    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
+    tokens = tokens.to(None)
+    if (DP := getenv("DP", 1)) > 1:
+      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
+      tokens = tokens.shard(device, 0)
+    if (MP := getenv("MP", 1)) > 1:
+      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
+      tokens = tokens.shard(device)
+    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
+    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, model_params["vocab_size"], size=(bs, SEQLEN + 1), dtype=np.int32)
-      yield Tensor(fake_data_np, device="NPY")
+      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=model_params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)

  def get_train_iter():
    if getenv("FAKEDATA", 0):
@@ -1453,53 +1463,49 @@ def train_llama3():
  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):
+      for _ in range(grad_acc):
        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
+        dt = time.perf_counter()
+        loss = minibatch(tokens)
      if stopped: break

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

-      loss = sum(losses) / len(losses)
-      optim_time = et - gt
-      dev_time += optim_time
+      loss = loss.float().item()
+      lr = lr.item()
+
+      et = time.perf_counter()
      step_time = et - st
      gbs_time = gt - st
+      optim_time = ot - gt
+      data_time = dt - ist
+      dev_time = step_time - data_time * grad_acc
      if BENCHMARK: step_times.append(step_time)

      i += 1
-      sequences_seen += actual_gbs
+      sequences_seen += GBS

      mem_gb = GlobalCounters.mem_used / 1e9
      gflops = GlobalCounters.global_ops / 1e9 / dev_time
-      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * 2.3e15)) * 100
+      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * max(getenv("DP", 1), getenv("MP", 1)) * 2.3e15)) * 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())))
+          f"{lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")

      if WANDB:
        wandb.log({
-          "train/loss": loss,
-          "train/lr": lr,
-          "train/grad_norm": grad_norm,
+          "lr": lr, "train/loss": loss,
          "train/step_time": step_time,
          "train/gbs_time": gbs_time,
          "train/optim_time": optim_time,
@@ -1528,7 +1534,7 @@ def train_llama3():
        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 (sequences_seen % EVAL_FREQ == 0 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}")
@@ -1536,7 +1542,7 @@ def train_llama3():
      # run eval
      eval_losses = []
      eval_iter = get_eval_iter()
-      tqdm.write(f"evaluating {EVAL_SAMPLES//EVAL_BS} batches of {EVAL_BS} sequences")
+      tqdm.write(f"evaluating {5760//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()
@@ -1544,7 +1550,7 @@ def train_llama3():
        if BENCHMARK and (j+1) == min(BENCHMARK, EVAL_SAMPLES//EVAL_BS):
          return

-      log_perplexity = sum(eval_losses) / len(eval_losses)
+      log_perplexity = Tensor(eval_losses).mean().float().item()

      tqdm.write(f"eval log perplexity: {log_perplexity:.4f}")

@@ -1,214 +0,0 @@
-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"
-  # CDNA
-  os.environ["EMULATE"] = "AMD_CDNA4"
-  os.environ["DEVICE_IN_FUNCTION_BUG"] = "1"
-  os.environ["ALL2ALL"] = "1"
-  os.environ["USE_ATOMICS"] = "1"
-  if "HK_FLASH_ATTENTION" not in os.environ:
-    os.environ["HK_FLASH_ATTENTION"] = "1"
-    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
-from tinygrad.uop.ops import Ops, UOp
-from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
-
-FP8 = getenv("FP8", 0)
-WQKV = getenv("WQKV", 0)
-
-FP8_DTYPE = dtypes.fp8e4m3
-FP8_MAX = 448.0
-
-def quantize_fp8(x:Tensor):
-  scale = FP8_MAX / (x.abs().max().detach() + 1e-8)
-  x_scaled = x * scale
-  x_clamped = x_scaled + (x_scaled.detach().clamp(-FP8_MAX, FP8_MAX) - x_scaled.detach())  # STE
-  return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal()
-
-def matmul(x:Tensor, w:Tensor) -> Tensor:
-  if not FP8: return x @ w.T
-  # weights are already FP8, just quantize activations
-  x_fp8, x_scale = quantize_fp8(x)
-  return x_fp8.dot(w.T, dtype=dtypes.float) * x_scale
-
-def rmsnorm(x_in:Tensor, eps:float):
-  x = x_in.float()
-  x = x * (x.square().mean(-1, keepdim=True) + eps).rsqrt()
-  return x.cast(x_in.dtype)
-
-class FlatTransformer:
-  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
-               rope_theta:int=10000, max_context:int=1024):
-    self.vocab_size = vocab_size
-    self.n_layers = n_layers
-    self.n_heads = n_heads
-    self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads # n_kv_heads != n_heads implies MQA [arxiv/2307.09288, A.2.1]
-    self.head_dim = dim // n_heads
-    self.n_rep = self.n_heads // self.n_kv_heads
-
-    scaled_std = 0.02 / math.sqrt(2 * n_layers)
-
-    # Attention
-    if WQKV:
-      self.wqkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
-    else:
-      self.wq = self.lin_per_layer(dim, self.n_heads * self.head_dim)
-      self.wk = self.lin_per_layer(dim, self.n_kv_heads * self.head_dim)
-      self.wv = self.lin_per_layer(dim, self.n_kv_heads * self.head_dim)
-    self.wo = self.lin_per_layer(self.n_heads * self.head_dim, dim, std=scaled_std)
-
-    # FeedForward
-    self.w1 = self.lin_per_layer(dim, hidden_dim)
-    self.w2 = self.lin_per_layer(hidden_dim, dim, std=scaled_std)
-    self.w3 = self.lin_per_layer(dim, hidden_dim)
-
-    self.norm_eps = norm_eps
-    self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
-    self.ffn_norm = Tensor.ones(n_layers, dim).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=0.02)
-    self.output = Tensor.normal(1, vocab_size, dim, mean=0.0, std=0.02)
-    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)
-
-  def lin_per_layer(self, in_features:int, out_features:int, std:float=0.02):
-    dt = FP8_DTYPE if FP8 else None
-    if getenv("ZEROS"): return Tensor.zeros(self.n_layers, out_features, in_features, dtype=dt)
-    return Tensor.normal(self.n_layers, out_features, in_features, mean=0.0, std=std, dtype=dt)
-
-  def attention(self, x:Tensor, freqs_cis:Tensor, attention_norm:Tensor, wo:Tensor, wqkv:Tensor|None=None,
-                wq:Tensor|None=None, wk:Tensor|None=None, wv:Tensor|None=None):
-    x = rmsnorm(x, self.norm_eps) * attention_norm
-    bsz, seqlen, _ = x.shape
-
-    if wqkv is not None:
-      xqkv = matmul(x, wqkv)
-      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)
-    else:
-      assert wq is not None and wk is not None and wv is not None
-      xq = matmul(x, wq).reshape(bsz, seqlen, self.n_heads, self.head_dim)
-      xk = matmul(x, wk).reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
-      xv = matmul(x, wv).reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
-
-    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
-    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)
-    return matmul(attn, wo)
-
-  def feed_forward(self, x:Tensor, ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor):
-    x = rmsnorm(x, self.norm_eps) * ffn_norm
-    x_w1 = matmul(x, w1).silu()
-    x_w3 = matmul(x.contiguous_backward(), w3)
-    return matmul(x_w1 * x_w3, w2)
-
-  @function(precompile=True, precompile_backward=True)
-  def run_layer(self, x:Tensor, freqs_cis:Tensor,
-                attention_norm:Tensor, wo:Tensor,
-                ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor,
-                wqkv:Tensor|None=None, wq:Tensor|None=None, wk:Tensor|None=None, wv:Tensor|None=None):
-    h = x + self.attention(x, freqs_cis, attention_norm, wo, wqkv=wqkv, wq=wq, wk=wk, wv=wv)
-    return h + self.feed_forward(h, ffn_norm, w1, w2, w3)
-
-  def shard(self, device:tuple[str, ...], mp:bool=False):
-    from tinygrad.nn.state import get_parameters
-    if not mp:
-      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
-      if WQKV:
-        self.wqkv.shard_(device, axis=1).realize()          # (n_layers, out, dim) shard out
-      else:
-        self.wq.shard_(device, axis=1).realize()            # (n_layers, n_heads*head_dim, dim) shard out
-        self.wk.shard_(device, axis=1).realize()            # (n_layers, n_kv_heads*head_dim, dim) shard out
-        self.wv.shard_(device, axis=1).realize()            # (n_layers, n_kv_heads*head_dim, dim) shard out
-      self.wo.shard_(device, axis=2).realize()             # (n_layers, dim, in) shard in
-      self.w1.shard_(device, axis=1).realize()             # (n_layers, hidden, dim) shard out
-      self.w2.shard_(device, axis=2).realize()             # (n_layers, dim, hidden) shard in
-      self.w3.shard_(device, axis=1).realize()             # (n_layers, hidden, dim) shard out
-      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()
-
-  def __call__(self, tokens:Tensor):
-    h = self.tok_embeddings(tokens)
-    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
-    for i in range(self.n_layers):
-      attn_kwargs = {"wqkv": self.wqkv[i]} if WQKV else {"wq": self.wq[i], "wk": self.wk[i], "wv": self.wv[i]}
-      h = self.run_layer(h, freqs_cis,
-                         self.attention_norm[i], self.wo[i],
-                         self.ffn_norm[i], self.w1[i], self.w2[i], self.w3[i], **attn_kwargs)
-    logits = self.norm(h) @ self.output[0].T
-    return logits
-
-# TODO: this shouldn't be needed, but it prevents a copy of the grads. CAT can help
-def apply_grad(old_grad:UOp, new_grad:UOp) -> list[UOp]:
-  if new_grad.op == Ops.ADD:
-    return apply_grad(old_grad, new_grad.src[0])+apply_grad(old_grad, new_grad.src[1])
-  elif new_grad.op == Ops.PAD:
-    grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(new_grad.src[0].shape, new_grad.marg)])
-    return apply_grad(old_grad.shrink(grad_shrink), new_grad.src[0])
-  else:
-    return [old_grad.store(old_grad + new_grad)]
-
-if __name__ == "__main__":
-  config = {}
-  BS                 = config["BS"]                     = getenv("BS", 16)
-  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
-
-  from examples.llama3 import MODEL_PARAMS
-  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
-  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
-  grads = {x:Tensor.zeros_like(x).contiguous() for x in state.values() if x.requires_grad is None}
-
-  # 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=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 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)):
-        grads[t] = Tensor(grads[t].uop.after(UOp.group(*apply_grad(grads[t].uop, g.uop))), device=t.device)
-    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")):
-      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())))
@@ -1,80 +0,0 @@
-from tinygrad import Tensor, nn
-from tinygrad.helpers import getenv
-from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
-
-class Attention:
-  def __init__(self, dim:int, n_heads:int, n_kv_heads:int|None=None, linear=nn.Linear):
-    self.n_heads = n_heads
-    self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads # n_kv_heads != n_heads implies MQA [arxiv/2307.09288, A.2.1]
-    self.head_dim = dim // n_heads
-    self.n_rep = self.n_heads // self.n_kv_heads
-
-    if getenv("WQKV"):
-      self.wqkv = linear(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2, bias=False)
-    else:
-      self.wq = linear(dim, self.n_heads * self.head_dim, bias=False)
-      self.wk = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
-      self.wv = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
-
-    self.wo = linear(self.n_heads * self.head_dim, dim, bias=False)
-
-  def __call__(self, x:Tensor, freqs_cis:Tensor) -> Tensor:
-    if getenv("WQKV"):
-      xqkv = self.wqkv(x)
-      xqkv = xqkv.reshape(xqkv.shape[0], xqkv.shape[1], self.n_kv_heads, self.n_rep + 2, self.head_dim)
-      xq = xqkv[:, :, :, :self.n_rep].reshape(xqkv.shape[0], xqkv.shape[1], -1)
-      xk = xqkv[:, :, :, self.n_rep:self.n_rep+1].reshape(xqkv.shape[0], xqkv.shape[1], -1)
-      xv = xqkv[:, :, :, self.n_rep+1:self.n_rep+2].reshape(xqkv.shape[0], xqkv.shape[1], -1)
-    else:
-      xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
-
-    xq = xq.reshape(xq.shape[0], xq.shape[1], self.n_heads, self.head_dim)
-    xk = xk.reshape(xk.shape[0], xk.shape[1], self.n_kv_heads, self.head_dim)
-    xv = xv.reshape(xv.shape[0], xv.shape[1], self.n_kv_heads, self.head_dim)
-
-    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
-    bsz, seqlen, _, _ = xq.shape
-
-    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)
-    return self.wo(attn)
-
-class FeedForward:
-  def __init__(self, dim:int, hidden_dim:int, linear=nn.Linear):
-    self.w1 = linear(dim, hidden_dim, bias=False)
-    self.w2 = linear(hidden_dim, dim, bias=False)
-    self.w3 = linear(dim, hidden_dim, bias=False) # the gate in Gated Linear Unit
-
-  def __call__(self, x:Tensor) -> Tensor:
-    w1 = self.w1(x).silu()
-    w3 = self.w3(x)
-    return self.w2(w1 * w3)
-
-class TransformerBlock:
-  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int|None, norm_eps:float, linear=nn.Linear):
-    self.attention = Attention(dim, n_heads, n_kv_heads, linear)
-    self.feed_forward = FeedForward(dim, hidden_dim, linear)
-    self.attention_norm = nn.RMSNorm(dim, norm_eps)
-    self.ffn_norm = nn.RMSNorm(dim, norm_eps)
-
-  def __call__(self, x:Tensor, freqs_cis:Tensor):
-    h = x + self.attention(self.attention_norm(x), freqs_cis)
-    return h + self.feed_forward(self.ffn_norm(h))
-
-class Transformer:
-  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
-               rope_theta:int=10000, max_context:int=1024, linear=nn.Linear, embedding=nn.Embedding):
-    self.layers = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, linear) for _ in range(n_layers)]
-    self.norm = nn.RMSNorm(dim, norm_eps)
-    self.tok_embeddings = embedding(vocab_size, dim)
-    self.output = nn.Linear(dim, vocab_size, bias=False) if embedding == nn.Embedding else linear(dim, vocab_size, bias=False)
-    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)
-
-  def __call__(self, tokens:Tensor):
-    h = self.tok_embeddings(tokens)
-    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
-    for layer in self.layers: h = layer(h, freqs_cis)
-    logits = self.output(self.norm(h))
-    return logits
@@ -1,140 +0,0 @@
-import os
-os.environ["WQKV"] = "1"
-import unittest
-import numpy as np
-from tinygrad import Tensor, nn, dtypes
-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
-
-def copy_weights(flat:FlatTransformer, ref:Transformer):
-  n_layers = flat.n_layers
-  Tensor.realize(*nn.state.get_state_dict(ref).values())
-  flat.wqkv.assign(Tensor(np.stack([ref.layers[i].attention.wqkv.weight.numpy() for i in range(n_layers)])))
-  flat.wo.assign(Tensor(np.stack([ref.layers[i].attention.wo.weight.numpy() for i in range(n_layers)])))
-  flat.w1.assign(Tensor(np.stack([ref.layers[i].feed_forward.w1.weight.numpy() for i in range(n_layers)])))
-  flat.w2.assign(Tensor(np.stack([ref.layers[i].feed_forward.w2.weight.numpy() for i in range(n_layers)])))
-  flat.w3.assign(Tensor(np.stack([ref.layers[i].feed_forward.w3.weight.numpy() for i in range(n_layers)])))
-  flat.attention_norm.assign(Tensor(np.stack([ref.layers[i].attention_norm.weight.numpy() for i in range(n_layers)])))
-  flat.ffn_norm.assign(Tensor(np.stack([ref.layers[i].ffn_norm.weight.numpy() for i in range(n_layers)])))
-  flat.norm.weight.assign(Tensor(ref.norm.weight.numpy()))
-  flat.tok_embeddings.weight.assign(Tensor(ref.tok_embeddings.weight.numpy()))
-  flat.output.weight.assign(Tensor(ref.output.weight.numpy()))
-
-class TestFlatLlama(unittest.TestCase):
-  def test_forward_match(self):
-    Tensor.manual_seed(42)
-    params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
-    ref = Transformer(**params)
-    flat = FlatTransformer(**params)
-    copy_weights(flat, ref)
-    Tensor.realize(*nn.state.get_state_dict(flat).values())
-
-    tokens = Tensor([[1, 50, 100, 999, 2]])
-    ref_logits = ref(tokens).realize()
-    flat_logits = flat(tokens).realize()
-    self.assertEqual(ref_logits.shape, flat_logits.shape)
-    diff = (ref_logits - flat_logits).abs().max().item()
-    self.assertLess(diff, 1e-5, f"forward mismatch: max abs diff {diff}")
-
-  def test_backward_match(self):
-    Tensor.manual_seed(42)
-    params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
-    ref = Transformer(**params)
-    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]])
-
-    ref_loss = ref(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
-    ref_loss.backward()
-    ref_grads = {k: v.grad.numpy() for k, v in nn.state.get_state_dict(ref).items() if v.grad is not None}
-
-    flat_loss = flat(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
-    flat_loss.backward()
-    flat_grads = {k: v.grad.numpy() for k, v in nn.state.get_state_dict(flat).items() if v.grad is not None}
-
-    # check loss matches
-    self.assertAlmostEqual(ref_loss.item(), flat_loss.item(), places=4)
-
-    # check output weight grad matches
-    diff = abs(ref_grads["output.weight"] - flat_grads["output.weight"]).max()
-    self.assertLess(diff, 1e-4, f"output.weight grad mismatch: max abs diff {diff}")
-
-    # check per-layer weight grads match
-    for i in range(params["n_layers"]):
-      for flat_key, ref_key in [
-        ("wqkv", f"layers.{i}.attention.wqkv.weight"),
-        ("wo", f"layers.{i}.attention.wo.weight"),
-        ("w1", f"layers.{i}.feed_forward.w1.weight"),
-        ("w2", f"layers.{i}.feed_forward.w2.weight"),
-        ("w3", f"layers.{i}.feed_forward.w3.weight"),
-      ]:
-        diff = abs(ref_grads[ref_key] - flat_grads[flat_key][i]).max()
-        self.assertLess(diff, 1e-4, f"layer {i} {flat_key} grad mismatch: max abs diff {diff}")
-
-  @unittest.skipUnless(Device.DEFAULT == "CPU", "multi-device CPU test")
-  def test_forward_match_mp(self):
-    Tensor.manual_seed(42)
-    params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
-    from tinygrad import Device
-    devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
-    ref = Transformer(**params)
-    flat = FlatTransformer(**params)
-    copy_weights(flat, ref)
-    Tensor.realize(*nn.state.get_state_dict(flat).values())
-    flat.shard(devices, mp=True)
-
-    tokens = Tensor([[1, 50, 100, 999, 2]], device=devices[0])
-    ref_logits = ref(tokens.to(devices[0])).numpy()
-    flat_logits = flat(tokens.shard(devices)).numpy()
-    self.assertEqual(ref_logits.shape, flat_logits.shape)
-    np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
-
-  @unittest.skipUnless(Device.DEFAULT == "CPU", "multi-device CPU test")
-  def test_forward_match_dp(self):
-    Tensor.manual_seed(42)
-    params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
-    from tinygrad import Device
-    devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
-    ref = Transformer(**params)
-    flat = FlatTransformer(**params)
-    copy_weights(flat, ref)
-    Tensor.realize(*nn.state.get_state_dict(flat).values())
-    flat.shard(devices)
-
-    tokens = Tensor([[1, 50, 100, 999, 2], [2, 100, 50, 1, 999]], device=devices[0])
-    ref_logits = ref(tokens.to(devices[0])).numpy()
-    flat_logits = flat(tokens.shard(devices, axis=0)).numpy()
-    self.assertEqual(ref_logits.shape, flat_logits.shape)
-    np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
-
-  @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
-    try:
-      flat_llama_mod.FP8 = 1
-      Tensor.manual_seed(42)
-      params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
-      ref = Transformer(**params)
-      flat = FlatTransformer(**params)
-      copy_weights(flat, ref)
-      Tensor.realize(*nn.state.get_state_dict(flat).values())
-
-      tokens = Tensor([[1, 50, 100, 999, 2]])
-      ref_logits = ref(tokens).numpy()
-      flat_logits = flat(tokens).numpy()
-      self.assertEqual(ref_logits.shape, flat_logits.shape)
-      # FP8 has lower precision, allow larger tolerance
-      np.testing.assert_allclose(flat_logits, ref_logits, atol=1.0, rtol=0.1)
-    finally:
-      flat_llama_mod.FP8 = old_fp8
-
-if __name__ == "__main__":
-  unittest.main()
@@ -1,80 +1,24 @@
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
-from tinygrad.nn.optim import Optimizer
-from tinygrad.helpers import FUSE_OPTIM, getenv
-from tinygrad.uop.ops import UOp, Ops
+from tinygrad.nn.optim import LAMB
+from tinygrad.helpers import FUSE_OPTIM

-STOCHASTIC_ROUND = getenv("STOCHASTIC_ROUND", 0)
-MASTER_WEIGHTS = getenv("MASTER_WEIGHTS", 0)
-
-def stochastic_round_bf16(x:Tensor) -> Tensor:
-  bits = x.bitcast(dtypes.uint32)
-  if isinstance(x.device, tuple):
-    shape = x.uop.shard_shape if x.uop.axis is not None else x.shape
-    noise = Tensor(UOp(Ops.MSTACK, dtypes.default_float, tuple(Tensor.rand(*shape, device=d).uop for d in x.device)))
-  else:
-    noise = x.rand_like()
-  noise = (noise * 0xFFFF).cast(dtypes.uint32)
-  return ((bits + noise) & 0xFFFF0000).bitcast(dtypes.float32).cast(dtypes.bfloat16)
-
-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, requires_grad=False) for _ in [b1, b2])
-    self.m = self._new_optim_param()
-    self.v = self._new_optim_param()
+class GradAccClipAdamW(LAMB):
+  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, fused=FUSE_OPTIM):
+    super().__init__(params, lr, b1, b2, eps, weight_decay, adam=True, fused=FUSE_OPTIM)
    self.grad_acc, self.clip_norm = grad_acc, clip_norm
-    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:
-      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))
-    to_realize = extra+self.params+self.buffers+(self.master_params or [])
-
-    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)
+      grads[0] = 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))
+      grads[0] = (grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype)
    else:
+      total_norm = Tensor.zeros((), dtype=dtypes.float32, device=self.device)
+      for g in grads:
+        total_norm += g.float().square().sum()
+      total_norm = total_norm.sqrt()
      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
-    for i, g in enumerate(grads):
-      m_new = self.b1 * self.m[i].float() + (1.0 - self.b1) * g.float()
-      v_new = self.b2 * self.v[i].float() + (1.0 - self.b2) * (g.float() * g.float())
-      self.m[i].assign(m_new.cast(self.m[i].dtype))
-      self.v[i].assign(v_new.cast(self.v[i].dtype))
-      m_hat = m_new / (1.0 - self.b1_t)
-      v_hat = v_new / (1.0 - self.b2_t)
-      up = m_hat / (v_hat.sqrt() + self.eps)
-      ret.append(self.lr * up)
-    return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]
-
-  def _apply_update(self, t:Tensor, up:Tensor, master:Tensor|None=None) -> Tensor:
-    w = master if master is not None else t
-    wd = self.wd if t.ndim >= 3 else 0.0
-    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 STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16: return stochastic_round_bf16(new_w)
-    return new_w.cast(t.dtype)
+        grads[i] = grads[i] / self.grad_acc
+        grads[i] = (grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype)
+    return super()._step(params, grads)
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128

@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
 export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024

@@ -1,7 +1,7 @@
 #!/bin/bash
 set -e  # Exit on any error

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_8xMI300X"
 export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

@@ -1,7 +1,7 @@
 #!/bin/bash
 set -e  # Exit on any error

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_green"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

@@ -1,7 +1,7 @@
 #!/bin/bash
 set -e  # Exit on any error

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_red"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192

@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192

@@ -1,7 +1,7 @@
 #!/bin/bash
 set -e  # Exit on any error

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="resnet"
 export SUBMISSION_PLATFORM="tinybox_green"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192

@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192

@@ -1,7 +1,7 @@
 #!/bin/bash
 set -e  # Exit on any error

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="resnet"
 export SUBMISSION_PLATFORM="tinybox_red"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="retinanet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
 export BASEDIR="/raid/datasets/openimages"
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="retinanet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
 export BASEDIR="/raid/datasets/openimages"
@@ -1,7 +1,7 @@
 #!/bin/bash
 set -e  # Exit on any error

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="retinanet"
 export SUBMISSION_PLATFORM="tinybox_green"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="retinanet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
 export BASEDIR="/raid/datasets/openimages"
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="retinanet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
 export BASEDIR="/raid/datasets/openimages"
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128

@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
 export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024

@@ -2,7 +2,7 @@
 set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_8xMI300X"
 export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90

@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90

@@ -2,7 +2,7 @@
 set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_green"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90

@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90

@@ -2,7 +2,7 @@
 set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_red"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192

@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192

@@ -2,7 +2,7 @@
 set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="resnet"
 export SUBMISSION_PLATFORM="tinybox_green"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192

@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192

@@ -2,7 +2,7 @@
 set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="resnet"
 export SUBMISSION_PLATFORM="tinybox_red"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="retinanet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
 export BASEDIR="/raid/datasets/openimages"
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="retinanet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
 export BASEDIR="/raid/datasets/openimages"
@@ -2,7 +2,7 @@
 set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="retinanet"
 export SUBMISSION_PLATFORM="tinybox_green"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="retinanet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
 export BASEDIR="/raid/datasets/openimages"
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="retinanet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
 export BASEDIR="/raid/datasets/openimages"
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128

@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
 export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024

@@ -2,7 +2,7 @@
 set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_8xMI300X"
 export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD DEBUG=0 JIT=1 FLASH_ATTENTION=1
+export PYTHONPATH="." AMD=1 DEBUG=0 JIT=1 FLASH_ATTENTION=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024

@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024

@@ -2,7 +2,7 @@
 set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_8xMI350X"
 export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72

@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72

@@ -2,7 +2,7 @@
 set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_green"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=72 EVAL_BS=72
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96

@@ -2,7 +2,7 @@
 set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_red"
 export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
@@ -1,38 +0,0 @@
-#!/usr/bin/env bash
-
-export PYTHONPATH="."
-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:-2}
-export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
-export ALL2ALL=${ALL2ALL:-1}
-export USE_ATOMICS=${USE_ATOMICS:-0}
-export ASM_GEMM=${ASM_GEMM:-1}
-export WQKV=${WQKV:-1}
-export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
-
-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-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/"
-export LLAMA3_SIZE=${LLAMA3_SIZE:-"405B"}
-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=1 BENCHMARK=10
-if [ -z "$FULL_LAYERS" ]; then
-  export LLAMA_LAYERS=2
-fi
-
-python3 examples/mlperf/model_train.py
@@ -1,32 +0,0 @@
-#!/usr/bin/env bash
-
-export PYTHONPATH="."
-export DEV=${DEV:-AMD}
-export EMULATE="AMD_CDNA4"
-export CHECK_OOB=0
-export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
-
-export DEBUG=${DEBUG:-0}
-export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
-export ALL2ALL=${ALL2ALL:-1}
-export USE_ATOMICS=${USE_ATOMICS:-0}
-export ASM_GEMM=${ASM_GEMM:-1}
-export WQKV=${WQKV:-1}
-export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
-
-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
-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/"
-export LLAMA3_SIZE=${LLAMA3_SIZE:-"405B"}
-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
@@ -5,17 +5,15 @@ 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:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
-export WQKV=${WQKV:-0}

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

 export MODEL="llama3"
@@ -23,7 +21,7 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
 export SEQLEN=${SEQLEN:-8192}
@@ -1,43 +0,0 @@
-#!/usr/bin/env bash
-
-export PYTHONPATH="."
-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:-2}
-export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
-export ALL2ALL=${ALL2ALL:-1}
-export USE_ATOMICS=${USE_ATOMICS:-0}
-export ASM_GEMM=${ASM_GEMM:-1}
-export WQKV=${WQKV:-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=1 BENCHMARK=10
-if [ -z "$FULL_LAYERS" ]; then
-  export LLAMA_LAYERS=2
-fi
-
-python3 examples/mlperf/model_train.py
@@ -5,17 +5,15 @@ 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 USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
-export WQKV=${WQKV:-0}

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

 export MODEL="llama3"
@@ -23,7 +21,7 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
 export SEQLEN=${SEQLEN:-8192}
@@ -1,38 +0,0 @@
-#!/usr/bin/env bash
-
-export PYTHONPATH="."
-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 USE_ATOMICS=${USE_ATOMICS:-0}
-export ASM_GEMM=${ASM_GEMM:-1}
-export WQKV=${WQKV:-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
@@ -1,5 +1,6 @@
 #!/bin/bash
 export BENCHMARK=5
 export EVAL_BS=0
-VIZ=${VIZ:--1} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
-extra/viz/cli.py --profile --device "AMD" | head -23
+export VIZ=${VIZ:--1}
+examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
+PYTHONPATH="." extra/viz/cli.py --profile --device "AMD" --top 20
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192

@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192

@@ -2,7 +2,7 @@
 set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="resnet"
 export SUBMISSION_PLATFORM="tinybox_green"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192

@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192

@@ -2,7 +2,7 @@
 set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="resnet"
 export SUBMISSION_PLATFORM="tinybox_red"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="retinanet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
 export BASEDIR="/raid/datasets/openimages"
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="retinanet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
 export BASEDIR="/raid/datasets/openimages"
@@ -2,7 +2,7 @@
 set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails

-export PYTHONPATH="." DEV=NV
+export PYTHONPATH="." NV=1
 export MODEL="retinanet"
 export SUBMISSION_PLATFORM="tinybox_green"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="retinanet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
 export BASEDIR="/raid/datasets/openimages"
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="." DEV=AMD
+export PYTHONPATH="." AMD=1
 export MODEL="retinanet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
 export BASEDIR="/raid/datasets/openimages"
@@ -1,11 +1,12 @@
 from pathlib import Path
+import numpy as np
 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
 from tinygrad.nn import optim
-from tinygrad.nn.datasets import mnist
+from extra.datasets import fetch_mnist

 class LinearGen:
  def __init__(self):
@@ -37,14 +38,14 @@ class LinearDisc:
    return x

 def make_batch(images):
-  sample = Tensor.randint(batch_size, low=0, high=images.shape[0])
-  return images[sample].reshape(batch_size, 28*28).cast('float').div(127.5).sub(1.0)
+  sample = np.random.randint(0, len(images), size=(batch_size))
+  image_b = images[sample].reshape(-1, 28*28).astype(np.float32) / 127.5 - 1.0
+  return Tensor(image_b)

 def make_labels(bs, col, val=-2.0):
-  y = Tensor.zeros(bs, 2)
-  if col == 0: y = y + Tensor([val, 0.0])
-  else: y = y + Tensor([0.0, val])
-  return y
+  y = np.zeros((bs, 2), np.float32)
+  y[range(bs), [col] * bs] = val  # Can we do label smoothing? i.e -2.0 changed to -1.98789.
+  return Tensor(y)

 def train_discriminator(optimizer, data_real, data_fake):
  real_labels = make_labels(batch_size, 1)
@@ -70,12 +71,12 @@ def train_generator(optimizer, data_fake):

 if __name__ == "__main__":
  # data for training and validation
-  X_train, _, _, _ = mnist()
+  images_real = np.vstack(fetch_mnist()[::2])
  ds_noise = Tensor.randn(64, 128, requires_grad=False)
  # parameters
  epochs, batch_size, k = 300, 512, 1
  sample_interval = epochs // 10
-  n_steps = X_train.shape[0] // batch_size
+  n_steps = len(images_real) // batch_size
  # models and optimizer
  generator = LinearGen()
  discriminator = LinearDisc()
@@ -83,24 +84,24 @@ if __name__ == "__main__":
  output_dir = Path(".").resolve() / "outputs"
  output_dir.mkdir(exist_ok=True)
  # optimizers
-  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)
+  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 Tensor.train():
-    for epoch in (t := trange(epochs)):
-      loss_g, loss_d = 0.0, 0.0
-      for _ in range(n_steps):
-        data_real = make_batch(X_train)
-        for step in range(k):  # Try with k = 5 or 7.
-          noise = Tensor.randn(batch_size, 128)
-          data_fake = generator.forward(noise).detach()
-          loss_d += train_discriminator(optim_d, data_real, data_fake)
+  Tensor.training = True
+  for epoch in (t := trange(epochs)):
+    loss_g, loss_d = 0.0, 0.0
+    for _ in range(n_steps):
+      data_real = make_batch(images_real)
+      for step in range(k):  # Try with k = 5 or 7.
        noise = Tensor.randn(batch_size, 128)
-        data_fake = generator.forward(noise)
-        loss_g += train_generator(optim_g, data_fake)
-      if (epoch + 1) % sample_interval == 0:
-        fake_images = generator.forward(ds_noise).detach().numpy()
-        fake_images = (fake_images.reshape(-1, 1, 28, 28) + 1) / 2  # 0 - 1 range.
-        save_image(make_grid(torch.tensor(fake_images)), output_dir / f"image_{epoch+1}.jpg")
-      t.set_description(f"Generator loss: {loss_g/n_steps}, Discriminator loss: {loss_d/n_steps}")
+        data_fake = generator.forward(noise).detach()
+        loss_d += train_discriminator(optim_d, data_real, data_fake)
+      noise = Tensor.randn(batch_size, 128)
+      data_fake = generator.forward(noise)
+      loss_g += train_generator(optim_g, data_fake)
+    if (epoch + 1) % sample_interval == 0:
+      fake_images = generator.forward(ds_noise).detach().numpy()
+      fake_images = (fake_images.reshape(-1, 1, 28, 28) + 1) / 2  # 0 - 1 range.
+      save_image(make_grid(torch.tensor(fake_images)), output_dir / f"image_{epoch+1}.jpg")
+    t.set_description(f"Generator loss: {loss_g/n_steps}, Discriminator loss: {loss_d/n_steps}")
  print("Training Completed!")
@@ -31,7 +31,7 @@ def compile(onnx_file):
  for i in range(3):
    GlobalCounters.reset()
    print(f"run {i}")
-    with Context(DEBUG=max(DEBUG.value, 2 if i == 2 else 1), OPENPILOT_HACKS=1):
+    with Context(DEBUG=max(DEBUG.value, 2 if i == 2 else 1)):
      ret = run_onnx_jit(**inputs).numpy()
    # copy i == 1 so use of JITBEAM is okay
    if i == 1: test_val = np.copy(ret)
@@ -1,16 +0,0 @@
-import sys, pickle
-from extra.bench_log import WallTimeEvent, BenchEvent
-from tinygrad.helpers import getenv
-
-PKL = sys.argv[1] if len(sys.argv) > 1 else "/tmp/openpilot.pkl"
-
-load_times = []
-
-for _ in range(10):
-  with WallTimeEvent(BenchEvent.STEP) as wte: pickle.load(open(PKL, 'rb'))
-  load_times.append(wte.time)
-  print(f"pickle load: {wte.time:6.2f} s")
-
-if (assert_time:=getenv("ASSERT_MIN_LOAD_TIME")):
-  min_time = min(load_times)
-  assert min_time < assert_time, f"Speed regression, expected min load time of < {assert_time} s but took: {min_time} s"
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	ca906371fa	fix assign copy	2026-02-18 12:14:18 +08:00
geohot	fdb0b48ce9	assign after copy shouldn't contig	2026-02-18 12:03:20 +08:00