variables should not be weakints

also no gradient for weak target

.github/actions/setup-tinygrad/action.yml

@@ -41,10 +41,6 @@ inputs:
     description: "Install LLVM?"
     required: false
     default: 'false'
-  tinydreno:
-    description: "Install tinydreno"
-    required: false
-    default: 'false'
   qemu:
     description: "Install qemu"
     required: false
@@ -277,12 +273,6 @@ runs:
       shell: bash
       run: brew install llvm@20
 
-    # *** tinydreno ***
-    - name: Install tinydreno (linux)
-      if: inputs.tinydreno == 'true' && runner.os == 'Linux'
-      shell: bash
-      run: sudo curl -fL https://github.com/sirhcm/tinydreno/raw/refs/heads/master/libllvm-qcom.so -o /usr/lib/libllvm-qcom.so
-
     # *** OpenCL ***
     - name: Install rusticl
       if: inputs.opencl == 'true'

.github/workflows/platform.yml

@@ -179,35 +179,3 @@ jobs:
       - 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

.github/workflows/test.yml

@@ -21,7 +21,7 @@ concurrency:
 jobs:
   docs:
     name: Docs
-    runs-on: &linux ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 10
     env:
       CHECK_OOB: 0
@@ -61,7 +61,7 @@ jobs:
 
   torchbackend:
     name: Torch Backend Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     steps:
     - name: Checkout Code
@@ -88,7 +88,7 @@ jobs:
 
   bepython:
     name: Python Backend
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     steps:
     - name: Checkout Code
@@ -126,7 +126,7 @@ jobs:
 
   linter:
     name: Linters
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 10
 
     steps:
@@ -157,7 +157,7 @@ jobs:
 
   nulltest:
     name: Null Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
 
     steps:
@@ -191,7 +191,7 @@ jobs:
 
   unittest:
     name: Unit Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
 
     steps:
@@ -228,7 +228,7 @@ jobs:
       matrix:
         group: [1, 2]
     name: SPEC=2 (${{ matrix.group }})
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     steps:
     - name: Checkout Code
@@ -244,7 +244,7 @@ jobs:
 
   fuzzing:
     name: Fuzzing
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 10
     steps:
     - name: Checkout Code
@@ -260,7 +260,7 @@ jobs:
 
   testopenclimage:
     name: CL IMAGE Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     steps:
       - name: Checkout Code
@@ -280,7 +280,7 @@ jobs:
 
   testopenpilot:
     name: openpilot Compile Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     steps:
       - name: Checkout Code
@@ -309,7 +309,7 @@ jobs:
 
   testonnxcpu:
     name: ONNX (CPU) Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 20
 
     steps:
@@ -328,7 +328,7 @@ jobs:
 
   testoptim:
     name: Optimization Tests
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 20
     steps:
       - name: Checkout Code
@@ -360,7 +360,7 @@ jobs:
 
   testllm:
     name: Test LLM
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     env:
       CHECK_OOB: 0
@@ -387,7 +387,7 @@ jobs:
 
   testmodels:
     name: Models
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     steps:
       - name: Checkout Code
@@ -407,7 +407,7 @@ jobs:
 
   testdsp:
     name: Linux (DSP)
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     steps:
     - name: Checkout Code
@@ -435,7 +435,7 @@ jobs:
           - 'WEBGPU'
 
     name: Linux (DEV=${{ matrix.dev }})
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 20
     steps:
     - name: Checkout Code
@@ -461,7 +461,7 @@ jobs:
 
   testamdasm:
     name: AMD ASM IDE
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 20
     env:
       DEV: MOCKKFD+AMD
@@ -507,7 +507,7 @@ jobs:
 
   hcq2:
     name: hcq2
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 5
     steps:
       - name: Checkout Code
@@ -521,16 +521,15 @@ jobs:
       - name: Run HCQ2 tests
         run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/test_tiny.py
       - name: Run HCQ2 multi-device tests
-        run: |
-          HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_multitensor.py \
-            TestMultiTensor.test_simple_add TestMultiTensor.test_shard_reduce \
-            TestMultiTensor.test_backward_sum TestMultiTensor.test_matmul_shard_0_0
+        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest -n=auto test/backend/test_multitensor.py
       - name: Run HCQ2 JIT tests
         run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_jit.py
+      - name: Run HCQ2 unit tests
+        run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest test/device/test_hcq2.py
 
   testmockam:
     name: Linux (am)
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     env:
       DEV: MOCKPCI+AMD
@@ -566,7 +565,7 @@ jobs:
         arch: [gfx1100, gfx1201, gfx950]
 
     name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
     env:
       DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
@@ -589,7 +588,7 @@ jobs:
         if: ${{ matrix.backend == 'amd' && matrix.arch == 'gfx950' }}
         run: PYTHONPATH=. DEV=NULL:HIP:gfx950 MXFP4=1 LLAMA_LAYERS=2 BENCHMARK=3 NULL_ALLOW_COPYOUT=1 NO_HIPCC=1 ROCM_PATH=/opt/rocm JITBEAM=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh
       - name: Run pytest (amd)
-        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM test/opt/test_tensor_cores.py --durations=20
+        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM --durations=20
       - name: Run disk copy tests
         run: python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk
       - name: Run TRANSCENDENTAL math
@@ -604,7 +603,7 @@ jobs:
         backend: [ptx, nv]
 
     name: Linux (${{ matrix.backend }})
-    runs-on: *linux
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 20
     env:
       FORWARD_ONLY: 1
@@ -638,10 +637,17 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        backend: [ir3, nak]
-    name: Compile-only (${{ matrix.backend }})
-    runs-on: *linux
+        dev:
+          - 'NULL:IR3:a630'
+          - 'NULL:QCOMCL:a630'
+          - 'NULL:NAK:sm_120'
+    name: Compile-only (DEV=${{ matrix.dev }})
+    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 15
+    env:
+      NULL_ALLOW_COPYOUT: 1
+      DEV: ${{ matrix.dev }}${{ contains(matrix.dev, 'a630') && ',IMAGE_PITCH_ALIGNMENT=64' || '' }}
+      IMAGE: ${{ contains(matrix.dev, 'a630') && '1' || '0' }}
     steps:
       - name: Checkout Code
         uses: actions/checkout@v6
@@ -650,21 +656,14 @@ jobs:
         with:
           key: compile-${{ matrix.backend }}
           deps: "testing_unit mesa"
-      - name: Set env
+          qemu: ${{ contains(matrix.dev, 'QCOMCL') }}
+      - name: Test IMAGE
         shell: bash
-        run: printf "NULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'DEV=NULL:IR3:a630' || matrix.backend == 'nak' && 'DEV=NULL:NAK:sm_120' }}" >> $GITHUB_ENV
+        if: contains(matrix.dev, 'a630')
+        run: DEBUG=7 python3 test/backend/test_ops.py TestOps.test_gemm | grep isam
       - name: Run test_ops
         shell: bash
         run: |
           python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
           DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
           python -m pytest -n=auto test/backend/test_ops.py --durations=20
-      - name: Run test_ops (IMAGE)
-        if: matrix.backend == 'ir3'
-        shell: bash
-        env:
-          IMAGE: 1
-          DEV: "NULL:IR3:a630,IMAGE_PITCH_ALIGNMENT=64"
-        run: |
-          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_gemm | grep image_load
-          python -m pytest -n=auto test/backend/test_ops.py --durations=20

docs/kimi_k3_mi350.md

@@ -1,224 +0,0 @@
-# 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.

examples/convert_kimi_linear.py

@@ -1,9 +0,0 @@
-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)

examples/kimi_k3_prepare.py

@@ -1,27 +0,0 @@
-#!/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()

examples/kimi_k3_smoke.py

@@ -1,27 +0,0 @@
-#!/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()

examples/mlperf/models/gpt_oss.py

@@ -12,7 +12,7 @@ from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb
 from extra.llama_kernels.rmsnorm import rmsnorm
-from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8
+from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8, asm_gemm, can_use_asm_gemm
 from extra.gemm.moe_gemm import grouped_mx_gemm
 from extra.gemm.moe_routing import route, dispatch, combine
 
@@ -182,12 +182,12 @@ class GPTOSS:
     xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
     xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)  # (B,N,H,D)/(B,N,KV,D)
 
-    if sliding:
-      attn = self._sliding_attention(xq, xk, xv, sinks)
-    elif getenv("HK_FLASH_ATTENTION"):
+    if getenv("HK_FLASH_ATTENTION"):
       from extra.thunder.amd.fa import flash_attention
-      attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks)
+      attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks, window=self.sliding_window if sliding else 0)
       attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
+    elif sliding:
+      attn = self._sliding_attention(xq, xk, xv, sinks)
     else:
       xqm = xq.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep, self.head_dim).permute(0, 2, 3, 1, 4)
       xkm, xvm = xk.permute(0, 2, 1, 3).unsqueeze(2), xv.permute(0, 2, 1, 3).unsqueeze(2)
@@ -263,7 +263,11 @@ class GPTOSS:
                         w_down=self.w_down[i], w_down_scale=self.w_down_scale[i], w_down_bias=self.w_down_bias[i])
       h, *_ = self.run_layer(h, freqs_cis, mask_full, i % 2 == 0, attn_kwargs, ffn_kwargs, save=save)
 
-    logits = self.norm(h) @ self.output.T
+    h_normed = self.norm(h)
+    pad = (-self.dim) % 256
+    h_padded, w_padded = h_normed.pad((None, None, (0, pad))), self.output.pad(((0, 0), (0, pad)))
+    if ASM_GEMM and can_use_asm_gemm(h_padded, w_padded.T): logits = asm_gemm(h_padded, w_padded.T)
+    else: logits = h_normed @ self.output.T
     return logits
 
 def _get_pads(uop:UOp) -> list[UOp]:

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh

@@ -26,7 +26,7 @@ export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
 export SPLIT_W13=${SPLIT_W13:-0}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
 
-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="float32"
 export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
@@ -46,7 +46,7 @@ export DATA_SEED=${DATA_SEED:-5760}
 export JITBEAM=${JITBEAM:-3}
 export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
 
-export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
+export FAKEDATA=${FAKEDATA:-$([[ "$DEV" == NULL:* ]] && echo 1 || echo 0)} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
   export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
 fi

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh

@@ -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
@@ -26,7 +26,7 @@ export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
 export SPLIT_W13=${SPLIT_W13:-0}
 export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
 
-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="float32"
 export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/gpt_oss/implementations/tinybox_8xMI350X/dev_beam.sh

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

examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/gpt_oss/implementations/tinybox_8xMI350X/dev_run.sh

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

extra/benchmark_kimi.py

@@ -1,79 +0,0 @@
-#!/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()

extra/benchmark_kimi_k3.py

@@ -1,83 +0,0 @@
-#!/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()

extra/benchmark_kimi_k3_fake.py

@@ -1,84 +0,0 @@
-#!/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()

extra/export_model.py

@@ -35,7 +35,7 @@ def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], Li
     return name
 
   for call in iter_kernel_calls(linear):
-    arg_uops = [b for b in call.src[1:] if b.op is not Ops.BIND]
+    arg_uops = [b for b in call.src[1:] if not b.is_bound_var]
     prg = to_program(call.src[0], Device[arg_uops[0].device].renderer)
     info = prg.arg
     functions[info.function_name] = prg.src[2].arg

extra/gemm/cdna_asm_gemm.py

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

extra/gemm/triton_nv_matmul.py

@@ -79,7 +79,7 @@ if __name__ == "__main__":
   linear, var_vals = C.linear_with_vars()
   last_call = linear.src[-1]
   ast = last_call.src[0]
-  bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]
+  bufs = [s.buffer for s in last_call.src[1:] if not s.is_bound_var]
 
   src = compiled.asm["ptx"]
   # specify the shared memory here so we don't need to do it dynamically

extra/hcq2/ops_amd2.py

@@ -3,7 +3,7 @@ from typing import cast, Any, Callable
 import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools, collections, atexit
 assert sys.platform != 'win32'
 from dataclasses import dataclass
-from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, encode_kernargs_clike, make_cmdbuf
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, encode_kernargs_clike, make_cmdbuf
 from tinygrad.runtime.support.hcq2 import make_binary_patch
 from tinygrad.uop.ops import sint, UOp
 from tinygrad.device import Compiled, BufferSpec, Buffer, Device
@@ -288,14 +288,16 @@ def amd_build_program(prg:UOp) -> UOp:
 
 class AMDAllocator(HCQAllocator['AMDDevice']):
   def __init__(self, dev:AMDDevice):
-    super().__init__(dev, supports_copy_from_disk=dev.has_sdma_queue, supports_transfer=dev.has_sdma_queue and not dev.is_usb())
+    super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb())
 
-  def _alloc(self, size:int, options:BufferSpec) -> HCQ2Buffer:
-    return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_sdma_queue)
+  def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
+    return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_copy_queue)
 
   def _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque)
 
-  def _do_map(self, buf:HCQ2Buffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+  def _do_map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+
+  def _do_unmap(self, buf:HCQBuffer): self.dev.iface.unmap(buf)
 
 @dataclass
 class AMDQueueDesc:
@@ -388,15 +390,24 @@ class KFDIface:
     return hcqbuf
 
   def free(self, mem):
+    self._unmap(mem)
+    if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
+    kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
+
+  def unmap(self, mem):
+    self._unmap(mem)
+    if getattr(mem, '_owns_kfd_handle', False): kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
+
+  def _unmap(self, mem):
     gpus = (ctypes.c_int32 * 1)(self.gpu_id)
     stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1)
     assert stm.n_success == 1
-    if mem.owner == self.dev:
-      if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
-      kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
 
   def map(self, mem):
-    if mem.owner is not None and mem.owner._is_cpu(): return self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
+    if mem.owner is not None and mem.owner._is_cpu():
+      mapped = self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
+      mapped._owns_kfd_handle = True
+      return mapped
 
     c_gpus = (ctypes.c_int32 * 1)(self.gpu_id)
     stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(c_gpus), n_devices=1)
@@ -468,6 +479,7 @@ class PCIIface(PCIIfaceBase):
 
   def require_profile_mode(self): return True
   def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return True # TODO: account for WGP disablement on some asics.
+  def unmap(self, mem): self.free(mem)
 
   def _compute_props(self):
     self.ip_versions = self.dev_impl.ip_ver
@@ -581,7 +593,7 @@ class AMDDevice(HCQ2Compiled):
 
     self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
     self.sdma_queues:dict = {}
-    self.has_sdma_queue = True # self.sdma_queue(0) is not None, TODO: think of this
+    self.has_copy_queue = not getenv("AMD_DISABLE_SDMA")
 
     super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, can_recover=self.is_am(), arch=self.arch)
 

extra/runbook_digitalocean_mi350x.md

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

extra/thunder/amd/fa.py

@@ -110,7 +110,49 @@ def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
   return _sharded_empty(ref.shape, ref, axis)
 
 @functools.cache
-def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink):
+def _windowed_lse(xq:Tensor, xk:Tensor, sinks, W:int) -> Tensor:
+  B, N, H, hd = xq.shape
+  H_KV = xk.shape[2]; R = H // H_KV; nb = N // W; sm = hd ** -0.5
+  q = xq.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
+  k = xk.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
+  k_prev = k.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
+  sc_d = (q @ k.transpose(-1, -2)) * sm
+  sc_p = (q @ k_prev.transpose(-1, -2)) * sm
+  li, lj = Tensor.arange(W).reshape(W, 1), Tensor.arange(W).reshape(1, W)
+  pv = (Tensor.arange(nb).reshape(nb, 1, 1) >= 1)
+  sc_d = (lj <= li).where(sc_d, -float("inf"))
+  sc_p = ((li < lj) & pv).where(sc_p, -float("inf"))
+  m = sc_d.max(-1, keepdim=True).maximum(sc_p.max(-1, keepdim=True))
+  if sinks is not None: m = m.maximum(sinks.reshape(1, H_KV, R, 1, 1, 1).float())
+  denom = (sc_d - m).exp().sum(-1, keepdim=True) + (sc_p - m).exp().sum(-1, keepdim=True)
+  if sinks is not None: denom = denom + (sinks.reshape(1, H_KV, R, 1, 1, 1).float() - m).exp()
+  return (m + denom.log()).reshape(B, H, N).unsqueeze(2)  # (B, H, 1, N), matches saved l_vec
+
+def _windowed_delta(xq:Tensor, xk:Tensor, xv:Tensor, do:Tensor, sinks, W:int) -> Tensor:
+  B, N, H, hd = xq.shape
+  H_KV = xk.shape[2]; R = H // H_KV; nb = N // W; sm = hd ** -0.5
+  q = xq.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
+  k = xk.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
+  v = xv.permute(0, 2, 1, 3).reshape(B, H_KV, 1, nb, W, hd).float()
+  dob = do.reshape(B, N, H_KV, R, hd).permute(0, 2, 3, 1, 4).reshape(B, H_KV, R, nb, W, hd).float()
+  k_prev = k.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
+  v_prev = v.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb]
+  sc_d = (q @ k.transpose(-1, -2)) * sm
+  sc_p = (q @ k_prev.transpose(-1, -2)) * sm
+  li, lj = Tensor.arange(W).reshape(W, 1), Tensor.arange(W).reshape(1, W)
+  pv = (Tensor.arange(nb).reshape(nb, 1, 1) >= 1)
+  sc_d = (lj <= li).where(sc_d, -float("inf"))
+  sc_p = ((li < lj) & pv).where(sc_p, -float("inf"))
+  m = sc_d.max(-1, keepdim=True).maximum(sc_p.max(-1, keepdim=True))
+  if sinks is not None: m = m.maximum(sinks.reshape(1, H_KV, R, 1, 1, 1).float())
+  e_d, e_p = (sc_d - m).exp(), (sc_p - m).exp()
+  denom = e_d.sum(-1, keepdim=True) + e_p.sum(-1, keepdim=True)
+  if sinks is not None: denom = denom + (sinks.reshape(1, H_KV, R, 1, 1, 1).float() - m).exp()
+  o = ((e_d / denom) @ v) + ((e_p / denom) @ v_prev)
+  delta = (dob * o).sum(-1)
+  return delta.reshape(B, H, N).unsqueeze(2)
+
+def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink, window=0):
   def grad(dou:UOp, ker:UOp) -> tuple:
     do = Tensor(dou, device=dou.device)
     attn = Tensor(ker.src[1].after(ker), device=ker.src[1].device)
@@ -118,6 +160,8 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
     xq = Tensor(ker.src[3], device=ker.src[3].device)
     xk = Tensor(ker.src[4], device=ker.src[4].device)
     xv = Tensor(ker.src[5], device=ker.src[5].device)
+    if window:
+      l_vec = _windowed_lse(xq, xk, Tensor(ker.src[6], device=ker.src[6].device) if has_sink else None, window)
 
     dq = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t)
     GROUP_SIZE = H_local // H_KV_local
@@ -128,8 +172,10 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
     # delta_vec = (do * attn).sum(-1, dtype=dtypes.float32).transpose(1, 2).unsqueeze(-2).detach()
     delta_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
     delta_vec, dq = Tensor.custom_kernel(delta_vec, dq, attn, do, fxn=functools.partial(custom_fa_backward_pre, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:2]
+    if window:
+      delta_vec = _windowed_delta(xq, xk, xv, do, Tensor(ker.src[6], device=ker.src[6].device) if has_sink else None, window)
 
-    dq, dk_partial, dv_partial = Tensor.custom_kernel(dq, dk_partial, dv_partial, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:3]
+    dq, dk_partial, dv_partial = Tensor.custom_kernel(dq, dk_partial, dv_partial, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, window=window))[:3]
 
     if D == 64:
       dq = dq.reshape(B, H, N//16, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2).permute(0, 1, 2, 8, 9, 10, 11, 3, 4, 6, 7, 5, 12).reshape(B, H, N, D).transpose(1, 2)
@@ -149,7 +195,7 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
   return grad
 
 # TODO: remove write_flat once scheduler can remove reshapes between custom_kernel. TestCustomKernel.test_simple_reshape
-def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False, write_flat:bool=False, sinks:Tensor|None=None):
+def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False, write_flat:bool=False, sinks:Tensor|None=None, window:int=0):
   assert attn_mask is None, "attn_mask not supported"
   assert is_causal, "only causal attention supported"
 
@@ -176,18 +222,18 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
   attn = _sharded_empty((B, N, H * D), xq, axis=shard_axis) if write_flat else _sharded_empty_like(xq, axis=shard_axis)
   l_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
 
-  grad = _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink)
+  grad = _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink, window=window)
 
   fwd_inputs = (attn, l_vec, xq, xk, xv) + ((sinks,) if has_sink else ())
-  attn, l_vec = Tensor.custom_kernel(*fwd_inputs, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, has_sink=has_sink), grad_fxn=grad)[:2]
+  attn, l_vec = Tensor.custom_kernel(*fwd_inputs, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, has_sink=has_sink, window=window), grad_fxn=grad)[:2]
 
   return attn, attn, l_vec
 
 @functools.cache
-def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None, *, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, has_sink:bool=True):
+def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None, *, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, has_sink:bool=True, window:int=0):
   code = (pathlib.Path(__file__).parent / "fa_fwd_causal.cpp").read_text()
   compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math",
-                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DATTN_SINK={int(has_sink)}"]
+                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DATTN_SINK={int(has_sink)}", f"-DWINDOW={window}"]
 
   Q_BLOCK_SIZE = 32
   NUM_WARPS = 8
@@ -247,10 +293,10 @@ def custom_fa_backward_pre(delta_vec:UOp, dq:UOp, o:UOp, do:UOp, device:str, arc
              src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
 
 @functools.cache
-def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_vec:UOp, delta_vec:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
+def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_vec:UOp, delta_vec:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, window:int=0):
   code = (pathlib.Path(__file__).parent / "fa_bwd_causal.cpp").read_text()
   compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math",
-                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}"]
+                  f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DWINDOW={window}"]
 
   BLOCK_SIZE_KV = 256
   GROUP_SIZE = H // H_KV

extra/thunder/amd/matvec_bf16_splitk.cpp

@@ -1,74 +0,0 @@
-#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);
-  }
-}

test/backend/test_custom_kernel.py

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

test/backend/test_jit.py

@@ -360,7 +360,7 @@ class TestJitGraphSplit(unittest.TestCase):
     self.expect(f, inp, inp_cpu,
       graph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()],
       multigraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()],
-      hcqgraph=[self.ji_graph(4)])
+      hcqgraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()]) # cpu is hcq2 now, it does not join hcq graphs
 
   def test_jit_cpu_several(self):
     if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
@@ -377,9 +377,9 @@ class TestJitGraphSplit(unittest.TestCase):
     inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
     inp_cpu = Tensor.randn(10, 10, device="CPU").realize()
     self.expect(f, inp, inp_cpu,
-      graph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
-      multigraph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
-      hcqgraph=[self.ji_graph(5)])
+      graph=[self.ji_graph(2), self.ji_comp(), self.ji_comp(), self.ji_comp()],
+      multigraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp(), self.ji_comp()],
+      hcqgraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp(), self.ji_comp()])
 
   def test_jit_multidev(self):
     if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")

test/backend/test_linearizer.py

@@ -30,7 +30,7 @@ class TestLinearizer(unittest.TestCase):
     c = ((a.shrink(((0, 2),)) - a.shrink(((2, 4),))) - (b.shrink(((0, 2),)) - b.shrink(((2, 4),))))
     linear = c.schedule_linear()
     run_linear(linear)
-    rawbufs = [s.buffer for s in linear.src[-1].src[1:] if s.op is not Ops.BIND]
+    rawbufs = [s.buffer for s in linear.src[-1].src[1:] if not s.is_bound_var]
     assert len(rawbufs) == 3 and set(rawbufs[1:]) == {a.uop.base.realized, b.uop.base.realized}
     np_c = (np_a[:2] - np_a[2:]) - (np_b[:2] - np_b[2:])
     np.testing.assert_allclose(np_c, c.numpy(), atol=1e-4, rtol=1e-4)
@@ -411,7 +411,7 @@ def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:
   last_call = linear.src[-1]
   ast = last_call.src[0]
   assert ast.op is Ops.SINK, f"helper_realized_ast expects a SINK {last_call}"
