cleaner

chat class
add olmo3 support to llm
2026-08-16 03:18:27 +00:00 · 2026-04-22 15:55:19 +08:00 · 2026-04-22 10:46:10 +08:00 · 2026-04-21 16:36:08 +08:00
312 changed files with 58951 additions and 73842 deletions
@@ -33,8 +33,12 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: 'autogen'
+        opencl: 'true'
        amd: 'true'
+        cuda: 'true'
        llvm: 'true'
+        webgpu: 'true'
+        mesa: 'true'
        pydeps: 'pyyaml mako'
    - name: Install autogen support packages
      run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev
@@ -44,7 +48,7 @@ jobs:
        python3 -c "from tinygrad.runtime.autogen import opencl"
        python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv"
        python3 -c "from tinygrad.runtime.autogen import comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
-        python3 -c "from tinygrad.runtime.autogen.am import *"
+        python3 -c "from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v13_0_6, smu_v13_0_12, smu_v14_0_2"
        python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
        python3 -c "from tinygrad.runtime.autogen import llvm"
        python3 -c "from tinygrad.runtime.autogen import webgpu"
@@ -51,35 +51,34 @@ jobs:
    - name: openpilot compile3 0.10.1 driving_vision
      run: FLOAT16=1 DEV=CL IMAGE=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx

-  # TODO: reenable when not flaky
-  #testframeworkpytest:
-  #  name: framework pytest
-  #  env:
-  #    CI: ""
-  #    CAPTURE_PROCESS_REPLAY: "0"
-  #  runs-on: [self-hosted, framework]
-  #  timeout-minutes: 10
-  #  defaults:
-  #    run:
-  #      shell: bash -e -o pipefail {0}
-  #  if: github.repository_owner == 'tinygrad'
-  #  steps:
-  #  - name: Checkout Code
-  #    uses: actions/checkout@v6
-  #  - name: setup python environment
-  #    run: |
-  #      rm -rf /tmp/tinygrad_pytest_ci
-  #      uv venv /tmp/tinygrad_pytest_ci
-  #      source /tmp/tinygrad_pytest_ci/bin/activate
-  #      uv pip install .[testing]
-  #  - name: setup staging db
-  #    run: |
-  #      echo "CACHEDB=/tmp/pytest-db-ci.db" >> $GITHUB_ENV
-  #      rm -f /tmp/pytest-db-ci*
-  #  - name: Run pytest -nauto
-  #    run: |
-  #      source /tmp/tinygrad_pytest_ci/bin/activate
-  #      pytest -nauto --durations=20
+  testframeworkpytest:
+    name: framework pytest
+    env:
+      CI: ""
+      CAPTURE_PROCESS_REPLAY: "0"
+    runs-on: [self-hosted, framework]
+    timeout-minutes: 10
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: setup python environment
+      run: |
+        rm -rf /tmp/tinygrad_pytest_ci
+        uv venv /tmp/tinygrad_pytest_ci
+        source /tmp/tinygrad_pytest_ci/bin/activate
+        uv pip install .[testing]
+    - name: setup staging db
+      run: |
+        echo "CACHEDB=/tmp/pytest-db-ci.db" >> $GITHUB_ENV
+        rm -f /tmp/pytest-db-ci*
+    - name: Run pytest -nauto
+      run: |
+        source /tmp/tinygrad_pytest_ci/bin/activate
+        pytest -nauto --durations=20

  testmacbenchmark:
    name: Mac Benchmark
@@ -190,7 +189,7 @@ jobs:
        path: |
          onnx_inference_speed.csv
    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
+      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3.11 process_replay.py

  testusbgpu:
    name: UsbGPU Benchmark
@@ -325,7 +324,7 @@ jobs:
        path: |
          onnx_inference_speed.csv
    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
+      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

  testmorenvidiabenchmark:
    name: tinybox green Training Benchmark
@@ -387,7 +386,7 @@ jobs:
      # TODO: remove BERT_LAYERS once scheduler is fast
      run: BENCHMARK_LOG=bert_10steps_6gpu DEV=NV CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
+      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

  testamdbenchmark:
    name: tinybox red Benchmark
@@ -499,7 +498,7 @@ jobs:
    - name: Run GPT2 w HALF/BEAM
      run: BENCHMARK_LOG=gpt2_half_beam DEV=AMD HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing
    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
+      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

  testmoreamdbenchmark:
    name: tinybox red Training Benchmark
@@ -556,7 +555,7 @@ jobs:
    #- name: Test full tinyfs load
    #  run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
+      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

  testmlperfamdbenchmark:
    name: tinybox red MLPerf Benchmark
@@ -602,7 +601,7 @@ jobs:
      # TODO: remove BERT_LAYERS once scheduler is fast
      run: BENCHMARK_LOG=bert_10steps_6gpu DEV=AMD CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
+      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

  testqualcommbenchmark:
    name: comma Benchmark
@@ -624,12 +623,10 @@ jobs:
      run: test/external/process_replay/reset.py
    - name: openpilot compile3 0.11.0 driving_vision
      run: BENCHMARK_LOG=openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot compile3 0.11.0 driving_vision (from pickle)
-      run: BENCHMARK_LOG=openpilot_0_11_0_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM taskset -c 4-7 python3 examples/openpilot/compile3.py
    - name: IR3 openpilot compile3 0.11.0 driving_vision
      run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.11.0 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
    - name: openpilot compile3 0.11.0 dmonitoring
      run: BENCHMARK_LOG=openpilot_0_11_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: DEBUG=2 openpilot compile3 0.10.1 driving_vision
@@ -637,7 +634,7 @@ jobs:
    - name: openpilot compile3 0.10.1 driving_vision
      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.10.1 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
    - name: openpilot compile3 0.10.1 dmonitoring
      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: benchmark MobileNetV2 on DSP
@@ -649,7 +646,7 @@ jobs:
        # benchmark on DSP with NOOPT=1, the devectorizer has issues
        PYTHONPATH=. CC=clang-19 DEV=DSP NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
+      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

  testcommausbgpubenchmark:
    name: UsbGPU Benchmark (comma)
@@ -671,8 +668,6 @@ jobs:
      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

  testreddriverbenchmark:
    name: AM Benchmark
@@ -746,7 +741,7 @@ jobs:
        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 DEV=PCI+AMD AMD_AQL=1 python3 test/test_tiny.py
        pkill -f 'extra/remote/serve.py' || true
    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
+      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

  testgreendriverbenchmark:
    name: NV Benchmark
@@ -809,4 +804,4 @@ jobs:
        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6483 DEV=NV python3 test/test_tiny.py
        pkill -f 'extra/remote/serve.py' || true
    - name: Run process replay tests
-      uses: ./.github/actions/process-replay
+      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
@@ -333,7 +333,7 @@ jobs:
        deps: testing_unit
        python-version: '3.14'
    - name: Test SPEC=2
-      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 test/unit test/backend test/opt --ignore test/backend/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" -k "not test_conv2d_ceildiv_edge_case" --splits 2 --group ${{ matrix.group }}
+      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 test/unit test/backend test/opt --ignore test/backend/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}

  fuzzing:
    name: Fuzzing
@@ -417,7 +417,7 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1486 ALLOWED_GATED_READ_IMAGE=18 FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1486 ALLOWED_GATED_READ_IMAGE=17 FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
      - name: Test openpilot CL compile fp16
        run: FLOAT16=1 DEV=CL IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
      - name: Test openpilot CL compile fp32 (test correctness)
@@ -676,8 +676,8 @@ jobs:
        env:
          AMD: 0
        run: |
-          PYTHONPATH=. DEV=NULL:HIP:gfx1100 python extra/mmapeak/mmapeak.py
-          PYTHONPATH=. DEV=NULL:HIP:gfx950 python3 -m pytest -n=auto test/testextra/test_tk.py test/backend/test_asm_gemm.py
+          PYTHONPATH=. DEV=NULL::gfx1100 python extra/mmapeak/mmapeak.py
+          PYTHONPATH=. DEV=NULL::gfx1201 python3 -m pytest -n=auto test/testextra/test_tk.py test/backend/test_asm_gemm.py
      - name: Run matmul on MOCKKFD
        run: |
          PYTHONPATH="." DEV=MOCKKFD+AMD N=256 python3 extra/gemm/amd_asm_matmul.py
@@ -773,7 +773,7 @@ jobs:
          cuda: 'true'
          ocelot: 'true'
      - name: Set env
-        run: printf "${{ matrix.backend == 'ptx' && 'DEV=MOCK+CUDA:PTX' || matrix.backend == 'nv' && 'DEV=MOCK+NV\nSKIP_SLOW_TEST=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'ptx' && 'DEV=MOCK+CUDA:PTX' || matrix.backend == 'nv' && 'DEV=MOCKNVK+NV\nSKIP_SLOW_TEST=1' }}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CUDA','NV'], Device.DEFAULT"
@@ -870,7 +870,7 @@ jobs:
        python -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py test/device/test_amd_llvm.py --durations=20
    - name: Run pytest (ptx)
      env:
-        DEV: "MOCK+NV:PTX"
+        DEV: "MOCKNVK+NV:PTX"
        FORWARD_ONLY: 1
        # TODO: failing due to library loading error
        CAPTURE_PROCESS_REPLAY: 0
@@ -68,4 +68,3 @@ mutants
 .mutmut-cache
 dagre/
 graphlib/
-uv.lock
@@ -1,4 +1,6 @@
 # abstractions2 goes from back to front, here we will go from front to back
+from typing import List
+from tinygrad.helpers import tqdm

 # *****
 # 0. Load mnist on the device
@@ -31,21 +33,21 @@ model(X).sparse_categorical_crossentropy(Y).backward()
 optim.schedule_step()   # this will step the optimizer without running realize

 # *****
-# 3. Create a schedule (linear uop).
+# 3. Create a schedule.

 # The weight Tensors have been assigned to, but not yet realized. Everything is still lazy at this point
 # l1.uop and l2.uop define a computation graph

-from tinygrad.engine.realize import run_linear
-linear = Tensor.schedule_linear(l1, l2)
+from tinygrad.schedule import ExecItem
+schedule: List[ExecItem] = Tensor.schedule(l1, l2)

-print(f"The schedule contains {len(linear.src)} items.")
-for call in linear.src: print(str(call)[:80])
+print(f"The schedule contains {len(schedule)} items.")
+for si in schedule: print(str(si)[:80])

 # *****
-# 4. Lower and run the schedule (linear uop).
+# 4. Lower and run the schedule.

-run_linear(linear)
+for si in tqdm(schedule): si.run()

 # *****
 # 5. Print the weight change
@@ -105,7 +105,7 @@ def example_3_custom_uop(a:Tensor, correct):
 def example_5_custom_assembly(a:Tensor, correct):
  # Kernel class copied from amd_asm_matmul
  class Kernel:
-    def __init__(self): self.instructions, self.labels, self.pos = [], {}, 0
+    def __init__(self, arch='gfx1100'): self.instructions, self.labels, self.pos, self.arch = [], {}, 0, arch
    def label(self, name): self.labels[name] = self.pos
    def emit(self, inst, target=None):
      self.instructions.append(inst)
@@ -17,13 +17,15 @@ The `UOp` graph specifies the compute in terms of low level tinygrad ops. Not al

 ## Scheduling

-The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/schedule/__init__.py) converts the graph of UOps into a `LINEAR` UOp whose `src` is a list of `CALL` UOps. One `CALL` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. The `CALL`'s `src[0]` (a `SINK` ast) specifies what compute to run, and the remaining `src` are the buffers to run it on.
+The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/schedule/__init__.py) converts the graph of UOps into a list of `ExecItem`. One `ExecItem` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. `ast` specifies what compute to run, and `bufs` specifies what buffers to run it on.
+
+::: tinygrad.schedule.ExecItem

 ## Lowering

-The code in [realize](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/realize.py) lowers each `CALL` by compiling its ast into a `PROGRAM` and running it.
+The code in [realize](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/realize.py) lowers `ExecItem` by populating its `prg` field with

-::: tinygrad.engine.realize.run_linear
+::: tinygrad.engine.realize.run_schedule

 There's a ton of complexity hidden behind this, see the `codegen/` directory.

@@ -33,7 +35,13 @@ Then we render the UOps into code with a `Renderer`, then we compile the code to

 ## Execution

-`run_linear` walks the `LINEAR` UOp, dispatching each `CALL` to a runner (kernel, copy, view, encdec, or graph).
+Creating `ExecItem`, which has a run method
+
+::: tinygrad.engine.realize.ExecItem
+    options:
+        members: true
+
+Lists of `ExecItem` can be condensed into a single ExecItem with the Graph API (rename to Queue?)

 ## Runtime

@@ -28,7 +28,7 @@ Transforms the ast into an optimized ast. This is where BEAM search and heuristi

 Transform the optimized ast into a linearized and rendered program.

-::: tinygrad.codegen.to_program
+::: tinygrad.codegen.get_program
    options:
        members: false
        show_labels: false
@@ -53,7 +53,7 @@ Transform the linearized list of UOps into a program, represented as a string.

 Abstracted high level interface to the runtimes.

-::: tinygrad.engine.realize.to_program
+::: tinygrad.engine.realize.get_program
    options:
        members: false
        show_labels: false
@@ -57,8 +57,6 @@ AMD:LLVM      | use the AMD device with the LLVM renderer
 NV:CUDA:sm_70 | use the NV device with the CUDA renderer targetting sm_70
 AMD::gfx950   | use the AMD device targetting gfx950
 USB+AMD       | use the AMD device over the USB interface
-CPU:LLVM      | use the CPU device with the LLVM renderer
-CPU:LLVM:x86_64,znver2,avx2,-avx512f | use the CPU device with the LLVM renderer, with [additional arch flags](runtime.md#cpu-arch)

 ### Debug breakdown

@@ -5,12 +5,12 @@ tinygrad supports various runtimes, enabling your code to scale across a wide ra
 | Runtime | Description | Compiler Options | Requirements |
 |---------|-------------|------------------|--------------|
 | [NV](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_nv.py) | Provides acceleration for NVIDIA GPUs | nvrtc (default)<br>PTX (`DEV=NV:PTX`) | Ampere/Ada/Blackwell series GPUs.<br>You can select an interface via [the `DEV` variable](env_vars.md#dev-variable). See [NV interfaces](#nv-interfaces) for details. |
-| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | LLVM (`DEV=AMD:LLVM`)<br>HIP/COMGR (`DEV=AMD:HIP`) | CDNA3, CDNA4, RDNA3 or RDNA4 GPUs.<br>You can select an interface via [the `DEV` variable](env_vars.md#dev-variable). See [AMD interfaces](#amd-interfaces) for details. |
+| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | LLVM (`DEV=AMD:LLVM`)<br>HIP/COMGR (`DEV=AMD:HIP`) | RDNA2 or newer GPUs.<br>You can select an interface via [the `DEV` variable](env_vars.md#dev-variable). See [AMD interfaces](#amd-interfaces) for details. |
 | [QCOM](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_qcom.py) | Provides acceleration for QCOM GPUs | - | 6xx series GPUs |
 | [METAL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_metal.py) | Utilizes Metal for acceleration on Apple devices | - | M1+ Macs; Metal 3.0+ for `bfloat` support |
 | [CUDA](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cuda.py) | Utilizes CUDA for acceleration on NVIDIA GPUs | nvrtc (default)<br> PTX (`DEV=CUDA:PTX`) | NVIDIA GPU with CUDA support |
 | [CL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cl.py) | Accelerates computations using OpenCL on GPUs | - | OpenCL 2.0 compatible device |
-| [CPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang or llvm compiler | Clang JIT (default)<br>LLVM IR (`DEV=CPU:LLVM`) | `clang` compiler in system `PATH`<br>You can specify additional arch parameters via [the `DEV` variable](env_vars.md#dev-variable). See [CPU arch](#cpu-arch) for details. |
+| [CPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang or llvm compiler | Clang JIT (default)<br>LLVM IR (`DEV=CPU:LLVM`) | `clang` compiler in system `PATH` |
 | [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | - | Dawn library installed and discoverable. Binaries: [pydawn v0.3.0](https://github.com/wpmed92/pydawn/releases/tag/v0.3.0) |


@@ -79,13 +79,3 @@ NV backend supports several interfaces for communicating with devices:

 * `NVK`: uses the nvidia driver
 * `PCI`: uses the [NV driver](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/support/nv/nvdev.py)
-
-## CPU Arch
-The CPU renderers may be additionally configured using the arch component of [the `DEV` environment variable](env_vars.md#dev-variable).
-CPU arch should be specified as a comma-separated list of parameters, and must contain at least two values: the architecture family (ie. x86_64, arm64, or riscv64) and the cpu type (as accepted by `clang`'s `-march`).
-If native is specified as the cpu type, tinygrad (or delegate compiler) will query the host cpu type. Additional comma-separated values may be specified as follows:
-
-* `AMX`: emit Apple silicon AMX instructions
-
-All other additional values are interpreted as cpu feature flags. When a value is preceded by a `-` character, the corresponding feature flag will be disabled, otherwise the flag will be enabled.
-Note that enabled feature flags should not be preceded by a `+`.
@@ -66,8 +66,8 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.sub
 ::: tinygrad.Tensor.mul
 ::: tinygrad.Tensor.div
+::: tinygrad.Tensor.idiv
 ::: tinygrad.Tensor.mod
-::: tinygrad.Tensor.fmod
 ::: tinygrad.Tensor.bitwise_xor
 ::: tinygrad.Tensor.bitwise_and
 ::: tinygrad.Tensor.bitwise_or
@@ -19,8 +19,8 @@

 ## tinygrad ops

-::: tinygrad.Tensor.linear_with_vars
-::: tinygrad.Tensor.schedule_linear
+::: tinygrad.Tensor.schedule_with_vars
+::: tinygrad.Tensor.schedule
 ::: tinygrad.Tensor.realize
 ::: tinygrad.Tensor.replace
 ::: tinygrad.Tensor.assign
@@ -4,7 +4,7 @@ TinyGPU app lets you use AMD and NVIDIA GPUs on macOS over USB4/Thunderbolt with

 ## Requirements

- macOS (13.0+)
+- macOS (12.1+)
 - USB4/Thunderbolt port
 - A supported GPU (AMD RDNA3+ or NVIDIA Ampere+)

@@ -113,7 +113,7 @@ class VLIWRenderer(Renderer):
        case Ops.GEP:
          # a GEP is just an alias to a special register in the vector
          r[u] = r[u.src[0]] + u.arg[0]
-        case Ops.STACK:
+        case Ops.VECTORIZE:
          if all(s == u.src[0] for s in u.src):
            # if all sources are the same, we can broadcast
            inst.append({"valu": [("vbroadcast", r[u], r[u.src[0]])]})
@@ -173,16 +173,16 @@ if __name__ == "__main__":

  # *** render to device ***

-  from tinygrad.codegen import to_program
+  from tinygrad.codegen import get_program
  with Context(PCONTIG=2, DEVECTORIZE=2, SPEC=0):
    out = tree_traversal(forest_t, val_t, height, rounds)
-    sink = out.schedule_linear().src[-1].src[0]
-    prg = to_program(sink, VLIWRenderer())
+    sink = out.schedule()[-1].ast
+    prg = get_program(sink, VLIWRenderer())

  # *** run on Machine and compare ***

  # NOTE: the scratch size needs to be reduced to 1536 when you have a register allocator
-  src = eval(prg.src[3].arg)
+  src = eval(prg.src)
  max_regs = max(t[1] for instr in src for v in instr.values() for t in v if len(t) > 1) + 8
  print(f"{max_regs:5d} regs used" + ("" if max_regs <= 1536 else "       <-- WARNING: TOO MANY REGISTERS, MUST BE <= 1536"))
  machine = problem.Machine(mem, src, problem.DebugInfo(scratch_map={}), n_cores=1, trace=False, scratch_size=max_regs)
@@ -123,7 +123,7 @@ def NF4Linear(block_size):
    def __call__(self, x: Tensor) -> Tensor:
      high_bits = self.weight
      low_bits = (self.weight * 2 ** 4).contiguous()
-      unpacked = Tensor.stack(high_bits, low_bits, dim=-1).div(2 ** 4, rounding_mode="trunc")
+      unpacked = Tensor.stack(high_bits, low_bits, dim=-1).idiv(2 ** 4)
      unscaled = CODE[unpacked].to(x.device).reshape(-1, block_size) * self.scale
      return x.linear(unscaled.reshape(self.out_features, self.in_features).T)

@@ -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
+  from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8, FP8_DTYPE
  from examples.llama3 import MODEL_PARAMS
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
  from examples.mlperf.optim import GradAccClipAdamW
@@ -1357,7 +1357,6 @@ def train_llama3():
      MLLOGGER.event(key=mllog_constants.OPT_LR_WARMUP_STEPS, value=WARMUP_STEPS)
      MLLOGGER.event(key=mllog_constants.NUM_WARMUP_STEPS, value=WARMUP_STEPS)
      MLLOGGER.event(key=mllog_constants.OPT_LR_DECAY_STEPS, value=MAX_STEPS - WARMUP_STEPS)
-      MLLOGGER.event(key=mllog_constants.OPT_LR_DECAY_SCHEDULE, value="cosine with linear warmup")
      MLLOGGER.event(key=mllog_constants.OPT_GRADIENT_CLIP_NORM, value=1.0)
  else:
    MLLOGGER = None
@@ -1396,7 +1395,7 @@ def train_llama3():

  params = get_parameters(model)

-  if getenv("EMPTYWEIGHT"):
+  if getenv("FAKEDATA"):
    for v in get_parameters(model):
      v = v.assign(Tensor.empty(v.shape, dtype=v.dtype))

@@ -1419,10 +1418,7 @@ def train_llama3():

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

  scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
@@ -1436,22 +1432,18 @@ def train_llama3():
    print(f"loading optim checkpoint from {fn}")
    load_state_dict(scheduler, safe_load(fn), realize=False)

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

-  from tinygrad.nn.state import get_state_dict
-  model_state = get_state_dict(model)
-  for wname in ["wqkv", "wo", "w13", "w2"]:
-    w = model_state[wname]
-    w._inv_scale = model._fp8_inv_scale[wname]
-    if optim.master_params:
-      idx = next(j for j, p in enumerate(optim.params) if p is w)
-      optim.master_params[idx].assign((optim.master_params[idx] * w._inv_scale.reshape(-1, *([1]*(w.ndim-1)))).contiguous())
-
-  # realize everything here
-  if optim.master_params: Tensor.realize(*optim.master_params)
-  Tensor.realize(*optim.params, *fp8_inv_scales, *fp8_amax, *fp8_grad_amax)
+  if FP8:
+    from tinygrad.nn.state import get_state_dict
+    model_state = get_state_dict(model)
+    for wname in ["wqkv", "wo", "w13", "w2"]:
+      w = model_state[wname]
+      w._inv_scale = model._fp8_inv_scale[wname]
+      if optim.master_params:
+        idx = next(j for j, p in enumerate(optim.params) if p is w)
+        optim.master_params[idx].assign((optim.master_params[idx] * w._inv_scale.reshape(-1, *([1]*(w.ndim-1)))).contiguous())

  @TinyJit
  def minibatch(tokens:Tensor):
@@ -1459,17 +1451,13 @@ def train_llama3():
    if is_mp: tokens = tokens.shard(device)
    if not is_sharding: tokens = tokens.to(None)
    logits:Tensor = model(tokens[:, :-1])
-    if getenv("FAST_CE", 0):
-      from extra.llama_kernels.fused_ce import fused_ce_loss
-      loss = fused_ce_loss(logits.cast(dtypes.bfloat16), tokens[:, 1:], label_smoothing=0.0)
-    else:
-      loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
+    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])

    for g, new_g in zip(grads, loss.gradient(*optim.params)):
      apply_grad(g, new_g.uop)

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

  @TinyJit
  def optim_step():
@@ -1567,7 +1555,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 * 4.6e15)) * 100
+      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * (4.6e15 if FP8 else 2.3e15))) * 100
      tqdm.write(
          f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
          f"{lr:.12f} LR, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
@@ -1644,6 +1632,7 @@ def train_llama3():
        tqdm.write(f"target achieved after {sequences_seen} sequences")
        if MLLOGGER and RUNMLPERF:
          MLLOGGER.end(key=mllog_constants.EPOCH_STOP, metadata={mllog_constants.SAMPLES_COUNT: sequences_seen})
+          MLLOGGER.event(key=mllog_constants.TRAIN_SAMPLES, value=sequences_seen)
          MLLOGGER.end(key=mllog_constants.RUN_STOP, metadata={mllog_constants.STATUS: mllog_constants.SUCCESS})
        if getenv("CKPT"):
          if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
@@ -1,4 +1,4 @@
-import math, os
+import math, os, functools
 if __name__ == "__main__":
  os.environ["DEFAULT_FLOAT"] = "bfloat16"
  os.environ["OPTIM_DTYPE"] = "bfloat16"
@@ -16,78 +16,65 @@ from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
 from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
-from extra.llama_kernels.rmsnorm import rmsnorm
-from extra.llama_kernels import FP8_MAX, local_abs_max

-ASM_GEMM = getenv("ASM_GEMM", 0)
-FUSED_INPUT_QUANTIZE = getenv("FUSED_INPUT_QUANTIZE", 0)
-FUSED_ADD_NORM_MUL_QUANTIZE = getenv("FUSED_ADD_NORM_MUL_QUANTIZE", 0)
-FUSED_SILU_W13 = getenv("FUSED_SILU_W13", 0)
+FP8 = getenv("FP8", 0)

 FP8_DTYPE = dtypes.fp8e4m3
 FP8_GRAD_DTYPE = dtypes.fp8e5m2
+FP8_MAX = 448.0
+
+# per-device abs max without allreduce (matches TE delayed scaling behavior)
+@functools.cache
+def _local_abs_max_fxn(x_p, device):
+  x = Tensor(x_p, device=device)
+  inner = Tensor(x.uop.src[0]) if x.uop.op is Ops.MULTI else x
+  return (inner.abs().max(),)
+
+def _local_abs_max(x:Tensor) -> Tensor:
+  param = x.as_param(0)
+  fxn = _local_abs_max_fxn(param.uop, x.device)
+  return Tensor(fxn[0].uop.call(x.uop).gettuple(0))

 def quantize_fp8(x:Tensor, amax_state:Tensor|None=None):
-  new_amax = (local_abs_max(x) if isinstance(x.device, tuple) else x.abs().max()).detach().cast(dtypes.float32)
+  new_amax = (_local_abs_max(x) if isinstance(x.device, tuple) else x.abs().max()).detach()
  scale = FP8_MAX / ((amax_state if amax_state is not None else new_amax) + 1e-8)
  x_scaled = x * scale
  x_clamped = x_scaled + (x_scaled.detach().clamp(-FP8_MAX, FP8_MAX) - x_scaled.detach())  # STE
  return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal(), new_amax

-def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_scale:Tensor|None=None,
-           x_fp8:Tensor|None=None, x_scale:Tensor|None=None, x_new_amax:Tensor|None=None,
-           grad_amax_state:Tensor|None=None) -> tuple[Tensor,...]:
+def matmul(x:Tensor, w:Tensor, fp8=FP8, amax_x:Tensor|None=None, w_inv_scale:Tensor|None=None,
+           x_fp8:Tensor|None=None, x_scale:Tensor|None=None, x_new_amax:Tensor|None=None) -> tuple[Tensor,...]:
  if not fp8:
-    if ASM_GEMM:
+    if getenv("ASM_GEMM"):
      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,)
  assert w_inv_scale is not None, "fp8 matmul requires w_inv_scale (weights must be stored in fp8 with per-tensor scale)"
-  if x_fp8 is None:
-    if FUSED_INPUT_QUANTIZE and amax_x is not None:
-      from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed
-      x_fp8, x_scale, x_new_amax, _ = quantize_fp8_delayed(x, amax_x, FP8_DTYPE)
-    else:
-      x_fp8, x_scale, x_new_amax = quantize_fp8(x, amax_state=amax_x)
-  if ASM_GEMM:
+  if x_fp8 is None: x_fp8, x_scale, x_new_amax = quantize_fp8(x, amax_state=amax_x)
+  if getenv("ASM_GEMM"):
    from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
-    if can_use_asm_gemm(x_fp8, w.T):
-      return asm_gemm(x_fp8, w.T, x_scale=x_scale, w_scale=w_inv_scale, grad_amax_state=grad_amax_state), x_new_amax, x_fp8, w
+    if can_use_asm_gemm(x_fp8, w.T): return asm_gemm(x_fp8, w.T, x_scale=x_scale, w_scale=w_inv_scale), x_new_amax, x_fp8, w
  return x_fp8.dot(w.T, dtype=dtypes.float) * x_scale * w_inv_scale, x_new_amax, x_fp8, w

-def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor, grad_amax_state:Tensor):
-  if FUSED_ADD_NORM_MUL_QUANTIZE:
-    from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_rmsnorm_mul_quantize_fp8
-    x_fp8, x_inv_scale, new_amax, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE)
-    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, x_scale=x_inv_scale, x_new_amax=new_amax, grad_amax_state=grad_amax_state)
-    return out, x_normed, rrms, ret
-  x_normed, rrms = rmsnorm(x, eps)
-  out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state)
-  return out, x_normed, rrms, ret
+def _rmsnorm_fwd(x_in:Tensor, eps:float) -> tuple[Tensor, Tensor]:
+  x = x_in.float()
+  rrms = (x.square().mean(-1, keepdim=True) + eps).rsqrt()
+  return (x * rrms).cast(x_in.dtype), rrms

-def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor):
-  if FUSED_ADD_NORM_MUL_QUANTIZE:
-    from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_add_rmsnorm_mul_quantize_fp8
-    x_fp8, x_inv_scale, new_amax, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE)
-    out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, x_scale=x_inv_scale, x_new_amax=new_amax)
-    return out, h, x_normed, rrms, ret
-  h = x + residual
-  x_normed, rrms = rmsnorm(h, eps)
-  out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale)
-  return out, h, x_normed, rrms, ret
+@functools.cache
+def _rmsnorm_fwd_fxn(x_in_p, eps, device):
+  return _rmsnorm_fwd(Tensor(x_in_p, device=device), eps)

-def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor,
-                             amax_x2:Tensor,
-                             grad_amax_xw13:Tensor, grad_amax_xout:Tensor):
-  if FUSED_SILU_W13:
-    from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
-    x2_fp8, x2_inv_scale, new_amax_x2 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13)
-    out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, x_scale=x2_inv_scale, x_new_amax=new_amax_x2, grad_amax_state=grad_amax_xout)
-    return out, ret
-  hidden = x_w13.shape[-1] // 2
-  x_w1, x_w3 = x_w13[..., :hidden], x_w13[..., hidden:]
-  out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout)
-  return out, ret
+def _rmsnorm_bwd(grad:UOp, call:UOp) -> tuple:
+  x_normed = Tensor(call.gettuple(0)).float()
+  do_float = Tensor(grad).float()
+  d_x = Tensor(call.gettuple(1)) * (do_float - x_normed * (do_float * x_normed).mean(-1, keepdim=True))
+  return (d_x.cast(call.src[1].dtype).uop,)
+
+def rmsnorm(x_in:Tensor, eps:float) -> tuple[Tensor, Tensor]:
+  fxn = _rmsnorm_fwd_fxn(x_in.as_param(0).uop, eps, x_in.device)
+  call = UOp.maketuple(fxn[0].uop, fxn[1].uop).call(x_in.uop, grad_fxn=_rmsnorm_bwd)
+  return Tensor(call.gettuple(0)), Tensor(call.gettuple(1))

 class FlatTransformer:
  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
@@ -103,7 +90,7 @@ class FlatTransformer:
    scaled_std = 0.02 / math.sqrt(2 * n_layers)

    # Attention
-    self._init_inv_scales = []  # populated by lin_per_layer
+    self._init_inv_scales = []  # populated by lin_per_layer when FP8
    self.wqkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
    self.wo = self.lin_per_layer(self.n_heads * self.head_dim, dim, std=scaled_std)

@@ -122,33 +109,37 @@ class FlatTransformer:
    self.output = Tensor.normal(1, vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)

-    def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().requires_grad_(False)
-    names = ["xqkv", "xo", "x13", "x2"]
-    self._fp8_amax = {name: [_amax() for _ in range(n_layers)] for name in names}
-    grad_names = ["xqkv", "xo", "xw13", "xout"]
-    self._fp8_grad_amax = {name: [_amax() for _ in range(n_layers)] for name in grad_names}
-    w_names = ["wqkv", "wo", "w13", "w2"]
-    self._fp8_inv_scale = {wname: inv_scales.float().contiguous().requires_grad_(False)
-                           for wname, inv_scales in zip(w_names, self._init_inv_scales)}
-    del self._init_inv_scales
+    if FP8:
+      def _amax(): return Tensor.full((), FP8_MAX).contiguous().requires_grad_(False)
+      names = ["xqkv", "xo", "x13", "x2"]
+      self._fp8_amax = {name: [_amax() for _ in range(n_layers)] for name in names}
+      # per-weight inv_scale: single (n_layers,) float32 tensor per weight (kernel reads float* pointers)
+      w_names = ["wqkv", "wo", "w13", "w2"]
+      self._fp8_inv_scale = {}
+      for wname, inv_scales in zip(w_names, self._init_inv_scales):
+        self._fp8_inv_scale[wname] = inv_scales.float().contiguous().requires_grad_(False)
+      del self._init_inv_scales

  def lin_per_layer(self, in_features:int, out_features:int, std:float=0.02):
    if getenv("ZEROS"): w = Tensor.zeros(self.n_layers, out_features, in_features)
    else: w = Tensor.normal(self.n_layers, out_features, in_features, mean=0.0, std=std)
+    if not FP8: return w
+    # per-layer scaled fp8 cast: fill the fp8 range for best precision
    amax = w.abs().flatten(1).max(1).detach()
    scale = FP8_MAX / (amax + 1e-8)
-    self._init_inv_scales.append((amax + 1e-8) / FP8_MAX)
+    self._init_inv_scales.append((amax + 1e-8) / FP8_MAX)  # save for inv_scale init
    return (w * scale.reshape(-1, 1, 1)).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE)

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

-    xqkv, x_normed, rrms, ret = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
-                                                     amax_x=amax_xqkv, grad_amax_state=grad_amax_xqkv)
-    saves.extend([x_normed, rrms])
+    x, rrms = rmsnorm(x, self.norm_eps)
+    saves.extend([x, rrms])
+    x = x * attention_norm
+
+    xqkv, *ret = matmul(x, wqkv, amax_x=amax_xqkv, w_inv_scale=s_qkv)
    new_amaxs.extend(ret[:1])
    saves.extend(ret[1:] + [xqkv])
    xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
