vf MI350X passthrough mode driver

mac/rdma imports lazy (#17496 )
* mac/rdma imports lazy * ish * fixes
2026-08-11 19:01:36 -07:00 · 2026-08-11 17:23:11 -07:00 · 2026-08-11 16:06:25 -07:00 · 2026-08-11 18:50:04 -04:00 · 2026-08-12 00:25:44 +03:00 · 2026-08-11 13:22:11 -07:00
92 changed files with 2049 additions and 833 deletions
@@ -0,0 +1,213 @@
+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
+
+
+  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
@@ -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: |
@@ -521,12 +521,11 @@ 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)
@@ -585,8 +584,11 @@ 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 test/opt/test_tensor_cores.py --durations=20
+        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM --durations=20
      - name: Run disk copy tests
        run: python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk
      - name: Run TRANSCENDENTAL math
@@ -629,165 +631,6 @@ 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:
@@ -824,33 +667,3 @@ jobs:
        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,
@@ -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))

@@ -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
@@ -7,7 +7,7 @@ from tinygrad.runtime.support.hcq import FileIOInterface
 from tinygrad.runtime.support.am.amdev import AMDev

 if __name__ == "__main__":
-  gpus = System.pci_scan_bus(0x1002, [(0xffff, [0x74a1, 0x75a0])])
+  gpus = System.pci_scan_bus(0x1002, [(0xffff, [0x74a1, 0x75a0, 0x75b0])])
  for gpu in gpus:
    drv_path = f"/sys/bus/pci/devices/{gpu}/driver"
    if FileIOInterface.exists(drv_path) and os.path.basename(os.readlink(drv_path)) == "amdgpu":
@@ -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)
@@ -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

@@ -125,6 +126,25 @@ def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp,
  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 +157,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 +381,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 +436,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:
@@ -288,15 +288,17 @@ 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)
+    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_unmap(self, buf:HCQ2Buffer): self.dev.iface.unmap(buf)
+
@dataclass
 class AMDQueueDesc:
  ring: Buffer; read_ptr: Buffer; write_ptr: Buffer; doorbell: Buffer; put_value: Buffer  # noqa: E702
@@ -388,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)
@@ -468,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
@@ -539,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

@@ -578,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)