-  last_bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]
+  last_bufs = [s.buffer for s in last_call.src[1:] if not s.is_bound_var]
   # now all input buffers in last_call should be realized
   # create fresh buffers for the outputs
   bufs = [Buffer(x.device, x.size, x.dtype).allocate() if i < len(ast.src) else x for i,x in enumerate(last_bufs)]
@@ -437,7 +437,7 @@ def reset_bufs(bufs:list[Buffer]):
   for buf in bufs: buf.copy_from(Buffer("PYTHON", buf.size, buf.dtype, opaque=memoryview(bytearray(buf.nbytes))))
 
 def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[],
-                               apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[]):
+                               apply_tc=False, atol=1e-4, rtol=1e-4, color_sizes=[], wanna_output=[], check_default_opt=True):
   outbufs = real_bufs[:len(realized_ast.src)]
   wanna_output = [np.array(x).flatten() for x in wanna_output]
   buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in real_bufs]
@@ -459,9 +459,7 @@ def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[]
     for buf,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(buf, want, atol=atol, rtol=rtol)
 
   # Check correctness of handcoded optimiztions.
-  reset_bufs(outbufs)
-  run_prg(opts=None)
-  for buf,want in zip(copyout_outputs(outbufs), wanna_output): np.testing.assert_allclose(buf, want, atol=atol, rtol=rtol)
+  if check_default_opt: check_opt(None)
   for x in opts: # Check custom transformations if any.
     check_opt(([Opt(OptOps.TC, 0, (TC_SELECT.value, TC_OPT.value, 1))] if apply_tc else [])+x)
 

test/backend/test_multitensor.py

@@ -1,9 +1,9 @@
 import unittest, random
 from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variable
-from tinygrad.uop.ops import Ops, UOp, AxisType
+from tinygrad.uop.ops import Ops, UOp, AxisType, graph_rewrite
 from tinygrad.helpers import getenv, prod, Context
 from tinygrad.nn.state import get_parameters
-from tinygrad.engine.realize import run_linear, compile_linear
+from tinygrad.engine.realize import run_linear, compile_linear, pm_beam, pm_compile
 import numpy as np
 from hypothesis import given, strategies as strat, settings
 from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule, assert_kernel_count
@@ -79,9 +79,9 @@ class TestMultiTensor(unittest.TestCase):
   def test_shard_beam(self):
     cpu_2 = ("CPU:1", "CPU:2")
     src = Tensor.ones(16).shard(cpu_2, 0).realize()
-    pad = src.to(cpu_2[::-1]).schedule_linear().src[0]
-    with Context(BEAM=1, IGNORE_BEAM_CACHE=1): prg = compile_linear(UOp(Ops.LINEAR, src=(pad,))).src[0].src[0]
-    self.assertNotEqual(prg.src[0].arg.applied_opts, ())
+    lin = UOp(Ops.LINEAR, src=(src.to(cpu_2[::-1]).schedule_linear().src[0],))
+    with Context(BEAM=1, IGNORE_BEAM_CACHE=1): call = graph_rewrite(graph_rewrite(lin, pm_beam, ctx=1, walk=True), pm_compile, walk=True).src[0]
+    self.assertNotEqual(call.src[0].src[0].arg.applied_opts, ())
 
   def test_shard_same_device(self):
     X = Tensor.ones(256).contiguous().realize()

test/backend/test_ops.py

@@ -720,10 +720,11 @@ class TestOps(unittest.TestCase):
       return torch.autograd.grad(t ** c, t)[0].item()
     for x in [-math.inf, 0, 1, math.inf]:
       for c in [-1, 0, 0.3, 1, 2]:
-        tiny_out = get_tiny_gradient(x, c)
         torch_out = get_torch_gradient(x, c)
+        # the pow backward routes through exp2/log2, whose 0/inf behavior is undefined on WEBGPU
+        if Device.DEFAULT == "WEBGPU" and not math.isfinite(torch_out): continue
+        tiny_out = get_tiny_gradient(x, c)
         if math.isnan(tiny_out):
-          if Device.DEFAULT == "WEBGPU": continue # TODO: WEBGPU issue with nan
           assert math.isnan(torch_out)
         else:
           self.assertAlmostEqual(tiny_out, torch_out, msg=f"{x}, {c}")
@@ -749,6 +750,7 @@ class TestOps(unittest.TestCase):
   def test_exp2_log2_zero_times_negative(self):
     # gallivm's exp2/log2 have "undefined behavior with infs, 0s and nans", so exp2(log2(0)*y) returns 0 instead of inf
     helper_test_op(None, lambda x,y: (x.log2()*y).exp2(), lambda x,y: (x.log2()*y).exp2(), vals=[[0.0], [-0.7]], forward_only=True)
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "pow at 0 routes through exp2/log2, whose 0/inf behavior is undefined on WEBGPU")
   def test_pow_zero_const(self):
     helper_test_op(None, lambda x: x**0.3, vals=[[0.0]])
     helper_test_op(None, lambda x: x**0.0, vals=[[0.0]])

test/backend/test_pickle.py

@@ -2,7 +2,7 @@ import unittest, pickle, types, tracemalloc
 import numpy as np
 from tinygrad import Tensor, Device, TinyJit, Variable, dtypes
 from tinygrad.helpers import GlobalCounters, ContextVar, Context, DEV
-from tinygrad.uop.ops import PatternMatcher, UPat, UOp
+from tinygrad.uop.ops import PatternMatcher, UPat, UOp, deconstruct_function
 
 class TestPickle(unittest.TestCase):
   def test_pickle_code_object(self):
@@ -11,6 +11,11 @@ class TestPickle(unittest.TestCase):
     fxn = types.FunctionType(pickle.loads(code_str), globals())
     self.assertEqual(fxn(2), 4)
 
+  def test_deconstruct_function_nested_comprehension(self):
+    # pre PEP 709, each comprehension is its own code object, so dtypes here is referenced two code objects deep
+    def fxn(): return [[dtypes.int for _ in range(2)] for _ in range(2)]
+    self.assertEqual(types.FunctionType(*deconstruct_function(fxn))(), fxn())
+
   def test_pickle_pattern_matcher(self):
     pm = PatternMatcher([(UPat.cvar('x'), lambda x: x*2)])
     sink = UOp.const(2)

test/device/test_hcq2.py

@@ -0,0 +1,14 @@
+import unittest, numpy as np
+from unittest.mock import patch
+from tinygrad import Device, Tensor
+from tinygrad.helpers import getenv
+from tinygrad.runtime.support.hcq2 import HCQ_DEVS, all_devices_in
+
+@unittest.skipUnless(getenv("HCQ2") and all_devices_in(Device.DEFAULT, HCQ_DEVS), "hcq2 device required")
+class TestHCQ2(unittest.TestCase):
+  def test_copy_without_copy_queue(self):
+    with patch.object(Device[Device.DEFAULT], "has_copy_queue", False):
+      np.testing.assert_equal(Tensor(np.arange(61, dtype=np.float32)).to(Device.DEFAULT).contiguous().realize().numpy(), np.arange(61))
+
+if __name__ == "__main__":
+  unittest.main()

test/external/external_benchmark_all2all.py

@@ -0,0 +1,17 @@
+from tinygrad import Device, Tensor, TinyJit, dtypes
+from tinygrad.helpers import Timing, Context
+
+GPUS, DEPTH, SZ = 8, 4, 128 * 2**20
+WARMUP, ITERS = 3, 5
+devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(GPUS))
+bufs = tuple(Tensor.empty(SZ, dtype=dtypes.uint8, device=dev).contiguous().realize() for _ in range(DEPTH) for dev in devs)
+
+@TinyJit
+def all_to_all(*srcs:Tensor): return Tensor.realize(*(src.to(dst) for i,src in enumerate(srcs) for j,dst in enumerate(devs) if i % GPUS != j))
+
+if __name__ == "__main__":
+  with Context(ALL2ALL=1, JIT_BATCH_SIZE=0):
+    for i in range(-WARMUP, ITERS):
+      with Timing("ALL2ALL ", lambda ns: f" {SZ*GPUS*(GPUS-1)*DEPTH/ns:.2f} GB/s", enabled=i>=0):
+        all_to_all(*bufs)
+        for dev in devs: Device[dev].synchronize()

test/external/external_test_schedule_scaling.py

@@ -1,5 +1,5 @@
-import unittest, time
-from tinygrad import Tensor
+import unittest, time, itertools
+from tinygrad import Tensor, Context
 
 class TestScheduleScaling(unittest.TestCase):
   """Test that .schedule() scales linearly with graph size (no O(n^2) behavior)."""
@@ -130,5 +130,18 @@ class TestScheduleScaling(unittest.TestCase):
       return parts[0].cat(*parts[1:])
     self._assert_linear(concat_chain)
 
+  @Context(DEV="NULL:HIP:gfx1100")
+  def test_custom_kernel_assign_scaling(self):
+    from tinygrad.uop.ops import UOp, Ops, KernelInfo
+    from tinygrad.runtime.autogen.amd.rdna3.ins import s_nop
+    count = itertools.count(0)
+    def custom_kernel_assign(n):
+      def custom_asm(out):
+        return UOp(Ops.PROGRAM, src=(UOp.sink(out, arg=KernelInfo(f"fxn_{next(count)}")),
+                                     UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=s_nop(i)) for i in range(n*8)))))
+      call = Tensor.custom_kernel(Tensor.empty(1), fxn=custom_asm)[0]
+      return Tensor.cat(*[Tensor.empty(1).assign(call+i) for i in range(n)])
+    self._assert_linear(custom_kernel_assign, n_small=50, n_large=500)
+
 if __name__ == '__main__':
   unittest.main(verbosity=2)

test/helpers.py

@@ -48,6 +48,8 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
   else:
     assert isinstance(t, UOp), f"can't schedule {t}"
     linear, var_vals = Tensor(t).linear_with_vars()
+  # test compiling the linear
+  compile_linear(linear)
   kernel_cnt = sum((len(call.device) if isinstance(call.device, tuple) else 1)
                    for call in linear.src if call.src[0].op is Ops.SINK or not filter_sink)
   if kernel_cnt != allowed:
@@ -57,8 +59,6 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
         print("kernel", i+1)
         print(call.src[0])
     raise KernelCountException(allowed, kernel_cnt)
-  # test compiling the linear
-  compile_linear(linear)
   return linear, var_vals
 
 def assert_kernel_count(expected:int):
@@ -86,7 +86,7 @@ def assert_jit_cache_len(fxn, expected_len):
   if linear is None or not linear.src:
     if expected_len != 0: raise KernelCountException(expected_len, 0)
     return
-  if expected_len and all(call_is_hcq(call) for call in linear.src): expected_len = 3 # HCQ2: merged same-queue calls + finalizer + bumps
+  if expected_len and all(call_is_hcq(call) for call in linear.src): expected_len = 4 # HCQ2: fence + reset + merged same-queue calls + finalizer
   if call_is_graph(linear.src[0]):
     if len(linear.src) != 1: raise KernelCountException(1, len(linear.src))
     inner = linear.src[0].src[0].src[0]  # LINEAR UOp inside CUSTOM_FUNCTION

test/mockgpu/amd/emu.py

@@ -260,19 +260,6 @@ def _cond(cond, if_true, if_false):
 def _cond_hi16(cond, val: UOp) -> UOp: return _cond(cond, _hi16(val), val)
 def _apply_opsel(val: UOp, sel_bit: int, opsel: int) -> UOp: return _hi16(val) if opsel & (1 << sel_bit) else val
 
-def _set_lane_bit(old: UOp, lane: UOp, val: UOp, exec_mask: UOp) -> UOp:
-  """Set/clear a single bit in a mask based on lane index, respecting exec mask."""
-  if old.dtype in (dtypes.uint64, dtypes.int64):
-    dt = dtypes.uint64
-    mask = UOp.const(1, dt) << lane.cast(dt)
-    new_bit = _to_u32(val).cast(dt) << lane.cast(dt)
-    cleared = old.cast(dt) & (mask ^ UOp.const(0xFFFFFFFFFFFFFFFF, dt))
-    return _lane_active(exec_mask, lane).where(cleared | new_bit, old.cast(dt))
-  mask = _c(1) << lane.cast(dtypes.uint32)
-  new_bit = _to_u32(val) << lane.cast(dtypes.uint32)
-  cleared = old & (mask ^ _c(MASK32))
-  return _lane_active(exec_mask, lane).where(cleared | new_bit, old)
-
 def _val_to_u32(val: UOp) -> UOp:
   """Convert any value to uint32 for storage (bitcast floats, cast ints)."""
   if val.dtype == dtypes.uint32: return val
@@ -532,6 +519,19 @@ class _Ctx:
       return [self.wsgpr_dyn(reg, lo), self.wsgpr_dyn(reg + _c(1), hi)]
     return [self.wsgpr_dyn(reg, val)]
 
+  def wmask_lane_bit(self, reg: UOp, lane: UOp, val: UOp, exec_mask: UOp) -> list[UOp]:
+    """Set/clear bit `lane` of the mask at `reg` from val for exec-active lanes, preserving memory for inactive lanes"""
+    active, bit = _lane_active(exec_mask, lane), _to_u32(val)
+    if self.wave_size <= 32:
+      old = self.rsgpr_dyn(reg)
+      mask = _c(1) << lane.cast(dtypes.uint32)
+      return [self.wsgpr_dyn(reg, active.where((old & (mask ^ _c(MASK32))) | (bit << lane.cast(dtypes.uint32)), old))]
+    off = (lane & _c(31, dtypes.int)).cast(dtypes.uint32)
+    mask = _c(1) << off
+    def half(old: UOp, sel: UOp) -> UOp: return sel.where(active.where((old & (mask ^ _c(MASK32))) | (bit << off), old), old)
+    return [self.wsgpr_dyn(reg, half(self.rsgpr_dyn(reg), lane < _c(32, dtypes.int))),
+            self.wsgpr_dyn(reg + _c(1), half(self.rsgpr_dyn(reg + _c(1)), _c(32, dtypes.int) <= lane))]
+
   def rmask(self, reg: UOp) -> UOp:
     """Read a lane mask (VCC/EXEC). Combines lo/hi for wave64."""
     if self.wave_size > 32: return _u64(self.rsgpr_dyn(reg), self.rsgpr_dyn(reg + _c(1)))
@@ -718,9 +718,7 @@ class _Ctx:
         raw_stores.append(('vgpr_direct', self.vgpr.index(val[0].valid(active)).store(new_val)))
         continue
       if 'D0' in dest and '[laneId]' in dest:
-        old_vcc = self.rmask(_c(VCC_LO.offset))
-        new_vcc = _set_lane_bit(old_vcc, lane, val, exec_mask)
-        raw_stores.extend([('vcc', s) for s in self.wmask(_c(VCC_LO.offset), new_vcc)])
+        raw_stores.extend([('vcc', s) for s in self.wmask_lane_bit(_c(VCC_LO.offset), lane, val, exec_mask)])
       elif dest.startswith('D0'):
         dest_suffix = re.match(r'D0\.(\w+)', dest)
         if dest_suffix is not None:
@@ -1039,13 +1037,11 @@ def _compile_sdwa(inst: irc.VOP1_SDWA | irc.VOP2_SDWA | irc.VOP2_SDWA_SDST | irc
         result = _sdwa_write(old, result, dst_sel, dst_unused)
       stores.append(ctx.wvgpr_dyn(vdst_reg, lane, result, exec_mask))
     elif dest.startswith('VCC'):
-      old_vcc = ctx.rmask(_c(VCC_LO.offset))
-      stores.extend(ctx.wmask(_c(VCC_LO.offset), _set_lane_bit(old_vcc, lane, val, exec_mask)))
+      stores.extend(ctx.wmask_lane_bit(_c(VCC_LO.offset), lane, val, exec_mask))
   if vcc_val is not None:
     # Initialize sdst to 0 before lane loop (old value may be unrelated data), then set lane bits in loop
     init_stores = [ctx.wsgpr_dyn(sdst_off, _c(0)), ctx.wsgpr_dyn(sdst_off + _c(1), _c(0))]
-    old_sdst = ctx.rmask(sdst_off)
-    stores.extend(ctx.wmask(sdst_off, _set_lane_bit(old_sdst, lane, vcc_val, exec_mask)))
+    stores.extend(ctx.wmask_lane_bit(sdst_off, lane, vcc_val, exec_mask))
     if stores:
       return UOp.sink(*init_stores, UOp.sink(*stores).end(lane), *ctx.inc_pc())
     return UOp.sink(*init_stores, *ctx.inc_pc())

test/models/test_whisper.py

@@ -74,6 +74,7 @@ class TestWhisper(unittest.TestCase):
         err
       )
 
+  @slow
   def test_transcribe_file1(self):
     self.assertEqual(transcribe_file(self.model, self.enc, TEST_FILE_1),  TRANSCRIPTION_1)
 
@@ -89,6 +90,7 @@ class TestWhisper(unittest.TestCase):
     self.assertEqual(TRANSCRIPTION_1,  transcriptions[0])
     self.assertEqual(TRANSCRIPTION_2,  transcriptions[1])
 
+  @slow
   def test_transcribe_batch21(self):
     waveforms = [load_file_waveform(TEST_FILE_2), load_file_waveform(TEST_FILE_1)]
     transcriptions = transcribe_waveform(self.model, self.enc, waveforms)

test/null/test_const_folding.py

@@ -1,6 +1,6 @@
 import unittest, itertools, math
 from tinygrad import Tensor, dtypes, Context
-from tinygrad.dtype import DType, ConstType, truncate
+from tinygrad.dtype import DType, ConstType
 from tinygrad.uop.ops import Ops, UOp
 from test.helpers import full_rewrite
 import numpy as np
@@ -51,16 +51,13 @@ 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())
+  def test_single_rounding_log10_backward(self):
+    # log10 backward folds log10(2)/log(2) = 1/log(10) in one rounding, not the double-rounded 1/float32(log(10))
+    x = Tensor([1.0, 2.0, 3.0])
+    ast = next(s.src[0] for s in x.log10().sum().gradient(x)[0].schedule_linear().src if s.src[0].op is Ops.SINK)
+    const = next(u.arg for u in full_rewrite(ast).toposort() if u.op is Ops.CONST and u.dtype is dtypes.float32)
+    # correctly rounded: within half a float32 ulp of the exact value (folding at float32 lands 0.66 ulp off)
+    self.assertLess(abs(const - 1/math.log(10)), 2**-26)
 
 class TestBinaryOpsConstFolding(unittest.TestCase):
   def test_add_literal_zero(self):

test/null/test_graph_rewrite.py

@@ -143,13 +143,13 @@ class TestModuloAndDivisionFolding(unittest.TestCase):
     self.assertIs(apply_rewrite(x_var_uop.cast(dtypes.weakint) % 10).render(simplify=False), x_var_uop.render(simplify=False))
 
   def test_full_graph_rewrite_division_with_remainder(self):
-    x_var_uop = UOp.variable('x', 7, 9)
+    x_var_uop = UOp.variable('x', 7, 9, param=True)
     optimized_sink = apply_rewrite(x_var_uop // 2)
     for x_value in range(7, 10):
       self.assertEqual(x_value // 2, evaluate_uop(optimized_sink, {'x': x_value}))
 
   def test_full_graph_rewrite_complex_mod_div_expression(self):
-    x_var_uop = UOp.variable('x', 1, 10)
+    x_var_uop = UOp.variable('x', 1, 10, param=True)
     optimized_sink = apply_rewrite(((x_var_uop * 5) % 3) // 2)
     for x_value in range(1, 11):
       original_result = ((x_value * 5) % 3) // 2

test/null/test_kimi.py

@@ -1,30 +0,0 @@
-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()

test/null/test_kimi_k3.py

@@ -1,29 +0,0 @@
-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()

test/null/test_memory_planner.py

@@ -27,10 +27,15 @@ def _make_linear(buffer_lists, copies=None):
     calls.append(UOp(Ops.CALL, src=(src0, *bufs)))
   return UOp(Ops.LINEAR, src=tuple(calls))
 
+def _get_planned_view(buf:UOp) -> tuple[UOp, int, int]|None:
+  view = buf.src[0] if buf.op is Ops.BITCAST else buf
+  if view.op is not Ops.SHRINK or view.src[0].op is not Ops.BUFFER: return None
+  return (arena:=view.src[0]), view.src[1].val * arena.dtype.itemsize, view.src[2].val * arena.dtype.itemsize
+
 def _get_arena(buf, linear, result):
   for orig_si, new_si in zip(linear.src, result.src):
     for orig, new in zip(orig_si.src[1:], new_si.src[1:]):
-      if orig is buf and new.op is Ops.SLICE: return new.src[0]
+      if orig is buf and (planned:=_get_planned_view(new)) is not None: return planned[0]
   return None
 
 def check_assign(buffer_lists, copies=None):
@@ -41,8 +46,8 @@ def check_assign(buffer_lists, copies=None):
   replace_map: dict[int, tuple[UOp, int, int]] = {}
   for orig_si, new_si in zip(linear.src, result.src):
     for orig, new in zip(orig_si.src[1:], new_si.src[1:]):
-      if new.op is Ops.SLICE and id(orig) not in replace_map:
-        replace_map[id(orig)] = (new.src[0], new.src[1].val * new.src[0].dtype.itemsize, new.arg * new.dtype.itemsize)
+      if (planned:=_get_planned_view(new)) is not None and id(orig) not in replace_map:
+        replace_map[id(orig)] = planned
 
   # verify pinned buffers are not planned
   for buf in held_bufs:

test/null/test_simplify_valid_idx.py

@@ -25,7 +25,7 @@ def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UO
   ))
 
 def Special(expr, nmax): return UOp(Ops.SPECIAL, src=(UOp.const(nmax),), arg=expr)
-def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax)
+def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax, param=True)
 def Range(n, nmax): return UOp.range(nmax, n)
 
 class TestValidIdxSimplification(unittest.TestCase):

test/null/test_uop_graph.py

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

test/null/test_uop_symbolic.py

@@ -3,10 +3,10 @@ import unittest, pickle, functools, math
 import z3
 
 from tinygrad.dtype import dtypes, ConstType, DType, Invalid
-from test.helpers import get_uops
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad.uop.spec import spec_shared, type_verify
-from tinygrad.uop.symbolic import sym, pm_fold_cast_const, commutative, pm_simplify_valid, pm_move_where_on_load
+from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load
+from tinygrad.uop.weak import pm_cast_weak
 from tinygrad.uop.validate import uops_to_z3
 
 def check_uop_against_string(self, v:UOp, s:str):
@@ -16,7 +16,8 @@ def check_uop_against_string(self, v:UOp, s:str):
   s_eval = graph_rewrite(s_eval, commutative, name="cannonicalize eval")
   self.assertIs(s_eval, v, f"eval did not match simplified: {s_eval} != {v.render()} for {s}")
 
-def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.weakint): return UOp.variable(name,min_val,max_val,dtype)
+def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.weakint):
+  return UOp.variable(name, min_val, max_val, dtype, param=True)
 def uconst(val): return UOp.const(val)
 def usum(ops): return functools.reduce(lambda x,y: x+y, ops)
 def uand(ops): return functools.reduce(lambda x,y: x*y, ops)
@@ -35,7 +36,7 @@ class TestSymbolic(unittest.TestCase):
     self.assertEqual(solver.check(expr1 != expr2), z3.unsat, "simplified expression not equal to original")
 
   def helper_test_variable(self, v, n, m, s, test_z3:bool=True):
-    v_simplified = graph_rewrite(v, sym+pm_fold_cast_const, name="simplify symbolic uop")
+    v_simplified = graph_rewrite(v, sym+pm_cast_weak, name="simplify symbolic uop")
     if test_z3: self.check_equal_z3(v, v_simplified)
     nmin, nmax = v_simplified.vmin, v_simplified.vmax
     check_uop_against_string(self, v_simplified, s)
@@ -442,7 +443,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(uand([uconst(1), Variable("a", 0, 1)]), 0, 1, "a")
 
   def test_masked_shr_fold(self):
-    x = UOp.variable('x', 0, 255, dtype=dtypes.uint32)
+    x = UOp.variable('x', 0, 255, dtype=dtypes.uint32, param=True)
     self.helper_test_variable((x & -4) >> 2, 0, 63, "(x>>2)")
 
   def test_bool_or_not_tautology(self):