@@ -157,53 +148,60 @@ class FlatTransformer:
    xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)

    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
-    xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
+    if FP8: xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
+    xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
    if getenv("HK_FLASH_ATTENTION"):
      from extra.thunder.amd.fa import flash_attention
      attn, *save = flash_attention(xq, xk, xv, is_causal=True)
      saves.extend(save)
    else:
-      xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
-      attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
-    attn = attn.reshape(bsz, seqlen, -1)
+      attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True)
+    attn = attn.transpose(1, 2).reshape(bsz, seqlen, -1)

-    out, *ret = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo)
+    out, *ret = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o)
    new_amaxs.extend(ret[:1])
    saves.extend(ret[1:] + [out])
    return (out, *new_amaxs, *saves)

-  def feed_forward(self, x:Tensor, residual:Tensor, ffn_norm:Tensor, w13:Tensor, w2:Tensor,
-                   amax_x13:Tensor, amax_x2:Tensor, s_13:Tensor, s_2:Tensor,
-                   grad_amax_xw13:Tensor, grad_amax_xout:Tensor):
+  def feed_forward(self, x:Tensor, ffn_norm:Tensor, w13:Tensor, w2:Tensor,
+                   amax_x13=None, amax_x2=None, s_13=None, s_2=None):
    new_amaxs, saves = [], []

-    x_w13, h, x_normed, rrms, ret = add_norm_quantize_matmul(x, residual, ffn_norm, w13, s_13, self.norm_eps,
-                                                             amax_x=amax_x13)
-    saves.extend([x_normed, rrms])
+    x, rrms = rmsnorm(x, self.norm_eps)
+    saves.extend([x, rrms])
+    x = x * ffn_norm
+
+    x_w13, *ret = matmul(x, w13, amax_x=amax_x13, w_inv_scale=s_13)
    new_amaxs.extend(ret[:1])
    saves.extend(ret[1:] + [x_w13])

-    out, ret = silu_w13_quantize_matmul(x_w13, w2, s_2, amax_x2=amax_x2, grad_amax_xw13=grad_amax_xw13, grad_amax_xout=grad_amax_xout)
+    if FP8 and getenv("FUSED_SILU_W13", 1):
+      from extra.amax.cast_amax import fused_quantize_fp8_w13
+      amax_s = amax_x2 if amax_x2 is not None else Tensor.full((), 1.0, dtype=dtypes.bfloat16, device=x_w13.device)
+      x2_fp8, x2_inv_scale, new_amax_x2 = fused_quantize_fp8_w13(x_w13, amax_s, FP8_DTYPE)
+      out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, x_scale=x2_inv_scale, x_new_amax=new_amax_x2)
+    else:
+      x_w1, x_w3 = x_w13[..., :self.hidden_dim], x_w13[..., self.hidden_dim:]
+      out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2)
    new_amaxs.extend(ret[:1])
    saves.extend(ret[1:] + [out])
-    return (out, h, *new_amaxs, *saves)
+    return (out, *new_amaxs, *saves)

  @function(precompile=True, precompile_backward=True)
  def run_layer(self, x:Tensor, freqs_cis:Tensor,
                attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
                ffn_norm:Tensor, w13:Tensor, w2:Tensor,
-                amax_xqkv:Tensor, amax_xo:Tensor,
-                amax_x13:Tensor, amax_x2:Tensor,
-                s_qkv:Tensor, s_o:Tensor, s_13:Tensor, s_2:Tensor,
-                grad_amax_xqkv:Tensor, grad_amax_xo:Tensor,
-                grad_amax_xw13:Tensor, grad_amax_xout:Tensor):
+                amax_xqkv=None, amax_xo=None,
+                amax_x13=None, amax_x2=None,
+                s_qkv=None, s_o=None, s_13=None, s_2=None):
    attn, *attn_ret = self.attention(x, freqs_cis, attention_norm, wqkv, wo,
-                                     amax_xqkv=amax_xqkv, amax_xo=amax_xo, s_qkv=s_qkv, s_o=s_o,
-                                     grad_amax_xqkv=grad_amax_xqkv, grad_amax_xo=grad_amax_xo)
+                                     amax_xqkv=amax_xqkv, amax_xo=amax_xo,
+                                     s_qkv=s_qkv, s_o=s_o)
    attn_amaxs, attn_saves = attn_ret[:2], attn_ret[2:]
-    ffn, h, *ffn_ret = self.feed_forward(x, attn, ffn_norm, w13, w2,
-                                                   amax_x13=amax_x13, amax_x2=amax_x2, s_13=s_13, s_2=s_2,
-                                                   grad_amax_xw13=grad_amax_xw13, grad_amax_xout=grad_amax_xout)
+    h = x + attn
+    ffn, *ffn_ret = self.feed_forward(h, ffn_norm, w13, w2,
+                                      amax_x13=amax_x13, amax_x2=amax_x2,
+                                      s_13=s_13, s_2=s_2)
    ffn_amaxs, ffn_saves = ffn_ret[:2], ffn_ret[2:]
    h = h + ffn
    return (h, *attn_amaxs, *ffn_amaxs, *attn_saves, *ffn_saves)
@@ -215,40 +213,43 @@ class FlatTransformer:
    else:
      # flat per-layer weights: axis 0 is n_layers, so shard axes are +1 vs per-layer Transformer
      self.wqkv.shard_(device, axis=1).realize()          # (n_layers, out, dim) shard out
-      self.wo.shard_(device, axis=2).realize()            # (n_layers, dim, in) shard in
-      self.w13.shard_(device, axis=1).realize()           # (n_layers, hidden*2, dim) shard out
-      self.w2.shard_(device, axis=2).realize()            # (n_layers, dim, hidden) shard in
+      self.wo.shard_(device, axis=2).realize()             # (n_layers, dim, in) shard in
+      self.w13.shard_(device, axis=1).realize()             # (n_layers, hidden*2, dim) shard out
+      self.w2.shard_(device, axis=2).realize()             # (n_layers, dim, hidden) shard in
      self.attention_norm.shard_(device, axis=None).realize()
      self.ffn_norm.shard_(device, axis=None).realize()
      self.norm.weight.shard_(device, axis=None).realize()
      self.tok_embeddings.weight.shard_(device, axis=0).realize()
      self.output.shard_(device, axis=1).realize()
      self.freqs_cis.shard_(device, axis=None).realize()
-      for amax_dict in (self._fp8_amax, self._fp8_grad_amax):
-        for name in amax_dict:
-          for i in range(len(amax_dict[name])):
-            amax_dict[name][i] = amax_dict[name][i].to(device).contiguous().requires_grad_(False)
-      for name in self._fp8_inv_scale:
-        self._fp8_inv_scale[name] = self._fp8_inv_scale[name].to(device).contiguous().requires_grad_(False)
+      if FP8:
+        for name in self._fp8_amax:
+          for i in range(len(self._fp8_amax[name])):
+            self._fp8_amax[name][i] = self._fp8_amax[name][i].to(device).contiguous().requires_grad_(False)
+        for name in self._fp8_inv_scale:
+          self._fp8_inv_scale[name] = self._fp8_inv_scale[name].to(device).contiguous().requires_grad_(False)

  def __call__(self, tokens:Tensor):
    h = self.tok_embeddings(tokens)
    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
-    a, ga, s = self._fp8_amax, self._fp8_grad_amax, self._fp8_inv_scale
+    a = self._fp8_amax if FP8 else None
+    s = self._fp8_inv_scale if FP8 else None
    for i in range(self.n_layers):
+      amax_layer = {"amax_xqkv": a["xqkv"][i], "amax_xo": a["xo"][i],
+                    "amax_x13": a["x13"][i], "amax_x2": a["x2"][i]} if a else {}
+      scale_layer = {"s_qkv": s["wqkv"][i], "s_o": s["wo"][i],
+                     "s_13": s["w13"][i], "s_2": s["w2"][i]} if s else {}
      h, *ret = self.run_layer(h, freqs_cis,
                               self.attention_norm[i], self.wqkv[i], self.wo[i],
                               self.ffn_norm[i], self.w13[i], self.w2[i],
-                               amax_xqkv=a["xqkv"][i], amax_xo=a["xo"][i],
-                               amax_x13=a["x13"][i], amax_x2=a["x2"][i],
-                               s_qkv=s["wqkv"][i], s_o=s["wo"][i],
-                               s_13=s["w13"][i], s_2=s["w2"][i],
-                               grad_amax_xqkv=ga["xqkv"][i], grad_amax_xo=ga["xo"][i],
-                               grad_amax_xw13=ga["xw13"][i], grad_amax_xout=ga["xout"][i])
-      for name, new_val in zip(["xqkv", "xo", "x13", "x2"], ret[:5]):
-        a[name][i].assign(new_val)
+                               **amax_layer, **scale_layer)
+      if a:
+        amaxs = ret[:5]
+        amax_names = ["xqkv", "xo", "x13", "x2"]
+        for name, new_val in zip(amax_names, amaxs):
+          a[name][i].assign(new_val)

-    logits = matmul(self.norm(h), self.output[0], fp8=False)[0]
+    logits = matmul(self.norm(h).contiguous().contiguous_backward(), self.output[0], fp8=False)[0].contiguous_backward()
    return logits

 def _get_pads(uop:UOp) -> list[UOp]:
@@ -257,19 +258,13 @@ def _get_pads(uop:UOp) -> list[UOp]:

 def apply_grad(grad_buf:Tensor, new_grad:UOp):
  pads = _get_pads(new_grad)
+  new_grad = new_grad.cast(grad_buf.dtype)
  if len(pads) <= 1:
-    new_grad = new_grad.cast(grad_buf.dtype)
    store = grad_buf.uop.store(grad_buf.uop + new_grad)
    grad_buf.uop = grad_buf.uop.after(store)
    return
  sorted_pads = sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0)
-  inners_raw = [Tensor(p.src[0] if p.op == Ops.PAD else p, device=grad_buf.device) for p in sorted_pads]
-  if getenv("FUSED_PAD_GRAD_ACCUM", 0):
-    from extra.llama_kernels.fused_pad_grad_accum import fused_pad_grad_accum, can_fused_pad_grad_accum
-    if can_fused_pad_grad_accum(grad_buf, inners_raw):
-      grad_buf.uop = fused_pad_grad_accum(grad_buf, inners_raw).uop
-      return
-  inners = [t.cast(grad_buf.dtype) for t in inners_raw]
+  inners = [Tensor(p.src[0] if p.op == Ops.PAD else p, device=grad_buf.device).cast(grad_buf.dtype) for p in sorted_pads]
  grad_buf.assign(grad_buf + inners[0].cat(*inners[1:], dim=0))

 if __name__ == "__main__":
@@ -81,7 +81,7 @@ class GradAccClipAdamW(Optimizer):
    if STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16: return stochastic_round_bf16(new_w)
    if t.dtype in dtypes.fp8s:
      from examples.mlperf.models.flat_llama import FP8_MAX
-      amax = new_w.float().abs().max(axis=tuple(range(1, new_w.ndim))).detach()  # per-layer amax for (n_layers, out, in)
+      amax = new_w.float().abs().flatten(1).max(1).detach()  # per-layer amax for (n_layers, out, in)
      scale = FP8_MAX / (amax + 1e-8)
      fp8_w = (new_w * scale.reshape(-1, *([1]*(new_w.ndim-1)))).clamp(-FP8_MAX, FP8_MAX).cast(t.dtype)
      if hasattr(t, '_inv_scale'):
@@ -2,6 +2,7 @@

 export PYTHONPATH="."
 export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
 export DEVICE_IN_FUNCTION_BUG=1
@@ -9,22 +10,14 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export USE_ATOMICS=${USE_ATOMICS:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export FP8=${FP8:-1}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-0}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
-export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
-export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
-export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}

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

 export MODEL="llama3"
 export BASEDIR="/raid/datasets/c4/"
@@ -37,7 +30,7 @@ export DATA_SEED=${DATA_SEED:-5760}
 export JITBEAM=${JITBEAM:-3}
 export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1

-export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+export FAKEDATA=1 BENCHMARK=10
 if [ -z "$FULL_LAYERS" ]; then
  export LLAMA_LAYERS=2
 fi
@@ -9,21 +9,15 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-1}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
-export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
-export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
-export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1}

 export 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 DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-16} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
 export GBS=$((BS * GRADIENT_ACC_STEPS))

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

-export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+export FAKEDATA=1 BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
  export LLAMA_LAYERS=2
 fi
@@ -2,6 +2,7 @@

 export PYTHONPATH="."
 export DEV=${DEV:-AMD}
+export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
 export DEVICE_IN_FUNCTION_BUG=1
@@ -9,17 +10,9 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export USE_ATOMICS=${USE_ATOMICS:-1}
+export USE_ATOMICS=${USE_ATOMICS:-0}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
-export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export FP8=${FP8:-1}
-export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-0}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
-export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
-export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
-export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
@@ -42,7 +35,7 @@ export DATA_SEED=${DATA_SEED:-5760}
 export JITBEAM=${JITBEAM:-3}
 export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1

-export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+export FAKEDATA=1 BENCHMARK=10
 if [ -z "$FULL_LAYERS" ]; then
  export LLAMA_LAYERS=2
 fi
@@ -9,21 +9,15 @@ export DEVICE_IN_FUNCTION_BUG=1
 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
 export ALL2ALL=${ALL2ALL:-1}
-export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
 export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
-export FAST_CE=${FAST_CE:-1}
-export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
-export FUSED_ADD_NORM_MUL_QUANTIZE=${FUSED_ADD_NORM_MUL_QUANTIZE:-1}
-export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
-export FUSED_PAD_GRAD_ACCUM=${FUSED_PAD_GRAD_ACCUM:-1}

 export 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 DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-16} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
 export GBS=$((BS * GRADIENT_ACC_STEPS))

 export MODEL="llama3"
@@ -3,4 +3,4 @@ export BENCHMARK=5
 export EVAL_BS=0
 VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh
 SRC="AMD"; [[ $DEV == NULL* ]] && SRC="NULL"
-python -m tinygrad.viz.cli -s "$SRC" -t
+python -m tinygrad.viz.cli --profile -s "$SRC" --top 20
@@ -10,21 +10,15 @@ export DEVICE_IN_FUNCTION_BUG=1

 export HK_FLASH_ATTENTION=1
 export ALL2ALL=1
-export LATE_ALLREDUCE=0
 export USE_ATOMICS=1
 export ASM_GEMM=1
 export WQKV=1
 export MASTER_WEIGHTS=1
 export FP8=1
 export ALLREDUCE_CAST=1
-export FAST_CE=1
-export FUSED_INPUT_QUANTIZE=1
-export FUSED_ADD_NORM_MUL_QUANTIZE=1
-export FUSED_SILU_W13=1
-export FUSED_PAD_GRAD_ACCUM=1

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

 export MODEL="llama3"
@@ -5,6 +5,7 @@ if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"
 from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device, dtypes
 from tinygrad.helpers import DEBUG, getenv
 from tinygrad.uop.ops import Ops
+from tinygrad.engine.realize import get_runner
 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"
@@ -36,7 +37,7 @@ def compile(onnx_file):
    # copy i == 1 so use of JITBEAM is okay
    if i == 1: test_val = np.copy(ret)
  # iterate kernel CALLs in the captured LINEAR UOp; toposort descends into batched graph CUSTOM_FUNCTIONs
-  kernel_asts = {Ops.PROGRAM}
+  kernel_asts = {Ops.SINK, Ops.PROGRAM, Ops.BEAM}
  kernel_calls = [u for u in run_onnx_jit.captured.linear.toposort(gate=lambda x: x.op not in kernel_asts)
                  if u.op is Ops.CALL and u.src[0].op in kernel_asts]
  print(f"captured {len(kernel_calls)} kernels")
@@ -48,8 +49,8 @@ def compile(onnx_file):
  read_image_count = 0
  gated_read_image_count = 0
  for call in kernel_calls:
-    _, _, _, source, _ = call.src[0].src
-    src = source.arg
+    device = next(b.device for b in call.src[1:] if b.op is not Ops.BIND)
+    src = get_runner(device, call.src[0]).p.src
    kernel_count += 1
    read_image_count += src.count("read_image")
    gated_read_image_count += src.count("?read_image")
@@ -133,20 +134,14 @@ def bench(run, inputs):
      run(**inputs).numpy()

 if __name__ == "__main__":
-  if getenv("RUN_PICKLE"):
-    with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
-    inputs = {name: Tensor(Tensor.randn(*[int(s) for s in view.src[1].arg], dtype=dtype).numpy(), device=device)
-              for name, (view, _vars, dtype, device) in zip(pickle_loaded.captured.expected_names, pickle_loaded.captured.expected_input_info)}
-    test_vs_compile(pickle_loaded, inputs)
-  else:
-    onnx_file = fetch(OPENPILOT_MODEL)
-    inputs, outputs = compile(onnx_file)
+  onnx_file = fetch(OPENPILOT_MODEL)
+  inputs, outputs = compile(onnx_file)

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

-    test_vs_compile(pickle_loaded, inputs, outputs)
-    if getenv("SELFTEST"):
-      test_vs_onnx(inputs, outputs, onnx_file, 1e-4)
+  test_vs_compile(pickle_loaded, inputs, outputs)
+  if getenv("SELFTEST"):
+    test_vs_onnx(inputs, outputs, onnx_file, 1e-4)

  if getenv("BENCHMARK_LOG", ""):
    bench(pickle_loaded, inputs)
@@ -1,6 +1,7 @@
 from tinygrad import Tensor, Device, TinyJit, dtypes
+from tinygrad.helpers import getenv

-GPUS = Device[Device.DEFAULT].count()
+GPUS = getenv("GPUS", 4) # TODO: expose a way in tinygrad to access this
 N = 6144

@TinyJit
@@ -114,7 +114,7 @@ if __name__ == "__main__":
    linear, output_bufs = jit_model(step, *step.input)
    functions, statements, bufs, _ = compile_net(linear, output_bufs)
    state = get_state_dict(model)
-    weights = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None}
+    weights = {id(x.uop.base.realized): name for name, x in state.items()}
    kernel_code = '\n\n'.join([f"const {key} = `{fixup_code(code, key)}`;" for key, code in functions.items()])
    kernel_names = ', '.join([name for (name, _, _, _) in statements])
    input_names = [f"input{i}" for i in range(len(step.input))]
@@ -0,0 +1,85 @@
+import functools, pathlib
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
+
+FP8_MAX = 448.0
+NUM_WG, THREADS_PER_WG = 1024, 256
+
+def _compile(cpp_name:str, n_elems:int, hidden:int):
+  src = (pathlib.Path(__file__).parent/cpp_name).read_text()
+  defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={hidden}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
+  return src, HIPCCCompiler("gfx950", ["-std=c++20", "-ffast-math", *defines]).compile_cached(src)
+
+def _shard_shape(shape:tuple, axis:int, ndev:int) -> list:
+  s = list(shape); s[axis] //= ndev; return s
+
+@functools.cache
+def _custom_fused_bwd_w13(grad_xw13:UOp, xw13:UOp, grad_x2:UOp, amax_state:UOp, dname:str) -> UOp:
+  hidden = xw13.shape[2] // 2
+  n_elems = xw13.shape[0] * xw13.shape[1] * hidden
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
+  # read 2*N bf16 (xw13) + N bf16 (grad_x2) + 1 scalar; write 2*N bf16 (grad_xw13)
+  mem = n_elems * 2 * 5
+  sink = UOp.sink(grad_xw13.base, xw13.base, grad_x2.base, amax_state.base, threads, workgroups,
+                  arg=KernelInfo(f"fused_silu_mul_bwd_w13_{n_elems}", estimates=Estimates(ops=8*n_elems, mem=mem)))
+  src, lib = _compile("cast_amax_bwd_w13.cpp", n_elems, hidden)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
+
+@functools.cache
+def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_buf:UOp, xw13:UOp, amax_state:UOp, dname:str) -> UOp:
+  hidden = xw13.shape[2] // 2
+  n_elems = xw13.shape[0] * xw13.shape[1] * hidden
+  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
+  # read 2*N bf16 + 1 scalar, write N fp8 + NUM_WG bf16
+  mem = n_elems * 2 * 2 + n_elems + NUM_WG * 2
+  sink = UOp.sink(fp8_out.base, amax_buf.base, xw13.base, amax_state.base, threads, workgroups,
+                  arg=KernelInfo(f"fused_silu_mul_cast_amax_w13_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=mem)))
+  src, lib = _compile("cast_amax_fwd_w13.cpp", n_elems, hidden)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
+                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
+
+def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
+  # kernel.src[1:] is (fp8_out, amax_buf, xw13, amax_state); only xw13 needs a grad
+  _, _, xw13, amax_state = kernel.src[1:]
+  device = xw13.device
+  if isinstance(device, tuple):
+    axis, ndev = xw13.axis, len(device)
+    assert axis in (0, 1), f"unsupported sharding axis={axis}"
+    grad_xw13 = Tensor(Tensor.invalid(*_shard_shape(xw13.shape, axis, ndev), dtype=dtypes.bfloat16, device=device).uop.multi(axis), device=device)
+    dname = device[0].split(":")[0]
+  else:
+    grad_xw13 = Tensor.invalid(*xw13.shape, dtype=dtypes.bfloat16, device=device)
+    dname = device.split(":")[0] if isinstance(device, str) else device
+  grad_x2_t = Tensor(gradient, device=device).cast(dtypes.bfloat16)
+  fxn = functools.partial(_custom_fused_bwd_w13, dname=dname)
+  grad_xw13, *_ = Tensor.custom_kernel(grad_xw13, Tensor(xw13, device=device), grad_x2_t, Tensor(amax_state, device=device), fxn=fxn)
+  return (None, None, grad_xw13.uop, None)
+
+def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype) -> tuple[Tensor, Tensor, Tensor]:
+  # silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns (fp8, inv_scale, new_amax).
+  assert xw13.dtype == dtypes.bfloat16, f"expected bf16, got {xw13.dtype}"
+  MBS, SEQ, H2 = xw13.shape
+  assert H2 % 2 == 0, f"w13 last-axis must be even, got {H2}"
+  HIDDEN = H2 // 2
+  if isinstance(xw13.device, tuple):
+    axis, ndev = xw13.uop.axis, len(xw13.device)
+    assert axis in (0, 1), f"unsupported sharding axis={axis}"
+    fp8_out = Tensor(Tensor.invalid(*_shard_shape((MBS, SEQ, HIDDEN), axis, ndev), dtype=fp8_dtype, device=xw13.device).uop.multi(axis), device=xw13.device)
+    amax_buf = Tensor(Tensor.invalid(NUM_WG, dtype=dtypes.bfloat16, device=xw13.device).uop.multi(0), device=xw13.device)
+    dname = xw13.device[0].split(":")[0]
+  else:
+    fp8_out = Tensor.invalid(MBS, SEQ, HIDDEN, dtype=fp8_dtype, device=xw13.device)
+    amax_buf = Tensor.invalid(NUM_WG, dtype=dtypes.bfloat16, device=xw13.device)
+    dname = xw13.device.split(":")[0] if isinstance(xw13.device, str) else xw13.device
+  fxn = functools.partial(_custom_fused_cast_amax_w13, dname=dname)
+  fp8_out, amax_buf, *_ = Tensor.custom_kernel(fp8_out, amax_buf, xw13, amax_state, fxn=fxn, grad_fxn=_fused_quantize_bwd_w13)
+  # per-device scalar amax (no cross-device allreduce, matches _local_abs_max semantics)
+  if isinstance(amax_buf.device, tuple):
+    from examples.mlperf.models.flat_llama import _local_abs_max
+    new_amax = _local_abs_max(amax_buf).detach()
+  else: new_amax = amax_buf.max().detach()
+  inv_scale = (FP8_MAX / (amax_state + 1e-8)).float().reciprocal()
+  return fp8_out, inv_scale, new_amax
@@ -0,0 +1,68 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_bf16.h>
+
+#ifndef N_ELEMS
+#define N_ELEMS 234881024
+#endif
+#ifndef HIDDEN
+#define HIDDEN 14336
+#endif
+#ifndef NUM_WG
+#define NUM_WG 1024
+#endif
+#ifndef THREADS_PER_WG
+#define THREADS_PER_WG 256
+#endif
+
+constexpr int VEC = 8;
+constexpr float FP8_MAX = 448.0f;
+
+static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC");
+static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC");
+
+extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
+fused_silu_mul_bwd_w13(
+    __hip_bfloat16*       __restrict__ grad_xw13_out,    // bf16, 2*N_ELEMS (interleaved layout)
+    const __hip_bfloat16* __restrict__ xw13,             // bf16, 2*N_ELEMS (interleaved)
+    const __hip_bfloat16* __restrict__ grad_x2,          // bf16, N_ELEMS
+    const __hip_bfloat16* __restrict__ amax_state)       // bf16 scalar
+{
+  const int tid = threadIdx.x;
+  const int wg  = blockIdx.x;
+  const int gid = wg * THREADS_PER_WG + tid;
+  const int stride_elems = NUM_WG * THREADS_PER_WG * VEC;
+
+  const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
+
+  for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) {
+    const int outer = base / HIDDEN;
+    const int inner = base % HIDDEN;
+    const int xw1_off = outer * 2 * HIDDEN + inner;
+    const int xw3_off = xw1_off + HIDDEN;
+
+    float4 x1_raw = *reinterpret_cast<const float4*>(&xw13[xw1_off]);
+    float4 x3_raw = *reinterpret_cast<const float4*>(&xw13[xw3_off]);
+    float4 g_raw  = *reinterpret_cast<const float4*>(&grad_x2[base]);
+
+    const __hip_bfloat16 *x1 = reinterpret_cast<const __hip_bfloat16*>(&x1_raw);
+    const __hip_bfloat16 *x3 = reinterpret_cast<const __hip_bfloat16*>(&x3_raw);
+    const __hip_bfloat16 *gv = reinterpret_cast<const __hip_bfloat16*>(&g_raw);
+
+    __hip_bfloat16 out1[VEC], out3[VEC];
+    #pragma unroll
+    for (int i = 0; i < VEC; i++) {
+      const float f1 = static_cast<float>(x1[i]);
+      const float f3 = static_cast<float>(x3[i]);
+      const float fg = static_cast<float>(gv[i]);
+      const float sig = 1.0f / (1.0f + __expf(-f1));
+      const float silu = f1 * sig;
+      const float silu_prime = sig + silu * (1.0f - sig);
+      const float gs = fg * scale;
+      out1[i] = static_cast<__hip_bfloat16>(gs * silu_prime * f3);
+      out3[i] = static_cast<__hip_bfloat16>(gs * silu);
+    }
+
+    *reinterpret_cast<float4*>(&grad_xw13_out[xw1_off]) = *reinterpret_cast<float4*>(out1);
+    *reinterpret_cast<float4*>(&grad_xw13_out[xw3_off]) = *reinterpret_cast<float4*>(out3);
+  }
+}
@@ -24,9 +24,9 @@ static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC (so VEC loads
 extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
 fused_silu_mul_cast_amax_w13(
    __hip_fp8_storage_t*  __restrict__ fp8_out,         // fp8, N_ELEMS
-    float*                __restrict__ amax_buf,        // fp32, NUM_WG (per-WG amaxes)
+    __hip_bfloat16*       __restrict__ amax_buf,        // bf16, NUM_WG (per-WG amaxes)
    const __hip_bfloat16* __restrict__ xw13,            // bf16, 2*N_ELEMS
-    const float*          __restrict__ amax_state)      // fp32 scalar
+    const __hip_bfloat16* __restrict__ amax_state)      // bf16 scalar
 {
  __shared__ float sdata[THREADS_PER_WG];

@@ -75,5 +75,5 @@ fused_silu_mul_cast_amax_w13(
    __syncthreads();
  }

-  if (tid == 0) amax_buf[wg] = sdata[0];
+  if (tid == 0) amax_buf[wg] = static_cast<__hip_bfloat16>(sdata[0]);
 }
@@ -64,7 +64,7 @@ def get_bar0_size(pcibus):

 class AMSMI(AMDev):
  def __init__(self, pcibus, vram_bar:MMIOInterface, doorbell_bar:MMIOInterface, mmio_bar:MMIOInterface):
-    self.pcibus, self.devfmt = pcibus, pcibus
+    self.pcibus = pcibus
    self.vram, self.doorbell64, self.mmio = vram_bar, doorbell_bar, mmio_bar
    self.pci_state = self.read_pci_state()
    if self.pci_state == "D0": self._init_from_d0()
@@ -91,7 +91,6 @@ class SMICtx:
    self.prev_lines_cnt = 0
    self.prev_terminal_width = 0
    self.prev_terminal_height = 0
-    self.prev_metrics = {}

    remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:", "Processing accelerators:"]
    lspci = subprocess.check_output(["lspci"]).decode("utf-8").splitlines()
@@ -236,29 +235,6 @@ class SMICtx:
      case (13,0,12): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.SocketPowerLimit)
      case _: return metrics.SmuMetrics.AverageSocketPower, metrics.SmuMetrics.dGPU_W_MAX

