deepsek work

llm server: temperature + test client (#15444 )
* improvements to the llm server * eval script * eval llm * better eval gets 58.71 * cleanups * add temperature, but multinomial is absurdly slow * claude is so smart * lint * remove slop * no more stop
2026-03-25 10:53:39 +08:00 · 2026-03-24 21:07:15 +08:00 · 2026-03-24 18:23:53 +08:00 · 2026-03-24 04:52:54 -04:00 · 2026-03-23 22:18:38 -04:00 · 2026-03-24 10:10:46 +08:00
252 changed files with 12447 additions and 19419 deletions
@@ -45,6 +45,10 @@ inputs:
    description: "Install mesa"
    required: false
    default: 'false'
+  tinydreno:
+    description: "Install tinydreno"
+    required: false
+    default: 'false'
 runs:
  using: "composite"
  steps:
@@ -62,14 +66,14 @@ runs:
      uses: actions/cache/restore@v4
      with:
        path: ${{ github.workspace }}/.venv
-        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
+        key: venv-${{ runner.os }}-${{ runner.arch }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
    - name: Cache Python packages
      if: github.event_name != 'pull_request'
      id: restore-venv
      uses: actions/cache@v4
      with:
        path: ${{ github.workspace }}/.venv
-        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
+        key: venv-${{ runner.os }}-${{ runner.arch }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}

    # **** Caching downloads ****

@@ -195,13 +199,13 @@ runs:
      uses: actions/cache/restore@v4
      with:
        path: /var/cache/apt/archives/
-        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
    - name: Cache apt
      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true') && github.event_name != 'pull_request'
      uses: actions/cache@v4
      with:
        path: /var/cache/apt/archives/
-        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}

    - name: Run apt Update + Install
      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
@@ -233,7 +237,7 @@ runs:
      shell: bash
      run: |
        sudo mkdir -p /usr/local/lib
-        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/nimlgen/amdcomgr_dylib/releases/latest | \
+        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/tinygrad/amdcomgr_dylib/releases/latest | \
          jq -r '.assets[] | select(.name == "libamd_comgr.dylib").browser_download_url' | \
          sudo xargs curl -fL -o /usr/local/lib/libamd_comgr.dylib
        cargo build --release --manifest-path ./extra/remu/Cargo.toml
@@ -283,6 +287,7 @@ runs:

        CMAKE_ARGS="-Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5"
        if [[ "${{ runner.os }}" == "macOS" ]]; then
+          sudo xcode-select -s /Applications/Xcode_16.2.app/Contents/Developer
          CMAKE_ARGS="$CMAKE_ARGS -DBoost_INCLUDE_DIR=$(brew --prefix boost)/include -DBoost_LIBRARY_DIR=$(brew --prefix boost)/lib"
        fi

@@ -326,3 +331,9 @@ runs:
      if: inputs.mesa == 'true' && runner.os == 'macOS'
      shell: bash
      run: brew install sirhcm/tinymesa/tinymesa_cpu
+
+    # *** tinydreno ***
+    - name: Install tinydreno (linux)
+      if: inputs.tinydreno == 'true' && runner.os == 'Linux'
+      shell: bash
+      run: sudo curl -fL https://github.com/sirhcm/tinydreno/raw/refs/heads/master/libllvm-qcom.so -o /usr/lib/libllvm-qcom.so
@@ -32,6 +32,7 @@ jobs:
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
+        key: 'autogen'
        opencl: 'true'
        amd: 'true'
        cuda: 'true'
@@ -55,6 +56,7 @@ jobs:
        python3 -c "from tinygrad.runtime.autogen import libusb"
        python3 -c "from tinygrad.runtime.autogen import mesa"
        python3 -c "from tinygrad.runtime.autogen import avcodec"
+        python3 -c "from tinygrad.runtime.autogen import llvm_qcom"
        REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
    - name: Check for differences
      run: |
@@ -81,6 +83,7 @@ jobs:
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
+        key: 'autogen-mac'
        llvm: 'true'
    - name: Regenerate autogen files
      run: |
@@ -110,6 +113,8 @@ jobs:
      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: 'autogen-comgr'
    - name: Install autogen support packages
      run: |
        wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
@@ -49,8 +49,6 @@ jobs:
        source /tmp/tinygrad_pytest_ci/bin/activate
        pytest -nauto --durations=20
    - name: openpilot compile3 0.10.1 driving_vision
-      run: FLOAT16=1 CL=1 IMAGE=2 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: IMAGE=1 openpilot compile3 0.10.1 driving_vision
      run: FLOAT16=1 CL=1 IMAGE=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx

  testmacbenchmark:
@@ -187,13 +185,13 @@ jobs:
        PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
        PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
    - name: UsbGPU boot time
-      run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
+      run: sudo -E PYTHONPATH=. GMMU=0 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
+      run: sudo -E PYTHONPATH=. GMMU=0 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
+      run: sudo -E PYTHONPATH=. GMMU=0 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 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+    #  run: sudo -E PYTHONPATH=. GMMU=0 AMD=1 AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
    - name: UsbGPU (USB4/TB) boot time
      run: PYTHONPATH=. DEBUG=3 NV=1 NV_IFACE=PCI NV_NAK=1 time python3.11 test/test_tiny.py TestTiny.test_plus
    - name: UsbGPU (USB4/TB) tiny tests
@@ -332,7 +330,7 @@ jobs:
    # - name: Fuzz Padded Tensor Core GEMM (PTX)
    #   run: NV=1 NV_PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
    - name: HEVC Decode Benchmark
-      run: VALIDATE=1 MAX_FRAMES=100 JITBEAM=1 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
+      run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
    - name: Train MNIST
      run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
    - name: Run 10 CIFAR training steps
@@ -520,7 +518,7 @@ jobs:
      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    - name: Run full CIFAR training steps w 6 GPUS
      run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
-    # this needs to be mocked and testable on a local machine
+    # TODO: broken on some of the machines
    #- name: Test full tinyfs load
    #  run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
    - name: Run process replay tests
@@ -590,22 +588,22 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: test/external/process_replay/reset.py
-    - name: openpilot compile3 0.10.0 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
-    - name: openpilot compile3 0.10.0 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: openpilot compile3 0.11.0 driving_vision
+      run: BENCHMARK_LOG=openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
+    - name: IR3 openpilot compile3 0.11.0 driving_vision
+      run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM QCOM_IR3=1 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot compile3 0.11.0 driving_policy
+      run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
+    - name: openpilot compile3 0.11.0 dmonitoring
+      run: BENCHMARK_LOG=openpilot_0_11_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: DEBUG=2 openpilot compile3 0.10.1 driving_vision
-      run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: DEBUG=2 IMAGE=1 openpilot compile3 0.10.1 driving_vision
      run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: IMAGE=1 openpilot compile3 0.10.1 driving_vision
-      run: BENCHMARK_LOG=image_1_openpilot_0_10_1_vision PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.10.1 driving_vision
-      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.10.1 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
    - name: openpilot compile3 0.10.1 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: benchmark MobileNetV2 on DSP
      run: |
        # generate quantized weights
@@ -617,6 +615,27 @@ jobs:
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

+  testcommausbgpubenchmark:
+    name: UsbGPU Benchmark (comma)
+    runs-on: [self-hosted, Linux, comma4]
+    timeout-minutes: 20
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: openpilot compile3 0.10.1 driving_vision
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." GMMU=0 DEV=AMD AMD_LLVM=1 AMD_IFACE=USB ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot load_pickle 0.10.1 driving_vision
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." GMMU=0 DEV=AMD AMD_IFACE=USB ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
+
  testreddriverbenchmark:
    name: AM Benchmark
    runs-on: [self-hosted, Linux, tinyboxrandom]
@@ -678,6 +697,14 @@ jobs:
    - name: Run 10 MLPerf Bert training steps (1 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
      run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+    - name: Remote
+      run: |
+        pkill -f 'extra/remote/serve.py' || true
+        PYTHONPATH=. python3 extra/remote/serve.py 6482 &
+        sleep 1
+        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 AMD=1 AMD_IFACE=PCI python3 test/test_tiny.py
+        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 AMD=1 AMD_AQL=1 AMD_IFACE=PCI python3 test/test_tiny.py
+        pkill -f 'extra/remote/serve.py' || true
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

@@ -734,5 +761,12 @@ jobs:
    - name: Run 10 MLPerf Bert training steps (1 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
      run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+    - name: Remote
+      run: |
+        pkill -f 'extra/remote/serve.py' || true
+        PYTHONPATH=. python3 extra/remote/serve.py 6483 &
+        sleep 1
+        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6483 NV=1 python3 test/test_tiny.py
+        pkill -f 'extra/remote/serve.py' || true
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
@@ -1,7 +1,7 @@
 name: Unit Tests
 env:
  # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '17'
+  CACHE_VERSION: '18'
  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
  PYTHONPATH: ${{ github.workspace }}
@@ -70,7 +70,7 @@ jobs:
        source venv/bin/activate
        pip install $GITHUB_WORKSPACE
        cp $GITHUB_WORKSPACE/examples/beautiful_mnist.py .
-        BS=2 STEPS=10 python beautiful_mnist.py
+        BS=2 STEPS=10 MAX_BUFFER_SIZE=0 python beautiful_mnist.py
    - name: Test Docs Build
      run: python -m mkdocs build --strict
    - name: Test Docs
@@ -141,7 +141,7 @@ jobs:
        sudo apt update || true
        sudo apt install -y --no-install-recommends ninja-build
    - name: Test beautiful_mnist in torch with TINY_BACKEND
-      run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+      run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.0 MAX_BUFFER_SIZE=0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
    - name: Test some torch tests (expect failure)
      run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true

@@ -167,10 +167,10 @@ jobs:
      run: PYTHON=1 python3 test/backend/test_symbolic_ops.py
    - name: test_renderer_failures with Python emulator
      run: PYTHON=1 python3 -m pytest -rA test/backend/test_renderer_failures.py::TestRendererFailures
-    - name: Test IMAGE=2 support
+    - name: Test IMAGE support
      run: |
-        IMAGE=2 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
-        IMAGE=2 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_simple_conv2d
+        IMAGE=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
+        IMAGE=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_simple_conv2d
    - name: Test emulated METAL tensor cores
      run: |
        DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_big_gemm
@@ -244,6 +244,37 @@ jobs:
    - name: Run TYPED=1
      run: CHECK_OOB=0 DEV=CPU TYPED=1 python test/test_tiny.py

+  nulltest:
+    name: Null Tests
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: unittest-13
+        pydeps: "pillow ftfy regex pre-commit"
+        deps: testing_unit
+        llvm: 'true'
+        amd: 'true'
+    - name: Run NULL backend tests
+      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
+    - name: Run targetted tests on NULL backend
+      run: NULL=1 python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
+    # TODO: too slow
+    # - name: Run SDXL on NULL backend
+    #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
+    - name: Run Clip tests for SD MLPerf on NULL backend
+      run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
+    - name: Run AMD emulated BERT training on NULL backend
+      run: EMULATE=AMD_RDNA4 NULL=1 NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+    # TODO: support fake weights
+    #- name: Run LLaMA 7B on 4 fake devices
+    #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
+
  unittest:
    name: Unit Tests
    runs-on: ubuntu-latest
@@ -268,20 +299,6 @@ jobs:
      run: |
        CPU=1 python test/null/test_device.py TestRunAsModule.test_module_runs
        CPU=1 python -m pytest -n=auto test/unit/ --durations=20
-    - name: Run NULL backend tests
-      run: NULL=1 python -m pytest -n=auto test/null/ --durations=20
-    - name: Run targetted tests on NULL backend
-      run: NULL=1 python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
-    # TODO: too slow
-    # - name: Run SDXL on NULL backend
-    #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
-    - name: Run Clip tests for SD MLPerf on NULL backend
-      run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
-    - name: Run AMD emulated BERT training on NULL backend
-      run: EMULATE=AMD_RDNA4 NULL=1 NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
-    # TODO: support fake weights
-    #- name: Run LLaMA 7B on 4 fake devices
-    #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
    - name: Run GC tests
      run: python test/external/external_uop_gc.py
    - name: External Benchmark Schedule
@@ -352,11 +369,11 @@ jobs:
          key: gpu-image
          deps: testing_unit
          opencl: 'true'
-      - name: Test CL IMAGE=2 ops
+      - name: Test CL IMAGE=1 ops
        run: |
-          CL=1 IMAGE=2 python -m pytest -n=auto test/backend/test_ops.py --durations=20
+          CL=1 IMAGE=1 python -m pytest -n=auto test/backend/test_ops.py --durations=20
        # TODO: training is broken
-        # CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
+        # CL=1 IMAGE=1 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -401,13 +418,13 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1397 ALLOWED_GATED_READ_IMAGE=94 FLOAT16=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1486 ALLOWED_GATED_READ_IMAGE=17 FLOAT16=1 CL=1 IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
      - name: Test openpilot CL compile fp16
-        run: FLOAT16=1 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+        run: FLOAT16=1 DEBUGCL=1 CL=1 IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
      - name: Test openpilot CL compile fp32 (test correctness)
-        run: DEBUGCL=1 CL=1 IMAGE=2 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
+        run: CL=1 IMAGE=1 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
      - name: Test openpilot LLVM compile fp16
-        run: FLOAT16=1 CPU=1 CPU_LLVM=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+        run: IMAGE=1 FLOAT16=1 CPU=1 CPU_LLVM=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

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

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

@@ -644,7 +666,7 @@ jobs:
          sudo apt-get update
          sudo apt-get install llvm-21 llvm-21-tools cloc
      - name: Install rocprof-trace-decoder
-        run: sudo PYTHONPATH="." ./extra/sqtt/install_sqtt_decoder.py
+        run: sudo PYTHONPATH="." ./extra/sqtt/install_rocprof_decoder.py
      - name: Run AMD renderer tests
        run: AMD_LLVM=0 python -m pytest -n=auto test/amd/ --durations 20
      - name: Run AMD renderer tests (AMD_LLVM=1)
@@ -682,7 +704,7 @@ jobs:
      - name: Run test_tiny on MOCKAM
        run: python test/test_tiny.py
      - name: Run test_tiny on MOCKAM USB
-        run: AMD_IFACE=USB python test/test_tiny.py
+        run: GMMU=0 AMD_IFACE=USB python test/test_tiny.py
      - name: Run test_hcq on MOCKAM
        run: python -m pytest test/device/test_hcq.py

@@ -691,8 +713,7 @@ jobs:
      fail-fast: false
      matrix:
        backend: [amd, amdllvm]
-        arch: [rdna3, rdna4]
-        #arch: [rdna3, rdna4, cdna4]
+        arch: [rdna3, rdna4, cdna4]

    name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
    runs-on: ubuntu-22.04
@@ -718,7 +739,7 @@ jobs:
          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['AMD'], Device.DEFAULT"
          DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
      - name: Run pytest (amd)
-        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/testextra/test_cfg_viz.py test/external/external_test_am.py --durations=20
+        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM --durations=20
      - name: Run TRANSCENDENTAL math
        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
      - name: Run process replay tests
@@ -994,3 +1015,26 @@ jobs:
          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
          python -m pytest -n=auto test/backend/test_ops.py --durations=20
+  qcomclcompiletests:
+    name: Compile-only (QCOM CL)
+    runs-on: ubuntu-24.04-arm
+    timeout-minutes: 15
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: compile-qcomcl
+          deps: testing_unit
+          tinydreno: 'true'
+          python-version: '3.12'
+      - name: Set env
+        shell: bash
+        run: printf "NULL=1\nNULL_ALLOW_COPYOUT=1\nNULL_QCOMCL=1" >> $GITHUB_ENV
+      - name: Run test_ops
+        shell: bash
+        run: |
+          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
+          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20
@@ -1,6 +1,6 @@
 # model based off https://medium.com/data-science/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
 from typing import Callable
-from tinygrad import Tensor, TinyJit, nn, GlobalCounters
+from tinygrad import Tensor, TinyJit, nn, GlobalCounters, function
 from tinygrad.helpers import getenv, colored, trange
 from tinygrad.nn.datasets import mnist

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

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

+  @TinyJit
+  @Tensor.train()
+  def train_step(self, X_train:Tensor, Y_train:Tensor) -> Tensor:
+    opt.zero_grad()
+    samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
+    loss = self(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
+    return loss.realize(*opt.schedule_step())
+
+  @TinyJit
+  def get_test_acc(self, X_test:Tensor, Y_test:Tensor) -> Tensor: return (self(X_test).argmax(axis=1) == Y_test).mean()*100
+
 if __name__ == "__main__":
  X_train, Y_train, X_test, Y_test = mnist(fashion=getenv("FASHION"))

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

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

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

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

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

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

 from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
-from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling, profile_marker
+from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling, profile_marker, DEBUG
 from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict, safe_load, safe_save
 from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW

@@ -13,6 +13,8 @@ from extra.bench_log import BenchEvent, WallTimeEvent
 # TODO: fix benchmark logging and use tinygrad tqdm
 from tqdm import tqdm

+from tinygrad.uop.ops import UOp
+
 def train_resnet():
  from extra.models import resnet
  from examples.mlperf.dataloader import batch_load_resnet
@@ -1282,7 +1284,7 @@ def train_bert():
        previous_step = i

 def train_llama3():
-  from extra.models.llama import Transformer
+  from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad
  from examples.llama3 import MODEL_PARAMS
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
  from examples.mlperf.optim import GradAccClipAdamW
@@ -1335,10 +1337,16 @@ def train_llama3():
  model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
  # vocab_size from the mixtral tokenizer
  if not SMALL: model_params |= {"vocab_size": 32000}
+  real_vocab_size = model_params['vocab_size']
  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
  print(f"model parameters: {model_params}")

-  model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
+  # pad vocab
+  if (MP := getenv("MP", 1)) > 1: model_params['vocab_size'] = round_up(model_params['vocab_size'], 256 * MP)
+  vocab_mask:Tensor = Tensor.arange(model_params['vocab_size']).reshape(1, 1, -1) >= real_vocab_size
+
+  model = FlatTransformer(**model_params, max_context=SEQLEN)
+
  params = get_parameters(model)
  # weights are all bfloat16 for now
  assert params and all(p.dtype == dtypes.bfloat16 for p in params)
@@ -1347,38 +1355,25 @@ def train_llama3():
    for v in get_parameters(model):
      v = v.assign(Tensor.empty(v.shape))

-  if (DP := getenv("DP", 1)) > 1:
-    device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-    for v in get_parameters(model):
-      v.shard_(device, axis=None)
+  is_dp = (DP := getenv("DP", 1)) > 1
+  is_mp = (MP := getenv("MP", 1)) > 1
+  is_sharding = is_dp or is_mp
+  device_count = max(DP, MP)
+  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))

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

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

  # init grads
-  for p in optim.params:
-    p.grad = p.zeros_like().contiguous().realize()
-  grads: list[Tensor] = [p.grad for p in optim.params]
+  grads = [Tensor.zeros_like(p).contiguous() for p in optim.params]

  scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)

@@ -1393,51 +1388,47 @@ def train_llama3():

  @TinyJit
  def minibatch(tokens:Tensor):
-    tokens = tokens.to(None)
-    if (DP := getenv("DP", 1)) > 1:
-      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-      tokens = tokens.shard(device, 0)
-    if (MP := getenv("MP", 1)) > 1:
-      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
-      tokens = tokens.shard(device)
-    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-    loss.backward()
-    assert all(p.grad is g for p,g in zip(optim.params, grads))
-    Tensor.realize(loss, *grads)
-    return loss.flatten().float().to("CPU")
+    if is_dp: tokens = tokens.to(None).shard(device, 0)
+    if is_mp: tokens = tokens.shard(device)
+    if not is_sharding: tokens = tokens.to(None)
+    logits:Tensor = model(tokens[:, :-1])
+    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
+
+    for i,(t,g) in enumerate(zip(optim.params, loss.gradient(*optim.params))):
+      grads[i].replace(Tensor(grads[i].uop.after(UOp.group(*apply_grad(grads[i].uop, g.uop))), device=t.device))
+
+    loss_cpu = loss.flatten().float().to("CPU")
+    return loss_cpu.realize(*grads)

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

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

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

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

  @TinyJit
  @Tensor.train(False)
  def eval_step(tokens:Tensor):
-    tokens = tokens.to(None)
-    if (DP := getenv("DP", 1)) > 1:
-      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-      tokens = tokens.shard(device, 0)
-    if (MP := getenv("MP", 1)) > 1:
-      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
-      tokens = tokens.shard(device)
-    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    if is_dp: tokens = tokens.to(None).shard(device, 0)
+    if is_mp: tokens = tokens.shard(device)
+    if not is_sharding: tokens = tokens.to(None)
+    logits:Tensor = model(tokens[:, :-1])
+    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
    return loss.flatten().float().to("CPU")

  # ** data iters **
  def fake_data(bs, samples):
+    import numpy as np
    for _ in range(samples // bs):
-      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=model_params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
+      fake_data_np = np.random.randint(0, model_params["vocab_size"], size=(bs, SEQLEN + 1), dtype=np.int32)
+      yield Tensor(fake_data_np, device="NPY")

  def get_train_iter():
    if getenv("FAKEDATA", 0):
@@ -1464,49 +1455,53 @@ def train_llama3():
  step_times = []
  while i < MAX_STEPS:
    GlobalCounters.reset()
+    actual_gbs = GBS if i >= 2 else BS
    if getenv("TRAIN", 1):
      profile_marker(f"train @ {i}")
      st = time.perf_counter()

      stopped = False
-      for _ in range(grad_acc):
+      losses, data_time, dev_time = [], 0, 0
+      for _ in range(grad_acc if i >= 2 else 1):
        ist = time.perf_counter()
        try: tokens = next(train_iter)
        except StopIteration:
          stopped = True
          break
-        dt = time.perf_counter()
-        loss = minibatch(tokens)
+        mst = time.perf_counter()
+        data_time += mst - ist
+        losses.append(minibatch(tokens).item())
+        dev_time += time.perf_counter() - mst
      if stopped: break

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

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

      mem_gb = GlobalCounters.mem_used / 1e9
      gflops = GlobalCounters.global_ops / 1e9 / dev_time
      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * max(getenv("DP", 1), getenv("MP", 1)) * 2.3e15)) * 100
      tqdm.write(
          f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
-          f"{lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
+          f"{lr:.12f} LR, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
+      if DEBUG >= 1: tqdm.write("  mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))

      if WANDB:
        wandb.log({
-          "lr": lr, "train/loss": loss,
+          "train/loss": loss,
+          "train/lr": lr,
+          "train/grad_norm": grad_norm,
          "train/step_time": step_time,
          "train/gbs_time": gbs_time,
          "train/optim_time": optim_time,
@@ -1535,7 +1530,7 @@ def train_llama3():
        print(f"epoch global_ops: {GlobalCounters.global_ops:_}, "
              f"epoch global_mem: {GlobalCounters.global_mem:_}")

-    if (sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1)) or (BENCHMARK and i == BENCHMARK):
+    if (sequences_seen // EVAL_FREQ != (sequences_seen - actual_gbs) // EVAL_FREQ and (i != 1 or EVAL_FREQ == 1)) or (BENCHMARK and i == BENCHMARK):
      if EVAL_BS == 0: return
      tqdm.write(f"evaluating after {sequences_seen} sequences")
      profile_marker(f"eval @ {i}")
@@ -1543,7 +1538,7 @@ def train_llama3():
      # run eval
      eval_losses = []
      eval_iter = get_eval_iter()
-      tqdm.write(f"evaluating {5760//EVAL_BS} batches of {EVAL_BS} sequences")
+      tqdm.write(f"evaluating {EVAL_SAMPLES//EVAL_BS} batches of {EVAL_BS} sequences")

      for j,tokens in tqdm(enumerate(eval_iter), total=EVAL_SAMPLES//EVAL_BS):
        eval_losses += eval_step(tokens).tolist()
@@ -0,0 +1,193 @@
+import math, os
+if __name__ == "__main__":
+  os.environ["DEFAULT_FLOAT"] = "bfloat16"
+  os.environ["OPTIM_DTYPE"] = "bfloat16"
+  if "DEV" not in os.environ: os.environ["DEV"] = "NULL"
+  # CDNA
+  os.environ["EMULATE"] = "AMD_CDNA4"
+  os.environ["DEVICE_IN_FUNCTION_BUG"] = "1"
+  os.environ["ALL2ALL"] = "1"
+  os.environ["USE_ATOMICS"] = "1"
+  if "HK_FLASH_ATTENTION" not in os.environ:
+    os.environ["HK_FLASH_ATTENTION"] = "1"
+    if "ASM_GEMM" not in os.environ:
+      os.environ["ASM_GEMM"] = "1"
+from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
+from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
+from tinygrad.uop.ops import Ops, UOp
+from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
+
+FP8 = getenv("FP8", 0)
+
+FP8_DTYPE = dtypes.fp8e4m3
+FP8_MAX = 448.0
+
+def quantize_fp8(x:Tensor):
+  scale = FP8_MAX / (x.abs().max().detach() + 1e-8)
+  x_scaled = x * scale
+  x_clamped = x_scaled + (x_scaled.detach().clamp(-FP8_MAX, FP8_MAX) - x_scaled.detach())  # STE
+  return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal()
+
+def matmul(x:Tensor, w:Tensor) -> Tensor:
+  if not FP8: return x @ w.T
+  # weights are already FP8, just quantize activations
+  x_fp8, x_scale = quantize_fp8(x)
+  return x_fp8.dot(w.T, dtype=dtypes.float) * x_scale
+
+def rmsnorm(x_in:Tensor, eps:float):
+  x = x_in.float()
+  x = x * (x.square().mean(-1, keepdim=True) + eps).rsqrt()
+  return x.cast(x_in.dtype)
+
+class FlatTransformer:
+  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
+               rope_theta:int=10000, max_context:int=1024):
+    self.vocab_size = vocab_size
+    self.n_layers = n_layers
+    self.n_heads = n_heads
+    self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads # n_kv_heads != n_heads implies MQA [arxiv/2307.09288, A.2.1]
+    self.head_dim = dim // n_heads
+    self.n_rep = self.n_heads // self.n_kv_heads
+
+    # Attention
+    self.wqkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
+    self.wo = self.lin_per_layer(self.n_heads * self.head_dim, dim)
+
+    # FeedForward
+    self.w1 = self.lin_per_layer(dim, hidden_dim)
+    self.w2 = self.lin_per_layer(hidden_dim, dim)
+    self.w3 = self.lin_per_layer(dim, hidden_dim)
+
+    self.norm_eps = norm_eps
+    self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
+    self.ffn_norm = Tensor.ones(n_layers, dim).contiguous()
+
+    # output
+    self.norm = nn.RMSNorm(dim, norm_eps)
+    self.tok_embeddings = nn.Embedding(vocab_size, dim)
+    self.output = nn.Linear(dim, vocab_size, bias=False)
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)
+
+  def lin_per_layer(self, in_features:int, out_features:int):
+    bound = 1 / math.sqrt(in_features)
+    dt = FP8_DTYPE if FP8 else None
+    if getenv("ZEROS"): return Tensor.zeros(self.n_layers, out_features, in_features, dtype=dt)
+    return Tensor.uniform(self.n_layers, out_features, in_features, low=-bound, high=bound, dtype=dt)
+
+  def attention(self, x:Tensor, freqs_cis:Tensor, attention_norm:Tensor, wqkv:Tensor, wo:Tensor):
+    x = rmsnorm(x, self.norm_eps) * attention_norm
+    xqkv = matmul(x, wqkv)
+
+    bsz, seqlen, _ = xqkv.shape
+    # interleaved layout: each kv group has [n_rep q heads, 1 k head, 1 v head] for clean MP sharding
+    xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
+    xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
+    xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
+    xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
+
+    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
+    xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
+    attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
+    attn = attn.reshape(bsz, seqlen, -1)
+    return matmul(attn, wo)
+
+  def feed_forward(self, x:Tensor, ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor):
+    x = rmsnorm(x, self.norm_eps) * ffn_norm
+    x_w1 = matmul(x, w1).silu()
+    x_w3 = matmul(x.contiguous_backward(), w3)
+    return matmul(x_w1 * x_w3, w2)
+
+  @function(precompile=True, precompile_backward=True)
+  def run_layer(self, x:Tensor, freqs_cis:Tensor,
+                attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
+                ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor):
+    h = x + self.attention(x, freqs_cis, attention_norm, wqkv, wo)
+    return h + self.feed_forward(h, ffn_norm, w1, w2, w3)
+
+  def shard(self, device:tuple[str, ...], mp:bool=False):
+    from tinygrad.nn.state import get_parameters
+    if not mp:
+      for v in get_parameters(self): v.shard_(device, axis=None)
+    else:
+      # flat per-layer weights: axis 0 is n_layers, so shard axes are +1 vs per-layer Transformer
+      self.wqkv.shard_(device, axis=1).realize()          # (n_layers, out, dim) shard out
+      self.wo.shard_(device, axis=2).realize()             # (n_layers, dim, in) shard in
+      self.w1.shard_(device, axis=1).realize()             # (n_layers, hidden, dim) shard out
+      self.w2.shard_(device, axis=2).realize()             # (n_layers, dim, hidden) shard in
+      self.w3.shard_(device, axis=1).realize()             # (n_layers, hidden, dim) shard out
+      self.attention_norm.shard_(device, axis=None).realize()
+      self.ffn_norm.shard_(device, axis=None).realize()
+      self.norm.weight.shard_(device, axis=None).realize()
+      self.tok_embeddings.weight.shard_(device, axis=0).realize()
+      self.output.weight.shard_(device, axis=0).realize()
+      self.freqs_cis.shard_(device, axis=None).realize()
+
+  def __call__(self, tokens:Tensor):
+    h = self.tok_embeddings(tokens)
+    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
+    for i in range(self.n_layers):
+      h = self.run_layer(h, freqs_cis,
+                         self.attention_norm[i], self.wqkv[i], self.wo[i],
+                         self.ffn_norm[i], self.w1[i], self.w2[i], self.w3[i])
+    logits = self.output(self.norm(h))
+    return logits
+
+# TODO: this shouldn't be needed, but it prevents a copy of the grads. CAT can help
+def apply_grad(old_grad:UOp, new_grad:UOp) -> list[UOp]:
+  if new_grad.op == Ops.ADD:
+    return apply_grad(old_grad, new_grad.src[0])+apply_grad(old_grad, new_grad.src[1])
+  elif new_grad.op == Ops.PAD:
+    grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(new_grad.src[0].shape, new_grad.marg)])
+    return apply_grad(old_grad.shrink(grad_shrink), new_grad.src[0])
+  else:
+    return [old_grad.store(old_grad + new_grad)]
+
+if __name__ == "__main__":
+  config = {}
+  BS                 = config["BS"]                     = getenv("BS", 16)
+  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
+
+  from examples.llama3 import MODEL_PARAMS
+  model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
+  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
+  model = FlatTransformer(**model_params, max_context=SEQLEN)
+  state = nn.state.get_state_dict(model)
+  print("tensor count:", len(state))
+
+  # shard the model
+  from tinygrad import Device
+  if (DP := getenv("DP", 1)) > 1:
+    model.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)))
+  if (MP := getenv("MP", 1)) > 1:
+    model.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)), mp=True)
+
+  # preallocate all the grad buffers and zero them out
+  grads = {x:Tensor.zeros_like(x).contiguous() for x in state.values() if x.requires_grad is None}
+
+  # print model size
+  sz = 0
+  for k,v in state.items():
+    print(f"{colored(k, 'green' if v in grads else 'white'):30s} {str(v.shape):30s} {str(v.dtype):20s} {v.device}  {v.nbytes()/1e9:.2f} GB")
+    sz += v.nbytes()
+  print(f"total sz: {sz/1e9:.2f} GB")
+
+  with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=model.vocab_size, dtype=dtypes.int)
+  with Timing("realize weights/grads/data: "): Tensor.realize(*state.values(), *grads.values(), tokens)
+  print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
+  if DP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)), axis=0)
+  if MP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)))
+
+  @TinyJit
+  def jit_step(tokens:Tensor):
+    with Timing("python forward: "): loss = model(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
+    with Timing("python backward: "):
+      for t,g in zip(grads, loss.gradient(*grads)):
+        grads[t] = Tensor(grads[t].uop.after(UOp.group(*apply_grad(grads[t].uop, g.uop))), device=t.device)
+    with Timing("run step: "): loss.realize(*grads.values())
+
+  for i in range(6):
+    GlobalCounters.reset()
+    profile_marker(f"step {i}")
+    with Timing(colored(f"*** step {i}: ", "red")):
+      jit_step(tokens)
+  print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
@@ -0,0 +1,80 @@
+from tinygrad import Tensor, nn
+from tinygrad.helpers import getenv
+from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
+
+class Attention:
+  def __init__(self, dim:int, n_heads:int, n_kv_heads:int|None=None, linear=nn.Linear):
+    self.n_heads = n_heads
+    self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads # n_kv_heads != n_heads implies MQA [arxiv/2307.09288, A.2.1]
+    self.head_dim = dim // n_heads
+    self.n_rep = self.n_heads // self.n_kv_heads
+
+    if getenv("WQKV"):
+      self.wqkv = linear(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2, bias=False)
+    else:
+      self.wq = linear(dim, self.n_heads * self.head_dim, bias=False)
+      self.wk = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
+      self.wv = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
+
+    self.wo = linear(self.n_heads * self.head_dim, dim, bias=False)
+
+  def __call__(self, x:Tensor, freqs_cis:Tensor) -> Tensor:
+    if getenv("WQKV"):
+      xqkv = self.wqkv(x)
+      xqkv = xqkv.reshape(xqkv.shape[0], xqkv.shape[1], self.n_kv_heads, self.n_rep + 2, self.head_dim)
+      xq = xqkv[:, :, :, :self.n_rep].reshape(xqkv.shape[0], xqkv.shape[1], -1)
+      xk = xqkv[:, :, :, self.n_rep:self.n_rep+1].reshape(xqkv.shape[0], xqkv.shape[1], -1)
+      xv = xqkv[:, :, :, self.n_rep+1:self.n_rep+2].reshape(xqkv.shape[0], xqkv.shape[1], -1)
+    else:
+      xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
+
+    xq = xq.reshape(xq.shape[0], xq.shape[1], self.n_heads, self.head_dim)
+    xk = xk.reshape(xk.shape[0], xk.shape[1], self.n_kv_heads, self.head_dim)
+    xv = xv.reshape(xv.shape[0], xv.shape[1], self.n_kv_heads, self.head_dim)
+
+    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
+    bsz, seqlen, _, _ = xq.shape
+
+    xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
+    attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
+
+    attn = attn.reshape(bsz, seqlen, -1)
+    return self.wo(attn)
+
+class FeedForward:
+  def __init__(self, dim:int, hidden_dim:int, linear=nn.Linear):
+    self.w1 = linear(dim, hidden_dim, bias=False)
+    self.w2 = linear(hidden_dim, dim, bias=False)
+    self.w3 = linear(dim, hidden_dim, bias=False) # the gate in Gated Linear Unit
+
+  def __call__(self, x:Tensor) -> Tensor:
+    w1 = self.w1(x).silu()
+    w3 = self.w3(x)
+    return self.w2(w1 * w3)
+
+class TransformerBlock:
+  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int|None, norm_eps:float, linear=nn.Linear):
+    self.attention = Attention(dim, n_heads, n_kv_heads, linear)
+    self.feed_forward = FeedForward(dim, hidden_dim, linear)
+    self.attention_norm = nn.RMSNorm(dim, norm_eps)
+    self.ffn_norm = nn.RMSNorm(dim, norm_eps)
+
+  def __call__(self, x:Tensor, freqs_cis:Tensor):
+    h = x + self.attention(self.attention_norm(x), freqs_cis)
+    return h + self.feed_forward(self.ffn_norm(h))
+
+class Transformer:
+  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
+               rope_theta:int=10000, max_context:int=1024, linear=nn.Linear, embedding=nn.Embedding):
+    self.layers = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, linear) for _ in range(n_layers)]
+    self.norm = nn.RMSNorm(dim, norm_eps)
+    self.tok_embeddings = embedding(vocab_size, dim)
+    self.output = nn.Linear(dim, vocab_size, bias=False) if embedding == nn.Embedding else linear(dim, vocab_size, bias=False)
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)
+
+  def __call__(self, tokens:Tensor):
+    h = self.tok_embeddings(tokens)
+    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
+    for layer in self.layers: h = layer(h, freqs_cis)
+    logits = self.output(self.norm(h))
+    return logits
@@ -0,0 +1,140 @@
+import os
+os.environ["WQKV"] = "1"
+import unittest
+import numpy as np
+from tinygrad import Tensor, nn, dtypes
+from tinygrad.nn.state import get_parameters
+from tinygrad.device import is_dtype_supported
+from examples.mlperf.models.llama import Transformer
+from examples.mlperf.models.flat_llama import FlatTransformer
+
+def copy_weights(flat:FlatTransformer, ref:Transformer):
+  n_layers = flat.n_layers
+  Tensor.realize(*nn.state.get_state_dict(ref).values())
+  flat.wqkv.assign(Tensor(np.stack([ref.layers[i].attention.wqkv.weight.numpy() for i in range(n_layers)])))
+  flat.wo.assign(Tensor(np.stack([ref.layers[i].attention.wo.weight.numpy() for i in range(n_layers)])))
+  flat.w1.assign(Tensor(np.stack([ref.layers[i].feed_forward.w1.weight.numpy() for i in range(n_layers)])))
+  flat.w2.assign(Tensor(np.stack([ref.layers[i].feed_forward.w2.weight.numpy() for i in range(n_layers)])))
+  flat.w3.assign(Tensor(np.stack([ref.layers[i].feed_forward.w3.weight.numpy() for i in range(n_layers)])))
+  flat.attention_norm.assign(Tensor(np.stack([ref.layers[i].attention_norm.weight.numpy() for i in range(n_layers)])))
+  flat.ffn_norm.assign(Tensor(np.stack([ref.layers[i].ffn_norm.weight.numpy() for i in range(n_layers)])))
+  flat.norm.weight.assign(Tensor(ref.norm.weight.numpy()))
+  flat.tok_embeddings.weight.assign(Tensor(ref.tok_embeddings.weight.numpy()))
+  flat.output.weight.assign(Tensor(ref.output.weight.numpy()))
+
+class TestFlatLlama(unittest.TestCase):
+  def test_forward_match(self):
+    Tensor.manual_seed(42)
+    params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
+    ref = Transformer(**params)
+    flat = FlatTransformer(**params)
+    copy_weights(flat, ref)
+    Tensor.realize(*nn.state.get_state_dict(flat).values())
+
+    tokens = Tensor([[1, 50, 100, 999, 2]])
+    ref_logits = ref(tokens).realize()
+    flat_logits = flat(tokens).realize()
+    self.assertEqual(ref_logits.shape, flat_logits.shape)
+    diff = (ref_logits - flat_logits).abs().max().item()
+    self.assertLess(diff, 1e-5, f"forward mismatch: max abs diff {diff}")
+
+  def test_backward_match(self):
+    Tensor.manual_seed(42)
+    params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
+    ref = Transformer(**params)
+    flat = FlatTransformer(**params)
+    copy_weights(flat, ref)
+
+    for p in get_parameters(ref): p.requires_grad_(True)
+    for p in get_parameters(flat): p.requires_grad_(True)
+    Tensor.realize(*nn.state.get_state_dict(flat).values())
+
+    tokens = Tensor([[1, 50, 100, 999, 2, 10]])
+
+    ref_loss = ref(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
+    ref_loss.backward()
+    ref_grads = {k: v.grad.numpy() for k, v in nn.state.get_state_dict(ref).items() if v.grad is not None}
+
+    flat_loss = flat(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
+    flat_loss.backward()
+    flat_grads = {k: v.grad.numpy() for k, v in nn.state.get_state_dict(flat).items() if v.grad is not None}
+
+    # check loss matches
+    self.assertAlmostEqual(ref_loss.item(), flat_loss.item(), places=4)
+
+    # check output weight grad matches
+    diff = abs(ref_grads["output.weight"] - flat_grads["output.weight"]).max()
+    self.assertLess(diff, 1e-4, f"output.weight grad mismatch: max abs diff {diff}")
+
+    # check per-layer weight grads match
+    for i in range(params["n_layers"]):
+      for flat_key, ref_key in [
+        ("wqkv", f"layers.{i}.attention.wqkv.weight"),
+        ("wo", f"layers.{i}.attention.wo.weight"),
+        ("w1", f"layers.{i}.feed_forward.w1.weight"),
+        ("w2", f"layers.{i}.feed_forward.w2.weight"),
+        ("w3", f"layers.{i}.feed_forward.w3.weight"),
+      ]:
+        diff = abs(ref_grads[ref_key] - flat_grads[flat_key][i]).max()
+        self.assertLess(diff, 1e-4, f"layer {i} {flat_key} grad mismatch: max abs diff {diff}")
+
+  @unittest.skipUnless(os.getenv("CPU", "") == "1", "multi-device CPU test")
+  def test_forward_match_mp(self):
+    Tensor.manual_seed(42)
+    params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
+    from tinygrad import Device
+    devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
+    ref = Transformer(**params)
+    flat = FlatTransformer(**params)
+    copy_weights(flat, ref)
+    Tensor.realize(*nn.state.get_state_dict(flat).values())
+    flat.shard(devices, mp=True)
+
+    tokens = Tensor([[1, 50, 100, 999, 2]], device=devices[0])
+    ref_logits = ref(tokens.to(devices[0])).numpy()
+    flat_logits = flat(tokens.shard(devices)).numpy()
+    self.assertEqual(ref_logits.shape, flat_logits.shape)
+    np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
+
+  @unittest.skipUnless(os.getenv("CPU", "") == "1", "multi-device CPU test")
+  def test_forward_match_dp(self):
+    Tensor.manual_seed(42)
+    params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
+    from tinygrad import Device
+    devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
+    ref = Transformer(**params)
+    flat = FlatTransformer(**params)
+    copy_weights(flat, ref)
+    Tensor.realize(*nn.state.get_state_dict(flat).values())
+    flat.shard(devices)
+
+    tokens = Tensor([[1, 50, 100, 999, 2], [2, 100, 50, 1, 999]], device=devices[0])
+    ref_logits = ref(tokens.to(devices[0])).numpy()
+    flat_logits = flat(tokens.shard(devices, axis=0)).numpy()
+    self.assertEqual(ref_logits.shape, flat_logits.shape)
+    np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
+
+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), "fp8 not supported on this device")
+  def test_forward_fp8(self):
+    import examples.mlperf.models.flat_llama as flat_llama_mod
+    old_fp8 = flat_llama_mod.FP8
+    try:
+      flat_llama_mod.FP8 = 1
+      Tensor.manual_seed(42)
+      params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
+      ref = Transformer(**params)
+      flat = FlatTransformer(**params)
+      copy_weights(flat, ref)
+      Tensor.realize(*nn.state.get_state_dict(flat).values())
+
+      tokens = Tensor([[1, 50, 100, 999, 2]])
+      ref_logits = ref(tokens).numpy()
+      flat_logits = flat(tokens).numpy()
+      self.assertEqual(ref_logits.shape, flat_logits.shape)
+      # FP8 has lower precision, allow larger tolerance
+      np.testing.assert_allclose(flat_logits, ref_logits, atol=1.0, rtol=0.1)
+    finally:
+      flat_llama_mod.FP8 = old_fp8
+
+if __name__ == "__main__":
+  unittest.main()
@@ -7,40 +7,53 @@ class GradAccClipAdamW(Optimizer):
  def __init__(self, params:list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, grad_acc=1, clip_norm=1.0, device=None, fused=FUSE_OPTIM):
    super().__init__(params, lr, device, fused)
    self.b1, self.b2, self.eps, self.wd = b1, b2, eps, weight_decay
