kimi linear working

llm: reuse exact recurrent prefixes
llm: fully warm recurrent serving at startup
2026-08-16 16:18:27 +00:00 · 2026-08-10 15:37:20 +00:00 · 2026-08-10 14:06:15 +00:00 · 2026-08-10 14:06:15 +00:00 · 2026-08-10 14:06:15 +00:00 · 2026-08-10 14:06:15 +00:00
104 changed files with 4916 additions and 1130 deletions
@@ -94,6 +94,7 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -148,6 +149,7 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -200,6 +202,7 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -249,6 +252,7 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -0,0 +1,213 @@
+name: Platform Tests
+env:
+  # increment this when downloads substantially change to avoid the internet
+  CACHE_VERSION: '19'
+  CAPTURE_PROCESS_REPLAY: ${{ github.event_name == 'pull_request' && contains(github.event.pull_request.title, '[pr]') && '1' || '0' }}
+  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+  PYTHONPATH: ${{ github.workspace }}
+  CHECK_OOB: 1
+
+on:
+  push:
+    branches:
+      - master
+  pull_request:
+  workflow_dispatch:
+
+concurrency:
+  group: platform-${{ github.event_name }}-${{ github.event_name == 'pull_request' && github.event.pull_request.number || github.run_id }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
+jobs:
+
+# ****** OSX Tests ******
+
+  unittestmacos:
+    name: MacOS (unit)
+    runs-on: macos-26
+    timeout-minutes: 20
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: unittest-macos
+        deps: testing_unit
+    - name: Run unit tests
+      run: DEV=METAL python -m pytest -n=auto test/unit/ --durations=20
+    - name: Test tensor core ops (fake)
+      run: DEV=METAL DEBUG=3 TC=2 python test/backend/test_ops.py TestOps.test_gemm
+    - name: Test tensor core ops (real)
+      run: DEV=METAL DEBUG=3 python test/backend/test_ops.py TestOps.test_big_gemm
+    - name: Test Beam Search
+      run: DEV=METAL IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
+    - name: Test Device Specific
+      run: DEV=METAL python3 -m pytest test/device/test_metal.py
+    #- name: Fuzz Test linearizer
+    #  run: DEV=METAL DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
+  unittestmacosmock:
+    name: MacOS (unit, mock)
+    runs-on: macos-26
+    timeout-minutes: 20
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: unittest-macos-mock
+        deps: testing_unit
+        amd: 'true'
+        ocelot: 'true'
+    - name: Run NULL backend tests
+      run: SPEC=2 DEV=NULL python -m pytest -n=auto test/null/ --durations=20
+    - name: Run pytest (amd)
+      env:
+        DEV: MOCKKFD+AMD
+        FORWARD_ONLY: 1
+      run: |
+        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
+    - name: Run pytest (ptx)
+      env:
+        DEV: "MOCK+NV:PTX"
+        FORWARD_ONLY: 1
+        # TODO: failing due to library loading error
+        CAPTURE_PROCESS_REPLAY: 0
+      run: |
+        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py \
+          test/testextra/test_hevc.py::TestHevc::test_hevc_decode_compile --durations=20
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
+  testmetal:
+    strategy:
+      fail-fast: false
+      matrix:
+        group: [1, 2]
+    name: MacOS (DEV=METAL) (${{ matrix.group }})
+    runs-on: macos-26
+    timeout-minutes: 20
+    env:
+      DEV: METAL
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: macos-metal
+        deps: testing_unit
+    - name: Check Device.DEFAULT and print some source
+      run: |
+        python -c "from tinygrad import Device; assert Device.DEFAULT == 'METAL'"
+        DEBUG=4 python test/test_tiny.py TestTiny.test_plus
+    - name: Run backend tests
+      run: python -m pytest -n=auto test/backend --durations=20 --splits 2 --group ${{ matrix.group }}
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
+  testmacos:
+    strategy:
+      fail-fast: false
+      matrix:
+        dev:
+          - 'CPU:CLANG'
+          - 'CPU:LLVM'
+          - 'CPU:LVP'
+          - 'WEBGPU'
+
+    name: MacOS (DEV=${{ matrix.dev }})
+    runs-on: macos-26
+    timeout-minutes: 20
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: macos-${{ matrix.dev }}
+        deps: "testing_unit${{ contains(matrix.dev, 'LVP') && ' mesa' || '' }}"
+        llvm: ${{ contains(matrix.dev, 'LLVM') || contains(matrix.dev, 'LVP') }}
+        webgpu: ${{ matrix.dev == 'WEBGPU' }}
+    - name: Set env
+      run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
+    - name: Check Device.DEFAULT and print some source
+      run: |
+        python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
+        DEBUG=4 python test/test_tiny.py TestTiny.test_plus
+    - name: Run test_tiny
+      run: python -m pytest -n=auto test/test_tiny.py --durations=20
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
+# ****** Windows Tests ******
+
+  testwindows:
+    strategy:
+      fail-fast: false
+      matrix:
+        dev:
+          - 'CPU:CLANG'
+          - 'CPU:LLVM'
+          - 'CPU:X86'
+          - 'WEBGPU'
+
+    name: Windows (DEV=${{ matrix.dev }})
+    runs-on: windows-2025
+    timeout-minutes: 15
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v6
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: windows-${{ matrix.dev }}-minimal
+          deps: testing_unit
+          pydeps: ${{ matrix.dev == 'WEBGPU' && 'dawn-python' || '' }}
+      - name: Set env
+        shell: bash
+        run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
+      - name: Check Device.DEFAULT and print some source
+        shell: bash
+        run: |
+          python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
+          DEBUG=4 python test/test_tiny.py TestTiny.test_plus
+      - name: Run test_tiny
+        shell: bash
+        run: python -m pytest -n=auto test/test_tiny.py --durations=20
+
+
+  qcomclcompiletests:
+    name: Compile-only (QCOM CL)
+    runs-on: ubuntu-24.04-arm
+    timeout-minutes: 15
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v6
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: compile-qcomcl
+          deps: testing_unit
+          tinydreno: 'true'
+      - name: Set env
+        shell: bash
+        run: printf "DEV=NULL:QCOMCL:a630\nNULL_ALLOW_COPYOUT=1" >> $GITHUB_ENV
+      - name: Run test_ops
+        shell: bash
+        run: |
+          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
+          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20
+      - name: Run test_ops (IMAGE)
+        shell: bash
+        env:
+          IMAGE: 1
+          DEV: "NULL:QCOMCL:a630,IMAGE_PITCH_ALIGNMENT=64"
+        run: |
+          DEBUG=4 python test/backend/test_ops.py TestOps.test_gemm | grep read_imagef
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20
@@ -294,7 +294,7 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1361 ALLOWED_GATED_READ_IMAGE=54 FLOAT16=1 DEV="CL::IMAGE_PITCH_ALIGNMENT=64" IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1361 ALLOWED_GATED_READ_IMAGE=38 FLOAT16=1 DEV="CL::IMAGE_PITCH_ALIGNMENT=64" IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
          # IMAGE_PITCH_ALIGNMENT=64 matches adreno 630
      - name: Test openpilot CL compile fp32 (test correctness)
        run: |
@@ -585,8 +585,11 @@ jobs:
        run: |
          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['AMD'], Device.DEFAULT"
          DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
+      - name: Run MXFP4 Llama training on NULL backend
+        if: ${{ matrix.backend == 'amd' && matrix.arch == 'gfx950' }}
+        run: PYTHONPATH=. DEV=NULL:HIP:gfx950 MXFP4=1 LLAMA_LAYERS=2 BENCHMARK=3 NULL_ALLOW_COPYOUT=1 NO_HIPCC=1 ROCM_PATH=/opt/rocm JITBEAM=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh
      - name: Run pytest (amd)
-        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM --durations=20
+        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM test/opt/test_tensor_cores.py --durations=20
      - name: Run disk copy tests
        run: python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk
      - name: Run TRANSCENDENTAL math
@@ -629,165 +632,6 @@ jobs:
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

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

  compiletests:
@@ -824,33 +668,3 @@ jobs:
        run: |
          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_gemm | grep image_load
          python -m pytest -n=auto test/backend/test_ops.py --durations=20
-  qcomclcompiletests:
-    name: Compile-only (QCOM CL)
-    runs-on: ubuntu-24.04-arm
-    timeout-minutes: 15
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v6
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: compile-qcomcl
-          deps: testing_unit
-          tinydreno: 'true'
-      - name: Set env
-        shell: bash
-        run: printf "DEV=NULL:QCOMCL:a630\nNULL_ALLOW_COPYOUT=1" >> $GITHUB_ENV
-      - name: Run test_ops
-        shell: bash
-        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
-          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
-          python -m pytest -n=auto test/backend/test_ops.py --durations=20
-      - name: Run test_ops (IMAGE)
-        shell: bash
-        env:
-          IMAGE: 1
-          DEV: "NULL:QCOMCL:a630,IMAGE_PITCH_ALIGNMENT=64"
-        run: |
-          DEBUG=4 python test/backend/test_ops.py TestOps.test_gemm | grep read_imagef
-          python -m pytest -n=auto test/backend/test_ops.py --durations=20
@@ -0,0 +1,224 @@
+# Kimi K3 on 8× MI350X
+
+This branch targets text generation directly from the official `moonshotai/Kimi-K3` checkpoint at `/raid/weights/kimi-k3`. It intentionally ignores the vision tower and multimodal projector. The checkpoint remains in its official 96-shard format; the loader never converts, rewrites, or creates a second 1.56 TB copy.
+
+The checked TP8 layout consumes 196.78 GB (183.27 GiB) of text weights per GPU. The compressed MLA cache adds 28.99 GB (27 GiB) per GPU at the full 1,048,576-token context, leaving approximately 62.23 GB of each nominal 288 GB MI350X for execution buffers and allocator overhead. Start much smaller.
+
+## Resume the current optimization session
+
+Work on branch `kimi_slop`. It was cleanly rebased onto `origin/kimi_slop` commit `553bdf68e` on 2026-08-10. The retained K3 commits after that base are `1b3732a6e`, `c6ac4961d`, `1d8620471`, `224bac031`, `f53f0e7e7`, and `2b1b8c22a`; verify the current hashes with `git log` because a later rebase may rewrite them. Before starting any benchmark, check that the worktree is clean and that no model process remains:
+
+```sh
+git status --short --branch
+git log --oneline --decorate -10
+pgrep -af 'tinygrad.llm.cli|benchmark_kimi_k3' || true
+```
+
+The active acceptance target is **more than 100 tok/s decode, more than 200 tok/s prefill, and less than 180 seconds cold startup** on TP8/gfx950. None is currently met. The authoritative official-checkpoint baseline is 389.84 seconds startup, 38.65 tok/s prefill, and 6.25 tok/s decode. The 1.56 TB checkpoint has a measured 6.9 GB/s single-XFS-NVMe read ceiling, giving a roughly 227-second physical cold-read floor; meeting the startup target therefore also requires a faster storage path, not only loader code.
+
+Use the fake-weight, one-layer loop for development. Do not repeatedly load the official checkpoint while optimizing:
+
+```sh
+DEV=AMD python extra/benchmark_kimi_k3_fake.py --mode attention --iterations 30
+DEV=AMD python extra/benchmark_kimi_k3_fake.py --mode block --iterations 30
+PROFILE=1 DEV=AMD python extra/benchmark_kimi_k3_fake.py --mode block --iterations 5
+```
+
+The clean retained baseline is about 0.630 ms per attention layer and 1.37 ms per complete block, with fake initialization taking about 0.9/2 seconds respectively after the rebase. Since K3 has 93 sequential blocks, a 100 tok/s projection requires at most approximately 0.108 ms per complete block. Only run another 96-shard official validation after a candidate produces a large whole-block gain, remains finite and deterministic, and passes a direct numerical comparison. Test one candidate at a time and remove failed experiments before moving on.
+
+The immediate bottleneck is launch and synchronization granularity: an official four-token decode profile contained 6,304 kernel events, while packed expert work was only a small fraction of total GPU time. Continue with whole-component or whole-block fusion/replay work, not isolated expert microkernels. The latest fake-loop A/B retested the previously rejected dual gate/up and weighted-down MFMA prototypes: 1.374 ms baseline versus 1.375 ms fused, so they were removed again. A fused whole-core KDA recurrence was also slower in the exact fake attention gate (0.665 versus 0.633 ms) and must not be restored unchanged.
+
+Preserve these invariants when official validation resumes: use `/raid/weights/kimi-k3` directly, keep all 96 shards byte-for-byte untouched, run only one model process, begin at context 128, verify all eight devices are `gfx950`, and preserve the first failure instead of retrying over it. The most recent preserved official failure from a rejected KDA experiment was the invalid sequence `[198, 163840, 163840, 163840]`; token 163840 is outside the valid vocabulary. The retained path before that experiment produced deterministic in-range replay.
+
+After a synthetic candidate passes, run correctness and performance in this order: NULL gfx950 compile coverage, focused tests with `-n12` where supported, TP8 fake numerical comparison, official context-128 deterministic tokens, load/prefill/decode timing, and then context admission at 4K, 32K, 131K, and 262K. Run `python -m mypy tinygrad/` and `python -m ruff check .` when those tools are installed. Read `tinygrad/viz/README.md` before inspecting rewrite or device profiles.
+
+## Before renting the machine
+
+- Keep the existing 96 shards in `/raid/weights/kimi-k3`; no additional model-sized free space is required. Leave ordinary headroom for logs and temporary files.
+- The host should have roughly 3 TB RAM, in line with AMD's MI350X platform guidance. The loader itself is streaming and must not need checkpoint-sized RAM.
+- Use a recent kernel/ROCm stack supported by the host vendor, although tinygrad uses its own AMD userspace driver when `DEV=AMD`.
+- Clone this exact commit/branch and keep the official checkpoint directory separate from the repository.
+
+Validate the existing directory without modifying it:
+
+```sh
+python examples/kimi_k3_prepare.py /raid/weights/kimi-k3 --context 4096
+```
+
+For a metadata-only preflight, place the official `config.json` and `model.safetensors.index.json` in a directory and run:
+
+```sh
+python examples/kimi_k3_prepare.py /raid/weights/kimi-k3 --metadata-only
+```
+
+## Hardware admission checks
+
+Do these before loading weights. Stop if any device is missing or reports a different architecture.
+
+```sh
+lspci -d 1002:75a0
+amd-smi list
+DEV=AMD DEBUG=2 python - <<'PY'
+from tinygrad import Device
+for i in range(8):
+  dev = Device[f"AMD:{i}"]
+  print(i, dev.arch)
+PY
+```
+
+Expected architecture: `gfx950` on all eight devices. Then run the small TP8 graph tests:
+
+```sh
+python -m pytest test/unit/test_llm_k3.py test/null/test_kimi_k3.py -q -n12
+DEV=NULL:HIP:gfx950 NULL_ALLOW_COPYOUT=1 python -m pytest \
+  test/unit/test_llm_k3.py::TestKimiK3::test_chunked_recurrent_generate -q -n1
+DEV=AMD python examples/kimi_k3_smoke.py --devices 8
+```
+
+The last two commands are deliberately small. They compile CDNA4 kernels and then exercise the complete TP8 topology without loading the checkpoint.
+
+For performance iteration, use the exact-width fake-weight harness before another official load:
+
+```sh
+DEV=AMD python extra/benchmark_kimi_k3_fake.py --mode attention --iterations 20
+DEV=AMD python extra/benchmark_kimi_k3_fake.py --mode block --iterations 20
+```
+
+It retains K3's 7,168-wide residual stream, 12,288-wide KDA state, 96 heads, 128×128 recurrent matrices, TP8 layouts, top-k 16 routing, packed MXFP4 expert shapes, collectives, and decode JIT, but uses one layer and 16 fake experts. Fake attention weights initialize in about 0.9 seconds and the full block in about 3 seconds. The retained path measured 0.630 ms per fake attention layer and 1.367 ms per complete fake block, projecting about 7.87 tok/s across 93 identical blocks versus 6.25 tok/s for the official heterogeneous model. Treat this as a candidate admission benchmark, not a correctness substitute for official weights.
+
+## First official load
+
+Start at a short context so cache allocation and compilation are bounded. The loader reads disk-backed safetensors, TP-shards every destination before realizing it, and drops each source shard/projection immediately afterward.
+
+```sh
+/usr/bin/time -v env DEV=AMD DEBUG=1 python -m tinygrad.llm.cli \
+  --model /raid/weights/kimi-k3 --devices 8 --max_context 128 </dev/null 2>&1 | tee kimi-k3-load.log
+```
+
+Watch host RAM, swap, HBM, temperatures, and XGMI traffic from a second terminal. Do not start with a one-million-token cache. If loading fails, preserve the first exception and the last loader progress line; do not retry with a larger host-side cache.
+
+## Correctness and performance sequence
+
+1. Load with context 128 and generate one token.
+2. Repeat a fixed prompt twice and confirm token-for-token deterministic greedy output.
+3. Compare the first several greedy tokens against the official Transformers implementation at temperature zero.
+4. Benchmark decode only after two warm-up tokens.
+5. Benchmark prefill at 128, 512, 2K, and 8K tokens. Increase context only while HBM and compile time remain healthy.
+6. Use `VIZ=1` plus `python -m tinygrad.viz.cli` to inspect kernels; use `VIZ=2` only for short SQTT captures because it adds overhead.
+
+Example decode benchmark:
+
+```sh
+DEV=AMD DEBUG=1 python -m tinygrad.llm.cli --model /raid/weights/kimi-k3 \
+  --devices 8 --max_context 4096 --warmup --benchmark 20
+```
+
+## MI350X validation results (2026-08-10)
+
+The official directory was audited in place: 96 shards, 497,220 indexed tensors, 497,052 language tensors, and 1,560,860,324,864 total bytes. All eight devices reported `gfx950`. No checkpoint file was converted, copied, or modified, and every model run used a single process. The actual text tower is 1,559,965,606,912 bytes; its checked TP8 layout is 196,784,397,312 bytes per GPU.
+
+The preserved first full-checkpoint error was an `A_log` shape mismatch, `(128,) -> (96, 1)`. K3 stores one decay value per 128-wide KDA channel, not one per head. The loader now keeps this field replicated and applies the official channel-wise broadcast. A numerical unit test covers the distinction from the older head-wise Kimi Linear behavior.
+
+Load speed was fixed before generation. The original loader opened thousands of individual expert tensors and independently realized eight strided TP slices. The MI350 path now does the following without changing the checkpoint:
+
+- parses safetensor headers selectively, constructing disk-backed tensors only for the 2,460 non-expert entries consumed by that pass instead of materializing metadata objects for every expert entry twice;
+- copies contiguous axis-zero shards and replicas directly into their final device buffers;
+- reads a replicated tensor once and fans it out over XGMI instead of issuing eight identical direct reads (14.31 GB less RAID traffic);
+- stages an inner-axis tensor once and schedules all eight TP slices together;
+- reads each layer's contiguous 15.72 GB expert region once, reorders its lexicographically stored expert records on GPU 0, and realizes all six packed/scale destinations together;
+- retains only final MultiBuffer identities, drops the reorder graph, and flushes the 15.72 GB staging allocation before the next layer.
+
+One real expert layer leaves exactly 1,965,293,568 bytes resident on each GPU and zero bytes in the GPU-0 allocator cache. Complete context-128 loads measured 527.20 seconds before the final staging cleanup and 490.05/489.59 seconds afterward. Peak host RSS for the unprofiled correctness run was 2.11 GiB with zero swap. RAID variability produced later loads from 489.06 to 532.85 seconds.
+
+The selective-metadata and bounded-GC pass reduced non-expert loading from 125.77 to 57.77 seconds. A subsequent full official context-128 load completed in 411.49 seconds, 78.10 seconds (16.0%) faster than the 489.59-second baseline. It read the 96 shards in place with 1,049,688 KiB peak host RSS and zero swap; no weight payload was converted, copied, or modified. Direct-I/O probes measured approximately 6.9 GB/s aggregate for both one and eight concurrent 1 GiB reads. At that rate the 1.56 TB checkpoint has a roughly 227-second cold-read lower bound, so this RAID cannot meet a true cold sub-three-minute startup regardless of loader overhead.
+
+Expert staging graphs are acyclic and are released by reference counting after each layer, so the loader now suppresses unnecessary cyclic-collector scans only around that loop and restores its prior state on every exit. A quiet context-128 load then completed in 391.54 seconds, 30.14 seconds (7.1%) faster than the immediately preceding 421.68-second run, with 1.04 GiB peak RSS and zero swap, although storage variability contributes to run-to-run timing. The host used for these measurements actually mounts `/raid` from one 3.5 TB XFS NVMe, not a multi-drive RAID; shard 28 has 218 extents and live reads fell to roughly 160 MB/s there. This storage layout, plus the physical checkpoint size, remains the limiting cold-start constraint. The weights were not defragmented, copied, or modified.
+
+The fixed XTML prompt `Reply with exactly: OK` encodes to 93 tokens. After excluding the cold JIT capture from replay comparison, two greedy runs produced the identical eight-token sequence:
+
+```text
+[9545, 59991, 10580, 14404, 9545, 59991, 9545, 59991]
+```
+
+At context 128, steady prefill was 14.32 seconds (6.49 tok/s) and eight-token decode was 2.27 seconds (3.53 tok/s, 283.3 ms/token). The same first tokens remained stable at every admitted context. These rates are much lower than the planning estimates below and should be treated as the current measured baseline.
+
+The retained gfx950 serving pass enables the validated wave64 recurrent prefill kernel with 128-token chunks, uses exact BF16 decode projections, combines the routed/shared final TP partials into one collective, and tiles four adjacent packed-expert outputs during multi-token execution. On the same 93-token prompt, two replay trials produced the identical sequence `[198, 92652, 220, 80225]`. Prefill replay measured 2.418--2.482 seconds (37.47--38.46 tok/s), and eight-token decode measured 1.294 seconds (6.18 tok/s, 161.81 ms/token). Peak RSS was 2.77 GiB with zero swap. The packed prefill tile changes floating-point reduction order: direct official-layer comparison against the original kernel had maximum differences of 0.015625 for gate and 0.0078125 for down, and the end-to-end greedy sequence was stable across replay.
+
+A subsequent gfx950 decode pass split the 7,168-wide replicated BF16 projections across eight waves per 16 output channels and used CDNA4 BF16 MFMA, with one FP32 LDS reduction at the end. It is enabled only for batch-one/token-one replicated projections whose dimensions satisfy the hardware tile; prefill, the FP32 router, and the output-sharded 12,288-wide KDA gate remain unchanged. The official retained path uses it for MLA q-a/kv-a and KDA f-a. Isolated TP8 measurements improved replicated 128/576-output projections by about 16--18%; applying it to the already output-sharded KDA gate was slower and was rejected. Random-shape comparison against the generic graph had maximum/mean absolute BF16 differences of 2.0/0.1114 because the split changes reduction order. Against a serial FP32 accumulation rounded once to BF16, the 7,168-to-1,536 kernel was bit-exact in the tested sample.
+
+The final official context-128 validation loaded in 389.84 seconds with 2.71 GiB peak RSS and zero swap. Two replay trials produced the identical four-token sequence `[198, 59675, 9817, 12519]`; prefill remained 2.406 seconds (38.65 tok/s), while eight-token decode improved to 1.280 seconds (6.25 tok/s, 160.00 ms/token). A one-wave MFMA variant and a full-wave fused decode recurrence were both rejected: the former delivered 6.02 tok/s, and the latter 6.179 tok/s, while both changed the greedy sequence without a useful speed gain.
+
+A final load-first experiment increased the disk-to-HBM io_uring queue depth from one to the 32 existing bounded 2 MiB staging buffers. On a direct 1 GiB read from fragmented shard 28 it measured 6.834 GB/s versus 6.832 GB/s for the original path, so the change was rejected. The subsequent unmodified official 96-shard load completed in 389.48 seconds, confirming both the prior result and the single-NVMe lower bound. Peak RSS was 2.75 GiB with zero swap.
+
+Two direct packed-expert MFMA prototypes were also rejected after that load. A fused gate/up kernel was about 29% faster in isolation at the TP8-local shape, and a routed-down kernel which combined projection, probability weighting, and route reduction measured 1.45 ms versus 2.42 ms in isolation. End-to-end, however, stable replay produced `[198, 2338, 2127, 148297]`, prefill measured 38.87 tok/s, and decode measured 6.263 tok/s. That is indistinguishable from the retained 38.65/6.25 tok/s path while changing floating-point reduction order, so neither kernel was retained.
+
+A whole-core KDA decode experiment fused convolution, Q/K normalization, channel decay, recurrence, RMS normalization, output gating, and four persistent state updates. Its raw kernel replayed in about 109 microseconds per local KDA layer and matched a one-step synthetic reference within `9.77e-4` output and `8.13e-4` state maximum error. The exact-width fake-layer gate caught that it was slower than the retained attention path (0.665 versus 0.633 ms/layer). The already-running official validation was stopped after its first invalid greedy sequence, `[198, 163840, 163840, 163840]`, where 163840 is outside the checkpoint's vocabulary. The kernel was rejected and removed.
+
+| Maximum context | Load | Short-prompt replay | Result |
+|---:|---:|---:|---|
+| 128 | 489.59s | 14.32s | stable 8-token replay |
+| 4,096 | 489.06s | 14.32s | stable replay, zero swap |
+| 32,768 | 532.85s | 14.33s | stable replay, zero swap |
+| 131,072 | 520.91s | 14.37s | stable first token, zero swap |
+| 262,144 | 497.34s | 14.41s | stable first token, zero swap |
+
+These are maximum-context/cache admission tests with the same 93-token prompt, not full-length 32K/131K/262K prefills. The full cache allocation path was exercised, but filling those contexts remains a separate long-running throughput test.
+
+Runtime profiling bracketed four steady decode tokens. It recorded 6,304 kernel events and about 474--478 ms of summed GPU work across the eight devices inside a roughly 1.5-second profiled wall interval. The packed `mxfp4_expert_linear_wave64` kernels accounted for only about 22.5 ms summed; the largest families were small 1,792-wide reductions. This identifies launch/synchronization granularity as the immediate MI350 bottleneck rather than packed-weight bandwidth. `JIT_BATCH_SIZE=64` produced the same original 3.53 tok/s as 32. A gfx950 fused MXFP8 QDQ experiment was bit-exact but slower on the real device (about 95 microseconds versus 57--64 microseconds), so it was rejected. Combining the routed and shared final TP partials removed one collective per routed decode layer and helped raise unprofiled decode to 6.18 tok/s, but the remaining sequential launch boundaries still dominate.
+
+The checkpoint's bundled Transformers code was used as the architectural reference for channel decay and tensor mapping. A full independent Transformers/vLLM token comparison was not run on this host because the required `compressed_tensors`/serving backend is not installed; deterministic tinygrad replay and the numerical KDA, loader-layout, NULL gfx950 compile, and real TP8 smoke tests are the completed correctness gates.
+
+## Known hardware-only gate
+
+The correctness path now consumes packed MXFP4 expert weights directly on gfx950 with a wave64 software-decode kernel, so it does not create selected-expert BF16 weight expansions. MXFP8 activation quantization is still emulated. tinygrad has gfx950/CDNA4 BF16 and FP8 matrix-core support, but this branch does not yet have a hardware-validated native MXFP4×MXFP8 expert GEMM. Expect the first run to be a correctness bring-up, not production throughput. Capture profiles on MI350X before changing the representation: native FP4 work cannot be validated faithfully on the available gfx1100 cards.
+
+Recurrent prefill is fused. The gfx950 wave-parallel kernel was compared directly with the portable graph at the official per-GPU shape through 128 tokens: maximum core/state differences remained below `8e-6`/`1e-6`, outputs were finite, and replay was about 2.7 ms versus about 8 ms for the portable kernel in the isolated test. Full K3 therefore uses 128-token recurrent chunks on gfx950. Chunk size remains part of the numerical configuration because different reduction orders can select different final greedy tokens.
+
+The following serving changes apply to the official K3 path: recurrent-state reset graph capture, direct AMD scalar readback without rebuilding a scheduler graph, materialized gate/up boundaries, separate greedy decode JITs, K3's uncorrected routed probability semantics, gfx950 KDA Q/K/V and exact BF16 partial projections, one combined routed/shared final collective, a gfx950 greedy output-head kernel, the wave64 packed-expert path, and the multi-token four-output packed tile. Software MXFP8 remains in use.
+
+After hardware admission on MI350X, profile before porting those kernels. The likely implementation order is:
+
+1. A native packed MXFP4×MXFP8 grouped expert GEMM using CDNA4 matrix instructions.
+2. A wave64/MFMA KDA Q/K/V decode projection.
+3. Combined routed/shared down-projection TP partials so each layer performs one XGMI all-reduce.
+4. A CDNA4 output-head matvec and router matvec if they remain visible in the profile.
+
+Every port needs a direct numerical comparison with the generic graph and an end-to-end greedy-token comparison before performance measurements. The wave64 packed-expert kernel has compile coverage through `NULL:HIP:gfx950`; numerical and performance validation still require real MI350X hardware. None of the remaining gfx11-only kernels should be enabled on gfx950 by changing only the architecture guard.
+
+## MI350X performance expectation
+
+Treat the first rental as bring-up, not a guaranteed throughput run. The loader reads every official expert tensor once into a transient GPU-0 staging buffer (at most one packed projection), then redistributes TP8 slices over the GPU fabric; it does not generate files or require checkpoint-sized host RAM. A reasonable planning range for the full text model on eight MI350X cards is 3–8 minutes to stream and TP-shard the 1.56 TB checkpoint, 150–400 tok/s for initial short/medium prefill, and 25–60 tok/s decode with the software packed-expert path. After a native CDNA4 MXFP4×MXFP8 grouped expert kernel, wave64/MFMA recurrent projections, and XGMI collective tuning, 500+ tok/s prefill and roughly 80–150 tok/s decode are plausible targets. These ranges are engineering estimates, not measurements.
+
+The nominal HBM bandwidth is not the main uncertainty: eight MI350X devices have enough aggregate bandwidth for K3's active weights. Utilization is limited by 93 sequential layers, small routed projections, and synchronization after TP input-sharded projections. Record actual HBM and XGMI counters before deciding whether the next port should target matrix instructions or collective count.
+
+The official checkpoint also contains MoonViT-V2 and multimodal projector weights. They are skipped by the text loader. Image input remains a separate implementation and validation task.
+
+## Local TP4 performance baseline
+
+The pre-rental benchmark uses the converted `Kimi-Linear-48B-A3B-Instruct-MXFP4-v2` checkpoint on four gfx1100 GPUs. It is a useful regression test for the KDA/MLA/MoE text path, not a projection of K3 throughput on MI350X.
+
+```sh
+DEV=AMD JIT_BATCH_SIZE=64 python extra/benchmark_kimi.py \
+  /raid/models/Kimi-Linear-48B-A3B-Instruct-MXFP4-v2 \
+  --devices 4 --max-context 128 --prompt-tokens 32 --decode-tokens 32 --chunk-size 32
+```
+
+Results from 2026-08-10:
+
+- load from RAID: 44.28s for the 29.27 GB checkpoint
+- first 32-token prefill includes roughly 10s of compilation/capture
+- steady fresh-prompt prefill replay: 0.118s, 270.20 tok/s
+- steady context-32 decode replay: 101.82 tok/s, 9.82 ms/token
+- peak host RSS: 729.9 MiB; swap was not used
+
+The load, prefill, and decode targets are all met in the bounded prompt-32 run. Decode improved from 23.03 tok/s to 101.82 tok/s. The retained greedy output was checked across 32 decode steps; rejected half-wave and unrounded recurrent reductions were faster but diverged and eventually collapsed to a repeated token.
+
+Fully warmed HTTP serving was also measured with `--max_context 4096`. Startup, including weight load, capture, and replay of both serving shapes, took 113.73s. After a two-turn cache test, the first aligned 64-token request reported 271 tok/s prefill and 101 tok/s decode over 64 generated tokens. A 99-token prompt reported 254 tok/s prefill and 99 tok/s decode; decode falls slightly as MLA context grows.
+
+Recurrent serving uses only the captured 32-token prefill graph and captured single-token graph. Warmup uses two consecutive chunks so both initial and nonzero-position prefill execution are ready before the socket opens. A prompt tail shorter than 32 tokens runs through the single-token graph instead of compiling a new static shape, so no request-time JIT capture is required. Exact extensions reuse recurrent and KV state—the live second turn logged `in: 18 + 15`—while divergent prompts reset both safely. Very short prompts can report less than 200 aggregate prefill tok/s because fixed reset and single-token costs dominate; aligned and medium/long prompts exercise the 200+ tok/s prefill path.
+
+Four 7900 XTX cards provide 96 GB aggregate VRAM and about 3.84 TB/s aggregate physical memory bandwidth. Their nominal aggregate vector FP16 rate is about 245.6 TFLOP/s, or about 492 TFLOP/s through matrix instructions. Kimi Linear activates roughly 3.107B parameters per token; a simple active-weight accounting gives approximately 4.05 GB/token and an optimistic bandwidth-only ceiling near 948 tok/s. The measured decode rate is much lower because this MoE decode workload is a collection of small matrix-vector operations plus PCIe collectives, not one ideal streaming kernel.
+
+The generic loader currently rereads logical TP shards and accounts for roughly 227 GB of disk traffic for a TP4 load. RAID bandwidth hides that inefficiency locally, but a direct one-pass shard loader remains worthwhile before slow remote storage is used. It was not retained here because the attempted direct-shard graph exposed an unresolved scheduler/renderer edge; correctness and bounded memory take priority over avoiding the redundant reads.
+
+Different chunk sizes can choose a different final token because their matrix kernels use different floating-point reduction orders. Each measured shape was repeatable between cold and captured execution. For official K3 validation, compare logits/tokens against the reference at one fixed chunk size and greedy settings rather than requiring bitwise agreement between performance shapes.
@@ -0,0 +1,9 @@
+import argparse
+from tinygrad.llm.kimi import convert_kimi
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser(description="Convert official Kimi-Linear-48B-A3B BF16 weights to tinygrad MXFP4/BF16")
+  parser.add_argument("source", help="downloaded moonshotai/Kimi-Linear-48B-A3B-Instruct directory")
+  parser.add_argument("output", help="output directory")
+  args = parser.parse_args()
+  convert_kimi(args.source, args.output)
@@ -0,0 +1,27 @@
+#!/usr/bin/env python3
+"""Cheap preflight for an official moonshotai/Kimi-K3 checkout. Does not load model weights."""
+import argparse, json, pathlib, shutil
+from tinygrad.llm.kimi_k3 import KIMI_K3_TP8_BYTES_PER_GPU, audit_kimi_k3_checkpoint
+
+def main() -> None:
+  parser = argparse.ArgumentParser()
+  parser.add_argument("model_dir", type=pathlib.Path)
+  parser.add_argument("--metadata-only", action="store_true", help="permit absent weight shards")
+  parser.add_argument("--context", type=int, default=4096, help="context length used for the memory estimate")
+  args = parser.parse_args()
+  stats = audit_kimi_k3_checkpoint(args.model_dir, require_shards=not args.metadata_only)
+  if not 1 <= args.context <= 1_048_576: raise ValueError("--context must be between 1 and 1048576")
+
+  # K3 has 24 MLA layers. Each token stores the 512-value compressed latent plus 64 RoPE values in BF16.
+  per_gpu_weights = KIMI_K3_TP8_BYTES_PER_GPU
+  mla_cache = 24 * args.context * (512 + 64) * 2
+  hbm = 288_000_000_000
+  print(json.dumps(stats, indent=2))
+  print(f"exact text weights/GPU under this TP8 layout: {per_gpu_weights/1e9:.2f} GB ({per_gpu_weights/2**30:.2f} GiB)")
+  print(f"replicated MLA cache/GPU at {args.context:,} tokens: {mla_cache/1e9:.2f} GB ({mla_cache/2**30:.2f} GiB)")
+  print(f"nominal MI350X headroom before runtime buffers: {(hbm-per_gpu_weights-mla_cache)/1e9:.2f} GB")
+  if not args.metadata_only:
+    usage = shutil.disk_usage(args.model_dir)
+    print(f"filesystem free space: {usage.free/1e9:.2f} GB")
+
+if __name__ == "__main__": main()
@@ -0,0 +1,27 @@
+#!/usr/bin/env python3
+"""Run a reduced, architecture-complete K3 prefill/decode on tensor-parallel devices."""
+import argparse, time
+from tinygrad import Tensor, Device, dtypes, nn
+from tinygrad.llm.kimi_k3 import _shard_kimi_k3, kimi_k3_smoke_config
+from tinygrad.llm.model import Transformer
+
+def main() -> None:
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--devices", type=int, default=8)
+  args = parser.parse_args()
+  if args.devices not in (1, 2, 4, 8): raise ValueError("the K3 admission smoke test supports 1, 2, 4, or 8 devices")
+  devices = tuple(f"AMD:{i}" for i in range(args.devices))
+  model = Transformer(kimi_k3_smoke_config())
+  for name,value in nn.state.get_state_dict(model).items():
+    fill = 127 if name.endswith("weight_scale") else 0
+    dtype = value.dtype if value.dtype is dtypes.uint8 else dtypes.bfloat16
+    value.replace(Tensor.full(value.shape, fill, dtype=dtype, device="CPU"))
+  _shard_kimi_k3(model, devices)
+  temperature = Tensor([0.0], device=devices)
+  for label,tokens,start in (("prefill", [[1, 2]], 0), ("decode", [[3]], 2), ("decode replay", [[4]], 3)):
+    begin = time.perf_counter()
+    out = model(Tensor(tokens, dtype=dtypes.int32, device=devices), start, temperature).realize()
+    for device in devices: Device[device].synchronize()
+    print(f"{label}: shape={out.shape}, {time.perf_counter()-begin:.3f}s")
+
+if __name__ == "__main__": main()
@@ -1458,7 +1458,8 @@ def train_llama3():

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

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

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

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

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

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

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

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

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

-      loss = sum(losses) / len(losses)
      optim_time = et - gt
      dev_time += optim_time
      step_time = et - st
@@ -114,6 +114,11 @@ def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor,
                             amax_x2:Tensor|None, next_amax_x2:Tensor|None,
                             grad_amax_xw13:Tensor|None, next_grad_amax_xw13:Tensor|None,
                             grad_amax_xout:Tensor|None, next_grad_amax_xout:Tensor|None):
+  if FUSED_SILU_W13 and MXFP4:
+    from extra.llama_kernels.swiglu import swiglu
+    out, *ret = matmul(swiglu(x_w13), w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout,
+                       next_grad_amax_state=next_grad_amax_xout, next_amax_x=next_amax_x2)
+    return out, ret
  if FUSED_SILU_W13 and not MXFP4:
    from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
    x2_fp8 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13,
@@ -1,8 +1,8 @@
 #!/usr/bin/env bash

 export PYTHONPATH="."
-export PATH="/opt/rocm-7.1.1/bin:$PATH"
-export ROCM_PATH="/opt/rocm-7.1.1"
+export ROCM_PATH=${ROCM_PATH:-/opt/rocm-7.1.1}
+export PATH="$ROCM_PATH/bin:$PATH"
 export DEV=${DEV:-AMD}
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -16,7 +16,7 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export FP8=${FP8:-1}
+export MXFP4=${MXFP4:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
 export FAST_CE=${FAST_CE:-1}
 export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
@@ -16,7 +16,7 @@ export USE_ATOMICS=${USE_ATOMICS:-1}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
-export FP8=${FP8:-1}
+export MXFP4=${MXFP4:-1}
 export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
 export FAST_CE=${FAST_CE:-1}
 export FUSED_INPUT_QUANTIZE=${FUSED_INPUT_QUANTIZE:-1}
@@ -1,4 +1,6 @@
 #!/bin/bash
+set -e
+
 export BENCHMARK=${BENCHMARK:-5}
 export EVAL_BS=0
 VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=${DEBUG:--0} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+"""Benchmark Kimi-Linear load, prefill, and decode on its TP4 checkpoint."""
+import argparse, resource, time
+from tinygrad import Device, TinyJit
+from tinygrad.helpers import profile_marker
+from tinygrad.llm.kimi import load_kimi
+
+def sync(devices:int) -> None:
+  for i in range(devices): Device[f"AMD:{i}"].synchronize()
+
+def timed_next(gen, devices:int) -> tuple[int, float]:
+  begin = time.perf_counter()
+  token = next(gen)
+  sync(devices)
+  return token, time.perf_counter()-begin
+
+def fresh_generate(model, prompt:list[int], chunk_size:int):
+  # Force recurrent/KV state reset so repeated runs and chunk sweeps measure the entire prompt,
+  # rather than silently reusing the prefix cached by the previous measurement.
+  model._cached_tokens = [-1] * len(prompt)
+  return model.generate(prompt.copy(), chunk_size=chunk_size)
+
+def main() -> None:
+  parser = argparse.ArgumentParser()
+  parser.add_argument("model", help="converted Kimi-Linear-48B-A3B MXFP4-v2 directory")
+  parser.add_argument("--devices", type=int, default=4)
+  parser.add_argument("--max-context", type=int, default=128)
+  parser.add_argument("--prompt-tokens", type=int, default=32)
+  parser.add_argument("--decode-tokens", type=int, default=8)
+  parser.add_argument("--chunk-size", type=int, default=32)
+  parser.add_argument("--sweep-chunks", help="comma-separated prefill chunk sizes; uses the fastest for decode")
+  args = parser.parse_args()
+  if args.prompt_tokens < 1 or args.prompt_tokens + args.decode_tokens + 1 > args.max_context:
+    raise ValueError("prompt and decode tokens must fit within --max-context")
+
+  begin = time.perf_counter()
+  model = load_kimi(args.model, max_context=args.max_context, devices=args.devices)
+  sync(args.devices)
+  print(f"load: {time.perf_counter()-begin:.3f}s", flush=True)
+
+  prompt = [1] + [1000+i%1000 for i in range(args.prompt_tokens-1)]
+  chunks = [int(x) for x in args.sweep_chunks.split(",")] if args.sweep_chunks else [args.chunk_size]
+  if any(x < 1 or x > args.prompt_tokens for x in chunks): raise ValueError("prefill chunks must be between 1 and --prompt-tokens")
+  timings:list[tuple[float, int]] = []
+  prefill_jits:dict[int, TinyJit] = {}
+  for chunk in chunks:
+    # Recurrent prefill has a static token dimension. Give each swept shape its own capture;
+    # the rollout JIT remains shared and independently benchmarks chunk 1/decode.
+    if chunk != 1: model.prefill_jit = TinyJit(model.forward)
+    cold = fresh_generate(model, prompt, chunk)
+    first, cold_prefill = timed_next(cold, args.devices)
+    print(f"chunk {chunk}: cold prefill {cold_prefill:.3f}s, token={first}", flush=True)
+    warm = fresh_generate(model, prompt, chunk)
+    warm_first, prefill = timed_next(warm, args.devices)
+    if first != warm_first: raise RuntimeError(f"chunk {chunk} is not repeatable: cold={first}, warm={warm_first}")
+    timings.append((prefill, chunk))
+    if chunk != 1: prefill_jits[chunk] = model.prefill_jit
+    print(f"chunk {chunk}: prefill {prefill:.3f}s ({args.prompt_tokens/prefill:.3f} tok/s), token={first}", flush=True)
+
+  prefill, best_chunk = min(timings)
+  if best_chunk != 1: model.prefill_jit = prefill_jits[best_chunk]
+  warm = fresh_generate(model, prompt, best_chunk)
+  first, replay_prefill = timed_next(warm, args.devices)
+  _, cold_decode = timed_next(warm, args.devices)
+  _, capture_decode = timed_next(warm, args.devices)
+  print(f"selected chunk: {best_chunk}; prefill replay {replay_prefill:.3f}s "
+        f"({args.prompt_tokens/replay_prefill:.3f} tok/s), token={first}", flush=True)
+  print(f"cold decode: {cold_decode:.3f}s", flush=True)
+  print(f"capture decode: {capture_decode:.3f}s", flush=True)
+  profile_marker("kimi decode steady start")
+  begin = time.perf_counter()
+  output = [next(warm) for _ in range(args.decode_tokens)]
+  sync(args.devices)
+  decode = time.perf_counter()-begin
+  profile_marker("kimi decode steady end")
+  print(f"decode: {decode:.3f}s ({args.decode_tokens/decode:.3f} tok/s, {decode/args.decode_tokens*1e3:.3f} ms/tok), output={output}", flush=True)
+  print(f"peak RSS: {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024:.1f} MiB", flush=True)
+
+if __name__ == "__main__": main()
@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+"""Bounded correctness and load/prefill/decode benchmark for the official TP8 Kimi K3 checkpoint."""
+import argparse, resource, time
+
+from tinygrad import Device
+from tinygrad.helpers import profile_marker
+from tinygrad.llm.cli import KimiK3Template, SimpleTokenizer
+from tinygrad.llm.kimi_k3 import load_kimi_k3, load_kimi_tokenizer_data
+
+def sync(devices:int) -> None:
+  for i in range(devices): Device[f"AMD:{i}"].synchronize()
+
+def fresh_generate(model, prompt:list[int], chunk_size:int):
+  # Never reuse a prefix or recurrent state across correctness/benchmark trials.
+  model._cached_tokens = [-1] * len(prompt)
+  return model.generate(prompt.copy(), chunk_size=chunk_size, temperature=0.0)
+
+def timed_next(gen, devices:int) -> tuple[int, float]:
+  begin = time.perf_counter()
+  token = next(gen)
+  sync(devices)
+  return token, time.perf_counter()-begin
+
+def main() -> None:
+  parser = argparse.ArgumentParser()
+  parser.add_argument("model", help="official unmodified Kimi K3 checkpoint directory")
+  parser.add_argument("--devices", type=int, default=8)
+  parser.add_argument("--max-context", type=int, default=128)
+  parser.add_argument("--prompt", default="Reply with exactly: OK")
+  parser.add_argument("--stable-tokens", type=int, default=8)
+  parser.add_argument("--decode-tokens", type=int, default=8)
+  parser.add_argument("--chunk-size", type=int, default=128)
+  args = parser.parse_args()
+
+  begin = time.perf_counter()
+  model = load_kimi_k3(args.model, max_context=args.max_context, devices=args.devices)
+  sync(args.devices)
+  load_time = time.perf_counter()-begin
+  print(f"load: {load_time:.3f}s", flush=True)
+
+  normal, special, bos, eos = load_kimi_tokenizer_data(args.model)
+  tok = SimpleTokenizer(normal, special, "kimi-k2", bos_id=bos, eos_id=eos, eot_id=eos)
+  rendered = KimiK3Template().render(messages=[{"role":"user", "content":args.prompt}], add_generation_prompt=True)
+  prompt = tok.encode(rendered)
+  needed = len(prompt) + max(args.stable_tokens, args.decode_tokens+3)
+  if needed > args.max_context: raise ValueError(f"prompt and output need {needed} tokens but max context is {args.max_context}")
+  print(f"prompt: {len(prompt)} tokens, chunk={args.chunk_size}", flush=True)
+
+  sequences:list[list[int]] = []
+  # TinyJit executes uncaptured once, captures the second call, and replays from the third call.
+  # Compare two replay paths rather than capture numerics/timing against replay.
+  for trial in range(4):
+    gen = fresh_generate(model, prompt, args.chunk_size)
+    sequence:list[int] = []
+    prefill = 0.0
+    for step in range(args.stable_tokens):
+      token, elapsed = timed_next(gen, args.devices)
+      sequence.append(token)
+      if step == 0: prefill = elapsed
+    if trial >= 2: sequences.append(sequence)
+    label = ("uncaptured warmup", "capture warmup", "stable trial 1", "stable trial 2")[trial]
+    print(f"{label}: prefill={prefill:.3f}s "
+          f"({len(prompt)/prefill:.3f} tok/s), tokens={sequence}", flush=True)
+  if sequences[0] != sequences[1]: raise RuntimeError(f"greedy output is not repeatable: {sequences}")
+  print(f"stable text: {tok.decode(sequences[0])!r}", flush=True)
+
+  gen = fresh_generate(model, prompt, args.chunk_size)
+  profile_marker("kimi k3 steady prefill start")
+  first, prefill = timed_next(gen, args.devices)
+  profile_marker("kimi k3 steady prefill end")
+  warmup = [timed_next(gen, args.devices)[0] for _ in range(2)]
+  profile_marker("kimi k3 steady decode start")
+  begin = time.perf_counter()
+  output = [next(gen) for _ in range(args.decode_tokens)]
+  sync(args.devices)
+  decode = time.perf_counter()-begin
+  profile_marker("kimi k3 steady decode end")
+  print(f"prefill replay: {prefill:.3f}s ({len(prompt)/prefill:.3f} tok/s), token={first}", flush=True)
+  print(f"decode after warmup {warmup}: {decode:.3f}s ({args.decode_tokens/decode:.3f} tok/s, "
+        f"{decode/args.decode_tokens*1e3:.3f} ms/tok), output={output}", flush=True)
+  print(f"peak RSS: {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024:.1f} MiB", flush=True)
+
+if __name__ == "__main__": main()
@@ -0,0 +1,84 @@
+#!/usr/bin/env python3
+"""Fast exact-shape K3 KDA/layer benchmark using bounded fake weights instead of the 1.56 TB checkpoint."""
+from __future__ import annotations
+import argparse, statistics, time
+from dataclasses import replace
+
+from tinygrad import Device, Tensor, TinyJit, dtypes, nn
+from tinygrad.helpers import profile_marker
+from tinygrad.llm.kimi_k3 import kimi_k3_config
+from tinygrad.llm.model import GatedDeltaNetBlock
+
+def tp_axis(name:str) -> int|None:
+  if "ffn_gate_exps.weight" in name or "ffn_up_exps.weight" in name: return 1
+  if "ffn_gate_exps.weight_scale" in name or "ffn_up_exps.weight_scale" in name: return 1
+  if "ffn_down_exps.weight" in name or "ffn_down_exps.weight_scale" in name: return 2
+  if name.endswith(("ffn_gate_shexp.weight", "ffn_up_shexp.weight")): return 0
+  if name.endswith(("ffn_down_shexp.weight", "ffn_routed_down.weight", "ffn_routed_up.weight", "ssm_out.weight")): return 1
+  if name.endswith(("attn_q.weight", "attn_k.weight", "attn_v.weight", "ssm_g_full.weight", "ssm_f_b.weight", "ssm_beta.weight")): return 0
+  if name.endswith(("ssm_q_conv1d.weight", "ssm_k_conv1d.weight", "ssm_v_conv1d.weight", "ssm_dt.bias")): return 0
+  return None
+
+def fake_value(name:str) -> tuple[int|float, object]:
+  if name.endswith("weight_scale"): return 120, dtypes.uint8
+  if name.endswith("_exps.weight"): return 0x11, dtypes.uint8
+  if name.endswith("ssm_a"): return -0.1, dtypes.float32
+  if name.endswith("ssm_dt.bias"): return 0.1, dtypes.float32
+  if "conv1d.weight" in name: return 0.1, dtypes.float32
+  if name.endswith("exp_probs_b.bias"): return 0.0, dtypes.float32
+  if name.endswith("norm.weight"): return 1.0, dtypes.bfloat16
+  return 0.001, dtypes.bfloat16
+
+def fake_tp_tensor(shape:tuple[int, ...], value:int|float, dtype, devices:tuple[str, ...], axis:int|None) -> Tensor:
+  if axis is not None and shape[axis] % len(devices): raise ValueError(f"shape {shape} is not TP{len(devices)} divisible on axis {axis}")
+  source = Tensor.full(shape, value, dtype=dtype, device=devices[0]).clone().realize()
+  return source.shard(devices, axis=axis).realize()
+
+def sync(devices:tuple[str, ...]) -> None:
+  for device in devices: Device[device].synchronize()
+
+def main() -> None:
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--devices", type=int, default=8)
+  parser.add_argument("--mode", choices=("attention", "block"), default="attention")
+  parser.add_argument("--iterations", type=int, default=20)
+  args = parser.parse_args()
+  devices = tuple(f"AMD:{i}" for i in range(args.devices))
+  # One exact-width KDA layer, but only 16 fake routed experts. This retains top-k 16 and every
+  # official per-GPU matrix/state shape while keeping fake expert storage below 300 MB per layer.
+  config = replace(kimi_k3_config(4), num_blocks=1, num_experts=16, num_experts_per_tok=16, ssm_layers=(True,),
+                   attn_res_block_size=0)
+  block = GatedDeltaNetBlock(config, config.ssm)
+  begin = time.perf_counter()
+  for name,tensor in nn.state.get_state_dict(block).items():
+    if args.mode == "attention" and name.startswith(("ffn_", "exp_probs_")): continue
+    value, dtype = fake_value(name)
+    tensor.replace(fake_tp_tensor(tuple(int(x) for x in tensor.shape), value, dtype, devices, tp_axis(name)))
+  sync(devices)
+  print(f"fake weights: {time.perf_counter()-begin:.3f}s", flush=True)
+  x_source = (((Tensor.arange(config.dim, dtype=dtypes.float32).reshape(1, 1, config.dim) % 31) / 31) \
+    .cast(dtypes.bfloat16).to(devices[0])).clone().realize()
+  x = x_source.shard(devices, axis=None).realize()
+  block._init_state(x)
+  # Use direct buffer-backed state shards. The production path reaches this form after prefill;
+  # the fake harness begins immediately at decode and must not feed lazy clone graphs to TinyJit.
+  for state,axis in ((block.conv_state_q, 2), (block.conv_state_k, 2), (block.conv_state_v, 2), (block.recurrent_state, 1)):
+    state.replace(Tensor.zeros(*state.shape, dtype=state.dtype, device=devices[0]).shard(devices, axis=axis).realize())
+
+  @TinyJit
+  def run(inp:Tensor) -> Tensor:
+    if args.mode == "attention": return block._attention(block.attn_norm(inp), 0).realize()
+    return block(inp, 0).realize()
+
+  # uncaptured, capture, then replay only
+  run(x); sync(devices)
+  run(x); sync(devices)
+  samples:list[float] = []
+  profile_marker(f"fake K3 {args.mode} start")
+  for _ in range(args.iterations):
+    begin = time.perf_counter(); out = run(x); sync(devices); samples.append((time.perf_counter()-begin)*1e3)
+  profile_marker(f"fake K3 {args.mode} end")
+  print(f"{args.mode}: median={statistics.median(samples):.3f} ms/layer, min={min(samples):.3f} ms/layer, "
+        f"projected_93_layer_rate={1000/(statistics.median(samples)*93):.3f} tok/s, finite={out.float().isfinite().all().item()}")
+
+if __name__ == "__main__": main()
@@ -0,0 +1,31 @@
+import argparse, time
+from tinygrad.llm.model import Transformer
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--model", required=True, help="path to gguf model")
+  parser.add_argument("--max-context", type=int, default=8192, help="max context length (default: %(default)s)")
+  parser.add_argument("--prompt-tokens", type=int, default=1024, help="number of prompt tokens (default: %(default)s)")
+  parser.add_argument("--decode-tokens", type=int, default=16, help="number of tokens to decode (default: %(default)s)")
+  parser.add_argument("--chunk-size", type=int, default=32, help="chunk size for prefill (default: %(default)s)")
+  args = parser.parse_args()
+
+  st = time.perf_counter()
+  model, _ = Transformer.from_gguf(args.model, args.max_context)
+  print(f"load {time.perf_counter()-st:.3f}s", flush=True)
+
+  st = time.perf_counter()
+  model.warmup()
+  print(f"warm {time.perf_counter()-st:.3f}s", flush=True)
+
+  prompt = [257] + [1000+i%1000 for i in range(args.prompt_tokens-1)]
+  gen = model.generate(prompt, chunk_size=args.chunk_size)
+  st = time.perf_counter()
+  # first token is time-to-first-token; counted as part of prefill
+  output = [next(gen)]
+  pt = time.perf_counter()
+  print(f"prefill {args.prompt_tokens/(pt-st):.3f} tok/s", flush=True)
+
+  for _ in range(args.decode_tokens): output.append(next(gen))
+  et = time.perf_counter()
+  print(f"decode {args.decode_tokens/(et-pt):.3f} tok/s output {output}", flush=True)
@@ -6,6 +6,7 @@ from tinygrad.renderer import Estimates
 from tinygrad.helpers import getenv, all_same, DEBUG, ceildiv
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8
+from extra.llama_kernels.quantize_mxfp4 import quantize_mxfp4

 TILE_M, TILE_N, TILE_K = 256, 256, 64

@@ -125,6 +126,25 @@ def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp,
  insts = build_kernel(M, N, K, tile_m, tile_n)
  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in insts))))

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

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

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

 # ** main gemm function

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

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

 def _atomic_add(device:str) -> str:
@@ -182,7 +182,8 @@ def sdma_copy(ctx, call):
  src_addr, dst_addr = call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs)
  return call.ins(SDMAOps.COPY, src=tuple(UOp.const(x, dtypes.uint32) for off in range(0, sz, ctx.max_copy_size) for x in (
    ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR),
-    ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0, *data64_le(src_addr+off), *data64_le(dst_addr+off))))
+    ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0,
+    *data64_le(src_addr+UOp.const(off, dtypes.uint64)), *data64_le(dst_addr+UOp.const(off, dtypes.uint64)))))

 def sdma_wait(ctx, ins, dst, val):
  op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \
@@ -507,11 +508,12 @@ class PCIIface(PCIIfaceBase):
    if drain_only: d.iface.dev_impl.ih.drain()
    else: d.iface.dev_impl.ih.interrupt_handler()

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

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

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

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

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

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

  def scratch_buffer(self, private_segment_size):
-    private_segment_size = max(private_segment_size, 128)
+    AMDDevice.max_scratch_psize = private_segment_size = max(private_segment_size, 128, AMDDevice.max_scratch_psize)
    if self.max_private_segment_size < private_segment_size:
      lanes_per_wave = 64 # wave64
      mem_alignment_size = 256 if self.target[0] != 9 else 1024
@@ -0,0 +1,36 @@
+import functools, math, pathlib
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.renderer import Estimates
+from extra.llama_kernels import alloc_like, compile_hip
+
+@functools.cache
+def _custom_quantize_mxfp4(row_fp4:UOp, row_scale:UOp, col_fp4:UOp, col_scale:UOp, x:UOp, *, shuffle_row:bool, shuffle_col:bool) -> UOp:
+  M, N = math.prod(x.shape[:-1]), x.shape[-1]
+  assert M % 256 == 0 and N % 256 == 0, f"MXFP4 quantization requires multiples of 256, got {x.shape}"
+  name = f"quantize_mxfp4_{int(shuffle_row)}_{int(shuffle_col)}_{M}_{N}"
+  mem = M*N*2 + M*N + M*N//16  # read bf16, write row+col fp4 + e8m0
+  outputs = (row_fp4, row_scale, col_fp4, col_scale)
+  sink = UOp.sink(*(o.base for o in outputs), x.base,
+                  *(UOp(Ops.CUSTOM, dtypes.void, (o.base.index(0),), arg="") for o in outputs),
+                  UOp.special(256, "lidx0"), UOp.special(M//128, "gidx0"), UOp.special(N//64, "gidx1"),
+                  arg=KernelInfo(name, estimates=Estimates(ops=12*M*N, mem=mem)))
+  src = (pathlib.Path(__file__).parent/"quantize_mxfp4.cpp").read_text()
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
+    UOp(Ops.BINARY, arg=compile_hip(src, [f"-DKERNEL_NAME={name}", f"-DM_DIM={M}", f"-DN_DIM={N}",
+                                             f"-DSHUFFLE_ROWWISE_FP4_VALUE={int(shuffle_row)}",
+                                             f"-DSHUFFLE_COLWISE_FP4_VALUE={int(shuffle_col)}"]))))
+
+def quantize_mxfp4(x:Tensor, *, shuffle_row:bool=False, shuffle_col:bool=False, flatten_row:bool=False) -> tuple[Tensor, Tensor, Tensor, Tensor]:
+  assert x.dtype == dtypes.bfloat16 and x.ndim >= 2, f"expected BF16 matrix, got {x.dtype} {x.shape}"
+  M, N = math.prod(x.shape[:-1]), x.shape[-1]
+  assert M % 256 == 0 and N % 256 == 0, f"MXFP4 quantization requires multiples of 256, got {x.shape}"
+  axis = x.uop.axis if isinstance(x.device, tuple) else None
+  row_axis = 0 if flatten_row and axis is not None else axis
+  col_axis = None if axis is None else (0 if axis == x.ndim-1 else 1)
+  outputs = (alloc_like((M, N//2) if flatten_row else (*x.shape[:-1], N//2), dtypes.uint8, x.device, row_axis),
+             alloc_like((M, N//32) if flatten_row else (*x.shape[:-1], N//32), dtypes.uint8, x.device, row_axis),
+             alloc_like((N, M//2), dtypes.uint8, x.device, col_axis),
+             alloc_like((N, M//32), dtypes.uint8, x.device, col_axis))
+  fxn = functools.partial(_custom_quantize_mxfp4, shuffle_row=shuffle_row, shuffle_col=shuffle_col)
+  return tuple(Tensor.custom_kernel(*outputs, x, fxn=fxn)[:4])
@@ -0,0 +1,226 @@
+// Copyright (c) 2025-2026, Advanced Micro Devices, Inc. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#include <hip/hip_runtime.h>
+#include <cstdint>
+
+#if !defined(KERNEL_NAME) || !defined(M_DIM) || !defined(N_DIM) || !defined(SHUFFLE_ROWWISE_FP4_VALUE) || \
+    !defined(SHUFFLE_COLWISE_FP4_VALUE)
+#error kernel dimensions and layouts must be defined
+#endif
+
+namespace {
+
+constexpr int BLOCK = 32;
+constexpr int TILE_M = 128;
+constexpr int TILE_N = 64;
+constexpr int THREADS = 256;
+constexpr int THREADS_PER_ROW = 8;
+constexpr int VALUES_PER_THREAD = 4;
+constexpr int SMEM_STRIDE = BLOCK + 2;
+constexpr int M = M_DIM;
+constexpr int N = N_DIM;
+constexpr int M_PACKED = M / 2;
+constexpr int N_PACKED = N / 2;
+constexpr int M_SCALES = M / BLOCK;
+constexpr int N_SCALES = N / BLOCK;
+constexpr bool SHUFFLE_ROWWISE_FP4 = SHUFFLE_ROWWISE_FP4_VALUE;
+constexpr bool SHUFFLE_COLWISE_FP4 = SHUFFLE_COLWISE_FP4_VALUE;
+
+static_assert(M % 256 == 0 && N % 256 == 0);
+
+struct Quantized4 {
+  uint16_t fp4;
+  uint8_t scale;
+};
+
+__device__ __forceinline__ float swizzle_xor1(float value) {
+  float result;
+  asm volatile("ds_swizzle_b32 %0, %1 offset:0x041f\n\ts_waitcnt lgkmcnt(0)" : "=v"(result) : "v"(value));
+  return result;
+}
+
+__device__ __forceinline__ float swizzle_xor2(float value) {
+  float result;
+  asm volatile("ds_swizzle_b32 %0, %1 offset:0x081f\n\ts_waitcnt lgkmcnt(0)" : "=v"(result) : "v"(value));
+  return result;
+}
+
+__device__ __forceinline__ float swizzle_xor4(float value) {
+  float result;
+  asm volatile("ds_swizzle_b32 %0, %1 offset:0x101f\n\ts_waitcnt lgkmcnt(0)" : "=v"(result) : "v"(value));
+  return result;
+}
+
+__device__ __forceinline__ float max8(float value) {
+  value = fmaxf(value, swizzle_xor4(value));
+  value = fmaxf(value, swizzle_xor2(value));
+  return fmaxf(value, swizzle_xor1(value));
+}
+
+__device__ __forceinline__ float4 load_bf16x4(const uint16_t* values) {
+  const uint32_t lo = *reinterpret_cast<const uint32_t*>(values);
+  const uint32_t hi = *reinterpret_cast<const uint32_t*>(values + 2);
+  return make_float4(__uint_as_float(lo << 16), __uint_as_float(lo & 0xffff0000u),
+                     __uint_as_float(hi << 16), __uint_as_float(hi & 0xffff0000u));
+}
+
+__device__ __forceinline__ void hadamard16(float4& value, int lane) {
+  const float a0 = value.x + value.y, a1 = value.x - value.y;
+  const float a2 = value.z + value.w, a3 = value.z - value.w;
+  value = make_float4(a0 + a2, a1 + a3, a0 - a2, a1 - a3);
+
+  const float4 xor1 = make_float4(swizzle_xor1(value.x), swizzle_xor1(value.y), swizzle_xor1(value.z), swizzle_xor1(value.w));
+  value = lane & 1 ? make_float4(xor1.x - value.x, xor1.y - value.y, xor1.z - value.z, xor1.w - value.w)
+                   : make_float4(xor1.x + value.x, xor1.y + value.y, xor1.z + value.z, xor1.w + value.w);
+
+  const float4 xor2 = make_float4(swizzle_xor2(value.x), swizzle_xor2(value.y), swizzle_xor2(value.z), swizzle_xor2(value.w));
+  value = lane & 2 ? make_float4(xor2.x - value.x, xor2.y - value.y, xor2.z - value.z, xor2.w - value.w)
+                   : make_float4(xor2.x + value.x, xor2.y + value.y, xor2.z + value.z, xor2.w + value.w);
+  value.x *= 0.25f;
+  value.y *= 0.25f;
+  value.z *= 0.25f;
+  value.w *= 0.25f;
+}
+
+__device__ __forceinline__ uint8_t e8m0_scale(float amax, float& scale) {
+  if (amax == 0.0f) {
+    scale = 1.0f;
+    return 127;
+  }
+
+  const uint32_t rounded = (__float_as_uint(amax) + 0x200000u) & 0xff800000u;
+  int exponent = static_cast<int>((rounded >> 23) & 0xff) - 129;
+  exponent = exponent < -127 ? -127 : exponent > 127 ? 127 : exponent;
+  scale = exponent == -127 ? __uint_as_float(0x00400000u) : __uint_as_float(static_cast<uint32_t>(exponent + 127) << 23);
+  return static_cast<uint8_t>(exponent + 127);
+}
+
+__device__ __forceinline__ uint16_t pack_fp4(float4 value, float scale) {
+  uint32_t lo = 0, hi = 0;
+  asm volatile("v_cvt_scalef32_pk_fp4_f32 %0, %1, %2, %3" : "+v"(lo) : "v"(value.x), "v"(value.y), "v"(scale));
+  asm volatile("v_cvt_scalef32_pk_fp4_f32 %0, %1, %2, %3" : "+v"(hi) : "v"(value.z), "v"(value.w), "v"(scale));
+  return static_cast<uint16_t>(lo | (hi << 8));
+}
+
+__device__ __forceinline__ Quantized4 quantize(float4 value, int lane) {
+  hadamard16(value, lane);
+  const float local_max = fmaxf(fmaxf(fabsf(value.x), fabsf(value.y)), fmaxf(fabsf(value.z), fabsf(value.w)));
+  float scale;
+  const uint8_t e8m0 = e8m0_scale(max8(local_max), scale);
+  return {pack_fp4(value, scale), e8m0};
+}
+
+__device__ __forceinline__ void store_scale(uint8_t* output, int row, int col, int cols, uint8_t value) {
+  const int tile = ((row >> 5) * (cols >> 3) + (col >> 3)) << 8;
+  const int offset = ((col & 3) << 6) + ((row & 15) << 2) + (((col >> 2) & 1) << 1) + ((row >> 4) & 1);
+  output[tile + offset] = value;
+}
+
+template<bool Shuffled>
+__device__ __forceinline__ void store_fp4(uint8_t* output, int row, int col, int packed_cols, uint16_t value) {
+  int index = row * packed_cols + col;
+  if constexpr (Shuffled) {
+    const int tile = (row >> 4) * (packed_cols << 4) + (col >> 5) * 512;
+    const int offset = ((col >> 4) & 1) * 256 + (row & 15) * 16 + (col & 15);
+    index = tile + offset;
+  }
+  *reinterpret_cast<uint16_t*>(output + index) = value;
+}
+
+__device__ __forceinline__ void load_tile(uint16_t* tile, const uint16_t* input, int tile_m, int tile_n) {
+  const int row = threadIdx.x / THREADS_PER_ROW;
+  const int col = threadIdx.x % THREADS_PER_ROW * VALUES_PER_THREAD;
+  const uint64_t packed = *reinterpret_cast<const uint64_t*>(input + (tile_m + row) * N + tile_n + col);
+  *reinterpret_cast<uint32_t*>(tile + row * SMEM_STRIDE + col) = static_cast<uint32_t>(packed);
+  *reinterpret_cast<uint32_t*>(tile + row * SMEM_STRIDE + col + 2) = static_cast<uint32_t>(packed >> 32);
+}
+
+__device__ __forceinline__ void quantize_row(uint16_t* tile, uint8_t* fp4_output, uint8_t* scale_output,
+                                             int tile_m, int tile_n, int local_row, int lane) {
+  const int row = tile_m + local_row;
+  const int col = lane * VALUES_PER_THREAD;
+  const Quantized4 result = quantize(load_bf16x4(tile + local_row * SMEM_STRIDE + col), lane);
+  store_fp4<SHUFFLE_ROWWISE_FP4>(fp4_output, row, (tile_n + col) / 2, N_PACKED, result.fp4);
+  if (lane == 0) store_scale(scale_output, row, tile_n / BLOCK, N_SCALES, result.scale);
+}
+
+__device__ __forceinline__ Quantized4 quantize_col(uint16_t* tile, int col, int lane) {
+  const int row = lane * VALUES_PER_THREAD;
+  return quantize(make_float4(
+    __uint_as_float(static_cast<uint32_t>(tile[(row + 0) * SMEM_STRIDE + col]) << 16),
+    __uint_as_float(static_cast<uint32_t>(tile[(row + 1) * SMEM_STRIDE + col]) << 16),
+    __uint_as_float(static_cast<uint32_t>(tile[(row + 2) * SMEM_STRIDE + col]) << 16),
+    __uint_as_float(static_cast<uint32_t>(tile[(row + 3) * SMEM_STRIDE + col]) << 16)), lane);
+}
+
+} // namespace
+
+extern "C" __global__ __launch_bounds__(THREADS, 8)
+void KERNEL_NAME(uint8_t* __restrict__ rowwise_fp4, uint8_t* __restrict__ rowwise_scale,
+                 uint8_t* __restrict__ colwise_fp4, uint8_t* __restrict__ colwise_scale,
+                 const uint16_t* __restrict__ input) {
+  __shared__ uint16_t tile[BLOCK * SMEM_STRIDE];
+  const int tid = threadIdx.x;
+  const int line = tid / THREADS_PER_ROW;
+  const int lane = tid % THREADS_PER_ROW;
+  const int block_m = blockIdx.x * TILE_M;
+  const int block_n = blockIdx.y * TILE_N;
+
+  if constexpr (!SHUFFLE_COLWISE_FP4) {
+    uint16_t col_fp4[TILE_N / BLOCK][TILE_M / BLOCK];
+    uint8_t col_scale[TILE_N / BLOCK][TILE_M / BLOCK];
+
+    for (int chunk_m = 0; chunk_m < TILE_M / BLOCK; chunk_m++) {
+      for (int chunk_n = 0; chunk_n < TILE_N / BLOCK; chunk_n++) {
+        const int tile_m = block_m + chunk_m * BLOCK;
+        const int tile_n = block_n + chunk_n * BLOCK;
+        load_tile(tile, input, tile_m, tile_n);
+        __syncthreads();
+
+        quantize_row(tile, rowwise_fp4, rowwise_scale, tile_m, tile_n, line, lane);
+        const Quantized4 result = quantize_col(tile, line, lane);
+        col_fp4[chunk_n][chunk_m] = result.fp4;
+        col_scale[chunk_n][chunk_m] = result.scale;
+        __syncthreads();
+      }
+    }
+
+    for (int chunk_n = 0; chunk_n < TILE_N / BLOCK; chunk_n++) {
+      for (int chunk_m = 0; chunk_m < TILE_M / BLOCK; chunk_m++)
+        tile[line * BLOCK + chunk_m * THREADS_PER_ROW + lane] = col_fp4[chunk_n][chunk_m];
+      __syncthreads();
+
+      for (int round = 0; round < BLOCK / THREADS_PER_ROW; round++) {
+        const int col = round * THREADS_PER_ROW + tid / BLOCK;
+        const int row_pair = tid % BLOCK;
+        *reinterpret_cast<uint16_t*>(colwise_fp4 + (block_n + chunk_n * BLOCK + col) * M_PACKED + block_m / 2 + row_pair * 2) =
+          tile[col * BLOCK + row_pair];
+      }
+
+      if (lane == 0) {
+        const int col = block_n + chunk_n * BLOCK + line;
+        for (int chunk_m = 0; chunk_m < TILE_M / BLOCK; chunk_m++)
+          store_scale(colwise_scale, col, block_m / BLOCK + chunk_m, M_SCALES, col_scale[chunk_n][chunk_m]);
+      }
+      __syncthreads();
+    }
+  } else {
+    for (int chunk_m = 0; chunk_m < TILE_M / BLOCK; chunk_m++) {
+      for (int chunk_n = 0; chunk_n < TILE_N / BLOCK; chunk_n++) {
+        const int tile_m = block_m + chunk_m * BLOCK;
+        const int tile_n = block_n + chunk_n * BLOCK;
+        load_tile(tile, input, tile_m, tile_n);
+        __syncthreads();
+
+        quantize_row(tile, rowwise_fp4, rowwise_scale, tile_m, tile_n, line, lane);
+        const int row = lane * VALUES_PER_THREAD;
+        const int col = tile_n + line;
+        const Quantized4 result = quantize_col(tile, line, lane);
+        store_fp4<true>(colwise_fp4, col, (tile_m + row) / 2, M_PACKED, result.fp4);
+        if (lane == 0) store_scale(colwise_scale, col, tile_m / BLOCK, M_SCALES, result.scale);
+        __syncthreads();
+      }
+    }
+  }
+}
@@ -0,0 +1,49 @@
+import functools, math
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, KernelInfo
+from tinygrad.renderer import Estimates
+from extra.llama_kernels import alloc_like
+
+LOG2E = 1.4426950408889634
+
+@functools.cache
+def _custom_swiglu(out:UOp, x_w13:UOp) -> UOp:
+  rows, hidden = math.prod(x_w13.shape[:-1]), x_w13.shape[-1]//2
+  n_elems = rows * hidden
+  out, x_w13 = out.reshape(n_elems), x_w13.reshape(rows, 2*hidden)
+  i = UOp.range(n_elems, 0)
+  row, col = i // hidden, i % hidden
+  act, gate = x_w13[row, col].cast(dtypes.float), x_w13[row, hidden+col].cast(dtypes.float)
+  sigmoid = (1.0 + (-LOG2E * act).exp2()).reciprocal()
+  store = out[i].store((act * sigmoid * gate).cast(out.dtype))
+  return store.end(i).sink(arg=KernelInfo(f"swiglu_fwd_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=6*n_elems)))
+
+@functools.cache
+def _custom_swiglu_bwd(grad_out:UOp, x_w13:UOp, grad_act:UOp) -> UOp:
+  rows, hidden = math.prod(x_w13.shape[:-1]), x_w13.shape[-1]//2
+  n_elems = rows * hidden
+  grad_out, x_w13, grad_act = grad_out.reshape(rows, 2*hidden), x_w13.reshape(rows, 2*hidden), grad_act.reshape(n_elems)
+  i = UOp.range(n_elems, 0)
+  row, col = i // hidden, i % hidden
+  act, gate = x_w13[row, col].cast(dtypes.float), x_w13[row, hidden+col].cast(dtypes.float)
+  grad = grad_act[i].cast(dtypes.float)
+  sigmoid = (1.0 + (-LOG2E * act).exp2()).reciprocal()
+  silu = act * sigmoid
+  dact = grad_out[row, col].store((grad * (sigmoid + silu * (1.0 - sigmoid)) * gate).cast(grad_out.dtype))
+  dgate = grad_out.after(dact)[row, hidden+col].store((grad * silu).cast(grad_out.dtype))
+  return dgate.end(i).sink(arg=KernelInfo(f"swiglu_bwd_{n_elems}", estimates=Estimates(ops=10*n_elems, mem=10*n_elems)))
+
+def _swiglu_bwd(gradient:UOp, kernel:UOp):
+  _, x_w13 = kernel.src[1:]
+  axis = x_w13.axis if isinstance(x_w13.device, tuple) else None
+  grad_out = alloc_like(x_w13.shape, dtypes.bfloat16, x_w13.device, axis)
+  grad_out, *_ = Tensor.custom_kernel(grad_out, Tensor(x_w13, device=x_w13.device), Tensor(gradient, device=x_w13.device),
+                                      fxn=_custom_swiglu_bwd)
+  return (None, grad_out.uop)
+
+def swiglu(x_w13:Tensor) -> Tensor:
+  assert x_w13.dtype == dtypes.bfloat16 and x_w13.ndim >= 2 and x_w13.shape[-1] % 32 == 0
+  *prefix, two_k = x_w13.shape
+  axis = x_w13.uop.axis if isinstance(x_w13.device, tuple) else None
+  out = alloc_like((*prefix, two_k//2), dtypes.bfloat16, x_w13.device, axis)
+  return Tensor.custom_kernel(out, x_w13, fxn=_custom_swiglu, grad_fxn=_swiglu_bwd)[0]
@@ -2,7 +2,7 @@ import math, pathlib, functools, struct

 from tinygrad import Device, Tensor
 from tinygrad.dtype import DTypeLike, dtypes
-from tinygrad.helpers import DEBUG
+from tinygrad.helpers import DEBUG, getenv
 from tinygrad.renderer import Estimates
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 from tinygrad.runtime.support.elf import elf_loader
@@ -206,10 +206,11 @@ def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None
                  arg=KernelInfo(name="custom_fa_forward", estimates=estimates))

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

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

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

  return UOp(Ops.PROGRAM,
             src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
@@ -268,10 +270,11 @@ def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_ve
                  arg=KernelInfo(name="custom_fa_backward", estimates=estimates))

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    // hot loop
-    for (int j = 3; j < max_num_tiles - 1; j += 2) {
+    for (int j = min_tile + 3; j < max_num_tiles - 1; j += 2) {
        // Cluster 0:
        //      QK1
        zero(att_block[1]);
        transpose(k_reg_transposed, k_reg);
        mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]);
+#if WINDOW
+        // window masks interior tiles that causal skips
+        mask_kv_tile(att_block[1], tile_idx, j - 2, neg_inf_v, lane);
+#endif
        //      Finish softmax for QK0
        exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]);
        mul(norm_vec, norm_vec, scale_vec);
@@ -379,7 +421,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
        load(v_reg, v_smem[1]);
        if constexpr (causal) {
            const int kv_end_pos = (j) * KV_BLOCK_SIZE;
-            if (q_start_pos < kv_end_pos) {  // Only mask if needed
+            if (WINDOW || q_start_pos < kv_end_pos) {
                mask_kv_tile(att_block[0], tile_idx, j - 1, neg_inf_v, lane);
            }
        }
@@ -447,7 +489,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    load(v_reg, v_smem[0]);
    if constexpr (causal) {
        const int kv_end_pos = (max_num_tiles - 2) * KV_BLOCK_SIZE;
-        if (__builtin_expect(q_start_pos < kv_end_pos, 0)) {  // Only mask if needed
+        if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, WINDOW ? 1 : 0)) {
            mask_kv_tile(att_block[1], tile_idx, max_num_tiles - 3, neg_inf_v, lane);
        }
    }
@@ -510,7 +552,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    load(v_reg, v_smem[1]);
    if constexpr (causal) {
        const int kv_end_pos = (max_num_tiles - 1) * KV_BLOCK_SIZE;
-        if (__builtin_expect(q_start_pos < kv_end_pos, 1)) {  // Only mask if needed
+        if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, 1)) {
            mask_kv_tile(att_block[0], tile_idx, max_num_tiles - 2, neg_inf_v, lane);
        }
    }
@@ -572,7 +614,7 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    load(v_reg, v_smem[0]);
    if constexpr (causal) {
        const int kv_end_pos = (max_num_tiles) * KV_BLOCK_SIZE;
-        if (__builtin_expect(q_start_pos < kv_end_pos, 1)) {  // Only mask if needed
+        if (__builtin_expect(WINDOW || q_start_pos < kv_end_pos, 1)) {
            mask_kv_tile(att_block[1], tile_idx, max_num_tiles - 1, neg_inf_v, lane);
        }
    }
@@ -97,4 +97,33 @@ __device__ inline static void atomic_pk_add_bf16_with_warpid(const GL &dst, cons
            }(std::make_index_sequence<RT::width>{});
        }.template operator()<Ns>(), ...);
    }(std::make_index_sequence<RT::height>{});
-}
+}
+// make_window: complement of make_causal for the window lower boundary (q_pos-k_pos == WINDOW). masks = ~(causal masks)
+template<int N, int M, int GPR, ducks::art::all T0, ducks::art::all T1>
+__device__ static inline void make_window(T0 &dst, const T1 &src) {
+    static_assert(std::is_same_v<typename T0::T, float> && std::is_same_v<typename T1::T, float>, "Only float to float window mask is supported");
+    static_assert(std::is_same_v<typename T0::layout, typename T1::layout>, "Only same layout is supported");
+    static_assert(std::is_same_v<typename T0::shape, typename T1::shape>, "Only same shape is supported");
+
+    if constexpr (std::is_same_v<typename T0::layout, typename ducks::rt_layout::col> && std::is_same_v<typename T0::shape, typename ducks::rt_shape::rt_16x16>) {
+        using range_type_T0 = ducks::art::get_nth_range_t<typename T0::register_ranges, N * T0::width + M>;
+        using registers_T0 = ducks::art::split_many_t<ducks::art::type_list<range_type_T0>, 1>;
+        using range_type_T1 = ducks::art::get_nth_range_t<typename T1::register_ranges, N * T1::width + M>;
+        using registers_T1 = ducks::art::split_many_t<ducks::art::type_list<range_type_T1>, 1>;
+        static_assert(registers_T0::size == registers_T1::size);
+
+        uint64_t window_mask = 0x1FFF01FF001F0001;
+        macros::v_cndmask_b32_e64<ducks::art::get_nth_range_t<registers_T0, 0>::lo, ducks::art::get_nth_range_t<registers_T1, 0>::lo, GPR>(window_mask);
+
+        window_mask = 0x3FFF03FF003F0003;
+        macros::v_cndmask_b32_e64<ducks::art::get_nth_range_t<registers_T0, 1>::lo, ducks::art::get_nth_range_t<registers_T1, 1>::lo, GPR>(window_mask);
+
+        window_mask = 0x7FFF07FF007F0007;
+        macros::v_cndmask_b32_e64<ducks::art::get_nth_range_t<registers_T0, 2>::lo, ducks::art::get_nth_range_t<registers_T1, 2>::lo, GPR>(window_mask);
+
+        window_mask = 0xFFFF0FFF00FF000F;
+        macros::v_cndmask_b32_e64<ducks::art::get_nth_range_t<registers_T0, 3>::lo, ducks::art::get_nth_range_t<registers_T1, 3>::lo, GPR>(window_mask);
+    } else {
+        static_assert(false, "Unsupported window mask");
+    }
+}
@@ -0,0 +1,74 @@
+#include "kittens.cuh"
+
+using namespace kittens;
+
+#ifndef MATVEC_N
+#define MATVEC_N 1536
+#endif
+#ifndef MATVEC_K
+#define MATVEC_K 7168
+#endif
+
+constexpr int SPLIT_WAVES = 8;
+
+template<int W>
+__device__ __forceinline__ float run_split(const bf16 *A_ptr, const bf16 *B_ptr, int out_base,
+                                           st_bf<16, 32, st_16x32_s> &As,
+                                           st_bf<16, 32, st_16x32_s> &Bs) {
+  constexpr int K = MATVEC_K;
+  rt_bf<16, 32, row_l, rt_16x32_s> A;
+  rt_bf<16, 32, row_l, rt_16x32_s> B;
+  rt_fl<16, 16, col_l, rt_16x16_s> C;
+  zero(C);
+  const int lane = laneid();
+  constexpr int k_begin = W * (K / SPLIT_WAVES), k_end = k_begin + K / SPLIT_WAVES;
+  #pragma unroll 1
+  for (int k = k_begin; k < k_end; k += 32) {
+    #pragma unroll
+    for (int idx = lane; idx < 16 * 32; idx += 64) {
+      const int row = idx / 32, col = idx % 32;
+      *reinterpret_cast<bf16 *>(reinterpret_cast<char *>(&As.data[0]) + As.swizzle({row, col})) = A_ptr[k + col];
+      *reinterpret_cast<bf16 *>(reinterpret_cast<char *>(&Bs.data[0]) + Bs.swizzle({row, col})) =
+        B_ptr[(out_base + row) * K + k + col];
+    }
+    asm volatile("s_waitcnt lgkmcnt(0)");
+    load(A, As);
+    load(B, Bs);
+    asm volatile("s_waitcnt lgkmcnt(0)");
+    mma_ABt(C, A, B, C);
+  }
+  return C.tiles[0][0].data[0].x;
+}
+
+// Eight waves split K for one 16-channel output tile.  Each wave uses MFMA on
+// a repeated activation row, then wave zero reduces the eight FP32 partials.
+__global__ __launch_bounds__(64 * SPLIT_WAVES, 1)
+void hk_bf16_matvec_splitk(bf16 *C_ptr, const bf16 *A_ptr, const bf16 *B_ptr, bf16 *unused) {
+  constexpr int N = MATVEC_N, K = MATVEC_K;
+  static_assert(N % 16 == 0 && K % (32 * SPLIT_WAVES) == 0);
+  __shared__ st_bf<16, 32, st_16x32_s> As[SPLIT_WAVES];
+  __shared__ st_bf<16, 32, st_16x32_s> Bs[SPLIT_WAVES];
+  __shared__ float partial[SPLIT_WAVES][16];
+  const int tid = threadIdx.x, wave = tid / 64, lane = tid & 63;
+  const int out_base = blockIdx.x * 16;
+  float result = 0.0f;
+  switch (wave) {
+    case 0: result = run_split<0>(A_ptr, B_ptr, out_base, As[0], Bs[0]); break;
+    case 1: result = run_split<1>(A_ptr, B_ptr, out_base, As[1], Bs[1]); break;
+    case 2: result = run_split<2>(A_ptr, B_ptr, out_base, As[2], Bs[2]); break;
+    case 3: result = run_split<3>(A_ptr, B_ptr, out_base, As[3], Bs[3]); break;
+    case 4: result = run_split<4>(A_ptr, B_ptr, out_base, As[4], Bs[4]); break;
+    case 5: result = run_split<5>(A_ptr, B_ptr, out_base, As[5], Bs[5]); break;
+    case 6: result = run_split<6>(A_ptr, B_ptr, out_base, As[6], Bs[6]); break;
+    case 7: result = run_split<7>(A_ptr, B_ptr, out_base, As[7], Bs[7]); break;
+  }
+  if (lane < 16) partial[wave][lane] = result;
+  asm volatile("s_waitcnt lgkmcnt(0)");
+  __builtin_amdgcn_s_barrier();
+  if (wave == 0 && lane < 16) {
+    float total = 0.0f;
+    #pragma unroll
+    for (int i = 0; i < SPLIT_WAVES; i++) total += partial[i][lane];
+    C_ptr[out_base + lane] = static_cast<bf16>(total);
+  }
+}
@@ -471,6 +471,20 @@ class TestCmpFloat(unittest.TestCase):
    st = run_program(instructions, n_lanes=1)
    self.assertEqual(st.vcc & 1, 1, "Expected vcc=1 (1.0 != 2.0)")

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

@@ -4,13 +4,13 @@ from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
 from tinygrad.helpers import Context, getenv, DEV
 from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear
 from tinygrad.renderer.ptx import PTXRenderer
-from test.helpers import needs_second_gpu, check_schedule, assert_kernel_count
+from test.helpers import needs_second_gpu, check_schedule, assert_kernel_count, KernelCountException

 class TestArange(unittest.TestCase):
  def _get_flops(self, tensor, desired):
    GlobalCounters.reset()
    linear = compile_linear(tensor.schedule_linear())
-    self.assertEqual(len(linear.src), 1)
+    if len(linear.src) != 1: raise KernelCountException(1, len(linear.src))
    run_linear(linear)
    np.testing.assert_equal(tensor.numpy(), desired)
    return estimate_uop(linear.src[-1]).ops
@@ -253,7 +253,7 @@ class TestIndexing(unittest.TestCase):
    xq_rope, _ = apply_rotary_emb(xq, xq, freqs_cis)
    xq_rope.sum().backward()
    linear = compile_linear(wq.grad.schedule_linear())
-    assert len(linear.src) == 1, f"expected one kernel for backward, got: {len(linear.src)}"
+    if len(linear.src) != 1: raise KernelCountException(1, len(linear.src))
    bwd_ops = estimate_uop(linear.src[0]).ops
    expected_ops = bs*seqlen*dim*dim*ops_scale
    print(f"rope matmul bwd ({dtype}): {GlobalCounters.kernel_count} kernels, {bwd_ops:,} ops")
@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.helpers import getenv, system, DEV
-from extra.gemm.cdna_asm_gemm import asm_gemm, hk_bf16_atb_gemm, quantize_mxfp4
+from extra.gemm.cdna_asm_gemm import asm_gemm, hk_bf16_atb_gemm
 from test.helpers import needs_second_gpu
 from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8, FP8_MAX

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

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

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

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

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

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

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

-  @unittest.expectedFailure
  def test_sliced_buffer_function(self):
    x = Tensor.arange(32).reshape(8, 4).clone().realize()
    from tinygrad import function
@@ -419,7 +423,8 @@ class TestCustomKernel(unittest.TestCase):
    GlobalCounters.reset()
    y = run(x[0]).realize()
    # it's copying the input and the output
-    assert_kernel_count(1)
+    # TODO: subbuffer usage has runtime specific behavior, this will be fixed after the removal of SLICE.
+    assert_kernel_count(2 if y.device in ("CL", "WEBGPU") else 1)
    self.assertEqual(y.tolist(), [1, 2, 3, 4])

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

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

  def test_reshape(self): self._test_mop(lambda x: x.reshape(16, 2), max_kernels=2)
  def test_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T, max_kernels=3)
-  def test_double_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T.T, max_kernels=3)
-  def test_shrink(self): self._test_mop(lambda x: x[:4], max_kernels=2)
+  def test_double_permute(self): self._test_mop(lambda x: x.reshape(4, 8).T.T, max_kernels=2)
+  def test_shrink(self): self._test_mop(lambda x: x[:4], max_kernels=1)
  def test_pad(self): self._test_mop(lambda x: x[:4].pad(((0, 4),)), max_kernels=2)
  def test_flip(self): self._test_mop(lambda x: x.flip(0), max_kernels=2)
  def test_offset_shrink(self): self._test_mop(lambda x: x[4:8], max_kernels=2)
@@ -169,6 +169,13 @@ class TestFp8sConversions(unittest.TestCase):
  def test_fp8e5m2fnuz_to_float(self, x):
    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2fnuz).float().item())

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

-  @unittest.expectedFailure
-  def test_unsafe_cast_float_to_int_failure(self):
-    val = float(dtypes.int32.max - 1)
+  def test_unsafe_cast_float_to_int(self):
+    # the value is off the float32 grid but rounds in-range: the buffer and const-fold paths must agree
+    # (out-of-range float->int cast stays undefined: hardware may saturate where the fold wraps)
+    val = 2147483000.0
    t1 = Tensor([val], dtype=dtypes.float32).cast(dtypes.int32)
    t2 = Tensor(val, dtype=dtypes.float32).cast(dtypes.int32)
    np.testing.assert_equal(t1.item(), t2.item())
@@ -5,6 +5,7 @@ from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8
 from extra.llama_kernels.fused_ce import fused_ce_loss
 from extra.llama_kernels import local_abs_max
 from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed, quantize_fp8_scalar
+from extra.llama_kernels.swiglu import swiglu
 from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
 from extra.thunder.amd.fa import custom_fused_qkv_rope_backward, fused_qkv_rope
 from test.helpers import needs_second_gpu, assert_kernel_count
@@ -161,5 +162,31 @@ class TestFusedQKVRoPE(unittest.TestCase):
    ref = Tensor.cat(dq_ref, dk_ref, dv_ref, dim=3).reshape(*dx.shape).realize()
    with Context(DEBUG=0): self.assertTrue(dx.allclose(ref, atol=2e-2, rtol=2e-2).item(), "backward mismatch")

+def run_swiglu(test:unittest.TestCase, shape:tuple[int, ...]) -> None:
+  Tensor.manual_seed(0)
+  x = (Tensor.randn(*shape) * 2).cast(dtypes.bfloat16).realize()
+  hidden = x.shape[-1] // 2
+  out, ref = swiglu(x), x[..., :hidden].silu() * x[..., hidden:]
+  Tensor.realize(out, ref)
+  with Context(DEBUG=0): test.assertTrue(out.allclose(ref, atol=2.5e-1, rtol=3e-2).item(), "SwiGLU forward mismatch")
+
+  grad = (Tensor.randn(*out.shape) * 2).cast(dtypes.bfloat16).realize()
+  grad_x, grad_ref = out.gradient(x, gradient=grad)[0], ref.gradient(x, gradient=grad)[0]
+  Tensor.realize(grad_x, grad_ref)
+  test.assertEqual(grad_x.shape, shape)
+  test.assertEqual(grad_x.dtype, dtypes.bfloat16)
+  with Context(DEBUG=0): test.assertTrue(grad_x.allclose(grad_ref, atol=2.5e-1, rtol=3e-2).item(), "SwiGLU backward mismatch")
+
+class TestSwiGLU(unittest.TestCase):
+  def setUp(self):
+    if dtypes.bfloat16 not in Device[Device.DEFAULT].renderer.supported_dtypes(): self.skipTest("need bfloat16")
+
+  def test_simple(self): run_swiglu(self, (2, 32, 64))
+
+  def test_llama_shape(self):
+    if Device.DEFAULT != "AMD" or not Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950"):
+      self.skipTest("only run on real machine for speed")
+    run_swiglu(self, (2, 8192, 28672))
+
 if __name__ == '__main__':
  unittest.main()
@@ -1535,6 +1535,8 @@ class TestOps(unittest.TestCase):

  def test_prod(self):
    helper_test_op(None, lambda x: x.prod(), vals=[[1.0, 2.0, 3.0]])
+    helper_test_op(None, lambda x: x.prod(), vals=[[0.0, 2.0, 3.0]])
+    helper_test_op(None, lambda x: x.prod(), vals=[[0.0, 0.0, 3.0]])
    with Context(NOOPT=1): helper_test_op(None, lambda x: x.prod(), vals=[[1.0, 2.0, 3.0]])
    helper_test_op([(3,4,5,6)], lambda x: x.prod(dim=3), lambda x: x.prod(axis=3))
    helper_test_op([(3,4,5,6)], lambda x: x.prod(dim=1), lambda x: x.prod(axis=1))
@@ -6,6 +6,15 @@ from examples.gpt2 import Attention
 import numpy as np

 class TestSymbolicOps(unittest.TestCase):
+  def test_negative_slice(self):
+    a = Tensor.rand(3, 10, 4)
+    for i in range(3, 10):
+      vi = Variable("i", 1, 10).bind(i)
+      # negative int bounds against a symbolic dim must resolve against the size, like slice.indices
+      np.testing.assert_allclose(a[:, :vi][:, -3:-1].numpy(), a[:, :i][:, -3:-1].numpy(), atol=1e-6, rtol=1e-6)
+      np.testing.assert_allclose(a[:, :vi][:, -1:].numpy(), a[:, :i][:, -1:].numpy(), atol=1e-6, rtol=1e-6)
+      np.testing.assert_allclose(a[:, :vi][:, -1].numpy(), a[:, :i][:, -1].numpy(), atol=1e-6, rtol=1e-6)
+
  def test_plus1(self):
    def f(a): return (a+1).realize()
    a = Tensor.rand(3, 10)
@@ -109,7 +109,7 @@ def _init_sqtt_encoder():
  _SMEM = (ir3.SMEM, ir4.SMEM, irc.SMEM)
  _VALU = (ir3.VOP1, ir3.VOP2, ir3.VOP3, ir3.VOP3P, ir3.VOPC, ir3.VOPD, ir3.VOP3SD, ir3.VOP3_SDST, ir3.VOP1_SDST,
           ir4.VOP1, ir4.VOP2, ir4.VOP3, ir4.VOP3P, ir4.VOPC, ir4.VOPD, ir4.VOP3SD, ir4.VOP3_SDST, ir4.VOP1_SDST,
-           irc.VOP1, irc.VOP2, irc.VOP3, irc.VOP3P, irc.VOPC, irc.VOP3SD, irc.VOP3_SDST)
+           irc.VOP1, irc.VOP2, irc.VOP3, irc.VOP3P, irc.VOP3PX2, irc.VOPC, irc.VOP3SD, irc.VOP3_SDST)
  _DS = (ir3.DS, ir4.DS, irc.DS)
  _GLOBAL = (ir3.GLOBAL, ir4.VGLOBAL, irc.GLOBAL)
  _FLAT = (ir3.FLAT, ir4.VFLAT, irc.FLAT)
@@ -1150,6 +1150,9 @@ def _compile_vopc(inst: ir3.VOPC|ir3.VOPC_DPP16|ir3.VOP3|ir4.VOPC|ir4.VOPC_DPP16
  def get_cmp_bit(lane) -> UOp:
    lc = lane.cast(dtypes.int) if isinstance(lane, UOp) else _c(lane, dtypes.int)
    s0 = _load_dpp16_src0(ctx, inst, lc, _c(0)) if is_dpp16 else ctx.rsrc_dyn(src0_off, lc, bits['s0'], literal, is_f64)
+    if is_vopc and not isinstance(inst, irc.VOPC) and bits['s0'] == 16 and not is_dpp16:
+      src0_hi = src0_off >= _c(384)
+      s0 = src0_hi.where(_hi16(ctx.rvgpr_dyn(src0_hi.where(src0_off - _c(384), _c(0)), lc)), s0)
    s1 = _cond_hi16(vsrc1_hi, ctx.rsrc_dyn(src1_off, lc, bits['s1'], literal, is_f64)) if bits['s0'] == 16 \
      else ctx.rsrc_dyn(src1_off, lc, bits['s1'], literal, is_f64)
    if bits['s0'] == 16 and opsel: s0, s1 = _apply_opsel(s0, 0, opsel), _apply_opsel(s1, 1, opsel)
@@ -1323,7 +1326,7 @@ def _compile_vop3sd(inst: ir3.VOP3SD | ir4.VOP3SD | irc.VOP3SD, ctx: _Ctx) -> UO
  else:
    return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask, sdst_reg=inst.sdst.offset)

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

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

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

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

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

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

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

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

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

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

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

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

  # ACCVGPR_WRITE/READ/MOV: copies between VGPR and ACCVGPR register files
  # Detect by checking operand types for ACCVGPR involvement
@@ -2044,7 +2093,7 @@ _INST_HANDLERS: dict[type, Callable[..., UOp]] = {
  irc.SOPP: _compile_sopp, irc.SMEM: _compile_smem, irc.SOP1: _compile_sop, irc.SOP2: _compile_sop, irc.SOPC: _compile_sop, irc.SOPK: _compile_sop,
  irc.VOP1: _compile_vop12, irc.VOP1_DPP16: _compile_vop12, irc.VOP2: _compile_vop12, irc.VOP2_DPP16: _compile_vop12,
  irc.VOPC: _compile_vopc, irc.VOP3: _compile_vop3,
-  irc.VOP3_SDST: _compile_vop3, irc.VOP3SD: _compile_vop3sd, irc.VOP3P: _compile_vop3p,
+  irc.VOP3_SDST: _compile_vop3, irc.VOP3SD: _compile_vop3sd, irc.VOP3P: _compile_vop3p, irc.VOP3PX2: _compile_vop3p,
  irc.VOP1_SDWA: _compile_sdwa, irc.VOP2_SDWA: _compile_sdwa, irc.VOP2_SDWA_SDST: _compile_sdwa, irc.VOPC_SDWA_SDST: _compile_sdwa,
  irc.DS: _compile_mem_op, irc.FLAT: _compile_mem_op, irc.GLOBAL: _compile_mem_op, irc.SCRATCH: _compile_mem_op,
  irc.MUBUF: _compile_mubuf,
@@ -1,6 +1,6 @@
 import unittest, itertools, math
 from tinygrad import Tensor, dtypes, Context
-from tinygrad.dtype import DType, ConstType
+from tinygrad.dtype import DType, ConstType, truncate
 from tinygrad.uop.ops import Ops, UOp
 from test.helpers import full_rewrite
 import numpy as np
@@ -51,6 +51,17 @@ class TestWeakConstFolding(unittest.TestCase):
  def test_invalid_poison(self):
    self.assertTrue(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().is_invalid)

+  def test_cast_commits_to_dtype_grid(self):
+    # committing a weak const to a stated width puts the value on that width's grid, same as storage packing and native compilers
+    v = 1/123008  # not representable in float16
+    out = UOp.const(v).cast(dtypes.half).simplify()
+    self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.half, truncate[dtypes.half](v)))
+    self.assertNotEqual(out.val, v)
+    # the grid commit preserves the sign of zero
+    self.assertEqual(math.copysign(1, UOp.const(-0.0).cast(dtypes.half).simplify().val), -1)
+    # observable at tensor level: the const-folded comparison agrees with the committed value
+    self.assertTrue((Tensor(-3.2).cast(dtypes.float32) <= truncate[dtypes.float32](-3.2)).item())
+
 class TestBinaryOpsConstFolding(unittest.TestCase):
  def test_add_literal_zero(self):
    _check_ast_count(0, Tensor([1.0, 2, 3, 4]) + 0)
@@ -208,7 +208,7 @@ class TestGEPAndVectorizeRewrite(unittest.TestCase):


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

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

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

@@ -0,0 +1,30 @@
+import unittest
+from tinygrad import Tensor, dtypes, nn
+from tinygrad.llm.kimi import _shard_kimi
+from tinygrad.llm.model import SSMConfig, Transformer, TransformerConfig
+
+class TestKimiTP4(unittest.TestCase):
+  def test_prefill_and_decode_graph(self):
+    devices = ("NULL:0", "NULL:1", "NULL:2", "NULL:3")
+    config = TransformerConfig(num_blocks=2, dim=32, hidden_dim=128, n_heads=4, n_kv_heads=1, norm_eps=1e-5,
+      vocab_size=64, head_dim=12, rope_theta=10000, rope_dim=4, v_head_dim=8, max_context=4, kv_lora_rank=16,
+      num_experts=8, num_experts_per_tok=2, norm_topk_prob=True, shared_expert_dim=32, ssm_layers=(True, False),
+      ssm=SSMConfig(4, 8, 4, 4, 32, True), shared_expert_gate=False, leading_dense_blocks=1, dense_hidden_dim=64,
+      routed_scaling_factor=2.446, expert_bias=True, expert_mxfp4=True, bf16_activations=True, kda_split_qkv=True)
+    model = Transformer(config)
+    for name, value in nn.state.get_state_dict(model).items():
+      fill = 127 if name.endswith("weight_scale") else 0
+      value.replace(Tensor.full(value.shape, fill, dtype=value.dtype if value.dtype is dtypes.uint8 else dtypes.bfloat16, device="NULL"))
+    _shard_kimi(model, devices)
+
+    temperature = Tensor([0.0], device=devices)
+    prefill = model(Tensor([[1, 2]], dtype=dtypes.int32, device=devices), 0, temperature).realize()
+    model(Tensor([[1, 2]], dtype=dtypes.int32, device=devices), 0, temperature).realize()  # replay prefill JIT
+    decode = model(Tensor([[3]], dtype=dtypes.int32, device=devices), 2, temperature).realize()
+    model(Tensor([[4]], dtype=dtypes.int32, device=devices), 3, temperature).realize()  # replay decode JIT
+    self.assertEqual(prefill.shape, (1, 1))
+    self.assertEqual(decode.shape, (1, 1))
+    self.assertEqual(model.blk[0].recurrent_state.uop.axis, 1)
+    self.assertEqual(model.blk[1].cache_k.dtype, dtypes.bfloat16)
+
+if __name__ == "__main__": unittest.main()
@@ -0,0 +1,29 @@
+import unittest
+from tinygrad import Tensor, dtypes, nn
+from tinygrad.llm.kimi_k3 import _shard_kimi_k3
+from test.unit.test_llm_k3 import small_k3_config
+from tinygrad.llm.model import Transformer
+
+class TestKimiK3TP8(unittest.TestCase):
+  @staticmethod
+  def _model():
+    model = Transformer(small_k3_config())
+    for name,value in nn.state.get_state_dict(model).items():
+      fill = 127 if name.endswith("weight_scale") else 0
+      dtype = value.dtype if value.dtype is dtypes.uint8 else dtypes.bfloat16
+      value.replace(Tensor.full(value.shape, fill, dtype=dtype, device="NULL"))
+    _shard_kimi_k3(model, tuple(f"NULL:{i}" for i in range(8)))
+    return model
+
+  def test_prefill_decode_and_jit_replay(self):
+    devices = tuple(f"NULL:{i}" for i in range(8))
+    model = self._model()
+    temperature = Tensor([0.0], device=devices)
+    self.assertEqual(model(Tensor([[1, 2]], dtype=dtypes.int32, device=devices), 0, temperature).realize().shape, (1, 1))
+    model(Tensor([[1, 2]], dtype=dtypes.int32, device=devices), 0, temperature).realize()
+    self.assertEqual(model(Tensor([[3]], dtype=dtypes.int32, device=devices), 2, temperature).realize().shape, (1, 1))
+    model(Tensor([[4]], dtype=dtypes.int32, device=devices), 3, temperature).realize()
+    self.assertEqual(model.blk[0].recurrent_state.uop.axis, 1)
+    self.assertEqual(model.blk[1].cache_k.dtype, dtypes.bfloat16)
+
+if __name__ == "__main__": unittest.main()
@@ -6,7 +6,7 @@ from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat, KernelInfo, AxisType
 from tinygrad.helpers import GlobalCounters, Context
 from tinygrad.engine.realize import run_linear, compile_linear
 from tinygrad.codegen import to_program, full_rewrite_to_sink
-from test.helpers import check_schedule, assert_kernel_count
+from test.helpers import check_schedule, assert_kernel_count, KernelCountException

 def _realize_weights(m):
  for p in nn.state.get_parameters(m): p.realize()
@@ -592,9 +592,7 @@ class TestSchedule(unittest.TestCase):
    img = Tensor.randn(BS, CIN, 64, 64).realize()
    w = Tensor.uniform(16, CIN, 3, 3).realize()
    ret = Tensor.conv2d(img, w).relu().mean().backward()
-    linear, var_vals = Tensor.linear_with_vars(ret, img.grad, w.grad)
-    cnt = len([call for call in linear.src if call.src[0].op is Ops.SINK])
-    assert cnt == allowed, f"expected {allowed} kernels, got {cnt}"
+    check_schedule([ret, img.grad, w.grad], allowed)

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

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

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

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

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

  def _test_fusion(self, shapes, f, cnt):
    with Context(DEBUG=0, TRACK_MATCH_STATS=0):
@@ -858,6 +858,65 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.rand(32)
    check_schedule(x, 1, [Tensor._device_rng_counters[x.device]])

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

 from tinygrad.codegen.late.coalesce import indexing_simplify
 from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import UOp, Ops, graph_rewrite, pm_lower_index_dtype
+from tinygrad.uop.ops import UOp, Ops, graph_rewrite
+from tinygrad.uop.weak import pm_lower_index_dtype
 from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load
 from tinygrad.helpers import Context
 from test.helpers import full_rewrite
@@ -332,7 +333,7 @@ class TestImageSimplification(unittest.TestCase):
    load = get_load_image_uop(shape, valid, idx)

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

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

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

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

 simple_pm = PatternMatcher([
  (UPat.cvar('x', dtypes.weakint), lambda x: UOp.const(1.0) + UOp.const(2.0)),
@@ -536,6 +537,15 @@ class TestReduceCollapse(unittest.TestCase):
    # Should become add of two separate reduces
    self.assertEqual(result.op, Ops.ADD)

+  def test_reduce_shapeless_const_unroll(self):
+    """a REDUCE over a shapeless CONST (e.g. x*0 folded late in codegen) must collapse before the expander"""
+    out = UOp.param(0, dtypes.float, (1,))
+    red = UOp.const(3.0).cast(dtypes.float).reduce(UOp.range(4, 0, AxisType.UNROLL), arg=(Ops.ADD, 0))
+    ast = UOp.sink(out.index(UOp.const(0)).store(red)).replace(arg=KernelInfo())
+    uops = full_rewrite_to_sink(ast, Device["CPU"].renderer, optimize=False).toposort()
+    self.assertNotIn(Ops.REDUCE, [u.op for u in uops])
+    self.assertIn(12.0, [u.val for u in uops if u.op is Ops.CONST])
+
 class TestMovementOps(unittest.TestCase):
  def test_pm_mops_partial_reshape_index_removes_reshape(self):
    from tinygrad.schedule.rangeify import pm_mops
@@ -6,7 +6,7 @@ from tinygrad.dtype import dtypes, ConstType, DType, Invalid
 from test.helpers import get_uops
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad.uop.spec import spec_shared, type_verify
-from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load
+from tinygrad.uop.symbolic import sym, pm_fold_cast_const, commutative, pm_simplify_valid, pm_move_where_on_load
 from tinygrad.uop.validate import uops_to_z3

 def check_uop_against_string(self, v:UOp, s:str):
@@ -35,7 +35,7 @@ class TestSymbolic(unittest.TestCase):
    self.assertEqual(solver.check(expr1 != expr2), z3.unsat, "simplified expression not equal to original")

  def helper_test_variable(self, v, n, m, s, test_z3:bool=True):
-    v_simplified = graph_rewrite(v, sym, name="simplify symbolic uop")
+    v_simplified = graph_rewrite(v, sym+pm_fold_cast_const, name="simplify symbolic uop")
    if test_z3: self.check_equal_z3(v, v_simplified)
    nmin, nmax = v_simplified.vmin, v_simplified.vmax
    check_uop_against_string(self, v_simplified, s)
@@ -1,6 +1,6 @@
 import unittest, math
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.dtype import dtypes, Invalid
+from tinygrad.dtype import dtypes, Invalid, truncate

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

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

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

-@track_rewrites()
+@rewrite_group()
 def do_compile(up):
  print("\n***** COMPILE", up)
  match_code = _get_code(up, False)
@@ -3,7 +3,7 @@ from pathlib import Path
 from dataclasses import dataclass
 from typing import Generator

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

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

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

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

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

-  def test_profile_matches_invalid_arg(self):
+  def test_rewrite_group_invalid_arg(self):
    with save_viz():
-      @profile_matches
+      @rewrite_group(new_ctx=False)
      def invalid_fxn(arg:str): return graph_rewrite(UOp(Ops.SINK), PatternMatcher([]))
      with self.assertRaisesRegex(AssertionError, "invalid match tracing input"):
        invalid_fxn("test")
@@ -395,7 +395,7 @@ class TestVizIntegration(unittest.TestCase):
    graph = next(viz.get_details(0, 0))["graph"]
    self.assertEqual(len([n for n in graph.values() if repr(metadata) in n["label"]]), 1)

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

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

-  # @track_rewrites always starts a new group.
+  # @rewrite_group always starts a new group.
  def test_group_combined(self):
    with save_viz() as viz:
      def default_test(root): return graph_rewrite(root, sym)
-      tracked_test = track_rewrites()(default_test)
+      tracked_test = rewrite_group()(default_test)
      c = UOp.const(1)
      default_test(c+1) # goes to the default group
      tracked_test(c)   # all rewrites after this go inside the second group.
@@ -101,7 +101,8 @@ class TestTensorCores(unittest.TestCase):
      if Device.DEFAULT == "CPU" and DEV.renderer == "LLVM":
        assert "0x201000" in prg.src[2].arg
      elif Device.DEFAULT == "AMD" and DEV.renderer == "LLVM":
-        assert "@llvm.amdgcn.wmma" in prg.src[2].arg
+        # RDNA emits wmma intrinsics, CDNA emits mfma intrinsics
+        assert ("@llvm.amdgcn.wmma" in prg.src[2].arg) or ("@llvm.amdgcn.mfma" in prg.src[2].arg)
      elif Device[Device.DEFAULT].renderer.suffix == "PTX":
        assert "mma.sync.aligned" in prg.src[2].arg
      else:
@@ -181,10 +182,12 @@ class TestTensorCores(unittest.TestCase):
  @unittest.skipIf(Device.DEFAULT == "PYTHON", "slow on EMULATED device")
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
  def test_tensor_cores_unroll_phi(self):
-    tc = Device[Device.DEFAULT].renderer.tensor_cores[0]
-    x, y = Tensor.rand(128, 128, dtype=tc.dtype_in), Tensor.rand(128, 128, dtype=tc.dtype_in)
+    # skip fp8 tcs: the unoptimized ALU baseline quantizes products to fp8 (JAX promotion), which legitimately
+    # differs from the MFMA path (f32 accumulation), so the baseline-vs-TC numerical gate can't hold for fp8.
+    tc = next(tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in not in dtypes.fp8s)
+    x, y = Tensor.rand(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in)
    r = x.matmul(y, dtype=tc.dtype_out)
-    opts = [Opt(OptOps.UNROLL, 0, 4)]
+    opts = [Opt(OptOps.UNROLL, 0, 2)]
    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
    for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src):
      if u.op is Ops.WMMA:
@@ -195,10 +198,10 @@ class TestTensorCores(unittest.TestCase):
  @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
  @unittest.skipIf(Device.DEFAULT in {"CPU"}, "CPU does not support using a different type for accumulation")
  def test_tensor_cores_unroll_casted_phi(self):
-    tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out][0]
-    x, y = Tensor.rand(128, 128, dtype=tc.dtype_in), Tensor.rand(128, 128, dtype=tc.dtype_in)
+    tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out and tc.dtype_in not in dtypes.fp8s][0]
+    x, y = Tensor.rand(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in)
    r = x.matmul(y, dtype=tc.dtype_out)
-    opts = [Opt(OptOps.UNROLL, 0, 4)]
+    opts = [Opt(OptOps.UNROLL, 0, 2)]
    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
    for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src):
      if u.op is Ops.WMMA:
@@ -211,10 +214,10 @@ class TestTensorCores(unittest.TestCase):
  @unittest.skipIf(Device.DEFAULT in {"CPU"}, "CPU does not support using a different type for accumulation")
  def test_tensor_cores_unroll_casted_phi_with_children(self):
    # all STORE children are outside the loop
-    tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out][0]
-    x, y = Tensor.rand(128, 128, dtype=tc.dtype_in), Tensor.rand(128, 128, dtype=tc.dtype_in)
+    tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out and tc.dtype_in not in dtypes.fp8s][0]
+    x, y = Tensor.rand(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in)
    r = x.matmul(y, dtype=tc.dtype_out).relu()
-    opts = [Opt(OptOps.UNROLL, 0, 4)]
+    opts = [Opt(OptOps.UNROLL, 0, 2)]
    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
    for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src):
      if u.op is Ops.WMMA:
@@ -1,7 +1,9 @@
 import unittest

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

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

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

+  @unittest.skipUnless(Device.DEFAULT == "NV", "NV only")
+  def test_hevc_decode_compile(self):
+    url = "https://github.com/haraschax/filedump/raw/09a497959f7fa6fd8dba501a25f2cdb3a41ecb12/comma_video.hevc"
+    dat = fetch(url, headers={"Range": f"bytes=0-{512<<10}"}).read_bytes()
+
+    opaque, frame_info, _, _, luma_w, luma_h, _ = parse_hevc_file_headers(dat)
+    offset, sz, frame_pos, max_hist, _ = frame_info[1]
+    out_image_size = luma_h + (luma_h + 1) // 2, round_up(luma_w, 64)
+    history = [Tensor.empty(*out_image_size, dtype=dtypes.uint8, device="NV") for _ in range(max_hist)]
+    decoded = Tensor(dat, device="NV")[offset:offset+sz].decode_hevc_frame(
+      Variable("pos", 0, max_hist + 1).bind(frame_pos), out_image_size, opaque[1], history)
+
+    compiled = compile_linear(decoded.linear_with_vars()[0])
+    self.assertTrue(any(call.src[0].op is Ops.PROGRAM for call in compiled.src))
+    encdec_calls = [call for call in compiled.src if call.src[0].op is Ops.CUSTOM_FUNCTION and call.src[0].arg == "encdec"]
+    self.assertEqual(len(encdec_calls), 1)
+
 if __name__ == "__main__":
  unittest.main()
@@ -1,9 +1,11 @@
 import unittest
+from types import SimpleNamespace
 import numpy as np
-from tinygrad import Tensor, dtypes
+from tinygrad import Tensor, dtypes, nn
+from tinygrad.llm.kimi import _shard_kimi
 from tinygrad.llm.model import (
  GatedDeltaNetBlock, SSMConfig, TransformerBlock, TransformerConfig,
-  apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk,
+  apply_rope as apply_rope_new, iterative_topk, l2norm, precompute_freqs_cis, pairwise_topk,
 )

 def apply_rope(x:Tensor, start_pos:int):
@@ -41,6 +43,11 @@ class TestAttention(unittest.TestCase):
    np.testing.assert_allclose(block.cache_kv[0, :, :, :seqlen, :].numpy(), expected.numpy(), rtol=1e-5, atol=1e-5)

 class TestGatedDeltaNetBlock(unittest.TestCase):
+  def test_kda_l2norm_matches_fla(self):
+    x = np.array([[1e-4, -2e-4, 3e-4], [1.0, 2.0, -3.0]], dtype=np.float32)
+    expected = x / np.sqrt((x*x).sum(axis=-1, keepdims=True) + 1e-6)
+    np.testing.assert_allclose(l2norm(Tensor(x)).numpy(), expected, rtol=1e-6, atol=1e-6)
+
  def _tensor_linspace(self, start:float, stop:float, shape:tuple[int, ...]) -> Tensor:
    return Tensor.linspace(start, stop, int(np.prod(shape)), dtype=dtypes.float32).reshape(*shape)

@@ -190,6 +197,87 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
    alpha = np.exp(-self._softplus_np(np.arange(1, 5)).reshape(1, 2, 1, 2))
    np.testing.assert_allclose(block.recurrent_state.numpy(), initial_state.numpy() * alpha, rtol=1e-5, atol=1e-5)

+  def test_kda_safe_gate_decay(self):
+    config = self._make_config(n_heads=2, kda_full_rank_gate=True, kda_gate_lower_bound=-5.0,
+      ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=2, time_step_rank=2, inner_size=4, kda=True))
+    block, x = GatedDeltaNetBlock(config, config.ssm), Tensor([[[1., 2., 0., 0.]]])
+    block.ssm_f_a.weight = Tensor([[1., 0., 0., 0.], [0., 1., 0., 0.]])
+    block.ssm_f_b.weight = Tensor([[1., 0.], [0., 1.], [1., 1.], [2., 1.]])
+    block.ssm_dt["bias"] = Tensor.zeros(4)
+    block.ssm_a = Tensor([[-2.], [-3.]])  # stores -exp(A_log)
+    block._init_state(x)
+    initial_state = Tensor.arange(8, dtype=dtypes.float32).reshape(1, 2, 2, 2)
+    block.recurrent_state.assign(initial_state).realize()
+    block._attention(x, 0).realize()
+    gate_logits = np.arange(1, 5, dtype=np.float32).reshape(1, 2, 2)
+    exp_a = np.array([2., 3.], dtype=np.float32).reshape(1, 2, 1)
+    alpha = np.exp(-5.0 / (1.0 + np.exp(-(exp_a * gate_logits)))).reshape(1, 2, 1, 2)
+    np.testing.assert_allclose(block.recurrent_state.numpy(), initial_state.numpy() * alpha, rtol=2e-5, atol=2e-5)
+
+  def test_kda_per_channel_a(self):
+    config = self._make_config(n_heads=2, kda_full_rank_gate=True, kda_gate_lower_bound=-5.0,
+      ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=2, time_step_rank=2, inner_size=4, kda=True, channel_decay=True))
+    block, x = GatedDeltaNetBlock(config, config.ssm), Tensor([[[1., 2., 0., 0.]]])
+    block.ssm_f_a.weight = Tensor([[1., 0., 0., 0.], [0., 1., 0., 0.]])
+    block.ssm_f_b.weight = Tensor([[1., 0.], [0., 1.], [1., 1.], [2., 1.]])
+    block.ssm_dt["bias"] = Tensor.zeros(4)
+    block.ssm_a = Tensor([[-2.], [-3.]])
+    block._init_state(x)
+    initial_state = Tensor.arange(8, dtype=dtypes.float32).reshape(1, 2, 2, 2)
+    block.recurrent_state.assign(initial_state).realize()
+    block._attention(x, 0).realize()
+    gate_logits = np.arange(1, 5, dtype=np.float32).reshape(1, 2, 2)
+    exp_a = np.array([2., 3.], dtype=np.float32).reshape(1, 1, 2)
+    alpha = np.exp(-5.0 / (1.0 + np.exp(-(exp_a * gate_logits)))).reshape(1, 2, 1, 2)
+    np.testing.assert_allclose(block.recurrent_state.numpy(), initial_state.numpy() * alpha, rtol=2e-5, atol=2e-5)
+
+  def test_kda_chunked_prefill_matches_decode(self):
+    config = self._make_config(max_context=4, n_heads=2,
+      ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=2, time_step_rank=2, inner_size=4, kda=True), kda_split_qkv=True)
+    x = Tensor.linspace(-1, 1, 4*config.dim, dtype=dtypes.float32).reshape(1, 4, config.dim).cast(dtypes.bfloat16)
+
+    chunked = GatedDeltaNetBlock(config, config.ssm)
+    sequential = GatedDeltaNetBlock(config, config.ssm)
+    for value in nn.state.get_state_dict(chunked).values(): value.replace(value.cast(dtypes.bfloat16).realize())
+    sequential_state = nn.state.get_state_dict(sequential)
+    for name, value in nn.state.get_state_dict(chunked).items(): sequential_state[name].replace(value)
+
+    chunked._init_state(x)
+    chunk_out = chunked._attention(x, 0).realize()
+    sequential._init_state(x)
+    seq_out = Tensor.cat(*[sequential._attention(x[:, t:t+1], t).realize() for t in range(x.shape[1])], dim=1).realize()
+
+    np.testing.assert_allclose(chunk_out.numpy(), seq_out.numpy(), rtol=1e-5, atol=1e-5)
+    for name in ("conv_state_q", "conv_state_k", "conv_state_v"):
+      np.testing.assert_allclose(getattr(chunked, name).numpy(), getattr(sequential, name).numpy(), rtol=2e-2, atol=4e-3)
+    np.testing.assert_allclose(chunked.recurrent_state.numpy(), sequential.recurrent_state.numpy(), rtol=2e-3, atol=2e-3)
+
+  def test_kda_tp_final_token_matches_unsharded(self):
+    config = self._make_config(dim=8, hidden_dim=16, n_heads=4, n_kv_heads=4, head_dim=2, rope_dim=2, v_head_dim=2,
+      ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=4, time_step_rank=4, inner_size=8, kda=True), kda_split_qkv=True)
+    single, tp = GatedDeltaNetBlock(config, config.ssm), GatedDeltaNetBlock(config, config.ssm)
+    for name, value in nn.state.get_state_dict(single).items():
+      data = np.full(value.shape, 1.0, np.float32) if "norm.weight" in name else \
+             np.linspace(-0.2, 0.2, value.numel(), dtype=np.float32).reshape(value.shape)
+      if name == "ssm_a": data.fill(-0.1)
+      value.replace(Tensor(data, device="CPU", dtype=dtypes.bfloat16).realize())
+    tp_state = nn.state.get_state_dict(tp)
+    for name, value in nn.state.get_state_dict(single).items(): tp_state[name].replace(value)
+    devices = ("CPU", "CPU:1")
+    _shard_kimi(SimpleNamespace(blk=[tp]), devices)
+
+    x = Tensor(np.linspace(-1, 1, 32, dtype=np.float32).reshape(1, 4, 8), device="CPU", dtype=dtypes.bfloat16)
+    single._init_state(x)
+    expected = single._attention(x, 0).realize()
+    x_tp = x.shard(devices, axis=None)
+    tp._init_state(x_tp)
+    actual = tp._attention(x_tp, 0).realize()
+
+    np.testing.assert_equal(actual.numpy(), expected.numpy())
+    np.testing.assert_equal(tp.recurrent_state.numpy(), single.recurrent_state.numpy())
+    for name in ("conv_state_q", "conv_state_k", "conv_state_v"):
+      np.testing.assert_equal(getattr(tp, name).numpy(), getattr(single, name).numpy())
+
 class TestPairwiseTopk(unittest.TestCase):
  def test_basic_topk(self):
    x = Tensor([[[1.0, 3.0, 2.0, 5.0, 4.0]]])
@@ -213,5 +301,13 @@ class TestPairwiseTopk(unittest.TestCase):
        self.assertEqual(set(sel.numpy()[b, t].tolist()), expected)
        np.testing.assert_allclose(vals.numpy()[b, t], data[b, t][sel.numpy()[b, t]])

+  def test_iterative_matches_numpy(self):
+    rng = np.random.default_rng(42)
+    data = rng.standard_normal((2, 3, 896), dtype=np.float32)
+    vals, sel = iterative_topk(Tensor(data), 16)
+    expected = np.argsort(-data, axis=-1, stable=True)[..., :16]
+    np.testing.assert_equal(sel.numpy(), expected)
+    np.testing.assert_allclose(vals.numpy(), np.take_along_axis(data, expected, axis=-1))
+
 if __name__ == '__main__':
  unittest.main()
@@ -212,6 +212,18 @@ class TestCallSchedule(unittest.TestCase):
    out = f(a, v.bind(5))
    np.testing.assert_allclose(out.numpy(), [5., 10., 15.])

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

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

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

+  def test_weak_expression_anchors_at_strong_lub(self):
+    # regression test for the HALF bert nan (#17408, reverted in #17409): lub(int32, weakfloat)==weakfloat makes
+    # `loss_mask.sum() + 1e-5` a weakfloat EXPRESSION. Meeting a strong float in a binop must pin it at the lub
+    denom = (Tensor.zeros(912, dtype=dtypes.int32) != Tensor.zeros(912, dtype=dtypes.float32)).sum() + 1e-5
+    self.assertIs(denom.dtype, dtypes.weakfloat)  # the setup: the denominator expression itself is weak
+    x, y = Tensor([2048.0], dtype=dtypes.float32)._broadcasted(denom)
+    self.assertIs(y.dtype, dtypes.float32)
+    recips = [u for u in (x / y)._uop.toposort() if u.op is Ops.RECIPROCAL]
+    self.assertEqual([(u.dtype, u.src[0].dtype) for u in recips], [(dtypes.float32, dtypes.float32)])
+    with Context(DEFAULT_FLOAT=dtypes.float16):
+      committed = graph_rewrite((UOp.const(1).cast(dtypes.int32) + UOp.const(1.0)).cast(dtypes.float32), pm_lower_index_dtype, ctx={})
+    self.assertEqual([u.dtype for u in committed.toposort() if u.op is Ops.ADD], [dtypes.float32])
+
+  def test_cast_weak_expression_commits_at_cast_floor(self):
+    # the floor never narrows: a cast BELOW the default does not pull the compute width down with it
+    with Context(DEFAULT_FLOAT=dtypes.float32):
+      narrowed = graph_rewrite((UOp.const(1.0) + UOp.const(2.0)).cast(dtypes.float16), pm_lower_index_dtype, ctx={})
+    self.assertEqual((narrowed.dtype, narrowed.src[0].dtype), (dtypes.float16, dtypes.float32))
+
+  def test_cast_weak_expression_value_uses_cast_floor(self):
+    with Context(DEFAULT_FLOAT=dtypes.float16):
+      denom = Tensor.ones(1, dtype=dtypes.int32, device="CPU").sum() * 70000 + 1e-5
+      out = Tensor(1.0, dtype=dtypes.float32, device="CPU") / denom
+      self.assertAlmostEqual(out.item(), 1 / (70000 + 1e-5), places=10)
+
  def test_uop_scalar_const_lifts_kind(self):
    for dtype, value, out_dtype, const_dtype in ((dtypes.weakint, 1, dtypes.weakint, dtypes.weakint),
                                                 (dtypes.int32, 1, dtypes.int32, dtypes.weakint),
@@ -0,0 +1,30 @@
+import unittest
+from tinygrad.llm.cli import KimiK3Template
+from tinygrad.llm.serve import StreamRouter
+
+class TestKimiK3Template(unittest.TestCase):
+  def test_simple_text_chat(self):
+    template = KimiK3Template()
+    got = template.render([{"role":"system", "content":"Be concise."}, {"role":"user", "content":"Hello"}])
+    self.assertTrue(got.startswith('<|open|>message role="system" type="thinking-effort"<|sep|>'))
+    self.assertIn('<|open|>message role="user"<|sep|>Hello<|close|>message<|sep|><|end_of_msg|>', got)
+    self.assertTrue(got.endswith('<|open|>message role="assistant"<|sep|><|open|>think<|sep|>'))
+
+  def test_preserves_assistant_thinking(self):
+    got = KimiK3Template().render([{"role":"assistant", "reasoning_content":"why", "content":"answer"}], add_generation_prompt=False)
+    self.assertIn('<|open|>think<|sep|>why<|close|>think<|sep|>', got)
+    self.assertIn('<|open|>response<|sep|>answer<|close|>response<|sep|>', got)
+
+  def test_rejects_unimplemented_modalities(self):
+    with self.assertRaisesRegex(ValueError, "text-only"):
+      KimiK3Template().render([{"role":"user", "content":[{"type":"image", "url":"x"}]}])
+    with self.assertRaisesRegex(ValueError, "tool rendering"):
+      KimiK3Template().render([{"role":"user", "content":"x"}], tools=[{"type":"function"}])
+
+  def test_xtml_stream_router(self):
+    router, routed = StreamRouter(reasoning=True, xtml=True), []
+    for piece in ("rea", "son<|close|>thi", "nk<|sep|><|open|>response<|sep|>ans", "wer<|close|>response<|sep|>"):
+      routed.extend(router.route(piece))
+    self.assertEqual(routed, [("reasoning_content", "rea"), ("reasoning_content", "son"), ("content", "ans"), ("content", "wer")])
+
+if __name__ == "__main__": unittest.main()
@@ -0,0 +1,139 @@
+import tempfile, unittest
+from pathlib import Path
+from dataclasses import replace
+import numpy as np
+from tinygrad import Tensor, dtypes, nn
+from tinygrad.helpers import getenv
+from tinygrad.llm.kernels import bf16_mfma_splitk
+from tinygrad.llm.kimi_k3 import KIMI_K3_FULL_ATTN_LAYERS, KIMI_K3_SSM_LAYERS, KIMI_K3_TEXT_SIZE, KIMI_K3_TP8_BYTES_PER_GPU, \
+  _layer_sources, _load_stacked_experts, _replace, _safe_load_selected, _shard_kimi_k3, _validate_config, kimi_k3_config, kimi_k3_smoke_config
+from tinygrad.llm.model import FFNBlock, Transformer
+
+def small_k3_config(max_context:int=4): return replace(kimi_k3_smoke_config(max_context), num_experts=8)
+
+class TestKimiK3(unittest.TestCase):
+  def test_selective_safetensor_load(self):
+    with tempfile.TemporaryDirectory() as tmp:
+      path = Path(tmp) / "weights.safetensors"
+      nn.state.safe_save({"keep":Tensor.arange(8), "skip":Tensor.arange(16)}, str(path))
+      selected = _safe_load_selected(path, ["keep"])
+      self.assertEqual(list(selected), ["keep"])
+      np.testing.assert_equal(selected["keep"].numpy(), np.arange(8))
+      with self.assertRaisesRegex(ValueError, "missing tensor absent"): _safe_load_selected(path, ["absent"])
+
+  def test_smoke_config_preserves_gfx950_expert_alignment(self):
+    c = kimi_k3_smoke_config()
+    self.assertEqual(c.routed_expert_dim % 64, 0)
+    self.assertEqual((c.hidden_dim // 8) % 64, 0)
+
+  @unittest.skipUnless(getenv("DEV", "") == "NULL:HIP:gfx950", "gfx950 compile coverage")
+  def test_gfx950_mfma_splitk_compile(self):
+    x = Tensor.zeros(1, 1, 256, dtype=dtypes.bfloat16, device="NULL:HIP:gfx950")
+    weight = Tensor.zeros(16, 256, dtype=dtypes.bfloat16, device="NULL:HIP:gfx950")
+    self.assertEqual(bf16_mfma_splitk(x, weight).realize().shape, (1, 1, 16))
+
+  def test_official_config(self):
+    c = kimi_k3_config(1_048_576)
+    self.assertEqual((c.num_blocks, c.dim, c.n_heads, c.num_experts, c.num_experts_per_tok), (93, 7168, 96, 896, 16))
+    self.assertEqual((sum(KIMI_K3_SSM_LAYERS), len(KIMI_K3_FULL_ATTN_LAYERS)), (69, 24))
+    self.assertEqual(KIMI_K3_FULL_ATTN_LAYERS, (*range(3, 93, 4), 92))
+    self.assertEqual((c.routed_expert_dim, c.hidden_dim, c.shared_expert_dim), (3584, 3072, 6144))
+    self.assertTrue(c.route_weights_uncorrected and c.kda_full_rank_gate and c.attn_output_gate)
+    self.assertTrue(c.ssm is not None and c.ssm.channel_decay)
+    self.assertEqual((c.activation_situ_beta, c.activation_situ_linear_beta, c.kda_gate_lower_bound), (4.0, 25.0, -5.0))
+
+  def test_config_rejects_wrong_checkpoint(self):
+    with self.assertRaisesRegex(ValueError, "not the supported official"):
+      _validate_config({"model_type":"kimi_linear", "hidden_size":2304})
+
+  def test_official_mapping_covers_model(self):
+    model = Transformer(kimi_k3_config(1))
+    state = nn.state.get_state_dict(model)
+    targets = {"token_embd.weight", "output_norm.weight", "output.weight", "output_attn_res_norm.weight", "output_attn_res_proj.weight"}
+    for i,is_kda in enumerate(KIMI_K3_SSM_LAYERS):
+      for target in _layer_sources(i, is_kda).values(): targets.update(target.split("|"))
+      if i:
+        for name in ("ffn_gate_exps.weight", "ffn_gate_exps.weight_scale", "ffn_up_exps.weight", "ffn_up_exps.weight_scale",
+                     "ffn_down_exps.weight", "ffn_down_exps.weight_scale"): targets.add(f"blk.{i}.{name}")
+    self.assertEqual(targets, set(state))
+    self.assertEqual(state["blk.1.ffn_gate_exps.weight"].shape, (896, 3072, 1792))
+    self.assertEqual(state["blk.1.ffn_gate_exps.weight_scale"].shape, (896, 3072, 112))
+    self.assertEqual(state["blk.0.ssm_a"].shape, (128, 1))
+    _shard_kimi_k3(model, tuple(f"NULL:{i}" for i in range(8)))
+    total, per_gpu = 0, 0
+    for name,value in state.items():
+      dtype = dtypes.uint8 if name.endswith(("weight_scale", "_exps.weight")) else dtypes.float32 if name.endswith(
+        ("exp_probs_b.bias", "ssm_q_conv1d.weight", "ssm_k_conv1d.weight", "ssm_v_conv1d.weight", "ssm_norm.weight", "ssm_a", "ssm_dt.bias")) \
+        else dtypes.bfloat16
+      size = value.numel() * dtype.itemsize
+      total += size
+      per_gpu += size if value.uop.axis is None else size//8
+    self.assertEqual((total, per_gpu), (KIMI_K3_TEXT_SIZE, KIMI_K3_TP8_BYTES_PER_GPU))
+
+  def test_situ_matches_reference(self):
+    block = FFNBlock(small_k3_config())
+    gate, up = Tensor([[-8., -1., 0., 3.]]), Tensor([[-30., -2., 5., 40.]])
+    got = block._activation(gate, up).numpy()
+    g, u = gate.numpy().astype(np.float32), up.numpy().astype(np.float32)
+    expected = (4*np.tanh(g/4)/(1+np.exp(-g))) * (25*np.tanh(u/25))
+    np.testing.assert_allclose(got, expected, rtol=1e-5, atol=1e-5)
+
+  def test_attention_residual_matches_reference(self):
+    block = FFNBlock(small_k3_config())
+    block.attn_res_norm.weight.assign([1.0+i/16 for i in range(32)])
+    block.attn_res_proj.weight.assign([[(-1.0)**i/8 for i in range(32)]])
+    prefix, residual = Tensor.arange(64).reshape(2, 32).float()/16, Tensor.arange(128).reshape(2, 2, 32).float()/32
+    got = block._apply_attn_res(prefix, residual, block.attn_res_proj, block.attn_res_norm).numpy()
+    v = np.concatenate((residual.numpy(), prefix.numpy()[:, None]), axis=1).astype(np.float32)
+    k = v / np.sqrt(np.mean(v*v, axis=-1, keepdims=True) + 1e-5)
+    scores = np.sum(k * block.attn_res_norm.weight.numpy() * block.attn_res_proj.weight.numpy()[0], axis=-1)
+    probs = np.exp(scores-scores.max(axis=-1, keepdims=True))
+    probs /= probs.sum(axis=-1, keepdims=True)
+    expected = np.matmul(probs[:, None], v).squeeze(1)
+    np.testing.assert_allclose(got, expected, rtol=1e-5, atol=1e-5)
+
+  def test_tp8_schema(self):
+    model = Transformer(small_k3_config())
+    _shard_kimi_k3(model, tuple(f"NULL:{i}" for i in range(8)))
+    state = nn.state.get_state_dict(model)
+    for name,axis in (("token_embd.weight",0), ("blk.1.ffn_gate_exps.weight",1), ("blk.1.ffn_down_exps.weight_scale",2),
+                      ("blk.1.ffn_routed_down.weight",1), ("blk.0.ssm_g_full.weight",0), ("blk.1.attn_q_b.weight",0)):
+      self.assertEqual(state[name].uop.axis, axis, name)
+    self.assertIsNone(state["blk.1.attn_res_norm.weight"].uop.axis)
+    self.assertIsNone(state["blk.1.ffn_routed_norm.weight"].uop.axis)
+    self.assertIsNone(state["blk.0.ssm_a"].uop.axis)
+
+  def test_direct_expert_staging(self):
+    devices = tuple(f"PYTHON:{i}" for i in range(4))
+    sources = [Tensor([[(e*40+r*4+c)&255 for c in range(4)] for r in range(8)], dtype=dtypes.uint8,
+                      device=devices[0]).realize() for e in range(8)]
+    expected = Tensor.stack(*sources).numpy()
+    for axis in (1, 2):
+      dst = Tensor.zeros(8, 8, 4, dtype=dtypes.uint8, device=devices[0]).shard(devices, axis=axis)
+      _load_stacked_experts(dst, sources)
+      np.testing.assert_equal(dst.numpy(), expected)
+
+  def test_direct_tp_replacement(self):
+    devices = tuple(f"PYTHON:{i}" for i in range(4))
+    source = Tensor.arange(64, dtype=dtypes.float32).reshape(8, 8).realize()
+    expected = source.numpy()
+    for axis in (None, 0, 1):
+      dst = Tensor.zeros(8, 8, device="PYTHON").shard(devices, axis=axis)
+      _replace(dst, source)
+      np.testing.assert_equal(dst.numpy(), expected)
+
+  def test_chunked_recurrent_generate(self):
+    model = Transformer(small_k3_config(max_context=8))
+    for name,value in nn.state.get_state_dict(model).items():
+      fill = 127 if name.endswith("weight_scale") else 0
+      value.replace(Tensor.full(value.shape, fill, dtype=value.dtype if value.dtype is dtypes.uint8 else dtypes.bfloat16, device="PYTHON"))
+    self.assertIsInstance(next(model.generate([1], chunk_size=2)), int)
+    prompt = [1, 2, 3, 4]
+    for _ in range(3): self.assertIsInstance(next(model.generate(prompt.copy(), chunk_size=2)), int)
+    self.assertEqual(model.get_start_pos(model._cached_tokens + [42]), len(prompt))
+    self.assertEqual(model.get_start_pos([9, 2, 3, 4, 42]), 0)
+    self.assertIsInstance(next(model.generate([1, 2, 3, 4, 5], chunk_size=3)), int)
+    self.assertEqual(set(model.recurrent_greedy_prefill_jits), {2})
+    self.assertEqual(model._cached_tokens[:4], [1, 2, 3, 4])
+
+if __name__ == "__main__": unittest.main()
@@ -0,0 +1,37 @@
+import unittest
+from tinygrad import dtypes, nn
+from tinygrad.llm.kimi import KIMI_LOGICAL_BYTES, KIMI_SSM_LAYERS, KIMI_TENSOR_COUNT, _shard_kimi, _validate_kimi_state, kimi_config
+from tinygrad.llm.model import Transformer
+
+class TestKimiLinear(unittest.TestCase):
+  def test_architecture_config(self):
+    config = kimi_config(4096)
+    self.assertEqual((config.num_blocks, config.dim, config.n_heads, config.vocab_size), (27, 2304, 32, 163840))
+    self.assertEqual(tuple(i for i, is_kda in enumerate(KIMI_SSM_LAYERS) if not is_kda), (3, 7, 11, 15, 19, 23, 26))
+    self.assertEqual((config.num_experts, config.num_experts_per_tok, config.shared_expert_dim), (256, 8, 1024))
+    self.assertTrue(config.expert_mxfp4 and config.bf16_activations and config.kda_split_qkv)
+    self.assertFalse(config.shared_expert_gate)
+
+  def test_tp4_schema_and_axes(self):
+    model = Transformer(kimi_config(32))
+    state = nn.state.get_state_dict(model)
+    self.assertEqual(len(state), KIMI_TENSOR_COUNT)
+    self.assertNotIn("blk.1.ffn_gate_inp_shexp.weight", state)
+    self.assertEqual(state["blk.1.ffn_gate_exps.weight"].dtype, dtypes.uint8)
+    self.assertEqual(state["blk.1.ffn_gate_exps.weight_scale"].dtype, dtypes.uint8)
+
+    _shard_kimi(model, ("NULL:0", "NULL:1", "NULL:2", "NULL:3"))
+    state = nn.state.get_state_dict(model)
+    for name, axis in (("token_embd.weight", 0), ("blk.1.ffn_gate_exps.weight", 1),
+                       ("blk.1.ffn_down_exps.weight_scale", 2), ("blk.3.attn_k_b.weight", 0)):
+      self.assertEqual(state[name].uop.axis, axis, name)
+    self.assertIsNone(state["blk.1.attn_norm.weight"].uop.axis)
+
+  def test_converted_schema_validation(self):
+    model = Transformer(kimi_config(1))
+    state = {name:value if value.dtype is dtypes.uint8 else value.cast(dtypes.bfloat16)
+             for name,value in nn.state.get_state_dict(model).items()}
+    _validate_kimi_state(model, state)
+    self.assertEqual(sum(value.nbytes() for value in state.values()), KIMI_LOGICAL_BYTES)
+
+if __name__ == "__main__": unittest.main()
@@ -72,6 +72,20 @@ class TestMoEFeedForward(unittest.TestCase):
    expected = (Tensor([1.0]).silu().item() + Tensor([3.0]).silu().item()) / 2
    np.testing.assert_allclose(out.numpy()[0, 0, 0], expected, rtol=1e-2)

+  def test_kimi_correction_bias_affects_route_weights(self):
+    dim, hidden, n_heads, num_experts, k = 8, 16, 2, 4, 2
+    config = replace(_moe_config(dim, hidden, n_heads, num_experts, k), norm_topk_prob=True, expert_bias=True)
+    block = TransformerBlock(config)
+    block.ffn_gate_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) * (i + 1) for i in range(num_experts)])
+    block.ffn_up_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) for _ in range(num_experts)])
+    block.ffn_down_exps.weight = Tensor.stack(*[Tensor.eye(dim, hidden) for _ in range(num_experts)])
+    block.ffn_gate_inp.weight = Tensor.zeros(num_experts, dim)
+    block.exp_probs_b["bias"] = Tensor([0.2, 0.1, 0.0, -0.1])
+
+    out = block._feed_forward(Tensor.ones(1, 1, dim))
+    expected = (Tensor([1.0]).silu().item() * 0.7 + Tensor([2.0]).silu().item() * 0.6) / 1.3
+    np.testing.assert_allclose(out.numpy()[0, 0, 0], expected, rtol=1e-2)
+
  def test_moe_feed_forward_shared_expert(self):
    dim, hidden, n_heads = 8, 16, 2
    num_experts, k = 4, 2
@@ -0,0 +1,68 @@
+import unittest
+import numpy as np
+from tinygrad import Tensor, dtypes
+from tinygrad.llm.quant import MXFP4_VALUES, dequantize_mxfp4, quantize_dequantize_mxfp8, quantize_mxfp4, quantize_mxfp4_cpu
+
+class TestMXFormats(unittest.TestCase):
+  def test_mxfp4_known_codes_and_scale(self):
+    values = np.array(MXFP4_VALUES * 2, dtype=np.float32)
+    packed, scale = quantize_mxfp4(Tensor(values))
+    # Positive and negative zero are numerically identical, so nearest-value encoding canonicalizes to +0.
+    np.testing.assert_array_equal(packed.numpy(), np.array([0x10, 0x32, 0x54, 0x76, 0x90, 0xba, 0xdc, 0xfe] * 2, dtype=np.uint8))
+    np.testing.assert_array_equal(scale.numpy(), np.array([127], dtype=np.uint8))
+    np.testing.assert_array_equal(dequantize_mxfp4(packed, scale, dtypes.float32).numpy(), values)
+
+  def test_mxfp4_block_scales_and_zero(self):
+    x = Tensor(np.array([0.0]*32 + [12.0, -12.0] + [0.0]*30, dtype=np.float32))
+    packed, scale = quantize_mxfp4(x)
+    np.testing.assert_array_equal(scale.numpy(), np.array([127, 128], dtype=np.uint8))
+    np.testing.assert_allclose(dequantize_mxfp4(packed, scale, dtypes.float32).numpy(), x.numpy())
+
+  def test_mxfp4_scale_rounds_amax_over_format_max(self):
+    # OCP E8M0 scale selection rounds log2(amax / 6), rather than flooring the
+    # input exponent. At this boundary the two rules differ by a factor of two.
+    x = Tensor(np.array([8.0] + [0.0]*31, dtype=np.float32))
+    packed, scale = quantize_mxfp4(x)
+    np.testing.assert_array_equal(scale.numpy(), np.array([127], dtype=np.uint8))
+    self.assertEqual(dequantize_mxfp4(packed, scale, dtypes.float32).numpy()[0], 6.0)
+
+  def test_mxfp4_cpu_converter_matches_tensor_path(self):
+    x = Tensor(np.linspace(-13, 13, 64*32, dtype=np.float32).reshape(64, 32))
+    packed, scale = quantize_mxfp4(x)
+    cpu_packed, cpu_scale = quantize_mxfp4_cpu(x)
+    np.testing.assert_array_equal(cpu_packed.numpy(), packed.numpy())
+    np.testing.assert_array_equal(cpu_scale.numpy(), scale.numpy())
+
+  def test_mxfp4_midpoints_round_to_even(self):
+    midpoints = np.array([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0], dtype=np.float32)
+    x = Tensor(np.pad(np.concatenate((midpoints, -midpoints)), (0, 18)))
+    packed, scale = quantize_mxfp4(x)
+    expected = np.pad(np.array([0, 1, 1, 2, 2, 4, 4, 0, -1, -1, -2, -2, -4, -4], dtype=np.float32), (0, 18))
+    np.testing.assert_array_equal(dequantize_mxfp4(packed, scale, dtypes.float32).numpy(), expected)
+
+  def test_mxfp8_roundtrip_and_dtype(self):
+    # All E4M3-exact values remain exact after extracting a shared exponent.
+    x = Tensor(np.array(([0.0, 0.5, 1.0, 1.5, 2.0, -3.0, 4.0, -6.0] * 4), dtype=np.float32))
+    out = quantize_dequantize_mxfp8(x)
+    self.assertEqual(out.dtype, dtypes.bfloat16)
+    np.testing.assert_array_equal(out.float().numpy(), x.numpy())
+
+  def test_mxfp8_subnormal_and_rounding(self):
+    x = np.zeros(32, dtype=np.float32)
+    x[:5] = [1.0, 1.0625, 1.07, 2**-9, 2**-10]
+    out = quantize_dequantize_mxfp8(Tensor(x), dtype=dtypes.float32).numpy()
+    # amax / 448 rounds to an E8M0 scale of 2**-9, saturating the largest
+    # values while retaining the E4M3 subnormal quantum for this block.
+    np.testing.assert_array_equal(out[:5], [0.875, 0.875, 0.875, 2**-9, 2**-10])
+
+  def test_mxfp8_uses_full_e4m3_range(self):
+    x = np.zeros(32, dtype=np.float32)
+    x[:4] = [448.0, 416.0, 400.0, -448.0]
+    np.testing.assert_array_equal(quantize_dequantize_mxfp8(Tensor(x), dtype=dtypes.float32).numpy()[:4], [448.0, 416.0, 384.0, -448.0])
+
+  def test_mxfp8_scale_rounds_amax_over_format_max(self):
+    x = np.zeros(32, dtype=np.float32)
+    x[0] = 512.0
+    self.assertEqual(quantize_dequantize_mxfp8(Tensor(x), dtype=dtypes.float32).numpy()[0], 448.0)
+
+if __name__ == "__main__": unittest.main()
@@ -1,9 +1,10 @@
 import unittest
+from dataclasses import replace
 from unittest.mock import patch
 from tinygrad import Tensor, UOp
 from tinygrad.schedule import schedule_cache
 from tinygrad.llm.model import Transformer, TransformerConfig
-from tinygrad.llm.serve import StreamRouter
+from tinygrad.llm.serve import StreamRouter, parse_kimi_tool_call

 TEST_CONFIG = TransformerConfig(num_blocks=1, dim=64, hidden_dim=128, n_heads=2, n_kv_heads=2,
                           norm_eps=1e-5, vocab_size=100, head_dim=32, rope_theta=10000.0, rope_dim=32, v_head_dim=32, max_context=32)
@@ -13,11 +14,29 @@ V_TOKS = UOp.variable("toks", 1, 32)  # 32 is the default chunk_size in generate
 class TestTransformerGenerate(unittest.TestCase):
  def test_warmup(self):
    model, calls = Transformer(TEST_CONFIG), []
-    def generate(tokens):
-      calls.append(tokens)
+    def generate(tokens, temperature):
+      calls.append((tokens, temperature))
      yield from (1, 2)
    with patch.object(model, "generate", generate): model.warmup()
-    self.assertEqual(calls, [[0], [0]])
+    self.assertEqual(calls, [([0], 0.0), ([0], 0.0)])
+
+  def test_recurrent_warmup_captures_reset_replay(self):
+    model, calls = Transformer(TEST_CONFIG), []
+    model.has_recurrent_block = True
+    state = Tensor.ones(4).realize()
+    model.blk[0]._state_reset_ops = lambda: [state.assign(state.const_like(0))]
+    def generate(tokens, temperature):
+      if calls: model.reset_jit()
+      calls.append((tokens.copy(), temperature))
+      tokens.append(42)
+      yield from (1, 2)
+    with patch.object(model, "generate", generate): model.warmup()
+    prompt = [0] * (TEST_CONFIG.max_context-2)
+    self.assertEqual(calls, [(prompt, 0.0)] * 3 + [(prompt, 1.0)] * 3 + [(prompt + list(range(1, i+1)), 0.0) for i in range(1, 4)])
+    self.assertEqual(model.reset_jit.cnt, 8)
+    cache_size = len(schedule_cache)
+    model.reset_jit()
+    self.assertEqual(len(schedule_cache), cache_size)

  def test_first_recurrent_generate_before_state_init(self):
    model = Transformer(TEST_CONFIG)
@@ -25,6 +44,17 @@ class TestTransformerGenerate(unittest.TestCase):
    with patch.object(Transformer, '__call__', return_value=Tensor([[42]])):
      self.assertEqual(next(model.generate([0])), 42)

+  def test_recurrent_prefill_tail_uses_rollout_shape(self):
+    model = Transformer(TEST_CONFIG)
+    model.has_recurrent_block = True
+    model.config = replace(model.config, recurrent_prefill_chunked=True)
+    calls = []
+    def mock_call(self, tokens, start_pos, temperature, **kwargs):
+      calls.append(tokens.shape)
+      return Tensor([[42]])
+    with patch.object(Transformer, '__call__', mock_call): next(model.generate([1, 2, 3, 4, 5, 6], chunk_size=4))
+    self.assertEqual(calls, [(1, 4), (1, 1), (1, 1)])
+
  def test_recurrent_live_state_reuse(self):
    model = Transformer(TEST_CONFIG)
    model.has_recurrent_block = True
@@ -38,11 +68,37 @@ class TestTransformerGenerate(unittest.TestCase):
      next(model.generate([1, 2, 3, 4, 5, 42, 10]))
    self.assertEqual(calls, [((1, 1), V_START_POS.bind(5)), ((1, 1), V_START_POS.bind(6))])

+  def test_recurrent_prompt_snapshot_reuse(self):
+    model = Transformer(TEST_CONFIG)
+    model.has_recurrent_block = True
+    state, calls = Tensor.ones(4).realize(), []
+    def mock_call(self, tokens, start_pos, temperature, **kwargs):
+      calls.append(start_pos)
+      return Tensor([[42]])
+    with patch.object(model, "_state_tensors", return_value=[state]), patch.object(model.blk[0], "_reusable_prefix_len", return_value=0), \
+         patch.object(Transformer, '__call__', mock_call):
+      next(model.generate([1, 2, 3]))
+      state.assign(state.const_like(5)).realize()
+      model._cached_tokens = [1, 2, 3, 9, 9]
+      calls.clear()
+      self.assertEqual(model.get_start_pos([1, 2, 3, 7, 8]), 3)
+      next(model.generate([1, 2, 3, 7, 8]))
+    self.assertEqual(calls, [V_START_POS.bind(3), V_START_POS.bind(4)])
+    self.assertEqual(state.tolist(), [1.0] * 4)
+
  def test_template_starts_reasoning(self):
    router = StreamRouter(reasoning=True)
    self.assertEqual(list(router.route("reasoning</think>answer")),
                     [("reasoning_content", "reasoning"), ("content", "answer")])

+  def test_kimi_tool_call_stream(self):
+    router = StreamRouter()
+    self.assertEqual(list(router.route("before<|tool_calls_section_beg")), [("content", "before")])
+    self.assertEqual(list(router.route("in|><|tool_call_begin|>functions.read:0<|tool_call_argument_begin|>"
+                                      '{"path":"/tmp/x"}<|tool_call_end|><|tool_calls_section_end|>')), [])
+    self.assertEqual(parse_kimi_tool_call("functions.read:0<|tool_call_argument_begin|>{\"path\":\"/tmp/x\"}"),
+                     ("read", {"path":"/tmp/x"}))
+
  def test_kv_cache_reuse(self):
    """Test that generate reuses the KV cache when tokens extend the cached prefix."""
    model = Transformer(TEST_CONFIG)
@@ -2,7 +2,7 @@ import unittest, numpy as np
 from tinygrad import Tensor, Variable, Context, Device, TinyJit, GlobalCounters, dtypes, UOp, nn, getenv
 from tinygrad.nn.state import get_parameters, get_state_dict
 from tinygrad.uop.ops import Ops
-from test.helpers import not_support_multi_device, needs_second_gpu, slow, assert_kernel_count
+from test.helpers import not_support_multi_device, needs_second_gpu, slow, assert_kernel_count, KernelCountException
 from hypothesis import given, strategies as strat, settings

 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
@@ -384,6 +384,12 @@ class TestMultiTensor(unittest.TestCase):
      np.testing.assert_allclose(r.numpy(), np.ones(256)+np.ones(256), atol=1e-4, rtol=1e-5)
    assert jf.captured is not None

+  def test_symbolic_broadcast_copy(self):
+    rows = Variable("rows", 1, 4).bind(3)
+    out = Tensor.ones(rows, 8).to(devices_2).realize()
+    self.assertEqual(out.shape, (rows, 8))
+    np.testing.assert_equal(out[:3].to(Device.DEFAULT).numpy(), np.ones((3, 8)))
+
  def test_multitensor_jit_in_list(self):
    # test MULTI tensor inside a list container - exercises the container unpacking + MULTI unpacking
    @TinyJit
@@ -583,7 +589,7 @@ class TestMultiTensor(unittest.TestCase):
      zeros = Tensor.zeros(3).realize()
    b = a.to(devices_2)*zeros.to(devices_2)
    sched = b.schedule_linear().src
-    self.assertEqual(len(sched), 0)
+    if len(sched) != 0: raise KernelCountException(0, len(sched))
    self.assertListEqual(b.tolist(), [0, 0, 0])

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

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

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

    # optimize (schedule) the AST
    sink = graph_rewrite(sink, pm_flatten_range+pm_simplify_ranges, ctx={}, name="simplify ranges")
@@ -310,7 +310,8 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
    sink = apply_opts(sink, ren, beam=ast.arg.beam)

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

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

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

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

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

  # final symbolic before decomp
  sink = graph_rewrite(sink, symbolic, name="final symbolic")
@@ -354,7 +357,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:

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

  # late decomps + move gates from unrenderable INVALID where
@@ -448,7 +451,7 @@ pm_to_program = PatternMatcher([
  (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.LINEAR), UPat(Ops.SOURCE, name="source")), name="prg"), do_compile),
 ])

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

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

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

@@ -1,8 +1,8 @@
 import itertools, functools
 from collections import defaultdict
 from tinygrad.dtype import dtypes, AddrSpace, Invalid, DType
-from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp, shape_to_shape_arg
-from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate
+from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp, shape_to_shape_arg, graph_rewrite
+from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate, sym
 from tinygrad.helpers import getenv, IMAGE, OSX, ceildiv, is_image_shape
 from tinygrad.renderer import Renderer

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

 def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
@@ -332,9 +332,9 @@ class Scheduler:
  @property
  def group_for_reduces(self) -> int: return len(self.axes_of(AxisType.GROUP_REDUCE))

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

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

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

  try:
    rawbufs = _ensure_buffer_alloc(rawbufs)
-    var_vals: dict[str, int] = {k.expr:int(k.vmax+k.vmin)//2 for k in s.ast.variables()}
    exiting, st = False, time.perf_counter()
    dev = Device[s.ren.target.device]
    while not exiting:
@@ -1,7 +1,7 @@
 import itertools
 from typing import Callable
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start, AxisType
-from tinygrad.uop.symbolic import symbolic, invalid_gate
+from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const, invalid_gate
 from tinygrad.helpers import partition
 from tinygrad.dtype import dtypes

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

        # check if it simplifies
@@ -80,7 +80,7 @@ class DType(metaclass=DTypeMetaClass):
    # NOTE: float('nan') != float('nan'), so we canonicalize here
    if isinstance(val, float) and math.isnan(val): val = math.nan
    # int is the default. wrap floats in ConstFloat to distinguish -0.0 from 0.0 in cache
-    return ConstFloat(float(val)) if dtypes.is_float(self) else bool(val) if dtypes.is_bool(self) else int(val)
+    return ConstFloat(truncate.get(self, float)(float(val))) if dtypes.is_float(self) else bool(val) if dtypes.is_bool(self) else int(val)


 class DTypes:
@@ -4,7 +4,7 @@ from tinygrad.tensor import Tensor, all_tensors
 from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv, JIT, JIT_BATCH_SIZE, dedup, pluralize, VIZ, disable_gc
 from tinygrad.device import Buffer, Compiled, Device, MultiBuffer, DepsTracker
 from tinygrad.dtype import DType
-from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, track_rewrites, graph_rewrite
+from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, rewrite_group, graph_rewrite
 from tinygrad.renderer import Estimates
 from tinygrad.engine.realize import capturing, compile_linear, link_linear, run_linear, graph_cache, estimate_uop, get_runtime
 from tinygrad.engine.realize import unwrap_multi, resolve_params, get_call_arg_uops, get_call_outs_ins
@@ -64,7 +64,7 @@ def _copy_input(u:UOp) -> UOp:
  run_linear(UOp(Ops.LINEAR, src=(u.copy_to_device(u.device).call(new:=UOp.new_buffer(u.device, u.max_numel(), u.dtype), u),)))
  return new

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

@@ -3,12 +3,12 @@ from typing import cast, Iterator, Any, Sequence
 import time, random, itertools, math, contextlib, weakref, array
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, prod, flatten, Context, getenv, to_tuple
-from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events
+from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, wait_cond
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, buffers, graph_rewrite
 from tinygrad.device import Device, Buffer, MultiBuffer
 from tinygrad.renderer import Estimates
 from tinygrad.codegen import to_program
-from tinygrad.codegen.opt.postrange import bufs_from_ast
+from tinygrad.codegen.opt.postrange import args_from_ast

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

@@ -33,7 +33,7 @@ def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|N
  if ast.op is Ops.COPY: return colored(f"copy {_uop_sz_to_str(arg_uops[0]):>10}, {_dev_str(bufs[0]):>7s} <- {_dev_str(bufs[1]):7s}", "yellow")
  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return colored(f"enc/dec {_uop_sz_to_str(arg_uops[0])}", "yellow")
  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return colored(f"batched {len(ast.src[0].src)}", "cyan")
-  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq": return call.arg.aux.name
+  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq": return cast(str, call.arg.name)
  raise NotImplementedError("get_call_name is not implemented")

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

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

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

  exec_kernel(replace(ctx, update_stats=False), call, ast)

-  st = time.perf_counter()
-  for d in call.arg.aux.device:
-    with track_stats(ctx, call, d, [], ctx.var_vals):
-      if ctx.wait: Device[d].synchronize()
-  return time.perf_counter() - st
+  tms:list[float|None] = []
+  for e in (aux:=call.arg.aux).prof: cast(Any, Device[e.device]).prof_ents[e.st_id] = e
+  for d in [cast(Any, Device[x]) for x in aux.device]:
+    with track_stats(ctx, call, d.device, [], ctx.var_vals) as et:
+      if ctx.wait:
+        d.synchronize(timeout=ctx.timeout)
+        ts = [d.signal(i)._buf.cpu_view().view(fmt='Q')[0] for e in aux.prof if e.device == d.device for i in (e.st_id, e.en_id)]
+        if ts: et[0] = float(max(ts)-min(ts))/d.timestamp_divider/1e6
+      tms += et
+  return tms[0]

 # flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src
 pm_flatten_linear = PatternMatcher([
@@ -265,11 +272,11 @@ pm_exec = PatternMatcher([

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

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

 def link_linear(linear:UOp, cache=True) -> UOp: return hcq_link(linear, cache=cache) if getenv("HCQ2") else linear
@@ -287,5 +294,5 @@ def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None
      from tinygrad.tensor import Tensor
      with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024, 1024).contiguous().realize(do_update_stats=False)
  ctx = ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False)
-  linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0), cache=ctx.cache)
+  linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0, profile=True), cache=ctx.cache)
  return max(pm_exec.rewrite(c, ctx) or 0.0 for c in linear.src)
@@ -251,7 +251,7 @@ DEFAULT_FLOAT, DEFAULT_INT = ContextVar("DEFAULT_FLOAT", "float32"), ContextVar(
 CAPTURE_PROCESS_REPLAY = ContextVar("CAPTURE_PROCESS_REPLAY", 0)
 def _get_cpu_count() -> int:
  # os.process_cpu_count (3.13+) respects cgroup limits
-  if hasattr(os, "process_cpu_count"): return max(1, os.process_cpu_count())
+  if hasattr(os, "process_cpu_count"): return max(1, os.process_cpu_count() or 1)
  # cgroup v2 (containers with --cpus=N)
  try:
    with open("/sys/fs/cgroup/cpu.max") as f:
@@ -21,7 +21,7 @@
    const d = document.createElement('div'); d.className = 'msg'; chat.appendChild(d);
    const r = await fetch('/v1/chat/completions', {method: 'POST', headers: {'Content-Type': 'application/json'},
      body: JSON.stringify({model: 'llama', messages: msgs, stream: true, temperature: 0.7})});
-    let buf = '';
+    let buf = '', txt = '', rsn = '';
    for (const rd = r.body.getReader(), dec = new TextDecoder();;) {
      const {done, value} = await rd.read();
      if (done) break;
@@ -30,9 +30,13 @@
      buf = lines.pop();
      for (const ln of lines)
        if (ln.startsWith('data: ') && !ln.includes('[DONE]'))
-          try { d.textContent += JSON.parse(ln.slice(6)).choices[0]?.delta?.content || '' } catch {}
+          try { const dl = JSON.parse(ln.slice(6)).choices[0]?.delta;
+            if (dl?.reasoning_content) { const s = document.createElement('span'); s.style.color = '#888';
+              s.textContent = dl.reasoning_content; rsn += dl.reasoning_content; d.appendChild(s) }
+            if (dl?.content) { const s = document.createElement('span');
+              s.textContent = dl.content; txt += dl.content; d.appendChild(s) } } catch {}
      chat.scrollTop = chat.scrollHeight;
    }
-    msgs.push({role: 'assistant', content: d.textContent});
+    const m = {role:'assistant', content:txt}; if (rsn) m.reasoning_content = rsn; msgs.push(m);
  }
 </script></body></html>
@@ -1,5 +1,5 @@
 from __future__ import annotations
-import sys, argparse, codecs, itertools, typing, re, unicodedata, json, time
+import sys, argparse, codecs, itertools, typing, re, unicodedata, json, time, pathlib
 from typing import TYPE_CHECKING
 from tinygrad import nn
 from tinygrad.uop.ops import UOp, Ops
@@ -127,30 +127,95 @@ class FallbackTemplate:
      out += self.end_turn()
    return out + self.role("assistant") if add_generation_prompt else out

-from tinygrad.llm.serve import LLMServer
+class KimiK3Template:
+  """Official K3 XTML envelope for text-only system/user/assistant conversations."""
+  OPEN, CLOSE, SEP, END = "<|open|>", "<|close|>", "<|sep|>", "<|end_of_msg|>"
+  def _open(self, tag:str, attrs:tuple[tuple[str, str], ...]=()) -> str:
+    escaped = ((k, str(v).replace("&", "&amp;").replace('"', "&quot;")) for k,v in attrs)
+    return self.OPEN + tag + "".join(f' {k}="{v}"' for k,v in escaped) + self.SEP
+  def _close(self, tag:str) -> str: return self.CLOSE + tag + self.SEP
+  def _message(self, role:str, content:str, name:str|None=None) -> str:
+    attrs = (("role", role),) + (() if name is None else (("name", name),))
+    return self._open("message", attrs) + content + self._close("message") + self.END
+  @staticmethod
+  def _content(message:dict) -> str:
+    content = message.get("content")
+    if content is None: return ""
+    if isinstance(content, str): return content
+    if isinstance(content, list):
+      if any(part.get("type") != "text" for part in content): raise ValueError("Kimi K3 native loader is text-only; image content is not implemented")
+      return "".join(part["text"] for part in content)
+    raise ValueError(f"unsupported Kimi K3 content type {type(content).__name__}")
+  def render(self, messages:list[dict], tools=None, add_generation_prompt:bool=True, preserve_thinking:bool=False, **kwargs) -> str:
+    if tools or any(m.get("role") == "tool" or m.get("tool_calls") for m in messages):
+      raise ValueError("Kimi K3 XTML tool rendering is not implemented in the native text loader")
+    effort = kwargs.get("thinking_effort", "max")
+    if effort not in ("low", "high", "max"): raise ValueError(f"invalid Kimi K3 thinking_effort {effort!r}")
+    body = "`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), " \
+           "supported values include `low`, `medium`, `high`, and `max`.\n" \
+           f"Now the system is invoked with `thinking_effort={effort}`."
+    out = self._open("message", (("role", "system"), ("type", "thinking-effort"))) + body + self._close("message") + self.END
+    for message in messages:
+      role = message["role"]
+      if role in ("user", "system"):
+        out += self._message(role, self._content(message), message.get("name"))
+      elif role == "assistant":
+        reasoning = message.get("reasoning_content") or message.get("reasoning") or ""
+        content = self._open("think") + str(reasoning) + self._close("think")
+        content += self._open("response") + self._content(message) + self._close("response")
+        out += self._message(role, content, message.get("name"))
+      else: raise ValueError(f"unsupported Kimi K3 role {role!r}")
+    if add_generation_prompt: out += self._open("message", (("role", "assistant"),)) + self._open("think")
+    return out
+
+from tinygrad.llm.serve import LLMServer, StreamRouter

 def main():
  parser = argparse.ArgumentParser()
-  parser.add_argument("--model", "-m", default=list(models.keys())[0], help=f"Model choice ({', '.join(models.keys())}) or path to a local GGUF file")
+  parser.add_argument("--model", "-m", default=list(models.keys())[0],
+                      help=f"Model choice ({', '.join(models.keys())}), local GGUF file, converted Kimi directory, or official Kimi K3 directory")
  parser.add_argument("--max_context", type=int, default=4096, help="Max Context Length")
  parser.add_argument("--serve", nargs='?', type=int, const=8000, metavar="PORT", help="Run OpenAI compatible API (optional port, default 8000)")
  parser.add_argument("--warmup", action="store_true", help="warmup the JIT")
  parser.add_argument("--benchmark", nargs='?', type=int, const=20, metavar="COUNT", help="Benchmark tok/s (optional count, default 20)")
+  parser.add_argument("--devices", type=int, default=1, help="Tensor-parallel device count (Kimi-Linear requires 4, Kimi K3 requires 8)")
  args = parser.parse_args()

  # load the model
-  model, kv = Transformer.from_gguf(fetch(models.get(args.model, args.model)), args.max_context)
-  model_name = kv.get('general.name') or kv.get('general.basename') or args.model
+  model_path = pathlib.Path(args.model)
+  kv:dict[str, typing.Any]
+  is_k3 = False
+  if model_path.is_dir() and (model_path / "config.json").exists():
+    raw_config = json.loads((model_path / "config.json").read_text())
+    is_k3 = raw_config.get("model_type") == "kimi_k3"
+  if is_k3:
+    from tinygrad.llm.kimi_k3 import load_kimi_k3, load_kimi_tokenizer_data
+    model, kv = load_kimi_k3(model_path, args.max_context, args.devices), {}
+    normal, special, bos, eos = load_kimi_tokenizer_data(model_path)
+    tok = SimpleTokenizer(normal, special, "kimi-k2", bos_id=bos, eos_id=eos, eot_id=eos)
+    model_name = "Kimi-K3"
+    tok_cfg = json.loads((model_path / "tokenizer_config.json").read_text())
+    ct = tok_cfg.get("chat_template")
+  elif model_path.is_dir() and (model_path / "tinygrad-kimi.json").exists():
+    from tinygrad.llm.kimi import load_kimi, load_kimi_tokenizer_data
+    model, kv = load_kimi(model_path, args.max_context, args.devices), {}
+    normal, special, bos, eos = load_kimi_tokenizer_data(model_path)
+    tok = SimpleTokenizer(normal, special, "kimi-k2", bos_id=bos, eos_id=eos, eot_id=eos)
+    model_name = "Kimi-Linear-48B-A3B-Instruct-MXFP4"
+    ct = (model_path / "chat_template.jinja").read_text() if (model_path / "chat_template.jinja").exists() else None
+  else:
+    if args.devices != 1: raise ValueError("--devices is currently supported by the native Kimi MXFP4 loader only")
+    model, kv = Transformer.from_gguf(fetch(models.get(args.model, args.model)), args.max_context)
+    model_name = kv.get('general.name') or kv.get('general.basename') or args.model
+    tok = SimpleTokenizer.from_gguf_kv(kv)
+    ct = kv.get('tokenizer.chat_template')
  file_sizes = [y.nbytes() for y in UOp.sink(*[x.uop for x in nn.state.get_parameters(model)]).toposort() if y.op is Ops.BUFFER]
  print(f"using model \"{model_name}\" with {sum(file_sizes):,} bytes and {sum(x.numel() for x in nn.state.get_parameters(model)):,} params, "
        f"max context {args.max_context} on {nn.state.get_parameters(model)[0].device}")

-  # get tokenizer
-  tok = SimpleTokenizer.from_gguf_kv(kv)
-
  # use the model's chat template if jinja2 is available (enables model-specific formatting)
-  template: jinja2.Template|FallbackTemplate = FallbackTemplate(tok)
-  if (ct := kv.get('tokenizer.chat_template')) is not None:
+  template: jinja2.Template|FallbackTemplate|KimiK3Template = KimiK3Template() if is_k3 else FallbackTemplate(tok)
+  if ct is not None:
    try:
      import jinja2
      env = jinja2.Environment()
@@ -162,9 +227,10 @@ def main():
      template = env.from_string(ct)
    except ImportError: print("warning: jinja2 is not installed, the model's chat template is disabled")

-  # warmup the JIT
-  if args.warmup or args.serve:
+  # Capture the default greedy serving shapes before accepting requests.
+  if args.warmup:
    with Context(DEBUG=max(DEBUG.value, 1)): model.warmup()
+  elif args.serve: model.warmup()

  # start server
  if args.serve: LLMServer(('', args.serve), model, model_name, tok, template).serve_forever()
@@ -189,15 +255,26 @@ def main():
  while 1:
    try: messages.append({"role":"user", "content":input('>>> ')})
    except EOFError: break
-    ids = tok.encode(template.render(messages=messages, add_generation_prompt=True))
-    reply, dec = "", tok.stream_decoder()
+    rendered = template.render(messages=messages, add_generation_prompt=True)
+    ids = tok.encode(rendered)
+    reply, reasoning_reply, dec = "", "", tok.stream_decoder()
+    xtml = rendered.rstrip().endswith("<|open|>think<|sep|>")
+    router = StreamRouter(reasoning=xtml or rendered.rstrip().endswith("<think>"), xtml=xtml)
    for next_id in model.generate(ids):
      if tok.is_end(next_id):
-        sys.stdout.write(dec() + "\n\n")
+        for field,text in router.route(dec(), final=True):
+          if field == "content": reply += text
+          elif field == "reasoning_content": reasoning_reply += text
+          sys.stdout.write(text)
+        sys.stdout.write("\n\n")
        break
-      reply += (piece := dec(next_id))
-      sys.stdout.write(piece)
-      sys.stdout.flush()
-    messages.append({"role":"assistant", "content":reply})
+      for field,text in router.route(dec(next_id)):
+        if field == "content": reply += text
+        elif field == "reasoning_content": reasoning_reply += text
+        sys.stdout.write(text)
+        sys.stdout.flush()
+    assistant = {"role":"assistant", "content":reply}
+    if reasoning_reply: assistant["reasoning_content"] = reasoning_reply
+    messages.append(assistant)

 if __name__ == "__main__": main()
@@ -0,0 +1,238 @@
+import functools
+from typing import cast
+from tinygrad import Tensor, UOp, Device, Context, dtypes
+from tinygrad.device import Buffer, MultiBuffer
+from tinygrad.dtype import AddrSpace
+from tinygrad.uop.ops import AxisType, KernelInfo
+
+def amd_custom_kernels_supported(device:str|tuple[str, ...]|None) -> bool:
+  """The hand-written wave32 kernel is intentionally limited to RDNA3/gfx11."""
+  if device is None: return False
+  device = device[0] if isinstance(device, tuple) else device
+  with Context(ALLOW_DEVICE_USAGE=1):
+    return (target:=getattr(Device[device], "target", None)) is not None and target[0] == 11
+
+def amd_wave64_custom_kernels_supported(device:str|tuple[str, ...]|None) -> bool:
+  """CDNA4 wave64 kernels used by the MI350X K3 path."""
+  if device is None: return False
+  device = device[0] if isinstance(device, tuple) else device
+  if device.startswith("NULL:HIP:gfx950"): return True  # compile-only CDNA4 coverage without MI350X hardware
+  with Context(ALLOW_DEVICE_USAGE=1):
+    return (target:=getattr(Device[device], "target", None)) is not None and target[:2] == (9, 5)
+
+def amd_packed_mxfp4_supported(device:str|tuple[str, ...]|None) -> bool:
+  return amd_custom_kernels_supported(device) or amd_wave64_custom_kernels_supported(device)
+
+def amd_exact_bf16_custom_kernels_supported(device:str|tuple[str, ...]|None) -> bool:
+  """The LDS-reduced exact BF16 pair kernels are portable across RDNA3 and CDNA4."""
+  return amd_custom_kernels_supported(device) or amd_wave64_custom_kernels_supported(device)
+
+def amd_int32_item(x:Tensor, host:memoryview) -> int:
+  """Copy a realized replicated AMD scalar without constructing a new scheduler graph."""
+  if x.numel() != 1 or x.dtype != dtypes.int32 or host.nbytes != 4: raise ValueError("expected one int32 and a four-byte host view")
+  buf = x.uop.buffer
+  if isinstance(buf, MultiBuffer): buf = buf.bufs[0]
+  if not isinstance(buf, Buffer) or not buf.device.startswith("AMD"): raise ValueError("expected a realized AMD buffer")
+  buf.allocator._copyout(host, buf._buf)
+  return int.from_bytes(host, byteorder="little", signed=True)
+
+def mxfp4_expert_linear(sel:Tensor, x:Tensor, weight:Tensor, scale:Tensor, partial:bool=False) -> Tensor:
+  """Run a TP routed projection without materializing selected BF16 weights."""
+  from tinygrad.llm.kernels.amd import (_mxfp4_expert_linear_kernel, _mxfp4_expert_linear_wave64_kernel,
+                                        _mxfp4_expert_linear_wave64_prefill_kernel)
+  batch, tokens, topk = sel.shape
+  out_features = weight.shape[1]
+  weight_axis = weight.uop.axis
+  if isinstance(weight.device, tuple):
+    devices = weight.device
+    # Gate/up shard their output dimension. Down shards its reduction dimension;
+    # represent each GPU's partial as a size-one device axis, then all-reduce it.
+    axis = 3 if weight_axis == 1 else 4
+    shard_shape: tuple[int|UOp, ...]
+    if weight_axis == 1:
+      if out_features % len(devices): raise ValueError(f"expert output {out_features} is not divisible by {len(devices)} devices")
+      shard_shape = (batch, tokens, topk, out_features//len(devices))
+    elif weight_axis == 2:
+      shard_shape = (batch, tokens, topk, out_features, 1)
+    else: raise ValueError(f"unsupported expert TP axis {weight_axis}")
+    partial_dtype = dtypes.float32 if weight_axis == 2 else dtypes.bfloat16
+    parts = [Tensor.empty(*shard_shape, dtype=partial_dtype, device=device).uop for device in devices]
+    out = Tensor(parts[0].mstack(*parts[1:]).unshard(axis))
+  else:
+    out = Tensor.empty(batch, tokens, topk, out_features, dtype=dtypes.bfloat16, device=weight.device)
+  if amd_wave64_custom_kernels_supported(weight.device):
+    kernel = _mxfp4_expert_linear_wave64_prefill_kernel if tokens > 1 else _mxfp4_expert_linear_wave64_kernel
+  else: kernel = _mxfp4_expert_linear_kernel
+  out = Tensor.custom_kernel(out, sel.contiguous(), x.contiguous(), weight, scale, fxn=kernel)[0]
+  return out if weight_axis == 2 and partial else out.sum(4).cast(dtypes.bfloat16) if weight_axis == 2 else out
+
+def bf16_partial_linear(x:Tensor, weight:Tensor) -> Tensor:
+  """Return output-shaped FP32 TP partials with a final device axis, without all-reduce."""
+  from tinygrad.llm.kernels.amd import _bf16_partial_linear_kernel
+  if not isinstance(weight.device, tuple) or weight.uop.axis != 1: raise ValueError("partial linear expects input-sharded TP weight")
+  batch, tokens, _ = x.shape
+  devices, out_features = weight.device, weight.shape[0]
+  shard_shape = (batch, tokens, out_features, 1)
+  parts = [Tensor.empty(*shard_shape, dtype=dtypes.float32, device=device).uop for device in devices]
+  out = Tensor(parts[0].mstack(*parts[1:]).unshard(3))
+  return Tensor.custom_kernel(out, x.contiguous(), weight, fxn=_bf16_partial_linear_kernel)[0]
+
+def bf16_matvec(x:Tensor, weight:Tensor) -> Tensor:
+  from tinygrad.llm.kernels.amd import _bf16_matvec_kernel
+  batch, tokens, _ = x.shape
+  out_features = weight.shape[0]
+  if isinstance(weight.device, tuple):
+    devices = weight.device
+    if weight.uop.axis == 0:
+      shard_shape = (batch, tokens, out_features//len(devices))
+      parts = [Tensor.empty(*shard_shape, dtype=dtypes.bfloat16, device=device).uop for device in devices]
+      out = Tensor(parts[0].mstack(*parts[1:]).unshard(2))
+    elif weight.uop.axis is None: out = Tensor.empty(batch, tokens, out_features, dtype=dtypes.bfloat16, device=weight.device)
+    else: raise ValueError("bf16_matvec expects output-sharded or replicated TP weight")
+  else: out = Tensor.empty(batch, tokens, out_features, dtype=dtypes.bfloat16, device=weight.device)
+  return Tensor.custom_kernel(out, x.contiguous(), weight, fxn=_bf16_matvec_kernel)[0]
+
+def bf16_mfma_splitk(x:Tensor, weight:Tensor) -> Tensor:
+  """gfx950 decode matvec for replicated or output-sharded BF16 weights."""
+  from tinygrad.llm.kernels.amd import _bf16_mfma_splitk_kernel
+  batch, tokens, in_features = x.shape
+  out_features = weight.shape[0]
+  if batch != 1 or tokens != 1 or weight.shape[1] != in_features or in_features % 256:
+    raise ValueError(f"unsupported MFMA split-K shapes {x.shape} {weight.shape}")
+  if not amd_wave64_custom_kernels_supported(weight.device): raise ValueError("MFMA split-K requires gfx950")
+  if isinstance(weight.device, tuple):
+    devices = weight.device
+    if weight.uop.axis == 0:
+      if out_features % (16*len(devices)): raise ValueError("local MFMA output must be divisible by 16")
+      shape = (batch, tokens, out_features//len(devices))
+      parts = [Tensor.empty(*shape, dtype=dtypes.bfloat16, device=device).uop for device in devices]
+      out = Tensor(parts[0].mstack(*parts[1:]).unshard(2))
+    elif weight.uop.axis is None:
+      if out_features % 16: raise ValueError("MFMA output must be divisible by 16")
+      out = Tensor.empty(batch, tokens, out_features, dtype=dtypes.bfloat16, device=devices)
+    else: raise ValueError("MFMA split-K requires replicated or output-sharded weight")
+  else:
+    if out_features % 16: raise ValueError("MFMA output must be divisible by 16")
+    out = Tensor.empty(batch, tokens, out_features, dtype=dtypes.bfloat16, device=weight.device)
+  return Tensor.custom_kernel(out, x.contiguous(), weight, fxn=_bf16_mfma_splitk_kernel)[0]
+
+def mxfp8_quantize_dequantize(x:Tensor) -> Tensor:
+  """gfx11 software MXFP8 round trip without a multi-kernel reduction graph."""
+  from tinygrad.llm.kernels.amd import _mxfp8_qdq_kernel
+  out = Tensor.empty_like(x, dtype=dtypes.bfloat16)
+  return Tensor.custom_kernel(out, x.contiguous(), fxn=_mxfp8_qdq_kernel)[0]
+
+def kda_qkv_linear(x:Tensor, qw:Tensor, kw:Tensor, vw:Tensor) -> tuple[Tensor, Tensor, Tensor]:
+  """Fuse equal-sized output-sharded KDA Q/K/V decode projections."""
+  from tinygrad.llm.kernels.amd import _kda_qkv_kernel
+  batch, tokens, _ = x.shape
+  out_features = qw.shape[0]
+  if not (qw.shape == kw.shape == vw.shape): raise ValueError("fused KDA Q/K/V weights must have equal shapes")
+  if isinstance(qw.device, tuple):
+    devices = qw.device
+    if qw.uop.axis != 0 or out_features % len(devices): raise ValueError("fused KDA Q/K/V expects output-sharded weights")
+    shard_shape = (batch, tokens, out_features//len(devices))
+    def make_out() -> Tensor:
+      parts = [Tensor.empty(*shard_shape, dtype=dtypes.bfloat16, device=device).uop for device in devices]
+      return Tensor(parts[0].mstack(*parts[1:]).unshard(2))
+    outs: tuple[Tensor, Tensor, Tensor] = (make_out(), make_out(), make_out())
+  else:
+    outs = (Tensor.empty(batch, tokens, out_features, dtype=dtypes.bfloat16, device=qw.device),
+            Tensor.empty(batch, tokens, out_features, dtype=dtypes.bfloat16, device=qw.device),
+            Tensor.empty(batch, tokens, out_features, dtype=dtypes.bfloat16, device=qw.device))
+  ret = Tensor.custom_kernel(*outs, x.contiguous(), qw, kw, vw, fxn=_kda_qkv_kernel)
+  return ret[0], ret[1], ret[2]
+
+def dual_bf16_matvec(x:Tensor, aw:Tensor, bw:Tensor, fast:bool=False) -> tuple[Tensor, Tensor]:
+  """Fuse two equal-shaped BF16 decode projections that consume the same input."""
+  from tinygrad.llm.kernels.amd import _dual_bf16_matvec_fast_kernel, _dual_bf16_matvec_kernel
+  batch, tokens, _ = x.shape
+  out_features = aw.shape[0]
+  if aw.shape != bw.shape: raise ValueError("fused BF16 weights must have equal shapes")
+  if isinstance(aw.device, tuple) and aw.uop.axis == 0:
+    devices = aw.device
+    if out_features % len(devices): raise ValueError("fused BF16 output is not divisible by the device count")
+    shard_shape = (batch, tokens, out_features//len(devices))
+    def make_out() -> Tensor:
+      parts = [Tensor.empty(*shard_shape, dtype=dtypes.bfloat16, device=device).uop for device in devices]
+      return Tensor(parts[0].mstack(*parts[1:]).unshard(2))
+    outs = (make_out(), make_out())
+  else:
+    outs = (Tensor.empty(batch, tokens, out_features, dtype=dtypes.bfloat16, device=aw.device),
+            Tensor.empty(batch, tokens, out_features, dtype=dtypes.bfloat16, device=aw.device))
+  ret = Tensor.custom_kernel(*outs, x.contiguous(), aw, bw, fxn=_dual_bf16_matvec_fast_kernel if fast else _dual_bf16_matvec_kernel)
+  return ret[0], ret[1]
+
+def dual_input_bf16_matvec(ax:Tensor, bx:Tensor, aw:Tensor, bw:Tensor) -> tuple[Tensor, Tensor]:
+  """Fuse equal-shaped BF16 projections with separate inputs and identical TP layouts."""
+  from tinygrad.llm.kernels.amd import _dual_input_bf16_matvec_kernel
+  if ax.shape != bx.shape or aw.shape != bw.shape: raise ValueError("fused BF16 inputs and weights must have equal shapes")
+  batch, tokens, _ = ax.shape
+  out_features = aw.shape[0]
+  if isinstance(aw.device, tuple) and aw.uop.axis == 0:
+    devices = aw.device
+    if out_features % len(devices): raise ValueError("fused BF16 output is not divisible by the device count")
+    shard_shape = (batch, tokens, out_features//len(devices))
+    def make_out() -> Tensor:
+      parts = [Tensor.empty(*shard_shape, dtype=dtypes.bfloat16, device=device).uop for device in devices]
+      return Tensor(parts[0].mstack(*parts[1:]).unshard(2))
+    outs = (make_out(), make_out())
+  else:
+    outs = (Tensor.empty(batch, tokens, out_features, dtype=dtypes.bfloat16, device=aw.device),
+            Tensor.empty(batch, tokens, out_features, dtype=dtypes.bfloat16, device=aw.device))
+  ret = Tensor.custom_kernel(*outs, ax.contiguous(), bx.contiguous(), aw, bw, fxn=_dual_input_bf16_matvec_kernel)
+  return ret[0], ret[1]
+
+def kda_fgb_linear(x:Tensor, gw:Tensor, fw:Tensor, bw:Tensor) -> tuple[Tensor, Tensor, Tensor]:
+  """Fuse replicated KDA g/f low-rank projections with its output-sharded beta projection."""
+  from tinygrad.llm.kernels.amd import _kda_fgb_kernel
+  if gw.shape != fw.shape or not isinstance(gw.device, tuple) or gw.uop.axis is not None or \
+     bw.device != gw.device or bw.uop.axis != 0: raise ValueError("unsupported KDA f/g/beta TP layout")
+  batch, tokens, _ = x.shape
+  devices, rank, beta_features = gw.device, gw.shape[0], bw.shape[0]
+  gout = Tensor.empty(batch, tokens, rank, dtype=dtypes.bfloat16, device=devices)
+  fout = Tensor.empty(batch, tokens, rank, dtype=dtypes.bfloat16, device=devices)
+  beta_shape = (batch, tokens, beta_features//len(devices))
+  parts = [Tensor.empty(*beta_shape, dtype=dtypes.bfloat16, device=device).uop for device in devices]
+  bout = Tensor(parts[0].mstack(*parts[1:]).unshard(2))
+  ret = Tensor.custom_kernel(gout, fout, bout, x.contiguous(), gw, fw, bw, fxn=_kda_fgb_kernel)
+  return ret[0], ret[1], ret[2]
+
+@functools.cache
+def _gated_delta_prefill_kernel(core:UOp, next_state:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:UOp, state:UOp, kq:UOp) -> UOp:
+  batch, heads, tokens, value_dim = cast(tuple[int, int, int, int], core.shape)
+  key_dim, alpha_dim = cast(int, q.shape[-1]), cast(int, alpha.shape[-1]) if len(alpha.shape) == 4 else 1
+  core, v = (x.reshape(batch*heads, tokens, value_dim) for x in (core, v))
+  q, k = (x.reshape(batch*heads, tokens, key_dim) for x in (q, k))
+  beta, kq = (x.reshape(batch*heads, tokens) for x in (beta, kq))
+  alpha = alpha.reshape(batch*heads, tokens, alpha_dim)
+  state, next_state = (x.reshape(batch*heads, value_dim, key_dim) for x in (state, next_state))
+  bh, row, cols = UOp.range(batch*heads, 0, AxisType.GLOBAL), UOp.range(value_dim, 2), tuple(range(key_dim))
+  current = UOp.placeholder((key_dim,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  current = current.after(UOp.group(*(current[col].store(state[bh, row, col].float()) for col in cols)))
+  token = UOp.range(tokens, 1, AxisType.REDUCE)
+  previous = tuple(current.after(token)[col].load() for col in cols)
+  keys, queries = (tuple(x[bh, token, col].load() for col in cols) for x in (k, q))
+  av = tuple(alpha[bh, token, col if alpha_dim > 1 else 0].load() for col in cols)
+  bv = beta[bh, token].load()
+  state_k = sum((x*a*y for x,a,y in zip(previous, av, keys)), UOp.const(0, dtypes.float32))
+  state_q = sum((x*a*y for x,a,y in zip(previous, av, queries)), UOp.const(0, dtypes.float32))
+  delta = (v[bh, token, row].load() - state_k) * bv
+  step = UOp.group(core[bh, token, row].store(state_q + delta*kq[bh, token]),
+                   *(current[col].store(x*a + delta*y) for col,x,a,y in zip(cols, previous, av, keys))).end(token)
+  stores = (next_state[bh, row, col].store(current.after(step)[col].load().cast(next_state.dtype)) for col in cols)
+  return UOp.group(*stores).end(row, bh).sink(arg=KernelInfo(name="gated_delta_prefill", opts_to_apply=()))
+
+def gated_delta_prefill(q:Tensor, k:Tensor, v:Tensor, beta:Tensor, alpha:Tensor, state:Tensor) -> tuple[Tensor, Tensor]:
+  batch, heads, tokens, key_dim = q.shape
+  value_dim = v.shape[-1]
+  assert q.shape == k.shape and v.shape[:3] == q.shape[:3] and beta.shape == (batch, heads, tokens)
+  assert alpha.shape in ((batch, heads, tokens), (batch, heads, tokens, key_dim))
+  assert state.shape == (batch, heads, value_dim, key_dim)
+  kernel = _gated_delta_prefill_kernel
+  if amd_exact_bf16_custom_kernels_supported(q.device) and key_dim % 32 == 0 and value_dim % 4 == 0:
+    from tinygrad.llm.kernels.amd import _gated_delta_prefill_kernel as kernel
+  core, next_state, kq = Tensor.empty_like(v), Tensor.empty_like(state), (q*k).sum(-1).contiguous()
+  result = Tensor.custom_kernel(core, next_state, q.contiguous(), k.contiguous(), v.contiguous(), beta.contiguous(), alpha.contiguous(), state, kq,
+                                fxn=kernel)
+  return result[0], result[1]
@@ -0,0 +1,353 @@
+from __future__ import annotations
+import functools
+import pathlib
+from typing import cast
+from tinygrad import UOp
+from tinygrad.uop.ops import AxisType, KernelInfo, Ops
+from tinygrad.dtype import AddrSpace, dtypes
+
+@functools.cache
+def _bf16_mfma_splitk_kernel(out:UOp, x:UOp, weight:UOp) -> UOp:
+  """CDNA4 BF16 matvec with eight waves splitting K per 16 output channels."""
+  from tinygrad.renderer import Estimates
+  from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
+  out_features, in_features = cast(tuple[int, int], weight.shape)
+  assert out.numel() == out_features and x.numel() == in_features and out_features % 16 == 0 and in_features % 256 == 0
+  threads, workgroups = UOp.special(512, "lidx0"), UOp.special(out_features//16, "gidx0")
+  sink = UOp.sink(out.base, x.base, weight.base, threads, workgroups,
+                  arg=KernelInfo(name=f"bf16_mfma_splitk_{out_features}_{in_features}",
+                                 estimates=Estimates(ops=2*out_features*in_features,
+                                                     mem=(out_features*in_features+in_features+out_features)*2)))
+  root = pathlib.Path(__file__).parents[3]/"extra"/"thunder"/"amd"
+  src = (root/"matvec_bf16_splitk.cpp").read_text()
+  lib = HIPCCCompiler("gfx950", [f"-I{(root/'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-ffast-math",
+                                   "-DHIP_ENABLE_WARP_SYNC_BUILTINS", f"-DMATVEC_N={out_features}",
+                                   f"-DMATVEC_K={in_features}"]).compile_cached(src)
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
+
+def warp_reduce(val:UOp, full_wave:bool=False, maximum:bool=False) -> UOp:
+  for offset in ((16, 8, 4, 2, 1) if full_wave else (8, 4, 2, 1)):
+    if val.op is Ops.INDEX and val.addrspace == AddrSpace.REG: val = val.load()
+    other = UOp(Ops.CUSTOM, dtypes.float, (val,), arg=
+      f"__builtin_bit_cast(float, __builtin_amdgcn_ds_swizzle(__builtin_bit_cast(int, {{0}}), {0x1f | offset<<10}))")
+    val = val.maximum(other) if maximum else val + other
+  return val
+
+@functools.cache
+def _mxfp8_qdq_kernel(out:UOp, x:UOp) -> UOp:
+  """Software OCP E4M3/E8M0 round trip, one wave per 32-value MX block."""
+  groups = cast(int, x.shape[-1])//32
+  outer = x.numel()//cast(int, x.shape[-1])
+  block, lane = UOp.range(outer*groups, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  value = x.reshape(outer, groups, 32)[block//groups, block%groups, lane].float()
+  amax = warp_reduce(value.abs(), full_wave=True, maximum=True)
+  exponent = (amax.maximum(1e-38)/448.0).log2().round().maximum(-127.0).minimum(127.0)
+  block_scale = amax.eq(0).where(1.0, exponent.exp2())
+  normalized = value/block_scale
+  magnitude = normalized.abs().minimum(448.0)
+  elem_exp = magnitude.maximum(2**-9).log2().floor().maximum(-6.0).minimum(8.0)
+  quantum = (elem_exp-3.0).exp2()
+  quantized = (magnitude/quantum).round()*quantum
+  quantized = (normalized < 0).where(-quantized, quantized).maximum(-448.0).minimum(448.0)
+  store = out.reshape(outer, groups, 32)[block//groups, block%groups, lane].store((quantized*block_scale).cast(out.dtype))
+  return store.end(lane, block).sink(arg=KernelInfo(name="mxfp8_qdq", opts_to_apply=()))
+
+def _mxfp4_value(code:UOp) -> UOp:
+  """Decode one OCP E2M1 nibble without a lookup-table memory access."""
+  magnitude = code & 7
+  value = magnitude.eq(7).where(6.0, magnitude.eq(6).where(4.0, magnitude.eq(5).where(3.0, magnitude.float()*0.5)))
+  return (code & 8).ne(0).where(-value, value)
+
+def _e8m0_value(scale:UOp) -> UOp:
+  """Decode an E8M0 byte with IEEE exponent bits instead of a transcendental exp2."""
+  bits = scale.cast(dtypes.uint32) << 23
+  # E8M0 byte zero denotes 2**-127, halfway through IEEE's subnormal exponent bin.
+  return scale.eq(0).where(UOp.const(0x00400000, dtypes.uint32).bitcast(dtypes.float32), bits.bitcast(dtypes.float32))
+
+@functools.cache
+def _kda_qkv_kernel(qout:UOp, kout:UOp, vout:UOp, x:UOp, qw:UOp, kw:UOp, vw:UOp) -> UOp:
+  """Fused BF16 decode projection for equal-sized KDA Q/K/V tensors."""
+  batch, tokens, out_features = cast(tuple[int, int, int], qout.shape)
+  in_features, output_tile = cast(int, x.shape[-1]), 1
+  assert qout.shape == kout.shape == vout.shape and out_features % output_tile == 0 and in_features % 32 == 0
+  row, lane = UOp.range(batch*tokens*(out_features//output_tile), 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  token, output_block = row // (out_features//output_tile), row % (out_features//output_tile)
+  outputs = tuple(output_block*output_tile+i for i in range(output_tile))
+  acc = UOp.placeholder((3, output_tile), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  acc = acc.after(acc.store(acc.const_like(0.0)))
+  group = UOp.range(in_features//32, 2, AxisType.REDUCE)
+  activation = x.reshape(batch*tokens, in_features)[token, group*32+lane].float()
+  updates = [acc.after(group)[p, i].load()+activation*w[output, group*32+lane].float()
+             for p,w in enumerate((qw, kw, vw)) for i,output in enumerate(outputs)]
+  update = acc.store(UOp.stack(*updates).reshape(3, output_tile)).end(group)
+  outs = (qout, kout, vout)
+  stores = (outs[p].reshape(batch*tokens, out_features)[token, output.valid(lane.eq(0))].store(
+    warp_reduce(acc.after(update)[p, i], full_wave=True).cast(outs[p].dtype))
+    for p in range(3) for i,output in enumerate(outputs))
+  return UOp.group(*stores).end(lane, row).sink(arg=KernelInfo(name="kda_qkv", opts_to_apply=()))
+
+@functools.cache
+def _dual_bf16_matvec_kernel(aout:UOp, bout:UOp, x:UOp, aw:UOp, bw:UOp) -> UOp:
+  """Exact pair of BF16 decode projections with one activation read."""
+  batch, tokens, out_features = cast(tuple[int, int, int], aout.shape)
+  in_features = cast(int, x.shape[-1])
+  assert aout.shape == bout.shape and out_features == aw.shape[0] == bw.shape[0] and in_features % 16 == 0
+  row, lane = UOp.range(batch*tokens*out_features, 0), UOp.range(16, 1, axis_type=AxisType.LOCAL)
+  token, output = row//out_features, row%out_features
+  acc = UOp.placeholder((2,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  acc = acc.after(acc.store(acc.const_like(0.0)))
+  chunk, group = in_features//16, UOp.range(in_features//16, 2, AxisType.REDUCE)
+  input_idx = lane*chunk+group
+  activation = x.reshape(batch*tokens, in_features)[token, input_idx].float()
+  update = acc.store(UOp.stack(*(acc.after(group)[i].load()+(activation*w[output, input_idx].float()).cast(dtypes.bfloat16).float()
+                                for i,w in enumerate((aw, bw))))).end(group)
+  local = UOp.placeholder((2, 16), dtypes.float32, slot=1, addrspace=AddrSpace.LOCAL)
+  barrier = UOp.group(*(local[i, lane].store(acc.after(update)[i]) for i in range(2))).barrier()
+  stores = (out.reshape(batch*tokens, out_features)[token, output.valid(lane.eq(0))].store(
+    sum((local.after(barrier)[i, j] for j in range(16)), UOp.const(0, dtypes.float32)).cast(out.dtype))
+            for i,out in enumerate((aout, bout)))
+  return UOp.group(*stores).end(lane, row).sink(arg=KernelInfo(name="dual_bf16_matvec", opts_to_apply=()))
+
+@functools.cache
+def _dual_bf16_matvec_fast_kernel(aout:UOp, bout:UOp, x:UOp, aw:UOp, bw:UOp) -> UOp:
+  """Coalesced pair used where one changed reduction boundary does not feed recurrent state."""
+  batch, tokens, out_features = cast(tuple[int, int, int], aout.shape)
+  in_features = cast(int, x.shape[-1])
+  assert aout.shape == bout.shape and out_features == aw.shape[0] == bw.shape[0] and in_features % 32 == 0
+  row, lane = UOp.range(batch*tokens*out_features, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  token, output = row//out_features, row%out_features
+  acc = UOp.placeholder((2,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  acc = acc.after(acc.store(acc.const_like(0.0)))
+  group = UOp.range(in_features//32, 2, AxisType.REDUCE)
+  input_idx = group*32+lane
+  activation = x.reshape(batch*tokens, in_features)[token, input_idx].float()
+  update = acc.store(UOp.stack(*(acc.after(group)[i].load()+(activation*w[output, input_idx].float()).cast(dtypes.bfloat16).float()
+                                for i,w in enumerate((aw, bw))))).end(group)
+  stores = (out.reshape(batch*tokens, out_features)[token, output.valid(lane.eq(0))].store(
+    warp_reduce(acc.after(update)[i], full_wave=True).cast(out.dtype)) for i,out in enumerate((aout, bout)))
+  return UOp.group(*stores).end(lane, row).sink(arg=KernelInfo(name="dual_bf16_matvec_fast", opts_to_apply=()))
+
+@functools.cache
+def _dual_input_bf16_matvec_kernel(aout:UOp, bout:UOp, ax:UOp, bx:UOp, aw:UOp, bw:UOp) -> UOp:
+  """Exact pair of equal-shaped BF16 projections with distinct inputs."""
+  batch, tokens, out_features = cast(tuple[int, int, int], aout.shape)
+  in_features = cast(int, ax.shape[-1])
+  assert aout.shape == bout.shape and ax.shape == bx.shape and out_features == aw.shape[0] == bw.shape[0] and in_features % 16 == 0
+  row, lane = UOp.range(batch*tokens*out_features, 0), UOp.range(16, 1, axis_type=AxisType.LOCAL)
+  token, output = row//out_features, row%out_features
+  acc = UOp.placeholder((2,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  acc = acc.after(acc.store(acc.const_like(0.0)))
+  chunk, group = in_features//16, UOp.range(in_features//16, 2, AxisType.REDUCE)
+  input_idx = lane*chunk+group
+  update = acc.store(UOp.stack(*(acc.after(group)[i].load()+(inp.reshape(batch*tokens, in_features)[token, input_idx].float()*
+    weight[output, input_idx].float()).cast(dtypes.bfloat16).float() for i,(inp,weight) in enumerate(((ax,aw), (bx,bw)))))).end(group)
+  local = UOp.placeholder((2, 16), dtypes.float32, slot=1, addrspace=AddrSpace.LOCAL)
+  barrier = UOp.group(*(local[i, lane].store(acc.after(update)[i]) for i in range(2))).barrier()
+  stores = (out.reshape(batch*tokens, out_features)[token, output.valid(lane.eq(0))].store(
+    sum((local.after(barrier)[i, j] for j in range(16)), UOp.const(0, dtypes.float32)).cast(out.dtype))
+            for i,out in enumerate((aout, bout)))
+  return UOp.group(*stores).end(lane, row).sink(arg=KernelInfo(name="dual_input_bf16_matvec", opts_to_apply=()))
+
+@functools.cache
+def _kda_fgb_kernel(gout:UOp, fout:UOp, bout:UOp, x:UOp, gw:UOp, fw:UOp, bw:UOp) -> UOp:
+  """Mixed-output wave32 KDA g/f/beta projection."""
+  batch, tokens, rank = cast(tuple[int, int, int], gout.shape)
+  beta_features, in_features = cast(int, bout.shape[-1]), cast(int, x.shape[-1])
+  assert gout.shape == fout.shape and rank == gw.shape[0] == fw.shape[0] and beta_features == bw.shape[0] and in_features % 32 == 0
+  rows = batch*tokens*(2*rank+beta_features)
+  row, lane = UOp.range(rows, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  token, projection_row = row//(2*rank+beta_features), row%(2*rank+beta_features)
+  is_g, is_f = projection_row < rank, (projection_row >= rank) & (projection_row < 2*rank)
+  g_row = projection_row.valid(is_g)
+  f_row = (projection_row-rank).valid(is_f)
+  b_row = (projection_row-2*rank).valid(~is_g & ~is_f)
+  acc = UOp.placeholder((), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  acc = acc.after(acc.store(0.0))
+  group = UOp.range(in_features//32, 2, AxisType.REDUCE)
+  input_idx = group*32+lane
+  weight = is_g.where(gw[g_row, input_idx], is_f.where(fw[f_row, input_idx], bw[b_row, input_idx])).float()
+  product = (x.reshape(batch*tokens, in_features)[token, input_idx].float()*weight).cast(dtypes.bfloat16).float()
+  update = acc.store(acc.after(group)+product).end(group)
+  total = warp_reduce(acc.after(update)[0], full_wave=True).cast(dtypes.bfloat16)
+  stores = (gout.reshape(batch*tokens, rank)[token, g_row.valid(lane.eq(0))].store(total),
+            fout.reshape(batch*tokens, rank)[token, f_row.valid(lane.eq(0))].store(total),
+            bout.reshape(batch*tokens, beta_features)[token, b_row.valid(lane.eq(0))].store(total))
+  return UOp.group(*stores).end(lane, row).sink(arg=KernelInfo(name="kda_fgb", opts_to_apply=()))
+
+def _mxfp4_expert_linear_impl(out:UOp, sel:UOp, x:UOp, weight:UOp, scale:UOp) -> UOp:
+  """Wave32 decode GEMM which consumes selected experts directly from packed MXFP4 storage."""
+  batch, tokens, topk, out_features = cast(tuple[int, int, int, int], out.shape[:4])
+  partials = cast(int, out.shape[4]) if len(out.shape) == 5 else 1
+  output_tile = 1
+  assert out_features % output_tile == 0
+  in_features = cast(int, weight.shape[-1])*2
+  assert in_features % 32 == 0 and x.shape[-1] == in_features and sel.shape == (batch, tokens, topk)
+  xchoices = cast(int, x.shape[-2])
+  assert xchoices in (1, topk)
+  total_rows = batch*tokens*topk*(out_features//output_tile)*partials
+  row, lane = UOp.range(total_rows, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  partial, output_block, route = row % partials, (row//partials) % (out_features//output_tile), \
+    row // ((out_features//output_tile)*partials)
+  outputs = tuple(output_block*output_tile+i for i in range(output_tile))
+  token, choice = route // topk, route % topk
+  expert = sel.reshape(batch*tokens, topk)[token, choice]
+  xv = x.reshape(batch*tokens, xchoices, in_features)
+  acc = UOp.placeholder((output_tile,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  acc = acc.after(acc.store(acc.const_like(0.0)))
+  group = UOp.range(in_features//32, 2, AxisType.REDUCE)
+  activation = xv[token, 0 if xchoices == 1 else choice, group*32+lane].float()
+  updates = []
+  for i,output in enumerate(outputs):
+    packed = weight[expert, output, group*16+lane//2]
+    code = (packed >> ((lane&1)*4).cast(dtypes.uint8)) & 15
+    w = _mxfp4_value(code) * _e8m0_value(scale[expert, output, group])
+    updates.append(acc.after(group)[i].load()+activation*w)
+  update = acc.store(UOp.stack(*updates)).end(group)
+  out = out.reshape(batch*tokens, topk, out_features, partials)
+  stores = (out[token, choice, output, partial.valid(lane.eq(0))].store(warp_reduce(acc.after(update)[i], full_wave=True).cast(out.dtype))
+            for i,output in enumerate(outputs))
+  return UOp.group(*stores).end(lane, row).sink(arg=KernelInfo(name="mxfp4_expert_linear", opts_to_apply=()))
+
+@functools.cache
+def _mxfp4_expert_linear_kernel(out:UOp, sel:UOp, x:UOp, weight:UOp, scale:UOp) -> UOp:
+  return _mxfp4_expert_linear_impl(out, sel, x, weight, scale)
+
+@functools.cache
+def _mxfp4_expert_linear_wave64_kernel(out:UOp, sel:UOp, x:UOp, weight:UOp, scale:UOp) -> UOp:
+  """Wave64 decode GEMM for CDNA, with a workgroup-wide reduction independent of local-id decomposition."""
+  batch, tokens, topk, out_features = cast(tuple[int, int, int, int], out.shape[:4])
+  partials = cast(int, out.shape[4]) if len(out.shape) == 5 else 1
+  in_features = cast(int, weight.shape[-1])*2
+  assert in_features % 64 == 0 and x.shape[-1] == in_features and sel.shape == (batch, tokens, topk)
+  xchoices = cast(int, x.shape[-2])
+  assert xchoices in (1, topk)
+  total_rows = batch*tokens*topk*out_features*partials
+  row, lane = UOp.range(total_rows, 0), UOp.range(64, 1, axis_type=AxisType.LOCAL)
+  partial, output, route = row%partials, (row//partials)%out_features, row//(out_features*partials)
+  token, choice = route//topk, route%topk
+  expert = sel.reshape(batch*tokens, topk)[token, choice]
+  xv = x.reshape(batch*tokens, xchoices, in_features)
+  acc = UOp.placeholder((), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  acc = acc.after(acc.store(0.0))
+  group = UOp.range(in_features//64, 2, AxisType.REDUCE)
+  activation = xv[token, 0 if xchoices == 1 else choice, group*64+lane].float()
+  packed = weight[expert, output, group*32+lane//2]
+  code = (packed >> ((lane&1)*4).cast(dtypes.uint8)) & 15
+  weight_value = _mxfp4_value(code) * _e8m0_value(scale[expert, output, group*2+lane//32])
+  update = acc.store(acc.after(group)+activation*weight_value).end(group)
+  local = UOp.placeholder((64,), dtypes.float32, slot=1, addrspace=AddrSpace.LOCAL)
+  barrier = local[lane].store(acc.after(update)[0]).barrier()
+  total = sum((local.after(barrier)[i] for i in range(64)), UOp.const(0, dtypes.float32))
+  out = out.reshape(batch*tokens, topk, out_features, partials)
+  store = out[token, choice, output, partial.valid(lane.eq(0))].store(total.cast(out.dtype))
+  return store.end(lane, row).sink(arg=KernelInfo(name="mxfp4_expert_linear_wave64", opts_to_apply=()))
+
+@functools.cache
+def _mxfp4_expert_linear_wave64_prefill_kernel(out:UOp, sel:UOp, x:UOp, weight:UOp, scale:UOp) -> UOp:
+  """Tiled CDNA4 prefill GEMM. Four adjacent outputs share activation loads and each wave half reduces in parallel."""
+  batch, tokens, topk, out_features = cast(tuple[int, int, int, int], out.shape[:4])
+  partials = cast(int, out.shape[4]) if len(out.shape) == 5 else 1
+  in_features, output_tile = cast(int, weight.shape[-1])*2, 4
+  assert tokens > 1 and in_features % 64 == 0 and out_features % output_tile == 0
+  assert x.shape[-1] == in_features and sel.shape == (batch, tokens, topk)
+  xchoices = cast(int, x.shape[-2])
+  assert xchoices in (1, topk)
+  total_rows = batch*tokens*topk*(out_features//output_tile)*partials
+  row, lane = UOp.range(total_rows, 0), UOp.range(64, 1, axis_type=AxisType.LOCAL)
+  partial, output_block, route = row%partials, (row//partials)%(out_features//output_tile), row//((out_features//output_tile)*partials)
+  outputs = tuple(output_block*output_tile+i for i in range(output_tile))
+  token, choice = route//topk, route%topk
+  expert = sel.reshape(batch*tokens, topk)[token, choice]
+  xv = x.reshape(batch*tokens, xchoices, in_features)
+  acc = UOp.placeholder((output_tile,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  acc = acc.after(acc.store(acc.const_like(0.0)))
+  group = UOp.range(in_features//64, 2, AxisType.REDUCE)
+  activation = xv[token, 0 if xchoices == 1 else choice, group*64+lane].float()
+  updates = []
+  for i,output in enumerate(outputs):
+    packed = weight[expert, output, group*32+lane//2]
+    code = (packed >> ((lane&1)*4).cast(dtypes.uint8)) & 15
+    weight_value = _mxfp4_value(code) * _e8m0_value(scale[expert, output, group*2+lane//32])
+    updates.append(acc.after(group)[i].load()+activation*weight_value)
+  update = acc.store(UOp.stack(*updates)).end(group)
+  half_totals = tuple(warp_reduce(acc.after(update)[i], full_wave=True) for i in range(output_tile))
+  local = UOp.placeholder((output_tile, 2), dtypes.float32, slot=1, addrspace=AddrSpace.LOCAL)
+  half = (lane//32).valid((lane&31).eq(0))
+  barrier = UOp.group(*(local[i, half].store(total) for i,total in enumerate(half_totals))).barrier()
+  out = out.reshape(batch*tokens, topk, out_features, partials)
+  stores = (out[token, choice, output, partial.valid(lane.eq(0))].store(
+    (local.after(barrier)[i, 0]+local.after(barrier)[i, 1]).cast(out.dtype)) for i,output in enumerate(outputs))
+  return UOp.group(*stores).end(lane, row).sink(arg=KernelInfo(name="mxfp4_expert_linear_wave64_prefill", opts_to_apply=()))
+
+@functools.cache
+def _bf16_partial_linear_kernel(out:UOp, x:UOp, weight:UOp) -> UOp:
+  """Per-device BF16 down projection; its dummy final axis is reduced after combining TP partials."""
+  batch, tokens, out_features, partials = cast(tuple[int, int, int, int], out.shape)
+  in_features, output_tile = cast(int, x.shape[-1]), 1
+  assert out_features % output_tile == 0 and in_features % 32 == 0
+  row, lane = UOp.range(batch*tokens*(out_features//output_tile)*partials, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  partial, output_block, token = row%partials, (row//partials)%(out_features//output_tile), row//(partials*(out_features//output_tile))
+  outputs = tuple(output_block*output_tile+i for i in range(output_tile))
+  acc = UOp.placeholder((output_tile,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  acc = acc.after(acc.store(acc.const_like(0.0)))
+  group = UOp.range(in_features//32, 2, AxisType.REDUCE)
+  input_idx = group*32+lane
+  activation = x.reshape(batch*tokens, in_features)[token, input_idx].float()
+  update = acc.store(UOp.stack(*(acc.after(group)[i].load()+(activation*weight[output, input_idx].float()).cast(dtypes.bfloat16).float()
+    for i,output in enumerate(outputs)))).end(group)
+  stores = (out.reshape(batch*tokens, out_features, partials)[token, output, partial.valid(lane.eq(0))].store(
+    warp_reduce(acc.after(update)[i], full_wave=True)) for i,output in enumerate(outputs))
+  return UOp.group(*stores).end(lane, row).sink(arg=KernelInfo(name="bf16_partial_linear", opts_to_apply=()))
+
+@functools.cache
+def _bf16_matvec_kernel(out:UOp, x:UOp, weight:UOp) -> UOp:
+  batch, tokens, out_features = cast(tuple[int, int, int], out.shape)
+  in_features = cast(int, x.shape[-1])
+  assert in_features % 32 == 0
+  row, lane = UOp.range(batch*tokens*out_features, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  token, output = row//out_features, row%out_features
+  acc = UOp.placeholder((), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  acc = acc.after(acc.store(0.0))
+  group = UOp.range(in_features//32, 2, AxisType.REDUCE)
+  input_idx = group*32+lane
+  product = (x.reshape(batch*tokens, in_features)[token, input_idx].float()*weight[output, input_idx].float()).cast(dtypes.bfloat16).float()
+  update = acc.store(acc.after(group)+product).end(group)
+  total = warp_reduce(acc.after(update)[0], full_wave=True)
+  return out.reshape(batch*tokens, out_features)[token, output.valid(lane.eq(0))].store(total.cast(out.dtype)).end(lane, row).sink(
+    arg=KernelInfo(name="bf16_matvec", opts_to_apply=()))
+
+@functools.cache
+def _gated_delta_prefill_kernel(core:UOp, next_state:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:UOp, state:UOp, kq:UOp) -> UOp:
+  batch, heads, tokens, value_dim, row_tile = *core.shape, 4
+  key_dim, alpha_dim = q.shape[-1], alpha.shape[-1] if len(alpha.shape) == 4 else 1
+  assert all(isinstance(x, int) for x in (batch, heads, tokens, value_dim, key_dim)) and key_dim % 32 == 0 and value_dim % row_tile == 0
+  batch, heads, tokens, value_dim, key_dim = cast(tuple[int, int, int, int, int], (batch, heads, tokens, value_dim, key_dim))
+  core, v = (x.reshape(batch*heads, tokens, value_dim) for x in (core, v))
+  q, k = (x.reshape(batch*heads, tokens, key_dim) for x in (q, k))
+  beta, kq = (x.reshape(batch*heads, tokens) for x in (beta, kq))
+  alpha = alpha.reshape(batch*heads, tokens, alpha_dim)
+  state, next_state = (x.reshape(batch*heads, value_dim, key_dim) for x in (state, next_state))
+  bh_row, lane = UOp.range(batch*heads*value_dim//row_tile, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  bh, row_base = bh_row // (value_dim//row_tile), (bh_row % (value_dim//row_tile))*row_tile
+  rows = tuple(row_base+i for i in range(row_tile))
+  cols = tuple(lane + i*32 for i in range(key_dim//32))
+  current = UOp.placeholder((row_tile*key_dim//32,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  current = current.after(current.store(UOp.stack(*(state[bh, row, col].float() for row in rows for col in cols))))
+  token = UOp.range(tokens, 2, AxisType.REDUCE)
+  keys = tuple(k[bh, token, col].load() for col in cols)
+  queries = tuple(q[bh, token, col].load() for col in cols)
+  updates:list[UOp] = []
+  stores:list[UOp] = []
+  for row_idx,row in enumerate(rows):
+    previous = tuple(current.after(token)[row_idx*key_dim//32+i].load() for i in range(key_dim//32))
+    av = tuple(alpha[bh, token, col if alpha_dim > 1 else 0].load() for col in cols)
+    bv = beta[bh, token].load()
+    state_k = warp_reduce(sum((x*a*y for x,a,y in zip(previous, av, keys)), UOp.const(0, dtypes.float32)), full_wave=True)
+    state_q = warp_reduce(sum((x*a*y for x,a,y in zip(previous, av, queries)), UOp.const(0, dtypes.float32)), full_wave=True)
+    delta = (v[bh, token, row].load() - state_k) * bv
+    updates += [x*a + delta*y for x,a,y in zip(previous, av, keys)]
+    stores.append(core[bh, token, row.valid(lane.eq(0))].store(state_q + delta*kq[bh, token]))
+  step = UOp.group(*stores, current.store(UOp.stack(*updates))).end(token)
+  state_stores = (next_state[bh, row, col].store(current.after(step)[row_idx*key_dim//32+i].load().cast(next_state.dtype))
+                  for row_idx,row in enumerate(rows) for i,col in enumerate(cols))
+  return UOp.group(*state_stores).end(lane, bh_row).sink(arg=KernelInfo(name="gated_delta_prefill", opts_to_apply=()))
@@ -0,0 +1,181 @@
+from __future__ import annotations
+import base64, gc, json, pathlib, shutil
+from tinygrad import Tensor, Device, dtypes, nn
+from tinygrad.nn.state import safe_load, safe_save
+from tinygrad.llm.model import Transformer, TransformerConfig, SSMConfig
+from tinygrad.llm.quant import quantize_mxfp4_cpu
+
+KIMI_SSM_LAYERS = tuple(i not in (3, 7, 11, 15, 19, 23, 26) for i in range(27))
+KIMI_TENSOR_COUNT, KIMI_LOGICAL_BYTES = 688, 29_051_930_368
+KIMI_CHECKPOINT_FORMAT = "tinygrad-kimi-mxfp4-v2"
+
+def kimi_config(max_context:int, expert_mxfp4:bool=True) -> TransformerConfig:
+  return TransformerConfig(num_blocks=27, dim=2304, hidden_dim=1024, n_heads=32, n_kv_heads=32, norm_eps=1e-5,
+    vocab_size=163840, head_dim=192, rope_theta=10000.0, rope_dim=64, v_head_dim=128, q_lora_rank=0, kv_lora_rank=512,
+    num_experts=256, num_experts_per_tok=8, norm_topk_prob=True, shared_expert_dim=1024, leading_dense_blocks=1,
+    dense_hidden_dim=9216, routed_scaling_factor=2.446, expert_bias=True, max_context=max_context, expert_mxfp4=expert_mxfp4,
+    shared_expert_gate=False, bf16_activations=True, kda_split_qkv=True,
+    recurrent_prefill_chunked=True, recurrent_prefill_chunk_size=32,
+    ssm=SSMConfig(conv_kernel=4, state_size=128, group_count=32, time_step_rank=32, inner_size=4096, kda=True),
+    ssm_layers=KIMI_SSM_LAYERS)
+
+def _shard_kimi(model:Transformer, devices:tuple[str, ...]) -> None:
+  """Tensor-parallel layout. Every GPU owns a slice of every expert (not a replicated expert set)."""
+  for name, value in nn.state.get_state_dict(model).items():
+    axis = None
+    if name in ("token_embd.weight", "output.weight"): axis = 0
+    elif ".ffn_gate_exps.weight" in name or ".ffn_up_exps.weight" in name: axis = 1
+    elif ".ffn_gate_exps.weight_scale" in name or ".ffn_up_exps.weight_scale" in name: axis = 1
+    elif ".ffn_down_exps.weight" in name or ".ffn_down_exps.weight_scale" in name: axis = 2
+    elif name.endswith((".ffn_gate.weight", ".ffn_up.weight", ".ffn_gate_shexp.weight", ".ffn_up_shexp.weight")): axis = 0
+    elif name.endswith((".ffn_down.weight", ".ffn_down_shexp.weight", ".attn_output.weight", ".ssm_out.weight")): axis = 1
+    elif name.endswith((".attn_q.weight", ".attn_k.weight", ".attn_v.weight", ".attn_qkv.weight",
+                        ".ssm_f_b.weight", ".ssm_g_b.weight", ".ssm_beta.weight")): axis = 0
+    elif name.endswith((".ssm_conv1d.weight", ".ssm_q_conv1d.weight", ".ssm_k_conv1d.weight", ".ssm_v_conv1d.weight")): axis = 0
+    elif name.endswith((".ssm_a", ".ssm_dt.bias")): axis = 0
+    elif name.endswith((".attn_k_b.weight", ".attn_v_b.weight")): axis = 0
+    value.shard_(devices, axis=axis)
+
+def _validate_kimi_state(model:Transformer, state:dict[str, Tensor]) -> None:
+  model_state = nn.state.get_state_dict(model)
+  missing, unexpected = set(model_state)-set(state), set(state)-set(model_state)
+  if missing or unexpected: raise ValueError(f"invalid Kimi tensor names: missing={sorted(missing)}, unexpected={sorted(unexpected)}")
+  for name, value in state.items():
+    if value.shape != model_state[name].shape: raise ValueError(f"invalid shape for {name}: expected {model_state[name].shape}, got {value.shape}")
+    expected_dtype = dtypes.uint8 if name.endswith((".weight_scale", "_exps.weight")) else dtypes.bfloat16
+    if value.dtype != expected_dtype: raise ValueError(f"invalid dtype for {name}: expected {expected_dtype}, got {value.dtype}")
+  if len(state) != KIMI_TENSOR_COUNT or (nbytes := sum(x.nbytes() for x in state.values())) != KIMI_LOGICAL_BYTES:
+    raise ValueError(f"invalid Kimi checkpoint size: {len(state)} tensors, {nbytes} bytes")
+
+def _load_converted_state(model_dir:pathlib.Path, files:list[str]) -> dict[str, Tensor]:
+  state:dict[str, Tensor] = {}
+  for filename in files:
+    part = safe_load(model_dir / filename)
+    if duplicates := set(state) & set(part): raise ValueError(f"duplicate Kimi tensors in {filename}: {sorted(duplicates)}")
+    state.update(part)
+  return state
+
+def load_kimi(model_dir:str|pathlib.Path, max_context:int=4096, devices:int=4) -> Transformer:
+  model_dir = pathlib.Path(model_dir)
+  manifest = json.loads((model_dir / "tinygrad-kimi.json").read_text())
+  if manifest.get("format") != KIMI_CHECKPOINT_FORMAT: raise ValueError("unsupported Kimi checkpoint format")
+  if devices != 4: raise ValueError("Kimi-Linear MXFP4 checkpoint currently requires TP4 (--devices 4)")
+  devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(devices))
+  model = Transformer(kimi_config(max_context, expert_mxfp4=True))
+  state = _load_converted_state(model_dir, manifest["files"])
+  _validate_kimi_state(model, state)
+  _shard_kimi(model, devs)
+  nn.state.load_state_dict(model, state, strict=True, consume=True, realize=True)
+  if state: raise ValueError(f"unexpected Kimi tensors: {sorted(state)}")
+  return model
+
+def _load_hf_state(src:pathlib.Path) -> dict[str, Tensor]:
+  index = json.loads((src / "model.safetensors.index.json").read_text())
+  state:dict[str, Tensor] = {}
+  for filename in sorted(set(index["weight_map"].values())): state.update(safe_load(src / filename))
+  return state
+
+def _layer_key(i:int, suffix:str) -> str: return f"model.layers.{i}.{suffix}"
+
+def _convert_attention(sd:dict[str, Tensor], i:int, is_kda:bool, consume:bool=False) -> dict[str, Tensor]:
+  p, out = f"blk.{i}.", {}
+  def get(suffix:str) -> Tensor:
+    key = _layer_key(i, suffix)
+    return (sd.pop(key) if consume else sd[key]).to("CPU")
+  if is_kda:
+    for src_name, dst_name in (("q_proj", "attn_q"), ("k_proj", "attn_k"), ("v_proj", "attn_v")):
+      out[p+dst_name+".weight"] = get(f"self_attn.{src_name}.weight")
+    for src_name, dst_name in (("q_conv1d", "ssm_q_conv1d"), ("k_conv1d", "ssm_k_conv1d"), ("v_conv1d", "ssm_v_conv1d")):
+      out[p+dst_name+".weight"] = get(f"self_attn.{src_name}.weight").squeeze(1)
+    for src_name, dst_name in (("f_a_proj", "ssm_f_a"), ("f_b_proj", "ssm_f_b"), ("g_a_proj", "ssm_g_a"),
+                               ("g_b_proj", "ssm_g_b"), ("b_proj", "ssm_beta"), ("o_proj", "ssm_out")):
+      out[p+dst_name+".weight"] = get(f"self_attn.{src_name}.weight")
+    out[p+"ssm_norm.weight"] = get("self_attn.o_norm.weight")
+    out[p+"ssm_dt.bias"] = get("self_attn.dt_bias")
+    out[p+"ssm_a"] = (-get("self_attn.A_log").float().exp()).reshape(32, 1)
+  else:
+    out[p+"attn_q.weight"] = get("self_attn.q_proj.weight")
+    out[p+"attn_kv_a_mqa.weight"] = get("self_attn.kv_a_proj_with_mqa.weight")
+    out[p+"attn_kv_a_norm.weight"] = get("self_attn.kv_a_layernorm.weight")
+    kv_b = get("self_attn.kv_b_proj.weight").reshape(32, 256, 512)
+    k_b, v_b = kv_b[:, :128], kv_b[:, 128:]
+    out[p+"attn_k_b.weight"], out[p+"attn_v_b.weight"] = k_b.transpose(1, 2), v_b
+    out[p+"attn_output.weight"] = get("self_attn.o_proj.weight")
+  return out
+
+def convert_kimi(src_dir:str|pathlib.Path, dst_dir:str|pathlib.Path) -> None:
+  """Stream the official BF16 checkpoint into the tinygrad TP4 MXFP4/BF16 representation."""
+  src, dst = pathlib.Path(src_dir), pathlib.Path(dst_dir)
+  dst.mkdir(parents=True, exist_ok=True)
+  config = json.loads((src / "config.json").read_text())
+  expected = {"hidden_size":2304, "num_hidden_layers":27, "num_attention_heads":32, "num_key_value_heads":32,
+              "vocab_size":163840, "intermediate_size":9216, "num_experts":256, "num_experts_per_token":8,
+              "moe_intermediate_size":1024, "num_shared_experts":1, "qk_nope_head_dim":128, "qk_rope_head_dim":64,
+              "v_head_dim":128, "kv_lora_rank":512, "first_k_dense_replace":1, "mla_use_nope":True}
+  if any(config.get(k) != v for k,v in expected.items()): raise ValueError(f"not Kimi-Linear-48B-A3B: expected {expected}")
+  sd, files = _load_hf_state(src), []
+
+  common = {"token_embd.weight":sd.pop("model.embed_tokens.weight").to("CPU"), "output_norm.weight":sd.pop("model.norm.weight").to("CPU"),
+            "output.weight":sd.pop("lm_head.weight").to("CPU")}
+  safe_save(common, str(dst / "model-common.safetensors"))
+  del common
+  gc.collect()
+  files.append("model-common.safetensors")
+  for i in range(27):
+    p = f"blk.{i}."
+    layer = _convert_attention(sd, i, KIMI_SSM_LAYERS[i], consume=True)
+    layer[p+"attn_norm.weight"] = sd.pop(_layer_key(i, "input_layernorm.weight")).to("CPU")
+    layer[p+"ffn_norm.weight"] = sd.pop(_layer_key(i, "post_attention_layernorm.weight")).to("CPU")
+    if i == 0:
+      for src_name, dst_name in (("gate_proj", "ffn_gate"), ("up_proj", "ffn_up"), ("down_proj", "ffn_down")):
+        layer[p+dst_name+".weight"] = sd.pop(_layer_key(i, f"mlp.{src_name}.weight")).to("CPU")
+    else:
+      base = _layer_key(i, "block_sparse_moe")
+      layer[p+"ffn_gate_inp.weight"] = sd.pop(base+".gate.weight").to("CPU")
+      # The official name is e_score_correction_bias; tolerate the early checkpoint spelling.
+      bias_name = next(k for k in (base+".gate.e_score_correction_bias", base+".gate.e_score_correction") if k in sd)
+      layer[p+"exp_probs_b.bias"] = sd.pop(bias_name).to("CPU")
+      for src_name, dst_name in (("gate_proj", "ffn_gate_shexp"), ("up_proj", "ffn_up_shexp"), ("down_proj", "ffn_down_shexp")):
+        layer[p+dst_name+".weight"] = sd.pop(base+f".shared_experts.{src_name}.weight").to("CPU")
+    layer_file = f"model-layer-{i:02d}.safetensors"
+    safe_save({k:v.cast(dtypes.bfloat16).contiguous() for k,v in layer.items()}, str(dst/layer_file))
+    del layer
+    gc.collect()
+    files.append(layer_file)
+
+    if i:
+      base = _layer_key(i, "block_sparse_moe.experts")
+      for wid, dst_name in (("w1", "ffn_gate_exps"), ("w3", "ffn_up_exps"), ("w2", "ffn_down_exps")):
+        packed, scales = [], []
+        for e in range(256):
+          q, s = quantize_mxfp4_cpu(sd.pop(f"{base}.{e}.{wid}.weight").to("CPU"))
+          packed.append(q)
+          scales.append(s)
+        expert_file = f"model-layer-{i:02d}-{wid}-mxfp4.safetensors"
+        safe_save({p+dst_name+".weight":Tensor.stack(*packed), p+dst_name+".weight_scale":Tensor.stack(*scales)}, str(dst/expert_file))
+        del packed, scales, q, s
+        gc.collect()
+        files.append(expert_file)
+
+  if sd: raise ValueError(f"unconverted Kimi source tensors: {sorted(sd)}")
+  for name in ("config.json", "tokenizer_config.json", "special_tokens_map.json", "tiktoken.model", "chat_template.jinja"):
+    if (src/name).exists(): shutil.copy2(src/name, dst/name)
+  converted = _load_converted_state(dst, files)
+  _validate_kimi_state(Transformer(kimi_config(max_context=1)), converted)
+  manifest = {"format":KIMI_CHECKPOINT_FORMAT, "tensor_count":KIMI_TENSOR_COUNT,
+              "logical_bytes":KIMI_LOGICAL_BYTES, "files":files}
+  (dst / "tinygrad-kimi.json").write_text(json.dumps(manifest, indent=2)+"\n")
+
+def load_kimi_tokenizer_data(model_dir:str|pathlib.Path) -> tuple[dict[str, int], dict[str, int], int, int]:
+  """Return byte-encoded normal tokens and specials for SimpleTokenizer without transformers/tiktoken."""
+  model_dir = pathlib.Path(model_dir)
+  normal:dict[str, int] = {}
+  bs = [*range(33, 127), *range(161, 173), *range(174, 256)]
+  byte_encoder = {b:chr(b) for b in bs} | {b:chr(256+i) for i,b in enumerate(b for b in range(256) if b not in bs)}
+  for line in (model_dir / "tiktoken.model").read_bytes().splitlines():
+    token, rank = line.split()
+    normal["".join(byte_encoder[b] for b in base64.b64decode(token))] = int(rank)
+  tc = json.loads((model_dir / "tokenizer_config.json").read_text())
+  specials = {v["content"]:int(k) for k,v in tc.get("added_tokens_decoder", {}).items()}
+  cfg = json.loads((model_dir / "config.json").read_text())
+  return normal, specials, cfg["bos_token_id"], cfg["eos_token_id"]
@@ -0,0 +1,337 @@
+from __future__ import annotations
+import gc, json, math, pathlib
+from dataclasses import replace
+from collections import defaultdict
+from typing import Callable, cast
+from tinygrad import Tensor, Device, dtypes, nn
+from tinygrad.device import Buffer
+from tinygrad.uop.ops import UOp
+from tinygrad.nn.state import safe_dtypes, safe_load_metadata
+from tinygrad.llm.kimi import load_kimi_tokenizer_data
+from tinygrad.llm.model import SSMConfig, Transformer, TransformerConfig
+
+KIMI_K3_TOTAL_SIZE = 1_560_860_324_864
+KIMI_K3_TEXT_SIZE = 1_559_965_606_912
+KIMI_K3_TP8_BYTES_PER_GPU = 196_784_397_312
+KIMI_K3_SHARDS = 96
+KIMI_K3_EXPERTS = 896
+KIMI_K3_LAYERS = 93
+KIMI_K3_FULL_ATTN_LAYERS = (*range(3, KIMI_K3_LAYERS, 4), 92)
+KIMI_K3_SSM_LAYERS = tuple(i not in KIMI_K3_FULL_ATTN_LAYERS for i in range(KIMI_K3_LAYERS))
+
+def kimi_k3_config(max_context:int) -> TransformerConfig:
+  """Official Kimi K3 text-tower configuration (zero-based full-attention layers)."""
+  return TransformerConfig(num_blocks=93, dim=7168, hidden_dim=3072, n_heads=96, n_kv_heads=96, norm_eps=1e-5,
+    vocab_size=163840, head_dim=192, rope_theta=10000.0, rope_dim=64, v_head_dim=128, max_context=max_context,
+    q_lora_rank=1536, kv_lora_rank=512, num_experts=896, num_experts_per_tok=16, norm_topk_prob=True,
+    shared_expert_dim=6144, leading_dense_blocks=1, dense_hidden_dim=33792, routed_scaling_factor=1.0,
+    expert_bias=True, expert_mxfp4=True, bf16_activations=True, kda_split_qkv=True,
+    ssm=SSMConfig(conv_kernel=4, state_size=128, group_count=96, time_step_rank=96, inner_size=12288, kda=True, channel_decay=True),
+    ssm_layers=KIMI_K3_SSM_LAYERS, shared_expert_gate=False, attn_output_gate=True,
+    activation_situ_beta=4.0, activation_situ_linear_beta=25.0, routed_expert_dim=3584, latent_moe_norm=True,
+    route_weights_uncorrected=True, attn_res_block_size=12, kda_full_rank_gate=True, kda_gate_lower_bound=-5.0,
+    recurrent_prefill_chunked=True, recurrent_prefill_chunk_size=128)
+
+def kimi_k3_smoke_config(max_context:int=4) -> TransformerConfig:
+  """Reduced K3 with every architectural feature retained for cheap compile/hardware admission tests."""
+  # Keep both the routed latent and the TP8-local expert hidden dimension wave64 aligned. The real
+  # gfx950 packed-expert kernel requires this, so the hardware smoke test must preserve the constraint.
+  return replace(kimi_k3_config(max_context), num_blocks=2, dim=32, hidden_dim=512, n_heads=8, n_kv_heads=8,
+    vocab_size=64, head_dim=8, rope_dim=4, v_head_dim=4, q_lora_rank=16, kv_lora_rank=8, num_experts=512,
+    num_experts_per_tok=2, shared_expert_dim=32, dense_hidden_dim=64, routed_expert_dim=64,
+    ssm=SSMConfig(4, 4, 8, 8, 32, True, True), ssm_layers=(True, False), attn_res_block_size=1)
+
+def _shard_kimi_k3(model:Transformer, devices:tuple[str, ...]) -> None:
+  """Tensor parallel layout for K3. The official dimensions are divisible by TP8."""
+  if len(devices) not in (1, 2, 4, 8): raise ValueError(f"Kimi K3 tensor parallelism requires 1, 2, 4, or 8 devices, got {len(devices)}")
+  for name, value in nn.state.get_state_dict(model).items():
+    axis = None
+    if name in ("token_embd.weight", "output.weight"): axis = 0
+    elif ".ffn_gate_exps.weight" in name or ".ffn_up_exps.weight" in name: axis = 1
+    elif ".ffn_gate_exps.weight_scale" in name or ".ffn_up_exps.weight_scale" in name: axis = 1
+    elif ".ffn_down_exps.weight" in name or ".ffn_down_exps.weight_scale" in name: axis = 2
+    elif name.endswith((".ffn_gate.weight", ".ffn_up.weight", ".ffn_gate_shexp.weight", ".ffn_up_shexp.weight")): axis = 0
+    elif name.endswith((".ffn_down.weight", ".ffn_down_shexp.weight", ".ffn_routed_down.weight", ".ffn_routed_up.weight",
+                        ".attn_output.weight", ".ssm_out.weight")): axis = 1
+    elif name.endswith((".attn_q_b.weight", ".attn_k_b.weight", ".attn_v_b.weight", ".attn_gate.weight",
+                        ".attn_q.weight", ".attn_k.weight", ".attn_v.weight", ".ssm_f_b.weight", ".ssm_g_full.weight", ".ssm_beta.weight")): axis = 0
+    elif name.endswith((".ssm_q_conv1d.weight", ".ssm_k_conv1d.weight", ".ssm_v_conv1d.weight", ".ssm_dt.bias")): axis = 0
+    value.shard_(devices, axis=axis)
+
+def _validate_config(config:dict) -> None:
+  text = config.get("text_config", config)
+  expected = {"model_type":"kimi_linear", "hidden_size":7168, "num_hidden_layers":93, "num_attention_heads":96,
+              "vocab_size":163840, "intermediate_size":33792, "num_experts":896, "num_experts_per_token":16,
+              "moe_intermediate_size":3072, "num_shared_experts":2, "q_lora_rank":1536, "kv_lora_rank":512,
+              "qk_nope_head_dim":128, "qk_rope_head_dim":64, "v_head_dim":128, "routed_expert_hidden_size":3584,
+              "attn_res_block_size":12, "hidden_act":"situ", "mla_use_nope":True, "mla_use_output_gate":True,
+              "activation_situ_beta":4.0, "activation_situ_linear_beta":25.0, "latent_moe_use_norm":True,
+              "moe_renormalize":True, "first_k_dense_replace":1, "num_expert_group":1, "topk_group":1}
+  bad = {k:(text.get(k), v) for k,v in expected.items() if text.get(k) != v}
+  linear = text.get("linear_attn_config", {})
+  linear_expected = {"head_dim":128, "num_heads":96, "short_conv_kernel_size":4, "use_full_rank_gate":True,
+                     "gate_lower_bound":-5.0, "full_attn_layers":[i+1 for i in KIMI_K3_FULL_ATTN_LAYERS],
+                     "kda_layers":[i+1 for i,x in enumerate(KIMI_K3_SSM_LAYERS) if x]}
+  bad.update({f"linear_attn_config.{k}":(linear.get(k), v) for k,v in linear_expected.items() if linear.get(k) != v})
+  quant = text.get("quantization_config", {})
+  if quant.get("format") != "mxfp4-pack-quantized": bad["quantization_config.format"] = (quant.get("format"), "mxfp4-pack-quantized")
+  if bad: raise ValueError(f"not the supported official Kimi K3 checkpoint: {bad}")
+
+def audit_kimi_k3_checkpoint(model_dir:str|pathlib.Path, require_shards:bool=True) -> dict[str, int]:
+  """Validate checkpoint metadata only. This never opens weight data and is safe on small hosts."""
+  root = pathlib.Path(model_dir)
+  _validate_config(json.loads((root / "config.json").read_text()))
+  index = json.loads((root / "model.safetensors.index.json").read_text())
+  weight_map, total = index.get("weight_map", {}), index.get("metadata", {}).get("total_size")
+  language = [k for k in weight_map if k.startswith("language_model.")]
+  experts = [k for k in language if ".block_sparse_moe.experts." in k]
+  missing_files = {fn for fn in weight_map.values() if not (root / fn).is_file()}
+  if total != KIMI_K3_TOTAL_SIZE: raise ValueError(f"unexpected checkpoint size {total}, expected {KIMI_K3_TOTAL_SIZE}")
+  if len(set(weight_map.values())) != KIMI_K3_SHARDS: raise ValueError("official Kimi K3 must contain 96 safetensor shards")
+  if len(experts) != 92 * KIMI_K3_EXPERTS * 3 * 2: raise ValueError(f"unexpected routed-expert tensor count {len(experts)}")
+  if require_shards and missing_files: raise FileNotFoundError(f"missing {len(missing_files)} checkpoint shards, first: {sorted(missing_files)[0]}")
+  return {"tensors":len(weight_map), "language_tensors":len(language), "expert_tensors":len(experts),
+          "shards":len(set(weight_map.values())), "missing_shards":len(missing_files), "total_size":total}
+
+def _layer_sources(i:int, is_kda:bool) -> dict[str, str]:
+  src, dst = f"language_model.model.layers.{i}.", f"blk.{i}."
+  out = {
+    src+"input_layernorm.weight":dst+"attn_norm.weight", src+"post_attention_layernorm.weight":dst+"ffn_norm.weight",
+    src+"self_attention_res_norm.weight":dst+"attn_res_norm.weight", src+"self_attention_res_proj.weight":dst+"attn_res_proj.weight",
+    src+"mlp_res_norm.weight":dst+"mlp_res_norm.weight", src+"mlp_res_proj.weight":dst+"mlp_res_proj.weight",
+  }
+  if is_kda:
+    for a,b in (("q_proj","attn_q"),("k_proj","attn_k"),("v_proj","attn_v"),("g_proj","ssm_g_full"),
+                ("f_a_proj","ssm_f_a"),("f_b_proj","ssm_f_b"),("b_proj","ssm_beta"),("o_proj","ssm_out")):
+      out[src+f"self_attn.{a}.weight"] = dst+b+".weight"
+    for a,b in (("q_conv1d","ssm_q_conv1d"),("k_conv1d","ssm_k_conv1d"),("v_conv1d","ssm_v_conv1d")):
+      out[src+f"self_attn.{a}.weight"] = dst+b+".weight"
+    out[src+"self_attn.o_norm.weight"], out[src+"self_attn.dt_bias"], out[src+"self_attn.A_log"] = \
+      dst+"ssm_norm.weight", dst+"ssm_dt.bias", dst+"ssm_a"
+  else:
+    for a,b in (("q_a_proj","attn_q_a"),("q_a_layernorm","attn_q_a_norm"),("q_b_proj","attn_q_b"),
+                ("kv_a_proj_with_mqa","attn_kv_a_mqa"),("kv_a_layernorm","attn_kv_a_norm"),
+                ("g_proj","attn_gate"),("o_proj","attn_output")):
+      out[src+f"self_attn.{a}.weight"] = dst+b+".weight"
+    # kv_b_proj is split into head-wise K and V tensors while loading.
+    out[src+"self_attn.kv_b_proj.weight"] = dst+"attn_k_b.weight|"+dst+"attn_v_b.weight"
+  if i == 0:
+    for a,b in (("gate_proj","ffn_gate"),("up_proj","ffn_up"),("down_proj","ffn_down")): out[src+f"mlp.{a}.weight"] = dst+b+".weight"
+  else:
+    base = src+"block_sparse_moe."
+    out[base+"gate.weight"], out[base+"gate.e_score_correction_bias"] = dst+"ffn_gate_inp.weight", dst+"exp_probs_b.bias"
+    for a,b in (("gate_proj","ffn_gate_shexp"),("up_proj","ffn_up_shexp"),("down_proj","ffn_down_shexp"),
+                ("routed_expert_down_proj","ffn_routed_down"),("routed_expert_up_proj","ffn_routed_up"),
+                ("routed_expert_norm","ffn_routed_norm")):
+      out[base+(f"shared_experts.{a}.weight" if a.endswith("_proj") and not a.startswith("routed_") else a+".weight")] = dst+b+".weight"
+  return out
+
+def _replace(dst:Tensor, src:Tensor) -> None:
+  if dst.shape != src.shape: raise ValueError(f"shape mismatch: expected {dst.shape}, got {src.shape}")
+  if not isinstance(dst.device, tuple):
+    dst.replace(src.to(dst.device)).realize()
+    return
+  if isinstance(src.device, tuple):
+    dst.replace(src.shard_like(dst)).realize()
+    return
+
+  # Build the final MultiBuffer directly. The generic shard().realize() path schedules several
+  # kernels per tensor and recompiles them for every DISK:<filename> device. Axis-0 shards and
+  # replicas are contiguous, so copy those bytes straight into their final device buffers.
+  devices, axis, shape = dst.device, dst.uop.axis, tuple(int(x) for x in dst.shape)
+  try: src_buffer = cast(Buffer, src.uop.buffer)
+  except (AssertionError, RuntimeError):
+    src = src.clone().realize()
+    src_buffer = cast(Buffer, src.uop.buffer)
+
+  if axis is None:
+    # Replicas are identical on every device. Read the disk tensor once, retain that allocation on
+    # GPU 0, and fan it out over XGMI instead of issuing eight identical direct reads.
+    staging = Tensor.empty(*shape, dtype=src.dtype, device=devices[0]).realize()
+    cast(Buffer, staging.uop.buffer).ensure_allocated().copy_from(src_buffer.ensure_allocated())
+    parts = [staging]
+    for device in devices[1:]:
+      part = Tensor.empty(*shape, dtype=src.dtype, device=device).realize()
+      cast(Buffer, part.uop.buffer).ensure_allocated().copy_from(cast(Buffer, staging.uop.buffer).ensure_allocated())
+      parts.append(part)
+    dst.replace(Tensor(parts[0].uop.mstack(*(x.uop for x in parts[1:]))))
+    return
+  if axis == 0:
+    part_shape = (shape[0]//len(devices), *shape[1:])
+    part_numel = math.prod(part_shape)
+    parts:list[Tensor] = []
+    for i,device in enumerate(devices):
+      part = Tensor.empty(*part_shape, dtype=src.dtype, device=device).realize()
+      source = src_buffer.view(part_numel, src.dtype, i*part_numel*src.dtype.itemsize)
+      cast(Buffer, part.uop.buffer).ensure_allocated().copy_from(source.ensure_allocated())
+      parts.append(part)
+  else:
+    # Inner-axis TP slices are strided in row-major safetensors. Stage one complete tensor on
+    # GPU 0, then schedule all slice kernels and peer copies as one multi-device graph.
+    staging = Tensor.empty(*shape, dtype=src.dtype, device=devices[0]).realize()
+    cast(Buffer, staging.uop.buffer).ensure_allocated().copy_from(src_buffer.ensure_allocated())
+    dst.replace(staging.shard(devices, axis=axis)).realize()
+    return
+  dst.replace(Tensor(parts[0].uop.mstack(*(x.uop for x in parts[1:])).unshard(axis)))
+
+def _safe_load_selected(fn:pathlib.Path, keys:tuple[str, ...]|list[str]) -> dict[str, Tensor]:
+  """Create disk-backed tensors only for selected safetensor entries, without touching payload data."""
+  source, data_start, metadata = safe_load_metadata(fn)
+  data = source[data_start:]
+  missing = [key for key in keys if key not in metadata]
+  if missing: raise ValueError(f"missing tensor {missing[0]} from {fn.name}")
+  out:dict[str, Tensor] = {}
+  for key in keys:
+    entry = metadata[key]
+    out[key] = data[entry["data_offsets"][0]:entry["data_offsets"][1]].bitcast(safe_dtypes[entry["dtype"]]).reshape(entry["shape"])
+  return out
+
+def _load_stacked_experts(dst:Tensor, sources:list[Tensor]) -> None:
+  """Read expert tensors once into a transient GPU staging buffer, then redistribute TP slices over the GPU fabric."""
+  if not sources or not isinstance(dst.device, tuple) or dst.uop.axis is None: raise ValueError("expected a TP-sharded expert destination")
+  devices, axis = dst.device, dst.uop.axis
+  shape = (len(sources), *sources[0].shape)
+  if dst.shape != shape or any(x.shape != sources[0].shape or x.dtype != dst.dtype for x in sources):
+    raise ValueError(f"expert source shape/dtype does not match destination {dst.shape} {dst.dtype}")
+
+  # Each official expert tensor is contiguous in its safetensor file. Assemble it once on GPU 0;
+  # Buffer.copy_from uses the AMD driver's bounded DISK->GPU staging path and never allocates host-sized storage.
+  staging = Tensor.empty(*shape, dtype=dst.dtype, device=devices[0]).realize()
+  staging_buffer = cast(Buffer, staging.uop.buffer)
+  def free_staging_cache() -> None:
+    if (free_cache:=getattr(Device[devices[0]].allocator, "free_cache", None)) is not None: free_cache()
+  offset = 0
+  for source in sources:
+    source_buffer = cast(Buffer, source.uop.buffer)
+    staging_buffer.view(cast(int, source.numel()), source.dtype, offset).ensure_allocated().copy_from(source_buffer.ensure_allocated())
+    offset += source.nbytes()
+
+  # Schedule all TP slices and peer copies together. This avoids eight independent realization
+  # passes and lets the runtime overlap the multi-device transfer graph.
+  dst.replace(staging.shard(devices, axis=axis)).realize()
+  del staging
+  gc.collect()
+  free_staging_cache()
+
+def _load_nonexperts(root:pathlib.Path, weight_map:dict[str, str], model:Transformer, progress:Callable[[str], None]) -> set[str]:
+  model_state, mappings = nn.state.get_state_dict(model), {
+    "language_model.model.embed_tokens.weight":"token_embd.weight", "language_model.model.norm.weight":"output_norm.weight",
+    "language_model.lm_head.weight":"output.weight", "language_model.model.output_attn_res_norm.weight":"output_attn_res_norm.weight",
+    "language_model.model.output_attn_res_proj.weight":"output_attn_res_proj.weight"}
+  for i,is_kda in enumerate(KIMI_K3_SSM_LAYERS): mappings.update(_layer_sources(i, is_kda))
+  by_file:dict[str, list[str]] = defaultdict(list)
+  for source in mappings:
+    if source not in weight_map: raise ValueError(f"missing Kimi K3 tensor {source}")
+    by_file[weight_map[source]].append(source)
+  consumed:set[str] = set()
+  for filename, sources in sorted(by_file.items()):
+    progress(f"loading non-expert tensors from {filename}")
+    shard = _safe_load_selected(root / filename, sources)
+    for source in sources:
+      value, targets = shard[source], mappings[source].split("|")
+      # A_log is the only checkpoint tensor requiring arithmetic during load. Realize its 128
+      # channel values on CPU so replicating it does not try to render the disk/PYTHON graph.
+      if source.endswith("A_log"): value = (-value.to("CPU").float().exp()).reshape(model_state[targets[0]].shape).realize()
+      if source.endswith("conv1d.weight"): value = value.squeeze(1)
+      if source.endswith("kv_b_proj.weight"):
+        # Splitting K/V includes a transpose, which cannot be rendered against a disk buffer.
+        # Materialize only this one 25 MiB projection on CPU, then release it with the shard.
+        value = value.to("CPU").realize().reshape(96, 256, 512)
+        values:tuple[Tensor, ...] = (value[:, :128].transpose(1, 2), value[:, 128:])
+      else: values = (value,)
+      for target,tensor in zip(targets, values): _replace(model_state[target], tensor)
+      consumed.add(source)
+    del shard
+  gc.collect()
+  return consumed
+
+def _load_experts(root:pathlib.Path, weight_map:dict[str, str], model:Transformer, progress:Callable[[str], None]) -> set[str]:
+  model_state, consumed = nn.state.get_state_dict(model), set[str]()
+  for i in range(1, KIMI_K3_LAYERS):
+    base = f"language_model.model.layers.{i}.block_sparse_moe.experts"
+    fields = tuple((wid, suffix, dst_name) for wid,dst_name in (("w1","ffn_gate_exps"),("w2","ffn_down_exps"),("w3","ffn_up_exps"))
+                   for suffix in ("weight_packed", "weight_scale"))
+    keys = {(e,wid,suffix):f"{base}.{e}.{wid}.{suffix}" for e in range(KIMI_K3_EXPERTS) for wid,suffix,_ in fields}
+    files = sorted({weight_map[k] for k in keys.values()})
+    progress(f"loading layer {i}/92 routed experts from {', '.join(files)}")
+    shards = {fn:_safe_load_selected(root / fn, [key for key in keys.values() if weight_map[key] == fn]) for fn in files}
+
+    # Official K3 stores all six tensors for an expert contiguously and all 896 experts for a
+    # layer in one contiguous shard region, ordered lexicographically by expert name. Read that
+    # region once, then reorder/split on GPU 0 directly into the six TP8 destinations.
+    blocks:list[tuple[int, int, int, list[Buffer]]] = []
+    for e in range(KIMI_K3_EXPERTS):
+      bufs = [cast(Buffer, shards[weight_map[keys[e,wid,suffix]]][keys[e,wid,suffix]].uop.buffer) for wid,suffix,_ in fields]
+      if not all(bufs[j].device == bufs[0].device and bufs[j].offset+bufs[j].nbytes == bufs[j+1].offset for j in range(len(bufs)-1)):
+        raise ValueError(f"layer {i} expert {e} tensors are not contiguous in the official shard")
+      blocks.append((bufs[0].offset, bufs[-1].offset+bufs[-1].nbytes, e, bufs))
+    blocks.sort()
+    if len(files) != 1 or not all(blocks[j][1] == blocks[j+1][0] for j in range(len(blocks)-1)):
+      raise ValueError(f"layer {i} routed experts are not one contiguous official-shard region")
+    row_bytes = blocks[0][1]-blocks[0][0]
+    if any(end-start != row_bytes for start,end,_,_ in blocks): raise ValueError(f"layer {i} expert records have inconsistent sizes")
+
+    devices = cast(tuple[str, ...], model_state[f"blk.{i}.ffn_gate_exps.weight"].device)
+    raw = Tensor.empty(KIMI_K3_EXPERTS, row_bytes, dtype=dtypes.uint8, device=devices[0]).realize()
+    raw_buffer, first_buffer = cast(Buffer, raw.uop.buffer), blocks[0][3][0]
+    raw_buffer.ensure_allocated().copy_from(first_buffer.base.view(KIMI_K3_EXPERTS*row_bytes, dtypes.uint8, blocks[0][0]).ensure_allocated())
+    lexpos = {expert:pos for pos,(_,_,expert,_) in enumerate(blocks)}
+    permutation = Tensor([lexpos[e] for e in range(KIMI_K3_EXPERTS)], device=devices[0])
+    field_offset, outputs = 0, []
+    for field_idx,(wid,suffix,dst_name) in enumerate(fields):
+      field_bytes = blocks[0][3][field_idx].nbytes
+      dst = model_state[f"blk.{i}.{dst_name}.weight" + ("_scale" if suffix == "weight_scale" else "")]
+      axis = dst.uop.axis
+      if axis is None: raise ValueError(f"layer {i} expert destination {dst_name} is not TP-sharded")
+      value = raw[:, field_offset:field_offset+field_bytes][permutation].reshape(dst.shape).shard(devices, axis=axis)
+      # Realize into a buffer-identity tensor, then retain only that identity in the model. Keeping
+      # value's arithmetic UOp would also keep the 15.7 GB raw staging tensor and its reorder graph
+      # alive for every loaded weight, wasting about 44 GB on GPU 0 after the load completes.
+      shard_shape = tuple(int(x) for x in value.uop.shard_shape)
+      storage = UOp.new_buffer(devices, math.prod(shard_shape), dst.dtype).reshape(shard_shape).unshard(axis)
+      final = Tensor(storage)
+      final.assign(value)
+      outputs.append((dst, final, storage))
+      field_offset += field_bytes
+    if field_offset != row_bytes: raise ValueError(f"layer {i} expert field sizes do not cover the contiguous record")
+    outputs[0][1].realize(*(value for _,value,_ in outputs[1:]))
+    for dst,_,storage in outputs: dst.replace(Tensor(storage))
+    consumed.update(keys.values())
+    # Drop the realized assignment graphs before flushing the allocator cache. Their final storage
+    # UOps remain in model_state, while the graphs themselves still reference raw and permutation.
+    del shards, raw, raw_buffer, permutation, outputs, value, final, storage, dst
+    if (free_cache:=getattr(Device[devices[0]].allocator, "free_cache", None)) is not None: free_cache()
+  return consumed
+
+def load_kimi_k3(model_dir:str|pathlib.Path, max_context:int=4096, devices:int=8,
+                 progress:Callable[[str], None]=print) -> Transformer:
+  """Load the official native K3 checkpoint without ever materializing it in host RAM.
+
+  Safetensor shards remain disk-backed. Expert tensors are read once into a bounded GPU staging
+  buffer, redistributed as TP slices, and discarded after every projection. Vision tensors are intentionally ignored.
+  """
+  root = pathlib.Path(model_dir)
+  _validate_config(json.loads((root / "config.json").read_text()))
+  index = json.loads((root / "model.safetensors.index.json").read_text())
+  weight_map = index["weight_map"]
+  if devices != 8: raise ValueError("official Kimi K3 currently requires --devices 8")
+  if index.get("metadata", {}).get("total_size") != KIMI_K3_TOTAL_SIZE or len(set(weight_map.values())) != KIMI_K3_SHARDS:
+    raise ValueError("checkpoint index does not match the official 96-shard Kimi K3 release")
+  missing_files = {fn for fn in weight_map.values() if not (root / fn).is_file()}
+  if missing_files: raise FileNotFoundError(f"missing {len(missing_files)} checkpoint shards, first: {sorted(missing_files)[0]}")
+  model = Transformer(kimi_k3_config(max_context))
+  _shard_kimi_k3(model, tuple(f"{Device.DEFAULT}:{i}" for i in range(devices)))
+  consumed = _load_nonexperts(root, weight_map, model, progress)
+  # Expert staging graphs are acyclic and released by reference counting after each layer. Avoid
+  # unnecessary cyclic-collector scans across the complete persistent model graph during this loop.
+  gc_was_enabled = gc.isenabled()
+  gc.disable()
+  try: consumed.update(_load_experts(root, weight_map, model, progress))
+  finally:
+    if gc_was_enabled: gc.enable()
+  unused_language = {k for k in weight_map if k.startswith("language_model.")} - consumed
+  if unused_language: raise ValueError(f"unmapped language tensors: {sorted(unused_language)[:20]}")
+  return model
+
+__all__ = ["KIMI_K3_FULL_ATTN_LAYERS", "KIMI_K3_SSM_LAYERS", "KIMI_K3_TEXT_SIZE", "KIMI_K3_TP8_BYTES_PER_GPU",
+           "audit_kimi_k3_checkpoint", "kimi_k3_config", "kimi_k3_smoke_config", "load_kimi_k3", "load_kimi_tokenizer_data"]
@@ -1,9 +1,16 @@
 from __future__ import annotations
-import functools, itertools, pathlib
+import array, functools, itertools, pathlib
 from dataclasses import dataclass, replace
-from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function
+from typing import Callable, cast
+from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, dtypes
+from tinygrad.device import MultiBuffer
 from tinygrad.nn import Linear
 from tinygrad.llm.gguf import gguf_load
+from tinygrad.llm.quant import dequantize_mxfp4, quantize_dequantize_mxfp8
+from tinygrad.llm.kernels import amd_custom_kernels_supported, amd_exact_bf16_custom_kernels_supported, amd_int32_item, \
+  amd_packed_mxfp4_supported, amd_wave64_custom_kernels_supported, bf16_matvec, bf16_mfma_splitk, bf16_partial_linear, \
+  dual_bf16_matvec, dual_input_bf16_matvec, \
+  gated_delta_prefill, kda_fgb_linear, kda_qkv_linear, mxfp4_expert_linear, mxfp8_quantize_dequantize
 from tinygrad.uop.ops import resolve

@functools.cache
@@ -20,12 +27,37 @@ class ExpertWeights:
    # sel: (B, T, k), x: (B, T, 1, in) or (B, T, k, in) -> output: (B, T, k, out)
    return (x.unsqueeze(-2) @ self.weight[sel].transpose(-1, -2)).contiguous().squeeze(-2)

+class MXFP4ExpertWeights:
+  """Routed-expert weights stored as packed OCP MXFP4 with one E8M0 scale per 32 values."""
+  def __init__(self, num_experts:int, in_features:int, out_features:int):
+    if in_features % 32: raise ValueError(f"MXFP4 expert input size must be divisible by 32, got {in_features}")
+    self.in_features, self.out_features = in_features, out_features
+    self.weight = Tensor.zeros(num_experts, out_features, in_features//2, dtype=dtypes.uint8)
+    self.weight_scale = Tensor.full((num_experts, out_features, in_features//32), 127, dtype=dtypes.uint8)
+  def __call__(self, sel:Tensor, x:Tensor, quantized:bool=False, partial:bool=False) -> Tensor:
+    # Only selected weights are expanded, so packed storage remains resident during generation.
+    if isinstance(self.weight.device, tuple) and not isinstance(sel.device, tuple): sel = sel.shard(self.weight.device, axis=None)
+    if not quantized:
+      x = mxfp8_quantize_dequantize(x.cast(dtypes.bfloat16)) if amd_custom_kernels_supported(x.device) else \
+        quantize_dequantize_mxfp8(x.cast(dtypes.bfloat16))
+    # gfx11 has no native FP4 instructions, but decoding nibbles inside the dot product still avoids
+    # the much larger selected-expert BF16 temporary. Gate/up weights are output-sharded in TP.
+    if amd_packed_mxfp4_supported(self.weight.device):
+      return mxfp4_expert_linear(sel, x, self.weight, self.weight_scale, partial=partial)
+    weight = dequantize_mxfp4(self.weight[sel], self.weight_scale[sel], dtype=dtypes.bfloat16)
+    return (x.unsqueeze(-2) @ weight.transpose(-1, -2)).contiguous().squeeze(-2)
+
 def apply_rope(x:Tensor, freqs_cis:Tensor) -> Tensor:
  assert x.shape[-1] % 2 == 0
  cos, sin = freqs_cis.reshape(1, 1, x.shape[2], -1).chunk(2, dim=-1)
  x1, x2 = x.chunk(2, dim=-1)
  return (x1 * cos - x2 * sin).cat(x2 * cos + x1 * sin, dim=-1)

+def l2norm(x:Tensor, eps:float=1e-6) -> Tensor:
+  """FLA-compatible L2 normalization: FP32 reduction and epsilon inside the square root."""
+  dtype, x = x.dtype, x.float()
+  return (x * (x.square().sum(axis=-1, keepdim=True, dtype=dtypes.float32) + eps).rsqrt()).cast(dtype)
+
 def pairwise_topk(x: Tensor, k: int) -> tuple[Tensor, Tensor]:
  n = x.shape[-1]
  vals = Tensor.arange(n).reshape(1,1,n).cast(x.dtype).expand(x.shape)
@@ -34,6 +66,16 @@ def pairwise_topk(x: Tensor, k: int) -> tuple[Tensor, Tensor]:
  sel = x.const_like(0).scatter(-1, cmp.sum(axis=-1).cast('int32'), vals)[:,:,n-k:].cast('int32')
  return x.gather(-1, sel), sel

+def iterative_topk(x:Tensor, k:int) -> tuple[Tensor, Tensor]:
+  """O(k*N) top-k for very wide MoE routers, with stable first-index tie breaking."""
+  work, values, indices = x, [], []
+  for _ in range(k):
+    sel = work.argmax(-1, keepdim=True)
+    values.append(x.gather(-1, sel))
+    indices.append(sel)
+    work = work.scatter(-1, sel, x.dtype.min)
+  return values[0].cat(*values[1:], dim=-1), indices[0].cat(*indices[1:], dim=-1)
+
@dataclass(frozen=True)
 class SSMConfig:
  conv_kernel: int
@@ -42,6 +84,7 @@ class SSMConfig:
  time_step_rank: int
  inner_size: int
  kda: bool = False
+  channel_decay: bool = False

@dataclass(frozen=True)
 class TransformerConfig:
@@ -73,6 +116,20 @@ class TransformerConfig:
  routed_scaling_factor: float = 1.0
  qkv_bias: bool = False
  expert_bias: bool = False
+  expert_mxfp4: bool = False
+  bf16_activations: bool = False
+  kda_split_qkv: bool = False
+  # Kimi K3 extensions. Defaults preserve all existing model behavior.
+  activation_situ_beta: float = 0.0
+  activation_situ_linear_beta: float = 0.0
+  routed_expert_dim: int = 0
+  latent_moe_norm: bool = False
+  route_weights_uncorrected: bool = False
+  attn_res_block_size: int = 0
+  kda_full_rank_gate: bool = False
+  kda_gate_lower_bound: float = 0.0
+  recurrent_prefill_chunked: bool = False
+  recurrent_prefill_chunk_size: int = 0

 class FFNBlock:
  def __init__(self, config:TransformerConfig):
@@ -86,9 +143,15 @@ class FFNBlock:
    if config.num_experts > 0:
      self.ffn_gate_inp = Linear(config.dim, config.num_experts, bias=False)  # router
      if config.expert_bias: self.exp_probs_b = {"bias": Tensor.zeros(config.num_experts)}
-      self.ffn_gate_exps = ExpertWeights(config.num_experts, config.dim, config.hidden_dim)
-      self.ffn_up_exps = ExpertWeights(config.num_experts, config.dim, config.hidden_dim)
-      self.ffn_down_exps = ExpertWeights(config.num_experts, config.hidden_dim, config.dim)
+      expert_cls = MXFP4ExpertWeights if config.expert_mxfp4 else ExpertWeights
+      expert_dim = config.routed_expert_dim or config.dim
+      self.ffn_gate_exps = expert_cls(config.num_experts, expert_dim, config.hidden_dim)
+      self.ffn_up_exps = expert_cls(config.num_experts, expert_dim, config.hidden_dim)
+      self.ffn_down_exps = expert_cls(config.num_experts, config.hidden_dim, expert_dim)
+      if config.routed_expert_dim:
+        self.ffn_routed_down = Linear(config.dim, expert_dim, bias=False)
+        self.ffn_routed_up = Linear(expert_dim, config.dim, bias=False)
+        if config.latent_moe_norm: self.ffn_routed_norm = nn.RMSNorm(expert_dim, config.norm_eps)
      if config.shared_expert_dim > 0:
        self.ffn_gate_shexp = Linear(config.dim, config.shared_expert_dim, bias=False)
        self.ffn_up_shexp = Linear(config.dim, config.shared_expert_dim, bias=False)
@@ -99,28 +162,85 @@ class FFNBlock:
      self.ffn_up      = Linear(config.dim, config.hidden_dim, bias=False)
      self.ffn_down    = Linear(config.hidden_dim, config.dim, bias=False)

+    if config.attn_res_block_size:
+      self.attn_res_norm, self.mlp_res_norm = nn.RMSNorm(config.dim, config.norm_eps), nn.RMSNorm(config.dim, config.norm_eps)
+      self.attn_res_proj, self.mlp_res_proj = Linear(config.dim, 1, bias=False), Linear(config.dim, 1, bias=False)
+
+  def _activation(self, gate:Tensor, up:Tensor) -> Tensor:
+    if not self.config.activation_situ_beta: return gate.silu() * up
+    gate32, up32, beta = gate.float(), up.float(), self.config.activation_situ_beta
+    gate32 = beta * (gate32 / beta).tanh() * gate32.sigmoid()
+    if (linear_beta := self.config.activation_situ_linear_beta): up32 = linear_beta * (up32 / linear_beta).tanh()
+    return (gate32 * up32).cast(gate.dtype)
+
  def _feed_forward(self, x:Tensor) -> Tensor:
    if hasattr(self, 'ffn_gate_exps'):
      h = x.unsqueeze(2)  # (B, T, 1, D) - add expert dim for broadcasting
-      logits = self.ffn_gate_inp(x)
+      # Kimi computes router logits in FP32 even though the residual stream and weights are BF16.
+      logits = x.float().linear(self.ffn_gate_inp.weight.float().transpose()) if self.config.bf16_activations else self.ffn_gate_inp(x)
      if hasattr(self, 'exp_probs_b'):
-        probs = logits.sigmoid()
-        _, sel = pairwise_topk(probs + self.exp_probs_b["bias"], self.config.num_experts_per_tok)
-        probs = probs.gather(-1, sel)
-        if self.config.norm_topk_prob: probs = probs / probs.sum(axis=-1, keepdim=True)
+        scores = logits.sigmoid()
+        adjusted_scores = scores + self.exp_probs_b["bias"]
+        topk = iterative_topk if self.config.num_experts >= 512 else pairwise_topk
+        _, sel = topk(adjusted_scores, self.config.num_experts_per_tok)
+        probs = (scores if self.config.route_weights_uncorrected else adjusted_scores).gather(-1, sel)
+        # Kimi-Linear-48B's older reference weights corrected scores. K3 selects with the correction
+        # but gathers the uncorrected sigmoid scores, so keep this an explicit compatibility switch.
+        if self.config.norm_topk_prob: probs = probs / (probs.sum(axis=-1, keepdim=True) + 1e-20)
      else:
        vals, sel = pairwise_topk(logits, self.config.num_experts_per_tok)
        probs = vals.softmax(-1) if self.config.norm_topk_prob else logits.softmax(-1).gather(-1, sel)
      probs = probs * self.config.routed_scaling_factor
-      x_down = self.ffn_down_exps(sel, (self.ffn_gate_exps(sel, h).silu() * self.ffn_up_exps(sel, h)).contiguous())  # (B, T, k, D)
-      out = (x_down * probs.unsqueeze(-1)).sum(axis=2)  # (B, T, D)
+      if hasattr(self, 'ffn_routed_down'): h = self.ffn_routed_down(x).unsqueeze(2)
+      if isinstance(self.ffn_gate_exps, MXFP4ExpertWeights) and amd_packed_mxfp4_supported(h.device):
+        hq = mxfp8_quantize_dequantize(h.cast(dtypes.bfloat16)) if amd_custom_kernels_supported(h.device) else \
+          quantize_dequantize_mxfp8(h.cast(dtypes.bfloat16))
+        gate = self.ffn_gate_exps(sel, hq, quantized=True)
+        up = cast(MXFP4ExpertWeights, self.ffn_up_exps)(sel, hq, quantized=True)
+      else: gate, up = self.ffn_gate_exps(sel, h), self.ffn_up_exps(sel, h)
+      routed_activation = self._activation(gate, up).contiguous()
+      combine_down = resolve(x.shape[1] == 1) and isinstance(self.ffn_down_exps, MXFP4ExpertWeights) and \
+        hasattr(self, 'ffn_gate_shexp') and not hasattr(self, 'ffn_routed_up') and amd_custom_kernels_supported(x.device)
+      x_down = cast(MXFP4ExpertWeights, self.ffn_down_exps)(sel, routed_activation, partial=True) if combine_down else \
+        self.ffn_down_exps(sel, routed_activation)
+      out = (x_down * probs.unsqueeze(-1).unsqueeze(-1)).sum(axis=2) if combine_down else \
+        (x_down * probs.unsqueeze(-1)).sum(axis=2).cast(x_down.dtype)  # (B, T, D[, devices])
+      combine_final = resolve(x.shape[1] == 1) and hasattr(self, 'ffn_routed_up') and hasattr(self, 'ffn_gate_shexp') and \
+        not hasattr(self, 'ffn_gate_inp_shexp') and isinstance(self.ffn_routed_up.weight.device, tuple) and \
+        isinstance(self.ffn_down_shexp.weight.device, tuple) and self.ffn_routed_up.weight.uop.axis == self.ffn_down_shexp.weight.uop.axis == 1 and \
+        self.ffn_routed_up.weight.shape[1] % (32*len(self.ffn_routed_up.weight.device)) == 0 and \
+        self.ffn_down_shexp.weight.shape[1] % (32*len(self.ffn_down_shexp.weight.device)) == 0 and \
+        amd_exact_bf16_custom_kernels_supported(x.device)
+      if hasattr(self, 'ffn_routed_up'):
+        if hasattr(self, 'ffn_routed_norm'): out = self.ffn_routed_norm(out)
+        out = bf16_partial_linear(out, self.ffn_routed_up.weight) if combine_final else self.ffn_routed_up(out)
      if hasattr(self, 'ffn_gate_shexp'):
-        shexp = self.ffn_down_shexp(self.ffn_gate_shexp(x).silu().contiguous() * self.ffn_up_shexp(x))
-        if hasattr(self, 'ffn_gate_inp_shexp'): shexp = shexp * (x * self.ffn_gate_inp_shexp["weight"]).sum(axis=-1, keepdim=True).sigmoid()
-        out = out + shexp
+        if resolve(x.shape[1] == 1) and amd_exact_bf16_custom_kernels_supported(x.device) and \
+           self.ffn_gate_shexp.weight.shape == self.ffn_up_shexp.weight.shape:
+          shared_gate, shared_up = dual_bf16_matvec(x, self.ffn_gate_shexp.weight, self.ffn_up_shexp.weight,
+                                                    fast=amd_custom_kernels_supported(x.device))
+        else: shared_gate, shared_up = self.ffn_gate_shexp(x).contiguous(), self.ffn_up_shexp(x).contiguous()
+        shared_activation = self._activation(shared_gate, shared_up).contiguous()
+        if combine_down:
+          out = (out + bf16_partial_linear(shared_activation, self.ffn_down_shexp.weight)).sum(3).cast(dtypes.bfloat16)
+        elif combine_final:
+          out = (out + bf16_partial_linear(shared_activation, self.ffn_down_shexp.weight)).sum(3).cast(dtypes.bfloat16)
+        else:
+          shexp = self.ffn_down_shexp(shared_activation)
+          if hasattr(self, 'ffn_gate_inp_shexp'):
+            shexp = shexp * (x * self.ffn_gate_inp_shexp["weight"]).sum(axis=-1, keepdim=True).sigmoid()
+          out = out + shexp
      return out
    # TODO: remove the need for this contiguous
-    return self.ffn_down(self.ffn_gate(x).silu().contiguous() * self.ffn_up(x))
+    if resolve(x.shape[1] == 1) and amd_exact_bf16_custom_kernels_supported(x.device) and \
+       self.ffn_gate.weight.shape == self.ffn_up.weight.shape:
+      dense_gate, dense_up = dual_bf16_matvec(x, self.ffn_gate.weight, self.ffn_up.weight, fast=amd_custom_kernels_supported(x.device))
+    else: dense_gate, dense_up = self.ffn_gate(x).contiguous(), self.ffn_up(x).contiguous()
+    dense_activation = self._activation(dense_gate, dense_up).contiguous()
+    if resolve(x.shape[1] == 1) and isinstance(self.ffn_down.weight.device, tuple) and self.ffn_down.weight.uop.axis == 1 and \
+       self.ffn_down.weight.shape[1] % (32*len(self.ffn_down.weight.device)) == 0 and amd_exact_bf16_custom_kernels_supported(x.device):
+      return bf16_partial_linear(dense_activation, self.ffn_down.weight).sum(3).cast(dtypes.bfloat16)
+    return self.ffn_down(dense_activation)

  # given the token-prefix match, return how much cached state this block can still reuse
  def _reusable_prefix_len(self, prefix_len:int, cached_len:int) -> int: return prefix_len
@@ -131,6 +251,11 @@ class FFNBlock:

  def __call__(self, x: Tensor, start_pos: int|UOp):
    self._init_state(x)
+    # Kimi's heterogeneous TP shards are captured by the outer TinyJit; per-block precompilation
+    # cannot represent their differently shaped local buffers as one implicit parameter bundle.
+    if self.config.bf16_activations:
+      h = x + self._attention(self.attn_norm(x), start_pos)
+      return (h + self._feed_forward(self.ffn_norm(h))).contiguous()
    # we pass in the weights implicitly so we unpack the GGUF on the fly
    @function(precompile=True, allow_implicit=True)
    def _run(x:Tensor, start_pos:int|UOp):
@@ -138,6 +263,32 @@ class FFNBlock:
      return (h + self._feed_forward(self.ffn_norm(h))).contiguous()
    return _run(x, start_pos)

+  @staticmethod
+  def _apply_attn_res(prefix_sum:Tensor, block_residual:Tensor, proj:Linear, norm:nn.RMSNorm) -> Tensor:
+    # Both inputs are flattened over B*T. Scoring is intentionally FP32, matching K3 eager inference.
+    v = block_residual.cat(prefix_sum.unsqueeze(1), dim=1)
+    vf = v.float()
+    k = vf * (vf.square().mean(axis=-1, keepdim=True) + norm.eps).rsqrt()
+    assert norm.weight is not None
+    scores = (k * (norm.weight.float() * proj.weight.squeeze(0).float())).sum(axis=-1)
+    return (scores.softmax(-1).unsqueeze(1) @ vf).squeeze(1).cast(v.dtype)
+
+  def attn_residual(self, x:Tensor, start_pos:int|UOp, block_residual:Tensor, layer_idx:int) -> tuple[Tensor, Tensor]:
+    self._init_state(x)
+    shape, prefix_sum = x.shape, x
+    prefix:Tensor|None = prefix_sum
+    if block_residual.shape[1]: x = self._apply_attn_res(x.reshape(-1, shape[-1]), block_residual,
+                                                        self.attn_res_proj, self.attn_res_norm).reshape(shape)
+    if layer_idx % self.config.attn_res_block_size == 0:
+      block_residual = block_residual.cat(prefix_sum.reshape(-1, shape[-1]).unsqueeze(1), dim=1)
+      prefix = None
+    attn = self._attention(self.attn_norm(x), start_pos)
+    prefix = attn if prefix is None else prefix + attn
+    x = self._apply_attn_res(prefix.reshape(-1, shape[-1]), block_residual,
+                             self.mlp_res_proj, self.mlp_res_norm).reshape(shape)
+    mlp = self._feed_forward(self.ffn_norm(x))
+    return (prefix + mlp).contiguous(), block_residual
+
 class TransformerBlock(FFNBlock):
  def __init__(self, config:TransformerConfig):
    super().__init__(config)
@@ -179,7 +330,9 @@ class TransformerBlock(FFNBlock):

    # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
    # TODO: this if statement should be removed and it shouldn't generate extra kernels
-    mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, buffer=False).triu(start_pos+1) \
+    # Build the static T×T causal corner on-device, then prepend the unmasked cached prefix.
+    # A broadcast const with symbolic width otherwise defaults to CPU in multi-device graphs.
+    mask = Tensor.full((1, 1, T, T), float("-inf"), dtype=x.dtype, device=x.device).triu(1).pad(((0, 0),)*3+((start_pos, 0),)) \
      if resolve(T != 1) else None
    attn = q.scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True)     # (B,H,T,Hd)
    attn = attn.transpose(1, 2).reshape(B, T, -1)                                    # back to (B,T,D)
@@ -188,7 +341,8 @@ class TransformerBlock(FFNBlock):
  def _init_state(self, x:Tensor):
    if not hasattr(self, "cache_kv"):
      # TODO: how is the dtype of this determined?
-      self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim, device=x.device)
+      self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim,
+                                   device=x.device, dtype=x.dtype)
      self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta, device=x.device)

 class MLATransformerBlock(FFNBlock):
@@ -206,17 +360,23 @@ class MLATransformerBlock(FFNBlock):
    self.attn_k_b = {"weight": Tensor.zeros(config.n_heads, config.kv_lora_rank, qk_nope_head_dim)}
    self.attn_v_b = {"weight": Tensor.zeros(config.n_heads, config.v_head_dim, config.kv_lora_rank)}
    self.attn_output = Linear(config.n_heads * config.v_head_dim, config.dim, bias=False)
+    if config.attn_output_gate: self.attn_gate = Linear(config.dim, config.n_heads * config.v_head_dim, bias=False)

  def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
    B, T, _ = x.shape
    q_nope_head_dim = self.config.head_dim - self.config.rope_dim
-    q_proj = self.attn_q_b(self.attn_q_a_norm(self.attn_q_a(x))) if self.config.q_lora_rank > 0 else self.attn_q(x)
+    mfma_decode = resolve(T == 1) and x.shape[-1] % 256 == 0 and amd_wave64_custom_kernels_supported(x.device)
+    q_a_mfma = self.config.q_lora_rank > 0 and mfma_decode and self.attn_q_a.weight.shape[0] % 16 == 0
+    q_a = bf16_mfma_splitk(x, self.attn_q_a.weight) if q_a_mfma else \
+      self.attn_q_a(x) if self.config.q_lora_rank > 0 else None
+    q_proj = self.attn_q_b(self.attn_q_a_norm(q_a)) if q_a is not None else self.attn_q(x)
    q = q_proj.reshape(B, T, self.config.n_heads, self.config.head_dim).transpose(1, 2)
    q_nope, q_rope = q[..., :q_nope_head_dim], q[..., q_nope_head_dim:]
    if not self.config.ssm or not self.config.ssm.kda: q_rope = apply_rope(q_rope, self.freqs_cis[start_pos:start_pos+T])
    q = (q_nope @ self.attn_k_b["weight"].transpose(-1, -2)).cat(q_rope, dim=-1)

-    kv_a = self.attn_kv_a_mqa(x)
+    kv_a_mfma = mfma_decode and self.attn_kv_a_mqa.weight.shape[0] % 16 == 0
+    kv_a = bf16_mfma_splitk(x, self.attn_kv_a_mqa.weight) if kv_a_mfma else self.attn_kv_a_mqa(x)
    c_kv = self.attn_kv_a_norm(kv_a[..., :self.config.kv_lora_rank])
    k_rope = kv_a[..., self.config.kv_lora_rank:].reshape(B, T, 1, self.config.rope_dim).transpose(1, 2)
    if not self.config.ssm or not self.config.ssm.kda: k_rope = apply_rope(k_rope, self.freqs_cis[start_pos:start_pos+T])
@@ -225,17 +385,23 @@ class MLATransformerBlock(FFNBlock):
    k = Tensor(self.cache_k.uop.after(self.cache_k[:, :, start_pos:start_pos+T, :].uop.store(k_store.uop)))[:, :, 0:start_pos+T, :]
    v = k[..., :self.config.kv_lora_rank]

-    mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, buffer=False).triu(start_pos+1) \
+    mask = Tensor.full((1, 1, T, T), float("-inf"), dtype=x.dtype, device=x.device).triu(1).pad(((0, 0),)*3+((start_pos, 0),)) \
      if resolve(T != 1) else None
    attn = q @ k.transpose(-1, -2) * (1.0 / self.config.head_dim ** 0.5)
    if mask is not None: attn = attn + mask
-    attn = attn.softmax(-1)
+    # Match eager Kimi MLA: normalize attention scores in FP32, then return to the query dtype.
+    attn = attn.softmax(-1, dtype=dtypes.float32).cast(q.dtype)
    attn = ((attn @ v) @ self.attn_v_b["weight"].transpose(-1, -2)).transpose(1, 2).reshape(B, T, -1)
+    if hasattr(self, "attn_gate"): attn = attn * self.attn_gate(x).sigmoid()
+    if resolve(T == 1) and isinstance(self.attn_output.weight.device, tuple) and \
+       self.attn_output.weight.shape[1] % (32*len(self.attn_output.weight.device)) == 0 and amd_exact_bf16_custom_kernels_supported(attn.device):
+      return bf16_partial_linear(attn, self.attn_output.weight).sum(3).cast(dtypes.bfloat16)
    return self.attn_output(attn)

  def _init_state(self, x:Tensor):
    if not hasattr(self, "cache_k"):
-      self.cache_k = Tensor.empty(x.shape[0], 1, self.config.max_context, self.config.kv_lora_rank + self.config.rope_dim, device=x.device)
+      self.cache_k = Tensor.empty(x.shape[0], 1, self.config.max_context, self.config.kv_lora_rank + self.config.rope_dim,
+                                  device=x.device, dtype=x.dtype)
      self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta, device=x.device)

 class GatedDeltaNetBlock(FFNBlock):
@@ -245,68 +411,144 @@ class GatedDeltaNetBlock(FFNBlock):
    assert self.num_v_heads % self.num_k_heads == 0
    self.head_v_dim, self.ssm_conv_kernel = ssm.inner_size // ssm.time_step_rank, ssm.conv_kernel
    self.conv_channels, self.q_dim = ssm.inner_size + 2*ssm.group_count*ssm.state_size, ssm.state_size*ssm.group_count
-    self.attn_qkv = Linear(config.dim, self.conv_channels, bias=False)
+    if ssm.kda and config.kda_split_qkv:
+      self.attn_q, self.attn_k = Linear(config.dim, self.q_dim, bias=False), Linear(config.dim, self.q_dim, bias=False)
+      self.attn_v = Linear(config.dim, ssm.inner_size, bias=False)
+      self.ssm_q_conv1d = {"weight": Tensor.zeros(self.q_dim, self.ssm_conv_kernel)}
+      self.ssm_k_conv1d = {"weight": Tensor.zeros(self.q_dim, self.ssm_conv_kernel)}
+      self.ssm_v_conv1d = {"weight": Tensor.zeros(ssm.inner_size, self.ssm_conv_kernel)}
+    else:
+      self.attn_qkv = Linear(config.dim, self.conv_channels, bias=False)
+      self.ssm_conv1d = {"weight": Tensor.zeros(self.conv_channels, self.ssm_conv_kernel)}
    if ssm.kda:
-      self.ssm_g_a, self.ssm_g_b = Linear(config.dim, self.head_v_dim, bias=False), Linear(self.head_v_dim, ssm.inner_size, bias=False)
+      if config.kda_full_rank_gate: self.ssm_g_full = Linear(config.dim, ssm.inner_size, bias=False)
+      else: self.ssm_g_a, self.ssm_g_b = Linear(config.dim, self.head_v_dim, bias=False), Linear(self.head_v_dim, ssm.inner_size, bias=False)
      self.ssm_f_a, self.ssm_f_b = Linear(config.dim, self.head_k_dim, bias=False), Linear(self.head_k_dim, ssm.inner_size, bias=False)
    else:
      self.attn_gate = Linear(config.dim, ssm.inner_size, bias=False)
      self.ssm_alpha = Linear(config.dim, self.num_v_heads, bias=False)
    self.ssm_beta = Linear(config.dim, self.num_v_heads, bias=False)
-    self.ssm_conv1d = {"weight": Tensor.zeros(self.conv_channels, self.ssm_conv_kernel)}
    self.ssm_dt = {"bias": Tensor.zeros(ssm.inner_size if ssm.kda else self.num_v_heads)}
-    self.ssm_a = Tensor.zeros(self.num_v_heads, 1) if ssm.kda else Tensor.zeros(self.num_v_heads)
+    self.ssm_a = Tensor.zeros(self.head_v_dim if ssm.channel_decay else self.num_v_heads, 1) if ssm.kda else Tensor.zeros(self.num_v_heads)
+    self.kda_channel_decay = ssm.channel_decay
    self.ssm_norm, self.ssm_out = nn.RMSNorm(self.head_v_dim, config.norm_eps), Linear(ssm.inner_size, config.dim, bias=False)

  def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
    B, T, _ = x.shape
-    assert T == 1, "GatedDeltaNetBlock currently only supports T=1"

    # input processing
-    x = x.half()
-    out_gate = self.ssm_g_b(self.ssm_g_a(x)) if hasattr(self, "ssm_g_a") else self.attn_gate(x)
-    out_gate = out_gate.reshape(B, 1, self.num_v_heads, self.head_v_dim)
-    beta = self.ssm_beta(x).sigmoid().reshape(B, self.num_v_heads, 1, 1)
-    alpha = self.ssm_f_b(self.ssm_f_a(x)) if hasattr(self, "ssm_f_a") else self.ssm_alpha(x)
-    alpha = ((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, self.num_v_heads, -1) *
-             self.ssm_a.reshape(1, self.num_v_heads, -1)).exp().unsqueeze(-2)
+    # Kimi-Linear is a BF16 model. Qwen 3.5 GGDN checkpoints historically use FP16 here.
+    x = x.cast(dtypes.bfloat16) if self.config.ssm and self.config.ssm.kda else x.half()
+    fused_fg = hasattr(self, "ssm_g_a") and resolve(T == 1) and amd_custom_kernels_supported(x.device) and \
+      self.ssm_g_a.weight.shape == self.ssm_f_a.weight.shape
+    mfma_decode = resolve(T == 1) and x.shape[-1] % 256 == 0 and amd_wave64_custom_kernels_supported(x.device)
+    if hasattr(self, "ssm_g_full"): out_gate = self.ssm_g_full(x)
+    elif hasattr(self, "ssm_g_a"):
+      if fused_fg:
+        gate_a, alpha_a, beta_logits = kda_fgb_linear(x, self.ssm_g_a.weight, self.ssm_f_a.weight, self.ssm_beta.weight)
+        out_gate, alpha_logits = dual_input_bf16_matvec(gate_a, alpha_a, self.ssm_g_b.weight, self.ssm_f_b.weight)
+      else: out_gate = self.ssm_g_b(self.ssm_g_a(x))
+    else: out_gate = self.attn_gate(x)
+    if not fused_fg: beta_logits = self.ssm_beta(x)
+    if not fused_fg:
+      if hasattr(self, "ssm_f_a"):
+        f_a = bf16_mfma_splitk(x, self.ssm_f_a.weight) if mfma_decode and self.ssm_f_a.weight.shape[0] % 16 == 0 else self.ssm_f_a(x)
+        alpha_logits = self.ssm_f_b(f_a)
+      else: alpha_logits = self.ssm_alpha(x)

-    # qkv conv
-    conv_window = self.conv_state.cat(self.attn_qkv(x), dim=1)
-    conv_out = (conv_window * self.ssm_conv1d["weight"].T.unsqueeze(0)).sum(1).silu()
-    q, k, v = conv_out.split([self.q_dim, self.q_dim, self.conv_channels - 2*self.q_dim], dim=-1)
-    q = q.reshape(B, self.num_k_heads, self.head_k_dim).normalize(dim=-1).repeat(1, self.num_v_heads//self.num_k_heads, 1)
-    k = k.reshape(B, self.num_k_heads, self.head_k_dim).normalize(dim=-1).repeat(1, self.num_v_heads//self.num_k_heads, 1)
-    v = v.reshape(B, self.num_v_heads, self.head_v_dim)
-    q, k, v = q.mul(self.head_k_dim**-0.5).unsqueeze(-1), k.unsqueeze(-1), v.unsqueeze(-1)
+    # Causal depthwise Q/K/V convolution. All tokens are projected together, then the recurrent
+    # update is fused into one kernel so prefill doesn't build a Python-unrolled graph.
+    split_qkv = hasattr(self, "attn_q")
+    if split_qkv:
+      if resolve(T == 1) and amd_packed_mxfp4_supported(x.device) and \
+         self.attn_q.weight.shape == self.attn_k.weight.shape == self.attn_v.weight.shape:
+        projected_q, projected_k, projected_v = kda_qkv_linear(x, self.attn_q.weight, self.attn_k.weight, self.attn_v.weight)
+      else: projected_q, projected_k, projected_v = self.attn_q(x), self.attn_k(x), self.attn_v(x)
+      # Snapshot mutable caches before constructing the recurrence. Otherwise the final store can
+      # overwrite their buffers before earlier outputs in a multi-token lazy graph consume them.
+      conv_state_q, conv_state_k, conv_state_v = self.conv_state_q.clone(), self.conv_state_k.clone(), self.conv_state_v.clone()
+    else: projected, conv_state = self.attn_qkv(x), self.conv_state
+    def causal_conv(projected:Tensor, state:Tensor, weight:Tensor) -> tuple[Tensor, Tensor]:
+      window = state.cat(projected, dim=1)
+      out = functools.reduce(lambda a,b: a+b, (window[:, i:i+T] * weight[:, i] for i in range(self.ssm_conv_kernel))).silu()
+      return out, window[:, T:T+self.ssm_conv_kernel-1]
+    if split_qkv:
+      q, conv_state_q = causal_conv(projected_q, conv_state_q, self.ssm_q_conv1d["weight"])
+      k, conv_state_k = causal_conv(projected_k, conv_state_k, self.ssm_k_conv1d["weight"])
+      v, conv_state_v = causal_conv(projected_v, conv_state_v, self.ssm_v_conv1d["weight"])
+    else:
+      conv_out, conv_state = causal_conv(projected, conv_state, self.ssm_conv1d["weight"])
+      q, k, v = conv_out.split([self.q_dim, self.q_dim, self.conv_channels - 2*self.q_dim], dim=-1)

-    # recurrent
-    recurrent_state = self.recurrent_state * alpha
-    recurrent_state = recurrent_state + ((v - recurrent_state@k) * beta)@k.transpose(-1, -2)
+    q, k = q.reshape(B, T, self.num_k_heads, self.head_k_dim), k.reshape(B, T, self.num_k_heads, self.head_k_dim)
+    q, k = (l2norm(q), l2norm(k)) if self.config.ssm and self.config.ssm.kda else (q.normalize(dim=-1), k.normalize(dim=-1))
+    q = q.repeat(1, 1, self.num_v_heads//self.num_k_heads, 1).transpose(1, 2).float() * self.head_k_dim**-0.5
+    k = k.repeat(1, 1, self.num_v_heads//self.num_k_heads, 1).transpose(1, 2).float()
+    v = v.reshape(B, T, self.num_v_heads, self.head_v_dim).transpose(1, 2).float()
+    beta = (beta_logits.float() if self.config.ssm and self.config.ssm.kda else beta_logits).sigmoid().transpose(1, 2)
+    gate_logits = (alpha_logits.float() + self.ssm_dt["bias"]).reshape(B, T, self.num_v_heads, -1)
+    a_shape = (1, 1, 1, self.head_v_dim) if self.kda_channel_decay else (1, 1, self.num_v_heads, 1)
+    if self.config.kda_gate_lower_bound:
+      log_alpha = self.config.kda_gate_lower_bound * ((-self.ssm_a).reshape(a_shape) * gate_logits).sigmoid()
+    else: log_alpha = gate_logits.softplus() * self.ssm_a.reshape(a_shape)
+    alpha = log_alpha.squeeze(-1).transpose(1, 2).exp() if log_alpha.shape[-1] == 1 else log_alpha.permute(0, 2, 1, 3).exp()
+    if T == 1:
+      # Keep decode on the small elementwise graph. The fused prefill kernel writes a temporary
+      # recurrent matrix, which is worthwhile for multiple tokens but needlessly copies state at T=1.
+      decay = alpha if len(alpha.shape) == 4 else alpha.unsqueeze(-1)
+      recurrent_state = self.recurrent_state * decay
+      k1, q1 = k[:, :, 0].unsqueeze(-1), q[:, :, 0].unsqueeze(-1)
+      recurrent_state = recurrent_state + ((v[:, :, 0].unsqueeze(-1) - recurrent_state@k1) * beta[:, :, 0].reshape(B, self.num_v_heads, 1, 1)) @ \
+                        k1.transpose(-1, -2)
+      core = (recurrent_state @ q1).squeeze(-1).unsqueeze(2)
+    else: core, recurrent_state = gated_delta_prefill(q, k, v, beta, alpha, self.recurrent_state)
+    core = core.transpose(1, 2)

-    # store the updated state
-    conv_state_store = self.conv_state.uop.store(conv_window[:, 1:, :].cast(self.conv_state.dtype).uop)
-    recurrent_state_store = self.recurrent_state.uop.store(recurrent_state.cast(self.recurrent_state.dtype).uop)
-    recurrent_state = Tensor(self.recurrent_state.uop.after(recurrent_state_store, conv_state_store))
+    # Store each cache with its own AFTER. Multi-device lowering handles one sharded STORE per
+    # AFTER; grouping these effects under one cache silently drops stores on the other shards.
+    state_updates:list[Tensor]
+    if split_qkv:
+      state_updates = [self.conv_state_q.assign(conv_state_q.cast(self.conv_state_q.dtype)),
+                       self.conv_state_k.assign(conv_state_k.cast(self.conv_state_k.dtype)),
+                       self.conv_state_v.assign(conv_state_v.cast(self.conv_state_v.dtype))]
+    else: state_updates = [self.conv_state.assign(conv_state.cast(self.conv_state.dtype))]
+    state_updates.append(self.recurrent_state.assign(recurrent_state.cast(self.recurrent_state.dtype)))
+    core_attn_out = self.ssm_norm(core.cast(x.dtype) if self.config.ssm and self.config.ssm.kda else core)
+    gate = out_gate.reshape(B, T, self.num_v_heads, self.head_v_dim)
+    gate = gate.float().sigmoid().cast(core_attn_out.dtype) if hasattr(self, "ssm_g_a") else gate.silu()
+    out = (core_attn_out * gate).reshape(B, T, -1)
+    out = out.cast(x.dtype)
+    ret = bf16_partial_linear(out, self.ssm_out.weight).sum(3).cast(dtypes.bfloat16) if resolve(T == 1) and \
+      isinstance(self.ssm_out.weight.device, tuple) and self.ssm_out.weight.shape[1] % (32*len(self.ssm_out.weight.device)) == 0 and \
+      amd_exact_bf16_custom_kernels_supported(out.device) else self.ssm_out(out)
+    return ret.realize(*state_updates)

-    # output
-    core_attn_out = self.ssm_norm((recurrent_state@q).squeeze(-1).reshape(B, 1, self.num_v_heads, self.head_v_dim))
-    out_gate = out_gate.sigmoid() if hasattr(self, "ssm_g_a") else out_gate.silu()
-    return self.ssm_out((core_attn_out * out_gate).reshape(B, 1, -1).cast(x.dtype))
-
-  # recurrent state can't be partially reused after divergence, force a full rebuild
-  def _state_reset_ops(self):
-    return [self.conv_state.assign(self.conv_state.const_like(0)),
-            self.recurrent_state.assign(self.recurrent_state.const_like(0))] if hasattr(self, "conv_state") else []
-  def _reusable_prefix_len(self, prefix_len:int, cached_len:int) -> int: return 0 if prefix_len != cached_len else prefix_len
+  # Recurrent state can be reused only when the new prompt exactly extends all currently valid state.
+  def _state_tensors(self) -> tuple[Tensor, ...]:
+    if hasattr(self, "conv_state_q"):
+      return self.conv_state_q, self.conv_state_k, self.conv_state_v, self.recurrent_state
+    return (self.conv_state, self.recurrent_state) if hasattr(self, "conv_state") else ()
+  def _state_reset_ops(self): return [s.assign(s.const_like(0)) for s in self._state_tensors()]
+  def _reusable_prefix_len(self, prefix_len:int, cached_len:int) -> int: return prefix_len if prefix_len == cached_len else 0

  def _init_state(self, x):
-    if not hasattr(self, "conv_state"):
-      self.conv_state = Tensor.zeros(x.shape[0], self.ssm_conv_kernel-1, self.conv_channels, device=x.device).clone()
-      self.recurrent_state = Tensor.zeros(x.shape[0], self.num_v_heads, self.head_v_dim, self.head_k_dim, device=x.device).clone()
+    if not hasattr(self, "conv_state") and not hasattr(self, "conv_state_q"):
+      if hasattr(self, "attn_q"):
+        device = x.device[0] if isinstance(x.device, tuple) else x.device
+        self.conv_state_q = Tensor.zeros(x.shape[0], self.ssm_conv_kernel-1, self.q_dim, device=device, dtype=x.dtype).clone()
+        self.conv_state_k = Tensor.zeros(x.shape[0], self.ssm_conv_kernel-1, self.q_dim, device=device, dtype=x.dtype).clone()
+        self.conv_state_v = Tensor.zeros(x.shape[0], self.ssm_conv_kernel-1, self.num_v_heads*self.head_v_dim, device=device, dtype=x.dtype).clone()
+        self.recurrent_state = Tensor.zeros(x.shape[0], self.num_v_heads, self.head_v_dim, self.head_k_dim, device=device).clone()
+        if isinstance(x.device, tuple):
+          for state in (self.conv_state_q, self.conv_state_k, self.conv_state_v): state.shard_(x.device, axis=2)
+          self.recurrent_state.shard_(x.device, axis=1)
+      else:
+        self.conv_state = Tensor.zeros(x.shape[0], self.ssm_conv_kernel-1, self.conv_channels, device=x.device, dtype=x.dtype).clone()
+        self.recurrent_state = Tensor.zeros(x.shape[0], self.num_v_heads, self.head_v_dim, self.head_k_dim, device=x.device).clone()

 class Transformer:
  def __init__(self, config:TransformerConfig):
+    self.config = config
    dense_config = replace(config, num_experts=0, num_experts_per_tok=0, shared_expert_dim=0, hidden_dim=config.dense_hidden_dim or config.hidden_dim)
    if config.ssm: config = replace(config, qk_norm=config.head_dim)
    block_cls = MLATransformerBlock if config.kv_lora_rank > 0 else TransformerBlock
@@ -316,22 +558,76 @@ class Transformer:
    self.token_embd  = nn.Embedding(config.vocab_size, config.dim)
    self.output_norm = nn.RMSNorm(config.dim, config.norm_eps)
    self.output = Linear(config.dim, config.vocab_size, bias=False)
+    if config.attn_res_block_size:
+      self.output_attn_res_norm = nn.RMSNorm(config.dim, config.norm_eps)
+      self.output_attn_res_proj = Linear(config.dim, 1, bias=False)
    self.max_context = config.max_context
    self.has_recurrent_block = any(isinstance(b, GatedDeltaNetBlock) for b in self.blk)
    self._cached_tokens: list[int] = []
+    self._snapshot_tokens: list[int] = []
+    self._state_snapshots:list[Tensor] = []
+    self._token_buffer:Tensor|None = None
+    self._temperature_buffer:Tensor|None = None
    # we specialize the JIT for prefill and rollout
    self.prefill_jit = TinyJit(self.forward)
    self.rollout_jit = TinyJit(self.forward)
+    self.greedy_prefill_jit = TinyJit(self.forward)
+    self.greedy_rollout_jit = TinyJit(self.forward)
+    self.recurrent_prefill_jits:dict[int, Callable[..., Tensor]] = {}
+    self.recurrent_greedy_prefill_jits:dict[int, Callable[..., Tensor]] = {}
+    self.reset_jit = TinyJit(self._reset_state)
+    self.save_state_jit = TinyJit(self._save_state)
+    self.restore_state_jit = TinyJit(self._restore_state)

-  def forward(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor) -> Tensor:
-    x = self.token_embd(tokens).float()                   # (B, T, D)
-    for block in self.blk: x = block(x, start_pos)
-    logits = self.output(self.output_norm(x))[:, -1, :]
+  def _reset_state(self) -> None:
+    if resets := [r for b in self.blk for r in b._state_reset_ops()]: Tensor.realize(*resets)
+
+  def _state_tensors(self) -> list[Tensor]:
+    return [s for block in self.blk if isinstance(block, GatedDeltaNetBlock) for s in block._state_tensors()]
+
+  def _init_state_snapshots(self) -> None:
+    if not self._state_snapshots: self._state_snapshots = [s.clone().realize() for s in self._state_tensors()]
+
+  def _save_state(self) -> None:
+    if writes := [dst.assign(src) for dst,src in zip(self._state_snapshots, self._state_tensors())]: Tensor.realize(*writes)
+
+  def _restore_state(self) -> None:
+    if writes := [dst.assign(src) for dst,src in zip(self._state_tensors(), self._state_snapshots)]: Tensor.realize(*writes)
+
+  def forward(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor|None) -> Tensor:
+    if len(tokens.shape) == 1: tokens = tokens.reshape(1, -1)
+    x = self.token_embd(tokens).cast(dtypes.bfloat16) if self.config.bf16_activations else self.token_embd(tokens).float()
+    block_residual = Tensor.zeros(x.shape[0]*x.shape[1], 0, x.shape[2], device=x.device, dtype=x.dtype) \
+      if self.config.attn_res_block_size else None
+    for i, block in enumerate(self.blk):
+      if block_residual is not None: x, block_residual = block.attn_residual(x, start_pos, block_residual, i)
+      else: x = block(x, start_pos)
+      # Tensor indexing lowers selected experts through a fused one-hot reduction. Keeping all 26
+      # of those high-level graphs alive until the final output is scheduled exhausts host memory.
+      # A realization boundary lowers one block at a time; TinyJit still captures and memory-plans
+      # the resulting schedules for rollout replay.
+      if self.config.expert_mxfp4: x.realize()
+    if block_residual is not None:
+      x = FFNBlock._apply_attn_res(x.reshape(-1, x.shape[-1]), block_residual,
+                                   self.output_attn_res_proj, self.output_attn_res_norm).reshape(x.shape)
+    final_x = self.output_norm(x)
+    if temperature is None and resolve(tokens.numel() == 1) and amd_exact_bf16_custom_kernels_supported(x.device):
+      return bf16_matvec(final_x, self.output.weight).argmax(-1, keepdim=True)
+    logits = self.output(final_x)[:, -1, :]
+    if temperature is None: return logits.argmax(-1, keepdim=True)
    # Gumbel-max trick: argmax(logits/temp - log(-log(uniform))) is equivalent to sampling from softmax(logits/temp)
    return (logits / temperature.maximum(1e-12) - (Tensor.rand_like(logits).maximum(1e-12).log().neg()).log()).argmax(-1, keepdim=True)

-  def __call__(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor) -> Tensor:
-    return (self.prefill_jit if resolve(tokens.shape[1] != 1) else self.rollout_jit)(tokens.contiguous(), start_pos, temperature)
+  def __call__(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor|None) -> Tensor:
+    token_count = tokens.numel()
+    if self.has_recurrent_block and resolve(token_count != 1):
+      assert isinstance(token_count, int)
+      cache = self.recurrent_greedy_prefill_jits if temperature is None else self.recurrent_prefill_jits
+      jit = cache.setdefault(token_count, TinyJit(self.forward))
+      return jit(tokens.flatten().contiguous(), start_pos, temperature)
+    if temperature is None:
+      return (self.greedy_prefill_jit if resolve(token_count != 1) else self.greedy_rollout_jit)(tokens.flatten().contiguous(), start_pos, None)
+    return (self.prefill_jit if resolve(token_count != 1) else self.rollout_jit)(tokens.flatten().contiguous(), start_pos, temperature)

  @staticmethod
  def from_gguf(gguf:Tensor|str|pathlib.Path, max_context:int|None=None,
@@ -417,31 +713,91 @@ class Transformer:
    return model, kv

  def warmup(self):
-    for _ in range(2): list(zip(range(2), self.generate([0])))
+    # Capture the only two shapes used by recurrent serving: a full prefill chunk and one-token rollout.
+    # Two chunks exercise both the initial and nonzero-position prefill paths before the server opens.
+    recurrent_chunk = self.config.recurrent_prefill_chunk_size or 32
+    prompt = [0] * max(1, min(recurrent_chunk*2, self.max_context-2)) if self.has_recurrent_block else [0]
+    # Recurrent serving captures both greedy and sampled graphs, then executes one replay so graph
+    # creation/lowering cannot leak into request latency for either HTTP temperature path.
+    # generate mutates its token list, so each pass needs a fresh prompt to exercise cache reset.
+    for temperature in ((0.0, 1.0) if self.has_recurrent_block else (0.0,)):
+      for _ in range(3 if self.has_recurrent_block else 2): list(zip(range(2), self.generate(prompt.copy(), temperature=temperature)))
+    # Capture prompt-boundary restore using successively extended prompts so every restore starts
+    # from the checkpoint made by the previous pass.
+    if self.has_recurrent_block:
+      for i in range(1, 4): list(zip(range(2), self.generate(prompt + list(range(1, i+1)), temperature=0.0)))

-  def get_start_pos(self, tokens:list[int]) -> int:
+  def _cache_start(self, tokens:list[int]) -> tuple[int, bool]:
    prefix_len = sum(1 for _ in itertools.takewhile(lambda ab: ab[0] == ab[1], zip(tokens[:-1], self._cached_tokens)))
-    return min(block._reusable_prefix_len(prefix_len, len(self._cached_tokens)) for block in self.blk)
+    live_start = min(block._reusable_prefix_len(prefix_len, len(self._cached_tokens)) for block in self.blk)
+    snapshot_prefix = sum(1 for _ in itertools.takewhile(lambda ab: ab[0] == ab[1], zip(tokens[:-1], self._snapshot_tokens)))
+    snapshot_start = len(self._snapshot_tokens) if snapshot_prefix == len(self._snapshot_tokens) else 0
+    return (snapshot_start, True) if snapshot_start > live_start else (live_start, False)
+
+  def get_start_pos(self, tokens:list[int]) -> int: return self._cache_start(tokens)[0]

  def generate(self, tokens:list[int], chunk_size:int=32, temperature:float=0.0):
-    if self.has_recurrent_block: chunk_size = 1
+    chunked_recurrent = self.has_recurrent_block and self.config.recurrent_prefill_chunked
+    if chunked_recurrent and self.config.recurrent_prefill_chunk_size:
+      chunk_size = min(chunk_size, self.config.recurrent_prefill_chunk_size)
+    if self.has_recurrent_block and not chunked_recurrent: chunk_size = 1
    v_start_pos = UOp.variable("start_pos", 0, self.max_context-1)
    v_toks = UOp.variable("toks", 1, chunk_size)
    # TODO: use UOp.variable for temperature once float variables are supported
-    temp = Tensor([temperature])
-    # assign all input tokens once, then slice from start_pos for the model call
-    t = Tensor(tokens + [0] * (self.max_context - len(tokens)), dtype="int32").reshape(1, self.max_context)
+    model_device = self.token_embd.weight.device
+    amd_tp = isinstance(model_device, tuple) and all(d.startswith("AMD") for d in model_device)
+    if temperature == 0.0: temp = None
+    elif amd_tp:
+      if self._temperature_buffer is None: self._temperature_buffer = Tensor.empty(1, device=model_device).realize()
+      temp_storage = self._temperature_buffer.uop.buf_uop.buffer
+      temp_buffers = temp_storage.bufs if isinstance(temp_storage, MultiBuffer) else [temp_storage]
+      temp_host = memoryview(array.array('f', [temperature])).cast('B')
+      for buf in temp_buffers: buf.ensure_allocated().allocator._copyin(buf._buf, temp_host)
+      temp = self._temperature_buffer
+    else: temp = Tensor([temperature], device=model_device)
+    # Keep the replicated AMD token buffer identity stable across HTTP requests so captured graphs
+    # see the same input topology. Updating this small int32 buffer is cheaper than rebuilding JITs.
+    if amd_tp:
+      if self._token_buffer is None:
+        self._token_buffer = Tensor.empty(1, self.max_context, dtype=dtypes.int32, device=model_device).realize()
+      token_storage = self._token_buffer.uop.buf_uop.buffer
+      token_buffers = token_storage.bufs if isinstance(token_storage, MultiBuffer) else [token_storage]
+      input_host = memoryview(array.array('i', tokens + [0] * (self.max_context-len(tokens)))).cast('B')
+      for buf in token_buffers: buf.ensure_allocated().allocator._copyin(buf._buf, input_host)
+      t = self._token_buffer
+    else: t = Tensor(tokens + [0] * (self.max_context - len(tokens)), dtype="int32", device=model_device).reshape(1, self.max_context)
    # recompute start_pos from what's currently valid in the caches
-    start_pos = self.get_start_pos(tokens)
-    if start_pos < len(self._cached_tokens) and (resets := [r for b in self.blk for r in b._state_reset_ops()]): Tensor.realize(*resets)
+    start_pos, restore_snapshot = self._cache_start(tokens)
+    # This graph is captured by warmup. Resetting on-device avoids hundreds of synchronous copies
+    # to sharded AMD state buffers and guarantees unrelated requests don't schedule new kernels.
+    if restore_snapshot: self.restore_state_jit()
+    elif start_pos < len(self._cached_tokens) and self.has_recurrent_block: self.reset_jit()
    out, prompt_len = None, len(tokens)
+    token_host = memoryview(bytearray(4)) if amd_tp else None
    while len(tokens) < self.max_context:
-      n_toks = min(chunk_size, len(tokens) - start_pos)
-      sp, nt = v_start_pos.bind(start_pos), v_toks.bind(n_toks)
-      out = self(t[:, sp:sp+nt] if start_pos < prompt_len or out is None else out, sp, temp).realize()
+      remaining = len(tokens) - start_pos
+      # Full recurrent chunks use the high-throughput prefill graph. Process the tail through the
+      # rollout graph so every request uses only the two shapes captured during server warmup.
+      n_toks = chunk_size if chunked_recurrent and remaining >= chunk_size else 1 if chunked_recurrent else min(chunk_size, remaining)
+      # Recurrent blocks execute an explicit recurrence over T. Give them a static chunk length so
+      # Python constructs the recurrence once per encountered size; decode remains the T=1 JIT.
+      if chunked_recurrent:
+        # Token count is static for the recurrent kernel, but cache position must remain a runtime
+        # variable so repeated chunks do not replay MLA stores at the capture position.
+        sp = v_start_pos.bind(start_pos)
+        model_input = t[:, sp:sp+n_toks] if start_pos < prompt_len or out is None else out
+      else:
+        sp = v_start_pos.bind(start_pos)
+        nt = v_toks.bind(n_toks)
+        model_input = t[:, sp:sp+nt] if start_pos < prompt_len or out is None else out
+      out = self(model_input, sp, temp).realize()
      start_pos += n_toks
      # chunked prefill: keep processing until all prompt tokens are consumed
      if start_pos < len(tokens): continue
-      tokens.append(int(out.item()))
+      if self.has_recurrent_block and len(tokens) == prompt_len and self._state_tensors():
+        self._init_state_snapshots()
+        self.save_state_jit()
+        self._snapshot_tokens = tokens.copy()
+      tokens.append(amd_int32_item(out, token_host) if token_host is not None else int(out.item()))
      self._cached_tokens = tokens[:-1]
      yield tokens[-1]
@@ -0,0 +1,81 @@
+from tinygrad import Tensor, dtypes
+
+MX_BLOCK_SIZE = 32
+MXFP4_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
+                -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0)
+
+def _e8m0_scale(scale:Tensor) -> Tensor:
+  """Decode an OCP E8M0 scale byte. 127 encodes 2**0."""
+  return (scale.cast(dtypes.float32) - 127.0).exp2()
+
+def quantize_mxfp4(x:Tensor) -> tuple[Tensor, Tensor]:
+  """Quantize the last dimension to OCP MXFP4 (E2M1 values, E8M0 scale, block size 32)."""
+  if x.shape[-1] % MX_BLOCK_SIZE: raise ValueError(f"MXFP4 requires a multiple-of-32 last dimension, got {x.shape}")
+  *outer, k = x.shape
+  blocks = x.float().reshape(*outer, k//MX_BLOCK_SIZE, MX_BLOCK_SIZE)
+  amax = blocks.abs().max(axis=-1)
+  # Match the OCP/MLX reference: quantize amax / E2M1_MAX to the nearest E8M0
+  # power of two. This deliberately differs from extracting exponent bits: blocks
+  # whose maximum is near a power-of-two boundary can select a scale 2x smaller.
+  exponent = (amax.maximum(1e-38).div(6.0).log2().round()).clamp(-127, 127)
+  scale = (amax == 0).where(127, exponent + 127).cast(dtypes.uint8)
+  normalized = blocks / _e8m0_scale(scale).unsqueeze(-1)
+  # Midpoint bins avoid materializing a 16x larger distance tensor while preserving nearest-value encoding.
+  magnitude = normalized.abs()
+  # OCP formats use round-to-nearest-even: at alternating midpoints, the upper code has an even mantissa LSB.
+  code = sum(((magnitude >= midpoint) if upper_even else (magnitude > midpoint)).cast(dtypes.uint8)
+             for midpoint, upper_even in ((0.25, False), (0.75, True), (1.25, False), (1.75, True),
+                                           (2.5, False), (3.5, True), (5.0, False)))
+  code = ((normalized < 0) & (code != 0)).where(code + 8, code).reshape(*outer, k)
+  # Safetensors has no nibble dtype. Store the earlier element in the low nibble.
+  packed = code[..., ::2] + code[..., 1::2] * 16
+  return packed.contiguous(), scale.contiguous()
+
+def quantize_mxfp4_cpu(x:Tensor) -> tuple[Tensor, Tensor]:
+  """CPU converter fast path for large checkpoints. Inference itself does not depend on numpy."""
+  import numpy as np
+  if x.shape[-1] % MX_BLOCK_SIZE: raise ValueError(f"MXFP4 requires a multiple-of-32 last dimension, got {x.shape}")
+  array = x.float().numpy()
+  blocks = array.reshape(*array.shape[:-1], array.shape[-1]//MX_BLOCK_SIZE, MX_BLOCK_SIZE)
+  amax = np.max(np.abs(blocks), axis=-1)
+  exponent = np.clip(np.rint(np.log2(np.maximum(amax, 1e-38) / 6.0)), -127, 127)
+  scale = np.where(amax == 0, 127, exponent + 127).astype(np.uint8)
+  normalized = blocks / np.exp2(scale.astype(np.float32) - 127)[..., None]
+  magnitude = np.abs(normalized)
+  code = sum(((magnitude >= midpoint) if upper_even else (magnitude > midpoint)).astype(np.uint8)
+             for midpoint, upper_even in ((0.25, False), (0.75, True), (1.25, False), (1.75, True),
+                                           (2.5, False), (3.5, True), (5.0, False)))
+  code = np.where((normalized < 0) & (code != 0), code + 8, code).astype(np.uint8).reshape(array.shape)
+  packed = code[..., ::2] + code[..., 1::2] * 16
+  return Tensor(packed), Tensor(scale)
+
+def dequantize_mxfp4(packed:Tensor, scale:Tensor, dtype=dtypes.bfloat16) -> Tensor:
+  """Decode the packed representation emitted by quantize_mxfp4."""
+  if packed.shape[-1] != scale.shape[-1] * 16:
+    raise ValueError(f"incompatible MXFP4 values/scales: {packed.shape} and {scale.shape}")
+  lo = packed - packed.div(16, rounding_mode="trunc") * 16
+  hi = packed.div(16, rounding_mode="trunc")
+  code = Tensor.stack(lo, hi, dim=-1).reshape(*packed.shape[:-1], packed.shape[-1]*2)
+  values = Tensor(MXFP4_VALUES, dtype=dtypes.float32, device=packed.device)[code]
+  scales = _e8m0_scale(scale).unsqueeze(-1).expand(*scale.shape, MX_BLOCK_SIZE).reshape(*scale.shape[:-1], scale.shape[-1]*MX_BLOCK_SIZE)
+  return (values * scales).cast(dtype)
+
+def quantize_dequantize_mxfp8(x:Tensor, dtype=dtypes.bfloat16) -> Tensor:
+  """Apply the Kimi expert-activation MXFP8 E4M3/E8M0 round trip in 32-value blocks."""
+  if x.shape[-1] % MX_BLOCK_SIZE: raise ValueError(f"MXFP8 requires a multiple-of-32 last dimension, got {x.shape}")
+  *outer, k = x.shape
+  blocks = x.float().reshape(*outer, k//MX_BLOCK_SIZE, MX_BLOCK_SIZE)
+  amax = blocks.abs().max(axis=-1)
+  # As for MXFP4, the E8M0 scale is nearest-power-of-two(amax / E4M3_MAX).
+  exponent = (amax.maximum(1e-38).div(448.0).log2().round()).clamp(-127, 127)
+  scale = (amax == 0).where(127, exponent + 127).cast(dtypes.uint8)
+  normalized = blocks / _e8m0_scale(scale).unsqueeze(-1)
+  # Software OCP E4M3 rounding is required on gfx1100 (RDNA3 has no native FP8 dtype).
+  # E4M3 has three explicit mantissa bits and a minimum normal exponent of -6;
+  # using e=-6 also gives the 2**-9 subnormal quantum.
+  magnitude = normalized.abs().clamp(max_=448.0)
+  elem_exp = magnitude.maximum(2**-9).log2().floor().clamp(-6, 8)
+  quantum = (elem_exp - 3).exp2()
+  quantized = (magnitude / quantum).round() * quantum
+  quantized = (normalized < 0).where(-quantized, quantized).clamp(-448.0, 448.0)
+  return (quantized * _e8m0_scale(scale).unsqueeze(-1)).reshape(*outer, k).cast(dtype)
@@ -24,6 +24,11 @@ def parse_tool_call(s:str) -> tuple[str, typing.Any]|None:
    return fm.group(1), args
  return None

+def parse_kimi_tool_call(s:str) -> tuple[str, typing.Any]|None:
+  if (m := re.match(r"\s*(?:functions\.)?([^:\s]+)(?::[^\s]+)?\s*<\|tool_call_argument_begin\|>\s*(.*?)\s*\Z", s, re.DOTALL)) is None: return None
+  try: return m.group(1), json.loads(m.group(2))
+  except json.JSONDecodeError: return None
+
 def normalize_messages(messages:list[dict]) -> None:
  # chat templates expect tool_call arguments as dicts (OpenAI clients send JSON strings)
  for m in messages:
@@ -34,9 +39,10 @@ def normalize_messages(messages:list[dict]) -> None:

 class StreamRouter:
  # routes streamed output text to (field, text) deltas, keeping tool_call regions in .buf for the final parse
-  def __init__(self, reasoning:bool=False):
+  def __init__(self, reasoning:bool=False, xtml:bool=False):
    self.buf = ""
    self.mode = "reasoning" if reasoning else "undecided"  # output inside a think block is sent as reasoning_content
+    self.xtml = xtml
  def split(self, tag:str, final:bool) -> tuple[str, bool]:
    # split buf on the first full tag, holding back a partial tag at the end unless final
    if tag in self.buf:
@@ -45,20 +51,39 @@ class StreamRouter:
    hold = max((i for i in range(1, min(len(self.buf), len(tag))+1) if tag.startswith(self.buf[-i:])), default=0) if not final else 0
    emit, self.buf = self.buf[:len(self.buf)-hold], self.buf[len(self.buf)-hold:]
    return emit, False
+  def split_any(self, tags:tuple[str, ...], final:bool) -> tuple[str, str|None]:
+    found = [(self.buf.index(tag), tag) for tag in tags if tag in self.buf]
+    if found:
+      pos, tag = min(found)
+      before, self.buf = self.buf[:pos], self.buf[pos+len(tag):]
+      return before, tag
+    hold = max((i for tag in tags for i in range(1, min(len(self.buf), len(tag))+1) if tag.startswith(self.buf[-i:])), default=0) if not final else 0
+    emit, self.buf = self.buf[:len(self.buf)-hold], self.buf[len(self.buf)-hold:]
+    return emit, None
  def route(self, piece:str, final:bool=False) -> typing.Iterator[tuple[str, str]]:
    self.buf += piece
    if self.mode == "undecided":  # decide whether the output starts with a think block
      if not final and len(self.buf) < len("<think>") and "<think>".startswith(self.buf): return
      self.mode, self.buf = ("reasoning", self.buf[len("<think>"):]) if self.buf.startswith("<think>") else ("content", self.buf)
    if self.mode == "reasoning":
-      emit, done = self.split("</think>", final)
+      emit, done = self.split("<|close|>think<|sep|>" if self.xtml else "</think>", final)
      if emit: yield "reasoning_content", emit
      if not done: return
+      self.mode = "content_open" if self.xtml else "content"
+    if self.mode == "content_open":
+      _, found = self.split("<|open|>response<|sep|>", final)
+      if not found: return
      self.mode = "content"
+    if self.mode == "done": return
+    if self.xtml and self.mode == "content":
+      emit, found = self.split("<|close|>response<|sep|>", final)
+      if emit: yield "content", emit
+      if found: self.mode = "done"
+      return
    if self.mode == "tool": return
-    emit, found = self.split("<tool_call>", final)
+    emit, tool_tag = self.split_any(("<tool_call>", "<|tool_calls_section_begin|>"), final)
    if emit: yield "content", emit
-    if found: self.mode, self.buf = "tool", "<tool_call>" + self.buf
+    if tool_tag: self.mode, self.buf = "tool", tool_tag + self.buf

 class Handler(HTTPRequestHandler):
  server: LLMServer
@@ -67,7 +92,7 @@ class Handler(HTTPRequestHandler):
    if self.path == "/v1/models": self.send_data(json.dumps({"object":"list","data":[{"id":self.server.model_name,"object":"model"}]}).encode())
    else: self.send_data((pathlib.Path(__file__).parent / "chat.html").read_bytes(), content_type="text/html")
  def run_model(self, ids:list[int], model_name:str, include_usage=False, max_tokens:int|None=None, temperature:float=0.0,
-                reasoning:bool=False):
+                reasoning:bool=False, xtml:bool=False):
    model, tok = self.server.model, self.server.tok
    prompt_tokens = len(ids)
    cache_start_pos = model.get_start_pos(ids)
@@ -78,7 +103,7 @@ class Handler(HTTPRequestHandler):
    finish_reason = "stop"
    st = pt = time.perf_counter()
    dec = tok.stream_decoder()
-    router = StreamRouter(reasoning)
+    router = StreamRouter(reasoning, xtml)
    def log_stats(interrupted:bool=False):
      et = time.perf_counter()
      total = f"total:{et-st:6.2f}s"
@@ -106,6 +131,14 @@ class Handler(HTTPRequestHandler):
          name, args = parsed
          tool_calls.append({"index":len(tool_calls), "id":f"call_{uuid.uuid4().hex[:24]}", "type":"function",
                             "function":{"name":name, "arguments":args if isinstance(args, str) else json.dumps(args)}})
+      for m in re.finditer(r"<\|tool_call_begin\|>(.*?)<\|tool_call_end\|>", router.buf, re.DOTALL):
+        if (parsed := parse_kimi_tool_call(m.group(1))) is None:
+          stderr_log(f"failed to parse Kimi tool call: {m.group(1)[:200]}")
+          yield chunk({"content":m.group(0)})
+        else:
+          name, args = parsed
+          tool_calls.append({"index":len(tool_calls), "id":f"call_{uuid.uuid4().hex[:24]}", "type":"function",
+                             "function":{"name":name, "arguments":json.dumps(args)}})
      if tool_calls:
        yield chunk({"tool_calls":tool_calls})
        if finish_reason == "stop": finish_reason = "tool_calls"
@@ -139,9 +172,10 @@ class Handler(HTTPRequestHandler):

      # reply
      max_tokens = body.get("max_completion_tokens") or body.get("max_tokens")
+      xtml = rendered.rstrip().endswith("<|open|>think<|sep|>")
      chunks = self.run_model(ids, body["model"], not body.get("stream") or body.get("stream_options",{}).get("include_usage", False),
                              max_tokens=max_tokens, temperature=float(body.get("temperature", 0.0)),
-                              reasoning=rendered.rstrip().endswith("<think>"))
+                              reasoning=xtml or rendered.rstrip().endswith("<think>"), xtml=xtml)
      if body.get("stream"): self.stream_json(chunks)
      else:
        out, reasoning, tool_calls, finish_reason = [], [], [], "stop"
@@ -221,7 +221,7 @@ class ElementwiseMixin(CreationMixin):
    if dtypes.is_int(a.dtype) and dtypes.is_int(b.dtype): return a.alu(Ops.CMOD, b)
    return a - a.div(b, rounding_mode="trunc") * b

-  def div(self, x: Self | ConstType, reverse: bool = False, rounding_mode: Literal["trunc", "floor"] | None = None) -> Self:
+  def div(self, x: 'Self|ConstType|UOp', reverse: bool = False, rounding_mode: Literal["trunc", "floor"] | None = None) -> Self:
    """
    Divides `self` by `x`.
    Equivalent to `self / x`.
@@ -7,7 +7,11 @@ from tinygrad.dtype import sum_acc_dtype
 def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
  if op == Ops.ADD: return (ctx._broadcast_to(ret.src[0].shape),)
  if op == Ops.MAX: return (((mask:=ret.src[0].eq(ret).cast(ctx.dtype))/mask._rop(Ops.ADD, tuple(range(ret.arg[1])))) * ctx,)
-  if op == Ops.MUL: return (ctx * ret / ret.src[0],)
+  if op == Ops.MUL:
+    # d(prod x)/dx_j = prod_{i!=j} x_i: ret/x_j whenever x_j != 0 (any zero makes ret 0), else the product of the others
+    safe_x, axes = (is_zero:=(x:=ret.src[0]).eq(0)).where(1, x), tuple(range(ret.arg[1]))
+    zero_count = is_zero.cast(sum_acc_dtype(is_zero.dtype))._rop(Ops.ADD, axes)
+    return (ctx * is_zero.where(zero_count.eq(1).where(safe_x._rop(Ops.MUL, axes), 0), ret/safe_x),)

 def _compact_params(body:UOp, all_args:tuple[UOp, ...]) -> tuple[UOp, tuple[UOp, ...]]:
  """Remove unused PARAMs from body and return compacted (body, args)."""
@@ -90,8 +90,11 @@ class MovementMixin:
        if resolve(index.step == 0, False): raise ValueError(f"{index=} cannot have 0 as step")
        start, stop = 0 if index.start is None else index.start, size if index.stop is None else index.stop
        step = 1 if index.step is None else index.step
+        # resolve negative int bounds against the (possibly symbolic) size, like slice.indices
+        if isinstance(start, int) and start < 0: start = start + size
+        if isinstance(stop, int) and stop < 0: stop = stop + size
        if all_int((start, stop, step)):
-          # handle int slicing (resolve negative bounds, clamp, stride)
+          # handle int slicing (clamp, stride)
          *bound, stride = index.indices(int(size.vmax) if isinstance(size, UOp) else size)
          bound = [0, 0] if stride * (bound[1] - bound[0]) < 0 else ([bound[1]+1, bound[0]+1] if stride < 0 else bound)
          return {"size":ceildiv(bound[1]-bound[0], abs(stride)), "boundary":tuple(bound), "stride":stride, "collapse_dim":False}
@@ -265,7 +268,8 @@ class MovementMixin:
    return self.shrink(tuple([None if ns is None else (0, ns) for ns in argfix(shape, *args)]))

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

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

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

  def var_mean(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> tuple[Self, Self]:
    """
@@ -1057,14 +1056,16 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
    assert not (align_corners and mode != "linear"), "align_corners option can only be set with the interpolating mode linear"
    x, expand = self, list(self.shape)
    for i in range(-1,-len(size)-1,-1):
-      scale = (int(self.shape[i]) - int(align_corners)) / (size[i] - int(align_corners))
-      arr, reshape = type(self).arange(size[i], dtype=dtypes.float32), [1] * self.ndim
+      in_sz, reshape = int(self.shape[i]), [1] * self.ndim
      reshape[i] = expand[i] = size[i]
      if mode == "linear":
-        index = (scale*arr if align_corners else (scale*(arr+0.5))-0.5).clip(0, self.shape[i]-1)
-        low, high, perc = [y.reshape(reshape).expand(expand) for y in (index.floor().int(), index.ceil().int(), index - index.floor())]
+        arr = type(self).arange(size[i])
+        num, den = (arr*(in_sz-1), size[i]-1) if align_corners else ((arr*2+1)*in_sz - size[i], size[i]*2)
+        num = num.clip(0, (in_sz-1)*den)
+        low, high, perc = [y.reshape(reshape).expand(expand) for y in (num//den, (num+den-1)//den, (num % den).cast(dtypes.float32)/den)]
        x = x.gather(i, low).lerp(x.gather(i, high), perc)
      else:
+        scale, arr = in_sz / size[i], type(self).arange(size[i], dtype=dtypes.float32)
        index = (scale*(arr+0.5) if mode=="nearest-exact" else scale*arr).cast(dtypes.int32).reshape(reshape).expand(expand)
        x = x.gather(i, index)
    return x.cast(self.dtype)
@@ -35,21 +35,36 @@ def lcast(input_type:DType, output_type:DType):
    if dtypes.is_int(output_type): return 'trunc' if output_type.itemsize < input_type.itemsize else 'sext'
  raise NotImplementedError(f"cast from {input_type} -> {output_type} not implemented")

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

 # llvm ops, lop[<dtype>][<op>]
 unsigned_lop = { Ops.ADD: "add", Ops.MUL: "mul", Ops.CDIV: "udiv", Ops.CMOD: "urem",
@@ -254,13 +269,21 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n  %trunc
    attributes = ["alwaysinline", "nounwind", '"no-builtins"',
                  f'"amdgpu-flat-work-group-size"="1,{requiredMaxThreadsPerBlock}"', '"no-trapping-math"="true"']
    return 'attributes #0 = { ' + ' '.join(attributes) + ' }'
+  @staticmethod
+  def is_rdna4(arch): return arch.split(':')[0] in {'gfx1200', 'gfx1201'}
  def __init__(self, target:Target):
    super().__init__(target)
    from tinygrad.runtime.support.compiler_llvm import AMDLLVMCompiler
    self.compiler, self.tensor_cores, self.is_cdna = AMDLLVMCompiler(target.arch), tc.get_amd(target.arch), HIPRenderer.is_cdna(target.arch)
-    self.string_rewrite += PatternMatcher([(UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, cdna=self.is_cdna: render_wmma_amd(ctx, wmma, cdna))])
+    self.string_rewrite += PatternMatcher([
+      (UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, rdna4=AMDLLVMRenderer.is_rdna4(target.arch), cdna=self.is_cdna:
+        render_wmma_amd(ctx, wmma, cdna, rdna4))
+    ])
    if self.is_cdna:
      self.extra_matcher += PatternMatcher([
+        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
+          if x.arg[0][2] == 128 and x.src[0].dtype.itemsize <= 8 else None),
        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
          if x.max_numel() == 4 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 4 else None),
@@ -274,9 +297,10 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n  %trunc
          src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
          if x.src[0].dtype == dtypes.int8 and x.src[0].max_numel() == 16 else None),
        (UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(x.replace(
-          src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(j//2) if j%2 == 0 else UOp.const(0.0, x.src[2].dtype)
+          src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(UOp.const(j//2, dtypes.int16))
+            if j%2 == 0 else UOp.const(0.0, x.src[2].dtype)
            for j in range(x.max_numel()*2)))),
-          arg=(*x.arg[:4], None)).index(i*2)
+          arg=(*x.arg[:4], None)).index(UOp.const(i*2, dtypes.int16))
          for i in range(x.max_numel()))) if x.max_numel() == 8 else None),
        (UPat(Ops.WMMA, name="x"), lambda x: x.replace(
          src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
@@ -285,6 +309,7 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n  %trunc
    if target.arch in {"gfx1200", "gfx1201"}:
      self.extra_matcher += PatternMatcher([
        (UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16), lambda x: x.replace(
+          dtype=dtypes.uint16,
          src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2].bitcast(dtypes.uint16)))
            .bitcast(dtypes.bfloat16) if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None),
        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
@@ -50,7 +50,7 @@ def worker_prog():

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

  entry = [ring.after(ready).index((cur % RING_SLOTS) * CMD_SIZE + i).load() for i in range(CMD_SIZE)]
  return entry[0].call(*entry[1:], ret_dtype=dtypes.void).end(cur)
@@ -167,10 +167,11 @@ class CPUDevice(HCQCompiled):
    (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx[0].signal("sentinel", (1 << 64) - 1)),
    (UPat(Ops.PARAM, tag="timeline_signal"), lambda ctx: ctx[0].signal("timeline")),
    (UPat(Ops.PARAM, tag="timeline_value"), lambda ctx: ctx[0].signal("value", 1)),
+    (UPat(Ops.PARAM, tag="signal", name="b"), lambda ctx, b: ctx[0].signal(b.arg.slot)),
  ])

  @functools.cache
-  def signal(self, name:str, init_value:int=0) -> Buffer:
+  def signal(self, name:str|int, init_value:int=0) -> Buffer:
    (buf:=Buffer(self.device, 1, dtypes.uint64, preallocate=True)).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = init_value
    return buf

@@ -89,9 +89,11 @@ class HIPCompiler(Compiler):

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

@dataclass(frozen=True)
 class HCQInfo:
-  name:str
-  estimates:Estimates
  device:tuple[str, ...]
-  queue:str
+  estimates:Estimates = Estimates()

  input_idxs:tuple[int, ...] = () # indexes into input_uops used by this call
  inputs:int|None = None
+  prof:tuple[ProfileGraphEntry, ...] = () # st_id/en_id are timestamp signal slots until collect

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

@@ -73,6 +73,8 @@ def make_submit(*cmds, devs:str|tuple[str, ...], queue:str) -> UOp:
  return UOp.custom_function("submit_cmdbuf", UOp(Ops.LINEAR, src=tuple(cmds), arg=(to_tuple(devs), queue)))
 def get_submit(ast:UOp) -> UOp: return next(u for u in ast.toposort() if u.op is Ops.CUSTOM_FUNCTION and u.arg == "submit_cmdbuf")

+def make_call(name:str, body:UOp, info:HCQInfo) -> UOp: return UOp.custom_function("hcq", body).call(name=name, aux=info)
+
 def encode_kernargs_clike(call:UOp, prg:UOp, devs:str|tuple[str, ...]) -> UOp:
  data, info = prg.arg
  buf = UOp.placeholder((data.kernargs_alloc_size // 4,), dtypes.uint32, next(UOp.unique_num), device=devs).rtag("kernargs")
@@ -137,12 +139,6 @@ def _build_wait_cmds(slots:dict[str, int], dep_lanes:list[tuple[tuple, int, int]
    waits.append(UOp(Ops.INS, arg="wait", src=(sig, UOp.const(dtag + 1, dtypes.uint64))))
  return waits, {dtag for _, _, dtag in deps}

-def make_fence(timeline:UOp, prev:UOp, sigs:list[UOp]) -> UOp:
-  free = (cur:=timeline.after(loop:=UOp.loop(0)).index(0).load()).end(loop, cur < prev.index(0).load())
-  return UOp.sink(*[s.after(free).index(0).store(0) for s in sigs])
-
-def _hcq_call(devs, name:str, body:UOp) -> UOp: return UOp.custom_function("hcq", body).call(aux=HCQInfo(name, Estimates(), devs, "COMPUTE:0"))
-
 def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[tuple[tuple[str, ...], str]],
                      tracker:HCQDepsTracker, slots:dict[str, int]) -> tuple[list[UOp], list[UOp], set[int]]:
  # collect all buffers which belong to devices
@@ -151,48 +147,50 @@ def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[t
    for b in itertools.chain.from_iterable(_get_call_bufs_by_lane(call, devices)):
      for bd in to_tuple(b.device): dev_bufs[bd][id(b)] = b

-  n, fences, fins, waited = len(batch_info), [], [], set()
+  n, fences, fins, signal_tags = len(batch_info), [], [], set()
  for _, devgroup in itertools.groupby(sorted(dev_bufs), key=lambda d: d.split(":")[0]):
    devs = tuple(devgroup)

    # to finalize the batch, sync all accesses from other devices to buffers that belong to this device
    fin_deps = [dl for dl in _get_deps(tracker, [list(dev_bufs[d].values()) for d in devs], None, key=(devs, "COMPUTE:0", n)) if dl[0][2] < n]
-    waits, cur_waited = _build_wait_cmds(slots, fin_deps, devs, "COMPUTE:0")
-    waited |= cur_waited
+    waits, cur_signal_tags = _build_wait_cmds(slots, fin_deps, devs, "COMPUTE:0")
+    signal_tags |= cur_signal_tags

    # wait the syncs and signal the device epoch, then bump the timeline on the host
-    timeline, tl = make_signal(devs, tag="timeline_signal"), make_signal(devs, tag="timeline_value")
-    submit = make_submit(*waits, UOp(Ops.INS, arg="store", src=(timeline, tl.index(0))), devs=devs, queue="COMPUTE:0")
-    cur = (bump:=tl.after(submit).index(0)).load()
-    bumps = [bump.store(cur + 1)]
+    tl_signal, tl_value = make_signal(devs, tag="timeline_signal"), make_signal(devs, tag="timeline_value")
+    fin_submit = make_submit(*waits, UOp(Ops.INS, arg="store", src=(tl_signal, tl_value.index(0))), devs=devs, queue="COMPUTE:0")
+    epoch = (epoch_slot:=tl_value.after(fin_submit).index(0)).load()

-    # devices running the batch reset their queue signals before each run, fencing on the epoch kept from the previous one
-    if qs:=dedup([qn for bdevs, qn in batch_info if set(bdevs) & set(devs)]):
-      prev = make_signal(devs, next(UOp.unique_num))
-      fences.append(_hcq_call(devs, "hcq_fence", make_fence(timeline, prev, [make_signal(devs, slots[q]) for q in qs])))
-      bumps.append(prev.after(submit).index(0).store(cur))
-    fins.append(_hcq_call(devs, "hcq_finalizer", UOp.sink(*bumps)))
-  return fences, fins, waited
+    # fence once per device group on this schedule's previous epoch, then reset any queue signals used by the group
+    qs = dedup([qn for bdevs, qn in batch_info if set(bdevs) & set(devs)])
+    sched_epoch = make_signal(devs, next(UOp.unique_num))

-def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]]) -> list[UOp]:
+    wait_device_epoch = (done:=tl_signal.after(loop:=UOp.loop(0)).index(0).load()).end(loop, done < sched_epoch.index(0).load())
+    resets = [make_signal(devs, slots[q]).after(wait_device_epoch).index(0).store(0) for q in qs]
+
+    fences.append(make_call("hcq_fence", UOp.sink(*(resets or [wait_device_epoch])), HCQInfo(devs)))
+    fins.append(make_call("hcq_finalizer", UOp.sink(epoch_slot.store(epoch + 1), sched_epoch.after(fin_submit).index(0).store(epoch)), HCQInfo(devs)))
+  return fences, fins, signal_tags
+
+def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> list[UOp]:
  batch_info = [(devices, "COMPUTE:0" if call.src[0].op is Ops.PROGRAM else "COPY:0") for call, devices in batch]

  # schedule deps
-  waited:set[int] = set()
+  signal_tags:set[int] = set()
  slots:dict[str, int] = collections.defaultdict(lambda: next(UOp.unique_num))
  deps_tracker = HCQDepsTracker()
  call_waits:list[list[UOp]] = []
  for tag, ((call, _), (devices, queue)) in enumerate(zip(batch, batch_info)):
    deps = _get_deps(deps_tracker, _get_call_bufs_by_lane(call, devices), get_call_outs_ins(call)[0], key=(devices, queue, tag))
-    cmds, cur_waited = _build_wait_cmds(slots, deps, devices, queue)
+    cmds, cur_signal_tags = _build_wait_cmds(slots, deps, devices, queue)
    call_waits.append(cmds)
-    waited |= cur_waited
+    signal_tags |= cur_signal_tags

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

-  src = []
+  src, prof = [], []
  for tag, ((call, _), (devices, queue), q) in enumerate(zip(batch, batch_info, call_waits)):
    # first queue use, sync prior device work with the device timeline
    if batch_info.index((devices, queue)) == tag:
@@ -200,31 +198,38 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]]) -> list[UOp]:
      q = [UOp(Ops.INS, arg="barrier", src=()), UOp(Ops.INS, arg="wait", src=(make_signal(devices, tag="timeline_signal"), epoch))] + q

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

    # signal the queue if someone waits for us
-    if tag in waited: q += [UOp(Ops.INS, arg="store", src=(make_signal(devices, slots[queue]), UOp.const(tag + 1, dtypes.uint64)))]
-    src.append(UOp.custom_function("hcq", make_submit(*q, devs=devices, queue=queue).sink()).call(name="hcq", aux=info))
+    if tag in signal_tags: q += [UOp(Ops.INS, arg="store", src=(make_signal(devices, slots[queue]), UOp.const(tag + 1, dtypes.uint64)))]
+    src.append(make_call(name, make_submit(*q, devs=devices, queue=queue).sink(), info))
+
+  # append batch timestamps to finalizers
+  finalizers = [f.replace(arg=replace(f.arg, aux=replace(a:=f.arg.aux, prof=tuple(e for e in prof if e.device in a.device)))) for f in finalizers]
  return fences + src + finalizers

-def sched_hcq_batches(l:UOp) -> UOp:
+def sched_hcq_batches(l:UOp, profile:bool) -> UOp:
  srcs:list[UOp] = []
  batch:list[tuple[UOp, tuple[str, ...]]] = []
  for call in l.src:
    if (devs:=next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None)) is not None: batch.append((call, to_tuple(devs)))
-    else: srcs, batch = srcs + _finalize_batch(batch) + [call], []
-  return l.replace(src=tuple(srcs + _finalize_batch(batch)))
+    else: srcs, batch = srcs + _finalize_batch(batch, profile) + [call], []
+  return l.replace(src=tuple(srcs + _finalize_batch(batch, profile)))

 # *****************
 # 3. merge into queues

 def _merged_hcq_call(calls:list[UOp]) -> UOp: # TODO: simplify?
  if len(calls) == 1: return calls[0]
-  info = replace(calls[0].arg.aux, name=f"submit {calls[0].arg.aux.queue} ({len(calls)})",
-                 estimates=sum((c.arg.aux.estimates for c in calls), start=Estimates()))
-  cmds = [cmd for c in calls for cmd in get_submit(c).src[0].src]
-  return UOp.custom_function("hcq", make_submit(*cmds, devs=info.device, queue=info.queue).sink()).call(name="hcq", aux=info)
+  devs, queue = get_submit(calls[0]).src[0].arg
+  body = make_submit(*[cmd for c in calls for cmd in get_submit(c).src[0].src], devs=devs, queue=queue).sink()
+  return make_call(f"submit {queue} ({len(calls)})", body,
+    replace(calls[0].arg.aux, estimates=sum((c.arg.aux.estimates for c in calls), start=Estimates())))

 def merge_queues(linear:UOp) -> UOp:
  new_src:list[UOp] = []
@@ -232,24 +237,25 @@ def merge_queues(linear:UOp) -> UOp:
  limits:dict[tuple[tuple[str, ...], str], int] = collections.defaultdict(lambda: JIT_BATCH_SIZE.value)

  for call in linear.src:
-    if not isinstance(info:=call.arg.aux, HCQInfo) or info.name.startswith("hcq_"): # non-hcq call, fence or finalizer: close all open queues
+    # non-hcq call, fence or finalizer: close all open queues
+    if not isinstance(call.arg.aux, HCQInfo) or (call.arg.name or "").startswith("hcq_"):
      new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in list(opened_qs)] + [call]
      continue

-    if (old:=opened_qs.pop(key:=(info.device, info.queue), None)) is not None:
+    devs, queue = get_submit(call).src[0].arg
+    if (old:=opened_qs.pop(key:=(devs, queue), None)) is not None:
      if limits[key] and len(old) >= limits[key]: new_src, old, limits[key] = new_src + [_merged_hcq_call(old)], [], limits[key] * 2
      new_rec = old + [call]
    else:
      # no such queue opened: close every open submit on this queue that shares a device, so submit order is kept
-      closing = [k for k in opened_qs if k[1] == info.queue and set(k[0]) & set(info.device)]
+      closing = [k for k in opened_qs if k[1] == queue and set(k[0]) & set(devs)]
      new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in closing]
      new_rec = [call]
-    opened_qs[(info.device, info.queue)] = new_rec
+    opened_qs[(devs, queue)] = new_rec
  return linear.replace(src=tuple(new_src + [_merged_hcq_call(c) for c in opened_qs.values()]))

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

 # *****************
 # 4.2. hcq lowering: ops to ir
@@ -285,21 +291,26 @@ def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[UOp, dict[UOp
  fills = (table.after(*make_patches(table, [(i*table.dtype.itemsize, addr) for addr, i in slots.items()])),) if slots else ()
  return table, reads, fills, {g:slots[bare[g]] for g in gaddrs}

-def make_scatter_loop(patches:list[UOp], inputs_table:tuple, lt_patches:list[UOp]) -> dict[UOp, UOp]:
-  (table, _, _, slots), dst, data, subs = inputs_table, patches[0].buf_uop, [], {}
-  for p in patches:
-    words = [(off, val, get_getaddrs(val)) for off,val in zip(p.src[0].src[1].src, p.src[1].src)]
-    data += [off.val << 32 | slots[gaddrs[0]] for off,_,gaddrs in words if gaddrs][::2]
-    scalars = [(off.val*dst.dtype.itemsize, val) for off,val,gaddrs in words if not gaddrs]
-    subs[p] = UOp.group(*make_patches(dst, scalars)) if scalars else UOp(Ops.NOOP)
+def is_bare_addr(val:UOp) -> bool: return val.op is Ops.CAST and val.src[0].op in (Ops.AND, Ops.SHR) and val.src[0].src[0].op is Ops.GETADDR

-  # plan entry: dst word offset << 32 | addr table slot
-  plan = UOp.placeholder((len(data),), dtypes.uint64, next(UOp.unique_num), device=dst.device).rtag("systems")
-  entry = plan.index(ridx:=UOp.range(len(data), next(UOp.unique_num), dtype=dtypes.int, src=(plan, dst))).load()
-  slot, widx = ((entry & 0xffffffff) % table.max_numel()).cast(dtypes.int), ((entry >> 32) % (dst.max_numel()-1)).cast(dtypes.int) # CHECK_OOB bounds
-  loop = UOp.group(*[dst.index(widx+i).store((table.index(slot).load() >> 32*i).cast(dtypes.uint32)) for i in range(2)]).end(ridx)
-  lt_patches.append(make_binary_patch(plan, struct.pack(f'<{len(data)}Q', *data)))
-  subs[patches[0]] = UOp.group(loop, subs[patches[0]])
+def make_scatter_loops(patches:list[UOp], inputs_table:tuple, lt_patches:list[UOp]) -> dict[UOp, UOp]:
+  table, _, _, slots = inputs_table
+  subs, by_dst = {}, collections.defaultdict(list)
+  for p in patches: by_dst[p.buf_uop].append(p)
+  for dst, patches in by_dst.items():
+    data = []
+    for p in patches:
+      words = [(off, val, get_getaddrs(val)) for off,val in zip(p.src[0].src[1].src, p.src[1].src)]
+      data += [(off.val, slots[gaddrs[0]]) for off,_,gaddrs in words if gaddrs][::2]
+      scalars = [(off.val*dst.dtype.itemsize, val) for off,val,gaddrs in words if not gaddrs]
+      subs[p] = UOp.group(*make_patches(dst, scalars)) if scalars else UOp(Ops.NOOP)
+
+    word_table, slot_table = (UOp.placeholder((len(data),), dtypes.uint32, next(UOp.unique_num), device=dst.device).rtag("systems") for _ in range(2))
+    ridx = UOp.range(len(data), next(UOp.unique_num), dtype=dtypes.int, src=(word_table, slot_table, dst))
+    widx, slot = ((p.index(ridx).load() % bound).cast(dtypes.int) for p,bound in ((word_table, dst.max_numel()-1), (slot_table, table.max_numel())))
+    loop = UOp.group(*[dst.index(widx+i).store((table.index(slot).load() >> 32*i).cast(dtypes.uint32)) for i in range(2)]).end(ridx)
+    lt_patches += [make_binary_patch(buf, struct.pack(f'<{len(data)}I', *vals)) for buf,vals in zip((word_table, slot_table), zip(*data))]
+    subs[patches[0]] = UOp.group(loop, subs[patches[0]])
  return subs

 def is_input_addr(g:UOp) -> bool: return all(x.op is Ops.PARAM and x.tag is None for x in unwrap_mstack(g.buf_uop))
@@ -307,17 +318,22 @@ def is_input_addr(g:UOp) -> bool: return all(x.op is Ops.PARAM and x.tag is None
 def split_patches(call:UOp) -> UOp|None:
  rt_patches:list[UOp] = []
  lt_patches:list[UOp] = []
-  body = graph_rewrite(call.src[0], pm_trim_link_patches, ctx=(rt_patches, lt_patches), name=f"trim link-time patches ({call.arg.aux.name})")
+  body = graph_rewrite(call.src[0], pm_trim_link_patches, ctx=(rt_patches, lt_patches), name=f"trim link-time patches ({call.arg.name})")

  # split patches
  inputs, internals = partition(dedup(g for p in rt_patches for g in get_getaddrs(p)), is_input_addr)
  runtimes, systems = partition(internals, lambda g: any(x.tag in {"program", "kernargs", "cmdbuf"} for x in unwrap_mstack(g.buf_uop)))
  tables = [make_addr_table(call, gs, n) for gs,n in ((inputs, "inputs"), (runtimes, "runtime"), (systems, "systems"))]
  reads, fills = {k:v for _,r,_,_ in tables for k,v in r.items()}, [f for t in tables[1:] for f in t[2]] # inputs table is filled by exec
-  input_patches = [p for p in rt_patches if (gs:=get_getaddrs(p)) and all(map(is_input_addr, gs))]
-  scatter = make_scatter_loop(input_patches, tables[0], lt_patches) if input_patches else {}
+  input_patches = [p for p in rt_patches if (gs:=get_getaddrs(p)) and all(map(is_input_addr, gs))
+    and all(is_bare_addr(v) for v in p.src[1].src if get_getaddrs(v))]
+  scatter = make_scatter_loops(input_patches, tables[0], lt_patches)
  body = body.substitute({p:p.substitute(scatter | reads) for p in rt_patches})

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

  sub = {(b:=u.without_after): UOp.param(i, u.dtype, shape=b.shape, device=HCQ_RUNTIME_DEV.value, volatile=b.op is Ops.PARAM and b.arg.volatile)
         for i,u in enumerate(c_args)} | {v: v.replace(arg=replace(v.arg, slot=-1)) for v in variables if v.op is Ops.PARAM}
-  info = replace(call.arg.aux, inputs=next((i for i,u in enumerate(c_args) if u.tag == "inputs"), None))
+  info = replace(call.arg.aux, inputs=next((i for i,u in enumerate(c_args) if u.without_after.tag == "inputs"), None))
  return call.replace(src=(body.substitute(sub).replace(arg="hcq_args"), *c_args, *refhold),
                      arg=replace(call.arg, aux=info)) # TODO: call.after(*refhold)?
 pm_replace_params = PatternMatcher([
@@ -388,27 +404,27 @@ def callify_hcq(call:UOp, cf:UOp) -> UOp:
 pm_callify_hcq = PatternMatcher([(UPat(Ops.CALL, src=(
  UPat(Ops.CUSTOM_FUNCTION, arg="hcq_args", src=(UPat(Ops.SINK),), name="cf"),), name="call", allow_any_len=True), callify_hcq)])

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

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

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

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

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

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

    # and compile it
@@ -474,14 +490,14 @@ def link_buf_key(a:UOp): return a.key, to_tuple(a.device)
 link_buf_cache:dict[tuple[bytes, tuple[str, ...]], UOp] = {}
 link_linear_cache:dict[bytes, UOp] = {}

-@track_rewrites(lambda _,cache,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}")
+@rewrite_group(lambda _,cache,ret: f"HCQ Link {pluralize('Kernel', len(ret.src))}")
 def hcq_link(linear:UOp, cache=True) -> UOp:
  if (linked:=link_linear_cache.get(linear_key:=linear.key)) is not None: return linked

  bufs = {(j,i):a for j,c in enumerate(linear.src) for i,a in enumerate(c.src[1:], 1)
          if a.op is Ops.AFTER and unwrap_mstack(a.src[0])[0].tag in HCQ_CACHE_TAGS}
  linear = linear.substitute({x:link_buf_cache[k] for a in bufs.values() if (k:=link_buf_key(a)) in link_buf_cache for x in (a, a.src[0])}, walk=True)
-  linear = graph_rewrite(linear, pm_resolve_patches+symbolic+pm_assert_no_afters, bpm=pm_bufferize, ctx=cache, bottom_up=False,
+  linear = graph_rewrite(linear, pm_resolve_patches+symbolic+pm_fold_cast_const+pm_assert_no_afters, bpm=pm_bufferize, ctx=cache, bottom_up=False,
                         name="resolve patches")
  for (j,i),a in bufs.items(): link_buf_cache.setdefault(link_buf_key(a), linear.src[j].src[i])
  if cache: link_linear_cache[linear_key] = linear
@@ -495,6 +511,7 @@ class HCQ2Compiled(Compiled):

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

    self.pm_bufferize = PatternMatcher([
      (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx[0].signal("sentinel", (1 << 64) - 1)),
@@ -508,6 +525,27 @@ class HCQ2Compiled(Compiled):

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

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

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

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

@@ -1,6 +1,6 @@
 import time, inspect
 from collections import deque
-from tinygrad.uop.ops import UOp, Ops, UOpMetaClass, track_rewrites, graph_rewrite, gate_kernel_sink, KernelInfo
+from tinygrad.uop.ops import UOp, Ops, UOpMetaClass, rewrite_group, graph_rewrite, gate_kernel_sink, KernelInfo
 from tinygrad.uop.spec import type_verify, spec_tensor
 from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR, partition, dedup

@@ -98,10 +98,14 @@ pm_post_sched_cache = PatternMatcher([
   create_new_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None),
 ])

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

 schedule_cache: dict[bytes, UOp] = {}
@@ -167,7 +171,7 @@ pm_copy_from_store = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.SINK, name="ast"),), allow_any_len=True), assert_all_same_devices),
 ])

-@track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[0].src))}")
+@rewrite_group(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[0].src))}")
 def create_linear_with_vars(big_sink:UOp) -> tuple[UOp, dict[str, int]]:
  # big_sink srcs are all the Tensors
  linear_call = graph_rewrite(big_sink, pm_schedule, name="schedule to linear", enter_calls=True)
@@ -15,14 +15,13 @@ def handle_allreduce(buf:UOp, red:UOp) -> UOp|None:
  use_ring = concrete and not use_all2all and (RING >= 2 or (ndev > 2 and numel > getenv("RING_ALLREDUCE_THRESHOLD", 256_000) and RING >= 1))
  if DEBUG >= 2: print(f"{'ALL2ALL' if use_all2all else 'RING' if use_ring else 'NAIVE'} ALLREDUCE {ndev}x{numel} | {buf.dtype}")

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

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

  # chunk data into ndev pieces
  assert isinstance(numel, int)
@@ -2,30 +2,55 @@ from typing import Iterator
 import functools, itertools
 from dataclasses import dataclass, field, replace
 from tinygrad.dtype import dtypes, AddrSpace
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches, broadcast_axes
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, rewrite_group, broadcast_axes
 from tinygrad.uop.ops import gate_kernel_sink
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
 from tinygrad.helpers import argsort, all_same, cpu_profile, PCONTIG, colored, Context, SPEC

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

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

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

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

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

-@dataclass
-class IndexingContext:
-  realize_map: dict[UOp, None|list[int]] = field(default_factory=dict)
-  range_map: dict[UOp, tuple[tuple[UOp, ...], tuple[UOp, ...]]] = field(default_factory=dict)
-  # loads reachable from each UOp memoized across matches
-  buf_cache: dict[UOp, frozenset[UOp]] = field(default_factory=dict)
-
-  # create ranges
-  range_idx: Iterator[int] = field(default_factory=itertools.count)
-  def new_range(self, s:sint, axistype:AxisType=AxisType.WEAK) -> UOp:
-    if isinstance(s, UOp) and s.op is Ops.RANGE: return s
-    # if a range has a 1 src, it's the same as UOp.const(0)
-    return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(0)
-
 def broadcast_rngs(x:UOp, src:UOp, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]:
  if x.op not in GroupOp.Broadcastable: return rngs
  baxes, nleft = broadcast_axes(src.shape, x.shape), len(x.shape)-len(src.shape)
@@ -84,7 +95,7 @@ def create_bufferize_and_index_srcs(ctx:IndexingContext, x:UOp) -> list[UOp]:
        new_src = s.end(*[r for r in closed_ranges if r.op is Ops.RANGE])
        del ctx.realize_map[s]
      else:
-        removable = s.op not in ALWAYS_CONTIGUOUS
+        removable = s.op not in ALWAYS_CONTIGUOUS and s not in ctx.non_removable
        # LOCAL: None in the device assigns it a number later
        opts = BufferizeOpts(device=s.device, removable=removable) if len(ctx.range_map[s][1]) == len(realized_ranges) else \
               BufferizeOpts(device=s.device, addrspace=AddrSpace.LOCAL, removable=removable)
@@ -105,6 +116,7 @@ def convert_pad_to_where_to_keep_behavior_local(ctx:IndexingContext, x:UOp):

 def convert_reduce_to_reduce_with_ranges(ctx:IndexingContext, x:UOp):
  if x.arg[1] == 0: return None
+  if x not in ctx.range_map: raise RuntimeError("REDUCE has no ranges in rangeify, UOp verification failed")
  bx = create_bufferize_and_index_based_on_ranges(ctx, x)
  # input ranges
  new_ranges = list(ctx.range_map[x][0][:x.arg[1]])
@@ -176,13 +188,13 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO
    case _: raise RuntimeError(f"{op} is not a MovementOp")
  return rngs

-@profile_matches
+@rewrite_group(new_ctx=False)
 def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
  if debug: print("**************************")
  rctx = IndexingContext()

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

  # get the consumer map
  with cpu_profile("consumer map in rangeify", "TINY"):
@@ -126,7 +126,7 @@ def reshape_multi(root:UOp, multi:UOp):
  new_shardings = []
  for ax, rng in multi.sharding:
    count = int(rng.vmax)+1
-    target = prod(multi.shape[:ax])
+    target = ssimplify(prod(multi.shape[:ax]))
    if target not in arg_acc: raise RuntimeError(f"reshape {multi.shape} -> {new_shape} moved items between shards")
    new_ax = len(arg_acc) - arg_acc[::-1].index(target) - 1
    if new_shape[new_ax] % count != 0: raise RuntimeError(f"reshape {multi.shape} -> {new_shape} moved items between shards")
@@ -3,10 +3,10 @@ from typing import cast
 import itertools
 from tinygrad.dtype import dtypes, AddrSpace, Invalid, to_dtype, strong_dtype
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg, shape_to_shape_arg
-from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches, identity_element
-from tinygrad.uop.symbolic import symbolic
+from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, rewrite_group, identity_element
+from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const
 from tinygrad.uop.movement import mop_cleanup
-from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS
+from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS, SPEC
 from tinygrad.helpers import PCONTIG, FLOAT16, OPENPILOT_HACKS, argsort, partition, get_single_element
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
@@ -141,7 +141,7 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),

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

  # ** copy rules **

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

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

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

+def no_indexing_calls(u:UOp):
+  new_srcs = []
+  for x in u.src:
+    if x.op is Ops.INDEX:
+      # sometimes if call srcs have children the call will get an INDEX. we remove it here.
+      # TODO: we should add safety checks here for contiguous
+      new_srcs.append(x.src[0])
+    elif x.op is Ops.SHRINK:
+      # SHRINK with offset 0 is fine
+      # TODO: check offset
+      new_srcs.append(x.src[0])
+    else:
+      # everything else we pass through
+      new_srcs.append(x)
+  return u.replace(src=tuple(new_srcs))
+
+pm_no_indexing_calls = PatternMatcher([
+  (UPat(Ops.CALL, name="u"), no_indexing_calls),
+])
+
 DEVICE_MAX_BUFS = {"METAL": 31, "WEBGPU": 8, "CPU": 31} # TODO: get from device?
 def limit_bufs(ctx:IndexingContext, root:UOp):
  if (device:=root.device) is None: return None # no device, index related calculations
@@ -551,7 +572,7 @@ pm_copy_to_store = PatternMatcher([
  (UPat(Ops.COPY, name="copy"), convert_copy_to_store),
 ])

-@profile_matches
+@rewrite_group(new_ctx=False)
 def get_kernel_graph(sink:UOp) -> UOp:
  tsink = graph_rewrite(sink, multi_pm, name="multi_pm")
  if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters")
@@ -562,7 +583,9 @@ def get_kernel_graph(sink:UOp) -> UOp:
  # convert movement ops to ranges
  tsink, rctx = run_rangeify(tsink, bool(DEBUG_RANGEIFY))

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

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

  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
+  if SPEC:
+    # validate the kernel graph
+    from tinygrad.uop.spec import type_verify, spec_kernel_graph
+    type_verify(tsink, spec_kernel_graph, enter_calls=False)
  return tsink
@@ -1,17 +1,242 @@
 # inspired by https://github.com/karpathy/micrograd/blob/master/micrograd/engine.py
 from __future__ import annotations
 import time, functools, sys, inspect, pathlib, hashlib, weakref
+from dataclasses import dataclass, field
 from typing import Any, Callable, cast, get_args, ParamSpec, TypeGuard, TypeVar, Generic, TYPE_CHECKING
 if TYPE_CHECKING: import numpy
-from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, least_upper_dtype, to_dtype, strong_dtype, _from_np_dtype, _to_np_dtype, PyConst
+from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, least_upper_dtype, to_dtype, strong_dtype, \
+  _from_np_dtype, _to_np_dtype, PyConst, AddrSpace
 from tinygrad.helpers import all_int, getenv, fetch, Metadata, TRACEMETA, TracingKey
-from tinygrad.helpers import cpu_profile, suppress_finalizing, disable_gc
-from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, Variable, ConstLike
+from tinygrad.helpers import cpu_profile, suppress_finalizing, disable_gc, VIZ, pluralize
+from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, Variable, ConstLike, UPat, PatternMatcher, GroupOp, ParamArg, graph_rewrite, rewrite_group
 from tinygrad.mixin.rand import RandMixin
 from tinygrad.schedule import create_linear_with_vars
 from tinygrad.device import Buffer, canonicalize_device
 from tinygrad.engine.realize import run_linear
-from tinygrad.callify import transform_to_call
+
+# *** callify: transform a tensor graph into a CALL UOp such that all state is properly scoped ***
+
+@dataclass
+class AllocCtx:
+  uop_list: list[UOp] = field(default_factory=list)
+  buffer_map: dict[UOp, UOp] = field(default_factory=dict)
+  bases: set[UOp] = field(default_factory=set)
+  assigns: list[UOp] = field(default_factory=list)
+  replacements: list[UOp] = field(default_factory=list)
+
+def tag_uop(ctx:AllocCtx, x:UOp):
+  if x.tag is not None: return None
+  ctx.uop_list.append(x)
+  return x.replace(tag=(len(ctx.uop_list)-1,))
+
+def disk_like(u:UOp): return isinstance(u.device, str) and u.device.startswith(("DISK", "TINYFS"))
+
+def disk_copy_is_buffer(ctx:AllocCtx, u:UOp):
+  # copies to disk are replaced with the disk buffer
+  if disk_like(u) and u.tag is None:
+    ctx.buffer_map[u] = u.empty_like()
+    return u.rtag(())
+  # all copies from disk/numpy are realized into a real buffer
+  from_creation = isinstance(u.src[0].device, str) and u.src[0].device.startswith(("NPY", "DISK", "PYTHON", "TINYFS"))
+  if from_creation: return tag_uop(ctx, u)
+
+# CONTIGUOUS and AFTER + parents are the only nodes that get updated
+add_tags = PatternMatcher([
+  (UPat(Ops.COPY, name="u"), disk_copy_is_buffer),
+  # no tag on copies that are assigned via STORE+AFTER — merge COPY tag into AFTER
+  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(name="dest"), UPat(Ops.COPY, name="c")))), name="a"),
+   lambda a,c,dest: a.replace(src=(a.src[0], a.src[1].replace(src=(dest, c.rtag(())))), tag=a.tag+c.tag) if a.tag and c.tag else None),
+  (UPat((Ops.CONTIGUOUS, Ops.AFTER), name="x"), tag_uop),
+  (UPat(GroupOp.All, name="x"), lambda ctx,x: tag_uop(ctx,x) if x in ctx.bases else None),
+])
+
+def replace_contig_with_store_after(u:UOp):
+  # can't allocate a buffer for a virtual value
+  if u.is_virtual: return None
+  # if size is 0, remove the contig
+  if 0 in u.shape: return u.src[0]
+  # no real contig for DISK/TINYFS tensors, they are left alone
+  if disk_like(u): return u.rtag(None)
+  buf = u.empty_like()
+  return buf.after(buf.store(u.src[0])).rtag(u.tag)
+
+def replace_store_after_with_contig(u:UOp, src:UOp):
+  assigned_to = u
+  while assigned_to.op in {Ops.BITCAST, Ops.AFTER, Ops.UNSHARD}: assigned_to = assigned_to.src[0].base
+  if assigned_to.op not in {Ops.BUFFER, Ops.SLICE}: return src.contiguous(tag=u.tag)
+
+def _make_buffer_view(src:UOp) -> UOp|None:
+  """If movement ops on src collapse to a contiguous range, return SLICE. Otherwise None."""
+  if (offset := src.contiguous_view_offset()) is None: return None
+  buf = src.base
+  if buf.op is Ops.SLICE:
+    byte_offset = buf.src[1].val * buf.src[0].dtype.itemsize + offset * src.dtype.itemsize
+    buf = buf.src[0]
+    if byte_offset % buf.dtype.itemsize != 0: return None
+    offset = byte_offset // buf.dtype.itemsize
+  return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(offset)), src.numel())
+
+def contiguous_mops_to_view(c:UOp, src:UOp):
+  """MOPS(BUFFER) → SLICE when movement ops collapse to a contiguous range."""
+  buf = src.base
+  if buf.op not in {Ops.BUFFER, Ops.SLICE, Ops.UNSHARD}: return None
+  if src.op is Ops.RESHAPE and src.src[0].op in {Ops.BUFFER, Ops.SLICE} and c.op is not Ops.BITCAST: return None
+  if c.op is not Ops.BITCAST and src.op is Ops.BUFFER: return None
+
+  # no symbolic shape
+  if not all_int(c.shape): return None
+
+  if buf.op is not Ops.UNSHARD and (view := _make_buffer_view(src)) is not None:
+    view = (view.replace(dtype=c.dtype, arg=c.numel()) if c.op is Ops.BITCAST else view).reshape(c.shape)
+    return c.replace(src=(view,)) if c.op is Ops.COPY else view
+
+  # for UNSHARD tensors, use multi_pm to resolve per-shard movement ops, then create SLICE on the resolved result
+  if not isinstance(c.device, str):
+    from tinygrad.schedule.multi import multi_pm
+    resolved = graph_rewrite(src, multi_pm, name="multi_buffer_view")
+    if resolved.op is not Ops.UNSHARD: return None
+    if (view := _make_buffer_view(resolved.src[0])) is None: return None
+    return view.reshape(resolved.src[0].shape).unshard(resolved.arg, resolved.src[1:]).contiguous(tag=c.tag)
+
+  return None
+
+def _precompiled_output_redirect(s:UOp, t:UOp) -> UOp|None:
+  # how output s lands in the caller's buffer t, or None if it must be copied into t
+  # materialize straight into t
+  if s.op is Ops.CONTIGUOUS: return t.after(t.store(s.src[0]))
+  # rebind output storage to t
+  if s.op in {Ops.BUFFER, Ops.UNSHARD} and s.has_buffer_identity(): return t
+  return None
+
+def transform_precompiled_call(c:UOp) -> UOp|None:
+  if not c.arg.precompile: return None
+  assert c.src[0].op is Ops.TUPLE, f"expected TUPLE body for precompiled FUNCTION, got {c.src[0].op}"
+  input_buffers = tuple(x.contiguous() if x.op not in {Ops.AFTER, Ops.BIND} else x for x in c.src[1:])
+
+  # add the outputs to the call
+  srcs = c.src[0].src
+  resolved = [c.gettuple(i) for i in range(len(srcs))]
+  outs = tuple(r.empty_like() for r in resolved)
+  targets = [o.param_like(len(c.src)-1+i).shrink_to(s.shape) for i,(o,s) in enumerate(zip(outs, srcs))]
+
+  subs:dict[UOp, UOp] = {}
+  items:list[UOp] = []
+  for s, t in zip(srcs, targets):
+    after_deps:list[UOp] = []
+    while s.op is Ops.AFTER:
+      after_deps.extend(s.src[1:])
+      s = s.src[0]
+    if (placed := _precompiled_output_redirect(s, t)) is not None and s not in subs:
+      subs[s] = placed
+      items.append(s.after(*after_deps) if after_deps else s)
+    else:
+      items.append(t.after(t.store(s.after(*after_deps))))
+  fxn = UOp.sink(*(x.substitute(subs) for x in items))
+
+  # body switches from TUPLE to SINK, so the node becomes an opaque CALL (not FUNCTION)
+  new_call = UOp(Ops.CALL, src=(fxn, *input_buffers, *outs), arg=c.arg)
+  rets = tuple(o.after(new_call) for o in outs)
+
+  # if the CALL has symbolic shapes, shrink the max-sized output to the actual symbolic shape
+  # NOTE: must use resolved shapes from the FUNCTION (which substitutes PARAMs with external args), not raw body shapes
+  rets = tuple(r.shrink_to(rs.shape) for r,rs in zip(rets, resolved))
+
+  return UOp.maketuple(*rets)
+
+# NOTE: adding rules to here is bad. these all need to run before the schedule cache
+pm_early_transform_tensor_graph = PatternMatcher([
+  # transform precompiled FUNCTIONs into CALLs (body becomes SINK with stores)
+  (UPat(Ops.FUNCTION, name="c"), transform_precompiled_call),
+
+  # resolve TUPLE+GETTUPLE (for precompiled calls)
+  (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]),
+
+  # fold MOPS+BITCAST over BUFFER/SLICE into SLICE when movement ops collapse to contiguous range
+  (UPat((Ops.BITCAST, Ops.COPY, Ops.CONTIGUOUS), src=(UPat(GroupOp.Movement|{Ops.BUFFER}, name="src"),), name="c"), contiguous_mops_to_view),
+
+  # remove contiguous on movement ops before a copy on disk
+  (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.CONTIGUOUS).f(Ops.COPY, name="copy"), lambda x,copy:
+   copy.replace(src=(x,), tag=None) if isinstance(x.device, str) and x.device.startswith("DISK") else None),
+  # push copy past movement ops to disk
+  (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.COPY, name="copy"), lambda x,copy:
+   x.replace(src=(copy.replace(src=(x.src[0],), tag=None),)+x.src[1:]) \
+   if isinstance(x.device, str) and x.device.startswith("DISK") else None),
+
+  # add CONTIGUOUS to tagged UOps
+  (UPat(GroupOp.All-{Ops.CONTIGUOUS, Ops.AFTER, Ops.STORE}, name="x"),
+   lambda x: None if x.tag is None else x.rtag(None).contiguous(tag=x.tag) if x.tag else x.replace(tag=None)),
+  # remove extra CONTIGUOUS on AFTER (only when target is contiguous)
+  (UPat(Ops.CONTIGUOUS, src=(UPat(Ops.AFTER, name="a"),), name="c"),
+   lambda a,c: a.replace(tag=(a.tag or ())+(c.tag or ())) if a.src[0].has_buffer_identity() else None),
+  # replace AFTER+STORE with CONTIGUOUS when target is not a buffer
+  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(), UPat(name="src")))), name="u"), replace_store_after_with_contig),
+  # replace CONTIGUOUS with STORE+AFTER
+  (UPat(Ops.CONTIGUOUS, name="u"), replace_contig_with_store_after),
+  # remove DETACH/CONTIGUOUS_BACKWARD (allows more contiguous removal)
+  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),
+])
+
+def finalize_after(ctx:AllocCtx, x:UOp):
+  # untagged: record as an assign for the call body
+  if x.tag is None:
+    ctx.assigns.append(x)
+    return None
+  # tagged: untag and map each original pre-rewrite UOp to the stripped buffer; the untagged result is reprocessed as untagged
+  ret = x.replace(tag=None)
+  replace_uop = ret
+  # then, add views back
+  views:list[UOp] = []
+  while replace_uop.op in GroupOp.Movement|{Ops.UNSHARD, Ops.BITCAST, Ops.AFTER}:
+    if replace_uop.op is not Ops.AFTER: views.append(replace_uop)
+    replace_uop = replace_uop.src[0]
+  for v in reversed(views): replace_uop = v.replace(src=(replace_uop,)+v.src[1:])
+  for t in x.tag:
+    original_uop: UOp = ctx.uop_list[t]
+    ctx.buffer_map[original_uop] = replace_uop.shrink_to(original_uop.shape)
+  return ret
+
+def replace_input_buffer(ctx:AllocCtx, b:UOp):
+  ctx.replacements.append(b)
+  if b.op is Ops.BIND: return b.param_like(len(ctx.replacements)-1)
+  return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b.device,
+                   addrspace=b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL)
+
+pm_finalize_call = PatternMatcher([
+  (UPat(Ops.AFTER, name="x"), finalize_after),
+  (UPat(Ops.COPY, name="x"), lambda ctx,x: ctx.assigns.append(x) if isinstance(x.device, str) and x.device.startswith(("DISK", "TINYFS")) else None),
+])
+
+pm_replace_buf = PatternMatcher([
+  # replace BUFFER with PARAM for cache key normalization
+  (UPat(Ops.BUFFER, src=(UPat(),), name="b"), lambda ctx,b:
+   replace_input_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None),
+  # replace SLICE with PARAM. this rewrite is bottom up so BUFFERs we don't need won't be in the input
+  (UPat(Ops.SLICE, src=(UPat(Ops.BUFFER), UPat(Ops.CONST, dtype=dtypes.weakint)), name="b"), replace_input_buffer),
+  # strip value from BIND for cache key normalization, so different values hit same cache
+  (UPat(Ops.BIND, src=(UPat(Ops.PARAM), UPat(Ops.CONST)), name="b"), replace_input_buffer),
+])
+
+@rewrite_group(lambda _,ret: f"Callify {pluralize('Buffer', len(ret[1]))}")
+def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
+  if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Tensor Graph")
+  # uop list is a list in the original_sink graph and we can map to the tags later
+  # same predicate as Tensor.realize
+  ctx = AllocCtx(bases={base for x in big_sink.src if not (base:=x.base).is_virtual and not base.has_buffer_identity()
+                        and base.op is not Ops.AFTER and base.addrspace is not AddrSpace.ALU})
+
+  # this rewrite is "read-only", it adds simple things to buffer_map and may sink things on big_sink, bottom_up
+  # this is the only one where we have to be careful to not break the tensor graph
+  big_sink = graph_rewrite(big_sink, add_tags, ctx=ctx, bottom_up=True, name="number the uops")
+
+  # here we can break the tensor graph. this is the only place you need to maintain numbered tags
+  big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, name="early transform tensor graph")
+
+  # here we construct the final buffer_map: as-built nodes -> their final storage. values are never keys
+  graph_rewrite(big_sink, pm_finalize_call, ctx=ctx, name="finalize call")
+  ret = graph_rewrite(UOp.sink(*ctx.assigns), pm_replace_buf, ctx=ctx, bottom_up=True, name="replace bufs").call(*ctx.replacements)
+  assert not any(x in ctx.buffer_map for x in ctx.buffer_map.values())
+  if VIZ: graph_rewrite(ret, PatternMatcher([]), name="View Call")
+  return ret, ctx.buffer_map

 # *** all in scope Tensors are here. this gets relevant UOps ***

@@ -528,9 +528,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    if self.op is Ops.CONST: return self
    if self.op is Ops.SINK and all(s.op is Ops.CONST or (s.op is Ops.STACK and len(s.src) == 0) for s in self.src): return self
    # late import!
-    from tinygrad.uop.symbolic import symbolic
+    from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const
    with Context(TRACK_MATCH_STATS=0 if not tracked else TRACK_MATCH_STATS.value):
-      return graph_rewrite(self, symbolic, name="simplify")
+      return graph_rewrite(self, symbolic+pm_fold_cast_const, name="simplify")
  def ssimplify(self) -> UOp|ConstType: return ret.val if (ret:=self.simplify()).op is Ops.CONST else ret
  def _eval(self, dtype, expected_type:Type[T]) -> T:
    assert self.dtype in dtype, f"eval with wrong dtype {self}"
@@ -758,7 +758,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    assert arg is None or isinstance(self.device, tuple)
    inp = self if arg is None else UOp(Ops.MSELECT, src=(self,), arg=arg)
    if inp.dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {inp.dtype}")
-    return UOp(Ops.COPY, src=(inp,), arg=device)
+    return UOp(Ops.COPY, src=(inp.pad_to(inp.max_shape),), arg=device).shrink_to(inp.shape)
  def mselect(self, arg:int) -> UOp: return UOp(Ops.MSELECT, src=(self,), arg=arg)
  def mstack(self, *srcs: UOp) -> UOp: return UOp(Ops.MSTACK, src=(self,)+srcs) if len(srcs) else self
  @property
@@ -772,6 +772,15 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    if self.op is Ops.DETACH: return self.src[0].base  # DETACH can't change base
    return self

+  # base with UNSHARD
+  @property
+  def unsharded_base(self) -> UOp:
+    if self.op in GroupOp.Movement: return self.src[0].base
+    if self.op is Ops.DETACH: return self.src[0].base  # DETACH can't change base
+    # TODO: why can't this be in normal base?
+    if self.op is Ops.UNSHARD: return self.src[0].base
+    return self
+
  # cached property here makes external_uop_gc fail, why?
  @property
  def as_shape(self) -> tuple[sint, ...]:
@@ -1097,12 +1106,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    if self.op is Ops.CONST and self.val is not Invalid: return self.val, self.val
    if self.op is Ops.INDEX: return self.src[0]._min_max
    if self.op is Ops.CAST:
-      # an int destination truncates a float source toward zero. trunc is monotone
+      # rounding is monotone (truncation toward zero into an int, to-nearest onto the value grid into a float)
      smin, smax = self.src[0]._min_max
-      if dtypes.is_int(self.dtype) and dtypes.is_float(self.src[0].dtype) and all(math.isfinite(v) for v in (smin, smax)):
-        smin, smax = math.trunc(smin), math.trunc(smax)
-      # a cast to unsigned keeps exact bounds when the source fits
-      # TODO: can do more based on new dtype window
+      trunc = truncate.get(self.dtype) if dtypes.is_float(self.dtype) else math.trunc if dtypes.is_int(self.dtype) else None
+      if trunc is not None and all(math.isfinite(v) for v in (smin, smax)): smin, smax = trunc(smin), trunc(smax)
      if dtypes.is_unsigned(self.dtype) and 0 <= smin and smax <= self.dtype.max: return smin, smax
      if self.dtype in dtypes.floats+dtypes.sints+(dtypes.weakint,): return max(self.dtype.min, smin), min(smax, self.dtype.max)
    return self.dtype.min, self.dtype.max
@@ -1166,8 +1173,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    src: tuple[UOp, ...] = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),)
    return UOp(Ops.PARAM, src=src, arg=ParamArg(slot, dtype, vmin_vmax, multiple_of, name, addrspace, axis, device, volatile))
  def param_like(self, slot:int):
+    if self.op is Ops.BIND: return self.src[0].replace(arg=replace(self.src[0].arg, slot=slot, name=f"p{slot}"))
    addrspace = self.addrspace if self.addrspace is not None else AddrSpace.GLOBAL
-    if self.op is Ops.BIND: return self.src[0].replace(arg=replace(self.src[0].arg, slot=slot, addrspace=addrspace))
    return UOp.param(slot, self.dtype, self.shard_shape if self.axis is not None else self._shape, self.device, addrspace=addrspace, axis=self.axis)

  @staticmethod
@@ -1187,10 +1194,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    body = self if self.op is Ops.TUPLE else UOp.maketuple(self)
    return UOp(Ops.FUNCTION, src=(body,)+srcs, arg=CallInfo(grad_fxn, name, precompile, precompile_backward, aux))
  def custom_kernel(*srcs:UOp, fxn:Callable, grad_fxn:Callable|None=None) -> list[UOp]:
-    contig_srcs = tuple(x.contiguous() if x.op is not Ops.AFTER else x for x in srcs)
-    placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(contig_srcs)]
-    kernel = fxn(*placeholders).call(*contig_srcs, grad_fxn=grad_fxn)
-    return [s.after(kernel) for s in contig_srcs]
+    placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(srcs)]
+    kernel = fxn(*placeholders).call(*srcs, grad_fxn=grad_fxn)
+    return [s.after(kernel) for s in srcs]

  def to_elf(self) -> TinyELF:
    assert self.op is Ops.PROGRAM and isinstance(self.arg, ProgramInfo), "to_elf should only be called on a PROGRAM ast"
@@ -1515,55 +1521,52 @@ def add_trace_group(kt:TracingKey) -> None:
  tracked_ctxs.append([])

 active_group:list[int] = []
-def track_rewrites(name:Callable[..., str|TracingKey]|bool=True, replay:bool=False):
+active_rewrites:list[TrackedGraphRewrite] = []
+def rewrite_group(name:Callable[..., str|TracingKey]|bool=True, replay:bool=False, new_ctx:bool=True):
+  if not new_ctx: assert not callable(name) and not replay, "name fxn and replay are only supported for new_ctx groups"
  def _decorator(func):
    def __wrapper(*args, **kwargs):
+      # without tracking, we just call the function (unless top-level, which always profiles)
+      if TRACK_MATCH_STATS < 2 and not new_ctx: return func(*args, **kwargs)
      fn = key = func.__name__
      idx = -1
      if TRACK_MATCH_STATS >= 2:
-        add_trace_group(key:=TracingKey(n:=f"{fn} n{next(_name_cnt.setdefault(fn, itertools.count(1)))}", (n,)))
-        active_group.append(idx:=len(tracked_keys)-1)
+        if new_ctx:
+          add_trace_group(key:=TracingKey(n:=f"{fn} n{next(_name_cnt.setdefault(fn, itertools.count(1)))}", (n,)))
+          active_group.append(idx:=len(tracked_keys)-1)
+        else:
+          rewrite_name = str(kwargs.get("name", None) or fn)
+          assert args and isinstance(args[0], UOp), f"invalid match tracing inputs for {rewrite_name} with {args}"
+          loc = ((frm:=sys._getframe(1)).f_code.co_filename, frm.f_lineno)
+          depth = len(active_rewrites)
+          if not tracked_ctxs: add_trace_group(TracingKey(f"default {fn}"))
+          dest_group = active_group[-1] if active_group else len(tracked_ctxs)-1
+          tracked_ctxs[dest_group].append(ctx:=TrackedGraphRewrite(loc, args[0].trace_num, [], rewrite_name, depth, kwargs.get("bottom_up", False),
+                                                                   kwargs.get("walk", False), kwargs.get("enter_calls", False)))
+          active_rewrites.append(ctx)
+          key = rewrite_name  # profile spans are named after the rewrite step
      with cpu_profile(key, "TINY") as e:
        ret = func(*args, **kwargs)
-      if TRACK_MATCH_STATS >= 2: active_group.pop()
-      if TRACK_MATCH_STATS >= 2 and callable(name):
-        name_ret = name(*args, **kwargs, ret=ret)
-        assert isinstance(name_ret, (TracingKey, str)), f"name function returned {type(name_ret)}"
-        tracked_keys[idx] = k = TracingKey(n:=tracked_keys[idx].display_name.replace(fn, name_ret), (n,)) if isinstance(name_ret, str) else name_ret
-        e.name = TracingKey(k.display_name if isinstance(name_ret, str) else f"{fn} for {k.display_name}", k.keys)
+      if TRACK_MATCH_STATS >= 2:
+        if new_ctx: active_group.pop()
+        else: active_rewrites.pop()
+        if callable(name):
+          name_ret = name(*args, **kwargs, ret=ret)
+          assert isinstance(name_ret, (TracingKey, str)), f"name function returned {type(name_ret)}"
+          tracked_keys[idx] = k = TracingKey(n:=tracked_keys[idx].display_name.replace(fn, name_ret), (n,)) if isinstance(name_ret, str) else name_ret
+          e.name = TracingKey(k.display_name if isinstance(name_ret, str) else f"{fn} for {k.display_name}", k.keys)
      if CAPTURE_PROCESS_REPLAY and replay:
        # find the unittest frame we're capturing in
        frm = sys._getframe(1)
        while (f_back:=frm.f_back) is not None and "unittest" not in f_back.f_code.co_filename: frm = f_back
-        loc = f"{frm.f_code.co_filename.split('/')[-1]}:{frm.f_lineno} {frm.f_code.co_name}"
+        replay_loc = f"{frm.f_code.co_filename.split('/')[-1]}:{frm.f_lineno} {frm.f_code.co_name}"
        # capture global context vars and all the args passed in
        inputs = (fn, args, kwargs, ContextVar._cache)
-        replay_capture.append(pickle.dumps(inputs+(loc, ret)))
+        replay_capture.append(pickle.dumps(inputs+(replay_loc, ret)))
      return ret
    return __wrapper
  return _decorator

-active_rewrites:list[TrackedGraphRewrite] = []
-def profile_matches(fxn:Callable):
-  def wrap_profile_matches(*args, **kwargs):
-    if TRACK_MATCH_STATS >= 2:
-      name = str(kwargs.get("name", None) or fxn.__name__)
-      assert args and isinstance(args[0], UOp), f"invalid match tracing inputs for {name} with {args}"
-      loc = ((frm:=sys._getframe(1)).f_code.co_filename, frm.f_lineno)
-      depth = len(active_rewrites)
-      if not tracked_ctxs: add_trace_group(TracingKey(f"default {fxn.__name__}"))
-      dest_group = active_group[-1] if active_group else len(tracked_ctxs)-1
-      tracked_ctxs[dest_group].append(ctx:=TrackedGraphRewrite(loc, args[0].trace_num, [], name, depth, kwargs.get("bottom_up", False),
-                                                               kwargs.get("walk", False), kwargs.get("enter_calls", False)))
-      active_rewrites.append(ctx)
-      with cpu_profile(name, "TINY"):
-        ret = fxn(*args, **kwargs)
-      active_rewrites.pop()
-      return ret
-    # without tracking, we just call the function
-    return fxn(*args, **kwargs)
-  return wrap_profile_matches
-
 class TrackedPatternMatcher(PatternMatcher):
  def rewrite(self, uop:UOp, ctx=None):
    if len(pats:=self.pdict.get(uop.op, [])):
@@ -1742,7 +1745,7 @@ class RewriteContext:
          if n in waitlist: stack.extend(waitlist.pop(n))
    return self.replace[root]

-@profile_matches
+@rewrite_group(new_ctx=False)
 def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=None, bpm=None, walk=False, enter_calls=False) -> UOp:
  rewrite_ctx = RewriteContext(pm if not bottom_up else None, pm if bottom_up else bpm, ctx, enter_calls)
  return rewrite_ctx.walk_rewrite(sink) if walk else rewrite_ctx.unified_rewrite(sink)
@@ -1755,73 +1758,6 @@ def _rebuild_dtype(n:UOp, new_src:tuple[UOp,...]) -> DType:
 def sint_to_uop(x:sint, dtype=dtypes.weakint) -> UOp: return UOp.const(x, dtype)
 def to_max_shape(shape:tuple[sint, ...]) -> tuple[int, ...]: return tuple(int(x.vmax) if isinstance(x, UOp) else x for x in shape)

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

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

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

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

 # late imports to avoid circular import
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.schedule.rangeify import BufferizeOpts
+from tinygrad.renderer import Estimates
 glbls:dict[str, Any] = {"inf": math.inf, "nan": math.nan, "KernelInfo": KernelInfo, "Metadata": Metadata,
                        "UOp": UOp, "dtypes": dtypes, "Ops": Ops, "AxisType": AxisType, "Invalid": Invalid,
                        "Opt": Opt, "OptOps": OptOps, "BufferizeOpts": BufferizeOpts, "AddrSpace": AddrSpace, "panic": panic,
-                        "ConstFloat": ConstFloat, "ParamArg": ParamArg}
+                        "ConstFloat": ConstFloat, "ParamArg": ParamArg, "Estimates": Estimates}
 def eval_pyrender(code:str) -> UOp:
  lcls:dict[str, Any] = {}
  exec(code, glbls, lcls)
@@ -68,7 +68,7 @@ invalid_pat = UPat(Ops.CONST, arg=Invalid, name="i")
 invalid_gate = UPat.var("cond").where(UPat.var("x"), invalid_pat)
 pm_data_invalid = PatternMatcher([
  (invalid_pat.broadcast(), lambda i: i),
-  (UPat(GroupOp.Unary|{Ops.BITCAST}, src=(invalid_pat,)), lambda i: i),
+  (UPat(GroupOp.Unary|{Ops.CAST, Ops.BITCAST}, src=(invalid_pat,)), lambda i: i),
  (UPat(GroupOp.Unary|{Ops.CAST, Ops.BITCAST}, src=(invalid_gate,), name="op"),
   lambda cond,x,op,i: cond.where(op.replace(src=(x,)), i)),
  # binary ops move inside the gate, with Invalid in the false branch
@@ -96,6 +96,10 @@ pm_remove_invalid = PatternMatcher([
   if any(x.is_invalid for x in s.src) else None),
 ])

+# the one rule that collapses the pair CAST(dt, CONST(v)) into a typed CONST
+# TODO: delete this once CONST has no dtype
+pm_fold_cast_const = PatternMatcher([(UPat(Ops.CAST, name="root", src=(UPat.cvar("c"),)), lambda root, c: root.const_like(c.val))])
+
 symbolic_simple = pm_data_invalid + PatternMatcher([
  # ** self folding **
  (UPat.var("x") + 0, lambda x: x),    # x+0 -> x
@@ -152,8 +156,6 @@ symbolic_simple = pm_data_invalid + PatternMatcher([
  (UPat.var("x") * 0, lambda x: x.const_like(float("nan") if x.op is Ops.CONST
                                             and isinstance(x.val, float) and (math.isnan(x.val) or math.isinf(x.val)) else 0)),
  # *** cast/bitcast ***
-  # TODO: delete this once CONST has no dtype
-  (UPat(Ops.CAST, name="root", src=(UPat.cvar("c"),)), lambda root, c: root.const_like(c.val)),
  (UPat((Ops.CAST, Ops.BITCAST), name="root"), lambda root: root.src[0] if root.dtype == root.src[0].dtype else None),
  (UPat(Ops.BITCAST, name="root", src=(UPat.cvar("c"),)), fold_bitcast),
  # b.cast(a).cast(b) -> b if a preserves all values in b
@@ -253,10 +255,11 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
  # ** two stage ALU folding **
  *((UPat.var("x").alu(op, UPat.cvar("c1")).alu(op, UPat.cvar("c2")).named("f"),
     lambda f,x,c1,c2: x.alu(f.op,c1.alu(f.op,c2))) for op in GroupOp.Associative),
-  ((UPat.cvar("c0") + UPat.var("x")) < UPat.cvar("c1"), lambda x,c0,c1: x<(c1-c0)),  # c0 + x < c1 -> x < c1 - c0
  # (x//c1)//c2 -> x//(c1*c2) for c2>0
  ((UPat.var("x") // UPat.cvar("c1")) // UPat.cvar("c2"), lambda x,c1,c2: x//(c1*c2) if c2.vmin>0 else None),
  # ** lt **
+  # c0+x<c1 -> x < c1-c0
+  ((UPat.cvar("c0") + UPat.var("x", dtype=dtypes.ints+(dtypes.weakint,))) < UPat.cvar("c1"), lambda x,c0,c1: x<(c1-c0)),
  # c0*x<c1 -> sign(c0)*x < ceil(c1/abs(c0))
  ((UPat.cvar("c0")*UPat.var("x", dtype=dtypes.weakint))<UPat.cvar("c1"),
   lambda x,c0,c1: (x if c0.val > 0 else -x)<-(-c1.val//abs(c0.val)) if abs(c0.val) > 1 else None),
@@ -285,7 +288,7 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
    (UOp.const(x.val) if x.op is Ops.CONST else x.cast(dtypes.int)).alu(u.op,
     UOp.const(y.val) if y.op is Ops.CONST else y.cast(dtypes.int)).cast(u.dtype)
    if not any(v.overflows(dtypes.int) for v in (u,x,y)) else None),
-  ((UPat.var("x", dtypes.weakint) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+c.cast(cast.dtype)),
+  ((UPat.var("x", dtypes.weakint) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+cast.const_like(c.val)),
  # only RANGE/IF/STORE/KERNEL have side effects
  (UPat(Ops.AFTER, name="x"), lambda x: x.replace(src=(x.src[0],)+
    tuple(dedup(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.CALL, Ops.FUNCTION, Ops.BARRIER, Ops.END, Ops.LINEAR, Ops.STAGE}
@@ -0,0 +1,78 @@
+from dataclasses import replace
+from tinygrad.dtype import dtypes, DType, AddrSpace, Invalid, least_upper_dtype, strong_dtype, weak_dtype
+from tinygrad.helpers import unwrap
+from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, GroupOp, graph_rewrite, dtype_from_uop
+
+def select_dtype(u:UOp):
+  if u.dtype is dtypes.weakfloat: return dtypes.default_float
+  return dtypes.long if u.overflows(dtypes.int32) else dtypes.int
+
+def lower_weak_node(u:UOp) -> UOp|None:
+  start, src = (1 if u.op is Ops.WHERE else 0), tuple(s.src[0] if s.op is Ops.CAST and s.dtype in dtypes.weaks else s for s in u.src)
+  if src == u.src or any(s.dtype in dtypes.weaks for s in src[start:]): return None
+  dt = strong_dtype(least_upper_dtype(select_dtype(u), *(s.dtype for s in src)) if u.op in GroupOp.Binary
+                    else unwrap(dtype_from_uop(u.op, src, u.arg)))
+  return u.replace(dtype=None, src=src[:start]+tuple(s if s.base.is_invalid else s.cast(dt) for s in src[start:])).cast(u.dtype)
+
+pm_lower_weak = PatternMatcher([
+  (UPat(Ops.CONST, dtype=dtypes.weaks, name="u"), lambda u: UOp.const(u.val, select_dtype(u)).cast(u.dtype)),
+  # two stacked weak casts are a weakint value used as weakfloat (or vice versa): resolve the inner one at the outer kind's default.
+  # a SINGLE weak cast is never rewritten here, each consumer absorbs it on its own edge (see lower_weak_srcs)
+  (UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat.var("x"),)),), name="u"),
+   lambda u,x: x.cast(select_dtype(u)).cast(u.dtype) if x.dtype not in dtypes.weaks else None),
+  # Binary can widen from the bounds, all other nodes derive from the lowered sources.
+  # a weakfloat Unary (sin/exp2/...) must resolve here, before the transcendental decomposition
+  (UPat(GroupOp.Binary|GroupOp.Unary|{Ops.WHERE, Ops.RANGE, Ops.STACK, Ops.SPECIAL}, name="u"), lower_weak_node),
+  (UPat(Ops.PARAM, dtype=dtypes.weakint, name="u"),
+    lambda u: u.replace(dtype=None, arg=replace(u.arg, dtype=select_dtype(u))).cast(dtypes.weakint) if u.addrspace == AddrSpace.ALU else None),
+])
+
+def lower_weak_srcs(ctx:dict[UOp, UOp]|None, u:UOp) -> UOp|None:
+  if ctx is None: ctx = {}
+  def lower(s:UOp) -> UOp:
+    if (r:=ctx.get(s)) is None:
+      r = graph_rewrite(s, pm_lower_weak)
+      # the consumer absorbs the cast on its own edge
+      ctx[s] = r = r.src[0] if r.op is Ops.CAST and r.dtype in dtypes.weaks else r
+    return r
+  # a comparison demands a common operand width: lower it whole so the Binary rule unifies its operands
+  ret = lower(u) if u.op in GroupOp.Comparison else u.replace(src=tuple(lower(s) if s.dtype in dtypes.weaks else s for s in u.src))
+  return None if ret is u else ret
+
+def commit_weak(s:UOp, dt:DType) -> UOp:
+  # a bare weak CONST commits directly (the value stays mathematical, emission truncates), a weak non-const src takes the demand cast
+  return UOp.const(s.val, dt) if s.op is Ops.CONST else s.cast(dt)
+
+def commit_weak_srcs(u:UOp) -> UOp|None:
+  if not any(s.dtype in dtypes.weaks for s in u.src): return None
+  if (dt:=least_upper_dtype(*(s.dtype for s in u.src))) in dtypes.weaks: return None
+  # the root re-derives: a shift's dtype is its lhs's, so committing the lhs commits the node too
+  return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src))
+
+# runs in index lowering and in the decomps: a rule that mints a weak const commits it in the same rewrite, so none reaches the renderer
+pm_commit_weak = PatternMatcher([
+  (UPat(GroupOp.Broadcastable, name="u"), commit_weak_srcs),
+  # demand from the destination: a STORE's weak value commits at the destination's dtype
+  (UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.weaks)), allow_any_len=True, name="u"),
+   lambda u: u.replace(src=(u.src[0], commit_weak(u.src[1], u.src[0].dtype), *u.src[2:]))),
+])
+
+# a concrete CAST over a weak node states the width the value will live at. that width is a floor, never a narrowing
+def cast_weak_srcs(c:UOp, u:UOp) -> UOp|None:
+  if c.dtype in dtypes.weaks or weak_dtype(c.dtype) is not u.dtype: return None
+  dt = least_upper_dtype(c.dtype, select_dtype(u))
+  return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src)).cast(c.dtype)
+
+pm_cast_weak = PatternMatcher([
+  (UPat(Ops.CAST, name="c", src=(UPat(GroupOp.ALU, dtype=dtypes.weaks, name="u"),)), cast_weak_srcs),
+])
+
+pm_lower_index_dtype = pm_commit_weak+pm_cast_weak+PatternMatcher([
+  (UPat(GroupOp.All, name="u"),
+   lambda ctx,u: lower_weak_srcs(ctx, u) if u.dtype not in dtypes.weaks and any(s.dtype in dtypes.weaks for s in u.src) else None),
+  # a valid index into an n-element buffer lives in [0,n): a gated long index narrows when n-1 fits int32 (out-of-gate wraps, discarded)
+  # TODO: more generic
+  (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"), UPat.var("gate").where(UPat.var("idx", dtypes.long), UPat(Ops.CONST, arg=Invalid))),
+        allow_any_len=True, name="u"),
+   lambda u,buf,gate,idx: u.replace(src=(buf, idx.cast(dtypes.int).valid(gate))+u.src[2:]) if buf.max_numel()-1 <= dtypes.int32.max else None),
+])
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	a0d4a2971b	kimi linear working	2026-08-10 15:37:20 +00:00
geohot	951aaf893b	llm: reuse exact recurrent prefixes	2026-08-10 14:06:15 +00:00
geohot	a455e17539	llm: fully warm recurrent serving at startup	2026-08-10 14:06:15 +00:00
nimlgenandgeohot	778f8aee59	fix hevc (#17477 ) * hevc tests * x	2026-08-10 14:06:15 +00:00
qazalandgeohot	9b347cc3e7	late loss.to("CPU") in llama (#17476 ) * late loss.to("CPU") in llama * acc = 0	2026-08-10 14:06:15 +00:00
qazalandgeohot	32b9149040	llama: custom silu kernels (#17462 ) * start by copying the C * uop kernel * cleanup tests * estimates is part of SPEC	2026-08-10 14:06:15 +00:00
George Hotzandgeohot	bbd4a77351	move platform tests to platform.yml (#17475 ) * ci: split mac/windows/qcom-cl tests into platform.yml Move the 6 jobs that don't run on Linux (4 macos, 1 windows, 1 QCOM CL compile test on arm) out of test.yml into a separate Platform Tests workflow so they run (and can be gated/runners-matched) independently. * ci: gate platform tests to the upstream repo Skip mac/windows/qcom-cl jobs anywhere but tinygrad/tinygrad, so the Platform Tests workflow is disabled on the gitea fork (and any fork). * ci: revert repo gate on platform tests Job-level if is only evaluated by gitea when a runner with matching labels fetches the task; with no mac/windows/arm runners the jobs queue forever. Disable the workflow on the instance instead.	2026-08-10 14:06:15 +00:00
geohot	e103421a12	document Kimi K3 optimization handoff	2026-08-10 06:14:30 -07:00
geohot	2b1b8c22a9	add fast exact-shape Kimi K3 benchmark	2026-08-10 06:13:24 -07:00
geohot	f53f0e7e79	document final Kimi K3 MI350 load profiling	2026-08-10 06:13:24 -07:00
geohot	224bac0318	speed up Kimi K3 decode projections on MI350X	2026-08-10 06:13:24 -07:00
geohot	1d86204718	optimize Kimi K3 serving on MI350X	2026-08-10 06:13:24 -07:00
geohot	c6ac4961d7	speed up Kimi K3 safetensor loading	2026-08-10 06:13:24 -07:00
geohot	1b3732a6ed	speed up Kimi K3 TP8 loading on MI350X	2026-08-10 06:13:24 -07:00
geohot	553bdf68e6	llm: defer recurrent server warmup	2026-08-10 06:59:49 +00:00
geohot	4e1c0166f8	llm: stabilize recurrent serving across requests	2026-08-10 06:43:27 +00:00
geohot	d28f5f261b	llm: keep recurrent single-token shapes static	2026-08-10 05:45:20 +00:00
geohot	14595b9ae8	llm: prepare direct Kimi K3 serving on MI350X	2026-08-10 05:28:18 +00:00
geohot	eaf7822239	llm: accelerate Kimi serving on gfx11	2026-08-10 03:58:44 +00:00
geohot	77e5be99bc	llm: prepare Kimi K3 and accelerate recurrent prefill	2026-08-10 02:03:14 +00:00
geohot	6edb5f9698	get kimi-linear running on 4x7900XTX (codex slop)	2026-08-10 00:38:38 +00:00
nimlgenandGitHub	8c8b43de62	hcq2: fix beam (#17467 ) * fix beam * x	2026-08-09 16:53:47 +03:00
nimlgenandGitHub	e17c21e102	hcq2: timings (#17464 ) * hcq2: timings * Dx * x * x * x * x * align * x	2026-08-08 22:00:32 +03:00
George HotzandGitHub	d4d537c8ae	add SPEC checking for the kernel graph (#17432 ) * add SPEC checking for the kernel graph * skip test * raise there, not None * handwritten * issue with unshard * multi works * and bitcast * fix new tests * fix linear * remove call index * fix shrink * fixes	2026-08-08 10:00:06 -07:00
b1tgandGitHub	abe2256299	fix symbolic sharded reshape (#17463 )	2026-08-08 09:18:02 -07:00
b1tgandGitHub	8c49a7a34b	support symbolic shapes in copy (#17461 ) * pad_to is no-op when same shape * support symbolic shapes in copy	2026-08-08 09:16:59 -07:00
qazalandGitHub	9dd3b8402e	default llama 8b to MXFP4=1 (#17465 )	2026-08-09 00:13:32 +08:00
sirhcmandGitHub	c0d2f9ac0c	nolocals supports variables (#17457 )	2026-08-07 17:54:38 -04:00
nimlgenandGitHub	4c206a52b1	fix ci emu (gpt) (#17437 ) * fix ci emu * x	2026-08-07 22:36:59 +03:00
chenyuandGitHub	4a3b8f6501	better _drop_valid_stmts [pr] (#17454 )	2026-08-07 15:18:35 -04:00
chenyuandGitHub	59b88ea5e2	move pm_fold_cast_const [pr] (#17453 ) move to lower index dtype	2026-08-07 13:31:56 -04:00
chenyuandGitHub	f76422b8af	fix cast to float _min_max [pr] (#17451 )	2026-08-07 11:56:19 -04:00
wozeparrotandGitHub	1827ec57f7	gptoss: fix sharded invalids (#17450 )	2026-08-07 08:42:02 -07:00
nimlgenandGitHub	b6189db8e9	cpu: fix eintr (#17449 )	2026-08-07 17:46:57 +03:00
chenyuandGitHub	73e670c10f	c0+x<c1 -> x < c1-c0 is ints only [pr] (#17448 )	2026-08-07 10:44:01 -04:00
chenyuandGitHub	fca695a36f	clean up reduce MUL gradient (#17447 )	2026-08-07 10:08:04 -04:00
Robert JosephandGitHub	0c96cdc300	fix prod gradients at zero (#17404 )	2026-08-07 09:56:52 -04:00
chenyuandGitHub	baa6148066	fix var of large half input (#17444 ) * fix var of large half input similar to mean, we use sum_acc_dtype for denominator * mypy	2026-08-06 23:14:32 -04:00
George Hotz GitHub qazal	1858f1fd9a	viz: collapse PROGRAM nodes like CALL (codex) (#17438 ) Co-authored-by: qazal <[email protected]>	2026-08-07 11:44:34 +09:00
qazal GitHub geohot	f253c4469d	remove contiguous from custom_kernel (#17149 ) * no user contig on custom_kernel * clean up * non removable * test MXFP4 llama without hipcc * use compiler-free HIPCC renderer in llama CI * move llama coverage to AMD tests * run llama coverage in AMD test matrix * respect configured ROCm path in llama profile * work * clean up * fix * add views back * remove that * update test * test_double_permute one less kernel * test_shrink less kernels --------- Co-authored-by: George Hotz <[email protected]>	2026-08-07 11:20:59 +09:00
chenyuandGitHub	28195d51fb	fix f2f from fp8e5m2fnuz to half (#17442 ) * fix f2f from fp8e5m2fnuz to half * it works if it's supported	2026-08-06 21:19:46 -04:00
chenyuandGitHub	9020a88f03	truncate float in DType.const [pr] (#17439 )	2026-08-06 20:19:08 -04:00
chenyuandGitHub	d8cbc11105	update linear interpolate to use int math for indices (#17441 )	2026-08-06 20:18:39 -04:00
wozeparrotandGitHub	1fd6b1035f	fa: swa support (#17367 )	2026-08-06 08:07:30 -07:00
nimlgenandGitHub	46230e9f17	hcq2: fence inputs (#17436 )	2026-08-06 16:27:12 +03:00
qazalandGitHub	9636dd1a25	test MXFP4 llama without hipcc (#17435 ) * test MXFP4 llama without hipcc * first pythonpath then dev	2026-08-06 17:31:40 +09:00
qazalandGitHub	f258708d7d	llama: custom quantize_mxfp4+transpose kernel (codex) (#17434 ) * llama: custom quantize_mxfp4+transpose kernel (codex) * rename to cpp * inline * cleanup * lds load_bf16x4 * more tests, add Estimates	2026-08-06 16:13:28 +09:00
chenyuandGitHub	28e6ef6937	fix postopt symbolic [pr] (#17433 ) REDUCE with src simplied to const would become unparented	2026-08-06 00:12:30 -04:00
chenyuandGitHub	969df866a3	one less strong dtype const in symbolic [pr] (#17431 )	2026-08-05 23:35:06 -04:00
chenyuandGitHub	7a9cd8e329	move weak function and pm to uop/weak [PR] (#17429 )	2026-08-05 22:27:40 -04:00
George HotzandGitHub	b4372df9c6	revert wrong custom kernel fix (#17427 )	2026-08-05 18:31:35 -07:00
chenyuandGitHub	d51e55aa17	remove some pm_fold_cast_const [pr] (#17426 )	2026-08-05 21:28:38 -04:00
sirhcmandGitHub	be25207a7a	scope variable names inside CALLs (#17424 )	2026-08-05 20:59:09 -04:00
chenyuandGitHub	d726e5f7f3	split pm_fold_cast_const [PR] (#17425 ) need to delete this rule that writes strong typed CONST	2026-08-05 19:56:04 -04:00
George HotzandGitHub	470c032a5e	fix slice + non contig kernels (#17423 ) * movement: resolve negative int slice bounds against symbolic sizes negative int bounds in a slice against a symbolic dim were passed through unresolved, giving wrong views. resolve them against the (possibly symbolic) size, like slice.indices does for int dims * schedule: realize custom kernel inputs that don't resolve to a buffer state rangeify assigns ranges backward from consumers and CALL contributes none, so the subgraph above a custom kernel input gets no ranges unless something in it is realized, and reduce conversion crashes with a KeyError. realize call inputs that don't resolve to a buffer state. only view-only movement ops preserve the underlying buffer: anything computed (ALU, REDUCE, ...) must be realized even if one of its sources resolves to a buffer, since the whole subgraph above the call has no ranges. unwrapping src[0] unconditionally missed const branches hanging off non-src[0] children and silently resolved REDUCEs to their source buffer. includes regression tests for pure const, mixed buffer+const, and view-over-buffer inputs	2026-08-05 16:18:17 -07:00
geohot	a8a8030bc9	add benchmark_llm script	2026-08-05 15:59:51 -07:00
George HotzandGitHub	581bfdd94f	merge track_rewrites and profile_matches into rewrite_group [PR] (#17420 ) * merge track_rewrites and profile_matches into rewrite_group * bug * flip ctx polarity	2026-08-05 14:41:38 -07:00
chenyuandGitHub	07ac911665	few weak and decomp tweaks [PR] (#17419 )	2026-08-05 15:55:04 -04:00
chenyuandGitHub	c2f1e5ae2a	fix weak cast to strong dtype [pr] (#17418 ) weak can mean higher than that strong dtype, so always use that strong dtype is wrong	2026-08-05 15:32:23 -04:00
George HotzandGitHub	757a727808	move callify into tensor (#17416 )	2026-08-05 11:48:56 -07:00
George HotzandGitHub	2cce85a606	chat: display reasoning_content from streamed responses (#17414 ) * chat: display reasoning_content from streamed responses The server's StreamRouter emits reasoning_content deltas for think blocks, but the chat UI was only reading delta.content, silently dropping all reasoning. Now reasoning is shown in gray (#888) and included in the message history sent back to the server. * fix	2026-08-05 10:49:24 -07:00
nimlgenandGitHub	9b27ea8523	hcq2: cleaner (#17413 ) * hcq2: cleaner * x	2026-08-05 19:45:04 +03:00
chenyuandGitHub	6cb419b9b7	regression test for bert nan with weak (#17412 )	2026-08-05 12:14:26 -04:00
nimlgenandGitHub	5b0b68ec55	remove debug from test (#17410 )	2026-08-05 15:47:40 +03:00
qazalandGitHub	ad32bd272b	viz/cli: faster and more complete rewrites print (#17411 ) * viz/cli: faster and more complete matches print * kwargs	2026-08-05 20:02:27 +09:00
nimlgenandGitHub	874d33128b	hcq2 benchmark (#17235 ) * hcq2 in ci? * fix * traning * x * x * x * recover * debug * impler * x * x * x * hcq2: group input scatter plans by destination * hcq2: simplify input scatter tables * x	2026-08-05 10:00:42 +03:00
chenyuandGitHub	3bf9e70b19	Revert "don't cast weak in _broadcasted [pr] (#17408 )" (#17409 ) This reverts commit `b45058b5ec`.	2026-08-05 02:40:33 -04:00
Raine GitHub George Hotz	77e124e455	fix AMD WMMA emulation and test in CI (#17184 ) * fix SPEC=1 test_tensor_cores * implement i32 WMMA for RDNA3, add regression test to CI * gfx950 scaled mfma llvmir fix * detect VOP3PX2 in emu * start cdna4 scaled mfam emu * fix gfx12 llvmir signatures * oops * fix * fix src2 const field extraction (gpt) * scaled mfma fixes cdna4 (kimi) * fp8 out breaks mfma (glm) * fix const signature * another * fix f string for linter * lint * clean * and a final lint * . * fix mypy * skip slow tests on ci * reduce unroll tensor shape -> 64x64 --------- Co-authored-by: George Hotz <[email protected]>	2026-08-04 23:25:32 -07:00
chenyuandGitHub	b45058b5ec	don't cast weak in _broadcasted [pr] (#17408 ) * don't cast weak in _broadcasted [pr] * fine now?	2026-08-05 02:19:31 -04:00
George HotzandGitHub	46f0003776	more KernelCountException (#17407 )	2026-08-04 22:55:38 -07:00