@@ -691,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,235 @@
+# 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 numpy tqdm wandb tiktoken sentencepiece
+```
+
+### 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 /root/datasets/c4-8b
+rclone copy mlc-training:mlcommons-training-wg-public/llama3_1/datasets/c4/llama3_1_8b/ /root/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`
+
+### Symlink for the submission script
+The `dev_run.sh` script hardcodes `BASEDIR="/raid/datasets/c4-8b/"`. Symlink:
+```bash
+mkdir -p /raid/datasets
+ln -s /root/datasets/c4-8b /raid/datasets/c4-8b
+```
+
+## Phase 4: wandb Login
+```bash
+wandb login
+```
+Enter API key from https://wandb.ai/authorize
+
+## Phase 5: Run Training
+
+### 5.1 Smoke test (beam search, 2 layers, fake data)
+Always run beam first to validate the pipeline:
+```bash
+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 BASEDIR=/root/datasets/c4-8b/ \
+  bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
+```
+
+### 5.2 Full training run
+```bash
+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 BASEDIR=/root/datasets/c4-8b/ \
+WANDB=1 \
+  bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh
+```
+
+## 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 |
+| `BASEDIR` | `/root/datasets/c4-8b/` | Where C4 dataset was downloaded (script hardcodes `/raid/datasets/c4-8b/`) |
+| `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.
+
+### `FileNotFoundError: '/raid/datasets/c4-8b/...'`
+Script hardcodes `BASEDIR`. Either symlink or edit the script:
+```bash
+mkdir -p /raid/datasets && ln -s /root/datasets/c4-8b /raid/datasets/c4-8b
+```
+
+### `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`.
+
+### PCI device ID
+`lspci -v` shows device ID `0x75b0` and subsystem ID `0x75a0`:
+```
+83:00.0 Processing accelerators: ... Device 75b0
+    Subsystem: ... Device 75a0
+```
+tinygrad's `PCIIface` in `ops_amd.py` and `hive_reset.py` did not list `0x75b0`, so the GPU was not found. Adding `0x75b0` to the device ID list in both files fixes the detection.
+
+### 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.
+
+### `/dev/kfd`
+`/dev/kfd` exists when amdgpu is loaded. Opening it returns `OSError: [Errno 22] Invalid argument`.
+
+### VRAM BAR reads all 0xFF
+After amdgpu initializes the GPU and is then unbound, reading the VRAM BAR (via `/sys/bus/pci/devices/0000:83:00.0/resource0`) returns all `0xFF` at all offsets — including the discovery table at `vram_size - 64KB`. tinygrad's `AMDev._run_discovery()` fails with `AssertionError: discovery signatures mismatch`.
+
+A PCI reset (`echo 1 > /sys/bus/pci/devices/0000:83:00.0/reset`) did not change the VRAM contents — still all `0xFF`.
+
+VRAM was also all `0xFF` when read via `/dev/mem` at the BAR physical address (`0xa0000000000`).
+
+### VFIO attempt
+Bound the GPU to `vfio-pci` with `enable_unsafe_noiommu_mode=1`. The GPU bound successfully and `/dev/vfio/noiommu-0` appeared. Running tinygrad with `VFIO=1` still failed with the same `discovery signatures mismatch` — VRAM BAR still reads all `0xFF`.
+
+### No IOMMU in guest
+`dmesg` has no `AMD-Vi` entries. PCI devices have no `iommu_group` symlink.
+
+### 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.
+
+### Current status: NOT WORKING
+tinygrad's `PCIIface` finds the GPU (after adding `0x75b0`) but `AMDev._run_discovery()` fails because the VRAM discovery table reads all `0xFF`. This was observed with the GPU unbound from any driver, after PCI reset, and with VFIO bound.
@@ -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
@@ -206,10 +206,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,10 +237,11 @@ 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)))
@@ -268,10 +270,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 +301,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
@@ -419,7 +423,8 @@ class TestCustomKernel(unittest.TestCase):
    GlobalCounters.reset()
    y = run(x[0]).realize()
    # it's copying the input and the output
-    assert_kernel_count(1)
+    # TODO: subbuffer usage has runtime specific behavior, this will be fixed after the removal of SLICE.
+    assert_kernel_count(2 if y.device in ("CL", "WEBGPU") else 1)
    self.assertEqual(y.tolist(), [1, 2, 3, 4])

  @Context(DEV="CPU")
@@ -429,12 +434,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 +492,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())
@@ -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()
@@ -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))
@@ -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()
@@ -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):
@@ -1150,6 +1150,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)
@@ -1,6 +1,6 @@
 import unittest, itertools, math
 from tinygrad import Tensor, dtypes, Context
-from tinygrad.dtype import DType, ConstType
+from tinygrad.dtype import DType, ConstType, truncate
 from tinygrad.uop.ops import Ops, UOp
 from test.helpers import full_rewrite
 import numpy as np
@@ -51,6 +51,17 @@ class TestWeakConstFolding(unittest.TestCase):
  def test_invalid_poison(self):
    self.assertTrue(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().is_invalid)

+  def test_cast_commits_to_dtype_grid(self):
+    # committing a weak const to a stated width puts the value on that width's grid, same as storage packing and native compilers
+    v = 1/123008  # not representable in float16
+    out = UOp.const(v).cast(dtypes.half).simplify()
+    self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.half, truncate[dtypes.half](v)))
+    self.assertNotEqual(out.val, v)
+    # the grid commit preserves the sign of zero
+    self.assertEqual(math.copysign(1, UOp.const(-0.0).cast(dtypes.half).simplify().val), -1)
+    # observable at tensor level: the const-folded comparison agrees with the committed value
+    self.assertTrue((Tensor(-3.2).cast(dtypes.float32) <= truncate[dtypes.float32](-3.2)).item())
+
 class TestBinaryOpsConstFolding(unittest.TestCase):
  def test_add_literal_zero(self):
    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + 0)
@@ -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
@@ -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)),
@@ -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
@@ -6,7 +6,7 @@ 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):
@@ -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)
@@ -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
@@ -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):
@@ -185,10 +186,10 @@ class TestTensorCores(unittest.TestCase):
    # 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(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in)
+    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, 2)]
-    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
+    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
@@ -199,10 +200,10 @@ 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(self):
    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(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in)
+    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, 2)]
-    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
+    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]))
@@ -215,10 +216,10 @@ class TestTensorCores(unittest.TestCase):
  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 and tc.dtype_in not in dtypes.fp8s][0]
-    x, y = Tensor.rand(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in)
+    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, 2)]
-    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
+    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()
@@ -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
@@ -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
@@ -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, rewrite_group, 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
@@ -301,7 +301,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
    sink = graph_rewrite(sink, pm_split_ranges+pm_flatten_range, ctx={}, name="split ranges")

    # symbolic (NOTE: this is a requirement for pm_simplify_ranges to be correct)
-    sink = graph_rewrite(sink, sym+pm_flatten_range, name="initial symbolic")
+    sink = graph_rewrite(sink, sym+pm_fold_cast_const+pm_flatten_range, name="initial symbolic")

    # optimize (schedule) the AST
    sink = graph_rewrite(sink, pm_flatten_range+pm_simplify_ranges, ctx={}, name="simplify ranges")
@@ -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")
@@ -354,7 +357,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:

  # floordiv+mod / dtype decomp (early)
  supported_ops = tuple(ren.code_for_op.keys())
-  pm_decomp = symbolic_simple+get_simplifying_rewrite_patterns(supported_ops)
+  pm_decomp = symbolic_simple+pm_fold_cast_const+get_simplifying_rewrite_patterns(supported_ops)
  sink = graph_rewrite(sink, pm_decomp, name="early decompositions")

  # late decomps + move gates from unrenderable INVALID where
@@ -99,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)
@@ -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

@@ -27,11 +27,14 @@ def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]:
    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
+      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:
@@ -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,
@@ -1,7 +1,7 @@
 import itertools
 from typing import Callable
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start, AxisType
-from tinygrad.uop.symbolic import symbolic, invalid_gate
+from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const, invalid_gate
 from tinygrad.helpers import partition
 from tinygrad.dtype import dtypes

@@ -32,7 +32,7 @@ def simplify_merge_adjacent(u:UOp) -> UOp|None:
        s0, s1 = r0.src[0], r1.src[0]
        # do the merge
        new_range = r0.replace(src=(s0*s1,))
-        nidx = graph_rewrite(u, _substitute+symbolic+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
+        nidx = graph_rewrite(u, _substitute+symbolic+pm_fold_cast_const+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
                             name=f"check_merge_{r0.arg[0]}_{r1.arg[0]}")

        # check if it simplifies
@@ -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:
@@ -44,8 +44,6 @@ 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,))])
    graph_t = graph_class(devs[0]) if devs[0].graph is not None else None

@@ -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,12 +3,12 @@ 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 ****************

@@ -17,7 +17,7 @@ def get_call_arg_uops(call:UOp) -> tuple[UOp, ...]: return tuple(s for s in call
 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,9 +27,6 @@ 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")
@@ -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: cast(Any, Device[d]).synchronize(timeout=ctx.timeout)
-  return time.perf_counter() - st
+  tms = []
+  for devices,name,estimates,prof in info.kernels:
+    for device in devices:
+      d, tm = cast(Any, Device[device]), None
+      if prof:
+        d.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)
+      with track_stats(ctx, call.replace(arg=replace(call.arg, name=name, aux=replace(info, estimates=estimates))), d.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,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)
@@ -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:
@@ -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`.
@@ -7,7 +7,11 @@ from tinygrad.dtype import sum_acc_dtype
 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)."""
@@ -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]
@@ -279,43 +279,9 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n  %trunc
      (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 += 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),
-      ])
-    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(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),
-      ])
-    if target.arch in {"gfx1200", "gfx1201"}:
-      self.extra_matcher += 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)
-      ])
+    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}
@@ -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)
@@ -842,7 +842,7 @@ class KFDIface:

 class PCIIface(PCIIfaceBase):
  def __init__(self, dev, dev_id):
-    super().__init__(dev, dev_id, vendor=0x1002, devices=((0xffff, (0x74a1,0x744c,0x7480,0x7550,0x7551,0x7590,0x75a0)),), vram_bar=0,
+    super().__init__(dev, dev_id, vendor=0x1002, devices=((0xffff, (0x74a1,0x744c,0x7480,0x7550,0x7551,0x7590,0x75a0,0x75b0)),), vram_bar=0,
      va_start=AMMemoryManager.va_allocator.base, va_size=AMMemoryManager.va_allocator.size, dev_impl_t=AMDev)
    self._compute_props()

@@ -1096,6 +1096,11 @@ class AMDDevice(HCQCompiled):

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

+  def finalize(self):
+    try: super().finalize()
+    finally:
+      if self.is_am(): self.iface.dev_impl.release_vf_access()
+
  def device_props(self): return self.iface.props

  def hw_copy_queues(self): return [(f"SDMA:{i}", functools.partial(unwrap(self.hw_copy_queue_t), queue_idx=i)) for i in self.sdma_queues]
@@ -1,5 +1,5 @@
 from __future__ import annotations
-import platform, sys, os, ctypes, functools, mmap, threading, array, itertools
+import platform, sys, os, ctypes, ctypes.util, 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
@@ -50,7 +50,7 @@ def worker_prog():

  # 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)
@@ -1,5 +1,5 @@
 from __future__ import annotations
-import ctypes, collections, dataclasses, functools, hashlib, array
+import ctypes, collections, dataclasses, functools, hashlib, array, time, contextlib
 from tinygrad.helpers import mv_address, getenv, DEBUG, lo32, hi32, fetch_fw
 from tinygrad.runtime.autogen import pci
 from tinygrad.runtime.autogen.am import am, fw
@@ -149,6 +149,14 @@ class AMDev:
    self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
    self.vram, self.doorbell64, self.mmio = self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I')

+    # MI350X VFs start with most MMIO and VRAM access gated by the host PF. Ask the PF for full access only when discovery isn't readable yet.
+    self.is_vf = bool(self.mmio[0xde5] & 1) # RCC_IOV_FUNC_IDENTIFIER.FUNC_IDENTIFIER
+    self.vf_access_acquired, self.vf_initialized = False, False
+    if self.is_vf:
+      self._vf_mailbox_request(6, 7, data2=2, retries=5, event_timeout=2) # IDH_REQ_GPU_INIT_DATA -> IDH_REQ_GPU_INIT_DATA_READY
+      self._vf_mailbox_request(1, 1, retries=5, event_timeout=2) # IDH_REQ_GPU_INIT_ACCESS -> IDH_READY_TO_ACCESS_GPU
+      self.vf_access_acquired = True
+
    self._run_discovery()
    self._build_regs()

@@ -165,21 +173,21 @@ class AMDev:
    self.is_booting = True # During boot only boot memory can be allocated. This flag is to validate this.
    self.init_sw(smi_dev=False)

-    self.partial_boot = (self.reg("regSCRATCH_REG7").read() == AMDev.Version) and (getenv("AM_RESET", 0) != 1)
+    self.partial_boot = not self.is_vf and (self.reg("regSCRATCH_REG7").read() == AMDev.Version) and (getenv("AM_RESET", 0) != 1)
    if self.partial_boot and (self.reg("regSCRATCH_REG6").read() != 0 or self.reg(self.gmc.pf_status_reg("GC")).read() != 0):
      if DEBUG >= 2: print(f"am {self.devfmt}: Malformed state. Issuing a full reset.")
      self.partial_boot = False

-    # Init hw for IP blocks where it is needed
+    # Init hw for IP blocks where it is needed. PSP and SMU are PF-owned on a VF and must not be reset or reloaded by the guest.
    if not self.partial_boot:
-      if self.psp.is_sos_alive() and self.smu.is_smu_alive():
+      if not self.is_vf and self.psp.is_sos_alive() and self.smu.is_smu_alive():
        self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) & ~pci.PCI_COMMAND_MASTER, 2)
        if self.is_hive():
          if reset_mode: return # in reset mode, do not raise
          raise RuntimeError("Malformed state. Use extra/amdpci/hive_reset.py to reset the hive")
        self.smu.mode1_reset()
      self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
-      self.init_hw(self.soc, self.gmc, self.ih, self.psp, self.smu)
+      self.init_hw(self.soc, self.gmc, self.ih, *(() if self.is_vf else (self.psp, self.smu)))

    # Booting done
    self.is_booting = False
@@ -187,13 +195,17 @@ class AMDev:
    # Re-initialize main blocks
    self.init_hw(self.gfx, self.sdma)

-    if (max_power:=getenv("AM_POWER_LIMIT", 0.0)) > 0:
-      self.smu.set_power_limit(max_power)
-      self.smu.set_clocks(level=None)
-    else: self.smu.set_clocks(level=-1) # last level, max perf.
-    for ip in [self.soc, self.gfx]: ip.set_clockgating_state()
-    self.reg("regSCRATCH_REG7").write(AMDev.Version)
-    self.reg("regSCRATCH_REG6").write(1) # set initialized state.
+    if not self.is_vf:
+      if (max_power:=getenv("AM_POWER_LIMIT", 0.0)) > 0:
+        self.smu.set_power_limit(max_power)
+        self.smu.set_clocks(level=None)
+      else: self.smu.set_clocks(level=-1) # last level, max perf.
+    if not self.is_vf:
+      for ip in [self.soc, self.gfx]: ip.set_clockgating_state()
+    if not self.is_vf:
+      self.reg("regSCRATCH_REG7").write(AMDev.Version)
+      self.reg("regSCRATCH_REG6").write(1) # set initialized state.
+    self.vf_initialized = self.is_vf
    if DEBUG >= 2: print(f"am {self.devfmt}: boot done")

  def init_sw(self, smi_dev=False):
@@ -202,7 +214,8 @@ class AMDev:
    # Memory manager & firmware
    self.mm = AMMemoryManager(self, self.vram_size - self.reserved_vram_size, boot_size=(32 << 20), pt_t=AMPageTableEntry, va_shifts=[12, 21, 30, 39],
      va_bits=48, first_lv=am.AMDGPU_VM_PDB2, va_base=AMMemoryManager.va_allocator.base, reserve_ptable=not self.large_bar,
-      palloc_ranges=[(1 << (i + 12), (2 << 20) if i >= 9 else 0x1000) for i in range(9 * (3 - am.AMDGPU_VM_PDB2), -1, -1)])
+      palloc_ranges=[(1 << (i + 12), (2 << 20) if i >= 9 else 0x1000) for i in range(9 * (3 - am.AMDGPU_VM_PDB2), -1, -1)],
+      paddr_base=(1 << 20) if self.is_vf else 0)
    self.fw = AMFirmware(self)

    # Initialize IP blocks
@@ -224,10 +237,52 @@ class AMDev:

  def fini(self):
    if DEBUG >= 2: print(f"am {self.devfmt}: Finalizing")
-    for ip in [self.sdma, self.gfx]: ip.fini_hw()
-    self.smu.set_clocks(level=0)
-    self.ih.interrupt_handler()
-    self.reg("regSCRATCH_REG6").write(self.is_err_state) # set finalized state.
+    try:
+      for ip in [self.sdma, self.gfx]: ip.fini_hw()
+      if not self.is_vf: self.smu.set_clocks(level=0)
+      self.ih.interrupt_handler()
+      if not self.is_vf: self.reg("regSCRATCH_REG6").write(self.is_err_state) # set finalized state.
+    finally: self.release_vf_access()
+
+  def release_vf_access(self):
+    if not getattr(self, "vf_access_acquired", False): return
+    # tinygrad retains IDH_REQ_GPU_INIT_ACCESS for direct MMIO/VRAM access, so always release that same lease.
+    with contextlib.suppress(Exception): self._vf_mailbox_request(2, None) # IDH_REL_GPU_INIT_ACCESS
+    self.vf_access_acquired = False
+
+  def __del__(self):
+    # Constructor failures do not reach HCQ finalization; return a partially acquired VF init lease to the PF.
+    self.release_vf_access()
+
+  def _vf_mailbox_request(self, req:int, event:int|None, data1=0, data2=0, data3=0, retries=1, event_timeout=2.0):
+    # Navi VF/PF mailbox protocol from the kernel's mxgpu_nv driver. This requests access only; it never requests a GPU or PCI reset.
+    mmio8, control, trn, rcv = self.mmio.view(fmt='B'), 0xe5e * 4, 0xe56, 0xe5a
+    if mmio8[control+1] & 1: mmio8[control+1] = 2 # acknowledge a stale PF event before transmitting a new request
+    for retry in range(retries):
+      deadline = time.monotonic() + 1
+      while True:
+        mmio8[control] = 0 # clear TRN_MSG_VALID and wait for the old PF acknowledgement to drop
+        if not (mmio8[control] & 2): break
+        if time.monotonic() > deadline: raise TimeoutError("VF mailbox acknowledgement did not clear")
+        time.sleep(0.001)
+
+      for i, val in enumerate((req, data1, data2, data3)): self.mmio[trn+i] = val
+      mmio8[control] = 1
+      deadline = time.monotonic() + 0.5
+      while not (mmio8[control] & 2):
+        if time.monotonic() > deadline: raise TimeoutError(f"VF mailbox request {req:#x} was not acknowledged")
+        time.sleep(0.005)
+      mmio8[control] = 0
+
+      if event is None: return
+      deadline = time.monotonic() + event_timeout
+      while time.monotonic() <= deadline:
+        if self.mmio[rcv] == event:
+          mmio8[control+1] = 2 # acknowledge RCV_MSG_VALID
+          return
+        time.sleep(0.01)
+      if DEBUG >= 2 and retry+1 < retries: print(f"am {self.devfmt}: retrying VF mailbox request {req:#x} ({retry+1}/{retries})")
+    raise TimeoutError(f"VF mailbox request {req:#x} did not receive event {event:#x}")

  def recover(self, force=False) -> bool:
    if not force and not self.is_err_state: return False
@@ -251,9 +251,10 @@ class AM_GFX(AM_IP):
    self.mqd_mc = [self.adev.paddr2mc(mqd_paddr) for mqd_paddr in self.mqd_paddr]

  def init_hw(self):
-    # Wait for RLC autoload to complete
-    wait_cond(lambda: self.adev.regCP_STAT.read() == 0 or self.adev.regRLC_RLCS_BOOTLOAD_STATUS.read_bitfields()['bootload_complete'] == 0,
-              value=True, msg="RLC autoload timeout")
+    # Wait for RLC autoload to complete on architectures that expose the bootload status register.
+    if hasattr(self.adev, "regRLC_RLCS_BOOTLOAD_STATUS"):
+      wait_cond(lambda: self.adev.regCP_STAT.read() == 0 or self.adev.regRLC_RLCS_BOOTLOAD_STATUS.read_bitfields()['bootload_complete'] == 0,
+                value=True, msg="RLC autoload timeout")

    self.adev.gmc.init_hub("GC", inst_cnt=self.xccs)
    if self.adev.partial_boot: return self.reset_mec()
@@ -297,8 +298,8 @@ class AM_GFX(AM_IP):

    self._enable_mec()

-    # Set 1 partition
-    if self.xccs > 1: self.adev.psp._spatial_partition_cmd(1)
+    # Set 1 partition on bare metal. A VF must use the spatial partition assigned by its host PF.
+    if self.xccs > 1 and not self.adev.is_vf: self.adev.psp._spatial_partition_cmd(1)

  def fini_hw(self): self._dequeue_hqds()

@@ -486,7 +487,7 @@ class AM_IH(AM_IP):
    if athub_err or cntlr_err:
      print(f"am {self.adev.devfmt}: fatal hardware error detected: {'RAS_ATHUB_ERR_EVENT ' if athub_err else ''}{'RAS_CNTLR' if cntlr_err else ''}")

-      acas = self.adev.smu._aca_read_banks(ue=True) + self.adev.smu._aca_read_banks(ue=False)
+      acas = [] if self.adev.is_vf else self.adev.smu._aca_read_banks(ue=True) + self.adev.smu._aca_read_banks(ue=False)
      for regs in acas:
        acatyp = 'Uncorrectable' if (regs[1] >> 61) & 1 and (regs[1] >> 57) & 1 else 'Correctable'
        hwname = f'{self.adev.hwid_names.get((regs[5] >> 32) & 0xFFF, "")} ({(regs[5] >> 32) & 0xFFF:#03x})'
@@ -89,9 +89,11 @@ class HIPCompiler(Compiler):

 class HIPCCCompiler(Compiler):
  def __init__(self, arch:str, extra_options:list[str]=[]):
-    self.arch, self.extra_options = arch, extra_options
-    super().__init__(f"compile_hipcc_{self.arch}_{hashlib.sha256(' '.join(extra_options).encode()).hexdigest()[:8]}")
+    self.arch, self.extra_options, self.no_hipcc = arch, extra_options, getenv("NO_HIPCC")
+    super().__init__(f"compile_hipcc_{self.arch}_{hashlib.sha256(' '.join(extra_options).encode()).hexdigest()[:8]}"+
+                     ("_nohipcc" if self.no_hipcc else ""))
  def compile(self, src:str) -> bytes:
+    if self.no_hipcc: return b""
    with tempfile.NamedTemporaryFile(suffix=".cpp") as srcf, tempfile.NamedTemporaryFile(suffix=".bc") as bcf:
      with tempfile.NamedTemporaryFile(suffix=".hsaco") as libf:
        srcf.write(src.encode())
@@ -1,12 +1,13 @@
 from __future__ import annotations
 from typing import cast, Callable, TypeVar, Generic, Any, Sequence
-import struct, functools, time, collections, itertools
+import struct, functools, time, collections, itertools, decimal, statistics
 from dataclasses import replace, dataclass
-from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, dedup, pluralize, JIT_BATCH_SIZE, unwrap
-from tinygrad.helpers import to_tuple, round_up, partition, data64_le, panic, ContextVar
+from tinygrad.helpers import DEV, getenv, select_first_inited, select_by_name, suppress_finalizing, dedup, pluralize, JIT_BATCH_SIZE, unwrap, PROFILE
+from tinygrad.helpers import to_tuple, round_up, partition, data64_le, panic, ContextVar, perf_counter_us, Context
 from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer, DepsTracker
+from tinygrad.device import ProfileDeviceEvent, ProfileGraphEntry, ProfileGraphEvent
 from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, rewrite_group, GroupOp
-from tinygrad.uop.symbolic import symbolic
+from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const
 from tinygrad.dtype import dtypes, truncate
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support.memory import BumpAllocator
@@ -32,12 +33,13 @@ class HCQInfo:

  input_idxs:tuple[int, ...] = () # indexes into input_uops used by this call
  inputs:int|None = None
+  kernels:tuple[tuple[tuple[str, ...], str, Estimates, tuple[int, ...]], ...] = ()

 def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c

 def unwrap_mstack(u):
  if u.op is Ops.MSTACK: return tuple(x for s in u.src for x in unwrap_mstack(s))
-  return unwrap_mstack(u.src[0]) if u.op in {Ops.MSELECT, Ops.SLICE} else (u,)
+  return unwrap_mstack(u.src[0]) if u.op is Ops.MSELECT else (u,)

 def is_value_known_at_link(val:UOp) -> bool:
  runtime_reads = [u for u in val.toposort() if u.op in (Ops.LOAD, Ops.INDEX)]
@@ -94,12 +96,24 @@ pm_replace_buffers = PatternMatcher([(UPat(Ops.CALL, name="call"), replace_call_

 def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS) and not all_devices_in(b.device, HCQ_P2P_DEVS)

+def hcq_call_devs(call:UOp) -> Any|None: return next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None)
+
 def stage_copy(dst:UOp, src:UOp) -> UOp|None:
  if not (_need_staging(src, dst) or _need_staging(dst, src)): return None

  stage = UOp.new_buffer("CPU", src.max_numel() * src.dtype.itemsize, dtypes.uint8)
  return UOp(Ops.LINEAR, src=(src.copy_to_device("CPU").call(stage, src), stage.copy_to_device(dst.device).call(dst, stage)))
-pm_insert_copy_staging = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy)])
+
+def kernel_copy(call:UOp, dst:UOp, src:UOp) -> UOp|None:
+  if (devs:=hcq_call_devs(call)) is None or Device[(dev:=to_tuple(devs)[0])].has_copy_queue: return None
+  d, s = (UOp.param(i, dst.dtype, (n:=dst.max_numel(),), device=devs) for i in range(2))
+  ast = d.index(r:=UOp.range(n, 0)).store(s.index(r).load()).end(r).sink(arg=KernelInfo(name="copy"), tag=1)
+  return call.replace(src=(to_program(ast, Device[dev].renderer), dst, src))
+
+pm_insert_copy_staging = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy),
+  (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src")), name="call"), kernel_copy)
+])

 # *****************
 # 2. deps
@@ -127,14 +141,14 @@ def _build_wait_cmds(slots:dict[str, int], dep_lanes:list[tuple[tuple, int, int]

  # opt2: keep latest dep per (dep device, queue, cur lane)
  latest = {((dep[0][dlane], dep[1]), lane): (dep, dlane) for dep, dlane, lane in sorted(dep_lanes, key=lambda x: x[0][2])}
-  deps:dict[tuple, list[int|None]] = collections.defaultdict(lambda: [None]*len(devices))
-  for (_, lane), (dep, dlane) in latest.items(): deps[dep][lane] = dlane
+  deps:dict[tuple, dict[int, list[int]]] = collections.defaultdict(lambda: collections.defaultdict(list))
+  for (_, lane), (dep, dlane) in latest.items(): deps[dep][lane].append(dlane)

  waits = []
-  for (ddevs, dqueue, dtag), lanes in deps.items():
-    sig = UOp.mstack(*[make_signal(d, tag="sentinel_signal") if dl is None else make_signal(ddevs[dl], slots[dqueue])
-                       for dl, d in zip(lanes, devices)])
-    waits.append(UOp(Ops.INS, arg="wait", src=(sig, UOp.const(dtag + 1, dtypes.uint64))))
+  for (ddevs, dqueue, dtag), by_lane in deps.items():
+    for ls in itertools.zip_longest(*(by_lane[lane] for lane in range(len(devices)))):
+      s = UOp.mstack(*[make_signal(d, tag="sentinel_signal") if dl is None else make_signal(ddevs[dl], slots[dqueue]) for dl, d in zip(ls, devices)])
+      waits.append(UOp(Ops.INS, arg="wait", src=(s, UOp.const(dtag + 1, dtypes.uint64))))
  return waits, {dtag for _, _, dtag in deps}

 def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[tuple[tuple[str, ...], str]],
@@ -170,7 +184,7 @@ def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[t
    fins.append(make_call("hcq_finalizer", UOp.sink(epoch_slot.store(epoch + 1), sched_epoch.after(fin_submit).index(0).store(epoch)), HCQInfo(devs)))
  return fences, fins, signal_tags

-def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]]) -> list[UOp]:
+def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> list[UOp]:
  batch_info = [(devices, "COMPUTE:0" if call.src[0].op is Ops.PROGRAM else "COPY:0") for call, devices in batch]

  # schedule deps
@@ -185,10 +199,10 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]]) -> list[UOp]:
    signal_tags |= cur_signal_tags

  # build fences and finalizers
-  fences, finalizers, finalizer_signal_tags = _build_finalizers(batch, batch_info, deps_tracker, slots)
+  fences, fins, finalizer_signal_tags = _build_finalizers(batch, batch_info, deps_tracker, slots)
  signal_tags |= finalizer_signal_tags

-  src = []
+  src, kerns = [], []
  for tag, ((call, _), (devices, queue), q) in enumerate(zip(batch, batch_info, call_waits)):
    # first queue use, sync prior device work with the device timeline
    if batch_info.index((devices, queue)) == tag:
@@ -197,20 +211,27 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]]) -> list[UOp]:

    # and make hcq call
    name, info = get_call_name(call, get_call_arg_uops(call)), HCQInfo(devices, estimate_uop(call))
-    q += [call.replace(arg=replace(call.arg, aux=info))]
+    ts_ids = [next(UOp.unique_num) for _ in range(2)] if profile else []
+    kerns.append((devices, name, info.estimates, tuple(ts_ids)))
+
+    ts_ins = [UOp(Ops.INS, arg="timestamp", src=(make_signal(devices, s),)) for s in ts_ids]
+    q += ts_ins[:1] + [call.replace(arg=replace(call.arg, aux=info))] + ts_ins[1:]

    # signal the queue if someone waits for us
    if tag in signal_tags: q += [UOp(Ops.INS, arg="store", src=(make_signal(devices, slots[queue]), UOp.const(tag + 1, dtypes.uint64)))]
-    src.append(make_call(name, make_submit(*q, devs=devices, queue=queue).sink(), info))
-  return fences + src + finalizers
+    src.append(make_call(f"submit {name}", make_submit(*q, devs=devices, queue=queue).sink(), info))

-def sched_hcq_batches(l:UOp) -> UOp:
+  # append batch timestamps to finalizers
+  fins = [f.replace(arg=replace(f.arg, aux=replace(a:=f.arg.aux, kernels=tuple(x for x in kerns if set(x[0]) & set(a.device))))) for f in fins]
+  return fences + src + fins
+
+def sched_hcq_batches(l:UOp, profile:bool) -> UOp:
  srcs:list[UOp] = []
  batch:list[tuple[UOp, tuple[str, ...]]] = []
  for call in l.src:
-    if (devs:=next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None)) is not None: batch.append((call, to_tuple(devs)))
-    else: srcs, batch = srcs + _finalize_batch(batch) + [call], []
-  return l.replace(src=tuple(srcs + _finalize_batch(batch)))
+    if (devs:=hcq_call_devs(call)) is not None: batch.append((call, to_tuple(devs)))
+    else: srcs, batch = srcs + _finalize_batch(batch, profile) + [call], []
+  return l.replace(src=tuple(srcs + _finalize_batch(batch, profile)))

 # *****************
 # 3. merge into queues
@@ -246,7 +267,7 @@ def merge_queues(linear:UOp) -> UOp:
  return linear.replace(src=tuple(new_src + [_merged_hcq_call(c) for c in opened_qs.values()]))

 pm_schedule_and_merge = PatternMatcher([(UPat(Ops.LINEAR, name="l"),
-  lambda ctx, l: merge_queues(sched_hcq_batches(l).substitute(ctx, walk=True, enter_calls=True)))])
+  lambda ctx, l: merge_queues(sched_hcq_batches(l, ctx[1]).substitute(ctx[0], walk=True, enter_calls=True)))])

 # *****************
 # 4.2. hcq lowering: ops to ir
@@ -321,6 +342,10 @@ def split_patches(call:UOp) -> UOp|None:
  scatter = make_scatter_loops(input_patches, tables[0], lt_patches)
  body = body.substitute({p:p.substitute(scatter | reads) for p in rt_patches})

+  if inputs: # fence inputs
+    fills.append((t:=tables[0][0]).after(make_binary_patch(t, bytes(t.max_numel() * 8)))) # zeroed at link, slot 0 is the host fence
+    body = body.replace(src=(UOp.sink(*body.src[0].src, t.after(*body.src[0].src).index(0).store(0)),)) # open it once consumed
+
  lt_srcs = collections.defaultdict(list)
  for p in lt_patches: lt_srcs[p.buf_uop].append(p)
  return call.replace(src=(body, *call.src[1:], *[b.after(*ps) for b,ps in lt_srcs.items()], *fills),
@@ -344,7 +369,7 @@ def replace_params(call:UOp) -> UOp|None:

  sub = {(b:=u.without_after): UOp.param(i, u.dtype, shape=b.shape, device=HCQ_RUNTIME_DEV.value, volatile=b.op is Ops.PARAM and b.arg.volatile)
         for i,u in enumerate(c_args)} | {v: v.replace(arg=replace(v.arg, slot=-1)) for v in variables if v.op is Ops.PARAM}
-  info = replace(call.arg.aux, inputs=next((i for i,u in enumerate(c_args) if u.tag == "inputs"), None))
+  info = replace(call.arg.aux, inputs=next((i for i,u in enumerate(c_args) if u.without_after.tag == "inputs"), None))
  return call.replace(src=(body.substitute(sub).replace(arg="hcq_args"), *c_args, *refhold),
                      arg=replace(call.arg, aux=info)) # TODO: call.after(*refhold)?
 pm_replace_params = PatternMatcher([
@@ -352,14 +377,15 @@ pm_replace_params = PatternMatcher([

 # *****************

-def resolve_getaddr_slice(bv:UOp, g:UOp) -> UOp:
+def resolve_getaddr_view(bv:UOp, g:UOp) -> UOp:
  base = bv.src[0].after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ())
-  itemsize = bv.src[0].dtype.itemsize if bv.src[0].without_after.op in (Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize
+  if bv.op is Ops.BITCAST: return UOp(Ops.GETADDR, src=(base,), arg=g.arg)
+  itemsize = bv.src[0].dtype.itemsize if bv.src[0].without_after.op in (Ops.BUFFER, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize
  return UOp(Ops.GETADDR, src=(base,), arg=g.arg) + UOp.const(bv.src[1].val * itemsize, dtypes.uint64)

 pm_early_simplify = PatternMatcher([
-  (UPat(Ops.GETADDR, src=(UPat.any(sl:=UPat(Ops.SLICE, name="bv"), sl.after(allow_any_len=True)),), name="g"), resolve_getaddr_slice),
-  (UPat(Ops.INDEX, src=(UPat(Ops.SLICE, name="bv"),), allow_any_len=True, name="x"),
+  (UPat(Ops.GETADDR, src=(UPat((Ops.SHRINK, Ops.BITCAST), name="bv").or_after(),), name="g"), resolve_getaddr_view),
+  (UPat(Ops.INDEX, src=(UPat(Ops.SHRINK, name="bv"),), allow_any_len=True, name="x"),
   lambda bv,x: x.replace(src=(bv.src[0], x.src[1] + bv.src[1].cast(x.src[1].dtype), *x.src[2:]))),
 ])

@@ -377,7 +403,7 @@ def pack_hcq_placeholders(call:UOp) -> UOp|None:
      sizes[b.tag] = offs[b] + b.max_numel()
  counts = collections.Counter(b.tag for b in bufs)
  bases = {b.tag:UOp.placeholder((sizes[b.tag],), b.dtype, next(UOp.unique_num), device=b.device).rtag(b.tag) for b in bufs if counts[b.tag] > 1}
-  subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(offs.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases}
+  subs = {b:bases[b.tag][(off:=offs.get(b, 0)):off+b.max_numel()] for b in bufs if b.tag in bases}
  return call.replace(src=(call.src[0].substitute(subs, walk=True), *call.src[1:])) if subs else None
 pm_pack_placeholders = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), pack_hcq_placeholders)])
@@ -391,27 +417,27 @@ def callify_hcq(call:UOp, cf:UOp) -> UOp:
 pm_callify_hcq = PatternMatcher([(UPat(Ops.CALL, src=(
  UPat(Ops.CUSTOM_FUNCTION, arg="hcq_args", src=(UPat(Ops.SINK),), name="cf"),), name="call", allow_any_len=True), callify_hcq)])

-hcq_compile_cache:dict[bytes, UOp] = {}
+hcq_compile_cache:dict[tuple[bytes, bool], UOp] = {}

-@rewrite_group(lambda linear,input_uops,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
-def hcq_compile(linear:UOp, input_uops:list[UOp]|None=None) -> UOp:
+@rewrite_group(lambda linear,input_uops,profile,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
+def hcq_compile(linear:UOp, input_uops:list[UOp]|None, profile:bool) -> UOp:
  if input_uops is not None:
    slots = {u:i for i,u in reversed(tuple(enumerate(input_uops)))}
    linear = graph_rewrite(linear, pm_replace_buffers, ctx=(input_uops, slots), walk=True, name="replace buffer")

-  if (final_linear:=(hcq_compile_cache.get(cache_key:=linear.key))) is None:
+  if (final_linear:=(hcq_compile_cache.get(cache_key:=(linear.key, profile)))) is None:
    # prep
    linear = linear.substitute(back_map:={s.param_like(i): s for i,s in enumerate(input_uops)} if input_uops is not None else {}, walk=True)
    linear = graph_rewrite(linear, pm_insert_copy_staging+pm_flatten_linear, name="insert copy staging")

    # schedule
-    linear = graph_rewrite(linear, pm_schedule_and_merge, ctx={s:p for p,s in back_map.items()}, walk=True, name="schedule and merge hcq")
+    linear = graph_rewrite(linear, pm_schedule_and_merge, ctx=({s:p for p,s in back_map.items()}, profile), walk=True, name="schedule and merge hcq")

    # lowering to hcq ir
    linear = graph_rewrite(linear, pm_encode_cmdbufs+pm_pack_placeholders, walk=True, name="encode and pack", enter_calls=True)

    # patches and runtime uops
-    linear = graph_rewrite(linear, pm_early_simplify+symbolic, bottom_up=False, name="simplify patches", enter_calls=True)
+    linear = graph_rewrite(linear, pm_early_simplify+symbolic+pm_fold_cast_const, bottom_up=False, name="simplify patches", enter_calls=True)
    linear = graph_rewrite(linear, pm_split_patches, walk=True, name="split patches")

    # and compile it
@@ -468,7 +494,7 @@ pm_resolve_patches = PatternMatcher([
  (UPat(name="buf").index(UPat(Ops.RANGE), allow_any_len=True)
    .store(UPat.any(UPat(Ops.BINARY, name="blob"), UPat(Ops.BINARY, name="blob").bitcast()).index(UPat(Ops.RANGE), allow_any_len=True).load())
    .end(UPat(Ops.RANGE)), fold_binary),
-  (UPat({Ops.BUFFER, Ops.SLICE, Ops.MSTACK}, name="buf").index(UPat(Ops.STACK, name="off")).store(UPat(Ops.STACK, name="val")), fold_const_store),
+  (UPat({Ops.BUFFER, Ops.MSTACK}, name="buf").index(UPat(Ops.STACK, name="off")).store(UPat(Ops.STACK, name="val")), fold_const_store),
 ])

 pm_assert_no_afters = PatternMatcher([(UPat(Ops.AFTER, name="a"), lambda a: panic(RuntimeError, f"AFTER left at hcq_link: {a.src[0].op}"))])
@@ -484,7 +510,7 @@ def hcq_link(linear:UOp, cache=True) -> UOp:
  bufs = {(j,i):a for j,c in enumerate(linear.src) for i,a in enumerate(c.src[1:], 1)
          if a.op is Ops.AFTER and unwrap_mstack(a.src[0])[0].tag in HCQ_CACHE_TAGS}
  linear = linear.substitute({x:link_buf_cache[k] for a in bufs.values() if (k:=link_buf_key(a)) in link_buf_cache for x in (a, a.src[0])}, walk=True)
-  linear = graph_rewrite(linear, pm_resolve_patches+symbolic+pm_assert_no_afters, bpm=pm_bufferize, ctx=cache, bottom_up=False,
+  linear = graph_rewrite(linear, pm_resolve_patches+symbolic+pm_fold_cast_const+pm_assert_no_afters, bpm=pm_bufferize, ctx=cache, bottom_up=False,
                         name="resolve patches")
  for (j,i),a in bufs.items(): link_buf_cache.setdefault(link_buf_key(a), linear.src[j].src[i])
  if cache: link_linear_cache[linear_key] = linear
@@ -512,6 +538,27 @@ class HCQ2Compiled(Compiled):

    self.rt_buffer = Buffer(self.device, 64 << 20, dtypes.uint8, options=BufferSpec(uncached=True, cpu_access=True))
    self.rt_allocator = BumpAllocator(64 << 20)
+    self.prof_ents:dict[int, ProfileGraphEntry] = {}
+
+  def collect_prof(self):
+    if PROFILE:
+      es = list(self.prof_ents.values())
+      sigs = [self.signal(i)._buf.cpu_view().view(fmt='Q')[0]/decimal.Decimal(self.timestamp_divider) for e in es for i in (e.st_id, e.en_id)]
+      Compiled.profile_events.append(ProfileGraphEvent([replace(e, st_id=2*i, en_id=2*i+1) for i,e in enumerate(es)], [], sigs))
+    self.prof_ents.clear()
+
+  def _at_profile_finalize(self):
+    from tinygrad.tensor import Tensor
+    tdiffs = []
+    for _ in range(5):
+      with Context(DEBUG=0, BEAM=0, TRACK_MATCH_STATS=0): Tensor.ones(1, device=self.device).contiguous().realize()
+      if not (ents:=list(self.prof_ents.values())): return
+      self.prof_ents.clear()
+      st = perf_counter_us()
+      self.synchronize()
+      gpu = max(self.signal(e.en_id)._buf.cpu_view().view(fmt='Q')[0] for e in ents)/decimal.Decimal(self.timestamp_divider)
+      tdiffs.append((st+perf_counter_us())/2 - gpu)
+    Compiled.profile_events.append(ProfileDeviceEvent(self.device, statistics.median(tdiffs), self.device_props()))

  def new_buffer(self, b:UOp, cache:bool) -> Buffer:
    if cache or b.tag in HCQ_CACHE_TAGS:
@@ -525,13 +572,15 @@ class HCQ2Compiled(Compiled):
    return buf

  def synchronize(self, timeout:int|None=None):
-    if not hasattr(self, 'iface'): return
+    if HCQ_RUNTIME_DEV.value != self.device: Device[HCQ_RUNTIME_DEV.value].synchronize()
+
    sig = self.signal("timeline").as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
    tl = self.signal("value", 1).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
    timeout = timeout if timeout is not None and self.can_recover else None
    st = time.perf_counter()
    while sig[0] < tl[0] - 1:
      if time.perf_counter() - st > (timeout or 3000) / 1000: self.on_device_hang()
+    if self.prof_ents: self.collect_prof()

  def on_device_hang(self): raise RuntimeError(f"{self.device} hang detected")

@@ -575,6 +624,8 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
    if not hasattr(self, '_do_map'): raise NotImplementedError("map failed: no method implemented")
    return self._do_map(buf)

+  def _do_unmap(self, mb): self.dev.iface.free(mb)
+
  @suppress_finalizing
  def _free(self, buf:HCQ2Buffer, options:BufferSpec|None=None):
    if options is not None and options.external_ptr is not None: return
@@ -583,6 +634,6 @@ class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):

  def _unmap(self, mb):
    self.dev.synchronize()
-    self.dev.iface.free(mb)
+    self._do_unmap(mb)

  def _offset(self, buf, size:int, offset:int) -> HCQ2Buffer: return buf.offset(offset=offset, size=size)
@@ -173,14 +173,16 @@ class MemoryManager:
  va_allocator: ClassVar[TLSFAllocator|None] = None

  def __init__(self, dev, vram_size:int, boot_size:int, pt_t, va_bits:int, va_shifts:list[int], va_base:int,
-               palloc_ranges:list[tuple[int, int]], first_lv:int=0, reserve_ptable=False):
+               palloc_ranges:list[tuple[int, int]], first_lv:int=0, reserve_ptable=False, paddr_base:int=0):
    self.dev, self.vram_size, self.va_shifts, self.va_base, lvl_msb = dev, vram_size, va_shifts, va_base, va_shifts + [va_bits + 1]
    self.pte_covers, self.pte_cnt = [1 << x for x in va_shifts][::-1], [1 << (lvl_msb[i+1] - lvl_msb[i]) for i in range(len(lvl_msb) - 1)][::-1]
    self.pt_t, self.palloc_ranges, self.level_cnt, self.va_bits, self.reserve_ptable = pt_t, palloc_ranges, len(va_shifts), va_bits, reserve_ptable

-    self.boot_allocator = TLSFAllocator(boot_size, base=0)
-    self.ptable_allocator = TLSFAllocator(round_up(vram_size // 512, 1 << 20) if self.reserve_ptable else 0, base=self.boot_allocator.size)
-    self.pa_allocator = TLSFAllocator(vram_size - (off_sz:=self.boot_allocator.size + self.ptable_allocator.size), base=off_sz)
+    self.boot_allocator = TLSFAllocator(boot_size, base=paddr_base)
+    self.ptable_allocator = TLSFAllocator(round_up(vram_size // 512, 1 << 20) if self.reserve_ptable else 0,
+                                          base=paddr_base + self.boot_allocator.size)
+    off_sz = paddr_base + self.boot_allocator.size + self.ptable_allocator.size
+    self.pa_allocator = TLSFAllocator(vram_size - off_sz, base=off_sz)
    self.root_page_table = pt_t(self.dev, self.palloc(0x1000, zero=not self.dev.smi_dev, boot=True), lv=first_lv)

  def _frag_size(self, va, sz, must_cover=True):
@@ -1,7 +1,8 @@
 from __future__ import annotations
-import os, mmap, array, functools, ctypes, select, contextlib, dataclasses, sys, itertools, struct, socket, subprocess, time, enum, atexit
+import os, mmap, array, functools, ctypes, ctypes.util, select, contextlib, dataclasses, sys, itertools, struct, socket
+import subprocess, time, enum, atexit
 from tinygrad.helpers import round_up, getenv, OSX, temp, ceildiv, unwrap, fetch, system, _ensure_downloads_dir, DEBUG, flatten, pluralize
-from tinygrad.runtime.autogen import libc, pci, vfio, iokit, corefoundation
+from tinygrad.runtime.autogen import libc, pci, vfio
 from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface, HCQBuffer, hcq_filter_visible_devices
 from tinygrad.runtime.support.memory import VirtMapping, AddrSpace, BumpAllocator
 from tinygrad.runtime.support.usb import USB3, CustomASM24Controller, USBMMIOInterface
@@ -55,6 +56,7 @@ class _System:
  def pci_scan_bus(self, vendor:int, devices:tuple[tuple[int, tuple[int, ...]], ...], base_class:int|None=None) -> list[str]:
    all_devs = []
    if OSX:
+      from tinygrad.runtime.autogen import iokit, corefoundation
      def read_prop(svc, key) -> int:
        cfkey = corefoundation.CFStringCreateWithCString(None, key.encode(), corefoundation.kCFStringEncodingUTF8)
        cfdata = ctypes.cast(iokit.IORegistryEntryCreateCFProperty(svc, ctypes.cast(cfkey, iokit.CFStringRef), None, 0), corefoundation.CFDataRef)
@@ -98,10 +98,14 @@ pm_post_sched_cache = PatternMatcher([
   create_new_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None),
 ])

+def resolve_linear_call(linear_call:UOp):
+  linear = graph_rewrite(linear_call.src[0], pm_post_sched_cache, ctx=({}, linear_call.src[1:]), walk=True, name="params to buffers")
+  binds = {f"p{i}":x.src[0] for i,x in enumerate(linear_call.src[1:]) if x.op is Ops.BIND}
+  return linear.substitute({v:binds[v.expr] for v in linear.variables() if v.expr in binds}, enter_calls=True, name="resolve scalar params")
+
 pm_resolve_linear_call = PatternMatcher([
  # call LINEAR is resolved here
-  (UPat(Ops.CALL, src=(UPat(Ops.LINEAR),), name="linear_call", allow_any_len=True), lambda linear_call:
-   graph_rewrite(linear_call.src[0], pm_post_sched_cache, ctx=({}, linear_call.src[1:]), walk=True, name="params to buffers")),
+  (UPat(Ops.CALL, src=(UPat(Ops.LINEAR),), name="linear_call", allow_any_len=True), resolve_linear_call),
 ])+pm_flatten_linear

 schedule_cache: dict[bytes, UOp] = {}
@@ -15,14 +15,13 @@ def handle_allreduce(buf:UOp, red:UOp) -> UOp|None:
  use_ring = concrete and not use_all2all and (RING >= 2 or (ndev > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and RING >= 1))
  if DEBUG >= 2: print(f"{'ALL2ALL' if use_all2all else 'RING' if use_ring else 'NAIVE'} ALLREDUCE {ndev}x{numel} | {buf.dtype}")

-  if not concrete: buf = buf.pad_to(buf.max_shape)
+  buf = buf.pad_to(buf.max_shape)
  # contiguous before we copy it
  buf = buf.contiguous()

  # naive: copy to all devices. if you shrink later, that'll be handled
  if not use_ring and not use_all2all:
-    out = functools.reduce(lambda x,y: x.alu(op, y), [buf.mselect(i).copy_to_device(device) for i in range(ndev)])
-    return out if concrete else out.shrink_to(shape)
+    return functools.reduce(lambda x,y: x.alu(op, y), [buf.mselect(i).copy_to_device(device) for i in range(ndev)]).shrink_to(shape)

  # chunk data into ndev pieces
  assert isinstance(numel, int)
@@ -7,25 +7,46 @@ from tinygrad.uop.ops import gate_kernel_sink
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
 from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored, Context, SPEC

-ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.BUFFER, Ops.SLICE,
+@dataclass
+class IndexingContext:
+  realize_map: dict[UOp, None|list[int]] = field(default_factory=dict)
+  non_removable: dict[UOp, None] = field(default_factory=dict)
+  range_map: dict[UOp, tuple[tuple[UOp, ...], tuple[UOp, ...]]] = field(default_factory=dict)
+  # loads reachable from each UOp memoized across matches
+  buf_cache: dict[UOp, frozenset[UOp]] = field(default_factory=dict)
+
+  # create ranges
+  range_idx: Iterator[int] = field(default_factory=itertools.count)
+  def new_range(self, s:sint, axistype:AxisType=AxisType.WEAK) -> UOp:
+    if isinstance(s, UOp) and s.op is Ops.RANGE: return s
+    # if a range has a 1 src, it's the same as UOp.const(0)
+    return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(0)
+
+
+ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.BUFFER,
                      Ops.CONST, Ops.BIND, Ops.MSELECT, Ops.MSTACK, Ops.PARAM,
                      Ops.LOAD, Ops.CALL, Ops.FUNCTION}

-def realize(ctx:dict[UOp, None], tr:UOp) -> None: ctx[tr] = None
+def realize(ctx:IndexingContext, tr:UOp) -> None: ctx.realize_map[tr] = None

-def realize_srcs(ctx:dict[UOp, None], rb:UOp) -> None:
+def realize_srcs(ctx:IndexingContext, rb:UOp) -> None:
  for s in rb.src:
-    if s.base.op not in ALWAYS_CONTIGUOUS: ctx[s] = None
+    if s.base.op not in ALWAYS_CONTIGUOUS: ctx.realize_map[s] = None

-def realize_store_after_src(ctx:dict[UOp, None], dest:UOp, src:UOp):
-  # don't realize SLICE when it's the direct source of STORE+AFTER — the target buffer is the output
-  if src.op is Ops.SLICE and src in ctx \
-     and not dest.op_in_backward_slice_with_self(Ops.SHRINK, Ops.PERMUTE, Ops.FLIP, Ops.PAD):
-    del ctx[src]
+def realize_store_after_src(ctx:IndexingContext, dest:UOp, src:UOp):
  # you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce
-  if dest.base in src.backward_slice_with_self: ctx[src] = None
+  if dest.base in src.backward_slice_with_self: ctx.realize_map[src] = None
+
+def realize_custom_kernel_srcs(ctx:IndexingContext, c:UOp) -> None:
+  for s in c.src[1:]:
+    while s.op is Ops.RESHAPE: s = s.src[0]
+    if s.op not in ALWAYS_CONTIGUOUS:
+      ctx.realize_map[s] = None
+      ctx.non_removable[s] = None

 pm_generate_realize_map = PatternMatcher([
+  # realize the inputs of custom kernel calls
+  (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.PROGRAM)),), name="c", allow_any_len=True), realize_custom_kernel_srcs),
  # always realize
  (UPat({Ops.CONTIGUOUS, Ops.STORE}, name="tr"), realize),
  # realize srcs of these
@@ -41,20 +62,6 @@ class BufferizeOpts:
  addrspace: AddrSpace = AddrSpace.GLOBAL
  removable: bool = True

-@dataclass
-class IndexingContext:
-  realize_map: dict[UOp, None|list[int]] = field(default_factory=dict)
-  range_map: dict[UOp, tuple[tuple[UOp, ...], tuple[UOp, ...]]] = field(default_factory=dict)
-  # loads reachable from each UOp memoized across matches
-  buf_cache: dict[UOp, frozenset[UOp]] = field(default_factory=dict)
-
-  # create ranges
-  range_idx: Iterator[int] = field(default_factory=itertools.count)
-  def new_range(self, s:sint, axistype:AxisType=AxisType.WEAK) -> UOp:
-    if isinstance(s, UOp) and s.op is Ops.RANGE: return s
-    # if a range has a 1 src, it's the same as UOp.const(0)
-    return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(0)
-
 def broadcast_rngs(x:UOp, src:UOp, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]:
  if x.op not in GroupOp.Broadcastable: return rngs
  baxes, nleft = broadcast_axes(src.shape, x.shape), len(x.shape)-len(src.shape)
@@ -63,7 +70,7 @@ def broadcast_rngs(x:UOp, src:UOp, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]:
 # TODO: srcs contain (real data srcs, something else, ranges) and the boundary is confusing. see range_start
 def data_srcs(op:Ops, src:tuple[UOp, ...]) -> tuple[UOp, ...]:
  if op in {Ops.PARAM, Ops.BUFFER, Ops.RANGE, Ops.SPECIAL, Ops.BIND}: return ()
-  if op in GroupOp.Movement|{Ops.INDEX, Ops.SLICE, Ops.STAGE, Ops.REDUCE, Ops.AFTER, Ops.END}: return src[:1]
+  if op in GroupOp.Movement|{Ops.INDEX, Ops.STAGE, Ops.REDUCE, Ops.AFTER, Ops.END}: return src[:1]
  return src

 def create_bufferize_and_index_srcs(ctx:IndexingContext, x:UOp) -> list[UOp]:
@@ -73,7 +80,7 @@ def create_bufferize_and_index_srcs(ctx:IndexingContext, x:UOp) -> list[UOp]:
  for i, s in enumerate(x.src):
    new_src = s
    src_rngs = broadcast_rngs(x, s, ctx.range_map[x][0]) if x in ctx.range_map else ()
-    if s.op in {Ops.PARAM, Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
+    if s.op in {Ops.PARAM, Ops.BUFFER, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
      if x in ctx.range_map and i < data_src_count: new_src = new_src.index(*src_rngs)
    elif s in ctx.realize_map:
      realized_ranges = ctx.realize_map[s]
@@ -84,7 +91,7 @@ def create_bufferize_and_index_srcs(ctx:IndexingContext, x:UOp) -> list[UOp]:
        new_src = s.end(*[r for r in closed_ranges if r.op is Ops.RANGE])
        del ctx.realize_map[s]
      else:
-        removable = s.op not in ALWAYS_CONTIGUOUS
+        removable = s.op not in ALWAYS_CONTIGUOUS and s not in ctx.non_removable
        # LOCAL: None in the device assigns it a number later
        opts = BufferizeOpts(device=s.device, removable=removable) if len(ctx.range_map[s][1]) == len(realized_ranges) else \
               BufferizeOpts(device=s.device, addrspace=AddrSpace.LOCAL, removable=removable)
@@ -105,6 +112,7 @@ def convert_pad_to_where_to_keep_behavior_local(ctx:IndexingContext, x:UOp):

 def convert_reduce_to_reduce_with_ranges(ctx:IndexingContext, x:UOp):
  if x.arg[1] == 0: return None
+  if x not in ctx.range_map: raise RuntimeError("REDUCE has no ranges in rangeify, UOp verification failed")
  bx = create_bufferize_and_index_based_on_ranges(ctx, x)
  # input ranges
  new_ranges = list(ctx.range_map[x][0][:x.arg[1]])
@@ -182,7 +190,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
  rctx = IndexingContext()

  # get ops to realize
-  graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx.realize_map, name="get realize")
+  graph_rewrite(tsink, pm_generate_realize_map, ctx=rctx, name="get realize")

  # get the consumer map
  with cpu_profile("consumer map in rangeify", "TINY"):
@@ -52,11 +52,9 @@ def memory_plan_rewrite(linear:UOp, held_bufs:set[UOp]|None=None) -> UOp:
    peaks[_key(buf)] = (max(peaks[_key(buf)][0], offsets[buf] + buf.max_numel() * buf.dtype.itemsize), peaks[_key(buf)][1])
  arena_sizes = {key: round_up(peak, block_size) for key, (peak, _) in peaks.items()}

-  # build replace_map: each buffer becomes a SLICE into a shared per-device-lane arena
+  # build replace_map: each buffer becomes a SHRINK/BITCAST into a shared per-device-lane arena
  arenas = {key: UOp.new_buffer(key[0], sz, dtypes.int8) for key, sz in arena_sizes.items()}
-  replace_map:dict[UOp, UOp] = {}
-  for buf_uop, offset in offsets.items():
-    replace_map[buf_uop] = UOp(Ops.SLICE, buf_uop.dtype, (arenas[_key(buf_uop)], UOp.const(offset)), buf_uop.max_numel())
+  replace_map = {buf_uop:arenas[_key(buf_uop)][offset:offset+buf_uop.nbytes()].bitcast(buf_uop.dtype) for buf_uop, offset in offsets.items()}

  if DEBUG >= 1 and (omem:=sum(nbytes.values()) / 1e6) != (nmem:=sum(arena_sizes.values()) / 1e6):
    print(f"memory reduced from {omem:.2f} MB -> {nmem:.2f} MB, {len(first_appearance)} -> {len(arenas)} bufs")
@@ -126,7 +126,7 @@ def reshape_multi(root:UOp, multi:UOp):
  new_shardings = []
  for ax, rng in multi.sharding:
    count = int(rng.vmax)+1
-    target = prod(multi.shape[:ax])
+    target = ssimplify(prod(multi.shape[:ax]))
    if target not in arg_acc: raise RuntimeError(f"reshape {multi.shape} -> {new_shape} moved items between shards")
    new_ax = len(arg_acc) - arg_acc[::-1].index(target) - 1
    if new_shape[new_ax] % count != 0: raise RuntimeError(f"reshape {multi.shape} -> {new_shape} moved items between shards")
@@ -4,9 +4,9 @@ import itertools
 from tinygrad.dtype import dtypes, AddrSpace, Invalid, to_dtype, strong_dtype
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg, shape_to_shape_arg
 from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, rewrite_group, identity_element
-from tinygrad.uop.symbolic import symbolic
+from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const
 from tinygrad.uop.movement import mop_cleanup
-from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS
+from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS, SPEC
 from tinygrad.helpers import PCONTIG, FLOAT16, OPENPILOT_HACKS, argsort, partition, get_single_element
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
@@ -141,7 +141,7 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),

  # SINK only ever references the base
-  (UPat(Ops.SINK, name="x"), lambda x: x.replace(src=tuple(y.base for y in x.src))),
+  (UPat(Ops.SINK, name="x"), lambda x: x.replace(src=tuple(y.unsharded_base for y in x.src))),

  # ** copy rules **

@@ -193,6 +193,7 @@ ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.NOOP}

 # you don't know in the first pass if axes are going to die, this happens if there's an EXPAND to the left
 def cleanup_dead_axes(b:UOp):
+  if not b.arg.removable: return None
  # don't optimize ALWAYS_RUN_OPS or AFTER (AFTER is a buffer identity — ranges define consumer access, not computation)
  if b.src[0].op in ALWAYS_RUN_OPS or b.src[0].op is Ops.AFTER: return None

@@ -326,6 +327,26 @@ pm_remove_bufferize = PatternMatcher([
  (UPat(Ops.END, src=(UPat(Ops.NOOP, name="x"),), allow_any_len=True), lambda x: x),
 ])

+def no_indexing_calls(u:UOp):
+  new_srcs = []
+  for x in u.src:
+    if x.op is Ops.INDEX:
+      # sometimes if call srcs have children the call will get an INDEX. we remove it here.
+      # TODO: we should add safety checks here for contiguous
+      new_srcs.append(x.src[0])
+    elif x.op is Ops.SHRINK:
+      # SHRINK with offset 0 is fine
+      # TODO: check offset
+      new_srcs.append(x.src[0])
+    else:
+      # everything else we pass through
+      new_srcs.append(x)
+  return u.replace(src=tuple(new_srcs))
+
+pm_no_indexing_calls = PatternMatcher([
+  (UPat(Ops.CALL, name="u"), no_indexing_calls),
+])
+
 DEVICE_MAX_BUFS = {"METAL": 31, "WEBGPU": 8, "CPU": 31} # TODO: get from device?
 def limit_bufs(ctx:IndexingContext, root:UOp):
  if (device:=root.device) is None: return None # no device, index related calculations
@@ -562,7 +583,9 @@ def get_kernel_graph(sink:UOp) -> UOp:
  # convert movement ops to ranges
  tsink, rctx = run_rangeify(tsink, bool(DEBUG_RANGEIFY))

-  tsink = graph_rewrite(tsink, symbolic+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize, name="symbolic+reduce_collapse+debuf")
+  tsink = graph_rewrite(tsink,
+                        symbolic+pm_fold_cast_const+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize+pm_no_indexing_calls,
+                        name="symbolic+reduce_collapse+debuf")
  tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")

  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Rangeify")
@@ -574,4 +597,8 @@ def get_kernel_graph(sink:UOp) -> UOp:
  tsink = graph_rewrite(tsink, split_kernels, bottom_up=True, name="split kernels")

  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
+  if SPEC:
+    # validate the kernel graph
+    from tinygrad.uop.spec import type_verify, spec_kernel_graph
+    type_verify(tsink, spec_kernel_graph, enter_calls=False)
  return tsink
@@ -23,6 +23,7 @@ class AllocCtx:
  bases: set[UOp] = field(default_factory=set)
  assigns: list[UOp] = field(default_factory=list)
  replacements: list[UOp] = field(default_factory=list)
+  views: set[UOp] = field(default_factory=set)

 def tag_uop(ctx:AllocCtx, x:UOp):
  if x.tag is not None: return None
@@ -63,40 +64,37 @@ def replace_contig_with_store_after(u:UOp):
 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)
+  if assigned_to.op is not Ops.BUFFER: 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())
+  if (cv := src.contiguous_view()) is None: return None
+  (buf, offset), size = cv, src.max_numel() * src.element_size() // cv[0].element_size()
+  if buf.op is not Ops.BUFFER: return None
+  # NB: make offset a UOp.variable here to do the offset computation in the kernels
+  return buf[offset:offset+size].bitcast(src.dtype)

-def contiguous_mops_to_view(c:UOp, src:UOp):
-  """MOPS(BUFFER) → SLICE when movement ops collapse to a contiguous range."""
+def contiguous_mops_to_view(ctx:AllocCtx, c:UOp, src:UOp):
+  """MOPS(BUFFER) → SHRINK 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
+  while buf.op is Ops.BITCAST: buf = buf.src[0].base
+  if buf.op not in {Ops.BUFFER, Ops.UNSHARD}: 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
+    ctx.views.add(view)
+    view = view.reshape(c.shape)
+    return c.replace(src=(view,)+c.src[1:]) if c.op in {Ops.COPY, Ops.STORE} else view