-    self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False).contiguous() for _ in [b1, b2])
+    self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False) for _ in [b1, b2])
    self.m = self._new_optim_param()
    self.v = self._new_optim_param()
    self.grad_acc, self.clip_norm = grad_acc, clip_norm

+  def fstep(self, grads:list[Tensor]):
+    if self.fused:
+      out, extra = self._step([], grads)
+      updates = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
+    else:
+      updates, extra = self._step([], grads)
+    for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i]))
+    to_realize = extra+self.params+self.buffers
+
+    Tensor.realize(*to_realize)
+    return extra[-1]
+
  def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
+    grads = list(grads)
+
    for i in range(len(grads)):
      if grads[i].device != self.m[i].device: grads[i] = grads[i].to(self.m[i].device)

    if self.fused:
-      grads[0] = grads[0] / self.grad_acc
+      grads[0].assign(grads[0] / self.grad_acc)
      total_norm = grads[0].float().square().sum().sqrt()
-      grads[0] = (grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype)
+      grads[0].assign((grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype))
    else:
-      total_norm = Tensor.zeros((), dtype=dtypes.float32, device=self.device)
-      for g in grads:
-        total_norm += g.float().square().sum()
-      total_norm = total_norm.sqrt()
      for i in range(len(grads)):
-        grads[i] = grads[i] / self.grad_acc
-        grads[i] = (grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype)
+        grads[i].assign(grads[i] / self.grad_acc)
+      total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
+      for i in range(len(grads)):
+        grads[i].assign((grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype))

    ret = []
    self.b1_t *= self.b1
    self.b2_t *= self.b2
-    for i, (t, g) in enumerate(zip(params, grads)):
+    for i, g in enumerate(grads):
      self.m[i].assign((self.b1 * self.m[i] + (1.0 - self.b1) * g).cast(self.m[i].dtype))
      self.v[i].assign((self.b2 * self.v[i] + (1.0 - self.b2) * (g * g)).cast(self.v[i].dtype))
-      m_hat = self.m[i] / (1.0 - self.b1_t)
-      v_hat = self.v[i] / (1.0 - self.b2_t)
+      m_hat = (self.m[i] / (1.0 - self.b1_t)).cast(self.m[i].dtype)
+      v_hat = (self.v[i] / (1.0 - self.b2_t)).cast(self.v[i].dtype)
      up = m_hat / (v_hat.sqrt() + self.eps)
-      ret.append((self.lr * up).cast(t.dtype))
-    return ret, [self.b1_t, self.b2_t] + self.m + self.v
+      ret.append((self.lr * up).cast(g.dtype))
+    return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]

  def _apply_update(self, t:Tensor, up:Tensor) -> Tensor:
-    up = up.shard_like(t) + self.lr.to(t.device) * self.wd * t.detach()
+    wd = self.wd if t.ndim >= 3 else 0.0
+    up = up.shard_like(t) + self.lr.to(t.device) * wd * t.detach()
    return t.detach() - up.cast(t.dtype)
@@ -0,0 +1,38 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-2}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-1} MP=${MP:-8}
+export BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4/"
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"405B"}
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-5760}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+export FAKEDATA=1 BENCHMARK=10
+if [ -z "$FULL_LAYERS" ]; then
+  export LLAMA_LAYERS=2
+fi
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,32 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+
+export DEBUG=${DEBUG:-0}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-1} MP=${MP:-8}
+export BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-1152}
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4/"
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"405B"}
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-$RANDOM}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+python3 examples/mlperf/model_train.py
@@ -5,6 +5,7 @@ export DEV=${DEV:-AMD}
 export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1

 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
@@ -14,7 +15,7 @@ export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-0}

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

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

 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
@@ -14,7 +15,7 @@ export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-0}

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

 export MODEL="llama3"
@@ -22,7 +23,7 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="4e-4" END_LR="4e-5" WARMUP_SAMPLES=256 MAX_STEPS=1200000
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
 export SEQLEN=${SEQLEN:-8192}
@@ -0,0 +1,38 @@
+#!/usr/bin/env bash
+
+export PYTHONPATH="."
+export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
+export CHECK_OOB=0
+export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
+
+export DEBUG=${DEBUG:-0}
+export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
+export ALL2ALL=${ALL2ALL:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
+export ASM_GEMM=${ASM_GEMM:-1}
+export WQKV=${WQKV:-1}
+export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}
+
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DP=${DP:-1} MP=${MP:-8} BS=${BS:-1} EVAL_BS=${EVAL_BS:-1} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-32}
+export GBS=$((BS * GRADIENT_ACC_STEPS))
+
+export MODEL="llama3"
+export BASEDIR="/raid/datasets/c4-8b/"
+export SMALL=1
+export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
+export EVAL_TARGET=3.3 EVAL_FREQ=12288
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
+export SAMPLES=$((MAX_STEPS * GBS))
+export SEQLEN=${SEQLEN:-8192}
+
+export SEED=${SEED:-$RANDOM}
+export DATA_SEED=${DATA_SEED:-5760}
+
+export JITBEAM=${JITBEAM:-3}
+export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
+
+python3 examples/mlperf/model_train.py
@@ -1,6 +1,5 @@
 #!/bin/bash
 export BENCHMARK=5
 export EVAL_BS=0
-export VIZ=${VIZ:--1}
-examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
-PYTHONPATH="." extra/viz/cli.py --profile --device "AMD" --top 20
+VIZ=${VIZ:--1} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
+extra/viz/cli.py --profile --device "AMD" --limit 20
@@ -1,12 +1,11 @@
 from pathlib import Path
-import numpy as np
 import torch
 from torchvision.utils import make_grid, save_image
 from tinygrad.nn.state import get_parameters
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import trange
 from tinygrad.nn import optim
-from extra.datasets import fetch_mnist
+from tinygrad.nn.datasets import mnist

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

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

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

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

 if __name__ == "__main__":
  # data for training and validation
-  images_real = np.vstack(fetch_mnist()[::2])
+  X_train, _, _, _ = mnist()
  ds_noise = Tensor.randn(64, 128, requires_grad=False)
  # parameters
  epochs, batch_size, k = 300, 512, 1
  sample_interval = epochs // 10
-  n_steps = len(images_real) // batch_size
+  n_steps = X_train.shape[0] // batch_size
  # models and optimizer
  generator = LinearGen()
  discriminator = LinearDisc()
@@ -84,24 +83,24 @@ if __name__ == "__main__":
  output_dir = Path(".").resolve() / "outputs"
  output_dir.mkdir(exist_ok=True)
  # optimizers
-  optim_g = optim.Adam(get_parameters(generator),lr=0.0002, b1=0.5)  # 0.0002 for equilibrium!
-  optim_d = optim.Adam(get_parameters(discriminator),lr=0.0002, b1=0.5)
+  optim_g = optim.Adam(get_parameters(generator), lr=0.0002, b1=0.5)  # 0.0002 for equilibrium!
+  optim_d = optim.Adam(get_parameters(discriminator), lr=0.0002, b1=0.5)
  # training loop
-  Tensor.training = True
-  for epoch in (t := trange(epochs)):
-    loss_g, loss_d = 0.0, 0.0
-    for _ in range(n_steps):
-      data_real = make_batch(images_real)
-      for step in range(k):  # Try with k = 5 or 7.
+  with Tensor.train():
+    for epoch in (t := trange(epochs)):
+      loss_g, loss_d = 0.0, 0.0
+      for _ in range(n_steps):
+        data_real = make_batch(X_train)
+        for step in range(k):  # Try with k = 5 or 7.
+          noise = Tensor.randn(batch_size, 128)
+          data_fake = generator.forward(noise).detach()
+          loss_d += train_discriminator(optim_d, data_real, data_fake)
        noise = Tensor.randn(batch_size, 128)
-        data_fake = generator.forward(noise).detach()
-        loss_d += train_discriminator(optim_d, data_real, data_fake)
-      noise = Tensor.randn(batch_size, 128)
-      data_fake = generator.forward(noise)
-      loss_g += train_generator(optim_g, data_fake)
-    if (epoch + 1) % sample_interval == 0:
-      fake_images = generator.forward(ds_noise).detach().numpy()
-      fake_images = (fake_images.reshape(-1, 1, 28, 28) + 1) / 2  # 0 - 1 range.
-      save_image(make_grid(torch.tensor(fake_images)), output_dir / f"image_{epoch+1}.jpg")
-    t.set_description(f"Generator loss: {loss_g/n_steps}, Discriminator loss: {loss_d/n_steps}")
+        data_fake = generator.forward(noise)
+        loss_g += train_generator(optim_g, data_fake)
+      if (epoch + 1) % sample_interval == 0:
+        fake_images = generator.forward(ds_noise).detach().numpy()
+        fake_images = (fake_images.reshape(-1, 1, 28, 28) + 1) / 2  # 0 - 1 range.
+        save_image(make_grid(torch.tensor(fake_images)), output_dir / f"image_{epoch+1}.jpg")
+      t.set_description(f"Generator loss: {loss_g/n_steps}, Discriminator loss: {loss_d/n_steps}")
  print("Training Completed!")
@@ -31,7 +31,7 @@ def compile(onnx_file):
  for i in range(3):
    GlobalCounters.reset()
    print(f"run {i}")
-    with Context(DEBUG=max(DEBUG.value, 2 if i == 2 else 1)):
+    with Context(DEBUG=max(DEBUG.value, 2 if i == 2 else 1), OPENPILOT_HACKS=1):
      ret = run_onnx_jit(**inputs).numpy()
    # copy i == 1 so use of JITBEAM is okay
    if i == 1: test_val = np.copy(ret)
@@ -0,0 +1,16 @@
+import sys, pickle
+from extra.bench_log import WallTimeEvent, BenchEvent
+from tinygrad.helpers import getenv
+
+PKL = sys.argv[1] if len(sys.argv) > 1 else "/tmp/openpilot.pkl"
+
+load_times = []
+
+for _ in range(10):
+  with WallTimeEvent(BenchEvent.STEP) as wte: pickle.load(open(PKL, 'rb'))
+  load_times.append(wte.time)
+  print(f"pickle load: {wte.time:6.2f} s")
+
+if (assert_time:=getenv("ASSERT_MIN_LOAD_TIME")):
+  min_time = min(load_times)
+  assert min_time < assert_time, f"Speed regression, expected min load time of < {assert_time} s but took: {min_time} s"
@@ -65,7 +65,7 @@ def get_bar0_size(pcibus):
 class AMSMI(AMDev):
  def __init__(self, pcibus, vram_bar:MMIOInterface, doorbell_bar:MMIOInterface, mmio_bar:MMIOInterface):
    self.pcibus = pcibus
-    self.vram, self.doorbell64, self.mmio, self.dma_regions = vram_bar, doorbell_bar, mmio_bar, None
+    self.vram, self.doorbell64, self.mmio = vram_bar, doorbell_bar, mmio_bar
    self.pci_state = self.read_pci_state()
    if self.pci_state == "D0": self._init_from_d0()

@@ -2,7 +2,7 @@

 import os
 from tinygrad.helpers import Context
-from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase
+from tinygrad.runtime.support.system import System, PCIDevice
 from tinygrad.runtime.support.hcq import FileIOInterface
 from tinygrad.runtime.support.am.amdev import AMDev

@@ -12,7 +12,7 @@ if __name__ == "__main__":
    drv_path = f"/sys/bus/pci/devices/{gpu}/driver"
    if FileIOInterface.exists(drv_path) and os.path.basename(os.readlink(drv_path)) == "amdgpu":
      raise RuntimeError(f"amdgpu is bound to {gpu}. Stopping...")
-  pcidevs = [PCIDevice("AM", gpu, bars=[0, 2, 5]) for gpu in gpus]
+  pcidevs = [PCIDevice("AM", gpu) for gpu in gpus]
  amdevs = []
  with Context(DEBUG=2):
    for pcidev in pcidevs:
@@ -7,8 +7,8 @@ class GFXFake:
  def __init__(self): self.xccs = 8

 class AMDFake(AMDev):
-  def __init__(self, pci_dev, dma_regions=None):
-    self.pci_dev, self.devfmt, self.dma_regions = pci_dev, pci_dev.pcibus, dma_regions
+  def __init__(self, pci_dev):
+    self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
    self.vram, self.doorbell64, self.mmio = self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I')
    self._run_discovery()
    self._build_regs()
@@ -34,7 +34,8 @@ class WallTimeEvent:
    self.start = time.monotonic()
    return self
  def __exit__(self, *_):
-    _events[self.event]["wall"].append(time.monotonic() - self.start)
+    self.time = time.monotonic() - self.start
+    _events[self.event]["wall"].append(self.time)
    return False

 class KernelTimeEvent:
@@ -1,11 +1,11 @@
 from typing import Tuple, Dict, List, Optional
-from tinygrad.dtype import DType
+from tinygrad.dtype import DType, dtypes
 from tinygrad.renderer import ProgramSpec
-from tinygrad.tensor import Device, Tensor
+from tinygrad.tensor import Tensor
+from tinygrad.device import Device
 from tinygrad.engine.jit import TinyJit
 from tinygrad.nn.state import get_state_dict
 from tinygrad.helpers import Context, to_mv
-from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import Ops
 import json
 from collections import OrderedDict
@@ -13,12 +13,20 @@ from collections import OrderedDict
 EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CPU", "CUDA", "CL"]

 def compile_net(run:TinyJit, special_names:Dict[int,str]) -> Tuple[Dict[str,str],List[Tuple[str,List[str],List[int]]],Dict[str,Tuple[int,DType,int]],Dict[str,Tensor]]:
+  # memory-planned subbuffers can have multiple Buffer objects for the same memory region
+  canon, _seen = {}, {}
+  for ji in run.jit_cache:
+    for b in ji.bufs:
+      if b is not None: canon[id(b)] = _seen.setdefault((id(b.base._buf), b.offset, b.size, b.dtype), b)
+  special_names = {id(canon[k]): v for k, v in special_names.items() if k in canon}
+
  functions, bufs, bufs_to_save, statements, bufnum = {}, {}, {}, [], 0
  for ji in run.jit_cache:
    fxn: ProgramSpec = ji.prg.p
    functions[fxn.function_name] = fxn.src   # NOTE: this assumes all with the same name are the same
    cargs = []
    for i,arg in enumerate(ji.bufs):
+      arg = canon[id(arg)]
      key = id(arg)
      if key not in bufs:
        if key in special_names:
@@ -291,7 +291,7 @@ def build_kernel(N, arch='gfx1100'):
  # MAIN GEMM LOOP
  # ===========================================================================

-  NO_DS, NO_GLOBAL = getenv("NO_DS", 0), getenv("NO_GLOBAL", 0)
+  NO_ALU, NO_DS, NO_GLOBAL = getenv("NO_ALU", 0), getenv("NO_DS", 0), getenv("NO_GLOBAL", 0)

  k.label('LOOP_INC')
  k.emit(s_add_i32(s[S_LOOP_CTR], s[S_LOOP_CTR], 8))
@@ -350,10 +350,11 @@ def build_kernel(N, arch='gfx1100'):

    # 64 dual FMACs
    k.waitcnt(lgkm=0)
-    k.emit(s_clause(simm16=len(FMAC_PATTERN)-1))
-    for i, (vdst_x, vdst_y, ax, bx, ay, by) in enumerate(FMAC_PATTERN):
-      k.emit(VOPD(VOPDOp.V_DUAL_FMAC_F32, VOPDOp.V_DUAL_FMAC_F32,
-                  vdstx=v[vdst_x], vdsty=v[vdst_y], srcx0=v[ax], vsrcx1=v[bx], srcy0=v[ay], vsrcy1=v[by]))
+    if not NO_ALU:
+      k.emit(s_clause(simm16=len(FMAC_PATTERN)-1))
+      for i, (vdst_x, vdst_y, ax, bx, ay, by) in enumerate(FMAC_PATTERN):
+        k.emit(VOPD(VOPDOp.V_DUAL_FMAC_F32, VOPDOp.V_DUAL_FMAC_F32,
+                    vdstx=v[vdst_x], vdsty=v[vdst_y], srcx0=v[ax], vsrcx1=v[bx], srcy0=v[ay], vsrcy1=v[by]))

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

 N = getenv("N", 4096)
-M = K = N
-run_count = getenv("CNT", 5)
+M = getenv("M", N)
+K = getenv("K", N)
+NUM_RUNS = getenv("CNT", 5)

 # ---------------------------
 # launch/config constants
 # ---------------------------

 WARP_SIZE = 32
-
-# Threadblock tile sizes (block-level tile of C that a block computes)
-BLOCK_N = 128   # columns of C (N-dim) per block
-BLOCK_M = 128   # rows of C (M-dim) per block
-BLOCK_K = 8     # K-slice per block iteration
-
-# Register tile sizes (per-thread accumulator tile of C)
-TN = 4     # columns per thread
-TM = 4     # rows per thread
+BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 8
+TM, TN = 4, 4
+LANES_PER_WAVE_M, LANES_PER_WAVE_N = 4, 8
+assert N % BLOCK_N == 0 and M % BLOCK_M == 0 and K % BLOCK_K == 0

 is_kernel5 = getenv("K5", 0)
 THREADS_PER_BLOCK = 128 if is_kernel5 else 256