@@ -483,12 +484,12 @@ class TestSymbolic(unittest.TestCase):
 
   def test_div_drop_small_terms(self):
     # from openpilot, shouldnt simplify
-    gidx0 = UOp.variable("gidx0", 0, 10)
-    gidx1 = UOp.variable("gidx1", 0, 10)
-    lidx0 = UOp.variable("lidx0", 0, 1)
-    lidx1 = UOp.variable("lidx1", 0, 1)
-    ridx1005 = UOp.variable("ridx1005", 0, 2)
-    ridx1006 = UOp.variable("ridx1006", 0, 2)
+    gidx0 = UOp.variable("gidx0", 0, 10, param=True)
+    gidx1 = UOp.variable("gidx1", 0, 10, param=True)
+    lidx0 = UOp.variable("lidx0", 0, 1, param=True)
+    lidx1 = UOp.variable("lidx1", 0, 1, param=True)
+    ridx1005 = UOp.variable("ridx1005", 0, 2, param=True)
+    ridx1006 = UOp.variable("ridx1006", 0, 2, param=True)
     self.helper_test_variable((lidx1+((gidx1*18)+(ridx1005*18)+(lidx0*162))+(gidx0*2)+(ridx1006*2)+-40)//18, -3, 20,
       "(gidx1+ridx1005+lidx0*9+(gidx0+ridx1006+7)//9+-3)")
 
@@ -948,6 +949,11 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(cond.where(u0, u1), 0, 1, "((a<2)!=True)")
     self.helper_test_variable(cond.where(u0, u1).where(u0, u1), 0, 1, "(a<2)")
 
+  def test_equivalent_const_max(self):
+    x = Variable("x", -10, 10)
+    self.helper_test_variable((x < 0).where(0, x), 0, 10, "x.maximum(0)")
+    self.helper_test_variable((0 < x).where(x, 0), 0, 10, "x.maximum(0)")
+
   def test_where_combine(self):
     cond = Variable("x", 0, 3) < 2
     a = Variable("a", 0, 3)
@@ -992,7 +998,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(cond.ne(False), 0, 1, "(x<2)")
 
   def test_bitcast_chain(self):
-    a = UOp.variable("a", 0, 3, dtype=dtypes.int32)
+    a = UOp.variable("a", 0, 3, dtype=dtypes.int32, param=True)
     self.assertIs(graph_rewrite(a.bitcast(dtypes.float32).bitcast(a.dtype), sym), a)
 
   def test_negation_in_where(self):
@@ -1007,22 +1013,6 @@ class TestSymbolic(unittest.TestCase):
     b = Variable("b", 0, 3)
     self.helper_test_variable(-a<-b, False, True, "(b<a)")
 
-  def test_where_cast(self):
-    s = Variable("s", 0, 3, dtypes.int)
-    cond = s < 2
-    a = Variable("a", 0, 3, dtypes.int)
-    b = Variable("b", 0, 3, dtypes.int)
-    expr = cond.where(a, b).cast(dtypes.half)
-
-    # TODO: copied from render, render does not support cast
-    glbl = UOp.param(0, dtypes.int, (1,))
-    uops = get_uops(UOp(Ops.STORE, src=(glbl.index(UOp.const(0, dtypes.int)), expr)).sink())
-    rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[1]
-
-    # the vars are now scalar PARAMs
-    pvar = {u.expr: u for u in rewritten_uop.toposort() if u.op is Ops.PARAM}
-    self.assertEqual(rewritten_uop, (pvar['s']<UOp.const(2, dtypes.int)).where(pvar['a'].cast(dtypes.half), pvar['b'].cast(dtypes.half)))
-
   def test_where_merge_branches(self):
     cond1 = Variable("s", 0, 10) < 6
     cond2 = Variable("s", 0, 10) > 2
@@ -1175,7 +1165,7 @@ class TestSymbolicVariables(unittest.TestCase):
     assert (a//4 + a//6).variables() == [a]
 
   def test_variable_min_eq_max_bind_folds(self):
-    b = Variable("x", 1, 1).bind(1)
+    b = UOp.variable("x", 1, 1).bind(1)
     s = b.simplify()
     self.assertEqual(s.op, Ops.CONST)
     self.assertEqual(s.val, 1)
@@ -1369,6 +1359,16 @@ class TestInvalidIndex(unittest.TestCase):
     c2 = UOp.const((1, Invalid, 1, 1))
     self.assertIs((c1+c2).simplify(), UOp.const((2, Invalid, Invalid, Invalid)))
 
+  def test_gated_load_keeps_index_valid(self):
+    # the load executes even on gated-off iterations: gated_given_valid must not erase its mask (PADTO OOB shape)
+    buf = UOp.param(0, dtypes.bool, (17,))
+    ridx = Variable("ridx", 0, 31)
+    cond = ridx < 17
+    load = buf.index(ridx.valid(cond))
+    out = graph_rewrite(cond.where(load.where(uconst(2), uconst(0)), UOp.invalid()), sym)
+    idx = next(u for u in out.toposort() if u.op is Ops.INDEX)
+    self.assertIs(idx.src[1].get_valid(), cond.simplify())
+
 class TestStoreLoadFolding(unittest.TestCase):
   """Tests for store(index, load(index)) -> NOOP rule. This rule matches patterns that EMERGE during simplification."""
   def test_store_load_folding(self):

test/null/test_validate_oob.py

@@ -1,10 +1,12 @@
 import unittest
-from tinygrad import dtypes, Variable
+from tinygrad import dtypes
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import Ops, UOp, AxisType
 from test.helpers import to_uops_list
 
+def Variable(name, nmin, nmax): return UOp.variable(name, nmin, nmax, param=True)
+
 class TestValidateOOB(unittest.TestCase):
   """Test z3 validation of index bounds for different ALU ops and patterns."""
 

test/null/test_viz.py

@@ -305,10 +305,10 @@ class TestVizTree(unittest.TestCase):
 
   def test_tree_view(self):
     with save_viz() as viz:
-      a = UOp.variable("a",0,10)
-      b = UOp.variable("b",0,10)
-      c = UOp.variable("c",0,10)
-      d = UOp.variable("d",0,10)
+      a = UOp.variable("a",0,10,param=True)
+      b = UOp.variable("b",0,10,param=True)
+      c = UOp.variable("c",0,10,param=True)
+      d = UOp.variable("d",0,10,param=True)
       sink = UOp.sink(a+b, c+d)
       def tree_rewrite(): return graph_rewrite(sink, root, name="root")
       tree_rewrite()

test/opt/test_kernel_opts.py

@@ -239,6 +239,15 @@ class TestKernelOpts(unittest.TestCase):
     helper_linearizer_opt(a.sum().exp(), [[Opt(OptOps.PADTO, 0, 32)],])
     helper_linearizer_opt(a.sum(0).exp(), [[Opt(OptOps.PADTO, 1, 32)],])
 
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
+  @unittest.expectedFailure
+  def test_padto_group_full_unroll_sum(self):
+    a = Tensor.ones(2, 28, 4096, dtype=dtypes.bfloat16).realize()
+    out = ((a * 0.5).float().square()).sum(axis=(0, 2))
+    opts_to_apply = [Opt(OptOps.GROUPTOP, 1, 256), Opt(OptOps.PADTO, 3, 32), Opt(OptOps.UNROLL, 2, 0), Opt(OptOps.UPCAST, 0, 7)]
+    helper_linearizer_opt(out, [opts_to_apply], check_default_opt=False)
+
   def test_padto_sum(self):
     N = 18
     # NOTE: this setup prevents 17 * 17 contiguous merged into one dimension

test/opt/test_tensor_cores.py

@@ -79,7 +79,8 @@ class TestTensorCores(unittest.TestCase):
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
   def test_tensor_cores(self):
     for tc in Device[Device.DEFAULT].renderer.tensor_cores:
-      helper_tc_allclose(tc.dims[0], tc.dims[1], tc.dims[2], tc.dtype_in, tc.dtype_out, axis=0, tc_opt=0)
+      with self.subTest(tc=tc):
+        helper_tc_allclose(tc.dims[0], tc.dims[1], tc.dims[2], tc.dtype_in, tc.dtype_out, axis=0, tc_opt=0)
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.tensor_cores, "test requires tensor cores")
   def test_tensor_cores_nested_reduce(self):
@@ -185,10 +186,10 @@ class TestTensorCores(unittest.TestCase):
     # skip fp8 tcs: the unoptimized ALU baseline quantizes products to fp8 (JAX promotion), which legitimately
     # differs from the MFMA path (f32 accumulation), so the baseline-vs-TC numerical gate can't hold for fp8.
     tc = next(tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in not in dtypes.fp8s)
-    x, y = Tensor.rand(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in)
+    x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in)
     r = x.matmul(y, dtype=tc.dtype_out)
     opts = [Opt(OptOps.UNROLL, 0, 2)]
-    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
+    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False)
     for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src):
       if u.op is Ops.WMMA:
         assert u.src[-1].src[0].op != Ops.STORE
@@ -199,10 +200,10 @@ class TestTensorCores(unittest.TestCase):
   @unittest.skipIf(Device.DEFAULT in {"CPU"}, "CPU does not support using a different type for accumulation")
   def test_tensor_cores_unroll_casted_phi(self):
     tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out and tc.dtype_in not in dtypes.fp8s][0]
-    x, y = Tensor.rand(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in)
+    x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in)
     r = x.matmul(y, dtype=tc.dtype_out)
     opts = [Opt(OptOps.UNROLL, 0, 2)]
-    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
+    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False)
     for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src):
       if u.op is Ops.WMMA:
         #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))
@@ -215,10 +216,10 @@ class TestTensorCores(unittest.TestCase):
   def test_tensor_cores_unroll_casted_phi_with_children(self):
     # all STORE children are outside the loop
     tc = [tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in != tc.dtype_out and tc.dtype_in not in dtypes.fp8s][0]
-    x, y = Tensor.rand(64, 64, dtype=tc.dtype_in), Tensor.rand(64, 64, dtype=tc.dtype_in)
+    x, y = Tensor.rand(16, 64, dtype=tc.dtype_in), Tensor.rand(64, 16, dtype=tc.dtype_in)
     r = x.matmul(y, dtype=tc.dtype_out).relu()
     opts = [Opt(OptOps.UNROLL, 0, 2)]
-    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3)
+    ast = helper_linearizer_opt(r, [opts], apply_tc=True, atol=3e-2, rtol=1e-3, check_default_opt=False)
     for u in tuple(to_program(replace_opts(ast, opts), Device[Device.DEFAULT].renderer).src[1].src):
       if u.op is Ops.WMMA:
         #assert u.src[-1].dtype == dtypes.float.vec(prod(tc.thread_local_sizes[2]))

test/unit/test_assign.py

@@ -540,6 +540,10 @@ class TestAssign(unittest.TestCase):
     c = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
     c[0:2].bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000], dtype=dtypes.uint32)).realize()
     np.testing.assert_allclose(c.numpy(), [4.0, 3.0, 3.0, 4.0])
+    # without .realize()
+    a = Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32).realize()
+    a.bitcast(dtypes.uint32).assign(Tensor([0x40800000, 0x40400000, 0x40000000, 0x3f800000], dtype=dtypes.uint32))
+    np.testing.assert_allclose(a.numpy(), [4.0, 3.0, 2.0, 1.0])
 
   def test_assign_bitcast_different_size(self):
     # assign to a shape-changing bitcast view (only works on DISK currently)

test/unit/test_attention.py

@@ -1,11 +1,9 @@
 import unittest
-from types import SimpleNamespace
 import numpy as np
-from tinygrad import Tensor, dtypes, nn
-from tinygrad.llm.kimi import _shard_kimi
+from tinygrad import Tensor, dtypes
 from tinygrad.llm.model import (
   GatedDeltaNetBlock, SSMConfig, TransformerBlock, TransformerConfig,
-  apply_rope as apply_rope_new, iterative_topk, l2norm, precompute_freqs_cis, pairwise_topk,
+  apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk,
 )
 
 def apply_rope(x:Tensor, start_pos:int):
@@ -43,11 +41,6 @@ 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)
 
@@ -197,87 +190,6 @@ 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]]])
@@ -301,13 +213,5 @@ 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()

test/unit/test_dtype_weak.py

@@ -76,6 +76,11 @@ class TestWeakPromotion(unittest.TestCase):
       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_div_sub_operand_kept_weak(self):
+    a = Tensor.empty(4, dtype=dtypes.float32)
+    for t in (a / 1, a - 0):
+      self.assertEqual(t.uop.src[1].dtype, dtypes.weakfloat)
+
   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):
@@ -88,6 +93,13 @@ class TestWeakPromotion(unittest.TestCase):
       out = Tensor(1.0, dtype=dtypes.float32, device="CPU") / denom
       self.assertAlmostEqual(out.item(), 1 / (70000 + 1e-5), places=10)
 
+  def test_stacked_weak_casts_convert_each_kind(self):
+    # each weak cast is a kind conversion: weakint truncates before weakfloat re-lifts (neither is only a marker)
+    x = Tensor([2.5, -3.7], dtype=dtypes.float32, device="CPU")
+    stacked = x.cast(dtypes.weakint).cast(dtypes.weakfloat)
+    self.assertIs(stacked.dtype, dtypes.weakfloat)
+    self.assertEqual(stacked.tolist(), [2.0, -3.0])
+
   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),

test/unit/test_gradient.py

@@ -51,6 +51,10 @@ class TestTensorGradient(unittest.TestCase):
     with self.assertRaises(RuntimeError): x.sum().gradient(x)
     with self.assertRaises(RuntimeError): x.float().sum().gradient(x)
 
+  def test_const_target_raise(self):
+    t = Tensor(2.0)
+    with self.assertRaises(RuntimeError): (t * 2.0).gradient(t)
+
   def test_copy_to_device_gradient(self):
     t = Tensor([1.0, 2, 3]).realize()
     t.to("CPU:1").square().sum().backward()
@@ -100,7 +104,7 @@ class TestTensorGradient(unittest.TestCase):
 
   def test_implicit_broadcast_where_gradient(self):
     # WHERE with a bare ()-shape branch: the scalar's gradient counts the positions where it is selected
-    cond, x, w = Tensor([True, False, True]), Tensor([1.0, 2.0, 3.0]), Tensor(4.0)
+    cond, x, w = Tensor([True, False, True]), Tensor([1.0, 2.0, 3.0]), Tensor(4.0, dtype=dtypes.float32)
     dw = Tensor(cond.uop.alu(Ops.WHERE, x.uop, w.uop)).sum().gradient(w)[0]
     self.assertEqual(dw.shape, ())
     self.assertEqual(dw.item(), 1.0)
@@ -109,7 +113,7 @@ class TestTensorGradient(unittest.TestCase):
 
   def test_implicit_broadcast_alu_gradient(self):
     # MUL with a bare ()-shape src, no EXPAND in the graph
-    x, w = Tensor([1.0, 2.0, 3.0]), Tensor(2.0)
+    x, w = Tensor([1.0, 2.0, 3.0]), Tensor(2.0, dtype=dtypes.float32)
     m = x.uop.alu(Ops.MUL, w.uop)
     self.assertIs(m.src[1], w.uop)
     dw = Tensor(m).sum().gradient(w)[0]
@@ -118,7 +122,7 @@ class TestTensorGradient(unittest.TestCase):
 
   def test_implicit_broadcast_intermediate_accumulation(self):
     # s is used directly and through an implicit broadcast edge, each edge's gradient reduces to s's shape before they sum
-    x, p = Tensor([1.0, 2.0, 3.0]), Tensor(0.5)
+    x, p = Tensor([1.0, 2.0, 3.0]), Tensor(0.5, dtype=dtypes.float32)
     s = p.sin()
     z = Tensor(x.uop.alu(Ops.MUL, s.uop)).sum() + s
     dp = z.gradient(p)[0]

test/unit/test_hcq_graph.py

@@ -25,10 +25,10 @@ class TestHCQUnit(unittest.TestCase):
     cpu_call = UOp(Ops.PROGRAM, src=(UOp.sink(),)).call(UOp.new_buffer("CPU", 1, dtypes.float))
     gpu_devs = [d0]
 
-    # local MMIO: GPU works alone and with CPU in batch (cpu_support=True)
+    # CPU uses HCQ2 and is no longer batched into legacy HCQ graphs.
     assert HCQGraph.supports_uop(gpu_devs, gpu_call) is True
-    assert HCQGraph.supports_uop(gpu_devs, cpu_call) is True
-    assert HCQGraph.supports_uop(gpu_devs + [cpu_dev], gpu_call) is True
+    assert HCQGraph.supports_uop(gpu_devs, cpu_call) is False
+    assert HCQGraph.supports_uop(gpu_devs + [cpu_dev], gpu_call) is False
 
     # USB MMIO: GPU-only still works, but CPU batching must be rejected (cpu_support=False)
     orig_view = d0.timeline_signal.base_buf.view

test/unit/test_llm_cli_k3.py

@@ -1,30 +0,0 @@
-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()

test/unit/test_llm_k3.py

@@ -1,139 +0,0 @@
-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()

test/unit/test_llm_kimi.py

@@ -1,37 +0,0 @@
-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()

test/unit/test_llm_moe.py

@@ -72,20 +72,6 @@ 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

test/unit/test_llm_quant.py

@@ -1,68 +0,0 @@
-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()

test/unit/test_llm_server.py

@@ -1,10 +1,9 @@
 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, parse_kimi_tool_call
+from tinygrad.llm.serve import StreamRouter
 
 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)
@@ -14,29 +13,17 @@ 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, temperature):
-      calls.append((tokens, temperature))
+    def generate(tokens, **kwargs):
+      calls.append(tokens)
       yield from (1, 2)
     with patch.object(model, "generate", generate): model.warmup()