-  def get_throttle_info(self, dev, metrics):
-    match dev.ip_ver[am.MP1_HWIP]:
-      case (13,0,6)|(13,0,12):
-        throttle_fields = [('ProchotResidencyAcc', 'Prochot'), ('PptResidencyAcc', 'PPT'),
-                           ('SocketThmResidencyAcc', 'Socket Thm'), ('VrThmResidencyAcc', 'VR Thm'), ('HbmThmResidencyAcc', 'HBM Thm')]
-        prev = self.prev_metrics.get(dev.pcibus)
-        active = []
-        if prev is not None:
-          acc_delta = metrics.AccumulationCounter - prev.AccumulationCounter
-          if acc_delta > 0:
-            for field, name in throttle_fields:
-              delta = getattr(metrics, field) - getattr(prev, field)
-              if delta > 0 and (pct := min(100, (delta * 100 + acc_delta // 2) // acc_delta)) > 0: active.append((name, pct))
-        return active
-      case _:
-        smu_mod = dev.smu.smu_mod
-        throttler_names = {getattr(smu_mod, a): a[len('THROTTLER_'):-len('_BIT')]
-                           for a in dir(smu_mod) if a.startswith('THROTTLER_') and a.endswith('_BIT')}
-        active = []
-        for i, pct in enumerate(metrics.SmuMetrics.ThrottlingPercentage):
-          if pct > 0: active.append((throttler_names.get(i, f"UNK_{i}"), int(pct)))
-        return active
-
  def get_mem_usage(self, dev):
    usage = 0
    pt_stack = [dev.mm.root_page_table]
@@ -305,13 +281,6 @@ class SMICtx:
                    + [f"MEM Activity {draw_bar(self.get_mem_activity(dev, metrics) / 100, activity_line_width)}"] \
                    + [f"MEM Usage    {draw_bar(mem_used / mem_total, activity_line_width, opt_text=mem_fmt)}"] \

-      throttle_info = self.get_throttle_info(dev, metrics)
-      if throttle_info:
-        throttle_text = colored(', '.join(f"{name} {pct}%" for name, pct in throttle_info), "red")
-      else:
-        throttle_text = colored("None", "green")
-      activity_line += [f"Throttle     {throttle_text}" + " " * (activity_line_width + 2)]
-
      temps_data, temps_data_compact = self.get_temps(dev, metrics), self.get_temps(dev, metrics, compact=True)
      temps_table = ["=== Temps (°C) ==="] + [f"{name:<16}: {color_temp(val)}" for name, val in temps_data.items()]
      temps_table_compact = ["Temps (°C):" + '/'.join([f"{color_temp(val)} {name}" for name, val in temps_data_compact.items()])]
@@ -355,8 +324,6 @@ class SMICtx:

      dev_content.append(device_line + activity_line + same_line([temps_table, power_table, frequency_table]))

-    self.prev_metrics = {dev.pcibus: m for dev, m in dev_metrics.items() if m is not None}
-
    raw_text = 'AM Monitor'.center(terminal_width) + "\n" + "=" * terminal_width + "\n\n"
    for i in range(0, len(dev_content), 2):
      if i + 1 < len(dev_content): raw_text += '\n'.join(same_line([dev_content[i], dev_content[i+1]], split=padding))
@@ -28,7 +28,15 @@
 // #include "soc15_ih_clientid.h"
 // #include "amdgpu_ih.h"

+#define int32_t int
+#define uint32_t unsigned int
+#define int8_t signed char
+#define uint8_t unsigned char
+#define uint16_t unsigned short
+#define int16_t short
+#define uint64_t unsigned long long
 #define bool _Bool
+#define u32 unsigned int

 #define AMDGPU_MAX_IRQ_SRC_ID		0x100
 #define AMDGPU_MAX_IRQ_CLIENT_ID	0x100
@@ -22,7 +22,15 @@
 #ifndef __AMDGPU_SMU_H__
 #define __AMDGPU_SMU_H__

+#define int32_t int
+#define uint32_t unsigned int
+#define int8_t signed char
+#define uint8_t unsigned char
+#define uint16_t unsigned short
+#define int16_t short
+#define uint64_t unsigned long long
 #define bool _Bool
+#define u32 unsigned int

 #define SMU_THERMAL_MINIMUM_ALERT_TEMP		0
 #define SMU_THERMAL_MAXIMUM_ALERT_TEMP		255
@@ -24,7 +24,15 @@
 #define __AMDGPU_UCODE_H__

 // #include "amdgpu_socbb.h"
+#define int32_t int
+#define uint32_t unsigned int
+#define int8_t signed char
+#define uint8_t unsigned char
+#define uint16_t unsigned short
+#define int16_t short
+#define uint64_t unsigned long long
 #define bool _Bool
+#define u32 unsigned int

 struct common_firmware_header {
 	uint32_t size_bytes; /* size of the entire header+image(s) in bytes */
@@ -6,13 +6,13 @@ from tinygrad.engine.jit import TinyJit
 from tinygrad.nn.state import get_state_dict
 from tinygrad.helpers import Context, to_mv, prod
 from tinygrad.uop.ops import Ops, UOp
-from tinygrad.codegen import to_program
+from tinygrad.codegen import get_program
 import json
 from collections import OrderedDict

 EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CPU", "CUDA", "CL"]

-_KERNEL_ASTS = {Ops.SINK, Ops.PROGRAM}
+_KERNEL_ASTS = {Ops.SINK, Ops.PROGRAM, Ops.BEAM}
 def iter_kernel_calls(linear:UOp):
  """Yield kernel CALLs from a LINEAR UOp. Toposort descends naturally into CUSTOM_FUNCTION graph batches; gate stops at kernel ASTs."""
  return (u for u in linear.toposort(gate=lambda x: x.op not in _KERNEL_ASTS) if u.op is Ops.CALL and u.src[0].op in _KERNEL_ASTS)
@@ -36,11 +36,10 @@ def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], Li

  for call in iter_kernel_calls(linear):
    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[3].arg
-    cargs = [name_of(bu, i == 0) for i, bu in enumerate(arg_uops)] + [v for v in info.vars if v.op is Ops.DEFINE_VAR]
-    statements.append((info.function_name, cargs, info.global_size, info.local_size))
+    prg = get_program(call.src[0], Device[arg_uops[0].device].renderer)
+    functions[prg.function_name] = prg.src
+    cargs = [name_of(bu, i == 0) for i, bu in enumerate(arg_uops)] + [v for v in prg.vars if v.op is Ops.DEFINE_VAR]
+    statements.append((prg.function_name, cargs, prg.global_size, prg.local_size))

  return functions, statements, {name:(size, dtype, key) for name, size, dtype, key in bufs.values()}, bufs_to_save

@@ -245,7 +244,7 @@ def export_model(model, target:str, *inputs, model_name: Optional[str] = "model"
  with Context(JIT=2, CPU_COUNT=1): linear, output_bufs = jit_model(model, *inputs)
  functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs)
  state = get_state_dict(model)
-  weight_names = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None}
+  weight_names = {id(x.uop.base.realized): name for name, x in state.items()}
  input_names = [f"input{i}" for i in range(len(inputs))]
  output_names = [f"output{i}" for i in range(len(output_bufs))]

@@ -13,7 +13,7 @@ from tinygrad import Tensor, Device, Context, GlobalCounters
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.helpers import getenv, colored
 from tinygrad.dtype import dtypes, AddrSpace
-from tinygrad.engine.realize import Estimates, run_linear
+from tinygrad.engine.realize import Estimates
 from tinygrad.renderer.amd.dsl import s, v, VCC_LO, NULL
 from tinygrad.runtime.autogen.amd.rdna3.ins import *

@@ -167,7 +167,7 @@ PREFETCH_LOADS = [(V_LDS_A_DATA[4+2*i], V_LDS_A_DATA[4+2*i+1], V_GLOBAL_B_ADDR,
 # =============================================================================

 class Kernel:
-  def __init__(self): self.instructions, self.labels, self.pos = [], {}, 0
+  def __init__(self, arch='gfx1100'): self.instructions, self.labels, self.pos, self.arch = [], {}, 0, arch
  def label(self, name): self.labels[name] = self.pos

  def emit(self, inst, target=None):
@@ -196,10 +196,10 @@ class Kernel:
 # Kernel builder
 # =============================================================================

-def build_kernel(N):
+def build_kernel(N, arch='gfx1100'):
  assert N % 128 == 0, f"N must be a multiple of 128 (tile size), got {N}"
  assert N >= 256, f"N must be >= 256 (prefetch pipeline requires at least 2 K-blocks), got {N}"
-  k = Kernel()
+  k = Kernel(arch)

  # ===========================================================================
  # PROLOGUE: Load kernel arguments, compute tile coordinates and addresses
@@ -443,7 +443,7 @@ def test_matmul():
  dev = Device[Device.DEFAULT]
  print(f"Device arch: {dev.renderer.target.arch}")

-  insts = build_kernel(N)
+  insts = build_kernel(N, dev.renderer.target.arch)

  rng = np.random.default_rng(42)
  a = Tensor(rng.random((N, N), dtype=np.float32) - 0.5)
@@ -463,14 +463,11 @@ def test_matmul():
                                                                                  estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
  c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
-  linear = c.schedule_linear()
+  ei = c.schedule()[0].lower()

  ets = []
  with Context(DEBUG=2):
-    for _ in range(getenv("CNT", 5)):
-      start = GlobalCounters.time_sum_s
-      run_linear(linear)
-      ets.append(GlobalCounters.time_sum_s - start)
+    for _ in range(getenv("CNT", 5)): ets.append(ei.run(wait=True))
  print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")

  if getenv("VERIFY", 1):
@@ -1,39 +1,31 @@
 # kernel8_batched_gmem.s from https://seb-v.github.io/optimization/update/2025/01/20/Fast-GPU-Matrix-multiplication.html
 # sudo PATH=/opt/homebrew/Cellar/llvm/20.1.6/bin:$PATH AMD_LLVM=0 AMD=1 DEBUG=2 python3 extra/gemm/amd_matmul.py
 import pathlib
+from dataclasses import replace
 from tinygrad import Tensor, Device, Context, GlobalCounters
 from tinygrad.helpers import getenv
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.renderer import Estimates
-from tinygrad.engine.realize import run_linear
+from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program

 N = 4096
 run_count = 5

-def make_matmul_kernel(name:str, src:str, local_size:int):
-  def fxn(a:UOp, b:UOp, c:UOp) -> UOp:
-    threads = UOp.special(local_size, "lidx0")
-    wg_x = UOp.special(N//128, "gidx0")
-    wg_y = UOp.special(N//128, "gidx1")
-    sink = UOp.sink(a.base, b.base, c.base, threads, wg_x, wg_y, arg=KernelInfo(name, estimates=Estimates(ops=2*N**3, mem=3*N*N*4)))
-    lib = Device[Device.DEFAULT].compiler.compile_cached(src)
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR, src=(*sink.src, sink)),
-                                 UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
-  return fxn
-
 if __name__ == "__main__":
+  ast = (Tensor.empty(N, N)@Tensor.empty(N, N)).schedule()[-1].ast
+  prg = get_program(ast, Device.default.renderer)
+
  if getenv("ASM") == 1:
    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel8_batched_gmem.s").read_text()
-    name, local_size = "kernel", 128
+    prgfast = replace(prg, name="kernel", src=src, global_size=[N//128, N//128, 1], local_size=[128, 1, 1])
  elif getenv("ASM") == -1:
    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel3_registers.cpp").read_text()
-    name, local_size = "kernel3_registers", 256
+    prgfast = replace(prg, name="kernel3_registers", src=src, global_size=[N//128, N//128, 1], local_size=[256, 1, 1])
  elif getenv("ASM") == -2:
    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel4_gmem_df.cpp").read_text()
-    name, local_size = "kernel4_gmem_db", 256
+    prgfast = replace(prg, name="kernel4_gmem_db", src=src, global_size=[N//128, N//128, 1], local_size=[256, 1, 1])
  else:
    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel5_lds_optim.cpp").read_text()
-    name, local_size = "kernel5_lds_optim", 128
+    prgfast = replace(prg, name="kernel5_lds_optim", src=src, global_size=[N//128, N//128, 1], local_size=[128, 1, 1])
+  runner = CompiledRunner(prgfast)

  a = Tensor.randn(N, N).realize()
  b = Tensor.randn(N, N).realize()
@@ -43,8 +35,8 @@ if __name__ == "__main__":
  with Context(DEBUG=2):
    for _ in range(run_count): tc = (a@b).realize()

-  linear = Tensor.custom_kernel(a, b, c, fxn=make_matmul_kernel(name, src, local_size))[2].schedule_linear()
  GlobalCounters.reset()
+  ei = ExecItem(ast, [a.uop.buffer, b.uop.buffer, c.uop.buffer], prg=runner)
  with Context(DEBUG=2):
-    for _ in range(run_count): run_linear(linear)
+    for _ in range(run_count): ei.run(wait=True)
  print(f"custom  {(c-tc).square().mean().item()}")
@@ -2628,24 +2628,20 @@ def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:
 # ** FP8 GEMM custom kernel

@functools.cache
-def custom_hk_fp8_gemm(C:UOp, A:UOp, B:UOp, *args:UOp, dname:str, scale_mode:int=3) -> UOp:
-  # scale_mode: 0=no scale, 1=x only, 2=w only, 3=both
-  n_scales = (1 if scale_mode & 1 else 0) + (1 if scale_mode & 2 else 0)
-  scales, extra = args[:n_scales], args[n_scales:]
+def custom_hk_fp8_gemm(C:UOp, A:UOp, B:UOp, X_s:UOp, W_s:UOp, dname:str) -> UOp:
+  # A is (batch, M, K), B is (N, K) transposed, X_s is x_scale, W_s is w_scale — kernel multiplies by both
  M, K = A.shape[0]*A.shape[1], A.shape[2]
  N, K2 = B.shape[(1 if B.ndim == 3 else 0):]
  assert K == K2, f"{A.shape} {B.shape}"
  block_size = 256
  threads = UOp.special(64 * 8, "lidx0")
  workgroups = UOp.special((M // block_size) * (N // block_size), "gidx0")
-  sink_inputs = (C.base, A.base, B.base) + tuple(s.base for s in scales) + (threads, workgroups)
-  sink = UOp.sink(*sink_inputs,
+  sink = UOp.sink(C.base, A.base, B.base, X_s.base, W_s.base, threads, workgroups,
                  arg=KernelInfo(f"hk_fp8_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K, mem=(M*K+N*K)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
  kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
  src = (kittens_path/"gemm_fp8.cpp").read_text()
  lib = HIPCCCompiler("gfx950", [f"-I{(kittens_path/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
-                                 "-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}",
-                                 f"-DSCALE_MODE={scale_mode}"]).compile_cached(src)
+                                 "-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DGEMM_M={M}", f"-DGEMM_N={N}", f"-DGEMM_K={K}"]).compile_cached(src)
  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
                               UOp(Ops.BINARY, arg=lib)))

@@ -2702,37 +2698,19 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:

 def custom_gemm_bw(gradient:UOp, kernel:UOp):
  inputs = kernel.src[1:]
-  if inputs[1].dtype == FP8_DTYPE:
-    grad_amax_state = inputs[5] if len(inputs) == 6 else None
-    out, a, b, s_x, s_w = inputs[:5]
+  # fp8 scaled gemm has 5 inputs (out, a, b, x_scale, w_scale), others have 3 (out, a, b)
+  if len(inputs) == 5:
+    out, a, b, s_x, s_w = inputs
    a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
    s_x_t, s_w_t = Tensor(s_x, device=a.device), Tensor(s_w, device=a.device)
    g_t = g_t[:a.shape[0]]
-    from extra.llama_kernels.cast_amax import _grad_fp8_mailbox
-    from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed
-    gbase = gradient.base if hasattr(gradient, "base") else gradient
-    mailbox_entry = _grad_fp8_mailbox.pop(gbase, None) or _grad_fp8_mailbox.pop(gradient, None)
-    if mailbox_entry is not None:
-      g_fp8_u, inv_scale_u, _new_amax_u, store_effect = mailbox_entry
-      g_fp8 = Tensor(g_fp8_u, device=a.device)[:a.shape[0]]
-      g_scale = Tensor(inv_scale_u, device=a.device)
-    else:
-      assert grad_amax_state is not None, "fp8 matmul bwd needs either a mailbox entry or a grad_amax_state"
-      g_fp8, g_scale, _, store_effect = quantize_fp8_delayed(g_t, Tensor(grad_amax_state, device=a.device))
+    g_fp8, g_scale, _ = quantize_fp8(g_t)
    # dgrad: uses g_scale * x_scale * w_scale
    grad_a = asm_gemm(g_fp8, b_t, x_scale=g_scale * s_x_t, w_scale=s_w_t)
    # wgrad: no w_scale
-    g_fp8_2d = g_fp8.reshape(-1, g_fp8.shape[-1])
-    if getenv("FAST_FP8_TRANSPOSE", 0) and g_fp8_2d.shape[0] % 64 == 0 and g_fp8_2d.shape[1] % 64 == 0:
-      from extra.llama_kernels.fp8_transpose import fast_fp8_transpose
-      g_fp8_T = fast_fp8_transpose(g_fp8_2d)
-    else:
-      g_fp8_T = g_fp8.permute(2, 0, 1).reshape(g_t.shape[-1], -1)
-    grad_b = asm_gemm(g_fp8_T, a_t.reshape(-1, a_t.shape[-1]), x_scale=g_scale * s_x_t)
-    # Attach the delayed-amax store effect (if any) to grad_a so realizing grads commits the amax update.
-    ret = (None, grad_a.uop.after(store_effect), grad_b.uop, None, None)
-    if len(inputs) == 6: ret = ret + (None,)
-    return ret
+    _one = Tensor(1.0, dtype=dtypes.float, device=a.device)
+    grad_b = asm_gemm(g_fp8.permute(2, 0, 1).reshape(g_t.shape[-1], -1), a_t.reshape(-1, a_t.shape[-1]), x_scale=g_scale * s_x_t, w_scale=_one)
+    return (None, grad_a.uop, grad_b.uop, None, None)
  else:
    out, a, b = inputs
    assert all_same([gradient.device, a.device, b.device, out.device])
@@ -2747,7 +2725,7 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp):

 # ** 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) -> Tensor:
+def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=None) -> Tensor:
  assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
  counters["used"] += 1
  unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0
@@ -2767,25 +2745,24 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N

  if is_multi:
    if n_sharded:
-      out = Tensor(Tensor.invalids(batch, M, N//len(a.device), dtype=out_dtype, device=a.device).uop.multi(2), device=a.device)
+      out = Tensor(Tensor.invalid(batch, M, N//len(a.device), dtype=out_dtype, device=a.device).uop.multi(2), device=a.device)
    elif m_sharded:
-      out = Tensor(Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device).uop.multi(1), device=a.device)
+      out = Tensor(Tensor.invalid(batch, M, N, dtype=out_dtype, device=a.device).uop.multi(1), device=a.device)
    else:
-      out = Tensor(Tensor.invalids(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=out_dtype, device=a.device).uop.multi(0),
+      out = Tensor(Tensor.invalid(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=out_dtype, device=a.device).uop.multi(0),
                   device=a.device)
  else:
-    out = Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device)
+    out = Tensor.invalid(batch, M, N, dtype=out_dtype, device=a.device)

  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):
    # fp8 gemm computes [email protected], kernel multiplies output by x_scale * w_scale before bf16 store
    if a.dtype == FP8_DTYPE:
-      scales = tuple(s for s in (x_scale, w_scale) if s is not None)
-      scale_mode = (1 if x_scale is not None else 0) | (2 if w_scale is not None else 0)
-      extra = [grad_amax_state] if grad_amax_state is not None else []
-      fxn = functools.partial(custom_hk_fp8_gemm, dname=dname, scale_mode=scale_mode)
-      out = Tensor.custom_kernel(out, a, b.T, *scales, *extra, fxn=fxn, grad_fxn=custom_gemm_bw)[0]
+      _one = lambda: Tensor(1.0, dtype=dtypes.float, device=a.device)
+      xs = x_scale if x_scale is not None else _one()
+      ws = w_scale if w_scale is not None else _one()
+      out = Tensor.custom_kernel(out, a, b.T, xs, ws, fxn=functools.partial(custom_hk_fp8_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
    else:
      out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
  else:
@@ -1,6 +1,7 @@
 import numpy as np, os
 from tinygrad.helpers import getenv, flat_mv
 from tinygrad import dtypes
+from tinygrad.engine.realize import get_program

 # for copied uops
 from tinygrad import dtypes
@@ -4,7 +4,7 @@ from tinygrad import Tensor, Device, Context, GlobalCounters
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.helpers import getenv, colored
 from tinygrad.dtype import dtypes, AddrSpace
-from tinygrad.engine.realize import Estimates, run_linear
+from tinygrad.engine.realize import Estimates
 from tinygrad.renderer.amd.dsl import s, v, VCC_LO, NULL, src, ttmp
 from tinygrad.runtime.autogen.amd.rdna4.ins import *

@@ -225,14 +225,11 @@ def test_matmul():
    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

  c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
-  linear = c.schedule_linear()
+  ei = c.schedule()[0].lower()

  ets = []
  with Context(DEBUG=2):
-    for _ in range(getenv("CNT", 5)):
-      start = GlobalCounters.time_sum_s
-      run_linear(linear)
-      ets.append(GlobalCounters.time_sum_s - start)
+    for _ in range(getenv("CNT", 5)): ets.append(ei.run(wait=True))
  print(f"REAL TFLOPS {N*N*N*2 / min(ets) * 1e-12:.2f}")

  if getenv("VERIFY", 1):
@@ -2,7 +2,6 @@ import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.helpers import getenv, get_single_element
 from tinygrad.dtype import _to_np_dtype
-from tinygrad.engine.realize import compile_linear
 from tinygrad.codegen.opt import OptOps

 dtype_in = (dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else
@@ -39,10 +38,10 @@ if __name__ == "__main__":
    c = a.matmul(b, dtype=acc_dtype).realize()

  if getenv("SHOULD_USE_TC"):
-    linear = compile_linear(a.matmul(b, dtype=acc_dtype).schedule_linear())
-    call = get_single_element(list(linear.src))
-    applied_opts = call.src[0].src[0].arg.applied_opts
-    assert any(opt.op is OptOps.TC for opt in applied_opts), f"TC not triggered, {applied_opts}"
+    sched = a.matmul(b, dtype=acc_dtype).schedule()
+    ei = get_single_element(sched)
+    ei.lower()
+    assert any(opt.op is OptOps.TC for opt in ei.prg.p.applied_opts), f"TC not triggered, {ei.prg.p.applied_opts}"

  ref = a.numpy().astype(np.float32) @ b.numpy().astype(np.float32)
  res = c.numpy()
@@ -1,7 +1,7 @@
-from tinygrad import Tensor, dtypes, Context
-from tinygrad.helpers import getenv
-from tinygrad.codegen.opt import Opt, OptOps
-from tinygrad.engine.realize import run_linear
+from tinygrad import Tensor, dtypes, Device
+from tinygrad.helpers import getenv, DEBUG
+from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
 from dataclasses import replace

 N = 4096
@@ -11,6 +11,9 @@ if __name__ == "__main__":
  else:
    A, B = Tensor.empty(N, N, dtype=dtypes.float16), Tensor.empty(N, N, dtype=dtypes.float16)
  C = A.matmul(B)
+  si = C.schedule()[-1]
+  ast = si.ast
+  k = Kernel(ast, opts=Device[Device.DEFAULT].renderer)
  if getenv("GEMV"):
    opts = [
      Opt(op=OptOps.UNROLL, axis=0, amt=8),
@@ -25,10 +28,10 @@ if __name__ == "__main__":
      Opt(op=OptOps.LOCAL, axis=1, amt=2),
      Opt(op=OptOps.LOCAL, axis=0, amt=2),
    ]
-  linear = C.schedule_linear()
-  call = linear.src[-1]
-  new_ast = call.src[0].replace(arg=replace(call.src[0].arg, opts_to_apply=tuple(opts)))
-  new_call = call.replace(src=(new_ast, *call.src[1:]))
-  linear = linear.replace(src=tuple(new_call if c is call else c for c in linear.src))
-  with Context(DEBUG=2):
-    for i in range(5): run_linear(linear)
+  k.apply_opts(opts)
+  prg = get_program(k.ast.replace(arg=replace(k.ast.arg, opts_to_apply=tuple(k.applied_opts))), k.opts)
+  new_src = prg.src
+  # can mod source here
+  prg = replace(prg, src=new_src)
+  ei = ExecItem(si.ast, [x.ensure_allocated() for x in si.bufs], si.metadata, prg=CompiledRunner(prg))
+  for i in range(5): ei.run(wait=True)
@@ -4,9 +4,7 @@ import triton.language as tl
 from triton.compiler import AttrsDescriptor, ASTSource, compile as triton_compile
 import numpy as np
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.engine.realize import get_runtime
-from tinygrad.codegen import to_program
-from tinygrad.uop.ops import Ops, UOp, KernelInfo, ProgramInfo
+from tinygrad.engine.realize import CompiledRunner, ExecItem, ProgramSpec
 from tinygrad.helpers import getenv
 np.set_printoptions(suppress=True)

@@ -75,11 +73,8 @@ if __name__ == "__main__":

  A, B = Tensor.normal(M, K, std=1e-1, dtype=dtypes.float16).realize(), Tensor.normal(K, N, std=1e-1, dtype=dtypes.float16).realize()
  C = A.matmul(B)
-  from tinygrad.uop.ops import Ops
-  linear, var_vals = C.linear_with_vars()
-  last_call = linear.src[-1]
-  ast = last_call.src[0]
-  bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]
+  sched = C.schedule()
+  si = sched[-1]

  src = compiled.asm["ptx"]
  # specify the shared memory here so we don't need to do it dynamically
@@ -90,27 +85,22 @@ if __name__ == "__main__":
  # remove debug sections
  src = src.split("\t.file")[0]
  assert '.extern .shared' not in src
-  info = ProgramInfo(name="matmul_kernel",
-                     global_size=(M//BLOCK_SIZE_M, N//BLOCK_SIZE_N, 1), local_size=(32*compiled.metadata.num_warps, 1, 1))
-  sink = UOp.sink(arg=KernelInfo(name="matmul_kernel"))
-  prg_uop = to_program(UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR), UOp(Ops.SOURCE, arg=src)), arg=info),
-                       Device.default.renderer)
-  rt = get_runtime(Device.DEFAULT, prg_uop)
-  all_bufs = [x.ensure_allocated() for x in bufs]
-  prg_bufs = [all_bufs[i] for i in info.globals]
-  gsize, lsize = info.launch_dims({})
+  prg = ProgramSpec("matmul_kernel", src, device=Device.DEFAULT,
+                global_size=[M//BLOCK_SIZE_M, N//BLOCK_SIZE_N, 1], local_size=[32*compiled.metadata.num_warps, 1, 1],
+                mem_estimate=A.nbytes() + B.nbytes() + C.nbytes())
+  ei = ExecItem(si.ast, [x.ensure_allocated() for x in si.bufs], si.metadata, prg=CompiledRunner(prg))
  tflops = []
  for i in range(5):
-    tm = rt(*[b._buf for b in prg_bufs], global_size=gsize, local_size=lsize, vals=info.vals({}), wait=True)
+    tm = ei.run(wait=True)
    tflops.append((2*M*K*N/tm)*1e-12)
  print(f"TFLOPS: {max(tflops):.2f}")

  # check correctness
  if getenv("VERIFY"):
-    from tinygrad.engine.realize import run_linear
+    from tinygrad.engine.realize import run_schedule
    triton_buf = np.frombuffer(si.bufs[0].as_memoryview(), np.float16).reshape(M,N)
    print(triton_buf)
-    run_linear(linear, var_vals)
+    run_schedule(sched)
    tinygrad_buf = np.frombuffer(si.bufs[0].as_memoryview(), np.float16).reshape(M,N)
    print(tinygrad_buf)
    np.testing.assert_allclose(triton_buf, tinygrad_buf)
@@ -36,10 +36,10 @@ A = Tensor.rand(M, K, device="CPU")
 B = Tensor.rand(K, N, device="CPU")
 C = (A.reshape(M, 1, K) * B.permute(1,0).reshape(1, N, K)).sum(axis=2)

-linear = C.schedule_linear()
+sched = C.schedule()
 from tinygrad.codegen.opt.kernel import Kernel
 from tinygrad.device import CompilerOptions
-lin = Kernel(linear.src[-1].src[0], CompilerOptions(has_local=False, supports_float4=False))
+lin = Kernel(sched[-1].ast, CompilerOptions(has_local=False, supports_float4=False))
 lin.to_program()
 from tinygrad.runtime.ops_cpu import renderer
 src = renderer("mmult", lin.uops)
@@ -1,51 +0,0 @@
-from __future__ import annotations
-import functools, pathlib
-from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import Ops
-from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
-
-FP8_MAX = 448.0
-NUM_WG, THREADS_PER_WG = 1024, 256
-
-# per-device abs max without allreduce
-@functools.cache
-def _local_abs_max_fxn(x_p, device):
-  x = Tensor(x_p, device=device)
-  inner = Tensor(x.uop.src[0]) if x.uop.op is Ops.MULTI else x
-  return (inner.abs().max(),)
-
-def local_abs_max(x:Tensor) -> Tensor:
-  param = x.as_param(0)
-  fxn = _local_abs_max_fxn(param.uop, x.device)
-  return Tensor(fxn[0].uop.call(x.uop).gettuple(0))
-
-def scalar_amax(amax_buf:Tensor) -> Tensor:
-  if isinstance(amax_buf.device, tuple):
-    return local_abs_max(amax_buf).detach()
-  return amax_buf.max().detach()
-
-def shard_shape(shape:tuple, axis:int, ndev:int) -> list:
-  s = list(shape)
-  s[axis] //= ndev
-  return s
-
-def dname_of(device) -> str:
-  if isinstance(device, tuple): return device[0].split(":")[0]
-  return device.split(":")[0] if isinstance(device, str) else device
-
-def alloc_like(shape, dtype, device, axis=None) -> Tensor:
-  if isinstance(device, tuple) and axis is not None:
-    return Tensor(Tensor.invalids(*shard_shape(shape, axis, len(device)), dtype=dtype, device=device).uop.multi(axis), device=device)
-  return Tensor.invalids(*shape, dtype=dtype, device=device)
-
-def alloc_local(shape, dtype, device, axis=None) -> Tensor:
-  if isinstance(device, tuple) and axis is not None:
-    return Tensor(Tensor.invalids(*shape, dtype=dtype, device=device).uop.multi(0), device=device)
-  return Tensor.invalids(*shape, dtype=dtype, device=device)
-
-def compile_hip(src:str, defines:list[str]):
-  return HIPCCCompiler("gfx950", ["-std=c++20", "-ffast-math", *defines]).compile_cached(src)
-
-def compile_cpp(cpp_dir:pathlib.Path, cpp_name:str, n_elems:int, hidden:int):
-  src = (cpp_dir/cpp_name).read_text()
-  return src, compile_hip(src, [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={hidden}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"])
@@ -1,74 +0,0 @@
-from __future__ import annotations
-import functools, pathlib
-from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.renderer import Estimates
-from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, compile_cpp, alloc_like, alloc_local, scalar_amax, dname_of
-
-# module-level mailbox: grad_xw13 UOp -> (grad_xw13_fp8 UOp, inv_scale UOp, new_amax UOp, store_effect)
-# lets cdna_asm_gemm's bwd reuse the fp8 companion produced by the fused silu_mul bwd kernel
-# instead of doing a redundant bf16 -> fp8 quantize.
-_grad_fp8_mailbox:dict = {}
-
-@functools.cache
-def _custom_fused_bwd_w13(grad_xw13:UOp, grad_xw13_fp8:UOp, grad_amax_buf:UOp,
-                          xw13:UOp, grad_x2:UOp, amax_state:UOp, grad_amax_state:UOp, dname:str) -> UOp:
-  hidden = xw13.shape[2] // 2
-  n_elems = xw13.shape[0] * xw13.shape[1] * hidden
-  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 * 5 + n_elems * 2 + NUM_WG * 4 + 4
-  sink = UOp.sink(grad_xw13.base, grad_xw13_fp8.base, grad_amax_buf.base,
-                  xw13.base, grad_x2.base, amax_state.base, grad_amax_state.base, threads, workgroups,
-                  arg=KernelInfo(f"fused_silu_mul_bwd_w13_{n_elems}", estimates=Estimates(ops=10*n_elems, mem=mem)))
-  src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_bwd_w13.cpp", n_elems, hidden)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
-                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
-
-@functools.cache
-def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_buf:UOp, xw13:UOp, amax_state:UOp, grad_amax_state:UOp, dname:str) -> UOp:
-  # NOTE: grad_amax_state is plumbed through as an unused fwd input so the bwd kernel can read it via kernel.src
-  hidden = xw13.shape[2] // 2
-  n_elems = xw13.shape[0] * xw13.shape[1] * hidden
-  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 * 2 + n_elems + NUM_WG * 4
-  sink = UOp.sink(fp8_out.base, amax_buf.base, xw13.base, amax_state.base, threads, workgroups,
-                  arg=KernelInfo(f"fused_silu_mul_cast_amax_w13_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=mem)))
-  src, lib = compile_cpp(pathlib.Path(__file__).parent, "cast_amax_fwd_w13.cpp", n_elems, hidden)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
-                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
-
-def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
-  _, _, xw13, amax_state, grad_amax_state = kernel.src[1:]
-  device = xw13.device
-  axis = xw13.axis if isinstance(device, tuple) else None
-  grad_xw13     = alloc_like(xw13.shape, dtypes.bfloat16, device, axis)
-  grad_xw13_fp8 = alloc_like(xw13.shape, dtypes.fp8e4m3,  device, axis)
-  grad_amax_buf = alloc_local((NUM_WG,), dtypes.float32,  device, axis)
-  grad_amax_state_t = Tensor(grad_amax_state, device=device)
-  fxn = functools.partial(_custom_fused_bwd_w13, dname=dname_of(device))
-  grad_xw13, grad_xw13_fp8, grad_amax_buf, *_ = Tensor.custom_kernel(
-    grad_xw13, grad_xw13_fp8, grad_amax_buf,
-    Tensor(xw13, device=device), Tensor(gradient, device=device).cast(dtypes.bfloat16),
-    Tensor(amax_state, device=device), grad_amax_state_t, fxn=fxn)
-  inv_scale = (grad_amax_state_t.float() + 1e-8) / FP8_MAX
-  new_grad_amax = scalar_amax(grad_amax_buf)
-  store_effect = grad_amax_state_t.uop.store(new_grad_amax.uop)
-  # Stash fp8 companion + amax store for cdna_asm_gemm's bwd to attach to grad_a.
-  _grad_fp8_mailbox[grad_xw13.uop] = (grad_xw13_fp8.uop, inv_scale.uop, new_grad_amax.uop, store_effect)
-  return (None, None, grad_xw13.uop, None, None)
-
-def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype, grad_amax_state:Tensor) -> tuple[Tensor, Tensor, Tensor]:
-  # NOTE: silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns (fp8, inv_scale, new_amax)
-  # grad_amax_state: delayed amax for grad_xw13 fp8 quantization in the backward.
-  assert xw13.dtype == dtypes.bfloat16, f"expected bf16, got {xw13.dtype}"
-  MBS, SEQ, H2 = xw13.shape
-  assert H2 % 2 == 0, f"w13 last-axis must be even, got {H2}"
-  HIDDEN = H2 // 2
-  axis = xw13.uop.axis if isinstance(xw13.device, tuple) else None
-  fp8_out  = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype,      xw13.device, axis)
-  amax_buf = alloc_local((NUM_WG,),         dtypes.float32, xw13.device, axis)
-  fxn = functools.partial(_custom_fused_cast_amax_w13, dname=dname_of(xw13.device))
-  fp8_out, amax_buf, *_ = Tensor.custom_kernel(fp8_out, amax_buf, xw13, amax_state, grad_amax_state,
-                                                fxn=fxn, grad_fxn=_fused_quantize_bwd_w13)
-  inv_scale = (amax_state.float() + 1e-8) / FP8_MAX
-  return fp8_out, inv_scale, scalar_amax(amax_buf)
@@ -1,98 +0,0 @@
-#include <hip/hip_runtime.h>
-#include <hip/hip_bf16.h>
-#include <hip/hip_fp8.h>
-
-#ifndef N_ELEMS
-#define N_ELEMS 234881024
-#endif
-#ifndef HIDDEN
-#define HIDDEN 14336
-#endif
-#ifndef NUM_WG
-#define NUM_WG 1024
-#endif
-#ifndef THREADS_PER_WG
-#define THREADS_PER_WG 256
-#endif
-
-constexpr int VEC = 8;
-constexpr float FP8_MAX = 448.0f;
-
-static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC");
-static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC");
-
-// fused silu*mul backward, three outputs in a single HBM pass:
-//   1) bf16 grad_xw13      — consumed by downstream bf16 autograd chain
-//   2) fp8  grad_xw13_fp8  — delayed-scale quantize using grad_amax_state (mailbox to matmul bwd)
-//   3) fp32 grad_amax_buf  — per-WG partial |grad_xw13|, reduced into next step's grad_amax_state
-// grad_amax_state is read for the fp8 scale. The store of new_grad_amax into grad_amax_state's
-// buffer is built in Python as a separate effect and threaded into grad_a via .after(store).
-extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
-fused_silu_mul_bwd_w13(
-    __hip_bfloat16*       __restrict__ grad_xw13_out,        // bf16, 2*N_ELEMS
-    __hip_fp8_storage_t*  __restrict__ grad_xw13_fp8_out,    // fp8,  2*N_ELEMS
-    float*                __restrict__ grad_amax_buf,        // fp32, NUM_WG per-WG partials
-    const __hip_bfloat16* __restrict__ xw13,                 // bf16, 2*N_ELEMS
-    const __hip_bfloat16* __restrict__ grad_x2,              // bf16, N_ELEMS
-    const float*          __restrict__ amax_state,           // fp32 scalar (fwd x2 amax)
-    const float*          __restrict__ grad_amax_state)      // fp32 scalar (delayed grad amax)
-{
-  __shared__ float sdata[THREADS_PER_WG];
-
-  const int tid = threadIdx.x;
-  const int wg  = blockIdx.x;
-  const int gid = wg * THREADS_PER_WG + tid;
-  const int stride_elems = NUM_WG * THREADS_PER_WG * VEC;
-
-  const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
-  const float g_scale = FP8_MAX / (static_cast<float>(*grad_amax_state) + 1e-8f);
-  float local_max = 0.0f;
-
-  for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) {
-    const int outer = base / HIDDEN;
-    const int inner = base % HIDDEN;
-    const int xw1_off = outer * 2 * HIDDEN + inner;
-    const int xw3_off = xw1_off + HIDDEN;
-
-    float4 x1_raw = *reinterpret_cast<const float4*>(&xw13[xw1_off]);
-    float4 x3_raw = *reinterpret_cast<const float4*>(&xw13[xw3_off]);
-    float4 g_raw  = *reinterpret_cast<const float4*>(&grad_x2[base]);
-
-    const __hip_bfloat16 *x1 = reinterpret_cast<const __hip_bfloat16*>(&x1_raw);
-    const __hip_bfloat16 *x3 = reinterpret_cast<const __hip_bfloat16*>(&x3_raw);
-    const __hip_bfloat16 *gv = reinterpret_cast<const __hip_bfloat16*>(&g_raw);
-
-    __hip_bfloat16 out1[VEC], out3[VEC];
-    __hip_fp8_storage_t fp8_1[VEC], fp8_3[VEC];
-    #pragma unroll
-    for (int i = 0; i < VEC; i++) {
-      const float f1 = static_cast<float>(x1[i]);
-      const float f3 = static_cast<float>(x3[i]);
-      const float fg = static_cast<float>(gv[i]);
-      const float sig = 1.0f / (1.0f + __expf(-f1));
-      const float silu = f1 * sig;
-      const float silu_prime = sig + silu * (1.0f - sig);
-      const float gs = fg * scale;
-      const float g1 = gs * silu_prime * f3;
-      const float g3 = gs * silu;
-      out1[i] = static_cast<__hip_bfloat16>(g1);
-      out3[i] = static_cast<__hip_bfloat16>(g3);
-      local_max = fmaxf(local_max, fmaxf(fabsf(g1), fabsf(g3)));
-      fp8_1[i] = __hip_cvt_float_to_fp8(fmaxf(-FP8_MAX, fminf(FP8_MAX, g1 * g_scale)), __HIP_SATFINITE, __HIP_E4M3);
-      fp8_3[i] = __hip_cvt_float_to_fp8(fmaxf(-FP8_MAX, fminf(FP8_MAX, g3 * g_scale)), __HIP_SATFINITE, __HIP_E4M3);
-    }
-
-    *reinterpret_cast<float4*>(&grad_xw13_out[xw1_off]) = *reinterpret_cast<float4*>(out1);
-    *reinterpret_cast<float4*>(&grad_xw13_out[xw3_off]) = *reinterpret_cast<float4*>(out3);
-    *reinterpret_cast<uint64_t*>(&grad_xw13_fp8_out[xw1_off]) = *reinterpret_cast<uint64_t*>(fp8_1);
-    *reinterpret_cast<uint64_t*>(&grad_xw13_fp8_out[xw3_off]) = *reinterpret_cast<uint64_t*>(fp8_3);
-  }
-
-  sdata[tid] = local_max;
-  __syncthreads();
-  for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
-    if (tid < s) sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]);
-    __syncthreads();
-  }
-  if (tid == 0) grad_amax_buf[wg] = sdata[0];
-}
@@ -1,41 +0,0 @@
-from __future__ import annotations
-import functools, pathlib
-from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.renderer import Estimates
-from extra.llama_kernels import THREADS_PER_WG, alloc_like, dname_of, compile_hip
-
-TILE = 64
-
-@functools.cache
-def _custom_fp8_transpose(out:UOp, inp:UOp, dname:str) -> UOp:
-  M, N = inp.shape
-  num_wg = (M // TILE) * (N // TILE)
-  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(num_wg, "gidx0")
-  mem = M * N * 2  # one byte read + one byte write per element
-  sink = UOp.sink(out.base, inp.base, threads, workgroups,
-                  arg=KernelInfo(f"fp8_transpose_{M}_{N}",
-                                 estimates=Estimates(ops=M*N, mem=mem)))
-  src = (pathlib.Path(__file__).parent/"fp8_transpose.cpp").read_text()
-  defines = [f"-DM_DIM={M}", f"-DN_DIM={N}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
-                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
-
-def fast_fp8_transpose(t:Tensor) -> Tensor:
-  assert t.ndim == 2, f"fast_fp8_transpose needs 2D input, got shape {t.shape}"
-  assert t.dtype in dtypes.fp8s, f"fast_fp8_transpose needs fp8 dtype, got {t.dtype}"
-  M, N = t.shape
-  assert M % TILE == 0 and N % TILE == 0, f"M={M}, N={N} must be multiples of {TILE}"
-
-  device = t.device
-  axis = t.uop.axis if isinstance(device, tuple) else None
-  out_axis = None
-  if axis == 0: out_axis = 1
-  elif axis == 1: out_axis = 0
-  elif axis is not None:
-    raise ValueError(f"fast_fp8_transpose: unsupported axis {axis}")
-
-  out = alloc_like((N, M), t.dtype, device, out_axis)
-  fxn = functools.partial(_custom_fp8_transpose, dname=dname_of(device))
-  out, _ = Tensor.custom_kernel(out, t, fxn=fxn)
-  return out
@@ -1,74 +0,0 @@
-#include <hip/hip_runtime.h>
-
-// LDS-staged 64x64 fp8 transpose.
-//   in : (M_DIM, N_DIM) fp8 contiguous
-//   out: (N_DIM, M_DIM) fp8 contiguous, out[c][r] = in[r][c]
-//
-// One WG processes one 64x64 output tile.  Each thread reads one uint4 (16 fp8) coalesced
-// from input rows, stages into LDS, then writes one uint4 coalesced to the output (whose
-// 16 fp8 come from 16 different input rows via in-LDS gather).
-//
-// LDS layout: lds[64][LDS_STRIDE] with LDS_STRIDE=65 (1 byte pad) to mitigate bank conflicts
-// during the column-direction read of the write phase.
-
-#ifndef M_DIM
-#define M_DIM 16384
-#endif
-#ifndef N_DIM
-#define N_DIM 28672
-#endif
-#ifndef THREADS_PER_WG
-#define THREADS_PER_WG 256
-#endif
-
-constexpr int TILE = 64;
-constexpr int VEC = 16;                    // fp8 per uint4 (128-bit) load/store
-constexpr int LDS_PAD = 1;
-constexpr int LDS_STRIDE = TILE + LDS_PAD; // 65 fp8 per row
-
-static_assert(THREADS_PER_WG * VEC == TILE * TILE, "256 threads * 16 fp8 = 64*64");
-static_assert(M_DIM % TILE == 0, "M_DIM must be a multiple of 64");
-static_assert(N_DIM % TILE == 0, "N_DIM must be a multiple of 64");
-
-constexpr int N_TILES_N = N_DIM / TILE;
-
-struct alignas(16) fp8x16 { uint8_t v[16]; };
-
-extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
-fp8_transpose(uint8_t* __restrict__ out,         // (N_DIM, M_DIM)
-              const uint8_t* __restrict__ in)    // (M_DIM, N_DIM)
-{
-  __shared__ uint8_t lds[TILE * LDS_STRIDE];
-
-  const int tid     = threadIdx.x;
-  const int wg_id   = blockIdx.x;
-  const int tile_r  = wg_id / N_TILES_N;       // tile index along M dim of input
-  const int tile_c  = wg_id % N_TILES_N;       // tile index along N dim of input
-
-  const int a = tid / (TILE / VEC);            // 0..63 (row within tile during read; col within tile during write)
-  const int b = tid % (TILE / VEC);            // 0..3
-  const int b16 = b * VEC;                     // 0,16,32,48
-
-  // ---- Read phase: input rows -> LDS rows
-  {
-    const long long src = (long long)(tile_r * TILE + a) * (long long)N_DIM
-                        + (long long)(tile_c * TILE + b16);
-    fp8x16 v = *reinterpret_cast<const fp8x16*>(&in[src]);
-    *reinterpret_cast<fp8x16*>(&lds[a * LDS_STRIDE + b16]) = v;
-  }
-  __syncthreads();
-
-  // ---- Write phase: LDS columns (gathered) -> output rows
-  // out[(tile_c*TILE + a)][(tile_r*TILE + b16 + i)] = in[(tile_r*TILE + b16 + i)][(tile_c*TILE + a)]
-  //                                                  = lds[b16 + i][a]
-  {
-    fp8x16 v;
-    #pragma unroll
-    for (int i = 0; i < VEC; ++i) {
-      v.v[i] = lds[(b16 + i) * LDS_STRIDE + a];
-    }
-    const long long dst = (long long)(tile_c * TILE + a) * (long long)M_DIM
-                        + (long long)(tile_r * TILE + b16);
-    *reinterpret_cast<fp8x16*>(&out[dst]) = v;
-  }
-}
@@ -1,96 +0,0 @@
-from __future__ import annotations
-import functools, pathlib
-from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.renderer import Estimates
-from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
-
-THREADS_PER_WG = 256
-
-@functools.cache
-def _custom_fused_ce_loss_fwd(loss_out:UOp, max_out:UOp, lse_out:UOp, logits:UOp, targets:UOp,
-                              dname:str, vocab:int, rows:int, label_smoothing:float) -> UOp:
-  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(rows, "gidx0")
-  mem = rows * vocab * 2 + rows * 12 + rows * 4
-  sink = UOp.sink(loss_out.base, max_out.base, lse_out.base, logits.base, targets.base,
-                  threads, workgroups,
-                  arg=KernelInfo(f"fused_ce_loss_fwd", estimates=Estimates(ops=6*rows*vocab, mem=mem)))
-  src = (pathlib.Path(__file__).parent/"fused_ce_loss.cpp").read_text()
-  defines = [f"-DVOCAB={vocab}", f"-DTHREADS_PER_WG={THREADS_PER_WG}",
-             f"-DLABEL_SMOOTHING={label_smoothing}f"]
-  lib = HIPCCCompiler("gfx950", ["-std=c++20", "-ffast-math", *defines]).compile_cached(src)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
-                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
-
-@functools.cache
-def _custom_fused_ce_loss_bwd(d_logits:UOp, logits:UOp, lse:UOp, targets:UOp, scale:UOp,
-                              dname:str, vocab:int, rows:int, label_smoothing:float) -> UOp:
-  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(rows, "gidx0")
-  mem = rows * vocab * 4 + rows * 8 + 4
-  sink = UOp.sink(d_logits.base, logits.base, lse.base, targets.base, scale.base,
-                  threads, workgroups,
-                  arg=KernelInfo(f"fused_ce_loss_bwd", estimates=Estimates(ops=4*rows*vocab, mem=mem)))
-  src = (pathlib.Path(__file__).parent/"fused_ce_loss_bwd.cpp").read_text()
-  defines = [f"-DVOCAB={vocab}", f"-DTHREADS_PER_WG={THREADS_PER_WG}",
-             f"-DLABEL_SMOOTHING={label_smoothing}f"]
-  lib = HIPCCCompiler("gfx950", ["-std=c++20", "-ffast-math", *defines]).compile_cached(src)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
-                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
-
-def _fused_ce_loss_bwd(gradient:UOp, kernel:UOp, label_smoothing:float):
-  # NOTE: forward inputs are (loss_out, max_out, lse_out, logits, targets)
-  # gradient is the upstream grad w.r.t. per-row loss (shape: (rows,) fp32)
-  _, _, lse_u, logits_u, targets_u = kernel.src[1:]
-  device = logits_u.device
-  rows, VOCAB = logits_u.shape  # (rows, VOCAB) after reshape
-  if isinstance(device, tuple):
-    axis = logits_u.axis
-    ndev = len(device)
-    d_logits = Tensor(Tensor.invalids(rows // ndev, VOCAB, dtype=dtypes.bfloat16, device=device).uop.multi(axis), device=device)
-    dname = device[0].split(":")[0]
-    rows_per_dev = rows // ndev
-  else:
-    d_logits = Tensor.invalids(rows, VOCAB, dtype=dtypes.bfloat16, device=device)
-    dname = device.split(":")[0] if isinstance(device, str) else device
-    rows_per_dev = rows
-  # NOTE: .mean() backward gives same grad per row (1/N), so broadcast is safe; take scalar
-  scale = Tensor(gradient, device=device).float().reshape(-1)[0:1].contiguous()
-  logits_t = Tensor(logits_u.after(kernel), device=device)
-  lse_t = Tensor(lse_u.after(kernel), device=device)
-  targets_t = Tensor(targets_u, device=device)
-  fxn = functools.partial(_custom_fused_ce_loss_bwd, dname=dname, vocab=VOCAB, rows=rows_per_dev, label_smoothing=label_smoothing)
-  d_logits, *_ = Tensor.custom_kernel(d_logits, logits_t, lse_t, targets_t, scale, fxn=fxn)
-  return (None, None, None, d_logits.uop, None)
-
-def fused_ce_loss(logits:Tensor, targets:Tensor, label_smoothing:float=0.1) -> Tensor:
-  # NOTE: fused sparse_categorical_crossentropy with label smoothing, returns mean loss scalar
-  assert logits.dtype == dtypes.bfloat16, f"expected bf16, got {logits.dtype}"
-  assert logits.ndim == 3, f"expected (MBS, SEQ, VOCAB), got {logits.shape}"
-  MBS, SEQ, VOCAB = logits.shape
-  rows = MBS * SEQ
-  if isinstance(logits.device, tuple):
-    axis = logits.uop.axis
-    assert axis in (0, 1), f"unsupported sharding axis={axis} for CE loss"
-    ndev = len(logits.device)
-    loss_out = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0),
-                      device=logits.device)
-    max_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0),
-                      device=logits.device)
-    lse_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0),
-                      device=logits.device)
-    dname = logits.device[0].split(":")[0]
-    rows_per_dev = rows // ndev
-  else:
-    loss_out = Tensor.invalids(rows, dtype=dtypes.float32, device=logits.device)
-    max_out  = Tensor.invalids(rows, dtype=dtypes.float32, device=logits.device)
-    lse_out  = Tensor.invalids(rows, dtype=dtypes.float32, device=logits.device)
-    dname = logits.device.split(":")[0] if isinstance(logits.device, str) else logits.device
-    rows_per_dev = rows
-  logits_flat = logits.reshape(rows, VOCAB)
-  targets_flat = targets.reshape(-1).cast(dtypes.int32)
-  fxn = functools.partial(_custom_fused_ce_loss_fwd, dname=dname, vocab=VOCAB, rows=rows_per_dev,
-                          label_smoothing=label_smoothing)
-  loss_out, max_out, lse_out, *_ = Tensor.custom_kernel(
-    loss_out, max_out, lse_out, logits_flat, targets_flat,
-    fxn=fxn, grad_fxn=functools.partial(_fused_ce_loss_bwd, label_smoothing=label_smoothing))
-  return loss_out.mean()
@@ -1,104 +0,0 @@
-#include <hip/hip_runtime.h>
-#include <hip/hip_bf16.h>
-
-// Fused forward sparse-CE with label smoothing.
-// SINGLE-PASS online softmax + vectorized 8-wide bf16 loads for HBM coalescing.
-
-#ifndef VOCAB
-#define VOCAB 128256
-#endif
-#ifndef THREADS_PER_WG
-#define THREADS_PER_WG 256
-#endif
-#ifndef LABEL_SMOOTHING
-#define LABEL_SMOOTHING 0.1f
-#endif
-
-constexpr int VEC = 8;
-
-extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
-fused_ce_loss_fwd(
-    float*                __restrict__ loss_out,  // out: fp32, ROWS
-    float*                __restrict__ max_out,   // out: fp32, ROWS
-    float*                __restrict__ lse_out,   // out: fp32, ROWS
-    const __hip_bfloat16* __restrict__ logits,    // in:  bf16, ROWS*VOCAB
-    const int*            __restrict__ targets)   // in:  int32, ROWS
-{
-  __shared__ float sdata_m[THREADS_PER_WG];
-  __shared__ float sdata_s[THREADS_PER_WG];
-  __shared__ float sdata_sumx[THREADS_PER_WG];
-  __shared__ float sdata_tgt[THREADS_PER_WG];
-
-  const int tid = threadIdx.x;
-  const int row = blockIdx.x;
-  const int target = targets[row];
-  const __hip_bfloat16* row_logits = logits + (size_t)row * VOCAB;
-
-  float m = -INFINITY;
-  float s = 0.0f;
-  float sum_x = 0.0f;
-  float target_logit = 0.0f;
-  constexpr bool needs_sum_x = (LABEL_SMOOTHING != 0.0f);
-
-  // Vectorized stride: each iter loads 8 bf16 = 16 bytes. Warp loads 32*16 = 512 bytes (4 cache lines).
-  const int VOCAB_VEC = VOCAB & ~(VEC - 1);  // round down to multiple of VEC
-  for (int i = tid * VEC; i < VOCAB_VEC; i += THREADS_PER_WG * VEC) {
-    float4 raw = *reinterpret_cast<const float4*>(&row_logits[i]);
-    const __hip_bfloat16* xi = reinterpret_cast<const __hip_bfloat16*>(&raw);
-    #pragma unroll
-    for (int k = 0; k < VEC; k++) {
-      const float x = static_cast<float>(xi[k]);
-      if constexpr (needs_sum_x) sum_x += x;
-      if (i + k == target) target_logit = x;
-      if (x > m) {
-        s = s * __expf(m - x) + 1.0f;
-        m = x;
-      } else {
-        s += __expf(x - m);
-      }
-    }
-  }
-  // tail (VOCAB not divisible by VEC):
-  for (int i = VOCAB_VEC + tid; i < VOCAB; i += THREADS_PER_WG) {
-    const float x = static_cast<float>(row_logits[i]);
-    if constexpr (needs_sum_x) sum_x += x;
-    if (i == target) target_logit = x;
-    if (x > m) { s = s * __expf(m - x) + 1.0f; m = x; }
-    else        { s += __expf(x - m); }
-  }
-
-  sdata_m[tid] = m;
-  sdata_s[tid] = s;
-  sdata_sumx[tid] = sum_x;
-  sdata_tgt[tid] = target_logit;
-  __syncthreads();
-
-  for (int step = THREADS_PER_WG / 2; step > 0; step >>= 1) {
-    if (tid < step) {
-      const float m1 = sdata_m[tid];
-      const float m2 = sdata_m[tid + step];
-      const float s1 = sdata_s[tid];
-      const float s2 = sdata_s[tid + step];
-      const float m_new = fmaxf(m1, m2);
-      const float s_new = s1 * __expf(m1 - m_new) + s2 * __expf(m2 - m_new);
-      sdata_m[tid] = m_new;
-      sdata_s[tid] = s_new;
-      sdata_sumx[tid] += sdata_sumx[tid + step];
-      sdata_tgt[tid]  += sdata_tgt[tid + step];
-    }
-    __syncthreads();
-  }
-
-  if (tid == 0) {
-    const float row_max = sdata_m[0];
-    const float row_sum_exp = sdata_s[0];
-    const float row_sum_x = sdata_sumx[0];
-    const float tgt = sdata_tgt[0];
-    const float row_lse = logf(row_sum_exp) + row_max;
-    const float mean_logits = row_sum_x / static_cast<float>(VOCAB);
-    const float loss = row_lse - (1.0f - LABEL_SMOOTHING) * tgt - LABEL_SMOOTHING * mean_logits;
-    loss_out[row] = loss;
-    max_out[row]  = row_max;
-    lse_out[row]  = row_lse;
-  }
-}
@@ -1,58 +0,0 @@
-#include <hip/hip_runtime.h>
-#include <hip/hip_bf16.h>
-
-// Vectorized CE bwd: 8-wide bf16 loads + stores.
-
-#ifndef VOCAB
-#define VOCAB 128256
-#endif
-#ifndef THREADS_PER_WG
-#define THREADS_PER_WG 256
-#endif
-#ifndef LABEL_SMOOTHING
-#define LABEL_SMOOTHING 0.1f
-#endif
-
-constexpr int VEC = 8;
-
-extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
-fused_ce_loss_bwd(
-    __hip_bfloat16*       __restrict__ d_logits,
-    const __hip_bfloat16* __restrict__ logits,
-    const float*          __restrict__ lse,
-    const int*            __restrict__ targets,
-    const float*          __restrict__ scale_in)
-{
-  const int tid = threadIdx.x;
-  const int row = blockIdx.x;
-  const int target = targets[row];
-  const float lse_r = lse[row];
-  const __hip_bfloat16* row_logits  = logits   + (size_t)row * VOCAB;
-  __hip_bfloat16*       row_dlogits = d_logits + (size_t)row * VOCAB;
-  const float inv_vocab = 1.0f / static_cast<float>(VOCAB);
-  const float scale = *scale_in;
-  const float ls_term = LABEL_SMOOTHING * inv_vocab;
-
-  const int VOCAB_VEC = VOCAB & ~(VEC - 1);
-  for (int i = tid * VEC; i < VOCAB_VEC; i += THREADS_PER_WG * VEC) {
-    float4 raw = *reinterpret_cast<const float4*>(&row_logits[i]);
-    const __hip_bfloat16* xi = reinterpret_cast<const __hip_bfloat16*>(&raw);
-    __hip_bfloat16 out[VEC];
-    #pragma unroll
-    for (int k = 0; k < VEC; k++) {
-      const float x = static_cast<float>(xi[k]);
-      float g = __expf(x - lse_r);
-      if (i + k == target) g -= (1.0f - LABEL_SMOOTHING);
-      g -= ls_term;
-      out[k] = static_cast<__hip_bfloat16>(g * scale);
-    }
-    *reinterpret_cast<float4*>(&row_dlogits[i]) = *reinterpret_cast<float4*>(out);
-  }
-  for (int i = VOCAB_VEC + tid; i < VOCAB; i += THREADS_PER_WG) {
-    const float x = static_cast<float>(row_logits[i]);
-    float g = __expf(x - lse_r);
-    if (i == target) g -= (1.0f - LABEL_SMOOTHING);
-    g -= ls_term;
-    row_dlogits[i] = static_cast<__hip_bfloat16>(g * scale);
-  }
-}
@@ -1,55 +0,0 @@
-from __future__ import annotations
-import functools, pathlib
-from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.renderer import Estimates
-from extra.llama_kernels import THREADS_PER_WG, dname_of, compile_hip
-
-ELEMS_PER_THREAD = 8  # vectorized 16-byte load (uint4 = 8 bf16)
-
-def _build_src(n_chunks:int) -> str:
-  template = (pathlib.Path(__file__).parent/"fused_pad_grad_accum.cpp").read_text()
-  params = "".join(f",\n    const __hip_bfloat16* __restrict__ chunk{i}" for i in range(n_chunks))
-  dispatch = "\n    ".join(f"case {i}: chunk_ptr = chunk{i}; break;" for i in range(n_chunks))
-  return (template.replace("__FUSED_PAD_GRAD_ACCUM_PARAMS", params)
-                  .replace("__FUSED_PAD_GRAD_ACCUM_DISPATCH", dispatch))
-
-@functools.cache
-def _custom_fused_pad_grad_accum(grad_buf:UOp, *chunk_uops, dname:str, n_chunks:int, chunk_size:int) -> UOp:
-  total = n_chunks * chunk_size
-  elems_per_block = THREADS_PER_WG * ELEMS_PER_THREAD
-  assert chunk_size % elems_per_block == 0, f"chunk_size {chunk_size} must be multiple of {elems_per_block}"
-  num_wg = total // elems_per_block
-  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(num_wg, "gidx0")
-  mem = total * 2 * 3
-  sink = UOp.sink(grad_buf.base, *(c.base for c in chunk_uops), threads, workgroups,
-                  arg=KernelInfo(f"fused_pad_grad_accum_n{n_chunks}_c{chunk_size}",
-                                 estimates=Estimates(ops=2*total, mem=mem)))
-  src = _build_src(n_chunks)
-  defines = [f"-DCHUNK_SIZE={chunk_size}", f"-DTHREADS_PER_WG={THREADS_PER_WG}", f"-DELEMS_PER_THREAD={ELEMS_PER_THREAD}"]
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
-                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
-
-def can_fused_pad_grad_accum(grad_buf:Tensor, chunks:list[Tensor]) -> bool:
-  if not chunks or grad_buf.dtype != dtypes.bfloat16: return False
-  if any(c.dtype != dtypes.bfloat16 for c in chunks): return False
-  chunk_shape = chunks[0].shape
-  if any(c.shape != chunk_shape for c in chunks): return False
-  chunk_size, total = 1, 1
-  for d in chunk_shape: chunk_size *= d
-  for d in grad_buf.shape: total *= d
-  return total == len(chunks) * chunk_size and chunk_size % (THREADS_PER_WG * ELEMS_PER_THREAD) == 0
-
-def fused_pad_grad_accum(grad_buf:Tensor, chunks:list[Tensor]) -> Tensor:
-  # NOTE: grad_buf += cat(*chunks, dim=0) in one HBM pass (in-place add). Returns new grad_buf Tensor.
-  # Requires uniform chunk shapes and chunk_size % (THREADS_PER_WG*ELEMS_PER_THREAD) == 0.
-  assert chunks and grad_buf.dtype == dtypes.bfloat16
-  for c in chunks: assert c.dtype == dtypes.bfloat16, f"chunk dtype must be bf16, got {c.dtype}"
-  chunk_size, total = 1, 1
-  for d in chunks[0].shape: chunk_size *= d
-  for d in grad_buf.shape: total *= d
-  assert total == len(chunks) * chunk_size, f"grad_buf size {total} != n_chunks {len(chunks)} * chunk_size {chunk_size}"
-  fxn = functools.partial(_custom_fused_pad_grad_accum, dname=dname_of(grad_buf.device),
-                          n_chunks=len(chunks), chunk_size=chunk_size)
-  out, *_ = Tensor.custom_kernel(grad_buf, *chunks, fxn=fxn)
-  return out
@@ -1,63 +0,0 @@
-// Fused custom kernel: grad_buf += cat(*chunks, dim=0) in one HBM pass.
-//
-// Template source — chunk parameter list and switch dispatch are filled by codegen
-// in cast_amax.py:_build_fused_pad_grad_accum_src to support arbitrary N.
-//
-// Defines required at compile time:
-//   CHUNK_SIZE       elements per chunk (must be multiple of THREADS_PER_WG * ELEMS_PER_THREAD)
-//   THREADS_PER_WG
-//   ELEMS_PER_THREAD (8 = one uint4 per thread = 16-byte vectorized load)
-//
-// Layout: one block-per-(slice-of-chunk) — blockIdx.x / BLOCKS_PER_CHUNK selects the chunk.
-// All threads in a block read the same chunk → switch is uniform → no warp divergence.
-
-#include <hip/hip_runtime.h>
-#include <hip/hip_bf16.h>
-
-#ifndef THREADS_PER_WG
-#define THREADS_PER_WG 256
-#endif
-#ifndef ELEMS_PER_THREAD
-#define ELEMS_PER_THREAD 8
-#endif
-
-#define ELEMS_PER_BLOCK (THREADS_PER_WG * ELEMS_PER_THREAD)
-#define BLOCKS_PER_CHUNK (CHUNK_SIZE / ELEMS_PER_BLOCK)
-
-extern "C" __attribute__((global))
-__attribute__((amdgpu_flat_work_group_size(1, THREADS_PER_WG)))
-void fused_pad_grad_accum(
-    __hip_bfloat16* __restrict__ grad_buf
-    __FUSED_PAD_GRAD_ACCUM_PARAMS
-) {
-  const int bid = blockIdx.x;
-  const int chunk_idx = bid / BLOCKS_PER_CHUNK;
-  const int block_in_chunk = bid - chunk_idx * BLOCKS_PER_CHUNK;
-  const int tid = threadIdx.x;
-
-  const __hip_bfloat16* chunk_ptr;
-  switch (chunk_idx) {
-    __FUSED_PAD_GRAD_ACCUM_DISPATCH
-    default: chunk_ptr = (const __hip_bfloat16*)0; break;  // unreachable
-  }
-
-  // int64 for global_offset: at 32 chunks × 117M elements = 3.6B, int32 overflows → MEMVIOL.
-  const int local_offset = block_in_chunk * ELEMS_PER_BLOCK + tid * ELEMS_PER_THREAD;
-  const long long global_offset = (long long)chunk_idx * (long long)CHUNK_SIZE + (long long)local_offset;
-
-  // Vectorized 16-byte load (uint4 = 8 bf16). Requires CHUNK_SIZE % 8 == 0 and 16-byte alignment.
-  const uint4 chunk_v = *reinterpret_cast<const uint4*>(&chunk_ptr[local_offset]);
-  const uint4 grad_v  = *reinterpret_cast<const uint4*>(&grad_buf[global_offset]);
-  uint4 out_v;
-
-  const __hip_bfloat16* chunk_bf = reinterpret_cast<const __hip_bfloat16*>(&chunk_v);
-  const __hip_bfloat16* grad_bf  = reinterpret_cast<const __hip_bfloat16*>(&grad_v);
-  __hip_bfloat16*       out_bf   = reinterpret_cast<__hip_bfloat16*>(&out_v);
-
-  #pragma unroll
-  for (int i = 0; i < ELEMS_PER_THREAD; i++) {
-    out_bf[i] = (__hip_bfloat16)((float)grad_bf[i] + (float)chunk_bf[i]);
-  }
-
-  *reinterpret_cast<uint4*>(&grad_buf[global_offset]) = out_v;
-}
@@ -1,153 +0,0 @@
-from __future__ import annotations
-import functools, pathlib
-from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.renderer import Estimates
-from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, scalar_amax, dname_of, compile_hip
-
-def _src() -> str: return (pathlib.Path(__file__).parent/"fused_rmsnorm_mul_quantize_fp8.cpp").read_text()
-def _src_bwd() -> str: return (pathlib.Path(__file__).parent/"fused_rmsnorm_mul_quantize_fp8_bwd.cpp").read_text()
-
-@functools.cache
-def _custom_fwd(fp8_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_buf:UOp,
-                x:UOp, weight:UOp, amax_state:UOp, dname:str, eps_val:float) -> UOp:
-  MBS, SEQ, HIDDEN = x.shape
-  n_elems = MBS * SEQ * HIDDEN
-  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 + n_elems + MBS * SEQ * 4 + n_elems + HIDDEN * 2 + NUM_WG * 4 + 4
-  sink = UOp.sink(fp8_out.base, x_normed_out.base, rrms_out.base, amax_buf.base,
-                  x.base, weight.base, amax_state.base, threads, workgroups,
-                  arg=KernelInfo(f"fused_rmsnorm_mul_quantize_fp8_{n_elems}_h{HIDDEN}_eps{eps_val:.0e}",
-                                 estimates=Estimates(ops=6*n_elems, mem=mem)))
-  defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={HIDDEN}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}",
-             f"-DEPS_LITERAL={eps_val}f"]
-  src = _src()
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
-                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
-
-@functools.cache
-def _custom_fwd_add(fp8_out:UOp, h_out:UOp, x_normed_out:UOp, rrms_out:UOp, amax_buf:UOp,
-                    x:UOp, residual:UOp, weight:UOp, amax_state:UOp, dname:str, eps_val:float) -> UOp:
-  MBS, SEQ, HIDDEN = x.shape
-  n_elems = MBS * SEQ * HIDDEN
-  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 * 4 + MBS * SEQ * 4 + HIDDEN * 2 + NUM_WG * 4 + 4
-  sink = UOp.sink(fp8_out.base, h_out.base, x_normed_out.base, rrms_out.base, amax_buf.base,
-                  x.base, residual.base, weight.base, amax_state.base, threads, workgroups,
-                  arg=KernelInfo(f"fused_add_rmsnorm_mul_quantize_fp8_{n_elems}_h{HIDDEN}_eps{eps_val:.0e}",
-                                 estimates=Estimates(ops=7*n_elems, mem=mem)))
-  defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={HIDDEN}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}",
-             f"-DEPS_LITERAL={eps_val}f", f"-DHAS_RESIDUAL=1"]
-  src = _src()
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
-                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
-
-@functools.cache
-def _custom_bwd(grad_x:UOp, grad_weight_partial:UOp,
-                grad_fp8:UOp, x_normed:UOp, rrms:UOp, weight:UOp, amax_state:UOp, dname:str) -> UOp:
-  MBS, SEQ, HIDDEN = x_normed.shape
-  n_elems = MBS * SEQ * HIDDEN
-  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 * 3 + NUM_WG * HIDDEN * 4 + MBS * SEQ * 4 + HIDDEN * 2 + 4
-  sink = UOp.sink(grad_x.base, grad_weight_partial.base,
-                  grad_fp8.base, x_normed.base, rrms.base, weight.base, amax_state.base, threads, workgroups,
-                  arg=KernelInfo(f"fused_rmsnorm_mul_quantize_fp8_bwd_{n_elems}_h{HIDDEN}",
-                                 estimates=Estimates(ops=8*n_elems, mem=mem)))
-  defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={HIDDEN}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
-  src = _src_bwd()
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
-                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
-
-def _bwd_common(fp8_grad_u, h_grad_u, x_u, x_normed_u, rrms_u, weight_u, amax_state_u, kernel:UOp):
-  device = x_u.device
-  MBS, SEQ, HIDDEN = x_normed_u.shape
-  axis = x_normed_u.axis if isinstance(device, tuple) else None
-  grad_x = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, device, axis)
-  grad_weight_partial = alloc_local((NUM_WG, HIDDEN), dtypes.float32, device, axis)
-  grad_h_from_fp8 = None
-  grad_weight_uop = None
-  if fp8_grad_u is not None:
-    fxn = functools.partial(_custom_bwd, dname=dname_of(device))
-    grad_x_t, grad_weight_partial_t, *_ = Tensor.custom_kernel(
-      grad_x, grad_weight_partial,
-      Tensor(fp8_grad_u, device=device).cast(dtypes.bfloat16),
-      Tensor(x_normed_u.after(kernel), device=device),
-      Tensor(rrms_u.after(kernel), device=device),
-      Tensor(weight_u, device=device),
-      Tensor(amax_state_u, device=device), fxn=fxn)
-    grad_h_from_fp8 = grad_x_t
-    grad_weight_uop = grad_weight_partial_t.sum(axis=0).cast(dtypes.bfloat16).uop
-  if h_grad_u is not None:
-    h_grad_t = Tensor(h_grad_u, device=device).cast(dtypes.bfloat16)
-    grad_total = (grad_h_from_fp8 + h_grad_t) if grad_h_from_fp8 is not None else h_grad_t
-  else:
-    grad_total = grad_h_from_fp8
-  return grad_total.uop, grad_weight_uop
-
-def _fused_bwd(gradient:UOp, kernel:UOp):
-  # NOTE: fwd inputs (fp8_out, x_normed_out, rrms_out, amax_buf, x, weight, amax_state)
-  _, x_normed_u, rrms_u, _, x_u, weight_u, amax_state_u = kernel.src[1:]
-  grad_x, grad_w = _bwd_common(gradient, None, x_u, x_normed_u, rrms_u, weight_u, amax_state_u, kernel)
-  return (None, None, None, None, grad_x, grad_w, None)
-
-def _fused_add_bwd(*args, **kwargs):
-  # Two invocation modes: 1 grad => positional; >1 grads => kwarg `call=`.
-  # Outputs: (fp8_out, h_out, x_normed_out, rrms_out, amax_buf). Both fp8 and h may be consumed
-  # downstream — TUPLE order in gradient.py preserves kernel-output slot order.
-  # Don't dispatch by dtype: matmul's bwd emits fp8 grad as bf16 (no explicit cast), so
-  # dtype-detection collapses both into h_grad and silently drops the rmsnorm-bwd path.
-  if 'call' in kwargs:
-    kernel, all_grads = kwargs['call'], list(args)
-  else:
-    gradient, kernel = args
-    all_grads = [gradient]
-  fp8_grad_u = h_grad_u = None
-  if len(all_grads) >= 2:
-    fp8_grad_u, h_grad_u = all_grads[0], all_grads[1]
-  elif len(all_grads) == 1:
-    g = all_grads[0]
-    if g.dtype == dtypes.bfloat16: h_grad_u = g
-    else: fp8_grad_u = g
-  _, _, x_normed_u, rrms_u, _, x_u, _, weight_u, amax_state_u = kernel.src[1:]
-  grad_h, grad_w = _bwd_common(fp8_grad_u, h_grad_u, x_u, x_normed_u, rrms_u, weight_u, amax_state_u, kernel)
-  return (None, None, None, None, None, grad_h, grad_h, grad_w, None)
-
-def fused_rmsnorm_mul_quantize_fp8(x:Tensor, weight:Tensor, amax_state:Tensor, eps:float, fp8_dtype) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
-  # NOTE: rmsnorm(x) * weight -> fp8 + amax. Returns (fp8, inv_scale, new_amax, x_normed, rrms).
-  # x_normed + rrms are saved for the rmsnorm backward (also recomputed here from x regs).
-  assert x.dtype == dtypes.bfloat16 and weight.dtype == dtypes.bfloat16
-  assert x.shape[-1] == weight.shape[-1], f"HIDDEN mismatch: x={x.shape}, weight={weight.shape}"
-  MBS, SEQ, HIDDEN = x.shape
-  axis = x.uop.axis if isinstance(x.device, tuple) else None
-  if isinstance(x.device, tuple): assert axis in (None, 0, 1), f"unsupported sharding axis={axis}"
-  fp8_out      = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype,       x.device, axis)
-  x_normed_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
-  rrms_out     = alloc_like((MBS, SEQ),         dtypes.float32,  x.device, axis)
-  amax_buf     = alloc_local((NUM_WG,),         dtypes.float32,  x.device, axis)
-  fxn = functools.partial(_custom_fwd, dname=dname_of(x.device), eps_val=eps)
-  fp8_out, x_normed_out, rrms_out, amax_buf, *_ = Tensor.custom_kernel(
-    fp8_out, x_normed_out, rrms_out, amax_buf, x, weight, amax_state, fxn=fxn, grad_fxn=_fused_bwd)
-  inv_scale = (amax_state.float() + 1e-8) / FP8_MAX
-  return fp8_out, inv_scale, scalar_amax(amax_buf), x_normed_out, rrms_out
-
-def fused_add_rmsnorm_mul_quantize_fp8(x:Tensor, residual:Tensor, weight:Tensor, amax_state:Tensor,
-                                       eps:float, fp8_dtype) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
-  # NOTE: h = x + residual; y_normed = rmsnorm(h); fp8 = quantize(y_normed * weight).
-  # Returns (fp8, inv_scale, new_amax, h, x_normed, rrms). h is also written so downstream can
-  # reuse it without recomputing x+residual — eliminates the separate residual-add kernel.
-  assert x.dtype == dtypes.bfloat16 and residual.dtype == dtypes.bfloat16 and weight.dtype == dtypes.bfloat16
-  assert x.shape == residual.shape
-  MBS, SEQ, HIDDEN = x.shape
-  axis = x.uop.axis if isinstance(x.device, tuple) else None
-  if isinstance(x.device, tuple): assert axis in (None, 0, 1), f"unsupported sharding axis={axis}"
-  fp8_out      = alloc_like((MBS, SEQ, HIDDEN), fp8_dtype,       x.device, axis)
-  h_out        = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
-  x_normed_out = alloc_like((MBS, SEQ, HIDDEN), dtypes.bfloat16, x.device, axis)
-  rrms_out     = alloc_like((MBS, SEQ),         dtypes.float32,  x.device, axis)
-  amax_buf     = alloc_local((NUM_WG,),         dtypes.float32,  x.device, axis)
-  fxn = functools.partial(_custom_fwd_add, dname=dname_of(x.device), eps_val=eps)
-  fp8_out, h_out, x_normed_out, rrms_out, amax_buf, *_ = Tensor.custom_kernel(
-    fp8_out, h_out, x_normed_out, rrms_out, amax_buf, x, residual, weight, amax_state,
-    fxn=fxn, grad_fxn=_fused_add_bwd)
-  inv_scale = (amax_state.float() + 1e-8) / FP8_MAX
-  return fp8_out, inv_scale, scalar_amax(amax_buf), h_out, x_normed_out, rrms_out
@@ -1,155 +0,0 @@
-#include <hip/hip_runtime.h>
-#include <hip/hip_bf16.h>
-#include <hip/hip_fp8.h>
-
-// Fuses the full pre-matmul preparation for a layer into a single HBM pass:
-//   y = rmsnorm(x) * weight            (reduce-mean-square + rsqrt + per-elem mul)
-//   fp8 = fp8_sat(y * (FP8_MAX / amax_state))
-// Also writes:
-//   rrms[row]        — saved for the rmsnorm backward
-//   amax_buf[wg]     — per-WG |y| partials, reduced later to update amax_state
-//
-// Layout: one WG per row, ROWS_PER_WG rows per WG via grid-stride (ROWS = N_ELEMS / HIDDEN).
-// Each thread handles HIDDEN / THREADS_PER_WG elements per row.
-
-#ifndef N_ELEMS
-#define N_ELEMS 67108864
-#endif
-#ifndef HIDDEN
-#define HIDDEN 4096
-#endif
-#ifndef NUM_WG
-#define NUM_WG 1024
-#endif
-#ifndef THREADS_PER_WG
-#define THREADS_PER_WG 256
-#endif
-#ifndef EPS_LITERAL
-#define EPS_LITERAL 1e-5f
-#endif
-#ifndef HAS_RESIDUAL
-#define HAS_RESIDUAL 0
-#endif
-
-constexpr int VEC = 8;
-constexpr float FP8_MAX = 448.0f;
-
-static_assert(N_ELEMS % HIDDEN == 0, "N_ELEMS must be a multiple of HIDDEN");
-static_assert(HIDDEN % (THREADS_PER_WG * VEC) == 0, "HIDDEN must be divisible by THREADS_PER_WG*VEC");
-
-constexpr int ROWS = N_ELEMS / HIDDEN;
-constexpr int ELEMS_PER_THREAD = HIDDEN / THREADS_PER_WG;  // each thread sees this many elems per row
-constexpr int VECS_PER_THREAD = ELEMS_PER_THREAD / VEC;    // number of 8-wide vec loads
-
-#if HAS_RESIDUAL
-extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
-fused_add_rmsnorm_mul_quantize_fp8(
-    __hip_fp8_storage_t*  __restrict__ fp8_out,         // fp8, ROWS*HIDDEN
-    __hip_bfloat16*       __restrict__ h_out,           // bf16, ROWS*HIDDEN — x + residual (saved for downstream)
-    __hip_bfloat16*       __restrict__ x_normed_out,    // bf16, ROWS*HIDDEN
-    float*                __restrict__ rrms_out,        // fp32, ROWS
-    float*                __restrict__ amax_buf,        // fp32, NUM_WG
-    const __hip_bfloat16* __restrict__ x,               // bf16, ROWS*HIDDEN
-    const __hip_bfloat16* __restrict__ residual,        // bf16, ROWS*HIDDEN — added into x before rmsnorm
-    const __hip_bfloat16* __restrict__ weight,          // bf16, HIDDEN
-    const float*          __restrict__ amax_state)      // fp32 scalar
-{
-#else
-extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
-fused_rmsnorm_mul_quantize_fp8(
-    __hip_fp8_storage_t*  __restrict__ fp8_out,         // fp8, ROWS*HIDDEN
-    __hip_bfloat16*       __restrict__ x_normed_out,    // bf16, ROWS*HIDDEN (saved for rmsnorm bwd)
-    float*                __restrict__ rrms_out,        // fp32, ROWS (fp32 to match rmsnorm_bwd.cpp expectation)
-    float*                __restrict__ amax_buf,        // fp32, NUM_WG per-WG partials
-    const __hip_bfloat16* __restrict__ x,               // bf16, ROWS*HIDDEN
-    const __hip_bfloat16* __restrict__ weight,          // bf16, HIDDEN (per-hidden scale)
-    const float*          __restrict__ amax_state)      // fp32 scalar
-{
-#endif
-  __shared__ float sdata[THREADS_PER_WG];
-
-  const int tid = threadIdx.x;
-  const int wg  = blockIdx.x;
-
-  const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
-  const float inv_hidden = 1.0f / static_cast<float>(HIDDEN);
-  float local_max = 0.0f;
-
-  // Grid-stride over rows. Each WG processes rows (wg, wg+NUM_WG, wg+2*NUM_WG, ...).
-  for (int row = wg; row < ROWS; row += NUM_WG) {
-    const int row_off = row * HIDDEN;
-
-    // Load row (+ residual if present) into registers.
-    float regs[ELEMS_PER_THREAD];
-    float sum_sq = 0.0f;
-    #pragma unroll
-    for (int v = 0; v < VECS_PER_THREAD; v++) {
-      const int h_base = tid * VEC + v * THREADS_PER_WG * VEC;
-      float4 raw = *reinterpret_cast<const float4*>(&x[row_off + h_base]);
-      const __hip_bfloat16 *xi = reinterpret_cast<const __hip_bfloat16*>(&raw);
-#if HAS_RESIDUAL
-      float4 res_raw = *reinterpret_cast<const float4*>(&residual[row_off + h_base]);
-      const __hip_bfloat16 *ri = reinterpret_cast<const __hip_bfloat16*>(&res_raw);
-      __hip_bfloat16 h_buf[VEC];
-#endif
-      #pragma unroll
-      for (int i = 0; i < VEC; i++) {
-#if HAS_RESIDUAL
-        const float f = static_cast<float>(xi[i]) + static_cast<float>(ri[i]);
-        h_buf[i] = static_cast<__hip_bfloat16>(f);
-#else
-        const float f = static_cast<float>(xi[i]);
-#endif
-        regs[v * VEC + i] = f;
-        sum_sq += f * f;
-      }
-#if HAS_RESIDUAL
-      *reinterpret_cast<float4*>(&h_out[row_off + h_base]) = *reinterpret_cast<float4*>(h_buf);
-#endif
-    }
-
-    // LDS tree-reduce sum_sq across the WG.
-    sdata[tid] = sum_sq;
-    __syncthreads();
-    for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
-      if (tid < s) sdata[tid] = sdata[tid] + sdata[tid + s];
-      __syncthreads();
-    }
-    const float mean_sq = sdata[0] * inv_hidden;
-    const float rrms = 1.0f / sqrtf(mean_sq + EPS_LITERAL);
-
-    if (tid == 0) rrms_out[row] = rrms;
-
-    // Normalize, multiply by weight, quantize. Also write x_normed (for rmsnorm bwd).
-    #pragma unroll
-    for (int v = 0; v < VECS_PER_THREAD; v++) {
-      const int h_base = tid * VEC + v * THREADS_PER_WG * VEC;
-      float4 w_raw = *reinterpret_cast<const float4*>(&weight[h_base]);
-      const __hip_bfloat16 *wi = reinterpret_cast<const __hip_bfloat16*>(&w_raw);
-
-      __hip_fp8_storage_t out[VEC];
-      __hip_bfloat16 xn[VEC];
-      #pragma unroll
-      for (int i = 0; i < VEC; i++) {
-        const float x_normed = regs[v * VEC + i] * rrms;
-        xn[i] = static_cast<__hip_bfloat16>(x_normed);
-        const float y = x_normed * static_cast<float>(wi[i]);
-        local_max = fmaxf(local_max, fabsf(y));
-        const float scaled = fmaxf(-FP8_MAX, fminf(FP8_MAX, y * scale));
-        out[i] = __hip_cvt_float_to_fp8(scaled, __HIP_SATFINITE, __HIP_E4M3);
-      }
-      *reinterpret_cast<uint64_t*>(&fp8_out[row_off + h_base]) = *reinterpret_cast<uint64_t*>(out);
-      *reinterpret_cast<float4*>(&x_normed_out[row_off + h_base]) = *reinterpret_cast<float4*>(xn);
-    }
-    __syncthreads();  // before next row's sum_sq reduce reuses sdata
-  }
-
-  // Final per-WG amax reduce.
-  sdata[tid] = local_max;
-  __syncthreads();
-  for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
-    if (tid < s) sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]);
-    __syncthreads();
-  }
-  if (tid == 0) amax_buf[wg] = sdata[0];
-}
@@ -1,147 +0,0 @@
-#include <hip/hip_runtime.h>
-#include <hip/hip_bf16.h>
-
-// Full backward for fused_rmsnorm_mul_quantize_fp8.cpp. One HBM pass per row produces:
-//   grad_x (bf16)              — gradient w.r.t. pre-rmsnorm x
-//   grad_weight_partial (fp32) — per-WG partial of the weight gradient, reduced later
-//
-// Input (all read):
-//   grad_fp8 (bf16)     — upstream grad w.r.t. fp8_out (bf16-typed gradient value)
-//   x_normed (bf16)     — saved from the fwd kernel, shape (ROWS, HIDDEN)
-//   rrms (fp32)         — saved rrms per row
-//   weight (bf16)       — per-HIDDEN rmsnorm weight
-//   amax_state (bf16)   — delayed amax used to compute the fp8 scale in fwd
-//
-// Chain: y = x_normed * weight; fp8 = sat(y * scale). Through STE: grad_y = grad_fp8 * scale.
-//        grad_x_normed = grad_y * weight.
-//        grad_weight   = sum_rows(grad_y * x_normed).
-//        grad_x        = rrms * (grad_x_normed - x_normed * mean(grad_x_normed * x_normed, last_dim)).
-
-#ifndef N_ELEMS
-#define N_ELEMS 67108864
-#endif
-#ifndef HIDDEN
-#define HIDDEN 4096
-#endif
-#ifndef NUM_WG
-#define NUM_WG 1024
-#endif
-#ifndef THREADS_PER_WG
-#define THREADS_PER_WG 256
-#endif
-
-constexpr int VEC = 8;
-constexpr float FP8_MAX = 448.0f;
-
-static_assert(N_ELEMS % HIDDEN == 0, "N_ELEMS must be a multiple of HIDDEN");
-static_assert(HIDDEN % (THREADS_PER_WG * VEC) == 0, "HIDDEN must be divisible by THREADS_PER_WG*VEC");
-
-constexpr int ROWS = N_ELEMS / HIDDEN;
-constexpr int ELEMS_PER_THREAD = HIDDEN / THREADS_PER_WG;
-constexpr int VECS_PER_THREAD = ELEMS_PER_THREAD / VEC;
-
-extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
-fused_rmsnorm_mul_quantize_fp8_bwd(
-    __hip_bfloat16*       __restrict__ grad_x,              // out: bf16, ROWS*HIDDEN
-    float*                __restrict__ grad_weight_partial, // out: fp32, NUM_WG*HIDDEN
-    const __hip_bfloat16* __restrict__ grad_fp8,            // in:  bf16, ROWS*HIDDEN (grad of fp8_out)
-    const __hip_bfloat16* __restrict__ x_normed,            // in:  bf16, ROWS*HIDDEN
-    const float*          __restrict__ rrms,                // in:  fp32, ROWS
-    const __hip_bfloat16* __restrict__ weight,              // in:  bf16, HIDDEN
-    const float*          __restrict__ amax_state)          // in:  fp32 scalar
-{
-  __shared__ float sdata[THREADS_PER_WG];
-
-  const int tid = threadIdx.x;
-  const int wg  = blockIdx.x;
-
-  const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
-  const float inv_hidden = 1.0f / static_cast<float>(HIDDEN);
-
-  // Per-thread accumulator for grad_weight (across all rows this WG touches).
-  float gw_accum[ELEMS_PER_THREAD];
-  #pragma unroll
-  for (int i = 0; i < ELEMS_PER_THREAD; i++) gw_accum[i] = 0.0f;
-
-  // Preload weight into registers (same across rows). Use ELEMS_PER_THREAD entries.
-  float w_regs[ELEMS_PER_THREAD];
-  #pragma unroll
-  for (int v = 0; v < VECS_PER_THREAD; v++) {
-    const int h_base = tid * VEC + v * THREADS_PER_WG * VEC;
-    float4 w_raw = *reinterpret_cast<const float4*>(&weight[h_base]);
-    const __hip_bfloat16 *wi = reinterpret_cast<const __hip_bfloat16*>(&w_raw);
-    #pragma unroll
-    for (int i = 0; i < VEC; i++) w_regs[v * VEC + i] = static_cast<float>(wi[i]);
-  }
-
-  for (int row = wg; row < ROWS; row += NUM_WG) {
-    const int row_off = row * HIDDEN;
-    const float rrms_v = rrms[row];
-
-    // Load grad_fp8 and x_normed rows into registers, compute grad_y and grad_x_normed.
-    float g_y_regs[ELEMS_PER_THREAD];
-    float xn_regs[ELEMS_PER_THREAD];
-    float g_xn_regs[ELEMS_PER_THREAD];  // grad_x_normed
-    float local_dot = 0.0f;  // sum(grad_x_normed * x_normed) for mean
-
-    #pragma unroll
-    for (int v = 0; v < VECS_PER_THREAD; v++) {
-      const int h_base = tid * VEC + v * THREADS_PER_WG * VEC;
-      float4 g_raw  = *reinterpret_cast<const float4*>(&grad_fp8[row_off + h_base]);
-      float4 xn_raw = *reinterpret_cast<const float4*>(&x_normed[row_off + h_base]);
-      const __hip_bfloat16 *gi  = reinterpret_cast<const __hip_bfloat16*>(&g_raw);
-      const __hip_bfloat16 *xni = reinterpret_cast<const __hip_bfloat16*>(&xn_raw);
-
-      #pragma unroll
-      for (int i = 0; i < VEC; i++) {
-        const int idx = v * VEC + i;
-        const float g_y = static_cast<float>(gi[i]) * scale;
-        const float xn  = static_cast<float>(xni[i]);
-        g_y_regs[idx] = g_y;
-        xn_regs[idx]  = xn;
-        g_xn_regs[idx] = g_y * w_regs[idx];           // grad_x_normed = grad_y * weight
-        gw_accum[idx] += g_y * xn;                    // grad_weight contrib
-        local_dot += g_xn_regs[idx] * xn;             // for mean
-      }
-    }
-
-    // LDS reduce local_dot to sdata[0].
-    sdata[tid] = local_dot;
-    __syncthreads();
-    for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
-      if (tid < s) sdata[tid] = sdata[tid] + sdata[tid + s];
-      __syncthreads();
-    }
-    const float mean_term = sdata[0] * inv_hidden;
-
-    // Compute grad_x = rrms * (grad_x_normed - x_normed * mean_term) and write.
-    #pragma unroll
-    for (int v = 0; v < VECS_PER_THREAD; v++) {
-      const int h_base = tid * VEC + v * THREADS_PER_WG * VEC;
-      __hip_bfloat16 out[VEC];
-      #pragma unroll
-      for (int i = 0; i < VEC; i++) {
-        const int idx = v * VEC + i;
-        const float dx = rrms_v * (g_xn_regs[idx] - xn_regs[idx] * mean_term);
-        out[i] = static_cast<__hip_bfloat16>(dx);
-      }
-      *reinterpret_cast<float4*>(&grad_x[row_off + h_base]) = *reinterpret_cast<float4*>(out);
-    }
-    __syncthreads();
-  }
-
-  // Write this WG's grad_weight partial to HBM (fp32, NUM_WG x HIDDEN layout).
-  const int gw_row_off = wg * HIDDEN;
-  #pragma unroll
-  for (int v = 0; v < VECS_PER_THREAD; v++) {
-    const int h_base = tid * VEC + v * THREADS_PER_WG * VEC;
-    // Write 8 fp32 values with two float4 stores.
-    float4 out_lo, out_hi;
-    out_lo.x = gw_accum[v * VEC + 0]; out_lo.y = gw_accum[v * VEC + 1];
-    out_lo.z = gw_accum[v * VEC + 2]; out_lo.w = gw_accum[v * VEC + 3];
-    out_hi.x = gw_accum[v * VEC + 4]; out_hi.y = gw_accum[v * VEC + 5];
-    out_hi.z = gw_accum[v * VEC + 6]; out_hi.w = gw_accum[v * VEC + 7];
-    *reinterpret_cast<float4*>(&grad_weight_partial[gw_row_off + h_base + 0]) = out_lo;
-    *reinterpret_cast<float4*>(&grad_weight_partial[gw_row_off + h_base + 4]) = out_hi;
-  }
-}
@@ -1,67 +0,0 @@
-from __future__ import annotations
-import functools, pathlib
-from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.renderer import Estimates
-from extra.llama_kernels import FP8_MAX, NUM_WG, THREADS_PER_WG, alloc_like, alloc_local, scalar_amax, dname_of, compile_hip
-
-@functools.cache
-def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_partial:UOp, x:UOp, amax_state:UOp, dname:str) -> UOp:
-  n_elems = 1
-  for d in x.shape: n_elems *= d
-  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 + n_elems + 4 + NUM_WG * 4
-  sink = UOp.sink(fp8_out.base, amax_partial.base, x.base, amax_state.base, threads, workgroups,
-                  arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", estimates=Estimates(ops=3*n_elems, mem=mem)))
-  src = (pathlib.Path(__file__).parent/"quantize_fp8_with_amax.cpp").read_text()
-  defines = [f"-DN_ELEMS={n_elems}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
-                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
-
-@functools.cache
-def _custom_quantize_fp8_scalar(fp8_out:UOp, x:UOp, amax_state:UOp, dname:str) -> UOp:
-  n_elems = 1
-  for d in x.shape: n_elems *= d
-  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  mem = n_elems * 2 + n_elems
-  sink = UOp.sink(fp8_out.base, x.base, amax_state.base, threads, workgroups,
-                  arg=KernelInfo(f"quantize_fp8_scalar_{n_elems}", estimates=Estimates(ops=2*n_elems, mem=mem)))
-  src = (pathlib.Path(__file__).parent/"quantize_fp8_scalar.cpp").read_text()
-  defines = [f"-DN_ELEMS={n_elems}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
-                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
-
-def _quantize_fp8_delayed_bwd(gradient:UOp, kernel:UOp):
-  # NOTE: STE-equivalent backward — grad_x = grad_fp8 * scale, scale = FP8_MAX / amax_state.
-  # `gradient` is bf16 grad w.r.t. fp8 output (asm_gemm bwd already applied x_scale).
-  _, _, x, amax_state = kernel.src[1:]
-  device = x.device
-  scale = FP8_MAX / (Tensor(amax_state, device=device).float() + 1e-8)
-  grad_x = (Tensor(gradient, device=device).float() * scale).cast(dtypes.bfloat16)
-  return (None, None, grad_x.uop, None)
-
-def quantize_fp8_delayed(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3) -> tuple[Tensor, Tensor, Tensor, UOp]:
-  # NOTE: one-pass bf16 -> fp8 quantize with delayed scaling. Returns (fp8, inv_scale, new_amax, store_effect).
-  # Fused kernel reads x once and writes fp8 + per-WG |x| partials (then a small reduce produces scalar new_amax).
-  # store_effect writes new_amax into amax_state's buffer — the caller must thread it into a realized
-  # output via `.after(store_effect)`. Calling `amax_state.assign(new_amax)` inside a grad_fxn does
-  # NOT work because .assign mutates only the temp Tensor's .uop, not the original layer-owned buffer.
-  assert x.dtype == dtypes.bfloat16, f"expected bf16, got {x.dtype}"
-  axis = x.uop.axis if isinstance(x.device, tuple) else None
-  fp8_out      = alloc_like(x.shape,  fp8_dtype,      x.device, axis)
-  amax_partial = alloc_local((NUM_WG,), dtypes.float32, x.device, axis)
-  fxn = functools.partial(_custom_quantize_fp8_with_amax, dname=dname_of(x.device))
-  fp8_out, amax_partial, *_ = Tensor.custom_kernel(fp8_out, amax_partial, x, amax_state,
-                                                    fxn=fxn, grad_fxn=_quantize_fp8_delayed_bwd)
-  new_amax = scalar_amax(amax_partial)
-  inv_scale = (amax_state.float() + 1e-8) / FP8_MAX
-  store_effect = amax_state.uop.store(new_amax.uop)
-  return fp8_out, inv_scale, new_amax, store_effect
-
-def quantize_fp8_scalar(x:Tensor, amax_state:Tensor, fp8_dtype=dtypes.fp8e4m3) -> Tensor:
-  # NOTE: pure one-pass bf16 -> fp8 quantize with delayed scalar scale. No amax computation.
-  axis = x.uop.axis if isinstance(x.device, tuple) else None
-  fp8_out = alloc_like(x.shape, fp8_dtype, x.device, axis)
-  fxn = functools.partial(_custom_quantize_fp8_scalar, dname=dname_of(x.device))
-  fp8_out, *_ = Tensor.custom_kernel(fp8_out, x, amax_state, fxn=fxn)
-  return fp8_out
@@ -1,48 +0,0 @@
-#include <hip/hip_runtime.h>
-#include <hip/hip_bf16.h>
-#include <hip/hip_fp8.h>
-
-// Pure one-pass bf16 -> fp8 quantize with delayed scalar scale. No amax computation.
-
-#ifndef N_ELEMS
-#define N_ELEMS 67108864
-#endif
-#ifndef NUM_WG
-#define NUM_WG 1024
-#endif
-#ifndef THREADS_PER_WG
-#define THREADS_PER_WG 256
-#endif
-
-constexpr int VEC = 8;
-constexpr float FP8_MAX = 448.0f;
-
-static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC");
-
-extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
-quantize_fp8_scalar(
-    __hip_fp8_storage_t*  __restrict__ fp8_out,         // fp8, N_ELEMS
-    const __hip_bfloat16* __restrict__ x,               // bf16, N_ELEMS
-    const float*          __restrict__ amax_state)      // fp32 scalar (delayed)
-{
-  const int tid = threadIdx.x;
-  const int wg  = blockIdx.x;
-  const int gid = wg * THREADS_PER_WG + tid;
-  const int stride_elems = NUM_WG * THREADS_PER_WG * VEC;
-
-  const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
-
-  for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) {
-    float4 x_raw = *reinterpret_cast<const float4*>(&x[base]);
-    const __hip_bfloat16 *xi = reinterpret_cast<const __hip_bfloat16*>(&x_raw);
-
-    __hip_fp8_storage_t out[VEC];
-    #pragma unroll
-    for (int i = 0; i < VEC; i++) {
-      const float v = static_cast<float>(xi[i]);
-      const float scaled = fmaxf(-FP8_MAX, fminf(FP8_MAX, v * scale));
-      out[i] = __hip_cvt_float_to_fp8(scaled, __HIP_SATFINITE, __HIP_E4M3);
-    }
-    *reinterpret_cast<uint64_t*>(&fp8_out[base]) = *reinterpret_cast<uint64_t*>(out);
-  }
-}
@@ -1,63 +0,0 @@
-#include <hip/hip_runtime.h>
-#include <hip/hip_bf16.h>
-#include <hip/hip_fp8.h>
-
-// One-pass bf16 -> fp8 quantize using a scalar delayed amax state,
-// AND simultaneously computes per-WG |x| max partials for the next step's amax state.
-// Saves one full HBM pass over the grad tensor vs. doing quantize + separate abs().max().
-
-#ifndef N_ELEMS
-#define N_ELEMS 67108864
-#endif
-#ifndef NUM_WG
-#define NUM_WG 1024
-#endif
-#ifndef THREADS_PER_WG
-#define THREADS_PER_WG 256
-#endif
-
-constexpr int VEC = 8;
-constexpr float FP8_MAX = 448.0f;
-
-static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC");
-
-extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
-quantize_fp8_with_amax(
-    __hip_fp8_storage_t*  __restrict__ fp8_out,       // out: fp8, N_ELEMS
-    float*                __restrict__ amax_partial,  // out: fp32, NUM_WG per-WG partials
-    const __hip_bfloat16* __restrict__ x,             // in:  bf16, N_ELEMS
-    const float*          __restrict__ amax_state)    // in:  fp32 scalar (delayed)
-{
-  __shared__ float sdata[THREADS_PER_WG];
-
-  const int tid = threadIdx.x;
-  const int wg  = blockIdx.x;
-  const int gid = wg * THREADS_PER_WG + tid;
-  const int stride_elems = NUM_WG * THREADS_PER_WG * VEC;
-
-  const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
-  float local_max = 0.0f;
-
-  for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) {
-    float4 x_raw = *reinterpret_cast<const float4*>(&x[base]);
-    const __hip_bfloat16 *xi = reinterpret_cast<const __hip_bfloat16*>(&x_raw);
-
-    __hip_fp8_storage_t out[VEC];
-    #pragma unroll
-    for (int i = 0; i < VEC; i++) {
-      const float v = static_cast<float>(xi[i]);
-      local_max = fmaxf(local_max, fabsf(v));
-      const float scaled = fmaxf(-FP8_MAX, fminf(FP8_MAX, v * scale));
-      out[i] = __hip_cvt_float_to_fp8(scaled, __HIP_SATFINITE, __HIP_E4M3);
-    }
-    *reinterpret_cast<uint64_t*>(&fp8_out[base]) = *reinterpret_cast<uint64_t*>(out);
-  }
-
-  sdata[tid] = local_max;
-  __syncthreads();
-  for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
-    if (tid < s) sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]);
-    __syncthreads();
-  }
-  if (tid == 0) amax_partial[wg] = sdata[0];
-}
@@ -1,24 +0,0 @@
-from __future__ import annotations
-import functools
-from tinygrad import Tensor
-from tinygrad.uop.ops import UOp
-
-def rmsnorm_fwd(x_in:Tensor, eps:float) -> tuple[Tensor, Tensor]:
-  x = x_in.float()
-  rrms = (x.square().mean(-1, keepdim=True) + eps).rsqrt()
-  return (x * rrms).cast(x_in.dtype), rrms
-
-@functools.cache
-def _rmsnorm_fwd_fxn(x_in_p, eps, device):
-  return rmsnorm_fwd(Tensor(x_in_p, device=device), eps)
-
-def _rmsnorm_bwd(grad:UOp, call:UOp) -> tuple:
-  x_normed = Tensor(call.gettuple(0)).float()
-  do_float = Tensor(grad).float()
-  d_x = Tensor(call.gettuple(1)) * (do_float - x_normed * (do_float * x_normed).mean(-1, keepdim=True))
-  return (d_x.cast(call.src[1].dtype).uop,)
-
-def rmsnorm(x_in:Tensor, eps:float) -> tuple[Tensor, Tensor]:
-  fxn = _rmsnorm_fwd_fxn(x_in.as_param(0).uop, eps, x_in.device)
-  call = UOp.maketuple(fxn[0].uop, fxn[1].uop).call(x_in.uop, grad_fxn=_rmsnorm_bwd)
-  return Tensor(call.gettuple(0)), Tensor(call.gettuple(1))
@@ -3,12 +3,11 @@ import os
 # TODO: there is a timing bug without this
 os.environ["AMD_AQL"] = "1"