-assert THREADS_PER_BLOCK % BLOCK_N == 0, "THREADS_PER_BLOCK must be divisible by BLOCK_N"
-assert THREADS_PER_BLOCK % BLOCK_K == 0, "THREADS_PER_BLOCK must be divisible by BLOCK_K"
-assert (BLOCK_N * BLOCK_K) % THREADS_PER_BLOCK == 0
-assert (BLOCK_M * BLOCK_K) % THREADS_PER_BLOCK == 0
+WAVES_PER_BLOCK_N = 1 if is_kernel5 else 2
+WAVES_PER_BLOCK_M = THREADS_PER_BLOCK // WARP_SIZE // WAVES_PER_BLOCK_N
+REG_TILES_PER_WAVE_N = BLOCK_N // (WAVES_PER_BLOCK_N * LANES_PER_WAVE_N * TN)
+REG_TILES_PER_WAVE_M = BLOCK_M // (WAVES_PER_BLOCK_M * LANES_PER_WAVE_M * TM)

-WARPS_PER_BLOCK = THREADS_PER_BLOCK // WARP_SIZE
-WAVE_TILE_N = 128 if is_kernel5 else 64
-WAVE_TILE_M = BLOCK_N * BLOCK_M // WARPS_PER_BLOCK // WAVE_TILE_N
-assert BLOCK_N % WAVE_TILE_N == 0, "BN must be a multiple of WN"
-assert BLOCK_M % WAVE_TILE_M == 0, "BM must be a multiple of WM"
-WAVES_IN_BLOCK_X = BLOCK_N // WAVE_TILE_N
-WAVES_IN_BLOCK_Y = BLOCK_M // WAVE_TILE_M
-assert WAVES_IN_BLOCK_X * WAVES_IN_BLOCK_Y == WARPS_PER_BLOCK, "wave grid must match warps/block"
-
-LANES_PER_WAVE_X = 8
-LANES_PER_WAVE_Y = 4
-ITERS_PER_WAVE_N = WAVE_TILE_N // (LANES_PER_WAVE_X * TN)
-ITERS_PER_WAVE_M = WAVE_TILE_M // (LANES_PER_WAVE_Y * TM)
-assert WAVE_TILE_N % (LANES_PER_WAVE_X * TN) == 0, "WAVE_TILE_N must be divisible by LANES_PER_WAVE_X*TN"
-assert WAVE_TILE_M % (LANES_PER_WAVE_Y * TM) == 0, "WAVE_TILE_M must be divisible by LANES_PER_WAVE_Y*TM"
+assert WAVES_PER_BLOCK_M*REG_TILES_PER_WAVE_M*LANES_PER_WAVE_M*TM == BLOCK_M, "M reshape is wrong"
+assert WAVES_PER_BLOCK_N*REG_TILES_PER_WAVE_N*LANES_PER_WAVE_N*TN == BLOCK_N, "N reshape is wrong"

 def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.LOOP): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
-def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=False):
+def copy(dest:UOp, src:UOp, rng:int, upcast=False):
  assert dest.shape == src.shape
  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.LOOP)
-  copy = dest[*rngs].store(src[*rngs]).end(*rngs)
-  return dest.after(copy) if set else copy
+  return dest[*rngs].store(src[*rngs]).end(*rngs)

-def hand_spec_kernel3():
+def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
  # ---------------------------
-  # block indices & placeholders
+  # block indices
  # ---------------------------
