direct lds

work
bench
2026-08-15 07:38:26 +00:00 · 2026-07-31 09:03:32 -07:00 · 2026-07-30 21:46:06 -07:00 · 2026-07-30 20:55:22 -07:00 · 2026-07-30 20:39:25 -07:00 · 2026-07-30 18:06:34 -07:00
191 changed files with 3861 additions and 9278 deletions
@@ -41,6 +41,10 @@ inputs:
    description: "Install LLVM?"
    required: false
    default: 'false'
+  tinydreno:
+    description: "Install tinydreno"
+    required: false
+    default: 'false'
  qemu:
    description: "Install qemu"
    required: false
@@ -273,6 +277,12 @@ runs:
      shell: bash
      run: brew install llvm@20

+    # *** tinydreno ***
+    - name: Install tinydreno (linux)
+      if: inputs.tinydreno == 'true' && runner.os == 'Linux'
+      shell: bash
+      run: sudo curl -fL https://github.com/sirhcm/tinydreno/raw/refs/heads/master/libllvm-qcom.so -o /usr/lib/libllvm-qcom.so
+
    # *** OpenCL ***
    - name: Install rusticl
      if: inputs.opencl == 'true'
@@ -94,7 +94,6 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -105,9 +104,6 @@ jobs:
        ./extra/amdpci/setup_python_cap.sh
        ./extra/hcq/hcq_smi.py amd rmmod
        ./extra/hcq/hcq_smi.py amd kill_pids
-    - name: Setup (NV)
-      if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -149,7 +145,6 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -160,9 +155,6 @@ jobs:
        ./extra/amdpci/setup_python_cap.sh
        ./extra/hcq/hcq_smi.py amd rmmod
        ./extra/hcq/hcq_smi.py amd kill_pids
-    - name: Setup (NV)
-      if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -202,7 +194,6 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -213,9 +204,6 @@ jobs:
        ./extra/amdpci/setup_python_cap.sh
        ./extra/hcq/hcq_smi.py amd rmmod
        ./extra/hcq/hcq_smi.py amd kill_pids
-    - name: Setup (NV)
-      if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
    - name: Symlink models and datasets
      run: |
        mkdir -p extra/datasets
@@ -252,7 +240,6 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -263,9 +250,6 @@ jobs:
        ./extra/amdpci/setup_python_cap.sh
        ./extra/hcq/hcq_smi.py amd rmmod
        ./extra/hcq/hcq_smi.py amd kill_pids
-    - name: Setup (NV)
-      if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -308,9 +292,6 @@ jobs:
        ./extra/amdpci/setup_python_cap.sh
        ./extra/hcq/hcq_smi.py amd rmmod
        ./extra/hcq/hcq_smi.py amd kill_pids
-    - name: Setup (NV)
-      if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -534,12 +515,14 @@ jobs:
      run: |
        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: openpilot compile3 big_driving_supercombo
-      run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo PICKLE_OOB=1 PYTHONPATH="." TC_OPT=2 GMMU=0 DEV=USB+AMD:LLVM ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/10926f2c0911821ca0e72439c1c3bf3ec11f0a08789aa14b7ee8f25379b2afa4 openpilot.pkl
-    - name: openpilot load_pickle big_driving_supercombo
-      run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_load_pickle PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_LOAD_TIME=25 python3 examples/openpilot/load_pickle.py openpilot.pkl
-    - name: openpilot run_pickle big_driving_supercombo
-      run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_run_pickle RUN_PICKLE=1 PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py - openpilot.pkl
+    - name: reset chestnut
+      run: python3 extra/usbgpu/debug.py -rn
+    - name: openpilot compile3 0.10.1 driving_vision
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." GMMU=0 DEV=USB+AMD:LLVM ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot load_pickle 0.10.1 driving_vision
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
+    - name: openpilot run_pickle 0.10.1 driving_vision
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py
    - name: Test copy speeds
      run: SIZE=64e6 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds

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

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

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

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

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

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

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

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

    steps:
@@ -219,8 +219,8 @@ jobs:
      run: python3 test/external/external_benchmark_schedule.py
    - name: Run process replay tests
      uses: ./.github/actions/process-replay
-    - name: Repo line count <= 26000 lines
-      run: MAX_LINE_COUNT=26000 python sz.py
+    - name: Repo line count < 25000 lines
+      run: MAX_LINE_COUNT=25000 python sz.py

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

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

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

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

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

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

  testoptim:
    name: Optimization Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 20
    steps:
      - name: Checkout Code
@@ -353,14 +353,12 @@ jobs:
        run: DEV=NULL NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
      - name: Test llama 3 training
        run: DEV=NULL NULL_ALLOW_COPYOUT=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
-      - name: Test gpt-oss training
-        run: DEV=NULL NULL_ALLOW_COPYOUT=1 SAMPLES=32 BS=2 SEQLEN=128 GRADIENT_ACC_STEPS=1 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 MXFP8=1 VOCAB_SIZE=32000 LAYERS=2 EXPERTS=4 MODEL=gptoss PYTHONPATH=. python3 examples/mlperf/model_train.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

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

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

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

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

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

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

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

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

    name: Linux (${{ matrix.backend }})
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
    timeout-minutes: 20
    env:
      FORWARD_ONLY: 1
@@ -631,23 +627,160 @@ 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
+        amd: 'true'
+        ocelot: 'true'
+    - name: Run unit tests
+      run: DEV=METAL python -m pytest -n=auto test/unit/ --durations=20
+    - name: Run NULL backend tests
+      run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --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 pytest (amd)
+      env:
+        DEV: MOCKKFD+AMD
+        FORWARD_ONLY: 1
+      run: |
+        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
+    - name: Run pytest (ptx)
+      env:
+        DEV: "MOCK+NV:PTX"
+        FORWARD_ONLY: 1
+        # TODO: failing due to library loading error
+        CAPTURE_PROCESS_REPLAY: 0
+      run: |
+        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
+  testmetal:
+    strategy:
+      fail-fast: false
+      matrix:
+        group: [1, 2]
+    name: MacOS (DEV=METAL) (${{ matrix.group }})
+    runs-on: macos-26
+    timeout-minutes: 20
+    env:
+      DEV: METAL
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: macos-metal
+        deps: testing_unit
+    - name: Check Device.DEFAULT and print some source
+      run: |
+        python -c "from tinygrad import Device; assert Device.DEFAULT == 'METAL'"
+        DEBUG=4 python test/test_tiny.py TestTiny.test_plus
+    - name: Run backend tests
+      run: python -m pytest -n=auto test/backend --durations=20 --splits 2 --group ${{ matrix.group }}
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
+  testmacos:
+    strategy:
+      fail-fast: false
+      matrix:
+        dev:
+          - 'CPU:CLANG'
+          - 'CPU:LLVM'
+          - 'CPU:LVP'
+          - 'WEBGPU'
+
+    name: MacOS (DEV=${{ matrix.dev }})
+    runs-on: macos-26
+    timeout-minutes: 20
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: macos-${{ matrix.dev }}
+        deps: "testing_unit${{ contains(matrix.dev, 'LVP') && ' mesa' || '' }}"
+        llvm: ${{ contains(matrix.dev, 'LLVM') || contains(matrix.dev, 'LVP') }}
+        webgpu: ${{ matrix.dev == 'WEBGPU' }}
+    - name: Set env
+      run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
+    - name: Check Device.DEFAULT and print some source
+      run: |
+        python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
+        DEBUG=4 python test/test_tiny.py TestTiny.test_plus
+    - name: Run test_tiny
+      run: python -m pytest -n=auto test/test_tiny.py --durations=20
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
+# ****** Windows Tests ******
+
+  testwindows:
+    strategy:
+      fail-fast: false
+      matrix:
+        dev:
+          - 'CPU:CLANG'
+          - 'CPU:LLVM'
+          - 'CPU:X86'
+          - 'WEBGPU'
+
+    name: Windows (DEV=${{ matrix.dev }})
+    runs-on: windows-2025
+    timeout-minutes: 15
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v6
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: windows-${{ matrix.dev }}-minimal
+          deps: testing_unit
+          pydeps: ${{ matrix.dev == 'WEBGPU' && 'dawn-python' || '' }}
+      - name: Set env
+        shell: bash
+        run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
+      - name: Check Device.DEFAULT and print some source
+        shell: bash
+        run: |
+          python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
+          DEBUG=4 python test/test_tiny.py TestTiny.test_plus
+      - name: Run test_tiny
+        shell: bash
+        run: python -m pytest -n=auto test/test_tiny.py --durations=20
+
 # ****** Compile-only Tests ******

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

  def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]) -> Tensor:
+    if mask is not None or start_pos.val == 0:
+      # no symbolic shape qkv when consuming prompts
+      start_pos = start_pos.val
+
    if HALF: x = x.half()
    xqkv = self.c_attn(x).reshape(None, None, 3, self.n_heads, self.head_dim)
    xq, xk, xv = [xqkv[:, :, i, :, :] for i in range(3)]
@@ -34,8 +38,12 @@ class Attention:
    # update the cache
    self.cache_kv[:, :, start_pos:start_pos+seqlen, :, :].assign(Tensor.stack(xk, xv)).realize()

-    keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
-    values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
+    if start_pos > 0:
+      keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
+      values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
+    else:
+      keys = xk
+      values = xv

    xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
    return self.c_proj(xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2).reshape(bsz, seqlen, self.dim))
@@ -78,14 +86,15 @@ class Transformer:
      seqlen = tokens.shape[1]
      tok_emb = self.wte(tokens)

-    # start_pos is a bound Variable, so everything below it stays symbolic
-    pos_emb = self.wpe(self.allpos.shrink((None, (start_pos, start_pos+seqlen))))
+    # not symbolic when consuming the prompt
+    selected_pos = (0, seqlen) if start_pos.val == 0 else (start_pos, start_pos+1)
+    pos_emb = self.wpe(self.allpos.shrink((None, selected_pos)))

    h = tok_emb + pos_emb

    if HALF: h = h.half()

-    mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype).triu(start_pos+1) if seqlen > 1 else None
+    mask = Tensor.full((1, 1, seqlen, start_pos.val+seqlen), float("-inf"), dtype=h.dtype).triu(start_pos.val+1) if seqlen > 1 else None

    for hi in self.h: h = hi(h, start_pos, mask)

@@ -1282,7 +1282,7 @@ def train_bert():
        previous_step = i

 def train_llama3():
-  from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8_DTYPE, MXFP8, MXFP4
+  from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8_DTYPE, MXFP8
  from examples.llama3 import MODEL_PARAMS
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
  from examples.mlperf.optim import GradAccClipAdamW, clip_grads
@@ -1434,9 +1434,9 @@ def train_llama3():
    load_state_dict(scheduler, safe_load(fn), realize=False)

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

  from tinygrad.nn.state import get_state_dict
@@ -1458,12 +1458,12 @@ def train_llama3():

  # realize everything here
  if optim.master_params: Tensor.realize(*optim.master_params)
-  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)
+  Tensor.realize(*optim.params, *fp8_inv_scales, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)

  @TinyJit
  def minibatch(tokens:Tensor):
-    model.reset_amax()
+    for nxt in fp8_next_amax: nxt.assign(0)
+    for nxt in fp8_next_grad_amax: nxt.assign(0)
    if is_dp: tokens = tokens.to(None).shard(device, 0)
    if is_mp: tokens = tokens.shard(device)
    if not is_sharding: tokens = tokens.to(None)
@@ -1477,8 +1477,8 @@ def train_llama3():
    for g, new_g in zip(grads, loss.gradient(*optim.params)):
      apply_grad(g, new_g.uop)

-    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)
+    loss_cpu = loss.flatten().float().to("CPU")
+    return loss_cpu.realize(*grads, *fp8_amax, *fp8_next_amax, *fp8_grad_amax, *fp8_next_grad_amax)

  @TinyJit
  def optim_step():
@@ -1487,14 +1487,14 @@ def train_llama3():
    scheduler.step()

    for g in grads: g.assign(0)
-    model.update_amax()
+    for cur, nxt in zip(fp8_amax, fp8_next_amax): cur.assign(nxt)
+    for cur, nxt in zip(fp8_grad_amax, fp8_next_grad_amax): cur.assign(nxt)

    lr_cpu = optim.lr.float().to("CPU")
    grad_norm_cpu = grad_norm.float().to("CPU")
-    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)
+    Tensor.realize(lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales, *fp8_amax, *fp8_grad_amax)

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

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

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

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

+      loss = sum(losses) / len(losses)
      optim_time = et - gt
      dev_time += optim_time
      step_time = et - st
@@ -1578,7 +1579,7 @@ def train_llama3():

      mem_gb = GlobalCounters.mem_used / 1e9
      gflops = GlobalCounters.global_ops / 1e9 / dev_time
-      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * (9.2e15 if MXFP4 else 4.6e15))) * 100
+      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * 4.6e15)) * 100
      tqdm.write(
          f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
          f"{lr:.12f} LR, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
@@ -1710,10 +1711,9 @@ def train_gptoss():
    wandb.init(config=config, **wandb_args, project="MLPerf-gpt-oss")

  model_params = GPT_OSS_20B
-  model_params['vocab_size'] = getenv("VOCAB_SIZE", 128256)
+  model_params['vocab_size'] = 128256
  real_vocab_size = model_params['vocab_size']
  if (layers:=getenv("LAYERS")) != 0: model_params['n_layers'] = layers
-  if (experts:=getenv("EXPERTS")) != 0: model_params['n_experts'] = experts
  print(f"model parameters: {model_params}")

  model = GPTOSS(**model_params, max_context=SEQLEN)
@@ -1748,10 +1748,7 @@ def train_gptoss():

  from extra.gemm.cdna_asm_gemm import _mx_block_scale
  model_state = get_state_dict(model)
-  def _scale_key(n):
-    if "." in n and (c:=f"{(b:=n.rsplit('.',1))[0]}_scale.{b[1]}") in model_state: return c
-    return f"{n}_scale"
-  fp8_scale_names = {n: _scale_key(n) for n, t in model_state.items() if t.dtype == FP8_DTYPE}
+  fp8_scale_names = {n: f"{n}_scale" for n, t in model_state.items() if t.dtype == FP8_DTYPE}
  fp8_inv_scales = [model_state[sname] for sname in fp8_scale_names.values()]
  for wname, sname in fp8_scale_names.items():
    w, scale = model_state[wname], model_state[sname]
@@ -25,7 +25,6 @@ FUSED_SILU_W13 = getenv("FUSED_SILU_W13", 0)
 SPLIT_W13 = getenv("SPLIT_W13", 0)
 COLUMNWISE_WEIGHT_SCALE = getenv("COLUMNWISE_WEIGHT_SCALE", 0)
 MXFP8 = getenv("MXFP8", 0)
-MXFP4 = getenv("MXFP4", 0)

 FP8_DTYPE = dtypes.fp8e4m3
 FP8_GRAD_DTYPE = dtypes.fp8e5m2
@@ -45,11 +44,6 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca
      from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
      if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T),)
    return (x @ w.T,)
-  if MXFP4:
-    assert x is not None, "MXFP4 matmul requires an unquantized input"
-    from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm
-    if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T, mxfp4=True),)
-    return (x @ w.T,)
  assert w_inv_scale is not None, "fp8 matmul requires w_inv_scale (weights must be stored in fp8 with per-tensor scale)"
  if MXFP8:
    from extra.gemm.cdna_asm_gemm import asm_gemm, quantize_mxfp8, mx_pack, can_use_asm_gemm, _mx_block_scale
@@ -83,9 +77,9 @@ def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_sca
      return out, x_fp8
  return (x_fp8.dot(w.T, dtype=dtypes.float) * ((amax_x.float() + 1e-8) / FP8_MAX) * w_inv_scale).cast(dtypes.bfloat16), x_fp8

-def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor|None,
-                         next_amax_x:Tensor|None, grad_amax_state:Tensor|None, next_grad_amax_state:Tensor|None):
-  if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
+def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
+                         next_amax_x:Tensor, grad_amax_state:Tensor, next_grad_amax_state:Tensor):
+  if FUSED_ADD_NORM_MUL_QUANTIZE:
    from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_rmsnorm_mul_quantize_fp8
    x_fp8, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
@@ -96,9 +90,9 @@ def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, ep
                     next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
  return out, x_normed, rrms, ret

-def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor|None,
-                             next_amax_x:Tensor|None, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
-  if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
+def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
+                             next_amax_x:Tensor, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
+  if FUSED_ADD_NORM_MUL_QUANTIZE:
    from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_add_rmsnorm_mul_quantize_fp8
    x_fp8, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
@@ -111,15 +105,10 @@ def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w
  return out, h, x_normed, rrms, ret

 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:
+                             amax_x2:Tensor, next_amax_x2:Tensor,
+                             grad_amax_xw13:Tensor, next_grad_amax_xw13:Tensor,
+                             grad_amax_xout:Tensor, next_grad_amax_xout:Tensor):
+  if FUSED_SILU_W13:
    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,
                                                 next_grad_amax_state=next_grad_amax_xw13, amax_out=next_amax_x2)
@@ -169,15 +158,14 @@ class FlatTransformer:
    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).clone().is_param_(False)

    def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().is_param_(False)
-    n_amax = 0 if MXFP4 else n_layers
    names = ["xqkv", "xo", "x2"]
    names += ["x1", "x3"] if SPLIT_W13 else ["x13"]
-    self._fp8_amax = {name: [_amax() for _ in range(n_amax)] for name in names}
-    self._fp8_next_amax = {name: [_amax() for _ in range(n_amax)] for name in names}
+    self._fp8_amax = {name: [_amax() for _ in range(n_layers)] for name in names}
+    self._fp8_next_amax = {name: [_amax() for _ in range(n_layers)] for name in names}
    grad_names = ["xqkv", "xo", "xout"]
    grad_names += ["xw1", "xw3"] if SPLIT_W13 else ["xw13"]
-    self._fp8_grad_amax = {name: [_amax() for _ in range(n_amax)] for name in grad_names}
-    self._fp8_next_grad_amax = {name: [_amax() for _ in range(n_amax)] for name in grad_names}
+    self._fp8_grad_amax = {name: [_amax() for _ in range(n_layers)] for name in grad_names}
+    self._fp8_next_grad_amax = {name: [_amax() for _ in range(n_layers)] for name in grad_names}
    w_scales = [("wqkv", s_qkv), ("wo", s_o), ("w2", s_2)]
    w_scales += [("w1", s_1), ("w3", s_3)] if SPLIT_W13 else [("w13", s_13)]
    self._fp8_inv_scale = {name: (s if MXFP8 else s.float()).contiguous().is_param_(False) for name, s in w_scales}
@@ -191,9 +179,6 @@ class FlatTransformer:
      from extra.gemm.cdna_asm_gemm import quantize_mxfp8
      w_q, w_e8, _ = quantize_mxfp8(w.reshape(self.n_layers * out_features, in_features))
      return w_q.reshape(self.n_layers, out_features, in_features), w_e8.reshape(self.n_layers, out_features, in_features // 32)
-    if MXFP4:
-      # FP4 is produced dynamically so optimizer updates always start from the current BF16 weight.
-      return w.cast(dtypes.bfloat16), Tensor.ones(self.n_layers)
    amax = (w.abs().max(axis=2) if COLUMNWISE_WEIGHT_SCALE else w.abs().flatten(1).max(1)).detach()
    scale = FP8_MAX / (amax + 1e-8)
    inv_scale = (amax + 1e-8) / FP8_MAX
@@ -201,10 +186,9 @@ class FlatTransformer:
    return (w * scale_b).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE), inv_scale

  def attention(self, x:Tensor, freqs_cis:Tensor, *, attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
-                amax_xqkv:Tensor|None, amax_xo:Tensor|None, s_qkv:Tensor, s_o:Tensor,
-                next_amax_xqkv:Tensor|None, next_amax_xo:Tensor|None,
-                grad_amax_xqkv:Tensor|None, grad_amax_xo:Tensor|None,
-                next_grad_amax_xqkv:Tensor|None, next_grad_amax_xo:Tensor|None):
+                amax_xqkv:Tensor, amax_xo:Tensor, s_qkv:Tensor, s_o:Tensor,
+                next_amax_xqkv:Tensor, next_amax_xo:Tensor,
+                grad_amax_xqkv:Tensor, grad_amax_xo:Tensor, next_grad_amax_xqkv:Tensor, next_grad_amax_xo:Tensor):
    bsz, seqlen, _ = x.shape
    saves = []

@@ -326,33 +310,28 @@ class FlatTransformer:
          for i in range(len(amax_dict[name])):
            amax_dict[name][i] = amax_dict[name][i].to(device).contiguous().is_param_(False)

-  def reset_amax(self):
-    for st in (self._fp8_next_amax, self._fp8_next_grad_amax):
-      for ts in st.values():
-        for t in ts: t.assign(0)
-
-  def update_amax(self):
-    for cur, nxt in ((self._fp8_amax, self._fp8_next_amax), (self._fp8_grad_amax, self._fp8_next_grad_amax)):
-      for name in cur:
-        for c, n in zip(cur[name], nxt[name]): c.assign(n)
-
  def __call__(self, tokens:Tensor, save:bool=True):
    h = self.tok_embeddings(tokens)
    freqs_cis = self.freqs_cis.cast(h.dtype)
    if not getenv("HK_FLASH_ATTENTION"): freqs_cis = freqs_cis[:, :tokens.shape[1], :, :, :]
    a, na, ga, nga, s = self._fp8_amax, self._fp8_next_amax, self._fp8_grad_amax, self._fp8_next_grad_amax, self._fp8_inv_scale
-    def amax_kwargs(i:int, act_names:tuple[str, ...], grad_names:tuple[str, ...]) -> dict[str, Tensor|None]:
-      specs = (("amax_", a, act_names), ("next_amax_", na, act_names), ("grad_amax_", ga, grad_names), ("next_grad_amax_", nga, grad_names))
-      if MXFP4: return dict.fromkeys(f"{prefix}{name}" for prefix, _, names in specs for name in names)
-      return {f"{prefix}{name}":val[name][i] for prefix, val, names in specs for name in names}
    for i in range(self.n_layers):
-      attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i], s_qkv=s["wqkv"][i], s_o=s["wo"][i],
-                         **amax_kwargs(i, ("xqkv", "xo"), ("xqkv", "xo")))
-      ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], w2=self.w2[i], s_2=s["w2"][i], **amax_kwargs(i, ("x2",), ("xout",)))
+      attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i],
+                         amax_xqkv=a["xqkv"][i], amax_xo=a["xo"][i], s_qkv=s["wqkv"][i], s_o=s["wo"][i],
+                         next_amax_xqkv=na["xqkv"][i], next_amax_xo=na["xo"][i],
+                         grad_amax_xqkv=ga["xqkv"][i], grad_amax_xo=ga["xo"][i],
+                         next_grad_amax_xqkv=nga["xqkv"][i], next_grad_amax_xo=nga["xo"][i])
+      ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], w2=self.w2[i],
+                        amax_x2=a["x2"][i], s_2=s["w2"][i], grad_amax_xout=ga["xout"][i], next_grad_amax_xout=nga["xout"][i],
+                        next_amax_x2=na["x2"][i])
      if SPLIT_W13:
-        ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i], s_1=s["w1"][i], s_3=s["w3"][i], **amax_kwargs(i, ("x1", "x3"), ("xw1", "xw3")))
+        ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i], amax_x1=a["x1"][i], amax_x3=a["x3"][i],
+                          next_amax_x1=na["x1"][i], next_amax_x3=na["x3"][i],
+                          s_1=s["w1"][i], s_3=s["w3"][i], grad_amax_xw1=ga["xw1"][i], grad_amax_xw3=ga["xw3"][i],
+                          next_grad_amax_xw1=nga["xw1"][i], next_grad_amax_xw3=nga["xw3"][i])
      else:
-        ffn_kwargs.update(w13=self.w13[i], s_13=s["w13"][i], **amax_kwargs(i, ("x13",), ("xw13",)))
+        ffn_kwargs.update(w13=self.w13[i], amax_x13=a["x13"][i], s_13=s["w13"][i], grad_amax_xw13=ga["xw13"][i],
+                          next_grad_amax_xw13=nga["xw13"][i], next_amax_x13=na["x13"][i])
      h, *_ = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)

    logits = matmul(self.norm(h), self.output[0], fp8=False)[0]
@@ -436,7 +415,9 @@ if __name__ == "__main__":
  @TinyJit
  def fwd_bwd(tokens:Tensor):
    with Timing("python forward: "):
-      model.reset_amax()
+      for amax_dict in (model._fp8_next_amax, model._fp8_next_grad_amax):
+        for ts in amax_dict.values():
+          for nxt in ts: nxt.assign(0)
      logits = model(tokens[:, :-1], save=llama_size=="8B")
      loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
    with Timing("python backward: "):
@@ -12,15 +12,11 @@ from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb
 from extra.llama_kernels.rmsnorm import rmsnorm
-from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8, asm_gemm, can_use_asm_gemm
-from extra.gemm.moe_gemm import grouped_mx_gemm
-from extra.gemm.moe_routing import route, dispatch, combine
+from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8

 FP8_DTYPE = dtypes.fp8e4m3
 FP8_MAX = 448.0
-INIT_STD = 0.02
-ASM_GEMM = getenv("ASM_GEMM", 0)
-
+INIT_STD = 0.008

 def _quant_dequant_fwd(x:Tensor) -> Tensor:
  # x (2d bf16) -> bf16 value after an mxfp8 round-trip (1x32 block scaling on the last axis)
@@ -63,34 +59,10 @@ def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:

 def matmul_mx(x:Tensor, w_q:Tensor, w_scale:Tensor) -> Tensor:
  l_shape = x.shape[:-1]
-  if ASM_GEMM:
-    from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm, mx_pack
-    x2, K, N = x.reshape(-1, x.shape[-1]), x.shape[-1], w_q.shape[0]
-    wq, ws = w_q, w_scale
-    if (pad := (-K) % 256):
-      x2 = x2.pad(((0, 0), (0, pad)))
-      wq = wq.pad(((0, 0), (0, pad)))
-      ws = ws.pad(((0, 0), (0, pad // 32)), value=127).cast(dtypes.uint8)
-    if (npad := (-N) % 256):
-      wq = wq.pad(((0, npad), (0, 0)))
-      ws = ws.pad(((0, npad), (0, 0)), value=127).cast(dtypes.uint8)
-    x_q, x_e8, x_si = quantize_mxfp8(x2)
-    if x_si is not None and can_use_asm_gemm(x_q, wq.T):
-      out = asm_gemm(x_q, wq.T, mx=True, mx_scales=(x_si, x_e8, mx_pack(ws), ws), mx_w_stored=True)
-      return (out[:, :N] if npad else out).reshape(*l_shape, N).cast(dtypes.bfloat16)
  x_phys = quant_dequant_mx(x.reshape(-1, x.shape[-1])).reshape(*l_shape, x.shape[-1])
  w_phys = dequant_weight(w_q, w_scale)
  return (x_phys @ w_phys.T).cast(dtypes.bfloat16)

-def _pad_to_mult(t:Tensor, axis:int, mult:int=256) -> Tensor:
-  if (r := (-t.shape[axis]) % mult) == 0: return t
-  pads = [(0, 0)] * t.ndim
-  pads[axis] = (0, r)
-  return t.pad(tuple(pads))
-
-def _pad_cols(t:Tensor) -> Tensor: return _pad_to_mult(t, -1)
-def _pad_rows(t:Tensor) -> Tensor: return _pad_to_mult(t, -2)
-
 def swiglu(x:Tensor, limit:float=7.0, alpha:float=1.702) -> Tensor:
  x_glu, x_linear = x[..., ::2], x[..., 1::2]
  x_glu = x_glu.clamp(max_=limit)
@@ -127,9 +99,9 @@ class GPTOSS:
    self.ffn_norm = Tensor.ones(n_layers, dim).contiguous()
    self.gate = Tensor.normal(n_layers, n_experts, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
    self.gate_bias = Tensor.zeros(n_layers, n_experts, dtype=dtypes.bfloat16).contiguous()
-    self.w_gate_up, self.w_gate_up_scale = self._quant_weight(n_layers, n_experts, intermediate_size * 2, dim, moe=True)
+    self.w_gate_up, self.w_gate_up_scale = self._quant_weight(n_layers, n_experts, intermediate_size * 2, dim)
    self.w_gate_up_bias = Tensor.zeros(n_layers, n_experts, intermediate_size * 2, dtype=dtypes.bfloat16).contiguous()
-    self.w_down, self.w_down_scale = self._quant_weight(n_layers, n_experts, dim, intermediate_size, std=scaled_std, moe=True)
+    self.w_down, self.w_down_scale = self._quant_weight(n_layers, n_experts, dim, intermediate_size, std=scaled_std)
    self.w_down_bias = Tensor.zeros(n_layers, n_experts, dim, dtype=dtypes.bfloat16).contiguous()

    # output
@@ -139,15 +111,10 @@ class GPTOSS:
    self.output = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
    self.freqs_cis = precompute_freqs_cis(head_dim, max_context * 2, rope_theta).contiguous().is_param_(False)

-  def _quant_weight(self, *shape:int, std:float=INIT_STD, moe:bool=False):
-    def _one(*s:int):
-      w = Tensor.zeros(*s) if getenv("ZEROS") else Tensor.normal(*s, mean=0.0, std=std)
-      w_q, w_e8, _ = quantize_mxfp8(_pad_cols(_pad_rows(w)) if moe else w)
-      return w_q, w_e8.is_param_(False)
-    if moe:
-      qs = [_one(*shape[1:]) for _ in range(shape[0])]
-      return [q[0] for q in qs], [q[1] for q in qs]
-    return _one(*shape)
+  def _quant_weight(self, *shape:int, std:float=INIT_STD):
+    w = Tensor.zeros(*shape) if getenv("ZEROS") else Tensor.normal(*shape, mean=0.0, std=std)
+    w_q, w_e8, _ = quantize_mxfp8(w)
+    return w_q, w_e8.is_param_(False)

  def _attn_mask(self, seqlen:int, dtype) -> Tensor:
    i, j = Tensor.arange(seqlen).reshape(seqlen, 1), Tensor.arange(seqlen).reshape(1, seqlen)
@@ -182,12 +149,12 @@ class GPTOSS:
    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
    xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)  # (B,N,H,D)/(B,N,KV,D)

-    if getenv("HK_FLASH_ATTENTION"):
-      from extra.thunder.amd.fa import flash_attention
-      attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks, window=self.sliding_window if sliding else 0)
-      attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
-    elif sliding:
+    if sliding:
      attn = self._sliding_attention(xq, xk, xv, sinks)
+    elif getenv("HK_FLASH_ATTENTION"):
+      from extra.thunder.amd.fa import flash_attention
+      attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks)
+      attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
    else:
      xqm = xq.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep, self.head_dim).permute(0, 2, 3, 1, 4)
      xkm, xvm = xk.permute(0, 2, 1, 3).unsqueeze(2), xv.permute(0, 2, 1, 3).unsqueeze(2)
@@ -206,32 +173,17 @@ class GPTOSS:
                   w_down:Tensor, w_down_scale:Tensor, w_down_bias:Tensor):
    x_normed, rrms = rmsnorm(x, self.norm_eps)
    inp = x_normed * ffn_norm
+
    logits = inp.float() @ gate.float().T + gate_bias.float()
-    dim, inter = self.dim, self.intermediate_size
+    thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
+    weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)

-    if getenv("GROUPED_MOE", 0):
-      bsz, seqlen = x.shape[:2]
-      inp, logits = inp.reshape(-1, dim), logits.reshape(-1, self.n_experts)
-      r = route(logits, self.experts_per_tok, self.n_experts)
-      onehot = r.rows_e.one_hot(self.n_experts).float()
-      xg = dispatch(_pad_cols(inp.cast(dtypes.bfloat16)), r)
-      h = grouped_mx_gemm(xg, (w_gate_up, w_gate_up_scale), r.off)[:, :2*inter] + (onehot @ w_gate_up_bias.float()).cast(dtypes.bfloat16)
-      y = swiglu(h, self.swiglu_limit)
-      z = grouped_mx_gemm(_pad_cols(y.cast(dtypes.bfloat16)), (w_down, w_down_scale), r.off)[:, :dim] \
-          + (onehot @ w_down_bias.float()).cast(dtypes.bfloat16)
-      out = combine(z, r, inp.shape[0], self.experts_per_tok).reshape(bsz, seqlen, dim)
-    else:
-      thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
-      weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)
-
-      out = None
-      for e in range(self.n_experts):
-        gu_q, gu_s = w_gate_up[e][:2*inter, :dim].contiguous(), w_gate_up_scale[e][:2*inter, :dim//32].contiguous()
-        dn_q, dn_s = w_down[e][:dim, :inter].contiguous(), w_down_scale[e][:dim, :inter//32].contiguous()
-        gate_up = matmul_mx(inp, gu_q, gu_s) + w_gate_up_bias[e]
-        y = (matmul_mx(swiglu(gate_up, self.swiglu_limit), dn_q, dn_s) + w_down_bias[e]).contiguous()
-        contrib = weights[..., e:e+1].cast(y.dtype) * y
-        out = contrib if out is None else out + contrib
+    out = None
+    for e in range(self.n_experts):
+      gate_up = matmul_mx(inp, w_gate_up[e], w_gate_up_scale[e]) + w_gate_up_bias[e]
+      y = (matmul_mx(swiglu(gate_up, self.swiglu_limit), w_down[e], w_down_scale[e]) + w_down_bias[e]).contiguous()
+      contrib = weights[..., e:e+1].cast(y.dtype) * y
+      out = contrib if out is None else out + contrib
    return out, [x_normed, rrms]

  @function(precompile=True, precompile_backward=True)
@@ -263,11 +215,7 @@ class GPTOSS:
                        w_down=self.w_down[i], w_down_scale=self.w_down_scale[i], w_down_bias=self.w_down_bias[i])
      h, *_ = self.run_layer(h, freqs_cis, mask_full, i % 2 == 0, attn_kwargs, ffn_kwargs, save=save)

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

 def _get_pads(uop:UOp) -> list[UOp]:
@@ -2,7 +2,7 @@ from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
 from tinygrad.nn.optim import Optimizer, OptimizerGroup
 from tinygrad.helpers import FUSE_OPTIM, getenv
-from tinygrad.uop.ops import UOp, Ops, AxisType
+from tinygrad.uop.ops import UOp, Ops

 STOCHASTIC_ROUND = getenv("STOCHASTIC_ROUND", 0)
 MASTER_WEIGHTS = getenv("MASTER_WEIGHTS", 0)
@@ -42,8 +42,8 @@ class GradAccClipAdamW(Optimizer):
      self.master_params = None

  def _zero_shard(self, t:Tensor) -> Tensor:
-    if not self.zero or t.ndim < 2 or (t.shape[0] % len(self.device)) != 0: return t
-    return Tensor(t.uop._shard(0, UOp.range(len(self.device), -1, AxisType.DEVICE)).unshard(0)).clone()
+    if not self.zero or (t.shape[0] % len(self.device)) != 0: return t
+    return Tensor(t.uop._shard(0, len(self.device)).unshard(0)).clone()

  def _zero_gather(self, t:Tensor) -> Tensor:
    if not isinstance(t.device, tuple) or t.uop.axis != 0: return t
@@ -123,9 +123,6 @@ class GradAccClipAdamW(Optimizer):
    return out.shard_like(t) if offloaded else out

 class GradAccClipAdamWGroup(OptimizerGroup):
-  def __init__(self, *optimizers:GradAccClipAdamW):
-    super().__init__(*optimizers)
-    for o in self.optimizers[1:]: o.lr = self.optimizers[0].lr
  def fstep(self, grads:list[Tensor], grad_norm:Tensor|None=None):
    offset = 0
    to_realize = []
@@ -1,8 +1,8 @@
 #!/usr/bin/env bash

 export PYTHONPATH="."
-export ROCM_PATH=${ROCM_PATH:-/opt/rocm-7.1.1}
-export PATH="$ROCM_PATH/bin:$PATH"
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=${DEV:-AMD}
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -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 MXFP4=${MXFP4:-1}
+export FP8=${FP8:-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="float32"
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
 export GBS=$((BS * GRADIENT_ACC_STEPS))

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

-export FAKEDATA=${FAKEDATA:-$([[ "$DEV" == NULL:* ]] && echo 1 || echo 0)} BENCHMARK=${BENCHMARK:-10}
+export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
 fi
@@ -1,8 +1,8 @@
 #!/usr/bin/env bash

 export PYTHONPATH="."
-export ROCM_PATH=${ROCM_PATH:-/opt/rocm-7.1.1}
-export PATH="$ROCM_PATH/bin:$PATH"
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=${DEV:-AMD}
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -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 MXFP4=${MXFP4:-1}
+export FP8=${FP8:-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="float32"
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
 export GBS=$((BS * GRADIENT_ACC_STEPS))

@@ -1,6 +1,4 @@
 #!/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
@@ -10,8 +10,6 @@ export DEVICE_IN_FUNCTION_BUG=1

 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
-export ASM_GEMM=${ASM_GEMM:-1}
-export GROUPED_MOE=${GROUPED_MOE:-1}
 export ALL2ALL=${ALL2ALL:-1}
 export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
@@ -10,8 +10,6 @@ export DEVICE_IN_FUNCTION_BUG=1

 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
-export ASM_GEMM=${ASM_GEMM:-1}
-export GROUPED_MOE=${GROUPED_MOE:-1}
 export ALL2ALL=${ALL2ALL:-1}
 export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
@@ -1,4 +1,4 @@
-import os, sys, pickle, time, re, tempfile, struct, shutil, io
+import os, sys, pickle, time, re
 import numpy as np
 if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"

@@ -9,39 +9,6 @@ from tinygrad.nn.onnx import OnnxRunner

 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
 OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
-PICKLE_OOB = getenv("PICKLE_OOB")
-
-def dump_pickle(obj, f):
-  if PICKLE_OOB:
-    # allows pickling when buffers don't fit in (CPU) RAM
-    # from openpilot/selfdrive/modeld/helpers.py
-    with tempfile.TemporaryFile(dir=".") as tmp:
-      def buffer_callback(pb: pickle.PickleBuffer):
-        m = pb.raw()
-        tmp.write(struct.pack('<q', m.nbytes))
-        tmp.write(m)
-        pb.release() # keep peak ram at ~1 buffer
-      stream = io.BytesIO()
-      pickle.Pickler(stream, protocol=5, buffer_callback=buffer_callback).dump(obj)
-      opcodes = stream.getvalue()
-      f.write(struct.pack('<q', len(opcodes)))
-      f.write(opcodes)
-      tmp.seek(0)
-      shutil.copyfileobj(tmp, f)
-  else: pickle.dump(obj, f)
-
-def load_pickle(f):
-  if PICKLE_OOB:
-    # allows unpickling when buffers don't fit in (CPU) RAM
-    # from openpilot/selfdrive/modeld/helpers.py
-    opcodes = f.read(struct.unpack('<q', f.read(8))[0])
-    def buffers():
-      while (h := f.read(8)):
-        pb = pickle.PickleBuffer(bytearray(struct.unpack('<q', h)[0]))
-        f.readinto(pb)
-        yield pb
-    return pickle.load(io.BytesIO(opcodes), buffers=buffers())
-  else: return pickle.load(f)

 def compile(onnx_file):
  run_onnx = OnnxRunner(onnx_file)
@@ -98,7 +65,8 @@ def compile(onnx_file):
  if (allowed_gated_read_image:=getenv("ALLOWED_GATED_READ_IMAGE", -1)) != -1:
    assert gated_read_image_count == allowed_gated_read_image, f"different gated read_image! {gated_read_image_count=}, {allowed_gated_read_image=}"

-  with open(OUTPUT, "wb") as f: dump_pickle(run_onnx_jit, f)
+  with open(OUTPUT, "wb") as f:
+    pickle.dump(run_onnx_jit, f)
  mdl_sz = os.path.getsize(onnx_file)
  pkl_sz = os.path.getsize(OUTPUT)
  print(f"mdl size is {mdl_sz/1e6:.2f}M")
@@ -168,7 +136,7 @@ def bench(run, inputs):

 if __name__ == "__main__":
  if getenv("RUN_PICKLE"):
-    with open(OUTPUT, "rb") as f: pickle_loaded = load_pickle(f)
+    with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
    inputs = {name: Tensor(Tensor.randn(*view.shape, dtype=dtype).numpy(), device=device)
              for name, (view, _vars, dtype, device) in zip(pickle_loaded.captured.expected_names, pickle_loaded.captured.expected_input_info)}
    test_vs_compile(pickle_loaded, inputs)
@@ -176,7 +144,7 @@ if __name__ == "__main__":
    onnx_file = fetch(OPENPILOT_MODEL)
    inputs, outputs = compile(onnx_file)

-    with open(OUTPUT, "rb") as f: pickle_loaded = load_pickle(f)
+    with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)

    test_vs_compile(pickle_loaded, inputs, outputs)
    if getenv("SELFTEST"):
@@ -1,6 +1,5 @@
-import sys
+import sys, pickle
 from extra.bench_log import WallTimeEvent, BenchEvent
-from examples.openpilot.compile3 import load_pickle
 from tinygrad.helpers import getenv

 PKL = sys.argv[1] if len(sys.argv) > 1 else "/tmp/openpilot.pkl"
@@ -8,7 +7,7 @@ PKL = sys.argv[1] if len(sys.argv) > 1 else "/tmp/openpilot.pkl"
 load_times = []

 for _ in range(10):
-  with WallTimeEvent(BenchEvent.STEP) as wte: load_pickle(open(PKL, 'rb'))
+  with WallTimeEvent(BenchEvent.STEP) as wte: pickle.load(open(PKL, 'rb'))
  load_times.append(wte.time)
  print(f"pickle load: {wte.time:6.2f} s")

@@ -1,31 +0,0 @@
-import argparse, time
-from tinygrad.llm.model import Transformer
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser()
-  parser.add_argument("--model", required=True, help="path to gguf model")
-  parser.add_argument("--max-context", type=int, default=8192, help="max context length (default: %(default)s)")
-  parser.add_argument("--prompt-tokens", type=int, default=1024, help="number of prompt tokens (default: %(default)s)")
-  parser.add_argument("--decode-tokens", type=int, default=16, help="number of tokens to decode (default: %(default)s)")
-  parser.add_argument("--chunk-size", type=int, default=32, help="chunk size for prefill (default: %(default)s)")
-  args = parser.parse_args()
-
-  st = time.perf_counter()
-  model, _ = Transformer.from_gguf(args.model, args.max_context)
-  print(f"load {time.perf_counter()-st:.3f}s", flush=True)
-
-  st = time.perf_counter()
-  model.warmup()
-  print(f"warm {time.perf_counter()-st:.3f}s", flush=True)
-
-  prompt = [257] + [1000+i%1000 for i in range(args.prompt_tokens-1)]
-  gen = model.generate(prompt, chunk_size=args.chunk_size)
-  st = time.perf_counter()
-  # first token is time-to-first-token; counted as part of prefill
-  output = [next(gen)]
-  pt = time.perf_counter()
-  print(f"prefill {args.prompt_tokens/(pt-st):.3f} tok/s", flush=True)
-
-  for _ in range(args.decode_tokens): output.append(next(gen))
-  et = time.perf_counter()
-  print(f"decode {args.decode_tokens/(et-pt):.3f} tok/s output {output}", flush=True)
@@ -35,7 +35,7 @@ def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], Li
    return name

  for call in iter_kernel_calls(linear):
-    arg_uops = [b for b in call.src[1:] if not b.is_bound_var]
+    arg_uops = [b for b in call.src[1:] if b.op is not Ops.BIND]
    prg = to_program(call.src[0], Device[arg_uops[0].device].renderer)
    info = prg.arg
    functions[info.function_name] = prg.src[2].arg
@@ -264,7 +264,7 @@ def export_model(model, target:str, *inputs, model_name: Optional[str] = "model"
        if getattr(dim, "op", None) is Ops.ADD and len(dim.src) == 2 and \
           any(s.op is Ops.PARAM and s.addrspace is AddrSpace.ALU for s in dim.src) and any(s.op is Ops.CONST for s in dim.src):
          name, val = dim.src if dim.src[1].op is Ops.CONST else reversed(dim.src)
-          global_size[j] = f"_{name.expr}[0] + {val.val}"
+          global_size[j] = f"_{name.expr}[0] + {val.arg}"

  prg = ""
  if target == "clang":
@@ -0,0 +1,203 @@
+from tinygrad import Tensor, UOp, getenv
+from tinygrad.uop.ops import AxisType, KernelInfo, Ops
+from tinygrad.dtype import AddrSpace, dtypes
+from tinygrad.helpers import DEBUG, GlobalCounters, Context
+import math
+
+BLOCK_M, BLOCK_N = 64, 64
+WARP_SIZE = 32
+WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
+WAVES_M, WAVES_N = 4, 1
+LANES_PER_WAVE_M, LANES_PER_WAVE_N = 2, 16
+WMMA_ACC = WMMA_M // LANES_PER_WAVE_M
+THREADS_PER_BLOCK = WARP_SIZE * WAVES_M * WAVES_N
+LDS_PAD = 4  # pad LDS rows to reduce bank conflicts
+
+WMMA_ARG = (WMMA_M, WMMA_N, WMMA_K), 'AMD', 32
+LOG2E = math.log2(math.e)
+
+def warp_shfl_xor(val, offset, lane):
+  """Read val from lane ^ offset using ds_bpermute."""
+  idx = ((lane ^ offset) * 4).cast(dtypes.int)
+  if val.op is Ops.INDEX and val.addrspace == AddrSpace.REG: val = val.load()
+  return UOp(Ops.CUSTOM, dtypes.float, (idx, val),
+             arg="__builtin_bit_cast(float, __builtin_amdgcn_ds_bpermute({0}, __builtin_bit_cast(int, {1})))")
+
+def warp_reduce_max(val, lane):
+  """Tree reduce MAX across LANES_PER_WAVE_N=16 lanes."""
+  for offset in [8, 4, 2, 1]:
+    val = UOp(Ops.MAX, dtypes.float, (val, warp_shfl_xor(val, offset, lane)))
+  return val
+
+def warp_reduce_sum(val, lane):
+  """Tree reduce SUM across LANES_PER_WAVE_N=16 lanes."""
+  for offset in [8, 4, 2, 1]:
+    val = val + warp_shfl_xor(val, offset, lane)
+  return val
+
+def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
+  # inputs are (B*H, N, D)
+  BH, N, D = q.shape
+  assert N % BLOCK_M == 0 and N % BLOCK_N == 0, f"N={N} must be divisible by BLOCK_M={BLOCK_M} and BLOCK_N={BLOCK_N}"
+  assert D % WMMA_K == 0 and D % LANES_PER_WAVE_N == 0, f"D={D} must be divisible by WMMA_K={WMMA_K} and LANES_PER_WAVE_N={LANES_PER_WAVE_N}"
+  assert BLOCK_M % (WAVES_M * WMMA_M) == 0 and BLOCK_N % LANES_PER_WAVE_N == 0
+  TM = BLOCK_M // (WAVES_M * LANES_PER_WAVE_M)
+  TN = BLOCK_N // (WAVES_N * LANES_PER_WAVE_N)
+  TD = D // (WAVES_N * LANES_PER_WAVE_N)
+  SCALE = 1.0 / math.sqrt(D)
+
+  block_bh = UOp.range(BH, 0, AxisType.GLOBAL)
+  block_m = UOp.range(N // BLOCK_M, 1, AxisType.GLOBAL)
+
+  q = q.reshape(BH, N//BLOCK_M, BLOCK_M, D)[block_bh, block_m]
+  k = k.reshape(BH, N//BLOCK_N, BLOCK_N, D)[block_bh]
+  v = v.reshape(BH, N//BLOCK_N, BLOCK_N, D)[block_bh]
+  o = o.reshape(BH, N//BLOCK_M, BLOCK_M, D)[block_bh, block_m]
+
+  wave_m = UOp.range(WAVES_M, 2, AxisType.LOCAL)
+  wave_n = UOp.range(WAVES_N, 3, AxisType.LOCAL)
+  lane = UOp.range(WARP_SIZE, -1, AxisType.WARP)
+  tid = (wave_m * WAVES_N + wave_n) * WARP_SIZE + lane
+  lane_m = lane // LANES_PER_WAVE_N
+  lane_n = lane % LANES_PER_WAVE_N
+
+  # LDS allocation: slot 0 = Q then P (shared), slot 1 = K then V
+  # TODO: the memory planner should be able to find this reuse
+  ELEMS_PER_THREAD = BLOCK_M * D // THREADS_PER_BLOCK
+  QP_lds = UOp.placeholder((BLOCK_M, D + LDS_PAD), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL)
+  KV_lds = UOp.placeholder((BLOCK_N, D + LDS_PAD), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL)[:, :D]
+
+  # register state
+  acc = UOp.placeholder((TM, TD), dtypes.float, slot=2, addrspace=AddrSpace.REG)
+  m_i = UOp.placeholder((TM,), dtypes.float, slot=3, addrspace=AddrSpace.REG)
+  l_i = UOp.placeholder((TM,), dtypes.float, slot=4, addrspace=AddrSpace.REG)
+  acc = acc.after(acc.store(acc.const_like(0)))
+  m_i = m_i.after(m_i.store(m_i.const_like(-math.inf)))
+  l_i = l_i.after(l_i.store(l_i.const_like(0)))
+
+  # ====== KV tile loop ======
+  n_tile = UOp.range(N // BLOCK_N, 100, AxisType.REDUCE)
+
+  # load Q + K into LDS (Q reloaded each iteration since P overwrites slot 0)
+  Q_lds = QP_lds[:, :D]
+  Q_store = Q_lds.after(n_tile).reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
+    q.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
+  K_store = KV_lds.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
+    k[n_tile].reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
+  # NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
+  Q_lds = Q_lds.after(UOp.group(Q_store, K_store))
+  KV_lds_k = KV_lds.after(UOp.group(Q_store, K_store))
+
+  # -- S = Q @ K^T via WMMA (re-init each n_tile) --
+  S_reg = UOp.placeholder((TM, TN), dtypes.float, slot=6, addrspace=AddrSpace.REG)
+  S_reg = S_reg.after(S_reg.after(n_tile).store(S_reg.const_like(0)))
+  k_qk = UOp.range(D // WMMA_K, 101, AxisType.REDUCE)
+  tm1 = UOp.range(TM // WMMA_ACC, 200)
+  tn1 = UOp.range(TN, 201)
+  S_frag = S_reg.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0, 2, 1)[tm1, tn1]
+  q_frag = Q_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, D // WMMA_K, WMMA_K)[wave_m, tm1, lane_n, k_qk]
+  k_frag = KV_lds_k.reshape(WAVES_N, TN, WMMA_N, D // WMMA_K, WMMA_K)[wave_n, tn1, lane_n, k_qk]
+  qk = UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), *WMMA_ARG)
+  qk_done = S_frag.store(qk).end(tm1, tn1).end(k_qk)
+  S_reg = S_reg.after(qk_done)
+
+  # -- softmax in registers with warp shuffles --
+  S_reg = S_reg.after(S_reg.store(S_reg * SCALE))
+
+  # per-thread local row max over TN=4 elements, then warp reduce across 16 lanes
+  m_ij = UOp.placeholder((TM,), dtypes.float, slot=7, addrspace=AddrSpace.REG)
+  m_ij = m_ij.after(m_ij.after(n_tile).store(m_ij.const_like(-math.inf)))
+  rm2 = UOp.range(TN, 261, AxisType.REDUCE)
+  m_ij = m_ij.after(m_ij.store(m_ij.after(rm2).maximum(S_reg[:, rm2])).end(rm2))
+  # warp reduce max (in-place)
+  ri_w = UOp.range(TM, 270)
+  m_ij = m_ij.after(m_ij[ri_w].store(warp_reduce_max(m_ij[ri_w], lane)).end(ri_w))
+
+  # compute P = exp(S - m_ij) in S_reg
+  S_reg = S_reg.after(S_reg.store(((S_reg - m_ij.reshape(TM, 1).expand(TM, TN)) * LOG2E).exp2()))
+
+  p_local = UOp.placeholder((TM,), dtypes.float, slot=8, addrspace=AddrSpace.REG)
+  p_local = p_local.after(p_local.after(n_tile).store(p_local.const_like(0)))
+  rp2 = UOp.range(TN, 291, AxisType.REDUCE)
+  p_local = p_local.after(p_local.store(p_local.after(rp2) + S_reg[:, rp2]).end(rp2))
+  ri_ws = UOp.range(TM, 295)
+  p_sum = p_local.after(p_local[ri_ws].store(warp_reduce_sum(p_local[ri_ws], lane)).end(ri_ws))
+
+  # write P = exp(S - m_ij) to P_lds (reuses slot 0, Q no longer needed)
+  P_lds = QP_lds[:, :BLOCK_N]
+  P_write = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_ACC, LANES_PER_WAVE_M, WAVES_N, TN, LANES_PER_WAVE_N)
+  P_write = P_write.permute((0, 4, 3, 6, 1, 2, 5)).reshape(THREADS_PER_BLOCK, TM, TN)
+  P_store = P_write[tid].store(S_reg.cast(dtypes.half))
+
+  # -- online softmax correction --
+  ri4 = UOp.range(TM, 330)
+  m_new_val = m_i[ri4].maximum(m_ij[ri4])
+  alpha_val = ((m_i[ri4] - m_new_val) * LOG2E).exp2()
+  beta_val = ((m_ij[ri4] - m_new_val) * LOG2E).exp2()
+  rj4 = UOp.range(TD, 331)
+  correction = UOp.group(
+    acc[ri4, rj4].store(alpha_val * acc[ri4, rj4]).end(rj4),
+    l_i[ri4].store(alpha_val * l_i[ri4] + beta_val * p_sum[ri4]),
+    m_i[ri4].store(m_new_val),
+  ).end(ri4)
+  acc = acc.after(correction)
+  l_i = l_i.after(correction)
+  m_i = m_i.after(correction)
+
+  # load V into KV_lds (must wait for QK WMMA to finish reading K from KV_lds)
+  V_store = KV_lds.after(qk_done).reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
+    v[n_tile].reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
+  # NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
+  P_lds = P_lds.after(UOp.group(P_store, V_store))
+  KV_lds_v = KV_lds.after(UOp.group(P_store, V_store))
+
+  # -- acc += P @ V via WMMA --
+  k_pv = UOp.range(BLOCK_N // WMMA_K, 400, AxisType.REDUCE)
+  tm2 = UOp.range(TM // WMMA_ACC, 401)
+  tn2 = UOp.range(TD, 402)
+  acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TD).permute(0, 2, 1)[tm2, tn2]
+  p_frag = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_N // WMMA_K, WMMA_K)[wave_m, tm2, lane_n, k_pv]
+  v_frag = KV_lds_v.reshape(WAVES_N, TD, WMMA_N, BLOCK_N // WMMA_K, WMMA_K)[wave_n, tn2, lane_n, k_pv]
+  pv = UOp.wmma(p_frag, v_frag, acc_frag.after(k_pv), *WMMA_ARG)
+
+  # end KV tile loop
+  n_tile_end = acc_frag.store(pv).end(tm2, tn2).end(k_pv).end(n_tile)
+  acc = acc.after(n_tile_end)
+  l_i = l_i.after(n_tile_end)
+  m_i = m_i.after(n_tile_end)
+
+  # normalize: acc /= l_i
+  acc = acc.after(acc.store(acc * (1 / l_i).reshape(TM, 1).expand(TM, TD)))
+
+  # store output
+  o = o.reshape(WAVES_M, TM // WMMA_ACC, WMMA_ACC, LANES_PER_WAVE_M, WAVES_N, TD, LANES_PER_WAVE_N)
+  o = o.permute((0, 4, 3, 6, 1, 2, 5)).reshape(THREADS_PER_BLOCK, TM, TD)
+  return o[tid].store(acc).end(wave_m, wave_n, lane).end(block_m, block_bh).sink(arg=KernelInfo(opts_to_apply=()))
+
+if __name__ == "__main__":
+  B, H, N, D = getenv("B", 1), getenv("H", 32), getenv("N", 1024), getenv("D", 64)
+  q = Tensor.rand(B, H, N, D).cast(dtypes.half)
+  k = Tensor.rand(B, H, N, D).cast(dtypes.half)
+  v = Tensor.rand(B, H, N, D).cast(dtypes.half)
+  o = Tensor.empty(B, H, N, D, dtype=dtypes.float)
+  with Context(DEBUG=0): Tensor.realize(q, k, v)
+
+  q_flat, k_flat, v_flat, o_flat = q.reshape(B*H, N, D), k.reshape(B*H, N, D), v.reshape(B*H, N, D), o.reshape(B*H, N, D)
+  NUM_RUNS = getenv("CNT", 5)
+  ets = []
+  with Context(DEBUG=2):
+    for _ in range(NUM_RUNS):
+      GlobalCounters.reset()
+      tst = Tensor.custom_kernel(o_flat, q_flat, k_flat, v_flat, fxn=amd_flash_attention)[0].realize()
+      ets.append(GlobalCounters.time_sum_s)
+  print(f"best time: {min(ets)*1e3:.2f}ms")
+
+  if getenv("VERIFY", 1):
+    with Context(DEBUG=0):
+      ref = q.float().scaled_dot_product_attention(k.float(), v.float()).reshape(B*H, N, D).realize()
+      err = (ref - tst).square().mean().item()
+    print(f"mean squared error {err}")
+    if err > 1e-2:
+      raise RuntimeError("flash attention is wrong!")
+    else:
+      print("flash attention is correct!")
@@ -1,12 +1,10 @@
-import atexit, functools, math, pathlib
+import atexit, functools, pathlib
 from tinygrad import Tensor, Device, dtypes
-from tinygrad.dtype import AddrSpace
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
 from tinygrad.renderer import Estimates
-from tinygrad.helpers import getenv, all_same, DEBUG, ceildiv
+from tinygrad.helpers import getenv, all_same, DEBUG
 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

@@ -109,43 +107,6 @@ def custom_hk_mxfp8_gemm(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, *extra:U
  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
                               UOp(Ops.BINARY, arg=lib)))

-# ** MXFP4 GEMM custom kernel
-
-@functools.cache
-def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp, tile_m:int, tile_n:int) -> UOp:
-  from extra.gemm.gemm_mxfp4 import build_kernel
-  M, half_k = math.prod(A.shape[:-1]), A.shape[-1]
-  N, half_k_b = math.prod(B.shape[:-1]), B.shape[-1]
-  K = half_k * 2
-  assert half_k == half_k_b and math.prod(C.shape[:-1]) == M and C.shape[-1] == N
-  threads = UOp.special(256, "lidx0")
-  groups_x, groups_y = UOp.special(ceildiv(N, tile_n), "gidx0"), UOp.special(ceildiv(M, tile_m), "gidx1")
-  lds = UOp.placeholder((163840,), dtypes.uint8, 0, AddrSpace.LOCAL)
-  sink = UOp.sink(C.base, A.base, B.base, scale_a.base, scale_b.base, *(x.base for x in extra), lds, threads, groups_x, groups_y,
-                  arg=KernelInfo(f"mxfp4_gemm_{M}_{N}_{K}",
-                                 estimates=Estimates(ops=2*M*N*K, mem=(M*half_k+N*half_k)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
-  insts = build_kernel(M, N, K, tile_m, tile_n)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in insts))))
-
-def _mxfp4_gemm_quantized(a_q:Tensor, b_q:Tensor, scale_a:Tensor, scale_b:Tensor) -> Tensor:
-  M, half_k = a_q.shape
-  N, half_k_b = b_q.shape
-  assert half_k == half_k_b
-  is_multi = isinstance(a_q.device, tuple)
-  reduce_out = is_multi and (a_q.uop.axis == 1 or b_q.uop.axis == 1)
-  if not is_multi: out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a_q.device)
-  elif reduce_out: out = Tensor(Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a_q.device).uop.unshard(0), device=a_q.device)
-  elif a_q.uop.axis == 0:
-    out = Tensor(Tensor.invalids(1, M//len(a_q.device), N, dtype=dtypes.bfloat16, device=a_q.device).uop.unshard(1), device=a_q.device)
-  elif b_q.uop.axis == 0:
-    out = Tensor(Tensor.invalids(1, M, N//len(a_q.device), dtype=dtypes.bfloat16, device=a_q.device).uop.unshard(2), device=a_q.device)
-  else: out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a_q.device)
-  tile_m, tile_n = next((tm, tn) for tm, tn in ((256, 256), (192, 256), (128, 512)) if M % tm == N % tn == 0)
-  out = Tensor.custom_kernel(out, a_q, b_q, scale_a, scale_b,
-                             fxn=functools.partial(custom_mxfp4_gemm, tile_m=tile_m, tile_n=tile_n))[0]
-  if reduce_out: out = out.sum(0)
-  return out.squeeze(0)
-
 def quantize_mxfp8(x:Tensor) -> tuple[Tensor, Tensor, Tensor]:
  # 1x32 block scaling along the last axis
  *batch, K = x.shape
@@ -213,10 +174,10 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
  m = UOp.range(M, 1)
  n = UOp.range(N, 2)
  k = UOp.range(K, 0, AxisType.REDUCE)
-  mul = (A.flatten().index((m*UOp.const(K)+k))*
-         B.flatten().index((k*UOp.const(N)+n))).cast(dtypes.float32)
+  mul = (A.flatten().index((m*UOp.const(dtypes.weakint, K)+k))*
+         B.flatten().index((k*UOp.const(dtypes.weakint, N)+n))).cast(dtypes.float32)
  red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype)
-  store = C.flatten().index((m*UOp.const(N)+n)).store(red).end(m, n)
+  store = C.flatten().index((m*UOp.const(dtypes.weakint, N)+n)).store(red).end(m, n)
  return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))

 # ** bf16 A @ B.T kernel in C
@@ -283,6 +244,7 @@ def hk_bf16_atb_gemm(a:Tensor, b:Tensor) -> Tensor:
  if reduce_out: out = out.sum(0)
  return out.squeeze(0) if out.ndim == 3 else out

+
 # ** backward gemm, might use the asm gemm

 def custom_gemm_bw(gradient:UOp, kernel:UOp, n_scales:int=2, has_grad_amax:bool=False, has_w_post:bool=False):
@@ -379,30 +341,13 @@ def custom_mx_gemm_bw(gradient:UOp, kernel:UOp, has_w_post:bool, w_stored:bool=F
  if wp is not None: grad_b = grad_b / wp.reshape(-1, 1)
  return (None, grad_a.uop, grad_b.uop) + tuple(None for _ in inputs[3:])

-# ** mxfp4 gemm backward
-
-def custom_mxfp4_gemm_bw(gradient:UOp, kernel:UOp):
-  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)
-  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

 def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=None, grad_amax_state:Tensor|None=None,
             next_grad_amax_state:Tensor|None=None,
             w_post_scale:Tensor|None=None, mx:bool=False, mx_scales:tuple|None=None, mx_w_stored:bool=False, g_amax:Tensor|None=None,
-             a_pretranspose:Tensor|None=None, mxfp4:bool=False) -> Tensor:
+             a_pretranspose:Tensor|None=None) -> Tensor:
  assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
-  if mxfp4:
-    assert not mx and mx_scales is None, "mxfp4 owns quantization; mx/mx_scales are for mxfp8"
-    assert a.dtype == dtypes.bfloat16, f"cannot quantize {a.dtype} to mxfp4"
  counters["used"] += 1
  unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0
  if unfold_batch:
@@ -410,7 +355,7 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N
    a = a.reshape(a.shape[0]*a.shape[1], a.shape[2])
  squeeze = a.ndim == 2
  if squeeze: a = a.unsqueeze(0)
-  out_dtype = dtypes.bfloat16 if a.dtype == FP8_DTYPE or mxfp4 else a.dtype
+  out_dtype = dtypes.bfloat16 if a.dtype == FP8_DTYPE else a.dtype

  batch, M, K = a.shape
  N = b.shape[1]
@@ -433,15 +378,7 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N
  renderer = Device[dname:=(a.device[0] if is_multi else a.device)].renderer
  dname, arch = dname.split(":")[0], renderer.target.arch
  if arch.startswith("gfx950") and getenv("USE_ASM", 1):
-    if mxfp4:
-      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
-      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:
+    if mx:
      # mxfp8 1x32 block scaling
      if mx_scales is not None:
        a_si, a_e8, b_si, b_e8 = mx_scales
@@ -0,0 +1,395 @@
+"""
+HipKittens hk_bf16_gemm (extra/thunder/amd/gemm_bf16.cpp) reimplemented with tinygrad UOps.
+
+C[M, N] (bf16) = A[M, K] @ B[N, K]^T, fp32 accumulation, exactly the kittens kernel shape:
+  - 256x256 output tile per workgroup, K_STEP=64
+  - 8 warps in a 2x4 grid, each warp owns a 128x64 warp-tile
+  - v_mfma_f32_16x16x32_bf16 on CDNA4 (gfx950) / v_wmma_f32_16x16x16_bf16 (wave32, gfx12) on RDNA4,
+    fp32 accumulators
+  - shared tiles As/Bs with the kittens st_16x32_s swizzle (16x32 subtiles of 1024B)
+  - K stages (STAGES=1: synchronous single buffer)
+
+Validated on gfx1201 hardware (exact for identity-B, rounding-level noise otherwise), rendered
+and compiled to gfx950 with comgr for assembly comparison against gemm_bf16.cpp.
+
+What is NOT expressible vs the kittens C++:
+  - explicit s_waitcnt vmcnt()/lgkmcnt() pipelining and s_setprio: tinygrad models async copy
+    overlap with slot dependencies and emits full workgroup barriers; instruction scheduling
+    is left to clang/LLVM
+  - direct-to-LDS global loads (buffer_load_lds): tinygrad goes global->reg->LDS
+
+Pipelining status: STAGES=2 gives the kittens-shaped double-buffered pipeline (2 x 64KB LDS
+like gemm_bf16.cpp, copies overlap the previous pair's mma's), written with FA/gemm_fragment
+conventions: LDS buffers are (2, tile) placeholders indexed by symbolic parity (ko % 2),
+which sidesteps static slot choice, fill iterations, predication and duplicate static stores.
+Validated on the CDNA4 emulator for all tile counts (amt = K//64 in {1..32}, odd/even),
+single- AND multi-workgroup (bit-close to stages=1 / to hippkittens at rounding level).
+
+Bug hunt notes (all fixed on this branch; they were entangled for a long time):
+  1. The double-buffered pipeline REGISTER-SPILLS (255+ VGPRs vs 166 for stages=1), and the
+     mock emulator aliased the spill (scratch) segment of ALL waves of a workgroup onto one
+     64-lane region. On real HW each wavefront owns a per-lane segment of the scratch ring
+     (indexed by (wave_id, lane)); waves trampled each other's spilled accumulators, giving
+     the "only the last wave's output survives" signature. emu.py now allocates per-wave
+     scratch buffers.
+  2. The remaining "shape-dependent" corruption (NaNs, mispositioned values in contiguous
+     copies feeding the GEMM) came from tinygrad's devectorizer fusing adjacent bf16 stores
+     into 32-bit stores with UNALIGNED (2-byte) granularity: legal on AMD FLAT/GLOBAL (the
+     hardware splits them), but the emulator floored misaligned addresses to the word below.
+     _mem_store now handles unaligned 32-bit (and wider) accesses byte-exactly.
+  3. memory_coalescing (late/coalesce.py) assumed a single static store per (buffer, index)
+     ("attempting multiple stores"); aliased stores (a double-buffered LDS slot written in a
+     prologue AND a loop body) are now simply kept scalar instead of asserting/merging.
+  4. pm_split_ranges may only split ranges WITHOUT hardware meaning (WEAK/REDUCE/LOOP);
+     splitting LOCAL/WARP/THREAD/GLOBAL/GROUP_REDUCE/UPCAST ranges scrambles the
+     logical<->hardware mapping of hand-written kernels such as this one.
+
+  RDNA4 (gfx12) uses 8-element accumulator fragments, so the 8x4 tile grid needs
+  256 fp32 acc registers per thread -> guaranteed spills (0.85 TF vs 96 TF default on
+  gfx1201). The kernel is right-sized for CDNA4 (fragsz 4 -> 128 acc regs).
+
+Lane layouts (RDNA4 verified with probing on gfx1201 hardware; CDNA from the mfma docs):
+  CDNA (64 thr/warp, 16x16x32):  A/B frag: tile-row = l%16, k = (l//16)*8+i (i in 0..7)
+  RDNA4 (32 thr/warp, 16x16x16): A/B frag: tile-row = l%16, k = (l//16)*4+(i%4)+8*(i//4) (i in 0..7)
+  both:                        acc frag: CDNA m=(l//16)*4+i (i<4) / RDNA4 m=(l//16)*8+i (i<8), n=l%16
+
+The RDNA4 fragment k-set {k0..3, k0+8..11} is not contiguous, so on RDNA4 the LDS column layout
+is block-permuted (4-element blocks within each 16-col group are stored as [0,2,1,3]) making
+every fragment 8 contiguous halves (one 16B chunk) on both archs; the copy path applies the
+same permutation.
+
+NOTE: thread ids come from UOp.special (like mi350x_uop_matmul.py), not an AxisType.LOCAL
+RANGE. (pm_split_ranges now only splits WEAK/REDUCE/LOOP ranges, so LOCAL ranges would
+survive too, but UOp.special is the sanctioned way to tag hardware lane ids.)
+NOTE 2: WMMA operand/accumulator fragments must carry the fragment length in their UOp shape.
+NOTE 3: swizzled addresses are written in provably-contiguous "base + vector-offset" form,
+otherwise the devectorizer emits scalar ds_read_u16/ds_write_b16.
+"""
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
+from tinygrad.dtype import AddrSpace
+from tinygrad.renderer import Estimates
+from tinygrad.helpers import getenv, cdiv
+
+# ---- tile shape (identical to gemm_bf16.cpp; HK_TILE=128 overrides for small-LDS devices) ----
+BLOCK_M = BLOCK_N = getenv("HK_TILE", 256)
+K_STEP = 64
+WARPS_M, WARPS_N = 2, 4
+NUM_WARPS = WARPS_M * WARPS_N                                # 8
+WARP_TILE_M, WARP_TILE_N = BLOCK_M // WARPS_M, BLOCK_N // WARPS_N   # 128 x 64 (64 x 32 at HK_TILE=128)
+
+def arch_params(arch:str):
+  is_cdna = arch.startswith("gfx9")
+  if is_cdna:   # CDNA mfma 16x16x32 bf16: in-frag 8, acc 4 (16,16,16 on the acc side)
+    return dict(warp_threads=64, dims=(16,16,32), frag_in=8, frag_out=4,
+                acc_m=lambda l, i: (l//16)*4 + i, kperm=None, copy_vec=8)
+  # RDNA4 wmma 16x16x16 (wave32, gfx12 layout): in-frag 8 (permuted in LDS), acc 8
+  return dict(warp_threads=32, dims=(16,16,16), frag_in=8, frag_out=8,
+              acc_m=lambda l, i: (l//16)*8 + i, kperm=(0,2,1,3), copy_vec=4)
+
+# ---- kittens st_16x32_s swizzle (byte offset: off ^ (((off % 1024) >> 9) << 5)) ----
+# halves-index form: within a 1024B (16x32) subtile, halves-index bit4 ^= row bit3,
+# written in "base + vector-offset" form so the devectorizer can prove contiguity.
+def st_half_base(r, c, tile_cols:int):
+  """swizzled halves-index pre-vector-offset; c is the logical (permuted) column, 4-aligned."""
+  subtile_id = (r//16) * (tile_cols//32) + (c//32)
+  r16 = r % 16
+  flip = (r16 >> 3) & 1
+  return subtile_id*512 + r16*32 + (((c % 32) >> 2) ^ (flip << 2)) * 4
+
+def hk_bf16_gemm_kernel(C:UOp, A:UOp, B:UOp, *, arch:str, stages:int=1) -> UOp:
+  """C = A @ B^T ; A is (M,K), B is (N,K), C is (M,N). HipKittens tile shape."""
+  M, K = A.shape
+  N, K2 = B.shape
+  assert K == K2 and A.dtype == B.dtype == dtypes.bfloat16 and C.dtype == dtypes.bfloat16
+  assert not (M % BLOCK_M or N % BLOCK_N or K % K_STEP), f"dims must be multiples of {(BLOCK_M, BLOCK_N, K_STEP)}"
+
+  ap = arch_params(arch)
+  warp_threads, dims = ap["warp_threads"], ap["dims"]
+  FRAG_IN, FRAG_OUT, kperm, acc_m, CPV = ap["frag_in"], ap["frag_out"], ap["kperm"], ap["acc_m"], ap["copy_vec"]
+  NUM_THREADS = NUM_WARPS * warp_threads
+  TC_M, TC_N, TC_K = dims
+  MT, NT = WARP_TILE_M // TC_M, WARP_TILE_N // TC_N               # 8, 4 tiles per warp
+  # permute 4-half blocks within each 16-col group (RDNA4: [0,2,1,3] = swap middle blocks)
+  def perm_col(c):
+    if kperm is None: return c
+    return (c & ~15) | ((((c>>2) & 1) << 1 | ((c>>3) & 1)) << 2) | (c & 3)
+
+  bx, by = UOp.special(N//BLOCK_N, "gidx0"), UOp.special(M//BLOCK_M, "gidx1")
+  lane = UOp.special(warp_threads, "lidx0")
+  warp = UOp.special(NUM_WARPS, "lidx1")
+  warp_row, warp_col = warp // WARPS_N, warp % WARPS_N
+  tid = warp*warp_threads + lane
+
+  def smem(slot) -> UOp: return UOp.placeholder((BLOCK_M*K_STEP,), dtypes.bfloat16, slot, AddrSpace.LOCAL)
+  As = [smem(2*i) for i in range(stages)]
+  Bs = [smem(2*i+1) for i in range(stages)]
+
+  # per-warp accumulator: (MT x NT) 16x16 tiles of FRAG_OUT fp32 per thread
+  acc = UOp.placeholder((MT, NT, FRAG_OUT), dtypes.float32, 12, AddrSpace.REG)
+  acc = acc.after(acc.store(acc.const_like(0.0)))   # FA-style init: self-store, keeps the value flow loop-carried
+
+  # global -> LDS copy: CPV halves per op (16B on CDNA, 8B on RDNA4), thread-major coalescing
+  OPS_PER_TILE = BLOCK_M*K_STEP//CPV
+  OPR = K_STEP//CPV
+  def copy_tile(dst:UOp, src:UOp, base_row:UOp, base_col:UOp, slot:int) -> UOp:
+    ir = UOp.range(cdiv(OPS_PER_TILE, NUM_THREADS), slot, AxisType.LOOP)
+    j = UOp.range(CPV, slot+1, AxisType.UPCAST)
+    chunk = ir*NUM_THREADS + tid
+    r, cb = chunk // OPR, chunk % OPR          # row, 4/8-col block
+    return dst[st_half_base(r, perm_col(cb*CPV), K_STEP) + j].store(src[base_row + r, base_col + cb*CPV + j]).end(ir, j)
+
+  def load_stage(sidx:int, ko, slot:int, barrier:bool) -> tuple[UOp, UOp]:
+    A_r = copy_tile(As[sidx], A, by*BLOCK_M, ko*K_STEP, slot)
+    B_r = copy_tile(Bs[sidx], B, bx*BLOCK_N, ko*K_STEP, slot+10)
+    bar = UOp.barrier(A_r, B_r) if barrier else UOp.group(A_r, B_r)
+    return As[sidx].after(bar), Bs[sidx].after(bar)
+
+  # ---- pipelined path (stages=2) ----
+  NIT = cdiv(OPS_PER_TILE, NUM_THREADS)      # copy ops per thread per tile
+  def setprio(n:int, slot:int) -> UOp:
+    """__builtin_amdgcn_s_setprio(n), like gemm_bf16.cpp: raise warp priority for the mma phase
+    so global/LDS traffic of the other waves doesn't starve issue slots."""
+    # distinct src slot per call site so identical-priority instructions at different k-tiles
+    # don't get UOp-hash-deduped into one placement (s_setprio is position-sensitive)
+    return UOp(Ops.CUSTOMI, dtypes.void, src=(UOp.const(dtypes.weakint, slot), UOp.const(dtypes.weakint, n)),
+               arg="__builtin_amdgcn_s_setprio({1}); // {0}")
+
+  def gload_write_tile(dst:UOp, src:UOp, base_row:UOp, kt, slot:int) -> UOp:
+    """store one global tile into an LDS slot (loads and stores share the vec range j)."""
+    j = UOp.range(CPV, slot, AxisType.UPCAST)
+    def one(ir:int) -> UOp:
+      chunk = ir*NUM_THREADS + tid
+      r, cb = chunk // OPR, chunk % OPR
+      return dst[st_half_base(r, perm_col(cb*CPV), K_STEP) + j].store(src[base_row + r, kt*K_STEP + cb*CPV + j])
+    return UOp.group(*[one(ir) for ir in range(NIT)]).end(j)
+
+  def compute(acc:UOp, A_l:UOp, B_l:UOp, afters:tuple[UOp, ...], pred:UOp|None=None, aoff:UOp=None, boff:UOp=None) -> UOp:
+    """One K_STEP=64 iteration: (K_STEP//TC_K) k-chunks unrolled, (MT x NT) mma each, like the kittens main loop.
+
+    pred (optional): a loop-range condition; accumulator stores are predicated on it so the
+    first (fill) iteration of a software pipeline can run the body with garbage LDS contents
+    without contaminating the accumulator."""
+    arow = warp_row*WARP_TILE_M + lane % 16    # fragment tile row in the LDS tile (m)
+    brow = warp_col*WARP_TILE_N + lane % 16    # (n)
+    ja = UOp.range(FRAG_IN, 701, AxisType.UPCAST)
+    jb = UOp.range(FRAG_IN, 702, AxisType.UPCAST)
+    acc_k = acc.after(*afters) if afters else acc
+    last_store = None
+    # in the permuted layout every fragment is 8 contiguous halves starting at an 8-aligned col
+    for kk in range(K_STEP//TC_K):
+      cc = kk*(TC_K//FRAG_IN) + (lane // 16)   # fragment chunk col (8 halves)
+      oa, ob = (aoff, boff) if aoff is not None else (None, None)
+      a_frags = [A_l[st_half_base(arow + mt*16, cc*8, K_STEP) + ja].contract(ja) if oa is None else
+                 A_l[oa + st_half_base(arow + mt*16, cc*8, K_STEP) + ja].contract(ja) for mt in range(MT)]
+      b_frags = [B_l[st_half_base(brow + nt*16, cc*8, K_STEP) + jb].contract(jb) if ob is None else
+                 B_l[ob + st_half_base(brow + nt*16, cc*8, K_STEP) + jb].contract(jb) for nt in range(NT)]
+      for mt in range(MT):
+        for nt in range(NT):
+          cur = acc_k[mt, nt]
+          out = UOp.wmma(a_frags[mt], b_frags[nt], cur, dims, 'AMD', warp_threads)
+          if pred is not None: out = pred.where(cur, out)
+          last_store = acc_k[mt, nt].store(out)
+          acc_k = acc_k.after(last_store)
+    return last_store
+
+  # ---- K loop ----
+  amt = cdiv(K, K_STEP)
+  _stages = stages
+  if _stages == 1:
+    ko = UOp.range(amt, 600, AxisType.LOOP)
+    A_l, B_l = load_stage(0, ko, 100, barrier=True)
+    last = compute(acc, A_l, B_l, afters=(ko,))
+    acc = acc.after(last.barrier().end(ko))
+  else:
+    # Double-buffered pipeline on FA/gemm_fragment conventions: each LDS buffer is a
+    # (2, tile) placeholder indexed by symbolic parity (ko % 2) -- no static slot choice,
+    # no duplicate static stores (memory_coalescing-safe), no fill iteration, no predication.
+    def smem2(slot) -> UOp: return UOp.placeholder((2*BLOCK_M*K_STEP,), dtypes.bfloat16, slot, AddrSpace.LOCAL)
+    A_l, B_l = smem2(0), smem2(1)
+
+    TILE_ELEMS = BLOCK_M * K_STEP
+    def copy_stage(dst:UOp, slot_off:UOp, src:UOp, base_row:UOp, kt, slot:int) -> UOp:
+      """store one global tile into dst + slot_off (flat element offset -- slot_off = parity*TILE_ELEMS).
+
+      Each thread's 8-element chunk is a single buffer_load_lds direct-to-LDS instruction
+      (the kittens '... offen lds' fill path), emitted via Ops.CUSTOMI so it bypasses the
+      devectorizer (a SHRINK store of a SHRINK load gets expanded to scalars before render)."""
+      ir = UOp.range(cdiv(OPS_PER_TILE, NUM_THREADS), slot+1, AxisType.LOOP)
+      chunk = ir*NUM_THREADS + tid
+      r, cc = chunk // OPR, chunk % OPR
+      if getenv("HK_G2L", 0) == 3:
+        # direct-to-LDS fill (kittens '... offen lds' path): the hardware writes each lane's
+        # chunk to the lane-linear LDS address (M0 + lane*size), so the swizzle is moved to
+        # the GLOBAL side: lane q's 16B chunk fetches the matrix element that st_half_base
+        # maps to the tile-linear position q. Verified bijective; the fragment-read layout
+        # (and therefore the read swizzle) is unchanged.
+        chunk = ir*NUM_THREADS + tid
+        p_ = chunk * CPV                             # tile-linear halves position of this lane's chunk
+        sub = p_ >> 9                                # 16x32 subtile id (512 halves)
+        r16 = (p_ & 511) >> 5
+        flip = (r16 >> 3) & 1
+        cb = ((p_ & 31) >> 2) ^ (flip << 2)
+        r_ = (sub >> 1) * 16 + r16
+        c_ = cb*4 + (sub & 1) * 32                   # global column (8-aligned)
+        off_g = (base_row + r_) * K + kt*K_STEP + c_
+        lds_el = slot_off + ir*NUM_THREADS*CPV       # elements; &buf[el*8] = chunk base byte addr
+        # feed the raw PARAM (unwrapping the scheduler's RESHAPE view, which would otherwise
+        # live unfused into the program and fail spec: 'movement ops not allowed in programs').
+        prm = src
+        while prm.op is not Ops.PARAM and len(prm.src): prm = prm.src[0]
+        nbytes = prm.max_numel() * prm.dtype.itemsize
+        gname = f"data{prm.arg.slot}_{prm.max_numel()}"
+        return UOp(Ops.CUSTOMI, dtypes.void, src=(prm, dst, lds_el, off_g),
+                   arg=(f"llvm_amdgcn_raw_buffer_load_lds(make_srsrc_((void*){gname}, {nbytes}), "
+                        f"(as3_uint32_ptr)(&({{1}}[({{2}})])), {CPV*2}, ((unsigned)({{3}}))*2U, 0, 0, 0);")).end(ir)
+      # default: elementwise global->LDS stores
+      off_l = slot_off + st_half_base(r, perm_col(cc*CPV), K_STEP)
+      off_g = (base_row + r) * K + kt*K_STEP + cc*CPV
+      j = UOp.range(CPV, slot, AxisType.UPCAST)
+      return dst[off_l + j].store(src[base_row + r, kt*K_STEP + cc*CPV + j]).end(ir, j)
+
+    ZERO = UOp.const(dtypes.weakint, 0)
+    # prologue: tile 0 into slot 0 of both buffers, barrier before first read
+    g0 = UOp.group(copy_stage(A_l, ZERO, A, by*BLOCK_M, ZERO, 100),
+                   copy_stage(B_l, ZERO, B, bx*BLOCK_N, ZERO, 110))
+    bar0 = UOp.barrier(g0)
+    # Double-buffered pipeline: slot ko%2 holds k-tile ko; the prefetch copy of tile ko+1
+    # (into the other slot) rides IN FRONT of the wmma's and overlaps them; one barrier per
+    # k-tile hand-off covers write(ko)->read(ko+1) [and read(ko)->write(ko+1) is closed by
+    # the ko-1 barrier already]. The parities/offsets are static python constants when
+    # HK_UNROLL (default on): straight-line like the kittens main loop; the rolled variant
+    # uses pm_split_ranges to split the ko LOOP range at the (ko % 2) boundary.
+    if getenv("HK_UNROLL", 1) and amt % (UN := getenv("HK_UNROLL_U", 8)) == 0:
+      # outer rolled loop of amt//U iterations, U python-unrolled k-tiles inside: nearly the
+      # kittens straight-line node shape (one barrier per k-tile) at a fraction of the
+      # full-unroll uop count (full unroll of amt=64 needs ~12 min of schedule time; U=8
+      # keeps every tile's prefetch + compute + hand-off barrier but stays seconds).
+      ko_o = UOp.range(amt // UN, 600, AxisType.LOOP)
+      pa, pb = A_l.after(bar0, ko_o), B_l.after(bar0, ko_o)
+      for i in range(UN):
+        kt = ko_o * UN + i
+        pr, pn = (i % 2) * TILE_ELEMS, ((i + 1) % 2) * TILE_ELEMS
+        kt_next = UOp.minimum(kt + 1, amt - 1)
+        ga0 = UOp.group(copy_stage(pa, UOp.const(dtypes.weakint, pn), A, by*BLOCK_M, kt_next, 300 + 4*i),
+                        copy_stage(pb, UOp.const(dtypes.weakint, pn), B, bx*BLOCK_N, kt_next, 302 + 4*i))
+        sp_hi = setprio(1, 300 + 4*i)                      # kittens: raised prio for the mma phase
+        last = compute(acc, pa, pb, afters=(ko_o, sp_hi), aoff=UOp.const(dtypes.weakint, pr), boff=UOp.const(dtypes.weakint, pr))
+        sp_lo = setprio(0, 301 + 4*i)
+        handoff = UOp.group(last, sp_lo, ga0).barrier()
+        acc = acc.after(handoff)
+        pa, pb = A_l.after(handoff, ga0), B_l.after(handoff, ga0)
+      acc = acc.after(UOp.group(handoff).end(ko_o))
+    else:
+      ko = UOp.range(amt, 600, AxisType.LOOP)
+      pr, pn = ko % 2, (ko+1) % 2                    # slot of the tile being computed / being prefetched
+      kt_next = UOp.minimum(ko+1, amt-1)             # clamped tail prefetch (its data is unused)
+      pa, pb = A_l.after(bar0, ko), B_l.after(bar0, ko)
+      ga = UOp.group(copy_stage(pa, pn*TILE_ELEMS, A, by*BLOCK_M, kt_next, 130),
+                     copy_stage(pb, pn*TILE_ELEMS, B, bx*BLOCK_N, kt_next, 140))
+      sp_hi = setprio(1, 150)
+      last = compute(acc, pa, pb, afters=(ko, sp_hi), aoff=pr*TILE_ELEMS, boff=pr*TILE_ELEMS)
+      acc = acc.after(UOp.group(last, setprio(0, 151), ga).barrier().end(ko))
+
+  # ---- epilogue: per-thread fragment stores, cast to bf16 (scalar per fragment element) ----
+  mt, nt = UOp.range(MT, 801, AxisType.LOOP), UOp.range(NT, 802, AxisType.LOOP)
+  def store_i(i:int) -> UOp:
+    crow = by*BLOCK_M + warp_row*WARP_TILE_M + mt*16 + acc_m(lane, i)
+    ccol = bx*BLOCK_N + warp_col*WARP_TILE_N + nt*16 + lane % 16
+    return C[crow, ccol].store(acc[mt, nt, i].cast(dtypes.bfloat16))
+  out_st = UOp.group(*[store_i(i) for i in range(FRAG_OUT)])
+  return out_st.end(mt, nt).sink(arg=KernelInfo(name="hk_bf16_gemm",
+    estimates=Estimates(ops=2*M*N*K, mem=(M*K+N*K+M*N)*2)))
+
+def hk_bf16_gemm_tiny(a:Tensor, b:Tensor, stages:int=1) -> Tensor:
+  """C = a @ b.T for bf16 a (M,K), b (N,K) with the HipKittens-shaped tinygrad kernel."""
+  arch = Device[a.device].renderer.target.arch
+  c = Tensor.empty(a.shape[0], b.shape[0], dtype=dtypes.bfloat16, device=a.device)
+  return c.custom_kernel(a, b, fxn=lambda C, A, B: hk_bf16_gemm_kernel(C, A, B, arch=arch, stages=stages))[0]
+
+if __name__ == "__main__":
+  import numpy as np
+  from tinygrad import Device
+  M = N = K = 512
+  # exact test: B = identity -> C must equal A bit-exactly
+  a = Tensor.randn(M, K, dtype=dtypes.bfloat16).contiguous()
+  bid = Tensor(np.eye(K, N, dtype=np.float32), dtype=dtypes.bfloat16).contiguous()
+  cid = hk_bf16_gemm_tiny(a, bid, stages=getenv("STAGES", 1)).realize()
+  assert np.array_equal(cid.float().numpy(), a.float().numpy()), "identity test failed"
+  # real test: bf16 gemm vs fp32 reference, rounding-level noise
+  b = Tensor.randn(N, K, dtype=dtypes.bfloat16).contiguous()
+  c = hk_bf16_gemm_tiny(a, b, stages=getenv("STAGES", 1)).realize()
+  ref = (a @ b.T).float().realize()
+  err = (c.float() - ref).abs().max().item()
+  print(f"identity exact, random max err: {err:.5f}")
+
+  # ---- benchmark mode: kittens hk_bf16_gemm vs tinygrad stages={1,2} vs the default scheduled gemm ----
+  # run on real hardware with:  DEV=AMD:HIP:gfx950 DEBUG=2 HK_BENCH=1 python extra/gemm/hk_gemm_frag.py
+  # sizes via HK_SIZES="2048x2048x2048,4096x4096x4096" (default 2048 cubed), iteration count via ITERS=20.
+  # timings come from GlobalCounters.time_sum_s (sum of kernel times; same source as the DEBUG=2 'tm' column).
+  if getenv("HK_BENCH"):
+    from tinygrad.helpers import GlobalCounters
+    from extra.gemm.cdna_asm_gemm import asm_gemm
+    from tinygrad import Device
+    dev, iters, warm = Device.DEFAULT, getenv("ITERS", 20), 3
+    arch = Device[dev].renderer.target.arch
+    assert arch.startswith("gfx9"), "CDNA only"
+    def bench(label:str, fn, M:int, N:int, K:int) -> float:
+      try:
+        for _ in range(warm): fn()
+        Device[dev].synchronize()
+        GlobalCounters.reset()
+        import time
+        t0 = time.perf_counter()
+        for _ in range(iters): fn()
+        Device[dev].synchronize()
+        wall = time.perf_counter() - t0
+      except Exception as e:
+        print(f"    {label:32s}  unsupported/failed: {type(e).__name__}: {e}")
+        return float('nan')
+      # prefer kernel-side time (GlobalCounters matches the DEBUG=2 'tm' column); fall back to wall clock
+      ms = (GlobalCounters.time_sum_s if GlobalCounters.time_sum_s > 0 else wall) * 1e3 / iters
+      tf = 2*M*N*K / (ms * 1e-3) / 1e12
+      print(f"    {label:32s} {ms:9.3f} ms  {tf:8.1f} TFLOPS")
+      return tf
+    for (M, N, K) in [tuple(map(int, s.split("x"))) for s in getenv("HK_SIZES", "2048x2048x2048").split(",")]:
+      print(f"  size ({M},{N},{K}), grid {M//BLOCK_M}x{N//BLOCK_N} WGs, amt={K//K_STEP} k-tiles/WG")
+      np.random.seed(0)
+      An, Bn = np.random.randn(M, K), np.random.randn(K, N)
+      A = Tensor(An, dtype=dtypes.bfloat16).contiguous().realize()          # (M,K)
+      Bk = Tensor(Bn, dtype=dtypes.bfloat16).contiguous().realize()          # (K,N) for kittens
+      Bt = Bk.T.contiguous().realize()                                       # (N,K) for ours
+      tf_kc = bench("kittens hk_bf16_gemm (asm_gemm)", lambda: asm_gemm(A, Bk).realize(), M, N, K)
+      tf_s1 = bench("tiny stages=1", lambda: hk_bf16_gemm_tiny(A, Bt, stages=1).realize(), M, N, K)
+      tf_s2 = bench("tiny stages=2", lambda: hk_bf16_gemm_tiny(A, Bt, stages=2).realize(), M, N, K)
+      tf_df = bench("tinygrad default (a @ Bt.T)", lambda: (A @ Bt.T).realize(), M, N, K)
+      err = (hk_bf16_gemm_tiny(A, Bt, stages=2).float() - asm_gemm(A, Bk).float()).abs().max().item()
+      print(f"    correctness tiny-s2 vs kittens max diff: {err:.5f}")
+      for nm, tf in [("s1", tf_s1), ("s2", tf_s2), ("default", tf_df)]:
+        if tf == tf and tf_kc == tf_kc: print(f"    tiny {nm:8s}/kittens: {tf/tf_kc:6.2%}")
+  # match the real HipKittens hk_bf16_gemm on the (mock) CDNA4 emulator at small sizes.
+  # run from the repo root with:  DEV=MOCK+AMD:HIP:gfx950 HK_COMPARE=1 python extra/gemm/hk_gemm_frag.py
+  if getenv("HK_COMPARE"):
+    from extra.gemm.cdna_asm_gemm import asm_gemm
+    assert Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950"), "needs CDNA4 (mock emulator or hardware)"
+    def compare(M:int, N:int, K:int, seed:int=0, identity:bool=False):
+      np.random.seed(seed)
+      An = np.random.randn(M, K)
+      Bn = np.eye(K, N) if identity else np.random.randn(K, N)      # (K,N) as expected by asm_gemm
+      A = Tensor(An, dtype=dtypes.bfloat16).contiguous()
+      B = Tensor(Bn, dtype=dtypes.bfloat16).contiguous()             # (K,N) for asm_gemm
+      c_hkc = asm_gemm(A, B).realize().float().numpy()               # real HipKittens hk_bf16_gemm
+      c_hkt = hk_bf16_gemm_tiny(A, B.T.contiguous(), stages=getenv("STAGES", 2)).realize().float().numpy()
+      ref64 = A.float().numpy().astype(np.float64) @ B.float().numpy()
+      tag = "ident" if identity else "rand "
+      print(f"({M},{N},{K}) {tag}: tiny-vs-kittens {np.abs(c_hkt-c_hkc).max():9.6f}  "
+            f"tiny-vs-fp64 {np.abs(c_hkt-ref64).max():9.6f}  kittens-vs-fp64 {np.abs(c_hkc-ref64).max():9.6f}")
+      assert np.abs(c_hkt - ref64).max() < 0.26, "tiny kernel must match fp64 at rounding level"
+      assert np.abs(c_hkt - c_hkc).max() < 0.51, "tiny kernel must match hipkittens"
+    # NOTE: hk_bf16_gemm requires K % 128 == 0 (its prologue+epilogue unconditionally touch
+    # k-tiles num_tiles-1 and num_tiles-2); at other K it reads wrong-but-in-bounds global
+    # memory on the emulator and on real hardware, so only K%128==0 sizes are checked here.
+    compare(256, 256, 128, seed=1)   # single workgroup
+    compare(256, 256, 256, seed=2)
+    compare(512, 512, 128, seed=3)   # multi workgroup
+    compare(256, 256, 128, identity=True)   # bit-exact check
@@ -79,7 +79,7 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
  # this is the big accumulator
  acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float, 0, AddrSpace.REG)
  assert acc.size*WARP_SIZE*WARPGROUP_SIZE*4 == BLOCK_M*BLOCK_N
-  acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const((0.0,)*4, dtypes.float), end=init_l)
+  acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const(dtypes.float, (0.0,)*4), end=init_l)

  # create locals (note A is permuted, and the stride is changed to avoid bank conflicts)
  def make_locals(slot) -> tuple[UOp, UOp]:
@@ -5,8 +5,7 @@ BLOCK_ROW = 256

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

 def _atomic_add(device:str) -> str:
@@ -35,7 +34,7 @@ def _ggather_fwd_kernel(out:UOp, table:UOp, idx:UOp) -> UOp:

 def _ggather_zero_kernel(out:UOp) -> UOp:
  i = UOp.range(out.numel(), 0)
-  return out.flatten().index(i).store(UOp.const(0.0, out.dtype)).end(i).sink(arg=KernelInfo(name="ggather_zero"))
+  return out.flatten().index(i).store(UOp.const(out.dtype, 0.0)).end(i).sink(arg=KernelInfo(name="ggather_zero"))

 def _sharded_zeros(shape:tuple[int, ...], dtype, device) -> Tensor:
  return Tensor.custom_kernel(_sharded_invalids(shape, dtype, device), fxn=_ggather_zero_kernel)[0]
@@ -79,7 +79,7 @@ if __name__ == "__main__":
  linear, var_vals = C.linear_with_vars()
  last_call = linear.src[-1]
  ast = last_call.src[0]
-  bufs = [s.buffer for s in last_call.src[1:] if not s.is_bound_var]
+  bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]

  src = compiled.asm["ptx"]
  # specify the shared memory here so we don't need to do it dynamically
@@ -1,25 +1,55 @@
 """
 tilelang-style matmul_relu written with tinygrad UOp APIs.

+Demonstrates that tilelang's T.alloc_fragment is expressible with existing
+tinygrad primitives: a per-thread REG buffer, wrapped in one Ops.UNSHARD per
+sharded axis over the LOCAL thread-grid ranges to form the full logical tile.
+Here the 64 threads are an 8x8 grid and each thread owns an 8x8 sub-tile --
+the 2-D fragment layout tilelang infers. The kernel is written against the
+full-tile UNSHARD view, and multi_pm (the same pass that lowers multi-device
+UNSHARDs) resolves it into per-thread shard code.
+
 Reference tilelang kernel:
-  @tilelang.jit
-  def matmul_relu(A, B, block_M=64, block_N=64, block_K=64,
-                  dtype=T.float16, accum_dtype=T.float32):
-    M, N, K = T.const('M, N, K')
-    C = T.empty([M, N], dtype)
-    with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
-      A_shared = T.alloc_shared((block_M, block_K), dtype)
-      B_shared = T.alloc_shared((block_K, block_N), dtype)
-      C_local  = T.alloc_fragment((block_M, block_N), accum_dtype)
-      T.clear(C_local)
-      for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
-        T.copy(A[by * block_M, ko * block_K], A_shared)
-        T.copy(B[ko * block_K, bx * block_N], B_shared)
-        T.gemm(A_shared, B_shared, C_local)
-      for i, j in T.Parallel(block_M, block_N):
-        C_local[i, j] = T.max(C_local[i, j], 0)
-      T.copy(C_local, C[by * block_M, bx * block_N])
-    return C
+
+    @tilelang.jit
+    def matmul_relu(A, B, block_M=64, block_N=64, block_K=64,
+                    dtype=T.float16, accum_dtype=T.float32):
+        M, N, K = T.const('M, N, K')
+        C = T.empty([M, N], dtype)
+        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
+            A_shared = T.alloc_shared((block_M, block_K), dtype)
+            B_shared = T.alloc_shared((block_K, block_N), dtype)
+            C_local  = T.alloc_fragment((block_M, block_N), accum_dtype)
+            T.clear(C_local)
+            for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
+                T.copy(A[by * block_M, ko * block_K], A_shared)
+                T.copy(B[ko * block_K, bx * block_N], B_shared)
+                T.gemm(A_shared, B_shared, C_local)
+            for i, j in T.Parallel(block_M, block_N):
+                C_local[i, j] = T.max(C_local[i, j], 0)
+            T.copy(C_local, C[by * block_M, bx * block_N])
+        return C
+
+API mapping (tilelang -> tinygrad UOps, idioms from test/backend/test_custom_kernel.py):
+
+    T.Kernel(gx, gy, threads=T)    -> AxisType.GLOBAL ranges (blocks) + AxisType.LOCAL ranges (thread grid)
+    T.alloc_shared(shape, dtype)   -> UOp.placeholder(shape, dtype, slot, AddrSpace.LOCAL)
+    T.alloc_fragment(shape, dt)    -> per-thread REG placeholder, wrapped in one Ops.UNSHARD per sharded axis over
+                                      the AxisType.LOCAL ranges: fragment.unshard((axis_y, axis_x), (ty, tx)).
+                                      The full logical tile is the shard with each sharded axis multiplied by its
+                                      range size, exactly like device sharding, but the sharding axes are thread
+                                      axes carried by the RANGE metadata instead of a device tuple. C_local[i, j]
+                                      with [i, j] in this thread's shard is INDEX on the UNSHARD, which multi_pm
+                                      resolves into INDEX on the per-thread REG shard, axis by axis.
+    T.copy(gmem_slice, smem)       -> smem[thread_idx].set(gmem_slice[thread_idx], end=copy_rng). set returns the
+                                      smem tile AFTER the copy; the implicit-barrier pass turns the store->load
+                                      dependency of the loop that consumes it into a workgroup barrier
+    T.gemm (no WMMA)               -> C_local[..].set(C_local.after(k)[..] + a_shared[..] * b_shared[..], end=k)
+                                      with k a loop-carried LOOP range (codegen builds the register accumulator
+                                      from this self-referential store automatically)
+    T.copy(fragment, gmem)         -> gmem.index(gidx).store(C_local[..]).end(all_ranges)
+    UNSHARD lowering               -> multi_pm in codegen (full_rewrite_to_sink): INDEX/AFTER/STORE ops on the
+                                      full-tile view become per-thread shard ops, no UNSHARD survives into the program.
 """

 from tinygrad.dtype import dtypes, AddrSpace, DType
@@ -31,17 +61,25 @@ from tinygrad.tensor import Tensor
 # tilelang builtins, expressed with tinygrad UOp APIs
 # ---------------------------------------------------------------------------

-def alloc_shared(shape:tuple[int, ...], dtype:DType, slot:int) -> UOp:
+def alloc_shared(shape:tuple[int, ...], dtype:DType) -> UOp:
  """T.alloc_shared: one LOCAL buffer shared by all threads in the block."""
-  return UOp.placeholder(tuple(shape), dtype, slot, AddrSpace.LOCAL)
+  return UOp.placeholder(tuple(shape), dtype, next(UOp.unique_num), AddrSpace.LOCAL)

-def alloc_fragment(shape:tuple[int, ...], dtype:DType, slot:int, axes:tuple[int, ...], rngs:tuple[UOp, ...]) -> UOp:
-  """T.alloc_fragment: per-thread REG fragment + UNSHARD over the LOCAL thread grid."""
+def alloc_fragment(shape:tuple[int, ...], dtype:DType, axes:tuple[int, ...], rngs:tuple[UOp, ...]) -> UOp:
+  """T.alloc_fragment: per-thread REG fragment + UNSHARD over the LOCAL thread grid.
+
+  Each thread privately owns shape[axis]//threads elements along every sharded
+  axis in a REG buffer. The UNSHARDs over the LOCAL thread ranges present the
+  full logical tile: full_shape = shard_shape with each sharded axis multiplied
+  by its range size. This is exactly how UNSHARD carries a DEVICE axis today,
+  except the sharding axes are thread axes carried by the RANGE metadata.
+  """
  assert len(axes) == len(rngs)
  assert all(tnum.op is Ops.RANGE and tnum.arg[-1] is AxisType.LOCAL for tnum in rngs), "fragments shard over LOCAL ranges"
+  assert all(shape[a] % (int(rng.vmax)+1) == 0 for a, rng in zip(axes, rngs))
  by_axis = dict(zip(axes, rngs))
  shard_shape = tuple(s // (int(by_axis[i].vmax)+1) if i in by_axis else s for i, s in enumerate(shape))
-  fragment = UOp.placeholder(shard_shape, dtype, slot, AddrSpace.REG)
+  fragment = UOp.placeholder(shard_shape, dtype, next(UOp.unique_num), AddrSpace.REG)
  return fragment.unshard(axes, rngs)

 # ---------------------------------------------------------------------------
@@ -68,45 +106,64 @@ def matmul_relu_kernel(c:UOp, a:UOp, b:UOp) -> UOp:
  # with T.Kernel(T.ceildiv(N, BLOCK_N), T.ceildiv(M, BLOCK_M), threads=128) as (bx, by):
  bx = UOp.range(cdiv(N, BLOCK_N), 0, AxisType.GLOBAL)
  by = UOp.range(cdiv(M, BLOCK_M), 1, AxisType.GLOBAL)
-
-  # 16*8 threads = 128 threads
+  # tx (N, 16) is the fast/inner LOCAL axis so a warp covers 16 cols x 2 rows --
+  # matching tilelang's (tidx>>4, tidx&15) warp composition. This keeps the 8 A_shared
+  # reads in a warp on only 2 row-groups (broadcast across 16 cols) instead of 8 rows
+  # (8-way bank conflict), since A_shared[row*512 + ...] all map to the same bank when 8
+  # distinct rows land in one warp.
  tx = UOp.range(TX, 2, AxisType.LOCAL)
  ty = UOp.range(TY, 3, AxisType.LOCAL)

-  # shared + fragment (regs)
-  A_shared = alloc_shared((BLOCK_M, BLOCK_K), a.dtype, 0)
-  B_shared = alloc_shared((BLOCK_K, BLOCK_N), b.dtype, 1)
-  C_local = alloc_fragment((TM, TY, TX, TN), dtypes.float32, 0, (1, 2), (ty, tx))
+  # A_shared = T.alloc_shared((BLOCK_M, BLOCK_K), dtype)
+  # B_shared = T.alloc_shared((BLOCK_K, BLOCK_N), dtype)
+  A_shared = alloc_shared((BLOCK_M, BLOCK_K), a.dtype)
+  B_shared = alloc_shared((BLOCK_K, BLOCK_N), b.dtype)

-  # zero out the regs to start. this is expanded by the devectorizer
-  C_local = C_local.after(C_local.store(0.0))
+  # C_local = T.alloc_fragment((BLOCK_M, BLOCK_N), accum_dtype) -- an 8x4 REG tile per thread of the 8x16 grid
+  C_local = alloc_fragment((BLOCK_M, BLOCK_N), dtypes.float32, (0, 1), (ty, tx))
+
+  # T.clear(C_local) -- each thread zeroes its own fragment sub-tile
+  ic, jc = UOp.range(TM, 4, AxisType.LOOP), UOp.range(TN, 5, AxisType.UPCAST)
+  C_loc = C_local[ic*TM + ty, tx*TN + jc].set(0.0, end=(ic, jc))

  # for ko in T.Pipelined(T.ceildiv(K, BLOCK_K), num_stages=3):
+  # (num_stages pipelining is async copy + multi-buffering; this is the synchronous single-buffer version)
  ko = UOp.range(cdiv(K, BLOCK_K), 6, AxisType.LOOP)

-  # index the outer matrices
-  a = a.rearrange("(m bm) (k bk) -> m k bm bk", bm=BLOCK_M, bk=BLOCK_K)[by, ko]
-  b = b.rearrange("(k bk) (n bn) -> k n bk bn", bk=BLOCK_K, bn=BLOCK_N)[ko, bx]
-  c = c.rearrange("(m bm) (n bn) -> m n bm bn", bm=BLOCK_M, bn=BLOCK_N)[by, bx]
+  # T.copy(A[by * BLOCK_M, ko * BLOCK_K], A_shared) -- each thread copies its own 8x4 sub-tile.
+  # Row index is iar*TM + ty (strided by TM across ty), matching tilelang's layout: thread ty owns
+  # rows {ty, ty+8, ..., ty+56} not {ty*8, ..., ty*8+7}.
+  iar, ka = UOp.range(TM, 7, AxisType.LOOP), UOp.range(TN, 8, AxisType.UPCAST)
+  A_store = A_shared[iar*TM + ty, tx*TN + ka].store(a[by*BLOCK_M + iar*TM + ty, ko*BLOCK_K + tx*TN + ka]).end(iar, ka)

-  # T.copy: A_shared <- a, B_shared <- b
-  def with_threads(x:UOp): return x.rearrange("(tm ty) (tx tn) -> ty tx tm tn", tm=TM, tn=TN)[ty, tx]
-  A_shared = A_shared.after(with_threads(A_shared).store(with_threads(a)))
-  B_shared = B_shared.after(with_threads(B_shared).store(with_threads(b)))
+  # T.copy(B[ko * BLOCK_K, bx * BLOCK_N], B_shared)
+  kb, ibr = UOp.range(TM, 9, AxisType.LOOP), UOp.range(TN, 10, AxisType.UPCAST)
+  B_store = B_shared[kb*TM + ty, tx*TN + ibr].store(b[ko*BLOCK_K + kb*TM + ty, bx*BLOCK_N + tx*TN + ibr]).end(kb, ibr)

-  # T.gemm(A_shared, B_shared, C_local), no WMMA
-  kk = UOp.range(BLOCK_K, 11, AxisType.LOOP)
-  ir = UOp.range(TM, 12, AxisType.LOOP)
+  # get the shared after the stores (single barrier)
+  A_shared = A_shared.after(A_store, B_store)
+  B_shared = B_shared.after(A_store, B_store)
+
+  # T.gemm(A_shared, B_shared, C_local), no WMMA -- per-thread accumulate over its fragment sub-tile.
+  # identical to custom_gemm: a self-referential store over the loop-carried kk range,
+  # which codegen turns into a register accumulator
+  # kk is the outer compute loop (axis 11) so that for each kk we read all 8 A rows and reuse
+  # the B[kk] read across them -- matching tilelang's ko > kk > row > col access order exactly.
+  kk, ir = UOp.range(BLOCK_K, 11, AxisType.LOOP), UOp.range(TM, 12, AxisType.LOOP)
  jj = UOp.range(TN, 13, AxisType.UPCAST)
-  acc = C_local.after(kk)[ir, ty, tx, jj] + A_shared[ir*TM + ty, kk].cast(dtypes.float32) * B_shared[kk, tx*TN + jj].cast(dtypes.float32)
+  acc = C_loc.after(kk)[ir*TM + ty, tx*TN + jj] + A_shared[ir*TM + ty, kk].cast(dtypes.float32) * B_shared[kk, tx*TN + jj].cast(dtypes.float32)
  # closing the ko loop here too; codegen adds the barrier so no thread overwrites the tiles while others still read them
-  C_local = C_local[ir, ty, tx, jj].set(acc, end=(kk, ir, jj, ko))
+  C_loc = C_loc[ir*TM + ty, tx*TN + jj].set(acc, end=(kk, ir, jj, ko))

-  # c <- C_local (with relu and cast): every thread stores its shard's sub-view of the output tile
-  c_st = c.reshape(C_local.shape).store(C_local.relu().cast(c.dtype))
+  # for i, j in T.Parallel(BLOCK_M, BLOCK_N): C_local[i, j] = T.max(C_local[i, j], 0)
+  # T.copy(C_local, C[by * BLOCK_M, bx * BLOCK_N]) -- per-thread store of the fragment shard (relu fused into it)
+  # LOOP: these loops are the per-thread output layout; convert_loop_to_global must not globalize them
+  ie, je = UOp.range(TM, 14, AxisType.LOOP), UOp.range(TN, 15, AxisType.UPCAST)
+  c_st = c[by*BLOCK_M + ie*TM + ty, bx*BLOCK_N + tx*TN + je].store(C_loc[ie*TM + ty, tx*TN + je].relu().cast(c.dtype))

-  # close the locals and globals
-  return c_st.end(tx, ty, bx, by).sink(arg=KernelInfo(name="matmul_relu", opts_to_apply=()))
+  # all open ranges are closed at the final store (ko was closed above).
+  # the fragment UNSHARDs go to codegen as is: multi_pm there resolves the full-tile view into per-thread shard code
+  return c_st.end(je, ie, tx, ty, bx, by).sink(arg=KernelInfo(name="matmul_relu", opts_to_apply=()))

 # ---------------------------------------------------------------------------
 # python wrapper: same signature as the tilelang function
@@ -131,8 +188,7 @@ if __name__ == "__main__":
  b = Tensor.randn(K, N, dtype=dtype_in).contiguous()
  ref = (a @ b).relu().realize()

-  for _ in range(10):
-    out = matmul_relu(a, b).realize()
+  out = matmul_relu(a, b).realize()

  import numpy as np
  np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=1e-1, rtol=1e-2)
@@ -3,8 +3,8 @@ from typing import cast, Any, Callable
 import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools, collections, atexit
 assert sys.platform != 'win32'
 from dataclasses import dataclass
-from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, encode_kernargs_clike, make_cmdbuf
-from tinygrad.runtime.support.hcq2 import make_binary_patch
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, encode_kernargs_clike, make_cmdbuf
+from tinygrad.runtime.support.hcq2 import make_binary_patch, make_patches
 from tinygrad.uop.ops import sint, UOp
 from tinygrad.device import Compiled, BufferSpec, Buffer, Device
 from tinygrad.dtype import dtypes
@@ -37,7 +37,7 @@ class PM4Ops(FastEnum):
  RELEASE_MEM = auto(); DISPATCH_DIRECT = auto(); EVENT_WRITE = auto()  # noqa: E702

 def pkt3(ctx, op:PM4Ops, *vals):
-  return UOp(Ops.INS, arg=op, src=tuple(UOp.const(x, dtypes.uint32)
+  return UOp(Ops.INS, arg=op, src=tuple(UOp.const(dtypes.uint32, x)
    for x in (ctx.pm4.PACKET3(getattr(ctx.pm4, f"PACKET3_{op.name}"), len(vals) - 1), *vals)))

 def wreg(ctx, reg:AMDReg, *args:sint, **kwargs:int):
@@ -152,22 +152,28 @@ def pm4_submit(ctx, lin):
  ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COMPUTE:0_{name}")
    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))

-  # the host fence at the start of the batch guarantees the ib is free to reuse
-  size_dw = sum(len(ins.src) for ins in lin.src)
+  # two tail dwords coordinate safe IB reuse: GPU completions and host submits
+  size_dw = sum(len(ins.src) for ins in lin.src) + len(release_mem(ctx, 0, 0).src)
  assert size_dw < (1 << 20), f"indirect buffer of {size_dw} dwords doesn't fit one packet"

-  ib = UOp.placeholder((size_dw,), dtypes.uint32, next(UOp.unique_num), device=devs, volatile=True).rtag("cmdbuf")
-  cmdbuf = make_cmdbuf(lin, devs, buf=ib)
+  ib = UOp.placeholder((size_dw + 2,), dtypes.uint32, next(UOp.unique_num), device=devs, volatile=True).rtag("cmdbuf")
+  done_idx, submit_idx = UOp.const(dtypes.int, size_dw + 0), UOp.const(dtypes.int, size_dw + 1)
+  submitted = (counter:=ib.after(make_patches(ib, [((size_dw + i) * 4, UOp.const(dtypes.uint32, 0)) for i in range(2)])).index(submit_idx)).load()
+  completed = ib.after(loop:=UOp.loop(0)).index(done_idx).load()
+  ib_free = completed.end(loop, completed != submitted)
+
+  bump_fence = pm4_store(ctx, UOp(Ops.SLICE, dtypes.uint32, (ib, UOp.const(dtypes.weakint, size_dw)), 2), (submitted + 1).cast(dtypes.uint64))
+  cmdbuf = make_cmdbuf(lin.replace(src=lin.src + (bump_fence,)), devs, buf=ib, dep=ib_free)

  # the ring itself only carries a packet pointing at the ib, wrapping the ring
-  put = put_ptr.index(zero:=UOp.const(0, dtypes.int))
+  put = put_ptr.index(zero:=UOp.const(dtypes.int, 0))
  pkt = (ctx.pm4.PACKET3(ctx.pm4.PACKET3_INDIRECT_BUFFER, 2), *data64_le(cmdbuf.getaddr(devs)), size_dw | ctx.pm4.INDIRECT_BUFFER_VALID)
-  write_pkt = UOp.barrier(*[ring.index(((put + off) % q.ring.size).cast(dtypes.int)).store(UOp.const(x, dtypes.uint32)) for off,x in enumerate(pkt)])
+  write_pkt = UOp.barrier(*[ring.index(((put + off) % q.ring.size).cast(dtypes.int)).store(UOp.const(dtypes.uint32, x)) for off,x in enumerate(pkt)])

  # advance the put/write pointers past the packet
  bump_put_ptr = put_ptr.index(zero).store(put + len(pkt))
  bump_wptr = wptr.index(zero).store(put + len(pkt))
-  flush = UOp.barrier(write_pkt, bump_put_ptr, bump_wptr)
+  flush = UOp.barrier(write_pkt, bump_put_ptr, bump_wptr, counter.store(submitted + 1))
  return doorbell.after(flush).index(zero).store(put + len(pkt))

 pm_pm4_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"), pm4_submit)])
@@ -180,27 +186,26 @@ class SDMAOps(FastEnum): COPY = auto(); POLL_REGMEM = auto(); FENCE = auto(); TR
 def sdma_copy(ctx, call):
  sz = call.src[2].max_numel() * call.src[2].dtype.itemsize
  src_addr, dst_addr = call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs)
-  return call.ins(SDMAOps.COPY, src=tuple(UOp.const(x, dtypes.uint32) for off in range(0, sz, ctx.max_copy_size) for x in (
+  return call.ins(SDMAOps.COPY, src=tuple(UOp.const(dtypes.uint32, x) for off in range(0, sz, ctx.max_copy_size) for x in (
    ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR),
-    ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0,
-    *data64_le(src_addr+UOp.const(off, dtypes.uint64)), *data64_le(dst_addr+UOp.const(off, dtypes.uint64)))))
+    ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0, *data64_le(src_addr+off), *data64_le(dst_addr+off))))

 def sdma_wait(ctx, ins, dst, val):
  op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \
     | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1)
-  return ins.ins(SDMAOps.POLL_REGMEM, src=tuple(UOp.const(x, dtypes.uint32) for x in (
+  return ins.ins(SDMAOps.POLL_REGMEM, src=tuple(UOp.const(dtypes.uint32, x) for x in (
    op, *data64_le(dst.getaddr(ctx.devs)), val, 0xffffffff,
    ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))))

 def sdma_store(ctx, ins, dst, val):
  op = ctx.sdma.SDMA_OP_FENCE | (ctx.sdma.SDMA_PKT_FENCE_HEADER_MTYPE(3) if ctx.target[0] != 9 else 0)
  return UOp(Ops.LINEAR, src=(
-    ins.ins(SDMAOps.FENCE, src=tuple(UOp.const(x, dtypes.uint32) for x in (op, *data64_le(dst.getaddr(ctx.devs)), val))),
-    ins.ins(SDMAOps.TRAP, src=tuple(UOp.const(x, dtypes.uint32) for x in (ctx.sdma.SDMA_OP_TRAP, 0)))))
+    ins.ins(SDMAOps.FENCE, src=tuple(UOp.const(dtypes.uint32, x) for x in (op, *data64_le(dst.getaddr(ctx.devs)), val))),
+    ins.ins(SDMAOps.TRAP, src=tuple(UOp.const(dtypes.uint32, x) for x in (ctx.sdma.SDMA_OP_TRAP, 0)))))

 def sdma_timestamp(ctx, ins, dst):
  op = ctx.sdma.SDMA_OP_TIMESTAMP | ctx.sdma.SDMA_PKT_TIMESTAMP_GET_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_TIMESTAMP_GET_GLOBAL)
-  return ins.ins(SDMAOps.TIMESTAMP, src=tuple(UOp.const(x, dtypes.uint32) for x in (op, *data64_le(dst.getaddr(ctx.devs)))))
+  return ins.ins(SDMAOps.TIMESTAMP, src=tuple(UOp.const(dtypes.uint32, x) for x in (op, *data64_le(dst.getaddr(ctx.devs)))))

 pm_sdma_opsel = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy),
@@ -213,7 +218,7 @@ pm_sdma_opsel = PatternMatcher([

 def sdma_submit(cmdbuf, devs):
  # the cmdbuf to submit + the patch writes that fill it
-  size_dw, zero = cmdbuf.nbytes() // dtypes.uint32.itemsize, UOp.const(0, dtypes.int)
+  size_dw, zero = cmdbuf.nbytes() // dtypes.uint32.itemsize, UOp.const(dtypes.int, 0)

  # the sdma queue's ring and its host-side ring/write/put pointers
  for d in devs: q = Device[d].sdma_queue(0)
@@ -229,8 +234,8 @@ def sdma_submit(cmdbuf, devs):

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

  # advance the put/write pointers past the zeroed tail and the cmdbuf
@@ -288,16 +293,14 @@ 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_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb())
+    super().__init__(dev, supports_copy_from_disk=dev.has_sdma_queue, supports_transfer=dev.has_sdma_queue and not dev.is_usb())

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

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

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

@dataclass
 class AMDQueueDesc:
@@ -390,24 +393,15 @@ 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():
-      mapped = self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
-      mapped._owns_kfd_handle = True
-      return mapped
+    if mem.owner is not None and mem.owner._is_cpu(): return self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)

    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)
@@ -479,7 +473,6 @@ 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
@@ -520,12 +513,12 @@ class PCIIface(PCIIfaceBase):
    if drain_only: d.iface.dev_impl.ih.drain()
    else: d.iface.dev_impl.ih.interrupt_handler()

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

  def sleep(self, timeout):
    if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
@@ -551,12 +544,9 @@ 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

@@ -593,7 +583,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_copy_queue = not getenv("AMD_DISABLE_SDMA")
+    self.has_sdma_queue = True # self.sdma_queue(0) is not None, TODO: think of this

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

@@ -648,6 +638,9 @@ class AMDDevice(HCQ2Compiled):
    qname = f"{'COPY' if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA else 'COMPUTE'}:{idx}"
    self.pm_bufferize = PatternMatcher([
      (UPat(Ops.PARAM, tag=f"{qname}_{name}"), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"]
+    ] + [
+      (UPat(Ops.PARAM, tag=f"{qname}_timeline_signal"), lambda ctx, q=qname: ctx[0].timeline_signal(q)),
+      (UPat(Ops.PARAM, tag=f"{qname}_timeline_value"), lambda ctx, q=qname: ctx[0].timeline_value(q)),
    ]) + self.pm_bufferize

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

  def scratch_buffer(self, private_segment_size):
-    AMDDevice.max_scratch_psize = private_segment_size = max(private_segment_size, 128, AMDDevice.max_scratch_psize)
+    private_segment_size = max(private_segment_size, 128)
    if self.max_private_segment_size < private_segment_size:
      lanes_per_wave = 64 # wave64
      mem_alignment_size = 256 if self.target[0] != 9 else 1024
@@ -48,7 +48,7 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:
  device = device[0].split(":")[0] if isinstance(device, tuple) else device.split(":")[0]
  if device in {"AMD", "NULL"}: atomic_arg = "if ({2} > {3}) __hip_atomic_fetch_max((int*){0}, {1}, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT);"
  else: raise NotImplementedError(f"no atomic max for device {device}")
-  amax_idx = amax_out.reshape((1,)).index(UOp.const(0))
+  amax_idx = amax_out.reshape((1,)).index(UOp.const(dtypes.weakint, 0))
  max_val = lds[0].load()
  atomic = UOp(Ops.CUSTOM, dtypes.void, (amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=atomic_arg)
  return atomic.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=()))
@@ -1,36 +0,0 @@
-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])
@@ -1,226 +0,0 @@
-// 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();
-      }
-    }
-  }
-}
@@ -1,49 +0,0 @@
-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]
@@ -1,217 +0,0 @@
-# Runbook: Llama 3 8B Training on DigitalOcean MI350X
-
-## Machine Specs
- 8x MI350X GPUs (gfx950, device ID 75b0), 288GB VRAM each
- 2TB RAM, 192 CPUs, 2TB disk
- ROCm 7.14 at `/opt/rocm` (NOT `/opt/rocm-7.1.1` like the submission scripts assume)
- Python 3.12
-
-## Phase 1: System Setup
-
-### 1.1 Install packages
-```bash
-apt-get update
-apt-get install -y python3-pip python3-venv git tmux rclone clang
-```
-
-### 1.2 Install Python deps
-```bash
-python3 -m pip install --break-system-packages --ignore-installed typing-extensions numpy tqdm wandb tiktoken sentencepiece
-```
-Note: `--ignore-installed typing-extensions` is needed because the base image ships typing-extensions 4.10.0 without a RECORD file, so pip cannot uninstall it.
-
-### 1.3 Install ROCm dev headers
-The base image has ROCm runtime but NOT the HIP dev headers. Need:
-```bash
-apt-get install -y amdrocm-core-dev
-```
-This installs `hip/hip_runtime.h` at `/opt/rocm/core-7.14/include/hip/hip_runtime.h`.
-The symlink `/opt/rocm/include` → `/opt/rocm/core-7.14/include` makes it available at `/opt/rocm/include/hip/hip_runtime.h`.
-
-### 1.4 Configure ROCm comgr
-ROCm 7.14 ships comgr 3.3 at `/opt/rocm/lib/libamd_comgr.so`. tinygrad's DLL loader needs explicit env vars to find it (it searches for `libcomgr.so*` by default, not `libamd_comgr.so*`). Set these in the run command:
-```bash
-export COMGR_PATH=/opt/rocm/lib/libamd_comgr.so
-export COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so
-```
-Also add ROCm libs to ldconfig so comgr's shared library dependencies resolve:
-```bash
-cat > /etc/ld.so.conf.d/rocm.conf << 'EOF'
-/opt/rocm/lib
-/opt/rocm/lib/llvm/lib
-/opt/rocm/lib/rocm_sysdeps/lib
-EOF
-ldconfig
-```
-
-### 1.5 Install geohot tmux config
-```bash
-curl -sL https://raw.githubusercontent.com/geohot/configuration/master/.tmux.conf -o ~/.tmux.conf
-```
-
-### 1.6 Reload amdgpu driver
-tinygrad's HCQ backend needs `/dev/kfd` which is created by the amdgpu kernel driver.
-If the driver was unloaded, reload it:
-```bash
-modprobe amdgpu
-ls /dev/kfd  # should exist
-```
-
-## Phase 2: Clone tinygrad
-```bash
-cd /root
-git clone https://github.com/tinygrad/tinygrad.git
-cd tinygrad
-python3 -m pip install --break-system-packages -e .
-```
-
-## Phase 3: Download C4 Dataset
-
-The C4 data is on the MLCommons Cloudflare R2 bucket in Megatron-LM indexed format.
-
-```bash
-rclone config create mlc-training s3 provider=Cloudflare \
-  access_key_id=76ea42eadb867e854061a1806220ee1e \
-  secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 \
-  endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
-
-mkdir -p /raid/datasets/c4-8b
-rclone copy mlc-training:mlcommons-training-wg-public/llama3_1/datasets/c4/llama3_1_8b/ /raid/datasets/c4-8b/ -P
-```
-
-Files downloaded (~85GB total, ~6 minutes):
- `c4-train.en_6_text_document.bin` (79 GB)
- `c4-train.en_6_text_document.idx` (870 MB)
- `c4-validation-91205-samples.en_text_document.bin` (159 MB)
- `c4-validation-91205-samples.en_text_document.idx` (1.8 MB)
- `LICENSE.txt`, `NOTICE.txt`
-
-**Wait for rclone to fully complete before starting training.** Starting training while the dataset is still downloading will read a truncated .bin file, causing `ValueError: all input arrays must have the same shape` in the dataloader. The stale `.index_cache` and `.blend_cache` files must also be deleted if this happens:
-```bash
-rm -f /raid/datasets/c4-8b/*.index_cache /raid/datasets/c4-8b/*.blend_cache
-```
-
-## Phase 4: wandb Login
-```bash
-wandb login
-```
-Enter API key from https://wandb.ai/authorize
-
-Alternatively, pass the key directly:
-```bash
-wandb login <API_KEY>
-```
-
-## Phase 5: Run Training
-
-Run training in tmux so it survives SSH disconnects:
-```bash
-tmux new-session -d -s train 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=AMD:HIP ROCM_PATH=/opt/rocm WANDB=1 bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh 2>&1 | tee /root/train.log'
-```
-Attach with `tmux attach -t train`.
-
-### 5.1 Smoke test (beam search, 2 layers, real data)
-Always run beam first to validate the pipeline:
-```bash
-tmux new-session -d -s beam 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=AMD:HIP ROCM_PATH=/opt/rocm bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh 2>&1 | tee /root/beam.log'
-```
-
-The beam test runs 10 training steps with 2 layers. Expected results:
- ~0.29s per step after warmup
- ~700K GFLOPS, ~7% MFU (low because only 2 layers)
- ~380 GB VRAM used
- Loss stable at ~12.55 with random init
-
-### 5.2 Full training run
-```bash
-tmux new-session -d -s train 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=AMD:HIP ROCM_PATH=/opt/rocm WANDB=1 bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh 2>&1 | tee /root/train.log'
-```
-
-## Environment Variable Reference
-
-| Variable | Value | Why |
-|---|---|---|
-| `COMGR_PATH` | `/opt/rocm/lib/libamd_comgr.so` | tinygrad's DLL loader needs explicit path to find comgr 3.3 |
-| `COMGR_3_PATH` | `/opt/rocm/lib/libamd_comgr.so` | comgr 3.x uses a separate `comgr_3` module with its own path var |
-| `CC` | `/opt/rocm/core-7.14/lib/llvm/bin/clang` | System clang doesn't know gfx950; must use ROCm's bundled clang |
-| `DEV` | `AMD:HIP` | Force HIPRenderer (comgr-based) over HIPCCRenderer (hipcc subprocess) |
-| `ROCM_PATH` | `/opt/rocm` | Script defaults to `/opt/rocm-7.1.1` which doesn't exist |
-| `WANDB` | `1` | Enable wandb logging (off by default) |
-
-## Architecture
-
-| Component | Source file |
-|---|---|
-| Model | `examples/mlperf/models/flat_llama.py` — FlatTransformer, FP8 MXFP4 weights, fused QKV, flash attention |
-| Trainer | `examples/mlperf/model_train.py` → `train_llama3()` |
-| Optimizer | `examples/mlperf/optim.py` — GradAccClipAdamW, master weights, FP8 re-quant |
-| LR schedule | `examples/mlperf/lr_schedulers.py` — CosineAnnealingLRWithWarmup |
-| Dataloader | `examples/mlperf/dataloader.py` — Megatron-LM indexed bin format |
-| ASM GEMM | `extra/gemm/cdna_asm_gemm.py` — gfx950 MFMA assembly, MXFP4 |
-| Flash attention | `extra/thunder/amd/fa.py` |
-| Fused kernels | `extra/llama_kernels/` — rmsnorm, silu, quantize, fused_ce |
-| GPU driver | `tinygrad/runtime/ops_amd.py` — HCQ, direct KFD ioctl |
-| Renderer | `tinygrad/renderer/cstyle.py` — HIPRenderer for gfx950 |
-| comgr compiler | `tinygrad/runtime/support/compiler_amd.py` — HIPCompiler using comgr 3.3 |
-
-## Troubleshooting
-
-### `'hip/hip_runtime.h' file not found`
-Install `amdrocm-core-dev`:
-```bash
-apt-get install -y amdrocm-core-dev
-```
-
-### `'gfx950' is not a recognized processor` + LLVM crash
-System clang doesn't know gfx950. Set `CC=/opt/rocm/core-7.14/lib/llvm/bin/clang`.
-
-### `comgr not available: try setting COMGR_PATH?`
-Add ROCm libs to ldconfig and set `COMGR_PATH` and `COMGR_3_PATH`:
-```bash
-# /etc/ld.so.conf.d/rocm.conf should contain /opt/rocm/lib paths
-ldconfig
-```
-
-### `comgr not available: try setting COMGR_3_PATH?`
-comgr 3.x uses a separate module. Set `COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so` too.
-
-### `No such file or directory: 'clang'`
-Install clang: `apt-get install -y clang` (for CPU compilation).
-For gfx950 HIP compilation, comgr (not clang) is used — ensure the ROCm 7.14 comgr 3.3 is properly loaded via `COMGR_PATH` and `COMGR_3_PATH`.
-
-## Appendix: KVM Virtualization Observations
-
-### Virtualization detection
-```
-$ systemd-detect-virt
-kvm
-$ lspci -nn | grep AMD
-83:00.0 ... Device [1002:75b0]
-```
-CPU flags include `hypervisor`. `dmesg` shows `Hypervisor detected: KVM`.
-
-### Working path: amdgpu driver (KFDIface)
-The amdgpu driver loads on boot and binds to all 8 GPUs, creating `/dev/kfd` and 64 renderD nodes (`/dev/dri/renderD128` through `/dev/dri/renderD191`). tinygrad's `KFDIface` enumerates GPUs through `/sys/devices/virtual/kfd/kfd/topology/nodes` and uses `/dev/kfd` for ioctl. No PCI device ID patching is needed — the KFD path does not use `PCIIface` or `AMDev._run_discovery()`.
-
-This is the working configuration. No code changes to tinygrad are required.
-
-### PCIIface path (does not work on this VM)
-For reference, the `PCIIface` path was also explored but does not work in this KVM guest:
-
- `PCIIface` in `ops_amd.py` does not list device ID `0x75b0`. Adding it allows PCI detection but `AMDev._run_discovery()` fails because the VRAM BAR reads all `0xFF`.
- This was observed with the GPU unbound from any driver, after PCI reset, and with VFIO bound.
- VFIO binding (`vfio-pci` with `enable_unsafe_noiommu_mode=1`) succeeded but VRAM BAR still reads all `0xFF`.
- No IOMMU in guest — `dmesg` has no `AMD-Vi` entries, PCI devices have no `iommu_group` symlink.
-
-### amdgpu driver behavior
-On first boot, amdgpu loaded and bound to all 8 GPUs. On one boot it failed to initialize:
-```
-[  799.780369] amdgpu 0000:83:00.0: Failed to alloc msi vectors
-[  799.781476] amdgpu 0000:83:00.0: sw_init of IP block <vega20_ih> failed -22
-[  799.782724] amdgpu 0000:83:00.0: amdgpu_device_ip_init failed
-[  799.793885] amdgpu 0000:83:00.0: Fatal error during GPU init
-```
-On a subsequent boot, amdgpu initialized successfully (SMU initialized, VRAM ready). After unbinding all 8 GPUs from amdgpu, `rmmod amdgpu` wedged the module (stuck in "Unloading" state in `/proc/modules`), requiring a full VM reboot.
-
-### No fan control
-No `fan*` or `pwm*` hwmon entries exist. Only `temp*`, `power*`, `freq*` are exposed. GPU temps read 56-63°C, power ~265W per GPU.
@@ -2,7 +2,7 @@ import math, pathlib, functools, struct

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

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

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

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

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

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

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

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

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

  return attn, attn, l_vec

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

  Q_BLOCK_SIZE = 32
  NUM_WARPS = 8
@@ -252,11 +206,10 @@ 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)
-  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)
+  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)))
@@ -283,20 +236,19 @@ 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)
-  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)
+  lib = bytearray(lib)
+  rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
+  struct.pack_into('<I', lib, rodata_off, 160000)
+  lib = bytes(lib)

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

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

  BLOCK_SIZE_KV = 256
  GROUP_SIZE = H // H_KV
@@ -316,11 +268,10 @@ 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)
-  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)
+  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)))
@@ -347,11 +298,10 @@ 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)
-  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)
+  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,10 +43,6 @@ 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)
@@ -92,12 +88,7 @@ __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;

@@ -389,13 +380,6 @@ __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);
@@ -549,8 +533,6 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
-        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
-        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -656,13 +638,6 @@ __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);
@@ -816,8 +791,6 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
-        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
-        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -922,14 +895,6 @@ __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);
@@ -1083,8 +1048,6 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
-        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
-        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -1188,15 +1151,6 @@ __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);
@@ -1349,8 +1303,6 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
-        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
-        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -1476,13 +1428,6 @@ __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);
@@ -1637,8 +1582,6 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
-        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
-        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -1746,13 +1689,6 @@ __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);
@@ -1906,8 +1842,6 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
-        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
-        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -2012,14 +1946,6 @@ __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);
@@ -2173,8 +2099,6 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[tic][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
-        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
-        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -2278,15 +2202,6 @@ __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);
@@ -2439,8 +2354,6 @@ __global__ void attend_bwd_combined_ker(bf16 *dQ_ptr, bf16 *dK_ptr, bf16 *dV_ptr
        if constexpr (D == 128) load<0, 2>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}), Q_i_addr);
        if constexpr (D == 128) load<0, 3>(Q_i, subtile_inplace<DOT_SLICE_QO, D>(Q_i_smem[toc][0], {0, 0}), Q_i_addr);
        mma_AtB<0, 0, 7>(dQ_i_T, K_j_col, dP_ij_bf16_col_T, dQ_i_T);
-        // D=64: wait out MFMA->VALU accumulator hazard on dQ_i_T
-        if constexpr (D == 64) asm volatile("s_nop 15");
        if constexpr (D == 128) mma_AtB<1, 0, 0>(dQ_i_T, K_j_col, dP_ij_bf16_col_T);
        // Load K_j from shared memory to registers
        // load(K_j, subtile_inplace<WARP_SIZE_KV, D>(K_j_smem, {warpid, 0}));
@@ -2558,12 +2471,6 @@ __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);
@@ -2825,13 +2732,6 @@ __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);
@@ -3088,14 +2988,6 @@ __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);
@@ -3352,15 +3244,6 @@ __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,10 +34,6 @@ 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)
@@ -86,26 +82,11 @@ 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 rel_hi_vgpr, uint32_t neg_inf_vgpr,
+__device__ inline void mask_vec2_imm(uint32_t rel_vgpr, uint32_t neg_inf_vgpr,
                                     uint32_t& x_ref, uint32_t& y_ref) {

    uint64_t x_mask, y_mask;
-#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
+    // uint32_t ox, oy;
    asm volatile(
        // x: rel < THR_X ?
        "v_cmp_lt_i32_e64 %0, %6, %7\n\t"
@@ -118,7 +99,7 @@ __device__ inline void mask_vec2_imm(uint32_t rel_vgpr, uint32_t rel_hi_vgpr, ui
          "n"(THR_X), "v"(neg_inf_vgpr), "n"(THR_Y)
        : "vcc"
    );
-#endif
+    // x_ref = ox; y_ref = oy;
 }

 template<ducks::rt::col_layout RT>
@@ -141,8 +122,6 @@ __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) {
@@ -166,14 +145,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, 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);
+            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);
        }
    }
 }
@@ -222,16 +201,6 @@ __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;

@@ -262,7 +231,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, min_tile, head_idx_kv, 0}, swizzled_offsets_K);
+    G::load<1, false>(k_smem[0], g.Kg, {batch_idx, 0, head_idx_kv, 0}, swizzled_offsets_K);
    __builtin_amdgcn_s_waitcnt(0);
    __builtin_amdgcn_sched_barrier(0);
    __builtin_amdgcn_s_barrier();
@@ -274,9 +243,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, min_tile + 1, head_idx_kv, 0}, swizzled_offsets_K);
+    G::load<1, false>(k_smem[1], g.Kg, {batch_idx, 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, min_tile, head_idx_kv, 0}, swizzled_offsets_V);
+    G::load<1, false>(v_smem[0], g.Vg, {batch_idx, 0, head_idx_kv, 0}, swizzled_offsets_V);
    load(k_reg, k_smem[0]);
    __builtin_amdgcn_sched_barrier(0);
    asm volatile("s_waitcnt lgkmcnt(0)");
@@ -290,20 +259,13 @@ __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 = (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);
+        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);
        }
    }
    // 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);
@@ -322,25 +284,21 @@ __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, min_tile + 2, head_idx_kv, 0}, swizzled_offsets_K);
+    G::load<1, false>(k_smem[0], g.Kg, {batch_idx, 2, head_idx_kv, 0}, swizzled_offsets_K);
    // All warps then collaboratively load in the second slice of V (V1) into shared memory
-    G::load<1, false>(v_smem[1], g.Vg, {batch_idx, min_tile + 1, head_idx_kv, 0}, swizzled_offsets_V);
+    G::load<1, false>(v_smem[1], g.Vg, {batch_idx, 1, head_idx_kv, 0}, swizzled_offsets_V);
    asm volatile("s_waitcnt lgkmcnt(0)");
    asm volatile("s_waitcnt vmcnt(" FA_VM4 ")");
    __builtin_amdgcn_sched_barrier(0);
    __builtin_amdgcn_s_barrier();

    // hot loop
-    for (int j = min_tile + 3; j < max_num_tiles - 1; j += 2) {
+    for (int j = 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);
@@ -421,7 +379,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 (WINDOW || q_start_pos < kv_end_pos) {
+            if (q_start_pos < kv_end_pos) {  // Only mask if needed
                mask_kv_tile(att_block[0], tile_idx, j - 1, neg_inf_v, lane);
            }
        }
@@ -489,7 +447,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(WINDOW || q_start_pos < kv_end_pos, WINDOW ? 1 : 0)) {
+        if (__builtin_expect(q_start_pos < kv_end_pos, 0)) {  // Only mask if needed
            mask_kv_tile(att_block[1], tile_idx, max_num_tiles - 3, neg_inf_v, lane);
        }
    }
@@ -552,7 +510,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(WINDOW || q_start_pos < kv_end_pos, 1)) {
+        if (__builtin_expect(q_start_pos < kv_end_pos, 1)) {  // Only mask if needed
            mask_kv_tile(att_block[0], tile_idx, max_num_tiles - 2, neg_inf_v, lane);
        }
    }
@@ -614,7 +572,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(WINDOW || q_start_pos < kv_end_pos, 1)) {
+        if (__builtin_expect(q_start_pos < kv_end_pos, 1)) {  // Only mask if needed
            mask_kv_tile(att_block[1], tile_idx, max_num_tiles - 1, neg_inf_v, lane);
        }
    }
@@ -97,33 +97,4 @@ __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,20 +471,6 @@ 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).

@@ -11,8 +11,8 @@ from tinygrad.runtime.autogen.amd.rdna3.enum import VOP1Op, VOP2Op, SOP2Op, DSOp

 def _srcs():
  """Create minimal source variables for pcode parsing."""
-  def u32(v=0): return UOp.const(v, dtypes.uint32)
-  return {'S0': u32(), 'S1': u32(), 'S2': u32(), 'SCC': u32(), 'VCC': UOp.const(0, dtypes.uint64), 'laneId': u32()}
+  def u32(v=0): return UOp.const(dtypes.uint32, v)
+  return {'S0': u32(), 'S1': u32(), 'S2': u32(), 'SCC': u32(), 'VCC': UOp.const(dtypes.uint64, 0), 'laneId': u32()}

 class TestBasicParsing(unittest.TestCase):
  """Test basic pcode parsing for common instruction patterns."""
@@ -44,8 +44,8 @@ class TestWithSources(unittest.TestCase):

  def test_v_add_f32_with_sources(self):
    """Test V_ADD_F32 with actual float constants."""
-    s0 = UOp.const(0x3f800000, dtypes.uint32)  # 1.0f
-    s1 = UOp.const(0x40000000, dtypes.uint32)  # 2.0f
+    s0 = UOp.const(dtypes.uint32, 0x3f800000)  # 1.0f
+    s1 = UOp.const(dtypes.uint32, 0x40000000)  # 2.0f
    _, assigns = parse_pcode(PCODE[VOP2Op.V_ADD_F32_E32], {'S0': s0, 'S1': s1})
    self.assertEqual(len(assigns), 1)
    dest, val = assigns[0]
@@ -55,8 +55,8 @@ class TestWithSources(unittest.TestCase):

  def test_v_mul_f32_with_sources(self):
    """Test V_MUL_F32 with actual float constants."""
-    s0 = UOp.const(0x40000000, dtypes.uint32)  # 2.0f
-    s1 = UOp.const(0x40400000, dtypes.uint32)  # 3.0f
+    s0 = UOp.const(dtypes.uint32, 0x40000000)  # 2.0f
+    s1 = UOp.const(dtypes.uint32, 0x40400000)  # 3.0f
    _, assigns = parse_pcode(PCODE[VOP2Op.V_MUL_F32_E32], {'S0': s0, 'S1': s1})
    self.assertEqual(len(assigns), 1)
    dest, val = assigns[0]
@@ -67,36 +67,36 @@ class TestParseExpr(unittest.TestCase):

  def test_integer_literals(self):
    """Test parsing integer literals."""
-    self.assertEqual(parse_expr('0', {}).val, 0)
-    self.assertEqual(parse_expr('42', {}).val, 42)
-    self.assertEqual(parse_expr('42U', {}).val, 42)
+    self.assertEqual(parse_expr('0', {}).arg, 0)
+    self.assertEqual(parse_expr('42', {}).arg, 42)
+    self.assertEqual(parse_expr('42U', {}).arg, 42)

  def test_negative_integers(self):
    """Test parsing negative integer literals."""
    result = parse_expr('-1', {})
-    self.assertEqual(result.val, -1)
+    self.assertEqual(result.arg, -1)
    self.assertEqual(result.dtype, dtypes.int)

  def test_float_literals(self):
    """Test parsing float literals."""
    result = parse_expr('1.0F', {})
-    self.assertEqual(result.val, 1.0)
+    self.assertEqual(result.arg, 1.0)
    self.assertEqual(result.dtype, dtypes.float32)

  def test_hex_literals(self):
    """Test parsing hex literals."""
    result = parse_expr('0xFF', {})
-    self.assertEqual(result.val, 255)
+    self.assertEqual(result.arg, 255)

  def test_variable_lookup(self):
    """Test variable lookup in parse_expr."""
-    vrs = {'x': UOp.const(42, dtypes.uint32)}
+    vrs = {'x': UOp.const(dtypes.uint32, 42)}
    result = parse_expr('x', vrs)
-    self.assertEqual(result.val, 42)
+    self.assertEqual(result.arg, 42)

  def test_binary_ops(self):
    """Test parsing binary operations."""
-    vrs = {'a': UOp.const(10, dtypes.uint32), 'b': UOp.const(5, dtypes.uint32)}
+    vrs = {'a': UOp.const(dtypes.uint32, 10), 'b': UOp.const(dtypes.uint32, 5)}

    # Addition
    result = parse_expr('a + b', vrs)
@@ -105,11 +105,11 @@ class TestParseExpr(unittest.TestCase):
    # Subtraction with constant folding
    result = parse_expr('10 - 5', {})
    self.assertEqual(result.op, Ops.CONST)
-    self.assertEqual(result.val, 5)
+    self.assertEqual(result.arg, 5)

  def test_ternary(self):
    """Test parsing ternary expressions."""
-    vrs = {'cond': UOp.const(True), 'a': UOp.const(1, dtypes.uint32), 'b': UOp.const(0, dtypes.uint32)}
+    vrs = {'cond': UOp.const(dtypes.bool, True), 'a': UOp.const(dtypes.uint32, 1), 'b': UOp.const(dtypes.uint32, 0)}
    result = parse_expr('cond ? a : b', vrs)
    self.assertEqual(result.op, Ops.WHERE)

@@ -127,7 +127,7 @@ class TestForLoopParsing(unittest.TestCase):
  def test_clz_parsing(self):
    """Test CLZ pcode parsing produces correct structure."""
    pcode = PCODE[VOP1Op.V_CLZ_I32_U32_E32]
-    S0 = UOp.const(0xFFFFFFFF, dtypes.uint32)  # All ones - CLZ should be 0
+    S0 = UOp.const(dtypes.uint32, 0xFFFFFFFF)  # All ones - CLZ should be 0
    _vrs, assigns = parse_pcode(pcode, {'S0': S0})

    self.assertEqual(len(assigns), 1)
@@ -139,7 +139,7 @@ class TestForLoopParsing(unittest.TestCase):
  def test_clz_with_zero(self):
    """Test CLZ with input 0 - should return -1."""
    pcode = PCODE[VOP1Op.V_CLZ_I32_U32_E32]
-    S0 = UOp.const(0, dtypes.uint32)
+    S0 = UOp.const(dtypes.uint32, 0)
    _vrs, assigns = parse_pcode(pcode, {'S0': S0})

    # Check that the innermost value (default) is -1 (may be wrapped in CAST)
@@ -150,7 +150,7 @@ class TestForLoopParsing(unittest.TestCase):
    # Unwrap CAST if present
    while val.op == Ops.CAST:
      val = val.src[0]
-    self.assertEqual(val.val, -1)
+    self.assertEqual(val.arg, -1)

  def test_ctz_parsing(self):
    """Test CTZ pcode parsing."""
@@ -158,7 +158,7 @@ class TestForLoopParsing(unittest.TestCase):
    if pcode is None:
      self.skipTest("V_CTZ_I32_B32_E32 pcode not available")

-    S0 = UOp.const(1, dtypes.uint32)  # LSB set - CTZ should be 0
+    S0 = UOp.const(dtypes.uint32, 1)  # LSB set - CTZ should be 0
    _vrs, assigns = parse_pcode(pcode, {'S0': S0})
    self.assertEqual(len(assigns), 1)

@@ -169,8 +169,8 @@ class TestDSPcodePatterns(unittest.TestCase):
    """Test GLOBAL_ATOMIC_ADD_F32 keeps memory values in float dtype."""
    vmem = UOp.param(2, dtypes.uint32, (1024,))
    srcs = {
-      'ADDR': UOp.const(0, dtypes.uint64),
-      'DATA': UOp.const(0x3f800000, dtypes.uint32),
+      'ADDR': UOp.const(dtypes.uint64, 0),
+      'DATA': UOp.const(dtypes.uint32, 0x3f800000),
      '_vmem': vmem,
    }

@@ -199,8 +199,8 @@ class TestDSPcodePatterns(unittest.TestCase):
    """Test MEM[addr].type read expression parsing."""
    # Create a mock LDS buffer
    lds = UOp.param(3, dtypes.uint32, (16384,))
-    addr = UOp.const(0, dtypes.uint32)
-    vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(0, dtypes.uint32)}
+    addr = UOp.const(dtypes.uint32, 0)
+    vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(dtypes.uint32, 0)}

    result = parse_expr('MEM[ADDR + OFFSET].b32', vrs)
    # Should be an INDEX operation into LDS
@@ -212,13 +212,13 @@ class TestDSPcodePatterns(unittest.TestCase):
    self.assertIsNotNone(pcode)
    assert pcode is not None
    srcs = {
-      'ADDR': UOp.const(0, dtypes.uint32),
-      'OFFSET0': UOp.const(0, dtypes.uint32),
-      'OFFSET1': UOp.const(1, dtypes.uint32),
-      'DATA': UOp.const(0xAAAAAAAA, dtypes.uint32),
-      'DATA2': UOp.const(0xBBBBBBBB, dtypes.uint32),
+      'ADDR': UOp.const(dtypes.uint32, 0),
+      'OFFSET0': UOp.const(dtypes.uint32, 0),
+      'OFFSET1': UOp.const(dtypes.uint32, 1),
+      'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA),
+      'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB),
    }
-    srcs['laneId'] = UOp.const(0, dtypes.uint32)
+    srcs['laneId'] = UOp.const(dtypes.uint32, 0)
    _, assigns = parse_pcode(pcode, srcs)
    # Should have 2 MEM write assignments
    self.assertEqual(len(assigns), 2)
@@ -235,12 +235,12 @@ class TestDSPcodePatterns(unittest.TestCase):
    assert pcode is not None
    lds = UOp.param(3, dtypes.uint32, (16384,))
    srcs = {
-      'ADDR': UOp.const(0, dtypes.uint32),
-      'OFFSET0': UOp.const(0, dtypes.uint32),
-      'OFFSET1': UOp.const(1, dtypes.uint32),
+      'ADDR': UOp.const(dtypes.uint32, 0),
+      'OFFSET0': UOp.const(dtypes.uint32, 0),
+      'OFFSET1': UOp.const(dtypes.uint32, 1),
      '_lds': lds,
    }
-    srcs['laneId'] = UOp.const(0, dtypes.uint32)
+    srcs['laneId'] = UOp.const(dtypes.uint32, 0)
    _, assigns = parse_pcode(pcode, srcs)
    # Should have 2 RETURN_DATA assignments
    self.assertEqual(len(assigns), 2)
@@ -252,36 +252,36 @@ class TestDSPcodePatterns(unittest.TestCase):
    pcode = PCODE.get(DSOp.DS_STORE_2ADDR_B32)
    assert pcode is not None
    srcs = {
-      'ADDR': UOp.const(100, dtypes.uint32),
-      'OFFSET0': UOp.const(2, dtypes.uint32),
-      'OFFSET1': UOp.const(5, dtypes.uint32),
-      'DATA': UOp.const(0xAAAAAAAA, dtypes.uint32),
-      'DATA2': UOp.const(0xBBBBBBBB, dtypes.uint32),
+      'ADDR': UOp.const(dtypes.uint32, 100),
+      'OFFSET0': UOp.const(dtypes.uint32, 2),
+      'OFFSET1': UOp.const(dtypes.uint32, 5),
+      'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA),
+      'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB),
    }
-    srcs['laneId'] = UOp.const(0, dtypes.uint32)
+    srcs['laneId'] = UOp.const(dtypes.uint32, 0)
    _, assigns = parse_pcode(pcode, srcs)
    # Check addresses: 100 + 2*4 = 108, 100 + 5*4 = 120
    # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp
-    self.assertEqual(assigns[0][1][0].simplify().val, 108)  # type: ignore[index]
-    self.assertEqual(assigns[1][1][0].simplify().val, 120)  # type: ignore[index]
+    self.assertEqual(assigns[0][1][0].simplify().arg, 108)  # type: ignore[index]
+    self.assertEqual(assigns[1][1][0].simplify().arg, 120)  # type: ignore[index]

  def test_ds_store_data_values(self):
    """Test DS_STORE_2ADDR_B32 uses correct data values."""
    pcode = PCODE.get(DSOp.DS_STORE_2ADDR_B32)
    assert pcode is not None
    srcs = {
-      'ADDR': UOp.const(0, dtypes.uint32),
-      'OFFSET0': UOp.const(0, dtypes.uint32),
-      'OFFSET1': UOp.const(1, dtypes.uint32),
-      'DATA': UOp.const(0xAAAAAAAA, dtypes.uint32),
-      'DATA2': UOp.const(0xBBBBBBBB, dtypes.uint32),
+      'ADDR': UOp.const(dtypes.uint32, 0),
+      'OFFSET0': UOp.const(dtypes.uint32, 0),
+      'OFFSET1': UOp.const(dtypes.uint32, 1),
+      'DATA': UOp.const(dtypes.uint32, 0xAAAAAAAA),
+      'DATA2': UOp.const(dtypes.uint32, 0xBBBBBBBB),
    }
-    srcs['laneId'] = UOp.const(0, dtypes.uint32)
+    srcs['laneId'] = UOp.const(dtypes.uint32, 0)
    _, assigns = parse_pcode(pcode, srcs)
    # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp
    # DATA[31:0] should preserve the value
-    self.assertEqual(assigns[0][1][1].simplify().val, 0xAAAAAAAA)  # type: ignore[index]
-    self.assertEqual(assigns[1][1][1].simplify().val, 0xBBBBBBBB)  # type: ignore[index]
+    self.assertEqual(assigns[0][1][1].simplify().arg, 0xAAAAAAAA)  # type: ignore[index]
+    self.assertEqual(assigns[1][1][1].simplify().arg, 0xBBBBBBBB)  # type: ignore[index]

 class TestConditionalParsing(unittest.TestCase):
  """Test conditional (if/elsif/else) pcode parsing."""
@@ -290,9 +290,9 @@ class TestConditionalParsing(unittest.TestCase):
    """Test parsing ternary expression (which becomes WHERE)."""
    # S_CSELECT_B32: D0.u32 = SCC ? S0.u32 : S1.u32
    pcode = PCODE[SOP2Op.S_CSELECT_B32]
-    s0 = UOp.const(10, dtypes.uint32)
-    s1 = UOp.const(20, dtypes.uint32)
-    scc = UOp.const(1, dtypes.uint32)
+    s0 = UOp.const(dtypes.uint32, 10)
+    s1 = UOp.const(dtypes.uint32, 20)
+    scc = UOp.const(dtypes.uint32, 1)
    _vrs, assigns = parse_pcode(pcode, {'S0': s0, 'S1': s1, 'SCC': scc})
    self.assertEqual(len(assigns), 1)
    dest, val = assigns[0]
@@ -305,26 +305,26 @@ class TestConcatWidthParsing(unittest.TestCase):

  def test_permlanex16_altrow_concat(self):
    for row, expected in [(0, 1), (1, 0), (2, 3), (3, 2)]:
-      parsed = parse_expr('{ row[1], ~row[0] }', {'row': UOp.const(row, dtypes.uint32)})
-      self.assertEqual(parsed.simplify().val, expected)
+      parsed = parse_expr('{ row[1], ~row[0] }', {'row': UOp.const(dtypes.uint32, row)})
+      self.assertEqual(parsed.simplify().arg, expected)

  def test_permlane64_altlane_concat(self):
    for lane, expected in [(0, 32), (1, 33), (31, 63), (32, 0), (63, 31)]:
-      parsed = parse_expr('{ ~lane[5], lane[4:0] }', {'lane': UOp.const(lane, dtypes.uint32)})
-      self.assertEqual(parsed.simplify().val, expected)
+      parsed = parse_expr('{ ~lane[5], lane[4:0] }', {'lane': UOp.const(dtypes.uint32, lane)})
+      self.assertEqual(parsed.simplify().arg, expected)

  def test_permlane64_wave64_pcode_indices(self):
    vgpr = UOp.param(0, dtypes.uint32, (256,))
    srcs = {
-      'SRC0': UOp.const(0, dtypes.uint32),
-      'VDST': UOp.const(1, dtypes.uint32),
-      'EXEC_LO': UOp.const(0xFFFFFFFF, dtypes.uint32),
-      'EXEC': UOp.const(0xFFFFFFFFFFFFFFFF, dtypes.uint64),
+      'SRC0': UOp.const(dtypes.uint32, 0),
+      'VDST': UOp.const(dtypes.uint32, 1),
+      'EXEC_LO': UOp.const(dtypes.uint32, 0xFFFFFFFF),
+      'EXEC': UOp.const(dtypes.uint64, 0xFFFFFFFFFFFFFFFF),
      '_vgpr': vgpr,
      '_wave_size': 64,
-      'S0': UOp.const(0, dtypes.uint32),
-      'S1': UOp.const(0, dtypes.uint32),
-      'S2': UOp.const(0, dtypes.uint32),
+      'S0': UOp.const(dtypes.uint32, 0),
+      'S1': UOp.const(dtypes.uint32, 0),
+      'S2': UOp.const(dtypes.uint32, 0),
    }

    def load_idx(v: UOp) -> int:
@@ -333,12 +333,12 @@ class TestConcatWidthParsing(unittest.TestCase):
      self.assertEqual(simp.src[0].op, Ops.INDEX)
      idx = simp.src[0].src[1].simplify()
      self.assertEqual(idx.op, Ops.CONST)
-      return idx.val
+      return idx.arg

    _, assigns = parse_pcode(PCODE[VOP1Op.V_PERMLANE64_B32_E32], srcs)
    self.assertEqual(len(assigns), 64)
    for lane, (dst_idx, src_idx) in {0: (64, 32), 31: (95, 63), 32: (96, 0), 63: (127, 31)}.items():
-      self.assertEqual(assigns[lane][1][0].simplify().val, dst_idx)  # type: ignore[index]
+      self.assertEqual(assigns[lane][1][0].simplify().arg, dst_idx)  # type: ignore[index]
      self.assertEqual(load_idx(assigns[lane][1][1]), src_idx)  # type: ignore[index]

 class TestAllPcode(unittest.TestCase):
@@ -346,7 +346,7 @@ class TestAllPcode(unittest.TestCase):

  def _make_srcs(self):
    """Create dummy source variables for pcode parsing."""
-    u32, u64 = lambda v=0: UOp.const(v, dtypes.uint32), lambda v=0: UOp.const(v, dtypes.uint64)
+    u32, u64 = lambda v=0: UOp.const(dtypes.uint32, v), lambda v=0: UOp.const(dtypes.uint64, v)
    lds = UOp.param(3, dtypes.uint32, (16384,))
    return {'laneId': u32(), 'laneID': u32(), 'S0': u32(), 'S1': u32(), 'S2': u32(), 'S3': u32(), 'SRC0': u32(),
            'D0': u32(), 'D1': u32(), 'DST': u32(), 'VDST': u32(), 'SDST': u32(),
@@ -358,7 +358,7 @@ class TestAllPcode(unittest.TestCase):
            'M0': u32(), 'PC': u64(), 'DENORM': u32(1), 'ROUND_MODE': u32(), 'ROUND_TOWARD_ZERO': u32(),
            'ROUND_NEAREST_EVEN': u32(), 'WAVE_STATUS': u32(),
            'MAX_FLOAT_F32': u32(0x7f7fffff), 'Unsigned': u32(1), 'clampedLOD': u32(),
-            '_lds': lds, '_vmem': lds, '_active': UOp.const(True)}
+            '_lds': lds, '_vmem': lds, '_active': UOp.const(dtypes.bool, True)}

  def _parse_all_pcode(self, pcode_dict, arch: str, min_pct: float):
    """Parse all pcode. RuntimeError = parser limitation (ok), other exceptions = real bugs."""
@@ -4,13 +4,13 @@ from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
 from tinygrad.helpers import Context, getenv, DEV
 from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear
 from tinygrad.renderer.ptx import PTXRenderer
-from test.helpers import needs_second_gpu, check_schedule, assert_kernel_count, KernelCountException
+from test.helpers import needs_second_gpu

 class TestArange(unittest.TestCase):
  def _get_flops(self, tensor, desired):
    GlobalCounters.reset()
    linear = compile_linear(tensor.schedule_linear())
-    if len(linear.src) != 1: raise KernelCountException(1, len(linear.src))
+    self.assertEqual(len(linear.src), 1)
    run_linear(linear)
    np.testing.assert_equal(tensor.numpy(), desired)
    return estimate_uop(linear.src[-1]).ops
@@ -55,7 +55,8 @@ class TestIndexing(unittest.TestCase):
    with Context(NOOPT=1):
      GlobalCounters.reset()
      out = ((Tensor.arange(1,16385)-1)*needle).sum()
-      linear, var_vals = check_schedule(out, 1)
+      linear, var_vals = out.linear_with_vars()
+      self.assertEqual(len(linear.src), 1)
      run_linear(linear, var_vals)
    self.assertEqual(out.item(), 1337)

@@ -71,7 +72,8 @@ class TestIndexing(unittest.TestCase):
      reshape_dataset = dataset.T.reshape(1, DDIM, DSET, 1).expand(4, DDIM, DSET, 1)
      full = (rng==idxs).where(reshape_dataset, Tensor.zeros(4, DDIM, DSET, 1, buffer=False))
      X = full.sum(axis=(2,3))
-      linear, var_vals = check_schedule(X, 1)
+      linear, var_vals = X.linear_with_vars()
+      self.assertEqual(len(linear.src), 1)
      run_linear(linear, var_vals)
      assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops}"
    np.testing.assert_allclose(real_index, X.numpy())
@@ -96,7 +98,8 @@ class TestIndexing(unittest.TestCase):
      GlobalCounters.reset()
      X = dataset[idxs]
      assert X.shape == (4,DDIM)
-      linear, var_vals = check_schedule(X, 1)
+      linear, var_vals = X.linear_with_vars()
+      self.assertEqual(len(linear.src), 1)
      run_linear(linear, var_vals)
      assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops}"
    np.testing.assert_allclose(real_index, X.numpy())
@@ -110,7 +113,8 @@ class TestIndexing(unittest.TestCase):
      GlobalCounters.reset()
      X = dataset[idxs]
      assert X.shape == (4,DDIM)
-      linear, var_vals = check_schedule(X, 1)
+      linear, var_vals = X.linear_with_vars()
+      self.assertEqual(len(linear.src), 1)
      run_linear(linear, var_vals)
      assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops} != {4*DSET}"
    np.testing.assert_allclose(real_index, X.numpy())
@@ -153,7 +157,7 @@ class TestIndexing(unittest.TestCase):
      GlobalCounters.reset()
      z = emb(x).realize()
      self.assertLessEqual(GlobalCounters.global_ops, op_limit)
-      assert_kernel_count(2)
+      self.assertEqual(GlobalCounters.kernel_count, 2)
    if getenv("CHECK", 1):
      import torch
      with torch.no_grad():
@@ -253,7 +257,7 @@ class TestIndexing(unittest.TestCase):
    xq_rope, _ = apply_rotary_emb(xq, xq, freqs_cis)
    xq_rope.sum().backward()
    linear = compile_linear(wq.grad.schedule_linear())
-    if len(linear.src) != 1: raise KernelCountException(1, len(linear.src))
+    assert len(linear.src) == 1, f"expected one kernel for backward, got: {len(linear.src)}"
    bwd_ops = estimate_uop(linear.src[0]).ops
    expected_ops = bs*seqlen*dim*dim*ops_scale
    print(f"rope matmul bwd ({dtype}): {GlobalCounters.kernel_count} kernels, {bwd_ops:,} ops")
@@ -150,44 +150,6 @@ class TestAsmGEMM(unittest.TestCase):
    with self.assertRaisesRegex(AssertionError, "not a multiple"):
      verify_asm_gemm(1, 256, 1000, 256)

-class TestMXFP4(unittest.TestCase):
-  def setUp(self):
-    if not is_cdna4() or DEV.interface.startswith("MOCK"):
-      self.skipTest("requires real amd machine")
-
-  def test_quantize(self):
-    import numpy as np
-    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
-    M = N = K = 256
-    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().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)
-    a = Tensor.empty(M, K, dtype=dtypes.bfloat16)
-    b = Tensor.empty(N, K, dtype=dtypes.bfloat16)
-    asm_gemm(a, b.T, mxfp4=True).realize()
-
 # test the Asm GEMM with Llama shapes, only run on the real machine for speed

@unittest.skipUnless(has_hipcc(), "requires hipcc to compile")
@@ -6,11 +6,11 @@ from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.uop.ops import KernelInfo

 def call_out_kernel(F:UOp, C:UOp) -> UOp:
-  call = F[0].load().call(UOp.const(3).cast(dtypes.int), C[0], ret_dtype=dtypes.void)
+  call = F[0].load().call(UOp.const(dtypes.int, 3), C[0], ret_dtype=dtypes.void)
  return C.after(call)[1].store(C.after(call)[0].load() + 1).sink(arg=KernelInfo(name="call_out"))

 def call_ret_kernel(F:UOp, C:UOp) -> UOp:
-  val = F[0].load().call(UOp.const(21).cast(dtypes.int), ret_dtype=dtypes.int)
+  val = F[0].load().call(UOp.const(dtypes.int, 21), ret_dtype=dtypes.int)
  return C[0].store(val * 2).sink(arg=KernelInfo(name="call_ret"))

@unittest.skipUnless(isinstance(Device["CPU"].renderer, CStyleLanguage), "TODO: CALL is rendered in C style only")
@@ -16,7 +16,7 @@ def _check_ast_count(desired_count:int, t:Tensor):

 class TestMovedConstFolding(unittest.TestCase):
  def test_contiguous_deviceless_const(self):
-    t = Tensor(UOp.const(2.0, dtypes.float)).contiguous()
+    t = Tensor(UOp.const(dtypes.float, 2.0)).contiguous()
    self.assertIs(t.uop.op, Ops.CONST)
    self.assertIsNone(t.uop.device)

@@ -169,7 +169,7 @@ class TestMultiConstFolding(unittest.TestCase):
 class TestThreefryConstFolding(unittest.TestCase):
  def test_threefry(self):
    # THREEFRY(const,const) folds to a const once decomposed
-    x = threefry2x32(UOp.const(5, dtypes.uint64), UOp.const(10, dtypes.uint64))
+    x = threefry2x32(UOp.const(dtypes.uint64, 5), UOp.const(dtypes.uint64, 10))
    self.assertIs(x.simplify().op, Ops.CONST)

 class TestTautologicalCompare(unittest.TestCase):
@@ -1,10 +1,8 @@
 import unittest
 from tinygrad import Tensor, UOp, GlobalCounters, Context, Device
-import numpy as np
 from tinygrad.dtype import AddrSpace, dtypes, Invalid
 from tinygrad.uop.ops import KernelInfo, AxisType, Ops
 from tinygrad.renderer.ptx import PTXRenderer
-from test.helpers import assert_kernel_count

 # **** kernels ****

@@ -190,12 +188,6 @@ 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)
@@ -283,7 +275,7 @@ class TestCustomKernel(unittest.TestCase):

    GlobalCounters.reset()
    out.realize()
-    assert_kernel_count(5)
+    self.assertEqual(GlobalCounters.kernel_count, 5)

  def test_simple_reshape(self):
    a = Tensor.ones(2,3,4).realize()
@@ -293,7 +285,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(2)
+    self.assertEqual(GlobalCounters.kernel_count, 3)

  def test_multi_after_schedule_order(self):
    """Test correct scheduling order when custom_kernel has multiple outputs.
@@ -343,12 +335,12 @@ class TestCustomKernel(unittest.TestCase):
    c = Tensor.custom_kernel(c, a, fxn=custom_add_one_kernel)[0]
    GlobalCounters.reset()
    c.realize()
-    assert_kernel_count(len(devs))
+    self.assertEqual(GlobalCounters.kernel_count, len(devs))
    self.assertTrue((c == 2).all().item())

  def test_partial_invalid_store_keeps_uncovered_reads(self):
    x = Tensor([10., 20., 30., 40.])
-    after = x.uop.after(x.uop.shrink(((0, 2),)).store(Invalid))
+    after = x.uop.after(x.uop.shrink(((0, 2),)).store(UOp.const(dtypes.float, Invalid, shape=(2,))))
    self.assertEqual(Tensor(after).contiguous().tolist(), [10., 20., 30., 40.])

  def test_multi_after_invalid_store_dep_removed(self):
@@ -408,11 +400,13 @@ class TestCustomKernel(unittest.TestCase):
    else: z = y.T.T+1
    GlobalCounters.reset()
    z.realize()
-    assert_kernel_count(2)
+    self.assertEqual(GlobalCounters.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
@@ -422,8 +416,8 @@ class TestCustomKernel(unittest.TestCase):
      return Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
    GlobalCounters.reset()
    y = run(x[0]).realize()
-    # backends that support contiguous views don't launch extra kernels
-    assert_kernel_count(2 if x[0].uop.contiguous_view() is None else 1)
+    # it's copying the input and the output
+    self.assertEqual(GlobalCounters.kernel_count, 1)
    self.assertEqual(y.tolist(), [1, 2, 3, 4])

  @Context(DEV="CPU")
@@ -433,37 +427,11 @@ 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, B:UOp) -> UOp:
+    def custom_src_kernel(A: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).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)
-      col = UOp.range(A.shape[1], 1)
-      return A[row, col].store(A[row].max(axis=0)).end(col).end(row).sink(arg=KernelInfo(name=f"assign_row_max_{A.numel()}"))
-    a = Tensor.arange(4).clone().realize()
-    a = Tensor.custom_kernel(a.reshape(2, 2).T, fxn=custom_assign_row_max_kernel)[0]
-    self.assertEqual(a.flatten().tolist(), [2, 2, 3, 3])
-    self.assertEqual(a.shape, (2, 2))
+    a = Tensor.custom_kernel(a.reshape(2, 2).T, fxn=custom_src_kernel)[0]
+    self.assertEqual(a.tolist(), [[1, 2], [1, 3]])

 class TestCustomKernelInput(unittest.TestCase):
  def _test_mop(self, mop_fxn, max_kernels):
@@ -491,8 +459,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=2)
-  def test_shrink(self): self._test_mop(lambda x: x[:4], max_kernels=1)
+  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_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)
@@ -559,77 +527,6 @@ class TestUnshardIndex(unittest.TestCase):
    with self.assertRaisesRegex(RuntimeError, "cannot shard index"):
      self._run(kernel, (64, 8))

-def _run_fragment_kernel(testcase, kernel, out_shape, inputs=()):
-  c = Tensor.empty(*out_shape)
-  out = Tensor.custom_kernel(c, *inputs, fxn=kernel)[0]
-  try: return out.numpy()
-  except RuntimeError as e:
-    if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and "dynamic register indexing" in str(e):
-      testcase.skipTest("PTX does not support dynamic register indexing")
-    raise
-
-class TestUnshardAlu(unittest.TestCase):
-  """Tests for ALU on (fragment) UNSHARD values in schedule/multi.py's alu_multi.
-
-  An ALU with UNSHARD srcs lowers to per-shard ops when every src is one of:
-    same sharding:  peel the UNSHARD, keep the layout
-    scalar:         broadcast to every shard
-    whole unsharded same-shape value: takes its per-shard sub-view (shard_subview)
-  """
-  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
-  def test_alu_scalar_broadcast(self):
-    # scalar srcs broadcast to every shard: frag*2.0 where frag is 1.5 per thread -> 3.0 everywhere
-    def kernel(C:UOp) -> UOp:
-      ty = UOp.range(8, 0, AxisType.LOCAL)
-      # 8 values per thread, 8 threads -> 64-value full view
-      frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,))
-      v = frag.after(frag.store(1.5)) * 2.0
-      return C.store(v).end(ty).sink(arg=KernelInfo(name="alu_scalar", opts_to_apply=()))
-    out = _run_fragment_kernel(self, kernel, (64,))
-    np.testing.assert_allclose(out, 3.0)
-
-  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
-  def test_alu_whole_value_subview(self):
-    # UNSHARD + whole unsharded same-shape value: each shard adds its own sub-view of A.
-    def kernel(C:UOp, A:UOp) -> UOp:
-      ty = UOp.range(8, 0, AxisType.LOCAL)
-      frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,))
-      v = frag.after(frag.store(0.0)) + A
-      return C.store(v).end(ty).sink(arg=KernelInfo(name="alu_subview", opts_to_apply=()))
-    a = Tensor(np.arange(64, dtype=np.float32))
-    out = _run_fragment_kernel(self, kernel, (64,), inputs=(a,))
-    np.testing.assert_allclose(out, a.numpy(), atol=1e-4)
-
-class TestUnshardStore(unittest.TestCase):
-  """Tests for STORE of a sharded value into an unsharded dest (store_value_multi in schedule/multi.py).
-
-  Every shard stores its value into its own contiguous sub-view of the dest, one SHRINK per sharded axis.
-  """
-  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
-  def test_store_unshard_value(self):
-    # single-axis: 8 threads each own 8 values of the 64-value output tile
-    def kernel(C:UOp) -> UOp:
-      ty = UOp.range(8, 0, AxisType.LOCAL)
-      frag = UOp.placeholder((8,), dtypes.float32, 0, AddrSpace.LOCAL).unshard((0,), (ty,))
-      v = frag.after(frag.store(0.0)) + 2.5
-      return C.store(v).end(ty).sink(arg=KernelInfo(name="store_unshard", opts_to_apply=()))
-    out = _run_fragment_kernel(self, kernel, (64,))
-    np.testing.assert_allclose(out, 2.5)
-
-  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
-  def test_store_unshard_value_2axis(self):
-    # two sharded axes (the gemm fragment layout): thread (ty, tx) owns the (2, 1, 1, 2) sub-view of the
-    # (2, 4, 2, 2) output tile; the store must SHRINK dest on both sharded axes
-    def kernel(C:UOp, A:UOp) -> UOp:
-      ty = UOp.range(4, 0, AxisType.LOCAL)
-      tx = UOp.range(2, 1, AxisType.LOCAL)
-      frag = UOp.placeholder((2, 1, 1, 2), dtypes.float32, 0, AddrSpace.REG).unshard((1, 2), (ty, tx))
-      v = frag.after(frag.store(0.0)) + A
-      return C.store(v).end(tx, ty).sink(arg=KernelInfo(name="store_unshard_2axis", opts_to_apply=()))
-    a = Tensor(np.arange(32, dtype=np.float32).reshape(2, 4, 2, 2))
-    out = _run_fragment_kernel(self, kernel, (2, 4, 2, 2), inputs=(a,))
-    np.testing.assert_allclose(out, a.numpy(), atol=1e-4)
-
 class TestUOpReduce(unittest.TestCase):
  def test_uop_sum(self):
    a = Tensor([1.0, 2, 3, 4, 5])
@@ -169,13 +169,6 @@ 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,10 +399,9 @@ class TestDTypeALU(unittest.TestCase):
    if float_dtype not in supported_dtypes: float_dtype = dtypes.float32
    universal_test_cast(a, float_dtype, unsigned_dtype)

-  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
+  @unittest.expectedFailure
+  def test_unsafe_cast_float_to_int_failure(self):
+    val = float(dtypes.int32.max - 1)
    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())
@@ -7,7 +7,7 @@ from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops
 from tinygrad.renderer.isa import IselContext

 # INDEX on a register value with a constant index extracts a single element (the old GEP)
-def lane(y:UOp, i:int) -> UOp: return y.index(UOp.const(i, dtypes.int), dtype=y.dtype.scalar())
+def lane(y:UOp, i:int) -> UOp: return y.index(UOp.const(dtypes.int, i), dtype=y.dtype.scalar())

@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only x86")
 class TestIselX86(unittest.TestCase):
@@ -49,7 +49,7 @@ class TestIselX86(unittest.TestCase):
    load = UOp.param(0, dtypes.int32, (16,)).index(a + 1).load()
    n = self.isel_rewrite(load)
    # displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32
-    self.assertTrue(n.src[2].op is Ops.CONST and n.src[2].dtype is dtypes.int8 and n.src[2].val == 4)
+    self.assertTrue(n.src[2].op is Ops.CONST and n.src[2].dtype is dtypes.int8 and n.src[2].arg == 4)

 if __name__ == "__main__":
  unittest.main()
@@ -2,7 +2,7 @@
 import unittest
 import numpy as np

-from test.helpers import assert_jit_cache_len, call_is_graph, not_support_multi_device, needs_second_gpu, KernelCountException
+from test.helpers import assert_jit_cache_len, call_is_graph, not_support_multi_device, needs_second_gpu
 from test.unit.test_jit import _simple_test
 from tinygrad import Tensor, Variable, TinyJit, Device, dtypes
 from tinygrad.engine.jit import graph_class
@@ -97,7 +97,7 @@ class TestJit(unittest.TestCase):
      prev = o

    # Checking that 2 graphs are inited.
-    if len(jf.captured.linear.src) != 2: raise KernelCountException(2, len(jf.captured.linear.src))
+    assert len(jf.captured.linear.src) == 2
    for si in jf.captured.linear.src:
      assert call_is_graph(si)

@@ -360,7 +360,7 @@ class TestJitGraphSplit(unittest.TestCase):
    self.expect(f, inp, inp_cpu,
      graph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()],
      multigraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()],
-      hcqgraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()]) # cpu is hcq2 now, it does not join hcq graphs
+      hcqgraph=[self.ji_graph(4)])

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

  def test_jit_multidev(self):
    if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
@@ -12,7 +12,7 @@ from tinygrad.dtype import DType, dtypes, AddrSpace
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.renderer.isa import ISARenderer
-from test.helpers import replace_opts, check_schedule
+from test.helpers import replace_opts
 from test.backend.test_softmax_fusion import single_kernel_softmax
 MOCKGPU = DEV.interface.startswith("MOCK")

@@ -30,7 +30,7 @@ class TestLinearizer(unittest.TestCase):
    c = ((a.shrink(((0, 2),)) - a.shrink(((2, 4),))) - (b.shrink(((0, 2),)) - b.shrink(((2, 4),))))
    linear = c.schedule_linear()
    run_linear(linear)
-    rawbufs = [s.buffer for s in linear.src[-1].src[1:] if not s.is_bound_var]
+    rawbufs = [s.buffer for s in linear.src[-1].src[1:] if s.op is not Ops.BIND]
    assert len(rawbufs) == 3 and set(rawbufs[1:]) == {a.uop.base.realized, b.uop.base.realized}
    np_c = (np_a[:2] - np_a[2:]) - (np_b[:2] - np_b[2:])
    np.testing.assert_allclose(np_c, c.numpy(), atol=1e-4, rtol=1e-4)
@@ -268,9 +268,9 @@ class TestLinearizer(unittest.TestCase):
    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
    idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
    idxs = sorted(idxs, key=lambda uop: uop.arg)
-    assert (idxs[0].arg, idxs[0].src[0].val) == ('gidx0', 6), idxs[0]
-    assert (idxs[1].arg, idxs[1].src[0].val) == ('gidx1', 5), idxs[1].arg
-    assert (idxs[2].arg, idxs[2].src[0].val) == ('gidx2', 4), idxs[2].arg
+    assert (idxs[0].arg, idxs[0].src[0].arg) == ('gidx0', 6), idxs[0]
+    assert (idxs[1].arg, idxs[1].src[0].arg) == ('gidx1', 5), idxs[1].arg
+    assert (idxs[2].arg, idxs[2].src[0].arg) == ('gidx2', 4), idxs[2].arg

  def test_sum_collapse(self):
    t = Tensor([2]).reshape(1, 1).expand(256, 256).sum()
@@ -293,7 +293,8 @@ class TestLinearizer(unittest.TestCase):
    a = Tensor.ones(4, 4).contiguous().realize()
    b = a.shrink(((1, 2), None)).pad(((1, 2), None)).bool()
    a.assign(b.where(2, a))
-    linear, var_vals = check_schedule(a, 1)
+    linear, var_vals = a.linear_with_vars()
+    assert len(linear.src) == 1
    run_linear(linear, var_vals)
    np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.])
    program = to_program(replace_opts(linear.src[-1].src[0], []), renderer=Device[Device.DEFAULT].renderer)
@@ -411,7 +412,7 @@ def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
  last_call = linear.src[-1]
  ast = last_call.src[0]
  assert ast.op is Ops.SINK, f"helper_realized_ast expects a SINK {last_call}"
-  last_bufs = [s.buffer for s in last_call.src[1:] if not s.is_bound_var]
+  last_bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]
  # now all input buffers in last_call should be realized
  # create fresh buffers for the outputs
  bufs = [Buffer(x.device, x.size, x.dtype).allocate() if i < len(ast.src) else x for i,x in enumerate(last_bufs)]
@@ -437,7 +438,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=[], check_default_opt=True):
+                               apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[]):
  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,7 +460,9 @@ 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.
-  if check_default_opt: check_opt(None)
+  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)
  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)

@@ -12,18 +12,18 @@ class TestLinearizerFailure(unittest.TestCase):
  @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
  def test_failure_beam_mnist(self):
    c0 = UOp.param(0, dtypes.uchar, (4014080,))
-    c1 = UOp.range(UOp.const(512), 0, AxisType.GLOBAL)
-    c2 = UOp.range(UOp.const(784), 1, AxisType.GLOBAL)
-    c3 = UOp.range(UOp.const(10), 3, AxisType.GLOBAL)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 0, AxisType.GLOBAL)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 784), 1, AxisType.GLOBAL)
+    c3 = UOp.range(UOp.const(dtypes.weakint, 10), 3, AxisType.GLOBAL)
    c4 = UOp.param(1, dtypes.int, (512,))
-    c5 = c4.index(c1.valid(UOp.const(True)))
-    c6 = UOp.range(UOp.const(6000), 1004, AxisType.REDUCE)
-    c7 = UOp.range(UOp.const(3750), 2006, AxisType.REDUCE)
-    c8 = UOp.range(UOp.const(16), 2007, AxisType.GROUP_REDUCE)
+    c5 = c4.index(c1.valid(UOp.const(dtypes.bool, True)))
+    c6 = UOp.range(UOp.const(dtypes.weakint, 6000), 1004, AxisType.REDUCE)
+    c7 = UOp.range(UOp.const(dtypes.weakint, 3750), 2006, AxisType.REDUCE)
+    c8 = UOp.range(UOp.const(dtypes.weakint, 16), 2007, AxisType.GROUP_REDUCE)
    c9 = UOp.param(2, dtypes.uchar, (47040000,))
-    c10 = c9.index((((c3*UOp.const(4704000))+c2)+(c6*UOp.const(784))).valid(UOp.const(True)))
-    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(6000))+c6)+((c7*UOp.const(16))+c8)).alu(Ops.CMPLT, UOp.const(59999)).where(UOp.const(0).cast(dtypes.int), UOp.const(1).cast(dtypes.int)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(-1).cast(dtypes.int))).where(UOp.const(0).cast(dtypes.uchar), c10).reduce(c6, arg=Ops.ADD)
-    c12 = c0.index((((c1*UOp.const(7840))+(c2*UOp.const(10)))+c3).valid(UOp.const(True))).store(c11).end(c1, c2, c3)
+    c10 = c9.index((((c3*UOp.const(dtypes.weakint, 4704000))+c2)+(c6*UOp.const(dtypes.weakint, 784))).valid(UOp.const(dtypes.bool, True)))
+    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.weakint, 6000))+c6)+((c7*UOp.const(dtypes.weakint, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.weakint, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
+    c12 = c0.index((((c1*UOp.const(dtypes.weakint, 7840))+(c2*UOp.const(dtypes.weakint, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11).end(c1, c2, c3)
    ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None))
    _ = to_program(ast, Device["METAL"].renderer)

@@ -5,10 +5,9 @@ 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
+from test.helpers import needs_second_gpu
 from test.backend.test_asm_gemm import has_hipcc

 def run_fused_ce(bs:int, seqlen:int, vocab:int, label_smoothing:float=0.0) -> None:
@@ -96,7 +95,7 @@ class TestLocalAmax(unittest.TestCase):
    x = Tensor.arange(16).reshape(4, 4).cast(dtypes.float).clone(devices[0]).realize().shard(devices, axis=0).realize()
    GlobalCounters.reset()
    out = (x * local_abs_max(x)).clone().realize()
-    assert_kernel_count(2)
+    self.assertEqual(GlobalCounters.kernel_count, 2)
    self.assertEqual(out.tolist(), [[0., 7., 14., 21.], [28., 35., 42., 49.], [120., 135., 150., 165.], [180., 195., 210., 225.]])

@unittest.skipUnless(has_hipcc() and Device.DEFAULT == "AMD", "requires hipcc to compile and amd device to run")
@@ -162,31 +161,5 @@ class TestFusedQKVRoPE(unittest.TestCase):
    ref = Tensor.cat(dq_ref, dk_ref, dv_ref, dim=3).reshape(*dx.shape).realize()
    with Context(DEBUG=0): self.assertTrue(dx.allclose(ref, atol=2e-2, rtol=2e-2).item(), "backward mismatch")

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

 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
@@ -62,7 +62,7 @@ class TestMultiTensor(unittest.TestCase):
  def test_shard_empty(self):
    GlobalCounters.reset()
    X = Tensor.empty(256).shard(devices_2, 0).realize()
-    assert_kernel_count(0)
+    assert GlobalCounters.kernel_count == 0
    (X + X).realize()

  # TODO: fix this to not copy on the src device
@@ -76,13 +76,6 @@ class TestMultiTensor(unittest.TestCase):
    run_linear(linear)
    self.assertEqual(len(set(names)), 1, "function was relinearized")

-  def test_shard_beam(self):
-    cpu_2 = ("CPU:1", "CPU:2")
-    src = Tensor.ones(16).shard(cpu_2, 0).realize()
-    lin = UOp(Ops.LINEAR, src=(src.to(cpu_2[::-1]).schedule_linear().src[0],))
-    with Context(BEAM=1, IGNORE_BEAM_CACHE=1): call = graph_rewrite(graph_rewrite(lin, pm_beam, ctx=1, walk=True), pm_compile, walk=True).src[0]
-    self.assertNotEqual(call.src[0].src[0].arg.applied_opts, ())
-
  def test_shard_same_device(self):
    X = Tensor.ones(256).contiguous().realize()
    X.shard_((d1, X.device), 0)
@@ -355,7 +348,8 @@ class TestMultiTensor(unittest.TestCase):
  def test_const_like_shrink_on_shard_axis(self):
    t = Tensor.ones(16, 16, dtype=dtypes.int).shard(devices_2, axis=0)
    out = t.const_like(2)[:, :8]
-    linear, var_vals = check_schedule(out, 0)
+    linear, var_vals = out.linear_with_vars()
+    self.assertEqual(len(linear.src), 0)
    run_linear(linear, var_vals)
    self.assertEqual(out.tolist(), [[2]*8]*16)

@@ -3,11 +3,11 @@ import unittest
 import numpy as np
 import torch
 from tinygrad import Tensor, Device, TinyJit, dtypes
+from tinygrad.uop.ops import Ops
 from tinygrad.helpers import GlobalCounters, Context
 from tinygrad.nn import Conv1d, ConvTranspose1d, Conv2d, ConvTranspose2d, Linear, Embedding
 from tinygrad.nn import BatchNorm, LayerNorm, LayerNorm2d, GroupNorm, InstanceNorm, RMSNorm, LSTMCell
 from tinygrad.nn.state import load_state_dict
-from test.helpers import check_schedule
 from tinygrad.engine.realize import run_linear
 from test.helpers import not_support_multi_device, needs_second_gpu, slow

@@ -428,14 +428,18 @@ class TestNN(unittest.TestCase):
    a = Tensor([[1, 5, 9, 11],
                [12, 19, 8, 1]])
    result = layer(a)
-    linear, var_vals = check_schedule(result, kcount)
+    linear, var_vals = result.linear_with_vars()
+    self.assertEqual(len([call for call in linear.src if call.src[0].op is Ops.SINK]), kcount,
+                     "first run realizes weight and embedding")
    run_linear(linear, var_vals)

    b = Tensor([[1, 2, 3],
                [4, 5, 6],
                [7, 8, 9]])
    result = layer(b)
-    linear, var_vals = check_schedule(result, 1)
+    linear, var_vals = result.linear_with_vars()
+    self.assertEqual(1, len([call for call in linear.src if call.src[0].op is Ops.SINK]),
+                     "second run realizes embedding only")
    run_linear(linear, var_vals)
    print(f"Embedding used {GlobalCounters.global_ops} ops")
    self.assertLessEqual(GlobalCounters.global_ops, ops)
@@ -728,17 +728,6 @@ class TestOps(unittest.TestCase):
        else:
          self.assertAlmostEqual(tiny_out, torch_out, msg=f"{x}, {c}")

-  def test_pow_neg_inf_frac_exponent(self):
-    # pow(-inf, 0.3) is +inf, so the gradient 0.3*pow(-inf, -0.7) is 0, never nan
-    helper_test_op(None, lambda x: x**0.3, vals=[[-math.inf]])
-    # is_odd truncates, so it calls 3.3 odd: only the non_int guard keeps pow(-inf, 3.3) from negating to -inf
-    helper_test_op(None, lambda x: x**3.3, vals=[[-math.inf]])
-
-  def test_pow_zero_exponent(self):
-    # x ** 0 is the constant 1 for every x, so the gradient with respect to the base is 0, never nan
-    # TODO: nan ** 0, failed on WEBGPU
-    helper_test_op(None, lambda x,y: x**y, vals=[[-math.inf, math.inf, 0.0], [0.0, 0.0, 0.0]])
-
  def test_pow_zero_tensor(self):
    helper_test_op(None, lambda x,y: x**y, vals=[[0.0], [0.0]])
    # TODO: fix WEBGPU
@@ -1535,8 +1524,6 @@ 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))
@@ -1741,8 +1728,6 @@ class TestOps(unittest.TestCase):
    helper_test_op([(10,10,10)], lambda x: x.log_softmax(0), atol=1e-7, grad_atol=1e-7)
    helper_test_op([(10,10,10)], lambda x: x.log_softmax(1), atol=1e-7, grad_atol=1e-7)
    helper_test_op([(10,10,10)], lambda x: x.log_softmax(2), atol=1e-7, grad_atol=1e-7)
-  def test_softmin(self):
-    helper_test_op([(45,65)], torch.nn.Softmin(dim=1), Tensor.softmin, atol=1e-7, grad_atol=1e-7)

  def test_normalize(self):
    helper_test_op([(45,65)], lambda x: torch.nn.functional.normalize(x), lambda x: x.normalize(), atol=1e-7, grad_atol=1e-7)
@@ -2,7 +2,7 @@ import unittest, pickle, types, tracemalloc
 import numpy as np
 from tinygrad import Tensor, Device, TinyJit, Variable, dtypes
 from tinygrad.helpers import GlobalCounters, ContextVar, Context, DEV
-from tinygrad.uop.ops import PatternMatcher, UPat, UOp, deconstruct_function
+from tinygrad.uop.ops import PatternMatcher, UPat, UOp

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

-  def test_deconstruct_function_nested_comprehension(self):
-    # pre PEP 709, each comprehension is its own code object, so dtypes here is referenced two code objects deep
-    def fxn(): return [[dtypes.int for _ in range(2)] for _ in range(2)]
-    self.assertEqual(types.FunctionType(*deconstruct_function(fxn))(), fxn())
-
  def test_pickle_pattern_matcher(self):
    pm = PatternMatcher([(UPat.cvar('x'), lambda x: x*2)])
-    sink = UOp.const(2)
+    sink = UOp.const(dtypes.int, 2)
    tt = pm.rewrite(sink)
    pm_str = pickle.dumps(pm)
    pm2 = pickle.loads(pm_str)
@@ -8,7 +8,6 @@ from tinygrad.helpers import prod
 from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.wgsl import WGSLRenderer
-from test.helpers import check_schedule
 from tinygrad.runtime.ops_python import PythonRenderer
 from tinygrad.uop.ops import UOp, Ops, KernelInfo, python_alu
 from tinygrad.tensor import Tensor
@@ -25,7 +24,7 @@ def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
  dtype = alu_src_uops[0].dtype
  a = UOp.param(0, dtype, (1,))
  b = UOp.param(1, dtype, (1,))
-  idx = UOp.const(0)
+  idx = UOp.const(dtypes.int, 0)
  ld = b.index(idx).load()
  alu = ld.alu(alu_op, *alu_src_uops)
  store = UOp.store(a.index(idx), alu)
@@ -36,7 +35,7 @@ class TestRendererFailures(unittest.TestCase):
  def test_gated_store_with_alu(self):
    a = UOp.param(0, dtypes.int, (4,))
    gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
-    gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0.valid(gate_alu)), UOp.const(1).cast(dtypes.int)))
+    gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0.valid(gate_alu)), UOp.const(dtypes.int, 1)))
    sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo())
    ret = _test_uop_result([], sink, local_size=[4, 1, 1])[0]
    np.testing.assert_equal(ret, [0, 1, 1, 1])
@@ -46,7 +45,7 @@ class TestRendererFailures(unittest.TestCase):
    a = UOp.param(0, dtypes.int, (8,))
    gate_alu_0 = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
    gate_alu_1 = (lidx1:=UOp.special(2, 'lidx1')).ne(0)
-    gated_alu_store = UOp(Ops.STORE, src=(a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(1).cast(dtypes.int)))
+    gated_alu_store = UOp(Ops.STORE, src=(a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(dtypes.int, 1)))
    sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo())
    ret = _test_uop_result([], sink, local_size=[4, 2, 1])[0]
    np.testing.assert_equal(ret, [0, 0, 0, 0, 0, 1, 1, 1])
@@ -55,14 +54,15 @@ class TestRendererFailures(unittest.TestCase):
 class TestCStyleFailures(unittest.TestCase):
  def test_inline_const_alu(self):
    # CPU doesn't use the max function
-    ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int.min+1).cast(dtypes.int))
+    ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int, dtypes.int.min+1))
    self.assertEqual(ret[0], 1)

  def _test_src_strip_paren(self, op: Ops, should_strip_paren:bool=True):
    dtype = "bool" if op in (Ops.OR, Ops.XOR, Ops.AND) else None
    ret = Tensor.empty(1, dtype=dtype)
    for _ in range(5): ret = python_alu[op](ret, Tensor.empty(1, dtype=dtype))
-    linear, _ = check_schedule(ret, 1)
+    linear = ret.schedule_linear()
+    assert len(linear.src) == 1
    src = to_program(linear.src[0].src[0], Device[Device.DEFAULT].renderer).src[2].arg
    self.assertEqual("("*5 not in src, should_strip_paren)

@@ -80,7 +80,7 @@ class TestWGSLFailures(unittest.TestCase):
  def test_multiply_infinity(self):
    # multiplying a positive constant by infinity should return infinity
    # WGSL pipelines do not handle this reliably, some of which return zero, unless infinity always comes from a read on a dynamic buffer
-    ret = _setup_and_test_alu(Ops.MUL, 5.0, UOp.const(float("inf")).cast(dtypes.float32))
+    ret = _setup_and_test_alu(Ops.MUL, 5.0, UOp.const(dtypes.float32, float("inf")))
    self.assertEqual(ret[0], float("inf"))

  # WGSL has a specific select(alt, val, gate) ternary operator instead of gate?val:alt
@@ -104,7 +104,7 @@ class TestPTXFailures(unittest.TestCase):
  def test_gated_store_with_if(self):
    a = UOp.param(0, dtypes.int, (4,))
    gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
-    val = UOp.const(1).cast(dtypes.int)
+    val = UOp.const(dtypes.int, 1)
    if_uop = UOp(Ops.IF, src=(gate_alu,))
    gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0, if_uop), val))
    sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo())
@@ -6,13 +6,34 @@ import unittest, time
 import numpy as np

 from tinygrad import nn, dtypes, Device, Tensor, Variable
-from tinygrad.uop.ops import Ops, UPat
-from tinygrad.helpers import DEV, GlobalCounters, Context, all_same, temp
-from tinygrad.engine.realize import run_linear
-from test.helpers import check_schedule, assert_kernel_count
+from tinygrad.uop.ops import UOp, Ops, UPat
+from tinygrad.helpers import DEBUG, DEV, GlobalCounters, Context, all_same, temp
+from tinygrad.engine.realize import compile_linear, run_linear

 supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()

+class KernelCountException(Exception): pass
+def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True):
+  if to_prerealize:
+    with Context(DEBUG=0, TRACK_MATCH_STATS=0): Tensor.realize(*to_prerealize)
+  if isinstance(t, Tensor): linear, var_vals = t.linear_with_vars()
+  elif isinstance(t, list) and isinstance(t[0], Tensor): linear, var_vals = Tensor.linear_with_vars(*t)
+  else:
+    assert isinstance(t, UOp), f"can't schedule {t}"
+    linear, var_vals = Tensor(t).linear_with_vars()
+  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:
+    print(f"SCHEDULE ISSUE, expecting {allowed} got {kernel_cnt}")
+    if DEBUG >= 3:
+      for i,call in enumerate(linear.src):
+        print("kernel", i+1)
+        print(call.src[0])
+    raise KernelCountException(f"{kernel_cnt} != {allowed}")
+  # test compiling the linear
+  compile_linear(linear)
+  return linear, var_vals
+
 def _realize_weights(m):
  for p in nn.state.get_parameters(m): p.realize()

@@ -92,9 +113,11 @@ class TestSchedule(unittest.TestCase):
    a2 = mop(a)
    expected = (a+a2).tolist()
    a.assign(a+a2)
-    linear, var_vals = check_schedule(a, expected_kcount)
+    linear, var_vals = a.linear_with_vars()
+    kcount = len(linear.src)
    run_linear(linear, var_vals)
    self.assertListEqual(a.tolist(), expected)
+    self.assertEqual(kcount, expected_kcount)
  def test_setitem_permuted_sched(self): self.test_setitem_sched(lambda x: x.T, 2)
  def test_setitem_paddded_sched(self): self.test_setitem_sched(lambda x: x.shrink_to(4, 1).pad_to(4, 4), 1)

@@ -103,9 +126,9 @@ class TestSchedule(unittest.TestCase):
    a = Tensor.arange(16).clone().realize()
    GlobalCounters.reset()
    a[4] = 3
-    assert_kernel_count(0)
+    self.assertEqual(GlobalCounters.kernel_count, 0)
    a.realize()
-    assert_kernel_count(1)
+    self.assertEqual(GlobalCounters.kernel_count, 1)
    self.assertListEqual(a.tolist(), [0, 1, 2, 3, 3, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15])

  def test_no_extra_contiguous_on_setitem_assign_back(self):
@@ -4,7 +4,6 @@ from tinygrad import Tensor, GlobalCounters, Context, Device
 from tinygrad.dtype import DTypeLike, dtypes
 from tinygrad.engine.realize import run_linear
 from tinygrad.helpers import DEBUG, get_single_element
-from test.helpers import check_schedule

 def single_kernel_softmax(x_in:Tensor, axis=-1, dtype:DTypeLike|None=None) -> Tensor:
  # only support axis =-1
@@ -104,7 +103,8 @@ class TestFuse(unittest.TestCase):
    k = (x @ wk).contiguous()
    v = (x @ wv).contiguous()
    attn = q.scaled_dot_product_attention(k, v)
-    check_schedule(attn, 4) # 3 matmul and 1 attention
+    s = attn.schedule_linear()
+    self.assertEqual(len(s.src), 4) # 3 matmul and 1 attention

  @unittest.skip("needs RANGEIFY>1")
  def test_flash_attention(self):
@@ -6,15 +6,6 @@ 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)
@@ -24,13 +24,13 @@ class TestTinygrad(unittest.TestCase):
    self.assertEqual(Tensor(3.14).shape, ())

  def test_deviceless_const_construct_device_repr(self):
-    t = Tensor(UOp.const(2.0).cast(dtypes.float))
+    t = Tensor(UOp.const(dtypes.float, 2.0))
    self.assertIsNone(t.uop.device)
    self.assertIsNone(t.device)
    self.assertIn("<UOp None", repr(t))

  def test_deviceless_const_realize_noop(self):
-    t = Tensor(UOp.const(2.0).cast(dtypes.float))
+    t = Tensor(UOp.const(dtypes.float, 2.0))
    uop = t.uop
    t.realize()
    self.assertIs(t.uop, uop)
@@ -728,12 +728,12 @@ class TestZeroShapeTensor(unittest.TestCase):
    self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)

  def test_clone_deviceless_const(self):
-    t = Tensor(UOp.const(2.0).cast(dtypes.float)).clone()
+    t = Tensor(UOp.const(dtypes.float, 2.0)).clone()
    np.testing.assert_equal(t.numpy(), 2.0)
    self.assertTrue(t.uop.has_buffer_identity())

  def test_numpy_deviceless_const(self):
-    np.testing.assert_equal(Tensor(UOp.const(2.0).cast(dtypes.float)).numpy(), 2.0)
+    np.testing.assert_equal(Tensor(UOp.const(dtypes.float, 2.0)).numpy(), 2.0)

  def test_clone_with_shrink(self):
    a = Tensor.rand(16, 16)
@@ -756,7 +756,7 @@ class TestZeroShapeTensor(unittest.TestCase):
    np.testing.assert_allclose(a.grad.numpy(), b.grad.numpy())

  def test_clone_deviceless_const_to_cpu(self):
-    t = Tensor(UOp.const(2.0).cast(dtypes.float)).clone(device="CPU")
+    t = Tensor(UOp.const(dtypes.float, 2.0)).clone(device="CPU")
    self.assertEqual(t.device, "CPU")
    np.testing.assert_equal(t.numpy(), 2.0)

@@ -2,7 +2,7 @@ import unittest
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.helpers import Context, getenv, DEV, OSX
-from test.helpers import check_schedule
+from test.backend.test_schedule import check_schedule
 from test.backend.test_dtype_alu import ht, dtypes_float
 import numpy as np
 import math
@@ -19,7 +19,7 @@ def run_uops(uops_list:list[UOp], bufs:list[Buffer]):
  run_linear(UOp(Ops.LINEAR, src=(UOp.sink(*uops_list, arg=KernelInfo()).call(*buf_uops),)))

 def uop(uops:list[UOp], op:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
-  if op is Ops.CONST: uops.append(UOp.const(arg).cast(dtype))
+  if op is Ops.CONST: uops.append(UOp.const(dtype, arg))
  elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, shape=(1,)))
  else: uops.append(UOp(op, dtype, tuple(src), arg))
  return uops[-1]
@@ -43,7 +43,7 @@ def _test_single_value_const(vals, op, dts):
  buf_store = uop(uops, Ops.PARAM, output_dtype, (), 0)
  loads = (uop(uops, Ops.CONST, dtype, [], a) for a,dtype in zip(vals, dts))
  alu = uop(uops, op, output_dtype, loads)
-  out = buf_store[UOp.const(0).cast(dtypes.int32)].store(alu)
+  out = buf_store[UOp.const(dtypes.int32, 0)].store(alu)
  buf = Buffer(Device.DEFAULT, 1, output_dtype).allocate()
  run_uops([out], [buf])
  return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0]
@@ -221,12 +221,12 @@ class TestAssembly(unittest.TestCase):
  def test_bitshift_left(self):
    g1 = UOp.param(0, dtypes.int32, shape=(3,))
    out = UOp.param(1, dtypes.int32, shape=(2,))
-    c1 = UOp.const(2)
-    c2 = UOp.const(3)
+    c1 = UOp.const(dtypes.int, 2)
+    c2 = UOp.const(dtypes.int, 3)
    l1 = g1.index(c1)
    a1 = UOp(Ops.MUL, src=(l1, c1))
    a2 = UOp(Ops.MUL, src=(l1, c2))
-    uops = to_uops_list([out.index(UOp.const(0)).store(a1), out.index(UOp.const(1)).store(a2)],
+    uops = to_uops_list([out.index(UOp.const(dtypes.int, 0)).store(a1), out.index(UOp.const(dtypes.int, 1)).store(a2)],
                        ren=Device[Device.DEFAULT].renderer)
    Device[Device.DEFAULT].renderer.render(uops)
    ops = [x.op for x in uops]
@@ -249,16 +249,16 @@ class TestAssembly(unittest.TestCase):

  def test_mulacc_shl(self):
    g1 = UOp.param(0, dtypes.int32, shape=(2,))
-    c1 = UOp.const(0)
-    c2 = UOp.const(1)
-    expr = g1.index(c1) * UOp.const(4096) + g1.index(c2)
+    c1 = UOp.const(dtypes.int, 0)
+    c2 = UOp.const(dtypes.int, 1)
+    expr = g1.index(c1) * UOp.const(dtypes.int, 4096) + g1.index(c2)
    uops = to_uops_list([expr], ren=Device[Device.DEFAULT].renderer)
    Device[Device.DEFAULT].renderer.render(uops)
    self.assertIn(Ops.MULACC, [x.op for x in uops])

  def test_use_cmpeq(self):
    g = UOp.param(0, dtypes.uint32, shape=(8,))
-    c = UOp.const(7)
+    c = UOp.const(dtypes.uint, 7)
    comp = g.index(c).ne(c).ne(True)
    uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer)
    Device[Device.DEFAULT].renderer.render(uops)
@@ -330,7 +330,6 @@ class TestHCQ(unittest.TestCase):
  # Test profile api
  def test_speed_exec_time(self):
    sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal()
-    st = time.perf_counter()
    TestHCQ.d0.hw_compute_queue_t().timestamp(sig_st) \
                                   .exec(TestHCQ.runtime, TestHCQ.kernargs_ba_ptr, TestHCQ.prg.arg.global_size, TestHCQ.prg.arg.local_size) \
                                   .timestamp(sig_en) \
@@ -338,13 +337,11 @@ class TestHCQ(unittest.TestCase):

    TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
    TestHCQ.d0.timeline_value += 1
-    host_us = (time.perf_counter() - st) * 1e6

    et = float(sig_en.timestamp - sig_st.timestamp)

    print(f"exec kernel time: {et:.2f} us")
-    # emulated devices are only bounded by the host window around submit+wait
-    assert 0.1 <= et <= (host_us if MOCKGPU or Device.DEFAULT in {"CPU"} else 100)
+    assert 0.1 <= et <= (3000000 if MOCKGPU or Device.DEFAULT in {"CPU"} else 100)

  def test_speed_copy_bandwidth(self):
    if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
@@ -355,7 +352,6 @@ class TestHCQ(unittest.TestCase):
    b = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()

    sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal()
-    st = time.perf_counter()
    TestHCQ.d0.hw_copy_queue_t().timestamp(sig_st) \
                                .copy(a._buf, b._buf, SZ) \
                                .timestamp(sig_en) \
@@ -363,14 +359,13 @@ class TestHCQ(unittest.TestCase):

    TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
    TestHCQ.d0.timeline_value += 1
-    host_ms = (time.perf_counter() - st) * 1e3

-    et_ms = float(sig_en.timestamp - sig_st.timestamp) / 1e3
-    assert 0 < et_ms <= host_ms # timestamps are in us and cover only the copy
+    et = float(sig_en.timestamp - sig_st.timestamp)
+    et_ms = et / 1e3

    gb_s = ((SZ / 1e9) / et_ms) * 1e3
    print(f"same device copy:  {et_ms:.2f} ms, {gb_s:.2f} GB/s")
-    assert (0 if MOCKGPU else 10) <= gb_s <= 1000
+    assert (0.2 if MOCKGPU else 10) <= gb_s <= 1000

  def test_speed_cross_device_copy_bandwidth(self):
    if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
@@ -384,7 +379,6 @@ class TestHCQ(unittest.TestCase):
    TestHCQ.d0.allocator._map(b._buf)

    sig_st, sig_en = TestHCQ.d0.new_signal(), TestHCQ.d0.new_signal()
-    st = time.perf_counter()
    TestHCQ.d0.hw_copy_queue_t().timestamp(sig_st) \
                                .copy(a._buf, b._buf, SZ) \
                                .timestamp(sig_en) \
@@ -392,14 +386,13 @@ class TestHCQ(unittest.TestCase):

    TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
    TestHCQ.d0.timeline_value += 1
-    host_ms = (time.perf_counter() - st) * 1e3

-    et_ms = float(sig_en.timestamp - sig_st.timestamp) / 1e3
-    assert 0 < et_ms <= host_ms # timestamps are in us and cover only the copy
+    et = float(sig_en.timestamp - sig_st.timestamp)
+    et_ms = et / 1e3

    gb_s = ((SZ / 1e9) / et_ms) * 1e3
    print(f"cross device copy: {et_ms:.2f} ms, {gb_s:.2f} GB/s")
-    assert (0 if MOCKGPU else 2) <= gb_s <= 100
+    assert (0.2 if MOCKGPU else 2) <= gb_s <= 100

  def test_timeline_signal_rollover(self):
    for queue_type in [TestHCQ.d0.hw_compute_queue_t, TestHCQ.d0.hw_copy_queue_t]:
@@ -1,14 +0,0 @@
-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()
@@ -1,17 +0,0 @@
-from tinygrad import Device, Tensor, TinyJit, dtypes
-from tinygrad.helpers import Timing, Context
-
-GPUS, DEPTH, SZ = 8, 4, 128 * 2**20
-WARMUP, ITERS = 3, 5
-devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(GPUS))
-bufs = tuple(Tensor.empty(SZ, dtype=dtypes.uint8, device=dev).contiguous().realize() for _ in range(DEPTH) for dev in devs)
-
-@TinyJit
-def all_to_all(*srcs:Tensor): return Tensor.realize(*(src.to(dst) for i,src in enumerate(srcs) for j,dst in enumerate(devs) if i % GPUS != j))
-
-if __name__ == "__main__":
-  with Context(ALL2ALL=1, JIT_BATCH_SIZE=0):
-    for i in range(-WARMUP, ITERS):
-      with Timing("ALL2ALL ", lambda ns: f" {SZ*GPUS*(GPUS-1)*DEPTH/ns:.2f} GB/s", enabled=i>=0):
-        all_to_all(*bufs)
-        for dev in devs: Device[dev].synchronize()
@@ -24,8 +24,8 @@ def vision_conv_143():
  c32 = ((c27<3)!=True)&(c27<67)
  c34 = UOp.param(1, dtypes.half, shape=(32, 1024, 4))
  c38 = c5//2
-  c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(Invalid, dtypes.weakint))
-  c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float))
+  c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.weakint, Invalid))
+  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
  c49 = UOp.param(2, dtypes.half, shape=(64, 49, 4))
  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
  c63 = UOp.param(3, dtypes.float, (128,))
@@ -50,8 +50,8 @@ def vision_conv_153():
  c32 = ((c27<3)!=True)&(c27<35)
  c34 = UOp.param(1, dtypes.half, shape=(16, 1024, 4))
  c38 = c5//2
-  c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(Invalid, dtypes.weakint))
-  c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float))
+  c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.weakint, Invalid))
+  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
  c49 = UOp.param(2, dtypes.half, shape=(128, 49, 4))
  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
  c63 = UOp.param(3, dtypes.float, (256,))
@@ -1,5 +1,5 @@
-import unittest, time, itertools
-from tinygrad import Tensor, Context
+import unittest, time
+from tinygrad import Tensor

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

-  @Context(DEV="NULL:HIP:gfx1100")
-  def test_custom_kernel_assign_scaling(self):
-    from tinygrad.uop.ops import UOp, Ops, KernelInfo
-    from tinygrad.runtime.autogen.amd.rdna3.ins import s_nop
-    count = itertools.count(0)
-    def custom_kernel_assign(n):
-      def custom_asm(out):
-        return UOp(Ops.PROGRAM, src=(UOp.sink(out, arg=KernelInfo(f"fxn_{next(count)}")),
-                                     UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=s_nop(i)) for i in range(n*8)))))
-      call = Tensor.custom_kernel(Tensor.empty(1), fxn=custom_asm)[0]
-      return Tensor.cat(*[Tensor.empty(1).assign(call+i) for i in range(n)])
-    self._assert_linear(custom_kernel_assign, n_small=50, n_large=500)
-
 if __name__ == '__main__':
  unittest.main(verbosity=2)
@@ -40,7 +40,7 @@ def random_int_expr(depth=10):
 def random_bool_expr(depth=10, expr1=None):
  if depth == 0: return True
  if expr1 is None: expr1 = random_int_expr(depth-1)
-  expr2 = random.choice([random_or_sub_expression_int(depth-1, expr1), UOp.const(random.randint(-10, 10))])
+  expr2 = random.choice([random_or_sub_expression_int(depth-1, expr1), UOp.const(dtypes.weakint, random.randint(-10, 10))])
  return random.choice(comp_ops)(expr1, expr2)


@@ -82,8 +82,8 @@ if __name__ == "__main__":
            f"v2=Variable(\"{u2.arg[0]}\", {u2.arg[1]}, {u2.arg[2]})\n" +\
            f"v3=Variable(\"{u3.arg[0]}\", {u3.arg[1]}, {u3.arg[2]})\n" +\
            f"expr = {expr}\n" +\
-            f"v1_val, v2_val, v3_val = UOp.const({n1.as_long()}), UOp.const({n2.as_long()})," +\
-                f"UOp.const({n3.as_long()})\n" +\
+            f"v1_val, v2_val, v3_val = UOp.const(dtypes.weakint, {n1.as_long()}), UOp.const(dtypes.weakint, {n2.as_long()})," +\
+                f"UOp.const(dtypes.weakint, {n3.as_long()})\n" +\
            "num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()\n" +\
            "rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()\n" +\
            "assert num==rn, f\"{num} != {rn}\"\n"
@@ -8,11 +8,10 @@ from tinygrad.tensor import _to_np_dtype
 from tinygrad.codegen import to_program
 from tinygrad.dtype import DType, truncate
 from tinygrad.nn.state import get_parameters
-from tinygrad.helpers import T, Target, DEV, DEBUG, Context, GlobalCounters
+from tinygrad.helpers import T, Target, DEV
 from tinygrad.renderer import Renderer
 from tinygrad.codegen import full_rewrite_to_sink, line_rewrite, pm_linearize_cleanups
 from tinygrad.codegen.late.linearizer import linearize
-from tinygrad.engine.realize import compile_linear

 # decorator to skip slow tests by default, run with RUN_SLOW=1 to include them
 slow = unittest.skipUnless(os.getenv("RUN_SLOW"), "slow test, set RUN_SLOW=1 to run")
@@ -35,36 +34,6 @@ def derandomize_model(model):
    p.replace(Tensor.empty(p.shape, device=p.device, dtype=p.dtype))
    p.realize()

-class KernelCountException(Exception):
-  def __init__(self, expected:int, got:int):
-    self.expected, self.got = expected, got
-    super().__init__(f"expected {expected}, got {got}")
-
-def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True):
-  if to_prerealize:
-    with Context(DEBUG=0, TRACK_MATCH_STATS=0): Tensor.realize(*to_prerealize)
-  if isinstance(t, Tensor): linear, var_vals = t.linear_with_vars()
-  elif isinstance(t, list) and isinstance(t[0], Tensor): linear, var_vals = Tensor.linear_with_vars(*t)
-  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:
-    print(f"SCHEDULE ISSUE, expecting {allowed} got {kernel_cnt}")
-    if DEBUG >= 3:
-      for i,call in enumerate(linear.src):
-        print("kernel", i+1)
-        print(call.src[0])
-    raise KernelCountException(allowed, kernel_cnt)
-  return linear, var_vals
-
-def assert_kernel_count(expected:int):
-  got = GlobalCounters.kernel_count
-  if got != expected: raise KernelCountException(expected, got)
-
 def call_is_graph(call:UOp) -> bool:
  ast = call.src[0]
  return ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph"
@@ -84,15 +53,15 @@ def jit_cache_count(linear:UOp) -> int:
 def assert_jit_cache_len(fxn, expected_len):
  linear = fxn.captured.linear if fxn.captured is not None else None
  if linear is None or not linear.src:
-    if expected_len != 0: raise KernelCountException(expected_len, 0)
+    assert expected_len == 0, expected_len
    return
-  if expected_len and all(call_is_hcq(call) for call in linear.src): expected_len = 4 # HCQ2: fence + reset + merged same-queue calls + finalizer
+  if expected_len and all(call_is_hcq(call) for call in linear.src): expected_len = 3 # HCQ2: merged same-queue calls + finalizer + bumps
  if call_is_graph(linear.src[0]):
-    if len(linear.src) != 1: raise KernelCountException(1, len(linear.src))
+    assert len(linear.src) == 1, len(linear.src)
    inner = linear.src[0].src[0].src[0]  # LINEAR UOp inside CUSTOM_FUNCTION
-    if len(inner.src) != expected_len: raise KernelCountException(expected_len, len(inner.src))
+    assert len(inner.src) == expected_len, f"expected {expected_len}, got {len(inner.src)}"
  else:
-    if len(linear.src) != expected_len: raise KernelCountException(expected_len, len(linear.src))
+    assert len(linear.src) == expected_len, f"expected {expected_len}, got {len(linear.src)}"

 def min_normal(dt:DType) -> float: return 2.0 ** (2 - (1 << (dtypes.finfo(dt)[0] - 1)))

@@ -121,7 +90,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple
    bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
    allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
  g = UOp.param(0, uop.dtype, (1,))
-  prg = to_program(UOp.store(g.index(UOp.const(0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
+  prg = to_program(UOp.store(g.index(UOp.const(dtypes.int, 0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
  prog = dev.runtime(prg.to_elf())
  prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs, vals=vals)
  return out_buf.cast(uop.dtype.fmt or "").tolist()[0]
@@ -6,7 +6,7 @@
 #   arg=3: lds - local data share
 #   arg=4: scratch - per-lane scratch memory
 from __future__ import annotations
-import ctypes, functools, re, platform, subprocess, tempfile
+import ctypes, functools, re, platform, subprocess, tempfile, os
 from typing import Callable

 # Set/restore DAZ+FTZ (denormals-are-zero + flush-to-zero) to match RDNA3 default float mode
@@ -109,7 +109,7 @@ def _init_sqtt_encoder():
  _SMEM = (ir3.SMEM, ir4.SMEM, irc.SMEM)
  _VALU = (ir3.VOP1, ir3.VOP2, ir3.VOP3, ir3.VOP3P, ir3.VOPC, ir3.VOPD, ir3.VOP3SD, ir3.VOP3_SDST, ir3.VOP1_SDST,
           ir4.VOP1, ir4.VOP2, ir4.VOP3, ir4.VOP3P, ir4.VOPC, ir4.VOPD, ir4.VOP3SD, ir4.VOP3_SDST, ir4.VOP1_SDST,
-           irc.VOP1, irc.VOP2, irc.VOP3, irc.VOP3P, irc.VOP3PX2, irc.VOPC, irc.VOP3SD, irc.VOP3_SDST)
+           irc.VOP1, irc.VOP2, irc.VOP3, irc.VOP3P, irc.VOPC, irc.VOP3SD, irc.VOP3_SDST)
  _DS = (ir3.DS, ir4.DS, irc.DS)
  _GLOBAL = (ir3.GLOBAL, ir4.VGLOBAL, irc.GLOBAL)
  _FLAT = (ir3.FLAT, ir4.VFLAT, irc.FLAT)
@@ -192,16 +192,16 @@ def _init_sqtt_encoder():

  return emit, finish, finalize

-def _c(val, dtype=dtypes.uint32): return UOp.const(val, dtype)
+def _c(val, dtype=dtypes.uint32): return UOp.const(dtype, val)

 def _u64(lo: UOp, hi: UOp) -> UOp:
  """Combine two 32-bit UOps into a 64-bit UOp."""
-  return lo.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64))
+  return lo.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32))

 def _split64(val: UOp) -> tuple[UOp, UOp]:
  """Split a 64-bit value into (lo, hi) 32-bit values."""
  v64 = val.bitcast(dtypes.uint64) if val.dtype == dtypes.float64 else val.cast(dtypes.uint64) if val.dtype != dtypes.uint64 else val
-  return v64.cast(dtypes.uint32), (v64 >> UOp.const(32, dtypes.uint64)).cast(dtypes.uint32)
+  return v64.cast(dtypes.uint32), (v64 >> UOp.const(dtypes.uint64, 32)).cast(dtypes.uint32)

 _SRC_MOD_TYPES = {16: (dtypes.uint16, dtypes.half, 0x7FFF), 32: (dtypes.uint32, dtypes.float32, 0x7FFFFFFF),
                  64: (dtypes.uint64, dtypes.float64, 0x7FFFFFFFFFFFFFFF)}
@@ -210,7 +210,7 @@ def _apply_src_mods(val: UOp, mod_bit: int, abs_bits: int, neg_bits: int, bits:
  if not (abs_bits & (1 << mod_bit)) and not (neg_bits & (1 << mod_bit)): return val
  ut, ft, mask = _SRC_MOD_TYPES[bits]
  fv = val.cast(ut).bitcast(ft) if bits == 16 else val.bitcast(ft) if val.dtype == ut else val
-  if abs_bits & (1 << mod_bit): fv = (fv.bitcast(ut) & UOp.const(mask, ut)).bitcast(ft)
+  if abs_bits & (1 << mod_bit): fv = (fv.bitcast(ut) & UOp.const(ut, mask)).bitcast(ft)
  if neg_bits & (1 << mod_bit): fv = fv.neg()
  return fv.bitcast(ut).cast(dtypes.uint32) if bits == 16 else fv.bitcast(ut)

@@ -251,7 +251,7 @@ def _to_u32(val: UOp) -> UOp:
  if val.dtype.itemsize == 4: return val.bitcast(dtypes.uint32)  # same size: bitcast (float32->uint32)
  return val.cast(dtypes.uint32)  # different size: cast (bool, int16, etc)
 def _lane_active(exec_mask: UOp, lane: UOp) -> UOp:
-  if exec_mask.dtype == dtypes.uint64: return ((exec_mask >> lane.cast(dtypes.uint64)) & UOp.const(1, dtypes.uint64)).ne(UOp.const(0, dtypes.uint64))
+  if exec_mask.dtype == dtypes.uint64: return ((exec_mask >> lane.cast(dtypes.uint64)) & UOp.const(dtypes.uint64, 1)).ne(UOp.const(dtypes.uint64, 0))
  return ((exec_mask >> lane.cast(dtypes.uint32)) & _c(1)).ne(_c(0))
 def _hi16(v: UOp) -> UOp: return (v >> _c(16)) & _c(0xFFFF)
 def _cond(cond, if_true, if_false):
@@ -260,6 +260,19 @@ def _cond(cond, if_true, if_false):
 def _cond_hi16(cond, val: UOp) -> UOp: return _cond(cond, _hi16(val), val)
 def _apply_opsel(val: UOp, sel_bit: int, opsel: int) -> UOp: return _hi16(val) if opsel & (1 << sel_bit) else val

+def _set_lane_bit(old: UOp, lane: UOp, val: UOp, exec_mask: UOp) -> UOp:
+  """Set/clear a single bit in a mask based on lane index, respecting exec mask."""
+  if old.dtype in (dtypes.uint64, dtypes.int64):
+    dt = dtypes.uint64
+    mask = UOp.const(dt, 1) << lane.cast(dt)
+    new_bit = _to_u32(val).cast(dt) << lane.cast(dt)
+    cleared = old.cast(dt) & (mask ^ UOp.const(dt, 0xFFFFFFFFFFFFFFFF))
+    return _lane_active(exec_mask, lane).where(cleared | new_bit, old.cast(dt))
+  mask = _c(1) << lane.cast(dtypes.uint32)
+  new_bit = _to_u32(val) << lane.cast(dtypes.uint32)
+  cleared = old & (mask ^ _c(MASK32))
+  return _lane_active(exec_mask, lane).where(cleared | new_bit, old)
+
 def _val_to_u32(val: UOp) -> UOp:
  """Convert any value to uint32 for storage (bitcast floats, cast ints)."""
  if val.dtype == dtypes.uint32: return val
@@ -352,18 +365,40 @@ def _write_64bit(val: UOp, wfn, reg_or_addr, is_mem: bool, *args) -> list[UOp]:
  """Write a 64-bit value as two 32-bit writes. args passed to wfn after reg/addr and lo/hi value."""
  lo, hi = _split64(val)
  incr = 4 if is_mem else 1  # 4 bytes for memory addresses, 1 for register indices
-  return [wfn(reg_or_addr, lo, *args), wfn(reg_or_addr + (UOp.const(incr, reg_or_addr.dtype) if isinstance(reg_or_addr, UOp) else incr), hi, *args)]
+  return [wfn(reg_or_addr, lo, *args), wfn(reg_or_addr + (UOp.const(reg_or_addr.dtype, incr) if isinstance(reg_or_addr, UOp) else incr), hi, *args)]

 def _write_val(bits: int, val: UOp, wfn, reg_or_addr, *args, is_mem: bool = False) -> list[UOp]:
  """Write value, splitting 64-bit if needed. bits=64 for 64-bit writes, otherwise 32-bit."""
  return _write_64bit(val, wfn, reg_or_addr, is_mem, *args) if bits == 64 else [wfn(reg_or_addr, _to_u32(val), *args)]

 def _mem_store(mem: UOp, addr: UOp, val: UOp, active: UOp, addr_bits: int = 32, data_bits: int = 32) -> list[UOp]:
-  """Conditional memory store with sub-word support. Returns list of store UOps."""
+  """Conditional memory store with sub-word and unaligned support. Returns list of store UOps.
+
+  FLAT/GLOBAL accesses on AMD hardware are allowed to be unaligned, so a 32-bit store at a
+  byte offset of 1-3 spans two words (handled in 64-bit domain to keep shifts in range)."""
  adt = dtypes.uint64 if addr_bits == 64 else dtypes.uint32
-  word_addr = addr >> UOp.const(2, adt)
+  if data_bits > 32:   # wider stores decompose into dwords; each dword handles its own alignment
+    ws = val.cast(dtypes.uint64) if data_bits > 64 else val
+    return [s for i in range(data_bits // 32)
+            for s in _mem_store(mem, addr + UOp.const(adt, i * 4), ws >> UOp.const(ws.dtype, 32 * i) if i else ws, active, addr_bits, 32)]
+  word_addr = addr >> UOp.const(adt, 2)
  idx = mem.index(word_addr.valid(active))
-  if data_bits == 32: return [idx.store(active.where(_to_u32(val), idx))]
+  if data_bits == 32:
+    byte_off = (addr & UOp.const(adt, 3)).cast(dtypes.uint32)
+    is_unaligned = byte_off.ne(UOp.const(dtypes.uint32, 0))
+    if addr.divides(4) is not None: return [idx.store(active.where(_to_u32(val), idx))]
+    shift = byte_off * UOp.const(dtypes.uint32, 8)
+    val64 = _to_u32(val).cast(dtypes.uint64)
+    # word0 keeps its low byte_off*8 bits, gets val's low bits shifted in; word1 gets the rest
+    low_keep = (UOp.const(dtypes.uint32, 1) << shift) - UOp.const(dtypes.uint32, 1)
+    lo_bits = ((val64 << shift.cast(dtypes.uint64)) & UOp.const(dtypes.uint64, 0xFFFFFFFF)).cast(dtypes.uint32)
+    new_word0 = (idx & low_keep) | lo_bits
+    store0 = idx.store(active.where(is_unaligned.where(new_word0, _to_u32(val)), idx))
+    idx1 = mem.index((word_addr + UOp.const(adt, 1)).cast(dtypes.int64).valid(active & is_unaligned))
+    spill = (val64 >> (UOp.const(dtypes.uint64, 32) - shift.cast(dtypes.uint64))).cast(dtypes.uint32)
+    keep = UOp.const(dtypes.uint32, 0xFFFFFFFF) << shift
+    new_word1 = (idx1 & keep) | spill
+    return [store0, idx1.store((active & is_unaligned).where(new_word1, idx1))]
  # Sub-word store: read-modify-write with mask
  byte_pos = addr.cast(dtypes.uint32) & _c(3)
  byte_shift = byte_pos * _c(8)
@@ -375,7 +410,7 @@ def _mem_store(mem: UOp, addr: UOp, val: UOp, active: UOp, addr_bits: int = 32,
  is_cross = byte_pos.eq(_c(3))
  cross_word0 = (idx & _c(0x00FFFFFF)) | ((val_u32 & _c(0xFF)) << _c(24))
  store0 = idx.store(active.where(is_cross.where(cross_word0, new_word), idx))
-  next_idx = mem.index((word_addr + UOp.const(1, adt)).valid(active & is_cross))
+  next_idx = mem.index((word_addr + UOp.const(adt, 1)).valid(active & is_cross))
  cross_word1 = (next_idx & _c(0xFFFFFF00)) | ((val_u32 >> _c(8)) & _c(0xFF))
  return [store0, next_idx.store((active & is_cross).where(cross_word1, next_idx))]

@@ -384,8 +419,8 @@ def _mem_store_bytes(mem: UOp, addr: UOp, val: UOp, active: UOp, data_bits: int
  stores = []
  val_u32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
  for i in range(data_bits // 8):
-    byte_val = (val_u32 >> UOp.const(i * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32)
-    stores.append(mem.index((addr + UOp.const(i, dtypes.uint64)).valid(active)).store(byte_val.cast(dtypes.uint8)))
+    byte_val = (val_u32 >> UOp.const(dtypes.uint32, i * 8)) & UOp.const(dtypes.uint32, 0xFF)
+    stores.append(mem.index((addr + UOp.const(dtypes.uint64, i)).valid(active)).store(byte_val.cast(dtypes.uint8)))
  return stores

 def _collect_data_slices(assigns: list[tuple[str, UOp]], data_prefix: str, pcode_vars: dict | None = None, op_name: str = "") -> dict[int, UOp]:
@@ -449,8 +484,8 @@ class _Ctx:
  def inst_word(self, dword_idx: int) -> UOp:
    """Read instruction dword from vmem at PC + dword_idx*4."""
    pc = self.rpc()
-    addr = pc if dword_idx == 0 else pc + UOp.const(dword_idx * 4, dtypes.uint64)
-    return self.vmem.index(addr >> UOp.const(2, dtypes.uint64)).load()
+    addr = pc if dword_idx == 0 else pc + UOp.const(dtypes.uint64, dword_idx * 4)
+    return self.vmem.index(addr >> UOp.const(dtypes.uint64, 2)).load()

  def inst_field(self, field) -> UOp:
    """Extract field bits from instruction encoding. Tracks field for canonical key computation."""
@@ -462,15 +497,15 @@ class _Ctx:
    word = self.inst_word(dword_idx)
    if lo // 32 == hi // 32:  # Same dword
      mask = (1 << (hi - lo + 1)) - 1
-      shifted = word if lo_in_dword == 0 else word >> UOp.const(lo_in_dword, dtypes.uint32)
-      return shifted & UOp.const(mask, dtypes.uint32)
+      shifted = word if lo_in_dword == 0 else word >> UOp.const(dtypes.uint32, lo_in_dword)
+      return shifted & UOp.const(dtypes.uint32, mask)
    else:  # Spans two dwords
      lo_bits = 32 - lo_in_dword
      lo_mask = (1 << lo_bits) - 1
      hi_mask = (1 << (hi_in_dword + 1)) - 1
-      lo_part = (word >> UOp.const(lo_in_dword, dtypes.uint32)) & UOp.const(lo_mask, dtypes.uint32)
-      hi_part = self.inst_word(dword_idx + 1) & UOp.const(hi_mask, dtypes.uint32)
-      return lo_part | (hi_part << UOp.const(lo_bits, dtypes.uint32))
+      lo_part = (word >> UOp.const(dtypes.uint32, lo_in_dword)) & UOp.const(dtypes.uint32, lo_mask)
+      hi_part = self.inst_word(dword_idx + 1) & UOp.const(dtypes.uint32, hi_mask)
+      return lo_part | (hi_part << UOp.const(dtypes.uint32, lo_bits))

  def inst_field_signed(self, field) -> UOp:
    """Extract field and sign-extend based on field width."""
@@ -519,19 +554,6 @@ class _Ctx:
      return [self.wsgpr_dyn(reg, lo), self.wsgpr_dyn(reg + _c(1), hi)]
    return [self.wsgpr_dyn(reg, val)]

-  def wmask_lane_bit(self, reg: UOp, lane: UOp, val: UOp, exec_mask: UOp) -> list[UOp]:
-    """Set/clear bit `lane` of the mask at `reg` from val for exec-active lanes, preserving memory for inactive lanes"""
-    active, bit = _lane_active(exec_mask, lane), _to_u32(val)
-    if self.wave_size <= 32:
-      old = self.rsgpr_dyn(reg)
-      mask = _c(1) << lane.cast(dtypes.uint32)
-      return [self.wsgpr_dyn(reg, active.where((old & (mask ^ _c(MASK32))) | (bit << lane.cast(dtypes.uint32)), old))]
-    off = (lane & _c(31, dtypes.int)).cast(dtypes.uint32)
-    mask = _c(1) << off
-    def half(old: UOp, sel: UOp) -> UOp: return sel.where(active.where((old & (mask ^ _c(MASK32))) | (bit << off), old), old)
-    return [self.wsgpr_dyn(reg, half(self.rsgpr_dyn(reg), lane < _c(32, dtypes.int))),
-            self.wsgpr_dyn(reg + _c(1), half(self.rsgpr_dyn(reg + _c(1)), _c(32, dtypes.int) <= lane))]
-
  def rmask(self, reg: UOp) -> UOp:
    """Read a lane mask (VCC/EXEC). Combines lo/hi for wave64."""
    if self.wave_size > 32: return _u64(self.rsgpr_dyn(reg), self.rsgpr_dyn(reg + _c(1)))
@@ -584,7 +606,7 @@ class _Ctx:
      inline = is_float_const.where(float_inline.bitcast(dtypes.uint64), int_inline.bitcast(dtypes.uint64))
      # Literal handling: F64 VOP puts literal in high 32 bits; B64/I64/U64 VOP and SOP zero-extend
      if literal is not None:
-        lit_val = literal.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64) if is_f64 else literal.cast(dtypes.uint64)
+        lit_val = literal.cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32) if is_f64 else literal.cast(dtypes.uint64)
        inline = off.eq(_c(255)).where(lit_val, inline)
      scalar_val = (off < _c(128)).where(sgpr_val, inline)
    else:
@@ -602,7 +624,7 @@ class _Ctx:

  def inc_pc(self) -> list[UOp]:
    """Increment PC by instruction size in bytes. Returns [store]."""
-    new_pc = self.rpc() + UOp.const(self.inst_size, dtypes.uint64)
+    new_pc = self.rpc() + UOp.const(dtypes.uint64, self.inst_size)
    lo, hi = _split64(new_pc)
    return [self.wsgpr_dyn(_c(PC_LO_IDX), lo), self.wsgpr_dyn(_c(PC_HI_IDX), hi)]

@@ -668,7 +690,7 @@ class _Ctx:
    if 'VCC' not in srcs: srcs['VCC'] = self.rmask(_c(vcc_reg))
    srcs.update({'EXEC': exec_mask, 'SCC': self.rsgpr_dyn(_c(SCC.offset)), 'laneId': lane, 'VDST': vdst_reg,
                  'ROUND_MODE': _c(0), 'ROUND_TOWARD_ZERO': _c(0), 'ROUND_NEAREST_EVEN': _c(0), '_vgpr': self.vgpr, '_wave_size': self.wave_size,
-                  'MAX_FLOAT_F32': UOp.const(3.4028234663852886e38, dtypes.float32),
+                  'MAX_FLOAT_F32': UOp.const(dtypes.float32, 3.4028234663852886e38),
                  # CDNA SDWA byte/word select constants (E32 always uses BYTE0/WORD0 defaults)
                  'SDWA_SRC0_SEL': _c(0), 'BYTE0': _c(0), 'BYTE1': _c(1), 'BYTE2': _c(2), 'BYTE3': _c(3),
                  'WORD0': _c(0), 'WORD1': _c(1)})  # rounding mode and SDWA constants
@@ -709,7 +731,7 @@ class _Ctx:
      # VGPR bit-slice assignment: VGPR[lane][reg][hi:lo] = (vgpr_idx, rhs_val, hi, lo[, cond]) -> read-modify-write
      if dest.startswith('VGPR[') and re.search(r'\[\d+:\d+\]', dest):
        # VGPR bit-slice: (vgpr_idx, rhs_val, hi_bit, lo_bit) - hi/lo are UOp constants
-        hi_bit, lo_bit = int(val[2].val), int(val[3].val)
+        hi_bit, lo_bit = int(val[2].arg), int(val[3].arg)
        width = hi_bit - lo_bit + 1
        old = self.vgpr.index(val[0]).load()
        new_val = _set_bits(old, _val_to_bits(val[1]), width, lo_bit).cast(dtypes.uint32)
@@ -718,7 +740,9 @@ class _Ctx:
        raw_stores.append(('vgpr_direct', self.vgpr.index(val[0].valid(active)).store(new_val)))
        continue
      if 'D0' in dest and '[laneId]' in dest:
-        raw_stores.extend([('vcc', s) for s in self.wmask_lane_bit(_c(VCC_LO.offset), lane, val, exec_mask)])
+        old_vcc = self.rmask(_c(VCC_LO.offset))
+        new_vcc = _set_lane_bit(old_vcc, lane, val, exec_mask)
+        raw_stores.extend([('vcc', s) for s in self.wmask(_c(VCC_LO.offset), new_vcc)])
      elif dest.startswith('D0'):
        dest_suffix = re.match(r'D0\.(\w+)', dest)
        if dest_suffix is not None:
@@ -730,24 +754,24 @@ class _Ctx:
            d0_width, slice_mask = d0_hi_bit - d0_lo_bit + 1, (1 << (d0_hi_bit - d0_lo_bit + 1)) - 1
            val_bits = val.bitcast(dtypes.uint16).cast(dtypes.uint32) if val.dtype == dtypes.half else \
                       val.cast(dtypes.uint32) if val.dtype in (dtypes.uint16, dtypes.int16) else \
-                       val.cast(dtypes.uint32) & UOp.const(slice_mask, dtypes.uint32)
+                       val.cast(dtypes.uint32) & UOp.const(dtypes.uint32, slice_mask)
            raw_stores.append(('vgpr_slice', (d0_lo_bit, d0_width, val_bits)))
            continue
        # For integer ops with clamp, use pre-computed saturated value; for floats, clamp to [0,1]
        if int_saturate is not None: val = int_saturate
        elif clmp and val.dtype in (dtypes.float32, dtypes.half, dtypes.float64):
-          clamped = val.maximum(UOp.const(0.0, val.dtype)).minimum(UOp.const(1.0, val.dtype))
-          val = _FUNCS['isNAN'](val).where(UOp.const(0.0, val.dtype), clamped)
+          clamped = val.maximum(UOp.const(val.dtype, 0.0)).minimum(UOp.const(val.dtype, 1.0))
+          val = _FUNCS['isNAN'](val).where(UOp.const(val.dtype, 0.0), clamped)
        if val.dtype in (dtypes.uint64, dtypes.int64, dtypes.float64):
          lo, hi = _split64(val)
          raw_stores.extend([('vgpr', self.wvgpr_dyn(vdst_reg, lane, lo, exec_mask)),
                             ('vgpr', self.wvgpr_dyn(vdst_reg + _c(1), lane, hi, exec_mask))])
        elif val.dtype in (dtypes.half, dtypes.uint16, dtypes.int16):
          result, old_val = _val_to_u32(val), self.rvgpr_dyn(vdst_reg, lane)
-          hi_result = (old_val & UOp.const(0xFFFF, dtypes.uint32)) | (result << UOp.const(16, dtypes.uint32))
+          hi_result = (old_val & UOp.const(dtypes.uint32, 0xFFFF)) | (result << UOp.const(dtypes.uint32, 16))
          # GFX9/CDNA zeroes upper 16 bits on lo-half write; RDNA preserves them
-          lo_result = (result & UOp.const(0xFFFF, dtypes.uint32)) if self.wave_size == 64 else \
-                      (old_val & UOp.const(0xFFFF0000, dtypes.uint32)) | (result & UOp.const(0xFFFF, dtypes.uint32))
+          lo_result = (result & UOp.const(dtypes.uint32, 0xFFFF)) if self.wave_size == 64 else \
+                      (old_val & UOp.const(dtypes.uint32, 0xFFFF0000)) | (result & UOp.const(dtypes.uint32, 0xFFFF))
          result = opsel_dst_hi.where(hi_result, lo_result) if isinstance(opsel_dst_hi, UOp) else hi_result if opsel_dst_hi else lo_result
          raw_stores.append(('vgpr', self.wvgpr_dyn(vdst_reg, lane, result, exec_mask)))
        else: raw_stores.append(('vgpr', self.wvgpr_dyn(vdst_reg, lane, _val_to_u32(val), exec_mask)))
@@ -765,8 +789,8 @@ class _Ctx:
    if slice_stores:
      result = self.rvgpr_dyn(vdst_reg, lane)
      for lo_bit, width, val_bits in slice_stores:
-        mask = UOp.const(((1 << width) - 1) << lo_bit, dtypes.uint32)
-        result = (result & (mask ^ UOp.const(0xFFFFFFFF, dtypes.uint32))) | (val_bits << UOp.const(lo_bit, dtypes.uint32))
+        mask = UOp.const(dtypes.uint32, ((1 << width) - 1) << lo_bit)
+        result = (result & (mask ^ UOp.const(dtypes.uint32, 0xFFFFFFFF))) | (val_bits << UOp.const(dtypes.uint32, lo_bit))
      lane_stores.append(self.wvgpr_dyn(vdst_reg, lane, result, exec_mask))
    # VCC/EXEC mask writes must be computed BEFORE VGPR stores to avoid reading modified VGPRs.
    # When vdst overlaps with src operands (e.g. v_add_co_u32 v[0], vcc, s[8], v[0]), the carry
@@ -788,8 +812,8 @@ class _Ctx:
 def _compile_sopp(inst: ir3.SOPP | ir4.SOPP, ctx: _Ctx) -> UOp:
  simm16 = ctx.inst_field_signed(type(inst).simm16).cast(dtypes.int16)
  if inst.op in (ir3.SOPPOp.S_ENDPGM, ir4.SOPPOp.S_ENDPGM, irc.SOPPOp.S_ENDPGM):
-    return UOp.sink(ctx.wsgpr_dyn(_c(PC_LO_IDX), UOp.const(0xFFFFFFFF, dtypes.uint32)),
-                          ctx.wsgpr_dyn(_c(PC_HI_IDX), UOp.const(0xFFFFFFFF, dtypes.uint32)))
+    return UOp.sink(ctx.wsgpr_dyn(_c(PC_LO_IDX), UOp.const(dtypes.uint32, 0xFFFFFFFF)),
+                          ctx.wsgpr_dyn(_c(PC_HI_IDX), UOp.const(dtypes.uint32, 0xFFFFFFFF)))
  # S_BARRIER: advance PC past the barrier instruction. The execution loop detects barriers before executing and handles synchronization.
  barrier_ops = {ir3.SOPPOp.S_BARRIER, irc.SOPPOp.S_BARRIER}
  if hasattr(ir4.SOPPOp, 'S_BARRIER_WAIT'): barrier_ops.add(ir4.SOPPOp.S_BARRIER_WAIT)
@@ -802,8 +826,8 @@ def _compile_sopp(inst: ir3.SOPP | ir4.SOPP, ctx: _Ctx) -> UOp:
    pc_bytes = ctx.rpc()  # PC is already 64-bit byte address
    vcc, exec_val = ctx.rmask(_c(VCC_LO.offset)), ctx.rexec()
    srcs: dict[str, UOp|int] = {'PC': pc_bytes.cast(dtypes.int64), 'SIMM16': simm16, 'SCC': ctx.rsgpr_dyn(_c(SCC.offset)), 'VCC': vcc,
-            'VCCZ': vcc.eq(UOp.const(0, vcc.dtype)).cast(dtypes.uint32),
-            'EXECZ': exec_val.eq(UOp.const(0, exec_val.dtype)).cast(dtypes.uint32)}
+            'VCCZ': vcc.eq(UOp.const(vcc.dtype, 0)).cast(dtypes.uint32),
+            'EXECZ': exec_val.eq(UOp.const(exec_val.dtype, 0)).cast(dtypes.uint32)}
    for dest, val in parse_pcode(pcode, srcs)[1]:
      if dest == 'PC' or dest.startswith('PC.'):
        lo, hi = _split64(val.cast(dtypes.uint64))
@@ -831,12 +855,12 @@ def _compile_smem(inst: ir3.SMEM | ir4.SMEM, ctx: _Ctx) -> UOp:
  part = op_name.rsplit('_', 1)[1]  # B32, DWORD, DWORDX2, U8, I8, etc.
  nval = int(part.removeprefix('DWORD').removeprefix('X') or '1') if 'DWORD' in part else int(part[1:]) / 32 * (-1 if part[0] == 'I' else 1)
  ndwords = max(1, int(abs(nval)))
-  dword_base = addr >> UOp.const(2, dtypes.uint64)
-  vals = [ctx.vmem.index(dword_base + UOp.const(i, dtypes.uint64)) for i in range(ndwords)]
+  dword_base = addr >> UOp.const(dtypes.uint64, 2)
+  vals = [ctx.vmem.index(dword_base + UOp.const(dtypes.uint64, i)) for i in range(ndwords)]
  if abs(nval) < 1:
    nbits = int(abs(nval) * 32)
-    byte_off = (addr & UOp.const(3, dtypes.uint64)).cast(dtypes.uint32) * UOp.const(8, dtypes.uint32)
-    extracted = (vals[0] >> byte_off) & UOp.const((1 << nbits) - 1, dtypes.uint32)
+    byte_off = (addr & UOp.const(dtypes.uint64, 3)).cast(dtypes.uint32) * UOp.const(dtypes.uint32, 8)
+    extracted = (vals[0] >> byte_off) & UOp.const(dtypes.uint32, (1 << nbits) - 1)
    vals[0] = extracted.cast({8: dtypes.int8, 16: dtypes.int16}[nbits]).cast(dtypes.int32).bitcast(dtypes.uint32) if nval < 0 else extracted
  stores = [ctx.wsgpr_dyn(sdata_reg + _c(i), vals[i]) for i in range(ndwords)]
  return UOp.sink(*stores, *ctx.inc_pc())
@@ -939,7 +963,7 @@ def _dpp16_ctrl(lane: UOp, dpp: int, row_mask: int, bank_mask: int, wave_size: i
  enabled = (((_c(row_mask) >> row.cast(dtypes.uint32)) & _c(1)).ne(_c(0)) &
             (((_c(bank_mask) >> bank.cast(dtypes.uint32)) & _c(1)).ne(_c(0))))
  op, arg = decode_dpp16(dpp)
-  src_lane, valid = lane_i, UOp.const(True)
+  src_lane, valid = lane_i, UOp.const(dtypes.bool, True)

  if op == 'quad_perm':
    assert isinstance(arg, tuple)
@@ -965,7 +989,7 @@ def _load_dpp16_src0(ctx: _Ctx, inst, lane: UOp, fallback: UOp) -> UOp:
                                         getattr(inst, 'bank_mask', 0xf) or 0xf, ctx.wave_size)
  safe_src_lane = (enabled & valid).where(src_lane, _c(0, dtypes.int))
  swizzled = ctx.rvgpr_dyn(ctx.inst_field(type(inst).vsrc0), safe_src_lane)
-  invalid = UOp.const(0, fallback.dtype) if getattr(inst, 'bc', 0) else fallback
+  invalid = UOp.const(fallback.dtype, 0) if getattr(inst, 'bc', 0) else fallback
  return enabled.where(valid.where(swizzled, invalid), fallback)

 def _compile_sdwa(inst: irc.VOP1_SDWA | irc.VOP2_SDWA | irc.VOP2_SDWA_SDST | irc.VOPC_SDWA_SDST, ctx: _Ctx) -> UOp:
@@ -1037,11 +1061,13 @@ def _compile_sdwa(inst: irc.VOP1_SDWA | irc.VOP2_SDWA | irc.VOP2_SDWA_SDST | irc
        result = _sdwa_write(old, result, dst_sel, dst_unused)
      stores.append(ctx.wvgpr_dyn(vdst_reg, lane, result, exec_mask))
    elif dest.startswith('VCC'):
-      stores.extend(ctx.wmask_lane_bit(_c(VCC_LO.offset), lane, val, exec_mask))
+      old_vcc = ctx.rmask(_c(VCC_LO.offset))
+      stores.extend(ctx.wmask(_c(VCC_LO.offset), _set_lane_bit(old_vcc, lane, val, exec_mask)))
  if vcc_val is not None:
    # Initialize sdst to 0 before lane loop (old value may be unrelated data), then set lane bits in loop
    init_stores = [ctx.wsgpr_dyn(sdst_off, _c(0)), ctx.wsgpr_dyn(sdst_off + _c(1), _c(0))]
-    stores.extend(ctx.wmask_lane_bit(sdst_off, lane, vcc_val, exec_mask))
+    old_sdst = ctx.rmask(sdst_off)
+    stores.extend(ctx.wmask(sdst_off, _set_lane_bit(old_sdst, lane, vcc_val, exec_mask)))
    if stores:
      return UOp.sink(*init_stores, UOp.sink(*stores).end(lane), *ctx.inc_pc())
    return UOp.sink(*init_stores, *ctx.inc_pc())
@@ -1146,9 +1172,6 @@ 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)
@@ -1158,7 +1181,7 @@ def _compile_vopc(inst: ir3.VOPC|ir3.VOPC_DPP16|ir3.VOP3|ir4.VOPC|ir4.VOPC_DPP16
        s1 = _apply_src_mods(s1, 0, 1 if getattr(inst, 'src1_abs', 0) else 0, 1 if getattr(inst, 'src1_neg', 0) else 0, bits['s1'])
      s0 = _apply_src_mods(s0, 0, abs_bits, neg_bits, bits['s0'])
      s1 = _apply_src_mods(s1, 1, abs_bits, neg_bits, bits['s1'])
-    for dest, val in parse_pcode(pcode, {'S0': s0, 'S1': s1, 'laneId': lc, 'D0': UOp.const(0, dtypes.uint64)})[1]:
+    for dest, val in parse_pcode(pcode, {'S0': s0, 'S1': s1, 'laneId': lc, 'D0': UOp.const(dtypes.uint64, 0)})[1]:
      if '[laneId]' in dest and ('D0' in dest or 'EXEC' in dest): return val.cast(dtypes.uint32)
    return _c(0)

@@ -1188,8 +1211,8 @@ def _compile_bitop3(inst, ctx: _Ctx, exec_mask: UOp, bits: dict, op_name: str) -
  is_16 = 'B16' in op_name
  dt, mask = (dtypes.uint16, 0xFFFF) if is_16 else (dtypes.uint32, 0xFFFFFFFF)
  s0, s1, s2 = src0.cast(dt), src1.cast(dt), src2.cast(dt)
-  def bnot(v): return v ^ UOp.const(mask, dt)
-  result = UOp.const(0, dt)
+  def bnot(v): return v ^ UOp.const(dt, mask)
+  result = UOp.const(dt, 0)
  for i in range(8):
    if not (ttbl & (1 << i)): continue
    result = result | ((s0 if i & 4 else bnot(s0)) & (s1 if i & 2 else bnot(s1)) & (s2 if i & 1 else bnot(s2)))
@@ -1243,7 +1266,7 @@ def _compile_vop3(inst: ir3.VOP3 | ir4.VOP3 | irc.VOP3, ctx: _Ctx) -> UOp:
  src0 = _apply_src_mods(src0, 0, abs_bits, neg_bits, bits['s0'])
  src1 = _apply_src_mods(src1, 1, abs_bits, neg_bits, bits['s1'])
  src2 = _apply_src_mods(src2, 2, abs_bits, neg_bits, bits['s2'])
-  srcs = {'S0': src0, 'S1': src1, 'S2': src2, 'OPSEL': UOp.const(opsel, dtypes.uint32)}
+  srcs = {'S0': src0, 'S1': src1, 'S2': src2, 'OPSEL': UOp.const(dtypes.uint32, opsel)}
  if 'CNDMASK' in op_name and src2 is not None: srcs['VCC'] = src2
  # FMAC instructions need D0 (accumulator) from destination register
  if 'FMAC' in op_name: srcs['D0'] = ctx.rvgpr_dyn(vdst_reg, lane)
@@ -1322,7 +1345,7 @@ def _compile_vop3sd(inst: ir3.VOP3SD | ir4.VOP3SD | irc.VOP3SD, ctx: _Ctx) -> UO
  else:
    return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask, sdst_reg=inst.sdst.offset)

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

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

-  scaled = isinstance(inst, irc.VOP3PX2)
-  if scaled:
-    assert isinstance(inst, irc.VOP3PX2)
-    # F8F6F4 input formats: 0=FP8(E4M3), 1=BF8(E5M2). FP6/FP4 (2-4) not emulated.
-    src0_fmt, src1_fmt = int(inst.cbsz), int(inst.blgp)
-    if src0_fmt > 1 or src1_fmt > 1: raise RuntimeError(f"unsupported scaled MFMA formats cbsz={src0_fmt} blgp={src1_fmt}")
-    # scale_src0/scale_src1 are source operands pointing at 32-bit registers holding 4 packed E8M0 scale exponents.
-    # The 2-bit opsel/opsel_hi select which byte applies to A/B for this instruction.
-    scale0_off = ctx.inst_field(type(inst).scale_src0)
-    scale1_off = ctx.inst_field(type(inst).scale_src1)
-    sel0, sel1 = int(inst.opsel) & 3, int(inst.opsel_hi) & 3
-    def _scale_exp(off: UOp, sel: int, lane: UOp) -> UOp:
-      sv = ctx.rsrc_dyn(off, lane, 32)
-      byte = (sv >> UOp.const(sel * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32)
-      return byte.cast(dtypes.int32) - UOp.const(127, dtypes.int32)
-    # combined A*B scale for this lane: 2^(ea-127) * 2^(eb-127)
-    def scale_factor(lane: UOp) -> UOp:
-      return UOp.exp2((_scale_exp(scale0_off, sel0, lane) + _scale_exp(scale1_off, sel1, lane)).cast(dtypes.float32))
-
  m = _re.search(r'(\d+)X(\d+)X(\d+)', op_name)
  if m is None: raise ValueError(f"could not parse MFMA dimensions from {op_name}")
  M, N, K = int(m.group(1)), int(m.group(2)), int(m.group(3))
@@ -1416,64 +1420,53 @@ def _compile_mfma(inst: irc.VOP3P|irc.VOP3PX2, ctx: _Ctx) -> UOp:
  # Layout: tmp[0..n_a_elems-1] = A[m][k], tmp[n_a_elems..n_a_elems+n_b_elems-1] = B[n][k]
  # Within each group of lanes, lane%grp_sub gives M/N index, lane//grp_sub gives sub-block
  grp_sub = min(M, 16)  # lanes within group mapped to M/N dimension
-  b_off = UOp.const(n_a_elems, dtypes.int)
+  b_off = UOp.const(dtypes.int, n_a_elems)
  acc_dt = dtypes.int32 if is_int_out else dtypes.float32
  # Use uint32 temp array to prevent optimizer from eliminating f16→f32 bitcast chains.
  # The optimizer folds bitcast(uint32→float32) stores to float32 arrays, losing the conversion.
  tmp = UOp.placeholder((n_a_elems + n_b_elems,), dtypes.uint32, slot=0, addrspace=AddrSpace.LOCAL)

-  # Per-operand fp8 format ("fp8"=E4M3, "bf8"=E5M2) for A and B
-  if 'F8F6F4' in op_name:
-    assert isinstance(inst, (irc.VOP3P_MFMA, irc.VOP3PX2))
-    _fmts = {0: "fp8", 1: "bf8"}
-    a_fmt, b_fmt = _fmts.get(int(inst.cbsz), "fp8"), _fmts.get(int(inst.blgp), "fp8")
-  elif is_fp8:
-    # A/B formats from name suffix, e.g. V_MFMA_F32_16X16X32_BF8_FP8
-    suffixes = op_name.rsplit('_', 2)[-2:]
-    a_fmt, b_fmt = ("bf8" if sfx == "BF8" else "fp8" for sfx in suffixes)
-  else: a_fmt = b_fmt = "fp8"
-
-  def cvt_elem(raw: UOp, sub_idx: int, fp8_fmt: str = "fp8") -> UOp:
+  def cvt_elem(raw: UOp, sub_idx: int) -> UOp:
    if is_i8:
      # Extract i8, sign-extend to i32
-      byte_val = (raw >> UOp.const(sub_idx * 8, dtypes.uint32)) & UOp.const(0xFF, dtypes.uint32)
-      return (byte_val.cast(dtypes.int32) ^ UOp.const(0x80, dtypes.int32)) - UOp.const(0x80, dtypes.int32)
+      byte_val = (raw >> UOp.const(dtypes.uint32, sub_idx * 8)) & UOp.const(dtypes.uint32, 0xFF)
+      return (byte_val.cast(dtypes.int32) ^ UOp.const(dtypes.int32, 0x80)) - UOp.const(dtypes.int32, 0x80)
    elif is_f32_src:
      return raw  # already uint32 (f32 bit pattern)
    elif is_fp8:
-      return _FUNCS[f"{fp8_fmt}_to_f32"](raw >> UOp.const(sub_idx * 8, dtypes.uint32)).bitcast(dtypes.uint32)
+      return ((raw >> UOp.const(dtypes.uint32, sub_idx * 8)) & UOp.const(dtypes.uint32, 0xFF)).cast(dtypes.uint32)
    elif is_bf16:
      # bf16→f32 bits: just shift left by 16 (bf16 is upper 16 bits of f32)
-      return ((raw >> UOp.const(sub_idx * 16, dtypes.uint32)) & UOp.const(0xFFFF, dtypes.uint32)) << UOp.const(16, dtypes.uint32)
+      return ((raw >> UOp.const(dtypes.uint32, sub_idx * 16)) & UOp.const(dtypes.uint32, 0xFFFF)) << UOp.const(dtypes.uint32, 16)
    else:
      # f16→f32 conversion using float arithmetic to avoid UOp optimizer eliminating the conversion.
      # The optimizer folds bitcast(uint32→float32) chains, so we compute the float value directly.
-      h = (raw >> UOp.const(sub_idx * 16, dtypes.uint32)) & UOp.const(0xFFFF, dtypes.uint32)
-      sign = (h >> UOp.const(15, dtypes.uint32)) & UOp.const(1, dtypes.uint32)
-      exp = (h >> UOp.const(10, dtypes.uint32)) & UOp.const(0x1F, dtypes.uint32)
-      mant = h & UOp.const(0x3FF, dtypes.uint32)
+      h = (raw >> UOp.const(dtypes.uint32, sub_idx * 16)) & UOp.const(dtypes.uint32, 0xFFFF)
+      sign = (h >> UOp.const(dtypes.uint32, 15)) & UOp.const(dtypes.uint32, 1)
+      exp = (h >> UOp.const(dtypes.uint32, 10)) & UOp.const(dtypes.uint32, 0x1F)
+      mant = h & UOp.const(dtypes.uint32, 0x3FF)
      # Use bf16 path: shift left by 16 to create bf16 bits, then shift mantissa and adjust exponent in float domain
      # bf16 bits = (sign << 15) | (exp_bf16 << 7) | mant_bf16 -- but f16 and bf16 have different formats
      # Instead: construct f32 bits properly, use a local uint32 array to force materialization
-      f32_bits = (sign << UOp.const(31, dtypes.uint32)) | \
-                 ((exp + UOp.const(112, dtypes.uint32)) << UOp.const(23, dtypes.uint32)) | \
-                 (mant << UOp.const(13, dtypes.uint32))
-      is_zero = exp.eq(UOp.const(0, dtypes.uint32))
+      f32_bits = (sign << UOp.const(dtypes.uint32, 31)) | \
+                 ((exp + UOp.const(dtypes.uint32, 112)) << UOp.const(dtypes.uint32, 23)) | \
+                 (mant << UOp.const(dtypes.uint32, 13))
+      is_zero = exp.eq(UOp.const(dtypes.uint32, 0))
      # Return uint32 (f32 bit pattern) — stored directly to uint32 temp array, bitcast to float on read
-      return is_zero.where(UOp.const(0, dtypes.uint32), f32_bits)
+      return is_zero.where(UOp.const(dtypes.uint32, 0), f32_bits)

  read_lane = ctx.range()
  # For 32x32: lane%16 = M/N index within 16-wide block, lane//16 = which of 4 quarter-waves
  # Groups: lanes 0-31 = group 0, lanes 32-63 = group 1
  # Within group: (lane%32)%16 = M/N[0-15], (lane%32)//16 selects M/N[0-15] or [16-31]
-  lane_in_grp = read_lane % UOp.const(grp_size, dtypes.int)
-  grp_idx = read_lane // UOp.const(grp_size, dtypes.int)
+  lane_in_grp = read_lane % UOp.const(dtypes.int, grp_size)
+  grp_idx = read_lane // UOp.const(dtypes.int, grp_size)

  if M == 32:
    # 32x32: lane_in_grp%16 = sub-row/col (0-15), lane_in_grp//16 = block (0=rows 0-15, 1=rows 16-31)
-    sub_mn = lane_in_grp % UOp.const(16, dtypes.int)
-    block_mn = lane_in_grp // UOp.const(16, dtypes.int)
-    mn_idx = block_mn * UOp.const(16, dtypes.int) + sub_mn  # actual M/N index (0-31)
+    sub_mn = lane_in_grp % UOp.const(dtypes.int, 16)
+    block_mn = lane_in_grp // UOp.const(dtypes.int, 16)
+    mn_idx = block_mn * UOp.const(dtypes.int, 16) + sub_mn  # actual M/N index (0-31)
  else:
    mn_idx = lane_in_grp  # for 16x16 and 4x4

@@ -1483,20 +1476,20 @@ def _compile_mfma(inst: irc.VOP3P|irc.VOP3PX2, ctx: _Ctx) -> UOp:
    # Read A/B sources. Use rsrc_dyn for inline constants/SGPRs (src_off < 256), rvgpr_dyn for VGPRs (src_off >= 256).
    a_raw = src0_is_vgpr.where(ctx.rvgpr_dyn(src0_r + _c(reg_idx), read_lane),
                                ctx.rsrc_dyn(src0_off, _c(0, dtypes.int), 32))
-    a_val = cvt_elem(a_raw, sub_idx, a_fmt)
+    a_val = cvt_elem(a_raw, sub_idx)
    if M == 4:
-      a_idx = grp_idx * UOp.const(M * K, dtypes.int) + mn_idx * UOp.const(K, dtypes.int) + UOp.const(kl, dtypes.int)
+      a_idx = grp_idx * UOp.const(dtypes.int, M * K) + mn_idx * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, kl)
    else:
-      a_idx = mn_idx * UOp.const(K, dtypes.int) + grp_idx * UOp.const(k_per_grp, dtypes.int) + UOp.const(kl, dtypes.int)
+      a_idx = mn_idx * UOp.const(dtypes.int, K) + grp_idx * UOp.const(dtypes.int, k_per_grp) + UOp.const(dtypes.int, kl)
    read_stores.append(tmp.index(a_idx).store(a_val))

    b_raw = src1_is_vgpr.where(ctx.rvgpr_dyn(src1_r + _c(reg_idx), read_lane),
                                ctx.rsrc_dyn(src1_off, _c(0, dtypes.int), 32))
-    b_val = cvt_elem(b_raw, sub_idx, b_fmt)
+    b_val = cvt_elem(b_raw, sub_idx)
    if M == 4:
-      b_idx = b_off + grp_idx * UOp.const(N * K, dtypes.int) + mn_idx * UOp.const(K, dtypes.int) + UOp.const(kl, dtypes.int)
+      b_idx = b_off + grp_idx * UOp.const(dtypes.int, N * K) + mn_idx * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, kl)
    else:
-      b_idx = b_off + mn_idx * UOp.const(K, dtypes.int) + grp_idx * UOp.const(k_per_grp, dtypes.int) + UOp.const(kl, dtypes.int)
+      b_idx = b_off + mn_idx * UOp.const(dtypes.int, K) + grp_idx * UOp.const(dtypes.int, k_per_grp) + UOp.const(dtypes.int, kl)
    read_stores.append(tmp.index(b_idx).store(b_val))

  read_phase = UOp.group(*read_stores).end(read_lane)
@@ -1509,17 +1502,6 @@ def _compile_mfma(inst: irc.VOP3P|irc.VOP3PX2, ctx: _Ctx) -> UOp:
  #   Actually: 16 ACCVGPRs per lane, organized as 4 groups (l//32 gives half, each half has 2 sub-groups) of 4 rows
  tmp2 = tmp.after(read_phase)

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

@@ -1528,11 +1510,11 @@ def _compile_mfma(inst: irc.VOP3P|irc.VOP3PX2, ctx: _Ctx) -> UOp:
    # Lane mapping: n = (lane%32)%16 + ((lane%32)//16)*16, gives column 0-31
    # Row groups: 4 groups of 4, covering rows 0-31. Group g (0-3): rows g*4 .. g*4+3
    # group assignment: lane//16 gives quarter (0-3), each quarter maps to 4 rows
-    c_lane_in_32 = compute_lane % UOp.const(32, dtypes.int)
-    c_sub = c_lane_in_32 % UOp.const(16, dtypes.int)
-    c_block = c_lane_in_32 // UOp.const(16, dtypes.int)
-    n_idx = c_block * UOp.const(16, dtypes.int) + c_sub
-    c_half = compute_lane // UOp.const(32, dtypes.int)  # 0 or 1
+    c_lane_in_32 = compute_lane % UOp.const(dtypes.int, 32)
+    c_sub = c_lane_in_32 % UOp.const(dtypes.int, 16)
+    c_block = c_lane_in_32 // UOp.const(dtypes.int, 16)
+    n_idx = c_block * UOp.const(dtypes.int, 16) + c_sub
+    c_half = compute_lane // UOp.const(dtypes.int, 32)  # 0 or 1

    for out_reg in range(16):
      # Each half covers 8 rows. out_reg 0-3: rows 0-3 (half0) or 16-19 (half1)
@@ -1543,14 +1525,17 @@ def _compile_mfma(inst: irc.VOP3P|irc.VOP3PX2, ctx: _Ctx) -> UOp:
      # acc[r] at lane l -> D[m][n] where n = (l%32)%16 + ((l%32)//16)*16
      # m = (l//32)*16 + (r//4)*4 + (r%4)  ... giving rows in blocks of 4
      # So: m_base = half * 16 + (out_reg // 4) * 4 + (out_reg % 4)
-      m_base = c_half * UOp.const(16, dtypes.int) + UOp.const((out_reg // 4) * 4 + (out_reg % 4), dtypes.int)
+      m_base = c_half * UOp.const(dtypes.int, 16) + UOp.const(dtypes.int, (out_reg // 4) * 4 + (out_reg % 4))

      acc_v = (ctx.raccvgpr_dyn if use_acc else ctx.rvgpr_dyn)(src2_r + _c(out_reg), compute_lane, src2_is_vgpr)
      if is_int_out: acc_v = acc_v.cast(dtypes.int32)
      else: acc_v = acc_v.bitcast(dtypes.float32)
      acc = src2_is_vgpr.where(acc_v, acc_scalar)

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

      if is_int_out:
        compute_stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)(
@@ -1560,8 +1545,8 @@ def _compile_mfma(inst: irc.VOP3P|irc.VOP3PX2, ctx: _Ctx) -> UOp:
          vdst_reg + _c(out_reg), compute_lane, acc.bitcast(dtypes.uint32), exec_mask))
  else:
    # 16x16 and 4x4: each lane computes out_per_lane outputs
-    n_idx = compute_lane % UOp.const(grp_sub, dtypes.int)
-    c_grp = compute_lane // UOp.const(grp_sub, dtypes.int)
+    n_idx = compute_lane % UOp.const(dtypes.int, grp_sub)
+    c_grp = compute_lane // UOp.const(dtypes.int, grp_sub)

    for out_reg in range(out_per_lane):
      acc_v = (ctx.raccvgpr_dyn if use_acc else ctx.rvgpr_dyn)(src2_r + _c(out_reg), compute_lane, src2_is_vgpr)
@@ -1571,14 +1556,18 @@ def _compile_mfma(inst: irc.VOP3P|irc.VOP3PX2, ctx: _Ctx) -> UOp:

      if M == 4:
        # 4x4: each group is independent. A/B indexed per-group.
-        m_base = c_grp * UOp.const(M * K, dtypes.int) + UOp.const(out_reg * K, dtypes.int)
-        b_base = b_off + c_grp * UOp.const(N * K, dtypes.int) + n_idx * UOp.const(K, dtypes.int)
+        m_base = c_grp * UOp.const(dtypes.int, M * K) + UOp.const(dtypes.int, out_reg * K)
+        for k in range(K):
+          a_val = tmp2.index(m_base + UOp.const(dtypes.int, k)).bitcast(acc_dt)
+          b_val = tmp2.index(b_off + c_grp * UOp.const(dtypes.int, N*K) + n_idx * UOp.const(dtypes.int, K)+UOp.const(dtypes.int, k)).bitcast(acc_dt)
+          acc = acc + a_val * b_val
      else:
        # 16x16: K is split across groups. Shared MxK/NxK arrays.
-        m_base = c_grp * UOp.const(out_per_lane, dtypes.int) + UOp.const(out_reg, dtypes.int)
-        b_base = b_off + n_idx * UOp.const(K, dtypes.int)
-
-      acc = _dot_accum(acc, m_base if M == 4 else m_base * UOp.const(K, dtypes.int), b_base, compute_lane)
+        m_base = c_grp * UOp.const(dtypes.int, out_per_lane) + UOp.const(dtypes.int, out_reg)
+        for k in range(K):
+          a_val = tmp2.index(m_base * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, k)).bitcast(acc_dt)
+          b_val = tmp2.index(b_off + n_idx * UOp.const(dtypes.int, K) + UOp.const(dtypes.int, k)).bitcast(acc_dt)
+          acc = acc + a_val * b_val

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

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

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

  # ACCVGPR_WRITE/READ/MOV: copies between VGPR and ACCVGPR register files
  # Detect by checking operand types for ACCVGPR involvement
@@ -1733,8 +1710,8 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P | irc.VOP3PX2, ctx: _
      scalar_hi_sel = src_lo if not opsel_hi_bit else is_sgpr_pair.where(sgpr_hi, src_lo)
      lo = is_vgpr.where(vgpr_hi if opsel_lo else vgpr_lo, scalar_lo_sel)
      hi = is_vgpr.where(vgpr_hi if opsel_hi_bit else vgpr_lo, scalar_hi_sel)
-      if neg_lo: lo = lo ^ UOp.const(0x80000000, dtypes.uint32)
-      if neg_hi_bit: hi = hi ^ UOp.const(0x80000000, dtypes.uint32)
+      if neg_lo: lo = lo ^ UOp.const(dtypes.uint32, 0x80000000)
+      if neg_hi_bit: hi = hi ^ UOp.const(dtypes.uint32, 0x80000000)
      return _u64(lo, hi)
    srcs = {'S0': build_pk_f32(src0, src_offs[0], opsel & 1, opsel_hi & 1, neg & 1, neg_hi & 1),
            'S1': build_pk_f32(src1, src_offs[1], opsel & 2, opsel_hi & 2, neg & 2, neg_hi & 2),
@@ -1745,35 +1722,35 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P | irc.VOP3PX2, ctx: _
    def apply_abs(v, bit, opsel_hi_bit, opsel_bit):
      if not (neg_hi & bit): return v
      # Apply abs based on whether source is f32 or f16
-      if not (combined_opsel_hi & opsel_hi_bit): return v & UOp.const(0x7FFFFFFF, dtypes.uint32)  # f32 abs
-      if opsel & opsel_bit: return v & UOp.const(0x7FFF0000, dtypes.uint32)  # f16 hi abs (preserve lo)
-      return v & UOp.const(0xFFFF7FFF, dtypes.uint32)  # f16 lo abs (preserve hi)
+      if not (combined_opsel_hi & opsel_hi_bit): return v & UOp.const(dtypes.uint32, 0x7FFFFFFF)  # f32 abs
+      if opsel & opsel_bit: return v & UOp.const(dtypes.uint32, 0x7FFF0000)  # f16 hi abs (preserve lo)
+      return v & UOp.const(dtypes.uint32, 0xFFFF7FFF)  # f16 lo abs (preserve hi)
    def apply_neg_mix(v, bit, opsel_hi_bit, opsel_bit):
      if not (neg & bit): return v
-      if not (combined_opsel_hi & opsel_hi_bit): return v ^ UOp.const(0x80000000, dtypes.uint32)  # f32 neg
-      if opsel & opsel_bit: return v ^ UOp.const(0x80000000, dtypes.uint32)  # f16 hi neg
-      return v ^ UOp.const(0x00008000, dtypes.uint32)  # f16 lo neg
+      if not (combined_opsel_hi & opsel_hi_bit): return v ^ UOp.const(dtypes.uint32, 0x80000000)  # f32 neg
+      if opsel & opsel_bit: return v ^ UOp.const(dtypes.uint32, 0x80000000)  # f16 hi neg
+      return v ^ UOp.const(dtypes.uint32, 0x00008000)  # f16 lo neg
    s0_mod = apply_neg_mix(apply_abs(src0, 1, 1, 1), 1, 1, 1)
    s1_mod = apply_neg_mix(apply_abs(src1, 2, 2, 2), 2, 2, 2)
    s2_mod = apply_neg_mix(apply_abs(src2, 4, 4, 4), 4, 4, 4)
    srcs = {'S@0': s0_mod, 'S@1': s1_mod, 'S@2': s2_mod,
-            'OPSEL_HI': UOp.const(combined_opsel_hi, dtypes.uint32), 'OPSEL': UOp.const(opsel, dtypes.uint32)}
+            'OPSEL_HI': UOp.const(dtypes.uint32, combined_opsel_hi), 'OPSEL': UOp.const(dtypes.uint32, opsel)}
  else:
    def get_half_bits(val: UOp, use_hi: bool, apply_neg: bool = False) -> UOp:
-      bits = ((val >> UOp.const(16, dtypes.uint32)) if use_hi else val) & UOp.const(0xFFFF, dtypes.uint32)
+      bits = ((val >> UOp.const(dtypes.uint32, 16)) if use_hi else val) & UOp.const(dtypes.uint32, 0xFFFF)
      if apply_neg: bits = bits.cast(dtypes.uint16).bitcast(dtypes.half).neg().bitcast(dtypes.uint16).cast(dtypes.uint32)
      return bits
    def build_remapped_src(src: UOp, opsel_lo_bit: int, opsel_hi_bit: int, neg_lo_bit: int, neg_hi_bit: int) -> UOp:
      lo = get_half_bits(src, bool(opsel_lo_bit), bool(neg_lo_bit))
      hi = get_half_bits(src, bool(opsel_hi_bit), bool(neg_hi_bit))
-      return lo | (hi << UOp.const(16, dtypes.uint32))
+      return lo | (hi << UOp.const(dtypes.uint32, 16))
    # DOT IU instructions use NEG bits for signed/unsigned selection, not fp16 negation
    is_dot_iu = 'DOT' in op_name and 'IU' in op_name
    n0, n1, n2, nh0, nh1, nh2 = (0, 0, 0, 0, 0, 0) if is_dot_iu else (neg & 1, neg & 2, neg & 4, neg_hi & 1, neg_hi & 2, neg_hi & 4)
    srcs = {'S0': build_remapped_src(src0, opsel & 1, opsel_hi & 1, n0, nh0),
            'S1': build_remapped_src(src1, opsel & 2, opsel_hi & 2, n1, nh1),
            'S2': build_remapped_src(src2, opsel & 4, 1 if opsel_hi2 else 0, n2, nh2)}
-    if is_dot_iu: srcs['NEG'] = UOp.const(neg, dtypes.uint32)
+    if is_dot_iu: srcs['NEG'] = UOp.const(dtypes.uint32, neg)
  return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask)

 def _compile_vopd(inst: ir3.VOPD | ir4.VOPD, ctx: _Ctx) -> UOp:
@@ -1826,7 +1803,7 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
  # CDNA acc bit: when set, VGPR operands (vdst/vdata) target ACCVGPR file instead of VGPR
  use_acc = bool(getattr(inst, 'acc', 0))
  mem = ctx.lds if is_lds else ctx.scratch if is_scratch else ctx.vmem
-  addr_shift = UOp.const(2, dtypes.uint32 if is_lds else dtypes.uint64)
+  addr_shift = UOp.const(dtypes.uint32 if is_lds else dtypes.uint64, 2)

  # Extract register info - all dynamic for deduplication
  if is_lds:
@@ -1876,20 +1853,20 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
      return addr
    offset64 = offset.cast(dtypes.uint64)
    # Dynamic saddr check: saddr < 124 means valid SGPR, otherwise use VGPR pair for address
-    use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(False)
+    use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(dtypes.bool, False)
    if is_scratch:
      scratch_stride = ctx.rsgpr_dyn(_c(SCRATCH_STRIDE_IDX)).cast(dtypes.uint64)
      base = lane.cast(dtypes.uint64) * scratch_stride
      # SVE (Scratch VGPR Enable): when SVE=1, VADDR is used as offset; when SVE=0, VADDR is ignored
      sve = getattr(inst, 'sve', 0)
      vaddr = ctx.rvgpr_dyn(addr_reg, lane).cast(dtypes.uint64)
-      addr_offset = vaddr if sve == 1 else UOp.const(0, dtypes.uint64)
+      addr_offset = vaddr if sve == 1 else UOp.const(dtypes.uint64, 0)
      # Add saddr value only if use_saddr is true (saddr < 124)
-      saddr_contrib = use_saddr.where(ctx.rsgpr_dyn(saddr_reg).cast(dtypes.uint64), UOp.const(0, dtypes.uint64)) \
-        if saddr_reg is not None else UOp.const(0, dtypes.uint64)
+      saddr_contrib = use_saddr.where(ctx.rsgpr_dyn(saddr_reg).cast(dtypes.uint64), UOp.const(dtypes.uint64, 0)) \
+        if saddr_reg is not None else UOp.const(dtypes.uint64, 0)
      return base + addr_offset + saddr_contrib + offset64
    # FLAT/GLOBAL: choose between SGPR base (saddr) or VGPR pair (addr) based on saddr validity
-    saddr_base = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(0, dtypes.uint64)
+    saddr_base = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(dtypes.uint64, 0)
    vaddr_base = _u64(ctx.rvgpr_dyn(addr_reg, lane), ctx.rvgpr_dyn(addr_reg + _c(1), lane))
    # When saddr is valid: base = saddr pair, vaddr is 32-bit offset; otherwise: base = 0, vaddr is 64-bit address
    base_addr = use_saddr.where(saddr_base + ctx.rvgpr_dyn(addr_reg, lane).cast(dtypes.uint64), vaddr_base)
@@ -1919,18 +1896,18 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
        data = {'DATA': ctx.rvgpr_dyn(vdata_reg, lane), 'DATA1': ctx.rvgpr_dyn(vdata_reg + _c(1), lane),
                'DATA2': ctx.rvgpr_dyn(vdata_reg + _c(2), lane)}
      elif data_bits_mem <= 32:
-        data = {'DATA': ctx.rvgpr_dyn(vdata_reg, lane), 'DATA2': ctx.rvgpr_dyn(data1_reg, lane) if has_data1 else UOp.const(0, dtypes.uint32)}
+        data = {'DATA': ctx.rvgpr_dyn(vdata_reg, lane), 'DATA2': ctx.rvgpr_dyn(data1_reg, lane) if has_data1 else UOp.const(dtypes.uint32, 0)}
      else:
        data = {'DATA': _u64(ctx.rvgpr_dyn(vdata_reg, lane), ctx.rvgpr_dyn(vdata_reg + _c(1), lane)),
-                'DATA2': _u64(ctx.rvgpr_dyn(data1_reg, lane), ctx.rvgpr_dyn(data1_reg + _c(1), lane)) if has_data1 else UOp.const(0, dtypes.uint64)}
+                'DATA2': _u64(ctx.rvgpr_dyn(data1_reg, lane), ctx.rvgpr_dyn(data1_reg + _c(1), lane)) if has_data1 else UOp.const(dtypes.uint64, 0)}
      # RDNA3 uses ADDR/OFFSET, RDNA4 uses vgpr_a/offset (lowercase) + CalcDsAddr function
      return {'ADDR': addr, 'ADDR_BASE': addr, 'OFFSET': offset, 'OFFSET0': offset0, 'OFFSET1': offset1, '_lds': mem, 'laneId': lane,
              'vgpr_a': ctx.rvgpr_dyn(addr_reg, lane), 'offset': offset, 'offset0': offset0, 'offset1': offset1, **data}
    active = _lane_active(exec_mask, lane)
    # saddr < 124 means valid SGPR pair, otherwise use 0 (NULL means no saddr contribution)
-    use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(False)
-    saddr_raw = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(0, dtypes.uint64)
-    saddr_base = use_saddr.where(saddr_raw, UOp.const(0, dtypes.uint64))
+    use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(dtypes.bool, False)
+    saddr_raw = _u64(ctx.rsgpr_dyn(saddr_reg), ctx.rsgpr_dyn(saddr_reg + _c(1))) if saddr_reg is not None else UOp.const(dtypes.uint64, 0)
+    saddr_base = use_saddr.where(saddr_raw, UOp.const(dtypes.uint64, 0))
    # Sign-extend offset to 64-bit for the final address calculation
    ioffset64 = offset.cast(dtypes.int64).cast(dtypes.uint64)
    # v_addr for CalcGlobalAddr: when saddr valid, use low 32 bits as offset; otherwise full 64-bit address. Include ioffset.
@@ -1945,13 +1922,13 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
    # acc bit: read/write ACCVGPR instead of VGPR for data operands
    _rvdata = (lambda r, l, *a: ctx.raccvgpr_dyn(r, l)) if use_acc else ctx.rvgpr_dyn
    vdata = _rvdata(vdata_reg, lane).cast(dtypes.uint64) if 'STORE' in op_name \
-      else _rvdata(vdst_reg, lane) if 'D16' in op_name else UOp.const(0, dtypes.uint32)
+      else _rvdata(vdst_reg, lane) if 'D16' in op_name else UOp.const(dtypes.uint32, 0)
    if 'STORE' in op_name and data_bits_mem >= 64:
-      vdata = vdata | (_rvdata(vdata_reg + _c(1), lane).cast(dtypes.uint64) << UOp.const(32, dtypes.uint64))
+      vdata = vdata | (_rvdata(vdata_reg + _c(1), lane).cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32))
    srcs = {'ADDR': addr, 'VDATA': vdata, '_vmem': mem, '_active': active,
            'laneId': lane, 'v_addr': vaddr_base, 's_saddr': saddr_base, 'SADDR': saddr_base, 'OFFSET': offset}
    for i in range(data_bits_mem // 32):
-      srcs[f'VDATA{i}'] = _rvdata(vdata_reg + _c(i), lane) if 'STORE' in op_name else UOp.const(0, dtypes.uint32)
+      srcs[f'VDATA{i}'] = _rvdata(vdata_reg + _c(i), lane) if 'STORE' in op_name else UOp.const(dtypes.uint32, 0)
    return srcs

  def make_stores(dest: str, val: UOp, lane: UOp, active: UOp, writes_return_data: bool) -> list[UOp]:
@@ -2028,7 +2005,7 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp:
  offset, offen, idxen = ctx.inst_field(type(inst).offset), ctx.inst_field(type(inst).offen), ctx.inst_field(type(inst).idxen)

  # V# descriptor: base[0:1], num_records[2], stride=word3[13:0]
-  base = _u64(ctx.rsgpr_dyn(srsrc), ctx.rsgpr_dyn(srsrc + _c(1))) & UOp.const(0xFFFFFFFFFFFF, dtypes.uint64)
+  base = _u64(ctx.rsgpr_dyn(srsrc), ctx.rsgpr_dyn(srsrc + _c(1))) & UOp.const(dtypes.uint64, 0xFFFFFFFFFFFF)
  num_records = ctx.rsgpr_dyn(srsrc + _c(2))
  stride = (ctx.rsgpr_dyn(srsrc + _c(3)) & _c(0x3FFF)).cast(dtypes.uint64)

@@ -2045,29 +2022,31 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp:
  buffer_offset = (stride * index + voff + offset.cast(dtypes.uint64)).cast(dtypes.uint32)
  in_bounds = active & buffer_offset.__lt__(num_records)
  addr = base + soff + buffer_offset.cast(dtypes.uint64)
-  addr = in_bounds.where(addr, UOp.const(0, dtypes.uint64))  # safe address when OOB
+  addr = in_bounds.where(addr, UOp.const(dtypes.uint64, 0))  # safe address when OOB
  mem = ctx.vmem

  stores: list[UOp] = []
  if is_lds and not is_store:
-    # LDS load: buffer -> LDS (bypass VGPRs), LDS addr = M0[17:0] + lane * elem_size
+    # LDS load: buffer -> LDS (bypass VGPRs), LDS addr = M0[17:0] + lane * elem_size.
+    # HW never takes a per-lane LDS address: kittens' direct fill sets M0 (s_mov_b32 m0, sN)
+    # before every lds instruction, giving lane-linear chunks with the swizzle on the GLOBAL side.
    lds_base = ctx.rsgpr_dyn(_c(124)) & _c(0x3FFFF)
    lds_addr = lds_base + lane.cast(dtypes.uint32) * _c(n_dwords * 4)
    for i in range(n_dwords):
-      word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64)
+      word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2)
      val = in_bounds.where(mem.index(word_addr.cast(dtypes.int64)).load(), _c(0))
      lds_idx = (lds_addr + _c(i * 4)) >> _c(2)
      lds_slot = ctx.lds.index(lds_idx.valid(active))
      stores.append(lds_slot.store(active.where(val, lds_slot)))
  elif is_store:
    for i in range(n_dwords):
-      word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64)
+      word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2)
      idx = mem.index(word_addr.cast(dtypes.int64).valid(in_bounds))
      val = (ctx.raccvgpr_dyn if use_acc else ctx.rvgpr_dyn)(vdata + _c(i), lane)
      stores.append(idx.store(in_bounds.where(_to_u32(val), idx)))
  else:
    for i in range(n_dwords):
-      word_addr = (addr + UOp.const(i * 4, dtypes.uint64)) >> UOp.const(2, dtypes.uint64)
+      word_addr = (addr + UOp.const(dtypes.uint64, i * 4)) >> UOp.const(dtypes.uint64, 2)
      val = in_bounds.where(mem.index(word_addr.cast(dtypes.int64).valid(in_bounds)).load(), _c(0))
      stores.append((ctx.waccvgpr_dyn if use_acc else ctx.wvgpr_dyn)(vdata + _c(i), lane, val, exec_mask))
  return UOp.sink(UOp.group(*stores).end(lane), *ctx.inc_pc())
@@ -2089,7 +2068,7 @@ _INST_HANDLERS: dict[type, Callable[..., UOp]] = {
  irc.SOPP: _compile_sopp, irc.SMEM: _compile_smem, irc.SOP1: _compile_sop, irc.SOP2: _compile_sop, irc.SOPC: _compile_sop, irc.SOPK: _compile_sop,
  irc.VOP1: _compile_vop12, irc.VOP1_DPP16: _compile_vop12, irc.VOP2: _compile_vop12, irc.VOP2_DPP16: _compile_vop12,
  irc.VOPC: _compile_vopc, irc.VOP3: _compile_vop3,
-  irc.VOP3_SDST: _compile_vop3, irc.VOP3SD: _compile_vop3sd, irc.VOP3P: _compile_vop3p, irc.VOP3PX2: _compile_vop3p,
+  irc.VOP3_SDST: _compile_vop3, irc.VOP3SD: _compile_vop3sd, irc.VOP3P: _compile_vop3p,
  irc.VOP1_SDWA: _compile_sdwa, irc.VOP2_SDWA: _compile_sdwa, irc.VOP2_SDWA_SDST: _compile_sdwa, irc.VOPC_SDWA_SDST: _compile_sdwa,
  irc.DS: _compile_mem_op, irc.FLAT: _compile_mem_op, irc.GLOBAL: _compile_mem_op, irc.SCRATCH: _compile_mem_op,
  irc.MUBUF: _compile_mubuf,
@@ -2159,6 +2138,8 @@ def _decode_at(pc: int, arch: str):
 F32_INLINE = {240: 0x3f000000, 241: 0xbf000000, 242: 0x3f800000, 243: 0xbf800000,  # 0.5, -0.5, 1.0, -1.0
              244: 0x40000000, 245: 0xc0000000, 246: 0x40800000, 247: 0xc0800000, 248: 0x3e22f983}  # 2.0, -2.0, 4.0, -4.0, 1/(2*pi)

+_inst_hist: dict = {}
+
 class WaveState:
  __slots__ = ('vgpr_buf', 'sgpr_buf', 'accvgpr_buf', '_vgpr_mv', '_sgpr_mv', 'n_lanes', 'wave_size')

@@ -2254,7 +2235,8 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
  # Use Buffer objects with external_ptr=0 for vmem
  vmem_buf = Buffer('CPU', 1 << 40, dtypes.uint32, options=BufferSpec(external_ptr=0)).ensure_allocated()
  lds_buf = Buffer('CPU', max(lds_size // 4, 1), dtypes.uint32).ensure_allocated()
-  scratch_buf = Buffer('CPU', scratch_size * wave_size, dtypes.uint8).ensure_allocated() if scratch_size else None
+  ctypes.memset(lds_buf._buf.va_addr, 0, max(lds_size, 4))
+  # NOTE: scratch (private/spill) memory is per-wavefront; buffers are allocated in the wave loop below

  # Initialize SQTT encoder — emits packets inline as instructions execute (only when profiling)
  if PROFILE:
@@ -2272,6 +2254,8 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
        print(colored(msg, 'green') if len(_canonical_runner_cache) > prev_len else msg)
    return program[pc]

+  if os.getenv("EMU_TRACE_INST") or os.getenv("EMU_WATCH_PC") is not None:
+    print(f"[emu-dispatch] gx={gx} gy={gy} gz={gz} lx={lx} ly={ly} lz={lz} scratch={scratch_size}", flush=True)
  # Set DAZ+FTZ during emulator execution, restore afterward to avoid breaking hypothesis tests
  # Only trace the first workgroup (like real HW traces one CU/SIMD), subsequent workgroups run but don't add to trace
  tracing = bool(PROFILE)
@@ -2282,18 +2266,27 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
        for gidx in range(gx):
          # Initialize all wavefronts for this workgroup
          waves: list[tuple[WaveState, list]] = []
+          wave_scratch_bufs: list[Buffer] = []   # keep alive for the dispatch
          for wave_start in range(0, total_threads, wave_size):
            st = _init_wave(lib, wave_start, total_threads, lx, ly, lz, args_ptr, rsrc2, scratch_size, arch, gidx, gidy, gidz, user_data,
                            wave_size)
+            # each wavefront owns its private (spill) scratch segment: on real HW the ring is indexed by
+            # (wave_id, lane), and scratch addresses here are lane*stride within the wave's segment.
+            if scratch_size:
+              wave_scratch_bufs.append(Buffer('CPU', scratch_size * wave_size, dtypes.uint8).ensure_allocated())
+              ctypes.memset(wave_scratch_bufs[-1]._buf.va_addr, 0, scratch_size * wave_size)
            c_bufs = [ctypes.c_uint64(st.sgpr_buf._buf.va_addr), ctypes.c_uint64(st.vgpr_buf._buf.va_addr),
                      ctypes.c_uint64(vmem_buf._buf.va_addr), ctypes.c_uint64(lds_buf._buf.va_addr),
-                      ctypes.c_uint64(scratch_buf._buf.va_addr if scratch_buf else 0),
+                      ctypes.c_uint64(wave_scratch_bufs[-1]._buf.va_addr if scratch_size else 0),
                      ctypes.c_uint64(st.accvgpr_buf._buf.va_addr)]
            waves.append((st, c_bufs))

          # Execute wavefronts with barrier synchronization
          # Each wave runs until it hits s_barrier or s_endpgm. When all waves have stopped, release barrier waves.
          done = [False] * len(waves)
+          _exec_dumped: set = set()
+          _trace_on = bool(os.getenv("EMU_TRACE_INST")) and total_threads >= int(os.getenv("EMU_TRACE_MIN_THREADS", "1"))
+          if _trace_on: _inst_hist.clear()
          for total_inst in range(10_000_000):
            if all(done): break
            for wi, (st, c_bufs) in enumerate(waves):
@@ -2304,9 +2297,23 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
                if pc == ENDPGM_PC:
                  done[wi] = True
                  if tracing: sqtt_finish(wi)
+                  if os.getenv("EMU_EXEC_DUMP") and wi not in _exec_dumped:
+                    _exec_dumped.add(wi)
+                    print(f"[exec-dump] gid=({gidx},{gidy},{gidz}) wave{wi} "
+                          f"exec_lo={st._read_sgpr(EXEC_LO.offset):08x} exec_hi={st._read_sgpr(EXEC_LO.offset+1):08x}")
                  break
                fxn, globals_list, is_barrier, inst = _ensure_compiled(pc)
                if DEBUG >= 5: print(f"  exec gid=({gidx},{gidy},{gidz}) w={wi} PC={pc - lib}: {inst!r}", flush=True)
+                if _trace_on:
+                  key = (gidx, gidy, gidz, wi)
+                  _inst_hist.setdefault(key, []).append((pc - lib, type(inst).__name__, getattr(inst, 'op', None) and inst.op.name))
+                wpc, wwave = os.getenv("EMU_WATCH_PC"), int(os.getenv("EMU_WATCH_WAVE", "0"))
+                if wpc is not None and (int(wpc) < 0 or (pc - lib) == int(wpc)) and wi == wwave and gidx == gidy == gidz == 0:
+                  if int(wpc) < 0 and total_inst < 800: print(f"[watch-stream] pc={pc-lib} {inst!r}", flush=True)
+                  for rv in os.getenv("EMU_WATCH_VGPR", "").split(","):
+                    if rv: print(f"[watch pc={pc-lib}] wave{wi} {rv}:", [st._read_vgpr(int(rv[1:]), l) for l in range(8)], flush=True)
+                  for rv in os.getenv("EMU_WATCH_SGPR", "").split(","):
+                    if rv: print(f"[watch pc={pc-lib}] wave{wi} {rv}:", [st._read_sgpr(int(rv[1:]))], flush=True)
                fxn(*[c_bufs[g] for g in globals_list])
                if tracing:
                  inst_op = inst.op.value if hasattr(inst, 'op') else 0
@@ -2314,9 +2321,26 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
                if is_barrier: break  # s_barrier hit: PC already advanced past it, pause this wave
              else: raise RuntimeError("exceeded 1M instructions in single wave, likely infinite loop")
            # All waves have either hit barrier or endpgm — release barrier waves for next round
+            if os.getenv("EMU_DUMP_ROUNDS") and lds_size > 0 and total_inst < int(os.getenv("EMU_DUMP_ROUNDS")):
+              import numpy as _np
+              lds_words = _np.frombuffer((ctypes.c_uint32 * (lds_size//4)).from_address(lds_buf._buf.va_addr), dtype=_np.uint32)
+              nz = _np.argwhere(lds_words != 0)
+              print(f"[emu-dump] round={total_inst} lds nonzero words={len(nz)}",
+                    (f"first16={[f'w{w}:0x{lds_words[w]:08x}' for w in nz[:16].flatten()]}" if len(nz) else ""), flush=True)
+              for wi2, (st2, _) in enumerate(waves):
+                v = _np.frombuffer((ctypes.c_uint32 * (256*st2.wave_size)).from_address(st2.vgpr_buf._buf.va_addr), dtype=_np.uint32)
+                if st2.wave_size == 64:
+                  av = _np.frombuffer((ctypes.c_uint32 * (256*st2.wave_size)).from_address(st2.accvgpr_buf._buf.va_addr), dtype=_np.uint32)
+                else: av = _np.zeros(1, dtype=_np.uint32)
+                print(f"[emu-dump]   wave{wi2} vgpr nonzero={int((v!=0).sum())} accvgpr nonzero={int((av!=0).sum())}", flush=True)
          else: raise RuntimeError("exceeded 10M total scheduling rounds")
          tracing = False  # only trace the first workgroup

+          if _trace_on:
+            import pickle
+            tag = os.environ["EMU_TRACE_INST"]
+            with open(f"/tmp/emu_trace_{tag}.pkl", "wb") as f: pickle.dump(dict(_inst_hist), f)
+            os.environ.pop("EMU_TRACE_INST", None)   # only dump for the first matching dispatch
          # Reset LDS for next workgroup
          if lds_size > 0: ctypes.memset(lds_buf._buf.va_addr, 0, max(lds_size, 4))

@@ -7,7 +7,7 @@ from tinygrad.codegen.decomp.dtype import f2f
 # Type alias for vars dict: stores UOps and tuples for lambda definitions
 VarVal = UOp | tuple[str, list[str], str]

-def _const(dt, v): return UOp.const(v, dt)
+def _const(dt, v): return UOp.const(dt, v)
 def _u32(v): return _const(dtypes.uint32, v)
 def _u64(v): return _const(dtypes.uint64, v)
 def _to_u32(v): return v if v.dtype == dtypes.uint32 else v.bitcast(dtypes.uint32) if v.dtype.itemsize == 4 else v.cast(dtypes.uint32)
@@ -55,8 +55,8 @@ def _expr_bits(v: UOp) -> int:
  if v.op in (Ops.AND, Ops.XOR):
    widths: list[int] = []
    for src in v.src:
-      if src.op == Ops.CONST and isinstance(src.val, int) and src.val > 0 and (src.val & (src.val + 1)) == 0:
-        widths.append(src.val.bit_length())
+      if src.op == Ops.CONST and isinstance(src.arg, int) and src.arg > 0 and (src.arg & (src.arg + 1)) == 0:
+        widths.append(src.arg.bit_length())
    if widths: return max(widths)
  return v.dtype.bitsize

@@ -144,9 +144,9 @@ def _minmax_reduce(is_max: bool, dt, *args: UOp) -> UOp:
 def _find_two_pi_mul(x):
  if x.op != Ops.MUL or len(x.src) != 2: return None
  for i, s in enumerate(x.src):
-    if s.op == Ops.CONST and abs(s.val - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586)
+    if s.op == Ops.CONST and abs(s.arg - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586)
    if s.op == Ops.MUL and len(s.src) == 2:
-      vals = [ss.val for ss in s.src if ss.op == Ops.CONST] + [ss.src[0].val for ss in s.src if ss.op == Ops.CAST and ss.src[0].op == Ops.CONST]
+      vals = [ss.arg for ss in s.src if ss.op == Ops.CONST] + [ss.src[0].arg for ss in s.src if ss.op == Ops.CAST and ss.src[0].op == Ops.CONST]
      if len(vals) == 2 and abs(vals[0] * vals[1] - 6.283185307179586) < 1e-5: return (x.src[1-i], vals[0] * vals[1])
  return None

@@ -163,7 +163,7 @@ def _trig_reduce(x, phase=0.0):

 def _signext(val: UOp) -> UOp:
  for bits, mask, ext in [(4, 0xF, 0xFFFFFFF0), (8, 0xFF, 0xFFFFFF00), (16, 0xFFFF, 0xFFFF0000)]:
-    if (val.op == Ops.AND and len(val.src) == 2 and val.src[1].op == Ops.CONST and val.src[1].val == mask) or val.dtype.itemsize == bits // 8:
+    if (val.op == Ops.AND and len(val.src) == 2 and val.src[1].op == Ops.CONST and val.src[1].arg == mask) or val.dtype.itemsize == bits // 8:
      v32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
      sb = (v32 >> _u32(bits - 1)) & _u32(1)
      return sb.ne(_u32(0)).where(v32 | _u32(ext), v32).cast(dtypes.int)
@@ -497,7 +497,7 @@ class Parser:
        if not dtypes.is_int(right.dtype): right = right.cast(dtypes.uint32)
        return (left >> right) if op == '>>' else (left << right)
      case '+' | '-':
-        if op == '-' and left.op == Ops.CONST and right.op == Ops.CONST: return _const(left.dtype, left.val - right.val)
+        if op == '-' and left.op == Ops.CONST and right.op == Ops.CONST: return _const(left.dtype, left.arg - right.arg)
        return (left + right) if op == '+' else (left - right)
      case '*' | '/':
        # Integer promotion: promote 16-bit integers to 32-bit before multiply to avoid overflow
@@ -507,7 +507,7 @@ class Parser:
          left, right = left.cast(pdt), right.cast(pdt)
        if op == '*': return left * right
        return (left // right) if dtypes.is_int(left.dtype) else (left / right)
-      case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if left.op == Ops.CONST and left.val == 2.0 else left
+      case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if left.op == Ops.CONST and left.arg == 2.0 else left

  _PREC = [('||',), ('&&',), ('|',), ('^',), ('&',), ('==', '!=', '<>'), ('>=', '<=', '>', '<'), ('>>', '<<'), ('+', '-'), ('*', '/'), ('**',)]

@@ -530,7 +530,7 @@ class Parser:
    if self.try_eat_val('-', 'OP'):
      inner = self.unary()
      if inner.op == Ops.CONST:
-        return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -inner.val)
+        return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -inner.arg)
      return inner.neg()
    if self.try_eat_val('+', 'OP'): return self.unary()
    return self.postfix()
@@ -670,14 +670,14 @@ class Parser:
      width = self.parse()
      self.eat('RBRACKET')
      if width.op == Ops.CONST:
-        w = int(width.val)
+        w = int(width.arg)
        return (base >> _to_u32(first)) & _const(base.dtype, (1 << w) - 1)
      return base
    if self.try_eat('COLON'):
      second = self.parse()
      self.eat('RBRACKET')
      if first.op == Ops.CONST and second.op == Ops.CONST:
-        a, b = int(first.val), int(second.val)
+        a, b = int(first.arg), int(second.arg)
        if a < b: return _bitreverse(base, b - a + 1)
        hi, lo = a, b
        if lo >= base.dtype.itemsize * 8:
@@ -699,7 +699,7 @@ class Parser:
    if var_name is None:
      var_name = self._find_var_name(base)
    if first.op == Ops.CONST:
-      idx = int(first.val)
+      idx = int(first.arg)
      # Check for array element (var@idx)
      if var_name and f'{var_name}@{idx}' in self.vars:
        v = self.vars[f'{var_name}@{idx}']
@@ -866,13 +866,13 @@ class Parser:
        idx_hi_native = ((addr + _const(adt, 4)) >> _const(adt, 2)).cast(dtypes.int64)
        safe_idx_hi = is_unaligned.where(idx_hi_native, idx_native)
        hi = mindex(safe_idx_hi)
-        combined = val.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(32, dtypes.uint64))
-        val = is_unaligned.where((combined >> (byte_off.cast(dtypes.uint64) * UOp.const(8, dtypes.uint64))).cast(dtypes.uint32), val)
+        combined = val.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32))
+        val = is_unaligned.where((combined >> (byte_off.cast(dtypes.uint64) * UOp.const(dtypes.uint64, 8))).cast(dtypes.uint32), val)
    return _cast_to(val, dt)

  def _coerce_cmp(self, l: UOp, r: UOp) -> tuple[UOp, UOp]:
    if l.dtype != r.dtype:
-      if r.dtype == dtypes.int and r.op == Ops.CONST and r.val < 0: l = l.cast(dtypes.int)
+      if r.dtype == dtypes.int and r.op == Ops.CONST and r.arg < 0: l = l.cast(dtypes.int)
      else: r = r.cast(l.dtype)
    return l, r

@@ -970,7 +970,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
        if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl'))
        expr = p.parse().simplify()
        assert expr.op == Ops.CONST, f"loop bound must be constant, got {expr}"
-        return int(expr.val)
+        return int(expr.arg)
      start_val = parse_bound()
      p.eat('COLON')
      end_val = parse_bound()
@@ -1258,7 +1258,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
      def parse_cond(s, kw):
        ll = s.lower()
        return _to_bool(parse_expr(s[ll.find(kw) + len(kw):ll.rfind('then')].strip(), env, funcs))
-      def is_const(c, v): return c.op == Ops.CONST and c.val is v
+      def is_const(c, v): return c.op == Ops.CONST and c.arg is v
      cond = parse_cond(line, 'if')
      conditions: list[tuple[UOp, UOp | dict[str, VarVal] | None]] = [(cond, None)] if not is_const(cond, False) else []
      branch_assigns: list[tuple[UOp, list]] = []  # (cond, assigns_list) for side-effect merging
@@ -1339,7 +1339,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
              return (dest, (val[0], cnd.where(val[1], val[1])))
            return (dest, val)
          # Build combined condition: each branch fires when its cond is true AND no earlier cond was true
-          remaining = UOp.const(True)
+          remaining = UOp.const(dtypes.bool, True)
          for bc, bse in branch_assigns:
            effective = remaining & bc if remaining.op != Ops.CONST else bc
            for dest, val in bse: assigns.append(_cond_side_effect(effective, dest, val))
@@ -74,7 +74,6 @@ class TestWhisper(unittest.TestCase):
        err
      )

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

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

-  @slow
  def test_transcribe_batch21(self):
    waveforms = [load_file_waveform(TEST_FILE_2), load_file_waveform(TEST_FILE_1)]
    transcriptions = transcribe_waveform(self.model, self.enc, waveforms)
@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Tensor, dtypes, TinyJit, UOp
 from tinygrad.llm.model import apply_rope as apply_rope_new, precompute_freqs_cis
-from test.helpers import assert_jit_cache_len, check_schedule
+from test.helpers import assert_jit_cache_len

 def apply_rope(x:Tensor, start_pos:int):
  B, H, T, Hd = x.shape
@@ -16,8 +16,9 @@ class TestAttention(unittest.TestCase):
    k = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize()
    v = Tensor.ones(BS, seqlen, dim, dtype=dtypes.half).contiguous().realize()
    attn = q.scaled_dot_product_attention(k, v)
+    sched = attn.schedule_linear()
    # attention has 4 kernels now
-    check_schedule(attn, 4)
+    self.assertEqual(len(sched.src), 4)

  def test_apply_rope_jit_prune(self):
    def rope_fn(x_in, pos): return apply_rope(x_in, pos)
@@ -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, Invalid
 from tinygrad.uop.ops import Ops, UOp
 from test.helpers import full_rewrite
 import numpy as np
@@ -36,28 +36,21 @@ class TestUnaryOpsConstFolding(unittest.TestCase):

 class TestWeakConstFolding(unittest.TestCase):
  def test_weakint_math(self):
-    out = (UOp.const(2**40) + UOp.const(2**40)).simplify()
-    self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.weakint, 2**41))
+    out = (UOp.const(dtypes.weakint, 2**40) + UOp.const(dtypes.weakint, 2**40)).simplify()
+    self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakint, 2**41))

  def test_float_unaries(self):
-    for op in (Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL):
-      out = UOp.const(4.0).alu(op).simplify()
-      self.assertEqual((out.op, out.dtype), (Ops.CONST, dtypes.weakfloat))
+    for dtype in (dtypes.weakfloat,):
+      for op in (Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL):
+        out = UOp.const(dtype, 4).alu(op).simplify()
+        self.assertEqual((out.op, out.dtype), (Ops.CONST, dtypes.weakfloat))

  def test_weakfloat_math(self):
-    out = (UOp.const(1.25) + UOp.const(2.5)).simplify()
-    self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.weakfloat, 3.75))
+    out = (UOp.const(dtypes.weakfloat, 1.25) + UOp.const(dtypes.weakfloat, 2.5)).simplify()
+    self.assertEqual((out.op, out.dtype, out.arg), (Ops.CONST, dtypes.weakfloat, 3.75))

  def test_invalid_poison(self):
-    self.assertTrue(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().is_invalid)
-
-  def test_single_rounding_log10_backward(self):
-    # log10 backward folds log10(2)/log(2) = 1/log(10) in one rounding, not the double-rounded 1/float32(log(10))
-    x = Tensor([1.0, 2.0, 3.0])
-    ast = next(s.src[0] for s in x.log10().sum().gradient(x)[0].schedule_linear().src if s.src[0].op is Ops.SINK)
-    const = next(u.arg for u in full_rewrite(ast).toposort() if u.op is Ops.CONST and u.dtype is dtypes.float32)
-    # correctly rounded: within half a float32 ulp of the exact value (folding at float32 lands 0.66 ulp off)
-    self.assertLess(abs(const - 1/math.log(10)), 2**-26)
+    self.assertIs(UOp.invalid().alu(Ops.CDIV, UOp.const(dtypes.weakint, 0)).simplify().arg, Invalid)

 class TestBinaryOpsConstFolding(unittest.TestCase):
  def test_add_literal_zero(self):
@@ -128,10 +121,10 @@ class TestBitcastConstFolding(unittest.TestCase):
    def t(cases: dict[DType, ConstType]):
      for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()):
        if not math.isnan(from_v):
-          r = full_rewrite(UOp.const(from_v, from_dt).bitcast(to_dt).sink()).src[0]
+          r = full_rewrite(UOp.const(from_dt, from_v).bitcast(to_dt).sink()).src[0]
          self.assertEqual(r.op, Ops.CONST, msg:=f"{from_dt} -> {to_dt} ({from_v} -> {to_v})")
          self.assertEqual(r.dtype, to_dt, msg)
-          np.testing.assert_equal(r.val, to_v, msg)
+          np.testing.assert_equal(r.arg, to_v, msg)

    t({dtypes.int8: 0, dtypes.uint8: 0, dtypes.bool: False})
    t({dtypes.int8: 1, dtypes.uint8: 1, dtypes.bool: True})
@@ -152,9 +145,9 @@ class TestBitcastConstFolding(unittest.TestCase):

  def test_vec_bitcast(self):
    with Context(SPEC=0):
-      srcs = full_rewrite(UOp.const((-1, -2**31, 75), dtypes.int32).bitcast(dtypes.uint32).sink()).src
+      srcs = full_rewrite(UOp.const(dtypes.int32, (-1, -2**31, 75)).bitcast(dtypes.uint32).sink()).src
    self.assertTrue(all(r.op is Ops.CONST and r.dtype == dtypes.uint32 for r in srcs))
-    self.assertEqual(tuple(x.val for x in srcs), (2**32-1, 2**31, 75))
+    self.assertEqual(tuple(x.arg for x in srcs), (2**32-1, 2**31, 75))

 # folds advance indexing into basic indexing
 class TestIndexingConstFolding(unittest.TestCase):
@@ -12,7 +12,7 @@ class TestGroupedDims(unittest.TestCase):
    idxs = get_grouped_dims(prefix, dims, max_sizes, reverse)
    loop_idxs = dedup(flatten([[y for y in x.toposort() if y.op is Ops.SPECIAL] for x in idxs]))
    loop_idxs = sorted(loop_idxs, key=lambda uop: uop.arg)
-    sizes = [x.src[0].val for x in loop_idxs]
+    sizes = [x.src[0].arg for x in loop_idxs]
    assert len(idxs) == len(dims), f"expected idxs to have same length as dims {len(dims)}, got {len(idxs)}"
    if assert_same_length:
      assert len(loop_idxs) == min(len(sizes), len(dims)), f"expected idxs to have length {min(len(sizes), len(dims))}, got {len(loop_idxs)}"
@@ -24,7 +24,7 @@ class TestGroupedDims(unittest.TestCase):
    total = math.prod(dims)
    specials = sorted(dedup(flatten([[y for y in x.toposort() if y.op is Ops.SPECIAL] for x in idxs])), key=lambda u: u.arg)
    # build flat index and primed flat (same expression with renamed SPECIALs)
-    flat = UOp.const(0)
+    flat = UOp.const(dtypes.weakint, 0)
    for i, idx in enumerate(idxs):
      flat = flat + idx * int(math.prod(dims[i+1:]))
    flat_p = flat.substitute({s: UOp(Ops.SPECIAL, src=s.src, arg=s.arg+"_p") for s in specials})
@@ -107,7 +107,7 @@ class TestGroupedDims(unittest.TestCase):

  def test_global_prod_max(self):
    g, l = UOp.range(256, 0, AxisType.GLOBAL), UOp.range(256, 1, AxisType.LOCAL)
-    sink = UOp.param(0, dtypes.float, (512,)).index(g + l).store(UOp.const(1.0)).end(g, l).sink(arg=KernelInfo())
+    sink = UOp.param(0, dtypes.float, (512,)).index(g + l).store(UOp.const(dtypes.float, 1.0)).end(g, l).sink(arg=KernelInfo())
    class R(Renderer): global_max, local_max, global_prod_max = (256, 256, 256), (128, 128, 128), (128, 128, 128)
    specials = [u for u in add_gpudims(R(Target()), sink).toposort() if u.op is Ops.SPECIAL]
    self.assertGreater(len([s for s in specials if "lidx" in s.arg]), 1)
@@ -14,10 +14,10 @@ class TestGradient(unittest.TestCase):
  def _test_one_input_function(self, f:Callable, jf:Callable|None=None):
    if jf is None: jf = f
    x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float)
-    gx = compute_gradient(f(x), UOp.const(1.0), set([x]))[x]
+    gx = compute_gradient(f(x), UOp.const(dtypes.float, 1.0), set([x]))[x]

    for val in [-5., -2.0, 0.0, 2.0, 5.]:
-      tg_out = gx.substitute({x: UOp.const(val)}).ssimplify()
+      tg_out = gx.substitute({x: x.const_like(val)}).ssimplify()
      tx = torch.tensor([val], dtype=torch.float, requires_grad=True)
      torch_out = torch.autograd.grad(jf(tx), tx)[0].item()
      self._cmp_nan_okay(tg_out, torch_out)
@@ -26,13 +26,13 @@ class TestGradient(unittest.TestCase):
    if jf is None: jf = f
    x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float)
    y = UOp.variable('y', -math.inf, math.inf, dtype=dtypes.float)
-    grads = compute_gradient(f(x, y), UOp.const(1.0), set([x, y]))
+    grads = compute_gradient(f(x, y), UOp.const(dtypes.float, 1.0), set([x, y]))
    gx, gy = grads[x], grads[y]

    for valx in [-5., -2.0, 0.0, 2.0, 5.]:
      for valy in [-5., -2.0, 0.0, 2.0, 5.]:
        # Substitute the values into the gradient expressions
-        substitutions = {x: UOp.const(valx), y: UOp.const(valy)}
+        substitutions = {x: x.const_like(valx), y: y.const_like(valy)}
        tg_out_x = gx.substitute(substitutions).ssimplify()
        tg_out_y = gy.substitute(substitutions).ssimplify()

@@ -15,13 +15,13 @@ def apply_rewrite(expr):
 def apply_rewrite_values(expr):
  srcs = full_rewrite(expr.sink()).src
  if len(srcs) == 1:
-    if srcs[0].op is Ops.CONST: return (srcs[0].val,)
-    if srcs[0].op is Ops.STACK: return tuple(s.val for s in srcs[0].src)
-  return tuple(s.val for s in srcs)
+    if srcs[0].op is Ops.CONST: return (srcs[0].arg,) if not isinstance(srcs[0].arg, tuple) else srcs[0].arg
+    if srcs[0].op is Ops.STACK: return tuple(s.arg for s in srcs[0].src)
+  return tuple(s.arg for s in srcs)

 def evaluate_uop(uop, variables):
  if uop.op == Ops.CONST:
-    return uop.val
+    return uop.arg
  elif uop.op == Ops.PARAM and uop.arg.addrspace is AddrSpace.ALU:
    return variables[uop.expr]
  elif uop.op in GroupOp.ALU:
@@ -32,67 +32,67 @@ def evaluate_uop(uop, variables):

 class TestArithmeticSimplifications(unittest.TestCase):
  def test_full_graph_rewrite_division_by_zero(self):
-    optimized_div_uop = apply_rewrite(UOp.const(10.0) / UOp.const(0.0))
+    optimized_div_uop = apply_rewrite(UOp.const(dtypes.float32, 10.0) / UOp.const(dtypes.float32, 0.0))
    self.assertEqual(optimized_div_uop.op, Ops.CONST)
-    self.assertTrue(math.isinf(optimized_div_uop.val) or math.isnan(optimized_div_uop.val))
+    self.assertTrue(math.isinf(optimized_div_uop.arg) or math.isnan(optimized_div_uop.arg))

  def test_full_graph_rewrite_redundant_operations(self):
-    optimized_uop = apply_rewrite((UOp.const(10.0) + UOp.const(0.0)) * UOp.const(1.0))
+    optimized_uop = apply_rewrite((UOp.const(dtypes.float32, 10.0) + UOp.const(dtypes.float32, 0.0)) * UOp.const(dtypes.float32, 1.0))
    self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.val, 10.0)
+    self.assertEqual(optimized_uop.arg, 10.0)

  def test_full_graph_rewrite_large_graph(self):
-    prev_uop = UOp.const(0)
+    prev_uop = UOp.const(dtypes.int32, 0)
    for i in range(1, 101):
-      prev_uop += UOp.const(i)
+      prev_uop += UOp.const(dtypes.int32, i)
    optimized_uop = apply_rewrite(prev_uop)
    self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.val, sum(range(1, 101)))
+    self.assertEqual(optimized_uop.arg, sum(range(1, 101)))

  def test_full_graph_rewrite_division_by_one(self):
-    optimized_uop = apply_rewrite(UOp.const(42.0) / UOp.const(1.0))
+    optimized_uop = apply_rewrite(UOp.const(dtypes.float32, 42.0) / UOp.const(dtypes.float32, 1.0))
    self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.val, 42.0)
+    self.assertEqual(optimized_uop.arg, 42.0)

  def test_full_graph_rewrite_modulo_by_one(self):
-    optimized_uop = apply_rewrite(UOp.const(42) % UOp.const(1))
+    optimized_uop = apply_rewrite(UOp.const(dtypes.int32, 42) % UOp.const(dtypes.int32, 1))
    self.assertEqual(optimized_uop.op, Ops.CONST)
-    self.assertEqual(optimized_uop.val, 0)
+    self.assertEqual(optimized_uop.arg, 0)


 class TestFoldingAndReduction(unittest.TestCase):
  @unittest.skip("reduce is removed now")
  def test_full_graph_rewrite_constant_reduction_folding(self):
-    const1 = UOp.const(5)
-    const2 = UOp.const(10)
-    const3 = UOp.const(20)
+    const1 = UOp.const(dtypes.int32, 5)
+    const2 = UOp.const(dtypes.int32, 10)
+    const3 = UOp.const(dtypes.int32, 20)
    optimized_sink = apply_rewrite((const1 + const2 + const3).reduce(Ops.ADD))
    expected_sum = 5 + 10 + 20
-    self.assertEqual(optimized_sink.val, expected_sum)
+    self.assertEqual(optimized_sink.arg, expected_sum)

  @unittest.skip("reduce is removed now")
  def test_full_graph_rewrite_reduction_with_unused_range(self):
-    const1 = UOp.const(15)
-    const2 = UOp.const(25)
+    const1 = UOp.const(dtypes.int32, 15)
+    const2 = UOp.const(dtypes.int32, 25)
    rng = UOp.range(10, idx=0)
    optimized_sink = apply_rewrite((const1 + const2).reduce(Ops.ADD, rng))
    expected_sum = 10 * (15 + 25)
-    self.assertEqual(optimized_sink.val, expected_sum)
+    self.assertEqual(optimized_sink.arg, expected_sum)

  @unittest.skip("currently failing")
  def test_full_graph_rewrite_range_reduction(self):
    simple_range = UOp.range(5, idx=0)
    optimized_sink = apply_rewrite(simple_range.reduce(Ops.ADD, simple_range))
    expected_sum = sum(range(5))
-    self.assertEqual(optimized_sink.val, expected_sum)
+    self.assertEqual(optimized_sink.arg, expected_sum)

  @unittest.skip("currently failing")
  def test_full_graph_rewrite_simple_reduction_folding(self):
    simple_range = UOp.range(4, idx=0)
-    add_uop = simple_range + UOp.const(1)
+    add_uop = simple_range + UOp.const(dtypes.int32, 1)
    optimized_sink = apply_rewrite(add_uop.reduce(Ops.ADD, simple_range))
    expected_sum = sum(i + 1 for i in range(4))
-    self.assertEqual(optimized_sink.val, expected_sum)
+    self.assertEqual(optimized_sink.arg, expected_sum)

  @unittest.skip("currently failing")
  def test_full_graph_rewrite_nested_loop_collapse(self):
@@ -101,7 +101,7 @@ class TestFoldingAndReduction(unittest.TestCase):
    expr = (outer_range * 10) + inner_range
    optimized_reduce_uop = apply_rewrite(expr.reduce(Ops.ADD, outer_range, inner_range))
    self.assertEqual(optimized_reduce_uop.op, Ops.CONST)
-    self.assertEqual(optimized_reduce_uop.val, sum((i * 10) + j for i in range(8) for j in range(4)))
+    self.assertEqual(optimized_reduce_uop.arg, sum((i * 10) + j for i in range(8) for j in range(4)))


 class TestModuloAndDivisionFolding(unittest.TestCase):
@@ -110,27 +110,27 @@ class TestModuloAndDivisionFolding(unittest.TestCase):
    x_var_uop = UOp.variable('x', 0, 100).cast(dtypes.weakint)
    optimized_mod_uop = apply_rewrite(((x_var_uop * 4) + 2) % 4)
    self.assertEqual(optimized_mod_uop.op, Ops.CONST)
-    self.assertEqual(optimized_mod_uop.val, 2)
+    self.assertEqual(optimized_mod_uop.arg, 2)

  def test_full_graph_rewrite_division_folding_with_define_var(self):
    # index dtype because div-mod rules only work on index
    n_var_uop = UOp.variable('n', 1, 1000).cast(dtypes.weakint)
    optimized_div_uop = apply_rewrite((n_var_uop * 6) // 3)
    self.assertEqual(optimized_div_uop.op, Ops.MUL)
-    self.assertEqual(optimized_div_uop.src[1].val, 2)
+    self.assertEqual(optimized_div_uop.src[1].arg, 2)

  def test_full_graph_rewrite_complex_mod_div_folding(self):
    # index dtype because div-mod rules only work on index
    k_var_uop = UOp.variable('k', 0, 50).cast(dtypes.weakint)
    optimized_div_uop = apply_rewrite(((k_var_uop * 12 + 8) % 6) // 2)
    self.assertEqual(optimized_div_uop.op, Ops.CONST)
-    self.assertEqual(optimized_div_uop.val, 1)
+    self.assertEqual(optimized_div_uop.arg, 1)

  def test_graph_rewrite_div_folding_bug(self):
    lhs = UOp(Ops.ADD, src=(
-      UOp(Ops.STACK, arg=None, src=(UOp(Ops.SPECIAL, src=(UOp.const(32),), arg='lidx0'),)*4),
-      UOp.const((0, 256, 512, 768))))
-    rhs = UOp.const((2,)*4)
+      UOp(Ops.STACK, arg=None, src=(UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 32),), arg='lidx0'),)*4),
+      UOp.const(dtypes.int, (0, 256, 512, 768))))
+    rhs = UOp.const(dtypes.int, (2,)*4)
    unopt = lhs<rhs
    opt = apply_rewrite(unopt)
    print(unopt)
@@ -143,13 +143,13 @@ class TestModuloAndDivisionFolding(unittest.TestCase):
    self.assertIs(apply_rewrite(x_var_uop.cast(dtypes.weakint) % 10).render(simplify=False), x_var_uop.render(simplify=False))

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

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

 class TestEdgeCasesAndSpecialOperations(unittest.TestCase):
  def test_full_graph_rewrite_transcendental_edge_cases(self):
-    optimized_sink = full_rewrite(UOp.const(-1.0).log2().sink(UOp.const(0.0).reciprocal()))
+    optimized_sink = full_rewrite(UOp.const(dtypes.float32, -1.0).log2().sink(UOp.const(dtypes.float32, 0.0).reciprocal()))
    optimized_log2_neg, optimized_recip_zero = optimized_sink.src
-    self.assertTrue(math.isnan(optimized_log2_neg.val), f"Expected NaN for log2(-1.0), got {optimized_log2_neg.val}")
-    self.assertTrue(math.isinf(optimized_recip_zero.val) and optimized_recip_zero.val > 0,
-                    f"Expected +inf for reciprocal(0.0), got {optimized_recip_zero.val}")
+    self.assertTrue(math.isnan(optimized_log2_neg.arg), f"Expected NaN for log2(-1.0), got {optimized_log2_neg.arg}")
+    self.assertTrue(math.isinf(optimized_recip_zero.arg) and optimized_recip_zero.arg > 0,
+                    f"Expected +inf for reciprocal(0.0), got {optimized_recip_zero.arg}")

  @unittest.skip("broken")
  def test_full_graph_rewrite_modulo_negative_dividend(self):
@@ -182,45 +182,45 @@ class TestEdgeCasesAndSpecialOperations(unittest.TestCase):
 class TestGEPAndVectorizeRewrite(unittest.TestCase):
  def test_gep_single_element_extraction(self):
    # GEP on a vector dtype to extract a single element
-    base_vector = UOp.const((1.0, 2.0, 3.0, 4.0))
-    self.assertEqual(apply_rewrite(base_vector.index(2)).val, 3.0)
+    base_vector = UOp.const(dtypes.float32, (1.0, 2.0, 3.0, 4.0))
+    self.assertEqual(apply_rewrite(base_vector.index(2)).arg, 3.0)

  def test_gep_tuple_extraction(self):
    # GEP on a vector dtype to extract multiple elements as a vector
-    base_vector = UOp.const((1.0, 2.0, 3.0, 4.0))
+    base_vector = UOp.const(dtypes.float32, (1.0, 2.0, 3.0, 4.0))
    self.assertEqual(list(apply_rewrite_values(UOp.stack(*[base_vector.index(i) for i in (2, 3)]))), [3.0, 4.0])

  def test_gep_on_const_stack(self):
    # GEP on a const STACK to extract a single element
-    const_stack = UOp.const((1.0, 2.0, 3.0, 4.0))
-    self.assertEqual(apply_rewrite(const_stack.index(2)).val, 3.0)
+    const_stack = UOp.const(dtypes.float32, (1.0, 2.0, 3.0, 4.0))
+    self.assertEqual(apply_rewrite(const_stack.index(2)).arg, 3.0)

  def test_gep_tuple_on_const_stack(self):
    # GEP on a const STACK using a tuple to extract multiple elements
-    const_stack = UOp.const((7.0, 8.0, 9.0, 10.0))
+    const_stack = UOp.const(dtypes.float32, (7.0, 8.0, 9.0, 10.0))
    self.assertEqual(list(apply_rewrite_values(UOp.stack(*[const_stack.index(i) for i in (1, 3)]))), [8.0, 10.0])

  def test_vectorize_multiple_elements(self):
    # Vectorizing multiple elements using GEP
-    base_vector = UOp.const((5.0, 10.0, 15.0, 20.0))
+    base_vector = UOp.const(dtypes.float32, (5.0, 10.0, 15.0, 20.0))
    vectorized_uop = UOp(Ops.STACK, src=tuple(base_vector.index(i) for i in range(4)))
    self.assertEqual(list(apply_rewrite_values(vectorized_uop)), [5.0, 10.0, 15.0, 20.0])


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

 class TestBottomUpRewrite(unittest.TestCase):
  def test_const_folding(self):
-    a = UOp.const(5)
+    a = UOp.const(dtypes.int, 5)
    ret = (a*3) + (a*7)
    gt = graph_rewrite(ret, symbolic_simple)
    ret = graph_rewrite(ret, symbolic_simple, bottom_up=True)
    self.assertIs(gt, ret)

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

@@ -305,28 +305,28 @@ class TestRecurse(unittest.TestCase):
    graph_rewrite(a, pm, bottom_up=True)

  def test_inf_loop(self):
-    a = UOp.const(3)
+    a = UOp.const(dtypes.int, 3)
    pm = PatternMatcher([
-      (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)),
-      (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)),
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
    ])
    with self.assertRaises(RuntimeError):
      graph_rewrite(a, pm)

  def test_inf_loop_bottom_up(self):
-    a = UOp.const(3)
+    a = UOp.const(dtypes.int, 3)
    pm = PatternMatcher([
-      (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)),
-      (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)),
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
    ])
    with self.assertRaises(RuntimeError):
      graph_rewrite(a, pm, bottom_up=True)

-def bidir_append(ctx, x, b): ctx.append((x.val if x.op is Ops.CONST else "+", b))
+def bidir_append(ctx, x, b): ctx.append((x.arg if x.op is Ops.CONST else "+", b))
 class TestBidirectional(unittest.TestCase):
  def test_simple(self):
-    a = UOp.const(1)
-    b = UOp.const(2)
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
    c = a + b
    pm = PatternMatcher([ (UPat(GroupOp.All, name="x"), lambda ctx,x: bidir_append(ctx, x, False)) ])
    bpm = PatternMatcher([ (UPat(GroupOp.All, name="x"), lambda ctx,x: bidir_append(ctx, x, True)) ])
@@ -336,14 +336,14 @@ class TestBidirectional(unittest.TestCase):

 class TestStopEarly(unittest.TestCase):
  def test_stop_early(self):
-    a = UOp.const(3)
-    b = UOp.const(4)
+    a = UOp.const(dtypes.int, 3)
+    b = UOp.const(dtypes.int, 4)
    c = a+b
-    cn = UOp.const(7)
-    d = UOp.const(2)
+    cn = UOp.const(dtypes.int, 7)
+    d = UOp.const(dtypes.int, 2)
    def visit_const(c:UOp):
-      print(f"visit {c.val}")
-      assert c.val not in (3,4)
+      print(f"visit {c.arg}")
+      assert c.arg not in (3,4)
    pm_cvisit = PatternMatcher([(UPat(Ops.CONST, name="c"), visit_const),])
    ret = (c+d).substitute({c:cn}, extra_pm=pm_cvisit)
    assert ret == cn+d
@@ -376,15 +376,15 @@ class TestWalkRewrite(unittest.TestCase):

  def test_walk_topdown_no_fixed_point(self):
    """A bouncing pattern applies once and stops instead of looping."""
-    a = UOp.const(3)
+    a = UOp.const(dtypes.int, 3)
    pm = PatternMatcher([
-      (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)),
-      (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)),
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
    ])
    with self.assertRaises(RuntimeError):
      graph_rewrite(a, pm, bottom_up=True)
    ret = graph_rewrite(a, pm, walk=True)
-    self.assertIs(ret, UOp.const(4))
+    self.assertIs(ret, UOp.const(dtypes.int, 4))

  def test_walk_topdown_rewrites_children(self):
    a = UOp.variable('a', 0, 10)
@@ -418,11 +418,11 @@ class TestWalkRewrite(unittest.TestCase):
    """Top-down walk fires pm after children are processed (post-order)."""
    visited = []
    def track_visit(ctx, x):
-      ctx.append(x.val if x.op is Ops.CONST else x.op)
+      ctx.append(x.arg if x.op is Ops.CONST else x.op)
      return None
    pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)])
-    a = UOp.const(1)
-    b = UOp.const(2)
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
    graph_rewrite(a + b, pm, ctx=visited, walk=True)
    self.assertEqual(visited, [1, 2, Ops.ADD])

@@ -454,23 +454,23 @@ class TestWalkRewrite(unittest.TestCase):

  def test_walk_bottomup_no_fixed_point(self):
    """Bottom-up walk also applies once per node, no fixed-point iteration."""
-    a = UOp.const(3)
+    a = UOp.const(dtypes.int, 3)
    pm = PatternMatcher([
-      (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)),
-      (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)),
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
    ])
    ret = graph_rewrite(a, pm, bottom_up=True, walk=True)
-    self.assertIs(ret, UOp.const(4))
+    self.assertIs(ret, UOp.const(dtypes.int, 4))

  def test_walk_bottomup_visit_order(self):
    """Bottom-up walk fires bpm before descending (pre-order)."""
    visited = []
    def track_visit(ctx, x):
-      ctx.append(x.val if x.op is Ops.CONST else x.op)
+      ctx.append(x.arg if x.op is Ops.CONST else x.op)
      return None
    pm = PatternMatcher([(UPat(GroupOp.All, name="x"), track_visit)])
-    a = UOp.const(1)
-    b = UOp.const(2)
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
    graph_rewrite(a + b, pm, ctx=visited, bottom_up=True, walk=True)
    # bpm fires on each node before children: +, 1, 2
    self.assertEqual(visited, [Ops.ADD, 1, 2])
@@ -490,15 +490,15 @@ class TestWalkRewrite(unittest.TestCase):
    """Bidirectional walk: bpm fires pre-order, pm fires post-order."""
    visited = []
    def bpm_visit(ctx, x):
-      ctx.append((x.val if x.op is Ops.CONST else x.op, "bpm"))
+      ctx.append((x.arg if x.op is Ops.CONST else x.op, "bpm"))
      return None
    def pm_visit(ctx, x):
-      ctx.append((x.val if x.op is Ops.CONST else x.op, "pm"))
+      ctx.append((x.arg if x.op is Ops.CONST else x.op, "pm"))
      return None
    bpm = PatternMatcher([(UPat(GroupOp.All, name="x"), bpm_visit)])
    pm = PatternMatcher([(UPat(GroupOp.All, name="x"), pm_visit)])
-    a = UOp.const(1)
-    b = UOp.const(2)
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
    graph_rewrite(a + b, pm, ctx=visited, bpm=bpm, walk=True)
    # bpm fires pre-order, pm fires post-order
    self.assertEqual(visited, [
@@ -509,23 +509,23 @@ class TestWalkRewrite(unittest.TestCase):
    """If bpm matches, children are skipped and pm never fires on that node."""
    visited = []
    def bpm_match(ctx, x):
-      ctx.append((x.val if x.op is Ops.CONST else x.op, "bpm"))
+      ctx.append((x.arg if x.op is Ops.CONST else x.op, "bpm"))
      # rewrite const(1) -> const(10), short-circuiting its subtree
-      if x.op is Ops.CONST and x.val == 1: return UOp.const(10, x.dtype)
+      if x.op is Ops.CONST and x.arg == 1: return x.replace(arg=10)
      return None
    def pm_match(ctx, x):
-      ctx.append((x.val if x.op is Ops.CONST else x.op, "pm"))
+      ctx.append((x.arg if x.op is Ops.CONST else x.op, "pm"))
      return None
    bpm = PatternMatcher([(UPat(GroupOp.All, name="x"), bpm_match)])
    pm = PatternMatcher([(UPat(GroupOp.All, name="x"), pm_match)])
-    a = UOp.const(1)
-    b = UOp.const(2)
+    a = UOp.const(dtypes.int, 1)
+    b = UOp.const(dtypes.int, 2)
    ret = graph_rewrite(a + b, pm, ctx=visited, bpm=bpm, walk=True)
    # bpm matches const(1) and short-circuits it, so pm never fires on const(1)
    self.assertNotIn((1, "pm"), visited)
    # but pm still fires on const(2) and the rebuilt ADD
    self.assertIn((2, "pm"), visited)
-    self.assertIs(ret, UOp.const(10) + b)
+    self.assertIs(ret, UOp.const(dtypes.int, 10) + b)

 if __name__ == '__main__':
  unittest.main()
@@ -297,10 +297,10 @@ class TestPolyN(unittest.TestCase):
    from tinygrad.dtype import dtypes
    from tinygrad.uop.ops import UOp
    from test.helpers import eval_uop
-    np.testing.assert_allclose(eval_uop(polyN(UOp.const(1.0).cast(dtypes.float), [1.0, -2.0, 1.0])), 0.0)
-    np.testing.assert_allclose(eval_uop(polyN(UOp.const(2.0).cast(dtypes.float), [1.0, -2.0, 1.0])), 1.0)
-    np.testing.assert_allclose(eval_uop(polyN(UOp.const(3.0).cast(dtypes.float), [1.0, -2.0, 1.0])), 4.0)
-    np.testing.assert_allclose(eval_uop(polyN(UOp.const(4.0).cast(dtypes.float), [1.0, -2.0, 1.0])), 9.0)
+    np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 1.0), [1.0, -2.0, 1.0])), 0.0)
+    np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 2.0), [1.0, -2.0, 1.0])), 1.0)
+    np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 3.0), [1.0, -2.0, 1.0])), 4.0)
+    np.testing.assert_allclose(eval_uop(polyN(UOp.const(dtypes.float, 4.0), [1.0, -2.0, 1.0])), 9.0)

 class TestTimeToStr(unittest.TestCase):
  def test_seconds(self):           self.assertEqual("   10.01s ", time_to_str(10.01))
@@ -8,15 +8,15 @@ from tinygrad.codegen import to_program
 class TestLinearizerFailures(unittest.TestCase):
  def test_fail_1(self):
    c0 = UOp.param(0, dtypes.float, (64,))
-    c1 = UOp.range(UOp.const(2), 1, AxisType.WEAK)
-    c2 = UOp.range(UOp.const(32), 2, AxisType.WEAK)
-    c3 = ((c1*UOp.const(32))+c2)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 2), 1, AxisType.WEAK)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 32), 2, AxisType.WEAK)
+    c3 = ((c1*UOp.const(dtypes.weakint, 32))+c2)
    c4 = UOp.param(1, dtypes.float, (163840,))
-    c5 = UOp.range(UOp.const(2560), 0, AxisType.REDUCE)
-    c6 = c4.index(((((((c5//UOp.const(8))%UOp.const(8))*UOp.const(8))+(c5%UOp.const(8)))+(((c2*UOp.const(40))+(c5//UOp.const(64)))*UOp.const(64)))+(c1*UOp.const(81920))))
+    c5 = UOp.range(UOp.const(dtypes.weakint, 2560), 0, AxisType.REDUCE)
+    c6 = c4.index(((((((c5//UOp.const(dtypes.weakint, 8))%UOp.const(dtypes.weakint, 8))*UOp.const(dtypes.weakint, 8))+(c5%UOp.const(dtypes.weakint, 8)))+(((c2*UOp.const(dtypes.weakint, 40))+(c5//UOp.const(dtypes.weakint, 64)))*UOp.const(dtypes.weakint, 64)))+(c1*UOp.const(dtypes.weakint, 81920))))
    c7 = UOp.param(2, dtypes.float, (64,))
    c8 = c7.index(c3)
-    c9 = ((((c6+(c8*UOp.const(-1.0)))*(c6+(c8*UOp.const(-1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(0.000390625))+UOp.const(1e-05)).sqrt().reciprocal()
+    c9 = ((((c6+(c8*UOp.const(dtypes.float, -1.0)))*(c6+(c8*UOp.const(dtypes.float, -1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(dtypes.float, 0.000390625))+UOp.const(dtypes.float, 1e-05)).sqrt().reciprocal()
    c10 = c0.index(c3).store(c9).end(c1, c2)
    ast = c10.sink(arg=KernelInfo())
    to_program(ast, renderer=Device[Device.DEFAULT].renderer)
@@ -109,28 +109,6 @@ class TestLLMServer(unittest.TestCase):

    self.assertGreater(len(contents), 0)

-  def test_interrupted_stream_logs_tokens(self):
-    with patch.object(self.mock_model, "generate", side_effect=lambda ids, **kwargs: iter([300, 301, 999])), \
-         patch("tinygrad.llm.serve.stderr_log") as log, patch("tinygrad.llm.serve.colored", side_effect=lambda text, color: text) as color:
-      stream = self.server.RequestHandlerClass.run_model(Mock(server=self.server), [200, 201, 202], "test")
-      next(stream)
-      next(stream)
-      stream.close()
-    interrupt = log.call_args.args[0]
-    self.assertFalse(interrupt.startswith("\n"))
-    self.assertTrue(interrupt.endswith("\n"))
-    self.assertIn("gen:", interrupt)
-    self.assertIn("out:    1", interrupt)
-    self.assertTrue(any(args[0].startswith("total:") and args[1] == "red" for args, _ in color.call_args_list))
-
-  def test_stream_disconnect_closes_source(self):
-    from tinygrad.llm.serve import Handler
-    source, handler = Mock(), Mock()
-    source.__iter__ = Mock(return_value=iter([{}]))
-    handler.wfile.write.side_effect = BrokenPipeError
-    Handler.stream_json(handler, source)
-    source.close.assert_called_once()
-
  def test_non_streaming(self):
    resp = self.client.chat.completions.create(
      model="test-model",
@@ -27,15 +27,10 @@ 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 (planned:=_get_planned_view(new)) is not None: return planned[0]
+      if orig is buf and new.op is Ops.SLICE: return new.src[0]
  return None

 def check_assign(buffer_lists, copies=None):
@@ -46,8 +41,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 (planned:=_get_planned_view(new)) is not None and id(orig) not in replace_map:
-        replace_map[id(orig)] = planned
+      if new.op is Ops.SLICE and id(orig) not in replace_map:
+        replace_map[id(orig)] = (new.src[0], new.src[1].arg * new.src[0].dtype.itemsize, new.arg * new.dtype.itemsize)

  # verify pinned buffers are not planned
  for buf in held_bufs:
@@ -1,5 +1,5 @@
 import unittest, time
-from tinygrad import Tensor, UOp, getenv
+from tinygrad import dtypes, Tensor, UOp, getenv
 from tinygrad.helpers import Profiling

 PYPROFILE = getenv("PYPROFILE")
@@ -27,29 +27,29 @@ class TestBench(unittest.TestCase):
    print(f"{self._testMethodName:30s} {et*1e6/self.N:.2f} us")

  def test_uop_instant_creation(self):
-    for i in range(self.N): UOp.const(100+i)
+    for i in range(self.N): UOp.const(dtypes.int, 100+i)

  def test_uop_list_creation(self):
-    [UOp.const(100+i) for i in range(self.N)]
+    [UOp.const(dtypes.int, 100+i) for i in range(self.N)]

  def test_uop_add_2n(self):
-    a = UOp.const(2)
+    a = UOp.const(dtypes.int, 2)
    for _ in range(self.N): a = a + a

  def test_uop_toposort(self):
-    a = UOp.const(0)
-    for i in range(self.N): a = a + UOp.const(100+i)
+    a = UOp.const(dtypes.int, 0)
+    for i in range(self.N): a = a + UOp.const(dtypes.int, 100+i)
    self.start_time()
    self.assertEqual(len(a.toposort()), 2*self.N+1)

  def test_uop_toposort_2n(self):
-    a = UOp.const(0)
+    a = UOp.const(dtypes.int, 0)
    for _ in range(self.N): a = a + a
    self.start_time()
    self.assertEqual(len(a.toposort()), self.N+1)

  def test_uop_simplify(self):
-    a = UOp.const(2)
+    a = UOp.const(dtypes.int, 2)
    for _ in range(self.N): (a+a).simplify()

  def test_uop_simplify_complex(self):
@@ -68,7 +68,7 @@ class TestBench(unittest.TestCase):
    for _ in range(self.N): expr.simplify()

  def test_uop_chain_free(self):
-    a = UOp.const(2)
+    a = UOp.const(dtypes.int, 2)
    for _ in range(self.N): a = a + a
    self.start_time()
    del a
@@ -5,9 +5,9 @@ from tinygrad.uop.ops import PatternMatcher, UPat

 class TestPatternMatcher(unittest.TestCase):
  def test_simple_match(self):
-    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.weakfloat), lambda x: x.rtag())])
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(1)
+    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.float), lambda x: x.rtag())])
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.int, 1)
    self.assertEqual(matcher.rewrite(c1), c1.rtag())
    self.assertEqual(matcher.rewrite(c2), None)

@@ -61,16 +61,16 @@ class TestPatternMatcher(unittest.TestCase):

  def test_uop(self):
    matcher = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.rtag())])
-    c1 = UOp.const(1.0)
+    c1 = UOp.const(dtypes.float, 1.0)
    c2 = UOp(Ops.ADD, src=(c1, c1))
    self.assertEqual(matcher.rewrite(c1), c1.rtag())
    self.assertEqual(matcher.rewrite(c2), None)

  def test_uop_set(self):
    matcher = PatternMatcher([(UPat((Ops.CONST, Ops.CAST), name="x"), lambda x: x.rtag())])
-    c1 = UOp.const(False)
+    c1 = UOp.const(dtypes.bool, False)
    c2 = UOp(Ops.CAST, arg=dtypes.int, src=(c1,))
-    c3 = UOp.const(1.0)
+    c3 = UOp.const(dtypes.float, 1.0)
    c4 = UOp(Ops.ADD, src=(c3, c3))
    self.assertEqual(matcher.rewrite(c1), c1.rtag())
    self.assertEqual(matcher.rewrite(c2), c2.rtag())
@@ -82,11 +82,11 @@ class TestPatternMatcher(unittest.TestCase):
      (UPat(Ops.CONST, arg=False, name="x"), lambda x: x.rtag()),
      (UPat(Ops.MAX, name="x"), lambda x: x.rtag()),
    ])
-    c1 = UOp.const(0.0)
-    c2 = UOp.const(False)
+    c1 = UOp.const(dtypes.float, 0.0)
+    c2 = UOp.const(dtypes.bool, False)
    c3 = UOp(Ops.MAX, src=(c1, c1))
    c4 = UOp(Ops.MUL, src=(c1, c1))
-    c5 = UOp.const(-1)
+    c5 = UOp.const(dtypes.int, -1)
    self.assertEqual(matcher.rewrite(c1), c1.rtag())
    self.assertEqual(matcher.rewrite(c2), c2.rtag())
    self.assertEqual(matcher.rewrite(c3), c3.rtag())
@@ -96,11 +96,11 @@ class TestPatternMatcher(unittest.TestCase):
  def test_filter_arg(self):
    matcher = PatternMatcher([
      (UPat(Ops.MUL, src=[UPat(Ops.CONST, name="c"), UPat(Ops.CONST, arg=2)], name="x"),
-       lambda x,c: x.rtag() if c.val in {1, -1} else None)
+       lambda x,c: x.rtag() if c.arg in {1, -1} else None)
    ])
-    y1 = UOp.const(1)
-    y2 = UOp.const(2)
-    y3 = UOp.const(-1)
+    y1 = UOp.const(dtypes.int, 1)
+    y2 = UOp.const(dtypes.int, 2)
+    y3 = UOp.const(dtypes.int, -1)
    c1 = UOp(Ops.MUL, src=(y1, y2))
    c2 = UOp(Ops.MUL, src=(y2, y2))
    c3 = UOp(Ops.MUL, src=(y3, y2))
@@ -114,27 +114,26 @@ class TestPatternMatcher(unittest.TestCase):

  def test_dup_name(self):
    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST, name="y"), UPat(Ops.CONST, name="y"))), lambda x, y: x.rtag())])
-    y1 = UOp.const(1.0)
-    y2 = UOp.const(1.0)
+    y1 = UOp.const(dtypes.float, 1.0)
+    y2 = UOp.const(dtypes.float, 1.0)
    c1 = UOp(Ops.ADD, src=(y1, y1))
    c2 = UOp(Ops.ADD, src=(y1, y2))
    self.assertEqual(matcher.rewrite(c1), c1.rtag())
    self.assertEqual(matcher.rewrite(c2), c1.rtag())

  def test_dtype(self):
-    # a concrete const dtype lives on the pair's CAST
-    matcher = PatternMatcher([(UPat(Ops.CAST, name="x", dtype=dtypes.float32), lambda x: x.rtag())])
-    c1 = UOp.const(1.0).cast(dtypes.float32)
-    c2 = UOp.const(1.0).cast(dtypes.float64)
+    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.float32), lambda x: x.rtag())])
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float64, 1.0)
    self.assertEqual(matcher.rewrite(c1), c1.rtag())
    self.assertEqual(matcher.rewrite(c2), None)

  def test_dtype_set(self):
-    matcher = PatternMatcher([(UPat(Ops.CAST, name="x", dtype={dtypes.float32, dtypes.float64}), lambda x: x.rtag())])
-    c1 = UOp.const(1.0).cast(dtypes.float32)
-    c2 = UOp.const(1.0).cast(dtypes.float64)
-    c3 = UOp.const(1.0).cast(dtypes.float16)
-    c4 = UOp.const(1).cast(dtypes.int)
+    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype={dtypes.float32, dtypes.float64}), lambda x: x.rtag())])
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float64, 1.0)
+    c3 = UOp.const(dtypes.float16, 1.0)
+    c4 = UOp.const(dtypes.int, 1)
    self.assertEqual(matcher.rewrite(c1), c1.rtag())
    self.assertEqual(matcher.rewrite(c2), c2.rtag())
    self.assertEqual(matcher.rewrite(c3), None)
@@ -142,8 +141,8 @@ class TestPatternMatcher(unittest.TestCase):

  def test_src_one(self):
    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST), UPat(Ops.CONST))), lambda x: x.rtag())])
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
    c3 = UOp(Ops.ADD, src=(c1,c2))
    self.assertEqual(matcher.rewrite(c3), c3.rtag())
    self.assertEqual(matcher.rewrite(c2), None)
@@ -159,8 +158,8 @@ class TestPatternMatcher(unittest.TestCase):

  def test_src_permutations(self):
    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=[UPat(Ops.CONST), UPat(GroupOp.ALU)]), lambda x: x.rtag())])
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
    c3 = UOp(Ops.ADD, src=(c1,c2))
    c4 = UOp(Ops.ADD, src=(c3,c2))
    c5 = UOp(Ops.ADD, src=(c2,c3))
@@ -172,8 +171,8 @@ class TestPatternMatcher(unittest.TestCase):

  def test_src_repeat(self):
    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=UPat(Ops.CONST)), lambda x: x.rtag())])
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
    c3 = UOp(Ops.ADD, src=(c1,c2))
    c4 = UOp(Ops.ADD, src=(c2,c3))
    self.assertEqual(matcher.rewrite(c3), c3.rtag())
@@ -181,9 +180,9 @@ class TestPatternMatcher(unittest.TestCase):

  def test_allow_len(self):
    matcher = PatternMatcher([(UPat(Ops.MULACC, name="x", src=(UPat(Ops.CONST),), allow_any_len=True), lambda x: x.rtag())])
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
-    c3 = UOp.const(3.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
    c4 = UOp(Ops.EXP2, src=(c1,))
    c5 = UOp(Ops.ADD, src=(c1,c2))
    c6 = UOp(Ops.MULACC, src=(c1,c2,c3))
@@ -192,8 +191,8 @@ class TestPatternMatcher(unittest.TestCase):
    self.assertEqual(matcher.rewrite(c6), c6.rtag())

  def test_deep_src_permutations(self):
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
    u1 = (c1 + c2) + c1
    u2 = (c2 + c1) + c1
    matcher = PatternMatcher([
@@ -5,7 +5,7 @@ from tinygrad.nn.state import get_parameters
 from tinygrad.engine.jit import TinyJit
 from tinygrad import Tensor, Device, GlobalCounters, dtypes, Variable
 from tinygrad.helpers import Context
-from test.helpers import slow, jit_cache_count, KernelCountException
+from test.helpers import slow, jit_cache_count
 from extra.lr_scheduler import OneCycleLR
 from test.helpers import derandomize_model

@@ -35,9 +35,8 @@ def helper_test(nm, gen, model, max_memory_allowed, max_kernels_allowed, all_jit
    assert mem_used < max_memory_allowed, f"{nm} used more than {max_memory_allowed:.3f} GB - {mem_used:.3} GB used"
    assert (max_memory_allowed - mem_used) / max_memory_allowed < 0.2, f"{max_memory_allowed:.3f} GB is too far from {mem_used:.3} GB used"
    if kernels_used:
-      if kernels_used > max_kernels_allowed: raise KernelCountException(max_kernels_allowed, kernels_used)
-      if (max_kernels_allowed - kernels_used) / max_kernels_allowed >= 0.2:
-        raise KernelCountException(max_kernels_allowed, kernels_used)
+      assert kernels_used <= max_kernels_allowed, f"{nm} used more than {max_kernels_allowed} kernels, it used {kernels_used}"
+      assert (max_kernels_allowed - kernels_used) / max_kernels_allowed < 0.2, f"{max_kernels_allowed=} is too far from {kernels_used=} used"
    if all_jitted:
      assert kernels_used > 0 and kernels_used == GlobalCounters.kernel_count or (kernels_used <= GlobalCounters.kernel_count and getattr(Device[Device.DEFAULT], "graph", None)), f"only {kernels_used} out of {GlobalCounters.kernel_count} were jitted"  # noqa: E501

@@ -2,11 +2,32 @@
 import gc, unittest, time
 from typing import cast
 from tinygrad import nn, dtypes, Device, Tensor, getenv
-from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat, KernelInfo, AxisType
-from tinygrad.helpers import GlobalCounters, Context
-from tinygrad.engine.realize import run_linear, compile_linear
-from tinygrad.codegen import to_program, full_rewrite_to_sink
-from test.helpers import check_schedule, assert_kernel_count, KernelCountException
+from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat, KernelInfo
+from tinygrad.helpers import DEBUG, GlobalCounters, Context
+from tinygrad.engine.realize import compile_linear, run_linear
+from tinygrad.codegen import to_program
+
+class KernelCountException(Exception): pass
+def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True):
+  if to_prerealize:
+    with Context(DEBUG=0, TRACK_MATCH_STATS=0): Tensor.realize(*to_prerealize)
+  if isinstance(t, Tensor): linear, var_vals = t.linear_with_vars()
+  elif isinstance(t, list) and isinstance(t[0], Tensor): linear, var_vals = Tensor.linear_with_vars(*t)
+  else:
+    assert isinstance(t, UOp), f"can't schedule {t}"
+    linear, var_vals = Tensor(t).linear_with_vars()
+  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:
+    print(f"SCHEDULE ISSUE, expecting {allowed} got {kernel_cnt}")
+    if DEBUG >= 3:
+      for i,call in enumerate(linear.src):
+        print("kernel", i+1)
+        print(call.src[0])
+    raise KernelCountException(f"{kernel_cnt} != {allowed}")
+  # test compiling the linear
+  compile_linear(linear)
+  return linear, var_vals

 def _realize_weights(m):
  for p in nn.state.get_parameters(m): p.realize()
@@ -122,7 +143,7 @@ class TestSimpleSchedule(unittest.TestCase):
    a = Tensor.empty(16,16).sum(axis=1)
    a1 = a.reshape(4,4)
    a2 = a.reshape(16,1,1)
-    check_schedule([a1, a2], 1)
+    self.assertEqual(len(Tensor.schedule_linear(a1, a2).src), 1)

 class TestSchedule(unittest.TestCase):
  def setUp(self):
@@ -134,7 +155,8 @@ class TestSchedule(unittest.TestCase):
  def test_arange_avgpool2d(self, kcount=1):
    x = Tensor.arange(25).reshape(1,1,5,5).cast(dtypes.float32)
    t = x.avg_pool2d(padding=1).clone()
-    check_schedule(t, kcount)
+    linear, var_vals = t.linear_with_vars()
+    self.assertEqual(len(linear.src), kcount)

  def test_arange_avgpool2d_fused_noopt(self):
    with Context(NOOPT=1): self.test_arange_avgpool2d(kcount=1)
@@ -202,7 +224,7 @@ class TestSchedule(unittest.TestCase):
    GlobalCounters.reset()
    expr = (a/b)/c
    expr.realize()
-    assert_kernel_count(1)
+    self.assertEqual(GlobalCounters.kernel_count, 1)
    self.assertLessEqual(GlobalCounters.global_ops, 4*3)

  # NOTE: this is causing "LAZYCACHE=1 incorrectly reuses contiguous const" #4562
@@ -335,11 +357,6 @@ class TestSchedule(unittest.TestCase):
    out1 = a.sum() + b
    check_schedule([out0, out1], 2)

-  def test_reduce_broadcast_not_recomputed(self):
-    a = Tensor.empty(32, 16).realize()
-    out = a-a.mean(axis=0, keepdim=True)
-    check_schedule(out, 2)
-
  def test_scaled_dot_product_attention_multireduce_fusion(self):
    q = Tensor.empty(32,8,16,8).realize()
    k = Tensor.empty(32,8,16,8).realize()
@@ -592,7 +609,9 @@ class TestSchedule(unittest.TestCase):
    img = Tensor.randn(BS, CIN, 64, 64).realize()
    w = Tensor.uniform(16, CIN, 3, 3).realize()
    ret = Tensor.conv2d(img, w).relu().mean().backward()
-    check_schedule([ret, img.grad, w.grad], allowed)
+    linear, var_vals = Tensor.linear_with_vars(ret, img.grad, w.grad)
+    cnt = len([call for call in linear.src if call.src[0].op is Ops.SINK])
+    assert cnt == allowed, f"expected {allowed} kernels, got {cnt}"

  def test_conv2d_half(self): self.test_conv2d(4, dtype=dtypes.half)

@@ -613,8 +632,7 @@ class TestSchedule(unittest.TestCase):
        return len([call for call in linear.src if call.src[0].op is Ops.PROGRAM])

      with Context(IMAGE=1):
-        got = cnt()
-        if got != 5: raise KernelCountException(5, got)
+        self.assertEqual(cnt(), 5)

  def test_image_f16_residual_fusion(self):
    with Context(FLOAT16=1, OPENPILOT_HACKS=1):
@@ -629,8 +647,7 @@ class TestSchedule(unittest.TestCase):
        return len([call for call in linear.src if call.src[0].op is Ops.PROGRAM])

      with Context(IMAGE=1):
-        got = cnt()
-        if got != 9: raise KernelCountException(9, got)
+        self.assertEqual(cnt(), 9)

  def _test_fusion(self, shapes, f, cnt):
    with Context(DEBUG=0, TRACK_MATCH_STATS=0):
@@ -697,19 +714,6 @@ class TestSchedule(unittest.TestCase):
    xt = X[[Tensor([2]), Tensor([1])]]
    check_schedule(xt, 1)

-  def test_split_advanced_indexing_not_recomputed(self):
-    with Context(SPLIT_REDUCEOP=1):
-      X = Tensor.empty(32768, 4).realize()
-      idx = Tensor.randint(4, high=X.shape[0])
-      linear, _ = check_schedule(X[idx], 3, [Tensor._device_rng_counters[idx.device]])
-      # The split's final reduction remains, but the one-hot gather should collapse into a direct indexed load.
-      reduce_kernels = 0
-      for call in linear.src:
-        if call.src[0].op is not Ops.SINK: continue
-        sink = full_rewrite_to_sink(call.src[0], Device[call.device].renderer)
-        reduce_kernels += any(u.op is Ops.RANGE and u.arg[-1] is AxisType.REDUCE for u in sink.toposort())
-      self.assertEqual(reduce_kernels, 1)
-
  def test_push_through_reshape(self):
    x = Tensor.empty(10, 20).realize()
    out = x.argmax(1)
@@ -858,65 +862,6 @@ 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
@@ -929,7 +874,8 @@ class TestSchedule(unittest.TestCase):
    t = Tensor.zeros((3, 3)).contiguous().realize()
    v = t[1]  # view - is_realized but not has_buffer_identity
    assert v.uop.is_realized
-    check_schedule(v, 0)
+    linear, _ = Tensor.linear_with_vars(v)
+    self.assertEqual(len(linear.src), 0)

  # NOTE: because empty does not have a lowered kernel if realize is called on a childless empty, it never gets allocated.
  def test_childless_empty_never_allocates(self):
@@ -1511,7 +1457,8 @@ class TestSchedule(unittest.TestCase):
    Tensor.manual_seed(0)
    x = Tensor.randn(4, 12, 64, 64, dtype=dtypes.half).realize()
    out = x.softmax(dtype=dtypes.float)
-    linear, _ = check_schedule(out, 3)
+    linear = out.schedule_linear()
+    self.assertEqual(len(linear.src), 3)
    # max reduction stays in input dtype (no numerical loss), upcast happens after subtracting max
    self.assertEqual(linear.src[0].src[1].dtype, dtypes.half)
    self.assertEqual(linear.src[1].src[1].dtype, dtypes.float)
@@ -1693,11 +1640,11 @@ class TestSchedule(unittest.TestCase):
    self.assertEqual(GlobalCounters.mem_used-base, 0)

  def test_const_schedule(self):
-    constv = Tensor.empty(2, 2).const_like(10).uop
+    constv = Tensor.empty(2, 2).uop.const_like(10)
    check_schedule(constv, 0)

  def test_const_schedule_contig(self):
-    constv = Tensor.empty(2, 2).const_like(10).uop.contiguous()
+    constv = Tensor.empty(2, 2).uop.const_like(10).contiguous()
    check_schedule(constv, 0)

  def test_advanced_simple_indexing_combined(self):
@@ -1926,7 +1873,8 @@ class TestFusionOp(unittest.TestCase):
    val = 1.0
    a = Tensor(val)
    for _ in range(24): a = Tensor.stack(a, a)[0]
-    check_schedule(a, 0)
+    linear = a.schedule_linear()
+    self.assertLessEqual(len(linear.src), 1)
    self.assertLess(time.perf_counter()-st, 2.0)

  def test_recursive_reshape(self):
@@ -1935,7 +1883,8 @@ class TestFusionOp(unittest.TestCase):
    b = Tensor.empty(16, 2).realize()
    r = a.sum(1)
    for _ in range(24): r = r.reshape(16, 2) + b
-    check_schedule(r, 1)
+    linear = r.schedule_linear()
+    self.assertEqual(len(linear.src), 1)
    self.assertLess(time.perf_counter()-st, 2.0)

 # NOTE: the NULL backend supports SLICE
@@ -3,7 +3,6 @@ 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
-from tinygrad.uop.weak import pm_lower_index_dtype
 from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load
 from tinygrad.helpers import Context
 from test.helpers import full_rewrite
@@ -24,8 +23,8 @@ def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UO
    UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)),
  ))

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

 class TestValidIdxSimplification(unittest.TestCase):
@@ -208,7 +207,7 @@ class TestImageSimplification(unittest.TestCase):
    if svalid is not None:
      check_uop_against_string(self, off.src[1].get_valid(), svalid)
    else:
-      self.assertEqual(off.src[1].get_valid(), UOp.const(True), "svalid is None but valid is not True")
+      self.assertEqual(off.src[1].get_valid(), UOp.const(dtypes.bool, True), "svalid is None but valid is not True")

  def test_idx_gt_c(self):
    # (idx1 < c+1).ne(True) ? (..., idx1-1+c) : 0 can drop the valid
@@ -333,7 +332,7 @@ class TestImageSimplification(unittest.TestCase):
    load = get_load_image_uop(shape, valid, idx)

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

@@ -456,12 +455,12 @@ class TestImageSimplification(unittest.TestCase):
    A1 = lidx0*32 + r0*32 + lidx1*4 - 99
    valid = ((lidx1 < 1).ne(True)) & ((lidx0 + r0) < 3).ne(True) & ((lidx0 + r0) < 19)
    alu0 = gidx0 + (A1 % 32)*32 + (A1 // 32 % 16)*1024
-    load = get_load_image_uop((1, 16384, 4), valid, (alu0, UOp.const(0)))
+    load = get_load_image_uop((1, 16384, 4), valid, (alu0, UOp.const(dtypes.weakint, 0)))
    try:
      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)<19))",
+      self.check(load, "(((lidx1<1)!=True)&(((lidx0+r0)<3)!=True)&((lidx0+r0)<19))",
                       "(gidx0+lidx1*128+(lidx0*1024+r0*1024)+-3168)", "0")

  def test_simplify10(self):
@@ -475,12 +474,12 @@ class TestImageSimplification(unittest.TestCase):
    A1 = lidx0*16 + r0*16 + lidx1*4 - 51
    valid = ((lidx1 < 1).ne(True)) & ((lidx0 + r0) < 3).ne(True) & ((lidx0 + r0) < 11)
    alu0 = lidx2 + gidx0*4 + (A1 % 16)*64 + (A1 // 16 % 8)*1024
-    load = get_load_image_uop((1, 8192, 4), valid, (alu0, UOp.const(0)))
+    load = get_load_image_uop((1, 8192, 4), valid, (alu0, UOp.const(dtypes.weakint, 0)))
    try:
      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)<11))",
+      self.check(load, "(((lidx1<1)!=True)&(((lidx0+r0)<3)!=True)&((lidx0+r0)<11))",
                       "(lidx2+gidx0*4+lidx1*256+(lidx0*1024+r0*1024)+-3264)", "0")

  def test_drop_non_monotonic_window(self):
@@ -489,20 +488,10 @@ class TestImageSimplification(unittest.TestCase):
    gidx0 = Special("gidx0", 1064)
    r12 = Range(12, 3)
    valid = ((gidx0 < 645).ne(True)) & (gidx0 < 653)
-    idx = (r12*4 + (gidx0+3)%4 + (gidx0+3)//4*24 - 3888, UOp.const(0))
+    idx = (r12*4 + (gidx0+3)%4 + (gidx0+3)//4*24 - 3888, UOp.const(dtypes.weakint, 0))
    load = get_load_image_uop((1, 48, 4), valid, idx)
    self.check(load, None, "(r12*4+(gidx0+3)%4+(gidx0+3)//4*24+-3888)", "0")

-  def test_drop_gate_committed_in_the_index_pass(self):
-    # the fused index pass runs without symbolic, so committing a weak src must not leave a CAST that
-    # symbolic later folds inside the index only: the gate's copy of the expression has to stay the same node
-    f = UOp.variable("f", 0.0, 9.0, dtypes.float)
-    idx_y = (f + UOp.const(1.0)).cast(dtypes.int)
-    load = get_load_image_uop((10, 10, 4), (UOp.const(-1) < idx_y) & (idx_y < UOp.const(10)),
-                              (Special("gidx0", 10), idx_y))
-    off = graph_rewrite(load.sink(), pm_lower_index_dtype+indexing_simplify, ctx={}).src[0].src[0]
-    self.assertEqual(off.src[1].get_valid(), UOp.const(True))
-
 class TestDropTrueGate(unittest.TestCase):
  def test_drop_true_gate_on_index(self):
    # test that INDEX with a constant True valid gets simplified to drop the valid
@@ -510,8 +499,8 @@ class TestDropTrueGate(unittest.TestCase):
    from tinygrad.uop.ops import graph_rewrite
    from tinygrad.uop.symbolic import sym
    buf = UOp.param(0, dtypes.int, (1,))
-    idx = UOp.const(0)
-    true_gate = UOp.const(True)
+    idx = UOp.const(dtypes.weakint, 0)
+    true_gate = UOp.const(dtypes.bool, True)
    index_with_gate = UOp(Ops.INDEX, src=(buf, idx.valid(true_gate)))
    # apply the optimization
    result = graph_rewrite(index_with_gate, sym+indexing_simplify)
@@ -527,50 +516,50 @@ class TestRangeShrink(unittest.TestCase):
  def test_range_shrink_single_guard(self):
    # range 0..203 guarded by r < 4 everywhere -> shrink to 0..3
    r = Range(0, 204)
-    load = get_gated_load_uop(r < UOp.const(4), r)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
    ranges = self.get_ranges(load.sink())
    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 4)
+    self.assertEqual(ranges[0].src[0].arg, 4)

  def test_range_shrink_picks_max_guard(self):
    # two loads guard the same range with r < 4 and r < 8 -> shrink to max(4, 8) = 8
    r = Range(0, 204)
-    load1 = get_gated_load_uop(r < UOp.const(4), r)
-    load2 = get_gated_load_uop(r < UOp.const(8), r)
+    load1 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
+    load2 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 8), r)
    ranges = self.get_ranges(UOp.sink(load1, load2))
    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 8)
+    self.assertEqual(ranges[0].src[0].arg, 8)

  def test_range_no_shrink_guard_ge_max(self):
    # guard r < 300 with range max 204 -> no shrink (guard doesn't constrain)
    r = Range(0, 204)
-    load = get_gated_load_uop(r < UOp.const(300), r)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 300), r)
    ranges = self.get_ranges(load.sink())
    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 204)
+    self.assertEqual(ranges[0].src[0].arg, 204)

  def test_range_no_shrink_when_unguarded_elsewhere(self):
    # one load guards r < 4, but another load uses r without a gate -> no shrink
    r = Range(0, 204)
-    load1 = get_gated_load_uop(r < UOp.const(4), r)
+    load1 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
    load2 = UOp(Ops.LOAD, src=(UOp.param(1, dtypes.float, (204,)).index(r),))
    ranges = self.get_ranges(UOp.sink(load1, load2))
    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 204)
+    self.assertEqual(ranges[0].src[0].arg, 204)

  def test_range_no_shrink_when_used_in_reduce(self):
    # range used in both a gated load AND directly in the reduce expression -> no shrink
    r = Range(0, 204)
-    gated_load = get_gated_load_uop(r < UOp.const(4), r)
+    gated_load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
    red = (r.cast(dtypes.float) + gated_load).reduce(r, arg=Ops.ADD)
    ranges = self.get_ranges(red.sink())
    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 204)
+    self.assertEqual(ranges[0].src[0].arg, 204)

  def test_range_shrink_to_single_iteration(self):
    # guard r < 1 shrinks range to 1 -> single iteration, range eliminated entirely
    r = Range(0, 204)
-    load = get_gated_load_uop(r < UOp.const(1), r)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 1), r)
    ranges = self.get_ranges(load.sink())
    self.assertEqual(len(ranges), 0)

@@ -578,19 +567,19 @@ class TestRangeShrink(unittest.TestCase):
    # emulates mask.where(x.pad_to(mask.shape), Invalid): range should shrink accordingly
    from tinygrad.dtype import Invalid
    r = Range(0, 204)
-    x = (r < 4).where(UOp.const(1.0), Invalid)
+    x = (r < 4).where(UOp.const(dtypes.float, 1), Invalid)
    ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, Invalid)).sink())
    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 4)
+    self.assertEqual(ranges[0].src[0].arg, 4)

  def test_range_shrink_store_where_invalid_flipped(self):
    # above, but flipped
    from tinygrad.dtype import Invalid
    r = Range(0, 204)
-    x = (r < 4).where(UOp.const(1.0), Invalid)
+    x = (r < 4).where(UOp.const(dtypes.float, 1), Invalid)
    ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r >= 4).where(Invalid, x)).sink())
    self.assertEqual(len(ranges), 1)
-    self.assertEqual(ranges[0].src[0].val, 4)
+    self.assertEqual(ranges[0].src[0].arg, 4)

 if __name__ == '__main__':
  unittest.main()
@@ -1,6 +1,5 @@
 import unittest
 from tinygrad import Variable
-from tinygrad.uop.ops import UOp


 class TestFuzzFailure(unittest.TestCase):
@@ -9,7 +8,7 @@ class TestFuzzFailure(unittest.TestCase):
    v2=Variable('v2', 0, 2)
    v3=Variable('v3', 0, 1)
    expr = (((((((((((((((((((((((0//4)%2)//8)+-2)+-4)+-3)+v1)+-4)+v2)+-2)+v3)+v2)//3)%7)*1)//2)+v2)*-1)+2)+1)+0)+-3)+v3)
-    v1_val, v2_val, v3_val = UOp.const(8), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(8), v2.const_like(0), v3.const_like(0)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -20,7 +19,7 @@ class TestFuzzFailure(unittest.TestCase):
    v3=Variable('v3', 0, 3)
    expr = (((((((((((((((((((((((((0*4)//5)*2)*-1)*-2)+-4)*4)*2)*3)*4)+-4)*4)+v2)+v2)+v3)//3)+v2)+v1)//9)+3)+1)//1)+-4)//4)*2)
    expr = (((((v1+(v2+(((v3+(v2*2))+1)//3)))+4)//9)+-57)//(9*4))
-    v1_val, v2_val, v3_val = UOp.const(6), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(6), v2.const_like(0), v3.const_like(0)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -30,7 +29,7 @@ class TestFuzzFailure(unittest.TestCase):
    v2=Variable('v2', 0, 1)
    v3=Variable('v3', 0, 2)
    expr = (((((((((((((((((((0//2)//3)+v3)+0)+-4)*-2)*-2)+-1)+2)+3)+v3)+0)//8)*-3)+0)*-2)*-4)*-2)//5)
-    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(0), v3.const_like(0)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -40,7 +39,7 @@ class TestFuzzFailure(unittest.TestCase):
    v2=Variable('v2', 0, 3)
    v3=Variable('v3', 0, 4)
    expr = (((((((((((((((((((((((((((((0*-2)+0)*-1)//9)//6)//8)+v1)*-4)+v2)//4)//8)+4)*3)+v1)+v3)//8)//7)+4)+v3)*-4)+1)+v1)*3)+4)*2)//5)//2)//3)*-4)
-    v1_val, v2_val, v3_val = UOp.const(2), UOp.const(0), UOp.const(2)
+    v1_val, v2_val, v3_val = v1.const_like(2), v2.const_like(0), v3.const_like(2)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -50,7 +49,7 @@ class TestFuzzFailure(unittest.TestCase):
    v2=Variable('v2', 0, 1)
    v3=Variable('v3', 0, 3)
    expr = ((((((((((((((0+v2)+v1)*0)+v2)//1)//7)+-2)+v2)+v1)*4)+-3)//5)+v2)+1)
-    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(0), v3.const_like(0)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -61,7 +60,7 @@ class TestFuzzFailure(unittest.TestCase):
    v3=Variable('v3', 0, 128)
    expr = (((((((((((((((((((((((((((((0//3)+4)+v1)//2)+-1)//1)*1)*-1)*4)//5)+v1)//6)+v1)*-1)+-4)+v2)+-2)*-3)+v3)+-4)+-2)*-1)//8)//4)*-4)+3)+v3)*
            -2)+v2)
-    v1_val, v2_val, v3_val = UOp.const(8), UOp.const(3), UOp.const(2)
+    v1_val, v2_val, v3_val = v1.const_like(8), v2.const_like(3), v3.const_like(2)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -71,7 +70,7 @@ class TestFuzzFailure(unittest.TestCase):
    v2=Variable('v2', 0, 5)
    v3=Variable('v3', 0, 128)
    expr = (((((((((((((((((((((((((((((0+v2)*-4)+0)//9)+-4)*-2)*3)*4)//9)+v3)+v1)//4)+v1)+v3)+-1)*4)//4)+v2)//7)//3)+v1)+v2)+v3)+1)*2)//4)*3)+-1)*1)
-    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(2), UOp.const(65)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(2), v3.const_like(65)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -81,7 +80,7 @@ class TestFuzzFailure(unittest.TestCase):
    v2=Variable('v2', 0, 8)
    v3=Variable('v3', 0, 9)
    expr = (((((((0+-1)+2)+v1)*-2)//3)+v1)*-4)
-    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(0), v3.const_like(0)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -91,7 +90,7 @@ class TestFuzzFailure(unittest.TestCase):
    v2=Variable('v2', 0, 1)
    v3=Variable('v3', 0, 8)
    expr = (((((((((((((((((((((((((((((0*-2)//1)+3)*-2)+-3)*-4)*1)+v1)+0)%2)%8)%9)+v2)%9)+-4)//4)+-1)*-2)+0)+v1)+v1)+3)+v1)+4)+-4)+0)*2)+-3)%6)
-    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(1), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(1), v3.const_like(0)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -102,7 +101,7 @@ class TestFuzzFailure(unittest.TestCase):
    v3=Variable("v3", 0, 32)
    x5 = (v1 <= 9).where(v1 * -4 - 4, v1 // 9) // 9
    expr = ((x5 >= -4).where(x5, (v2 % 3 + v2) // 5) * -1).maximum(((v1 * -2) % 6 + v3 % 1) * -1) * -1
-    v1_val, v2_val, v3_val = UOp.const(9), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(9), v2.const_like(0), v3.const_like(0)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -112,7 +111,7 @@ class TestFuzzFailure(unittest.TestCase):
    v2=Variable("v2", 0, 128)
    v3=Variable("v3", 0, 5)
    expr = (((v2 * 0).maximum(8) - v2 * 2) % 5 + v1 // 6 + v1 + 5) % 5
-    v1_val, v2_val, v3_val = UOp.const(0), UOp.const(7), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(0), v2.const_like(7), v3.const_like(0)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    self.assertEqual(num, rn)
@@ -60,7 +60,7 @@ class TestTensorUOpClone(unittest.TestCase):
    t = _t(3, 4).float()
    self.assertIs(_strip_unique(t.clone().uop), _strip_unique(t.uop.clone()))
  def test_clone_deviceless_const(self):
-    u = UOp.const(2.0)
+    u = UOp.const(dtypes.float, 2.0)
    self.assertIs(_strip_unique(Tensor(u).clone().uop), _strip_unique(u.clone()))

 class TestTensorUOpGradient(unittest.TestCase):
@@ -382,7 +382,7 @@ class TestTensorUOpStack(unittest.TestCase):
    self.assertIs(_t(2, 3).uop.stack(w.uop).dtype, dtypes.float32)
  def test_stack_index_dtype(self):
    # index is outside the promotion lattice, equal dtypes bypass promotion
-    self.assertEqual(UOp.const(1).stack(UOp.const(2)).shape, (2,))
+    self.assertEqual(UOp.const(dtypes.weakint, 1).stack(UOp.const(dtypes.weakint, 2)).shape, (2,))

 class TestTensorUOpConv2d(unittest.TestCase):
  def test_conv2d_basic(self):
@@ -10,7 +10,7 @@ class TestTranscendentalFunctions(unittest.TestCase):
    # TODO: Test constant input when constant folding is fixed (or maybe test both variants)
    # Load input value from a buffer to prevent constant folding
    input_buf = UOp.param(1, dtypes.double, (1,))
-    loaded_value = input_buf.index(UOp.const(0)).load()
+    loaded_value = input_buf.index(UOp.const(dtypes.int, 0)).load()
    def eval_payne_hanek_reduction(v:float) -> tuple[float, int]:
      return tuple(eval_uop(u, [(dtypes.float64, [v])]) for u in payne_hanek_reduction(loaded_value))

@@ -27,48 +27,48 @@ class TestTranscendentalFunctions(unittest.TestCase):
    np.testing.assert_equal(q, 4)

  def test_cody_waite_reduction(self):
-    r, q = (eval_uop(u) for u in cody_waite_reduction(UOp.const(12 * math.pi + 0.1).cast(dtypes.float64)))
+    r, q = (eval_uop(u) for u in cody_waite_reduction(UOp.const(dtypes.float64, 12 * math.pi + 0.1)))
    np.testing.assert_allclose(r, 0.1)
    np.testing.assert_equal(q, 12)

  def test_frexp(self):
    for x in (1, -1):
-      mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(float(x)).cast(dtypes.float64)))
+      mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, x)))
      np.testing.assert_equal(mantissa, 0.5)
      np.testing.assert_equal(exponent, 1)

    for x in (2, -2):
-      mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(2.0).cast(dtypes.float64)))
+      mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, 2.0)))
      np.testing.assert_equal(mantissa, 0.5)
      np.testing.assert_equal(exponent, 2)

-    mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(5.0).cast(dtypes.float64)))
+    mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, 5.0)))
    np.testing.assert_equal(mantissa, 0.625)
    np.testing.assert_equal(exponent, 3)

-    mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(1000.0).cast(dtypes.float64)))
+    mantissa, exponent = (eval_uop(u) for u in frexp(UOp.const(dtypes.float64, 1000.0)))
    np.testing.assert_allclose(mantissa, 0.9765625)
    np.testing.assert_equal(exponent, 10)

  def test_rintk(self):
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(0.0).cast(dtypes.float))), 0)
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.0).cast(dtypes.float))), 5)
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.5).cast(dtypes.float))), 6)
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(5.999).cast(dtypes.float))), 6)
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.0).cast(dtypes.float))), -5)
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.5).cast(dtypes.float))), -6)
-    np.testing.assert_allclose(eval_uop(rintk(UOp.const(-5.999).cast(dtypes.float))), -6)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, 0.0))), 0)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, 5.0))), 5)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, 5.5))), 6)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, 5.999))), 6)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, -5.0))), -5)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, -5.5))), -6)
+    np.testing.assert_allclose(eval_uop(rintk(UOp.const(dtypes.float, -5.999))), -6)

  def test_pow2if(self):
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(0).cast(dtypes.int), dtypes.float)), 1.0)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(1).cast(dtypes.int), dtypes.float)), 2.0)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(2).cast(dtypes.int), dtypes.float)), 4.0)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(10).cast(dtypes.int), dtypes.float)), 1024.0)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(63).cast(dtypes.int), dtypes.float)), 2**63)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-1).cast(dtypes.int), dtypes.float)), 0.5)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-2).cast(dtypes.int), dtypes.float)), 0.25)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-10).cast(dtypes.int), dtypes.float)), 2**-10)
-    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(-63).cast(dtypes.int), dtypes.float)), 2**-63)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 0), dtypes.float)), 1.0)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 1), dtypes.float)), 2.0)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 2), dtypes.float)), 4.0)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 10), dtypes.float)), 1024.0)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, 63), dtypes.float)), 2**63)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -1), dtypes.float)), 0.5)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -2), dtypes.float)), 0.25)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -10), dtypes.float)), 2**-10)
+    np.testing.assert_allclose(eval_uop(pow2if(UOp.const(dtypes.int, -63), dtypes.float)), 2**-63)

 if __name__ == '__main__':
  unittest.main()
@@ -1,42 +1,41 @@
 import unittest, pytest
-from tinygrad import dtypes, Variable, Device
+from tinygrad import dtypes, Variable
 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, KernelInfo
+from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, GroupOp, AxisType, broadcast_axes
 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)),
-  (UPat.cvar('x') + UPat.cvar('y'), lambda x,y: UOp.const(x.val+y.val)),
-  (UPat.cvar('x') * UPat.cvar('y') * UPat.cvar('z'), lambda x,y,z: UOp.const(x.val*y.val*z.val)),
-  ((UPat.var('x') + UPat.cvar('c1')) + UPat.cvar('c2'), lambda x,c1,c2: x + (c1.val+c2.val)),
+  (UPat.cvar('x', dtypes.int), lambda x: UOp.const(dtypes.float, 1.0) + UOp.const(dtypes.float, 2.0)),
+  (UPat.cvar('x') + UPat.cvar('y'), lambda x,y: UOp.const(dtypes.float, x.arg+y.arg)),
+  (UPat.cvar('x') * UPat.cvar('y') * UPat.cvar('z'), lambda x,y,z: UOp.const(dtypes.float, x.arg*y.arg*z.arg)),
+  ((UPat.var('x') + UPat.cvar('c1')) + UPat.cvar('c2'), lambda x,c1,c2: x + (c1.arg+c2.arg)),
 ])

 def const_values(u:UOp):
-  if u.op is Ops.CONST: return (u.val,)
-  if u.op is Ops.STACK: return tuple(x.val for x in u.src)
+  if u.op is Ops.CONST: return (u.arg,) if not isinstance(u.arg, tuple) else u.arg
+  if u.op is Ops.STACK: return tuple(x.arg for x in u.src)
  raise AssertionError(f"expected const-like UOp, got {u.op}")

 class TestGraphRewriteConst(unittest.TestCase):
  def test_gep_const(self):
-    v1 = UOp.const((0,1,2), dtypes.int)
+    v1 = UOp.const(dtypes.int, (0,1,2))
    v2 = v1.index(1)
    ret = graph_rewrite(v2, sym)
    self.assertEqual(ret.dtype, dtypes.int)
-    self.assertEqual(ret.val, 1)
+    self.assertEqual(ret.arg, 1)

  def test_add_const(self):
-    v1 = UOp.const((0,1,2))
-    v2 = UOp.const((5,6,7))
+    v1 = UOp.const(dtypes.int, (0,1,2))
+    v2 = UOp.const(dtypes.int, (5,6,7))
    ret = graph_rewrite(v1+v2, sym)
    self.assertEqual(ret.op, Ops.STACK)
    self.assertEqual(const_values(ret), (5,7,9))

  def test_add_const_lose_v(self):
-    v1 = UOp.const((0,1,2))
-    v2 = UOp.const((2,1,0))
+    v1 = UOp.const(dtypes.int, (0,1,2))
+    v2 = UOp.const(dtypes.int, (2,1,0))
    ret = graph_rewrite(v1+v2, sym)
    self.assertEqual(ret.op, Ops.STACK)
    self.assertEqual(const_values(ret), (2,2,2))
@@ -50,47 +49,47 @@ class TestModularWraparound(unittest.TestCase):
    self.assertEqual(len(results), 2)  # +1 for SINK
    self.assertEqual(results[0].op, Ops.CONST)
    self.assertEqual(results[0].dtype, uop.dtype)
-    self.assertEqual(results[0].val, expected)
+    self.assertEqual(results[0].arg, expected)

  @xfail_broken_const_wraparound
  def test_cast(self):
    t = self._test
-    t(UOp.const(0xABCD17D6, dtypes.uint).cast(dtypes.uint8), 0xD6)
-    t(UOp.const(0xABCD17D6, dtypes.uint).cast(dtypes.uint8).cast(dtypes.uint), 0xD6)
+    t(UOp.const(dtypes.uint, 0xABCD17D6).cast(dtypes.uint8), 0xD6)
+    t(UOp.const(dtypes.uint, 0xABCD17D6).cast(dtypes.uint8).cast(dtypes.uint), 0xD6)

  @xfail_broken_const_wraparound
  def test_mul(self):
    t = self._test
-    t(UOp.const(0xABCD17D6, dtypes.uint) * 0xAABBCCDD, 1147018174)
-    t(UOp.const(0xABCD17D6, dtypes.int) * 10, -1241321892)
+    t(UOp.const(dtypes.uint, 0xABCD17D6) * 0xAABBCCDD, 1147018174)
+    t(UOp.const(dtypes.int, 0xABCD17D6) * 10, -1241321892)

  @xfail_broken_const_wraparound
  def test_div(self):
    t = self._test
-    t(UOp.const(0xABCD17D6, dtypes.uint) * 0xAABBCCDD // 11, 104274379)
-    t(UOp.const(0xABCD17D6, dtypes.int) * 10 // 11, -112847444)
+    t(UOp.const(dtypes.uint, 0xABCD17D6) * 0xAABBCCDD // 11, 104274379)
+    t(UOp.const(dtypes.int, 0xABCD17D6) * 10 // 11, -112847444)

  @xfail_broken_const_wraparound
  def test_neg(self):
    t = self._test
-    t(-UOp.const(1, dtypes.uint8), 0xFF)
-    t(-UOp.const(1, dtypes.uint16), 0xFFFF)
-    t(-UOp.const(1, dtypes.uint32), 0xFFFFFFFF)
-    t(-UOp.const(1, dtypes.uint64), 0xFFFFFFFFFFFFFFFF)
+    t(-UOp.const(dtypes.uint8, 1), 0xFF)
+    t(-UOp.const(dtypes.uint16, 1), 0xFFFF)
+    t(-UOp.const(dtypes.uint32, 1), 0xFFFFFFFF)
+    t(-UOp.const(dtypes.uint64, 1), 0xFFFFFFFFFFFFFFFF)

  @xfail_broken_const_wraparound
  def test_neg_min_int(self):
    t = self._test
-    t(-UOp.const(-2**7, dtypes.int8), -2**7)
-    t(-UOp.const(-2**15, dtypes.int16), -2**15)
-    t(-UOp.const(-2**31, dtypes.int32), -2**31)
-    t(-UOp.const(-2**63, dtypes.int64), -2**63)
+    t(-UOp.const(dtypes.int8, -2**7), -2**7)
+    t(-UOp.const(dtypes.int16, -2**15), -2**15)
+    t(-UOp.const(dtypes.int32, -2**31), -2**31)
+    t(-UOp.const(dtypes.int64, -2**63), -2**63)

  @xfail_broken_const_wraparound
  def test_payne_hanek_reduction_bug(self):
    t = self._test
-    a = (UOp.const(43748177600, dtypes.uint).cast(dtypes.uint) | 36).cast(dtypes.ulong)
-    b = 2536655455 * a + 4294967296 * UOp.const(25366554550, dtypes.ulong)
+    a = (UOp.const(dtypes.uint, 43748177600).cast(dtypes.uint) | 36).cast(dtypes.ulong)
+    b = 2536655455 * a + 4294967296 * UOp.const(dtypes.ulong, 25366554550)
    c = (b + 2261737165) // 4611686018427387904
    t(c, 0)

@@ -104,67 +103,67 @@ class TestGraphRewrite(unittest.TestCase):
  # NOTE: this shows why we can't have a UOp in arg
  @unittest.expectedFailure
  def test_no_dedup_args(self):
-    a1 = UOp.variable("a1", UOp.const(0), UOp.const(11), dtypes.int)
-    a2 = UOp.variable("a2", UOp.const(0), UOp.const(11), dtypes.int)
+    a1 = UOp.variable("a1", UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 11), dtypes.int)
+    a2 = UOp.variable("a2", UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 11), dtypes.int)
    sink = a1.sink(a2)
    variables = [x for x in graph_rewrite(sink, PatternMatcher([])).toposort() if x.op is Ops.PARAM and x.addrspace is AddrSpace.ALU]
    self.assertEqual(len(variables), 1)

  def test_simple(self):
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
    nout = graph_rewrite(c1+c2, simple_pm)
    self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.val, 3.0)
+    self.assertEqual(nout.arg, 3.0)

  def test_depth_2_late(self):
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
-    c3 = UOp.const(3.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
    nout = graph_rewrite(c1*c2*(c3+c3), simple_pm)
    self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.val, 12.0)
+    self.assertEqual(nout.arg, 12.0)

  def test_double(self):
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
-    c3 = UOp.const(3.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
    nout = graph_rewrite(c1+c2+c3, simple_pm)
    self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.val, 6.0)
+    self.assertEqual(nout.arg, 6.0)

  def test_triple(self):
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
-    c3 = UOp.const(3.0)
-    c4 = UOp.const(4.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
+    c4 = UOp.const(dtypes.float, 4.0)
    nout = graph_rewrite(c1+c2+c3+c4, simple_pm)
    self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.val, 10.0)
+    self.assertEqual(nout.arg, 10.0)

  def test_diamond(self):
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
-    c3 = UOp.const(3.0)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
+    c3 = UOp.const(dtypes.float, 3.0)
    nout = graph_rewrite((c1+c2)+(c1+c3), simple_pm)
    self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.val, 7.0)
+    self.assertEqual(nout.arg, 7.0)

  def test_magic_4(self):
-    c1 = UOp.const(4)
+    c1 = UOp.const(dtypes.int, 4.0)
    nout = graph_rewrite(c1, simple_pm)
    self.assertEqual(nout.op, Ops.CONST)
-    self.assertEqual(nout.val, 3.0)
+    self.assertEqual(nout.arg, 3.0)

  def test_depth_2_fold(self):
-    v = UOp.variable("v", 0, 1, dtypes.float, param=True)
-    c1 = UOp.const(1.0)
-    c2 = UOp.const(2.0)
+    v = UOp.variable("v", 0, 1, dtypes.float)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
    nout = graph_rewrite(v+c1+c2, simple_pm)
    self.assertEqual(nout.op, Ops.ADD)
    self.assertEqual(nout.src[0].op, Ops.PARAM)
    self.assertEqual(nout.src[1].op, Ops.CONST)
-    self.assertEqual(nout.src[1].val, 3.0)
+    self.assertEqual(nout.src[1].arg, 3.0)

  def test_commutative_work(self):
    a = UOp.variable('a', 0, 1)
@@ -175,14 +174,14 @@ class TestGraphRewrite(unittest.TestCase):
    a = UOp.variable('a', 0, 1)
    tst = (2+a).simplify()
    self.assertIs(tst.src[0], a)
-    self.assertIs(tst.src[1], UOp.const(2))
+    self.assertIs(tst.src[1], a.const_like(2))

  def test_consts_go_last(self):
    a = UOp.variable('a', 0, 1)
    b = UOp.variable('b', 0, 1)
    c = UOp.variable('c', 0, 1)
    d = UOp.variable('d', 0, 1)
-    outs = [2+a, 2+a+d+3+b+c+4, UOp.const(2)+a, (4+d)+c+(2+a)+b]
+    outs = [2+a, 2+a+d+3+b+c+4, a.const_like(2)+a, (4+d)+c+(2+a)+b]
    for out in outs:
      sink = graph_rewrite(out, sym)
      print(sink.render())
@@ -192,67 +191,85 @@ class TestGraphRewrite(unittest.TestCase):

 class TestUOpGraph(unittest.TestCase):
  def test_add_constant_fold(self):
-    c1 = UOp.const(1.0, dtypes.float)
-    c2 = UOp.const(2.0, dtypes.float)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
    out = c1+c2
    uops = to_uops_list([out])
    self.assertEqual(len(uops), 2)  # +1 for SINK
    out = uops[-2]
    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 3.0)
+    self.assertEqual(out.arg, 3.0)

  def test_where_same_fold(self):
    v = UOp.variable('tmp', 0, 1)
-    c0 = UOp.const(0)
+    c0 = UOp.const(dtypes.weakint, 0)
    vc = v != c0
-    c1 = UOp.const(1.0, dtypes.float)
+    c1 = UOp.const(dtypes.float, 1.0)
    out = vc.where(c1, c1)
    uops = to_uops_list([out])
    self.assertEqual(len(uops), 2)  # +1 for SINK
    out = uops[-2]
    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 1.0)
+    self.assertEqual(out.arg, 1.0)

  def test_where_const_fold(self):
-    bf = UOp.const(False)
-    c1 = UOp.const(1.0, dtypes.float)
-    c2 = UOp.const(2.0, dtypes.float)
+    bf = UOp.const(dtypes.bool, False)
+    c1 = UOp.const(dtypes.float, 1.0)
+    c2 = UOp.const(dtypes.float, 2.0)
    out = bf.where(c1, c2)
    uops = to_uops_list([out])
    self.assertEqual(len(uops), 2)  # +1 for SINK
    out = uops[-2]
    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 2.0)
+    self.assertEqual(out.arg, 2.0)

  def test_const_cast(self):
-    bf = UOp.const(False)
+    bf = UOp.const(dtypes.bool, False)
    out = bf.cast(dtypes.int)
    uops = to_uops_list([out])
    self.assertEqual(len(uops), 2)  # +1 for SINK
    out = uops[-2]
    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 0)
+    self.assertEqual(out.arg, 0)

  def test_const_bitcast(self):
-    bf = UOp.const(1.0, dtypes.float)
+    bf = UOp.const(dtypes.float, 1.0)
    out = bf.bitcast(dtypes.uint32)
    uops = to_uops_list([out])
    self.assertEqual(len(uops), 2)  # +1 for SINK
    out = uops[-2]
    self.assertEqual(out.op, Ops.CONST)
-    self.assertEqual(out.val, 0x3F800000)
+    self.assertEqual(out.arg, 0x3F800000)

  @unittest.expectedFailure
  def test_const_shape_change_bitcast(self):
-    bf = UOp.const(0x3F).cast(dtypes.uint8)
+    bf = UOp.const(dtypes.uint8, 0x3F)
    out = bf.bitcast(dtypes.half)
    uops = to_uops_list([out])
    self.assertEqual(len(uops), 2)  # +1 for SINK

+  def test_coalesce_aliased_stores(self):
+    from tinygrad.codegen.late.coalesce import memory_coalescing
+    from tinygrad import Device
+    # two distinct stores to the same INDEX (e.g. a double-buffered LDS slot written in a
+    # prologue and a loop body) can't be merged: merging would drop one write
+    lbuf = UOp.placeholder((64,), dtypes.half, 0, AddrSpace.LOCAL)
+    r = UOp.range(4, 0, AxisType.LOOP)
+    gbuf = UOp.placeholder((64,), dtypes.half, 1, AddrSpace.GLOBAL)
+    s_a = lbuf.index(r*2).store(gbuf.index(r*2).load())
+    s_b = lbuf.index(r*2).store(gbuf.index(r*2+64).load())   # aliases s_a
+    s_c = lbuf.index(r*2+1).store(gbuf.index(r*2+32).load()) # adjacent, unique -> coalesceable with nothing
+    out = memory_coalescing(UOp.sink(s_a, s_b, s_c).end(r), Device["NULL"].renderer)
+    stores = [u for u in out.toposort() if u.op is Ops.STORE]
+    # both aliased stores survive (both datas preserved), nothing merged across them
+    self.assertEqual(len([u for u in stores if u.src[0].src[1] is not None and u.src[0].op is Ops.INDEX]), 3)
+    datas = sorted([u.src[1] for u in stores], key=str)
+    self.assertEqual(len(datas), 3)
+
  def test_devectorize_derives_lane_dtype(self):
    from tinygrad.codegen import do_devectorize
    # an Invalid lane derives bool while the value lane derives float: the lane rebuild must derive, not inherit
-    lhs = UOp.stack(UOp.invalid(), UOp.const(1.0).cast(dtypes.float))
+    lhs = UOp.stack(UOp.invalid(), UOp.const(None, 1.0).cast(dtypes.float))
    out = do_devectorize(lhs * lhs)
    invalid_lane_mul = next(u for u in out.src[0].toposort() if u.op is Ops.MUL)
    self.assertIs(invalid_lane_mul.dtype, dtypes.bool)
@@ -260,7 +277,7 @@ class TestUOpGraph(unittest.TestCase):
  @unittest.skip("this test isn't valid uops")
  def test_noop_vectorize_fold(self):
    d0 = UOp.param(0, dtypes.float, (1,))
-    idx = UOp.const(0)
+    idx = UOp.const(dtypes.int, 0)
    ld = d0.load(idx, dtype=dtypes.float)
    vec = UOp(Ops.STACK, dtypes.float, (ld,))
    x = vec.index(0)
@@ -274,7 +291,7 @@ class TestUOpGraph(unittest.TestCase):
    d0 = UOp.param(0, dtypes.float, (1,))
    d1 = UOp.param(1, dtypes.float, (1,))
    d2 = UOp.param(2, dtypes.float, (1,))
-    idx = UOp.const(0)
+    idx = UOp.const(dtypes.int, 0)
    def _test_vec(geps, count=4):
      vec = UOp(Ops.STACK, dtypes.float, geps)
      out = d0.index(idx).store(vec)
@@ -311,7 +328,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_gep_vec_const_fold(self):
    for vec_size in [2, 4, 8]:
-      consts = [UOp.const(float(i), dtypes.float) for i in range(vec_size)]
+      consts = [UOp.const(dtypes.float, float(i)) for i in range(vec_size)]
      vec = UOp(Ops.STACK, src=tuple(consts))
      with Context(SPEC=0):
        uops = to_uops_list([vec.index(i) for i in range(vec_size)])
@@ -321,7 +338,7 @@ class TestUOpGraph(unittest.TestCase):
  def test_cast_alu_fold(self):
    d0 = UOp.param(0, dtypes.bool, (1,))
    d1 = UOp.param(1, dtypes.int, (1,))
-    idx = UOp.const(0)
+    idx = UOp.const(dtypes.int, 0)
    ld = d1.index(idx)
    alu = (ld<1).cast(dtypes.bool)
    out = d0.index(idx).store(alu)
@@ -331,7 +348,7 @@ class TestUOpGraph(unittest.TestCase):
  def test_double_cast_fold(self):
    d0 = UOp.param(0, dtypes.float, (1,))
    d1 = UOp.param(1, dtypes.int, (1,))
-    idx = UOp.const(0, dtypes.int)
+    idx = UOp.const(dtypes.int, 0)
    ld = d1.index(idx)
    alu = ld.cast(dtypes.float).cast(dtypes.float)
    out = d0.index(idx).store(alu)
@@ -339,9 +356,9 @@ class TestUOpGraph(unittest.TestCase):
    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)

  def test_depth_2_const_fold(self):
-    v = UOp.variable("tmp", 0, 1, dtypes.int, param=True)
-    c2 = UOp.const(2, dtypes.int)
-    c4 = UOp.const(4, dtypes.int)
+    v = UOp.variable("tmp", 0, 1, dtypes.int)
+    c2 = UOp.const(dtypes.int, 2)
+    c4 = UOp.const(dtypes.int, 4)
    vc = v+c2
    out = vc+c4
    uops = to_uops_list([out])
@@ -349,19 +366,19 @@ class TestUOpGraph(unittest.TestCase):
    out = uops[-2]  # -2 to skip SINK
    self.assertEqual(out.op, Ops.ADD)
    self.assertEqual(out.src[1].op, Ops.CONST)
-    self.assertEqual(out.src[1].val, 6)
+    self.assertEqual(out.src[1].arg, 6)

  def test_bitcast_to_same_dtype_fold(self):
    for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
      d0 = UOp.param(0, dt, (1,))
-      v = d0.index(UOp.const(0))
+      v = d0.index(UOp.const(dtypes.int, 0))
      uops = to_uops_list([v.bitcast(dt)])
      self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST and x.dtype is dt]), 0, f"dtype = {dt}")

  def test_sub_with_cast_folds(self):
    a = Variable("a", 0, 5)
    uops = to_uops_list([a.cast(dtypes.int)+(-a).cast(dtypes.int)])
-    assert uops[0] == UOp.const(0, dtypes.int)
+    assert uops[0] == UOp.const(dtypes.int, 0)
    assert uops[-1].op == Ops.SINK

  def test_where_on_gated_load_fold(self):
@@ -373,7 +390,7 @@ class TestUOpGraph(unittest.TestCase):
    uops = to_uops_list([out.index(ridx0).store(w)])
    for u in uops:
      assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].val==5
+      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].arg==5

  def test_where_on_gated_load_folds_swapped_branches(self):
    ridx0 = UOp.range(100, 0)
@@ -383,7 +400,7 @@ class TestUOpGraph(unittest.TestCase):
    uops = to_uops_list([w])
    for u in uops:
      assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD: assert u.src[1].val==5
+      if u.op is Ops.LOAD: assert u.src[1].arg==5

  def test_where_on_gated_load_with_cast(self):
    ridx0 = UOp.range(100, 0)
@@ -395,13 +412,13 @@ class TestUOpGraph(unittest.TestCase):
    uops = to_uops_list([out.index(ridx0).store(w)])
    for u in uops:
      assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].val == 5
+      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: assert u.src[1].arg == 5

  def test_where_on_casted_gated_load_extra_cond(self):
    ridx0 = UOp.range(100, 0)
    d0 = UOp.param(0, dtypes.float, (100,))
    ld = d0.index(ridx0.valid(ridx0<50))
-    w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(0.0)).cast(dtypes.half)
+    w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(dtypes.float, 0)).cast(dtypes.half)
    out = UOp.param(1, dtypes.half, (100,))
    uops = to_uops_list([out.index(ridx0).store(w)])
    for u in uops:
@@ -411,7 +428,7 @@ class TestUOpGraph(unittest.TestCase):
    ridx0 = UOp.range(100, 0)
    d0 = UOp.param(0, dtypes.float, (100,))
    ld = d0.index(ridx0.valid(ridx0<50))
-    w = ((ridx0<50) & (ridx0>30)).where(UOp.const(0.0), ld).cast(dtypes.half)
+    w = ((ridx0<50) & (ridx0>30)).where(UOp.const(dtypes.float, 0), ld).cast(dtypes.half)
    out = UOp.param(1, dtypes.half, (100,))
    uops = to_uops_list([out.index(ridx0).store(w)])
    for u in uops:
@@ -427,22 +444,22 @@ class TestUOpGraph(unittest.TestCase):
    uops = to_uops_list([st])
    for u in uops:
      assert u.op is not Ops.WHERE
-      if u.op is Ops.STORE: assert u.src[1].val==5
+      if u.op is Ops.STORE: assert u.src[1].arg==5

  def test_load_idx_becomes_int(self):
    # mnist indexing with split reduceop
    # Make sure we are not doign math on the loaded index, which would promote it to long
    c0 = UOp.param(0, dtypes.uchar, (128000,))
-    c1 = UOp.range(UOp.const(512), 1, AxisType.WEAK)
-    c2 = UOp.range(UOp.const(250), 2, AxisType.WEAK)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.WEAK)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.WEAK)
    c3 = UOp.param(1, dtypes.int, (512,))
    c4 = c3.index(c1)
-    c5 = UOp.range(UOp.const(240), 0, AxisType.REDUCE)
-    c6 = ((c2*UOp.const(240))+c5)
+    c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
+    c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5)
    c7 = UOp.param(2, dtypes.uchar, (60000,))
    c8 = c7.index(c6)
    c9 = ((c4<0).where((c4+60000), c4)!=c6.cast(dtypes.int)).where(0, c8.cast(dtypes.uint).cast(dtypes.uchar)).reduce(c5, arg=Ops.ADD)
-    c10 = c0.index(((c1*UOp.const(250))+c2)).store(c9).end(c1, c2)
+    c10 = c0.index(((c1*UOp.const(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2)
    uops = to_uops_list([c10])
    for u in uops:
      self.assertNotEqual(u.dtype, dtypes.long)
@@ -450,19 +467,19 @@ class TestUOpGraph(unittest.TestCase):
  def test_load_idx_no_math_on_loaded(self):
    # test the (x+y)<c pattern where x has loads - we shouldn't do math on loaded indices
    c0 = UOp.param(0, dtypes.uchar, (128000,))
-    c1 = UOp.range(UOp.const(512), 1, AxisType.WEAK)
-    c2 = UOp.range(UOp.const(250), 2, AxisType.WEAK)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.WEAK)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.WEAK)
    c3 = UOp.param(1, dtypes.int, (512,))
    c4 = c3.index(c1)  # c4 is a load
-    c5 = UOp.range(UOp.const(240), 0, AxisType.REDUCE)
-    c6 = ((c2*UOp.const(240))+c5)
+    c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
+    c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5)
    c7 = UOp.param(2, dtypes.uchar, (60000,))
    c8 = c7.index(c6)
    # (loaded + range) < const pattern - loaded value shouldn't be promoted to long
    loaded_idx = c4.cast(dtypes.weakint)
-    comparison = (loaded_idx + c5) < UOp.const(60000)
+    comparison = (loaded_idx + c5) < UOp.const(dtypes.weakint, 60000)
    c9 = comparison.where(c8.cast(dtypes.uint).cast(dtypes.uchar), 0).reduce(c5, arg=Ops.ADD)
-    c10 = c0.index(((c1*UOp.const(250))+c2)).store(c9).end(c1, c2)
+    c10 = c0.index(((c1*UOp.const(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2)
    uops = to_uops_list([c10])
    for u in uops:
      self.assertNotEqual(u.dtype, dtypes.long)
@@ -471,9 +488,9 @@ class TestUOpGraph(unittest.TestCase):
    glbl0 = UOp.param(0, dtypes.int, (1,))
    glbl1 = UOp.param(1, dtypes.int, (1,))
    glbl2 = UOp.param(2, dtypes.int, (1,))
-    idx = UOp.const(0)
+    idx = UOp.const(dtypes.int, 0)
    ld0 = glbl1.index(UOp.invalid())
-    ld1 = glbl2.index(idx.valid(UOp.const(True)))
+    ld1 = glbl2.index(idx.valid(UOp.const(dtypes.bool, True)))
    uops = to_uops_list([glbl0.index(idx).store(ld1+ld0)])
    # the gate and invalid value are deleted from ld1
    self.assertEqual(len([u for u in uops if u.op is Ops.LOAD]), 1)
@@ -485,7 +502,7 @@ class TestUOpGraph(unittest.TestCase):
    st = smem.index(lidx).store(glbl0.index(lidx).load())
    barrier = st.barrier()
    ld0 = smem.after(barrier).index(UOp.invalid())
-    ld1 = smem.after(barrier).index((lidx+2).valid(UOp.const(True)))
+    ld1 = smem.after(barrier).index((lidx+2).valid(UOp.const(dtypes.bool, True)))
    uops = to_uops_list([glbl0.index(lidx).store(ld1+ld0)])

    # the gate and invalid value are deleted from ld1
@@ -493,10 +510,10 @@ class TestUOpGraph(unittest.TestCase):

  def test_fold_gated_store(self):
    glbl = UOp.param(0, dtypes.int, (1,))
-    idx0 = UOp.const(0)
-    val = UOp.const(42)
+    idx0 = UOp.const(dtypes.int, 0)
+    val = UOp.const(dtypes.int, 42)
    st0 = glbl.index(UOp.invalid()).store(val)
-    st1 = glbl.index(idx0.valid(UOp.const(True))).store(val)
+    st1 = glbl.index(idx0.valid(UOp.const(dtypes.bool, True))).store(val)
    uops = to_uops_list([st0, st1])
    # only the second store happens
    self.assertEqual(len([u for u in uops if u.op is Ops.STORE]), 1)
@@ -504,9 +521,9 @@ class TestUOpGraph(unittest.TestCase):
  @unittest.skip("this is a uop type error")
  def test_asserts_bad_gate(self):
    glbl0 = UOp.param(0, dtypes.int, (1,))
-    idx = UOp.const(0)
-    bad_gate = UOp.const(1)
-    with self.assertRaises(AssertionError): to_uops_list([UOp(Ops.STORE, src=(glbl0, idx, UOp.const(42), bad_gate))])
+    idx = UOp.const(dtypes.int, 0)
+    bad_gate = UOp.const(dtypes.int, 1)
+    with self.assertRaises(AssertionError): to_uops_list([UOp(Ops.STORE, src=(glbl0, idx, UOp.const(dtypes.int, 42), bad_gate))])

  def test_after_end(self):
    r = UOp.range(10, 0)
@@ -514,7 +531,7 @@ class TestUOpGraph(unittest.TestCase):
    c = r + 1
    self.assertIn(r, c.ranges)

-    e = UOp.const(1).end(r)
+    e = UOp.const(dtypes.int, 1).end(r)
    self.assertNotIn(r, e.ranges)

    a = c.after(e)
@@ -537,15 +554,6 @@ 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
@@ -573,7 +581,7 @@ class TestConstBufferize(unittest.TestCase):
    CONST doesn't depend on ranges (constant is same value everywhere).
    """
    from tinygrad.schedule.rangeify import pm_const_buffer_folding, BufferizeOpts
-    c = UOp.const(42.0)
+    c = UOp.const(dtypes.float, 42.0)
    r1 = UOp.range(3, 0)
    bufferize_with_range = UOp(Ops.STAGE, src=(c, r1), arg=BufferizeOpts(device="CPU"))
    self.assertEqual(len(bufferize_with_range.src), 2)  # const + 1 range
@@ -581,13 +589,13 @@ class TestConstBufferize(unittest.TestCase):
    result = graph_rewrite(bufferize_with_range, pm_const_buffer_folding, name='test')
    # BUFFERIZE should be removed, result is const broadcast to shape
    self.assertNotEqual(result.op, Ops.STAGE)
-    const_vals = [u.val for u in result.toposort() if u.op is Ops.CONST and u.dtype is dtypes.weakfloat]
+    const_vals = [u.arg for u in result.toposort() if u.op is Ops.CONST and u.dtype == dtypes.float]
    self.assertIn(42.0, const_vals)

  def test_const_bufferize_with_multiple_ranges(self):
    """Test CONST.BUFFERIZE with multiple ranges is also folded."""
    from tinygrad.schedule.rangeify import pm_const_buffer_folding, BufferizeOpts
-    c = UOp.const(3.14)
+    c = UOp.const(dtypes.float, 3.14)
    r1 = UOp.range(3, 0)
    r2 = UOp.range(4, 1)
    bufferize_with_ranges = UOp(Ops.STAGE, src=(c, r1, r2), arg=BufferizeOpts(device="CPU"))
@@ -596,14 +604,14 @@ class TestConstBufferize(unittest.TestCase):
    result = graph_rewrite(bufferize_with_ranges, pm_const_buffer_folding, name='test')
    # BUFFERIZE should be removed
    self.assertNotEqual(result.op, Ops.STAGE)
-    const_vals = [u.val for u in result.toposort() if u.op is Ops.CONST and u.dtype is dtypes.weakfloat]
+    const_vals = [u.arg for u in result.toposort() if u.op is Ops.CONST and u.dtype == dtypes.float]
    self.assertIn(3.14, const_vals)

 class TestUOpTags(unittest.TestCase):
  def test_inc_by_one(self):
-    g = UOp.const(1) + UOp.const(1)
+    g = UOp.const(dtypes.int, 1) + UOp.const(dtypes.int, 1)
    assert g.ssimplify() == 2
-    pm_plus_1 = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: UOp.const(x.val+1, x.dtype).rtag(1) if x.tag is None else None)])
+    pm_plus_1 = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.replace(arg=x.arg+1, tag=1) if x.tag is None else None)])
    pm_strip_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
    g = graph_rewrite(g, pm_plus_1)
    assert g.ssimplify() == 4
@@ -692,36 +700,36 @@ class TestUOpGetItem(unittest.TestCase):

 class TestUOpBroadcast(unittest.TestCase):
  def test_broadcast_row(self):
-    a = UOp.const(1.0).expand((4, 8))
-    b = UOp.const(2.0).expand((4, 1))
+    a = UOp.const(dtypes.float, 1, shape=(4, 8))
+    b = UOp.const(dtypes.float, 2, shape=(4, 1))
    c = a + b
    self.assertEqual(c.shape, (4, 8))
    self.assertEqual(c.op, Ops.ADD)

  def test_broadcast_col(self):
-    a = UOp.const(1.0).expand((4, 8))
-    b = UOp.const(2.0).expand((1, 8))
+    a = UOp.const(dtypes.float, 1, shape=(4, 8))
+    b = UOp.const(dtypes.float, 2, shape=(1, 8))
    c = a + b
    self.assertEqual(c.shape, (4, 8))
    self.assertEqual(c.op, Ops.ADD)

  def test_broadcast_lower_dim(self):
-    a = UOp.const(1.0).expand((4, 8))
-    b = UOp.const(2.0).expand((8,))
+    a = UOp.const(dtypes.float, 1, shape=(4, 8))
+    b = UOp.const(dtypes.float, 2, shape=(8,))
    c = a * b
    self.assertEqual(c.shape, (4, 8))
    self.assertEqual(c.op, Ops.MUL)

  def test_broadcast_scalar(self):
-    a = UOp.const(1.0).expand((4, 8))
+    a = UOp.const(dtypes.float, 1, shape=(4, 8))
    c = a * 2
    self.assertEqual(c.shape, (4, 8))
    self.assertEqual(c.op, Ops.MUL)

  def test_broadcast_symbolic_same_shape(self):
    t = Variable("t", 1, 10)
-    a = UOp.const(1.0).expand((1, 1, t))
-    b = UOp.const(2.0).expand((1, 1, t))
+    a = UOp.const(dtypes.float, 1, shape=(1, 1, t))
+    b = UOp.const(dtypes.float, 2, shape=(1, 1, t))
    c = a + b
    self.assertEqual(c.op, Ops.ADD)

@@ -1,34 +1,34 @@
 import unittest
-from tinygrad import UOp
+from tinygrad import UOp, dtypes

 class TestUOpRepr(unittest.TestCase):
  def test_simple_const(self):
-    a = UOp.const(42)
-    self.assertEqual(repr(a), "UOp(Ops.CONST, dtypes.weakint, arg=42, src=())")
+    a = UOp.const(dtypes.int, 42)
+    self.assertEqual(repr(a), "UOp(Ops.CONST, dtypes.int, arg=42, src=())")
  def test_different_consts(self):
-    a, b = UOp.const(42), UOp.const(3)
+    a, b = UOp.const(dtypes.int, 42), UOp.const(dtypes.int, 3)
    expected = (
-      "UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
-      "  UOp(Ops.CONST, dtypes.weakint, arg=42, src=()),\n" +
-      "  UOp(Ops.CONST, dtypes.weakint, arg=3, src=()),))"
+      "UOp(Ops.ADD, dtypes.int, arg=None, src=(\n" +
+      "  UOp(Ops.CONST, dtypes.int, arg=42, src=()),\n" +
+      "  UOp(Ops.CONST, dtypes.int, arg=3, src=()),))"
    )
    self.assertEqual(repr(a+b), expected)
  def test_walrus_operator_indentation(self):
    # The reference should have the same indentation as the definition
-    a = UOp.const(42)
+    a = UOp.const(dtypes.int, 42)
    expected = (
-      "UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
-      "  x0:=UOp(Ops.CONST, dtypes.weakint, arg=42, src=()),\n" +
+      "UOp(Ops.ADD, dtypes.int, arg=None, src=(\n" +
+      "  x0:=UOp(Ops.CONST, dtypes.int, arg=42, src=()),\n" +
      "  x0,))"
    )
    self.assertEqual(repr(a+a), expected)
  def test_nested_walrus_indentation(self):
    # Ensure indentation is consistent at multiple levels
-    b = (a:=UOp.const(1)) + a
+    b = (a:=UOp.const(dtypes.int, 1)) + a
    expected = (
-      "UOp(Ops.MUL, dtypes.weakint, arg=None, src=(\n" +
-      "  x0:=UOp(Ops.ADD, dtypes.weakint, arg=None, src=(\n" +
-      "    x1:=UOp(Ops.CONST, dtypes.weakint, arg=1, src=()),\n" +
+      "UOp(Ops.MUL, dtypes.int, arg=None, src=(\n" +
+      "  x0:=UOp(Ops.ADD, dtypes.int, arg=None, src=(\n" +
+      "    x1:=UOp(Ops.CONST, dtypes.int, arg=1, src=()),\n" +
      "    x1,)),\n" +
      "  x0,))"
    )
@@ -4,48 +4,44 @@ from tinygrad.uop.ops import UOp, resolve

 class TestUOpResolve(unittest.TestCase):
  def test_simple_int(self):
-    u = UOp.const(4, dtypes.int)
+    u = UOp.const(dtypes.int, 4)
    self.assertEqual(int(u), 4)

-  def test_weak_const(self):
-    self.assertEqual(int(UOp.const(5)), 5)
-    self.assertEqual(float(UOp.const(1.5)), 1.5)
-
  def test_int_add(self):
-    u = UOp.const(4, dtypes.int) + 7
+    u = UOp.const(dtypes.int, 4) + 7
    self.assertEqual(int(u), 11)

  def test_lt(self):
-    u = UOp.const(4) < 7
+    u = UOp.const(dtypes.int, 4) < 7
    self.assertTrue(u)

  def test_rfloordiv(self):
-    u = 8 // UOp.const(4, dtypes.int)
+    u = 8 // UOp.const(dtypes.int, 4)
    self.assertEqual(int(u), 2)

  def test_rtruediv(self):
-    u = 9 / UOp.const(4, dtypes.float)
+    u = 9 / UOp.const(dtypes.float, 4)
    self.assertEqual(float(u), 2.25)

  def test_leq(self):
-    u = UOp.const(4) <= 4
+    u = UOp.const(dtypes.int, 4) <= 4
    self.assertTrue(u)

  def test_ne(self):
-    u = UOp.const(4) != 7
+    u = UOp.const(dtypes.int, 4) != 7
    self.assertTrue(u)

  def test_ne_f(self):
-    u = UOp.const(4) != 4
+    u = UOp.const(dtypes.int, 4) != 4
    self.assertFalse(u)

  def test_ngt(self):
-    u = UOp.const(4) > 7
+    u = UOp.const(dtypes.int, 4) > 7
    self.assertFalse(u)

  def test_ssimplify(self):
-    self.assertEqual((8 % UOp.const(4)).ssimplify(), 0)
-    self.assertEqual((8 * UOp.const(4)).ssimplify(), 32)
+    self.assertEqual((8 % UOp.const(dtypes.int, 4)).ssimplify(), 0)
+    self.assertEqual((8 * UOp.const(dtypes.int, 4)).ssimplify(), 32)

  def test_ambiguous_less_than(self):
    u = UOp.variable("i", 1, 10)
@@ -56,7 +52,7 @@ class TestUOpResolve(unittest.TestCase):
    self.assertFalse(resolve(u < -1, True))

  def test_float_direct(self):
-    u = UOp.const(4.5, dtypes.float) + 7
+    u = UOp.const(dtypes.float, 4.5) + 7
    self.assertEqual(float(u), 11.5)

  def test_var_cmp_t(self):
@@ -3,21 +3,22 @@ import unittest, pickle, functools, math
 import z3

 from tinygrad.dtype import dtypes, ConstType, DType, Invalid
+from test.helpers import get_uops
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad.uop.spec import spec_shared, type_verify
-from tinygrad.uop.symbolic import sym, pm_fold_cast_const, commutative, pm_simplify_valid, pm_move_where_on_load
+from tinygrad.uop.symbolic import sym, 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):
  sym_vars = {v.render():v for v in v.toposort() if v.op in (Ops.RANGE, Ops.SPECIAL, Ops.PARAM)}
  s_eval = eval(s, sym_vars)
-  if isinstance(s_eval, (bool, int, float)): s_eval = UOp.const(s_eval)
+  if isinstance(s_eval, int) and v.dtype==dtypes.weakint: s_eval = UOp.const(dtypes.weakint, s_eval)
+  elif isinstance(s_eval, (bool, int, float)): s_eval = UOp.const(dtypes.from_py(s_eval), s_eval)
  s_eval = graph_rewrite(s_eval, commutative, name="cannonicalize eval")
  self.assertIs(s_eval, v, f"eval did not match simplified: {s_eval} != {v.render()} for {s}")

-def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.weakint):
-  return UOp.variable(name, min_val, max_val, dtype, param=True)
-def uconst(val): return UOp.const(val)
+def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.weakint): return UOp.variable(name,min_val,max_val,dtype)
+def uconst(val): return UOp.const(dtypes.weakint, val)
 def usum(ops): return functools.reduce(lambda x,y: x+y, ops)
 def uand(ops): return functools.reduce(lambda x,y: x*y, ops)

@@ -35,7 +36,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+pm_fold_cast_const, name="simplify symbolic uop")
+    v_simplified = graph_rewrite(v, sym, 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)
@@ -103,16 +104,16 @@ class TestSymbolic(unittest.TestCase):
    self.assertEqual(UOp.gcd(a, a*b, a*3).simplify(), a)
    self.assertEqual(UOp.gcd(a*a*a, a*b*a, a*3*a).simplify(), a*a)
    self.assertEqual(UOp.gcd(a*a*10, b*a*5, a*a*5).simplify(), a*5)
-    self.assertEqual(UOp.gcd(a*10, b*5, a*5).simplify(), uconst(5))
-    self.assertEqual(UOp.gcd(a, b*5, a*5).simplify(), uconst(1))
+    self.assertEqual(UOp.gcd(a*10, b*5, a*5).simplify(), a.const_like(5))
+    self.assertEqual(UOp.gcd(a, b*5, a*5).simplify(), a.const_like(1))

  def test_divides_exact(self):
    a = Variable("a", 1, 8)
    b = Variable("b", 1, 8)
    self.assertEqual((a*a*3).divide_exact(a).simplify(), a*3)
-    self.assertEqual((a*a*3).divide_exact(a*a*3).simplify(), uconst(1))
-    self.assertEqual((a*a*6).divide_exact(a*a*3).simplify(), uconst(2))
-    self.assertEqual((a*b*3).divide_exact(uconst(3)).simplify(), a*b)
+    self.assertEqual((a*a*3).divide_exact(a*a*3).simplify(), a.const_like(1))
+    self.assertEqual((a*a*6).divide_exact(a*a*3).simplify(), a.const_like(2))
+    self.assertEqual((a*b*3).divide_exact(a.const_like(3)).simplify(), a*b)
    self.assertEqual((a*a*3).divide_exact(a*(-3)).simplify(), a*-1)
    self.assertEqual((a*a*b*3).divide_exact(a*b).simplify(), a*3)
    self.assertEqual((a*3+a*b).divide_exact(a).simplify(), b+3)
@@ -442,7 +443,7 @@ class TestSymbolic(unittest.TestCase):
    self.helper_test_variable(uand([uconst(1), Variable("a", 0, 1)]), 0, 1, "a")

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

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

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

@@ -882,7 +883,7 @@ class TestSymbolic(unittest.TestCase):
    idx = Variable("idx", 0, 24)
    self.helper_test_variable(idx//4, 0, 6, "(idx//4)")
    # TODO: simplify the true branch
-    self.helper_test_variable((idx<4).where(idx//4, uconst(-1)), -1, 6, "(idx<4).where((idx//4), -1)")
+    self.helper_test_variable((idx<4).where(idx//4, idx.const_like(-1)), -1, 6, "(idx<4).where((idx//4), -1)")

  def test_floordiv_lt(self):
    # x//d<c <=> x<c*d for d>0, and <=> c*d<x for d<0
@@ -941,18 +942,13 @@ class TestSymbolic(unittest.TestCase):

  def test_where_removal(self):
    cond = Variable("a", 0, 3) < 2
-    u1, u0 = UOp.const(True), UOp.const(False)
+    u1, u0 = cond.const_like(True), cond.const_like(False)
    self.helper_test_variable(cond, 0, 1, "(a<2)")
    self.helper_test_variable(cond.where(u1, u0), 0, 1, "(a<2)")
    self.helper_test_variable(cond.where(u1, u0).where(u1, u0), 0, 1, "(a<2)")
    self.helper_test_variable(cond.where(u0, u1), 0, 1, "((a<2)!=True)")
    self.helper_test_variable(cond.where(u0, u1).where(u0, u1), 0, 1, "(a<2)")

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

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

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

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

  def test_variable_min_eq_max_bind_folds(self):
-    b = UOp.variable("x", 1, 1).bind(1)
+    b = Variable("x", 1, 1).bind(1)
    s = b.simplify()
    self.assertEqual(s.op, Ops.CONST)
-    self.assertEqual(s.val, 1)
+    self.assertEqual(s.arg, 1)

 class TestSymInfer(unittest.TestCase):
  def test_sym_infer(self):
@@ -1191,13 +1203,9 @@ class TestSymInfer(unittest.TestCase):
    # floor: 1 % -1000 = -999, 1 // -1000 = -1
    assert sym_infer(a%b, var_vals) == -999
    assert sym_infer(a//b, var_vals) == -1
-  def test_sym_infer_with_cast(self):
-    a = Variable("a", 0, 100, dtypes.int)
-    assert sym_infer(a.cast(dtypes.long) + 1, {a.expr: 5}) == 6
-    assert sym_infer(a.cast(dtypes.float) * 0.5, {a.expr: 5}) == 2.5
  def test_sym_infer_with_bitcast(self):
    a = Variable("a", 1, 10, dtypes.int)
-    expr = ((a.bitcast(dtypes.uint) << UOp.const(1)).bitcast(dtypes.int) + 2)
+    expr = ((a.bitcast(dtypes.uint) << UOp.const(dtypes.uint, 1)).bitcast(dtypes.int) + 2)
    ret = sym_infer(expr, {a.expr: 2})
    assert isinstance(ret, int)
    assert ret == 6
@@ -1208,7 +1216,7 @@ class TestSymInfer(unittest.TestCase):
    c = Variable("c", 0, 0xFFFFFFFF, dtypes.uint)
    assert sym_infer(c.bitcast(dtypes.int), {c.expr: 0xFFFFFFFF}) == -1

-    assert sym_infer(UOp.const(1.5).cast(dtypes.float).bitcast(dtypes.uint), {}) == 1069547520
+    assert sym_infer(UOp.const(dtypes.float, 1.5).bitcast(dtypes.uint), {}) == 1069547520

  def test_sym_infer_deeply_nested(self):
    # build an expression that exceeds Python's nested parentheses limit for eval
@@ -1329,7 +1337,7 @@ class TestInvalidIndex(unittest.TestCase):
  def test_invalid_times_0(self):
    ridx = Variable("ridx", 0, 10)
    idx = (ridx<5).where(ridx, UOp.invalid())*0
-    self.assertIs(idx.simplify(), (ridx<5).where(uconst(0), UOp.invalid()),
+    self.assertIs(idx.simplify(), (ridx<5).where(UOp.const(dtypes.weakint, 0), UOp.invalid()),
                  "multiplying an index by 0 should preserve the invalid")

  def test_alu_moves_inside_invalid(self):
@@ -1354,34 +1362,24 @@ class TestInvalidIndex(unittest.TestCase):
    self.assertIs((UOp.invalid()<Variable("a",0,10)).simplify().dtype, dtypes.bool)

  def test_alu_invalid_vconst(self):
-    c1 = UOp.const((1, 1, Invalid, Invalid))
-    c2 = UOp.const((1, Invalid, 1, 1))
-    self.assertIs((c1+c2).simplify(), UOp.const((2, Invalid, Invalid, Invalid)))
-
-  def test_gated_load_keeps_index_valid(self):
-    # the load executes even on gated-off iterations: gated_given_valid must not erase its mask (PADTO OOB shape)
-    buf = UOp.param(0, dtypes.bool, (17,))
-    ridx = Variable("ridx", 0, 31)
-    cond = ridx < 17
-    load = buf.index(ridx.valid(cond))
-    out = graph_rewrite(cond.where(load.where(uconst(2), uconst(0)), UOp.invalid()), sym)
-    idx = next(u for u in out.toposort() if u.op is Ops.INDEX)
-    self.assertIs(idx.src[1].get_valid(), cond.simplify())
+    c1 = UOp.const(dtypes.weakint, (1, 1, Invalid, Invalid))
+    c2 = UOp.const(dtypes.weakint, (1, Invalid, 1, 1))
+    self.assertIs((c1+c2).simplify(), UOp.const(dtypes.weakint, (2, Invalid, Invalid, Invalid)))

 class TestStoreLoadFolding(unittest.TestCase):
  """Tests for store(index, load(index)) -> NOOP rule. This rule matches patterns that EMERGE during simplification."""
  def test_store_load_folding(self):
    # store(idx, load(idx)) -> NOOP, including emergent patterns like store(idx, load(idx) + 0)
    buf = UOp.param(0, dtypes.int, (1,))
-    index = buf.index(UOp.const(0))
+    index = buf.index(UOp.const(dtypes.weakint, 0))
    # Direct: store(idx, load(idx)) -> NOOP
    self.assertEqual(graph_rewrite(index.store(index.load()), sym).op, Ops.NOOP)
    # Emergent: store(idx, load(idx) + 0) -> store(idx, load(idx)) -> NOOP
-    self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(0)), sym).op, Ops.NOOP)
+    self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(dtypes.int, 0)), sym).op, Ops.NOOP)
    # Emergent: store(idx, load(idx) * 1) -> store(idx, load(idx)) -> NOOP
-    self.assertEqual(graph_rewrite(index.store(index.load() * UOp.const(1)), sym).op, Ops.NOOP)
+    self.assertEqual(graph_rewrite(index.store(index.load() * UOp.const(dtypes.int, 1)), sym).op, Ops.NOOP)
    # Negative: store(idx, load(idx) + 1) should NOT fold
-    self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(1)), sym).op, Ops.STORE)
+    self.assertEqual(graph_rewrite(index.store(index.load() + UOp.const(dtypes.int, 1)), sym).op, Ops.STORE)

 class TestMoveWhereOnLoad(unittest.TestCase):
  def test_bool_index_preserves_dtype(self):
@@ -1392,7 +1390,7 @@ class TestMoveWhereOnLoad(unittest.TestCase):
    cond = (a < 4) & (r < 2)
    valid = (a < 2)  # pre-existing valid on the load (to pass can_move check for the r-only clause)
    idx = buf.index(a.valid(valid))
-    expr = cond.where(idx, UOp.const(0))
+    expr = cond.where(idx, idx.const_like(0))
    out = graph_rewrite(expr, pm_move_where_on_load)
    type_verify(out, spec_shared)  # Invalid matches any dtype

@@ -1477,7 +1475,7 @@ class TestFuzzFailure(unittest.TestCase):
    v2=Variable('v2', 0, 2)
    v3=Variable('v3', 0, 1)
    expr = (((((((((((((((((((((((0//4)%2)//8)+-2)+-4)+-3)+v1)+-4)+v2)+-2)+v3)+v2)//3)%7)*1)//2)+v2)*-1)+2)+1)+0)+-3)+v3)
-    v1_val, v2_val, v3_val = UOp.const(8), UOp.const(0), UOp.const(0)
+    v1_val, v2_val, v3_val = v1.const_like(8), v2.const_like(0), v3.const_like(0)
    num = expr.simplify().substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    rn = expr.substitute({v1:v1_val, v2:v2_val, v3:v3_val}).ssimplify()
    assert num==rn, f"{num} != {rn}"
@@ -1,16 +1,16 @@
 import unittest, math
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.dtype import dtypes, Invalid, truncate
+from tinygrad.dtype import dtypes, Invalid

 class TestVminVmaxProperties(unittest.TestCase):
  def test_vmin_vmax_constant(self):
    # vmin and vmax for a constant
-    uop = UOp.const(42)
+    uop = UOp.const(dtypes.int32, 42)
    self.assertEqual(uop.vmin, 42)
    self.assertEqual(uop.vmax, 42)

  def test_vmin_vmax_cmpne(self):
-    uop = UOp.const(42)
+    uop = UOp.const(dtypes.int32, 42)
    def test_bool(u, x):
      self.assertEqual(u.vmin, x)
      self.assertEqual(u.vmax, x)
@@ -81,8 +81,8 @@ class TestVminVmaxProperties(unittest.TestCase):

  def test_vmin_vmax_multiplication_0_inf(self):
    # vmin and vmax for multiplication with a variable
-    x = UOp.const(0.0)
-    y = UOp.load(UOp.param(0, dtypes.float, (1,)), UOp.const(0), dtype=dtypes.float)
+    x = UOp.const(dtypes.float, 0.0)
+    y = UOp.load(UOp.param(0, dtypes.float, (1,)), UOp.const(dtypes.int, 0), dtype=dtypes.float)
    uop = x * y
    # TODO: these should be 0, but definitely should not be nan
    self.assertEqual(uop.vmin, -math.inf)
@@ -162,22 +162,12 @@ class TestVminVmaxProperties(unittest.TestCase):
    self.assertEqual(x_uint.vmin, dtypes.uint.min)
    self.assertEqual(x_uint.vmax, dtypes.uint.max)

-  def test_vmin_vmax_cast_float_to_int(self):
-    self.assertEqual(UOp.variable('x', -4.5, 4.5, dtypes.float).cast(dtypes.int)._min_max, (-4, 4))
-    self.assertEqual(UOp.const(4.5).cast(dtypes.float).cast(dtypes.int)._min_max, (4, 4))
-    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)

  def test_vmin_vmax_invalid_vconst(self):
-    x = UOp.const((0, 4, Invalid, Invalid))
+    x = UOp.const(dtypes.weakint, (0, 4, Invalid, Invalid))
    self.assertEqual((x.vmin, x.vmax), (0, 4))

 class TestVminVmaxDivMod(unittest.TestCase):
@@ -208,14 +198,14 @@ class TestVminVmaxDivMod(unittest.TestCase):

  def test_vmin_vmax_floordiv_floormod(self):
    x = UOp.variable('x', -7, 7)
-    floordiv = x.alu(Ops.FLOORDIV, UOp.const(3))
+    floordiv = x.alu(Ops.FLOORDIV, x.const_like(3))
    self.assertEqual(floordiv.vmin, -3)
    self.assertEqual(floordiv.vmax, 2)
-    floormod = x.alu(Ops.FLOORMOD, UOp.const(3))
+    floormod = x.alu(Ops.FLOORMOD, x.const_like(3))
    self.assertEqual(floormod.vmin, 0)
    self.assertEqual(floormod.vmax, 2)
    # negative const divisor: floormod range is [c+1, 0]
-    floormod_neg = x.alu(Ops.FLOORMOD, UOp.const(-3))
+    floormod_neg = x.alu(Ops.FLOORMOD, x.const_like(-3))
    self.assertEqual(floormod_neg.vmin, -2)
    self.assertEqual(floormod_neg.vmax, 0)

@@ -296,44 +286,44 @@ class TestVminVmaxDivMod(unittest.TestCase):
 class TestVminVmaxVConst(unittest.TestCase):
  def test_vmin_vmax_vconst_single_element(self):
    # vmin and vmax for a single-element vector constant
-    uop = UOp.const((42,))
+    uop = UOp.const(dtypes.int32, (42,))
    self.assertEqual(uop.vmin, 42)
    self.assertEqual(uop.vmax, 42)

  def test_vmin_vmax_vconst_multiple_elements(self):
    # vmin and vmax for a multi-element vector constant
-    uop = UOp.const((10, 20, -5, 7))
+    uop = UOp.const(dtypes.int32, (10, 20, -5, 7))
    self.assertEqual(uop.vmin, -5)
    self.assertEqual(uop.vmax, 20)

  def test_vmin_vmax_vconst_all_equal(self):
    # vmin and vmax for a vector where all elements are equal
-    uop = UOp.const((7, 7, 7))
+    uop = UOp.const(dtypes.int32, (7, 7, 7))
    self.assertEqual(uop.vmin, 7)
    self.assertEqual(uop.vmax, 7)

  def test_vmin_vmax_vconst_with_negative_values(self):
    # vmin and vmax for a vector constant containing negative values
-    uop = UOp.const((-10, -20, -5, -15))
+    uop = UOp.const(dtypes.int32, (-10, -20, -5, -15))
    self.assertEqual(uop.vmin, -20)
    self.assertEqual(uop.vmax, -5)

  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, truncate[dtypes.default_float](-3.2))
-    self.assertEqual(uop.vmax, truncate[dtypes.default_float](1.5))
+    uop = UOp.const(dtypes.float32, (1.5, -3.2, 0.0))
+    self.assertEqual(uop.vmin, -3.2)
+    self.assertEqual(uop.vmax, 1.5)

  def test_vmin_vmax_vconst_with_bools(self):
    # vmin and vmax for a vector constant of bool values
-    uop = UOp.const((True, False, False))
+    uop = UOp.const(dtypes.bool, (True, False, False))
    self.assertIs(uop.vmin, False)
    self.assertIs(uop.vmax, True)

  def test_vmin_vmax_vector_with_gep(self):
    # vmin and vmax for a vector constant of bool values
    d1 = UOp.param(1, dtypes.int, (1,))
-    idx = UOp.const(0)
+    idx = UOp.const(dtypes.int, 0)
    val = UOp(Ops.LOAD, src=(d1.index(idx),))
    uop = (val // 32)
    self.assertEqual(uop.vmin, -67108864)
@@ -342,17 +332,17 @@ class TestVminVmaxVConst(unittest.TestCase):
 class TestConstFactor(unittest.TestCase):
  def test_const_factor_constant(self):
    # const_factor for a constant
-    uop = UOp.const(42)
+    uop = UOp.const(dtypes.int32, 42)
    self.assertEqual(uop.const_factor(), 42)

  def test_const_factor_addition(self):
    # const_factor for an addition of constants
-    uop = UOp.const(30) + UOp.const(12)
+    uop = UOp.const(dtypes.int32, 30) + UOp.const(dtypes.int32, 12)
    self.assertEqual(uop.const_factor(), 6)  # GCD(30, 12) = 6

  def test_const_factor_multiplication(self):
    # const_factor for a multiplication of constants
-    uop = UOp.const(5) * UOp.const(7)
+    uop = UOp.const(dtypes.int32, 5) * UOp.const(dtypes.int32, 7)
    self.assertEqual(uop.const_factor(), 5)  # For multiplication, it's one of the factors

  def test_const_factor_with_variable(self):
@@ -387,14 +377,14 @@ class TestConstFactor(unittest.TestCase):
 class TestDivides(unittest.TestCase):
  def test_divides_constant_exact(self):
    # Divides a constant by an exact divisor
-    uop = UOp.const(42)
+    uop = UOp.const(dtypes.int32, 42)
    result = uop.divides(7)
    self.assertIsNotNone(result)
    self.assertEqual(result.const_factor(), 6)  # 42 / 7 = 6

  def test_divides_constant_inexact(self):
    # Try to divide a constant by a non-exact divisor
-    uop = UOp.const(42)
+    uop = UOp.const(dtypes.int32, 42)
    result = uop.divides(5)
    self.assertIsNone(result)  # 42 is not divisible by 5

@@ -5,25 +5,23 @@ 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  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
-from tinygrad.uop.weak import pm_lower_index_dtype
+from tinygrad.uop.ops import Ops, ParamArg, 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.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

 class TestDTypeFromUOp(unittest.TestCase):
  def test_broadcastable_promotion(self):
-    self.assertEqual(dtype_from_uop(Ops.ADD, (UOp.const(1.0).cast(dtypes.float32), UOp.const(1.0).cast(dtypes.float16)), None), dtypes.float32)
-    self.assertEqual(dtype_from_uop(Ops.MUL, (UOp.const(1).cast(dtypes.int8), UOp.const(1).cast(dtypes.int32)), None), dtypes.int32)
+    self.assertEqual(dtype_from_uop(Ops.ADD, (UOp.const(dtypes.float32, 1.0), UOp.const(dtypes.float16, 1.0)), None), dtypes.float32)
+    self.assertEqual(dtype_from_uop(Ops.MUL, (UOp.const(dtypes.int8, 1), UOp.const(dtypes.int32, 1)), None), dtypes.int32)

  def test_same_dtype_fast_path(self):
-    src = (UOp.const(1), UOp.const(2))
+    src = (UOp.const(dtypes.weakint, 1), UOp.const(dtypes.weakint, 2))
    self.assertEqual(dtype_from_uop(Ops.ADD, src, None), dtypes.weakint)

  def test_where_promotion(self):
-    cond = UOp.const(True)
-    srcs = (cond, UOp.const(1.0).cast(dtypes.float32), UOp.const(1.0).cast(dtypes.float16))
-    self.assertEqual(dtype_from_uop(Ops.WHERE, srcs, None), dtypes.float32)
+    cond = UOp.const(dtypes.bool, True)
+    self.assertEqual(dtype_from_uop(Ops.WHERE, (cond, UOp.const(dtypes.float32, 1.0), UOp.const(dtypes.float16, 1.0)), None), dtypes.float32)
    idx = UOp.range(4, 0)
    self.assertEqual(idx.valid(idx < 4).dtype, dtypes.weakint)

@@ -39,39 +37,33 @@ class TestDTypeFromUOp(unittest.TestCase):
    self.assertEqual(UOp(Ops.CONST, arg=True).dtype, dtypes.bool)
    self.assertEqual(UOp(Ops.CONST, arg=Invalid).dtype, dtypes.bool)
    # an explicit (strong) const dtype is legal until the field is removed
-    self.assertEqual(UOp.const(3, dtypes.int32).dtype, dtypes.int32)
+    self.assertEqual(UOp.const(dtypes.int32, 3).dtype, dtypes.int32)

  def test_weak_dtype_rejected_by_program_spec(self):
    for weak, concrete, value in ((dtypes.weakint, dtypes.int32, 1), (dtypes.weakfloat, dtypes.float32, 1.0)):
-      with self.assertRaises(RuntimeError): type_verify(UOp.const(value, weak).sink(), spec_program)
-      type_verify(UOp.const(value, concrete).sink(), spec_program)
-
-  def test_invalid_stated_dtype(self):
-    # UOp.const normalizes a stated dtype away (const_like/full pass their position's); the core constructor does not,
-    # and the spec is what rejects a non-bool Invalid
-    self.assertIs(UOp.const(Invalid, dtypes.float32), UOp.invalid())
-    with self.assertRaises(RuntimeError): type_verify(UOp(Ops.CONST, dtypes.float32, arg=Invalid), spec_shared)
+      with self.assertRaises(RuntimeError): type_verify(UOp.const(weak, value).sink(), spec_program)
+      type_verify(UOp.const(concrete, value).sink(), spec_program)

  def test_invalid_dtype_and_consumers(self):
    invalid = UOp.invalid()
    self.assertIs(invalid.dtype, dtypes.bool)
-    self.assertIs(UOp.const(Invalid, dtypes.float32), invalid)
+    self.assertIs(UOp.const(dtypes.float32, Invalid), invalid)
+    self.assertIs((moved:=invalid.reshape((1,))).cast(dtypes.float32), moved)
    scratch = Tensor.invalids(4, dtype=dtypes.float32)
    self.assertEqual((scratch.dtype, next(u.dtype for u in scratch.uop.toposort() if u.op is Ops.BUFFER), next(u.dtype for u in scratch.uop.toposort()
-      if u.is_invalid)), (dtypes.float32, dtypes.float32, dtypes.bool))
-    invalid, value = UOp.invalid(), UOp.const(1, dtypes.float32)
-    for u in (UOp.param(0, dtypes.bool, ()).where(value, invalid), value+invalid, UOp.stack(value, invalid)): self.assertIs(u.src[-1], invalid)
+      if u.arg is Invalid)), (dtypes.float32, dtypes.float32, dtypes.bool))
+    invalid, value = UOp.invalid(), UOp.const(dtypes.float32, 1)
    for u in (UOp(Ops.STACK, dtypes.float32, src=(value, invalid)), UOp(Ops.ADD, dtypes.float32, src=(value, invalid)),
-              UOp.const(True).where(value, invalid), UOp(Ops.CMPLT, src=(invalid, value)), UOp(Ops.CMPLT, src=(value, invalid)),
+              UOp.const(dtypes.bool, True).where(value, invalid), UOp(Ops.CMPLT, src=(invalid, value)), UOp(Ops.CMPLT, src=(value, invalid)),
              UOp.param(0, dtypes.float32, (4,)).index(invalid)): type_verify(u, spec_shared)
    gate, value = UOp.param(0, dtypes.bool, ()), UOp.param(1, dtypes.float, ())
-    self.assertIs((out:=graph_rewrite(gate.where(value, UOp.invalid()), pm_remove_invalid)).src[2], UOp.const(0, dtypes.float))
+    self.assertIs((out:=graph_rewrite(gate.where(value, UOp.invalid()), pm_remove_invalid)).src[2], UOp.const(dtypes.float, 0))
    type_verify(out.sink(), spec_program)

  def test_remove_invalid_stack_lanes(self):
-    stack = UOp(Ops.STACK, dtypes.half, (UOp.const(1, dtypes.half), UOp.invalid()))
+    stack = UOp(Ops.STACK, dtypes.half, (UOp.const(dtypes.half, 1), UOp.invalid()))
    out = graph_rewrite(stack, pm_remove_invalid)
-    self.assertEqual(out.src, (UOp.const(1, dtypes.half), UOp.const(0, dtypes.half)))
+    self.assertEqual(out.src, (UOp.const(dtypes.half, 1), UOp.const(dtypes.half, 0)))
    type_verify(out.sink(), spec_program)

 class TestLowerIndexDtype(unittest.TestCase):
@@ -80,7 +72,7 @@ class TestLowerIndexDtype(unittest.TestCase):
    # width the offset bounds select (this one needs long)
    buf = UOp.param(0, dtypes.float, (2**31+64,))
    i = UOp.variable("i", 0, 2**28)
-    shrink = UOp(Ops.SHRINK, src=(buf, (i*24).valid(i < 2**28), UOp.const(4)))
+    shrink = UOp(Ops.SHRINK, src=(buf, (i*24).valid(i < 2**28), UOp.const(dtypes.weakint, 4)))
    lowered = graph_rewrite(shrink.sink(), pm_lower_index_dtype)
    self.assertTrue(all(u.dtype != dtypes.weakint for u in lowered.backward_slice_with_self), "lowering must resolve all weakint")
    sh = next(u for u in lowered.backward_slice_with_self if u.op is Ops.SHRINK)
@@ -119,26 +111,6 @@ class TestSafeCast(unittest.TestCase):
    self.assertEqual(a.cast(dtypes.int8).cast(dtypes.int64).simplify(), a.cast(dtypes.int64))
    self.assertEqual(a.cast(dtypes.int8).cast(dtypes.float).simplify(), a.cast(dtypes.float))

-class TestConstFloatEq(unittest.TestCase):
-  def test_nan_eq_ne_agree(self):
-    nan = dtypes.float32.const(math.nan)
-    self.assertTrue(nan == math.nan)
-    self.assertFalse(nan != math.nan)  # float.__ne__ would say True here
-    self.assertFalse(nan == Invalid)
-    self.assertTrue(nan != Invalid)    # __ne__ must defer to the reflected eq, not swallow NotImplemented
-
-  def test_invalid_eq_defers_to_reflected(self):
-    class HoldsInvalid:  # a carrier that knows it holds Invalid. returning False for foreign types would silence its eq
-      def __eq__(self, other): return other is Invalid
-    self.assertTrue(Invalid == HoldsInvalid())
-    self.assertFalse(Invalid != HoldsInvalid())
-
-  def test_matchers_agree_on_nan(self):
-    n = UOp.const(math.nan, dtypes.float32)
-    for compiled in (False, True):
-      pm = PatternMatcher([(UPat(Ops.CONST, arg=math.nan), lambda: True)], compiled=compiled)
-      self.assertTrue(pm.rewrite(n), f"{compiled=}")
-
 class TestExecALU(unittest.TestCase):
  def test_sqrt(self):
    self.assertEqual(exec_alu(Ops.SQRT, dtypes.float, (0.0,)), 0.0)
@@ -226,9 +198,9 @@ class TestGatedStoreRewrite(unittest.TestCase):
  def test_tiny_gate_store(self):
    gmem = UOp.param(0, dtypes.float, (8,))
    gidx0 = UOp.special(4, 'gidx0')
-    gate = gidx0<UOp.const(1)
-    idx = UOp(Ops.INDEX, src=(gmem, (gidx0 * UOp.const(2)).valid(gate)))
-    val = UOp.const(42.0).cast(dtypes.float)
+    gate = gidx0<UOp.const(dtypes.weakint, 1)
+    idx = UOp(Ops.INDEX, src=(gmem, (gidx0 * UOp.const(dtypes.weakint, 2)).valid(gate)))
+    val = UOp.const(dtypes.float, 42.0)
    store = UOp(Ops.STORE, src=(idx, val))
    uops = to_uops_list([store])
    if_uop = next(u for u in uops if u.op is Ops.IF)
@@ -243,10 +215,10 @@ class TestGatedStoreRewrite(unittest.TestCase):
    gmem0 = UOp.param(0, dtypes.float, (8,))
    gmem1 = UOp.param(1, dtypes.float, (8,))
    gidx0 = UOp.special(4, 'gidx0')
-    idx = gidx0 * UOp.const(2)
-    idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gidx0<UOp.const(1))))
+    idx = gidx0 * UOp.const(dtypes.weakint, 2)
+    idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gidx0<UOp.const(dtypes.weakint, 1))))
    idx1 = UOp(Ops.INDEX, src=(gmem1, idx))
-    val = UOp.const(42.0).cast(dtypes.float)
+    val = UOp.const(dtypes.float, 42.0)
    stores = [UOp.store(idx0, val), UOp.store(idx1, val)]
    uops = to_uops_list(stores)
    if_uop = next(u for u in uops if u.op is Ops.IF)
@@ -263,11 +235,11 @@ class TestGatedStoreRewrite(unittest.TestCase):
    gmem0 = UOp.param(0, dtypes.float, (8,))
    gmem1 = UOp.param(1, dtypes.float, (8,))
    gidx0 = UOp.special(4, 'gidx0')
-    idx = gidx0*UOp.const(2)
-    gate = gidx0<UOp.const(1)
+    idx = gidx0*UOp.const(dtypes.weakint, 2)
+    gate = gidx0<UOp.const(dtypes.weakint, 1)
    idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gate)))
    idx1 = UOp(Ops.INDEX, src=(gmem1, idx.valid(gate)))
-    val = UOp.const(42.0).cast(dtypes.float)
+    val = UOp.const(dtypes.float, 42.0)
    stores = [UOp.store(idx0, val), UOp.store(idx1, val)]
    uops = to_uops_list(stores)
    ifs = [u for u in uops if u.op is Ops.IF]
@@ -285,7 +257,7 @@ class TestFastIdiv(unittest.TestCase):
  def test_division_power_of_two(self):
    for dt in (dtypes.int32, dtypes.uint32):
      g = UOp.param(0, dt, (3,))
-      c = UOp.const(2).cast(dt)
+      c = UOp.const(dt, 2)
      l = g.index(c)
      a = UOp(Ops.CDIV, dt, (l, c))
      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -298,7 +270,7 @@ class TestFastIdiv(unittest.TestCase):
    # FLOORMOD by a power of two lowers to AND (correct floor mod for any sign in two's complement)
    for dt in (dtypes.int32, dtypes.uint32):
      g = UOp.param(0, dt, (9,))
-      c = UOp.const(8).cast(dt)
+      c = UOp.const(dt, 8)
      a = UOp(Ops.FLOORMOD, dt, (g.index(c), c))
      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
      ops = [x.op for x in uops]
@@ -310,7 +282,7 @@ class TestFastIdiv(unittest.TestCase):
    # uint FLOORDIV by a power of two lowers to a shift, leaving no IDIV/FLOORDIV in the kernel
    for dt in (dtypes.uint32, dtypes.uint64):
      g = UOp.param(0, dt, (3,))
-      c = UOp.const(2).cast(dt)
+      c = UOp.const(dt, 2)
      a = UOp(Ops.FLOORDIV, dt, (g.index(c), c))
      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
      ops = [x.op for x in uops]
@@ -322,7 +294,7 @@ class TestFastIdiv(unittest.TestCase):
  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
  def test_fast_idiv_and_mod(self):
    g = UOp.param(0, dtypes.uint32, (4,))
-    c = UOp.const(3).cast(dtypes.uint)
+    c = UOp.const(dtypes.uint, 3)
    l = g.index(c)
    a = UOp(Ops.CDIV, src=(l, c))
    uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -342,7 +314,7 @@ class TestFastIdiv(unittest.TestCase):
  def test_fast_idiv_bounded_numerator_zero(self):
    x = UOp.variable("x", 0, 1, dtype=dtypes.int32)
    for val in range(2):
-      self.assertEqual(eval_uop(x.alu(Ops.CDIV, UOp.const(3).cast(x.dtype)), vals=(val,)), cdiv(val, 3))
+      self.assertEqual(eval_uop(x.alu(Ops.CDIV, x.const_like(3)), vals=(val,)), cdiv(val, 3))

  @Context(DISABLE_FAST_IDIV=0)
  def test_fast_idiv_remove_powers_of_two(self):
@@ -357,7 +329,7 @@ class TestFastIdiv(unittest.TestCase):
  def test_fast_idiv_overflow(self):
    # This will be possible with a slightly different method for fast_idiv
    g = UOp.param(0, dtypes.uint32, (8,))
-    c = UOp.const(7).cast(dtypes.uint)
+    c = UOp.const(dtypes.uint, 7)
    l = UOp(Ops.LOAD, src=(g.index(c),))
    a = UOp(Ops.CDIV, src=(l, c))
    uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -368,7 +340,7 @@ class TestFastIdiv(unittest.TestCase):

  def test_disable_fast_idiv(self):
    g = UOp.param(0, dtypes.uint32, (4,))
-    c = UOp.const(3).cast(dtypes.uint)
+    c = UOp.const(dtypes.uint, 3)
    l = g.index(c)
    a = UOp(Ops.CDIV, src=(l, c))
    with Context(DISABLE_FAST_IDIV=1):
@@ -380,8 +352,8 @@ class TestFastIdiv(unittest.TestCase):
 class TestUOpMethod(unittest.TestCase):
  @unittest.skip("uops lt no longer ordered")
  def test_compare_alu_same_src_different_arg(self):
-    a = UOp.const(2.0)
-    b = UOp.const(3.0)
+    a = UOp.const(dtypes.float, 2.0)
+    b = UOp.const(dtypes.float, 3.0)

    add = UOp(Ops.ADD, src=(a, b))
    mul = UOp(Ops.MUL, src=(a, b))
@@ -396,8 +368,8 @@ class TestUOpMethod(unittest.TestCase):
    self.assertEqual(list(var_vals)[0], a.expr)

  def test_const_factor(self):
-    gidx0 = UOp(Ops.SPECIAL, src=(UOp.const(8),), arg='gidx0')
-    self.assertEqual(UOp.const(17).const_factor(), 17)
+    gidx0 = UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 8),), arg='gidx0')
+    self.assertEqual(UOp.const(dtypes.int, 17).const_factor(), 17)
    self.assertEqual(gidx0.const_factor(), 1)
    self.assertEqual((gidx0*3).const_factor(), 3)
    self.assertEqual((gidx0*3+6).const_factor(), 3)
@@ -416,26 +388,26 @@ class TestUOpMethod(unittest.TestCase):

  def test_const_zero_neg_zero_different(self):
    # -0.0 and 0.0 must be different UOps (for IEEE754 correctness, e.g. 1/-0.0 = -inf)
-    pos_zero = UOp.const(0.0)
-    neg_zero = UOp.const(-0.0)
+    pos_zero = UOp.const(dtypes.float, 0.0)
+    neg_zero = UOp.const(dtypes.float, -0.0)
    self.assertIsNot(pos_zero, neg_zero)
    self.assertNotEqual(hash(pos_zero.arg), hash(neg_zero.arg))

  def test_const_nan_same(self):
    # nan constants should be deduplicated
-    nan1 = UOp.const(float('nan'))
-    nan2 = UOp.const(float('nan'))
+    nan1 = UOp.const(dtypes.float, float('nan'))
+    nan2 = UOp.const(dtypes.float, float('nan'))
    self.assertIs(nan1, nan2)

 class TestUOpStr(unittest.TestCase):
  def test_uop_str(self):
-    a = UOp.const(2.0) + UOp.const(3.0)
+    a = UOp.const(dtypes.float, 2.0) + UOp.const(dtypes.float, 3.0)
    for _ in range(20): a = a + a
    assert len(str(a)) < 10_000, "exponential string growth"
    assert str(eval(str(a))) == str(a)

  def test_vectorized_str(self):
-    vec = UOp(Ops.STACK, src=tuple(UOp.const(x) for x in range(4)))
+    vec = UOp(Ops.STACK, src=tuple(UOp.const(dtypes.int, x) for x in range(4)))
    assert str(eval(str(vec))) == str(vec)

  def test_reduceop_arg(self):
@@ -454,7 +426,7 @@ class TestUPatHelpers(unittest.TestCase):

 class TestUopsObject(unittest.TestCase):
  def test_timing(self):
-    with Timing("create 10k uops:"): ret = [UOp.const(10000000+i, dtypes.int) for i in range(10000)]
+    with Timing("create 10k uops:"): ret = [UOp(Ops.CONST, dtypes.int, arg=10000000+i) for i in range(10000)]
    assert len(ret) == 10000

  def test_nested(self):
@@ -470,16 +442,16 @@ class TestUOpRender(unittest.TestCase):
    u = UOp(Ops.STACK, dtype=dtypes.void, src=())
    self.assertEqual(u.render(), "{}")
  def test_render_vectorize_same(self):
-    u = UOp(Ops.STACK, src=(UOp.const(0),)*3)
+    u = UOp(Ops.STACK, dtype=dtypes.int, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0)))
    self.assertEqual(u.render(simplify=False), "{0,0,0}")
  def test_render_vectorize_different(self):
-    u = UOp(Ops.STACK, src=tuple(UOp.const(i) for i in range(3)))
+    u = UOp(Ops.STACK, dtype=dtypes.int, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)))
    self.assertEqual(u.render(simplify=False), "{0,1,2}")
  def test_render_vectorize_same_simplified(self):
-    u = UOp(Ops.STACK, src=(UOp.const(0),)*3)
+    u = UOp(Ops.STACK, dtype=dtypes.int, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0)))
    self.assertEqual(u.render(), "{0,0,0}")
  def test_render_vectorize_different_simplified(self):
-    u = UOp(Ops.STACK, src=tuple(UOp.const(i) for i in range(3)))
+    u = UOp(Ops.STACK, dtype=dtypes.int, src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2)))
    self.assertEqual(u.render(), "{0,1,2}")

 class TestContiguousViewOffset(unittest.TestCase):
@@ -147,21 +147,21 @@ class TestUOpsStats(unittest.TestCase):
  #MULACC should have the same stats as MUL + ADD
  def test_mulacc(self):
    globl = UOp.param(0, dtypes.int, (3,))
-    o1 = UOp.const(1, dtypes.int)
-    o2 = UOp.const(2, dtypes.int)
+    o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1)
+    o2 = UOp(Ops.CONST, dtypes.int, tuple(), 2)
    u1 = globl.index(o1)
    u2 = globl.index(o2)
-    u3 = UOp.const(3, dtypes.int)
+    u3 = UOp(Ops.CONST, dtypes.int, tuple(), 3)
    u4 = UOp(Ops.MUL, src=(u1,u2))
    u5 = UOp(Ops.ADD, src=(u4,u3))
    uops = tuple(u5.toposort())

    globl = UOp.param(0, dtypes.int, (3,))
-    o1 = UOp.const(1, dtypes.int)
-    o2 = UOp.const(2, dtypes.int)
+    o1 = UOp(Ops.CONST, dtypes.int, tuple(), 1)
+    o2 = UOp(Ops.CONST, dtypes.int, tuple(), 2)
    u1 = globl.index(o1)
    u2 = globl.index(o2)
-    u3 = UOp.const(3, dtypes.int)
+    u3 = UOp(Ops.CONST, dtypes.int, tuple(), 3)
    u4 = UOp(Ops.MULACC, src=(u1,u2,u3))
    uops_fma = tuple(u4.toposort())

@@ -1,11 +1,11 @@
 import unittest
 from tinygrad.helpers import DEBUG, Context
 from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import UPat, rewrite_group, GroupOp, Ops
+from tinygrad.uop.ops import UPat, track_rewrites, GroupOp, Ops
 from tinygrad.uop.upat import _get_code, upat_compile
 import dis

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

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

@@ -14,12 +12,12 @@ class TestValidateOOB(unittest.TestCase):
  def test_const_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (16,))
-      to_uops_list([buf.index(UOp.const(0)).load(dtype=dtypes.int)])  # valid
-      to_uops_list([buf.index(UOp.const(15)).load(dtype=dtypes.int)])  # valid (last element)
+      to_uops_list([buf.index(UOp.const(dtypes.int, 0)).load(dtype=dtypes.int)])  # valid
+      to_uops_list([buf.index(UOp.const(dtypes.int, 15)).load(dtype=dtypes.int)])  # valid (last element)
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(UOp.const(16)).load(dtype=dtypes.int)])  # off by one
+        to_uops_list([buf.index(UOp.const(dtypes.int, 16)).load(dtype=dtypes.int)])  # off by one
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(UOp.const(42)).load(dtype=dtypes.int)])  # way out
+        to_uops_list([buf.index(UOp.const(dtypes.int, 42)).load(dtype=dtypes.int)])  # way out

  def test_variable_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
@@ -140,19 +138,11 @@ class TestValidateOOB(unittest.TestCase):
      with self.assertRaises(RuntimeError):
        to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 64))).load(dtype=dtypes.int)])  # oob

-  def test_load_from_shrink_as_index(self):
-    with Context(CHECK_OOB=1, SPEC=2):
-      buf0 = UOp.param(0, dtypes.int, (16,))
-      buf1 = UOp.param(1, dtypes.int, (64,))
-      shrink = UOp(Ops.SHRINK, src=(buf0, UOp.const(0, dtypes.int), UOp.const(4)))
-      ld0 = shrink.load(dtype=dtypes.int).index(0)
-      to_uops_list([buf1.index(ld0.valid((ld0 >= 0) & (ld0 < 64))).load(dtype=dtypes.int)])
-
  def test_load_bool_as_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf_bool = UOp.param(0, dtypes.bool, (16,))
      buf_int = UOp.param(1, dtypes.int, (8,))
-      gidx = UOp(Ops.SPECIAL, src=(UOp.const(16),), arg="gidx0")
+      gidx = UOp(Ops.SPECIAL, src=(UOp.const(dtypes.weakint, 16),), arg="gidx0")
      ld_bool = buf_bool.index(gidx).load()
      with self.assertRaises(RuntimeError):
        to_uops_list([buf_int.index(gidx.valid(ld_bool)).load()])  # gidx 0..15, buf_int size 8
@@ -166,12 +156,12 @@ class TestValidateOOB(unittest.TestCase):
      sbuf = UOp.placeholder((8,), dtypes.uint, slot=0, addrspace=AddrSpace.LOCAL)

      # Define indices, valids and barrier
-      gidx = UOp(Ops.SPECIAL, src=(UOp.const(416),), arg="gidx0")
-      lidx = UOp(Ops.SPECIAL, src=(UOp.const(10),), arg="lidx0")
+      gidx = UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 416),), arg="gidx0")
+      lidx = UOp(Ops.SPECIAL, src=(UOp.const(dtypes.int, 10),), arg="lidx0")

      gate = (gidx<400) & (lidx<8)

-      local_store = sbuf.index(lidx.valid(lidx<8)).store(UOp.const(1))
+      local_store = sbuf.index(lidx.valid(lidx<8)).store(UOp.const(dtypes.uint, 1))

      barrier = UOp(Ops.BARRIER, src=(local_store,))
      if_barrier = UOp(Ops.IF, src=(gate, barrier))
@@ -189,7 +179,7 @@ class TestValidateOOB(unittest.TestCase):
      glbl0 = UOp.param(0, dtypes.int, (16,))
      mask = UOp.param(0, dtypes.bool, (16,))
      ridx = UOp.range(20, 0)
-      ld0 = UOp(Ops.LOAD, src=(glbl0.index(UOp.const(ridx<16&mask, ridx))))
+      ld0 = UOp(Ops.LOAD, src=(glbl0.index(UOp.const(ridx, ridx<16&mask))))
      to_uops_list([ld0])

 if __name__ == "__main__":
@@ -3,7 +3,7 @@ from pathlib import Path
 from dataclasses import dataclass
 from typing import Generator

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

-@rewrite_group(name=True)
+@track_rewrites(name=True)
 def exec_rewrite(sink:UOp, pm_lst:list[PatternMatcher], names:None|list[str]=None) -> UOp:
  for i,pm in enumerate(pm_lst):
    sink = graph_rewrite(sink, TrackedPatternMatcher(pm.patterns), name=names[i] if names else None)
@@ -96,10 +96,10 @@ class TestViz(unittest.TestCase):
  def test_exceptions(self):
    # VIZ tracks rewrites up to and including the error
    def count_3(x:UOp):
-      assert x.val <= 3
-      return UOp.const(x.val+1, x.dtype)
+      assert x.arg <= 3
+      return x.replace(arg=x.arg+1)
    err_pm = PatternMatcher([(UPat.cvar("x"), count_3),])
-    a = UOp.const(1)
+    a = UOp.const(dtypes.int, 1)
    with save_viz() as viz:
      with self.assertRaises(AssertionError): exec_rewrite(a, [err_pm])
    lst = viz.list_items()
@@ -109,7 +109,7 @@ class TestViz(unittest.TestCase):
  def test_default_name(self):
    with save_viz() as viz:
      a = UOp.variable("a", 1, 10)
-      @rewrite_group()
+      @track_rewrites()
      def name_default(): return graph_rewrite(a, PatternMatcher([]))
      name_default()
    lst = viz.list_items()
@@ -118,7 +118,7 @@ class TestViz(unittest.TestCase):
  # name can also come from a function that returns a string
  def test_dyn_name_fxn(self):
    with save_viz() as viz:
-      @rewrite_group(name=lambda *args,ret,**kwargs: ret.render())
+      @track_rewrites(name=lambda *args,ret,**kwargs: ret.render())
      def name_from_fxn(s:UOp, arg:list|None=None): return graph_rewrite(s, PatternMatcher([]))
      name_from_fxn(UOp.variable("a", 1, 10)+1, arg=["test"])
    lst = viz.list_items()
@@ -128,18 +128,18 @@ class TestViz(unittest.TestCase):
  # name can also come from a function that returns a TracingKey
  def test_tracing_key(self):
    with save_viz() as viz:
-      @rewrite_group(name=lambda inp,ret: TracingKey("custom_name", (inp,)))
+      @track_rewrites(name=lambda inp,ret: TracingKey("custom_name", (inp,)))
      def test(s:UOp): return graph_rewrite(s, PatternMatcher([]))
      test(UOp.variable("a", 1, 10)+1)
    lst = viz.list_items()
    # NOTE: names from TracingKey do not get deduped
    self.assertEqual(lst[0]["name"], "custom_name")

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

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

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

-  def test_rewrite_group_invalid_arg(self):
+  def test_profile_matches_invalid_arg(self):
    with save_viz():
-      @rewrite_group(new_ctx=False)
+      @profile_matches
      def invalid_fxn(arg:str): return graph_rewrite(UOp(Ops.SINK), PatternMatcher([]))
      with self.assertRaisesRegex(AssertionError, "invalid match tracing input"):
        invalid_fxn("test")
@@ -199,11 +199,11 @@ class TestViz(unittest.TestCase):
    self.assertEqual(ansistrip(a2["label"]), "CUSTOM\nx\nyzww\nw")

  def test_inf_loop(self):
-    a = UOp.const(3)
-    b = UOp.const(4)
+    a = UOp.const(dtypes.int, 3)
+    b = UOp.const(dtypes.int, 4)
    pm = PatternMatcher([
-      (UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype)),
-      (UPat(Ops.CONST, arg=4, name="x"), lambda x: UOp.const(3, x.dtype)),
+      (UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4)),
+      (UPat(Ops.CONST, arg=4, name="x"), lambda x: x.replace(arg=3)),
    ])
    with save_viz() as viz:
      # use smaller stack limit for faster test (default is 250000)
@@ -224,9 +224,9 @@ class TestViz(unittest.TestCase):
    list(viz.get_details(0, 0))

  def test_enter_calls_rewrite(self):
-    pm = PatternMatcher([(UPat(Ops.CONST, arg=3, name="x"), lambda x: UOp.const(4, x.dtype))])
+    pm = PatternMatcher([(UPat(Ops.CONST, arg=3, name="x"), lambda x: x.replace(arg=4))])
    with save_viz() as viz:
-      inner = UOp.const(3)
+      inner = UOp.const(dtypes.int, 3)
      call = UOp(Ops.CALL, src=(UOp(Ops.SINK, src=(inner,)),))
      func = UOp(Ops.FUNCTION, src=(UOp(Ops.TUPLE, src=(call,)),))
      graph_rewrite(func, TrackedPatternMatcher(pm.patterns), enter_calls=True)
@@ -236,8 +236,8 @@ class TestViz(unittest.TestCase):
  def test_const_node_visibility(self):
    with save_viz() as viz:
      a = UOp.variable("a", 0, 10, dtype=dtypes.int)
-      z = UOp.const(0, a.dtype)
-      y = UOp.const(math.pi, dtypes.float)
+      z = UOp.const(a.dtype, 0)
+      y = UOp.const(dtypes.float, math.pi)
      alu = a*z
      ret = exec_rewrite(sink:=UOp.sink(alu, y), [sym])
    lst = viz.list_items()
@@ -253,7 +253,7 @@ class TestViz(unittest.TestCase):

  def test_const_reshape_expand_folded(self):
    # CONST->EXPAND should be folded into the ALU node, not shown as separate EXPAND nodes
-    c = UOp.const(1.0).expand((3,4))  # creates CONST->EXPAND chain
+    c = UOp.const(dtypes.float, 1.0, shape=(3,4))  # creates CONST->EXPAND chain
    a = UOp.variable("a", 0.0, 10.0, dtypes.float)
    alu = a + c
    with save_viz() as viz:
@@ -267,13 +267,13 @@ class TestViz(unittest.TestCase):

  def test_stack_movement_not_folded_unless_all_const(self):
    a = UOp.variable("a", 0, 10, dtype=dtypes.int)
-    c = UOp.const(1)
+    c = UOp.const(dtypes.int, 1)
    stack = a.stack(c)
    reshaped = stack.reshape((1, 2))
    graph = uop_to_json(VizData(), reshaped)
    self.assertFalse(graph[id(stack)]["exclude"])

-    const_stack = c.stack(UOp.const(2))
+    const_stack = c.stack(UOp.const(dtypes.int, 2))
    const_reshaped = const_stack.reshape((1, 2))
    const_graph = uop_to_json(VizData(), const_reshaped)
    self.assertTrue(const_graph[id(const_stack)]["exclude"])
@@ -305,10 +305,10 @@ class TestVizTree(unittest.TestCase):

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

-  # tracing also works without a rewrite_group context
+  # tracing also works without a track_rewrites context
  # all graph_rewrites get put into the default group
  def test_default_tracing(self):
    with save_viz() as viz:
      def test(root):
        return graph_rewrite(root, sym)
-      test(c:=UOp.const(1))
+      test(c:=UOp.const(dtypes.int, 1))
      test(c+1)
    ls = viz.list_items()
    self.assertEqual(len(ls), 1)
    self.assertEqual(ls[0]["name"], "default graph_rewrite")

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

-  # @rewrite_group always starts a new group.
+  # @track_rewrites always starts a new group.
  def test_group_combined(self):
    with save_viz() as viz:
      def default_test(root): return graph_rewrite(root, sym)
-      tracked_test = rewrite_group()(default_test)
-      c = UOp.const(1)
+      tracked_test = track_rewrites()(default_test)
+      c = UOp.const(dtypes.int, 1)
      default_test(c+1) # goes to the default group
      tracked_test(c)   # all rewrites after this go inside the second group.
      default_test(c+2)
    ls = viz.list_items()
    self.assertEqual(len(ls), 2)
    graph = next(viz.get_details(0, 0))["graph"]
-    # both operands of c+1 are the same bare weak CONST, so the graph has two nodes
    self.assertEqual(list(graph), [id(c), id(c+1)])
    self.assertTrue(graph[id(c)]["exclude"])
    self.assertFalse(graph[id(c+1)]["exclude"])
    self.assertEqual(list(next(viz.get_details(1, 0))["graph"]), [id(c)])
    graph = next(viz.get_details(1, 1))["graph"]
-    self.assertEqual(list(graph), [id(c), id((c+2).src[1]), id(c+2)])
+    self.assertEqual(list(graph), [id(c), id(c.const_like(2)), id(c+2)])
    self.assertTrue(graph[id(c)]["exclude"])
-    self.assertTrue(graph[id((c+2).src[1])]["exclude"])
+    self.assertTrue(graph[id(c.const_like(2))]["exclude"])
    self.assertFalse(graph[id(c+2)]["exclude"])

  def test_recurse(self):
@@ -1,7 +1,6 @@
 import unittest, sys
 from tinygrad import Tensor, GlobalCounters, dtypes, Context
 from tinygrad.helpers import WINO
-from test.helpers import check_schedule

@unittest.skipIf(sys.platform.startswith("win"), "flaky on Windows")
 class TestWinograd(unittest.TestCase):
@@ -14,7 +13,7 @@ class TestWinograd(unittest.TestCase):
  def test_forward_kernels(self):
    x,w = Tensor.rand(1,4,9,9).realize(), Tensor.rand(4,4,3,3).realize()
    out = Tensor.conv2d(x,w)
-    check_schedule(out, 4)
+    self.assertEqual(len(out.schedule_linear().src), 4)

  def test_backward_counters(self):
    # contiguous_backward on the pooled input keeps the input-transform adjoint out of the overlap accumulation, so
@@ -239,15 +239,6 @@ class TestKernelOpts(unittest.TestCase):
    helper_linearizer_opt(a.sum().exp(), [[Opt(OptOps.PADTO, 0, 32)],])
    helper_linearizer_opt(a.sum(0).exp(), [[Opt(OptOps.PADTO, 1, 32)],])

-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
-  @unittest.expectedFailure
-  def test_padto_group_full_unroll_sum(self):
-    a = Tensor.ones(2, 28, 4096, dtype=dtypes.bfloat16).realize()
-    out = ((a * 0.5).float().square()).sum(axis=(0, 2))
-    opts_to_apply = [Opt(OptOps.GROUPTOP, 1, 256), Opt(OptOps.PADTO, 3, 32), Opt(OptOps.UNROLL, 2, 0), Opt(OptOps.UPCAST, 0, 7)]
-    helper_linearizer_opt(out, [opts_to_apply], check_default_opt=False)
-
  def test_padto_sum(self):
    N = 18
    # NOTE: this setup prevents 17 * 17 contiguous merged into one dimension
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	9540700bbc	direct lds	2026-07-31 09:03:32 -07:00
geohot	1cfaa385d6	work	2026-07-30 21:46:06 -07:00
geohot	36bf7cf65e	bench	2026-07-30 20:55:22 -07:00
geohot	d9ec3d7282	bugfixes	2026-07-30 20:39:25 -07:00
geohot	679faeacc7	hipkittens style gemm (kimi)	2026-07-30 18:06:34 -07:00