-  # for UNSHARD tensors, use multi_pm to resolve per-shard movement ops, then create SLICE on the resolved result
+  # for UNSHARD tensors, use multi_pm to resolve per-shard movement ops, then create SHRINK 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)
+    ctx.views.add(view)
+    return view.reshape(resolved.src[0].shape).unshard(resolved.arg, resolved.src[1:])

  return None

@@ -151,8 +149,9 @@ pm_early_transform_tensor_graph = PatternMatcher([
  # 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),
+  # fold MOPS+BITCAST over BUFFER into SHRINK when movement ops collapse to contiguous range
+  (UPat((Ops.COPY, Ops.CONTIGUOUS), src=(UPat(GroupOp.Movement|{Ops.BITCAST}, name="src"),), name="c"), contiguous_mops_to_view),
+  (UPat(Ops.STORE, src=(UPat(Ops.BITCAST, name="src"), UPat()), name="c", allow_any_len=True), 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:
@@ -184,7 +183,12 @@ def finalize_after(ctx:AllocCtx, x:UOp):
  # 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]
+  # then, add views back
+  views:list[UOp] = []
+  while replace_uop.op in GroupOp.Movement|{Ops.UNSHARD, Ops.BITCAST, Ops.AFTER}:
+    if replace_uop.op is not Ops.AFTER: views.append(replace_uop)
+    replace_uop = replace_uop.src[0]
+  for v in reversed(views): replace_uop = v.replace(src=(replace_uop,)+v.src[1:])
  for t in x.tag:
    original_uop: UOp = ctx.uop_list[t]
    ctx.buffer_map[original_uop] = replace_uop.shrink_to(original_uop.shape)
@@ -196,6 +200,8 @@ def replace_input_buffer(ctx:AllocCtx, b:UOp):
  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)

+def replace_input_view(ctx:AllocCtx, b:UOp): return replace_input_buffer(ctx, b) if b in ctx.views else None
+
 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),
@@ -205,8 +211,9 @@ 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),
+  # replace SHRINK with PARAM
+  (UPat(Ops.SHRINK, src=(UPat(Ops.BUFFER),), name="b", allow_any_len=True), replace_input_view),
+  (UPat(Ops.BITCAST, src=(UPat.any(UPat(Ops.SHRINK, src=(UPat(Ops.BUFFER),), allow_any_len=True), UPat(Ops.BUFFER)),), name="b"), replace_input_view),
  # 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),
 ])