-  blockIdx_x = UOp.special(N // BLOCK_N, "gidx0")
-  blockIdx_y = UOp.special(N // BLOCK_M, "gidx1")
-
-  a = UOp.placeholder((N, N), dtypes.float, slot=1)
-  b = UOp.placeholder((N, N), dtypes.float, slot=2)
-  c = UOp.placeholder((N, N), dtypes.float, slot=0)
+  block_id_n = UOp.special(N // BLOCK_N, "gidx0")
+  block_id_m = UOp.special(M // BLOCK_M, "gidx1")

  # index the output with the globals
-  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[blockIdx_y, :, blockIdx_x, :]
+  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[block_id_m, :, block_id_n, :]

  # open the main reduction range
-  k_tile_range = UOp.range(N // BLOCK_K, 0, AxisType.REDUCE)
-  a = a.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_K, BLOCK_K)[blockIdx_y, :, k_tile_range, :]
-  b = b.reshape(N // BLOCK_K, BLOCK_K, N // BLOCK_N, BLOCK_N)[k_tile_range, :, blockIdx_x, :]
+  k_tile_range = UOp.range(K // BLOCK_K, 0, AxisType.REDUCE)
+  a = a.reshape(M // BLOCK_M, BLOCK_M, K // BLOCK_K, BLOCK_K)[block_id_m, :, k_tile_range, :]
+  b = b.reshape(K // BLOCK_K, BLOCK_K, N // BLOCK_N, BLOCK_N)[k_tile_range, :, block_id_n, :]

  # globals are no longer used, they are already in the indexes
-  del blockIdx_y, blockIdx_x
+  del block_id_m, block_id_n

  # ---------------------------
-  # GLOBAL -> LOCAL (As, Bs)
+  # GLOBAL -> LOCAL (A_local, B_local)
  # ---------------------------
  tid = UOp.special(THREADS_PER_BLOCK, "lidx0")

  # A: read BM x BK tiles (permute on store into locals)
-  BM_As_stride = (BLOCK_M + 4) if is_kernel5 else BLOCK_M
-  As = UOp.placeholder((BLOCK_K, BM_As_stride), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL).shrink_to((BLOCK_K, BLOCK_M))
-  As_store = copy(As.permute((1,0)).reshape(-1, THREADS_PER_BLOCK)[:, tid], a.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=100)
+  BM_A_local_stride = (BLOCK_M + 4) if is_kernel5 else BLOCK_M
+  A_local = UOp.placeholder((BLOCK_K, BM_A_local_stride), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL).shrink_to((BLOCK_K, BLOCK_M))
+  A_local_store = copy(A_local.permute((1,0)).reshape(-1, THREADS_PER_BLOCK)[:, tid], a.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=100)

  # B: read BK x BN tiles
-  Bs = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
-  Bs_store = copy(Bs.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)
+  B_local = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
+  B_local_store = copy(B_local.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)

  # TODO: can we automate barrier?
-  barrier = UOp.barrier(As_store, Bs_store)
-  As, Bs = As.after(barrier), Bs.after(barrier)
+  barrier = UOp.barrier(A_local_store, B_local_store)
+  A_local, B_local = A_local.after(barrier), B_local.after(barrier)

  # open inner k range
  k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)
@@ -100,31 +76,30 @@ def hand_spec_kernel3():
  # ---------------------------
  # LOCAL -> REG (per-wave tiles)
  # ---------------------------
-  waveIdx = (tid // WARP_SIZE) % WAVES_IN_BLOCK_X
-  waveIdy = (tid // WARP_SIZE) // WAVES_IN_BLOCK_X
-  assert waveIdy.vmax+1 == WAVES_IN_BLOCK_Y
+  warp, lane = tid // WARP_SIZE, tid % WARP_SIZE
+  waveIdx, waveIdy = warp % WAVES_PER_BLOCK_N, warp // WAVES_PER_BLOCK_N
+  laneIdx, laneIdy = lane % LANES_PER_WAVE_N, lane // LANES_PER_WAVE_N
+  assert waveIdy.vmax+1 == WAVES_PER_BLOCK_M and laneIdy.vmax+1 == LANES_PER_WAVE_M

-  laneIdx = (tid % WARP_SIZE) % LANES_PER_WAVE_X
-  laneIdy = (tid % WARP_SIZE) // LANES_PER_WAVE_X
-  assert laneIdy.vmax+1 == LANES_PER_WAVE_Y
+  A_col = UOp.placeholder((REG_TILES_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+  A_local_slice = A_local[k, :].reshape(WAVES_PER_BLOCK_M, REG_TILES_PER_WAVE_M, LANES_PER_WAVE_M, TM)[waveIdy, :, laneIdy, :]
+  A_col = A_col.after(copy(A_col, A_local_slice, 300, upcast=True))

-  A_col = UOp.placeholder((ITERS_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
-  A_col = copy(A_col, As[k, :].reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM)[waveIdy, :, laneIdy, :], 300, set=True, upcast=True)
-
-  B_row = UOp.placeholder((ITERS_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
-  B_row = copy(B_row, Bs[k, :].reshape(WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)[waveIdx, :, laneIdx, :], 400, set=True, upcast=True)
+  B_row = UOp.placeholder((REG_TILES_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+  B_local_slice = B_local[k, :].reshape(WAVES_PER_BLOCK_N, REG_TILES_PER_WAVE_N, LANES_PER_WAVE_N, TN)[waveIdx, :, laneIdx, :]
+  B_row = B_row.after(copy(B_row, B_local_slice, 400, upcast=True))

  # ---------------------------
  # FMA: c_regs += A_col * B_row
  # ---------------------------
-  c_regs = UOp.placeholder((ITERS_PER_WAVE_M, TM, ITERS_PER_WAVE_N, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
+  c_regs = UOp.placeholder((REG_TILES_PER_WAVE_M, TM, REG_TILES_PER_WAVE_N, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
  i = UOp.range(c_regs.size, 16)
  c_regs = c_regs.after(c_regs.flatten()[i].store(0.0).end(i))

  # TODO: why don't these work as upcast?
  # why if the ranges merge is it slow?!? (if you change the order on end, they will merge. big slowdown on METAL)
-  iterWaveM, yt, iterWaveN, xt = rngs = rngs_for_shape(c_regs.shape, 500)
-  sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iterWaveM, yt] * B_row[iterWaveN, xt]).end(iterWaveM, iterWaveN, yt, xt)
+  iter_m, t_m, iter_n, t_n = rngs = rngs_for_shape(c_regs.shape, 500)
+  sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iter_m, t_m] * B_row[iter_n, t_n]).end(iter_m, iter_n, t_m, t_n)

  # Close k, sync, and close K tiles
  sink = sink.end(k).barrier().end(k_tile_range)
@@ -132,38 +107,37 @@ def hand_spec_kernel3():
  # ---------------------------
  # REG -> GLOBAL (epilogue)
  # ---------------------------
-  c = c.reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM,
-                WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)
+  c = c.reshape(WAVES_PER_BLOCK_M, REG_TILES_PER_WAVE_M, LANES_PER_WAVE_M, TM,
+                WAVES_PER_BLOCK_N, REG_TILES_PER_WAVE_N, LANES_PER_WAVE_N, TN)
  c = c[waveIdy, :, laneIdy, :,
        waveIdx, :, laneIdx, :]
  sink = copy(c, c_regs.after(sink), rng=600)

  return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()

-def test_matmul(sink:UOp, dtype=dtypes.float32, N=N):
-  rng = np.random.default_rng()
-  a = Tensor(rng.random((N, N), dtype=np.float32)-0.5, dtype=dtype)
-  b = Tensor(rng.random((N, N), dtype=np.float32)-0.5, dtype=dtype)
-  hc = Tensor.empty(N, N, dtype=dtype)
-  Tensor.realize(a, b, hc)
-
-  ei = ExecItem(sink, [t.uop.buffer for t in [hc, a, b]], prg=get_runner(Device.DEFAULT, sink))
+def eval_custom_matmul(fxn, dt=dtypes.float):
+  a = Tensor.randn(M, K, dtype=dt)
+  b = Tensor.randn(K, N, dtype=dt)
+  c = Tensor.empty(M, N, dtype=dtypes.float)
+  with Context(DEBUG=0): Tensor.realize(a, b)

  ets = []
-  with Context(DEBUG=2):
-    for _ in range(run_count):
-      ets.append(ei.run(wait=True))
-  print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")
+  with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2 if dt == dtypes.half else 0):
+    for _ in range(NUM_RUNS):
+      GlobalCounters.reset()
+      tst = Tensor.custom_kernel(c, a, b, fxn=fxn)[0].realize()
+      ets.append(GlobalCounters.time_sum_s)
+  print(f"REAL TFLOPS {M * N * K * 2 / min(ets) * 1e-12:.2f}")

  if getenv("VERIFY", 1):
    GlobalCounters.reset()
    with Context(DEBUG=2):
-      tc = (a @ b).realize()
+      tc = (a.float() @ b.float()).realize()
    with Context(DEBUG=0):
-      err = (hc - tc).square().mean().item()
+      err = (tc - tst).square().mean().item()
    print(f"mean squared error {err}")
-    if err > 1e-06:
+    if err > (1e-2 if dt == dtypes.half else 1e-6):
      raise RuntimeError("matmul is wrong!")

 if __name__ == "__main__":
-  test_matmul(hand_spec_kernel3(), N=N)
+  eval_custom_matmul(hand_spec_kernel3)
@@ -1,101 +0,0 @@
-import atexit, functools
-from tinygrad import Tensor, Device, dtypes
-from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
-from tinygrad.renderer import Estimates
-from tinygrad.helpers import getenv, all_same, dedup
-from extra.gemm.asm.cdna.asm import build_kernel, TILE_M, TILE_N, TILE_K, NUM_WG
-
-# ** CDNA4 assembly gemm
-
-WORKGROUP_SIZE = 256
-
-@functools.cache
-def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str, arch:str, wg:int) -> UOp:
-  batch, M, K = A.shape
-  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
-  assert K == K2
-  lidx = UOp.special(WORKGROUP_SIZE, "lidx0")
-  gidx = UOp.special(wg, "gidx0")
-  insts = build_kernel(batch, M, N, K, A.dtype.base)
-  lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=133_120, addrspace=AddrSpace.LOCAL), (), 'lds')
-  sink = UOp.sink(C.base, A.base, B.base, lds, lidx, gidx,
-                  arg=KernelInfo(name=f"gemm_{batch}_{M}_{N}_{K}", estimates=Estimates(ops=2*batch*M*N*K, mem=(batch*M*K + K*N + batch*M*N)*2)))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname),
-                                UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
-
-counters = {"used":0, "todos":[]}
-def todo(msg:str) -> bool: counters["todos"].append(msg); return False
-atexit.register(lambda: print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used'))
-
-def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
-  if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}")
-  if a.dtype not in {dtypes.bfloat16, dtypes.float16}: return todo(f"only bfloat16/float16, got {a.dtype}")
-  batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
-  N = b.shape[1]
-  # only sharding on the batch or K is tested, others might work too
-  if isinstance(a.device, tuple):
-    if a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
-    elif a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None: batch //= len(a.device)
-    else: return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
-    dname = a.device[0]
-  else: dname = a.device
-  arch = getattr(Device[dname].renderer, "arch", "")
-  if batch not in {1, 2}: return todo(f"GEMM batch size {batch}")
-  if (M % TILE_M != 0 or N % TILE_N != 0 or K % TILE_K != 0) and arch == "gfx950":
-    return todo(f"GEMM shape ({M},{N},{K}) not a multiple of ({TILE_M},{TILE_N},{TILE_K})")
-  return True
-
-# ** UOp gemm to test Tensor.custom_kernel multi and backward correctness on non cdna4
-# note: this can be removed after we have GEMM on mixins
-
-def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
-  M, K = A.shape[0]*A.shape[1], A.shape[2]
-  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
-  assert K == K2
-  m = UOp.range(M, 1, AxisType.LOOP)
-  n = UOp.range(N, 2, AxisType.LOOP)
-  k = UOp.range(K, 0, AxisType.REDUCE)
-  mul = (A.index((m*UOp.const(dtypes.index, K)+k))*B.index((k*UOp.const(dtypes.index, N)+n))).cast(dtypes.float32)
-  red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype.base)
-  store = C.index((m*UOp.const(dtypes.index, N)+n), ptr=True).store(red).end(m, n)
-  return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
-
-# ** backward gemm, might use the asm gemm
-
-def custom_gemm_bw(gradient:UOp, kernel:UOp):
-  out, a, b = kernel.src[1:]
-  assert all_same([gradient.device, a.device, b.device, out.device])
-  a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
-  # TODO: this needs to be cleaned up and done properly, the batch dim of grad and a multi need to align
-  g_t = g_t[:a.shape[0]]
-  grad_a = (g_t @ b_t.T).uop
-  grad_b = (a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1) @ g_t.reshape(-1, g_t.shape[-1])).uop
-  return (None, grad_a, grad_b)
-
-# ** main gemm function
-
-def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
-  assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
-  counters["used"] += 1
-  squeeze = a.ndim == 2
-  if squeeze: a = a.unsqueeze(0)
-
-  batch, M, K = a.shape
-  N = b.shape[1]
-  is_multi = isinstance(a.device, tuple)
-  if (k_sharded:=is_multi and a.uop.axis == 2): K //= len(a.device)
-
-  if is_multi:
-    out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
-  else:
-    out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)
-
-  renderer = Device[a.device[0] if is_multi else a.device].renderer
-  dname, arch = renderer.device, getattr(renderer, "arch", "")
-  if arch.startswith("gfx950") and getenv("USE_ASM", 1):
-    out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname, wg=NUM_WG, arch=arch), grad_fxn=custom_gemm_bw)[0]
-  else:
-    out = Tensor.custom_kernel(out, a, b, fxn=custom_uop_gemm, grad_fxn=custom_gemm_bw)[0]
-  if k_sharded: out = out.sum(0)
-  return out.squeeze(0) if squeeze else out
@@ -1,76 +0,0 @@
-.text
-.section	.text.
-.global	gemm
-.p2align	8
-.type	gemm,@function
-
-gemm:
-INSTRUCTIONS
-
-.section .rodata,"a",@progbits
-.p2align 6, 0x0
-.amdhsa_kernel gemm
-  # basic memory requirements
-  .amdhsa_group_segment_fixed_size 30336
-  .amdhsa_private_segment_fixed_size 0
-  .amdhsa_kernarg_size 32
-  # register usage (RSRC1)
-  .amdhsa_next_free_vgpr 256
-  .amdhsa_next_free_sgpr 100
-  # workgroup / workitem IDs (RSRC2)
-  .amdhsa_system_sgpr_workgroup_id_x 1
-  .amdhsa_system_sgpr_workgroup_id_y 1
-  .amdhsa_system_sgpr_workgroup_id_z 1
-  # user SGPRs: kernarg ptr in s[0:1]
-  .amdhsa_user_sgpr_kernarg_segment_ptr 1
-  .amdhsa_user_sgpr_count 2
-  # gfx10+ / gfx11 specifics (RSRC1[29..31])
-  .amdhsa_wavefront_size32 1
-  .amdhsa_workgroup_processor_mode 1
-  .amdhsa_memory_ordered 1
-  .amdhsa_forward_progress 1
-  # misc for gfx11
-  .amdhsa_dx10_clamp 1
-  .amdhsa_ieee_mode 1
-  .amdhsa_uses_dynamic_stack 0
-.end_amdhsa_kernel
-
-.amdgpu_metadata
---
-amdhsa.kernels:
-  - .args:
-      - .address_space:  generic
-        .name:           C
-        .offset:         0
-        .size:           8
-        .value_kind:     global_buffer
-        .value_type:     f16
-      - .address_space:  generic
-        .name:           A
-        .offset:         8
-        .size:           8
-        .value_kind:     global_buffer
-        .value_type:     f16
-      - .address_space:  generic
-        .name:           B
-        .offset:         16
-        .size:           8
-        .value_kind:     global_buffer
-        .value_type:     f16
-    .group_segment_fixed_size: 30336
-    .kernarg_segment_align: 8
-    .kernarg_segment_size: 32
-    .max_flat_workgroup_size: 128
-    .name:           gemm
-    .private_segment_fixed_size: 0
-    .sgpr_count:     70
-    .sgpr_spill_count: 0
-    .symbol:         gemm.kd
-    .vgpr_count:     256
-    .vgpr_spill_count: 0
-    .wavefront_size: 32
-amdhsa.version:
-  - 1
-  - 1
-...
-.end_amdgpu_metadata
@@ -1,30 +0,0 @@
-import math, pathlib
-
-from tinygrad import Device, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-
-from extra.gemm.amd_uop_matmul import test_matmul
-
-N = 4096
-TN = 96
-THREADS_PER_WG = 128
-NUM_WG = math.ceil(N / TN) * math.ceil(N / TN)
-
-dname:str = Device.DEFAULT
-template:str = (pathlib.Path(__file__).parent/"template.s").read_text()
-
-def asm_kernel() -> UOp:
-  lidx = UOp.special(THREADS_PER_WG, "lidx0")
-  gidx = UOp.special(NUM_WG, "gidx0")
-
-  a = UOp.placeholder((N*N,), dtypes.half, slot=1)
-  b = UOp.placeholder((N*N,), dtypes.half, slot=2)
-  c = UOp.placeholder((N*N,), dtypes.half, slot=0)
-
-  src = template.replace("INSTRUCTIONS", (pathlib.Path(__file__).parent/"gemm.s").read_text())
-
-  sink = UOp.sink(a, b, c, lidx, gidx, arg=KernelInfo(name="gemm"))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src)))
-
-if __name__ == "__main__":
-  test_matmul(asm_kernel(), dtype=dtypes.half, N=N)
@@ -1,179 +0,0 @@
-# unpack the complete kernel descriptor of an amdgpu ELF
-# https://rocm.docs.amd.com/projects/llvm-project/en/latest/LLVM/llvm/html/AMDGPUUsage.html#code-object-v3-kernel-descriptor
-import struct, pathlib, sys
-from tinygrad.runtime.support.elf import elf_loader
-
-def bits(x, lo, hi): return (x >> lo) & ((1 << (hi - lo + 1)) - 1)
-def assert_zero(x, lo, hi): assert bits(x, lo, hi) == 0
-
-with open(sys.argv[1], "rb") as f:
-  lib = f.read()
-
-image, sections, relocs = elf_loader(lib)
-rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"))
-
-# rodata is exactly 64 bytes
-kd = image[rodata_entry:rodata_entry+64]
-desc = int.from_bytes(kd, byteorder="little")
-
-group_segment_fixed_size = bits(desc, 0, 31)
-private_segment_fixed_size = bits(desc, 32, 63)
-kernarg_size = bits(desc, 64, 95)
-reserved_127_96 = bits(desc, 96, 127)
-assert reserved_127_96 == 0
-
-print("GROUP_SEGMENT_FIXED_SIZE:", group_segment_fixed_size)
-print("PRIVATE_SEGMENT_FIXED_SIZE:", private_segment_fixed_size)
-print("KERNARG_SIZE:", kernarg_size)
-print("RESERVED 127:96:", reserved_127_96)
-
-entry_off = bits(desc, 128, 191)
-
-# sign-extend manually if needed
-if entry_off & (1 << 63):
-  entry_off -= 1 << 64
-
-print("KERNEL_CODE_ENTRY_BYTE_OFFSET:", entry_off)
-
-kd_addr = 0x1840
-entry_addr = kd_addr + entry_off
-
-print("Computed entry address: 0x%016x" % entry_addr)
-print("256B aligned:", entry_addr % 256 == 0)
-
-pgm_rsrc3 = bits(desc, 352, 383)
-pgm_rsrc1 = bits(desc, 384, 415)
-pgm_rsrc2 = bits(desc, 416, 447)
-
-print("COMPUTE_PGM_RSRC3: 0x%08x" % pgm_rsrc3)
-print("COMPUTE_PGM_RSRC1: 0x%08x" % pgm_rsrc1)
-print("COMPUTE_PGM_RSRC2: 0x%08x" % pgm_rsrc2)
-
-# rsrc 3 (gfx950)
-
-accum_offset_raw = bits(pgm_rsrc3, 0, 5)
-assert_zero(pgm_rsrc3, 6, 15)
-tg_split = bits(pgm_rsrc3, 16, 16)
-accum_offset_vgprs = (accum_offset_raw + 1) * 4
-print("RSRC3.ACCUM_OFFSET (AccVGPR index):", accum_offset_vgprs)
-print("RSRC3.TG_SPLIT:", tg_split)
-
-# rsrc 1
-
-vgpr_gran = bits(pgm_rsrc1, 0, 5)
-sgpr_gran = bits(pgm_rsrc1, 6, 9)
-assert_zero(pgm_rsrc1, 27, 28)
-
-# NOTE: this is vgprs + agprs
-vgprs_used = (vgpr_gran + 1) * 8
-assert 0 <= vgprs_used <= 512
-
-k = sgpr_gran // 2
-sgprs_used = (k + 1) * 16
-
-print("RSRC1.VGPRS:", vgprs_used)
-print("RSRC1.SGPRS:", sgprs_used)
-
-assert_zero(pgm_rsrc1, 10, 11)
-
-float_round_mode_32 = bits(pgm_rsrc1, 12, 13)
-float_round_mode_16_64 = bits(pgm_rsrc1, 15, 14)
-float_denorm_mode_32 = bits(pgm_rsrc1, 16, 17)
-float_denorm_mode_16_64 = bits(pgm_rsrc1, 18, 19)
-
-priv = bits(pgm_rsrc1, 20, 20)
-assert priv == 0
-enable_dx10_clamp_wg_rr_en = bits(pgm_rsrc1, 21, 21)
-debug_mode = bits(pgm_rsrc1, 22, 22)
-enable_ieee_mode = bits(pgm_rsrc1, 23, 23)
-bulky = bits(pgm_rsrc1, 24, 24)
-assert bulky == 0
-cdbg_user = bits(pgm_rsrc1, 25, 25)
-assert cdbg_user == 0
-fp16_ovfl = bits(pgm_rsrc1, 26, 26)
-assert_zero(pgm_rsrc1, 27, 28)  # reserved
-assert_zero(pgm_rsrc1, 29, 29)  # WGP_MODE (reserved on gfx9)
-assert_zero(pgm_rsrc1, 30, 30)  # MEM_ORDERED (reserved on gfx9)
-assert_zero(pgm_rsrc1, 31, 31)  # FWD_PROGRESS (reserved on gfx9)
-
-# rsrc 2
-
-enable_private_segment = bits(pgm_rsrc2, 0, 0)  # SCRATCH_EN
-user_sgpr_count = bits(pgm_rsrc2, 1, 5)         # USER_SGPR
-enable_trap_handler = bits(pgm_rsrc2, 6, 6)     # TRAP_PRESENT (must be 0 here)
-assert enable_trap_handler == 0
-
-enable_sgpr_workgroup_id_x = bits(pgm_rsrc2, 7, 7)
-enable_sgpr_workgroup_id_y = bits(pgm_rsrc2, 8, 8)
-enable_sgpr_workgroup_id_z = bits(pgm_rsrc2, 9, 9)
-enable_sgpr_workgroup_info = bits(pgm_rsrc2, 10, 10)
-
-enable_vgpr_workitem_id = bits(pgm_rsrc2, 11, 12)  # TIDIG_CMP_CNT enum (0..3)
-
-enable_exception_address_watch = bits(pgm_rsrc2, 13, 13)
-assert enable_exception_address_watch == 0
-enable_exception_memory = bits(pgm_rsrc2, 14, 14)
-assert enable_exception_memory == 0
-
-granulated_lds_size = bits(pgm_rsrc2, 15, 23)
-assert granulated_lds_size == 0  # spec: must be 0; CP uses dispatch packet rounding
-
-enable_exception_fp_invalid = bits(pgm_rsrc2, 24, 24)
-enable_exception_fp_denorm_src = bits(pgm_rsrc2, 25, 25)
-enable_exception_fp_div0 = bits(pgm_rsrc2, 26, 26)
-enable_exception_fp_overflow = bits(pgm_rsrc2, 27, 27)
-enable_exception_fp_underflow = bits(pgm_rsrc2, 28, 28)
-enable_exception_fp_inexact = bits(pgm_rsrc2, 29, 29)
-enable_exception_int_div0 = bits(pgm_rsrc2, 30, 30)
-
-assert_zero(pgm_rsrc2, 31, 31)
-
-print("RSRC2.ENABLE_PRIVATE_SEGMENT:", enable_private_segment)
-print("RSRC2.USER_SGPR_COUNT:", user_sgpr_count)
-print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_X:", enable_sgpr_workgroup_id_x)
-print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_Y:", enable_sgpr_workgroup_id_y)
-print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_Z:", enable_sgpr_workgroup_id_z)
-print("RSRC2.ENABLE_SGPR_WORKGROUP_INFO:", enable_sgpr_workgroup_info)
-print("RSRC2.ENABLE_VGPR_WORKITEM_ID (enum):", enable_vgpr_workitem_id)
-
-print("RSRC2.EXC_FP_INVALID:", enable_exception_fp_invalid)
-print("RSRC2.EXC_FP_DENORM_SRC:", enable_exception_fp_denorm_src)
-print("RSRC2.EXC_FP_DIV0:", enable_exception_fp_div0)
-print("RSRC2.EXC_FP_OVERFLOW:", enable_exception_fp_overflow)
-print("RSRC2.EXC_FP_UNDERFLOW:", enable_exception_fp_underflow)
-print("RSRC2.EXC_FP_INEXACT:", enable_exception_fp_inexact)
-print("RSRC2.EXC_INT_DIV0:", enable_exception_int_div0)
-
-# user sgprs
-
-enable_sgpr_private_segment_buffer = bits(desc, 448, 448)
-enable_sgpr_dispatch_ptr = bits(desc, 449, 449)
-enable_sgpr_queue_ptr = bits(desc, 450, 450)
-enable_sgpr_kernarg_segment_ptr = bits(desc, 451, 451)
-enable_sgpr_dispatch_id = bits(desc, 452, 452)
-enable_sgpr_flat_scratch_init = bits(desc, 453, 453)
-enable_sgpr_private_segment_size = bits(desc, 454, 454)
-
-assert_zero(desc, 455, 457)
-
-print("DESC.ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER:", enable_sgpr_private_segment_buffer)
-print("DESC.ENABLE_SGPR_DISPATCH_PTR:", enable_sgpr_dispatch_ptr)
-print("DESC.ENABLE_SGPR_QUEUE_PTR:", enable_sgpr_queue_ptr)
-print("DESC.ENABLE_SGPR_KERNARG_SEGMENT_PTR:", enable_sgpr_kernarg_segment_ptr)
-print("DESC.ENABLE_SGPR_DISPATCH_ID:", enable_sgpr_dispatch_id)
-print("DESC.ENABLE_SGPR_FLAT_SCRATCH_INIT:", enable_sgpr_flat_scratch_init)
-print("DESC.ENABLE_SGPR_PRIVATE_SEGMENT_SIZE:", enable_sgpr_private_segment_size)
-
-assert_zero(desc, 458, 459)
-
-uses_dynamic_stack = bits(desc, 459, 460)
-print("DESC.USES_DYNAMIC_STACK:", uses_dynamic_stack)
-
-# gfx950 only
-assert_zero(desc, 460, 463)
-kernarg_preload_spec_length = bits(desc, 464, 470)
-print("DESC.KERNARG_PRELOAD_SPEC_LENGTH:", kernarg_preload_spec_length)
-kernarg_preload_spec_offset = bits(desc, 471, 479)
-print("DESC.KERNARG_PRELOAD_SPEC_OFFSET:", kernarg_preload_spec_offset)
-
-assert_zero(desc, 480, 511)
@@ -10,9 +10,9 @@ HEVC_ROUNDUP = getenv("DATA_ROUNDUP", 32)
@functools.cache
 def _hevc_jitted_decoder(out_image_size:tuple[int, int], max_hist:int, inplace:bool):
  def hevc_decode_frame(pos:Variable, hevc_tensor:Tensor, offset:Variable, sz:Variable, opaque:Tensor, i:Variable, *hist:Tensor, outbuf:Tensor|None=None):
-    x = hevc_tensor[offset:offset+sz*HEVC_ROUNDUP].decode_hevc_frame(pos, out_image_size, opaque[i], hist)
+    x = hevc_tensor[offset:offset+sz*HEVC_ROUNDUP].decode_hevc_frame(pos, out_image_size, opaque[i], hist).realize()
    if outbuf is not None: outbuf.assign(x).realize()
-    return x.realize()
+    return x
  return TinyJit(hevc_decode_frame)

 def hevc_decode(hevc_tensor:Tensor, opaque:Tensor, frame_info:list, luma_h:int, luma_w:int,
@@ -74,10 +74,14 @@ if __name__ == "__main__":
  Device.default.synchronize()

  # decode all frames using the iterator
-  with Timing("decoding whole file: ", on_exit=(lambda et: f", {len(frame_info)} frames, {len(frame_info)/(et/1e9):.2f} fps")):
+  tm = Timing("decoding whole file: ", on_exit=(lambda et: f", {len(frame_info)} frames, {len(frame_info)/(et/1e9):.2f} fps"))
+  with tm:
    images = list(hevc_decode(hevc_tensor, opaque_nv, frame_info, luma_h, luma_w, history=hist, preallocated_outputs=out_images))
    Device.default.synchronize()

+  fps = len(frame_info)/(tm.et/1e9)
+  assert fps >= getenv("ASSERT_FPS", 0), f"HEVC decode too slow: {fps:.2f} fps"
+
  # validation
  if getenv("VALIDATE", 0):
    import pickle
@@ -56,7 +56,10 @@ class Attention:
  def __call__(self, x:Tensor, start_pos:Union[Variable,int], freqs_cis:Tensor, mask:Optional[Tensor]=None) -> Tensor:
    if getenv("WQKV"):
      xqkv = self.wqkv(x)
-      xq, xk, xv = xqkv.split([self.n_heads * self.head_dim, self.n_kv_heads * self.head_dim, self.n_kv_heads * self.head_dim], dim=2)
+      xqkv = xqkv.reshape(xqkv.shape[0], xqkv.shape[1], self.n_kv_heads, self.n_rep + 2, self.head_dim)
+      xq = xqkv[:, :, :, :self.n_rep].reshape(xqkv.shape[0], xqkv.shape[1], -1)
+      xk = xqkv[:, :, :, self.n_rep:self.n_rep+1].reshape(xqkv.shape[0], xqkv.shape[1], -1)
+      xv = xqkv[:, :, :, self.n_rep+1:self.n_rep+2].reshape(xqkv.shape[0], xqkv.shape[1], -1)
    else:
      xq, xk, xv = self.wq(x), self.wk(x.contiguous_backward()), self.wv(x)

@@ -0,0 +1,46 @@
+#!/usr/bin/env python3
+import os, sys, time
+from tinygrad.runtime.support.system import RemotePCIDevice
+
+LAT_N_RUNS = 500
+THROUGHPUT_N_RUNS = 8
+SIZES = [4, 1 << 10, 8 << 20]
+
+if __name__ == "__main__":
+  os.environ["REMOTE"] = sys.argv[1] if len(sys.argv) > 1 else os.environ.get("REMOTE", "127.0.0.1:6667")
+
+  # choose any amd/nv gpu.
+  devs = RemotePCIDevice.remote_list(0x1002, ((0, (0,)),), 0) or RemotePCIDevice.remote_list(0x10de, ((0, (0,)),), 0x03)
+  if not devs: raise RuntimeError("no GPU found on remote")
+
+  pci = RemotePCIDevice("BN", devs[0])
+  print(f"connected to {os.environ['REMOTE']}, device: {devs[0]}\n")
+
+  # ping (minimal server round-trip, no device I/O)
+  from tinygrad.runtime.support.system import RemoteCmd
+  sock = pci.sock
+  for _ in range(10): RemotePCIDevice._rpc(sock, 0, RemoteCmd.PING)
+  st = time.perf_counter()
+  for _ in range(LAT_N_RUNS): RemotePCIDevice._rpc(sock, 0, RemoteCmd.PING)
+  ping_lat = (time.perf_counter() - st) / LAT_N_RUNS
+  print(f"PING latency: {ping_lat*1e6:.1f} us ({1/ping_lat:,.0f} ops/sec)\n")
+
+  # throughput
+  sysmem, _ = pci.alloc_sysmem(max(SIZES))
+  print(f"{'size':>10s}  {'write MB/s':>10s}  {'read MB/s':>10s}")
+  for sz in SIZES:
+    data = b'\x01' * sz
+
+    for _ in range(5): sysmem[0:sz] = data
+    st = time.perf_counter()
+    for _ in range(THROUGHPUT_N_RUNS): sysmem[0:sz] = data
+    pci.read_config(0, 4) # flush, since writes are posted
+    w = (time.perf_counter() - st) / THROUGHPUT_N_RUNS
+
+    for _ in range(5): sysmem[0:sz]
+    st = time.perf_counter()
+    for _ in range(THROUGHPUT_N_RUNS): sysmem[0:sz]
+    r = (time.perf_counter() - st) / THROUGHPUT_N_RUNS
+
+    sfx, div = [('B',1),('K',1<<10),('M',1<<20)][[sz>=1<<10,sz>=1<<20,sz>=1<<30].count(True)]
+    print(f"{sz/div:>9.4g}{sfx}  {sz/w/1e6:>10.1f}  {sz/r/1e6:>10.1f}")
@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+import socket, struct, sys
+from tinygrad.runtime.support.system import PCIDevice, RemoteCmd, System
+from tinygrad.helpers import DEBUG, OSX
+
+def resp(resp0=0, resp1=0, status=0): return struct.pack('<BQQ', status, resp0, resp1)
+def resp_err(msg): return struct.pack('<BQQ', 1, len(err:=msg.encode()), 0) + err
+
+discovered_devices: list[str] = []
+opened_devices: dict[int, PCIDevice] = {}
+mapped_bars: dict[tuple[int, int], object] = {}
+sysmem_allocs: list[tuple] = []
+
+def handle(conn, cmd, dev_id, bar, arg0, arg1, arg2):
+  if cmd == RemoteCmd.PING:
+    return conn.sendall(resp())
+
+  if cmd == RemoteCmd.PROBE:
+    payload = conn.recv(arg1, socket.MSG_WAITALL) if arg1 > 0 else b""
+    filter_devices: dict[int, list[int]] = {}
+    for i in range(0, len(payload), 8):
+      mask, dev = struct.unpack('<II', payload[i:i+8])
+      filter_devices.setdefault(mask, []).append(dev)
+    base_class = None if arg0 == 0 else int(arg0)
+    devs = System.list_devices(arg2, tuple([(x, tuple(y)) for x,y in filter_devices.items()]), base_class)
+    for p in devs:
+      if p not in discovered_devices: discovered_devices.append(p)
+    data = "\n".join(f"{p[1]}:{discovered_devices.index(p)}" for p in devs).encode()
+    return conn.sendall(resp(len(data), len(devs)) + data)
+
+  # lazy device open
+  if dev_id not in opened_devices:
+    if dev_id >= len(discovered_devices): raise RuntimeError(f"device {dev_id} not probed")
+    cl, pcibus = discovered_devices[dev_id]
+    opened_devices[dev_id] = cl("SV", pcibus)
+  pci_dev = opened_devices[dev_id]
+
+  if cmd == RemoteCmd.MAP_BAR:
+    if (dev_id, bar) not in mapped_bars: mapped_bars[(dev_id, bar)] = pci_dev.map_bar(bar)
+    conn.sendall(resp(*pci_dev.bar_info(bar)))
+  elif cmd == RemoteCmd.CFG_READ:
+    conn.sendall(resp(pci_dev.read_config(arg0, arg1)))
+  elif cmd == RemoteCmd.CFG_WRITE:
+    pci_dev.write_config(arg0, arg2, arg1)
+    conn.sendall(resp())
+  elif cmd == RemoteCmd.RESIZE_BAR:
+    pci_dev.resize_bar(bar)
+    conn.sendall(resp())
+  elif cmd == RemoteCmd.RESET:
+    pci_dev.reset()
+    conn.sendall(resp())
+  elif cmd == RemoteCmd.MMIO_READ:
+    bar_view = mapped_bars[(dev_id, bar)]
+    if arg0 % 4 == 0 and arg1 == 4: conn.sendmsg([resp(arg1), struct.pack(f'<{arg1 // 4}I', bar_view.view(arg0, arg1, fmt='I')[0])])
+    else: conn.sendmsg([resp(arg1), bar_view[arg0:arg0+arg1]])
+  elif cmd == RemoteCmd.MMIO_WRITE:
+    data = conn.recv(arg1, socket.MSG_WAITALL)
+    bar_view = mapped_bars[(dev_id, bar)]
+    if arg0 % 4 == 0 and arg1 == 4: bar_view.view(arg0, arg1, fmt='I')[0] = struct.unpack(f'<{arg1 // 4}I', data)[0]
+    else: bar_view[arg0:arg0+arg1] = data
+  elif cmd == RemoteCmd.MAP_SYSMEM:
+    memview, paddrs = pci_dev.alloc_sysmem(arg0, contiguous=bool(arg1))
+    sysmem_allocs.append((memview, paddrs))
+    paddrs_bytes = struct.pack(f'<{len(paddrs)}Q', *paddrs)
+    conn.sendall(resp(len(paddrs_bytes), len(sysmem_allocs) - 1) + paddrs_bytes)
+  elif cmd == RemoteCmd.SYSMEM_READ:
+    conn.sendmsg([resp(arg1), sysmem_allocs[bar][0][arg0:arg0+arg1]])
+  elif cmd == RemoteCmd.SYSMEM_WRITE:
+    sysmem_allocs[bar][0][arg0:arg0+arg1] = conn.recv(arg1, socket.MSG_WAITALL)
+  else: raise RuntimeError(f"unknown command {cmd}")
+
+def serve(conn:socket.socket):
+  REQ = '<BIIQQQ'
+  while True:
+    hdr = conn.recv(struct.calcsize(REQ), socket.MSG_WAITALL)
+    if len(hdr) < struct.calcsize(REQ): raise ConnectionError("client disconnected")
+    cmd, dev_id, bar, arg0, arg1, arg2 = struct.unpack(REQ, hdr)
+    if DEBUG >= 4: print(f"cmd={RemoteCmd(cmd).name} dev={dev_id} bar={bar} arg0={arg0:#x} arg1={arg1:#x} arg2={arg2:#x}")
+    try: handle(conn, cmd, dev_id, bar, arg0, arg1, arg2)
+    except ConnectionError: raise
+    except Exception as e:
+      if cmd in {RemoteCmd.MMIO_WRITE, RemoteCmd.SYSMEM_WRITE}: raise ConnectionError(f"write failed: {e}")
+      print(f"ERROR: {e}")
+      conn.sendall(resp_err(str(e)))
+
+if __name__ == "__main__":
+  if not OSX: System.reserve_hugepages(128) # for sysmem allocations
+
+  port = int(sys.argv[1]) if len(sys.argv) > 1 else 6667
+  server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+  server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+  server.bind(("0.0.0.0", port))
+  server.listen(1)
+  s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+  try: s.connect(("8.8.8.8", 80)); ip = s.getsockname()[0]
+  finally: s.close()
+  print(f"listening on {ip}:{port}")
+  while True:
+    conn, addr = server.accept()
+    conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+    for bt in [socket.SO_SNDBUF, socket.SO_RCVBUF]: conn.setsockopt(socket.SOL_SOCKET, bt, 64 << 20)
+    try: serve(conn)
+    except ConnectionError: print("disconnected")
@@ -0,0 +1,31 @@
+#!/bin/sh
+install_loc="$HOME/.local/bin"
+docker build --platform=linux/amd64 -t rocm-hipcc:7.2 - <<'EOF'
+FROM ubuntu:22.04
+ENV DEBIAN_FRONTEND=noninteractive
+ENV TZ=Etc/UTC
+RUN apt-get update && apt-get install -y --no-install-recommends wget ca-certificates gnupg tzdata && \
+  wget https://repo.radeon.com/amdgpu-install/7.2/ubuntu/jammy/amdgpu-install_7.2.70200-1_all.deb && \
+  apt-get install -y ./amdgpu-install_7.2.70200-1_all.deb && \
+  amdgpu-install -y --usecase=rocm --no-dkms --no-32 && \
+  rm -rf /var/lib/apt/lists/*
+ENV PATH=/opt/rocm/bin:$PATH
+EOF
+
+mkdir -p "$install_loc"
+tee "$install_loc/hipccshim" >/dev/null <<'EOF'
+#!/bin/sh
+set -eu
+cname="rocm-hipcc-persistent"
+if ! docker inspect --format='{{.State.Running}}' "$cname" 2>/dev/null | grep -q true; then
+  docker rm -f "$cname" 2>/dev/null || true
+  docker run -d --platform=linux/amd64 --name "$cname" \
+    -v /var/folders:/var/folders -v "$HOME":"$HOME" \
+    rocm-hipcc:7.2 sleep 300 >/dev/null
+fi
+exec docker exec "$cname" "$(basename "$0")" "$@"
+EOF
+chmod +x "$install_loc/hipccshim"
+for t in hipcc hipconfig; do
+  ln -sf "$install_loc/hipccshim" "$install_loc/$t"
+done
@@ -1,9 +1,9 @@
 #!/bin/sh
 install_loc="$HOME/.local/bin"
-docker build --platform=linux/amd64 -t cuda-nvcc:12.8 - <<'EOF'
+docker build --platform=linux/arm64 -t cuda-nvcc:12.8 - <<'EOF'
 FROM ubuntu:22.04
 RUN apt-get update && apt-get install -y --no-install-recommends wget ca-certificates && \
-  wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb && \
+  wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-keyring_1.1-1_all.deb && \
  dpkg -i cuda-keyring_1.1-1_all.deb && \
  apt-get update && apt-get install -y --no-install-recommends cuda-nvcc-12-8 cuda-nvdisasm-12-8 cuda-cuobjdump-12-8 && rm -rf /var/lib/apt/lists/*
 ENV PATH=/usr/local/cuda/bin:$PATH
@@ -13,10 +13,14 @@ mkdir -p "$install_loc"
 tee "$install_loc/nvccshim" >/dev/null <<'EOF'
 #!/bin/sh
 set -eu
-# assume the final arg is the input path
-# mount it so that container can read it
-dir=$(dirname "${@: -1}")
-exec docker run --rm --platform=linux/amd64 -v "$dir":"$dir" cuda-nvcc:12.8 "$(basename "$0")" "$@"
+cname="cuda-nvcc-persistent"
+if ! docker inspect --format='{{.State.Running}}' "$cname" 2>/dev/null | grep -q true; then
+  docker rm -f "$cname" 2>/dev/null || true
+  docker run -d --platform=linux/arm64 --name "$cname" \
+    -v /var/folders:/var/folders -v "$HOME":"$HOME" \
+    cuda-nvcc:12.8 sleep 300 >/dev/null
+fi
+exec docker exec "$cname" "$(basename "$0")" "$@"
 EOF
 chmod +x "$install_loc/nvccshim"
 for t in nvcc nvdisasm; do
@@ -2,30 +2,13 @@

 ## Getting SQ Thread Trace

-SQTT is implemented on top of normal tinygrad profiling, `VIZ=1 SQTT=1` to get profile pickle with sqtt data embedded in it.
+`VIZ=2` to enable SQTT profiling.
+
+`SQTT_ITRACE_SE_MASK=X` to select shader engines for instruction tracing, -1 = all, 0 = disabled, >0 = SE bitmask, default 0b11.

 `SQTT_BUFFER_SIZE=X` to change size of SQTT buffer (per shader engine, 6 SEs on 7900xtx) in megabytes, default 256.

-`SQTT_ITRACE_SE_MASK=X` to select for which shader engines instruction tracing will be enabled, -1 is all, 0 is none (instruction tracing disabled), >0 is
-bitfield/mask for SEs to enable instruction tracing on. Masking shader engines will give smaller file sizes at a cost of less hits and kernels that
-don't have any wavefront on first simd of shader engine with instruction tracing enabled will not have instruction timings.
-The default is 2 (second shader engine only), only one for file size reasons, second instead of first because dispatch starts from it so there is
-greater chance that kernels with small global size will have instruction tracing data.
- 
-Note that instruction tracing might not be available for kernels with small global dims, this is not a bug, but it can be improved with various hacks
-to the point where it can reliably trace a kernel consisting of a single wavefront (am only, not quite reliable under amdgpu due to waves sometimes
-being dispatched starting from different simds). More info in comments in ops_amd.py
+## Viewing the traces

-## Converting pickled profile with SQTT data into RGP file
-
-```bash
-extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
-```
-
-Then load gpu0.rgp into Radeon GPU Profiler. It works just fine both in wine (macos, native version available for linux) and via ssh X forwarding
-
-If multiple gpus are used you can select which one to export with `-d` like this:
-
-```bash
-extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -d 'AMD:5' -o /tmp/gpu5.rgp
-```
+- Web UI: `tinygrad/viz/serve.py`
+- Command line: `python -m tinygrad.renderer.amd.sqtt`
@@ -1,152 +0,0 @@
-import os
-os.environ["PYTHONPATH"] = "."
-os.environ["SQTT"] = "1"
-if "DEV" not in os.environ: os.environ["DEV"] = "AMD"
-os.environ["PROFILE"] = "1"
-os.environ["AMD_LLVM"] = "0"
-
-from dataclasses import replace
-import atexit, contextlib
-from tinygrad import Tensor
-from tinygrad.helpers import system, OSX
-from tinygrad.runtime.ops_amd import AMDProgram
-from extra.sqtt.roc import decode, WaveExec, ProfileSQTTEvent
-from tinygrad.device import Device
-
-from extra.sqtt.attempt_sqtt_parse import parse_sqtt_print_packets
-
-dev = Device["AMD"]
-
-@contextlib.contextmanager
-def save_sqtt():
-  # clear the old traces
-  dev.profile_events.clear()
-  sqtt:dict[str, list[WaveExec]] = {}
-  yield sqtt
-  events = dev.profile_events
-
-  #rctx = decode(events)
-  #assert len(rctx.inst_execs) > 0, "empty sqtt output"
-  #sqtt.update(rctx.inst_execs)
-
-  for e in events:
-    if isinstance(e, ProfileSQTTEvent):
-      print(replace(e, blob=b''))
-      if e.se == 0:
-        parse_sqtt_print_packets(e.blob)
-
-template = """.text
-.globl matmul
-.p2align 8
-.type matmul,@function
-matmul:
-  INSTRUCTION
-
-.rodata
-.p2align 6
-.amdhsa_kernel matmul
-  .amdhsa_kernarg_size 8
-  .amdhsa_user_sgpr_kernarg_segment_ptr 1
-  .amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
-  .amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
-  .amdhsa_wavefront_size32 1
-.end_amdhsa_kernel
-
-.amdgpu_metadata
---
-amdhsa.version:
-  - 1
-  - 0
-amdhsa.kernels:
-  - .name: matmul
-    .symbol: matmul.kd
-    .group_segment_fixed_size: 0
-    .private_segment_fixed_size: 0
-    .wavefront_size: 32
-    .sgpr_count: 8
-    .vgpr_count: 8
-    .max_flat_workgroup_size: 1024
-    .kernarg_segment_align: 8
-    .kernarg_segment_size: 8
-    .args:
-      - .address_space:  global
-        .name:           a
-        .offset:         0
-        .size:           8
-        .type_name:      'float*'
-        .value_kind:     global_buffer
-...
-.end_amdgpu_metadata
-"""
-
-def run_asm(src, num_workgroups=1, num_waves=1):
-  WAVE_SIZE = 32
-  t = Tensor.empty(0x1000).realize()
-  buf = t.uop.buffer.ensure_allocated()
-  lib = dev.compiler.compile(template.replace("INSTRUCTION", '\n'.join(src)))
-  dev.compiler.disassemble(lib)
-  fxn = AMDProgram(dev, "matmul", lib)
-  fxn(buf._buf, global_size=(num_workgroups,1,1), local_size=(WAVE_SIZE*num_waves,1,1), wait=True)
-
-if __name__ == "__main__":
-  with save_sqtt() as sqtt:
-    run_asm([
-      "s_nop 100",
-      "s_nop 100",
-      "s_load_b64 s[0:1], s[0:1], null",
-      "s_waitcnt lgkmcnt(0)",
-      "s_nop 100",
-      "s_nop 100",
-      "s_add_i32 s2, s2, 10",
-      "s_add_i32 s2, s2, 10",
-      "s_nop 100",
-      "s_nop 100",
-      "v_mov_b32_e32 v0, 0",
-      "v_mov_b32_e32 v0, 0",
-      "s_nop 100",
-      "s_nop 100",
-      "v_dual_fmac_f32 v2, v48, v24 :: v_dual_fmac_f32 v9, v37, v51",
-      "v_dual_fmac_f32 v2, v48, v24 :: v_dual_fmac_f32 v9, v37, v51",
-      "s_nop 100",
-      "s_nop 100",
-      "global_load_b128 v[2:5], v0, s[0:1]",
-      "global_load_b128 v[2:5], v0, s[0:1]",
-      "s_nop 100",
-      "s_nop 100",
-      "s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)",
-      "s_endpgm",
-    ], num_workgroups=1, num_waves=1)
-  exit(0)
-
-  with save_sqtt() as sqtt:
-    #(Tensor.empty(16,16) @ Tensor.empty(16,16)).elu().realize()
-    #Tensor.empty(1, 64).sum(axis=1).realize()
-    Tensor.empty(1).log2().realize()
-  exit(0)
-
-  with save_sqtt() as sqtt:
-    # what's in v0?
-    run_asm([
-      "v_mov_b32_e32 v0, 0",
-      "v_mov_b32_e32 v1, 0",
-      "s_clause 0x1",
-      "s_load_b64 s[0:1], s[0:1], null",
-      "s_waitcnt lgkmcnt(0)",
-    ]+[
-      "global_load_b32 v1, v0, s[0:1]",
-    ]*10+[
-      "global_load_b32 v10, v1, s[0:1]",
-      "s_waitcnt vmcnt(0)",
-
-      #"v_rcp_f32 v1, v0"
-      #"v_add_f32_e32 v1 v0 v0",
-      #"v_add_f32_e32 v5 v4 v4",
-      #"v_add_f32_e32 v7 v6 v6",
-      #"v_add_f32_e32 v1 v0 v0",
-      #"v_add_f32_e32 v2 v1 v1",
-      #"s_nop 1"
-    ]*5+[
-      "v_add_f32_e32 v3 v2 v2",
-    ]*5+[
-      "v_mul_f32_e32 v3 v2 v2",
-    ]*7)
@@ -1,548 +0,0 @@
-import pickle, sys
-from tinygrad.helpers import getenv, Timing, colored
-from extra.sqtt.roc import decode, ProfileSQTTEvent
-
-# do these enums match fields in the packets?
-#from tinygrad.runtime.support.amd import import_soc
-#soc = import_soc([11])
-#perf_sel = {getattr(soc, k):k for k in dir(soc) if k.startswith("SQ_PERF_")}
-
-# Instruction packets (one per ISA op)
-# NOTE: these are bad guesses and may be wrong! feel free to update if you know better
-# some names were taken from SQ_TT_TOKEN_MASK_TOKEN_EXCLUDE_SHIFT
-
-# we see 18 opcodes
-# opcodes(18):  1  2  3  4  5  6  8  9  F 10 11 12 14 15 16 17 18 19
-# if you exclude everything, you are left with 6
-# opcodes( 6): 10 11 14 15 16 17
-# sometimes we see a lot of B, but not repeatable
-
-# not seen
-# 7 A C
-
-# NOTE: INST runs before EXEC
-
-OPCODE_COLORS = {
-  # dispatches are BLACK
-  0x1: "BLACK",
-  0x18: "BLACK",
-
-  # execs are yellow
-  0x2: "yellow",
-  0x3: "yellow",
-  0x4: "YELLOW",
-  0x5: "YELLOW",
-
-  # waves are blue
-  0x8: "blue",
-  0x9: "blue",
-  0x6: "cyan",
-  0xb: "cyan",
-}
-
-OPCODE_NAMES = {
-  # gated by SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT (but others must be enabled for it to show)
-  0x01: "VALUINST",
-  # gated by SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT
-  0x02: "VMEMEXEC",
-  # gated by SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT
-  0x03: "ALUEXEC",
-  # gated by SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT
-  0x04: "IMMEDIATE",
-  0x05: "IMMEDIATE_MASK",
-
-  # gated by SQ_TT_TOKEN_EXCLUDE_WAVERDY_SHIFT
-  0x06: "WAVERDY",
-  # gated by SQ_TT_TOKEN_EXCLUDE_WAVESTARTEND_SHIFT
-  0x08: "WAVEEND",
-  0x09: "WAVESTART",
-  # gated by SQ_TT_TOKEN_EXCLUDE_WAVEALLOC_SHIFT
-  0x0B: "WAVEALLOC",  # FFF00
-
-  # gated by NOT SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT
-  0x0D: "PERF",
-  # gated by SQ_TT_TOKEN_EXCLUDE_EVENT_SHIFT
-  0x12: "EVENT",
-  0x13: "EVENT_BIG",  # FFFFF800
-  # some gated by SQ_TT_TOKEN_EXCLUDE_REG_SHIFT, some always there. something is broken with the timing on this
-  0x14: "REG",
-  # gated by SQ_TT_TOKEN_EXCLUDE_INST_SHIFT
-  0x18: "INST",
-  # gated by SQ_TT_TOKEN_EXCLUDE_UTILCTR_SHIFT
-  0x19: "UTILCTR",
-
-  # this is the first (8 byte) packet in the bitstream
-  0x17: "LAYOUT_HEADER",       # layout/mode/group + selectors A/B (reversed)
-
-  # pure time (no extra bits)
-  0x0F: "TS_DELTA_SHORT",
-  0x10: "NOP",
-  0x11: "TS_WAVE_STATE",     # almost pure time, has a small flag
-
-  # not a good name, but seen and understood mostly
-  0x15: "SNAPSHOT",          # small delta + 50-ish bits of snapshot
-  0x16: "TS_DELTA_OR_MARK",  # 36-bit long delta or 36-bit marker
-
-  # packets we haven't seen / rarely see 0x0b
-  0x07: "TS_DELTA_S8_W3_7",         # shift=8,  width=3  (small delta)
-  0x0A: "TS_DELTA_S5_W2_A",         # shift=5,  width=2
-  0x0C: "TS_DELTA_S5_W3_B",         # shift=5,  width=3 (different consumer)
-}
-
-#  SALU    =  0x0 / s_mov_b32
-#  SMEM    =  0x1 / s_load_b*
-#  JUMP    =  0x3 / s_cbranch_scc0
-#  NEXT    =  0x4 / s_cbranch_execz
-#  MESSAGE =  0x9 / s_sendmsg
-#  VALU    =  0xb / v_(exp,log)_f32_e32
-#  VALU    =  0xd / v_lshlrev_b64
-#  VALU    =  0xe / v_mad_u64_u32
-#  VMEM    = 0x21 / global_load_b32
-#  VMEM    = 0x22 / global_load_b32
-#  VMEM    = 0x24 / global_store_b32
-#  VMEM    = 0x25 / global_store_b64
-#  VMEM    = 0x27 / global_store
-#  VMEM    = 0x28 / global_store_b64
-#  LDS     = 0x29 / ds_load_b128
-#  LDS     = 0x2b / ds_store_b32
-#  LDS     = 0x2e / ds_store_b128
-#  ????    = 0x5a / hidden global_load  instruction
-#  ????    = 0x5b / hidden global_load  instruction
-#  ????    = 0x5c / hidden global_store instruction
-#  VALU    = 0x73 / v_cmpx_eq_u32_e32 (not normal VALUINST)
-OPNAME = {
-  0x0: "SALU",
-  0x1: "SMEM",
-  0x3: "JUMP",
-  0x4: "NEXT",
-  0x9: "MESSAGE",
-  0xb: "VALU",
-  0xd: "VALU",
-  0xe: "VALU",
-  0x21: "VMEM_LOAD",
-  0x22: "VMEM_LOAD",
-  0x24: "VMEM_STORE",
-  0x25: "VMEM_STORE",
-  0x26: "VMEM_STORE",
-  0x27: "VMEM_STORE",
-  0x28: "VMEM_STORE",
-  0x29: "LDS_LOAD",
-  0x2b: "LDS_STORE",
-  0x2e: "LDS_STORE",
-  0x50: "__SIMD_LDS_LOAD",
-  0x51: "__SIMD_LDS_LOAD",
-  0x54: "__SIMD_LDS_STORE",
-  0x5a: "__SIMD_VMEM_LOAD",
-  0x5b: "__SIMD_VMEM_LOAD",
-  0x5c: "__SIMD_VMEM_STORE",
-  0x5d: "__SIMD_VMEM_STORE",
-  0x5e: "__SIMD_VMEM_STORE",
-  0x5f: "__SIMD_VMEM_STORE",
-  0x72: "SALU_OR",
-  0x73: "VALU_CMPX",
-}
-
-ALUSRC = {
-  1: "SALU",
-  2: "VALU",
-  3: "VALU_SALU",
-}
-
-MEMSRC = {
-  0: "LDS",
-  1: "__LDS",
-  2: "VMEM",
-  3: "__VMEM",
-}
-
-
-# these tables are from rocprof trace decoder
-# rocprof_trace_decoder_parse_data-0x11c6a0
-# parse_sqtt_180 = b *rocprof_trace_decoder_parse_data-0x11c6a0+0x110040
-
-# ---------- 1. local_138: 256-byte state->opcode table ----------
-
-STATE_TO_OPCODE: bytes = bytes([
-  0x10, 0x16, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x17, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x07, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x19, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x00, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x11, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x12, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x15, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x16, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x17, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x07, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x19, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x00, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x11, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-  0x10, 0x13, 0x18, 0x01, 0x05, 0x0b, 0x0c, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x09, 0x04, 0x03, 0x02,
-  0x10, 0x15, 0x18, 0x01, 0x06, 0x08, 0x0d, 0x00, 0x0f, 0x14, 0x18, 0x01, 0x0a, 0x04, 0x03, 0x02,
-])
-
-# opcode mask (the bits used to determine the opcode, worked out by looking at the repeats in STATE_TO_OPCODE)
-
-opcode_mask = {
-  0x10: 0b1111,
-
-  0x16: 0b1111111,
-  0x17: 0b1111111,
-  0x07: 0b1111111,
-  0x19: 0b1111111,
-  0x11: 0b1111111,
-  0x12: 0b11111111,
-  0x13: 0b11111111,
-  0x15: 0b1111111,
-
-  0x18: 0b111,
-  0x1: 0b111,
-
-  0x5: 0b11111,
-  0x6: 0b11111,
-  0xb: 0b11111,
-  0x8: 0b11111,
-  0xc: 0b11111,
-  0xd: 0b11111,
-
-  0xf: 0b1111,
-  0x14: 0b1111,
-
-  0x9: 0b11111,
-  0xa: 0b11111,
-
-  0x4: 0b1111,
-  0x3: 0b1111,
-  0x2: 0b1111,
-}
-
-# ---------- 2. DAT_0012e280: nibble budget per opcode&0x1F ----------
-
-NIBBLE_BUDGET = [
-  0x08, 0x0C, 0x08, 0x08, 0x0C, 0x18, 0x18, 0x40, 0x14, 0x20, 0x30, 0x14, 0x34, 0x1C, 0x30, 0x08,
-  0x04, 0x18, 0x18, 0x20, 0x40, 0x40, 0x30, 0x40, 0x14, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
-]
-
-# ---------- 3. delta_map from your hash nodes ----------
-
-# opcode -> (shift, width)
-DELTA_MAP_DEFAULT = {
-  0x01: (3,  3),   # shift=3,  end=6
-  0x02: (4,  2),   # shift=4,  end=6
-  0x03: (4,  2),   # shift=4,  end=6
-  0x04: (4,  3),   # shift=4,  end=7
-  0x05: (5,  3),   # shift=5,  end=8
-  0x06: (5,  3),   # shift=5,  end=8
-  0x07: (8,  3),   # shift=8,  end=11
-  0x08: (5,  3),   # shift=5,  end=8
-  0x09: (5,  2),   # shift=5,  end=7
-  0x0A: (5,  2),   # shift=5,  end=7
-  0x0B: (5,  3),   # shift=5,  end=8
-  0x0C: (5,  3),   # shift=5,  end=8
-  0x0D: (5,  3),   # shift=5,  end=8
-  # NOTE: 0x0e can never be decoded, it's not in the STATE_TO_OPCODE table
-  #0x0E: (7,  2),   # shift=7,  end=9
-  0x0F: (4,  4),   # shift=4,  end=8
-  0x10: (0,  0),   # shift=0,  end=0  (no delta)
-  0x11: (7,  9),   # shift=7,  end=16
-  0x12: (8,  3),   # shift=8,  end=11
-  0x13: (8,  3),   # shift=8,  end=11
-  0x14: (4,  3),   # shift=4,  end=7
-  0x15: (7,  3),   # shift=7,  end=10
-  0x16: (12, 36),  # shift=12, end=48 (36-bit field, matches the 0x16 special-case)
-  0x17: (0,  0),   # shift=0,  end=0  (no delta)
-  0x18: (4,  3),   # shift=4,  end=7
-  0x19: (7,  2),   # shift=7,  end=9
-}
-
-# ---------- 4. One-line-per-packet parser ----------
-
-def reg_mask(opcode):
-  nb_bits = NIBBLE_BUDGET[opcode & 0x1F]
-  shift, width = DELTA_MAP_DEFAULT[opcode]
-  delta_mask = ((1 << width) - 1) << shift
-  assert delta_mask & opcode_mask[opcode] == 0, "masks shouldn't overlap"
-  return ((1 << nb_bits) - 1) & ~(delta_mask | opcode_mask[opcode])
-
-def decode_packet_fields(opcode: int, reg: int) -> str:
-  """
-  Decode packet payloads conservatively, using:
-    - NIBBLE_BUDGET[opcode & 0x1F] to mask reg down to true width.
-    - DELTA_MAP_DEFAULT[opcode] to expose the "primary" field (often delta).
-    - Per-opcode layouts derived from rocprof's decompiled consumers.
-  """
-  # --- 0. Restrict to real packet bits not used in delta ---------------------------------
-  pkt = reg & reg_mask(opcode)
-  fields: list[str] = []
-
-  match opcode:
-    case 0x01: # VALUINST
-      # 6 bit field
-      flag = (pkt >> 6) & 1
-      wave = pkt >> 7
-      fields.append(f"wave={wave:x}")
-      if flag: fields.append("flag")
-    case 0x02: # VMEMEXEC
-      # 2 bit field (pipe is a guess)
-      src = pkt>>6
-      fields.append(f"src={src} [{MEMSRC.get(src, '')}]")
-    case 0x03: # ALUEXEC
-      # 2 bit field
-      src = pkt>>6
-      fields.append(f"src={src} [{ALUSRC.get(src, '')}]")
-    case 0x04: # IMMEDIATE_4
-      # 5 bit field (actually 4)
-      wave = pkt >> 7
-      fields.append(f"wave={wave:x}")
-    case 0x05: # IMMEDIATE_5
-      # 16 bit field
-      # 1 bit per wave
-      fields.append(f"mask={pkt>>8:016b}")
-    case 0x6:
-      # wave ready FFFF00
-      # 16 bit field
-      # 1 bit per wave
-      fields.append(f"mask={pkt>>8:016b}")
-    case 0x0d:
-      # 20 bit field
-      fields.append(f"arg = {pkt>>8:X}")
-    case 0x12:
-      fields.append(f"event = {pkt>>11:X}")
-    case 0x15:
-      fields.append(f"snap = {pkt>>10:X}")
-    case 0x19:
-      # wave end
-      fields.append(f"ctr = {pkt>>9:X}")
-    case 0xf:
-      extracted_delta = (reg >> 4) & 0xF
-      fields.append(f"strange_delta=0x{extracted_delta:x}")
-    case 0x11:
-      # DELTA_MAP_DEFAULT: shift=7, width=9 -> small delta.
-      # FF0000 is the mask
-      coarse    = pkt >> 16
-      fields.append(f"coarse=0x{coarse:02x}")
-      # From decomp:
-      #  - when layout<3 and coarse&1, it sets a "has interesting wave" flag
-      #  - when coarse&8, it marks all live waves as "terminated"
-      if coarse & 0x01:
-        fields.append("flag_wave_interest=1")
-      if coarse & 0x08:
-        fields.append("flag_terminate_all=1")
-    case 0x8:
-      # wave end, this is 20 bits (FFF00)
-      flag7   = (pkt >> 8) & 1
-      simd    = (pkt >> 9) & 3
-      cu      = ((pkt >> 11) & 0x7) | (flag7 << 3)
-      wave    = (pkt >> 15) & 0x1f
-      fields.append(f"wave={wave:x}")
-      fields.append(f"simd={simd}")
-      fields.append(f"cu={cu}")
-    case 0x9:
-      # From case 9 (WAVESTART) in multiple consumers:
-      #   flag7  = (w >> 7) & 1        (low bit of uVar41)
-      #   cls2   = (w >> 8) & 3        (class / group)
-      #   slot4  = (w >> 10) & 0xf     (slot / group index)
-      #   idx_lo = (w >> 0xd) & 0x1f   (low index, layout<4 path)
-      #   idx_hi = (w >> 0xf) & 0x1f   (high index, layout>=4 path)
-      #   id7    = (w >> 0x19) & 0x7f  (7-bit id)
-      flag7   = (pkt >> 7) & 1
-      simd    = (pkt >> 8) & 3
-      cu      = ((pkt >> 10) & 0x7) | (flag7 << 3)
-      wave    = (pkt >> 13) & 0x1F
-      id7     = (pkt >> 17)
-      fields.append(f"wave={wave:x}")
-      fields.append(f"simd={simd}")
-      fields.append(f"cu={cu}")
-      fields.append(f"id7=0x{id7:x}")
-    case 0x18:
-      # FFF88 is the mask
-      # From case 0x18:
-      #   low3   = w & 7
-      #   grp3   = (w >> 3) or (w >> 4) & 7   (layout-dependent)
-      #   flags  = bits 6 (B6) and 7 (B7)
-      #   hi8    = (w >> 0xc) & 0xff   (layout 4 path)
-      #   hi7    = (w >> 0xd) & 0x7f   (other layouts)
-      #   idx5   = (w >> 7) or (w >> 8) & 0x1f, used as wave index
-      flag1 = (pkt >> 3) & 1
-      flag2 = (pkt >> 7) & 1
-      wave = (pkt >> 8) & 0x1F
-      op = (pkt >> 13)
-      fields.append(f"wave={wave:x}")
-      fields.append(f"op=0x{op:02x} [{OPNAME.get(op, '')}]")
-      if flag1: fields.append("flag1")
-      if flag2: fields.append("flag2")
-    case 0x14:
-      subop   = (pkt >> 16) & 0xFFFF       # (short)(w >> 0x10)
-      val32   = (pkt >> 32) & 0xFFFFFFFF   # (uint)(w >> 0x20)
-      slot    = (pkt >> 7) & 0x7           # index in local_168[...] tables
-      hi_byte = (pkt >> 8) & 0xFF          # determines config vs marker
-
-      fields.append(f"subop=0x{subop:04x}")
-      fields.append(f"slot={slot}")
-      fields.append(f"val32=0x{val32:08x}")
-
-      if hi_byte & 0x80:
-        # Config flavour: writes config words into per-slot state arrays.
-        fields.append("kind=config")
-        if subop == 0x000C:
-          fields.append("slot=lo")
-        elif subop == 0x000D:
-          fields.append("slot=hi")
-      else:
-        # COR marker: subop 0xC342, payload "COR\0" → start of a COR region.
-        if subop == 0xC342:
-          fields.append("kind=cor_stream")
-          if val32 == 0x434F5200:
-            fields.append("cor_magic='COR\\0'")
-    case 0x16:
-      # Bits:
-      #   bit8  -> 0x100
-      #   bit9  -> 0x200
-      #   bits 12..47 -> 36-bit field used as delta or marker
-      bit8 = bool(pkt & 0x100)
-      bit9 = bool(pkt & 0x200)
-      if not bit9:
-        mode = "delta"
-      elif not bit8:
-        mode = "marker"
-      else:
-        mode = "other"
-      # need to use reg here
-      val36 = (reg >> 12) & ((1 << 36) - 1)
-      fields.append(f"mode={mode}")
-      if mode != "delta":
-        fields.append(f"val36=0x{val36:x}")
-    case 0x17:
-      # From decomp (two sites with identical logic):
-      #   layout = (w >> 7) & 0x3f
-      #   mode   = (w >> 0xd) & 3
-      #   group  = (w >> 0xf) & 7
-      #   sel_a  = (w >> 0x1c) & 0xf
-      #   sel_b  = (w >> 0x21) & 7
-      #   flag4  = (w >> 0x3b) & 1  (only meaningful when layout == 4)
-      layout = (pkt >> 7)  & 0x3F
-      simd   = (pkt >> 13) & 0x3  # you can change this by changing traced simd
-      group  = (pkt >> 15) & 0x7
-      sel_a  = (pkt >> 0x1C) & 0xF
-      sel_b  = (pkt >> 0x21) & 0x7
-      flag4  = (pkt >> 0x3B) & 0x1
-
-      fields.append(f"layout={layout}")
-      fields.append(f"group={group}")
-      fields.append(f"simd={simd}")
-      fields.append(f"sel_a={sel_a}")
-      fields.append(f"sel_b={sel_b}")
-      if layout == 4:
-        fields.append(f"layout4_flag={flag4}")
-    case _:
-      fields.append(f"{pkt:X} & {reg_mask(opcode):X}")
-  return ",".join(fields)
-
-FILTER_LEVEL = getenv("FILTER", 1)
-
-DEFAULT_FILTER: tuple[int, ...] = tuple()
-# NOP + pure time + "sample"
-if FILTER_LEVEL >= 0: DEFAULT_FILTER += (0x10, 0xf, 0x11)
-# reg + event + sample + marker
-# TODO: events are probably good
-if FILTER_LEVEL >= 1: DEFAULT_FILTER += (0x14, 0x12, 0x16)
-# instruction runs + valuinst
-if FILTER_LEVEL >= 2: DEFAULT_FILTER += (0x01, 0x02, 0x03)
-# instructions dispatch (inst, immed)
-if FILTER_LEVEL >= 3: DEFAULT_FILTER += (0x4, 0x5, 0x18)
-# waves
-if FILTER_LEVEL >= 4: DEFAULT_FILTER += (0x6, 0x8, 0x9)
-
-def parse_sqtt_print_packets(data: bytes, filter=DEFAULT_FILTER, verbose=True) -> None:
-  """
-  Minimal debug: print ONE LINE per decoded token (packet).
-
-  Now prints only the actual nibbles that belong to each packet, instead of
-  the full 64-bit shift register.
-  """
-  n = len(data)
-  time = 0
-  last_printed_time = 0
-  reg = 0          # shift register
-  offset = 0       # bit offset, in steps of 4 (one nibble)
-  nib_budget = 0x40
-  flags = 0
-  token_index = 0
-  opcodes_seen = set()
-
-  while (offset >> 3) < n:
-    # 1) Fill register with nibbles according to nib_budget
-    if nib_budget != 0:
-      target = offset + 4 + ((nib_budget - 1) & ~3)
-      while offset != target and (offset >> 3) < n:
-        byte = data[offset >> 3]
-        nib = (byte >> (offset & 4)) & 0xF
-        reg = ((reg >> 4) | (nib << 60)) & ((1 << 64) - 1)
-        offset += 4
-      if offset != target: break  # don't parse past the end
-
-    # 2) Decode token from low 8 bits
-    opcode = STATE_TO_OPCODE[reg & 0xFF]
-    opcodes_seen.add(opcode)
-
-    # 4) Set next nibble budget based on opcode
-    nib_budget = NIBBLE_BUDGET[opcode & 0x1F]
-
-    # 5) Get delta
-    shift, width = DELTA_MAP_DEFAULT[opcode]
-    delta = (reg >> shift) & ((1 << width) - 1)
-
-    # 6) Update time and handle special opcodes 0xF/0x16
-    if opcode == 0x16:
-      two_bits = (reg >> 8) & 0x3
-      if two_bits == 1:
-        flags |= 0x01
-
-      # Common 36-bit field at bits [12..47]
-      if (reg & 0x200) == 0:
-        # delta mode: add 36-bit delta to time
-        pass
-      elif (reg & 0x100) == 0:
-        # marker / other modes: no time advance
-        # real marker: bit9=1, bit8=0, non-zero payload
-        # "other" 0x16 variants, ignored for timing
-        delta = 0
-      else:
-        raise RuntimeError("unknown 0x16 delta")
-    elif opcode == 0x0F:
-      # opcode 0x0F has an offset of 4 to the delta
-      # update: it's actually computed to be 8 to match WAVESTART
-      delta = delta + 8
-
-    # Append extra decoded fields into the note string
-    note = decode_packet_fields(opcode, reg)
-
-    # this delta happens before the instruction
-    time += delta
-    token_index += 1
-
-    if verbose and (filter is None or opcode not in filter):
-      print(f"{time:8d} +{time-last_printed_time:8d} : "+colored(f"{OPCODE_NAMES[opcode]:18s} ", OPCODE_COLORS.get(opcode, "white"))+f"{note}")
-      last_printed_time = time
-
-  # Optional summary at the end
-  print(f"# done: tokens={token_index:_}, final_time={time}, flags=0x{flags:02x}")
-  if verbose:
-    print(f"opcodes({len(opcodes_seen):2d}):",
-          ' '.join([colored(f"{op:2X}", "WHITE" if op in opcodes_seen else "BLACK") for op in sorted(opcode_mask)]))
-
-
-def parse(fn:str):
-  with Timing(f"unpickle {fn}: "): dat = pickle.load(open(fn, "rb"))
-  #if getenv("ROCM", 0):
-  #  with Timing(f"decode {fn}: "): ctx = decode(dat)
-  dat_sqtt = [x for x in dat if isinstance(x, ProfileSQTTEvent)]
-  print(f"got {len(dat_sqtt)} SQTT events in {fn}")
-  return dat_sqtt
-
-if __name__ == "__main__":
-  fn = "extra/sqtt/examples/profile_gemm_run_0.pkl"
-  dat_sqtt = parse(sys.argv[1] if len(sys.argv) > 1 else fn)
-  for i,dat in enumerate(dat_sqtt):
-    with Timing(f"decode pkt {i} with len {len(dat.blob):_}: "):
-      parse_sqtt_print_packets(dat.blob, verbose=getenv("V", 1))
@@ -1,24 +1,25 @@
-import os, subprocess, sys
+import os, subprocess, sys, shlex
 from pathlib import Path
 from tinygrad.helpers import temp

 EXAMPLES_DIR = Path(__file__).parent
 PROFILE_PATH = Path(temp("profile.pkl", append_user=True))

-EXAMPLES = [
-  "test/backend/test_custom_kernel.py TestCustomKernel.test_empty",
-  "test/test_tiny.py TestTiny.test_plus",
-  "test/test_tiny.py TestTiny.test_gemm",
-]
+EXAMPLES = {
+  "empty":"test/backend/test_custom_kernel.py TestCustomKernel.test_empty",
+  "plus":"test/test_tiny.py TestTiny.test_plus",
+  "gemm":"-c \"from tinygrad import Tensor; (Tensor.empty(N:=32, N)@Tensor.empty(N, N)).realize()\"",
+  "sync":"test/amd/test_custom_kernel.py TestCustomKernel.test_lds_sync",
+}

 if __name__ == "__main__":
  arch = subprocess.check_output(["python", "-c", "from tinygrad import Device; print(Device['AMD'].arch)"], text=True,
                                 env={**os.environ, "DEBUG":"0"}).rstrip()
  (EXAMPLES_DIR/arch).mkdir(exist_ok=True)
-  for test in EXAMPLES:
+  for name,test in EXAMPLES.items():
    for i in range(2):
      # AM_RESET=1 gets a clear trace, does not work on mi300 machines
-      subprocess.run([sys.executable, *test.split()], cwd=EXAMPLES_DIR.parent.parent.parent,
+      subprocess.run([sys.executable, *shlex.split(test)], cwd=EXAMPLES_DIR.parent.parent.parent,
                     env={**os.environ, "AMD":"1", "AM_RESET":"1" if not arch.startswith("gfx9") else "0", "VIZ":"-2", "PYTHONPATH":"."})
-      PROFILE_PATH.rename(dest:=EXAMPLES_DIR/arch/f"profile_{test.split('.')[-1].replace('test_', '')}_run_{i}.pkl")
+      PROFILE_PATH.rename(dest:=EXAMPLES_DIR/arch/f"profile_{name}_run_{i}.pkl")
      print(f"saved SQTT trace to {dest}")
@@ -118,7 +118,7 @@ def decode(sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, Inst]])
    nonlocal exc
    try: rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
    except AttributeError as e:
-      exc = RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_sqtt_decoder.py to install")
+      exc = RuntimeError("Failed to find rocprof-trace-decoder. Run sudo ./extra/sqtt/install_rocprof_decoder.py to install")
      exc.__cause__ = e
  (t:=threading.Thread(target=worker, daemon=True)).start()
  t.join()
@@ -10,15 +10,47 @@ from tinygrad.uop.ops import UOp, Ops, KernelInfo

 def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None=None) -> Tensor:
  dtype = dtype or ref.dtype
-  if not isinstance(ref.device, tuple): return Tensor.empty(*shape, dtype=dtype, device=ref.device)
+  if not isinstance(ref.device, tuple): return Tensor.invalid(*shape, dtype=dtype, device=ref.device)
  shard_axis = ref.uop.axis if axis is None else axis
  shape = tuple(s // len(ref.device) if i == shard_axis else s for i, s in enumerate(shape))
  axis = ref.uop.axis if axis is None else axis
-  return Tensor(Tensor.empty(*shape, dtype=dtype, device=ref.device).uop.multi(axis), dtype=dtype, device=ref.device)
+  return Tensor(Tensor.invalid(*shape, dtype=dtype, device=ref.device).uop.multi(axis), dtype=dtype, device=ref.device)

 def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
  return _sharded_empty(ref.shape, ref, axis)

+@functools.cache
+def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch):
+  def grad(dou:UOp, ker:UOp) -> tuple[None, None, UOp, UOp, UOp]:
+    do = Tensor(dou, device=dou.device)
+    attn = Tensor(ker.src[1].after(ker), device=ker.src[1].device)
+    l_vec = Tensor(ker.src[2].after(ker), device=ker.src[2].device)
+    xq = Tensor(ker.src[3], device=ker.src[3].device)
+    xk = Tensor(ker.src[4], device=ker.src[4].device)
+    xv = Tensor(ker.src[5], device=ker.src[5].device)
+
+    dq = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t)
+    GROUP_SIZE = H_local // H_KV_local
+    dk_partial = _sharded_empty((B * GROUP_SIZE, N, H_KV, D), xk, axis=shard_axis)
+    dv_partial = _sharded_empty((B * GROUP_SIZE, N, H_KV, D), xv, axis=shard_axis)
+
+    # delta_vec = (do * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
+    delta_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
+    delta_vec, dq = Tensor.custom_kernel(delta_vec, dq, attn, do, fxn=functools.partial(custom_fa_backward_pre, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:2]
+
+    dq, dk_partial, dv_partial = Tensor.custom_kernel(dq, dk_partial, dv_partial, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:3]
+
+    # unshuffle dq: atomic_pk_add_bf16_with_warpid creates a shuffled layout within each 16x128 tile
+    # decompose each tile into (j=4, a=2, b=2, d=4, e=4, k=4, c=2) and permute to (e, k, j, a, d, b, c) = standard row-major
+    dq = dq.reshape(B, H, N//16, 4, 2, 2, 4, 4, 4, 2).permute(0, 1, 2, 7, 8, 3, 4, 6, 5, 9).reshape(B, H, N, D).transpose(1, 2)
+
+    # reduce partial dK/dV across GROUP_SIZE query heads
+    dk = dk_partial.reshape(B, GROUP_SIZE, N, H_KV, D).sum(1)
+    dv = dv_partial.reshape(B, GROUP_SIZE, N, H_KV, D).sum(1)
+
+    return None, None, dq.uop, dk.uop, dv.uop
+  return grad
+
 def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False):
  assert attn_mask is None, "attn_mask not supported"
  assert is_causal, "only causal attention supported"
@@ -45,23 +77,7 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
  attn = _sharded_empty_like(xq, axis=shard_axis)
  l_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)

-  def grad(dou:UOp, _) -> tuple[None, None, UOp, UOp, UOp]:
-    do = Tensor(dou, device=dou.device)
-    dq_in = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t)
-    dq = _sharded_empty_like(xq, axis=shard_axis)
-    dk = _sharded_empty_like(xk, axis=shard_axis)
-    dv = _sharded_empty_like(xv, axis=shard_axis)
-
-    # delta_vec = (do * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
-    delta_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
-    delta_vec, dq_in = Tensor.custom_kernel(delta_vec, dq_in, attn, do, fxn=functools.partial(custom_fa_backward_pre, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:2]
-
-    dq_in, dk, dv = Tensor.custom_kernel(dq_in, dk, dv, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:3]
-
-    # unshuffle dq
-    dq = Tensor.custom_kernel(dq, dq_in, fxn=functools.partial(custom_fa_backward_post, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[0]
-
-    return None, None, dq.uop, dk.uop, dv.uop
+  grad = _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch)

  attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D), grad_fxn=grad)[:2]

@@ -89,7 +105,6 @@ def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, device:str, arch:st
                  arg=KernelInfo(name="custom_fa_forward", estimates=estimates))

  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-
  lib = bytearray(lib)
  rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
  struct.pack_into('<I', lib, rodata_off, 160000)
@@ -120,7 +135,6 @@ def custom_fa_backward_pre(delta_vec:UOp, dq:UOp, o:UOp, do:UOp, device:str, arc
                  arg=KernelInfo(name="custom_fa_backward_pre", estimates=estimates))

  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-
  lib = bytearray(lib)
  rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
  struct.pack_into('<I', lib, rodata_off, 160000)
@@ -138,7 +152,7 @@ def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_ve
  BLOCK_SIZE_KV = 256
  NUM_WARPS = 4
  NUM_THREADS = 64 * NUM_WARPS
-  gsz = (H_KV, N // BLOCK_SIZE_KV, B)
+  gsz = (H, N // BLOCK_SIZE_KV, B)
  lsz = (NUM_THREADS, 1, 1)
  threadIdx_x = UOp.special(lsz[0], "lidx0")
  blockIdx_x, blockIdx_y, blockIdx_z = UOp.special(gsz[0], "gidx0"), UOp.special(gsz[1], "gidx1"), UOp.special(gsz[2], "gidx2")
@@ -151,7 +165,6 @@ def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_ve
                  arg=KernelInfo(name="custom_fa_backward", estimates=estimates))

  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-
  lib = bytearray(lib)
  rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
  struct.pack_into('<I', lib, rodata_off, 160000)
@@ -182,7 +195,6 @@ def custom_fa_backward_post(dq_out:UOp, dq_in:UOp, device:str, arch:str, B:int,
                  arg=KernelInfo(name="custom_fa_backward_post", estimates=estimates))

  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-
  lib = bytearray(lib)
  rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
  struct.pack_into('<I', lib, rodata_off, 160000)
@@ -37,7 +37,7 @@ using namespace kittens;
 using _gl_QdO  = gl<bf16, ATTN_B, ATTN_N, ATTN_H, ATTN_D>;
 using _gl_KV   = gl<bf16, ATTN_B, ATTN_N, ATTN_H_KV, ATTN_D>;
 using _gl_dQ   = gl<bf16, ATTN_B, ATTN_H, ATTN_N, ATTN_D>;
-using _gl_dKV  = gl<bf16, ATTN_B, ATTN_N, ATTN_H_KV, ATTN_D>;
+using _gl_dKV  = gl<bf16, ATTN_B * GROUP_SIZE, ATTN_N, ATTN_H_KV, ATTN_D>;
 using _gl_Lvec = gl<float, ATTN_B, ATTN_H, 1, ATTN_N>;

 template<int D> struct attn_bwd_combined_globals {
@@ -47,7 +47,7 @@ template<int D> struct attn_bwd_combined_globals {
  _gl_dQ dQg;
  _gl_dKV dKg, dVg;
  _gl_Lvec L_vec, delta_vec;
-  dim3 grid() { return dim3(ATTN_H_KV, (ATTN_N / BLOCK_SIZE_KV), ATTN_B); }
+  dim3 grid() { return dim3(ATTN_H, (ATTN_N / BLOCK_SIZE_KV), ATTN_B); }
  dim3 block() { return dim3(NUM_THREADS); }
  size_t dynamic_shared_memory() { return MAX_SHARED_MEMORY; }
 };
@@ -55,10 +55,12 @@ template<int D> struct attn_bwd_combined_globals {
 template<int D> __launch_bounds__(NUM_THREADS, 1)
 __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr, bf16 *dO_ptr, bf16 *Q_ptr, bf16 *K_ptr, bf16 *V_ptr, float *L_vec_ptr, float *delta_vec_ptr) {

-  const int kv_head_idx = blockIdx.x;  // This is the KV head index
+  const int q_head_idx_fixed = blockIdx.x;  // This is the query head index [0, ATTN_H)
+  const int kv_head_idx = q_head_idx_fixed / GROUP_SIZE;
+  const int q_head_in_group = q_head_idx_fixed % GROUP_SIZE;
  const int seq_idx = blockIdx.y;
  const int batch_idx = blockIdx.z;
-  const int first_q_head = kv_head_idx * GROUP_SIZE;
+  const int first_q_head = q_head_idx_fixed;

  const int warpid = kittens::warpid();
  const int j = seq_idx * NUM_WARPS + warpid;
@@ -70,7 +72,7 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
  // first Q step that can overlap this K_span:
  const int first_step = max(0, k_start_min / STEP_QO);
  const int num_steps_per_head = total_steps_per_head - first_step;
-  const int num_steps = num_steps_per_head * GROUP_SIZE;
+  const int num_steps = num_steps_per_head;
  const int k_pos = j * WARP_SIZE_KV;

  constexpr float L_SCALE_FACTOR = 1.44269504089f;
@@ -3355,14 +3357,14 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
    }
  }

-  store<1>(g.dVg, dV_j, {batch_idx, 0, kv_head_idx, 0}, {0, j, 0, 0});
+  store<1>(g.dVg, dV_j, {batch_idx * GROUP_SIZE + q_head_in_group, 0, kv_head_idx, 0}, {0, j, 0, 0});
  __builtin_amdgcn_s_waitcnt(0);
  __builtin_amdgcn_s_barrier();

  // We first copy dV_j_T from accumulator GPRs to vector GPRs and then perform the store
  accvgpr_read(dV_j_T, dK_j_T);
  mul(dV_j_T, dV_j_T, dP_SCALE_FACTOR);
-  store<1>(g.dKg, dV_j, {batch_idx, 0, kv_head_idx, 0}, {0, j, 0, 0});
+  store<1>(g.dKg, dV_j, {batch_idx * GROUP_SIZE + q_head_in_group, 0, kv_head_idx, 0}, {0, j, 0, 0});

  // Write out final dQ_i slice
  mul(dQ_i_T, dQ_i_T, dP_SCALE_FACTOR);
@@ -66,7 +66,7 @@ template<int D, typename T=bf16, typename L=row_l, typename S=rt_32x16_s> using
 template<int D, typename T=bf16, typename L=col_l, typename S=rt_16x32_s> using qo_tile_transposed = rt<T, D, Q_BLOCK_SIZE, L, S>;
 template<int D, typename T=bf16, typename L=row_l, typename S=rt_32x16_s> using kv_tile = rt<T, KV_BLOCK_SIZE, D, L, S>;
 template<int D, typename T=bf16, typename L=col_l, typename S=rt_16x32_s> using kv_tile_transposed = rt<T, D, KV_BLOCK_SIZE, L, S>;
-template<int D, typename T=float, typename L=col_l, typename S=rt_16x32_4_s> using attn_tile = rt<T, KV_BLOCK_SIZE, Q_BLOCK_SIZE, L, S>;
+template<typename T=float, typename L=col_l, typename S=rt_16x32_4_s> using attn_tile = rt<T, KV_BLOCK_SIZE, Q_BLOCK_SIZE, L, S>;

 /**********************************************************/
 template<int THR_X, int THR_Y>
@@ -103,7 +103,7 @@ __device__ inline void mask_kv_tile(RT &dst, int q_abs, int k_abs, uint32_t neg_

    #pragma unroll
    for (int i = 0; i < dst.height; ++i) {
-        // Row base of the 32x* chunk produced by MFMA 
+        // Row base of the 32x* chunk produced by MFMA
        const int row_base = (i * 32) + ((lane >> 5) << 2); // multiplesof 4

        // Relative index of the FIRST element in this row-chunk w.r.t. q_pos
@@ -148,7 +148,7 @@ __device__ inline void mask_kv_tile(RT &dst, int q_abs, int k_abs, uint32_t neg_
 /**********************************************************/

 template<int D> struct attn_globals {
-    _gl_QKVO Qg, Kg, Vg, Og; 
+    _gl_QKVO Qg, Kg, Vg, Og;
    gl<float, -1, -1, -1, -1> L_vec;
    dim3 grid() { return dim3(ATTN_H, ((ATTN_N / Q_BLOCK_SIZE + NUM_WARPS - 1) / NUM_WARPS), ATTN_B); }
    dim3 block() { return dim3(NUM_THREADS); }
@@ -196,10 +196,10 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p

    kv_tile<D, bf16, col_l, rt_16x32_4_s> v_reg;
    qo_tile_transposed<D, float, col_l, rt_32x32_s> o_reg; // Output tile.
-    attn_tile<D, float, col_l, rt_32x32_s> att_block[2]; // attention tile, in float.
-    attn_tile<D, bf16, col_l, rt_32x32_s> att_block_bf16;
-    attn_tile<D, bf16, col_l, rt_16x32_4_s> att_block_bf16_in;
-    typename attn_tile<D, float, col_l, rt_32x32_s>::row_vec max_vec, norm_vec, max_vec_prev, scale_vec;
+    attn_tile<float, col_l, rt_32x32_s> att_block[2]; // attention tile, in float.
+    attn_tile<bf16, col_l, rt_32x32_s> att_block_bf16;
+    attn_tile<bf16, col_l, rt_16x32_4_s> att_block_bf16_in;
+    typename attn_tile<float, col_l, rt_32x32_s>::row_vec max_vec, norm_vec, max_vec_prev, scale_vec;

    zero(o_reg);
    zero(norm_vec);
@@ -241,8 +241,8 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    zero(att_block[0]);
    transpose(k_reg_transposed, k_reg);
    mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
-    __builtin_amdgcn_sched_barrier(0); 
-    if constexpr (causal) { 
+    __builtin_amdgcn_sched_barrier(0);
+    if constexpr (causal) {
        const int kv_end_pos = (1) * KV_BLOCK_SIZE;
        if (__builtin_expect(q_start_pos < kv_end_pos, 0)) {  // Only mask if needed
            mask_kv_tile(att_block[0], tile_idx, 0, neg_inf_v, lane);
@@ -269,7 +269,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    load(k_reg, k_smem[1]);
    // All warps then collaboratively load in the third slice of K (K2) into shared memory
    G::load<1, false>(k_smem[0], g.Kg, {batch_idx, 2, head_idx_kv, 0}, swizzled_offsets_K);
-    // All warps then collaboratively load in the second slice of V (V1) into shared memory 
+    // All warps then collaboratively load in the second slice of V (V1) into shared memory
    G::load<1, false>(v_smem[1], g.Vg, {batch_idx, 1, head_idx_kv, 0}, swizzled_offsets_V);
    asm volatile("s_waitcnt lgkmcnt(0)");
    asm volatile("s_waitcnt vmcnt(4)");
@@ -288,7 +288,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
        mul(norm_vec, norm_vec, scale_vec);
        col_sum(norm_vec, att_block[0], norm_vec);
        copy(att_block_bf16, att_block[0]);
-        att_block_bf16_in = *reinterpret_cast<attn_tile<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+        att_block_bf16_in = *reinterpret_cast<attn_tile< bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
        sched_barrier_exp_pairs<6, 3, 1>();
        sched_barrier_pairs<10, 5, 1>();
        __builtin_amdgcn_sched_barrier(0);
@@ -296,7 +296,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
        __builtin_amdgcn_sched_barrier(0);

        // Cluster 1:
-        //      Load K3 into shared 
+        //      Load K3 into shared
        G::load<1, false>(k_smem[1], g.Kg, {batch_idx, j, head_idx_kv, 0}, swizzled_offsets_K);
        //      Load V0 into registers
        load(v_reg, v_smem[0]);
@@ -348,7 +348,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
        mul(norm_vec, norm_vec, scale_vec);
        col_sum(norm_vec, att_block[1], norm_vec);
        copy(att_block_bf16, att_block[1]);
-        att_block_bf16_in = *reinterpret_cast<attn_tile<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+        att_block_bf16_in = *reinterpret_cast<attn_tile<bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
        sched_barrier_exp_pairs<6, 3, 3>();
        sched_barrier_pairs<10, 5, 3>();
        __builtin_amdgcn_s_setprio(0);
@@ -417,7 +417,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p

    col_sum(norm_vec, att_block[0], norm_vec);
    copy(att_block_bf16, att_block[0]);
-    att_block_bf16_in = *reinterpret_cast<attn_tile<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+    att_block_bf16_in = *reinterpret_cast<attn_tile<bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
    sched_barrier_exp_pairs<6, 3, 5>();
    sched_barrier_pairs<10, 5, 5>();
    __builtin_amdgcn_sched_barrier(0);
@@ -482,7 +482,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    mul(norm_vec, norm_vec, scale_vec);
    col_sum(norm_vec, att_block[1], norm_vec);
    copy(att_block_bf16, att_block[1]);
-    att_block_bf16_in = *reinterpret_cast<attn_tile<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+    att_block_bf16_in = *reinterpret_cast<attn_tile<bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
    sched_barrier_exp_pairs<6, 3, 7>();
    sched_barrier_pairs<10, 5, 7>();
    __builtin_amdgcn_sched_barrier(0);
@@ -544,7 +544,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    mul(norm_vec, norm_vec, scale_vec);
    col_sum(norm_vec, att_block[0], norm_vec);
    copy(att_block_bf16, att_block[0]);
-    att_block_bf16_in = *reinterpret_cast<attn_tile<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+    att_block_bf16_in = *reinterpret_cast<attn_tile<bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
    sched_barrier_exp_pairs<6, 3, 9>();
    sched_barrier_pairs<10, 5, 9>();
    __builtin_amdgcn_sched_barrier(0);
@@ -586,7 +586,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p

    col_sum(norm_vec, att_block[1], norm_vec);
    copy(att_block_bf16, att_block[1]);
-    att_block_bf16_in = *reinterpret_cast<attn_tile<D, bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);
+    att_block_bf16_in = *reinterpret_cast<attn_tile<bf16, col_l, rt_16x32_4_s>*>(&att_block_bf16);

    __builtin_amdgcn_sched_barrier(0);
    mul_col(o_reg, o_reg, scale_vec);
@@ -0,0 +1,113 @@
+#ifndef _TINYDRENO_H
+#define _TINYDRENO_H
+
+#include <stdint.h>
+#include <stddef.h>
+
+typedef void * cl_llvm_instance;
+
+cl_llvm_instance cl_compiler_create_llvm_instance(void);
+void cl_compiler_destroy_llvm_instance(cl_llvm_instance inst);
+
+enum cl_handle_type {
+  CL_HANDLE_COMPILED = 1,
+  CL_HANDLE_LIBRARY,
+  CL_HANDLE_LINKED
+};
+
+// handle->data for CL_HANDLE_COMPILED and CL_HANDLE_LIBRARY
+struct cl_compiled_data {
+  uint64_t chip_id;
+  uint32_t mode;
+  void    *llvm_bitcode;
+  uint64_t llvm_bitcode_size;
+  char    *build_log;
+  uint32_t build_log_len;
+  uint32_t error_code;
+};
+
+// handle->data for CL_HANDLE_LINKED
+struct cl_executable_data {
+  int32_t  num_kernels;
+  void    *kernel_props;
+  uint32_t error_code;
+  char    *build_log;
+  char     _unk0[0x20];
+  uint64_t chip_id;
+  uint32_t mode;
+};
+
+typedef struct {
+  enum cl_handle_type type;
+  union {
+    struct cl_compiled_data *compiled;
+    struct cl_executable_data *executable;
+  };
+} cl_handle;
+
+
+#define CL_MODE_32BIT 0
+#define CL_MODE_64BIT 1
+
+#define CL_SRC_STR  0
+#define CL_SRC_BLOB 1
+
+cl_handle *cl_compiler_compile_source(cl_llvm_instance inst, uint64_t chip_id, int mode, const char *options, int p5, uint64_t p6, uint64_t p7,
+                                      const char *source, uint64_t source_len, uint64_t source_type, void *p11);
+
+cl_handle *cl_compiler_link_program(cl_llvm_instance inst, uint64_t chip_id, int mode, const char *options, int num_handles,
+                                    cl_handle **input_handles);
+
+
+void cl_compiler_handle_create_binary(cl_handle *handle, void **out_ptr, size_t *out_size);
+
+// lib binary format (output of handle_create_binary for type 3)
+// layout: cl_lib_header, then cl_lib_section[num_sections], then data
+
+#define CL_LIB_PROGRAM     0
+#define CL_LIB_CONSTS      6
+#define CL_LIB_IMAGE       7
+#define CL_LIB_CODE       10
+#define CL_LIB_IMAGE_DESC 11
+
+typedef struct {
+  uint32_t id;
+  uint32_t offset;
+  uint32_t size;
+  uint32_t count;
+  uint32_t entry_size;
+} cl_lib_section;
+
+typedef struct {
+  uint32_t _unk0[6];
+  uint32_t num_sections;
+  uint32_t _unk1[5];
+  cl_lib_section sections[];
+} cl_lib_header;
+
+// at sections[CL_LIB_PROGRAM].offset
+typedef struct {
+  char     name[8];
+  uint32_t _unk0[3];
+  uint32_t fregs;
+  uint32_t hregs;
+} cl_lib_prog;
+
+// at sections[CL_LIB_IMAGE_DESC].offset
+typedef struct {
+  char     _unk0[0xc4];
+  uint32_t prg_offset;
+  uint32_t pvtmem;
+  char     _unk1[0x0c];
+  uint32_t shmem;
+  uint32_t samp_cnt;
+  char     _unk2[0x28];
+  uint32_t brnchstck;
+  char     _unk4[0x4c];
+  char     kernel_name[];
+} cl_lib_img_desc;
+
+void cl_compiler_free_handle(cl_handle *handle);
+void cl_compiler_free_assembly(void *ptr);
+
+#endif
@@ -505,7 +505,9 @@ tiny_backend_out = {**{f"aten.{x}.out":getattr(Tensor,x) for x in simple_tensor_
  "aten.lt.Tensor_out": Tensor.__lt__, "aten.lt.Scalar_out": Tensor.__lt__,
  "aten.le.Tensor_out": Tensor.__le__, "aten.le.Scalar_out": Tensor.__le__,
  "aten.clamp_max.Tensor_out": lambda input,max_: input.clamp(max_=max_),
+  "aten.clamp_max.out": lambda input,max_: input.clamp(max_=max_),
  "aten.clamp_min.Tensor_out": lambda input,min_: input.clamp(min_=min_),
+  "aten.clamp_min.out": lambda input,min_: input.clamp(min_=min_),
  "aten.fmod.Tensor_out": lambda input,other: input-input.div(other, rounding_mode="trunc")*other,
  # TODO: this might result in overflow issues
  "aten.round.decimals_out": lambda self,decimals: (self*10**decimals).round()/10**decimals,
@@ -588,7 +590,7 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
  "aten.repeat": lambda x,*repeats: Tensor.repeat(x,*repeats).contiguous(), # not a view
  "aten._softmax": lambda self,dim,half_to_float: self.softmax(dim),
  "aten._log_softmax": lambda self,dim,half_to_float: self.log_softmax(dim),
-  "aten.random_": lambda self: Tensor.randint(*self.shape, low=dtypes.min(self.dtype), high=dtypes.max(self.dtype), device=self.device, dtype=self.dtype),
+  "aten.random_": lambda self: Tensor.randint(*self.shape, low=self.dtype.min, high=self.dtype.max, device=self.device, dtype=self.dtype),
  "aten.random_.from": lambda self, from_, to: Tensor.randint(*self.shape, low=from_, high=to, device=self.device, dtype=self.dtype),
  "aten.uniform_": lambda self, low=0, high=1: Tensor.uniform(*self.shape, low=low, high=high, dtype=self.dtype),
  "aten.normal_": lambda self, mean=0, std=1: Tensor.normal(*self.shape, mean=mean, std=std, dtype=self.dtype),
@@ -24,7 +24,7 @@ if __name__ == "__main__":
  kernel_count = GlobalCounters.kernel_count
  assert kernel_count > 0, "No kernels, test failed"
  # NOTE: this is 124 on torch 2.10.0
-  expected_kernels = 332
+  expected_kernels = 334
  expectation = f"ResNet18 kernels are {kernel_count} vs {expected_kernels} expected."
  if kernel_count < expected_kernels: warnings.warn(f"{expectation} Expectation can be lowered.", UserWarning)
-  assert kernel_count <= expected_kernels, f"{expectation}"
+  assert kernel_count <= expected_kernels, f"{expectation}"
@@ -1,7 +1,6 @@
 # simple tests
 import unittest
 import torch
-import warnings
 from tinygrad.helpers import getenv, GlobalCounters
 if getenv("TINY_BACKEND2"):
  import extra.torch_backend.backend2
@@ -18,15 +17,13 @@ class TestKernelFusionRegression(unittest.TestCase):
    torch.manual_seed(42)
    GlobalCounters.reset()
    fn().detach().cpu().numpy()
-    expectation = f"{GlobalCounters.kernel_count} vs {expected_kernels} expected."
-    if GlobalCounters.kernel_count < expected_kernels: warnings.warn(f"{expectation} Expectation can be lowered.", UserWarning)
-    self.assertLessEqual(GlobalCounters.kernel_count, expected_kernels, f"{expectation}")
+    self.assertEqual(GlobalCounters.kernel_count, expected_kernels)

  def test_elementwise_fusion(self):
    def fn():
      x = torch.randn(128, 128, device=device)
      return (x + 1.0) * 2.0 - 0.5
-    self._check_kernel_count(fn, 5)
+    self._check_kernel_count(fn, 7)

  def test_relu_fusion(self):
    def fn():
@@ -44,26 +41,26 @@ class TestKernelFusionRegression(unittest.TestCase):
      bn.eval()
      with torch.no_grad():
        return torch.nn.functional.relu(bn(conv(x)))
-    self._check_kernel_count(fn, 16)
+    self._check_kernel_count(fn, 12)

  def test_reduce_fusion(self):
    def fn():
      x = torch.randn(64, 64, device=device)
      return (x * 2.0).sum()
-    self._check_kernel_count(fn, 5)
+    self._check_kernel_count(fn, 7)

  def test_matmul_elementwise_fusion(self):
    def fn():
      x = torch.randn(32, 32, device=device)
      w = torch.randn(32, 32, device=device)
      return torch.nn.functional.relu(x @ w + 1.0)
-    self._check_kernel_count(fn, 7)
+    self._check_kernel_count(fn, 9)

  def test_pooling_fusion(self):
    def fn():
      x = torch.randn(1, 8, 16, 16, device=device)
      return torch.nn.functional.max_pool2d(x * 2.0, 2)
-    self._check_kernel_count(fn, 5)
+    self._check_kernel_count(fn, 7)

  def test_residual_add_relu_fusion(self):
    def fn():
@@ -71,7 +68,7 @@ class TestKernelFusionRegression(unittest.TestCase):
      identity = torch.randn(1, 8, 16, 16, device=device)
      out = x + identity
      return torch.nn.functional.relu(out)
-    self._check_kernel_count(fn, 7)
+    self._check_kernel_count(fn, 9)

  def test_inplace_add_relu_fusion(self):
    def fn():
@@ -79,7 +76,7 @@ class TestKernelFusionRegression(unittest.TestCase):
      y = torch.randn(1, 16, 32, 32, device=device)
      x += y
      return torch.nn.functional.relu(x)
-    self._check_kernel_count(fn, 7)
+    self._check_kernel_count(fn, 9)

  def test_conv_bn_add_relu_fusion(self):
    def fn():
@@ -92,7 +89,7 @@ class TestKernelFusionRegression(unittest.TestCase):
        out = bn(conv(x))
        out += identity
        return torch.nn.functional.relu(out)
-    self._check_kernel_count(fn, 17)
+    self._check_kernel_count(fn, 14)

  def test_multiple_inplace_ops_fusion(self):
    def fn():
@@ -100,7 +97,7 @@ class TestKernelFusionRegression(unittest.TestCase):
      x += 1.0
      x *= 2.0
      return torch.nn.functional.relu(x)
-    self._check_kernel_count(fn, 4)
+    self._check_kernel_count(fn, 6)

  def test_view_inplace_no_fusion_break(self):
    def fn():
@@ -108,7 +105,7 @@ class TestKernelFusionRegression(unittest.TestCase):
      view = x[1:3]
      view += 1.0
      return x.sum()
-    self._check_kernel_count(fn, 8)
+    self._check_kernel_count(fn, 10)

  def test_batchnorm_running_stats_update(self):
    def fn():
@@ -138,7 +135,7 @@ class TestKernelFusionRegression(unittest.TestCase):
      loss.backward()
      optimizer.step()
      return loss
-    self._check_kernel_count(fn, 28)
+    self._check_kernel_count(fn, 26)

 if __name__ == "__main__":
  unittest.main()
@@ -17,6 +17,7 @@
 // Protocol

 enum {
+  CMD_PROBE = 0,          // probe devices, returns count
  CMD_MAP_BAR = 1,        // map PCI BAR, returns size
  CMD_MAP_SYSMEM_FD = 2,  // alloc DMA memory, returns fd via SCM_RIGHTS
  CMD_CFG_READ = 3,       // read PCI config space
@@ -24,11 +25,15 @@ enum {
  CMD_RESET = 5,          // reset device
  CMD_MMIO_READ = 6,      // bulk read from BAR
  CMD_MMIO_WRITE = 7,     // bulk write to BAR
+  CMD_MAP_SYSMEM = 8,     // map system memory
+  CMD_SYSMEM_READ = 9,    // bulk read from system memory
+  CMD_SYSMEM_WRITE = 10,  // bulk write to system memory
+  CMD_RESIZE_BAR = 11,    // resize bar (noop)
  RESP_OK = 0, RESP_ERR = 1,
 };

-typedef struct { uint8_t cmd, bar; uint64_t offset, size, value; } __attribute__((packed)) request_t;
-typedef struct { uint8_t status; uint64_t value, addr; } __attribute__((packed)) response_t;
+typedef struct { uint8_t cmd; uint32_t dev_id, bar; uint64_t arg0, arg1, arg2; } __attribute__((packed)) request_t;
+typedef struct { uint8_t status; uint64_t resp0, resp1; } __attribute__((packed)) response_t;

 // Constants and state

@@ -77,7 +82,7 @@ static int send_response(int fd, response_t *resp, int send_fd) {
 }

 static void send_error(int fd, const char *msg) {
-  response_t resp = {.status = RESP_ERR, .value = strlen(msg)};
+  response_t resp = {.status = RESP_ERR, .resp0 = strlen(msg)};
  send_response(fd, &resp, -1);
  send(fd, msg, strlen(msg), 0);
 }
@@ -116,12 +121,12 @@ static int dext_rpc(uint32_t sel, uint64_t *in, uint32_t in_cnt, uint64_t *out_v
 static int map_bar(uint32_t bar, response_t *resp) {
  if (bar >= MAX_BARS) return -1;
  if (!g_bars[bar].addr && IOConnectMapMemory64(g_conn, bar, mach_task_self(), &g_bars[bar].addr, &g_bars[bar].size, kIOMapAnywhere)) return -1;
-  resp->addr = g_bars[bar].addr;
-  resp->value = g_bars[bar].size;
+  resp->resp0 = g_bars[bar].addr;
+  resp->resp1 = g_bars[bar].size;
  return 0;
 }

-static int map_sysmem_fd(uint64_t size, response_t *resp, int *out_fd) {
+static int map_sysmem_fd(uint64_t size, int contiguous, response_t *resp, int *out_fd) {
  if (g_sysmem_count >= MAX_SYSMEM) return -1;
  int idx = g_sysmem_count;
  int fd = -1;
@@ -148,7 +153,7 @@ static int map_sysmem_fd(uint64_t size, response_t *resp, int *out_fd) {
  strncpy(g_sysmem[idx].shm_name, shm_name, sizeof(g_sysmem[idx].shm_name));
  g_sysmem_count++;

-  *resp = (response_t){.addr = idx, .value = alloc_sz};
+  *resp = (response_t){.resp0 = alloc_sz, .resp1 = idx};
  *out_fd = fd;
  return 0;

@@ -203,39 +208,42 @@ static void handle_client(int fd) {

    case CMD_MAP_SYSMEM_FD: {
      int shm_fd = -1;
-      resp.status = map_sysmem_fd(req.size, &resp, &shm_fd) ? 1 : 0;
+      resp.status = map_sysmem_fd(req.arg0, (int)req.arg1, &resp, &shm_fd) ? 1 : 0;
      send_response(fd, &resp, shm_fd);
      continue;
    }

    case CMD_CFG_READ: {
-      uint64_t in[2] = {req.offset, req.size};
-      resp.status = dext_rpc(0, in, 2, &resp.value) ? 1 : 0;
+      uint64_t in[2] = {req.arg0, req.arg1};
+      resp.status = dext_rpc(0, in, 2, &resp.resp0) ? 1 : 0;
      break;
    }

    case CMD_CFG_WRITE: {
-      uint64_t in[3] = {req.offset, req.size, req.value};
+      uint64_t in[3] = {req.arg0, req.arg1, req.arg2};
      resp.status = dext_rpc(1, in, 3, NULL) ? 1 : 0;
      break;
    }

+    case CMD_RESIZE_BAR:
+      break;
+
    case CMD_RESET:
      resp.status = dext_rpc(2, NULL, 0, NULL) ? 1 : 0;
      break;

    case CMD_MMIO_READ:
-      if (validate_bar(req.bar, req.offset, req.size)) { resp.status = 1; break; }
-      mmio_copy(g_bulk_buf, (void*)(g_bars[req.bar].addr + req.offset), req.size);
-      resp.value = req.size;
+      if (validate_bar(req.bar, req.arg0, req.arg1)) { resp.status = 1; break; }
+      mmio_copy(g_bulk_buf, (void*)(g_bars[req.bar].addr + req.arg0), req.arg1);
+      resp.resp0 = req.arg1;
      send_response(fd, &resp, -1);
-      send(fd, g_bulk_buf, req.size, 0);
+      send(fd, g_bulk_buf, req.arg1, 0);
      continue;

    case CMD_MMIO_WRITE:
-      recvall(fd, g_bulk_buf, req.size);
-      if (!validate_bar(req.bar, req.offset, req.size))
-        mmio_copy((void*)(g_bars[req.bar].addr + req.offset), g_bulk_buf, req.size);
+      recvall(fd, g_bulk_buf, req.arg1);
+      if (!validate_bar(req.bar, req.arg0, req.arg1))
+        mmio_copy((void*)(g_bars[req.bar].addr + req.arg0), g_bulk_buf, req.arg1);
      continue;

    default:
@@ -0,0 +1,17 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>com.apple.application-identifier</key>
+	<string>9YG3G8543N.org.tinygrad.tinygpu.edriver</string>
+	<key>com.apple.developer.driverkit</key>
+	<true/>
+	<key>com.apple.developer.driverkit.transport.pci</key>
+	<array>
+		<dict>
+			<key>IOPCIPrimaryMatch</key>
+			<string>0x000010de&amp;0x0000FFFF</string>
+		</dict>
+	</array>
+</dict>
+</plist>
@@ -0,0 +1,33 @@
+#!/bin/bash
+set -e
+
+xcodebuild clean build CODE_SIGN_IDENTITY="" CODE_SIGNING_REQUIRED=NO -alltargets -configuration Release build
+
+cp "../profiles/edriver_rel_2.provisionprofile" "./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext/embedded.provisionprofile"
+cp "../profiles/installer_provisioning.provisionprofile" "./build/Release/TinyGPU.app/Contents/embedded.provisionprofile"
+
+codesign \
+    --sign "Developer ID Application: tinygrad, Corp. (9YG3G8543N)" \
+    --entitlements ./TinyGPUDriverExtension/TinyGPUDriver.NV.Release.entitlements \
+    --verbose \
+    --options runtime \
+    --timestamp \
+    --force \
+    ./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext
+
+codesign \
+    --sign "Developer ID Application: tinygrad, Corp. (9YG3G8543N)" \
+    --entitlements ./macOS/macOS.entitlements \
+    --options runtime \
+    --verbose \
+    --timestamp \
+    --force \
+    ./build/Release/TinyGPU.app
+
+codesign --verify --deep --strict --verbose=4 ./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext
+
+codesign --verify --deep --strict --verbose=4 ./build/Release/TinyGPU.app
+
+spctl -a -vv ./build/Release/TinyGPU.app
+
+spctl -a -vv ./build/Release/TinyGPU.app/Contents/Library/SystemExtensions/org.tinygrad.tinygpu.edriver.dext
@@ -3,3 +3,7 @@ set -e

 ditto -c -k --keepParent ./build/Release/TinyGPU.app ./build/Release/TinyGPU.zip
 xcrun notarytool submit ./build/Release/TinyGPU.zip --keychain-profile "hgwJFhdheiIEy82nDN" --wait
+
+rm ./build/Release/TinyGPU.zip
+xcrun stapler staple ./build/Release/TinyGPU.app
+ditto -c -k --keepParent ./build/Release/TinyGPU.app ./build/Release/TinyGPU.zip
@@ -1,17 +1,18 @@
 A command line tool for exploring the VIZ trace.

-After running with VIZ=-1, use `PYTHONPATH=. extra/viz/cli.py` to explore the saved trace files.
+1. Set VIZ to -1 to save the trace.
+2. Use `extra/viz/cli.py` to inspect the trace files.

 ## Inspect runtime profiling

-Use `PYTHONPATH=. extra/viz/cli.py --profile` to list all traced devices.
+Use `extra/viz/cli.py --profile` to list all traced devices.

 List top slowest kernels on a device: `--profile --device "AMD"`
 List samples of a kernel on a device: `--profile --device "AMD" --kernel E_3`

 ## Inspect codegen and PatternMatcher

-Use `PYTHONPATH=. extra/viz/cli.py --rewrites` to list all traced kernels.
+Use `extra/viz/cli.py --rewrites` to list all traced kernels.

 List all codegen steps for a kernel: `--rewrites --kernel E_3`
 Get source code: `--rewrites --kernel E_3 --select "View Source"`
@@ -1,44 +1,79 @@
 #!/usr/bin/env python3
-import os
-os.environ["VIZ"] = "0"
-import argparse, pathlib
+import argparse, pathlib, sys, struct, json, itertools
 from typing import Iterator
 from tinygrad.viz import serve as viz
 from tinygrad.uop.ops import RewriteTrace
 from tinygrad.helpers import temp, ansistrip, colored, time_to_str, ansilen
-from test.null.test_viz import load_profile
+
+# ** generic helpers

 def optional_eq(val:dict, arg:str|None) -> bool: return arg is None or ansistrip(val["name"]) == arg

 def print_data(data:dict) -> None:
  if isinstance(data.get("value"), Iterator):
    for m in data["value"]:
-      if m.get("uop"):
-        print("Input UOp:")
-        print(m["uop"])
-      if not m["diff"]: continue
-      print("Rewrites:")
-      fp = pathlib.Path(m["upat"][0][0])
-      print(f"{fp.parent.name}/{fp.name}:{m['upat'][0][1]}")
-      print(m["upat"][1])
-      for line in m["diff"]:
-        color = "red" if line.startswith("-") else "green" if line.startswith("+") else None
-        print(colored(line, color))
+      if m.get("uop"): print(f"Input UOp:\n{m['uop']}")
+      if m.get("diff"):
+        loc = pathlib.Path(m["upat"][0][0])
+        print(f"Rewrite at {loc.parent.name}/{loc.name}:{m['upat'][0][1]}\n{m['upat'][1]}")
+        for line in m["diff"]: print(colored(line, "red" if line.startswith("-") else "green" if line.startswith("+") else None))
  if data.get("src") is not None: print(data["src"])

+# ** Profiler trace decoder
+
+# 0 means None, otherwise it's an enum value
+def option(i:int) -> int|None: return None if i == 0 else i-1
+
+def decode_profile(data:bytes) -> dict:
+  ret, off = data, 0
+  def u(fmt:str) -> tuple:
+    nonlocal off
+    vals = struct.unpack_from(fmt, ret, off)
+    off += struct.calcsize(fmt)
+    return vals
+  total_dur, global_peak, index_len, layout_len = u("<IQII")
+  strings, dtypes, markers = json.loads(ret[off:off+index_len]).values()
+  off += index_len
+  layout:dict[str, dict] = {}
+  for _ in range(layout_len):
+    klen = u("<B")[0]
+    k = ret[off:off+klen].decode()
+    off += klen
+    layout[k] = v = {"events":[]}
+    event_type, event_count = u("<BI")
+    if event_type == 0:
+      for _ in range(event_count):
+        name, ref, key, st, dur, fmt = u("<IIIIfI")
+        v["events"].append({"name":strings[name], "ref":option(ref), "key":option(key), "st":st, "dur":dur, "fmt":strings[fmt]})
+    else:
+      v["linear"] = u("<B")[0]
+      v["peak"] = u("<Q")[0]
+      for _ in range(event_count):
+        if v["linear"]:
+          ts, value = u("<IQ")
+          v["events"].append({"event":"freq", "ts":ts, "value":value})
+        else:
+          alloc, ts, key = u("<BII")
+          if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
+          else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IIIB") for _ in range(u("<I")[0])]}})
+  return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}
+
 if __name__ == "__main__":
  parser = argparse.ArgumentParser()
  g_mode = parser.add_argument_group("mode")
  g_mode.add_argument("--profile", action="store_true", help="View profile trace")
  g_mode.add_argument("--rewrites", action="store_true", help="View rewrites trace")
+  g_common = parser.add_argument_group("common options")
+  g_common.add_argument("--kernel", type=str, default=None, metavar="NAME", help="Select a kernel by name (optional name, default: only list names)")
+  g_common.add_argument("--no-color", action="store_true", default=not (sys.stdin.isatty() and sys.stdout.isatty()),
+                        help="Disable colored output (default: true in non-interactive mode)")
  g_profile = parser.add_argument_group("profile options")
  g_profile.add_argument("--device", type=str, default=None, metavar="NAME", help="Select a device (optional name, default: only list names)")
-  g_profile.add_argument("--top", type=int, default=10, metavar="N", help="Number of top kernels to show (-1 for all, default: 10)")
+  g_profile.add_argument("--offset", type=int, default=0, metavar="N", help="event offset (default: 0)")
+  g_profile.add_argument("--limit", type=int, default=10, metavar="N", help="events to display (-1 for all, default: 10)")
  g_rewrites = parser.add_argument_group("rewrites options")
  g_rewrites.add_argument("--select", type=str, default=None, metavar="NAME",
                          help="Select an item within the chosen kernel (optional name, default: only list names)")
-  g_common = parser.add_argument_group("common options")
-  g_common.add_argument("--kernel", type=str, default=None, metavar="NAME", help="Select a kernel by name (optional name, default: only list names)")
  parser.add_argument("--profile-path", type=pathlib.Path, metavar="PATH", help="Path to profile (optional file, default: latest profile)",
                        default=pathlib.Path(temp("profile.pkl", append_user=True)))
  parser.add_argument("--rewrites-path", type=pathlib.Path, metavar="PATH", help="Path to rewrites (optional file, default: latest rewrites)",
@@ -46,27 +81,77 @@ if __name__ == "__main__":
  args = parser.parse_args()
  if not args.profile and not args.rewrites:
    parser.print_help()
-    exit(0)
+    sys.exit(0)

  viz.trace = viz.load_pickle(args.rewrites_path, default=RewriteTrace([], [], {}))
  viz.ctxs = viz.get_rewrites(viz.trace)

+  def format_colored(s:str) -> str: return ansistrip(s) if args.no_color else s
+
  if args.profile:
    from tabulate import tabulate
-    profile = load_profile(viz.load_pickle(args.profile_path, default=[]))
+    profile = decode_profile(viz.get_profile(profile_data:=viz.load_pickle(args.profile_path, default=[])))
+    viz.load_amd_counters(viz.ctxs, profile_data)
+    counters = {f'{c["name"]} SQTT {s["name"]}': s["data"] for c in viz.ctxs if c["name"].startswith("Exec") for s in c["steps"]
+                if s["name"].startswith("PKTS")}
+    if args.device is None:
+      print("Select a device:")
+      for k in (*profile["layout"], *counters):
+        print(f"  {format_colored(k)}")
+      sys.exit(0)
+
+    # ** SQTT printer
+    if args.device is not None and (sqtt_data:=next((v for k,v in counters.items() if ansistrip(k) == args.device), None)) is not None:
+      assert args.limit > 1, f"SQTT limit must be greater than 1, got {args.limit}"
+      sqtt_events, has_more = viz.sqtt_timeline(*sqtt_data, max_pkts=args.offset+args.limit)
+      sqtt_pkts = [e for e in sqtt_events if type(e).__name__ == "ProfileRangeEvent"]
+      pc_map = next((e.arg for e in sqtt_events if type(e).__name__ == "ProfilePointEvent" and e.key == 'pcMap'), None)
+      if pc_map is None:
+        print(f"No SQTT instruction trace data for {args.device}")
+        sys.exit(0)
+      # modern terminals support 24-bit color
+      def hex_colored(st:str, color:str) -> str: return f"\x1b[38;2;{int(color[1:3],16)};{int(color[3:5],16)};{int(color[5:7],16)}m{st}\x1b[0m"
+      WAVE_COLORS = ((('VALU', 'VINTERP'), '#ffffc0'), (('SALU',), '#cef263'), (('VMEM',), '#b2b7c9'), (('LOAD', 'SMEM'), '#ffc0c0'),
+                     (('STORE',), '#4fa3cc'), (('IMMEDIATE',), '#f3b44a'), (('BARRIER',), '#d00000'), (('LDS',), '#9fb4a6'), (('JUMP',), '#ffb703'),
+                     (('JUMP_NO',), '#fb8500'), (('MESSAGE',), '#90dbf4'), (('WAVERDY',), '#1a2a2a'))
+      print(f"{'Clk':<12} {'Unit':<20} {'Op':<22} {'Dur':<4} {'Info'}")
+      print("-" * 90)
+      # start from the first packet in trace, prepare packet indexes and map dispatches
+      pkt_idxs:dict[str, itertools.count] = {}
+      dispatch_to_pc:dict[str, int] = {}
+      for e in sqtt_pkts[:-args.limit]:
+        idx = next(pkt_idxs.setdefault(e.device, itertools.count()))
+        if e.name.ret is not None and e.name.ret.startswith("PC:"): dispatch_to_pc[f"{e.device}-{idx}"] = int(e.name.ret.replace("PC:", ""))
+      # start printing from the offset point
+      for e in sqtt_pkts[-args.limit:]:
+        op_name, info = e.name.display_name, e.name.ret or ""
+        color = next((c for p, c in WAVE_COLORS if any(x in op_name for x in p)), None)
+        op_str = hex_colored(op_name, color) if color and not args.no_color else op_name
+        phase, pc = None, None
+        idx = next(pkt_idxs.setdefault(e.device, itertools.count()))
+        if info.startswith("PC:"):
+          dispatch_to_pc[f"{e.device}-{idx}"] = pc = int(info.replace("PC:", ""))
+          phase = "DISPATCH"
+        if info.startswith("LINK:"): phase, pc = "EXEC", dispatch_to_pc[info.replace("LINK:", "")]
+        if pc and phase: info = f"{phase:<8} 0x{pc:05x} {pc_map[pc]}"
+        print(f"{int(e.st):<12} {e.device:<20} {op_str}{' '*(22-ansilen(op_str))} {int(e.en-e.st):<4} {info}")
+      # note: we only print the important packets and skip the rest
+      if has_more: print(f"Selected packets {args.offset:,}-{args.offset + args.limit:,}. Use --offset and --limit to see others")
+      sys.exit(0)
+
+    # ** Profiler printer
    agg, total, n = {}, 0, 0
-    if args.device is None: print("Select a device:")
    for k,v in profile["layout"].items():
      if not optional_eq({"name":k}, args.device): continue
-      print(f"  {k}")
+      print(f"  {format_colored(k)}")
      if args.device is None: continue
      for e in v.get("events", []):
        et = e["dur"]*1e-6
        if args.kernel is not None:
-          if ansistrip(e["name"]) == args.kernel and n < 10:
+          if optional_eq(e, args.kernel) and n < 10:
            ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
            name = e["name"]+(" " * (46 - ansilen(e["name"])))
-            print(f"{name} {ptm}/{(et or 0)*1e3:9.2f}ms  "+e['fmt'].replace('\n', ' | ')+"  ")
+            print(f"{name} {ptm}/{(et or 0)*1e3:9.2f}ms  "+e.get('fmt', '').replace('\n', ' | ')+"  ")
            n += 1
        else:
          a = agg.setdefault(e["name"], [0.0, 0])
@@ -75,14 +160,15 @@ if __name__ == "__main__":
          total += et
    if agg and total > 0:
      items = sorted(agg.items(), key=lambda kv:kv[1][0], reverse=True)
-      sel = items if args.top == -1 else items[:args.top]
+      sel = items if args.limit == -1 else items[args.offset:args.offset+args.limit]
      table = [[name, time_to_str(t, w=9), c, f"{(t/total*100.0):.2f}%"] for name,(t,c) in sel]
-      if args.top != -1 and (other:=items[len(sel):]):
+      if args.limit != -1 and (other:=items[len(sel):]):
        other_t = total-sum(t for _, (t, _) in sel)
        table.append([f"Other ({len(other)} unique)", time_to_str(other_t, w=9), sum(c for _,(_,c) in other), f"{other_t/total*100.0:.2f}%"])
      print(tabulate(table, headers=["name", "total", "count", "pct"], tablefmt="github"))
-    exit(0)
+    sys.exit(0)

+  # ** Graph rewrites printer
  for k in viz.ctxs:
    if not optional_eq(k, args.kernel): continue
    print(k["name"])
@@ -74,7 +74,7 @@ testing_minimal = [
  "hypothesis>=6.148.9",
  "z3-solver<4.15.4",  # 4.15.4 has a segfault when creating many z3.Context()
 ]
-testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate", "openai", "gguf"]
+testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate", "openai", "gguf>=0.18"]
 testing = [
  "tinygrad[testing_unit]",
  "pillow",
@@ -324,6 +324,12 @@ def _disasm_smem(inst: SMEM) -> str:
  if name in ('s_memrealtime', 's_memtime'): return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}"
  return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}, {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (getattr(inst, 'dlc', 0), " dlc"))

+R4_TH_LOAD = {1: 'TH_LOAD_NT', 2: 'TH_LOAD_HT', 3: 'TH_LOAD_LU', 4: 'TH_LOAD_RT_WB', 5: 'TH_LOAD_NT_WB'}
+R4_TH_STORE = {1: 'TH_STORE_NT', 2: 'TH_STORE_HT', 3: 'TH_STORE_ST', 4: 'TH_STORE_RT_WB', 5: 'TH_STORE_NT_WB'}
+R4_TH_ATOMIC = {1: 'TH_ATOMIC_RETURN', 2: 'TH_ATOMIC_NT', 3: 'TH_ATOMIC_RETURN_NT',
+                4: 'TH_ATOMIC_CASCADE_RT', 5: 'TH_ATOMIC_CASCADE_RETURN', 6: 'TH_ATOMIC_CASCADE_NT', 7: 'TH_ATOMIC_CASCADE_RETURN_NT'}
+R4_SCOPE = {1: 'SCOPE_SE', 2: 'SCOPE_DEV', 3: 'SCOPE_SYS'}
+
 def _disasm_flat(inst: FLAT) -> str:
  name, cdna, r4 = inst.op_name.lower(), _is_cdna(inst), _is_r4(inst)
  acc = getattr(inst, 'acc', 0)
@@ -331,9 +337,10 @@ def _disasm_flat(inst: FLAT) -> str:
  if r4: seg = 'flat' if (cls_name:=inst.__class__.__name__) == 'VFLAT' else ('global' if cls_name == 'VGLOBAL' else 'scratch')
  else: seg = ['flat', 'scratch', 'global'][inst.seg] if inst.seg < 3 else 'flat'
  instr = f"{seg}_{name.split('_', 1)[1] if '_' in name else name}"
-  # Global/scratch uses 13-bit signed offset
+  # Global/scratch uses 13-bit signed offset (RDNA3/CDNA), 24-bit signed offset (RDNA4)
  offset = inst.ioffset if r4 else inst.offset  # type: ignore[attr-defined]
-  if seg != 'flat':
+  if r4: off_val = offset if offset < (1 << 23) else offset - (1 << 24)  # sign extend 24-bit
+  elif seg != 'flat':
    if cdna:
      # CDNA: bit 12 is sign bit but not in offset field
      raw = int.from_bytes(inst.to_bytes(), 'little')
@@ -348,7 +355,9 @@ def _disasm_flat(inst: FLAT) -> str:
  w = regs.get('data', regs.get('d', 1)) if 'store' in name or 'atomic' in name else regs.get('d', 1)
  off_s = f" offset:{off_val}" if off_val else ""
  if cdna: mods = f"{off_s}{' sc0' if inst.sc0 else ''}{' nt' if inst.nt else ''}{' sc1' if getattr(inst, 'sc1', 0) else ''}"  # type: ignore[attr-defined]
-  elif r4: mods = f"{off_s}{' scope' if inst.scope else ''}{' th' if inst.th else ''}"  # type: ignore[attr-defined]
+  elif r4:
+    th_names = R4_TH_ATOMIC if 'atomic' in name else (R4_TH_STORE if 'store' in name else R4_TH_LOAD)
+    mods = off_s + (f" th:{th_names[inst.th]}" if inst.th in th_names else "") + (f" scope:{R4_SCOPE[inst.scope]}" if inst.scope in R4_SCOPE else "")
  else: mods = f"{off_s}{' glc' if inst.glc else ''}{' slc' if inst.slc else ''}{' dlc' if inst.dlc else ''}"
  if seg == 'flat': saddr_s = ""
  elif _unwrap(inst.saddr) in (0x7F, 124): saddr_s = ", off"
@@ -357,7 +366,7 @@ def _disasm_flat(inst: FLAT) -> str:
    saddr_s = f", {(SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS)[_unwrap(inst.saddr)]}"
  elif t := _ttmp(inst.saddr, 2): saddr_s = f", {t}"
  else: saddr_s = f", {_sreg(inst.saddr, 2) if _unwrap(inst.saddr) < 106 else decode_src(_unwrap(inst.saddr), cdna)}"
-  if 'addtid' in name: return f"{instr} {reg_fn(inst.data if 'store' in name else inst.vdst)}{saddr_s}{mods}"
+  if 'addtid' in name: return f"{instr} {reg_fn((inst.vsrc if r4 else inst.data) if 'store' in name else inst.vdst)}{saddr_s}{mods}"
   # RDNA4: vaddr instead of addr, vsrc instead of data
  addr = inst.vaddr if r4 else inst.addr  # type: ignore[attr-defined]
  data = inst.vsrc if r4 else inst.data  # type: ignore[attr-defined]
@@ -372,7 +381,7 @@ def _disasm_flat(inst: FLAT) -> str:
  addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(addr, addr_w)
  data_s, vdst_s = reg_fn(data, w), reg_fn(inst.vdst, w // 2 if 'cmpswap' in name else w)
  if 'atomic' in name:
-    glc_or_sc0 = inst.sc0 if cdna else inst.glc  # type: ignore[attr-defined]
+    glc_or_sc0 = inst.sc0 if cdna else (inst.th & 1 if r4 else inst.glc)  # type: ignore[attr-defined]
    sfx = f"{saddr_s if seg != 'flat' else ''}{mods}"
    return f"{instr} {vdst_s}, {addr_s}, {data_s}{sfx}" if glc_or_sc0 else f"{instr} {addr_s}, {data_s}{sfx}"
  if 'store' in name: return f"{instr} {addr_s}, {data_s}{saddr_s}{mods}"
@@ -5,8 +5,8 @@ from tinygrad.runtime.autogen import llvm
 from tinygrad.runtime.support.elf import elf_loader

 ARCH_TO_TARGET:dict[str, list[str]] = {
-  "rdna3":["gfx1100"],
-  "rdna4":["gfx1200"],
+  "rdna3":["gfx1100", "gfx1151"],
+  "rdna4":["gfx1200", "gfx1201"],
  "cdna":["gfx950", "gfx942"],
 }

@@ -117,6 +117,58 @@ class TestDS2AddrMore(unittest.TestCase):
    self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
    self.assertEqual(st.vgpr[0][4], 0xDEADBEEF, "v4 should be untouched")

+  def test_ds_load_2addr_b64_addr_overlaps_vdst(self):
+    """DS_LOAD_2ADDR_B64 where addr register overlaps vdst range.
+
+    Hardware reads the address before writing any results, so addr=v[4]
+    with vdst=v[4:7] must load all 4 dwords using the original v[4] value.
+    """
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 0xAAAAAAAA),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
+      s_mov_b32(s[2], 0xBBBBBBBB),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=4),
+      s_mov_b32(s[2], 0xCCCCCCCC),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=8),
+      s_mov_b32(s[2], 0xDDDDDDDD),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=12),
+      s_waitcnt(lgkmcnt=0),
+      # addr=v[4] overlaps vdst=v[4:7]
+      v_mov_b32_e32(v[4], 0),
+      DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[4], vdst=v[4:7], offset0=0, offset1=1),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA, "v4 = LDS[0:4]")
+    self.assertEqual(st.vgpr[0][5], 0xBBBBBBBB, "v5 = LDS[4:8]")
+    self.assertEqual(st.vgpr[0][6], 0xCCCCCCCC, "v6 = LDS[8:12]")
+    self.assertEqual(st.vgpr[0][7], 0xDDDDDDDD, "v7 = LDS[12:16]")
+
+  def test_ds_load_2addr_b32_addr_overlaps_vdst(self):
+    """DS_LOAD_2ADDR_B32 where addr register overlaps vdst range."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 0xAAAAAAAA),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
+      s_mov_b32(s[2], 0xBBBBBBBB),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=4),
+      s_waitcnt(lgkmcnt=0),
+      # addr=v[2] overlaps vdst=v[2:3]
+      v_mov_b32_e32(v[2], 0),
+      DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[2], vdst=v[2:3], offset0=0, offset1=1),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA, "v2 = LDS[0:4]")
+    self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB, "v3 = LDS[4:8]")
+
  def test_ds_load_b64_no_overwrite(self):
    """DS_LOAD_B64 should only write 2 VGPRs."""
    instructions = [
@@ -759,6 +811,47 @@ class TestDsPermute(unittest.TestCase):
    # Lane 0 receives data (highest numbered active lane wins)
    self.assertEqual(st.vgpr[0][2], 0x11111111)

+  def test_ds_bpermute_b32_xor_swap(self):
+    """DS_BPERMUTE_B32 with XOR-1 pattern — each lane reads from lane^1.
+
+    This is the pattern used by warp_shfl_xor in flash attention for reduce_max/reduce_sum.
+    Each lane has a unique value (lane_id + 100), and reads from the adjacent lane.
+    """
+    instructions = [
+      # v[0] = (lane_id ^ 1) * 4  (byte offset for bpermute)
+      v_xor_b32_e32(v[0], 1, v[255]),
+      v_lshlrev_b32_e32(v[0], 2, v[0]),
+      # v[1] = lane_id + 100 (unique per-lane value)
+      s_mov_b32(s[0], 100),
+      v_add_nc_u32_e32(v[1], s[0], v[255]),
+      # ds_bpermute: v[2] = v[1] from lane (lane_id ^ 1)
+      ds_bpermute_b32(vdst=v[2], addr=v[0], data0=v[1]),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=32)
+    for lane in range(32):
+      src_lane = lane ^ 1
+      expected = src_lane + 100
+      self.assertEqual(st.vgpr[lane][2], expected, f"lane {lane}: expected v[1] from lane {src_lane} = {expected}, got {st.vgpr[lane][2]}")
+
+
+class TestDSSubDword(unittest.TestCase):
+  """Tests for sub-dword DS operations (ds_store_b16, ds_store_b16_d16_hi)."""
+
+  def test_ds_store_b16_and_d16_hi(self):
+    """DS_STORE_B16 stores low 16 bits, DS_STORE_B16_D16_HI stores high 16 bits to adjacent LDS half-words."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0),
+      v_mov_b32_e32(v[1], 0xBEEF1234),
+      DS(DSOp.DS_STORE_B16, addr=v[0], data0=v[1], offset0=0),
+      DS(DSOp.DS_STORE_B16_D16_HI, addr=v[0], data0=v[1], offset0=2),
+      s_waitcnt(lgkmcnt=0),
+      ds_load_b32(vdst=v[2], addr=v[0], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xBEEF1234, "lo=0x1234 at byte 0, hi=0xBEEF at byte 2")
+

 class TestDSLargeOffset(unittest.TestCase):
  """Tests for DS instructions with offsets > 255 (offset1 > 0).
@@ -408,6 +408,23 @@ class TestVOP3P(unittest.TestCase):
    self.assertEqual(lo, 0x0005, f"lo: expected 0x0005, got 0x{lo:04x}")
    self.assertEqual(hi, 0x4003, f"hi: expected 0x4003, got 0x{hi:04x}")

+  def test_v_pk_add_u16_literal_constant(self):
+    """V_PK_ADD_U16 with a literal constant (value > 64, requires VOP3P_LIT encoding).
+    Regression test: VOP3P literal constants were not passed to rsrc_dyn, so literal src read as 0.
+    """
+    instructions = [
+      s_mov_b32(s[0], 0x1C001C00),  # packed u16: hi=0x1C00, lo=0x1C00 (f16 for 2^-8)
+      v_mov_b32_e32(v[0], s[0]),
+      v_pk_add_u16(v[1], 0x2000, v[0], opsel_hi=2, opsel_hi2=1),  # add 0x2000 bias to both halves
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][1]
+    lo = result & 0xffff
+    hi = (result >> 16) & 0xffff
+    # lo = 0x1C00 + 0x2000 = 0x3C00 (f16 1.0), hi = 0x1C00 + 0x2000 = 0x3C00 (f16 1.0)
+    self.assertEqual(lo, 0x3C00, f"lo: expected 0x3C00, got 0x{lo:04x}")
+    self.assertEqual(hi, 0x3C00, f"hi: expected 0x3C00, got 0x{hi:04x}")
+

 class TestWMMAF16(unittest.TestCase):
  """Tests for WMMA F16 output variant (V_WMMA_F16_16X16X16_F16).
@@ -104,6 +104,34 @@ class TestCmpClass(unittest.TestCase):
    st = run_program(instructions, n_lanes=1)
    self.assertEqual(st.vcc & 1, 0, "Signaling NaN should not match quiet mask")

+  def test_v_cmp_lg_f32_nan(self):
+    """v_cmp_lg_f32 is ordered not-equal (<>): NaN <> x should be False per IEEE 754."""
+    quiet_nan = 0x7fc00000
+    one_f32 = 0x3f800000  # 1.0f
+    instructions = [
+      s_mov_b32(s[0], quiet_nan),
+      v_mov_b32_e32(v[0], s[0]),
+      s_mov_b32(s[1], one_f32),
+      v_mov_b32_e32(v[1], s[1]),
+      v_cmp_lg_f32_e32(v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vcc & 1, 0, "v_cmp_lg_f32(NaN, 1.0) should be 0")
+
+  def test_v_cmp_neq_f32_nan(self):
+    """v_cmp_neq_f32 is unordered not-equal (!=): NaN != x should be True per IEEE 754."""
+    quiet_nan = 0x7fc00000
+    one_f32 = 0x3f800000  # 1.0f
+    instructions = [
+      s_mov_b32(s[0], quiet_nan),
+      v_mov_b32_e32(v[0], s[0]),
+      s_mov_b32(s[1], one_f32),
+      v_mov_b32_e32(v[1], s[1]),
+      v_cmp_neq_f32_e32(v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vcc & 1, 1, "v_cmp_neq_f32(NaN, 1.0) should be 1")
+
  def test_v_cmp_sets_vcc_bits(self):
    """V_CMP_EQ sets VCC bits based on per-lane comparison."""
    instructions = [
@@ -138,6 +138,47 @@ class TestVOPDLiterals(unittest.TestCase):
    self.assertAlmostEqual(i2f(st.vgpr[0][2]), 13.0, places=5, msg="fma(2.0, 5.0, 3.0) should be 13.0")


+class TestVOPDDot2Acc(unittest.TestCase):
+  """Tests for V_DUAL_DOT2ACC_F32_F16 - packed f16 dot product accumulate."""
+
+  def test_vopd_dot2acc_f32_f16_basic(self):
+    """V_DUAL_DOT2ACC_F32_F16: D += lo(S0)*lo(S1) + hi(S0)*hi(S1).
+
+    S0 = pack(1.0h, 2.0h), S1 = pack(3.0h, 4.0h), D = 10.0f
+    result = 10.0 + 1.0*3.0 + 2.0*4.0 = 10.0 + 3.0 + 8.0 = 21.0
+    """
+    from test.amd.hw.helpers import f2i, i2f, f32_to_f16
+    pk_s0 = f32_to_f16(1.0) | (f32_to_f16(2.0) << 16)  # lo=1.0h, hi=2.0h
+    pk_s1 = f32_to_f16(3.0) | (f32_to_f16(4.0) << 16)  # lo=3.0h, hi=4.0h
+    instructions = [
+      v_mov_b32_e32(v[0], pk_s0),
+      v_mov_b32_e32(v[1], pk_s1),
+      v_mov_b32_e32(v[3], f2i(10.0)),  # accumulator in v[3] (vdsty with vdstx=v[4])
+      # X: v[4] = MOV v[0] (don't care), Y: v[3] += dot2(v[0], v[1])
+      VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_DOT2ACC_F32_F16, v[4], v[3], v[0], v[0], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertAlmostEqual(i2f(st.vgpr[0][3]), 21.0, places=2, msg="10.0 + 1.0*3.0 + 2.0*4.0 = 21.0")
+
+  def test_vopd_dot2acc_f32_f16_zero_accum(self):
+    """V_DUAL_DOT2ACC_F32_F16 with zero accumulator — pure dot product.
+
+    S0 = pack(0.5h, -1.0h), S1 = pack(2.0h, 3.0h), D = 0.0f
+    result = 0.0 + 0.5*2.0 + (-1.0)*3.0 = 1.0 - 3.0 = -2.0
+    """
+    from test.amd.hw.helpers import f2i, i2f, f32_to_f16
+    pk_s0 = f32_to_f16(0.5) | (f32_to_f16(-1.0) << 16)
+    pk_s1 = f32_to_f16(2.0) | (f32_to_f16(3.0) << 16)
+    instructions = [
+      v_mov_b32_e32(v[0], pk_s0),
+      v_mov_b32_e32(v[1], pk_s1),
+      v_mov_b32_e32(v[3], f2i(0.0)),  # zero accumulator in v[3]
+      VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_DOT2ACC_F32_F16, v[4], v[3], v[0], v[0], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertAlmostEqual(i2f(st.vgpr[0][3]), -2.0, places=2, msg="0.5*2.0 + (-1.0)*3.0 = -2.0")
+
+
 class TestVOPDMultilane(unittest.TestCase):
  """Tests for VOPD with multiple lanes."""

@@ -1,10 +1,14 @@
 import unittest
+import functools
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
-
+from tinygrad.dtype import AddrSpace
 from tinygrad.runtime.autogen.amd.rdna3.ins import *
+import tinygrad.runtime.autogen.amd.rdna3.ins as r3
+import tinygrad.runtime.autogen.amd.rdna4.ins as r4
 from tinygrad.renderer.amd.dsl import s, v
+from test.amd.helpers import TARGET_TO_ARCH

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

+def custom_wave_sync(A:UOp, arch:str) -> UOp:
+  # 4 waves across 1024 WG — enough to saturate a SIMD with many concurrent WGs
+  # s_sleep yields the SIMD so waves from different WGs interleave, causing barrier packet reordering
+  threads = UOp.special(128, "lidx0")
+  wg = UOp.special(1024, "gidx0")
+  insts = []
+  for _ in range(4):
+    insts.append(s_sleep(4))
+    insts += [s_barrier()] if arch == "rdna3" else [r4.s_barrier_signal(), r4.s_barrier_wait()]
+    insts += [s_nop(0)]*4
+  insts.append(s_endpgm())
+  sink = UOp.sink(A.base, threads, wg, arg=KernelInfo("custom_wave_sync"))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+
+def custom_lds_sync(A:UOp, arch:str) -> UOp:
+  A = A.flatten()
+  num_threads = A.shape[0]
+  threads = UOp.special(num_threads, "lidx0")
+  wg = UOp.special(1, "gidx0")
+  lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=512, addrspace=AddrSpace.LOCAL), (), 'lds')  # 128 * 4 bytes
+  isa = r4 if arch == "rdna4" else r3
+  wait_kmcnt = [isa.s_wait_kmcnt(simm16=0)] if arch == "rdna4" else [isa.s_waitcnt(lgkmcnt=0)]
+  wait_dscnt = [isa.s_wait_dscnt(simm16=0)] if arch == "rdna4" else [isa.s_waitcnt(lgkmcnt=0)]
+  barrier = [isa.s_barrier_signal(ssrc0=-1), isa.s_barrier_wait(simm16=-1)] if arch == "rdna4" else [isa.s_barrier()]
+  global_store = [isa.global_store_b32(vaddr=v[6:7], saddr=s[0:1], vsrc=v[5])] if arch == "rdna4" \
+      else [isa.global_store_b32(addr=v[6], data=v[5], saddr=s[0:1])]
+  insts = [
+    isa.s_load_b64(s[0:1], s[0:1], soffset=NULL),
+    *wait_kmcnt,
+    isa.v_lshlrev_b32_e32(v[1], 2, v[0]),
+    # lds[thread_idx] = thread_idx
+    isa.ds_store_b32(addr=v[1], data0=v[0]),
+    *wait_dscnt,
+    *barrier,
+    # out[threaed_idx] = thread_idx == num_threads ? -1 : lds[thread_idx + 1]
+    isa.v_add_nc_u32_e32(v[2], 4, v[1]),
+    isa.v_cmp_gt_u32_e32(num_threads-1, v[0]),
+    isa.ds_load_b32(vdst=v[3], addr=v[2]),
+    *wait_dscnt,
+    isa.v_mov_b32_e32(v[4], -1),
+    isa.v_cndmask_b32_e32(v[5], v[4], v[3]),
+    isa.v_lshlrev_b32_e32(v[6], 2, v[0]),
+    *global_store,
+    isa.s_endpgm(),
+  ]
+  sink = UOp.sink(A.base, lds, threads, wg, arg=KernelInfo("custom_lds_sync"))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+
@unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
 class TestCustomKernel(unittest.TestCase):
+  def setUp(self): self.arch = TARGET_TO_ARCH[Device["AMD"].arch]
+
  def test_simple(self):
+    if self.arch != "rdna3": self.skipTest("only rdna3")
    a = Tensor.full((16, 16), 1.).contiguous().realize()
    a = Tensor.custom_kernel(a, fxn=custom_add_one)[0]
    ei = a.schedule()[-1].lower()
@@ -55,6 +110,7 @@ class TestCustomKernel(unittest.TestCase):
    self.assertTrue((a.numpy() == 2.).all())

  def test_variable(self):
+    if self.arch != "rdna3": self.skipTest("only rdna3")
    b = Tensor.full((16, 16), 1, dtype=dtypes.uint32).contiguous().realize()
    a = Tensor.zeros_like(b).contiguous().realize()
    a = Tensor.custom_kernel(a, b, fxn=custom_add_var)[0]
@@ -63,5 +119,14 @@ class TestCustomKernel(unittest.TestCase):
      ei.run({"var":i})
      self.assertTrue((a.numpy() == 1+i).all())

+  def test_lds_sync(self):
+    if self.arch not in ("rdna3", "rdna4"): self.skipTest("only rdna3/rdna4")
+    a = Tensor.empty(128, dtype=dtypes.int32).contiguous().realize()
+    a = Tensor.custom_kernel(a, fxn=functools.partial(custom_lds_sync, arch=self.arch))[0]
+    a.realize()
+    ref = Tensor.arange(1, 129, dtype=dtypes.int32)
+    ref[127] = -1
+    self.assertListEqual(a.tolist(), ref.tolist())
+
 if __name__ == "__main__":
  unittest.main()
@@ -40,7 +40,7 @@ RDNA4_FILES = ['gfx12_asm_sop1.s', 'gfx12_asm_sop2.s', 'gfx12_asm_sopp.s', 'gfx1
  'gfx12_asm_vop1.s', 'gfx12_asm_vop2.s', 'gfx12_asm_vopc.s', 'gfx12_asm_vopcx.s', 'gfx12_asm_vop3.s', 'gfx12_asm_vop3c.s',
  'gfx12_asm_vop3cx.s', 'gfx12_asm_vop3p.s', 'gfx12_asm_vop3_from_vop1.s', 'gfx12_asm_vop3_from_vop2.s',
  'gfx12_asm_vop3p_features.s', 'gfx12_asm_vopd.s', 'gfx12_asm_vopd_features.s',
-  'gfx12_asm_ds.s', 'gfx12_asm_smem.s',
+  'gfx12_asm_ds.s', 'gfx12_asm_smem.s', 'gfx12_asm_vflat.s',
  'gfx12_asm_wmma_w32.s']

 def _parse_llvm_tests(text: str, pattern: str) -> list[tuple[str, bytes]]:
@@ -9,19 +9,12 @@ from tinygrad.renderer.amd import decode_inst
 from tinygrad.runtime.autogen.amd.rdna3.ins import SOPP
 from tinygrad.runtime.autogen.amd.rdna3.enum import SOPPOp
 from tinygrad.renderer.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_RDNA4, WAVEEND, INST, INST_RDNA4, VALUINST,
-                                     IMMEDIATE, IMMEDIATE_MASK, PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4,
-                                     InstOp, InstOpRDNA4, print_packets)
+                                     IMMEDIATE, IMMEDIATE_MASK, PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4, PACKET_TYPES_CDNA, CDNA_WAVESTART,
+                                     print_packets, CDNA_WAVEEND, CDNA_INST)
 from test.amd.helpers import TARGET_TO_ARCH

 import tinygrad
 EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"
-# INST ops for non-traced SIMDs (excluded from instruction count)
-OTHER_SIMD_OPS = {InstOp.OTHER_LDS_LOAD, InstOp.OTHER_LDS_STORE, InstOp.OTHER_LDS_STORE_64, InstOp.OTHER_LDS_STORE_128,
-                  InstOp.OTHER_FLAT_LOAD, InstOp.OTHER_FLAT_STORE, InstOp.OTHER_FLAT_STORE_64, InstOp.OTHER_FLAT_STORE_96,
-                  InstOp.OTHER_FLAT_STORE_128, InstOp.OTHER_GLOBAL_LOAD, InstOp.OTHER_GLOBAL_LOAD_VADDR,
-                  InstOp.OTHER_GLOBAL_STORE_64, InstOp.OTHER_GLOBAL_STORE_96, InstOp.OTHER_GLOBAL_STORE_128,
-                  InstOp.OTHER_GLOBAL_STORE_VADDR_128}
-OTHER_SIMD_OPS_RDNA4 = {InstOpRDNA4.OTHER_VMEM, InstOpRDNA4.OTHER_VMEM_STORE}

 # ═══════════════════════════════════════════════════════════════════════════════
 # ROCPROF DECODER
@@ -125,7 +118,7 @@ class SQTTExamplesTestBase(unittest.TestCase):
          self.assertIsInstance(packets[0], LAYOUT_HEADER, f"first packet should be LAYOUT_HEADER in {name}")

  def test_packet_types_valid(self):
-    all_classes = set(PACKET_TYPES_RDNA3.values()) | set(PACKET_TYPES_RDNA4.values())
+    all_classes = set(PACKET_TYPES_RDNA3.values()) | set(PACKET_TYPES_RDNA4.values()) | set(PACKET_TYPES_CDNA.values())
    for name, (events, *_) in self.examples.items():
      for i, event in enumerate(events):
        with self.subTest(example=name, event=i):
@@ -138,8 +131,8 @@ class SQTTExamplesTestBase(unittest.TestCase):
      if "empty" in name: continue
      with self.subTest(example=name):
        all_packets = [p for e in events for p in decode(e.blob)]
-        self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVESTART, WAVESTART_RDNA4))]), 0, f"no WAVESTART in {name}")
-        self.assertGreater(len([p for p in all_packets if isinstance(p, WAVEEND)]), 0, f"no WAVEEND in {name}")
+        self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVESTART, WAVESTART_RDNA4, CDNA_WAVESTART))]), 0, f"no WAVESTART in {name}")
+        self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVEEND, CDNA_WAVEEND))]), 0, f"no WAVEEND in {name}")

  def test_time_monotonic(self):
    for name, (events, *_) in self.examples.items():
