feat: tag 0.10.3 (#10310 )

small changes from O(1) multi [pr] (#10309 )
hotfix: filter_too_much (#10308 )
2026-08-24 23:06:07 +00:00 · 2025-05-14 15:45:13 -07:00 · 2025-05-14 15:34:07 -07:00 · 2025-05-14 15:31:51 -07:00 · 2025-05-14 15:16:40 -07:00 · 2025-05-14 14:50:37 -07:00
453 changed files with 1356046 additions and 249159 deletions
@@ -44,6 +44,9 @@ runs:
      uses: actions/setup-python@v5
      with:
        python-version: ${{ inputs.python-version }}
+    - name: Upgrade pip
+      shell: bash
+      run: python -m pip install --upgrade pip

    # **** Caching packages ****
    # TODO: key should include input.deps, but it can't since it can't contain commas
@@ -121,9 +124,8 @@ runs:
        echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
        sudo apt update || true
        sudo apt install --no-install-recommends --allow-unauthenticated -y hsa-rocr comgr hsa-rocr-dev liburing-dev libc6-dev
-        curl -s https://api.github.com/repos/Qazalin/remu/releases/latest | \
-        jq -r '.assets[] | select(.name == "libremu.so").browser_download_url' | \
-        sudo xargs curl -L -o /usr/local/lib/libremu.so
+        cargo build --release --manifest-path ./extra/remu/Cargo.toml
+        sudo ln -sf ${{ github.workspace }}/extra/remu/target/release/libremu.so /usr/local/lib/libremu.so
        sudo tee --append /etc/ld.so.conf.d/rocm.conf <<'EOF'
          /opt/rocm/lib
          /opt/rocm/lib64
@@ -137,9 +139,7 @@ runs:
        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 -L -o /usr/local/lib/libamd_comgr.dylib
-        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/Qazalin/remu/releases/latest | \
-          jq -r '.assets[] | select(.name == "libremu.dylib").browser_download_url' | \
-          sudo xargs curl -L -o /usr/local/lib/libremu.dylib
+        cargo build --release --manifest-path ./extra/remu/Cargo.toml

    # **** CUDA ****

@@ -175,7 +175,7 @@ runs:
        git checkout b16039dc940dc6bc4ea0a98380495769ff35ed99
        mkdir build
        cd build
-        cmake .. -Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF
+        cmake .. -Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5
        ninja
    - name: Install gpuocelot
      if: inputs.cuda == 'true'
@@ -202,8 +202,8 @@ runs:
      if: inputs.webgpu == 'true' && runner.os == 'macOS'
      shell: bash
      run: |
-        sudo mkdir -p /usr/local/lib
-        sudo curl -L https://github.com/wpmed92/pydawn/releases/download/v0.1.6/libwebgpu_dawn.dylib -o /usr/local/lib/libwebgpu_dawn.dylib
+        brew tap wpmed92/dawn
+        brew install dawn

    # **** LLVM ****

@@ -212,11 +212,13 @@ runs:
      shell: bash
      run: |
        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
+        wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
+        echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-19 main" | sudo tee /etc/apt/sources.list.d/llvm.list
        sudo apt update -y || true
-        sudo apt install -y --no-install-recommends llvm-dev
+        sudo apt install -y --no-install-recommends libllvm19 clang-19 lld-19

    - name: Install LLVM (macOS)
      if: inputs.llvm == 'true' && runner.os == 'macOS'
      shell: bash
      run: |
-        brew install llvm
+        brew install llvm@19
@@ -24,6 +24,9 @@ on:
 jobs:
  testmacbenchmark:
    name: Mac Benchmark
+    env:
+      # since sudo is required for usbgpu on macos, move the cache to a new location, as some of the files are owned by root
+      PYTHONPYCACHEPREFIX: /tmp/tiny_python_pycache
    runs-on: [self-hosted, macOS]
    timeout-minutes: 20
    defaults:
@@ -54,24 +57,25 @@ jobs:
      run: JIT=1 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
    - name: Run Stable Diffusion v2
      run: JIT=1 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
+    # process replay can't capture this, the graph is too large
    - name: Run SDXL
-      run: JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+      run: CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run model inference benchmark
      run: METAL=1 python3.11 test/external/external_model_benchmark.py
    - name: Test speed vs torch
      run: BIG=2 MPS=1 python3.11 test/test_speed_v_torch.py | tee torch_speed.txt
    - name: Test tensor cores
-      run: METAL=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
+      run: METAL=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Test AMX tensor cores
      run: |
-        DEBUG=2 CPU=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
-        DEBUG=2 LLVM=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
+        DEBUG=2 CPU=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        DEBUG=2 LLVM=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Run Tensor Core GEMM (float)
-      run: DEBUG=2 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
+      run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
    - name: Run Tensor Core GEMM (half)
-      run: DEBUG=2 HALF=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_half.txt
+      run: DEBUG=2 SHOULD_USE_TC=1 HALF=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_half.txt
    - name: Run Tensor Core GEMM (bfloat16)
-      run: DEBUG=2 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
+      run: DEBUG=2 SHOULD_USE_TC=1 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
    - name: Fuzz Padded Tensor Core GEMM
      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
@@ -84,6 +88,10 @@ jobs:
      run: |
        python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize int8 | tee llama_int8.txt
        python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize nf4 | tee llama_nf4.txt
+    - name: Run quantized LLaMA3
+      run: |
+        python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize int8 | tee llama3_int8.txt
+        python3.11 examples/llama3.py --size 8B --temperature 0 --benchmark --quantize nf4 | tee llama3_nf4.txt
    #- name: Run LLaMA 7B on 4 (virtual) GPUs
    #  run: python3.11 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_four_gpu.txt
    - name: Run GPT2
@@ -94,6 +102,8 @@ jobs:
      run: HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
    - name: Run GPT2 w HALF/BEAM
      run: HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
+    - name: Run OLMoE
+      run: python3.11 examples/olmoe.py
    - name: Train MNIST
      run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py | tee beautiful_mnist.txt
    - name: Run 10 CIFAR training steps
@@ -104,6 +114,14 @@ jobs:
    #  run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
    - name: Run 10 CIFAR training steps w winograd
      run: JIT=1 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+    - name: UsbGPU boot time
+      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=. AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
+    - name: UsbGPU copy speeds
+      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=. AMD=1 AMD_IFACE=USB AMD_LLVM=1 NOLOCALS=0 IMAGE=0 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (Mac)
@@ -115,6 +133,8 @@ jobs:
          llama_beam.txt
          llama_int8.txt
          llama_nf4.txt
+          llama3_int8.txt
+          llama3_nf4.txt
          llama_four_gpu.txt
          gpt2_unjitted.txt
          gpt2_jitted.txt
@@ -174,17 +194,17 @@ jobs:
      run: NV=1 python test/external/external_benchmark_multitensor_allreduce.py
    - name: Test tensor cores
      run: |
-        NV=1 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
-        PTX=1 ALLOW_TF32=1 NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
+        NV=1 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        PTX=1 ALLOW_TF32=1 NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Run Tensor Core GEMM (CUDA)
      run: |
-        CUDA=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
-        CUDA=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
-        CUDA=1 ALLOW_TF32=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_tf32.txt
+        CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
+        CUDA=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
+        CUDA=1 SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_tf32.txt
    - name: Run Tensor Core GEMM (PTX)
-      run: NV=1 PTX=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
+      run: NV=1 PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
    - name: Run Tensor Core GEMM (NV)
-      run: NV=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_nv.txt
+      run: NV=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_nv.txt
    - name: Test NV=1
      run: DEBUG=2 NV=1 python -m pytest -rA test/test_tiny.py
    - name: Test CUDA=1
@@ -192,7 +212,7 @@ jobs:
    - name: Run Stable Diffusion
      run: NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
    - name: Run SDXL
-      run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+      run: CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run LLaMA
      run: |
        NV=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
@@ -303,7 +323,7 @@ jobs:
      run: NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
    - name: Run 10 MLPerf Bert training steps (6 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+      run: NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (NVIDIA Training)
@@ -369,9 +389,12 @@ jobs:
    - name: Test speed vs theoretical
      run: AMD=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
    - name: Test tensor cores
-      run: AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
+      run: |
+        AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+        AMD=1 AMD_LLVM=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+        AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
    - name: Run Tensor Core GEMM (AMD)
-      run: AMD=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
+      run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
    - name: Test AMD=1
      run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
    - name: Test HIP=1
@@ -380,7 +403,7 @@ jobs:
    #- name: Fuzz Padded Tensor Core GEMM
    #  run: HSA=1 M_START=12 M_STOP=20 M_STEP=1 N_START=12 N_STOP=20 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 DEBUG=2 python3 ./extra/gemm/fuzz_matmul.py
    - name: Remove amdgpu
-      run: sleep 5 && sudo rmmod amdgpu # sleep a bit to let the driver unload the prev pid.
+      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 AMD=1 AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Test AM warm start time
@@ -388,7 +411,7 @@ jobs:
    - name: Run Stable Diffusion
      run: AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
    - name: Run SDXL
-      run: AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+      run: CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run LLaMA 7B
      run: |
        AMD=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
@@ -447,7 +470,7 @@ jobs:
  testmoreamdbenchmark:
    name: tinybox red Training Benchmark
    runs-on: [self-hosted, Linux, tinybox]
-    timeout-minutes: 20
+    timeout-minutes: 30
    defaults:
      run:
        shell: bash -o pipefail {0}
@@ -455,8 +478,8 @@ jobs:
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
-    - name: Insert amdgpu
-      run: sudo modprobe amdgpu
+    - name: Remove amdgpu
+      run: sudo rmmod amdgpu || true
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -474,10 +497,6 @@ 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: setup perflevel
-      run: |
-        examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/setup.sh
-        rocm-smi
    - name: Train MNIST
      run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
    - name: Run 10 CIFAR training steps
@@ -500,7 +519,7 @@ jobs:
      run: AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
    - name: Run 10 MLPerf Bert training steps (6 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+      run: AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (AMD Training)
@@ -29,9 +29,9 @@ jobs:
    - name: External Benchmark Schedule
      run: PYTHONPATH="." python3 test/external/external_benchmark_schedule.py
    - name: Speed Test
-      run: LLVM=1 LLVMOPT=1 python3 test/test_speed_v_torch.py
+      run: LLVM=1 python3 test/test_speed_v_torch.py
    - name: Speed Test (BEAM=2)
-      run: BEAM=2 LLVM=1 LLVMOPT=1 python3 test/test_speed_v_torch.py
+      run: BEAM=2 LLVM=1 python3 test/test_speed_v_torch.py

  docs:
    name: Docs
@@ -88,7 +88,7 @@ jobs:
  autogen:
    name: Autogen
    runs-on: ubuntu-22.04
-    timeout-minutes: 10
+    timeout-minutes: 15
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -116,14 +116,20 @@ jobs:
    - name: Verify AMD autogen
      run: |
        cp tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
+        cp tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
        cp tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
        cp tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
+        cp tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
        ./autogen_stubs.sh hsa
+        ./autogen_stubs.sh kfd
        ./autogen_stubs.sh comgr
        ./autogen_stubs.sh amd
+        ./autogen_stubs.sh sqtt
        diff /tmp/hsa.py.bak tinygrad/runtime/autogen/hsa.py
+        diff /tmp/kfd.py.bak tinygrad/runtime/autogen/kfd.py
        diff /tmp/comgr.py.bak tinygrad/runtime/autogen/comgr.py
        diff /tmp/amd_gpu.py.bak tinygrad/runtime/autogen/amd_gpu.py
+        diff /tmp/sqtt.py.bak tinygrad/runtime/autogen/sqtt.py
    - name: Verify Linux autogen
      run: |
        cp tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
@@ -143,10 +149,75 @@ jobs:
        ./autogen_stubs.sh llvm
        diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py

+  torchbackend:
+    name: Torch Backend Tests
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    env:
+      IGNORE_OOB: 0
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: torch-backend-pillow-torchvision-et-pt
+        deps: testing_minimal
+        pydeps: "pillow torchvision expecttest"
+        llvm: 'true'
+    - name: Install ninja
+      run: |
+        sudo apt update || true
+        sudo apt install -y --no-install-recommends ninja-build
+    - name: Lint with ruff
+      run: |
+        pip3 install --upgrade --force-reinstall ruff==0.11.0
+        python3 -m ruff check extra/torch_backend/backend.py
+    - name: Test one op
+      run: PYTHONPATH=. FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
+    - name: Test ResNet-18
+      run: PYTHONPATH=. DEBUG=2 python3 extra/torch_backend/example.py
+    - name: My (custom) tests
+      run: PYTHONPATH=. python3 extra/torch_backend/test.py
+    - name: Test one op in torch tests
+      run: PYTHONPATH=. DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
+    - name: Test Ops with TINY_BACKEND
+      run: PYTHONPATH=. LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
+    - name: Test in-place operations on views
+      run: PYTHONPATH=. TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
+    - name: Test multi-gpu
+      run: PYTHONPATH=. LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
+
+  torchbackendmore:
+    name: Torch Backend Tests More
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    env:
+      IGNORE_OOB: 0
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: torch-backend-pillow-torchvision-et-pt
+        deps: testing_minimal
+        llvm: 'true'
+    - name: Install ninja
+      run: |
+        sudo apt update || true
+        sudo apt install -y --no-install-recommends ninja-build
+    - name: Test beautiful_mnist in torch with TINY_BACKEND
+      run: SPLIT_REDUCEOP=0 FUSE_ARANGE=1 PYTHONPATH=. LLVM=1 TARGET_EVAL_ACC_PCT=96.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+    - name: Test some torch tests (expect failure)
+      run: PYTHONPATH=. python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
+
  tc:
    name: Tensor Core tests
    runs-on: ubuntu-latest
    timeout-minutes: 10
+    env:
+      IGNORE_OOB: 0
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -162,7 +233,8 @@ jobs:
    - name: Test emulated METAL tensor cores
      run: |
        DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_big_gemm
-        PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
+        PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
+        PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Test emulated AMX tensor cores
      run: PYTHONPATH=. DEBUG=2 AMX=1 EMULATE_AMX=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
    - name: Test emulated AMD tensor cores
@@ -171,13 +243,28 @@ jobs:
        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 python3 ./extra/gemm/simple_matmul.py
        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+    - name: Test emulated AMD MFMA tensor cores
+      run: |
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_MFMA=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_MFMA=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_MFMA=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+    - name: Test emulated AMD RDNA4 tensor cores
+      run: |
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Test emulated CUDA tensor cores
      run: |
        DEBUG=2 EMULATE_CUDA=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
        DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
        DEBUG=2 EMULATE_CUDA_SM75=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
-        PYTHONPATH="." DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
+        PYTHONPATH="." DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
+        PYTHONPATH="." DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Test emulated INTEL OpenCL tensor cores
      run: DEBUG=2 EMULATE_INTEL=1 FORWARD_ONLY=1 PYTHON=1 HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
    - name: Full test tensor cores
@@ -189,11 +276,12 @@ jobs:
        PYTHONPATH=. DEBUG=2 AMX=1 EMULATE_AMX=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
    - name: Test tensor cores (TC=3)
      run: |
-        TC=3 DEBUG=3 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
-        TC=3 PYTHONPATH=. DEBUG=3 EMULATE_AMD=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        TC=3 DEBUG=3 EMULATE_CUDA=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
-        TC=3 PYTHONPATH=. DEBUG=3 EMULATE_INTEL=1 PYTHON=1 N=16 HALF=1 python3 ./extra/gemm/simple_matmul.py
-        TC=3 PYTHONPATH=. DEBUG=3 AMX=1 EMULATE_AMX=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
+        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_METAL=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
+        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_AMD=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
+        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_AMD_MFMA=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
+        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_CUDA=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
+        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_INTEL=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
+        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_AMX=1 AMX=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
    - name: Test device flop counts
      run: |
        PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
@@ -206,6 +294,8 @@ jobs:
    name: Python Backend
    runs-on: ubuntu-latest
    timeout-minutes: 10
+    env:
+      IGNORE_OOB: 0
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -217,13 +307,15 @@ jobs:
    - name: Test dtype with Python emulator
      run: DEBUG=1 PYTHONPATH=. PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
    - name: Test ops with Python emulator
-      run: DEBUG=2 PYTHON=1 python3 -m pytest -n=auto test/test_ops.py -k "not (test_split or test_simple_cumsum or test_cumsum or test_einsum or test_dot or test_dot_1d or test_big_gemm or test_broadcastdot or test_multidot or test_var_axis or test_std_axis or test_broadcast_full or test_broadcast_partial or test_simple_conv3d or test_dilated_conv_transpose2d or test_simple_conv_transpose3d or test_large_input_conv2d or test_max_pool2d or test_max_pool2d_simple or test_max_pool2d_bigger_stride or test_avg_pool2d or test_cat or test_scaled_product_attention or test_scaled_product_attention_causal or test_slice_fancy_indexing_dim_inject_none or test_slice_fancy_indexing_list_indices or test_slice_fancy_indexing_no_dim_collapse or test_slice_fancy_indexing_tuple_indices or test_slice_fancy_indexing_list_with_tensors or test_slice_fancy_indexing_dim_collapse_int or test_interpolate_bilinear or test_interpolate_bilinear_corners_aligned or test_scaled_dot_product_attention or test_cummax)" --durations=20
+      run: DEBUG=2 PYTHON=1 python3 -m pytest -n=auto test/test_ops.py -k "not (test_split or test_simple_cumsum or test_cumsum or test_einsum or test_dot or test_dot_1d or test_big_gemm or test_broadcastdot or test_multidot or test_var_axis or test_std_axis or test_broadcast_full or test_broadcast_partial or test_simple_conv3d or test_dilated_conv_transpose2d or test_simple_conv_transpose3d or test_large_input_conv2d or test_max_pool2d or test_max_pool2d_simple or test_max_pool2d_bigger_stride or test_avg_pool2d or test_cat or test_scaled_product_attention or test_scaled_product_attention_causal or test_slice_fancy_indexing_dim_inject_none or test_slice_fancy_indexing_list_indices or test_slice_fancy_indexing_no_dim_collapse or test_slice_fancy_indexing_tuple_indices or test_slice_fancy_indexing_list_with_tensors or test_slice_fancy_indexing_dim_collapse_int or test_interpolate_bilinear or test_interpolate_bilinear_corners_aligned or test_scaled_dot_product_attention or test_cummax or test_simple_cummax or test_logcumsumexp or test_sort or test_cumprod)" --durations=20
    - name: Test uops with Python emulator
      run: PYTHON=1 python3 -m pytest test/test_uops.py --durations=20
    - name: Test symbolic with Python emulator
      run: PYTHONPATH=. PYTHON=1 python3 test/test_symbolic_ops.py
    - name: test_linearizer_failures with Python emulator
      run: PYTHONPATH=. PYTHON=1 python3 -m pytest -rA test/test_linearizer_failures.py::TestLinearizerFailures::test_failure_1
+    - name: test_renderer_failures with Python emulator
+      run: PYTHONPATH=. PYTHON=1 python3 -m pytest -rA test/test_renderer_failures.py::TestRendererFailures

  linter:
    name: Linters
@@ -244,8 +336,9 @@ jobs:
      run: python -m pylint --disable=all -e W0311 -e C0303 --jobs=0 --indent-string='  ' --recursive=y .
    - name: Lint with ruff
      run: |
-        pip3 install --upgrade --force-reinstall ruff
+        pip3 install --upgrade --force-reinstall ruff==0.11.0
        python3 -m ruff check .
+        python3 -m ruff check examples/mlperf/model_train.py --ignore E501
    - name: Lint tinygrad with pylint
      run: python -m pylint tinygrad/
    - name: Run mypy
@@ -268,8 +361,15 @@ jobs:
      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' README.md > README.py &&  PYTHONPATH=. python README.py
    - name: Run unit tests
      run: PYTHONPATH="." python -m pytest -n=auto test/unit/
-    - name: Repo line count < 11300 lines
-      run: MAX_LINE_COUNT=11300 python sz.py
+    - name: Run targetted tests on NULL backend
+      run: PYTHONPATH="." NULL=1 python3 test/test_multitensor.py TestMultiTensor.test_data_parallel_resnet_train_step
+    # 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: PYTHONPATH="." python test/external/external_uop_gc.py
+    - name: Repo line count < 13000 lines
+      run: MAX_LINE_COUNT=13000 python sz.py

  fuzzing:
    name: Fuzzing
@@ -285,6 +385,8 @@ jobs:
        deps: testing_minimal
    - name: Fuzz Test symbolic
      run: python test/external/fuzz_symbolic.py
+    - name: Fuzz Test fast idiv
+      run: python test/external/fuzz_fast_idiv.py
    - name: Fuzz Test shapetracker
      run: |
        PYTHONPATH="." python test/external/fuzz_shapetracker.py
@@ -292,8 +394,10 @@ jobs:

  testgpuimage:
    name: 'GPU IMAGE Tests'
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
    timeout-minutes: 10
+    env:
+      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -316,8 +420,10 @@ jobs:

  testopenpilot:
    name: 'openpilot Compile Tests'
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
    timeout-minutes: 10
+    env:
+      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -327,13 +433,16 @@ jobs:
          key: openpilot-compile
          deps: testing
          opencl: 'true'
+          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          PYTHONPATH="." ALLOWED_KERNEL_COUNT=209 ALLOWED_READ_IMAGE=2105 ALLOWED_GATED_READ_IMAGE=29 FLOAT16=0 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
+          PYTHONPATH="." ALLOWED_KERNEL_COUNT=209 ALLOWED_READ_IMAGE=2138 ALLOWED_GATED_READ_IMAGE=29 FLOAT16=0 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot alt model correctness (float32)
        run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot fastvits model correctness (float32)
        run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+      - name: Test openpilot LLVM compile
+        run: PYTHONPATH="." LLVM=1 LLVMOPT=1 JIT=2 BEAM=0 IMAGE=0 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -341,6 +450,8 @@ jobs:
    name: 'ONNX+Optimization Tests'
    runs-on: ubuntu-22.04
    timeout-minutes: 20
+    env:
+      IGNORE_OOB: 0

    steps:
      - name: Checkout Code
@@ -359,11 +470,15 @@ jobs:
        run: CPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
      - name: Test ONNX (LLVM)
        run: LLVM=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
-      - name: Run CLOUD=1 Test
+      - name: Test Additional ONNX Ops (CPU)
+        run: CPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_ops.py
+      - name: Test Quantize ONNX
+        run: CPU=1 PYTHONPATH=. python3 test/test_quantize_onnx.py
+      - name: Run REMOTE=1 Test
        run: |
-          CLOUDDEV=CPU CLOUD=1 python3 test/test_tiny.py
-          CLOUDDEV=GPU CLOUD=1 python3 test/test_tiny.py
-          CLOUDDEV=GPU IMAGE=2 CLOUD=1 python3 test/test_tiny.py
+          REMOTEDEV=CPU REMOTE=1 python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py
+          REMOTEDEV=GPU REMOTE=1 python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py
+          REMOTEDEV=GPU IMAGE=2 REMOTE=1 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
      - name: Test Optimization Helpers
        run: PYTHONPATH="." DEBUG=1 python3 extra/optimization/test_helpers.py
      - name: Test Action Space
@@ -381,6 +496,8 @@ jobs:
    name: Models (llvm+cpu+gpu)
    runs-on: ubuntu-22.04
    timeout-minutes: 10
+    env:
+      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -403,7 +520,9 @@ jobs:
  testdsp:
    name: Linux (DSP)
    runs-on: ubuntu-24.04
-    timeout-minutes: 10
+    timeout-minutes: 15
+    env:
+      IGNORE_OOB: 0
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -412,7 +531,7 @@ jobs:
      with:
        key: dsp-minimal
        deps: testing_minimal
-        pydeps: "onnx==1.16.0 onnxruntime"
+        pydeps: "onnx==1.17.0 onnxruntime pillow"
        llvm: "true"
    - name: Set up Docker Buildx
      uses: docker/setup-buildx-action@v3
@@ -427,10 +546,14 @@ jobs:
        cache-to: type=gha,mode=min
    - name: Run test_tiny on DSP
      run: DEBUG=2 DSP=1 python test/test_tiny.py
+    - name: Test transcendentals
+      run: CC=clang-19 PYTHONPATH="." DEBUG=2 DSP=1 python test/test_transcendental.py TestTranscendentalVectorized
    - name: Test quantize onnx
      run: PYTHONPATH="." DEBUG=2 DSP=1 python3 test/test_quantize_onnx.py
    - name: Test LLVM=1 DEVECTORIZE=0
      run: LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
+    - name: Test LLVM=1 DEVECTORIZE=0 for model
+      run: PYTHONPATH="." LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
    - name: Test CPU=1 DEVECTORIZE=0
      run: CPU=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"

@@ -460,11 +583,50 @@ jobs:
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

-  tests:
+  testamd:
    strategy:
      fail-fast: false
      matrix:
-        backend: [llvm, cpu, gpu, ptx, amd, nv] #, triton]
+        backend: [amd, amdllvm]
+
+    name: Linux (${{ matrix.backend }})
+    runs-on: ubuntu-22.04
+    timeout-minutes: 20
+    env:
+      IGNORE_OOB: 0
+      AMD: 1
+      MOCKGPU: 1
+      FORWARD_ONLY: 1
+      AMD_LLVM: ${{ matrix.backend == 'amdllvm' && '1' || matrix.backend != 'amdllvm' && '0' }}
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: ${{ matrix.backend }}-minimal
+          deps: testing_minimal
+          amd: 'true'
+          llvm: ${{ matrix.backend == 'amdllvm' && 'true' }}
+      - name: Check Device.DEFAULT and print some source
+        run: |
+          PYTHONPATH=${{ github.workspace }} python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['AMD'], Device.DEFAULT"
+          DEBUG=5 PYTHONPATH=${{ github.workspace }} FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+      - name: Run LLVM test
+        if: matrix.backend=='amdllvm'
+        run: PYTHONPATH="." python test/test_amd_llvm.py
+      - name: Run pytest (amd)
+        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/imported/test_indexing.py test/test_hcq.py test/external/external_test_am.py --durations=20
+      - name: Run TRANSCENDENTAL math
+        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
+      - name: Run process replay tests
+        uses: ./.github/actions/process-replay
+
+  testnvidia:
+    strategy:
+      fail-fast: false
+      matrix:
+        backend: [ptx, nv]

    name: Linux (${{ matrix.backend }})
    runs-on: ubuntu-22.04
@@ -477,25 +639,49 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: ${{ matrix.backend }}-minimal
-          deps: testing_minimal${{matrix.backend=='ptx'&&',cuda'||matrix.backend=='triton'&&',triton'||''}}
-          opencl: ${{ matrix.backend == 'gpu' && 'true' }}
-          amd: ${{ matrix.backend == 'amd' && 'true' }}
-          cuda: ${{ (matrix.backend == 'ptx' || matrix.backend == 'triton' || matrix.backend == 'nv') && 'true' }}
+          deps: testing_minimal
+          cuda: 'true'
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'gpu' && 'GPU=1' || matrix.backend == 'PTX' && 'FORWARD_ONLY=1\nJIT=1\nOPT=2\nCUDA=1\nPTX=1\nMOCKGPU=1' || matrix.backend == 'triton' && 'FORWARD_ONLY=1\nJIT=1\nOPT=2\nNV=1\nMOCKGPU=1\nTRITON=1\nTRITON_PTXAS_PATH=/usr/bin/ptxas' || matrix.backend == 'amd' && 'AMD=1\nMOCKGPU=1\nFORWARD_ONLY=1' || matrix.backend == 'nv' && 'NV=1\nMOCKGPU=1\nFORWARD_ONLY=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'PTX' && 'FORWARD_ONLY=1\nJIT=1\nOPT=2\nCUDA=1\nPTX=1\nMOCKGPU=1' || matrix.backend == 'nv' && 'NV=1\nMOCKGPU=1\nFORWARD_ONLY=1' }}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
-          PYTHONPATH=${{ github.workspace }} python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['LLVM','CPU','CUDA','GPU','AMD','NV'], Device.DEFAULT"
+          PYTHONPATH=${{ github.workspace }} python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CUDA','NV'], Device.DEFAULT"
          DEBUG=5 PYTHONPATH=${{ github.workspace }} FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
-      - name: Run pytest (not cuda or amd)
-        if: matrix.backend!='ptx' && matrix.backend!='triton' && matrix.backend != 'amd' && matrix.backend != 'nv'
-        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
      - name: Run pytest (cuda)
-        if: matrix.backend=='ptx'||matrix.backend=='triton'||matrix.backend=='nv'
        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore test/test_gc.py --durations=20
-      - name: Run pytest (amd)
-        if: matrix.backend=='amd'
-        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/imported/test_indexing.py test/test_hcq.py test/external/external_test_am.py --durations=20
+      - name: Run process replay tests
+        uses: ./.github/actions/process-replay
+
+  tests:
+    strategy:
+      fail-fast: false
+      matrix:
+        backend: [llvm, cpu, gpu]
+
+    name: Linux (${{ matrix.backend }})
+    runs-on: ubuntu-22.04
+    timeout-minutes: 20
+    env:
+      IGNORE_OOB: 0
+
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: ${{ matrix.backend }}-minimal
+          deps: testing_minimal
+          opencl: ${{ matrix.backend == 'gpu' && 'true' }}
+          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
+      - name: Set env
+        run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'gpu' && 'GPU=1' }}" >> $GITHUB_ENV
+      - name: Check Device.DEFAULT and print some source
+        run: |
+          PYTHONPATH=${{ github.workspace }} python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['LLVM','CPU','GPU'], Device.DEFAULT"
+          DEBUG=5 PYTHONPATH=${{ github.workspace }} FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+      - name: Run pytest (not cuda)
+        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
      - name: Run TRANSCENDENTAL math
        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
      - name: Run process replay tests
@@ -506,7 +692,9 @@ jobs:
  testmetal2:
    name: MacOS (unit)
    runs-on: macos-14
-    timeout-minutes: 10
+    timeout-minutes: 20
+    env:
+      IGNORE_OOB: 0

    steps:
    - name: Checkout Code
@@ -519,12 +707,13 @@ jobs:
        python-version: '3.11'
        amd: 'true'
        cuda: 'true'
+        llvm: 'true'
    - name: Run real world test
-      run: JIT=2 METAL=1 python -m pytest -n=auto test/models/test_real_world.py --durations=20
+      run: METAL=1 python -m pytest -n=auto test/models/test_real_world.py --durations=20
    - name: Test models (Metal)
-      run: JIT=2 METAL=1 python -m pytest -n=auto test/models --durations=20
+      run: METAL=1 python -m pytest test/models -v --durations=20
    - name: Run ONNX
-      run: JIT=2 METAL=1 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: TC=2 METAL=1 DEBUG=3 python test/test_ops.py TestOps.test_gemm
    - name: Test tensor core ops (real)
@@ -535,8 +724,6 @@ jobs:
      run: PYTHONPATH="." METAL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
    - name: Fuzz Test linearizer
      run: PYTHONPATH="." METAL=1 DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
-    # - name: Fuzz Test models schedule
-    # run: FUZZ_SCHEDULE=1 FUZZ_SCHEDULE_MAX_PATHS=5 python -m pytest test/models/test_train.py test/models/test_end2end.py
    - name: Run TRANSCENDENTAL math
      run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
    - name: Run pytest (amd)
@@ -546,6 +733,14 @@ jobs:
        FORWARD_ONLY: 1
      run: |
        python3 -m pytest -n=auto test/test_hcq.py test/test_tiny.py --durations=20
+    - name: Run pytest (amd with llvm backend)
+      env:
+        MOCKGPU: 1
+        AMD: 1
+        AMD_LLVM: 1
+        FORWARD_ONLY: 1
+      run: |
+        python -m pytest -n=auto test/test_hcq.py test/test_tiny.py test/test_amd_llvm.py --durations=20
    - name: Run pytest (ptx)
      env:
        MOCKGPU: 1
@@ -568,7 +763,10 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: osx-webgpu
+        deps: testing_minimal
        webgpu: 'true'
+    - name: Test infinity math in WGSL
+      run: WEBGPU=1 python -m pytest -n=auto test/test_renderer_failures.py::TestWGSLFailures::test_multiply_infinity --durations=20
    - name: Build WEBGPU Efficientnet
      run: WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m examples.compile_efficientnet
    - name: Clean npm cache
@@ -578,6 +776,30 @@ jobs:
    - name: Run WEBGPU Efficientnet
      run: node test/web/test_webgpu.js

+  osxremote:
+    name: MacOS (remote)
+    runs-on: macos-15
+    timeout-minutes: 10
+    env:
+      REMOTE: 1
+      REMOTEDEV: METAL
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: macos-remote
+          deps: testing_minimal
+      - name: Check Device.DEFAULT and print some source
+        run: |
+          python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
+          python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'METAL', Device.default.properties.real_device"
+          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
+      - name: Run REMOTE=1 Test
+        run: |
+          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py
+
  osxtests:
    strategy:
      fail-fast: false
@@ -585,7 +807,9 @@ jobs:
        backend: [metal, llvm, cpu]
    name: MacOS (${{ matrix.backend }})
    runs-on: macos-15
-    timeout-minutes: 10
+    timeout-minutes: 20
+    env:
+      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -596,7 +820,7 @@ jobs:
          deps: testing_minimal
          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'metal' && 'METAL=1\nJIT=2'}}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'metal' && 'METAL=1'}}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
          python -c "from tinygrad import Device; assert Device.DEFAULT == '${{ matrix.backend }}'.upper(), Device.DEFAULT"
@@ -605,6 +829,9 @@ jobs:
        run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
      - name: Run process replay tests
        uses: ./.github/actions/process-replay
+      - name: Run macOS-specific unit test
+        if: matrix.backend == 'cpu'
+        run: python3 -m pytest test/unit/test_disk_tensor.py::TestDiskTensor::test_copy_to_cpu_not_truncated

 # ****** Windows Tests ******

@@ -616,7 +843,9 @@ jobs:

    name: Windows (${{ matrix.backend }})
    runs-on: windows-latest
-    timeout-minutes: 10
+    timeout-minutes: 15
+    env:
+      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -10,6 +10,7 @@ notebooks
 *.so
 *.txt
 build
+!examples/tinychat/assets/cdn.jsdelivr.net/npm/[email protected]/build/
 /dist
 *.egg-info
 /env
@@ -33,6 +34,8 @@ extra/datasets/open-images-v6-mlperf
 extra/datasets/kits/
 extra/datasets/COCO/
 extra/datasets/audio*
+extra/huggingface_onnx/models/*
+extra/huggingface_onnx/*.yaml
 extra/weights
 venv
 examples/**/net.*[js,json]
@@ -55,5 +58,7 @@ weights
 comgr_*
 *.pkl
 site/
-master_schedule.py
 profile_stats
+*.log
+target
+.mypy_cache
@@ -35,7 +35,7 @@ def _try_dlopen_$name():
  for candidate in PATHS_TO_TRY:
    try: return ctypes.CDLL(candidate)
    except OSError: pass
-  raise RuntimeError("library $name not found")
+  return None
 EOF
 }

@@ -78,11 +78,11 @@ generate_kfd() {
  clang2py /usr/include/linux/kfd_ioctl.h -o $BASE/kfd.py -k cdefstum

  fixup $BASE/kfd.py
-  sed -i "s\import ctypes\import ctypes, os\g" $BASE/kfd.py
-  sed -i "s\import fcntl, functools\import functools" $BASE/kfd.py
-  sed -i "s\import ctypes,os\a from tinygrad.runtime.support import HWInterface\g" $BASE/kfd.py
-  sed -i "s\def _do_ioctl(__idir, __base, __nr, __user_struct, __fd, **kwargs):\def _do_ioctl(__idir, __base, __nr, __user_struct, __fd:HWInterface, **kwargs):\g" $BASE/kfd.py
-  sed -i "s\fcntl.ioctl(__fd, (__idir<<30)\__fd.ioctl((__idir<<30)\g" $BASE/kfd.py
+  sed -i "s/import ctypes/import ctypes, os/g" $BASE/kfd.py
+  sed -i "s/import fcntl, functools/import functools/g" $BASE/kfd.py
+  sed -i "/import functools/a from tinygrad.runtime.support.hcq import FileIOInterface" $BASE/kfd.py
+  sed -i "s/def _do_ioctl(__idir, __base, __nr, __user_struct, __fd, \*\*kwargs):/def _do_ioctl(__idir, __base, __nr, __user_struct, __fd:FileIOInterface, \*\*kwargs):/g" $BASE/kfd.py
+  sed -i "s/fcntl.ioctl(__fd, (__idir<<30)/__fd.ioctl((__idir<<30)/g" $BASE/kfd.py
  python3 -c "import tinygrad.runtime.autogen.kfd"
 }

@@ -171,6 +171,7 @@ generate_amd() {
    extra/hip_gpu_driver/sdma_v6_0_0_pkt_open.h \
    extra/hip_gpu_driver/gc_11_0_0_offset.h \
    extra/hip_gpu_driver/gc_10_3_0_offset.h \
+    extra/hip_gpu_driver/sienna_cichlid_ip_offset.h \
    --clang-args="-I/opt/rocm/include -x c++" \
    -o $BASE/amd_gpu.py

@@ -286,83 +287,260 @@ generate_vfio() {
  fixup $BASE/vfio.py
  sed -i "s\import ctypes\import ctypes, os\g" $BASE/vfio.py
  sed -i "s\import fcntl, functools\import functools" $BASE/vfio.py
-  sed -i "s\import ctypes,os\a from tinygrad.runtime.support import HWInterface\g" $BASE/vfio.py
+  sed -i "s\import ctypes,os\a from tinygrad.runtime.support import FileIOInterface\g" $BASE/vfio.py
  sed -i "s\fcntl.ioctl(__fd, (__idir<<30)\return __fd.ioctl((__idir<<30)\g" $BASE/vfio.py
 }

 generate_am() {
+  AMKERN_COMMIT_HASH=ceb12c04e2b5b53ec0779362831f5ee40c4921e4
+  AMKERN_SRC=/tmp/ROCK-Kernel-Driver-$AMKERN_COMMIT_HASH
+  if [ ! -d "$AMKERN_SRC" ]; then
+    git clone https://github.com/ROCm/ROCK-Kernel-Driver $AMKERN_SRC --depth 1
+  fi
+  AMKERN_AMD=$AMKERN_SRC/drivers/gpu/drm/amd/
+  AMKERN_INC=$AMKERN_AMD/include/
+
  clang2py -k cdefstum \
    extra/amdpci/headers/v11_structs.h \
+    extra/amdpci/headers/v12_structs.h \
    extra/amdpci/headers/amdgpu_vm.h \
    extra/amdpci/headers/discovery.h \
    extra/amdpci/headers/amdgpu_ucode.h \
-    extra/amdpci/headers/soc21_enum.h \
    extra/amdpci/headers/psp_gfx_if.h \
    extra/amdpci/headers/amdgpu_psp.h \
    extra/amdpci/headers/amdgpu_irq.h \
    extra/amdpci/headers/amdgpu_doorbell.h \
-    extra/amdpci/headers/soc15_ih_clientid.h \
+    $AMKERN_INC/soc15_ih_clientid.h \
+    --clang-args="-include stdint.h" \
    -o $BASE/am/am.py
  fixup $BASE/am/am.py
+  sed -i "s\(int64_t)\ \g" $BASE/am/am.py
+  sed -i "s\AMDGPU_PTE_MTYPE_VG10(2)\AMDGPU_PTE_MTYPE_VG10(0, 2)\g" $BASE/am/am.py # incorrect parsing (TODO: remove when clang2py is gone).

  clang2py -k cdefstum \
-    extra/amdpci/headers/mp_13_0_0_offset.h \
-    extra/amdpci/headers/mp_13_0_0_sh_mask.h \
+    $AMKERN_AMD/amdkfd/kfd_pm4_headers_ai.h \
+    $AMKERN_AMD/amdgpu/soc15d.h \
+    -o $BASE/am/pm4_soc15.py
+  fixup $BASE/am/pm4_soc15.py
+
+  clang2py -k cdefstum \
+    $AMKERN_AMD/amdkfd/kfd_pm4_headers_ai.h \
+    $AMKERN_AMD/amdgpu/nvd.h \
+    -o $BASE/am/pm4_nv.py
+  fixup $BASE/am/pm4_nv.py
+
+  clang2py -k cdefstum \
+    $AMKERN_INC/vega10_enum.h \
+    -o $BASE/am/vega10.py
+  fixup $BASE/am/vega10.py
+
+  clang2py -k cdefstum \
+    $AMKERN_INC/navi10_enum.h \
+    -o $BASE/am/navi10.py
+  fixup $BASE/am/navi10.py
+
+  clang2py -k cdefstum \
+    $AMKERN_INC/soc21_enum.h \
+    -o $BASE/am/soc21.py
+  fixup $BASE/am/soc21.py
+
+  clang2py -k cdefstum \
+    $AMKERN_INC/soc24_enum.h \
+    -o $BASE/am/soc24.py
+  fixup $BASE/am/soc24.py
+
+  clang2py -k cdefstum \
+    $AMKERN_INC/asic_reg/mp/mp_13_0_0_offset.h \
+    $AMKERN_INC/asic_reg/mp/mp_13_0_0_sh_mask.h \
    -o $BASE/am/mp_13_0_0.py
  fixup $BASE/am/mp_13_0_0.py

+  # 14_0_3 reuses 14_0_2
  clang2py -k cdefstum \
-    extra/amdpci/headers/mp_11_0_offset.h \
-    extra/amdpci/headers/mp_11_0_sh_mask.h \
+    $AMKERN_INC/asic_reg/mp/mp_14_0_2_offset.h \
+    $AMKERN_INC/asic_reg/mp/mp_14_0_2_sh_mask.h \
+    -o $BASE/am/mp_14_0_3.py
+  fixup $BASE/am/mp_14_0_3.py
+
+  clang2py -k cdefstum \
+    $AMKERN_INC/asic_reg/mp/mp_11_0_offset.h \
+    $AMKERN_INC/asic_reg/mp/mp_11_0_sh_mask.h \
    -o $BASE/am/mp_11_0.py
  fixup $BASE/am/mp_11_0.py

  clang2py -k cdefstum \
-    extra/amdpci/headers/gc_11_0_0_offset.h \
-    extra/amdpci/headers/gc_11_0_0_sh_mask.h \
+    $AMKERN_INC/asic_reg/gc/gc_9_4_3_offset.h \
+    $AMKERN_INC/asic_reg/gc/gc_9_4_3_sh_mask.h \
+    extra/amdpci/overlay/gc_9_4_3.h \
+    -o $BASE/am/gc_9_4_3.py
+  fixup $BASE/am/gc_9_4_3.py
+
+  clang2py -k cdefstum \
+    $AMKERN_INC/asic_reg/gc/gc_10_3_0_offset.h \
+    $AMKERN_INC/asic_reg/gc/gc_10_3_0_sh_mask.h \
+    -o $BASE/am/gc_10_3_0.py
+  fixup $BASE/am/gc_10_3_0.py
+
+  clang2py -k cdefstum \
+    $AMKERN_INC/asic_reg/gc/gc_11_0_0_offset.h \
+    $AMKERN_INC/asic_reg/gc/gc_11_0_0_sh_mask.h \
    -o $BASE/am/gc_11_0_0.py
  fixup $BASE/am/gc_11_0_0.py

  clang2py -k cdefstum \
-    extra/amdpci/headers/mmhub_3_0_0_offset.h \
-    extra/amdpci/headers/mmhub_3_0_0_sh_mask.h \
+    $AMKERN_INC/asic_reg/gc/gc_12_0_0_offset.h \
+    $AMKERN_INC/asic_reg/gc/gc_12_0_0_sh_mask.h \
+    -o $BASE/am/gc_12_0_0.py
+  fixup $BASE/am/gc_12_0_0.py
+
+  clang2py -k cdefstum \
+    extra/hip_gpu_driver/sdma_registers.h \
+    $AMKERN_AMD/amdgpu/vega10_sdma_pkt_open.h \
+    --clang-args="-I/opt/rocm/include -x c++" \
+    -o $BASE/am/sdma_4_0_0.py
+  fixup $BASE/am/sdma_4_0_0.py
+
+  clang2py -k cdefstum \
+    extra/hip_gpu_driver/sdma_registers.h \
+    $AMKERN_AMD/amdgpu/navi10_sdma_pkt_open.h \
+    --clang-args="-I/opt/rocm/include -x c++" \
+    -o $BASE/am/sdma_5_0_0.py
+  fixup $BASE/am/sdma_5_0_0.py
+
+  clang2py -k cdefstum \
+    extra/hip_gpu_driver/sdma_registers.h \
+    $AMKERN_AMD/amdgpu/sdma_v6_0_0_pkt_open.h \
+    --clang-args="-I/opt/rocm/include -x c++" \
+    -o $BASE/am/sdma_6_0_0.py
+  fixup $BASE/am/sdma_6_0_0.py
+
+  clang2py -k cdefstum \
+    $AMKERN_INC/asic_reg/mmhub/mmhub_3_0_0_offset.h \
+    $AMKERN_INC/asic_reg/mmhub/mmhub_3_0_0_sh_mask.h \
    -o $BASE/am/mmhub_3_0_0.py
  fixup $BASE/am/mmhub_3_0_0.py

  clang2py -k cdefstum \
-    extra/amdpci/headers/mmhub_3_0_2_offset.h \
-    extra/amdpci/headers/mmhub_3_0_2_sh_mask.h \
+    $AMKERN_INC/asic_reg/mmhub/mmhub_3_0_2_offset.h \
+    $AMKERN_INC/asic_reg/mmhub/mmhub_3_0_2_sh_mask.h \
    -o $BASE/am/mmhub_3_0_2.py
  fixup $BASE/am/mmhub_3_0_2.py

  clang2py -k cdefstum \
-    extra/amdpci/headers/nbio_4_3_0_offset.h \
-    extra/amdpci/headers/nbio_4_3_0_sh_mask.h \
+    $AMKERN_INC/asic_reg/nbio/nbio_2_3_offset.h \
+    $AMKERN_INC/asic_reg/nbio/nbio_2_3_sh_mask.h \
+    -o $BASE/am/nbio_2_3_0.py
+  fixup $BASE/am/nbio_2_3_0.py
+
+  clang2py -k cdefstum \
+    $AMKERN_INC/asic_reg/nbio/nbio_7_2_0_offset.h \
+    $AMKERN_INC/asic_reg/nbio/nbio_7_2_0_sh_mask.h \
+    -o $BASE/am/nbio_7_2_0.py
+  fixup $BASE/am/nbio_7_2_0.py
+
+  clang2py -k cdefstum \
+    $AMKERN_INC/asic_reg/mmhub/mmhub_4_1_0_offset.h \
+    $AMKERN_INC/asic_reg/mmhub/mmhub_4_1_0_sh_mask.h \
+    -o $BASE/am/mmhub_4_1_0.py
+  fixup $BASE/am/mmhub_4_1_0.py
+
+  clang2py -k cdefstum \
+    $AMKERN_INC/asic_reg/nbio/nbio_4_3_0_offset.h \
+    $AMKERN_INC/asic_reg/nbio/nbio_4_3_0_sh_mask.h \
    -o $BASE/am/nbio_4_3_0.py
  fixup $BASE/am/nbio_4_3_0.py

  clang2py -k cdefstum \
-    extra/amdpci/headers/osssys_6_0_0_offset.h \
-    extra/amdpci/headers/osssys_6_0_0_sh_mask.h \
+    $AMKERN_INC/asic_reg/nbif/nbif_6_3_1_offset.h \
+    $AMKERN_INC/asic_reg/nbif/nbif_6_3_1_sh_mask.h \
+    -o $BASE/am/nbif_6_3_1.py
+  fixup $BASE/am/nbif_6_3_1.py
+
+  clang2py -k cdefstum \
+    $AMKERN_INC/asic_reg/nbio/nbio_7_9_0_offset.h \
+    $AMKERN_INC/asic_reg/nbio/nbio_7_9_0_sh_mask.h \
+    -o $BASE/am/nbio_7_9_0.py
+  fixup $BASE/am/nbio_7_9_0.py
+
+  clang2py -k cdefstum \
+    $AMKERN_INC/asic_reg/nbio/nbio_7_11_0_offset.h \
+    $AMKERN_INC/asic_reg/nbio/nbio_7_11_0_sh_mask.h \
+    -o $BASE/am/nbio_7_11_0.py
+  fixup $BASE/am/nbio_7_11_0.py
+
+  clang2py -k cdefstum \
+    $AMKERN_INC/asic_reg/oss/osssys_6_0_0_offset.h \
+    $AMKERN_INC/asic_reg/oss/osssys_6_0_0_sh_mask.h \
    -o $BASE/am/osssys_6_0_0.py
  fixup $BASE/am/osssys_6_0_0.py

  clang2py -k cdefstum \
-    extra/amdpci/headers/smu_v13_0_0_ppsmc.h \
-    extra/amdpci/headers/smu13_driver_if_v13_0_0.h \
+    $AMKERN_INC/asic_reg/oss/osssys_7_0_0_offset.h \
+    $AMKERN_INC/asic_reg/oss/osssys_7_0_0_sh_mask.h \
+    -o $BASE/am/osssys_7_0_0.py
+  fixup $BASE/am/osssys_7_0_0.py
+
+  clang2py -k cdefstum \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu_v13_0_0_ppsmc.h \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu13_driver_if_v13_0_0.h \
    extra/amdpci/headers/amdgpu_smu.h \
    -o $BASE/am/smu_v13_0_0.py
  fixup $BASE/am/smu_v13_0_0.py
+
+  clang2py -k cdefstum \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu_v14_0_0_pmfw.h \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu_v14_0_2_ppsmc.h \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu14_driver_if_v14_0.h \
+    extra/amdpci/headers/amdgpu_smu.h \
+    --clang-args="-include stdint.h" \
+    -o $BASE/am/smu_v14_0_3.py
+  fixup $BASE/am/smu_v14_0_3.py
+
+  clang2py -k cdefstum \
+    $AMKERN_INC/asic_reg/hdp/hdp_6_0_0_offset.h \
+    $AMKERN_INC/asic_reg/hdp/hdp_6_0_0_sh_mask.h \
+    -o $BASE/am/hdp_6_0_0.py
+  fixup $BASE/am/hdp_6_0_0.py
+
+  clang2py -k cdefstum \
+    $AMKERN_INC/asic_reg/hdp/hdp_7_0_0_offset.h \
+    $AMKERN_INC/asic_reg/hdp/hdp_7_0_0_sh_mask.h \
+    -o $BASE/am/hdp_7_0_0.py
+  fixup $BASE/am/hdp_7_0_0.py
+}
+
+generate_sqtt() {
+  clang2py -k cdefstum \
+    extra/sqtt/sqtt.h \
+    -o $BASE/sqtt.py
+
+  fixup $BASE/sqtt.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/sqtt.py
+  python3 -c "import tinygrad.runtime.autogen.sqtt"
 }

 generate_webgpu() {
-  clang2py -l /usr/local/lib/libwebgpu_dawn.so extra/webgpu/webgpu.h -o $BASE/webgpu.py
+  clang2py extra/webgpu/webgpu.h -o $BASE/webgpu.py
  fixup $BASE/webgpu.py
-  sed -i 's/import ctypes/import ctypes, ctypes.util/g' $BASE/webgpu.py
-  sed -i "s|ctypes.CDLL('/usr/local/lib/libwebgpu_dawn.so')|ctypes.CDLL(ctypes.util.find_library('webgpu_dawn'))|g" $BASE/webgpu.py
+  sed -i "s/FIXME_STUB/webgpu/g" "$BASE/webgpu.py"
+  sed -i "s/FunctionFactoryStub()/ctypes.CDLL(webgpu_support.WEBGPU_PATH)/g" "$BASE/webgpu.py"
+  sed -i "s/import ctypes/import ctypes, tinygrad.runtime.support.webgpu as webgpu_support/g" "$BASE/webgpu.py"
  python3 -c "import tinygrad.runtime.autogen.webgpu"
 }

+generate_libusb() {
+  clang2py -k cdefstum \
+    /usr/include/libusb-1.0/libusb.h \
+    -o $BASE/libusb.py
+
+  fixup $BASE/libusb.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/libusb.py
+  sed -i "s/FIXME_STUB/libusb/g" "$BASE/libusb.py"
+  sed -i "s/libusb_le16_to_cpu = libusb_cpu_to_le16//g" "$BASE/libusb.py"
+  sed -i "s/FunctionFactoryStub()/None if (lib_path:=os.getenv('LIBUSB_PATH', ctypes.util.find_library('usb-1.0'))) is None else ctypes.CDLL(lib_path)/g" "$BASE/libusb.py"
+  python3 -c "import tinygrad.runtime.autogen.libusb"
+}
+
 if [ "$1" == "opencl" ]; then generate_opencl
 elif [ "$1" == "hip" ]; then generate_hip
 elif [ "$1" == "comgr" ]; then generate_comgr
@@ -373,6 +551,7 @@ elif [ "$1" == "kfd" ]; then generate_kfd
 elif [ "$1" == "nv" ]; then generate_nv
 elif [ "$1" == "amd" ]; then generate_amd
 elif [ "$1" == "am" ]; then generate_am
+elif [ "$1" == "sqtt" ]; then generate_sqtt
 elif [ "$1" == "qcom" ]; then generate_qcom
 elif [ "$1" == "io_uring" ]; then generate_io_uring
 elif [ "$1" == "libc" ]; then generate_libc
@@ -382,6 +561,7 @@ elif [ "$1" == "adreno" ]; then generate_adreno
 elif [ "$1" == "pci" ]; then generate_pci
 elif [ "$1" == "vfio" ]; then generate_vfio
 elif [ "$1" == "webgpu" ]; then generate_webgpu
+elif [ "$1" == "libusb" ]; then generate_libusb
 elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_libc; generate_am; generate_webgpu
 else echo "usage: $0 <type>"
 fi
@@ -1,7 +1,7 @@
 # tinygrad is a tensor library, and as a tensor library it has multiple parts
 # 1. a "runtime". this allows buffer management, compilation, and running programs
 # 2. a "Device" that uses the runtime but specifies compute in an abstract way for all
-# 3. a "LazyBuffer" that fuses the compute into kernels, using memory only when needed
+# 3. a "UOp" that fuses the compute into kernels, using memory only when needed
 # 4. a "Tensor" that provides an easy to use frontend with autograd ".backward()"


@@ -74,30 +74,47 @@ fxn.exec([out, a, b])
 assert out.as_buffer().cast('I')[0] == 5


-print("******** third, the LazyBuffer ***********")
+print("******** third, the UOp ***********")

 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.schedule import create_schedule_with_vars
+from tinygrad.engine.grouper import get_becomes_map

 # allocate some values + load in values
-a = UOp.metaop(Ops.EMPTY, (1,), dtypes.int32, DEVICE)
-b = UOp.metaop(Ops.EMPTY, (1,), dtypes.int32, DEVICE)
+a = UOp.new_buffer(DEVICE, 1, dtypes.int32)
+b = UOp.new_buffer(DEVICE, 1, dtypes.int32)
 a.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 2))))
 b.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 3))))

 # describe the computation
-out = a.alu(Ops.ADD, b)
+out = a + b
+s = UOp(Ops.SINK, dtypes.void, (out,))

-# schedule the computation as a list of kernels
-sched, _, becomes_map = create_schedule_with_vars(out.sink())
-for si in sched: print(si.ast.op)  # NOTE: the first two convert it to CPU
-# NOTE: UOps are no longer mutable, the scheduler gives you a map to lookup which BUFFER the result was written to
-out = becomes_map[out]
+# group the computation into kernels
+becomes_map = get_becomes_map(s)
+
+# the compute maps to an assign
+assign = becomes_map[a+b]
+
+# the first source is the output buffer (data)
+assert assign.src[0].op is Ops.BUFFER
+# the second source is the kernel (compute)
+assert assign.src[1].op is Ops.KERNEL
+
+# schedule the kernel graph in a linear list
+s = UOp(Ops.SINK, dtypes.void, (assign,))
+sched, _, becomes_map = create_schedule_with_vars(s)
+assert len(sched) == 1

 # DEBUGGING: print the compute ast
 print(sched[-1].ast)
 # NOTE: sched[-1].ast is the same as st_0 above

+# the output will be stored in a new buffer
+out = becomes_map[assign]
+assert out.op is Ops.BUFFER and not out.buffer.is_allocated()
+print(out)
+
 # run that schedule
 run_schedule(sched)

@@ -1,6 +1,6 @@
 # AM Driver

-AM driver is a userspace driver targeting AMD's 7900XTX. You only need tinygrad to send compute tasks to your GPU!
+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 `AMD=1`!
@@ -115,9 +115,8 @@ HCQ-compatible devices use a global timeline signal for synchronizing all operat
 ```python
 HWQueue().wait(your_device.timeline_signal, your_device.timeline_value - 1) \
         .exec(...)
-         .signal(your_device.timeline_signal, your_device.timeline_value) \
+         .signal(your_device.timeline_signal, your_device.next_timeline()) \
         .submit(your_device)
-your_device.timeline_value += 1

 # Optionally wait for execution
 your_device.timeline_signal.wait(your_device.timeline_value - 1)
@@ -0,0 +1,109 @@
+# Kernel Creation
+
+Tinygrad lazily builds up a graph of Tensor operations. The Tensor graph includes a mix of:
+
+- Buffer and Assignment Ops: `BUFFER`, `BUFFER_VIEW`, `COPY`, `ASSIGN`
+- Movement Ops: `RESHAPE`, `EXPAND`, `PERMUTE`, `PAD`, `SHRINK`, `FLIP`
+- Compute Ops: `ADD`, `MUL`, `REDUCE_AXIS`, ...
+
+`Tensor.kernelize` creates the kernels and buffers needed to realize the output Tensor(s).
+
+## Kernelize flow
+
+Let's see how a multiply add Tensor graph becomes a fused elementwise kernel.
+
+```py
+# initialize 3 input buffers on the device
+a = Tensor([1]).realize()
+b = Tensor([2]).realize()
+c = Tensor([3]).realize()
+
+# create the Tensor graph
+mul = a*b
+out = mul+c
+
+print(mul) # <Tensor <UOp METAL (1,) int (<Ops.MUL: 48>, None)> on METAL with grad None>
+print(out) # <Tensor <UOp METAL (1,) int (<Ops.ADD: 52>, None)> on METAL with grad None>
+
+out.kernelize()
+
+print(mul) # <Tensor <UOp METAL (1,) int (<Ops.MUL: 48>, None)> on METAL with grad None>
+print(out) # <Tensor <UOp METAL (1,) int (<Ops.ASSIGN: 66>, None)> on METAL with grad None>
+```
+
+The multiply Tensor stays the same because it is fused. The output Tensor's UOp becomes a new ASSIGN UOp:
+
+```py
+print(out.lazydata)
+```
+
+The first source is the output BUFFER:
+
+```
+UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+  UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),
+  UOp(Ops.UNIQUE, dtypes.void, arg=6, src=()),))
+```
+
+And the second source is the KERNEL and its 4 buffer edges (output_buffer, a, b, c):
+
+```
+UOp(Ops.KERNEL, dtypes.void, arg=<Kernel 12 SINK(<Ops.STORE: 45>,) (__add__, __mul__)>, src=(
+  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+    x1:=UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),
+    UOp(Ops.UNIQUE, dtypes.void, arg=6, src=()),)),
+  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+     x1,
+    UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),)),
+  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+     x1,
+    UOp(Ops.UNIQUE, dtypes.void, arg=3, src=()),)),
+  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+     x1,
+    UOp(Ops.UNIQUE, dtypes.void, arg=5, src=()),)),))
+```
+
+KERNEL describes the compute AST, metadata and memory dependencies.
+
+BUFFER holds a reference to the device memory where the output will be stored.
+
+Once a Tensor is kernelized, all children will LOAD its BUFFER, instead of fusing it:
+
+```py
+child = out+2
+child.kernelize()
+print(child.lazydata.src[1].arg.ast)
+```
+
+```
+UOp(Ops.SINK, dtypes.void, arg=None, src=(
+  UOp(Ops.STORE, dtypes.void, arg=None, src=(
+    UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), arg=0, src=()),
+    x2:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True),)), src=()),
+    UOp(Ops.ADD, dtypes.int, arg=None, src=(
+      UOp(Ops.LOAD, dtypes.int, arg=None, src=(
+        UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), arg=1, src=()),
+         x2,)),
+      UOp(Ops.CONST, dtypes.int, arg=2, src=(
+         x2,)),)),)),))
+```
+
+`Tensor.realize` will execute the kernels and write outputs to memory:
+
+```py
+Tensor.realize(out)
+print(out)        # <Tensor <UOp METAL (1,) int (<Ops.BUFFER: 23>, <buf real:True device:METAL size:1 dtype:dtypes.int offset:0>)> on METAL with grad None>
+print(out.item()) # 5
+```
+
+<hr />
+
+**Summary**
+
+- The large Tensor graph is built from a mix of data, compute and movement Ops.
+
+- `Tensor.kernelize` splits the Tensor graph into data (BUFFER), compute (KERNEL) and links dependencies with ASSIGN.
+
+- `Tensor.realize` executes KERNELs on device and replaces the Tensor graph with just a BUFFER.
+
+- Kernelize can be called multiple times on a Tensor. This allows for incrementally building the kernel fusion layout of a large Tensor graph, without having to call `realize` or `schedule`.
@@ -30,7 +30,7 @@ These control the behavior of core tinygrad even when used as a library.

 Variable | Possible Value(s) | Description
 ---|---|---
-DEBUG               | [1-6]      | enable debugging output, with 4 you get operations, timings, speed, generated code and more
+DEBUG               | [1-7]      | enable debugging output (operations, timings, speed, generated code and more)
 GPU                 | [1]        | enable the GPU (OpenCL) backend
 CUDA                | [1]        | enable CUDA backend
 AMD                 | [1]        | enable AMD backend
@@ -49,4 +49,17 @@ VISIBLE_DEVICES     | [list[int]]| restricts the NV/AMD devices that are availab
 JIT                 | [0-2]      | 0=disabled, 1=[jit enabled](quickstart.md#jit) (default), 2=jit enabled, but graphs are disabled
 VIZ                 | [1]        | 0=disabled, 1=[viz enabled](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/viz)
 ALLOW_TF32          | [1]        | enable TensorFloat-32 tensor cores on Ampere or newer GPUs.
-WEBGPU_BACKEND      | [WGPUBackendType_Metal, ...]          | Force select a backend for WebGPU (Metal, DirectX, OpenGL, Vulkan...)
+WEBGPU_BACKEND      | [WGPUBackendType_Metal, ...]          | Force select a backend for WebGPU (Metal, DirectX, OpenGL, Vulkan...)
+CUDA_PATH           | str        | Use `CUDA_PATH/include` for CUDA headers for CUDA and NV backends. If not set, TinyGrad will use `/usr/local/cuda/include`, `/usr/include` and `/opt/cuda/include`.
+
+## Debug breakdown
+
+Variable | Value | Description
+---|---|---
+DEBUG               | >= 1       | Enables debugging and lists devices being used
+DEBUG               | >= 2       | Provides performance metrics for operations, including timing, memory usage, bandwidth for each kernel execution
+DEBUG               | >= 3       | Outputs buffers used for each kernel (shape, dtype and strides) and the applied optimizations at a kernel level
+DEBUG               | >= 4       | Outputs the generated kernel code
+DEBUG               | >= 5       | Displays the intermediate representation of the computation UOps (AST)
+DEBUG               | >= 6       | Displays the intermediate representation of the computation UOps in a linearized manner, detailing the operation sequence
+DEBUG               | >= 7       | Outputs the assembly code generated for the target hardware
@@ -29,5 +29,12 @@
 ::: tinygrad.nn.state.get_state_dict
 ::: tinygrad.nn.state.get_parameters
 ::: tinygrad.nn.state.load_state_dict
+::: tinygrad.nn.state.tar_extract
+    options:
+        show_signature: false
+        separate_signature: false
 ::: tinygrad.nn.state.torch_load
+    options:
+        show_signature: false
+        separate_signature: false
 ::: tinygrad.nn.state.gguf_load
@@ -110,7 +110,7 @@ class TinyNet:

  def __call__(self, x):
    x = self.l1(x)
-    x = x.leakyrelu()
+    x = x.leaky_relu()
    x = self.l2(x)
    return x

@@ -118,7 +118,7 @@ net = TinyNet()
 ```

 We can see that the forward pass of our neural network is just the sequence of operations performed on the input tensor `x`.
-We can also see that functional operations like `leakyrelu` are not defined as classes and instead are just methods we can just call.
+We can also see that functional operations like `leaky_relu` are not defined as classes and instead are just methods we can just call.
 Finally, we just initialize an instance of our neural network, and we are ready to start training it.

 ## Training
@@ -5,7 +5,7 @@ tinygrad supports various runtimes, enabling your code to scale across a wide ra
 | Runtime | Description | Requirements |
 |---------|-------------|--------------|
 | [NV](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_nv.py) | Provides acceleration for NVIDIA GPUs | Ampere/Ada series GPUs |
-| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | RDNA2/RDNA3 series GPUs |
+| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | RDNA2/RDNA3/RDNA4 series GPUs. You can select one of the interfaces for communication by setting `AMD_IFACE=(KFD|PCI)`. See [AMD interfaces](#amd-interfaces) for more details. |
 | [QCOM](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_qcom.py) | Provides acceleration for QCOM GPUs | 6xx series GPUs |
 | [METAL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_metal.py) | Utilizes Metal for acceleration on Apple devices | M1+ Macs; Metal 3.0+ for `bfloat` support |
 | [CUDA](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cuda.py) | Utilizes CUDA for acceleration on NVIDIA GPUs | NVIDIA GPU with CUDA support |
@@ -13,3 +13,62 @@ tinygrad supports various runtimes, enabling your code to scale across a wide ra
 | [CPU (C Code)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang compiler | `clang` compiler in system `PATH` |
 | [LLVM (LLVM IR)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_llvm.py) | Runs on CPU using the LLVM compiler infrastructure | llvm libraries installed and findable |
 | [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | Dawn library installed and findable. Download binaries [here](https://github.com/wpmed92/pydawn/releases/tag/v0.1.6). |
+
+## Interoperability
+
+tinygrad provides interoperability with OpenCL and PyTorch, allowing efficient tensor data sharing between frameworks through the `Tensor.from_blob` API. This enables zero-copy operations by working directly with external memory pointers.
+
+**Important**: When using external memory pointers with tinygrad tensors, you must ensure these pointers remain valid throughout the entire lifetime of the tinygrad tensor to prevent memory corruption.
+
+### `CUDA`/`METAL` PyTorch Interoperability
+
+You can seamlessly work with CUDA/MPS tensors between PyTorch and tinygrad without data copying:
+```python
+from tinygrad.dtype import _from_torch_dtype
+tensor1 = torch.tensor([1.0, 2.0, 3.0], device=torch.device("cuda"))
+tiny_tensor1 = Tensor.from_blob(tensor1.data_ptr(), tensor1.shape, dtype=_from_torch_dtype(tensor1.dtype), device='CUDA')
+
+# Before tinygrad calculations, mps needs to be synchronized to make sure data is valid.
+if data.device.type == "mps": torch.mps.synchronize()
+else: torch.cuda.synchronize()
+
+x = (tiny_tensor1 + 1).realize()
+```
+
+### `QCOM` OpenCL Interoperability
+
+tinygrad supports OpenCL interoperability on `QCOM` backend.
+
+Buffer interop allows direct access to OpenCL memory buffers:
+```python
+# create raw opencl buffer.
+cl_buf = cl.clCreateBuffer(cl_context, cl.CL_MEM_READ_WRITE, 0x100, None, status := ctypes.c_int32())
+
+# extract pointers
+cl_buf_desc_ptr = to_mv(ctypes.addressof(cl_buf), 8).cast('Q')[0]
+rawbuf_ptr = to_mv(cl_buf_desc_ptr, 0x100).cast('Q')[20] # offset 0xA0 is a raw gpu pointer.
+
+# create tiny tensor
+tiny = Tensor.from_blob(rawbuf_ptr, (8, 8), dtype=dtypes.int, device='QCOM')
+```
+
+And the same for the images:
+```python
+# create cl image.
+cl_img = cl.clCreateImage2D(cl_context, cl.CL_MEM_READ_WRITE, cl.cl_image_format(cl.CL_RGBA, cl.CL_FLOAT), w, h, 0, None, status := ctypes.c_int32())
+
+# extract pointers
+cl_buf_desc_ptr = to_mv(ctypes.addressof(cl_img), 8).cast('Q')[0]
+rawbuf_ptr = to_mv(cl_buf_desc_ptr, 0x100).cast('Q')[20] # offset 0xA0 is a raw gpu pointer.
+
+# create tiny tensor
+tiny = Tensor.from_blob(rawbuf_ptr, (h*w*4,), dtype=dtypes.imagef((h,w)), device='QCOM')
+```
+
+## AMD Interfaces
+AMD backend supports several interfaces for communicating with devices:
+
+* `KFD`: uses the amdgpu driver
+* `PCI`: uses the [AM driver](developer/am.md)
+
+You can force an interface by setting `AMD_IFACE` to one of these values. In the case of `AMD_IFACE=PCI`, this may unbind your GPU from the amdgpu driver.
@@ -20,7 +20,9 @@

 ::: tinygrad.Tensor.manual_seed
 ::: tinygrad.Tensor.rand
+::: tinygrad.Tensor.rand_like
 ::: tinygrad.Tensor.randn
+::: tinygrad.Tensor.randn_like
 ::: tinygrad.Tensor.randint
 ::: tinygrad.Tensor.normal
 ::: tinygrad.Tensor.uniform
@@ -22,6 +22,7 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.round
 ::: tinygrad.Tensor.isinf
 ::: tinygrad.Tensor.isnan
+::: tinygrad.Tensor.isfinite
 ::: tinygrad.Tensor.lerp
 ::: tinygrad.Tensor.square
 ::: tinygrad.Tensor.clamp
@@ -52,7 +53,7 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.erf
 ::: tinygrad.Tensor.gelu
 ::: tinygrad.Tensor.quick_gelu
-::: tinygrad.Tensor.leakyrelu
+::: tinygrad.Tensor.leaky_relu
 ::: tinygrad.Tensor.mish
 ::: tinygrad.Tensor.softplus
 ::: tinygrad.Tensor.softsign
@@ -65,13 +66,17 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.div
 ::: tinygrad.Tensor.idiv
 ::: tinygrad.Tensor.mod
-::: tinygrad.Tensor.xor
+::: tinygrad.Tensor.bitwise_xor
+::: tinygrad.Tensor.bitwise_and
+::: tinygrad.Tensor.bitwise_or
+::: tinygrad.Tensor.bitwise_not
 ::: tinygrad.Tensor.lshift
 ::: tinygrad.Tensor.rshift
 ::: tinygrad.Tensor.pow
 ::: tinygrad.Tensor.maximum
 ::: tinygrad.Tensor.minimum
 ::: tinygrad.Tensor.where
+::: tinygrad.Tensor.copysign

 ## Casting Ops

@@ -9,6 +9,7 @@
 ::: tinygrad.Tensor.isclose
 ::: tinygrad.Tensor.mean
 ::: tinygrad.Tensor.var
+::: tinygrad.Tensor.var_mean
 ::: tinygrad.Tensor.std
 ::: tinygrad.Tensor.std_mean
 ::: tinygrad.Tensor.softmax
@@ -22,6 +23,7 @@

 ::: tinygrad.Tensor.avg_pool2d
 ::: tinygrad.Tensor.max_pool2d
+::: tinygrad.Tensor.max_unpool2d
 ::: tinygrad.Tensor.conv2d
 ::: tinygrad.Tensor.conv_transpose2d
 ::: tinygrad.Tensor.dot
@@ -34,6 +36,9 @@
 ::: tinygrad.Tensor.interpolate
 ::: tinygrad.Tensor.scatter
 ::: tinygrad.Tensor.scatter_reduce
+::: tinygrad.Tensor.masked_select
+::: tinygrad.Tensor.sort
+::: tinygrad.Tensor.topk

 ## Neural Network (functional)

@@ -25,6 +25,7 @@
 ::: tinygrad.Tensor.replace
 ::: tinygrad.Tensor.assign
 ::: tinygrad.Tensor.detach
+::: tinygrad.Tensor.clone
 ::: tinygrad.Tensor.to
 ::: tinygrad.Tensor.to_
 ::: tinygrad.Tensor.shard
@@ -1,4 +1,4 @@
-# model based off https://towardsdatascience.com/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
+# model based off https://medium.com/data-science/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
 from typing import List, Callable
 from tinygrad import Tensor, TinyJit, nn, GlobalCounters
 from tinygrad.helpers import getenv, colored, trange
@@ -1,12 +1,12 @@
-import sys, onnx, time
-from tinygrad import TinyJit, Device, GlobalCounters, fetch, getenv
-from extra.onnx import OnnxRunner
+import sys, onnx, time, pickle
+from tinygrad import TinyJit, GlobalCounters, fetch, getenv
+from tinygrad.frontend.onnx import OnnxRunner
 from extra.onnx_helpers import get_example_inputs, validate

 def load_onnx_model(onnx_file):
  onnx_model = onnx.load(onnx_file)
  run_onnx = OnnxRunner(onnx_model)
-  run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(Device.DEFAULT) 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__":
@@ -33,4 +33,4 @@ if __name__ == "__main__":

  if getenv("ORT"):
    validate(onnx_file, new_inputs, rtol=1e-3, atol=1e-3)
-    print("model validated")
+    print("model validated")
@@ -34,8 +34,8 @@ if __name__ == "__main__":
    part2 = nn.state.torch_load(fetch("https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/resolve/main/pytorch_model-00002-of-00002.bin?download=true"))

  with Timing("weights -> model: "):
-    nn.state.load_state_dict(model, fix_bf16(convert_from_huggingface(part1, model, 32, 8)), strict=False)
-    nn.state.load_state_dict(model, fix_bf16(convert_from_huggingface(part2, model, 32, 8)), strict=False)
+    nn.state.load_state_dict(model, fix_bf16(convert_from_huggingface(part1, 32, 32, 8)), strict=False)
+    nn.state.load_state_dict(model, fix_bf16(convert_from_huggingface(part2, 32, 32, 8)), strict=False)

  if not os.path.isfile("/tmp/tokenizer.model"): create_fixed_tokenizer("/tmp/tokenizer.model")
  spp = SentencePieceProcessor(model_file="/tmp/tokenizer.model")
@@ -8,7 +8,7 @@ import numpy as np
 import subprocess
 import tensorflow as tf
 import tf2onnx
-from extra.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner
 from tinygrad.tensor import Tensor
 from extra.export_model import export_model_clang, compile_net, jit_model

@@ -189,7 +189,7 @@ class GPT2:
          tokens = Variable("tokens", 0, VOCAB_SIZE).bind(toks[0][start_pos])
        else:
          tokens = Tensor([x[start_pos:] for x in toks])
-        tok = self.model(tokens, Variable("start_pos", 1 if start_pos else 0, MAX_CONTEXT).bind(start_pos), temperature).tolist()
+        tok = self.model(tokens, Variable("start_pos", 1 if start_pos else 0, MAX_CONTEXT-1).bind(start_pos), temperature).tolist()
      start_pos = len(toks[0])
      for i,t in enumerate(tok): toks[i].append(t)
    return [self.tokenizer.decode(x) for x in toks]
@@ -3,6 +3,7 @@ from extra.mcts_search import mcts_search
 from examples.mlperf.helpers import get_mlperf_bert_model
 from tinygrad import Tensor, Device, dtypes, nn
 from tinygrad.codegen.kernel import Kernel
+from tinygrad.codegen.heuristic import hand_coded_optimizations
 from tinygrad.ops import Ops, sym_infer
 from tinygrad.device import Compiled
 from tinygrad.engine.search import beam_search, bufs_from_lin
@@ -83,7 +84,7 @@ if __name__ == "__main__":

    # always try hand coded opt
    lin = Kernel(si.ast, opts=device.renderer)
-    lin.hand_coded_optimizations()
+    lin.apply_opts(hand_coded_optimizations(lin))
    lins.append((lin, "HC"))

    # maybe try tensor cores
@@ -211,11 +211,13 @@ class LLaMa:
    else:
      weights = load(str(model_path))
    if "model.embed_tokens.weight" in weights:
-      weights = convert_from_huggingface(weights, model, params["args"]["n_heads"], params["args"].get("n_kv_heads", params["args"]["n_heads"]))
+      weights = convert_from_huggingface(weights, params["args"]["n_layers"], params["args"]["n_heads"], params["args"].get("n_kv_heads", params["args"]["n_heads"]))

    weights = fix_bf16(weights)

-    with Context(BEAM=0):
+    # prevent tracking model weights
+    # this is a part of a larger problem with BUFFER UOps and gc in TRACK_MATCH_STATS=2
+    with Context(BEAM=0, TRACK_MATCH_STATS=0):
      # quantize
      if quantize is not None:
        weights = linear.quantize(weights, device)
@@ -47,7 +47,7 @@ def concat_weights(models, device=None):
    disk_tensors: List[Tensor] = [model[name] for model in models]
    if len(disk_tensors) == 1 or len(disk_tensors[0].shape) == 1:
      return disk_tensors[0].to(device=device)
-    axis = 1 if name.endswith(".attention.wo.weight") or name.endswith(".feed_forward.w2.weight") else 0
+    axis = 1 if name.endswith((".attention.wo.weight", ".feed_forward.w2.weight")) else 0
    lazy_tensors = [data.to(device=device) for data in disk_tensors]
    return lazy_tensors[0].cat(*lazy_tensors[1:], dim=axis)
  return {name: convert(name) for name in {name: None for model in models for name in model}}
@@ -73,16 +73,17 @@ class Int8Linear:
    self.scale = Tensor.ones(out_features, dtype=dtypes.half)

  def __call__(self, x):
-    return x.dot(self.weight.cast(dtype=dtypes.half).T*self.scale)
+    return x.dot(self.weight.cast(self.scale.dtype).T*self.scale)

  @staticmethod
-  def quantize(tensors, device):
+  def quantize(tensors, device, scale_dtype=dtypes.float16, quantize_embeds=False):
    new_tensors = {}
    for name,v in tensors.items():
-      if "feed_forward" in name or "attention.w" in name:
+      if "feed_forward" in name or "attention.w" in name or (quantize_embeds and "tok_embeddings.weight" in name):
        assert "weight" in name, name
+        v = v.cast(scale_dtype)
        scale = v.abs().max(axis=1) / 127.0
-        int8_weight = (v.T/scale).T.cast(dtype=dtypes.int8)
+        int8_weight = (v.T/scale).T.round().cast(dtype=dtypes.int8) # without round(), cast truncates -34.9 to -34
        new_tensors[name] = int8_weight
        new_tensors[name.replace('weight', 'scale')] = scale
        if isinstance(device, tuple):
@@ -90,8 +91,20 @@ class Int8Linear:
          new_tensors[name.replace('weight', 'scale')].shard_(device, axis=None)
      else:
        new_tensors[name] = v
+    if quantize_embeds: new_tensors.update({"output.weight": new_tensors["tok_embeddings.weight"], "output.scale": new_tensors["tok_embeddings.scale"]})
    return new_tensors

+class Int8Embedding:
+  def __init__(self, vocab_size:int, embed_size:int):
+    self.vocab_sz, self.embed_sz = vocab_size, embed_size
+    self.weight, self.scale = Tensor.ones(vocab_size, embed_size, dtype=dtypes.int8), Tensor.ones(vocab_size, dtype=dtypes.half)
+
+  def __call__(self, idx:Tensor) -> Tensor:
+    if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, requires_grad=False, device=self.weight.device).unsqueeze(-1)
+    big_shp = idx.shape+(self.vocab_sz, self.embed_sz)
+    arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1)).expand(big_shp), (self.weight.cast(self.scale.dtype).T*self.scale).T
+    return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
+
 def NF4Linear(block_size):
  _CODE = [
    -1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453, -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0,
@@ -113,7 +126,8 @@ def NF4Linear(block_size):
      return x.linear(unscaled.reshape(self.out_features, self.in_features).T)

    @staticmethod
-    def quantize(state_dict: dict[str, Tensor], device) -> dict[str, Tensor]:
+    def quantize(state_dict: dict[str, Tensor], device, scale_dtype=dtypes.float16, quantize_embeds=False) -> dict[str, Tensor]:
+      assert not quantize_embeds  # TODO: support this?
      new_state_dict = {}
      for k, v in state_dict.items():
        if "feed_forward" in k or "attention.w" in k:
@@ -121,7 +135,7 @@ def NF4Linear(block_size):
          scale = (grouped.abs().max(axis=1, keepdim=True))
          coded = ((grouped / scale).unsqueeze(-1) - CODE.to(v.device)).abs().argmin(axis=-1).cast(dtypes.uint8).flatten()
          new_state_dict[k] = coded[::2] * 2 ** 4 + coded[1::2]
-          new_state_dict[k.replace(".weight", ".scale")] = scale.cast(dtypes.float16)
+          new_state_dict[k.replace(".weight", ".scale")] = scale.cast(scale_dtype)
          if isinstance(device, tuple):
            new_state_dict[k].shard_(device, axis=-1)
            new_state_dict[k.replace('weight', 'scale')].shard_(device, axis=None)
@@ -144,13 +158,14 @@ MODEL_PARAMS = {
    "files": 8
  }
 }
-def build_transformer(model_path: Path, model_size="8B", quantize=None, device=None):
+def build_transformer(model_path: Path, model_size="8B", quantize=None, scale_dtype=dtypes.float16, device=None, max_context=8192, load_weights=True):
  # build model
-  if quantize == "int8": linear = Int8Linear
-  elif quantize == "nf4": linear = NF4Linear(64)
-  else: linear = nn.Linear
-  model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=linear, max_context=8192, jit=True)
+  if quantize == "int8": linear, embedding, quantize_embeds = Int8Linear, Int8Embedding, True
+  elif quantize == "nf4": linear, embedding, quantize_embeds = NF4Linear(64), nn.Embedding, False
+  else: linear, embedding, quantize_embeds = nn.Linear, nn.Embedding, False
+  model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=linear, embedding=embedding, max_context=max_context, jit=True)

+  if not load_weights: return model
  # load weights
  if model_path.is_dir():
    if (model_path / "model.safetensors.index.json").exists(): weights = load(str(model_path / "model.safetensors.index.json"))
@@ -159,16 +174,16 @@ def build_transformer(model_path: Path, model_size="8B", quantize=None, device=N
  else:
    weights = load(str(model_path))
  if "model.embed_tokens.weight" in weights:
-    weights = convert_from_huggingface(weights, model, MODEL_PARAMS[model_size]["args"]["n_heads"], MODEL_PARAMS[model_size]["args"]["n_kv_heads"])
+    weights = convert_from_huggingface(weights, MODEL_PARAMS[model_size]["args"]["n_layers"], MODEL_PARAMS[model_size]["args"]["n_heads"], MODEL_PARAMS[model_size]["args"]["n_kv_heads"])
  elif "token_embd.weight" in weights:
-    weights = convert_from_gguf(weights, model)
+    weights = convert_from_gguf(weights, MODEL_PARAMS[model_size]["args"]["n_layers"])
  weights = fix_bf16(weights)

  with Context(BEAM=0):
    # quantize
    if quantize == "float16": weights = {k:v.cast(quantize).contiguous() for k,v in weights.items()}
    elif quantize is not None:
-      weights = linear.quantize(weights, device)
+      weights = linear.quantize(weights, device, scale_dtype, quantize_embeds)
      for _,v in weights.items(): v.realize()

    # shard
@@ -25,7 +25,7 @@ if __name__ == "__main__":
  Tensor.training = True
  optimizer = nn.optim.Adam(nn.state.get_parameters(model), lr=1e-4)
  warmup_count = getenv("WARMUP", 3)
-  for i in range(warmup_count):  # TODO: why does it take three and not two to stablize
+  for i in range(warmup_count):  # TODO: why does it take three and not two to stabilize
    GlobalCounters.reset()
    X = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
    Y = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
@@ -0,0 +1,157 @@
+# much taken from https://github.com/cloneofsimo/minRF
+from tinygrad import Tensor, nn, GlobalCounters, TinyJit
+from tinygrad.helpers import getenv, trange
+from extra.models.llama import Attention, FeedForward, precompute_freqs_cis
+
+def modulate(x:Tensor, shift:Tensor, scale:Tensor) -> Tensor: return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+# TODO: why doesn't the TimestepEmbedder from minRF work?
+class TimestepEmbedder:
+  def __init__(self, hidden_size): self.mlp = [nn.Linear(1, hidden_size), Tensor.silu, nn.Linear(hidden_size, hidden_size)]
+  def __call__(self, t:Tensor): return t.reshape(-1, 1).sequential(self.mlp)
+
+class TransformerBlock:
+  def __init__(self, dim, n_heads, norm_eps=1e-5):
+    self.attention = Attention(dim, n_heads)
+    self.feed_forward = FeedForward(dim, 4*dim)
+    self.attention_norm = nn.LayerNorm(dim, eps=norm_eps)
+    self.ffn_norm = nn.LayerNorm(dim, eps=norm_eps)
+    self.adaLN_modulation = nn.Linear(dim, 6 * dim, bias=True)
+
+  def __call__(self, x:Tensor, freqs_cis:Tensor, adaln_input:Tensor):
+    shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(adaln_input.silu()).chunk(6, dim=1)
+    x = x + gate_msa.unsqueeze(1) * self.attention(modulate(self.attention_norm(x), shift_msa, scale_msa), 0, freqs_cis)
+    x = x + gate_mlp.unsqueeze(1) * self.feed_forward(modulate(self.ffn_norm(x), shift_mlp, scale_mlp))
+    return x.contiguous().contiguous_backward()
+
+class FinalLayer:
+  def __init__(self, dim, patch_size, out_channels):
+    self.norm_final = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+    self.linear = nn.Linear(dim, patch_size*patch_size*out_channels, bias=True)
+    self.adaLN_modulation = nn.Linear(dim, 2 * dim, bias=True)
+
+    # init weights/bias to 0
+    self.linear.weight.replace(self.linear.weight.zeros_like().contiguous())
+    self.linear.bias.replace(self.linear.bias.zeros_like().contiguous())
+
+  def __call__(self, x:Tensor, c:Tensor):
+    shift, scale = self.adaLN_modulation(c.silu()).chunk(2, dim=1)
+    x = modulate(self.norm_final(x), shift, scale)
+    return self.linear(x)
+
+# channels=1, input_size=32, dim=64, n_layers=6, n_heads=4, num_classes=10
+class DiT_Llama:
+  def __init__(self, in_channels=1, dim=64, n_layers=6, n_heads=4, num_classes=10, patch_size=2):
+    self.patch_size = patch_size
+    self.out_channels = in_channels
+    self.num_classes = num_classes
+
+    self.init_conv_seq = [
+      nn.Conv2d(in_channels, dim // 2, kernel_size=5, padding=2, stride=1), Tensor.silu, nn.GroupNorm(32, dim//2),
+      nn.Conv2d(dim //2, dim // 2, kernel_size=5, padding=2, stride=1), Tensor.silu, nn.GroupNorm(32, dim//2),
+    ]
+
+    self.x_embedder = nn.Linear(self.patch_size * self.patch_size * dim // 2, dim, bias=True)
+    self.t_embedder = TimestepEmbedder(dim)
+    self.y_embedder = nn.Embedding(num_classes+1, dim)
+    self.final_layer = FinalLayer(dim, self.patch_size, self.out_channels)
+
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, 4096)
+    self.layers = [TransformerBlock(dim, n_heads) for _ in range(n_layers)]
+
+  def unpatchify(self, x:Tensor):
+    c, p = self.out_channels, self.patch_size
+    h = w = int(x.shape[1] ** 0.5)
+    x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
+    x = x.rearrange("n h w p q c -> n c h p w q")
+    return x.reshape(shape=(x.shape[0], c, h * p, h * p))
+
+  def patchify(self, x:Tensor):
+    B, C, H, W = x.shape
+    x = x.reshape(B, C, H // self.patch_size, self.patch_size, W // self.patch_size, self.patch_size)
+    x = x.permute(0, 2, 4, 1, 3, 5).flatten(-3).flatten(1, 2)
+    return x  # B <H*W ish> <C*patch_size*patch_size>
+
+  def __call__(self, x:Tensor, t:Tensor, y:Tensor) -> Tensor:
+    x = x.sequential(self.init_conv_seq)
+    x = self.patchify(x)
+    x = self.x_embedder(x)
+    adaln_input = self.t_embedder(t) + self.y_embedder(y)
+    adaln_input = adaln_input.contiguous()
+    for layer in self.layers:
+      x = layer(x, self.freqs_cis[:, :x.size(1)], adaln_input=adaln_input)
+    x = self.final_layer(x, adaln_input)
+    return self.unpatchify(x)
+
+  def rf(self, x:Tensor, cond:Tensor):
+    b = x.shape[0]
+    # self.ln is True
+    t = Tensor.randn((b,)).sigmoid()
+    texp = t.view([b, *([1] * len(x.shape[1:]))])
+
+    # conditional dropout
+    dropout_prob = 0.1
+    cond = (Tensor.rand(cond.shape[0]) < dropout_prob).where(cond.full_like(self.num_classes), cond)
+
+    # this is rectified flow
+    z1 = x.randn_like()
+    zt = (1 - texp) * x + texp * z1
+    vtheta = self(zt, t, cond)
+
+    # MSE loss
+    return ((z1 - x) - vtheta).square().mean()
+
+  def sample(self, z, cond, null_cond, sample_steps=50, cfg=2.0):
+    b = z.size(0)
+    dt = Tensor.full((b,)+(1,)*len(z.shape[1:]), fill_value=1.0/sample_steps).contiguous()
+    images = [z]
+    for i in range(sample_steps, 0, -1):
+      t = Tensor.full((b,), fill_value=i/sample_steps).contiguous()
+      vc = self(z, t, cond)
+      vu = self(z, t, null_cond)
+      vc = vu + cfg * (vc - vu)
+      z = z - dt * vc
+      z = z.contiguous()
+      images.append(z)
+    return images
+
+def mviz(t:Tensor):
+  assert len(t.shape) == 4 and t.shape[1] == 1
+  ft = t.permute(1,2,0,3).reshape(32, -1)
+  assert ft.shape[-1]%32 == 0
+  print("")
+  for y in ((ft+1)/2).clamp(0,1).tolist():
+    ln = [f"\033[38;5;{232+int(x*23)}m██" for x in y]
+    print(''.join(ln) + "\033[0m")
+
+if __name__ == "__main__":
+  X_train, Y_train, X_test, Y_test = nn.datasets.mnist()
+  X_train = X_train.pad((2,2,2,2))
+  X_train = ((X_train.float()/255)-0.5)/0.5
+  Y_train = Y_train.int()
+
+  model = DiT_Llama(patch_size=getenv("PATCH_SIZE", 2))
+  for r in nn.state.get_parameters(model): r.realize()
+  optimizer = nn.optim.Adam(nn.state.get_parameters(model), lr=5e-4)
+
+  @TinyJit
+  @Tensor.train()
+  def train_step():
+    if getenv("OVERFIT"): samples = Tensor.zeros(getenv("BS", 256), dtype='int')
+    else: samples = Tensor.randint(getenv("BS", 256), high=X_train.shape[0])
+    optimizer.zero_grad()
+    loss = model.rf(X_train[samples], Y_train[samples])
+    loss.backward()
+    optimizer.step()
+    return loss
+
+  @TinyJit
+  @Tensor.test()
+  def sample(z:Tensor, cond:Tensor) -> Tensor:
+    return model.sample(z, cond, Tensor.full_like(cond, 10), sample_steps=getenv("SAMPLE_STEPS", 20))[-1]
+
+  for steps in (t:=trange(getenv("STEPS", 5000))):
+    if steps%10 == 0: mviz(sample(Tensor.randn(3, 1, 32, 32), Tensor([5,0,4], dtype='int')))
+    GlobalCounters.reset()
+    loss = train_step()
+    t.set_description(f"loss: {loss.item():9.2f}")
@@ -348,6 +348,167 @@ def batch_load_unet3d(preprocessed_dataset_dir:Path, batch_size:int=6, val:bool=
      # happens with BENCHMARK set
      pass

+### RetinaNet
+
+def load_retinanet_data(base_dir:Path, val:bool, queue_in:Queue, queue_out:Queue,
+                        imgs:Tensor, boxes:Tensor, labels:Tensor, matches:Tensor|None=None,
+                        anchors:Tensor|None=None, seed:int|None=None):
+  from extra.datasets.openimages import image_load, random_horizontal_flip, resize
+  from examples.mlperf.helpers import box_iou, find_matches, generate_anchors
+  import torch
+
+  while (data:=queue_in.get()) is not None:
+    idx, img, tgt = data
+    img = image_load(base_dir, img["subset"], img["file_name"])
+
+    if val:
+      img = resize(img)[0]
+    else:
+      if seed is not None:
+        np.random.seed(seed)
+        random.seed(seed)
+        torch.manual_seed(seed)
+
+      img, tgt = random_horizontal_flip(img, tgt)
+      img, tgt, _ = resize(img, tgt=tgt)
+      match_quality_matrix = box_iou(tgt["boxes"], (anchor := np.concatenate(generate_anchors((800, 800)))))
+      match_idxs = find_matches(match_quality_matrix, allow_low_quality_matches=True)
+      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].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_boxes.tobytes()
+      labels[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_labels.tobytes()
+      matches[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = match_idxs.tobytes()
+      anchors[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = anchor.tobytes()
+
+    imgs[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
+
+    queue_out.put(idx)
+  queue_out.put(None)
+
+def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, shuffle:bool=True, seed:int|None=None):
+  def _enqueue_batch(bc):
+    from extra.datasets.openimages import prepare_target
+    for idx in range(bc * batch_size, (bc+1) * batch_size):
+      img = dataset.loadImgs(next(dataset_iter))[0]
+      ann = dataset.loadAnns(dataset.getAnnIds(img_id:=img["id"]))
+      tgt = prepare_target(ann, img_id, (img["height"], img["width"]))
+
+      if img_ids is not None:
+        img_ids[idx] = img_id
+
+      if img_sizes is not None:
+        img_sizes[idx] = tgt["image_size"]
+
+      queue_in.put((idx, img, tgt))
+
+  def _setup_shared_mem(shm_name:str, size:tuple[int, ...], dtype:dtypes) -> tuple[shared_memory.SharedMemory, Tensor]:
+    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}")
+    return shm, shm_tensor
+
+  image_ids = sorted(dataset.imgs.keys())
+  batch_count = min(32, len(image_ids) // batch_size)
+
+  queue_in, queue_out = Queue(), Queue()
+  procs, data_out_count = [], [0] * batch_count
+
+  shm_imgs, imgs = _setup_shared_mem("retinanet_imgs", (batch_size * batch_count, 800, 800, 3), dtypes.uint8)
+
+  if val:
+    boxes, labels, matches, anchors = None, None, None, None
+    img_ids, img_sizes = [None] * (batch_size * batch_count), [None] * (batch_size * batch_count)
+  else:
+    img_ids, img_sizes = None, None
+    shm_boxes, boxes = _setup_shared_mem("retinanet_boxes", (batch_size * batch_count, 120087, 4), dtypes.float32)
+    shm_labels, labels = _setup_shared_mem("retinanet_labels", (batch_size * batch_count, 120087), dtypes.int64)
+    shm_matches, matches = _setup_shared_mem("retinanet_matches", (batch_size * batch_count, 120087), dtypes.int64)
+    shm_anchors, anchors = _setup_shared_mem("retinanet_anchors", (batch_size * batch_count, 120087, 4), dtypes.float64)
+
+  shutdown = False
+  class Cookie:
+    def __init__(self, bc):
+      self.bc = bc
+    def __del__(self):
+      if not shutdown:
+        try: _enqueue_batch(self.bc)
+        except StopIteration: pass
+
+  def shuffle_indices(indices, seed):
+    rng = random.Random(seed)
+    rng.shuffle(indices)
+
+  if shuffle: shuffle_indices(image_ids, seed=seed)
+  dataset_iter = iter(image_ids)
+
+  try:
+    for _ in range(cpu_count()):
+      proc = Process(
+        target=load_retinanet_data,
+        args=(base_dir, val, queue_in, queue_out, imgs, boxes, labels),
+        kwargs={"matches": matches, "anchors": anchors, "seed": seed}
+      )
+      proc.daemon = True
+      proc.start()
+      procs.append(proc)
+
+    for bc in range(batch_count):
+      _enqueue_batch(bc)
+
+    for _ in range(len(image_ids) // batch_size):
+      while True:
+        bc = queue_out.get() // batch_size
+        data_out_count[bc] += 1
+        if data_out_count[bc] == batch_size: break
+
+      data_out_count[bc] = 0
+
+      if val:
+        yield (imgs[bc * batch_size:(bc + 1) * batch_size],
+               img_ids[bc * batch_size:(bc + 1) * batch_size],
+               img_sizes[bc * batch_size:(bc + 1) * batch_size],
+               Cookie(bc))
+      else:
+        yield (imgs[bc * batch_size:(bc + 1) * batch_size],
+               boxes[bc * batch_size:(bc + 1) * batch_size],
+               labels[bc * batch_size:(bc + 1) * batch_size],
+               matches[bc * batch_size:(bc + 1) * batch_size],
+               anchors[bc * batch_size:(bc + 1) * batch_size],
+               Cookie(bc))
+  finally:
+    shutdown = True
+
+    for _ in procs: queue_in.put(None)
+    queue_in.close()
+
+    for _ in procs:
+      while queue_out.get() is not None: pass
+    queue_out.close()
+
+    # shutdown processes
+    for proc in procs: proc.join()
+
+    shm_imgs.close()
+
+    if not val:
+      shm_boxes.close()
+      shm_labels.close()
+      shm_matches.close()
+      shm_anchors.close()
+
+    try:
+      shm_imgs.unlink()
+
+      if not val:
+        shm_boxes.unlink()
+        shm_labels.unlink()
+        shm_matches.unlink()
+        shm_anchors.unlink()
+    except FileNotFoundError:
+      # happens with BENCHMARK set
+      pass
+
 if __name__ == "__main__":
  def load_unet3d(val):
    assert not val, "validation set is not supported due to different sizes on inputs"
@@ -368,6 +529,14 @@ if __name__ == "__main__":
      for x,y,c in batch_load_resnet(val=val):
        pbar.update(x.shape[0])

+  def load_retinanet(val):
+    from extra.datasets.openimages import BASEDIR, download_dataset
+    from pycocotools.coco import COCO
+    dataset = COCO(download_dataset(base_dir:=getenv("BASEDIR", BASEDIR), "validation" if val else "train"))
+    with tqdm(total=len(dataset.imgs.keys())) as pbar:
+      for x in batch_load_retinanet(dataset, val, base_dir):
+        pbar.update(x[0].shape[0])
+
  load_fn_name = f"load_{getenv('MODEL', 'resnet')}"
  if load_fn_name in globals():
    globals()[load_fn_name](getenv("VAL", 1))
@@ -1,6 +1,7 @@
 from collections import OrderedDict
 import unicodedata
 from typing import Optional
+import math
 import numpy as np
 from tinygrad.nn import state
 from tinygrad.tensor import Tensor, dtypes
@@ -230,3 +231,126 @@ def get_fake_data_bert(BS:int):
    "masked_lm_weights": Tensor.empty((BS, 76), dtype=dtypes.float32, device="CPU"),
    "next_sentence_labels": Tensor.empty((BS, 1), dtype=dtypes.int32, device="CPU"),
  }
+
+def find_matches(match_quality_matrix:np.ndarray, high_threshold:float=0.5, low_threshold:float=0.4, allow_low_quality_matches:bool=False) -> np.ndarray:
+  BELOW_LOW_THRESHOLD, BETWEEN_THRESHOLDS = -1, -2
+
+  def _set_low_quality_matches_(matches:np.ndarray, all_matches:np.ndarray, match_quality_matrix:np.ndarray):
+    highest_quality_foreach_gt = np.max(match_quality_matrix, axis=1)
+    pred_inds_to_update = np.nonzero(match_quality_matrix == highest_quality_foreach_gt[:, None])[1]
+    matches[pred_inds_to_update] = all_matches[pred_inds_to_update]
+
+  assert low_threshold <= high_threshold
+
+  matched_vals, matches = match_quality_matrix.max(axis=0), match_quality_matrix.argmax(axis=0)
+  all_matches = np.copy(matches) if allow_low_quality_matches else None
+  below_low_threshold = matched_vals < low_threshold
+  between_thresholds = (matched_vals >= low_threshold) & (matched_vals < high_threshold)
+  matches[below_low_threshold] = BELOW_LOW_THRESHOLD
+  matches[between_thresholds] = BETWEEN_THRESHOLDS
+
+  if allow_low_quality_matches:
+    assert all_matches is not None
+    _set_low_quality_matches_(matches, all_matches, match_quality_matrix)
+
+  return matches
+
+def box_iou(boxes1:np.ndarray, boxes2:np.ndarray) -> np.ndarray:
+  def _box_area(boxes:np.ndarray) -> np.ndarray: return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
+
+  def _box_inter_union(boxes1:np.ndarray, boxes2:np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    area1, area2 = _box_area(boxes1), _box_area(boxes2)
+    lt, rb = np.maximum(boxes1[:, None, :2], boxes2[:, :2]), np.minimum(boxes1[:, None, 2:], boxes2[:, 2:])
+    wh = np.clip(rb - lt, a_min=0, a_max=None)
+    inter = wh[:, :, 0] * wh[:, :, 1]
+    union = area1[:, None] + area2 - inter
+    return inter, union
+
+  inter, union = _box_inter_union(boxes1, boxes2)
+  return inter / union
+
+def generate_anchors(input_size:tuple[int, int], scales:Optional[tuple[Tensor, ...]]=None, aspect_ratios:Optional[tuple[Tensor, ...]]=None) -> list[np.ndarray]:
+  def _compute_grid_sizes(input_size:tuple[int, int]) -> np.ndarray:
+    return np.ceil(np.array(input_size)[None, :] / 2 ** np.arange(3, 8)[:, None])
+
+  scales = tuple((i, int(i * 2 ** (1/3)), int(i * 2 ** (2/3))) for i in 2 ** np.arange(5, 10)) if scales is None else scales
+  aspect_ratios = ((0.5, 1.0, 2.0),) * len(scales) if aspect_ratios is None else aspect_ratios
+  aspect_ratios = tuple(ar for ar in aspect_ratios)
+  grid_sizes = _compute_grid_sizes(input_size)
+
+  assert len(scales) == len(aspect_ratios) == len(grid_sizes), "scales, aspect_ratios, and grid_sizes must have the same length"
+
+  anchors = []
+  for s, ar, gs in zip(scales, aspect_ratios, grid_sizes):
+    s, ar = np.array(s), np.array(ar)
+    h_ratios = np.sqrt(ar)
+    w_ratios = 1 / h_ratios
+    ws = (w_ratios[:, None] * s[None, :]).reshape(-1)
+    hs = (h_ratios[:, None] * s[None, :]).reshape(-1)
+    base_anchors = (np.stack([-ws, -hs, ws, hs], axis=1) / 2).round()
+    stride_h, stride_w = input_size[0] // gs[0], input_size[1] // gs[1]
+    shifts_x, shifts_y = np.meshgrid(np.arange(gs[1]) * stride_w, np.arange(gs[0]) * stride_h)
+    shifts_x, shifts_y = shifts_x.reshape(-1), shifts_y.reshape(-1)
+    shifts = np.stack([shifts_x, shifts_y, shifts_x, shifts_y], axis=1, dtype=np.float32)
+    anchors.append((shifts[:, None] + base_anchors[None, :]).reshape(-1, 4))
+
+  return anchors
+
+
+class BoxCoder(object):
+  def __init__(self, weights, bbox_xform_clip=math.log(1000. / 16), apply_to_remove=True):
+    self.weights = weights
+    self.bbox_xform_clip = bbox_xform_clip
+    self.apply_to_remove = apply_to_remove
+
+  def encode(self, reference_boxes, proposals):
+    TO_REMOVE = self.apply_to_remove  # TODO remove
+    ex_widths = proposals[..., 2] - proposals[..., 0] + TO_REMOVE
+    ex_heights = proposals[..., 3] - proposals[..., 1] + TO_REMOVE
+    ex_ctr_x = proposals[..., 0] + 0.5 * ex_widths
+    ex_ctr_y = proposals[..., 1] + 0.5 * ex_heights
+
+    gt_widths = reference_boxes[..., 2] - reference_boxes[..., 0] + TO_REMOVE
+    gt_heights = reference_boxes[..., 3] - reference_boxes[..., 1] + TO_REMOVE
+    gt_ctr_x = reference_boxes[..., 0] + 0.5 * gt_widths
+    gt_ctr_y = reference_boxes[..., 1] + 0.5 * gt_heights
+
+    wx, wy, ww, wh = self.weights
+    targets_dx = wx * (gt_ctr_x - ex_ctr_x) / ex_widths
+    targets_dy = wy * (gt_ctr_y - ex_ctr_y) / ex_heights
+    targets_dw = ww * Tensor.log(gt_widths / ex_widths)
+    targets_dh = wh * Tensor.log(gt_heights / ex_heights)
+
+    targets = Tensor.stack(targets_dx, targets_dy, targets_dw, targets_dh, dim=-1)
+    return targets
+
+  def decode(self, rel_codes, boxes):
+    boxes = boxes.cast(rel_codes.dtype)
+    rel_codes = rel_codes
+
+    TO_REMOVE = self.apply_to_remove  # TODO remove
+    widths = boxes[:, 2] - boxes[:, 0] + TO_REMOVE
+    heights = boxes[:, 3] - boxes[:, 1] + TO_REMOVE
+    ctr_x = boxes[:, 0] + 0.5 * widths
+    ctr_y = boxes[:, 1] + 0.5 * heights
+
+    wx, wy, ww, wh = self.weights
+    dx = rel_codes[:, 0::4] / wx
+    dy = rel_codes[:, 1::4] / wy
+    dw = rel_codes[:, 2::4] / ww
+    dh = rel_codes[:, 3::4] / wh
+
+    # Prevent sending too large values into Tensor.exp()
+    dw = dw.clip(min_=dw.min(), max_=self.bbox_xform_clip)
+    dh = dh.clip(min_=dh.min(), max_=self.bbox_xform_clip)
+
+    pred_ctr_x = dx * widths[:, None] + ctr_x[:, None]
+    pred_ctr_y = dy * heights[:, None] + ctr_y[:, None]
+    pred_w = dw.exp() * widths[:, None]
+    pred_h = dh.exp() * heights[:, None]
+    x = pred_ctr_x - 0.5 * pred_w
+    y = pred_ctr_y - 0.5 * pred_h
+    w = pred_ctr_x + 0.5 * pred_w - 1
+    h = pred_ctr_y + 0.5 * pred_h - 1
+    pred_boxes = Tensor.stack(x, y, w, h).permute(1,2,0).reshape(rel_codes.shape[0], rel_codes.shape[1])
+    return pred_boxes
@@ -53,7 +53,9 @@ class EmbeddingBert(nn.Embedding):
    arange_shp, weight_shp, big_shp = (1, 1, self.vocab_sz, 1), (1, 1, self.vocab_sz, self.embed_sz), idx.shape+(self.vocab_sz, self.embed_sz,)
    if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, requires_grad=False, device=self.weight.device).reshape(arange_shp)
    arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1,)).expand(big_shp), self.weight.cast(dtypes.default_float).reshape(weight_shp).expand(big_shp)
-    return (arange == idx).mul(vals).sum(2, acc_dtype=vals.dtype)
+    # TODO: contiguous() here because the embedding dropout creates different asts on each device, and search becomes very slow.
+    # Should fix with fixing random ast on multi device, and fuse arange to make embedding fast.
+    return (arange == idx).mul(vals).sum(2, dtype=vals.dtype).contiguous()

 class LayerNormBert:
  def __init__(self, normalized_shape:Union[int, tuple[int, ...]], eps:float=1e-12, elementwise_affine:bool=True):
@@ -66,3 +68,62 @@ class LayerNormBert:
    xn = x.cast(dtypes.float32).layernorm(eps=self.eps, axis=self.axis).cast(x.dtype)
    if not self.elementwise_affine: return xn
    return (xn * self.weight.cast(dtypes.default_float) + self.bias.cast(dtypes.default_float))
+
+class FrozenBatchNorm2dRetinaNet(nn.BatchNorm2d):
+  def __init__(self, sz:int, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1):
+    self.eps, self.track_running_stats, self.momentum = eps, track_running_stats, momentum
+
+    self.weight = Tensor.ones(sz, dtype=dtypes.float32, requires_grad=False) if affine else None
+    self.bias = Tensor.zeros(sz, dtype=dtypes.float32, requires_grad=False) if affine else None
+
+    if track_running_stats: self.running_mean, self.running_var = Tensor.zeros(sz, dtype=dtypes.float32, requires_grad=False), Tensor.ones(sz, dtype=dtypes.float32, requires_grad=False)
+    self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.long, requires_grad=False)
+
+  def __call__(self, x:Tensor) -> Tensor:
+    batch_mean, batch_var = super().calc_stats(x.cast(dtypes.float32))
+    if self.track_running_stats and Tensor.training:
+      self.running_mean.assign((1-self.momentum) * self.running_mean + self.momentum * batch_mean.detach().cast(self.running_mean.dtype))
+      self.running_var.assign((1-self.momentum) * self.running_var + self.momentum * x.numel()/(x.numel()-x.shape[1]) * batch_var.detach().cast(self.running_var.dtype))
+      self.num_batches_tracked += 1
+    return x.cast(dtypes.float32).batchnorm(self.weight, self.bias, batch_mean, batch_var.add(self.eps).rsqrt()).cast(x.dtype)
+
+class Conv2dNormalRetinaNet(nn.Conv2d):
+  def __init__(self, in_channels:int, out_channels:int, kernel_size:int|tuple[int, ...],
+               stride:int=1, padding:int|tuple[int, ...]|str=0, dilation:int=1, groups:int=1,
+               bias:bool=True, prior_prob:float|None=None):
+    super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
+    self.weight = Tensor.normal(*self.weight.shape, std=0.01, dtype=dtypes.float32)
+    if bias:
+      if prior_prob:
+        prior_prob = Tensor(prior_prob, device=self.bias.device, dtype=dtypes.float32).expand(*self.bias.shape)
+        self.bias = -(((1 - prior_prob) / prior_prob).log())
+      else: self.bias = Tensor.zeros_like(self.bias, dtype=dtypes.float32)
+
+  def __call__(self, x:Tensor) -> Tensor:
+    return x.conv2d(self.weight.cast(dtypes.default_float), self.bias.cast(dtypes.default_float) if self.bias is not None else None,
+                    groups=self.groups, stride=self.stride, padding=self.padding)
+
+class Conv2dKaimingUniformRetinaNet(nn.Conv2d):
+  def __init__(self, in_channels:int, out_channels:int, kernel_size:int|tuple[int, ...],
+               stride:int=1, padding:int|tuple[int, ...]|str=0, dilation:int=1, groups:int=1,
+               bias:bool=True):
+    super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
+    self.weight = Tensor.kaiming_uniform(*self.weight.shape, a=1, dtype=dtypes.float32)
+    if bias: self.bias = Tensor.zeros_like(self.bias, dtype=dtypes.float32)
+
+  def __call__(self, x:Tensor) -> Tensor:
+    return x.conv2d(self.weight.cast(dtypes.default_float), self.bias.cast(dtypes.default_float) if self.bias is not None else None,
+                    groups=self.groups, stride=self.stride, padding=self.padding)
+
+class Conv2dRetinaNet(nn.Conv2d):
+  def __init__(self, in_channels:int, out_channels:int, kernel_size:int|tuple[int, ...],
+               stride:int=1, padding:int|tuple[int, ...]|str=0, dilation:int=1, groups:int=1,
+               bias:bool=True):
+    super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
+    scale = 1 / math.sqrt(in_channels * prod(self.kernel_size))
+    self.weight = Tensor.uniform(out_channels, in_channels//groups, *self.kernel_size, low=-scale, high=scale, dtype=dtypes.float32)
+    self.bias: Tensor|None = Tensor.uniform(out_channels, low=-scale, high=scale, dtype=dtypes.float32) if bias else None
+
+  def __call__(self, x:Tensor) -> Tensor:
+    return x.conv2d(self.weight.cast(dtypes.default_float), self.bias.cast(dtypes.default_float) if self.bias is not None else None,
+                    groups=self.groups, stride=self.stride, dilation=self.dilation, padding=self.padding)
@@ -1,6 +1,29 @@
 from examples.mlperf.metrics import dice_score
+from tinygrad import Tensor

 def dice_ce_loss(pred, tgt):
  ce = pred.permute(0, 2, 3, 4, 1).sparse_categorical_crossentropy(tgt.squeeze(1))
  dice = (1.0 - dice_score(pred, tgt, argmax=False, to_one_hot_x=False)).mean()
  return (dice + ce) / 2
+
+def sigmoid_focal_loss(pred:Tensor, tgt:Tensor, alpha:float=0.25, gamma:float=2.0, reduction:str="none") -> Tensor:
+  assert reduction in ["mean", "sum", "none"], f"unsupported reduction {reduction}"
+  p, ce_loss = pred.sigmoid(), pred.binary_crossentropy_logits(tgt, reduction="none")
+  p_t = p * tgt + (1 - p) * (1 - tgt)
+  loss = ce_loss * ((1 - p_t) ** gamma)
+
+  if alpha >= 0:
+    alpha_t = alpha * tgt + (1 - alpha) * (1 - tgt)
+    loss = loss * alpha_t
+
+  if reduction == "mean": loss = loss.mean()
+  elif reduction == "sum": loss = loss.sum()
+  return loss
+
+def l1_loss(pred:Tensor, tgt:Tensor, reduction:str="none") -> Tensor:
+  assert reduction in ["mean", "sum", "none"], f"unsupported reduction {reduction}"
+  loss = (pred - tgt).abs()
+
+  if reduction == "mean": loss = loss.mean()
+  elif reduction == "sum": loss = loss.sum()
+  return loss
@@ -81,47 +81,43 @@ def eval_unet3d():

 def eval_retinanet():
  # RetinaNet with ResNeXt50_32X4D
+  from examples.mlperf.dataloader import batch_load_retinanet
+  from extra.datasets.openimages import normalize, download_dataset, BASEDIR
  from extra.models.resnet import ResNeXt50_32X4D
  from extra.models.retinanet import RetinaNet
-  mdl = RetinaNet(ResNeXt50_32X4D())
-  mdl.load_from_pretrained()
-
-  input_mean = Tensor([0.485, 0.456, 0.406]).reshape(1, -1, 1, 1)
-  input_std = Tensor([0.229, 0.224, 0.225]).reshape(1, -1, 1, 1)
-  def input_fixup(x):
-    x = x.permute([0,3,1,2]) / 255.0
-    x -= input_mean
-    x /= input_std
-    return x
-
-  from extra.datasets.openimages import download_dataset, iterate, BASEDIR
  from pycocotools.coco import COCO
  from pycocotools.cocoeval import COCOeval
  from contextlib import redirect_stdout
-  coco = COCO(download_dataset(base_dir:=getenv("BASE_DIR", BASEDIR), 'validation'))
+  tlog("imports")
+
+  mdl = RetinaNet(ResNeXt50_32X4D())
+  mdl.load_from_pretrained()
+  tlog("loaded models")
+
+  coco = COCO(download_dataset(base_dir:=getenv("BASEDIR", BASEDIR), 'validation'))
  coco_eval = COCOeval(coco, iouType="bbox")
  coco_evalimgs, evaluated_imgs, ncats, narea = [], [], len(coco_eval.params.catIds), len(coco_eval.params.areaRng)
+  tlog("loaded dataset")

-  from tinygrad.engine.jit import TinyJit
-  mdlrun = TinyJit(lambda x: mdl(input_fixup(x)).realize())
-
-  n, bs = 0, 8
+  iterator = batch_load_retinanet(coco, True, Path(base_dir), getenv("BS", 8), shuffle=False)
+  def data_get():
+    x, img_ids, img_sizes, cookie = next(iterator)
+    return x.to(Device.DEFAULT).realize(), img_ids, img_sizes, cookie
+  n = 0
+  proc = data_get()
+  tlog("loaded initial data")
  st = time.perf_counter()
-  for x, targets in iterate(coco, base_dir, bs):
-    dat = Tensor(x.astype(np.float32))
-    mt = time.perf_counter()
-    if dat.shape[0] == bs:
-      outs = mdlrun(dat).numpy()
-    else:
-      mdlrun._jit_cache = []
-      outs =  mdl(input_fixup(dat)).numpy()
-    et = time.perf_counter()
-    predictions = mdl.postprocess_detections(outs, input_size=dat.shape[1:3], orig_image_sizes=[t["image_size"] for t in targets])
-    ext = time.perf_counter()
-    n += len(targets)
-    print(f"[{n}/{len(coco.imgs)}] == {(mt-st)*1000:.2f} ms loading data, {(et-mt)*1000:.2f} ms to run model, {(ext-et)*1000:.2f} ms for postprocessing")
-    img_ids = [t["image_id"] for t in targets]
-    coco_results  = [{"image_id": targets[i]["image_id"], "category_id": label, "bbox": box.tolist(), "score": score}
+  while proc is not None:
+    GlobalCounters.reset()
+    proc = (mdl(normalize(proc[0])), proc[1], proc[2], proc[3])
+    run = time.perf_counter()
+    # load the next data here
+    try: next_proc = data_get()
+    except StopIteration: next_proc = None
+    nd = time.perf_counter()
+    predictions, img_ids = mdl.postprocess_detections(proc[0].numpy(), orig_image_sizes=proc[2]), proc[1]
+    pd = time.perf_counter()
+    coco_results  = [{"image_id": img_ids[i], "category_id": label, "bbox": box.tolist(), "score": score}
      for i, prediction in enumerate(predictions) for box, score, label in zip(*prediction.values())]
    with redirect_stdout(None):
      coco_eval.cocoDt = coco.loadRes(coco_results)
@@ -129,13 +125,18 @@ def eval_retinanet():
      coco_eval.evaluate()
    evaluated_imgs.extend(img_ids)
    coco_evalimgs.append(np.array(coco_eval.evalImgs).reshape(ncats, narea, len(img_ids)))
-    st = time.perf_counter()
+    n += len(proc[0])
+    et = time.perf_counter()
+    tlog(f"****** {(run-st)*1000:7.2f} ms to enqueue, {(et-run)*1000:7.2f} ms to realize ({(nd-run)*1000:7.2f} ms fetching, {(pd-run)*1000:4.2f} ms postprocess_detections). {(len(proc))/(et-st):8.2f} examples/sec. {GlobalCounters.global_ops*1e-12/(et-st):5.2f} TFLOPS")
+    st = et
+    proc, next_proc = next_proc, None

  coco_eval.params.imgIds = evaluated_imgs
  coco_eval._paramsEval.imgIds = evaluated_imgs
  coco_eval.evalImgs = list(np.concatenate(coco_evalimgs, -1).flatten())
  coco_eval.accumulate()
  coco_eval.summarize()
+  tlog("done")

 def eval_rnnt():
  # RNN-T
@@ -1,11 +1,11 @@
-import os, time, math, functools
+import os, time, math, functools, random
 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, FUSE_CONV_BW
+from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, FUSE_CONV_BW, Profiling
 from tinygrad.nn.state import get_parameters, get_state_dict, safe_load, safe_save
-from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup
+from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam

 from extra.lr_scheduler import LRSchedulerGroup
 from examples.mlperf.helpers import get_training_state, load_training_state
@@ -273,7 +273,7 @@ def train_resnet():
      else:
        it = iter(tqdm(batch_load_resnet(batch_size=EVAL_BS, val=True, shuffle=False, pad_first_batch=True), total=steps_in_val_epoch))
        i, proc = 0, data_get(it)
-        
+
      prev_cookies = []
      while proc is not None:
        GlobalCounters.reset()
@@ -343,8 +343,347 @@ def train_resnet():
        safe_save(get_training_state(model, optimizer_group, scheduler_group), fn)

 def train_retinanet():
-  # TODO: Retinanet
-  pass
+  from contextlib import redirect_stdout
+  from examples.mlperf.dataloader import batch_load_retinanet
+  from examples.mlperf.initializers import FrozenBatchNorm2dRetinaNet, Conv2dNormalRetinaNet, Conv2dKaimingUniformRetinaNet, Linear, Conv2dRetinaNet
+  from extra.datasets.openimages import MLPERF_CLASSES, BASEDIR, download_dataset, normalize, get_dataset_count
+  from extra.models import resnet, retinanet
+  from pycocotools.coco import COCO
+  from pycocotools.cocoeval import COCOeval
+  from tinygrad.helpers import colored
+  from typing import Iterator
+
+  import numpy as np
+
+  config, target_metric = {}, 0.34
+
+  config["SEED"] = SEED = getenv("SEED", random.SystemRandom().randint(0, 2**32 - 1))
+  Tensor.manual_seed(SEED)
+
+  NUM_CLASSES = len(MLPERF_CLASSES)
+  BASEDIR = getenv("BASEDIR", BASEDIR)
+  BENCHMARK = getenv("BENCHMARK")
+  INITMLPERF = getenv("INITMLPERF")
+  RUNMLPERF = getenv("RUNMLPERF")
+
+  if getenv("LOGMLPERF"):
+    from mlperf_logging import mllog
+    import mlperf_logging.mllog.constants as mllog_constants
+
+    mllog.config(filename=f"result_retinanet_{SEED}.log")
+    mllog.config(root_dir=Path(__file__).parents[3].as_posix())
+    MLLOGGER = mllog.get_mllogger()
+    MLLOGGER.logger.propagate = False
+
+    if INITMLPERF:
+      assert BENCHMARK, "BENCHMARK must be set for INITMLPERF"
+      MLLOGGER.event(key=mllog_constants.SUBMISSION_ORG, value="tinycorp")
+      MLLOGGER.event(key=mllog_constants.SUBMISSION_PLATFORM, value=getenv("SUBMISSION_PLATFORM", "tinybox"))
+      MLLOGGER.event(key=mllog_constants.SUBMISSION_DIVISION, value=mllog_constants.CLOSED)
+      MLLOGGER.event(key=mllog_constants.SUBMISSION_STATUS, value=mllog_constants.ONPREM)
+
+      MLLOGGER.event(key=mllog_constants.SUBMISSION_BENCHMARK, value=mllog_constants.RETINANET)
+
+      diskcache_clear()
+      MLLOGGER.event(key=mllog_constants.CACHE_CLEAR, value=True)
+      MLLOGGER.start(key=mllog_constants.INIT_START)
+
+    if RUNMLPERF:
+      MLLOGGER.start(key=mllog_constants.RUN_START)
+      MLLOGGER.event(key=mllog_constants.SEED, value=SEED)
+  else:
+    MLLOGGER = None
+
+  config["gpus"] = GPUS = [f"{Device.DEFAULT}:{i}" for i in range(getenv("GPUS", 6))]
+
+  for x in GPUS: Device[x]
+  print(f"training on {GPUS}")
+
+  def _freeze_backbone_layers(backbone:resnet.ResNet, trainable_layers:int):
+    layers_to_train = ["layer4", "layer3", "layer2", "layer1", "conv1"][:trainable_layers]
+    for k, v in get_state_dict(backbone).items():
+      if all([not k.startswith(layer) for layer in layers_to_train]):
+        v.requires_grad = False
+
+  def _data_get(it:Iterator[tuple[Tensor, ...]], val:bool=False):
+    if val:
+      x, img_ids, img_sizes, cookie = next(it)
+      return x.shard(GPUS, axis=0), img_ids, img_sizes, cookie
+
+    x, y_boxes, y_labels, matches, anchors, cookie = next(it)
+    return x.shard(GPUS, axis=0), y_boxes.shard(GPUS, axis=0), y_labels.shard(GPUS, axis=0), matches.shard(GPUS, axis=0), anchors.shard(GPUS, axis=0), cookie
+
+  def _fake_data_get(bs:int, val:bool=False):
+    x = Tensor.empty(bs, 800, 800, 3, dtype=dtypes.uint8)
+    if val:
+      img_ids, img_sizes = [0] * bs, [(800, 800)] * bs
+      return x.shard(GPUS, axis=0), img_ids, img_sizes, None
+
+    y_boxes = Tensor.empty(bs, 120087, 4, dtype=dtypes.float32)
+    y_labels = Tensor.empty(bs, 120087, dtype=dtypes.int64)
+    matches = Tensor.empty(bs, 120087, dtype=dtypes.int64)
+    anchors = Tensor.empty(bs, 120087, 4, dtype=dtypes.float64)
+    return x.shard(GPUS, axis=0), y_boxes.shard(GPUS, axis=0), y_labels.shard(GPUS, axis=0), matches.shard(GPUS, axis=0), anchors.shard(GPUS, axis=0), None
+
+  @TinyJit
+  def _train_step(model, optim, loss_scaler, x, **kwargs):
+    optim.zero_grad()
+
+    losses = model(normalize(x, GPUS), **kwargs)
+    loss = sum(losses.values())
+
+    (loss * loss_scaler).backward()
+    for t in optim.params: t.grad = t.grad / loss_scaler
+
+    optim.step()
+
+    return loss.realize(), losses
+
+  @TinyJit
+  def _eval_step(model, x, **kwargs):
+    out = model(normalize(x, GPUS), **kwargs)
+    # reassemble on GPUS[0] before sending back to CPU for speed
+    return out.to(GPUS[0]).realize()
+
+  # ** hyperparameters **
+  config["BS"] = BS = getenv("BS", 16 * len(GPUS) if dtypes.default_float == dtypes.float16 else 12 * len(GPUS))
+  config["EVAL_BS"] = EVAL_BS = getenv("EVAL_BS", BS)
+  config["EPOCHS"] = EPOCHS = getenv("EPOCHS", 4)
+  config["TRAIN_BEAM"] = TRAIN_BEAM = getenv("TRAIN_BEAM", BEAM.value)
+  config["EVAL_BEAM"] = EVAL_BEAM = getenv("EVAL_BEAM", BEAM.value)
+  config["LR"] = lr = getenv("LR", 9.5e-5 * (BS / 96))
+  config["LOSS_SCALER"] = loss_scaler = getenv("LOSS_SCALER", 2**11 if dtypes.default_float == dtypes.float16 else 1.0)
+  config["DEFAULT_FLOAT"] = dtypes.default_float.name
+  config["EVAL_FREQ"] = eval_freq = getenv("EVAL_FREQ", 1)
+
+  # ** initialize wandb **
+  if (WANDB:=getenv("WANDB")):
+    import wandb
+    wandb.init(config=config, project="MLPerf-RetinaNet")
+
+  # ** model initializers **
+  resnet.BatchNorm = FrozenBatchNorm2dRetinaNet
+  resnet.Linear = Linear
+  resnet.Conv2d = Conv2dRetinaNet
+
+  retinanet.ConvHead = Conv2dNormalRetinaNet
+  retinanet.ConvClassificationHeadLogits = functools.partial(Conv2dNormalRetinaNet, prior_prob=0.01)
+  retinanet.ConvFPN = Conv2dKaimingUniformRetinaNet
+
+  # ** model setup **
+  backbone = resnet.ResNeXt50_32X4D(num_classes=None)
+  if RUNMLPERF:
+    backbone.load_from_pretrained()
+  _freeze_backbone_layers(backbone, 3)
+
+  model = retinanet.RetinaNet(backbone, num_classes=NUM_CLASSES)
+  params = get_parameters(model)
+
+  if not RUNMLPERF:
+    # for init, zero out all weights
+    for p in params:
+      p = p.assign(Tensor.zeros_like(p).contiguous()).realize()
+
+  if len(GPUS) > 1:
+    for p in params: p.to_(GPUS)
+
+  step_times, start_epoch = [], 0
+
+  # ** optimizer **
+  optim = Adam(params, lr=lr)
+
+  # ** dataset **
+  config["STEPS_IN_TRAIN_EPOCH"] = steps_in_train_epoch = round_up(get_dataset_count((base_dir_path:=Path(BASEDIR)), False), BS) // BS
+  config["STEPS_IN_VAL_EPOCH"] = steps_in_val_epoch = (round_up(get_dataset_count(base_dir_path, True), EVAL_BS) // EVAL_BS)
+
+  # log mlperf hparams
+  if MLLOGGER:
+    if RUNMLPERF:
+      MLLOGGER.event(key=mllog_constants.GLOBAL_BATCH_SIZE, value=config["BS"])
+      MLLOGGER.event(key=mllog_constants.TRAIN_SAMPLES, value=config["STEPS_IN_TRAIN_EPOCH"])
+      MLLOGGER.event(key=mllog_constants.EVAL_SAMPLES, value=config["STEPS_IN_VAL_EPOCH"])
+      MLLOGGER.event(key=mllog_constants.EPOCH_COUNT, value=config["EPOCHS"])
+      MLLOGGER.event(key=mllog_constants.FIRST_EPOCH_NUM, value=start_epoch)
+
+      MLLOGGER.event(key=mllog_constants.OPT_NAME, value=mllog_constants.ADAM)
+      MLLOGGER.event(key=mllog_constants.OPT_BASE_LR, value=config["LR"])
+      MLLOGGER.event(key=mllog_constants.OPT_WEIGHT_DECAY, value=0)
+      MLLOGGER.event(key=mllog_constants.OPT_LR_WARMUP_EPOCHS, value=0)
+      MLLOGGER.event(key=mllog_constants.OPT_LR_WARMUP_FACTOR, value=0)
+      MLLOGGER.event(key=mllog_constants.GRADIENT_ACCUMULATION_STEPS, value=1)
+
+  if RUNMLPERF:
+    train_dataset = COCO(download_dataset(BASEDIR, "train"))
+    val_dataset = COCO(download_dataset(BASEDIR, "validation"))
+    coco_val = COCOeval(cocoGt=val_dataset, iouType="bbox")
+
+  print(f"training with batch size {BS} for {EPOCHS} epochs")
+
+  for e in range(start_epoch, EPOCHS):
+    # ** training loop **
+    if MLLOGGER and RUNMLPERF:
+      MLLOGGER.start(key=mllog_constants.EPOCH_START, value=e + 1, metadata={"epoch_num": e + 1})
+
+    BEAM.value = TRAIN_BEAM
+
+    if not RUNMLPERF:
+      i, proc = 0, _fake_data_get(BS)
+    else:
+      train_dataloader = batch_load_retinanet(train_dataset, False, base_dir_path, batch_size=BS, seed=SEED)
+      it = iter(tqdm(train_dataloader, total=steps_in_train_epoch, desc=f"epoch {e + 1}", disable=BENCHMARK))
+      i, proc = 0, _data_get(it)
+
+    prev_cookies = []
+    st = time.perf_counter()
+
+    while proc is not None:
+      GlobalCounters.reset()
+
+      x, y_bboxes, y_labels, matches, anchors, proc = proc
+      loss, losses = _train_step(model, optim, loss_scaler, x, labels=y_labels, matches=matches, anchors=anchors, bboxes=y_bboxes)
+
+      pt = time.perf_counter()
+
+      if len(prev_cookies) == getenv("STORE_COOKIES", 1): prev_cookies = []  # free previous cookies after gpu work has been enqueued
+      try:
+        if not RUNMLPERF:
+          next_proc = _fake_data_get(BS)
+        else:
+          next_proc = _data_get(it)
+      except StopIteration:
+        next_proc = None
+
+      dt = time.perf_counter()
+
+      device_str = loss.device if isinstance(loss.device, str) else f"{loss.device[0]} * {len(loss.device)}"
+      loss = loss.item()
+
+      cl = time.perf_counter()
+      if BENCHMARK: step_times.append(cl - st)
+
+      if not math.isfinite(loss):
+        print("loss is nan")
+        return
+
+      tqdm.write(
+        f"{i:5} {((cl - st)) * 1000.0:7.2f} ms run, {(pt - st) * 1000.0:7.2f} ms python, {(dt - pt) * 1000.0:6.2f} ms fetch data, "
+        f"{(cl - dt) * 1000.0:7.2f} ms {device_str}, {loss:5.2f} loss, {losses['classification_loss'].item():5.4f} classification loss, {losses['regression_loss'].item():5.4f} regression loss, "
+        f"{optim.lr.numpy()[0]:.6f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, {GlobalCounters.global_ops * 1e-9 / (cl - st):9.2f} GFLOPS"
+      )
+
+      if WANDB:
+        wandb.log({"lr": optim.lr.numpy(), "train/loss": loss, "train/classification_loss": losses["classification_loss"].item(), "train/regression_loss": losses["regression_loss"].item(),
+                  "train/step_time": cl - st, "train/python_time": pt - st, "train/data_time": dt - pt, "train/cl_time": cl - dt,
+                  "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": e + (i + 1) / steps_in_train_epoch})
+
+      st = cl
+      prev_cookies.append(proc)
+      proc, next_proc = next_proc, None  # return old cookie
+      i += 1
+
+      if i == BENCHMARK:
+        assert not math.isnan(loss)
+        median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]  # in seconds
+        estimated_total_minutes = int(median_step_time * steps_in_train_epoch * EPOCHS / 60)
+        print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
+        print(f"epoch global_ops: {steps_in_train_epoch * GlobalCounters.global_ops:_}, "
+              f"epoch global_mem: {steps_in_train_epoch * GlobalCounters.global_mem:_}")
+        # if we are doing beam search, run the first eval too
+        if (TRAIN_BEAM or EVAL_BEAM) and e == start_epoch: break
+        return
+
+    if MLLOGGER and RUNMLPERF:
+      MLLOGGER.event(key=mllog_constants.EPOCH_STOP, value=e + 1, metadata={"epoch_num": e + 1})
+
+    # ** eval loop **
+    if (e + 1) % eval_freq == 0:
+      if MLLOGGER and RUNMLPERF:
+        MLLOGGER.start(key=mllog_constants.EVAL_START, value=e + 1, metadata={"epoch_num": e + 1})
+
+      BEAM.value = EVAL_BEAM
+
+      if getenv("RESET_STEP", 1): _train_step.reset()
+
+      with Tensor.train(mode=False), Tensor.test():
+        if not RUNMLPERF:
+          i, proc = 0, _fake_data_get(EVAL_BS, val=(val:=True))
+        else:
+          val_dataloader = batch_load_retinanet(val_dataset, (val:=True), Path(BASEDIR), batch_size=EVAL_BS, shuffle=False, seed=SEED)
+          it = iter(tqdm(val_dataloader, total=steps_in_val_epoch))
+          i, proc = 0, _data_get(it, val=val)
+          val_img_ids, val_imgs, ncats, narea = [], [], len(coco_val.params.catIds), len(coco_val.params.areaRng)
+
+        eval_times, prev_cookies = [], []
+
+        while proc is not None:
+          GlobalCounters.reset()
+          st = time.time()
+
+          out, img_ids, img_sizes, proc = _eval_step(model, (x:=proc[0])).numpy(), proc[1], proc[2], proc[3]
+
+          if RUNMLPERF:
+            out = model.postprocess_detections(out, input_size=x.shape[1:3], orig_image_sizes=img_sizes)
+            coco_results  = [{"image_id": img_ids[i], "category_id": label, "bbox": box.tolist(), "score": score}
+              for i, prediction in enumerate(out) for box, score, label in zip(*prediction.values())]
+
+            with redirect_stdout(None):
+              coco_val.cocoDt = val_dataset.loadRes(coco_results)
+              coco_val.params.imgIds = img_ids
+              coco_val.evaluate()
+
+            val_img_ids.extend(img_ids)
+            val_imgs.append(np.array(coco_val.evalImgs).reshape(ncats, narea, len(img_ids)))
+
+          if len(prev_cookies) == getenv("STORE_COOKIES", 1): prev_cookies = []  # free previous cookies after gpu work has been enqueued
+          try:
+            if not RUNMLPERF:
+              next_proc = _fake_data_get(EVAL_BS, val=val)
+            else:
+              next_proc = _data_get(it, val=val)
+          except StopIteration:
+            next_proc = None
+
+          prev_cookies.append(proc)
+          proc, next_proc = next_proc, None
+          i += 1
+
+          et = time.time()
+          eval_times.append(et - st)
+
+          if i == BENCHMARK:
+            # assume INITMLPERF has BENCHMARK set
+            if MLLOGGER and INITMLPERF:
+              MLLOGGER.event(key=mllog_constants.INIT_STOP)
+            return
+
+        if getenv("RESET_STEP", 1): _eval_step.reset()
+        total_fw_time = sum(eval_times) / len(eval_times)
+
+        if RUNMLPERF:
+          coco_val.params.imgIds = val_img_ids
+          coco_val._paramsEval.imgIds = val_img_ids
+          coco_val.evalImgs = list(np.concatenate(val_imgs, -1).flatten())
+          coco_val.accumulate()
+          coco_val.summarize()
+
+          val_metric = coco_val.stats[0]
+
+          tqdm.write(f"eval time: {total_fw_time:.2f}, eval metric: {val_metric:.4f}")
+
+          if WANDB:
+            wandb.log({"eval/forward_time": total_fw_time, "eval/metric": val_metric, "epoch": e + 1})
+
+          if MLLOGGER:
+            MLLOGGER.event(key=mllog_constants.EVAL_ACCURACY, value=val_metric, metadata={"epoch_num": e + 1}, clear_line=True)
+            MLLOGGER.end(key=mllog_constants.EVAL_STOP, value=e + 1, metadata={"epoch_num": e + 1})
+
+          if val_metric >= target_metric:
+            print(colored(f"target metric reached: {val_metric:.2f}/{target_metric:.2f}", color="green"))
+
+            if MLLOGGER:
+              MLLOGGER.end(key=mllog_constants.RUN_STOP, metadata={"status": mllog_constants.SUCCESS})
+
+            break

 def train_unet3d():
  """
@@ -446,7 +785,7 @@ def train_unet3d():
    loss.backward()
    optim.step()
    return loss.realize()
-  
+
  @Tensor.train(mode=False)
  @Tensor.test()
  def eval_step(model, x, y):
@@ -455,7 +794,7 @@ def train_unet3d():
    loss = dice_ce_loss(y_hat, y)
    score = dice_score(y_hat, y)
    return loss.realize(), score.realize()
-  
+
  if WANDB: wandb.init(config=config, project=PROJ_NAME)

  step_times, start_epoch = [], 1
@@ -464,7 +803,7 @@ def train_unet3d():
  next_eval_at = start_eval_at

  print(f"Training on {GPUS}")
-  
+
  if BENCHMARK: print("Benchmarking UNet3D")
  else: print(f"Start evaluation at epoch {start_eval_at} and every {evaluate_every} epoch(s) afterwards")

@@ -551,7 +890,7 @@ def train_unet3d():

        if mean_dice >= TARGET_METRIC:
          is_successful = True
-          save_checkpoint(get_state_dict(model), f"./ckpts/unet3d.safe")
+          save_checkpoint(get_state_dict(model), "./ckpts/unet3d.safe")
        elif mean_dice < 1e-6:
          print("Model diverging. Aborting.")
          diverged = True
@@ -575,33 +914,41 @@ def train_rnnt():
 def train_step_bert(model, optimizer, scheduler, loss_scaler:float, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor,
                    masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor, GPUS):
  for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
-    t.shard_(GPUS, axis=0)
+    if len(GPUS) > 1: t.shard_(GPUS, axis=0)
+    else: t.to_(GPUS[0])
  optimizer.zero_grad()

  lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
  loss = model.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
  (loss * loss_scaler).backward()

-  global_norm = Tensor([0.0], dtype=dtypes.float32, device=optimizer[0].device).realize()
-  for p in optimizer.params: 
+  global_norm = Tensor([0.0], dtype=dtypes.float32, device=optimizer[0].device)
+  for p in optimizer.params:
    p.grad = p.grad / loss_scaler
    global_norm += p.grad.float().square().sum()
-  global_norm = global_norm.sqrt()
-  for p in optimizer.params: p.grad = (p.grad / Tensor.where(global_norm > 1.0, global_norm, 1.0)).cast(p.grad.dtype)
+  global_norm = global_norm.sqrt().contiguous()
+  for p in optimizer.params:
+    p.grad = (global_norm > 1.0).where((p.grad/global_norm).cast(p.grad.dtype), p.grad)

  optimizer.step()
  scheduler.step()
-  return loss.realize(), global_norm.realize()
+  # TODO: no to("CPU") here because it blocks and messes the python time
+  Tensor.realize(loss, global_norm, optimizer.optimizers[0].lr)
+  return loss, global_norm, optimizer.optimizers[0].lr

@TinyJit
 def eval_step_bert(model, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor, masked_positions:Tensor, masked_lm_ids:Tensor,
                   masked_lm_weights:Tensor, next_sentence_labels:Tensor, GPUS):
  for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
-    t.shard_(GPUS, axis=0)
+    if len(GPUS) > 1: t.shard_(GPUS, axis=0)
+    else: t.to_(GPUS[0])
  lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
  masked_lm_accuracy, seq_relationship_accuracy, masked_lm_loss, next_sentence_loss = \
    model.accuracy(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
-  return masked_lm_accuracy.realize(), seq_relationship_accuracy.realize(), masked_lm_loss.realize(), next_sentence_loss.realize()
+  for t in [masked_lm_accuracy, seq_relationship_accuracy, masked_lm_loss, next_sentence_loss]:
+    t.to_("CPU")
+  Tensor.realize(masked_lm_accuracy, seq_relationship_accuracy, masked_lm_loss, next_sentence_loss)
+  return masked_lm_accuracy, seq_relationship_accuracy, masked_lm_loss, next_sentence_loss

 def train_bert():
  # NOTE: pip install tensorflow, wandb required
@@ -630,7 +977,7 @@ def train_bert():
    MLLOGGER.logger.propagate = False

    if INITMLPERF:
-      assert BENCHMARK, f"BENCHMARK must be set for INITMLPERF"
+      assert BENCHMARK, "BENCHMARK must be set for INITMLPERF"
      MLLOGGER.event(key=mllog_constants.SUBMISSION_ORG, value="tinycorp")
      MLLOGGER.event(key=mllog_constants.SUBMISSION_PLATFORM, value=getenv("SUBMISSION_PLATFORM", "tinybox"))
      MLLOGGER.event(key=mllog_constants.SUBMISSION_DIVISION, value=mllog_constants.CLOSED)
@@ -651,9 +998,11 @@ def train_bert():
  # ** hyperparameters **
  BS                 = config["GLOBAL_BATCH_SIZE"]      = getenv("BS", 11 * len(GPUS) if dtypes.default_float in (dtypes.float16, dtypes.bfloat16) else 8 * len(GPUS))
  EVAL_BS            = config["EVAL_BS"]                = getenv("EVAL_BS", 1 * len(GPUS))
-  max_lr             = config["OPT_BASE_LEARNING_RATE"] = getenv("OPT_BASE_LEARNING_RATE", 0.00011 * math.sqrt(BS/66))
+  max_lr             = config["OPT_BASE_LEARNING_RATE"] = getenv("OPT_BASE_LEARNING_RATE", 0.000175 * math.sqrt(BS/96))
+  opt_lamb_beta_1    = config["OPT_LAMB_BETA_1"]        = getenv("OPT_LAMB_BETA_1", 0.9)
+  opt_lamb_beta_2    = config["OPT_LAMB_BETA_2"]        = getenv("OPT_LAMB_BETA_2", 0.999)

-  train_steps        = config["TRAIN_STEPS"]            = getenv("TRAIN_STEPS", 3630000 // BS)
+  train_steps        = config["TRAIN_STEPS"]            = getenv("TRAIN_STEPS", 3600000 // BS)
  warmup_steps       = config["NUM_WARMUP_STEPS"]       = getenv("NUM_WARMUP_STEPS", 1)
  max_eval_steps     = config["MAX_EVAL_STEPS"]         = getenv("MAX_EVAL_STEPS", (10000 + EVAL_BS - 1) // EVAL_BS) # EVAL_BS * MAX_EVAL_STEPS >= 10000
  eval_step_freq     = config["EVAL_STEP_FREQ"]         = getenv("EVAL_STEP_FREQ", int((math.floor(0.05 * (230.23 * BS + 3000000) / 25000) * 25000) / BS)) # Round down
@@ -662,7 +1011,7 @@ def train_bert():
  save_ckpt_dir      = config["SAVE_CKPT_DIR"]          = getenv("SAVE_CKPT_DIR", "./ckpts")
  init_ckpt          = config["INIT_CKPT_DIR"]          = getenv("INIT_CKPT_DIR", BASEDIR)

-  loss_scaler        = config["LOSS_SCALER"]            = getenv("LOSS_SCALER", 2.0**10 if dtypes.default_float == dtypes.float16 else 1.0)
+  loss_scaler        = config["LOSS_SCALER"]            = getenv("LOSS_SCALER", 2.0**11 if dtypes.default_float == dtypes.float16 else 1.0)
  decay              = config["DECAY"]                  = getenv("DECAY", 0.01)
  epsilon            = config["EPSILON"]                = getenv("EPSILON", 1e-6)
  poly_power         = config["POLY_POWER"]             = getenv("POLY_POWER", 1.0)
@@ -696,8 +1045,9 @@ def train_bert():
      p = p.assign(Tensor.zeros_like(p).contiguous()).realize()

  parameters = get_parameters(model)
-  for p in parameters:
-    p.to_(GPUS)
+  if len(GPUS) > 1:
+    for p in parameters:
+      p.to_(GPUS)

  # ** Log run config **
  for key, value in config.items(): print(f'HParam: "{key}": {value}')
@@ -705,8 +1055,8 @@ def train_bert():
  # ** Optimizer **
  parameters_no_wd = [v for k, v in get_state_dict(model).items() if "bias" in k or "LayerNorm" in k]
  parameters = [x for x in parameters if x not in set(parameters_no_wd)]
-  optimizer_wd = LAMB(parameters, lr=max_lr, eps=epsilon, weight_decay=decay, adam=False)
-  optimizer_no_wd = LAMB(parameters_no_wd, lr=max_lr, eps=epsilon, weight_decay=0.0, adam=False)
+  optimizer_wd = LAMB(parameters, lr=max_lr, b1=opt_lamb_beta_1, b2=opt_lamb_beta_2, eps=epsilon, weight_decay=decay, adam=False)
+  optimizer_no_wd = LAMB(parameters_no_wd, lr=max_lr, b1=opt_lamb_beta_1, b2=opt_lamb_beta_2, eps=epsilon, weight_decay=0.0, adam=False)
  optimizer_group = OptimizerGroup(optimizer_wd, optimizer_no_wd)

  # ** LR scheduler **
@@ -725,8 +1075,8 @@ def train_bert():
      MLLOGGER.event(key=mllog_constants.OPT_NAME, value="LAMB")
      MLLOGGER.event(key=mllog_constants.OPT_BASE_LR, value=config["OPT_BASE_LEARNING_RATE"])
      MLLOGGER.event(key=mllog_constants.OPT_LAMB_WEIGHT_DECAY, value=config["DECAY"])
-      MLLOGGER.event(key=mllog_constants.OPT_LAMB_BETA_1, value=optimizer_wd.b1)
-      MLLOGGER.event(key=mllog_constants.OPT_LAMB_BETA_2, value=optimizer_wd.b2)
+      MLLOGGER.event(key=mllog_constants.OPT_LAMB_BETA_1, value=config["OPT_LAMB_BETA_1"])
+      MLLOGGER.event(key=mllog_constants.OPT_LAMB_BETA_2, value=config["OPT_LAMB_BETA_2"])
      MLLOGGER.event(key=mllog_constants.OPT_LAMB_LR_DECAY_POLY_POWER, value=config["POLY_POWER"])
      MLLOGGER.event(key=mllog_constants.OPT_LAMB_EPSILON, value=config["EPSILON"])

@@ -769,54 +1119,57 @@ def train_bert():
      MLLOGGER.start(key=mllog_constants.EPOCH_START, value=i*BS, metadata={"epoch_num": i*BS})

  while train_data is not None and i < train_steps and not achieved:
-    Tensor.training = True
-    BEAM.value = TRAIN_BEAM
-    st = time.perf_counter()
-    GlobalCounters.reset()
-    loss, global_norm = train_step_bert(model, optimizer_group, scheduler_group, loss_scaler,
-      train_data["input_ids"], train_data["segment_ids"], train_data["input_mask"], train_data["masked_lm_positions"], \
-      train_data["masked_lm_ids"], train_data["masked_lm_weights"], train_data["next_sentence_labels"], GPUS)
+    if getenv("TRAIN", 1):
+      Tensor.training = True
+      BEAM.value = TRAIN_BEAM
+      st = time.perf_counter()
+      GlobalCounters.reset()
+      loss, global_norm, lr = train_step_bert(model, optimizer_group, scheduler_group, loss_scaler,
+        train_data["input_ids"], train_data["segment_ids"], train_data["input_mask"], train_data["masked_lm_positions"], \
+        train_data["masked_lm_ids"], train_data["masked_lm_weights"], train_data["next_sentence_labels"], GPUS)

-    pt = time.perf_counter()
+      pt = time.perf_counter()

-    try:
-      next_data = next(train_it)
-    except StopIteration:
-      next_data = None
+      try:
+        next_data = next(train_it)
+      except StopIteration:
+        next_data = None

-    dt = time.perf_counter()
+      dt = time.perf_counter()

-    device_str = loss.device if isinstance(loss.device, str) else f"{loss.device[0]} * {len(loss.device)}"
-    loss = loss.item()
+      device_str = parameters[0].device if isinstance(parameters[0].device, str) else f"{parameters[0].device[0]} * {len(parameters[0].device)}"
+      loss = loss.item()
+      assert not math.isnan(loss)
+      lr = lr.item()

-    cl = time.perf_counter()
-    if BENCHMARK: step_times.append(cl - st)
+      cl = time.perf_counter()
+      if BENCHMARK: step_times.append(cl - st)

-    tqdm.write(
-      f"{i:5} {((cl - st)) * 1000.0:7.2f} ms run, {(pt - st) * 1000.0:7.2f} ms python, {(dt - pt) * 1000.0:6.2f} ms fetch data, "
-      f"{(cl - dt) * 1000.0:7.2f} ms {device_str}, {loss:5.2f} loss, {optimizer_wd.lr.numpy()[0]:.6f} LR, "
-      f"{GlobalCounters.mem_used / 1e9:.2f} GB used, {GlobalCounters.global_ops * 1e-9 / (cl - st):9.2f} GFLOPS")
-    if WANDB:
-      wandb.log({"lr": optimizer_wd.lr.numpy(), "train/loss": loss, "train/global_norm": global_norm.item(), "train/step_time": cl - st,
-                  "train/python_time": pt - st, "train/data_time": dt - pt, "train/cl_time": cl - dt,
-                  "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*BS})
+      tqdm.write(
+        f"{i:5} {((cl - st)) * 1000.0:7.2f} ms run, {(pt - st) * 1000.0:7.2f} ms python, {(dt - pt) * 1000.0:6.2f} ms fetch data, "
+        f"{(cl - dt) * 1000.0:7.2f} ms {device_str}, {loss:5.2f} loss, {lr:.6f} LR, "
+        f"{GlobalCounters.mem_used / 1e9:.2f} GB used, {GlobalCounters.global_ops * 1e-9 / (cl - st):9.2f} GFLOPS")
+      if WANDB:
+        wandb.log({"lr": lr, "train/loss": loss, "train/global_norm": global_norm.item(), "train/step_time": cl - st,
+                    "train/python_time": pt - st, "train/data_time": dt - pt, "train/cl_time": cl - dt,
+                    "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*BS})

-    train_data, next_data = next_data, None
-    i += 1
+      train_data, next_data = next_data, None
+      i += 1

-    if i == BENCHMARK:
-      median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]  # in seconds
-      estimated_total_minutes = int(median_step_time * train_steps / 60)
-      print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
-      print(f"epoch global_ops: {train_steps * GlobalCounters.global_ops:_}, "
-            f"epoch global_mem: {train_steps * GlobalCounters.global_mem:_}")
+      if i == BENCHMARK:
+        median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]  # in seconds
+        estimated_total_minutes = int(median_step_time * train_steps / 60)
+        print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
+        print(f"epoch global_ops: {train_steps * GlobalCounters.global_ops:_}, "
+              f"epoch global_mem: {train_steps * GlobalCounters.global_mem:_}")

    # ** eval loop **
    if i % eval_step_freq == 0 or (BENCHMARK and i == BENCHMARK) or i == train_steps:
      if MLLOGGER and RUNMLPERF:
        MLLOGGER.start(key=mllog_constants.EVAL_START, value=None, metadata={"epoch_num": i*BS, "step_num": i})
-      if getenv("RESET_STEP", 0) or INITMLPERF: train_step_bert.reset()
-      else: train_step_bert.captured.free_intermediates()
+      if getenv("RESET_STEP"): train_step_bert.reset()
+      elif getenv("FREE_INTERMEDIATE", 1) and train_step_bert.captured is not None: train_step_bert.captured.free_intermediates()
      eval_lm_losses = []
      eval_clsf_losses = []
      eval_lm_accs = []
@@ -843,14 +1196,15 @@ def train_bert():
        et = time.time()
        eval_times.append(et - st)

-        if BENCHMARK and j == BENCHMARK:
+        if BENCHMARK and (j+1) == min(BENCHMARK, max_eval_steps):
          # assume INITMLPERF has BENCHMARK set
          if MLLOGGER and INITMLPERF:
            MLLOGGER.event(key=mllog_constants.INIT_STOP, value=None)
          return

-      if getenv("RESET_STEP", 0): eval_step_bert.reset()
-      else: eval_step_bert.captured.free_intermediates()
+      if getenv("RESET_STEP"): eval_step_bert.reset()
+      elif getenv("FREE_INTERMEDIATE", 1) and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()
+
      del eval_data
      avg_lm_loss = sum(eval_lm_losses) / len(eval_lm_losses)
      avg_clsf_loss = sum(eval_clsf_losses) / len(eval_clsf_losses)
@@ -863,7 +1217,7 @@ def train_bert():

      if WANDB:
        wandb.log({"eval/lm_loss": avg_lm_loss, "eval/clsf_loss": avg_clsf_loss, "eval/lm_accuracy": avg_lm_acc, \
-                    "eval/clsf_accuracy": avg_clsf_acc, "eval/forward_time": avg_fw_time})
+                    "eval/clsf_accuracy": avg_clsf_acc, "eval/forward_time": avg_fw_time, "epoch": (i+1)*BS})

      if MLLOGGER and RUNMLPERF:
        MLLOGGER.end(key=mllog_constants.EVAL_STOP, value=i*BS, metadata={"epoch_count": i*BS, "step_num": i, "samples_count": config["EVAL_BS"] * config["MAX_EVAL_STEPS"]})
@@ -890,11 +1244,14 @@ def train_bert():
        # stop once hitting the target
        break

+    # should not happen, BENCHMARK not properly terminated
+    if BENCHMARK: assert i < BENCHMARK, i
+
    if getenv("CKPT") and i % save_ckpt_freq == 0:
      if MLLOGGER and RUNMLPERF:
        if previous_step:
          MLLOGGER.end(key=mllog_constants.BLOCK_STOP, value=None, metadata={"first_epoch_num": 1, "epoch_num": 1, "first_step_num": i, "step_num": i, "step_count": i - previous_step})
-        MLLOGGER.start(key="checkpoint_start", value=None, metadata={"step_num" : i})
+        MLLOGGER.start(key="checkpoint_start", value=None, metadata={"step_num": i})
      if not os.path.exists(ckpt_dir := save_ckpt_dir): os.mkdir(ckpt_dir)
      if WANDB and wandb.run is not None:
        fn = f"{ckpt_dir}/{time.strftime('%Y%m%d_%H%M%S')}_{wandb.run.id}.safe"
@@ -924,4 +1281,4 @@ if __name__ == "__main__":
      nm = f"train_{m}"
      if nm in globals():
        print(f"training {m}")
-        globals()[nm]()
+        with Profiling(enabled=getenv("PYPROFILE")): globals()[nm]()
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128
+
+export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+# export BEAM_LOG_SURPASS_MAX=1
+# export BASEDIR="/raid/datasets/wiki"
+
+export RESET_STEP=1
+export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,69 @@
+# 1. Problem
+
+This problem uses BERT for NLP.
+
+## Requirements
+
+Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+Also install gdown (for dataset), numpy, tqdm and tensorflow.
+```
+pip install gdown numpy tqdm tensorflow
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+# 2. Directions
+
+## Steps to download and verify data
+
+### 1. Download raw data
+
+```
+BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
+```
+
+### 2. Preprocess train and validation data
+
+Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. 
+
+#### Training:
+```
+BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
+```
+
+Generating a specific topic (Between 0 and 499)
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
+```
+
+#### Validation:
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
+```
+## Running
+
+### tinybox_green
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
+```
+
+### tinybox_red
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
+```
+### tinybox_8xMI300X
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
+```
@@ -0,0 +1,14 @@
+#!/bin/bash
+
+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
+
+export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
+export BASEDIR="/raid/datasets/wiki"
+
+export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,17 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
+
+# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
+export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
+export TRAIN_STEPS=3900
+
+export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
+export BASEDIR="/raid/datasets/wiki"
+
+export WANDB=1 PARALLEL=0
+
+RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -0,0 +1,27 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export SUBMISSION_PLATFORM="tinybox_8xMI300X"
+export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
+
+# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
+export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
+export TRAIN_STEPS=3900
+
+export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
+export BASEDIR="/raid/datasets/wiki"
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=1
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="bert_8xMI300x_${DATETIME}_${SEED}.log"
+
+# init  # TODO: without DEBUG=2 it hangs
+BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 DEBUG=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -4,24 +4,20 @@ This problem uses BERT for NLP.

 ## Requirements

-Install tinygrad and mlperf-logging from master.
+Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
 ```
 git clone https://github.com/tinygrad/tinygrad.git
 python3 -m pip install -e ".[mlperf]"
 ```
-Also install tqdm and tensorflow.
+Also install gdown (for dataset), numpy, tqdm and tensorflow.
 ```
-pip install tqdm tensorflow
+pip install gdown numpy tqdm tensorflow
 ```

 ### tinybox_green
 Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
 This is the default on production tinybox green.

-### tinybox_red
-Disable cwsr + increase mes timeout.
-Install the custom amdgpu driver per [README](https://github.com/nimlgen/amdgpu_ubuntu_22_04/blob/v6.1.3/readme.md)
-
 # 2. Directions

 ## Steps to download and verify data
@@ -56,18 +52,18 @@ BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval

 #### Steps to run benchmark
 ```
-examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
 ```

 ### tinybox_red

-#### One time setup
-
+#### Steps to run benchmark
 ```
-examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/setup.sh
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
 ```
+### tinybox_8xMI300X

 #### Steps to run benchmark
 ```
-examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
 ```
@@ -1,13 +1,16 @@
 #!/bin/bash

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

-export BEAM=4 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+
+export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
+export BEAM_LOG_SURPASS_MAX=1
 export BASEDIR="/raid/datasets/wiki"

-export BENCHMARK=10 DEBUG=2
+export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2

 python3 examples/mlperf/model_train.py
@@ -1,10 +1,12 @@
 #!/bin/bash

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

-export BEAM=4 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+
+export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"

@@ -1,11 +1,13 @@
 #!/bin/bash

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

-export BEAM=4 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+
+export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"

@@ -17,7 +19,7 @@ DATETIME=$(date "+%m%d%H%M")
 LOGFILE="bert_green_${DATETIME}_${SEED}.log"

 # init
-BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
+BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE

 # run
 PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -4,24 +4,20 @@ This problem uses BERT for NLP.

 ## Requirements

-Install tinygrad and mlperf-logging from master.
+Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
 ```
 git clone https://github.com/tinygrad/tinygrad.git
 python3 -m pip install -e ".[mlperf]"
 ```
-Also install tqdm and tensorflow.
+Also install gdown (for dataset), numpy, tqdm and tensorflow.
 ```
-pip install tqdm tensorflow
+pip install gdown numpy tqdm tensorflow
 ```

 ### tinybox_green
 Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
 This is the default on production tinybox green.

-### tinybox_red
-Disable cwsr + increase mes timeout.
-Install the custom amdgpu driver per [README](https://github.com/nimlgen/amdgpu_ubuntu_22_04/blob/v6.1.3/readme.md)
-
 # 2. Directions

 ## Steps to download and verify data
@@ -56,18 +52,18 @@ BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval

 #### Steps to run benchmark
 ```
-examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
 ```

 ### tinybox_red

-#### One time setup
-
+#### Steps to run benchmark
 ```
-examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/setup.sh
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
 ```
+### tinybox_8xMI300X

 #### Steps to run benchmark
 ```
-examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
 ```
@@ -1,13 +1,17 @@
 #!/bin/bash

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

-export BEAM=3 BEAM_UOPS_MAX=3000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+
+export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
+export BEAM_LOG_SURPASS_MAX=1
 export BASEDIR="/raid/datasets/wiki"

-export BENCHMARK=10 DEBUG=2
+export RESET_STEP=1
+export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2

 python3 examples/mlperf/model_train.py
@@ -1,10 +1,12 @@
 #!/bin/bash

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

-export BEAM=3 BEAM_UOPS_MAX=3000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+
+export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"

@@ -1,11 +1,13 @@
 #!/bin/bash

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

-export BEAM=3 BEAM_UOPS_MAX=3000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
+
+export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"

@@ -16,8 +18,13 @@ export SEED=$RANDOM
 DATETIME=$(date "+%m%d%H%M")
 LOGFILE="bert_red_${DATETIME}_${SEED}.log"

+export HCQDEV_WAIT_TIMEOUT_MS=100000  # prevents hang?
+
 # init
-BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
+sleep 5 && sudo rmmod amdgpu || true
+BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE

 # run
+# TODO: AM driver resulted in nan
+sudo modprobe amdgpu
 PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,8 +0,0 @@
-#!/bin/bash
-
-rocm-smi --setprofile compute
-rocm-smi --setmclk 3
-rocm-smi --setperflevel high
-
-# power cap to 350W
-# echo "350000000" | sudo tee /sys/class/drm/card{1..6}/device/hwmon/hwmon*/power1_cap
@@ -1,6 +1,6 @@
 #!/bin/bash

-export PYTHONPATH="."
+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="."
+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="."
+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="."
+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="."
+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="."
+export PYTHONPATH="." AMD=1
 export MODEL="resnet"
 export SUBMISSION_PLATFORM="tinybox_red"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
@@ -0,0 +1,38 @@
+# 1. Problem
+
+This problem uses RetinaNet for SSD.
+
+## Requirements
+
+Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+
+Also install the following dependencies:
+```
+pip install tqdm numpy pycocotools boto3 pandas torch torchvision
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+# 2. Directions
+
+## Steps to download data
+
+Run the following:
+```
+BASEDIR=/raid/datasets/openimages python3 extra/datasets/openimages.py
+```
+
+## Running
+
+### tinybox_green
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/run_and_time.sh
+```
@@ -0,0 +1,14 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="retinanet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+export BASEDIR="/raid/datasets/openimages"
+
+# export RESET_STEP=0
+
+export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export BENCHMARK=5 DEBUG=2
+
+python examples/mlperf/model_train.py
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="retinanet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+export BASEDIR="/raid/datasets/openimages"
+
+# export RESET_STEP=0
+
+export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export WANDB=1 PARALLEL=0
+export RUNMLPERF=1
+
+python examples/mlperf/model_train.py
@@ -0,0 +1,23 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="retinanet"
+export SUBMISSION_PLATFORM="tinybox_green"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+
+export TRAIN_BEAM=2 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+export IGNORE_JIT_FIRST_BEAM=1
+export BASEDIR="/raid/datasets/openimages"
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=1
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="retinanet_green_${DATETIME}_${SEED}.log"
+
+# init
+BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -0,0 +1,14 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="retinanet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+export BASEDIR="/raid/datasets/openimages"
+
+# export RESET_STEP=0
+
+export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export BENCHMARK=5 DEBUG=2
+
+python examples/mlperf/model_train.py
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="retinanet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+export BASEDIR="/raid/datasets/openimages"
+
+# export RESET_STEP=0
+
+export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export WANDB=1 PARALLEL=0
+export RUNMLPERF=1
+
+python examples/mlperf/model_train.py
@@ -0,0 +1,39 @@
+{
+    "submitter": "tinycorp",
+    "division": "closed",
+    "status": "Available on-premise",
+    "system_name": "tinybox 8xMI300X",
+    "number_of_nodes": "1",
+    "host_processors_per_node": "2",
+    "host_processor_model_name": "AMD EPYC 9354 32-Core Processor",
+    "host_processor_core_count": "32",
+    "host_processor_vcpu_count": "64",
+    "host_processor_frequency": "",
+    "host_processor_caches": "",
+    "host_processor_interconnect": "",
+    "host_memory_capacity": "2304GB",
+    "host_storage_type": "NVMe SSD",
+    "host_storage_capacity": "3x 4TB raid array",
+    "host_networking": "",
+    "host_networking_topology": "",
+    "host_memory_configuration": "24x 96GB DDR5",
+    "accelerators_per_node": "8",
+    "accelerator_model_name": "AMD Instinct MI300X",
+    "accelerator_host_interconnect": "PCIe 5.0 x16",
+    "accelerator_frequency": "",
+    "accelerator_on-chip_memories": "",
+    "accelerator_memory_configuration": "HBM3",
+    "accelerator_memory_capacity": "192GB",
+    "accelerator_interconnect": "",
+    "accelerator_interconnect_topology": "",
+    "cooling": "air",
+    "hw_notes": "",
+    "framework": "tinygrad, branch mlperf_training_v5.0",
+    "other_software_stack": {
+      "python": "3.10.16",
+      "ROCm": "3.0.0+94441cb"
+    },
+    "operating_system": "Ubuntu 24.04.1 LTS",
+    "sw_notes": ""
+  }
+  
@@ -28,7 +28,7 @@
  "accelerator_interconnect_topology": "",
  "cooling": "air",
  "hw_notes": "",
-  "framework": "tinygrad, commit b5546912e24e0a864b35924da4efa5d71cfe368b",
+  "framework": "tinygrad, branch mlperf_training_v5.0",
  "other_software_stack": {
    "python": "3.10.12",
    "CUDA": "12.4"
@@ -28,10 +28,9 @@
  "accelerator_interconnect_topology": "",
  "cooling": "air",
  "hw_notes": "",
-  "framework": "tinygrad, commit b5546912e24e0a864b35924da4efa5d71cfe368b",
+  "framework": "tinygrad, branch mlperf_training_v5.0",
  "other_software_stack": {
-    "python": "3.10.12",
-    "ROCm": "6.1.3"
+    "python": "3.10.12"
  },
  "operating_system": "Ubuntu 22.04.4",
  "sw_notes": ""
@@ -16,9 +16,9 @@ class LinearGen:
    self.l4 = Tensor.scaled_uniform(1024, 784)

  def forward(self, x):
-    x = x.dot(self.l1).leakyrelu(0.2)
-    x = x.dot(self.l2).leakyrelu(0.2)
-    x = x.dot(self.l3).leakyrelu(0.2)
+    x = x.dot(self.l1).leaky_relu(0.2)
+    x = x.dot(self.l2).leaky_relu(0.2)
+    x = x.dot(self.l3).leaky_relu(0.2)
    x = x.dot(self.l4).tanh()
    return x

@@ -31,9 +31,9 @@ class LinearDisc:

  def forward(self, x):
    # balance the discriminator inputs with const bias (.add(1))
-    x = x.dot(self.l1).add(1).leakyrelu(0.2).dropout(0.3)
-    x = x.dot(self.l2).leakyrelu(0.2).dropout(0.3)
-    x = x.dot(self.l3).leakyrelu(0.2).dropout(0.3)
+    x = x.dot(self.l1).add(1).leaky_relu(0.2).dropout(0.3)
+    x = x.dot(self.l2).leaky_relu(0.2).dropout(0.3)
+    x = x.dot(self.l3).leaky_relu(0.2).dropout(0.3)
    x = x.dot(self.l4).log_softmax()
    return x

@@ -0,0 +1,94 @@
+# https://arxiv.org/pdf/2409.02060
+import time
+import numpy as np
+np.set_printoptions(suppress=True, linewidth=1000)
+import functools
+from tinygrad import Tensor, nn, Device, GlobalCounters
+from tinygrad.helpers import Timing, getenv
+from extra.models.llama import Transformer, convert_from_huggingface
+
+class MixtureFeedForward:
+  def __init__(self, num_experts:int, activated_experts:int, dim:int, hidden_dim:int, linear=nn.Linear):
+    self.activated_experts = activated_experts
+    self.gate = nn.Linear(dim, num_experts, bias=False)
+    self.up_proj = Tensor.zeros(num_experts, hidden_dim, dim, dtype='bfloat16')
+    self.down_proj = Tensor.zeros(num_experts, dim, hidden_dim, dtype='bfloat16')
+    self.gate_proj = Tensor.zeros(num_experts, hidden_dim, dim, dtype='bfloat16')
+  def __call__(self, x:Tensor) -> Tensor:
+    assert x.shape[0] == 1, "only BS=1"
+    assert x.shape[1] == 1, "only length=1"
+    g = self.gate(x).float().softmax(-1)
+
+    g = g.squeeze() # (BS, length, num_experts) -> (num_experts,)
+    probs, sel = g.topk(self.activated_experts)
+
+    # run MoE
+    x_up_gate = x.dot(self.gate_proj[sel].permute(0,2,1)).silu() * x.dot(self.up_proj[sel].permute(0,2,1))
+    x_down = x_up_gate.dot(self.down_proj[sel].permute(0,2,1))
+    return (x_down.float() * probs.reshape(self.activated_experts, 1, 1)).sum(axis=0)
+
+# model is bf16, 1.3B active, 6.9B total
+# M3 Max is 400 GB/s, so 400/2.6 = ~154 tok/s
+
+def fetch_weights() -> dict[str, Tensor]:
+  # TODO: make this lazy so the 3 fetches can happen in parallel
+  m1 = Tensor.from_url("https://huggingface.co/allenai/OLMoE-1B-7B-0924/resolve/main/model-00001-of-00003.safetensors").to(Device.DEFAULT)
+  m2 = Tensor.from_url("https://huggingface.co/allenai/OLMoE-1B-7B-0924/resolve/main/model-00002-of-00003.safetensors").to(Device.DEFAULT)
+  m3 = Tensor.from_url("https://huggingface.co/allenai/OLMoE-1B-7B-0924/resolve/main/model-00003-of-00003.safetensors").to(Device.DEFAULT)
+  return {**nn.state.safe_load(m1), **nn.state.safe_load(m2), **nn.state.safe_load(m3)}
+
+if __name__ == "__main__":
+  if getenv("TORCH"):
+    from transformers import OlmoeForCausalLM, AutoTokenizer
+    model = OlmoeForCausalLM.from_pretrained("allenai/OLMoE-1B-7B-0924")
+    tokenizer = AutoTokenizer.from_pretrained("allenai/OLMoE-1B-7B-0924")
+    inputs = tokenizer("Hello", return_tensors="pt")
+    generate_ids = model.generate(inputs.input_ids, max_length=30)
+    out = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+    print(out)
+    exit(0)
+
+  with Timing("create model: "):
+    model = Transformer(n_layers=16, dim=2048, hidden_dim=1024, n_heads=16, norm_eps=1e-5, qk_norm=1e-5, max_context=1024,
+                        vocab_size=50304, feed_forward=functools.partial(MixtureFeedForward, 64, 8))
+    model_state_dict = nn.state.get_state_dict(model)
+    del model_state_dict['freqs_cis']
+
+  with Timing("load weights to GPU: "):
+    nhf_state = convert_from_huggingface(fetch_weights(), 16, 16, 16)
+    # NOTE: i'm not sure this actually needs float32, it may just change the type of things downstream from it. but doesn't match torch w/o this
+    for needs_float32 in ['tok_embeddings.weight']: nhf_state[needs_float32] = nhf_state[needs_float32].float()
+  print(f"ram used: {GlobalCounters.mem_used/1e9:.2f} GB")
+
+  with Timing("unpack weights: "):
+    nn.state.load_state_dict(model, nhf_state, verbose=False, strict=False, consume=True, realize=False)
+    assert len(nhf_state) == 0
+    Tensor.realize(*list(nn.state.get_state_dict(model).values()))
+  print(f"ram used: {GlobalCounters.mem_used/1e9:.2f} GB")
+
+  count = 30
+  temperature = 0
+
+  with Timing("load tokenizer: "):
+    from transformers import AutoTokenizer
+    tokenizer = AutoTokenizer.from_pretrained("allenai/OLMoE-1B-7B-0924")
+
+  toks = [12092]
+  start_pos = 0
+  timings = []
+  for i in range(count):
+    GlobalCounters.reset()
+    st = time.perf_counter()
+    tok = model(Tensor([toks[start_pos:]]), start_pos, temperature).item()
+    timings.append(time.perf_counter()-st)
+    toks.append(tok)
+    start_pos += 1
+    print(toks)
+    print(tokenizer.decode(toks))
+  print(f"fastest token {min(timings)*1e3:.2f} ms, {1/min(timings):.1f} tok/s")
+
+  if temperature == 0:
+    # Hello, I am a newbie to this forum and I am trying to get a better understanding of the different types of data that can be stored in a
+    assert toks == [12092, 13, 309, 717, 247, 747, 17782, 281, 436, 12209, 285, 309, 717, 2820, 281, 755,
+                    247, 1805, 4685, 273, 253, 1027, 3510, 273, 941, 326, 476, 320, 7141, 275, 247], "BAD OUTPUT!"
+
@@ -12,7 +12,7 @@ from tinygrad.engine.realize import CompiledRunner

 import onnx
 from onnx.helper import tensor_dtype_to_np_dtype
-from extra.onnx import OnnxRunner   # TODO: port to main tinygrad
+from tinygrad.frontend.onnx import OnnxRunner

 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
 OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
@@ -92,10 +92,11 @@ def test_vs_compile(run, new_inputs, test_val=None):
  print("**** test done ****")

  # test that changing the numpy changes the model outputs
-  for v in new_inputs_numpy.values(): v *= 2
-  out = run(**inputs)
-  changed_val = out.numpy()
-  np.testing.assert_raises(AssertionError, np.testing.assert_array_equal, val, changed_val)
+  if any([x.device == 'NPY' for x in inputs.values()]):
+    for v in new_inputs_numpy.values(): v *= 2
+    out = run(**inputs)
+    changed_val = out.numpy()
+    np.testing.assert_raises(AssertionError, np.testing.assert_array_equal, val, changed_val)
  return val

 def test_vs_onnx(new_inputs, test_val, onnx_file, ort=False):
@@ -1,5 +1,5 @@
-from tinygrad import dtypes
-from tinygrad.helpers import trange
+from tinygrad import dtypes, getenv, Device
+from tinygrad.helpers import trange, colored, DEBUG, temp
 from tinygrad.nn.datasets import mnist
 import torch
 from torch import nn, optim
@@ -26,14 +26,20 @@ class Model(nn.Module):
    return self.lin(torch.flatten(x, 1))

 if __name__ == "__main__":
-  mps_device = torch.device("mps")
+  if getenv("TINY_BACKEND"):
+    import tinygrad.frontend.torch
+    device = torch.device("tiny")
+  else:
+    device = torch.device({"METAL":"mps","NV":"cuda"}.get(Device.DEFAULT, "cpu"))
+  if DEBUG >= 1: print(f"using torch backend {device}")
  X_train, Y_train, X_test, Y_test = mnist()
-  X_train = torch.tensor(X_train.float().numpy(), device=mps_device)
-  Y_train = torch.tensor(Y_train.cast(dtypes.int64).numpy(), device=mps_device)
-  X_test = torch.tensor(X_test.float().numpy(), device=mps_device)
-  Y_test = torch.tensor(Y_test.cast(dtypes.int64).numpy(), device=mps_device)
+  X_train = torch.tensor(X_train.float().numpy(), device=device)
+  Y_train = torch.tensor(Y_train.cast(dtypes.int64).numpy(), device=device)
+  X_test = torch.tensor(X_test.float().numpy(), device=device)
+  Y_test = torch.tensor(Y_test.cast(dtypes.int64).numpy(), device=device)

-  model = Model().to(mps_device)
+  if getenv("TORCHVIZ"): torch.cuda.memory._record_memory_history()
+  model = Model().to(device)
  optimizer = optim.Adam(model.parameters(), 1e-3)

  loss_fn = nn.CrossEntropyLoss()
@@ -48,8 +54,16 @@ if __name__ == "__main__":
    return loss

  test_acc = float('nan')
-  for i in (t:=trange(70)):
+  for i in (t:=trange(getenv("STEPS", 70))):
    samples = torch.randint(0, X_train.shape[0], (512,))  # putting this in JIT didn't work well
    loss = step(samples)
    if i%10 == 9: test_acc = ((model(X_test).argmax(axis=-1) == Y_test).sum() * 100 / X_test.shape[0]).item()
    t.set_description(f"loss: {loss.item():6.2f} test_accuracy: {test_acc:5.2f}%")
+
+  # verify eval acc
+  if target := getenv("TARGET_EVAL_ACC_PCT", 0.0):
+    if test_acc >= target and test_acc != 100.0: print(colored(f"{test_acc=} >= {target}", "green"))
+    else: raise ValueError(colored(f"{test_acc=} < {target}", "red"))
+  if getenv("TORCHVIZ"):
+    torch.cuda.memory._dump_snapshot(fp:=temp("torchviz.pkl", append_user=True))
+    print(f"saved torch memory snapshot to {fp}, view in https://pytorch.org/memory_viz")
@@ -44,7 +44,7 @@ def load_model(model_path:Path, model_params:Dict[str, Union[int, float]]) -> Tr
  model.layers = updated_layers

  # load weights
-  weights = fix_bf16(convert_from_huggingface(load(str(model_path / "model.safetensors.index.json")), model, model_params["n_heads"], model_params["n_kv_heads"], permute_layers=False))
+  weights = fix_bf16(convert_from_huggingface(load(str(model_path / "model.safetensors.index.json")), model_params["n_layers"], model_params["n_heads"], model_params["n_kv_heads"], permute_layers=False))

  # replace weights in model
  load_state_dict(model, weights, strict=False, consume=True)
@@ -5,7 +5,7 @@

 from tinygrad import Tensor, TinyJit, dtypes, GlobalCounters
 from tinygrad.nn import Conv2d, GroupNorm
-from tinygrad.nn.state import safe_load, load_state_dict
+from tinygrad.nn.state import safe_load, load_state_dict, get_state_dict
 from tinygrad.helpers import fetch, trange, colored, Timing
 from extra.models.clip import Embedder, FrozenClosedClipEmbedder, FrozenOpenClipEmbedder
 from extra.models.unet import UNetModel, Upsample, Downsample, timestep_embedding
@@ -356,8 +356,7 @@ class DPMPP2MSampler:
          c=c,
          uc=uc,
        )
-        x.realize()
-        old_denoised.realize()
+        x.realize(old_denoised)

    return x

@@ -385,7 +384,12 @@ if __name__ == "__main__":

  default_weight_url = 'https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors'
  weights = args.weights if args.weights else fetch(default_weight_url, 'sd_xl_base_1.0.safetensors')
-  load_state_dict(model, safe_load(weights), strict=False)
+  loaded_weights = load_state_dict(model, safe_load(weights), strict=False, verbose=False, realize=False)
+
+  start_mem_used = GlobalCounters.mem_used
+  with Timing("loaded weights in ", lambda et_ns: f", {(B:=(GlobalCounters.mem_used-start_mem_used))/1e9:.2f} GB loaded at {B/et_ns:.2f} GB/s"):
+    Tensor.realize(*loaded_weights)
+    del loaded_weights

  N = 1
  C = 4
@@ -396,8 +400,7 @@ if __name__ == "__main__":

  c, uc = model.create_conditioning([args.prompt], args.width, args.height)
  del model.conditioner
-  for v in c .values(): v.realize()
-  for v in uc.values(): v.realize()
+  Tensor.realize(*c.values(), *uc.values())
  print("created batch")

  # https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/inference/helpers.py#L101
@@ -437,14 +437,14 @@ class Generator:
    x = self.conv_pre(x)
    if g is not None:  x = x + self.cond(g)
    for i in range(self.num_upsamples):
-      x, xs = self.ups[i](x.leakyrelu(LRELU_SLOPE)), None
+      x, xs = self.ups[i](x.leaky_relu(LRELU_SLOPE)), None
      x_source = self.noise_convs[i](har_source)
      x = x + x_source
      for j in range(self.num_kernels):
        if xs is None: xs = self.resblocks[i * self.num_kernels + j].forward(x)
        else: xs += self.resblocks[i * self.num_kernels + j].forward(x)
      x = xs / self.num_kernels
-    return self.conv_post(x.leakyrelu()).tanh()
+    return self.conv_post(x.leaky_relu()).tanh()

 # **** helpers ****

@@ -504,7 +504,7 @@ def load_checkpoint_enc(checkpoint_path, model: ContentVec, optimizer=None, skip
        obj, v = getattr(parent, "weight"), weight_norm(weight_v, weight_g, 0)
        weight_g, weight_v, parent, skip = None, None, None, False
      if not skip and obj.shape == v.shape:
-        if "feature_extractor" in key and (isinstance(parent, nn.GroupNorm) or isinstance(parent, nn.LayerNorm)):  # cast
+        if "feature_extractor" in key and (isinstance(parent, (nn.GroupNorm, nn.LayerNorm))):  # cast
          obj.assign(v.to(obj.device).float())
        else:
          obj.assign(v.to(obj.device))
@@ -5,10 +5,13 @@
 #  - symbolic removal

 from examples.beautiful_mnist import Model
-from tinygrad import Tensor, nn, getenv, GlobalCounters
+from tinygrad import Tensor, nn, getenv, GlobalCounters, Variable
 from tinygrad.nn.datasets import mnist
 from tinygrad.helpers import trange, DEBUG

+# STEPS=70 python3 examples/stunning_mnist.py
+# NOTE: it's broken with STACK=1, why?
+
 if __name__ == "__main__":
  X_train, Y_train, X_test, Y_test = mnist()
  print("*** got data")
@@ -24,19 +27,21 @@ if __name__ == "__main__":
  print("*** got samples")

  with Tensor.train():
-    # TODO: this shouldn't be a for loop. something like: (contract is still up in the air)
    """
    i = UOp.range(samples.shape[0])  # TODO: fix range function on UOp
    losses = model(X_samp[i]).sparse_categorical_crossentropy(Y_samp[i]).backward().contract(i)
    opt.schedule_steps(i)
    """
+    # TODO: this shouldn't be a for loop. something like: (contract is still up in the air)
+    vi = Variable('i', 0, samples.shape[0]-1)
    losses = []
    for i in range(samples.shape[0]):
+      vib = vi.bind(i)
      opt.zero_grad()
-      losses.append(model(X_samp[i]).sparse_categorical_crossentropy(Y_samp[i]).backward())
+      losses.append(model(X_samp[vib]).sparse_categorical_crossentropy(Y_samp[vib]).backward())
      opt.schedule_step()
    # TODO: this stack currently breaks the "generator" aspect of losses. it probably shouldn't
-    #losses = Tensor.stack(*losses)
+    if getenv("STACK", 0): losses = Tensor.stack(*losses)
  print("*** scheduled training")

  # evaluate the model
@@ -49,5 +54,8 @@ if __name__ == "__main__":

  # only actually do anything at the end
  if getenv("LOSS", 1):
-    for i in (t:=trange(len(losses))): t.set_description(f"loss: {losses[i].item():6.2f}")
-  print(f"test_accuracy: {test_acc.item():5.2f}%")
+    for i in (t:=trange(len(losses))):
+      GlobalCounters.reset()
+      t.set_description(f"loss: {losses[i].item():6.2f}")
+  if getenv("TEST", 1):
+    print(f"test_accuracy: {test_acc.item():5.2f}%")
@@ -58,7 +58,7 @@ if __name__ == "__main__":
            return None
          return {"input": img.numpy()}
      quantize_static(model_fp32, fn, ImagenetReader(), quant_format=QuantFormat.QDQ, per_channel=False,
-                      activation_type=QuantType.QUInt8, weight_type=QuantType.QInt8,
+                      activation_type=QuantType.QUInt8, weight_type=QuantType.QUInt8,
                      extra_options={"ActivationSymmetric": False})

  run_onnx_jit, input_specs = load_onnx_model(fetch(fn))
@@ -70,9 +70,12 @@ if __name__ == "__main__":
    GlobalCounters.reset()
    p = run_onnx_jit(**{t_name:img})
    assert p.shape == (1,1000)
-    t = p.argmax().item()
+    t = p.to('cpu').argmax().item()
    hit += y==t
    print(f"target: {y:3d}  pred: {t:3d}  acc: {hit/(i+1)*100:.2f}%")

+  MS_TARGET = 13.4
+  print(f"need {GlobalCounters.global_ops/1e9*(1000/MS_TARGET):.2f} GFLOPS for {MS_TARGET:.2f} ms")
+
  import pickle
  with open("/tmp/im.pkl", "wb") as f: pickle.dump(run_onnx_jit, f)
@@ -0,0 +1,19 @@
+import sys, pickle
+from tinygrad import GlobalCounters
+from tinygrad.helpers import fetch, getenv
+from examples.test_onnx_imagenet import imagenet_dataloader
+
+if __name__ == "__main__":
+  with open(fetch(sys.argv[1]), "rb") as f:
+    run_onnx_jit = pickle.load(f)
+  input_name = run_onnx_jit.captured.expected_names[0]
+  device = run_onnx_jit.captured.expected_st_vars_dtype_device[0][-1]
+  print(f"input goes into {input_name=} on {device=}")
+  hit = 0
+  for i,(img,y) in enumerate(imagenet_dataloader(cnt=getenv("CNT", 100))):
+    GlobalCounters.reset()
+    p = run_onnx_jit(**{input_name:img.to(device)})
+    assert p.shape == (1,1000)
+    t = p.to('cpu').argmax().item()
+    hit += y==t
+    print(f"target: {y:3d}  pred: {t:3d}  acc: {hit/(i+1)*100:.2f}%")
@@ -0,0 +1,11 @@
+/*!
+Pure v3.0.0
+Copyright 2013 Yahoo!
+Licensed under the BSD License.
+https://github.com/pure-css/pure/blob/master/LICENSE
+*/
+/*!
+normalize.css v | MIT License | https://necolas.github.io/normalize.css/
+Copyright (c) Nicolas Gallagher and Jonathan Neal
+*/
+/*! normalize.css v8.0.1 | MIT License | github.com/necolas/normalize.css */html{line-height:1.15;-webkit-text-size-adjust:100%}body{margin:0}main{display:block}h1{font-size:2em;margin:.67em 0}hr{box-sizing:content-box;height:0;overflow:visible}pre{font-family:monospace,monospace;font-size:1em}a{background-color:transparent}abbr[title]{border-bottom:none;text-decoration:underline;-webkit-text-decoration:underline dotted;text-decoration:underline dotted}b,strong{font-weight:bolder}code,kbd,samp{font-family:monospace,monospace;font-size:1em}small{font-size:80%}sub,sup{font-size:75%;line-height:0;position:relative;vertical-align:baseline}sub{bottom:-.25em}sup{top:-.5em}img{border-style:none}button,input,optgroup,select,textarea{font-family:inherit;font-size:100%;line-height:1.15;margin:0}button,input{overflow:visible}button,select{text-transform:none}[type=button],[type=reset],[type=submit],button{-webkit-appearance:button}[type=button]::-moz-focus-inner,[type=reset]::-moz-focus-inner,[type=submit]::-moz-focus-inner,button::-moz-focus-inner{border-style:none;padding:0}[type=button]:-moz-focusring,[type=reset]:-moz-focusring,[type=submit]:-moz-focusring,button:-moz-focusring{outline:1px dotted ButtonText}fieldset{padding:.35em .75em .625em}legend{box-sizing:border-box;color:inherit;display:table;max-width:100%;padding:0;white-space:normal}progress{vertical-align:baseline}textarea{overflow:auto}[type=checkbox],[type=radio]{box-sizing:border-box;padding:0}[type=number]::-webkit-inner-spin-button,[type=number]::-webkit-outer-spin-button{height:auto}[type=search]{-webkit-appearance:textfield;outline-offset:-2px}[type=search]::-webkit-search-decoration{-webkit-appearance:none}::-webkit-file-upload-button{-webkit-appearance:button;font:inherit}details{display:block}summary{display:list-item}template{display:none}[hidden]{display:none}html{font-family:sans-serif}.hidden,[hidden]{display:none!important}.pure-img{max-width:100%;height:auto;display:block}
@@ -0,0 +1,5 @@
+net_*
+llama3-2.tiktoken
+tiktoken.js
+tiktoken_bg.wasm
+transformer*
@@ -0,0 +1,8 @@
+# How to build and run tinychat in browser (WebGPU and WASM)
+- `PYTHONPATH=. python examples/tinychat/tinychat-browser/compile.py`
+- `./examples/tinychat/tinychat-browser/compile_wasm.sh`
+    - Prerequisite: [install emscripten](https://emscripten.org/docs/getting_started/downloads.html). This script looks for `~/emsdk/emsdk_env.sh`, adjust this based on your installation.
+- `./examples/tinychat/tinychat-browser/make_tiktoken_js.sh`
+    - Prerequisite: install `npm`, `webpack`.
+- `cd examples/tinychat && python -m http.server 7776`
+- In browser: open either `localhost:7776/tinychat-browser` (WebGPU), or `localhost:7776/tinychat-browser/?backend=wasm` (WASM)
@@ -0,0 +1,149 @@
+import os, json, hashlib, math
+from extra.export_model import export_model
+from examples.llama3 import build_transformer, Tokenizer
+from tinygrad.nn.state import get_state_dict, load_state_dict
+from tinygrad import Device, Variable, Tensor, dtypes, TinyJit
+from tinygrad.helpers import fetch, Context
+from tiktoken.load import load_tiktoken_bpe, dump_tiktoken_bpe
+
+def prepare_browser_chunks(model):
+  # split weights into browser-friendly chunks
+  state_dict = get_state_dict(model)
+  del state_dict['output.weight'], state_dict['output.scale'] # same as tok_embeddings; ensures consistency with model export
+  chunk_size = 16 * 1024 * 1024 # small chunks based on iphone browser constraints
+  metadata = {}
+  # We won't export cache_kv bytes (because we start inference on client at start_pos=0), but we will tell the client how big cache_kv needs to be
+  t_infos = [(v.lazydata.base.realized.nbytes, k, v.dtype) for k,v in state_dict.items() if "cache_kv" not in k]
+  empty_t_infos = [(v.lazydata.base.realized.nbytes, k, v.dtype) for k,v in state_dict.items() if "cache_kv" in k]
+
+  split_t_infos = []
+  for size, name, dtype in t_infos:
+    if size <= chunk_size:
+      split_t_infos.append((size, name, dtype, ()))
+    else: # split large weights into multiple parts
+      for i in range(0, size, chunk_size):
+        split_t_infos.append((min(chunk_size, size-i), f"{name}_part{math.ceil(i/chunk_size)}", dtype, (i, min(i+chunk_size, size))))
+
+  files = []
+  # pack weights into files with FFD bin packing
+  split_t_infos = sorted(split_t_infos, reverse=True)
+  for info in split_t_infos:
+    placed = False
+    for file in files:
+      if sum(i[0] for i in file) + info[0] <= chunk_size:
+        if info[3] and any(i[3] for i in file): continue # no two split tensors can touch the same file, due to wasm loading constraints
+        file.append(info)
+        placed = True
+        break
+    if not placed:
+      files.append([info])
+
+  tinygrad_dtypes = {dtypes.float32: "float32", dtypes.float16: "float16", dtypes.int8: "int8", dtypes.int32: "int32"}
+  for i, file in enumerate(files):
+    cursor = 0
+    with open(os.path.join(os.path.dirname(__file__), f'./net_part{i}.chunk'), "wb+") as writer:
+      for size, name, dtype, offsets in file:
+        name, part_num = (name, 0) if "_part" not in name else (name.split("_part")[0], int(name.split("_part")[1]))
+        default = {"parts": {}, "dtype": tinygrad_dtypes[dtype]}
+        weight_metadata = metadata.get(name, default)
+        weight_metadata["parts"][part_num] = {"file": i, "file_start_pos": cursor, "size": size}
+        metadata[name] = weight_metadata
+        data = bytes(state_dict[name].lazydata.base.realized.as_buffer())
+        data = data if not offsets else data[offsets[0]:offsets[1]]
+        writer.write(data)
+        cursor += size
+
+  metadata.update({name: {"parts": {0: {"empty": True, "size": size}}, "dtype": tinygrad_dtypes[dtype]} for size, name, dtype in empty_t_infos})
+
+  for k in metadata:
+    metadata[k]["parts"] = [part for part_num, part in sorted(metadata[k]["parts"].items(), key = lambda x: x[0])]
+    cursor = 0
+    for i, part in enumerate(metadata[k]["parts"]):
+      metadata[k]["parts"][i]["target_start_pos"] = cursor
+      cursor += part["size"]
+    metadata[k]["size"] = cursor
+
+  # compute hashes, which client app will check to determine whether to update with new weights and/or detect integrity issues
+  state_dict_hash = hashlib.sha256(json.dumps(metadata, sort_keys=True).encode("utf-8")).hexdigest()
+  metadata = {"state_dict": metadata, "state_dict_hash": state_dict_hash, "files": []}
+  hashes = set()
+  for i in range(len(files)):
+    with open(os.path.join(os.path.dirname(__file__), f'./net_part{i}.chunk'), "rb") as reader:
+      hash = hashlib.sha256(reader.read()).hexdigest()
+      hashes.add(hash)
+      metadata["files"].append({"name": f'net_part{i}.chunk', "hash": hash})
+  if len(hashes) != len(files): print(f"WARNING: {len(files)} files were exported, but only {len(hashes)} are unique: something may have gone wrong")
+  metadata_hash = hashlib.sha256(json.dumps(metadata, sort_keys=True).encode("utf-8")).hexdigest()
+  metadata = {"metadata": metadata, "metadata_hash": metadata_hash}
+
+  with open(os.path.join(os.path.dirname(__file__), f'./net_metadata.json'), "w") as writer: json.dump(metadata, writer, indent=4)
+  return metadata
+
+def validate_model(model, tokenizer):
+  prompt = "yo"
+  toks = [tokenizer.bos_id]
+  toks += [tokenizer.special_tokens["<|start_header_id|>"]] + tokenizer.encode("user") + [tokenizer.special_tokens["<|end_header_id|>"]] + tokenizer.encode("\n\n")
+  toks += tokenizer.encode(prompt) + [tokenizer.special_tokens["<|eot_id|>"]]
+  toks += [tokenizer.special_tokens["<|start_header_id|>"]] + tokenizer.encode("assistant") + [tokenizer.special_tokens["<|end_header_id|>"]] + tokenizer.encode("\n\n")
+  start_pos = 0
+  run = TinyJit(model.forward)
+  for tok in toks[:-1]:
+    run(Tensor([[tok]]), Variable("start_pos", 0, model.max_context).bind(start_pos), 0.0, 0, 0.0, 0.0, 0.0).realize()
+    start_pos += 1
+  tok = toks[-1]
+  result = ""
+  expected = "How's it going?"
+  while True:
+    tok = run(Tensor([[tok]]), Variable("start_pos", 0, model.max_context).bind(start_pos), 0.0, 0, 0.0, 0.0, 0.0).item()
+    start_pos += 1
+    if tok in tokenizer.stop_tokens or len(result) > len(expected): break
+    result += tokenizer.decode([tok])
+  assert result == expected, f"Model validation failed, expected output: {expected}, actual output: {result}"
+
+if __name__=="__main__":
+  # Export BPE data for use with tiktoken.js
+  tokenizer_path = fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model", "tokenizer.model", subdir="llama3-1b-instruct")
+  mergeable_ranks = load_tiktoken_bpe(str(tokenizer_path))
+  bpe_path = os.path.join(os.path.dirname(__file__), "llama3-2.tiktoken")
+  dump_tiktoken_bpe(mergeable_ranks, bpe_path)
+  tokenizer = Tokenizer(str(tokenizer_path))
+
+  model_path = fetch("https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-f16.gguf", "Llama-3.2-1B-Instruct-f16.gguf", subdir="llama3-1b-instruct")
+  Tensor.no_grad = True
+  max_context=1024
+  tok = 128000
+  TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P = 0.95, 0, 0.0, 0.0, 0.0
+  start_pos = Variable("start_pos", 0, max_context).bind(0)
+  model_input = lambda: [Tensor([[tok]]), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P]
+
+  Device.DEFAULT="CPU"
+  model = build_transformer(model_path, model_size="1B", quantize="int8", scale_dtype=dtypes.float32, device=Device.DEFAULT, max_context=max_context)
+  state_dict = get_state_dict(model)
+  validate_model(model, tokenizer)
+  model_name = "transformer"
+
+  with Context(BEAM=3):
+    cprog, js_wrapper = export_model(model, "wasm", *model_input(), model_name=model_name)
+    # ensure consistency with exported weights
+    js_wrapper = js_wrapper.replace("output.weight", "tok_embeddings.weight").replace("output.scale", "tok_embeddings.scale")
+
+  with open(os.path.join(os.path.dirname(__file__), f"{model_name}.c"), "w") as f: f.write(cprog)
+  with open(os.path.join(os.path.dirname(__file__), "net_clang.js"), "w") as f: f.write(js_wrapper)
+
+  Device.DEFAULT="WEBGPU"
+  # float16 is not yet supported for dawn/Vulkan/NVIDIA stack, see: https://issues.chromium.org/issues/42251215
+  # therefore for now, we used CLANG to quantize the float16 llama to int8 with float32 scales, then load to WEBGPU
+  model = build_transformer(model_path, model_size="1B", quantize="int8", max_context=max_context, load_weights=False)
+  load_state_dict(model, state_dict)
+  # these were the same before load_state_dict
+  model.output.weight, model.output.scale = model.tok_embeddings.weight, model.tok_embeddings.scale
+
+  validate_model(model, tokenizer)
+  metadata = prepare_browser_chunks(model) # export weights to disk
+
+  with Context(BEAM=3):
+    prg, input_sizes, output_sizes, state = export_model(model, "webgpu", *model_input(), model_name=model_name, stream_weights=True)
+    # ensure consistency with exported weights
+    prg = prg.replace("output.weight", "tok_embeddings.weight").replace("output.scale", "tok_embeddings.scale")
+
+  with open(os.path.join(os.path.dirname(__file__), "net.js"), "w") as f: f.write(prg)
@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+cd "$(dirname "$0")"
+
+# prereq: install emscripten: https://emscripten.org/docs/getting_started/downloads.html
+EMSCRIPTEN_PATH=~/emsdk/emsdk_env.sh
+source $EMSCRIPTEN_PATH
+step="transformer"
+initial_memory=6553600
+max_memory=1500053504
+exported_functions='["_net", "_malloc", "_free", "_set_buf"]'
+
+emcc "${step}.c" \
+  -O3 -msimd128 -ffast-math -flto \
+  -o "${step}.js" \
+  -s MODULARIZE=1 \
+  -s EXPORT_ES6=1 \
+  -s EXPORTED_FUNCTIONS="${exported_functions}" \
+  -s ENVIRONMENT='worker' \
+  -s FILESYSTEM=0 \
+  -s EVAL_CTORS \
+  -s ALLOW_MEMORY_GROWTH=1 \
+  -s INITIAL_MEMORY="$initial_memory" \
+  -s MAXIMUM_MEMORY="$max_memory"
@@ -0,0 +1,322 @@
+/* define colors */
+:root {
+  --primary-color: #fff;
+  --secondary-color: #2a2a2a;
+  --secondary-color-transparent: #ffffff66;
+  --primary-bg-color: #1a1a1a;
+  --foreground-color: #f0f0f0;
+}
+
+main {
+  width: 100%;
+  height: 100%;
+
+  display: flex;
+  flex-direction: column;
+
+  place-items: center;
+}
+
+.home {
+  width: 100%;
+  height: 90%;
+
+  margin-bottom: 10rem;
+}
+
+.title {
+  font-size: 3rem;
+  margin: 1rem 0;
+  margin-top: 3rem;
+}
+
+.histories-container-container {
+  width: 100%;
+  max-height: 75%;
+
+  position: relative;
+}
+
+.histories-container {
+  overflow-y: auto;
+  overflow-x: hidden;
+  width: 100%;
+  height: 100%;
+
+  display: flex;
+  flex-direction: column;
+  gap: 1rem;
+  align-items: center;
+
+  margin: 0;
+  padding: 3rem 1rem;
+}
+
+.histories-start {
+  height: 3rem;
+  width: 100%;
+
+  z-index: 999;
+  top: 0;
+  position: absolute;
+
+  background: linear-gradient(
+    180deg,
+    var(--primary-bg-color) 0%,
+    transparent 100%
+  );
+}
+.histories-end {
+  height: 3rem;
+  width: 100%;
+
+  z-index: 999;
+  bottom: 0;
+  position: absolute;
+
+  background: linear-gradient(
+    0deg,
+    var(--primary-bg-color) 0%,
+    transparent 100%
+  );
+}
+
+.history {
+  padding: 1rem;
+  width: 100%;
+  max-width: 40rem;
+
+  background-color: var(--secondary-color);
+  border-radius: 10px;
+  border-left: 2px solid var(--primary-color);
+
+  cursor: pointer;
+
+  transform: translateX(calc(1px * var(--tx, 0)));
+  opacity: var(--opacity, 1);
+}
+.history:hover {
+  background-color: var(--secondary-color);
+}
+
+.history-delete-button {
+  position: absolute;
+  top: 0;
+  right: 0;
+  padding: 0.5rem;
+  margin: 0;
+  outline: none;
+  border: none;
+  background-color: var(--secondary-color);
+  color: var(--foreground-color);
+  border-radius: 0 0 0 10px;
+  cursor: pointer;
+  transition: 0.2s;
+}
+.history-delete-button:hover {
+  background-color: var(--secondary-color);
+  padding: 0.75rem;
+}
+
+.messages {
+  overflow-y: auto;
+  height: 100%;
+  width: 100%;
+  max-width: 1200px;
+
+  display: flex;
+  flex-direction: column;
+  gap: 1rem;
+  align-items: center;
+  padding-top: 1rem;
+  padding-bottom: 11rem;
+}
+
+.message {
+  max-width: 75%;
+  padding: 0.5rem 1rem;
+  border-radius: 20px;
+}
+.message-role-assistant {
+  background-color: var(--secondary-color);
+  margin-right: auto;
+  color: #fff;
+}
+.message-role-user {
+  margin-left: auto;
+  background-color: var(--primary-color);
+  color: #000;
+}
+
+.message > pre {
+  white-space: pre-wrap;
+}
+
+.hljs {
+  width: 100%;
+  position: relative;
+  border-radius: 10px;
+  /* wrap code blocks */
+  white-space: pre-wrap;
+}
+/* put clipboard button in the top right corner of the code block */
+.clipboard-button {
+  position: absolute;
+  top: 0;
+  right: 0;
+  padding: 0.5rem;
+  margin: 0;
+  outline: none;
+  border: none;
+  background-color: var(--secondary-color);
+  color: var(--foreground-color);
+  border-radius: 0 0 0 10px;
+  cursor: pointer;
+  transition: 0.2s;
+}
+.clipboard-button:hover {
+  background-color: var(--secondary-color);
+  padding: 0.75rem;
+}
+
+.input-container {
+  position: absolute;
+  bottom: 0;
+
+  /* linear gradient from background-color to transparent on the top */
+  background: linear-gradient(
+    0deg,
+    var(--primary-bg-color) 55%,
+    transparent 100%
+  );
+
+  width: 100%;
+  max-width: 1200px;
+  display: flex;
+  flex-direction: column;
+  justify-content: center;
+  align-items: center;
+  z-index: 999;
+}
+
+.input-performance {
+  margin-top: 4rem;
+
+  display: flex;
+  flex-direction: row;
+  gap: 1rem;
+}
+
+.input-performance-point {
+  display: flex;
+  flex-direction: row;
+  place-items: center;
+  gap: 0.5rem;
+}
+.input-performance-point > p {
+  height: 1rem;
+  line-height: normal;
+}
+
+.input {
+  width: 90%;
+  min-height: 3rem;
+  flex-shrink: 0;
+
+  display: flex;
+  flex-direction: row;
+  justify-content: center;
+  gap: 0.5rem;
+
+  align-items: flex-end;
+  margin-bottom: 2rem;
+}
+
+.input-form {
+  width: 100%;
+  padding: 1rem;
+  min-height: 3rem;
+  max-height: 8rem;
+
+  background-color: var(--secondary-color);
+  color: var(--foreground-color);
+  border-radius: 10px;
+  border: none;
+  resize: none;
+  outline: none;
+}
+.mobile .input-form { /* prevent auto-zoom on touching prompt box */
+  font-size: 16px;
+}
+
+.input-button {
+  height: 3rem;
+  width: 4rem;
+
+  background-color: var(--primary-color);
+  color: var(--secondary-color);
+  border-radius: 10px;
+  padding: 0.5rem;
+  cursor: pointer;
+}
+.input-button:hover {
+  background-color: var(--secondary-color-transparent);
+}
+.input-button:disabled {
+  background-color: var(--secondary-color);
+  cursor: not-allowed;
+}
+
+/* wrap text */
+p {
+  white-space: pre-wrap;
+}
+
+/* fonts */
+.megrim-regular {
+  font-family: monospace;
+  font-weight: 400;
+  font-style: normal;
+}
+
+.monospace {
+  font-family: monospace;
+}
+
+.loading-bar {
+  display: flex;
+  flex-direction: row;
+  align-items: center;
+  gap: 0.5rem;
+  width: 100%;
+  min-height: 3rem;
+  margin-bottom: 2rem;
+}
+
+.loading-text {
+  color: var(--foreground-color);
+  font-size: 1rem;
+  white-space: nowrap;
+}
+
+#progress-percentage {
+  color: var(--foreground-color);
+  font-size: 1rem;
+  white-space: nowrap;
+}
+
+.progress-bar {
+  flex-grow: 1;
+  height: 0.5rem;
+  background-color: var(--secondary-color);
+  border-radius: 5px;
+  overflow: hidden;
+  position: relative;
+}
+
+.progress {
+  width: 0%;
+  height: 100%;
+  background-color: var(--primary-color);
+  transition: width 0.2s ease-in-out;
+}
@@ -0,0 +1,182 @@
+<!DOCTYPE html>
+
+<head>
+  <title>tinychat</title>
+  <meta name="viewport" content="width=device-width, initial-scale=1">
+  <link rel="icon" href="../favicon.svg" type="image/svg+xml">
+
+  <script defer src="../assets/cdn.jsdelivr.net/npm/@alpine-collective/[email protected]/dist/cdn.min.js"></script>
+  <script defer src="../assets/cdn.jsdelivr.net/npm/@alpinejs/[email protected]/dist/cdn.min.js"></script>
+  <script defer src="../assets/cdn.jsdelivr.net/npm/@alpinejs/[email protected]/dist/cdn.min.js"></script>
+  <script defer src="../assets/unpkg.com/@marcreichel/[email protected]/dist/alpine-autosize.min.js"></script>
+  <script defer src="../assets/unpkg.com/[email protected]/dist/cdn.min.js"></script>
+
+  <script src="../assets/unpkg.com/[email protected]/dist/purify.min.js"></script>
+  <script src="../assets/unpkg.com/[email protected]/marked.min.js"></script>
+  <script src="../assets/unpkg.com/[email protected]/lib/index.umd.js"></script>
+  <script src="../assets/unpkg.com/@highlightjs/[email protected]/highlight.min.js"></script>
+
+  <script src="index.js"></script>
+
+  <link rel="stylesheet" href="../assets/cdn.jsdelivr.net/npm/[email protected]/build/base-min.css">
+  <link rel="stylesheet" href="../assets/cdnjs.cloudflare.com/ajax/libs/font-awesome/6.5.2/css/all.min.css"
+    integrity="sha512-SnH5WK+bZxgPHs44uWIX+LLJAJ9/2PkPKZ5QiAj6Ta86w+fsb2TkcmfRyVX3pBnMFcV7oQPJkl9QevSCWr3W6A=="
+    crossorigin="anonymous" referrerpolicy="no-referrer" />
+  <link rel="stylesheet" href="../assets/unpkg.com/@highlightjs/[email protected]/styles/vs2015.min.css">
+
+  <link rel="stylesheet" href="index.css">
+  <link rel="stylesheet" href="../common.css">
+</head>
+
+<body>
+  <main x-data="state" x-init="console.log(endpoint)">
+    <div class="home centered" x-show="home === 0" x-transition x-effect="
+      $refs.inputForm.focus();
+      if (home === 1) setTimeout(() => home = 2, 100);
+      if (home === -1) setTimeout(() => home = 0, 100);
+    " @popstate.window="
+      if (home === 2) {
+        cancelGeneration = true;
+        if (maxContextReached) generating = false;
+        if (!generating) cstate = { time: null, messages: [] };
+        home = -1;
+        time_till_first = 0;
+        tokens_per_second = 0;
+        total_tokens = 0;
+      }
+    ">
+      <h1 class="title megrim-regular">tinychat</h1>
+      <div class="histories-container-container">
+        <template x-if="histories.length">
+          <div class="histories-start"></div>
+        </template>
+        <div class="histories-container" x-intersect="
+          $el.scrollTo({ top: 0, behavior: 'smooth' });
+        ">
+          <template x-for="_state in histories.toSorted((a, b) => b.time - a.time)">
+            <div x-data="{ otx: 0, trigger: 75 }" class="history" @click="
+            cstate = _state;
+            updateTotalTokens(cstate.messages);
+            home = 1;
+            // ensure that going back in history will go back to home
+            window.history.pushState({}, '', window.TINYCHAT_ROOT || '/');
+          " @touchstart="
+            otx = $event.changedTouches[0].clientX;
+          " @touchmove="
+            $el.style.setProperty('--tx', $event.changedTouches[0].clientX - otx);
+            $el.style.setProperty('--opacity', 1 - (Math.abs($event.changedTouches[0].clientX - otx) / trigger));
+          " @touchend="
+            if (Math.abs($event.changedTouches[0].clientX - otx) > trigger) removeHistory(_state);
+            $el.style.setProperty('--tx', 0);
+            $el.style.setProperty('--opacity', 1);
+          ">
+              <h3 x-text="new Date(_state.time).toLocaleString()"></h3>
+              <p x-text="$truncate(_state.messages[0].content, 80)"></p>
+              <!-- delete button -->
+              <button class="history-delete-button" @click.stop="removeHistory(_state);">
+                <i class=" fas fa-trash"></i>
+              </button>
+            </div>
+          </template>
+        </div>
+        <template x-if="histories.length">
+          <div class="histories-end"></div>
+        </template>
+      </div>
+    </div>
+    <div x-ref="messages" class="messages" x-init="
+      $watch('cstate', value => {
+        $el.innerHTML = '';
+        value.messages.forEach(({ role, content }) => {
+          const div = document.createElement('div');
+          div.className = `message message-role-${role}`;
+          try {
+            div.innerHTML = DOMPurify.sanitize(marked.parse(content));
+          } catch (e) {
+            console.log(content);
+            console.error(e);
+          }
+
+          // add a clipboard button to all code blocks
+          const codeBlocks = div.querySelectorAll('.hljs');
+          codeBlocks.forEach(codeBlock => {
+            const button = document.createElement('button');
+            button.className = 'clipboard-button';
+            button.innerHTML = '<i class=\'fas fa-clipboard\'></i>';
+            button.onclick = () => {
+              // navigator.clipboard.writeText(codeBlock.textContent);
+              const range = document.createRange();
+              range.setStartBefore(codeBlock);
+              range.setEndAfter(codeBlock);
+              window.getSelection()?.removeAllRanges();
+              window.getSelection()?.addRange(range);
+              document.execCommand('copy');
+              window.getSelection()?.removeAllRanges();
+
+              button.innerHTML = '<i class=\'fas fa-check\'></i>';
+              setTimeout(() => button.innerHTML = '<i class=\'fas fa-clipboard\'></i>', 1000);
+            };
+            codeBlock.appendChild(button);
+          });
+
+          $el.appendChild(div);
+        });
+
+        $el.scrollTo({ top: $el.scrollHeight, behavior: 'smooth' });
+      });
+    " x-intersect="
+      $el.scrollTo({ top: $el.scrollHeight, behavior: 'smooth' });
+    " x-show="home === 2" x-transition>
+    </div>
+    <div class="input-container">
+      <div class="input-performance">
+        <span class="input-performance-point">
+          <p class="monospace" x-text="(time_till_first / 1000).toFixed(2)"></p>
+          <p class="megrim-regular">SEC TO FIRST TOKEN</p>
+        </span>
+        <span class="input-performance-point">
+          <p class="monospace" x-text="tokens_per_second.toFixed(1)"></p>
+          <p class="megrim-regular">TOKENS/SEC</p>
+        </span>
+        <span class="input-performance-point">
+          <p class="monospace" x-text="total_tokens"></p>
+          <p class="megrim-regular">TOKENS</p>
+        </span>
+      </div>
+      <div class="loading-bar" x-show="loadingMessage !== ''">
+        <p class="loading-text" id="loading-message">Loading:</p>
+        <span id="progress-percentage">0%</span>
+        <div class="progress-bar">
+          <div class="progress"></div>
+        </div>
+      </div>
+      <div class="input" x-show="loadingMessage === ''">
+        <textarea x-ref="inputForm" id="input-form" class="input-form" autofocus rows=1 x-autosize
+          :placeholder="generating ? placeholderText : 'Say something'" :disabled="generating" @input="
+            home = (home === 0) ? 1 : home
+            if (cstate.messages.length === 0 && $el.value === '') home = -1;
+
+            if ($el.value !== '') {
+              const messages = [...cstate.messages];
+              messages.push({ role: 'user', content: $el.value });
+              updateTotalTokens(messages);
+            } else {
+              if (cstate.messages.length === 0) total_tokens = 0;
+              else updateTotalTokens(cstate.messages);
+            }
+          " x-effect="
+            console.log(generating);
+            if (!generating) $nextTick(() => {
+              $el.focus();
+              setTimeout(() => $refs.messages.scrollTo({ top: $refs.messages.scrollHeight, behavior: 'smooth' }), 100);
+            });
+          " @keydown.enter="await handleEnter($event)" @keydown.escape.window="$focus.focus($el)"></textarea>
+        <button class="input-button" :disabled="generating" @click="await handleSend()">
+          <i class="fas" :class="generating ? 'fa-spinner fa-spin' : 'fa-paper-plane'"></i>
+        </button>
+      </div>
+    </div>
+  </main>
+</body>
+
+</html>
@@ -0,0 +1,927 @@
+window.TINYCHAT_ROOT = "/tinychat-browser/";
+const queryParams = new URLSearchParams(window.location.search);
+const normalizedParams = Object.fromEntries([...queryParams].map(([key, value]) => [key.toUpperCase(), value.toUpperCase()]));
+window.BACKEND = (normalizedParams["BACKEND"] === "WASM") ? "WASM" : "WebGPU";
+const isMobileAgent = /Mobi|Android|iPhone|iPad|iPod/i.test(navigator.userAgent);
+const hasTouchScreen = 'ontouchstart' in window || navigator.maxTouchPoints > 0;
+window.isMobile = isMobileAgent || hasTouchScreen;
+if (window.isMobile) document.documentElement.classList.add('mobile'); // prevent annoying auto-zoom when entering prompt on mobile
+// MODEL_BASE_URL is where the weights are hosted, WEBGPU_EXPORT is the JS-wrapped WebGPU code exported from tinygrad
+window.PC_MODEL_BASE_URL = ".";
+window.PC_WEBGPU_EXPORT = './net.js'
+window.PC_MAX_CONTEXT = 1024;
+window.MOBILE_MODEL_BASE_URL = ".";
+window.MOBILE_WEBGPU_EXPORT = './net.js'
+window.MOBILE_MAX_CONTEXT = 1024;
+
+const tiktokenReady = (async () => {
+  const { init, get_encoding, Tiktoken, load } = await import('./tiktoken.js');
+  window.Tiktoken = Tiktoken;
+  window.tiktokenInit = init;
+  window.tiktokenGetEncoding = get_encoding;
+  window.tiktokenLoad = load;
+})();
+
+async function getDevice() {
+  let adapter;
+  try {
+    adapter = await navigator.gpu.requestAdapter();
+    if (!adapter) {
+      this.loadingMessage = "Loading WASM (WebGPU not enabled):";
+      throw new Error("No WebGPU adapter found");
+    }
+  } catch(error) {
+      this.loadingMessage = "Loading WASM (WebGPU not enabled):";
+    throw error;
+  }
+  const requiredLimits = {};
+  const maxBufferSize = 322122544;
+  requiredLimits.maxStorageBufferBindingSize = maxBufferSize;
+  requiredLimits.maxBufferSize = maxBufferSize;
+  requiredLimits.maxComputeInvocationsPerWorkgroup = 512; // may need to vary based on what the WEBGPU backend produces
+            
+  try {
+    return await adapter.requestDevice({ requiredLimits });
+  } catch(error) {
+    this.loadingMessage = "Loading WASM (WebGPU error):";
+    throw error;
+  }
+};
+
+// copied from examples/webgpu/stable_diffusion/index.html 
+function initDb() {
+  return new Promise((resolve, reject) => {
+    let db;
+    const request = indexedDB.open('tinydb', 1);
+    request.onerror = (event) => {
+      console.error('Database error:', event.target.error);
+      resolve(null);
+    };
+
+    request.onsuccess = (event) => {
+      db = event.target.result;
+      console.log("Db initialized.");
+      resolve(db);
+    };
+
+    request.onupgradeneeded = (event) => {
+      db = event.target.result;
+      if (!db.objectStoreNames.contains('tensors')) {
+        db.createObjectStore('tensors', { keyPath: 'id' });
+      }
+    };
+  });
+}
+
+// copied from examples/webgpu/stable_diffusion/index.html 
+function readTensorFromDb(db, id) {
+  return new Promise((resolve, reject) => {
+    if (db == null) {
+      resolve(null);
+    }
+            
+      const transaction = db.transaction(['tensors'], 'readonly');
+      const store = transaction.objectStore('tensors');
+      const request = store.get(id);
+
+      transaction.onabort = (event) => {
+        console.log("Transaction error while reading tensor: " + event.target.error);
+        resolve(null);
+      };
+
+      request.onsuccess = (event) => {
+        const result = event.target.result;
+        if (result) {
+          resolve(result);
+        } else {
+          resolve(null);
+        }
+      };
+
+      request.onerror = (event) => {
+        console.error('Tensor retrieve failed: ', event.target.error);
+        resolve(null);
+      };
+  });
+}
+
+function getAllKeysFromDb(db) {
+  return new Promise((resolve, reject) => {
+    if (db == null) {resolve([]);}
+    const transaction = db.transaction(['tensors'], 'readonly');
+    const store = transaction.objectStore('tensors');
+    const request = store.getAllKeys();
+    transaction.onabort = (event) => {
+      console.log("Transaction error while reading IndexedDB keys: " + event.target.error);
+      resolve([]);
+    };
+    request.onsuccess = function (event) {resolve(event.target.result);};
+    request.onerror = (event) => {
+      console.error('Retrieval of IndexedDB keys failed: ', event.target.error);
+      resolve([]);
+    };
+  });
+}
+
+// modified from examples/webgpu/stable_diffusion/index.html 
+function saveTensorToDb(db, id, tensor) {
+  return readTensorFromDb(db, id).then((result) => {
+    if (!result) {
+      new Promise((resolve, reject) => {
+        if (db == null) {
+          resolve(null);
+        }
+
+        const transaction = db.transaction(['tensors'], 'readwrite');
+        const store = transaction.objectStore('tensors');
+        const request = store.put({ id: id, content: tensor });
+
+        transaction.onabort = (event) => {
+          console.log("Transaction error while saving tensor: " + event.target.error);
+          resolve(null);
+        };
+
+        request.onsuccess = () => {
+          console.log('Tensor saved successfully.');
+          resolve();
+        };
+
+        request.onerror = (event) => {
+          console.error('Tensor save failed:', event.target.error);
+          resolve(null);
+        };
+      });
+    } else {
+      return null;
+    }
+  }).catch(()=> null);
+}
+
+function deleteTensorFromDb(db, id) {
+  return new Promise((resolve, reject) => {
+    if (db == null) {
+      console.error("Database is not initialized.");
+      resolve(null);
+      return;
+    }
+
+    const transaction = db.transaction(['tensors'], 'readwrite');
+    const store = transaction.objectStore('tensors');
+    const request = store.delete(id);
+
+    transaction.oncomplete = () => {
+      console.log(`Tensor with ID '${id}' deleted successfully.`);
+      resolve();
+    };
+
+    transaction.onerror = (event) => {
+      console.error("Transaction error while deleting tensor:", event.target.error);
+      resolve(null);
+    };
+
+    request.onerror = (event) => {
+      console.error('Tensor deletion failed:', event.target.error);
+      resolve(null);
+    };
+
+    request.onsuccess = () => {
+      console.log(`Delete request for tensor with ID '${id}' succeeded.`);
+    };
+  });
+}
+
+function makeProgress(total) {
+  let acc = 0;
+  const ret = function progress(amount, message) {
+    if (amount >= 0) { // allow updating message only
+      acc += amount;
+      const percentage = total ? Math.trunc((acc / total) * 100) : 0;
+      document.querySelector('.progress').style.width = `${percentage}%`;
+      document.getElementById('progress-percentage').textContent = `${percentage}%`;
+    }
+    if (message) {
+      this.loadingMessage = message;
+      document.getElementById('loading-message').textContent = this.loadingMessage;
+    }
+  }.bind(this);
+  ret.total = total;
+  return ret;
+}
+
+function sendMessageToWorker(worker, message) {
+  return new Promise((resolve, reject) => {
+    const onMessage = (event) => {
+      resolve(event.data);
+      worker.removeEventListener('message', onMessage);
+      worker.removeEventListener('error', onError);
+    };
+
+    const onError = (error) => {
+      reject(error);
+      worker.removeEventListener('message', onMessage);
+      worker.removeEventListener('error', onError);
+    };
+
+    worker.addEventListener('message', onMessage);
+    worker.addEventListener('error', onError);
+
+    if (message.header === "token") worker.postMessage(message.data);
+    else if (message.header === "load_state_dict") {
+      if (message.data === "done") worker.postMessage(message.data);
+      else worker.postMessage(message.data, message.data.map(file => file.bytes.buffer));
+    }
+    else if (message.header === "init") worker.postMessage("init");
+  });
+}
+
+async function load_state_dict (data, device, progress) {
+  let state_dict = data.metadata.state_dict;
+  let completed = 0;
+
+  // modified from examples/webgpu/stable_diffusion/index.html getProgressDlForPart
+  const loadPart = async (part) => {
+      const response = await fetch(part);
+      const res = new Response(new ReadableStream({
+          async start(controller) {
+              const reader = response.body.getReader();
+              for (;;) {
+                  const { done, value } = await reader.read();
+                  if (done) break;
+                  progress(value.byteLength);
+                  controller.enqueue(value);
+              }
+              controller.close();
+          },
+      }));
+        
+      return res.arrayBuffer();
+  };
+
+  let db = await initDb();
+
+  const getPart = async(filename, hash) => {
+    let part = await readTensorFromDb(db, hash);
+
+    if (part) {
+      console.log(`Cache hit: ${filename}, hash: ${hash}`);
+      progress(part.content.byteLength);
+      return Promise.resolve(part.content);
+    } else {
+      console.log(`Cache miss: ${filename}, hash: ${hash}`);
+      return loadPart(`${window.MODEL_BASE_URL}/${filename}`);
+    }
+  }
+
+  const correctHashes = data.metadata.files.map(file => file.hash)
+  // delete unused cached buffers to free disk space -- if we update weights, user will otherwise have obsolete cached buffers
+  const dbKeys = await getAllKeysFromDb(db);
+  const correctHashesSet = new Set(correctHashes);
+  const notInCorrectHashes = dbKeys.filter(key => !correctHashesSet.has(key));
+  // await these right before starting to save new stuff
+  const deletionPromises = notInCorrectHashes.map(async (hash) => deleteTensorFromDb(db, hash));
+
+  // instantiates empty weight buffers on WebGPU, attaches buffers to state_dict
+  let model;
+  if (window.BACKEND === "WebGPU") {
+    //model = await transformer().setup(device, state_dict, progress);
+    model = await transformer.setupNet(device, state_dict);
+    progress(0.15 * progress.total);
+
+  }
+  else if (window.BACKEND === "WASM") {
+    progress(0.02 * progress.total);
+    model = new Worker(`./worker.js?version=${Date.now()}`);
+    await sendMessageToWorker(model, {header: "init"});
+    progress(0.02 * progress.total);
+    progress(0.11 * progress.total);
+  }
+
+  const downloaded = [];
+  const triggerChainDownload = async (toDownload) => {
+    const numDownloaders = window.isMobile ? 4 : toDownload.length; // TODO: dynamically base this on DL file size? current assumption is 16 MiB chunks
+
+    const chainDownload = async() => {
+      const file = toDownload.shift();
+      loadPart(`${window.MODEL_BASE_URL}/${file.name}`) // triggers download
+      .then(async (arraybuf) => { 
+        downloaded.push({ ...file, bytes: new Uint8Array(arraybuf)});
+        // pause downloads if further processing is a bottleneck
+        while (toDownload.length && downloaded.length >= numDownloaders) await new Promise(resolve => setTimeout(resolve, 5));
+        if (toDownload.length && downloaded.length < numDownloaders) chainDownload(); // start next download
+      })
+    }
+    for (let i=0; i<numDownloaders; i++) if (toDownload.length) chainDownload();
+  }
+
+  const loadFileToStateDict = async(file) => {
+    if (window.BACKEND === "WebGPU") {
+      for (const part of file.parts) {
+        if (part.empty) continue;
+        part.bytes = (part.size === file.bytes.length) ? file.bytes : file.bytes.slice(part.file_start_pos, part.file_start_pos + part.size);
+        device.queue.writeBuffer(state_dict[part.key].bytes, part.target_start_pos, part.bytes); // improves stability over mappedAtCreation writing
+        part.bytes = null;
+      }
+    }
+    else if (window.BACKEND === "WASM") {
+      await sendMessageToWorker(model, {header: "load_state_dict", data: [file]});
+    }
+    file.bytes = null;
+  }
+
+  if (window.BACKEND === "WebGPU") { // contiguous loading not needed for WebGPU stability
+    const files = data.tensor_file_groups.flatMap(obj => obj.files);
+    data.tensor_file_groups = [{contiguous: false, files: files}];
+  }
+
+  for (const group of data.tensor_file_groups) {
+    const contiguous = group.contiguous;
+    const files = group.files;
+    const tensor_file_indices = files.map(file => file.index);
+    const contiguousFiles = [];
+    const fileHashes = new Set(files.map(file => file.hash));
+    const cachedFileHashes = new Set(dbKeys.filter(key => fileHashes.has(key)));
+    const cachedFiles = files.filter(file => cachedFileHashes.has(file.hash));
+    const toDownload = files.filter(file => !cachedFileHashes.has(file.hash));
+    triggerChainDownload(toDownload);
+
+    const loadDelay = 5;
+    await Promise.all(deletionPromises);
+
+    while (completed < files.length) {
+      const start = performance.now();
+      // prioritize files from downloaded queue, so we can continue downloading more files
+      if (downloaded.length) {
+        const file = downloaded.shift();
+        await saveTensorToDb(db, file.hash, file.bytes); // for wasm, must await to prevent race between indexedDB and transfer to worker
+        if (!contiguous) await loadFileToStateDict(file);
+        else contiguousFiles.push(file);
+        completed += 1;
+      }
+      else if (!downloaded.length && cachedFiles.length) {
+        const file = cachedFiles.shift();
+        file.bytes = await getPart(file.name, file.hash); // reads data from IndexedDB
+        if (!contiguous) await loadFileToStateDict(file);
+        else contiguousFiles.push(file);
+        completed += 1;
+      }
+      const end = performance.now();
+      const elapsed = end - start;
+      if (elapsed < loadDelay) await new Promise(resolve => setTimeout(resolve, loadDelay - elapsed));
+    }
+    if (contiguous) {
+      const orderMap = tensor_file_indices.reduce((acc, id, index) => {acc[id] = index; return acc;}, {});
+      contiguousFiles.sort((a, b) => orderMap[a.index] - orderMap[b.index]); // glue files together in the right order
+      await sendMessageToWorker(model, {header: "load_state_dict", data: contiguousFiles});
+    }
+    completed = 0;
+  }
+
+  // initialize empty kv_caches, which were part of exported model's state_dict, but which we didn't want to package/download
+  if (window.BACKEND === "WASM") {
+    for (const [k, v] of Object.entries(state_dict).filter(([_, v]) => v.empty === true)) {
+      v.parts[0].file_start_pos = 0;
+      const file = { parts: v.parts, size: v.size, bytes: new Uint8Array(v.size).fill(0) };
+      await loadFileToStateDict(file);
+    }
+  }
+
+  return model;
+};
+
+document.addEventListener("alpine:init", () => {
+  Alpine.data("state", () => ({
+    // loadingMessage updates the user on page load progress, including weights download and decompression
+    // if loadingMessage is not '', then prompt box will be hidden: this is default behavior on page load
+    placeholderText: "Generating...",
+    loadingMessage: `Loading ${window.BACKEND} model:`,
+    // model
+    nets: {},
+    tokenizer: null,
+    max_context: 1024,
+    lastSeenToks: [],
+
+    progress: null,
+
+    async init() {
+      var device = null;
+      var webgpuErrorMessage = null;
+      if (window.BACKEND === "WebGPU") {
+        try {
+          device = await getDevice.call(this);
+          console.log("WebGPU device initialized");
+        } catch (error) {
+          window.BACKEND = "WASM";
+          console.log(`error: ${error}\nFailed to launch WebGPU. Loading WASM model instead...`); // return;
+          webgpuErrorMessage = this.loadingMessage;
+        }
+      }
+
+      window.MODEL_BASE_URL = (window.BACKEND === "WebGPU" && !window.isMobile) ? window.PC_MODEL_BASE_URL : window.MOBILE_MODEL_BASE_URL;
+      this.max_context = (window.BACKEND === "WebGPU" && !window.isMobile) ? window.PC_MAX_CONTEXT : window.MOBILE_MAX_CONTEXT;
+
+      const kernelsReady = (async () => {
+        if (window.BACKEND === "WASM") {var exports = await import(`./net_clang.js?version=${Date.now()}`);}
+        else if (window.BACKEND === "WebGPU" && !window.isMobile) {var exports = await import(`${PC_WEBGPU_EXPORT}?version=${Date.now()}`);}
+        else if (window.BACKEND === "WebGPU" && window.isMobile) {var exports = await import(`${MOBILE_WEBGPU_EXPORT}?version=${Date.now()}`);}
+        self.transformer = exports.default;
+      })();
+
+      const response = await fetch(`${window.MODEL_BASE_URL}/net_metadata.json`);
+      // TODO: cache metadata (and everything else, including tokenizer)
+      // TODO: use service worker to reload page when offline
+      const data = await response.json();
+      data.metadata.files = data.metadata.files.map((file, index) => ({...file, index}));
+      const state_dict = data.metadata.state_dict;
+
+      /*
+      - allocating memory to WASM on mobile has longstanding issues: https://github.com/WebAssembly/design/issues/1397
+
+      - the below pattern, while yielding a succesfully-functioning model when it doesn't crash, causes regular crashes on iOS Safari (iphone 15 iOS 18.3):
+        - call WASM malloc (to fit all tensors, or one per tensor) for all tensors up front, then load tensor byte chunks into the buffers in random order
+
+      - the below pattern has been stable on iOS Safari (iphone 15 iOS 18.3):
+        - call only one WASM malloc at a time before filling the allocated bytes, as small as possible (malloc up to 256 MiB has been tested)
+        - fill the malloc'd memory in linear order from start to end (what has been tested is calling wasm.HEAPU8.set on 16 MiB chunks from start to end)
+        - use ALLOW_MEMORY_GROWTH=1 in wasm compilation, minimize initial memory
+
+      - additional considerations affecting loading design, for WASM:
+        - it seems that copying bytes into wasm memory cannot be zero-copy without sharedarraybuffer, which isn't currently used due to increased hosting complexity
+        - non-zero copies create memory pressure, which is not reliably capped because of lack of control over garbage collection
+        - to minimize peak memory pressure if GC is delayed, we process (i.e. download + copy into WASM) large tensors (> 16 MiB) one at a time, in descending size order
+      */
+      data.tensor_file_groups = []; // see above: for WASM, limit processing of multi-file Tensors to one at a time, in descending order based on Tensor size
+      const unsplit_tensors = [];
+      const sortedEntries = Object.entries(state_dict).sort(([, objA], [, objB]) => objB.size - objA.size);
+
+      let totalSize = 0;
+      const seen = new Set();
+      for (const [k,v] of sortedEntries) {
+        const files_in_tensor = [];
+        for (const part of v.parts) {
+          part.key = k;
+          if (part.empty) state_dict[k].empty = true; // assumes no other parts of this weight exist and are non-empty
+          else {
+            const file = data.metadata.files[part.file];
+            if (!seen.has(file.index)) {
+              seen.add(file.index);
+              files_in_tensor.push(file);
+            }
+            totalSize += part.size;
+            part.dtype = v.dtype;
+            if (!data.metadata.files[part.file].parts) data.metadata.files[part.file].parts = [];
+            data.metadata.files[part.file].size ??= 0;
+            data.metadata.files[part.file].size += part.size;
+            data.metadata.files[part.file].parts.push(part);
+          }
+        }
+        if (files_in_tensor.length > 1) data.tensor_file_groups.push({contiguous: true, files: files_in_tensor}); // [tensorN_file0, tensorN_file1, ...]
+        else if (files_in_tensor.length > 0) unsplit_tensors.push(files_in_tensor[0]);
+      }
+      data.tensor_file_groups.push({contiguous: false, files: unsplit_tensors});
+
+      data.totalSize = totalSize;
+      totalSize = totalSize / 0.8; // give space in progress bar for initializing model bufs, and tokenizer
+      this.progress = makeProgress.call(this, totalSize); // creates closure with totalSize
+
+      try {
+        this.progress(0.01 * totalSize, "Loading tokenizer:");
+        const wasmResponse = await fetch(`${window.MODEL_BASE_URL}/tiktoken_bg.wasm`);
+        this.progress(0.01 * totalSize);
+        const wasmBytes = await wasmResponse.arrayBuffer();
+        await tiktokenReady;
+        await window.tiktokenInit((imports) => WebAssembly.instantiate(wasmBytes, imports));
+        this.progress(0.01 * totalSize);
+
+        this.tokenizer = await createTokenizer(`${window.MODEL_BASE_URL}/llama3-2.tiktoken`);
+        const tokenizer_works = (new TextDecoder().decode(this.tokenizer.decode(this.tokenizer.encode("hello world"))) === "hello world");
+        console.log("tokenizer works:", tokenizer_works)
+        this.progress(0.01 * totalSize);
+      } catch (error) {this.progress(-1, `Error launching tokenizer: ${error}`); console.log(error); return;}
+
+      try {
+        const loadModelMessage = (webgpuErrorMessage) ? webgpuErrorMessage : `Loading ${window.BACKEND} model:`
+        this.progress(0, loadModelMessage);
+        await kernelsReady;
+        const model = await load_state_dict(data, device, this.progress);
+
+        if (window.BACKEND === "WebGPU") {
+          this.nets = {"transformer": model};
+        }
+        else if (window.BACKEND === "WASM") {
+          const msg = await sendMessageToWorker(model, {header: "load_state_dict", data: "done"});
+          this.nets = {"transformer": async (tok, start_pos) => sendMessageToWorker(model, {header: "token", data: [tok, start_pos]})};
+        }
+        this.progress(0.01 * totalSize, `Launching ${window.BACKEND} model:`);
+        this.loadingMessage = ""; // Triggers removal of loading bar, display of prompt box
+      } catch (error) {this.progress(-1, `Error launching model: ${error}`); console.log(error); return;}
+    },
+
+    // current state
+    cstate: {
+      time: null,
+      messages: [],
+    },
+
+    // historical state
+    histories: JSON.parse(localStorage.getItem("histories")) || [],
+
+    home: 0,
+    generating: false,
+    maxContextReached: false,
+    cancelGeneration: false,
+    endpoint: `${window.location.origin}/v1`,
+
+    // performance tracking
+    time_till_first: 0,
+    tokens_per_second: 0,
+    total_tokens: 0,
+    max_context: 0,
+
+    removeHistory(cstate) {
+      const index = this.histories.findIndex((state) => {
+        return state.time === cstate.time;
+      });
+      if (index !== -1) {
+        this.histories.splice(index, 1);
+        localStorage.setItem("histories", JSON.stringify(this.histories));
+      }
+    },
+
+    async handleSend() {
+      const el = document.getElementById("input-form");
+      const value = el.value.trim();
+      if (!value) return;
+
+      if (this.generating) return;
+      this.maxContextReached = false;
+      this.placeholderText = "Generating...";
+      this.generating = true;
+      this.cancelGeneration = false;
+      if (this.home === 0) this.home = 1;
+
+      // ensure that going back in history will go back to home
+      window.history.pushState({}, "", window.TINYCHAT_ROOT || "/");
+
+      // add message to list
+      this.cstate.messages.push({ role: "user", content: value });
+
+      // clear textarea
+      el.value = "";
+      el.style.height = "auto";
+      el.style.height = el.scrollHeight + "px";
+
+      // reset performance tracking
+      const prefill_start = Date.now();
+      let start_time = 0;
+      let tokens = 0;
+      this.tokens_per_second = 0;
+
+      let gottenFirstChunk = false;
+      try {
+        for await (
+          const chunk of this.openaiChatCompletion(this.cstate.messages)
+        ) {
+          if (!gottenFirstChunk) {
+            this.cstate.messages.push({ role: "assistant", content: "" });
+            gottenFirstChunk = true;
+          }
+
+          // add chunk to the last message
+          // TODO: handle localStorage overflow
+          //   possible example: this.cstate.messages[...] was undefined when trying to prompt within an old cstate (chat session)
+          this.cstate.messages[this.cstate.messages.length - 1].content += chunk;
+
+          // calculate performance tracking
+          tokens += 1;
+          this.total_tokens += 1;
+          if (start_time === 0) {
+            start_time = Date.now();
+            this.time_till_first = start_time - prefill_start;
+          } else {
+            const diff = Date.now() - start_time;
+            if (diff > 0) {
+              this.tokens_per_second = tokens / (diff / 1000);
+            }
+          }
+          this.checkMaxContext(this.total_tokens);
+          if (this.cancelGeneration) break;
+        }
+      } finally {
+        // update the state in histories or add it if it doesn't exist
+        const index = this.histories.findIndex((cstate) => {
+          return cstate.time === this.cstate.time;
+        });
+        this.cstate.time = Date.now();
+        if (index !== -1) {
+          // update the time
+          this.histories[index] = this.cstate;
+        } else {
+          this.histories.push(this.cstate);
+        }
+        // update in local storage
+        localStorage.setItem("histories", JSON.stringify(this.histories));
+
+        if (!this.maxContextReached) this.generating = false;
+        if (this.cancelGeneration && !this.maxContextReached) this.cstate = { time: null, messages: [] };
+      }
+    },
+
+    async handleEnter(event) {
+      // if shift is not pressed
+      if (!event.shiftKey) {
+        event.preventDefault();
+        await this.handleSend();
+      }
+    },
+
+    updateTotalTokens(messages) {
+      try {
+        let toks = [this.tokenizer.bos_id];
+        messages.forEach((message) => {
+          if (!message.role || !message.content) {
+            throw new Error("Each message must have a 'role' and 'content' property.");
+          }
+          toks = toks.concat(this.tokenizer.encodeMessage(message.role, message.content));
+
+          if (messages.length > 0 && messages[messages.length - 1].role === "user") {
+            toks = toks.concat(this.tokenizer.encodeRole("assistant"));
+          }
+          this.total_tokens = toks.length;
+        });
+      } catch (error) {
+        console.error("Error updating total tokens:", error);
+      }
+    },
+
+    checkMaxContext(num_tokens) {
+      if (num_tokens >= this.max_context) {
+        this.cancelGeneration = true;
+        this.maxContextReached = true;
+        this.placeholderText = `Max context reached: ${this.max_context} tokens`;
+      }
+    },
+
+    async *openaiChatCompletion(messages) {
+      let tokens = [this.tokenizer.bos_id];
+      for (const message of messages) {
+        tokens = tokens.concat(this.tokenizer.encodeMessage(message.role, message.content));
+      }
+      tokens = tokens.concat(this.tokenizer.encodeRole("assistant"));
+      this.checkMaxContext(tokens.length); // don't waste time prefilling if we know we're over the token limit
+      let startPos = 0
+      const prefillToks = tokens.slice(0, -1);
+
+      // Skip the largest possible sequence of tokens already represented at the beginning of the model's kv caches
+      for (let i=0; i <= prefillToks.length; i++) {
+        startPos = i;
+        if (i == prefillToks.length) break;
+        if (i == this.lastSeenToks.length) break;
+        if (prefillToks[i] !== this.lastSeenToks[i]) break;
+      }
+      //this.lastSeenToks = prefillToks;
+      //prefillToks = prefillToks.slice(startPos);
+      const unprocessedPrefillToks = prefillToks.slice(startPos);
+      this.lastSeenToks = prefillToks.slice(0, startPos);
+
+      this.progress = makeProgress(unprocessedPrefillToks.length);
+      this.loadingMessage = (window.BACKEND === "WebGPU") ? "Reading input:" : "Loading (enable WebGPU for speed):";
+      this.progress(0, this.loadingMessage);
+      for (const tok of unprocessedPrefillToks) {
+        if (this.cancelGeneration) {this.loadingMessage=""; return;}
+        if (window.BACKEND === "WebGPU") {await this.nets["transformer"](new Int32Array([tok]), new Int32Array([startPos]));}
+        else {await this.nets["transformer"](tok, startPos);}
+        this.lastSeenToks.push(tok)
+        startPos += 1;
+        this.progress(1);
+      }
+      this.loadingMessage = ""; // hides progress bar
+
+      let lastTok = tokens[tokens.length - 1];
+      while (true) {
+        if (window.BACKEND === "WebGPU") {var tok = await this.nets["transformer"](new Int32Array([lastTok]), new Int32Array([startPos])); tok = tok[0][0];}
+        else {var tok = await this.nets["transformer"](lastTok, startPos);}
+        this.lastSeenToks.push(lastTok); // lets us skip prefilling with these tokens at the next prompt in this chain
+        startPos += 1;
+        lastTok = tok;
+        if (this.tokenizer.stop_tokens.has(lastTok)) break;
+        yield new TextDecoder().decode(this.tokenizer.decode([lastTok]));
+      }
+    },
+  }));
+});
+
+const { markedHighlight } = globalThis.markedHighlight;
+marked.use(markedHighlight({
+  langPrefix: "hljs language-",
+  highlight(code, lang, _info) {
+    const language = hljs.getLanguage(lang) ? lang : "plaintext";
+    return hljs.highlight(code, { language }).value;
+  },
+}));
+
+// **** eventsource-parser ****
+class EventSourceParserStream extends TransformStream {
+  constructor() {
+    let parser;
+
+    super({
+      start(controller) {
+        parser = createParser((event) => {
+          if (event.type === "event") {
+            controller.enqueue(event);
+          }
+        });
+      },
+
+      transform(chunk) {
+        parser.feed(chunk);
+      },
+    });
+  }
+}
+
+function createParser(onParse) {
+  let isFirstChunk;
+  let buffer;
+  let startingPosition;
+  let startingFieldLength;
+  let eventId;
+  let eventName;
+  let data;
+  reset();
+  return {
+    feed,
+    reset,
+  };
+  function reset() {
+    isFirstChunk = true;
+    buffer = "";
+    startingPosition = 0;
+    startingFieldLength = -1;
+    eventId = void 0;
+    eventName = void 0;
+    data = "";
+  }
+  function feed(chunk) {
+    buffer = buffer ? buffer + chunk : chunk;
+    if (isFirstChunk && hasBom(buffer)) {
+      buffer = buffer.slice(BOM.length);
+    }
+    isFirstChunk = false;
+    const length = buffer.length;
+    let position = 0;
+    let discardTrailingNewline = false;
+    while (position < length) {
+      if (discardTrailingNewline) {
+        if (buffer[position] === "\n") {
+          ++position;
+        }
+        discardTrailingNewline = false;
+      }
+      let lineLength = -1;
+      let fieldLength = startingFieldLength;
+      let character;
+      for (
+        let index = startingPosition;
+        lineLength < 0 && index < length;
+        ++index
+      ) {
+        character = buffer[index];
+        if (character === ":" && fieldLength < 0) {
+          fieldLength = index - position;
+        } else if (character === "\r") {
+          discardTrailingNewline = true;
+          lineLength = index - position;
+        } else if (character === "\n") {
+          lineLength = index - position;
+        }
+      }
+      if (lineLength < 0) {
+        startingPosition = length - position;
+        startingFieldLength = fieldLength;
+        break;
+      } else {
+        startingPosition = 0;
+        startingFieldLength = -1;
+      }
+      parseEventStreamLine(buffer, position, fieldLength, lineLength);
+      position += lineLength + 1;
+    }
+    if (position === length) {
+      buffer = "";
+    } else if (position > 0) {
+      buffer = buffer.slice(position);
+    }
+  }
+  function parseEventStreamLine(lineBuffer, index, fieldLength, lineLength) {
+    if (lineLength === 0) {
+      if (data.length > 0) {
+        onParse({
+          type: "event",
+          id: eventId,
+          event: eventName || void 0,
+          data: data.slice(0, -1),
+          // remove trailing newline
+        });
+
+        data = "";
+        eventId = void 0;
+      }
+      eventName = void 0;
+      return;
+    }
+    const noValue = fieldLength < 0;
+    const field = lineBuffer.slice(
+      index,
+      index + (noValue ? lineLength : fieldLength),
+    );
+    let step = 0;
+    if (noValue) {
+      step = lineLength;
+    } else if (lineBuffer[index + fieldLength + 1] === " ") {
+      step = fieldLength + 2;
+    } else {
+      step = fieldLength + 1;
+    }
+    const position = index + step;
+    const valueLength = lineLength - step;
+    const value = lineBuffer.slice(position, position + valueLength).toString();
+    if (field === "data") {
+      data += value ? "".concat(value, "\n") : "\n";
+    } else if (field === "event") {
+      eventName = value;
+    } else if (field === "id" && !value.includes("\0")) {
+      eventId = value;
+    } else if (field === "retry") {
+      const retry = parseInt(value, 10);
+      if (!Number.isNaN(retry)) {
+        onParse({
+          type: "reconnect-interval",
+          value: retry,
+        });
+      }
+    }
+  }
+}
+const BOM = [239, 187, 191];
+function hasBom(buffer) {
+  return BOM.every((charCode, index) => buffer.charCodeAt(index) === charCode);
+}
+
+const PAT_STR = "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+";
+
+async function createTokenizer(bpeUrl) {
+  const num_base_tokens = 128000;
+  const special_tokens = {
+    "<|begin_of_text|>": 128000,
+    "<|end_of_text|>": 128001,
+    "<|start_header_id|>": 128006,
+    "<|end_header_id|>": 128007,
+    "<|eot_id|>": 128009
+  };
+  const model = await window.tiktokenLoad({
+        "load_tiktoken_bpe": bpeUrl,
+        "special_tokens": special_tokens,
+        "pat_str": PAT_STR
+    });
+  const tokenizer = new window.Tiktoken(model.bpe_ranks, model.special_tokens, model.pat_str)
+
+  return {
+    get bos_id() {
+      return special_tokens["<|begin_of_text|>"];
+    },
+
+    get stop_tokens() {
+      return new Set([
+        special_tokens["<|end_of_text|>"],
+        special_tokens["<|eot_id|>"],
+      ]);
+    },
+
+    decode(toks) {
+      const filtered = toks.filter((t) => t < num_base_tokens);
+      return tokenizer.decode(filtered);
+    },
+
+    encode(text, allow_special = false) {
+      const allowedSpecial = allow_special ? "all" : new Set();
+      const disallowedSpecial = new Set();
+      return tokenizer.encode(text, allowedSpecial, disallowedSpecial);
+    },
+
+    encodeRole(role) {
+      const tokens = [];
+      tokens.push(special_tokens["<|start_header_id|>"]);
+      tokens.push(...this.encode(role));
+      tokens.push(special_tokens["<|end_header_id|>"]);
+      tokens.push(...this.encode("\n\n"));
+      return tokens;
+    },
+
+    encodeMessage(role, content) {
+      const roleTokens = this.encodeRole(role);
+      const contentTokens = this.encode(content.trim());
+      return [...roleTokens, ...contentTokens, special_tokens["<|eot_id|>"]];
+    },
+  };
+}
@@ -0,0 +1,11 @@
+#!/usr/bin/env bash
+cd "$(dirname "$0")"
+npm init -y && \
+npm install --save-dev webpack webpack-cli && \
+npm install tiktoken && \
+jq '.scripts.build = "webpack"' package.json > package.tmp.json && \
+mv package.tmp.json package.json && \
+npm run build && \
+mv dist/*.wasm ./tiktoken_bg.wasm && \
+mv dist/* ./ && \
+rm -rf dist node_modules package-lock.json package.json
@@ -0,0 +1,5 @@
+// Force Webpack to copy the WASM
+import 'tiktoken/tiktoken_bg.wasm'; 
+import { init, get_encoding, encoding_for_model, Tiktoken } from 'tiktoken/init';
+import { load } from 'tiktoken/load';
+export { init, get_encoding, encoding_for_model, Tiktoken, load };
@@ -0,0 +1,25 @@
+const path = require("path");
+
+module.exports = {
+  mode: "production",
+  entry: "./tiktoken-export.js",
+  output: {
+    filename: "tiktoken.js",
+    path: path.resolve(__dirname, "dist"),
+    library: {
+      type: "module"
+    }
+  },
+  experiments: {
+    outputModule: true,
+    asyncWebAssembly: true
+  },
+  module: {
+    rules: [
+      {
+        test: /\.wasm$/,
+        type: "asset/resource",
+      }
+    ]
+  }
+};
@@ -0,0 +1,62 @@
+const kernelsReady = (async () => {
+  // can't get browser to use updated versions except with cache-busting query string
+  const exports = await import(`./net_clang.js?version=${Date.now()}`);
+  Object.assign(self, exports);
+})();
+
+async function init(event) {
+  await kernelsReady;
+  self.model = await self.transformer();
+  self.addEventListener("message", loadStateDict);
+  self.removeEventListener("message", init);
+  self.postMessage("success");
+}
+
+function loadStateDict(event) {
+  if (event.data === "done") {
+    self.addEventListener("message", inference);
+    self.removeEventListener("message", loadStateDict);
+  }
+  else {
+    if (event.data.length > 1) {
+      // the bytes from files are set contiguously in WASM memory
+      const malloc_size = event.data.reduce((sum, file) => sum + file.bytes.length, 0);
+      const malloc_ptr = self.model.wasm._malloc(malloc_size);
+      let cursor = 0;
+      for (const file of event.data) {
+        self.model.wasm.HEAPU8.set(file.bytes, malloc_ptr + cursor);
+        for (const part of file.parts) {
+          if (part.target_start_pos === 0) {
+            // tell WASM code where the tensor is in memory
+            self.model.wasm._set_buf(self.transformer_name_to_id[part.key], malloc_ptr + cursor);
+          }
+          cursor += part.size;
+        }
+        file.bytes = null;
+      }
+    }
+    else {
+      // the bytes from files are not guaranteed to be set contiguously in WASM memory
+      const file = event.data[0];
+      const malloc_ptr = self.model.wasm._malloc(file.size);
+      self.model.wasm.HEAPU8.set(file.bytes, malloc_ptr);
+      for (const part of file.parts) {
+        if (part.target_start_pos === 0) {
+          self.model.wasm._set_buf(self.transformer_name_to_id[part.key], malloc_ptr + part.file_start_pos);
+        }
+      }
+      file.bytes = null;
+    }
+  }
+  self.postMessage("success");
+}
+
+function inference(event) {
+  const [tok, start_pos] = event.data;
+  const int32tok = new Int32Array([tok]);
+  const model_out = self.model.run(new Uint8Array(int32tok.buffer), start_pos);
+  const int32nextTok = new Int32Array(model_out[0].buffer);
+  self.postMessage(int32nextTok[0]);
+}
+
+self.addEventListener("message", init);
@@ -0,0 +1,38 @@
+#!POPCORN leaderboard grayscale
+#!POPCORN gpu A100
+# not a stable API, but works
+
+import torch, functools
+from tinygrad import Tensor, TinyJit, Device
+from tinygrad.engine.realize import CompiledRunner
+from tinygrad.helpers import get_single_element, Context, OSX
+from tinygrad.dtype import _from_torch_dtype
+
+@TinyJit
+def f(tg_out, tg_data): return tg_out.assign(tg_data[:, :, 0] * 0.2989 + tg_data[:, :, 1] * 0.5870 + tg_data[:, :, 2] * 0.1140).realize()
+
+def custom_kernel(data: torch.Tensor, device="CUDA") -> torch.Tensor:
+  assert data.dtype == torch.float32
+  tg_data = Tensor.from_blob(data.data_ptr(), data.shape, dtype=_from_torch_dtype(data.dtype), device=device)
+
+  out = torch.empty((data.shape[0], data.shape[1]), dtype=data.dtype, device=data.device)
+  tg_out = Tensor.from_blob(out.data_ptr(), out.shape, dtype=_from_torch_dtype(out.dtype), device=device)
+
+  # Need to sync torch to make sure the data is valid.
+  if data.device.type == "mps": torch.mps.synchronize()
+  else: torch.cuda.synchronize()
+
+  with Context(BEAM=2): f(tg_out, tg_data)
+
+  # Wait for computation to finish and the data is valid.
+  Device[device].synchronize()
+
+  return out
+
+if __name__ == "__main__":
+  for i in range(3):
+    if OSX:
+      out = custom_kernel(inp:=torch.rand(16, 16, 3, device=torch.device("mps")), device="METAL")
+    else:
+      out = custom_kernel(inp:=torch.rand(16, 16, 3, device=torch.device("cuda")), device="CUDA")
+    assert torch.allclose(out, inp[:, :, 0] * 0.2989 + inp[:, :, 1] * 0.5870 + inp[:, :, 2] * 0.1140)
@@ -17,7 +17,7 @@ def make_dataset():
  random.shuffle(ds)
  ds = np.array(ds).astype(np.float32)
  ds_X = ds[:, 0:6]
-  ds_Y = np.copy(ds[:, 1:])
+  ds_Y = np.copy(ds[:, 1:]).astype(np.int32)
  ds_X_train, ds_X_test = ds_X[0:8000], ds_X[8000:]
  ds_Y_train, ds_Y_test = ds_Y[0:8000], ds_Y[8000:]
  return ds_X_train, ds_Y_train, ds_X_test, ds_Y_test
@@ -105,12 +105,12 @@ class Vgg7:
    Output format: (1, 3, Y - 14, X - 14)
    (the - 14 represents the 7-pixel context border that is lost)
    """
-    x = self.conv1.forward(x).leakyrelu(0.1)
-    x = self.conv2.forward(x).leakyrelu(0.1)
-    x = self.conv3.forward(x).leakyrelu(0.1)
-    x = self.conv4.forward(x).leakyrelu(0.1)
-    x = self.conv5.forward(x).leakyrelu(0.1)
-    x = self.conv6.forward(x).leakyrelu(0.1)
+    x = self.conv1.forward(x).leaky_relu(0.1)
+    x = self.conv2.forward(x).leaky_relu(0.1)
+    x = self.conv3.forward(x).leaky_relu(0.1)
+    x = self.conv4.forward(x).leaky_relu(0.1)
+    x = self.conv5.forward(x).leaky_relu(0.1)
+    x = self.conv6.forward(x).leaky_relu(0.1)
    x = self.conv7.forward(x)
    return x

@@ -193,10 +193,10 @@ class Generator:
    x = self.conv_pre(x)
    if g is not None:  x = x + self.cond(g)
    for i in range(self.num_upsamples):
-      x = self.ups[i](x.leakyrelu(LRELU_SLOPE))
+      x = self.ups[i](x.leaky_relu(LRELU_SLOPE))
      xs = sum(self.resblocks[i * self.num_kernels + j].forward(x) for j in range(self.num_kernels))
      x = (xs / self.num_kernels).realize()
-    res = self.conv_post(x.leakyrelu()).tanh().realize()
+    res = self.conv_post(x.leaky_relu()).tanh().realize()
    return res

 class LayerNorm(nn.LayerNorm):
@@ -238,8 +238,8 @@ class ResBlock1:
    self.convs2 = [nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)) for _ in range(3)]
  def forward(self, x: Tensor, x_mask=None):
    for c1, c2 in zip(self.convs1, self.convs2):
-      xt = x.leakyrelu(LRELU_SLOPE)
-      xt = c1(xt if x_mask is None else xt * x_mask).leakyrelu(LRELU_SLOPE)
+      xt = x.leaky_relu(LRELU_SLOPE)
+      xt = c1(xt if x_mask is None else xt * x_mask).leaky_relu(LRELU_SLOPE)
      x = c2(xt if x_mask is None else xt * x_mask) + x
    return x if x_mask is None else x * x_mask

@@ -18,10 +18,6 @@ canvas {  display: none; }
 </style>
 <title>tinygrad has WebGPU</title>
 <link rel="icon" type="image/x-icon" href="https://raw.githubusercontent.com/tinygrad/tinygrad/master/docs/logo.png">
-<script type="module">
-	import model from "../../net.js";
-	window.model = model;
-</script>
 </head>
 <body>
 <h1>WebGPU <a href="https://github.com/geohot/tinygrad">tinygrad</a> EfficientNet!</h1>
@@ -102,6 +98,7 @@ canvas {  display: none; }
 			resultText.innerHTML = "loading..."
 			labels = await getLabels();
 			const device = await getDevice();
+			const model = (await import("../../net.js")).default;
 			net = await timer(() => model.load(device, '../../net.safetensors'), "(compilation)");
 			resultText.innerHTML = "ready"
 		} catch (e) {
@@ -12,7 +12,7 @@ if __name__ == "__main__":
    yolo_infer = YOLOv8(w=0.25, r=2.0, d=0.33, num_classes=80)
    state_dict = safe_load(get_weights_location(yolo_variant))
    load_state_dict(yolo_infer, state_dict)
-    prg, inp_sizes, out_sizes, state = export_model(yolo_infer, Device.DEFAULT.lower(), Tensor.randn(1,3,416,416), model_name="yolov8")
+    prg, inp_sizes, out_sizes, state = export_model(yolo_infer, Device.DEFAULT.lower(), Tensor.randn(1,3,640,640), model_name="yolov8")
    dirname = Path(__file__).parent
    safe_save(state, (dirname / "net.safetensors").as_posix())
    with open(dirname / f"net.js", "w") as text_file:
@@ -111,7 +111,7 @@
    </div>
    <script>
        let net = null;
-        const modelInputSize = 416;
+        const modelInputSize = 640;
        let lastCalledTime;
        let fps = 0, accumFps = 0, frameCounter = 0;

@@ -127,14 +127,22 @@
        const offscreenContext = offscreenCanvas.getContext('2d');


-        if (navigator.mediaDevices && navigator.mediaDevices.getUserMedia) {
-            navigator.mediaDevices.getUserMedia({ audio: false,  video: { facingMode: { ideal: "environment" }}}).then(function (stream) {
+        if (navigator.mediaDevices?.getUserMedia) {
+            const tryCamera = facing => navigator.mediaDevices.getUserMedia({ audio: false, video: { facingMode: { ideal: facing } } });
+            const handle = stream => {
                video.srcObject = stream;
                video.onloadedmetadata = function() {
                    canvas.width = video.clientWidth;
                    canvas.height = video.clientHeight;
-                }
-            });
+                };
+            };
+            tryCamera("environment").then(handle).catch(() => 
+                tryCamera("user").then(handle).catch(e => {
+                    wgpuError.textContent = "Error: Could not access camera."
+                    wgpuError.style.display = "block";
+                    loadingContainer.style.display = "none";
+                })
+            );
        }

        async function processFrame() {
@@ -176,10 +184,12 @@
            offscreenContext.clearRect(0, 0, modelInputSize, modelInputSize);
            offscreenContext.drawImage(video, offsetX, offsetY, targetWidth, targetHeight);
            const boxes = await detectObjectsOnFrame(offscreenContext);
-            drawBoxes(offscreenCanvas, boxes, targetWidth, targetHeight, offsetX, offsetY);
+            const validBoxes = [];
+            for (let i = 0; i < boxes.length; i += 6)
+                if (boxes[i + 4] > 0) validBoxes.push([boxes[i], boxes[i + 1], boxes[i + 2], boxes[i + 3], boxes[i + 5]]);
+            drawBoxes(offscreenCanvas, validBoxes, targetWidth, targetHeight, offsetX, offsetY);
            requestAnimationFrame(processFrame);
        }
-
        requestAnimationFrame(processFrame);

        function drawBoxes(offscreenCanvas, boxes, targetWidth, targetHeight, offsetX, offsetY) {
@@ -190,8 +200,8 @@
            const scaleX = canvas.width / targetWidth;
            const scaleY = canvas.height / targetHeight;

-            boxes.forEach(([x1, y1, x2, y2, label]) => {
-                const classIndex = yolo_classes.indexOf(label);
+            boxes.forEach(([x1, y1, x2, y2, classIndex]) => {
+                const label = yolo_classes[classIndex];
                const color = classColors[classIndex];
                ctx.strokeStyle = color;
                ctx.fillStyle = color;
@@ -219,21 +229,13 @@
                net = await yolov8.load(device, "./net.safetensors");
                loadingContainer.style.display = "none";
            }
-            let start = performance.now();
-            const [input,img_width,img_height] = await prepareInput(offscreenContext);
-            console.log("Preprocess took: " + (performance.now() - start) + " ms");
-            start = performance.now();
+            const input = await prepareInput(offscreenContext);
            const output = await net(new Float32Array(input));
-            console.log("Inference took: " + (performance.now() - start) + " ms");
-            start = performance.now();
-            let out = processOutput(output[0],img_width,img_height);
-            console.log("Postprocess took: " + (performance.now() - start) + " ms");
-            return out;
+            return output[0];
        }

        async function prepareInput(offscreenContext) {
            return new Promise(resolve => {
-                const [img_width,img_height] = [modelInputSize, modelInputSize]
                const imgData = offscreenContext.getImageData(0,0,modelInputSize,modelInputSize);
                const pixels = imgData.data;
                const red = [], green = [], blue = [];
@@ -244,7 +246,7 @@
                    blue.push(pixels[index+2]/255.0);
                }
                const input = [...red, ...green, ...blue];
-                resolve([input, img_width, img_height])
+                resolve(input)
            })
        }

@@ -252,62 +254,10 @@
            if (!navigator.gpu) return false;
            const adapter = await navigator.gpu.requestAdapter();
            return await adapter.requestDevice({
-                requiredFeatures: ["shader-f16"],
                powerPreference: "high-performance"
 		    });
        };
        
-        function processOutput(output, img_width, img_height) {
-            let boxes = [];
-            const numPredictions = Math.pow(modelInputSize/32, 2) * 21;
-            for (let index=0;index<numPredictions;index++) {
-                const [class_id,prob] = [...Array(80).keys()]
-                    .map(col => [col, output[numPredictions*(col+4)+index]])
-                    .reduce((accum, item) => item[1]>accum[1] ? item : accum,[0,0]);
-
-                if (prob < 0.25) continue;
-                const label = yolo_classes[class_id];
-                const xc = output[index];
-                const yc = output[numPredictions+index];
-                const w = output[2*numPredictions+index];
-                const h = output[3*numPredictions+index];
-                const x1 = (xc-w/2)/modelInputSize*img_width;
-                const y1 = (yc-h/2)/modelInputSize*img_height;
-                const x2 = (xc+w/2)/modelInputSize*img_width;
-                const y2 = (yc+h/2)/modelInputSize*img_height;
-                boxes.push([x1,y1,x2,y2,label,prob]);
-            }
-
-            boxes = boxes.sort((box1,box2) => box2[5]-box1[5])
-            const result = [];
-            while (boxes.length>0) {
-                result.push(boxes[0]);
-                boxes = boxes.filter(box => iou(boxes[0],box)<0.7);
-            }
-            return result;
-        }
-
-        function iou(box1,box2) {
-            return intersection(box1,box2)/union(box1,box2);
-        }
-
-        function union(box1,box2) {
-            const [box1_x1,box1_y1,box1_x2,box1_y2] = box1;
-            const [box2_x1,box2_y1,box2_x2,box2_y2] = box2;
-            const box1_area = (box1_x2-box1_x1)*(box1_y2-box1_y1)
-            const box2_area = (box2_x2-box2_x1)*(box2_y2-box2_y1)
-            return box1_area + box2_area - intersection(box1,box2)
-        }
-
-        function intersection(box1,box2) {
-            const [box1_x1,box1_y1,box1_x2,box1_y2] = box1;
-            const [box2_x1,box2_y1,box2_x2,box2_y2] = box2;
-            const x1 = Math.max(box1_x1,box2_x1);
-            const y1 = Math.max(box1_y1,box2_y1);
-            const x2 = Math.min(box1_x2,box2_x2);
-            const y2 = Math.min(box1_y2,box2_y2);
-            return (x2-x1)*(y2-y1)
-        }

        const yolo_classes = [
            'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat',
@@ -228,7 +228,7 @@ class Darknet:
          module.append(BatchNorm2d(filters, eps=1e-05, track_running_stats=True))
        # LeakyReLU activation
        if activation == "leaky":
-          module.append(lambda x: x.leakyrelu(0.1))
+          module.append(lambda x: x.leaky_relu(0.1))
      elif module_type == "maxpool":
        size, stride = int(x["size"]), int(x["stride"])
        module.append(lambda x: x.max_pool2d(kernel_size=(size, size), stride=stride))
@@ -3,7 +3,7 @@ import os
 from ultralytics import YOLO
 import onnx
 from pathlib import Path
-from extra.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner
 from extra.onnx_helpers import get_example_inputs
 from tinygrad.tensor import Tensor

@@ -42,69 +42,8 @@ def preprocess(im, imgsz=640, model_stride=32, model_pt=True):
  im = im / 255.0  # 0 - 255 to 0.0 - 1.0
  return im

-# Post Processing functions
-def box_area(box):
-  return (box[:, 2] - box[:, 0]) * (box[:, 3] - box[:, 1])

-def box_iou(box1, box2):
-  lt = np.maximum(box1[:, None, :2], box2[:, :2])
-  rb = np.minimum(box1[:, None, 2:], box2[:, 2:])
-  wh = np.clip(rb - lt, 0, None)
-  inter = wh[:, :, 0] * wh[:, :, 1]
-  area1 = box_area(box1)[:, None]
-  area2 = box_area(box2)[None, :]
-  iou = inter / (area1 + area2 - inter)
-  return iou
-
-def compute_nms(boxes, scores, iou_threshold):
-  order, keep = scores.argsort()[::-1], []
-  while order.size > 0:
-    i = order[0]
-    keep.append(i)
-    if order.size == 1:
-      break
-    iou = box_iou(boxes[i][None, :], boxes[order[1:]])
-    inds = np.where(np.atleast_1d(iou.squeeze()) <= iou_threshold)[0]
-    order = order[inds + 1]
-  return np.array(keep)
-
-def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, agnostic=False, max_det=300, nc=0, max_wh=7680):
-  prediction = prediction[0] if isinstance(prediction, (list, tuple)) else prediction
-  bs, nc = prediction.shape[0], nc or (prediction.shape[1] - 4)
-  xc = np.amax(prediction[:, 4:4 + nc], axis=1) > conf_thres
-  nm = prediction.shape[1] - nc - 4
-  output = [np.zeros((0, 6 + nm))] * bs
-
-  for xi, x in enumerate(prediction):
-    x = x.swapaxes(0, -1)[xc[xi]]
-    if not x.shape[0]: continue
-    box, cls, mask = np.split(x, [4, 4 + nc], axis=1)
-    conf, j = np.max(cls, axis=1, keepdims=True), np.argmax(cls, axis=1, keepdims=True)
-    x = np.concatenate((xywh2xyxy(box), conf, j.astype(np.float32), mask), axis=1)
-    x = x[conf.ravel() > conf_thres]
-    if not x.shape[0]: continue
-    x = x[np.argsort(-x[:, 4])]
-    c = x[:, 5:6] * (0 if agnostic else max_wh)
-    boxes, scores = x[:, :4] + c, x[:, 4]
-    i = compute_nms(boxes, scores, iou_thres)[:max_det]
-    output[xi] = x[i]
-  return output
-
-def postprocess(preds, img, orig_imgs):
-  print('copying to CPU now for post processing')
-  #if you are on CPU, this causes an overflow runtime error. doesn't "seem" to make any difference in the predictions though.
-  # TODO: make non_max_suppression in tinygrad - to make this faster
-  preds = preds.numpy() if isinstance(preds, Tensor) else preds
-  preds = non_max_suppression(prediction=preds, conf_thres=0.25, iou_thres=0.7, agnostic=False, max_det=300)
-  all_preds = []
-  for i, pred in enumerate(preds):
-    orig_img = orig_imgs[i] if isinstance(orig_imgs, list) else orig_imgs
-    if not isinstance(orig_imgs, Tensor):
-      pred[:, :4] = scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
-      all_preds.append(pred)
-  return all_preds
-
-def draw_bounding_boxes_and_save(orig_img_paths, output_img_paths, all_predictions, class_labels, iou_threshold=0.5):
+def draw_bounding_boxes_and_save(orig_img_path, output_img_path, predictions, class_labels):
  color_dict = {label: tuple((((i+1) * 50) % 256, ((i+1) * 100) % 256, ((i+1) * 150) % 256)) for i, label in enumerate(class_labels)}
  font = cv2.FONT_HERSHEY_SIMPLEX

@@ -113,52 +52,32 @@ def draw_bounding_boxes_and_save(orig_img_paths, output_img_paths, all_predictio
    brightness = (r * 299 + g * 587 + b * 114) / 1000
    return brightness > 127

-  for img_idx, (orig_img_path, output_img_path, predictions) in enumerate(zip(orig_img_paths, output_img_paths, all_predictions)):
-    predictions = np.array(predictions)
-    orig_img = cv2.imread(orig_img_path) if not isinstance(orig_img_path, np.ndarray) else cv2.imdecode(orig_img_path, 1)
-    height, width, _ = orig_img.shape
-    box_thickness = int((height + width) / 400)
-    font_scale = (height + width) / 2500
+  orig_img = cv2.imread(orig_img_path) if not isinstance(orig_img_path, np.ndarray) else cv2.imdecode(orig_img_path, 1)
+  height, width, _ = orig_img.shape
+  box_thickness = int((height + width) / 400)
+  font_scale = (height + width) / 2500
+  object_count = defaultdict(int)

-    grouped_preds = defaultdict(list)
-    object_count = defaultdict(int)
+  for pred in predictions:
+    x1, y1, x2, y2, conf, class_id = pred
+    if conf == 0: continue
+    x1, y1, x2, y2, class_id = map(int, (x1, y1, x2, y2, class_id))
+    color = color_dict[class_labels[class_id]]
+    cv2.rectangle(orig_img, (x1, y1), (x2, y2), color, box_thickness)
+    label = f"{class_labels[class_id]} {conf:.2f}"
+    text_size, _ = cv2.getTextSize(label, font, font_scale, 1)
+    label_y, bg_y = (y1 - 4, y1 - text_size[1] - 4) if y1 - text_size[1] - 4 > 0 else (y1 + text_size[1], y1)
+    cv2.rectangle(orig_img, (x1, bg_y), (x1 + text_size[0], bg_y + text_size[1]), color, -1)
+    font_color = (0, 0, 0) if is_bright_color(color) else (255, 255, 255)
+    cv2.putText(orig_img, label, (x1, label_y), font, font_scale, font_color, 1, cv2.LINE_AA)
+    object_count[class_labels[class_id]] += 1

-    for pred_np in predictions:
-      grouped_preds[int(pred_np[-1])].append(pred_np)
+  print("Objects detected:")
+  for obj, count in object_count.items():
+    print(f"- {obj}: {count}")

-    def draw_box_and_label(pred, color):
-      x1, y1, x2, y2, conf, _ = pred
-      x1, y1, x2, y2 = map(int, (x1, y1, x2, y2))
-      cv2.rectangle(orig_img, (x1, y1), (x2, y2), color, box_thickness)
-      label = f"{class_labels[class_id]} {conf:.2f}"
-      text_size, _ = cv2.getTextSize(label, font, font_scale, 1)
-      label_y, bg_y = (y1 - 4, y1 - text_size[1] - 4) if y1 - text_size[1] - 4 > 0 else (y1 + text_size[1], y1)
-      cv2.rectangle(orig_img, (x1, bg_y), (x1 + text_size[0], bg_y + text_size[1]), color, -1)
-      font_color = (0, 0, 0) if is_bright_color(color) else (255, 255, 255)
-      cv2.putText(orig_img, label, (x1, label_y), font, font_scale, font_color, 1, cv2.LINE_AA)
-
-    for class_id, pred_list in grouped_preds.items():
-      pred_list = np.array(pred_list)
-      while len(pred_list) > 0:
-        max_conf_idx = np.argmax(pred_list[:, 4])
-        max_conf_pred = pred_list[max_conf_idx]
-        pred_list = np.delete(pred_list, max_conf_idx, axis=0)
-        color = color_dict[class_labels[class_id]]
-        draw_box_and_label(max_conf_pred, color)
-        object_count[class_labels[class_id]] += 1
-        iou_scores = box_iou(np.array([max_conf_pred[:4]]), pred_list[:, :4])
-        low_iou_indices = np.where(iou_scores[0] < iou_threshold)[0]
-        pred_list = pred_list[low_iou_indices]
-        for low_conf_pred in pred_list:
-          draw_box_and_label(low_conf_pred, color)
-
-    print(f"Image {img_idx + 1}:")
-    print("Objects detected:")
-    for obj, count in object_count.items():
-      print(f"- {obj}: {count}")
-
-    cv2.imwrite(output_img_path, orig_img)
-    print(f'saved detections at {output_img_path}')
+  cv2.imwrite(output_img_path, orig_img)
+  print(f'saved detections at {output_img_path}')

 # utility functions for forward pass.
 def dist2bbox(distance, anchor_points, xywh=True, dim=-1):
@@ -202,34 +121,26 @@ def clip_boxes(boxes, shape):
  boxes[..., [1, 3]] = np.clip(boxes[..., [1, 3]], 0, shape[0])  # y1, y2
  return boxes

-def scale_boxes(img1_shape, boxes, img0_shape, ratio_pad=None):
+def scale_boxes(img1_shape, predictions, img0_shape, ratio_pad=None):
  gain = ratio_pad if ratio_pad else min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])
  pad = ((img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2)
-  boxes_np = boxes.numpy() if isinstance(boxes, Tensor) else boxes
-  boxes_np[..., [0, 2]] -= pad[0]
-  boxes_np[..., [1, 3]] -= pad[1]
-  boxes_np[..., :4] /= gain
-  boxes_np = clip_boxes(boxes_np, img0_shape)
-  return boxes_np
-
-def xywh2xyxy(x):
-  xy = x[..., :2]  # center x, y
-  wh = x[..., 2:4]  # width, height
-  xy1 = xy - wh / 2  # top left x, y
-  xy2 = xy + wh / 2  # bottom right x, y
-  result = np.concatenate((xy1, xy2), axis=-1)
-  return Tensor(result) if isinstance(x, Tensor) else result
+  for pred in predictions:
+    boxes_np = pred[:4].numpy() if isinstance(pred[:4], Tensor) else pred[:4]
+    boxes_np[..., [0, 2]] -= pad[0]
+    boxes_np[..., [1, 3]] -= pad[1]
+    boxes_np[..., :4] /= gain
+    boxes_np = clip_boxes(boxes_np, img0_shape)
+    pred[:4] = boxes_np
+  return predictions

 def get_variant_multiples(variant):
  return {'n':(0.33, 0.25, 2.0), 's':(0.33, 0.50, 2.0), 'm':(0.67, 0.75, 1.5), 'l':(1.0, 1.0, 1.0), 'x':(1, 1.25, 1.0) }.get(variant, None)

 def label_predictions(all_predictions):
  class_index_count = defaultdict(int)
-  for predictions in all_predictions:
-    predictions = np.array(predictions)
-    for pred_np in predictions:
-      class_id = int(pred_np[-1])
-      class_index_count[class_id] += 1
+  for pred in all_predictions:
+    class_id = int(pred[-1])
+    if pred[-2] != 0: class_index_count[class_id] += 1

  return dict(class_index_count)

@@ -380,7 +291,9 @@ class YOLOv8:
  def __call__(self, x):
    x = self.net(x)
    x = self.fpn(*x)
-    return self.head(x)
+    x = self.head(x)
+    # TODO: postprocess needs to be in the model to be compiled to webgpu
+    return postprocess(x)

  def return_all_trainable_modules(self):
    backbone_modules = [*range(10)]
@@ -403,16 +316,45 @@ def convert_f16_safetensor_to_f32(input_file: Path, output_file: Path):
    f.write(new_metadata_bytes)
    float32_values.tofile(f)

+def compute_iou_matrix(boxes):
+  x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
+  areas = (x2 - x1) * (y2 - y1)
+  x1 = Tensor.maximum(x1[:, None], x1[None, :])
+  y1 = Tensor.maximum(y1[:, None], y1[None, :])
+  x2 = Tensor.minimum(x2[:, None], x2[None, :])
+  y2 = Tensor.minimum(y2[:, None], y2[None, :])
+  w = Tensor.maximum(Tensor(0), x2 - x1)
+  h = Tensor.maximum(Tensor(0), y2 - y1)
+  intersection = w * h
+  union = areas[:, None] + areas[None, :] - intersection
+  return intersection / union
+
+def postprocess(output, max_det=300, conf_threshold=0.25, iou_threshold=0.45):
+  xc, yc, w, h, class_scores = output[0][0], output[0][1], output[0][2], output[0][3], output[0][4:]
+  class_ids = Tensor.argmax(class_scores, axis=0)
+  probs = Tensor.max(class_scores, axis=0)
+  probs = Tensor.where(probs >= conf_threshold, probs, 0)
+  x1 = xc - w / 2
+  y1 = yc - h / 2
+  x2 = xc + w / 2
+  y2 = yc + h / 2
+  boxes = Tensor.stack(x1, y1, x2, y2, probs, class_ids, dim=1)
+  order = Tensor.topk(probs, max_det)[1]
+  boxes = boxes[order]
+  iou = compute_iou_matrix(boxes[:, :4])
+  iou = Tensor.triu(iou, diagonal=1)
+  same_class_mask = boxes[:, -1][:, None] == boxes[:, -1][None, :]
+  high_iou_mask = (iou > iou_threshold) & same_class_mask
+  no_overlap_mask = high_iou_mask.sum(axis=0) == 0
+  boxes = boxes * no_overlap_mask.unsqueeze(-1)
+  return boxes
+
 def get_weights_location(yolo_variant: str) -> Path:
  weights_location = Path(__file__).parents[1] / "weights" / f'yolov8{yolo_variant}.safetensors'
  fetch(f'https://gitlab.com/r3sist/yolov8_weights/-/raw/master/yolov8{yolo_variant}.safetensors', weights_location)
-
-  if not is_dtype_supported(dtypes.half):
-    f32_weights = weights_location.with_name(f"{weights_location.stem}_f32.safetensors")
-    if not f32_weights.exists(): convert_f16_safetensor_to_f32(weights_location, f32_weights)
-    weights_location = f32_weights
-
-  return weights_location
+  f32_weights = weights_location.with_name(f"{weights_location.stem}_f32.safetensors")
+  if not f32_weights.exists(): convert_f16_safetensor_to_f32(weights_location, f32_weights)
+  return f32_weights

 if __name__ == '__main__':

@@ -428,14 +370,13 @@ if __name__ == '__main__':
  output_folder_path = Path('./outputs_yolov8')
  output_folder_path.mkdir(parents=True, exist_ok=True)
  #absolute image path or URL
-  image_location = [np.frombuffer(fetch(img_path).read_bytes(), np.uint8)]
-  image = [cv2.imdecode(image_location[0], 1)]
-  out_paths = [(output_folder_path / f"{Path(img_path).stem}_output{Path(img_path).suffix or '.png'}").as_posix()]
+  image_location = np.frombuffer(fetch(img_path).read_bytes(), np.uint8)
+  image = [cv2.imdecode(image_location, 1)]
+  out_path = (output_folder_path / f"{Path(img_path).stem}_output{Path(img_path).suffix or '.png'}").as_posix()
  if not isinstance(image[0], np.ndarray):
    print('Error in image loading. Check your image file.')
    sys.exit(1)
  pre_processed_image = preprocess(image)
-
  # Different YOLOv8 variants use different w , r, and d multiples. For a list , refer to this yaml file (the scales section) https://github.com/ultralytics/ultralytics/blob/main/ultralytics/cfg/models/v8/yolov8.yaml
  depth, width, ratio = get_variant_multiples(yolo_variant)
  yolo_infer = YOLOv8(w=width, r=ratio, d=depth, num_classes=80)
@@ -443,15 +384,13 @@ if __name__ == '__main__':
  load_state_dict(yolo_infer, state_dict)

  st = time.time()
-  predictions = yolo_infer(pre_processed_image)
+  predictions = yolo_infer(pre_processed_image).numpy()
+
  print(f'did inference in {int(round(((time.time() - st) * 1000)))}ms')
-
-  post_predictions = postprocess(preds=predictions, img=pre_processed_image, orig_imgs=image)
-
  #v8 and v3 have same 80 class names for Object Detection
  class_labels = fetch('https://raw.githubusercontent.com/pjreddie/darknet/master/data/coco.names').read_text().split("\n")
-
-  draw_bounding_boxes_and_save(orig_img_paths=image_location, output_img_paths=out_paths, all_predictions=post_predictions, class_labels=class_labels)
+  predictions = scale_boxes(pre_processed_image.shape[2:], predictions, image[0].shape)
+  draw_bounding_boxes_and_save(orig_img_path=image_location, output_img_path=out_path, predictions=predictions, class_labels=class_labels)

 # TODO for later:
 #  1. Fix SPPF minor difference due to maxpool
@@ -1,16 +1,25 @@
 #!/usr/bin/env python3

 import time, mmap, sys, shutil, os, glob, subprocess
-from tinygrad.helpers import to_mv, DEBUG, colored, ansilen
+from tinygrad.helpers import DEBUG, colored, ansilen
 from tinygrad.runtime.autogen import libc
-from tinygrad.runtime.autogen.am import smu_v13_0_0
+from tinygrad.runtime.autogen.am import am
+from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
-from tinygrad.runtime.support.am.ip import AM_SOC21, AM_GMC, AM_IH, AM_PSP, AM_SMU, AM_GFX, AM_SDMA
+from tinygrad.runtime.support.am.ip import AM_SOC, AM_GMC, AM_IH, AM_PSP, AM_SMU, AM_GFX, AM_SDMA

-AM_VERSION = 0xA0000002
+AM_VERSION = 0xA0000004

 def bold(s): return f"\033[1m{s}\033[0m"

+def trim(s:str, length:int) -> str:
+  if len(s) > length: return s[:length-3] + "..."
+  return s
+
+def pad(x:str, length:int) -> str:
+  if len(x) < length: return x + " " * (length - len(x))
+  return x
+
 def color_temp(temp):
  if temp >= 87: return colored(f"{temp:>3}", "red")
  elif temp >= 80: return colored(f"{temp:>3}", "yellow")
@@ -23,14 +32,17 @@ def draw_bar(percentage, width=40, fill='█', empty='░'):
  bar = fill * filled_width + empty * (width - filled_width)
  return f'[{bar}] {percentage*100:5.1f}%'

-def same_line(strs:list[list[str]], split=8) -> list[str]:
+def same_line(strs:list[list[str]|None], split=8) -> list[str]:
+  strs = [s for s in strs if s is not None]
+  if len(strs) == 0: return []
+
  ret = []
  max_width_in_block = [max(ansilen(line) for line in block) for block in strs]
  max_height = max(len(block) for block in strs)
  for i in range(max_height):
    line = []
    for bid, block in enumerate(strs):
-      if i < len(block): line.append(block[i] + ' ' * (split + max_width_in_block[bid] - ansilen(block[i])))
+      if i < len(block): line.append(block[i] + (' ' * (split + max_width_in_block[bid] - ansilen(block[i])) if bid != len(strs) - 1 else ''))
      else: line.append(' ' * (split + max_width_in_block[bid]))
    ret.append(' '.join(line))
  return ret
@@ -47,9 +59,9 @@ def get_bar0_size(pcibus):
  return int(end_hex, 16) - int(start_hex, 16) + 1

 class AMSMI(AMDev):
-  def __init__(self, pcibus, vram_bar:memoryview, doorbell_bar:memoryview, mmio_bar:memoryview):
+  def __init__(self, pcibus, vram_bar:MMIOInterface, doorbell_bar:MMIOInterface, mmio_bar:MMIOInterface):
    self.pcibus = pcibus
-    self.vram, self.doorbell64, self.mmio = vram_bar, doorbell_bar, mmio_bar
+    self.vram, self.doorbell64, self.mmio, self.dma_regions = vram_bar, doorbell_bar, mmio_bar, None
    self.pci_state = self.read_pci_state()
    if self.pci_state == "D0": self._init_from_d0()

@@ -65,12 +77,14 @@ class AMSMI(AMDev):
    self.mm = AMMemoryManager(self, self.vram_size)

    # Initialize IP blocks
-    self.soc21:AM_SOC21 = AM_SOC21(self)
+    self.soc:AM_SOC = AM_SOC(self)
    self.gmc:AM_GMC = AM_GMC(self)
    self.ih:AM_IH = AM_IH(self)
    self.psp:AM_PSP = AM_PSP(self)
    self.smu:AM_SMU = AM_SMU(self)

+    for ip in [self.soc, self.gmc, self.ih, self.psp, self.smu]: ip.init_sw()
+
  def read_pci_state(self):
    with open(f"/sys/bus/pci/devices/{self.pcibus}/power_state", "r") as f: return f.read().strip().rstrip()

@@ -80,6 +94,7 @@ class SMICtx:
    self.opened_pcidevs = []
    self.opened_pci_resources = {}
    self.prev_lines_cnt = 0
+    self.prev_terminal_width = 0

    remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:"]
    lspci = subprocess.check_output(["lspci"]).decode("utf-8").splitlines()
@@ -92,9 +107,9 @@ class SMICtx:
      bar_fds = {bar: os.open(f"/sys/bus/pci/devices/{pcibus}/resource{bar}", os.O_RDWR | os.O_SYNC) for bar in [0, 2, 5]}
      bar_size = {0: get_bar0_size(pcibus), 2: os.fstat(bar_fds[2]).st_size, 5: os.fstat(bar_fds[5]).st_size}

-      def map_pci_range(bar):
-        return to_mv(libc.mmap(0, bar_size[bar], mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED, bar_fds[bar], 0), bar_size[bar])
-      self.opened_pci_resources[pcibus] = (map_pci_range(0), None, map_pci_range(5).cast('I'))
+      def map_pci_range(bar, fmt='B'):
+        return MMIOInterface(libc.mmap(0, bar_size[bar], mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED, bar_fds[bar], 0), bar_size[bar], fmt)
+      self.opened_pci_resources[pcibus] = (map_pci_range(0), None, map_pci_range(5, 'I'))

    try:
      self.devs.append(AMSMI(pcibus, *self.opened_pci_resources[pcibus]))
@@ -125,52 +140,107 @@ class SMICtx:

  def collect(self): return {d: d.smu.read_metrics() if d.pci_state == "D0" else None for d in self.devs}

+  def get_gfx_activity(self, dev, metrics): return metrics.SmuMetrics.AverageGfxActivity
+  def get_mem_activity(self, dev, metrics): return metrics.SmuMetrics.AverageUclkActivity
+
+  def get_temps(self, dev, metrics, compact=False):
+    temps_keys = [(k, name) for k, name in dev.smu.smu_mod.c__EA_TEMP_e__enumvalues.items()
+                  if k < dev.smu.smu_mod.TEMP_COUNT and metrics.SmuMetrics.AvgTemperature[k] != 0]
+    if compact: temps_keys = [(k, name) for k, name in temps_keys if k in (dev.smu.smu_mod.TEMP_HOTSPOT, dev.smu.smu_mod.TEMP_MEM)]
+    return {name: metrics.SmuMetrics.AvgTemperature[k] for k, name in temps_keys}
+
+  def get_voltage(self, dev, metrics, compact=False):
+    voltage_keys = [(k, name) for k, name in dev.smu.smu_mod.c__EA_SVI_PLANE_e__enumvalues.items()
+                        if k < dev.smu.smu_mod.SVI_PLANE_COUNT and metrics.SmuMetrics.AvgVoltage[k] != 0]
+    return {name: metrics.SmuMetrics.AvgVoltage[k] for k, name in voltage_keys}
+
+  def get_busy_threshold(self, dev):
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (14, 0, 2): return 5
+      case _: return 15
+
+  def get_gfx_freq(self, dev, metrics):
+    return metrics.SmuMetrics.AverageGfxclkFrequencyPostDs if self.get_gfx_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
+          metrics.SmuMetrics.AverageGfxclkFrequencyPreDs
+
+  def get_mem_freq(self, dev, metrics):
+    return metrics.SmuMetrics.AverageMemclkFrequencyPostDs if self.get_mem_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
+           metrics.SmuMetrics.AverageMemclkFrequencyPreDs
+
+  def get_fckl_freq(self, dev, metrics):
+    return metrics.SmuMetrics.AverageFclkFrequencyPostDs if self.get_mem_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
+           metrics.SmuMetrics.AverageFclkFrequencyPreDs
+
+  def get_fan_rpm_pwm(self, dev, metrics): return metrics.SmuMetrics.AvgFanRpm, metrics.SmuMetrics.AvgFanPwm
+
+  def get_power(self, dev, metrics): return metrics.SmuMetrics.AverageSocketPower, metrics.SmuMetrics.dGPU_W_MAX
+
  def draw(self):
-    terminal_width, _ = shutil.get_terminal_size()
+    terminal_width, terminal_height = shutil.get_terminal_size()
+    if self.prev_terminal_width != terminal_width or self.prev_terminal_height != terminal_height:
+      os.system('clear')
+    self.prev_terminal_width, self.prev_terminal_height = terminal_width, terminal_height
+
+    padding = 8
+    col_size = (terminal_width) // 2 - padding - 2
+    activity_line_width = 50 if terminal_width > 170 else \
+                         (30 if terminal_width > 130 else \
+                         (16 if terminal_width > 92 else \
+                         max(0, terminal_width - 77)))

    dev_metrics = self.collect()
    dev_content = []
    for dev, metrics in dev_metrics.items():
      if dev.pci_state != "D0":
-        dev_content.append([f"{colored('(sleep)', 'yellow')} {bold(dev.pcibus)}: {self.lspci[dev.pcibus[5:]]}"] +
-                           [f"PCI State: {dev.pci_state}"] + [" "*107])
+        dev_content.append([f"{colored('(sleep)', 'yellow')} {bold(dev.pcibus)}: {trim(self.lspci[dev.pcibus[5:]], col_size - 20)}"] +
+                           [pad(f"PCI State: {dev.pci_state}", col_size)])
        continue

-      device_line = [f"{bold(dev.pcibus)}: {self.lspci[dev.pcibus[5:]]}"] + [""]
-      activity_line = [f"GFX Activity {draw_bar(metrics.SmuMetrics.AverageGfxActivity / 100, 50)}"] \
-                    + [f"MEM Activity {draw_bar(metrics.SmuMetrics.AverageUclkActivity / 100, 50)}"] + [""]
+      device_line = [f"{bold(dev.pcibus)} {trim(self.lspci[dev.pcibus[5:]], col_size - 20)}"] + [pad("", col_size)]
+      activity_line = [f"GFX Activity {draw_bar(self.get_gfx_activity(dev, metrics) / 100, activity_line_width)}"] \
+                    + [f"MEM Activity {draw_bar(self.get_mem_activity(dev, metrics) / 100, activity_line_width)}"]

-      # draw_metrics_table(metrics, dev)
-      temps_keys = [(k, name) for k, name in smu_v13_0_0.c__EA_TEMP_e__enumvalues.items()
-                                if k < smu_v13_0_0.TEMP_COUNT and metrics.SmuMetrics.AvgTemperature[k] != 0]
-      temps_table = ["=== Temps (C) ==="] + [f"{name:<15}: {color_temp(metrics.SmuMetrics.AvgTemperature[k])}" for k, name in temps_keys]
+      temps_data, temps_data_compact = self.get_temps(dev, metrics), self.get_temps(dev, metrics, compact=True)
+      temps_table = ["=== Temps (°C) ==="] + [f"{name:<16}: {color_temp(val)}" for name, val in temps_data.items()]
+      temps_table_compact = ["Temps (°C):" + '/'.join([f"{color_temp(val)} {name}" for name, val in temps_data_compact.items()])]

-      voltage_keys = [(k, name) for k, name in smu_v13_0_0.c__EA_SVI_PLANE_e__enumvalues.items() if k < smu_v13_0_0.SVI_PLANE_COUNT]
-      power_table = ["=== Power ==="] \
-                  + [f"Fan Speed: {metrics.SmuMetrics.AvgFanRpm} RPM"] \
-                  + [f"Fan Power: {metrics.SmuMetrics.AvgFanPwm}%"] \
-                  + [f"Power: {metrics.SmuMetrics.AverageSocketPower:>3}W " +
-                       draw_bar(metrics.SmuMetrics.AverageSocketPower / metrics.SmuMetrics.dGPU_W_MAX, 16)] \
-                  + ["", "=== Voltages ==="] + [f"{name:<20}: {color_voltage(metrics.SmuMetrics.AvgVoltage[k])}" for k, name in voltage_keys]
+      fan_rpm, fan_pwm = self.get_fan_rpm_pwm(dev, metrics)
+      power_table = ["=== Power ==="] + [f"Fan Speed: {fan_rpm} RPM"] + [f"Fan Power: {fan_pwm}%"]

-      frequency_table = ["=== Frequencies ===",
-        f"GFXCLK Target : {metrics.SmuMetrics.AverageGfxclkFrequencyTarget:>4} MHz",
-        f"GFXCLK PreDs  : {metrics.SmuMetrics.AverageGfxclkFrequencyPreDs:>4} MHz",
-        f"GFXCLK PostDs : {metrics.SmuMetrics.AverageGfxclkFrequencyPostDs:>4} MHz",
-        f"FCLK PreDs    : {metrics.SmuMetrics.AverageFclkFrequencyPreDs:>4} MHz",
-        f"FCLK PostDs   : {metrics.SmuMetrics.AverageFclkFrequencyPostDs:>4} MHz",
-        f"MCLK PreDs    : {metrics.SmuMetrics.AverageMemclkFrequencyPreDs:>4} MHz",
-        f"MCLK PostDs   : {metrics.SmuMetrics.AverageMemclkFrequencyPostDs:>4} MHz",
-        f"VCLK0         : {metrics.SmuMetrics.AverageVclk0Frequency:>4} MHz",
-        f"DCLK0         : {metrics.SmuMetrics.AverageDclk0Frequency:>4} MHz",
-        f"VCLK1         : {metrics.SmuMetrics.AverageVclk1Frequency:>4} MHz",
-        f"DCLK1         : {metrics.SmuMetrics.AverageDclk1Frequency:>4} MHz"]
+      total_power, max_power = self.get_power(dev, metrics)
+      power_line = [f"Power: {total_power:>3}W " + draw_bar(total_power / max_power, 16)]
+      power_line_compact = [f"Power:  {total_power:>3}W " + draw_bar(total_power / max_power, activity_line_width)]
+
+      voltage_data = self.get_voltage(dev, metrics)
+      voltage_table = ["=== Voltages ==="] + [f"{name:<20}: {color_voltage(voltage)}" for name, voltage in voltage_data.items()]
+
+      gfx_freq = self.get_gfx_freq(dev, metrics)
+      mclk_freq = self.get_mem_freq(dev, metrics)
+      fclk_freq = self.get_fckl_freq(dev, metrics)
+
+      frequency_table = ["=== Frequencies ===", f"GFXCLK: {gfx_freq:>4} MHz", f"FCLK  : {fclk_freq:>4} MHz", f"MCLK  : {mclk_freq:>4} MHz"]
+
+      if self.prev_terminal_width >= 231:
+        power_table += power_line + [""] + voltage_table
+        activity_line += [""]
+      elif self.prev_terminal_width >= 171:
+        power_table += power_line + [""] + frequency_table
+        activity_line += [""]
+        frequency_table = None
+      elif self.prev_terminal_width >= 121:
+        temps_table = None
+        activity_line += power_line_compact
+      else:
+        temps_table = None
+        power_table = None
+        frequency_table = None
+        activity_line += power_line_compact

      dev_content.append(device_line + activity_line + same_line([temps_table, power_table, frequency_table]))

    raw_text = 'AM Monitor'.center(terminal_width) + "\n" + "=" * terminal_width + "\n\n"
    for i in range(0, len(dev_content), 2):
-      if i + 1 < len(dev_content): raw_text += '\n'.join(same_line([dev_content[i], dev_content[i+1]]))
+      if i + 1 < len(dev_content): raw_text += '\n'.join(same_line([dev_content[i], dev_content[i+1]], split=padding))
      else: raw_text += '\n'.join(dev_content[i])
      if i + 2 < len(dev_content): raw_text += "\n" + "=" * terminal_width + "\n\n"

@@ -21,395 +21,426 @@
 *
 */

-#ifndef _DISCOVERY_H_
-#define _DISCOVERY_H_
-
-#define uint32_t unsigned int
-#define uint8_t unsigned char
-#define uint16_t unsigned short
-#define uint64_t unsigned long long
-#define u32 unsigned int
-#define u8 unsigned char
-#define u16 unsigned short
-#define u64 unsigned long long
-#define bool unsigned char
-
-#define PSP_HEADER_SIZE                 256
-#define BINARY_SIGNATURE                0x28211407
-#define DISCOVERY_TABLE_SIGNATURE       0x53445049
-#define GC_TABLE_ID                     0x4347
-#define HARVEST_TABLE_SIGNATURE         0x56524148
-#define VCN_INFO_TABLE_ID               0x004E4356
-#define MALL_INFO_TABLE_ID              0x4C4C414D
-#define NPS_INFO_TABLE_ID 0x0053504E
-
-typedef enum {
-	IP_DISCOVERY = 0,
-	GC,
-	HARVEST_INFO,
-	VCN_INFO,
-	MALL_INFO,
-	NPS_INFO,
-	TOTAL_TABLES = 6
-} table;
-
-#pragma pack(1)
-
-typedef struct table_info
-{
-	uint16_t offset;   /* Byte offset */
-	uint16_t checksum; /* Byte sum of the table */
-	uint16_t size;     /* Table size */
-	uint16_t padding;
-} table_info;
-
-typedef struct binary_header
-{
-	/* psp structure should go at the top of this structure */
-	uint32_t binary_signature; /* 0x7, 0x14, 0x21, 0x28 */
-	uint16_t version_major;
-	uint16_t version_minor;
-	uint16_t binary_checksum;  /* Byte sum of the binary after this field */
-	uint16_t binary_size;      /* Binary Size*/
-	table_info table_list[TOTAL_TABLES];
-} binary_header;
-
-typedef struct die_info
-{
-	uint16_t die_id;
-	uint16_t die_offset; /* Points to the corresponding die_header structure */
-} die_info;
-
-
-typedef struct ip_discovery_header
-{
-	uint32_t signature;    /* Table Signature */
-	uint16_t version;      /* Table Version */
-	uint16_t size;         /* Table Size */
-	uint32_t id;           /* Table ID */
-	uint16_t num_dies;     /* Number of Dies */
-	die_info die_info[16]; /* list die information for up to 16 dies */
-	union {
-		uint16_t padding[1];	/* version <= 3 */
-		struct {		/* version == 4 */
-			uint8_t base_addr_64_bit : 1; /* ip structures are using 64 bit base address */
-			uint8_t reserved : 7;
-			uint8_t reserved2;
-		};
-	};
-} ip_discovery_header;
-
-typedef struct ip
-{
-	uint16_t hw_id;           /* Hardware ID */
-	uint8_t number_instance;  /* instance of the IP */
-	uint8_t num_base_address; /* Number of Base Addresses */
-	uint8_t major;            /* HCID Major */
-	uint8_t minor;            /* HCID Minor */
-	uint8_t revision;         /* HCID Revision */
-#if defined(__BIG_ENDIAN)
-	uint8_t reserved : 4;     /* Placeholder field */
-	uint8_t harvest : 4;      /* Harvest */
-#else
-	uint8_t harvest : 4;      /* Harvest */
-	uint8_t reserved : 4;     /* Placeholder field */
-#endif
-	uint32_t base_address[1]; /* variable number of Addresses */
-} ip;
-
-typedef struct ip_v3
-{
-	uint16_t hw_id;                         /* Hardware ID */
-	uint8_t instance_number;                /* Instance number for the IP */
-	uint8_t num_base_address;               /* Number of base addresses*/
-	uint8_t major;                          /* Hardware ID.major version */
-	uint8_t minor;                          /* Hardware ID.minor version */
-	uint8_t revision;                       /* Hardware ID.revision version */
-#if defined(__BIG_ENDIAN)
-	uint8_t variant : 4;                    /* HW variant */
-	uint8_t sub_revision : 4;               /* HCID Sub-Revision */
-#else
-	uint8_t sub_revision : 4;               /* HCID Sub-Revision */
-	uint8_t variant : 4;                    /* HW variant */
-#endif
-	uint32_t base_address[1];		/* Base Address list. Corresponds to the num_base_address field*/
-} ip_v3;
-
-typedef struct ip_v4 {
-	uint16_t hw_id;                         /* Hardware ID */
-	uint8_t instance_number;                /* Instance number for the IP */
-	uint8_t num_base_address;               /* Number of base addresses*/
-	uint8_t major;                          /* Hardware ID.major version */
-	uint8_t minor;                          /* Hardware ID.minor version */
-	uint8_t revision;                       /* Hardware ID.revision version */
-#if defined(LITTLEENDIAN_CPU)
-	uint8_t sub_revision : 4;               /* HCID Sub-Revision */
-	uint8_t variant : 4;                    /* HW variant */
-#elif defined(BIGENDIAN_CPU)
-	uint8_t variant : 4;                    /* HW variant */
-	uint8_t sub_revision : 4;               /* HCID Sub-Revision */
-#endif
-	uint64_t base_address_64[1];
-} ip_v4;
-
-typedef struct die_header
-{
-	uint16_t die_id;
-	uint16_t num_ips;
-} die_header;
-
-typedef struct ip_structure
-{
-	ip_discovery_header* header;
-	struct die
-	{
-		die_header *die_header;
-		union
-		{
-			ip *ip_list;
-			ip_v3 *ip_v3_list;
-			ip_v4 *ip_v4_list;
-		};                                  /* IP list. Variable size*/
-	} die;
-} ip_structure;
-
-struct gpu_info_header {
-	uint32_t table_id;      /* table ID */
-	uint16_t version_major; /* table version */
-	uint16_t version_minor; /* table version */
-	uint32_t size;          /* size of the entire header+data in bytes */
-};
-
-struct gc_info_v1_0 {
-	struct gpu_info_header header;
-
-	uint32_t gc_num_se;
-	uint32_t gc_num_wgp0_per_sa;
-	uint32_t gc_num_wgp1_per_sa;
-	uint32_t gc_num_rb_per_se;
-	uint32_t gc_num_gl2c;
-	uint32_t gc_num_gprs;
-	uint32_t gc_num_max_gs_thds;
-	uint32_t gc_gs_table_depth;
-	uint32_t gc_gsprim_buff_depth;
-	uint32_t gc_parameter_cache_depth;
-	uint32_t gc_double_offchip_lds_buffer;
-	uint32_t gc_wave_size;
-	uint32_t gc_max_waves_per_simd;
-	uint32_t gc_max_scratch_slots_per_cu;
-	uint32_t gc_lds_size;
-	uint32_t gc_num_sc_per_se;
-	uint32_t gc_num_sa_per_se;
-	uint32_t gc_num_packer_per_sc;
-	uint32_t gc_num_gl2a;
-};
-
-struct gc_info_v1_1 {
-	struct gpu_info_header header;
-
-	uint32_t gc_num_se;
-	uint32_t gc_num_wgp0_per_sa;
-	uint32_t gc_num_wgp1_per_sa;
-	uint32_t gc_num_rb_per_se;
-	uint32_t gc_num_gl2c;
-	uint32_t gc_num_gprs;
-	uint32_t gc_num_max_gs_thds;
-	uint32_t gc_gs_table_depth;
-	uint32_t gc_gsprim_buff_depth;
-	uint32_t gc_parameter_cache_depth;
-	uint32_t gc_double_offchip_lds_buffer;
-	uint32_t gc_wave_size;
-	uint32_t gc_max_waves_per_simd;
-	uint32_t gc_max_scratch_slots_per_cu;
-	uint32_t gc_lds_size;
-	uint32_t gc_num_sc_per_se;
-	uint32_t gc_num_sa_per_se;
-	uint32_t gc_num_packer_per_sc;
-	uint32_t gc_num_gl2a;
-	uint32_t gc_num_tcp_per_sa;
-	uint32_t gc_num_sdp_interface;
-	uint32_t gc_num_tcps;
-};
-
-struct gc_info_v1_2 {
-	struct gpu_info_header header;
-	uint32_t gc_num_se;
-	uint32_t gc_num_wgp0_per_sa;
-	uint32_t gc_num_wgp1_per_sa;
-	uint32_t gc_num_rb_per_se;
-	uint32_t gc_num_gl2c;
-	uint32_t gc_num_gprs;
-	uint32_t gc_num_max_gs_thds;
-	uint32_t gc_gs_table_depth;
-	uint32_t gc_gsprim_buff_depth;
-	uint32_t gc_parameter_cache_depth;
-	uint32_t gc_double_offchip_lds_buffer;
-	uint32_t gc_wave_size;
-	uint32_t gc_max_waves_per_simd;
-	uint32_t gc_max_scratch_slots_per_cu;
-	uint32_t gc_lds_size;
-	uint32_t gc_num_sc_per_se;
-	uint32_t gc_num_sa_per_se;
-	uint32_t gc_num_packer_per_sc;
-	uint32_t gc_num_gl2a;
-	uint32_t gc_num_tcp_per_sa;
-	uint32_t gc_num_sdp_interface;
-	uint32_t gc_num_tcps;
-	uint32_t gc_num_tcp_per_wpg;
-	uint32_t gc_tcp_l1_size;
-	uint32_t gc_num_sqc_per_wgp;
-	uint32_t gc_l1_instruction_cache_size_per_sqc;
-	uint32_t gc_l1_data_cache_size_per_sqc;
-	uint32_t gc_gl1c_per_sa;
-	uint32_t gc_gl1c_size_per_instance;
-	uint32_t gc_gl2c_per_gpu;
-};
-
-struct gc_info_v2_0 {
-	struct gpu_info_header header;
-
-	uint32_t gc_num_se;
-	uint32_t gc_num_cu_per_sh;
-	uint32_t gc_num_sh_per_se;
-	uint32_t gc_num_rb_per_se;
-	uint32_t gc_num_tccs;
-	uint32_t gc_num_gprs;
-	uint32_t gc_num_max_gs_thds;
-	uint32_t gc_gs_table_depth;
-	uint32_t gc_gsprim_buff_depth;
-	uint32_t gc_parameter_cache_depth;
-	uint32_t gc_double_offchip_lds_buffer;
-	uint32_t gc_wave_size;
-	uint32_t gc_max_waves_per_simd;
-	uint32_t gc_max_scratch_slots_per_cu;
-	uint32_t gc_lds_size;
-	uint32_t gc_num_sc_per_se;
-	uint32_t gc_num_packer_per_sc;
-};
-
-struct gc_info_v2_1 {
-	struct gpu_info_header header;
-
-	uint32_t gc_num_se;
-	uint32_t gc_num_cu_per_sh;
-	uint32_t gc_num_sh_per_se;
-	uint32_t gc_num_rb_per_se;
-	uint32_t gc_num_tccs;
-	uint32_t gc_num_gprs;
-	uint32_t gc_num_max_gs_thds;
-	uint32_t gc_gs_table_depth;
-	uint32_t gc_gsprim_buff_depth;
-	uint32_t gc_parameter_cache_depth;
-	uint32_t gc_double_offchip_lds_buffer;
-	uint32_t gc_wave_size;
-	uint32_t gc_max_waves_per_simd;
-	uint32_t gc_max_scratch_slots_per_cu;
-	uint32_t gc_lds_size;
-	uint32_t gc_num_sc_per_se;
-	uint32_t gc_num_packer_per_sc;
-	/* new for v2_1 */
-	uint32_t gc_num_tcp_per_sh;
-	uint32_t gc_tcp_size_per_cu;
-	uint32_t gc_num_sdp_interface;
-	uint32_t gc_num_cu_per_sqc;
-	uint32_t gc_instruction_cache_size_per_sqc;
-	uint32_t gc_scalar_data_cache_size_per_sqc;
-	uint32_t gc_tcc_size;
-};
-
-typedef struct harvest_info_header {
-	uint32_t signature; /* Table Signature */
-	uint32_t version;   /* Table Version */
-} harvest_info_header;
-
-typedef struct harvest_info {
-	uint16_t hw_id;          /* Hardware ID */
-	uint8_t number_instance; /* Instance of the IP */
-	uint8_t reserved;        /* Reserved for alignment */
-} harvest_info;
-
-typedef struct harvest_table {
-	harvest_info_header header;
-	harvest_info list[32];
-} harvest_table;
-
-struct mall_info_header {
-	uint32_t table_id; /* table ID */
-	uint16_t version_major; /* table version */
-	uint16_t version_minor; /* table version */
-	uint32_t size_bytes; /* size of the entire header+data in bytes */
-};
-
-struct mall_info_v1_0 {
-	struct mall_info_header header;
-	uint32_t mall_size_per_m;
-	uint32_t m_s_present;
-	uint32_t m_half_use;
-	uint32_t m_mall_config;
-	uint32_t reserved[5];
-};
-
-struct mall_info_v2_0 {
-	struct mall_info_header header;
-	uint32_t mall_size_per_umc;
-	uint32_t reserved[8];
-};
-
-#define VCN_INFO_TABLE_MAX_NUM_INSTANCES 4
-
-struct vcn_info_header {
-    uint32_t table_id; /* table ID */
-    uint16_t version_major; /* table version */
-    uint16_t version_minor; /* table version */
-    uint32_t size_bytes; /* size of the entire header+data in bytes */
-};
-
-struct vcn_instance_info_v1_0
-{
-	uint32_t instance_num; /* VCN IP instance number. 0 - VCN0; 1 - VCN1 etc*/
-	union _fuse_data {
-		struct {
-			uint32_t av1_disabled : 1;
-			uint32_t vp9_disabled : 1;
-			uint32_t hevc_disabled : 1;
-			uint32_t h264_disabled : 1;
-			uint32_t reserved : 28;
-		} bits;
-		uint32_t all_bits;
-	} fuse_data;
-	uint32_t reserved[2];
-};
-
-struct vcn_info_v1_0 {
-	struct vcn_info_header header;
-	uint32_t num_of_instances; /* number of entries used in instance_info below*/
-	struct vcn_instance_info_v1_0 instance_info[VCN_INFO_TABLE_MAX_NUM_INSTANCES];
-	uint32_t reserved[4];
-};
-
-#define NPS_INFO_TABLE_MAX_NUM_INSTANCES 12
-
-struct nps_info_header {
-	uint32_t table_id; /* table ID */
-	uint16_t version_major; /* table version */
-	uint16_t version_minor; /* table version */
-	uint32_t size_bytes; /* size of the entire header+data in bytes = 0x000000D4 (212) */
-};
-
-struct nps_instance_info_v1_0 {
-	uint64_t base_address;
-	uint64_t limit_address;
-};
-
-struct nps_info_v1_0 {
-	struct nps_info_header header;
-	uint32_t nps_type;
-	uint32_t count;
-	struct nps_instance_info_v1_0
-		instance_info[NPS_INFO_TABLE_MAX_NUM_INSTANCES];
-};
-
-enum amd_hw_ip_block_type {
+ #ifndef _DISCOVERY_H_
+ #define _DISCOVERY_H_
+ 
+ #define PSP_HEADER_SIZE                 256
+ #define BINARY_SIGNATURE                0x28211407
+ #define DISCOVERY_TABLE_SIGNATURE       0x53445049
+ #define GC_TABLE_ID                     0x4347
+ #define HARVEST_TABLE_SIGNATURE         0x56524148
+ #define VCN_INFO_TABLE_ID               0x004E4356
+ #define MALL_INFO_TABLE_ID              0x4C4C414D
+ #define NPS_INFO_TABLE_ID 0x0053504E
+ 
+ typedef enum {
+	 IP_DISCOVERY = 0,
+	 GC,
+	 HARVEST_INFO,
+	 VCN_INFO,
+	 MALL_INFO,
+	 NPS_INFO,
+	 TOTAL_TABLES = 6
+ } table;
+ 
+ #pragma pack(1)
+ 
+ typedef struct table_info
+ {
+	 uint16_t offset;   /* Byte offset */
+	 uint16_t checksum; /* Byte sum of the table */
+	 uint16_t size;     /* Table size */
+	 uint16_t padding;
+ } table_info;
+ 
+ typedef struct binary_header
+ {
+	 /* psp structure should go at the top of this structure */
+	 uint32_t binary_signature; /* 0x7, 0x14, 0x21, 0x28 */
+	 uint16_t version_major;
+	 uint16_t version_minor;
+	 uint16_t binary_checksum;  /* Byte sum of the binary after this field */
+	 uint16_t binary_size;      /* Binary Size*/
+	 table_info table_list[TOTAL_TABLES];
+ } binary_header;
+ 
+ typedef struct die_info
+ {
+	 uint16_t die_id;
+	 uint16_t die_offset; /* Points to the corresponding die_header structure */
+ } die_info;
+ 
+ 
+ typedef struct ip_discovery_header
+ {
+	 uint32_t signature;    /* Table Signature */
+	 uint16_t version;      /* Table Version */
+	 uint16_t size;         /* Table Size */
+	 uint32_t id;           /* Table ID */
+	 uint16_t num_dies;     /* Number of Dies */
+	 die_info die_info[16]; /* list die information for up to 16 dies */
+	 union {
+		 uint16_t padding[1];	/* version <= 3 */
+		 struct {		/* version == 4 */
+			 uint8_t base_addr_64_bit : 1; /* ip structures are using 64 bit base address */
+			 uint8_t reserved : 7;
+			 uint8_t reserved2;
+		 };
+	 };
+ } ip_discovery_header;
+ 
+ typedef struct ip
+ {
+	 uint16_t hw_id;           /* Hardware ID */
+	 uint8_t number_instance;  /* instance of the IP */
+	 uint8_t num_base_address; /* Number of Base Addresses */
+	 uint8_t major;            /* HCID Major */
+	 uint8_t minor;            /* HCID Minor */
+	 uint8_t revision;         /* HCID Revision */
+ #if defined(__BIG_ENDIAN)
+	 uint8_t reserved : 4;     /* Placeholder field */
+	 uint8_t harvest : 4;      /* Harvest */
+ #else
+	 uint8_t harvest : 4;      /* Harvest */
+	 uint8_t reserved : 4;     /* Placeholder field */
+ #endif
+	 uint32_t base_address[]; /* variable number of Addresses */
+ } ip;
+ 
+ typedef struct ip_v3
+ {
+	 uint16_t hw_id;                         /* Hardware ID */
+	 uint8_t instance_number;                /* Instance number for the IP */
+	 uint8_t num_base_address;               /* Number of base addresses*/
+	 uint8_t major;                          /* Hardware ID.major version */
+	 uint8_t minor;                          /* Hardware ID.minor version */
+	 uint8_t revision;                       /* Hardware ID.revision version */
+ #if defined(__BIG_ENDIAN)
+	 uint8_t variant : 4;                    /* HW variant */
+	 uint8_t sub_revision : 4;               /* HCID Sub-Revision */
+ #else
+	 uint8_t sub_revision : 4;               /* HCID Sub-Revision */
+	 uint8_t variant : 4;                    /* HW variant */
+ #endif
+	 uint32_t base_address[];		/* Base Address list. Corresponds to the num_base_address field*/
+ } ip_v3;
+ 
+ typedef struct ip_v4 {
+	 uint16_t hw_id;                         /* Hardware ID */
+	 uint8_t instance_number;                /* Instance number for the IP */
+	 uint8_t num_base_address;               /* Number of base addresses*/
+	 uint8_t major;                          /* Hardware ID.major version */
+	 uint8_t minor;                          /* Hardware ID.minor version */
+	 uint8_t revision;                       /* Hardware ID.revision version */
+ #if defined(LITTLEENDIAN_CPU)
+	 uint8_t sub_revision : 4;               /* HCID Sub-Revision */
+	 uint8_t variant : 4;                    /* HW variant */
+ #elif defined(BIGENDIAN_CPU)
+	 uint8_t variant : 4;                    /* HW variant */
+	 uint8_t sub_revision : 4;               /* HCID Sub-Revision */
+ #endif
+ } ip_v4;
+ 
+ typedef struct die_header
+ {
+	 uint16_t die_id;
+	 uint16_t num_ips;
+ } die_header;
+ 
+ typedef struct ip_structure
+ {
+	 ip_discovery_header* header;
+	 struct die
+	 {
+		 die_header *die_header;
+		 union
+		 {
+			 ip *ip_list;
+			 ip_v3 *ip_v3_list;
+			 ip_v4 *ip_v4_list;
+		 };                                  /* IP list. Variable size*/
+	 } die;
+ } ip_structure;
+ 
+ struct gpu_info_header {
+	 uint32_t table_id;      /* table ID */
+	 uint16_t version_major; /* table version */
+	 uint16_t version_minor; /* table version */
+	 uint32_t size;          /* size of the entire header+data in bytes */
+ };
+ 
+ struct gc_info_v1_0 {
+	 struct gpu_info_header header;
+ 
+	 uint32_t gc_num_se;
+	 uint32_t gc_num_wgp0_per_sa;
+	 uint32_t gc_num_wgp1_per_sa;
+	 uint32_t gc_num_rb_per_se;
+	 uint32_t gc_num_gl2c;
+	 uint32_t gc_num_gprs;
+	 uint32_t gc_num_max_gs_thds;
+	 uint32_t gc_gs_table_depth;
+	 uint32_t gc_gsprim_buff_depth;
+	 uint32_t gc_parameter_cache_depth;
+	 uint32_t gc_double_offchip_lds_buffer;
+	 uint32_t gc_wave_size;
+	 uint32_t gc_max_waves_per_simd;
+	 uint32_t gc_max_scratch_slots_per_cu;
+	 uint32_t gc_lds_size;
+	 uint32_t gc_num_sc_per_se;
+	 uint32_t gc_num_sa_per_se;
+	 uint32_t gc_num_packer_per_sc;
+	 uint32_t gc_num_gl2a;
+ };
+ 
+ struct gc_info_v1_1 {
+	 struct gpu_info_header header;
+ 
+	 uint32_t gc_num_se;
+	 uint32_t gc_num_wgp0_per_sa;
+	 uint32_t gc_num_wgp1_per_sa;
+	 uint32_t gc_num_rb_per_se;
+	 uint32_t gc_num_gl2c;
+	 uint32_t gc_num_gprs;
+	 uint32_t gc_num_max_gs_thds;
+	 uint32_t gc_gs_table_depth;
+	 uint32_t gc_gsprim_buff_depth;
+	 uint32_t gc_parameter_cache_depth;
+	 uint32_t gc_double_offchip_lds_buffer;
+	 uint32_t gc_wave_size;
+	 uint32_t gc_max_waves_per_simd;
+	 uint32_t gc_max_scratch_slots_per_cu;
+	 uint32_t gc_lds_size;
+	 uint32_t gc_num_sc_per_se;
+	 uint32_t gc_num_sa_per_se;
+	 uint32_t gc_num_packer_per_sc;
+	 uint32_t gc_num_gl2a;
+	 uint32_t gc_num_tcp_per_sa;
+	 uint32_t gc_num_sdp_interface;
+	 uint32_t gc_num_tcps;
+ };
+ 
+ struct gc_info_v1_2 {
+	 struct gpu_info_header header;
+	 uint32_t gc_num_se;
+	 uint32_t gc_num_wgp0_per_sa;
+	 uint32_t gc_num_wgp1_per_sa;
+	 uint32_t gc_num_rb_per_se;
+	 uint32_t gc_num_gl2c;
+	 uint32_t gc_num_gprs;
+	 uint32_t gc_num_max_gs_thds;
+	 uint32_t gc_gs_table_depth;
+	 uint32_t gc_gsprim_buff_depth;
+	 uint32_t gc_parameter_cache_depth;
+	 uint32_t gc_double_offchip_lds_buffer;
+	 uint32_t gc_wave_size;
+	 uint32_t gc_max_waves_per_simd;
+	 uint32_t gc_max_scratch_slots_per_cu;
+	 uint32_t gc_lds_size;
+	 uint32_t gc_num_sc_per_se;
+	 uint32_t gc_num_sa_per_se;
+	 uint32_t gc_num_packer_per_sc;
+	 uint32_t gc_num_gl2a;
+	 uint32_t gc_num_tcp_per_sa;
+	 uint32_t gc_num_sdp_interface;
+	 uint32_t gc_num_tcps;
+	 uint32_t gc_num_tcp_per_wpg;
+	 uint32_t gc_tcp_l1_size;
+	 uint32_t gc_num_sqc_per_wgp;
+	 uint32_t gc_l1_instruction_cache_size_per_sqc;
+	 uint32_t gc_l1_data_cache_size_per_sqc;
+	 uint32_t gc_gl1c_per_sa;
+	 uint32_t gc_gl1c_size_per_instance;
+	 uint32_t gc_gl2c_per_gpu;
+ };
+ 
+ struct gc_info_v1_3 {
+	 struct gpu_info_header header;
+	 uint32_t gc_num_se;
+	 uint32_t gc_num_wgp0_per_sa;
+	 uint32_t gc_num_wgp1_per_sa;
+	 uint32_t gc_num_rb_per_se;
+	 uint32_t gc_num_gl2c;
+	 uint32_t gc_num_gprs;
+	 uint32_t gc_num_max_gs_thds;
+	 uint32_t gc_gs_table_depth;
+	 uint32_t gc_gsprim_buff_depth;
+	 uint32_t gc_parameter_cache_depth;
+	 uint32_t gc_double_offchip_lds_buffer;
+	 uint32_t gc_wave_size;
+	 uint32_t gc_max_waves_per_simd;
+	 uint32_t gc_max_scratch_slots_per_cu;
+	 uint32_t gc_lds_size;
+	 uint32_t gc_num_sc_per_se;
+	 uint32_t gc_num_sa_per_se;
+	 uint32_t gc_num_packer_per_sc;
+	 uint32_t gc_num_gl2a;
+	 uint32_t gc_num_tcp_per_sa;
+	 uint32_t gc_num_sdp_interface;
+	 uint32_t gc_num_tcps;
+	 uint32_t gc_num_tcp_per_wpg;
+	 uint32_t gc_tcp_l1_size;
+	 uint32_t gc_num_sqc_per_wgp;
+	 uint32_t gc_l1_instruction_cache_size_per_sqc;
+	 uint32_t gc_l1_data_cache_size_per_sqc;
+	 uint32_t gc_gl1c_per_sa;
+	 uint32_t gc_gl1c_size_per_instance;
+	 uint32_t gc_gl2c_per_gpu;
+	 uint32_t gc_tcp_size_per_cu;
+	 uint32_t gc_tcp_cache_line_size;
+	 uint32_t gc_instruction_cache_size_per_sqc;
+	 uint32_t gc_instruction_cache_line_size;
+	 uint32_t gc_scalar_data_cache_size_per_sqc;
+	 uint32_t gc_scalar_data_cache_line_size;
+	 uint32_t gc_tcc_size;
+	 uint32_t gc_tcc_cache_line_size;
+ };
+ 
+ struct gc_info_v2_0 {
+	 struct gpu_info_header header;
+ 
+	 uint32_t gc_num_se;
+	 uint32_t gc_num_cu_per_sh;
+	 uint32_t gc_num_sh_per_se;
+	 uint32_t gc_num_rb_per_se;
+	 uint32_t gc_num_tccs;
+	 uint32_t gc_num_gprs;
+	 uint32_t gc_num_max_gs_thds;
+	 uint32_t gc_gs_table_depth;
+	 uint32_t gc_gsprim_buff_depth;
+	 uint32_t gc_parameter_cache_depth;
+	 uint32_t gc_double_offchip_lds_buffer;
+	 uint32_t gc_wave_size;
+	 uint32_t gc_max_waves_per_simd;
+	 uint32_t gc_max_scratch_slots_per_cu;
+	 uint32_t gc_lds_size;
+	 uint32_t gc_num_sc_per_se;
+	 uint32_t gc_num_packer_per_sc;
+ };
+ 
+ struct gc_info_v2_1 {
+	 struct gpu_info_header header;
+ 
+	 uint32_t gc_num_se;
+	 uint32_t gc_num_cu_per_sh;
+	 uint32_t gc_num_sh_per_se;
+	 uint32_t gc_num_rb_per_se;
+	 uint32_t gc_num_tccs;
+	 uint32_t gc_num_gprs;
+	 uint32_t gc_num_max_gs_thds;
+	 uint32_t gc_gs_table_depth;
+	 uint32_t gc_gsprim_buff_depth;
+	 uint32_t gc_parameter_cache_depth;
+	 uint32_t gc_double_offchip_lds_buffer;
+	 uint32_t gc_wave_size;
+	 uint32_t gc_max_waves_per_simd;
+	 uint32_t gc_max_scratch_slots_per_cu;
+	 uint32_t gc_lds_size;
+	 uint32_t gc_num_sc_per_se;
+	 uint32_t gc_num_packer_per_sc;
+	 /* new for v2_1 */
+	 uint32_t gc_num_tcp_per_sh;
+	 uint32_t gc_tcp_size_per_cu;
+	 uint32_t gc_num_sdp_interface;
+	 uint32_t gc_num_cu_per_sqc;
+	 uint32_t gc_instruction_cache_size_per_sqc;
+	 uint32_t gc_scalar_data_cache_size_per_sqc;
+	 uint32_t gc_tcc_size;
+ };
+ 
+ typedef struct harvest_info_header {
+	 uint32_t signature; /* Table Signature */
+	 uint32_t version;   /* Table Version */
+ } harvest_info_header;
+ 
+ typedef struct harvest_info {
+	 uint16_t hw_id;          /* Hardware ID */
+	 uint8_t number_instance; /* Instance of the IP */
+	 uint8_t reserved;        /* Reserved for alignment */
+ } harvest_info;
+ 
+ typedef struct harvest_table {
+	 harvest_info_header header;
+	 harvest_info list[32];
+ } harvest_table;
+ 
+ struct mall_info_header {
+	 uint32_t table_id; /* table ID */
+	 uint16_t version_major; /* table version */
+	 uint16_t version_minor; /* table version */
+	 uint32_t size_bytes; /* size of the entire header+data in bytes */
+ };
+ 
+ struct mall_info_v1_0 {
+	 struct mall_info_header header;
+	 uint32_t mall_size_per_m;
+	 uint32_t m_s_present;
+	 uint32_t m_half_use;
+	 uint32_t m_mall_config;
+	 uint32_t reserved[5];
+ };
+ 
+ struct mall_info_v2_0 {
+	 struct mall_info_header header;
+	 uint32_t mall_size_per_umc;
+	 uint32_t reserved[8];
+ };
+ 
+ #define VCN_INFO_TABLE_MAX_NUM_INSTANCES 4
+ 
+ struct vcn_info_header {
+	 uint32_t table_id; /* table ID */
+	 uint16_t version_major; /* table version */
+	 uint16_t version_minor; /* table version */
+	 uint32_t size_bytes; /* size of the entire header+data in bytes */
+ };
+ 
+ struct vcn_instance_info_v1_0
+ {
+	 uint32_t instance_num; /* VCN IP instance number. 0 - VCN0; 1 - VCN1 etc*/
+	 union _fuse_data {
+		 struct {
+			 uint32_t av1_disabled : 1;
+			 uint32_t vp9_disabled : 1;
+			 uint32_t hevc_disabled : 1;
+			 uint32_t h264_disabled : 1;
+			 uint32_t reserved : 28;
+		 } bits;
+		 uint32_t all_bits;
+	 } fuse_data;
+	 uint32_t reserved[2];
+ };
+ 
+ struct vcn_info_v1_0 {
+	 struct vcn_info_header header;
+	 uint32_t num_of_instances; /* number of entries used in instance_info below*/
+	 struct vcn_instance_info_v1_0 instance_info[VCN_INFO_TABLE_MAX_NUM_INSTANCES];
+	 uint32_t reserved[4];
+ };
+ 
+ #define NPS_INFO_TABLE_MAX_NUM_INSTANCES 12
+ 
+ struct nps_info_header {
+	 uint32_t table_id; /* table ID */
+	 uint16_t version_major; /* table version */
+	 uint16_t version_minor; /* table version */
+	 uint32_t size_bytes; /* size of the entire header+data in bytes = 0x000000D4 (212) */
+ };
+ 
+ struct nps_instance_info_v1_0 {
+	 uint64_t base_address;
+	 uint64_t limit_address;
+ };
+ 
+ struct nps_info_v1_0 {
+	 struct nps_info_header header;
+	 uint32_t nps_type;
+	 uint32_t count;
+	 struct nps_instance_info_v1_0
+		 instance_info[NPS_INFO_TABLE_MAX_NUM_INSTANCES];
+ };
+ 
+ enum amd_hw_ip_block_type {
 	GC_HWIP = 1,
 	HDP_HWIP,
 	SDMA0_HWIP,
@@ -565,3 +596,5 @@ static int hw_id_map[MAX_HWIP] = {
 };

 #endif
+
+ 
--- a/Show More
+++ b/Show More