@@ -224,7 +231,7 @@ def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
  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")
+  big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, ctx=ctx, 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")
@@ -93,7 +93,7 @@ class Ops(FastEnum):
  CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto()

  # buffer ops
-  STAGE = auto(); COPY = auto(); SLICE = auto(); MSELECT = auto(); MSTACK = auto(); CUSTOM_FUNCTION = auto()
+  STAGE = auto(); COPY = auto(); MSELECT = auto(); MSTACK = auto(); CUSTOM_FUNCTION = auto()

  # the core 6 movement ops! these only exist in the tensor graph
  RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); FLIP = auto()
@@ -45,8 +45,7 @@ axis_colors = {AxisType.DEVICE: "green", AxisType.GLOBAL: "blue", AxisType.THREA
 axis_to_pos = {AxisType.DEVICE: -2, AxisType.WEAK: -1, AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1,
               AxisType.LOCAL: 2, AxisType.UPCAST: 3, AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}

-range_start = {Ops.STAGE: 1, Ops.REDUCE: 1, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.FUNCTION: 1,
-               Ops.SLICE: 2, Ops.LINEAR: 0}
+range_start = {Ops.STAGE: 1, Ops.REDUCE: 1, Ops.WMMA: 3, Ops.END: 1, Ops.CALL: 1, Ops.FUNCTION: 1, Ops.LINEAR: 0}

 # https://en.wikipedia.org/wiki/Identity_element
 def identity_element(op:Ops, dt:DType) -> PyConst: return dt.const({Ops.ADD:0, Ops.MUL:1, Ops.MAX:dt.min}[op])
@@ -171,9 +170,6 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None:
      return arg.dtype
    case Ops.BINARY:
      return dtypes.uint8
-    case Ops.SLICE:
-      # TODO: slice just shouldn't exist
-      return None
    case Ops.CAST | Ops.BITCAST:
      assert isinstance(arg, DType), f"CAST/BITCAST arg must be DType, got {arg}"
      return arg
@@ -221,7 +217,7 @@ class UOpMetaClass(type):
    return created

 # some uops map to other stuff
-buffers:weakref.WeakKeyDictionary[UOp, Buffer|MultiBuffer] = weakref.WeakKeyDictionary() # this maps BUFFER/SLICE uops to their device Buffers
+buffers:weakref.WeakKeyDictionary[UOp, Buffer|MultiBuffer] = weakref.WeakKeyDictionary() # this maps BUFFER/view uops to their device Buffers
 all_metadata:weakref.WeakKeyDictionary[UOp, tuple[Metadata, ...]] = weakref.WeakKeyDictionary() # TODO: should this be here?

 # recursive_property replaces functools.cached_property in recursive UOp functions to prevent RecursionError
@@ -386,10 +382,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
      case Ops.BUFFER:
        if len(self.src): return self.src[0].as_shape
        return ()
-      case Ops.SLICE:
-        # HACK: SLICE is used inside kernels, so we set the shape to () if it's on an INDEX
-        if self.src[0].op is Ops.INDEX: return ()
-        return (self.arg,)
      case Ops.CUSTOM | Ops.CUSTOMI:
        if self.dtype is dtypes.void: return None
        input_shapes = [x._shape for x in self.src if x._shape is not None]
@@ -528,9 +520,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    if self.op is Ops.CONST: return self
    if self.op is Ops.SINK and all(s.op is Ops.CONST or (s.op is Ops.STACK and len(s.src) == 0) for s in self.src): return self
    # late import!
-    from tinygrad.uop.symbolic import symbolic
+    from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const
    with Context(TRACK_MATCH_STATS=0 if not tracked else TRACK_MATCH_STATS.value):
-      return graph_rewrite(self, symbolic, name="simplify")
+      return graph_rewrite(self, symbolic+pm_fold_cast_const, name="simplify")
  def ssimplify(self) -> UOp|ConstType: return ret.val if (ret:=self.simplify()).op is Ops.CONST else ret
  def _eval(self, dtype, expected_type:Type[T]) -> T:
    assert self.dtype in dtype, f"eval with wrong dtype {self}"
@@ -568,9 +560,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    in_tuple = self.src[0] if self.op is Ops.FUNCTION else self
    assert in_tuple.op is Ops.TUPLE, f"gettuple requires FUNCTION or TUPLE source, got {self.op}"
    return UOp(Ops.GETTUPLE, src=(self,), arg=idx)
-  def group(*srcs:UOp|None):  # pylint: disable=no-self-argument
+  def group(*srcs:UOp|None, **kwargs):  # pylint: disable=no-self-argument
    if len(srcs) == 1 and isinstance(srcs[0], UOp): return srcs[0]
-    return UOp(Ops.GROUP, src=tuple([x for x in srcs if x is not None]))
+    return UOp(Ops.GROUP, src=tuple([x for x in srcs if x is not None]), **kwargs)
  def index(self, *srcs:UOp|int|None, **kwargs):
    new_srcs: list[UOp] = [UOp.const(x) if isinstance(x, int) else x for x in srcs if x is not None]
    if len(new_srcs) == 1 and new_srcs[0].op is Ops.CONST and self.op is Ops.STACK: return self.src[new_srcs[0].val]
@@ -758,7 +750,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    assert arg is None or isinstance(self.device, tuple)
    inp = self if arg is None else UOp(Ops.MSELECT, src=(self,), arg=arg)
    if inp.dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {inp.dtype}")
-    return UOp(Ops.COPY, src=(inp,), arg=device)
+    return UOp(Ops.COPY, src=(inp.pad_to(inp.max_shape),), arg=device).shrink_to(inp.shape)
  def mselect(self, arg:int) -> UOp: return UOp(Ops.MSELECT, src=(self,), arg=arg)
  def mstack(self, *srcs: UOp) -> UOp: return UOp(Ops.MSTACK, src=(self,)+srcs) if len(srcs) else self
  @property
@@ -772,6 +764,15 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    if self.op is Ops.DETACH: return self.src[0].base  # DETACH can't change base
    return self

+  # base with UNSHARD
+  @property
+  def unsharded_base(self) -> UOp:
+    if self.op in GroupOp.Movement: return self.src[0].base
+    if self.op is Ops.DETACH: return self.src[0].base  # DETACH can't change base
+    # TODO: why can't this be in normal base?
+    if self.op is Ops.UNSHARD: return self.src[0].base
+    return self
+
  # cached property here makes external_uop_gc fail, why?
  @property
  def as_shape(self) -> tuple[sint, ...]:
@@ -813,7 +814,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
  unique_num = itertools.count(0)

  def getaddr(self, device=None) -> UOp:
-    if self.without_after.op not in {Ops.BUFFER, Ops.SLICE, Ops.BINARY, Ops.MSTACK, Ops.MSELECT, Ops.PARAM}: return self
+    if self.without_after.op not in {Ops.BUFFER, Ops.SHRINK, Ops.BITCAST, Ops.BINARY, Ops.MSTACK, Ops.MSELECT, Ops.PARAM}: return self
    return UOp(Ops.GETADDR, src=(self,), arg=device or to_tuple(self.device)[0])
  @staticmethod
  def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType, num=None):