@@ -153,7 +146,10 @@ class SQTTExamplesTestBase(unittest.TestCase):
      if "gemm" not in name: continue
      with self.subTest(example=name):
        all_packets = [p for e in events for p in decode(e.blob)]
-        self.assertGreater(len([p for p in all_packets if isinstance(p, (INST, INST_RDNA4))]), 0, f"no INST packets in {name}")
+        inst_packets = [p for p in all_packets if isinstance(p, (INST, INST_RDNA4, CDNA_INST))]
+        self.assertGreater(len(inst_packets), 0, f"no INST packets in {name}")
+        if isinstance(inst_packets[0], (INST, INST_RDNA4)):
+          self.assertGreater(len([p for p in inst_packets if p.op.name.startswith("JUMP")]), 0, f"no JUMP packets in {name}")

  expected: dict[str, list[int]] = {}  # override in subclasses
  def test_packet_counts(self):
@@ -180,11 +176,18 @@ class SQTTExamplesTestBase(unittest.TestCase):
        our_waves: list[tuple[int, int]] = []
        for event in events:
          wave_starts: dict[tuple[int, int, int], int] = {}
+          first_timestamp:int|None = None
          for p in decode(event.blob):
-            if isinstance(p, (WAVESTART, WAVESTART_RDNA4)): wave_starts[(p.wave, p.simd, p.cu)] = p._time
-            elif isinstance(p, WAVEEND) and (key := (p.wave, p.simd, p.cu)) in wave_starts:
+            if first_timestamp is None: first_timestamp = p._time
+            if isinstance(p, (WAVESTART, CDNA_WAVESTART, WAVESTART_RDNA4)): wave_starts[(p.wave, p.simd, p.cu)] = p._time
+            elif isinstance(p, (WAVEEND, CDNA_WAVEEND)) and (key := (p.wave, p.simd, p.cu)) in wave_starts:
              our_waves.append((wave_starts[key], p._time))