-from tinygrad import Tensor, Device, GlobalCounters, Context
+from tinygrad import Tensor, Device
 from tinygrad.helpers import getenv, DEV
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
 from tinygrad.renderer.amd.dsl import Reg, Inst, s, v
-from tinygrad.engine.realize import run_linear

 NUM_WORKGROUPS = 96
 WAVE_SIZE = 32
@@ -37,17 +36,11 @@ def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, **kwargs)
    gidx = UOp.special(NUM_WORKGROUPS, "gidx0")
    FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
    sink = UOp.sink(A.base, threads, gidx, arg=KernelInfo(inst.op.name.lower(), estimates=Estimates(ops=FLOPs, mem=0)))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
  dummy = Tensor.zeros(1).contiguous().realize()
  out = Tensor.custom_kernel(dummy, fxn=fxn)[0]
-  linear = out.schedule_linear()
-  ets = []
-  with Context(DEBUG=2):
-    for _ in range(2):
-      start = GlobalCounters.time_sum_s
-      run_linear(linear)
-      ets.append(GlobalCounters.time_sum_s - start)
-  elapsed = min(ets)
+  ei = out.schedule()[-1].lower()
+  elapsed = min([ei.run(wait=True) for _ in range(2)])
  FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
  print(f"{inst.op_name.lower():<29} : {FLOPs/elapsed/10**12:.2f} T(FL)OPS")