@@ -892,8 +893,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    while len(s.src) and s.op not in {Ops.BUFFER, Ops.PARAM, Ops.STAGE, Ops.MSTACK}: s = s.src[0]
    return s

-  def contiguous_view_offset(self) -> int|None:
-    """If movement ops on a BUFFER collapse to a contiguous range, return `offset` in elements. Otherwise None."""
+  def contiguous_view(self) -> tuple[UOp, int]|None:
    from tinygrad.schedule.rangeify import pm_mops
    from tinygrad.uop.symbolic import symbolic

@@ -906,14 +906,17 @@ class UOp(RandMixin, metaclass=UOpMetaClass):

    idx = self.flatten().index(UOp.range(self.numel(), 0))
    out = graph_rewrite(idx, pm_mops+symbolic+pm_contiguous_view_offset, ctx=self, name="contiguous_view_offset")
-    return out.val if out.op is Ops.CONST and isinstance(out.val, int) else None
+    if out.op is not Ops.INDEX or not (b:=out.src[0]).tag or (c:=out.src[1]).op is not Ops.CONST or not isinstance(c.val, int): return None
+    return b.rtag(None), c.val
+
+  def contiguous_view_offset(self) -> int|None: return None if (view := self.contiguous_view()) is None else view[1]

  def has_buffer_identity(self, after_ok=False):
    """Check if this UOp has a concrete buffer identity in the graph (RESHAPE/UNSHARD -> BUFFER chain)."""
    # TODO: this is confusing because UOp.variable('v', 0, 1, dtypes.weakfloat) is True for jit to work, but it doesn't have a buffer
    if self.op in {Ops.RESHAPE, Ops.UNSHARD, Ops.MSELECT}: return self.src[0].has_buffer_identity(after_ok)
    if after_ok and self.op == Ops.AFTER: return self.src[0].has_buffer_identity(after_ok)
