broadcast

broadcast shape
2026-05-14 17:04:37 -07:00 · 2026-05-14 17:01:20 -07:00
257 changed files with 3704 additions and 5707 deletions
@@ -5,7 +5,6 @@ runs:
  steps:
    - name: Run process replay tests
      shell: bash
-      if: env.CAPTURE_PROCESS_REPLAY == '1'
      run: |
        export PR_TITLE=$(jq -r .pull_request.title "$GITHUB_EVENT_PATH")
        export CURRENT_SHA=${{ github.event.pull_request && github.event.pull_request.head.sha || github.sha }}
@@ -42,7 +42,7 @@ inputs:
    required: false
    default: 'false'
  mesa:
-    description: "Install mesa (true, false, cpu)"
+    description: "Install mesa"
    required: false
    default: 'false'
  tinydreno:
@@ -228,11 +228,6 @@ runs:

        sudo chown -R $USER:$USER /var/cache/apt/archives/

-    - name: Add clang to PATH (Linux)
-      if: inputs.llvm == 'true' && runner.os == 'Linux'
-      shell: bash
-      run: echo "/usr/lib/llvm-20/bin" >> "$GITHUB_PATH"
-
    # **** AMD ****
    - name: Setup AMD (Linux)
      if: inputs.amd == 'true' && runner.os == 'Linux'
@@ -286,7 +281,7 @@ runs:
        cache-name: cache-gpuocelot-build-1
      with:
        path: ${{ github.workspace }}/gpuocelot/ocelot
-        key: ${{ runner.os }}-gpuocelot-f463259669c69abce7b3a0567b6c284f348d0f32-rebuild-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
    - name: Cache gpuocelot
      if: inputs.ocelot == 'true' && github.event_name != 'pull_request'
      id: cache-build
@@ -295,14 +290,14 @@ runs:
        cache-name: cache-gpuocelot-build-1
      with:
        path: ${{ github.workspace }}/gpuocelot/ocelot
-        key: ${{ runner.os }}-gpuocelot-f463259669c69abce7b3a0567b6c284f348d0f32-rebuild-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
    - name: Clone/compile gpuocelot
      if: inputs.ocelot == 'true' && steps.cache-build-pr.outputs.cache-hit != 'true' && steps.cache-build.outputs.cache-hit != 'true'
      shell: bash
      run: |
-        git clone --recurse-submodules https://github.com/tinygrad/gpuocelot.git ${{ github.workspace }}/gpuocelot
+        git clone --recurse-submodules https://github.com/gpuocelot/gpuocelot.git ${{ github.workspace }}/gpuocelot
        cd ${{ github.workspace }}/gpuocelot/ocelot
-        git checkout f463259669c69abce7b3a0567b6c284f348d0f32
+        git checkout b16039dc940dc6bc4ea0a98380495769ff35ed99
        mkdir build
        cd build

@@ -311,7 +306,10 @@ runs:
          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"
        else
-          CMAKE_ARGS="$CMAKE_ARGS -DLLVM_DIR=$(llvm-config-15 --cmakedir)"
+          curl -fL https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvcc/linux-x86_64/cuda_nvcc-linux-x86_64-11.5.119-archive.tar.xz \
+            | sudo tar -xJ -C /usr/ --strip-components=1
+          curl -fL https://developer.download.nvidia.com/compute/cuda/redist/cuda_cudart/linux-x86_64/cuda_cudart-linux-x86_64-11.5.117-archive.tar.xz \
+            | sudo tar -xJ -C /usr/ --strip-components=1
        fi

        cmake .. $CMAKE_ARGS
@@ -347,13 +345,13 @@ runs:

    # **** mesa ****
    - name: Install mesa (linux)
-      if: inputs.mesa != 'false' && runner.os == 'Linux'
+      if: inputs.mesa == 'true' && runner.os == 'Linux'
      shell: bash
-      run: sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/v1/libtinymesa${{ inputs.mesa == 'cpu' && '_cpu' || '' }}-mesa-25.2.7-linux-amd64.so -o /usr/lib/libtinymesa${{ inputs.mesa == 'cpu' && '_cpu' || '' }}.so
+      run: sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/v1/libtinymesa_cpu-mesa-25.2.7-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
    - name: Install mesa (macOS)
-      if: inputs.mesa != 'false' && runner.os == 'macOS'
+      if: inputs.mesa == 'true' && runner.os == 'macOS'
      shell: bash
-      run: brew install sirhcm/tinymesa/tinymesa${{ inputs.mesa == 'cpu' && '_cpu' || '' }}
+      run: brew install sirhcm/tinymesa/tinymesa_cpu

    # *** tinydreno ***
    - name: Install tinydreno (linux)
@@ -25,7 +25,7 @@ jobs:
      CI: ""
      CAPTURE_PROCESS_REPLAY: "0"
    runs-on: [self-hosted, macOS]
-    timeout-minutes: 4
+    timeout-minutes: 3
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -625,7 +625,7 @@ jobs:
    - 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:IR3 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.2 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
+      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
@@ -633,7 +633,7 @@ jobs:
    - name: openpilot compile3 0.10.1 driving_vision
      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.10.1 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3.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_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=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
@@ -806,16 +806,3 @@ jobs:
        pkill -f 'extra/remote/serve.py' || true
    - name: Run process replay tests
      uses: ./.github/actions/process-replay
-
-  llvmspeed:
-    name: LLVM Speed
-    runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 20
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Speed Test
-      run: DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py
-    - name: Speed Test (BEAM=2)
-      run: BEAM=2 DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py
@@ -2,7 +2,7 @@ name: Unit Tests
 env:
  # increment this when downloads substantially change to avoid the internet
  CACHE_VERSION: '19'
-  CAPTURE_PROCESS_REPLAY: ${{ github.event_name == 'pull_request' && contains(github.event.pull_request.title, '[pr]') && '1' || '0' }}
+  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
  PYTHONPATH: ${{ github.workspace }}
  CHECK_OOB: 1
@@ -14,14 +14,28 @@ on:
  pull_request:
  workflow_dispatch:

-concurrency:
-  group: test-${{ github.event_name }}-${{ github.event_name == 'pull_request' && github.event.pull_request.number || github.run_id }}
-  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
-
 jobs:
+  llvmspeed:
+    name: LLVM Speed
+    runs-on: ubuntu-24.04
+    timeout-minutes: 20
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: llvm-speed
+        deps: testing_unit
+        llvm: 'true'
+    - name: Speed Test
+      run: DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py
+    - name: Speed Test (BEAM=2)
+      run: BEAM=2 DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py
+
  docs:
    name: Docs
-    runs-on: &linux ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: ubuntu-22.04
    timeout-minutes: 10
    env:
      CHECK_OOB: 0
@@ -75,7 +89,7 @@ jobs:

  torchbackend:
    name: Torch Backend Tests
-    runs-on: *linux
+    runs-on: ubuntu-latest
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -111,7 +125,7 @@ jobs:

  torchbackendmore:
    name: Torch Backend Tests More
-    runs-on: *linux
+    runs-on: ubuntu-latest
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -133,7 +147,7 @@ jobs:

  bepython:
    name: Python Backend
-    runs-on: *linux
+    runs-on: ubuntu-latest
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -201,7 +215,7 @@ jobs:

  linter:
    name: Linters
-    runs-on: *linux
+    runs-on: ubuntu-latest
    timeout-minutes: 10

    steps:
@@ -232,7 +246,7 @@ jobs:

  nulltest:
    name: Null Tests
-    runs-on: *linux
+    runs-on: ubuntu-latest
    timeout-minutes: 15

    steps:
@@ -263,7 +277,7 @@ jobs:

  unittest:
    name: Unit Tests
-    runs-on: *linux
+    runs-on: ubuntu-latest
    timeout-minutes: 15

    steps:
@@ -276,6 +290,7 @@ jobs:
        pydeps: "pillow ftfy regex pre-commit"
        deps: testing_unit
        llvm: 'true'
+        amd: 'true'
    - name: Run pre-commit test hooks
      run: SKIP=ruff,mypy pre-commit run --all-files
    - name: Check Device.DEFAULT
@@ -297,8 +312,8 @@ jobs:
        python extra/optimization/extract_dataset.py
        gzip -c /tmp/sops > extra/datasets/sops.gz
        #DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py
-    - name: Repo line count < 25000 lines
-      run: MAX_LINE_COUNT=25000 python sz.py
+    - name: Repo line count < 24000 lines
+      run: MAX_LINE_COUNT=24000 python sz.py

  spec:
    strategy:
@@ -306,7 +321,7 @@ jobs:
      matrix:
        group: [1, 2]
    name: SPEC=2 (${{ matrix.group }})
-    runs-on: *linux
+    runs-on: ubuntu-latest
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -317,13 +332,12 @@ jobs:
        key: spec-unit
        deps: testing_unit
        python-version: '3.14'
-        llvm: 'true'
    - name: Test SPEC=2
      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 test/unit test/backend test/opt --ignore test/backend/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" -k "not test_conv2d_ceildiv_edge_case" --splits 2 --group ${{ matrix.group }}

  fuzzing:
    name: Fuzzing
-    runs-on: *linux
+    runs-on: ubuntu-latest
    timeout-minutes: 10
    steps:
    - name: Checkout Code
@@ -344,7 +358,7 @@ jobs:

  testopenclimage:
    name: CL IMAGE Tests
-    runs-on: *linux
+    runs-on: ubuntu-22.04
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -364,7 +378,7 @@ jobs:

  testgpumisc:
    name: CL Misc tests
-    runs-on: *linux
+    runs-on: ubuntu-22.04
    timeout-minutes: 10
    steps:
      - name: Checkout Code
@@ -389,7 +403,7 @@ jobs:

  testopenpilot:
    name: openpilot Compile Tests
-    runs-on: *linux
+    runs-on: ubuntu-22.04
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -403,11 +417,11 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1468 ALLOWED_GATED_READ_IMAGE=18 FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1486 ALLOWED_GATED_READ_IMAGE=18 FLOAT16=1 DEV=CL 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 DEV=CL 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: |
-          DEV=CL IMAGE=1 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
-          DEV=CL IMAGE=1 SELFTEST=1 RUN_PICKLE=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
+        run: DEV=CL IMAGE=1 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
      - name: Test openpilot LLVM compile fp16
        run: IMAGE=1 FLOAT16=1 DEV=CPU:LLVM python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
      - name: Run process replay tests
@@ -417,7 +431,7 @@ jobs:

  testonnxcpu:
    name: ONNX (CPU) Tests
-    runs-on: *linux
+    runs-on: ubuntu-22.04
    timeout-minutes: 20

    steps:
@@ -445,7 +459,7 @@ jobs:

  testopencl:
    name: ONNX (CL)+Optimization Tests
-    runs-on: *linux
+    runs-on: ubuntu-22.04
    timeout-minutes: 20
    steps:
      - name: Checkout Code
@@ -479,7 +493,7 @@ jobs:

  testllm:
    name: Test LLM
-    runs-on: *linux
+    runs-on: ubuntu-24.04
    timeout-minutes: 15
    env:
      CHECK_OOB: 0
@@ -504,7 +518,7 @@ jobs:

  testmodels:
    name: Models (llvm+cpu+gpu)
-    runs-on: *linux
+    runs-on: ubuntu-22.04
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -545,9 +559,30 @@ jobs:

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

+  testdevectorize:
+    name: Linux (devectorize)
+    runs-on: ubuntu-24.04
+    timeout-minutes: 15
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: devectorize-minimal
+        deps: testing_unit
+        pydeps: "pillow"
+        llvm: "true"
+    - name: Test LLVM=1 DEVECTORIZE=0
+      run: DEV=CPU:LLVM DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/backend/test_ops.py
+    - name: Test LLVM=1 DEVECTORIZE=0 for model
+      run: DEV=CPU:LLVM DEVECTORIZE=0 python3 test/models/test_efficientnet.py
+    - name: Test DEV=CPU DEVECTORIZE=0
+      run: DEV=CPU DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/backend/test_ops.py
+
  testdsp:
    name: Linux (DSP)
-    runs-on: *linux
+    runs-on: ubuntu-24.04
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -556,7 +591,8 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: dsp-minimal
-        deps: testing
+        deps: testing_unit
+        pydeps: "onnx==1.18.0 onnxruntime ml_dtypes"
        llvm: "true"
        qemu: "true"
    - name: Set MOCKDSP env
@@ -564,24 +600,13 @@ jobs:
    - name: Run test_tiny on DSP
      run: DEBUG=2 DEV=DSP python test/test_tiny.py
    - name: Test transcendentals
-      run: DEBUG=2 DEV=DSP python test/backend/test_transcendental.py TestTranscendentalVectorized
+      run: CC=clang-20 DEBUG=2 DEV=DSP python test/backend/test_transcendental.py TestTranscendentalVectorized
    - name: Test quantize onnx
      run: DEBUG=2 DEV=DSP python3 test/backend/test_quantize_onnx.py

-  testlinux:
-    strategy:
-      fail-fast: false
-      matrix:
-        dev:
-          - 'CPU:CLANG'
-          - 'CPU:LLVM'
-          - 'CPU:LVP'
-          - 'CPU:X86'
-          - 'CL'
-          - 'WEBGPU'
-
-    name: Linux (DEV=${{ matrix.dev }})
-    runs-on: *linux
+  testwebgpu:
+    name: Linux (WebGPU)
+    runs-on: ubuntu-22.04
    timeout-minutes: 20
    steps:
    - name: Checkout Code
@@ -589,27 +614,23 @@ jobs:
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
-        key: linux-${{ matrix.dev }}
+        key: webgpu-minimal
        deps: testing_unit
        python-version: '3.12'
-        llvm: ${{ contains(matrix.dev, 'LLVM') || contains(matrix.dev, 'LVP') || contains(matrix.dev, 'CLANG') }}
-        mesa: ${{ contains(matrix.dev, 'LVP') && 'cpu' || 'false' }}
-        webgpu: ${{ matrix.dev == 'WEBGPU' }}
-        opencl: ${{ matrix.dev == 'CL' }}
-    - name: Set env
-      run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
-    - name: Check Device.DEFAULT and print some source
+        webgpu: 'true'
+    - name: Check Device.DEFAULT (WEBGPU) and print some source
      run: |
-        python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
-        DEBUG=4 python test/test_tiny.py TestTiny.test_plus
-    - name: Run backend tests
-      run: python -m pytest -n=auto test/backend --durations=20
+        DEV=WEBGPU python -c "from tinygrad import Device; assert Device.DEFAULT == 'WEBGPU', Device.DEFAULT"
+        DEV=WEBGPU DEBUG=4 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
+    - name: Run selected webgpu tests
+      run: |
+          DEV=WEBGPU WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/backend --durations=20
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

  testamdasm:
    name: AMD ASM IDE
-    runs-on: *linux
+    runs-on: ubuntu-24.04
    timeout-minutes: 20
    env:
      DEV: MOCKKFD+AMD
@@ -656,7 +677,7 @@ jobs:

  testmockam:
    name: Linux (am)
-    runs-on: *linux
+    runs-on: ubuntu-24.04
    timeout-minutes: 15
    env:
      DEV: MOCKPCI+AMD
@@ -692,7 +713,7 @@ jobs:
        arch: [gfx1100, gfx1201, gfx950]

    name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
-    runs-on: *linux
+    runs-on: ubuntu-22.04
    timeout-minutes: 15
    env:
      DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
@@ -727,7 +748,7 @@ jobs:
        backend: [ptx, nv]

    name: Linux (${{ matrix.backend }})
-    runs-on: *linux
+    runs-on: ubuntu-22.04
    timeout-minutes: 20
    env:
      FORWARD_ONLY: 1
@@ -755,6 +776,39 @@ jobs:
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

+  testcpuopencl:
+    strategy:
+      fail-fast: false
+      matrix:
+        backend: [llvm, cpu, opencl, lvp]
+
+    name: Linux (${{ matrix.backend }})
+    runs-on: ubuntu-22.04
+    timeout-minutes: 20
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v6
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: ${{ matrix.backend }}-minimal
+          deps: testing_unit
+          opencl: ${{ matrix.backend == 'opencl' && 'true' }}
+          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
+          mesa: ${{ matrix.backend == 'lvp' && 'true' }}
+      - name: Set env
+        run: printf "${{ matrix.backend == 'llvm' && 'DEV=CPU:LLVM' || matrix.backend == 'cpu' && 'DEV=CPU\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'DEV=CL' || matrix.backend == 'lvp' && 'DEV=CPU:LVP' }}" >> $GITHUB_ENV
+      - name: Check Device.DEFAULT and print some source
+        run: |
+          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT"
+          DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
+      - name: Run pytest (${{ matrix.backend }})
+        run: python -m pytest -n=auto test/backend --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
+        uses: ./.github/actions/process-replay
+
 # ****** OSX Tests ******

  testmetal:
@@ -814,56 +868,84 @@ jobs:
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

-  testmacos:
-    strategy:
-      fail-fast: false
-      matrix:
-        dev:
-          - 'CPU:CLANG'
-          - 'CPU:LLVM'
-          - 'CPU:LVP'
-          - 'METAL'
-          - 'WEBGPU'
-
-    name: MacOS (DEV=${{ matrix.dev }})
-    runs-on: macos-15
-    timeout-minutes: 20
+  osxwebgpu:
+    name: MacOS (WebGPU)
+    runs-on: macos-14
+    timeout-minutes: 10
    steps:
    - name: Checkout Code
      uses: actions/checkout@v6
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
-        key: macos-${{ matrix.dev }}
-        deps: testing_unit
-        python-version: '3.12'
-        llvm: ${{ contains(matrix.dev, 'LLVM') || contains(matrix.dev, 'LVP') }}
-        mesa: ${{ contains(matrix.dev, 'LVP') && 'cpu' || 'false' }}
-        webgpu: ${{ matrix.dev == 'WEBGPU' }}
-    - name: Set env
-      run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
-    - name: Check Device.DEFAULT and print some source
-      run: |
-        python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
-        DEBUG=4 python test/test_tiny.py TestTiny.test_plus
-    - name: Run backend tests
-      run: python -m pytest -n=auto test/backend --durations=20
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
+        key: osx-webgpu
+        deps: testing
+        webgpu: 'true'
+    - name: Build WEBGPU Efficientnet
+      run: DEV=WEBGPU WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m examples.compile_efficientnet
+    - name: Run selected webgpu tests
+      run: DEV=WEBGPU WEBGPU_BACKEND="WGPUBackendType_Metal" python3 -m pytest -n=auto test/backend --durations=20
+    #- name: Clean npm cache
+    #  run: npm cache clean --force
+    #- name: Install Puppeteer
+    #  run: npm install puppeteer
+    # this is also flaky
+    #- name: Run WEBGPU Efficientnet
+    #  run: node test/web/test_webgpu.js
+    # this is flaky
+    #- name: Run VIZ tests as external package
+    #  run: |
+    #    mkdir $GITHUB_WORKSPACE/test_dir
+    #    cd $GITHUB_WORKSPACE/test_dir
+    #    python -m venv venv
+    #    source venv/bin/activate
+    #    pip install $GITHUB_WORKSPACE
+    #    cp $GITHUB_WORKSPACE/test/web/test_viz.js .
+    #    node test_viz.js
+    - name: Test ONNX Runner (WEBGPU)
+      run: DEV=WEBGPU python3 test/external/external_test_onnx_runner.py

-# ****** Windows Tests ******
-
-  testwindows:
+  osxtests:
    strategy:
      fail-fast: false
      matrix:
-        dev:
-          - 'CPU:CLANG'
-          - 'CPU:LLVM'
-          - 'CPU:X86'
-          - 'WEBGPU'
+        backend: [metal, llvm, cpu, lvp]
+    name: MacOS (${{ matrix.backend }})
+    runs-on: macos-15
+    timeout-minutes: 20
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v6
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: macos-${{ matrix.backend }}-minimal
+          deps: testing_unit
+          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
+          mesa: ${{ matrix.backend == 'lvp' && 'true' }}
+      - name: Set env
+        run: printf "${{ matrix.backend == 'llvm' && 'DEV=CPU:LLVM' || matrix.backend == 'cpu' && 'DEV=CPU\nCPU_COUNT=2' || matrix.backend == 'metal' && 'DEV=METAL' || matrix.backend == 'lvp' && 'DEV=CPU:LVP' }}" >> $GITHUB_ENV
+      - name: Check Device.DEFAULT and print some source
+        run: |
+          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU','LVP':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
+          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
+      - name: Run pytest (${{ matrix.backend }})
+        run: python3 -m pytest -n=auto test/backend --durations=20
+      - name: Run process replay tests
+        uses: ./.github/actions/process-replay
+      - name: Run macOS-specific unit test
+        if: matrix.backend == 'llvm'
+        run: python3 -m pytest test/unit/test_disk_tensor.py::TestDiskTensor::test_copy_to_cpu_not_truncated test/unit/test_cpu.py

-    name: Windows (DEV=${{ matrix.dev }})
+# ****** Windows Tests ******
+
+  wintests:
+    strategy:
+      fail-fast: false
+      matrix:
+        backend: [llvm, cpu, webgpu]
+
+    name: Windows (${{ matrix.backend }})
    runs-on: windows-latest
    timeout-minutes: 15
    steps:
@@ -872,20 +954,25 @@ jobs:
      - name: Setup Environment
        uses: ./.github/actions/setup-tinygrad
        with:
-          key: windows-${{ matrix.dev }}-minimal
+          key: windows-${{ matrix.backend }}-minimal
          deps: testing_unit
-          pydeps: ${{ matrix.dev == 'WEBGPU' && 'dawn-python' || '' }}
+          pydeps: ${{ matrix.backend == 'webgpu' && 'dawn-python' || '' }}
      - name: Set env
        shell: bash
-        run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
-      - name: Check Device.DEFAULT and print some source
+        run:  printf "${{ matrix.backend == 'llvm' && 'DEV=CPU:LLVM' || matrix.backend == 'cpu' && 'DEV=CPU\nCPU_COUNT=2' || matrix.backend == 'webgpu' && 'DEV=WEBGPU'}}" >> $GITHUB_ENV
+      - name: Run unit tests
+        if: matrix.backend=='llvm'
+        # test_newton_schulz hits RecursionError
+        run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_tar.py --ignore=test/unit/test_linalg.py --durations=20
+      - name: Run NULL backend tests
+        if: matrix.backend=='llvm'
+        shell: bash
+        run: DEV=NULL python -m pytest -n=auto test/null/ --ignore=test/null/test_elf.py --durations=20
+      - name: Run pytest (${{ matrix.backend }})
        shell: bash
        run: |
-          python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
-          DEBUG=4 python test/test_tiny.py TestTiny.test_plus
-      - name: Run test_tiny
-        shell: bash
-        run: python -m pytest -n=auto test/test_tiny.py --durations=20
+          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
+          python -m pytest -n=auto test/test_tiny.py test/backend/test_ops.py --durations=20

 # ****** Compile-only Tests ******

@@ -895,7 +982,7 @@ jobs:
      matrix:
        backend: [ir3, nak]
    name: Compile-only (${{ matrix.backend }})
-    runs-on: *linux
+    runs-on: ubuntu-24.04
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -927,7 +1014,7 @@ jobs:
          python -m pytest -n=auto test/backend/test_ops.py --durations=20
  qcomclcompiletests:
    name: Compile-only (QCOM CL)
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'MEMBER' && 'namespace-profile-tinygrad-arm64' || 'ubuntu-24.04-arm' }}
+    runs-on: ubuntu-24.04-arm
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -140,8 +140,8 @@ Documentation along with a quick start guide can be found on the [docs website](
 ```python
 from tinygrad import Tensor

-x = Tensor.eye(3)
-y = Tensor([[2.0,0,-2.0]])
+x = Tensor.eye(3, requires_grad=True)
+y = Tensor([[2.0,0,-2.0]], requires_grad=True)
 z = y.matmul(x).sum()
 z.backward()

@@ -133,7 +133,7 @@ For our loss function we will be using sparse categorical cross entropy loss. Th
 ```python
 def sparse_categorical_crossentropy(self, Y, ignore_index=-1) -> Tensor:
    loss_mask = Y != ignore_index
-    y_counter = Tensor.arange(self.shape[-1], dtype=dtypes.int32, device=self.device).unsqueeze(0).expand(Y.numel(), self.shape[-1])
+    y_counter = Tensor.arange(self.shape[-1], dtype=dtypes.int32, requires_grad=False, device=self.device).unsqueeze(0).expand(Y.numel(), self.shape[-1])
    y = ((y_counter == Y.flatten().reshape(-1, 1)).where(-1.0, 0) * loss_mask.reshape(-1, 1)).reshape(*Y.shape, self.shape[-1])
    return self.log_softmax().mul(y).sum() / loss_mask.sum()
 ```
@@ -175,7 +175,7 @@ with Tensor.train():
  for step in range(1000):
    # random sample a batch
    samp = np.random.randint(0, X_train.shape[0], size=(64))
-    batch = Tensor(X_train[samp])
+    batch = Tensor(X_train[samp], requires_grad=False)
    # get the corresponding labels
    labels = Tensor(Y_train[samp])

@@ -213,7 +213,7 @@ with Timing("Time: "):
  for step in range(1000):
    # random sample a batch
    samp = np.random.randint(0, X_test.shape[0], size=(64))
-    batch = Tensor(X_test[samp])
+    batch = Tensor(X_test[samp], requires_grad=False)
    # get the corresponding labels
    labels = Y_test[samp]

@@ -257,7 +257,7 @@ with Timing("Time: "):
  for step in range(1000):
    # random sample a batch
    samp = np.random.randint(0, X_test.shape[0], size=(64))
-    batch = Tensor(X_test[samp])
+    batch = Tensor(X_test[samp], requires_grad=False)
    # get the corresponding labels
    labels = Y_test[samp]

@@ -174,7 +174,7 @@ if __name__ == "__main__":
  # *** render to device ***

  from tinygrad.codegen import to_program
-  with Context(PCONTIG=2, SPEC=0):
+  with Context(PCONTIG=2, DEVECTORIZE=2, SPEC=0):
    out = tree_traversal(forest_t, val_t, height, rounds)
    sink = out.schedule_linear().src[-1].src[0]
    prg = to_program(sink, VLIWRenderer())
@@ -67,8 +67,8 @@ class ConvGroup:
    self.conv2 = nn.Conv2d(channels_out, channels_out, kernel_size=3, padding=1, bias=False)
    self.norm1 = nn.BatchNorm(channels_out, track_running_stats=False, eps=1e-12, momentum=hyp['net']['batch_norm_momentum'])
    self.norm2 = nn.BatchNorm(channels_out, track_running_stats=False, eps=1e-12, momentum=hyp['net']['batch_norm_momentum'])
-    cast(Tensor, self.norm1.weight).is_param_(False)
-    cast(Tensor, self.norm2.weight).is_param_(False)
+    cast(Tensor, self.norm1.weight).requires_grad = False
+    cast(Tensor, self.norm2.weight).requires_grad = False
  def __call__(self, x:Tensor) -> Tensor:
    x =    self.norm1(self.conv1(x).max_pool2d().float()).cast(dtypes.default_float).quick_gelu()
    return self.norm2(self.conv2(x).float()).cast(dtypes.default_float).quick_gelu() + x
@@ -35,21 +35,22 @@ if __name__ == "__main__":

  params = nn.state.get_parameters(model)

-  # init params
+  # init params, set requires grad on the ones we need gradients of
  for x in params:
+    if x.requires_grad is None: x.requires_grad_()
    x.replace(x.contiguous())
  Tensor.realize(*params)

  # split params (with grads) and buffers (without)
-  params, buffers = partition(params, lambda x: x.is_param)
+  params, buffers = partition(params, lambda x: x.requires_grad)
  print(f"params: {len(params)} buffers: {len(buffers)}")

  # optim params
  pos_params = list(itertools.accumulate(params, lambda x,y: x+y.numel(), initial=0))
  adam_m = Tensor.zeros(pos_params[-1], device="CPU").contiguous()
  adam_v = Tensor.zeros(pos_params[-1], device="CPU").contiguous()
-  adam_b1_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU").contiguous()
-  adam_b2_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU").contiguous()
+  adam_b1_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU", requires_grad=False).contiguous()
+  adam_b2_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU", requires_grad=False).contiguous()
  adam_params = [adam_m, adam_v, adam_b1_t, adam_b2_t]

  # create loss and grads. init all state so the JIT works on microbatch
@@ -30,9 +30,9 @@ class UnsyncedBatchNorm:
    if affine: self.weight, self.bias = Tensor.ones(sz, dtype=dtypes.float32), Tensor.zeros(sz, dtype=dtypes.float32)
    else: self.weight, self.bias = None, None

-    self.running_mean = Tensor.zeros(num_devices, sz, dtype=dtypes.float32).is_param_(False)
-    self.running_var = Tensor.ones(num_devices, sz, dtype=dtypes.float32).is_param_(False)
-    self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.int).is_param_(False)
+    self.running_mean = Tensor.zeros(num_devices, sz, dtype=dtypes.float32, requires_grad=False)
+    self.running_var = Tensor.ones(num_devices, sz, dtype=dtypes.float32, requires_grad=False)
+    self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.int, requires_grad=False)

  def __call__(self, x:Tensor):
    xr = x.reshape(self.num_devices, -1, *x.shape[1:]).cast(dtypes.float32)
@@ -68,7 +68,8 @@ class UnsyncedBatchNorm:
 class BatchNorm(nn.BatchNorm2d if getenv("SYNCBN") else UnsyncedBatchNorm):
  def __init__(self, num_features):
    super().__init__(num_features, track_running_stats=False, eps=1e-12, momentum=0.85, affine=True)
-    self.weight.is_param_(False)
+    self.weight.requires_grad = False
+    self.bias.requires_grad = True

 class ConvGroup:
  def __init__(self, channels_in, channels_out):
@@ -171,7 +172,7 @@ def train_cifar():
    Λ, V = _eigens(_patches(X.float().numpy()))
    W = V/np.sqrt(Λ+1e-2)[:,None,None,None]

-    return Tensor(W.astype(np.float32)).cast(dtypes.default_float).is_param_(False)
+    return Tensor(W.astype(np.float32), requires_grad=False).cast(dtypes.default_float)

  # ========== Loss ==========
  def cross_entropy(x:Tensor, y:Tensor, reduction:str='mean', label_smoothing:float=0.0) -> Tensor:
@@ -263,6 +264,7 @@ def train_cifar():
      # self.model_ema = copy.deepcopy(net) # won't work for opencl due to unpickeable pyopencl._cl.Buffer
      self.net_ema = SpeedyResNet(w)
      for net_ema_param, net_param in zip(get_state_dict(self.net_ema).values(), get_state_dict(net).values()):
+        net_ema_param.requires_grad = False
        net_ema_param.assign(net_param.numpy())

    @TinyJit
@@ -305,7 +307,7 @@ def train_cifar():
  params_bias = []
  params_non_bias = []
  for params in params_dict:
-    if params_dict[params].is_param:
+    if params_dict[params].requires_grad is not False:
      if 'bias' in params:
        params_bias.append(params_dict[params])
      else:
@@ -102,7 +102,7 @@ class Int8Embedding:
    self.weight, self.scale = Tensor.ones(vocab_size, embed_size, dtype=dtypes.int8), Tensor.ones(vocab_size, dtype=dtypes.half)

  def __call__(self, idx:Tensor) -> Tensor:
-    if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, device=self.weight.device).unsqueeze(-1)
+    if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, requires_grad=False, device=self.weight.device).unsqueeze(-1)
    big_shp = idx.shape+(self.vocab_sz, self.embed_sz)
    arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1)).expand(big_shp), (self.weight.cast(self.scale.dtype).T*self.scale).T
    return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
@@ -25,7 +25,7 @@ class CausalSelfAttention:
    self.n_embd = config.n_embd
    # not really a 'bias', more of a mask, but following the OpenAI/HF naming though
    self.bias = Tensor.ones(1, 1, config.block_size, config.block_size).tril()
-    self.bias.is_param_(False)
+    self.bias.requires_grad = False

  def __call__(self, x:Tensor):
    B, T, C = x.shape
@@ -1,6 +1,6 @@
 import functools, argparse, pathlib
 from tinygrad import Tensor, nn, Device, GlobalCounters, Variable
-from tinygrad.helpers import Timing, Profiling, tqdm
+from tinygrad.helpers import Timing, Profiling, CI, tqdm
 from tinygrad.nn.state import torch_load, get_state_dict
 from extra.models.llama import FeedForward, Transformer
 from extra.bench_log import BenchEvent, WallTimeEvent
@@ -36,7 +36,7 @@ if __name__ == "__main__":
    model = Transformer(n_layers=32, dim=4096, hidden_dim=14336, n_heads=32, n_kv_heads=8, norm_eps=1e-5, vocab_size=32000, feed_forward=functools.partial(MixtureFeedForward, 8), jit=False)
    model_state_dict = get_state_dict(model)

-    for k in (t := tqdm(state, disable=None)):
+    for k in (t := tqdm(state, disable=CI)):
      if 'feed_forward.experts.' in k:
        expert_no = int(k.split('feed_forward.experts.')[1].split('.')[0])
        device = Device.DEFAULT + ":" + str((expert_no//2)+1)
@@ -44,7 +44,7 @@ if __name__ == "__main__":
        device = Device.DEFAULT
      t.set_description(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB, loading {k} to {device}")
      model_state_dict[k].replace(state[k].to(device).half()).realize()
-  if t.disable: print(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB")
+  if CI: print(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB")

  from sentencepiece import SentencePieceProcessor
  spp = SentencePieceProcessor(model_file=args.weights + "/tokenizer.model")
@@ -57,7 +57,7 @@ class EmbeddingBert(nn.Embedding):
  def __call__(self, idx:Tensor) -> Tensor:
    if idx.numel() == 0: return Tensor.empty(idx.shape+(self.embed_sz,), dtype=self.weight.dtype, device=self.weight.device)
    arange_shp, weight_shp, big_shp = (1, 1, self.vocab_sz, 1), (1, 1, self.vocab_sz, self.embed_sz), idx.shape+(self.vocab_sz, self.embed_sz,)
-    if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, device=self.weight.device).reshape(arange_shp)
+    if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, requires_grad=False, device=self.weight.device).reshape(arange_shp)
    arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1,)).expand(big_shp), self.weight.cast(dtypes.default_float).reshape(weight_shp).expand(big_shp)
    return (arange == idx).where(vals, 0).sum(2, dtype=vals.dtype)

@@ -77,11 +77,11 @@ class FrozenBatchNorm2dRetinaNet(nn.BatchNorm2d):
  def __init__(self, sz:int, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1):
    self.eps, self.track_running_stats, self.momentum = eps, track_running_stats, momentum

-    self.weight = Tensor.ones(sz, dtype=dtypes.float32).is_param_(False) if affine else None
-    self.bias = Tensor.zeros(sz, dtype=dtypes.float32).is_param_(False) if affine else None
+    self.weight = Tensor.ones(sz, dtype=dtypes.float32, requires_grad=False) if affine else None
+    self.bias = Tensor.zeros(sz, dtype=dtypes.float32, requires_grad=False) if affine else None

-    if track_running_stats: self.running_mean, self.running_var = Tensor.zeros(sz, dtype=dtypes.float32).is_param_(False), Tensor.ones(sz, dtype=dtypes.float32).is_param_(False)
-    self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.long).is_param_(False)
+    if track_running_stats: self.running_mean, self.running_var = Tensor.zeros(sz, dtype=dtypes.float32, requires_grad=False), Tensor.ones(sz, dtype=dtypes.float32, requires_grad=False)
+    self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.long, requires_grad=False)

  def __call__(self, x:Tensor) -> Tensor:
    batch_mean, batch_var = super().calc_stats(x.cast(dtypes.float32))
@@ -180,11 +180,11 @@ def train_resnet():
  def fake_data_get(batch_size):
    x = Tensor.zeros(batch_size, 224, 224, 3, dtype=dtypes.uchar).contiguous()
    y = [0] * batch_size
-    return x.shard(GPUS, axis=0).realize(), Tensor(y).shard(GPUS, axis=0), y, None
+    return x.shard(GPUS, axis=0).realize(), Tensor(y, requires_grad=False).shard(GPUS, axis=0), y, None

  def data_get(it):
    x, y, cookie = next(it)
-    return x.shard(GPUS, axis=0).realize(), Tensor(y).shard(GPUS, axis=0), y, cookie
+    return x.shard(GPUS, axis=0).realize(), Tensor(y, requires_grad=False).shard(GPUS, axis=0), y, cookie

  # ** epoch loop **
  step_times = []
@@ -413,7 +413,7 @@ def train_retinanet():
    layers_to_train = ["layer4", "layer3", "layer2", "layer1", "conv1"][:trainable_layers]
    for k, v in get_state_dict(backbone).items():
      if all([not k.startswith(layer) for layer in layers_to_train]):
-        v.is_param_(False)
+        v.requires_grad = False

  def _data_get(it:Iterator[tuple[Tensor, ...]], val:bool=False):
    if val:
@@ -798,7 +798,7 @@ def train_unet3d():
  @Tensor.train(mode=False)
  def eval_step(model, x, y):
    y_hat, y = sliding_window_inference(model, x, y, gpus=GPUS)
-    y_hat, y = Tensor(y_hat), Tensor(y)
+    y_hat, y = Tensor(y_hat), Tensor(y, requires_grad=False)
    loss = dice_ce_loss(y_hat, y)
    score = dice_score(y_hat, y)
    return loss.realize(), score.realize()
@@ -1419,7 +1419,10 @@ def train_llama3():

  for p in optim.params:
    grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
-    p.grad = p.zeros_like(dtype=grad_dtype).contiguous()
+    if isinstance(p.device, tuple) and p.uop.axis is not None:
+      p.grad = Tensor.zeros(p.shape, dtype=grad_dtype, device=p.device[0]).shard_(p.device, axis=p.uop.axis).contiguous()
+    else:
+      p.grad = Tensor.zeros(p.shape, dtype=grad_dtype, device=p.device).contiguous()
  grads = [p.grad for p in optim.params]

  scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
@@ -1435,19 +1438,16 @@ def train_llama3():

  fp8_amax = [t for ts in model._fp8_amax.values() for t in ts]
  fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts] if hasattr(model, "_fp8_grad_amax") else []
-  fp8_inv_scales = list(model._fp8_inv_scale.values()) + list(model._fp8_next_inv_scale.values())
+  fp8_inv_scales = list(model._fp8_inv_scale.values())

  from tinygrad.nn.state import get_state_dict
  model_state = get_state_dict(model)
  for wname in model._fp8_inv_scale:
    w = model_state[wname]
    w._inv_scale = model._fp8_inv_scale[wname]
-    w._next_inv_scale = model._fp8_next_inv_scale[wname]
    if optim.master_params:
      idx = next(j for j, p in enumerate(optim.params) if p is w)
-      master = optim.master_params[idx]
-      inv = w._inv_scale if w._inv_scale.device == master.device else w._inv_scale.to(master.device)
-      master.assign((master * inv.reshape(-1, *([1]*(w.ndim-1)))).contiguous())
+      optim.master_params[idx].assign((optim.master_params[idx] * w._inv_scale.reshape(-1, *([1]*(w.ndim-1)))).contiguous())

  # realize everything here
  if optim.master_params: Tensor.realize(*optim.master_params)
@@ -1458,7 +1458,7 @@ def train_llama3():
    if is_dp: tokens = tokens.to(None).shard(device, 0)
    if is_mp: tokens = tokens.shard(device)
    if not is_sharding: tokens = tokens.to(None)
-    logits:Tensor = model(tokens[:, :-1], save=bool(SMALL))
+    logits:Tensor = model(tokens[:, :-1])
    if getenv("FAST_CE", 0):
      from extra.llama_kernels.fused_ce import fused_ce_loss
      loss = fused_ce_loss(logits.cast(dtypes.bfloat16), tokens[:, 1:], label_smoothing=0.0)
@@ -1498,7 +1498,7 @@ def train_llama3():
  def fake_data(bs, samples):
    import numpy as np
    for _ in range(samples // bs):
-      fake_data_np = np.random.randint(0, real_vocab_size, size=(bs, SEQLEN + 1), dtype=np.int32)
+      fake_data_np = np.random.randint(0, model_params["vocab_size"], size=(bs, SEQLEN + 1), dtype=np.int32)
      yield Tensor(fake_data_np, device="NPY")

  def get_train_iter():
@@ -2,8 +2,9 @@ import math, os
 if __name__ == "__main__":
  os.environ["DEFAULT_FLOAT"] = "bfloat16"
  os.environ["OPTIM_DTYPE"] = "bfloat16"
-  if "DEV" not in os.environ: os.environ["DEV"] = "NULL::gfx950"
+  if "DEV" not in os.environ: os.environ["DEV"] = "NULL"
  # CDNA
+  os.environ["EMULATE"] = "AMD_CDNA4"
  os.environ["DEVICE_IN_FUNCTION_BUG"] = "1"
  os.environ["ALL2ALL"] = "1"
  os.environ["USE_ATOMICS"] = "1"
@@ -12,7 +13,7 @@ if __name__ == "__main__":
    if "ASM_GEMM" not in os.environ:
      os.environ["ASM_GEMM"] = "1"
 from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
-from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker, round_up
+from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
 from extra.llama_kernels.rmsnorm import rmsnorm
@@ -52,8 +53,8 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca
  if ASM_GEMM:
    from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
    if can_use_asm_gemm(x_fp8, w.T):
-      return asm_gemm(x_fp8, w.T, x_scale=x_scale, w_scale=w_inv_scale, grad_amax_state=grad_amax_state), x_new_amax, x_fp8
-  return (x_fp8.dot(w.T, dtype=dtypes.float) * x_scale * w_inv_scale).cast(dtypes.bfloat16), x_new_amax, x_fp8
+      return asm_gemm(x_fp8, w.T, x_scale=x_scale, w_scale=w_inv_scale, grad_amax_state=grad_amax_state), x_new_amax, x_fp8, w
+  return (x_fp8.dot(w.T, dtype=dtypes.float) * x_scale * w_inv_scale).cast(dtypes.bfloat16), x_new_amax, x_fp8, w

 def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor, grad_amax_state:Tensor):
  if FUSED_ADD_NORM_MUL_QUANTIZE:
@@ -124,9 +125,9 @@ class FlatTransformer:
    self.tok_embeddings = nn.Embedding(vocab_size, dim)
    self.tok_embeddings.weight = Tensor.normal(vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
    self.output = Tensor.normal(1, vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
-    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().is_param_(False)
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)

-    def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().is_param_(False)
+    def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().requires_grad_(False)
    names = ["xqkv", "xo", "x2"]
    names += ["x1", "x3"] if SPLIT_W13 else ["x13"]
    self._fp8_amax = {name: [_amax() for _ in range(n_layers)] for name in names}
@@ -135,8 +136,7 @@ class FlatTransformer:
    self._fp8_grad_amax = {name: [_amax() for _ in range(n_layers)] for name in grad_names}
    w_scales = [("wqkv", s_qkv), ("wo", s_o), ("w2", s_2)]
    w_scales += [("w1", s_1), ("w3", s_3)] if SPLIT_W13 else [("w13", s_13)]
-    self._fp8_inv_scale = {name: s.float().contiguous().is_param_(False) for name, s in w_scales}
-    self._fp8_next_inv_scale = {name: s.float().contiguous().is_param_(False) for name, s in w_scales}
+    self._fp8_inv_scale = {name: s.float().contiguous().requires_grad_(False) for name, s in w_scales}

  def lin_per_layer(self, in_features:int, out_features:int, std:float=0.02):
    if getenv("ZEROS"): w = Tensor.zeros(self.n_layers, out_features, in_features)
@@ -208,13 +208,11 @@ class FlatTransformer:
    return out, h, amaxs, saves

  @function(precompile=True, precompile_backward=True)
-  def run_layer(self, x:Tensor, freqs_cis:Tensor, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
+  def run_layer(self, x:Tensor, freqs_cis:Tensor, attn_kwargs:dict, ffn_kwargs:dict):
    attn, attn_amaxs, attn_saves = self.attention(x, freqs_cis, **attn_kwargs)
    ffn, h, ffn_amaxs, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs)
    h = h + ffn
-    amaxs = tuple(a.detach() for a in (*attn_amaxs, *ffn_amaxs))
-    if save: return (h, *amaxs, *attn_saves, *ffn_saves)
-    else: return (h, *amaxs)
+    return (h, *attn_amaxs, *ffn_amaxs, *attn_saves, *ffn_saves)

  def shard(self, device:tuple[str, ...], mp:bool=False):
    from tinygrad.nn.state import get_parameters
@@ -222,19 +220,14 @@ class FlatTransformer:
      for v in get_parameters(self): v.shard_(device, axis=None)
    else:
      # flat per-layer weights: axis 0 is n_layers, so shard axes are +1 vs per-layer Transformer
-      def _shard_fp8(name:str, axis:int):
-        getattr(self, name).shard_(device, axis=axis)
-        self._fp8_inv_scale[name] = self._fp8_inv_scale[name].to(device).contiguous().is_param_(False)
-        self._fp8_next_inv_scale[name] = self._fp8_next_inv_scale[name].to(device).contiguous().is_param_(False)
-        Tensor.realize(getattr(self, name), self._fp8_inv_scale[name], self._fp8_next_inv_scale[name])
-      _shard_fp8("wqkv", 1)          # (n_layers, out, dim) shard out
-      _shard_fp8("wo", 2)            # (n_layers, dim, in) shard in
+      self.wqkv.shard_(device, axis=1).realize()          # (n_layers, out, dim) shard out
+      self.wo.shard_(device, axis=2).realize()            # (n_layers, dim, in) shard in
      if SPLIT_W13:
-        _shard_fp8("w1", 1)
-        _shard_fp8("w3", 1)
+        self.w1.shard_(device, axis=1).realize()
+        self.w3.shard_(device, axis=1).realize()
      else:
-        _shard_fp8("w13", 1)         # (n_layers, hidden*2, dim) shard out
-      _shard_fp8("w2", 2)            # (n_layers, dim, hidden) shard in
+        self.w13.shard_(device, axis=1).realize()           # (n_layers, hidden*2, dim) shard out
+      self.w2.shard_(device, axis=2).realize()            # (n_layers, dim, hidden) shard in
      self.attention_norm.shard_(device, axis=None).realize()
      self.ffn_norm.shard_(device, axis=None).realize()
      self.norm.weight.shard_(device, axis=None).realize()
@@ -244,9 +237,11 @@ class FlatTransformer:
      for amax_dict in (self._fp8_amax, self._fp8_grad_amax):
        for name in amax_dict:
          for i in range(len(amax_dict[name])):
-            amax_dict[name][i] = amax_dict[name][i].to(device).contiguous().is_param_(False)
+            amax_dict[name][i] = amax_dict[name][i].to(device).contiguous().requires_grad_(False)
+      for name in self._fp8_inv_scale:
+        self._fp8_inv_scale[name] = self._fp8_inv_scale[name].to(device).contiguous().requires_grad_(False)

-  def __call__(self, tokens:Tensor, save:bool=True):
+  def __call__(self, tokens:Tensor):
    h = self.tok_embeddings(tokens)
    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
    a, ga, s = self._fp8_amax, self._fp8_grad_amax, self._fp8_inv_scale
@@ -261,7 +256,7 @@ class FlatTransformer:
                          s_1=s["w1"][i], s_3=s["w3"][i], grad_amax_xw1=ga["xw1"][i], grad_amax_xw3=ga["xw3"][i])
      else:
        ffn_kwargs.update(w13=self.w13[i], amax_x13=a["x13"][i], s_13=s["w13"][i], grad_amax_xw13=ga["xw13"][i])
-      h, *ret = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)
+      h, *ret = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs)
      amax_names = ["xqkv", "xo"] + (["x1", "x3"] if SPLIT_W13 else ["x13"]) + ["x2"]
      for name, new_val in zip(amax_names, ret[:len(amax_names)]):
        a[name][i].assign(new_val)
@@ -277,59 +272,41 @@ def apply_grad(grad_buf:Tensor, new_grad:UOp):
  pads = _get_pads(new_grad)
  if len(pads) <= 1:
    new_grad = new_grad.cast(grad_buf.dtype)
-    grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(grad_buf.uop + new_grad))
+    store = grad_buf.uop.store(grad_buf.uop + new_grad)
+    grad_buf.uop = grad_buf.uop.after(store)
    return
-  cur = grad_buf.uop
-  for pad in sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0, reverse=True):
-    if pad.op == Ops.PAD:
-      grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg)])
-      buf_slice = cur.shrink(grad_shrink)
-      cur = cur.after(buf_slice.store(buf_slice + pad.src[0].cast(cur.dtype)))
-    else:
-      cur = cur.after(cur.store(cur + pad.cast(cur.dtype)))
-  grad_buf.uop = cur
+  sorted_pads = sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0)
+  inners_raw = [Tensor(p.src[0] if p.op == Ops.PAD else p, device=grad_buf.device) for p in sorted_pads]
+  if getenv("FUSED_PAD_GRAD_ACCUM", 0):
+    from extra.llama_kernels.fused_pad_grad_accum import fused_pad_grad_accum, can_fused_pad_grad_accum
+    if can_fused_pad_grad_accum(grad_buf, inners_raw):
+      grad_buf.uop = fused_pad_grad_accum(grad_buf, inners_raw).uop
+      return
+  inners = [t.cast(grad_buf.dtype) for t in inners_raw]
+  grad_buf.assign(grad_buf + inners[0].cat(*inners[1:], dim=0))

 if __name__ == "__main__":
  config = {}
  BS                 = config["BS"]                     = getenv("BS", 16)
  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
-  SMALL              = config["SMALL"]                  = getenv("SMALL", 0)

  from examples.llama3 import MODEL_PARAMS
-  model_params = MODEL_PARAMS[llama_size:=getenv("LLAMA3_SIZE", "8B")]["args"]
-  # vocab_size from mixtral tokenizer
-  if not SMALL: model_params |= {"vocab_size": 32000}
-  real_vocab_size = model_params['vocab_size']
-  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params["n_layers"] = llama_layers
-
-  # pad vocab
-  if (MP := getenv("MP", 1)) > 1: model_params["vocab_size"] = round_up(model_params["vocab_size"], 256 * MP)
-  vocab_mask:Tensor = Tensor.arange(model_params["vocab_size"]).reshape(1, 1, -1) >= real_vocab_size
-
+  model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
+  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
  model = FlatTransformer(**model_params, max_context=SEQLEN)
-
  state = nn.state.get_state_dict(model)
  print("tensor count:", len(state))

  # shard the model
  from tinygrad import Device
-  is_dp = (DP := getenv("DP", 1)) > 1
-  is_mp = (MP := getenv("MP", 1)) > 1
-  is_sharding = is_dp or is_mp
-  device_count = max(DP, MP)
-  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))
-
-  model.shard(device, is_mp)
-
-  if is_dp: vocab_mask.shard_(device, axis=None).realize()
-  if is_mp: vocab_mask.shard_(device, axis=2).realize()
+  if (DP := getenv("DP", 1)) > 1:
+    model.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)))
+  if (MP := getenv("MP", 1)) > 1:
+    model.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)), mp=True)

  # preallocate all the grad buffers and zero them out
-  grad_dtype = lambda x: dtypes.bfloat16 if x.dtype in dtypes.fp8s else x.dtype
-  grads = {x:x.zeros_like(dtype=grad_dtype(x)).contiguous() for x in state.values() if x.is_param}
-
-  fp8_amax = [t for ts in model._fp8_amax.values() for t in ts]
-  fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts]
+  grads = {x:Tensor.zeros(x.shape, dtype=x.dtype, device=x.device).contiguous()
+           for x in state.values() if x.requires_grad is None}

  # print model size
  sz = 0
@@ -338,31 +315,23 @@ if __name__ == "__main__":
    sz += v.nbytes()
  print(f"total sz: {sz/1e9:.2f} GB")

-  with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=real_vocab_size, dtype=dtypes.int)
+  with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=model.vocab_size, dtype=dtypes.int)
  with Timing("realize weights/grads/data: "): Tensor.realize(*state.values(), *grads.values(), tokens)
  print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
  if DP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)), axis=0)
  if MP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)))

  @TinyJit
-  def fwd_bwd(tokens:Tensor):
-    with Timing("python forward: "):
-      logits = model(tokens[:, :-1], save=llama_size=="8B")
-      loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
+  def jit_step(tokens:Tensor):
+    with Timing("python forward: "): loss = model(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
    with Timing("python backward: "):
      for t,g in zip(grads, loss.gradient(*grads)):
        apply_grad(grads[t], g.uop)
-    with Timing("run fwd_bwd: "): loss.realize(*grads.values(), *fp8_amax, *fp8_grad_amax)
-
-  @TinyJit
-  def optim_step():
-    for g in grads.values(): g.assign(g.zeros_like())
-    Tensor.realize(*grads.values())
+    with Timing("run step: "): loss.realize(*grads.values())

  for i in range(6):
    GlobalCounters.reset()
    profile_marker(f"step {i}")
    with Timing(colored(f"*** step {i}: ", "red")):
-      fwd_bwd(tokens)
-      optim_step()
+      jit_step(tokens)
  print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
@@ -1,68 +0,0 @@
-import unittest
-from tinygrad import Tensor, TinyJit
-from tinygrad.nn.state import get_parameters
-from examples.mlperf.models.flat_llama import apply_grad
-
-class FlatModel:
-  def __init__(self, n_layers:int, dim:int, hidden:int):
-    self.n_layers = n_layers
-    self.w1 = Tensor.uniform(n_layers, dim, hidden, low=-0.1, high=0.1)
-    self.w2 = Tensor.uniform(n_layers, hidden, dim, low=-0.1, high=0.1)
-    self.scale = Tensor.uniform(dim, low=0.9, high=1.1)
-    self.bias = Tensor.zeros(dim).contiguous()
-
-  def __call__(self, x:Tensor) -> Tensor:
-    h = x
-    for i in range(self.n_layers):
-      h = (h @ self.w1[i]).relu() @ self.w2[i] + h
-    return (h * self.scale + self.bias).sum()
-
-class TestApplyGradE2E(unittest.TestCase):
-  def _run_with_apply_grad(self, model, xs):
-    grads = {p: Tensor.zeros(p.shape, dtype=p.dtype).contiguous().realize() for p in get_parameters(model)}
-    for x in xs:
-      loss = model(x)
-      for p, g in zip(grads, loss.gradient(*grads)):
-        apply_grad(grads[p], g.uop)
-      Tensor.realize(loss, *grads.values())
-    return [grads[p] for p in get_parameters(model)]
-
-  def _run_reference(self, model, xs):
-    for x in xs: model(x).backward()
-    return [p.grad for p in get_parameters(model)]
-
-  def _assert_close(self, got, expected, atol, rtol):
-    for g, e in zip(got, expected):
-      self.assertTrue(g.allclose(e, atol=atol, rtol=rtol).item(), f"grad mismatch (max abs diff {(g - e).abs().max().item()})")
-
-  def _assert_match(self, model, xs, atol, rtol):
-    self._assert_close(self._run_with_apply_grad(model, xs), self._run_reference(model, xs), atol, rtol)
-
-  def test_e2e_single_step(self):
-    model = FlatModel(n_layers=3, dim=8, hidden=16)
-    Tensor.realize(*get_parameters(model))
-    self._assert_match(model, [Tensor.randn(2, 8).realize()], atol=1e-4, rtol=1e-4)
-
-  def test_e2e_multi_step_accumulation(self):
-    model = FlatModel(n_layers=4, dim=8, hidden=16)
-    Tensor.realize(*get_parameters(model))
-    self._assert_match(model, [Tensor.randn(2, 8).realize() for _ in range(3)], atol=1e-4, rtol=1e-4)
-
-  def test_e2e_jit(self):
-    model = FlatModel(n_layers=3, dim=8, hidden=16)
-    Tensor.realize(*get_parameters(model))
-    grads = {p: Tensor.zeros(p.shape, dtype=p.dtype).contiguous().realize() for p in get_parameters(model)}
-
-    @TinyJit
-    def fwd_bwd(x:Tensor):
-      loss = model(x)
-      for p, g in zip(grads, loss.gradient(*grads)): apply_grad(grads[p], g.uop)
-      Tensor.realize(loss, *grads.values())
-
-    xs = [Tensor.randn(2, 8).realize() for _ in range(3)]
-    for x in xs: fwd_bwd(x)
-    self._assert_close([grads[p] for p in get_parameters(model)], self._run_reference(model, xs), atol=1e-3, rtol=1e-3)
-
-
-if __name__ == "__main__":
-  unittest.main()
@@ -3,7 +3,8 @@ os.environ["WQKV"] = "1"
 import unittest
 import numpy as np
 from tinygrad import Tensor, nn, dtypes
-from tinygrad.device import Device
+from tinygrad.nn.state import get_parameters
+from tinygrad.device import is_dtype_supported, Device
 from examples.mlperf.models.llama import Transformer
 from examples.mlperf.models.flat_llama import FlatTransformer

@@ -44,6 +45,8 @@ class TestFlatLlama(unittest.TestCase):
    flat = FlatTransformer(**params)
    copy_weights(flat, ref)

+    for p in get_parameters(ref): p.requires_grad_(True)
+    for p in get_parameters(flat): p.requires_grad_(True)
    Tensor.realize(*nn.state.get_state_dict(flat).values())

    tokens = Tensor([[1, 50, 100, 999, 2, 10]])
@@ -111,7 +114,7 @@ class TestFlatLlama(unittest.TestCase):
    self.assertEqual(ref_logits.shape, flat_logits.shape)
    np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)

-  @unittest.skipUnless(dtypes.fp8e4m3 in Device[Device.DEFAULT].renderer.supported_dtypes(), "fp8 not supported on this device")
+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), "fp8 not supported on this device")
  def test_forward_fp8(self):
    import examples.mlperf.models.flat_llama as flat_llama_mod
    old_fp8 = flat_llama_mod.FP8
@@ -6,7 +6,6 @@ from tinygrad.uop.ops import UOp, Ops

 STOCHASTIC_ROUND = getenv("STOCHASTIC_ROUND", 0)
 MASTER_WEIGHTS = getenv("MASTER_WEIGHTS", 0)
-FP8_AMAX_MARGIN = getenv("FP8_AMAX_MARGIN", 1.1)

 def stochastic_round_bf16(x:Tensor) -> Tensor:
  bits = x.bitcast(dtypes.uint32)
@@ -22,14 +21,11 @@ class GradAccClipAdamW(Optimizer):
  def __init__(self, params:list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, grad_acc=1, clip_norm=1.0, device=None, fused=FUSE_OPTIM):
    super().__init__(params, lr, device, fused)
    self.b1, self.b2, self.eps, self.wd = b1, b2, eps, weight_decay
-    self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device) for _ in [b1, b2])
+    self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device, requires_grad=False) for _ in [b1, b2])
    self.m = self._new_optim_param()
    self.v = self._new_optim_param()
    self.grad_acc, self.clip_norm = grad_acc, clip_norm
-    if MASTER_WEIGHTS and self.params[0].dtype != dtypes.float32:
-      self.master_params:list[Tensor]|None = [p.to(self.device).float().contiguous() for p in self.params]
-    else:
-      self.master_params = None
+    self.master_params:list[Tensor]|None = [p.float().contiguous() for p in self.params] if MASTER_WEIGHTS and self.params[0].dtype != dtypes.float32 else None

  def fstep(self, grads:list[Tensor]):
    if self.fused:
@@ -40,8 +36,7 @@ class GradAccClipAdamW(Optimizer):
    for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i], self.master_params[i] if self.master_params else None))
    # collect inv_scale tensors attached to fp8 params (set by _apply_update)
    fp8_inv_scales = [tt._inv_scale for tt in self.params if hasattr(tt, '_inv_scale')]
-    fp8_next_inv_scales = [tt._next_inv_scale for tt in self.params if hasattr(tt, '_next_inv_scale')]
-    to_realize = extra+self.params+self.buffers+(self.master_params or [])+fp8_inv_scales+fp8_next_inv_scales
+    to_realize = extra+self.params+self.buffers+(self.master_params or [])+fp8_inv_scales

    Tensor.realize(*to_realize)
    return extra[-1]
@@ -83,23 +78,13 @@ class GradAccClipAdamW(Optimizer):
    up = up.float().shard_like(w) + self.lr.to(w.device) * wd * w.detach()
    new_w = w.detach() - up
    if master is not None: master.assign(new_w)
-    # when master is offloaded to a different device than the param, results are resharded back onto the param's (sharded) device
-    offloaded = master is not None and master.device != t.device
-    if STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16:
-      out = stochastic_round_bf16(new_w)
-      return out.shard_like(t) if offloaded else out
+    if STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16: return stochastic_round_bf16(new_w)
    if t.dtype in dtypes.fp8s:
      from examples.mlperf.models.flat_llama import FP8_MAX
-      # delayed scaling: reuse previous step's inv_scale
-      t._inv_scale.assign(t._next_inv_scale)
-      inv_scale = t._inv_scale.to(new_w.device) if offloaded else t._inv_scale
-      scale = inv_scale.reciprocal().reshape(-1, *([1]*(new_w.ndim-1)))
-      scaled = (new_w * scale).clamp(-FP8_MAX, FP8_MAX)
-      ret = scaled.cast(t.dtype)
-      # update inv_scale for next step from quantized result
-      new_amax = (ret.float().abs().max(axis=tuple(range(1, ret.ndim))) * inv_scale * FP8_AMAX_MARGIN).detach()
-      new_inv = ((new_amax + 1e-8) / FP8_MAX).cast(t._inv_scale.dtype)
-      t._next_inv_scale.assign(new_inv.shard_like(t._next_inv_scale) if offloaded else new_inv)
-      return ret.shard_like(t) if offloaded else ret
-    out = new_w.cast(t.dtype)
-    return out.shard_like(t) if offloaded else out
+      amax = new_w.float().abs().max(axis=tuple(range(1, new_w.ndim))).detach()  # per-layer amax for (n_layers, out, in)
+      scale = FP8_MAX / (amax + 1e-8)
+      fp8_w = (new_w * scale.reshape(-1, *([1]*(new_w.ndim-1)))).clamp(-FP8_MAX, FP8_MAX).cast(t.dtype)
+      if hasattr(t, '_inv_scale'):
+        t._inv_scale.assign(((amax + 1e-8) / FP8_MAX).cast(t._inv_scale.dtype))
+      return fp8_w
+    return new_w.cast(t.dtype)
@@ -1,8 +1,6 @@
 #!/usr/bin/env bash

 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=${DEV:-AMD}
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -18,12 +16,10 @@ export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
 export FAST_CE=${FAST_CE:-0}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
-export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
-export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
-export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
-export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-0}
-export SPLIT_W13=${SPLIT_W13:-1}
+export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
+export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
+export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
+export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
@@ -1,34 +1,22 @@
 #!/usr/bin/env bash

 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
-export DEVICE_IN_FUNCTION_BUG=1

 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export USE_ATOMICS=${USE_ATOMICS:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export FP8=${FP8:-1}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-0}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
-export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
-export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
-export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
-export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-0}
-export SPLIT_W13=${SPLIT_W13:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}

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

 export MODEL="llama3"
 export BASEDIR="/raid/datasets/c4/"
@@ -1,8 +1,6 @@
 #!/usr/bin/env bash

 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=${DEV:-AMD}
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -23,7 +21,7 @@ export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
 export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
 export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
 export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
-export SPLIT_W13=${SPLIT_W13:-0}
+export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1}

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
@@ -1,8 +1,6 @@
 #!/usr/bin/env bash

 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=${DEV:-AMD}
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -22,6 +20,7 @@ export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
 export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
 export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
 export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
+export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-0}
 export SPLIT_W13=${SPLIT_W13:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}

@@ -1,8 +1,6 @@
 #!/usr/bin/env bash

 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=${DEV:-AMD}
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -23,7 +21,7 @@ export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
 export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-1}
 export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
 export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
-export SPLIT_W13=${SPLIT_W13:-0}
+export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1}

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
@@ -1,9 +1,8 @@
 #!/usr/bin/env bash

 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
 export DEVICE_IN_FUNCTION_BUG=1
@@ -22,6 +21,7 @@ export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-0}
 export FUSED_GRAD_QUANTIZE=${FUSED_GRAD_QUANTIZE:-0}
 export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-0}
 export FUSED_SILU_W13=${FUSED_SILU_W13:-0}
+export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-0}
 export SPLIT_W13=${SPLIT_W13:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}

@@ -1,6 +1,6 @@
 #!/bin/bash
 export BENCHMARK=5
 export EVAL_BS=0
-VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
+VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh
 SRC="AMD"; [[ $DEV == NULL* ]] && SRC="NULL"
-python -m tinygrad.viz.cli -s "$SRC" -t --interval "train @ 2" "train @ 3"
+python -m tinygrad.viz.cli -s "$SRC" -t
@@ -3,8 +3,6 @@ set -e  # Exit on any error
 set -o pipefail  # Make pipeline fail if any command fails

 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=AMD
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -24,7 +22,7 @@ export FUSED_INPUT_QUANTIZE=1
 export FUSED_GRAD_QUANTIZE=1
 export FUSED_ADD_NORM_MUL_QUANTIZE=1
 export FUSED_SILU_W13=1
-export SPLIT_W13=0
+export FUSED_PAD_GRAD_ACCUM=1

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=8 MP=1 BS=16 EVAL_BS=8 GRADIENT_ACC_STEPS=2
@@ -4,7 +4,7 @@ export EVAL_BS=0
 export FAKEDATA=1
 export NULL_ALLOW_COPYOUT=1
 export HIP_VISIBLE_DEVICES=""
-export DEV=NULL:HIP:gfx950
+export DEV=NULL
 export JITBEAM=0
 export LLAMA_LAYERS=${LLAMA_LAYERS:-"2"}
-time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh
+time examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
@@ -71,7 +71,7 @@ def train_generator(optimizer, data_fake):
 if __name__ == "__main__":
  # data for training and validation
  X_train, _, _, _ = mnist()
-  ds_noise = Tensor.randn(64, 128)
+  ds_noise = Tensor.randn(64, 128, requires_grad=False)
  # parameters
  epochs, batch_size, k = 300, 512, 1
  sample_interval = epochs // 10
@@ -21,8 +21,6 @@ def compile(onnx_file):
  # TODO this seems dumb
  input_types = {k:(dtypes.float32 if v is dtypes.float16 else v) for k,v in input_types.items()}
  Tensor.manual_seed(100)
-  # replace symbolic dimensions (e.g. 'b' for dynamic batch) with 1
-  input_shapes = {k:tuple(s if isinstance(s, int) else 1 for s in shp) for k,shp in input_shapes.items()}
  inputs = {k:Tensor(Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize().numpy(), device='NPY') for k,shp in sorted(input_shapes.items())}
  if not getenv("NPY_IMG"):
    inputs = {k:Tensor(v.numpy(), device=Device.DEFAULT).realize() if 'img' in k else v for k,v in inputs.items()}
@@ -87,7 +85,7 @@ def test_vs_compile(run, inputs, test_val=None):
    step_times.append((et-st)*1e3)
    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")

-  if (assert_time:=getenv("ASSERT_MIN_STEP_TIME", 0.0)):
+  if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
    min_time = min(step_times)
    assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"

@@ -104,7 +102,7 @@ def test_vs_compile(run, inputs, test_val=None):
 def test_vs_onnx(new_inputs, test_val, onnx_file, tol):
  import onnx
  import onnxruntime as ort
-
+  
  onnx_inputs = {k:v.numpy() for k,v in new_inputs.items()}
  onnx_model = onnx.load(onnx_file)

@@ -137,7 +135,7 @@ def bench(run, inputs):
 if __name__ == "__main__":
  if getenv("RUN_PICKLE"):
    with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
-    inputs = {name: Tensor(Tensor.randn(*view.shape, dtype=dtype).numpy(), device=device)
+    inputs = {name: Tensor(Tensor.randn(*[int(s) for s in view.src[1].arg], dtype=dtype).numpy(), device=device)
              for name, (view, _vars, dtype, device) in zip(pickle_loaded.captured.expected_names, pickle_loaded.captured.expected_input_info)}
    test_vs_compile(pickle_loaded, inputs)
  else:
@@ -164,8 +164,8 @@ elif cmd == "train":
    x_img = image_load(samples_base + "/" + str(sample_idx) + "a.png")
    y_img = image_load(samples_base + "/" + str(sample_idx) + "b.png")

-    sample_x = Tensor(x_img)
-    sample_y = Tensor(y_img)
+    sample_x = Tensor(x_img, requires_grad = False)
+    sample_y = Tensor(y_img, requires_grad = False)

    # magic code roughly from readme example
    # An explanation, in case anyone else has to go down this path:
@@ -23,7 +23,7 @@ def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], Li

  def name_of(bu:UOp, is_out:bool) -> str:
    nonlocal n
-    if bu.op is Ops.PARAM: key, name, size = ("in", bu.arg.slot), f"input{bu.arg.slot}", prod(bu.shape)*bu.dtype.itemsize
+    if bu.op is Ops.PARAM: key, name, size = ("in", bu.arg), f"input{bu.arg}", prod(bu.shape)*bu.dtype.itemsize
    else:
      b = bu.buffer
      key, size = (id(b.base), b.offset, b.size, b.dtype), b.size*b.dtype.itemsize
@@ -122,7 +122,7 @@ def eval_custom_matmul(fxn, dt=dtypes.float):
  with Context(DEBUG=0): Tensor.realize(a, b)

  ets = []
-  with Context(DEBUG=max(2, DEBUG.value)):
+  with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2 if dt == dtypes.half else 0):
    for _ in range(NUM_RUNS):
      GlobalCounters.reset()
      tst = Tensor.custom_kernel(c, a, b, fxn=fxn)[0].realize()
@@ -218,7 +218,7 @@ if __name__ == "__main__":
  ref.realize()

  GlobalCounters.reset()
-  with Context(DEBUG=max(2, DEBUG.value)):
+  with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2):
    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
    tst.realize()
  print(f"{(N*M*K*2 / GlobalCounters.time_sum_s)*1e-12:.2f} REAL TFLOPS")
@@ -127,7 +127,7 @@ if __name__ == "__main__":


  GlobalCounters.reset()
-  with Context(DEBUG=max(2, DEBUG.value)):
+  with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2):
    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
    tst.realize()
  print(f"{(N*M*K*2 / GlobalCounters.time_sum_s)*1e-12:.2f} REAL TFLOPS")
@@ -5,13 +5,9 @@ from tinygrad.device import Buffer, BufferSpec, Compiled, Device, MultiBuffer
 from tinygrad.dtype import dtypes
 from tinygrad.engine.jit import GraphRunner
 from tinygrad.engine.realize import get_call_outs_ins, get_runtime
-from tinygrad.helpers import round_up, ceildiv
-from tinygrad.runtime.support.memory import BumpAllocator
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, graph_rewrite
-from extra.hcq2.hcq2 import HCQ2Compiled, HCQ2DeviceCtx, HCQ2LowerCtx, pm_prep_runtime, pm_lower_ops
-from extra.hcq2.hcq2 import pm_split_into_queues, pm_add_barriers, pm_add_signals
-from extra.hcq2.hcq2 import pm_bufferize, pm_lift_patches_to_cmdbuf, pm_resolve_patches, pm_parametrize_host_buffers
-from extra.hcq2.hcq2 import pm_add_timeline_inc, pm_callify, pm_calc_kernargs_sizes
+from extra.hcq2.hcq2 import HCQ2Compiled, HCQ2DeviceCtx, HCQ2LowerCtx, prep_runtime, pm_lower_kernargs, pm_lower_ops
+from extra.hcq2.hcq2 import pm_split_into_queues, pm_add_barriers, pm_add_signals, build_host_program

 # **************** insert deps ****************

@@ -26,20 +22,18 @@ def insert_deps(ctx:HCQ2Graph, linear:UOp) -> UOp:
  return linear.replace(src=tuple(src))
 pm_insert_deps = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), insert_deps)])

-pm_replace_params = PatternMatcher([
-  (UPat(Ops.PARAM, name="p"), lambda ctx, p: ctx.input_addrs_uop.index(UOp.const(dtypes.int, p.arg))),
-  (UPat(Ops.SLICE, src=(UPat(Ops.INDEX, name="addr"), UPat(Ops.CONST, dtype=dtypes.weakint, name="off")), name="bv"),
-    lambda ctx, bv, addr, off: addr.cast(dtypes.uint64) + UOp.const(dtypes.uint64, off.arg * ctx.input_uops[addr.src[1].arg].dtype.itemsize)),
-])
+def replace_params(ctx:HCQ2Graph, call:UOp) -> UOp|None:
+  if not any(x.op is Ops.PARAM for x in call.src[1:]): return None
+  return call.replace(src=tuple(ctx.input_addrs_uop[x.arg] if x.op is Ops.PARAM else x for x in call.src))
+pm_replace_params = PatternMatcher([(UPat(Ops.CALL, name="call", allow_any_len=True), replace_params)])

 # **************** graph-only passes ****************

 def alloc_queue_sig(ctx:HCQ2Graph, q:UOp) -> None:
  if q.arg in ctx.queue_sigs: return None
-  dev = q.arg[0][0]  # TODO: multi device
-  buf = Buffer(dev, 0x100, dtypes.uint8, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
+  buf = Buffer(q.arg[0], 0x100, dtypes.uint8, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
  ctx.queue_sig_bufs.append(buf)
-  ctx.queue_sigs[q.arg] = UOp.from_buffer(buf, dev)
+  ctx.queue_sigs[q.arg] = UOp.from_buffer(buf, q.arg[0])
  return None
 pm_alloc_queue_sigs = PatternMatcher([(UPat(Ops.LINEAR, src=UPat({Ops.PROGRAM, Ops.COPY}), name="q"), alloc_queue_sig)])

@@ -75,13 +69,12 @@ def drop_dead_stores(ctx:HCQ2Graph, outer:UOp) -> UOp:
  return outer.replace(src=tuple(q.replace(src=tuple(x for x in q.src if x.op is not Ops.STORE or x in live)) for q in outer.src))
 pm_drop_dead_stores = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR), name="outer"), drop_dead_stores)])

-def add_queue_sig_resets(ctx:HCQ2Graph, x:UOp, cmdbuf:UOp) -> UOp|None:
-  if not ctx.queue_sig_bufs or cmdbuf.tag not in ("compute", "copy"): return None
-  resets = tuple((b:=UOp.from_buffer(sig)).index(UOp.const(dtypes.int, 0), dtype=b.dtype.ptr())
-                 .cast(dtypes.uint64.ptr()).store(UOp.const(dtypes.uint64, 0)) for sig in ctx.queue_sig_bufs)
-  return x.replace(src=x.src + resets)
-pm_add_queue_sig_resets = PatternMatcher([(UPat(Ops.AFTER, src=(UPat(Ops.BUFFER, name="cmdbuf"),), allow_any_len=True, name="x"),
-                                           add_queue_sig_resets)])
+def add_queue_sig_resets(ctx:HCQ2Graph, outer:UOp) -> UOp|None:
+  if not ctx.queue_sig_bufs: return None
+  resets = tuple(ctx.hcq_ctx.host_param(sig).index(UOp.const(dtypes.int, 0), ptr=True).cast(dtypes.uint64.ptr())
+                 .store(UOp.const(dtypes.uint64, 0)) for sig in ctx.queue_sig_bufs)
+  return outer.replace(src=tuple(c.replace(src=c.src + resets) if c.op is Ops.AFTER else c.after(*resets) for c in outer.src))
+pm_add_queue_sig_resets = PatternMatcher([(UPat(Ops.LINEAR, name="outer"), add_queue_sig_resets)])

 # **************** Graph ****************

@@ -92,12 +85,16 @@ class HCQ2Graph(GraphRunner):
    self.hcq_ctx = HCQ2LowerCtx(name="hcq_graph")

    self.input_addrs = Buffer("CPU", max(len(input_uops), 1), dtypes.uint64, preallocate=True)
-    self.input_addrs_uop = UOp.from_buffer(self.input_addrs, "CPU")
+    self.input_addrs_uop = self.hcq_ctx.host_param(self.input_addrs)

    self.linear = graph_rewrite(self.linear, pm_insert_deps, ctx=self, name="hcq: insert deps", walk=True)
+    self.linear, sizes = prep_runtime(self.hcq_ctx, self.linear)
+    for dev_name, sz in sizes.items():
+      buf = Buffer(dev_name, sz, dtypes.uint8, options=BufferSpec(cpu_access=True), preallocate=True)
+      self.hcq_ctx.devs[dev_name] = HCQ2DeviceCtx(dev_name, UOp.from_buffer(buf, dev_name), UOp.const(dtypes.uint64, buf._buf.va_addr))
+
    self.linear = graph_rewrite(self.linear, pm_replace_params, ctx=self, name="hcq: replace params", walk=True)
-    self.linear = graph_rewrite(self.linear, pm_prep_runtime, ctx=self.hcq_ctx, name="hcq: prepare runtime")
-    self.linear = graph_rewrite(self.linear, pm_lower_ops, ctx=self.hcq_ctx, name="hcq: lower ops")
+    self.linear = graph_rewrite(self.linear, pm_lower_kernargs + pm_lower_ops, ctx=self.hcq_ctx, name="hcq: lower ops")

    # per-queue signal state — populated as a side-effect by pm_alloc_queue_sigs walking the lowered linear.
    self.queue_sig_bufs:list[Buffer] = []
@@ -110,20 +107,9 @@ class HCQ2Graph(GraphRunner):
    self.linear = graph_rewrite(self.linear, pm_optimize_queue_deps, ctx=self, name="hcq: optimize queue deps", walk=True)
    self.linear = graph_rewrite(self.linear, pm_drop_dead_stores, ctx=self, name="hcq: drop dead stores")
    self.linear = graph_rewrite(self.linear, pm_add_signals, ctx=self.hcq_ctx, name="hcq: add signals", walk=True)
-    self.linear = graph_rewrite(self.linear, pm_add_timeline_inc, ctx=self.hcq_ctx, name="hcq: add submit", walk=True)
    self.linear = graph_rewrite(self.linear, self.dev.pm_lower, ctx=self.hcq_ctx, name=f"hcq: encode cmdbuf {self.dev.device}", walk=True)
-
-    graph_rewrite(self.linear, pm_calc_kernargs_sizes, ctx=(sizes:={}), name=None)
-    for dev_name, sz in sizes.items():
-      buf = Buffer(dev_name, sz, dtypes.uint8, options=BufferSpec(cpu_access=True), preallocate=True)
-      self.hcq_ctx.dev_ctx[dev_name] = HCQ2DeviceCtx(dev_name, UOp.from_buffer(buf, dev_name), UOp.const(dtypes.uint64, buf._buf.va_addr))
-
-    self.linear = graph_rewrite(self.linear, pm_bufferize, ctx=self.hcq_ctx, bottom_up=True, name="realize binaries")
-    self.linear = graph_rewrite(self.linear, pm_lift_patches_to_cmdbuf, ctx=self.hcq_ctx, bottom_up=False, name="lift patches to cmdbuf")
-    self.linear = graph_rewrite(self.linear, pm_resolve_patches, ctx=self.hcq_ctx, bottom_up=False, name="simplify patches")
    self.linear = graph_rewrite(self.linear, pm_add_queue_sig_resets, ctx=self, name="hcq: add queue sig resets", walk=True)
-    self.linear = graph_rewrite(self.linear, pm_parametrize_host_buffers, ctx=self.hcq_ctx, bottom_up=True, name="parametrize host buffers")
-    self.host_call = graph_rewrite(self.linear, pm_callify, ctx=self.hcq_ctx, name="hcq: callify")
+    self.host_call = build_host_program(self.hcq_ctx, self.linear, None, self.dev)

    self.host_rt, self.host_globals = get_runtime("CPU", self.host_call.src[0]), self.host_call.src[0].arg.globals

@@ -1,35 +1,39 @@
 from __future__ import annotations
 from typing import cast, Callable, TypeVar, Generic, Any, TYPE_CHECKING
-import struct, functools, time, collections, importlib, itertools
+import struct, functools, time, collections
 from dataclasses import replace
 if TYPE_CHECKING: from tinygrad.engine.realize import ExecContext
-from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, mv_address, round_up, DEBUG, dedup, pluralize
-from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer
-from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, track_rewrites, GroupOp
-from tinygrad.uop.symbolic import symbolic_simple, symbolic
-from tinygrad.dtype import dtypes, DType
+from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, mv_address, round_up, DEBUG, dedup
+from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator
+from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, track_rewrites
+from tinygrad.dtype import dtypes
 from dataclasses import dataclass, field
 from tinygrad.runtime.support.memory import BumpAllocator
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.renderer import Renderer, Estimates
-from tinygrad.engine.realize import to_program, track_stats, get_call_arg_uops, resolve_params, pm_flatten_linear
+from tinygrad.engine.realize import to_program, track_stats, get_call_arg_uops, resolve_params

 HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQ2Compiled')

 class HCQ2Compiled(Compiled):
-  timestamp_divider: float = 1000.0 # GPU timestamp counter ticks per microsecond; override per device
+  """
+  A base class for devices compatible with the HCQ (Hardware Command Queue) API.
+  """
+  timestamp_divider: float = 1000.0  # GPU timestamp counter ticks per microsecond; override per device

-  def __init__(self, device:str, allocator:'HCQAllocator', compilers:list[type[Renderer]], runtime, can_recover:bool=False, arch=None):
+  def __init__(self, device:str, allocator:'HCQAllocator', compilers:list[type[Renderer]], runtime,
+               kernargs_size=(16 << 20), can_recover:bool=False, arch=None):
    self.device_id:int = int(device.split(":")[1]) if ":" in device else 0

-    # default pm bufferize
-    self.pm_bufferize = PatternMatcher([
-      (UPat(Ops.BUFFER, tag="timeline_signal"), lambda ctx: ctx.timeline_signal),
-      (UPat(Ops.BUFFER, tag="timeline_value"), lambda ctx: ctx.timeline_value),
-      (UPat(Ops.BUFFER, name="b"), lambda ctx, b: Buffer(ctx.device, b.arg, b.dtype, options=BufferSpec(host=True, uncached=True, cpu_access=True))),
-    ])
+    from extra.hcq2.graph.hcq import HCQ2Graph
+    super().__init__(device, allocator, compilers, lambda *a, **kw: None, HCQ2Graph, arch=arch)

-    super().__init__(device, allocator, compilers, lambda *a, **kw: None, None, arch=arch)
+    self.kernargs_size = kernargs_size
+    self.kernargs_offset_allocator:BumpAllocator = BumpAllocator(kernargs_size, wrap=True)
+
+  @functools.cached_property
+  def kernargs_buf(self) -> Buffer:
+    return Buffer(self.device, self.kernargs_size, dtypes.uint8, options=BufferSpec(cpu_access=True), preallocate=True)

  @functools.cached_property
  def timeline_signal(self) -> Buffer:
@@ -55,6 +59,14 @@ class HCQ2Compiled(Compiled):

  def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent.

+  def _realloc(self, oldbuf:HCQ2Buffer|None, new_size:int, options:BufferSpec|None=None, force=False) -> tuple[HCQ2Buffer, bool]:
+    if oldbuf is not None: self.allocator.free(oldbuf, oldbuf.size, options=options)
+    try: buf, realloced = self.allocator.alloc(new_size, options=options), True
+    except MemoryError:
+      if force: raise
+      buf, realloced = self.allocator.alloc(oldbuf.size if oldbuf is not None else new_size, options=options), False
+    return buf, realloced
+
  def count(self) -> int: return self.iface.count if hasattr(self, 'iface') else 1

  def _select_iface(self):
@@ -113,7 +125,7 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
  def _copy(self, dst:Buffer, src:Buffer):
    from tinygrad.engine.realize import run_linear
    su = UOp.from_buffer(src)
-    run_linear(UOp(Ops.LINEAR, dtypes.void, (su.copy_to_device(dst.device).call(UOp.from_buffer(dst), su),)), update_stats=False)
+    run_linear(UOp(Ops.LINEAR, dtypes.void, (su.copy_to_device(dst.device).call(UOp.from_buffer(dst), su),)), jit=True, update_stats=False)

  def _copyin(self, dest:HCQ2Buffer, src:memoryview):
    s = Buffer(self.dev.device, len(src), dtypes.uint8, options=BufferSpec(host=True), preallocate=True)
@@ -126,270 +138,240 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
    self.dev.synchronize()
    dest[:] = d._buf.cpu_view()[:len(dest)]

-  # def _as_buffer(self, buf): return buf.cpu_view().mv
+  def _as_buffer(self, buf): return buf.cpu_view().mv

-def unwrap_after(uop):
-  while uop.op is Ops.AFTER: uop = uop.src[0]
-  return uop
+# **************** lower context ****************
+
+@dataclass
+class HCQ2DeviceCtx:
+  device:str                       # device name; resolve to instance via Device[device]
+  kernargs_host:UOp                # UOp whose .buffer is dev.kernargs_buf (BUFFER UOp in runtime, PARAM in graph)
+  kernargs_gpu:UOp                 # va_addr const of dev.kernargs_buf
+  kernargs_allocator:BumpAllocator = field(default_factory=lambda: BumpAllocator(2 << 20, wrap=False))
+
+@dataclass
+class HCQ2LowerCtx:
+  name:str
+  inputs:list[Buffer] = field(default_factory=list)
+  holds:list[UOp] = field(default_factory=list)
+  devs:dict[str, HCQ2DeviceCtx] = field(default_factory=dict)
+
+  def host_param(self, buf:Buffer) -> UOp:
+    if buf not in self.inputs: self.inputs.append(buf)
+    return UOp.placeholder((buf.size,), buf.dtype, self.inputs.index(buf))

 class HCQEncoder:
-  def __init__(self): self.blob, self.patches = b'', []
+  def __init__(self, ctx:HCQ2LowerCtx, dev:HCQ2Compiled): self.ctx, self.dev, self.blob, self.patches, self.deps = ctx, dev, b'', [], []

-  def get_dev_addr(self, uop:UOp) -> UOp:
-    if unwrap_after(uop).op not in (Ops.BUFFER, Ops.SLICE, Ops.BINARY, Ops.MSTACK, Ops.MSELECT): return uop
-    return UOp(Ops.GETADDR, dtypes.uint64, src=(uop, UOp(Ops.DEVICE, arg=self.dev.device)))
+  @property
+  def src(self) -> tuple[UOp, ...]: return tuple(self.patches + dedup(self.deps))
+
+  def get_dev_addr(self, uop:UOp) -> sint|UOp:
+    while uop.op is Ops.AFTER:
+      self.deps.extend(uop.src[1:])
+      uop = uop.src[0]
+    if isinstance(val:=uop.ssimplify(), UOp): self.deps.append(uop)
+    return uop.buffer.get_buf(self.dev.device).va_addr if uop.op in (Ops.BUFFER, Ops.BUFFER_VIEW) else val

  def append(self, *data, dtype=dtypes.uint32):
    for d in data:
      if isinstance(d, int): self.blob += struct.pack(f'<{dtype.fmt}', d)
+      elif d.op is Ops.CONST: self.blob += struct.pack(f'<{dtype.fmt}', d.arg)
      else:
-        self.patches.append((len(self.blob), self.get_dev_addr(d), dtype))
+        self.patches.append(UOp(Ops.PATCH, dtype, src=(d,), arg=len(self.blob)))
        self.blob += struct.pack(f'<{dtype.fmt}', 0)

  def q(self, *values): self.append(*values)

-  def uop(self, dev:str|tuple[str, ...], tag:str|None=None) -> UOp:
-    buf = UOp.new_buffer(dev, len(self.blob), dtypes.uint8)
-    if tag: buf = buf.rtag(tag)
-    blob_uop = UOp(Ops.BINARY, dtypes.void, src=(), arg=self.blob)
-    stores = [buf.index(UOp.const(dtypes.int, off), dtype=buf.dtype.ptr()).cast(dt.ptr()).store(val.cast(dt)) for off, val, dt in self.patches]
-    return buf.after(buf.store(blob_uop), *stores)
-
-# *****************
-# 0. helpers
-
-HCQ_DEVS = frozenset(("AMD",))
-HCQ_P2P_DEVS = HCQ_DEVS | frozenset(("CPU",))
-
-def to_tuple(d): return d if isinstance(d, tuple) else (d,)
-
-def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c
-
-# *****************
-# 1.1. prep runtimes: staging copies
-
-def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS) and not all_devices_in(b.device, HCQ_P2P_DEVS)
-
-def stage_copy(dst:UOp, src:UOp) -> UOp|None:
-  if not (_need_staging(src, dst) or _need_staging(dst, src)): return None
-
-  stage = UOp.new_buffer("CPU", src.buffer.nbytes, dtypes.uint8)
-  return UOp(Ops.LINEAR, dtypes.void, (src.copy_to_device("CPU").call(stage, src), stage.copy_to_device(dst.device).call(dst, stage)))
-pm_insert_copy_staging = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy)])
-
-# *****************
-# 1.2. prep runtimes: programs/kernargs
-
-@functools.cache
-def get_pm_prep_program(name:str) -> PatternMatcher|None:
-  try:
-    importlib.import_module(f'tinygrad.runtime.ops_{name.lower()}') # TODO: remove that
-    return importlib.import_module(f'extra.hcq2.ops_{name.lower()}2').pm_prep_program
-  except ImportError: return None
-
-def prep_program(call:UOp, prg:UOp) -> UOp|None:
-  dev = call.src[1].device
-  if (pm:=get_pm_prep_program(to_tuple(dev)[0].split(":")[0])) is None or (lowered:=pm.rewrite(prg)) is None: return None
-
-  data, image_bytes = lowered
-  buf = UOp.new_buffer(dev, len(image_bytes), dtypes.uint8).rtag("program")
-  blob = UOp(Ops.BINARY, dtypes.void, src=(), arg=image_bytes)
-  return call.replace(src=(prg.replace(src=(buf.after(buf.store(blob)),), arg=(data, prg.arg)),) + call.src[1:])
-
-def prep_kernargs(call:UOp, prg:UOp) -> UOp:
-  data, info = prg.arg
-  patches = [(i*dtypes.uint64.itemsize, UOp(Ops.GETADDR, dtypes.uint64, src=(call.src[1+gi], UOp(Ops.DEVICE, arg=call.src[1+gi].device))),
-              dtypes.uint64) for i,gi in enumerate(info.globals)] \
-          + [(len(info.globals)*dtypes.uint64.itemsize + i*dtypes.uint32.itemsize, v, dtypes.uint32) for i,v in enumerate(info.vars)]
-
-  buf = UOp.new_buffer(call.src[1].device, data.kernargs_alloc_size, dtypes.uint8).rtag("kernargs")
-  kernargs = buf.after(*tuple(buf.index(UOp.const(dtypes.int, o), dtype=buf.dtype.ptr()).cast(dt.ptr()).store(val.cast(dt)) for o, val, dt in patches))
-
-  return call.replace(src=(prg.replace(src=prg.src + (kernargs,), arg=(data, info)),) + call.src[1:])
+# **************** prep runtime ****************

 pm_prep_runtime = PatternMatcher([
-  # bind generic PROGRAM device to the call's actual dev(s), then run device-specific lowering
-  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"),),
-    name="call", allow_any_len=True), prep_program),
-
-  # lower kernargs (PROGRAM.src[0] is now AFTER(BUFFER, COPY) — the lowered program image)
-  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.AFTER),), name="prg"),), name="call", allow_any_len=True), prep_kernargs),
+  # device-specific lowering of the program
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE), UPat(), UPat(), UPat(Ops.BINARY)), name="p"),), name="c", allow_any_len=True),
+    lambda ctx, c, p: c.replace(src=(Device[p.src[1].arg].pm_lower.rewrite(p, ctx),) + c.src[1:])),
 ])

-# *****************
-# 2.1. lowering to hcq ir
+def calc_kernargs_sizes(ctx:dict[str,int], u:UOp) -> None:
+  d = u.src[0].buffer.device
+  ctx[d] = ctx.get(d, 0) + round_up(u.arg[0].kernargs_alloc_size, 16)
+pm_calc_kernargs_sizes = PatternMatcher([(UPat(Ops.PROGRAM, name="u"), calc_kernargs_sizes)])

-def lower_program(call:UOp, prg:UOp) -> UOp:
-  q = UOp(Ops.LINEAR, dtypes.void, (prg,), arg=(call.src[1].device, "COMPUTE"))
-  return call.replace(src=(q,) + call.src[1:]).rtag('hcq')
+# **************** lower kernargs ****************

-def lower_copy(call:UOp, copy:UOp) -> UOp|None:
+def lower_kernargs(ctx:HCQ2LowerCtx, call:UOp, prg:UOp) -> UOp:
+  data, info = prg.arg
+  # after amd_build_program, prg.src is (BUFFER_lib_gpu,); the buffer's device names the device
+  dctx = ctx.devs[prg.src[0].buffer.device]
+
+  enc = HCQEncoder(ctx, Device[dctx.device])
+  for gi in info.globals: enc.append(enc.get_dev_addr(call.src[1+gi]), dtype=dtypes.uint64)
+  for v in info.vars: enc.append(v, dtype=dtypes.uint32)
+
+  args_off = dctx.kernargs_allocator.alloc(data.kernargs_alloc_size, 16)
+  dctx.kernargs_host.buffer.view(len(enc.blob), dtypes.uint8, args_off).ensure_allocated().as_memoryview(force_zero_copy=True)[:] = enc.blob
+
+  args_uop = (dctx.kernargs_gpu + args_off).after(dctx.kernargs_host.after(*tuple(p.replace(arg=p.arg+args_off) for p in enc.patches)))
+  return call.replace(src=(prg.replace(src=prg.src + (args_uop,), arg=(data, info)),) + call.src[1:])
+
+pm_lower_kernargs = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.BUFFER),), name="prg"),), name="call", allow_any_len=True), lower_kernargs),
+])
+
+# **************** lower ops ****************
+
+def lower_program(ctx:HCQ2LowerCtx, call:UOp, prg:UOp) -> UOp:
+  q = UOp(Ops.LINEAR, dtypes.void, (prg,), arg=(prg.src[0].buffer.device, "COMPUTE"))
+  return UOp(Ops.LINEAR, dtypes.void, (q,), tag=call.tag)
+
+def lower_copy(ctx:HCQ2LowerCtx, call:UOp, copy:UOp) -> UOp:
  dst, src = call.src[1], call.src[2]
-  if (hcq_dev:=next((b.device for b in (dst, src) if b.device.split(":")[0] in HCQ_DEVS), None)) is None: return None
-
-  q = UOp(Ops.LINEAR, dtypes.void, (UOp(Ops.COPY, dtypes.void, src=(dst, src), arg=src.buffer.nbytes),), arg=(hcq_dev, "COPY"))
-  return call.replace(src=(q,) + call.src[1:]).rtag('hcq')
+  q = UOp(Ops.LINEAR, dtypes.void, (UOp(Ops.COPY, dtypes.void, src=(dst, src), arg=src.buffer.nbytes),), arg=(dst.buffer.device, "COPY"))
+  return UOp(Ops.LINEAR, dtypes.void, (q,), tag=call.tag)

 pm_lower_ops = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.AFTER), UPat(Ops.AFTER)), name="prg"),), name="call", allow_any_len=True), lower_program),
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, src=(UPat(Ops.BUFFER), UPat()), name="prg"),), name="call", allow_any_len=True), lower_program),
  (UPat(Ops.CALL, src=(UPat(Ops.COPY, name="copy"),), name="call", allow_any_len=True), lower_copy),
 ])

-# *****************
-# 2.2. queue split
+# **************** split into queues ****************

-# def split_into_queues(linear:UOp) -> UOp:
-#   out = []
-#   for k, grp in itertools.groupby(linear.src, lambda c: c.src[0].arg if c.op is Ops.CALL and c.src[0].op is Ops.LINEAR else None):
-#     if k is None: out.extend(grp)
-#     else:
-#       calls = list(grp)
-#       items = tuple(x for c in calls for x in c.src[0].src)
-#       args = tuple(a for c in calls for a in c.src[1:])
-#       out.append(calls[0].replace(src=(UOp(Ops.LINEAR, dtypes.void, items, arg=k),) + args))
-#   return linear.replace(src=tuple(out))
-# pm_split_into_queues = PatternMatcher([(UPat(Ops.LINEAR, name="linear"), split_into_queues)])
+def split_into_queues(ctx:HCQ2LowerCtx, outer:UOp) -> UOp:
+  groups:dict[tuple, list[UOp]] = collections.defaultdict(list)
+  for child in outer.src:
+    wrapper = child.src[0] if child.op is Ops.AFTER else child
+    for q in wrapper.src: groups[q.arg].extend(q.src)
+  return outer.replace(src=tuple(UOp(Ops.LINEAR, dtypes.void, tuple(cmds), arg=k) for k, cmds in groups.items()))
+pm_split_into_queues = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR, src=UPat(Ops.LINEAR)).or_after(), name="outer"), split_into_queues)])

-# *****************
-# 2.3. barriers / signals / timeline inc
+# **************** add signals (runtime) ****************

-def add_barriers(call:UOp, q:UOp) -> UOp:
-  return call.replace(src=(q.replace(src=(UOp(Ops.BARRIER, dtypes.void), *q.src)),) + call.src[1:])
-pm_add_barriers = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.LINEAR, name="q"),), name="call", allow_any_len=True), add_barriers)])
+def add_signals(ctx:HCQ2LowerCtx, outer:UOp) -> UOp:
+  def wrap(q:UOp) -> UOp:
+    (dev_name, qname), devs = q.arg, {q.arg[0]} | {u.buffer.device for u in q.toposort() if u.op in (Ops.BUFFER, Ops.BUFFER_VIEW)}
+    sigs_tls = [(UOp.from_buffer(Device[d].timeline_signal), ctx.host_param(Device[d].timeline_value)) for d in sorted(devs) if d.startswith("AMD")]
+    return q.replace(src=(*(s.wait(t[0]-1) for s,t in sigs_tls), *q.src, *(s.store(t[0]) for s,t in sigs_tls)), arg=qname)
+  return outer.replace(src=tuple(wrap(q) for q in outer.src))

-def add_signals(call:UOp, q:UOp) -> UOp:
-  sig = UOp.new_buffer(q.arg[0], 0x100, dtypes.uint8).rtag("timeline_signal")
-  tl = UOp.new_buffer(q.arg[0], 1, dtypes.uint64).rtag("timeline_value").index(UOp.const(dtypes.int, 0))
-  return call.replace(src=(q.replace(src=(sig.wait(tl-1), *q.src, sig.store(tl)), arg=q.arg),) + call.src[1:])
-pm_add_signals = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.LINEAR, name="q"),), name="call", allow_any_len=True), add_signals)])
+pm_add_barriers = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR), name="outer"),
+  lambda ctx, outer: outer.replace(src=tuple(q.replace(src=(UOp(Ops.BARRIER, dtypes.void), *q.src)) for q in outer.src)))])

-# *****************
-# 3.1. encode cmdbufs
+pm_add_signals = PatternMatcher([(UPat(Ops.LINEAR, src=UPat(Ops.LINEAR), name="outer"), add_signals)])

-@functools.cache
-def get_pm_lower(name:str) -> PatternMatcher|None:
-  try:
-    importlib.import_module(f'tinygrad.runtime.ops_{name.lower()}') # TODO: remove that
-    return importlib.import_module(f'extra.hcq2.ops_{name.lower()}2').pm_lower
-  except ImportError: return None
+# **************** build host program ****************

-def encode_cmdbuf(call:UOp, q:UOp) -> UOp|None:
-  if (pm:=get_pm_lower(to_tuple(q.arg[0])[0].split(":")[0])) is None or (encoded:=pm.rewrite(q)) is None: return None
-  return call.replace(src=(encoded,) + call.src[1:])
-pm_encode_cmdbufs = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.LINEAR, name="q"),), name="call", allow_any_len=True), encode_cmdbuf)])
+def resolve_cmdbuf(ctx:HCQ2LowerCtx, blob:UOp) -> UOp:
+  inner = blob.src[0] if blob.op is Ops.AFTER else blob
+  dev_name, qtype = inner.tag

-# *****************
-# 3.2. add timeline inc
+  # prepare the cmdbuf and make it a param
+  bb = Buffer("CPU", len(inner.arg)//4, dtypes.uint32, preallocate=True)
+  bb.copyin(memoryview(bytearray(inner.arg)))
+  bb_param = ctx.host_param(bb)

-def add_timeline_inc(call:UOp, s:UOp) -> UOp:
-  tl = UOp.new_buffer(s.device, 1, dtypes.uint64).rtag("timeline_value")
-  return call.replace(src=(tl.after(s).index(UOp.const(dtypes.int, 0), dtype=tl.dtype.ptr()).store(tl.index(UOp.const(dtypes.int, 0)) + 1),) + call.src[1:])
-pm_add_timeline_inc = PatternMatcher([(UPat(Ops.CALL, tag="hcq", src=(UPat(name="s"),), name="call", allow_any_len=True), add_timeline_inc)])
+  submit_cf = UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(bb_param.after(*(blob.src[1:] if blob.op is Ops.AFTER else ())),),
+                  arg=f"submit_{qtype.lower()}", tag=dev_name)

-# *****************
-# 3.3. lift patches to the command buffer (root)
+  # increment the timeline value
+  tl = ctx.host_param(Device[dev_name].timeline_value)
+  return tl.after(UOp(Ops.BARRIER, dtypes.void, src=(submit_cf,))).index(UOp.const(dtypes.int, 0), ptr=True).store(tl[0] + 1)

-def lift_patches_to_cmdbuf(cmdbuf:UOp) -> UOp|None:
-  if not (patches:=dedup(u for store in cmdbuf.src[1:] for u in store.toposort() if u.op is Ops.AFTER)): return None
-  deps = tuple(d for p in patches for d in p.src[1:])
-  return cmdbuf.replace(src=cmdbuf.src + deps).substitute({p: p.src[0] for p in patches})
-pm_lift_patches_to_cmdbuf = PatternMatcher([
-  (UPat(Ops.AFTER, src=(UPat(Ops.BUFFER, tag={"compute", "copy"}),), allow_any_len=True, name="cmdbuf"), lift_patches_to_cmdbuf),
+def resolve_patches(ctx:HCQ2LowerCtx, buf:UOp) -> UOp|None:
+  inner = buf.src[0]
+
+  # buffer is accessed from the launcher, so transform it to a host param
+  if inner.op is Ops.BUFFER: inner = ctx.host_param(inner.buffer)
+
+  return inner.after(*(inner.index(UOp.const(dtypes.int, p.arg//inner.dtype.base.itemsize), ptr=True).cast(p.dtype.ptr()).store(p.src[0].cast(p.dtype))
+                           if p.op is Ops.PATCH else p for p in buf.src[1:]))
+
+def resolve_ref_buffers(ctx:HCQ2LowerCtx, buf:UOp) -> UOp:
+  if buf not in ctx.holds: ctx.holds.append(buf)
+  return UOp(Ops.NOOP)
+
+def hcq_callify(ctx:HCQ2LowerCtx, sink:UOp) -> UOp:
+  call = to_program(sink, Device["CPU"].renderer).call(*[UOp.from_buffer(b, "CPU") if isinstance(b, Buffer) else b for b in ctx.inputs])
+  return call.replace(src=call.src + (UOp(Ops.BIND, dtypes.void, src=tuple(ctx.holds)),)) if ctx.holds else call
+
+pm_create_host_sink = PatternMatcher([
+  (UPat(Ops.LINEAR, name="l", allow_any_len=True), lambda ctx, l: UOp.sink(*l.src, arg=KernelInfo(name=ctx.name, estimates=Estimates()), tag=1))
 ])

-# *****************
-# 4. bufferize placeholders: replace placeholders with real buffers.
-
-def bufferize_buf(buf:UOp) -> UOp|None:
-  if buf.tag is None: return None
-  uops = tuple(UOp.from_buffer((dv:=Device[dev]).pm_bufferize.rewrite(buf, ctx=dv), dev) for dev in to_tuple(buf.src[1].arg))
-  return uops[0] if len(uops) == 1 else UOp(Ops.MSTACK, uops[0].dtype, uops)
-pm_bufferize = PatternMatcher([(UPat(Ops.BUFFER, name="buf"), bufferize_buf)])
-
-# *****************
-# 5.1. capture buffers reachable from each hcq call as BIND, so we don't drop their refs
-
-def hold_call_buffers(call:UOp) -> UOp|None:
-  if not (bufs:=tuple(dedup(u for u in call.src[0].toposort() if u.op is Ops.BUFFER and u not in call.src))): return None
-  return call.replace(src=call.src + (UOp(Ops.BIND, dtypes.void, src=bufs),))
-pm_hold_call_buffers = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), hold_call_buffers)])
-
-# *****************
-# 5.2. resolve patches
-
-def push_stack(op, s): return UOp(Ops.STACK, op.dtype.scalar().vec(len(s.src)),
-  tuple(op.replace(dtype=op.dtype.scalar(), src=tuple(x if y is s else y for y in op.src)) for x in s.src))
-
-def fold_blob_store(buf:UOp, blob:UOp) -> UOp:
-  for b in (buf.src if buf.op is Ops.MSTACK else (buf,)): b.buffer.ensure_allocated()._buf.cpu_view().mv.cast('B')[:len(blob.arg)] = blob.arg
-  return UOp(Ops.NOOP)
-
-def fold_const_store(buf:UOp, off:UOp, val:UOp) -> UOp:
-  for b, v in zip((buf.src if buf.op is Ops.MSTACK else (buf,)), (val.src if val.op is Ops.STACK else (val,))):
-    struct.pack_into(f'<{v.dtype.fmt}', b.buffer.ensure_allocated()._buf.cpu_view().mv.cast('B'), off.arg * b.dtype.base.itemsize, v.arg)
-  return UOp(Ops.NOOP)
-
-def resolve_getaddr(buf:UOp, g:UOp) -> UOp:
-  if isinstance(b:=buf.buffer, Buffer): return UOp.const(dtypes.uint64, b.get_buf(g.src[1].arg).va_addr)
-  return UOp(Ops.STACK, dtypes.uint64.vec(len(b.bufs)), tuple(UOp.const(dtypes.uint64, x.ensure_allocated()._buf.va_addr) for x in b.bufs))
+# lower cmdbuf submits
+pm_lower_cmdbufs = PatternMatcher([
+  (UPat(Ops.AFTER, src=(UPat(Ops.BINARY),), name="blob", allow_any_len=True), resolve_cmdbuf),
+  (UPat(Ops.BINARY, name="blob"), resolve_cmdbuf),
+])

+# transform patches attached to buffers and params
 pm_resolve_patches = PatternMatcher([
-  # multi
-  (UPat(GroupOp.ALU, src=[UPat(Ops.STACK, name="s"), UPat(Ops.CONST)], name="op"), push_stack),
-  (UPat(Ops.CAST, src=(UPat(Ops.STACK, name="s"),), name="op"), push_stack),
-
-  # getaddr
-  (UPat(Ops.GETADDR, src=(UPat(Ops.SLICE, name="bv"), UPat(Ops.DEVICE, name="dev"))), # getaddr(slice(x)) -> offset+getaddr(x)
-    lambda bv, dev: UOp(Ops.GETADDR, dtypes.uint64, src=(bv.src[0], dev)) + UOp.const(dtypes.uint64, bv.src[1].arg * bv.src[0].dtype.itemsize)),
-  (UPat(Ops.GETADDR, src=(UPat({Ops.BUFFER, Ops.MSTACK, Ops.MSELECT}, name="buf"), UPat(Ops.DEVICE)), name="g"), resolve_getaddr),
-
-  # folders
-  (UPat({Ops.BUFFER, Ops.MSTACK}, name="buf").store(UPat(Ops.BINARY, name="blob")), fold_blob_store),
-  (UPat({Ops.BUFFER, Ops.MSTACK}, name="buf").index(UPat.cvar("off")).or_casted().store(UPat.any(UPat.cvar("val"), UPat(Ops.STACK, name="val"))),
-    fold_const_store),
-]) + symbolic_simple
-
-# *****************
-# 6. callify hcq programs
-
-pm_fixup = PatternMatcher([ # TODO: this should gone?
-  (UPat(Ops.CONST, name="c"), lambda c: c.replace(src=()) if len(c.src) else None),
+  (UPat(Ops.AFTER, src=(UPat((Ops.BUFFER, Ops.PARAM)),), name="buf", allow_any_len=True), resolve_patches)
 ])

-def to_param(bufs:list[UOp], ref:UOp) -> UOp:
-  bufs.append(ref)
-  return UOp.placeholder((ref.buffer.size,), ref.dtype, len(bufs)-1)
-pm_to_param = PatternMatcher([(UPat({Ops.MSELECT, Ops.MSTACK, Ops.BUFFER}, name="r"), lambda ctx, r: to_param(ctx, r))])
+# replace referenced buffers with noops
+pm_resolve_ref_buffers = PatternMatcher([(UPat((Ops.BUFFER, Ops.BUFFER_VIEW), name="buf"), resolve_ref_buffers)])

-def parametrize_host_buffers(call:UOp) -> UOp:
-  body = graph_rewrite(call.src[0], pm_to_param, ctx=(bufs:=[]), bottom_up=True, name="parametrize host buffers")
-  return call.replace(src=(body, *bufs) + call.src[1:], tag="hcq_param")
-pm_parametrize_host_buffers = PatternMatcher([(UPat(Ops.CALL, tag="hcq", name="call"), parametrize_host_buffers)])
+pm_callify = PatternMatcher([(UPat(Ops.SINK, name="sink"), hcq_callify)])

-def callify_hcq(call:UOp) -> UOp:
-  sink = UOp.sink(call.src[0], arg=KernelInfo(name="hcq_submit", estimates=Estimates()), tag=1)
-  return to_program(sink, Device["CPU"].renderer).call(*call.src[1:])
-pm_callify_hcq = PatternMatcher([(UPat(Ops.CALL, tag="hcq_param", name="call"), callify_hcq)])
+# **************** schedule ****************

-@track_rewrites(lambda _,ret: f"HCQ Schedule {pluralize('Kernel', len(ret.src))}")
-def hcq_schedule(linear:UOp) -> UOp:
-  linear = graph_rewrite(linear, pm_insert_copy_staging + pm_flatten_linear, name="insert copy staging")
-  linear = graph_rewrite(linear, pm_prep_runtime, name="prepare runtime")
+def prep_runtime(ctx:HCQ2LowerCtx, linear:UOp) -> tuple[UOp, dict[str,int]]:
+  linear = graph_rewrite(linear, pm_prep_runtime, ctx=ctx, name="hcq: prepare runtime")
+  graph_rewrite(linear, pm_calc_kernargs_sizes, ctx=(sizes:={}), enter_calls=True)
+  return linear, sizes

-  linear = graph_rewrite(linear, pm_lower_ops, name="lower ops into hcq ir")
-  # linear = graph_rewrite(linear, pm_split_into_queues, name="split into queues")
-  linear = graph_rewrite(linear, pm_add_barriers, walk=True, name="add barriers")
-  linear = graph_rewrite(linear, pm_add_signals, walk=True, name="add signals")
-  linear = graph_rewrite(linear, pm_encode_cmdbufs, walk=True, name="encode cmdbufs")
-  linear = graph_rewrite(linear, pm_add_timeline_inc, walk=True, name="add timeline inc")
-  linear = graph_rewrite(linear, pm_lift_patches_to_cmdbuf, name="lift patches to cmdbuf", enter_calls=True)
+def build_host_program(ctx:HCQ2LowerCtx, linear:UOp, ast:UOp, dev:HCQ2Compiled) -> UOp:
+  sink = graph_rewrite(linear, pm_create_host_sink, ctx=ctx, name="hcq: create host sink", walk=True)
+  sink = graph_rewrite(sink, pm_lower_cmdbufs, ctx=ctx, bottom_up=True, name="hcq: lower cmdbufs")
+  sink = graph_rewrite(sink, pm_resolve_patches, ctx=ctx, bottom_up=True, name="hcq: resolve patches")
+  sink = graph_rewrite(sink, pm_resolve_ref_buffers, ctx=ctx, bottom_up=True, name="hcq: resolve ref buffers")
+  sink = graph_rewrite(sink, dev.pm_lower, ctx=ctx, name=f"hcq: device lower {dev.device}", walk=True)
+  return graph_rewrite(sink, pm_callify, ctx=ctx, name="hcq: callify")

-  # realize starts from here
-  linear = graph_rewrite(linear, pm_bufferize, bottom_up=True, name="bufferize placeholders", enter_calls=True)
-  linear = graph_rewrite(linear, pm_hold_call_buffers, walk=True, name="hold call buffers")
-  linear = graph_rewrite(linear, pm_resolve_patches, bottom_up=False, name="simplify patches", enter_calls=True)
-  linear = graph_rewrite(linear, pm_fixup, bottom_up=False, name="fixup", enter_calls=True)
-  linear = graph_rewrite(linear, pm_parametrize_host_buffers, name="parametrize host buffers")
-  linear = graph_rewrite(linear, pm_callify_hcq, name="callify hcq")
+@track_rewrites(name=lambda ctx,linear,ast,dev,**kw: f"hcq schedule {getattr(ast.arg, 'name', ast.op.name.lower())}")
+def hcq_schedule(ctx:HCQ2LowerCtx, linear:UOp, ast:UOp, dev:HCQ2Compiled) -> UOp:
+  linear, sizes = prep_runtime(ctx, linear)
+  for dev_name, sz in sizes.items():
+    off = dev.kernargs_offset_allocator.alloc(sz, 16)
+    ctx.devs[dev_name] = HCQ2DeviceCtx(dev_name, UOp.from_buffer(dev.kernargs_buf.view(sz, dtypes.uint8, off), dev_name),
+                                       UOp.const(dtypes.uint64, dev.kernargs_buf.get_buf(dev_name).va_addr + off))
+  linear = graph_rewrite(linear, pm_lower_kernargs + pm_lower_ops, ctx=ctx, name="hcq: lower ops")
+  linear = graph_rewrite(linear, pm_split_into_queues, ctx=ctx, name="hcq: split into queues")
+  linear = graph_rewrite(linear, pm_add_barriers, ctx=ctx, name="hcq: add barriers", walk=True)
+  linear = graph_rewrite(linear, pm_add_signals, ctx=ctx, name="hcq: add signals", walk=True)
+  linear = graph_rewrite(linear, dev.pm_lower, ctx=ctx, name=f"hcq: encode cmdbuf {dev.device}", walk=True)
+  return build_host_program(ctx, linear, ast, dev)

-  return linear
+def ensure_accessible(ctx:HCQ2LowerCtx, call:UOp, copy:UOp) -> UOp|None:
+  src_buf = call.src[2].buffer # TODO: cleanup
+  dev = call.src[1].buffer.device
+  try: src_buf.get_buf(dev)
+  except Exception:
+    (cpubuf := Buffer("CPU", src_buf.nbytes, dtypes.uint8, preallocate=True)).copyin(src_buf.ensure_allocated().as_memoryview())
+    ctx.holds.append(buf_uop:=UOp.from_buffer(cpubuf, dev))
+    return call.replace(src=call.src[:2] + (buf_uop,) + call.src[3:])
+pm_ensure_bufs_accessible = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY, name="copy"),), name="call", allow_any_len=True), ensure_accessible)])
+
+def hcq_exec(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
+  from tinygrad.engine.realize import run_linear
+
+  if ast.src[1].arg.split(":")[0] != "AMD": return None
+
+  # TODO: this mess should gone
+  resolved_call = call.replace(src=(ast,) + tuple(resolve_params(call, ctx.input_uops)) + tuple(s for s in call.src[1:] if s.op is Ops.BIND))
+  bufs = [cast(Buffer, resolved_call.src[1+gi].buffer) for gi in ast.arg.globals] if ast.op is Ops.PROGRAM \
+    else [cast(Buffer, resolved_call.src[i].buffer) for i in range(1, len(resolved_call.src))]
+  dev = cast(HCQ2Compiled, Device[bufs[0].device])
+  hcq_ctx = HCQ2LowerCtx(name="submit")
+  linear = graph_rewrite(UOp(Ops.LINEAR, dtypes.void, (resolved_call,)), pm_ensure_bufs_accessible, ctx=hcq_ctx)
+  host_call = hcq_schedule(hcq_ctx, linear, ast, dev)
+  with track_stats(ctx, call, dev.device, bufs, ctx.var_vals) as tm:
+    st = time.perf_counter() if ctx.wait else 0.0
+    run_linear(UOp(Ops.LINEAR, dtypes.void, (host_call,)), var_vals=ctx.var_vals, jit=True, update_stats=DEBUG>=3)
+    if ctx.wait:
+      dev.synchronize()
+      tm[0] = time.perf_counter() - st
+  return tm[0] if tm[0] is not None else 0.0
+
+pm_hcq_exec = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat({Ops.PROGRAM, Ops.COPY}, name="ast"),), name="call", allow_any_len=True), hcq_exec),
+])
@@ -23,14 +23,14 @@ from tinygrad.runtime.ops_amd import SQTT, SQTT_ITRACE_SE_MASK, SQTT_LIMIT_SE, S
 from tinygrad.runtime.ops_amd import EVENT_INDEX_PARTIAL_FLUSH, WAIT_REG_MEM_FUNCTION_EQ, WAIT_REG_MEM_FUNCTION_NEQ, WAIT_REG_MEM_FUNCTION_GEQ
 if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import

+from extra.hcq2.hcq2 import HCQ2LowerCtx
 from tinygrad.engine.realize import get_runtime
 from tinygrad.uop.ops import Ops, UPat, PatternMatcher, graph_rewrite

 class AMDComputeQueue(HCQEncoder):
-  def __init__(self, dev:AMDDevice, devs:tuple[str, ...]|None=None):
-    super().__init__()
-    self.dev, self.devs = dev, devs or (dev.device,)
-    self.pm4, self.gc, self.nbio, self.soc = dev.pm4, dev.gc, dev.nbio, dev.soc
+  def __init__(self, ctx:HCQ2LowerCtx, dev:AMDDevice):
+    super().__init__(ctx, dev)
+    self.pm4, self.gc, self.nbio, self.soc = self.dev.pm4, self.dev.gc, self.dev.nbio, self.dev.soc

  def pkt3(self, cmd, *vals): self.q(self.pm4.PACKET3(cmd, len(vals) - 1), *vals)

@@ -105,12 +105,10 @@ class AMDComputeQueue(HCQEncoder):

    self.acquire_mem(gli=0, gl2=0)

-    scratch_addr = self.get_dev_addr(UOp.new_buffer(self.devs, data.private_segment_size, dtypes.uint8).rtag("scratch"))
    args_addr = self.get_dev_addr(args)
-
    user_regs = []
    if data.enable_private_segment_sgpr:
-      scratch_hilo = data64_le(scratch_addr)
+      scratch_hilo = data64_le(self.dev.scratch.va_addr)
      user_regs = [scratch_hilo[0], scratch_hilo[1] | 1 << 31, 0xffffffff, 0x20c14000]
    if data.enable_dispatch_ptr: user_regs += [*data64_le(args_addr + data.kernargs_segment_size)]
    user_regs += [*data64_le(args_addr)]
@@ -118,10 +116,10 @@ class AMDComputeQueue(HCQEncoder):
    self.wreg(self.gc.regCOMPUTE_PGM_LO, *data64_le(prog_addr >> 8))
    self.wreg(self.gc.regCOMPUTE_PGM_RSRC1, data.rsrc1, data.rsrc2)
    self.wreg(self.gc.regCOMPUTE_PGM_RSRC3, data.rsrc3)
-    self.wreg(self.gc.regCOMPUTE_TMPRING_SIZE, self.dev.tmpring_size(data.private_segment_size))
+    self.wreg(self.gc.regCOMPUTE_TMPRING_SIZE, self.dev.tmpring_size)

    for xcc_id in range(self.dev.xccs):
-      scratch_base = scratch_addr + (data.private_segment_size // self.dev.xccs * xcc_id)
+      scratch_base = self.dev.scratch.va_addr + (self.dev.scratch.size // self.dev.xccs * xcc_id)
      self.wreg(self.gc.regCOMPUTE_DISPATCH_SCRATCH_BASE_LO, *data64_le(scratch_base >> 8))

    self.wreg(self.gc.regCOMPUTE_RESTART_X, 0, 0, 0)
@@ -135,47 +133,42 @@ class AMDComputeQueue(HCQEncoder):
    self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.CS_PARTIAL_FLUSH) | self.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH))

 amd_inner_pm = PatternMatcher([
-  (UPat(Ops.LINEAR, src=(UPat(Ops.WAIT, name="x"),)),    lambda ctx, x: ctx.wait(x)),
-  (UPat(Ops.LINEAR, src=(UPat(Ops.BARRIER, name="x"),)), lambda ctx, x: ctx.barrier(x)),
-  (UPat(Ops.LINEAR, src=(UPat(Ops.PROGRAM, name="x"),)), lambda ctx, x: ctx.program(x)),
-  (UPat(Ops.LINEAR, src=(UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", name="x"),)), lambda ctx, x: ctx.timestamp(x)),
-  (UPat(Ops.LINEAR, src=(UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM)), UPat()), name="x"),)), lambda ctx, x: ctx.store(x)),
+  (UPat(Ops.WAIT, name="x"),    lambda ctx, x: ctx.wait(x)),
+  (UPat(Ops.BARRIER, name="x"), lambda ctx, x: ctx.barrier(x)),
+  (UPat(Ops.PROGRAM, name="x"), lambda ctx, x: ctx.program(x)),
+  (UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", name="x"), lambda ctx, x: ctx.timestamp(x)),
+  (UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM)), UPat()), name="x"), lambda ctx, x: ctx.store(x)),
 ])

-def amd_lower_pm4(linear, devs):
-  enc = AMDComputeQueue(Device[devs[0]], devs)
-  graph_rewrite(linear.replace(src=tuple(UOp(Ops.LINEAR, dtypes.void, (cmd,)) for cmd in linear.src)), amd_inner_pm, ctx=enc, name="amd: encode")
-  return enc.uop(dev=devs if len(devs) > 1 else devs[0], tag="compute")
+def amd_lower_pm4(ctx, linear):
+  prg = next(s for s in linear.src if s.op is Ops.PROGRAM)
+  dev = Device[prg.src[1].arg]
+  enc = AMDComputeQueue(ctx, dev)
+  graph_rewrite(linear, amd_inner_pm, ctx=enc, name="amd: encode")
+  return UOp(Ops.BINARY, dtypes.void, arg=enc.blob).rtag((dev.device, "COMPUTE")).after(*enc.src)

-def amd_submit_pm4(cmdbuf, devs):
-  size, zero = UOp.const(dtypes.uint32, cmdbuf.src[0].arg // dtypes.uint32.itemsize), UOp.const(dtypes.int, 0)
+def amd_submit_pm4(ctx, cf):
+  dev = Device[cf.tag]
+  bb_param = cf.src[0]
+  q = dev.compute_queue
+  ring, wptr, doorbell, put_ptr = (ctx.host_param(b) for b in (q.ring, q.write_ptr, q.doorbell, q.put_value))
+  size, ring_dwords = UOp.const(dtypes.uint32, bb_param.dtype.size), q.ring.size

-  # the compute queue's ring and its host-side ring/write/put pointers (placeholders, resolved in pm_bufferize)
-  q = Device['AMD'].compute_queue
-  ring, wptr, doorbell, put_ptr = (UOp.new_buffer(devs, b.size, b.dtype).rtag(("COMPUTE:0", name))
-    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
-
-  # place the cmdbuf at the ring's write offset, wrapping the ring
-  put = put_ptr.index(zero)
+  put = put_ptr[0]
+  i = UOp.range(size, 0, dtype=dtypes.int)
  next_put = put + size.cast(put.dtype)
-  i = UOp.range(size, 0, dtype=dtypes.int, src=(cmdbuf,))
-  ring_idx = ((put + i.cast(put.dtype)) % q.ring.size).cast(dtypes.int)
+  ring_idx = ((put + i.cast(put.dtype)) % ring_dwords).cast(dtypes.int)

-  # copy the cmdbuf into the ring and advance the put/write pointers
-  copy_to_ring = ring.index(ring_idx, dtype=ring.dtype.ptr()).store(
-    cmdbuf.index(i*4, dtype=cmdbuf.dtype.ptr()).cast(dtypes.uint32.ptr()).load()).end(i)
-  bump_put_ptr = put_ptr.index(zero, dtype=put_ptr.dtype.ptr()).store(next_put)
-  bump_wptr = wptr.index(zero, dtype=wptr.dtype.ptr()).store(next_put)
-
-  # ring the doorbell once the copy and pointer bumps have landed
+  copy_to_ring = ring[ring_idx].store(bb_param[i]).end(i)
+  bump_put_ptr = put_ptr[0].store(next_put)
+  bump_wptr = wptr[0].store(next_put)
  flush = UOp.barrier(copy_to_ring, bump_put_ptr, bump_wptr)
-  return doorbell.after(flush).index(zero, dtype=doorbell.dtype.ptr()).store(next_put)
+  return doorbell.after(flush)[0].store(next_put)

 class AMDCopyQueue(HCQEncoder):
-  def __init__(self, dev:AMDDevice, queue_idx=0):
-    super().__init__()
-    self.dev = dev
-    self.sdma, self.queue_idx, self.max_copy_size = dev.sdma, queue_idx, dev.max_copy_size
+  def __init__(self, ctx:HCQ2LowerCtx, dev:AMDDevice, queue_idx=0):
+    super().__init__(ctx, dev)
+    self.sdma, self.queue_idx, self.max_copy_size = self.dev.sdma, queue_idx, self.dev.max_copy_size

  def copy(self, x):
    dest, src, copy_size = self.get_dev_addr(x.src[0]), self.get_dev_addr(x.src[1]), x.arg
@@ -200,93 +193,89 @@ class AMDCopyQueue(HCQEncoder):
    self.q(self.sdma.SDMA_OP_TIMESTAMP | self.sdma.SDMA_PKT_TIMESTAMP_GET_HEADER_SUB_OP(self.sdma.SDMA_SUBOP_TIMESTAMP_GET_GLOBAL),
           *data64_le(self.get_dev_addr(x.src[0])))

-def amd_lower_sdma(linear, devs):
-  enc = AMDCopyQueue(Device[devs[0]])
-  graph_rewrite(linear.replace(src=tuple(UOp(Ops.LINEAR, dtypes.void, (cmd,)) for cmd in linear.src)), amd_inner_sdma_pm, ctx=enc, name="amd: encode sdma")
-  return enc.uop(dev=devs if len(devs) > 1 else devs[0], tag="copy")
+def amd_lower_sdma(ctx, linear):
+  copy = next(s for s in linear.src if s.op is Ops.COPY)
+  dev = Device[copy.src[0].buffer.device]
+  enc = AMDCopyQueue(ctx, dev)
+  graph_rewrite(linear, amd_inner_sdma_pm, ctx=enc, name="amd: encode sdma")
+  return UOp(Ops.BINARY, dtypes.void, arg=enc.blob).rtag((dev.device, "COPY")).after(*enc.src)

 amd_inner_sdma_pm = PatternMatcher([
-  (UPat(Ops.LINEAR, src=(UPat(Ops.WAIT, name="x"),)), lambda ctx, x: ctx.wait(x)),
-  (UPat(Ops.LINEAR, src=(UPat(Ops.BARRIER, name="x"),)), lambda ctx, x: None),
-  (UPat(Ops.LINEAR, src=(UPat(Ops.COPY, name="x"),)), lambda ctx, x: ctx.copy(x)),
-  (UPat(Ops.LINEAR, src=(UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", name="x"),)), lambda ctx, x: ctx.timestamp(x)),
-  (UPat(Ops.LINEAR, src=(UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM)), UPat()), name="x"),)), lambda ctx, x: ctx.store(x)),
+  (UPat(Ops.WAIT,  name="x"), lambda ctx, x: ctx.wait(x)),
+  (UPat(Ops.BARRIER, name="x"), lambda ctx, x: None),
+  (UPat(Ops.COPY,  name="x"), lambda ctx, x: ctx.copy(x)),
+  (UPat(Ops.CUSTOM_FUNCTION, arg="timestamp", name="x"), lambda ctx, x: ctx.timestamp(x)),
+  (UPat(Ops.STORE, src=(UPat((Ops.BUFFER, Ops.PARAM)), UPat()), name="x"), lambda ctx, x: ctx.store(x)),
 ])

-def amd_submit_sdma(cmdbuf, devs):
-  # the cmdbuf to submit + the patch writes that fill it
-  size_dw, zero = cmdbuf.src[0].arg // dtypes.uint32.itemsize, UOp.const(dtypes.int, 0)
+def amd_submit_sdma(ctx, cf):
+  dev = Device[cf.tag]
+  bb_param = cf.src[0]
+  q = dev.sdma_queue(0)
+  ring, wptr, doorbell, put_ptr = (ctx.host_param(b) for b in (q.ring, q.write_ptr, q.doorbell, q.put_value))
+  size_dw, ring_bytes = bb_param.dtype.size, q.ring.size * 4

-  # the sdma queue's ring and its host-side ring/write/put pointers
-  q = Device['AMD'].sdma_queue(0)
-  ring, wptr, doorbell, put_ptr = (UOp.new_buffer(devs, b.size, b.dtype).rtag(("SDMA:0", name))
-    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
-
-  # sdma needs the cmdbuf contiguous: if it won't fit before the ring end, restart at 0 and zero the tail
-  put_b = put_ptr.index(zero)
-  tail_off_dw = ((put_b % (q.ring.size * 4)) // 4).cast(dtypes.int)
+  put_b = put_ptr[0]
+  tail_off_dw = ((put_b % ring_bytes) // 4).cast(dtypes.int)
  fits = (size_dw <= q.ring.size - tail_off_dw).cast(dtypes.int)
  start_dw = fits * tail_off_dw
  zero_amt_dw = (1 - fits) * (q.ring.size - tail_off_dw)

-  # zero the wrapped tail, then copy the cmdbuf into the ring
-  zi = UOp.range(zero_amt_dw, 0, dtype=dtypes.int, src=(cmdbuf,))
-  zero_tail = ring.index(tail_off_dw + zi, dtype=ring.dtype.ptr()).store(UOp.const(dtypes.uint32, 0)).end(zi)
-  i = UOp.range(UOp.const(dtypes.int, size_dw), 0, dtype=dtypes.int, src=(cmdbuf,))
-  copy_to_ring = ring.index(start_dw + i, dtype=ring.dtype.ptr()).store(
-    cmdbuf.index(i*4, dtype=cmdbuf.dtype.ptr()).cast(dtypes.uint32.ptr()).load()).end(i)
+  zi = UOp.range(zero_amt_dw, 0, dtype=dtypes.int)
+  zero_tail = ring[tail_off_dw + zi].store(UOp.const(dtypes.uint32, 0)).end(zi)
+
+  i = UOp.range(UOp.const(dtypes.int, size_dw), 0, dtype=dtypes.int)
+  copy_to_ring = ring[start_dw + i].store(bb_param[i]).end(i)

-  # advance the put/write pointers past the zeroed tail and the cmdbuf
  next_put_b = put_b + ((zero_amt_dw + size_dw) * 4).cast(put_b.dtype)
-  bump_put_ptr = put_ptr.index(zero, dtype=put_ptr.dtype.ptr()).store(next_put_b)
-  bump_wptr = wptr.index(zero, dtype=wptr.dtype.ptr()).store(next_put_b)
-
-  # ring the doorbell once the writes have landed
+  bump_put_ptr = put_ptr[0].store(next_put_b)
+  bump_wptr = wptr[0].store(next_put_b)
  flush = UOp.barrier(zero_tail, copy_to_ring, bump_put_ptr, bump_wptr)
-  return doorbell.after(flush).index(zero, dtype=doorbell.dtype.ptr()).store(next_put_b)
+  return doorbell.after(flush)[0].store(next_put_b)

@dataclass(frozen=True)
 class AMDProgramData:
  entry_point_offset:int; rsrc1:int; rsrc2:int; rsrc3:int; wave32:bool
-  private_segment_size:int; kernargs_segment_size:int; kernargs_alloc_size:int
+  kernargs_segment_size:int; kernargs_alloc_size:int
  enable_dispatch_ptr:int; enable_private_segment_sgpr:int

-_amd_program_cache:dict[tuple[bytes,str], tuple[AMDProgramData,bytes]] = {}
+_amd_program_cache:dict[tuple[bytes,str], tuple[AMDProgramData,Buffer]] = {}

-def amd_build_program(prg:UOp) -> UOp:
-  dev = Device[prg.src[1].arg] # TODO: rm this
+def amd_build_program(ctx:HCQ2LowerCtx, prg:UOp) -> UOp:
+  dev = Device[prg.src[1].arg]
  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[4].arg, dev.device))) is None:
    image, sections, relocs = elf_loader(lib)
    rodata = next(sh.header.sh_addr for sh in sections if sh.name == ".rodata")
    for off, sym, typ, addent in relocs:
      assert typ == 5, f"unknown AMD reloc {typ}"  # R_AMDGPU_REL64
      image[off:off+8] = struct.pack('<q', sym - off + addent)
+    lib_gpu = Buffer(dev.device, round_up(image.nbytes, 0x1000), dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)
+    dev.allocator._copyin(lib_gpu._buf, image)
+    dev.synchronize()
    desc = amdgpu_kd.llvm_amdhsa_kernel_descriptor_t.from_buffer_copy(bytes(image[rodata:rodata+ctypes.sizeof(amdgpu_kd.llvm_amdhsa_kernel_descriptor_t)]))
    if (lds:=((desc.group_segment_fixed_size+511)//512)&0x1FF) > (dev.iface.props['lds_size_in_kb']*1024)//512:
      raise RuntimeError("Too many resources requested: group_segment_size")
+    dev._ensure_has_local_memory(desc.private_segment_fixed_size)
    edp = desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_DISPATCH_PTR
    cached = _amd_program_cache[key] = (AMDProgramData(
      entry_point_offset=rodata + desc.kernel_code_entry_byte_offset,
      rsrc1=desc.compute_pgm_rsrc1 | ((1<<20) if dev.target[0]==11 else 0),  # priv=1 on gfx11 for cwsr
      rsrc2=desc.compute_pgm_rsrc2 | (lds<<15), rsrc3=desc.compute_pgm_rsrc3,
      wave32=bool(desc.kernel_code_properties & 0x400),
-      private_segment_size=desc.private_segment_fixed_size,
      kernargs_segment_size=desc.kernarg_size,
      kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0),
      enable_dispatch_ptr=edp,
-      enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER), bytes(image))
-  return cached
-
-pm_prep_program = PatternMatcher([
-  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.DEVICE, arg="AMD"), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"), amd_build_program),
-])
+      enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER,
+    ), lib_gpu)
+  data, lib_gpu = cached
+  return prg.replace(src=(UOp.from_buffer(lib_gpu, dev.device),), arg=(data, prg.arg))

 class AMDAllocator(HCQAllocator['AMDDevice']):
  def __init__(self, dev:AMDDevice):
    super().__init__(dev, supports_copy_from_disk=dev.has_sdma_queue, supports_transfer=dev.has_sdma_queue and not dev.is_usb())

  def _alloc(self, size:int, options:BufferSpec) -> HCQ2Buffer:
-    return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_sdma_queue)
+    return self.dev.iface.alloc(size, host=True, uncached=options.uncached, cpu_access=True)

  def _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque)

@@ -373,18 +362,17 @@ class PCIIface(PCIIfaceBase):

 def _mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface,), {})

-def encode_queue(q:UOp) -> UOp|None:
-  if not (isinstance(q.arg, tuple) and len(q.arg) == 2 and q.arg[1] in ("COMPUTE", "COPY")): return None
-  devs = (q.arg[0],) if isinstance(q.arg[0], str) else q.arg[0]  # TODO: make this prettier
-  return amd_submit_pm4(amd_lower_pm4(q, devs), devs) if q.arg[1] == "COMPUTE" else amd_submit_sdma(amd_lower_sdma(q, devs), devs)
-
-pm_lower = PatternMatcher([
-  (UPat(Ops.LINEAR, name="q"), encode_queue),
-])
-
 class AMDDevice(HCQ2Compiled):
  timestamp_divider = 100.0  # AMD GPU clock: ticks/us

+  pm_lower = PatternMatcher([
+    (UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"), amd_build_program),
+    (UPat(Ops.LINEAR, arg="COMPUTE", name="linear"), amd_lower_pm4),
+    (UPat(Ops.LINEAR, arg="COPY", name="linear"), amd_lower_sdma),
+    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_compute", name="cf"), amd_submit_pm4),
+    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_copy", name="cf"), amd_submit_sdma),
+  ])
+
  ifaces = [PCIIface]

  def is_am(self) -> bool: return isinstance(self.iface, (PCIIface,))
@@ -423,13 +411,14 @@ class AMDDevice(HCQ2Compiled):

    self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
    self.sdma_queues:dict = {}
-    self.has_sdma_queue = True # self.sdma_queue(0) is not None, TODO: think of this
+    self.has_sdma_queue = self.sdma_queue(0) is not None

-    super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, can_recover=self.is_am(), arch=self.arch)
+    super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None,
+                     kernargs_size=16 << 20, can_recover=self.is_am(), arch=self.arch)

    # Scratch setup
    self.max_private_segment_size = 0
-    self.pm_bufferize = PatternMatcher([(UPat(Ops.BUFFER, tag="scratch", name="b"), lambda ctx, b: ctx.scratch_buffer(b.arg))]) + self.pm_bufferize
+    self._ensure_has_local_memory(128) # set default scratch size to 128 bytes per thread

    self.pmc_enabled:bool = PROFILE > 0 and PMC > 0
    if self.pmc_enabled:
@@ -456,6 +445,19 @@ class AMDDevice(HCQ2Compiled):
      self.sqtt_wptrs = self.allocator.alloc(round_up(self.se_cnt * self.xccs * 4, 0x1000), BufferSpec(cpu_access=True, nolru=True))
      self.sqtt_next_cmd_id = itertools.count(0)

+  @functools.cached_property
+  def compute_queue(self) -> AMDQueueDesc:
+    # https://gitlab.freedesktop.org/agd5f/linux/-/blob/a1fc9f584c4aaf8bc1ebfa459fc57a3f26a290d8/drivers/gpu/drm/amd/amdkfd/kfd_queue.c#L391
+    sgrp_size_per_cu, hwreg_size_per_cu = 0x4000, 0x1000
+    lds_size_per_cu = self.iface.props["lds_size_in_kb"] << 10 if self.target[:2] == (9,5) else 0x10000
+    vgpr_size_per_cu = 0x60000 if self.target in {(11,0,0), (11,0,1), (11,5,1), (12,0,0), (12,0,1)} else 0x80000 if self.target[0] == 9 else 0x40000
+    wg_data_size = round_up((vgpr_size_per_cu + sgrp_size_per_cu + lds_size_per_cu + hwreg_size_per_cu) * self.cu_cnt, mmap.PAGESIZE)
+    ctl_stack_size = round_up((12 if self.target[0] != 9 else 8) * self.wave_cnt + 8 + 40, mmap.PAGESIZE)
+    return self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL if self.is_aql else kfd.KFD_IOC_QUEUE_TYPE_COMPUTE,
+      0x2000 if self.is_usb() else (16 << 20), eop_buffer_size=0x1000,
+      ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size,
+      debug_memory_size=round_up(self.wave_cnt * 32, 64))
+
  def create_queue(self, queue_type, ring_size, ctx_save_restore_size=0, eop_buffer_size=0, ctl_stack_size=0, debug_memory_size=0, idx=0):
    ring = self.iface.alloc(ring_size, uncached=True, cpu_access=True)
    gart = self.iface.alloc(0x100, uncached=True, cpu_access=True)
@@ -471,29 +473,9 @@ class AMDDevice(HCQ2Compiled):
    cwsr_buffer = self.iface.alloc(cwsr_buffer_size) if ctx_save_restore_size else None
    eop_buffer = self.iface.alloc(eop_buffer_size) if eop_buffer_size else None

-    queue = (self.iface.create_queue(queue_type, ring, gart, rptr=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset,
-             wptr=getattr(hsa.amd_queue_t, 'write_dispatch_id').offset, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer,
-             ctx_save_restore_size=ctx_save_restore_size, ctl_stack_size=ctl_stack_size, idx=idx))
-
-    qname = f"{'SDMA' if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA else 'COMPUTE'}:{idx}"
-    self.pm_bufferize = PatternMatcher([
-      (UPat(Ops.BUFFER, tag={(qname, name)}), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"]
-    ]) + self.pm_bufferize
-
-    return queue
-
-  @functools.cached_property
-  def compute_queue(self) -> AMDQueueDesc:
-    # https://gitlab.freedesktop.org/agd5f/linux/-/blob/a1fc9f584c4aaf8bc1ebfa459fc57a3f26a290d8/drivers/gpu/drm/amd/amdkfd/kfd_queue.c#L391
-    sgrp_size_per_cu, hwreg_size_per_cu = 0x4000, 0x1000
-    lds_size_per_cu = self.iface.props["lds_size_in_kb"] << 10 if self.target[:2] == (9,5) else 0x10000
-    vgpr_size_per_cu = 0x60000 if self.target in {(11,0,0), (11,0,1), (11,5,1), (12,0,0), (12,0,1)} else 0x80000 if self.target[0] == 9 else 0x40000
-    wg_data_size = round_up((vgpr_size_per_cu + sgrp_size_per_cu + lds_size_per_cu + hwreg_size_per_cu) * self.cu_cnt, mmap.PAGESIZE)
-    ctl_stack_size = round_up((12 if self.target[0] != 9 else 8) * self.wave_cnt + 8 + 40, mmap.PAGESIZE)
-    return self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL if self.is_aql else kfd.KFD_IOC_QUEUE_TYPE_COMPUTE,
-      0x2000 if self.is_usb() else (16 << 20), eop_buffer_size=0x1000,
-      ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size,
-      debug_memory_size=round_up(self.wave_cnt * 32, 64))
+    return (self.iface.create_queue(queue_type, ring, gart, rptr=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset,
+            wptr=getattr(hsa.amd_queue_t, 'write_dispatch_id').offset, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer,
+            ctx_save_restore_size=ctx_save_restore_size, ctl_stack_size=ctl_stack_size, idx=idx))

  def sdma_queue(self, idx:int):
    if getenv("AMD_DISABLE_SDMA"): return None
@@ -502,49 +484,38 @@ class AMDDevice(HCQ2Compiled):
      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20), idx=idx)
    return self.sdma_queues.get(idx, None)

-  def tmpring_size(self, private_segment_size):
-    private_segment_size = max(private_segment_size, 128)
+  def _ensure_has_local_memory(self, private_segment_size):
+    if self.max_private_segment_size >= private_segment_size: return

    lanes_per_wave = 64 # wave64
    mem_alignment_size = 256 if self.target[0] != 9 else 1024
    size_per_thread = round_up(private_segment_size, mem_alignment_size // lanes_per_wave)
    size_per_xcc = size_per_thread * lanes_per_wave * self.iface.props['max_slots_scratch_cu'] * self.cu_cnt
+    self.scratch, ok = self._realloc(getattr(self, 'scratch', None), size_per_xcc * self.xccs)
+    if ok:
+      # NOTE: xcc logic is correct only for GFX9.
+      max_scratch_waves = self.cu_cnt * self.iface.props['max_slots_scratch_cu'] * self.xccs
+      wave_scratch = ceildiv(lanes_per_wave * size_per_thread, mem_alignment_size)
+      num_waves = (size_per_xcc // (wave_scratch * mem_alignment_size)) // (self.se_cnt if self.target[0] != 9 else 1)

-    # NOTE: xcc logic is correct only for GFX9.
-    max_scratch_waves = self.cu_cnt * self.iface.props['max_slots_scratch_cu'] * self.xccs
-    wave_scratch = ceildiv(lanes_per_wave * size_per_thread, mem_alignment_size)
-    num_waves = (size_per_xcc // (wave_scratch * mem_alignment_size)) // (self.se_cnt if self.target[0] != 9 else 1)
-
-    tmpring_t = getattr(hsa, f'union_COMPUTE_TMPRING_SIZE{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
-    tmpring = int.from_bytes(tmpring_t(WAVES=min(num_waves, max_scratch_waves), WAVESIZE=wave_scratch), 'little')
-
-    if hasattr(self, 'aql_desc'):
-      gfx9_rsrc = {'NUM_FORMAT':hsa.BUF_NUM_FORMAT_UINT, 'DATA_FORMAT':hsa.BUF_DATA_FORMAT_32, 'ELEMENT_SIZE':1, 'INDEX_STRIDE':3}
-      rsrc = {'DST_SEL_X':hsa.SQ_SEL_X, 'DST_SEL_Y':hsa.SQ_SEL_Y, 'DST_SEL_Z':hsa.SQ_SEL_Z, 'DST_SEL_W':hsa.SQ_SEL_W, 'ADD_TID_ENABLE':1,
-              'TYPE':hsa.SQ_RSRC_BUF, **(gfx9_rsrc if self.target[0] == 9 else {'FORMAT':hsa.BUF_FORMAT_32_UINT, 'OOB_SELECT':2})}
-      rsrc1_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD1{"_GFX11" if self.target[0] != 9 else ""}_bitfields')
-      rsrc3_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD3{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
-
-      self.aql_desc.scratch_backing_memory_location = int(self.scratch.get_buf().va_addr)
-      self.aql_desc.scratch_wave64_lane_byte_size = self.max_private_segment_size * lanes_per_wave // 64
-      self.aql_desc.scratch_resource_descriptor[:] = [lo32(self.scratch.get_buf().va_addr),
-        int.from_bytes(rsrc1_t(BASE_ADDRESS_HI=hi32(self.scratch.get_buf().va_addr), SWIZZLE_ENABLE=1), 'little'),
-        lo32(size_per_xcc), int.from_bytes(bytes(rsrc3_t(**rsrc)), 'little')]
-      self.aql_desc.compute_tmpring_size = tmpring
-      self.aql_gart.cpu_view()[:ctypes.sizeof(self.aql_desc)] = bytes(self.aql_desc)
-
-    return tmpring
-
-  def scratch_buffer(self, private_segment_size):
-    private_segment_size = max(private_segment_size, 128)
-    if self.max_private_segment_size < private_segment_size:
-      lanes_per_wave = 64 # wave64
-      mem_alignment_size = 256 if self.target[0] != 9 else 1024
-      size_per_thread = round_up(private_segment_size, mem_alignment_size // lanes_per_wave)
-      size_per_xcc = size_per_thread * lanes_per_wave * self.iface.props['max_slots_scratch_cu'] * self.cu_cnt
-      self.scratch = Buffer(self.device, size_per_xcc * self.xccs, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)
+      tmpring_t = getattr(hsa, f'union_COMPUTE_TMPRING_SIZE{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
+      self.tmpring_size = int.from_bytes(tmpring_t(WAVES=min(num_waves, max_scratch_waves), WAVESIZE=wave_scratch), 'little')
      self.max_private_segment_size = private_segment_size
-    return self.scratch
+
+      if hasattr(self, 'aql_desc'):
+        gfx9_rsrc = {'NUM_FORMAT':hsa.BUF_NUM_FORMAT_UINT, 'DATA_FORMAT':hsa.BUF_DATA_FORMAT_32, 'ELEMENT_SIZE':1, 'INDEX_STRIDE':3}
+        rsrc = {'DST_SEL_X':hsa.SQ_SEL_X, 'DST_SEL_Y':hsa.SQ_SEL_Y, 'DST_SEL_Z':hsa.SQ_SEL_Z, 'DST_SEL_W':hsa.SQ_SEL_W, 'ADD_TID_ENABLE':1,
+                'TYPE':hsa.SQ_RSRC_BUF, **(gfx9_rsrc if self.target[0] == 9 else {'FORMAT':hsa.BUF_FORMAT_32_UINT, 'OOB_SELECT':2})}
+        rsrc1_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD1{"_GFX11" if self.target[0] != 9 else ""}_bitfields')
+        rsrc3_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD3{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
+
+        self.aql_desc.scratch_backing_memory_location = int(self.scratch.va_addr)
+        self.aql_desc.scratch_wave64_lane_byte_size = self.max_private_segment_size * lanes_per_wave // 64
+        self.aql_desc.scratch_resource_descriptor[:] = [lo32(self.scratch.va_addr),
+          int.from_bytes(rsrc1_t(BASE_ADDRESS_HI=hi32(self.scratch.va_addr), SWIZZLE_ENABLE=1), 'little'),
+          lo32(size_per_xcc), int.from_bytes(bytes(rsrc3_t(**rsrc)), 'little')]
+        self.aql_desc.compute_tmpring_size = self.tmpring_size
+        self.aql_gart.cpu_view()[:ctypes.sizeof(self.aql_desc)] = bytes(self.aql_desc)

  def on_device_hang(self): self.iface.on_device_hang()

@@ -5,19 +5,19 @@ from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
 from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, compile_cpp, alloc_like, alloc_local, scalar_amax, dname_of

-# module-level mailbox: grad_xw13 UOp -> (grad_xw13_fp8 UOp, inv_scale UOp)
+# module-level mailbox: grad_xw13 UOp -> (grad_xw13_fp8 UOp, inv_scale UOp, new_amax UOp, store_effect)
 # lets cdna_asm_gemm's bwd reuse the fp8 companion produced by the fused silu_mul bwd kernel
 # instead of doing a redundant bf16 -> fp8 quantize.
-_grad_fp8_mailbox:dict[UOp, tuple[UOp, UOp]] = {}
+_grad_fp8_mailbox:dict = {}

@functools.cache
-def _custom_fused_bwd_w13(grad_xw13_fp8:UOp, grad_amax_buf:UOp,
+def _custom_fused_bwd_w13(grad_xw13:UOp, grad_xw13_fp8:UOp, grad_amax_buf:UOp,
                          xw13:UOp, grad_x2:UOp, amax_state:UOp, grad_amax_state:UOp, dname:str) -> UOp:
  hidden = xw13.shape[2] // 2
  n_elems = xw13.shape[0] * xw13.shape[1] * hidden
  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 * 3 + n_elems * 2 + NUM_WG * 4 + 4
-  sink = UOp.sink(grad_xw13_fp8.base, grad_amax_buf.base,
+  mem = n_elems * 2 * 5 + n_elems * 2 + NUM_WG * 4 + 4
+  sink = UOp.sink(grad_xw13.base, grad_xw13_fp8.base, grad_amax_buf.base,
                  xw13.base, grad_x2.base, amax_state.base, grad_amax_state.base, threads, workgroups,
                  arg=KernelInfo(f"fused_silu_mul_bwd_w13_{n_elems}", estimates=Estimates(ops=10*n_elems, mem=mem)))
  src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_bwd_w13.cpp", n_elems, hidden)
@@ -41,23 +41,23 @@ def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
  _, _, xw13, amax_state, grad_amax_state = kernel.src[1:]
  device = xw13.device
  axis = xw13.axis if isinstance(device, tuple) else None
+  grad_xw13     = alloc_like(xw13.shape, dtypes.bfloat16, device, axis)
  grad_xw13_fp8 = alloc_like(xw13.shape, dtypes.fp8e4m3,  device, axis)
  grad_amax_buf = alloc_local((NUM_WG,), dtypes.float32,  device, axis)
  grad_amax_state_t = Tensor(grad_amax_state, device=device)
  fxn = functools.partial(_custom_fused_bwd_w13, dname=dname_of(device))
-  grad_xw13_fp8, grad_amax_buf, *_ = Tensor.custom_kernel(
-    grad_xw13_fp8, grad_amax_buf,
+  grad_xw13, grad_xw13_fp8, grad_amax_buf, *_ = Tensor.custom_kernel(
+    grad_xw13, grad_xw13_fp8, grad_amax_buf,
    Tensor(xw13, device=device), Tensor(gradient, device=device).cast(dtypes.bfloat16),
    Tensor(amax_state, device=device), grad_amax_state_t, fxn=fxn)
-  grad_xw13_uop = grad_xw13_fp8.uop.cast(dtypes.bfloat16)
  inv_scale = (grad_amax_state_t.float() + 1e-8) / FP8_MAX
  new_grad_amax = scalar_amax(grad_amax_buf)
  store_effect = grad_amax_state_t.uop.store(new_grad_amax.uop)
  assert grad_xw13_fp8.uop.op is Ops.AFTER, f"expected AFTER, got {grad_xw13_fp8.uop.op}"
  grad_xw13_fp8_uop = grad_xw13_fp8.uop.replace(src=grad_xw13_fp8.uop.src + (store_effect,))
  # Stash fp8 companion for cdna_asm_gemm's bwd to attach to grad_a.
-  _grad_fp8_mailbox[grad_xw13_uop] = (grad_xw13_fp8_uop, inv_scale.uop)
-  return (None, None, grad_xw13_uop, None, None)
+  _grad_fp8_mailbox[grad_xw13.uop] = (grad_xw13_fp8_uop, inv_scale.uop)
+  return (None, None, grad_xw13.uop, None, None)

 def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_state:Tensor) -> tuple[Tensor, Tensor, Tensor]:
  # NOTE: silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns (fp8, inv_scale, new_amax)
@@ -21,13 +21,15 @@ constexpr float FP8_MAX = 448.0f;
 static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC");
 static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC");

-// fused silu*mul backward, two outputs in a single HBM pass:
-//   1) fp8  grad_xw13_fp8  — delayed-scale quantize using grad_amax_state (mailbox to matmul bwd)
-//   2) fp32 grad_amax_buf  — per-WG partial |grad_xw13|, reduced into next step's grad_amax_state
+// fused silu*mul backward, three outputs in a single HBM pass:
+//   1) bf16 grad_xw13      — consumed by downstream bf16 autograd chain
+//   2) fp8  grad_xw13_fp8  — delayed-scale quantize using grad_amax_state (mailbox to matmul bwd)
+//   3) fp32 grad_amax_buf  — per-WG partial |grad_xw13|, reduced into next step's grad_amax_state
 // grad_amax_state is read for the fp8 scale. The store of new_grad_amax into grad_amax_state's
 // buffer is built in Python as a separate effect and threaded into grad_a via .after(store).
 extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
 fused_silu_mul_bwd_w13(
+    __hip_bfloat16*       __restrict__ grad_xw13_out,        // bf16, 2*N_ELEMS
    __hip_fp8_storage_t*  __restrict__ grad_xw13_fp8_out,    // fp8,  2*N_ELEMS
    float*                __restrict__ grad_amax_buf,        // fp32, NUM_WG per-WG partials
    const __hip_bfloat16* __restrict__ xw13,                 // bf16, 2*N_ELEMS
@@ -60,6 +62,7 @@ fused_silu_mul_bwd_w13(
    const __hip_bfloat16 *x3 = reinterpret_cast<const __hip_bfloat16*>(&x3_raw);
    const __hip_bfloat16 *gv = reinterpret_cast<const __hip_bfloat16*>(&g_raw);

+    __hip_bfloat16 out1[VEC], out3[VEC];
    __hip_fp8_storage_t fp8_1[VEC], fp8_3[VEC];
    #pragma unroll
    for (int i = 0; i < VEC; i++) {
@@ -72,11 +75,15 @@ fused_silu_mul_bwd_w13(
      const float gs = fg * scale;
      const float g1 = gs * silu_prime * f3;
      const float g3 = gs * silu;
+      out1[i] = static_cast<__hip_bfloat16>(g1);
+      out3[i] = static_cast<__hip_bfloat16>(g3);
      local_max = fmaxf(local_max, fmaxf(fabsf(g1), fabsf(g3)));
      fp8_1[i] = __hip_cvt_float_to_fp8(fmaxf(-FP8_MAX, fminf(FP8_MAX, g1 * g_scale)), __HIP_SATFINITE, __HIP_E4M3);
      fp8_3[i] = __hip_cvt_float_to_fp8(fmaxf(-FP8_MAX, fminf(FP8_MAX, g3 * g_scale)), __HIP_SATFINITE, __HIP_E4M3);
    }

+    *reinterpret_cast<float4*>(&grad_xw13_out[xw1_off]) = *reinterpret_cast<float4*>(out1);
+    *reinterpret_cast<float4*>(&grad_xw13_out[xw3_off]) = *reinterpret_cast<float4*>(out3);
    *reinterpret_cast<uint64_t*>(&grad_xw13_fp8_out[xw1_off]) = *reinterpret_cast<uint64_t*>(fp8_1);
    *reinterpret_cast<uint64_t*>(&grad_xw13_fp8_out[xw3_off]) = *reinterpret_cast<uint64_t*>(fp8_3);
  }
@@ -1,38 +1,41 @@
-import functools
+from __future__ import annotations
+import functools, pathlib
 from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
+
+THREADS_PER_WG = 256

@functools.cache
 def _custom_fused_ce_loss_fwd(loss_out:UOp, max_out:UOp, lse_out:UOp, logits:UOp, targets:UOp,
-                              vocab:int, rows:int, label_smoothing:float) -> UOp:
-  row = UOp.range(rows, 0)
-
-  v_max = UOp.range(vocab, 1, axis_type=AxisType.REDUCE)
-  row_max = logits[row, v_max].cast(dtypes.float).reduce(v_max, arg=Ops.MAX)
-
-  v_lse = UOp.range(vocab, 2, axis_type=AxisType.REDUCE)
-  row_lse = (logits[row, v_lse].cast(dtypes.float) - row_max).exp().reduce(v_lse, arg=Ops.ADD).log() + row_max
-
-  v_smooth = UOp.range(vocab, 3, axis_type=AxisType.REDUCE)
-  target = logits[row, targets[row].cast(dtypes.weakint)].cast(dtypes.float)
-  mean_logits = logits[row, v_smooth].cast(dtypes.float).reduce(v_smooth, arg=Ops.ADD) / vocab
-  loss = row_lse - (1.0 - label_smoothing) * target - label_smoothing * mean_logits
-  stores = UOp.group(loss_out[row].store(loss), max_out[row].store(row_max), lse_out[row].store(row_lse))
-
-  return stores.end(row).sink(arg=KernelInfo(f"fused_ce_loss_fwd_{rows}_{vocab}"))
+                              dname:str, vocab:int, rows:int, label_smoothing:float) -> UOp:
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(rows, "gidx0")
+  mem = rows * vocab * 2 + rows * 12 + rows * 4
+  sink = UOp.sink(loss_out.base, max_out.base, lse_out.base, logits.base, targets.base,
+                  threads, workgroups,
+                  arg=KernelInfo(f"fused_ce_loss_fwd", estimates=Estimates(ops=6*rows*vocab, mem=mem)))
+  src = (pathlib.Path(__file__).parent/"fused_ce_loss.cpp").read_text()
+  defines = [f"-DVOCAB={vocab}", f"-DTHREADS_PER_WG={THREADS_PER_WG}",
+             f"-DLABEL_SMOOTHING={label_smoothing}f"]
+  lib = HIPCCCompiler("gfx950", ["-std=c++20", "-ffast-math", *defines]).compile_cached(src)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))

@functools.cache
 def _custom_fused_ce_loss_bwd(d_logits:UOp, logits:UOp, lse:UOp, targets:UOp, scale:UOp,
-                              vocab:int, rows:int, label_smoothing:float) -> UOp:
-  row = UOp.range(rows, 0)
-  v = UOp.range(vocab, 1)
-
-  prob = (logits[row, v].cast(dtypes.float) - lse[row]).exp()
-  target = v.eq(targets[row].cast(dtypes.weakint)).where(1.0 - label_smoothing, 0.0)
-  smooth = label_smoothing / vocab
-  grad = (prob - target - smooth) * scale[0]
-
-  return d_logits[row, v].store(grad.cast(d_logits.dtype.base)).end(v, row).sink(arg=KernelInfo(f"fused_ce_loss_bwd_{rows}_{vocab}"))
+                              dname:str, vocab:int, rows:int, label_smoothing:float) -> UOp:
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(rows, "gidx0")
+  mem = rows * vocab * 4 + rows * 8 + 4
+  sink = UOp.sink(d_logits.base, logits.base, lse.base, targets.base, scale.base,
+                  threads, workgroups,
+                  arg=KernelInfo(f"fused_ce_loss_bwd", estimates=Estimates(ops=4*rows*vocab, mem=mem)))
+  src = (pathlib.Path(__file__).parent/"fused_ce_loss_bwd.cpp").read_text()
+  defines = [f"-DVOCAB={vocab}", f"-DTHREADS_PER_WG={THREADS_PER_WG}",
+             f"-DLABEL_SMOOTHING={label_smoothing}f"]
+  lib = HIPCCCompiler("gfx950", ["-std=c++20", "-ffast-math", *defines]).compile_cached(src)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))

 def _fused_ce_loss_bwd(gradient:UOp, kernel:UOp, label_smoothing:float):
  # NOTE: forward inputs are (loss_out, max_out, lse_out, logits, targets)
@@ -44,16 +47,18 @@ def _fused_ce_loss_bwd(gradient:UOp, kernel:UOp, label_smoothing:float):
    axis = logits_u.axis
    ndev = len(device)
    d_logits = Tensor(Tensor.invalids(rows // ndev, VOCAB, dtype=dtypes.bfloat16, device=device).uop.multi(axis), device=device)
+    dname = device[0].split(":")[0]
    rows_per_dev = rows // ndev
  else:
    d_logits = Tensor.invalids(rows, VOCAB, dtype=dtypes.bfloat16, device=device)
+    dname = device.split(":")[0] if isinstance(device, str) else device
    rows_per_dev = rows
  # NOTE: .mean() backward gives same grad per row (1/N), so broadcast is safe; take scalar
  scale = Tensor(gradient, device=device).float().reshape(-1)[0:1].contiguous()
  logits_t = Tensor(logits_u.after(kernel), device=device)
  lse_t = Tensor(lse_u.after(kernel), device=device)
  targets_t = Tensor(targets_u, device=device)
-  fxn = functools.partial(_custom_fused_ce_loss_bwd, vocab=VOCAB, rows=rows_per_dev, label_smoothing=label_smoothing)
+  fxn = functools.partial(_custom_fused_ce_loss_bwd, dname=dname, vocab=VOCAB, rows=rows_per_dev, label_smoothing=label_smoothing)
  d_logits, *_ = Tensor.custom_kernel(d_logits, logits_t, lse_t, targets_t, scale, fxn=fxn)
  return (None, None, None, d_logits.uop, None)

@@ -73,15 +78,17 @@ def fused_ce_loss(logits:Tensor, targets:Tensor, label_smoothing:float=0.1) -> T
                      device=logits.device)
    lse_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0),
                      device=logits.device)
+    dname = logits.device[0].split(":")[0]
    rows_per_dev = rows // ndev
  else:
    loss_out = Tensor.invalids(rows, dtype=dtypes.float32, device=logits.device)
    max_out  = Tensor.invalids(rows, dtype=dtypes.float32, device=logits.device)
    lse_out  = Tensor.invalids(rows, dtype=dtypes.float32, device=logits.device)
+    dname = logits.device.split(":")[0] if isinstance(logits.device, str) else logits.device
    rows_per_dev = rows
  logits_flat = logits.reshape(rows, VOCAB)
  targets_flat = targets.reshape(-1).cast(dtypes.int32)
-  fxn = functools.partial(_custom_fused_ce_loss_fwd, vocab=VOCAB, rows=rows_per_dev,
+  fxn = functools.partial(_custom_fused_ce_loss_fwd, dname=dname, vocab=VOCAB, rows=rows_per_dev,
                          label_smoothing=label_smoothing)
  loss_out, max_out, lse_out, *_ = Tensor.custom_kernel(
    loss_out, max_out, lse_out, logits_flat, targets_flat,
@@ -0,0 +1,104 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+
+// Fused forward sparse-CE with label smoothing.
+// SINGLE-PASS online softmax + vectorized 8-wide bf16 loads for HBM coalescing.
+
+#ifndef VOCAB
+#define VOCAB 128256
+#endif
+#ifndef THREADS_PER_WG
+#define THREADS_PER_WG 256
+#endif
+#ifndef LABEL_SMOOTHING
+#define LABEL_SMOOTHING 0.1f
+#endif
+
+constexpr int VEC = 8;
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
+fused_ce_loss_fwd(
+    float*                __restrict__ loss_out,  // out: fp32, ROWS
+    float*                __restrict__ max_out,   // out: fp32, ROWS
+    float*                __restrict__ lse_out,   // out: fp32, ROWS
+    const __hip_bfloat16* __restrict__ logits,    // in:  bf16, ROWS*VOCAB
+    const int*            __restrict__ targets)   // in:  int32, ROWS
+{
+  __shared__ float sdata_m[THREADS_PER_WG];
+  __shared__ float sdata_s[THREADS_PER_WG];
+  __shared__ float sdata_sumx[THREADS_PER_WG];
+  __shared__ float sdata_tgt[THREADS_PER_WG];
+
+  const int tid = threadIdx.x;
+  const int row = blockIdx.x;
+  const int target = targets[row];
+  const __hip_bfloat16* row_logits = logits + (size_t)row * VOCAB;
+
+  float m = -INFINITY;
+  float s = 0.0f;
+  float sum_x = 0.0f;
+  float target_logit = 0.0f;
+  constexpr bool needs_sum_x = (LABEL_SMOOTHING != 0.0f);
+
+  // Vectorized stride: each iter loads 8 bf16 = 16 bytes. Warp loads 32*16 = 512 bytes (4 cache lines).
+  const int VOCAB_VEC = VOCAB & ~(VEC - 1);  // round down to multiple of VEC
+  for (int i = tid * VEC; i < VOCAB_VEC; i += THREADS_PER_WG * VEC) {
+    float4 raw = *reinterpret_cast<const float4*>(&row_logits[i]);
+    const __hip_bfloat16* xi = reinterpret_cast<const __hip_bfloat16*>(&raw);
+    #pragma unroll
+    for (int k = 0; k < VEC; k++) {
+      const float x = static_cast<float>(xi[k]);
+      if constexpr (needs_sum_x) sum_x += x;
+      if (i + k == target) target_logit = x;
+      if (x > m) {
+        s = s * __expf(m - x) + 1.0f;
+        m = x;
+      } else {
+        s += __expf(x - m);
+      }
+    }
+  }
+  // tail (VOCAB not divisible by VEC):
+  for (int i = VOCAB_VEC + tid; i < VOCAB; i += THREADS_PER_WG) {
+    const float x = static_cast<float>(row_logits[i]);
+    if constexpr (needs_sum_x) sum_x += x;
+    if (i == target) target_logit = x;
+    if (x > m) { s = s * __expf(m - x) + 1.0f; m = x; }
+    else        { s += __expf(x - m); }
+  }
+
+  sdata_m[tid] = m;
+  sdata_s[tid] = s;
+  sdata_sumx[tid] = sum_x;
+  sdata_tgt[tid] = target_logit;
+  __syncthreads();
+
+  for (int step = THREADS_PER_WG / 2; step > 0; step >>= 1) {
+    if (tid < step) {
+      const float m1 = sdata_m[tid];
+      const float m2 = sdata_m[tid + step];
+      const float s1 = sdata_s[tid];
+      const float s2 = sdata_s[tid + step];
+      const float m_new = fmaxf(m1, m2);
+      const float s_new = s1 * __expf(m1 - m_new) + s2 * __expf(m2 - m_new);
+      sdata_m[tid] = m_new;
+      sdata_s[tid] = s_new;
+      sdata_sumx[tid] += sdata_sumx[tid + step];
+      sdata_tgt[tid]  += sdata_tgt[tid + step];
+    }
+    __syncthreads();
+  }
+
+  if (tid == 0) {
+    const float row_max = sdata_m[0];
+    const float row_sum_exp = sdata_s[0];
+    const float row_sum_x = sdata_sumx[0];
+    const float tgt = sdata_tgt[0];
+    const float row_lse = logf(row_sum_exp) + row_max;
+    const float mean_logits = row_sum_x / static_cast<float>(VOCAB);
+    const float loss = row_lse - (1.0f - LABEL_SMOOTHING) * tgt - LABEL_SMOOTHING * mean_logits;
+    loss_out[row] = loss;
+    max_out[row]  = row_max;
+    lse_out[row]  = row_lse;
+  }
+}
@@ -0,0 +1,58 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+
+// Vectorized CE bwd: 8-wide bf16 loads + stores.
+
+#ifndef VOCAB
+#define VOCAB 128256
+#endif
+#ifndef THREADS_PER_WG
+#define THREADS_PER_WG 256
+#endif
+#ifndef LABEL_SMOOTHING
+#define LABEL_SMOOTHING 0.1f
+#endif
+
+constexpr int VEC = 8;
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
+fused_ce_loss_bwd(
+    __hip_bfloat16*       __restrict__ d_logits,
+    const __hip_bfloat16* __restrict__ logits,
+    const float*          __restrict__ lse,
+    const int*            __restrict__ targets,
+    const float*          __restrict__ scale_in)
+{
+  const int tid = threadIdx.x;
+  const int row = blockIdx.x;
+  const int target = targets[row];
+  const float lse_r = lse[row];
+  const __hip_bfloat16* row_logits  = logits   + (size_t)row * VOCAB;
+  __hip_bfloat16*       row_dlogits = d_logits + (size_t)row * VOCAB;
+  const float inv_vocab = 1.0f / static_cast<float>(VOCAB);
+  const float scale = *scale_in;
+  const float ls_term = LABEL_SMOOTHING * inv_vocab;
+
+  const int VOCAB_VEC = VOCAB & ~(VEC - 1);
+  for (int i = tid * VEC; i < VOCAB_VEC; i += THREADS_PER_WG * VEC) {
+    float4 raw = *reinterpret_cast<const float4*>(&row_logits[i]);
+    const __hip_bfloat16* xi = reinterpret_cast<const __hip_bfloat16*>(&raw);
+    __hip_bfloat16 out[VEC];
+    #pragma unroll
+    for (int k = 0; k < VEC; k++) {
+      const float x = static_cast<float>(xi[k]);
+      float g = __expf(x - lse_r);
+      if (i + k == target) g -= (1.0f - LABEL_SMOOTHING);
+      g -= ls_term;
+      out[k] = static_cast<__hip_bfloat16>(g * scale);
+    }
+    *reinterpret_cast<float4*>(&row_dlogits[i]) = *reinterpret_cast<float4*>(out);
+  }
+  for (int i = VOCAB_VEC + tid; i < VOCAB; i += THREADS_PER_WG) {
+    const float x = static_cast<float>(row_logits[i]);
+    float g = __expf(x - lse_r);
+    if (i == target) g -= (1.0f - LABEL_SMOOTHING);
+    g -= ls_term;
+    row_dlogits[i] = static_cast<__hip_bfloat16>(g * scale);
+  }
+}
@@ -0,0 +1,55 @@
+from __future__ import annotations
+import functools, pathlib
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from extra.llama_kernels import THREADS_PER_WG, dname_of, compile_hip
+
+ELEMS_PER_THREAD = 8  # vectorized 16-byte load (uint4 = 8 bf16)
+
+def _build_src(n_chunks:int) -> str:
+  template = (pathlib.Path(__file__).parent/"fused_pad_grad_accum.cpp").read_text()
+  params = "".join(f",\n    const __hip_bfloat16* __restrict__ chunk{i}" for i in range(n_chunks))
+  dispatch = "\n    ".join(f"case {i}: chunk_ptr = chunk{i}; break;" for i in range(n_chunks))
+  return (template.replace("__FUSED_PAD_GRAD_ACCUM_PARAMS", params)
+                  .replace("__FUSED_PAD_GRAD_ACCUM_DISPATCH", dispatch))
+
+@functools.cache
+def _custom_fused_pad_grad_accum(grad_buf:UOp, *chunk_uops, dname:str, n_chunks:int, chunk_size:int) -> UOp:
+  total = n_chunks * chunk_size
+  elems_per_block = THREADS_PER_WG * ELEMS_PER_THREAD
+  assert chunk_size % elems_per_block == 0, f"chunk_size {chunk_size} must be multiple of {elems_per_block}"
+  num_wg = total // elems_per_block
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(num_wg, "gidx0")
+  mem = total * 2 * 3
+  sink = UOp.sink(grad_buf.base, *(c.base for c in chunk_uops), threads, workgroups,
+                  arg=KernelInfo(f"fused_pad_grad_accum_n{n_chunks}_c{chunk_size}",
+                                 estimates=Estimates(ops=2*total, mem=mem)))
+  src = _build_src(n_chunks)
+  defines = [f"-DCHUNK_SIZE={chunk_size}", f"-DTHREADS_PER_WG={THREADS_PER_WG}", f"-DELEMS_PER_THREAD={ELEMS_PER_THREAD}"]
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
+
+def can_fused_pad_grad_accum(grad_buf:Tensor, chunks:list[Tensor]) -> bool:
+  if not chunks or grad_buf.dtype != dtypes.bfloat16: return False
+  if any(c.dtype != dtypes.bfloat16 for c in chunks): return False
+  chunk_shape = chunks[0].shape
+  if any(c.shape != chunk_shape for c in chunks): return False
+  chunk_size, total = 1, 1
+  for d in chunk_shape: chunk_size *= d
+  for d in grad_buf.shape: total *= d
+  return total == len(chunks) * chunk_size and chunk_size % (THREADS_PER_WG * ELEMS_PER_THREAD) == 0
+
+def fused_pad_grad_accum(grad_buf:Tensor, chunks:list[Tensor]) -> Tensor:
+  # NOTE: grad_buf += cat(*chunks, dim=0) in one HBM pass (in-place add). Returns new grad_buf Tensor.
+  # Requires uniform chunk shapes and chunk_size % (THREADS_PER_WG*ELEMS_PER_THREAD) == 0.
+  assert chunks and grad_buf.dtype == dtypes.bfloat16
+  for c in chunks: assert c.dtype == dtypes.bfloat16, f"chunk dtype must be bf16, got {c.dtype}"
+  chunk_size, total = 1, 1
+  for d in chunks[0].shape: chunk_size *= d
+  for d in grad_buf.shape: total *= d
+  assert total == len(chunks) * chunk_size, f"grad_buf size {total} != n_chunks {len(chunks)} * chunk_size {chunk_size}"
+  fxn = functools.partial(_custom_fused_pad_grad_accum, dname=dname_of(grad_buf.device),
+                          n_chunks=len(chunks), chunk_size=chunk_size)
+  out, *_ = Tensor.custom_kernel(grad_buf, *chunks, fxn=fxn)
+  return out
@@ -0,0 +1,63 @@
+// Fused custom kernel: grad_buf += cat(*chunks, dim=0) in one HBM pass.
+//
+// Template source — chunk parameter list and switch dispatch are filled by codegen
+// in cast_amax.py:_build_fused_pad_grad_accum_src to support arbitrary N.
+//
+// Defines required at compile time:
+//   CHUNK_SIZE       elements per chunk (must be multiple of THREADS_PER_WG * ELEMS_PER_THREAD)
+//   THREADS_PER_WG
+//   ELEMS_PER_THREAD (8 = one uint4 per thread = 16-byte vectorized load)
+//
+// Layout: one block-per-(slice-of-chunk) — blockIdx.x / BLOCKS_PER_CHUNK selects the chunk.
+// All threads in a block read the same chunk → switch is uniform → no warp divergence.
+
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+
+#ifndef THREADS_PER_WG
+#define THREADS_PER_WG 256
+#endif
+#ifndef ELEMS_PER_THREAD
+#define ELEMS_PER_THREAD 8
+#endif
+
+#define ELEMS_PER_BLOCK (THREADS_PER_WG * ELEMS_PER_THREAD)
+#define BLOCKS_PER_CHUNK (CHUNK_SIZE / ELEMS_PER_BLOCK)
+
+extern "C" __attribute__((global))
+__attribute__((amdgpu_flat_work_group_size(1, THREADS_PER_WG)))
+void fused_pad_grad_accum(
+    __hip_bfloat16* __restrict__ grad_buf
+    __FUSED_PAD_GRAD_ACCUM_PARAMS
+) {
+  const int bid = blockIdx.x;
+  const int chunk_idx = bid / BLOCKS_PER_CHUNK;
+  const int block_in_chunk = bid - chunk_idx * BLOCKS_PER_CHUNK;
+  const int tid = threadIdx.x;
+
+  const __hip_bfloat16* chunk_ptr;
+  switch (chunk_idx) {
+    __FUSED_PAD_GRAD_ACCUM_DISPATCH
+    default: chunk_ptr = (const __hip_bfloat16*)0; break;  // unreachable
+  }
+
+  // int64 for global_offset: at 32 chunks × 117M elements = 3.6B, int32 overflows → MEMVIOL.
+  const int local_offset = block_in_chunk * ELEMS_PER_BLOCK + tid * ELEMS_PER_THREAD;
+  const long long global_offset = (long long)chunk_idx * (long long)CHUNK_SIZE + (long long)local_offset;
+
+  // Vectorized 16-byte load (uint4 = 8 bf16). Requires CHUNK_SIZE % 8 == 0 and 16-byte alignment.
+  const uint4 chunk_v = *reinterpret_cast<const uint4*>(&chunk_ptr[local_offset]);
+  const uint4 grad_v  = *reinterpret_cast<const uint4*>(&grad_buf[global_offset]);
+  uint4 out_v;
+
+  const __hip_bfloat16* chunk_bf = reinterpret_cast<const __hip_bfloat16*>(&chunk_v);
+  const __hip_bfloat16* grad_bf  = reinterpret_cast<const __hip_bfloat16*>(&grad_v);
+  __hip_bfloat16*       out_bf   = reinterpret_cast<__hip_bfloat16*>(&out_v);
+
+  #pragma unroll
+  for (int i = 0; i < ELEMS_PER_THREAD; i++) {
+    out_bf[i] = (__hip_bfloat16)((float)grad_bf[i] + (float)chunk_bf[i]);
+  }
+
+  *reinterpret_cast<uint4*>(&grad_buf[global_offset]) = out_v;
+}
@@ -1,64 +1,35 @@
-import functools
+from __future__ import annotations
+import functools, pathlib
 from tinygrad import Tensor, dtypes
-from tinygrad.dtype import AddrSpace
-from tinygrad.helpers import prod
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
-from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, scalar_amax
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, scalar_amax, dname_of, compile_hip

@functools.cache
-def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_partial:UOp, x:UOp, amax_state:UOp) -> UOp:
-  VEC = 8
-  n_elems = prod(x.shape)
-  assert n_elems % (NUM_WG * THREADS_PER_WG * VEC) == 0
-  assert amax_partial.shape[0] == NUM_WG
-
-  x = x.reshape(n_elems)
-  fp8_out = fp8_out.reshape(n_elems)
-
-  wg = UOp.range(NUM_WG, 0, AxisType.GLOBAL)
-  tid = UOp.range(THREADS_PER_WG, 1, AxisType.LOCAL)
-  it = UOp.range((n_elems // VEC) // (NUM_WG * THREADS_PER_WG), 2, AxisType.LOOP)
-  lane = UOp.range(VEC, 3, AxisType.UNROLL)
-
-  idx = (((it * NUM_WG + wg) * THREADS_PER_WG + tid) * VEC) + lane
-
-  scale = FP8_MAX / (amax_state[0].cast(dtypes.float) + 1e-8)
-  x_f = x[idx].cast(dtypes.float)
-  abs_x = (x_f < 0.0).where(-x_f, x_f)
-  scaled = (x_f * scale).maximum(-FP8_MAX).minimum(FP8_MAX)
-
-  fp8_store = fp8_out[idx].store(scaled.cast(fp8_out.dtype.base)).end(lane)
-  lane_max = abs_x.reduce(lane, arg=Ops.MAX)
-
-  lmax = UOp.placeholder((1,), dtypes.float, slot=1, addrspace=AddrSpace.REG)
-  lmax_init = lmax.after(wg, tid)[0].store(0.0)
-  lmax_prev = lmax.after(lmax_init, it)[0]
-  lmax_store = lmax.after(fp8_store)[0].store(lmax_prev.maximum(lane_max))
-  lmax_val = lmax.after(lmax_store.end(it))[0]
-
-  lds = UOp.placeholder((THREADS_PER_WG,), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
-  lds = lds.after(lds[tid].store(lmax_val).barrier())
-
-  step = THREADS_PER_WG // 2
-  while step:
-    active = tid < step
-    other = lds[tid + step].load(UOp.const(dtypes.float, 0.0), active)
-    lds = lds.after(lds[tid].store(lds[tid].maximum(other), gate=active).barrier())
-    step //= 2
-
-  amax_store = amax_partial[tid.eq(0).where(wg, UOp.invalid())].store(lds[0])
-  return amax_store.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=()))
+def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_partial:UOp, x:UOp, amax_state:UOp, dname:str) -> UOp:
+  n_elems = 1
+  for d in x.shape: n_elems *= d
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
+  mem = n_elems * 2 + n_elems + 4 + NUM_WG * 4
+  sink = UOp.sink(fp8_out.base, amax_partial.base, x.base, amax_state.base, threads, workgroups,
+                  arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", estimates=Estimates(ops=3*n_elems, mem=mem)))
+  src = (pathlib.Path(__file__).parent/"quantize_fp8_with_amax.cpp").read_text()
+  defines = [f"-DN_ELEMS={n_elems}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))

@functools.cache
-def _custom_quantize_fp8_scalar(fp8_out:UOp, x:UOp, amax_state:UOp) -> UOp:
-  n_elems = prod(x.shape)
-  i = UOp.range(n_elems, 0)
-
-  x_f = x.reshape(n_elems)[i].cast(dtypes.float)
-  scale = FP8_MAX / (amax_state[0].cast(dtypes.float) + 1e-8)
-  store = fp8_out.reshape(n_elems)[i].store((x_f * scale).cast(fp8_out.dtype.base))
-
-  return store.end(i).sink(arg=KernelInfo(f"quantize_fp8_scalar_{n_elems}"))
+def _custom_quantize_fp8_scalar(fp8_out:UOp, x:UOp, amax_state:UOp, dname:str) -> UOp:
+  n_elems = 1
+  for d in x.shape: n_elems *= d
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
+  mem = n_elems * 2 + n_elems
+  sink = UOp.sink(fp8_out.base, x.base, amax_state.base, threads, workgroups,
+                  arg=KernelInfo(f"quantize_fp8_scalar_{n_elems}", estimates=Estimates(ops=2*n_elems, mem=mem)))
+  src = (pathlib.Path(__file__).parent/"quantize_fp8_scalar.cpp").read_text()
+  defines = [f"-DN_ELEMS={n_elems}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))

 def _quantize_fp8_delayed_bwd(gradient:UOp, kernel:UOp):
  # NOTE: STE-equivalent backward — grad_x = grad_fp8 * scale, scale = FP8_MAX / amax_state.
@@ -78,10 +49,8 @@ def quantize_fp8_delayed(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3)
  assert x.dtype == dtypes.bfloat16, f"expected bf16, got {x.dtype}"
  axis = x.uop.axis if isinstance(x.device, tuple) else None
  fp8_out      = alloc_like(x.shape,  fp8_dtype,      x.device, axis)
-  n_elems = prod(x.uop.shard_shape)
-  assert n_elems % NUM_WG == 0, f"{n_elems=} must divide over {NUM_WG=}"
  amax_partial = alloc_local((NUM_WG,), dtypes.float32, x.device, axis)
-  fxn = _custom_quantize_fp8_with_amax
+  fxn = functools.partial(_custom_quantize_fp8_with_amax, dname=dname_of(x.device))
  fp8_out, amax_partial, *_ = Tensor.custom_kernel(fp8_out, amax_partial, x, amax_state,
                                                    fxn=fxn, grad_fxn=_quantize_fp8_delayed_bwd)
  new_amax = scalar_amax(amax_partial)
@@ -93,6 +62,6 @@ def quantize_fp8_scalar(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3) -
  # NOTE: pure one-pass bf16 -> fp8 quantize with delayed scalar scale. No amax computation.
  axis = x.uop.axis if isinstance(x.device, tuple) else None
  fp8_out = alloc_like(x.shape, fp8_dtype, x.device, axis)
-  fxn = _custom_quantize_fp8_scalar
+  fxn = functools.partial(_custom_quantize_fp8_scalar, dname=dname_of(x.device))
  fp8_out, *_ = Tensor.custom_kernel(fp8_out, x, amax_state, fxn=fxn)
  return fp8_out
@@ -0,0 +1,48 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+#include <hip/hip_fp8.h>
+
+// Pure one-pass bf16 -> fp8 quantize with delayed scalar scale. No amax computation.
+
+#ifndef N_ELEMS
+#define N_ELEMS 67108864
+#endif
+#ifndef NUM_WG
+#define NUM_WG 1024
+#endif
+#ifndef THREADS_PER_WG
+#define THREADS_PER_WG 256
+#endif
+
+constexpr int VEC = 8;
+constexpr float FP8_MAX = 448.0f;
+
+static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC");
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
+quantize_fp8_scalar(
+    __hip_fp8_storage_t*  __restrict__ fp8_out,         // fp8, N_ELEMS
+    const __hip_bfloat16* __restrict__ x,               // bf16, N_ELEMS
+    const float*          __restrict__ amax_state)      // fp32 scalar (delayed)
+{
+  const int tid = threadIdx.x;
+  const int wg  = blockIdx.x;
+  const int gid = wg * THREADS_PER_WG + tid;
+  const int stride_elems = NUM_WG * THREADS_PER_WG * VEC;
+
+  const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
+
+  for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) {
+    float4 x_raw = *reinterpret_cast<const float4*>(&x[base]);
+    const __hip_bfloat16 *xi = reinterpret_cast<const __hip_bfloat16*>(&x_raw);
+
+    __hip_fp8_storage_t out[VEC];
+    #pragma unroll
+    for (int i = 0; i < VEC; i++) {
+      const float v = static_cast<float>(xi[i]);
+      const float scaled = fmaxf(-FP8_MAX, fminf(FP8_MAX, v * scale));
+      out[i] = __hip_cvt_float_to_fp8(scaled, __HIP_SATFINITE, __HIP_E4M3);
+    }
+    *reinterpret_cast<uint64_t*>(&fp8_out[base]) = *reinterpret_cast<uint64_t*>(out);
+  }
+}
@@ -0,0 +1,63 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+#include <hip/hip_fp8.h>
+
+// One-pass bf16 -> fp8 quantize using a scalar delayed amax state,
+// AND simultaneously computes per-WG |x| max partials for the next step's amax state.
+// Saves one full HBM pass over the grad tensor vs. doing quantize + separate abs().max().
+
+#ifndef N_ELEMS
+#define N_ELEMS 67108864
+#endif
+#ifndef NUM_WG
+#define NUM_WG 1024
+#endif
+#ifndef THREADS_PER_WG
+#define THREADS_PER_WG 256
+#endif
+
+constexpr int VEC = 8;
+constexpr float FP8_MAX = 448.0f;
+
+static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC");
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
+quantize_fp8_with_amax(
+    __hip_fp8_storage_t*  __restrict__ fp8_out,       // out: fp8, N_ELEMS
+    float*                __restrict__ amax_partial,  // out: fp32, NUM_WG per-WG partials
+    const __hip_bfloat16* __restrict__ x,             // in:  bf16, N_ELEMS
+    const float*          __restrict__ amax_state)    // in:  fp32 scalar (delayed)
+{
+  __shared__ float sdata[THREADS_PER_WG];
+
+  const int tid = threadIdx.x;
+  const int wg  = blockIdx.x;
+  const int gid = wg * THREADS_PER_WG + tid;
+  const int stride_elems = NUM_WG * THREADS_PER_WG * VEC;
+
+  const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
+  float local_max = 0.0f;
+
+  for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) {
+    float4 x_raw = *reinterpret_cast<const float4*>(&x[base]);
+    const __hip_bfloat16 *xi = reinterpret_cast<const __hip_bfloat16*>(&x_raw);
+
+    __hip_fp8_storage_t out[VEC];
+    #pragma unroll
+    for (int i = 0; i < VEC; i++) {
+      const float v = static_cast<float>(xi[i]);
+      local_max = fmaxf(local_max, fabsf(v));
+      const float scaled = fmaxf(-FP8_MAX, fminf(FP8_MAX, v * scale));
+      out[i] = __hip_cvt_float_to_fp8(scaled, __HIP_SATFINITE, __HIP_E4M3);
+    }
+    *reinterpret_cast<uint64_t*>(&fp8_out[base]) = *reinterpret_cast<uint64_t*>(out);
+  }
+
+  sdata[tid] = local_max;
+  __syncthreads();
+  for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
+    if (tid < s) sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]);
+    __syncthreads();
+  }
+  if (tid == 0) amax_partial[wg] = sdata[0];
+}
@@ -6,7 +6,7 @@ from tinygrad.tensor import Tensor
 class LR_Scheduler:
  def __init__(self, optimizer: Optimizer):
    self.optimizer = optimizer
-    self.epoch_counter = Tensor([0], device=self.optimizer.device)
+    self.epoch_counter = Tensor([0], requires_grad=False, device=self.optimizer.device)

  def get_lr(self): pass

@@ -52,7 +52,7 @@ class BertForPretraining:
  # Reference has residual on denominator: https://github.com/mlcommons/training/blob/master/language_model/tensorflow/bert/run_pretraining.py#L315
  def sparse_categorical_crossentropy(self, predictions:Tensor, labels:Tensor, ignore_index=-1):
    log_probs, loss_mask = predictions.log_softmax(dtype=dtypes.float), (labels != ignore_index)
-    y_counter = Tensor.arange(predictions.shape[-1], device=predictions.device).unsqueeze(0).expand(labels.numel(), predictions.shape[-1])
+    y_counter = Tensor.arange(predictions.shape[-1], requires_grad=False, device=predictions.device).unsqueeze(0).expand(labels.numel(), predictions.shape[-1])
    y = ((y_counter == labels.flatten().reshape(-1, 1)) * loss_mask.reshape(-1, 1)).reshape(*labels.shape, predictions.shape[-1])
    return -((log_probs * y).sum()) / (loss_mask.sum() + 1e-5) # Small constant to avoid division by zero

@@ -159,7 +159,7 @@ class BertPooler:
    return self.dense(hidden_states[:, 0]).tanh()

 def gather(prediction_logits:Tensor, masked_lm_positions:Tensor):
-  counter = Tensor.arange(prediction_logits.shape[1], device=prediction_logits.device).reshape(1, 1, prediction_logits.shape[1]).expand(*masked_lm_positions.shape, prediction_logits.shape[1])
+  counter = Tensor.arange(prediction_logits.shape[1], device=prediction_logits.device, requires_grad=False).reshape(1, 1, prediction_logits.shape[1]).expand(*masked_lm_positions.shape, prediction_logits.shape[1])
  onehot = counter == masked_lm_positions.unsqueeze(2).expand(*masked_lm_positions.shape, prediction_logits.shape[1])
  return onehot @ prediction_logits

@@ -189,7 +189,7 @@ class BertEmbeddings:
    input_shape = input_ids.shape
    seq_length = input_shape[1]

-    position_ids = Tensor.arange(seq_length, device=input_ids.device).unsqueeze(0).expand(*input_shape)
+    position_ids = Tensor.arange(seq_length, requires_grad=False, device=input_ids.device).unsqueeze(0).expand(*input_shape)
    words_embeddings = self.word_embeddings(input_ids)
    position_embeddings = self.position_embeddings(position_ids)
    token_type_embeddings = self.token_type_embeddings(token_type_ids)
@@ -201,7 +201,7 @@ class Transformer:
    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.max_context = max_context
-    self.freqs_cis = precompute_freqs_cis(dim // n_heads, self.max_context * 2, rope_theta).contiguous().is_param_(False)
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, self.max_context * 2, rope_theta).contiguous().requires_grad_(False)
    self.forward_jit = TinyJit(self.forward) if jit else None

  def forward(self, tokens:Tensor, start_pos:Union[Variable,int], temperature:float, top_k:int, top_p:float, alpha_f:float, alpha_p:float):
@@ -78,7 +78,7 @@ def tensor_getitem(tensor, *keys):
 # for gather with indicies only on axis=0
 def tensor_gather(tensor, indices):
  if not isinstance(indices, Tensor):
-    indices = Tensor(indices)
+    indices = Tensor(indices, requires_grad=False)
  if len(tensor.shape) > 2:
    rem_shape = list(tensor.shape)[1:]
    tensor = tensor.reshape(tensor.shape[0], -1)
@@ -15,7 +15,7 @@ class RNNT:
  @TinyJit
  def __call__(self, x, y, hc=None):
    f, _ = self.encoder(x, None)
-    g, _ = self.prediction(y, hc, Tensor.ones(1))
+    g, _ = self.prediction(y, hc, Tensor.ones(1, requires_grad=False))
    out = self.joint(f, g)
    return out.realize()

@@ -30,10 +30,10 @@ class RNNT:
    return outputs

  def _greedy_decode(self, logits, logit_len):
-    hc = Tensor.zeros(self.prediction.rnn.layers, 2, self.prediction.hidden_size)
+    hc = Tensor.zeros(self.prediction.rnn.layers, 2, self.prediction.hidden_size, requires_grad=False)
    labels = []
-    label = Tensor.zeros(1, 1)
-    mask = Tensor.zeros(1)
+    label = Tensor.zeros(1, 1, requires_grad=False)
+    mask = Tensor.zeros(1, requires_grad=False)
    for time_idx in range(logit_len):
      logit = logits[time_idx, :, :].unsqueeze(0)
      not_blank = True
@@ -41,7 +41,7 @@ class RNNT:
      while not_blank and added < 30:
        if len(labels) > 0:
          mask = (mask + 1).clip(0, 1)
-          label = Tensor([[labels[-1] if labels[-1] <= 28 else labels[-1] - 1]]) + 1 - 1
+          label = Tensor([[labels[-1] if labels[-1] <= 28 else labels[-1] - 1]], requires_grad=False) + 1 - 1
        jhc = self._pred_joint(Tensor(logit.numpy()), label, hc, mask)
        k = jhc[0, 0, :29].argmax(axis=0).numpy()
        not_blank = k != 28
@@ -129,7 +129,7 @@ class LSTM:
      return self.do_step(x_, hc_)

    if hc is None:
-      hc = Tensor.zeros(self.layers, 2 * x.shape[1], self.hidden_size).contiguous().realize()
+      hc = Tensor.zeros(self.layers, 2 * x.shape[1], self.hidden_size, requires_grad=False).contiguous().realize()

    output = None
    for t in range(x.shape[0]):
@@ -41,7 +41,7 @@ class TransformerBlock:
 class Transformer:
  def __init__(self, syms, maxlen, layers, embed_dim, num_heads, ff_dim):
    self.maxlen, self.syms = maxlen, syms
-    self.embed = Tensor.scaled_uniform(maxlen+syms, embed_dim).is_param_(False)
+    self.embed = Tensor.scaled_uniform(maxlen+syms, embed_dim, requires_grad=False)
    self.tbs = [TransformerBlock(embed_dim, num_heads, ff_dim) for _ in range(layers)]
    self.final = Tensor.scaled_uniform(embed_dim, syms)

@@ -1,4 +1,5 @@
-from tinygrad import Tensor, Device, dtypes, nn
+from tinygrad import Tensor, dtypes, nn
+from tinygrad.device import is_dtype_supported
 from typing import Optional, Union, List, Any, Tuple, Callable
 import math

@@ -12,7 +13,7 @@ def timestep_embedding(timesteps:Tensor, dim:int, max_period=10000):
  freqs = (-math.log(max_period) * Tensor.arange(half, device=timesteps.device) / half).exp()
  args = timesteps.unsqueeze(1) * freqs.unsqueeze(0)
  out = Tensor.cat(args.cos(), args.sin(), dim=-1)
-  return out.cast(mixed_precision_dtype) if mixed_precision_dtype in Device[Device.DEFAULT].renderer.supported_dtypes() else out
+  return out.cast(mixed_precision_dtype) if is_dtype_supported(mixed_precision_dtype) else out

 class ResBlock:
  def __init__(self, channels:int, emb_channels:int, out_channels:int, num_groups:int=32):
@@ -237,7 +238,7 @@ class UNetModel:
      assert y.shape[0] == x.shape[0]
      emb = emb + y.sequential(self.label_emb[0])

-    if mixed_precision_dtype in Device[Device.DEFAULT].renderer.supported_dtypes():
+    if is_dtype_supported(mixed_precision_dtype):
      emb = emb.cast(mixed_precision_dtype)
      ctx = ctx.cast(mixed_precision_dtype)
      x   = x  .cast(mixed_precision_dtype)
@@ -1,7 +1,7 @@
 import unittest
 import numpy as np

-from tinygrad.helpers import BEAM, Timing, prod
+from tinygrad.helpers import BEAM, Timing, CI, prod
 from tinygrad import Variable, Device, Tensor
 from tinygrad.nn import Conv2d
 from tinygrad.uop.ops import AxisType, Ops
@@ -64,7 +64,7 @@ class TestBeamSearch(unittest.TestCase):
    actual = a.numpy()
    np.testing.assert_allclose(actual, desired)

-  @unittest.skip("flaky. CL_OUT_OF_RESOURCES")
+  @unittest.skipIf(CI, "flaky. CL_OUT_OF_RESOURCES")
  def test_conv_beam(self):
    c = Conv2d(3, 16, (3,3))
    x = rand(1,3,32,32)
@@ -84,6 +84,8 @@ def serve(conn:socket.socket):
      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)
@@ -1,7 +1,16 @@
 #!/bin/sh
 install_loc="$HOME/.local/bin"
-docker pull --platform=linux/amd64 rocm/dev-ubuntu-22.04:7.1.1
-docker tag rocm/dev-ubuntu-22.04:7.1.1 rocm-hipcc:7.1.1
+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'
@@ -12,7 +21,7 @@ if ! docker inspect --format='{{.State.Running}}' "$cname" 2>/dev/null | grep -q
  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.1.1 sleep 300 >/dev/null
+    rocm-hipcc:7.2 sleep 300 >/dev/null
 fi
 exec docker exec "$cname" "$(basename "$0")" "$@"
 EOF
@@ -1,294 +0,0 @@
-#!/usr/bin/env python3
-"""
-Self-contained Qualcomm Hexagon **testsig** generator.
-
-Replicates: python2 elfsigner.py -t 0x67489311 -o .
-Dependencies: standard library + cryptography (pip install cryptography).
-Multiple serial numbers: use -t multiple times.
-"""
-
-import argparse, base64, hashlib, os, struct
-from datetime import datetime, timedelta, timezone
-from cryptography import x509
-from cryptography.hazmat.primitives import hashes, serialization
-from cryptography.hazmat.primitives.asymmetric import rsa
-from cryptography.x509.name import _ASN1Type
-from cryptography.x509.oid import ExtensionOID, NameOID, ObjectIdentifier
-
-# Embedded assets (raw base64) --- no external files needed
-# Compact test_elf_nop.so contents used for signing. Only bytes covered by
-# program headers are needed; section headers and alignment gaps are not.
-ORIG_PHDRS = [
-    {'t': 1, 'o': 0x0000, 'v': 0x0000, 'p': 0x0000, 'fs': 0x02fc, 'ms': 0x02fc, 'fl': 0x4, 'al': 0x1000},
-    {'t': 1, 'o': 0x1000, 'v': 0x1000, 'p': 0x1000, 'fs': 0x0104, 'ms': 0x0104, 'fl': 0x5, 'al': 0x1000},
-    {'t': 1, 'o': 0x2000, 'v': 0x2000, 'p': 0x2000, 'fs': 0x0004, 'ms': 0x0004, 'fl': 0x4, 'al': 0x1000},
-    {'t': 1, 'o': 0x3000, 'v': 0x4000, 'p': 0x4000, 'fs': 0x00d0, 'ms': 0x0100, 'fl': 0x6, 'al': 0x1000},
-    {'t': 2, 'o': 0x3010, 'v': 0x4010, 'p': 0x4010, 'fs': 0x00a8, 'ms': 0x00a8, 'fl': 0x6, 'al': 0x4},
-]
-ORIG_SEGS = {
-    0x0000: base64.b64decode("""
-f0VMRgEBAQAAAAAAAAAAAAMApAABAAAAsBAAADQAAACIMQAAAwAAADQAIAAFACgAFQASAAEAAAAAAAAAAAAAAAAAAAD8AgAA/AIA
-AAQAAAAAEAAAAQAAAAAQAAAAEAAAABAAAAQBAAAEAQAABQAAAAAQAAABAAAAACAAAAAgAAAAIAAABAAAAAQAAAAEAAAAABAAAAEA
-AAAAMAAAAEAAAABAAADQAAAAAAEAAAYAAAAAEAAAAgAAABAwAAAQQAAAEEAAAKgAAACoAAAABgAAAAQAAAADAAAAEwAAABIAAAAR
-AAAADgAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAACgAAAAAAAAAJAAAACwAAAAwAAAANAAAA
-DwAAABAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAQAAAAAAAAAwAFAAAAAACwEAAAAAAAAAMABwAAAAAAwBAAAAAAAAADAAgAAAAA
-AAAgAAAAAAAAAwAJAAAAAAAAQAAAAAAAAAMACgAAAAAACEAAAAAAAAADAAsAAAAAAMxAAAAAAAAAAwAPAAAAAAAAQQAAAAAAAAMA
-EAAPAAAAABAAAGwAAAASAAUAFQAAAAAAAAAAAAAAEgAAAGkAAAAAQQAAAAAAABAA8f9KAAAAzEAAAAQAAAARAA8AMQAAAARAAAAA
-AAAAEAAKAAEAAAAIQAAAAAAAABAACwBcAAAAwBAAAEQAAAASAAgAYgAAANBAAAAAAAAAEADx/3UAAAAAQQAAAAAAABAA8f9GAAAA
-sBAAAAQAAAASAAcAAF9fRFRPUl9MSVNUX18AX2luaXQAX19yZWdpc3Rlcl9mcmFtZV9pbmZvX2Jhc2VzAF9fQ1RPUl9FTkRfXwBs
-aWJjLnNvAG5vcABub3BfdmFyAGxpYmdjYy5zbwBfZmluaQBfZWRhdGEAX19ic3Nfc3RhcnQAX2VuZAB0ZXN0X2VsZl9ub3Auc28A
-AADIQAAAIgoAAAAAAAA=
-"""),
-    0x1000: base64.b64decode("""
-AcCdoADbnaEB2J2hGMAJalTP6nH//+pyGNgq8///4HJI3+BxAMAY8wHAgJEIwAEQAkAAeAEoAyg0wABa///7ckz/+3Eb2xjzm//7
-vwDAm5EGwAAQAMCgUPj//1k4wJ2RG0CdkR7AHpAAwJ9SAAAAAMFAAAAcxElqDkKc4k9AnJE8wJyRDkIOjADAnFIAAAAAAAAAAAAA
-AAAAAAAAAAAAAMBAAAAO1ElqHMCOkQDAnFIAwJ9SAAAAAAAAAAAAAAAAAcCdoADbnaEPwAlqENDqcf//6nIPzyrz///7clD/+3Eb
-2w/zm8AbsADAm5EGwAAQAMCgUPj//1kbQJ2RHsAekADAn1I=
-"""),
-    0x2000: base64.b64decode("AAAAAA=="),
-    0x3000: base64.b64decode("""
-AAAAAAAAAAAAAAAAAAAAAAEAAAA+AAAAAQAAAFIAAAAOAAAAegAAAAwAAAAAEAAADQAAAMAQAAAEAAAA1AAAAAUAAABkAgAABgAA
-ADQBAAAKAAAAigAAAAsAAAAQAAAAAwAAALhAAAACAAAADAAAABQAAAAHAAAAFwAAAPACAAABAABwAwAAAAAAAAAAAAAAAAAAAAAA
-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABBAAAAAAAAAAAAAAAAAAABwEAAACgAAAA==
-"""),
-}
-ORIG_EHDR = ORIG_SEGS[0][:0x34]
-
-ATTESTCA_KEY = base64.b64decode("""LS0tLS1CRUdJTiBSU0EgUFJJVkFURSBLRVktLS0tLQpNSUlFb3dJQkFBS0NBUUVBeGpLNkRZbVZxbGIvZHFtb1BtY1FsNkV6Vmx3S3dIUnFCQmdsZzBkQnd6djJMR1NKCi9xWnlpbkFWZlloaEFkKzZwbm1kYk1vMVMzaXdjNXUrWUtRYlFOSGpVMUZrZDN4c29weHA3YTZHNnZUbSsyNVEKVnRHRkdLRHk3dm9mRWRKNTE2NVpQU0JtZ0thcWFQcFIxRlNHQnRDaEV3bzhkVEFZdjBaeXFYeGlmeDU4ejcvbgozT1BBbDJQdkpjZWpzd1J2ekRvcVRDR0xlMjVMbzFiYlQ1RGM4bFhZeWFmcjJVcGpXc1ZBaGlHQVNaK2dIRzYrCk1xWDFUdUlvM3FJTGNOcnRNUVh6K2NMamZaSnYzZExqNmRFSlZ0c3MvT3lMdUQ2bVozMXZUdzMxakhXUGRLeDIKNHh3dXVrbkp5OUVrUElQem5ORTZGcFZ4dmdHQWN4aEhKY0pLdndJQkF3S0NBUUVBaENIUlhsdTVIRG4vcEhFYQoxRVMxdW10M2p1Z0hLdmhHcldWdVY0VFdnaWY1Y3UyeFZHNzNCdlZqcVFXV0FUL1J4RkVUbmR3amg2WEs5NzBwCmxjSzgxZUZDTjR1WVQ2aEliR2hHbm5SWjhmaVovUFExanpaWXV4WDNTZndVdG94Uk9uUTdmaFdacXhuRzhLYmgKT0RoWldlQnJZZ2JTK01xN0tpNzNHNmhCcWhNbU1xcFJoK2dUNUVZbnBHZmJUbVRlUUJveXBrTEtkZm1PUytsWQpDM0lFaXJ1LzRhTXhpc0FDNzMvY3VWenQ1T2k5S3BaVVNhWTJkK1pINWtyWWtEMWhMVDFFRTJxbDhuTS9YZ2NlCk81bDkxeVMxZXg5Z2pzZ0lwTnc5elF0U05qMW80SUFraWRnWDhCak5MdmJSYThWdzBibDgxZXp2UEFoSE9tM0wKazkxTGV3S0JnUUR3MkYrUUZxK2xWcEtzL3pFNTlyUjNCdHBBTFVIVEFqWW1rbTVNUHRLeWtYVUQyQ2lNV3lYVgpxNGNaRnBaV2tUY1YzKzN2Q2FQcnNWdFBlcFZsaTZvM3VyMmR2VVFyQ2IxQUh0d2E4R21Bb2VseUJFSzZSeUdPCjh0dW16aTZqZDEyM0xQQTgvQ1JzL01ncmtjM2QxWTVSSDhzS3F1OGZCdlRncjNId2JTYkRPd0tCZ1FEU3EyRGQKZWxmOWFtZ0dyL3FneHJpclpUaWVKWHVJeUVIUG53cUsvNU1qa01jMUh3cVZMZ1NKUi81Y1dYL0U0UXR1Zko5UQp1b2kzNmExdGlweGhETlNqc214d0gxNi9pQ2hRc2V3Ym9BRkw4anFxdi9yQW12RUtmVmRyelU1V3c3dTY5dndECkdET2QxVnhxYzVFTVJVMTVRbS8wS3NMcjNRQTA3QUdsV3MrV1RRS0JnUUNna0QrMVpIVVk1R0hJcWlEUlR5TDYKQkpHQUhpdmlBWGx2REVtSUtlSE1ZUGl0T3NXeTUyNk9jbG9RdWJtUEMzb09sVWxLQm0xSHk1STAvR09aQjhiUAowZE8rZmkxeUJuNHF2ejFuU3ZFQWEvRDJyWUhSaE1FSjl6MFozc25DVDVQUGMwclRVc0x6VXpBZEM5NlQ0N1EyCkZUSUhISjlxQktOQWRQYWdTTVNDSndLQmdRQ01ja0NUcHVWVG5FVlp5cWNWMmRCeVEzc1VHUDBGMnRhS2FnY0gKVlF6Q1lJVE9GTEc0eVZoYmhWUTlrUC9ZbGdlZS9iK0xKd1hQOFI1SkJ4THJYZU1YekVoS3Y1Ui9zQnJnZHAxbgp3QURkVENjY2YveUFaMHRjVTQrZE00bVBMU2ZSK2YxWFpYZStqajJjVFF0ZGc0ajdnWi80SElIeWsxVjRuVlp1ClBJcGtNd0tCZ0FrZ0ljYlh0VHlDbGxJWnFWVTJseUUrcng3Wm5RNGs5ZTBaSENhaXFBT3JDNVBseUtwQ0hRZEgKWGYxcVhtRTBPeFhBeVEycU8wbWJWMzgza2ovU3E3b0p0RHRveS9Bc3ZiRG1vNjZjTzJRSXRSRmZPazJ6Q3UycQpDMFpraE9nazNGWUo4aXloV3pPV3VDWExKck9QVVVqekFkSDJTYWloVy9KY0hVUmhNTnBTCi0tLS0tRU5EIFJTQSBQUklWQVRFIEtFWS0tLS0tCg==""")
-ATTESTCA_CERT = base64.b64decode("""MIIEIDCCAwigAwIBAgIBBTANBgkqhkiG9w0BAQsFADCBsjELMAkGA1UEBhMCVVMxEzARBgNVBAgTCkNhbGlmb3JuaWExEjAQBgNVBAcTCVNhbiBEaWVnbzEwMC4GA1UECxMnR2VuZXJhbCBVc2UgVGVzdCBLZXkgKGZvciB0ZXN0aW5nIG9ubHkpMRowGAYDVQQLExFDRE1BIFRlY2hub2xvZ2llczERMA8GA1UEChMIUVVBTENPTU0xGTAXBgNVBAMTEFFQU0EgU0hBMjU2IFJvb3QwHhcNMTMwNDIyMjIwNjE4WhcNMzMwNDE3MjIwNjE4WjB8MQswCQYDVQQGEwJVUzELMAkGA1UECBMCQ0ExEjAQBgNVBAcTCVNhbiBEaWVnbzEaMBgGA1UECxMRQ0RNQSBUZWNobm9sb2dpZXMxETAPBgNVBAoTCFFVQUxDT01NMR0wGwYDVQQDExRRUFNBIE9QRU5EU1AgVEVTVCBDQTCCASAwDQYJKoZIhvcNAQEBBQADggENADCCAQgCggEBAMYyug2JlapW/3apqD5nEJehM1ZcCsB0agQYJYNHQcM79ixkif6mcopwFX2IYQHfuqZ5nWzKNUt4sHObvmCkG0DR41NRZHd8bKKcae2uhur05vtuUFbRhRig8u76HxHSedeuWT0gZoCmqmj6UdRUhgbQoRMKPHUwGL9Gcql8Yn8efM+/59zjwJdj7yXHo7MEb8w6Kkwhi3tuS6NW20+Q3PJV2Mmn69lKY1rFQIYhgEmfoBxuvjKl9U7iKN6iC3Da7TEF8/nC432Sb93S4+nRCVbbLPzsi7g+pmd9b08N9Yx1j3SsduMcLrpJycvRJDyD85zROhaVcb4BgHMYRyXCSr8CAQOjeDB2MB8GA1UdIwQYMBaAFElk8+VAE1VZc2dnWT99Qreru/tXMB0GA1UdDgQWBBQnxAfEeRhBNAnuLkunmI4I+aSyOzAPBgNVHRMECDAGAQH/AgEAMAsGA1UdDwQEAwIBBjAWBgorBgEEAYspCQYDBAgAAeJAAAn78TANBgkqhkiG9w0BAQsFAAOCAQEAYHPAAlh+ezXdqUDIptraYfoiVxw2YsX++Ytg2eJ69YFVlCo33bLJFwQMj+zTMauRgvLew2cZTK47ghVV7130M13E53aN49p/DTOe3u5OFGA+z+ZLrqhraUPT+UhaAuVO9Yu9eOLudsPvgJTeD1a7RaC6PmPsUFPxLUlmlJn3lSXjlYe98+hittLnJ9gTnjdTVH/PgEJhMvUcjjyBWdRsog54VpyqesqLJedC4OF7fHJZ4S7rxDAINI15aDBQrOW/LD6HsBdr4WikS5Lnmecaw+2Um/ge/3Jl/kFBgh8EyORmSzaN4q1OoPYykxTGxenP3Z6D9WJurPd0d0fnuf+bNw==""")
-ROOTCA_CERT = base64.b64decode("""MIIEGzCCAwOgAwIBAgIBATANBgkqhkiG9w0BAQsFADCBsjELMAkGA1UEBhMCVVMxEzARBgNVBAgTCkNhbGlmb3JuaWExEjAQBgNVBAcTCVNhbiBEaWVnbzEwMC4GA1UECxMnR2VuZXJhbCBVc2UgVGVzdCBLZXkgKGZvciB0ZXN0aW5nIG9ubHkpMRowGAYDVQQLExFDRE1BIFRlY2hub2xvZ2llczERMA8GA1UEChMIUVVBTENPTU0xGTAXBgNVBAMTEFFQU0EgU0hBMjU2IFJvb3QwHhcNMTMwMzI4MjMxOTA4WhcNMzMwMzIzMjMxOTA4WjCBsjELMAkGA1UEBhMCVVMxEzARBgNVBAgTCkNhbGlmb3JuaWExEjAQBgNVBAcTCVNhbiBEaWVnbzEwMC4GA1UECxMnR2VuZXJhbCBVc2UgVGVzdCBLZXkgKGZvciB0ZXN0aW5nIG9ubHkpMRowGAYDVQQLExFDRE1BIFRlY2hub2xvZ2llczERMA8GA1UEChMIUVVBTENPTU0xGTAXBgNVBAMTEFFQU0EgU0hBMjU2IFJvb3QwggEgMA0GCSqGSIb3DQEBAQUAA4IBDQAwggEIAoIBAQC3mmlyc5XmZ4nQcUf8gXoHX82fCU12SW6VJdlz5IyKOJzl+IeYs2ArpkDHXaF2NwYvS4cJVBHtvx5TLbsBMAF9teFORqSs6wl+r+3nQwCogNOn/8JZrcPdxkjA8cVAkydxSK0jPxAdAGf8vGXD7tKDWWZyHquPoWqNVG/P4OyHAWMKCYg/w7/5MYTOcV1TXW2BraH7dztGkS4ey2hsOPlJzxP74cN1WyXjLPkn5CZWkx+95CKN5i+T9S+FeKD/1zbxuNlwv4x3x1Ohw9KBJYQzrB/wP9wrsVEnh2K9jy7rapKFFWOgQj8omg1EbIMqdOHuSZYcexFvAqN233xxluDBAgEDozwwOjAdBgNVHQ4EFgQUSWTz5UATVVlzZ2dZP31Ct6u7+1cwDAYDVR0TBAUwAwEB/zALBgNVHQ8EBAMCAQYwDQYJKoZIhvcNAQELBQADggEBAJ31kC2nKcTK1XrArhpkmnAX8zCPZkl+Azm7qF/Mr09h8FJiXJ7hBBoLHD+/+DifgUkLF4EjLOUnTTUPLPPKS5KgvuDkgJKvAMvv/GXxuGabdl4EebYCbJvnEgwkyG4pYVd5pGHQ0z2Md4nI6aMYco2X44bIjdqxFJwWOgPhioef1KbO/6CzykG0cPvpJB0XkWR8QGztFt9HofF+uVgpY2t1dL4/SuM/wJSeS8rdqstk0BYq/EDFFs99r1wP2R6hjJgCPkMvo7WiSE2yRrEkMNhgyEKrrD4pr7YWtsFPkYeTMXOvYoY16aOfvcrw0FfH+yATdn/OdAQ2saJISmilWh4=""")
-
-PT_NULL = 0
-PF_OS_PHDR = 0x07000000
-PF_OS_HASH = 0x02200000
-MBN_V3 = 3
-SIG_SIZE = 256
-CC_SIZE = 0x1800
-PAD_FF = b'\xff'
-IPAD = 0x3636363636363636
-OPAD = 0x5C5C5C5C5C5C5C5C
-
-
-def _pad(data, size, pad=PAD_FF):
-  if len(data) > size:
-    raise ValueError("data too large: %d > %d" % (len(data), size))
-  return data + pad * (size - len(data))
-
-
-def _orig_segment_data(ph):
-  off, size = ph['o'], ph['fs']
-  for base_off, data in ORIG_SEGS.items():
-    rel = off - base_off
-    if 0 <= rel and rel + size <= len(data):
-      return data[rel:rel + size]
-  raise KeyError("missing original segment at 0x%x" % off)
-
-
-def _build_ehdr(base_ehdr, num_phdrs):
-  e = bytearray(base_ehdr)
-  struct.pack_into('<I', e, 0x20, 0)
-  struct.pack_into('<H', e, 0x2c, num_phdrs)
-  struct.pack_into('<H', e, 0x30, 0)
-  struct.pack_into('<H', e, 0x32, 0)
-  return bytes(e)
-
-
-def _build_elf(ehdr, phs, segs):
-  phoff = struct.unpack_from('<I', ehdr, 0x1c)[0]
-  d = bytes(ehdr)
-  if len(d) < phoff:
-    d += b'\x00' * (phoff - len(d))
-  for ph in phs:
-    d += struct.pack('<IIIIIIII', ph['t'], ph['o'], ph['v'],
-             ph['p'], ph['fs'], ph['ms'], ph['fl'], ph['al'])
-  for off, sdata in sorted(segs.items()):
-    if len(d) < off:
-      d += b'\x00' * (off - len(d))
-    d = d[:off] + sdata + d[off + len(sdata):]
-  return d
-
-
-def _qti_hmac(data, msm=0, sw=0):
-  def _u(v):
-    return bytes.fromhex(format(v, 'x').zfill(16))
-  Si, So = _u(sw ^ IPAD), _u(msm ^ OPAD)
-  a = hashlib.sha256(data).digest()
-  b = hashlib.sha256(Si + a).digest()
-  c = hashlib.sha256(So + b).digest()
-  return c
-
-
-def _raw_pkcs1_sign(private_key, data):
-  """Raw RSA-PKCS1-v1_5 signing WITHOUT DigestInfo wrapper.
-  This matches OpenSSL:  pkeyutl -sign -pkeyopt rsa_padding_mode:pkcs1"""
-  numbers = private_key.private_numbers()
-  d = numbers.d
-  n = numbers.public_numbers.n
-  key_len = (n.bit_length() + 7) // 8
-  pad_len = key_len - 3 - len(data)
-  if pad_len < 8:
-    raise ValueError("data too long for key size")
-  em = b'\x00\x01' + b'\xff' * pad_len + b'\x00' + data
-  m_int = int.from_bytes(em, 'big')
-  sig_int = pow(m_int, d, n)
-  return sig_int.to_bytes(key_len, 'big')
-
-
-def _new_cert(ca_key, ca_cert, attrs):
-  k = rsa.generate_private_key(public_exponent=3, key_size=2048)
-  def _na(oid, value, typ):
-    return x509.NameAttribute(oid, value, _type=typ)
-
-  # Match the SecTools/OpenSSL attestation cert profile accepted by DSP loaders.
-  n = [
-    _na(NameOID.COUNTRY_NAME, "US", _ASN1Type.PrintableString),
-    _na(NameOID.COMMON_NAME, "SecTools Test User", _ASN1Type.PrintableString),
-    _na(NameOID.LOCALITY_NAME, "San Diego", _ASN1Type.PrintableString),
-    _na(NameOID.ORGANIZATION_NAME, "SecTools", _ASN1Type.PrintableString),
-    _na(NameOID.STATE_OR_PROVINCE_NAME, "California", _ASN1Type.PrintableString),
-    _na(NameOID.ORGANIZATIONAL_UNIT_NAME, "01 %.16X SW_ID" % attrs['sw'], _ASN1Type.T61String),
-    _na(NameOID.ORGANIZATIONAL_UNIT_NAME, "02 %.16X HW_ID" % attrs['hw'], _ASN1Type.T61String),
-    _na(NameOID.ORGANIZATIONAL_UNIT_NAME, "04 %.4X OEM_ID" % attrs['oid'], _ASN1Type.T61String),
-    _na(NameOID.ORGANIZATIONAL_UNIT_NAME, "05 %.8X SW_SIZE" % attrs['sz'], _ASN1Type.T61String),
-    _na(NameOID.ORGANIZATIONAL_UNIT_NAME, "06 %.4X MODEL_ID" % attrs['mid'], _ASN1Type.T61String),
-    _na(NameOID.ORGANIZATIONAL_UNIT_NAME, "07 0001 %s" % attrs['ha'], _ASN1Type.PrintableString),
-    _na(NameOID.ORGANIZATIONAL_UNIT_NAME, "03 %.16X DEBUG" % attrs['dbg'], _ASN1Type.PrintableString),
-  ]
-
-  nvb = datetime.now(timezone.utc).replace(microsecond=0)
-  nva = nvb + timedelta(days=20 * 365)
-  b = (x509.CertificateBuilder()
-     .subject_name(x509.Name(n))
-     .issuer_name(ca_cert.subject)
-     .public_key(k.public_key())
-     .serial_number(1)
-     .not_valid_before(nvb)
-     .not_valid_after(nva))
-  b = b.add_extension(x509.AuthorityKeyIdentifier.from_issuer_public_key(ca_key.public_key()), False)
-  b = b.add_extension(x509.UnrecognizedExtension(ExtensionOID.BASIC_CONSTRAINTS, b'\x30\x03\x02\x01\x00'), False)
-  b = b.add_extension(x509.KeyUsage(
-    digital_signature=True,
-    content_commitment=False,
-    key_encipherment=False,
-    data_encipherment=False,
-    key_agreement=False,
-    key_cert_sign=False,
-    crl_sign=False,
-    encipher_only=False,
-    decipher_only=False,
-  ), False)
-  b = b.add_extension(x509.UnrecognizedExtension(ObjectIdentifier("1.3.6.1.4.1.1449.9.6.3"), b'\x00\x01\xe2\x40\x00\x01\xe2\x40'), False)
-  return k, b.sign(ca_key, hashes.SHA256())
-
-
-def _build_chain(attest, ca, root):
-  c = (attest.public_bytes(serialization.Encoding.DER) +
-     ca.public_bytes(serialization.Encoding.DER) +
-     root.public_bytes(serialization.Encoding.DER))
-  return _pad(c, CC_SIZE, b'\xff')
-
-
-def _sign(serial_num, out_dir):
-  orig = ORIG_PHDRS
-  base_ehdr = ORIG_EHDR
-
-  nph = len(orig) + 2
-  phsz = nph * 32
-  ehsz = len(base_ehdr)
-  phoff = struct.unpack_from('<I', base_ehdr, 0x1c)[0]
-  hv = 0x5000
-  ho = 0x1000
-  hfs = 0x1a08
-
-  prog = {'t': PT_NULL, 'o': 0, 'v': 0, 'p': 0,
-      'fs': ehsz + phsz, 'ms': 0, 'fl': PF_OS_PHDR, 'al': 0}
-  hph = {'t': PT_NULL, 'o': ho, 'v': hv, 'p': hv,
-       'fs': hfs, 'ms': 0x2000, 'fl': PF_OS_HASH, 'al': 0x1000}
-  shifted = [dict(ph, o=ph['o'] + 0x3000) for ph in orig]
-  allph = [prog, hph] + shifted
-
-  # Temporary ELF for hash 0
-  tmp_segs = {}
-  for i, ph in enumerate(shifted):
-    tmp_segs[ph['o']] = _orig_segment_data(orig[i])
-  ehdr = _build_ehdr(base_ehdr, nph)
-  tmp = _build_elf(ehdr, allph, tmp_segs)
-
-  phb = tmp[phoff:phoff + phsz]
-  hash0 = hashlib.sha256(ehdr + phb).digest()
-  hash1 = struct.pack('<I', serial_num) + b'\x00' * 28
-  hs = [hash0, hash1]
-  for ph in orig:
-    hs.append(hashlib.sha256(_orig_segment_data(ph)).digest())
-  ht = b''.join(hs)
-
-  cs, ss, ccs = len(ht), SIG_SIZE, CC_SIZE
-  dst = hv + 40
-  sp = dst + cs
-  cp = sp + ss
-  isz = cs + ss + ccs
-
-  mbn = struct.pack('<IIIIIIIIII',
-            0, MBN_V3, 0, dst, isz, cs, sp, ss, cp, ccs)
-  dts = mbn + ht
-  hmac = _qti_hmac(dts)
-
-  ca_key = serialization.load_pem_private_key(ATTESTCA_KEY, password=None)
-  ca_cert = x509.load_der_x509_certificate(ATTESTCA_CERT)
-  root_cert = x509.load_der_x509_certificate(ROOTCA_CERT)
-
-  attrs = {'sw': 0, 'hw': 0, 'oid': 0, 'mid': 0,
-       'sz': len(dts), 'ha': 'SHA256', 'dbg': 2}
-  nk, ac = _new_cert(ca_key, ca_cert, attrs)
-  sig = _raw_pkcs1_sign(nk, hmac)
-  sig = _pad(sig, SIG_SIZE, b'\x00')
-  cc = _build_chain(ac, ca_cert, root_cert)
-
-  hseg = mbn + ht + sig + cc
-  if len(hseg) != hfs:
-    raise RuntimeError("hash seg size mismatch %d vs %d" % (len(hseg), hfs))
-
-  segs = {ho: hseg}
-  for i, ph in enumerate(shifted):
-    segs[ph['o']] = _orig_segment_data(orig[i])
-
-  final = _build_elf(ehdr, allph, segs)
-
-  os.makedirs(out_dir, exist_ok=True)
-  out = os.path.join(out_dir, "testsig-0x%08X.so" % serial_num)
-  with open(out, 'wb') as f:
-    f.write(final)
-  print("Signing complete! Output saved at %s" % out)
-  return out
-
-
-def main():
-  p = argparse.ArgumentParser(description="Qualcomm testsig generator")
-  p.add_argument("-t", "--testsig", action="append", default=None,
-           help="serial number (e.g. 0x67489311); repeatable. Defaults to value from /sys/devices/soc0/serial_number")
-  p.add_argument("-o", "--output_dir", default=".")
-  args = p.parse_args()
-
-  serials = args.testsig if args.testsig else []
-
-  if not serials:
-    # Read default serial number from device
-    try:
-      with open('/sys/devices/soc0/serial_number', 'r') as f:
-        serial_str = f.read().strip()
-      serials = [serial_str]
-      print("Using serial number from /sys/devices/soc0/serial_number: %s" % serial_str)
-    except FileNotFoundError:
-      raise SystemExit("Error: No serial number provided (-t) and /sys/devices/soc0/serial_number not found.")
-    except PermissionError:
-      raise SystemExit("Error: Cannot read /sys/devices/soc0/serial_number (permission denied).")
-
-  for s in serials:
-    v = int(s.strip(), 0)
-    if not (0 <= v <= 0xFFFFFFFF):
-      raise ValueError("bad serial %r" % s)
-    _sign(v, args.output_dir)
-
-
-if __name__ == "__main__":
-  main()
@@ -376,7 +376,7 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
    if isinstance(xq.device, tuple) and not isinstance(attn_mask.device, tuple):
      attn_mask = attn_mask.shard(xq.device, axis=0)
  else:
-    attn_mask = Tensor.zeros((B, 1, N, N), device=single_device, dtype=dtypes.float32)
+    attn_mask = Tensor.zeros((B, 1, N, N), requires_grad=False, device=single_device, dtype=dtypes.float32)
    if isinstance(xq.device, tuple):
      attn_mask = attn_mask.shard(xq.device, axis=0)

@@ -564,8 +564,8 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
  "aten.__rshift__.Scalar": lambda x,y: x>>y,
  "aten.__irshift__.Scalar": lambda x,y: x>>y,
  # inplace ops using replace for fusion
-  "aten.zero_": lambda x: x.const_like(0),
-  "aten.fill_.Scalar": lambda x, y: x.const_like(y),
+  "aten.zero_": lambda x: x.zeros_like(),
+  "aten.fill_.Scalar": lambda x, y: x.full_like(y),
  "aten.add_.Tensor": lambda self, other, alpha=1.0: self + other * alpha,
  "aten.add_.Scalar": lambda self, other, alpha=1.0: self + other * alpha,
  "aten.mul_.Tensor": lambda self, other: self * other,
@@ -617,7 +617,7 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
  "aten.asinh": Tensor.asinh,
  "aten.mul": Tensor.mul,
  "aten.atanh": Tensor.atanh,
-  "aten.fill_.Tensor": lambda self, value: self.const_like(value.reshape(()).item()),
+  "aten.fill_.Tensor": lambda self, value: Tensor.full(self.shape, value.reshape(()).item(), device=self.device, dtype=self.dtype),
  "aten.flip": Tensor.flip,
  "aten.scatter_reduce.two": Tensor.scatter_reduce,
  "aten.squeeze_.dim": Tensor.squeeze,
@@ -709,7 +709,7 @@ def wrap_inplace_view_op(k,f):
      views = derived_views(base)
      if views:
        old_base = Tensor(base.uop, device=base.device)
-        old_base.is_param = base.is_param
+        old_base.requires_grad = base.requires_grad
        old_base._views = getattr(base, "_views", set())
        for v in views: v._view_base = old_base
        base._views = set()
@@ -240,9 +240,9 @@ class TestTorchBackend(unittest.TestCase):
    np.testing.assert_equal(result.cpu().numpy(), [3., 3., 2.])

  def test_mnist_index(self):
-    # from tinygrad.nn.datasets import mnist
-    X_train, Y_train = Tensor.randint(60000, 1, 28, 28, dtype='uchar').realize(), Tensor.randint(60000, dtype='uchar').realize()
    GlobalCounters.reset()
+    from tinygrad.nn.datasets import mnist
+    X_train, Y_train, _, _ = mnist()
    X_train = torch.tensor(X_train.float().numpy(), device=device)
    Y_train = torch.tensor(Y_train.cast('int64').numpy(), device=device)
    samples = torch.randint(0, X_train.shape[0], (32,))
@@ -105,7 +105,7 @@ class TestKernelFusionRegression(unittest.TestCase):
      view = x[1:3]
      view += 1.0
      return x.sum()
-    self._check_kernel_count(fn, 7)
+    self._check_kernel_count(fn, 8)

  def test_batchnorm_running_stats_update(self):
    def fn():
@@ -1,6 +1,6 @@
 import numpy as np
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import trange
+from tinygrad.helpers import CI, trange
 from tinygrad.engine.jit import TinyJit


@@ -24,9 +24,9 @@ def train(model, X_train, Y_train, optim, steps, BS=128, lossfn=lambda out,y: ou

  with Tensor.train():
    losses, accuracies = [], []
-    for i in (t := trange(steps, disable=None)):
+    for i in (t := trange(steps, disable=CI)):
      samp = np.random.randint(0, X_train.shape[0], size=(BS))
-      x = Tensor(transform(X_train[samp]))
+      x = Tensor(transform(X_train[samp]), requires_grad=False)
      y = Tensor(target_transform(Y_train[samp]))
      loss, accuracy = train_step(x, y)
      # printing
@@ -43,7 +43,7 @@ def evaluate(model, X_test, Y_test, num_classes=None, BS=128, return_predict=Fal
  Tensor.training = False
  def numpy_eval(Y_test, num_classes):
    Y_test_preds_out = np.zeros(list(Y_test.shape)+[num_classes])
-    for i in trange((len(Y_test)-1)//BS+1, disable=None):
+    for i in trange((len(Y_test)-1)//BS+1, disable=CI):
      x = Tensor(transform(X_test[i*BS:(i+1)*BS]))
      out = model.forward(x) if hasattr(model, 'forward') else model(x)
      Y_test_preds_out[i*BS:(i+1)*BS] = out.numpy()
@@ -7,15 +7,15 @@ def get_node(graph:dict, key): return graph[str(key)]

 if __name__ == "__main__":
  parser = argparse.ArgumentParser(description="print CALL graph from DEBUG=5 tinygrad.viz.cli --json output")
-  parser.add_argument("kernel", type=str, nargs="?", default="ALL", metavar="NAME", help="Kernel name to stop at (default: print all kernels)")
+  parser.add_argument("kernel", type=str, default=None, help="Kernel name to stop at (default: print all kernels)")
  args = parser.parse_args()
  ref:int|None = None
  for line in sys.stdin:
    if not line.strip(): continue
    graph = json.loads(line)
-    if graph.get("ref") is not None and (args.kernel == "ALL" or graph["ref"] == ref):
+    if ref is not None and graph.get("ref") == ref:
      print(graph)
-      if (v:=json.loads(next(sys.stdin, "{}")).get("value")): print(v)
+      if (v:=json.loads(next(sys.stdin)).get("value")): print(v)
    if ref is not None or not isinstance(rec:=next(iter(graph.values()), {}), dict) or "label" not in rec: continue
    for v in graph.values():
      if not v["label"].startswith("CALL"): continue
@@ -39,6 +39,6 @@ if __name__ == "__main__":
          src_str = ["SRC"]+get_node(graph, get_node(graph, s)["src"][0][1])["label"].splitlines()[1:]
          print(" ".join(idx_str+src_str))
        ss += [x[1] for x in get_node(graph, s)["src"]]
-      if args.kernel != "ALL" and args.kernel in ansistrip(v["label"]):
+      if args.kernel is not None and args.kernel in ansistrip(v["label"]):
        ref = v["ref"]
        break
@@ -1,6 +1,6 @@
 [project]
 name = "tinygrad"
-version = "0.13.0"
+version = "0.12.0"
 description = "You like pytorch? You like micrograd? You love tinygrad! <3"
 authors = [{ name = "George Hotz" }]

@@ -28,7 +28,6 @@ packages = [
  'tinygrad.nn',
  'tinygrad.renderer',
  'tinygrad.renderer.amd',
-  'tinygrad.renderer.isa',
  'tinygrad.runtime',
  'tinygrad.runtime.autogen',
  'tinygrad.runtime.autogen.am',
@@ -36,7 +35,6 @@ packages = [
  'tinygrad.runtime.autogen.amd.rdna3',
  'tinygrad.runtime.autogen.amd.rdna4',
  'tinygrad.runtime.autogen.amd.cdna',
-  'tinygrad.runtime.autogen.nv_regs',
  'tinygrad.runtime.graph',
  'tinygrad.runtime.support',
  'tinygrad.runtime.support.am',
@@ -49,7 +47,6 @@ packages = [

 [tool.setuptools.package-data]
 tinygrad = ["py.typed"]
-"tinygrad.llm" = ["chat.html"]
 "tinygrad.viz" = ["index.html", "assets/**/*", "js/*"]


@@ -136,7 +133,7 @@ debug = true

 [tool.mypy]
 warn_unused_configs = true
-files = ["tinygrad", "test/mockgpu"]
+files = ["tinygrad"]
 ignore_missing_imports = true
 check_untyped_defs = true
 explicit_package_bases = true
@@ -16,7 +16,6 @@
 \definecolor{elwyellow}{HTML}{F9A825}
 \definecolor{callblue}{HTML}{1565C0}
 \definecolor{assignbrown}{HTML}{795548}
-\definecolor{loadred}{HTML}{c08080}
 \definecolor{multipurple}{HTML}{7B1FA2}
 \definecolor{markerorange}{HTML}{E65100}
 % AxisType colors (from tinygrad)
@@ -49,14 +48,16 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op
 \toprule
 \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
 \midrule
-\op{Param} & $(\mathbf{s})$ & slot, dtype, device?, addrspace? &
-  Placeholder with shape $\mathbf{s}$. Substituted in \op{Function}. \\[4pt]
 \op{Buffer} & () & size, dtype, device, addrspace &
  Shape $(n \cdot \textit{size},)$ if device is $n$-tuple, else $(\textit{size},)$. \\
+\op{BufferView} & (buf,) & size, dtype, offset &
+  Typed access into a buffer. Zero-copy $(\textit{size},)$ slice at offset; inherits addrspace. \\
+\op{Param} & $(\mathbf{s})$ or $(\mathbf{s}, \text{min}, \text{max})$ & slot, dtype, device? &
+  Placeholder with shape $\mathbf{s}$. Substituted in \op{Function}. \\[4pt]
 \op{Const} & () & value, dtype &
  A scalar constant with shape $(\ )$. \\
-  & & & Form vector consts with \op{Stack} \\
-\op{Binary} & () & data & Raw binary data, has dtype uint8 and shape len($data$)  \\
+\op{Vconst} & () & values, dtype &
+  A vector constant with shape $(n,)$. \\
 \bottomrule
 \end{tabular}

@@ -64,7 +65,7 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op
 A \op{Buffer}'s \textbf{addrspace} is \texttt{GLOBAL}, \texttt{LOCAL}, or \texttt{REG}.

 %% ============================================================
-\subsection*{{\color{movgreen}Movement Ops} \normalfont\small--- no arithmetic; view, indexing, and reinterpretation only}
+\subsection*{{\color{movgreen}Movement Ops} \normalfont\small--- no arithmetic, shapes are $(k,)$-shaped UOps with dtype \texttt{index} in src}

 \begin{tabular}{@{}l l l l@{}}
 \toprule
@@ -74,13 +75,11 @@ A \op{Buffer}'s \textbf{addrspace} is \texttt{GLOBAL}, \texttt{LOCAL}, or \textt
 \op{Flip} & $(T,)$ & bools $\mathbf{f}$ & Reverse along flagged axes. \\
 \op{Reshape} & $(T, \mathbf{s'})$ & --- & Reinterpret in row-major order. $\prod s_k = \prod s'_k$. \\
 \op{Expand} & $(T, \mathbf{s'})$ & --- & Broadcast size-1 axes. $s_k \in \{1, s'_k\}$. \\
-\op{Pad} & $(T, \mathbf{o}, \mathbf{s'})$ & --- & Place $T$ at offset $o_k$ in a zero-filled output of shape $s'_k$. \\
-\op{Shrink} & $(T, \mathbf{o}, \mathbf{s'})$ & --- & Keep $s'_k$ elements starting at offset $o_k$ per axis. Inverse of \op{Pad}. \\
+\op{Pad} & $(T, \mathbf{b}, \mathbf{e})$ & --- & Pad with $0$s: $b_k$ before, $e_k$ after each axis. \\
+\op{Shrink} & $(T, \mathbf{b}, \mathbf{e})$ & --- & Keep $[b_k, e_k)$ per axis. Inverse of \op{Pad}. \\
 \op{Index} & $(T, i_0, i_1, \ldots)$ & --- & Index from left. $()$-shaped $i$ removes dim; $(k,)$-shaped makes it $k$. \\
 \op{Stack} & $(T_0, T_1, \ldots)$ & --- & Join along a newly created leading axis. All shapes must match. \\
 \op{Replicated} & $(T,)$ & axes & Mark $T$ as replicated along axes. Collapse axes to $1$. \\
-\op{Slice} & $(T, \mathrm{offset})$ & size, dtype & Zero-copy \textit{size} elems of dtype; offset is elems of $T$ dtype. \\
-\op{Bitcast} & $(T,)$ & dtype & Reinterpret storage as target dtype; preserve total bytes. \\
 \bottomrule
 \end{tabular}

@@ -91,7 +90,7 @@ A \op{Buffer}'s \textbf{addrspace} is \texttt{GLOBAL}, \texttt{LOCAL}, or \textt
 \toprule
 \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
 \midrule
-\op{Reduce} & ($T$, $r_0$, $r_1$, \ldots) & op, axes & Reduce $T$ along axes or ranges. Op is \op{Add}, \op{Max}, or \op{Mul}. \\
+\op{Reduce} & $(T,)$ & op, axes & Reduce $T$ along axes. Op is \op{Add}, \op{Max}, or \op{Mul}. \\
 \bottomrule
 \end{tabular}

@@ -110,25 +109,13 @@ A \op{Buffer}'s \textbf{addrspace} is \texttt{GLOBAL}, \texttt{LOCAL}, or \textt
 \end{tabular}

 %% ============================================================
-\subsection*{{\color{loadred}Load Ops} \normalfont\small--- can change device or addrspace}
+\subsection*{{\color{multipurple}Store Ops} \normalfont\small--- side effects}

 \begin{tabular}{@{}l l l l@{}}
 \toprule
 \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
 \midrule
-\op{Load} & (buf, alt?, gate?) & device, addrspace & Read (pull) from buffer into a new anonymous buffer. \\
-  & & & Note: this replaces \op{Copy} and \op{Contiguous}. \\
-\bottomrule
-\end{tabular}
-
-%% ============================================================
-\subsection*{{\color{multipurple}Store Ops} \normalfont\small--- the only op with observable side effects}
-
-\begin{tabular}{@{}l l l l@{}}
-\toprule
-\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
-\midrule
-\op{Store} & (buf, val, gate?) & --- & Write (push) val into buf. buf.shape $=$ val.shape. \\
+\op{Store} & (buf, val, gate?) & --- & Write val into buf. buf.shape $=$ val.shape. \\
  & & & If gate is present, write only when gate is true. Output is void. \\
 \bottomrule
 \end{tabular}
@@ -166,7 +153,8 @@ Unary & $(T,)$
  & $\mathrm{trunc}(x)$: round toward zero. \\
 & & \op{Cast}
  & Convert to target dtype (specified in arg). \\
-  \\[4pt]
+& & \op{Bitcast}
+  & Reinterpret bits as target dtype. Must be same size. \\[4pt]
 Binary & $(A, B)$
  & \op{Add}, \op{Mul}, \op{Max}, \op{Mod}, \op{Idiv}
  & $a+b$, $a \cdot b$, $\max(a,b)$, $a \bmod b$, $\lfloor a/b \rfloor$ \\
@@ -217,6 +205,7 @@ Ternary & $(P, A, B)$
 \op{Contiguous} & $(T,)$ & --- & Force contiguous memory layout. \\
 \op{ContiguousBackward} & $(T,)$ & --- & Force contiguous in backward pass. \\
 \op{Detach} & $(T,)$ & --- & Stops gradient propagation. \\
+\op{Copy} & $(T,)$ & device & Copy to target device. \\
 \bottomrule
 \end{tabular}

@@ -227,6 +216,8 @@ Ternary & $(P, A, B)$
 \toprule
 \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
 \midrule
+\op{Load} & (idx,alt?,gate?) & --- & Dereference: read element at index from buffer. \\
+ & & & All loads will be replaced by \op{Store}. \\
 \op{Barrier} & (deps\ldots) & --- & Synchronize threads within a workgroup. \\
 \op{Ins} & \ldots & \ldots & A single machine instruction (e.g.\ AMD ISA). \\
 \op{Special} & (bound,) & name & GPU thread/workgroup index (e.g.\ \texttt{gidx0}, \texttt{lidx1}). \\
@@ -1,20 +0,0 @@
-#!/usr/bin/env python3
-import unittest
-import tinygrad.runtime.autogen.amd.cdna.ins as cdna
-from test.amd.hw.test_cdna_vop3 import run_cdna
-
-class TestCDNASDWA(unittest.TestCase):
-  def test_v_add_co_u32_e32_writes_vcc(self):
-    out = run_cdna([
-      cdna.s_mov_b32(cdna.s[0], 0xffffffff),
-      cdna.v_mov_b32_e32(cdna.v[0], cdna.s[0]),
-      cdna.v_mov_b32_e32(cdna.v[13], 1),
-      cdna.v_add_co_u32_e32(cdna.v[0], cdna.SDWA, cdna.v[13], vsrc0=cdna.v[0], dst_sel=6, src0_sel=6),
-      cdna.v_mov_b32_e32(cdna.v[2], cdna.VCC_LO),
-      cdna.v_lshlrev_b32_e32(cdna.v[2], 31, cdna.v[2]),
-      cdna.v_or_b32_e32(cdna.v[2], cdna.v[2], cdna.v[0]),
-    ])
-    self.assertEqual(out, 0x80000000)
-
-if __name__ == "__main__":
-  unittest.main()
@@ -1,130 +0,0 @@
-"""CDNA VOP3 instruction coverage.
-
-Exercises generated CDNA pcode end-to-end in the emulator and compares against
-gfx950 hardware when USE_HW=1.
-"""
-import ctypes, struct, unittest
-import tinygrad.runtime.autogen.amd.cdna.ins as cdna
-from tinygrad.helpers import flat_mv
-from tinygrad.renderer.amd.dsl import NULL
-from test.amd.hw.helpers import USE_HW, assemble
-from test.mockgpu.amd.emu import run_asm
-
-LANES = 1
-
-def _code(instructions: list, out_reg: int = 2, out_addr: int | None = None) -> bytes:
-  load_out_addr = [
-    cdna.s_mov_b32(cdna.s[92], out_addr & 0xffffffff),
-    cdna.s_mov_b32(cdna.s[93], out_addr >> 32),
-  ] if out_addr is not None else [
-    cdna.s_load_dwordx2(cdna.s[92:93], cdna.s[80:81], 0, soffset=NULL),
-    cdna.s_waitcnt(0),
-  ]
-  return assemble([
-    cdna.s_mov_b32(cdna.s[80], cdna.s[0]),
-    cdna.s_mov_b32(cdna.s[81], cdna.s[1]),
-    cdna.v_mov_b32_e32(cdna.v[255], cdna.v[0]),
-    *instructions,
-    *load_out_addr,
-    cdna.v_lshlrev_b32_e32(cdna.v[240], 2, cdna.v[255]),
-    cdna.global_store_dword(addr=cdna.v[240], data=cdna.v[out_reg], saddr=cdna.s[92:93], offset=0),
-    cdna.s_endpgm(),
-  ])
-
-def _run_emu(instructions: list, out_reg: int = 2) -> int:
-  out_buf = (ctypes.c_uint32 * LANES)(*([0] * LANES))
-  args = (ctypes.c_uint64 * 1)(ctypes.addressof(out_buf))
-  code = _code(instructions, out_reg)
-  kernel_buf = (ctypes.c_char * len(code)).from_buffer_copy(code)
-  result = run_asm(ctypes.addressof(kernel_buf), len(code), 1, 1, 1, LANES, 1, 1, ctypes.addressof(args),
-                   0x19c | (128 << 15), 0x10000, arch="cdna")
-  assert result == 0, f"run_asm failed with {result}"
-  return out_buf[0]
-
-def _run_hw(instructions: list, out_reg: int = 2) -> int:
-  from tinygrad.device import Device
-  from tinygrad.runtime.ops_amd import AMDProgram
-  from tinygrad.runtime.support.compiler_amd import HIPCompiler
-
-  dev = Device["AMD"]
-  if dev.arch != "gfx950": raise unittest.SkipTest("requires gfx950 hardware")
-  out_gpu = dev.allocator.alloc(LANES * 4)
-  code = _code(instructions, out_reg, out_gpu.va_addr)
-  byte_str = ", ".join(f"0x{b:02x}" for b in code)
-  asm_src = f""".text
-.globl test
-.p2align 8
-.type test,@function
-test:
-.byte {byte_str}
-
-.rodata
-.p2align 6
-.amdhsa_kernel test
-  .amdhsa_next_free_vgpr 256
-  .amdhsa_next_free_sgpr 96
-  .amdhsa_accum_offset 256
-  .amdhsa_kernarg_size 0
-.end_amdhsa_kernel
-
-.amdgpu_metadata
---
-amdhsa.version:
-  - 1
-  - 0
-amdhsa.kernels:
-  - .name: test
-    .symbol: test.kd
-    .kernarg_segment_size: 0
-    .group_segment_fixed_size: 0
-    .private_segment_fixed_size: 0
-    .kernarg_segment_align: 8
-    .wavefront_size: 64
-    .sgpr_count: 96
-    .vgpr_count: 256
-    .max_flat_workgroup_size: 1024
-...
-.end_amdgpu_metadata
-"""
-  prg = AMDProgram(dev, "test", HIPCompiler(dev.arch).compile(asm_src))
-  prg(global_size=(1, 1, 1), local_size=(LANES, 1, 1), wait=True)
-  out = bytearray(LANES * 4)
-  dev.allocator._copyout(flat_mv(memoryview(out)), out_gpu)
-  return struct.unpack("<I", out)[0]
-
-def run_cdna(instructions: list, out_reg: int = 2) -> int:
-  emu = _run_emu(instructions, out_reg)
-  if not USE_HW: return emu
-  hw = _run_hw(instructions, out_reg)
-  if emu != hw: raise AssertionError(f"Emulator vs Hardware mismatch: emu=0x{emu:08x} hw=0x{hw:08x}")
-  return hw
-
-class TestCDNAVOP3(unittest.TestCase):
-  def test_cvt_pk_fp8_f32_preserves_upper_half(self):
-    """V_CVT_PK_FP8_F32 with OPSEL[3]=0 writes only D[15:0]."""
-    out = run_cdna([
-      cdna.s_mov_b32(cdna.s[0], 0xdeadbeef),
-      cdna.v_mov_b32_e32(cdna.v[2], cdna.s[0]),
-      cdna.v_mov_b32_e32(cdna.v[0], 1.0),
-      cdna.v_mov_b32_e32(cdna.v[1], 2.0),
-      cdna.v_cvt_pk_fp8_f32(cdna.v[2], cdna.v[0], cdna.v[1]),
-    ])
-    self.assertEqual(out, 0xdead4038)
-
-  def test_cvt_pk_bf8_f32_overflow_and_inf(self):
-    """V_CVT_PK_BF8_F32 converts finite overflow and infinities to E5M2 infinities."""
-    for name, bits, expected in [
-      ("finite_overflow", 0x47700000, 0x7c),
-      ("pos_inf", 0x7f800000, 0x7c),
-      ("neg_inf", 0xff800000, 0xfc),
-    ]:
-      with self.subTest(name=name):
-        out = run_cdna([
-          cdna.s_mov_b32(cdna.s[0], 0xdeadbeef),
-          cdna.v_mov_b32_e32(cdna.v[2], cdna.s[0]),
-          cdna.s_mov_b32(cdna.s[0], bits),
-          cdna.v_mov_b32_e32(cdna.v[0], cdna.s[0]),
-          cdna.v_mov_b32_e32(cdna.v[1], 1.0),
-          cdna.v_cvt_pk_bf8_f32(cdna.v[2], cdna.v[0], cdna.v[1]),
-        ])
-        self.assertEqual(out, 0xdead3c00 | expected)
@@ -167,7 +167,7 @@ class TestDSPcodePatterns(unittest.TestCase):

  def test_global_atomic_add_f32_parsing(self):
    """Test GLOBAL_ATOMIC_ADD_F32 keeps memory values in float dtype."""
-    vmem = UOp.param(2, dtypes.uint32.ptr(1024))
+    vmem = UOp(Ops.PARAM, dtypes.uint32.ptr(1024), arg=2)
    srcs = {
      'ADDR': UOp.const(dtypes.uint64, 0),
      'DATA': UOp.const(dtypes.uint32, 0x3f800000),
@@ -198,7 +198,7 @@ class TestDSPcodePatterns(unittest.TestCase):
  def test_mem_read_parsing(self):
    """Test MEM[addr].type read expression parsing."""
    # Create a mock LDS buffer
-    lds = UOp.param(3, dtypes.uint32.ptr(16384))
+    lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
    addr = UOp.const(dtypes.uint32, 0)
    vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(dtypes.uint32, 0)}

@@ -233,7 +233,7 @@ class TestDSPcodePatterns(unittest.TestCase):
    pcode = PCODE.get(DSOp.DS_LOAD_2ADDR_B32)
    self.assertIsNotNone(pcode)
    assert pcode is not None
-    lds = UOp.param(3, dtypes.uint32.ptr(16384))
+    lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
    srcs = {
      'ADDR': UOp.const(dtypes.uint32, 0),
      'OFFSET0': UOp.const(dtypes.uint32, 0),
@@ -314,7 +314,7 @@ class TestConcatWidthParsing(unittest.TestCase):
      self.assertEqual(parsed.simplify().arg, expected)

  def test_permlane64_wave64_pcode_indices(self):
-    vgpr = UOp.param(0, dtypes.uint32.ptr(256))
+    vgpr = UOp(Ops.PARAM, dtypes.uint32.ptr(256), arg=0)
    srcs = {
      'SRC0': UOp.const(dtypes.uint32, 0),
      'VDST': UOp.const(dtypes.uint32, 1),
@@ -347,7 +347,7 @@ class TestAllPcode(unittest.TestCase):
  def _make_srcs(self):
    """Create dummy source variables for pcode parsing."""
    u32, u64 = lambda v=0: UOp.const(dtypes.uint32, v), lambda v=0: UOp.const(dtypes.uint64, v)
-    lds = UOp.param(3, dtypes.uint32.ptr(16384))
+    lds = UOp(Ops.PARAM, dtypes.uint32.ptr(16384), arg=3)
    return {'laneId': u32(), 'laneID': u32(), 'S0': u32(), 'S1': u32(), 'S2': u32(), 'S3': u32(), 'SRC0': u32(),
            'D0': u32(), 'D1': u32(), 'DST': u32(), 'VDST': u32(), 'SDST': u32(),
            'VCC': u64(), 'VCCZ': u32(), 'EXEC': u64(), 'EXEC_LO': u32(), 'EXECZ': u32(), 'SCC': u32(),
@@ -19,7 +19,7 @@ class TestArange(unittest.TestCase):
    self.assertEqual(self._get_flops(Tensor.arange(256), np.arange(256)), 0)
    self.assertEqual(self._get_flops(Tensor.arange(2560), np.arange(2560)), 0)

-  @unittest.skipIf(Device.DEFAULT == "CL", "flaky in CI")
+  @unittest.skipIf(Device.DEFAULT == "CL", "TODO: fails on CI CL")
  def test_arange_cumsum(self):
    np.testing.assert_equal(Tensor.arange(513).cumsum(0).numpy(), np.arange(513).cumsum())

@@ -64,7 +64,7 @@ class TestIndexing(unittest.TestCase):
      rng = Tensor.arange(DSET, dtype=dtypes.int).reshape(1, 1, DSET, 1).expand(4, DDIM, DSET, 1)
      idxs = idxs.reshape(4,1,1,1).expand(4, DDIM, DSET, 1)
      reshape_dataset = dataset.T.reshape(1, DDIM, DSET, 1).expand(4, DDIM, DSET, 1)
-      full = (rng==idxs).where(reshape_dataset, Tensor.zeros(4, DDIM, DSET, 1, buffer=False))
+      full = (rng==idxs).where(reshape_dataset, Tensor.zeros(4, DDIM, DSET, 1))
      X = full.sum(axis=(2,3))
      linear, var_vals = X.linear_with_vars()
      self.assertEqual(len(linear.src), 1)
@@ -125,8 +125,8 @@ class TestIndexing(unittest.TestCase):
  def test_index_mnist(self, noopt=1, op_limit=512*784*13, split_reduceop=0):
    # WEBGPU generates more ops due to bitpacking of < 4-byte dtypes
    if Device.DEFAULT == "WEBGPU": op_limit *= 15
-    # from tinygrad.nn.datasets import mnist
-    X_train, Y_train = Tensor.randint(DSET, 1, 28, 28, dtype='uchar').realize(), Tensor.randint(DSET, dtype='uchar').realize()
+    from tinygrad.nn.datasets import mnist
+    X_train, Y_train, _, _ = mnist()
    with Context(NOOPT=noopt, SPLIT_REDUCEOP=split_reduceop):
      samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0]).realize()
      GlobalCounters.reset()
@@ -172,7 +172,7 @@ class TestIndexing(unittest.TestCase):
    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)
+    emb.weight = Tensor.ones(vocab_size, embed_size, requires_grad=True)
    gt = Tensor.zeros(bs, seqlen, embed_size)
    Tensor.realize(idx, emb.weight, gt)
    GlobalCounters.reset()
@@ -198,14 +198,14 @@ class TestIndexing(unittest.TestCase):
    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)
+    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)
+    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)
@@ -215,12 +215,12 @@ class TestIndexing(unittest.TestCase):
    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 DEV.arch.startswith("gfx")), "tests AMD bf16 cast overhead")
-  def base_test_llama_8b_rope_backward(self, dtype, ops_scale=1):
+  def base_test_llama_8b_rope_backward(self, dtype):
    from extra.models.llama import precompute_freqs_cis, apply_rotary_emb
    bs, seqlen, dim, n_heads = 1, 512, 256, 4
    head_dim = dim // n_heads
    x = Tensor.randn(bs, seqlen, dim, dtype=dtype)
-    wq = Tensor.randn(dim, dim, dtype=dtype)
+    wq = Tensor.randn(dim, dim, dtype=dtype, requires_grad=True)
    freqs_cis = precompute_freqs_cis(head_dim, seqlen).cast(dtype)
    Tensor.realize(x, wq, freqs_cis)
    xq = (x @ wq.T)
@@ -232,15 +232,15 @@ class TestIndexing(unittest.TestCase):
    linear = compile_linear(wq.grad.schedule_linear())
    assert len(linear.src) == 1, f"expected one kernel for backward, got: {len(linear.src)}"
    bwd_ops = estimate_uop(linear.src[0]).ops
+    # bfloat16 on non CDNA4 has ~10x ops overhead because of the software emulation
+    if dtype == dtypes.bfloat16 and not Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950"): ops_scale = 10
+    else: ops_scale = 1
    expected_ops = bs*seqlen*dim*dim*ops_scale
    print(f"rope matmul bwd ({dtype}): {GlobalCounters.kernel_count} kernels, {bwd_ops:,} ops")
    self.assertLess(bwd_ops, expected_ops, f"rope bwd ops {bwd_ops:,} should be < {ops_scale} per (got {bwd_ops/(bs*seqlen*dim*dim):.1f})")

-  def test_llama_8b_rope_backward_f16(self):
-    self.base_test_llama_8b_rope_backward(dtypes.float16, ops_scale=2)
-  # bfloat16 on non CDNA4 has ~10x ops overhead because of the software emulation
-  def test_llama_8b_rope_backward_bf16(self):
-    self.base_test_llama_8b_rope_backward(dtypes.bfloat16, ops_scale=2 if Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950") else 25)
+  def test_llama_8b_rope_backward_f16(self): self.base_test_llama_8b_rope_backward(dtypes.float16)
+  def test_llama_8b_rope_backward_bf16(self): self.base_test_llama_8b_rope_backward(dtypes.bfloat16)

 if __name__ == "__main__":
  unittest.main()
@@ -1,9 +1,10 @@
 import unittest
 from tinygrad import Tensor, Device, dtypes, Context
+from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import getenv, system, DEV
 from extra.gemm.cdna_asm_gemm import asm_gemm
 from test.helpers import needs_second_gpu
-from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8, FP8_MAX
+from examples.mlperf.models.flat_llama import FP8_DTYPE

 # On non CDNA4 it will only validate the Tensor.custom_kernel integration
 # Use DEV=NULL:HIP:gfx950 to also test the assembly
@@ -11,48 +12,33 @@ def is_cdna4(): return Device[Device.DEFAULT].renderer.target.arch.startswith("g

 def run_asm_gemm(a_shape, b_shape, dtype=dtypes.float16, a_shard=None, b_shard=None, gpus:int=1) -> None:
  Tensor.manual_seed(0)
-  input_dtype = dtypes.bfloat16 if dtype == FP8_DTYPE else dtype
-  a_rand = Tensor.randn(a_shape, dtype=dtypes.float).sub(0.5).cast(input_dtype)
-  b_rand = Tensor.randn(b_shape, dtype=dtypes.float).sub(0.5).cast(input_dtype)
+  a_rand = Tensor.randn(a_shape, dtype=dtypes.float).sub(0.5).cast(dtype)
+  b_rand = Tensor.randn(b_shape, dtype=dtypes.float).sub(0.5).cast(dtype)
  with Context(DEBUG=0):
    Tensor.realize(a_rand, b_rand)

  devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(gpus)) if (multi:=gpus>1) else None

-  if dtype == FP8_DTYPE:
-    a_rand, x_scale, _ = quantize_fp8(a_rand)
-    b_rand, w_scale, _ = quantize_fp8(b_rand)
-    grad_amax_state = Tensor.full((), FP8_MAX, dtype=dtypes.float32, device=devs).contiguous()
-    with Context(DEBUG=0):
-      Tensor.realize(a_rand, x_scale, b_rand, w_scale, grad_amax_state)
-
-  # clone all inputs before any backward: a clone copies the source's current .grad
-  a, b = a_rand.clone(), b_rand.clone()
-  if dtype == FP8_DTYPE:
-    a_ref, b_ref = a_rand.detach().cast(dtypes.bfloat16), b_rand.detach().cast(dtypes.bfloat16)
-  else:
-    a_ref, b_ref = a_rand.clone(), b_rand.clone()
+  a, b = a_rand.clone().requires_grad_(), b_rand.clone().requires_grad_()
  if multi: a, b = a.shard(devs, axis=a_shard), b.shard(devs, axis=b_shard)
-  if dtype == FP8_DTYPE:
-    tst = asm_gemm(a, b, x_scale=x_scale, w_scale=w_scale, grad_amax_state=grad_amax_state)
-  else:
-    tst = asm_gemm(a, b)
+  tst = asm_gemm(a, b)
  tst.sum().backward()
  Tensor.realize(tst, a.grad, b.grad)

+  a_ref, b_ref = a_rand.clone().requires_grad_(), b_rand.clone().requires_grad_()
+  # do reference gemm in bf16 for fp8, adjusting atol for quantization effects
+  if a_ref.dtype == FP8_DTYPE:
+    a_ref = a_ref.cast(dtypes.bfloat16)
+    b_ref = b_ref.cast(dtypes.bfloat16)
  if multi: a_ref, b_ref = a_ref.shard(devs, axis=a_shard), b_ref.shard(devs, axis=b_shard)
-  if dtype == FP8_DTYPE:
-    ref = ((a_ref @ b_ref) * x_scale * w_scale).cast(dtypes.bfloat16)
-  else:
-    ref = a_ref @ b_ref
+  ref = a_ref @ b_ref
  ref.sum().backward()
  Tensor.realize(ref, a_ref.grad, b_ref.grad)

  # no validation on the NULL device
  if a_rand.device.startswith("NULL"): return None
  atol, rtol = (2e-1, 1e-2) if dtype == dtypes.bfloat16 else (256, 1e-2) if dtype == FP8_DTYPE else (1e-2, 1e-3)
-  # allow more rtol for multi because of ALLREDUCE_CAST
-  grad_atol, grad_rtol = (16895, 0.125) if dtype == FP8_DTYPE else (atol, 2e-2 if multi else rtol)
+  grad_atol, grad_rtol = (16895, 0.125) if dtype == FP8_DTYPE else (atol, rtol)
  with Context(DEBUG=0):
    # enable for debugging, slow for larger gemms
    if getenv("USE_NPY"):
@@ -84,7 +70,7 @@ def verify_asm_gemm_k_sharded_3d(batch:int, M:int, N:int, K:int, dtype=dtypes.fl

 # 128x smaller than usual
 # uses the UOp GEMM, runs on non CDNA4 and CI
-@unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half")
+@unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
 class TestGemm(unittest.TestCase):
  def setUp(self):
    if is_cdna4(): self.skipTest("shapes are too small for the assembly GEMM")
@@ -151,19 +137,13 @@ class TestGemmLlama(unittest.TestCase):
  def test_empty(self): asm_gemm(Tensor.empty(N:=getenv("N", 4096), N, dtype=self.dtype), Tensor.empty(N, N, dtype=self.dtype)).realize()

  def test_empty_bw(self):
-    x = Tensor.empty(1, N:=getenv("N", 4096), N, dtype=self.dtype)
-    y = Tensor.empty((N, N), dtype=self.dtype)
-    if self.dtype == FP8_DTYPE:
-      x_scale = Tensor.empty((), dtype=dtypes.float32)
-      w_scale = Tensor.empty((), dtype=dtypes.float32)
-      grad_amax_state = Tensor.empty((), dtype=dtypes.float32).contiguous()
-      z = asm_gemm(x, y, x_scale=x_scale, w_scale=w_scale, grad_amax_state=grad_amax_state)
-    else:
-      z = asm_gemm(x, y)
+    x = Tensor.empty(1, N:=getenv("N", 4096), N, dtype=self.dtype, requires_grad=True)
+    y = Tensor.empty((N, N), dtype=self.dtype, requires_grad=True)
+    z = asm_gemm(x, y)
    z.sum().backward()
    Tensor.realize(z, x.grad, y.grad)
-    # FP8 GEMM stores bf16 output and its backward produces bf16 gradients.
-    grad_dtype = dtypes.bfloat16 if self.dtype == FP8_DTYPE else self.dtype
+    # FP8 forward output is bf16, gradients use fp8e5m2 (aka bf8)
+    grad_dtype = dtypes.fp8e5m2 if self.dtype == FP8_DTYPE else self.dtype
    assert z.dtype == dtypes.bfloat16
    assert x.grad.dtype == y.grad.dtype == grad_dtype

@@ -2,6 +2,7 @@ import unittest, math
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.dtype import DTYPES_DICT
 from tinygrad.uop.ops import Ops, UOp
+from tinygrad.device import is_dtype_supported
 import numpy as np
 from test.helpers import not_support_multi_device

@@ -14,11 +15,6 @@ def _check_ast_count(desired_count:int, t:Tensor):
  #assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"

 class TestMovedConstFolding(unittest.TestCase):
-  def test_contiguous_deviceless_const(self):
-    t = Tensor(UOp.const(dtypes.float, 2.0)).contiguous()
-    self.assertIs(t.uop.op, Ops.CONST)
-    self.assertIsNone(t.uop.device)
-
  def test_add_shrunk_zero(self):
    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + Tensor.zeros(6).shrink(((1, 5),)))

@@ -32,16 +28,18 @@ class TestMovedConstFolding(unittest.TestCase):
    _check_ast_count(1, Tensor([1.0, 2, 3, 4]) * Tensor.ones(2).pad(((1, 1),)))

  def test_copy_padded_const(self):
-    schedule = Tensor.ones(4, device="CPU:0", buffer=False).pad(((1, 1),)).to("CPU:1").schedule_linear()
+    schedule = Tensor.ones(4, device="CPU:0").pad(((1, 1),)).to("CPU:1").schedule_linear()
    assert not any(si.src[0].op is Ops.COPY for si in schedule.src), "const copy should be folded"
-    np.testing.assert_equal(Tensor.ones(4, device="CPU:0", buffer=False).pad(((1, 1),)).to("CPU:1").numpy(), [0, 1, 1, 1, 1, 0])
+    np.testing.assert_equal(Tensor.ones(4, device="CPU:0").pad(((1, 1),)).to("CPU:1").numpy(), [0, 1, 1, 1, 1, 0])

  def test_cast_padded(self):
    # NOTE: it's always 1 kernel when calling .numpy, limitation of _check_ast_count
-    _check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16))
-    np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16).numpy(), [0, 1, 1, 1, 1, 0])
-    _check_ast_count(1, Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16))
-    np.testing.assert_equal(Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16).numpy(), [0, 65535, 65535, 65535, 65535, 0])
+    if is_dtype_supported(dtypes.int16):
+      _check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16))
+      np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16).numpy(), [0, 1, 1, 1, 1, 0])
+    if is_dtype_supported(dtypes.uint16):
+      _check_ast_count(1, Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16))
+      np.testing.assert_equal(Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16).numpy(), [0, 65535, 65535, 65535, 65535, 0])
    # folded
    _check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int64))
    np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int64).numpy(), [0, 1, 1, 1, 1, 0])
@@ -111,7 +109,7 @@ class TestReduceOpsConstFolding(unittest.TestCase):
  def test_sum_output_dtype(self):
    # sum output dtype can be different from input
    for dt in DTYPES_DICT.values():
-      if dt in Device[Device.DEFAULT].renderer.supported_dtypes():
+      if is_dtype_supported(dt):
        t = Tensor.ones(16, dtype=dt).reshape(4, 4)
        assert t.sum().dtype == t.contiguous().sum().dtype

@@ -119,7 +117,7 @@ class TestReduceOpsConstFolding(unittest.TestCase):
 class TestMultiConstFolding(unittest.TestCase):
  def test_multi_const_folding_literal(self):
    ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
-    t = Tensor.arange(16).float().clone().to(ds).realize()
+    t = Tensor.arange(16).float().to(ds).realize()

    # non const folding case creates one ast on each shard
    _check_ast_count(4, t + 1)
@@ -144,7 +142,7 @@ class TestMultiConstFolding(unittest.TestCase):

  def test_multi_const_folding_tensor(self):
    ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
-    t = Tensor.arange(16).float().clone().to(ds).realize()
+    t = Tensor.arange(16).float().to(ds).realize()
    zero = Tensor.zeros(16).to(ds).realize()
    one = Tensor.ones(16).to(ds).realize()

@@ -1,5 +1,5 @@
 import unittest
-from tinygrad import Tensor, UOp, GlobalCounters, Context, Device
+from tinygrad import Tensor, UOp, GlobalCounters, Context
 from tinygrad.dtype import AddrSpace, dtypes
 from tinygrad.uop.ops import KernelInfo, AxisType, Ops

@@ -155,7 +155,7 @@ class TestCustomKernel(unittest.TestCase):
    self.assertTrue((ref == tst).all().item())

  def test_eye(self):
-    ref = Tensor.eye(1024).clone().realize()
+    ref = Tensor.eye(1024).contiguous().realize()
    tst = Tensor.empty_like(ref)
    tst = tst.custom_kernel(fxn=custom_eye_kernel)[0]
    self.assertTrue((ref == tst).all().item())
@@ -220,14 +220,14 @@ class TestCustomKernel(unittest.TestCase):
    b_rand = Tensor.randn(8, N)
    Tensor.realize(a_rand, b_rand)

-    a, b = Tensor(a_rand.numpy()), Tensor(b_rand.numpy())
+    a, b = Tensor(a_rand.numpy(), requires_grad=True), Tensor(b_rand.numpy(), requires_grad=True)
    c = Tensor.empty(N, N)
    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm, grad_fxn=backward_gemm_custom if custom_backward_gemm else backward_gemm)[0]
    tst.sum().backward()
    grad_a, grad_b = a.grad, b.grad
    Tensor.realize(tst, grad_a, grad_b)

-    a, b = Tensor(a_rand.numpy()), Tensor(b_rand.numpy())
+    a, b = Tensor(a_rand.numpy(), requires_grad=True), Tensor(b_rand.numpy(), requires_grad=True)
    ref = (a@b)
    ref.sum().backward()
    real_grad_a, real_grad_b = a.grad, b.grad
@@ -284,30 +284,6 @@ class TestCustomKernel(unittest.TestCase):
    self.assertIsNotNone(custom_idx, "custom_addmul kernel not found in schedule")
    self.assertEqual(custom_idx, 3, f"custom_addmul should be at index 3, got {custom_idx}")

-  def test_invalids_into_custom_kernel_no_empty_kernel(self):
-    from tinygrad.engine.realize import compile_linear
-    a = Tensor.full((4, 4), 3.).contiguous()
-    b = Tensor.full((4, 4), 2.).contiguous()
-    Tensor.realize(a, b)
-    out = Tensor.invalids(*a.shape, dtype=a.dtype)
-    out, *_ = Tensor.custom_kernel(out, a, b, fxn=custom_elementwise_add_kernel)
-    compiled = compile_linear(out.schedule_linear())
-    for call in compiled.src:
-      prg = call.src[0]
-      if prg.op is not Ops.PROGRAM: continue
-      self.assertTrue(len(prg.arg.globals) > 0, f"empty kernel compiled (no globals): name={prg.arg.name}")
-
-  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "kernel timing not supported")
-  def test_invalids_into_custom_kernel_with_beam(self):
-    a = Tensor.full((4, 4), 3.).contiguous()
-    b = Tensor.full((4, 4), 2.).contiguous()
-    Tensor.realize(a, b)
-    with Context(BEAM=1, IGNORE_BEAM_CACHE=1):
-      out = Tensor.invalids(*a.shape, dtype=a.dtype)
-      out, *_ = Tensor.custom_kernel(out, a, b, fxn=custom_elementwise_add_kernel)
-      result = out.flatten().tolist()
-    self.assertTrue(all(x == 5 for x in result), f"expected all 5.0, got {result}")
-
  @unittest.skip("what are anonymous buffers?")
  def test_anonymous_buffers_in_function(self):
    """Test that custom kernels with anonymous output buffers work inside @function."""
@@ -335,7 +311,7 @@ class TestCustomKernel(unittest.TestCase):
    assert all(x == expected for x in result), f"expected all {expected}, got {result}"

  def test_custom_kernel_sched(self, use_custom=False):
-    x = Tensor.arange(32).reshape(8, 4).clone().realize()
+    x = Tensor.arange(32).reshape(8, 4).realize()
    y = Tensor.empty_like(x)
    y = Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
    if use_custom:
@@ -352,7 +328,7 @@ class TestCustomKernel(unittest.TestCase):

  @unittest.expectedFailure
  def test_sliced_buffer_function(self):
-    x = Tensor.arange(32).reshape(8, 4).clone().realize()
+    x = Tensor.arange(32).reshape(8, 4).realize()
    from tinygrad import function
    @function(precompile=True)
    def run(x:Tensor) -> Tensor:
@@ -2,7 +2,8 @@ import contextlib, unittest, math
 import numpy as np
 import torch
 from typing import Any, List
-from tinygrad.helpers import getenv, DEBUG, EMULATED_DTYPES, DEV
+from tinygrad.device import is_dtype_supported
+from tinygrad.helpers import getenv, DEBUG, CI, EMULATED_DTYPES
 from tinygrad.dtype import DType, DTYPES_DICT, least_upper_dtype, fp8_to_float, float_to_fp8, _to_np_dtype, _to_torch_dtype, truncate
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
@@ -16,13 +17,11 @@ pytestmark = pytest.mark.filterwarnings("ignore")
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")

-supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
-
 def get_available_cast_dtypes(dtype: DType) -> List[DType]:
-  dts = [v for k, v in DTYPES_DICT.items() if v != dtype and v in supported_dtypes or v in dtypes.fp8s+(dtypes.half,dtypes.bfloat16,dtypes.long)]
-  if dtype in (dtypes.long, dtypes.ulong) and (dtype not in supported_dtypes or dtypes.long in EMULATED_DTYPES.tolist(dtypes)):
+  dts = [v for k, v in DTYPES_DICT.items() if v != dtype and is_dtype_supported(v) or v in dtypes.fp8s+(dtypes.half,dtypes.bfloat16,dtypes.long)]
+  if dtype in (dtypes.long, dtypes.ulong) and (not is_dtype_supported(dtype) or dtypes.long in EMULATED_DTYPES.tolist(dtypes)):
    return [dt for dt in dts if dt != dtypes.double] # can't bitcast with no 64-bit support
-  if dtype not in supported_dtypes and dtype not in dtypes.fp8s+(dtypes.half,dtypes.bfloat16): return []
+  if not is_dtype_supported(dtype) and dtype not in dtypes.fp8s+(dtypes.half,dtypes.bfloat16): return []
  return dts

 def _to_torch_storage_type(dtype:DType):
@@ -61,7 +60,7 @@ class TestDType(unittest.TestCase):
  @classmethod
  def setUpClass(cls):
    if cls.DTYPE is None: raise unittest.SkipTest("base class")
-    cls.DATA = rand_for_dtype(cls.DTYPE, 0x10, allow_subnormal=cls.DTYPE in supported_dtypes)
+    cls.DATA = rand_for_dtype(cls.DTYPE, 0x10, allow_subnormal=is_dtype_supported(cls.DTYPE))

  def test_to_np(self):
    _test_to_np(Tensor(self.DATA, dtype=self.DTYPE), _to_np_dtype(self.DTYPE), np.array(self.DATA, dtype=_to_np_dtype(self.DTYPE)))
@@ -103,6 +102,31 @@ class TestDType(unittest.TestCase):
    _test_to_np(Tensor(v, dtype=self.DTYPE)+2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))+2)
    _test_to_np(Tensor(v, dtype=self.DTYPE)*2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))*2)

+  def test_dtypes_DTYPES_DICT(self):
+    self.assertIn("float", DTYPES_DICT)
+    self.assertIn("float32", DTYPES_DICT)
+    self.assertEqual(len(DTYPES_DICT), 28)
+    self.assertTrue(all(isinstance(value, DType) for value in DTYPES_DICT.values()))
+    self.assertTrue(all(issubclass(_to_np_dtype(value), np.generic) for value in DTYPES_DICT.values() if _to_np_dtype(value) is not None))
+
+  def test_resulting_and_init_dtypes_match(self):
+    dtypes = list(map(np.dtype, ["bool", "uint8", "int8", "int16", "int32", "int64", "float32", "float64"]))
+    data = [1., 2., 0., 0.5, -1.5, 5.25]
+    for dt in dtypes:
+      arr = np.asarray(data).astype(dt)
+      tensor = Tensor(arr)
+      if not is_dtype_supported(tensor.dtype): continue
+      tin = tensor.numpy()
+      tor = torch.as_tensor(arr).detach().numpy()
+      assert dt == tin.dtype == tor.dtype, f"dtype mismatch: expected={dt} | tinygrad={tin.dtype} | torch={tor.dtype}"
+      np.testing.assert_allclose(tin, tor, atol=1e-6, rtol=1e-3)
+
+  def test_finfo(self):
+    if self.DTYPE not in [dtypes.float16, dtypes.float32, dtypes.float64]: return
+    info = np.finfo(_to_np_dtype(self.DTYPE))
+    self.assertEqual(info.bits, self.DTYPE.bitsize)
+    self.assertEqual((info.nexp, info.nmant), dtypes.finfo(self.DTYPE))
+
 def _test_ops(a_dtype:DType, b_dtype:DType, target_dtype=None):
  target_dtype = target_dtype or least_upper_dtype(a_dtype, b_dtype)
  if a_dtype == dtypes.bool or b_dtype == dtypes.bool: return
@@ -110,9 +134,13 @@ def _test_ops(a_dtype:DType, b_dtype:DType, target_dtype=None):
  _assert_eq((Tensor([1], dtype=a_dtype).cast(b_dtype)+Tensor([1], dtype=a_dtype).cast(b_dtype)).cast(a_dtype), a_dtype, [2])
  _assert_eq(Tensor([1,2,3,4], dtype=a_dtype)*Tensor([1,2,3,4], dtype=b_dtype), target_dtype, [1,4,9,16])
  _assert_eq(Tensor([[1,2],[3,4]], dtype=a_dtype)@Tensor.eye(2, dtype=b_dtype), target_dtype, [[1,2],[3,4]])
-  _assert_eq(Tensor([1,1,1,1], dtype=a_dtype)+Tensor.ones((4,4), dtype=b_dtype), target_dtype, 2*np.ones((4,4)))
-  _assert_eq(Tensor([1,1,1,1], dtype=a_dtype)+Tensor.ones((4,4), dtype=b_dtype).clone(), target_dtype, 2*np.ones((4,4)))
-  _assert_eq(Tensor.ones((4,4), dtype=b_dtype).clone(), b_dtype, np.ones((4,4)))
+  _assert_eq(Tensor([1,1,1,1], dtype=a_dtype)+Tensor.ones((4,4), dtype=b_dtype), target_dtype, 2*Tensor.ones(4,4).numpy())
+
+class TestFp8s(unittest.TestCase):
+  def test_fp8e4m3_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e4m3).dtype == dtypes.fp8e4m3
+  def test_fp8e5m2_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e5m2).dtype == dtypes.fp8e5m2
+  def test_fp8e4m3fnuz_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e4m3fnuz).dtype == dtypes.fp8e4m3fnuz
+  def test_fp8e5m2fnuz_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e5m2fnuz).dtype == dtypes.fp8e5m2fnuz

 class TestFp8sConversions(unittest.TestCase):
  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3_MAX, max_value=FP8E4M3_MAX))
@@ -163,6 +191,25 @@ class TestFp8sConversions(unittest.TestCase):
  def test_fp8e5m2fnuz_to_float(self, x):
    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2fnuz).float().item())

+class TestBFloat16(unittest.TestCase):
+  def test_bf16_creation_numpy(self):
+    data = [-1, 1, 2]
+    t = Tensor(data, dtype=dtypes.bfloat16)
+    assert t.dtype == dtypes.bfloat16
+    tnp = t.numpy()
+    assert tnp.dtype == np.float32
+    np.testing.assert_allclose(tnp, np.array(data))
+
+  def test_bf16_ones(self):
+    t = Tensor.ones(3, 5, dtype=dtypes.bfloat16)
+    assert t.dtype == dtypes.bfloat16
+    np.testing.assert_allclose(t.numpy(), np.ones((3, 5)))
+
+  def test_bf16_eye(self):
+    t = Tensor.eye(3, dtype=dtypes.bfloat16)
+    assert t.dtype == dtypes.bfloat16
+    np.testing.assert_allclose(t.numpy(), np.eye(3))
+
 class TestBFloat16DType(unittest.TestCase):
  def test_bf16_to_float(self):
    _test_cast(Tensor([100000], dtype=dtypes.bfloat16), dtypes.float32)
@@ -224,10 +271,10 @@ class TestFloatDType(TestDType):
    _test_op(lambda: Tensor([-0.9, -0.3, 1.2], dtype=dtypes.float32).cast(dtypes.uint32), dtypes.uint32,
             [0, 0, 1])

-@unittest.skipUnless(dtypes.double in supported_dtypes, f"no double on {Device.DEFAULT}")
+@unittest.skipUnless(is_dtype_supported(dtypes.double), f"no double on {Device.DEFAULT}")
 class TestDoubleDType(TestDType):
  DTYPE = dtypes.double
-  @unittest.skipIf((DEV.interface.startswith("MOCK") and Device.DEFAULT in {"CUDA", "NV"}) or \
+  @unittest.skipIf((CI and Device.DEFAULT in {"CUDA", "NV"}) or \
   isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "conversion not supported on CI CUDA, PTX, and NIR")  # TODO: why not?
  def test_float64_increased_precision(self):
    for func in [
@@ -283,6 +330,10 @@ class TestBitCast(unittest.TestCase):
      # should fail because 3 int8 is 3 bytes but float16 is two and 3 isn't a multiple of 2
      Tensor.empty((3,), dtype=dtypes.int8).bitcast(dtypes.float16)

+    with self.assertRaises(RuntimeError):
+      # should fail because backprop through bitcast is undefined
+      Tensor.empty((4,), dtype=dtypes.int8, requires_grad=True).bitcast(dtypes.float16)
+
  def test_bitcast_float_to_int32(self):
    a = Tensor([1.,2,3])
    b = a.bitcast(dtypes.int32)
@@ -371,6 +422,48 @@ class TestEmulatedFp8e5m2(TestFp8e5m2):
  @classmethod
  def tearDownClass(cls): cls.stack.close()

+class TestPtrDType(unittest.TestCase):
+  def test_vec_double(self):
+    dt1 = dtypes.float.vec(4).ptr().vec(4)
+    dt2 = dtypes.float.vec(4).ptr().vec(4)
+    self.assertEqual(dt1, dt2)
+    self.assertEqual(str(dt1), str(dt2))
+
+  def test_scalar(self):
+    dt = dtypes.float.vec(4).ptr().scalar()
+    self.assertEqual(dt.base, dtypes.float.vec(4))
+
+    dt = dtypes.float.vec(4).ptr().vec(4).scalar()
+    self.assertEqual(dt.base, dtypes.float.vec(4))
+
+    dt = dtypes.float.vec(4).scalar()
+    self.assertEqual(dt, dtypes.float)
+
+  def test_serialize(self):
+    dt = dtypes.float.vec(4).ptr().vec(4)
+    self.assertEqual(dt, eval(str(dt)))
+
+  def test_vec_ptr_sz(self):
+    dt = dtypes.float.ptr(1024).vec(4)
+    self.assertEqual(dt, eval(str(dt)))
+    self.assertEqual(str(dt), "dtypes.float.ptr(1024).vec(4)")
+
+  def test_vcount(self):
+    dt = dtypes.float.ptr().vec(4)
+    self.assertEqual(dt.vcount, 4)
+    self.assertEqual(dt.v, 4)
+    self.assertEqual(dt.count, 1)
+
+    dt = dtypes.float.vec(4).ptr()
+    self.assertEqual(dt.vcount, 1)
+    self.assertEqual(dt.v, 1)
+    self.assertEqual(dt.count, 4)
+
+    dt = dtypes.float.vec(4).ptr().vec(4)
+    self.assertEqual(dt.vcount, 4)
+    self.assertEqual(dt.v, 4)
+    self.assertEqual(dt.count, 4)
+
 class TestImplicitFunctionTypeChange(unittest.TestCase):
  def test_functions(self):
    result = []
@@ -389,7 +482,7 @@ class TestImplicitFunctionTypeChange(unittest.TestCase):
 class TestTensorMethod(unittest.TestCase):
  @given(strat.sampled_from(core_dtypes))
  def test_abs_diff(self, dt):
-    if dt == dtypes.bool or dt not in supported_dtypes: return
+    if dt == dtypes.bool or not is_dtype_supported(dt): return
    a, b = Tensor([2], dtype=dt), Tensor([1], dtype=dt)
    ret = (a - b).abs()
    np.testing.assert_allclose(ret.numpy(), np.abs(a.numpy()-b.numpy()))
@@ -397,11 +490,11 @@ class TestTensorMethod(unittest.TestCase):
 class TestDtypeUsage(unittest.TestCase):
  def test_max_w_alu(self):
    for d in dtypes.ints:
-      if d in supported_dtypes:
+      if is_dtype_supported(d):
        t = Tensor([[1, 2], [3, 4]], dtype=d)
        (t*t).max().item()

-@unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, f"no bfloat16 on {Device.DEFAULT}")
+@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
 class TestOpsBFloat16(unittest.TestCase):
  def test_cast(self):
    # TODO: helper_test_op breaks in unrelated part
@@ -1,12 +1,12 @@
 import unittest, operator, math
 from tinygrad import Context, Tensor, dtypes, Device
 from tinygrad.dtype import DType, truncate, fp8_to_float
-from tinygrad.helpers import EMULATED_DTYPES, DEV, getenv
+from tinygrad.helpers import CI, EMULATED_DTYPES, DEV, getenv
 from tinygrad.tensor import _to_np_dtype
+from tinygrad.device import is_dtype_supported
 from tinygrad.runtime.ops_python import from_storage_scalar
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
-from tinygrad.uop import Ops
 import numpy as np
 import pytest
 from hypothesis import assume, given, strategies as strat, settings
@@ -37,11 +37,6 @@ if ((DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"})
  unary_operations.remove((Tensor.sin, np.sin))
  unary_operations.remove((Tensor.cos, np.cos))

-# transcendental isn't accurate enough
-if Ops.SQRT not in Device[Device.DEFAULT].renderer.code_for_op: unary_operations.remove((Tensor.sqrt, np.sqrt))
-
-supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
-
 class ht:
  float64 = strat.floats(width=64, allow_subnormal=False)
  float32 = strat.floats(width=32, allow_subnormal=False)
@@ -72,7 +67,7 @@ def universal_test(a, b, dtype, op):
    numpy_value = truncate[dtype](op[1](ta.numpy(), tb.numpy()).item())
  else: tensor_value, numpy_value = (op[0](ta, tb)).numpy(), op[1](ta.numpy(), tb.numpy())
  if dtype in dtypes.floats:
-    if dtype not in supported_dtypes or dtype in EMULATED_DTYPES.tolist(dtypes): # denormals are zero
+    if not is_dtype_supported(dtype) or dtype in EMULATED_DTYPES.tolist(dtypes): # denormals are zero
      fe, fm = dtypes.finfo(dtype)
      atol, rtol = 2 ** (2 - (1 << (fe - 1))), 2 ** (-fm)
    else: atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1),
@@ -88,8 +83,7 @@ def universal_test_unary(a, dtype, op):
  if op[0] == Tensor.log and a <= 0: return
  if dtype in dtypes.fp8s:
    # denormals are zero
-    if (dtype in EMULATED_DTYPES.tolist(dtypes) or dtype not in supported_dtypes
-        and abs(ta.numpy().item()) < 0.015625): return
+    if dtype in EMULATED_DTYPES.tolist(dtypes) or not is_dtype_supported(dtype) and abs(ta.numpy().item()) < 0.015625: return
    tensor_value = fp8_to_float(op[0](ta.realize()).bitcast(dtypes.uint8).item(), dtype)
    numpy_value = truncate[dtype](v:=op[1](ta.numpy()).item())
    # cuda cast f32 inf to f8 MAX, amd cast it to nan(E4M3)/inf(E5M2)
@@ -111,7 +105,6 @@ def universal_test_cast(a, in_dtype, dtype):
 def universal_test_midcast(a, b, c, op1, op2, d1:DType, d2:DType):
  if not isinstance(op1, tuple): op1 = (op1, op1)
  if not isinstance(op2, tuple): op2 = (op2, op2)
-  if op1[0] == operator.mod and b == 0: return
  # lt and max with nan is undefined in tinygrad
  if op1[0] in (operator.lt, Tensor.maximum) and (math.isnan(a) or math.isnan(b)): return
  if op2[0] in (operator.lt, Tensor.maximum) and math.isnan(c): return
@@ -122,14 +115,14 @@ def universal_test_midcast(a, b, c, op1, op2, d1:DType, d2:DType):
  np.testing.assert_allclose(tensor_value, numpy_value, rtol=1e-6 if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) else 1e-7)

 class TestDTypeALU(unittest.TestCase):
-  @unittest.skipUnless(dtypes.float64 in supported_dtypes, f"no float64 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.float64), f"no float64 on {Device.DEFAULT}")
  @given(ht.float64, ht.float64, strat.sampled_from(binary_operations))
  def test_float64(self, a, b, op): universal_test(a, b, dtypes.float64, op)

  @given(ht.float32, ht.float32, strat.sampled_from(binary_operations))
  def test_float32(self, a, b, op): universal_test(a, b, dtypes.float32, op)

-  @unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.float16), f"no float16 on {Device.DEFAULT}")
  @given(ht.float16, ht.float16, strat.sampled_from(binary_operations))
  def test_float16(self, a, b, op): universal_test(a, b, dtypes.float16, op)

@@ -137,17 +130,17 @@ class TestDTypeALU(unittest.TestCase):
  @Context(EMULATED_DTYPES="half")
  def test_emulated_float16(self, a, b, op): universal_test(a, b, dtypes.float16, op)

-  @unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, f"no bfloat16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
  @given(ht.bfloat16, ht.bfloat16, strat.sampled_from(binary_operations))
  def test_bfloat16(self, a, b, op):
-    universal_test(from_storage_scalar(a, dtypes.bfloat16), from_storage_scalar(b, dtypes.bfloat16), dtypes.bfloat16, op)
+    universal_test(from_storage_scalar(a, dtypes.bfloat16), from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)

  @given(ht.bfloat16, ht.bfloat16, strat.sampled_from(binary_operations))
  @Context(EMULATED_DTYPES="bfloat16")
  def test_emulated_bfloat16(self, a, b, op):
-    universal_test(from_storage_scalar(a, dtypes.bfloat16), from_storage_scalar(b, dtypes.bfloat16), dtypes.bfloat16, op)
+    universal_test(from_storage_scalar(a, dtypes.bfloat16), from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)

-  @unittest.skipUnless(dtypes.fp8e4m3 in supported_dtypes, f"no fp8e4m3 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), f"no fp8e4m3 on {Device.DEFAULT}")
  @given(ht.fp8e4m3, ht.fp8e4m3, strat.sampled_from(binary_operations))
  def test_fp8e4m3(self, a, b, op):
    universal_test(from_storage_scalar(a, dtypes.fp8e4m3), from_storage_scalar(b, dtypes.fp8e4m3), dtypes.fp8e4m3, op)
@@ -157,7 +150,7 @@ class TestDTypeALU(unittest.TestCase):
  def test_emulated_fp8e4m3(self, a, b, op):
    universal_test(from_storage_scalar(a, dtypes.fp8e4m3), from_storage_scalar(b, dtypes.fp8e4m3), dtypes.fp8e4m3, op)

-  @unittest.skipUnless(dtypes.fp8e5m2 in supported_dtypes, f"no fp8e5m2 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2), f"no fp8e5m2 on {Device.DEFAULT}")
  @given(ht.fp8e5m2, ht.fp8e5m2, strat.sampled_from(binary_operations))
  def test_fp8e5m2(self, a, b, op):
    universal_test(from_storage_scalar(a, dtypes.fp8e5m2), from_storage_scalar(b, dtypes.fp8e5m2), dtypes.fp8e5m2, op)
@@ -167,12 +160,12 @@ class TestDTypeALU(unittest.TestCase):
  def test_emulated_fp8e5m2(self, a, b, op):
    universal_test(from_storage_scalar(a, dtypes.fp8e5m2), from_storage_scalar(b, dtypes.fp8e5m2), dtypes.fp8e5m2, op)

-  @unittest.skipUnless(dtypes.fp8e4m3fnuz in supported_dtypes, f"no fp8e4m3fnuz on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3fnuz), f"no fp8e4m3fnuz on {Device.DEFAULT}")
  @given(ht.fp8e4m3fnuz, ht.fp8e4m3fnuz, strat.sampled_from(binary_operations))
  def test_fp8e4m3fnuz(self, a, b, op):
    universal_test(from_storage_scalar(a, dtypes.fp8e4m3fnuz), from_storage_scalar(b, dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)

-  @unittest.skipUnless(dtypes.fp8e5m2fnuz in supported_dtypes, f"no fp8e5m2fnuz on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2fnuz), f"no fp8e5m2fnuz on {Device.DEFAULT}")
  @given(ht.fp8e5m2fnuz, ht.fp8e5m2fnuz, strat.sampled_from(binary_operations))
  def test_fp8e5m2fnuz(self, a, b, op):
    universal_test(from_storage_scalar(a, dtypes.fp8e5m2fnuz), from_storage_scalar(b, dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
@@ -190,7 +183,7 @@ class TestDTypeALU(unittest.TestCase):
  @given(ht.float32, strat.sampled_from(unary_operations))
  def test_float32_unary(self, a, op): universal_test_unary(a, dtypes.float32, op)

-  @unittest.skipUnless(dtypes.float16 in supported_dtypes, f"no float16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.float16), f"no float16 on {Device.DEFAULT}")
  @given(ht.float16, strat.sampled_from(unary_operations))
  def test_float16_unary(self, a, op): universal_test_unary(a, dtypes.float16, op)

@@ -198,7 +191,7 @@ class TestDTypeALU(unittest.TestCase):
  @Context(EMULATED_DTYPES="half")
  def test_emulated_float16_unary(self, a, op): universal_test_unary(a, dtypes.float16, op)

-  @unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, f"no bfloat16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), f"no bfloat16 on {Device.DEFAULT}")
  @given(ht.bfloat16, strat.sampled_from(unary_operations))
  def test_bfloat16_unary(self, a, op): universal_test_unary(from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)

@@ -206,7 +199,7 @@ class TestDTypeALU(unittest.TestCase):
  @Context(EMULATED_DTYPES="bfloat16")
  def test_emulated_bfloat16_unary(self, a, op): universal_test_unary(from_storage_scalar(a, dtypes.bfloat16), dtypes.bfloat16, op)

-  @unittest.skipUnless(dtypes.fp8e4m3 in supported_dtypes, f"no fp8e4m3 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), f"no fp8e4m3 on {Device.DEFAULT}")
  @given(ht.fp8e4m3, strat.sampled_from(unary_operations))
  def test_fp8e4m3_unary(self, a, op):
    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3) != 0.0)
@@ -218,7 +211,7 @@ class TestDTypeALU(unittest.TestCase):
    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3) != 0.0)
    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3), dtypes.fp8e4m3, op)

-  @unittest.skipUnless(dtypes.fp8e5m2 in supported_dtypes, f"no fp8e5m2 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2), f"no fp8e5m2 on {Device.DEFAULT}")
  @given(ht.fp8e5m2, strat.sampled_from(unary_operations))
  def test_fp8e5m2_unary(self, a, op):
    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2) != 0.0)
@@ -230,13 +223,13 @@ class TestDTypeALU(unittest.TestCase):
    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2) != 0.0)
    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2), dtypes.fp8e5m2, op)

-  @unittest.skipUnless(dtypes.fp8e4m3fnuz in supported_dtypes, f"no fp8e4m3fnuz on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3fnuz), f"no fp8e4m3fnuz on {Device.DEFAULT}")
  @given(ht.fp8e4m3fnuz, strat.sampled_from(unary_operations))
  def test_fp8e4m3fnuz_unary(self, a, op):
    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz) != 0.0)
    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)

-  @unittest.skipUnless(dtypes.fp8e5m2fnuz in supported_dtypes, f"no fp8e5m2fnuz on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2fnuz), f"no fp8e5m2fnuz on {Device.DEFAULT}")
  @given(ht.fp8e5m2fnuz, strat.sampled_from(unary_operations))
  def test_fp8e5m2fnuz_unary(self, a, op):
    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz) != 0.0)
@@ -257,15 +250,15 @@ class TestDTypeALU(unittest.TestCase):
  @given(ht.uint8, ht.uint8, strat.sampled_from(integer_binary_operations))
  def test_uint8(self, a, b, op): universal_test(a, b, dtypes.uint8, op)

-  @unittest.skipUnless(dtypes.uint16 in supported_dtypes, f"no uint16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.uint16), f"no uint16 on {Device.DEFAULT}")
  @given(ht.uint16, ht.uint16, strat.sampled_from(integer_binary_operations))
  def test_uint16(self, a, b, op): universal_test(a, b, dtypes.uint16, op)

-  @unittest.skipUnless(dtypes.uint32 in supported_dtypes, f"no uint32 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.uint32), f"no uint32 on {Device.DEFAULT}")
  @given(ht.uint32, ht.uint32, strat.sampled_from(integer_binary_operations))
  def test_uint32(self, a, b, op): universal_test(a, b, dtypes.uint32, op)

-  @unittest.skipUnless(dtypes.uint64 in supported_dtypes, f"no uint64 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.uint64), f"no uint64 on {Device.DEFAULT}")
  @given(ht.uint64, ht.uint64, strat.sampled_from(integer_binary_operations))
  def test_uint64(self, a, b, op): universal_test(a, b, dtypes.uint64, op)

@@ -294,15 +287,15 @@ class TestDTypeALU(unittest.TestCase):
  @given(ht.uint8, strat.sampled_from(integer_unary_operations))
  def test_uint8_unary(self, a, op): universal_test_unary(a, dtypes.uint8, op)

-  @unittest.skipUnless(dtypes.uint16 in supported_dtypes, f"no uint16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.uint16), f"no uint16 on {Device.DEFAULT}")
  @given(ht.uint16, strat.sampled_from(integer_unary_operations))
  def test_uint16_unary(self, a, op): universal_test_unary(a, dtypes.uint16, op)

-  @unittest.skipUnless(dtypes.uint32 in supported_dtypes, f"no uint32 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.uint32), f"no uint32 on {Device.DEFAULT}")
  @given(ht.uint32, strat.sampled_from(integer_unary_operations))
  def test_uint32_unary(self, a, op): universal_test_unary(a, dtypes.uint32, op)

-  @unittest.skipUnless(dtypes.uint64 in supported_dtypes, f"no uint64 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.uint64), f"no uint64 on {Device.DEFAULT}")
  @given(ht.uint64, strat.sampled_from(integer_unary_operations))
  def test_uint64_unary(self, a, op): universal_test_unary(a, dtypes.uint64, op)

@@ -331,12 +324,12 @@ class TestDTypeALU(unittest.TestCase):
  @given(ht.bool, ht.bool, strat.sampled_from(((operator.add, operator.add), (operator.mul, operator.mul))))
  def test_bool(self, a, b, op): universal_test(a, b, dtypes.bool, op)

+  @unittest.skipIf(not CI and Device.DEFAULT == "METAL", "broken on local M3")
  @given(ht.int32, ht.int32, ht.float32, strat.sampled_from(integer_binary_operations), strat.sampled_from(binary_operations))
  def test_int32_midcast_float(self, a, b, c, op1, op2): universal_test_midcast(a, b, c, op1, op2, dtypes.int32, dtypes.float32)

-  # Metal and (MOCK)CUDA and HIP and NIR behave differently than numpy for overflows
-  skip_overflow = ((DEV.interface.startswith("MOCK") and Device.DEFAULT in {"AMD", "NV", "CUDA"})
-                   or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer))
+  # Metal and CUDA and HIP and NIR behave differently than numpy in CI for overflows
+  skip_overflow = (CI and Device.DEFAULT in {"AMD", "NV", "CUDA"}) or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer)
  @given(strat.floats(width=32, min_value=0, max_value=10.0) if skip_overflow else ht.float32,
         strat.floats(width=32, min_value=0, max_value=10.0) if skip_overflow else ht.float32,
         ht.int32, strat.sampled_from(binary_operations), strat.sampled_from(integer_binary_operations))
@@ -355,21 +348,21 @@ class TestDTypeALU(unittest.TestCase):
  @given(strat.floats(width=32, min_value=1.0, max_value=254.0, allow_subnormal=False),
         strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
  def test_float_cast_to_unsigned(self, a, float_dtype, unsigned_dtype):
-    if float_dtype not in supported_dtypes: float_dtype = dtypes.float32
+    if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32
    universal_test_cast(a, float_dtype, unsigned_dtype)

  @unittest.skip("relied on hacks")
  @given(strat.floats(width=32, min_value=256.0, max_value=65000.0, allow_subnormal=False),
         strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
  def test_float_cast_to_unsigned_overflow(self, a, float_dtype, unsigned_dtype):
-    if float_dtype not in supported_dtypes: float_dtype = dtypes.float32
+    if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32
    universal_test_cast(a, float_dtype, unsigned_dtype)

  @unittest.skip("relied on hacks")
  @given(strat.floats(width=32, min_value=-65000.0, max_value=-1.0, allow_subnormal=False),
         strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
  def test_float_cast_to_unsigned_underflow(self, a, float_dtype, unsigned_dtype):
-    if float_dtype not in supported_dtypes: float_dtype = dtypes.float32
+    if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32
    universal_test_cast(a, float_dtype, unsigned_dtype)

  @unittest.expectedFailure
@@ -127,19 +127,19 @@ class TestInputValidation(unittest.TestCase):
    with self.assertRaises(ValueError):
      torch.optim.AdamW([torch.tensor([1.], requires_grad=True)], lr=0.1, weight_decay=-0.1)
    with self.assertRaises(ValueError):
-      nn.optim.AdamW([Tensor([1.])], lr=0.1, weight_decay=-0.1)
+      nn.optim.AdamW([Tensor([1.], requires_grad=True)], lr=0.1, weight_decay=-0.1)

  def test_negative_lr(self):
    with self.assertRaises(ValueError):
      torch.optim.SGD([torch.tensor([1.], requires_grad=True)], lr=-0.1)
    with self.assertRaises(ValueError):
-      nn.optim.SGD([Tensor([1.])], lr=-0.1)
+      nn.optim.SGD([Tensor([1.], requires_grad=True)], lr=-0.1)

  def test_negative_momentum(self):
    with self.assertRaises(ValueError):
      torch.optim.SGD([torch.tensor([1.], requires_grad=True)], lr=0.1, momentum=-0.1)
    with self.assertRaises(ValueError):
-      nn.optim.SGD([Tensor([1.])], lr=0.1, momentum=-0.1)
+      nn.optim.SGD([Tensor([1.], requires_grad=True)], lr=0.1, momentum=-0.1)

 class TestZeroFolding(unittest.TestCase):
  # we don't need more of these
@@ -1,149 +0,0 @@
-import unittest
-from tinygrad import Device
-from tinygrad.uop.ops import UOp, Ops
-from tinygrad.dtype import dtypes
-from tinygrad.renderer.isa.x86 import X86Ops, X86Renderer, RBP, RDI, RSP, RSI, RAX, RDX, XMM, GPR, imm, def_reg
-
-def ins(op, dt, src, tag=None): return UOp(Ops.INS, arg=op, dtype=dt, src=src, tag=tag)
-
-@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only on x86")
-class TestEncodingsX86(unittest.TestCase):
-  # NOTE: x86 supports a single displacement as memory address and index without base memory address
-  # these have no use cases so they aren't supported
-  def encode(self, u:UOp): return Device[Device.DEFAULT].renderer.render([u])
-
-  # displacement of 0 isn't emitted
-  def test_base_address(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RDI), UOp(Ops.NOOP), imm(dtypes.int8, 0)), RDI)
-    # mov edi, dword ptr [rdi]
-    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 3F"))
-
-  # rsp/r12 require a sib byte when used as base memory address
-  def test_rsp_base_address(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RSP), UOp(Ops.NOOP), imm(dtypes.int8, 0)), RSP)
-    # mov esp, dword ptr [rsp]
-    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 24 24"))
-
-  # rbp/r13 require a displacement when used as base memory address
-  def test_rbp_base_address(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RBP), UOp(Ops.NOOP), imm(dtypes.int8, 0)), RBP)
-    # mov ebp, dword ptr [rbp + 0]
-    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 6D 00"))
-
-  # test [base + index*scale]
-  def test_base_index_address(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RAX), def_reg(dtypes.int32, RDX), imm(dtypes.int8, 0)), RAX)
-    # mov eax, dword ptr [rax + rdx*4]
-    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 04 90"))
-
-  # rsp as index means no index
-  def test_rsp_index_address(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RAX), def_reg(dtypes.int32, RSP), imm(dtypes.int8, 0)), RAX)
-    # mov eax, dword ptr [rax]
-    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 00"))
-
-  # however r12 is a valid index
-  def test_r12_index_address(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RAX), def_reg(dtypes.int32, GPR[12]), imm(dtypes.int8, 0)), RAX)
-    # mov eax, dword ptr [rax + r12*4]
-    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("42 8B 04 A0"))
-
-  # test [base + index*scale + 8bit disp]
-  def test_complex_address_8bit_disp(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10)), RDI)
-    # mov edi, dword ptr [rdi + rsi*4 + 0xa]
-    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 7C B7 0A"))
-
-  # test [base + index*scale + 32bit disp]
-  def test_complex_address_32bit_disp(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int32, 10000)), RDI)
-    # mov edi, dword ptr [rdi + rsi*4 + 0x2710]
-    self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B BC B7 10 27 00 00"))
-
-  # 8bit variants of legacy instructions subtract 1 from opcode
-  def test_8bit_legacy_encoding(self):
-    cast = ins(X86Ops.MOVSX, dtypes.int32, (def_reg(dtypes.int8, RDX),), RAX)
-    # movsx eax, dl
-    self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("0F BE C2"))
-
-  # accessing lower 8 bits of rsp, rbp, rsi, rdi requires rex prefix
-  def test_lower_8bits_reg(self):
-    cast = ins(X86Ops.MOVSX, dtypes.int32, (def_reg(dtypes.int8, RDI),), RAX)
-    # movsx eax, dil
-    self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("40 0F BE C7"))
-
-  # test 16 bit variant of legacy instruction
-  def test_16bit_legacy_encoding(self):
-    cast = ins(X86Ops.MOVSX, dtypes.int16, (def_reg(dtypes.int8, RDX),), RAX)
-    # movsx ax, dl
-    self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("66 0F BE C2"))
-
-  # test 64 bit variant of legacy instruction
-  def test_64bit_legacy_encoding(self):
-    cast = ins(X86Ops.MOVSX, dtypes.int64, (def_reg(dtypes.int8, RDX),), RAX)
-    # movsx rax, dl
-    self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("48 0F BE C2"))
-
-  # test compact vex encoding
-  def test_compact_vex_encoding(self):
-    xmm0, xmm1 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[1])
-    add = ins(X86Ops.VADDSS, dtypes.float32, (xmm0, xmm1), XMM[0])
-    # vaddss xmm0, xmm0, xmm1
-    self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C5 FA 58 C1"))
-
-  # test long vex encoding
-  def test_long_vex_encoding(self):
-    xmm0, xmm8 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[8])
-    add = ins(X86Ops.VADDSS, dtypes.float32, (xmm0, xmm8), XMM[0])
-    # vaddss xmm0, xmm0, xmm8
-    self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C4 C1 7A 58 C0"))
-
-  # test ymm encoding
-  def test_ymm_encoding(self):
-    xmm0, xmm1 = def_reg(dtypes.float32.vec(8), XMM[0]), def_reg(dtypes.float32.vec(8), XMM[1])
-    add = ins(X86Ops.VADDPS, dtypes.float32.vec(8), (xmm0, xmm1), XMM[0])
-    # vaddps ymm0, ymm0, ymm1
-    self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C5 FC 58 C1"))
-
-  # test encoding where register is in the immediate field
-  def test_reg_in_imm_field(self):
-    xmm0, xmm1, xmm2 = def_reg(dtypes.float32, XMM[0]), def_reg(dtypes.float32, XMM[1]), def_reg(dtypes.float32, XMM[2])
-    blend = ins(X86Ops.VBLENDVPS, dtypes.float32, (xmm0, xmm1, xmm2), XMM[0])
-    # vblendvps xmm0, xmm0, xmm1, xmm2
-    self.assertEqual(bytes.fromhex(self.encode(blend)), bytes.fromhex("C4 E3 79 4A C1 20"))
-
-  # when writting to mem the uop takes the store form where dtype is void and there's no definition
-  def test_write_mem(self):
-    base, index, disp = def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10)
-    xmm0 = def_reg(dtypes.float32, XMM[0])
-    extr = ins(X86Ops.VPEXTRD, dtypes.void, (base, index, disp, xmm0, imm(dtypes.uint8, 0)))
-    # vpextrd dword ptr [rdi + rsi*4 + 0xa], xmm0, 0
-    self.assertEqual(bytes.fromhex(self.encode(extr)), bytes.fromhex("C4 E3 79 16 44 B7 0A 00"))
-
-  # test two address instruction with fused load works
-  def test_two_address_load(self):
-    base, index, disp = def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int8, 10)
-    cmove = ins(X86Ops.CMOVE, dtypes.int32,  (base, index, disp), RAX)
-    # cmove eax, dword ptr [rdi + rsi*4 + 0xa]
-    self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 44 B7 0A"))
-
-  # test instruction where displacement and imm have the same value
-  def test_disp_imm_same_value(self):
-    base, index, disp = def_reg(dtypes.int8.ptr(), RDI), def_reg(dtypes.int8, RSI), imm(dtypes.int8, 10)
-    mov = ins(X86Ops.MOVi, dtypes.void, (base, index, disp, disp))
-    # mov byte ptr [rdi + rsi + 0xa], 0xa
-    self.assertEqual(bytes.fromhex(self.encode(mov)), bytes.fromhex("40 C6 44 37 0A 0A"))
-
-    base, index, disp = def_reg(dtypes.int32.ptr(), RDI), def_reg(dtypes.int32, RSI), imm(dtypes.int32, 10)
-    imul = ins(X86Ops.IMULi, dtypes.int32, (base, index, disp) + (imm(dtypes.int32, 10),), RDI)
-    # imul edi, dword ptr [rdi + rsi*4 + 0xa], 0xa
-    self.assertEqual(bytes.fromhex(self.encode(imul)), bytes.fromhex("69 BC B7 0A 00 00 00 0A 00 00 00"))
-
-  # cmoves have the cmp as the last src even though it is not explicitly used, the cmp doesn't define a reg and is ignored in the encoding
-  def test_cmove_ignore_cmp(self):
-    cmove = ins(X86Ops.CMOVE, dtypes.int32, (def_reg(dtypes.int32, RAX), UOp(Ops.INS, arg=X86Ops.CMP)), RDX)
-    # cmove edx, eax
-    self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 D0"))
-
-if __name__ == "__main__":
-  unittest.main()
@@ -1,9 +1,9 @@
 #!/usr/bin/env python
-import unittest, os
+import unittest
 import torch
 import numpy as np

-from tinygrad.helpers import DEV
+from tinygrad.helpers import CI, DEV
 from tinygrad.tensor import Tensor
 from tinygrad.device import Device
 from tinygrad.dtype import _from_torch_dtype, _to_torch_dtype
@@ -27,7 +27,7 @@ class TestInterop(unittest.TestCase):
    tg_out = tg_data[:, :, 0] * 0.2989 + tg_data[:, :, 1] * 0.5870 + tg_data[:, :, 2] * 0.1140
    tg_res = tg_out.numpy()

-    if self.torch_device == "mps" and os.getenv("CI", "") != "":
+    if self.torch_device == "mps" and CI:
      # MPS backend out of memory: https://discuss.pytorch.org/t/mps-back-end-out-of-memory-on-github-action/189773
      # Calculate expected value on cpu.
      inp = inp.cpu()
@@ -1,147 +0,0 @@
-import unittest
-from typing import cast
-from tinygrad import Device
-from tinygrad.uop import Ops
-from tinygrad.uop.ops import UOp, dtypes, graph_rewrite
-from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops
-from tinygrad.renderer.isa import IselContext
-
-@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only x86")
-class TestIselX86(unittest.TestCase):
-  def isel_rewrite(self, x:UOp):
-    return graph_rewrite(x, cast(X86Renderer, Device[Device.DEFAULT].renderer).isel_matcher, IselContext(x), bottom_up=True)
-
-  def _check_op(self, dt_op, expr):
-    nargs = expr.__code__.co_argcount
-    for dt,op in dt_op:
-      with self.subTest(dtype=dt):
-        v = [UOp.variable(str(i), 0, 0, dt) for i in range(nargs)]
-        n = self.isel_rewrite(expr(*v))
-        self.assertIs(n.arg, op)
-
-  def test_cmove(self):
-    a = UOp.variable("a", 0, 0, dtypes.int32)
-    b = UOp.variable("b", 0, 0, dtypes.int32)
-    c = (a < b).where(a, b)
-    d = (a != b).where(a, b)
-    f = c + d
-    n = self.isel_rewrite(f)
-    self.assertTrue(n.src[0].arg is X86Ops.CMOVL and n.src[1].arg is X86Ops.CMOVNE)
-    # both comparisons become the same instruction
-    self.assertTrue(n.src[0].src[2] == n.src[1].src[2] and n.src[0].src[2].arg is X86Ops.CMP)
-
-  def test_vmax(self):
-    dt_op = [(dtypes.float32, X86Ops.VMAXSS), (dtypes.float64, X86Ops.VMAXSD),
-             (dtypes.float32.vec(4), X86Ops.VMAXPS), (dtypes.float64.vec(4), X86Ops.VMAXPD)]
-    self._check_op(dt_op, lambda a,b: (a < b).where(b, a))
-
-  def test_vmin(self):
-    dt_op = [(dtypes.float32, X86Ops.VMINSS), (dtypes.float64, X86Ops.VMINSD),
-             (dtypes.float32.vec(4), X86Ops.VMINPS), (dtypes.float64.vec(4), X86Ops.VMINPD)]
-    self._check_op(dt_op, lambda a,b: (a < b).where(a, b))
-
-  def test_vfmadd(self):
-    dt_op = [(dtypes.float32, X86Ops.VFMADD213SS), (dtypes.float64, X86Ops.VFMADD213SD),
-             (dtypes.float32.vec(4), X86Ops.VFMADD213PS), (dtypes.float64.vec(4), X86Ops.VFMADD213PD)]
-    self._check_op(dt_op, lambda a,b,c: a * b + c)
-
-  # don't use fmadd if op being fused (mul) is used multiple times
-  def test_no_vfmadd(self):
-    dt_op = [(dtypes.float32, X86Ops.VADDSS), (dtypes.float64, X86Ops.VADDSD),
-             (dtypes.float32.vec(4), X86Ops.VADDPS), (dtypes.float64.vec(4), X86Ops.VADDPD)]
-    self._check_op(dt_op, lambda a,b: a * b + a * b)
-
-  def test_vpbroadcast(self):
-    a = UOp.variable("a", 0, 0, dtypes.int32)
-    n = self.isel_rewrite(a.broadcast(4))
-    # need to move src from gpr to xmm before broadcasting
-    self.assertTrue(n.arg is X86Ops.VPBROADCASTD and n.src[0].arg is X86Ops.VMOVD)
-    # if we can fuse a load we can skip the move and access memory directly
-    load = UOp.param(0, dtypes.int32.ptr()).index(UOp.const(dtypes.int32, 0), ptr=True).load()
-    n = self.isel_rewrite(load.broadcast(4))
-    self.assertTrue(n.arg is X86Ops.VPBROADCASTD and len(n.src) == 3)
-
-  def test_vbroadcastss(self):
-    a = UOp.variable("a", 0, 0, dtypes.float32)
-    valid = [UOp.vectorize(a, a, a, a), UOp.vectorize(a, a, a, a, a, a, a, a)]
-    for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VBROADCASTSS)
-
-  def test_vshufps(self):
-    a = UOp.variable("a", 0, 0, dtypes.float32.vec(8))
-    b = UOp.variable("b", 0, 0, dtypes.float32.vec(8))
-    c = UOp.variable("c", 0, 0, dtypes.float32)
-    d = UOp.variable("d", 0, 0, dtypes.float32)
-
-    valid = [UOp.vectorize(c, c, d, d),
-             UOp.vectorize(a.gep(0), a.gep(1), c, c),
-             UOp.vectorize(a.gep(0), a.gep(1), b.gep(2), b.gep(3)),
-             UOp.vectorize(a.gep(1), a.gep(2), a.gep(3), a.gep(0)),
-             UOp.vectorize(a.gep(3), a.gep(2), a.gep(1), a.gep(0), a.gep(7), a.gep(6), a.gep(5), a.gep(4)),
-             UOp.vectorize(a.gep(0), a.gep(0), b.gep(1), b.gep(1), a.gep(4), a.gep(4), b.gep(5), b.gep(5))]
-    for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VSHUFPS)
-
-    invalid = [UOp.vectorize(a.gep(0), a.gep(1), b.gep(4), b.gep(5)),
-               UOp.vectorize(a.gep(0), a.gep(5), b.gep(2), b.gep(3)),
-               UOp.vectorize(a.gep(0), a.gep(0), a.gep(0), a.gep(0), a.gep(4), a.gep(4), a.gep(4), a.gep(5)),
-               UOp.vectorize(a.gep(0), a.gep(0), b.gep(0), b.gep(0), a.gep(4), a.gep(4), b.gep(4), a.gep(4))]
-    for shuf in invalid: self.assertIsNot(self.isel_rewrite(shuf).arg, X86Ops.VSHUFPS)
-
-  def test_vshufpd(self):
-    a = UOp.variable("a", 0, 0, dtypes.float64.vec(4))
-    b = UOp.variable("b", 0, 0, dtypes.float64.vec(4))
-    c = UOp.variable("c", 0, 0, dtypes.float64)
-    d = UOp.variable("d", 0, 0, dtypes.float64)
-
-    valid = [UOp.vectorize(c, d),
-             UOp.vectorize(a.gep(0), c),
-             UOp.vectorize(a.gep(1), b.gep(1)),
-             UOp.vectorize(a.gep(0), b.gep(1), a.gep(2), b.gep(3)),
-             UOp.vectorize(a.gep(1), a.gep(1), a.gep(3), a.gep(3))]
-    for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VSHUFPD)
-
-    invalid = [UOp.vectorize(c, c, c, c),
-               UOp.vectorize(a.gep(0), a.gep(1), b.gep(2), b.gep(3)),
-               UOp.vectorize(a.gep(2), b.gep(3), a.gep(2), b.gep(3)),
-               UOp.vectorize(a.gep(0), b.gep(1), a.gep(0), b.gep(1))]
-    for shuf in invalid: self.assertIsNot(self.isel_rewrite(shuf).arg, X86Ops.VSHUFPD)
-
-  def test_vinsertps(self):
-    a = UOp.variable("a", 0, 0, dtypes.float32.vec(4))
-    b = UOp.variable("b", 0, 0, dtypes.float32.vec(4))
-    c = UOp.variable("c", 0, 0, dtypes.float32.vec(4))
-    d = UOp.variable("e", 0, 0, dtypes.float32)
-    # moving 0th element to position 0 does nothing so only 1 vinsertps is generated
-    n = self.isel_rewrite(UOp.vectorize(a.gep(0), d))
-    self.assertIs(n.arg, X86Ops.VINSERTPS)
-    self.assertIsNot(n.src[0].arg, X86Ops.VINSERTPS)
-
-    valid = [UOp.vectorize(a.gep(0), b.gep(1), a.gep(2), b.gep(3)),
-             UOp.vectorize(a.gep(3), b.gep(2), c.gep(1), d)]
-    for shuf in valid: self.assertIs(self.isel_rewrite(shuf).arg, X86Ops.VINSERTPS)
-
-  # complex address is [base + index*scale + displacement]
-  def test_complex_address(self):
-    a = UOp.variable("a", 0, 0, dtypes.int32)
-    load = UOp.param(0, dtypes.int32.ptr()).index(a + 1, ptr=True).load()
-    n = self.isel_rewrite(load)
-    # displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32
-    self.assertTrue(n.src[2].op is Ops.CONST and n.src[2].dtype is dtypes.int8 and n.src[2].arg == 4)
-
-  def test_fold_load(self):
-    load1 = UOp.param(0, dtypes.int32.ptr()).index(UOp.const(dtypes.int32, 0), ptr=True).load()
-    load2 = UOp.param(0, dtypes.int32.ptr()).index(UOp.const(dtypes.int32, 1), ptr=True).load()
-    n = self.isel_rewrite(load1 + load2)
-    self.assertTrue(len(n.src) == 4)
-
-  # don't fold when used multiple times
-  def test_dont_fold_load(self):
-    load = UOp.param(0, dtypes.int32.ptr()).index(UOp.const(dtypes.int32, 0), ptr=True).load()
-    # used by multiple users
-    n = self.isel_rewrite(load + 1 + load)
-    self.assertTrue(len(n.src) == 2)
-    # used mutiple times by same user
-    n = self.isel_rewrite(load * load)
-    self.assertTrue(len(n.src) == 2)
-
-if __name__ == "__main__":
-  unittest.main()
@@ -10,9 +10,8 @@ from tinygrad.engine.jit import TinyJit, JitError, graph_class
 from tinygrad.device import Device
 from tinygrad.helpers import Context, JIT, DEV, GlobalCounters
 from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import Ops, UOp
+from tinygrad.uop.ops import Ops
 from extra.models.unet import ResBlock
-from tinygrad.renderer.isa.x86 import X86Renderer

 def _simple_test(add, extract=lambda x: x, N=10):
  for _ in range(5):
@@ -49,7 +48,7 @@ class TestJit(unittest.TestCase):
      y = (x + 1).contiguous().realize()
      z = x.shrink(((st, st + N),)).contiguous().realize()
      return y, z
-    x = Tensor.arange(2*N).clone().realize()
+    x = Tensor.arange(2*N).contiguous().realize()
    for _ in range(3): y, z = f(x, Variable("a", 0, N).bind(0))
    self.assertEqual(y.shape, (2*N,))
    self.assertEqual(z.shape, (N,))
@@ -92,7 +91,7 @@ class TestJit(unittest.TestCase):
    @TinyJit
    def f(x): return (x[2:5].contiguous() + 1).realize()
    for i in range(5):
-      x = (Tensor.arange(10).float() + i * 10).clone().realize()
+      x = (Tensor.arange(10).float() + i * 10).contiguous().realize()
      np.testing.assert_allclose(f(x).numpy(), x.numpy()[2:5] + 1)

  def test_jit_multiple_outputs(self):
@@ -107,7 +106,6 @@ class TestJit(unittest.TestCase):
      np.testing.assert_allclose(e.numpy(), a.numpy()*b.numpy(), atol=1e-4, rtol=1e-5)
    assert_jit_cache_len(f, 3)

-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "estimates are wrong for x86")
  def test_global_counters_jit(self):
    @TinyJit
    def f(a, b):
@@ -535,18 +533,6 @@ class TestJit(unittest.TestCase):
      f(Tensor(2.0)).item()
    # self.assertEqual(f(Tensor([2.0])).item(), 1.0) # TODO: wrong output, should be 3.0. currently depends on empty value

-  def test_jit_const_input(self):
-    @TinyJit
-    def f(x:Tensor) -> Tensor: return (x + 1).realize()
-    with self.assertRaises(JitError):
-      f(Tensor(UOp.const(dtypes.float, 2.0))).item()
-
-  def test_jit_deviceless_compute_input(self):
-    @TinyJit
-    def f(x:Tensor) -> Tensor: return (x + 1).realize()
-    with self.assertRaises(JitError):
-      f(Tensor(UOp.const(dtypes.float, 2.0) + UOp.const(dtypes.float, 1.0))).item()
-
  def test_jit_init_empty_alt(self):
    @TinyJit
    def f(a:Tensor, b:Tensor) -> Tensor: return b.assign(a+1)
@@ -8,10 +8,10 @@ class TestKernelCache(unittest.TestCase):
    if Device.DEFAULT not in ["CPU"]:
      self.skipTest("No custom kernel cache is implemented")

-    const_value = 0.6765677269
+    unique_const = 0.6765677269
    a = Tensor.rand(4,4).realize()
    b = Tensor.rand(4,4).realize()
-    x = a + b + const_value
+    x = a + b + unique_const
    x.realize()

    a1 = Tensor.rand(4,4).realize()
@@ -20,7 +20,7 @@ class TestKernelCache(unittest.TestCase):
    Device['CPU'].compiler.compile_cached = None # making it not callable

    try:
-      x1 = a1 + b1 + const_value
+      x1 = a1 + b1 + unique_const
      x1.realize() # Same kernel should be from cache.
    finally:
      Device['CPU'].compiler.compile_cached = orig_compile_func
@@ -3,7 +3,7 @@ import unittest

 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.uop.ops import UOp, Ops, GroupOp, AxisType, buffers
-from tinygrad.device import Device, Buffer
+from tinygrad.device import Device, Buffer, is_dtype_supported
 from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
@@ -11,13 +11,11 @@ from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT, DEV
 from tinygrad.dtype import DType, dtypes, PtrDType, AddrSpace
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.cstyle import CUDARenderer
-from tinygrad.renderer.isa import ISARenderer
 from test.helpers import replace_opts
 MOCKGPU = DEV.interface.startswith("MOCK")

 from tinygrad.uop.render import print_uops # noqa: F401 # pylint: disable=unused-import

-@unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, ISARenderer), "isa backends don't preserve the op spec when lowering")
 class TestLinearizer(unittest.TestCase):
  def test_arg_dedup(self):
    # NOTE: this realize exists because Tensor.numpy calls .contiguous() internally
@@ -206,7 +204,7 @@ class TestLinearizer(unittest.TestCase):
  def test_sum_acc_dtype(self):
    for tensor_dtype, acc_dtype in (
      (dtypes.bool, dtypes.int), (dtypes.int16, dtypes.int), (dtypes.float16, dtypes.float), (dtypes.bfloat16, dtypes.float)):
-      if tensor_dtype in (dts:=Device[Device.DEFAULT].renderer.supported_dtypes()) and acc_dtype in dts:
+      if is_dtype_supported(tensor_dtype) and is_dtype_supported(acc_dtype):
        a = Tensor([1, 2, 3], dtype=tensor_dtype).sum()
        realized_ast = a.schedule_linear().src[-1].src[0]
        program = to_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
@@ -229,7 +227,7 @@ class TestLinearizer(unittest.TestCase):
      (dtypes.float, dtypes.float16, dtypes.float16),
    )
    for tensor_dtype, acc_dtype, expected_dtype in tests:
-      if tensor_dtype in (dts:=Device[Device.DEFAULT].renderer.supported_dtypes()) and acc_dtype in dts and expected_dtype in dts:
+      if is_dtype_supported(tensor_dtype) and is_dtype_supported(acc_dtype) and is_dtype_supported(expected_dtype):
        a, b = Tensor.rand(8, 8, dtype=tensor_dtype), Tensor.rand(8, 8, dtype=tensor_dtype)
        helper_arg_acc_dtype(a.sum(dtype=acc_dtype), expected_dtype)
        helper_arg_acc_dtype(a.matmul(b, dtype=acc_dtype), expected_dtype)
@@ -318,7 +316,6 @@ class TestLinearizer(unittest.TestCase):
    #helper(Tensor.arange(256), max_ops=2)
    helper(Tensor.arange(255), max_ops=2)

-  @unittest.skip("test implicitly depends on certain optimizations")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx for some reason")
  def test_grouped_store_phis(self):
@@ -341,7 +338,6 @@ class TestLinearizer(unittest.TestCase):
    for val in store_vals:
      assert val.dtype == dtypes.float.vec(4) # and val.op is not Ops.VECTORIZE

-  @unittest.skip("test implicitly depends on certain optimizations")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
  def test_grouped_store_values(self):
    x = Tensor.randn((4,3,6,6)).realize()
@@ -373,7 +369,6 @@ class TestLinearizer(unittest.TestCase):
    # assert barrier.src == tuple(local_stores)
    assert len([u for u in uops if u.op is Ops.IF])

-  @unittest.skip("test implicitly depends on certain optimizations")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
@@ -11,16 +11,16 @@ from tinygrad.codegen import to_program
 class TestLinearizerFailure(unittest.TestCase):
  @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
  def test_failure_beam_mnist(self):
-    c0 = UOp.param(0, dtypes.uchar.ptr(4014080))
+    c0 = UOp(Ops.PARAM, dtypes.uchar.ptr(4014080), arg=0, src=())
    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 0, AxisType.GLOBAL)
    c2 = UOp.range(UOp.const(dtypes.weakint, 784), 1, AxisType.GLOBAL)
    c3 = UOp.range(UOp.const(dtypes.weakint, 10), 3, AxisType.GLOBAL)
-    c4 = UOp.param(1, dtypes.int.ptr(512))
+    c4 = UOp(Ops.PARAM, dtypes.int.ptr(512), arg=1, src=())
    c5 = c4.index(c1.valid(UOp.const(dtypes.bool, True)))
    c6 = UOp.range(UOp.const(dtypes.weakint, 6000), 1004, AxisType.REDUCE)
    c7 = UOp.range(UOp.const(dtypes.weakint, 3750), 2006, AxisType.REDUCE)
    c8 = UOp.range(UOp.const(dtypes.weakint, 16), 2007, AxisType.GROUP_REDUCE)
-    c9 = UOp.param(2, dtypes.uchar.ptr(47040000))
+    c9 = UOp(Ops.PARAM, dtypes.uchar.ptr(47040000), arg=2, src=())
    c10 = c9.index((((c3*UOp.const(dtypes.weakint, 4704000))+c2)+(c6*UOp.const(dtypes.weakint, 784))).valid(UOp.const(dtypes.bool, True)))
    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.weakint, 6000))+c6)+((c7*UOp.const(dtypes.weakint, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.weakint, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
    c12 = c0.index((((c1*UOp.const(dtypes.weakint, 7840))+(c2*UOp.const(dtypes.weakint, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11).end(c1, c2, c3)
@@ -1,86 +0,0 @@
-import unittest
-from tinygrad import Tensor, Device, dtypes, Context
-from tinygrad.helpers import getenv
-from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8
-from extra.llama_kernels.fused_ce import fused_ce_loss
-from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed, quantize_fp8_scalar
-from test.helpers import needs_second_gpu
-
-def run_fused_ce(bs:int, seqlen:int, vocab:int, label_smoothing:float=0.0) -> None:
-  Tensor.manual_seed(0)
-  logits_rand = Tensor.randn(bs, seqlen, vocab).cast(dtypes.bfloat16)
-  targets = Tensor.randint(bs, seqlen, high=vocab, dtype=dtypes.int32)
-  logits, logits_ref = logits_rand.clone(), logits_rand.detach().float().contiguous()
-  with Context(DEBUG=0):
-    Tensor.realize(logits, logits_ref, targets)
-
-  loss = fused_ce_loss(logits, targets, label_smoothing=label_smoothing)
-  loss.backward()
-  Tensor.realize(loss, logits.grad)
-
-  ref = logits_ref.sparse_categorical_crossentropy(targets, label_smoothing=label_smoothing)
-  ref.backward()
-  Tensor.realize(ref, logits_ref.grad)
-
-  assert logits.grad.shape == (bs, seqlen, vocab)
-  with Context(DEBUG=0):
-    assert loss.allclose(ref, atol=2e-3, rtol=2e-3).item(), "forward mismatch"
-    assert logits.grad.allclose(logits_ref.grad, atol=2e-3, rtol=2e-3).item(), "grad mismatch"
-
-class TestFusedCE(unittest.TestCase):
-  def setUp(self):
-    if dtypes.bfloat16 not in Device[Device.DEFAULT].renderer.supported_dtypes(): self.skipTest("need bfloat16")
-
-  def test_fused_ce_1_2_16(self): run_fused_ce(1, 2, 16, label_smoothing=0.2)
-  def test_fused_ce_2_16_128(self): run_fused_ce(2, 16, 128)
-  def test_fused_ce_4_128_1024(self): run_fused_ce(4, 128, 1024, label_smoothing=0.2)
-
-  # note: this is the shape used in llama 8b
-  #def test_fused_ce_smoothing_16_1024_128256(self): run_fused_ce(16, 1024, 128256, label_smoothing=0.2)
-
-def run_quantize_fp8(shape:tuple[int, ...], delayed:bool=True) -> None:
-  Tensor.manual_seed(0)
-  x = Tensor.randn(*shape).cast(dtypes.bfloat16).contiguous()
-  amax_state = Tensor.full((), 2.0, dtype=dtypes.float32).contiguous()
-  with Context(DEBUG=0): Tensor.realize(x, amax_state)
-
-  if delayed:
-    fp8, inv_scale, new_amax, _ = quantize_fp8_delayed(x, amax_state, FP8_DTYPE)
-    ref_fp8, ref_inv_scale, ref_new_amax = quantize_fp8(x, amax_state=amax_state)
-    Tensor.realize(fp8, inv_scale, new_amax)
-    Tensor.realize(ref_fp8, ref_inv_scale, ref_new_amax)
-  else:
-    fp8 = quantize_fp8_scalar(x, amax_state, FP8_DTYPE)
-    ref_fp8, _, _ = quantize_fp8(x, amax_state=amax_state)
-    Tensor.realize(fp8)
-    Tensor.realize(ref_fp8)
-
-  with Context(DEBUG=0):
-    assert fp8.cast(dtypes.float).allclose(ref_fp8.cast(dtypes.float), atol=0, rtol=0).item(), "fp8 mismatch"
-    if delayed:
-      assert inv_scale.allclose(ref_inv_scale, atol=0, rtol=0).item(), "inv_scale mismatch"
-      assert new_amax.allclose(ref_new_amax, atol=0, rtol=0).item(), \
-        f"amax mismatch: got={new_amax.item()} ref={ref_new_amax.item()} diff={abs(new_amax.item()-ref_new_amax.item())}"
-
-class TestQuantizeFP8(unittest.TestCase):
-  def setUp(self):
-    ren = Device[Device.DEFAULT].renderer
-    if dtypes.bfloat16 not in ren.supported_dtypes(): self.skipTest("need bfloat16")
-    if not ren.has_local or not ren.has_shared: self.skipTest("need local/shared")
-
-  def test_scalar(self): run_quantize_fp8((getenv("N", 1024), 32), delayed=False)
-  def test_delayed(self): run_quantize_fp8((getenv("N", 2048), 1024))
-
-  @needs_second_gpu
-  def test_multi(self):
-    devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(8))
-    x = Tensor.empty(2048*8, 1024, dtype=dtypes.bfloat16, device=devs).uop.multi(0)
-    x = Tensor(x, device=devs)
-    amax_state = Tensor.full((), 2.0, dtype=dtypes.float32, device=devs).contiguous()
-    fp8, _, new_amax, _ = quantize_fp8_delayed(x, amax_state, FP8_DTYPE)
-    Tensor.realize(fp8, new_amax)
-    assert fp8.uop.shape == x.uop.shape
-    assert new_amax.shape == ()
-
-if __name__ == '__main__':
-  unittest.main()
@@ -1,5 +1,6 @@
 import unittest, random
 from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variable
+from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.helpers import getenv, prod, Context
 from tinygrad.nn.state import get_parameters, get_state_dict
@@ -97,11 +98,6 @@ class TestMultiTensor(unittest.TestCase):
    self.assertEqual(r.tolist(), out)
  def test_shard_reshape(self): self._test_shard_op(lambda t:t.reshape(2, 2), [[1.,1.],[1.,1.]])
  def test_shard_elementwise(self): self._test_shard_op(lambda t:(t+t).reshape(2, 2), [[2.,2.],[2.,2.]])
-  def test_alu_deviceless_const(self):
-    s = Tensor([1.0, 2, 3, 4]).shard((f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"), axis=0)
-    np.testing.assert_equal((s + Tensor(UOp.const(dtypes.float, 1.0))).numpy(), [2, 3, 4, 5])
-    np.testing.assert_equal((s + Tensor(UOp.const(dtypes.float, 1.0)).reshape((1,)).expand((4,))).numpy(), [2, 3, 4, 5])
-
  def test_shard_reduce(self):
    self._test_shard_op(lambda t:t.reshape(2, 3).sum(axis=1), [3.,3.], n=6)
    self._test_shard_op(lambda t:t.reshape(2, 3).sum(axis=0), [2.,2.,2.], n=6)
@@ -121,7 +117,7 @@ class TestMultiTensor(unittest.TestCase):
      _ = Tensor(X.uop, dtype=dtypes.float)

  def test_sharded_arange(self):
-    sharded_arange = Tensor.arange(1000).clone().shard(devices_2, 0)
+    sharded_arange = Tensor.arange(1000).shard(devices_2, 0)
    sharded_arange.realize()
    np.testing.assert_equal(sharded_arange.numpy(), np.arange(1000))

@@ -187,7 +183,7 @@ class TestMultiTensor(unittest.TestCase):
    np.testing.assert_allclose(O.numpy(), X.numpy()[0:2]*W.numpy()[0:2] < 2)

  def test_shrink_on_shard_axis(self):
-    X = Tensor.arange(4*4).reshape(4,4).clone().realize()
+    X = Tensor.arange(4*4).reshape(4,4).realize()
    X_np = X.numpy()
    X.shard_(devices_2, 0)
    # only shrink on the device that owns the shard, this is enabled by the mselect simplifier
@@ -234,33 +230,12 @@ class TestMultiTensor(unittest.TestCase):
    for ring in (0, 2):
      GlobalCounters.reset()
      with Context(RING=ring, SCACHE=0):
-        t = Tensor.arange(32).clone().shard(devices_4, 0).to(Device.DEFAULT)
+        t = Tensor.arange(32).contiguous().shard(devices_4, 0).to(Device.DEFAULT)
        t.realize()
      kernel_counts[ring] = GlobalCounters.kernel_count
      self.assertEqual(t.device, Device.DEFAULT)
      np.testing.assert_equal(t.numpy(), np.arange(32))
-    self.assertEqual(kernel_counts[0], kernel_counts[2])
-
-  def test_to_single_device_gather_memory(self):
-    nrows, ncols = 64, 1024
-    nbytes = nrows*ncols*4
-    for devs in (devices_2, devices_4):
-      ndev = len(devs)
-      for axis in (0, 1):
-        sh = Tensor.arange(nrows*ncols).reshape(nrows, ncols).clone().shard(devs, axis).realize()
-        kernels, mem = {}, {}
-        for ring in (0, 2):
-          GlobalCounters.reset()
-          with Context(RING=ring, SCACHE=0):
-            t = sh.to(Device.DEFAULT)
-            t.realize()
-          kernels[ring], mem[ring] = GlobalCounters.kernel_count, GlobalCounters.global_mem
-          self.assertEqual(t.device, Device.DEFAULT)
-          np.testing.assert_equal(t.numpy(), np.arange(nrows*ncols).reshape(nrows, ncols))
-        self.assertEqual(kernels[0], kernels[2])
-        self.assertEqual(mem[0], mem[2])
-        self.assertLess(kernels[0], 2*ndev)
-        self.assertLessEqual(mem[0], 4*nbytes)
+    self.assertNotEqual(kernel_counts[0], kernel_counts[2])

  def test_allreduce_all2all(self):
    with Context(ALL2ALL=2):
@@ -293,7 +268,7 @@ class TestMultiTensor(unittest.TestCase):
    @TinyJit
    def f(x): return (x+1).contiguous().sum()
    for _ in range(5):
-      tt = Tensor.arange(0, 4).clone().realize().shard((d1,d2), 0).realize()
+      tt = Tensor.arange(0, 4).contiguous().realize().shard((d1,d2), 0).realize()
      out = f(tt)
      assert out.item() == 1+2+3+4

@@ -309,7 +284,7 @@ class TestMultiTensor(unittest.TestCase):
    @TinyJit
    def f(x): return (x.shard((d1,d2), 0)+1).contiguous().sum()
    for _ in range(5):
-      tt = Tensor.arange(0, 4).clone().realize()
+      tt = Tensor.arange(0, 4).contiguous().realize()
      out = f(tt)
      assert out.item() == 1+2+3+4

@@ -397,7 +372,7 @@ class TestMultiTensor(unittest.TestCase):

  def test_backward_sum(self):
    x = Tensor([[1.,2,3,4], [5,6,7,8]]).shard(devices_2, axis=0)
-    w = Tensor([1.,2,3,4]).shard(devices_2)
+    w = Tensor([1.,2,3,4], requires_grad=True).shard(devices_2)
    out = x * w
    out.mean().backward()
    tst = w.grad.numpy()
@@ -429,6 +404,7 @@ class TestMultiTensor(unittest.TestCase):
    B, T, embed_size, vocab_size = 4, 10, 20, 28

    layer = nn.Embedding(vocab_size, embed_size)
+    layer.weight.requires_grad = True
    x = Tensor(np.random.randint(0, vocab_size, (B, T), dtype=np.int32))
    z = layer(x)
    z.sum().backward()
@@ -436,6 +412,7 @@ class TestMultiTensor(unittest.TestCase):

    layer_sharded = nn.Embedding(vocab_size, embed_size)
    layer_sharded.weight.replace(layer.weight.shard(devices_2, axis=shard_weight_axis)).realize()
+    layer_sharded.weight.requires_grad = True
    x_sharded = x.shard(devices_2, axis=None)
    z_shard = layer_sharded(x_sharded)
    z_shard.sum().backward()
@@ -498,7 +475,7 @@ class TestMultiTensor(unittest.TestCase):

  def _test_model_train_step(self, m, fake_image, labels):
    from tinygrad.nn.optim import LARS
-    optimizer = LARS(get_parameters(m), 0.1)
+    optimizer = LARS(get_parameters(m), 0.1)  # set requires_grad for all params

    optimizer.zero_grad()
    m.load_from_pretrained()
@@ -577,7 +554,7 @@ class TestMultiTensor(unittest.TestCase):
  def test_multi_tensor_jit_graph_assign_updates_each_shard(self):
    @TinyJit
    def jf(out: Tensor) -> Tensor:
-      tmp = (Tensor.arange(4, dtype=dtypes.float).clone().shard(devices_2, 0) + 1).contiguous().realize()
+      tmp = (Tensor.arange(4, dtype=dtypes.float).shard(devices_2, 0) + 1).contiguous().realize()
      out.assign((tmp + 1).contiguous()).realize()
      return out

@@ -769,6 +746,11 @@ class TestMultiTensor(unittest.TestCase):
    t2.realize()
  def test_rand_like_on_shard_axis(self): self.test_rand_like_on_shard(0)

+  def test_rand_like_on_shard_axis_requires_grad(self):
+    t = Tensor.empty((16, 16)).shard(devices_2, axis=0)
+    self.assertIs(t.rand_like(requires_grad=True).requires_grad, True)
+    self.assertIs(t.rand_like(requires_grad=False).requires_grad, False)
+
  def test_rand_like_from_alu(self):
    a = Tensor.ones(4, 4).shard(devices_4, axis=0)
    aa = a + a
@@ -821,9 +803,9 @@ class TestMultiTensor(unittest.TestCase):
    t2.realize()
  def test_full_like_on_shard_axis(self): self.test_full_like_on_shard(0)

-  def test_const_like_shrink_on_shard_axis(self):
+  def test_full_like_shrink_on_shard_axis(self):
    t = Tensor.ones(16, 16, dtype=dtypes.int).shard(devices_2, axis=0)
-    out = t.const_like(2)[:, :8]
+    out = Tensor.full_like(t, 2)[:, :8]
    linear, var_vals = out.linear_with_vars()
    self.assertEqual(len(linear.src), 0)
    run_linear(linear, var_vals)
@@ -835,7 +817,7 @@ class TestMultiTensor(unittest.TestCase):
      output = X.dropout(0.5).numpy()
      unique, counts = np.unique(output, return_counts=True)
      assert set(unique) == {0, 2}, unique
-      assert 96 < counts[0] < 160, counts[0]
+      assert 100 < counts[0] < 156, counts[0]

  def test_dropout_on_shard_axis(self):
    with Tensor.train():
@@ -843,7 +825,7 @@ class TestMultiTensor(unittest.TestCase):
      output = X.dropout(0.5).numpy()
      unique, counts = np.unique(output, return_counts=True)
      assert set(unique) == {0, 2}, unique
-      assert 192 < counts[0] < 320, counts[0]
+      assert 200 < counts[0] < 312, counts[0]

  @unittest.skip("TODO: this requires forced_realize to be deleted.")
  def test_shard_memory(self):
@@ -854,7 +836,7 @@ class TestMultiTensor(unittest.TestCase):

  def test_clone(self):
    for axis in (None, 0):
-      t = Tensor.arange(16).reshape(4, 4).clone().shard(devices_2, axis=axis).contiguous().realize()
+      t = Tensor.arange(16).reshape(4, 4).shard(devices_2, axis=axis).contiguous().realize()
      t_clone = t.clone().realize()
      self.assertEqual(t_clone.device, t.device)
      self.assertEqual(t_clone.uop.axis, axis)
@@ -865,7 +847,7 @@ class TestMultiTensor(unittest.TestCase):
  @unittest.skip("RANGEIFY doesn't support multi const folding")
  def test_multi_const_folding(self):
    with Context(TRACK_MATCH_STATS=0):
-      a = Tensor.arange(3).clone().realize()
+      a = Tensor.arange(3).realize()
      zeros = Tensor.zeros(3).realize()
    b = a.to(devices_2)*zeros.to(devices_2)
    sched = b.schedule_linear().src
@@ -904,7 +886,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):

  # shrink a multitensor on sharded axis
  def test_shrink_bad_args(self):
-    t = Tensor.arange(64).reshape(8, 8).clone().realize()
+    t = Tensor.arange(64).reshape(8, 8).contiguous().realize()
    t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0)

    with self.assertRaises(AssertionError):
@@ -926,8 +908,8 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):

  @given(strat.sampled_from([dtypes.float, dtypes.int, dtypes.int64, dtypes.int16]))
  def test_ops(self, dtype):
-    if dtype not in Device[Device.DEFAULT].renderer.supported_dtypes(): return
-    t = Tensor.arange(64).reshape(8, 8).clone().realize()
+    if not is_dtype_supported(dtype): return
+    t = Tensor.arange(64).reshape(8, 8).contiguous().realize()
    t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0)
    for i in range(4):
      print(f"{i=}")
@@ -971,7 +953,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
      np.testing.assert_allclose(a.flip(-1).numpy(), b.flip(-1).numpy(), rtol=1e-7, atol=1e-3)

  def test_add_two_partitions(self):
-    t = Tensor.arange(64).reshape(8, 8).clone().realize()
+    t = Tensor.arange(64).reshape(8, 8).contiguous().realize()
    t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0)

    a = t.shrink(((2, 4), None))
@@ -988,7 +970,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):

  def test_add_different_tensors(self):
    devices = [f"{Device.DEFAULT}:{i}" for i in range(4)]
-    x = Tensor.arange(64).reshape(8, 8).clone().realize().shard(devices, axis=0)
+    x = Tensor.arange(64).reshape(8, 8).contiguous().realize().shard(devices, axis=0)

    to_add = []
    for i in range(len(devices)):
@@ -1098,7 +1080,7 @@ class TestBatchNorm(unittest.TestCase):
  @given(strat.sampled_from((False, True)))
  def test_batchnorm(self, is_training):
    devices = [f"{Device.DEFAULT}:{i}" for i in range(4)]
-    x = Tensor.arange(4096).reshape(8, 8, 8, 8).clone().realize().shard(devices, axis=0)
+    x = Tensor.arange(4096).reshape(8, 8, 8, 8).contiguous().realize().shard(devices, axis=0)

    with Tensor.train(is_training):
      bns = []
@@ -1106,6 +1088,8 @@ class TestBatchNorm(unittest.TestCase):
        bn = nn.BatchNorm2d(8)
        for p in get_parameters(bn):
          p.shard_(devices)
+        bn.weight.requires_grad = True
+        bn.bias.requires_grad = True
        bns.append(bn)

      bn_ts = []
@@ -1184,28 +1168,28 @@ class TestMultiBufferView(unittest.TestCase):

  @unittest.skip("flaky on LLVM")
  def test_shrink_non_shard_axis(self):
-    ref = Tensor.arange(8*4*10).reshape(8, 4, 10).clone().realize()
-    a = Tensor.arange(8*4*10).reshape(8, 4, 10).clone().shard(devices_2, axis=1).realize()
+    ref = Tensor.arange(8*4*10).reshape(8, 4, 10).contiguous().realize()
+    a = Tensor.arange(8*4*10).reshape(8, 4, 10).contiguous().shard(devices_2, axis=1).realize()
    self._check(ref, a, lambda t: t[3])

  def test_shrink_2d(self):
-    ref = Tensor.arange(6*4).reshape(6, 4).clone().realize()
-    a = Tensor.arange(6*4).reshape(6, 4).clone().shard(devices_2, axis=1).realize()
+    ref = Tensor.arange(6*4).reshape(6, 4).contiguous().realize()
+    a = Tensor.arange(6*4).reshape(6, 4).contiguous().shard(devices_2, axis=1).realize()
    self._check(ref, a, lambda t: t.shrink(((1, 4), None)))

  def test_reshape_then_shrink(self):
-    ref = Tensor.arange(8*6).reshape(8, 6).clone().realize()
-    a = Tensor.arange(8*6).reshape(8, 6).clone().shard(devices_2, axis=1).realize()
+    ref = Tensor.arange(8*6).reshape(8, 6).contiguous().realize()
+    a = Tensor.arange(8*6).reshape(8, 6).contiguous().shard(devices_2, axis=1).realize()
    self._check(ref, a, lambda t: t.reshape(4, 2, 6)[1])

  def test_chained_shrink(self):
-    ref = Tensor.arange(10*8).reshape(10, 8).clone().realize()
-    a = Tensor.arange(10*8).reshape(10, 8).clone().shard(devices_2, axis=1).realize()
+    ref = Tensor.arange(10*8).reshape(10, 8).contiguous().realize()
+    a = Tensor.arange(10*8).reshape(10, 8).contiguous().shard(devices_2, axis=1).realize()
    self._check(ref, a, lambda t: t.shrink(((2, 8), None)).shrink(((1, 4), None)))

  def test_4_devices(self):
-    ref = Tensor.arange(8*12).reshape(8, 12).clone().realize()
-    a = Tensor.arange(8*12).reshape(8, 12).clone().shard(devices_4, axis=1).realize()
+    ref = Tensor.arange(8*12).reshape(8, 12).contiguous().realize()
+    a = Tensor.arange(8*12).reshape(8, 12).contiguous().shard(devices_4, axis=1).realize()
    out = a[5].contiguous()
    linear, var_vals = out.linear_with_vars()
    if all(hasattr(Device[d].allocator, "_offset") for d in out.device):
@@ -1296,7 +1280,7 @@ class TestMultiSetitem(unittest.TestCase):
  @needs_second_gpu
  def setUp(self): pass

-  def _t(self, axis): return Tensor.arange(16).clone().realize().shard(self.device, axis=axis)
+  def _t(self, axis): return Tensor.arange(16).contiguous().realize().shard(self.device, axis=axis)

  def test_setitem_scalar_axis0(self):
    t = self._t(0)
@@ -149,6 +149,8 @@ class TestNN(unittest.TestCase):

    # create in tinygrad
    layer = Conv2d(C1, C2, kernel_size=K, stride=S, padding=P)
+    layer.weight.requires_grad = True
+    layer.bias.requires_grad = True

    # create in torch
    torch_layer = torch.nn.Conv2d(C1, C2, kernel_size=K, stride=S, padding=P).eval()
@@ -156,7 +158,7 @@ class TestNN(unittest.TestCase):
    torch_layer.bias = torch.nn.Parameter(torch.tensor(layer.bias.numpy(), dtype=torch.float32))

    # test
-    x = Tensor.uniform(BS, C1, H, W)
+    x = Tensor.uniform(BS, C1, H, W, requires_grad=True)

    with Context(WINO=1):
      z = layer(x)
@@ -190,12 +192,12 @@ class TestNN(unittest.TestCase):

    # create in tinygrad
    layer = GroupNorm(G, C)
-    layer.weight = Tensor(torch_layer.weight.detach().numpy())
-    layer.bias = Tensor(torch_layer.bias.detach().numpy())
+    layer.weight = Tensor(torch_layer.weight.detach().numpy(), requires_grad=True)
+    layer.bias = Tensor(torch_layer.bias.detach().numpy(), requires_grad=True)

    for _ in range(10):
      # forward
-      x = Tensor.randn(BS, C, H, W)
+      x = Tensor.randn(BS, C, H, W, requires_grad=True)
      z = layer(x)
      z.sum().backward()

@@ -216,10 +218,10 @@ class TestNN(unittest.TestCase):

    # create in tinygrad
    layer = LayerNorm([H, W])
-    layer.weight = Tensor(torch_layer.weight.detach().numpy())
-    layer.bias = Tensor(torch_layer.bias.detach().numpy())
+    layer.weight = Tensor(torch_layer.weight.detach().numpy(), requires_grad=True)
+    layer.bias = Tensor(torch_layer.bias.detach().numpy(), requires_grad=True)

-    x = Tensor.empty(N, C, H, W)
+    x = Tensor.empty(N, C, H, W, requires_grad=True)
    z = layer(x)
    z.realize()

@@ -238,12 +240,12 @@ class TestNN(unittest.TestCase):

    # create in tinygrad
    layer = LayerNorm([H, W])
-    layer.weight = Tensor(torch_layer.weight.detach().numpy())
-    layer.bias = Tensor(torch_layer.bias.detach().numpy())
+    layer.weight = Tensor(torch_layer.weight.detach().numpy(), requires_grad=True)
+    layer.bias = Tensor(torch_layer.bias.detach().numpy(), requires_grad=True)

    for _ in range(10):
      # forward
-      x = Tensor.randn(N, C, H, W)
+      x = Tensor.randn(N, C, H, W, requires_grad=True)
      z = layer(x)
      z.sum().backward()

@@ -264,12 +266,12 @@ class TestNN(unittest.TestCase):

    # create in tinygrad
    layer = LayerNorm2d(C)
-    layer.weight = Tensor(torch_layer.weight.detach().numpy())
-    layer.bias = Tensor(torch_layer.bias.detach().numpy())
+    layer.weight = Tensor(torch_layer.weight.detach().numpy(), requires_grad=True)
+    layer.bias = Tensor(torch_layer.bias.detach().numpy(), requires_grad=True)

    for _ in range(10):
      # forward
-      x = Tensor.randn(N, C, H, W)
+      x = Tensor.randn(N, C, H, W, requires_grad=True)
      z = layer(x)
      z.sum().backward()

@@ -290,12 +292,12 @@ class TestNN(unittest.TestCase):

    # create in tinygrad
    layer = InstanceNorm(C)
-    layer.weight = Tensor(torch_layer.weight.detach().numpy())
-    layer.bias = Tensor(torch_layer.bias.detach().numpy())
+    layer.weight = Tensor(torch_layer.weight.detach().numpy(), requires_grad=True)
+    layer.bias = Tensor(torch_layer.bias.detach().numpy(), requires_grad=True)

    for _ in range(10):
      # forward
-      x = Tensor.randn(N, C, H, W)
+      x = Tensor.randn(N, C, H, W, requires_grad=True)
      z = layer(x)
      z.sum().backward()

@@ -316,12 +318,12 @@ class TestNN(unittest.TestCase):

    # create in tinygrad
    layer = InstanceNorm(C)
-    layer.weight = Tensor(torch_layer.weight.detach().numpy())
-    layer.bias = Tensor(torch_layer.bias.detach().numpy())
+    layer.weight = Tensor(torch_layer.weight.detach().numpy(), requires_grad=True)
+    layer.bias = Tensor(torch_layer.bias.detach().numpy(), requires_grad=True)

    for _ in range(10):
      # forward
-      x = Tensor.randn(N, C, D, H, W)
+      x = Tensor.randn(N, C, D, H, W, requires_grad=True)
      z = layer(x)
      z.sum().backward()

@@ -354,10 +356,11 @@ class TestNN(unittest.TestCase):
    B, T, embed_size = 4, 10, 20
    torch_layer = TorchRMSNorm(embed_size)
    layer = RMSNorm(embed_size)
+    layer.weight.requires_grad = True

    for _ in range(10):
      # forward
-      x = Tensor.randn(B, T, embed_size)
+      x = Tensor.randn(B, T, embed_size, requires_grad=True)
      z = layer(x)
      z.sum().backward()

@@ -374,7 +377,7 @@ class TestNN(unittest.TestCase):

    for _ in range(10):
      # forward
-      x = Tensor.randn(B, T, embed_size)
+      x = Tensor.randn(B, T, embed_size, requires_grad=True)
      z = layer(x)
      z.sum().backward()

@@ -1,10 +1,11 @@
-import time, math, unittest, functools, platform, warnings, sys
+import time, math, unittest, functools, platform, warnings
 import numpy as np
 from typing import List, Callable
 import torch
-from tinygrad.helpers import getenv, DEBUG, DEV, IMAGE, Context
+from tinygrad.helpers import getenv, CI, DEBUG, DEV, IMAGE, Context
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
+from tinygrad.device import is_dtype_supported
 from tinygrad.renderer.cstyle import QCOMCLRenderer
 from tinygrad.renderer.nir import NIRRenderer

@@ -73,7 +74,7 @@ def helper_test_op(shps, torch_fxn, tinygrad_fxn=None, atol=1e-6, rtol=1e-3, gra
    for i, (t, torch_grad) in enumerate(zip(tiny_grads, torch_grads)):
      compare(f"backward pass tensor {i}", t.numpy(), torch_grad.detach().cpu().numpy(), atol=grad_atol, rtol=grad_rtol)

-  if sys.stdout.isatty():
+  if not CI:
    print("\ntesting %40r   torch/tinygrad fp: %.2f / %.2f ms  bp: %.2f / %.2f ms " % \
          (shps, torch_fp*1000, tinygrad_fp*1000, torch_fbp*1000, tinygrad_fbp*1000), end="")

@@ -87,7 +88,7 @@ def prepare_test_op(low, high, shps, vals, forward_only=False):
  for i in range(len(ts)):
    # NOTE: torch default int64 for python ints input
    if ts[i].dtype == torch.int64: ts[i] = ts[i].type(torch.int32)
-  tst = [Tensor(x.detach().cpu().numpy()) for x in ts]
+  tst = [Tensor(x.detach().cpu().numpy(), requires_grad=(not forward_only and not FORWARD_ONLY)) for x in ts]
  return ts, tst

 class TestOps(unittest.TestCase):
@@ -102,7 +103,7 @@ class TestOps(unittest.TestCase):
    with self.assertRaises(expected) as tinygrad_cm:
      tinygrad_fxn(*tst)
    if exact: self.assertEqual(str(torch_cm.exception), str(tinygrad_cm.exception))
-    if sys.stdout.isatty(): print("\ntesting %40r   torch/tinygrad exception: %s / %s" % (shps, torch_cm.exception, tinygrad_cm.exception), end="")
+    if not CI: print("\ntesting %40r   torch/tinygrad exception: %s / %s" % (shps, torch_cm.exception, tinygrad_cm.exception), end="")

  def test_full_like(self):
    a = Tensor([[1,2,3],[4,5,6]], dtype=dtypes.float32)
@@ -248,9 +249,9 @@ class TestOps(unittest.TestCase):
    self.helper_test_exception([(8,)], lambda x: x.unfold(0, 1, -1), expected=RuntimeError)

  def test_meshgrid(self):
-    x, xt = torch.tensor([0.,1.,2.], requires_grad=True), Tensor([0.,1.,2.])
-    y, yt = torch.tensor([3.,4.,5.,6.], requires_grad=True), Tensor([3.,4.,5.,6.])
-    z, zt = torch.tensor([7.,8.,9.], requires_grad=True), Tensor([7.,8.,9.])
+    x, xt = torch.tensor([0.,1.,2.], requires_grad=True), Tensor([0.,1.,2.], requires_grad=True)
+    y, yt = torch.tensor([3.,4.,5.,6.], requires_grad=True), Tensor([3.,4.,5.,6.], requires_grad=True)
+    z, zt = torch.tensor([7.,8.,9.], requires_grad=True), Tensor([7.,8.,9.], requires_grad=True)
    for indexing in ("ij", "xy"):
      tor = torch.meshgrid(x, indexing=indexing)
      ten = xt.meshgrid(indexing=indexing)
@@ -263,7 +264,7 @@ class TestOps(unittest.TestCase):
      for tor_i, ten_i in zip(tor, ten):
        helper_test_op([], lambda: tor_i, lambda: ten_i)
      tor = torch.meshgrid(x, torch.tensor(10., requires_grad=True), y, z, indexing=indexing)
-      ten = xt.meshgrid(Tensor(10.), yt, zt, indexing=indexing)
+      ten = xt.meshgrid(Tensor(10., requires_grad=True), yt, zt, indexing=indexing)
      self.assertEqual(len(tor), len(ten))
      for tor_i, ten_i in zip(tor, ten):
        helper_test_op([], lambda: tor_i, lambda: ten_i)
@@ -385,11 +386,11 @@ class TestOps(unittest.TestCase):
    t1 = torch.ones(4, requires_grad=True)
    t2 = torch.ones(4, requires_grad=True)
    self.assertRaises(RuntimeError, (t1 != t2).sum().backward)
-    tt1 = Tensor.ones(4)
-    tt2 = Tensor.ones(4)
+    tt1 = Tensor.ones(4, requires_grad=True)
+    tt2 = Tensor.ones(4, requires_grad=True)
    self.assertRaises(RuntimeError, (tt1 != tt2).sum().backward)
    """
-    tt = Tensor.randn(4)
+    tt = Tensor.randn(4, requires_grad=True)
    (tt*(tt != 0)).sum().backward()
    t = torch.tensor(tt.numpy(), requires_grad=True)
    (t*(t != 0)).sum().backward()
@@ -401,11 +402,11 @@ class TestOps(unittest.TestCase):
    t1 = torch.ones(4, requires_grad=True)
    t2 = torch.ones(4, requires_grad=True)
    self.assertRaises(RuntimeError, (t1 < t2).sum().backward)
-    tt1 = Tensor.ones(4)
-    tt2 = Tensor.ones(4)
+    tt1 = Tensor.ones(4, requires_grad=True)
+    tt2 = Tensor.ones(4, requires_grad=True)
    self.assertRaises(RuntimeError, (tt1 < tt2).sum().backward)
    """
-    tt = Tensor.randn(4)
+    tt = Tensor.randn(4, requires_grad=True)
    (tt*(tt < 0)).sum().backward()
    t = torch.tensor(tt.numpy(), requires_grad=True)
    (t*(t < 0)).sum().backward()
@@ -682,7 +683,7 @@ class TestOps(unittest.TestCase):
    helper_test_op([(45,65)], lambda x: x**1.2, low=-30, high=-27)
    helper_test_op([()], lambda x: x**0.2, low=-30, high=-27)
    helper_test_op([()], lambda x: x**1.2, low=-30, high=-27)
-    a, b = Tensor([0.0]), torch.tensor([0.0], requires_grad=True)
+    a, b = Tensor([0.0], requires_grad=True), torch.tensor([0.0], requires_grad=True)
    helper_test_op([], lambda: b**1.1, lambda: a**1.1)

  def test_pow_const(self):
@@ -1448,7 +1449,9 @@ class TestOps(unittest.TestCase):
                                                                         np.arange(64,128,dtype=np.float32).reshape(8,8)])
  def test_small_gemm_eye(self):
    helper_test_op(None, lambda x,y: x.matmul(y), lambda x,y: x@y, vals=[np.eye(8).astype(np.float32), np.eye(8).astype(np.float32)])
-  @unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "not precise enough when emulating")
+  @unittest.skipIf(CI and Device.DEFAULT in ["NV", "CL", "CUDA"] or (Device.DEFAULT == "CPU" and DEV.renderer == "LLVM") or IMAGE
+  or (Device.DEFAULT == "WEBGPU" and platform.system() == "Windows"), "not supported on these in CI/IMAGE")
+  @unittest.skipIf(Device.DEFAULT == "QCOM", "not precise enough")
  def test_gemm_fp16(self):
    helper_test_op([(64,64), (64,64)], lambda x,y: x.half().matmul(y.half()), atol=5e-3, rtol=5e-3, grad_atol=5e-3, grad_rtol=5e-3)
  def test_gemm(self):
@@ -1505,8 +1508,7 @@ class TestOps(unittest.TestCase):

  def test_sum_dtype_arg(self):
    helper_test_op([(45,3)], lambda x: x.sum(), lambda x: x.sum(dtype=dtypes.float32))
-    if dtypes.float64 in Device[Device.DEFAULT].renderer.supported_dtypes():
-      helper_test_op([(45,3)], lambda x: x.sum(dtype=torch.float64), lambda x: x.sum(dtype=dtypes.float64))
+    if is_dtype_supported(dtypes.float64): helper_test_op([(45,3)], lambda x: x.sum(dtype=torch.float64), lambda x: x.sum(dtype=dtypes.float64))

    with self.assertRaises(AttributeError): Tensor([1.0, 2.0]).sum(dtype="")

@@ -2447,7 +2449,7 @@ class TestOps(unittest.TestCase):

  @unittest.skipUnless(Device.DEFAULT == "CPU" and DEV.renderer == "LLVM", "DEVECTORIZE=0 only for LLVM")
  def test_strided_conv2d_simple_vec(self):
-    self.test_strided_conv2d_simple()
+    with Context(DEVECTORIZE=0): self.test_strided_conv2d_simple()

  @slow_test
  def test_strided_conv2d(self):
@@ -2845,7 +2847,6 @@ class TestOps(unittest.TestCase):
    for dim in range(-1, 2):
      helper_test_op([(45,65), (45,65), (45,65)], lambda x,y,z: torch.cat((x,y,z), dim), lambda x,y,z: x.cat(y, z, dim=dim))

-  @unittest.skipIf(COMPILE_ONLY, "test requires runtime")
  def test_stack(self):
    for dim in range(-1, 3):
      helper_test_op([(5,6,3), (5,6,3), (5,6,3)], lambda x, y, z: torch.stack((x, y, z), dim), lambda x, y, z: Tensor.stack(x, y, z, dim=dim))
@@ -2856,7 +2857,8 @@ class TestOps(unittest.TestCase):
    with self.assertRaises(ValueError):
      Tensor.stack((Tensor([1, 2]), Tensor([3, 4])), Tensor([5, 6]))

-    np.testing.assert_allclose(Tensor.stack(Tensor(3.14), Tensor(3.14)).numpy(), np.array([3.14, 3.14]))
+    a = Tensor(3.14)
+    np.testing.assert_allclose(Tensor.stack(a, a).numpy(), Tensor([3.14, 3.14]).numpy())

  def test_stack_max(self):
    helper_test_op(None, lambda x, y: torch.stack((x, y)).max(axis=0)[0], lambda x, y: Tensor.stack(x, y).max(axis=0), vals=[[1.], [2.]])
@@ -2924,7 +2926,7 @@ class TestOps(unittest.TestCase):
    c = torch.randint(low=-5, high=5, size=(1,1,4,1,1,1), dtype=torch.int64, requires_grad=False)
    d = torch.randint(high=4, size=(2,1,1,5,1,1), dtype=torch.int64, requires_grad=False)
    e = torch.randint(high=1, size=(1,1,1,1,6,1), dtype=torch.int64, requires_grad=False)
-    i, j, k, o, p = [Tensor(tor.detach().cpu().numpy().astype(np.int32)) for tor in [a,b,c,d,e]]
+    i, j, k, o, p = [Tensor(tor.detach().cpu().numpy().astype(np.int32), requires_grad=False) for tor in [a,b,c,d,e]]
    return a,b,c,d,e,i,j,k,o,p

  def test_fancy_indexing_inf(self):
@@ -3035,7 +3037,7 @@ class TestOps(unittest.TestCase):
    # indices cannot have gradient
    # indices cannot be negative (torch gather)
    b = torch.randint(3, size=[3,4,5], dtype=torch.int64, requires_grad=False)
-    a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32)
+    a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32, requires_grad=False)
    helper_test_op([(4,5,6)], lambda x: x.gather(dim=0, index=b), lambda x: x.gather(dim=0, index=a))
    helper_test_op([(4,5,6)], lambda x: x.gather(dim=1, index=b), lambda x: x.gather(dim=1, index=a))
    helper_test_op([(4,5,6)], lambda x: x.gather(dim=2, index=b), lambda x: x.gather(dim=2, index=a))
@@ -3057,7 +3059,7 @@ class TestOps(unittest.TestCase):

  def test_scatter(self):
    b = torch.randint(3, size=[3,4,5], dtype=torch.int64, requires_grad=False)
-    a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32)
+    a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32, requires_grad=False)
    for dim in (0,1,2,-1,-2,-3):
      helper_test_op([(4,5,6), (4,5,6)], lambda x,src: x.scatter(dim=dim, index=b, src=src),
                                         lambda x,src: x.scatter(dim=dim, index=a, src=src), forward_only=True)
@@ -3082,7 +3084,7 @@ class TestOps(unittest.TestCase):

    # overlapping indices with 0s
    b = torch.tensor([0,0], requires_grad=False)
-    a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32)
+    a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32, requires_grad=False)
    helper_test_op(None,
      lambda x,src: x.scatter(0, b, src),
      lambda x,src: x.scatter(0, a, src), forward_only=True,
@@ -3090,7 +3092,7 @@ class TestOps(unittest.TestCase):

  def test_scatter_add(self):
    b = torch.randint(3, size=[3,4,5], dtype=torch.int64, requires_grad=False)
-    a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32)
+    a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32, requires_grad=False)
    helper_test_op([(4,5,6)], lambda x: x.scatter(dim=1, index=b, value=float("inf"), reduce="add"),
      lambda x: x.scatter(dim=1, index=a, src=float("inf"), reduce="add"), forward_only=True)

@@ -3102,7 +3104,7 @@ class TestOps(unittest.TestCase):

  def test_scatter_mul(self):
    b = torch.randint(3, size=[3,4,5], dtype=torch.int64, requires_grad=False)
-    a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32)
+    a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32, requires_grad=False)
    helper_test_op([(4,5,6)], lambda x: x.scatter(dim=1, index=b, value=float("inf"), reduce="multiply"),
      lambda x: x.scatter(dim=1, index=a, src=float("inf"), reduce="multiply"), forward_only=True)

@@ -3119,7 +3121,7 @@ class TestOps(unittest.TestCase):
  @slow_test
  def test_scatter_reduce(self):
    b = torch.randint(3, size=[3,4,5], dtype=torch.int64, requires_grad=False)
-    a = Tensor(b.detach().cpu().numpy().astype(np.int32))
+    a = Tensor(b.detach().cpu().numpy().astype(np.int32), requires_grad=False)
    for reduce in ("sum", "prod", "mean", "amin", "amax"):
      for dim in (-1,1,-3):
        helper_test_op([(3,4,5), (3,4,5)],
@@ -3131,7 +3133,7 @@ class TestOps(unittest.TestCase):

  def test_scatter_reduce_prod_zeros(self):
    b = torch.randint(3, size=[3,4,5], dtype=torch.int64, requires_grad=False)
-    a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32)
+    a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32, requires_grad=False)
    x = Tensor.zeros([4,5,6]).float()
    y = torch.zeros([4,5,6]).float()
    helper_test_op([(4,5,6)],
@@ -3140,7 +3142,7 @@ class TestOps(unittest.TestCase):

  def test_scatter_reduce_errors(self):
    b = torch.randint(3, size=[3,4,5], dtype=torch.int64, requires_grad=False)
-    a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32)
+    a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32, requires_grad=False)
    # invalid reduce arg
    self.helper_test_exception([(4,5,6), (4,5,6)],
      lambda x,src: x.scatter_reduce(dim=0, index=b, src=src, reduce="INVALID"),
@@ -3184,7 +3186,7 @@ class TestOps(unittest.TestCase):
    self.helper_test_exception([(32,31,16,64), (32,8,16,64), (32,8,16,64)],
      lambda x,y,z: torch.nn.functional.scaled_dot_product_attention(x,y,z),
      lambda x,y,z: Tensor.scaled_dot_product_attention(x,y,z,enable_gqa=True),
-      expected=(AssertionError, RuntimeError, ValueError, IndexError))
+      expected=(AssertionError, RuntimeError, ValueError))

  def test_binary_crossentropy(self):
    helper_test_op([(32,10), (32,10)], lambda x,y: torch.nn.functional.binary_cross_entropy(x.sigmoid(),y.clip(0,1)),
@@ -3346,7 +3348,6 @@ class TestOps(unittest.TestCase):
    # fill_value must not alter output dtype
    self.assertEqual(Tensor([1.0, 2.0]).masked_select(Tensor([True, False]), size=3, fill_value=-1).dtype, dtypes.default_float)

-  @slow_test
  def test_nonzero(self):
    helper_test_op([(32, 10)], lambda x: (x>0.5).nonzero().int(), lambda x: (x>0.5).nonzero(), forward_only=True)
    helper_test_op([(20,)], lambda x: (x>0.5).nonzero().int(), lambda x: (x>0.5).nonzero(), forward_only=True)
@@ -3378,6 +3379,7 @@ class TestOps(unittest.TestCase):
    t = (Tensor([0], dtype='int') | 0xFFFFFFFF).item()
    if not COMPILE_ONLY: assert t == -1

+@unittest.skipUnless(is_dtype_supported(dtypes.uchar), f"no uint8 on {Device.DEFAULT}")
 class TestOpsUint8(unittest.TestCase):
  def test_cast(self):
    helper_test_op([(2,3,64,64)], lambda x: x.type(torch.uint8), lambda x: x.cast('uint8'), forward_only=True, low=0, high=255)
@@ -3,6 +3,7 @@ import torch
 import unittest
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.nn.optim import Adam, SGD, AdamW, Muon, LAMB
+from tinygrad.device import is_dtype_supported
 from test.helpers import needs_second_gpu, slow

 np.random.seed(1337)
@@ -10,20 +11,17 @@ x_init = np.random.randn(1,4).astype(np.float32)
 W_init = np.random.randn(4,4).astype(np.float32)
 m_init = np.random.randn(1,4).astype(np.float32)

-def _param(tensor, val):
-  return tensor(val, requires_grad=True) if tensor is torch.tensor else tensor(val)
-
 class TeenyNet:
  def __init__(self, tensor):
-    self.x = _param(tensor, x_init.copy())
-    self.W = _param(tensor, W_init.copy())
+    self.x = tensor(x_init.copy(), requires_grad=True)
+    self.W = tensor(W_init.copy(), requires_grad=True)
  def forward(self):
    return (self.x * self.W).sum()

 class TinyNet:
  def __init__(self, tensor):
-    self.x = _param(tensor, x_init.copy())
-    self.W = _param(tensor, W_init.copy())
+    self.x = tensor(x_init.copy(), requires_grad=True)
+    self.W = tensor(W_init.copy(), requires_grad=True)
    self.m = tensor(m_init.copy())

  def forward(self):
@@ -144,7 +142,7 @@ class TestOptim(unittest.TestCase):

      np.testing.assert_allclose(losses[0], losses[1], atol=1e-4, rtol=0)

-  @unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half")
+  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
  def test_mixed_precision(self):
    old_default_float, dtypes.default_float = dtypes.default_float, dtypes.half
    # weight update would overflow without upcasting
@@ -154,7 +152,7 @@ class TestOptim(unittest.TestCase):
    dtypes.default_float = old_default_float

  def test_assert_tensor_train(self):
-    t = Tensor.ones((1,1))
+    t = Tensor.ones((1,1), requires_grad=True)
    optimizer = Adam([t])
    optimizer.zero_grad()
    old_state = Tensor.training
@@ -167,7 +165,7 @@ class TestOptim(unittest.TestCase):

  def test_lamb_cpu_offload(self):
    # test that LAMB works when optimizer params (m, v, b1_t, b2_t) are moved to CPU
-    t = Tensor(x_init.copy())
+    t = Tensor(x_init.copy(), requires_grad=True)
    opt = LAMB([t])
    # move optimizer state to CPU
    for p in opt.m + opt.v + [opt.b1_t, opt.b2_t]: p.to_("CPU")
@@ -180,7 +178,7 @@ class TestOptim(unittest.TestCase):
  @needs_second_gpu
  def test_lamb_cpu_offload_multi(self):
    ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
-    t = Tensor(x_init.copy()).shard(ds, axis=1)
+    t = Tensor(x_init.copy(), requires_grad=True).shard(ds, axis=1)
    ds = t.device
    opt = LAMB([t])
    # move optimizer state to CPU
@@ -67,7 +67,7 @@ class TestPickle(unittest.TestCase):

  # NOTE: currently Buffer exists on the uop, not tensor
  def test_pickle_buffer_uop(self):
-    t = Tensor.arange(4).clone().realize()
+    t = Tensor.arange(4).realize()
    a = t.uop
    assert a.is_realized
    self.assertIsNotNone(buffer:=a.base.realized)
@@ -95,7 +95,7 @@ class TestPickle(unittest.TestCase):
    np.testing.assert_equal(vt2.numpy(), 20)

  def test_pickle_buffer_view(self):
-    t = Tensor.arange(10).clone(device="CPU").realize()
+    t = Tensor.arange(10, device="CPU").contiguous().realize()
    vt = t[3:5].contiguous().realize()
    assert hasattr(vt.uop.buffer, 'base')
    ref_value = vt.tolist()
@@ -1,6 +1,6 @@
 import unittest, struct, contextlib, statistics, gc
 from tinygrad import Device, Tensor, dtypes, TinyJit
-from tinygrad.helpers import DEV, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup
+from tinygrad.helpers import CI, DEV, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup
 from tinygrad.device import Buffer, BufferSpec, Compiled, ProfileDeviceEvent, ProfileGraphEvent
 from tinygrad.runtime.support.hcq import HCQCompiled
 from tinygrad.engine.realize import get_runtime
@@ -144,8 +144,7 @@ class TestProfiler(unittest.TestCase):
    assert len(graph_evs) == 2, "2 graph events are expected"
    assert len(graph_evs[0].ents) == 2, "two entities are expected"

-  @unittest.skipIf(MOCKGPU, "skip MOCKGPU")
-  @unittest.skipUnless(issubclass(type(Device[Device.DEFAULT]), HCQCompiled), "must be HCQ")
+  @unittest.skipIf(CI or not issubclass(type(Device[Device.DEFAULT]), HCQCompiled), "skip CI")
  def test_dev_jitter_matrix(self):
    dev_cnt = 6
    try: devs = [Device[f"{Device.DEFAULT}:{i}"] for i in range(dev_cnt)]
@@ -1,15 +1,15 @@
 import unittest, math
+from functools import partial

-from tinygrad import dtypes, Tensor, Device
-from tinygrad.helpers import getenv, DEV
+from tinygrad import nn, dtypes, Tensor, Device, TinyJit, Variable
+from tinygrad.helpers import getenv, CI, OSX
+from tinygrad.device import is_dtype_supported
 from tinygrad.codegen import to_program

 from tinygrad.uop.ops import Ops
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
-from tinygrad.renderer.isa.x86 import X86Renderer
 from test.helpers import not_support_multi_device, needs_second_gpu
-from test.unit.test_randomness import equal_distribution, normal_test

 import numpy as np
 import torch
@@ -18,6 +18,54 @@ from hypothesis import given, settings, strategies as strat
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")

+# https://gist.github.com/devries/11405101
+def ksprob(a):
+  fac, total, termbf = 2.0, 0.0, 0.0
+  a2 = -2.0 * a * a
+  for j in range(1, 101):
+    term = fac * math.exp(a2 * j * j)
+    total += term
+    if math.fabs(term) <= 0.001 * termbf or math.fabs(term) <= 1e-8 * total:
+      return total
+    fac = -fac
+    termbf = math.fabs(term)
+  return 1.0
+
+def kstest(l1, l2):
+  n1, n2 = len(l1), len(l2)
+  l1.sort()
+  l2.sort()
+  j1, j2, d, fn1, fn2 = 0, 0, 0.0, 0.0, 0.0
+  while j1 < n1 and j2 < n2:
+    d1, d2 = l1[j1], l2[j2]
+    if d1 <= d2:
+      fn1 = (float(j1) + 1.0) / float(n1)
+      j1 += 1
+    if d2 <= d1:
+      fn2 = (float(j2) + 1.0) / float(n2)
+      j2 += 1
+    dtemp = math.fabs(fn2 - fn1)
+    if dtemp > d:
+      d = dtemp
+  ne = float(n1 * n2) / float(n1 + n2)
+  nesq = math.sqrt(ne)
+  prob = ksprob((nesq + 0.12 + 0.11 / nesq) * d)
+  return prob
+
+def equal_distribution(tiny_func, torch_func=None, numpy_func=None, shape=(40, 43), alpha=0.04):
+  Tensor.manual_seed(1337)
+  torch.manual_seed(1337)
+  np.random.seed(1337)
+  assert not (torch_func is None and numpy_func is None), "no function to compare with"
+  x1 = tiny_func(*shape).numpy().flatten()
+  x2 = tiny_func(shape).numpy().flatten()
+  if numpy_func is not None: y = numpy_func(shape).flatten()
+  if torch_func is not None: z = torch_func(shape).numpy().flatten()
+  return (numpy_func is None or (kstest(x1, y) >= alpha and kstest(x2, y) >= alpha)) and \
+    (torch_func is None or (kstest(x1, z) >= alpha and kstest(x2, z) >= alpha))
+
+def normal_test(func, shape=(20, 45), alpha=0.05): return equal_distribution(func, numpy_func=lambda x: np.random.randn(*x), shape=shape, alpha=alpha)
+
 class TestRandomness(unittest.TestCase):
  def test_rand(self):
    self.assertFalse(normal_test(Tensor.rand))
@@ -37,7 +85,7 @@ class TestRandomness(unittest.TestCase):
    self.assertTrue(r1.uop.is_realized, "tensor should be realized after .realize()")
    self.assertTrue(r2.uop.is_realized, "tensor should be realized after .realize()")

-  @unittest.skipUnless(dtypes.float16 in Device[Device.DEFAULT].renderer.supported_dtypes(), "need float16 support")
+  @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need float16 support")
  def test_rand_float16(self):
    N = 128
    x = Tensor.rand((2, N, N), dtype=dtypes.float16)
@@ -48,7 +96,7 @@ class TestRandomness(unittest.TestCase):
    assert nx[nx == 0].size > 0
    equal_distribution(lambda *x: Tensor.rand(*x, dtype=dtypes.float16), torch.rand, lambda x: np.random.rand(*x), shape=(2, N, N))

-  @unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}, "gpuocelot doesn't support certain ops needed for threefry")
+  @unittest.skipIf(CI and Device.DEFAULT in {"NV", "CUDA"}, "gpuocelot doesn't support certain ops needed for threefry")
  def test_threefry_against_reference(self):
    Tensor.manual_seed(1337)

@@ -70,7 +118,6 @@ class TestRandomness(unittest.TestCase):
    np.testing.assert_allclose(jr, r)

  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "PTX and NIR use pointer arithmetic")
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "X86 callee saved registers have ulong dtype")
  def test_threefry_doesnt_use_long(self):
    linear = Tensor.rand(20).schedule_linear()
    for call in linear.src:
@@ -162,7 +209,7 @@ class TestRandomness(unittest.TestCase):
      if not (x.src[0] == y.src[0]):
        print(f"{x.src[0]} != {y.src[0]}")

-  @unittest.skipUnless(dtypes.bfloat16 in Device[Device.DEFAULT].renderer.supported_dtypes(), "need bfloat16 support")
+  @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "need bfloat16 support")
  def test_rand_bfloat16(self):
    N = 128
    x = Tensor.rand((2, N, N), dtype=dtypes.bfloat16)
@@ -226,12 +273,17 @@ class TestRandomness(unittest.TestCase):
    assert rand.dtype == dtypes.float16
    assert rand.device == empty.device

+  def test_randn(self):
+    self.assertEqual(Tensor.randn(3,3,dtype=dtypes.half).dtype, dtypes.half)
+    self.assertTrue(normal_test(Tensor.randn))
+    self.assertTrue(equal_distribution(Tensor.randn, torch.randn, lambda x: np.random.randn(*x)))
+
  def test_randn_device(self):
    self.assertEqual(Tensor.randn(3,3,device="CPU").device, "CPU")

  @given(strat.sampled_from([dtypes.float, dtypes.float16, dtypes.bfloat16]))
  def test_randn_finite(self, default_float):
-    if default_float not in Device[Device.DEFAULT].renderer.supported_dtypes(): return
+    if not is_dtype_supported(default_float): return
    old_default_float = dtypes.default_float
    # low precision can result in inf from randn
    dtypes.default_float = default_float
@@ -243,6 +295,128 @@ class TestRandomness(unittest.TestCase):
    assert math.isfinite(mn), mn
    dtypes.default_float = old_default_float

+  def test_randint(self):
+    self.assertFalse(normal_test(Tensor.randint))
+    self.assertTrue(equal_distribution(partial(Tensor.randint, low=-2, high=5),
+                                       numpy_func=lambda x: np.random.randint(low=-2, high=5, size=x)))
+    self.assertTrue(equal_distribution(partial(Tensor.randint, low=-2, high=5, dtype="int32"),
+                                       numpy_func=lambda x: np.random.randint(low=-2, high=5, size=x)))
+    self.assertTrue(Tensor.randint(1, device="CPU").device=="CPU")
+    # check types of args
+    with self.assertRaises(TypeError): Tensor.randint((3, 4), low=0.1, high=3)
+    with self.assertRaises(TypeError): Tensor.randint((3, 4), low=0, high=3.5)
+    with self.assertRaises(TypeError): Tensor.randint((3, 4), low=1, high=3, dtype="float")
+    with self.assertRaises(TypeError): Tensor.randint((3, 4), low=0, high=3, dtype=dtypes.float32)
+    # check low < high
+    with self.assertRaises(ValueError): Tensor.randint((3, 4), low=10, high=5)
+    with self.assertRaises(ValueError): Tensor.randint((3, 4), low=10, high=10)
+    np.testing.assert_array_equal(Tensor.randint(16, low=5, high=6).numpy(), 5)
+
+  def test_normal(self):
+    self.assertTrue(normal_test(Tensor.normal))
+    self.assertTrue(equal_distribution(Tensor.normal, lambda x: torch.nn.init.normal_(torch.empty(x), mean=0, std=1),
+                                                      lambda x: np.random.normal(loc=0, scale=1, size=x)))
+    # check std >= 0
+    with self.assertRaises(ValueError): Tensor.normal((3, 4), mean=0, std=-1)
+
+  def test_uniform(self):
+    self.assertFalse(normal_test(Tensor.uniform))
+    self.assertTrue(equal_distribution(Tensor.uniform, lambda x: torch.nn.init.uniform_(torch.empty(x)), lambda x: np.random.uniform(size=x)))
+    self.assertTrue(equal_distribution(partial(Tensor.uniform, low=-100, high=100, dtype=dtypes.int32),
+                                       numpy_func=lambda x: np.random.randint(low=-100, high=100, size=x)))
+    # check low < high
+    with self.assertRaises(ValueError): Tensor.uniform((3, 4), low=5.0, high=3.0)
+    with self.assertRaises(ValueError): Tensor.uniform((3, 4), low=1.0, high=1.0)
+
+  def test_scaled_uniform(self):
+    self.assertFalse(normal_test(Tensor.scaled_uniform))
+    self.assertTrue(equal_distribution(Tensor.scaled_uniform, lambda x: torch.nn.init.uniform_(torch.empty(x), a=-1, b=1) / math.sqrt(math.prod(x)),
+                                                              lambda x: np.random.uniform(-1, 1, size=x) / math.sqrt(math.prod(x))))
+
+  def test_glorot_uniform(self):
+    self.assertFalse(normal_test(Tensor.glorot_uniform))
+    self.assertTrue(equal_distribution(Tensor.glorot_uniform, lambda x: torch.nn.init.xavier_uniform_(torch.empty(x)),
+                                                              lambda x: np.random.uniform(-1, 1, size=x) * math.sqrt(6 / (x[0] + math.prod(x[1:])))))
+
+  def test_kaiming_uniform(self):
+    for shape in [(32, 16, 3, 3), (20, 44), (5, 15, 35)]:
+      self.assertTrue(equal_distribution(Tensor.kaiming_uniform, lambda x: torch.nn.init.kaiming_uniform_(torch.empty(x)), shape=shape))
+
+  def test_kaiming_normal(self):
+    for shape in [(32, 16, 3, 3), (20, 44), (3, 15, 35)]:
+      self.assertTrue(equal_distribution(Tensor.kaiming_normal, lambda x: torch.nn.init.kaiming_normal_(torch.empty(x)), shape=shape))
+
+  def test_multinomial(self):
+    self.assertRaises(AssertionError, lambda: Tensor(2).multinomial(1, replacement=False))
+    self.assertRaises(AssertionError, lambda: Tensor([1, 9]).multinomial(0, replacement=False))
+    def _check_with_torch(w, num_samples, replacement):
+      tiny_res = Tensor(w).multinomial(num_samples, replacement=replacement)
+      torch_res = torch.tensor(w).multinomial(num_samples, replacement=replacement)
+      self.assertEqual(tiny_res.shape, torch_res.shape)
+      if torch_res.ndim == 1:
+        tiny_res = tiny_res.unsqueeze(0)
+        torch_res = torch_res.unsqueeze(0)
+      for i in range(torch_res.shape[0]):
+        self.assertTrue(equal_distribution(lambda *_: tiny_res[i], lambda _: torch_res[i]))
+    _check_with_torch(w=[0.231, 0., 1., 0.5], num_samples=300, replacement=True)
+    _check_with_torch(w=[[0.2, 0.8]], num_samples=300, replacement=True)  # 2D but only 1 row
+    _check_with_torch(w=[[0.453, 0., 1., 0.81], [0.1, 0.8, 0., 0.1]], num_samples=300, replacement=True)
+    # no-replacement
+    w = [0.1, 0.9]
+    self.assertRaises(AssertionError, lambda: Tensor(w).multinomial(100, replacement=False))
+
+    @TinyJit
+    def sample_one(): return Tensor(w).multinomial(1, replacement=False).realize()
+
+    tiny_samples = [sample_one().item() for _ in range(1000)]
+    torch_samples = [torch.tensor(w).multinomial(1, replacement=False).item() for _ in range(1000)]
+    self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_samples), lambda _: torch.tensor(torch_samples)))
+
+    w = list(range(32))
+    s1 = Tensor(w).multinomial(5, replacement=False).numpy()
+    self.assertEqual(len(set(s1.tolist())), 5)
+    s2 = Tensor(w).multinomial(5, replacement=False).numpy()
+    self.assertFalse(np.array_equal(s1, s2))
+    full = Tensor(w).multinomial(len(w), replacement=False).numpy()
+    self.assertEqual(sorted(full.tolist()), w)
+
+    w = [0.1, 0.2, 0.3, 0.4]
+    @TinyJit
+    def sample_three(): return Tensor(w).multinomial(3, replacement=False).realize()
+
+    tiny_draws = np.array([sample_three().numpy() for _ in range(1000)])
+    torch_draws = np.array([torch.tensor(w).multinomial(3, replacement=False).numpy() for _ in range(1000)])
+    for pos in range(3):
+      self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_draws[:, pos]), lambda _: torch.tensor(torch_draws[:, pos])))
+
+  @unittest.skip("this test is flaky")
+  def test_multinomial_counterexample(self):
+    tiny_res = Tensor([0.3, 0.6, 0.1]).multinomial(4000, replacement=True)
+    torch_res = torch.tensor([0.3, 0.6, 0.1]).multinomial(4000, replacement=True)
+    self.assertTrue(equal_distribution(lambda *_: tiny_res, lambda _: torch_res))
+    torch_res = torch.tensor([0.2, 0.7, 0.1]).multinomial(4000, replacement=True)
+    self.assertFalse(equal_distribution(lambda *_: tiny_res, lambda _: torch_res))
+
+  def test_conv2d_init(self):
+    params = (128, 256, (3,3))
+    assert equal_distribution(lambda *_: nn.Conv2d(*params).weight, lambda _: torch.nn.Conv2d(*params).weight.detach())
+    assert equal_distribution(lambda *_: nn.Conv2d(*params).bias, lambda _: torch.nn.Conv2d(*params).bias.detach())
+
+  def test_linear_init(self):
+    params = (64, 256)
+    assert equal_distribution(lambda *_: nn.Linear(*params).weight, lambda _: torch.nn.Linear(*params).weight.detach())
+    assert equal_distribution(lambda *_: nn.Linear(*params).bias, lambda _: torch.nn.Linear(*params).bias.detach())
+
+  def test_bn_init(self):
+    params = (64,)
+    assert equal_distribution(lambda *_: nn.BatchNorm2d(*params).weight, lambda _: torch.nn.BatchNorm2d(*params).weight.detach())
+    assert equal_distribution(lambda *_: nn.BatchNorm2d(*params).bias, lambda _: torch.nn.BatchNorm2d(*params).bias.detach())
+
+  def test_rand_chain(self):
+    # NOTE: this fails if property propagates deeper than stack limit
+    for _ in range(833): Tensor.rand(1)
+    Tensor.rand(1).realize()
+
  def test_random_counter_overflow(self):
    device = Device.DEFAULT
    Tensor.manual_seed(1337)
@@ -258,5 +432,20 @@ class TestRandomness(unittest.TestCase):
    c = Tensor._device_rng_counters[device].numpy()
    np.testing.assert_allclose(c, [14, 1])

+# TODO: still fails with MAX_KERNEL_BUFFERS
+@unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
+class TestSample(unittest.TestCase):
+  def test_sample(self):
+    X = Tensor.rand(1000, 50).realize()
+    BS = 16
+    idxs = np.random.randint(0, X.shape[0], size=(BS))
+    # this uncovered a bug with arg sort order
+    batch = [Variable(f'idx{i}', 0, X.shape[0]-1).bind(s) for i,s in enumerate(idxs.tolist())]
+    x = Tensor.cat(*[X.shrink(((batch[i], batch[i]+1), None)) for i in range(BS)])
+    print(idxs)
+    ret = x.numpy()
+    base = X.numpy()[idxs]
+    np.testing.assert_equal(ret, base)
+
 if __name__ == "__main__":
  unittest.main()
@@ -109,9 +109,9 @@ def fa():
 def fa_bw():
  Tensor.manual_seed(1337)
  with Context(DEBUG=0):
-    q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize() for _ in range(3)]
+    q,k,v = [Tensor.rand(BS, HEADS, SEQLEN, EMB).contiguous().realize().requires_grad_() for _ in range(3)]
    attn_output = nn.Linear(HEADS*EMB, HEADS*EMB, bias=False)
-    attn_output.weight.realize()
+    attn_output.weight.requires_grad_().realize()
    target = Tensor.rand(BS, SEQLEN, HEADS*EMB).contiguous().realize()

  GlobalCounters.reset()
@@ -1,6 +1,6 @@
 import unittest
 import numpy as np
-from tinygrad.device import Device
+from tinygrad.device import Device, is_dtype_supported
 from tinygrad.dtype import dtypes, ConstType
 from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
@@ -22,8 +22,8 @@ def _test_uop_result(inputs:list[Tensor], sink:UOp, local_size=None):

 def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
  dtype = alu_src_uops[0].dtype
-  a = UOp.param(0, dtype.ptr())
-  b = UOp.param(1, dtype.ptr())
+  a = UOp(Ops.PARAM, dtype.ptr(), (), 0)
+  b = UOp(Ops.PARAM, dtype.ptr(), (), 1)
  idx = UOp.const(dtypes.int, 0)
  ld = b.index(idx)
  alu = ld.alu(alu_op, *alu_src_uops)
@@ -33,7 +33,7 @@ def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
 class TestRendererFailures(unittest.TestCase):
  @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
  def test_gated_store_with_alu(self):
-    a = UOp.param(0, dtypes.int.ptr())
+    a = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
    gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
    gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index(lidx0.valid(gate_alu)), UOp.const(dtypes.int, 1)))
    sink = UOp(Ops.SINK, dtypes.void, (gated_alu_store,), arg=KernelInfo())
@@ -42,7 +42,7 @@ class TestRendererFailures(unittest.TestCase):

  @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
  def test_gated_store_with_alu_2d(self):
-    a = UOp.param(0, dtypes.int.ptr())
+    a = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
    gate_alu_0 = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
    gate_alu_1 = (lidx1:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 2),), 'lidx1')).ne(0)
    gated_alu_store = UOp(Ops.STORE, dtypes.void, (a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(dtypes.int, 1)))
@@ -87,7 +87,7 @@ class TestWGSLFailures(unittest.TestCase):
 class TestPTXFailures(unittest.TestCase):
  @unittest.skip("INDEX can only have a gate ALU parent, not an IF")
  def test_gated_store_with_if(self):
-    a = UOp.param(0, dtypes.int.ptr())
+    a = UOp(Ops.PARAM, dtypes.int.ptr(), (), 0)
    gate_alu = (lidx0:=UOp(Ops.SPECIAL, dtypes.int, (UOp.const(dtypes.int, 4),), 'lidx0')).ne(0)
    val = UOp.const(dtypes.int, 1)
    if_uop = UOp(Ops.IF, dtypes.void, (gate_alu,))
@@ -96,7 +96,7 @@ class TestPTXFailures(unittest.TestCase):
    ret = _test_uop_result([], sink, local_size=[4, 1, 1])[0]
    np.testing.assert_equal(ret, [0, 1, 1, 1])

-  @unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half")
+  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
  def test_gated_define_acc_with_half_dtype(self):
    a = Tensor.randn(32, 32, dtype=dtypes.half).realize()
    b = Tensor.randn(34, 32, dtype=dtypes.half).realize()
@@ -5,16 +5,15 @@
 import gc, unittest, functools
 import numpy as np
 from typing import cast
-from hypothesis import assume, given, strategies as strat
+from hypothesis import assume, given, settings, strategies as strat

 from tinygrad import nn, dtypes, Device, Tensor, Variable
+from tinygrad.device import is_dtype_supported
 from tinygrad.dtype import DType
 from tinygrad.uop.ops import UOp, Ops, UPat
-from tinygrad.helpers import DEBUG, DEV, OSX, GlobalCounters, Context, getenv, all_same, temp
+from tinygrad.helpers import CI, DEBUG, OSX, GlobalCounters, Context, getenv, all_same, temp
 from tinygrad.engine.realize import compile_linear, run_linear

-supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
-
 class KernelCountException(Exception): pass
 def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True):
  if to_prerealize:
@@ -45,8 +44,8 @@ def _test_conv2d(allowed:int, dtype:DType=dtypes.float):
  dtypes.default_float = dtype
  Tensor.manual_seed(0)
  BS, CIN = 2, 3
-  img = Tensor.randn(BS, CIN, 64, 64).realize()
-  w = Tensor.uniform(16, CIN, 3, 3).realize()
+  img = Tensor.randn(BS, CIN, 64, 64, requires_grad=True).realize()
+  w = Tensor.uniform(16, CIN, 3, 3, requires_grad=True).realize()
  ret = Tensor.conv2d(img, w).relu().mean().backward()
  dtypes.default_float = old_default_float
  linear, var_vals = Tensor.linear_with_vars(ret, img.grad, w.grad)
@@ -82,6 +81,13 @@ class TestSchedule(unittest.TestCase):
  def test_arange_avgpool2d_fused_noopt(self):
    with Context(NOOPT=1): self.test_arange_avgpool2d(kcount=1)

+  # linearizer error
+  @unittest.skip("recursion error no longer raised")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "needs supports_float4 to fail")
+  def test_arange_avgpool2d_fused(self):
+    with self.assertRaises(RecursionError):
+      with Context(NOOPT=0): self.test_arange_avgpool2d(kcount=1)
+
  # when we're fusing a reduce, all ReduceOps must have the same N in the dimensions
  # all permutes, reshapes, expands and shrinks push through the reduce
  def test_arange_sum(self):
@@ -99,7 +105,7 @@ class TestSchedule(unittest.TestCase):
    run_linear(*check_schedule(a, 1))
    self.assertListEqual(a.tolist(), [[15]])

-  @unittest.skipUnless(dtypes.half in supported_dtypes, "need half")
+  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and OSX, "WEBGPU Metal backend is not accurate enough")
  def test_expand_buffer_before_cast(self):
    a = Tensor.randn(4, 2, 1).realize().permute((1, 0, 2))
@@ -107,6 +113,13 @@ class TestSchedule(unittest.TestCase):
    run_linear(*check_schedule(b, 1))
    np.testing.assert_allclose(b.numpy(), np.broadcast_to(a.numpy().astype(np.float16), (2, 4, 4))+2, rtol=1e-3)

+  def test_indexing_scalars_simple(self):
+    X = Tensor.randn(2, 2).realize()
+    xt = X[Tensor(1)][Tensor(0)]
+    run_linear(*check_schedule(xt, 1))
+    np.testing.assert_equal(xt.numpy(), X.numpy()[1][0])
+
+  @unittest.skipIf(CI and Device.DEFAULT == "NV", "crashes on NV CI")
  def test_add_chain_buffers(self):
    N = 31
    with Context(TRACK_MATCH_STATS=0, DEBUG=0):
@@ -117,14 +130,14 @@ class TestSchedule(unittest.TestCase):
        root = root + functools.reduce(lambda a,b:a+b, bufs[i:i+X])
      self.assertEqual(root.item(), sum(range(N)))

-  def test_indexing_scalars(self):
-    # cover each shape at all index corners
-    for x, y in [(2,2), (2,3), (3,2), (3,3)]:
-      for a, b in [(0,0), (0,y-1), (x-1,0), (x-1,y-1)]:
-        X = Tensor.randn(x, y).realize()
-        xt = X[Tensor(a)][Tensor(b)]
-        run_linear(*check_schedule(xt, 1))
-        np.testing.assert_equal(xt.numpy(), X.numpy()[a][b])
+  @given(strat.sampled_from(range(2,4)), strat.sampled_from(range(2,4)), strat.sampled_from(range(0,4)), strat.sampled_from(range(0,4)))
+  @settings(deadline=None)
+  def test_indexing_scalars(self, x, y, a, b):
+    assume(a<x and b<y)
+    X = Tensor.randn(x, y).realize()
+    xt = X[Tensor(a)][Tensor(b)]
+    run_linear(*check_schedule(xt, 1))
+    np.testing.assert_equal(xt.numpy(), X.numpy()[a][b])

  def test_push_pads_elementwise(self):
    x = Tensor.full((4,4), 2.).contiguous().realize()
@@ -140,6 +153,11 @@ class TestSchedule(unittest.TestCase):
    run_linear(*check_schedule(z, 1, [x,y]))
    self.assertEqual(z.item(), 32)

+  def test_constants_can_store(self):
+    a = Tensor(2).contiguous()
+    run_linear(*check_schedule(a, 1))
+    np.testing.assert_equal(a.numpy(), 2)
+
  def test_allow_push_permutes(self):
    a = Tensor.randn(10,10,10).realize()
    b = Tensor.randn(10,10,1).realize()
@@ -220,9 +238,19 @@ class TestSchedule(unittest.TestCase):
      run_linear(*check_schedule(out, 4))
    np.testing.assert_allclose(out.numpy(), (x.numpy() - x.numpy().max(keepdims=True)).max())

+  @unittest.skip("these two Tensors are the same")
+  def test_example_matmul(self):
+    x = Tensor.eye(64, requires_grad=True)
+    y = Tensor.eye(64, requires_grad=True)
+    z = y.matmul(x).sum()
+    z.backward()
+    out = x.grad.contiguous()
+    run_linear(*check_schedule(out, 1))
+    np.testing.assert_allclose(out.numpy(), np.ones((64,64)))
+
  def test_example_matmul_contig(self):
-    x = Tensor.eye(64).clone().realize()
-    y = Tensor.eye(64).clone().realize()
+    x = Tensor.eye(64, requires_grad=True).contiguous().realize()
+    y = Tensor.eye(64, requires_grad=True).contiguous().realize()
    z = y.matmul(x).sum()
    z.backward()
    out = x.grad.contiguous()
@@ -230,7 +258,7 @@ class TestSchedule(unittest.TestCase):
    np.testing.assert_allclose(out.numpy(), np.ones((64,64)))

  def test_example_matmul_same(self):
-    x = Tensor.eye(64)
+    x = Tensor.eye(64, requires_grad=True)
    z = x.matmul(x).sum()
    z.backward()
    out = x.grad.contiguous()
@@ -695,15 +723,15 @@ class TestSchedule(unittest.TestCase):
    np.testing.assert_equal(d.numpy(), np.pad(np.exp2(a.numpy())[:, None, :], ((0, 0), (1, 1), (0, 0)))*2)

  def test_fuse_arange_pad_replicate_mode(self):
-    x = Tensor.empty(3,3,3,3)
+    x = Tensor.empty(3,3,3,3, requires_grad=True)
    y = x.pad((-1,2,2,-1), mode="replicate")
    dx = y.sum().gradient(x)[0]
-    sched = check_schedule(dx, 0)
+    sched = check_schedule(dx, 1)
    run_linear(*sched)
    np.testing.assert_allclose(dx.numpy(), [[[[0.,3.,9.],[0,1.,3.],[0.,0.,0.]]]*3]*3)

  # TODO like openpilot with imagef
-  @unittest.skipUnless(dtypes.half in supported_dtypes, "need half")
+  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
  def test_base_change_expand_expand(self):
    a = Tensor.ones(4, 4).contiguous().realize()
    b = a.cast(dtypes.half).expand(2, 4, 4)
@@ -749,9 +777,9 @@ class TestSchedule(unittest.TestCase):
  def test_conv2d(self): _test_conv2d(4)
  def test_conv2d_fused(self): _test_conv2d(4)

-  @unittest.skipUnless(dtypes.half in supported_dtypes, "need half")
+  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
  def test_conv2d_half(self): _test_conv2d(4, dtype=dtypes.half)
-  @unittest.skipUnless(dtypes.half in supported_dtypes, "need half")
+  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "Causes other tests to fail")
  def test_conv2d_fused_half(self): _test_conv2d(4, dtype=dtypes.half)

@@ -763,6 +791,7 @@ class TestSchedule(unittest.TestCase):
    gc.collect()
    self.assertEqual(GlobalCounters.mem_used-base, 1024)

+  @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
  def test_image_dot_f16_fusion(self):
    with Context(FLOAT16=1, OPENPILOT_HACKS=1):
      def cnt():
@@ -774,6 +803,7 @@ class TestSchedule(unittest.TestCase):
      with Context(IMAGE=1):
        self.assertEqual(cnt(), 5)

+  @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
  def test_image_f16_residual_fusion(self):
    with Context(FLOAT16=1, OPENPILOT_HACKS=1):
      def cnt():
@@ -789,6 +819,7 @@ class TestSchedule(unittest.TestCase):
      with Context(IMAGE=1):
        self.assertEqual(cnt(), 9)

+  @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
  def test_image_conv_fusion(self):
    with Context(OPENPILOT_HACKS=1):
      def cnt():
@@ -826,7 +857,7 @@ class TestSchedule(unittest.TestCase):
    self._test_fusion([(32, 32)], lambda a:a-a.sum(1), 2)

  def test_cast_padded_view(self):
-    a = Tensor.arange(4).reshape(1, 4).clone().realize()
+    a = Tensor.arange(4).reshape(1, 4)
    casted_view = a.pad(((0, 1), (0, 0))).cast(dtypes.float)
    casted_view.realize()
    self.assertEqual(casted_view.uop.base.realized.size, 8)
@@ -836,7 +867,7 @@ class TestSchedule(unittest.TestCase):

  # NOTE: we only reorder CAST if it's an EXPAND
  def test_cast_after_shrink(self):
-    a = Tensor.arange(4).reshape(1, 4).clone().realize()
+    a = Tensor.arange(4).reshape(1, 4)
    casted_view = a.shrink(((0, 1), (0, 2))).cast(dtypes.float)
    casted_view.realize()
    self.assertEqual(casted_view.uop.base.realized.size, 2)
@@ -845,7 +876,7 @@ class TestSchedule(unittest.TestCase):
    self.assertListEqual(realized_view.tolist(), [[0, 1]])

  def test_cast_const_view(self):
-    a = Tensor.ones((4, 4), dtype=dtypes.float32, buffer=False)
+    a = Tensor.ones((4, 4), dtype=dtypes.float32)
    casted_view = a.cast(dtypes.int32)
    run_linear(*check_schedule(casted_view, 1))
    realized_const_view = casted_view.contiguous()
@@ -855,7 +886,7 @@ class TestSchedule(unittest.TestCase):
  @given(strat.sampled_from(dtypes.all), strat.sampled_from(dtypes.all))
  @unittest.skip("kernel count depends on input")
  def test_cast_padded_const(self, dt1, dt2):
-    assume(dt1 in supported_dtypes and dt2 in supported_dtypes)
+    assume(is_dtype_supported(dt1) and is_dtype_supported(dt2))
    a = Tensor(1, dtype=dt1).reshape(1, 1).pad(((1, 1), None))
    casted_view = a.cast(dt2)
    run_linear(*check_schedule(casted_view, 0))
@@ -925,7 +956,7 @@ class TestSchedule(unittest.TestCase):
    np.testing.assert_equal(a.numpy(), (np.arange(4)*x.numpy()).T.sum())

  def test_div_padded_arange(self):
-    x = Tensor.full((2,2), 16, buffer=False)
+    x = Tensor.full((2,2), 16)
    y = x.div(Tensor.linspace(2, 8, steps=4, dtype=dtypes.int).reshape(2,2), rounding_mode="trunc").pad(((1,1), (1,1)))
    out = y.sum(axis=1)
    run_linear(*check_schedule(out, 1))
@@ -951,7 +982,7 @@ class TestSchedule(unittest.TestCase):
  def test_arange_index_contiguous(self):
    Tensor.manual_seed(0)
    x = Tensor.randn(5, 2).realize()
-    a = Tensor.arange(10).clone()
+    a = Tensor.arange(10).contiguous()
    out = (x + a[2]).sum()
    run_linear(*check_schedule(out, 2))
    np.testing.assert_allclose(out.numpy(), (x.numpy()+np.arange(10)[2]).sum(), atol=1e-5, rtol=1e-6)
@@ -967,12 +998,19 @@ class TestSchedule(unittest.TestCase):
  def test_user_contiguous(self):
    Tensor.manual_seed(0)
    x = Tensor.randn(5, 2).realize()
-    a = (Tensor.arange(10)+1).clone()
+    a = (Tensor.arange(10)+1).contiguous()
    out = (x + a[2]).sum()
    run_linear(*check_schedule(out, 2))
    np.testing.assert_allclose(out.numpy(), (x.numpy()+(np.arange(10)+1)[2]).sum(), atol=1e-5, rtol=1e-6)

-  @unittest.skipUnless(dtypes.half in supported_dtypes, "need half")
+  @unittest.skip("BUFFER_VIEW no longer supported on non-disk devices")
+  def test_arange_view_op(self):
+    a = Tensor.arange(12).reshape(4, 3).shrink(((1, 2), (1, 3))).contiguous()
+    sched = run_linear(*check_schedule(a, 1))
+    self.assertIs(sched[1].ast.op, Ops.BUFFER_VIEW)
+    np.testing.assert_equal(a.numpy(), [[4, 5]])
+
+  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
  def test_precompute_freqs_cis(self):
    from extra.models.llama import precompute_freqs_cis
    args = {"dim":32, "end":2048, "theta":10000}
@@ -986,12 +1024,12 @@ class TestSchedule(unittest.TestCase):
  def test_fuse_assign_contiguous(self):
    x = Tensor.zeros(4, 4, dtype=dtypes.int).contiguous().realize()
    a = Tensor.arange(8).reshape(4, 2)
-    run_linear(*check_schedule(x.shrink((None, (0, 2))).assign(a.clone()), 2))
+    run_linear(*check_schedule(x.shrink((None, (0, 2))).assign(a.contiguous()), 2))
    np.testing.assert_equal(x.numpy(), [[0, 1, 0, 0], [2, 3, 0, 0], [4, 5, 0, 0], [6, 7, 0, 0]])

  def test_assign_non_contiguous_alt(self): self.test_assign_non_contiguous(alt=True)
  def test_assign_non_contiguous(self, alt=False):
-    x = (Tensor.arange(16)-100).reshape(4,4).clone().realize()
+    x = (Tensor.arange(16)-100).reshape(4,4).contiguous().realize()
    xref = x.numpy()
    if alt:
      y = Tensor.randint(2, 4).contiguous().realize()
@@ -1007,7 +1045,7 @@ class TestSchedule(unittest.TestCase):
    np.testing.assert_equal(tst.numpy(), a.numpy())

  def test_setitem_sched(self, mop=lambda x:x, expected_kcount=1):
-    a = Tensor.arange(16).reshape(4, 4).clone(device="CPU").realize()
+    a = Tensor.arange(16, device="CPU").reshape(4, 4).contiguous().realize()
    a2 = mop(a)
    expected = (a+a2).tolist()
    a.assign(a+a2)
@@ -1021,7 +1059,7 @@ class TestSchedule(unittest.TestCase):

  def test_setitem_const_fused(self):
    # https://github.com/tinygrad/tinygrad/issues/10690
-    a = Tensor.arange(16).clone().realize()
+    a = Tensor.arange(16).contiguous().realize()
    GlobalCounters.reset()
    a[4] = 3
    self.assertEqual(GlobalCounters.kernel_count, 0)
@@ -1031,7 +1069,7 @@ class TestSchedule(unittest.TestCase):

  def test_no_extra_contiguous_on_setitem_assign_back(self):
    # pattern: contiguous copy, advanced setitem, assign back (e.g. torch backend _view_write)
-    base = Tensor.arange(16).reshape(4, 4).clone()
+    base = Tensor.arange(16).reshape(4, 4).contiguous()
    flat_base = base.reshape(16).contiguous()
    idx = Tensor([1,2,5,6], dtype=dtypes.int32)
    flat_base[idx] = Tensor([99,99,99,99])
@@ -1058,10 +1096,11 @@ class TestSchedule(unittest.TestCase):
    self.assertEqual(a.tolist(), [0., 0.])
    self.assertEqual(b.tolist(), [False, False])

+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "Validation error on WebGPU")
  def test_mnist_val(self):
-    # from tinygrad.nn.datasets import mnist
+    from tinygrad.nn.datasets import mnist
    import torch
-    Y_train = Tensor.randint(60000, dtype='uchar').realize()
+    _, Y_train, _, _ = mnist()
    samples = Tensor.randint(BS:=getenv("BS", 512), high=cast(int,Y_train.shape[-1])).realize()
    yt = Tensor.randn(BS, 10).realize()
    loss = yt.sparse_categorical_crossentropy(Y_train[samples])
@@ -1094,7 +1133,7 @@ class TestSchedule(unittest.TestCase):
      self.assertListEqual(a.tolist(), [[1.]*shape[1]]*shape[0])

 class TestLimitBufs(unittest.TestCase):
-  @unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT == "NV", "crashes in ocelot")
+  @unittest.skipIf(CI and Device.DEFAULT == "NV", "crashes on NV CI")
  def test_limit_bufs_with_var(self):
    N = 31
    with Context(TRACK_MATCH_STATS=0, DEBUG=0):
@@ -1230,7 +1269,7 @@ class TestView(unittest.TestCase):
  # x collapses along with its children
  def test_parent_view_collapses(self):
    a = Tensor([1, 2])
-    b = Tensor.arange(3).clone()
+    b = Tensor.arange(3).contiguous()
    bv = b.pad(((0, 2),))[-2:]
    # this becomes a late a*0
    late_mul = a*bv
@@ -1247,7 +1286,7 @@ class TestView(unittest.TestCase):
  # as long as one child realizes, x does not collapse
  def test_parent_multiple_children_no_collapse(self):
    a = Tensor([1, 2])
-    b = Tensor.arange(3).clone()
+    b = Tensor.arange(3).contiguous()
    bv = b.pad(((0, 2),))[-2:]
    late_mul = a*bv
    other_child = b+2
@@ -1264,7 +1303,7 @@ class TestView(unittest.TestCase):
 class TestCopyFolding(unittest.TestCase):
  def test_const_copy_is_free(self):
    b = Tensor(1).to("CPU") * 4
-    run_linear(*check_schedule(b, 0, filter_sink=False))
+    run_linear(*check_schedule(b, 1, filter_sink=False))
    assert b.item() == 4

  def test_one_hot_with_copy(self):
@@ -1273,13 +1312,13 @@ class TestCopyFolding(unittest.TestCase):
    check_schedule(x, 3, filter_sink=False)

  def test_const_copy_multi(self):
-    x = Tensor.ones(1, device="CPU", buffer=False).to_(["CPU", "CPU:1"]) * 2
+    x = Tensor.ones(1, device="CPU").to_(["CPU", "CPU:1"]) * 2
    run_linear(*check_schedule(x, 2, filter_sink=False))
    self.assertEqual(x.item(), 2.0)

  def test_late_const_copy_folding(self):
-    a = Tensor.arange(3).clone().realize()
-    zeros = Tensor.zeros(3, buffer=False).realize()
+    a = Tensor.arange(3).realize()
+    zeros = Tensor.zeros(3).realize()
    b = (a*zeros).to("CPU") + 1
    run_linear(*check_schedule(b, 1, filter_sink=False))
    self.assertListEqual(b.tolist(), [1, 1, 1])
@@ -1353,14 +1392,14 @@ class TestCopyFolding(unittest.TestCase):
    self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])

  def test_permute_on_disk(self):
-    with open(temp('dt_arange_4_permute'), "wb") as f: f.write(Tensor.arange(4).clone().realize().uop.base.buffer.as_memoryview())
+    with open(temp('dt_arange_4_permute'), "wb") as f: f.write(Tensor.arange(4).realize().uop.base.buffer.as_memoryview())
    a = Tensor.empty(4, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_4_permute')}")
    b = a.reshape(2, 2).permute(1, 0).to("CPU")
    b.realize()
    self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])

  def test_permute_on_disk_contiguous(self):
-    with open(temp('dt_arange_4_permute_contig'), "wb") as f: f.write(Tensor.arange(4).clone().realize().uop.base.buffer.as_memoryview())
+    with open(temp('dt_arange_4_permute_contig'), "wb") as f: f.write(Tensor.arange(4).realize().uop.base.buffer.as_memoryview())
    a = Tensor.empty(4, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_4_permute_contig')}")
    b = a.reshape(2, 2).permute(1, 0).contiguous().to("CPU")
    b.realize()
@@ -1374,7 +1413,7 @@ class TestCopyFolding(unittest.TestCase):

  # NOTE: disk permute must come after COPY
  def test_permute_after_shrink_on_disk(self):
-    with open(temp('dt_arange_5_permute'), "wb") as f: f.write(Tensor.arange(5).clone().realize().uop.base.buffer.as_memoryview())
+    with open(temp('dt_arange_5_permute'), "wb") as f: f.write(Tensor.arange(5).realize().uop.base.buffer.as_memoryview())
    a = Tensor.empty(5, dtype=dtypes.int32, device=f"disk:{temp('dt_arange_5_permute')}")
    b = a.shrink(((0, 4),)).reshape(2, 2).permute(1, 0).to("CPU")
    b.realize()
@@ -32,14 +32,13 @@ class TestSetitem(unittest.TestCase):
    self.assertListEqual(t.tolist(), [0, 1, 11, 3, 11, 5, 6, 7, 8, 9])

  def test_setitem_inplace_mul(self):
-    t = Tensor.arange(10).clone().realize()
+    t = Tensor.arange(10).realize()
    t[:3] *= 10
    self.assertListEqual(t.tolist(), [0, 10, 20, 3, 4, 5, 6, 7, 8, 9])

-  @unittest.skip("crashed in LLVM CI")
  def test_setitem_fancy_on_unrealized_view(self):
    # fancy indexing setitem on unrealized SHRINK view (triggered infinite loop in graph_rewrite)
-    base = Tensor.arange(20, dtype=dtypes.float).reshape(4, 5).clone().realize()
+    base = Tensor.arange(20, dtype=dtypes.float).reshape(4, 5)
    sub = base[1:3]
    flat = sub.reshape(sub.numel()).contiguous()
    idx = Tensor([0, 3, 7, 9])
@@ -230,7 +229,7 @@ class TestSetitem(unittest.TestCase):
    np.testing.assert_equal(t.numpy(), n)

  def test_setitem_swap_rows(self):
-    t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).clone().realize()
+    t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).contiguous().realize()
    tmp = t[0]
    t[0] = t[1]
    t[2] = tmp
@@ -238,7 +237,7 @@ class TestSetitem(unittest.TestCase):
    np.testing.assert_allclose(t.numpy(), [[2, 3], [2, 3], [2, 3]])

    # eager version
-    t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).clone().realize()
+    t = Tensor.arange(6, dtype=dtypes.float).reshape(3, 2).contiguous().realize()
    tmp = t[0].realize()
    t[0] = t[1].realize()
    t[2] = tmp.realize()
@@ -270,8 +269,8 @@ class TestSetitem(unittest.TestCase):
  def test_cross_assign_independence(self):
    # when assigning to two tensors using computations from both,
    # both assigns should see the OLD values of both tensors
-    a = Tensor.arange(4, dtype=dtypes.float).clone().realize()
-    b = Tensor.arange(4, 8, dtype=dtypes.float).clone().realize()
+    a = Tensor.arange(4, dtype=dtypes.float).contiguous().realize()
+    b = Tensor.arange(4, 8, dtype=dtypes.float).contiguous().realize()
    new_a = a + b    # [4, 6, 8, 10]
    new_b = a * 2    # [0, 2, 4, 6] -- should use OLD a
    a.assign(new_a)
@@ -284,8 +283,8 @@ class TestSetitem(unittest.TestCase):
      np.testing.assert_allclose(b.numpy(), [8, 12, 16, 20])

    # eager version
-    a = Tensor.arange(4, dtype=dtypes.float).clone().realize()
-    b = Tensor.arange(4, 8, dtype=dtypes.float).clone().realize()
+    a = Tensor.arange(4, dtype=dtypes.float).contiguous().realize()
+    b = Tensor.arange(4, 8, dtype=dtypes.float).contiguous().realize()
    new_a = (a + b).realize()
    new_b = (a * 2).realize()
    a.assign(new_a).realize()
@@ -302,71 +301,49 @@ class TestSetitem(unittest.TestCase):
    self.assertListEqual(z[6:7].tolist(), [3])

 class TestWithGrad(unittest.TestCase):
-  def test_basic_setitem_works(self):
+  def test_no_requires_grad_works(self):
    z = Tensor.rand(8, 8)
    x = Tensor.rand(8)
    z[:3] = x

-  def test_set_backward(self):
+  def test_set_with_requires_grad(self):
    z = Tensor.ones(8, 8)
-    x = Tensor.rand(8, 8)
+    x = Tensor.rand(8, 8, requires_grad=True)
    z[:] = x
    z.sum().backward()
    np.testing.assert_allclose(x.grad.numpy(), np.ones((8, 8)))

-  def test_set_nonleaf_backward(self):
-    x = Tensor([1.0, 2.0, 3.0, 4.0])
+  def test_set_nonleaf_requires_grad(self):
+    x = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
    z = x * 2
    z[:2] = Tensor([10.0, 20.0])
    z.sum().backward()
    np.testing.assert_allclose(x.grad.numpy(), [0, 0, 2, 2])

-  def test_set_overlapping_backward(self):
-    z = Tensor.zeros(6)
-    x = Tensor.ones(4).contiguous()
-    y = Tensor.ones(4) * 2
+  def test_set_overlapping_requires_grad(self):
+    z = Tensor.zeros(6, requires_grad=True)
+    x = Tensor.ones(4, requires_grad=True)
+    y = Tensor.ones(4, requires_grad=True) * 2
    z[:4] = x
    z[2:] = y
    z.sum().backward()
    np.testing.assert_allclose(x.grad.numpy(), [1, 1, 0, 0])
    np.testing.assert_allclose(y.grad.numpy(), np.ones(4))

-  def test_set_iadd_backward(self):
-    z = Tensor([1.0, 2.0, 3.0, 4.0])
-    x = Tensor([10.0, 20.0])
+  def test_set_iadd_requires_grad(self):
+    z = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
+    x = Tensor([10.0, 20.0], requires_grad=True)
    z[:2] += x
    z.sum().backward()
    np.testing.assert_allclose(z.grad.numpy(), np.ones(4))
    np.testing.assert_allclose(x.grad.numpy(), np.ones(2))

  def test_set_used_before_setitem(self):
-    z = Tensor([1.0, 2.0, 3.0, 4.0])
+    z = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
    _ = z.sum()
    with self.assertRaises(RuntimeError):
      z[:2] = Tensor([0.0, 0.0])

-  def test_setitem_raises_with_unrealized_downstream(self):
-    x = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
-    _y = x * 2.0
-    with self.assertRaises(RuntimeError):
-      x[0] = 99.0
-
-  def test_setitem_raises_on_unrealized_compute_base(self):
-    # y has a compute (unrealized) base; tmp is a view of y. eager: tmp would follow y's mutation. lazy: tmp keeps the old MUL graph.
-    x = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
-    y = x * 2.0
-    _tmp = y[:1]
-    with self.assertRaises(RuntimeError):
-      y[0] = 99.0
-
-  def test_setitem_raises_on_aliased_uop(self):
-    # two Tensor objects sharing the exact same unrealized uop. setitem on one updates its uop, the other keeps the stale graph reference.
-    x = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
-    y = x * 2.0
-    _z = Tensor(y.uop)
-    with self.assertRaises(RuntimeError):
-      y[0] = 99.0
-
 class TestSetitemLoop(unittest.TestCase):
  def test_arange(self):
    N = 10
@@ -4,6 +4,7 @@ from tinygrad import Tensor, GlobalCounters, Context, Device
 from tinygrad.dtype import DTypeLike, dtypes
 from tinygrad.engine.realize import run_linear
 from tinygrad.helpers import DEBUG, get_single_element
+from tinygrad.device import is_dtype_supported

 def single_kernel_softmax(x_in:Tensor, axis=-1, dtype:DTypeLike|None=None) -> Tensor:
  # only support axis =-1
@@ -61,7 +62,7 @@ class TestFuse(unittest.TestCase):
    b = Tensor.rand(50,50).realize()
    self._test_fuse(lambda a,b: ((a@b).relu()+a).contiguous().softmax(axis=-1), a,b, allow_multiple=True)

-  @unittest.skipUnless(dtypes.float16 in Device[Device.DEFAULT].renderer.supported_dtypes(), f"no float16 on {Device.DEFAULT}")
+  @unittest.skipUnless(is_dtype_supported(dtypes.float16), f"no float16 on {Device.DEFAULT}")
  @unittest.skip("needs RANGEIFY>1")
  def test_fuse_softmax_dtype(self):
    a = Tensor.rand(50,50).realize()
@@ -189,6 +190,7 @@ class TestSoftmaxFusion(unittest.TestCase):

  def test_softmax_bw(self):
    print("*** softmax bw ***")
+    self.test.requires_grad_()
    with Context(NOOPT=1, DEBUG=max(DEBUG.value, 2)):
      self.test.softmax(-1).sum().backward()
      sg = self.test.grad.realize()
@@ -36,7 +36,7 @@ class TestSubBuffer(unittest.TestCase):
    assert len(mv) == 5

  def test_subbuffer_used(self):
-    t = Tensor.arange(0, 10, dtype=dtypes.uint8).clone().realize()
+    t = Tensor.arange(0, 10, dtype=dtypes.uint8).realize()
    vt = t[2:4].realize()
    out = (vt + 100).tolist()
    assert out == [102, 103]
@@ -44,7 +44,7 @@ class TestSubBuffer(unittest.TestCase):
  @needs_second_gpu
  @unittest.skipIf(Device.DEFAULT not in {"CUDA", "NV", "AMD"} or DEV.interface.startswith("MOCK"), "only NV, AMD, CUDA")
  def test_subbuffer_transfer(self):
-    t = Tensor.arange(0, 10, dtype=dtypes.uint8).clone().realize()
+    t = Tensor.arange(0, 10, dtype=dtypes.uint8).realize()
    vt = t[2:5].contiguous().realize()
    out = vt.to(f"{Device.DEFAULT}:1").realize().tolist()
    assert out == [2, 3, 4]
@@ -5,8 +5,8 @@ from tinygrad import Tensor, Device, dtypes, nn
 from tinygrad.helpers import getenv, temp, mv_address
 from extra.gradcheck import numerical_jacobian, jacobian, gradcheck
 from hypothesis import given, settings, strategies as strat
+from tinygrad.device import is_dtype_supported
 from tinygrad.dtype import DTYPES_DICT
-from tinygrad.uop.ops import UOp

 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
@@ -23,19 +23,6 @@ class TestTinygrad(unittest.TestCase):
    self.assertEqual(Tensor(55).shape, ())
    self.assertEqual(Tensor(3.14).shape, ())

-  def test_deviceless_const_construct_device_repr(self):
-    t = Tensor(UOp.const(dtypes.float, 2.0))
-    self.assertIsNone(t.uop.device)
-    self.assertIsNone(t.device)
-    self.assertIn("<UOp None", repr(t))
-
-  def test_deviceless_const_realize_noop(self):
-    t = Tensor(UOp.const(dtypes.float, 2.0))
-    uop = t.uop
-    t.realize()
-    self.assertIs(t.uop, uop)
-    self.assertIsNone(t.uop.device)
-
  def test_plus_equals(self):
    a = Tensor.randn(10,10)
    b = Tensor.randn(10,10)
@@ -47,8 +34,8 @@ class TestTinygrad(unittest.TestCase):

  def test_backward_pass(self):
    def test_tinygrad():
-      x = Tensor(x_init)
-      W = Tensor(W_init)
+      x = Tensor(x_init, requires_grad=True)
+      W = Tensor(W_init, requires_grad=True)
      m = Tensor(m_init)
      out = x.dot(W).relu()
      out = out.log_softmax()
@@ -71,8 +58,8 @@ class TestTinygrad(unittest.TestCase):

  # A simple test is to check that we can accumulate gradients (run backward twice or more times)
  def test_accumulate_gradients(self):
-    x = Tensor(x_init)
-    W = Tensor(W_init)
+    x = Tensor(x_init, requires_grad=True)
+    W = Tensor(W_init, requires_grad=True)
    m = Tensor(m_init)
    out = x.dot(W).relu()
    out = out.log_softmax()
@@ -98,10 +85,10 @@ class TestTinygrad(unittest.TestCase):
      return second_derivative.numpy()

    def test_tinygrad():
-      x_val = Tensor([2.0])
+      x_val = Tensor(2.0)
      f = x_val**3
-      first_derivative = f.sum().gradient(x_val)[0]
-      second_derivative = first_derivative.sum().gradient(x_val)[0]
+      first_derivative = f.gradient(x_val)[0]
+      second_derivative = first_derivative.gradient(x_val)[0]
      return second_derivative.numpy()

    np.testing.assert_allclose(test_tinygrad(), test_pytorch(), atol=1e-5)
@@ -109,8 +96,8 @@ class TestTinygrad(unittest.TestCase):
  # passing `gradient` to backward
  def test_backward_pass_vjp(self):
    def test_tinygrad():
-      x = Tensor(x_init)
-      W = Tensor(W_init)
+      x = Tensor(x_init, requires_grad=True)
+      W = Tensor(W_init, requires_grad=True)
      m = Tensor(m_init)
      out = x.dot(W).relu()
      out = out.log_softmax()
@@ -133,9 +120,9 @@ class TestTinygrad(unittest.TestCase):

  def test_backward_pass_diamond_model(self):
    def test_tinygrad():
-      u = Tensor(U_init)
-      v = Tensor(V_init)
-      w = Tensor(W_init)
+      u = Tensor(U_init, requires_grad=True)
+      v = Tensor(V_init, requires_grad=True)
+      w = Tensor(W_init, requires_grad=True)
      x = u.mul(v).relu()
      y = u.mul(w).relu()
      out = x.add(y).mul(y).relu()
@@ -170,8 +157,8 @@ class TestTinygrad(unittest.TestCase):
      return w1.grad, w2.grad

    def test_tinygrad():
-      w1 = Tensor(init).clone()
-      w2 = Tensor(init).clone()
+      w1 = Tensor(init, requires_grad=True)
+      w2 = Tensor(init, requires_grad=True)
      out = w1.add(w2)
      out.backward()
      return w1.grad.numpy(), w2.grad.numpy()
@@ -190,11 +177,12 @@ class TestTinygrad(unittest.TestCase):
      return w1.grad.numpy(), w2.grad.numpy()

    def test_tinygrad():
-      w1 = Tensor(init).clone()
-      w2 = Tensor(init).clone()
-      assert w1.is_param is True and w2.is_param is True
+      w1 = Tensor(init)
+      w2 = Tensor(init)
+      assert w1.requires_grad is None and w2.requires_grad is None
+      # optimizer sets requires_grad=True for params with requires_grad=None
      nn.optim.SGD([w1, w2], lr=0.01)
-      assert w1.is_param is True and w2.is_param is True
+      assert w1.requires_grad is True and w2.requires_grad is True
      out = w1.add(w2)
      out.backward()
      return w1.grad.numpy(), w2.grad.numpy()
@@ -202,6 +190,21 @@ class TestTinygrad(unittest.TestCase):
    for x, y in zip(test_tinygrad(), test_pytorch()):
      np.testing.assert_allclose(x, y, atol=1e-5)

+  def test_nograd(self):
+    x = Tensor(x_init, requires_grad=False)
+    m = Tensor(m_init, requires_grad=False)
+    W = Tensor(W_init, requires_grad=True)
+    tmp = x.mul(m)
+    mm = tmp.matmul(W)
+    out = mm.relu()
+    out = out.sum()
+    out.backward()
+    assert x.grad is None
+    assert m.grad is None
+    assert tmp.grad is None
+    assert mm.grad is not None
+    assert W.grad is not None
+
  def test_dropout(self):
    with Tensor.train():
      n, rate = 1_000_000, 0.1
@@ -219,8 +222,8 @@ class TestTinygrad(unittest.TestCase):
    def torch_func(x): return torch.nn.functional.log_softmax(x.matmul(torch_W).relu(), dim=1)
    PJ = torch.autograd.functional.jacobian(torch_func, torch_x).squeeze().numpy()

-    tiny_x = Tensor(x)
-    tiny_W = Tensor(W)
+    tiny_x = Tensor(x, requires_grad=True)
+    tiny_W = Tensor(W, requires_grad=True)
    def tiny_func(x): return x.dot(tiny_W).relu().log_softmax()
    J = jacobian(tiny_func, tiny_x)
    NJ = numerical_jacobian(tiny_func, tiny_x)
@@ -232,8 +235,8 @@ class TestTinygrad(unittest.TestCase):
    W = np.random.RandomState(1337).random((10, 5)).astype(np.float32)
    x = np.random.RandomState(7331).random((1, 10)).astype(np.float32)

-    tiny_x = Tensor(x)
-    tiny_W = Tensor(W)
+    tiny_x = Tensor(x, requires_grad=True)
+    tiny_W = Tensor(W, requires_grad=True)
    def tiny_func(x): return x.dot(tiny_W).relu().log_softmax()

    self.assertTrue(gradcheck(tiny_func, tiny_x, eps = 1e-3))
@@ -260,6 +263,10 @@ class TestTinygrad(unittest.TestCase):
  def test_rand_rejects_unknown_kwargs(self):
    with self.assertRaises(TypeError): Tensor.rand(5, generator="foo")

+  def test_randperm_requires_grad(self):
+    self.assertIs(Tensor.randperm(5, requires_grad=True).requires_grad, True)
+    self.assertIs(Tensor.randperm(5, requires_grad=False).requires_grad, False)
+
  def test_randn_isnt_inf_on_zero(self):
    # simulate failure case of rand handing a zero to randn
    original_rand, Tensor.rand = Tensor.rand, Tensor.zeros
@@ -352,7 +359,7 @@ class TestTinygrad(unittest.TestCase):
      assert dtype.itemsize == Tensor.randn(3, dtype=dtype).element_size(), f"Tensor.element_size() not matching Tensor.dtype.itemsize for {dtype}"

  def test_deepwalk_ctx_check(self):
-    layer = Tensor.uniform(1, 1)
+    layer = Tensor.uniform(1, 1, requires_grad=True)
    x = Tensor.randn(1, 1, 1)
    x.dot(layer).mean().backward()
    x = Tensor.randn(1, 1, 1)
@@ -449,7 +456,7 @@ class TestTinygrad(unittest.TestCase):
    np.testing.assert_equal(Tensor(data, dtype=dtypes.float).numpy(), torch.tensor(data, dtype=torch.float).numpy())

  def test_tensor_list_special_values(self):
-    if dtypes.float16 in Device[Device.DEFAULT].renderer.supported_dtypes():
+    if is_dtype_supported(dtypes.float16):
      data = [math.nan, -math.inf, 65504, 65519, 65519.999, 65520, 65520.1]
      data = data + [-x for x in data]
      with np.errstate(over='ignore'): np.testing.assert_allclose(Tensor(data, dtype=dtypes.float16).numpy(), np.array(data).astype(np.float16))
@@ -538,7 +545,7 @@ class TestTinygrad(unittest.TestCase):
      _a = Tensor([3]) in [Tensor([3]), Tensor([4]), Tensor([5])]

  def test_repr_with_grad(self):
-    a = Tensor([1.0])
+    a = Tensor([1.0], requires_grad=True)
    b = Tensor([1])
    c = (a + b).sum().backward()
    print(a)
@@ -551,7 +558,7 @@ class TestTinygrad(unittest.TestCase):
      Tensor.zeros(2, 2).realize()

  def test_shrink(self):
-    t = Tensor.arange(32).clone().realize()
+    t = Tensor.arange(32).contiguous().realize()
    self.assertListEqual(t[16:20].tolist(), [16,17,18,19])
    self.assertListEqual(t.shrink_to(16).tolist(), list(range(16)))
    t = t.reshape(4, 8).contiguous().realize()
@@ -564,26 +571,26 @@ class TestTinygrad(unittest.TestCase):
 class TestMoveTensor(unittest.TestCase):
  d0, d1 = f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"
  @given(strat.sampled_from([d0, d1]), strat.sampled_from([d0, d1]),
-         strat.sampled_from([dtypes.float16, dtypes.float32]), strat.sampled_from([True, False]))
-  def test_to_preserves(self, src, dest, dtype, is_param):
-    if dtype not in Device[Device.DEFAULT].renderer.supported_dtypes():
+         strat.sampled_from([dtypes.float16, dtypes.float32]), strat.sampled_from([True, False, None]))
+  def test_to_preserves(self, src, dest, dtype, requires_grad):
+    if not is_dtype_supported(dtype):
      return
-    s = Tensor([1, 2, 3], device=src, dtype=dtype).is_param_(is_param)
-    if is_param: s.sum().backward()
+    s = Tensor([1, 2, 3], device=src, dtype=dtype, requires_grad=requires_grad)
+    if requires_grad: s.sum().backward()
    t = s.to(dest)
    np.testing.assert_equal(s.numpy(), t.numpy())
    assert s.dtype == t.dtype
-    assert s.is_param == t.is_param
-    if is_param:
+    assert s.requires_grad == t.requires_grad
+    if requires_grad:
      np.testing.assert_equal(s.grad.numpy(), t.grad.numpy())

-  @given(strat.sampled_from([dtypes.float16, dtypes.float32]), strat.sampled_from([True, False]))
-  def test_shard_preserves(self, dtype, is_param):
-    s = Tensor([1, 2, 3], dtype=dtype).is_param_(is_param)
+  @given(strat.sampled_from([dtypes.float16, dtypes.float32]), strat.sampled_from([True, False, None]))
+  def test_shard_preserves(self, dtype, requires_grad):
+    s = Tensor([1, 2, 3], dtype=dtype, requires_grad=requires_grad)
    t = s.shard((f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"))
    np.testing.assert_equal(s.numpy(), t.numpy())
    assert s.dtype == t.dtype
-    assert s.is_param == t.is_param
+    assert s.requires_grad == t.requires_grad

  @given(strat.sampled_from([d0, d1]))
  def test_same_dev(self, dev):
@@ -592,8 +599,8 @@ class TestMoveTensor(unittest.TestCase):
    assert x is y

  def test_to_grad(self):
-    x = Tensor.eye(3, device=self.d0)
-    y = Tensor([[2.0,0,-2.0]], device=self.d0)
+    x = Tensor.eye(3, requires_grad=True, device=self.d0)
+    y = Tensor([[2.0,0,-2.0]], requires_grad=True, device=self.d0)
    z = y.matmul(x).to(self.d1).sum()
    z.backward()
    np.testing.assert_equal(x.grad.numpy(), [[2,2,2],[0,0,0],[-2,-2,-2]])
@@ -726,14 +733,6 @@ class TestZeroShapeTensor(unittest.TestCase):
    np.testing.assert_allclose(a.numpy(), b.numpy())
    self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)

-  def test_clone_deviceless_const(self):
-    t = Tensor(UOp.const(dtypes.float, 2.0)).clone()
-    np.testing.assert_equal(t.numpy(), 2.0)
-    self.assertTrue(t.uop.has_buffer_identity())
-
-  def test_numpy_deviceless_const(self):
-    np.testing.assert_equal(Tensor(UOp.const(dtypes.float, 2.0)).numpy(), 2.0)
-
  def test_clone_with_shrink(self):
    a = Tensor.rand(16, 16)
    b = a.shrink(((2, 10), None)).clone()
@@ -747,18 +746,13 @@ class TestZeroShapeTensor(unittest.TestCase):
    self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)

  def test_clone_with_grad(self):
-    a = Tensor.rand(16, 16)
+    a = Tensor.rand(16, 16, requires_grad=True)
    a.mul(5.0).add(5.0).mean().backward()
    b = a.clone()
    assert a.grad is not None
    assert b.grad is not None
    np.testing.assert_allclose(a.grad.numpy(), b.grad.numpy())

-  def test_clone_deviceless_const_to_cpu(self):
-    t = Tensor(UOp.const(dtypes.float, 2.0)).clone(device="CPU")
-    self.assertEqual(t.device, "CPU")
-    np.testing.assert_equal(t.numpy(), 2.0)
-
  def test_reduce_default(self):
    np.testing.assert_equal(Tensor([]).max().numpy(), -float("inf"))
    np.testing.assert_equal(Tensor([]).min().numpy(), float("inf"))
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	770dac0e0d	broadcast	2026-05-14 17:04:37 -07:00
geohot	b827858479	broadcast shape	2026-05-14 17:01:20 -07:00