@@ -1,25 +0,0 @@
-/* adapted from linux/drivers/gpu/drm/nouveau/include/nvfw/fw.h */
-/* SPDX-License-Identifier: MIT */
-#ifndef __NVFW_FW_H__
-#define __NVFW_FW_H__
-typedef unsigned int u32;
-
-struct nvfw_bin_hdr {
-	u32 bin_magic;
-	u32 bin_ver;
-	u32 bin_size;
-	u32 header_offset;
-	u32 data_offset;
-	u32 data_size;
-};
-
-struct nvfw_bl_desc {
-	u32 start_tag;
-	u32 dmem_load_off;
-	u32 code_off;
-	u32 code_size;
-	u32 data_off;
-	u32 data_size;
-};
-
-#endif
@@ -1,52 +0,0 @@
-/* adapted from linux/drivers/gpu/drm/nouveau/include/nvfw/hs.h */
-/* SPDX-License-Identifier: MIT */
-#ifndef __NVFW_HS_H__
-#define __NVFW_HS_H__
-typedef unsigned int u32;
-
-struct nvfw_hs_header {
-	u32 sig_dbg_offset;
-	u32 sig_dbg_size;
-	u32 sig_prod_offset;
-	u32 sig_prod_size;
-	u32 patch_loc;
-	u32 patch_sig;
-	u32 hdr_offset;
-	u32 hdr_size;
-};
-
-struct nvfw_hs_header_v2 {
-	u32 sig_prod_offset;
-	u32 sig_prod_size;
-	u32 patch_loc;
-	u32 patch_sig;
-	u32 meta_data_offset;
-	u32 meta_data_size;
-	u32 num_sig;
-	u32 header_offset;
-	u32 header_size;
-};
-
-struct nvfw_hs_load_header {
-	u32 non_sec_code_off;
-	u32 non_sec_code_size;
-	u32 data_dma_base;
-	u32 data_size;
-	u32 num_apps;
-	u32 apps[];
-};
-
-struct nvfw_hs_load_header_v2 {
-	u32 os_code_offset;
-	u32 os_code_size;
-	u32 os_data_offset;
-	u32 os_data_size;
-	u32 num_apps;
-	struct {
-		u32 offset;
-		u32 size;
-		u32 data_offset;
-		u32 data_size;
-	} app[];
-};
-#endif
@@ -10,4 +10,4 @@ def extract_ast(*args) -> None:
  return None

 if __name__ == "__main__":