-    return self.op in {Ops.BUFFER, Ops.SLICE, Ops.PARAM}
+    return self.op in {Ops.BUFFER, Ops.PARAM}

  def _base_buffer_is_realized(self) -> bool:
    """Walk through AFTER chain to find if the underlying buffer is realized (has allocated memory)."""
@@ -923,30 +926,19 @@ class UOp(RandMixin, metaclass=UOpMetaClass):

  @property
  def buffer(self) -> Buffer|MultiBuffer:
-    if self.op in {Ops.CONTIGUOUS, Ops.RESHAPE, Ops.UNSHARD, Ops.DETACH, Ops.AFTER}: return self.src[0].buffer
+    if self.op in {Ops.CONTIGUOUS, Ops.CONTIGUOUS_BACKWARD, Ops.RESHAPE, Ops.UNSHARD, Ops.DETACH, Ops.AFTER}: return self.src[0].buffer
    # this buffer can process disk tensors and simple movement ops
-    if self is not self.base:
-      buf = self.base.buffer
-      assert isinstance(buf, Buffer), "must be a Buffer for movement ops"
-      offset = self.contiguous_view_offset()
-      if offset is None: raise RuntimeError(f"non-contiguous view is not supported for {buf.device} buffer")
-      return buf.view(prod(self.max_shape), self.dtype, offset*self.dtype.itemsize)
-    if self.op is Ops.BITCAST:
-      buf = self.src[0].buffer
-      assert isinstance(buf, Buffer), "must be a Buffer for BITCAST"
-      return buf.view(prod(self.max_shape), self.dtype, 0)
-    if self.op is Ops.SLICE:
+    if self is not self.base or self.op is Ops.BITCAST:
      if (cret:=buffers.get(self)) is not None: return cret
-      buf = self.src[0].buffer
-      offset = self.src[1].val
+      if (cv := self.contiguous_view()) is None: raise RuntimeError(f"non-contiguous view is not supported for {self.device} buffer")
+      buf, offset = (b:=cv[0]).base.buffer, cv[1]
      if isinstance(buf, MultiBuffer):
        mbuf = MultiBuffer.__new__(MultiBuffer)
-        mbuf.bufs = [b.view(self.arg, self.dtype, offset * self.src[0].dtype.itemsize) for b in buf.bufs]
+        mbuf.bufs = [x.view(prod(self.max_shape), self.dtype, offset*b.dtype.itemsize) for x in buf.bufs]
        buffers[self] = mbuf
        return mbuf
-      assert isinstance(buf, Buffer), "must be a Buffer for SLICE"
-      buffers[self] = bv = buf.view(self.arg, self.dtype, offset * self.src[0].dtype.itemsize)
-      return bv
+      buffers[self] = buf.view(prod(self.max_shape), self.dtype, offset*b.dtype.itemsize)
+      return buffers[self]
    if self.op is Ops.MSELECT:
      ret = self.src[0].buffer
      assert isinstance(ret, MultiBuffer)
@@ -1097,12 +1089,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    if self.op is Ops.CONST and self.val is not Invalid: return self.val, self.val
    if self.op is Ops.INDEX: return self.src[0]._min_max
    if self.op is Ops.CAST:
-      # an int destination truncates a float source toward zero. trunc is monotone
+      # rounding is monotone (truncation toward zero into an int, to-nearest onto the value grid into a float)
      smin, smax = self.src[0]._min_max
-      if dtypes.is_int(self.dtype) and dtypes.is_float(self.src[0].dtype) and all(math.isfinite(v) for v in (smin, smax)):
-        smin, smax = math.trunc(smin), math.trunc(smax)
-      # a cast to unsigned keeps exact bounds when the source fits
-      # TODO: can do more based on new dtype window
+      trunc = truncate.get(self.dtype) if dtypes.is_float(self.dtype) else math.trunc if dtypes.is_int(self.dtype) else None
+      if trunc is not None and all(math.isfinite(v) for v in (smin, smax)): smin, smax = trunc(smin), trunc(smax)
      if dtypes.is_unsigned(self.dtype) and 0 <= smin and smax <= self.dtype.max: return smin, smax
      if self.dtype in dtypes.floats+dtypes.sints+(dtypes.weakint,): return max(self.dtype.min, smin), min(smax, self.dtype.max)
    return self.dtype.min, self.dtype.max
@@ -1166,15 +1156,15 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    src: tuple[UOp, ...] = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),)
    return UOp(Ops.PARAM, src=src, arg=ParamArg(slot, dtype, vmin_vmax, multiple_of, name, addrspace, axis, device, volatile))
  def param_like(self, slot:int):
+    if self.op is Ops.BIND: return self.src[0].replace(arg=replace(self.src[0].arg, slot=slot, name=f"p{slot}"))
    addrspace = self.addrspace if self.addrspace is not None else AddrSpace.GLOBAL
-    if self.op is Ops.BIND: return self.src[0].replace(arg=replace(self.src[0].arg, slot=slot, addrspace=addrspace))
    return UOp.param(slot, self.dtype, self.shard_shape if self.axis is not None else self._shape, self.device, addrspace=addrspace, axis=self.axis)

  @staticmethod
  def custom_function(name:str, *src:UOp) -> UOp: return UOp(Ops.CUSTOM_FUNCTION, src=src, arg=name)

  # opaque bodies stay as Ops.CALL; value-producing bodies become Ops.FUNCTION (wrapped in TUPLE)
-  _OPAQUE_CALL_BODIES = {Ops.SINK, Ops.PROGRAM, Ops.LINEAR, Ops.COPY, Ops.SLICE, Ops.CUSTOM_FUNCTION}
+  _OPAQUE_CALL_BODIES = {Ops.SINK, Ops.PROGRAM, Ops.LINEAR, Ops.COPY, Ops.CUSTOM_FUNCTION}
  def call(self, *srcs:UOp, ret_dtype:DType|None=None, grad_fxn:Callable|None=None,
           name:str|None=None, precompile:bool=False, precompile_backward:bool=False, aux:Any=None) -> UOp:
    if ret_dtype is not None: return UOp(Ops.CALL, ret_dtype, src=(self,)+srcs)
@@ -1187,10 +1177,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    body = self if self.op is Ops.TUPLE else UOp.maketuple(self)
    return UOp(Ops.FUNCTION, src=(body,)+srcs, arg=CallInfo(grad_fxn, name, precompile, precompile_backward, aux))
  def custom_kernel(*srcs:UOp, fxn:Callable, grad_fxn:Callable|None=None) -> list[UOp]:
-    contig_srcs = tuple(x.contiguous() if x.op is not Ops.AFTER else x for x in srcs)
-    placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(contig_srcs)]
-    kernel = fxn(*placeholders).call(*contig_srcs, grad_fxn=grad_fxn)
-    return [s.after(kernel) for s in contig_srcs]
+    placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(srcs)]
+    kernel = fxn(*placeholders).call(*srcs, grad_fxn=grad_fxn)
+    return [s.after(kernel) for s in srcs]

  def to_elf(self) -> TinyELF:
    assert self.op is Ops.PROGRAM and isinstance(self.arg, ProgramInfo), "to_elf should only be called on a PROGRAM ast"