-        self.assertEqual(sorted(our_waves), sorted(roc_waves), f"wave times mismatch in {name}")
+          for st in wave_starts.values():
+            self.assertGreater(st, first_timestamp, "wave start must be after the first packet")
+        # rocprof fails non deterministically and gives inaccurate timestamps.
+        #self.assertEqual(sorted(our_waves), sorted(roc_waves), f"wave times mismatch in {name}")
+        for st, et in our_waves:
+          self.assertGreater(et, st, "wave end must be after start")

  def test_rocprof_inst_times_match(self):
    """Instruction times must match rocprof exactly (excluding s_endpgm)."""
@@ -197,8 +200,8 @@ class SQTTExamplesTestBase(unittest.TestCase):
        our_insts: list[int] = []
        for event in events:
          for p in decode(event.blob):
-            if isinstance(p, INST) and p.op not in OTHER_SIMD_OPS: our_insts.append(p._time)
-            elif isinstance(p, INST_RDNA4) and p.op not in OTHER_SIMD_OPS_RDNA4: our_insts.append(p._time)
+            # INST ops for non-traced SIMDs (excluded from instruction count)
+            if isinstance(p, (INST, INST_RDNA4)) and not p.op.name.startswith("OTHER_"): our_insts.append(p._time)
            elif isinstance(p, VALUINST): our_insts.append(p._time)
            elif isinstance(p, IMMEDIATE): our_insts.append(p._time)
            elif isinstance(p, IMMEDIATE_MASK):