-  _pmap({"do_to_program":extract_ast})
+  _pmap({"get_program":extract_ast})
@@ -84,7 +84,8 @@ class TestBeamSearch(unittest.TestCase):
    tc = Device[Device.DEFAULT].renderer.tensor_cores[0]
    size = max(tc.dims[0], tc.dims[1]) * 8
    a, b = Tensor.rand(size, size, dtype=tc.dtype_in), Tensor.rand(size, size, dtype=tc.dtype_in)
-    ast = a.matmul(b, dtype=tc.dtype_out).schedule_linear().src[-1].src[0]
+    ast = a.matmul(b, dtype=tc.dtype_out).schedule()[-1].ast
+    if ast.op is Ops.BEAM: ast = ast.src[0]
    s = Scheduler(ast, Device[Device.DEFAULT].renderer)
    s.apply_opt(Opt(OptOps.TC, 0, (-1, 0, 1)))
    up = prod([x for x, t in zip(s.full_shape, s.axis_types) if t in (AxisType.UPCAST, AxisType.UNROLL)])
@@ -94,7 +95,8 @@ class TestBeamSearch(unittest.TestCase):

  def test_max_up(self):
    a = Tensor.rand(16, 16)
-    ast = a.schedule_linear().src[-1].src[0]
+    ast = a.schedule()[-1].ast
+    if ast.op is Ops.BEAM: ast = ast.src[0]
    s = Scheduler(ast, Device[Device.DEFAULT].renderer)
    for max_up in (2, 4):
      actions = get_kernel_actions(s, include_0=False, max_up=max_up)
@@ -1,23 +1,20 @@
 #!/usr/bin/env python3
-import os, platform, shutil, subprocess
+import os, shutil
 from pathlib import Path
 from tinygrad.helpers import fetch, OSX

-VERSION = "0.1.6"
 DEST = Path("/usr/local/lib")
 DEST.mkdir(exist_ok=True)

 if __name__ == "__main__":
  if OSX:
-    arch = "arm64" if platform.machine() == "arm64" else "x86_64"
-    dmg = fetch(f"https://github.com/ROCm/rocprof-trace-decoder/releases/download/{VERSION}/rocprof-trace-decoder-macos-{arch}-{VERSION}-Darwin.dmg")
-    mnt = Path(subprocess.check_output(["hdiutil", "attach", "-nobrowse", "-readonly", "-mountrandom", "/tmp", str(dmg)],
-                                       text=True).split("\t")[-1].strip())
-    try: shutil.copy2(next(mnt.rglob("librocprof-trace-decoder.dylib")), DEST)
-    finally: subprocess.run(["hdiutil", "detach", str(mnt)], check=True)
-    lib = DEST/"librocprof-trace-decoder.dylib"
+    fp = fetch("https://github.com/ROCm/rocprof-trace-decoder/releases/download/0.1.4/rocprof-trace-decoder-macos-arm64-0.1.4-Darwin.sh")
+    lib = fp.parent/"rocprof-trace-decoder-macos-arm64-0.1.4-Darwin"/"lib"/"librocprof-trace-decoder.dylib"
+    os.chmod(fp, 0o755)
+    os.system(f"sudo {fp} --prefix={fp.parent} --include-subdir")
+    shutil.copy2(lib, DEST)
  else:
    lib = DEST/"librocprof-trace-decoder.so"
-    os.system(f"sudo curl -L https://github.com/ROCm/rocprof-trace-decoder/raw/{VERSION}/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so -o {lib}")
+    os.system("sudo curl -L https://github.com/ROCm/rocprof-trace-decoder/raw/43bf0fef74a83c3c25badfc5a09c0bd39ed8c6f9/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so -o"+str(lib))
    os.system("sudo ldconfig")
-  print(f"Installed {lib.name} ({VERSION}) to", DEST)
+  print(f"Installed {lib.name} to", DEST)
@@ -1,13 +1,10 @@
 #!/usr/bin/env python3
-import ctypes, pathlib, argparse, pickle, dataclasses, threading, itertools
-from decimal import Decimal
+import ctypes, pathlib, argparse, pickle, dataclasses, threading
 from typing import Generator
 from tinygrad.helpers import temp, unwrap, DEBUG
 from tinygrad.runtime.ops_amd import ProfileSQTTEvent
 from tinygrad.runtime.autogen import rocprof
 from tinygrad.renderer.amd.dsl import Inst
-from tinygrad.helpers import ProfileEvent, ProfileRangeEvent, ProfilePointEvent
-from tinygrad.device import ProfileProgramEvent
 from test.amd.disasm import disasm

@dataclasses.dataclass(frozen=True)
@@ -129,71 +126,6 @@ def decode(sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, Inst]])
    raise exc
  return ROCParseCtx

-def unpack_occ(viz_data, i:int, j:int, key:tuple[str, int], data:list, p:ProfileProgramEvent, target:str) -> dict:
-  from tinygrad.viz.serve import amd_decode, create_step, row_tuple
-  steps = viz_data.ctxs[i]["steps"]
-  if len(steps[j+1:]) > 0: return {"steps":[{k:v for k,v in s.items() if k != "data"} for s in steps[j+1:]]}
-  base = unwrap(p.base)
-  disasm:dict[int, Inst] = {addr+base:inst for addr,inst in amd_decode(unwrap(p.lib), target).items()}
-  rctx = decode(data, {p.tag:disasm})
-  cu_events:dict[str, list[ProfileEvent]] = {}
-  # ** inst traces
-  wave_insts:dict[str, dict[str, dict]] = {}
-  inst_units:dict[str, itertools.count] = {}
-  for w in rctx.inst_execs.get(key, []):
-    if (u:=w.wave_loc) not in inst_units: inst_units[u] = itertools.count(0)
-    n = next(inst_units[u])
-    if (events:=cu_events.get(w.cu_loc)) is None: cu_events[w.cu_loc] = events = []
-    events.append(ProfileRangeEvent(f"SIMD:{w.simd}", loc:=f"INST WAVE:{w.wave_id} N:{n}", Decimal(w.begin_time), Decimal(w.end_time)))
-    wave_insts.setdefault(w.cu_loc, {})[f"{u} N:{n}"] = {"wave":w, "disasm":disasm, "prg":p, "run_number":n, "loc":loc}
-  # ** occ traces (only WAVESTART/WAVEEND)
-  units:dict[str, itertools.count] = {}
-  wave_start:dict[str, int] = {}
-  for occ in rctx.occ_events.get(key, []):
-    if (u:=occ.wave_loc) not in units: units[u] = itertools.count(0)
-    if u in inst_units: continue
-    if occ.start: wave_start[u] = occ.time
-    else:
-      if (events:=cu_events.get(occ.cu_loc)) is None: cu_events[occ.cu_loc] = events = []
-      events.append(ProfileRangeEvent(f"SIMD:{occ.simd}", f"OCC WAVE:{occ.wave_id} N:{next(units[u])}", Decimal(wave_start.pop(u)),Decimal(occ.time)))
-  # ** split graph by CU
-  for cu in sorted(cu_events, key=row_tuple):
-    steps.append(create_step(f"{cu} {len(cu_events[cu])}", ("/cu-sqtt", i, len(steps)), depth=1,
-                             data=[ProfilePointEvent(unit, "start", unit, ts=Decimal(0)) for unit in units]+cu_events[cu]))
-    for k in sorted(wave_insts.get(cu, []), key=row_tuple):
-      wd = wave_insts[cu][k]
-      steps.append(create_step(k.replace(cu, ""), ("/amd-sqtt-insts", i, len(steps)), loc=wd["loc"], depth=2,
-                               data={"fxn":unpack_insts, "args":(wd,)}))
-  return {"steps":[{k:v for k,v in s.items() if k != "data"} for s in steps[j+1:]]}
-
-def unpack_insts(viz_data, i:int, j:int, data:dict) -> dict:
-  columns = ["PC", "Instruction", "Hits", "Cycles", "Stall", "Type"]
-  inst_columns = ["N", "Clk", "Idle", "Dur", "Stall"]
-  # Idle:     The total time gap between the completion of previous instruction and the beginning of the current instruction.
-  #           The idle time can be caused by:
-  #             * Arbiter loss
-  #             * Source or destination register dependency
-  #             * Instruction cache miss
-  # Stall:    The total number of cycles the hardware pipe couldn't issue an instruction.
-  # Duration: Total latency in cycles, defined as "Stall time + Issue time" for gfx9 or "Stall time + Execute time" for gfx10+.
-  prev_instr = (w:=data["wave"]).begin_time
-  pc_to_inst = data["disasm"]
-  start_pc = None
-  rows:dict[int, dict] = {}
-  for pc, inst in pc_to_inst.items():
-    if start_pc is None: start_pc = pc
-    rows[pc] = {"pc":pc-start_pc, "inst":str(inst), "hit_count":0, "dur":0, "stall":0, "type":"", "hits":{"cols":inst_columns, "rows":[]}}
-  for e in w.unpack_insts():
-    if not (inst:=rows[e.pc]).get("type"): inst["type"] = str(e.typ).split("_")[-1]
-    inst["hit_count"] += 1
-    inst["dur"] += e.dur
-    inst["stall"] += e.stall
-    inst["hits"]["rows"].append((inst["hit_count"]-1, e.time, max(0, e.time-prev_instr), e.dur, e.stall))
-    prev_instr = max(prev_instr, e.time + e.dur)
-  summary = [{"label":"Total Cycles", "value":w.end_time-w.begin_time}, {"label":"SE", "value":w.se}, {"label":"CU", "value":w.cu},
-             {"label":"SIMD", "value":w.simd}, {"label":"Wave ID", "value":w.wave_id}, {"label":"Run number", "value":data["run_number"]}]
-  return {"rows":[tuple(v.values()) for v in rows.values()], "cols":columns, "metadata":[summary], "ref":viz_data.ref_map.get(data["prg"].name)}
-
 def print_data(data:dict) -> None:
  from tabulate import tabulate
  # plaintext
@@ -10,11 +10,11 @@ from tinygrad.uop.ops import UOp, Ops, KernelInfo

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

 def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
  return _sharded_empty(ref.shape, ref, axis)
@@ -55,6 +55,8 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
  assert attn_mask is None, "attn_mask not supported"
  assert is_causal, "only causal attention supported"

+  xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
+
  B, N, H, D = xq.shape
  H_KV = xk.shape[2]
  assert D == 128, "only D=128 supported"
@@ -79,7 +81,7 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False

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

-  return attn, attn, l_vec
+  return attn.transpose(1, 2), attn, l_vec

@functools.cache
 def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
@@ -93,20 +93,7 @@ constexpr int NUM_WARPS = 8;

 using G = kittens::group<NUM_WARPS>;

-// scale_mode: 0=no scale, 1=x only, 2=w only, 3=both
-#ifndef SCALE_MODE
-#define SCALE_MODE 3
-#endif
-
-__global__ __launch_bounds__(512, 2) void hk_fp8_gemm(bf16 *C_ptr, fp8e4m3 *A_ptr, fp8e4m3 *B_ptr
-#if SCALE_MODE == 1
-    , float *x_scale_ptr
-#elif SCALE_MODE == 2
-    , float *w_scale_ptr
-#elif SCALE_MODE == 3
-    , float *x_scale_ptr, float *w_scale_ptr
-#endif
-) {
+__global__ __launch_bounds__(512, 2) void hk_fp8_gemm(bf16 *C_ptr, fp8e4m3 *A_ptr, fp8e4m3 *B_ptr, float *x_scale_ptr, float *w_scale_ptr) {
    constexpr int M = GEMM_M, N = GEMM_N, K = GEMM_K;

    kittens::gl<fp8e4m3, 1, 1, M, K> A{A_ptr, nullptr, nullptr, nullptr, nullptr};
@@ -346,25 +333,11 @@ __global__ __launch_bounds__(512, 2) void hk_fp8_gemm(bf16 *C_ptr, fp8e4m3 *A_pt
    }

    // apply x_scale * w_scale before bf16 store to prevent overflow
-#if SCALE_MODE == 1
-    float scale = *x_scale_ptr;
-    mul(cA, cA, scale);
-    mul(cB, cB, scale);
-    mul(cC, cC, scale);
-    mul(cD, cD, scale);
-#elif SCALE_MODE == 2
-    float scale = *w_scale_ptr;
-    mul(cA, cA, scale);
-    mul(cB, cB, scale);
-    mul(cC, cC, scale);
-    mul(cD, cD, scale);
-#elif SCALE_MODE == 3
    float scale = *x_scale_ptr * *w_scale_ptr;
    mul(cA, cA, scale);
    mul(cB, cB, scale);
    mul(cC, cC, scale);
    mul(cD, cD, scale);
-#endif

    store(C, cA, {0, 0, block_row * WARPS_ROW * 2 + warp_m, block_col * WARPS_COL * 2 + warp_n});
    store(C, cB, {0, 0, block_row * WARPS_ROW * 2 + warp_m, block_col * WARPS_COL * 2 + WARPS_COL + warp_n});
@@ -359,7 +359,7 @@
 					"$(inherited)",
 					"@executable_path/../Frameworks",
 				);
-				MACOSX_DEPLOYMENT_TARGET = 13.0;
+				MACOSX_DEPLOYMENT_TARGET = 12.1;
 				MARKETING_VERSION = 1.0.0;
 				PRODUCT_BUNDLE_IDENTIFIER = org.tinygrad.tinygpu.installer;
 				PRODUCT_NAME = TinyGPU;
@@ -397,7 +397,7 @@
 					"$(inherited)",
 					"@executable_path/../Frameworks",
 				);
-				MACOSX_DEPLOYMENT_TARGET = 13.0;
+				MACOSX_DEPLOYMENT_TARGET = 12.1;
 				MARKETING_VERSION = 1.0.0;
 				PRODUCT_BUNDLE_IDENTIFIER = org.tinygrad.tinygpu.installer;
 				PRODUCT_NAME = TinyGPU;
@@ -446,7 +446,7 @@
 				CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
 				COPY_PHASE_STRIP = NO;
 				DEBUG_INFORMATION_FORMAT = dwarf;
-				DRIVERKIT_DEPLOYMENT_TARGET = 22.0;
+				DRIVERKIT_DEPLOYMENT_TARGET = 21.0;
 				ENABLE_STRICT_OBJC_MSGSEND = YES;
 				ENABLE_TESTABILITY = YES;
 				GCC_C_LANGUAGE_STANDARD = gnu11;
@@ -506,7 +506,7 @@
 				CODE_SIGN_IDENTITY = "Apple Development";
 				COPY_PHASE_STRIP = NO;
 				DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym";
-				DRIVERKIT_DEPLOYMENT_TARGET = 22.0;
+				DRIVERKIT_DEPLOYMENT_TARGET = 21.0;
 				ENABLE_NS_ASSERTIONS = NO;
 				ENABLE_STRICT_OBJC_MSGSEND = YES;
 				GCC_C_LANGUAGE_STANDARD = gnu11;
@@ -533,7 +533,7 @@
 				CODE_SIGN_STYLE = Automatic;
 				CURRENT_PROJECT_VERSION = 3;
 				DEVELOPMENT_TEAM = 9YG3G8543N;