-    self.assertEqual(calls, [([0], 0.0), ([0], 0.0)])
+    self.assertEqual(calls, [[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_warmup_then_generate_with_default_chunk(self):
+    # warmup must not capture JIT graphs that generate()'s default chunk_size then rejects
+    model = Transformer(TEST_CONFIG)
+    model.warmup()
+    self.assertIsInstance(next(model.generate([5, 6, 7, 8])), int)
 
   def test_first_recurrent_generate_before_state_init(self):
     model = Transformer(TEST_CONFIG)
@@ -44,17 +31,6 @@ 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
@@ -68,37 +44,20 @@ 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):
+  def test_recurrent_divergent_prompt_restarts(self):
+    model, calls = Transformer(TEST_CONFIG), []
+    model.has_recurrent_block, model._cached_tokens = True, [1, 2, 9]
+    def mock_call(self, tokens, start_pos, temperature):
       calls.append(start_pos)
       return Tensor([[42]])
-    with patch.object(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)
+    with patch.object(Transformer, '__call__', mock_call): next(model.generate([1, 2, 10, 11]))
+    self.assertEqual(calls[0], V_START_POS.bind(0))
 
   def test_template_starts_reasoning(self):
     router = StreamRouter(reasoning=True)
     self.assertEqual(list(router.route("reasoning</think>answer")),
                      [("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)
@@ -234,6 +193,12 @@ class TestTransformerGenerate(unittest.TestCase):
     # with temperature=2.0, we should see at least 2 distinct outputs across 5 runs
     self.assertGreater(len(runs), 1, "high temperature should produce varied outputs")
 
+  def test_recurrent_temperature_high_produces_variety(self):
+    model = Transformer(TEST_CONFIG)
+    model.has_recurrent_block = True
+    outputs = {model.forward(Tensor([[1]]), 0, Tensor([2.0])).item() for _ in range(5)}
+    self.assertGreater(len(outputs), 1)
+
   def test_temperature_passed_to_forward(self):
     """Temperature from generate should be passed through to __call__."""
     model = Transformer(TEST_CONFIG)

test/unit/test_metal_graph.py

@@ -1,5 +1,4 @@
 import unittest
-from unittest.mock import MagicMock
 from tinygrad import Device
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.dtype import dtypes
@@ -11,36 +10,27 @@ class TestMetalGraph(unittest.TestCase):
     self.MetalGraph = MetalGraph
     self.dev = Device[Device.DEFAULT]
 
-  def metal_buf(self, offset):
-    buf = MagicMock()
-    if offset > 0:
-      buf.op = Ops.SLICE
-      src = MagicMock()
-      src.dtype = dtypes.uint8
-      buf.src = (src, UOp.const(offset))
-      buf.dtype = dtypes.uint8
-    else:
-      buf.op = Ops.BUFFER
-    buf.device = Device.DEFAULT
-    return buf
+  def metal_buf(self, offset, bitcast=False):
+    size = 4 if bitcast else 1
+    buf = UOp.new_buffer(Device.DEFAULT, offset+size, dtypes.uint8)
+    if offset: buf = buf[offset:offset+size]
+    return buf.bitcast(dtypes.float32) if bitcast else buf
 
-  def call(self, *bufs):
-    c = MagicMock()
-    c.src = (MagicMock(op=Ops.PROGRAM),) + tuple(bufs)
-    return c
+  def supports_uop(self, *bufs):
+    return self.MetalGraph.supports_uop([self.dev], UOp(Ops.PROGRAM, src=(UOp.sink(),)).call(*bufs))
 
   def test_supports_uop_normal_offset(self):
-    assert self.MetalGraph.supports_uop([self.dev], self.call(self.metal_buf(0), self.metal_buf(100), self.metal_buf(0xFFFFFFFF))) is True
+    assert self.supports_uop(self.metal_buf(0), self.metal_buf(100), self.metal_buf(0xFFFFFFFF)) is True
 
   def test_supports_uop_overflow_offset(self):
-    assert self.MetalGraph.supports_uop([self.dev], self.call(self.metal_buf(0), self.metal_buf(0x100000000))) is False
+    assert self.supports_uop(self.metal_buf(0), self.metal_buf(0x100000000)) is False
 
-  def test_supports_uop_nonmetal_buf(self):
-    # non-SLICE ops should not be checked for offset
-    buf = MagicMock()
-    buf.op = Ops.BUFFER
-    buf.device = Device.DEFAULT
-    self.MetalGraph.supports_uop([self.dev], self.call(buf))
+  def test_supports_uop_non_view_buf(self):
+    assert self.supports_uop(self.metal_buf(0)) is True
+
+  def test_supports_uop_bitcast(self):
+    assert self.supports_uop(self.metal_buf(0xFFFFFFFF, bitcast=True)) is True
+    assert self.supports_uop(self.metal_buf(0x100000000, bitcast=True)) is False
 
 if __name__ == "__main__":
   unittest.main()

test/unit/test_schedule_cache.py

@@ -1,11 +1,16 @@
 import unittest
 import functools
-from tinygrad import Tensor, Variable, UOp
+from tinygrad import Tensor, Variable, UOp, function
 from tinygrad.uop.ops import KernelInfo
 from tinygrad.schedule import schedule_cache
 
-def custom_set0_kernel(A:UOp, num:int) -> UOp:
-  return A[0].set(num).sink(arg=KernelInfo(f"custom_set0_{num}"))
+def custom_add_kernel(A:UOp, B:UOp, num:int=0) -> UOp:
+  return A[0].set(B[0] + num).sink(arg=KernelInfo(f"custom_add_{num}"))
+
+def custom_add_backward(grad_output:UOp, _) -> tuple[None, UOp]:
+  grad = Tensor.invalids(*grad_output.shape, dtype=grad_output.dtype, device=grad_output.device)
+  grad = Tensor.custom_kernel(grad, Tensor(grad_output, device=grad_output.device), fxn=functools.partial(custom_add_kernel, num=0))[0]
+  return None, grad.uop
 
 class TestScheduleCache(unittest.TestCase):
   def test_bound_variable_reuses_cache(self):
@@ -25,27 +30,27 @@ class TestScheduleCache(unittest.TestCase):
 
   def test_custom_kernel(self):
     for i in range(4):
-      a = Tensor.empty(1)
-      a = Tensor.custom_kernel(a, fxn=functools.partial(custom_set0_kernel, num=i))[0]
+      a, b = Tensor.empty(1), Tensor.ones(1)
+      a = Tensor.custom_kernel(a, b, fxn=functools.partial(custom_add_kernel, num=i))[0]
       a.realize()
-      self.assertEqual(a.item(), i)
+      self.assertEqual(a.item(), i+1)
 
   def test_same_custom_function_reuses_cache(self):
     schedule_cache.clear()
-    fxn = functools.partial(custom_set0_kernel, num=10)
+    fxn = functools.partial(custom_add_kernel, num=10)
 
     # first run
-    a = Tensor.empty(1)
-    a = Tensor.custom_kernel(a, fxn=fxn)[0]
+    a, x = Tensor.empty(1), Tensor.ones(1)
+    a = Tensor.custom_kernel(a, x, fxn=fxn)[0]
     a.realize()
-    self.assertEqual(a.item(), 10)
+    self.assertEqual(a.item(), 11)
     cache_size_after_first = len(schedule_cache)
 
     # second run with same function should reuse cache
-    b = Tensor.empty(1)
-    b = Tensor.custom_kernel(b, fxn=fxn)[0]
+    b, x = Tensor.empty(1), Tensor.ones(1)
+    b = Tensor.custom_kernel(b, x, fxn=fxn)[0]
     b.realize()
-    self.assertEqual(b.item(), 10)
+    self.assertEqual(b.item(), 11)
     self.assertEqual(len(schedule_cache), cache_size_after_first)
 
   def test_simple(self):
@@ -65,5 +70,29 @@ class TestScheduleCache(unittest.TestCase):
       print(num)
     self.assertEqual(len(schedule_cache), start_len_schedule_cache)
 
+  def test_simple_precompile(self):
+    @function(precompile=True, precompile_backward=True)
+    def f(x:Tensor) -> Tensor:
+      out = Tensor.invalids(*x.shape, dtype=x.dtype, device=x.device)
+      out = Tensor.custom_kernel(out, x, fxn=functools.partial(custom_add_kernel, num=10), grad_fxn=custom_add_backward)[0]
+      return out + x
+
+    # warmup
+    x = Tensor.ones(1).realize()
+    out = f(x)
+    out.backward(x)
+    self.assertEqual(out.item(), 12)
+    self.assertEqual(x.grad.item(), 2)
+
+    # use the cache next time function is called
+    start_len_schedule_cache = len(schedule_cache)
+    for _ in range(3):
+      x = Tensor.ones(1).realize()
+      out = f(x)
+      out.backward(x)
+      self.assertEqual(out.item(), 12)
+      self.assertEqual(x.grad.item(), 2)
+    self.assertEqual(len(schedule_cache), start_len_schedule_cache)
+
 if __name__ == "__main__":
   unittest.main()

tinygrad/codegen/__init__.py

@@ -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_fold_cast_const, pm_move_where_on_load, pm_clean_up_group_sink, pm_remove_invalid
+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.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
@@ -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_fold_cast_const+pm_flatten_range, name="initial symbolic")
+    sink = graph_rewrite(sink, sym+pm_flatten_range, name="initial symbolic")
 
     # optimize (schedule) the AST
     sink = graph_rewrite(sink, pm_flatten_range+pm_simplify_ranges, ctx={}, name="simplify ranges")
@@ -346,7 +346,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
 
   # lower index dtype
   # NOTE: we need indexing_simplify to remove the cast to long using the Invalid
-  sink = graph_rewrite(sink, symbolic_simple+pm_fold_cast_const+pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes")
+  sink = graph_rewrite(sink, symbolic_simple+pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes")
 
   # final symbolic before decomp
   sink = graph_rewrite(sink, symbolic, name="final symbolic")
@@ -357,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+pm_fold_cast_const+get_simplifying_rewrite_patterns(supported_ops)
+  pm_decomp = symbolic_simple+get_simplifying_rewrite_patterns(supported_ops)
   sink = graph_rewrite(sink, pm_decomp, name="early decompositions")
 
   # late decomps + move gates from unrenderable INVALID where

tinygrad/codegen/gpudims.py

@@ -57,7 +57,7 @@ def add_gpudims(ctx:Renderer, s:UOp):
 
   # get the idxs
   ki: KernelInfo = s.arg
-  if ctx.has_threads: idxs = [UOp.variable("core_id", 0, int(global_shape[0])-1, dtypes.int).cast(dtypes.weakint)]
+  if ctx.has_threads: idxs = [UOp.variable("core_id", 0, int(global_shape[0])-1, dtypes.int, param=True).cast(dtypes.weakint)]
   elif ki.dont_use_locals:
     assert not local_dims, "can't use locals if there's no local dims"
     idxs = get_grouped_dims("idx", global_shape, ctx.global_max, reverse=True)
@@ -89,7 +89,7 @@ def add_gpudims(ctx:Renderer, s:UOp):
 
 pm_device_to_var = PatternMatcher([
   # the DEVICE axis is not a program axis, it's bound per device at launch. lower it to the _device_num variable (like SPECIAL for devices)
-  (UPat(Ops.RANGE, name="r"), lambda r: UOp.variable("_device_num", 0, r.vmax, dtype=r.dtype) if r.arg[-1] is AxisType.DEVICE else None),
+  (UPat(Ops.RANGE, name="r"), lambda r: UOp.variable("_device_num", 0, r.vmax, dtype=r.dtype, param=True) if r.arg[-1] is AxisType.DEVICE else None),
   # ENDs that closed a DEVICE range no longer close it
   (UPat(Ops.END, name="e"), lambda e: e.replace(src=(e.src[0],)+tuple(s for s in e.src[1:] if s.op is not Ops.PARAM))
    if any(s.op is Ops.PARAM and s.arg.name == '_device_num' for s in e.src[1:]) else None),

tinygrad/codegen/late/coalesce.py

@@ -26,7 +26,7 @@ def _drop_valid_stmts(valid:UOp, idx:UOp, height:int, width:int) -> list[UOp]:
     # check if idx is out of bound when X is on the wrong side of the bound: X in [c+1, vmax] or [vmin, c-1]
     lo, hi = (c + 1, X.vmax) if is_upper_bound else (X.vmin, c - 1)
     if lo <= hi:
-      fake = UOp.variable(f"fake{i}", lo, hi, X.dtype)
+      fake = UOp.variable(f"fake{i}", lo, hi, X.dtype, param=True)
       subs = [{X: fake}]
       # idx may not have X itself, so also substitute a term of X: v -> fake - (X - v)
       terms = list(X.split_uop(Ops.ADD))

tinygrad/codegen/late/gater.py

@@ -3,7 +3,7 @@ from tinygrad.uop.ops import PatternMatcher, UPat, Ops
 from tinygrad.dtype import Invalid, dtypes
 
 def move_where_load(gate, l, a, w):
-  return l.replace(src=(l.src[0], l.vconst_like(0) if a.is_invalid else
+  return l.replace(src=(l.src[0], l.vconst_like(0) if a.is_invalid else l.const_like(a.val) if a.op is Ops.CONST else
     a.src[0] if a.op is Ops.CAST and a.src[0].dtype == l.dtype else a.cast(l.dtype), l.src[2])).cast(w.dtype)
 
 pm_move_gates_from_index = PatternMatcher([

tinygrad/codegen/opt/tc.py

@@ -1,6 +1,7 @@
 import math, functools
 from dataclasses import dataclass
 from tinygrad.dtype import DType, dtypes
+from tinygrad.uop.ops import PatternMatcher, UOp, UPat, Ops
 
 @dataclass(frozen=True)
 class TensorCore: # D = A * B + C, A is (M x K), B is (K x N), C and D are (M x N)
@@ -135,6 +136,42 @@ amd_cdna4 = amd_cdna_1616128 + amd_cdna_161632 + amd_cdna_161616
 
 def get_amd(arch): return {"gfx942": amd_cdna3, "gfx950": amd_cdna4, "gfx1200": amd_rdna4, "gfx1201": amd_rdna4}.get(arch, amd_rdna3)
 
+pm_validate_wmma_rdna3 = PatternMatcher([
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.int32), lambda x: x.replace(
+    src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
+    if x.src[0].dtype == dtypes.int8 and x.src[0].max_numel() == 16 else None),
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(x.replace(
+      src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(UOp.const(j//2, dtypes.int16))
+      if j%2 == 0 else UOp.const(0.0, x.src[2].dtype)
+      for j in range(x.max_numel()*2)))),
+      arg=(*x.arg[:4], None)).index(UOp.const(i*2, dtypes.int16))
+      for i in range(x.max_numel()))) if x.max_numel() == 8 else None),
+  (UPat(Ops.WMMA, name="x"), lambda x: x.replace(
+    src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
+    if x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 16 else None),
+])
+
+pm_validate_wmma_rdna4 = PatternMatcher([
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16), lambda x: x.replace(
+    dtype=dtypes.uint16,
+    src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2].bitcast(dtypes.uint16)))
+      .bitcast(dtypes.bfloat16) if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None),
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+    lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
+    if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None)
+])
+
+pm_validate_wmma_cdna = PatternMatcher([
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+    lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
+    if x.arg[0][2] == 128 and x.src[0].dtype.itemsize <= 8 else None),
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+    lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
+    if x.max_numel() == 4 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 4 else None),
+  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
+    lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2]))
+    if x.max_numel() == 4 and x.src[0].dtype in dtypes.fp8_ocp and x.src[0].max_numel() == 8 else None),
+])
 # ***** Apple Metal *****
 
 metal = [TensorCore(dims=(8,8,8), threads=32, elements_per_thread=(2,2,2), dtype_in=di, dtype_out=do,

tinygrad/codegen/simplify.py

@@ -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, pm_fold_cast_const, invalid_gate
+from tinygrad.uop.symbolic import symbolic, 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_fold_cast_const+pm_flatten_range, ctx={r0:new_range//s1, r1:new_range%s1},
+        nidx = graph_rewrite(u, _substitute+symbolic+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
@@ -137,7 +137,7 @@ def reduce_collapse(red:UOp, u:UOp, pm:PatternMatcher=pm_reduce_collapse) -> UOp
     for u in included:
       for s in u.src:
         if s in included or s in replaces or s.op in {Ops.CONST, Ops.PARAM, Ops.BUFFER}: continue
-        replaces[s] = UOp.variable(f'in{len(replaces)}', s.vmin, s.vmax, s.dtype)
+        replaces[s] = UOp.variable(f'in{len(replaces)}', s.vmin, s.vmax, s.dtype, param=True)
     collapse_fxn = u.substitute(replaces).reduce(r, arg=Ops.ADD)
     sink = graph_rewrite(collapse_fxn, pm, name="reduce_collapse")
     if not no_range(sink): return None

tinygrad/device.py

@@ -336,6 +336,8 @@ class Compiled:
   pm_lower:Any = None
   pm_bufferize:Any = None
 
+  has_copy_queue:bool = True
+
   def __init__(self, device:str, allocator:Allocator, renderers:list[type[Renderer]], runtime:type[Program[Self]]|None, graph=None, arch=None):
     from tinygrad.renderer import Renderer
     self.device, self.allocator, self.runtime_t, self.graph, self.renderers = device, allocator, runtime, graph, renderers or [Renderer]

tinygrad/engine/jit.py

@@ -44,9 +44,7 @@ def graph_split_rewrite(linear:UOp, max_batch_size:int=0) -> UOp:
     current_batch, current_batch_devs = [], []
 
   for si in linear.src:
-    if si.src[0].op is Ops.SLICE: continue
-
-    devs = dedup([Device[x] for b in si.src[1:] if b.op is not Ops.BIND for x in (b.device if isinstance(b.device, tuple) else (b.device,))])
+    devs = dedup([Device[x] for b in si.src[1:] if not b.is_bound_var for x in (b.device if isinstance(b.device, tuple) else (b.device,))])
     graph_t = graph_class(devs[0]) if devs[0].graph is not None else None
 
     can_graph = graph_t is not None and graph_t.supports_uop(devs, si)
@@ -180,7 +178,7 @@ class CapturedJit(Generic[ReturnType]):
       if call.op is not Ops.CALL: continue
       arg_uops = get_call_arg_uops(call)
       outs, ins = get_call_outs_ins(call)
-      out |= {arg_uops[k] for k in set(outs) - set(ins) if arg_uops[k].op in (Ops.BUFFER, Ops.SLICE)}
+      out |= {b for k in set(outs) - set(ins) if (b:=u if (cv:=(u:=arg_uops[k]).contiguous_view()) is None else cv[0]).op is Ops.BUFFER}
     return out
 
   def __call__(self, input_uops:list[UOp], var_vals:dict[str, int]) -> ReturnType:

tinygrad/engine/realize.py

@@ -4,20 +4,20 @@ 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, 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.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite
+from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
 from tinygrad.renderer import Estimates
 from tinygrad.codegen import to_program
 from tinygrad.codegen.opt.postrange import args_from_ast
 
 # **************** Helpers ****************
 
-def get_call_arg_uops(call:UOp) -> tuple[UOp, ...]: return tuple(s for s in call.src[1:] if s.op is not Ops.BIND)
+def get_call_arg_uops(call:UOp) -> tuple[UOp, ...]: return tuple(s for s in call.src[1:] if not s.is_bound_var)
 
 def get_call_outs_ins(call:UOp) -> tuple[tuple[int, ...], tuple[int, ...]]:
   ast = call.src[0]
   if ast.op is Ops.PROGRAM: return tuple(ast.arg.outs), tuple(ast.arg.ins)
-  if ast.op in (Ops.COPY, Ops.SLICE): return (0,), (1,)
+  if ast.op is Ops.COPY: return (0,), (1,)
   if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return (0,), tuple(range(1, len(get_call_arg_uops(call))))
   return (), ()
 
@@ -27,9 +27,6 @@ def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|N
 
   ast, arg_uops = call.src[0], get_call_arg_uops(call)
   if ast.op is Ops.PROGRAM: return ast.arg.name
-  if ast.op is Ops.SLICE:
-    offset = ast.src[1].val * arg_uops[1].dtype.itemsize
-    return colored(f"view {_uop_sz_to_str(arg_uops[0]):>10} @ {offset:<10d}", "yellow")
   if ast.op is Ops.COPY: return colored(f"copy {_uop_sz_to_str(arg_uops[0]):>10}, {_dev_str(bufs[0]):>7s} <- {_dev_str(bufs[1]):7s}", "yellow")
   if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return colored(f"enc/dec {_uop_sz_to_str(arg_uops[0])}", "yellow")
   if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return colored(f"batched {len(ast.src[0].src)}", "cyan")
@@ -140,7 +137,7 @@ class ExecContext:
   cache: bool = True
 
 def _resolve(b:UOp, inputs:tuple[UOp, ...]) -> UOp:
-  if b.op in (Ops.SLICE, Ops.MSELECT) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg.slot], *b.src[1:]))
+  if b.op in (Ops.MSELECT, Ops.SHRINK) and b.src[0].op is Ops.PARAM: return b.replace(src=(inputs[b.src[0].arg.slot], *b.src[1:]))
   if b.op is Ops.MSTACK: return b.replace(src=tuple(_resolve(x, inputs) for x in b.src))
   return inputs[b.arg.slot] if b.op is Ops.PARAM else b
 def resolve_params(call:UOp, inputs:tuple[UOp, ...]) -> list[UOp]: return [_resolve(b, inputs) for b in get_call_arg_uops(call)]
@@ -154,13 +151,6 @@ def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], d
                    for x in call.src[0].toposort())
     for j, per_dev in enumerate(zip(*[cast(MultiBuffer, b).bufs for b in bufs])): yield list(per_dev), {"_device_num": j} if has_dnum else {}
 
-def exec_view(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
-  resolved = resolve_params(call, ctx.input_uops)
-  bufs = [cast(Buffer, b.buffer) for b in resolved]
-  bv = bufs[1].view(resolved[0].max_numel(), ast.dtype, ast.src[1].val*bufs[1].dtype.itemsize)
-  with track_stats(ctx, call, bv.device, [bv, bufs[1]], ctx.var_vals): buffers[resolved[0]] = bv
-  return None
-
 def exec_copy(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
   for bufs, device_vars in unwrap_multi(call, resolve_params(call, ctx.input_uops)):
     dest, src = bufs[0].ensure_allocated(), bufs[1].ensure_allocated()
@@ -210,27 +200,30 @@ def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
   return t[0]
 
 def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
-  if (inputs:=call.arg.aux.inputs) is not None:
+  if (info:=call.arg.aux).inputs is not None:
     bufs = [_resolve(ctx.input_uops[i], ctx.input_uops).buffer for i in call.arg.aux.input_idxs]
-    table = call.src[1+inputs].buffer
+    table = call.src[1+info.inputs].buffer
     for j,dev in enumerate(call.arg.aux.device):
       addrs = array.array('Q', [(b.bufs[j] if isinstance(b, MultiBuffer) else b).get_buf(dev).va_addr for b in bufs])
       mv = (table.bufs[j] if isinstance(table, MultiBuffer) else table).ensure_allocated()._buf.cpu_view().view(fmt='Q')
       wait_cond(lambda: mv[0], value=0, timeout_ms=ctx.timeout or getenv("HCQDEV_WAIT_TIMEOUT_MS", 30000), msg=f"{dev} hang detected")
       mv[:len(addrs)] = addrs
 
-  exec_kernel(replace(ctx, update_stats=False), call, ast)
+  exec_kernel(replace(ctx, update_stats=DEBUG>=3), call, ast)
 
-  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]
+  tms = []
+  for devices,name,estimates,prof in info.kernels:
+    for device in devices:
+      tm = None
+      if prof:
+        (d:=cast(Any, Device[device])).prof_ents[prof[0]] = ProfileGraphEntry(device, name, *prof)
+        if ctx.wait:
+          d.synchronize(timeout=ctx.timeout)
+          st, en = (d.signal(x)._buf.cpu_view().view(fmt='Q')[0] for x in prof)
+          tms.append(tm:=float(en-st)/d.timestamp_divider/1e6)
+      stat_call = call.replace(arg=replace(call.arg, name=name, aux=replace(info, estimates=estimates, kernels=())))
+      with track_stats(ctx, stat_call, device, [], ctx.var_vals) as et: et[0] = tm
+  return max(tms) if tms else None
 
 # flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src
 pm_flatten_linear = PatternMatcher([
@@ -261,7 +254,6 @@ pm_optimize_local_size = PatternMatcher([
 ])
 
 pm_exec = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.SLICE, name="ast"),), name="call", allow_any_len=True), exec_view),
   (UPat(Ops.CALL, src=(UPat(Ops.COPY, name="ast"),), name="call", allow_any_len=True), exec_copy),
   (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="ast"),), name="call", allow_any_len=True), exec_kernel),
   (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="encdec", name="ast"),), name="call", allow_any_len=True), exec_encdec),
@@ -276,7 +268,7 @@ def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:li
   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, bool(PROFILE) if profile is None else profile)
+  if getenv("HCQ2"): linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile)
   return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)
 
 def link_linear(linear:UOp, cache=True) -> UOp: return hcq_link(linear, cache=cache) if getenv("HCQ2") else linear

tinygrad/function.py

@@ -1,4 +1,5 @@
 import functools, time
+from dataclasses import replace
 from typing import Generic, TypeVar, Callable, cast, overload
 from tinygrad.helpers import Context, dedup, getenv, DEBUG
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, PatternMatcher, UPat
@@ -12,7 +13,7 @@ def add_to_ctx(ctx, x:UOp):
   return ret
 
 pm_ctx = PatternMatcher([
-  (UPat((Ops.BUFFER, Ops.BIND), name="x"), add_to_ctx),
+  (UPat(Ops.BUFFER, name="x"), add_to_ctx),
   (UPat((Ops.AFTER, Ops.CONTIGUOUS), name="x"),
    lambda ctx,x: add_to_ctx(ctx,x) if not x.op_in_backward_slice_with_self(Ops.PARAM) and x.op_in_backward_slice_with_self(Ops.BUFFER) else None),
 ])
@@ -23,6 +24,10 @@ def invalid_outputs(uret:UOp) -> set[UOp]:
   return {u.src[0].buf_uop for u in uret.backward_slice_with_self
           if u.op is Ops.STORE and u.src[1].base.is_invalid and not u.src[0].buf_uop.is_realized}
 
+def renumber_invalid_outputs(uret:UOp) -> UOp:
+  return uret.substitute({b:b.replace(arg=replace(b.arg, slot=i))
+                          for i,b in enumerate(x for x in uret.toposort(enter_calls=False) if x in invalid_outputs(uret))})
+
 ReturnType = TypeVar('ReturnType')
 class _function(Generic[ReturnType]):
   depth = 0
@@ -65,6 +70,7 @@ class _function(Generic[ReturnType]):
     # the BUFFERs that are left are the implicit inputs
     num_explicit = len(call_uops)
     uret = graph_rewrite(uret, pm_ctx, (call_uops, invalid_outputs(uret)), bottom_up=True, name="get_implicit_inputs")
+    uret = renumber_invalid_outputs(uret)
     name = getattr(self.fxn, '__qualname__', None) or type(self.fxn).__qualname__
     if not self.allow_implicit:
       implicit_buffers = [x for x in call_uops[num_explicit:] if x.op is Ops.BUFFER]

tinygrad/helpers.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 import time
 START_TIME = time.perf_counter()
-import os, functools, platform, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc
+import os, functools, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc
 from collections import defaultdict
-import subprocess, shutil, math, types, copyreg, inspect, importlib, decimal, itertools, difflib
+import shutil, math, types, copyreg, inspect, importlib, decimal, itertools, difflib
 from dataclasses import dataclass, field, replace
 from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator, cast, overload
 
@@ -13,8 +13,7 @@ U = TypeVar("U")
 def prod(x:Iterable[T]) -> T|int: return functools.reduce(operator.mul, x, 1)
 
 # NOTE: helpers is not allowed to import from anything else in tinygrad
-OSX, WIN = platform.system() == "Darwin", sys.platform == "win32"
-ARCH_X86 = any(x in platform.processor() for x in ("Intel", "i386", "x86_64"))
+OSX, WIN = sys.platform == "darwin", sys.platform == "win32"
 BASEDIR = pathlib.Path(__file__).parent
 
 # fix colors on Windows, https://stackoverflow.com/questions/12492810/python-how-can-i-make-the-ansi-escape-codes-to-work-also-in-windows
@@ -231,7 +230,7 @@ class _DEV(ContextVar):
 
 DEV, DEBUG, BEAM, NOOPT = _DEV("DEV", ""), ContextVar("DEBUG", 0), ContextVar("BEAM", 0), ContextVar("NOOPT", 0)
 IMAGE, FLOAT16, OPENPILOT_HACKS = ContextVar("IMAGE", 0), ContextVar("FLOAT16", 0), ContextVar("OPENPILOT_HACKS", 0)
-JIT, JIT_BATCH_SIZE = ContextVar("JIT", 2 if OSX and ARCH_X86 else 1), ContextVar("JIT_BATCH_SIZE", 32)
+JIT, JIT_BATCH_SIZE = ContextVar("JIT", 1), ContextVar("JIT_BATCH_SIZE", 32)
 CHUNK_SIZE = 2**20  # TinyFS content-addressed store: blob chunk + hash-tree node granularity
 WINO, CAPTURING, TRACEMETA, NO_COLOR = ContextVar("WINO", 0), ContextVar("CAPTURING", 1), ContextVar("TRACEMETA", 1), ContextVar("NO_COLOR", 0)
 TRAINING = ContextVar("TRAINING", 0)
@@ -454,9 +453,9 @@ def _ensure_downloads_dir() -> pathlib.Path:
   if pathlib.Path("/etc/tinybox-release").is_file():
     # try creating dir with sudo
     if not (downloads_dir := pathlib.Path("/raid/downloads")).exists():
-      subprocess.run(["sudo", "mkdir", "-p", downloads_dir], check=True)
-      subprocess.run(["sudo", "chown", "tiny:root", downloads_dir], check=True)
-      subprocess.run(["sudo", "chmod", "775", downloads_dir], check=True)
+      system(f"sudo mkdir -p {downloads_dir}")
+      system(f"sudo chown tiny:root {downloads_dir}")
+      system(f"sudo chmod 775 {downloads_dir}")
     return downloads_dir
   return pathlib.Path(cache_dir) / "downloads"
 
@@ -497,6 +496,7 @@ def fetch_fw(path:str, name:str, sha256:str) -> bytes:
 # *** Exec helpers
 
 def system(cmd:str, **kwargs) -> str:
+  import subprocess
   st = time.perf_counter()
   try: ret = subprocess.check_output(cmd.split(), stderr=subprocess.STDOUT, **kwargs).decode().strip()
   except subprocess.CalledProcessError as e:

tinygrad/llm/cli.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-import sys, argparse, codecs, itertools, typing, re, unicodedata, json, time, pathlib
+import sys, argparse, codecs, itertools, typing, re, unicodedata, json, time
 from typing import TYPE_CHECKING
 from tinygrad import nn
 from tinygrad.uop.ops import UOp, Ops
@@ -127,95 +127,30 @@ class FallbackTemplate:
       out += self.end_turn()
     return out + self.role("assistant") if add_generation_prompt else out
 
-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
+from tinygrad.llm.serve import LLMServer
 
 def main():
   parser = argparse.ArgumentParser()
-  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("--model", "-m", default=list(models.keys())[0], help=f"Model choice ({', '.join(models.keys())}) or path to a local GGUF file")
   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_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')
+  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
   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|KimiK3Template = KimiK3Template() if is_k3 else FallbackTemplate(tok)
-  if ct is not None:
+  template: jinja2.Template|FallbackTemplate = FallbackTemplate(tok)
+  if (ct := kv.get('tokenizer.chat_template')) is not None:
     try:
       import jinja2
       env = jinja2.Environment()
@@ -227,10 +162,9 @@ def main():
       template = env.from_string(ct)
     except ImportError: print("warning: jinja2 is not installed, the model's chat template is disabled")
 
-  # Capture the default greedy serving shapes before accepting requests.
-  if args.warmup:
+  # warmup the JIT
+  if args.warmup or args.serve:
     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()
@@ -255,26 +189,15 @@ def main():
   while 1:
     try: messages.append({"role":"user", "content":input('>>> ')})
     except EOFError: break
-    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)
+    ids = tok.encode(template.render(messages=messages, add_generation_prompt=True))
+    reply, dec = "", tok.stream_decoder()
     for next_id in model.generate(ids):
       if tok.is_end(next_id):
-        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")
+        sys.stdout.write(dec() + "\n\n")
         break
-      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)
+      reply += (piece := dec(next_id))
+      sys.stdout.write(piece)
+      sys.stdout.flush()
+    messages.append({"role":"assistant", "content":reply})
 
 if __name__ == "__main__": main()

tinygrad/llm/kernels/__init__.py

@@ -1,238 +0,0 @@
-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]

tinygrad/llm/kernels/amd.py

@@ -1,353 +0,0 @@
-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=()))

tinygrad/llm/kimi.py

@@ -1,181 +0,0 @@
-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"]

tinygrad/llm/kimi_k3.py

@@ -1,337 +0,0 @@
-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"]

tinygrad/llm/model.py

@@ -1,16 +1,9 @@
 from __future__ import annotations
-import array, functools, itertools, pathlib
+import functools, itertools, pathlib
 from dataclasses import dataclass, replace
-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
@@ -27,37 +20,12 @@ 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)
@@ -66,16 +34,6 @@ 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
@@ -84,7 +42,6 @@ class SSMConfig:
   time_step_rank: int
   inner_size: int
   kda: bool = False
-  channel_decay: bool = False
 
 @dataclass(frozen=True)
 class TransformerConfig:
@@ -116,20 +73,6 @@ 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):
@@ -143,15 +86,9 @@ 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)}
-      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)
+      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)
       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)
@@ -162,85 +99,28 @@ 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
-      # 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)
+      logits = self.ffn_gate_inp(x)
       if hasattr(self, 'exp_probs_b'):
-        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)
+        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)
       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
-      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)
+      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_gate_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
+        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
       return out
     # TODO: remove the need for this contiguous
-    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)
+    return self.ffn_down(self.ffn_gate(x).silu().contiguous() * self.ffn_up(x))
 
   # 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
@@ -251,11 +131,6 @@ 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):
@@ -263,32 +138,6 @@ 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)
@@ -330,9 +179,7 @@ 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
-    # 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),)) \
+    mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, buffer=False).triu(start_pos+1) \
       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)
@@ -340,9 +187,8 @@ class TransformerBlock(FFNBlock):
 
   def _init_state(self, x:Tensor):
     if not hasattr(self, "cache_kv"):
-      # TODO: how is the dtype of this determined?
       self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim,
-                                   device=x.device, dtype=x.dtype)
+                                   dtype=dtypes.default_float, device=x.device)
       self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta, device=x.device)
 
 class MLATransformerBlock(FFNBlock):
@@ -360,23 +206,17 @@ 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
-    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_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)
     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_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)
+    kv_a = 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])
@@ -385,23 +225,17 @@ 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, T), float("-inf"), dtype=x.dtype, device=x.device).triu(1).pad(((0, 0),)*3+((start_pos, 0),)) \
+    mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, buffer=False).triu(start_pos+1) \
       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
-    # 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.softmax(-1)
     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, dtype=x.dtype)
+      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.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta, device=x.device)
 
 class GatedDeltaNetBlock(FFNBlock):
@@ -411,144 +245,68 @@ 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
-    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)}
+    self.attn_qkv = Linear(config.dim, self.conv_channels, bias=False)
     if ssm.kda:
-      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_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.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_a = Tensor.zeros(self.num_v_heads, 1) if ssm.kda else Tensor.zeros(self.num_v_heads)
     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"
+    is_kda = hasattr(self, "ssm_g_a")
 
     # input processing
-    # 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)
+    x = x.half()
+    out_gate = self.ssm_g_b(self.ssm_g_a(x)) if is_kda else self.attn_gate(x)
+    out_gate = out_gate.reshape(B, 1, self.num_v_heads, self.head_v_dim)
+    beta = self.ssm_beta(x).sigmoid().reshape(B, self.num_v_heads, 1, 1)
+    alpha = self.ssm_f_b(self.ssm_f_a(x)) if is_kda else self.ssm_alpha(x)
+    alpha = ((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, self.num_v_heads, -1) *
+             self.ssm_a.reshape(1, self.num_v_heads, -1)).exp().unsqueeze(-2)
 
-    # 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)
+    # 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)
 
-    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)
+    # recurrent
+    recurrent_state = self.recurrent_state * alpha
+    recurrent_state = recurrent_state + ((v - recurrent_state@k) * beta)@k.transpose(-1, -2)
 
-    # 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)
+    # 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))
 
-  # 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
+    # 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 is_kda else out_gate.silu()
+    return self.ssm_out((core_attn_out * out_gate).reshape(B, 1, -1).cast(x.dtype))
+
+  # recurrent state can't be partially reused after divergence, force a full rebuild
+  def _state_reset_ops(self):
+    return [self.conv_state.assign(self.conv_state.const_like(0)),
+            self.recurrent_state.assign(self.recurrent_state.const_like(0))] if hasattr(self, "conv_state") else []
 
   def _init_state(self, x):
-    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()
+    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()
 
 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
@@ -558,76 +316,23 @@ 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 _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)
+  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)
+    # only run the output projection on the last token
+    logits = self.output(self.output_norm(x[:, -1:]))[:, -1, :]
     # Gumbel-max trick: argmax(logits/temp - log(-log(uniform))) is equivalent to sampling from softmax(logits/temp)
     return (logits / temperature.maximum(1e-12) - (Tensor.rand_like(logits).maximum(1e-12).log().neg()).log()).argmax(-1, keepdim=True)
 
-  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)
+  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)
 
   @staticmethod
   def from_gguf(gguf:Tensor|str|pathlib.Path, max_context:int|None=None,
@@ -713,91 +418,35 @@ class Transformer:
     return model, kv
 
   def warmup(self):
-    # 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.
+    for _ in range(2): list(zip(range(2), self.generate([0])))
+
+  def get_start_pos(self, tokens:list[int]) -> int:
+    # recurrent state can't be partially reused after divergence: reuse it only when tokens extend the cached prefix
     if self.has_recurrent_block:
-      for i in range(1, 4): list(zip(range(2), self.generate(prompt + list(range(1, i+1)), temperature=0.0)))
-
-  def _cache_start(self, tokens:list[int]) -> tuple[int, bool]:
+      return len(self._cached_tokens) if self._cached_tokens and len(self._cached_tokens) < len(tokens) \
+        and tokens[:len(self._cached_tokens)] == self._cached_tokens else 0
     prefix_len = sum(1 for _ in itertools.takewhile(lambda ab: ab[0] == ab[1], zip(tokens[:-1], self._cached_tokens)))
-    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]
+    return min(block._reusable_prefix_len(prefix_len, len(self._cached_tokens)) for block in self.blk)
 
   def generate(self, tokens:list[int], chunk_size:int=32, temperature:float=0.0):
-    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
+    if self.has_recurrent_block: 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
-    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)
+    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)
     # recompute start_pos from what's currently valid in the caches
-    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()
+    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)
     out, prompt_len = None, len(tokens)
-    token_host = memoryview(bytearray(4)) if amd_tp else None
     while len(tokens) < self.max_context:
-      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()
+      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()
       start_pos += n_toks
       # chunked prefill: keep processing until all prompt tokens are consumed
       if start_pos < len(tokens): continue
-      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()))
+      tokens.append(int(out.item()))
       self._cached_tokens = tokens[:-1]
       yield tokens[-1]

tinygrad/llm/quant.py

@@ -1,81 +0,0 @@
-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)

tinygrad/llm/serve.py

@@ -24,11 +24,6 @@ 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:
@@ -39,10 +34,9 @@ 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, xtml:bool=False):
+  def __init__(self, reasoning: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:
@@ -51,39 +45,20 @@ 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("<|close|>think<|sep|>" if self.xtml else "</think>", final)
+      emit, done = self.split("</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, tool_tag = self.split_any(("<tool_call>", "<|tool_calls_section_begin|>"), final)
+    emit, found = self.split("<tool_call>", final)
     if emit: yield "content", emit
-    if tool_tag: self.mode, self.buf = "tool", tool_tag + self.buf
+    if found: self.mode, self.buf = "tool", "<tool_call>" + self.buf
 
 class Handler(HTTPRequestHandler):
   server: LLMServer
@@ -92,7 +67,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, xtml:bool=False):
+                reasoning:bool=False):
     model, tok = self.server.model, self.server.tok
     prompt_tokens = len(ids)
     cache_start_pos = model.get_start_pos(ids)
@@ -103,7 +78,7 @@ class Handler(HTTPRequestHandler):
     finish_reason = "stop"
     st = pt = time.perf_counter()
     dec = tok.stream_decoder()
-    router = StreamRouter(reasoning, xtml)
+    router = StreamRouter(reasoning)
     def log_stats(interrupted:bool=False):
       et = time.perf_counter()
       total = f"total:{et-st:6.2f}s"
@@ -131,14 +106,6 @@ 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"
@@ -172,10 +139,9 @@ 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=xtml or rendered.rstrip().endswith("<think>"), xtml=xtml)
+                              reasoning=rendered.rstrip().endswith("<think>"))
       if body.get("stream"): self.stream_json(chunks)
       else:
         out, reasoning, tool_calls, finish_reason = [], [], [], "stop"

tinygrad/mixin/elementwise.py