@@ -208,17 +211,22 @@ class SQTTExamplesTestBase(unittest.TestCase):
 class TestSQTTExamplesRDNA3(SQTTExamplesTestBase):
  target = "gfx1100"
  expected = {
-    "profile_empty_run_0": [1744, 1801, 1854, 1890, 1917, 1822],
-    "profile_empty_run_1": [1744, 1801, 1854, 1886, 1921, 1906],
-    "profile_gemm_run_0": [1800, 1867, 1899, 1898, 1914, 1895, 1694, 1779, 1819, 1872, 1877, 1858, 1750, 1834, 1866, 1834, 1911, 1796],
-    "profile_gemm_run_1": [1806, 1874, 1837, 1885, 1907, 1906, 1694, 1778, 1810, 1873, 1885, 1867, 1750, 1834, 1866, 1856, 1903, 1897],
-    "profile_plus_run_0": [1744, 1878, 1854, 1890, 1878, 1910],
-    "profile_plus_run_1": [1744, 1878, 1854, 1886, 1921, 1909],
+    "profile_empty_run_0": [1880, 1867, 1920, 1971, 1998, 1904],
+    "profile_empty_run_1": [1880, 1867, 1920, 1971, 1998, 1904],
+    "profile_gemm_run_0": [3275, 3278, 2426, 2475, 2511, 2431],
+    "profile_gemm_run_1": [3264, 3268, 2420, 2469, 2504, 2401],
+    "profile_ops_run_0": [1944, 4903, 1984, 2035, 2062, 1968],
+    "profile_ops_run_1": [1944, 4918, 1984, 2035, 2062, 1968],
+    "profile_plus_run_0": [1938, 1932, 1978, 2029, 2056, 1962],
+    "profile_plus_run_1": [1891, 1874, 1931, 1982, 2009, 1915],
  }

 class TestSQTTExamplesRDNA4(SQTTExamplesTestBase): target = "gfx1200"