@@ -1752,77 +1741,6 @@ def _rebuild_dtype(n:UOp, new_src:tuple[UOp,...]) -> DType:
 def sint_to_uop(x:sint, dtype=dtypes.weakint) -> UOp: return UOp.const(x, dtype)
 def to_max_shape(shape:tuple[sint, ...]) -> tuple[int, ...]: return tuple(int(x.vmax) if isinstance(x, UOp) else x for x in shape)

-def select_dtype(u:UOp):
-  if u.dtype is dtypes.weakfloat: return dtypes.default_float
-  return dtypes.long if u.overflows(dtypes.int32) else dtypes.int
-def lower_weak_node(u:UOp) -> UOp|None:
-  start, src = (1 if u.op is Ops.WHERE else 0), tuple(s.src[0] if s.op is Ops.CAST and s.dtype in dtypes.weaks else s for s in u.src)
-  if src == u.src or any(s.dtype in dtypes.weaks for s in src[start:]): return None
-  dt = strong_dtype(least_upper_dtype(select_dtype(u), *(s.dtype for s in src)) if u.op in GroupOp.Binary
-                    else unwrap(dtype_from_uop(u.op, src, u.arg)))
-  return u.replace(dtype=None, src=src[:start]+tuple(s if s.base.is_invalid else s.cast(dt) for s in src[start:])).cast(u.dtype)
-pm_lower_weak = PatternMatcher([
-  (UPat(Ops.CONST, dtype=dtypes.weaks, name="u"), lambda u: UOp.const(u.val, select_dtype(u)).cast(u.dtype)),
-  # two stacked weak casts are a weakint value used as weakfloat (or vice versa): resolve the inner one at the outer kind's default.
-  # a SINGLE weak cast is never rewritten here, each consumer absorbs it on its own edge (see lower_weak_srcs)
-  (UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat.var("x"),)),), name="u"),
-   lambda u,x: x.cast(select_dtype(u)).cast(u.dtype) if x.dtype not in dtypes.weaks else None),
-  # Binary can widen from the bounds, all other nodes derive from the lowered sources.
-  # a weakfloat Unary (sin/exp2/...) must resolve here, before the transcendental decomposition
-  (UPat(GroupOp.Binary|GroupOp.Unary|{Ops.WHERE, Ops.RANGE, Ops.STACK, Ops.SPECIAL}, name="u"), lower_weak_node),
-  (UPat(Ops.PARAM, dtype=dtypes.weakint, name="u"),
-    lambda u: u.replace(dtype=None, arg=replace(u.arg, dtype=select_dtype(u))).cast(dtypes.weakint) if u.addrspace == AddrSpace.ALU else None),
-])
-def lower_weak_srcs(ctx:dict[UOp, UOp]|None, u:UOp) -> UOp|None:
-  if ctx is None: ctx = {}
-  def lower(s:UOp) -> UOp:
-    if (r:=ctx.get(s)) is None:
-      r = graph_rewrite(s, pm_lower_weak)
-      # the consumer absorbs the cast on its own edge
-      ctx[s] = r = r.src[0] if r.op is Ops.CAST and r.dtype in dtypes.weaks else r
-    return r
-  # a comparison demands a common operand width: lower it whole so the Binary rule unifies its operands
-  ret = lower(u) if u.op in GroupOp.Comparison else u.replace(src=tuple(lower(s) if s.dtype in dtypes.weaks else s for s in u.src))
-  return None if ret is u else ret
-
-def commit_weak(s:UOp, dt:DType) -> UOp:
-  # a bare weak CONST commits directly (the value stays mathematical, emission truncates), a weak non-const src takes the demand cast
-  return UOp.const(s.val, dt) if s.op is Ops.CONST else s.cast(dt)
-
-def commit_weak_srcs(u:UOp) -> UOp|None:
-  if not any(s.dtype in dtypes.weaks for s in u.src): return None
-  if (dt:=least_upper_dtype(*(s.dtype for s in u.src))) in dtypes.weaks: return None
-  # the root re-derives: a shift's dtype is its lhs's, so committing the lhs commits the node too
-  return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src))
-
-# runs in index lowering and in the decomps: a rule that mints a weak const commits it in the same rewrite, so none reaches the renderer
-pm_commit_weak = PatternMatcher([
-  (UPat(GroupOp.Broadcastable, name="u"), commit_weak_srcs),
-  # demand from the destination: a STORE's weak value commits at the destination's dtype
-  (UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.weaks)), allow_any_len=True, name="u"),
-   lambda u: u.replace(src=(u.src[0], commit_weak(u.src[1], u.src[0].dtype), *u.src[2:]))),
-])
-
-# a concrete CAST over a weak node states the width the value will live at. that width is a floor, never a narrowing
-def cast_weak_srcs(c:UOp, u:UOp) -> UOp|None:
-  if c.dtype in dtypes.weaks or weak_dtype(c.dtype) is not u.dtype: return None
-  dt = least_upper_dtype(c.dtype, select_dtype(u))
-  return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src)).cast(c.dtype)
-
-pm_cast_weak = PatternMatcher([
-  (UPat(Ops.CAST, name="c", src=(UPat(GroupOp.ALU, dtype=dtypes.weaks, name="u"),)), cast_weak_srcs),
-])
-
-pm_lower_index_dtype = pm_commit_weak+pm_cast_weak+PatternMatcher([
-  (UPat(GroupOp.All, name="u"),
-   lambda ctx,u: lower_weak_srcs(ctx, u) if u.dtype not in dtypes.weaks and any(s.dtype in dtypes.weaks for s in u.src) else None),
-  # a valid index into an n-element buffer lives in [0,n): a gated long index narrows when n-1 fits int32 (out-of-gate wraps, discarded)
-  # TODO: more generic
-  (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"), UPat.var("gate").where(UPat.var("idx", dtypes.long), UPat(Ops.CONST, arg=Invalid))),
-        allow_any_len=True, name="u"),
-   lambda u,buf,gate,idx: u.replace(src=(buf, idx.cast(dtypes.int).valid(gate))+u.src[2:]) if buf.max_numel()-1 <= dtypes.int32.max else None),
-])
-
 _substitute = PatternMatcher([(UPat(tuple(Ops), name="x"), lambda ctx,x: ctx.get(x,None))])
 _pm_resolve_params = PatternMatcher([(UPat(Ops.PARAM, name="p"), lambda ctx,p: ctx[p.arg.slot])])
 remove_all_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
@@ -1840,10 +1758,14 @@ pm_unbind = PatternMatcher([(UPat(Ops.BIND, name="x"), do_unbind)])

 # ctx is source UOp for which we are finding a contiguous view for. used in contiguous_view_offset
 pm_contiguous_view_offset = PatternMatcher([
-  (UPat(Ops.INDEX, src=(UPat(),)), lambda: UOp.const(0)),
-  (UPat(Ops.INDEX, src=(UPat(), UPat(Ops.RANGE))), lambda: UOp.const(0)),
-  (UPat(Ops.INDEX, src=(UPat(), UPat(Ops.RANGE)+UPat.cvar('c'))), lambda c: c),
-  (UPat(Ops.INDEX, src=(UPat(), UPat.cvar('c'))), lambda ctx, c: c if resolve(ctx.numel() == 1, False) else None),
+  # normalize to 1d bitcasts
+  (UPat(Ops.BITCAST, name="b"), lambda b: b.src[0].flatten().bitcast(b.dtype).reshape(b.shape) if len(b.shape) != 1 else None),
+  (UPat(Ops.BITCAST, name="b").index(UPat.cvar("c")), lambda ctx, b, c:
+   b.src[0].flatten().index(UOp.range(ctx.numel() * (osz:=b.element_size())//(isz:=b.src[0].element_size()), 0) + (c * osz//isz)) if b.tag else None),
+  (UPat(Ops.INDEX, src=(UPat.var("b"),)), lambda b: b.rtag().index(0)),
+  (UPat(Ops.INDEX, src=(UPat.var("b"), UPat(Ops.RANGE))), lambda b: b.rtag().index(0)),
+  (UPat(Ops.INDEX, src=(UPat.var("b"), UPat(Ops.RANGE)+UPat.cvar('c'))), lambda ctx, b, c: b.rtag().index(c)),
+  (UPat(Ops.INDEX, src=(UPat.var("b"), UPat.cvar('c'))), lambda ctx, b, c: b.rtag().index(c) if resolve(ctx.numel() == 1, False) else None),
 ])

 # *** what was symbolic.py ***
@@ -32,8 +32,8 @@ def validate_index(uidx:UOp, gate:UOp|None=None):
  from tinygrad.uop.validate import validate_index_with_z3
  return validate_index_with_z3(sz, idx, gate)

-def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher):
-  lst = list(ast.toposort()) if isinstance(ast, UOp) else ast
+def type_verify(ast:UOp|list[UOp], check_spec:PatternMatcher, enter_calls=True):
+  lst = list(ast.toposort(enter_calls=enter_calls)) if isinstance(ast, UOp) else ast
  if SPEC > 1: test_pyrender(lst[-1])  # assume this is the sink

  with Context(TRACK_MATCH_STATS=0):
@@ -233,13 +233,6 @@ spec_hcq = PatternMatcher([
 spec_full = PatternMatcher([
  (UPat(Ops.REWRITE_ERROR, dtypes.void, name="x"), lambda x: isinstance(x.arg, str)),

-  # SLICE on BUFFER is allowed if BUFFER is
-  (UPat(Ops.SLICE, src=(UPat(GroupOp.Movement.union({Ops.BUFFER, Ops.PARAM, Ops.STAGE, Ops.AFTER})),
-                        UPat(Ops.CONST, dtype=dtypes.weakint)), allow_any_len=True, name="bv"),
-   lambda bv: isinstance(bv.arg, int)),
-
-  (UPat(Ops.CALL, dtypes.void, src=(UPat((Ops.SLICE,)),), allow_any_len=True), lambda: True),
-
  # codegen may end ranges after gpudims has replaced RANGE with SPECIAL.
  (UPat(Ops.END, src=(UPat(), UPat()), allow_any_len=True), lambda: True),

@@ -253,15 +246,44 @@ spec_full = PatternMatcher([
  (UPat(Ops.BIND, (dtypes.int, dtypes.weakint), (UPat(), UPat()), arg=None), lambda: True),
 ])+spec_tensor+spec_program+spec_hcq

+# ***** kernel graph spec *****
+
+spec_kernel_graph = PatternMatcher([
+  # sink
+  (UPat(Ops.SINK, dtypes.void), lambda: True),
+  # bind
+  (UPat(Ops.BIND), lambda: True),
+  # const + stack to make vconsts
+  (UPat(Ops.CONST, src=()), lambda: True),
+  (UPat(Ops.STACK, src=()), lambda: True),
+  (UPat(Ops.STACK, src=UPat((Ops.CONST, Ops.BIND, Ops.PARAM))), lambda: True),
+  # linear for more kernels (TODO: we should enter non sink calls)
+  #(UPat(Ops.LINEAR), lambda: True),
+  # param is outside buffer, buffer is local buffer
+  (UPat(Ops.PARAM, name="x"), lambda x: isinstance(x.arg, ParamArg)),
+  (UPat(Ops.BUFFER, name="x"), lambda x: isinstance(x.arg, ParamArg) and x.addrspace == AddrSpace.GLOBAL),
+  # RESHAPE/BITCAST are NOOPs in the kernel graph (do we need them?)
+  (UPat((Ops.RESHAPE, Ops.BITCAST)), lambda: True),
+  # mstack/mselect
+  (UPat(Ops.MSTACK, name="x"), lambda x: all(isinstance(s.device, str) for s in x.src) or (all_same(x.src) and x.src[0].device is None)),
+  (UPat(Ops.MSELECT, name="x"), lambda x: isinstance(x.src[0].device, tuple) and x.arg < len(x.src[0].device)),
+  # all calls are on various sinks
+  (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.LINEAR, Ops.PROGRAM, Ops.CUSTOM_FUNCTION)),), allow_any_len=True), lambda: True),
+  # after on PARAM or AFTER
+  (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.PARAM, Ops.AFTER, Ops.BUFFER, Ops.MSTACK, Ops.MSELECT, Ops.BITCAST, Ops.RESHAPE})),),
+        allow_any_len=True, name="x"), lambda x: matches_dtype(x.src[0], x.dtype)),
+])
+
 # **** pyrender (move this) ****

 # late imports to avoid circular import
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.schedule.rangeify import BufferizeOpts
+from tinygrad.renderer import Estimates
 glbls:dict[str, Any] = {"inf": math.inf, "nan": math.nan, "KernelInfo": KernelInfo, "Metadata": Metadata,
                        "UOp": UOp, "dtypes": dtypes, "Ops": Ops, "AxisType": AxisType, "Invalid": Invalid,
                        "Opt": Opt, "OptOps": OptOps, "BufferizeOpts": BufferizeOpts, "AddrSpace": AddrSpace, "panic": panic,
-                        "ConstFloat": ConstFloat, "ParamArg": ParamArg}
+                        "ConstFloat": ConstFloat, "ParamArg": ParamArg, "Estimates": Estimates}
 def eval_pyrender(code:str) -> UOp:
  lcls:dict[str, Any] = {}
  exec(code, glbls, lcls)
@@ -96,6 +96,10 @@ pm_remove_invalid = PatternMatcher([
   if any(x.is_invalid for x in s.src) else None),
 ])