-				DRIVERKIT_DEPLOYMENT_TARGET = 22.0;
+				DRIVERKIT_DEPLOYMENT_TARGET = 21.0;
 				ENABLE_USER_SCRIPT_SANDBOXING = YES;
 				EXCLUDED_ARCHS = "";
 				FRAMEWORK_SEARCH_PATHS = (
@@ -566,7 +566,7 @@
 				CURRENT_PROJECT_VERSION = 3;
 				DEVELOPMENT_TEAM = "";
 				"DEVELOPMENT_TEAM[sdk=driverkit*]" = 9YG3G8543N;
-				DRIVERKIT_DEPLOYMENT_TARGET = 22.0;
+				DRIVERKIT_DEPLOYMENT_TARGET = 21.0;
 				ENABLE_USER_SCRIPT_SANDBOXING = YES;
 				EXCLUDED_ARCHS = "";
 				FRAMEWORK_SEARCH_PATHS = (
@@ -188,8 +188,8 @@ kern_return_t TinyGPUDriver::CfgWrite(uint32_t off, uint32_t size, uint32_t val)
 kern_return_t TinyGPUDriver::ResetDevice()
 {
 	if (!ivars->pci) return kIOReturnNotReady;
-	kern_return_t ret = ivars->pci->Reset(kIOPCIDeviceResetTypeFunctionReset);
-	return ret == kIOReturnSuccess ? ret : ivars->pci->Reset(kIOPCIDeviceResetTypeHotReset);
+	ivars->pci->Reset(kIOPCIDeviceResetTypeFunctionReset);
+	return 0;
 }

 IOPCIDevice* TinyGPUDriver::GetPCI()
@@ -251,7 +251,8 @@ select = [
  "F541",
  "F841",
 ]
-"tinygrad/runtime/autogen/**/*.py" = ["E501", "F401", "E731", "F821", "A006", "A002", "F811", "F822"]
+"tinygrad/runtime/autogen/**/*.py" = ["E501", "F401", "E722", "E731", "F821", "A006", "A002", "F811"]
+"tinygrad/runtime/autogen/amd/**/*.py" = ["E501"]
 "test/amd/**/*.py" = ["F403", "F405"]

 [tool.ruff.format]
@@ -1 +0,0 @@
-Run `./render.sh` whenever you update tinyspec.tex to regenerate tinyspec.pdf.
@@ -1,10 +0,0 @@
-#!/bin/bash
-set -e
-
-if ! command -v tectonic &>/dev/null; then
-  echo "tectonic not found, installing..."
-  sudo pacman -S --noconfirm tectonic
-fi
-
-tectonic tinyspec.tex
-echo "done: tinyspec.pdf"
@@ -1,450 +0,0 @@
-\documentclass[10pt,letterpaper]{article}
-
-\usepackage[margin=0.75in]{geometry}
-\usepackage{amsmath,amssymb}
-\usepackage{booktabs}
-\usepackage{array}
-\usepackage[dvipsnames]{xcolor}
-\usepackage{enumitem}
-\usepackage{listings}
-\lstset{language=Python, basicstyle=\ttfamily\small, columns=fullflexible, keepspaces=true}
-
-\newcommand{\op}[1]{\textsc{#1}}
-
-\definecolor{movgreen}{HTML}{2E7D32}
-\definecolor{reducered}{HTML}{C62828}
-\definecolor{elwyellow}{HTML}{F9A825}
-\definecolor{callblue}{HTML}{1565C0}
-\definecolor{assignbrown}{HTML}{795548}
-\definecolor{multipurple}{HTML}{7B1FA2}
-\definecolor{markerorange}{HTML}{E65100}
-% AxisType colors (from tinygrad)
-\definecolor{axblue}{HTML}{1565C0}       % GLOBAL
-\definecolor{axcyan}{HTML}{00838F}       % LOCAL
-\definecolor{axbrcyan}{HTML}{00ACC1}     % WARP
-\definecolor{axbrblue}{HTML}{42A5F5}     % THREAD
-\definecolor{axwhite}{HTML}{616161}      % LOOP (gray on white paper)
-\definecolor{axred}{HTML}{C62828}        % REDUCE
-\definecolor{axbrred}{HTML}{E53935}      % GROUP_REDUCE
-\definecolor{axyellow}{HTML}{F9A825}     % UPCAST
-\definecolor{axmagenta}{HTML}{7B1FA2}    % UNROLL
-
-\title{tinygrad: a single dialect from Tensor programs to Command Buffers}
-\author{tinygrad, Corp. \\ \texttt{[email protected]}}
-\date{}
-
-\begin{document}
-\maketitle
-\thispagestyle{empty}
-
-\section*{UOps}
-
-All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op},\;\mathrm{src},\;\mathrm{arg},\;\mathrm{tag})$ where $\mathrm{op}$ is from the set below, $\mathrm{src}$ is a tuple of input UOps, $\mathrm{arg}$ is op-dependent, and $\mathrm{tag}$ is for temporary processing. The full program is a DAG of UOps. Each UOp has five derived properties --- \textbf{dtype}, \textbf{shape}, \textbf{device}, \textbf{min\_max}, and \textbf{axis} --- determined by the rules at the end of this document.
-
-%% ============================================================
-\subsection*{Source Ops \normalfont\small--- leaf nodes}
-
-\begin{tabular}{@{}l p{3.2cm} p{3.0cm} p{6.2cm}@{}}
-\toprule
-\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
-\midrule
-\op{Buffer} & () & size, dtype, device, addrspace &
-  Shape $(n \cdot \textit{size},)$ if device is $n$-tuple, else $(\textit{size},)$. \\
-\op{BufferView} & (buf,) & size, dtype, offset &
-  Typed access into a buffer. Zero-copy $(\textit{size},)$ slice at offset; inherits addrspace. \\
-\op{Param} & $(\mathbf{s})$ or $(\mathbf{s}, \text{min}, \text{max})$ & slot, dtype, device? &
-  Placeholder with shape $\mathbf{s}$. Substituted in \op{Function}. \\[4pt]
-\op{Const} & () & value, dtype &
-  A scalar constant with shape $(\ )$. \\
-\op{Vconst} & () & values, dtype &
-  A vector constant with shape $(n,)$. \\
-\bottomrule
-\end{tabular}
-
-\smallskip
-A \op{Buffer}'s \textbf{addrspace} is \texttt{GLOBAL}, \texttt{LOCAL}, or \texttt{REG}.
-
-%% ============================================================
-\subsection*{{\color{movgreen}Movement Ops} \normalfont\small--- no arithmetic, shapes are $(k,)$-shaped UOps with dtype \texttt{index} in src}
-
-\begin{tabular}{@{}l l l l@{}}
-\toprule
-\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
-\midrule
-\op{Permute} & $(T,)$ & axis order $\pi$ & Reorder axes. $\pi = (1,0)$ is transpose. \\
-\op{Flip} & $(T,)$ & bools $\mathbf{f}$ & Reverse along flagged axes. \\
-\op{Reshape} & $(T, \mathbf{s'})$ & --- & Reinterpret in row-major order. $\prod s_k = \prod s'_k$. \\
-\op{Expand} & $(T, \mathbf{s'})$ & --- & Broadcast size-1 axes. $s_k \in \{1, s'_k\}$. \\
-\op{Pad} & $(T, \mathbf{b}, \mathbf{e})$ & --- & Pad with $0$s: $b_k$ before, $e_k$ after each axis. \\
-\op{Shrink} & $(T, \mathbf{b}, \mathbf{e})$ & --- & Keep $[b_k, e_k)$ per axis. Inverse of \op{Pad}. \\
-\op{Index} & $(T, i_0, i_1, \ldots)$ & --- & Index from left. $()$-shaped $i$ removes dim; $(k,)$-shaped makes it $k$. \\
-\op{Stack} & $(T_0, T_1, \ldots)$ & --- & Join along a newly created leading axis. All shapes must match. \\
-\op{Replicated} & $(T,)$ & axes & Mark $T$ as replicated along axes. Collapse axes to $1$. \\
-\bottomrule
-\end{tabular}
-
-%% ============================================================
-\subsection*{{\color{reducered}Reduce Ops} \normalfont\small--- collapse axes to size $1$}
-
-\begin{tabular}{@{}l l l l@{}}
-\toprule
-\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
-\midrule
-\op{Reduce} & $(T,)$ & op, axes & Reduce $T$ along axes. Op is \op{Add}, \op{Max}, or \op{Mul}. \\
-\bottomrule
-\end{tabular}
-
-%% ============================================================
-\subsection*{{\color{callblue}Call Ops} \normalfont\small--- function abstraction, like the lambda calculus}
-
-\begin{tabular}{@{}l l l l@{}}
-\toprule
-\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
-\midrule
-\op{Function} & (body, $a_0$, $a_1$, \ldots) & --- & Substitute each \op{Param} $k$ in \op{Tuple} body with $a_k$. Gradient-able. \\
-\op{Call} & (body, $a_0$, $a_1$, \ldots) & --- & Opaque invocation of a compiled kernel or custom function. \\
-\op{Tuple} & $(v_0, v_1, \ldots)$ & --- & Pack values; required as \op{Function} body to return a value. \\
-\op{GetTuple} & $(T,)$ & idx & Extract element at idx from a \op{Tuple}. \\
-\bottomrule
-\end{tabular}
-
-%% ============================================================
-\subsection*{{\color{multipurple}Store Ops} \normalfont\small--- side effects}
-
-\begin{tabular}{@{}l l l l@{}}
-\toprule
-\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
-\midrule
-\op{Store} & (buf, val, gate?) & --- & Write val into buf. buf.shape $=$ val.shape. \\
-  & & & If gate is present, write only when gate is true. Output is void. \\
-\bottomrule
-\end{tabular}
-
-%% ============================================================
-\subsection*{{\color{assignbrown}Ordering Ops} \normalfont\small--- execution order}
-
-\begin{tabular}{@{}l l l p{6.0cm}@{}}
-\toprule
-\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
-\midrule
-\op{Range} & $(\text{bound},)$ & type & Iterator from $0$ to bound. \\
-\op{End} & (body, range) & --- & Close a \op{Range} loop. \\
-\op{After} & (buf, deps\ldots) & --- & Passthrough of buf; guarantees deps execute first. \\
-\op{Group} & $(u_0, u_1, \ldots)$ & --- & Void no-op that merges multiple \op{Store}s into one node, unordered. \\
-\op{Sink} & $(s_0, s_1, \ldots)$ & --- & Collect side effects into a single root node. \\
-\op{Linear} & (uops\ldots) & --- & Linearized (toposorted) instruction sequence. \\
-\bottomrule
-\end{tabular}
-
-\smallskip
-Assign is \op{Store} followed by \op{After}: write the value, then return the buffer with an ordering dependency.
-
-%% ============================================================
-\subsection*{{\color{elwyellow}Elementwise Ops} \normalfont\small--- all inputs same shape, output same shape, applied per-element}
-
-\begin{tabular}{@{}l l l l@{}}
-\toprule
-\textbf{Arity} & \textbf{src} & \textbf{Op} & \textbf{Semantics} \\
-\midrule
-Unary & $(T,)$
-  & \op{Recip}
-  & $1/x$ \\
-& & \op{Trunc}
-  & $\mathrm{trunc}(x)$: round toward zero. \\
-& & \op{Cast}
-  & Convert to target dtype (specified in arg). \\
-& & \op{Bitcast}
-  & Reinterpret bits as target dtype. Must be same size. \\[4pt]
-Binary & $(A, B)$
-  & \op{Add}, \op{Mul}, \op{Max}, \op{Mod}, \op{Idiv}
-  & $a+b$, $a \cdot b$, $\max(a,b)$, $a \bmod b$, $\lfloor a/b \rfloor$ \\
-& & \op{CmpLt}, \op{CmpNe}
-  & $[a < b]$, $[a \ne b]$ \\
-& & \op{Xor}, \op{Or}, \op{And}, \op{Shr}, \op{Shl}
-  & $a \oplus b$, $a \mid b$, $a \mathbin{\&} b$, $a \gg b$, $a \ll b$ \\[4pt]
-Ternary & $(P, A, B)$
-  & \op{Where}
-  & $A[\mathbf{i}]$ if $P[\mathbf{i}] \ne 0$, else $B[\mathbf{i}]$ \\
-
-\bottomrule
-\end{tabular}
-
-\medskip
-\textbf{Decomposed elementwise ops} --- defined in terms of the primitives above.
-
-\smallskip
-\begin{tabular}{@{}l l l@{}}
-\toprule
-\textbf{Op} & \textbf{Decomposition} & \textbf{Semantics} \\
-\midrule
-\op{Neg} & \op{Mul}($A$, $-1$) & $-x$ \\
-\op{Sub} & \op{Add}($A$, \op{Neg}($B$)) & $a - b$ \\
-\op{Div} & \op{Mul}($A$, \op{Recip}($B$)) & $a / b$ \\
-\op{CmpGt} & \op{CmpLt}($B$, $A$) & $[a > b]$ \\
-\op{CmpGe} & \op{CmpNe}(\op{CmpLt}($A$, $B$),\, $1$) & $[a \ge b]$ \\
-\op{CmpLe} & \op{CmpNe}(\op{CmpLt}($B$, $A$),\, $1$) & $[a \le b]$ \\
-\op{CmpEq} & \op{CmpNe}(\op{CmpNe}($A$, $B$),\, $1$) & $[a = b]$ \\
-\op{Not} & \op{CmpNe}($A$, $1$) & $\lnot a$ \\[4pt]
-\op{Exp2} & polynomial approx + \op{Mul}, \op{Add} & $2^x$ \\
-\op{Log2} & exponent extract + polynomial approx & $\log_2 x$ \\
-\op{Sin} & argument reduction + polynomial approx & $\sin x$ \\
-\op{Sqrt} & \op{Exp2}($0.5 \cdot$ \op{Log2}($A$)) & $\sqrt{x}$ \\
-\op{Pow} & \op{Exp2}(\op{Log2}($A$) $\cdot\, B$) & $a^b$ \\
-\op{Mulacc} & \op{Add}(\op{Mul}($A$, $B$),\, $C$) & $a \cdot b + c$ \\
-\op{Threefry} & 5 rounds of add-rotate-xor (ARX) & Threefry 2x32 PRNG \\
-\bottomrule
-\end{tabular}
-
-%% ============================================================
-\subsection*{{\color{markerorange}Marker Ops} \normalfont\small--- identity on data}
-
-\begin{tabular}{@{}l l l l@{}}
-\toprule
-\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
-\midrule
-\op{Contiguous} & $(T,)$ & --- & Force contiguous memory layout. \\
-\op{ContiguousBackward} & $(T,)$ & --- & Force contiguous in backward pass. \\
-\op{Detach} & $(T,)$ & --- & Stops gradient propagation. \\
-\op{Copy} & $(T,)$ & device & Copy to target device. \\
-\bottomrule
-\end{tabular}
-
-%% ============================================================
-\subsection*{Codegen Ops \normalfont\small--- generated code primitives, these do not appear in the main graph}
-
-\begin{tabular}{@{}l l l l@{}}
-\toprule
-\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
-\midrule
-\op{Load} & (idx,alt?,gate?) & --- & Dereference: read element at index from buffer. \\
- & & & All loads will be replaced by \op{Store}. \\
-\op{Barrier} & (deps\ldots) & --- & Synchronize threads within a workgroup. \\
-\op{Ins} & \ldots & \ldots & A single machine instruction (e.g.\ AMD ISA). \\
-\op{Special} & (bound,) & name & GPU thread/workgroup index (e.g.\ \texttt{gidx0}, \texttt{lidx1}). \\
-\op{If} & (gate,) & --- & Begin conditional execution block. \\
-\op{Endif} & (if,) & --- & End conditional execution block. \\
-\op{Wmma} & (A, B, acc) & config & Warp matrix multiply-accumulate (tensor cores). \\
-\op{Custom} & (args\ldots) & fmt & Inject custom code string into generated source. \\
-\op{AtomicAdd} & (idx, val) & --- & Atomic read-modify-write: \texttt{buf[idx] += val}. \\[4pt]
-\op{CustomFunction} & (meta\ldots) & name & Opaque device function (e.g.\ HW decode). Via \op{Call}. \\
-\op{Program} & (linear, source, binary) & --- & Compiled kernel: instructions, source, and machine code. \\
-\op{Source} & () & str & Human-readable rendered source code. \\
-\op{Binary} & () & bytes & Compiled machine code. \\
-\bottomrule
-\end{tabular}
-
-\smallskip
-These ops are not part of the core specification and are subject to change.
-
-%% ============================================================
-\subsection*{Derived Properties}
-
-Every UOp has a \textbf{dtype}, \textbf{shape}, \textbf{device}, \textbf{min\_max}, and \textbf{axis}, derived from its op, src, and arg:
-
-\medskip
-\begin{tabular}{@{}l l l l l@{}}
-\toprule
-\textbf{Op} & \textbf{dtype} & \textbf{shape} & \textbf{device} & \textbf{min\_max} \\
-\midrule
-\op{Buffer} & from arg & $(\text{size},)$ from arg & from arg & dtype range \\
-\op{Const} & from arg & $()$ & \textsc{null} & $[v, v]$ \\
-\op{Param} & from arg & from $\mathrm{src}[0]$ & from arg & from src or dtype range \\[3pt]
-Movement ops & $\mathrm{src}[0].\mathrm{dtype}$ & (see op) & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
-\op{Reduce} & $\mathrm{src}[0].\mathrm{dtype}$ & collapse axes to $1$ & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\[3pt]
-\op{Cast} & from arg & $\mathrm{src}[0].\mathrm{shape}$ & $\mathrm{src}[0].\mathrm{device}$ & clamped to dtype \\
-\op{Bitcast} & from arg & $\mathrm{src}[0].\mathrm{shape}$ & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\
-\op{Copy} & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0].\mathrm{shape}$ & from arg & $\mathrm{src}[0]$ \\
-ALU unary & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0].\mathrm{shape}$ & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\
-\op{Add} & $\mathrm{src}[0].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & $[a+b,\, A+B]$ \\
-\op{Mul} & $\mathrm{src}[0].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & $[\min,\max]$ of products \\
-\op{Max} & $\mathrm{src}[0].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & $[\max(a,b),\, \max(A,B)]$ \\
-Other binary & $\mathrm{src}[0].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\
-\op{CmpLt}, \op{CmpNe} & bool & broadcast & $\mathrm{src}[0].\mathrm{device}$ & from intervals \\
-\op{Where} & $\mathrm{src}[1].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & $[\min(b,c),\, \max(B,C)]$ \\[3pt]
-\op{Function}, \op{Call} & $\mathrm{src}[0].\mathrm{dtype}$ & substitute \op{Param} shapes & $\mathrm{src}[1].\mathrm{device}$ & dtype range \\
-\op{Range} & index & $()$ & \textsc{null} & $[0,\, n{-}1]$ \\
-\op{Index} & $\mathrm{src}[0].\mathrm{dtype}$ & remaining dims & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
-\op{Store} & void & $()$ & $\mathrm{src}[0].\mathrm{device}$ & --- \\
-\op{After} & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0].\mathrm{shape}$ & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
-\bottomrule
-\end{tabular}
-
-\smallskip
-$\mathrm{broadcast}$: right-align shapes, element-wise max; each axis must be equal or $1$.
-$[a,A]$, $[b,B]$, $[c,C]$ denote min\_max of $\mathrm{src}[0]$, $\mathrm{src}[1]$, $\mathrm{src}[2]$.
-Default \emph{dtype range}: $[\mathrm{dtype\_min},\, \mathrm{dtype\_max}]$.
-
-\medskip
-\textbf{axis} tracks the multi-device sharding dimension. \op{Buffer} with $n$-tuple device: axis $= 0$ (device dim).
-\op{Reshape} remaps axis to preserve the shard boundary. \op{Permute} follows the permutation.
-\op{Reduce} on the shard axis $\to$ \textsc{null}. \op{Replicated} on the shard axis $\to$ \textsc{null}. \op{Copy} $\to$ \textsc{null}. ALU ops inherit from sources. Default: \textsc{null}.
-
-%% ============================================================
-\subsection*{Kernel Optimizations (OptOps) \normalfont\small--- schedule-level transforms on kernel ranges}
-
-Each kernel's iteration space is a set of \op{Range} axes. Every range has an \textbf{AxisType}:
-
-\medskip
-\begin{tabular}{@{}l l l l l@{}}
-\toprule
-\textbf{AxisType} & \textbf{Letter} & \textbf{Split from} & \textbf{Direction} & \textbf{Semantics} \\
-\midrule
-{\color{axblue}\texttt{GLOBAL}} & \texttt{g} & --- & --- & GPU global workgroup dimension. \\
-{\color{axcyan}\texttt{LOCAL}} & \texttt{l} & g, L & inner & Workgroup local dimension (shared memory). \\
-{\color{axbrcyan}\texttt{WARP}} & \texttt{w} & \multicolumn{2}{l}{(created by \op{TC})} & Warp-level lanes for tensor cores. \\
-{\color{axbrblue}\texttt{THREAD}} & \texttt{t} & g & outer & CPU thread parallelism. \\
-{\color{axwhite}\texttt{LOOP}} & \texttt{L} & --- & --- & Generic sequential loop (initial state). \\
-{\color{axred}\texttt{REDUCE}} & \texttt{R} & --- & --- & Reduction axis. \\
-{\color{axbrred}\texttt{GROUP\_REDUCE}} & \texttt{G} & R & inner/outer & Shared-memory group reduction. \\
-{\color{axyellow}\texttt{UPCAST}} & \texttt{u} & g, l, L & inner & Register-level vectorization. \\
-{\color{axmagenta}\texttt{UNROLL}} & \texttt{r} & R, G & inner & Fully unrolled loop. \\
-\bottomrule
-\end{tabular}
-
-\medskip
-An optimization is a triple $(\mathrm{op},\;\mathrm{axis},\;\mathrm{arg})$:
-
-\smallskip
-\begin{tabular}{@{}l l l p{6.5cm}@{}}
-\toprule
-\textbf{OptOp} & \textbf{axis} & \textbf{arg} & \textbf{Semantics} \\
-\midrule
-\op{Split} & any & (factor $k$, target, top?) &
-  Split axis $n$ by $k$ into $(n/k, k)$ or $(k, n/k)$ if top. New sub-axis gets target AxisType (see table above). \\
-\op{Padto} & any & multiple $m$ &
-  Pad axis to next multiple of $m$ with validity masks. \\[4pt]
-\op{Swap} & axis$_i$ & axis$_j$ &
-  Swap two axes $i \leftrightarrow j$. \\
-\op{Nolocals} & --- & --- &
-  Disable local memory; no workgroup dims emitted. \\
-\op{TC} & reduce idx & (tc, opt, mode) &
-  Apply tensor core \op{Wmma}: split reduce/output axes into \texttt{WARP}, \texttt{UPCAST}, and \texttt{UNROLL} dims. \\
-\bottomrule
-\end{tabular}
-
-\smallskip
-Optimizations compose left-to-right. \op{TC} must be first. The search space is explored by BEAM search or hand-coded heuristics.
-
-%% ============================================================
-\subsection*{Common Ops as Compositions}
-
-All high-level tensor operations decompose into the primitives above.
-
-\begin{lstlisting}
-# gemm: C[M,N] = A[M,K] @ B[K,N]
-def gemm(A, B):
-  M,K = A.shape; _,N = B.shape
-  return (A.reshape(M,K,1) * B.reshape(1,K,N)).sum(1)
-
-# prefix_sum: cumulative sum via repeat+reshape sliding window trick
-def prefix_sum(T):
-  n = T.shape[0]
-  x = T.pad((n-1, 0))                            # (2n-1,)
-  x = x.reshape(1,2*n-1).expand(n+1,2*n-1)       # tile
-  x = x.reshape((n+1)*(2*n-1)).shrink_to(2*n*n)  # trim
-  x = x.reshape(n,2*n).shrink_to(n,n)            # windows
-  return x.sum(-1)                               # reduce
-
-# arange: prefix_sum of all 1s gives [1,2,...,n], subtract 1 for [0,1,...,n-1]
-def arange(n):
-  return prefix_sum(Tensor(1).reshape(1).expand(n)) - 1
-
-# gather: out[i] = T[idx[i]]. one-hot mask along gather axis, then reduce
-def gather(T, idx):
-  K = T.shape[0]
-  pos = arange(K).reshape(K, 1)                     # (K, 1)
-  mask = (pos == idx.reshape(1, -1)).cast(T.dtype)  # (K, D)
-  return (T.reshape(K, 1) * mask).sum(0)            # (D,)
-
-# scatter_add: T[idx[i]] += val[i]
-def scatter_add(T, idx, val):
-  K, D = T.shape[0], idx.shape[0]
-  pos = arange(K).reshape(K, 1)                     # (K, 1)
-  mask = (pos == idx.reshape(1, D)).cast(T.dtype)   # (K, D)
-  return T + (mask * val.reshape(1, D)).sum(1)      # (K,)
-\end{lstlisting}
-
-%% ============================================================
-\subsection*{{\color{multipurple}Multi-Device Collectives} \normalfont\small--- derived from primitives}
-
-Let $D = (d_0, \ldots, d_{n-1})$ be an $n$-tuple device.
-\op{Copy} to an $n$-tuple device reshards with axis $= 0$. \op{Copy} never changes shape.
-
-\begin{lstlisting}
-# T has shape (s,) on a single device.
-
-# broadcast: replicate T to all n devices
-def broadcast(T):
-  return T.reshape(1, s).expand(n, s).copy(D).replicated(0)        # (s,) on D, axis=null
-
-# scatter: split T into n chunks, one per device
-def scatter(T):
-  return T.copy(D)                                                  # (s,) on D, axis=0
-
-# T has shape (n*s,) on D with axis=0, so each device holds (s,) elements.
-
-# gather: collect all shards onto one device
-def gather(T):
-  return T.copy(D[0])                                               # (n*s,) on D[0], axis=null
-
-# reduce: gather + sum
-def reduce(T):
-  return gather(T).reshape(n, s).sum(0)                             # (s,) on D[0], axis=null
-
-# allgather: collect all shards, replicate to all devices
-def allgather(T):
-  return T.reshape(1, n*s).expand(n, n*s).copy(D).replicated(0)    # (n*s,) on D, axis=null
-
-# reduce_scatter: reduce across devices, scatter result
-def reduce_scatter(T):
-  return T.reshape(n, n, s//n).permute(1, 0, 2).copy(D).sum(1).reshape(s)  # (s,) on D, axis=0
-
-# allreduce: reduce_scatter + allgather
-def allreduce(T):
-  return allgather(reduce_scatter(T))                                       # (s,) on D, axis=null
-\end{lstlisting}
-
-%% ============================================================
-\subsection*{{\color{callblue}The \texttt{@function} Decorator} \normalfont\small--- graph capture via tracing}
-
-The \texttt{@function} decorator transforms a Python function on Tensors into a single \op{Function} node.
-
-\begin{lstlisting}
-@function
-def f(a: Tensor, b: Tensor) -> Tensor:
-  return a + b
-\end{lstlisting}
-
-When \texttt{f(x, y)} is called, the decorator:
-
-\begin{enumerate}[leftmargin=1.5em, itemsep=2pt]
-  \item \textbf{Extracts inputs}: walks all arguments to find every Tensor, deduplicates by identity.
-  \item \textbf{Runs the function} lazily (no device execution), building a UOp graph from the result.
-  \item \textbf{Parameterizes}: replaces each input UOp with a \op{Param}$(k)$ placeholder.
-  \item \textbf{Wraps the body} in a \op{Tuple} (even for single returns) and creates\\
-        \op{Function}(\op{Tuple}(body), $x$, $y$).
-  \item \textbf{Returns} the result via \op{GetTuple}$(0)$, or one \op{GetTuple} per element for tuple returns.
-\end{enumerate}
-
-The result is a reusable graph fragment: the body contains only \op{Param} references, not concrete buffers. At schedule time, the \op{Function} is resolved by substituting each \op{Param}$(k)$ back with its corresponding argument $a_k$, or lowered into an opaque \op{Call} if it is to be compiled as a reusable kernel.
-
-%% ============================================================
-\subsection*{Lowering Pipeline \normalfont\small--- from Tensor graph to machine code}
-
-\begin{tabular}{@{}l p{9.7cm}@{}}
-\toprule
-\textbf{Stage} & \textbf{Semantics} \\
-\midrule
-\textbf{Callify} & Transform the Tensor graph into a single stateless function. \\
-\textbf{Rangeify} & Determine the kernel split of the function. Break everything down to shape () \\
-\textbf{Optimize} & Insert local buffers. Swap and split ranges, and determine which axes are parallel and which are serial. \\
-\textbf{Expand} & Expand the parallel ranges into shape. \\
-\textbf{Instruction Selection} & Select target instructions, including WMMA and devectorization. \\
-\textbf{Linearize} & Topologically sort the graph and determine execution order. \\
-\textbf{Register/Memory Plan} & Allocate and reuse \texttt{GLOBAL}, \texttt{LOCAL}, and \texttt{REG} storage for values with non-overlapping lifetimes. \\
-\textbf{Render} & Output the machine code. \\
-\bottomrule
-\end{tabular}
-
-
-\end{document}
@@ -3,14 +3,12 @@ import functools
 import numpy as np
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear
 from tinygrad.renderer import Estimates
 from tinygrad.dtype import AddrSpace
-from tinygrad.helpers import getenv
 from tinygrad.runtime.autogen.amd.rdna3.ins import *
 import tinygrad.runtime.autogen.amd.rdna3.ins as r3
 import tinygrad.runtime.autogen.amd.rdna4.ins as r4
-from tinygrad.renderer.amd.dsl import s, v, NULL
+from tinygrad.renderer.amd.dsl import s, v
 from test.amd.helpers import TARGET_TO_ARCH
 from extra.gemm.amd_asm_matmul import Kernel

@@ -99,56 +97,43 @@ def custom_lds_sync(A:UOp, arch:str) -> UOp:
  sink = UOp.sink(A.base, lds, threads, wg, arg=KernelInfo("custom_lds_sync"))
  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

-def custom_handwritten(A:UOp) -> UOp:
+def custom_handwritten(A:UOp, arch:str) -> UOp:
  A = A.flatten()
  threads = UOp.special(128, "lidx0")
-  wg = UOp.special(1, "gidx0")
+  wg = UOp.special(256, "gidx0")
  lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=512, addrspace=AddrSpace.LOCAL), (), 'lds')  # 128 * 4 bytes
-  pipes = {getenv("PIPE", "")} if getenv("PIPE", "") else {"SALU", "VALU", "TRANSCENDENTAL", "WMMA"}
-  k = Kernel()
-  # wrap in loop to filter out icache misses
-  LOOP_N, UNROLL_N = 8, 5
-  k.emit(r4.s_mov_b32(s[1], LOOP_N))
-  k.label("loop")
-  if "SALU" in pipes:
-    for i in range(UNROLL_N):
+  k = Kernel(arch)
+  k.emit(r4.s_nop(0))
+  k.emit(r4.v_mov_b32_e32(v[1], 4))
+  def emit_alt():
+    for i in range(2):
+      k.emit(r4.v_mov_b32_e32(v[20+i], 4.0))
+      k.emit(r4.v_rcp_f32_e32(v[22+i], v[20+i]))
      k.emit(r4.s_mov_b32(s[20+i], i))
-      k.emit(r4.s_min_i32(s[30+i], i))
-      k.emit(r4.s_mov_b32(s[40+i], i))
      k.emit(r4.s_mul_i32(s[14+i], s[12+i], 32))
-  if "VALU" in pipes:
-    for i in range(UNROLL_N):
-      k.emit(r4.v_mov_b32_e32(v[20+i], i))
-      k.emit(r4.v_lshlrev_b64_e32(v[30+2*i:31+2*i], 2, v[12+i:13+i]))
-      k.emit(r4.v_mad_co_u64_u32(v[40+2*i:41+2*i], NULL, v[12+i], v[13+i], v[14+i:15+i]))
-  if "TRANSCENDENTAL" in pipes:
-    # transcendental VALU runs on the TFU, it can run regular VALU at the same time
-    for i in range(UNROLL_N):
-      k.emit(r4.v_mov_b32_e32(v[20+i], i))
-      k.emit(r4.v_s_rcp_f32(s[20+i], s[12+i]))
-      k.emit(r4.v_rcp_f32_e32(v[30+i], v[12+i]))
-      k.emit(r4.v_s_exp_f32(s[30+i], s[12+i]))
-  if "WMMA" in pipes:
-    base = 30
-    for i in range(UNROLL_N):
-      a = base + i*40
-      b, cd = a + 4, a + 8
-      k.emit(r4.v_wmma_f32_16x16x16_f16(v[cd:cd+7], v[a:a+3], v[b:b+3], v[cd:cd+7]))
-      a = base + i*40 + 16
-      b, cd = a + 2, a + 4
-      k.emit(r4.v_wmma_i32_16x16x16_iu8(v[cd:cd+7], v[a:a+1], v[b:b+1], v[cd:cd+7]))
-  k.emit(r4.s_add_co_i32(s[1], s[1], -1))
-  k.emit(r4.s_cmp_eq_i32(s[1], 0))
-  k.emit(r4.s_cbranch_scc0(), target="loop")
+  def emit_wmma():
+    for _ in range(2):
+      k.emit(r4.v_wmma_f32_16x16x16_f16(v[0:7], v[8:11], v[8:11], 1))
+  k.label("start")
+  k.emit(s_mov_b32(s[1], 10))
+  k.label("loop")
+  # wmma should've overlapped here if it was a different unit?
+  for _ in range(2):
+    emit_wmma()
+    emit_alt()
+  for _ in range(8): k.emit(s_nop(1))
+  k.emit(s_add_u32(s[1], s[1], -1))
+  k.emit(s_cmp_eq_i32(s[1], 0))
+  k.emit(s_cbranch_scc0(), target="loop")
  k.emit(r4.s_endpgm())
  insts = k.finalize()
  sink = UOp.sink(A.base, threads, wg, lds, arg=KernelInfo("custom_handwritten"))
  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

-def custom_data_deps(A:UOp) -> UOp:
+def custom_data_deps(A:UOp, arch:str) -> UOp:
  A = A.flatten()
  threads = UOp.special(A.numel(), "lidx0")
-  k = Kernel()
+  k = Kernel(arch)
  k.emit(s_load_b64(s[0:1], s[0:1], soffset=NULL))
  k.emit(s_waitcnt_lgkmcnt(sdst=NULL, simm16=0))
  k.emit(v_lshlrev_b32_e32(v[0], 2, v[0]))
@@ -169,11 +154,10 @@ class TestCustomKernel(unittest.TestCase):
    if self.arch != "rdna3": self.skipTest("only rdna3")
    a = Tensor.full((16, 16), 1.).contiguous().realize()
    a = Tensor.custom_kernel(a, fxn=custom_add_one)[0]
-    linear = compile_linear(a.schedule_linear())
-    est = estimate_uop(linear.src[-1])
-    self.assertEqual(est.ops, a.numel())
-    self.assertEqual(est.mem, a.nbytes()*2)
-    run_linear(linear)
+    ei = a.schedule()[-1].lower()
+    self.assertEqual(ei.prg.estimates.ops, a.numel())
+    self.assertEqual(ei.prg.estimates.mem, a.nbytes()*2)
+    ei.run()
    self.assertTrue((a.numpy() == 2.).all())

  def test_variable(self):
@@ -181,9 +165,9 @@ class TestCustomKernel(unittest.TestCase):
    b = Tensor.full((16, 16), 1, dtype=dtypes.uint32).contiguous().realize()
    a = Tensor.zeros_like(b).contiguous().realize()
    a = Tensor.custom_kernel(a, b, fxn=custom_add_var)[0]
-    linear = a.schedule_linear()
+    ei = a.schedule()[-1].lower()
    for i in range(4):
-      run_linear(linear, var_vals={"var":i})
+      ei.run({"var":i})
      self.assertTrue((a.numpy() == 1+i).all())

  def test_lds_sync(self):
@@ -198,13 +182,13 @@ class TestCustomKernel(unittest.TestCase):
  def test_handwritten(self):
    if self.arch != "rdna4": self.skipTest("only tested on rdna4")
    a = Tensor.empty(1024, dtype=dtypes.int32).contiguous().realize()
-    a = Tensor.custom_kernel(a, fxn=custom_handwritten)[0]
+    a = Tensor.custom_kernel(a, fxn=functools.partial(custom_handwritten, arch=self.arch))[0]
    a.realize()

  def test_data_deps(self):
    if self.arch != "rdna3": self.skipTest("only tested on rdna3")
    a = Tensor(np.full(32, 5.0, dtype=np.float32)).realize()
-    a = Tensor.custom_kernel(a, fxn=custom_data_deps)[0]
+    a = Tensor.custom_kernel(a, fxn=functools.partial(custom_data_deps, arch=self.arch))[0]
    a.realize()
    self.assertTrue((a.numpy() == 6.0).all())

@@ -78,18 +78,18 @@ class TestTinygradIntegration(unittest.TestCase):
  def _get_kernel_code(self, op_fn) -> bytes:
    from tinygrad import Tensor
    from tinygrad.helpers import Target
-    from tinygrad.codegen import to_program
+    from tinygrad.codegen import get_program
    from tinygrad.renderer.llvmir import AMDLLVMRenderer
    from tinygrad.runtime.support.elf import elf_loader
    from tinygrad.uop.ops import Ops

    result = op_fn(Tensor)
-    linear = result.schedule_linear()
-    sink_items = [call for call in linear.src if call.src[0].op == Ops.SINK]
+    schedule = result.schedule()
+    sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
    assert len(sink_items) > 0, "No SINK in schedule"
    renderer = AMDLLVMRenderer(Target("AMD", arch='gfx1100'))
-    prg = to_program(sink_items[0].src[0], renderer)
-    lib = renderer.compiler.compile(prg.src[3].arg)
+    prg = get_program(sink_items[0].ast, renderer)
+    lib = renderer.compiler.compile(prg.src)
    return next(s.content for s in elf_loader(lib)[1] if s.name == ".text")

  def test_simple_add_kernel(self):
@@ -8,14 +8,17 @@ class TestMockGPUInvalidInstruction(unittest.TestCase):
    test_code = '''
 import struct
 from tinygrad import Device, Tensor
-from tinygrad.engine.realize import compile_linear
+from tinygrad.engine.realize import get_runner
 from tinygrad.runtime.ops_amd import AMDProgram

 dev = Device["AMD"]
 a = Tensor([1.0]).realize()
 b = a + 1
-linear = compile_linear(b.schedule_linear())
-lib = bytearray(linear.src[-1].src[0].src[4].arg)
+si = b.schedule()[-1]
+runner = get_runner(dev.device, si.ast)
+
+prg = runner._prg
+lib = bytearray(prg.lib)

 # Find s_endpgm (0xBFB00000) and replace with V_MOVRELD_B32 (op=66) which has no pcode
 # VOP1 encoding: bits[31:25]=0x7E, op=bits[16:9], so op=66 -> 66<<9 = 0x8400
@@ -57,53 +57,49 @@ class KernelSnapshot:
 def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelSnapshot], dict[int, int], dict[int, bytes]]:
  """Compile a tinygrad operation and extract all kernels with their buffer mappings."""
  from tinygrad import Tensor
-  from tinygrad.uop.ops import Ops
-  from tinygrad.engine.realize import compile_linear, resolve_params, unwrap_multi
  from tinygrad.runtime.support.elf import elf_loader

  out = op_fn(Tensor)
-  linear = compile_linear(out.schedule_linear())
+  sched = out.schedule()
  kernels = []
  buf_pool: dict[int, int] = {}  # buffer id -> size
  buf_data: dict[int, bytes] = {}  # buffer id -> initial data from COPY

-  for call in linear.src:
-    ast = call.src[0]
-    for bufs, _ in unwrap_multi(call, resolve_params(call, ())):
-      if ast.op is Ops.COPY:
-        # Handle COPY: extract source data to initialize destination buffer
-        if len(bufs) >= 2:
-          dst_buf, src_buf = bufs[0], bufs[1]
-          dst_id = id(dst_buf)
-          if dst_id not in buf_pool:
-            buf_pool[dst_id] = dst_buf.nbytes
-          # Get source data if it's from numpy/CPU
-          if hasattr(src_buf, 'base') and src_buf.base is not None and hasattr(src_buf.base, '_buf'):
-            src_data = bytes(src_buf.base._buf)
-            buf_data[dst_id] = src_data
-      elif ast.op is Ops.PROGRAM:
-        info = ast.arg
-        if len(ast.src) > 4 and ast.src[4].op is Ops.BINARY:
-          lib = bytes(ast.src[4].arg)
-          _, sections, _ = elf_loader(lib)
-          for sec in sections:
-            if sec.name == '.text':
-              buf_idxs = []
-              buf_sizes = []
-              for b in bufs:
-                buf_id = id(b)
-                if buf_id not in buf_pool:
-                  buf_pool[buf_id] = b.nbytes
-                buf_idxs.append(buf_id)
-                buf_sizes.append(b.nbytes)
-              kernels.append(KernelSnapshot(
-                code=bytes(sec.content),
-                src=ast.src[3].arg,
-                global_size=tuple(info.global_size),
-                local_size=tuple(info.local_size),
-                buf_idxs=buf_idxs,
-                buf_sizes=buf_sizes
-              ))
+  for ei in sched:
+    lowered = ei.lower()
+    if ei.ast.op.name == 'COPY':
+      # Handle COPY: extract source data to initialize destination buffer
+      if len(lowered.bufs) >= 2:
+        dst_buf, src_buf = lowered.bufs[0], lowered.bufs[1]
+        dst_id = id(dst_buf)
+        if dst_id not in buf_pool:
+          buf_pool[dst_id] = dst_buf.nbytes
+        # Get source data if it's from numpy/CPU
+        if hasattr(src_buf, 'base') and src_buf.base is not None and hasattr(src_buf.base, '_buf'):
+          src_data = bytes(src_buf.base._buf)
+          buf_data[dst_id] = src_data
+    elif ei.ast.op.name == 'SINK':
+      if lowered.prg and lowered.prg.p.lib:
+        lib = bytes(lowered.prg.p.lib)
+        _, sections, _ = elf_loader(lib)
+        for sec in sections:
+          if sec.name == '.text':
+            buf_idxs = []
+            buf_sizes = []
+            for b in lowered.bufs:
+              buf_id = id(b)
+              if buf_id not in buf_pool:
+                buf_pool[buf_id] = b.nbytes
+              buf_idxs.append(buf_id)
+              buf_sizes.append(b.nbytes)
+            kernels.append(KernelSnapshot(
+              code=bytes(sec.content),
+              src=lowered.prg.p.src,
+              global_size=tuple(lowered.prg.p.global_size),
+              local_size=tuple(lowered.prg.p.local_size),
+              buf_idxs=buf_idxs,
+              buf_sizes=buf_sizes
+            ))
  if not kernels: raise RuntimeError("No kernel found")
  return kernels, buf_pool, buf_data

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 """Tests for SQTT encoder: verifies the emulator produces correct SQTT traces for known kernels.

-Run with: DEV=MOCK+AMD python -m pytest test/amd/test_sqtt_encoder.py -v
+Run with: DEV=MOCKKFD+AMD python -m pytest test/amd/test_sqtt_encoder.py -v
 """
 import ctypes, unittest
 from tinygrad.helpers import Context
@@ -8,11 +8,10 @@ from tinygrad.runtime.support.elf import elf_loader
 from tinygrad.renderer.amd import decode_inst
 from tinygrad.runtime.autogen.amd.rdna3.ins import SOPP
 from tinygrad.runtime.autogen.amd.rdna3.enum import SOPPOp
-from tinygrad.renderer.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_RDNA4, WAVEEND, WAVEEND_RDNA4, INST, INST_RDNA4, VALUINST,
+from tinygrad.renderer.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_RDNA4, WAVEEND, INST, INST_RDNA4, VALUINST,
                                     IMMEDIATE, IMMEDIATE_MASK, PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4, PACKET_TYPES_CDNA, CDNA_WAVESTART,
                                     print_packets, CDNA_WAVEEND, CDNA_INST)
 from test.amd.helpers import TARGET_TO_ARCH
-from test.amd.test_sqttmap import needs_rocprof

 import tinygrad
 EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"
@@ -133,7 +132,7 @@ class SQTTExamplesTestBase(unittest.TestCase):
      with self.subTest(example=name):
        all_packets = [p for e in events for p in decode(e.blob)]
        self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVESTART, WAVESTART_RDNA4, CDNA_WAVESTART))]), 0, f"no WAVESTART in {name}")
-        self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVEEND, WAVEEND_RDNA4, CDNA_WAVEEND))]), 0, f"no WAVEEND in {name}")
+        self.assertGreater(len([p for p in all_packets if isinstance(p, (WAVEEND, CDNA_WAVEEND))]), 0, f"no WAVEEND in {name}")

  def test_time_monotonic(self):
    for name, (events, *_) in self.examples.items():
@@ -161,7 +160,6 @@ class SQTTExamplesTestBase(unittest.TestCase):
        counts = [len(list(decode(e.blob))) for e in events]
        self.assertEqual(counts, self.expected[name], f"packet count mismatch in {name}")

-  @needs_rocprof
  def test_rocprof_wave_times_match(self):
    """Wave start/end times must match rocprof exactly."""
    for name, (events, lib, base) in self.examples.items():
@@ -182,7 +180,7 @@ class SQTTExamplesTestBase(unittest.TestCase):
          for p in decode(event.blob):
            if first_timestamp is None: first_timestamp = p._time
            if isinstance(p, (WAVESTART, CDNA_WAVESTART, WAVESTART_RDNA4)): wave_starts[(p.wave, p.simd, p.cu)] = p._time
-            elif isinstance(p, (WAVEEND, WAVEEND_RDNA4, CDNA_WAVEEND)) and (key := (p.wave, p.simd, p.cu)) in wave_starts:
+            elif isinstance(p, (WAVEEND, CDNA_WAVEEND)) and (key := (p.wave, p.simd, p.cu)) in wave_starts:
              our_waves.append((wave_starts[key], p._time))
          for st in wave_starts.values():
            self.assertGreater(st, first_timestamp, "wave start must be after the first packet")
@@ -191,7 +189,6 @@ class SQTTExamplesTestBase(unittest.TestCase):
        for st, et in our_waves:
          self.assertGreater(et, st, "wave end must be after start")

-  @needs_rocprof
  def test_rocprof_inst_times_match(self):
    """Instruction times must match rocprof exactly (excluding s_endpgm)."""
    for name, (events, lib, base) in self.examples.items():
@@ -1,8 +1,6 @@
 import unittest, contextlib
 from tinygrad import Device, Tensor, Context, TinyJit
 from tinygrad.device import Compiled, ProfileProgramEvent, ProfileDeviceEvent
-from tinygrad.engine.realize import run_linear
-from tinygrad.codegen import to_program
 from tinygrad.viz.serve import load_amd_counters, VizData

@contextlib.contextmanager
@@ -28,41 +26,39 @@ class TestSQTTProfiler(unittest.TestCase):
  def test_simple(self):
    t = Tensor.empty(1) + 1
    with save_sqtt() as sqtt:
-      linear = t.schedule_linear()
-      run_linear(linear)
-    fn_name = to_program(linear.src[0].src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name
+      ei = t.schedule()[0].lower()
+      ei.run()
    self.assertEqual(len(sqtt), 1)
-    self.assertEqual(sqtt[0]["name"], f"SQTT {fn_name}")
+    self.assertEqual(sqtt[0]["name"], f"SQTT {ei.prg.p.function_name}")

  def test_multiple_runs(self):
    t = Tensor.empty(1) + 1
    with save_sqtt() as sqtt:
-      linear = t.schedule_linear()
-      for _ in range(N:=3): run_linear(linear)
-    fn_name = to_program(linear.src[0].src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name
+      ei = t.schedule()[0].lower()
+      for _ in range(N:=3):
+        ei.run()
    self.assertEqual(len(sqtt), N)
    for i in range(1, N):
-      self.assertEqual(sqtt[i]["name"], f"SQTT {fn_name} n{i+1}")
+      self.assertEqual(sqtt[i]["name"], f"SQTT {ei.prg.p.function_name} n{i+1}")

  def test_multiple_kernels(self):
    t = ((Tensor.empty(1) + 1).contiguous() + 2)
-    linear = t.schedule_linear()
+    sched = t.schedule()
    with save_sqtt() as sqtt:
-      run_linear(linear)
-    self.assertEqual(len(sqtt), len(linear.src))
-    for i,call in enumerate(linear.src):
-      fn_name = to_program(call.src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name
-      self.assertEqual(sqtt[i]["name"], f"SQTT {fn_name}")
+      for si in sched: si.lower().run()
+    self.assertEqual(len(sqtt), len(sched))
+    for i,k in enumerate(sched):
+      self.assertEqual(sqtt[i]["name"], f"SQTT {k.lower().prg.p.function_name}")

  def test_multiple_kernels_lower(self):
    t = ((Tensor.empty(1) + 1).contiguous() + 2)
-    linear = t.schedule_linear()
+    sched = t.schedule()
    with save_sqtt() as sqtt:
-      run_linear(linear)
-    self.assertEqual(len(sqtt), len(linear.src))
-    for i,call in enumerate(linear.src):
-      fn_name = to_program(call.src[0], renderer=Device[Device.DEFAULT].renderer).arg.function_name
-      self.assertEqual(sqtt[i]["name"], f"SQTT {fn_name}")
+      prgs = [si.lower() for si in sched]
+      for p in prgs: p.run()
+    self.assertEqual(len(sqtt), len(sched))
+    for i,ei in enumerate(prgs):
+      self.assertEqual(sqtt[i]["name"], f"SQTT {ei.prg.p.function_name}")

  def test_jit(self):
    @TinyJit
@@ -1,36 +1,19 @@
 # test to compare every packet with the rocprof decoder
-import unittest, pickle, functools, json
+import unittest, pickle
 from typing import Iterator
 from pathlib import Path
-from tinygrad.helpers import DEBUG, getenv, temp, ansistrip, Context
+from tinygrad.helpers import DEBUG, getenv, temp, ansistrip
 from tinygrad.renderer.amd.sqtt import print_packets, map_insts
 from tinygrad.runtime.autogen.amd.rdna3.ins import s_endpgm
-from tinygrad.viz.serve import sqtt_timeline, amd_decode
+from tinygrad.viz.serve import sqtt_timeline
 from test.amd.disasm import disasm
 from test.null.test_viz import run_cli

 import tinygrad
 EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"

-def needs_rocprof(fn):
-  @functools.wraps(fn)
-  def wrapper(self, *args, **kwargs):
-    # check if latest rocprof is available, if not, skip rocprof comparison tests
-    # rocprof doesn't have a version string, decode a known pickle to validate it's the latest
-    try:
-      from extra.sqtt.roc import decode as roc_decode
-      with open(EXAMPLES_DIR/"gfx1200"/"profile_plus_run_0.pkl", "rb") as f:
-        data = pickle.load(f)
-      sqtt = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"][1]
-      kern = {e.tag:e for e in data if type(e).__name__ == "ProfileProgramEvent"}[sqtt.kern]
-      rctx = roc_decode([sqtt], {kern.tag:{addr+kern.base:inst for addr,inst in amd_decode(kern.lib, "gfx1200").items()}})
-      insts = [e.time for e in list(rctx.inst_execs.values())[0][0].unpack_insts()]
-      self.assertListEqual(insts, [28178, 28179, 28180, 28181, 28182, 29882, 29883, 29884, 29885, 30966, 30983, 30985, 30992, 30993])
-    except Exception as e: self.skipTest(f"latest rocprof not available, install with extra/sqtt/install_rocprof_decoder.py: {e}")
-    return fn(self, *args, **kwargs)
-  return wrapper
-
 def rocprof_inst_traces_match(sqtt, prg, target):
+  from tinygrad.viz.serve import amd_decode
  from extra.sqtt.roc import decode as roc_decode, InstExec
  addr_table = amd_decode(prg.lib, target)
  disasm_map = {addr+prg.base:inst for addr,inst in addr_table.items()}
@@ -80,7 +63,6 @@ class TestSQTTMapBase(unittest.TestCase):
      if sqtt_events and kern_events:
        cls.examples[pkl_path.stem] = (sqtt_events, kern_events, cls.target)

-  @needs_rocprof
  def test_rocprof_inst_traces_match(self):
    for name, (events, kern_events, target) in self.examples.items():
      if "sync" in name and self.target.startswith("gfx12"):
@@ -112,7 +94,7 @@ class TestSQTTMapBase(unittest.TestCase):
            elif "WAVE" in e.device:
              # sopk/immediates don't get ALU/MEM EXEC
              if e.name.display_name not in {"IMMEDIATE", "IMMEDIATE_MASK", "JUMP", "JUMP_NO", "MESSAGE", "BARRIER", "BARRIER_SIGNAL",
-                                             "WAVEEND", "WAVEEND_RDNA4", "WAVERDY"} and not e.name.display_name.startswith("OTHER_"): insts += 1
+                                             "WAVEEND", "WAVERDY"} and not e.name.display_name.startswith("OTHER_"): insts += 1
            else: raise Exception(f"timeline row must be INST or EXEC, got {e.device}")
          self.assertEqual(execs, insts)

@@ -129,18 +111,15 @@ class TestSQTTMapBase(unittest.TestCase):

  def test_sqtt_cli(self):
    for pkl_path in sorted((EXAMPLES_DIR/self.target).glob("*.pkl")):
-      out = run_cli("--profile-path", str(pkl_path), "--ls")
+      out = run_cli("--profile", "--profile-path", str(pkl_path))
      sqtt_traces = [l.strip() for l in out.split("\n") if "SQTT" in l]
      for name in sqtt_traces:
-        out = run_cli("--profile-path", str(pkl_path), "-s", ansistrip(name))
+        out = run_cli("--profile", "--profile-path", str(pkl_path), "-s", ansistrip(name))
        lines = out.split("\n")
        self.assertIn("Clk", lines[0])
        for r in lines[2:]:
          parts = r.split()
          self.assertTrue(parts[0].isdigit(), f"expected clock timestamp, got {parts[0]}")
-      with Context(DEBUG=2):
-        kernels = run_cli("--profile-path", str(pkl_path), "-s", "AMD").split("\n")
-      self.assertEqual(len(kernels), len(self.examples[pkl_path.stem][1]))

 class TestSQTTMapRDNA3(TestSQTTMapBase): target = "gfx1100"

@@ -148,25 +127,14 @@ class TestSQTTMapRDNA4(TestSQTTMapBase):
  target = "gfx1200"

  @unittest.expectedFailure
-  def test_pipes(self):
+  def test_rdna4_wmma(self):
    events, kernels, target = self.examples["profile_handwritten_run_0"]
-    lib = list(kernels.values())[0].lib
-    dispatch_st:dict[str, int] = {}
-    row_ends:dict[str, int] = {}
-    row_counts:dict[str, int] = {}
-    for e in sqtt_timeline(events[1].blob, lib, target):
-      if type(e).__name__ != "ProfileRangeEvent": continue
-      info = json.loads(e.name.ret) if e.name.ret else {}
-      if e.device.startswith("WAVE"):
-        idx = row_counts.get(e.device, 0)
-        dispatch_st[f"{e.device}-{idx}"] = int(e.st)
-        row_counts[e.device] = idx + 1
-      elif info.startswith("LINK:"):
-        delay = int(e.st) - dispatch_st[info[len("LINK:"):]]
-        self.assertGreaterEqual(delay, 1, f"EXEC {e.device} starts before DISPATCH: delay={delay}")
-        if (prev_en:=row_ends.get(e.device)) is not None:
-          self.assertGreaterEqual(e.st, prev_en, f"EXEC overlap in {e.device}: {e.st} < prev end {prev_en}")
-        row_ends[e.device] = int(e.en)
+    row_ends = {}
+    for e in sqtt_timeline(events[0].blob, list(kernels.values())[0].lib, target):
+      if type(e).__name__ != "ProfileRangeEvent" or e.device != "ALUEXEC:0 WMMA": continue
+      if (et:=row_ends.get(e.device)) is not None and e.st < et:
+        raise RuntimeError(f"WMMA exec overlaps in {e.device}: {e.st} {et}.")
+      row_ends[e.device] = e.en

 class TestSQTTMapCDNA(TestSQTTMapBase):
  target = "gfx950"
@@ -2,18 +2,22 @@ import unittest
 import numpy as np
 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.engine.realize import run_schedule
+from tinygrad.engine.realize import CompiledRunner, get_program
+from tinygrad.schedule import ExecItem
+from tinygrad.renderer import Estimates
 from tinygrad.renderer.ptx import PTXRenderer
 from test.helpers import needs_second_gpu

 class TestArange(unittest.TestCase):
  def _get_flops(self, tensor, desired):
    GlobalCounters.reset()
-    linear = compile_linear(tensor.schedule_linear())
-    self.assertEqual(len(linear.src), 1)
-    run_linear(linear)
+    sched = tensor.schedule()
+    self.assertEqual(len(sched), 1)
+    p = get_program(sched[-1].ast, renderer=Device[Device.DEFAULT].renderer)
+    ExecItem(sched[-1].ast, [tensor.uop.buffer], prg=CompiledRunner(p)).run()
    np.testing.assert_equal(tensor.numpy(), desired)
-    return estimate_uop(linear.src[-1]).ops
+    return p.estimates.ops

  def test_arange_complexity(self):
    self.assertEqual(self._get_flops(Tensor.arange(256), np.arange(256)), 0)
@@ -36,8 +40,9 @@ class TestArange(unittest.TestCase):
  def test_tri_complexity(self):
    with Context(NOOPT=1):
      t = Tensor.ones(256, 256).contiguous().realize()
-      linear = compile_linear(t.triu().schedule_linear())
-      self.assertLessEqual(estimate_uop(linear.src[-1]).ops, 4 * 256 * 256)
+      sched = t.triu().schedule()
+      p = get_program(sched[-1].ast, renderer=Device[Device.DEFAULT].renderer)
+      self.assertLessEqual(Estimates.from_uops(p.uops).ops, 4 * 256 * 256)

 DSET, DDIM = 2048, 32

@@ -49,9 +54,9 @@ class TestIndexing(unittest.TestCase):
    with Context(NOOPT=1):
      GlobalCounters.reset()
      out = ((Tensor.arange(1,16385)-1)*needle).sum()
-      linear, var_vals = out.linear_with_vars()
-      self.assertEqual(len(linear.src), 1)
-      run_linear(linear, var_vals)
+      sched = out.schedule()
+      self.assertEqual(len(sched), 1)
+      run_schedule(sched)
    self.assertEqual(out.item(), 1337)

  def test_manual_index(self):
@@ -66,9 +71,9 @@ 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))
      X = full.sum(axis=(2,3))
-      linear, var_vals = X.linear_with_vars()
-      self.assertEqual(len(linear.src), 1)
-      run_linear(linear, var_vals)
+      sched = X.schedule()
+      self.assertEqual(len(sched), 1)
+      run_schedule(sched)
      assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops}"
    np.testing.assert_allclose(real_index, X.numpy())

@@ -92,9 +97,9 @@ class TestIndexing(unittest.TestCase):
      GlobalCounters.reset()
      X = dataset[idxs]
      assert X.shape == (4,DDIM)
-      linear, var_vals = X.linear_with_vars()
-      self.assertEqual(len(linear.src), 1)
-      run_linear(linear, var_vals)
+      sched = X.schedule()
+      self.assertEqual(len(sched), 1)
+      run_schedule(sched)
      assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops}"
    np.testing.assert_allclose(real_index, X.numpy())

@@ -107,9 +112,9 @@ class TestIndexing(unittest.TestCase):
      GlobalCounters.reset()
      X = dataset[idxs]
      assert X.shape == (4,DDIM)
-      linear, var_vals = X.linear_with_vars()
-      self.assertEqual(len(linear.src), 1)
-      run_linear(linear, var_vals)
+      sched = X.schedule()
+      self.assertEqual(len(sched), 1)
+      run_schedule(sched)
      assert GlobalCounters.global_ops < 4*DSET, f"too many ops {GlobalCounters.global_ops} != {4*DSET}"
    np.testing.assert_allclose(real_index, X.numpy())
  @unittest.skip("not ready")
@@ -229,9 +234,10 @@ class TestIndexing(unittest.TestCase):
    xq = xq.reshape(bs, seqlen, n_heads, head_dim)
    xq_rope, _ = apply_rotary_emb(xq, xq, freqs_cis)
    xq_rope.sum().backward()
-    linear = compile_linear(wq.grad.schedule_linear())
-    assert len(linear.src) == 1, f"expected one kernel for backward, got: {len(linear.src)}"
-    bwd_ops = estimate_uop(linear.src[0]).ops
+    sched = wq.grad.schedule()
+    assert len(sched) == 1, f"expected one kernel for backward, got: {len(sched)}"
+    prg = sched[0].lower().prg.p
+    bwd_ops = prg.estimates.ops
    # bfloat16 on non CDNA4 has ~10x ops overhead because of the software emulation
    if dtype == dtypes.bfloat16 and not Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950"): ops_scale = 10
    else: ops_scale = 1
@@ -46,9 +46,9 @@ def run_asm_gemm(a_shape, b_shape, dtype=dtypes.float16, a_shard=None, b_shard=N
      np.testing.assert_allclose(tst.numpy(), ref.numpy(), atol=atol, rtol=rtol)
      np.testing.assert_allclose(a.grad.numpy(), a_ref.grad.numpy(), atol=grad_atol, rtol=grad_rtol)
      np.testing.assert_allclose(b.grad.numpy(), b_ref.grad.numpy(), atol=grad_atol, rtol=grad_rtol)
-    assert tst.allclose(ref, atol=atol, rtol=rtol).item(), "forward mismatch"
-    assert a.grad.allclose(a_ref.grad, atol=grad_atol, rtol=grad_rtol).item(), "grad_a mismatch"
-    assert b.grad.allclose(b_ref.grad, atol=grad_atol, rtol=grad_rtol).item(), "grad_b mismatch"
+    assert tst.allclose(ref, atol=atol, rtol=rtol), "forward mismatch"
+    assert a.grad.allclose(a_ref.grad, atol=grad_atol, rtol=grad_rtol), "grad_a mismatch"
+    assert b.grad.allclose(b_ref.grad, atol=grad_atol, rtol=grad_rtol), "grad_b mismatch"

 def verify_asm_gemm(batch:int, M:int, N:int, K:int, dtype=dtypes.float16, gpus:int=1) -> None:
  run_asm_gemm((batch, M, K), (K, N), dtype=dtype, a_shard=0, b_shard=None, gpus=gpus)
@@ -8,8 +8,8 @@ from test.helpers import not_support_multi_device

 def _check_ast_count(desired_count:int, t:Tensor):
  # NOTE: this has side effect because everything can be scheduled only once
-  schedule = t.schedule_linear()
-  asts = [s for s in schedule.src if s.src[0].op is Ops.SINK]
+  schedule = t.schedule()
+  asts = [s for s in schedule if s.ast.op is Ops.SINK]
  len(asts)
  # NOT SUPPORTED ANYMORE
  #assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"
@@ -28,8 +28,8 @@ class TestMovedConstFolding(unittest.TestCase):
    _check_ast_count(1, Tensor([1.0, 2, 3, 4]) * Tensor.ones(2).pad(((1, 1),)))

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

  def test_cast_padded(self):
@@ -1,5 +1,5 @@
 import unittest
-from tinygrad import Tensor, UOp, GlobalCounters
+from tinygrad import Tensor, UOp
 from tinygrad.dtype import AddrSpace, dtypes
 from tinygrad.uop.ops import KernelInfo, AxisType

@@ -189,7 +189,7 @@ class TestCustomKernel(unittest.TestCase):
    A = Tensor.randn(16, 16).contiguous()
    B = Tensor.empty(16)
    B = Tensor.custom_kernel(B, A, fxn=slice_sum_kernel)[0]
-    self.assertTrue(B.allclose(A.sum(1)).item())
+    self.assertTrue(B.allclose(A.sum(1)))

  def test_gemm(self):
    N = 16
@@ -273,12 +273,12 @@ class TestCustomKernel(unittest.TestCase):
    C, D, _, _ = Tensor.custom_kernel(C, D, A2, B2, fxn=custom_elementwise_addmul_kernel)  # depends on A2 AND B2
    E = (A2 * 3).contiguous()                      # kernel 2: depends only on A2
    result = (C + D + E).sum()                     # kernel 3: custom_addmul, then kernel 4: sum
-    schedule = result.schedule_linear().src
+    schedule = result.schedule()

    # Find the custom_addmul kernel position
    custom_idx = next((i for i, item in enumerate(schedule)
-                       if hasattr(item.src[0], "arg") and hasattr(item.src[0].arg, "name")
-                       and "custom_addmul" in item.src[0].arg.name), None)
+                       if hasattr(item.ast, "arg") and hasattr(item.ast.arg, "name")
+                       and "custom_addmul" in item.ast.arg.name), None)

    self.assertIsNotNone(custom_idx, "custom_addmul kernel not found in schedule")
    self.assertEqual(custom_idx, 3, f"custom_addmul should be at index 3, got {custom_idx}")
@@ -299,8 +299,8 @@ class TestCustomKernel(unittest.TestCase):
    from tinygrad import function
    @function(precompile=True)
    def run(x:Tensor, w:Tensor) -> Tensor:
-      out = Tensor.invalids(*x.shape, dtype=x.dtype)
-      tmp = Tensor.invalids(*x.shape, dtype=x.dtype)
+      out = Tensor.invalid(*x.shape, dtype=x.dtype)
+      tmp = Tensor.invalid(*x.shape, dtype=x.dtype)
      out, tmp = Tensor.custom_kernel(out, tmp, x, w, fxn=custom_add_with_tmp)[:2]
      return out+tmp

@@ -308,36 +308,6 @@ class TestCustomKernel(unittest.TestCase):
    expected = (3+2)*2+2
    assert all(x == expected for x in result), f"expected all {expected}, got {result}"

-  def test_custom_kernel_sched(self, use_custom=False):
-    x = Tensor.arange(32).reshape(8, 4).realize()
-    y = Tensor.empty_like(x)
-    y = Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
-    if use_custom:
-      z = Tensor.empty_like(x)
-      z = Tensor.custom_kernel(y, y.T.T, fxn=custom_add_one_kernel)[0]
-    else: z = y.T.T+1
-    GlobalCounters.reset()
-    z.realize()
-    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).realize()
-    from tinygrad import function
-    @function(precompile=True)
-    def run(x:Tensor) -> Tensor:
-      y = Tensor.invalids(*x.shape, dtype=x.dtype)
-      return Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
-    GlobalCounters.reset()
-    y = run(x[0]).realize()
-    # it's copying the input and the output
-    self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertEqual(y.tolist(), [1, 2, 3, 4])
-
 class TestUOpReduce(unittest.TestCase):
  def test_uop_sum(self):
    a = Tensor([1.0, 2, 3, 4, 5])
@@ -1,9 +0,0 @@
-import unittest
-from tinygrad import Device
-
-class TestDeviceCount(unittest.TestCase):
-  def test_count(self):
-    self.assertGreaterEqual(Device[Device.DEFAULT].count(), 1)
-
-if __name__ == "__main__":
-  unittest.main()
@@ -6,7 +6,7 @@ from tinygrad.tensor import Tensor
 from tinygrad.helpers import Context, from_mv
 from tinygrad.dtype import dtypes
 from tinygrad.engine.jit import MultiGraphRunner
-from tinygrad.engine.realize import run_linear, compile_linear
+from tinygrad.schedule import linear_to_schedule
 from tinygrad.uop.ops import UOp, Ops, buffers

 from test.helpers import needs_second_gpu
@@ -24,7 +24,7 @@ def get_ast(device:str, num_inputs:int) -> UOp:
      fst = [Tensor.randn(BUF_SIZE, dtype=dtypes.int).realize() for _ in range(num_inputs)]
      s = fst[0]
      for i in range(1, num_inputs): s = s.bitwise_xor(fst[i])
-      cached_asts[(device, num_inputs)] = s.schedule_linear().src[-1].src[0]
+      cached_asts[(device, num_inputs)] = s.schedule()[-1].ast
  return cached_asts[(device, num_inputs)]

 def make_buffer(device, size=BUF_SIZE, fill=False):
@@ -44,12 +44,12 @@ def get_buf_uop(buf:Buffer, cache:dict[Buffer,UOp]) -> UOp:
  return cache[buf]

 def make_graph(graph_cls, calls:list[UOp]):
-  linear = compile_linear(UOp(Ops.LINEAR, src=tuple(calls)))
+  linear = UOp(Ops.LINEAR, src=tuple(calls))
  cf = UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(linear,), arg="graph")
  return graph_cls(cf, [])

 def run_schedule(calls:list[UOp]):
-  run_linear(UOp(Ops.LINEAR, src=tuple(calls)))
+  for ei in linear_to_schedule(UOp(Ops.LINEAR, src=tuple(calls))): ei.lower().run({})

 def zero_bufs(bufs):
  for b in bufs:
@@ -4,7 +4,6 @@ import numpy as np

 from hypothesis import given, settings, strategies as strat
 from test.helpers import assert_jit_cache_len, call_is_graph, not_support_multi_device, needs_second_gpu
-from tinygrad import Variable
 from tinygrad.tensor import Tensor
 from tinygrad.engine.jit import TinyJit, JitError, graph_class
 from tinygrad.device import Device
@@ -40,19 +39,6 @@ class TestJit(unittest.TestCase):
    def add(a, b): return (a+b).realize()
    _simple_test(add)

-  @unittest.skipUnless(Device.DEFAULT == "CPU", "core_id is a CPU runtimevar")
-  def test_hcq_core_id_runtimevar_merge(self):
-    N = 262144
-    @TinyJit
-    def f(x, st):
-      y = (x + 1).contiguous().realize()
-      z = x.shrink(((st, st + N),)).contiguous().realize()
-      return y, z
-    x = Tensor.arange(2*N).contiguous().realize()
-    for _ in range(3): y, z = f(x, Variable("a", 0, N).bind(0))
-    self.assertEqual(y.shape, (2*N,))
-    self.assertEqual(z.shape, (N,))
-
  def test_jitbeam_triggers_beam(self):
    from unittest.mock import patch
    from tinygrad.helpers import getenv as _getenv
@@ -106,20 +92,6 @@ class TestJit(unittest.TestCase):
      np.testing.assert_allclose(e.numpy(), a.numpy()*b.numpy(), atol=1e-4, rtol=1e-5)
    assert_jit_cache_len(f, 3)

-  def test_global_counters_jit(self):
-    @TinyJit
-    def f(a, b):
-      c = (a + b).realize()
-      d = (c * 2).realize()
-      return (d - a).realize()
-    a, b = Tensor.randn(64, 64).realize(), Tensor.randn(64, 64).realize()
-    for _ in range(4):
-      GlobalCounters.reset()
-      f(a, b)
-      Device[a.device].synchronize()
-      self.assertGreater(GlobalCounters.global_mem, 0)
-      self.assertGreater(GlobalCounters.global_ops, 0)
-
  def test_nothing_jitted(self):
    @TinyJit
    def add(a, b): return None
@@ -734,7 +706,7 @@ class TestJitGraphSplit(unittest.TestCase):
        inner_cnt = len(ast.src[0].src)
        assert inner_cnt == expected["cnt"], f"Expected {expected['cnt']} operations in graph, got {inner_cnt}"
      elif expected["type"] == "comp":
-        assert ast.op in (Ops.SINK, Ops.PROGRAM), f"Expected kernel, got {ast.op}"
+        assert ast.op in (Ops.SINK, Ops.PROGRAM, Ops.BEAM), f"Expected kernel, got {ast.op}"
      elif expected["type"] in ("copy", "xfer"):
        assert ast.op is Ops.COPY, f"Expected COPY, got {ast.op}"

@@ -1,12 +1,12 @@
 import numpy as np
 import unittest
+from dataclasses import replace

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

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

 class TestLinearizer(unittest.TestCase):
  def test_arg_dedup(self):
@@ -25,9 +25,9 @@ class TestLinearizer(unittest.TestCase):
    a, b = Tensor.randn(4).realize(), Tensor.randn(4).realize()
    np_a, np_b = a.numpy(), b.numpy()
    c = ((a.shrink(((0, 2),)) - a.shrink(((2, 4),))) - (b.shrink(((0, 2),)) - b.shrink(((2, 4),))))
-    linear = c.schedule_linear()
-    run_linear(linear)
-    rawbufs = [s.buffer for s in linear.src[-1].src[1:] if s.op is not Ops.BIND]
+    sched = c.schedule()
+    for si in sched: si.run()
+    rawbufs = sched[-1].bufs
    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)
@@ -45,7 +45,7 @@ class TestLinearizer(unittest.TestCase):
    tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
    out = tst.neg().cast(dtypes.char).cast(dtypes.int).cast(dtypes.char) * 2
    ast = helper_linearizer_opt(out)
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)

  @unittest.expectedFailure
@@ -53,7 +53,7 @@ class TestLinearizer(unittest.TestCase):
    tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
    out = tst.neg().cast(dtypes.char).cast(dtypes.int) * 2
    ast = helper_linearizer_opt(out)
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
    self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)

  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx")
@@ -63,7 +63,7 @@ class TestLinearizer(unittest.TestCase):
    b = Tensor.empty(16)
    out = img.conv2d(w, b)
    ast = helper_linearizer_opt(out)
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
    # slice at the last loop end
    uslice = [i for i,u in enumerate(uops) if u.op == Ops.END][-1]
    # only valid test if outermost range is the reduce
@@ -84,7 +84,7 @@ class TestLinearizer(unittest.TestCase):
    a = Tensor.randn(2, ).realize()
    out = a.reshape(2, 1).expand(2, 3).sum()
    ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)).sum()])
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    assert len(ranges) == 1 # NOTE: it collapses now

@@ -92,7 +92,7 @@ class TestLinearizer(unittest.TestCase):
    a = Tensor.randn(2, ).realize()
    out = a.reshape(2, 1).expand(2, 3).expand(2, 2, 3).sum()
    ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)), (2, 2, 3)).sum()])
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    assert len(ranges) == 1 # NOTE: it collapses now

@@ -100,7 +100,7 @@ class TestLinearizer(unittest.TestCase):
    a = Tensor([2, 2]).realize()
    out = a.reshape(2, 1).pad(((1, 1), (1, 1)), value=2).sum()
    ast = helper_linearizer_opt(out, wanna_output=[24])
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    # RANGE -> ALU -> RANGE -> ALU + LOAD -> STORE
    assert any(x.op in GroupOp.ALU for x in uops[ranges[0]:ranges[1]])
@@ -113,7 +113,7 @@ class TestLinearizer(unittest.TestCase):
    b = Tensor.randn(1, 1).realize()
    out = (a + b[0]).sum() + b[0]
    ast = helper_linearizer_opt(out, wanna_output=[(a.numpy()+b.numpy()[0]).sum()+b.numpy()])
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    # LOAD -> RANGE -> LOAD -> STORE
    assert len([x for x in uops[:ranges[0]] if x.op is Ops.LOAD]) == 1
@@ -123,7 +123,7 @@ class TestLinearizer(unittest.TestCase):
    b = Tensor.randn(1, 1).realize()
    out = (a.reshape(2, 1).expand(2, 3) + b[0]).sum() + b[0]
    ast = helper_linearizer_opt(out, wanna_output=[(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)) + b.numpy()[0]).sum() + b.numpy()])
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    assert len(ranges) == 1 # NOTE: it collapses now

@@ -134,8 +134,7 @@ class TestLinearizer(unittest.TestCase):
    # these are of size 3 to avoid float4 coalesce
    r = a[:-1] + a[1:]

-    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
-                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).uops
    num_loads = len([uop for uop in uops if uop.op is Ops.LOAD])
    assert num_loads <= 4, "more load uops than needed"
    assert num_loads >= 4, "unexpected number of uops, maybe this test needs updating?"
@@ -147,8 +146,7 @@ class TestLinearizer(unittest.TestCase):
    a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
    r = a.expand([2]) + b.expand([2])

-    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
-                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).uops
    num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
    assert num_ops <= 1, "more alu uops than needed"

@@ -157,8 +155,8 @@ class TestLinearizer(unittest.TestCase):
    x, w = Tensor.randn((1,1,3)).realize(), Tensor.randn((1,1,2)).realize()
    r = Tensor.conv2d(x,w,padding=1).relu()

-    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0],
-      [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]),
+                       renderer=Device[Device.DEFAULT].renderer).uops
    accs = [u for u in uops if u.op is Ops.DEFINE_REG]
    stores = [u for u in uops if u.op is Ops.STORE]
    assert len(accs) == 0  # it's removed now
@@ -170,9 +168,8 @@ class TestLinearizer(unittest.TestCase):
  @unittest.skipUnless(Device.DEFAULT == "CPU", "test only for CPU")
  def test_upcast_with_locals_cpu(self):
    out = Tensor.ones(64,64).contiguous() @ Tensor.ones(64,64).contiguous()
-    prg = to_program(replace_opts(out.schedule_linear().src[-1].src[0], [Opt(OptOps.LOCAL, axis=0, arg=4)]),
-                      renderer=Device[Device.DEFAULT].renderer)
-    self.assertEqual(len(prg.src[3].arg.split("for")), 5)
+    prg = get_program(replace_opts(out.schedule()[-1].ast, [Opt(OptOps.LOCAL, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).uops
+    self.assertEqual(len(prg.src.split("for")), 5)

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@@ -182,9 +179,9 @@ class TestLinearizer(unittest.TestCase):
    x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
    r = (x@y).relu()
    opts_to_apply = [Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
-    program = to_program(replace_opts(r.schedule_linear().src[-1].src[0], opts_to_apply), renderer=Device[Device.DEFAULT].renderer)
+    program = get_program(replace_opts(r.schedule()[-1].ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer)

-    stores = [u for u in tuple(program.src[2].src) if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
+    stores = [u for u in program.uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]

    # the first store is to lds and can be upcasted
    assert stores[0].src[1].dtype == dtypes.float.vec(4)
@@ -196,8 +193,7 @@ class TestLinearizer(unittest.TestCase):
  def test_zero_fold(self):
    a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
    r = Tensor.stack(a, b)
-    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], [Opt(op=OptOps.UPCAST, axis=0, arg=0)]),
-                       renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).uops
    num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
    assert num_ops == 0, "more alu uops than needed"

@@ -206,16 +202,16 @@ class TestLinearizer(unittest.TestCase):
      (dtypes.bool, dtypes.int), (dtypes.int16, dtypes.int), (dtypes.float16, dtypes.float), (dtypes.bfloat16, dtypes.float)):
      if is_dtype_supported(tensor_dtype) and is_dtype_supported(acc_dtype):
        a = Tensor([1, 2, 3], dtype=tensor_dtype).sum()
-        realized_ast = a.schedule_linear().src[-1].src[0]
-        program = to_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
-        local = [uop for uop in tuple(program.src[2].src) if uop.op is Ops.DEFINE_REG]
+        realized_ast = a.schedule()[-1].ast
+        program = get_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
+        local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG]
        assert local[0].dtype.base == acc_dtype

  def test_arg_acc_dtype(self):
    def helper_arg_acc_dtype(c: Tensor, expected_dtype:DType):
-      realized_ast = c.schedule_linear().src[-1].src[0]
-      program = to_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
-      local = [uop for uop in tuple(program.src[2].src) if uop.op is Ops.DEFINE_REG]
+      realized_ast = c.schedule()[-1].ast
+      program = get_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
+      local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG]
      self.assertEqual(local[0].dtype.base, expected_dtype)

    tests = (
@@ -242,7 +238,7 @@ class TestLinearizer(unittest.TestCase):
    opt = [Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4)]
    ast = helper_linearizer_opt(r, [opt])
    # the uops graph is DEFINE_REG -> 4x STORE 0.0 -> RANGE -> 4x ALU -> 4x STORE -> ENDRANGE
-    uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = get_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).uops
    begin_range = [i for i, x in enumerate(uops) if x.op is Ops.RANGE][-1]
    end_range = [i for i, x in enumerate(uops) if x.op is Ops.END][0]
    for i,u in enumerate(uops): print(i, u.op, [uops.index(s) for s in u.src], u.arg, u.dtype)
@@ -262,7 +258,7 @@ class TestLinearizer(unittest.TestCase):
    # shrink so that the dims do not collapse
    t = Tensor.ones(5, 6, 7).contiguous().realize().shrink(((0, 4), (0, 5), (0, 6)))
    ast = helper_linearizer_opt(t+1)
-    uops = tuple(to_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
    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].arg) == ('gidx0', 6), idxs[0]
@@ -271,10 +267,10 @@ class TestLinearizer(unittest.TestCase):

  def test_sum_collapse(self):
    t = Tensor([2]).reshape(1, 1).expand(256, 256).sum()
-    sched = [si for si in t.schedule_linear().src if si.src[0].op is Ops.SINK]
+    sched = [si for si in t.schedule() if si.ast.op is Ops.SINK]
    # sum_collapse is a full collapse now
    assert len(sched) == 1
-    assert not any(u.op is Ops.REDUCE and len(u.arg[1]) > 0 for u in sched[0].src[0].toposort()), "found reduce in sum collapse"
+    assert not any(u.op is Ops.REDUCE_AXIS for u in sched[0].ast.toposort()), "found reduce in sum collapse"
    #lin = Kernel(sched[0].ast)
    #assert not any(u.op is Ops.RANGE for u in lin.linearize().uops), "found loop in sum collapse"

@@ -290,17 +286,18 @@ class TestLinearizer(unittest.TestCase):
    a = Tensor.ones(4, 4).contiguous().realize()
    b = a.shrink(((1, 2), None)).pad(((1, 2), None))
    a.assign(b.where(2, a))
-    linear, var_vals = a.linear_with_vars()
-    assert len(linear.src) == 1
-    run_linear(linear, var_vals)
+    sched = a.schedule()
+    assert len(sched) == 1
+    sched_copy = sched[:]
+    run_schedule(sched)
    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)
-    assert not any(u.op == Ops.WHERE for u in tuple(program.src[2].src)), "found where where where should be folded"
+    program = get_program(replace_opts(sched_copy[-1].ast, []), renderer=Device[Device.DEFAULT].renderer)
+    assert not any(u.op == Ops.WHERE for u in program.uops), "found where where where should be folded"

  def test_phi_simplification(self):
    def helper(t, max_ops=0):
      ast = helper_linearizer_opt(t)
-      uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src)
+      uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer).uops
      # ignore kernel optimized IF statements for now
      if if_op:=next((u for u in uops if u.op is Ops.IF), None):
        uops = uops[:uops.index(if_op)]
@@ -332,7 +329,7 @@ class TestLinearizer(unittest.TestCase):
    out = x.matmul(y)
    with Context(TC=0):
      ast = helper_linearizer_opt(out)
-      uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src)
+      uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer).uops
    # check that the float4 cast collapses
    store_vals = [u.src[1] for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
    for val in store_vals:
@@ -343,8 +340,8 @@ class TestLinearizer(unittest.TestCase):
    x = Tensor.randn((4,3,6,6)).realize()
    out = x.flip((0,1)).contiguous()
    ast = helper_linearizer_opt(out)
-    store_val = [u.src[1] for u in tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src) if u.op is Ops.STORE][0]
-    assert store_val.dtype == dtypes.float.vec(4) and store_val.op is not Ops.STACK
+    store_val = [u.src[1] for u in get_program(ast, renderer=Device[Device.DEFAULT].renderer).uops if u.op is Ops.STORE][0]
+    assert store_val.dtype == dtypes.float.vec(4) and store_val.op is not Ops.VECTORIZE

  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
@@ -356,7 +353,7 @@ class TestLinearizer(unittest.TestCase):
            Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 2)] # upcast accs in both reduces
    ast = helper_linearizer_opt(out, opts=[opt])
    def get_recursive(uop): return set.union(set(uop.src), [uop], *[get_recursive(v) for v in uop.src])
-    uops = tuple(to_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = get_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).uops
    local_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_LOCAL for x in get_recursive(u.src[0]))]
    global_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.PARAM for x in get_recursive(u.src[0]))]
    barrier = [u for u in uops if u.op is Ops.BARRIER]
@@ -377,7 +374,7 @@ class TestLinearizer(unittest.TestCase):
    x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
    r = (x@y).relu()
    ast = helper_linearizer_opt(r)
-    uops = tuple(to_program(ast, renderer=Device[Device.DEFAULT].renderer).src[2].src)
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer).uops
    stores = [u for u in uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]

    # the float4 value stores directly in lds and we skip upcast
@@ -391,18 +388,15 @@ class TestLinearizer(unittest.TestCase):

 def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
  if isinstance(r, Tensor): r = [r]
-  linear, var_vals = Tensor.linear_with_vars(*r)
-  run_linear(UOp(Ops.LINEAR, src=linear.src[:-1]), var_vals)  # run all kernels except the last one
-  last_call = linear.src[-1]
-  ast = last_call.src[0]
-  assert ast.op is Ops.SINK, f"helper_realized_ast expects a SINK {last_call}"
-  last_bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]
-  # now all input buffers in last_call should be realized
+  s = Tensor.schedule(*r)
+  run_schedule(s[:-1])  # run all kernels except the last one
+  assert s[-1].ast.op is Ops.SINK, f"helper_realized_ast expects a SINK {s[-1]}"
+  # now all input buffers in s[-1] 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)]
+  bufs = [Buffer(x.device, x.size, x.dtype).allocate() if i < len(s[-1].ast.src) else x for i,x in enumerate(s[-1].bufs)]
  # ensure buffers are allocated
  for b in bufs: b.ensure_allocated()
-  return ast, bufs
+  return s[-1].ast, bufs

 def helper_linearizer_ast(ast:UOp, inputs:list[Tensor], *args, **kwargs):
  assert isinstance(ast, UOp), "ast must be UOp"
@@ -424,28 +418,30 @@ def reset_bufs(bufs:list[Buffer]):
 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=[]):
  outbufs = real_bufs[:len(realized_ast.src)]
+  device = real_bufs[0].device
  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]
-  for u,b in zip(buf_uops, real_bufs): buffers[u] = b

-  def run_prg(opts):
+  def get_prg(opts):
    ast = realized_ast if opts is None else replace_opts(realized_ast, list(opts))
-    run_linear(UOp(Ops.LINEAR, src=(ast.call(*buf_uops),)))
+    return CompiledRunner(replace(get_program(ast, renderer=Device[Device.DEFAULT].renderer), device=device))

  def check_opt(opts):
+    prg = get_prg(opts=opts)
    reset_bufs(outbufs)
-    run_prg(opts)
+    prg.exec(real_bufs)
    for x,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(x, want, atol=atol, rtol=rtol)

  # Get baseline if it is not provided, which is not optimized at all.
-  run_prg(opts=())
+  prg = get_prg(opts=())
+  prg.exec(real_bufs)
  if len(wanna_output) == 0: wanna_output = copyout_outputs(outbufs)
  else:
    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.
+  prg = get_prg(opts=None)
  reset_bufs(outbufs)
-  run_prg(opts=None)
+  prg.exec(real_bufs)
  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)
@@ -6,7 +6,7 @@ import unittest
 from tinygrad import Device, dtypes
 from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
 from tinygrad.codegen.opt.search import Opt, OptOps
-from tinygrad.codegen import to_program
+from tinygrad.engine.realize import get_program

 class TestLinearizerFailure(unittest.TestCase):
  @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
@@ -25,7 +25,7 @@ class TestLinearizerFailure(unittest.TestCase):
    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)
+    _ = get_program(ast, Device["METAL"].renderer)

 if __name__ == '__main__':
  unittest.main()
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	72fccb5f5c	cleaner	2026-04-22 15:55:19 +08:00
geohot	6077327b80	chat class	2026-04-22 10:46:10 +08:00
geohot	79079ca3a9	add olmo3 support to llm	2026-04-21 16:36:08 +08:00
				`@@ -1 +0,0 @@`
				Run `./render.sh` whenever you update tinyspec.tex to regenerate tinyspec.pdf.