-@unittest.skip("TODO: fix CDNA")
-class TestSQTTExamplesCDNA(SQTTExamplesTestBase): target = "gfx950"
+
+class TestSQTTExamplesCDNA(SQTTExamplesTestBase):
+  target = "gfx950"
+  def test_rocprof_wave_times_match(self): self.skipTest("TODO: requires timestamp patching")
+  def test_rocprof_inst_times_match(self): self.skipTest("TODO: requires timestamp patching")

 if __name__ == "__main__":
  unittest.main()
@@ -1,183 +0,0 @@
-"""Tests comparing sqtt.py PACKET_TYPES_RDNA3/RDNA4 against AMD's rocprof-trace-decoder binary."""
-import unittest, struct, ctypes, pickle
-from pathlib import Path
-
-ROCPROF_LIB = Path("/usr/lib/librocprof-trace-decoder.so")
-import tinygrad
-EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"
-
-# CDNA pkt_fmt -> size in bytes (extracted from rocprof hash table)
-CDNA_PKT_SIZES = {0: 2, 1: 8, 2: 8, 3: 4, 4: 2, 5: 6, 6: 2, 7: 2, 8: 2, 9: 2, 10: 2, 11: 8, 12: 6, 13: 4, 14: 8, 15: 6}
-
-def _find_segment(perms: str):
-  """Find a segment of the loaded library with given permissions (e.g. 'rw-p', 'r--p')."""
-  with open('/proc/self/maps', 'r') as f:
-    for line in f:
-      if 'librocprof-trace-decoder.so' in line and f' {perms} ' in line:
-        parts = line.split()
-        return int(parts[0].split('-')[0], 16), int(parts[2], 16)
-  return None, None
-
-def _read_array(file_offset: int, count: int):
-  """Read an array of uint8 at file_offset from the loaded library."""
-  base, seg_offset = _find_segment('rw-p')
-  if base is None: return None
-  return list((ctypes.c_uint8 * count).from_address(base + (file_offset - seg_offset)))
-
-def _load_lib():
-  if not ROCPROF_LIB.exists(): return False
-  ctypes.CDLL(str(ROCPROF_LIB))
-  return True
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# RDNA EXTRACTION (nibble-based format)
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def extract_bit_tables():
-  """Extract bit budget tables. Returns (layout2, layout3, layout4) or None."""
-  if not _load_lib(): return None
-  return _read_array(0x2d220, 32), _read_array(0x2d280, 32), _read_array(0x2d2c0, 32)
-
-def extract_delta_fields():
-  """Extract delta bitfield tables. Returns (layout2, layout3, layout4) dicts mapping type_id -> (lo, hi)."""
-  if not _load_lib(): return None
-  ro_base, ro_offset = _find_segment('r--p')
-  if ro_base is None: return None
-
-  def read_table(file_offset, num_entries):
-    addr = ro_base + (file_offset - ro_offset)
-    data = bytes((ctypes.c_uint8 * (num_entries * 12)).from_address(addr))
-    return {type_id: (lo, hi) for j in range(0, len(data), 12)
-            for type_id, lo, hi in [struct.unpack('<III', data[j:j+12])] if type_id < 32}
-
-  return read_table(0x26800, 24), read_table(0x26dc0, 25), read_table(0x27300, 27)
-
-def extract_packet_encodings():
-  """Extract packet encodings. Returns (L2, L3, L4) dicts mapping type_id -> (mask, value)."""
-  if not _load_lib(): return None
-  rw_base, rw_offset = _find_segment('rw-p')
-  if rw_base is None: return None
-
-  # Read base encodings from registration vector at 0x2d340
-  vec_start = ctypes.c_void_p.from_address(rw_base + (0x2d340 - rw_offset)).value
-  vec_end = ctypes.c_void_p.from_address(rw_base + (0x2d348 - rw_offset)).value
-  base = {}
-  if vec_start and vec_end:
-    for i in range((vec_end - vec_start) // 32):
-      addr = vec_start + i * 32
-      type_id = ctypes.c_uint8.from_address(addr).value
-      pat_start = ctypes.c_void_p.from_address(addr + 8).value
-      pat_end = ctypes.c_void_p.from_address(addr + 16).value
-      if pat_start and pat_end and 0 < (n := pat_end - pat_start) <= 8:
-        pat = list((ctypes.c_uint8 * n).from_address(pat_start))
-        base[type_id] = (sum(1 << j for j in range(n)), sum(b << j for j, b in enumerate(pat)))
-
-  return {**base, 17: (0x7f, 0x51), 25: (0x7f, 0x31)}, base, {**base}  # L2 has overrides
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# CDNA EXTRACTION (16-bit header format)
-# ═══════════════════════════════════════════════════════════════════════════════
-
-def extract_cdna_packet_sizes():
-  """Extract CDNA pkt_fmt -> size mapping by running rocprof decoder to populate its hash table."""
-  if not _load_lib(): return None
-  from test.amd.test_sqtt_examples import run_rocprof_decoder
-
-  if not (pkl_path := next((EXAMPLES_DIR / "gfx950").glob("*.pkl"), None)): return None
-  with open(pkl_path, "rb") as f: data = pickle.load(f)
-  sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
-  prg = next((e for e in data if type(e).__name__ == "ProfileProgramEvent"), None)
-  if not sqtt_events or not prg: return None
-
-  # Run decoder to trigger hash table initialization
-  run_rocprof_decoder([e.blob for e in sqtt_events], prg.lib, prg.base, "gfx950")
-
-  # Extract hash table: head at 0x2d4f0, nodes are 16 bytes (next[8], key[4], value[4])
-  rw_base, rw_offset = _find_segment('rw-p')
-  if not (head := ctypes.c_void_p.from_address(rw_base + (0x2d4f0 - rw_offset)).value if rw_base else None): return None
-
-  pkt_sizes: dict[int, int] = {}
-  node, seen = head, set()
-  while node and node not in seen and len(pkt_sizes) < 20:
-    seen.add(node)
-    key, val = ctypes.c_uint32.from_address(node + 8).value, ctypes.c_uint32.from_address(node + 12).value
-    if key < 16 and val in (0x10, 0x20, 0x30, 0x40): pkt_sizes[key] = {0x10: 2, 0x20: 4, 0x30: 6, 0x40: 8}[val]
-    node = ctypes.c_void_p.from_address(node).value  # type: ignore[assignment]
-  return pkt_sizes if len(pkt_sizes) == 16 else None
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# TESTS
-# ═══════════════════════════════════════════════════════════════════════════════
-
-class TestSQTTMatchesBinary(unittest.TestCase):
-  def test_bit_counts_match_layout3(self): self._test_bit_counts(3)
-  def test_bit_counts_match_layout4(self): self._test_bit_counts(4)
-  def test_encodings_match_layout3(self): self._test_encodings(3)
-  def test_encodings_match_layout4(self): self._test_encodings(4)
-  def test_delta_fields_match_layout3(self): self._test_delta_fields(3)
-  def test_delta_fields_match_layout4(self): self._test_delta_fields(4)
-
-  def test_cdna_packet_sizes(self):
-    """Extract and verify CDNA pkt_fmt -> size mapping from rocprof's hash table."""
-    if not (EXAMPLES_DIR / "gfx950").exists(): self.skipTest("no CDNA examples")
-    if not (pkt_sizes := extract_cdna_packet_sizes()): self.skipTest("rocprof-trace-decoder not installed")
-    for pkt_fmt, size in CDNA_PKT_SIZES.items():
-      with self.subTest(pkt_fmt=pkt_fmt): self.assertEqual(pkt_sizes.get(pkt_fmt), size)
-
-  def test_cdna_packet_definitions(self):
-    from tinygrad.renderer.amd.sqtt import PACKET_TYPES_CDNA
-    for pkt_fmt, pkt_cls in PACKET_TYPES_CDNA.items():
-      with self.subTest(packet=pkt_cls.__name__):
-        self.assertEqual(pkt_cls.encoding.default, pkt_fmt)
-        self.assertEqual(CDNA_PKT_SIZES[pkt_fmt] * 2, pkt_cls._size_nibbles)  # type: ignore[attr-defined]
-
-  def _test_bit_counts(self, layout: int):
-    if not (tables := extract_bit_tables()): self.skipTest("rocprof-trace-decoder not installed")
-    from tinygrad.renderer.amd.sqtt import PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4
-    for type_id, pkt_cls in {3: PACKET_TYPES_RDNA3, 4: PACKET_TYPES_RDNA4}[layout].items():
-      with self.subTest(packet=pkt_cls.__name__):
-        self.assertEqual(pkt_cls._size_nibbles * 4, tables[layout - 2][type_id])  # type: ignore[attr-defined]
-
-  def _test_encodings(self, layout: int):
-    if not (encodings := extract_packet_encodings()): self.skipTest("rocprof-trace-decoder not installed")
-    from tinygrad.renderer.amd.sqtt import PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4
-    for type_id, pkt_cls in {3: PACKET_TYPES_RDNA3, 4: PACKET_TYPES_RDNA4}[layout].items():
-      with self.subTest(packet=pkt_cls.__name__):
-        self.assertEqual((pkt_cls.encoding.mask, pkt_cls.encoding.default), encodings[layout - 2][type_id])
-
-  def _test_delta_fields(self, layout: int):
-    if not (deltas := extract_delta_fields()): self.skipTest("rocprof-trace-decoder not installed")
-    from tinygrad.renderer.amd.sqtt import PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4
-    for type_id, pkt_cls in {3: PACKET_TYPES_RDNA3, 4: PACKET_TYPES_RDNA4}[layout].items():
-      if type_id not in deltas[layout - 2]: continue
-      delta = getattr(pkt_cls, 'delta', None)
-      actual = (0, 0) if delta is None else (delta.lo, delta.hi + 1)
-      with self.subTest(packet=pkt_cls.__name__): self.assertEqual(actual, deltas[layout - 2][type_id])
-
-if __name__ == "__main__":
-  tables = extract_bit_tables()
-  encodings = extract_packet_encodings()
-  deltas = extract_delta_fields()
-
-  TYPE_NAMES = {1: 'VALUINST', 2: 'VMEMEXEC', 3: 'ALUEXEC', 4: 'IMMEDIATE', 5: 'IMMEDIATE_MASK', 6: 'WAVERDY',
-    7: 'TS_DELTA_S8_W3', 8: 'WAVEEND', 9: 'WAVESTART', 10: 'TS_DELTA_S5_W2', 11: 'WAVEALLOC', 12: 'TS_DELTA_S5_W3',
-    13: 'PERF', 14: 'UTILCTR', 15: 'TS_DELTA_SHORT', 16: 'NOP', 17: 'TS_WAVE_STATE', 18: 'EVENT', 19: 'EVENT_BIG',
-    20: 'REG', 21: 'SNAPSHOT', 22: 'TS_DELTA_OR_MARK', 23: 'LAYOUT_HEADER', 24: 'INST', 25: 'UNK_25'}
-
-  print("L2:", tables[0], "\nL3:", tables[1], "\nL4:", tables[2])
-  if encodings and tables:
-    print(f"\n{'TypeID':>6} {'Name':>18} {'L2 enc':>12} {'L3 enc':>12} {'L4 enc':>12}"
-          f" {'L2':>4} {'L3':>4} {'L4':>4} {'L2 delta':>12} {'L3 delta':>12} {'L4 delta':>12}")
-    print("-" * 140)
-    for type_id in sorted(set(encodings[0]) | set(encodings[1]) | set(encodings[2])):
-      name = TYPE_NAMES.get(type_id, f'UNK_{type_id}')
-      bits = [tables[i][type_id] if type_id < len(tables[i]) else 0 for i in range(3)]
-      enc_strs = [f"0x{encodings[i][type_id][0]:02x}/0x{encodings[i][type_id][1]:02x}" if type_id in encodings[i] else "-" for i in range(3)]
-      delta_strs = [f"[{d[1]-1}:{d[0]}]" if (d := deltas[i].get(type_id, (0, 0)))[1] > d[0] else "-" for i in range(3)]
-      print(f"{type_id:6d} {name:>18} {enc_strs[0]:>12} {enc_strs[1]:>12} {enc_strs[2]:>12}"
-            f" {bits[0]:4d} {bits[1]:4d} {bits[2]:4d} {delta_strs[0]:>12} {delta_strs[1]:>12} {delta_strs[2]:>12}")
-
-  cdna = extract_cdna_packet_sizes()
-  if cdna: print(f"\nCDNA packet sizes: {cdna}")
-
-  unittest.main()
@@ -2,9 +2,10 @@
 import unittest, pickle
 from typing import Iterator
 from pathlib import Path
-from tinygrad.helpers import DEBUG
+from tinygrad.helpers import DEBUG, getenv, temp
 from tinygrad.renderer.amd.sqtt import print_packets, map_insts
 from tinygrad.runtime.autogen.amd.rdna3.ins import s_endpgm
+from tinygrad.viz.serve import sqtt_timeline
 from test.amd.disasm import disasm

 import tinygrad
@@ -24,7 +25,7 @@ def rocprof_inst_traces_match(sqtt, prg, target):

  passed_insts = 0
  for pkt, info in map_insts(sqtt.blob, prg.lib, target):
-    if DEBUG >= 2: print_packets([pkt])
+    if DEBUG >= 2: print_packets([(pkt, info)])
    if info is None: continue
    if DEBUG >= 2: print(f"{' '*29}{disasm(info.inst)}")
    rocprof_inst = next(rwaves_iter[info.wave][0])
@@ -53,7 +54,7 @@ class TestSQTTMapBase(unittest.TestCase):
  def setUpClass(cls):
    if cls is TestSQTTMapBase: raise unittest.SkipTest("base class")
    cls.examples = {}
-    for pkl_path in sorted((EXAMPLES_DIR/cls.target).glob("*.pkl")):
+    for pkl_path in ([Path(temp("profile.pkl", append_user=True))] if getenv("LOAD_PROFILE") else sorted((EXAMPLES_DIR/cls.target).glob("*.pkl"))):
      with open(pkl_path, "rb") as f:
        data = pickle.load(f)
      sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
@@ -63,6 +64,8 @@ class TestSQTTMapBase(unittest.TestCase):

  def test_rocprof_inst_traces_match(self):
    for name, (events, kern_events, target) in self.examples.items():
+      if "sync" in name and self.target.startswith("gfx12"):
+        self.skipTest("our timestamps are off by a few cycles because rocprof patches timestamps for rdna4 barriers")
      for event in events:
        if not event.itrace: continue
        if event.kern not in kern_events: continue
@@ -70,9 +73,48 @@ class TestSQTTMapBase(unittest.TestCase):
          passed_insts, n_waves, n_units = rocprof_inst_traces_match(event, kern_events[event.kern], target)
          if n_waves: print(f"{name}: passed for {passed_insts} instructions across {n_waves} waves scheduled on {n_units} wave units")

+  def test_sqtt_timeline(self):
+    for name, (events, kern_events, target) in self.examples.items():
+      for event in events:
+        if (p:=kern_events.get(event.kern)) is None: continue
+        with self.subTest(example=name, kern=event.kern):
+          # skip if there's no SQTT frequency data
+          if not (timeline:=sqtt_timeline(event.blob, p.lib, target)[0]): continue
+          if not (frequency:=[e.key for e in timeline if type(e).__name__ == "ProfilePointEvent" and e.name == "freq_hz"]): continue
+          mean = sum(frequency) / len(frequency)
+          variance = sum((v - mean) ** 2 for v in frequency) / len(frequency)
+          self.assertGreater(mean, 0)
+          if DEBUG >= 2: print(f"{name:20s} SE:{event.se} {mean/1e9:.2f} GHz mean, {variance/1e18:.2f} GHz^2 variance")
+          events = [e for e in timeline if type(e).__name__ == "ProfileRangeEvent"]
+          insts, execs = 0, 0
+          for e in events:
+            if "EXEC" in e.device:
+              if "ALT" not in e.name.display_name: execs += 1
+            elif "WAVE" in e.device:
+              # sopk/immediates don't get ALU/MEM EXEC
+              if e.name.display_name not in {"IMMEDIATE", "IMMEDIATE_MASK", "JUMP", "JUMP_NO", "MESSAGE", "BARRIER", "BARRIER_SIGNAL",
+                                             "WAVEEND", "WAVERDY"}: insts += 1
+            else: raise Exception(f"timeline row must be INST or EXEC, got {e.device}")
+          self.assertEqual(execs, insts)
+
+  def test_wave_sync(self):
+    for name, (events, kern_events, target) in self.examples.items():
+      for event in events:
+        wave_barriers = {}
+        for e in sqtt_timeline(event.blob, kern_events[event.kern].lib, target)[0]:
+          if type(e).__name__ == "ProfileRangeEvent" and e.name.display_name == "BARRIER": wave_barriers.setdefault(e.device, []).append(e)
+        if not wave_barriers: continue
+        for row, events in wave_barriers.items():
+          for e in events:
+            assert e.en-e.st > 1, f"all barriers must have a duration greater than 1, got {e}"
+
 class TestSQTTMapRDNA3(TestSQTTMapBase): target = "gfx1100"

 class TestSQTTMapRDNA4(TestSQTTMapBase): target = "gfx1200"

+class TestSQTTMapCDNA(TestSQTTMapBase):
+  target = "gfx950"
+  def test_rocprof_inst_traces_match(self): self.skipTest("requires timestamp patching to match rocprof, currently it's off by a few cycles")
+
 if __name__ == "__main__":
  unittest.main()
@@ -8,6 +8,7 @@ from tinygrad.engine.schedule import ExecItem
 from tinygrad.uop.ops import Ops
 from tinygrad.renderer import Estimates
 from tinygrad.renderer.ptx import PTXRenderer
+from test.helpers import needs_second_gpu

 class TestArange(unittest.TestCase):
  def _get_flops(self, tensor, desired):
@@ -188,6 +189,33 @@ class TestIndexing(unittest.TestCase):
    for i in idx.flatten().numpy(): expected_grad[i] += 2
    np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)

+  @needs_second_gpu
+  @unittest.skipIf(Device.DEFAULT not in ("CPU", "AMD"), "atomics only on AMD/CPU")
+  @Context(USE_ATOMICS=1, SPEC=1)
+  def test_embedding_backward_vocab_sharded(self):
+    from tinygrad.renderer.cstyle import CStyleLanguage
+    if Device.DEFAULT == "CPU" and not isinstance(Device["CPU"].renderer, CStyleLanguage): self.skipTest("CPU needs Clang renderer")
+    devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
+    vocab_size, embed_size = 1000, 128
+    bs, seqlen = 4, 256
+    idx = Tensor.randint(bs, seqlen, high=vocab_size)
+    emb = nn.Embedding(vocab_size, embed_size)
+    emb.weight = Tensor.ones(vocab_size, embed_size, requires_grad=True)
+    gt = Tensor.zeros(bs, seqlen, embed_size)
+    Tensor.realize(idx, emb.weight, gt)
+    # compute expected grad on single device
+    expected_grad = np.zeros((vocab_size, embed_size), dtype=np.float32)
+    for i in idx.flatten().numpy(): expected_grad[i] += 2
+    # now shard the embedding weight on vocab axis and recompute
+    emb.weight = Tensor.ones(vocab_size, embed_size, requires_grad=True)
+    emb.weight.shard_(devices, axis=0)
+    idx = idx.shard(devices, axis=None)
+    gt = gt.shard(devices, axis=None)
+    Tensor.realize(idx, emb.weight, gt)
+    loss = (emb(idx)-gt).square().sum()
+    loss.backward()
+    np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)
+
  @unittest.skipUnless(Device.DEFAULT == "AMD" or (Device.DEFAULT == "NULL" and EMULATE.value.startswith("AMD")), "tests AMD bf16 cast overhead")
  def base_test_llama_8b_rope_backward(self, dtype):
    from extra.models.llama import precompute_freqs_cis, apply_rotary_emb
--- a/Show More
+++ b/Show More