+# the one rule that collapses the pair CAST(dt, CONST(v)) into a typed CONST
+# TODO: delete this once CONST has no dtype
+pm_fold_cast_const = PatternMatcher([(UPat(Ops.CAST, name="root", src=(UPat.cvar("c"),)), lambda root, c: root.const_like(c.val))])
+
 symbolic_simple = pm_data_invalid + PatternMatcher([
  # ** self folding **
  (UPat.var("x") + 0, lambda x: x),    # x+0 -> x
@@ -152,8 +156,6 @@ symbolic_simple = pm_data_invalid + PatternMatcher([
  (UPat.var("x") * 0, lambda x: x.const_like(float("nan") if x.op is Ops.CONST
                                             and isinstance(x.val, float) and (math.isnan(x.val) or math.isinf(x.val)) else 0)),
  # *** cast/bitcast ***
-  # TODO: delete this once CONST has no dtype
-  (UPat(Ops.CAST, name="root", src=(UPat.cvar("c"),)), lambda root, c: root.const_like(c.val)),
  (UPat((Ops.CAST, Ops.BITCAST), name="root"), lambda root: root.src[0] if root.dtype == root.src[0].dtype else None),
  (UPat(Ops.BITCAST, name="root", src=(UPat.cvar("c"),)), fold_bitcast),
  # b.cast(a).cast(b) -> b if a preserves all values in b
@@ -253,10 +255,11 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
  # ** two stage ALU folding **
  *((UPat.var("x").alu(op, UPat.cvar("c1")).alu(op, UPat.cvar("c2")).named("f"),
     lambda f,x,c1,c2: x.alu(f.op,c1.alu(f.op,c2))) for op in GroupOp.Associative),
-  ((UPat.cvar("c0") + UPat.var("x")) < UPat.cvar("c1"), lambda x,c0,c1: x<(c1-c0)),  # c0 + x < c1 -> x < c1 - c0
  # (x//c1)//c2 -> x//(c1*c2) for c2>0
  ((UPat.var("x") // UPat.cvar("c1")) // UPat.cvar("c2"), lambda x,c1,c2: x//(c1*c2) if c2.vmin>0 else None),
  # ** lt **
+  # c0+x<c1 -> x < c1-c0
+  ((UPat.cvar("c0") + UPat.var("x", dtype=dtypes.ints+(dtypes.weakint,))) < UPat.cvar("c1"), lambda x,c0,c1: x<(c1-c0)),
  # c0*x<c1 -> sign(c0)*x < ceil(c1/abs(c0))
  ((UPat.cvar("c0")*UPat.var("x", dtype=dtypes.weakint))<UPat.cvar("c1"),
   lambda x,c0,c1: (x if c0.val > 0 else -x)<-(-c1.val//abs(c0.val)) if abs(c0.val) > 1 else None),
@@ -285,13 +288,13 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
    (UOp.const(x.val) if x.op is Ops.CONST else x.cast(dtypes.int)).alu(u.op,
     UOp.const(y.val) if y.op is Ops.CONST else y.cast(dtypes.int)).cast(u.dtype)
    if not any(v.overflows(dtypes.int) for v in (u,x,y)) else None),
-  ((UPat.var("x", dtypes.weakint) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+c.cast(cast.dtype)),
+  ((UPat.var("x", dtypes.weakint) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+cast.const_like(c.val)),
  # only RANGE/IF/STORE/KERNEL have side effects
  (UPat(Ops.AFTER, name="x"), lambda x: x.replace(src=(x.src[0],)+
    tuple(dedup(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.CALL, Ops.FUNCTION, Ops.BARRIER, Ops.END, Ops.LINEAR, Ops.STAGE}
                        else y.src for y in x.src[1:]]))))),
-  # after with 1 src is just src[0]
-  (UPat(Ops.AFTER, src=(UPat.var("s"),)), lambda s: s),
+  # after/end with 1 src is just src[0]
+  (UPat((Ops.AFTER, Ops.END), src=(UPat.var("s"),)), lambda s: s),
 ])+div_and_mod_symbolic

 # ******** we take a small aside to "simplify_valid" to rewrite valids ********
@@ -0,0 +1,78 @@
+from dataclasses import replace
+from tinygrad.dtype import dtypes, DType, AddrSpace, Invalid, least_upper_dtype, strong_dtype, weak_dtype
+from tinygrad.helpers import unwrap
+from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, GroupOp, graph_rewrite, dtype_from_uop
+
+def select_dtype(u:UOp):
+  if u.dtype is dtypes.weakfloat: return dtypes.default_float
+  return dtypes.long if u.overflows(dtypes.int32) else dtypes.int
+
+def lower_weak_node(u:UOp) -> UOp|None:
+  start, src = (1 if u.op is Ops.WHERE else 0), tuple(s.src[0] if s.op is Ops.CAST and s.dtype in dtypes.weaks else s for s in u.src)
+  if src == u.src or any(s.dtype in dtypes.weaks for s in src[start:]): return None
+  dt = strong_dtype(least_upper_dtype(select_dtype(u), *(s.dtype for s in src)) if u.op in GroupOp.Binary
+                    else unwrap(dtype_from_uop(u.op, src, u.arg)))
+  return u.replace(dtype=None, src=src[:start]+tuple(s if s.base.is_invalid else s.cast(dt) for s in src[start:])).cast(u.dtype)
+
+pm_lower_weak = PatternMatcher([
+  (UPat(Ops.CONST, dtype=dtypes.weaks, name="u"), lambda u: UOp.const(u.val, select_dtype(u)).cast(u.dtype)),
+  # two stacked weak casts are a weakint value used as weakfloat (or vice versa): resolve the inner one at the outer kind's default.
+  # a SINGLE weak cast is never rewritten here, each consumer absorbs it on its own edge (see lower_weak_srcs)
+  (UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat.var("x"),)),), name="u"),
+   lambda u,x: x.cast(select_dtype(u)).cast(u.dtype) if x.dtype not in dtypes.weaks else None),
+  # Binary can widen from the bounds, all other nodes derive from the lowered sources.
+  # a weakfloat Unary (sin/exp2/...) must resolve here, before the transcendental decomposition
+  (UPat(GroupOp.Binary|GroupOp.Unary|{Ops.WHERE, Ops.RANGE, Ops.STACK, Ops.SPECIAL}, name="u"), lower_weak_node),
+  (UPat(Ops.PARAM, dtype=dtypes.weakint, name="u"),
+    lambda u: u.replace(dtype=None, arg=replace(u.arg, dtype=select_dtype(u))).cast(dtypes.weakint) if u.addrspace == AddrSpace.ALU else None),
+])
+
+def lower_weak_srcs(ctx:dict[UOp, UOp]|None, u:UOp) -> UOp|None:
+  if ctx is None: ctx = {}
+  def lower(s:UOp) -> UOp:
+    if (r:=ctx.get(s)) is None:
+      r = graph_rewrite(s, pm_lower_weak)
+      # the consumer absorbs the cast on its own edge
+      ctx[s] = r = r.src[0] if r.op is Ops.CAST and r.dtype in dtypes.weaks else r
+    return r
+  # a comparison demands a common operand width: lower it whole so the Binary rule unifies its operands
+  ret = lower(u) if u.op in GroupOp.Comparison else u.replace(src=tuple(lower(s) if s.dtype in dtypes.weaks else s for s in u.src))
+  return None if ret is u else ret
+
+def commit_weak(s:UOp, dt:DType) -> UOp:
+  # a bare weak CONST commits directly (the value stays mathematical, emission truncates), a weak non-const src takes the demand cast
+  return UOp.const(s.val, dt) if s.op is Ops.CONST else s.cast(dt)
+
+def commit_weak_srcs(u:UOp) -> UOp|None:
+  if not any(s.dtype in dtypes.weaks for s in u.src): return None
+  if (dt:=least_upper_dtype(*(s.dtype for s in u.src))) in dtypes.weaks: return None
+  # the root re-derives: a shift's dtype is its lhs's, so committing the lhs commits the node too
+  return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src))
+
+# runs in index lowering and in the decomps: a rule that mints a weak const commits it in the same rewrite, so none reaches the renderer
+pm_commit_weak = PatternMatcher([
+  (UPat(GroupOp.Broadcastable, name="u"), commit_weak_srcs),
+  # demand from the destination: a STORE's weak value commits at the destination's dtype
+  (UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.weaks)), allow_any_len=True, name="u"),
+   lambda u: u.replace(src=(u.src[0], commit_weak(u.src[1], u.src[0].dtype), *u.src[2:]))),
+])
+
+# a concrete CAST over a weak node states the width the value will live at. that width is a floor, never a narrowing
+def cast_weak_srcs(c:UOp, u:UOp) -> UOp|None:
+  if c.dtype in dtypes.weaks or weak_dtype(c.dtype) is not u.dtype: return None
+  dt = least_upper_dtype(c.dtype, select_dtype(u))
+  return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src)).cast(c.dtype)
+
+pm_cast_weak = PatternMatcher([
+  (UPat(Ops.CAST, name="c", src=(UPat(GroupOp.ALU, dtype=dtypes.weaks, name="u"),)), cast_weak_srcs),
+])
+
+pm_lower_index_dtype = pm_commit_weak+pm_cast_weak+PatternMatcher([
+  (UPat(GroupOp.All, name="u"),
+   lambda ctx,u: lower_weak_srcs(ctx, u) if u.dtype not in dtypes.weaks and any(s.dtype in dtypes.weaks for s in u.src) else None),
+  # a valid index into an n-element buffer lives in [0,n): a gated long index narrows when n-1 fits int32 (out-of-gate wraps, discarded)
+  # TODO: more generic
+  (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"), UPat.var("gate").where(UPat.var("idx", dtypes.long), UPat(Ops.CONST, arg=Invalid))),
+        allow_any_len=True, name="u"),
+   lambda u,buf,gate,idx: u.replace(src=(buf, idx.cast(dtypes.int).valid(gate))+u.src[2:]) if buf.max_numel()-1 <= dtypes.int32.max else None),
+])
@@ -120,9 +120,9 @@ const drawGraph = (data) => {
    .attr("transform", d => `translate(${d.width/2-8}, ${-d.height/2+8})`).datum(e => ({ rect:true, width:10, height:10, fill:e.addrspace, stroke:"none" })));
  const CALL_TAG_WIDTH = 14;
  addTags(nodes.selectAll("g.type").data(d => d.collapsible ? [d] : []).join("g").attr("class", d => `tag clickable ${d.collapsed ? 'collapsed' : 'expanded'}`)
-    .attr("transform", d => d.callNode ? `translate(${CALL_TAG_WIDTH/2-d.width/2}, ${0})` : `translate(${-d.width/2}, ${0})`)
-    .datum(d => ({ ...d, text:d.collapsed ? "+" : "−", fill:d.callNode ? null : d.color,
-      ...(d.callNode && { rect:true, width:CALL_TAG_WIDTH }) })).on("click", (e,d) => {
+    .attr("transform", d => d.collapsePorts != null ? `translate(${CALL_TAG_WIDTH/2-d.width/2}, ${0})` : `translate(${-d.width/2}, ${0})`)
+    .datum(d => ({ ...d, text:d.collapsed ? "+" : "−", fill:d.collapsePorts != null ? null : d.color,
+      ...(d.collapsePorts != null && { rect:true, width:CALL_TAG_WIDTH }) })).on("click", (e,d) => {
      e.stopPropagation();
      const t = d3.zoomTransform(document.getElementById("graph-svg"));
      const [x, y] = t.apply([d.x, d.y]);
@@ -54,15 +54,17 @@ const layoutUOp = (g, { graph, change }, opts) => {
      width = Math.max(width, ctx.measureText(line).width);
      height += lineHeight;
    }
-    const callNode = label.startsWith("CALL\n") || label.startsWith("FUNCTION\n");
+    const op = label.split("\n", 1)[0];
+    const callNode = op === "CALL" || op === "FUNCTION", programNode = op === "PROGRAM";
+    const collapsePorts = callNode ? [0] : programNode ? [0, 1] : null;
    if (callNode) callCount++;
-    g.setNode(k, {...rectDims(width, height), label, labelX:0, ref, id:k, color, tag, callNode, exclude, addrspace,
+    g.setNode(k, {...rectDims(width, height), label, labelX:0, ref, id:k, color, tag, callNode, collapsePorts, exclude, addrspace,
      className:label.startsWith("REWRITE_ERROR") ? "err" : null});
    // add edges
    const edgeCounts = {};
    for (const [_, s] of src) edgeCounts[s] = (edgeCounts[s] || 0)+1;
    for (const [port, s] of src) g.setEdge(s, k, { label: edgeCounts[s] > 1 ? {type:"tag", text:edgeCounts[s]} : {type:"port", text:port},
-      ...(callNode && port === 0 && {color:"#a0a1b8"})});
+      ...(collapsePorts?.includes(port) && {color:"#a0a1b8"})});
    if (change?.includes(parseInt(k))) g.setParent(k, "overlay");
  }
  // optionally hide nodes from the layout
@@ -87,11 +89,11 @@ const layoutUOp = (g, { graph, change }, opts) => {
      const consumer = g.node(consumerId);
      // add +- toggle if this consumer has collapsible sources
      const edge = g.edge(n, consumerId);
-      const collapsible = consumer.callNode ? edge?.label?.text === 0 : node.exclude;
+      const collapsible = consumer.collapsePorts != null ? consumer.collapsePorts.includes(edge?.label?.text) : node.exclude;
      if (!collapsible) continue;
      consumer.collapsible = true;
-      // increase width of call/function nodes to make space for a toggle
-      if (consumer.callNode) { consumer.width = consumer.labelWidth+NODE_PADDING*2+CALL_TAG_WIDTH; consumer.labelX = CALL_TAG_WIDTH/2; }
+      // increase width of call/function/program nodes to make space for a toggle
+      if (consumer.collapsePorts != null) { consumer.width = consumer.labelWidth+NODE_PADDING*2+CALL_TAG_WIDTH; consumer.labelX = CALL_TAG_WIDTH/2; }
      // make sources invisible if UI has toggled it off
      const collapsed = consumer.callNode ? opts.showCallSrc === opts.callSrcMask.has(consumerId) : !opts.expandedNodes.has(consumerId);
      if (!collapsed) continue;
@@ -50,7 +50,7 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
               Ops.INDEX: "#CEF9B7", Ops.STACK: "#D8F9E4",
               Ops.WMMA: "#efefc0", Ops.UNSHARD: "#f6ccff", Ops.INS: "#eec4ff",
               **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
-               Ops.SLICE: "#E5EAFF", Ops.BUFFER: "#B0BDFF", Ops.GETADDR: "#9DB1F0", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6",
+               Ops.BUFFER: "#B0BDFF", Ops.GETADDR: "#9DB1F0", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6",
               Ops.CALL: "#00B7C8", Ops.FUNCTION: "#C07788", Ops.PARAM: "#14686F", Ops.SOURCE: "#c0c0c0", Ops.BINARY: "#404040",
               Ops.LINEAR: "#7DF4FF",
               Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",
Author	SHA1	Message	Date
geohot	ad7e744382	vf MI350X passthrough mode driver	2026-08-11 19:01:36 -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