@@ -115,7 +115,8 @@ class ElementwiseMixin(CreationMixin):
     ```
     """
     a, b = self._broadcasted(x, reverse)
-    return a + (-b)
+    # alu, not +: _broadcasted already promoted these, and a second promote would cast -b (only a bare weak CONST is kept weak)
+    return a.alu(Ops.ADD, -b)
 
   def mul(self, x: Self | ConstType, reverse: bool = False) -> Self:
     """
@@ -245,8 +246,9 @@ class ElementwiseMixin(CreationMixin):
     if dtypes.is_int(a.dtype) and dtypes.is_int(b.dtype):
       if rounding_mode == "trunc": return a.alu(Ops.CDIV, b)
       if rounding_mode == "floor": return a.alu(Ops.FLOORDIV, b)
-      a = a.cast(dtypes.default_float)
-    d = a * b.reciprocal()
+    if dtypes.is_int(a.dtype) or a.dtype == dtypes.bool: a = a.cast(dtypes.default_float)
+    # alu, not *: _broadcasted already promoted these, and a second promote would cast 1/b (only a bare weak CONST is kept weak)
+    d = a.alu(Ops.MUL, b.reciprocal())
     if rounding_mode is None: return d
     if rounding_mode == "trunc": return d.trunc()
     if rounding_mode == "floor": return d.floor()

tinygrad/mixin/gradient.py

@@ -3,6 +3,7 @@ import math, dataclasses
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata, broadcast_axes
 from tinygrad.helpers import argsort
 from tinygrad.dtype import sum_acc_dtype
+from tinygrad.function import renumber_invalid_outputs
 
 def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
   if op == Ops.ADD: return (ctx._broadcast_to(ret.src[0].shape),)
@@ -40,6 +41,7 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
   # collect needed gradient bodies, compact unused params, create a single backward CALL
   grad_bodies = [(i, grads[p]) for i in needed if (p:=params.get(i)) is not None and p in grads]
   bwd_body = UOp.maketuple(*(gb for _, gb in grad_bodies)).substitute(fwd_subs, walk=True)
+  bwd_body = renumber_invalid_outputs(bwd_body)
   bwd_body, compact_args = _compact_params(bwd_body, (*args, *grad_args, *fwd_outs))
   bwd_call = bwd_body.call(*compact_args, name=(k.arg.name or "")+"_backward", precompile=k.arg.precompile_backward)
   gb_map = {i: idx for idx, (i, _) in enumerate(grad_bodies)}

tinygrad/mixin/op.py

@@ -460,6 +460,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     """
     assert gradient is not None or self.shape == tuple(), "when no gradient is provided, backward must be called on a scalar tensor"
     if not (self.is_floating_point() and all(t.is_floating_point() for t in targets)): raise RuntimeError("only float Tensors have gradient")
+    if any(t.dtype in dtypes.weaks for t in targets): raise RuntimeError("cannot take gradient wrt a weak Tensor")
     from tinygrad.mixin.gradient import compute_gradient
     if gradient is None: gradient = self.const_like(1.0)
     target_uops = [t._uop for t in targets]

tinygrad/nn/state.py

@@ -1,4 +1,4 @@
-import json, math, pathlib, zipfile, pickle, tarfile, struct, functools, io, zlib
+import json, math, pathlib, struct, functools, io, zlib
 from collections import OrderedDict
 from typing import Any, Callable, BinaryIO, Iterable, cast
 from tinygrad.tensor import Tensor
@@ -219,6 +219,7 @@ def load_state_dict(model, state_dict:dict[str, Tensor], strict=True, verbose=Tr
 
 @accept_filename
 def zip_extract(t: Tensor) -> dict[str, Tensor]:
+  import zipfile
   files: dict[str, Tensor] = {}
   with zipfile.ZipFile(TensorIO(t), "r") as myzip:
     # sadly, the extra length needs to be read from the local header of each file.
@@ -249,6 +250,7 @@ def tar_extract(t: Tensor) -> dict[str, Tensor]:
   tensors = nn.state.tar_extract(Tensor(pathlib.Path("archive.tar")))
   ```
   """
+  import tarfile
   with tarfile.open(fileobj=TensorIO(t), mode="r") as tar:
     return {member.name:t[member.offset_data:member.offset_data+member.size] for member in tar if member.type == tarfile.REGTYPE}
 
@@ -303,6 +305,7 @@ def torch_load(t:Tensor) -> dict[str, Tensor]:
                "FloatTensor": None, "Parameter": Parameter}
   whitelist = {"torch", "collections", "numpy", "_codecs"}  # NOTE: this is not for security, only speed
   class Dummy: pass
+  import pickle, zipfile, tarfile
   class TorchPickle(pickle.Unpickler):
     def find_class(self, module, name):
       module_root = module.split(".")[0]

tinygrad/renderer/cstyle.py

@@ -264,6 +264,7 @@ class ClangRenderer(CStyleLanguage):
   nan = '__builtin_nanf("")'
 
   # language options
+  barrier = "__atomic_thread_fence(__ATOMIC_SEQ_CST);"
   buffer_suffix = " restrict"
   type_map = {dtypes.bool:"_Bool", dtypes.half:"__fp16"}
   code_for_op = {**({k:v for k,v in CStyleLanguage.code_for_op.items() if k not in [Ops.EXP2, Ops.SIN, Ops.LOG2, Ops.TRUNC, Ops.RECIPROCAL]}),

tinygrad/renderer/llvmir.py

@@ -142,7 +142,7 @@ base_rewrite = PatternMatcher([
   (UPat(Ops.IF, name="x"), lambda ctx,x: f"  br i1 {ctx[x.src[0]]}, label %ifbody_{ctx[x][1:]}, label %ifskip_{ctx[x][1:]}\nifbody_{ctx[x][1:]}:"),
   (UPat(Ops.ENDIF, name="x"), lambda ctx,x: f"  br label %ifskip_{ctx[x.src[0]][1:]}\nifskip_{ctx[x.src[0]][1:]}:"),
 
-  (UPat(Ops.BARRIER), lambda ctx: "")
+  (UPat(Ops.BARRIER), lambda ctx: "  fence seq_cst")
 ])
 
 class LLVMRenderer(Renderer):
@@ -238,8 +238,8 @@ class AMDLLVMRenderer(LLVMRenderer):
     (UPat(Ops.CAST, dtypes.fp8s, (UPat(dtype=dtypes.float),), name="x",), lambda ctx,x:
       f"  {ctx[x]} = call i8 @f32_to_fp8({ldt(x.src[0].dtype)}  {ctx[x.src[0]]}, i1 {'1' if x.dtype == dtypes.fp8e5m2 else '0'})"),
     (UPat(Ops.CAST, dtypes.float, (UPat.var("y", dtypes.fp8s),), name="x",), lambda ctx,x,y:
-      f"  {ctx[x.src[0]]}_i32 = zext i8 {ctx[x.src[0]]} to i32\n"
-      f"  {ctx[x]} = call float @llvm.amdgcn.cvt.f32.{'bf8' if y.dtype == dtypes.fp8e5m2 else 'fp8'}(i32 {ctx[x.src[0]]}_i32, i32 0)"),
+      f"  {ctx[x]}_i32 = zext i8 {ctx[x.src[0]]} to i32\n"
+      f"  {ctx[x]} = call float @llvm.amdgcn.cvt.f32.{'bf8' if y.dtype == dtypes.fp8e5m2 else 'fp8'}(i32 {ctx[x]}_i32, i32 0)"),
   ]) + base_rewrite
   extra_matcher = LLVMRenderer.extra_matcher + create_non_native_float_pats(dtypes.fp8s) + PatternMatcher([
     # amd llvm intrinsics llvm.log2/llvm.exp2 don't support double
@@ -279,43 +279,9 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n  %trunc
       (UPat(Ops.WMMA, name="wmma"), lambda ctx, wmma, rdna4=AMDLLVMRenderer.is_rdna4(target.arch), cdna=self.is_cdna:
         render_wmma_amd(ctx, wmma, cdna, rdna4))
     ])
-    if self.is_cdna:
-      self.extra_matcher += PatternMatcher([
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
-          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
-          if x.arg[0][2] == 128 and x.src[0].dtype.itemsize <= 8 else None),
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
-          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
-          if x.max_numel() == 4 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 4 else None),
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
-          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2]))
-          if x.max_numel() == 4 and x.src[0].dtype in dtypes.fp8_ocp and x.src[0].max_numel() == 8 else None),
-      ])
-    if target.arch in {"gfx1100", "gfx1151"}:
-      self.extra_matcher += PatternMatcher([
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.int32), lambda x: x.replace(
-          src=(x.src[0].bitcast(dtypes.uint32), x.src[1].bitcast(dtypes.uint32), x.src[2]))
-          if x.src[0].dtype == dtypes.int8 and x.src[0].max_numel() == 16 else None),
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(x.replace(
-          src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(UOp.const(j//2, dtypes.int16))
-            if j%2 == 0 else UOp.const(0.0, x.src[2].dtype)
-            for j in range(x.max_numel()*2)))),
-          arg=(*x.arg[:4], None)).index(UOp.const(i*2, dtypes.int16))
-          for i in range(x.max_numel()))) if x.max_numel() == 8 else None),
-        (UPat(Ops.WMMA, name="x"), lambda x: x.replace(
-          src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
-          if x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 16 else None),
-      ])
-    if target.arch in {"gfx1200", "gfx1201"}:
-      self.extra_matcher += PatternMatcher([
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16), lambda x: x.replace(
-          dtype=dtypes.uint16,
-          src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2].bitcast(dtypes.uint16)))
-            .bitcast(dtypes.bfloat16) if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None),
-        (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
-          lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))
-          if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None)
-      ])
+    if self.is_cdna: self.extra_matcher += tc.pm_validate_wmma_cdna
+    if target.arch in {"gfx1100", "gfx1151"}: self.extra_matcher += tc.pm_validate_wmma_rdna3
+    if target.arch in {"gfx1200", "gfx1201"}: self.extra_matcher += tc.pm_validate_wmma_rdna4
 
   def supported_dtypes(self): return {d for d in super().supported_dtypes()
                                       if (d not in dtypes.fp8_ocp or self.target.arch == "gfx950") and d not in dtypes.fp8_fnuz}

tinygrad/renderer/nir.py

@@ -137,7 +137,8 @@ class NIRRenderer(Renderer):
     (UPat(Ops.CAST, (dtypes.uchar, dtypes.ushort), src=(UPat.var("x", dtypes.floats),), name="c"), lambda x,c: x.cast(dtypes.int32).cast(c.dtype)),
     # load/store use pointer arithmetic, and the cast does nothing. NOTE: this doesn't apply to image indexing cause it's 1-D
     (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"), UPat.var("off")), allow_any_len=True, name="x"), lambda x,buf,off: x.replace(
-      src=(buf,off.cast(dtypes.long))+x.src[2:]) if buf.addrspace != AddrSpace.REG and not is_image_shape(buf._shape) else None),
+      src=(buf,UOp.const(off.val, dtypes.long) if off.op is Ops.CONST else off.cast(dtypes.long))+x.src[2:])
+      if buf.addrspace != AddrSpace.REG and not is_image_shape(buf._shape) else None),
     # images need index to be int for nir (coordinates only: the INDEX keeps its access dtype)
     (UPat.var("buf").index(UPat.var("idx_y"), UPat.var("idx_x"), name="x"),
      lambda x,buf,idx_y,idx_x: x.replace(src=(buf, idx_y.cast(dtypes.int), idx_x.cast(dtypes.int)))),

tinygrad/runtime/autogen/__init__.py

@@ -1,4 +1,4 @@
-import glob, importlib, os, pathlib, shutil, subprocess, tarfile, tempfile
+import glob, importlib, os, pathlib, subprocess
 from tinygrad.helpers import fetch, flatten, system, getenv
 
 root = (here:=pathlib.Path(__file__).parent).parents[2]
@@ -31,6 +31,7 @@ def load(name, files, **kwargs):
   if not (f:=(root/(path:=kwargs.pop("path", __name__)).replace('.','/')/f"{name}.py")).exists() or getenv('REGEN'):
     files, kwargs['args'] = files() if callable(files) else files, args() if callable(args:=kwargs.get('args', [])) else args
     if (srcs:=kwargs.pop('srcs', None)):
+      import tempfile, tarfile
       srcpath = (td:=tempfile.TemporaryDirectory(f"autogen-src-{name.replace('/','-')}")).name + "/"
       for src in (srcs if isinstance(srcs, list) else [srcs]):
         if 'tar' in src:
@@ -157,7 +158,7 @@ def __getattr__(nm):
           *[f"python3 src/compiler/nir/nir_{s}_h.py --outdir gen" for s in ["intrinsics", "intrinsics_indices"]]]), cwd=path, shell=True, check=True),
   srcs="https://gitlab.freedesktop.org/mesa/mesa/-/archive/mesa-25.2.7/mesa-25.2.7.tar.gz",
   dll=f"'tinymesa_cpu' if DEV.renderer == 'LVP' else 'tinymesa', {tinymesa_path}, emsg='pip install tinymesa==25.2.7.2'",
-  prolog=["from tinygrad.helpers import DEV", "import gzip, base64, platform, sysconfig, os"],
+  prolog=["from tinygrad.helpers import DEV", "import gzip, base64, sysconfig, os"],
   epilog=lambda path: [system(f"{root}/extra/mesa/lvp_nir_options.sh {path}")])
     case "libclang":
       return load("libclang",

tinygrad/runtime/autogen/mesa.py

@@ -5,7 +5,7 @@ from typing import Literal, TypeAlias
 from tinygrad.runtime.support.c import _IO, _IOW, _IOR, _IOWR
 from tinygrad.runtime.support import c
 from tinygrad.helpers import DEV
-import gzip, base64, platform, sysconfig, os
+import gzip, base64, sysconfig, os
 dll = c.DLL('mesa', 'tinymesa_cpu' if DEV.renderer == 'LVP' else 'tinymesa', os.path.join(sysconfig.get_paths()['platlib'], 'tinymesa'), emsg='pip install tinymesa==25.2.7.2')
 class struct_u_printf_info(c.Struct): pass
 u_printf_info: TypeAlias = struct_u_printf_info

tinygrad/runtime/graph/hcq.py

@@ -6,7 +6,6 @@ from tinygrad.device import Buffer, BufferSpec, Compiled, Device, MultiBuffer, P
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops, Variable
 from tinygrad.engine.jit import GraphRunner, MultiGraphRunner
-from tinygrad.runtime.ops_rdma import RDMACopyQueue
 
 class HCQGraph(MultiGraphRunner):
   def __init__(self, *args, **kwargs):
@@ -50,7 +49,7 @@ class HCQGraph(MultiGraphRunner):
 
     self.comp_queues: dict[HCQCompiled, HWQueue] = {dev: unwrap(dev.hw_compute_queue_t)() for dev in self.devices}
     self.copy_queues: dict[tuple[HCQCompiled, int], HWQueue] = {} # lazy allocation, keyed by (device, queue_idx)
-    self.rdma_queues: dict[tuple[HCQCompiled, HCQCompiled], RDMACopyQueue] = {} # lazy allocation, keyed by device pair
+    self.rdma_queues: dict[tuple[HCQCompiled, HCQCompiled], "RDMACopyQueue"] = {} # lazy allocation, keyed by device pair
     self.num_copy_queues: int = getenv("HCQ_NUM_SDMA", min(len(self.devices), 8) if ALL2ALL >= 1 else 1)
     self.num_rdma_ops: dict[tuple[HCQCompiled, HCQCompiled], int] = collections.defaultdict(int)
 
@@ -104,6 +103,7 @@ class HCQGraph(MultiGraphRunner):
       elif is_rdma:
         enqueue_queue = self.comp_queues[enqueue_dev]
         rdma_key = (cast(HCQCompiled, Device[bufs[0].device]).rdma_dev(), enqueue_dev.rdma_dev())
+        from tinygrad.runtime.ops_rdma import RDMACopyQueue
         self.rdma_queues.setdefault(rdma_key, RDMACopyQueue(enqueue_dev.rdma_dev()))
       else:
         assert (enqueue_dev.hw_copy_queue_t is not None), "device must implement a copy queue"

tinygrad/runtime/graph/metal.py

@@ -113,5 +113,6 @@ class MetalGraph(GraphRunner):
   @staticmethod
   def supports_uop(batch_devs, new_call:UOp) -> bool:
     # Metal ICB replay encodes offsets as uint32; reject if any Metal buffer offset exceeds 32-bit range.
-    if any(b.op is Ops.SLICE and b.src[1].val * b.src[0].dtype.itemsize > 0xFFFFFFFF for b in new_call.src[1:]): return False
+    for shrink in [s for src in new_call.src[1:] if (s:=src.src[0] if src.op is Ops.BITCAST else src).op is Ops.SHRINK]:
+      if shrink.src[1].val * shrink.src[0].dtype.itemsize > 0xFFFFFFFF: return False
     return GraphRunner.supports_uop(batch_devs, new_call)

tinygrad/runtime/ops_cpu.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
-import platform, sys, os, ctypes, functools, mmap, threading, array, itertools
-from dataclasses import replace
-from typing import cast
-from tinygrad.helpers import to_mv, OSX, WIN, Context, mv_address, suppress_finalizing, unwrap, data64_le, partition
-from tinygrad.device import Buffer, BufferSpec, TinyELF
-from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, HCQArgsState, HCQSignal, HCQProgram, MMIOInterface
-from tinygrad.runtime.support.hcq import CLikeArgsState
+import platform, sys, os, ctypes, functools, mmap, threading, array, struct, time
+from dataclasses import dataclass, replace
+from typing import cast, Callable
+from tinygrad.helpers import to_mv, from_mv, OSX, WIN, Context, mv_address, suppress_finalizing, unwrap, data64_le, to_tuple
+from tinygrad.device import Buffer, BufferSpec, TinyELF, Program, Device
+from tinygrad.runtime.support.hcq import HCQBuffer, MMIOInterface
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, make_cmdbuf, make_signal
 from tinygrad.renderer.cstyle import ClangRenderer
 from tinygrad.renderer.llvmir import CPULLVMRenderer
 from tinygrad.renderer.nir import LVPRenderer
@@ -13,11 +13,15 @@ from tinygrad.renderer.isa.x86 import X86Renderer
 from tinygrad.runtime.support.elf import jit_loader
 from tinygrad.runtime.autogen import libc
 from tinygrad.codegen import do_to_program
+from tinygrad.engine.realize import pm_flatten_linear, get_call_arg_uops, get_runtime
 from tinygrad import UOp, dtypes
 from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import sint, KernelInfo, Ops, UPat, PatternMatcher, graph_rewrite
+from tinygrad.uop.ops import KernelInfo, Ops, UPat, PatternMatcher, graph_rewrite
 
-MAX_ARGS, CMD_SIZE, RING_SLOTS = 63, 64, (16 << 10)
+MAX_ARGS, CMD_SIZE, RING_SLOTS, FUNCS = 63, 64, (16 << 10), (() if WIN else ('clock_gettime', 'sem_wait', 'sem_post'))
+
+# *****************
+# 1. workers
 
 def signal_prog():
   val = UOp.param(1, dtypes.int, (), vmin_vmax=(0, dtypes.int.max), name="value", addrspace=AddrSpace.ALU)
@@ -35,79 +39,86 @@ def timestamp_prog():
     val = ts.after(call)[0].load() * 1_000_000_000 + ts.after(call)[1].load()
   return UOp.param(0, dtypes.uint64, (1,))[0].store(val)
 
-def quit_prog():
-  fn = UOp.param(0, dtypes.uint64, (1 if WIN else 3,))
-  if WIN: return fn[0].load().call(UOp.const(0, dtypes.uint64), ret_dtype=dtypes.void) # ExitThread(0)
-  sem = UOp.param(1, dtypes.uint64, (1,))
-
-  close = fn[2].load().call(sem[0], ret_dtype=dtypes.void) # sem_close(sem)
-  return fn.after(close)[0].load().call(UOp.const(0, dtypes.uint64), ret_dtype=dtypes.void) # pthread_exit(0)
-
 def worker_prog():
   ring = UOp.param(0, dtypes.uint64, (RING_SLOTS * CMD_SIZE,), volatile=True)
-  wait, sem = UOp.param(1, dtypes.uint64, (1,), volatile=True), UOp.param(2, dtypes.uint64, (1,))
-  cur = UOp.range(2**64-1, 0, dtype=dtypes.uint64)
+  wait, done = UOp.param(1, dtypes.uint64, (1,), volatile=True), UOp.param(2, dtypes.uint64, (1,), volatile=True)
+  sem, cur = UOp.param(3, dtypes.uint64, (1,)), UOp.range(2**64-1, 0, dtype=dtypes.uint64) # sem is unused on windows, it has to come last
 
   # spin on windows, sem_wait to sleep on posix
   if WIN: ready = (v:=wait.after(lw:=UOp.loop(1), cur)[0].load()).end(lw, v <= cur)
   else: ready = (rv:=wait.after(lw:=UOp.loop(1), cur)[0].load().call(sem.after(cur)[0], ret_dtype=dtypes.int)).end(lw, rv != 0)
 
   entry = [ring.after(ready).index((cur % RING_SLOTS) * CMD_SIZE + i).load() for i in range(CMD_SIZE)]
-  return entry[0].call(*entry[1:], ret_dtype=dtypes.void).end(cur)
+  return done.after(entry[0].call(*entry[1:], ret_dtype=dtypes.void)).index(0).store(cur + 1).end(cur)
 
-def host_wait(ctx, dst:UOp, val:UOp) -> UOp:
-  return (cur:=dst.after(loop:=UOp.loop(next(ctx))).index(UOp.const(0, dtypes.int)).load()).end(loop, cur < val)
+@dataclass
+class CPUWorker: ring:Buffer; put:Buffer; sem:Buffer; sys:Buffer; done:Buffer; thread:threading.Thread  # noqa: E702
 
-pm_host_opsel = PatternMatcher([(UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val"))), host_wait)])
+# *****************
+# 2. queue encoders
 
-def encode_host_queue(q:UOp) -> UOp:
-  # TODO: subset of hcq2 for now
-  spins, (store,) = partition(graph_rewrite(q, pm_host_opsel, ctx=itertools.count(), walk=True, name="host opsel").src, lambda u: u.op is Ops.END)
-  assert store.op is Ops.INS and store.arg == "store", f"host queue cannot encode {store.op} {store.arg}"
-  return store.src[0].after(*spins).index(UOp.const(0, dtypes.int)).store(store.src[1])
+def cpu_cmd(devs:tuple[str, ...], prog, *args:UOp) -> UOp:
+  progs = [get_runtime(d, prog) if isinstance(prog, UOp) else cast(CPUDevice, Device[d]).prgs[prog] for d in devs]
+  addrs = tuple(UOp.const(p.addr, dtypes.uint64) for p in progs)
+  words = ((addrs[0] if len(addrs) == 1 else UOp(Ops.STACK, dtypes.uint64, addrs)),) + args
+  return UOp(Ops.INS, dtypes.void, words + (UOp.const(0, dtypes.uint64),) * (CMD_SIZE - len(words)), arg="cmd")
 
-class CPUComputeQueue(HWQueue):
-  def __init__(self, dev):
-    super().__init__()
-    self.dev = dev
-  def _cmd(self, prog, args=(), vals=()): return self.exec(prg:=self.dev.prgs[prog], prg.fill_kernargs(args, vals), None, None)
-  def memory_barrier(self): return self
-  def exec(self, prg:CPUProgram, args_state:HCQArgsState, global_size, local_size):
-    if (lvp:=isinstance(args_state, LVPArgsState)): self.bind_args_state(args_state)
-    args:list[sint|None] = [args_state.buf.va_addr] if lvp else [*[x.va_addr for x in args_state.bufs], *args_state.vals]
-    assert len(args) <= MAX_ARGS, f"CPU programs support at most {MAX_ARGS} arguments, got {len(args)}"
-    for tid in range(1 if lvp else (global_size or (1,))[0]):
-      if not lvp and 'core_id' in prg.runtimevars: args[prg.runtimevars['core_id']] = tid
-      self.q(prg, *[unwrap(x) for x in args], *([0] * (MAX_ARGS - len(args))))
-    return self
-  def wait(self, signal, value=0): return self._cmd(wait_prog, (signal.base_buf,), (value,))
-  def timestamp(self, signal): return self._cmd(timestamp_prog, (signal.base_buf.offset(8, 8), self.dev.func_table._buf.offset(0, 8)))
-  def signal(self, signal, value:sint=0): return self._cmd(signal_prog, (signal.base_buf,), (value,))
-  def _submit(self, dev):
-    dev.ensure_worker()
-    ring_view = dev.ring.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
-    for off in range(0, len(self._q), CMD_SIZE):
-      entry = [self._q[off].addr, *self._q[off+1:off+CMD_SIZE]]
-      ring_view[(base:=(dev.ring_pos % RING_SLOTS) * CMD_SIZE):base+CMD_SIZE] = array.array('Q', (int(x) & ((1<<64)-1) for x in entry))
-      dev.ring_pos += 1
-      if WIN: dev.sys.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = dev.ring_pos
-      else: assert libc.sem_post(dev.sem) == 0
+def cpu_exec(ctx:tuple[str, ...], call:UOp, prg:UOp) -> UOp:
+  args = [get_call_arg_uops(call)[i].getaddr(ctx) for i in prg.arg.globals] + [v.cast(dtypes.uint64) for v in prg.arg.vars]
+  if (core:=prg.arg.runtimevars.get('core_id')) is None: return cpu_cmd(ctx, prg, *args)
 
-class LVPArgsState(CLikeArgsState):
-  def __init__(self, buf, prg, bufs, vals=()): super().__init__(buf, prg, bufs, vals, [*data64_le(buf.va_addr + 12), (len(bufs) + len(vals)) * 2])
+  la = [cpu_cmd(ctx,prg,*args[:(cid:=(len(prg.arg.globals)+core))],UOp.const(t, dtypes.uint64),*args[cid+1:]) for t in range(prg.arg.global_size[0])]
+  return UOp(Ops.LINEAR, dtypes.void, tuple(la))
+
+pm_cpu_opsel = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), cpu_exec),
+
+  (UPat(Ops.INS, arg="barrier"), lambda: UOp(Ops.NOOP, dtypes.void, ())),
+  (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val"))),
+   lambda ctx, dst, val: cpu_cmd(ctx, wait_prog, dst.getaddr(ctx), val.cast(dtypes.uint64))),
+  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))),
+   lambda ctx, dst, val: cpu_cmd(ctx, signal_prog, dst.getaddr(ctx), val.cast(dtypes.uint64))),
+  (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),)),
+   lambda ctx, dst: cpu_cmd(ctx, timestamp_prog, dst.getaddr(ctx), *(() if WIN else (make_signal(ctx, tag="func:clock_gettime").getaddr(ctx),)))),
+])
+
+def encode_queue(q:UOp) -> UOp:
+  devs, queue = to_tuple(q.arg[0]), q.arg[1]
+  lin = graph_rewrite(q, pm_cpu_opsel+pm_flatten_linear, ctx=devs, walk=True, name=f"{queue} opsel")
+
+  cnt = sum(len(ins.src) for ins in lin.src) // CMD_SIZE
+  assert cnt < RING_SLOTS, f"submit of {cnt} entries doesn't fit the ring"
+  cmdbuf = make_cmdbuf(lin, devs, buf=UOp.placeholder((cnt*CMD_SIZE,), dtypes.uint64, next(UOp.unique_num), device=devs).rtag("cmdbuf"))
+  ring = UOp.placeholder((ring_words:=RING_SLOTS*CMD_SIZE,), dtypes.uint64, 0, device=devs, volatile=True).rtag(f"{queue}_ring")
+  put, done, sem, sysbuf = (make_signal(devs, tag=f"{queue}_{name}") for name in ("put", "done", "sem", "sys"))
+
+  # submits are serialized on the submitter, so they can bump put without atomics
+  ran = done.after(l:=UOp.loop(next(UOp.unique_num))).index(0).load()
+  room = ran.end(l, put.index(0).load() - ran > RING_SLOTS - cnt) # wait until cnt entries fit in the ring
+  base = ((put.after(room).index(0).load() % RING_SLOTS) * CMD_SIZE).cast(dtypes.int)
+  e = UOp.range(cnt, next(UOp.unique_num), dtype=dtypes.int, src=(cmdbuf, ring))
+  copy = UOp.group(*[ring.index((base + e*CMD_SIZE + w) % ring_words).store(cmdbuf.index(e*CMD_SIZE + w).load()) for w in range(CMD_SIZE)])
+
+  # wake the worker after each entry, keeping the post with the stores stops it from hoisting out of the loop
+  wake = copy.end(e) if WIN else make_signal(devs, tag="func:sem_post").after(copy).index(0).load().call(sem.index(0), ret_dtype=dtypes.void).end(e)
+  bumped = put.after(wake).index(0).store(put.index(0).load() + cnt)
+  return sysbuf.after(bumped).index(0).store(put.index(0).load() + cnt) if WIN else bumped
+
+# *****************
 
 # NOTE: MAP_JIT is added to mmap module in python 3.13
 MAP_JIT = 0x0800
 
-class CPUProgram(HCQProgram['CPUDevice']):
+class CPUProgram(Program['CPUDevice']):
   rt_lib = None
   try: rt_lib = ctypes.CDLL(ctypes.util.find_library('System' if OSX else 'kernel32') if OSX or WIN else 'libgcc_s.so.1')
   except OSError: pass
 
   def __init__(self, dev:CPUDevice, obj:TinyELF):
-    self.signature, self.runtimevars = obj.signature, {name:slot for name,slot,*_ in obj.signature if name == 'core_id'}
+    self.dev, self.name, self.signature = dev, obj.name, obj.signature
+    self.runtimevars = {name:slot for name,slot,*_ in obj.signature if name == 'core_id'}
+    self.lvp = obj.target.renderer == "LVP"
 
-    LVP = obj.target.renderer == "LVP"
     if sys.platform == "win32": # mypy doesn't understand when WIN is used here
       PAGE_EXECUTE_READWRITE, MEM_COMMIT, MEM_RESERVE = 0x40, 0x1000, 0x2000
       ctypes.windll.kernel32.VirtualAlloc.restype = ctypes.c_void_p
@@ -117,7 +128,7 @@ class CPUProgram(HCQProgram['CPUDevice']):
       ctypes.windll.kernel32.GetCurrentProcess.restype = ctypes.c_void_p
       proc = ctypes.windll.kernel32.GetCurrentProcess()
       ctypes.windll.kernel32.FlushInstructionCache(ctypes.c_void_p(proc), ctypes.c_void_p(self.addr), ctypes.c_size_t(len(obj.lib)))
-      self.fxn = ctypes.CFUNCTYPE(None)(self.addr)
+      self.fxn = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(self.addr) if self.lvp else ctypes.CFUNCTYPE(None)(self.addr)
     else:
       # On apple silicon with SPRR enabled (it always is in macos) RWX pages are unrepresentable: https://blog.svenpeter.dev/posts/m1_sprr_gxf/
       # MAP_JIT allows us to easily flip pages from RW- to R-X and vice versa. It is a noop on intel cpus. (man pthread_jit_write_protect_np)
@@ -125,7 +136,7 @@ class CPUProgram(HCQProgram['CPUDevice']):
       self.addr = mv_address(self.mem)
 
       if OSX: unwrap(CPUProgram.rt_lib).pthread_jit_write_protect_np(False)
-      lib = jit_loader(obj.lib, base=ctypes.addressof(ctypes.c_void_p.from_buffer(self.mem)), link_libs=['m']) if LVP else obj.lib
+      lib = jit_loader(obj.lib, base=ctypes.addressof(ctypes.c_void_p.from_buffer(self.mem)), link_libs=['m']) if self.lvp else obj.lib
       self.mem.write(lib)
       if OSX: unwrap(CPUProgram.rt_lib).pthread_jit_write_protect_np(True)
 
@@ -138,15 +149,30 @@ class CPUProgram(HCQProgram['CPUDevice']):
         # msync should be a universal POSIX way to do this
         libc.msync(ctypes.c_void_p(self.addr), len(lib), libc.MS_SYNC | libc.MS_INVALIDATE)
 
-      self.fxn = ctypes.CFUNCTYPE(None)(self.addr)
+      self.fxn = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(self.addr) if self.lvp else ctypes.CFUNCTYPE(None)(self.addr)
 
-    super().__init__(LVPArgsState if LVP else HCQArgsState, dev, obj, kernargs_alloc_size=12+256 if LVP else 0)
+  def __call__(self, *bufs:HCQBuffer, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),
+               vals:tuple[int|None, ...]=(), wait:bool=False, timeout:int|None=None) -> float|None:
+    st = time.perf_counter()
+    if self.lvp:
+      lvp_args = bytearray(12 + (len(bufs) + len(vals)) * 8)
+      addr = mv_address(lvp_args)
+      struct.pack_into(f'<3I{len(bufs)}Q', lvp_args, 0, *data64_le(addr+12), (len(bufs)+len(vals))*2, *[b.va_addr for b in bufs])
+      for v,(off,dt) in zip(vals, TinyELF.iter_sig(self.signature[-len(vals):], len(bufs)*8)): struct.pack_into(f'<{dt.fmt}', lvp_args, 12+off, v)
+      self.fxn(addr)
+    else:
+      args = [*[cast(int, b.va_addr) for b in bufs], *cast(tuple[int, ...], vals)]
+      assert len(args) <= MAX_ARGS, f"CPU programs support at most {MAX_ARGS} arguments, got {len(args)}"
+      for tid in range(global_size[0]):
+        if 'core_id' in self.runtimevars: args[self.runtimevars['core_id']] = tid
+        self.fxn(*[ctypes.c_uint64(x) for x in args])
+    return time.perf_counter() - st if wait else None
 
   @suppress_finalizing
   def __del__(self):
     if sys.platform == 'win32': ctypes.windll.kernel32.VirtualFree(ctypes.c_void_p(self.addr), ctypes.c_size_t(0), 0x8000) #0x8000 - MEM_RELEASE
 
-class CPUAllocator(HCQAllocator):
+class CPUAllocator(HCQAllocator['CPUDevice']):
   def __init__(self, dev:CPUDevice): super().__init__(dev, supports_copy_from_disk=False, supports_transfer=False)
   def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
     if options.external_ptr is not None: addr, buf = options.external_ptr, None
@@ -154,68 +180,55 @@ class CPUAllocator(HCQAllocator):
     else: addr = mv_address(buf:=mmap.mmap(-1, size, mmap.MAP_ANON | mmap.MAP_SHARED, mmap.PROT_READ | mmap.PROT_WRITE))
     return HCQBuffer(va:=addr, sz:=size, meta=buf, view=MMIOInterface(va, sz, fmt='B'), owner=self.dev)
   def _as_buffer(self, src) -> memoryview: return to_mv(src.va_addr, src.size)
+  def _copyin(self, dest:HCQBuffer, src:memoryview):
+    self.dev.synchronize()
+    ctypes.memmove(int(dest.va_addr), from_mv(src), len(src))
+  def _copyout(self, dest:memoryview, src:HCQBuffer):
+    self.dev.synchronize()
+    ctypes.memmove(from_mv(dest), int(src.va_addr), len(dest))
   def _do_map(self, buf:HCQBuffer):
     if buf.view is None or not isinstance(buf.view, MMIOInterface): raise RuntimeError("Cannot map buffer without view to cpu")
     return HCQBuffer(buf.view.addr, buf.size, view=buf.view, owner=buf.owner)
   def _unmap(self, mb): pass  # CPU _do_map returns a view wrapper, nothing to release
 
-class CPUDevice(HCQCompiled):
-  pm_lower = PatternMatcher([
-    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),)), encode_host_queue)])
-
-  pm_bufferize = PatternMatcher([
-    (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx[0].signal("sentinel", (1 << 64) - 1)),
-    (UPat(Ops.PARAM, tag="timeline_signal"), lambda ctx: ctx[0].signal("timeline")),
-    (UPat(Ops.PARAM, tag="timeline_value"), lambda ctx: ctx[0].signal("value", 1)),
-    (UPat(Ops.PARAM, tag="signal", name="b"), lambda ctx, b: ctx[0].signal(b.arg.slot)),
-  ])
-
-  @functools.cache
-  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
+class CPUDevice(HCQ2Compiled):
+  wait_timeout_ms, has_copy_queue = 30000, False
+  pm_lower = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),)), encode_queue)])
 
   def __init__(self, device:str=""):
-    super().__init__(device, CPUAllocator(self), [ClangRenderer, CPULLVMRenderer, LVPRenderer, X86Renderer], CPUProgram, HCQSignal,
-      functools.partial(CPUComputeQueue, self), arch={'amd64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine().lower(), m)+",native")
+    super().__init__(device, CPUAllocator(self), [ClangRenderer, CPULLVMRenderer, LVPRenderer, X86Renderer], CPUProgram,
+      arch={'amd64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine().lower(), m)+",native")
 
-    self.ring_pos = 0
+    self.pm_bufferize = PatternMatcher(
+      [(UPat(Ops.PARAM, tag=f"COMPUTE:0_{n}"), lambda ctx, n=n: getattr(ctx[0].worker, n)) for n in ("ring", "put", "sem", "sys", "done")] +
+      [(UPat(Ops.PARAM, tag=f"func:{f}"), lambda ctx, f=f: ctx[0].func_ptr(f)) for f in FUNCS]) + self.pm_bufferize
 
-    # posix uses sem to put cpus into sleep
-    self.sem_addr = 0
-    if not WIN:
-      self.sem = libc.sem_open(sem_name:=f"/tinygrad-{os.getpid()}-{id(self):x}".encode(), os.O_CREAT|os.O_EXCL, 0o600, 0) # type: ignore[call-arg]
-      self.sem_addr = unwrap(ctypes.cast(self.sem, ctypes.c_void_p).value)
-      if self.sem_addr == ctypes.c_void_p(-1).value or libc.sem_unlink(sem_name): raise OSError(ctypes.get_errno(), "semaphore")
-
-    # TODO: move to hcq2
     with Context(EMULATED_DTYPES="", TRACK_MATCH_STATS=0):
-      prgs = {f: f().sink(arg=KernelInfo(f.__name__), tag=1) for f in (signal_prog, wait_prog, timestamp_prog, quit_prog, worker_prog)}
-      self.prgs = {f: self.runtime(do_to_program(v, ClangRenderer(replace(self.renderer.target, renderer="CLANG"))).to_elf()) for f,v in prgs.items()}
+      clang = ClangRenderer(replace(self.renderer.target, renderer="CLANG"))
+      self.prgs:dict[Callable, CPUProgram] = {f: CPUProgram(self, do_to_program(f().sink(arg=KernelInfo(f.__name__), tag=1), clang).to_elf())
+                                              for f in (signal_prog, wait_prog, timestamp_prog, worker_prog)}
 
-  @functools.cached_property
-  def ring(self) -> Buffer: return Buffer(self.device, RING_SLOTS * CMD_SIZE, dtypes.uint64, preallocate=True)
-  @functools.cached_property
-  def sys(self) -> Buffer: return Buffer(self.device, 1, dtypes.uint64, preallocate=True)
-  @functools.cached_property
-  def sem_buf(self) -> Buffer: return Buffer(self.device, 1, dtypes.uint8, options=BufferSpec(external_ptr=self.sem_addr), preallocate=True)
+  def func_ptr(self, name:str) -> Buffer: return self.func_table.view(1, dtypes.uint64, FUNCS.index(name)*8).ensure_allocated()
 
-  # TODO: move to hcq2 infra
   @functools.cached_property
   def func_table(self) -> Buffer:
-    fns = ([0, ctypes.windll.kernel32.ExitThread, 0, 0] if WIN else  # type: ignore[attr-defined]
-           [libc.dll.clock_gettime, libc.dll.pthread_exit, libc.dll.sem_wait, libc.dll.sem_close])
-    addrs = array.array('Q', [unwrap(ctypes.cast(f, ctypes.c_void_p).value) if f else 0 for f in fns])
-    (ft:=Buffer(self.device, len(fns), dtypes.uint64, preallocate=True)).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[:] = addrs
+    lib = ctypes.windll.kernel32 if sys.platform == "win32" else libc.dll # type: ignore[attr-defined]
+    (ft:=Buffer(self.device, len(FUNCS), dtypes.uint64, preallocate=True))._buf.cpu_view().view(fmt='Q')[:] = \
+      array.array('Q', [unwrap(ctypes.cast(getattr(lib, f), ctypes.c_void_p).value) for f in FUNCS])
     return ft
 
-  @functools.cache
-  def ensure_worker(self):
-    threading.Thread(target=cast(CPUProgram, self.prgs[worker_prog]).fxn, daemon=True, args=[ctypes.c_uint64(x) for x in
-      [self.ring._buf.va_addr, self.sys._buf.va_addr if WIN else self.func_table._buf.va_addr+16, self.sem_addr]]).start()
+  @functools.cached_property
+  def worker(self) -> CPUWorker:
+    ring, put, sysbuf, done = (Buffer(self.device, sz, dtypes.uint64, preallocate=True) for sz in (RING_SLOTS*CMD_SIZE, 1, 1, 1))
+    addr, hsem = 0, None
 
-  def finalize(self):
-    if self.ring_pos == 0: return # the worker starts with the first submit
-    ft = self.func_table._buf
-    CPUComputeQueue(self)._cmd(quit_prog, (ft.offset(8, 8),) if WIN else (ft.offset(8, 24), self.sem_buf._buf)).submit(self)
-    self.ring_pos = 0
+    # sem are posix-only
+    if not WIN:
+      hsem = libc.sem_open(nm:=f"/tinygrad-{os.getpid()}-{id(ring):x}".encode(), os.O_CREAT|os.O_EXCL, 0o600, 0) # type: ignore[call-arg]
+      if (addr:=unwrap(ctypes.cast(hsem, ctypes.c_void_p).value)) == ctypes.c_void_p(-1).value or libc.sem_unlink(nm):
+        raise OSError(ctypes.get_errno(), "semaphore")
+    sem = Buffer(self.device, 1, dtypes.uint64, options=BufferSpec(external_ptr=addr), preallocate=True)
+
+    worker_args = [ring._buf.va_addr, sysbuf._buf.va_addr if WIN else self.func_ptr('sem_wait')._buf.va_addr, done._buf.va_addr, addr]
+    (worker:=threading.Thread(target=self.prgs[worker_prog].fxn, daemon=True, args=[ctypes.c_uint64(x) for x in worker_args])).start()
+    return CPUWorker(ring, put, sem, sysbuf, done, worker)

tinygrad/runtime/ops_disk.py

@@ -49,7 +49,8 @@ class DiskDevice(Compiled):
     DiskDevice._tried_io_uring_init = True
 
     if sys.platform == 'linux' and not hasattr(sys, "getandroidapilevel"):
-      fd = libc.syscall(io_uring.NR_io_uring_setup, 4096, ctypes.byref(p:=io_uring.struct_io_uring_params()))
+      p = io_uring.struct_io_uring_params(flags=io_uring.IORING_SETUP_SQPOLL, sq_thread_idle=0xffffffff)
+      fd = libc.syscall(io_uring.NR_io_uring_setup, 4096, ctypes.byref(p))
       if fd < 0: return
 
       sq_ptr = libc.mmap(0, p.sq_off.array + p.sq_entries * 4, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | MAP_POPULATE, fd, 0)
@@ -67,6 +68,7 @@ class DiskDevice(Compiled):
         kring_mask=u32ptr(sq_ptr+p.cq_off.ring_mask), cqes=ctypes.cast(cq_ptr+p.cq_off.cqes, ctypes.POINTER(io_uring.struct_io_uring_cqe)))
 
       DiskDevice.io_uring = io_uring.struct_io_uring(ring_fd=fd, sq=sqdesc, cq=cqdesc) # type: ignore
+      libc.syscall(io_uring.NR_io_uring_enter, fd, 0, 0, io_uring.IORING_ENTER_SQ_WAKEUP)
 
 class DiskBuffer:
   def __init__(self, device:DiskDevice, size:int, offset=0):
@@ -124,7 +126,6 @@ class DiskAllocator(Allocator):
         # Send sqe
         DiskDevice.io_uring.sq.array[sqe_index] = sqe_index
         DiskDevice.io_uring.sq.ktail[0] = tail + 1
-        libc.syscall(io_uring.NR_io_uring_enter, DiskDevice.io_uring.ring_fd, 1, 1, io_uring.IORING_ENTER_GETEVENTS)
 
         reqs.append((copy_batch, copied_in, minor_offset, real_copy_size:=min(sqe.len - minor_offset, size - copied_in)))
         next_read_offset += sqe.len

tinygrad/runtime/ops_python.py

@@ -3,10 +3,10 @@
 # works to test the tensor cores, and all the uops in general
 # this is the (living) definition of uops
 from typing import Any, TYPE_CHECKING
-import pickle, base64, itertools, time, sys, functools
+import pickle, base64, itertools, time, sys, functools, ctypes
 from dataclasses import replace
 from tinygrad.dtype import bitcast, DType, dtypes, AddrSpace, truncate, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar
-from tinygrad.helpers import all_same, getenv, flatten, Target, IMAGE, is_image_shape, cpu_profile
+from tinygrad.helpers import all_same, getenv, flatten, Target, IMAGE, is_image_shape, cpu_profile, mv_address
 from tinygrad.device import Buffer, Compiled, Compiler, Allocator, Program, TinyELF
 from tinygrad.codegen.opt import tc
 from tinygrad.uop.ops import exec_alu, python_alu, Ops, UOp, GroupOp
@@ -134,6 +134,13 @@ class PythonProgram(Program['PythonDevice']):
                                for k in range(len(src_values))], j, u.dtype) for j in range(load_sz)]
           else:
             values[u] = load(src_values, 0, u.dtype)
+        elif u.op is Ops.CALL:
+          assert u.dtype is dtypes.void
+          cfunc = ctypes.CFUNCTYPE(None, *[ctypes.c_uint64] * (len(src_values)-1))
+          values[u] = []
+          for args,gate in zip(zip(*src_values), exec_masks[-1]):
+            call_args = [(mv_address(x[0]) + x[1]*dt.itemsize) if isinstance(x, tuple) else x for x,dt in zip(args, src_dtypes)]
+            values[u].append(cfunc(call_args[0])(*call_args[1:]) if gate else None)
         elif u.op is Ops.WMMA:
           first_src_dtype = u.src[0].dtype
           assert isinstance(first_src_dtype, DType) # mypy

tinygrad/runtime/support/compiler_qcom.py

@@ -1,6 +1,6 @@
-import ctypes, struct
+import ctypes, struct, platform, pathlib, shutil, subprocess, sys, tarfile, tempfile
 from tinygrad.device import Compiler
-from tinygrad.helpers import DEBUG, system
+from tinygrad.helpers import DEBUG, system, fetch, unwrap
 from tinygrad.runtime.support.compiler_mesa import disas_adreno
 # see https://github.com/sirhcm/tinydreno
 from tinygrad.runtime.autogen import llvm_qcom
@@ -10,10 +10,17 @@ def _read_lib(lib, off) -> int: return struct.unpack("I", lib[off:off+4])[0]
 class QCOMCompiler(Compiler):
   def __init__(self, arch:str):
     assert arch.split(',')[0] == "a630", "only a630 supported"
-    self.arch, self.chip_id, self.llvm_inst = arch, 0x6030001, llvm_qcom.cl_compiler_create_llvm_instance()
+    if platform.machine() == "aarch64": self.arch, self.chip_id, self.llvm_inst = arch, 0x6030001, llvm_qcom.cl_compiler_create_llvm_instance()
+    else:
+      self.arch, self.chip_id, self.fs = arch, 0x6030001, tempfile.TemporaryDirectory()
+      with tarfile.open(fetch('https://git.tinygrad.win/sirhcm/images/releases/download/v2/qcomcl.tar.gz')) as t: t.extractall(fs:=self.fs.name)
+      if (qemu:=shutil.which("qemu-aarch64-static")): argv = f"{qemu} -cpu max,pauth=off -L {fs} {fs}/usr/bin/python3 {__file__} {arch}"
+      else: argv = (f"docker run --rm -i --platform linux/aarch64 -v {fs}/usr:/usr -v {pathlib.Path(__file__).parents[2]}:/tinygrad "
+                    f"-e PYTHONPATH=/ -e QEMU_CPU=max,pauth=off gcr.io/distroless/static python3 /tinygrad/runtime/support/compiler_qcom.py {arch}")
+      self.compiler_process = subprocess.Popen(argv.split(), stdout=subprocess.PIPE, stdin=subprocess.PIPE, bufsize=0)
     super().__init__(f"compile_qcomcl_{arch}")
 
-  def __del__(self): llvm_qcom.cl_compiler_destroy_llvm_instance(self.llvm_inst)
+  def __del__(self): llvm_qcom.cl_compiler_destroy_llvm_instance(self.llvm_inst) if platform.machine() == "aarch64" else self.compiler_process.kill()
 
   def __reduce__(self): return QCOMCompiler, (self.arch,)
 
@@ -25,6 +32,10 @@ class QCOMCompiler(Compiler):
     return handle
 
   def compile(self, src) -> bytes:
+    if platform.machine() != "aarch64":
+      unwrap(self.compiler_process.stdin).write(struct.pack("I", len(src.encode())) + src.encode())
+      if (lib:=unwrap(self.compiler_process.stdout).read(struct.unpack("I", unwrap(self.compiler_process.stdout).read(4))[0])): return lib
+      raise RuntimeError("QCOM Compilation Error")
     ch = self.checked(llvm_qcom.cl_compiler_compile_source(self.llvm_inst, self.chip_id, llvm_qcom.CL_MODE_64BIT, b"", 0, 0, 0, src.encode(), 0,
                                                            llvm_qcom.CL_SRC_STR, None))
     if DEBUG >= 8: print(system("llvm-dis", input=ctypes.string_at((comp:=ch.contents.compiled.contents).llvm_bitcode, comp.llvm_bitcode_size)))
@@ -36,3 +47,13 @@ class QCOMCompiler(Compiler):
     return ret
 
   def disassemble(self, lib: bytes): disas_adreno(lib[(ofs:=_read_lib(lib, 0xc0)):ofs+_read_lib(lib, 0x100)], self.chip_id)
+
+if __name__ == "__main__":
+  compiler = QCOMCompiler(sys.argv[1])
+  while (amt:=sys.stdin.buffer.read(4)):
+    try: lib = compiler.compile(sys.stdin.buffer.read(struct.unpack("I", amt)[0]).decode())
+    except Exception as e:
+      lib = b""
+      print(e, file=sys.stderr, flush=True)
+    sys.stdout.buffer.write(struct.pack("I", len(lib)) + lib)
+    sys.stdout.buffer.flush()

tinygrad/runtime/support/hcq.py

@@ -517,8 +517,7 @@ class HCQCompiled(Compiled, Generic[SignalType]):
     if hasattr(self, 'iface') and hasattr(self.iface, 'device_fini'): self.iface.device_fini()
 
 class HCQBuffer:
-  def __init__(self, va_addr:sint, size:int, meta:Any=None, _base:HCQBuffer|None=None, view:MMIOInterface|None=None,
-               owner:HCQCompiled|None=None):
+  def __init__(self, va_addr:sint, size:int, meta:Any=None, _base:HCQBuffer|None=None, view:MMIOInterface|None=None, owner:Any=None):
     self.va_addr, self.size, self.meta, self._base, self.view = va_addr, size, meta, _base, view
     self._devs, self.owner = ([owner] if owner is not None else []), owner
     self._mappings:dict[HCQCompiled, HCQBuffer] = {} # mapping to the other devices

tinygrad/runtime/support/hcq2.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import cast, Callable, TypeVar, Generic, Any, Sequence
+from typing import cast, Callable, TypeVar, Generic, Any, Sequence, Iterable
 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, PROFILE
@@ -7,9 +7,9 @@ from tinygrad.helpers import to_tuple, round_up, partition, data64_le, panic, Co
 from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator, MultiBuffer, DepsTracker
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEntry, ProfileGraphEvent
 from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, rewrite_group, GroupOp
-from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const
+from tinygrad.uop.symbolic import symbolic
 from tinygrad.dtype import dtypes, truncate
-from tinygrad.runtime.support.hcq import MMIOInterface
+from tinygrad.runtime.support.hcq import MMIOInterface, HCQBuffer
 from tinygrad.runtime.support.memory import BumpAllocator
 from tinygrad.renderer import Renderer, Estimates
 from tinygrad.engine.realize import to_program, get_call_arg_uops, get_call_name, get_call_outs_ins, estimate_uop
@@ -22,7 +22,7 @@ HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQ2Compiled')
 
 HCQ_RUNTIME_DEV = ContextVar("HCQ_RUNTIME_DEV", "CPU")
 
-HCQ_DEVS = frozenset(("AMD",))
+HCQ_DEVS = frozenset(("AMD", "CPU"))
 HCQ_P2P_DEVS = HCQ_DEVS | frozenset(("CPU",))
 HCQ_CACHE_TAGS = frozenset(("program", "systems", "template"))
 
@@ -33,13 +33,13 @@ class HCQInfo:
 
   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
+  kernels:tuple[tuple[tuple[str, ...], str, Estimates, tuple[int, ...]], ...] = ()
 
 def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c
 
 def unwrap_mstack(u):
   if u.op is Ops.MSTACK: return tuple(x for s in u.src for x in unwrap_mstack(s))
-  return unwrap_mstack(u.src[0]) if u.op in {Ops.MSELECT, Ops.SLICE} else (u,)
+  return unwrap_mstack(u.src[0]) if u.op is Ops.MSELECT else (u,)
 
 def is_value_known_at_link(val:UOp) -> bool:
   runtime_reads = [u for u in val.toposort() if u.op in (Ops.LOAD, Ops.INDEX)]
@@ -87,21 +87,36 @@ def encode_kernargs_clike(call:UOp, prg:UOp, devs:str|tuple[str, ...]) -> UOp:
 def replace_call_buffers(ctx:tuple[list[UOp], dict[UOp, int]], call:UOp) -> UOp|None:
   bufs, slots = ctx
   for s in call.src[1:]:
-    if s.op not in (Ops.PARAM, Ops.BIND) and slots.setdefault(s, len(bufs)) == len(bufs): bufs.append(s)
-  return call.replace(src=call.src[:1] + tuple(s if s.op in (Ops.PARAM, Ops.BIND) else s.param_like(slots[s]) for s in call.src[1:]))
+    if s.op is not Ops.PARAM and not s.is_bound_var and slots.setdefault(s, len(bufs)) == len(bufs): bufs.append(s)
+  return call.replace(src=call.src[:1] + tuple(s if s.op is Ops.PARAM or s.is_bound_var else s.param_like(slots[s]) for s in call.src[1:]))
 pm_replace_buffers = PatternMatcher([(UPat(Ops.CALL, name="call"), replace_call_buffers)])
 
 # *****************
 # 1.1. prep: staging copies
 
-def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS) and not all_devices_in(b.device, HCQ_P2P_DEVS)
+def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS - {"CPU"}) and not all_devices_in(b.device, HCQ_P2P_DEVS)
+
+def _get_enqueue_devs(call:UOp) -> Any|None:
+  if not (bufs:=call.src[1:]) or not all(all_devices_in(b.device, HCQ_P2P_DEVS) for b in bufs): return None
+  devs = min(bufs, key=lambda b: to_tuple(b.device)[0].startswith("CPU")).device # prio to enqueue on not CPU device
+  return devs if all_devices_in(devs, HCQ_DEVS) else None
 
 def stage_copy(dst:UOp, src:UOp) -> UOp|None:
   if not (_need_staging(src, dst) or _need_staging(dst, src)): return None
 
   stage = UOp.new_buffer("CPU", src.max_numel() * src.dtype.itemsize, dtypes.uint8)
   return UOp(Ops.LINEAR, src=(src.copy_to_device("CPU").call(stage, src), stage.copy_to_device(dst.device).call(dst, stage)))
-pm_insert_copy_staging = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy)])
+
+def kernel_copy(call:UOp, dst:UOp, src:UOp) -> UOp|None:
+  if (devs:=_get_enqueue_devs(call)) is None or Device[(dev:=to_tuple(devs)[0])].has_copy_queue: return None
+  d, s = (UOp.param(i, dst.dtype, (n:=dst.max_numel(),), device=devs) for i in range(2))
+  ast = d.index(r:=UOp.range(n, 0)).store(s.index(r).load()).end(r).sink(arg=KernelInfo(name="copy"), tag=1)
+  return call.replace(src=(to_program(ast, Device[dev].renderer), dst, src))
+
+pm_insert_copy_staging = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy),
+  (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src")), name="call"), kernel_copy)
+])
 
 # *****************
 # 2. deps
@@ -124,19 +139,19 @@ def _get_deps(ctx:DepsTracker, bufs_by_lane:list[list[Any]], write, key:tuple[tu
 
 def _build_wait_cmds(slots:dict[str, int], dep_lanes:list[tuple[tuple, int, int]], devices:tuple[str, ...], queue:str) -> tuple[list[UOp], set[int]]:
   # opt1: same-queue ops are fifo-ordered
-  if devices[0].split(":")[0] in {"AMD", "QCOM"} or queue.startswith("COPY"):
+  if devices[0].split(":")[0] in {"AMD", "QCOM", "CPU"} or queue.startswith("COPY"):
     dep_lanes = [(dep, dlane, lane) for dep, dlane, lane in dep_lanes if (dep[0][dlane], dep[1]) != (devices[lane], queue)]
 
   # opt2: keep latest dep per (dep device, queue, cur lane)
   latest = {((dep[0][dlane], dep[1]), lane): (dep, dlane) for dep, dlane, lane in sorted(dep_lanes, key=lambda x: x[0][2])}
-  deps:dict[tuple, list[int|None]] = collections.defaultdict(lambda: [None]*len(devices))
-  for (_, lane), (dep, dlane) in latest.items(): deps[dep][lane] = dlane
+  deps:dict[tuple, dict[int, list[int]]] = collections.defaultdict(lambda: collections.defaultdict(list))
+  for (_, lane), (dep, dlane) in latest.items(): deps[dep][lane].append(dlane)
 
   waits = []
-  for (ddevs, dqueue, dtag), lanes in deps.items():
-    sig = UOp.mstack(*[make_signal(d, tag="sentinel_signal") if dl is None else make_signal(ddevs[dl], slots[dqueue])
-                       for dl, d in zip(lanes, devices)])
-    waits.append(UOp(Ops.INS, arg="wait", src=(sig, UOp.const(dtag + 1, dtypes.uint64))))
+  for (ddevs, dqueue, dtag), by_lane in deps.items():
+    for ls in itertools.zip_longest(*(by_lane[lane] for lane in range(len(devices)))):
+      s = UOp.mstack(*[make_signal(d, tag="sentinel_signal") if dl is None else make_signal(ddevs[dl], slots[dqueue]) for dl, d in zip(ls, devices)])
+      waits.append(UOp(Ops.INS, arg="wait", src=(s, UOp.const(dtag + 1, dtypes.uint64))))
   return waits, {dtag for _, _, dtag in deps}
 
 def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[tuple[tuple[str, ...], str]],
@@ -147,7 +162,7 @@ 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, signal_tags = len(batch_info), [], [], set()
+  n, fences, resets, fins, signal_tags = len(batch_info), [], [], [], set()
   for _, devgroup in itertools.groupby(sorted(dev_bufs), key=lambda d: d.split(":")[0]):
     devs = tuple(devgroup)
 
@@ -161,16 +176,17 @@ def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[t
     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()
 
-    # fence once per device group on this schedule's previous epoch, then reset any queue signals used by the group
+    # fence once per device group on this schedule's previous epoch
     qs = dedup([qn for bdevs, qn in batch_info if set(bdevs) & set(devs)])
     sched_epoch = make_signal(devs, next(UOp.unique_num))
 
     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(wait_device_epoch), HCQInfo(devs)))
 
-    fences.append(make_call("hcq_fence", UOp.sink(*(resets or [wait_device_epoch])), HCQInfo(devs)))
+    # queues of other groups wait on these signals, so reset them only after every group reached its epoch
+    if qs: resets.append(make_call("hcq_reset", UOp.sink(*[make_signal(devs, slots[q]).index(0).store(0) for q in qs]), 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
+  return fences + resets, 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]
@@ -187,10 +203,10 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> li
     signal_tags |= cur_signal_tags
 
   # build fences and finalizers
-  fences, finalizers, finalizer_signal_tags = _build_finalizers(batch, batch_info, deps_tracker, slots)
+  fences, fins, finalizer_signal_tags = _build_finalizers(batch, batch_info, deps_tracker, slots)
   signal_tags |= finalizer_signal_tags
 
-  src, prof = [], []
+  src, kerns = [], []
   for tag, ((call, _), (devices, queue), q) in enumerate(zip(batch, batch_info, call_waits)):
     # first queue use, sync prior device work with the device timeline
     if batch_info.index((devices, queue)) == tag:
@@ -200,24 +216,24 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> li
     # and make hcq call
     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]
+    kerns.append((devices, name, info.estimates, tuple(ts_ids)))
 
     ts_ins = [UOp(Ops.INS, arg="timestamp", src=(make_signal(devices, s),)) for s in ts_ids]
     q += ts_ins[:1] + [call.replace(arg=replace(call.arg, aux=info))] + ts_ins[1:]
 
     # signal the queue if someone waits for us
     if tag in signal_tags: q += [UOp(Ops.INS, arg="store", src=(make_signal(devices, slots[queue]), UOp.const(tag + 1, dtypes.uint64)))]
-    src.append(make_call(name, make_submit(*q, devs=devices, queue=queue).sink(), info))
+    src.append(make_call(f"submit {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
+  fins = [f.replace(arg=replace(f.arg, aux=replace(a:=f.arg.aux, kernels=tuple(x for x in kerns if set(x[0]) & set(a.device))))) for f in fins]
+  return fences + src + fins
 
 def sched_hcq_batches(l:UOp, profile:bool) -> UOp:
   srcs:list[UOp] = []
   batch:list[tuple[UOp, tuple[str, ...]]] = []
   for call in l.src:
-    if (devs:=next((b.device for b in call.src[1:] if all_devices_in(b.device, HCQ_DEVS)), None)) is not None: batch.append((call, to_tuple(devs)))
+    if (devs:=_get_enqueue_devs(call)) is not None: batch.append((call, to_tuple(devs)))
     else: srcs, batch = srcs + _finalize_batch(batch, profile) + [call], []
   return l.replace(src=tuple(srcs + _finalize_batch(batch, profile)))
 
@@ -342,9 +358,13 @@ pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION
 
 # *****************
 
+def _rank_ranges(uops:Iterable[UOp]) -> dict[UOp, UOp]:
+  return {r: r.replace(arg=(i,)+r.arg[1:]) for i,r in enumerate(sorted([u for u in uops if u.op is Ops.RANGE], key=lambda r: r.arg))}
+
 def replace_params(call:UOp) -> UOp|None:
   body, variables, param_ops = call.src[0], call.src[0].variables(), {Ops.PARAM, Ops.MSTACK}
-  args = dedup([s for u in body.toposort(gate=lambda u: u.op not in param_ops) for s in u.src if s.op in param_ops and s not in variables])
+  tops = body.toposort(gate=lambda u: u.op not in param_ops)
+  args = dedup([s for u in tops for s in u.src if s.op in param_ops and s not in variables])
 
   patched, refhold = partition(call.src[1:], lambda x: x.src[0] in args)
   by_root = {p.src[0]: p for p in patched}
@@ -356,7 +376,7 @@ def replace_params(call:UOp) -> UOp|None:
   refhold += [a for a in addrs if a not in held and all(b.op is not Ops.PARAM or b.tag is not None for b in unwrap_mstack(a))]
 
   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}
+         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} | _rank_ranges(tops)
   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)?
@@ -365,14 +385,15 @@ pm_replace_params = PatternMatcher([
 
 # *****************
 
-def resolve_getaddr_slice(bv:UOp, g:UOp) -> UOp:
+def resolve_getaddr_view(bv:UOp, g:UOp) -> UOp:
   base = bv.src[0].after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ())
-  itemsize = bv.src[0].dtype.itemsize if bv.src[0].without_after.op in (Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize
+  if bv.op is Ops.BITCAST: return UOp(Ops.GETADDR, src=(base,), arg=g.arg)
+  itemsize = bv.src[0].dtype.itemsize if bv.src[0].without_after.op in (Ops.BUFFER, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize
   return UOp(Ops.GETADDR, src=(base,), arg=g.arg) + UOp.const(bv.src[1].val * itemsize, dtypes.uint64)
 
 pm_early_simplify = PatternMatcher([
-  (UPat(Ops.GETADDR, src=(UPat.any(sl:=UPat(Ops.SLICE, name="bv"), sl.after(allow_any_len=True)),), name="g"), resolve_getaddr_slice),
-  (UPat(Ops.INDEX, src=(UPat(Ops.SLICE, name="bv"),), allow_any_len=True, name="x"),
+  (UPat(Ops.GETADDR, src=(UPat((Ops.SHRINK, Ops.BITCAST), name="bv").or_after(),), name="g"), resolve_getaddr_view),
+  (UPat(Ops.INDEX, src=(UPat(Ops.SHRINK, name="bv"),), allow_any_len=True, name="x"),
    lambda bv,x: x.replace(src=(bv.src[0], x.src[1] + bv.src[1].cast(x.src[1].dtype), *x.src[2:]))),
 ])
 
@@ -390,7 +411,7 @@ def pack_hcq_placeholders(call:UOp) -> UOp|None:
       sizes[b.tag] = offs[b] + b.max_numel()
   counts = collections.Counter(b.tag for b in bufs)
   bases = {b.tag:UOp.placeholder((sizes[b.tag],), b.dtype, next(UOp.unique_num), device=b.device).rtag(b.tag) for b in bufs if counts[b.tag] > 1}
-  subs = {b:UOp(Ops.SLICE, b.dtype, (bases[b.tag], UOp.const(offs.get(b, 0))), b.max_numel()) for b in bufs if b.tag in bases}
+  subs = {b:bases[b.tag][(off:=offs.get(b, 0)):off+b.max_numel()] for b in bufs if b.tag in bases}
   return call.replace(src=(call.src[0].substitute(subs, walk=True), *call.src[1:])) if subs else None
 pm_pack_placeholders = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), pack_hcq_placeholders)])
@@ -424,7 +445,7 @@ def hcq_compile(linear:UOp, input_uops:list[UOp]|None, profile:bool) -> UOp:
     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+pm_fold_cast_const, bottom_up=False, name="simplify patches", enter_calls=True)
+    linear = graph_rewrite(linear, pm_early_simplify+symbolic, bottom_up=False, name="simplify patches", enter_calls=True)
     linear = graph_rewrite(linear, pm_split_patches, walk=True, name="split patches")
 
     # and compile it
@@ -456,7 +477,7 @@ def fold_binary(buf:UOp, blob:UOp) -> UOp:
 def fold_const_store(buf:UOp, off:UOp, val:UOp) -> UOp:
   for off,val in zip(off.src, val.src):
     for b,v in zip((bs:=mb.bufs if isinstance((mb:=buf.buffer), MultiBuffer) else (mb,)), val.src if val.op is Ops.STACK else (val,)*len(bs)):
-      data = struct.pack(f'<{v.dtype.fmt}', truncate[v.dtype](v.val))
+      data = struct.pack(f'<{v.dtype.fmt}', truncate[v.dtype]((v.src[0] if v.op is Ops.CAST else v).val))
       b.ensure_allocated().as_memoryview(force_zero_copy=True, no_sync=True).cast('B')[(bo:=off.val*buf.dtype.itemsize):bo+len(data)] = data
   return UOp(Ops.NOOP)
 
@@ -471,7 +492,7 @@ def resolve_getaddr(buf:UOp, g:UOp) -> UOp:
 
 pm_resolve_patches = PatternMatcher([
   # multi
-  (UPat(GroupOp.ALU, src=[UPat(Ops.STACK, name="s"), UPat(Ops.CONST)], name="op"), push_stack),
+  (UPat(GroupOp.ALU, src=[UPat(Ops.STACK, name="s"), UPat.any(UPat(Ops.CONST), UPat(Ops.CAST, src=(UPat(Ops.CONST),)))], name="op"), push_stack),
   (UPat(Ops.CAST, src=(UPat(Ops.STACK, name="s"),), name="op"), push_stack),
 
   # getaddr
@@ -481,7 +502,7 @@ pm_resolve_patches = PatternMatcher([
   (UPat(name="buf").index(UPat(Ops.RANGE), allow_any_len=True)
     .store(UPat.any(UPat(Ops.BINARY, name="blob"), UPat(Ops.BINARY, name="blob").bitcast()).index(UPat(Ops.RANGE), allow_any_len=True).load())
     .end(UPat(Ops.RANGE)), fold_binary),
-  (UPat({Ops.BUFFER, Ops.SLICE, Ops.MSTACK}, name="buf").index(UPat(Ops.STACK, name="off")).store(UPat(Ops.STACK, name="val")), fold_const_store),
+  (UPat({Ops.BUFFER, Ops.MSTACK}, name="buf").index(UPat(Ops.STACK, name="off")).store(UPat(Ops.STACK, name="val")), fold_const_store),
 ])
 
 pm_assert_no_afters = PatternMatcher([(UPat(Ops.AFTER, name="a"), lambda a: panic(RuntimeError, f"AFTER left at hcq_link: {a.src[0].op}"))])
@@ -497,7 +518,7 @@ def hcq_link(linear:UOp, cache=True) -> UOp:
   bufs = {(j,i):a for j,c in enumerate(linear.src) for i,a in enumerate(c.src[1:], 1)
           if a.op is Ops.AFTER and unwrap_mstack(a.src[0])[0].tag in HCQ_CACHE_TAGS}
   linear = linear.substitute({x:link_buf_cache[k] for a in bufs.values() if (k:=link_buf_key(a)) in link_buf_cache for x in (a, a.src[0])}, walk=True)
-  linear = graph_rewrite(linear, pm_resolve_patches+symbolic+pm_fold_cast_const+pm_assert_no_afters, bpm=pm_bufferize, ctx=cache, bottom_up=False,
+  linear = graph_rewrite(linear, pm_resolve_patches+symbolic+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
@@ -508,6 +529,7 @@ def hcq_link(linear:UOp, cache=True) -> UOp:
 
 class HCQ2Compiled(Compiled):
   timestamp_divider: float = 1000.0
+  wait_timeout_ms: float = 30000.0
 
   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
@@ -564,9 +586,10 @@ class HCQ2Compiled(Compiled):
     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()
+    st, done = time.perf_counter(), sig[0]
+    while done < tl[0] - 1:
+      if done != (done:=sig[0]): st = time.perf_counter()
+      elif time.perf_counter() - st > (timeout or self.wait_timeout_ms) / 1000: self.on_device_hang()
     if self.prof_ents: self.collect_prof()
 
   def on_device_hang(self): raise RuntimeError(f"{self.device} hang detected")
@@ -604,21 +627,23 @@ class HCQ2Buffer:
     return HCQ2Buffer(self.va_addr+offset, meta=self.meta, view=(self.view.view(offset=offset, size=size) if self.view is not None else None))
 
 class HCQAllocator(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
-  def _as_buffer(self, buf:HCQ2Buffer) -> memoryview:
+  def _as_buffer(self, buf:HCQBuffer) -> memoryview:
     return unwrap(buf.view).mv
 
-  def _map(self, buf:HCQ2Buffer) -> HCQ2Buffer:
+  def _map(self, buf:HCQBuffer) -> HCQBuffer:
     if not hasattr(self, '_do_map'): raise NotImplementedError("map failed: no method implemented")
     return self._do_map(buf)
 
+  def _do_unmap(self, mb): self.dev.iface.free(mb)
+
   @suppress_finalizing
-  def _free(self, buf:HCQ2Buffer, options:BufferSpec|None=None):
+  def _free(self, buf:HCQBuffer, options:BufferSpec|None=None):
     if options is not None and options.external_ptr is not None: return
     self.dev.synchronize()
     if hasattr(self, '_do_free'): self._do_free(buf, options)
 
   def _unmap(self, mb):
     self.dev.synchronize()
-    self.dev.iface.free(mb)
+    self._do_unmap(mb)
 
-  def _offset(self, buf, size:int, offset:int) -> HCQ2Buffer: return buf.offset(offset=offset, size=size)
+  def _offset(self, buf, size:int, offset:int) -> HCQBuffer: return buf.offset(offset=offset, size=size)

tinygrad/runtime/support/system.py

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

tinygrad/schedule/__init__.py

@@ -8,14 +8,13 @@ from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SC
 
 # unwrap VIEW/CAST/etc to find the actual data source (kernel output, buffer, or multi-device op)
 def _unwrap_src(s: UOp) -> UOp:
-  while len(s.src) and s.op not in {Ops.AFTER, Ops.BUFFER, Ops.PARAM, Ops.MSELECT, Ops.MSTACK, Ops.BIND}: s = s.src[0]
+  while len(s.src) and s.op not in {Ops.AFTER, Ops.BUFFER, Ops.PARAM, Ops.MSELECT, Ops.MSTACK}: s = s.src[0]
   return s
 
-# a buffer state is AFTER | BUFFER | PARAM. MSELECT/MSTACK join per-device states, BIND is not a buffer dependency
+# a buffer state is AFTER | BUFFER | PARAM. MSELECT/MSTACK join per-device states
 def _states(s: UOp) -> list[UOp]:
   s = _unwrap_src(s)
   if s.op in {Ops.MSELECT, Ops.MSTACK}: return [st for ss in s.src for st in _states(ss)]
-  if s.op is Ops.BIND: return []
   assert s.op in {Ops.AFTER, Ops.BUFFER, Ops.PARAM}, f"input to kernel must resolve to a buffer state, not {s.op}"
   return [s]
 
@@ -71,7 +70,7 @@ def create_schedule(sched_sink:UOp) -> UOp:
       else:
         k = rk.src[0] if rk.op is Ops.END else rk
         assert k.op is Ops.CALL, f"unexpected op in queue: {k.op}"
-        buf_uops = tuple(_unwrap_src(s).buf_uop for s in k.src[1:] if s.op is not Ops.BIND)
+        buf_uops = tuple(_unwrap_src(s).buf_uop for s in k.src[1:] if not s.is_bound_var)
         linearized.append(k.src[0].call(*buf_uops))
       for x in children.get(rk, []):
         in_degree[x] -= 1
@@ -100,7 +99,8 @@ pm_post_sched_cache = PatternMatcher([
 
 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}
+  # map the call body params back to the original Variables stored in the call args
+  binds = {f"p{i}":x.src[0].replace(op=Ops.PARAM) for i,x in enumerate(linear_call.src[1:]) if x.is_bound_var}
   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([
@@ -184,13 +184,13 @@ def create_linear_with_vars(big_sink:UOp) -> tuple[UOp, dict[str, int]]:
 
   # vars used in the schedule
   used_vars = set().union(*[{v.expr for v in si.src[0].variables()} for si in linear.src])
-  # get var_vals
+  # get var_vals from the bound Variables in the call args
   var_vals: dict[str, int] = {}
   for b in big_sink.src[1:]:
-    if b.op is Ops.BIND:
-      nm = b.src[0].expr
+    if b.is_bound_var:
+      v, val = b.unbind()
+      nm = v.expr
       if nm not in used_vars: continue
-      val = b.src[1].val
       if var_vals.get(nm, val) != val: raise RuntimeError(f"bind mismatch on {nm}, {var_vals[nm]} != {val}")
       var_vals[nm] = val
 

tinygrad/schedule/indexing.py

@@ -23,8 +23,8 @@ class IndexingContext:
     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,
+ALWAYS_CONTIGUOUS: set[Ops] = {Ops.CONTIGUOUS, Ops.AFTER, Ops.BUFFER,
+                      Ops.CONST, Ops.MSELECT, Ops.MSTACK, Ops.PARAM,
                       Ops.LOAD, Ops.CALL, Ops.FUNCTION}
 
 def realize(ctx:IndexingContext, tr:UOp) -> None: ctx.realize_map[tr] = None
@@ -34,12 +34,8 @@ def realize_srcs(ctx:IndexingContext, rb:UOp) -> None:
     if s.base.op not in ALWAYS_CONTIGUOUS: ctx.realize_map[s] = None
 
 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.realize_map \
-     and not dest.op_in_backward_slice_with_self(Ops.SHRINK, Ops.PERMUTE, Ops.FLIP, Ops.PAD):
-    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.realize_map[src] = None
+  if dest.base in src.toposort(enter_calls=False): ctx.realize_map[src] = None
 
 def realize_custom_kernel_srcs(ctx:IndexingContext, c:UOp) -> None:
   for s in c.src[1:]:
@@ -73,8 +69,10 @@ def broadcast_rngs(x:UOp, src:UOp, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]:
 
 # TODO: srcs contain (real data srcs, something else, ranges) and the boundary is confusing. see range_start
 def data_srcs(op:Ops, src:tuple[UOp, ...]) -> tuple[UOp, ...]:
-  if op in {Ops.PARAM, Ops.BUFFER, Ops.RANGE, Ops.SPECIAL, Ops.BIND}: return ()
-  if op in GroupOp.Movement|{Ops.INDEX, Ops.SLICE, Ops.STAGE, Ops.REDUCE, Ops.AFTER, Ops.END}: return src[:1]
+  if op in {Ops.PARAM, Ops.BUFFER, Ops.RANGE, Ops.SPECIAL}: return ()
+  # the store of a bound Variable only carries the input value, it has no data srcs
+  if op is Ops.STORE and src[0].is_variable: return ()
+  if op in GroupOp.Movement|{Ops.INDEX, Ops.STAGE, Ops.REDUCE, Ops.AFTER, Ops.END}: return src[:1]
   return src
 
 def create_bufferize_and_index_srcs(ctx:IndexingContext, x:UOp) -> list[UOp]:
@@ -84,7 +82,7 @@ def create_bufferize_and_index_srcs(ctx:IndexingContext, x:UOp) -> list[UOp]:
   for i, s in enumerate(x.src):
     new_src = s
     src_rngs = broadcast_rngs(x, s, ctx.range_map[x][0]) if x in ctx.range_map else ()
-    if s.op in {Ops.PARAM, Ops.BUFFER, Ops.SLICE, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
+    if s.op in {Ops.PARAM, Ops.BUFFER, Ops.MSTACK, Ops.MSELECT, Ops.AFTER}:
       if x in ctx.range_map and i < data_src_count: new_src = new_src.index(*src_rngs)
     elif s in ctx.realize_map:
       realized_ranges = ctx.realize_map[s]

tinygrad/schedule/memory.py

@@ -52,11 +52,9 @@ def memory_plan_rewrite(linear:UOp, held_bufs:set[UOp]|None=None) -> UOp:
     peaks[_key(buf)] = (max(peaks[_key(buf)][0], offsets[buf] + buf.max_numel() * buf.dtype.itemsize), peaks[_key(buf)][1])
   arena_sizes = {key: round_up(peak, block_size) for key, (peak, _) in peaks.items()}
 
-  # build replace_map: each buffer becomes a SLICE into a shared per-device-lane arena
+  # build replace_map: each buffer becomes a SHRINK/BITCAST into a shared per-device-lane arena
   arenas = {key: UOp.new_buffer(key[0], sz, dtypes.int8) for key, sz in arena_sizes.items()}
-  replace_map:dict[UOp, UOp] = {}
-  for buf_uop, offset in offsets.items():
-    replace_map[buf_uop] = UOp(Ops.SLICE, buf_uop.dtype, (arenas[_key(buf_uop)], UOp.const(offset)), buf_uop.max_numel())
+  replace_map = {buf_uop:arenas[_key(buf_uop)][offset:offset+buf_uop.nbytes()].bitcast(buf_uop.dtype) for buf_uop, offset in offsets.items()}
 
   if DEBUG >= 1 and (omem:=sum(nbytes.values()) / 1e6) != (nmem:=sum(arena_sizes.values()) / 1e6):
     print(f"memory reduced from {omem:.2f} MB -> {nmem:.2f} MB, {len(first_appearance)} -> {len(arenas)} bufs")

tinygrad/schedule/rangeify.py

@@ -4,7 +4,7 @@ import itertools
 from tinygrad.dtype import dtypes, AddrSpace, Invalid, to_dtype, strong_dtype
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg, shape_to_shape_arg
 from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, rewrite_group, identity_element
-from tinygrad.uop.symbolic import symbolic, pm_fold_cast_const
+from tinygrad.uop.symbolic import symbolic
 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, SPEC
 from tinygrad.helpers import PCONTIG, FLOAT16, OPENPILOT_HACKS, argsort, partition, get_single_element
@@ -60,7 +60,7 @@ pm_mops = PatternMatcher([
 # 0. do some cleanup rewrites, mostly copied from the old stuff
 
 def fix_store_hazard(target:UOp, src:UOp):
-  if (base:=target.base) not in src.backward_slice_with_self: return None
+  if (base:=target.base) not in src.toposort(enter_calls=False): return None
   # PERMUTE and FLIP reorder indices, SHRINK can have overlapping regions when dest is also shrunk
   unsafe = {Ops.PERMUTE, Ops.FLIP} | ({Ops.SHRINK} if target.op_in_backward_slice_with_self(Ops.SHRINK) else set())
   reaches_base: dict[UOp, bool] = {}
@@ -461,11 +461,12 @@ pm_add_buffers = pm_mops+pm_flatten_bufferize+PatternMatcher([
 class LocalAddBufferContext:
   dg:int = 0
   map:dict = field(default_factory=dict)
-  vars:dict = field(default_factory=dict)
   range:int = 0
   opts:tuple|None = None
 
 def debuf(ctx:LocalAddBufferContext, buf:UOp):
+  # Variables (ALU buffers with a value range) are scalar symbolic values, not real buffers: they become ALU params with no slot
+  if buf.is_variable: return buf.replace(op=Ops.PARAM)
   param = UOp(Ops.PARAM, src=(UOp.const(prod(buf.max_shape)),),
               arg=ParamArg(ctx.dg, buf.dtype, addrspace=buf.addrspace, device=buf.device))
   ret = param.reshape(buf.max_shape)
@@ -475,10 +476,6 @@ def debuf(ctx:LocalAddBufferContext, buf:UOp):
   ctx.dg += 1
   return ret
 
-def unbind_kernel(ctx:LocalAddBufferContext, b:UOp):
-  ctx.vars[b] = None
-  return b.src[0]
-
 def handle_after(ctx:LocalAddBufferContext, after:UOp):
   if after.addrspace == AddrSpace.LOCAL: return None
   buf = after.buf_uop
@@ -502,8 +499,7 @@ to_define_global = PatternMatcher([
   (UPat(Ops.STORE, name="x"), find_bufs),
   (UPat((Ops.BUFFER, Ops.MSTACK, Ops.MSELECT), name="buf"), debuf),
   (UPat(Ops.PARAM, name="v"), lambda v:
-   UOp.variable(v.arg.name, v.arg.vmin_vmax[0], v.arg.vmin_vmax[1], v.dtype, multiple_of=v.arg.multiple_of)
-   if v.arg.name is not None and v.arg.vmin_vmax is not None else None),
+   v.replace(arg=replace(v.arg, slot=-1)) if v.arg.name is not None and v.arg.vmin_vmax is not None and v.arg.slot != -1 else None),
 
   # this renumbers the params
   (UPat(Ops.PARAM, name="buf"), lambda ctx, buf:
@@ -512,7 +508,8 @@ to_define_global = PatternMatcher([
   # ALU params are scalar symbolic values, not buffers.
   (UPat(Ops.INDEX, src=(UPat(Ops.PARAM, name="v"),)), lambda v: v if v.addrspace == AddrSpace.ALU else None),
 
-  (UPat(Ops.BIND, name="b"), unbind_kernel),
+  # bound Variables are stores into Variable buffers: strip the store, the buffer becomes an ALU param via debuf
+  (UPat(Ops.AFTER, name="b"), lambda b: b.src[0] if b.is_bound_var else None),
   (UPat(Ops.AFTER, name="after"), handle_after),
 
   # remove device from local BUFFERIZE
@@ -541,13 +538,16 @@ pm_add_param_range_tags = PatternMatcher([
 def split_store(x:UOp) -> UOp|None:
   # if we have any open ranges here, we don't split. open DEVICE ranges are fine, they are bound per device at launch
   if any(r.arg[-1] is not AxisType.DEVICE for r in x.ranges): return None
+  # the store of a bound Variable is an input value, not a kernel
+  st = x.src[0] if x.op is Ops.END else x
+  if st.op is Ops.STORE and st.src[0].is_variable: return None
 
   # local kernel rewrite
   lctx = LocalAddBufferContext()
   ret = graph_rewrite(x, to_define_global+pm_flatten_range+rangeify_codegen, ctx=lctx, name="kernel split", bottom_up=True)
 
   # create the Kernel. NOTE: buffers can be on different devices here now, they are compiled to SDMA copies later by schedule
-  return ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts)).call(*lctx.map.values(), *lctx.vars.keys())
+  return ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts)).call(*lctx.map.values())
 
 split_kernels = PatternMatcher([
   (UPat((Ops.STORE, Ops.END), name="x"), split_store),
@@ -584,7 +584,7 @@ def get_kernel_graph(sink:UOp) -> UOp:
   tsink, rctx = run_rangeify(tsink, bool(DEBUG_RANGEIFY))
 
   tsink = graph_rewrite(tsink,
-                        symbolic+pm_fold_cast_const+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize+pm_no_indexing_calls,
+                        symbolic+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")