renumber invalids before callify (#17542 )

* renumber invalids before callify * change * Revert "change" This reverts commit 6f4df1541e79721a85ee3f5801114454f264c973. * renumber in tensor * scope renumber_invalid_outputs * cleanup
precompile_backward tests for sched_cache (#17544 )
2026-08-15 18:00:54 +09:00 · 2026-08-15 15:25:31 +09:00 · 2026-08-15 11:00:21 +09:00 · 2026-08-14 21:39:47 -04:00 · 2026-08-14 21:07:01 -04:00 · 2026-08-14 16:24:51 -04:00
123 changed files with 3018 additions and 1582 deletions
@@ -41,10 +41,6 @@ inputs:
    description: "Install LLVM?"
    required: false
    default: 'false'
-  tinydreno:
-    description: "Install tinydreno"
-    required: false
-    default: 'false'
  qemu:
    description: "Install qemu"
    required: false
@@ -277,12 +273,6 @@ runs:
      shell: bash
      run: brew install llvm@20

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

  torchbackend:
    name: Torch Backend Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -88,7 +88,7 @@ jobs:

  bepython:
    name: Python Backend
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -126,7 +126,7 @@ jobs:

  linter:
    name: Linters
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 10

    steps:
@@ -157,7 +157,7 @@ jobs:

  nulltest:
    name: Null Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15

    steps:
@@ -191,7 +191,7 @@ jobs:

  unittest:
    name: Unit Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15

    steps:
@@ -228,7 +228,7 @@ jobs:
      matrix:
        group: [1, 2]
    name: SPEC=2 (${{ matrix.group }})
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -244,7 +244,7 @@ jobs:

  fuzzing:
    name: Fuzzing
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 10
    steps:
    - name: Checkout Code
@@ -260,7 +260,7 @@ jobs:

  testopenclimage:
    name: CL IMAGE Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -280,7 +280,7 @@ jobs:

  testopenpilot:
    name: openpilot Compile Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -294,7 +294,7 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1361 ALLOWED_GATED_READ_IMAGE=54 FLOAT16=1 DEV="CL::IMAGE_PITCH_ALIGNMENT=64" 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=1361 ALLOWED_GATED_READ_IMAGE=38 FLOAT16=1 DEV="CL::IMAGE_PITCH_ALIGNMENT=64" IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
          # IMAGE_PITCH_ALIGNMENT=64 matches adreno 630
      - name: Test openpilot CL compile fp32 (test correctness)
        run: |
@@ -309,7 +309,7 @@ jobs:

  testonnxcpu:
    name: ONNX (CPU) Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 20

    steps:
@@ -328,7 +328,7 @@ jobs:

  testoptim:
    name: Optimization Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 20
    steps:
      - name: Checkout Code
@@ -360,7 +360,7 @@ jobs:

  testllm:
    name: Test LLM
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
    env:
      CHECK_OOB: 0
@@ -387,7 +387,7 @@ jobs:

  testmodels:
    name: Models
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
    steps:
      - name: Checkout Code
@@ -407,7 +407,7 @@ jobs:

  testdsp:
    name: Linux (DSP)
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
    steps:
    - name: Checkout Code
@@ -435,7 +435,7 @@ jobs:
          - 'WEBGPU'

    name: Linux (DEV=${{ matrix.dev }})
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 20
    steps:
    - name: Checkout Code
@@ -461,7 +461,7 @@ jobs:

  testamdasm:
    name: AMD ASM IDE
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 20
    env:
      DEV: MOCKKFD+AMD
@@ -507,7 +507,7 @@ jobs:

  hcq2:
    name: hcq2
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 5
    steps:
      - name: Checkout Code
@@ -521,16 +521,15 @@ jobs:
      - name: Run HCQ2 tests
        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/test_tiny.py
      - name: Run HCQ2 multi-device tests
-        run: |
-          HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_multitensor.py \
-            TestMultiTensor.test_simple_add TestMultiTensor.test_shard_reduce \
-            TestMultiTensor.test_backward_sum TestMultiTensor.test_matmul_shard_0_0
+        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest -n=auto test/backend/test_multitensor.py
      - name: Run HCQ2 JIT tests
        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_jit.py
+      - name: Run HCQ2 unit tests
+        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest test/device/test_hcq2.py

  testmockam:
    name: Linux (am)
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
    env:
      DEV: MOCKPCI+AMD
@@ -566,7 +565,7 @@ jobs:
        arch: [gfx1100, gfx1201, gfx950]

    name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
    env:
      DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
@@ -585,6 +584,9 @@ jobs:
        run: |
          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['AMD'], Device.DEFAULT"
          DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
+      - name: Run MXFP4 Llama training on NULL backend
+        if: ${{ matrix.backend == 'amd' && matrix.arch == 'gfx950' }}
+        run: PYTHONPATH=. DEV=NULL:HIP:gfx950 MXFP4=1 LLAMA_LAYERS=2 BENCHMARK=3 NULL_ALLOW_COPYOUT=1 NO_HIPCC=1 ROCM_PATH=/opt/rocm JITBEAM=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh
      - name: Run pytest (amd)
        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM --durations=20
      - name: Run disk copy tests
@@ -601,7 +603,7 @@ jobs:
        backend: [ptx, nv]

    name: Linux (${{ matrix.backend }})
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 20
    env:
      FORWARD_ONLY: 1
@@ -629,175 +631,23 @@ jobs:
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

-# ****** OSX Tests ******
-
-  unittestmacos:
-    name: MacOS (unit)
-    runs-on: macos-26
-    timeout-minutes: 20
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: unittest-macos
-        deps: testing_unit
-    - name: Run unit tests
-      run: DEV=METAL python -m pytest -n=auto test/unit/ --durations=20
-    - name: Test tensor core ops (fake)
-      run: DEV=METAL DEBUG=3 TC=2 python test/backend/test_ops.py TestOps.test_gemm
-    - name: Test tensor core ops (real)
-      run: DEV=METAL DEBUG=3 python test/backend/test_ops.py TestOps.test_big_gemm
-    - name: Test Beam Search
-      run: DEV=METAL IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
-    - name: Test Device Specific
-      run: DEV=METAL python3 -m pytest test/device/test_metal.py
-    #- name: Fuzz Test linearizer
-    #  run: DEV=METAL DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
-  unittestmacosmock:
-    name: MacOS (unit, mock)
-    runs-on: macos-26
-    timeout-minutes: 20
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: unittest-macos-mock
-        deps: testing_unit
-        amd: 'true'
-        ocelot: 'true'
-    - name: Run NULL backend tests
-      run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --durations=20
-    - name: Run pytest (amd)
-      env:
-        DEV: MOCKKFD+AMD
-        FORWARD_ONLY: 1
-      run: |
-        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
-    - name: Run pytest (ptx)
-      env:
-        DEV: "MOCK+NV:PTX"
-        FORWARD_ONLY: 1
-        # TODO: failing due to library loading error
-        CAPTURE_PROCESS_REPLAY: 0
-      run: |
-        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
-  testmetal:
-    strategy:
-      fail-fast: false
-      matrix:
-        group: [1, 2]
-    name: MacOS (DEV=METAL) (${{ matrix.group }})
-    runs-on: macos-26
-    timeout-minutes: 20
-    env:
-      DEV: METAL
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: macos-metal
-        deps: testing_unit
-    - name: Check Device.DEFAULT and print some source
-      run: |
-        python -c "from tinygrad import Device; assert Device.DEFAULT == 'METAL'"
-        DEBUG=4 python test/test_tiny.py TestTiny.test_plus
-    - name: Run backend tests
-      run: python -m pytest -n=auto test/backend --durations=20 --splits 2 --group ${{ matrix.group }}
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
-  testmacos:
-    strategy:
-      fail-fast: false
-      matrix:
-        dev:
-          - 'CPU:CLANG'
-          - 'CPU:LLVM'
-          - 'CPU:LVP'
-          - 'WEBGPU'
-
-    name: MacOS (DEV=${{ matrix.dev }})
-    runs-on: macos-26
-    timeout-minutes: 20
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: macos-${{ matrix.dev }}
-        deps: "testing_unit${{ contains(matrix.dev, 'LVP') && ' mesa' || '' }}"
-        llvm: ${{ contains(matrix.dev, 'LLVM') || contains(matrix.dev, 'LVP') }}
-        webgpu: ${{ matrix.dev == 'WEBGPU' }}
-    - name: Set env
-      run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
-    - name: Check Device.DEFAULT and print some source
-      run: |
-        python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
-        DEBUG=4 python test/test_tiny.py TestTiny.test_plus
-    - name: Run test_tiny
-      run: python -m pytest -n=auto test/test_tiny.py --durations=20
-    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
-
-# ****** Windows Tests ******
-
-  testwindows:
-    strategy:
-      fail-fast: false
-      matrix:
-        dev:
-          - 'CPU:CLANG'
-          - 'CPU:LLVM'
-          - 'CPU:X86'
-          - 'WEBGPU'
-
-    name: Windows (DEV=${{ matrix.dev }})
-    runs-on: windows-2025
-    timeout-minutes: 15
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v6
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: windows-${{ matrix.dev }}-minimal
-          deps: testing_unit
-          pydeps: ${{ matrix.dev == 'WEBGPU' && 'dawn-python' || '' }}
-      - name: Set env
-        shell: bash
-        run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
-      - name: Check Device.DEFAULT and print some source
-        shell: bash
-        run: |
-          python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
-          DEBUG=4 python test/test_tiny.py TestTiny.test_plus
-      - name: Run test_tiny
-        shell: bash
-        run: python -m pytest -n=auto test/test_tiny.py --durations=20
-
 # ****** Compile-only Tests ******

  compiletests:
    strategy:
      fail-fast: false
      matrix:
-        backend: [ir3, nak]
-    name: Compile-only (${{ matrix.backend }})
-    runs-on: *linux
+        dev:
+          - 'NULL:IR3:a630'
+          - 'NULL:QCOMCL:a630'
+          - 'NULL:NAK:sm_120'
+    name: Compile-only (DEV=${{ matrix.dev }})
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
    timeout-minutes: 15
+    env:
+      NULL_ALLOW_COPYOUT: 1
+      DEV: ${{ matrix.dev }}${{ contains(matrix.dev, 'a630') && ',IMAGE_PITCH_ALIGNMENT=64' || '' }}
+      IMAGE: ${{ contains(matrix.dev, 'a630') && '1' || '0' }}
    steps:
      - name: Checkout Code
        uses: actions/checkout@v6
@@ -806,51 +656,14 @@ jobs:
        with:
          key: compile-${{ matrix.backend }}
          deps: "testing_unit mesa"
-      - name: Set env
+          qemu: ${{ contains(matrix.dev, 'QCOMCL') }}
+      - name: Test IMAGE
        shell: bash
-        run: printf "NULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'DEV=NULL:IR3:a630' || matrix.backend == 'nak' && 'DEV=NULL:NAK:sm_120' }}" >> $GITHUB_ENV
+        if: contains(matrix.dev, 'a630')
+        run: DEBUG=7 python3 test/backend/test_ops.py TestOps.test_gemm | grep isam
      - name: Run test_ops
        shell: bash
        run: |
          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
          python -m pytest -n=auto test/backend/test_ops.py --durations=20
-      - name: Run test_ops (IMAGE)
-        if: matrix.backend == 'ir3'
-        shell: bash
-        env:
-          IMAGE: 1
-          DEV: "NULL:IR3:a630,IMAGE_PITCH_ALIGNMENT=64"
-        run: |
-          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_gemm | grep image_load
-          python -m pytest -n=auto test/backend/test_ops.py --durations=20
-  qcomclcompiletests:
-    name: Compile-only (QCOM CL)
-    runs-on: ubuntu-24.04-arm
-    timeout-minutes: 15
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v6
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: compile-qcomcl
-          deps: testing_unit
-          tinydreno: 'true'
-      - name: Set env
-        shell: bash
-        run: printf "DEV=NULL:QCOMCL:a630\nNULL_ALLOW_COPYOUT=1" >> $GITHUB_ENV
-      - name: Run test_ops
-        shell: bash
-        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
-          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
-          python -m pytest -n=auto test/backend/test_ops.py --durations=20
-      - name: Run test_ops (IMAGE)
-        shell: bash
-        env:
-          IMAGE: 1
-          DEV: "NULL:QCOMCL:a630,IMAGE_PITCH_ALIGNMENT=64"
-        run: |
-          DEBUG=4 python test/backend/test_ops.py TestOps.test_gemm | grep read_imagef
-          python -m pytest -n=auto test/backend/test_ops.py --durations=20
@@ -1458,7 +1458,8 @@ def train_llama3():

  # realize everything here
  if optim.master_params: Tensor.realize(*optim.master_params)
-  Tensor.realize(*optim.params, *fp8_inv_scales, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)
+  loss_acc = Tensor.zeros(1, dtype=dtypes.float32, device=device)
+  Tensor.realize(loss_acc, *optim.params, *fp8_inv_scales, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)

  @TinyJit
  def minibatch(tokens:Tensor):
@@ -1476,8 +1477,8 @@ def train_llama3():
    for g, new_g in zip(grads, loss.gradient(*optim.params)):
      apply_grad(g, new_g.uop)

-    loss_cpu = loss.flatten().float().to("CPU")
-    return loss_cpu.realize(*grads, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)
+    loss_acc.assign(loss_acc + loss.flatten().float())
+    return loss_acc.realize(*grads, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)

  @TinyJit
  def optim_step():
@@ -1490,9 +1491,10 @@ def train_llama3():

    lr_cpu = optim.lr.float().to("CPU")
    grad_norm_cpu = grad_norm.float().to("CPU")
-    Tensor.realize(lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales, *fp8_amax, *fp8_grad_amax)
+    loss_cpu = loss_acc.to("CPU")
+    Tensor.realize(lr_cpu, grad_norm_cpu, loss_cpu, loss_acc.assign(0), *grads, *fp8_inv_scales, *fp8_amax, *fp8_grad_amax)

-    return lr_cpu, grad_norm_cpu
+    return lr_cpu, grad_norm_cpu, loss_cpu

  @TinyJit
  @Context(TRAINING=0)
@@ -1547,8 +1549,8 @@ def train_llama3():
      st = time.perf_counter()

      stopped = False
-      losses, data_time, dev_time = [], 0, 0
-      for _ in range(grad_acc if i >= 2 else 1):
+      data_time, dev_time = 0, 0
+      for _ in range(accum_steps:=grad_acc if i >= 2 else 1):
        ist = time.perf_counter()
        try: tokens = next(train_iter)
        except StopIteration:
@@ -1556,16 +1558,15 @@ def train_llama3():
          break
        mst = time.perf_counter()
        data_time += mst - ist
-        losses.append(minibatch(tokens).item())
+        minibatch(tokens)
        dev_time += time.perf_counter() - mst
      if stopped: break

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

-      loss = sum(losses) / len(losses)
      optim_time = et - gt
      dev_time += optim_time
      step_time = et - st
@@ -114,6 +114,11 @@ def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor,
                             amax_x2:Tensor|None, next_amax_x2:Tensor|None,
                             grad_amax_xw13:Tensor|None, next_grad_amax_xw13:Tensor|None,
                             grad_amax_xout:Tensor|None, next_grad_amax_xout:Tensor|None):
+  if FUSED_SILU_W13 and MXFP4:
+    from extra.llama_kernels.swiglu import swiglu
+    out, *ret = matmul(swiglu(x_w13), w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout,
+                       next_grad_amax_state=next_grad_amax_xout, next_amax_x=next_amax_x2)
+    return out, ret
  if FUSED_SILU_W13 and not MXFP4:
    from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
    x2_fp8 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13,
@@ -12,7 +12,7 @@ from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb
 from extra.llama_kernels.rmsnorm import rmsnorm
-from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8
+from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8, asm_gemm, can_use_asm_gemm
 from extra.gemm.moe_gemm import grouped_mx_gemm
 from extra.gemm.moe_routing import route, dispatch, combine

@@ -182,12 +182,12 @@ class GPTOSS:
    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
    xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)  # (B,N,H,D)/(B,N,KV,D)

-    if sliding:
-      attn = self._sliding_attention(xq, xk, xv, sinks)
-    elif getenv("HK_FLASH_ATTENTION"):
+    if getenv("HK_FLASH_ATTENTION"):
      from extra.thunder.amd.fa import flash_attention
-      attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks)
+      attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks, window=self.sliding_window if sliding else 0)
      attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
+    elif sliding:
+      attn = self._sliding_attention(xq, xk, xv, sinks)
    else:
      xqm = xq.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep, self.head_dim).permute(0, 2, 3, 1, 4)
      xkm, xvm = xk.permute(0, 2, 1, 3).unsqueeze(2), xv.permute(0, 2, 1, 3).unsqueeze(2)
@@ -263,7 +263,11 @@ class GPTOSS:
                        w_down=self.w_down[i], w_down_scale=self.w_down_scale[i], w_down_bias=self.w_down_bias[i])
      h, *_ = self.run_layer(h, freqs_cis, mask_full, i % 2 == 0, attn_kwargs, ffn_kwargs, save=save)

-    logits = self.norm(h) @ self.output.T
+    h_normed = self.norm(h)
+    pad = (-self.dim) % 256
+    h_padded, w_padded = h_normed.pad((None, None, (0, pad))), self.output.pad(((0, 0), (0, pad)))
+    if ASM_GEMM and can_use_asm_gemm(h_padded, w_padded.T): logits = asm_gemm(h_padded, w_padded.T)
+    else: logits = h_normed @ self.output.T
    return logits

 def _get_pads(uop:UOp) -> list[UOp]:
@@ -1,8 +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 ROCM_PATH=${ROCM_PATH:-/opt/rocm-7.1.1}
+export PATH="$ROCM_PATH/bin:$PATH"
 export DEV=${DEV:-AMD}
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -16,7 +16,7 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export FP8=${FP8:-1}
+export MXFP4=${MXFP4:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
 export FAST_CE=${FAST_CE:-1}
 export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
@@ -26,7 +26,7 @@ export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
 export SPLIT_W13=${SPLIT_W13:-0}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}

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

@@ -46,7 +46,7 @@ export DATA_SEED=${DATA_SEED:-5760}
 export JITBEAM=${JITBEAM:-3}
 export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1

-export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+export FAKEDATA=${FAKEDATA:-$([[ "$DEV" == NULL:* ]] && echo 1 || echo 0)} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
 fi
@@ -1,8 +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 ROCM_PATH=${ROCM_PATH:-/opt/rocm-7.1.1}
+export PATH="$ROCM_PATH/bin:$PATH"
 export DEV=${DEV:-AMD}
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -16,7 +16,7 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export FP8=${FP8:-1}
+export MXFP4=${MXFP4:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
 export FAST_CE=${FAST_CE:-1}
 export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
@@ -26,7 +26,7 @@ export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
 export SPLIT_W13=${SPLIT_W13:-0}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}

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

@@ -1,4 +1,6 @@
 #!/bin/bash
+set -e
+
 export BENCHMARK=${BENCHMARK:-5}
 export EVAL_BS=0
 VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=${DEBUG:--0} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
@@ -11,6 +11,7 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
+export GROUPED_MOE=${GROUPED_MOE:-1}
 export ALL2ALL=${ALL2ALL:-1}
 export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
@@ -11,6 +11,7 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
+export GROUPED_MOE=${GROUPED_MOE:-1}
 export ALL2ALL=${ALL2ALL:-1}
 export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
@@ -0,0 +1,31 @@
+import argparse, time
+from tinygrad.llm.model import Transformer
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--model", required=True, help="path to gguf model")
+  parser.add_argument("--max-context", type=int, default=8192, help="max context length (default: %(default)s)")
+  parser.add_argument("--prompt-tokens", type=int, default=1024, help="number of prompt tokens (default: %(default)s)")
+  parser.add_argument("--decode-tokens", type=int, default=16, help="number of tokens to decode (default: %(default)s)")
+  parser.add_argument("--chunk-size", type=int, default=32, help="chunk size for prefill (default: %(default)s)")
+  args = parser.parse_args()
+
+  st = time.perf_counter()
+  model, _ = Transformer.from_gguf(args.model, args.max_context)
+  print(f"load {time.perf_counter()-st:.3f}s", flush=True)
+
+  st = time.perf_counter()
+  model.warmup()
+  print(f"warm {time.perf_counter()-st:.3f}s", flush=True)
+
+  prompt = [257] + [1000+i%1000 for i in range(args.prompt_tokens-1)]
+  gen = model.generate(prompt, chunk_size=args.chunk_size)
+  st = time.perf_counter()
+  # first token is time-to-first-token; counted as part of prefill
+  output = [next(gen)]
+  pt = time.perf_counter()
+  print(f"prefill {args.prompt_tokens/(pt-st):.3f} tok/s", flush=True)
+
+  for _ in range(args.decode_tokens): output.append(next(gen))
+  et = time.perf_counter()
+  print(f"decode {args.decode_tokens/(et-pt):.3f} tok/s output {output}", flush=True)
@@ -35,7 +35,7 @@ def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], Li
    return name

  for call in iter_kernel_calls(linear):
-    arg_uops = [b for b in call.src[1:] if b.op is not Ops.BIND]
+    arg_uops = [b for b in call.src[1:] if not b.is_bound_var]
    prg = to_program(call.src[0], Device[arg_uops[0].device].renderer)
    info = prg.arg
    functions[info.function_name] = prg.src[2].arg
@@ -6,6 +6,7 @@ from tinygrad.renderer import Estimates
 from tinygrad.helpers import getenv, all_same, DEBUG, ceildiv
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8
+from extra.llama_kernels.quantize_mxfp4 import quantize_mxfp4

 TILE_M, TILE_N, TILE_K = 256, 256, 64

@@ -121,10 +122,30 @@ def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp,
  groups_x, groups_y = UOp.special(ceildiv(N, tile_n), "gidx0"), UOp.special(ceildiv(M, tile_m), "gidx1")
  lds = UOp.placeholder((163840,), dtypes.uint8, 0, AddrSpace.LOCAL)
  sink = UOp.sink(C.base, A.base, B.base, scale_a.base, scale_b.base, *(x.base for x in extra), lds, threads, groups_x, groups_y,
-                  arg=KernelInfo(f"custom_mxfp4_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K)))
+                  arg=KernelInfo(f"mxfp4_gemm_{M}_{N}_{K}",
+                                 estimates=Estimates(ops=2*M*N*K, mem=(M*half_k+N*half_k)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
  insts = build_kernel(M, N, K, tile_m, tile_n)
  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in insts))))

+def _mxfp4_gemm_quantized(a_q:Tensor, b_q:Tensor, scale_a:Tensor, scale_b:Tensor) -> Tensor:
+  M, half_k = a_q.shape
+  N, half_k_b = b_q.shape
+  assert half_k == half_k_b
+  is_multi = isinstance(a_q.device, tuple)
+  reduce_out = is_multi and (a_q.uop.axis == 1 or b_q.uop.axis == 1)
+  if not is_multi: out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a_q.device)
+  elif reduce_out: out = Tensor(Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a_q.device).uop.unshard(0), device=a_q.device)
+  elif a_q.uop.axis == 0:
+    out = Tensor(Tensor.invalids(1, M//len(a_q.device), N, dtype=dtypes.bfloat16, device=a_q.device).uop.unshard(1), device=a_q.device)
+  elif b_q.uop.axis == 0:
+    out = Tensor(Tensor.invalids(1, M, N//len(a_q.device), dtype=dtypes.bfloat16, device=a_q.device).uop.unshard(2), device=a_q.device)
+  else: out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a_q.device)
+  tile_m, tile_n = next((tm, tn) for tm, tn in ((256, 256), (192, 256), (128, 512)) if M % tm == N % tn == 0)
+  out = Tensor.custom_kernel(out, a_q, b_q, scale_a, scale_b,
+                             fxn=functools.partial(custom_mxfp4_gemm, tile_m=tile_m, tile_n=tile_n))[0]
+  if reduce_out: out = out.sum(0)
+  return out.squeeze(0)
+
 def quantize_mxfp8(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
  # 1x32 block scaling along the last axis
  *batch, K = x.shape
@@ -137,50 +158,6 @@ def quantize_mxfp8(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
  packed = mx_pack(e8) if len(batch) == 1 and scale_K % 4 == 0 else None
  return x_clamped.cast(FP8_DTYPE), e8, packed

-def _mxfp4_shuffle_weight(x:Tensor) -> Tensor:
-  # shuffle_weight(x, layout=(16, 16)) on the packed uint8 buffer.
-  if x.ndim == 3:
-    ndev, rows, half_k = x.shape
-    return x.reshape(ndev, rows//16, 16, half_k//32, 2, 16).permute(0, 1, 3, 4, 2, 5).reshape(ndev, rows, half_k).contiguous()
-  rows, half_k = x.shape
-  return x.reshape(rows//16, 16, half_k//32, 2, 16).permute(0, 2, 3, 1, 4).reshape(rows, half_k).contiguous()
-
-def _mxfp4_shuffle_scales(x:Tensor) -> Tensor:
-  # e8m0_shuffle: each 256x8 scale tile is arranged for the raw MFMA scale loads.
-  if x.ndim == 3:
-    ndev, rows, scale_k = x.shape
-    return x.reshape(ndev, rows//32, 2, 16, scale_k//8, 2, 4).permute(0, 1, 4, 6, 3, 5, 2).reshape(ndev, rows, scale_k).contiguous()
-  rows, scale_k = x.shape
-  return x.reshape(rows//32, 2, 16, scale_k//8, 2, 4).permute(0, 3, 5, 2, 4, 1).reshape(rows, scale_k).contiguous()
-
-def quantize_mxfp4(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
-  # OCP MXFP4: 1x32 blocks, E2M1 values packed low-nibble first, and E8M0 scales.
-  *batch, K = x.shape
-  rows = math.prod(batch)
-  assert x.ndim >= 2 and K % 256 == 0 and rows % 32 == 0, \
-    f"mxfp4 quantization needs rows%32 and K%256, got {x.shape}"
-  xb = x.float().reshape(*batch, K//32, 32)
-  amax = xb.abs().max(axis=-1)
-
-  # even scale rounding: round the fp32 significand before choosing 2^(floor(log2)-2).
-  amax_rounded = ((amax.bitcast(dtypes.uint32) + 0x200000) & 0xFF800000).bitcast(dtypes.float32)
-  scale_exp = (amax_rounded.maximum(2**-126).log2().floor() - 2).clamp(-127, 127)
-  e8 = (scale_exp + 127).cast(dtypes.uint8)
-  scaled = xb * (-scale_exp).exp2().reshape(*batch, K//32, 1)
-
-  mag = scaled.abs()
-  code = sum(x.cast(dtypes.uint8) for x in
-             (mag > .25, mag >= .75, mag > 1.25, mag >= 1.75, mag > 2.5, mag >= 3.5, mag > 5.0))
-  code = code | ((scaled < 0).cast(dtypes.uint8) << 3)
-  code = code.reshape(*batch, K)
-  packed = code[..., 0::2] | (code[..., 1::2] << 4)
-  if isinstance(x.device, tuple) and x.uop.axis == x.ndim-2 and x.shape[x.uop.axis] == len(x.device):
-    axis = x.uop.axis
-    order = (axis, *range(axis), *range(axis+1, e8.ndim))
-    e8_local = e8.permute(order)
-    return packed, e8, _mxfp4_shuffle_scales(e8_local.reshape(e8_local.shape[0], -1, K//32))
-  return packed, e8, _mxfp4_shuffle_scales(e8.reshape(rows, K//32))
-
 def mx_pack(e8:Tensor) -> Tensor:
  rows, scale_K = e8.shape
  return e8.reshape(rows, scale_K // 4, 4).bitcast(dtypes.uint32).reshape(rows, scale_K // 4).permute(1, 0).contiguous()
@@ -405,15 +382,16 @@ def custom_mx_gemm_bw(gradient:UOp, kernel:UOp, has_w_post:bool, w_stored:bool=F
 # ** mxfp4 gemm backward

 def custom_mxfp4_gemm_bw(gradient:UOp, kernel:UOp):
-  # The raw kernel consumes quantized buffers, while the final two inputs retain the BF16 operands for STE gradients.
-  inputs = kernel.src[1:]  # (out, a_q, b_q, scale_a, scale_b, a, w)
-  assert len(inputs) == 7
+  inputs = kernel.src[1:]  # out, row operands/scales, BF16 operands, column operands/scales
+  assert len(inputs) == 11
  a, w = Tensor(inputs[5], device=inputs[5].device), Tensor(inputs[6], device=inputs[6].device)
+  a_col, scale_a_col = Tensor(inputs[7], device=a.device), Tensor(inputs[8], device=a.device)
+  w_col, scale_w_col = Tensor(inputs[9], device=a.device), Tensor(inputs[10], device=a.device)
  g = Tensor(gradient, device=a.device)[:a.shape[0]].cast(dtypes.bfloat16)
-  grad_a = asm_gemm(g, w, mxfp4=True)
-  a_flat, g_flat = a.reshape(-1, a.shape[-1]), g.reshape(-1, g.shape[-1])
-  grad_w = asm_gemm(g_flat.T, a_flat, mxfp4=True)
-  return (None, None, None, None, None, grad_a.uop, grad_w.uop)
+  g_row, scale_g_row, g_col, scale_g_col = quantize_mxfp4(g, flatten_row=True)
+  grad_a = _mxfp4_gemm_quantized(g_row, w_col, scale_g_row, scale_w_col).reshape(*a.shape[:-1], w.shape[-1])
+  grad_w = _mxfp4_gemm_quantized(g_col, a_col, scale_g_col, scale_a_col).reshape(w.shape)
+  return (None, None, None, None, None, grad_a.uop, grad_w.uop, None, None, None, None)

 # ** main gemm function

@@ -459,16 +437,10 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N
      tile_m, tile_n = next((tm, tn) for tm, tn in ((256, 256), (192, 256), (128, 512)) if (batch*M) % tm == N % tn == 0)
      fxn = functools.partial(custom_mxfp4_gemm, tile_m=tile_m, tile_n=tile_n)
      w = b.T
-      if k_sharded:
-        ndev = len(a.device)
-        a_q, _, scale_a = quantize_mxfp4(a.reshape(batch, M, ndev, K))
-        b_q, _, scale_b = quantize_mxfp4(w.reshape(w.shape[0], ndev, K))
-        b_q = _mxfp4_shuffle_weight(b_q.permute(1, 0, 2))
-      else:
-        a_q, _, scale_a = quantize_mxfp4(a.reshape(batch*M, K))
-        b_q, _, scale_b = quantize_mxfp4(w)
-        a_q, b_q = a_q.reshape(batch, M, K//2).contiguous(), _mxfp4_shuffle_weight(b_q)
-      out = Tensor.custom_kernel(out, a_q, b_q, scale_a, scale_b, a, w, fxn=fxn, grad_fxn=custom_mxfp4_gemm_bw)[0]
+      a_q, scale_a, a_col, scale_a_col = quantize_mxfp4(a, shuffle_col=True)
+      b_q, scale_b, b_col, scale_b_col = quantize_mxfp4(w, shuffle_row=True, shuffle_col=True)
+      out = Tensor.custom_kernel(out, a_q, b_q, scale_a, scale_b, a, w,
+                                 a_col, scale_a_col, b_col, scale_b_col, fxn=fxn, grad_fxn=custom_mxfp4_gemm_bw)[0]
    elif mx:
      # mxfp8 1x32 block scaling
      if mx_scales is not None:
@@ -5,7 +5,8 @@ BLOCK_ROW = 256

 def _sharded_invalids(shape:tuple[int, ...], dtype, device) -> Tensor:
  if isinstance(device, tuple):
-    return Tensor.invalids(*shape, dtype=dtype, device=device[0]).shard(device, axis=0)
+    per = Tensor.invalids(shape[0]//len(device), *shape[1:], dtype=dtype, device=device)
+    return Tensor(per.uop.unshard(0), device=device)
  return Tensor.invalids(*shape, dtype=dtype, device=device)

 def _atomic_add(device:str) -> str:
@@ -79,7 +79,7 @@ if __name__ == "__main__":
  linear, var_vals = C.linear_with_vars()
  last_call = linear.src[-1]
  ast = last_call.src[0]
-  bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]
+  bufs = [s.buffer for s in last_call.src[1:] if not s.is_bound_var]

  src = compiled.asm["ptx"]
  # specify the shared memory here so we don't need to do it dynamically
@@ -3,7 +3,7 @@ from typing import cast, Any, Callable
 import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools, collections, atexit
 assert sys.platform != 'win32'
 from dataclasses import dataclass
-from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, encode_kernargs_clike, make_cmdbuf
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, encode_kernargs_clike, make_cmdbuf
 from tinygrad.runtime.support.hcq2 import make_binary_patch
 from tinygrad.uop.ops import sint, UOp
 from tinygrad.device import Compiled, BufferSpec, Buffer, Device
@@ -182,7 +182,8 @@ def sdma_copy(ctx, call):
  src_addr, dst_addr = call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs)
  return call.ins(SDMAOps.COPY, src=tuple(UOp.const(x, dtypes.uint32) for off in range(0, sz, ctx.max_copy_size) for x in (
    ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR),
-    ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0, *data64_le(src_addr+off), *data64_le(dst_addr+off))))
+    ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0,
+    *data64_le(src_addr+UOp.const(off, dtypes.uint64)), *data64_le(dst_addr+UOp.const(off, dtypes.uint64)))))

 def sdma_wait(ctx, ins, dst, val):
  op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \
@@ -287,14 +288,16 @@ def amd_build_program(prg:UOp) -> UOp:

 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())
+    super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb())

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

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

-  def _do_map(self, buf:HCQ2Buffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+  def _do_map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+
+  def _do_unmap(self, buf:HCQBuffer): self.dev.iface.unmap(buf)

@dataclass
 class AMDQueueDesc:
@@ -387,15 +390,24 @@ class KFDIface:
    return hcqbuf

  def free(self, mem):
+    self._unmap(mem)
+    if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
+    kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
+
+  def unmap(self, mem):
+    self._unmap(mem)
+    if getattr(mem, '_owns_kfd_handle', False): kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
+
+  def _unmap(self, mem):
    gpus = (ctypes.c_int32 * 1)(self.gpu_id)
    stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1)
    assert stm.n_success == 1
-    if mem.owner == self.dev:
-      if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
-      kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)

  def map(self, mem):
-    if mem.owner is not None and mem.owner._is_cpu(): return self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
+    if mem.owner is not None and mem.owner._is_cpu():
+      mapped = self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
+      mapped._owns_kfd_handle = True
+      return mapped

    c_gpus = (ctypes.c_int32 * 1)(self.gpu_id)
    stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(c_gpus), n_devices=1)
@@ -467,6 +479,7 @@ class PCIIface(PCIIfaceBase):

  def require_profile_mode(self): return True
  def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return True # TODO: account for WGP disablement on some asics.
+  def unmap(self, mem): self.free(mem)

  def _compute_props(self):
    self.ip_versions = self.dev_impl.ip_ver
@@ -507,11 +520,12 @@ class PCIIface(PCIIfaceBase):
    if drain_only: d.iface.dev_impl.ih.drain()
    else: d.iface.dev_impl.ih.interrupt_handler()

-    if reset and d.iface.dev_impl.recover():
+    if reset and d.iface.dev_impl.recover(force=True):
      cq = d.compute_queue
      for b in (cq.put_value, cq.read_ptr, cq.write_ptr): b._buf.view.view(fmt='Q')[0] = 0
      d.iface.dev_impl.gfx.setup_ring(*cq.params)
-      d.signal('timeline')._buf.cpu_view().mv.cast('Q')[0] = d.signal('value', 1).as_memoryview(force_zero_copy=True).cast('Q')[0] - 1
+      d.signal('timeline')._buf.cpu_view().mv.cast('Q')[0] = \
+        d.signal('value', 1).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] - 1

  def sleep(self, timeout):
    if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
@@ -537,9 +551,12 @@ class AMDDevice(HCQ2Compiled):
  ])

  timestamp_divider = 100.0  # AMD GPU clock: ticks/us
+  max_scratch_psize = 0

  ifaces = [KFDIface, PCIIface, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface)]

+  def device_props(self): return self.iface.props
+
  def is_am(self) -> bool: return isinstance(self.iface, (PCIIface,))
  def is_usb(self) -> bool: return False

@@ -576,7 +593,7 @@ 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_copy_queue = not getenv("AMD_DISABLE_SDMA")

    super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, can_recover=self.is_am(), arch=self.arch)

@@ -689,7 +706,7 @@ class AMDDevice(HCQ2Compiled):
    return tmpring

  def scratch_buffer(self, private_segment_size):
-    private_segment_size = max(private_segment_size, 128)
+    AMDDevice.max_scratch_psize = private_segment_size = max(private_segment_size, 128, AMDDevice.max_scratch_psize)
    if self.max_private_segment_size < private_segment_size:
      lanes_per_wave = 64 # wave64
      mem_alignment_size = 256 if self.target[0] != 9 else 1024
@@ -0,0 +1,36 @@
+import functools, math, pathlib
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from extra.llama_kernels import alloc_like, compile_hip
+
+@functools.cache
+def _custom_quantize_mxfp4(row_fp4:UOp, row_scale:UOp, col_fp4:UOp, col_scale:UOp, x:UOp, *, shuffle_row:bool, shuffle_col:bool) -> UOp:
+  M, N = math.prod(x.shape[:-1]), x.shape[-1]
+  assert M % 256 == 0 and N % 256 == 0, f"MXFP4 quantization requires multiples of 256, got {x.shape}"
+  name = f"quantize_mxfp4_{int(shuffle_row)}_{int(shuffle_col)}_{M}_{N}"
+  mem = M*N*2 + M*N + M*N//16  # read bf16, write row+col fp4 + e8m0
+  outputs = (row_fp4, row_scale, col_fp4, col_scale)
+  sink = UOp.sink(*(o.base for o in outputs), x.base,
+                  *(UOp(Ops.CUSTOM, dtypes.void, (o.base.index(0),), arg="") for o in outputs),
+                  UOp.special(256, "lidx0"), UOp.special(M//128, "gidx0"), UOp.special(N//64, "gidx1"),
+                  arg=KernelInfo(name, estimates=Estimates(ops=12*M*N, mem=mem)))
+  src = (pathlib.Path(__file__).parent/"quantize_mxfp4.cpp").read_text()
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
+    UOp(Ops.BINARY, arg=compile_hip(src, [f"-DKERNEL_NAME={name}", f"-DM_DIM={M}", f"-DN_DIM={N}",
+                                             f"-DSHUFFLE_ROWWISE_FP4_VALUE={int(shuffle_row)}",
+                                             f"-DSHUFFLE_COLWISE_FP4_VALUE={int(shuffle_col)}"]))))
+
+def quantize_mxfp4(x:Tensor, *, shuffle_row:bool=False, shuffle_col:bool=False, flatten_row:bool=False) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+  assert x.dtype == dtypes.bfloat16 and x.ndim >= 2, f"expected BF16 matrix, got {x.dtype} {x.shape}"
+  M, N = math.prod(x.shape[:-1]), x.shape[-1]
+  assert M % 256 == 0 and N % 256 == 0, f"MXFP4 quantization requires multiples of 256, got {x.shape}"
+  axis = x.uop.axis if isinstance(x.device, tuple) else None
+  row_axis = 0 if flatten_row and axis is not None else axis
+  col_axis = None if axis is None else (0 if axis == x.ndim-1 else 1)
+  outputs = (alloc_like((M, N//2) if flatten_row else (*x.shape[:-1], N//2), dtypes.uint8, x.device, row_axis),
+             alloc_like((M, N//32) if flatten_row else (*x.shape[:-1], N//32), dtypes.uint8, x.device, row_axis),
+             alloc_like((N, M//2), dtypes.uint8, x.device, col_axis),
+             alloc_like((N, M//32), dtypes.uint8, x.device, col_axis))
+  fxn = functools.partial(_custom_quantize_mxfp4, shuffle_row=shuffle_row, shuffle_col=shuffle_col)
+  return tuple(Tensor.custom_kernel(*outputs, x, fxn=fxn)[:4])
@@ -0,0 +1,226 @@
+// Copyright (c) 2025-2026, Advanced Micro Devices, Inc. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#include <hip/hip_runtime.h>
+#include <cstdint>
+
+#if !defined(KERNEL_NAME) || !defined(M_DIM) || !defined(N_DIM) || !defined(SHUFFLE_ROWWISE_FP4_VALUE) || \
+    !defined(SHUFFLE_COLWISE_FP4_VALUE)
+#error kernel dimensions and layouts must be defined
+#endif
+
+namespace {
+
+constexpr int BLOCK = 32;
+constexpr int TILE_M = 128;
+constexpr int TILE_N = 64;
+constexpr int THREADS = 256;
+constexpr int THREADS_PER_ROW = 8;
+constexpr int VALUES_PER_THREAD = 4;
+constexpr int SMEM_STRIDE = BLOCK + 2;
+constexpr int M = M_DIM;
+constexpr int N = N_DIM;
+constexpr int M_PACKED = M / 2;
+constexpr int N_PACKED = N / 2;
+constexpr int M_SCALES = M / BLOCK;
+constexpr int N_SCALES = N / BLOCK;
+constexpr bool SHUFFLE_ROWWISE_FP4 = SHUFFLE_ROWWISE_FP4_VALUE;
+constexpr bool SHUFFLE_COLWISE_FP4 = SHUFFLE_COLWISE_FP4_VALUE;
+
+static_assert(M % 256 == 0 && N % 256 == 0);
+
+struct Quantized4 {
+  uint16_t fp4;
+  uint8_t scale;
+};
+
+__device__ __forceinline__ float swizzle_xor1(float value) {
+  float result;
+  asm volatile("ds_swizzle_b32 %0, %1 offset:0x041f\n\ts_waitcnt lgkmcnt(0)" : "=v"(result) : "v"(value));
+  return result;
+}
+
+__device__ __forceinline__ float swizzle_xor2(float value) {
+  float result;
+  asm volatile("ds_swizzle_b32 %0, %1 offset:0x081f\n\ts_waitcnt lgkmcnt(0)" : "=v"(result) : "v"(value));
+  return result;
+}
+
+__device__ __forceinline__ float swizzle_xor4(float value) {
+  float result;
+  asm volatile("ds_swizzle_b32 %0, %1 offset:0x101f\n\ts_waitcnt lgkmcnt(0)" : "=v"(result) : "v"(value));
+  return result;
+}
+
+__device__ __forceinline__ float max8(float value) {
+  value = fmaxf(value, swizzle_xor4(value));
+  value = fmaxf(value, swizzle_xor2(value));
+  return fmaxf(value, swizzle_xor1(value));
+}
+
+__device__ __forceinline__ float4 load_bf16x4(const uint16_t* values) {
+  const uint32_t lo = *reinterpret_cast<const uint32_t*>(values);
+  const uint32_t hi = *reinterpret_cast<const uint32_t*>(values + 2);
+  return make_float4(__uint_as_float(lo << 16), __uint_as_float(lo & 0xffff0000u),
+                     __uint_as_float(hi << 16), __uint_as_float(hi & 0xffff0000u));
+}
+
+__device__ __forceinline__ void hadamard16(float4& value, int lane) {
+  const float a0 = value.x + value.y, a1 = value.x - value.y;
+  const float a2 = value.z + value.w, a3 = value.z - value.w;
+  value = make_float4(a0 + a2, a1 + a3, a0 - a2, a1 - a3);
+
+  const float4 xor1 = make_float4(swizzle_xor1(value.x), swizzle_xor1(value.y), swizzle_xor1(value.z), swizzle_xor1(value.w));
+  value = lane & 1 ? make_float4(xor1.x - value.x, xor1.y - value.y, xor1.z - value.z, xor1.w - value.w)
+                   : make_float4(xor1.x + value.x, xor1.y + value.y, xor1.z + value.z, xor1.w + value.w);
+
+  const float4 xor2 = make_float4(swizzle_xor2(value.x), swizzle_xor2(value.y), swizzle_xor2(value.z), swizzle_xor2(value.w));
+  value = lane & 2 ? make_float4(xor2.x - value.x, xor2.y - value.y, xor2.z - value.z, xor2.w - value.w)
+                   : make_float4(xor2.x + value.x, xor2.y + value.y, xor2.z + value.z, xor2.w + value.w);
+  value.x *= 0.25f;
+  value.y *= 0.25f;
+  value.z *= 0.25f;
+  value.w *= 0.25f;
+}
+
+__device__ __forceinline__ uint8_t e8m0_scale(float amax, float& scale) {
+  if (amax == 0.0f) {
+    scale = 1.0f;
+    return 127;
+  }
+
+  const uint32_t rounded = (__float_as_uint(amax) + 0x200000u) & 0xff800000u;
+  int exponent = static_cast<int>((rounded >> 23) & 0xff) - 129;
+  exponent = exponent < -127 ? -127 : exponent > 127 ? 127 : exponent;
+  scale = exponent == -127 ? __uint_as_float(0x00400000u) : __uint_as_float(static_cast<uint32_t>(exponent + 127) << 23);
+  return static_cast<uint8_t>(exponent + 127);
+}
+
+__device__ __forceinline__ uint16_t pack_fp4(float4 value, float scale) {
+  uint32_t lo = 0, hi = 0;
+  asm volatile("v_cvt_scalef32_pk_fp4_f32 %0, %1, %2, %3" : "+v"(lo) : "v"(value.x), "v"(value.y), "v"(scale));
+  asm volatile("v_cvt_scalef32_pk_fp4_f32 %0, %1, %2, %3" : "+v"(hi) : "v"(value.z), "v"(value.w), "v"(scale));
+  return static_cast<uint16_t>(lo | (hi << 8));
+}
+
+__device__ __forceinline__ Quantized4 quantize(float4 value, int lane) {
+  hadamard16(value, lane);
+  const float local_max = fmaxf(fmaxf(fabsf(value.x), fabsf(value.y)), fmaxf(fabsf(value.z), fabsf(value.w)));
+  float scale;
+  const uint8_t e8m0 = e8m0_scale(max8(local_max), scale);
+  return {pack_fp4(value, scale), e8m0};
+}
+
+__device__ __forceinline__ void store_scale(uint8_t* output, int row, int col, int cols, uint8_t value) {
+  const int tile = ((row >> 5) * (cols >> 3) + (col >> 3)) << 8;
+  const int offset = ((col & 3) << 6) + ((row & 15) << 2) + (((col >> 2) & 1) << 1) + ((row >> 4) & 1);
+  output[tile + offset] = value;
+}
+
+template<bool Shuffled>
+__device__ __forceinline__ void store_fp4(uint8_t* output, int row, int col, int packed_cols, uint16_t value) {
+  int index = row * packed_cols + col;
+  if constexpr (Shuffled) {
+    const int tile = (row >> 4) * (packed_cols << 4) + (col >> 5) * 512;
+    const int offset = ((col >> 4) & 1) * 256 + (row & 15) * 16 + (col & 15);
+    index = tile + offset;
+  }
+  *reinterpret_cast<uint16_t*>(output + index) = value;
+}
+
+__device__ __forceinline__ void load_tile(uint16_t* tile, const uint16_t* input, int tile_m, int tile_n) {
+  const int row = threadIdx.x / THREADS_PER_ROW;
+  const int col = threadIdx.x % THREADS_PER_ROW * VALUES_PER_THREAD;
+  const uint64_t packed = *reinterpret_cast<const uint64_t*>(input + (tile_m + row) * N + tile_n + col);
+  *reinterpret_cast<uint32_t*>(tile + row * SMEM_STRIDE + col) = static_cast<uint32_t>(packed);
+  *reinterpret_cast<uint32_t*>(tile + row * SMEM_STRIDE + col + 2) = static_cast<uint32_t>(packed >> 32);
+}
+
+__device__ __forceinline__ void quantize_row(uint16_t* tile, uint8_t* fp4_output, uint8_t* scale_output,
+                                             int tile_m, int tile_n, int local_row, int lane) {
+  const int row = tile_m + local_row;
+  const int col = lane * VALUES_PER_THREAD;
+  const Quantized4 result = quantize(load_bf16x4(tile + local_row * SMEM_STRIDE + col), lane);
+  store_fp4<SHUFFLE_ROWWISE_FP4>(fp4_output, row, (tile_n + col) / 2, N_PACKED, result.fp4);
+  if (lane == 0) store_scale(scale_output, row, tile_n / BLOCK, N_SCALES, result.scale);
+}
+
+__device__ __forceinline__ Quantized4 quantize_col(uint16_t* tile, int col, int lane) {
+  const int row = lane * VALUES_PER_THREAD;
+  return quantize(make_float4(
+    __uint_as_float(static_cast<uint32_t>(tile[(row + 0) * SMEM_STRIDE + col]) << 16),
+    __uint_as_float(static_cast<uint32_t>(tile[(row + 1) * SMEM_STRIDE + col]) << 16),
+    __uint_as_float(static_cast<uint32_t>(tile[(row + 2) * SMEM_STRIDE + col]) << 16),
+    __uint_as_float(static_cast<uint32_t>(tile[(row + 3) * SMEM_STRIDE + col]) << 16)), lane);
+}
+
+} // namespace
+
+extern "C" __global__ __launch_bounds__(THREADS, 8)
+void KERNEL_NAME(uint8_t* __restrict__ rowwise_fp4, uint8_t* __restrict__ rowwise_scale,
+                 uint8_t* __restrict__ colwise_fp4, uint8_t* __restrict__ colwise_scale,
+                 const uint16_t* __restrict__ input) {
+  __shared__ uint16_t tile[BLOCK * SMEM_STRIDE];
+  const int tid = threadIdx.x;
+  const int line = tid / THREADS_PER_ROW;
+  const int lane = tid % THREADS_PER_ROW;
+  const int block_m = blockIdx.x * TILE_M;
+  const int block_n = blockIdx.y * TILE_N;
+
+  if constexpr (!SHUFFLE_COLWISE_FP4) {
+    uint16_t col_fp4[TILE_N / BLOCK][TILE_M / BLOCK];
+    uint8_t col_scale[TILE_N / BLOCK][TILE_M / BLOCK];
+
+    for (int chunk_m = 0; chunk_m < TILE_M / BLOCK; chunk_m++) {
+      for (int chunk_n = 0; chunk_n < TILE_N / BLOCK; chunk_n++) {
+        const int tile_m = block_m + chunk_m * BLOCK;
+        const int tile_n = block_n + chunk_n * BLOCK;
+        load_tile(tile, input, tile_m, tile_n);
+        __syncthreads();
+
+        quantize_row(tile, rowwise_fp4, rowwise_scale, tile_m, tile_n, line, lane);
+        const Quantized4 result = quantize_col(tile, line, lane);
+        col_fp4[chunk_n][chunk_m] = result.fp4;
+        col_scale[chunk_n][chunk_m] = result.scale;
+        __syncthreads();
+      }
+    }
+
+    for (int chunk_n = 0; chunk_n < TILE_N / BLOCK; chunk_n++) {
+      for (int chunk_m = 0; chunk_m < TILE_M / BLOCK; chunk_m++)
+        tile[line * BLOCK + chunk_m * THREADS_PER_ROW + lane] = col_fp4[chunk_n][chunk_m];
+      __syncthreads();
+
+      for (int round = 0; round < BLOCK / THREADS_PER_ROW; round++) {
+        const int col = round * THREADS_PER_ROW + tid / BLOCK;
+        const int row_pair = tid % BLOCK;
+        *reinterpret_cast<uint16_t*>(colwise_fp4 + (block_n + chunk_n * BLOCK + col) * M_PACKED + block_m / 2 + row_pair * 2) =
+          tile[col * BLOCK + row_pair];
+      }
+
+      if (lane == 0) {
+        const int col = block_n + chunk_n * BLOCK + line;
+        for (int chunk_m = 0; chunk_m < TILE_M / BLOCK; chunk_m++)
+          store_scale(colwise_scale, col, block_m / BLOCK + chunk_m, M_SCALES, col_scale[chunk_n][chunk_m]);
+      }
+      __syncthreads();
+    }
+  } else {
+    for (int chunk_m = 0; chunk_m < TILE_M / BLOCK; chunk_m++) {
+      for (int chunk_n = 0; chunk_n < TILE_N / BLOCK; chunk_n++) {
+        const int tile_m = block_m + chunk_m * BLOCK;
+        const int tile_n = block_n + chunk_n * BLOCK;
+        load_tile(tile, input, tile_m, tile_n);
+        __syncthreads();
+
+        quantize_row(tile, rowwise_fp4, rowwise_scale, tile_m, tile_n, line, lane);
+        const int row = lane * VALUES_PER_THREAD;
+        const int col = tile_n + line;
+        const Quantized4 result = quantize_col(tile, line, lane);
+        store_fp4<true>(colwise_fp4, col, (tile_m + row) / 2, M_PACKED, result.fp4);
+        if (lane == 0) store_scale(colwise_scale, col, tile_m / BLOCK, M_SCALES, result.scale);
+        __syncthreads();
+      }
+    }
+  }
+}
@@ -0,0 +1,49 @@
+import functools, math
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, KernelInfo
+from tinygrad.renderer import Estimates
+from extra.llama_kernels import alloc_like
+
+LOG2E = 1.4426950408889634
+
+@functools.cache
+def _custom_swiglu(out:UOp, x_w13:UOp) -> UOp:
+  rows, hidden = math.prod(x_w13.shape[:-1]), x_w13.shape[-1]//2
+  n_elems = rows * hidden
+  out, x_w13 = out.reshape(n_elems), x_w13.reshape(rows, 2*hidden)
+  i = UOp.range(n_elems, 0)
+  row, col = i // hidden, i % hidden
+  act, gate = x_w13[row, col].cast(dtypes.float), x_w13[row, hidden+col].cast(dtypes.float)
+  sigmoid = (1.0 + (-LOG2E * act).exp2()).reciprocal()
+  store = out[i].store((act * sigmoid * gate).cast(out.dtype))
+  return store.end(i).sink(arg=KernelInfo(f"swiglu_fwd_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=6*n_elems)))
+
+@functools.cache
+def _custom_swiglu_bwd(grad_out:UOp, x_w13:UOp, grad_act:UOp) -> UOp:
+  rows, hidden = math.prod(x_w13.shape[:-1]), x_w13.shape[-1]//2
+  n_elems = rows * hidden
+  grad_out, x_w13, grad_act = grad_out.reshape(rows, 2*hidden), x_w13.reshape(rows, 2*hidden), grad_act.reshape(n_elems)
+  i = UOp.range(n_elems, 0)
+  row, col = i // hidden, i % hidden
+  act, gate = x_w13[row, col].cast(dtypes.float), x_w13[row, hidden+col].cast(dtypes.float)
+  grad = grad_act[i].cast(dtypes.float)
+  sigmoid = (1.0 + (-LOG2E * act).exp2()).reciprocal()
+  silu = act * sigmoid
+  dact = grad_out[row, col].store((grad * (sigmoid + silu * (1.0 - sigmoid)) * gate).cast(grad_out.dtype))
+  dgate = grad_out.after(dact)[row, hidden+col].store((grad * silu).cast(grad_out.dtype))
+  return dgate.end(i).sink(arg=KernelInfo(f"swiglu_bwd_{n_elems}", estimates=Estimates(ops=10*n_elems, mem=10*n_elems)))
+
+def _swiglu_bwd(gradient:UOp, kernel:UOp):
+  _, x_w13 = kernel.src[1:]
+  axis = x_w13.axis if isinstance(x_w13.device, tuple) else None
+  grad_out = alloc_like(x_w13.shape, dtypes.bfloat16, x_w13.device, axis)
+  grad_out, *_ = Tensor.custom_kernel(grad_out, Tensor(x_w13, device=x_w13.device), Tensor(gradient, device=x_w13.device),
+                                      fxn=_custom_swiglu_bwd)
+  return (None, grad_out.uop)
+
+def swiglu(x_w13:Tensor) -> Tensor:
+  assert x_w13.dtype == dtypes.bfloat16 and x_w13.ndim >= 2 and x_w13.shape[-1] % 32 == 0
+  *prefix, two_k = x_w13.shape
+  axis = x_w13.uop.axis if isinstance(x_w13.device, tuple) else None
+  out = alloc_like((*prefix, two_k//2), dtypes.bfloat16, x_w13.device, axis)
+  return Tensor.custom_kernel(out, x_w13, fxn=_custom_swiglu, grad_fxn=_swiglu_bwd)[0]
@@ -0,0 +1,217 @@
+# Runbook: Llama 3 8B Training on DigitalOcean MI350X
+
+## Machine Specs
+- 8x MI350X GPUs (gfx950, device ID 75b0), 288GB VRAM each
+- 2TB RAM, 192 CPUs, 2TB disk
+- ROCm 7.14 at `/opt/rocm` (NOT `/opt/rocm-7.1.1` like the submission scripts assume)
+- Python 3.12
+
+## Phase 1: System Setup
+
+### 1.1 Install packages
+```bash
+apt-get update
+apt-get install -y python3-pip python3-venv git tmux rclone clang
+```
+
+### 1.2 Install Python deps
+```bash
+python3 -m pip install --break-system-packages --ignore-installed typing-extensions numpy tqdm wandb tiktoken sentencepiece
+```
+Note: `--ignore-installed typing-extensions` is needed because the base image ships typing-extensions 4.10.0 without a RECORD file, so pip cannot uninstall it.
+
+### 1.3 Install ROCm dev headers
+The base image has ROCm runtime but NOT the HIP dev headers. Need:
+```bash
+apt-get install -y amdrocm-core-dev
+```
+This installs `hip/hip_runtime.h` at `/opt/rocm/core-7.14/include/hip/hip_runtime.h`.
+The symlink `/opt/rocm/include` → `/opt/rocm/core-7.14/include` makes it available at `/opt/rocm/include/hip/hip_runtime.h`.
+
+### 1.4 Configure ROCm comgr
+ROCm 7.14 ships comgr 3.3 at `/opt/rocm/lib/libamd_comgr.so`. tinygrad's DLL loader needs explicit env vars to find it (it searches for `libcomgr.so*` by default, not `libamd_comgr.so*`). Set these in the run command:
+```bash
+export COMGR_PATH=/opt/rocm/lib/libamd_comgr.so
+export COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so
+```
+Also add ROCm libs to ldconfig so comgr's shared library dependencies resolve:
+```bash
+cat > /etc/ld.so.conf.d/rocm.conf << 'EOF'
+/opt/rocm/lib
+/opt/rocm/lib/llvm/lib
+/opt/rocm/lib/rocm_sysdeps/lib
+EOF
+ldconfig
+```
+
+### 1.5 Install geohot tmux config
+```bash
+curl -sL https://raw.githubusercontent.com/geohot/configuration/master/.tmux.conf -o ~/.tmux.conf
+```
+
+### 1.6 Reload amdgpu driver
+tinygrad's HCQ backend needs `/dev/kfd` which is created by the amdgpu kernel driver.
+If the driver was unloaded, reload it:
+```bash
+modprobe amdgpu
+ls /dev/kfd  # should exist
+```
+
+## Phase 2: Clone tinygrad
+```bash
+cd /root
+git clone https://github.com/tinygrad/tinygrad.git
+cd tinygrad
+python3 -m pip install --break-system-packages -e .
+```
+
+## Phase 3: Download C4 Dataset
+
+The C4 data is on the MLCommons Cloudflare R2 bucket in Megatron-LM indexed format.
+
+```bash
+rclone config create mlc-training s3 provider=Cloudflare \
+  access_key_id=76ea42eadb867e854061a1806220ee1e \
+  secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 \
+  endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
+
+mkdir -p /raid/datasets/c4-8b
+rclone copy mlc-training:mlcommons-training-wg-public/llama3_1/datasets/c4/llama3_1_8b/ /raid/datasets/c4-8b/ -P
+```
+
+Files downloaded (~85GB total, ~6 minutes):
+- `c4-train.en_6_text_document.bin` (79 GB)
+- `c4-train.en_6_text_document.idx` (870 MB)
+- `c4-validation-91205-samples.en_text_document.bin` (159 MB)
+- `c4-validation-91205-samples.en_text_document.idx` (1.8 MB)
+- `LICENSE.txt`, `NOTICE.txt`
+
+**Wait for rclone to fully complete before starting training.** Starting training while the dataset is still downloading will read a truncated .bin file, causing `ValueError: all input arrays must have the same shape` in the dataloader. The stale `.index_cache` and `.blend_cache` files must also be deleted if this happens:
+```bash
+rm -f /raid/datasets/c4-8b/*.index_cache /raid/datasets/c4-8b/*.blend_cache
+```
+
+## Phase 4: wandb Login
+```bash
+wandb login
+```
+Enter API key from https://wandb.ai/authorize
+
+Alternatively, pass the key directly:
+```bash
+wandb login <API_KEY>
+```
+
+## Phase 5: Run Training
+
+Run training in tmux so it survives SSH disconnects:
+```bash
+tmux new-session -d -s train 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=AMD:HIP ROCM_PATH=/opt/rocm WANDB=1 bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh 2>&1 | tee /root/train.log'
+```
+Attach with `tmux attach -t train`.
+
+### 5.1 Smoke test (beam search, 2 layers, real data)
+Always run beam first to validate the pipeline:
+```bash
+tmux new-session -d -s beam 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=AMD:HIP ROCM_PATH=/opt/rocm bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh 2>&1 | tee /root/beam.log'
+```
+
+The beam test runs 10 training steps with 2 layers. Expected results:
+- ~0.29s per step after warmup
+- ~700K GFLOPS, ~7% MFU (low because only 2 layers)
+- ~380 GB VRAM used
+- Loss stable at ~12.55 with random init
+
+### 5.2 Full training run
+```bash
+tmux new-session -d -s train 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=AMD:HIP ROCM_PATH=/opt/rocm WANDB=1 bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh 2>&1 | tee /root/train.log'
+```
+
+## Environment Variable Reference
+
+| Variable | Value | Why |
+|---|---|---|
+| `COMGR_PATH` | `/opt/rocm/lib/libamd_comgr.so` | tinygrad's DLL loader needs explicit path to find comgr 3.3 |
+| `COMGR_3_PATH` | `/opt/rocm/lib/libamd_comgr.so` | comgr 3.x uses a separate `comgr_3` module with its own path var |
+| `CC` | `/opt/rocm/core-7.14/lib/llvm/bin/clang` | System clang doesn't know gfx950; must use ROCm's bundled clang |
+| `DEV` | `AMD:HIP` | Force HIPRenderer (comgr-based) over HIPCCRenderer (hipcc subprocess) |
+| `ROCM_PATH` | `/opt/rocm` | Script defaults to `/opt/rocm-7.1.1` which doesn't exist |
+| `WANDB` | `1` | Enable wandb logging (off by default) |
+
+## Architecture
+
+| Component | Source file |
+|---|---|
+| Model | `examples/mlperf/models/flat_llama.py` — FlatTransformer, FP8 MXFP4 weights, fused QKV, flash attention |
+| Trainer | `examples/mlperf/model_train.py` → `train_llama3()` |
+| Optimizer | `examples/mlperf/optim.py` — GradAccClipAdamW, master weights, FP8 re-quant |
+| LR schedule | `examples/mlperf/lr_schedulers.py` — CosineAnnealingLRWithWarmup |
+| Dataloader | `examples/mlperf/dataloader.py` — Megatron-LM indexed bin format |
+| ASM GEMM | `extra/gemm/cdna_asm_gemm.py` — gfx950 MFMA assembly, MXFP4 |
+| Flash attention | `extra/thunder/amd/fa.py` |
+| Fused kernels | `extra/llama_kernels/` — rmsnorm, silu, quantize, fused_ce |
+| GPU driver | `tinygrad/runtime/ops_amd.py` — HCQ, direct KFD ioctl |
+| Renderer | `tinygrad/renderer/cstyle.py` — HIPRenderer for gfx950 |
+| comgr compiler | `tinygrad/runtime/support/compiler_amd.py` — HIPCompiler using comgr 3.3 |
+
+## Troubleshooting
+
+### `'hip/hip_runtime.h' file not found`
+Install `amdrocm-core-dev`:
+```bash
+apt-get install -y amdrocm-core-dev
+```
+
+### `'gfx950' is not a recognized processor` + LLVM crash
+System clang doesn't know gfx950. Set `CC=/opt/rocm/core-7.14/lib/llvm/bin/clang`.
+
+### `comgr not available: try setting COMGR_PATH?`
+Add ROCm libs to ldconfig and set `COMGR_PATH` and `COMGR_3_PATH`:
+```bash
+# /etc/ld.so.conf.d/rocm.conf should contain /opt/rocm/lib paths
+ldconfig
+```
+
+### `comgr not available: try setting COMGR_3_PATH?`
+comgr 3.x uses a separate module. Set `COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so` too.
+
+### `No such file or directory: 'clang'`
+Install clang: `apt-get install -y clang` (for CPU compilation).
+For gfx950 HIP compilation, comgr (not clang) is used — ensure the ROCm 7.14 comgr 3.3 is properly loaded via `COMGR_PATH` and `COMGR_3_PATH`.
+
+## Appendix: KVM Virtualization Observations
+
+### Virtualization detection
+```
+$ systemd-detect-virt
+kvm
+$ lspci -nn | grep AMD
+83:00.0 ... Device [1002:75b0]
+```
+CPU flags include `hypervisor`. `dmesg` shows `Hypervisor detected: KVM`.
+
+### Working path: amdgpu driver (KFDIface)
+The amdgpu driver loads on boot and binds to all 8 GPUs, creating `/dev/kfd` and 64 renderD nodes (`/dev/dri/renderD128` through `/dev/dri/renderD191`). tinygrad's `KFDIface` enumerates GPUs through `/sys/devices/virtual/kfd/kfd/topology/nodes` and uses `/dev/kfd` for ioctl. No PCI device ID patching is needed — the KFD path does not use `PCIIface` or `AMDev._run_discovery()`.
+
+This is the working configuration. No code changes to tinygrad are required.
+
+### PCIIface path (does not work on this VM)
+For reference, the `PCIIface` path was also explored but does not work in this KVM guest:
+
+- `PCIIface` in `ops_amd.py` does not list device ID `0x75b0`. Adding it allows PCI detection but `AMDev._run_discovery()` fails because the VRAM BAR reads all `0xFF`.
+- This was observed with the GPU unbound from any driver, after PCI reset, and with VFIO bound.
+- VFIO binding (`vfio-pci` with `enable_unsafe_noiommu_mode=1`) succeeded but VRAM BAR still reads all `0xFF`.
+- No IOMMU in guest — `dmesg` has no `AMD-Vi` entries, PCI devices have no `iommu_group` symlink.
+
+### amdgpu driver behavior
+On first boot, amdgpu loaded and bound to all 8 GPUs. On one boot it failed to initialize:
+```
+[  799.780369] amdgpu 0000:83:00.0: Failed to alloc msi vectors
+[  799.781476] amdgpu 0000:83:00.0: sw_init of IP block <vega20_ih> failed -22
+[  799.782724] amdgpu 0000:83:00.0: amdgpu_device_ip_init failed
+[  799.793885] amdgpu 0000:83:00.0: Fatal error during GPU init
+```
+On a subsequent boot, amdgpu initialized successfully (SMU initialized, VRAM ready). After unbinding all 8 GPUs from amdgpu, `rmmod amdgpu` wedged the module (stuck in "Unloading" state in `/proc/modules`), requiring a full VM reboot.
+
+### No fan control
+No `fan*` or `pwm*` hwmon entries exist. Only `temp*`, `power*`, `freq*` are exposed. GPU temps read 56-63°C, power ~265W per GPU.
@@ -2,7 +2,7 @@ import math, pathlib, functools, struct

 from tinygrad import Device, Tensor
 from tinygrad.dtype import DTypeLike, dtypes
-from tinygrad.helpers import DEBUG
+from tinygrad.helpers import DEBUG, getenv
 from tinygrad.renderer import Estimates
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 from tinygrad.runtime.support.elf import elf_loader
@@ -110,7 +110,49 @@ def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
  return _sharded_empty(ref.shape, ref, axis)

@functools.cache
-def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink):
+def _windowed_lse(xq:Tensor, xk:Tensor, sinks, W:int) -> Tensor:
+  B, N, H, hd = xq.shape
+  H_KV = xk.shape[2]; R = H // H_KV; nb = N // W; sm = hd ** -0.5
+  q = xq.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
+  k = xk.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
+  k_prev = k.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
+  sc_d = (q @ k.transpose(-1, -2)) * sm
+  sc_p = (q @ k_prev.transpose(-1, -2)) * sm
+  li, lj = Tensor.arange(W).reshape(W, 1), Tensor.arange(W).reshape(1, W)
+  pv = (Tensor.arange(nb).reshape(nb, 1, 1) >= 1)
+  sc_d = (lj <= li).where(sc_d, -float("inf"))
+  sc_p = ((li < lj) & pv).where(sc_p, -float("inf"))
+  m = sc_d.max(-1, keepdim=True).maximum(sc_p.max(-1, keepdim=True))
+  if sinks is not None: m = m.maximum(sinks.reshape(1, H_KV, R, 1, 1, 1).float())
+  denom = (sc_d - m).exp().sum(-1, keepdim=True) + (sc_p - m).exp().sum(-1, keepdim=True)
+  if sinks is not None: denom = denom + (sinks.reshape(1, H_KV, R, 1, 1, 1).float() - m).exp()
+  return (m + denom.log()).reshape(B, H, N).unsqueeze(2)  # (B, H, 1, N), matches saved l_vec
+
+def _windowed_delta(xq:Tensor, xk:Tensor, xv:Tensor, do:Tensor, sinks, W:int) -> Tensor:
+  B, N, H, hd = xq.shape
+  H_KV = xk.shape[2]; R = H // H_KV; nb = N // W; sm = hd ** -0.5
+  q = xq.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
+  k = xk.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
+  v = xv.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
+  dob = do.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
+  k_prev = k.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
+  v_prev = v.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
+  sc_d = (q @ k.transpose(-1, -2)) * sm
+  sc_p = (q @ k_prev.transpose(-1, -2)) * sm
+  li, lj = Tensor.arange(W).reshape(W, 1), Tensor.arange(W).reshape(1, W)
+  pv = (Tensor.arange(nb).reshape(nb, 1, 1) >= 1)
+  sc_d = (lj <= li).where(sc_d, -float("inf"))
+  sc_p = ((li < lj) & pv).where(sc_p, -float("inf"))
+  m = sc_d.max(-1, keepdim=True).maximum(sc_p.max(-1, keepdim=True))
+  if sinks is not None: m = m.maximum(sinks.reshape(1, H_KV, R, 1, 1, 1).float())
+  e_d, e_p = (sc_d - m).exp(), (sc_p - m).exp()
+  denom = e_d.sum(-1, keepdim=True) + e_p.sum(-1, keepdim=True)
+  if sinks is not None: denom = denom + (sinks.reshape(1, H_KV, R, 1, 1, 1).float() - m).exp()
+  o = ((e_d / denom) @ v) + ((e_p / denom) @ v_prev)
+  delta = (dob * o).sum(-1)
+  return delta.reshape(B, H, N).unsqueeze(2)
+
+def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink, window=0):
  def grad(dou:UOp, ker:UOp) -> tuple:
    do = Tensor(dou, device=dou.device)
    attn = Tensor(ker.src[1].after(ker), device=ker.src[1].device)
@@ -118,6 +160,8 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
    xq = Tensor(ker.src[3], device=ker.src[3].device)
    xk = Tensor(ker.src[4], device=ker.src[4].device)
    xv = Tensor(ker.src[5], device=ker.src[5].device)
+    if window:
+      l_vec = _windowed_lse(xq, xk, Tensor(ker.src[6], device=ker.src[6].device) if has_sink else None, window)

    dq = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t)
    GROUP_SIZE = H_local // H_KV_local
@@ -128,8 +172,10 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
    # delta_vec = (do * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
    delta_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
    delta_vec, dq = Tensor.custom_kernel(delta_vec, dq, attn, do, fxn=functools.partial(custom_fa_backward_pre, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:2]
+    if window:
+      delta_vec = _windowed_delta(xq, xk, xv, do, Tensor(ker.src[6], device=ker.src[6].device) if has_sink else None, window)

-    dq, dk_partial, dv_partial = Tensor.custom_kernel(dq, dk_partial, dv_partial, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:3]
+    dq, dk_partial, dv_partial = Tensor.custom_kernel(dq, dk_partial, dv_partial, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, window=window))[:3]

    if D == 64:
      dq = dq.reshape(B, H, N//16, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2).permute(0, 1, 2, 8, 9, 10, 11, 3, 4, 6, 7, 5, 12).reshape(B, H, N, D).transpose(1, 2)
@@ -149,7 +195,7 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
  return grad

 # TODO: remove write_flat once scheduler can remove reshapes between custom_kernel. TestCustomKernel.test_simple_reshape
-def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False, write_flat:bool=False, sinks:Tensor|None=None):
+def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False, write_flat:bool=False, sinks:Tensor|None=None, window:int=0):
  assert attn_mask is None, "attn_mask not supported"
  assert is_causal, "only causal attention supported"

@@ -176,18 +222,18 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
  attn = _sharded_empty((B, N, H * D), xq, axis=shard_axis) if write_flat else _sharded_empty_like(xq, axis=shard_axis)
  l_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)

-  grad = _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink)
+  grad = _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink, window=window)

  fwd_inputs = (attn, l_vec, xq, xk, xv) + ((sinks,) if has_sink else ())
-  attn, l_vec = Tensor.custom_kernel(*fwd_inputs, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, has_sink=has_sink), grad_fxn=grad)[:2]
+  attn, l_vec = Tensor.custom_kernel(*fwd_inputs, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, has_sink=has_sink, window=window), grad_fxn=grad)[:2]

  return attn, attn, l_vec

@functools.cache
-def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None, *, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, has_sink:bool=True):
+def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None, *, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, has_sink:bool=True, window:int=0):
  code = (pathlib.Path(__file__).parent / "fa_fwd_causal.cpp").read_text()
  compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math",
-                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DATTN_SINK={int(has_sink)}"]
+                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DATTN_SINK={int(has_sink)}", f"-DWINDOW={window}"]

  Q_BLOCK_SIZE = 32
  NUM_WARPS = 8
@@ -206,10 +252,11 @@ def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None
                  arg=KernelInfo(name="custom_fa_forward", estimates=estimates))

  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-  lib = bytearray(lib)
-  rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
-  struct.pack_into('<I', lib, rodata_off, 160000)
-  lib = bytes(lib)
+  if not getenv("NO_HIPCC"):
+    lib = bytearray(lib)
+    rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
+    struct.pack_into('<I', lib, rodata_off, 160000)
+    lib = bytes(lib)

  return UOp(Ops.PROGRAM,
             src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
@@ -236,19 +283,20 @@ def custom_fa_backward_pre(delta_vec:UOp, dq:UOp, o:UOp, do:UOp, device:str, arc
                  arg=KernelInfo(name="custom_fa_backward_pre", estimates=estimates))

  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-  lib = bytearray(lib)
-  rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
-  struct.pack_into('<I', lib, rodata_off, 160000)
-  lib = bytes(lib)
+  if not getenv("NO_HIPCC"):
+    lib = bytearray(lib)
+    rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
+    struct.pack_into('<I', lib, rodata_off, 160000)
+    lib = bytes(lib)

  return UOp(Ops.PROGRAM,
             src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))

@functools.cache
-def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_vec:UOp, delta_vec:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
+def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_vec:UOp, delta_vec:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, window:int=0):
  code = (pathlib.Path(__file__).parent / "fa_bwd_causal.cpp").read_text()
  compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math",
-                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}"]
+                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DWINDOW={window}"]

  BLOCK_SIZE_KV = 256
  GROUP_SIZE = H // H_KV
@@ -268,10 +316,11 @@ def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_ve
                  arg=KernelInfo(name="custom_fa_backward", estimates=estimates))

  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-  lib = bytearray(lib)
-  rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
-  struct.pack_into('<I', lib, rodata_off, 160000)
-  lib = bytes(lib)
+  if not getenv("NO_HIPCC"):
+    lib = bytearray(lib)
+    rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
+    struct.pack_into('<I', lib, rodata_off, 160000)
+    lib = bytes(lib)

  return UOp(Ops.PROGRAM,
             src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
@@ -298,10 +347,11 @@ def custom_fa_backward_post(dq_out:UOp, dq_in:UOp, device:str, arch:str, B:int,
                  arg=KernelInfo(name="custom_fa_backward_post", estimates=estimates))

  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-  lib = bytearray(lib)
-  rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
-  struct.pack_into('<I', lib, rodata_off, 160000)
-  lib = bytes(lib)
+  if not getenv("NO_HIPCC"):
+    lib = bytearray(lib)
+    rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
+    struct.pack_into('<I', lib, rodata_off, 160000)
+    lib = bytes(lib)

  return UOp(Ops.PROGRAM,
             src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
@@ -43,6 +43,10 @@ constexpr int SLICE_QO = 32;
 constexpr int DOT_SLICE_QO = 16;
 constexpr int WARP_SIZE_KV = 64; // warp size for KV
 constexpr bool causal = true;
+// WINDOW>0: sliding-window backward (query i sees keys in [i-WINDOW+1, i])
+#ifndef WINDOW
+#define WINDOW 0
+#endif

 #define NUM_WARPS 4
 #define NUM_THREADS (kittens::WARP_THREADS * NUM_WARPS)
@@ -88,7 +92,12 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
  const int k_start_min = j_min * WARP_SIZE_KV;
  // first Q step that can overlap this K_span:
  const int first_step = max(0, k_start_min / STEP_QO);
+#if WINDOW
+  // cap the Q loop, padded by 2 masked steps: the epilogue's deferred dq path miscomputes in-window tail queries
+  const int num_steps_per_head = min(total_steps_per_head - first_step, (BLOCK_SIZE_KV + WINDOW) / STEP_QO + 2);
+#else
  const int num_steps_per_head = total_steps_per_head - first_step;
+#endif
  const int num_steps = num_steps_per_head * HEADS_PER_WG;
  const int k_pos = j * WARP_SIZE_KV;

@@ -380,6 +389,13 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
            mov<0, 1, neg_inf_v>(P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
+#if WINDOW
+          // window lower boundary, mirror of the causal edge
+          } else if (q_pos - k_pos == WINDOW) {
+            make_window<0, 0, neg_inf_v>(P_ij, P_ij);
+          } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -640,6 +656,13 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
            make_causal<0, 1, neg_inf_v>(P_ij, P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
+#if WINDOW
+          } else if (q_pos - k_pos == WINDOW) {
+            mov<0, 0, neg_inf_v>(P_ij);
+            make_window<0, 1, neg_inf_v>(P_ij, P_ij);
+          } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -899,6 +922,14 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
            // Apply the causal mask to [0, 2] and set [0, 3:4] to -inf
            make_causal<0, 2, neg_inf_v>(P_ij, P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
+#if WINDOW
+          } else if (q_pos - k_pos == WINDOW) {
+            mov<0, 0, neg_inf_v>(P_ij);
+            mov<0, 1, neg_inf_v>(P_ij);
+            make_window<0, 2, neg_inf_v>(P_ij, P_ij);
+          } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -1157,6 +1188,15 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
          } else if (q_pos == k_pos) {
            // Apply the causal mask to [0, 3]
            make_causal<0, 3, neg_inf_v>(P_ij, P_ij);
+#if WINDOW
+          } else if (q_pos - k_pos == WINDOW) {
+            mov<0, 0, neg_inf_v>(P_ij);
+            mov<0, 1, neg_inf_v>(P_ij);
+            mov<0, 2, neg_inf_v>(P_ij);
+            make_window<0, 3, neg_inf_v>(P_ij, P_ij);
+          } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -1436,6 +1476,13 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
            mov<0, 1, neg_inf_v>(P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
+#if WINDOW
+          // window lower boundary, mirror of the causal edge
+          } else if (q_pos - k_pos == WINDOW) {
+            make_window<0, 0, neg_inf_v>(P_ij, P_ij);
+          } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -1699,6 +1746,13 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
            make_causal<0, 1, neg_inf_v>(P_ij, P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
+#if WINDOW
+          } else if (q_pos - k_pos == WINDOW) {
+            mov<0, 0, neg_inf_v>(P_ij);
+            make_window<0, 1, neg_inf_v>(P_ij, P_ij);
+          } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -1958,6 +2012,14 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
            // Apply the causal mask to [0, 2] and set [0, 3:4] to -inf
            make_causal<0, 2, neg_inf_v>(P_ij, P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
+#if WINDOW
+          } else if (q_pos - k_pos == WINDOW) {
+            mov<0, 0, neg_inf_v>(P_ij);
+            mov<0, 1, neg_inf_v>(P_ij);
+            make_window<0, 2, neg_inf_v>(P_ij, P_ij);
+          } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -2216,6 +2278,15 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
          } else if (q_pos == k_pos) {
            // Apply the causal mask to [0, 3]
            make_causal<0, 3, neg_inf_v>(P_ij, P_ij);
+#if WINDOW
+          } else if (q_pos - k_pos == WINDOW) {
+            mov<0, 0, neg_inf_v>(P_ij);
+            mov<0, 1, neg_inf_v>(P_ij);
+            mov<0, 2, neg_inf_v>(P_ij);
+            make_window<0, 3, neg_inf_v>(P_ij, P_ij);
+          } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
          }
        }
        mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -2487,6 +2558,12 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
            mov<0, 1, neg_inf_v>(P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
+#if WINDOW
+        } else if (q_pos - k_pos == WINDOW) {
+            make_window<0, 0, neg_inf_v>(P_ij, P_ij);
+        } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
        }
      }
      mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -2748,6 +2825,13 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
            make_causal<0, 1, neg_inf_v>(P_ij, P_ij);
            mov<0, 2, neg_inf_v>(P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
+#if WINDOW
+        } else if (q_pos - k_pos == WINDOW) {
+            mov<0, 0, neg_inf_v>(P_ij);
+            make_window<0, 1, neg_inf_v>(P_ij, P_ij);
+        } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
        }
      }
      mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -3004,6 +3088,14 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
            // Apply the causal mask to [0, 2] and set [0, 3:4] to -inf
            make_causal<0, 2, neg_inf_v>(P_ij, P_ij);
            mov<0, 3, neg_inf_v>(P_ij);
+#if WINDOW
+        } else if (q_pos - k_pos == WINDOW) {
+            mov<0, 0, neg_inf_v>(P_ij);
+            mov<0, 1, neg_inf_v>(P_ij);
+            make_window<0, 2, neg_inf_v>(P_ij, P_ij);
+        } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
        }
      }
      mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -3260,6 +3352,15 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
        } else if (q_pos == k_pos) {
            // Apply the causal mask to [0, 3]
            make_causal<0, 3, neg_inf_v>(P_ij, P_ij);
+#if WINDOW
+        } else if (q_pos - k_pos == WINDOW) {
+            mov<0, 0, neg_inf_v>(P_ij);
+            mov<0, 1, neg_inf_v>(P_ij);
+            mov<0, 2, neg_inf_v>(P_ij);
+            make_window<0, 3, neg_inf_v>(P_ij, P_ij);
+        } else if (q_pos - k_pos > WINDOW) {
+            mov<neg_inf_v>(P_ij);
+#endif
        }
      }
      mul<0, 2>(P_ij, P_ij, P_SCALE_FACTOR);
@@ -34,6 +34,10 @@ constexpr int ATTN_D = 128; // dimension
 constexpr int Q_BLOCK_SIZE = 32; // q block size
 constexpr int KV_BLOCK_SIZE = 64; // kv block size
 constexpr bool causal = true;
+// WINDOW>0: sliding-window attention, query i attends keys in [i-WINDOW+1, i]
+#ifndef WINDOW
+#define WINDOW 0
+#endif

 #define NUM_WARPS 8
 #define NUM_THREADS (kittens::WARP_THREADS * NUM_WARPS)
@@ -82,11 +86,26 @@ template<typename T=float, typename L=col_l, typename S=rt_16x32_4_s> using attn

 /**********************************************************/
 template<int THR_X, int THR_Y>
-__device__ inline void mask_vec2_imm(uint32_t rel_vgpr, uint32_t neg_inf_vgpr,
+__device__ inline void mask_vec2_imm(uint32_t rel_vgpr, uint32_t rel_hi_vgpr, uint32_t neg_inf_vgpr,
                                     uint32_t& x_ref, uint32_t& y_ref) {

    uint64_t x_mask, y_mask;
-    // uint32_t ox, oy;
+#if WINDOW
+    // causal+window in one asm block to not disturb register allocation
+    asm volatile(
+        "v_cmp_lt_i32_e64 %0, %4, %5\n\t"
+        "v_cmp_lt_i32_e64 %1, %4, %7\n\t"
+        "v_cndmask_b32_e64 %2, %2, %6, %0\n\t"
+        "v_cndmask_b32_e64 %3, %3, %6, %1\n\t"
+        "v_cmp_ge_i32_e64 %0, %8, %5\n\t"
+        "v_cmp_ge_i32_e64 %1, %8, %7\n\t"
+        "v_cndmask_b32_e64 %2, %2, %6, %0\n\t"
+        "v_cndmask_b32_e64 %3, %3, %6, %1\n\t"
+        : "=s"(x_mask), "=s"(y_mask), "+v"(x_ref), "+v"(y_ref)
+        : "v"(rel_vgpr), "n"(THR_X), "v"(neg_inf_vgpr), "n"(THR_Y), "v"(rel_hi_vgpr)
+        : "vcc"
+    );
+#else
    asm volatile(
        // x: rel < THR_X ?
        "v_cmp_lt_i32_e64 %0, %6, %7\n\t"
@@ -99,7 +118,7 @@ __device__ inline void mask_vec2_imm(uint32_t rel_vgpr, uint32_t neg_inf_vgpr,
          "n"(THR_X), "v"(neg_inf_vgpr), "n"(THR_Y)
        : "vcc"
    );
-    // x_ref = ox; y_ref = oy;
+#endif
 }

 template<ducks::rt::col_layout RT>
@@ -122,6 +141,8 @@ __device__ inline void mask_kv_tile(RT &dst, int q_abs, int k_abs, uint32_t neg_
        // (smaller rel ⇒ more "future" keys that must be -inf)
        const int rel0 = q_pos - (k_base + row_base);
        const uint32_t rel = static_cast<uint32_t>(rel0);
+        // rel-WINDOW keeps THR within the inline-constant range
+        const uint32_t rel_hi = static_cast<uint32_t>(rel0 - WINDOW);

        #pragma unroll
        for (int j = 0; j < dst.width; ++j) {
@@ -145,14 +166,14 @@ __device__ inline void mask_kv_tile(RT &dst, int q_abs, int k_abs, uint32_t neg_
            //  - reuse a single neg_inf register
            //  - keep VCC live across the pair
            //  - avoid reloading -inf or recomputing rel
-            mask_vec2_imm< 0, 1 >(rel, neg_inf_v, d0x, d0y);
-            mask_vec2_imm< 2, 3 >(rel, neg_inf_v, d1x, d1y);
-            mask_vec2_imm< 8, 9 >(rel, neg_inf_v, d2x, d2y);
-            mask_vec2_imm<10,11 >(rel, neg_inf_v, d3x, d3y);
-            mask_vec2_imm<16,17 >(rel, neg_inf_v, d4x, d4y);
-            mask_vec2_imm<18,19 >(rel, neg_inf_v, d5x, d5y);
-            mask_vec2_imm<24,25 >(rel, neg_inf_v, d6x, d6y);
-            mask_vec2_imm<26,27 >(rel, neg_inf_v, d7x, d7y);
+            mask_vec2_imm< 0, 1 >(rel, rel_hi, neg_inf_v, d0x, d0y);
+            mask_vec2_imm< 2, 3 >(rel, rel_hi, neg_inf_v, d1x, d1y);
+            mask_vec2_imm< 8, 9 >(rel, rel_hi, neg_inf_v, d2x, d2y);
+            mask_vec2_imm<10,11 >(rel, rel_hi, neg_inf_v, d3x, d3y);
+            mask_vec2_imm<16,17 >(rel, rel_hi, neg_inf_v, d4x, d4y);
+            mask_vec2_imm<18,19 >(rel, rel_hi, neg_inf_v, d5x, d5y);
+            mask_vec2_imm<24,25 >(rel, rel_hi, neg_inf_v, d6x, d6y);
+            mask_vec2_imm<26,27 >(rel, rel_hi, neg_inf_v, d7x, d7y);
        }
    }
 }
@@ -201,6 +222,16 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    else max_num_tiles = num_tiles;
    const int q_start_pos = tile_idx * Q_BLOCK_SIZE;

+#if WINDOW
+    // start at the first in-window tile; clamp keeps >=4 tiles for the pipeline unroll
+    const int block_min_q = block_tile_idx * NUM_WARPS * Q_BLOCK_SIZE;
+    int min_tile = (block_min_q - WINDOW + 1) / KV_BLOCK_SIZE;
+    if (min_tile < 0) min_tile = 0;
+    if (min_tile > max_num_tiles - 4) min_tile = max(0, max_num_tiles - 4);
+#else
+    constexpr int min_tile = 0;
+#endif
+
    constexpr float TEMPERATURE_SCALE = (D == 128) ? 0.08838834764f*1.44269504089f : 0.125f*1.44269504089f;
    uint32_t neg_inf_v = 0xff800000;

@@ -231,7 +262,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    G::prefill_swizzled_offsets<1, false>(k_smem[0], g.Kg, swizzled_offsets_K);
    G::prefill_swizzled_offsets<1, false>(v_smem[0], g.Vg, swizzled_offsets_V);

-    G::load<1, false>(k_smem[0], g.Kg, {batch_idx, 0, head_idx_kv, 0}, swizzled_offsets_K);
+    G::load<1, false>(k_smem[0], g.Kg, {batch_idx, min_tile, head_idx_kv, 0}, swizzled_offsets_K);
    __builtin_amdgcn_s_waitcnt(0);
    __builtin_amdgcn_sched_barrier(0);
    __builtin_amdgcn_s_barrier();
@@ -243,9 +274,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    transpose(q_reg_transposed, q_reg);

    // All warps then collaboratively load in the first slice of V (V0) and the second slice of K (K1) into shared memory
-    G::load<1, false>(k_smem[1], g.Kg, {batch_idx, 1, head_idx_kv, 0}, swizzled_offsets_K);
+    G::load<1, false>(k_smem[1], g.Kg, {batch_idx, min_tile + 1, head_idx_kv, 0}, swizzled_offsets_K);
    // All warps then load in the first slice of K (K0)
-    G::load<1, false>(v_smem[0], g.Vg, {batch_idx, 0, head_idx_kv, 0}, swizzled_offsets_V);
+    G::load<1, false>(v_smem[0], g.Vg, {batch_idx, min_tile, head_idx_kv, 0}, swizzled_offsets_V);
    load(k_reg, k_smem[0]);
    __builtin_amdgcn_sched_barrier(0);
    asm volatile("s_waitcnt lgkmcnt(0)");
@@ -259,13 +290,20 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
    __builtin_amdgcn_sched_barrier(0);
    if constexpr (causal) {
-        const int kv_end_pos = (1) * KV_BLOCK_SIZE;
-        if (__builtin_expect(q_start_pos < kv_end_pos, 0)) {  // Only mask if needed
-            mask_kv_tile(att_block[0], tile_idx, 0, neg_inf_v, lane);
+        const int kv_end_pos = (min_tile + 1) * KV_BLOCK_SIZE;
+        if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, WINDOW ? 1 : 0)) {
+            mask_kv_tile(att_block[0], tile_idx, min_tile, neg_inf_v, lane);
        }
    }
    // Each warp performs a partial softmax of QK0 (i.e. some of the online softmax up until but not including the second exponential scaling of the attention block likely)
+#if WINDOW
+    // floor the max: min_tile can be fully masked, which would NaN via exp2(-inf - -inf)
+    zero(max_vec_prev);
+    add(max_vec_prev, max_vec_prev, -1e4f);
+    col_max(max_vec, att_block[0], max_vec_prev);
+#else
    col_max(max_vec, att_block[0]);
+#endif

    copy(max_vec_prev, max_vec);
    exp2(scale_vec, scale_vec);
@@ -284,21 +322,25 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    // All warps then load in the second slice of K (K1)
    load(k_reg, k_smem[1]);
    // All warps then collaboratively load in the third slice of K (K2) into shared memory
-    G::load<1, false>(k_smem[0], g.Kg, {batch_idx, 2, head_idx_kv, 0}, swizzled_offsets_K);
+    G::load<1, false>(k_smem[0], g.Kg, {batch_idx, min_tile + 2, head_idx_kv, 0}, swizzled_offsets_K);
    // All warps then collaboratively load in the second slice of V (V1) into shared memory
-    G::load<1, false>(v_smem[1], g.Vg, {batch_idx, 1, head_idx_kv, 0}, swizzled_offsets_V);
+    G::load<1, false>(v_smem[1], g.Vg, {batch_idx, min_tile + 1, head_idx_kv, 0}, swizzled_offsets_V);
    asm volatile("s_waitcnt lgkmcnt(0)");
    asm volatile("s_waitcnt vmcnt(" FA_VM4 ")");
    __builtin_amdgcn_sched_barrier(0);
    __builtin_amdgcn_s_barrier();

    // hot loop
-    for (int j = 3; j < max_num_tiles - 1; j += 2) {
+    for (int j = min_tile + 3; j < max_num_tiles - 1; j += 2) {
        // Cluster 0:
        //      QK1
        zero(att_block[1]);
        transpose(k_reg_transposed, k_reg);
        mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]);
+#if WINDOW
+        // window masks interior tiles that causal skips
+        mask_kv_tile(att_block[1], tile_idx, j - 2, neg_inf_v, lane);
+#endif
        //      Finish softmax for QK0
        exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]);
        mul(norm_vec, norm_vec, scale_vec);
@@ -379,7 +421,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
        load(v_reg, v_smem[1]);
        if constexpr (causal) {
            const int kv_end_pos = (j) * KV_BLOCK_SIZE;
-            if (q_start_pos < kv_end_pos) {  // Only mask if needed
+            if (WINDOW || q_start_pos < kv_end_pos) {
                mask_kv_tile(att_block[0], tile_idx, j - 1, neg_inf_v, lane);
            }
        }
@@ -447,7 +489,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    load(v_reg, v_smem[0]);
    if constexpr (causal) {
        const int kv_end_pos = (max_num_tiles - 2) * KV_BLOCK_SIZE;
-        if (__builtin_expect(q_start_pos < kv_end_pos, 0)) {  // Only mask if needed
+        if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, WINDOW ? 1 : 0)) {
            mask_kv_tile(att_block[1], tile_idx, max_num_tiles - 3, neg_inf_v, lane);
        }
    }
@@ -510,7 +552,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    load(v_reg, v_smem[1]);
    if constexpr (causal) {
        const int kv_end_pos = (max_num_tiles - 1) * KV_BLOCK_SIZE;
-        if (__builtin_expect(q_start_pos < kv_end_pos, 1)) {  // Only mask if needed
+        if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, 1)) {
            mask_kv_tile(att_block[0], tile_idx, max_num_tiles - 2, neg_inf_v, lane);
        }
    }
@@ -572,7 +614,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    load(v_reg, v_smem[0]);
    if constexpr (causal) {
        const int kv_end_pos = (max_num_tiles) * KV_BLOCK_SIZE;
-        if (__builtin_expect(q_start_pos < kv_end_pos, 1)) {  // Only mask if needed
+        if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, 1)) {
            mask_kv_tile(att_block[1], tile_idx, max_num_tiles - 1, neg_inf_v, lane);
        }
    }
@@ -97,4 +97,33 @@ __device__ inline static void atomic_pk_add_bf16_with_warpid(const GL &dst, cons
            }(std::make_index_sequence<RT::width>{});
        }.template operator()<Ns>(), ...);
    }(std::make_index_sequence<RT::height>{});
-}
+}
+// make_window: complement of make_causal for the window lower boundary (q_pos-k_pos == WINDOW). masks = ~(causal masks)
+template<int N, int M, int GPR, ducks::art::all T0, ducks::art::all T1>
+__device__ static inline void make_window(T0 &dst, const T1 &src) {
+    static_assert(std::is_same_v<typename T0::T, float> && std::is_same_v<typename T1::T, float>, "Only float to float window mask is supported");
+    static_assert(std::is_same_v<typename T0::layout, typename T1::layout>, "Only same layout is supported");
+    static_assert(std::is_same_v<typename T0::shape, typename T1::shape>, "Only same shape is supported");
+
+    if constexpr (std::is_same_v<typename T0::layout, typename ducks::rt_layout::col> && std::is_same_v<typename T0::shape, typename ducks::rt_shape::rt_16x16>) {
+        using range_type_T0 = ducks::art::get_nth_range_t<typename T0::register_ranges, N * T0::width + M>;
+        using registers_T0 = ducks::art::split_many_t<ducks::art::type_list<range_type_T0>, 1>;
+        using range_type_T1 = ducks::art::get_nth_range_t<typename T1::register_ranges, N * T1::width + M>;
+        using registers_T1 = ducks::art::split_many_t<ducks::art::type_list<range_type_T1>, 1>;
+        static_assert(registers_T0::size == registers_T1::size);
+
+        uint64_t window_mask = 0x1FFF01FF001F0001;
+        macros::v_cndmask_b32_e64<ducks::art::get_nth_range_t<registers_T0, 0>::lo, ducks::art::get_nth_range_t<registers_T1, 0>::lo, GPR>(window_mask);
+
+        window_mask = 0x3FFF03FF003F0003;
+        macros::v_cndmask_b32_e64<ducks::art::get_nth_range_t<registers_T0, 1>::lo, ducks::art::get_nth_range_t<registers_T1, 1>::lo, GPR>(window_mask);
+
+        window_mask = 0x7FFF07FF007F0007;
+        macros::v_cndmask_b32_e64<ducks::art::get_nth_range_t<registers_T0, 2>::lo, ducks::art::get_nth_range_t<registers_T1, 2>::lo, GPR>(window_mask);
+
+        window_mask = 0xFFFF0FFF00FF000F;
+        macros::v_cndmask_b32_e64<ducks::art::get_nth_range_t<registers_T0, 3>::lo, ducks::art::get_nth_range_t<registers_T1, 3>::lo, GPR>(window_mask);
+    } else {
+        static_assert(false, "Unsupported window mask");
+    }
+}
@@ -471,6 +471,20 @@ class TestCmpFloat(unittest.TestCase):
    st = run_program(instructions, n_lanes=1)
    self.assertEqual(st.vcc & 1, 1, "Expected vcc=1 (1.0 != 2.0)")

+  def test_v_cmp_eq_f16_src0_hi(self):
+    """v_cmp_eq_f16 with src0 from high half (true16 384+n encoding)."""
+    cmp = v_cmp_eq_f16_e32(v[0], v[1])
+    cmp._raw += 128  # src0 v[0] -> v[0].h, the dsl can't encode hi-half src0 yet
+    instructions = [
+      s_mov_b32(s[0], 0x42003c00),  # hi=3.0, lo=1.0
+      v_mov_b32_e32(v[0], s[0]),
+      s_mov_b32(s[0], 0x47004200),  # hi=7.0, lo=3.0
+      v_mov_b32_e32(v[1], s[0]),
+      cmp,
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vcc & 1, 1, "Expected vcc=1 (v0.hi 3.0 == v1.lo 3.0)")
+
  def test_v_cmp_nge_f16_inf_self(self):
    """v_cmp_nge_f16 comparing -inf with itself (unordered less than).

@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.helpers import getenv, system, DEV
-from extra.gemm.cdna_asm_gemm import asm_gemm, hk_bf16_atb_gemm, quantize_mxfp4
+from extra.gemm.cdna_asm_gemm import asm_gemm, hk_bf16_atb_gemm
 from test.helpers import needs_second_gpu
 from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8, FP8_MAX

@@ -157,13 +157,20 @@ class TestMXFP4(unittest.TestCase):

  def test_quantize(self):
    import numpy as np
-    block = np.array([0, .26, .74, .75, 1.26, 1.75, 2.51, 3.5, 5.1, 6, -6] + [0] * 21, dtype=np.float32)
-    x = Tensor(np.tile(block, (32, 8)), dtype=dtypes.bfloat16)
-    packed, scale, _ = quantize_mxfp4(x)
-    p = packed.numpy()
-    codes = np.stack((p & 0xF, p >> 4), axis=-1).reshape(32, 256)
-    np.testing.assert_array_equal(codes[0, :11], [0, 1, 1, 2, 3, 4, 5, 6, 7, 7, 15])
-    np.testing.assert_array_equal(scale.numpy(), np.full((32, 8), 127, dtype=np.uint8))
+    from extra.llama_kernels.quantize_mxfp4 import quantize_mxfp4
+    rng = np.random.default_rng(0)
+    x = np.triu(rng.standard_normal((256, 256), dtype=np.float32))
+    x += np.triu(x, 1).T
+    x[:32, :32] = 0
+    row, row_scale, col, col_scale = quantize_mxfp4(Tensor(x, dtype=dtypes.bfloat16))
+    Tensor.realize(row, row_scale, col, col_scale)
+    row, row_scale = row.numpy(), row_scale.numpy()
+    col, col_scale = col.numpy(), col_scale.numpy()
+    np.testing.assert_array_equal(row, col)
+    np.testing.assert_array_equal(row_scale, col_scale)
+    self.assertTrue(row.any())
+    self.assertTrue((row_scale == 127).any())
+    self.assertTrue((row_scale != 127).any())

  def test_correctness(self):
    import numpy as np
@@ -171,17 +178,9 @@ class TestMXFP4(unittest.TestCase):
    rng = np.random.default_rng(1)
    a = Tensor(rng.standard_normal((M, K), dtype=np.float32), dtype=dtypes.bfloat16)
    b = Tensor(rng.standard_normal((N, K), dtype=np.float32), dtype=dtypes.bfloat16)
-    out = asm_gemm(a, b.T, mxfp4=True).realize()
-    # reference gemm
-    a_packed, scale_a, _ = quantize_mxfp4(a)
-    b_packed, scale_b, _ = quantize_mxfp4(b)
-    def unpack(x): return np.stack((x & 0xF, x >> 4), axis=-1).reshape(x.shape[0], -1)
-    code_a, code_b = unpack(a_packed.numpy()), unpack(b_packed.numpy())
-    lut = np.array([0, .5, 1, 1.5, 2, 3, 4, 6, -0., -.5, -1, -1.5, -2, -3, -4, -6], dtype=np.float32)
-    a_dequant = lut[code_a] * np.repeat(np.exp2(scale_a.numpy().astype(np.int16)-127), 32, axis=1)
-    b_dequant = lut[code_b] * np.repeat(np.exp2(scale_b.numpy().astype(np.int16)-127), 32, axis=1)
-    ref = Tensor(a_dequant @ b_dequant.T, dtype=dtypes.bfloat16).realize().numpy()
-    np.testing.assert_array_equal(out.numpy(), ref)
+    out = asm_gemm(a, b.T, mxfp4=True).realize().numpy().astype(np.float32)
+    ref = a.numpy().astype(np.float32) @ b.numpy().astype(np.float32).T
+    self.assertLess(np.linalg.norm(out-ref) / np.linalg.norm(ref), 0.2)

  def test_empty(self):
    M, N, K = getenv("M", 16384), getenv("N", 4096), getenv("K", 14336)
@@ -190,6 +190,12 @@ class TestCustomKernel(unittest.TestCase):
    b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
    self.assertEqual(b.item(), 15)

+  def test_sum_outside(self):
+    a = Tensor([1.0, 2, 3, 4, 5])+1
+    tst = Tensor.empty(1)
+    b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0]
+    self.assertEqual(b.item(), 20)
+
  def test_sum_int(self):
    a = Tensor([1, 2, 3, 4, 5])
    tst = Tensor.empty(1, dtype=a.dtype)
@@ -287,7 +293,7 @@ class TestCustomKernel(unittest.TestCase):
    GlobalCounters.reset()
    c.realize()
    assert all(i == 3. for i in c.flatten().tolist()), f"all 3 {c.tolist()}"
-    assert_kernel_count(3)
+    assert_kernel_count(2)

  def test_multi_after_schedule_order(self):
    """Test correct scheduling order when custom_kernel has multiple outputs.
@@ -405,10 +411,8 @@ class TestCustomKernel(unittest.TestCase):
    assert_kernel_count(2)
    self.assertEqual(z.tolist(), x.add(2).tolist())

-  @unittest.expectedFailure
  def test_custom_kernel_sched_copy(self): self.test_custom_kernel_sched(use_custom=True)

-  @unittest.expectedFailure
  def test_sliced_buffer_function(self):
    x = Tensor.arange(32).reshape(8, 4).clone().realize()
    from tinygrad import function
@@ -418,8 +422,8 @@ class TestCustomKernel(unittest.TestCase):
      return Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
    GlobalCounters.reset()
    y = run(x[0]).realize()
-    # it's copying the input and the output
-    assert_kernel_count(1)
+    # backends that support contiguous views don't launch extra kernels
+    assert_kernel_count(2 if x[0].uop.contiguous_view() is None else 1)
    self.assertEqual(y.tolist(), [1, 2, 3, 4])

  @Context(DEV="CPU")
@@ -429,12 +433,28 @@ class TestCustomKernel(unittest.TestCase):
    # TODO: it currently requires a compiler for Ops.BINARY
    from tinygrad.device import Device
    binary = Device[a.device].renderer.compiler.compile(src)
-    def custom_src_kernel(A:UOp) -> UOp:
+    def custom_src_kernel(A:UOp, B:UOp) -> UOp:
      sink = UOp.sink(A, arg=KernelInfo(name="test_src"))
      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
-    a = Tensor.custom_kernel(a.reshape(2, 2).T, fxn=custom_src_kernel)[0]
-    self.assertEqual(a.tolist(), [[1, 2], [1, 3]])
+    a = Tensor.custom_kernel(a.reshape(2, 2).clone(), a.reshape(2, 2).T, fxn=custom_src_kernel)[0]
+    self.assertEqual(a.tolist(), [[1, 1], [2, 3]])

+  @Context(DEV="CPU")
+  def test_simple_from_source_alt(self):
+    a = Tensor.arange(4).clone().realize()
+    src = "void copy(int* restrict out, int* restrict in) { for (int i = 0; i < 4; i++) out[i] = in[i]; }"
+    from tinygrad.device import Device
+    binary = Device[a.device].renderer.compiler.compile(src)
+    def custom_src_kernel(out:UOp, inp:UOp) -> UOp:
+      sink = UOp.sink(out, inp, arg=KernelInfo(name="copy"))
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
+    out = Tensor.custom_kernel(Tensor.empty_like(a), a+1, fxn=custom_src_kernel)[0]
+    GlobalCounters.reset()
+    out.realize()
+    assert_kernel_count(2)
+    self.assertEqual(out.tolist(), [1, 2, 3, 4])
+
+  @unittest.skip("this shouldn't be expected to work")
  def test_inplace_transpose(self):
    def custom_assign_row_max_kernel(A:UOp) -> UOp:
      row = UOp.range(A.shape[0], 0)
@@ -471,8 +491,8 @@ class TestCustomKernelInput(unittest.TestCase):

  def test_reshape(self): self._test_mop(lambda x: x.reshape(16, 2), max_kernels=2)
  def test_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T, max_kernels=3)
-  def test_double_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T.T, max_kernels=3)
-  def test_shrink(self): self._test_mop(lambda x: x[:4], max_kernels=2)
+  def test_double_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T.T, max_kernels=2)
+  def test_shrink(self): self._test_mop(lambda x: x[:4], max_kernels=1)
  def test_pad(self): self._test_mop(lambda x: x[:4].pad(((0, 4),)), max_kernels=2)
  def test_flip(self): self._test_mop(lambda x: x.flip(0), max_kernels=2)
  def test_offset_shrink(self): self._test_mop(lambda x: x[4:8], max_kernels=2)
@@ -169,6 +169,13 @@ 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())

+  def test_fp8e5m2fnuz_to_float_smallest_normals(self):
+    # fnuz bias exceeds half's, so exp-1 normals land below half's normal range: they flush to zero like denormals
+    if dtypes.half not in supported_dtypes or dtypes.half in EMULATED_DTYPES.tolist(dtypes) or dtypes.fp8e5m2fnuz in supported_dtypes:
+      self.skipTest("needs the emulated fp8 with a native half intermediate")
+    vals = Tensor([0x04, 0x05, 0x06, 0x07], dtype=dtypes.uint8).bitcast(dtypes.fp8e5m2fnuz).float().numpy()
+    np.testing.assert_equal(vals, [0., 0., 0., 0.])
+
 class TestBFloat16DType(unittest.TestCase):
  def test_bf16_to_float(self):
    _test_cast(Tensor([100000], dtype=dtypes.bfloat16), dtypes.float32)
@@ -399,9 +399,10 @@ class TestDTypeALU(unittest.TestCase):
    if float_dtype not in supported_dtypes: float_dtype = dtypes.float32
    universal_test_cast(a, float_dtype, unsigned_dtype)

-  @unittest.expectedFailure
-  def test_unsafe_cast_float_to_int_failure(self):
-    val = float(dtypes.int32.max - 1)
+  def test_unsafe_cast_float_to_int(self):
+    # the value is off the float32 grid but rounds in-range: the buffer and const-fold paths must agree
+    # (out-of-range float->int cast stays undefined: hardware may saturate where the fold wraps)
+    val = 2147483000.0
    t1 = Tensor([val], dtype=dtypes.float32).cast(dtypes.int32)
    t2 = Tensor(val, dtype=dtypes.float32).cast(dtypes.int32)
    np.testing.assert_equal(t1.item(), t2.item())
@@ -360,7 +360,7 @@ class TestJitGraphSplit(unittest.TestCase):
    self.expect(f, inp, inp_cpu,
      graph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()],
      multigraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()],
-      hcqgraph=[self.ji_graph(4)])
+      hcqgraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()]) # cpu is hcq2 now, it does not join hcq graphs

  def test_jit_cpu_several(self):
    if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
@@ -377,9 +377,9 @@ class TestJitGraphSplit(unittest.TestCase):
    inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
    inp_cpu = Tensor.randn(10, 10, device="CPU").realize()
    self.expect(f, inp, inp_cpu,
-      graph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
-      multigraph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
-      hcqgraph=[self.ji_graph(5)])
+      graph=[self.ji_graph(2), self.ji_comp(), self.ji_comp(), self.ji_comp()],
+      multigraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp(), self.ji_comp()],
+      hcqgraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp(), self.ji_comp()])

  def test_jit_multidev(self):
    if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
@@ -30,7 +30,7 @@ class TestLinearizer(unittest.TestCase):
    c = ((a.shrink(((0, 2),)) - a.shrink(((2, 4),))) - (b.shrink(((0, 2),)) - b.shrink(((2, 4),))))
    linear = c.schedule_linear()
    run_linear(linear)
-    rawbufs = [s.buffer for s in linear.src[-1].src[1:] if s.op is not Ops.BIND]
+    rawbufs = [s.buffer for s in linear.src[-1].src[1:] if not s.is_bound_var]
    assert len(rawbufs) == 3 and set(rawbufs[1:]) == {a.uop.base.realized, b.uop.base.realized}
    np_c = (np_a[:2] - np_a[2:]) - (np_b[:2] - np_b[2:])
    np.testing.assert_allclose(np_c, c.numpy(), atol=1e-4, rtol=1e-4)
@@ -411,7 +411,7 @@ def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
  last_call = linear.src[-1]
  ast = last_call.src[0]
  assert ast.op is Ops.SINK, f"helper_realized_ast expects a SINK {last_call}"
-  last_bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]
+  last_bufs = [s.buffer for s in last_call.src[1:] if not s.is_bound_var]
  # now all input buffers in last_call should be realized
  # create fresh buffers for the outputs
  bufs = [Buffer(x.device, x.size, x.dtype).allocate() if i < len(ast.src) else x for i,x in enumerate(last_bufs)]
@@ -437,7 +437,7 @@ def reset_bufs(bufs:list[Buffer]):
  for buf in bufs: buf.copy_from(Buffer("PYTHON", buf.size, buf.dtype, opaque=memoryview(bytearray(buf.nbytes))))

 def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[],
-                               apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[]):
+                               apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[], check_default_opt=True):
  outbufs = real_bufs[:len(realized_ast.src)]
  wanna_output = [np.array(x).flatten() for x in wanna_output]
  buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in real_bufs]
@@ -459,9 +459,7 @@ def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[]
    for buf,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(buf, want, atol=atol, rtol=rtol)

  # Check correctness of handcoded optimiztions.
-  reset_bufs(outbufs)
-  run_prg(opts=None)
-  for buf,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(buf, want, atol=atol, rtol=rtol)
+  if check_default_opt: check_opt(None)
  for x in opts: # Check custom transformations if any.
    check_opt(([Opt(OptOps.TC, 0, (TC_SELECT.value, TC_OPT.value, 1))] if apply_tc else [])+x)

@@ -5,6 +5,7 @@ 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 import local_abs_max
 from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed, quantize_fp8_scalar
+from extra.llama_kernels.swiglu import swiglu
 from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
 from extra.thunder.amd.fa import custom_fused_qkv_rope_backward, fused_qkv_rope
 from test.helpers import needs_second_gpu, assert_kernel_count
@@ -161,5 +162,31 @@ class TestFusedQKVRoPE(unittest.TestCase):
    ref = Tensor.cat(dq_ref, dk_ref, dv_ref, dim=3).reshape(*dx.shape).realize()
    with Context(DEBUG=0): self.assertTrue(dx.allclose(ref, atol=2e-2, rtol=2e-2).item(), "backward mismatch")

+def run_swiglu(test:unittest.TestCase, shape:tuple[int, ...]) -> None:
+  Tensor.manual_seed(0)
+  x = (Tensor.randn(*shape) * 2).cast(dtypes.bfloat16).realize()
+  hidden = x.shape[-1] // 2
+  out, ref = swiglu(x), x[..., :hidden].silu() * x[..., hidden:]
+  Tensor.realize(out, ref)
+  with Context(DEBUG=0): test.assertTrue(out.allclose(ref, atol=2.5e-1, rtol=3e-2).item(), "SwiGLU forward mismatch")
+
+  grad = (Tensor.randn(*out.shape) * 2).cast(dtypes.bfloat16).realize()
+  grad_x, grad_ref = out.gradient(x, gradient=grad)[0], ref.gradient(x, gradient=grad)[0]
+  Tensor.realize(grad_x, grad_ref)
+  test.assertEqual(grad_x.shape, shape)
+  test.assertEqual(grad_x.dtype, dtypes.bfloat16)
+  with Context(DEBUG=0): test.assertTrue(grad_x.allclose(grad_ref, atol=2.5e-1, rtol=3e-2).item(), "SwiGLU backward mismatch")
+
+class TestSwiGLU(unittest.TestCase):
+  def setUp(self):
+    if dtypes.bfloat16 not in Device[Device.DEFAULT].renderer.supported_dtypes(): self.skipTest("need bfloat16")
+
+  def test_simple(self): run_swiglu(self, (2, 32, 64))
+
+  def test_llama_shape(self):
+    if Device.DEFAULT != "AMD" or not Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950"):
+      self.skipTest("only run on real machine for speed")
+    run_swiglu(self, (2, 8192, 28672))
+
 if __name__ == '__main__':
  unittest.main()
@@ -1,9 +1,9 @@
 import unittest, random
 from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variable
-from tinygrad.uop.ops import Ops, UOp, AxisType
+from tinygrad.uop.ops import Ops, UOp, AxisType, graph_rewrite
 from tinygrad.helpers import getenv, prod, Context
 from tinygrad.nn.state import get_parameters
-from tinygrad.engine.realize import run_linear, compile_linear
+from tinygrad.engine.realize import run_linear, compile_linear, pm_beam, pm_compile
 import numpy as np
 from hypothesis import given, strategies as strat, settings
 from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule, assert_kernel_count
@@ -79,9 +79,9 @@ class TestMultiTensor(unittest.TestCase):
  def test_shard_beam(self):
    cpu_2 = ("CPU:1", "CPU:2")
    src = Tensor.ones(16).shard(cpu_2, 0).realize()
-    pad = src.to(cpu_2[::-1]).schedule_linear().src[0]
-    with Context(BEAM=1, IGNORE_BEAM_CACHE=1): prg = compile_linear(UOp(Ops.LINEAR, src=(pad,))).src[0].src[0]
-    self.assertNotEqual(prg.src[0].arg.applied_opts, ())
+    lin = UOp(Ops.LINEAR, src=(src.to(cpu_2[::-1]).schedule_linear().src[0],))
+    with Context(BEAM=1, IGNORE_BEAM_CACHE=1): call = graph_rewrite(graph_rewrite(lin, pm_beam, ctx=1, walk=True), pm_compile, walk=True).src[0]
+    self.assertNotEqual(call.src[0].src[0].arg.applied_opts, ())

  def test_shard_same_device(self):
    X = Tensor.ones(256).contiguous().realize()
@@ -1535,6 +1535,8 @@ class TestOps(unittest.TestCase):

  def test_prod(self):
    helper_test_op(None, lambda x: x.prod(), vals=[[1.0, 2.0, 3.0]])
+    helper_test_op(None, lambda x: x.prod(), vals=[[0.0, 2.0, 3.0]])
+    helper_test_op(None, lambda x: x.prod(), vals=[[0.0, 0.0, 3.0]])
    with Context(NOOPT=1): helper_test_op(None, lambda x: x.prod(), vals=[[1.0, 2.0, 3.0]])
    helper_test_op([(3,4,5,6)], lambda x: x.prod(dim=3), lambda x: x.prod(axis=3))
    helper_test_op([(3,4,5,6)], lambda x: x.prod(dim=1), lambda x: x.prod(axis=1))
@@ -2,7 +2,7 @@ import unittest, pickle, types, tracemalloc
 import numpy as np
 from tinygrad import Tensor, Device, TinyJit, Variable, dtypes
 from tinygrad.helpers import GlobalCounters, ContextVar, Context, DEV
-from tinygrad.uop.ops import PatternMatcher, UPat, UOp
+from tinygrad.uop.ops import PatternMatcher, UPat, UOp, deconstruct_function

 class TestPickle(unittest.TestCase):
  def test_pickle_code_object(self):
@@ -11,6 +11,11 @@ class TestPickle(unittest.TestCase):
    fxn = types.FunctionType(pickle.loads(code_str), globals())
    self.assertEqual(fxn(2), 4)

+  def test_deconstruct_function_nested_comprehension(self):
+    # pre PEP 709, each comprehension is its own code object, so dtypes here is referenced two code objects deep
+    def fxn(): return [[dtypes.int for _ in range(2)] for _ in range(2)]
+    self.assertEqual(types.FunctionType(*deconstruct_function(fxn))(), fxn())
+
  def test_pickle_pattern_matcher(self):
    pm = PatternMatcher([(UPat.cvar('x'), lambda x: x*2)])
    sink = UOp.const(2)
@@ -6,6 +6,15 @@ from examples.gpt2 import Attention
 import numpy as np

 class TestSymbolicOps(unittest.TestCase):
+  def test_negative_slice(self):
+    a = Tensor.rand(3, 10, 4)
+    for i in range(3, 10):
+      vi = Variable("i", 1, 10).bind(i)
+      # negative int bounds against a symbolic dim must resolve against the size, like slice.indices
+      np.testing.assert_allclose(a[:, :vi][:, -3:-1].numpy(), a[:, :i][:, -3:-1].numpy(), atol=1e-6, rtol=1e-6)
+      np.testing.assert_allclose(a[:, :vi][:, -1:].numpy(), a[:, :i][:, -1:].numpy(), atol=1e-6, rtol=1e-6)
+      np.testing.assert_allclose(a[:, :vi][:, -1].numpy(), a[:, :i][:, -1].numpy(), atol=1e-6, rtol=1e-6)
+
  def test_plus1(self):
    def f(a): return (a+1).realize()
    a = Tensor.rand(3, 10)
@@ -0,0 +1,14 @@
+import unittest, numpy as np
+from unittest.mock import patch
+from tinygrad import Device, Tensor
+from tinygrad.helpers import getenv
+from tinygrad.runtime.support.hcq2 import HCQ_DEVS, all_devices_in
+
+@unittest.skipUnless(getenv("HCQ2") and all_devices_in(Device.DEFAULT, HCQ_DEVS), "hcq2 device required")
+class TestHCQ2(unittest.TestCase):
+  def test_copy_without_copy_queue(self):
+    with patch.object(Device[Device.DEFAULT], "has_copy_queue", False):
+      np.testing.assert_equal(Tensor(np.arange(61, dtype=np.float32)).to(Device.DEFAULT).contiguous().realize().numpy(), np.arange(61))
+
+if __name__ == "__main__":
+  unittest.main()
@@ -0,0 +1,17 @@
+from tinygrad import Device, Tensor, TinyJit, dtypes
+from tinygrad.helpers import Timing, Context
+
+GPUS, DEPTH, SZ = 8, 4, 128 * 2**20
+WARMUP, ITERS = 3, 5
+devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(GPUS))
+bufs = tuple(Tensor.empty(SZ, dtype=dtypes.uint8, device=dev).contiguous().realize() for _ in range(DEPTH) for dev in devs)
+
+@TinyJit
+def all_to_all(*srcs:Tensor): return Tensor.realize(*(src.to(dst) for i,src in enumerate(srcs) for j,dst in enumerate(devs) if i % GPUS != j))
+
+if __name__ == "__main__":
+  with Context(ALL2ALL=1, JIT_BATCH_SIZE=0):
+    for i in range(-WARMUP, ITERS):
+      with Timing("ALL2ALL ", lambda ns: f" {SZ*GPUS*(GPUS-1)*DEPTH/ns:.2f} GB/s", enabled=i>=0):
+        all_to_all(*bufs)
+        for dev in devs: Device[dev].synchronize()
@@ -1,5 +1,5 @@
-import unittest, time
-from tinygrad import Tensor
+import unittest, time, itertools
+from tinygrad import Tensor, Context

 class TestScheduleScaling(unittest.TestCase):
  """Test that .schedule() scales linearly with graph size (no O(n^2) behavior)."""
@@ -130,5 +130,18 @@ class TestScheduleScaling(unittest.TestCase):
      return parts[0].cat(*parts[1:])
    self._assert_linear(concat_chain)

+  @Context(DEV="NULL:HIP:gfx1100")
+  def test_custom_kernel_assign_scaling(self):
+    from tinygrad.uop.ops import UOp, Ops, KernelInfo
+    from tinygrad.runtime.autogen.amd.rdna3.ins import s_nop
+    count = itertools.count(0)
+    def custom_kernel_assign(n):
+      def custom_asm(out):
+        return UOp(Ops.PROGRAM, src=(UOp.sink(out, arg=KernelInfo(f"fxn_{next(count)}")),
+                                     UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=s_nop(i)) for i in range(n*8)))))
+      call = Tensor.custom_kernel(Tensor.empty(1), fxn=custom_asm)[0]
+      return Tensor.cat(*[Tensor.empty(1).assign(call+i) for i in range(n)])
+    self._assert_linear(custom_kernel_assign, n_small=50, n_large=500)
+
 if __name__ == '__main__':
  unittest.main(verbosity=2)
@@ -48,6 +48,8 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
  else:
    assert isinstance(t, UOp), f"can't schedule {t}"
    linear, var_vals = Tensor(t).linear_with_vars()
+  # test compiling the linear
+  compile_linear(linear)
  kernel_cnt = sum((len(call.device) if isinstance(call.device, tuple) else 1)
                   for call in linear.src if call.src[0].op is Ops.SINK or not filter_sink)
  if kernel_cnt != allowed:
@@ -57,8 +59,6 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
        print("kernel", i+1)
        print(call.src[0])
    raise KernelCountException(allowed, kernel_cnt)
-  # test compiling the linear
-  compile_linear(linear)
  return linear, var_vals

 def assert_kernel_count(expected:int):
@@ -86,7 +86,7 @@ def assert_jit_cache_len(fxn, expected_len):
  if linear is None or not linear.src:
    if expected_len != 0: raise KernelCountException(expected_len, 0)
    return
-  if expected_len and all(call_is_hcq(call) for call in linear.src): expected_len = 3 # HCQ2: merged same-queue calls + finalizer + bumps
+  if expected_len and all(call_is_hcq(call) for call in linear.src): expected_len = 4 # HCQ2: fence + reset + merged same-queue calls + finalizer
  if call_is_graph(linear.src[0]):
    if len(linear.src) != 1: raise KernelCountException(1, len(linear.src))
    inner = linear.src[0].src[0].src[0]  # LINEAR UOp inside CUSTOM_FUNCTION
@@ -109,7 +109,7 @@ def _init_sqtt_encoder():
  _SMEM = (ir3.SMEM, ir4.SMEM, irc.SMEM)
  _VALU = (ir3.VOP1, ir3.VOP2, ir3.VOP3, ir3.VOP3P, ir3.VOPC, ir3.VOPD, ir3.VOP3SD, ir3.VOP3_SDST, ir3.VOP1_SDST,
           ir4.VOP1, ir4.VOP2, ir4.VOP3, ir4.VOP3P, ir4.VOPC, ir4.VOPD, ir4.VOP3SD, ir4.VOP3_SDST, ir4.VOP1_SDST,
-           irc.VOP1, irc.VOP2, irc.VOP3, irc.VOP3P, irc.VOPC, irc.VOP3SD, irc.VOP3_SDST)
+           irc.VOP1, irc.VOP2, irc.VOP3, irc.VOP3P, irc.VOP3PX2, irc.VOPC, irc.VOP3SD, irc.VOP3_SDST)
  _DS = (ir3.DS, ir4.DS, irc.DS)
  _GLOBAL = (ir3.GLOBAL, ir4.VGLOBAL, irc.GLOBAL)
  _FLAT = (ir3.FLAT, ir4.VFLAT, irc.FLAT)
@@ -260,19 +260,6 @@ def _cond(cond, if_true, if_false):
 def _cond_hi16(cond, val: UOp) -> UOp: return _cond(cond, _hi16(val), val)
 def _apply_opsel(val: UOp, sel_bit: int, opsel: int) -> UOp: return _hi16(val) if opsel & (1 << sel_bit) else val

-def _set_lane_bit(old: UOp, lane: UOp, val: UOp, exec_mask: UOp) -> UOp:
-  """Set/clear a single bit in a mask based on lane index, respecting exec mask."""
-  if old.dtype in (dtypes.uint64, dtypes.int64):
-    dt = dtypes.uint64
-    mask = UOp.const(1, dt) << lane.cast(dt)
-    new_bit = _to_u32(val).cast(dt) << lane.cast(dt)
-    cleared = old.cast(dt) & (mask ^ UOp.const(0xFFFFFFFFFFFFFFFF, dt))
-    return _lane_active(exec_mask, lane).where(cleared | new_bit, old.cast(dt))
-  mask = _c(1) << lane.cast(dtypes.uint32)
-  new_bit = _to_u32(val) << lane.cast(dtypes.uint32)
-  cleared = old & (mask ^ _c(MASK32))
-  return _lane_active(exec_mask, lane).where(cleared | new_bit, old)
-
 def _val_to_u32(val: UOp) -> UOp:
  """Convert any value to uint32 for storage (bitcast floats, cast ints)."""
  if val.dtype == dtypes.uint32: return val
@@ -532,6 +519,19 @@ class _Ctx:
      return [self.wsgpr_dyn(reg, lo), self.wsgpr_dyn(reg + _c(1), hi)]
    return [self.wsgpr_dyn(reg, val)]

+  def wmask_lane_bit(self, reg: UOp, lane: UOp, val: UOp, exec_mask: UOp) -> list[UOp]:
+    """Set/clear bit `lane` of the mask at `reg` from val for exec-active lanes, preserving memory for inactive lanes"""
+    active, bit = _lane_active(exec_mask, lane), _to_u32(val)
+    if self.wave_size <= 32:
+      old = self.rsgpr_dyn(reg)
+      mask = _c(1) << lane.cast(dtypes.uint32)
+      return [self.wsgpr_dyn(reg, active.where((old & (mask ^ _c(MASK32))) | (bit << lane.cast(dtypes.uint32)), old))]
+    off = (lane & _c(31, dtypes.int)).cast(dtypes.uint32)
+    mask = _c(1) << off
+    def half(old: UOp, sel: UOp) -> UOp: return sel.where(active.where((old & (mask ^ _c(MASK32))) | (bit << off), old), old)
+    return [self.wsgpr_dyn(reg, half(self.rsgpr_dyn(reg), lane < _c(32, dtypes.int))),
+            self.wsgpr_dyn(reg + _c(1), half(self.rsgpr_dyn(reg + _c(1)), _c(32, dtypes.int) <= lane))]
+
  def rmask(self, reg: UOp) -> UOp:
    """Read a lane mask (VCC/EXEC). Combines lo/hi for wave64."""
    if self.wave_size > 32: return _u64(self.rsgpr_dyn(reg), self.rsgpr_dyn(reg + _c(1)))
@@ -718,9 +718,7 @@ class _Ctx:
        raw_stores.append(('vgpr_direct', self.vgpr.index(val[0].valid(active)).store(new_val)))
        continue
      if 'D0' in dest and '[laneId]' in dest:
-        old_vcc = self.rmask(_c(VCC_LO.offset))
-        new_vcc = _set_lane_bit(old_vcc, lane, val, exec_mask)
-        raw_stores.extend([('vcc', s) for s in self.wmask(_c(VCC_LO.offset), new_vcc)])
+        raw_stores.extend([('vcc', s) for s in self.wmask_lane_bit(_c(VCC_LO.offset), lane, val, exec_mask)])
      elif dest.startswith('D0'):
        dest_suffix = re.match(r'D0\.(\w+)', dest)
        if dest_suffix is not None:
@@ -1039,13 +1037,11 @@ def _compile_sdwa(inst: irc.VOP1_SDWA | irc.VOP2_SDWA | irc.VOP2_SDWA_SDST | irc
        result = _sdwa_write(old, result, dst_sel, dst_unused)
      stores.append(ctx.wvgpr_dyn(vdst_reg, lane, result, exec_mask))
    elif dest.startswith('VCC'):
-      old_vcc = ctx.rmask(_c(VCC_LO.offset))
-      stores.extend(ctx.wmask(_c(VCC_LO.offset), _set_lane_bit(old_vcc, lane, val, exec_mask)))
+      stores.extend(ctx.wmask_lane_bit(_c(VCC_LO.offset), lane, val, exec_mask))
  if vcc_val is not None:
    # Initialize sdst to 0 before lane loop (old value may be unrelated data), then set lane bits in loop
    init_stores = [ctx.wsgpr_dyn(sdst_off, _c(0)), ctx.wsgpr_dyn(sdst_off + _c(1), _c(0))]
-    old_sdst = ctx.rmask(sdst_off)
-    stores.extend(ctx.wmask(sdst_off, _set_lane_bit(old_sdst, lane, vcc_val, exec_mask)))
+    stores.extend(ctx.wmask_lane_bit(sdst_off, lane, vcc_val, exec_mask))
    if stores:
      return UOp.sink(*init_stores, UOp.sink(*stores).end(lane), *ctx.inc_pc())
    return UOp.sink(*init_stores, *ctx.inc_pc())
@@ -1150,6 +1146,9 @@ def _compile_vopc(inst: ir3.VOPC|ir3.VOPC_DPP16|ir3.VOP3|ir4.VOPC|ir4.VOPC_DPP16
  def get_cmp_bit(lane) -> UOp:
    lc = lane.cast(dtypes.int) if isinstance(lane, UOp) else _c(lane, dtypes.int)
    s0 = _load_dpp16_src0(ctx, inst, lc, _c(0)) if is_dpp16 else ctx.rsrc_dyn(src0_off, lc, bits['s0'], literal, is_f64)
+    if is_vopc and not isinstance(inst, irc.VOPC) and bits['s0'] == 16 and not is_dpp16:
+      src0_hi = src0_off >= _c(384)
+      s0 = src0_hi.where(_hi16(ctx.rvgpr_dyn(src0_hi.where(src0_off - _c(384), _c(0)), lc)), s0)
    s1 = _cond_hi16(vsrc1_hi, ctx.rsrc_dyn(src1_off, lc, bits['s1'], literal, is_f64)) if bits['s0'] == 16 \
      else ctx.rsrc_dyn(src1_off, lc, bits['s1'], literal, is_f64)
    if bits['s0'] == 16 and opsel: s0, s1 = _apply_opsel(s0, 0, opsel), _apply_opsel(s1, 1, opsel)
@@ -1323,7 +1322,7 @@ def _compile_vop3sd(inst: ir3.VOP3SD | ir4.VOP3SD | irc.VOP3SD, ctx: _Ctx) -> UO
  else:
    return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask, sdst_reg=inst.sdst.offset)

-def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
+def _compile_mfma(inst: irc.VOP3P|irc.VOP3PX2, ctx: _Ctx) -> UOp:
  """CDNA MFMA matrix multiply-accumulate emulation.

  Uses local temp arrays to cache inputs, avoiding aliasing issues when vdst overlaps src0/src1.
@@ -1349,6 +1348,25 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
  src0_is_vgpr = src0_off >= _c(256)
  src1_is_vgpr = src1_off >= _c(256)

+  scaled = isinstance(inst, irc.VOP3PX2)
+  if scaled:
+    assert isinstance(inst, irc.VOP3PX2)
+    # F8F6F4 input formats: 0=FP8(E4M3), 1=BF8(E5M2). FP6/FP4 (2-4) not emulated.
+    src0_fmt, src1_fmt = int(inst.cbsz), int(inst.blgp)
+    if src0_fmt > 1 or src1_fmt > 1: raise RuntimeError(f"unsupported scaled MFMA formats cbsz={src0_fmt} blgp={src1_fmt}")
+    # scale_src0/scale_src1 are source operands pointing at 32-bit registers holding 4 packed E8M0 scale exponents.
+    # The 2-bit opsel/opsel_hi select which byte applies to A/B for this instruction.
+    scale0_off = ctx.inst_field(type(inst).scale_src0)
+    scale1_off = ctx.inst_field(type(inst).scale_src1)
+    sel0, sel1 = int(inst.opsel) & 3, int(inst.opsel_hi) & 3
+    def _scale_exp(off: UOp, sel: int, lane: UOp) -> UOp:
+      sv = ctx.rsrc_dyn(off, lane, 32)
+      byte = (sv >> UOp.const(sel * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32)
+      return byte.cast(dtypes.int32) - UOp.const(127, dtypes.int32)
+    # combined A*B scale for this lane: 2^(ea-127) * 2^(eb-127)
+    def scale_factor(lane: UOp) -> UOp:
+      return UOp.exp2((_scale_exp(scale0_off, sel0, lane) + _scale_exp(scale1_off, sel1, lane)).cast(dtypes.float32))
+
  m = _re.search(r'(\d+)X(\d+)X(\d+)', op_name)
  if m is None: raise ValueError(f"could not parse MFMA dimensions from {op_name}")
  M, N, K = int(m.group(1)), int(m.group(2)), int(m.group(3))
@@ -1404,7 +1422,18 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
  # The optimizer folds bitcast(uint32→float32) stores to float32 arrays, losing the conversion.
  tmp = UOp.placeholder((n_a_elems + n_b_elems,), dtypes.uint32, slot=0, addrspace=AddrSpace.LOCAL)

-  def cvt_elem(raw: UOp, sub_idx: int) -> UOp:
+  # Per-operand fp8 format ("fp8"=E4M3, "bf8"=E5M2) for A and B
+  if 'F8F6F4' in op_name:
+    assert isinstance(inst, (irc.VOP3P_MFMA, irc.VOP3PX2))
+    _fmts = {0: "fp8", 1: "bf8"}
+    a_fmt, b_fmt = _fmts.get(int(inst.cbsz), "fp8"), _fmts.get(int(inst.blgp), "fp8")
+  elif is_fp8:
+    # A/B formats from name suffix, e.g. V_MFMA_F32_16X16X32_BF8_FP8
+    suffixes = op_name.rsplit('_', 2)[-2:]
+    a_fmt, b_fmt = ("bf8" if sfx == "BF8" else "fp8" for sfx in suffixes)
+  else: a_fmt = b_fmt = "fp8"
+
+  def cvt_elem(raw: UOp, sub_idx: int, fp8_fmt: str = "fp8") -> UOp:
    if is_i8:
      # Extract i8, sign-extend to i32
      byte_val = (raw >> UOp.const(sub_idx * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32)
@@ -1412,7 +1441,7 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
    elif is_f32_src:
      return raw  # already uint32 (f32 bit pattern)
    elif is_fp8:
-      return ((raw >> UOp.const(sub_idx * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32)).cast(dtypes.uint32)
+      return _FUNCS[f"{fp8_fmt}_to_f32"](raw >> UOp.const(sub_idx * 8, dtypes.uint32)).bitcast(dtypes.uint32)
    elif is_bf16:
      # bf16→f32 bits: just shift left by 16 (bf16 is upper 16 bits of f32)
      return ((raw >> UOp.const(sub_idx * 16, dtypes.uint32)) & UOp.const(0xFFFF, dtypes.uint32)) << UOp.const(16, dtypes.uint32)
@@ -1454,7 +1483,7 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
    # Read A/B sources. Use rsrc_dyn for inline constants/SGPRs (src_off < 256), rvgpr_dyn for VGPRs (src_off >= 256).
    a_raw = src0_is_vgpr.where(ctx.rvgpr_dyn(src0_r + _c(reg_idx), read_lane),
                                ctx.rsrc_dyn(src0_off, _c(0, dtypes.int), 32))
-    a_val = cvt_elem(a_raw, sub_idx)
+    a_val = cvt_elem(a_raw, sub_idx, a_fmt)
    if M == 4:
      a_idx = grp_idx * UOp.const(M * K, dtypes.int) + mn_idx * UOp.const(K, dtypes.int) + UOp.const(kl, dtypes.int)
    else:
@@ -1463,7 +1492,7 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:

    b_raw = src1_is_vgpr.where(ctx.rvgpr_dyn(src1_r + _c(reg_idx), read_lane),
                                ctx.rsrc_dyn(src1_off, _c(0, dtypes.int), 32))
-    b_val = cvt_elem(b_raw, sub_idx)
+    b_val = cvt_elem(b_raw, sub_idx, b_fmt)
    if M == 4:
      b_idx = b_off + grp_idx * UOp.const(N * K, dtypes.int) + mn_idx * UOp.const(K, dtypes.int) + UOp.const(kl, dtypes.int)
    else:
@@ -1480,6 +1509,17 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
  #   Actually: 16 ACCVGPRs per lane, organized as 4 groups (l//32 gives half, each half has 2 sub-groups) of 4 rows
  tmp2 = tmp.after(read_phase)

+  def _dot_accum(acc: UOp, a_row: UOp, b_row: UOp, lane: UOp) -> UOp:
+    """acc += sum_k A[a_row+k] * B[b_row+k]. For scaled MFMA, only the dot product is scaled: D = dot*scale + C."""
+    def prod(k: int) -> UOp:
+      return tmp2.index(a_row + UOp.const(k, dtypes.int)).bitcast(acc_dt) * tmp2.index(b_row + UOp.const(k, dtypes.int)).bitcast(acc_dt)
+    if not scaled:
+      for k in range(K): acc = acc + prod(k)
+      return acc
+    dot = prod(0)
+    for k in range(1, K): dot = dot + prod(k)
+    return acc + dot * scale_factor(lane)
+
  compute_lane = ctx.range()
  compute_stores = []

@@ -1510,10 +1550,7 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
      else: acc_v = acc_v.bitcast(dtypes.float32)
      acc = src2_is_vgpr.where(acc_v, acc_scalar)

-      for k in range(K):
-        a_val = tmp2.index(m_base * UOp.const(K, dtypes.int) + UOp.const(k, dtypes.int)).bitcast(acc_dt)
-        b_val = tmp2.index(b_off + n_idx * UOp.const(K, dtypes.int) + UOp.const(k, dtypes.int)).bitcast(acc_dt)
-        acc = acc + a_val * b_val
+      acc = _dot_accum(acc, m_base * UOp.const(K, dtypes.int), b_off + n_idx * UOp.const(K, dtypes.int), compute_lane)

      if is_int_out:
        compute_stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)(
@@ -1535,17 +1572,13 @@ def _compile_mfma(inst: irc.VOP3P, ctx: _Ctx) -> UOp:
      if M == 4:
        # 4x4: each group is independent. A/B indexed per-group.
        m_base = c_grp * UOp.const(M * K, dtypes.int) + UOp.const(out_reg * K, dtypes.int)
-        for k in range(K):
-          a_val = tmp2.index(m_base + UOp.const(k, dtypes.int)).bitcast(acc_dt)
-          b_val = tmp2.index(b_off + c_grp * UOp.const(N*K, dtypes.int) + n_idx * UOp.const(K, dtypes.int)+UOp.const(k, dtypes.int)).bitcast(acc_dt)
-          acc = acc + a_val * b_val
+        b_base = b_off + c_grp * UOp.const(N * K, dtypes.int) + n_idx * UOp.const(K, dtypes.int)
      else:
        # 16x16: K is split across groups. Shared MxK/NxK arrays.
        m_base = c_grp * UOp.const(out_per_lane, dtypes.int) + UOp.const(out_reg, dtypes.int)
-        for k in range(K):
-          a_val = tmp2.index(m_base * UOp.const(K, dtypes.int) + UOp.const(k, dtypes.int)).bitcast(acc_dt)
-          b_val = tmp2.index(b_off + n_idx * UOp.const(K, dtypes.int) + UOp.const(k, dtypes.int)).bitcast(acc_dt)
-          acc = acc + a_val * b_val
+        b_base = b_off + n_idx * UOp.const(K, dtypes.int)
+
+      acc = _dot_accum(acc, m_base if M == 4 else m_base * UOp.const(K, dtypes.int), b_base, compute_lane)

      if is_int_out:
        compute_stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)(
@@ -1563,33 +1596,41 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
  vdst_reg = ctx.inst_field(type(inst).vdst)
  src0_r = ctx.inst_field(type(inst).src0) - _c(256)
  src1_r = ctx.inst_field(type(inst).src1) - _c(256)
-  src2_r = ctx.inst_field(type(inst).src2) - _c(256)
-  is_f16_output = 'F16_16X16X16_F16' in op_name or 'BF16_16X16X16_BF16' in op_name  # F16/BF16 output vs F32 output
+  src2_r = ctx.inst_field(type(inst).src2)
+  src2_r = (src2_r >= 256).where(src2_r - _c(256), src2_r)
+  output_type = op_name.split("WMMA_", 1)[1].split("_", 1)[0]
  is_bf16 = 'BF16' in op_name
  cvt = _FUNCS['bf16_to_f32'] if is_bf16 else _FUNCS['f16_to_f32']
  is_rdna4 = isinstance(inst, ir4.VOP3P)
-  # read 16x16 F16/BF16 matrix from VGPRs → flat f32 array[row*16+k]
-  def read_f16_val(src, lane, vgpr, half):
+  sz = 8 if "8" in op_name else 16
+  # read matrix from VGPRs → flat f32/i32 array[row*16+k]
+  def gval(src, lane, vgpr, ridx):
    v = ctx.rvgpr_dyn(src + _c(vgpr), UOp.const(lane, dtypes.int))
-    return cvt((v >> UOp.const(16, dtypes.uint32)) if half else (v & UOp.const(0xFFFF, dtypes.uint32)))
+    pkd = v >> UOp.const(ridx * sz, dtypes.uint32) if ridx > 0 else v
+    pkd = pkd & UOp.const((1 << sz) - 1, dtypes.uint32)
+    if "F" in output_type: return cvt(pkd)
+    return (pkd << _c(24, dtypes.uint)).bitcast(dtypes.int32) >> _c(24, dtypes.int32) # sign extend

-  # RDNA3: 16 lanes × 8 VGPRs × 2 halves, k maps linearly
-  # RDNA4: 32 lanes × 4 VGPRs × 2 halves, k bits are scrambled (k[2] goes to lane bit 4)
-  def read_f16_mat(src):
-  # (row, k) → (lane, vgpr, half)
+  # RDNA3 f16/bf16: 16 lanes × 8 VGPRs × 2 halves,    k maps linearly
+  # RDNA3 iu8:      16 lanes × 4 VGPRs × 4 quarters,  k maps linearly
+  # RDNA4:          32 lanes x 4 VGPRS x 2 halves, k bits are scrambled (k[2] goes to lane bit 4)
+  def read_mat(src):
+    n = 32 // sz # values per vgpr
+    # (row, k) → (lane, vgpr, row index)
    def ab_map(i, k):
      elem, lane = ((k & 3) | ((k >> 1) & 4), i + ((k >> 2) & 1) * 16) if is_rdna4 else (k, i)
-      return lane, elem // 2, elem % 2
-    return [read_f16_val(src, *ab_map(row, k)) for row in range(16) for k in range(16)]
-  mat_a, mat_b = read_f16_mat(src0_r), read_f16_mat(src1_r)
+      return lane, elem // n, elem % n
+    return [gval(src, *ab_map(row, k)) for row in range(16) for k in range(16)]
+
+  mat_a, mat_b = read_mat(src0_r), read_mat(src1_r)
  # (row, col) -> (lane, vgpr)
  def d_map(m, n):
    lane_bit, vgpr = (m >> 3, m & 7) if is_rdna4 else (m & 1, m >> 1)
    return n + lane_bit * 16, vgpr
-  if is_f16_output:
+  if output_type in ["F16", "BF16"]:
    # read accumulator C with f16 layout: for RDNA4, pairs of f32 vgprs pack into one f16 vgpr
    # for RDNA3, same layout as f32 but only lo 16 bits used
-    mat_c = [read_f16_val(src2_r, *((lane, vgpr // 2, vgpr % 2) if is_rdna4 else (lane, vgpr, 0)))
+    mat_c = [gval(src2_r, *((lane, vgpr // 2, vgpr % 2) if is_rdna4 else (lane, vgpr, 0)))
             for m in range(16) for n in range(16) for lane, vgpr in [d_map(m, n)]]
    mat_d = [sum(mat_a[r*16+k] * mat_b[c*16+k] for k in range(16)) + mat_c[r*16+c] for r in range(16) for c in range(16)]
    def f32_to_f16_bits(v: UOp) -> UOp: return v.cast(dtypes.half).bitcast(dtypes.uint16).cast(dtypes.uint32)
@@ -1602,18 +1643,22 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
    else:  # (rdna3) 1 f16 per VGPR (lo half only)
      stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1]), UOp.const(d_map(m, n)[0], dtypes.int), out_cvt(mat_d[m*16+n]), exec_mask)
                for m in range(16) for n in range(16)]
-  else: # f32
-    mat_c = [ctx.rvgpr_dyn(src2_r + _c(d_map(m, n)[1]), UOp.const(d_map(m, n)[0], dtypes.int)).bitcast(dtypes.float32)
+  else: # f32/i32
+    out_dt = dtypes.float32 if output_type == "F32" else dtypes.int32
+    mat_c = [ctx.rvgpr_dyn(src2_r + _c(d_map(m, n)[1]), UOp.const(d_map(m, n)[0], dtypes.int)).bitcast(out_dt)
             for m in range(16) for n in range(16)]
    mat_d = [sum(mat_a[r*16+k] * mat_b[c*16+k] for k in range(16)) + mat_c[r*16+c] for r in range(16) for c in range(16)]
    stores = [ctx.wvgpr_dyn(vdst_reg + _c(d_map(m, n)[1]), UOp.const(d_map(m, n)[0], dtypes.int), mat_d[m*16+n].bitcast(dtypes.uint32), exec_mask)
              for m in range(16) for n in range(16)]
  return UOp.sink(*stores, *ctx.inc_pc())

-def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
+def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P | irc.VOP3PX2, ctx: _Ctx) -> UOp:
  op_name = _op_name(inst)
-  if 'WMMA' in op_name and ('16X16X16_F16' in op_name or '16X16X16_BF16' in op_name): return _compile_wmma(inst, ctx)
-  if 'MFMA' in op_name and any(f'{s}X{s}X' in op_name for s in ('4', '16', '32')) and isinstance(inst, irc.VOP3P): return _compile_mfma(inst, ctx)
+  if 'WMMA' in op_name:
+    assert not isinstance(inst, irc.VOP3PX2)
+    return _compile_wmma(inst, ctx)
+  if 'MFMA' in op_name and any(f'{s}X{s}X' in op_name for s in ('4', '16', '32')) and isinstance(inst, (irc.VOP3P, irc.VOP3PX2)):
+    return _compile_mfma(inst, ctx)

  # ACCVGPR_WRITE/READ/MOV: copies between VGPR and ACCVGPR register files
  # Detect by checking operand types for ACCVGPR involvement
@@ -2044,7 +2089,7 @@ _INST_HANDLERS: dict[type, Callable[..., UOp]] = {
  irc.SOPP: _compile_sopp, irc.SMEM: _compile_smem, irc.SOP1: _compile_sop, irc.SOP2: _compile_sop, irc.SOPC: _compile_sop, irc.SOPK: _compile_sop,
  irc.VOP1: _compile_vop12, irc.VOP1_DPP16: _compile_vop12, irc.VOP2: _compile_vop12, irc.VOP2_DPP16: _compile_vop12,
  irc.VOPC: _compile_vopc, irc.VOP3: _compile_vop3,
-  irc.VOP3_SDST: _compile_vop3, irc.VOP3SD: _compile_vop3sd, irc.VOP3P: _compile_vop3p,
+  irc.VOP3_SDST: _compile_vop3, irc.VOP3SD: _compile_vop3sd, irc.VOP3P: _compile_vop3p, irc.VOP3PX2: _compile_vop3p,
  irc.VOP1_SDWA: _compile_sdwa, irc.VOP2_SDWA: _compile_sdwa, irc.VOP2_SDWA_SDST: _compile_sdwa, irc.VOPC_SDWA_SDST: _compile_sdwa,
  irc.DS: _compile_mem_op, irc.FLAT: _compile_mem_op, irc.GLOBAL: _compile_mem_op, irc.SCRATCH: _compile_mem_op,
  irc.MUBUF: _compile_mubuf,
@@ -74,6 +74,7 @@ class TestWhisper(unittest.TestCase):
        err
      )

+  @slow
  def test_transcribe_file1(self):
    self.assertEqual(transcribe_file(self.model, self.enc, TEST_FILE_1),  TRANSCRIPTION_1)

@@ -89,6 +90,7 @@ class TestWhisper(unittest.TestCase):
    self.assertEqual(TRANSCRIPTION_1,  transcriptions[0])
    self.assertEqual(TRANSCRIPTION_2,  transcriptions[1])

+  @slow
  def test_transcribe_batch21(self):
    waveforms = [load_file_waveform(TEST_FILE_2), load_file_waveform(TEST_FILE_1)]
    transcriptions = transcribe_waveform(self.model, self.enc, waveforms)
@@ -51,6 +51,14 @@ class TestWeakConstFolding(unittest.TestCase):
  def test_invalid_poison(self):
    self.assertTrue(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().is_invalid)

+  def test_single_rounding_log10_backward(self):
+    # log10 backward folds log10(2)/log(2) = 1/log(10) in one rounding, not the double-rounded 1/float32(log(10))
+    x = Tensor([1.0, 2.0, 3.0])
+    ast = next(s.src[0] for s in x.log10().sum().gradient(x)[0].schedule_linear().src if s.src[0].op is Ops.SINK)
+    const = next(u.arg for u in full_rewrite(ast).toposort() if u.op is Ops.CONST and u.dtype is dtypes.float32)
+    # correctly rounded: within half a float32 ulp of the exact value (folding at float32 lands 0.66 ulp off)
+    self.assertLess(abs(const - 1/math.log(10)), 2**-26)
+
 class TestBinaryOpsConstFolding(unittest.TestCase):
  def test_add_literal_zero(self):
    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + 0)
@@ -143,13 +143,13 @@ class TestModuloAndDivisionFolding(unittest.TestCase):
    self.assertIs(apply_rewrite(x_var_uop.cast(dtypes.weakint) % 10).render(simplify=False), x_var_uop.render(simplify=False))

  def test_full_graph_rewrite_division_with_remainder(self):
-    x_var_uop = UOp.variable('x', 7, 9)
+    x_var_uop = UOp.variable('x', 7, 9, param=True)
    optimized_sink = apply_rewrite(x_var_uop // 2)
    for x_value in range(7, 10):
      self.assertEqual(x_value // 2, evaluate_uop(optimized_sink, {'x': x_value}))

  def test_full_graph_rewrite_complex_mod_div_expression(self):
-    x_var_uop = UOp.variable('x', 1, 10)
+    x_var_uop = UOp.variable('x', 1, 10, param=True)
    optimized_sink = apply_rewrite(((x_var_uop * 5) % 3) // 2)
    for x_value in range(1, 11):
      original_result = ((x_value * 5) % 3) // 2
@@ -208,7 +208,7 @@ class TestGEPAndVectorizeRewrite(unittest.TestCase):


 import inspect
-from tinygrad.uop.ops import graph_rewrite, _substitute, track_rewrites
+from tinygrad.uop.ops import graph_rewrite, _substitute, rewrite_group
 from tinygrad.uop.symbolic import symbolic_simple

 class TestBottomUpRewrite(unittest.TestCase):
@@ -220,7 +220,7 @@ class TestBottomUpRewrite(unittest.TestCase):
    self.assertIs(gt, ret)

 # normally .substitute would be fine, but it's not tracked
-@track_rewrites()
+@rewrite_group()
 def named_substitute(name:str, uop:UOp, rel:dict[UOp, UOp]): return graph_rewrite(uop, _substitute, rel, bottom_up=True)
 def substitute(uop:UOp, rel:dict[UOp, UOp]): return named_substitute(inspect.stack()[1].function, uop, rel)

@@ -27,10 +27,15 @@ def _make_linear(buffer_lists, copies=None):
    calls.append(UOp(Ops.CALL, src=(src0, *bufs)))
  return UOp(Ops.LINEAR, src=tuple(calls))

+def _get_planned_view(buf:UOp) -> tuple[UOp, int, int]|None:
+  view = buf.src[0] if buf.op is Ops.BITCAST else buf
+  if view.op is not Ops.SHRINK or view.src[0].op is not Ops.BUFFER: return None
+  return (arena:=view.src[0]), view.src[1].val * arena.dtype.itemsize, view.src[2].val * arena.dtype.itemsize
+
 def _get_arena(buf, linear, result):
  for orig_si, new_si in zip(linear.src, result.src):
    for orig, new in zip(orig_si.src[1:], new_si.src[1:]):
-      if orig is buf and new.op is Ops.SLICE: return new.src[0]
+      if orig is buf and (planned:=_get_planned_view(new)) is not None: return planned[0]
  return None

 def check_assign(buffer_lists, copies=None):
@@ -41,8 +46,8 @@ def check_assign(buffer_lists, copies=None):
  replace_map: dict[int, tuple[UOp, int, int]] = {}
  for orig_si, new_si in zip(linear.src, result.src):
    for orig, new in zip(orig_si.src[1:], new_si.src[1:]):
-      if new.op is Ops.SLICE and id(orig) not in replace_map:
-        replace_map[id(orig)] = (new.src[0], new.src[1].val * new.src[0].dtype.itemsize, new.arg * new.dtype.itemsize)
+      if (planned:=_get_planned_view(new)) is not None and id(orig) not in replace_map:
+        replace_map[id(orig)] = planned

  # verify pinned buffers are not planned
  for buf in held_bufs:
@@ -858,6 +858,65 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.rand(32)
    check_schedule(x, 1, [Tensor._device_rng_counters[x.device]])

+  # **** custom kernel realize tests
+
+  @staticmethod
+  def _copy_fxn(name:str="copy"):
+    def copy_kernel(out:UOp, inp:UOp) -> UOp:
+      i = UOp.range(inp.numel(), 0)
+      return UOp.group(out[i].store(inp[i])).end(i).sink(arg=KernelInfo(name=name))
+    return copy_kernel
+
+  def _copy_call(self, out:Tensor, expr:Tensor, name:str="copy") -> Tensor:
+    # forge a custom kernel call with params and call args, like llm/kernels does (no Tensor.custom_kernel contiguous)
+    params = tuple(UOp.placeholder_like(u, slot=i) for i,u in enumerate((out.uop, expr.uop)))
+    return Tensor(out.uop.after(self._copy_fxn(name)(*params).call(out.uop, expr.uop)))
+
+  def test_custom_kernel_buffer_src(self):
+    # custom kernels need buffers: a buffer input must never add a realize kernel
+    y = Tensor.ones(64).contiguous().realize()
+    out = Tensor.empty_like(y)
+    check_schedule(self._copy_call(out, y), 1)
+
+  def test_custom_kernel_view_src(self):
+    # a RESHAPE over a buffer resolves to the buffer state (RESHAPEs on call args are stripped), no realize kernel
+    y = Tensor.ones(64).contiguous().realize()
+    out = Tensor.empty_like(y)
+    check_schedule(self._copy_call(out, y.reshape(8, 8).reshape(64)), 1)
+
+  def test_custom_kernel_elementwise_src(self):
+    # a computed input is not a buffer state: the call args are unwrapped to their base buffer,
+    # so the compute would be silently dropped. this must raise instead of producing wrong results
+    y = Tensor.ones(64).contiguous().realize()
+    out = Tensor.empty_like(y)
+    check_schedule(self._copy_call(out, y + y), 2)
+
+  def test_custom_kernel_lazy_const_src(self):
+    # a lazy const expression above the call has no buffer at all. this used to crash rangeify with a KeyError
+    x = Tensor.linspace(-1.0, 1.0, 64)
+    out = Tensor.empty_like(x)
+    check_schedule(self._copy_call(out, x), 2)
+
+  def test_custom_kernel_offset_view_src(self):
+    # a SHRINK with an offset over a buffer is not a buffer state either, the offset would be silently dropped
+    y = Tensor.ones(128).contiguous().realize()
+    out = Tensor.empty(64)
+    check_schedule(self._copy_call(out, y[16:80]), 2)
+
+  def test_custom_kernel_computed_src_api(self):
+    # the supported way to pass computed inputs: Tensor.custom_kernel makes inputs contiguous (one realize kernel)
+    y = Tensor.ones(64).contiguous().realize()
+    out = Tensor.empty_like(y)
+    check_schedule(Tensor.custom_kernel(out, y + y, fxn=self._copy_fxn())[0], 2)
+
+  def test_custom_kernel_on_custom_kernel(self):
+    # the output of a custom kernel is a buffer state, chaining custom kernels must not add kernels
+    y = Tensor.ones(64).contiguous().realize()
+    k1 = self._copy_call(Tensor.empty_like(y), y, name="k1")
+    k2 = self._copy_call(Tensor.empty_like(y), k1, name="k2")
+    sched, _ = check_schedule(k2, 2)
+    self.assertEqual([call.src[0].arg.name for call in sched.src], ["k1", "k2"])
+
  def test_empty_is_not_realized(self):
    a = Tensor.empty(10)
    child = a+2
@@ -2,7 +2,8 @@ import unittest, itertools

 from tinygrad.codegen.late.coalesce import indexing_simplify
 from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import UOp, Ops, graph_rewrite, pm_lower_index_dtype
+from tinygrad.uop.ops import UOp, Ops, graph_rewrite
+from tinygrad.uop.weak import pm_lower_index_dtype
 from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load
 from tinygrad.helpers import Context
 from test.helpers import full_rewrite
@@ -24,7 +25,7 @@ def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UO
  ))

 def Special(expr, nmax): return UOp(Ops.SPECIAL, src=(UOp.const(nmax),), arg=expr)
-def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax)
+def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax, param=True)
 def Range(n, nmax): return UOp.range(nmax, n)

 class TestValidIdxSimplification(unittest.TestCase):
@@ -332,7 +333,7 @@ class TestImageSimplification(unittest.TestCase):
    load = get_load_image_uop(shape, valid, idx)

    self.check(load,
-               "((((idx2*2)+r0)<11)&((((idx1*8)+r1)<3)!=True))",
+               "(((idx2*2)+r0)<11)",
               "(idx0+(idx1*512+r1*64)+-192)",
               "((((idx2*2)+r0)+(((idx1+((r1+5)//8))+1)//2))+-4)")

@@ -460,7 +461,7 @@ class TestImageSimplification(unittest.TestCase):
      self.check(load, None, "(gidx0+lidx0*1024+r0*1024+lidx1*128+-3168)", "0")
    except AssertionError:
      # TODO: fold valid
-      self.check(load, "(((lidx1<1)!=True)&(((lidx0+r0)<3)!=True)&((lidx0+r0)<19))",
+      self.check(load, "(((lidx1<1)!=True)&((lidx0+r0)<19))",
                       "(gidx0+lidx1*128+(lidx0*1024+r0*1024)+-3168)", "0")

  def test_simplify10(self):
@@ -479,7 +480,7 @@ class TestImageSimplification(unittest.TestCase):
      self.check(load, None, "(lidx2+gidx0*4+lidx0*1024+r0*1024+lidx1*256+-3264)", "0")
    except AssertionError:
      # TODO: fold valid
-      self.check(load, "(((lidx1<1)!=True)&(((lidx0+r0)<3)!=True)&((lidx0+r0)<11))",
+      self.check(load, "(((lidx1<1)!=True)&((lidx0+r0)<11))",
                       "(lidx2+gidx0*4+lidx1*256+(lidx0*1024+r0*1024)+-3264)", "0")

  def test_drop_non_monotonic_window(self):
@@ -1,10 +1,11 @@
 import unittest, pytest
-from tinygrad import dtypes, Variable
+from tinygrad import dtypes, Variable, Device
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import DEBUG, Context
-from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, GroupOp, AxisType, broadcast_axes
+from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, GroupOp, AxisType, broadcast_axes, KernelInfo
 from tinygrad.uop.symbolic import sym
 from test.helpers import to_uops_list
+from tinygrad.codegen import full_rewrite_to_sink

 simple_pm = PatternMatcher([
  (UPat.cvar('x', dtypes.weakint), lambda x: UOp.const(1.0) + UOp.const(2.0)),
@@ -156,7 +157,7 @@ class TestGraphRewrite(unittest.TestCase):
    self.assertEqual(nout.val, 3.0)

  def test_depth_2_fold(self):
-    v = UOp.variable("v", 0, 1, dtypes.float)
+    v = UOp.variable("v", 0, 1, dtypes.float, param=True)
    c1 = UOp.const(1.0)
    c2 = UOp.const(2.0)
    nout = graph_rewrite(v+c1+c2, simple_pm)
@@ -338,7 +339,7 @@ class TestUOpGraph(unittest.TestCase):
    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)

  def test_depth_2_const_fold(self):
-    v = UOp.variable("tmp", 0, 1, dtypes.int)
+    v = UOp.variable("tmp", 0, 1, dtypes.int, param=True)
    c2 = UOp.const(2, dtypes.int)
    c4 = UOp.const(4, dtypes.int)
    vc = v+c2
@@ -536,6 +537,15 @@ class TestReduceCollapse(unittest.TestCase):
    # Should become add of two separate reduces
    self.assertEqual(result.op, Ops.ADD)

+  def test_reduce_shapeless_const_unroll(self):
+    """a REDUCE over a shapeless CONST (e.g. x*0 folded late in codegen) must collapse before the expander"""
+    out = UOp.param(0, dtypes.float, (1,))
+    red = UOp.const(3.0).cast(dtypes.float).reduce(UOp.range(4, 0, AxisType.UNROLL), arg=(Ops.ADD, 0))
+    ast = UOp.sink(out.index(UOp.const(0)).store(red)).replace(arg=KernelInfo())
+    uops = full_rewrite_to_sink(ast, Device["CPU"].renderer, optimize=False).toposort()
+    self.assertNotIn(Ops.REDUCE, [u.op for u in uops])
+    self.assertIn(12.0, [u.val for u in uops if u.op is Ops.CONST])
+
 class TestMovementOps(unittest.TestCase):
  def test_pm_mops_partial_reshape_index_removes_reshape(self):
    from tinygrad.schedule.rangeify import pm_mops
@@ -3,10 +3,9 @@ import unittest, pickle, functools, math
 import z3

 from tinygrad.dtype import dtypes, ConstType, DType, Invalid
-from test.helpers import get_uops
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad.uop.spec import spec_shared, type_verify
-from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load
+from tinygrad.uop.symbolic import sym, pm_fold_cast_const, commutative, pm_simplify_valid, pm_move_where_on_load
 from tinygrad.uop.validate import uops_to_z3

 def check_uop_against_string(self, v:UOp, s:str):
@@ -16,7 +15,8 @@ def check_uop_against_string(self, v:UOp, s:str):
  s_eval = graph_rewrite(s_eval, commutative, name="cannonicalize eval")
  self.assertIs(s_eval, v, f"eval did not match simplified: {s_eval} != {v.render()} for {s}")

-def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.weakint): return UOp.variable(name,min_val,max_val,dtype)
+def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.weakint):
+  return UOp.variable(name, min_val, max_val, dtype, param=True)
 def uconst(val): return UOp.const(val)
 def usum(ops): return functools.reduce(lambda x,y: x+y, ops)
 def uand(ops): return functools.reduce(lambda x,y: x*y, ops)
@@ -35,7 +35,7 @@ class TestSymbolic(unittest.TestCase):
    self.assertEqual(solver.check(expr1 != expr2), z3.unsat, "simplified expression not equal to original")

  def helper_test_variable(self, v, n, m, s, test_z3:bool=True):
-    v_simplified = graph_rewrite(v, sym, name="simplify symbolic uop")
+    v_simplified = graph_rewrite(v, sym+pm_fold_cast_const, name="simplify symbolic uop")
    if test_z3: self.check_equal_z3(v, v_simplified)
    nmin, nmax = v_simplified.vmin, v_simplified.vmax
    check_uop_against_string(self, v_simplified, s)
@@ -442,7 +442,7 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable(uand([uconst(1), Variable("a", 0, 1)]), 0, 1, "a")

  def test_masked_shr_fold(self):
-    x = UOp.variable('x', 0, 255, dtype=dtypes.uint32)
+    x = UOp.variable('x', 0, 255, dtype=dtypes.uint32, param=True)
    self.helper_test_variable((x & -4) >> 2, 0, 63, "(x>>2)")

  def test_bool_or_not_tautology(self):
@@ -483,12 +483,12 @@ class TestSymbolic(unittest.TestCase):

  def test_div_drop_small_terms(self):
    # from openpilot, shouldnt simplify
-    gidx0 = UOp.variable("gidx0", 0, 10)
-    gidx1 = UOp.variable("gidx1", 0, 10)
-    lidx0 = UOp.variable("lidx0", 0, 1)
-    lidx1 = UOp.variable("lidx1", 0, 1)
-    ridx1005 = UOp.variable("ridx1005", 0, 2)
-    ridx1006 = UOp.variable("ridx1006", 0, 2)
+    gidx0 = UOp.variable("gidx0", 0, 10, param=True)
+    gidx1 = UOp.variable("gidx1", 0, 10, param=True)
+    lidx0 = UOp.variable("lidx0", 0, 1, param=True)
+    lidx1 = UOp.variable("lidx1", 0, 1, param=True)
+    ridx1005 = UOp.variable("ridx1005", 0, 2, param=True)
+    ridx1006 = UOp.variable("ridx1006", 0, 2, param=True)
    self.helper_test_variable((lidx1+((gidx1*18)+(ridx1005*18)+(lidx0*162))+(gidx0*2)+(ridx1006*2)+-40)//18, -3, 20,
      "(gidx1+ridx1005+lidx0*9+(gidx0+ridx1006+7)//9+-3)")

@@ -948,6 +948,11 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable(cond.where(u0, u1), 0, 1, "((a<2)!=True)")
    self.helper_test_variable(cond.where(u0, u1).where(u0, u1), 0, 1, "(a<2)")

+  def test_equivalent_const_max(self):
+    x = Variable("x", -10, 10)
+    self.helper_test_variable((x < 0).where(0, x), 0, 10, "x.maximum(0)")
+    self.helper_test_variable((0 < x).where(x, 0), 0, 10, "x.maximum(0)")
+
  def test_where_combine(self):
    cond = Variable("x", 0, 3) < 2
    a = Variable("a", 0, 3)
@@ -992,7 +997,7 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable(cond.ne(False), 0, 1, "(x<2)")

  def test_bitcast_chain(self):
-    a = UOp.variable("a", 0, 3, dtype=dtypes.int32)
+    a = UOp.variable("a", 0, 3, dtype=dtypes.int32, param=True)
    self.assertIs(graph_rewrite(a.bitcast(dtypes.float32).bitcast(a.dtype), sym), a)

  def test_negation_in_where(self):
@@ -1007,22 +1012,6 @@ class TestSymbolic(unittest.TestCase):
    b = Variable("b", 0, 3)
    self.helper_test_variable(-a<-b, False, True, "(b<a)")

-  def test_where_cast(self):
-    s = Variable("s", 0, 3, dtypes.int)
-    cond = s < 2
-    a = Variable("a", 0, 3, dtypes.int)
-    b = Variable("b", 0, 3, dtypes.int)
-    expr = cond.where(a, b).cast(dtypes.half)
-
-    # TODO: copied from render, render does not support cast
-    glbl = UOp.param(0, dtypes.int, (1,))
-    uops = get_uops(UOp(Ops.STORE, src=(glbl.index(UOp.const(0, dtypes.int)), expr)).sink())
-    rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[1]
-
-    # the vars are now scalar PARAMs
-    pvar = {u.expr: u for u in rewritten_uop.toposort() if u.op is Ops.PARAM}
-    self.assertEqual(rewritten_uop, (pvar['s']<UOp.const(2, dtypes.int)).where(pvar['a'].cast(dtypes.half), pvar['b'].cast(dtypes.half)))
-
  def test_where_merge_branches(self):
    cond1 = Variable("s", 0, 10) < 6
    cond2 = Variable("s", 0, 10) > 2
@@ -1175,7 +1164,7 @@ class TestSymbolicVariables(unittest.TestCase):
    assert (a//4 + a//6).variables() == [a]

  def test_variable_min_eq_max_bind_folds(self):
-    b = Variable("x", 1, 1).bind(1)
+    b = UOp.variable("x", 1, 1).bind(1)
    s = b.simplify()
    self.assertEqual(s.op, Ops.CONST)
    self.assertEqual(s.val, 1)
@@ -1369,6 +1358,16 @@ class TestInvalidIndex(unittest.TestCase):
    c2 = UOp.const((1, Invalid, 1, 1))
    self.assertIs((c1+c2).simplify(), UOp.const((2, Invalid, Invalid, Invalid)))

+  def test_gated_load_keeps_index_valid(self):
+    # the load executes even on gated-off iterations: gated_given_valid must not erase its mask (PADTO OOB shape)
+    buf = UOp.param(0, dtypes.bool, (17,))
+    ridx = Variable("ridx", 0, 31)
+    cond = ridx < 17
+    load = buf.index(ridx.valid(cond))
+    out = graph_rewrite(cond.where(load.where(uconst(2), uconst(0)), UOp.invalid()), sym)
+    idx = next(u for u in out.toposort() if u.op is Ops.INDEX)
+    self.assertIs(idx.src[1].get_valid(), cond.simplify())
+
 class TestStoreLoadFolding(unittest.TestCase):
  """Tests for store(index, load(index)) -> NOOP rule. This rule matches patterns that EMERGE during simplification."""
  def test_store_load_folding(self):
@@ -1,6 +1,6 @@
 import unittest, math
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.dtype import dtypes, Invalid
+from tinygrad.dtype import dtypes, Invalid, truncate

 class TestVminVmaxProperties(unittest.TestCase):
  def test_vmin_vmax_constant(self):
@@ -168,6 +168,10 @@ class TestVminVmaxProperties(unittest.TestCase):
    x = UOp.const(4.5).cast(dtypes.float)
    self.assertIs(x.ne(x.cast(dtypes.int).cast(dtypes.float)).simplify().arg, True)

+  def test_vmin_vmax_cast_int_to_float_grid(self):
+    # a cast to float only takes values on the float grid, so its bounds are the source bounds rounded at the destination
+    self.assertEqual(UOp.variable('x', 0, 16777219, dtypes.int).cast(dtypes.float)._min_max, (0.0, 16777220.0))
+
  def test_vmin_vmax_invalid(self):
    i = UOp.invalid()
    self.assertNotEqual(i.vmin, i.vmax)
@@ -317,8 +321,8 @@ class TestVminVmaxVConst(unittest.TestCase):
  def test_vmin_vmax_vconst_with_floats(self):
    # vmin and vmax for a vector constant of float values
    uop = UOp.const((1.5, -3.2, 0.0))
-    self.assertEqual(uop.vmin, -3.2)
-    self.assertEqual(uop.vmax, 1.5)
+    self.assertEqual(uop.vmin, truncate[dtypes.default_float](-3.2))
+    self.assertEqual(uop.vmax, truncate[dtypes.default_float](1.5))

  def test_vmin_vmax_vconst_with_bools(self):
    # vmin and vmax for a vector constant of bool values
@@ -5,7 +5,8 @@ from tinygrad.tensor import Tensor
 from tinygrad.helpers import Timing, Context, cdiv
 from tinygrad.dtype import dtypes, AddrSpace, ConstFloat, Invalid  # noqa: F401
 from tinygrad.device import Device
-from tinygrad.uop.ops import Ops, ParamArg, PatternMatcher, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite, pm_lower_index_dtype  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
+from tinygrad.uop.ops import Ops, ParamArg, PatternMatcher, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
+from tinygrad.uop.weak import pm_lower_index_dtype
 from tinygrad.uop.spec import spec_program, spec_shared, type_verify
 from tinygrad.uop.symbolic import sym, pm_remove_invalid
 from test.helpers import eval_uop, to_uops_list
@@ -1,11 +1,11 @@
 import unittest
 from tinygrad.helpers import DEBUG, Context
 from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import UPat, track_rewrites, GroupOp, Ops
+from tinygrad.uop.ops import UPat, rewrite_group, GroupOp, Ops
 from tinygrad.uop.upat import _get_code, upat_compile
 import dis

-@track_rewrites()
+@rewrite_group()
 def do_compile(up):
  print("\n***** COMPILE", up)
  match_code = _get_code(up, False)
@@ -1,10 +1,12 @@
 import unittest
-from tinygrad import dtypes, Variable
+from tinygrad import dtypes
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import Ops, UOp, AxisType
 from test.helpers import to_uops_list

+def Variable(name, nmin, nmax): return UOp.variable(name, nmin, nmax, param=True)
+
 class TestValidateOOB(unittest.TestCase):
  """Test z3 validation of index bounds for different ALU ops and patterns."""

@@ -3,7 +3,7 @@ from pathlib import Path
 from dataclasses import dataclass
 from typing import Generator

-from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, track_rewrites, profile_matches
+from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, TrackedPatternMatcher, graph_rewrite, rewrite_group
 from tinygrad.uop.symbolic import sym
 from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.helpers import colored, ansistrip, flatten, TracingKey, ProfileRangeEvent, ProfileEvent, Context, cpu_events, profile_marker
@@ -14,7 +14,7 @@ from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewr
 from tinygrad.viz.serve import load_rewrites, get_full_rewrite, uop_to_json, VizData, get_render, addrspace_colors
 from tinygrad.codegen import do_to_program

-@track_rewrites(name=True)
+@rewrite_group(name=True)
 def exec_rewrite(sink:UOp, pm_lst:list[PatternMatcher], names:None|list[str]=None) -> UOp:
  for i,pm in enumerate(pm_lst):
    sink = graph_rewrite(sink, TrackedPatternMatcher(pm.patterns), name=names[i] if names else None)
@@ -109,7 +109,7 @@ class TestViz(unittest.TestCase):
  def test_default_name(self):
    with save_viz() as viz:
      a = UOp.variable("a", 1, 10)
-      @track_rewrites()
+      @rewrite_group()
      def name_default(): return graph_rewrite(a, PatternMatcher([]))
      name_default()
    lst = viz.list_items()
@@ -118,7 +118,7 @@ class TestViz(unittest.TestCase):
  # name can also come from a function that returns a string
  def test_dyn_name_fxn(self):
    with save_viz() as viz:
-      @track_rewrites(name=lambda *args,ret,**kwargs: ret.render())
+      @rewrite_group(name=lambda *args,ret,**kwargs: ret.render())
      def name_from_fxn(s:UOp, arg:list|None=None): return graph_rewrite(s, PatternMatcher([]))
      name_from_fxn(UOp.variable("a", 1, 10)+1, arg=["test"])
    lst = viz.list_items()
@@ -128,18 +128,18 @@ class TestViz(unittest.TestCase):
  # name can also come from a function that returns a TracingKey
  def test_tracing_key(self):
    with save_viz() as viz:
-      @track_rewrites(name=lambda inp,ret: TracingKey("custom_name", (inp,)))
+      @rewrite_group(name=lambda inp,ret: TracingKey("custom_name", (inp,)))
      def test(s:UOp): return graph_rewrite(s, PatternMatcher([]))
      test(UOp.variable("a", 1, 10)+1)
    lst = viz.list_items()
    # NOTE: names from TracingKey do not get deduped
    self.assertEqual(lst[0]["name"], "custom_name")

-  def test_nested_track_rewrites(self):
+  def test_nested_rewrite_group(self):
    with save_viz() as viz:
-      @track_rewrites(name=lambda x,ret: TracingKey(f"inner fxn for {x.render()}", (ret,)))
+      @rewrite_group(name=lambda x,ret: TracingKey(f"inner fxn for {x.render()}", (ret,)))
      def inner(x:UOp): return graph_rewrite(x, PatternMatcher([]), name="each")
-      @track_rewrites(name=lambda *args,ret: f"outer rewrite of {len(args)} inputs")
+      @rewrite_group(name=lambda *args,ret: f"outer rewrite of {len(args)} inputs")
      def outer(*xs:tuple[UOp, ...]): return graph_rewrite(UOp.sink(*[inner(x) for x in xs]), PatternMatcher([]), name="all")
      items = ["a", "b", "c"]
      outer(*[UOp.variable(x, 1, 10) for x in items])
@@ -156,13 +156,13 @@ class TestViz(unittest.TestCase):
      self.assertEqual(len(steps), 1)
      self.assertEqual(steps[0]["name"], "each")

-  def test_profile_matches(self):
+  def test_rewrite_group_nested(self):
    with save_viz() as viz:
-      @profile_matches
+      @rewrite_group(new_ctx=False)
      def nested_function(u:UOp):
        for i in range(2): graph_rewrite(u, PatternMatcher([]), name=f"step {i+1}")

-      @track_rewrites()
+      @rewrite_group()
      def main_rewrite(u:UOp):
        graph_rewrite(u, PatternMatcher([]), name="init")
        nested_function(u)
@@ -173,9 +173,9 @@ class TestViz(unittest.TestCase):
    self.assertEqual(steps[1]["name"], "nested_function")
    self.assertEqual(len(steps), 4)

-  def test_profile_matches_invalid_arg(self):
+  def test_rewrite_group_invalid_arg(self):
    with save_viz():
-      @profile_matches
+      @rewrite_group(new_ctx=False)
      def invalid_fxn(arg:str): return graph_rewrite(UOp(Ops.SINK), PatternMatcher([]))
      with self.assertRaisesRegex(AssertionError, "invalid match tracing input"):
        invalid_fxn("test")
@@ -305,10 +305,10 @@ class TestVizTree(unittest.TestCase):

  def test_tree_view(self):
    with save_viz() as viz:
-      a = UOp.variable("a",0,10)
-      b = UOp.variable("b",0,10)
-      c = UOp.variable("c",0,10)
-      d = UOp.variable("d",0,10)
+      a = UOp.variable("a",0,10,param=True)
+      b = UOp.variable("b",0,10,param=True)
+      c = UOp.variable("c",0,10,param=True)
+      d = UOp.variable("d",0,10,param=True)
      sink = UOp.sink(a+b, c+d)
      def tree_rewrite(): return graph_rewrite(sink, root, name="root")
      tree_rewrite()
@@ -395,7 +395,7 @@ class TestVizIntegration(unittest.TestCase):
    graph = next(viz.get_details(0, 0))["graph"]
    self.assertEqual(len([n for n in graph.values() if repr(metadata) in n["label"]]), 1)

-  # tracing also works without a track_rewrites context
+  # tracing also works without a rewrite_group context
  # all graph_rewrites get put into the default group
  def test_default_tracing(self):
    with save_viz() as viz:
@@ -407,11 +407,11 @@ class TestVizIntegration(unittest.TestCase):
    self.assertEqual(len(ls), 1)
    self.assertEqual(ls[0]["name"], "default graph_rewrite")

-  # using @track_rewrites organizes function calls into groups
+  # using @rewrite_group organizes function calls into groups
  # and nicely counts function calls.
  def test_group_traces(self):
    with save_viz() as viz:
-      @track_rewrites()
+      @rewrite_group()
      def test(root):
        return graph_rewrite(root, sym)
      test(c:=UOp.const(1))
@@ -420,11 +420,11 @@ class TestVizIntegration(unittest.TestCase):
    self.assertEqual(len(ls), 2)
    for i in range(2): self.assertEqual(ls[i]["name"], f"test n{i+1}")

-  # @track_rewrites always starts a new group.
+  # @rewrite_group always starts a new group.
  def test_group_combined(self):
    with save_viz() as viz:
      def default_test(root): return graph_rewrite(root, sym)
-      tracked_test = track_rewrites()(default_test)
+      tracked_test = rewrite_group()(default_test)
      c = UOp.const(1)
      default_test(c+1) # goes to the default group
      tracked_test(c)   # all rewrites after this go inside the second group.
@@ -239,6 +239,15 @@ class TestKernelOpts(unittest.TestCase):
    helper_linearizer_opt(a.sum().exp(), [[Opt(OptOps.PADTO, 0, 32)],])
    helper_linearizer_opt(a.sum(0).exp(), [[Opt(OptOps.PADTO, 1, 32)],])

+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
+  @unittest.expectedFailure
+  def test_padto_group_full_unroll_sum(self):
+    a = Tensor.ones(2, 28, 4096, dtype=dtypes.bfloat16).realize()
+    out = ((a * 0.5).float().square()).sum(axis=(0, 2))
+    opts_to_apply = [Opt(OptOps.GROUPTOP, 1, 256), Opt(OptOps.PADTO, 3, 32), Opt(OptOps.UNROLL, 2, 0), Opt(OptOps.UPCAST, 0, 7)]
+    helper_linearizer_opt(out, [opts_to_apply], check_default_opt=False)
+
  def test_padto_sum(self):
    N = 18
    # NOTE: this setup prevents 17 * 17 contiguous merged into one dimension
@@ -79,7 +79,8 @@ class TestTensorCores(unittest.TestCase):
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
  def test_tensor_cores(self):
    for tc in Device[Device.DEFAULT].renderer.tensor_cores:
-      helper_tc_allclose(tc.dims[0], tc.dims[1], tc.dims[2], tc.dtype_in, tc.dtype_out, axis=0, tc_opt=0)
+      with self.subTest(tc=tc):
+        helper_tc_allclose(tc.dims[0], tc.dims[1], tc.dims[2], tc.dtype_in, tc.dtype_out, axis=0, tc_opt=0)

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
  def test_tensor_cores_nested_reduce(self):
@@ -101,7 +102,8 @@ class TestTensorCores(unittest.TestCase):
      if Device.DEFAULT == "CPU" and DEV.renderer == "LLVM":
        assert "0x201000" in prg.src[2].arg
      elif Device.DEFAULT == "AMD" and DEV.renderer == "LLVM":
-        assert "@llvm.amdgcn.wmma" in prg.src[2].arg
+        # RDNA emits wmma intrinsics, CDNA emits mfma intrinsics
+        assert ("@llvm.amdgcn.wmma" in prg.src[2].arg) or ("@llvm.amdgcn.mfma" in prg.src[2].arg)
      elif Device[Device.DEFAULT].renderer.suffix == "PTX":
        assert "mma.sync.aligned" in prg.src[2].arg
      else:
@@ -181,11 +183,13 @@ class TestTensorCores(unittest.TestCase):
  @unittest.skipIf(Device.DEFAULT == "PYTHON", "slow on EMULATED device")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
  def test_tensor_cores_unroll_phi(self):
-    tc = Device[Device.DEFAULT].renderer.tensor_cores[0]
-    x, y = Tensor.rand(128, 128, dtype=tc.dtype_in), Tensor.rand(128, 128, dtype=tc.dtype_in)
+    # skip fp8 tcs: the unoptimized ALU baseline quantizes products to fp8 (JAX promotion), which legitimately
+    # differs from the MFMA path (f32 accumulation), so the baseline-vs-TC numerical gate can't hold for fp8.
+    tc = next(tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in not in dtypes.fp8s)
+    x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in)
    r = x.matmul(y, dtype=tc.dtype_out)
-    opts = [Opt(OptOps.UNROLL, 0, 4)]
-    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
+    opts = [Opt(OptOps.UNROLL, 0, 2)]
+    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False)
    for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src):
      if u.op is Ops.WMMA:
        assert u.src[-1].src[0].op != Ops.STORE
@@ -195,11 +199,11 @@ class TestTensorCores(unittest.TestCase):
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
  @unittest.skipIf(Device.DEFAULT in {"CPU"}, "CPU does not support using a different type for accumulation")
  def test_tensor_cores_unroll_casted_phi(self):
-    tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out][0]
-    x, y = Tensor.rand(128, 128, dtype=tc.dtype_in), Tensor.rand(128, 128, dtype=tc.dtype_in)
+    tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out and tc.dtype_in not in dtypes.fp8s][0]
+    x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in)
    r = x.matmul(y, dtype=tc.dtype_out)
-    opts = [Opt(OptOps.UNROLL, 0, 4)]
-    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
+    opts = [Opt(OptOps.UNROLL, 0, 2)]
+    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False)
    for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src):
      if u.op is Ops.WMMA:
        #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
@@ -211,11 +215,11 @@ class TestTensorCores(unittest.TestCase):
  @unittest.skipIf(Device.DEFAULT in {"CPU"}, "CPU does not support using a different type for accumulation")
  def test_tensor_cores_unroll_casted_phi_with_children(self):
    # all STORE children are outside the loop
-    tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out][0]
-    x, y = Tensor.rand(128, 128, dtype=tc.dtype_in), Tensor.rand(128, 128, dtype=tc.dtype_in)
+    tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out and tc.dtype_in not in dtypes.fp8s][0]
+    x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in)
    r = x.matmul(y, dtype=tc.dtype_out).relu()
-    opts = [Opt(OptOps.UNROLL, 0, 4)]
-    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
+    opts = [Opt(OptOps.UNROLL, 0, 2)]
+    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False)
    for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src):
      if u.op is Ops.WMMA:
        #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
@@ -1,7 +1,9 @@
 import unittest

-from tinygrad import Tensor, Device, dtypes
-from tinygrad.helpers import fetch, round_up
+from tinygrad import Tensor, Device, Variable, dtypes
+from tinygrad.helpers import DEV, fetch, round_up
+from tinygrad.engine.realize import compile_linear
+from tinygrad.uop.ops import Ops
 from extra.hevc.hevc import parse_hevc_file_headers, nv_gpu
 from extra.hevc.decode import hevc_decode

@@ -63,7 +65,7 @@ class TestHevc(unittest.TestCase):
    self.assertEqual(list(frame3.initreflistidxl1), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
    self.assertEqual(list(frame3.RefDiffPicOrderCnts), [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])

-  @unittest.skipUnless(Device.DEFAULT == "NV", "NV only")
+  @unittest.skipUnless(Device.DEFAULT == "NV" and not DEV.interface.startswith("MOCK"), "real NV only")
  def test_hevc_decode(self):
    url = "https://github.com/haraschax/filedump/raw/09a497959f7fa6fd8dba501a25f2cdb3a41ecb12/comma_video.hevc"
    dat = fetch(url, headers={"Range": f"bytes=0-{512<<10}"}).read_bytes()
@@ -83,5 +85,22 @@ class TestHevc(unittest.TestCase):
      self.assertEqual(f.dtype, dtypes.uint8)
      self.assertEqual(f.device, "NV")

+  @unittest.skipUnless(Device.DEFAULT == "NV", "NV only")
+  def test_hevc_decode_compile(self):
+    url = "https://github.com/haraschax/filedump/raw/09a497959f7fa6fd8dba501a25f2cdb3a41ecb12/comma_video.hevc"
+    dat = fetch(url, headers={"Range": f"bytes=0-{512<<10}"}).read_bytes()
+
+    opaque, frame_info, _, _, luma_w, luma_h, _ = parse_hevc_file_headers(dat)
+    offset, sz, frame_pos, max_hist, _ = frame_info[1]
+    out_image_size = luma_h + (luma_h + 1) // 2, round_up(luma_w, 64)
+    history = [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV") for _ in range(max_hist)]
+    decoded = Tensor(dat, device="NV")[offset:offset+sz].decode_hevc_frame(
+      Variable("pos", 0, max_hist + 1).bind(frame_pos), out_image_size, opaque[1], history)
+
+    compiled = compile_linear(decoded.linear_with_vars()[0])
+    self.assertTrue(any(call.src[0].op is Ops.PROGRAM for call in compiled.src))
+    encdec_calls = [call for call in compiled.src if call.src[0].op is Ops.CUSTOM_FUNCTION and call.src[0].arg == "encdec"]
+    self.assertEqual(len(encdec_calls), 1)
+
 if __name__ == "__main__":
  unittest.main()
@@ -540,6 +540,10 @@ class TestAssign(unittest.TestCase):
    c = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
    c[0:2].bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000], dtype=dtypes.uint32)).realize()
    np.testing.assert_allclose(c.numpy(), [4.0, 3.0, 3.0, 4.0])
+    # without .realize()
+    a = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
+    a.bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.uint32))
+    np.testing.assert_allclose(a.numpy(), [4.0, 3.0, 2.0, 1.0])

  def test_assign_bitcast_different_size(self):
    # assign to a shape-changing bitcast view (only works on DISK currently)
@@ -212,6 +212,18 @@ class TestCallSchedule(unittest.TestCase):
    out = f(a, v.bind(5))
    np.testing.assert_allclose(out.numpy(), [5., 10., 15.])

+  def test_precompile_scoped_bind_arg(self):
+    @function(precompile=True)
+    def f(x:Tensor, scale:UOp) -> Tensor: return x * scale
+    a = Tensor.ones(3)
+    x = f(a, UOp.variable("scale_a", 1, 100).bind(2))
+    y = f(a, UOp.variable("scale_b", 1, 100).bind(3))
+    fx = next(u for u in x.uop.toposort() if u.op is Ops.FUNCTION)
+    fy = next(u for u in y.uop.toposort() if u.op is Ops.FUNCTION)
+    self.assertEqual(fx.src[0].key, fy.src[0].key)
+    np.testing.assert_equal(x.numpy(), [2, 2, 2])
+    np.testing.assert_equal(y.numpy(), [3, 3, 3])
+
  def test_precompile_schedule_cache_hit(self):
    """two instances of the same @function should produce identical function body keys (schedule cache hit)"""
    @function(precompile=True)
@@ -347,5 +359,15 @@ class TestCallMultiSharded(unittest.TestCase):
    np.testing.assert_allclose(a.grad.numpy(), b.numpy(), rtol=1e-5)
    np.testing.assert_allclose(b.grad.numpy(), a.numpy(), rtol=1e-5)

+  def test_symbolic_reshape_shard_axis(self):
+    toks = UOp.variable("toks", 1, 2).bind(2)
+    devs = ("CPU:0", "CPU:1")
+    x = Tensor(np.arange(16, dtype=np.float32).reshape(1, 2, 8)).shard(devs, axis=2).realize()
+    @function
+    def f(x:Tensor) -> Tensor: return x.reshape(1, x.shape[1], 2, 4)
+    out = f(x[:, :toks]).realize()
+    self.assertEqual(out.uop.axis, 2)
+    np.testing.assert_equal(out[:1, :2].to(devs[0]).numpy(), np.arange(16, dtype=np.float32).reshape(1, 2, 2, 4))
+
 if __name__ == '__main__':
  unittest.main()
@@ -222,6 +222,12 @@ class TestAutoCastType(unittest.TestCase):
    t.square().mean().backward()
    np.testing.assert_allclose(t.grad.numpy().flatten(), [60000 * 2 / (N*N)] * N*N)

+  @unittest.skipUnless(dtypes.half in supported_dtypes, "need half")
+  def test_var_half_precision_large_n(self):
+    # the element count (70000) exceeds half max (65504): the denominator must not be materialized in half
+    t = Tensor([[0.0, 1.0]], dtype=dtypes.half).expand(35000, 2).contiguous()
+    np.testing.assert_allclose(t.var().numpy(), 0.25, rtol=1e-3)
+
  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "Precision error")
  @unittest.skipUnless(dtypes.half in supported_dtypes, "need half")
  def test_softmax_dtype(self):
@@ -3,7 +3,8 @@ import tempfile, unittest, math
 from tinygrad import Tensor, dtypes, TinyJit
 from tinygrad.helpers import Context
 from tinygrad.dtype import least_upper_float
-from tinygrad.uop.ops import UOp, Ops, dtype_from_uop, graph_rewrite, pm_lower_index_dtype, pm_commit_weak
+from tinygrad.uop.ops import UOp, Ops, dtype_from_uop, graph_rewrite
+from tinygrad.uop.weak import pm_lower_index_dtype, pm_commit_weak
 from tinygrad.uop.symbolic import symbolic_simple
 from tinygrad.uop.spec import spec_shared, type_verify
 from tinygrad.engine.jit import JitError
@@ -62,6 +63,31 @@ class TestWeakPromotion(unittest.TestCase):
    self.assertEqual((x._uop.base.op, x._uop.base.val, x.dtype, x.shape, y.dtype),
                     (Ops.CONST, 1, dtypes.weakfloat, (1,), dtypes.float32))

+  def test_weak_expression_anchors_at_strong_lub(self):
+    # regression test for the HALF bert nan (#17408, reverted in #17409): lub(int32, weakfloat)==weakfloat makes
+    # `loss_mask.sum() + 1e-5` a weakfloat EXPRESSION. Meeting a strong float in a binop must pin it at the lub
+    denom = (Tensor.zeros(912, dtype=dtypes.int32) != Tensor.zeros(912, dtype=dtypes.float32)).sum() + 1e-5
+    self.assertIs(denom.dtype, dtypes.weakfloat)  # the setup: the denominator expression itself is weak
+    x, y = Tensor([2048.0], dtype=dtypes.float32)._broadcasted(denom)
+    self.assertIs(y.dtype, dtypes.float32)
+    recips = [u for u in (x / y)._uop.toposort() if u.op is Ops.RECIPROCAL]
+    self.assertEqual([(u.dtype, u.src[0].dtype) for u in recips], [(dtypes.float32, dtypes.float32)])
+    with Context(DEFAULT_FLOAT=dtypes.float16):
+      committed = graph_rewrite((UOp.const(1).cast(dtypes.int32) + UOp.const(1.0)).cast(dtypes.float32), pm_lower_index_dtype, ctx={})
+    self.assertEqual([u.dtype for u in committed.toposort() if u.op is Ops.ADD], [dtypes.float32])
+
+  def test_cast_weak_expression_commits_at_cast_floor(self):
+    # the floor never narrows: a cast BELOW the default does not pull the compute width down with it
+    with Context(DEFAULT_FLOAT=dtypes.float32):
+      narrowed = graph_rewrite((UOp.const(1.0) + UOp.const(2.0)).cast(dtypes.float16), pm_lower_index_dtype, ctx={})
+    self.assertEqual((narrowed.dtype, narrowed.src[0].dtype), (dtypes.float16, dtypes.float32))
+
+  def test_cast_weak_expression_value_uses_cast_floor(self):
+    with Context(DEFAULT_FLOAT=dtypes.float16):
+      denom = Tensor.ones(1, dtype=dtypes.int32, device="CPU").sum() * 70000 + 1e-5
+      out = Tensor(1.0, dtype=dtypes.float32, device="CPU") / denom
+      self.assertAlmostEqual(out.item(), 1 / (70000 + 1e-5), places=10)
+
  def test_uop_scalar_const_lifts_kind(self):
    for dtype, value, out_dtype, const_dtype in ((dtypes.weakint, 1, dtypes.weakint, dtypes.weakint),
                                                 (dtypes.int32, 1, dtypes.int32, dtypes.weakint),
@@ -25,10 +25,10 @@ class TestHCQUnit(unittest.TestCase):
    cpu_call = UOp(Ops.PROGRAM, src=(UOp.sink(),)).call(UOp.new_buffer("CPU", 1, dtypes.float))
    gpu_devs = [d0]

-    # local MMIO: GPU works alone and with CPU in batch (cpu_support=True)
+    # CPU uses HCQ2 and is no longer batched into legacy HCQ graphs.
    assert HCQGraph.supports_uop(gpu_devs, gpu_call) is True
-    assert HCQGraph.supports_uop(gpu_devs, cpu_call) is True
-    assert HCQGraph.supports_uop(gpu_devs + [cpu_dev], gpu_call) is True
+    assert HCQGraph.supports_uop(gpu_devs, cpu_call) is False
+    assert HCQGraph.supports_uop(gpu_devs + [cpu_dev], gpu_call) is False

    # USB MMIO: GPU-only still works, but CPU batching must be rejected (cpu_support=False)
    orig_view = d0.timeline_signal.base_buf.view
@@ -13,12 +13,18 @@ V_TOKS = UOp.variable("toks", 1, 32)  # 32 is the default chunk_size in generate
 class TestTransformerGenerate(unittest.TestCase):
  def test_warmup(self):
    model, calls = Transformer(TEST_CONFIG), []
-    def generate(tokens):
+    def generate(tokens, **kwargs):
      calls.append(tokens)
      yield from (1, 2)
    with patch.object(model, "generate", generate): model.warmup()
    self.assertEqual(calls, [[0], [0]])

+  def test_warmup_then_generate_with_default_chunk(self):
+    # warmup must not capture JIT graphs that generate()'s default chunk_size then rejects
+    model = Transformer(TEST_CONFIG)
+    model.warmup()
+    self.assertIsInstance(next(model.generate([5, 6, 7, 8])), int)
+
  def test_first_recurrent_generate_before_state_init(self):
    model = Transformer(TEST_CONFIG)
    model.has_recurrent_block = True
@@ -38,6 +44,15 @@ class TestTransformerGenerate(unittest.TestCase):
      next(model.generate([1, 2, 3, 4, 5, 42, 10]))
    self.assertEqual(calls, [((1, 1), V_START_POS.bind(5)), ((1, 1), V_START_POS.bind(6))])

+  def test_recurrent_divergent_prompt_restarts(self):
+    model, calls = Transformer(TEST_CONFIG), []
+    model.has_recurrent_block, model._cached_tokens = True, [1, 2, 9]
+    def mock_call(self, tokens, start_pos, temperature):
+      calls.append(start_pos)
+      return Tensor([[42]])
+    with patch.object(Transformer, '__call__', mock_call): next(model.generate([1, 2, 10, 11]))
+    self.assertEqual(calls[0], V_START_POS.bind(0))
+
  def test_template_starts_reasoning(self):
    router = StreamRouter(reasoning=True)
    self.assertEqual(list(router.route("reasoning</think>answer")),
@@ -178,6 +193,12 @@ class TestTransformerGenerate(unittest.TestCase):
    # with temperature=2.0, we should see at least 2 distinct outputs across 5 runs
    self.assertGreater(len(runs), 1, "high temperature should produce varied outputs")

+  def test_recurrent_temperature_high_produces_variety(self):
+    model = Transformer(TEST_CONFIG)
+    model.has_recurrent_block = True
+    outputs = {model.forward(Tensor([[1]]), 0, Tensor([2.0])).item() for _ in range(5)}
+    self.assertGreater(len(outputs), 1)
+
  def test_temperature_passed_to_forward(self):
    """Temperature from generate should be passed through to __call__."""
    model = Transformer(TEST_CONFIG)
@@ -1,5 +1,4 @@
 import unittest
-from unittest.mock import MagicMock
 from tinygrad import Device
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.dtype import dtypes
@@ -11,36 +10,27 @@ class TestMetalGraph(unittest.TestCase):
    self.MetalGraph = MetalGraph
    self.dev = Device[Device.DEFAULT]

-  def metal_buf(self, offset):
-    buf = MagicMock()
-    if offset > 0:
-      buf.op = Ops.SLICE
-      src = MagicMock()
-      src.dtype = dtypes.uint8
-      buf.src = (src, UOp.const(offset))
-      buf.dtype = dtypes.uint8
-    else:
-      buf.op = Ops.BUFFER
-    buf.device = Device.DEFAULT
-    return buf
+  def metal_buf(self, offset, bitcast=False):
+    size = 4 if bitcast else 1
+    buf = UOp.new_buffer(Device.DEFAULT, offset+size, dtypes.uint8)
+    if offset: buf = buf[offset:offset+size]
+    return buf.bitcast(dtypes.float32) if bitcast else buf

-  def call(self, *bufs):
-    c = MagicMock()
-    c.src = (MagicMock(op=Ops.PROGRAM),) + tuple(bufs)
-    return c
+  def supports_uop(self, *bufs):
+    return self.MetalGraph.supports_uop([self.dev], UOp(Ops.PROGRAM, src=(UOp.sink(),)).call(*bufs))

  def test_supports_uop_normal_offset(self):
-    assert self.MetalGraph.supports_uop([self.dev], self.call(self.metal_buf(0), self.metal_buf(100), self.metal_buf(0xFFFFFFFF))) is True
+    assert self.supports_uop(self.metal_buf(0), self.metal_buf(100), self.metal_buf(0xFFFFFFFF)) is True

  def test_supports_uop_overflow_offset(self):
-    assert self.MetalGraph.supports_uop([self.dev], self.call(self.metal_buf(0), self.metal_buf(0x100000000))) is False
+    assert self.supports_uop(self.metal_buf(0), self.metal_buf(0x100000000)) is False

-  def test_supports_uop_nonmetal_buf(self):
-    # non-SLICE ops should not be checked for offset
-    buf = MagicMock()
-    buf.op = Ops.BUFFER
-    buf.device = Device.DEFAULT
-    self.MetalGraph.supports_uop([self.dev], self.call(buf))
+  def test_supports_uop_non_view_buf(self):
+    assert self.supports_uop(self.metal_buf(0)) is True
+
+  def test_supports_uop_bitcast(self):
+    assert self.supports_uop(self.metal_buf(0xFFFFFFFF, bitcast=True)) is True
+    assert self.supports_uop(self.metal_buf(0x100000000, bitcast=True)) is False

 if __name__ == "__main__":
  unittest.main()
@@ -384,6 +384,12 @@ class TestMultiTensor(unittest.TestCase):
      np.testing.assert_allclose(r.numpy(), np.ones(256)+np.ones(256), atol=1e-4, rtol=1e-5)
    assert jf.captured is not None

+  def test_symbolic_broadcast_copy(self):
+    rows = Variable("rows", 1, 4).bind(3)
+    out = Tensor.ones(rows, 8).to(devices_2).realize()
+    self.assertEqual(out.shape, (rows, 8))
+    np.testing.assert_equal(out[:3].to(Device.DEFAULT).numpy(), np.ones((3, 8)))
+
  def test_multitensor_jit_in_list(self):
    # test MULTI tensor inside a list container - exercises the container unpacking + MULTI unpacking
    @TinyJit
@@ -1,11 +1,16 @@
 import unittest
 import functools
-from tinygrad import Tensor, Variable, UOp
+from tinygrad import Tensor, Variable, UOp, function
 from tinygrad.uop.ops import KernelInfo
 from tinygrad.schedule import schedule_cache

-def custom_set0_kernel(A:UOp, num:int) -> UOp:
-  return A[0].set(num).sink(arg=KernelInfo(f"custom_set0_{num}"))
+def custom_add_kernel(A:UOp, B:UOp, num:int=0) -> UOp:
+  return A[0].set(B[0] + num).sink(arg=KernelInfo(f"custom_add_{num}"))
+
+def custom_add_backward(grad_output:UOp, _) -> tuple[None, UOp]:
+  grad = Tensor.invalids(*grad_output.shape, dtype=grad_output.dtype, device=grad_output.device)
+  grad = Tensor.custom_kernel(grad, Tensor(grad_output, device=grad_output.device), fxn=functools.partial(custom_add_kernel, num=0))[0]
+  return None, grad.uop

 class TestScheduleCache(unittest.TestCase):
  def test_bound_variable_reuses_cache(self):
@@ -25,27 +30,27 @@ class TestScheduleCache(unittest.TestCase):

  def test_custom_kernel(self):
    for i in range(4):
-      a = Tensor.empty(1)
-      a = Tensor.custom_kernel(a, fxn=functools.partial(custom_set0_kernel, num=i))[0]
+      a, b = Tensor.empty(1), Tensor.ones(1)
+      a = Tensor.custom_kernel(a, b, fxn=functools.partial(custom_add_kernel, num=i))[0]
      a.realize()
-      self.assertEqual(a.item(), i)
+      self.assertEqual(a.item(), i+1)

  def test_same_custom_function_reuses_cache(self):
    schedule_cache.clear()
-    fxn = functools.partial(custom_set0_kernel, num=10)
+    fxn = functools.partial(custom_add_kernel, num=10)

    # first run
-    a = Tensor.empty(1)
-    a = Tensor.custom_kernel(a, fxn=fxn)[0]
+    a, x = Tensor.empty(1), Tensor.ones(1)
+    a = Tensor.custom_kernel(a, x, fxn=fxn)[0]
    a.realize()
-    self.assertEqual(a.item(), 10)
+    self.assertEqual(a.item(), 11)
    cache_size_after_first = len(schedule_cache)

    # second run with same function should reuse cache
-    b = Tensor.empty(1)
-    b = Tensor.custom_kernel(b, fxn=fxn)[0]
+    b, x = Tensor.empty(1), Tensor.ones(1)
+    b = Tensor.custom_kernel(b, x, fxn=fxn)[0]
    b.realize()
-    self.assertEqual(b.item(), 10)
+    self.assertEqual(b.item(), 11)
    self.assertEqual(len(schedule_cache), cache_size_after_first)

  def test_simple(self):
@@ -65,5 +70,29 @@ class TestScheduleCache(unittest.TestCase):
      print(num)
    self.assertEqual(len(schedule_cache), start_len_schedule_cache)

+  def test_simple_precompile(self):
+    @function(precompile=True, precompile_backward=True)
+    def f(x:Tensor) -> Tensor:
+      out = Tensor.invalids(*x.shape, dtype=x.dtype, device=x.device)
+      out = Tensor.custom_kernel(out, x, fxn=functools.partial(custom_add_kernel, num=10), grad_fxn=custom_add_backward)[0]
+      return out + x
+
+    # warmup
+    x = Tensor.ones(1).realize()
+    out = f(x)
+    out.backward(x)
+    self.assertEqual(out.item(), 12)
+    self.assertEqual(x.grad.item(), 2)
+
+    # use the cache next time function is called
+    start_len_schedule_cache = len(schedule_cache)
+    for _ in range(3):
+      x = Tensor.ones(1).realize()
+      out = f(x)
+      out.backward(x)
+      self.assertEqual(out.item(), 12)
+      self.assertEqual(x.grad.item(), 2)
+    self.assertEqual(len(schedule_cache), start_len_schedule_cache)
+
 if __name__ == "__main__":
  unittest.main()
@@ -1,221 +0,0 @@
-from dataclasses import dataclass, field
-from tinygrad.dtype import dtypes, AddrSpace
-from tinygrad.uop.ops import UOp, UPat, PatternMatcher, Ops, GroupOp, ParamArg, graph_rewrite, track_rewrites
-from tinygrad.helpers import VIZ, pluralize, all_int
-
-@dataclass
-class AllocCtx:
-  uop_list: list[UOp] = field(default_factory=list)
-  buffer_map: dict[UOp, UOp] = field(default_factory=dict)
-  bases: set[UOp] = field(default_factory=set)
-  assigns: list[UOp] = field(default_factory=list)
-  replacements: list[UOp] = field(default_factory=list)
-
-def tag_uop(ctx:AllocCtx, x:UOp):
-  if x.tag is not None: return None
-  ctx.uop_list.append(x)
-  return x.replace(tag=(len(ctx.uop_list)-1,))
-
-def disk_like(u:UOp): return isinstance(u.device, str) and u.device.startswith(("DISK", "TINYFS"))
-
-def disk_copy_is_buffer(ctx:AllocCtx, u:UOp):
-  # copies to disk are replaced with the disk buffer
-  if disk_like(u) and u.tag is None:
-    ctx.buffer_map[u] = u.empty_like()
-    return u.rtag(())
-  # all copies from disk/numpy are realized into a real buffer
-  from_creation = isinstance(u.src[0].device, str) and u.src[0].device.startswith(("NPY", "DISK", "PYTHON", "TINYFS"))
-  if from_creation: return tag_uop(ctx, u)
-
-# CONTIGUOUS and AFTER + parents are the only nodes that get updated
-add_tags = PatternMatcher([
-  (UPat(Ops.COPY, name="u"), disk_copy_is_buffer),
-  # no tag on copies that are assigned via STORE+AFTER — merge COPY tag into AFTER
-  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(name="dest"), UPat(Ops.COPY, name="c")))), name="a"),
-   lambda a,c,dest: a.replace(src=(a.src[0], a.src[1].replace(src=(dest, c.rtag(())))), tag=a.tag+c.tag) if a.tag and c.tag else None),
-  (UPat((Ops.CONTIGUOUS, Ops.AFTER), name="x"), tag_uop),
-  (UPat(GroupOp.All, name="x"), lambda ctx,x: tag_uop(ctx,x) if x in ctx.bases else None),
-])
-
-def replace_contig_with_store_after(u:UOp):
-  # can't allocate a buffer for a virtual value
-  if u.is_virtual: return None
-  # if size is 0, remove the contig
-  if 0 in u.shape: return u.src[0]
-  # no real contig for DISK/TINYFS tensors, they are left alone
-  if disk_like(u): return u.rtag(None)
-  buf = u.empty_like()
-  return buf.after(buf.store(u.src[0])).rtag(u.tag)
-
-def replace_store_after_with_contig(u:UOp, src:UOp):
-  assigned_to = u
-  while assigned_to.op in {Ops.BITCAST, Ops.AFTER, Ops.UNSHARD}: assigned_to = assigned_to.src[0].base
-  if assigned_to.op not in {Ops.BUFFER, Ops.SLICE}: return src.contiguous(tag=u.tag)
-
-def _make_buffer_view(src:UOp) -> UOp|None:
-  """If movement ops on src collapse to a contiguous range, return SLICE. Otherwise None."""
-  if (offset := src.contiguous_view_offset()) is None: return None
-  buf = src.base
-  if buf.op is Ops.SLICE:
-    byte_offset = buf.src[1].val * buf.src[0].dtype.itemsize + offset * src.dtype.itemsize
-    buf = buf.src[0]
-    if byte_offset % buf.dtype.itemsize != 0: return None
-    offset = byte_offset // buf.dtype.itemsize
-  return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(offset)), src.numel())
-
-def contiguous_mops_to_view(c:UOp, src:UOp):
-  """MOPS(BUFFER) → SLICE when movement ops collapse to a contiguous range."""
-  buf = src.base
-  if buf.op not in {Ops.BUFFER, Ops.SLICE, Ops.UNSHARD}: return None
-  if src.op is Ops.RESHAPE and src.src[0].op in {Ops.BUFFER, Ops.SLICE} and c.op is not Ops.BITCAST: return None
-  if c.op is not Ops.BITCAST and src.op is Ops.BUFFER: return None
-
-  # no symbolic shape
-  if not all_int(c.shape): return None
-
-  if buf.op is not Ops.UNSHARD and (view := _make_buffer_view(src)) is not None:
-    view = (view.replace(dtype=c.dtype, arg=c.numel()) if c.op is Ops.BITCAST else view).reshape(c.shape)
-    return c.replace(src=(view,)) if c.op is Ops.COPY else view
-
-  # for UNSHARD tensors, use multi_pm to resolve per-shard movement ops, then create SLICE on the resolved result
-  if not isinstance(c.device, str):
-    from tinygrad.schedule.multi import multi_pm
-    resolved = graph_rewrite(src, multi_pm, name="multi_buffer_view")
-    if resolved.op is not Ops.UNSHARD: return None
-    if (view := _make_buffer_view(resolved.src[0])) is None: return None
-    return view.reshape(resolved.src[0].shape).unshard(resolved.arg, resolved.src[1:]).contiguous(tag=c.tag)
-
-  return None
-
-def _precompiled_output_redirect(s:UOp, t:UOp) -> UOp|None:
-  # how output s lands in the caller's buffer t, or None if it must be copied into t
-  # materialize straight into t
-  if s.op is Ops.CONTIGUOUS: return t.after(t.store(s.src[0]))
-  # rebind output storage to t
-  if s.op in {Ops.BUFFER, Ops.UNSHARD} and s.has_buffer_identity(): return t
-  return None
-
-def transform_precompiled_call(c:UOp) -> UOp|None:
-  if not c.arg.precompile: return None
-  assert c.src[0].op is Ops.TUPLE, f"expected TUPLE body for precompiled FUNCTION, got {c.src[0].op}"
-  input_buffers = tuple(x.contiguous() if x.op not in {Ops.AFTER, Ops.BIND} else x for x in c.src[1:])
-
-  # add the outputs to the call
-  srcs = c.src[0].src
-  resolved = [c.gettuple(i) for i in range(len(srcs))]
-  outs = tuple(r.empty_like() for r in resolved)
-  targets = [o.param_like(len(c.src)-1+i).shrink_to(s.shape) for i,(o,s) in enumerate(zip(outs, srcs))]
-
-  subs:dict[UOp, UOp] = {}
-  items:list[UOp] = []
-  for s, t in zip(srcs, targets):
-    after_deps:list[UOp] = []
-    while s.op is Ops.AFTER:
-      after_deps.extend(s.src[1:])
-      s = s.src[0]
-    if (placed := _precompiled_output_redirect(s, t)) is not None and s not in subs:
-      subs[s] = placed
-      items.append(s.after(*after_deps) if after_deps else s)
-    else:
-      items.append(t.after(t.store(s.after(*after_deps))))
-  fxn = UOp.sink(*(x.substitute(subs) for x in items))
-
-  # body switches from TUPLE to SINK, so the node becomes an opaque CALL (not FUNCTION)
-  new_call = UOp(Ops.CALL, src=(fxn, *input_buffers, *outs), arg=c.arg)
-  rets = tuple(o.after(new_call) for o in outs)
-
-  # if the CALL has symbolic shapes, shrink the max-sized output to the actual symbolic shape
-  # NOTE: must use resolved shapes from the FUNCTION (which substitutes PARAMs with external args), not raw body shapes
-  rets = tuple(r.shrink_to(rs.shape) for r,rs in zip(rets, resolved))
-
-  return UOp.maketuple(*rets)
-
-# NOTE: adding rules to here is bad. these all need to run before the schedule cache
-pm_early_transform_tensor_graph = PatternMatcher([
-  # transform precompiled FUNCTIONs into CALLs (body becomes SINK with stores)
-  (UPat(Ops.FUNCTION, name="c"), transform_precompiled_call),
-
-  # resolve TUPLE+GETTUPLE (for precompiled calls)
-  (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]),
-
-  # fold MOPS+BITCAST over BUFFER/SLICE into SLICE when movement ops collapse to contiguous range
-  (UPat((Ops.BITCAST, Ops.COPY, Ops.CONTIGUOUS), src=(UPat(GroupOp.Movement|{Ops.BUFFER}, name="src"),), name="c"), contiguous_mops_to_view),
-
-  # remove contiguous on movement ops before a copy on disk
-  (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.CONTIGUOUS).f(Ops.COPY, name="copy"), lambda x,copy:
-   copy.replace(src=(x,), tag=None) if isinstance(x.device, str) and x.device.startswith("DISK") else None),
-  # push copy past movement ops to disk
-  (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.COPY, name="copy"), lambda x,copy:
-   x.replace(src=(copy.replace(src=(x.src[0],), tag=None),)+x.src[1:]) \
-   if isinstance(x.device, str) and x.device.startswith("DISK") else None),
-
-  # add CONTIGUOUS to tagged UOps
-  (UPat(GroupOp.All-{Ops.CONTIGUOUS, Ops.AFTER, Ops.STORE}, name="x"),
-   lambda x: None if x.tag is None else x.rtag(None).contiguous(tag=x.tag) if x.tag else x.replace(tag=None)),
-  # remove extra CONTIGUOUS on AFTER (only when target is contiguous)
-  (UPat(Ops.CONTIGUOUS, src=(UPat(Ops.AFTER, name="a"),), name="c"),
-   lambda a,c: a.replace(tag=(a.tag or ())+(c.tag or ())) if a.src[0].has_buffer_identity() else None),
-  # replace AFTER+STORE with CONTIGUOUS when target is not a buffer
-  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(), UPat(name="src")))), name="u"), replace_store_after_with_contig),
-  # replace CONTIGUOUS with STORE+AFTER
-  (UPat(Ops.CONTIGUOUS, name="u"), replace_contig_with_store_after),
-  # remove DETACH/CONTIGUOUS_BACKWARD (allows more contiguous removal)
-  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),
-])
-
-def finalize_after(ctx:AllocCtx, x:UOp):
-  # untagged: record as an assign for the call body
-  if x.tag is None:
-    ctx.assigns.append(x)
-    return None
-  # tagged: untag and map each original pre-rewrite UOp to the stripped buffer; the untagged result is reprocessed as untagged
-  ret = x.replace(tag=None)
-  replace_uop = ret
-  while replace_uop.op is Ops.AFTER: replace_uop = replace_uop.src[0]
-  for t in x.tag:
-    original_uop: UOp = ctx.uop_list[t]
-    ctx.buffer_map[original_uop] = replace_uop.shrink_to(original_uop.shape)
-  return ret
-
-def replace_input_buffer(ctx:AllocCtx, b:UOp):
-  ctx.replacements.append(b)
-  if b.op is Ops.BIND: return b.param_like(len(ctx.replacements)-1)
-  return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b.device,
-                   addrspace=b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL)
-
-pm_finalize_call = PatternMatcher([
-  (UPat(Ops.AFTER, name="x"), finalize_after),
-  (UPat(Ops.COPY, name="x"), lambda ctx,x: ctx.assigns.append(x) if isinstance(x.device, str) and x.device.startswith(("DISK", "TINYFS")) else None),
-])
-
-pm_replace_buf = PatternMatcher([
-  # replace BUFFER with PARAM for cache key normalization
-  (UPat(Ops.BUFFER, src=(UPat(),), name="b"), lambda ctx,b:
-   replace_input_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None),
-  # replace SLICE with PARAM. this rewrite is bottom up so BUFFERs we don't need won't be in the input
-  (UPat(Ops.SLICE, src=(UPat(Ops.BUFFER), UPat(Ops.CONST, dtype=dtypes.weakint)), name="b"), replace_input_buffer),
-  # strip value from BIND for cache key normalization, so different values hit same cache
-  (UPat(Ops.BIND, src=(UPat(Ops.PARAM), UPat(Ops.CONST)), name="b"), replace_input_buffer),
-])
-
-@track_rewrites(lambda _,ret: f"Callify {pluralize('Buffer', len(ret[1]))}")
-def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
-  if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Tensor Graph")
-  # uop list is a list in the original_sink graph and we can map to the tags later
-  # same predicate as Tensor.realize
-  ctx = AllocCtx(bases={base for x in big_sink.src if not (base:=x.base).is_virtual and not base.has_buffer_identity()
-                        and base.op is not Ops.AFTER and base.addrspace is not AddrSpace.ALU})
-
-  # this rewrite is "read-only", it adds simple things to buffer_map and may sink things on big_sink, bottom_up
-  # this is the only one where we have to be careful to not break the tensor graph
-  big_sink = graph_rewrite(big_sink, add_tags, ctx=ctx, bottom_up=True, name="number the uops")
-
-  # here we can break the tensor graph. this is the only place you need to maintain numbered tags
-  big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, name="early transform tensor graph")
-
-  # here we construct the final buffer_map: as-built nodes -> their final storage. values are never keys
-  graph_rewrite(big_sink, pm_finalize_call, ctx=ctx, name="finalize call")
-  ret = graph_rewrite(UOp.sink(*ctx.assigns), pm_replace_buf, ctx=ctx, bottom_up=True, name="replace bufs").call(*ctx.replacements)
-  assert not any(x in ctx.buffer_map for x in ctx.buffer_map.values())
-  if VIZ: graph_rewrite(ret, PatternMatcher([]), name="View Call")
-  return ret, ctx.buffer_map
@@ -2,8 +2,8 @@ from dataclasses import replace, dataclass
 import itertools, functools
 from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC
 from tinygrad.helpers import ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT, TracingKey, Context, panic
-from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, GroupOp
-from tinygrad.uop.ops import AxisType, pm_commit_weak, pm_cast_weak
+from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, Ops, UPat, rewrite_group, KernelInfo, ProgramInfo, GroupOp, AxisType
+from tinygrad.uop.weak import pm_lower_index_dtype, pm_commit_weak, pm_cast_weak
 from tinygrad.uop.render import pyrender
 from tinygrad.uop.spec import type_verify, spec_tensor, spec_program
 from tinygrad.renderer import Renderer, Estimates
@@ -12,7 +12,7 @@ from tinygrad.dtype import dtypes, AddrSpace

 # import all pattern matchers here
 from tinygrad.codegen.gpudims import pm_add_gpudims
-from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid
+from tinygrad.uop.symbolic import sym, symbolic_simple, symbolic, pm_fold_cast_const, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid
 from tinygrad.uop.movement import mop_cleanup
 from tinygrad.codegen.decomp.dtype import pm_dtype_decomps
 from tinygrad.codegen.decomp.op import get_late_rewrite_patterns, get_simplifying_rewrite_patterns
@@ -20,7 +20,7 @@ from tinygrad.codegen.decomp.transcendental import get_transcendental_patterns
 from tinygrad.codegen.late.coalesce import indexing_simplify
 from tinygrad.codegen.opt.postrange import apply_opts
 from tinygrad.codegen.late.gater import pm_move_gates_from_index
-from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse
+from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse, pm_reduce_unparented
 from tinygrad.schedule.multi import multi_pm
 from tinygrad.schedule.rangeify import pm_mops
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
@@ -310,7 +310,8 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
    sink = apply_opts(sink, ren, beam=ast.arg.beam)

  # ** expander (expand_rewrite) **
-  sink = graph_rewrite(sink, sym+pm_move_where_on_load+pm_flatten_range, name="postopt symbolic")
+  # reduce_unparented: a REDUCE whose src folded to a CONST (e.g. x*0) has no parented ranges, collapse it before the expander
+  sink = graph_rewrite(sink, sym+pm_move_where_on_load+pm_flatten_range+pm_reduce_unparented, name="postopt symbolic")

  # expand
  sink = graph_rewrite(sink, expander2, ctx=build_range_map(sink), name="expander")
@@ -336,14 +337,16 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:

  # do memory coalescing (late)
  sink = memory_coalescing(sink, ren)
-  sink = graph_rewrite(sink, symbolic_simple+ew_devectorizer+pm_simplify_add_image, name="add images", ctx=({}, ren), bottom_up=True)
+  sink = graph_rewrite(sink, symbolic_simple+ew_devectorizer+pm_simplify_add_image,
+                       name="add images", ctx=({}, ren), bottom_up=True)

  # extra symbolic before decomp. crashes without this?
-  sink = graph_rewrite(sink, sym, name="extra symbolic")
+  # NOTE: also run indexing_simplify here, while the index is still weakint and (x+y)*c -> x*c+y*c applies
+  sink = graph_rewrite(sink, sym+indexing_simplify, name="extra symbolic")

  # lower index dtype
  # NOTE: we need indexing_simplify to remove the cast to long using the Invalid
-  sink = graph_rewrite(sink, pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes")
+  sink = graph_rewrite(sink, symbolic_simple+pm_fold_cast_const+pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes")

  # final symbolic before decomp
  sink = graph_rewrite(sink, symbolic, name="final symbolic")
@@ -448,7 +451,7 @@ pm_to_program = PatternMatcher([
  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.LINEAR), UPat(Ops.SOURCE, name="source")), name="prg"), do_compile),
 ])

-@track_rewrites(name=lambda ast,renderer,ret,**kwargs: TracingKey(ret.src[0].arg.name,(ret.src[0].arg.function_name, ast), ret=renderer), replay=True)
+@rewrite_group(name=lambda ast,renderer,ret,**kwargs: TracingKey(ret.src[0].arg.name,(ret.src[0].arg.function_name, ast), ret=renderer), replay=True)
@Context(ALLOW_DEVICE_USAGE=0)
 def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
  """
@@ -78,9 +78,11 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
    case Ops.MAX: return l2i(Ops.WHERE, dt, l2i(Ops.CMPLT, dt, *uops), b0, b1, a0, a1)
    case _: raise NotImplementedError(f"long decomposition of {op} unsupported")

-def split_l2i(op: Ops, dt: DType, *uops:UOp):
-  # l2i does arithmetic on its inputs; rules enter here to split them to 32-bit words first, l2i recurses on itself
-  return l2i(op, dt, *graph_rewrite(UOp.sink(*uops), pm_long_decomp, bottom_up=True).src)
+def split_l2i(ctx:dict, op: Ops, dt: DType, *uops:UOp):
+  # l2i does arithmetic on its inputs; rules enter here to split them to 32-bit words first, l2i recurses on itself.
+  # both word halves of a node ask for the same split, so ctx memos it for the pass
+  if (key:=(op, dt, uops)) not in ctx: ctx[key] = l2i(op, dt, *graph_rewrite(UOp.sink(*uops), pm_long_decomp, ctx=ctx, bottom_up=True).src)
+  return ctx[key]

 # ***** floats *****
 f2f_dt = { f:getattr(dtypes, f"uint{f.bitsize}") for f in dtypes.floats }
@@ -97,7 +99,8 @@ def f2f(v, fr:DType, to:DType, sat=True):
    if fr in dtypes.fp8_fnuz:
      fnuz_nan = sign.ne(0) & nosign.eq(0)
      qnan = shl(shl(1, te) - 1, tm) | shl(1, tm - 1)
-      return fnuz_nan.where(qnan, sign | exp.eq(0).where(0, norm)).bitcast(to)
+      # the fnuz bias can exceed the target's: exp in [1, fb-tb] is normal in fr but lands below to's normal range, so it flushes like a denormal
+      return fnuz_nan.where(qnan, sign | (exp < max(fb - tb, 0) + 1).where(0, norm)).bitcast(to)
    # fp8e4m3 has only one nan
    is_nan = (nosign.eq(shl(1, fm + fe) - 1) if fr == dtypes.fp8e4m3 else exp.eq(shl(1, fe) - 1))
    return (sign | exp.eq(0).where(0, is_nan.where(nan, norm))).bitcast(to)
@@ -139,21 +142,21 @@ pm_long_decomp = PatternMatcher([
  (UPat(Ops.STORE, src=(UPat.var('idx', tuple(l2i_dt.keys())), UPat.var('val')), name='st'), lambda st,idx,val:
   st.replace(src=(idx.rtag((0, dt:=l2i_dt[idx.dtype])), val.rtag((0, dt)))).group(
     st.replace(src=(idx.rtag((1, dt)), val.rtag((1, dt))))) if val.tag is None else None),
-  (UPat(GroupOp.Comparison, src=[UPat.var('a', tuple(l2i_dt.keys())), UPat()], name="x"), lambda a,x:
-   split_l2i(x.op, dt:=l2i_dt[a.dtype], *flatten((s.rtag((0, dt)), s.rtag((1, dt))) for s in x.src))),
-  (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda a,x:
-   split_l2i(Ops.BITCAST, l2i_dt[x.dtype], a.rtag((0, dt:=l2i_dt[a.dtype])), a.rtag((1, dt)))[x.tag[0]]),
-  (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a'),), name="x"), lambda a,x:
-   split_l2i(x.op, x.dtype, a)[x.tag[0]] if x.tag is not None else None),
-  (UPat(Ops.CAST, src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda a,x:
-   split_l2i(x.op, x.dtype, a.rtag((0, dt:=l2i_dt[a.dtype])), a.rtag((1, dt))) if x.dtype not in l2i_dt and a.tag is None else None),
-  (UPat((Ops.SHL, Ops.SHR), tuple(l2i_dt.keys()), src=(UPat.var('a'), UPat.var('b')), name="x"), lambda a,b,x:
-   split_l2i(x.op, dt:=l2i_dt[x.dtype], a.rtag((0, dt)), a.rtag((1, dt)), b.rtag((0, dt)))[x.tag[0]] if x.tag is not None else None),
-  (UPat(Ops.WHERE, tuple(l2i_dt.keys()), src=(UPat.var('c'), UPat.var('a'), UPat.var('b')), name="x"), lambda a,b,c,x:
-   split_l2i(x.op, dt:=l2i_dt[x.dtype], c, a.rtag((0, dt)), a.rtag((1, dt)), b.rtag((0, dt)), b.rtag((1, dt)))[x.tag[0]]
+  (UPat(GroupOp.Comparison, src=[UPat.var('a', tuple(l2i_dt.keys())), UPat()], name="x"), lambda ctx,a,x:
+   split_l2i(ctx, x.op, dt:=l2i_dt[a.dtype], *flatten((s.rtag((0, dt)), s.rtag((1, dt))) for s in x.src))),
+  (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda ctx,a,x:
+   split_l2i(ctx, Ops.BITCAST, l2i_dt[x.dtype], a.rtag((0, dt:=l2i_dt[a.dtype])), a.rtag((1, dt)))[x.tag[0]]),
+  (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a'),), name="x"), lambda ctx,a,x:
+   split_l2i(ctx, x.op, x.dtype, a)[x.tag[0]] if x.tag is not None else None),
+  (UPat(Ops.CAST, src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda ctx,a,x:
+   split_l2i(ctx, x.op, x.dtype, a.rtag((0, dt:=l2i_dt[a.dtype])), a.rtag((1, dt))) if x.dtype not in l2i_dt and a.tag is None else None),
+  (UPat((Ops.SHL, Ops.SHR), tuple(l2i_dt.keys()), src=(UPat.var('a'), UPat.var('b')), name="x"), lambda ctx,a,b,x:
+   split_l2i(ctx, x.op, dt:=l2i_dt[x.dtype], a.rtag((0, dt)), a.rtag((1, dt)), b.rtag((0, dt)))[x.tag[0]] if x.tag is not None else None),
+  (UPat(Ops.WHERE, tuple(l2i_dt.keys()), src=(UPat.var('c'), UPat.var('a'), UPat.var('b')), name="x"), lambda ctx,a,b,c,x:
+   split_l2i(ctx, x.op, dt:=l2i_dt[x.dtype], c, a.rtag((0, dt)), a.rtag((1, dt)), b.rtag((0, dt)), b.rtag((1, dt)))[x.tag[0]]
   if x.tag is not None else None),
-  (UPat((*(GroupOp.ALU - GroupOp.Comparison - {Ops.SHL, Ops.SHR, Ops.WHERE}), Ops.BITCAST), tuple(l2i_dt.keys()), name="x"), lambda x:
-   split_l2i(x.op, l2i_dt[x.dtype], *flatten((a.rtag((0, l2i_dt[x.dtype])), a.rtag((1, l2i_dt[x.dtype]))) for a in x.src))[x.tag[0]]
+  (UPat((*(GroupOp.ALU - GroupOp.Comparison - {Ops.SHL, Ops.SHR, Ops.WHERE}), Ops.BITCAST), tuple(l2i_dt.keys()), name="x"), lambda ctx,x:
+   split_l2i(ctx, x.op, l2i_dt[x.dtype], *flatten((a.rtag((0, l2i_dt[x.dtype])), a.rtag((1, l2i_dt[x.dtype]))) for a in x.src))[x.tag[0]]
   if x.tag is not None else None),
  (UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda x,idx:
   x.replace(dtype=l2i_dt[x.dtype], src=(reindex(idx, x.tag[0]).replace(dtype=l2i_dt[x.dtype], tag=None),), tag=None) if x.tag is not None else None),
@@ -197,7 +200,7 @@ def do_dtype_decomps(sink:UOp, ctx:tuple[set[DType], Renderer]) -> UOp:
      to = dtypes.int if fr == dtypes.long else dtypes.half if not _should_emulate(dtypes.half) and fr in dtypes.fp8s else dtypes.float
      if DEBUG >= 2: print(f"emulating {fr} as {to}")
      pm = pm_float_decomp if fr in dtypes.floats else pm_long_decomp
-      sink = graph_rewrite(sink, pm, name=f"decomp {fr} -> {to}", ctx=(fr, to), bottom_up=True)
+      sink = graph_rewrite(sink, pm, name=f"decomp {fr} -> {to}", ctx={} if pm is pm_long_decomp else (fr, to), bottom_up=True)
  ctx[0].clear()
  return sink

@@ -57,7 +57,7 @@ def add_gpudims(ctx:Renderer, s:UOp):

  # get the idxs
  ki: KernelInfo = s.arg
-  if ctx.has_threads: idxs = [UOp.variable("core_id", 0, int(global_shape[0])-1, dtypes.int).cast(dtypes.weakint)]
+  if ctx.has_threads: idxs = [UOp.variable("core_id", 0, int(global_shape[0])-1, dtypes.int, param=True).cast(dtypes.weakint)]
  elif ki.dont_use_locals:
    assert not local_dims, "can't use locals if there's no local dims"
    idxs = get_grouped_dims("idx", global_shape, ctx.global_max, reverse=True)
@@ -89,7 +89,7 @@ def add_gpudims(ctx:Renderer, s:UOp):

 pm_device_to_var = PatternMatcher([
  # the DEVICE axis is not a program axis, it's bound per device at launch. lower it to the _device_num variable (like SPECIAL for devices)
-  (UPat(Ops.RANGE, name="r"), lambda r: UOp.variable("_device_num", 0, r.vmax, dtype=r.dtype) if r.arg[-1] is AxisType.DEVICE else None),
+  (UPat(Ops.RANGE, name="r"), lambda r: UOp.variable("_device_num", 0, r.vmax, dtype=r.dtype, param=True) if r.arg[-1] is AxisType.DEVICE else None),
  # ENDs that closed a DEVICE range no longer close it
  (UPat(Ops.END, name="e"), lambda e: e.replace(src=(e.src[0],)+tuple(s for s in e.src[1:] if s.op is not Ops.PARAM))
   if any(s.op is Ops.PARAM and s.arg.name == '_device_num' for s in e.src[1:]) else None),
@@ -1,8 +1,8 @@
 import itertools, functools
 from collections import defaultdict
 from tinygrad.dtype import dtypes, AddrSpace, Invalid, DType
-from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp, shape_to_shape_arg
-from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate
+from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp, shape_to_shape_arg, graph_rewrite
+from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate, sym
 from tinygrad.helpers import getenv, IMAGE, OSX, ceildiv, is_image_shape
 from tinygrad.renderer import Renderer

@@ -26,12 +26,15 @@ def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]:
    # check if idx is out of bound when X is on the wrong side of the bound: X in [c+1, vmax] or [vmin, c-1]
    lo, hi = (c + 1, X.vmax) if is_upper_bound else (X.vmin, c - 1)
    if lo <= hi:
-      fake = UOp.variable(f"fake{i}", lo, hi, X.dtype)
-      for coord,b in zip(idx.src, (width, height)):
-        rw = coord.substitute({X:fake}).simplify()
-        if rw.vmin >= b or rw.vmax < 0:
-          drop_stmt.append(stmt)
-          break
+      fake = UOp.variable(f"fake{i}", lo, hi, X.dtype, param=True)
+      subs = [{X: fake}]
+      # idx may not have X itself, so also substitute a term of X: v -> fake - (X - v)
+      terms = list(X.split_uop(Ops.ADD))
+      v = next((u for u in terms if u.op in GroupOp.Irreducible and u.op is not Ops.CONST), None)
+      if v is not None and (rest:=[u for u in terms if u is not v]): subs.append({v: fake - UOp.usum(*rest)})
+      if any((testidx:=graph_rewrite(coord.substitute(sub), sym)).vmin >= b or testidx.vmax < 0
+             for sub in subs for coord,b in zip(idx.src, (width, height))):
+        drop_stmt.append(stmt)
  return drop_stmt

 def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
@@ -3,7 +3,7 @@ from tinygrad.uop.ops import PatternMatcher, UPat, Ops
 from tinygrad.dtype import Invalid, dtypes

 def move_where_load(gate, l, a, w):
-  return l.replace(src=(l.src[0], l.vconst_like(0) if a.is_invalid else
+  return l.replace(src=(l.src[0], l.vconst_like(0) if a.is_invalid else l.const_like(a.val) if a.op is Ops.CONST else
    a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype), l.src[2])).cast(w.dtype)

 pm_move_gates_from_index = PatternMatcher([
@@ -332,9 +332,9 @@ class Scheduler:
  @property
  def group_for_reduces(self) -> int: return len(self.axes_of(AxisType.GROUP_REDUCE))

-def bufs_from_ast(ast:UOp, dname:str) -> list[Buffer]:
+def args_from_ast(ast:UOp, dname:str) -> tuple[list[Buffer], dict[str, int]]:
  glbls = sorted([x for x in ast.backward_slice if x.op is Ops.PARAM and x.arg.slot >= 0], key=lambda x: x.arg.slot)
-  return [Buffer(dname, x.max_numel(), x.dtype) for x in glbls]
+  return [Buffer(dname, x.max_numel(), x.dtype) for x in glbls], {k.expr:int(k.vmax+k.vmin)//2 for k in ast.variables()}

 def apply_opts(ast:UOp, ren:Renderer, beam:int=0) -> UOp:
  if ast.tag is not None: return ast
@@ -344,10 +344,10 @@ def apply_opts(ast:UOp, ren:Renderer, beam:int=0) -> UOp:
    for opt in ast.arg.opts_to_apply: k.apply_opt(opt)
  elif beam >= 1:
    from tinygrad.codegen.opt.search import beam_search
-    rawbufs = bufs_from_ast(ast, ren.target.device)
+    rawbufs, var_vals = args_from_ast(ast, ren.target.device)
    # beam search may open devices
    with Context(ALLOW_DEVICE_USAGE=1):
-      k = beam_search(k, rawbufs, beam, bool(getenv("BEAM_ESTIMATE", 1)))
+      k = beam_search(k, rawbufs, var_vals, beam, bool(getenv("BEAM_ESTIMATE", 1)))
  elif not NOOPT and (ast.arg is None or ast.arg.applied_opts == ()):
    from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
    # NOTE: hand_coded_optimizations doesn't support multiblock opts yet
@@ -1,6 +1,6 @@
 import math, time, multiprocessing, traceback, signal, atexit
 from dataclasses import replace
-from tinygrad.uop.ops import sym_infer, AxisType, UOp
+from tinygrad.uop.ops import sym_infer, AxisType, UOp, Ops
 from tinygrad.uop.render import pyrender
 from tinygrad.device import Device, Buffer
 from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str
@@ -62,7 +62,8 @@ def _try_compile(x:tuple[int,Scheduler]) -> tuple[int, tuple[UOp, float]|None]:
  ret = None
  try:
    st = time.perf_counter()
-    prg = to_program(x[1].copy().get_optimized_ast(name_override="test"), x[1].ren)
+    ast, dev = x[1].copy().get_optimized_ast(name_override="test"), x[1].ren.target.device
+    prg = to_program(ast.substitute({p: p.replace(arg=replace(p.arg, device=dev)) for p in ast.toposort() if p.op is Ops.PARAM}), x[1].ren)
    et = time.perf_counter() - st
    uops = prg.src[1].src
    if len(uops) >= (uops_max:=getenv("BEAM_UOPS_MAX", 3000)) > 0:
@@ -111,7 +112,7 @@ def get_kernel_actions(s:Scheduler, include_0=True, max_up:int|None=None) -> dic
  return acted

 beam_pool, BEAM_DEBUG = None, getenv("BEAM_DEBUG")
-def beam_search(s:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True, disable_cache=IGNORE_BEAM_CACHE.value):
+def beam_search(s:Scheduler, rawbufs:list[Buffer], var_vals:dict[str,int], amt:int, allow_test_size=True, disable_cache=IGNORE_BEAM_CACHE.value):
  global beam_pool
  key = {"ast": s.ast.key, "amt": amt, "allow_test_size": allow_test_size, "device": s.ren.target.device, "suffix": s.ren.suffix}
  if not disable_cache and CACHELEVEL >= 1 and (val:=diskcache_get("beam_search", key)) is not None:
@@ -136,7 +137,6 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=True

  try:
    rawbufs = _ensure_buffer_alloc(rawbufs)
-    var_vals: dict[str, int] = {k.expr:int(k.vmax+k.vmin)//2 for k in s.ast.variables()}
    exiting, st = False, time.perf_counter()
    dev = Device[s.ren.target.device]
    while not exiting:
@@ -1,6 +1,7 @@
 import math, functools
 from dataclasses import dataclass
 from tinygrad.dtype import DType, dtypes
+from tinygrad.uop.ops import PatternMatcher, UOp, UPat, Ops

@dataclass(frozen=True)
 class TensorCore: # D = A * B + C, A is (M x K), B is (K x N), C and D are (M x N)
@@ -135,6 +136,42 @@ amd_cdna4 = amd_cdna_1616128 + amd_cdna_161632 + amd_cdna_161616

 def get_amd(arch): return {"gfx942": amd_cdna3, "gfx950": amd_cdna4, "gfx1200": amd_rdna4, "gfx1201": amd_rdna4}.get(arch, amd_rdna3)

+pm_validate_wmma_rdna3 = PatternMatcher([
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.int32), lambda x: x.replace(
+    src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
+    if x.src[0].dtype == dtypes.int8 and x.src[0].max_numel() == 16 else None),
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(x.replace(
+      src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(UOp.const(j//2, dtypes.int16))
+      if j%2 == 0 else UOp.const(0.0, x.src[2].dtype)
+      for j in range(x.max_numel()*2)))),
+      arg=(*x.arg[:4], None)).index(UOp.const(i*2, dtypes.int16))
+      for i in range(x.max_numel()))) if x.max_numel() == 8 else None),
+  (UPat(Ops.WMMA, name="x"), lambda x: x.replace(
+    src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
+    if x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 16 else None),
+])
+
+pm_validate_wmma_rdna4 = PatternMatcher([
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16), lambda x: x.replace(
+    dtype=dtypes.uint16,
+    src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2].bitcast(dtypes.uint16)))
+      .bitcast(dtypes.bfloat16) if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None),
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+    lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
+    if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None)
+])
+
+pm_validate_wmma_cdna = PatternMatcher([
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+    lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
+    if x.arg[0][2] == 128 and x.src[0].dtype.itemsize <= 8 else None),
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+    lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
+    if x.max_numel() == 4 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 4 else None),
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+    lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2]))
+    if x.max_numel() == 4 and x.src[0].dtype in dtypes.fp8_ocp and x.src[0].max_numel() == 8 else None),
+])
 # ***** Apple Metal *****

 metal = [TensorCore(dims=(8,8,8), threads=32, elements_per_thread=(2,2,2), dtype_in=di, dtype_out=do,
@@ -137,7 +137,7 @@ def reduce_collapse(red:UOp, u:UOp, pm:PatternMatcher=pm_reduce_collapse) -> UOp
    for u in included:
      for s in u.src:
        if s in included or s in replaces or s.op in {Ops.CONST, Ops.PARAM, Ops.BUFFER}: continue
-        replaces[s] = UOp.variable(f'in{len(replaces)}', s.vmin, s.vmax, s.dtype)
+        replaces[s] = UOp.variable(f'in{len(replaces)}', s.vmin, s.vmax, s.dtype, param=True)
    collapse_fxn = u.substitute(replaces).reduce(r, arg=Ops.ADD)
    sink = graph_rewrite(collapse_fxn, pm, name="reduce_collapse")
    if not no_range(sink): return None
@@ -336,6 +336,8 @@ class Compiled:
  pm_lower:Any = None
  pm_bufferize:Any = None

+  has_copy_queue:bool = True
+
  def __init__(self, device:str, allocator:Allocator, renderers:list[type[Renderer]], runtime:type[Program[Self]]|None, graph=None, arch=None):
    from tinygrad.renderer import Renderer
    self.device, self.allocator, self.runtime_t, self.graph, self.renderers = device, allocator, runtime, graph, renderers or [Renderer]
@@ -80,7 +80,7 @@ class DType(metaclass=DTypeMetaClass):
    # NOTE: float('nan') != float('nan'), so we canonicalize here
    if isinstance(val, float) and math.isnan(val): val = math.nan
    # int is the default. wrap floats in ConstFloat to distinguish -0.0 from 0.0 in cache
-    return ConstFloat(float(val)) if dtypes.is_float(self) else bool(val) if dtypes.is_bool(self) else int(val)
+    return ConstFloat(truncate.get(self, float)(float(val))) if dtypes.is_float(self) else bool(val) if dtypes.is_bool(self) else int(val)


 class DTypes:
@@ -4,7 +4,7 @@ from tinygrad.tensor import Tensor, all_tensors
 from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv, JIT, JIT_BATCH_SIZE, dedup, pluralize, VIZ, disable_gc
 from tinygrad.device import Buffer, Compiled, Device, MultiBuffer, DepsTracker
 from tinygrad.dtype import DType
-from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, track_rewrites, graph_rewrite
+from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, rewrite_group, graph_rewrite
 from tinygrad.renderer import Estimates
 from tinygrad.engine.realize import capturing, compile_linear, link_linear, run_linear, graph_cache, estimate_uop, get_runtime
 from tinygrad.engine.realize import unwrap_multi, resolve_params, get_call_arg_uops, get_call_outs_ins
@@ -44,9 +44,7 @@ def graph_split_rewrite(linear:UOp, max_batch_size:int=0) -> UOp:
    current_batch, current_batch_devs = [], []

  for si in linear.src:
-    if si.src[0].op is Ops.SLICE: continue
-
-    devs = dedup([Device[x] for b in si.src[1:] if b.op is not Ops.BIND for x in (b.device if isinstance(b.device, tuple) else (b.device,))])
+    devs = dedup([Device[x] for b in si.src[1:] if not b.is_bound_var for x in (b.device if isinstance(b.device, tuple) else (b.device,))])
    graph_t = graph_class(devs[0]) if devs[0].graph is not None else None

    can_graph = graph_t is not None and graph_t.supports_uop(devs, si)
@@ -64,7 +62,7 @@ def _copy_input(u:UOp) -> UOp:
  run_linear(UOp(Ops.LINEAR, src=(u.copy_to_device(u.device).call(new:=UOp.new_buffer(u.device, u.max_numel(), u.dtype), u),)))
  return new

-@track_rewrites(lambda linear,held_bufs,input_uops,ret=(): f"JIT {pluralize('call', len(linear.src))}")
+@rewrite_group(lambda linear,held_bufs,input_uops,ret=(): f"JIT {pluralize('call', len(linear.src))}")
 def jit_lower(linear:UOp, held_bufs:set[UOp], input_uops:list[UOp]) -> UOp:
  if VIZ: graph_rewrite(linear, PatternMatcher([]), name="View captured linear")

@@ -180,7 +178,7 @@ class CapturedJit(Generic[ReturnType]):
      if call.op is not Ops.CALL: continue
      arg_uops = get_call_arg_uops(call)
      outs, ins = get_call_outs_ins(call)
-      out |= {arg_uops[k] for k in set(outs) - set(ins) if arg_uops[k].op in (Ops.BUFFER, Ops.SLICE)}
+      out |= {b for k in set(outs) - set(ins) if (b:=u if (cv:=(u:=arg_uops[k]).contiguous_view()) is None else cv[0]).op is Ops.BUFFER}
    return out

  def __call__(self, input_uops:list[UOp], var_vals:dict[str, int]) -> ReturnType:
@@ -3,21 +3,21 @@ from typing import cast, Iterator, Any, Sequence
 import time, random, itertools, math, contextlib, weakref, array
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, prod, flatten, Context, getenv, to_tuple
-from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events
-from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, buffers, graph_rewrite
-from tinygrad.device import Device, Buffer, MultiBuffer
+from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, wait_cond
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite
+from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
 from tinygrad.renderer import Estimates
 from tinygrad.codegen import to_program
-from tinygrad.codegen.opt.postrange import bufs_from_ast
+from tinygrad.codegen.opt.postrange import args_from_ast

 # **************** Helpers ****************

-def get_call_arg_uops(call:UOp) -> tuple[UOp, ...]: return tuple(s for s in call.src[1:] if s.op is not Ops.BIND)
+def get_call_arg_uops(call:UOp) -> tuple[UOp, ...]: return tuple(s for s in call.src[1:] if not s.is_bound_var)

 def get_call_outs_ins(call:UOp) -> tuple[tuple[int, ...], tuple[int, ...]]:
  ast = call.src[0]
  if ast.op is Ops.PROGRAM: return tuple(ast.arg.outs), tuple(ast.arg.ins)
-  if ast.op in (Ops.COPY, Ops.SLICE): return (0,), (1,)
+  if ast.op is Ops.COPY: return (0,), (1,)
  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return (0,), tuple(range(1, len(get_call_arg_uops(call))))
  return (), ()

@@ -27,13 +27,10 @@ def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|N

  ast, arg_uops = call.src[0], get_call_arg_uops(call)
  if ast.op is Ops.PROGRAM: return ast.arg.name
-  if ast.op is Ops.SLICE:
-    offset = ast.src[1].val * arg_uops[1].dtype.itemsize
-    return colored(f"view {_uop_sz_to_str(arg_uops[0]):>10} @ {offset:<10d}", "yellow")
  if ast.op is Ops.COPY: return colored(f"copy {_uop_sz_to_str(arg_uops[0]):>10}, {_dev_str(bufs[0]):>7s} <- {_dev_str(bufs[1]):7s}", "yellow")
  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return colored(f"enc/dec {_uop_sz_to_str(arg_uops[0])}", "yellow")
  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return colored(f"batched {len(ast.src[0].src)}", "cyan")
-  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq": return call.arg.aux.name
+  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq": return cast(str, call.arg.name)
  raise NotImplementedError("get_call_name is not implemented")

 # **************** Stat ****************
@@ -90,12 +87,13 @@ def optimize_local_size(call:UOp, prg:UOp) -> UOp|None:

  if (local_size:=local_size_cache.get(prg.key)) is None:
    # reuse one loaded runtime across candidates, only launch dims vary
-    bufs, runtime = [b.allocate() for b in bufs_from_ast(prg.src[0], device)], get_runtime(device, prg, cache=False)
+    (bufs, var_vals), runtime = args_from_ast(prg.src[0], device), get_runtime(device, prg, cache=False)
+    bufs  = [b.allocate() for b in bufs]
    def try_exec(local_size):
      try:
        new_gs = tuple(g//l if g%l == 0 else g/l for g,l in zip(prg.arg.global_size, local_size))
        return runtime(*[bufs[i].get_buf(device) for i in prg.arg.globals], global_size=new_gs, local_size=(*local_size,),
-                       vals=prg.arg.vals({}), wait=True)
+                       vals=prg.arg.vals(var_vals), wait=True)
      except Exception: return float('inf')

    MAX_WORKGROUP = 1024
@@ -139,7 +137,7 @@ class ExecContext:
  cache: bool = True

 def _resolve(b:UOp, inputs:tuple[UOp, ...]) -> UOp:
-  if b.op in (Ops.SLICE, Ops.MSELECT) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg.slot], *b.src[1:]))
+  if b.op in (Ops.MSELECT, Ops.SHRINK) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg.slot], *b.src[1:]))
  if b.op is Ops.MSTACK: return b.replace(src=tuple(_resolve(x, inputs) for x in b.src))
  return inputs[b.arg.slot] if b.op is Ops.PARAM else b
 def resolve_params(call:UOp, inputs:tuple[UOp, ...]) -> list[UOp]: return [_resolve(b, inputs) for b in get_call_arg_uops(call)]
@@ -153,13 +151,6 @@ def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], d
                   for x in call.src[0].toposort())
    for j, per_dev in enumerate(zip(*[cast(MultiBuffer, b).bufs for b in bufs])): yield list(per_dev), {"_device_num": j} if has_dnum else {}

-def exec_view(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
-  resolved = resolve_params(call, ctx.input_uops)
-  bufs = [cast(Buffer, b.buffer) for b in resolved]
-  bv = bufs[1].view(resolved[0].max_numel(), ast.dtype, ast.src[1].val*bufs[1].dtype.itemsize)
-  with track_stats(ctx, call, bv.device, [bv, bufs[1]], ctx.var_vals): buffers[resolved[0]] = bv
-  return None
-
 def exec_copy(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
  for bufs, device_vars in unwrap_multi(call, resolve_params(call, ctx.input_uops)):
    dest, src = bufs[0].ensure_allocated(), bufs[1].ensure_allocated()
@@ -209,21 +200,30 @@ def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
  return t[0]

 def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
-  if (inputs:=call.arg.aux.inputs) is not None:
+  if (info:=call.arg.aux).inputs is not None:
    bufs = [_resolve(ctx.input_uops[i], ctx.input_uops).buffer for i in call.arg.aux.input_idxs]
-    table = call.src[1+inputs].buffer
+    table = call.src[1+info.inputs].buffer
    for j,dev in enumerate(call.arg.aux.device):
      addrs = array.array('Q', [(b.bufs[j] if isinstance(b, MultiBuffer) else b).get_buf(dev).va_addr for b in bufs])
-      buf = table.bufs[j] if isinstance(table, MultiBuffer) else table
-      buf.ensure_allocated()._buf.cpu_view().view(fmt='Q')[:len(addrs)] = addrs
+      mv = (table.bufs[j] if isinstance(table, MultiBuffer) else table).ensure_allocated()._buf.cpu_view().view(fmt='Q')
+      wait_cond(lambda: mv[0], value=0, timeout_ms=ctx.timeout or getenv("HCQDEV_WAIT_TIMEOUT_MS", 30000), msg=f"{dev} hang detected")
+      mv[:len(addrs)] = addrs

-  exec_kernel(replace(ctx, update_stats=False), call, ast)
+  exec_kernel(replace(ctx, update_stats=DEBUG>=3), call, ast)

-  st = time.perf_counter()
-  for d in call.arg.aux.device:
-    with track_stats(ctx, call, d, [], ctx.var_vals):
-      if ctx.wait: Device[d].synchronize()
-  return time.perf_counter() - st
+  tms = []
+  for devices,name,estimates,prof in info.kernels:
+    for device in devices:
+      tm = None
+      if prof:
+        (d:=cast(Any, Device[device])).prof_ents[prof[0]] = ProfileGraphEntry(device, name, *prof)
+        if ctx.wait:
+          d.synchronize(timeout=ctx.timeout)
+          st, en = (d.signal(x)._buf.cpu_view().view(fmt='Q')[0] for x in prof)
+          tms.append(tm:=float(en-st)/d.timestamp_divider/1e6)
+      stat_call = call.replace(arg=replace(call.arg, name=name, aux=replace(info, estimates=estimates, kernels=())))
+      with track_stats(ctx, stat_call, device, [], ctx.var_vals) as et: et[0] = tm
+  return max(tms) if tms else None

 # flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src
 pm_flatten_linear = PatternMatcher([
@@ -254,7 +254,6 @@ pm_optimize_local_size = PatternMatcher([
 ])

 pm_exec = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.SLICE, name="ast"),), name="call", allow_any_len=True), exec_view),
  (UPat(Ops.CALL, src=(UPat(Ops.COPY, name="ast"),), name="call", allow_any_len=True), exec_copy),
  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="ast"),), name="call", allow_any_len=True), exec_kernel),
  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="encdec", name="ast"),), name="call", allow_any_len=True), exec_encdec),
@@ -265,11 +264,11 @@ pm_exec = PatternMatcher([

 if getenv("HCQ2"): from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link # noqa: E402 # down here, hcq2 imports the helpers above

-def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None) -> UOp:
+def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, profile:bool|None=None) -> UOp:
  if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
  if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
  linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True)
-  if getenv("HCQ2"): linear = hcq_compile(linear, input_uops)
+  if getenv("HCQ2"): linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile)
  return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)

 def link_linear(linear:UOp, cache=True) -> UOp: return hcq_link(linear, cache=cache) if getenv("HCQ2") else linear
@@ -287,5 +286,5 @@ def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None
      from tinygrad.tensor import Tensor
      with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024, 1024).contiguous().realize(do_update_stats=False)
  ctx = ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False)
-  linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0), cache=ctx.cache)
+  linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0, profile=True), cache=ctx.cache)
  return max(pm_exec.rewrite(c, ctx) or 0.0 for c in linear.src)
@@ -1,4 +1,5 @@
 import functools, time
+from dataclasses import replace
 from typing import Generic, TypeVar, Callable, cast, overload
 from tinygrad.helpers import Context, dedup, getenv, DEBUG
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, PatternMatcher, UPat
@@ -12,7 +13,7 @@ def add_to_ctx(ctx, x:UOp):
  return ret

 pm_ctx = PatternMatcher([
-  (UPat((Ops.BUFFER, Ops.BIND), name="x"), add_to_ctx),
+  (UPat(Ops.BUFFER, name="x"), add_to_ctx),
  (UPat((Ops.AFTER, Ops.CONTIGUOUS), name="x"),
   lambda ctx,x: add_to_ctx(ctx,x) if not x.op_in_backward_slice_with_self(Ops.PARAM) and x.op_in_backward_slice_with_self(Ops.BUFFER) else None),
 ])
@@ -23,6 +24,10 @@ def invalid_outputs(uret:UOp) -> set[UOp]:
  return {u.src[0].buf_uop for u in uret.backward_slice_with_self
          if u.op is Ops.STORE and u.src[1].base.is_invalid and not u.src[0].buf_uop.is_realized}

+def renumber_invalid_outputs(uret:UOp) -> UOp:
+  return uret.substitute({b:b.replace(arg=replace(b.arg, slot=i))
+                          for i,b in enumerate(x for x in uret.toposort(enter_calls=False) if x in invalid_outputs(uret))})
+
 ReturnType = TypeVar('ReturnType')
 class _function(Generic[ReturnType]):
  depth = 0
@@ -65,6 +70,7 @@ class _function(Generic[ReturnType]):
    # the BUFFERs that are left are the implicit inputs
    num_explicit = len(call_uops)
    uret = graph_rewrite(uret, pm_ctx, (call_uops, invalid_outputs(uret)), bottom_up=True, name="get_implicit_inputs")
+    uret = renumber_invalid_outputs(uret)
    name = getattr(self.fxn, '__qualname__', None) or type(self.fxn).__qualname__
    if not self.allow_implicit:
      implicit_buffers = [x for x in call_uops[num_explicit:] if x.op is Ops.BUFFER]
@@ -1,9 +1,9 @@
 from __future__ import annotations
 import time
 START_TIME = time.perf_counter()
-import os, functools, platform, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc
+import os, functools, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc
 from collections import defaultdict
-import subprocess, shutil, math, types, copyreg, inspect, importlib, decimal, itertools, difflib
+import shutil, math, types, copyreg, inspect, importlib, decimal, itertools, difflib
 from dataclasses import dataclass, field, replace
 from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator, cast, overload

@@ -13,8 +13,7 @@ U = TypeVar("U")
 def prod(x:Iterable[T]) -> T|int: return functools.reduce(operator.mul, x, 1)

 # NOTE: helpers is not allowed to import from anything else in tinygrad
-OSX, WIN = platform.system() == "Darwin", sys.platform == "win32"
-ARCH_X86 = any(x in platform.processor() for x in ("Intel", "i386", "x86_64"))
+OSX, WIN = sys.platform == "darwin", sys.platform == "win32"
 BASEDIR = pathlib.Path(__file__).parent

 # fix colors on Windows, https://stackoverflow.com/questions/12492810/python-how-can-i-make-the-ansi-escape-codes-to-work-also-in-windows
@@ -231,7 +230,7 @@ class _DEV(ContextVar):

 DEV, DEBUG, BEAM, NOOPT = _DEV("DEV", ""), ContextVar("DEBUG", 0), ContextVar("BEAM", 0), ContextVar("NOOPT", 0)
 IMAGE, FLOAT16, OPENPILOT_HACKS = ContextVar("IMAGE", 0), ContextVar("FLOAT16", 0), ContextVar("OPENPILOT_HACKS", 0)
-JIT, JIT_BATCH_SIZE = ContextVar("JIT", 2 if OSX and ARCH_X86 else 1), ContextVar("JIT_BATCH_SIZE", 32)
+JIT, JIT_BATCH_SIZE = ContextVar("JIT", 1), ContextVar("JIT_BATCH_SIZE", 32)
 CHUNK_SIZE = 2**20  # TinyFS content-addressed store: blob chunk + hash-tree node granularity
 WINO, CAPTURING, TRACEMETA, NO_COLOR = ContextVar("WINO", 0), ContextVar("CAPTURING", 1), ContextVar("TRACEMETA", 1), ContextVar("NO_COLOR", 0)
 TRAINING = ContextVar("TRAINING", 0)
@@ -251,7 +250,7 @@ DEFAULT_FLOAT, DEFAULT_INT = ContextVar("DEFAULT_FLOAT", "float32"), ContextVar(
 CAPTURE_PROCESS_REPLAY = ContextVar("CAPTURE_PROCESS_REPLAY", 0)
 def _get_cpu_count() -> int:
  # os.process_cpu_count (3.13+) respects cgroup limits
-  if hasattr(os, "process_cpu_count"): return max(1, os.process_cpu_count())
+  if hasattr(os, "process_cpu_count"): return max(1, os.process_cpu_count() or 1)
  # cgroup v2 (containers with --cpus=N)
  try:
    with open("/sys/fs/cgroup/cpu.max") as f:
@@ -454,9 +453,9 @@ def _ensure_downloads_dir() -> pathlib.Path:
  if pathlib.Path("/etc/tinybox-release").is_file():
    # try creating dir with sudo
    if not (downloads_dir := pathlib.Path("/raid/downloads")).exists():
-      subprocess.run(["sudo", "mkdir", "-p", downloads_dir], check=True)
-      subprocess.run(["sudo", "chown", "tiny:root", downloads_dir], check=True)
-      subprocess.run(["sudo", "chmod", "775", downloads_dir], check=True)
+      system(f"sudo mkdir -p {downloads_dir}")
+      system(f"sudo chown tiny:root {downloads_dir}")
+      system(f"sudo chmod 775 {downloads_dir}")
    return downloads_dir
  return pathlib.Path(cache_dir) / "downloads"

@@ -497,6 +496,7 @@ def fetch_fw(path:str, name:str, sha256:str) -> bytes:
 # *** Exec helpers

 def system(cmd:str, **kwargs) -> str:
+  import subprocess
  st = time.perf_counter()
  try: ret = subprocess.check_output(cmd.split(), stderr=subprocess.STDOUT, **kwargs).decode().strip()
  except subprocess.CalledProcessError as e:
@@ -21,7 +21,7 @@
    const d = document.createElement('div'); d.className = 'msg'; chat.appendChild(d);
    const r = await fetch('/v1/chat/completions', {method: 'POST', headers: {'Content-Type': 'application/json'},
      body: JSON.stringify({model: 'llama', messages: msgs, stream: true, temperature: 0.7})});
-    let buf = '';
+    let buf = '', txt = '', rsn = '';
    for (const rd = r.body.getReader(), dec = new TextDecoder();;) {
      const {done, value} = await rd.read();
      if (done) break;
@@ -30,9 +30,13 @@
      buf = lines.pop();
      for (const ln of lines)
        if (ln.startsWith('data: ') && !ln.includes('[DONE]'))
-          try { d.textContent += JSON.parse(ln.slice(6)).choices[0]?.delta?.content || '' } catch {}
+          try { const dl = JSON.parse(ln.slice(6)).choices[0]?.delta;
+            if (dl?.reasoning_content) { const s = document.createElement('span'); s.style.color = '#888';
+              s.textContent = dl.reasoning_content; rsn += dl.reasoning_content; d.appendChild(s) }
+            if (dl?.content) { const s = document.createElement('span');
+              s.textContent = dl.content; txt += dl.content; d.appendChild(s) } } catch {}
      chat.scrollTop = chat.scrollHeight;
    }
-    msgs.push({role: 'assistant', content: d.textContent});
+    const m = {role:'assistant', content:txt}; if (rsn) m.reasoning_content = rsn; msgs.push(m);
  }
 </script></body></html>
@@ -1,7 +1,7 @@
 from __future__ import annotations
 import functools, itertools, pathlib
 from dataclasses import dataclass, replace
-from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function
+from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, dtypes
 from tinygrad.nn import Linear
 from tinygrad.llm.gguf import gguf_load
 from tinygrad.uop.ops import resolve
@@ -187,8 +187,8 @@ class TransformerBlock(FFNBlock):

  def _init_state(self, x:Tensor):
    if not hasattr(self, "cache_kv"):
-      # TODO: how is the dtype of this determined?
-      self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim, device=x.device)
+      self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim,
+                                   dtype=dtypes.default_float, device=x.device)
      self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta, device=x.device)

 class MLATransformerBlock(FFNBlock):
@@ -261,13 +261,14 @@ class GatedDeltaNetBlock(FFNBlock):
  def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
    B, T, _ = x.shape
    assert T == 1, "GatedDeltaNetBlock currently only supports T=1"
+    is_kda = hasattr(self, "ssm_g_a")

    # input processing
    x = x.half()
-    out_gate = self.ssm_g_b(self.ssm_g_a(x)) if hasattr(self, "ssm_g_a") else self.attn_gate(x)
+    out_gate = self.ssm_g_b(self.ssm_g_a(x)) if is_kda else self.attn_gate(x)
    out_gate = out_gate.reshape(B, 1, self.num_v_heads, self.head_v_dim)
    beta = self.ssm_beta(x).sigmoid().reshape(B, self.num_v_heads, 1, 1)
-    alpha = self.ssm_f_b(self.ssm_f_a(x)) if hasattr(self, "ssm_f_a") else self.ssm_alpha(x)
+    alpha = self.ssm_f_b(self.ssm_f_a(x)) if is_kda else self.ssm_alpha(x)
    alpha = ((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, self.num_v_heads, -1) *
             self.ssm_a.reshape(1, self.num_v_heads, -1)).exp().unsqueeze(-2)

@@ -291,14 +292,13 @@ class GatedDeltaNetBlock(FFNBlock):

    # output
    core_attn_out = self.ssm_norm((recurrent_state@q).squeeze(-1).reshape(B, 1, self.num_v_heads, self.head_v_dim))
-    out_gate = out_gate.sigmoid() if hasattr(self, "ssm_g_a") else out_gate.silu()
+    out_gate = out_gate.sigmoid() if is_kda else out_gate.silu()
    return self.ssm_out((core_attn_out * out_gate).reshape(B, 1, -1).cast(x.dtype))

  # recurrent state can't be partially reused after divergence, force a full rebuild
  def _state_reset_ops(self):
    return [self.conv_state.assign(self.conv_state.const_like(0)),
            self.recurrent_state.assign(self.recurrent_state.const_like(0))] if hasattr(self, "conv_state") else []
-  def _reusable_prefix_len(self, prefix_len:int, cached_len:int) -> int: return 0 if prefix_len != cached_len else prefix_len

  def _init_state(self, x):
    if not hasattr(self, "conv_state"):
@@ -326,7 +326,8 @@ class Transformer:
  def forward(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor) -> Tensor:
    x = self.token_embd(tokens).float()                   # (B, T, D)
    for block in self.blk: x = block(x, start_pos)
-    logits = self.output(self.output_norm(x))[:, -1, :]
+    # only run the output projection on the last token
+    logits = self.output(self.output_norm(x[:, -1:]))[:, -1, :]
    # Gumbel-max trick: argmax(logits/temp - log(-log(uniform))) is equivalent to sampling from softmax(logits/temp)
    return (logits / temperature.maximum(1e-12) - (Tensor.rand_like(logits).maximum(1e-12).log().neg()).log()).argmax(-1, keepdim=True)

@@ -420,6 +421,10 @@ class Transformer:
    for _ in range(2): list(zip(range(2), self.generate([0])))

  def get_start_pos(self, tokens:list[int]) -> int:
+    # recurrent state can't be partially reused after divergence: reuse it only when tokens extend the cached prefix
+    if self.has_recurrent_block:
+      return len(self._cached_tokens) if self._cached_tokens and len(self._cached_tokens) < len(tokens) \
+        and tokens[:len(self._cached_tokens)] == self._cached_tokens else 0
    prefix_len = sum(1 for _ in itertools.takewhile(lambda ab: ab[0] == ab[1], zip(tokens[:-1], self._cached_tokens)))
    return min(block._reusable_prefix_len(prefix_len, len(self._cached_tokens)) for block in self.blk)

@@ -221,7 +221,7 @@ class ElementwiseMixin(CreationMixin):
    if dtypes.is_int(a.dtype) and dtypes.is_int(b.dtype): return a.alu(Ops.CMOD, b)
    return a - a.div(b, rounding_mode="trunc") * b

-  def div(self, x: Self | ConstType, reverse: bool = False, rounding_mode: Literal["trunc", "floor"] | None = None) -> Self:
+  def div(self, x: 'Self|ConstType|UOp', reverse: bool = False, rounding_mode: Literal["trunc", "floor"] | None = None) -> Self:
    """
    Divides `self` by `x`.
    Equivalent to `self / x`.
@@ -3,11 +3,16 @@ import math, dataclasses
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata, broadcast_axes
 from tinygrad.helpers import argsort
 from tinygrad.dtype import sum_acc_dtype
+from tinygrad.function import renumber_invalid_outputs

 def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
  if op == Ops.ADD: return (ctx._broadcast_to(ret.src[0].shape),)
  if op == Ops.MAX: return (((mask:=ret.src[0].eq(ret).cast(ctx.dtype))/mask._rop(Ops.ADD, tuple(range(ret.arg[1])))) * ctx,)
-  if op == Ops.MUL: return (ctx * ret / ret.src[0],)
+  if op == Ops.MUL:
+    # d(prod x)/dx_j = prod_{i!=j} x_i: ret/x_j whenever x_j != 0 (any zero makes ret 0), else the product of the others
+    safe_x, axes = (is_zero:=(x:=ret.src[0]).eq(0)).where(1, x), tuple(range(ret.arg[1]))
+    zero_count = is_zero.cast(sum_acc_dtype(is_zero.dtype))._rop(Ops.ADD, axes)
+    return (ctx * is_zero.where(zero_count.eq(1).where(safe_x._rop(Ops.MUL, axes), 0), ret/safe_x),)

 def _compact_params(body:UOp, all_args:tuple[UOp, ...]) -> tuple[UOp, tuple[UOp, ...]]:
  """Remove unused PARAMs from body and return compacted (body, args)."""
@@ -36,6 +41,7 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
  # collect needed gradient bodies, compact unused params, create a single backward CALL
  grad_bodies = [(i, grads[p]) for i in needed if (p:=params.get(i)) is not None and p in grads]
  bwd_body = UOp.maketuple(*(gb for _, gb in grad_bodies)).substitute(fwd_subs, walk=True)
+  bwd_body = renumber_invalid_outputs(bwd_body)
  bwd_body, compact_args = _compact_params(bwd_body, (*args, *grad_args, *fwd_outs))
  bwd_call = bwd_body.call(*compact_args, name=(k.arg.name or "")+"_backward", precompile=k.arg.precompile_backward)
  gb_map = {i: idx for idx, (i, _) in enumerate(grad_bodies)}
@@ -90,8 +90,11 @@ class MovementMixin:
        if resolve(index.step == 0, False): raise ValueError(f"{index=} cannot have 0 as step")
        start, stop = 0 if index.start is None else index.start, size if index.stop is None else index.stop
        step = 1 if index.step is None else index.step
+        # resolve negative int bounds against the (possibly symbolic) size, like slice.indices
+        if isinstance(start, int) and start < 0: start = start + size
+        if isinstance(stop, int) and stop < 0: stop = stop + size
        if all_int((start, stop, step)):
-          # handle int slicing (resolve negative bounds, clamp, stride)
+          # handle int slicing (clamp, stride)
          *bound, stride = index.indices(int(size.vmax) if isinstance(size, UOp) else size)
          bound = [0, 0] if stride * (bound[1] - bound[0]) < 0 else ([bound[1]+1, bound[0]+1] if stride < 0 else bound)
          return {"size":ceildiv(bound[1]-bound[0], abs(stride)), "boundary":tuple(bound), "stride":stride, "collapse_dim":False}
@@ -265,7 +268,8 @@ class MovementMixin:
    return self.shrink(tuple([None if ns is None else (0, ns) for ns in argfix(shape, *args)]))

  def pad_to(self, shape, *args) -> Self:
-    return self._mop(Ops.PAD, tuple((0, s if ns is None else ns) for s,ns in zip(self.shape, argfix(shape, *args), strict=True)))
+    ret = self._mop(Ops.PAD, tuple((0, s if ns is None else ns) for s,ns in zip(self.shape, argfix(shape, *args), strict=True)))
+    return self if ret.shape == self.shape else ret

  def view(self, shape, *args) -> Self:
    """`.view` is an alias for `.reshape`."""
@@ -514,7 +514,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
    output_dtype = self.dtype if dtypes.is_float(self.dtype) else dtypes.float32
    numerator = self.cast(sum_acc_dtype(self.dtype)).sum(axis=axis, keepdim=keepdim)
    denominator = prod([si for si, so in zip(self.shape, self.sum(axis=axis, keepdim=True).shape) if resolve(si != so)])
-    return numerator.div(denominator).cast(output_dtype)  # type: ignore[arg-type]
+    return numerator.div(denominator).cast(output_dtype)

  def var(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> Self:
    """
@@ -538,12 +538,11 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
    print(t.var(axis=1).numpy())
    ```
    """
+    output_dtype = self.dtype if dtypes.is_float(self.dtype) else dtypes.float32
    squares = (self - self.mean(axis=axis, keepdim=True)).square()
    n = prod([si for si, so in zip(self.shape, squares.sum(axis=axis, keepdim=True).shape) if resolve(si != so)])
-    reduced = squares.sum(axis=axis, keepdim=keepdim)
-    denominator = reduced.const_like(n) - correction  # type: ignore[arg-type]
-    # TODO: remove relu?
-    return reduced.div(denominator.relu())
+    numerator = squares.cast(sum_acc_dtype(self.dtype)).sum(axis=axis, keepdim=keepdim)
+    return numerator.div(smax(n - correction, 0)).cast(output_dtype)

  def var_mean(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> tuple[Self, Self]:
    """
@@ -1057,14 +1056,16 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
    assert not (align_corners and mode != "linear"), "align_corners option can only be set with the interpolating mode linear"
    x, expand = self, list(self.shape)
    for i in range(-1,-len(size)-1,-1):
-      scale = (int(self.shape[i]) - int(align_corners)) / (size[i] - int(align_corners))
-      arr, reshape = type(self).arange(size[i], dtype=dtypes.float32), [1] * self.ndim
+      in_sz, reshape = int(self.shape[i]), [1] * self.ndim
      reshape[i] = expand[i] = size[i]
      if mode == "linear":
-        index = (scale*arr if align_corners else (scale*(arr+0.5))-0.5).clip(0, self.shape[i]-1)
-        low, high, perc = [y.reshape(reshape).expand(expand) for y in (index.floor().int(), index.ceil().int(), index - index.floor())]
+        arr = type(self).arange(size[i])
+        num, den = (arr*(in_sz-1), size[i]-1) if align_corners else ((arr*2+1)*in_sz - size[i], size[i]*2)
+        num = num.clip(0, (in_sz-1)*den)
+        low, high, perc = [y.reshape(reshape).expand(expand) for y in (num//den, (num+den-1)//den, (num % den).cast(dtypes.float32)/den)]
        x = x.gather(i, low).lerp(x.gather(i, high), perc)
      else:
+        scale, arr = in_sz / size[i], type(self).arange(size[i], dtype=dtypes.float32)
        index = (scale*(arr+0.5) if mode=="nearest-exact" else scale*arr).cast(dtypes.int32).reshape(reshape).expand(expand)
        x = x.gather(i, index)
    return x.cast(self.dtype)
@@ -1,4 +1,4 @@
-import json, math, pathlib, zipfile, pickle, tarfile, struct, functools, io, zlib
+import json, math, pathlib, struct, functools, io, zlib
 from collections import OrderedDict
 from typing import Any, Callable, BinaryIO, Iterable, cast
 from tinygrad.tensor import Tensor
@@ -219,6 +219,7 @@ def load_state_dict(model, state_dict:dict[str, Tensor], strict=True, verbose=Tr

@accept_filename
 def zip_extract(t: Tensor) -> dict[str, Tensor]:
+  import zipfile
  files: dict[str, Tensor] = {}
  with zipfile.ZipFile(TensorIO(t), "r") as myzip:
    # sadly, the extra length needs to be read from the local header of each file.
@@ -249,6 +250,7 @@ def tar_extract(t: Tensor) -> dict[str, Tensor]:
  tensors = nn.state.tar_extract(Tensor(pathlib.Path("archive.tar")))
  ```
  """
+  import tarfile
  with tarfile.open(fileobj=TensorIO(t), mode="r") as tar:
    return {member.name:t[member.offset_data:member.offset_data+member.size] for member in tar if member.type == tarfile.REGTYPE}

@@ -303,6 +305,7 @@ def torch_load(t:Tensor) -> dict[str, Tensor]:
               "FloatTensor": None, "Parameter": Parameter}
  whitelist = {"torch", "collections", "numpy", "_codecs"}  # NOTE: this is not for security, only speed
  class Dummy: pass
+  import pickle, zipfile, tarfile
  class TorchPickle(pickle.Unpickler):
    def find_class(self, module, name):
      module_root = module.split(".")[0]
@@ -264,6 +264,7 @@ class ClangRenderer(CStyleLanguage):
  nan = '__builtin_nanf("")'

  # language options
+  barrier = "__atomic_thread_fence(__ATOMIC_SEQ_CST);"
  buffer_suffix = " restrict"
  type_map = {dtypes.bool:"_Bool", dtypes.half:"__fp16"}
  code_for_op = {**({k:v for k,v in CStyleLanguage.code_for_op.items() if k not in [Ops.EXP2, Ops.SIN, Ops.LOG2, Ops.TRUNC, Ops.RECIPROCAL]}),
@@ -35,21 +35,36 @@ def lcast(input_type:DType, output_type:DType):
    if dtypes.is_int(output_type): return 'trunc' if output_type.itemsize < input_type.itemsize else 'sext'
  raise NotImplementedError(f"cast from {input_type} -> {output_type} not implemented")

-def render_wmma_amd(ctx, wmma: UOp, cdna=False) -> str:
+def render_wmma_amd(ctx, wmma: UOp, cdna=False, rdna4=False) -> str:
  dt_map = {dtypes.half: "f16", dtypes.float: "f32", dtypes.ushort: "bf16.1k" if cdna else "bf16", dtypes.bfloat16: "bf16.1k" if cdna else "bf16",
            dtypes.fp8e4m3: ".fp8.fp8", dtypes.fp8e5m2: ".bf8.bf8", dtypes.int8: "iu8", dtypes.int32: "i32"}
  # https://github.com/llvm/llvm-project/blob/main/clang/test/CodeGenOpenCL/builtins-amdgcn-mfma.cl
  N,M,K = wmma.arg[0]
  if cdna:
    if K == 32: dt_map.update({dtypes.half: ".f16", dtypes.bfloat16: ".bf16"})
-    return f"  {ctx[wmma]} = call {ldt(wmma.dtype, wmma.max_numel())} @llvm.amdgcn.mfma.{dt_map[wmma.src[-1].dtype]}" + \
-           f".{N}x{M}x{K}{dt_map[wmma.arg[1]]}(" + ", ".join([f"{ldt(w.dtype, w.max_numel())} {ctx[w]}" for w in wmma.src]) + ", i32 0, i32 0, i32 0)"
+    scaled = K == 128
+    args = [f"{ldt(w.dtype, w.max_numel())} {ctx[w]}" for w in wmma.src]
+    # scaled mfma call require E8M0 scale args, byte = 0x7F = 127, scale = 2^(127 - 127) = 1.0
+    if scaled:
+      _fmt = { dtypes.fp8e5m2:1, dtypes.fp8e4m3:0 }
+      # (a_fp8_fmt, b_fp8_fmt, opsel, scale_a, opsel, scale_b)
+      args.extend([f"i32 {_fmt[wmma.arg[1]]}", f"i32 {_fmt[wmma.arg[1]]}", "i32 0", "i32 127", "i32 0", "i32 127"])
+    else: args.extend(["i32 0", "i32 0", "i32 0"]) # (cbsz, blgp, ?)
+
+    scale = "scale." if scaled else ""
+    dt_in = dt_map[wmma.arg[1]] if not scaled else ".f8f6f4"
+    return f"  {ctx[wmma]} = call {ldt(wmma.dtype, wmma.max_numel())} @llvm.amdgcn.mfma.{scale}{dt_map[wmma.src[-1].dtype]}" + \
+           f".{N}x{M}x{K}{dt_in}(" + ", ".join(args) + ")"
  # https://github.com/llvm/llvm-project/blob/main/llvm/test/CodeGen/AMDGPU/GlobalISel/llvm.amdgcn.wmma_32.ll
  # example: %wmma0 = call <8 x float> @llvm.amdgcn.wmma.f32.16x16x16.f16(<16 x half> %v99,<16 x half> %v100,<8 x float> %v101)
  args = [f"{ldt(w.dtype, w.max_numel())} {ctx[w]}" for w in wmma.src]
  if wmma.arg[1] == dtypes.int8: args = ["i1 true", args[0], "i1 true", args[1], args[2]]  # iu8 flags A/B signed
-  return f"  {ctx[wmma]} = call {ldt(wmma.dtype, wmma.max_numel())} @llvm.amdgcn.wmma.{dt_map[wmma.src[-1].dtype]}.16x16x16." + \
-    f"{dt_map[wmma.arg[1]]}(" + ", ".join(args) + (", i1 false)" if wmma.dtype != dtypes.float else ")")
+  if wmma.dtype != dtypes.float: args.append("i1 false") # opsel
+  def _bf16(dt:DType): return dtypes.ushort if dt is dtypes.bfloat16 else dt
+  suffix = f".v{wmma.max_numel()}{dt_map[_bf16(wmma.dtype)]}.v{wmma.src[0].max_numel()}{dt_map[_bf16(wmma.arg[1])]}" if rdna4 else ""
+  # bfloat treated as i16 in LLVM call
+  return f"  {ctx[wmma]} = call {ldt(_bf16(wmma.dtype), wmma.max_numel())} @llvm.amdgcn.wmma.{dt_map[wmma.src[-1].dtype]}.16x16x16." + \
+    f"{dt_map[wmma.arg[1]]}{suffix}(" + ", ".join(args) + ")"

 # llvm ops, lop[<dtype>][<op>]
 unsigned_lop = { Ops.ADD: "add", Ops.MUL: "mul", Ops.CDIV: "udiv", Ops.CMOD: "urem",
@@ -127,7 +142,7 @@ base_rewrite = PatternMatcher([
  (UPat(Ops.IF, name="x"), lambda ctx,x: f"  br i1 {ctx[x.src[0]]}, label %ifbody_{ctx[x][1:]}, label %ifskip_{ctx[x][1:]}\nifbody_{ctx[x][1:]}:"),
  (UPat(Ops.ENDIF, name="x"), lambda ctx,x: f"  br label %ifskip_{ctx[x.src[0]][1:]}\nifskip_{ctx[x.src[0]][1:]}:"),

-  (UPat(Ops.BARRIER), lambda ctx: "")
+  (UPat(Ops.BARRIER), lambda ctx: "  fence seq_cst")
 ])

 class LLVMRenderer(Renderer):
@@ -223,8 +238,8 @@ class AMDLLVMRenderer(LLVMRenderer):
    (UPat(Ops.CAST, dtypes.fp8s, (UPat(dtype=dtypes.float),), name="x",), lambda ctx,x:
      f"  {ctx[x]} = call i8 @f32_to_fp8({ldt(x.src[0].dtype)}  {ctx[x.src[0]]}, i1 {'1' if x.dtype == dtypes.fp8e5m2 else '0'})"),
    (UPat(Ops.CAST, dtypes.float, (UPat.var("y", dtypes.fp8s),), name="x",), lambda ctx,x,y:
-      f"  {ctx[x.src[0]]}_i32 = zext i8 {ctx[x.src[0]]} to i32\n"
-      f"  {ctx[x]} = call float @llvm.amdgcn.cvt.f32.{'bf8' if y.dtype == dtypes.fp8e5m2 else 'fp8'}(i32 {ctx[x.src[0]]}_i32, i32 0)"),
+      f"  {ctx[x]}_i32 = zext i8 {ctx[x.src[0]]} to i32\n"
+      f"  {ctx[x]} = call float @llvm.amdgcn.cvt.f32.{'bf8' if y.dtype == dtypes.fp8e5m2 else 'fp8'}(i32 {ctx[x]}_i32, i32 0)"),
  ]) + base_rewrite
  extra_matcher = LLVMRenderer.extra_matcher + create_non_native_float_pats(dtypes.fp8s) + PatternMatcher([
    # amd llvm intrinsics llvm.log2/llvm.exp2 don't support double
@@ -254,43 +269,19 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n  %trunc
    attributes = ["alwaysinline", "nounwind", '"no-builtins"',
                  f'"amdgpu-flat-work-group-size"="1,{requiredMaxThreadsPerBlock}"', '"no-trapping-math"="true"']
    return 'attributes #0 = { ' + ' '.join(attributes) + ' }'
+  @staticmethod
+  def is_rdna4(arch): return arch.split(':')[0] in {'gfx1200', 'gfx1201'}
  def __init__(self, target:Target):
    super().__init__(target)
    from tinygrad.runtime.support.compiler_llvm import AMDLLVMCompiler
    self.compiler, self.tensor_cores, self.is_cdna = AMDLLVMCompiler(target.arch), tc.get_amd(target.arch), HIPRenderer.is_cdna(target.arch)
-    self.string_rewrite += PatternMatcher([(UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, cdna=self.is_cdna: render_wmma_amd(ctx, wmma, cdna))])
-    if self.is_cdna:
-      self.extra_matcher += PatternMatcher([
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
-          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
-          if x.max_numel() == 4 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 4 else None),
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
-          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2]))
-          if x.max_numel() == 4 and x.src[0].dtype in dtypes.fp8_ocp and x.src[0].max_numel() == 8 else None),
-      ])
-    if target.arch in {"gfx1100", "gfx1151"}:
-      self.extra_matcher += PatternMatcher([
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.int32), lambda x: x.replace(
-          src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
-          if x.src[0].dtype == dtypes.int8 and x.src[0].max_numel() == 16 else None),
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(x.replace(
-          src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(j//2) if j%2 == 0 else UOp.const(0.0, x.src[2].dtype)
-            for j in range(x.max_numel()*2)))),
-          arg=(*x.arg[:4], None)).index(i*2)
-          for i in range(x.max_numel()))) if x.max_numel() == 8 else None),
-        (UPat(Ops.WMMA, name="x"), lambda x: x.replace(
-          src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
-          if x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 16 else None),
-      ])
-    if target.arch in {"gfx1200", "gfx1201"}:
-      self.extra_matcher += PatternMatcher([
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16), lambda x: x.replace(
-          src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2].bitcast(dtypes.uint16)))
-            .bitcast(dtypes.bfloat16) if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None),
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
-          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
-          if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None)
-      ])
+    self.string_rewrite += PatternMatcher([
+      (UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, rdna4=AMDLLVMRenderer.is_rdna4(target.arch), cdna=self.is_cdna:
+        render_wmma_amd(ctx, wmma, cdna, rdna4))
+    ])
+    if self.is_cdna: self.extra_matcher += tc.pm_validate_wmma_cdna
+    if target.arch in {"gfx1100", "gfx1151"}: self.extra_matcher += tc.pm_validate_wmma_rdna3
+    if target.arch in {"gfx1200", "gfx1201"}: self.extra_matcher += tc.pm_validate_wmma_rdna4

  def supported_dtypes(self): return {d for d in super().supported_dtypes()
                                      if (d not in dtypes.fp8_ocp or self.target.arch == "gfx950") and d not in dtypes.fp8_fnuz}
@@ -137,7 +137,8 @@ class NIRRenderer(Renderer):
    (UPat(Ops.CAST, (dtypes.uchar, dtypes.ushort), src=(UPat.var("x", dtypes.floats),), name="c"), lambda x,c: x.cast(dtypes.int32).cast(c.dtype)),
    # load/store use pointer arithmetic, and the cast does nothing. NOTE: this doesn't apply to image indexing cause it's 1-D
    (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"), UPat.var("off")), allow_any_len=True, name="x"), lambda x,buf,off: x.replace(
-      src=(buf,off.cast(dtypes.long))+x.src[2:]) if buf.addrspace != AddrSpace.REG and not is_image_shape(buf._shape) else None),
+      src=(buf,UOp.const(off.val, dtypes.long) if off.op is Ops.CONST else off.cast(dtypes.long))+x.src[2:])
+      if buf.addrspace != AddrSpace.REG and not is_image_shape(buf._shape) else None),
    # images need index to be int for nir (coordinates only: the INDEX keeps its access dtype)
    (UPat.var("buf").index(UPat.var("idx_y"), UPat.var("idx_x"), name="x"),
     lambda x,buf,idx_y,idx_x: x.replace(src=(buf, idx_y.cast(dtypes.int), idx_x.cast(dtypes.int)))),
@@ -1,4 +1,4 @@
-import glob, importlib, os, pathlib, shutil, subprocess, tarfile, tempfile
+import glob, importlib, os, pathlib, subprocess
 from tinygrad.helpers import fetch, flatten, system, getenv

 root = (here:=pathlib.Path(__file__).parent).parents[2]
@@ -31,6 +31,7 @@ def load(name, files, **kwargs):
  if not (f:=(root/(path:=kwargs.pop("path", __name__)).replace('.','/')/f"{name}.py")).exists() or getenv('REGEN'):
    files, kwargs['args'] = files() if callable(files) else files, args() if callable(args:=kwargs.get('args', [])) else args
    if (srcs:=kwargs.pop('srcs', None)):
+      import tempfile, tarfile
      srcpath = (td:=tempfile.TemporaryDirectory(f"autogen-src-{name.replace('/','-')}")).name + "/"
      for src in (srcs if isinstance(srcs, list) else [srcs]):
        if 'tar' in src:
@@ -157,7 +158,7 @@ def __getattr__(nm):
          *[f"python3 src/compiler/nir/nir_{s}_h.py --outdir gen" for s in ["intrinsics", "intrinsics_indices"]]]), cwd=path, shell=True, check=True),
  srcs="https://gitlab.freedesktop.org/mesa/mesa/-/archive/mesa-25.2.7/mesa-25.2.7.tar.gz",
  dll=f"'tinymesa_cpu' if DEV.renderer == 'LVP' else 'tinymesa', {tinymesa_path}, emsg='pip install tinymesa==25.2.7.2'",
-  prolog=["from tinygrad.helpers import DEV", "import gzip, base64, platform, sysconfig, os"],
+  prolog=["from tinygrad.helpers import DEV", "import gzip, base64, sysconfig, os"],
  epilog=lambda path: [system(f"{root}/extra/mesa/lvp_nir_options.sh {path}")])
    case "libclang":
      return load("libclang",
@@ -5,7 +5,7 @@ from typing import Literal, TypeAlias
 from tinygrad.runtime.support.c import _IO, _IOW, _IOR, _IOWR
 from tinygrad.runtime.support import c
 from tinygrad.helpers import DEV
-import gzip, base64, platform, sysconfig, os
+import gzip, base64, sysconfig, os
 dll = c.DLL('mesa', 'tinymesa_cpu' if DEV.renderer == 'LVP' else 'tinymesa', os.path.join(sysconfig.get_paths()['platlib'], 'tinymesa'), emsg='pip install tinymesa==25.2.7.2')
 class struct_u_printf_info(c.Struct): pass
 u_printf_info: TypeAlias = struct_u_printf_info
@@ -6,7 +6,6 @@ from tinygrad.device import Buffer, BufferSpec, Compiled, Device, MultiBuffer, P
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops, Variable
 from tinygrad.engine.jit import GraphRunner, MultiGraphRunner
-from tinygrad.runtime.ops_rdma import RDMACopyQueue

 class HCQGraph(MultiGraphRunner):
  def __init__(self, *args, **kwargs):
@@ -50,7 +49,7 @@ class HCQGraph(MultiGraphRunner):

    self.comp_queues: dict[HCQCompiled, HWQueue] = {dev: unwrap(dev.hw_compute_queue_t)() for dev in self.devices}
    self.copy_queues: dict[tuple[HCQCompiled, int], HWQueue] = {} # lazy allocation, keyed by (device, queue_idx)
-    self.rdma_queues: dict[tuple[HCQCompiled, HCQCompiled], RDMACopyQueue] = {} # lazy allocation, keyed by device pair
+    self.rdma_queues: dict[tuple[HCQCompiled, HCQCompiled], "RDMACopyQueue"] = {} # lazy allocation, keyed by device pair
    self.num_copy_queues: int = getenv("HCQ_NUM_SDMA", min(len(self.devices), 8) if ALL2ALL >= 1 else 1)
    self.num_rdma_ops: dict[tuple[HCQCompiled, HCQCompiled], int] = collections.defaultdict(int)

@@ -104,6 +103,7 @@ class HCQGraph(MultiGraphRunner):
      elif is_rdma:
        enqueue_queue = self.comp_queues[enqueue_dev]
        rdma_key = (cast(HCQCompiled, Device[bufs[0].device]).rdma_dev(), enqueue_dev.rdma_dev())
+        from tinygrad.runtime.ops_rdma import RDMACopyQueue
        self.rdma_queues.setdefault(rdma_key, RDMACopyQueue(enqueue_dev.rdma_dev()))
      else:
        assert (enqueue_dev.hw_copy_queue_t is not None), "device must implement a copy queue"
@@ -113,5 +113,6 @@ class MetalGraph(GraphRunner):
  @staticmethod
  def supports_uop(batch_devs, new_call:UOp) -> bool:
    # Metal ICB replay encodes offsets as uint32; reject if any Metal buffer offset exceeds 32-bit range.
-    if any(b.op is Ops.SLICE and b.src[1].val * b.src[0].dtype.itemsize > 0xFFFFFFFF for b in new_call.src[1:]): return False
+    for shrink in [s for src in new_call.src[1:] if (s:=src.src[0] if src.op is Ops.BITCAST else src).op is Ops.SHRINK]:
+      if shrink.src[1].val * shrink.src[0].dtype.itemsize > 0xFFFFFFFF: return False
    return GraphRunner.supports_uop(batch_devs, new_call)
@@ -1,11 +1,11 @@
 from __future__ import annotations
-import platform, sys, os, ctypes, functools, mmap, threading, array, itertools
-from dataclasses import replace
-from typing import cast
-from tinygrad.helpers import to_mv, OSX, WIN, Context, mv_address, suppress_finalizing, unwrap, data64_le, partition
-from tinygrad.device import Buffer, BufferSpec, TinyELF
-from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, HCQArgsState, HCQSignal, HCQProgram, MMIOInterface
-from tinygrad.runtime.support.hcq import CLikeArgsState
+import platform, sys, os, ctypes, functools, mmap, threading, array, struct, time
+from dataclasses import dataclass, replace
+from typing import cast, Callable
+from tinygrad.helpers import to_mv, from_mv, OSX, WIN, Context, mv_address, suppress_finalizing, unwrap, data64_le, to_tuple
+from tinygrad.device import Buffer, BufferSpec, TinyELF, Program, Device
+from tinygrad.runtime.support.hcq import HCQBuffer, MMIOInterface
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, make_cmdbuf, make_signal
 from tinygrad.renderer.cstyle import ClangRenderer
 from tinygrad.renderer.llvmir import CPULLVMRenderer
 from tinygrad.renderer.nir import LVPRenderer
@@ -13,11 +13,15 @@ from tinygrad.renderer.isa.x86 import X86Renderer
 from tinygrad.runtime.support.elf import jit_loader
 from tinygrad.runtime.autogen import libc
 from tinygrad.codegen import do_to_program
+from tinygrad.engine.realize import pm_flatten_linear, get_call_arg_uops, get_runtime
 from tinygrad import UOp, dtypes
 from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import sint, KernelInfo, Ops, UPat, PatternMatcher, graph_rewrite
+from tinygrad.uop.ops import KernelInfo, Ops, UPat, PatternMatcher, graph_rewrite

-MAX_ARGS, CMD_SIZE, RING_SLOTS = 63, 64, (16 << 10)
+MAX_ARGS, CMD_SIZE, RING_SLOTS, FUNCS = 63, 64, (16 << 10), (() if WIN else ('clock_gettime', 'sem_wait', 'sem_post'))
+
+# *****************
+# 1. workers

 def signal_prog():
  val = UOp.param(1, dtypes.int, (), vmin_vmax=(0, dtypes.int.max), name="value", addrspace=AddrSpace.ALU)
@@ -35,79 +39,86 @@ def timestamp_prog():
    val = ts.after(call)[0].load() * 1_000_000_000 + ts.after(call)[1].load()
  return UOp.param(0, dtypes.uint64, (1,))[0].store(val)

-def quit_prog():
-  fn = UOp.param(0, dtypes.uint64, (1 if WIN else 3,))
-  if WIN: return fn[0].load().call(UOp.const(0, dtypes.uint64), ret_dtype=dtypes.void) # ExitThread(0)
-  sem = UOp.param(1, dtypes.uint64, (1,))
-
-  close = fn[2].load().call(sem[0], ret_dtype=dtypes.void) # sem_close(sem)
-  return fn.after(close)[0].load().call(UOp.const(0, dtypes.uint64), ret_dtype=dtypes.void) # pthread_exit(0)
-
 def worker_prog():
  ring = UOp.param(0, dtypes.uint64, (RING_SLOTS * CMD_SIZE,), volatile=True)
-  wait, sem = UOp.param(1, dtypes.uint64, (1,), volatile=True), UOp.param(2, dtypes.uint64, (1,))
-  cur = UOp.range(2**64-1, 0, dtype=dtypes.uint64)
+  wait, done = UOp.param(1, dtypes.uint64, (1,), volatile=True), UOp.param(2, dtypes.uint64, (1,), volatile=True)
+  sem, cur = UOp.param(3, dtypes.uint64, (1,)), UOp.range(2**64-1, 0, dtype=dtypes.uint64) # sem is unused on windows, it has to come last

  # spin on windows, sem_wait to sleep on posix
  if WIN: ready = (v:=wait.after(lw:=UOp.loop(1), cur)[0].load()).end(lw, v <= cur)
-  else: ready = wait.after(cur)[0].load().call(sem.after(cur)[0], ret_dtype=dtypes.void)
+  else: ready = (rv:=wait.after(lw:=UOp.loop(1), cur)[0].load().call(sem.after(cur)[0], ret_dtype=dtypes.int)).end(lw, rv != 0)

  entry = [ring.after(ready).index((cur % RING_SLOTS) * CMD_SIZE + i).load() for i in range(CMD_SIZE)]
-  return entry[0].call(*entry[1:], ret_dtype=dtypes.void).end(cur)
+  return done.after(entry[0].call(*entry[1:], ret_dtype=dtypes.void)).index(0).store(cur + 1).end(cur)

-def host_wait(ctx, dst:UOp, val:UOp) -> UOp:
-  return (cur:=dst.after(loop:=UOp.loop(next(ctx))).index(UOp.const(0, dtypes.int)).load()).end(loop, cur < val)
+@dataclass
+class CPUWorker: ring:Buffer; put:Buffer; sem:Buffer; sys:Buffer; done:Buffer; thread:threading.Thread  # noqa: E702

-pm_host_opsel = PatternMatcher([(UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val"))), host_wait)])
+# *****************
+# 2. queue encoders

-def encode_host_queue(q:UOp) -> UOp:
-  # TODO: subset of hcq2 for now
-  spins, (store,) = partition(graph_rewrite(q, pm_host_opsel, ctx=itertools.count(), walk=True, name="host opsel").src, lambda u: u.op is Ops.END)
-  assert store.op is Ops.INS and store.arg == "store", f"host queue cannot encode {store.op} {store.arg}"
-  return store.src[0].after(*spins).index(UOp.const(0, dtypes.int)).store(store.src[1])
+def cpu_cmd(devs:tuple[str, ...], prog, *args:UOp) -> UOp:
+  progs = [get_runtime(d, prog) if isinstance(prog, UOp) else cast(CPUDevice, Device[d]).prgs[prog] for d in devs]
+  addrs = tuple(UOp.const(p.addr, dtypes.uint64) for p in progs)
+  words = ((addrs[0] if len(addrs) == 1 else UOp(Ops.STACK, dtypes.uint64, addrs)),) + args
+  return UOp(Ops.INS, dtypes.void, words + (UOp.const(0, dtypes.uint64),) * (CMD_SIZE - len(words)), arg="cmd")

-class CPUComputeQueue(HWQueue):
-  def __init__(self, dev):
-    super().__init__()
-    self.dev = dev
-  def _cmd(self, prog, args=(), vals=()): return self.exec(prg:=self.dev.prgs[prog], prg.fill_kernargs(args, vals), None, None)
-  def memory_barrier(self): return self
-  def exec(self, prg:CPUProgram, args_state:HCQArgsState, global_size, local_size):
-    if (lvp:=isinstance(args_state, LVPArgsState)): self.bind_args_state(args_state)
-    args:list[sint|None] = [args_state.buf.va_addr] if lvp else [*[x.va_addr for x in args_state.bufs], *args_state.vals]
-    assert len(args) <= MAX_ARGS, f"CPU programs support at most {MAX_ARGS} arguments, got {len(args)}"
-    for tid in range(1 if lvp else (global_size or (1,))[0]):
-      if not lvp and 'core_id' in prg.runtimevars: args[prg.runtimevars['core_id']] = tid
-      self.q(prg, *[unwrap(x) for x in args], *([0] * (MAX_ARGS - len(args))))
-    return self
-  def wait(self, signal, value=0): return self._cmd(wait_prog, (signal.base_buf,), (value,))
-  def timestamp(self, signal): return self._cmd(timestamp_prog, (signal.base_buf.offset(8, 8), self.dev.func_table._buf.offset(0, 8)))
-  def signal(self, signal, value:sint=0): return self._cmd(signal_prog, (signal.base_buf,), (value,))
-  def _submit(self, dev):
-    dev.ensure_worker()
-    ring_view = dev.ring.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
-    for off in range(0, len(self._q), CMD_SIZE):
-      entry = [self._q[off].addr, *self._q[off+1:off+CMD_SIZE]]
-      ring_view[(base:=(dev.ring_pos % RING_SLOTS) * CMD_SIZE):base+CMD_SIZE] = array.array('Q', (int(x) & ((1<<64)-1) for x in entry))
-      dev.ring_pos += 1
-      if WIN: dev.sys.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = dev.ring_pos
-      else: assert libc.sem_post(dev.sem) == 0
+def cpu_exec(ctx:tuple[str, ...], call:UOp, prg:UOp) -> UOp:
+  args = [get_call_arg_uops(call)[i].getaddr(ctx) for i in prg.arg.globals] + [v.cast(dtypes.uint64) for v in prg.arg.vars]
+  if (core:=prg.arg.runtimevars.get('core_id')) is None: return cpu_cmd(ctx, prg, *args)

-class LVPArgsState(CLikeArgsState):
-  def __init__(self, buf, prg, bufs, vals=()): super().__init__(buf, prg, bufs, vals, [*data64_le(buf.va_addr + 12), (len(bufs) + len(vals)) * 2])
+  la = [cpu_cmd(ctx,prg,*args[:(cid:=(len(prg.arg.globals)+core))],UOp.const(t, dtypes.uint64),*args[cid+1:]) for t in range(prg.arg.global_size[0])]
+  return UOp(Ops.LINEAR, dtypes.void, tuple(la))
+
+pm_cpu_opsel = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), cpu_exec),
+
+  (UPat(Ops.INS, arg="barrier"), lambda: UOp(Ops.NOOP, dtypes.void, ())),
+  (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val"))),
+   lambda ctx, dst, val: cpu_cmd(ctx, wait_prog, dst.getaddr(ctx), val.cast(dtypes.uint64))),
+  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))),
+   lambda ctx, dst, val: cpu_cmd(ctx, signal_prog, dst.getaddr(ctx), val.cast(dtypes.uint64))),
+  (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),)),
+   lambda ctx, dst: cpu_cmd(ctx, timestamp_prog, dst.getaddr(ctx), *(() if WIN else (make_signal(ctx, tag="func:clock_gettime").getaddr(ctx),)))),
+])
+
+def encode_queue(q:UOp) -> UOp:
+  devs, queue = to_tuple(q.arg[0]), q.arg[1]
+  lin = graph_rewrite(q, pm_cpu_opsel+pm_flatten_linear, ctx=devs, walk=True, name=f"{queue} opsel")
+
+  cnt = sum(len(ins.src) for ins in lin.src) // CMD_SIZE
+  assert cnt < RING_SLOTS, f"submit of {cnt} entries doesn't fit the ring"
+  cmdbuf = make_cmdbuf(lin, devs, buf=UOp.placeholder((cnt*CMD_SIZE,), dtypes.uint64, next(UOp.unique_num), device=devs).rtag("cmdbuf"))
+  ring = UOp.placeholder((ring_words:=RING_SLOTS*CMD_SIZE,), dtypes.uint64, 0, device=devs, volatile=True).rtag(f"{queue}_ring")
+  put, done, sem, sysbuf = (make_signal(devs, tag=f"{queue}_{name}") for name in ("put", "done", "sem", "sys"))
+
+  # submits are serialized on the submitter, so they can bump put without atomics
+  ran = done.after(l:=UOp.loop(next(UOp.unique_num))).index(0).load()
+  room = ran.end(l, put.index(0).load() - ran > RING_SLOTS - cnt) # wait until cnt entries fit in the ring
+  base = ((put.after(room).index(0).load() % RING_SLOTS) * CMD_SIZE).cast(dtypes.int)
+  e = UOp.range(cnt, next(UOp.unique_num), dtype=dtypes.int, src=(cmdbuf, ring))
+  copy = UOp.group(*[ring.index((base + e*CMD_SIZE + w) % ring_words).store(cmdbuf.index(e*CMD_SIZE + w).load()) for w in range(CMD_SIZE)])
+
+  # wake the worker after each entry, keeping the post with the stores stops it from hoisting out of the loop
+  wake = copy.end(e) if WIN else make_signal(devs, tag="func:sem_post").after(copy).index(0).load().call(sem.index(0), ret_dtype=dtypes.void).end(e)
+  bumped = put.after(wake).index(0).store(put.index(0).load() + cnt)
+  return sysbuf.after(bumped).index(0).store(put.index(0).load() + cnt) if WIN else bumped
+
+# *****************

 # NOTE: MAP_JIT is added to mmap module in python 3.13
 MAP_JIT = 0x0800

-class CPUProgram(HCQProgram['CPUDevice']):
+class CPUProgram(Program['CPUDevice']):
  rt_lib = None
  try: rt_lib = ctypes.CDLL(ctypes.util.find_library('System' if OSX else 'kernel32') if OSX or WIN else 'libgcc_s.so.1')
  except OSError: pass

  def __init__(self, dev:CPUDevice, obj:TinyELF):
-    self.signature, self.runtimevars = obj.signature, {name:slot for name,slot,*_ in obj.signature if name == 'core_id'}
+    self.dev, self.name, self.signature = dev, obj.name, obj.signature
+    self.runtimevars = {name:slot for name,slot,*_ in obj.signature if name == 'core_id'}
+    self.lvp = obj.target.renderer == "LVP"

-    LVP = obj.target.renderer == "LVP"
    if sys.platform == "win32": # mypy doesn't understand when WIN is used here
      PAGE_EXECUTE_READWRITE, MEM_COMMIT, MEM_RESERVE = 0x40, 0x1000, 0x2000
      ctypes.windll.kernel32.VirtualAlloc.restype = ctypes.c_void_p
@@ -117,7 +128,7 @@ class CPUProgram(HCQProgram['CPUDevice']):
      ctypes.windll.kernel32.GetCurrentProcess.restype = ctypes.c_void_p
      proc = ctypes.windll.kernel32.GetCurrentProcess()
      ctypes.windll.kernel32.FlushInstructionCache(ctypes.c_void_p(proc), ctypes.c_void_p(self.addr), ctypes.c_size_t(len(obj.lib)))
-      self.fxn = ctypes.CFUNCTYPE(None)(self.addr)
+      self.fxn = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(self.addr) if self.lvp else ctypes.CFUNCTYPE(None)(self.addr)
    else:
      # On apple silicon with SPRR enabled (it always is in macos) RWX pages are unrepresentable: https://blog.svenpeter.dev/posts/m1_sprr_gxf/
      # MAP_JIT allows us to easily flip pages from RW- to R-X and vice versa. It is a noop on intel cpus. (man pthread_jit_write_protect_np)
@@ -125,7 +136,7 @@ class CPUProgram(HCQProgram['CPUDevice']):
      self.addr = mv_address(self.mem)

      if OSX: unwrap(CPUProgram.rt_lib).pthread_jit_write_protect_np(False)
-      lib = jit_loader(obj.lib, base=ctypes.addressof(ctypes.c_void_p.from_buffer(self.mem)), link_libs=['m']) if LVP else obj.lib
+      lib = jit_loader(obj.lib, base=ctypes.addressof(ctypes.c_void_p.from_buffer(self.mem)), link_libs=['m']) if self.lvp else obj.lib
      self.mem.write(lib)
      if OSX: unwrap(CPUProgram.rt_lib).pthread_jit_write_protect_np(True)

@@ -138,15 +149,30 @@ class CPUProgram(HCQProgram['CPUDevice']):
        # msync should be a universal POSIX way to do this
        libc.msync(ctypes.c_void_p(self.addr), len(lib), libc.MS_SYNC | libc.MS_INVALIDATE)

-      self.fxn = ctypes.CFUNCTYPE(None)(self.addr)
+      self.fxn = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(self.addr) if self.lvp else ctypes.CFUNCTYPE(None)(self.addr)

-    super().__init__(LVPArgsState if LVP else HCQArgsState, dev, obj, kernargs_alloc_size=12+256 if LVP else 0)
+  def __call__(self, *bufs:HCQBuffer, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),
+               vals:tuple[int|None, ...]=(), wait:bool=False, timeout:int|None=None) -> float|None:
+    st = time.perf_counter()
+    if self.lvp:
+      lvp_args = bytearray(12 + (len(bufs) + len(vals)) * 8)
+      addr = mv_address(lvp_args)
+      struct.pack_into(f'<3I{len(bufs)}Q', lvp_args, 0, *data64_le(addr+12), (len(bufs)+len(vals))*2, *[b.va_addr for b in bufs])
+      for v,(off,dt) in zip(vals, TinyELF.iter_sig(self.signature[-len(vals):], len(bufs)*8)): struct.pack_into(f'<{dt.fmt}', lvp_args, 12+off, v)
+      self.fxn(addr)
+    else:
+      args = [*[cast(int, b.va_addr) for b in bufs], *cast(tuple[int, ...], vals)]
+      assert len(args) <= MAX_ARGS, f"CPU programs support at most {MAX_ARGS} arguments, got {len(args)}"
+      for tid in range(global_size[0]):
+        if 'core_id' in self.runtimevars: args[self.runtimevars['core_id']] = tid
+        self.fxn(*[ctypes.c_uint64(x) for x in args])
+    return time.perf_counter() - st if wait else None

  @suppress_finalizing
  def __del__(self):
    if sys.platform == 'win32': ctypes.windll.kernel32.VirtualFree(ctypes.c_void_p(self.addr), ctypes.c_size_t(0), 0x8000) #0x8000 - MEM_RELEASE

-class CPUAllocator(HCQAllocator):
+class CPUAllocator(HCQAllocator['CPUDevice']):
  def __init__(self, dev:CPUDevice): super().__init__(dev, supports_copy_from_disk=False, supports_transfer=False)
  def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
    if options.external_ptr is not None: addr, buf = options.external_ptr, None
@@ -154,67 +180,55 @@ class CPUAllocator(HCQAllocator):
    else: addr = mv_address(buf:=mmap.mmap(-1, size, mmap.MAP_ANON | mmap.MAP_SHARED, mmap.PROT_READ | mmap.PROT_WRITE))
    return HCQBuffer(va:=addr, sz:=size, meta=buf, view=MMIOInterface(va, sz, fmt='B'), owner=self.dev)
  def _as_buffer(self, src) -> memoryview: return to_mv(src.va_addr, src.size)
+  def _copyin(self, dest:HCQBuffer, src:memoryview):
+    self.dev.synchronize()
+    ctypes.memmove(int(dest.va_addr), from_mv(src), len(src))
+  def _copyout(self, dest:memoryview, src:HCQBuffer):
+    self.dev.synchronize()
+    ctypes.memmove(from_mv(dest), int(src.va_addr), len(dest))
  def _do_map(self, buf:HCQBuffer):
    if buf.view is None or not isinstance(buf.view, MMIOInterface): raise RuntimeError("Cannot map buffer without view to cpu")
    return HCQBuffer(buf.view.addr, buf.size, view=buf.view, owner=buf.owner)
  def _unmap(self, mb): pass  # CPU _do_map returns a view wrapper, nothing to release

-class CPUDevice(HCQCompiled):
-  pm_lower = PatternMatcher([
-    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),)), encode_host_queue)])
-
-  pm_bufferize = PatternMatcher([
-    (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx[0].signal("sentinel", (1 << 64) - 1)),
-    (UPat(Ops.PARAM, tag="timeline_signal"), lambda ctx: ctx[0].signal("timeline")),
-    (UPat(Ops.PARAM, tag="timeline_value"), lambda ctx: ctx[0].signal("value", 1)),
-  ])
-
-  @functools.cache
-  def signal(self, name:str, init_value:int=0) -> Buffer:
-    (buf:=Buffer(self.device, 1, dtypes.uint64, preallocate=True)).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = init_value
-    return buf
+class CPUDevice(HCQ2Compiled):
+  wait_timeout_ms, has_copy_queue = 30000, False
+  pm_lower = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),)), encode_queue)])

  def __init__(self, device:str=""):
-    super().__init__(device, CPUAllocator(self), [ClangRenderer, CPULLVMRenderer, LVPRenderer, X86Renderer], CPUProgram, HCQSignal,
-      functools.partial(CPUComputeQueue, self), arch={'amd64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine().lower(), m)+",native")
+    super().__init__(device, CPUAllocator(self), [ClangRenderer, CPULLVMRenderer, LVPRenderer, X86Renderer], CPUProgram,
+      arch={'amd64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine().lower(), m)+",native")

-    self.ring_pos = 0
+    self.pm_bufferize = PatternMatcher(
+      [(UPat(Ops.PARAM, tag=f"COMPUTE:0_{n}"), lambda ctx, n=n: getattr(ctx[0].worker, n)) for n in ("ring", "put", "sem", "sys", "done")] +
+      [(UPat(Ops.PARAM, tag=f"func:{f}"), lambda ctx, f=f: ctx[0].func_ptr(f)) for f in FUNCS]) + self.pm_bufferize

-    # posix uses sem to put cpus into sleep
-    self.sem_addr = 0
-    if not WIN:
-      self.sem = libc.sem_open(sem_name:=f"/tinygrad-{os.getpid()}-{id(self):x}".encode(), os.O_CREAT|os.O_EXCL, 0o600, 0) # type: ignore[call-arg]
-      self.sem_addr = unwrap(ctypes.cast(self.sem, ctypes.c_void_p).value)
-      if self.sem_addr == ctypes.c_void_p(-1).value or libc.sem_unlink(sem_name): raise OSError(ctypes.get_errno(), "semaphore")
-
-    # TODO: move to hcq2
    with Context(EMULATED_DTYPES="", TRACK_MATCH_STATS=0):
-      prgs = {f: f().sink(arg=KernelInfo(f.__name__), tag=1) for f in (signal_prog, wait_prog, timestamp_prog, quit_prog, worker_prog)}
-      self.prgs = {f: self.runtime(do_to_program(v, ClangRenderer(replace(self.renderer.target, renderer="CLANG"))).to_elf()) for f,v in prgs.items()}
+      clang = ClangRenderer(replace(self.renderer.target, renderer="CLANG"))
+      self.prgs:dict[Callable, CPUProgram] = {f: CPUProgram(self, do_to_program(f().sink(arg=KernelInfo(f.__name__), tag=1), clang).to_elf())
+                                              for f in (signal_prog, wait_prog, timestamp_prog, worker_prog)}

-  @functools.cached_property
-  def ring(self) -> Buffer: return Buffer(self.device, RING_SLOTS * CMD_SIZE, dtypes.uint64, preallocate=True)
-  @functools.cached_property
-  def sys(self) -> Buffer: return Buffer(self.device, 1, dtypes.uint64, preallocate=True)
-  @functools.cached_property
-  def sem_buf(self) -> Buffer: return Buffer(self.device, 1, dtypes.uint8, options=BufferSpec(external_ptr=self.sem_addr), preallocate=True)
+  def func_ptr(self, name:str) -> Buffer: return self.func_table.view(1, dtypes.uint64, FUNCS.index(name)*8).ensure_allocated()

-  # TODO: move to hcq2 infra
  @functools.cached_property
  def func_table(self) -> Buffer:
-    fns = ([0, ctypes.windll.kernel32.ExitThread, 0, 0] if WIN else  # type: ignore[attr-defined]
-           [libc.dll.clock_gettime, libc.dll.pthread_exit, libc.dll.sem_wait, libc.dll.sem_close])
-    addrs = array.array('Q', [unwrap(ctypes.cast(f, ctypes.c_void_p).value) if f else 0 for f in fns])
-    (ft:=Buffer(self.device, len(fns), dtypes.uint64, preallocate=True)).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[:] = addrs
+    lib = ctypes.windll.kernel32 if sys.platform == "win32" else libc.dll # type: ignore[attr-defined]
+    (ft:=Buffer(self.device, len(FUNCS), dtypes.uint64, preallocate=True))._buf.cpu_view().view(fmt='Q')[:] = \
+      array.array('Q', [unwrap(ctypes.cast(getattr(lib, f), ctypes.c_void_p).value) for f in FUNCS])
    return ft

-  @functools.cache
-  def ensure_worker(self):
-    threading.Thread(target=cast(CPUProgram, self.prgs[worker_prog]).fxn, daemon=True, args=[ctypes.c_uint64(x) for x in
-      [self.ring._buf.va_addr, self.sys._buf.va_addr if WIN else self.func_table._buf.va_addr+16, self.sem_addr]]).start()
+  @functools.cached_property
+  def worker(self) -> CPUWorker:
+    ring, put, sysbuf, done = (Buffer(self.device, sz, dtypes.uint64, preallocate=True) for sz in (RING_SLOTS*CMD_SIZE, 1, 1, 1))
+    addr, hsem = 0, None

-  def finalize(self):
-    if self.ring_pos == 0: return # the worker starts with the first submit
-    ft = self.func_table._buf
-    CPUComputeQueue(self)._cmd(quit_prog, (ft.offset(8, 8),) if WIN else (ft.offset(8, 24), self.sem_buf._buf)).submit(self)
-    self.ring_pos = 0
+    # sem are posix-only
+    if not WIN:
+      hsem = libc.sem_open(nm:=f"/tinygrad-{os.getpid()}-{id(ring):x}".encode(), os.O_CREAT|os.O_EXCL, 0o600, 0) # type: ignore[call-arg]
+      if (addr:=unwrap(ctypes.cast(hsem, ctypes.c_void_p).value)) == ctypes.c_void_p(-1).value or libc.sem_unlink(nm):
+        raise OSError(ctypes.get_errno(), "semaphore")
+    sem = Buffer(self.device, 1, dtypes.uint64, options=BufferSpec(external_ptr=addr), preallocate=True)
+
+    worker_args = [ring._buf.va_addr, sysbuf._buf.va_addr if WIN else self.func_ptr('sem_wait')._buf.va_addr, done._buf.va_addr, addr]
+    (worker:=threading.Thread(target=self.prgs[worker_prog].fxn, daemon=True, args=[ctypes.c_uint64(x) for x in worker_args])).start()
+    return CPUWorker(ring, put, sem, sysbuf, done, worker)
@@ -3,10 +3,10 @@
 # works to test the tensor cores, and all the uops in general
 # this is the (living) definition of uops
 from typing import Any, TYPE_CHECKING
-import pickle, base64, itertools, time, sys, functools
+import pickle, base64, itertools, time, sys, functools, ctypes
 from dataclasses import replace
 from tinygrad.dtype import bitcast, DType, dtypes, AddrSpace, truncate, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar
-from tinygrad.helpers import all_same, getenv, flatten, Target, IMAGE, is_image_shape, cpu_profile
+from tinygrad.helpers import all_same, getenv, flatten, Target, IMAGE, is_image_shape, cpu_profile, mv_address
 from tinygrad.device import Buffer, Compiled, Compiler, Allocator, Program, TinyELF
 from tinygrad.codegen.opt import tc
 from tinygrad.uop.ops import exec_alu, python_alu, Ops, UOp, GroupOp
@@ -134,6 +134,13 @@ class PythonProgram(Program['PythonDevice']):
                               for k in range(len(src_values))], j, u.dtype) for j in range(load_sz)]
          else:
            values[u] = load(src_values, 0, u.dtype)
+        elif u.op is Ops.CALL:
+          assert u.dtype is dtypes.void
+          cfunc = ctypes.CFUNCTYPE(None, *[ctypes.c_uint64] * (len(src_values)-1))
+          values[u] = []
+          for args,gate in zip(zip(*src_values), exec_masks[-1]):
+            call_args = [(mv_address(x[0]) + x[1]*dt.itemsize) if isinstance(x, tuple) else x for x,dt in zip(args, src_dtypes)]
+            values[u].append(cfunc(call_args[0])(*call_args[1:]) if gate else None)
        elif u.op is Ops.WMMA:
          first_src_dtype = u.src[0].dtype
          assert isinstance(first_src_dtype, DType) # mypy
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
qazalandGitHub	a57569349c	renumber invalids before callify (#17542 ) * renumber invalids before callify * change * Revert "change" This reverts commit 6f4df1541e79721a85ee3f5801114454f264c973. * renumber in tensor * scope renumber_invalid_outputs * cleanup	2026-08-15 18:00:54 +09:00
qazalandGitHub	5c43a89fb1	precompile_backward tests for sched_cache (#17544 ) * work * back * work * keep +	2026-08-15 15:25:31 +09:00
qazalandGitHub	e6f5bb9c09	simple test for Invalid clone cache miss regression (#17541 ) * simple test for Invalid clone cache miss regression * xfail * _	2026-08-15 11:00:21 +09:00
chenyuandGitHub	4b0525e594	no pm_fold_cast_const in UOp.simplify and hcq2 [pr] (#17540 )	2026-08-14 21:39:47 -04:00
chenyuandGitHub	64ccbde3bb	clean up STACK with a const [PR] (#17539 )	2026-08-14 21:07:01 -04:00
chenyuandGitHub	6ea665ed66	remove pm_fold_cast_const from dtype decomp [pr] (#17536 ) * remove pm_fold_cast_const from dtype decomp [pr] * fix	2026-08-14 16:24:51 -04:00
nimlgenandGitHub	0725acc392	reenable hcq2 ci (#17532 )	2026-08-14 23:13:18 +03:00
wozeparrotandGitHub	4a1f32977c	gptoss: default GROUPED_MOE=1 (#17537 )	2026-08-14 12:24:37 -07:00
chenyuandGitHub	13c381b0c0	remove pm_fold_cast_const from initial symbolic [pr] (#17535 ) interestingly it gives more accurate numerics when composing const like log10	2026-08-14 14:14:40 -04:00
chenyuandGitHub	ac7067ac60	fix deconstruct_function for python 3.11 (#17534 )	2026-08-14 13:38:34 -04:00
chenyuandGitHub	80169c6758	remove where push cast to branches from sym [pr] (#17533 ) * remove where push cast to branches from sym [pr] not really needed and one less place that generates casted weak const when it's not needed * fix	2026-08-14 12:57:56 -04:00
nimlgenandGitHub	adacaa3e17	hcq2 fix hangs (#17529 )	2026-08-14 16:23:59 +03:00
chenyuandGitHub	89ab344c42	fix assign into bitcast with no explicit realize (#17531 )	2026-08-14 09:23:01 -04:00
chenyuandGitHub	25c3bd027b	remove pm_fold_cast_const in simplify_merge_adjacent [pr] (#17530 )	2026-08-14 09:10:05 -04:00
nimlgenandGitHub	6b35220622	cpu hcq2 (#17503 ) * cpu hcq2 * temp * slop * test with backpressure * x * x * x * x * x * x * Dx * save reverts * um? * x * x * call from py * x? * x * submitters gone * x * x * z * Dx * Dx * x * x * fixes * repl * x * f * for now keep hcqbuffer	2026-08-14 15:06:03 +03:00
George HotzandGitHub	b1859805b1	remove Ops.BIND (#17511 ) * remove Ops.BIND * param arg * simplify that * simplify * cleaner * props, not functions * param and buffer can share	2026-08-13 23:52:16 -07:00
qazalandGitHub	faba071b1d	don't enter CALL body in assign fixups (#17527 ) * fix python time regression in mxfp4 * this saves even more time * s_nop test * itertools count * cleanup	2026-08-14 15:05:04 +09:00
qazalandGitHub	81dc8ec232	Revert "amd: fix ALL2ALL speed on amdgpu (gpt) (#17518 )" (#17526 ) This reverts commit `cc6d33bde7`.	2026-08-14 10:39:14 +09:00
geohot	95ca5081fe	hotfix: update extra/runbook_digitalocean_mi350x	2026-08-13 17:26:39 -07:00
sirhcmandGitHub	303d1677b3	qcomcl: use qemu for compilation (#17524 )	2026-08-13 20:16:02 -04:00
George HotzandGitHub	673c6463f9	disable HCQ2 for AMD CI (#17523 )	2026-08-13 13:14:40 -07:00
wozeparrotandGitHub	849074f0db	gptoss: use fa swa (#17522 )	2026-08-13 10:39:22 -07:00
George HotzandGitHub	0252cb8fa7	remove anchors from CI flow (#17521 ) * ci: remove yaml anchors from test.yml for gitea actions compatibility Gitea Actions does not support YAML anchors/aliases, which causes the workflow to fail parsing. Replace the &linux/linux anchor with a plain runs-on: ubuntu-24.04. ci: keep runner selection, inline expression instead of anchors Instead of replacing the anchored runs-on with a plain ubuntu-24.04 (which drops the namespace-profile-tinygrad routing for collaborator PRs), inline the full ${{ }} expression at every job. No YAML anchors, works with runners that can't parse them (gitea runner), and identical behavior on GitHub Actions.	2026-08-13 10:17:22 -07:00
qazalandGitHub	cc6d33bde7	amd: fix ALL2ALL speed on amdgpu (gpt) (#17518 )	2026-08-13 16:24:45 +09:00
qazalandGitHub	16c5ff2490	add external_benchmark_all2all.py (#17507 ) * add external_benchmark_all2all.py * mv * more minimal * less * fix space	2026-08-13 15:50:43 +09:00
wozeparrotandGitHub	1b7f040984	fa: paas through window (#17517 )	2026-08-13 14:45:57 +08:00
qazalandGitHub	39d144546e	fix mxfp4 mem estimate (#17515 ) * add mem estimates * rename * move	2026-08-13 14:48:15 +09:00
George HotzandGitHub	e103fb2a10	more lil llm improvements (#17514 ) * more lil llm improvements * default float	2026-08-12 20:13:20 -07:00
George HotzandGitHub	2297118541	lil llm improvements (#17513 )	2026-08-12 19:29:32 -07:00
sirhcmandGitHub	cd6d0d6ee3	allow running QCOMCL compiler in docker (#17499 )	2026-08-12 22:27:55 -04:00
geohot	ff0cb28c21	skip slow whisper tests	2026-08-12 13:12:37 -07:00
qazalandGitHub	ed8297a102	kerenl opts test from nan in llama 8b (#17510 ) * all2all * nan * remove that * less * has_local * only the nan change here * use nice getitem syntax for INDEX * work * remove * even simpler	2026-08-13 04:07:30 +09:00
RaineandGitHub	de04781b36	simplify equivalent const max (#17505 ) * add const max folds * add regression test * move	2026-08-12 08:39:29 -07:00
nimlgenandGitHub	4f106ebe87	hcq2: enqueue speed (#17504 )	2026-08-12 13:08:53 +03:00
qazalandGitHub	04c271ac41	simplify digitalocean_mi350x (#17502 ) * simplify digitalocean_mi350x * no hardcoded rocm path	2026-08-12 16:00:04 +09:00
qazalandGitHub	2e5a9a4121	no hardcoded device names in test_sliced_buffer_function (#17501 )	2026-08-12 15:19:25 +09:00
qazalandGitHub	e1013a6356	llama: create dataset cache by default in dev_beam (#17500 )	2026-08-12 15:03:39 +09:00
wozeparrotandGitHub	f891f5ffd0	gptoss: route lm_head thru asm_gemm (#17497 )	2026-08-11 18:32:50 -07:00
George HotzandGitHub	3686a1758f	mac/rdma imports lazy (#17496 ) * mac/rdma imports lazy * ish * fixes	2026-08-11 17:23:11 -07:00
George HotzandGitHub	4a253db9b4	minor cleanups to improve import speed (#17495 ) * minor cleanups to improve import speed * dumb	2026-08-11 16:06:25 -07:00
sirhcmandGitHub	479ffb0cda	remove Ops.SLICE (#17492 )	2026-08-11 18:50:04 -04:00
nimlgenandGitHub	2b5018e86a	hcq2: fix debug 2 info (#17491 ) * hcq2: fix debug 2 info * x * x * x	2026-08-12 00:25:44 +03:00
George HotzandGitHub	e11df72e0f	notes from digitalocean_mi350x (#17494 ) * notes from digitalocean_mi350x * cleanup * revert non-doc changes on digitalocean_mi350x branch	2026-08-11 13:22:11 -07:00
nimlgenandGitHub	a8c84ab34e	hcq2: enable all multitesnor tests (#17490 )	2026-08-11 17:47:33 +03:00
nimlgenandGitHub	ffef35c53e	hcq2: fix deps (#17481 ) * hcq2: proper unmap * hcq2: fix deps * x * x	2026-08-11 16:22:43 +03:00
nimlgenandGitHub	55e4f9d4f3	hcq2: proper unmap (#17489 )	2026-08-11 15:41:28 +03:00
sirhcmandGitHub	0c6a2c7dd6	slice is just shrink (#17483 )	2026-08-10 23:37:17 -04:00
RaineandGitHub	ad2fdeae69	move WMMA pms to codegen (#17485 ) * move wmma pms to codegen * lint tabs	2026-08-10 17:05:21 -07:00
RaineandGitHub	115bf9940f	add kwargs to group (#17484 )	2026-08-10 17:04:38 -07:00
George HotzandGitHub	7edf80a48a	small changes from new rangeify + remove flaky tc tests (#17486 ) * small changes from new rangeify * remove test/opt/test_tensor_cores.py	2026-08-10 15:58:49 -07:00
qazalandGitHub	22722ea2e7	llama: correct optim_dtype for mxfp4 (#17482 )	2026-08-11 03:01:28 +09:00
RaineandGitHub	d41ca5e60f	Fix WMMA CI (#17479 ) * init * split into sub tests * trigger ci	2026-08-10 08:39:14 -07:00
nimlgenandGitHub	e29606f07e	hcq2: copy kernel (#17480 ) * hcq2: copy with kernel * test * x	2026-08-10 17:28:46 +03:00
nimlgenandGitHub	8611fe22a7	fix hevc (#17477 ) * hevc tests * x	2026-08-10 13:33:49 +03:00
qazalandGitHub	2821bd646f	late loss.to("CPU") in llama (#17476 ) * late loss.to("CPU") in llama * acc = 0	2026-08-10 17:31:50 +09:00
qazalandGitHub	44f1f45cd5	llama: custom silu kernels (#17462 ) * start by copying the C * uop kernel * cleanup tests * estimates is part of SPEC	2026-08-10 16:43:01 +09:00
George HotzandGitHub	566f32fe9f	move platform tests to platform.yml (#17475 ) * ci: split mac/windows/qcom-cl tests into platform.yml Move the 6 jobs that don't run on Linux (4 macos, 1 windows, 1 QCOM CL compile test on arm) out of test.yml into a separate Platform Tests workflow so they run (and can be gated/runners-matched) independently. * ci: gate platform tests to the upstream repo Skip mac/windows/qcom-cl jobs anywhere but tinygrad/tinygrad, so the Platform Tests workflow is disabled on the gitea fork (and any fork). * ci: revert repo gate on platform tests Job-level if is only evaluated by gitea when a runner with matching labels fetches the task; with no mac/windows/arm runners the jobs queue forever. Disable the workflow on the instance instead.	2026-08-09 23:25:18 -07:00
nimlgenandGitHub	8c8b43de62	hcq2: fix beam (#17467 ) * fix beam * x	2026-08-09 16:53:47 +03:00
nimlgenandGitHub	e17c21e102	hcq2: timings (#17464 ) * hcq2: timings * Dx * x * x * x * x * align * x	2026-08-08 22:00:32 +03:00
George HotzandGitHub	d4d537c8ae	add SPEC checking for the kernel graph (#17432 ) * add SPEC checking for the kernel graph * skip test * raise there, not None * handwritten * issue with unshard * multi works * and bitcast * fix new tests * fix linear * remove call index * fix shrink * fixes	2026-08-08 10:00:06 -07:00
b1tgandGitHub	abe2256299	fix symbolic sharded reshape (#17463 )	2026-08-08 09:18:02 -07:00
b1tgandGitHub	8c49a7a34b	support symbolic shapes in copy (#17461 ) * pad_to is no-op when same shape * support symbolic shapes in copy	2026-08-08 09:16:59 -07:00
qazalandGitHub	9dd3b8402e	default llama 8b to MXFP4=1 (#17465 )	2026-08-09 00:13:32 +08:00
sirhcmandGitHub	c0d2f9ac0c	nolocals supports variables (#17457 )	2026-08-07 17:54:38 -04:00
nimlgenandGitHub	4c206a52b1	fix ci emu (gpt) (#17437 ) * fix ci emu * x	2026-08-07 22:36:59 +03:00
chenyuandGitHub	4a3b8f6501	better _drop_valid_stmts [pr] (#17454 )	2026-08-07 15:18:35 -04:00
chenyuandGitHub	59b88ea5e2	move pm_fold_cast_const [pr] (#17453 ) move to lower index dtype	2026-08-07 13:31:56 -04:00
chenyuandGitHub	f76422b8af	fix cast to float _min_max [pr] (#17451 )	2026-08-07 11:56:19 -04:00
wozeparrotandGitHub	1827ec57f7	gptoss: fix sharded invalids (#17450 )	2026-08-07 08:42:02 -07:00
nimlgenandGitHub	b6189db8e9	cpu: fix eintr (#17449 )	2026-08-07 17:46:57 +03:00
chenyuandGitHub	73e670c10f	c0+x<c1 -> x < c1-c0 is ints only [pr] (#17448 )	2026-08-07 10:44:01 -04:00
chenyuandGitHub	fca695a36f	clean up reduce MUL gradient (#17447 )	2026-08-07 10:08:04 -04:00
Robert JosephandGitHub	0c96cdc300	fix prod gradients at zero (#17404 )	2026-08-07 09:56:52 -04:00
chenyuandGitHub	baa6148066	fix var of large half input (#17444 ) * fix var of large half input similar to mean, we use sum_acc_dtype for denominator * mypy	2026-08-06 23:14:32 -04:00
George Hotz GitHub qazal	1858f1fd9a	viz: collapse PROGRAM nodes like CALL (codex) (#17438 ) Co-authored-by: qazal <[email protected]>	2026-08-07 11:44:34 +09:00
qazal GitHub geohot	f253c4469d	remove contiguous from custom_kernel (#17149 ) * no user contig on custom_kernel * clean up * non removable * test MXFP4 llama without hipcc * use compiler-free HIPCC renderer in llama CI * move llama coverage to AMD tests * run llama coverage in AMD test matrix * respect configured ROCm path in llama profile * work * clean up * fix * add views back * remove that * update test * test_double_permute one less kernel * test_shrink less kernels --------- Co-authored-by: George Hotz <[email protected]>	2026-08-07 11:20:59 +09:00
chenyuandGitHub	28195d51fb	fix f2f from fp8e5m2fnuz to half (#17442 ) * fix f2f from fp8e5m2fnuz to half * it works if it's supported	2026-08-06 21:19:46 -04:00
chenyuandGitHub	9020a88f03	truncate float in DType.const [pr] (#17439 )	2026-08-06 20:19:08 -04:00
chenyuandGitHub	d8cbc11105	update linear interpolate to use int math for indices (#17441 )	2026-08-06 20:18:39 -04:00
wozeparrotandGitHub	1fd6b1035f	fa: swa support (#17367 )	2026-08-06 08:07:30 -07:00
nimlgenandGitHub	46230e9f17	hcq2: fence inputs (#17436 )	2026-08-06 16:27:12 +03:00
qazalandGitHub	9636dd1a25	test MXFP4 llama without hipcc (#17435 ) * test MXFP4 llama without hipcc * first pythonpath then dev	2026-08-06 17:31:40 +09:00
qazalandGitHub	f258708d7d	llama: custom quantize_mxfp4+transpose kernel (codex) (#17434 ) * llama: custom quantize_mxfp4+transpose kernel (codex) * rename to cpp * inline * cleanup * lds load_bf16x4 * more tests, add Estimates	2026-08-06 16:13:28 +09:00
chenyuandGitHub	28e6ef6937	fix postopt symbolic [pr] (#17433 ) REDUCE with src simplied to const would become unparented	2026-08-06 00:12:30 -04:00
chenyuandGitHub	969df866a3	one less strong dtype const in symbolic [pr] (#17431 )	2026-08-05 23:35:06 -04:00
chenyuandGitHub	7a9cd8e329	move weak function and pm to uop/weak [PR] (#17429 )	2026-08-05 22:27:40 -04:00
George HotzandGitHub	b4372df9c6	revert wrong custom kernel fix (#17427 )	2026-08-05 18:31:35 -07:00
chenyuandGitHub	d51e55aa17	remove some pm_fold_cast_const [pr] (#17426 )	2026-08-05 21:28:38 -04:00
sirhcmandGitHub	be25207a7a	scope variable names inside CALLs (#17424 )	2026-08-05 20:59:09 -04:00
chenyuandGitHub	d726e5f7f3	split pm_fold_cast_const [PR] (#17425 ) need to delete this rule that writes strong typed CONST	2026-08-05 19:56:04 -04:00
George HotzandGitHub	470c032a5e	fix slice + non contig kernels (#17423 ) * movement: resolve negative int slice bounds against symbolic sizes negative int bounds in a slice against a symbolic dim were passed through unresolved, giving wrong views. resolve them against the (possibly symbolic) size, like slice.indices does for int dims * schedule: realize custom kernel inputs that don't resolve to a buffer state rangeify assigns ranges backward from consumers and CALL contributes none, so the subgraph above a custom kernel input gets no ranges unless something in it is realized, and reduce conversion crashes with a KeyError. realize call inputs that don't resolve to a buffer state. only view-only movement ops preserve the underlying buffer: anything computed (ALU, REDUCE, ...) must be realized even if one of its sources resolves to a buffer, since the whole subgraph above the call has no ranges. unwrapping src[0] unconditionally missed const branches hanging off non-src[0] children and silently resolved REDUCEs to their source buffer. includes regression tests for pure const, mixed buffer+const, and view-over-buffer inputs	2026-08-05 16:18:17 -07:00
geohot	a8a8030bc9	add benchmark_llm script	2026-08-05 15:59:51 -07:00
George HotzandGitHub	581bfdd94f	merge track_rewrites and profile_matches into rewrite_group [PR] (#17420 ) * merge track_rewrites and profile_matches into rewrite_group * bug * flip ctx polarity	2026-08-05 14:41:38 -07:00
chenyuandGitHub	07ac911665	few weak and decomp tweaks [PR] (#17419 )	2026-08-05 15:55:04 -04:00
chenyuandGitHub	c2f1e5ae2a	fix weak cast to strong dtype [pr] (#17418 ) weak can mean higher than that strong dtype, so always use that strong dtype is wrong	2026-08-05 15:32:23 -04:00
George HotzandGitHub	757a727808	move callify into tensor (#17416 )	2026-08-05 11:48:56 -07:00
George HotzandGitHub	2cce85a606	chat: display reasoning_content from streamed responses (#17414 ) * chat: display reasoning_content from streamed responses The server's StreamRouter emits reasoning_content deltas for think blocks, but the chat UI was only reading delta.content, silently dropping all reasoning. Now reasoning is shown in gray (#888) and included in the message history sent back to the server. * fix	2026-08-05 10:49:24 -07:00
nimlgenandGitHub	9b27ea8523	hcq2: cleaner (#17413 ) * hcq2: cleaner * x	2026-08-05 19:45:04 +03:00
chenyuandGitHub	6cb419b9b7	regression test for bert nan with weak (#17412 )	2026-08-05 12:14:26 -04:00
nimlgenandGitHub	5b0b68ec55	remove debug from test (#17410 )	2026-08-05 15:47:40 +03:00
qazalandGitHub	ad32bd272b	viz/cli: faster and more complete rewrites print (#17411 ) * viz/cli: faster and more complete matches print * kwargs	2026-08-05 20:02:27 +09:00
nimlgenandGitHub	874d33128b	hcq2 benchmark (#17235 ) * hcq2 in ci? * fix * traning * x * x * x * recover * debug * impler * x * x * x * hcq2: group input scatter plans by destination * hcq2: simplify input scatter tables * x	2026-08-05 10:00:42 +03:00
chenyuandGitHub	3bf9e70b19	Revert "don't cast weak in _broadcasted [pr] (#17408 )" (#17409 ) This reverts commit `b45058b5ec`.	2026-08-05 02:40:33 -04:00
Raine GitHub George Hotz	77e124e455	fix AMD WMMA emulation and test in CI (#17184 ) * fix SPEC=1 test_tensor_cores * implement i32 WMMA for RDNA3, add regression test to CI * gfx950 scaled mfma llvmir fix * detect VOP3PX2 in emu * start cdna4 scaled mfam emu * fix gfx12 llvmir signatures * oops * fix * fix src2 const field extraction (gpt) * scaled mfma fixes cdna4 (kimi) * fp8 out breaks mfma (glm) * fix const signature * another * fix f string for linter * lint * clean * and a final lint * . * fix mypy * skip slow tests on ci * reduce unroll tensor shape -> 64x64 --------- Co-authored-by: George Hotz <[email protected]>	2026-08-04 23:25:32 -07:00
chenyuandGitHub	b45058b5ec	don't cast weak in _broadcasted [pr] (#17408 ) * don't cast weak in _broadcasted [pr] * fine now?	2026-08-05 02:19:31 -04:00
George HotzandGitHub	46f0003776	more KernelCountException (#17407 )	2026-08-04 22:55:38 -07:00