deepsek work

* improvements to the llm server

* eval script

* eval llm

* better eval gets 58.71

* cleanups

* add temperature, but multinomial is absurdly slow

* claude is so smart

* lint

* remove slop

* no more stop

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

@@ -287,6 +287,7 @@ runs:
 
         CMAKE_ARGS="-Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5"
         if [[ "${{ runner.os }}" == "macOS" ]]; then
+          sudo xcode-select -s /Applications/Xcode_16.2.app/Contents/Developer
           CMAKE_ARGS="$CMAKE_ARGS -DBoost_INCLUDE_DIR=$(brew --prefix boost)/include -DBoost_LIBRARY_DIR=$(brew --prefix boost)/lib"
         fi
 

.github/workflows/autogen.yml

@@ -56,6 +56,7 @@ jobs:
         python3 -c "from tinygrad.runtime.autogen import libusb"
         python3 -c "from tinygrad.runtime.autogen import mesa"
         python3 -c "from tinygrad.runtime.autogen import avcodec"
+        python3 -c "from tinygrad.runtime.autogen import llvm_qcom"
         REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
     - name: Check for differences
       run: |

.github/workflows/benchmark.yml

@@ -49,8 +49,6 @@ jobs:
         source /tmp/tinygrad_pytest_ci/bin/activate
         pytest -nauto --durations=20
     - name: openpilot compile3 0.10.1 driving_vision
-      run: FLOAT16=1 CL=1 IMAGE=2 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: IMAGE=1 openpilot compile3 0.10.1 driving_vision
       run: FLOAT16=1 CL=1 IMAGE=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
 
   testmacbenchmark:
@@ -187,13 +185,13 @@ jobs:
         PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
         PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
     - name: UsbGPU boot time
-      run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
+      run: sudo -E PYTHONPATH=. GMMU=0 DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
     - name: UsbGPU tiny tests
-      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
+      run: sudo -E PYTHONPATH=. GMMU=0 AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
     - name: UsbGPU copy speeds
-      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+      run: sudo -E PYTHONPATH=. GMMU=0 AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
     #- name: UsbGPU openpilot test
-    #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+    #  run: sudo -E PYTHONPATH=. GMMU=0 AMD=1 AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
     - name: UsbGPU (USB4/TB) boot time
       run: PYTHONPATH=. DEBUG=3 NV=1 NV_IFACE=PCI NV_NAK=1 time python3.11 test/test_tiny.py TestTiny.test_plus
     - name: UsbGPU (USB4/TB) tiny tests
@@ -590,22 +588,22 @@ jobs:
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
     - name: reset process replay
       run: test/external/process_replay/reset.py
-    - name: openpilot compile3 0.10.0 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
-    - name: openpilot compile3 0.10.0 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: openpilot compile3 0.11.0 driving_vision
+      run: BENCHMARK_LOG=openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
+    - name: IR3 openpilot compile3 0.11.0 driving_vision
+      run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM QCOM_IR3=1 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot compile3 0.11.0 driving_policy
+      run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
+    - name: openpilot compile3 0.11.0 dmonitoring
+      run: BENCHMARK_LOG=openpilot_0_11_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
     - name: DEBUG=2 openpilot compile3 0.10.1 driving_vision
-      run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: DEBUG=2 IMAGE=1 openpilot compile3 0.10.1 driving_vision
       run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: IMAGE=1 openpilot compile3 0.10.1 driving_vision
-      run: BENCHMARK_LOG=image_1_openpilot_0_10_1_vision PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
     - name: openpilot compile3 0.10.1 driving_vision
-      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
     - name: openpilot compile3 0.10.1 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
     - name: openpilot compile3 0.10.1 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
     - name: benchmark MobileNetV2 on DSP
       run: |
         # generate quantized weights
@@ -634,9 +632,9 @@ jobs:
         echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
     - name: openpilot compile3 0.10.1 driving_vision
-      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." DEV=AMD AMD_LLVM=1 AMD_IFACE=USB ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." GMMU=0 DEV=AMD AMD_LLVM=1 AMD_IFACE=USB ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
     - name: openpilot load_pickle 0.10.1 driving_vision
-      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." DEV=AMD AMD_IFACE=USB ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." GMMU=0 DEV=AMD AMD_IFACE=USB ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
 
   testreddriverbenchmark:
     name: AM Benchmark
@@ -699,6 +697,14 @@ jobs:
     - name: Run 10 MLPerf Bert training steps (1 gpu)
       # TODO: remove BERT_LAYERS once scheduler is fast
       run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+    - name: Remote
+      run: |
+        pkill -f 'extra/remote/serve.py' || true
+        PYTHONPATH=. python3 extra/remote/serve.py 6482 &
+        sleep 1
+        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 AMD=1 AMD_IFACE=PCI python3 test/test_tiny.py
+        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 AMD=1 AMD_AQL=1 AMD_IFACE=PCI python3 test/test_tiny.py
+        pkill -f 'extra/remote/serve.py' || true
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 
@@ -755,5 +761,12 @@ jobs:
     - name: Run 10 MLPerf Bert training steps (1 gpu)
       # TODO: remove BERT_LAYERS once scheduler is fast
       run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+    - name: Remote
+      run: |
+        pkill -f 'extra/remote/serve.py' || true
+        PYTHONPATH=. python3 extra/remote/serve.py 6483 &
+        sleep 1
+        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6483 NV=1 python3 test/test_tiny.py
+        pkill -f 'extra/remote/serve.py' || true
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py

.github/workflows/test.yml

@@ -70,7 +70,7 @@ jobs:
         source venv/bin/activate
         pip install $GITHUB_WORKSPACE
         cp $GITHUB_WORKSPACE/examples/beautiful_mnist.py .
-        BS=2 STEPS=10 python beautiful_mnist.py
+        BS=2 STEPS=10 MAX_BUFFER_SIZE=0 python beautiful_mnist.py
     - name: Test Docs Build
       run: python -m mkdocs build --strict
     - name: Test Docs
@@ -141,7 +141,7 @@ jobs:
         sudo apt update || true
         sudo apt install -y --no-install-recommends ninja-build
     - name: Test beautiful_mnist in torch with TINY_BACKEND
-      run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+      run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.0 MAX_BUFFER_SIZE=0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
     - name: Test some torch tests (expect failure)
       run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
 
@@ -167,10 +167,10 @@ jobs:
       run: PYTHON=1 python3 test/backend/test_symbolic_ops.py
     - name: test_renderer_failures with Python emulator
       run: PYTHON=1 python3 -m pytest -rA test/backend/test_renderer_failures.py::TestRendererFailures
-    - name: Test IMAGE=2 support
+    - name: Test IMAGE support
       run: |
-        IMAGE=2 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
-        IMAGE=2 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_simple_conv2d
+        IMAGE=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
+        IMAGE=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_simple_conv2d
     - name: Test emulated METAL tensor cores
       run: |
         DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_big_gemm
@@ -369,11 +369,11 @@ jobs:
           key: gpu-image
           deps: testing_unit
           opencl: 'true'
-      - name: Test CL IMAGE=2 ops
+      - name: Test CL IMAGE=1 ops
         run: |
-          CL=1 IMAGE=2 python -m pytest -n=auto test/backend/test_ops.py --durations=20
+          CL=1 IMAGE=1 python -m pytest -n=auto test/backend/test_ops.py --durations=20
         # TODO: training is broken
-        # CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
+        # CL=1 IMAGE=1 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -418,13 +418,13 @@ jobs:
           llvm: 'true'
       - name: Test openpilot model kernel count and gate usage
         run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1397 ALLOWED_GATED_READ_IMAGE=94 FLOAT16=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1486 ALLOWED_GATED_READ_IMAGE=17 FLOAT16=1 CL=1 IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
       - name: Test openpilot CL compile fp16
-        run: FLOAT16=1 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+        run: FLOAT16=1 DEBUGCL=1 CL=1 IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
       - name: Test openpilot CL compile fp32 (test correctness)
-        run: DEBUGCL=1 CL=1 IMAGE=2 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
+        run: CL=1 IMAGE=1 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
       - name: Test openpilot LLVM compile fp16
-        run: FLOAT16=1 CPU=1 CPU_LLVM=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+        run: IMAGE=1 FLOAT16=1 CPU=1 CPU_LLVM=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -505,8 +505,13 @@ jobs:
         uses: ./.github/actions/setup-tinygrad
         with:
           key: apps_llm
-      - name: Test 1B LLM
-        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm | grep -i rooster
+      - name: Test 1B LLM (llama)
+        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model llama3.2:1b | tee /dev/stderr | grep -i rooster
+      - name: Test 1B LLM (llama q4)
+        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster
+      - name: Test 1B LLM (qwen)
+        # NOTE: qwen is dumb and only knows about female chickens
+        run: echo "What's a female chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model qwen3:0.6b | tee /dev/stderr | grep -i hen
 
 # ****** Models Tests ******
 
@@ -699,7 +704,7 @@ jobs:
       - name: Run test_tiny on MOCKAM
         run: python test/test_tiny.py
       - name: Run test_tiny on MOCKAM USB
-        run: AMD_IFACE=USB python test/test_tiny.py
+        run: GMMU=0 AMD_IFACE=USB python test/test_tiny.py
       - name: Run test_hcq on MOCKAM
         run: python -m pytest test/device/test_hcq.py
 
@@ -708,8 +713,7 @@ jobs:
       fail-fast: false
       matrix:
         backend: [amd, amdllvm]
-        arch: [rdna3, rdna4]
-        #arch: [rdna3, rdna4, cdna4]
+        arch: [rdna3, rdna4, cdna4]
 
     name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
     runs-on: ubuntu-22.04
@@ -735,7 +739,7 @@ jobs:
           python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['AMD'], Device.DEFAULT"
           DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
       - name: Run 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/testextra/test_cfg_viz.py test/external/external_test_am.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 TRANSCENDENTAL math
         run: TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
       - name: Run process replay tests

examples/beautiful_mnist.py

@@ -1,6 +1,6 @@
 # model based off https://medium.com/data-science/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
 from typing import Callable
-from tinygrad import Tensor, TinyJit, nn, GlobalCounters
+from tinygrad import Tensor, TinyJit, nn, GlobalCounters, function
 from tinygrad.helpers import getenv, colored, trange
 from tinygrad.nn.datasets import mnist
 
@@ -15,30 +15,31 @@ class Model:
       nn.BatchNorm(64), Tensor.max_pool2d,
       lambda x: x.flatten(1), nn.Linear(576, 10)]
 
+  @function
   def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers)
 
+  @TinyJit
+  @Tensor.train()
+  def train_step(self, X_train:Tensor, Y_train:Tensor) -> Tensor:
+    opt.zero_grad()
+    samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
+    loss = self(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
+    return loss.realize(*opt.schedule_step())
+
+  @TinyJit
+  def get_test_acc(self, X_test:Tensor, Y_test:Tensor) -> Tensor: return (self(X_test).argmax(axis=1) == Y_test).mean()*100
+
 if __name__ == "__main__":
   X_train, Y_train, X_test, Y_test = mnist(fashion=getenv("FASHION"))
 
   model = Model()
   opt = (nn.optim.Muon if getenv("MUON") else nn.optim.SGD if getenv("SGD") else nn.optim.Adam)(nn.state.get_parameters(model))
 
-  @TinyJit
-  @Tensor.train()
-  def train_step() -> Tensor:
-    opt.zero_grad()
-    samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
-    loss = model(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
-    return loss.realize(*opt.schedule_step())
-
-  @TinyJit
-  def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100
-
   test_acc = float('nan')
   for i in (t:=trange(getenv("STEPS", 70))):
     GlobalCounters.reset()   # NOTE: this makes it nice for DEBUG=2 timing
-    loss = train_step()
-    if i%10 == 9: test_acc = get_test_acc().item()
+    loss = model.train_step(X_train, Y_train)
+    if i%10 == 9: test_acc = model.get_test_acc(X_test, Y_test).item()
     t.set_description(f"loss: {loss.item():6.2f} test_accuracy: {test_acc:5.2f}%")
 
   # verify eval acc

examples/benchmark_onnx.py

@@ -5,7 +5,7 @@ from extra.onnx_helpers import get_example_inputs, validate
 
 def load_onnx_model(onnx_file):
   run_onnx = OnnxRunner(onnx_file)
-  run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(None) for k,v in kwargs.items()}).values())), prune=True, optimize=True)
+  run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(None) for k,v in kwargs.items()}).values())), prune=True)
   return run_onnx_jit, run_onnx.graph_inputs
 
 if __name__ == "__main__":

examples/mlperf/model_train.py

@@ -13,6 +13,8 @@ from extra.bench_log import BenchEvent, WallTimeEvent
 # TODO: fix benchmark logging and use tinygrad tqdm
 from tqdm import tqdm
 
+from tinygrad.uop.ops import UOp
+
 def train_resnet():
   from extra.models import resnet
   from examples.mlperf.dataloader import batch_load_resnet
@@ -1282,7 +1284,7 @@ def train_bert():
         previous_step = i
 
 def train_llama3():
-  from examples.mlperf.models.llama import Transformer
+  from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad
   from examples.llama3 import MODEL_PARAMS
   from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
   from examples.mlperf.optim import GradAccClipAdamW
@@ -1343,7 +1345,7 @@ def train_llama3():
   if (MP := getenv("MP", 1)) > 1: model_params['vocab_size'] = round_up(model_params['vocab_size'], 256 * MP)
   vocab_mask:Tensor = Tensor.arange(model_params['vocab_size']).reshape(1, 1, -1) >= real_vocab_size
 
-  model = Transformer(**model_params, max_context=SEQLEN)
+  model = FlatTransformer(**model_params, max_context=SEQLEN)
 
   params = get_parameters(model)
   # weights are all bfloat16 for now
@@ -1353,49 +1355,25 @@ def train_llama3():
     for v in get_parameters(model):
       v = v.assign(Tensor.empty(v.shape))
 
-  if (DP := getenv("DP", 1)) > 1:
-    device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-    for v in get_parameters(model):
-      v.shard_(device, axis=None)
+  is_dp = (DP := getenv("DP", 1)) > 1
+  is_mp = (MP := getenv("MP", 1)) > 1
+  is_sharding = is_dp or is_mp
+  device_count = max(DP, MP)
+  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))
 
-    vocab_mask.shard_(device, axis=None)
+  model.shard(device, is_mp)
 
-  if (MP := getenv("MP", 1)) > 1:
-    device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
-    for k,v in get_state_dict(model).items():
-      if 'scale' in k: v.shard_(device, axis=None)  # from quantized
-      elif '.attention.wq' in k: v.shard_(device, axis=0)
-      elif '.attention.wk' in k: v.shard_(device, axis=0)
-      elif '.attention.wv' in k: v.shard_(device, axis=0)
-      elif '.attention.wqkv' in k: v.shard_(device, axis=0)
-      elif '.attention.wo' in k: v.shard_(device, axis=1)
-      elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
-      elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
-      elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
-      elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
-      elif 'output.weight' in k: v.shard_(device, axis=0)
-      else:
-        # attention_norm, ffn_norm, norm
-        v.shard_(device, axis=None)
-      # prevents memory spike on device 0
-      v.realize()
-
-    vocab_mask.shard_(device, axis=2).realize()
+  if is_dp: vocab_mask.shard_(device, axis=None).realize()
+  if is_mp: vocab_mask.shard_(device, axis=2).realize()
 
   is_offload_optim = bool(getenv("OFFLOAD_OPTIM"))
   is_fake_offload = Device.DEFAULT == "NULL"
   optim_device = ("CPU" if not is_fake_offload else "NULL:99") if is_offload_optim else None
-  optim = GradAccClipAdamW(get_parameters(model), lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2,
+  optim = GradAccClipAdamW(params, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2,
                            eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)
 
   # init grads
-  if is_offload_optim:
-    for p in optim.params:
-      p.grad = Tensor.zeros(p.shape, dtype=p.dtype, device=optim_device, requires_grad=False).contiguous().realize()
-  else:
-    for p in optim.params:
-      p.grad = p.zeros_like().contiguous().realize()
-  grads: list[Tensor] = [p.grad for p in optim.params]
+  grads = [Tensor.zeros_like(p).contiguous() for p in optim.params]
 
   scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
 
@@ -1410,28 +1388,24 @@ def train_llama3():
 
   @TinyJit
   def minibatch(tokens:Tensor):
-    if (DP := getenv("DP", 1)) > 1:
-      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-      tokens = tokens.to(None).shard(device, 0)
-    if (MP := getenv("MP", 1)) > 1:
-      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
-      tokens = tokens.shard(device)
-    if DP == 1 and MP == 1: tokens = tokens.to(None)
+    if is_dp: tokens = tokens.to(None).shard(device, 0)
+    if is_mp: tokens = tokens.shard(device)
+    if not is_sharding: tokens = tokens.to(None)
     logits:Tensor = model(tokens[:, :-1])
     loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
-    loss.backward()
-    assert all(p.grad is g for p,g in zip(optim.params, grads))
+
+    for i,(t,g) in enumerate(zip(optim.params, loss.gradient(*optim.params))):
+      grads[i].replace(Tensor(grads[i].uop.after(UOp.group(*apply_grad(grads[i].uop, g.uop))), device=t.device))
+
     loss_cpu = loss.flatten().float().to("CPU")
-    Tensor.realize(loss_cpu, *grads)
-    return loss_cpu
+    return loss_cpu.realize(*grads)
 
   @TinyJit
   def optim_step():
     grad_norm = optim.fstep(grads)
     scheduler.step()
 
-    for g in grads:
-      g.assign(g.zeros_like())
+    for g in grads: g.assign(g.zeros_like())
 
     lr_cpu = optim.lr.float().to("CPU")
     grad_norm_cpu = grad_norm.float().to("CPU")
@@ -1442,13 +1416,9 @@ def train_llama3():
   @TinyJit
   @Tensor.train(False)
   def eval_step(tokens:Tensor):
-    if (DP := getenv("DP", 1)) > 1:
-      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
-      tokens = tokens.to(None).shard(device, 0)
-    if (MP := getenv("MP", 1)) > 1:
-      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
-      tokens = tokens.shard(device)
-    if DP == 1 and MP == 1: tokens = tokens.to(None)
+    if is_dp: tokens = tokens.to(None).shard(device, 0)
+    if is_mp: tokens = tokens.shard(device)
+    if not is_sharding: tokens = tokens.to(None)
     logits:Tensor = model(tokens[:, :-1])
     loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
     return loss.flatten().float().to("CPU")

examples/mlperf/models/flat_llama.py

@@ -0,0 +1,193 @@
+import math, os
+if __name__ == "__main__":
+  os.environ["DEFAULT_FLOAT"] = "bfloat16"
+  os.environ["OPTIM_DTYPE"] = "bfloat16"
+  if "DEV" not in os.environ: os.environ["DEV"] = "NULL"
+  # CDNA
+  os.environ["EMULATE"] = "AMD_CDNA4"
+  os.environ["DEVICE_IN_FUNCTION_BUG"] = "1"
+  os.environ["ALL2ALL"] = "1"
+  os.environ["USE_ATOMICS"] = "1"
+  if "HK_FLASH_ATTENTION" not in os.environ:
+    os.environ["HK_FLASH_ATTENTION"] = "1"
+    if "ASM_GEMM" not in os.environ:
+      os.environ["ASM_GEMM"] = "1"
+from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
+from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
+from tinygrad.uop.ops import Ops, UOp
+from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
+
+FP8 = getenv("FP8", 0)
+
+FP8_DTYPE = dtypes.fp8e4m3
+FP8_MAX = 448.0
+
+def quantize_fp8(x:Tensor):
+  scale = FP8_MAX / (x.abs().max().detach() + 1e-8)
+  x_scaled = x * scale
+  x_clamped = x_scaled + (x_scaled.detach().clamp(-FP8_MAX, FP8_MAX) - x_scaled.detach())  # STE
+  return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal()
+
+def matmul(x:Tensor, w:Tensor) -> Tensor:
+  if not FP8: return x @ w.T
+  # weights are already FP8, just quantize activations
+  x_fp8, x_scale = quantize_fp8(x)
+  return x_fp8.dot(w.T, dtype=dtypes.float) * x_scale
+
+def rmsnorm(x_in:Tensor, eps:float):
+  x = x_in.float()
+  x = x * (x.square().mean(-1, keepdim=True) + eps).rsqrt()
+  return x.cast(x_in.dtype)
+
+class FlatTransformer:
+  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
+               rope_theta:int=10000, max_context:int=1024):
+    self.vocab_size = vocab_size
+    self.n_layers = n_layers
+    self.n_heads = n_heads
+    self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads # n_kv_heads != n_heads implies MQA [arxiv/2307.09288, A.2.1]
+    self.head_dim = dim // n_heads
+    self.n_rep = self.n_heads // self.n_kv_heads
+
+    # Attention
+    self.wqkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
+    self.wo = self.lin_per_layer(self.n_heads * self.head_dim, dim)
+
+    # FeedForward
+    self.w1 = self.lin_per_layer(dim, hidden_dim)
+    self.w2 = self.lin_per_layer(hidden_dim, dim)
+    self.w3 = self.lin_per_layer(dim, hidden_dim)
+
+    self.norm_eps = norm_eps
+    self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
+    self.ffn_norm = Tensor.ones(n_layers, dim).contiguous()
+
+    # output
+    self.norm = nn.RMSNorm(dim, norm_eps)
+    self.tok_embeddings = nn.Embedding(vocab_size, dim)
+    self.output = nn.Linear(dim, vocab_size, bias=False)
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)
+
+  def lin_per_layer(self, in_features:int, out_features:int):
+    bound = 1 / math.sqrt(in_features)
+    dt = FP8_DTYPE if FP8 else None
+    if getenv("ZEROS"): return Tensor.zeros(self.n_layers, out_features, in_features, dtype=dt)
+    return Tensor.uniform(self.n_layers, out_features, in_features, low=-bound, high=bound, dtype=dt)
+
+  def attention(self, x:Tensor, freqs_cis:Tensor, attention_norm:Tensor, wqkv:Tensor, wo:Tensor):
+    x = rmsnorm(x, self.norm_eps) * attention_norm
+    xqkv = matmul(x, wqkv)
+
+    bsz, seqlen, _ = xqkv.shape
+    # interleaved layout: each kv group has [n_rep q heads, 1 k head, 1 v head] for clean MP sharding
+    xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
+    xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
+    xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
+    xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
+
+    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
+    xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
+    attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
+    attn = attn.reshape(bsz, seqlen, -1)
+    return matmul(attn, wo)
+
+  def feed_forward(self, x:Tensor, ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor):
+    x = rmsnorm(x, self.norm_eps) * ffn_norm
+    x_w1 = matmul(x, w1).silu()
+    x_w3 = matmul(x.contiguous_backward(), w3)
+    return matmul(x_w1 * x_w3, w2)
+
+  @function(precompile=True, precompile_backward=True)
+  def run_layer(self, x:Tensor, freqs_cis:Tensor,
+                attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
+                ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor):
+    h = x + self.attention(x, freqs_cis, attention_norm, wqkv, wo)
+    return h + self.feed_forward(h, ffn_norm, w1, w2, w3)
+
+  def shard(self, device:tuple[str, ...], mp:bool=False):
+    from tinygrad.nn.state import get_parameters
+    if not mp:
+      for v in get_parameters(self): v.shard_(device, axis=None)
+    else:
+      # flat per-layer weights: axis 0 is n_layers, so shard axes are +1 vs per-layer Transformer
+      self.wqkv.shard_(device, axis=1).realize()          # (n_layers, out, dim) shard out
+      self.wo.shard_(device, axis=2).realize()             # (n_layers, dim, in) shard in
+      self.w1.shard_(device, axis=1).realize()             # (n_layers, hidden, dim) shard out
+      self.w2.shard_(device, axis=2).realize()             # (n_layers, dim, hidden) shard in
+      self.w3.shard_(device, axis=1).realize()             # (n_layers, hidden, dim) shard out
+      self.attention_norm.shard_(device, axis=None).realize()
+      self.ffn_norm.shard_(device, axis=None).realize()
+      self.norm.weight.shard_(device, axis=None).realize()
+      self.tok_embeddings.weight.shard_(device, axis=0).realize()
+      self.output.weight.shard_(device, axis=0).realize()
+      self.freqs_cis.shard_(device, axis=None).realize()
+
+  def __call__(self, tokens:Tensor):
+    h = self.tok_embeddings(tokens)
+    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
+    for i in range(self.n_layers):
+      h = self.run_layer(h, freqs_cis,
+                         self.attention_norm[i], self.wqkv[i], self.wo[i],
+                         self.ffn_norm[i], self.w1[i], self.w2[i], self.w3[i])
+    logits = self.output(self.norm(h))
+    return logits
+
+# TODO: this shouldn't be needed, but it prevents a copy of the grads. CAT can help
+def apply_grad(old_grad:UOp, new_grad:UOp) -> list[UOp]:
+  if new_grad.op == Ops.ADD:
+    return apply_grad(old_grad, new_grad.src[0])+apply_grad(old_grad, new_grad.src[1])
+  elif new_grad.op == Ops.PAD:
+    grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(new_grad.src[0].shape, new_grad.marg)])
+    return apply_grad(old_grad.shrink(grad_shrink), new_grad.src[0])
+  else:
+    return [old_grad.store(old_grad + new_grad)]
+
+if __name__ == "__main__":
+  config = {}
+  BS                 = config["BS"]                     = getenv("BS", 16)
+  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
+
+  from examples.llama3 import MODEL_PARAMS
+  model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
+  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
+  model = FlatTransformer(**model_params, max_context=SEQLEN)
+  state = nn.state.get_state_dict(model)
+  print("tensor count:", len(state))
+
+  # shard the model
+  from tinygrad import Device
+  if (DP := getenv("DP", 1)) > 1:
+    model.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)))
+  if (MP := getenv("MP", 1)) > 1:
+    model.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)), mp=True)
+
+  # preallocate all the grad buffers and zero them out
+  grads = {x:Tensor.zeros_like(x).contiguous() for x in state.values() if x.requires_grad is None}
+
+  # print model size
+  sz = 0
+  for k,v in state.items():
+    print(f"{colored(k, 'green' if v in grads else 'white'):30s} {str(v.shape):30s} {str(v.dtype):20s} {v.device}  {v.nbytes()/1e9:.2f} GB")
+    sz += v.nbytes()
+  print(f"total sz: {sz/1e9:.2f} GB")
+
+  with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=model.vocab_size, dtype=dtypes.int)
+  with Timing("realize weights/grads/data: "): Tensor.realize(*state.values(), *grads.values(), tokens)
+  print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
+  if DP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)), axis=0)
+  if MP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)))
+
+  @TinyJit
+  def jit_step(tokens:Tensor):
+    with Timing("python forward: "): loss = model(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
+    with Timing("python backward: "):
+      for t,g in zip(grads, loss.gradient(*grads)):
+        grads[t] = Tensor(grads[t].uop.after(UOp.group(*apply_grad(grads[t].uop, g.uop))), device=t.device)
+    with Timing("run step: "): loss.realize(*grads.values())
+
+  for i in range(6):
+    GlobalCounters.reset()
+    profile_marker(f"step {i}")
+    with Timing(colored(f"*** step {i}: ", "red")):
+      jit_step(tokens)
+  print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))

examples/mlperf/models/test_flat_llama.py

@@ -0,0 +1,140 @@
+import os
+os.environ["WQKV"] = "1"
+import unittest
+import numpy as np
+from tinygrad import Tensor, nn, dtypes
+from tinygrad.nn.state import get_parameters
+from tinygrad.device import is_dtype_supported
+from examples.mlperf.models.llama import Transformer
+from examples.mlperf.models.flat_llama import FlatTransformer
+
+def copy_weights(flat:FlatTransformer, ref:Transformer):
+  n_layers = flat.n_layers
+  Tensor.realize(*nn.state.get_state_dict(ref).values())
+  flat.wqkv.assign(Tensor(np.stack([ref.layers[i].attention.wqkv.weight.numpy() for i in range(n_layers)])))
+  flat.wo.assign(Tensor(np.stack([ref.layers[i].attention.wo.weight.numpy() for i in range(n_layers)])))
+  flat.w1.assign(Tensor(np.stack([ref.layers[i].feed_forward.w1.weight.numpy() for i in range(n_layers)])))
+  flat.w2.assign(Tensor(np.stack([ref.layers[i].feed_forward.w2.weight.numpy() for i in range(n_layers)])))
+  flat.w3.assign(Tensor(np.stack([ref.layers[i].feed_forward.w3.weight.numpy() for i in range(n_layers)])))
+  flat.attention_norm.assign(Tensor(np.stack([ref.layers[i].attention_norm.weight.numpy() for i in range(n_layers)])))
+  flat.ffn_norm.assign(Tensor(np.stack([ref.layers[i].ffn_norm.weight.numpy() for i in range(n_layers)])))
+  flat.norm.weight.assign(Tensor(ref.norm.weight.numpy()))
+  flat.tok_embeddings.weight.assign(Tensor(ref.tok_embeddings.weight.numpy()))
+  flat.output.weight.assign(Tensor(ref.output.weight.numpy()))
+
+class TestFlatLlama(unittest.TestCase):
+  def test_forward_match(self):
+    Tensor.manual_seed(42)
+    params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
+    ref = Transformer(**params)
+    flat = FlatTransformer(**params)
+    copy_weights(flat, ref)
+    Tensor.realize(*nn.state.get_state_dict(flat).values())
+
+    tokens = Tensor([[1, 50, 100, 999, 2]])
+    ref_logits = ref(tokens).realize()
+    flat_logits = flat(tokens).realize()
+    self.assertEqual(ref_logits.shape, flat_logits.shape)
+    diff = (ref_logits - flat_logits).abs().max().item()
+    self.assertLess(diff, 1e-5, f"forward mismatch: max abs diff {diff}")
+
+  def test_backward_match(self):
+    Tensor.manual_seed(42)
+    params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
+    ref = Transformer(**params)
+    flat = FlatTransformer(**params)
+    copy_weights(flat, ref)
+
+    for p in get_parameters(ref): p.requires_grad_(True)
+    for p in get_parameters(flat): p.requires_grad_(True)
+    Tensor.realize(*nn.state.get_state_dict(flat).values())
+
+    tokens = Tensor([[1, 50, 100, 999, 2, 10]])
+
+    ref_loss = ref(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
+    ref_loss.backward()
+    ref_grads = {k: v.grad.numpy() for k, v in nn.state.get_state_dict(ref).items() if v.grad is not None}
+
+    flat_loss = flat(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
+    flat_loss.backward()
+    flat_grads = {k: v.grad.numpy() for k, v in nn.state.get_state_dict(flat).items() if v.grad is not None}
+
+    # check loss matches
+    self.assertAlmostEqual(ref_loss.item(), flat_loss.item(), places=4)
+
+    # check output weight grad matches
+    diff = abs(ref_grads["output.weight"] - flat_grads["output.weight"]).max()
+    self.assertLess(diff, 1e-4, f"output.weight grad mismatch: max abs diff {diff}")
+
+    # check per-layer weight grads match
+    for i in range(params["n_layers"]):
+      for flat_key, ref_key in [
+        ("wqkv", f"layers.{i}.attention.wqkv.weight"),
+        ("wo", f"layers.{i}.attention.wo.weight"),
+        ("w1", f"layers.{i}.feed_forward.w1.weight"),
+        ("w2", f"layers.{i}.feed_forward.w2.weight"),
+        ("w3", f"layers.{i}.feed_forward.w3.weight"),
+      ]:
+        diff = abs(ref_grads[ref_key] - flat_grads[flat_key][i]).max()
+        self.assertLess(diff, 1e-4, f"layer {i} {flat_key} grad mismatch: max abs diff {diff}")
+
+  @unittest.skipUnless(os.getenv("CPU", "") == "1", "multi-device CPU test")
+  def test_forward_match_mp(self):
+    Tensor.manual_seed(42)
+    params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
+    from tinygrad import Device
+    devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
+    ref = Transformer(**params)
+    flat = FlatTransformer(**params)
+    copy_weights(flat, ref)
+    Tensor.realize(*nn.state.get_state_dict(flat).values())
+    flat.shard(devices, mp=True)
+
+    tokens = Tensor([[1, 50, 100, 999, 2]], device=devices[0])
+    ref_logits = ref(tokens.to(devices[0])).numpy()
+    flat_logits = flat(tokens.shard(devices)).numpy()
+    self.assertEqual(ref_logits.shape, flat_logits.shape)
+    np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
+
+  @unittest.skipUnless(os.getenv("CPU", "") == "1", "multi-device CPU test")
+  def test_forward_match_dp(self):
+    Tensor.manual_seed(42)
+    params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
+    from tinygrad import Device
+    devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
+    ref = Transformer(**params)
+    flat = FlatTransformer(**params)
+    copy_weights(flat, ref)
+    Tensor.realize(*nn.state.get_state_dict(flat).values())
+    flat.shard(devices)
+
+    tokens = Tensor([[1, 50, 100, 999, 2], [2, 100, 50, 1, 999]], device=devices[0])
+    ref_logits = ref(tokens.to(devices[0])).numpy()
+    flat_logits = flat(tokens.shard(devices, axis=0)).numpy()
+    self.assertEqual(ref_logits.shape, flat_logits.shape)
+    np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
+
+  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), "fp8 not supported on this device")
+  def test_forward_fp8(self):
+    import examples.mlperf.models.flat_llama as flat_llama_mod
+    old_fp8 = flat_llama_mod.FP8
+    try:
+      flat_llama_mod.FP8 = 1
+      Tensor.manual_seed(42)
+      params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
+      ref = Transformer(**params)
+      flat = FlatTransformer(**params)
+      copy_weights(flat, ref)
+      Tensor.realize(*nn.state.get_state_dict(flat).values())
+
+      tokens = Tensor([[1, 50, 100, 999, 2]])
+      ref_logits = ref(tokens).numpy()
+      flat_logits = flat(tokens).numpy()
+      self.assertEqual(ref_logits.shape, flat_logits.shape)
+      # FP8 has lower precision, allow larger tolerance
+      np.testing.assert_allclose(flat_logits, ref_logits, atol=1.0, rtol=0.1)
+    finally:
+      flat_llama_mod.FP8 = old_fp8
+
+if __name__ == "__main__":
+  unittest.main()

examples/mlperf/optim.py

@@ -25,8 +25,10 @@ class GradAccClipAdamW(Optimizer):
     return extra[-1]
 
   def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
+    grads = list(grads)
+
     for i in range(len(grads)):
-      if grads[i].device != self.m[i].device: grads[i].assign(grads[i].to(self.m[i].device))
+      if grads[i].device != self.m[i].device: grads[i] = grads[i].to(self.m[i].device)
 
     if self.fused:
       grads[0].assign(grads[0] / self.grad_acc)
@@ -52,6 +54,6 @@ class GradAccClipAdamW(Optimizer):
     return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]
 
   def _apply_update(self, t:Tensor, up:Tensor) -> Tensor:
-    wd = self.wd if t.ndim >= 2 else 0.0
+    wd = self.wd if t.ndim >= 3 else 0.0
     up = up.shard_like(t) + self.lr.to(t.device) * wd * t.detach()
     return t.detach() - up.cast(t.dtype)

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

@@ -5,6 +5,7 @@ export DEV=${DEV:-AMD}
 export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
 
 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh

@@ -1,6 +1,5 @@
 #!/bin/bash
 export BENCHMARK=5
 export EVAL_BS=0
-export VIZ=${VIZ:--1}
-examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
-extra/viz/cli.py --profile --device "AMD" --top 20
+VIZ=${VIZ:--1} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
+extra/viz/cli.py --profile --device "AMD" --limit 20

examples/mnist_gan.py

@@ -1,12 +1,11 @@
 from pathlib import Path
-import numpy as np
 import torch
 from torchvision.utils import make_grid, save_image
 from tinygrad.nn.state import get_parameters
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import trange
 from tinygrad.nn import optim
-from extra.datasets import fetch_mnist
+from tinygrad.nn.datasets import mnist
 
 class LinearGen:
   def __init__(self):
@@ -38,14 +37,14 @@ class LinearDisc:
     return x
 
 def make_batch(images):
-  sample = np.random.randint(0, len(images), size=(batch_size))
-  image_b = images[sample].reshape(-1, 28*28).astype(np.float32) / 127.5 - 1.0
-  return Tensor(image_b)
+  sample = Tensor.randint(batch_size, low=0, high=images.shape[0])
+  return images[sample].reshape(batch_size, 28*28).cast('float').div(127.5).sub(1.0)
 
 def make_labels(bs, col, val=-2.0):
-  y = np.zeros((bs, 2), np.float32)
-  y[range(bs), [col] * bs] = val  # Can we do label smoothing? i.e -2.0 changed to -1.98789.
-  return Tensor(y)
+  y = Tensor.zeros(bs, 2)
+  if col == 0: y = y + Tensor([val, 0.0])
+  else: y = y + Tensor([0.0, val])
+  return y
 
 def train_discriminator(optimizer, data_real, data_fake):
   real_labels = make_labels(batch_size, 1)
@@ -71,12 +70,12 @@ def train_generator(optimizer, data_fake):
 
 if __name__ == "__main__":
   # data for training and validation
-  images_real = np.vstack(fetch_mnist()[::2])
+  X_train, _, _, _ = mnist()
   ds_noise = Tensor.randn(64, 128, requires_grad=False)
   # parameters
   epochs, batch_size, k = 300, 512, 1
   sample_interval = epochs // 10
-  n_steps = len(images_real) // batch_size
+  n_steps = X_train.shape[0] // batch_size
   # models and optimizer
   generator = LinearGen()
   discriminator = LinearDisc()
@@ -84,24 +83,24 @@ if __name__ == "__main__":
   output_dir = Path(".").resolve() / "outputs"
   output_dir.mkdir(exist_ok=True)
   # optimizers
-  optim_g = optim.Adam(get_parameters(generator),lr=0.0002, b1=0.5)  # 0.0002 for equilibrium!
-  optim_d = optim.Adam(get_parameters(discriminator),lr=0.0002, b1=0.5)
+  optim_g = optim.Adam(get_parameters(generator), lr=0.0002, b1=0.5)  # 0.0002 for equilibrium!
+  optim_d = optim.Adam(get_parameters(discriminator), lr=0.0002, b1=0.5)
   # training loop
-  Tensor.training = True
-  for epoch in (t := trange(epochs)):
-    loss_g, loss_d = 0.0, 0.0
-    for _ in range(n_steps):
-      data_real = make_batch(images_real)
-      for step in range(k):  # Try with k = 5 or 7.
+  with Tensor.train():
+    for epoch in (t := trange(epochs)):
+      loss_g, loss_d = 0.0, 0.0
+      for _ in range(n_steps):
+        data_real = make_batch(X_train)
+        for step in range(k):  # Try with k = 5 or 7.
+          noise = Tensor.randn(batch_size, 128)
+          data_fake = generator.forward(noise).detach()
+          loss_d += train_discriminator(optim_d, data_real, data_fake)
         noise = Tensor.randn(batch_size, 128)
-        data_fake = generator.forward(noise).detach()
-        loss_d += train_discriminator(optim_d, data_real, data_fake)
-      noise = Tensor.randn(batch_size, 128)
-      data_fake = generator.forward(noise)
-      loss_g += train_generator(optim_g, data_fake)
-    if (epoch + 1) % sample_interval == 0:
-      fake_images = generator.forward(ds_noise).detach().numpy()
-      fake_images = (fake_images.reshape(-1, 1, 28, 28) + 1) / 2  # 0 - 1 range.
-      save_image(make_grid(torch.tensor(fake_images)), output_dir / f"image_{epoch+1}.jpg")
-    t.set_description(f"Generator loss: {loss_g/n_steps}, Discriminator loss: {loss_d/n_steps}")
+        data_fake = generator.forward(noise)
+        loss_g += train_generator(optim_g, data_fake)
+      if (epoch + 1) % sample_interval == 0:
+        fake_images = generator.forward(ds_noise).detach().numpy()
+        fake_images = (fake_images.reshape(-1, 1, 28, 28) + 1) / 2  # 0 - 1 range.
+        save_image(make_grid(torch.tensor(fake_images)), output_dir / f"image_{epoch+1}.jpg")
+      t.set_description(f"Generator loss: {loss_g/n_steps}, Discriminator loss: {loss_d/n_steps}")
   print("Training Completed!")

extra/export_model.py

@@ -1,11 +1,11 @@
 from typing import Tuple, Dict, List, Optional
-from tinygrad.dtype import DType
+from tinygrad.dtype import DType, dtypes
 from tinygrad.renderer import ProgramSpec
-from tinygrad.tensor import Device, Tensor
+from tinygrad.tensor import Tensor
+from tinygrad.device import Device
 from tinygrad.engine.jit import TinyJit
 from tinygrad.nn.state import get_state_dict
 from tinygrad.helpers import Context, to_mv
-from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import Ops
 import json
 from collections import OrderedDict
@@ -13,12 +13,20 @@ from collections import OrderedDict
 EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CPU", "CUDA", "CL"]
 
 def compile_net(run:TinyJit, special_names:Dict[int,str]) -> Tuple[Dict[str,str],List[Tuple[str,List[str],List[int]]],Dict[str,Tuple[int,DType,int]],Dict[str,Tensor]]:
+  # memory-planned subbuffers can have multiple Buffer objects for the same memory region
+  canon, _seen = {}, {}
+  for ji in run.jit_cache:
+    for b in ji.bufs:
+      if b is not None: canon[id(b)] = _seen.setdefault((id(b.base._buf), b.offset, b.size, b.dtype), b)
+  special_names = {id(canon[k]): v for k, v in special_names.items() if k in canon}
+
   functions, bufs, bufs_to_save, statements, bufnum = {}, {}, {}, [], 0
   for ji in run.jit_cache:
     fxn: ProgramSpec = ji.prg.p
     functions[fxn.function_name] = fxn.src   # NOTE: this assumes all with the same name are the same
     cargs = []
     for i,arg in enumerate(ji.bufs):
+      arg = canon[id(arg)]
       key = id(arg)
       if key not in bufs:
         if key in special_names:

extra/gemm/amd_asm_matmul.py

@@ -291,7 +291,7 @@ def build_kernel(N, arch='gfx1100'):
   # MAIN GEMM LOOP
   # ===========================================================================
 
-  NO_DS, NO_GLOBAL = getenv("NO_DS", 0), getenv("NO_GLOBAL", 0)
+  NO_ALU, NO_DS, NO_GLOBAL = getenv("NO_ALU", 0), getenv("NO_DS", 0), getenv("NO_GLOBAL", 0)
 
   k.label('LOOP_INC')
   k.emit(s_add_i32(s[S_LOOP_CTR], s[S_LOOP_CTR], 8))
@@ -350,10 +350,11 @@ def build_kernel(N, arch='gfx1100'):
 
     # 64 dual FMACs
     k.waitcnt(lgkm=0)
-    k.emit(s_clause(simm16=len(FMAC_PATTERN)-1))
-    for i, (vdst_x, vdst_y, ax, bx, ay, by) in enumerate(FMAC_PATTERN):
-      k.emit(VOPD(VOPDOp.V_DUAL_FMAC_F32, VOPDOp.V_DUAL_FMAC_F32,
-                  vdstx=v[vdst_x], vdsty=v[vdst_y], srcx0=v[ax], vsrcx1=v[bx], srcy0=v[ay], vsrcy1=v[by]))
+    if not NO_ALU:
+      k.emit(s_clause(simm16=len(FMAC_PATTERN)-1))
+      for i, (vdst_x, vdst_y, ax, bx, ay, by) in enumerate(FMAC_PATTERN):
+        k.emit(VOPD(VOPDOp.V_DUAL_FMAC_F32, VOPDOp.V_DUAL_FMAC_F32,
+                    vdstx=v[vdst_x], vdsty=v[vdst_y], srcx0=v[ax], vsrcx1=v[bx], srcy0=v[ay], vsrcy1=v[by]))
 
   # wait for all global loads to finish
   # then sync the warp so it's safe to store local

extra/gemm/amd_copy_matmul.py

@@ -7,73 +7,99 @@ M = getenv("M", N)
 K = getenv("K", N)
 
 WARP_SIZE = 32
-BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 8
-TM, TN = 4, 4
-LANES_PER_WAVE_M, LANES_PER_WAVE_N = 4, 8
+BLOCK_M, BLOCK_N = 128, 128
+BLOCK_K = getenv("BK", 16)
 assert N % BLOCK_N == 0 and M % BLOCK_M == 0 and K % BLOCK_K == 0
 
-is_kernel5 = getenv("K5", 0)
-THREADS_PER_BLOCK = 128 if is_kernel5 else 256
-WAVES_PER_BLOCK_N = 1 if is_kernel5 else 2
-WAVES_PER_BLOCK_M = THREADS_PER_BLOCK // WARP_SIZE // WAVES_PER_BLOCK_N
-REG_TILES_PER_WAVE_N = BLOCK_N // (WAVES_PER_BLOCK_N * LANES_PER_WAVE_N * TN)
-REG_TILES_PER_WAVE_M = BLOCK_M // (WAVES_PER_BLOCK_M * LANES_PER_WAVE_M * TM)
+use_wmma = getenv("WMMA")
+if use_wmma:
+  WAVES_M, WAVES_N = 2, 2
+  LANES_PER_WAVE_M, LANES_PER_WAVE_N = 2, 16
+  UNROLL_M, UNROLL_N = 1, 1
 
-assert WAVES_PER_BLOCK_M*REG_TILES_PER_WAVE_M*LANES_PER_WAVE_M*TM == BLOCK_M, "M reshape is wrong"
-assert WAVES_PER_BLOCK_N*REG_TILES_PER_WAVE_N*LANES_PER_WAVE_N*TN == BLOCK_N, "N reshape is wrong"
+  # wmma params
+  WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
+  WMMA_ACC = WMMA_M // LANES_PER_WAVE_M
+else:
+  WAVES_M, WAVES_N = 4, 1
+  LANES_PER_WAVE_M, LANES_PER_WAVE_N = 4, 8
+  UNROLL_M, UNROLL_N = 4, 4
 
-consts = {"wpb_m":WAVES_PER_BLOCK_M, "lpw_m":LANES_PER_WAVE_M, "rt_m":REG_TILES_PER_WAVE_M, "t_m": TM,
-          "wpb_n":WAVES_PER_BLOCK_N, "lpw_n":LANES_PER_WAVE_N, "rt_n":REG_TILES_PER_WAVE_N, "t_n": TN}
+# WARP_SIZE * total waves
+THREADS_PER_BLOCK = WARP_SIZE * WAVES_M * WAVES_N
+
+# accumulator size
+TM = BLOCK_M // (WAVES_M * LANES_PER_WAVE_M)
+TN = BLOCK_N // (WAVES_N * LANES_PER_WAVE_N)
 
-# 128x128 out, kx128, kx128 in
 def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
-  tid = UOp.range(THREADS_PER_BLOCK, 2, AxisType.LOCAL)
+  wave_m = UOp.range(WAVES_M, 2, AxisType.LOCAL)
+  wave_n = UOp.range(WAVES_N, 3, AxisType.LOCAL)
+  lane = UOp.range(WARP_SIZE, -1, AxisType.WARP)
+  tid = (wave_m * WAVES_N + wave_n) * WARP_SIZE + lane
 
-  #tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
-  warp, lane = tid // WARP_SIZE, tid % WARP_SIZE
-  wave_n, wave_m = warp % WAVES_PER_BLOCK_N, warp // WAVES_PER_BLOCK_N
-  lane_n, lane_m = lane % LANES_PER_WAVE_N, lane // LANES_PER_WAVE_N
+  # -- GLOBAL -> LOCAL --
+  # wmma: spatial outer, k inner (k contiguous for vectorized WMMA tile loads)
+  # gemm: k outer, spatial inner
+  A_local = UOp.placeholder((BLOCK_M, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_M), a.dtype.base, slot=0, addrspace=AddrSpace.LOCAL)
+  B_local = UOp.placeholder((BLOCK_N, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_N), b.dtype.base, slot=1, addrspace=AddrSpace.LOCAL)
 
-  # define locals
-  A_local = UOp.placeholder((BLOCK_K, BLOCK_M), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
-  B_local = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
-
-  # open the main reduction range and copy in GLOBAL -> LOCAL
   a = a.reshape(K // BLOCK_K, BLOCK_K, BLOCK_M)
   b = b.reshape(K // BLOCK_K, BLOCK_K, BLOCK_N)
-  k_tile_range = UOp.range(K // BLOCK_K, 3, AxisType.REDUCE)
-  A_store = A_local.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(a[k_tile_range].reshape(-1, THREADS_PER_BLOCK)[:, tid])
-  B_store = B_local.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(b[k_tile_range].reshape(-1, THREADS_PER_BLOCK)[:, tid])
+  k_tile = UOp.range(K // BLOCK_K, 100, AxisType.REDUCE)
+
+  # copy with transpose for wmma (input is k×spatial, LDS is spatial×k)
+  A_copy = A_local.permute((1,0)) if use_wmma else A_local
+  B_copy = B_local.permute((1,0)) if use_wmma else B_local
+  A_store = A_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(a[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
+  B_store = B_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(b[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
   barrier = UOp.barrier(A_store, B_store)
   A_local, B_local = A_local.after(barrier), B_local.after(barrier)
 
-  # define accumulator (128x128), but broadcast across tid
-  c_regs = UOp.placeholder((REG_TILES_PER_WAVE_M*TM, REG_TILES_PER_WAVE_N*TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
-  c_regs = c_regs.after(c_regs.store(UOp.const(dtypes.float, 0).reshape((1,)*len(c_regs.shape)).expand(c_regs.shape)))
+  # -- COMPUTE --
+  lane_m, lane_n = lane // LANES_PER_WAVE_N, lane % LANES_PER_WAVE_N
 
-  # define registers (NOTE: the thread count is the device count for this multi, it's sharded across the THREADS_PER_BLOCK)
-  A_col = UOp.placeholder((REG_TILES_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
-  B_row = UOp.placeholder((REG_TILES_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+  # accumulator (unified: both paths use (TM, TN) with scalar dtypes.float)
+  acc = UOp.placeholder((TM, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
+  acc = acc.after(acc.store(UOp.const(dtypes.float, 0).reshape((1,)*len(acc.shape)).expand(acc.shape)))
 
-  # LOCAL -> REGS
-  k = UOp.range(BLOCK_K, 4, AxisType.REDUCE)
-  A_col = A_col.after(A_col.store(A_local[k].reshape(WAVES_PER_BLOCK_M, REG_TILES_PER_WAVE_M, LANES_PER_WAVE_M, TM)[wave_m, :, lane_m, :]))
-  B_row = B_row.after(B_row.store(B_local[k].reshape(WAVES_PER_BLOCK_N, REG_TILES_PER_WAVE_N, LANES_PER_WAVE_N, TN)[wave_n, :, lane_n, :]))
+  if use_wmma:
+    k = UOp.range(BLOCK_K // WMMA_K, 101, AxisType.REDUCE)
+    tile_m = UOp.range(TM // WMMA_ACC, 200, AxisType.LOOP)
+    tile_n = UOp.range(TN, 201, AxisType.LOOP)
 
-  # do FMA
-  A_col = A_col.reshape(REG_TILES_PER_WAVE_M*TM, 1).expand(REG_TILES_PER_WAVE_M*TM, REG_TILES_PER_WAVE_N*TN)
-  B_row = B_row.reshape(1, REG_TILES_PER_WAVE_N*TN).expand(REG_TILES_PER_WAVE_M*TM, REG_TILES_PER_WAVE_N*TN)
-  c_regs = c_regs.after(c_regs.store(c_regs.after(k) + (A_col * B_row)).end(k).barrier().end(k_tile_range))
+    acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0,2,1)[tile_m, tile_n]
+    a_frag = A_local.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_K // WMMA_K, WMMA_K)[wave_m, tile_m, lane_n, k]
+    b_frag = B_local.reshape(WAVES_N, TN, WMMA_N, BLOCK_K // WMMA_K, WMMA_K)[wave_n, tile_n, lane_n, k]
 
-  # store back to c
-  c_store = c.rearrange("(wpb_m rt_m lpw_m t_m) (wpb_n rt_n lpw_n t_n) -> (wpb_m wpb_n lpw_m lpw_n) (rt_m t_m) (rt_n t_n)", **consts)
-  return c_store[tid].store(c_regs).end(tid)
+    wmma = UOp(Ops.SHAPED_WMMA, dtypes.float, (a_frag, b_frag, acc_frag.after(k)), arg=((16, 16, 16), 'AMD', 32))
+    acc_store = acc_frag.store(wmma).end(tile_m, tile_n)
+  else:
+    # registers for LOCAL -> REG
+    a_frag = UOp.placeholder((TM//UNROLL_M, UNROLL_M), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+    b_frag = UOp.placeholder((TN//UNROLL_N, UNROLL_N), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+
+    k = UOp.range(BLOCK_K, 101, AxisType.REDUCE)
+    a_frag = a_frag.after(a_frag.store(A_local[k].reshape(WAVES_M, TM//UNROLL_M, LANES_PER_WAVE_M, UNROLL_M)[wave_m, :, lane_m, :]))
+    b_frag = b_frag.after(b_frag.store(B_local[k].reshape(WAVES_N, TN//UNROLL_N, LANES_PER_WAVE_N, UNROLL_N)[wave_n, :, lane_n, :]))
+
+    # FMA
+    a_frag = a_frag.reshape(TM, 1).expand(TM, TN)
+    b_frag = b_frag.reshape(1, TN).expand(TM, TN)
+    acc_store = acc.store(acc.after(k) + (a_frag * b_frag))
+
+  # store accumulator and loop
+  acc = acc.after(acc_store.end(k).barrier().end(k_tile))
+
+  # store accumulator to output (unified)
+  c = c.reshape(WAVES_M, TM//UNROLL_M, LANES_PER_WAVE_M, UNROLL_M,
+                WAVES_N, TN//UNROLL_N, LANES_PER_WAVE_N, UNROLL_N)
+  c = c.permute((0,4,2,6, 1,3,5,7)).reshape(THREADS_PER_BLOCK, TM, TN)
+  return c[tid].store(acc).end(wave_m, wave_n, lane)
 
 def amd_copy_matmul(c:UOp, a:UOp, b:UOp) -> UOp:
-  block_id_n = UOp.range(N // BLOCK_N, 0, AxisType.GLOBAL)
-  block_id_m = UOp.range(M // BLOCK_M, 1, AxisType.GLOBAL)
-
-  # index the output with the globals
+  block_id_m = UOp.range(M // BLOCK_M, 0, AxisType.GLOBAL)
+  block_id_n = UOp.range(N // BLOCK_N, 1, AxisType.GLOBAL)
   c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[block_id_m, :, block_id_n, :]
   a = a.T.reshape(K, M // BLOCK_M, BLOCK_M)[:, block_id_m, :]
   b = b.reshape(K, N // BLOCK_N, BLOCK_N)[:, block_id_n, :]
@@ -81,4 +107,4 @@ def amd_copy_matmul(c:UOp, a:UOp, b:UOp) -> UOp:
 
 if __name__ == "__main__":
   from amd_uop_matmul import eval_custom_matmul
-  eval_custom_matmul(amd_copy_matmul)
+  eval_custom_matmul(amd_copy_matmul, dtypes.half if use_wmma else dtypes.float)

extra/gemm/amd_flash_attention.py

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

extra/gemm/amd_uop_matmul.py

@@ -115,14 +115,14 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
 
   return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()
 
-def eval_custom_matmul(fxn):
-  a = Tensor.randn(M, K, dtype=dtypes.float)
-  b = Tensor.randn(K, N, dtype=dtypes.float)
+def eval_custom_matmul(fxn, dt=dtypes.float):
+  a = Tensor.randn(M, K, dtype=dt)
+  b = Tensor.randn(K, N, dtype=dt)
   c = Tensor.empty(M, N, dtype=dtypes.float)
   with Context(DEBUG=0): Tensor.realize(a, b)
 
   ets = []
-  with Context(DEBUG=max(2, DEBUG.value)):
+  with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2 if dt == dtypes.half else 0):
     for _ in range(NUM_RUNS):
       GlobalCounters.reset()
       tst = Tensor.custom_kernel(c, a, b, fxn=fxn)[0].realize()
@@ -132,11 +132,11 @@ def eval_custom_matmul(fxn):
   if getenv("VERIFY", 1):
     GlobalCounters.reset()
     with Context(DEBUG=2):
-      tc = (a @ b).realize()
+      tc = (a.float() @ b.float()).realize()
     with Context(DEBUG=0):
       err = (tc - tst).square().mean().item()
     print(f"mean squared error {err}")
-    if err > 1e-06:
+    if err > (1e-2 if dt == dtypes.half else 1e-6):
       raise RuntimeError("matmul is wrong!")
 
 if __name__ == "__main__":

extra/gemm/cdna_asm_gemm.py

@@ -2649,6 +2649,9 @@ def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
   else: dname = a.device
   arch = getattr(Device[dname].renderer, "arch", "")
   if batch not in {1, 2}: return todo(f"GEMM batch size {batch}")
+  # blacklist slow matmul
+  # TODO: why is this slow?
+  if (M,N,K) == (8192, 2304, 16384): return todo("blacklisted slow matmul")
   if (M % TILE_M != 0 or N % TILE_N != 0 or K % TILE_K != 0) and arch == "gfx950":
     return todo(f"GEMM shape ({M},{N},{K}) not a multiple of ({TILE_M},{TILE_N},{TILE_K})")
   return True
@@ -2663,10 +2666,10 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
   m = UOp.range(M, 1, AxisType.LOOP)
   n = UOp.range(N, 2, AxisType.LOOP)
   k = UOp.range(K, 0, AxisType.REDUCE)
-  mul = (A.flatten().index((m*UOp.const(dtypes.index, K)+k))*
-         B.flatten().index((k*UOp.const(dtypes.index, N)+n))).cast(dtypes.float32)
+  mul = (A.flatten().index((m*UOp.const(dtypes.weakint, K)+k))*
+         B.flatten().index((k*UOp.const(dtypes.weakint, N)+n))).cast(dtypes.float32)
   red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype.base)
-  store = C.flatten().index((m*UOp.const(dtypes.index, N)+n), ptr=True).store(red).end(m, n)
+  store = C.flatten().index((m*UOp.const(dtypes.weakint, N)+n), ptr=True).store(red).end(m, n)
   return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
 
 # ** backward gemm, might use the asm gemm
@@ -2702,13 +2705,14 @@ def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
 
   if is_multi:
     if n_sharded:
-      out = Tensor(Tensor.empty(batch, M, N//len(a.device), dtype=a.dtype, device=a.device).uop.multi(2), device=a.device)
+      out = Tensor(Tensor.invalid(batch, M, N//len(a.device), dtype=a.dtype, device=a.device).uop.multi(2), device=a.device)
     elif m_sharded:
-      out = Tensor(Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device).uop.multi(1), device=a.device)
+      out = Tensor(Tensor.invalid(batch, M, N, dtype=a.dtype, device=a.device).uop.multi(1), device=a.device)
     else:
-      out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
+      out = Tensor(Tensor.invalid(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0),
+                   device=a.device)
   else:
-    out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)
+    out = Tensor.invalid(batch, M, N, dtype=a.dtype, device=a.device)
 
   renderer = Device[a.device[0] if is_multi else a.device].renderer
   dname, arch = renderer.device, getattr(renderer, "arch", "")

extra/remote/bench.py

@@ -0,0 +1,46 @@
+#!/usr/bin/env python3
+import os, sys, time
+from tinygrad.runtime.support.system import RemotePCIDevice
+
+LAT_N_RUNS = 500
+THROUGHPUT_N_RUNS = 8
+SIZES = [4, 1 << 10, 8 << 20]
+
+if __name__ == "__main__":
+  os.environ["REMOTE"] = sys.argv[1] if len(sys.argv) > 1 else os.environ.get("REMOTE", "127.0.0.1:6667")
+
+  # choose any amd/nv gpu.
+  devs = RemotePCIDevice.remote_list(0x1002, ((0, (0,)),), 0) or RemotePCIDevice.remote_list(0x10de, ((0, (0,)),), 0x03)
+  if not devs: raise RuntimeError("no GPU found on remote")
+
+  pci = RemotePCIDevice("BN", devs[0])
+  print(f"connected to {os.environ['REMOTE']}, device: {devs[0]}\n")
+
+  # ping (minimal server round-trip, no device I/O)
+  from tinygrad.runtime.support.system import RemoteCmd
+  sock = pci.sock
+  for _ in range(10): RemotePCIDevice._rpc(sock, 0, RemoteCmd.PING)
+  st = time.perf_counter()
+  for _ in range(LAT_N_RUNS): RemotePCIDevice._rpc(sock, 0, RemoteCmd.PING)
+  ping_lat = (time.perf_counter() - st) / LAT_N_RUNS
+  print(f"PING latency: {ping_lat*1e6:.1f} us ({1/ping_lat:,.0f} ops/sec)\n")
+
+  # throughput
+  sysmem, _ = pci.alloc_sysmem(max(SIZES))
+  print(f"{'size':>10s}  {'write MB/s':>10s}  {'read MB/s':>10s}")
+  for sz in SIZES:
+    data = b'\x01' * sz
+
+    for _ in range(5): sysmem[0:sz] = data
+    st = time.perf_counter()
+    for _ in range(THROUGHPUT_N_RUNS): sysmem[0:sz] = data
+    pci.read_config(0, 4) # flush, since writes are posted
+    w = (time.perf_counter() - st) / THROUGHPUT_N_RUNS
+
+    for _ in range(5): sysmem[0:sz]
+    st = time.perf_counter()
+    for _ in range(THROUGHPUT_N_RUNS): sysmem[0:sz]
+    r = (time.perf_counter() - st) / THROUGHPUT_N_RUNS
+
+    sfx, div = [('B',1),('K',1<<10),('M',1<<20)][[sz>=1<<10,sz>=1<<20,sz>=1<<30].count(True)]
+    print(f"{sz/div:>9.4g}{sfx}  {sz/w/1e6:>10.1f}  {sz/r/1e6:>10.1f}")

extra/remote/serve.py

@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+import socket, struct, sys
+from tinygrad.runtime.support.system import PCIDevice, RemoteCmd, System
+from tinygrad.helpers import DEBUG, OSX
+
+def resp(resp0=0, resp1=0, status=0): return struct.pack('<BQQ', status, resp0, resp1)
+def resp_err(msg): return struct.pack('<BQQ', 1, len(err:=msg.encode()), 0) + err
+
+discovered_devices: list[str] = []
+opened_devices: dict[int, PCIDevice] = {}
+mapped_bars: dict[tuple[int, int], object] = {}
+sysmem_allocs: list[tuple] = []
+
+def handle(conn, cmd, dev_id, bar, arg0, arg1, arg2):
+  if cmd == RemoteCmd.PING:
+    return conn.sendall(resp())
+
+  if cmd == RemoteCmd.PROBE:
+    payload = conn.recv(arg1, socket.MSG_WAITALL) if arg1 > 0 else b""
+    filter_devices: dict[int, list[int]] = {}
+    for i in range(0, len(payload), 8):
+      mask, dev = struct.unpack('<II', payload[i:i+8])
+      filter_devices.setdefault(mask, []).append(dev)
+    base_class = None if arg0 == 0 else int(arg0)
+    devs = System.list_devices(arg2, tuple([(x, tuple(y)) for x,y in filter_devices.items()]), base_class)
+    for p in devs:
+      if p not in discovered_devices: discovered_devices.append(p)
+    data = "\n".join(f"{p[1]}:{discovered_devices.index(p)}" for p in devs).encode()
+    return conn.sendall(resp(len(data), len(devs)) + data)
+
+  # lazy device open
+  if dev_id not in opened_devices:
+    if dev_id >= len(discovered_devices): raise RuntimeError(f"device {dev_id} not probed")
+    cl, pcibus = discovered_devices[dev_id]
+    opened_devices[dev_id] = cl("SV", pcibus)
+  pci_dev = opened_devices[dev_id]
+
+  if cmd == RemoteCmd.MAP_BAR:
+    if (dev_id, bar) not in mapped_bars: mapped_bars[(dev_id, bar)] = pci_dev.map_bar(bar)
+    conn.sendall(resp(*pci_dev.bar_info(bar)))
+  elif cmd == RemoteCmd.CFG_READ:
+    conn.sendall(resp(pci_dev.read_config(arg0, arg1)))
+  elif cmd == RemoteCmd.CFG_WRITE:
+    pci_dev.write_config(arg0, arg2, arg1)
+    conn.sendall(resp())
+  elif cmd == RemoteCmd.RESIZE_BAR:
+    pci_dev.resize_bar(bar)
+    conn.sendall(resp())
+  elif cmd == RemoteCmd.RESET:
+    pci_dev.reset()
+    conn.sendall(resp())
+  elif cmd == RemoteCmd.MMIO_READ:
+    bar_view = mapped_bars[(dev_id, bar)]
+    if arg0 % 4 == 0 and arg1 == 4: conn.sendmsg([resp(arg1), struct.pack(f'<{arg1 // 4}I', bar_view.view(arg0, arg1, fmt='I')[0])])
+    else: conn.sendmsg([resp(arg1), bar_view[arg0:arg0+arg1]])
+  elif cmd == RemoteCmd.MMIO_WRITE:
+    data = conn.recv(arg1, socket.MSG_WAITALL)
+    bar_view = mapped_bars[(dev_id, bar)]
+    if arg0 % 4 == 0 and arg1 == 4: bar_view.view(arg0, arg1, fmt='I')[0] = struct.unpack(f'<{arg1 // 4}I', data)[0]
+    else: bar_view[arg0:arg0+arg1] = data
+  elif cmd == RemoteCmd.MAP_SYSMEM:
+    memview, paddrs = pci_dev.alloc_sysmem(arg0, contiguous=bool(arg1))
+    sysmem_allocs.append((memview, paddrs))
+    paddrs_bytes = struct.pack(f'<{len(paddrs)}Q', *paddrs)
+    conn.sendall(resp(len(paddrs_bytes), len(sysmem_allocs) - 1) + paddrs_bytes)
+  elif cmd == RemoteCmd.SYSMEM_READ:
+    conn.sendmsg([resp(arg1), sysmem_allocs[bar][0][arg0:arg0+arg1]])
+  elif cmd == RemoteCmd.SYSMEM_WRITE:
+    sysmem_allocs[bar][0][arg0:arg0+arg1] = conn.recv(arg1, socket.MSG_WAITALL)
+  else: raise RuntimeError(f"unknown command {cmd}")
+
+def serve(conn:socket.socket):
+  REQ = '<BIIQQQ'
+  while True:
+    hdr = conn.recv(struct.calcsize(REQ), socket.MSG_WAITALL)
+    if len(hdr) < struct.calcsize(REQ): raise ConnectionError("client disconnected")
+    cmd, dev_id, bar, arg0, arg1, arg2 = struct.unpack(REQ, hdr)
+    if DEBUG >= 4: print(f"cmd={RemoteCmd(cmd).name} dev={dev_id} bar={bar} arg0={arg0:#x} arg1={arg1:#x} arg2={arg2:#x}")
+    try: handle(conn, cmd, dev_id, bar, arg0, arg1, arg2)
+    except ConnectionError: raise
+    except Exception as e:
+      if cmd in {RemoteCmd.MMIO_WRITE, RemoteCmd.SYSMEM_WRITE}: raise ConnectionError(f"write failed: {e}")
+      print(f"ERROR: {e}")
+      conn.sendall(resp_err(str(e)))
+
+if __name__ == "__main__":
+  if not OSX: System.reserve_hugepages(128) # for sysmem allocations
+
+  port = int(sys.argv[1]) if len(sys.argv) > 1 else 6667
+  server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+  server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+  server.bind(("0.0.0.0", port))
+  server.listen(1)
+  s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+  try: s.connect(("8.8.8.8", 80)); ip = s.getsockname()[0]
+  finally: s.close()
+  print(f"listening on {ip}:{port}")
+  while True:
+    conn, addr = server.accept()
+    conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+    for bt in [socket.SO_SNDBUF, socket.SO_RCVBUF]: conn.setsockopt(socket.SOL_SOCKET, bt, 64 << 20)
+    try: serve(conn)
+    except ConnectionError: print("disconnected")

extra/setup_hipcc_osx.sh

@@ -0,0 +1,31 @@
+#!/bin/sh
+install_loc="$HOME/.local/bin"
+docker build --platform=linux/amd64 -t rocm-hipcc:7.2 - <<'EOF'
+FROM ubuntu:22.04
+ENV DEBIAN_FRONTEND=noninteractive
+ENV TZ=Etc/UTC
+RUN apt-get update && apt-get install -y --no-install-recommends wget ca-certificates gnupg tzdata && \
+  wget https://repo.radeon.com/amdgpu-install/7.2/ubuntu/jammy/amdgpu-install_7.2.70200-1_all.deb && \
+  apt-get install -y ./amdgpu-install_7.2.70200-1_all.deb && \
+  amdgpu-install -y --usecase=rocm --no-dkms --no-32 && \
+  rm -rf /var/lib/apt/lists/*
+ENV PATH=/opt/rocm/bin:$PATH
+EOF
+
+mkdir -p "$install_loc"
+tee "$install_loc/hipccshim" >/dev/null <<'EOF'
+#!/bin/sh
+set -eu
+cname="rocm-hipcc-persistent"
+if ! docker inspect --format='{{.State.Running}}' "$cname" 2>/dev/null | grep -q true; then
+  docker rm -f "$cname" 2>/dev/null || true
+  docker run -d --platform=linux/amd64 --name "$cname" \
+    -v /var/folders:/var/folders -v "$HOME":"$HOME" \
+    rocm-hipcc:7.2 sleep 300 >/dev/null
+fi
+exec docker exec "$cname" "$(basename "$0")" "$@"
+EOF
+chmod +x "$install_loc/hipccshim"
+for t in hipcc hipconfig; do
+  ln -sf "$install_loc/hipccshim" "$install_loc/$t"
+done

extra/setup_nvcc_osx.sh

@@ -1,9 +1,9 @@
 #!/bin/sh
 install_loc="$HOME/.local/bin"
-docker build --platform=linux/amd64 -t cuda-nvcc:12.8 - <<'EOF'
+docker build --platform=linux/arm64 -t cuda-nvcc:12.8 - <<'EOF'
 FROM ubuntu:22.04
 RUN apt-get update && apt-get install -y --no-install-recommends wget ca-certificates && \
-  wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb && \
+  wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-keyring_1.1-1_all.deb && \
   dpkg -i cuda-keyring_1.1-1_all.deb && \
   apt-get update && apt-get install -y --no-install-recommends cuda-nvcc-12-8 cuda-nvdisasm-12-8 cuda-cuobjdump-12-8 && rm -rf /var/lib/apt/lists/*
 ENV PATH=/usr/local/cuda/bin:$PATH
@@ -13,10 +13,14 @@ mkdir -p "$install_loc"
 tee "$install_loc/nvccshim" >/dev/null <<'EOF'
 #!/bin/sh
 set -eu
-# assume the final arg is the input path
-# mount it so that container can read it
-dir=$(dirname "${@: -1}")
-exec docker run --rm --platform=linux/amd64 -v "$dir":"$dir" cuda-nvcc:12.8 "$(basename "$0")" "$@"
+cname="cuda-nvcc-persistent"
+if ! docker inspect --format='{{.State.Running}}' "$cname" 2>/dev/null | grep -q true; then
+  docker rm -f "$cname" 2>/dev/null || true
+  docker run -d --platform=linux/arm64 --name "$cname" \
+    -v /var/folders:/var/folders -v "$HOME":"$HOME" \
+    cuda-nvcc:12.8 sleep 300 >/dev/null
+fi
+exec docker exec "$cname" "$(basename "$0")" "$@"
 EOF
 chmod +x "$install_loc/nvccshim"
 for t in nvcc nvdisasm; do

extra/sqtt/examples/generate_examples.py

@@ -9,7 +9,7 @@ EXAMPLES = {
   "empty":"test/backend/test_custom_kernel.py TestCustomKernel.test_empty",
   "plus":"test/test_tiny.py TestTiny.test_plus",
   "gemm":"-c \"from tinygrad import Tensor; (Tensor.empty(N:=32, N)@Tensor.empty(N, N)).realize()\"",
-  "sync":"test/amd/test_custom_kernel.py TestCustomKernel.test_wave_sync",
+  "sync":"test/amd/test_custom_kernel.py TestCustomKernel.test_lds_sync",
 }
 
 if __name__ == "__main__":

extra/thunder/amd/fa.py

@@ -10,11 +10,11 @@ from tinygrad.uop.ops import UOp, Ops, KernelInfo
 
 def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None=None) -> Tensor:
   dtype = dtype or ref.dtype
-  if not isinstance(ref.device, tuple): return Tensor.empty(*shape, dtype=dtype, device=ref.device)
+  if not isinstance(ref.device, tuple): return Tensor.invalid(*shape, dtype=dtype, device=ref.device)
   shard_axis = ref.uop.axis if axis is None else axis
   shape = tuple(s // len(ref.device) if i == shard_axis else s for i, s in enumerate(shape))
   axis = ref.uop.axis if axis is None else axis
-  return Tensor(Tensor.empty(*shape, dtype=dtype, device=ref.device).uop.multi(axis), dtype=dtype, device=ref.device)
+  return Tensor(Tensor.invalid(*shape, dtype=dtype, device=ref.device).uop.multi(axis), dtype=dtype, device=ref.device)
 
 def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
   return _sharded_empty(ref.shape, ref, axis)

extra/tinydreno.h

@@ -0,0 +1,113 @@
+#ifndef _TINYDRENO_H
+#define _TINYDRENO_H
+
+#include <stdint.h>
+#include <stddef.h>
+
+typedef void * cl_llvm_instance;
+
+cl_llvm_instance cl_compiler_create_llvm_instance(void);
+void cl_compiler_destroy_llvm_instance(cl_llvm_instance inst);
+
+enum cl_handle_type {
+  CL_HANDLE_COMPILED = 1,
+  CL_HANDLE_LIBRARY,
+  CL_HANDLE_LINKED
+};
+
+// handle->data for CL_HANDLE_COMPILED and CL_HANDLE_LIBRARY
+struct cl_compiled_data {
+  uint64_t chip_id;
+  uint32_t mode;
+  void    *llvm_bitcode;
+  uint64_t llvm_bitcode_size;
+  char    *build_log;
+  uint32_t build_log_len;
+  uint32_t error_code;
+};
+
+// handle->data for CL_HANDLE_LINKED
+struct cl_executable_data {
+  int32_t  num_kernels;
+  void    *kernel_props;
+  uint32_t error_code;
+  char    *build_log;
+  char     _unk0[0x20];
+  uint64_t chip_id;
+  uint32_t mode;
+};
+
+typedef struct {
+  enum cl_handle_type type;
+  union {
+    struct cl_compiled_data *compiled;
+    struct cl_executable_data *executable;
+  };
+} cl_handle;
+
+
+#define CL_MODE_32BIT 0
+#define CL_MODE_64BIT 1
+
+#define CL_SRC_STR  0
+#define CL_SRC_BLOB 1
+
+cl_handle *cl_compiler_compile_source(cl_llvm_instance inst, uint64_t chip_id, int mode, const char *options, int p5, uint64_t p6, uint64_t p7,
+                                      const char *source, uint64_t source_len, uint64_t source_type, void *p11);
+
+cl_handle *cl_compiler_link_program(cl_llvm_instance inst, uint64_t chip_id, int mode, const char *options, int num_handles,
+                                    cl_handle **input_handles);
+
+
+void cl_compiler_handle_create_binary(cl_handle *handle, void **out_ptr, size_t *out_size);
+
+// lib binary format (output of handle_create_binary for type 3)
+// layout: cl_lib_header, then cl_lib_section[num_sections], then data
+
+#define CL_LIB_PROGRAM     0
+#define CL_LIB_CONSTS      6
+#define CL_LIB_IMAGE       7
+#define CL_LIB_CODE       10
+#define CL_LIB_IMAGE_DESC 11
+
+typedef struct {
+  uint32_t id;
+  uint32_t offset;
+  uint32_t size;
+  uint32_t count;
+  uint32_t entry_size;
+} cl_lib_section;
+
+typedef struct {
+  uint32_t _unk0[6];
+  uint32_t num_sections;
+  uint32_t _unk1[5];
+  cl_lib_section sections[];
+} cl_lib_header;
+
+// at sections[CL_LIB_PROGRAM].offset
+typedef struct {
+  char     name[8];
+  uint32_t _unk0[3];
+  uint32_t fregs;
+  uint32_t hregs;
+} cl_lib_prog;
+
+// at sections[CL_LIB_IMAGE_DESC].offset
+typedef struct {
+  char     _unk0[0xc4];
+  uint32_t prg_offset;
+  uint32_t pvtmem;
+  char     _unk1[0x0c];
+  uint32_t shmem;
+  uint32_t samp_cnt;
+  char     _unk2[0x28];
+  uint32_t brnchstck;
+  char     _unk4[0x4c];
+  char     kernel_name[];
+} cl_lib_img_desc;
+
+void cl_compiler_free_handle(cl_handle *handle);
+void cl_compiler_free_assembly(void *ptr);
+
+#endif

extra/torch_backend/example.py

@@ -24,7 +24,7 @@ if __name__ == "__main__":
   kernel_count = GlobalCounters.kernel_count
   assert kernel_count > 0, "No kernels, test failed"
   # NOTE: this is 124 on torch 2.10.0
-  expected_kernels = 332
+  expected_kernels = 334
   expectation = f"ResNet18 kernels are {kernel_count} vs {expected_kernels} expected."
   if kernel_count < expected_kernels: warnings.warn(f"{expectation} Expectation can be lowered.", UserWarning)
-  assert kernel_count <= expected_kernels, f"{expectation}"
+  assert kernel_count <= expected_kernels, f"{expectation}"

extra/torch_backend/test_kernel_fusion.py

@@ -23,7 +23,7 @@ class TestKernelFusionRegression(unittest.TestCase):
     def fn():
       x = torch.randn(128, 128, device=device)
       return (x + 1.0) * 2.0 - 0.5
-    self._check_kernel_count(fn, 5)
+    self._check_kernel_count(fn, 7)
 
   def test_relu_fusion(self):
     def fn():
@@ -31,7 +31,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       conv = torch.nn.Conv2d(3, 16, 3, padding=1).to(device)
       with torch.no_grad():
         return torch.nn.functional.relu(conv(x))
-    self._check_kernel_count(fn, 6)
+    self._check_kernel_count(fn, 8)
 
   def test_batchnorm_fusion(self):
     def fn():
@@ -41,26 +41,26 @@ class TestKernelFusionRegression(unittest.TestCase):
       bn.eval()
       with torch.no_grad():
         return torch.nn.functional.relu(bn(conv(x)))
-    self._check_kernel_count(fn, 10)
+    self._check_kernel_count(fn, 12)
 
   def test_reduce_fusion(self):
     def fn():
       x = torch.randn(64, 64, device=device)
       return (x * 2.0).sum()
-    self._check_kernel_count(fn, 5)
+    self._check_kernel_count(fn, 7)
 
   def test_matmul_elementwise_fusion(self):
     def fn():
       x = torch.randn(32, 32, device=device)
       w = torch.randn(32, 32, device=device)
       return torch.nn.functional.relu(x @ w + 1.0)
-    self._check_kernel_count(fn, 7)
+    self._check_kernel_count(fn, 9)
 
   def test_pooling_fusion(self):
     def fn():
       x = torch.randn(1, 8, 16, 16, device=device)
       return torch.nn.functional.max_pool2d(x * 2.0, 2)
-    self._check_kernel_count(fn, 5)
+    self._check_kernel_count(fn, 7)
 
   def test_residual_add_relu_fusion(self):
     def fn():
@@ -68,7 +68,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       identity = torch.randn(1, 8, 16, 16, device=device)
       out = x + identity
       return torch.nn.functional.relu(out)
-    self._check_kernel_count(fn, 7)
+    self._check_kernel_count(fn, 9)
 
   def test_inplace_add_relu_fusion(self):
     def fn():
@@ -76,7 +76,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       y = torch.randn(1, 16, 32, 32, device=device)
       x += y
       return torch.nn.functional.relu(x)
-    self._check_kernel_count(fn, 7)
+    self._check_kernel_count(fn, 9)
 
   def test_conv_bn_add_relu_fusion(self):
     def fn():
@@ -89,7 +89,7 @@ class TestKernelFusionRegression(unittest.TestCase):
         out = bn(conv(x))
         out += identity
         return torch.nn.functional.relu(out)
-    self._check_kernel_count(fn, 12)
+    self._check_kernel_count(fn, 14)
 
   def test_multiple_inplace_ops_fusion(self):
     def fn():
@@ -97,7 +97,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       x += 1.0
       x *= 2.0
       return torch.nn.functional.relu(x)
-    self._check_kernel_count(fn, 4)
+    self._check_kernel_count(fn, 6)
 
   def test_view_inplace_no_fusion_break(self):
     def fn():
@@ -105,7 +105,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       view = x[1:3]
       view += 1.0
       return x.sum()
-    self._check_kernel_count(fn, 8)
+    self._check_kernel_count(fn, 10)
 
   def test_batchnorm_running_stats_update(self):
     def fn():
@@ -114,7 +114,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       bn.train()
       with torch.no_grad():
         return bn(x)
-    self._check_kernel_count(fn, 8)
+    self._check_kernel_count(fn, 10)
 
   # this is a minimal extra/other_mnist/beautiful_mnist_torch.py to cover fusion for training with optimizer
   def test_mnist_training_fusion(self):
@@ -135,7 +135,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       loss.backward()
       optimizer.step()
       return loss
-    self._check_kernel_count(fn, 24)
+    self._check_kernel_count(fn, 26)
 
 if __name__ == "__main__":
   unittest.main()

extra/usbgpu/tbgpu/installer/Shared/server.c

@@ -28,6 +28,7 @@ enum {
   CMD_MAP_SYSMEM = 8,     // map system memory
   CMD_SYSMEM_READ = 9,    // bulk read from system memory
   CMD_SYSMEM_WRITE = 10,  // bulk write to system memory
+  CMD_RESIZE_BAR = 11,    // resize bar (noop)
   RESP_OK = 0, RESP_ERR = 1,
 };
 
@@ -125,7 +126,7 @@ static int map_bar(uint32_t bar, response_t *resp) {
   return 0;
 }
 
-static int map_sysmem_fd(uint64_t size, response_t *resp, int *out_fd) {
+static int map_sysmem_fd(uint64_t size, int contiguous, response_t *resp, int *out_fd) {
   if (g_sysmem_count >= MAX_SYSMEM) return -1;
   int idx = g_sysmem_count;
   int fd = -1;
@@ -207,7 +208,7 @@ static void handle_client(int fd) {
 
     case CMD_MAP_SYSMEM_FD: {
       int shm_fd = -1;
-      resp.status = map_sysmem_fd(req.arg0, &resp, &shm_fd) ? 1 : 0;
+      resp.status = map_sysmem_fd(req.arg0, (int)req.arg1, &resp, &shm_fd) ? 1 : 0;
       send_response(fd, &resp, shm_fd);
       continue;
     }
@@ -224,6 +225,9 @@ static void handle_client(int fd) {
       break;
     }
 
+    case CMD_RESIZE_BAR:
+      break;
+
     case CMD_RESET:
       resp.status = dext_rpc(2, NULL, 0, NULL) ? 1 : 0;
       break;

extra/viz/README

@@ -1,6 +1,7 @@
 A command line tool for exploring the VIZ trace.
 
-After running with VIZ=-1, use `extra/viz/cli.py` to explore the saved trace files.
+1. Set VIZ to -1 to save the trace.
+2. Use `extra/viz/cli.py` to inspect the trace files.
 
 ## Inspect runtime profiling
 

extra/viz/cli.py

@@ -1,7 +1,5 @@
 #!/usr/bin/env python3
-import os
-os.environ["VIZ"] = "0"
-import argparse, pathlib, sys, struct, json
+import argparse, pathlib, sys, struct, json, itertools
 from typing import Iterator
 from tinygrad.viz import serve as viz
 from tinygrad.uop.ops import RewriteTrace
@@ -67,9 +65,12 @@ if __name__ == "__main__":
   g_mode.add_argument("--rewrites", action="store_true", help="View rewrites trace")
   g_common = parser.add_argument_group("common options")
   g_common.add_argument("--kernel", type=str, default=None, metavar="NAME", help="Select a kernel by name (optional name, default: only list names)")
+  g_common.add_argument("--no-color", action="store_true", default=not (sys.stdin.isatty() and sys.stdout.isatty()),
+                        help="Disable colored output (default: true in non-interactive mode)")
   g_profile = parser.add_argument_group("profile options")
   g_profile.add_argument("--device", type=str, default=None, metavar="NAME", help="Select a device (optional name, default: only list names)")
-  g_profile.add_argument("--top", type=int, default=10, metavar="N", help="Number of top kernels to show (-1 for all, default: 10)")
+  g_profile.add_argument("--offset", type=int, default=0, metavar="N", help="event offset (default: 0)")
+  g_profile.add_argument("--limit", type=int, default=10, metavar="N", help="events to display (-1 for all, default: 10)")
   g_rewrites = parser.add_argument_group("rewrites options")
   g_rewrites.add_argument("--select", type=str, default=None, metavar="NAME",
                           help="Select an item within the chosen kernel (optional name, default: only list names)")
@@ -85,14 +86,64 @@ if __name__ == "__main__":
   viz.trace = viz.load_pickle(args.rewrites_path, default=RewriteTrace([], [], {}))
   viz.ctxs = viz.get_rewrites(viz.trace)
 
+  def format_colored(s:str) -> str: return ansistrip(s) if args.no_color else s
+
   if args.profile:
     from tabulate import tabulate
-    profile = decode_profile(viz.get_profile(viz.load_pickle(args.profile_path, default=[])))
+    profile = decode_profile(viz.get_profile(profile_data:=viz.load_pickle(args.profile_path, default=[])))
+    viz.load_amd_counters(viz.ctxs, profile_data)
+    counters = {f'{c["name"]} SQTT {s["name"]}': s["data"] for c in viz.ctxs if c["name"].startswith("Exec") for s in c["steps"]
+                if s["name"].startswith("PKTS")}
+    if args.device is None:
+      print("Select a device:")
+      for k in (*profile["layout"], *counters):
+        print(f"  {format_colored(k)}")
+      sys.exit(0)
+
+    # ** SQTT printer
+    if args.device is not None and (sqtt_data:=next((v for k,v in counters.items() if ansistrip(k) == args.device), None)) is not None:
+      assert args.limit > 1, f"SQTT limit must be greater than 1, got {args.limit}"
+      sqtt_events, has_more = viz.sqtt_timeline(*sqtt_data, max_pkts=args.offset+args.limit)
+      sqtt_pkts = [e for e in sqtt_events if type(e).__name__ == "ProfileRangeEvent"]
+      pc_map = next((e.arg for e in sqtt_events if type(e).__name__ == "ProfilePointEvent" and e.key == 'pcMap'), None)
+      if pc_map is None:
+        print(f"No SQTT instruction trace data for {args.device}")
+        sys.exit(0)
+      # modern terminals support 24-bit color
+      def hex_colored(st:str, color:str) -> str: return f"\x1b[38;2;{int(color[1:3],16)};{int(color[3:5],16)};{int(color[5:7],16)}m{st}\x1b[0m"
+      WAVE_COLORS = ((('VALU', 'VINTERP'), '#ffffc0'), (('SALU',), '#cef263'), (('VMEM',), '#b2b7c9'), (('LOAD', 'SMEM'), '#ffc0c0'),
+                     (('STORE',), '#4fa3cc'), (('IMMEDIATE',), '#f3b44a'), (('BARRIER',), '#d00000'), (('LDS',), '#9fb4a6'), (('JUMP',), '#ffb703'),
+                     (('JUMP_NO',), '#fb8500'), (('MESSAGE',), '#90dbf4'), (('WAVERDY',), '#1a2a2a'))
+      print(f"{'Clk':<12} {'Unit':<20} {'Op':<22} {'Dur':<4} {'Info'}")
+      print("-" * 90)
+      # start from the first packet in trace, prepare packet indexes and map dispatches
+      pkt_idxs:dict[str, itertools.count] = {}
+      dispatch_to_pc:dict[str, int] = {}
+      for e in sqtt_pkts[:-args.limit]:
+        idx = next(pkt_idxs.setdefault(e.device, itertools.count()))
+        if e.name.ret is not None and e.name.ret.startswith("PC:"): dispatch_to_pc[f"{e.device}-{idx}"] = int(e.name.ret.replace("PC:", ""))
+      # start printing from the offset point
+      for e in sqtt_pkts[-args.limit:]:
+        op_name, info = e.name.display_name, e.name.ret or ""
+        color = next((c for p, c in WAVE_COLORS if any(x in op_name for x in p)), None)
+        op_str = hex_colored(op_name, color) if color and not args.no_color else op_name
+        phase, pc = None, None
+        idx = next(pkt_idxs.setdefault(e.device, itertools.count()))
+        if info.startswith("PC:"):
+          dispatch_to_pc[f"{e.device}-{idx}"] = pc = int(info.replace("PC:", ""))
+          phase = "DISPATCH"
+        if info.startswith("LINK:"): phase, pc = "EXEC", dispatch_to_pc[info.replace("LINK:", "")]
+        if pc and phase: info = f"{phase:<8} 0x{pc:05x} {pc_map[pc]}"
+        print(f"{int(e.st):<12} {e.device:<20} {op_str}{' '*(22-ansilen(op_str))} {int(e.en-e.st):<4} {info}")
+      # note: we only print the important packets and skip the rest
+      if has_more: print(f"Selected packets {args.offset:,}-{args.offset + args.limit:,}. Use --offset and --limit to see others")
+      sys.exit(0)
+
+    # ** Profiler printer
     agg, total, n = {}, 0, 0
-    if args.device is None: print("Select a device:")
     for k,v in profile["layout"].items():
       if not optional_eq({"name":k}, args.device): continue
-      print(f"  {k}")
+      print(f"  {format_colored(k)}")
       if args.device is None: continue
       for e in v.get("events", []):
         et = e["dur"]*1e-6
@@ -100,7 +151,7 @@ if __name__ == "__main__":
           if optional_eq(e, args.kernel) and n < 10:
             ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
             name = e["name"]+(" " * (46 - ansilen(e["name"])))
-            print(f"{name} {ptm}/{(et or 0)*1e3:9.2f}ms  "+e['fmt'].replace('\n', ' | ')+"  ")
+            print(f"{name} {ptm}/{(et or 0)*1e3:9.2f}ms  "+e.get('fmt', '').replace('\n', ' | ')+"  ")
             n += 1
         else:
           a = agg.setdefault(e["name"], [0.0, 0])
@@ -109,14 +160,15 @@ if __name__ == "__main__":
           total += et
     if agg and total > 0:
       items = sorted(agg.items(), key=lambda kv:kv[1][0], reverse=True)
-      sel = items if args.top == -1 else items[:args.top]
+      sel = items if args.limit == -1 else items[args.offset:args.offset+args.limit]
       table = [[name, time_to_str(t, w=9), c, f"{(t/total*100.0):.2f}%"] for name,(t,c) in sel]
-      if args.top != -1 and (other:=items[len(sel):]):
+      if args.limit != -1 and (other:=items[len(sel):]):
         other_t = total-sum(t for _, (t, _) in sel)
         table.append([f"Other ({len(other)} unique)", time_to_str(other_t, w=9), sum(c for _,(_,c) in other), f"{other_t/total*100.0:.2f}%"])
       print(tabulate(table, headers=["name", "total", "count", "pct"], tablefmt="github"))
     sys.exit(0)
 
+  # ** Graph rewrites printer
   for k in viz.ctxs:
     if not optional_eq(k, args.kernel): continue
     print(k["name"])

test/amd/hw/test_ds.py

@@ -117,6 +117,58 @@ class TestDS2AddrMore(unittest.TestCase):
     self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
     self.assertEqual(st.vgpr[0][4], 0xDEADBEEF, "v4 should be untouched")
 
+  def test_ds_load_2addr_b64_addr_overlaps_vdst(self):
+    """DS_LOAD_2ADDR_B64 where addr register overlaps vdst range.
+
+    Hardware reads the address before writing any results, so addr=v[4]
+    with vdst=v[4:7] must load all 4 dwords using the original v[4] value.
+    """
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 0xAAAAAAAA),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
+      s_mov_b32(s[2], 0xBBBBBBBB),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=4),
+      s_mov_b32(s[2], 0xCCCCCCCC),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=8),
+      s_mov_b32(s[2], 0xDDDDDDDD),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=12),
+      s_waitcnt(lgkmcnt=0),
+      # addr=v[4] overlaps vdst=v[4:7]
+      v_mov_b32_e32(v[4], 0),
+      DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[4], vdst=v[4:7], offset0=0, offset1=1),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA, "v4 = LDS[0:4]")
+    self.assertEqual(st.vgpr[0][5], 0xBBBBBBBB, "v5 = LDS[4:8]")
+    self.assertEqual(st.vgpr[0][6], 0xCCCCCCCC, "v6 = LDS[8:12]")
+    self.assertEqual(st.vgpr[0][7], 0xDDDDDDDD, "v7 = LDS[12:16]")
+
+  def test_ds_load_2addr_b32_addr_overlaps_vdst(self):
+    """DS_LOAD_2ADDR_B32 where addr register overlaps vdst range."""
+    instructions = [
+      v_mov_b32_e32(v[10], 0),
+      s_mov_b32(s[2], 0xAAAAAAAA),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
+      s_mov_b32(s[2], 0xBBBBBBBB),
+      v_mov_b32_e32(v[0], s[2]),
+      ds_store_b32(addr=v[10], data0=v[0], offset0=4),
+      s_waitcnt(lgkmcnt=0),
+      # addr=v[2] overlaps vdst=v[2:3]
+      v_mov_b32_e32(v[2], 0),
+      DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[2], vdst=v[2:3], offset0=0, offset1=1),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA, "v2 = LDS[0:4]")
+    self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB, "v3 = LDS[4:8]")
+
   def test_ds_load_b64_no_overwrite(self):
     """DS_LOAD_B64 should only write 2 VGPRs."""
     instructions = [
@@ -759,6 +811,47 @@ class TestDsPermute(unittest.TestCase):
     # Lane 0 receives data (highest numbered active lane wins)
     self.assertEqual(st.vgpr[0][2], 0x11111111)
 
+  def test_ds_bpermute_b32_xor_swap(self):
+    """DS_BPERMUTE_B32 with XOR-1 pattern — each lane reads from lane^1.
+
+    This is the pattern used by warp_shfl_xor in flash attention for reduce_max/reduce_sum.
+    Each lane has a unique value (lane_id + 100), and reads from the adjacent lane.
+    """
+    instructions = [
+      # v[0] = (lane_id ^ 1) * 4  (byte offset for bpermute)
+      v_xor_b32_e32(v[0], 1, v[255]),
+      v_lshlrev_b32_e32(v[0], 2, v[0]),
+      # v[1] = lane_id + 100 (unique per-lane value)
+      s_mov_b32(s[0], 100),
+      v_add_nc_u32_e32(v[1], s[0], v[255]),
+      # ds_bpermute: v[2] = v[1] from lane (lane_id ^ 1)
+      ds_bpermute_b32(vdst=v[2], addr=v[0], data0=v[1]),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=32)
+    for lane in range(32):
+      src_lane = lane ^ 1
+      expected = src_lane + 100
+      self.assertEqual(st.vgpr[lane][2], expected, f"lane {lane}: expected v[1] from lane {src_lane} = {expected}, got {st.vgpr[lane][2]}")
+
+
+class TestDSSubDword(unittest.TestCase):
+  """Tests for sub-dword DS operations (ds_store_b16, ds_store_b16_d16_hi)."""
+
+  def test_ds_store_b16_and_d16_hi(self):
+    """DS_STORE_B16 stores low 16 bits, DS_STORE_B16_D16_HI stores high 16 bits to adjacent LDS half-words."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0),
+      v_mov_b32_e32(v[1], 0xBEEF1234),
+      DS(DSOp.DS_STORE_B16, addr=v[0], data0=v[1], offset0=0),
+      DS(DSOp.DS_STORE_B16_D16_HI, addr=v[0], data0=v[1], offset0=2),
+      s_waitcnt(lgkmcnt=0),
+      ds_load_b32(vdst=v[2], addr=v[0], offset0=0),
+      s_waitcnt(lgkmcnt=0),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xBEEF1234, "lo=0x1234 at byte 0, hi=0xBEEF at byte 2")
+
 
 class TestDSLargeOffset(unittest.TestCase):
   """Tests for DS instructions with offsets > 255 (offset1 > 0).

test/amd/hw/test_vop3p.py

@@ -408,6 +408,23 @@ class TestVOP3P(unittest.TestCase):
     self.assertEqual(lo, 0x0005, f"lo: expected 0x0005, got 0x{lo:04x}")
     self.assertEqual(hi, 0x4003, f"hi: expected 0x4003, got 0x{hi:04x}")
 
+  def test_v_pk_add_u16_literal_constant(self):
+    """V_PK_ADD_U16 with a literal constant (value > 64, requires VOP3P_LIT encoding).
+    Regression test: VOP3P literal constants were not passed to rsrc_dyn, so literal src read as 0.
+    """
+    instructions = [
+      s_mov_b32(s[0], 0x1C001C00),  # packed u16: hi=0x1C00, lo=0x1C00 (f16 for 2^-8)
+      v_mov_b32_e32(v[0], s[0]),
+      v_pk_add_u16(v[1], 0x2000, v[0], opsel_hi=2, opsel_hi2=1),  # add 0x2000 bias to both halves
+    ]
+    st = run_program(instructions, n_lanes=1)
+    result = st.vgpr[0][1]
+    lo = result & 0xffff
+    hi = (result >> 16) & 0xffff
+    # lo = 0x1C00 + 0x2000 = 0x3C00 (f16 1.0), hi = 0x1C00 + 0x2000 = 0x3C00 (f16 1.0)
+    self.assertEqual(lo, 0x3C00, f"lo: expected 0x3C00, got 0x{lo:04x}")
+    self.assertEqual(hi, 0x3C00, f"hi: expected 0x3C00, got 0x{hi:04x}")
+
 
 class TestWMMAF16(unittest.TestCase):
   """Tests for WMMA F16 output variant (V_WMMA_F16_16X16X16_F16).

test/amd/hw/test_vopd.py

@@ -138,6 +138,47 @@ class TestVOPDLiterals(unittest.TestCase):
     self.assertAlmostEqual(i2f(st.vgpr[0][2]), 13.0, places=5, msg="fma(2.0, 5.0, 3.0) should be 13.0")
 
 
+class TestVOPDDot2Acc(unittest.TestCase):
+  """Tests for V_DUAL_DOT2ACC_F32_F16 - packed f16 dot product accumulate."""
+
+  def test_vopd_dot2acc_f32_f16_basic(self):
+    """V_DUAL_DOT2ACC_F32_F16: D += lo(S0)*lo(S1) + hi(S0)*hi(S1).
+
+    S0 = pack(1.0h, 2.0h), S1 = pack(3.0h, 4.0h), D = 10.0f
+    result = 10.0 + 1.0*3.0 + 2.0*4.0 = 10.0 + 3.0 + 8.0 = 21.0
+    """
+    from test.amd.hw.helpers import f2i, i2f, f32_to_f16
+    pk_s0 = f32_to_f16(1.0) | (f32_to_f16(2.0) << 16)  # lo=1.0h, hi=2.0h
+    pk_s1 = f32_to_f16(3.0) | (f32_to_f16(4.0) << 16)  # lo=3.0h, hi=4.0h
+    instructions = [
+      v_mov_b32_e32(v[0], pk_s0),
+      v_mov_b32_e32(v[1], pk_s1),
+      v_mov_b32_e32(v[3], f2i(10.0)),  # accumulator in v[3] (vdsty with vdstx=v[4])
+      # X: v[4] = MOV v[0] (don't care), Y: v[3] += dot2(v[0], v[1])
+      VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_DOT2ACC_F32_F16, v[4], v[3], v[0], v[0], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertAlmostEqual(i2f(st.vgpr[0][3]), 21.0, places=2, msg="10.0 + 1.0*3.0 + 2.0*4.0 = 21.0")
+
+  def test_vopd_dot2acc_f32_f16_zero_accum(self):
+    """V_DUAL_DOT2ACC_F32_F16 with zero accumulator — pure dot product.
+
+    S0 = pack(0.5h, -1.0h), S1 = pack(2.0h, 3.0h), D = 0.0f
+    result = 0.0 + 0.5*2.0 + (-1.0)*3.0 = 1.0 - 3.0 = -2.0
+    """
+    from test.amd.hw.helpers import f2i, i2f, f32_to_f16
+    pk_s0 = f32_to_f16(0.5) | (f32_to_f16(-1.0) << 16)
+    pk_s1 = f32_to_f16(2.0) | (f32_to_f16(3.0) << 16)
+    instructions = [
+      v_mov_b32_e32(v[0], pk_s0),
+      v_mov_b32_e32(v[1], pk_s1),
+      v_mov_b32_e32(v[3], f2i(0.0)),  # zero accumulator in v[3]
+      VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_DOT2ACC_F32_F16, v[4], v[3], v[0], v[0], v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertAlmostEqual(i2f(st.vgpr[0][3]), -2.0, places=2, msg="0.5*2.0 + (-1.0)*3.0 = -2.0")
+
+
 class TestVOPDMultilane(unittest.TestCase):
   """Tests for VOPD with multiple lanes."""
 

test/amd/test_custom_kernel.py

@@ -3,8 +3,10 @@ import functools
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
+from tinygrad.dtype import AddrSpace
 from tinygrad.runtime.autogen.amd.rdna3.ins import *
-from tinygrad.runtime.autogen.amd.rdna4.ins import s_barrier_wait, s_barrier_signal
+import tinygrad.runtime.autogen.amd.rdna3.ins as r3
+import tinygrad.runtime.autogen.amd.rdna4.ins as r4
 from tinygrad.renderer.amd.dsl import s, v
 from test.amd.helpers import TARGET_TO_ARCH
 
@@ -53,12 +55,46 @@ def custom_wave_sync(A:UOp, arch:str) -> UOp:
   insts = []
   for _ in range(4):
     insts.append(s_sleep(4))
-    insts += [s_barrier()] if arch == "rdna3" else [s_barrier_signal(), s_barrier_wait()]
+    insts += [s_barrier()] if arch == "rdna3" else [r4.s_barrier_signal(), r4.s_barrier_wait()]
     insts += [s_nop(0)]*4
   insts.append(s_endpgm())
   sink = UOp.sink(A.base, threads, wg, arg=KernelInfo("custom_wave_sync"))
   return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
 
+def custom_lds_sync(A:UOp, arch:str) -> UOp:
+  A = A.flatten()
+  num_threads = A.shape[0]
+  threads = UOp.special(num_threads, "lidx0")
+  wg = UOp.special(1, "gidx0")
+  lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=512, addrspace=AddrSpace.LOCAL), (), 'lds')  # 128 * 4 bytes
+  isa = r4 if arch == "rdna4" else r3
+  wait_kmcnt = [isa.s_wait_kmcnt(simm16=0)] if arch == "rdna4" else [isa.s_waitcnt(lgkmcnt=0)]
+  wait_dscnt = [isa.s_wait_dscnt(simm16=0)] if arch == "rdna4" else [isa.s_waitcnt(lgkmcnt=0)]
+  barrier = [isa.s_barrier_signal(ssrc0=-1), isa.s_barrier_wait(simm16=-1)] if arch == "rdna4" else [isa.s_barrier()]
+  global_store = [isa.global_store_b32(vaddr=v[6:7], saddr=s[0:1], vsrc=v[5])] if arch == "rdna4" \
+      else [isa.global_store_b32(addr=v[6], data=v[5], saddr=s[0:1])]
+  insts = [
+    isa.s_load_b64(s[0:1], s[0:1], soffset=NULL),
+    *wait_kmcnt,
+    isa.v_lshlrev_b32_e32(v[1], 2, v[0]),
+    # lds[thread_idx] = thread_idx
+    isa.ds_store_b32(addr=v[1], data0=v[0]),
+    *wait_dscnt,
+    *barrier,
+    # out[threaed_idx] = thread_idx == num_threads ? -1 : lds[thread_idx + 1]
+    isa.v_add_nc_u32_e32(v[2], 4, v[1]),
+    isa.v_cmp_gt_u32_e32(num_threads-1, v[0]),
+    isa.ds_load_b32(vdst=v[3], addr=v[2]),
+    *wait_dscnt,
+    isa.v_mov_b32_e32(v[4], -1),
+    isa.v_cndmask_b32_e32(v[5], v[4], v[3]),
+    isa.v_lshlrev_b32_e32(v[6], 2, v[0]),
+    *global_store,
+    isa.s_endpgm(),
+  ]
+  sink = UOp.sink(A.base, lds, threads, wg, arg=KernelInfo("custom_lds_sync"))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+
 @unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
 class TestCustomKernel(unittest.TestCase):
   def setUp(self): self.arch = TARGET_TO_ARCH[Device["AMD"].arch]
@@ -83,9 +119,14 @@ class TestCustomKernel(unittest.TestCase):
       ei.run({"var":i})
       self.assertTrue((a.numpy() == 1+i).all())
 
-  def test_wave_sync(self):
-    if self.arch not in {"rdna3", "rdna4"}: self.skipTest("only rdna3 or rdna4")
-    Tensor.empty(1).custom_kernel(fxn=functools.partial(custom_wave_sync, arch=self.arch))[0].realize()
+  def test_lds_sync(self):
+    if self.arch not in ("rdna3", "rdna4"): self.skipTest("only rdna3/rdna4")
+    a = Tensor.empty(128, dtype=dtypes.int32).contiguous().realize()
+    a = Tensor.custom_kernel(a, fxn=functools.partial(custom_lds_sync, arch=self.arch))[0]
+    a.realize()
+    ref = Tensor.arange(1, 129, dtype=dtypes.int32)
+    ref[127] = -1
+    self.assertListEqual(a.tolist(), ref.tolist())
 
 if __name__ == "__main__":
   unittest.main()

test/amd/test_sqtt_examples.py

@@ -10,18 +10,11 @@ from tinygrad.runtime.autogen.amd.rdna3.ins import SOPP
 from tinygrad.runtime.autogen.amd.rdna3.enum import SOPPOp
 from tinygrad.renderer.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_RDNA4, WAVEEND, INST, INST_RDNA4, VALUINST,
                                      IMMEDIATE, IMMEDIATE_MASK, PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4, PACKET_TYPES_CDNA, CDNA_WAVESTART,
-                                     InstOp, InstOpRDNA4, print_packets, CDNA_WAVEEND, CDNA_INST)
+                                     print_packets, CDNA_WAVEEND, CDNA_INST)
 from test.amd.helpers import TARGET_TO_ARCH
 
 import tinygrad
 EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"
-# INST ops for non-traced SIMDs (excluded from instruction count)
-OTHER_SIMD_OPS = {InstOp.OTHER_LDS_LOAD, InstOp.OTHER_LDS_STORE, InstOp.OTHER_LDS_STORE_64, InstOp.OTHER_LDS_STORE_128,
-                  InstOp.OTHER_FLAT_LOAD, InstOp.OTHER_FLAT_STORE, InstOp.OTHER_FLAT_STORE_64, InstOp.OTHER_FLAT_STORE_96,
-                  InstOp.OTHER_FLAT_STORE_128, InstOp.OTHER_GLOBAL_LOAD, InstOp.OTHER_GLOBAL_LOAD_VADDR,
-                  InstOp.OTHER_GLOBAL_STORE_64, InstOp.OTHER_GLOBAL_STORE_96, InstOp.OTHER_GLOBAL_STORE_128,
-                  InstOp.OTHER_GLOBAL_STORE_VADDR_128}
-OTHER_SIMD_OPS_RDNA4 = {InstOpRDNA4.OTHER_VMEM, InstOpRDNA4.OTHER_VMEM_5, InstOpRDNA4.OTHER_LDS_1, InstOpRDNA4.OTHER_LDS_2}
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # ROCPROF DECODER
@@ -183,11 +176,18 @@ class SQTTExamplesTestBase(unittest.TestCase):
         our_waves: list[tuple[int, int]] = []
         for event in events:
           wave_starts: dict[tuple[int, int, int], int] = {}
+          first_timestamp:int|None = None
           for p in decode(event.blob):
+            if first_timestamp is None: first_timestamp = p._time
             if isinstance(p, (WAVESTART, CDNA_WAVESTART, WAVESTART_RDNA4)): wave_starts[(p.wave, p.simd, p.cu)] = p._time
             elif isinstance(p, (WAVEEND, CDNA_WAVEEND)) and (key := (p.wave, p.simd, p.cu)) in wave_starts:
               our_waves.append((wave_starts[key], p._time))
-        self.assertEqual(sorted(our_waves), sorted(roc_waves), f"wave times mismatch in {name}")
+          for st in wave_starts.values():
+            self.assertGreater(st, first_timestamp, "wave start must be after the first packet")
+        # rocprof fails non deterministically and gives inaccurate timestamps.
+        #self.assertEqual(sorted(our_waves), sorted(roc_waves), f"wave times mismatch in {name}")
+        for st, et in our_waves:
+          self.assertGreater(et, st, "wave end must be after start")
 
   def test_rocprof_inst_times_match(self):
     """Instruction times must match rocprof exactly (excluding s_endpgm)."""
@@ -200,8 +200,8 @@ class SQTTExamplesTestBase(unittest.TestCase):
         our_insts: list[int] = []
         for event in events:
           for p in decode(event.blob):
-            if isinstance(p, INST) and p.op not in OTHER_SIMD_OPS: our_insts.append(p._time)
-            elif isinstance(p, INST_RDNA4) and p.op not in OTHER_SIMD_OPS_RDNA4: our_insts.append(p._time)
+            # INST ops for non-traced SIMDs (excluded from instruction count)
+            if isinstance(p, (INST, INST_RDNA4)) and not p.op.name.startswith("OTHER_"): our_insts.append(p._time)
             elif isinstance(p, VALUINST): our_insts.append(p._time)
             elif isinstance(p, IMMEDIATE): our_insts.append(p._time)
             elif isinstance(p, IMMEDIATE_MASK):

test/amd/test_sqttmap.py

@@ -79,7 +79,7 @@ class TestSQTTMapBase(unittest.TestCase):
         if (p:=kern_events.get(event.kern)) is None: continue
         with self.subTest(example=name, kern=event.kern):
           # skip if there's no SQTT frequency data
-          if not (timeline:=sqtt_timeline(event.blob, p.lib, target)): continue
+          if not (timeline:=sqtt_timeline(event.blob, p.lib, target)[0]): continue
           if not (frequency:=[e.key for e in timeline if type(e).__name__ == "ProfilePointEvent" and e.name == "freq_hz"]): continue
           mean = sum(frequency) / len(frequency)
           variance = sum((v - mean) ** 2 for v in frequency) / len(frequency)
@@ -92,7 +92,8 @@ class TestSQTTMapBase(unittest.TestCase):
               if "ALT" not in e.name.display_name: execs += 1
             elif "WAVE" in e.device:
               # sopk/immediates don't get ALU/MEM EXEC
-              if e.name.display_name not in {"IMMEDIATE", "IMMEDIATE_MASK", "JUMP", "JUMP_NO", "MESSAGE", "BARRIER", "BARRIER_SIGNAL"}: insts += 1
+              if e.name.display_name not in {"IMMEDIATE", "IMMEDIATE_MASK", "JUMP", "JUMP_NO", "MESSAGE", "BARRIER", "BARRIER_SIGNAL",
+                                             "WAVEEND", "WAVERDY"}: insts += 1
             else: raise Exception(f"timeline row must be INST or EXEC, got {e.device}")
           self.assertEqual(execs, insts)
 
@@ -100,7 +101,7 @@ class TestSQTTMapBase(unittest.TestCase):
     for name, (events, kern_events, target) in self.examples.items():
       for event in events:
         wave_barriers = {}
-        for e in sqtt_timeline(event.blob, kern_events[event.kern].lib, target):
+        for e in sqtt_timeline(event.blob, kern_events[event.kern].lib, target)[0]:
           if type(e).__name__ == "ProfileRangeEvent" and e.name.display_name == "BARRIER": wave_barriers.setdefault(e.device, []).append(e)
         if not wave_barriers: continue
         for row, events in wave_barriers.items():

test/backend/test_asm_gemm.py

@@ -81,16 +81,37 @@ class TestGemm(unittest.TestCase):
   @needs_second_gpu
   def test_gemm_k_sharded_3d(self): verify_asm_gemm_k_sharded_3d(1, 64, 32, 2*64, gpus=2)
 
-# uses the Asm GEMM on CDNA4 only for speed reasons
-class TestGemmLarge(unittest.TestCase):
+# uses the smallest size for the cdna assembly gemm
+class TestAsmGEMM(unittest.TestCase):
   def setUp(self):
     if not is_cdna4():
+      self.skipTest("assembly gemm is only for cdna4")
+
+  def test_tiny(self): verify_asm_gemm(1, 256, 256, 64)
+
+  def test_verify_with_numpy(self):
+    import numpy as np
+    M, N, K = 256, 256, 64
+    rng = np.random.default_rng(0)
+    a_np = (rng.random((M, K), dtype=np.float32) - 0.5).astype(np.half)
+    b_np = (rng.random((K, N), dtype=np.float32) - 0.5).astype(np.half)
+    c_np = a_np @ b_np
+    a, b = Tensor(a_np), Tensor(b_np)
+    c = asm_gemm(a, b)
+    c.realize()
+    # no validation on the NULL device
+    if a.device.startswith("NULL"): return None
+    np.testing.assert_allclose(c.numpy(), c_np, atol=2e-3, rtol=5e-2)
+
+# test the Asm GEMM with Llama shapes, only run on the real machine for speed
+class TestGemmLlama(unittest.TestCase):
+  def setUp(self):
+    if not is_cdna4() or getenv("MOCKGPU"):
       self.skipTest("very slow on non mi350x")
 
   @Context(ASM_GEMM=1)
   def test_empty(self): (Tensor.empty(N:=getenv("N", 4096), N, dtype=dtypes.half)@Tensor.empty(N, N, dtype=dtypes.half)).realize()
 
-  def test_tiny(self): verify_asm_gemm(1, 256, 256, 64)
   def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 4096), N, N, dtype=dtypes.half)
   def test_gemm(self): verify_asm_gemm(1, 8192, 4096, 14336)
   def test_gemm_batched(self): verify_asm_gemm(2, 8192, 4096, 4096)

test/backend/test_custom_kernel.py

@@ -283,5 +283,30 @@ class TestCustomKernel(unittest.TestCase):
     self.assertIsNotNone(custom_idx, "custom_addmul kernel not found in schedule")
     self.assertEqual(custom_idx, 3, f"custom_addmul should be at index 3, got {custom_idx}")
 
+  def test_anonymous_buffers_in_function(self):
+    """Test that custom kernels with anonymous output buffers work inside @function."""
+    a = Tensor.full((4, 4), 3.).contiguous()
+    b = Tensor.full((4, 4), 2.).contiguous()
+    Tensor.realize(a, b)
+
+    def custom_add_with_tmp(o1:UOp, o2:UOp, A:UOp, B:UOp) -> UOp:
+      o1,o2,A,B = o1.flatten(), o2.flatten(), A.flatten(), B.flatten()
+      i = UOp.range(o1.size, 0)
+      store_o1 = o1[i].store(A[i]+B[i])
+      store_o2 = o2[i].store(A[i]+B[i]+2)
+      return UOp.group(store_o1, store_o2).end(i).sink(arg=KernelInfo(name=f"add_with_tmp_{o1.size}")).simplify()
+
+    from tinygrad import function
+    @function(precompile=True)
+    def run(x:Tensor, w:Tensor) -> Tensor:
+      out = Tensor.invalid(*x.shape, dtype=x.dtype)
+      tmp = Tensor.invalid(*x.shape, dtype=x.dtype)
+      out, tmp = Tensor.custom_kernel(out, tmp, x, w, fxn=custom_add_with_tmp)[:2]
+      return out+tmp
+
+    result = run(a, b).flatten().tolist()
+    expected = (3+2)*2+2
+    assert all(x == expected for x in result), f"expected all {expected}, got {result}"
+
 if __name__ == '__main__':
   unittest.main()

test/backend/test_dtype.py

@@ -18,11 +18,10 @@ settings.register_profile("my_profile", max_examples=200, deadline=None, derando
 settings.load_profile("my_profile")
 
 def get_available_cast_dtypes(dtype: DType) -> List[DType]:
-  # dont cast internal dtypes
-  dts = [v for k, v in DTYPES_DICT.items() if v != dtype and is_dtype_supported(v) and not k.startswith("_")]
-  if not is_dtype_supported(dtype) or dtypes.long in EMULATED_DTYPES.tolist(dtypes):
-    if dtype in (dtypes.long, dtypes.ulong): return [dt for dt in dts if dt != dtypes.double] # can't bitcast with no 64-bit support
-    else: return []
+  dts = [v for k, v in DTYPES_DICT.items() if v != dtype and is_dtype_supported(v) or v in dtypes.fp8s+(dtypes.half,dtypes.bfloat16,dtypes.long)]
+  if dtype in (dtypes.long, dtypes.ulong) and (not is_dtype_supported(dtype) or dtypes.long in EMULATED_DTYPES.tolist(dtypes)):
+    return [dt for dt in dts if dt != dtypes.double] # can't bitcast with no 64-bit support
+  if not is_dtype_supported(dtype) and dtype not in dtypes.fp8s+(dtypes.half,dtypes.bfloat16): return []
   return dts
 
 def _to_torch_storage_type(dtype:DType):
@@ -60,10 +59,8 @@ class TestDType(unittest.TestCase):
   DATA: Any = None
   @classmethod
   def setUpClass(cls):
-    if not cls.DTYPE or not is_dtype_supported(cls.DTYPE): raise unittest.SkipTest("dtype not supported")
-    cls.DATA = rand_for_dtype(cls.DTYPE, 10)
-  def setUp(self):
-    if self.DTYPE is None: raise unittest.SkipTest("base class")
+    if cls.DTYPE is None: raise unittest.SkipTest("base class")
+    cls.DATA = rand_for_dtype(cls.DTYPE, 0x10, allow_subnormal=is_dtype_supported(cls.DTYPE))
 
   def test_to_np(self):
     _test_to_np(Tensor(self.DATA, dtype=self.DTYPE), _to_np_dtype(self.DTYPE), np.array(self.DATA, dtype=_to_np_dtype(self.DTYPE)))
@@ -132,7 +129,6 @@ class TestDType(unittest.TestCase):
 
 def _test_ops(a_dtype:DType, b_dtype:DType, target_dtype=None):
   target_dtype = target_dtype or least_upper_dtype(a_dtype, b_dtype)
-  if not is_dtype_supported(a_dtype) or not is_dtype_supported(b_dtype) or not is_dtype_supported(target_dtype): return
   if a_dtype == dtypes.bool or b_dtype == dtypes.bool: return
   _assert_eq(Tensor([1,2,3,4], dtype=a_dtype)+Tensor([1,2,3,4], dtype=b_dtype), target_dtype, [2,4,6,8])
   _assert_eq((Tensor([1], dtype=a_dtype).cast(b_dtype)+Tensor([1], dtype=a_dtype).cast(b_dtype)).cast(a_dtype), a_dtype, [2])
@@ -195,7 +191,6 @@ class TestFp8sConversions(unittest.TestCase):
   def test_fp8e5m2fnuz_to_float(self, x):
     np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2fnuz).float().item())
 
-@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "bfloat16 not supported")
 class TestBFloat16(unittest.TestCase):
   def test_bf16_creation_numpy(self):
     data = [-1, 1, 2]
@@ -215,7 +210,6 @@ class TestBFloat16(unittest.TestCase):
     assert t.dtype == dtypes.bfloat16
     np.testing.assert_allclose(t.numpy(), np.eye(3))
 
-@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "bfloat16 not supported")
 class TestBFloat16DType(unittest.TestCase):
   def test_bf16_to_float(self):
     _test_cast(Tensor([100000], dtype=dtypes.bfloat16), dtypes.float32)
@@ -229,7 +223,6 @@ class TestBFloat16DType(unittest.TestCase):
     back = t.cast(dtypes.float32)
     assert tuple(back.numpy().tolist()) == (9984., -1, -1000, -9984, 20)
 
-@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16) and is_dtype_supported(dtypes.float16), "bfloat16 or float16 not supported")
 class TestBFloat16DTypeCast(unittest.TestCase):
   def test_f16_to_bf16_conversion(self):
     original_tensor = Tensor([1.0, 2.0, 3.0], dtype=dtypes.float16)
@@ -278,6 +271,7 @@ class TestFloatDType(TestDType):
     _test_op(lambda: Tensor([-0.9, -0.3, 1.2], dtype=dtypes.float32).cast(dtypes.uint32), dtypes.uint32,
              [0, 0, 1])
 
+@unittest.skipUnless(is_dtype_supported(dtypes.double), f"no double on {Device.DEFAULT}")
 class TestDoubleDType(TestDType):
   DTYPE = dtypes.double
   @unittest.skipIf((CI and Device.DEFAULT in {"CUDA", "NV"}) or \

test/backend/test_dtype_alu.py

@@ -81,8 +81,8 @@ def universal_test_unary(a, dtype, op):
   if op[0] == Tensor.cos and abs(a) > 30: return
   if op[0] == Tensor.log and a <= 0: return
   if dtype in dtypes.fp8s:
-    # normals are zero
-    if dtype in EMULATED_DTYPES.tolist(dtypes) and abs(ta.numpy().item()) < 0.015625: return
+    # denormals are zero
+    if dtype in EMULATED_DTYPES.tolist(dtypes) or not is_dtype_supported(dtype) and abs(ta.numpy().item()) < 0.015625: return
     tensor_value = fp8_to_float(op[0](ta.realize()).bitcast(dtypes.uint8).item(), dtype)
     numpy_value = truncate[dtype](v:=op[1](ta.numpy()).item())
     # cuda cast f32 inf to f8 MAX, amd cast it to nan(E4M3)/inf(E5M2)

test/backend/test_image_dtype.py

@@ -1,186 +1,17 @@
 import unittest
 import numpy as np
 from tinygrad import Device, dtypes, Tensor, Context
-from tinygrad.device import LRUAllocator, is_dtype_supported
-from tinygrad.dtype import ImageDType
-from tinygrad.helpers import prod, unwrap
+from tinygrad.helpers import unwrap
 
 IMAGE_SUPPORTED_DEVICES = ("QCOM", "CL")
 
-@unittest.skipUnless(Device.DEFAULT in IMAGE_SUPPORTED_DEVICES, "Images not supported")
-class TestImageCopy(unittest.TestCase):
-  def test_image_copyout_1x8(self, img_type=dtypes.imagef):
-    it = Tensor.arange(32).cast(img_type((1,8,4))).realize()
-    buf = it.uop.buffer
-    out = buf.as_memoryview()
-    np.testing.assert_equal(out.cast(it.dtype.fmt).tolist(), np.arange(32))
-
-  @unittest.skipUnless(is_dtype_supported(dtypes.half, device="PYTHON"), "need half")
-  def test_imageh_copyout_1x8(self): self.test_image_copyout_1x8(img_type=dtypes.imageh)
-
-  def test_image_numpy_1x8(self, img_type=dtypes.imagef):
-    it = Tensor.arange(32).cast(img_type((1,8,4))).realize()
-    np.testing.assert_equal(it.numpy(), np.arange(32))
-  def test_imageh_numpy_1x8(self): self.test_image_numpy_1x8(img_type=dtypes.imageh)
-
-  def test_image_copyout_2x4(self):
-    it = Tensor.arange(2*4*4).cast(dtypes.imagef((2,4,4))).realize()
-    buf = it.uop.buffer
-    out = buf.as_memoryview()
-    np.testing.assert_equal(out.cast('f').tolist(), np.arange(2*4*4))
-
-  def test_image_roundtrip(self):
-    sz = (4,2,4)
-    it = Tensor.rand(prod(sz)).cast(dtypes.imagef(sz)).realize()
-    buf = it.uop.buffer
-    out = buf.as_memoryview()
-
-    it2 = Tensor.rand(prod(sz)).cast(dtypes.imagef(sz)).realize()
-    buf2 = it2.uop.buffer
-    buf2.copyin(out)
-
-    assert (it == it2).sum().item() == prod(sz)
-
 @unittest.skipUnless(Device.DEFAULT in IMAGE_SUPPORTED_DEVICES, "Images not supported")
 class TestImageDType(unittest.TestCase):
-  def test_image_pitch(self):
-    def __validate(imgdt, expected_pitch):
-      assert imgdt.pitch == expected_pitch, f"Failed pitch for image: {imgdt}. Got 0x{imgdt.pitch:X}, expected 0x{expected_pitch:X}"
-
-    # Match opencl pitches for perf
-    __validate(dtypes.imageh((1, 201)), 0x680)
-    __validate(dtypes.imageh((16, 216)), 0x700)
-    __validate(dtypes.imageh((16, 9)), 0x80)
-    __validate(dtypes.imageh((48, 64)), 0x200)
-    __validate(dtypes.imageh((32, 128)), 0x400)
-    __validate(dtypes.imageh((96, 128)), 0x400)
-    __validate(dtypes.imageh((64, 256)), 0x840)
-    __validate(dtypes.imageh((64, 9)), 0x80)
-    __validate(dtypes.imageh((192, 256)), 0x840)
-    __validate(dtypes.imageh((64, 768)), 0x1840)
-    __validate(dtypes.imageh((256, 49)), 0x1C0)
-    __validate(dtypes.imageh((128, 9)), 0x80)
-    __validate(dtypes.imageh((16, 1024)), 0x2080)
-    __validate(dtypes.imageh((64, 512)), 0x1040)
-    __validate(dtypes.imageh((16, 512)), 0x1080)
-    __validate(dtypes.imageh((132, 64)), 0x200)
-    __validate(dtypes.imageh((4, 512)), 0x1200)
-    __validate(dtypes.imageh((8, 512)), 0x1100)
-    __validate(dtypes.imageh((128, 128)), 0x400)
-    __validate(dtypes.imageh((32, 512)), 0x1040)
-    __validate(dtypes.imageh((26, 64)), 0x200)
-    __validate(dtypes.imageh((32, 516)), 0x1040)
-    __validate(dtypes.imageh((32, 1024)), 0x2040)
-    __validate(dtypes.imageh((16, 2048)), 0x4080)
-    __validate(dtypes.imageh((8, 2048)), 0x4100)
-    __validate(dtypes.imageh((4, 4096)), 0x8200)
-
-    __validate(dtypes.imagef((16, 49)), 0x380)
-    __validate(dtypes.imagef((16, 1024)), 0x4080)
-    __validate(dtypes.imagef((256, 64)), 0x400)
-    __validate(dtypes.imagef((64, 512)), 0x2040)
-    __validate(dtypes.imagef((16, 512)), 0x2080)
-    __validate(dtypes.imagef((132, 64)), 0x400)
-    __validate(dtypes.imagef((4, 512)), 0x2200)
-    __validate(dtypes.imagef((4, 16)), 0x200)
-    __validate(dtypes.imagef((2, 16)), 0x400)
-    __validate(dtypes.imagef((8, 512)), 0x2100)
-    __validate(dtypes.imagef((12, 64)), 0x400)
-    __validate(dtypes.imagef((3, 32)), 0x400)
-    __validate(dtypes.imagef((128, 128)), 0x840)
-    __validate(dtypes.imagef((32, 512)), 0x2040)
-    __validate(dtypes.imagef((8, 3072)), 0xC100)
-    __validate(dtypes.imagef((4, 2048)), 0x8200)
-    __validate(dtypes.imagef((4, 1024)), 0x4200)
-    __validate(dtypes.imagef((4, 4096)), 0x10200)
-    __validate(dtypes.imagef((10, 384)), 0x1900)
-    __validate(dtypes.imagef((24, 64)), 0x400)
-    __validate(dtypes.imagef((128, 12)), 0xC0)
-    __validate(dtypes.imagef((10, 24)), 0x200)
-    __validate(dtypes.imagef((1, 129)), 0x840)
-    __validate(dtypes.imagef((1, 32)), 0x200)
-    __validate(dtypes.imagef((1, 64)), 0x400)
-    __validate(dtypes.imagef((1, 1239)), 0x4D80)
-    __validate(dtypes.imagef((1, 1)), 0x40)
-
-  def test_image_and_back(self):
-    data = Tensor.randn(9*32*4).realize()
-    tst = data.numpy()
-    it = data.cast(dtypes.imagef((9,32,4))).contiguous().realize()
-    assert isinstance(it.uop.base.realized.dtype, ImageDType)
-    np.testing.assert_equal(tst, it.numpy())
-
-  def test_image_cast_and_back_collapses(self):
-    data = Tensor.randn(9*27*4).realize()
-    tst = data.numpy()
-    it = data.cast(dtypes.imagef((9,27,4))).realize()
-    # the underlying UOp is identical
-    #self.assertIs(it.uop.base.realized, data.uop.base.realized)
-    np.testing.assert_equal(tst, it.numpy())
-
-  def test_image_and_back_wrong_shape(self):
-    data = Tensor.randn(9*27*4).realize()
-    tst = data.numpy()
-    it = data.cast(dtypes.imagef((9,12,4))).realize()
-    assert not isinstance(it.uop.base.realized.dtype, ImageDType)
-    np.testing.assert_equal(tst, it.numpy())
-
-  def test_shrink_load_float(self):
-    it = Tensor.randn(16).cast(dtypes.imagef((1,4,4))).realize()
-    imgv = it.numpy()
-    np.testing.assert_equal(imgv[0:2], it[0:2].numpy())
-
-  def test_mul_stays_image(self):
-    # NOTE: contiguous is needed otherwise this folds
-    it = Tensor.randn(16).cast(dtypes.imagef((1,4,4))).contiguous().realize()
-    out = (it*2).realize()
-    assert isinstance(out.uop.base.realized.dtype, ImageDType)
-
-  def test_sum(self):
-    it = Tensor.rand(8).cast(dtypes.imagef((1,2,4))).realize()
-    itn = it.numpy()
-    np.testing.assert_allclose(np.sum(itn), it.sum().numpy(), rtol=1e-6)
-
-  def test_shrink_max(self):
-    it = Tensor.randn(16).cast(dtypes.imagef((1,4,4))).realize()
-    imgv = it.numpy()
-    np.testing.assert_equal(np.maximum(imgv[0:3], 0), it[0:3].relu().numpy())
-
-  def test_shrink_to_float(self):
-    it = Tensor.randn(4, 4).cast(dtypes.imagef((1,4,4))).realize()
-    imgv = it.numpy()
-    np.testing.assert_equal(np.maximum(imgv[:, 0], 0), it[:, 0].relu().numpy())
-
-  @unittest.skipUnless(isinstance(Device.default.allocator, LRUAllocator), "Requires LRU")
-  def test_lru_alloc(self):
-    data = Tensor.randn(9*27*4).realize()
-    it = data.cast(dtypes.imagef((9,27,4))).realize()
-    b1 = it.uop.base.realized._buf
-    del it
-    it = data.cast(dtypes.imagef((9,27,4))).realize()
-    assert it.uop.base.realized._buf == b1
-
-  def test_no_lru_alloc(self):
-    data = Tensor.randn(9*32*4).realize()
-    it = data.cast(dtypes.imagef((9,32,4))).contiguous().realize()
-    b1 = it.uop.base.realized._buf
-    del it
-    it = data.reshape(9,32,4).pad_to(10, None, None).cast(dtypes.imagef((10,32,4))).contiguous().realize()
-    assert it.uop.base.realized._buf != b1
-
-  def test_no_lru_alloc_dtype(self):
-    data = Tensor.randn(9*32*4).realize()
-    it = data.cast(dtypes.imagef((9,32,4))).contiguous().realize()
-    b1 = it.uop.base.realized._buf
-    del it
-    it = data.cast(dtypes.imageh((9,32,4))).realize()
-    assert it.uop.base.realized._buf != b1
-
   # issue caused by: don't realize image to image casts. this is part of a larger problem
   #@unittest.expectedFailure
   # update: passing after tensor_map
   def test_lil_model(self):
-    with Context(IMAGE=2):
+    with Context(IMAGE=1):
       x = Tensor.zeros(1, 1)
       w1 = Tensor.zeros(1, 8, requires_grad=True)
       w2 = Tensor.zeros(8, 2)
@@ -198,39 +29,5 @@ class TestImageDType(unittest.TestCase):
       self.assertEqual(w1.grad.uop.base.buffer.dtype, dtypes.float32)
       self.assertEqual(len(sched), 9)
 
-@unittest.skipUnless(Device.DEFAULT in IMAGE_SUPPORTED_DEVICES, "Images not supported")
-class TestImageRealization(unittest.TestCase):
-  def test_image_dtype_expand(self):
-    data = Tensor.randn(9*32*4).realize()
-    it = data.cast(dtypes.imagef((9,32,4))).contiguous().realize()
-    self.assertEqual(it.dtype, dtypes.imagef((9,32,4)))
-    it_expanded = it.reshape((9,32,4,1)).expand((9,32,4,4)).contiguous().realize()
-    self.assertEqual(it_expanded.dtype, dtypes.float32)
-
-  def test_image_dtype_expand_and_back(self):
-    data = Tensor.randn(9*32*4).realize()
-    it = data.cast(dtypes.imagef((9,32,4))).contiguous().realize()
-    self.assertEqual(it.dtype, dtypes.imagef((9,32,4)))
-    it_expanded = it.reshape((9,32,4,1)).expand((9,32,4,4))
-    it2 = it_expanded.sum(3).realize()
-    self.assertEqual(it2.dtype, dtypes.imagef((9,32,4)))
-
-  def test_image_alu_children(self):
-    data = Tensor.randn(9*32*4).realize()
-    it = data.cast(dtypes.imagef((9,32,4))).contiguous().realize()
-    self.assertEqual(it.dtype, dtypes.imagef((9,32,4)))
-    it_expanded = it.reshape((9,32,4,1)).expand((9,32,4,4)).contiguous()
-    alu1 = it_expanded+1
-    alu2 = it_expanded.sum(3)
-    it_expanded.realize()
-    # NOTE: the parent becomes float, but the alu child will stay image until its output cannot fit the image
-    self.assertEqual(alu1.dtype, dtypes.imagef((9,32,4)))
-    alu1.realize()
-    self.assertEqual(alu1.dtype, dtypes.float32)
-    # alu2 is back in image because it fits the dtype again
-    self.assertEqual(alu2.dtype, dtypes.imagef((9,32,4)))
-    alu2.realize()
-    self.assertEqual(alu2.dtype, dtypes.imagef((9,32,4)))
-
 if __name__ == '__main__':
   unittest.main()

test/backend/test_jit.py

@@ -39,6 +39,18 @@ class TestJit(unittest.TestCase):
     def add(a, b): return (a+b).realize()
     _simple_test(add)
 
+  def test_jitbeam_triggers_beam(self):
+    from unittest.mock import patch
+    from tinygrad.helpers import getenv as _getenv
+    @TinyJit
+    def add(a, b): return (a+b).realize()
+    a, b = Tensor.ones(10, 10).contiguous().realize(), Tensor.ones(10, 10).contiguous().realize()
+    with patch("tinygrad.codegen.opt.search.beam_search", wraps=lambda k,*a,**kw: k) as mock_beam:
+      add(a, b)
+      assert mock_beam.call_count == 0
+      with patch("tinygrad.engine.jit.getenv", side_effect=lambda k, d=0: 1 if k == "JITBEAM" else _getenv(k, d)): add(a, b)
+      assert mock_beam.call_count == 1
+
   def test_simple_jit_reset(self):
     @TinyJit
     def add(a, b): return (a+b).realize()
@@ -618,7 +630,7 @@ class TestJitFree(unittest.TestCase):
 
     expected_savings = (len(inp) * inp.dtype.itemsize * 2) + dtypes.float32.itemsize # (t1 and t2) + out
 
-    self.assertEqual(savings_after_free, expected_savings)
+    self.assertGreaterEqual(savings_after_free, expected_savings)
     out = fxn(Tensor([11,1,2,3,4]))
     self.assertEqual(out.item(), 136)
 
@@ -628,7 +640,7 @@ class TestJitFree(unittest.TestCase):
     fxn.captured.free_intermediates() # 2nd time to validate
     savings_after_free = pre_free - GlobalCounters.mem_used
 
-    self.assertEqual(savings_after_free, expected_savings)
+    self.assertGreaterEqual(savings_after_free, expected_savings)
     out = fxn(Tensor([11,1,2,3,4]))
     self.assertEqual(out.item(), 136)
 
@@ -648,25 +660,6 @@ class TestJitFree(unittest.TestCase):
     fxn(Tensor([2]))
     self.assertEqual(x.item(), 8)
 
-  def test_replan_buffers_memory_layout(self):
-    if not hasattr(Device[Device.DEFAULT].allocator, '_offset'): raise unittest.SkipTest("replan_buffers_memory_layout useless")
-
-    ext_tensor = Tensor([1,24,23,45,1]).contiguous()
-    ext_tensor_2 = Tensor([2,2,2,2,2]).contiguous()
-    @TinyJit
-    def fxn(x:Tensor):
-      out = (x*ext_tensor_2+ext_tensor).reshape(5,1).expand(5, 100).contiguous()
-      return out.sum()
-    for i in range(5):
-      out = fxn(Tensor([i,1,2,3,4]))
-      self.assertEqual(out.item(), 11400+200*i)
-    self.assertEqual(len(set([b.base for item in fxn.captured.jit_cache for b in item.bufs if b is not None])), 4)
-    fxn.captured.replan_buffers_memory_layout()
-    self.assertEqual(len(set([b.base for item in fxn.captured.jit_cache for b in item.bufs if b is not None])), 2)
-
-    out = fxn(Tensor([11,1,2,3,4]))
-    self.assertEqual(out.item(), 13600)
-
 class TestJitGraphSplit(unittest.TestCase):
   def compute(self, device, inp):
     assert inp.device == device, f"Input device {inp.device} does not match expected {device}"

test/backend/test_linearizer_dumb.py

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

test/backend/test_multitensor.py

@@ -110,7 +110,7 @@ class TestMultiTensor(unittest.TestCase):
   def test_tensor_from_multi(self):
     X = Tensor([1, 2], dtype=dtypes.int).shard_(devices_2, 0)
     Y = Tensor(X.uop)
-    self.assertEqual(Y.device, Device.DEFAULT)
+    self.assertEqual(Y.device, devices_2)
     np.testing.assert_equal(X.numpy(), Y.numpy())
 
     with self.assertRaises(AssertionError):
@@ -645,6 +645,7 @@ class TestMultiTensor(unittest.TestCase):
     out = t0.flip(0) + 1
     self.assertTrue((rng.flip(0)+1).allclose(out.to(rng.device)))
 
+  @unittest.skip("flaky")
   def test_reshape_on_axis(self):
     t0 = Tensor.rand((26, 15, 7)).shard(devices_3, axis=1)
 
@@ -1129,6 +1130,51 @@ class TestTensorOps(unittest.TestCase):
   def test_bitcast(self):
     helper_test_shard_op([(256,), (256,)], lambda x: x.bitcast(dtypes.int))
 
+@unittest.skipIf(not_support_multi_device(), "need multi")
+class TestMultiBufferView(unittest.TestCase):
+  @needs_second_gpu
+  def setUp(self): pass
+
+  def _check(self, a_ref:Tensor, a_multi:Tensor, view_fn):
+    """Apply view_fn to both, verify zero compiled kernels and matching values."""
+    b_ref = view_fn(a_ref)
+    b_multi = view_fn(a_multi).contiguous()
+    sched = b_multi.schedule()
+    compiled = [si for si in sched if isinstance(si.prg, CompiledRunner)]
+    self.assertEqual(len(compiled), 0, f"expected zero compiled kernels, got {len(compiled)}")
+    run_schedule(sched)
+    np.testing.assert_equal(b_multi.numpy(), b_ref.numpy())
+
+  @unittest.skip("flaky on LLVM")
+  def test_shrink_non_shard_axis(self):
+    ref = Tensor.arange(8*4*10).reshape(8, 4, 10).contiguous().realize()
+    a = Tensor.arange(8*4*10).reshape(8, 4, 10).contiguous().shard(devices_2, axis=1).realize()
+    self._check(ref, a, lambda t: t[3])
+
+  def test_shrink_2d(self):
+    ref = Tensor.arange(6*4).reshape(6, 4).contiguous().realize()
+    a = Tensor.arange(6*4).reshape(6, 4).contiguous().shard(devices_2, axis=1).realize()
+    self._check(ref, a, lambda t: t.shrink(((1, 4), None)))
+
+  def test_reshape_then_shrink(self):
+    ref = Tensor.arange(8*6).reshape(8, 6).contiguous().realize()
+    a = Tensor.arange(8*6).reshape(8, 6).contiguous().shard(devices_2, axis=1).realize()
+    self._check(ref, a, lambda t: t.reshape(4, 2, 6)[1])
+
+  def test_chained_shrink(self):
+    ref = Tensor.arange(10*8).reshape(10, 8).contiguous().realize()
+    a = Tensor.arange(10*8).reshape(10, 8).contiguous().shard(devices_2, axis=1).realize()
+    self._check(ref, a, lambda t: t.shrink(((2, 8), None)).shrink(((1, 4), None)))
+
+  def test_4_devices(self):
+    ref = Tensor.arange(8*12).reshape(8, 12).contiguous().realize()
+    a = Tensor.arange(8*12).reshape(8, 12).contiguous().shard(devices_4, axis=1).realize()
+    sched = a[5].contiguous().schedule()
+    compiled = [si for si in sched if isinstance(si.prg, CompiledRunner)]
+    self.assertEqual(len(compiled), 0)
+    run_schedule(sched)
+    np.testing.assert_equal(a[5].contiguous().numpy(), ref[5].numpy())
+
 @unittest.skipIf(not_support_multi_device(), "need multi")
 class TestMultiFromUnrenderable(unittest.TestCase):
   @needs_second_gpu

test/backend/test_randomness.py

@@ -62,7 +62,7 @@ def equal_distribution(tiny_func, torch_func=None, numpy_func=None, shape=(40, 4
   return (numpy_func is None or (kstest(x1, y) >= alpha and kstest(x2, y) >= alpha)) and \
     (torch_func is None or (kstest(x1, z) >= alpha and kstest(x2, z) >= alpha))
 
-def normal_test(func, shape=(20, 23), alpha=0.05): return equal_distribution(func, numpy_func=lambda x: np.random.randn(*x), shape=shape, alpha=alpha)
+def normal_test(func, shape=(20, 45), alpha=0.05): return equal_distribution(func, numpy_func=lambda x: np.random.randn(*x), shape=shape, alpha=alpha)
 
 class TestRandomness(unittest.TestCase):
   def test_rand(self):
@@ -131,29 +131,32 @@ class TestRandomness(unittest.TestCase):
     """
     key0 = 1337
     key1 = int.from_bytes(hashlib.sha256(int(0).to_bytes(4)).digest(), "big") & 0xffffffff
-    values = jax.extend.random.threefry_2x32((np.uint32(key1), np.uint32(key0)), np.arange(20, dtype=np.uint32))
+    # derive new key for the counter offset (c_low=0, c_high=0 for first call)
+    new_key_values = jax.extend.random.threefry_2x32((np.uint32(key1), np.uint32(key0)), np.array([0, 0], dtype=np.uint32))
+    new_key = (np.uint32(new_key_values[0]), np.uint32(new_key_values[1]))
+    values = jax.extend.random.threefry_2x32(new_key, np.arange(20, dtype=np.uint32))
     values = (values >> (32 - 23)) | np.array(1, dtype=np.float32).view(np.uint32)
-    values =  values.view(np.float32) - 1
+    values = values.view(np.float32) - 1
     print(f"[{', '.join(f'{v}' for v in values)}]")
     """
-    jr = np.array([0.9073467254638672, 0.8235964775085449, 0.6872662305831909, 0.9920015335083008, 0.4941047430038452,
-                   0.3108327388763428, 0.09639489650726318, 0.004686474800109863, 0.8435229063034058, 0.824237585067749,
-                   0.5873836278915405, 0.4232727289199829, 0.2530076503753662, 0.40300023555755615, 0.03966474533081055,
-                   0.27904558181762695, 0.9150195121765137, 0.48057758808135986, 0.23821306228637695, 0.7676635980606079], dtype=np.float32)
+    jr = np.array([0.45735931396484375, 0.6311527490615845, 0.15571284294128418, 0.8149417638778687, 0.7862188816070557,
+                   0.8008807897567749, 0.568588376045227, 0.9852620363235474, 0.42314577102661133, 0.9811755418777466,
+                   0.38059568405151367, 0.09186363220214844, 0.9497315883636475, 0.5826880931854248, 0.3796330690383911,
+                   0.5610522031784058, 0.16122901439666748, 0.3732343912124634, 0.9795231819152832, 0.3280656337738037], dtype=np.float32)
     r = Tensor.rand(20).numpy()
     np.testing.assert_allclose(r, jr, atol=1e-5, rtol=1e-5)
 
-    # next 20, np.arange(20, 40, dtype=np.uint32)
-    jr = np.array([0.7444133758544922, 0.7713677883148193, 0.8233780860900879, 0.43871235847473145, 0.517757773399353,
-                   0.6437174081802368, 0.967403769493103, 0.26167726516723633, 0.6825339794158936, 0.14966607093811035,
-                   0.28920769691467285, 0.017063498497009277, 0.2627382278442383, 0.9525482654571533, 0.9351049661636353,
-                   0.43904995918273926, 0.043945908546447754, 0.6616791486740112, 0.6667773723602295, 0.5228077173233032], dtype=np.float32)
+    # next 20 (c_low=20, c_high=0)
+    jr = np.array([0.09199333190917969, 0.9130761623382568, 0.7048608064651489, 0.22254979610443115, 0.0014830827713012695,
+                   0.37023448944091797, 0.7790107727050781, 0.7484984397888184, 0.7524604797363281, 0.19875383377075195,
+                   0.48537540435791016, 0.10002851486206055, 0.5369305610656738, 0.3294715881347656, 0.5246957540512085,
+                   0.7659651041030884, 0.7949080467224121, 0.34988296031951904, 0.9798505306243896, 0.2599533796310425], dtype=np.float32)
     r = Tensor.rand(20).numpy()
     np.testing.assert_allclose(r, jr, atol=1e-5, rtol=1e-5)
 
-    # next 10, np.arange(40, 50, dtype=np.uint32)
-    jr = np.array([0.9614430665969849, 0.059279561042785645, 0.01909029483795166, 0.47882091999053955, 0.9677121639251709,
-                   0.36863112449645996, 0.3102607727050781, 0.06608951091766357, 0.35329878330230713, 0.26518797874450684], dtype=np.float32)
+    # next 10 (c_low=40, c_high=0)
+    jr = np.array([0.3198714256286621, 0.7984923124313354, 0.320881724357605, 0.4716068506240845, 0.7323365211486816,
+                   0.9663800001144409, 0.13873648643493652, 0.16062307357788086, 0.49300849437713623, 0.10077548027038574], dtype=np.float32)
     r = Tensor.rand(10).numpy()
     np.testing.assert_allclose(r, jr, atol=1e-5, rtol=1e-5)
 
@@ -324,7 +327,7 @@ class TestRandomness(unittest.TestCase):
                                                               lambda x: np.random.uniform(-1, 1, size=x) * math.sqrt(6 / (x[0] + math.prod(x[1:])))))
 
   def test_kaiming_uniform(self):
-    for shape in [(32, 16, 3, 3), (20, 44), (3, 15, 35)]:
+    for shape in [(32, 16, 3, 3), (20, 44), (5, 15, 35)]:
       self.assertTrue(equal_distribution(Tensor.kaiming_uniform, lambda x: torch.nn.init.kaiming_uniform_(torch.empty(x)), shape=shape))
 
   def test_kaiming_normal(self):
@@ -384,6 +387,21 @@ class TestRandomness(unittest.TestCase):
     for _ in range(833): Tensor.rand(1)
     Tensor.rand(1).realize()
 
+  def test_random_counter_overflow(self):
+    device = Device.DEFAULT
+    Tensor.manual_seed(1337)
+    Tensor.rand(1).realize()
+
+    Tensor._device_rng_counters[device].assign(Tensor([dtypes.uint32.max - 5, 0], device=device, dtype=dtypes.uint32)).realize()
+
+    Tensor.rand(10).realize()
+    c = Tensor._device_rng_counters[device].numpy()
+    np.testing.assert_allclose(c, [4, 1])
+
+    Tensor.rand(10).realize()
+    c = Tensor._device_rng_counters[device].numpy()
+    np.testing.assert_allclose(c, [14, 1])
+
 # TODO: still fails with MAX_KERNEL_BUFFERS
 @unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
 class TestSample(unittest.TestCase):

test/backend/test_schedule.py

@@ -9,7 +9,7 @@ from hypothesis import assume, given, settings, strategies as strat
 
 from tinygrad import nn, dtypes, Device, Tensor, Variable
 from tinygrad.device import is_dtype_supported
-from tinygrad.dtype import DType, ImageDType
+from tinygrad.dtype import DType
 from tinygrad.uop.ops import UOp, Ops, UPat
 from tinygrad.helpers import CI, DEBUG, OSX, GlobalCounters, Context, getenv, all_same, temp
 from tinygrad.engine.realize import CompiledRunner, run_schedule
@@ -783,18 +783,6 @@ class TestSchedule(unittest.TestCase):
     gc.collect()
     self.assertEqual(GlobalCounters.mem_used-base, 1024)
 
-  @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
-  def test_image_matmul(self):
-    with Context(IMAGE=2):
-      x = Tensor.randn((9, 9)).realize()
-      y = Tensor.randn((9, 9)).realize()
-      out = x@y
-      run_schedule(check_schedule(out, 3))
-      np.testing.assert_allclose(out.numpy(), x.numpy()@y.numpy(), atol=1e-4, rtol=1e-4)
-      self.assertIsInstance(out.dtype, ImageDType)
-      self.assertIsNotNone(out.uop.base.realized)
-      self.assertIsInstance(out.uop.base.realized.dtype, ImageDType)
-
   @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
   def test_image_dot_f16_fusion(self):
     with Context(FLOAT16=1, OPENPILOT_HACKS=1):
@@ -805,11 +793,8 @@ class TestSchedule(unittest.TestCase):
         for si in sched: si.lower()
         return len([si for si in sched if isinstance(si.prg, CompiledRunner)])
 
-      with Context(IMAGE=1): cnt1 = cnt()
-      with Context(IMAGE=2): cnt2 = cnt()
-
-      self.assertEqual(cnt1, 5)
-      self.assertEqual(cnt2, 5)
+      with Context(IMAGE=1):
+        self.assertEqual(cnt(), 5)
 
   @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
   def test_image_f16_residual_fusion(self):
@@ -825,14 +810,10 @@ class TestSchedule(unittest.TestCase):
         for si in sched: si.lower()
         return len([si for si in sched if isinstance(si.prg, CompiledRunner)])
 
-      with Context(IMAGE=1): cnt1 = cnt()
-      with Context(IMAGE=2): cnt2 = cnt()
-
-      self.assertEqual(cnt1, 9)
-      self.assertEqual(cnt2, 9)
+      with Context(IMAGE=1):
+        self.assertEqual(cnt(), 9)
 
   @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
-  @unittest.expectedFailure
   def test_image_conv_fusion(self):
     with Context(OPENPILOT_HACKS=1):
       def cnt():
@@ -843,10 +824,8 @@ class TestSchedule(unittest.TestCase):
         for si in sched: si.lower()
         return len([si for si in sched if isinstance(si.prg, CompiledRunner)])
 
-      with Context(IMAGE=1): cnt1 = cnt()
-      with Context(IMAGE=2): cnt2 = cnt()
-
-      self.assertEqual(cnt1, cnt2)
+      with Context(IMAGE=1):
+        self.assertEqual(cnt(), 5)
 
   def _test_fusion(self, shapes, f, cnt):
     with Context(DEBUG=0, TRACK_MATCH_STATS=0): args = [Tensor.randn(s).realize() for s in shapes]

test/backend/test_setitem.py

@@ -292,6 +292,13 @@ class TestSetitem(unittest.TestCase):
     np.testing.assert_allclose(a.numpy(), [4, 6, 8, 10])
     np.testing.assert_allclose(b.numpy(), [0, 2, 4, 6])
 
+  def test_setitem_multiple_disjoint_on_invalid(self):
+    z = Tensor.invalid(10, dtype="int").realize()
+    z[2:5] = 2
+    z[6:7] = 3
+    z.realize()
+    self.assertListEqual(z[2:5].tolist(), [2, 2, 2])
+    self.assertListEqual(z[6:7].tolist(), [3])
 
 class TestWithGrad(unittest.TestCase):
   def test_no_requires_grad_works(self):

test/backend/test_tensor.py

@@ -256,7 +256,7 @@ class TestTinygrad(unittest.TestCase):
   def test_randperm(self):
     Tensor.manual_seed(0)
     a = Tensor.randperm(10).realize()
-    np.testing.assert_equal(a.numpy(), [5, 2, 8, 1, 3, 7, 9, 6, 0, 4])
+    np.testing.assert_equal(a.numpy(), [8, 9, 4, 3, 6, 1, 7, 5, 2, 0])
     b = Tensor.randperm(1000).realize()
     np.testing.assert_equal(set(b.numpy()), set(range(1000)))
 
@@ -493,6 +493,17 @@ class TestTinygrad(unittest.TestCase):
     dev = a.to(Device.DEFAULT)
     np.testing.assert_allclose(a.numpy(), dev.numpy())
 
+  def test_copy_from_numpy_dtype(self):
+    data = np.array([1.0, 2, 3], dtype=np.float32)
+    t = Tensor(data, dtype=dtypes.bfloat16)
+    try:
+      # TODO: fix dtype in tinygrad space
+      assert t.dtype == dtypes.bfloat16
+    except AssertionError:
+      assert t.dtype == dtypes.float32
+    np.testing.assert_equal(t.tolist(), data)
+    np.testing.assert_equal((t+1).tolist(), data+1)
+
   # Regression test for https://github.com/tinygrad/tinygrad/issues/1751
   def test_copy_from_numpy_unaligned(self):
     # 2**15 is the minimum for repro

test/backend/test_uops.py

@@ -77,15 +77,15 @@ class TestUOps(unittest.TestCase):
   def _test_uop_fxn(self, op, fxn, dts=(dtypes.float32, )):
     for f in [_test_single_value, _test_single_value_const]:
       for a in [-2.0, 0.0, 1.0]:
-        a = dtypes.as_const(a, dts[0])
+        a = dts[0].const(a)
         self._equal(f([a], op, dts), fxn(a))
 
   def _test_bop_fxn(self, op, fxn, dts=(dtypes.float32, )*2, no_b_zero=False, no_b_neg=False):
     for f in [_test_single_value, _test_single_value_const]:
       for a in [-2.0, 0.0, 1.0]:
         for b in [-3.0, 1.0] + ([] if no_b_zero else [0.0]):
-          a = dtypes.as_const(a, dts[0])
-          b = dtypes.as_const(abs(b) if no_b_neg else b, dts[1])
+          a = dts[0].const(a)
+          b = dts[1].const(abs(b) if no_b_neg else b)
           self._equal(f([a,b], op, dts), fxn(a,b))
 
   def _test_top_fxn(self, op, fxn, dts=(dtypes.float32, )*3):
@@ -93,9 +93,9 @@ class TestUOps(unittest.TestCase):
       for a in [-2.0, 0, 1]:
         for b in [-3.0, 3.0]:
           for c in [-4.0, 4.0]:
-            a = dtypes.as_const(a, dts[0])
-            b = dtypes.as_const(b, dts[1])
-            c = dtypes.as_const(c, dts[2])
+            a = dts[0].const(a)
+            b = dts[1].const(b)
+            c = dts[2].const(c)
             self._equal(f([a,b,c], op, dts), fxn(a,b,c))
 
 class TestFloatUOps(TestUOps):
@@ -117,7 +117,7 @@ class TestFloatUOps(TestUOps):
   def test_cmpne_nan(self):  # NaN != x for any x (IEEE 754)
     for a, b in [(math.nan, 1.0), (1.0, math.nan), (math.nan, math.nan)]:
       self.assertTrue(_test_single_value(
-        [dtypes.as_const(a, dtypes.float32), dtypes.as_const(b, dtypes.float32)],
+        [dtypes.float32.const(a), dtypes.float32.const(b)],
         Ops.CMPNE, (dtypes.float32, dtypes.float32)))
   # MOD isn't tested on floats
 

test/external/external_benchmark_op_conv.py

@@ -23,7 +23,7 @@ def vision_conv_143():
   c32 = ((c27<3)!=True)&(c27<67)
   c34 = UOp(Ops.PARAM, dtypes.imageh((32, 1024, 4)), (), 1)
   c38 = c5//2
-  c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.index, Invalid))
+  c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.weakint, Invalid))
   c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
   c49 = UOp(Ops.PARAM, dtypes.imageh((64, 49, 4)), (), 2)
   c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
@@ -49,7 +49,7 @@ def vision_conv_153():
   c32 = ((c27<3)!=True)&(c27<35)
   c34 = UOp(Ops.PARAM, dtypes.imageh((16, 1024, 4)), (), 1)
   c38 = c5//2
-  c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.index, Invalid))
+  c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.weakint, Invalid))
   c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
   c49 = UOp(Ops.PARAM, dtypes.imageh((128, 49, 4)), (), 2)
   c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))

test/external/external_llm_eval.py

@@ -1,32 +1,45 @@
-# eval for tinygrad.apps.llm
-import pyarrow.parquet as pq
+# eval for tinygrad.apps.llm -- hits the server via OpenAI API
+# uses Meta's exact ARC-Challenge prompt template from lm-evaluation-harness llama3 tasks
+import argparse, re, pyarrow.parquet as pq
+from openai import OpenAI
 from tinygrad.helpers import fetch, colored
-from tinygrad.apps.llm import Transformer, SimpleTokenizer, models
-from tinygrad import Tensor
+
+LABEL = ["A", "B", "C", "D"]
 
 if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--port", "-p", type=int, default=11434)
+  parser.add_argument("--limit", "-L", type=int, default=None)
+  parser.add_argument("--max_tokens", "-T", type=int, default=4096)
+  parser.add_argument("--offset", "-O", type=int, default=0)
+  parser.add_argument("--temperature", "-t", type=float, default=0.0)
+  parser.add_argument("--no_think", action="store_true", help="append /no_think to disable thinking (for Qwen3)")
+  parser.add_argument("--debug", action="store_true")
+  args = parser.parse_args()
+
+  client = OpenAI(base_url=f"http://127.0.0.1:{args.port}/v1", api_key="tinygrad")
   dat = fetch("https://huggingface.co/datasets/allenai/ai2_arc/resolve/main/ARC-Challenge/test-00000-of-00001.parquet")
   table = pq.read_table(dat)
 
-  model, kv = Transformer.from_gguf(Tensor.from_url(models["1B"]), max_context=4096)
-
-  tok = SimpleTokenizer.from_gguf_kv(kv)
-  bos_id: int = kv['tokenizer.ggml.bos_token_id']
-  eos_id: int = kv['tokenizer.ggml.eos_token_id']
-
   num_correct, num_answered = 0, 0
-  total_questions = len(table["question"])
-  for question, choices, answer in zip(table["question"], table["choices"], table["answerKey"]):
-    phrasing = f"Question: {question}\n\n" + \
-               '\n'.join([f"{k}) {v}" for k,v in zip(choices['label'], choices['text'])]) +\
-               "\n\nReply with the letter of the correct answer only."
-    try:
-      ids = [bos_id] + tok.role("user") + tok.encode(phrasing) + [eos_id] + tok.role("assistant") + tok.encode("Answer: ")
-    except RuntimeError:
-      # TODO: fix the tokenizer
-      pass
-    next_id = next(model.generate(ids))
-    correct, given = answer.as_py().strip(), tok.decode([next_id]).strip()
+  # filter to 4-choice questions and normalize labels to A/B/C/D (matches Meta's eval)
+  rows = [(q, c, a) for q, c, a in zip(table["question"], table["choices"], table["answerKey"]) if len(c["label"]) == 4]
+  total_questions = min(len(rows), args.offset + args.limit) if args.limit else len(rows)
+  for question, choices, answer in rows[args.offset:total_questions]:
+    phrasing = "Given the following question and four candidate answers (A, B, C and D), choose the best answer.\n" +\
+               f"Question: {question}\n" + '\n'.join([f"{l}. {t}" for l, t in zip(LABEL, choices['text'])]) +\
+               '\nYour response should end with "The best answer is [the_answer_letter]"' +\
+               " where the [the_answer_letter] is one of A, B, C or D." + (" /no_think" if args.no_think else "")
+    resp = client.chat.completions.create(model="test", messages=[{"role": "user", "content": phrasing}],
+                                          max_tokens=args.max_tokens, temperature=args.temperature)
+    # normalize answer key (some use 1/2/3/4 instead of A/B/C/D)
+    correct = answer.as_py().strip()
+    if correct not in LABEL: correct = LABEL[int(correct) - 1]
+    # extract answer: take last single capital letter A-D from response (prompt asks model to end with the answer)
+    text = resp.choices[0].message.content.strip()
+    if args.debug: print(f"\n--- PROMPT ---\n{phrasing}\n--- RESPONSE ---\n{text}\n---")
+    m = re.findall(r'\b([A-D])\b', text)
+    given = m[-1] if m else text[:1]
     num_correct += correct == given
     num_answered += 1
     print(f"{num_answered:4d}/{total_questions:4d}  "+\

test/external/external_test_am.py

@@ -176,6 +176,34 @@ class TestAMPageTable(unittest.TestCase):
       mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)], aspace=AddrSpace.PHYS)
       mm0.unmap_range(helper_va(0x1000000), 2 << 20)
 
+  def test_inspect_mode(self):
+    mm0 = self.d[0].mm
+
+    # Map a few disjoint ranges inside a larger region.
+    mappings = [(0x10000, 0x3000), (0x20000, 0x2000), (0x1000000, 2 << 20)]
+    for va, sz in mappings:
+      mm0.map_range(helper_va(va), sz, paddrs=[(va, sz)], aspace=AddrSpace.PHYS)
+
+    # Inspect over the whole region: should visit all mapped pages.
+    ctx = PageTableTraverseContext(self.d[0], mm0.root_page_table, helper_va(0x0), inspect=True)
+    visited = set()
+    for _off, pt, pte_idx, n_ptes, pte_covers in ctx.next(0x4000000):
+      for i in range(n_ptes):
+        pte = helper_read_entry_components(pt.entries[pte_idx + i])
+        if pte['valid']:
+          for p in range(0, pte_covers, 0x1000): visited.add(pte['paddr'] + p)
+
+    expected_pages = {va + off for va, sz in mappings for off in range(0, sz, 0x1000)}
+    assert visited == expected_pages
+
+    for va, sz in mappings:
+      mm0.unmap_range(helper_va(va), sz)
+
+    # Inspect after unmap: should find no valid entries.
+    ctx = PageTableTraverseContext(self.d[0], mm0.root_page_table, helper_va(0x0), inspect=True)
+    for _off, pt, pte_idx, n_ptes, pte_covers in ctx.next(0x4000000):
+      for i in range(n_ptes): assert not pt.valid(pte_idx + i)
+
   def test_frag_size(self):
     mm0 = self.d[0].mm
 

test/external/external_test_onnx_runner.py

@@ -3,7 +3,8 @@ import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.uop.ops import Ops
 from tinygrad.device import is_dtype_supported
-from tinygrad.nn.onnx import OnnxRunner, OnnxDataType
+from typing import Any
+from tinygrad.nn.onnx import OnnxRunner, OnnxPBParser, OnnxDataType
 from hypothesis import given, strategies as st
 
 # copied from test_const_folding.py
@@ -136,5 +137,40 @@ class TestOnnxRunnerDtypes(unittest.TestCase):
         from_disk=False)
     self.assertEqual(runner.graph_nodes[0].opts['value'].dtype, expected_dtype)
 
+# from openpilot selfdrive/modeld/get_model_metadata.py
+class MetadataOnnxPBParser(OnnxPBParser):
+  def _parse_ModelProto(self) -> dict:
+    obj: dict[str, Any] = {"graph": {"input": [], "output": []}, "metadata_props": []}
+    for fid, wire_type in self._parse_message(self.reader.len):
+      match fid:
+        case 7: obj["graph"] = self._parse_GraphProto()
+        case 14: obj["metadata_props"].append(self._parse_StringStringEntryProto())
+        case _: self.reader.skip_field(wire_type)
+    return obj
+
+class TestOnnxMetadata(unittest.TestCase):
+  def test_metadata_props(self):
+    graph = onnx.helper.make_graph(
+      nodes=[onnx.helper.make_node('Identity', ['input'], ['output'])],
+      name='test',
+      inputs=[onnx.helper.make_tensor_value_info('input', onnx.TensorProto.FLOAT, (1, 3))],
+      outputs=[onnx.helper.make_tensor_value_info('output', onnx.TensorProto.FLOAT, (1, 3))],
+    )
+    model = onnx.helper.make_model(graph)
+    model.metadata_props.append(onnx.StringStringEntryProto(key="model_checkpoint", value="v1.0"))
+    model.metadata_props.append(onnx.StringStringEntryProto(key="output_slices", value="dGVzdA=="))
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+      model_path = pathlib.Path(tmpdir) / "model.onnx"
+      onnx.save(model, model_path)
+      parsed = MetadataOnnxPBParser(model_path).parse()
+
+    # metadata_props should be accessible as dicts with "key" and "value"
+    self.assertEqual(len(parsed["metadata_props"]), 2)
+    self.assertEqual(parsed["metadata_props"][0]["key"], "model_checkpoint")
+    self.assertEqual(parsed["metadata_props"][0]["value"], "v1.0")
+    self.assertEqual(parsed["metadata_props"][1]["key"], "output_slices")
+    self.assertEqual(parsed["metadata_props"][1]["value"], "dGVzdA==")
+
 if __name__ == '__main__':
   unittest.main()

test/external/external_test_opt.py

@@ -6,14 +6,18 @@ import torch
 from tinygrad import GlobalCounters, Tensor, Device
 from tinygrad.helpers import getenv
 from tinygrad.nn.state import get_parameters
-from tinygrad.engine.realize import capturing
+from tinygrad.engine.realize import capturing, run_schedule
+from tinygrad.engine.schedule import linear_to_schedule
 from tinygrad.tensor import _to_np_dtype
 
 class CLCache:
   def __init__(self, allowed=None, strict=False, preclear=True, var_vals=None):
     self.allowed, self.strict, self.preclear, self.var_vals = allowed, strict, preclear, var_vals if var_vals is not None else {}
     self.count = 0
-  def add(self, ei): self.count += 1
+  def add_linear(self, linear, var_vals):
+    schedule = linear_to_schedule(linear)
+    self.count += len(schedule)
+    run_schedule(schedule, var_vals)
   def __enter__(self):
     if self.preclear:
       gc.collect()

test/mockgpu/amd/amdgpu.py

@@ -17,7 +17,7 @@ regCOMPUTE_USER_DATA_0 = 0x1be0 + amd_gpu.GC_BASE__INST0_SEG0
 regCOMPUTE_NUM_THREAD_X = 0x1ba7 + amd_gpu.GC_BASE__INST0_SEG0
 regGRBM_GFX_INDEX = 0x2200 + amd_gpu.GC_BASE__INST0_SEG1
 regSQ_THREAD_TRACE_BUF0_BASE = 0x39e8 + amd_gpu.GC_BASE__INST0_SEG1
-regSQ_THREAD_TRACE_BUF0_SIZE = {"rdna3": 0x39e9, "rdna4": 0x39e6}[MOCKGPU_ARCH] + amd_gpu.GC_BASE__INST0_SEG1
+regSQ_THREAD_TRACE_BUF0_SIZE = {"rdna3": 0x39e9, "rdna4": 0x39e6, "cdna4": 0x39e9}[MOCKGPU_ARCH] + amd_gpu.GC_BASE__INST0_SEG1
 regSQ_THREAD_TRACE_WPTR = 0x39ef + amd_gpu.GC_BASE__INST0_SEG1
 regSQ_THREAD_TRACE_STATUS = 0x39f4 + amd_gpu.GC_BASE__INST0_SEG1
 regCP_PERFMON_CNTL = 0x3808 + amd_gpu.GC_BASE__INST0_SEG1
@@ -200,11 +200,13 @@ class PM4Executor(AMDQueue):
       if st <= prg_addr < st+sz: prg_sz = sz - (prg_addr - st)
 
     # Get scratch size from COMPUTE_TMPRING_SIZE register
-    # For gfx11: WAVESIZE = ceildiv(64 * size_per_thread, 256), so size_per_thread ≈ WAVESIZE * 256 / 64 = WAVESIZE * 4
+    # WAVESIZE = ceildiv(lanes * size_per_thread, mem_alignment_size)
+    # GFX11+: mem_alignment_size=256, so size_per_thread = WAVESIZE * 256 / 64 = WAVESIZE * 4
+    # GFX9:   mem_alignment_size=1024, so size_per_thread = WAVESIZE * 1024 / 64 = WAVESIZE * 16
     try: tmpring_size = self.gpu.regs[regCOMPUTE_TMPRING_SIZE]
     except KeyError: tmpring_size = 0
-    wavesize = (tmpring_size >> 12) & 0x3FFF  # WAVESIZE field is bits 12:25 for gfx11
-    scratch_size = wavesize * 4  # This gives the scratch size per thread (lane)
+    wavesize = (tmpring_size >> 12) & 0x3FFF  # WAVESIZE field is bits 12:25
+    scratch_size = wavesize * (16 if self.gpu.arch == "cdna" else 4)  # per-thread scratch size in bytes
 
     assert prg_sz > 0, "Invalid prg ptr (not found in mapped ranges)"
     # Pass valid memory ranges, rsrc2, scratch_size, arch, and user data registers to Python emulator
@@ -402,7 +404,7 @@ p2p_links_count 5
 cpu_core_id_base 0
 simd_id_base 2147488032
 max_waves_per_simd 16
-lds_size_in_kb 128
+lds_size_in_kb 160
 gds_size_in_kb 0
 num_gws 64
 wave_front_size 64

test/mockgpu/amd/pcode.py

@@ -2,6 +2,7 @@
 from typing import Any, Callable
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import Ops, UOp
+from tinygrad.uop.decompositions import f2f
 
 # Type alias for vars dict: stores UOps for variables and tuples for lambda definitions
 VarVal = UOp | tuple[str, list[str], str]
@@ -40,9 +41,13 @@ def _bitreverse(v: UOp, bits: int) -> UOp:
 
 def _extract_bits(val: UOp, hi: int, lo: int) -> UOp:
   dt = dtypes.uint64 if val.dtype in (dtypes.uint64, dtypes.int64) else dtypes.uint32
-  result = ((val >> _const(dt, lo)) if lo > 0 else val) & _const(val.dtype, (1 << (hi - lo + 1)) - 1)
-  # Downcast to uint32 when extracting <=32 bits from a 64-bit value, so .f32 bitcast works correctly
-  if dt == dtypes.uint64 and (hi - lo + 1) <= 32: result = result.cast(dtypes.uint32)
+  width = hi - lo + 1
+  # Cast to dt first to ensure shift operands have matching types
+  val_cast = val.cast(dt) if val.dtype != dt else val
+  result = ((val_cast >> _const(dt, lo)) if lo > 0 else val_cast) & _const(dt, (1 << width) - 1)
+  # Downcast to match extracted bit width so brace-concat { hi, lo } computes correct output dtype
+  target_dt = _BITS_DT.get(width) or (dtypes.uint32 if width <= 32 else dtypes.uint64 if width <= 64 else dt)
+  if result.dtype != target_dt: result = result.cast(target_dt)
   return result
 
 def _set_bit(old, pos, val):
@@ -60,6 +65,48 @@ def _floor(x):
   return ((x < _const(x.dtype, 0)) & x.ne(t)).where(t - _const(x.dtype, 1), t)
 def _f16_extract(v): return (v & _u32(0xFFFF)).cast(dtypes.uint16).bitcast(dtypes.half) if v.dtype == dtypes.uint32 else v
 
+# ═════ FP8 (E4M3) and BF8 (E5M2) conversion helpers ═════
+# f32→fp8/bf8 uses f2f decomposition directly. fp8/bf8→f32 wraps f2f with subnormal handling
+# (f2f flushes denormals to zero, but AMD V_CVT_F32_FP8/BF8 preserves subnormals).
+def _fp8_to_f32(v: UOp) -> UOp:
+  b = (v.cast(dtypes.uint32) & _u32(0xFF)).cast(dtypes.uint8)
+  # E4M3 subnormal: exp==0, mant!=0 -> (-1)^sign * 2^(1-7) * (mant/8) = (-1)^sign * mant * 2^(-9)
+  bu = b.cast(dtypes.uint32)
+  sign, exp, mant = (bu >> _u32(7)) << _u32(31), (bu >> _u32(3)) & _u32(0xF), bu & _u32(0x7)
+  is_sub = exp.eq(_u32(0)) & mant.ne(_u32(0))
+  sub_f32 = (mant.cast(dtypes.float32) * _const(dtypes.float32, 1.0/512.0)).bitcast(dtypes.uint32) | sign
+  normal = f2f(b, dtypes.fp8e4m3, dtypes.float32)
+  return is_sub.where(sub_f32.bitcast(dtypes.float32), normal)
+
+def _bf8_to_f32(v: UOp) -> UOp:
+  b = (v.cast(dtypes.uint32) & _u32(0xFF)).cast(dtypes.uint8)
+  # E5M2 subnormal: exp==0, mant!=0 -> (-1)^sign * 2^(1-15) * (mant/4) = (-1)^sign * mant * 2^(-16)
+  bu = b.cast(dtypes.uint32)
+  sign, exp, mant = (bu >> _u32(7)) << _u32(31), (bu >> _u32(2)) & _u32(0x1F), bu & _u32(0x3)
+  is_sub = exp.eq(_u32(0)) & mant.ne(_u32(0))
+  sub_f32 = (mant.cast(dtypes.float32) * _const(dtypes.float32, 1.0/65536.0)).bitcast(dtypes.uint32) | sign
+  normal = f2f(b, dtypes.fp8e5m2, dtypes.float32)
+  return is_sub.where(sub_f32.bitcast(dtypes.float32), normal)
+
+def _f32_to_fp8(v: UOp) -> UOp:
+  return f2f((v.bitcast(dtypes.float32) if v.dtype != dtypes.float32 else v).bitcast(dtypes.uint32), dtypes.float32, dtypes.fp8e4m3)
+def _f32_to_bf8(v: UOp) -> UOp:
+  return f2f((v.bitcast(dtypes.float32) if v.dtype != dtypes.float32 else v).bitcast(dtypes.uint32), dtypes.float32, dtypes.fp8e5m2)
+def _f32_to_bf16(v: UOp) -> UOp:
+  """Convert f32 to bf16 with round-to-nearest-even. BF16 is the upper 16 bits of F32 with rounding."""
+  bits = (v.bitcast(dtypes.float32) if v.dtype != dtypes.float32 else v).bitcast(dtypes.uint32)
+  # Round-to-nearest-even: add rounding bias. If the bit just below the truncation point is 1 and the rest are 0, round to even.
+  round_bit = (bits >> _u32(16)) & _u32(1)  # bit 16 (LSB of kept part)
+  rounding = _u32(0x7FFF) + round_bit  # 0x7FFF + bit16: rounds to even
+  rounded = bits + rounding
+  return (rounded >> _u32(16)).cast(dtypes.uint16)
+def _f32_to_bf16_sr(v: UOp, stoch: UOp) -> UOp:
+  """Convert f32 to bf16 with stochastic rounding."""
+  bits = (v.bitcast(dtypes.float32) if v.dtype != dtypes.float32 else v).bitcast(dtypes.uint32)
+  # Stochastic rounding: add lower 16 bits of stochastic value to lower 16 bits of f32
+  rounded = bits + (stoch & _u32(0xFFFF))
+  return (rounded >> _u32(16)).cast(dtypes.uint16)
+
 def _check_nan(v: UOp, quiet: bool) -> UOp:
   if v.op == Ops.CAST and v.dtype == dtypes.float64: v = v.src[0]
   bits, exp_m, mant_m, qb, _ = _float_info(v)
@@ -119,7 +166,10 @@ def _abs(val: UOp) -> UOp:
   bt, ft = {10: (dtypes.uint16, dtypes.half), 23: (dtypes.uint32, dtypes.float32), 52: (dtypes.uint64, dtypes.float64)}[shift]
   return (val.bitcast(bt) & _const(bt, sign_mask)).bitcast(ft)
 
-def _f_to_u(f, dt): return UOp(Ops.TRUNC, f.dtype, ((f < _const(f.dtype, 0.0)).where(_const(f.dtype, 0.0), f),)).cast(dt)
+def _f_to_u(f, dt):
+  clamped = (f < _const(f.dtype, 0.0)).where(_const(f.dtype, 0.0), f)
+  truncated = UOp(Ops.TRUNC, f.dtype, (clamped,))
+  return (truncated >= _const(f.dtype, 2**(dt.itemsize*8))).where(_const(dt, dt.max), truncated.cast(dt))
 
 def _cvt_quiet(val: UOp) -> UOp:
   bits, _, _, qb, _ = _float_info(val)
@@ -289,6 +339,9 @@ _FUNCS: dict[str, Callable[..., UOp]] = {
   'CalcDsAddr': lambda a, o, *r: a.cast(dtypes.uint32) + o.cast(dtypes.uint32),
   'CalcGlobalAddr': lambda v, s, *r: v.cast(dtypes.uint64) + s.cast(dtypes.uint64),
   'CalcScratchAddr': lambda v, s, *r: v.cast(dtypes.uint64) + s.cast(dtypes.uint64),
+  # FP8/BF8/BF16 conversion functions
+  'fp8_to_f32': _fp8_to_f32, 'bf8_to_f32': _bf8_to_f32, 'f32_to_fp8': _f32_to_fp8, 'f32_to_bf8': _f32_to_bf8,
+  'f32_to_bf16': _f32_to_bf16, 'f32_to_bf16_SR': _f32_to_bf16_sr, 'f32_to_bf16_sr': _f32_to_bf16_sr,
 }
 for is_max, name in [(False, 'min'), (True, 'max')]:
   for dt, sfx in [(dtypes.float32, 'f32'), (dtypes.int, 'i32'), (dtypes.uint32, 'u32'), (dtypes.int16, 'i16'), (dtypes.uint16, 'u16')]:
@@ -313,7 +366,8 @@ for is_max, name in [(False, 'min'), (True, 'max')]:
 
 DTYPES = {'u32': dtypes.uint32, 'i32': dtypes.int, 'f32': dtypes.float32, 'b32': dtypes.uint32, 'u64': dtypes.uint64, 'i64': dtypes.int64,
           'f64': dtypes.float64, 'b64': dtypes.uint64, 'u16': dtypes.uint16, 'i16': dtypes.short, 'f16': dtypes.half, 'b16': dtypes.uint16,
-          'u8': dtypes.uint8, 'i8': dtypes.int8, 'b8': dtypes.uint8, 'u4': dtypes.uint8, 'i4': dtypes.int8, 'u1': dtypes.uint32}
+          'u8': dtypes.uint8, 'i8': dtypes.int8, 'b8': dtypes.uint8, 'u4': dtypes.uint8, 'i4': dtypes.int8, 'u1': dtypes.uint32,
+          'fp8': dtypes.uint8, 'bf8': dtypes.uint8, 'b3': dtypes.uint8, 'b2': dtypes.uint8}
 _BITS_DT = {8: dtypes.uint8, 16: dtypes.uint16, 32: dtypes.uint32, 64: dtypes.uint64}
 _NUM_SUFFIXES = ('ULL', 'LL', 'UL', 'U', 'L', 'F', 'f')
 def _strip_suffix(num: str) -> tuple[str, str]:
@@ -425,7 +479,14 @@ class Parser:
       case '+' | '-':
         if op == '-' and left.op == Ops.CONST and right.op == Ops.CONST: return _const(left.dtype, left.arg - right.arg)
         return (left + right) if op == '+' else (left - right)
-      case '*' | '/': return (left * right) if op == '*' else (left / right)
+      case '*' | '/':
+        # Integer promotion: promote 16-bit integers to 32-bit before multiply to avoid overflow
+        # (e.g. SOPP branch offset: SIMM16.i16 * 16'4 can exceed int16 range)
+        if op == '*' and left.dtype.itemsize == 2 and left.dtype in (dtypes.int16, dtypes.short, dtypes.uint16, dtypes.ushort):
+          pdt = dtypes.int if left.dtype in (dtypes.int16, dtypes.short) else dtypes.uint
+          left, right = left.cast(pdt), right.cast(pdt)
+        if op == '*': return left * right
+        return (left // right) if dtypes.is_int(left.dtype) else (left / right)
       case '**': return UOp(Ops.EXP2, left.dtype, (right.cast(left.dtype),)) if left.op == Ops.CONST and left.arg == 2.0 else left
 
   _PREC = [('||',), ('&&',), ('|',), ('^',), ('&',), ('==', '!=', '<>'), ('>=', '<=', '>', '<'), ('>>', '<<'), ('+', '-'), ('*', '/'), ('**',)]
@@ -502,7 +563,8 @@ class Parser:
         self.eat('RBRACKET')
         vgpr = self.vars.get('_vgpr')
         if vgpr is None: return _u32(0)
-        return vgpr.index(_to_u32(reg) * _u32(32) + _to_u32(lane), ptr=True).load()
+        ws = self.vars.get('_wave_size', 32)
+        return vgpr.index(_to_u32(reg) * _u32(ws) + _to_u32(lane), ptr=True).load()
       if self.try_eat('LPAREN'):
         args = self._parse_args()
         self.eat('RPAREN')
@@ -514,8 +576,8 @@ class Parser:
       if name == 'OVERFLOW_F32': return _const(dtypes.uint32, 0x7F7FFFFF).bitcast(dtypes.float32)
       if name == 'UNDERFLOW_F64': return _const(dtypes.uint64, 1).bitcast(dtypes.float64)
       if name == 'OVERFLOW_F64': return _const(dtypes.uint64, 0x7FEFFFFFFFFFFFFF).bitcast(dtypes.float64)
-      if name == 'WAVE32': return _const(dtypes.bool, True)
-      if name == 'WAVE64': return _const(dtypes.bool, False)
+      if name == 'WAVE32': return _const(dtypes.bool, self.vars.get('_wave_size', 32) <= 32)
+      if name == 'WAVE64': return _const(dtypes.bool, self.vars.get('_wave_size', 32) > 32)
       if name == 'WAVE_MODE' and self.try_eat('DOT') and self.try_eat_val('IEEE', 'IDENT'): return _u32(1)
       if self.try_eat('LBRACE'):
         idx = self.eat('NUM').val
@@ -527,7 +589,8 @@ class Parser:
           self.eat('RBRACKET')
           vgpr = self.vars.get('_vgpr')
           if vgpr is None: return _u32(0)
-          return vgpr.index(_to_u32(reg) * _u32(32) + _u32(int(idx)), ptr=True).load()
+          ws = self.vars.get('_wave_size', 32)
+          return vgpr.index(_to_u32(reg) * _u32(ws) + _u32(int(idx)), ptr=True).load()
         elem = self.vars.get(f'{name}@{idx}', self.vars.get(f'{name}{idx}'))
         if elem is None:
           # Extract bit idx from base variable (like var[idx])
@@ -660,13 +723,14 @@ class Parser:
     return None
 
   def _sized_literal(self, bits: int) -> UOp:
-    if self.at('IDENT') and self.peek().val in ('U', 'I', 'F', 'B'):
+    if self.at('IDENT') and self.peek().val in ('U', 'I', 'F', 'B', 'BF'):
       type_char = self.eat('IDENT').val
       self.eat('LPAREN')
       inner = self.parse()
       self.eat('RPAREN')
       dt = {('U',32): dtypes.uint32, ('U',64): dtypes.uint64, ('I',32): dtypes.int, ('I',64): dtypes.int64,
             ('F',16): dtypes.half, ('F',32): dtypes.float32, ('F',64): dtypes.float64,
+            ('BF',16): dtypes.bfloat16,
             ('B',32): dtypes.uint32, ('B',64): dtypes.uint64}.get((type_char, bits), dtypes.uint64 if bits > 32 else dtypes.uint32)
       if type_char == 'F' and inner.dtype in (dtypes.uint32, dtypes.uint64, dtypes.ulong, dtypes.int, dtypes.int64):
         if inner.dtype.itemsize != dt.itemsize: inner = inner.cast(dtypes.uint32 if dt.itemsize == 4 else dtypes.uint64)
@@ -769,6 +833,18 @@ class Parser:
       elif dt in (dtypes.uint8, dtypes.int8): val = (val >> ((addr & _const(adt, 3)).cast(dtypes.uint32) * _u32(8))) & _u32(0xFF)
       elif dt in (dtypes.uint16, dtypes.int16):
         val = (val >> (((addr >> _const(adt, 1)) & _const(adt, 1)).cast(dtypes.uint32) * _u32(16))) & _u32(0xFFFF)
+      else:
+        # Handle unaligned 32-bit loads: combine two consecutive dwords and shift.
+        # To avoid OOB at buffer boundaries for aligned loads, clamp idx_hi to idx (safe).
+        # Use int64 for the WHERE to avoid 32-bit int overflow in C pointer arithmetic (addr can be >8GB).
+        byte_off = (addr & _const(adt, 3)).cast(dtypes.uint32)
+        is_unaligned = byte_off.ne(_u32(0))
+        idx_native = (addr >> _const(adt, 2)).cast(dtypes.int64)
+        idx_hi_native = ((addr + _const(adt, 4)) >> _const(adt, 2)).cast(dtypes.int64)
+        safe_idx_hi = is_unaligned.where(idx_hi_native, idx_native)
+        hi = mem.index(safe_idx_hi, *gate)
+        combined = val.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32))
+        val = is_unaligned.where((combined >> (byte_off.cast(dtypes.uint64) * UOp.const(dtypes.uint64, 8))).cast(dtypes.uint32), val)
     return val
 
   def _coerce_cmp(self, l: UOp, r: UOp) -> tuple[UOp, UOp]:
@@ -980,17 +1056,34 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
       i += 1
       continue
 
-    # VGPR assignment: VGPR[lane][reg] = value
+    # VGPR assignment: VGPR[lane][reg] = value  or  VGPR[lane][reg][hi:lo].type = { ... }
     if first == 'vgpr' and toks[1].type == 'LBRACKET':
       j, lane_toks = _match_bracket(toks, 1)
       if j < len(toks) and toks[j].type == 'LBRACKET':
         j, reg_toks = _match_bracket(toks, j)
+        # Check for bit-slice: VGPR[lane][reg][hi:lo].type = value (read-modify-write)
+        if j < len(toks) and toks[j].type == 'LBRACKET':
+          j, slice_toks = _match_bracket(toks, j)
+          slice_str = _tok_str(slice_toks)
+          hi_str, lo_str = slice_str.split(':')
+          hi_val, lo_val = int(eval(hi_str.strip())), int(eval(lo_str.strip()))
+          if j < len(toks) and toks[j].type == 'DOT': j += 2  # skip .type suffix
+          if j < len(toks) and toks[j].type == 'EQUALS': j += 1
+          ln = parse_tokens(lane_toks, env, funcs)
+          rg, val = parse_tokens(reg_toks, env, funcs), parse_tokens(toks[j:], env, funcs)
+          ws = env.get('_wave_size', 32)
+          vgpr_idx = _to_u32(rg) * _u32(ws) + _to_u32(ln)
+          if assigns is not None:
+            assigns.append((f'VGPR[{_tok_str(lane_toks)}][{_tok_str(reg_toks)}][{hi_val}:{lo_val}]', (vgpr_idx, val, _u32(hi_val), _u32(lo_val))))
+          i += 1
+          continue
         if j < len(toks) and toks[j].type == 'DOT': j += 2  # skip .type suffix
         if j < len(toks) and toks[j].type == 'EQUALS': j += 1
         ln = parse_tokens(lane_toks, env, funcs)
         rg, val = parse_tokens(reg_toks, env, funcs), parse_tokens(toks[j:], env, funcs)
         if assigns is not None:
-          assigns.append((f'VGPR[{_tok_str(lane_toks)}][{_tok_str(reg_toks)}]', (_to_u32(rg) * _u32(32) + _to_u32(ln), val)))
+          ws = env.get('_wave_size', 32)
+          assigns.append((f'VGPR[{_tok_str(lane_toks)}][{_tok_str(reg_toks)}]', (_to_u32(rg) * _u32(ws) + _to_u32(ln), val)))
         i += 1
         continue
 
@@ -1143,12 +1236,16 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
       def is_const(c, v): return c.op == Ops.CONST and c.arg is v
       cond = parse_cond(line, 'if')
       conditions: list[tuple[UOp, UOp | dict[str, VarVal] | None]] = [(cond, None)] if not is_const(cond, False) else []
+      branch_assigns: list[tuple[UOp, list]] = []  # (cond, assigns_list) for side-effect merging
       else_branch: tuple[UOp | None, dict[str, VarVal]] = (None, {})
+      else_side_effects: list = []
       env_snap = dict(env)
       static_true = is_const(cond, True)  # track if any condition is statically true
       i += 1
-      i, branch, ret = parse_block(lines, i, env, funcs, assigns if not is_const(cond, False) else None)
+      if_side: list = [] if assigns is not None and not is_const(cond, False) else []
+      i, branch, ret = parse_block(lines, i, env, funcs, if_side if assigns is not None and not is_const(cond, False) else None)
       if conditions: conditions[0] = (cond, ret if ret is not None else branch)
+      if assigns is not None and not is_const(cond, False): branch_assigns.append((cond, if_side))
       env.clear()
       env.update(env_snap)
       while i < len(lines):
@@ -1159,16 +1256,21 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
           c = parse_cond(lines[i], 'elsif')
           take = not static_true and not is_const(c, False)
           i += 1
-          i, branch, ret = parse_block(lines, i, env, funcs, assigns if take else None)
+          br_side: list = [] if assigns is not None and take else []
+          i, branch, ret = parse_block(lines, i, env, funcs, br_side if assigns is not None and take else None)
           if take:
             conditions.append((c, ret if ret is not None else branch))
             if is_const(c, True): static_true = True
+            if assigns is not None: branch_assigns.append((c, br_side))
           env.clear()
           env.update(env_snap)
         elif lf == 'else':
           i += 1
-          i, branch, ret = parse_block(lines, i, env, funcs, assigns if not static_true else None)
-          if not static_true: else_branch = (ret, branch)
+          el_side: list = [] if assigns is not None and not static_true else []
+          i, branch, ret = parse_block(lines, i, env, funcs, el_side if assigns is not None and not static_true else None)
+          if not static_true:
+            else_branch = (ret, branch)
+            if assigns is not None: else_side_effects = el_side
           env.clear()
           env.update(env_snap)
         elif lf == 'endif':
@@ -1188,6 +1290,10 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
         ba = next((b for c, b in conditions if is_const(c, True) and isinstance(b, dict)), {})
         block_assigns.update(ba)
         env.update(ba)
+        # For static true, forward side effects unconditionally
+        if assigns is not None:
+          for bc, bse in branch_assigns:
+            if is_const(bc, True): assigns.extend(bse)
       else:
         else_assigns = else_branch[1]
         all_vars = set().union(*[ba.keys() for _, ba in conditions if isinstance(ba, dict)], else_assigns.keys())
@@ -1199,6 +1305,21 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
               if isinstance(tv, UOp) and isinstance(res, UOp):
                 res = cond.where(tv, res.cast(tv.dtype) if tv.dtype != res.dtype and tv.dtype.itemsize == res.dtype.itemsize else res)
           block_assigns[var] = env[var] = res
+        # Merge side effects from branches with conditions
+        if assigns is not None:
+          def _cond_side_effect(cnd, dest, val):
+            if isinstance(val, tuple) and len(val) == 4:  # VGPR bit-slice: (idx, rhs, hi, lo) -> add condition
+              return (dest, (val[0], val[1], val[2], val[3], cnd))
+            if isinstance(val, tuple) and len(val) == 2:  # VGPR/MEM write: (addr, rhs) -> condition rhs
+              return (dest, (val[0], cnd.where(val[1], val[1])))
+            return (dest, val)
+          # Build combined condition: each branch fires when its cond is true AND no earlier cond was true
+          remaining = UOp.const(dtypes.bool, True)
+          for bc, bse in branch_assigns:
+            effective = remaining & bc if remaining.op != Ops.CONST else bc
+            for dest, val in bse: assigns.append(_cond_side_effect(effective, dest, val))
+            remaining = remaining & bc.logical_not() if remaining.op != Ops.CONST else bc.logical_not()
+          for dest, val in else_side_effects: assigns.append(_cond_side_effect(remaining, dest, val))
       continue
 
     # Regular assignment: var = value

test/mockgpu/cuda/cuda.py

@@ -154,7 +154,7 @@ def cuMemHostAlloc(pp, bytesize: int, flags: int) -> int:
 def cuMemFreeHost(p: ctypes.c_void_p) -> int: return cuMemFree_v2(p)
 
 def cuMemcpyDtoDAsync_v2(dst, src, bytesize: int, stream: Any) -> int:
-  ctypes.memmove(dst.value, src.value, bytesize)
+  ctypes.memmove(dst if isinstance(dst, int) else dst.value, src if isinstance(src, int) else src.value, bytesize)
   return orig_cuda.CUDA_SUCCESS
 
 def cuFuncSetAttribute(hfunc, attrib: int, value: int) -> int:

test/null/test_device.py

@@ -1,8 +1,10 @@
 #!/usr/bin/env python
 import unittest, os, subprocess
+from unittest.mock import patch
 from tinygrad import Tensor
 from tinygrad.device import Device, Compiler, enumerate_devices_str
-from tinygrad.helpers import diskcache_get, diskcache_put, getenv, Context, WIN, CI
+from tinygrad.helpers import diskcache_get, diskcache_put, getenv, Context, WIN, CI, OSX
+from tinygrad.runtime.support.c import DLL
 
 class TestDevice(unittest.TestCase):
   def test_canonicalize(self):
@@ -75,6 +77,18 @@ class TestDevice(unittest.TestCase):
       self.assertIsInstance(Device["CPU"].compiler, CPULLVMCompiler)
       assert inst is Device["CPU"].compiler  # cached
 
+  @unittest.skipIf(Device.DEFAULT != "CPU", "only run on CPU")
+  def test_compiler_autodetect_fallback(self):
+    from tinygrad.runtime.support.compiler_cpu import CPULLVMCompiler
+
+    try: CPULLVMCompiler()
+    except Exception as e: self.skipTest(f"skipping: LLVM not available: {e}")
+
+    dev = Device["CPU"]
+    dev.cached_pair.clear()
+    with patch("tinygrad.renderer.cstyle.ClangJITRenderer.__init__", side_effect=RuntimeError("broken")):
+      self.assertIsInstance(dev.renderer.compiler, CPULLVMCompiler)
+
 class MockCompiler(Compiler):
   def __init__(self, key): super().__init__(key)
   def compile(self, src) -> bytes: return src.encode()
@@ -100,6 +114,7 @@ class TestCompiler(unittest.TestCase):
       a = Tensor([0.,1.], device=Device.DEFAULT).realize()
       (a + 1).realize()
 
+@unittest.skipIf(OSX and 'libclang' in DLL._loaded_, "MTLCompiler can't be loaded after libclang on OSX")
 class TestRunAsModule(unittest.TestCase):
   def test_module_runs(self):
     cpu_line = [l for l in enumerate_devices_str() if "CPU" in l][0]

test/null/test_gpudims.py

@@ -1,9 +1,10 @@
 import unittest, math
 import z3
-from tinygrad.codegen.gpudims import get_grouped_dims
-from tinygrad.uop.ops import UOp, Ops
+from tinygrad.codegen.gpudims import get_grouped_dims, add_gpudims
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
 from tinygrad.uop.validate import uops_to_z3
 from tinygrad.dtype import dtypes
+from tinygrad.renderer import Renderer
 from tinygrad.helpers import flatten, dedup
 
 class TestGroupedDims(unittest.TestCase):
@@ -23,7 +24,7 @@ class TestGroupedDims(unittest.TestCase):
     total = math.prod(dims)
     specials = sorted(dedup(flatten([[y for y in x.toposort() if y.op is Ops.SPECIAL] for x in idxs])), key=lambda u: u.arg)
     # build flat index and primed flat (same expression with renamed SPECIALs)
-    flat = UOp.const(dtypes.index, 0)
+    flat = UOp.const(dtypes.weakint, 0)
     for i, idx in enumerate(idxs):
       flat = flat + idx * int(math.prod(dims[i+1:]))
     flat_p = flat.substitute({s: UOp(Ops.SPECIAL, s.dtype, s.src, s.arg+"_p") for s in specials})
@@ -93,6 +94,14 @@ class TestGroupedDims(unittest.TestCase):
     assert idxs[2].op is Ops.SPECIAL, f"expected SPECIAL for direct-mapped dim, got {idxs[2].op}"
     assert idxs[3].op is Ops.SPECIAL, f"expected SPECIAL for direct-mapped dim, got {idxs[3].op}"
 
+  def test_global_prod_max(self):
+    g, l = UOp.range(256, 0, AxisType.GLOBAL), UOp.range(256, 1, AxisType.LOCAL)
+    sink = UOp(Ops.PARAM, dtypes.float.ptr(), (), 0).index(g + l).store(UOp.const(dtypes.float, 1.0)).end(g, l).sink(arg=KernelInfo())
+    class R(Renderer): global_max, local_max, global_prod_max = (256, 256, 256), (128, 128, 128), (128, 128, 128)
+    specials = [u for u in add_gpudims(R(), sink).toposort() if u.op is Ops.SPECIAL]
+    self.assertGreater(len([s for s in specials if "lidx" in s.arg]), 1)
+    self.assertGreater(len([s for s in specials if "gidx" in s.arg]), 1)
+
   def test_max_sizes_none(self):
     self._check_grouped_dims("gidx", (2,3,4), None, False, [2,3,4])
     self._check_grouped_dims("gidx", (100,), None, False, [100])

test/null/test_graph_rewrite.py

@@ -99,21 +99,21 @@ class TestFoldingAndReduction(unittest.TestCase):
 class TestModuloAndDivisionFolding(unittest.TestCase):
   def test_full_graph_rewrite_modulo_folding_with_define_var(self):
     # index dtype because div-mod rules only work on index
-    x_var_uop = UOp.variable('x', 0, 100).cast(dtypes.index)
+    x_var_uop = UOp.variable('x', 0, 100).cast(dtypes.weakint)
     optimized_mod_uop = apply_rewrite(((x_var_uop * 4) + 2) % 4)
     self.assertEqual(optimized_mod_uop.op, Ops.CONST)
     self.assertEqual(optimized_mod_uop.arg, 2)
 
   def test_full_graph_rewrite_division_folding_with_define_var(self):
     # index dtype because div-mod rules only work on index
-    n_var_uop = UOp.variable('n', 1, 1000).cast(dtypes.index)
+    n_var_uop = UOp.variable('n', 1, 1000).cast(dtypes.weakint)
     optimized_div_uop = apply_rewrite((n_var_uop * 6) // 3)
     self.assertEqual(optimized_div_uop.op, Ops.MUL)
     self.assertEqual(optimized_div_uop.src[1].arg, 2)
 
   def test_full_graph_rewrite_complex_mod_div_folding(self):
     # index dtype because div-mod rules only work on index
-    k_var_uop = UOp.variable('k', 0, 50).cast(dtypes.index)
+    k_var_uop = UOp.variable('k', 0, 50).cast(dtypes.weakint)
     optimized_div_uop = apply_rewrite(((k_var_uop * 12 + 8) % 6) // 2)
     self.assertEqual(optimized_div_uop.op, Ops.CONST)
     self.assertEqual(optimized_div_uop.arg, 1)
@@ -132,7 +132,7 @@ class TestModuloAndDivisionFolding(unittest.TestCase):
   def test_full_graph_rewrite_modulo_large_divisor(self):
     # index dtype because div-mod rules only work on index
     x_var_uop = UOp.variable('x', 1, 5)
-    self.assertIs(apply_rewrite(x_var_uop.cast(dtypes.index) % 10).render(simplify=False), x_var_uop.render(simplify=False))
+    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)

test/null/test_linearizer_failures.py

@@ -8,12 +8,12 @@ from tinygrad.device import Device
 class TestLinearizerFailures(unittest.TestCase):
   def test_fail_1(self):
     c0 = UOp(Ops.PARAM, dtypes.float.ptr(64), arg=0, src=())
-    c1 = UOp.range(UOp.const(dtypes.index, 2), 1, AxisType.LOOP)
-    c2 = UOp.range(UOp.const(dtypes.index, 32), 2, AxisType.LOOP)
-    c3 = ((c1*UOp.const(dtypes.index, 32))+c2)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 2), 1, AxisType.LOOP)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 32), 2, AxisType.LOOP)
+    c3 = ((c1*UOp.const(dtypes.weakint, 32))+c2)
     c4 = UOp(Ops.PARAM, dtypes.float.ptr(163840), arg=1, src=())
-    c5 = UOp.range(UOp.const(dtypes.index, 2560), 0, AxisType.REDUCE)
-    c6 = c4.index(((((((c5//UOp.const(dtypes.index, 8))%UOp.const(dtypes.index, 8))*UOp.const(dtypes.index, 8))+(c5%UOp.const(dtypes.index, 8)))+(((c2*UOp.const(dtypes.index, 40))+(c5//UOp.const(dtypes.index, 64)))*UOp.const(dtypes.index, 64)))+(c1*UOp.const(dtypes.index, 81920))))
+    c5 = UOp.range(UOp.const(dtypes.weakint, 2560), 0, AxisType.REDUCE)
+    c6 = c4.index(((((((c5//UOp.const(dtypes.weakint, 8))%UOp.const(dtypes.weakint, 8))*UOp.const(dtypes.weakint, 8))+(c5%UOp.const(dtypes.weakint, 8)))+(((c2*UOp.const(dtypes.weakint, 40))+(c5//UOp.const(dtypes.weakint, 64)))*UOp.const(dtypes.weakint, 64)))+(c1*UOp.const(dtypes.weakint, 81920))))
     c7 = UOp(Ops.PARAM, dtypes.float.ptr(64), arg=2, src=())
     c8 = c7.index(c3)
     c9 = ((((c6+(c8*UOp.const(dtypes.float, -1.0)))*(c6+(c8*UOp.const(dtypes.float, -1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(dtypes.float, 0.000390625))+UOp.const(dtypes.float, 1e-05)).sqrt().reciprocal()

test/null/test_llm_server.py

@@ -21,6 +21,7 @@ class TestLLMServer(unittest.TestCase):
 
     import tinygrad.apps.llm as llm_module
     llm_module.model = cls.mock_model
+    llm_module.model_name = "test-model"
     llm_module.tok = cls.mock_tok
     llm_module.bos_id = cls.bos_id
     llm_module.eos_id = cls.eos_id
@@ -134,5 +135,33 @@ class TestLLMServer(unittest.TestCase):
     self.assertIsNotNone(resp.usage.prompt_tokens)
     self.assertIsNotNone(resp.usage.completion_tokens)
 
+  def test_max_tokens_streaming(self):
+    self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 302, 303, 999]))
+    stream = self.client.chat.completions.create(
+      model="test", messages=[{"role": "user", "content": "Hello"}], stream=True, max_tokens=2
+    )
+    chunks = list(stream)
+    content_chunks = [c for c in chunks if c.choices and c.choices[0].delta.content]
+    self.assertEqual(len(content_chunks), 2)
+    self.assertEqual(chunks[-1].choices[0].finish_reason, "length")
+
+  def test_max_tokens_non_streaming(self):
+    self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 302, 303, 999]))
+    resp = self.client.chat.completions.create(
+      model="test", messages=[{"role": "user", "content": "Hello"}], stream=False, max_tokens=2
+    )
+    self.assertEqual(resp.choices[0].finish_reason, "length")
+    self.assertEqual(resp.usage.completion_tokens, 2)
+
+  def test_models_endpoint(self):
+    import requests as req
+    resp = req.get(f"http://127.0.0.1:{self.port}/v1/models")
+    self.assertEqual(resp.status_code, 200)
+    data = resp.json()
+    self.assertEqual(data["object"], "list")
+    self.assertEqual(len(data["data"]), 1)
+    self.assertEqual(data["data"][0]["id"], "test-model")
+    self.assertEqual(data["data"][0]["object"], "model")
+
 if __name__ == '__main__':
   unittest.main()

test/null/test_memory_planner.py

@@ -1,42 +1,73 @@
 import unittest
 from tinygrad import dtypes
-from tinygrad.device import Buffer
-from tinygrad.engine.memory import _internal_memory_planner
+from tinygrad.uop.ops import UOp, Ops
+from tinygrad.engine.memory import memory_plan_rewrite
 
 global_map = {}
+held_bufs: set[UOp] = set()
 def b(i, base=None, offset=0, pin=False, size=16):
   global global_map
   if i in global_map: return global_map[i]
-  global_map[i] = Buffer("NULL", size, dtypes.int8, base=global_map[base] if base is not None else None, offset=offset)
-  if pin: global_map[i].ref(1)
+  if base is not None:
+    global_map[i] = global_map[base]
+    return global_map[i]
+  global_map[i] = UOp.new_buffer("NULL", size, dtypes.int8)
+  if pin: held_bufs.add(global_map[i])
   return global_map[i]
 
-def check_assign(buffers:list[list[Buffer]|tuple[Buffer, ...]], copies:list[tuple[Buffer, Buffer]]|None=None):
-  assigned = _internal_memory_planner(buffers, copies=copies)
+def _make_linear(buffer_lists, copies=None):
+  copy_pairs = {frozenset((id(dst), id(src))) for dst, src in copies} if copies else set()
+  calls = []
+  for bufs in buffer_lists:
+    is_copy = len(bufs) == 2 and frozenset((id(bufs[0]), id(bufs[1]))) in copy_pairs
+    calls.append(UOp(Ops.CALL, dtypes.void, (UOp(Ops.COPY if is_copy else Ops.SINK), *bufs)))
+  return UOp(Ops.LINEAR, src=tuple(calls))
 
-  taken_parts = set()
+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.BUFFER_VIEW: return new.src[0]
+  return None
+
+def check_assign(buffer_lists, copies=None):
+  linear = _make_linear(buffer_lists, copies)
+  result = memory_plan_rewrite(linear, held_bufs)
+
+  # build mapping: original buf -> (arena, offset_bytes, nbytes) from the result
+  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.BUFFER_VIEW and id(orig) not in replace_map:
+        replace_map[id(orig)] = (new.src[0], new.arg[1] * new.dtype.itemsize, new.arg[0] * new.dtype.itemsize)
+
+  # verify pinned buffers are not planned
+  for buf in held_bufs:
+    assert id(buf) not in replace_map, "pinned buffer was planned"
+
+  # compute lifetimes
   first_appearance, last_appearance = {}, {}
-  for i,u in enumerate(buffers):
-    for buf in u:
-      if buf.is_allocated() or buf.base.is_allocated() or buf.uop_refcount > 0: continue
-      if buf.base not in first_appearance: first_appearance[buf.base] = i
-      last_appearance[buf.base] = i
+  for i, bufs in enumerate(buffer_lists):
+    for buf in bufs:
+      if buf in held_bufs: continue
+      if id(buf) not in first_appearance: first_appearance[id(buf)] = i
+      last_appearance[id(buf)] = i
 
-  for i,u in enumerate(buffers):
-    for buf in u:
-      if buf.is_allocated() or buf.base.is_allocated() or buf.uop_refcount > 0: continue
-      cur, base = assigned.get(buf, buf), assigned.get(buf.base, buf.base)
-      if buf._base is not None:
-        assert cur.base == base.base and cur.offset == buf.offset + base.offset, f"failed: {buf} {cur} {base} {buf.offset} {base.offset}"
-      else:
-        for part in taken_parts:
-          assert buf.base == part[3] or part[0] != cur.base or part[1] + part[2] <= cur.offset or part[1] >= cur.offset + buf.nbytes
-        if first_appearance[buf.base] == i: taken_parts.add((cur.base, cur.offset, buf.nbytes, buf.base))
-        if last_appearance[buf.base] == i: taken_parts.remove((cur.base, cur.offset, buf.nbytes, buf.base))
+  # verify non-overlapping: no two live buffers share the same arena region
+  taken_parts: set[tuple[int, int, int, int]] = set()  # (id(arena), offset, nbytes, id(buf))
+  for i, bufs in enumerate(buffer_lists):
+    for buf in bufs:
+      if buf in held_bufs or id(buf) not in replace_map: continue
+      arena, off, nb = replace_map[id(buf)]
+      for part in taken_parts:
+        assert id(buf) == part[3] or part[0] != id(arena) or part[1] + part[2] <= off or part[1] >= off + nb, \
+          f"overlap at step {i}: [{off}, {off+nb}) conflicts with [{part[1]}, {part[1]+part[2]})"
+      if first_appearance.get(id(buf)) == i: taken_parts.add((id(arena), off, nb, id(buf)))
+      if last_appearance.get(id(buf)) == i: taken_parts.discard((id(arena), off, nb, id(buf)))
 
 class TestMemoryPlanner(unittest.TestCase):
   def setUp(self):
     global global_map
+    held_bufs.clear()
     global_map = {}
 
   def test_simple_buffer(self):
@@ -140,9 +171,11 @@ class TestMemoryPlanner(unittest.TestCase):
       [b(1), b(2)],
       [b(3), b(2)],
     ]
-    assigned = _internal_memory_planner(bs, copies=[(b(1), b(0))])
-    r1, r2 = assigned.get(b(1), b(1)), assigned.get(b(2), b(2))
-    assert r1.base != r2.base
+    linear = _make_linear(bs, copies=[(b(1), b(0))])
+    result = memory_plan_rewrite(linear)
+    r1_arena, r2_arena = _get_arena(b(1), linear, result), _get_arena(b(2), linear, result)
+    assert r1_arena is not None and r2_arena is not None
+    assert r1_arena is not r2_arena
 
   def test_copy_bufs_reuse_among_copies(self):
     bs = [
@@ -150,9 +183,11 @@ class TestMemoryPlanner(unittest.TestCase):
       [b(2), b(1)],
       [b(3), b(2)],
     ]
-    assigned = _internal_memory_planner(bs, copies=[(b(1), b(0)), (b(2), b(1))])
-    r1, r2 = assigned.get(b(1), b(1)), assigned.get(b(2), b(2))
-    assert r1.base == r2.base
+    linear = _make_linear(bs, copies=[(b(1), b(0)), (b(2), b(1))])
+    result = memory_plan_rewrite(linear)
+    r1_arena, r2_arena = _get_arena(b(1), linear, result), _get_arena(b(2), linear, result)
+    assert r1_arena is not None and r2_arena is not None
+    assert r1_arena is r2_arena
 
   def test_compute_bufs_reuse_among_compute(self):
     bs = [
@@ -161,9 +196,11 @@ class TestMemoryPlanner(unittest.TestCase):
       [b(3), b(2)],
       [b(4), b(3)],
     ]
-    assigned = _internal_memory_planner(bs, copies=[(b(1), b(0))])
-    r2, r3 = assigned.get(b(2), b(2)), assigned.get(b(3), b(3))
-    assert r2.base == r3.base
+    linear = _make_linear(bs, copies=[(b(1), b(0))])
+    result = memory_plan_rewrite(linear)
+    r2_arena, r3_arena = _get_arena(b(2), linear, result), _get_arena(b(3), linear, result)
+    assert r2_arena is not None and r3_arena is not None
+    assert r2_arena is r3_arena
 
   def test_copy_and_compute_no_cross_reuse(self):
     bs = [
@@ -171,9 +208,11 @@ class TestMemoryPlanner(unittest.TestCase):
       [b(2), b(1)],
       [b(3), b(2)],
     ]
-    assigned = _internal_memory_planner(bs, copies=[(b(2), b(1))])
-    r0, r2 = assigned.get(b(0), b(0)), assigned.get(b(2), b(2))
-    assert r0.base != r2.base
+    linear = _make_linear(bs, copies=[(b(2), b(1))])
+    result = memory_plan_rewrite(linear)
+    r0_arena, r2_arena = _get_arena(b(0), linear, result), _get_arena(b(2), linear, result)
+    assert r0_arena is not None and r2_arena is not None
+    assert r0_arena is not r2_arena
 
   def test_multiple_copy_bufs_with_offsets(self):
     bs = [

test/null/test_real_world.py

@@ -74,7 +74,10 @@ class TestRealWorld(unittest.TestCase):
     def test(t, t2):
       for l in model: t = l(t, t2)
       return t.realize()
-    helper_test("test_unet_resblock", lambda: (Tensor.empty(4, 16, 8, 8), Tensor.empty(1, 24)), test, 0.0002, 37)
+
+    # TODO: support _offset on CL to get mem down to 0.0002
+    exp_mem = 0.00037 if Device.DEFAULT == "CL" else 0.0002
+    helper_test("test_unet_resblock", lambda: (Tensor.empty(4, 16, 8, 8), Tensor.empty(1, 24)), test, exp_mem, 37)
 
   @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need dtypes.float16")
   def test_llama(self):

test/null/test_schedule.py

@@ -1189,6 +1189,61 @@ class TestBufferView(unittest.TestCase):
     b = a.shrink(((200, 800),)).shrink(((0, 300),)).reshape((30, 10)).shrink(((20, 25), (0, 10))).contiguous()
     run_schedule(check_schedule(b, 0))
 
+  def test_shrink_non_shard_axis_is_buffer_view_multi(self):
+    # indexing a non-shard axis of a realized sharded tensor should be BUFFER_VIEW on each device, not copy kernels
+    # this is the flat_llama pattern: weight[layer_idx] where weight is (n_layers, out, dim) sharded on axis=1
+    devices = ("NULL:1", "NULL:2")
+    a = Tensor.arange(8*4*10).reshape(8, 4, 10).contiguous().shard(devices, axis=1).realize()
+    run_schedule(check_schedule(a[3].contiguous(), 0))
+
+  def test_shrink_2d_non_shard_axis_multi(self):
+    devices = ("NULL:1", "NULL:2")
+    a = Tensor.arange(6*4).reshape(6, 4).contiguous().shard(devices, axis=1).realize()
+    run_schedule(check_schedule(a.shrink(((1, 4), None)).contiguous(), 0))
+
+  def test_shrink_shard_axis_0_multi(self):
+    # shrinking a middle dim is not contiguous per shard, so this needs copy kernels
+    devices = ("NULL:1", "NULL:2")
+    a = Tensor.arange(4*6*2).reshape(4, 6, 2).contiguous().shard(devices, axis=0).realize()
+    run_schedule(check_schedule(a.shrink((None, (2, 5), None)).contiguous(), 2))
+
+  def test_reshape_then_shrink_multi(self):
+    devices = ("NULL:1", "NULL:2")
+    a = Tensor.arange(8*6).reshape(8, 6).contiguous().shard(devices, axis=1).realize()
+    run_schedule(check_schedule(a.reshape(4, 2, 6)[1].contiguous(), 0))
+
+  def test_permute_then_shrink_multi(self):
+    # permute makes per-shard view non-contiguous, needs copy kernels
+    devices = ("NULL:1", "NULL:2")
+    a = Tensor.arange(4*6*2).reshape(4, 6, 2).contiguous().shard(devices, axis=1).realize()
+    run_schedule(check_schedule(a.permute(1, 0, 2).shrink(((0, 6), (1, 3), None)).contiguous(), 2))
+
+  def test_multi_buffer_view_4_devices(self):
+    devices = tuple(f"NULL:{i}" for i in range(4))
+    a = Tensor.arange(8*12).reshape(8, 12).contiguous().shard(devices, axis=1).realize()
+    run_schedule(check_schedule(a[5].contiguous(), 0))
+
+  def test_chained_shrink_multi(self):
+    devices = ("NULL:1", "NULL:2")
+    a = Tensor.arange(10*8).reshape(10, 8).contiguous().shard(devices, axis=1).realize()
+    run_schedule(check_schedule(a.shrink(((2, 8), None)).shrink(((1, 4), None)).contiguous(), 0))
+
+  # negative tests: these should NOT become BUFFER_VIEW (non-contiguous per shard)
+  def test_expand_multi_not_buffer_view(self):
+    devices = ("NULL:1", "NULL:2")
+    a = Tensor.arange(4*2).reshape(4, 1, 2).contiguous().shard(devices, axis=2).realize()
+    run_schedule(check_schedule(a.expand(4, 3, 2).contiguous(), 2))
+
+  def test_pad_multi_not_buffer_view(self):
+    devices = ("NULL:1", "NULL:2")
+    a = Tensor.arange(4*2).reshape(4, 2).contiguous().shard(devices, axis=1).realize()
+    run_schedule(check_schedule(a.pad(((1, 1), (0, 0))).contiguous(), 2))
+
+  def test_flip_multi_not_buffer_view(self):
+    devices = ("NULL:1", "NULL:2")
+    a = Tensor.arange(4*2).reshape(4, 2).contiguous().shard(devices, axis=1).realize()
+    run_schedule(check_schedule(a.flip(0).contiguous(), 2))
+
 class TestInvalidTensor(unittest.TestCase):
   def test_full_invalid_is_zero_kernels(self):
     from tinygrad.dtype import Invalid

test/null/test_simplify_valid_idx.py

@@ -15,11 +15,11 @@ def get_gated_load_uop(valid:UOp, idx:UOp):
 
 def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]):
   return UOp(Ops.LOAD, dtypes.float.vec(4), (
-    UOp(Ops.PARAM, dtypes.imagef(image_shape), arg=0).index(UOp(Ops.VECTORIZE, dtypes.index.vec(2), idx).valid(valid), ptr=True),
+    UOp(Ops.PARAM, dtypes.imagef(image_shape), arg=0).index(UOp(Ops.VECTORIZE, dtypes.weakint.vec(2), idx).valid(valid), ptr=True),
     UOp(Ops.VECTORIZE, dtypes.float.vec(4), src=(UOp.const(dtypes.float, 0.0),) * 4)
   ))
 
-def Special(expr, nmax): return UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, nmax),), expr)
+def Special(expr, nmax): return UOp(Ops.SPECIAL, dtypes.weakint, (UOp.const(dtypes.weakint, nmax),), expr)
 def Variable(expr, nmin, nmax): return UOp.variable(expr, nmin, nmax)
 def Range(n, nmax): return UOp.range(nmax, n)
 
@@ -454,23 +454,13 @@ class TestImageSimplification(unittest.TestCase):
     load = get_load_image_uop((32, 1024, 4), valid, (alu0, alu1))
     self.check(load, "(lidx1<7)", "((gidx0*2+lidx1*512+(lidx0*8192+r0*4096)+-11711)//4%1024)", "(lidx0*2+r0+-3)")
 
-class TestUnfoldableImageChannelSelection(unittest.TestCase):
-  def _count_nans(self, load):
-    with Context(NOOPT=1, SPEC=0):
-      result = full_rewrite_to_sink(load.sink()).src[0]
-    return sum(1 for u in result.toposort() if u.op is Ops.CONST and u.arg != u.arg)
-
-  def test_bounded_channel_no_nan(self):
-    # unfoldable image load with bounded idx % 4 range [0,1] -> no NAN fallback needed
-    lidx = Special("lidx", 2)
-    load = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.PARAM, dtypes.imagef((10, 10, 4)), arg=0).index(lidx, ptr=True), UOp.const(dtypes.float, 0)))
-    self.assertEqual(self._count_nans(load), 0)
-
-  def test_unbounded_channel_has_nan(self):
-    # variable with negative range -> x % 4 can be negative -> needs NAN fallback
-    x = Variable("x", -10, 10)
-    load = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.PARAM, dtypes.imagef((10, 10, 4)), arg=0).index(x, ptr=True), UOp.const(dtypes.float, 0)))
-    self.assertEqual(self._count_nans(load), 1)
+class TestUnfoldableImage(unittest.TestCase):
+  def test_unfoldable_becomes_buffer(self):
+    with Context(SPEC=0):
+      lidx = Special("lidx", 2)
+      load = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.PARAM, dtypes.imagef((10, 10, 4)), arg=0).index(lidx, ptr=True), UOp.const(dtypes.float, 0)))
+      res = full_rewrite_to_sink(load.sink()).src[0]
+      self.assertEqual(res.src[0].src[0].dtype, dtypes.float.ptr(400))
 
 class TestDropTrueGate(unittest.TestCase):
   def test_drop_true_gate_on_index(self):
@@ -478,7 +468,7 @@ class TestDropTrueGate(unittest.TestCase):
     from tinygrad.codegen.late.devectorizer import load_store_indexing
     from tinygrad.uop.ops import graph_rewrite
     buf = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0)
-    idx = UOp.const(dtypes.index, 0)
+    idx = UOp.const(dtypes.weakint, 0)
     true_gate = UOp.const(dtypes.bool, True)
     index_with_gate = UOp(Ops.INDEX, dtypes.int.ptr(), (buf, idx, true_gate))
     # apply the optimization
@@ -495,7 +485,7 @@ class TestRangeShrink(unittest.TestCase):
   def test_range_shrink_single_guard(self):
     # range 0..203 guarded by r < 4 everywhere -> shrink to 0..3
     r = Range(0, 204)
-    load = get_gated_load_uop(r < UOp.const(dtypes.index, 4), r)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
     ranges = self.get_ranges(load.sink())
     self.assertEqual(len(ranges), 1)
     self.assertEqual(ranges[0].src[0].arg, 4)
@@ -503,8 +493,8 @@ class TestRangeShrink(unittest.TestCase):
   def test_range_shrink_picks_max_guard(self):
     # two loads guard the same range with r < 4 and r < 8 -> shrink to max(4, 8) = 8
     r = Range(0, 204)
-    load1 = get_gated_load_uop(r < UOp.const(dtypes.index, 4), r)
-    load2 = get_gated_load_uop(r < UOp.const(dtypes.index, 8), r)
+    load1 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
+    load2 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 8), r)
     ranges = self.get_ranges(UOp.sink(load1, load2))
     self.assertEqual(len(ranges), 1)
     self.assertEqual(ranges[0].src[0].arg, 8)
@@ -512,7 +502,7 @@ class TestRangeShrink(unittest.TestCase):
   def test_range_no_shrink_guard_ge_max(self):
     # guard r < 300 with range max 204 -> no shrink (guard doesn't constrain)
     r = Range(0, 204)
-    load = get_gated_load_uop(r < UOp.const(dtypes.index, 300), r)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 300), r)
     ranges = self.get_ranges(load.sink())
     self.assertEqual(len(ranges), 1)
     self.assertEqual(ranges[0].src[0].arg, 204)
@@ -520,7 +510,7 @@ class TestRangeShrink(unittest.TestCase):
   def test_range_no_shrink_when_unguarded_elsewhere(self):
     # one load guards r < 4, but another load uses r without a gate -> no shrink
     r = Range(0, 204)
-    load1 = get_gated_load_uop(r < UOp.const(dtypes.index, 4), r)
+    load1 = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
     load2 = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.PARAM, dtypes.float.ptr(), arg=1).index(r, ptr=True),))
     ranges = self.get_ranges(UOp.sink(load1, load2))
     self.assertEqual(len(ranges), 1)
@@ -529,7 +519,7 @@ class TestRangeShrink(unittest.TestCase):
   def test_range_no_shrink_when_used_in_reduce(self):
     # range used in both a gated load AND directly in the reduce expression -> no shrink
     r = Range(0, 204)
-    gated_load = get_gated_load_uop(r < UOp.const(dtypes.index, 4), r)
+    gated_load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 4), r)
     red = UOp(Ops.REDUCE, dtypes.float, (r.cast(dtypes.float) + gated_load, r), Ops.ADD)
     ranges = self.get_ranges(red.sink())
     self.assertEqual(len(ranges), 1)
@@ -538,7 +528,7 @@ class TestRangeShrink(unittest.TestCase):
   def test_range_shrink_to_single_iteration(self):
     # guard r < 1 shrinks range to 1 -> single iteration, range eliminated entirely
     r = Range(0, 204)
-    load = get_gated_load_uop(r < UOp.const(dtypes.index, 1), r)
+    load = get_gated_load_uop(r < UOp.const(dtypes.weakint, 1), r)
     ranges = self.get_ranges(load.sink())
     self.assertEqual(len(ranges), 0)
 

test/null/test_tensor.py

@@ -162,5 +162,17 @@ class TestTensorUnique(unittest.TestCase):
     Tensor.realize(b,c)
     self.assertIs(b.uop.buffer, c.uop.buffer)
 
+class TestRand(unittest.TestCase):
+  def test_rand_large_tensor(self):
+    # large tensor rand (num > uint32.max) should not crash in frontend
+    Tensor.manual_seed(0)
+    Tensor.rand(2**17, 2**17).schedule()
+    Tensor.rand(2**17, 2**17).schedule()
+    Tensor.rand(2**17, 2**17).schedule()
+
+class TestTensorDevice(unittest.TestCase):
+  def test_create_from_single_device_tuple(self):
+    (Tensor([1.0], device=(Device.DEFAULT,)) + Tensor([2.0])).realize()
+
 if __name__ == '__main__':
   unittest.main()

test/null/test_tensor_metadata.py

@@ -2,6 +2,7 @@ import unittest
 from tinygrad import Tensor, dtypes
 from tinygrad.tensor import _METADATA
 from tinygrad.engine.realize import capturing
+from tinygrad.engine.schedule import linear_to_schedule
 from tinygrad.helpers import Context
 
 @unittest.skip("tensor metadata is no longer supported")
@@ -94,10 +95,11 @@ class TestTensorMetadata(unittest.TestCase):
       self.assertEqual(si.metadata, ())
 
   def _has_metadata(self, h, name):
-    items = []
-    capturing.append(type("", (), {"add": lambda _, ei: items.append(ei)})())
+    linears = []
+    capturing.append(type("", (), {"add_linear": lambda _, linear, var_vals: linears.append(linear)})())
     try: h.realize()
     finally: capturing.clear()
+    items = [ei for linear in linears for ei in linear_to_schedule(linear)]
     return any(m.name == name for ei in items for m in ei.metadata)
 
   def test_metadata_survives_realize_pending_assign(self):

test/null/test_uop_graph.py

@@ -205,7 +205,7 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_where_same_fold(self):
     v = UOp.variable('tmp', 0, 1)
-    c0 = UOp.const(dtypes.index, 0)
+    c0 = UOp.const(dtypes.weakint, 0)
     vc = UOp(Ops.CMPNE, dtypes.bool, (v, c0))
     c1 = UOp.const(dtypes.float, 1.0)
     out = UOp(Ops.WHERE, dtypes.float, (vc, c1, c1))
@@ -410,7 +410,7 @@ class TestUOpGraph(unittest.TestCase):
       d0 = UOp(Ops.PARAM, dt.ptr(), arg=0)
       v = d0.index(UOp.const(dtypes.int, 0))
       uops = to_uops_list([v.bitcast(dt)])
-      self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST]), 0, f"dtype = {dt}")
+      self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST and x.dtype is dt]), 0, f"dtype = {dt}")
 
   def test_sub_with_cast_folds(self):
     a = Variable("a", 0, 5)
@@ -469,16 +469,16 @@ class TestUOpGraph(unittest.TestCase):
     # mnist indexing with split reduceop
     # Make sure we are not doign math on the loaded index, which would promote it to long
     c0 = UOp(Ops.PARAM, dtypes.uchar.ptr(128000), arg=0, src=())
-    c1 = UOp.range(UOp.const(dtypes.index, 512), 1, AxisType.LOOP)
-    c2 = UOp.range(UOp.const(dtypes.index, 250), 2, AxisType.LOOP)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.LOOP)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.LOOP)
     c3 = UOp(Ops.PARAM, dtypes.int.ptr(512), arg=1, src=())
     c4 = c3.index(c1)
-    c5 = UOp.range(UOp.const(dtypes.index, 240), 0, AxisType.REDUCE)
-    c6 = ((c2*UOp.const(dtypes.index, 240))+c5)
+    c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
+    c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5)
     c7 = UOp(Ops.PARAM, dtypes.uchar.ptr(60000), arg=2, src=())
     c8 = c7.index(c6)
     c9 = ((c4<0).where((c4+60000), c4)!=c6.cast(dtypes.int)).where(0, c8.cast(dtypes.uint).cast(dtypes.uchar)).reduce(c5, arg=Ops.ADD)
-    c10 = c0.index(((c1*UOp.const(dtypes.index, 250))+c2)).store(c9).end(c1, c2)
+    c10 = c0.index(((c1*UOp.const(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2)
     uops = to_uops_list([c10])
     for u in uops:
       self.assertNotEqual(u.dtype, dtypes.long)
@@ -486,19 +486,19 @@ class TestUOpGraph(unittest.TestCase):
   def test_load_idx_no_math_on_loaded(self):
     # test the (x+y)<c pattern where x has loads - we shouldn't do math on loaded indices
     c0 = UOp(Ops.PARAM, dtypes.uchar.ptr(128000), arg=0, src=())
-    c1 = UOp.range(UOp.const(dtypes.index, 512), 1, AxisType.LOOP)
-    c2 = UOp.range(UOp.const(dtypes.index, 250), 2, AxisType.LOOP)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.LOOP)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.LOOP)
     c3 = UOp(Ops.PARAM, dtypes.int.ptr(512), arg=1, src=())
     c4 = c3.index(c1)  # c4 is a load
-    c5 = UOp.range(UOp.const(dtypes.index, 240), 0, AxisType.REDUCE)
-    c6 = ((c2*UOp.const(dtypes.index, 240))+c5)
+    c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
+    c6 = ((c2*UOp.const(dtypes.weakint, 240))+c5)
     c7 = UOp(Ops.PARAM, dtypes.uchar.ptr(60000), arg=2, src=())
     c8 = c7.index(c6)
     # (loaded + range) < const pattern - loaded value shouldn't be promoted to long
-    loaded_idx = c4.cast(dtypes.index)
-    comparison = (loaded_idx + c5) < UOp.const(dtypes.index, 60000)
+    loaded_idx = c4.cast(dtypes.weakint)
+    comparison = (loaded_idx + c5) < UOp.const(dtypes.weakint, 60000)
     c9 = comparison.where(c8.cast(dtypes.uint).cast(dtypes.uchar), 0).reduce(c5, arg=Ops.ADD)
-    c10 = c0.index(((c1*UOp.const(dtypes.index, 250))+c2)).store(c9).end(c1, c2)
+    c10 = c0.index(((c1*UOp.const(dtypes.weakint, 250))+c2)).store(c9).end(c1, c2)
     uops = to_uops_list([c10])
     for u in uops:
       self.assertNotEqual(u.dtype, dtypes.long)
@@ -816,5 +816,68 @@ class TestUOpTags(unittest.TestCase):
     g = graph_rewrite(g, pm_plus_1)
     assert g.ssimplify() == 6
 
+class TestUOpGetItem(unittest.TestCase):
+  def _placeholder(self, shape, dtype=dtypes.half):
+    return UOp.placeholder(shape, dtype, slot=0, addrspace=AddrSpace.LOCAL)
+
+  # full slices (no shrink)
+  def test_full_slice(self):
+    p = self._placeholder((64, 64))
+    self.assertEqual(p[:, :].shape, (64, 64))
+  def test_full_slice_explicit(self):
+    p = self._placeholder((64, 64))
+    self.assertEqual(p[0:64, 0:64].shape, (64, 64))
+
+  # partial slices (shrink)
+  def test_shrink_cols(self):
+    p = self._placeholder((64, 80))
+    self.assertEqual(p[:, :64].shape, (64, 64))
+  def test_shrink_rows(self):
+    p = self._placeholder((80, 64))
+    self.assertEqual(p[:64, :].shape, (64, 64))
+  def test_shrink_both(self):
+    p = self._placeholder((80, 80))
+    self.assertEqual(p[:64, :64].shape, (64, 64))
+  def test_shrink_start(self):
+    p = self._placeholder((64, 64))
+    self.assertEqual(p[8:, :].shape, (56, 64))
+  def test_shrink_start_and_end(self):
+    p = self._placeholder((64, 64))
+    self.assertEqual(p[8:56, 4:60].shape, (48, 56))
+
+  # mixed slice and index
+  def test_index_and_slice(self):
+    p = self._placeholder((64, 80))
+    r = UOp.range(64, 100)
+    result = p[r, :64]
+    self.assertEqual(result.shape, (64,))
+  def test_slice_and_index(self):
+    p = self._placeholder((80, 64))
+    r = UOp.range(64, 100)
+    result = p[:64, r]
+    self.assertEqual(result.shape, (64,))
+  def test_shrink_then_index(self):
+    p = self._placeholder((64, 80))
+    s = p[:, :64]
+    r = UOp.range(64, 100)
+    result = s[r]
+    self.assertEqual(result.shape, (64,))
+
+  # integer index (no slice)
+  def test_int_index(self):
+    p = self._placeholder((64, 64))
+    result = p[0]
+    self.assertEqual(result.shape, (64,))
+
+  # all slices should not create a bare INDEX
+  def test_all_slices_no_index(self):
+    p = self._placeholder((64, 80))
+    result = p[:, :64]
+    self.assertNotEqual(result.op, Ops.INDEX)
+  def test_all_full_slices_no_index(self):
+    p = self._placeholder((64, 64))
+    result = p[:, :]
+    self.assertNotEqual(result.op, Ops.INDEX)
+
 if __name__ == '__main__':
   unittest.main(verbosity=2)

test/null/test_uop_symbolic.py

@@ -12,13 +12,13 @@ from tinygrad.uop.validate import uops_to_z3
 def check_uop_against_string(self, v:UOp, s:str):
   sym_vars = {v.render():v for v in v.toposort() if v.op in (Ops.DEFINE_VAR, Ops.RANGE, Ops.SPECIAL)}
   s_eval = eval(s, sym_vars)
-  if isinstance(s_eval, int) and v.dtype==dtypes.index: s_eval = UOp.const(dtypes.index, s_eval)
+  if isinstance(s_eval, int) and v.dtype==dtypes.weakint: s_eval = UOp.const(dtypes.weakint, s_eval)
   elif isinstance(s_eval, (bool, int, float)): s_eval = UOp.const(dtypes.from_py(s_eval), s_eval)
   s_eval = graph_rewrite(s_eval, commutative, name="cannonicalize eval")
   self.assertIs(s_eval, v, f"eval did not match simplified: {s_eval} != {v.render()} for {s}")
 
-def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.index): return UOp.variable(name,min_val,max_val,dtype)
-def uconst(val): return UOp.const(dtypes.index, val)
+def Variable(name: str, min_val: ConstType, max_val: ConstType, dtype: DType=dtypes.weakint): return UOp.variable(name,min_val,max_val,dtype)
+def uconst(val): return UOp.const(dtypes.weakint, val)
 def usum(ops): return functools.reduce(lambda x,y: x+y, ops)
 def uand(ops): return functools.reduce(lambda x,y: x*y, ops)
 
@@ -245,12 +245,12 @@ class TestSymbolic(unittest.TestCase):
     self.assertEqual((Variable("x", -10, 0)%Variable("y", 1, 10))._min_max, (-9, 0))
 
   def test_range_div_its_symbolic_bound(self):
-    a = Variable("a", 1, 10, dtypes.index)
+    a = Variable("a", 1, 10, dtypes.weakint)
     ridx0 = UOp.range(a+2, 0)
     self.helper_test_variable(ridx0//(a+2), 0, 0, "0")
 
   def test_range_mod_its_symbolic_bound(self):
-    a = Variable("a", 1, 10, dtypes.index)
+    a = Variable("a", 1, 10, dtypes.weakint)
     ridx = UOp.range(a+2, 0)
     self.helper_test_variable(ridx%(a+2), 0, 11, "r0")
 
@@ -941,7 +941,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((numerator//denominator)<=0, 1, 1, "True")
 
   def test_symbolic_range_doesnt_collapse(self):
-    r0 = UOp.range((Variable("a", 1, 10)<5).cast(dtypes.index), 0)
+    r0 = UOp.range((Variable("a", 1, 10)<5).cast(dtypes.weakint), 0)
     self.helper_test_variable(r0, 0, 0, "r0")
 
   def test_const_reciprocal(self):
@@ -1202,16 +1202,16 @@ class TestInvalidIndex(unittest.TestCase):
     self.assertIs((UOp.invalid()<Variable("a",0,10)).simplify().dtype, dtypes.bool)
 
   def test_alu_invalid_vconst(self):
-    c1 = UOp.const(dtypes.index.vec(4), (1, 1, Invalid, Invalid))
-    c2 = UOp.const(dtypes.index.vec(4), (1, Invalid, 1, 1))
-    self.assertIs((c1+c2).simplify(), UOp.const(dtypes.index.vec(4), (2, Invalid, Invalid, Invalid)))
+    c1 = UOp.const(dtypes.weakint.vec(4), (1, 1, Invalid, Invalid))
+    c2 = UOp.const(dtypes.weakint.vec(4), (1, Invalid, 1, 1))
+    self.assertIs((c1+c2).simplify(), UOp.const(dtypes.weakint.vec(4), (2, Invalid, Invalid, Invalid)))
 
 class TestStoreLoadFolding(unittest.TestCase):
   """Tests for store(index, load(index)) -> NOOP rule. This rule matches patterns that EMERGE during simplification."""
   def test_store_load_folding(self):
     # store(idx, load(idx)) -> NOOP, including emergent patterns like store(idx, load(idx) + 0)
     buf = UOp(Ops.PARAM, dtypes.int.ptr(), arg=0)
-    index = buf.index(UOp.const(dtypes.index, 0))
+    index = buf.index(UOp.const(dtypes.weakint, 0))
     # Direct: store(idx, load(idx)) -> NOOP
     self.assertEqual(graph_rewrite(index.store(index.load()), sym).op, Ops.NOOP)
     # Emergent: store(idx, load(idx) + 0) -> store(idx, load(idx)) -> NOOP
@@ -1250,10 +1250,10 @@ class TestGatedUopGivenValid(unittest.TestCase):
 
     idx0 = (r0 + uconst(-1)) // uconst(3)
     idx1 = r0 % uconst(3)
-    idx:UOp = (r0 < 3).where(UOp(Ops.VECTORIZE, dtypes.index.vec(2), (idx0, idx1)), UOp.invalid())
+    idx:UOp = (r0 < 3).where(UOp(Ops.VECTORIZE, dtypes.weakint.vec(2), (idx0, idx1)), UOp.invalid())
     idx = graph_rewrite(idx, pm_simplify_valid)
     # NOTE: independent simplification: (r0-1)//3 -> 0, r0%3 -> r0 when r0 in [0,2]
-    expected_vec = UOp(Ops.VECTORIZE, dtypes.index.vec(2), (uconst(0), r0))
+    expected_vec = UOp(Ops.VECTORIZE, dtypes.weakint.vec(2), (uconst(0), r0))
     self.assertEqual(idx, (r0 < 3).where(expected_vec, UOp.invalid()))
 
 class TestRangeSplitting(unittest.TestCase):

test/null/test_uop_vmin_vmax.py

@@ -144,7 +144,7 @@ class TestVminVmaxProperties(unittest.TestCase):
     self.assertNotEqual(i.vmin, i.vmax)
 
   def test_vmin_vmax_invalid_vconst(self):
-    x = UOp.const(dtypes.index.vec(4), (0, 4, Invalid, Invalid))
+    x = UOp.const(dtypes.weakint.vec(4), (0, 4, Invalid, Invalid))
     self.assertLess(x.vmin, 0)
     self.assertGreater(x.vmax, 4)
 

test/null/test_validate_oob.py

@@ -126,7 +126,7 @@ class TestValidateOOB(unittest.TestCase):
       buf0 = UOp(Ops.PARAM, dtypes.int.ptr(16), (), 0)
       buf1 = UOp(Ops.PARAM, dtypes.int.ptr(64), (), 1)
       r = UOp.range(42, 0, AxisType.GLOBAL)
-      ld0 = buf0.index(r.valid(r < 8), ptr=True).load(dtype=dtypes.int).cast(dtypes.index)
+      ld0 = buf0.index(r.valid(r < 8), ptr=True).load(dtype=dtypes.int).cast(dtypes.weakint)
       to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32)), ptr=True).load(dtype=dtypes.int)])  # valid
       with self.assertRaises(RuntimeError):
         to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 64)), ptr=True).load(dtype=dtypes.int)])  # oob
@@ -135,7 +135,7 @@ class TestValidateOOB(unittest.TestCase):
     with Context(CHECK_OOB=1, SPEC=2):
       buf_bool = UOp(Ops.PARAM, dtypes.bool.ptr(16), (), 0)
       buf_int = UOp(Ops.PARAM, dtypes.int.ptr(8), (), 1)
-      gidx = UOp(Ops.SPECIAL, dtypes.index, (UOp.const(dtypes.index, 16),), "gidx0")
+      gidx = UOp(Ops.SPECIAL, dtypes.weakint, (UOp.const(dtypes.weakint, 16),), "gidx0")
       ld_bool = buf_bool.index(gidx, ptr=True).load()
       with self.assertRaises(RuntimeError):
         to_uops_list([buf_int.index(gidx.valid(ld_bool), ptr=True).load()])  # gidx 0..15, buf_int size 8

test/null/test_viz.py

@@ -642,5 +642,194 @@ class TestVizMemoryLayout(BaseTestViz):
     users = profile["layout"][f"{a.device} Memory"]["events"].pop()["arg"]["users"]
     self.assertEqual(len(programs), len(set(users)), n)
 
+from tinygrad.uop.ops import KernelInfo
+from tinygrad.viz.serve import amdgpu_cfg
+from tinygrad.renderer.amd.dsl import s
+from tinygrad.runtime.autogen.amd.rdna3.ins import (s_add_u32, s_branch, s_cbranch_execz, s_cbranch_scc0, s_cbranch_scc1, s_cmp_eq_i32,
+                                                    s_cmp_eq_u64, s_code_end, s_endpgm, s_mov_b32, s_nop)
+from extra.gemm.amd_asm_matmul import Kernel
+from tinygrad.renderer.cstyle import AMDHIPRenderer
+
+class TestCfg(unittest.TestCase):
+  def setUp(self): self.arch = "gfx1100"
+
+  def get_cfg(self, name:str, k:Kernel):
+    insts = k.finalize()
+    def fxn(out:UOp) -> UOp:
+      lidx = UOp.special(1, "lidx0")
+      gidx = UOp.special(1, "gidx0")
+      sink = UOp.sink(out.base, lidx, gidx, arg=KernelInfo(name=name))
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="NULL"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+    with Context(EMULATE="AMD"):
+      out = Tensor.custom_kernel(Tensor.empty(1, device="NULL"), fxn=fxn)[0]
+      # TODO: uncomment the better version once EMULATE works in Context
+      #prg = out.schedule()[-1].lower().prg.p
+      prg = get_program(out.schedule()[-1].ast, AMDHIPRenderer(self.arch))
+      return amdgpu_cfg(prg.lib, self.arch)
+
+  def test_simple(self):
+    k = Kernel(arch=self.arch)
+    k.label("entry")
+    k.emit(s_branch(), target="bb1")
+    k.label("bb1")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    cfg = self.get_cfg("simple", k)["data"]
+    self.assertEqual(len(cfg["blocks"]), 2)
+
+  def test_diamond(self):
+    k = Kernel(arch=self.arch)
+    k.label("entry")
+    k.emit(s_mov_b32(s[0], 0))
+    k.emit(s_mov_b32(s[1], 0))
+    k.emit(s_cmp_eq_u64(s[0:1], 0))
+    k.emit(s_cbranch_scc1(), target="if")
+    k.emit(s_branch(), target="else")
+    k.label("if")
+    k.emit(s_nop(1))
+    k.emit(s_branch(), target="end")
+    k.label("else")
+    k.emit(s_nop(0))
+    k.label("end")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    ret = self.get_cfg("diamond", k)
+    cfg = ret["data"]
+    self.assertEqual(len(cfg["blocks"]), 5)
+    edge_count = sum(len(v) for v in cfg["paths"].values())
+    self.assertEqual(edge_count, 5)
+    references:dict[str, list[str]] = {}
+    for pc, tokens in cfg["pc_tokens"].items():
+      for t in tokens:
+        for key in t["keys"]: references.setdefault(key, []).append(pc)
+    self.assertEqual(len(references["r0"]), 2)
+    insts = [cfg["pc_tokens"][pc][0]["st"] for pc in references["r0"]]
+    self.assertEqual(insts, ['s_mov_b32', 's_cmp_eq_u64'])
+    end_block = [" ".join(t["st"] for t in cfg["pc_tokens"][pc]) for pc in list(cfg["blocks"].values())[-1]]
+    code_line = ret["src"].splitlines()[-1]
+    self.assertEqual(len(end_block), 2)
+    for st in [end_block[-1], code_line]:
+      assert st.startswith("s_code_end") and st.endswith("x)"), st
+
+  def test_loop(self):
+    k = Kernel(arch=self.arch)
+    k.label("entry")
+    k.emit(s_mov_b32(s[1], 4))
+    k.label("loop")
+    k.emit(s_add_u32(s[1], s[1], -1))
+    k.emit(s_cmp_eq_i32(s[1], 0))
+    k.emit(s_cbranch_scc0(), target="loop")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    self.get_cfg("simple_loop", k)
+
+  def test_loop_branch(self):
+    k = Kernel(arch=self.arch)
+    k.label("entry")
+    k.emit(s_mov_b32(s[1], 4))
+    k.label("loop")
+    k.emit(s_add_u32(s[1], s[1], -1))
+    k.emit(s_cmp_eq_i32(s[1], 2))
+    k.emit(s_cbranch_scc1(), target="cond")
+    k.emit(s_branch(), target="cont")
+    k.label("cond")
+    k.emit(s_add_u32(s[1], s[1], -2))
+    k.label("cont")
+    k.emit(s_cmp_eq_i32(s[1], 0))
+    k.emit(s_cbranch_scc0(), target="loop")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    self.get_cfg("loop_if", k)
+
+  def test_loop_break(self):
+    k = Kernel(arch=self.arch)
+    k.label("entry")
+    k.emit(s_mov_b32(s[1], 8))
+    k.label("loop")
+    k.emit(s_add_u32(s[1], s[1], -1))
+    k.emit(s_cmp_eq_i32(s[1], 5))
+    k.emit(s_cbranch_scc1(), target="break")
+    k.emit(s_cmp_eq_i32(s[1], 0))
+    k.emit(s_cbranch_scc0(), target="loop")
+    k.label("break")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    self.get_cfg("loop_break", k)
+
+  def test_switch(self):
+    k = Kernel(arch=self.arch)
+    k.label("entry")
+    k.emit(s_cmp_eq_i32(s[0], 0))
+    k.emit(s_cbranch_scc1(), target="case0")
+    k.emit(s_cmp_eq_i32(s[0], 1))
+    k.emit(s_cbranch_scc1(), target="case1")
+    k.emit(s_branch(), target="case2")
+    k.label("case0")
+    k.emit(s_nop(0))
+    k.emit(s_branch(), target="join")
+    k.label("case1")
+    k.emit(s_nop(1))
+    k.emit(s_branch(), target="join")
+    k.label("case2")
+    k.emit(s_nop(2))
+    k.emit(s_branch(), target="join")
+    k.label("join")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    self.get_cfg("switch_case", k)
+
+  def test_ping_pong(self):
+    k = Kernel(arch=self.arch)
+    k.label("entry")
+    k.emit(s_cmp_eq_i32(s[0], 0))
+    k.emit(s_cbranch_scc1(), target="ping")
+    k.emit(s_branch(), target="pong")
+    k.label("ping")
+    k.emit(s_cmp_eq_i32(s[1], 0))
+    k.emit(s_cbranch_scc1(), target="pong")
+    k.emit(s_branch(), target="end")
+    k.label("pong")
+    k.emit(s_cmp_eq_i32(s[2], 0))
+    k.emit(s_cbranch_scc1(), target="ping")
+    k.label("end")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    self.get_cfg("ping_pong", k)
+
+  def test_colored_blocks(self):
+    N = 10
+    k = Kernel(arch=self.arch)
+    k.label("entry")
+    k.emit(s_branch(), target="init0")
+    for i in range(N):
+      loop = f"loop{i}"
+      k.label(f"init{i}")
+      k.emit(s_mov_b32(s[1], i + 1))
+      k.emit(s_branch(), target=loop)
+      k.label(loop)
+      k.emit(s_nop(i & 7))
+      k.emit(s_add_u32(s[1], s[1], -1))
+      k.emit(s_cmp_eq_i32(s[1], 0))
+      k.emit(s_cbranch_scc0(), target=loop)
+      k.emit(s_branch(), target=f"init{i+1}" if i + 1 < N else "end")
+    k.label("end")
+    k.emit(s_endpgm())
+    k.emit(s_code_end())
+    self.get_cfg("test_colored_blocks", k)
+
+  def test_jump_back_to_end(self):
+    k = Kernel(arch=self.arch)
+    k.label("entry")
+    k.emit(s_mov_b32(s[1], 2))
+    k.emit(s_cbranch_execz(), target="loop")
+    k.label("end")
+    k.emit(s_endpgm())
+    k.label("loop")
+    k.emit(s_add_u32(s[1], s[1], -1))
+    k.emit(s_cmp_eq_i32(s[1], 0))
+    k.emit(s_branch(), target="end")
+    k.emit(s_code_end())
+    self.get_cfg("jump_back_to_end", k)
+
 if __name__ == "__main__":
   unittest.main()

test/test_tiny.py

@@ -1,7 +1,7 @@
 # basic self-contained tests of the external functionality of tinygrad
 import unittest, random
 from tinygrad import Tensor, Context, Variable, TinyJit, dtypes, Device, nn
-from tinygrad.helpers import IMAGE, CI, getenv
+from tinygrad.helpers import CI, getenv
 
 class TestTiny(unittest.TestCase):
 
@@ -39,14 +39,14 @@ class TestTiny(unittest.TestCase):
     out = Tensor.ones(N).contiguous().sum()
     self.assertEqual(out.item(), N)
 
-  def test_gemm(self, N=getenv("GEMM_N", 64), out_dtype=dtypes.float):
+  def test_gemm(self, N=getenv("GEMM_N", 64)):
     a = Tensor.ones(N,N).contiguous()
     b = Tensor.eye(N).contiguous()
     lst = (out:=a@b).tolist()
     for y in range(N):
       for x in range(N):
         self.assertEqual(lst[y][x], 1.0, msg=f"mismatch at ({y},{x})")
-    if IMAGE < 2: self.assertEqual(out.dtype, out_dtype)
+    self.assertEqual(out.dtype, dtypes.float)
 
   def test_gemv(self, N=getenv("GEMV_N", 64), out_dtype=dtypes.float):
     a = Tensor.ones(1,N).contiguous()
@@ -54,7 +54,7 @@ class TestTiny(unittest.TestCase):
     lst = (out:=a@b).tolist()
     for x in range(N):
       self.assertEqual(lst[0][x], 1.0, msg=f"mismatch at {x}")
-    if IMAGE < 2: self.assertEqual(out.dtype, out_dtype)
+    self.assertEqual(out.dtype, out_dtype)
 
   # *** randomness ***
 
@@ -149,7 +149,7 @@ class TestTiny(unittest.TestCase):
 
   @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
   def test_image(self):
-    with Context(IMAGE=2): self.test_gemm(N=4, out_dtype=dtypes.imagef((4, 1, 4)))
+    with Context(IMAGE=1): self.test_gemm(N=64)
 
   def test_beam_image(self):
     with Context(BEAM=1, IGNORE_BEAM_CACHE=1): self.test_image()

test/testextra/test_cfg_viz.py

@@ -1,214 +0,0 @@
-# ruff: noqa: F405, F403
-# allow define from star imports
-
-import unittest
-
-from tinygrad import Device, Tensor
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.viz.serve import amdgpu_cfg
-
-from tinygrad.runtime.autogen.amd.rdna3.ins import *
-from tinygrad.renderer.amd.dsl import s
-
-# TODO: this belongs to the dsl infrastructure
-from extra.gemm.amd_asm_matmul import Kernel
-
-def run_asm(name:str, k:Kernel):
-  insts = k.finalize()
-  def fxn(out:UOp) -> UOp:
-    lidx = UOp.special(1, "lidx0")
-    gidx = UOp.special(1, "gidx0")
-    sink = UOp.sink(out.base, lidx, gidx, arg=KernelInfo(name=name))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
-  out = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
-  ei = out.schedule()[-1].lower()
-  ei.run()
-  return ei
-
-@unittest.skipUnless(Device.DEFAULT == "AMD", "only on AMD")
-class TestCfg(unittest.TestCase):
-  def setUp(self):
-    self.arch = Device["AMD"].arch
-    if not any(self.arch.startswith(a) for a in {"gfx11", "gfx12"}):
-      self.skipTest(f"tests written for RDNA, got arch {self.arch}")
-
-  def test_simple(self):
-    k = Kernel(arch=Device["AMD"].arch)
-    k.label("entry")
-    k.emit(s_branch(), target="bb1")
-    k.label("bb1")
-    k.emit(s_endpgm())
-    k.emit(s_code_end())
-    run_asm("simple", k)
-
-  def test_diamond(self):
-    k = Kernel(arch=Device["AMD"].arch)
-    k.label("entry")
-    k.emit(s_mov_b32(s[0], 0))
-    k.emit(s_mov_b32(s[1], 0))
-    k.emit(s_cmp_eq_u64(s[0:1], 0))
-    k.emit(s_cbranch_scc1(), target="if")
-    k.emit(s_branch(), target="else")
-    k.label("if")
-    k.emit(s_nop(1))
-    k.emit(s_branch(), target="end")
-    k.label("else")
-    k.emit(s_nop(0))
-    k.label("end")
-    k.emit(s_endpgm())
-    k.emit(s_code_end())
-    ei = run_asm("diamond", k)
-    ret = amdgpu_cfg(ei.prg.p.lib, self.arch)
-    cfg = ret["data"]
-    self.assertEqual(len(cfg["blocks"]), 5)
-    edge_count = sum(len(v) for v in cfg["paths"].values())
-    self.assertEqual(edge_count, 5)
-    references:dict[str, list[str]] = {}
-    for pc, tokens in cfg["pc_tokens"].items():
-      for t in tokens:
-        for key in t["keys"]: references.setdefault(key, []).append(pc)
-    self.assertEqual(len(references["r0"]), 2)
-    insts = [cfg["pc_tokens"][pc][0]["st"] for pc in references["r0"]]
-    self.assertEqual(insts, ['s_mov_b32', 's_cmp_eq_u64'])
-    end_block = [" ".join(t["st"] for t in cfg["pc_tokens"][pc]) for pc in list(cfg["blocks"].values())[-1]]
-    code_line = ret["src"].splitlines()[-1]
-    self.assertEqual(len(end_block), 2)
-    for st in [end_block[-1], code_line]:
-      assert st.startswith("s_code_end") and st.endswith("x)"), st
-
-  def test_loop(self):
-    k = Kernel(arch=Device["AMD"].arch)
-    k.label("entry")
-    k.emit(s_mov_b32(s[1], 4))
-    k.label("loop")
-    k.emit(s_add_u32(s[1], s[1], -1))
-    k.emit(s_cmp_eq_i32(s[1], 0))
-    k.emit(s_cbranch_scc0(), target="loop")
-    k.emit(s_endpgm())
-    k.emit(s_code_end())
-    run_asm("simple_loop", k)
-
-  def test_loop_branch(self):
-    k = Kernel(arch=Device["AMD"].arch)
-    k.label("entry")
-    k.emit(s_mov_b32(s[1], 4))
-    k.label("loop")
-    k.emit(s_add_u32(s[1], s[1], -1))
-    k.emit(s_cmp_eq_i32(s[1], 2))
-    k.emit(s_cbranch_scc1(), target="cond")
-    k.emit(s_branch(), target="cont")
-    k.label("cond")
-    k.emit(s_add_u32(s[1], s[1], -2))
-    k.label("cont")
-    k.emit(s_cmp_eq_i32(s[1], 0))
-    k.emit(s_cbranch_scc0(), target="loop")
-    k.emit(s_endpgm())
-    k.emit(s_code_end())
-    run_asm("loop_if", k)
-
-  def test_loop_break(self):
-    k = Kernel(arch=Device["AMD"].arch)
-    k.label("entry")
-    k.emit(s_mov_b32(s[1], 8))
-    k.label("loop")
-    k.emit(s_add_u32(s[1], s[1], -1))
-    k.emit(s_cmp_eq_i32(s[1], 5))
-    k.emit(s_cbranch_scc1(), target="break")
-    k.emit(s_cmp_eq_i32(s[1], 0))
-    k.emit(s_cbranch_scc0(), target="loop")
-    k.label("break")
-    k.emit(s_endpgm())
-    k.emit(s_code_end())
-    run_asm("loop_break", k)
-
-  def test_switch(self):
-    k = Kernel(arch=Device["AMD"].arch)
-    k.label("entry")
-    k.emit(s_cmp_eq_i32(s[0], 0))
-    k.emit(s_cbranch_scc1(), target="case0")
-    k.emit(s_cmp_eq_i32(s[0], 1))
-    k.emit(s_cbranch_scc1(), target="case1")
-    k.emit(s_branch(), target="case2")
-    k.label("case0")
-    k.emit(s_nop(0))
-    k.emit(s_branch(), target="join")
-    k.label("case1")
-    k.emit(s_nop(1))
-    k.emit(s_branch(), target="join")
-    k.label("case2")
-    k.emit(s_nop(2))
-    k.emit(s_branch(), target="join")
-    k.label("join")
-    k.emit(s_endpgm())
-    k.emit(s_code_end())
-    run_asm("switch_case", k)
-
-  def test_ping_pong(self):
-    k = Kernel(arch=Device["AMD"].arch)
-    k.label("entry")
-    k.emit(s_cmp_eq_i32(s[0], 0))
-    k.emit(s_cbranch_scc1(), target="ping")
-    k.emit(s_branch(), target="pong")
-    k.label("ping")
-    k.emit(s_cmp_eq_i32(s[1], 0))
-    k.emit(s_cbranch_scc1(), target="pong")
-    k.emit(s_branch(), target="end")
-    k.label("pong")
-    k.emit(s_cmp_eq_i32(s[2], 0))
-    k.emit(s_cbranch_scc1(), target="ping")
-    k.label("end")
-    k.emit(s_endpgm())
-    k.emit(s_code_end())
-    run_asm("ping_pong", k)
-
-  def test_colored_blocks(self):
-    N = 10
-    k = Kernel(arch=Device["AMD"].arch)
-    k.label("entry")
-    k.emit(s_branch(), target="init0")
-    for i in range(N):
-      loop = f"loop{i}"
-      k.label(f"init{i}")
-      k.emit(s_mov_b32(s[1], i + 1))
-      k.emit(s_branch(), target=loop)
-      k.label(loop)
-      k.emit(s_nop(i & 7))
-      k.emit(s_add_u32(s[1], s[1], -1))
-      k.emit(s_cmp_eq_i32(s[1], 0))
-      k.emit(s_cbranch_scc0(), target=loop)
-      k.emit(s_branch(), target=f"init{i+1}" if i + 1 < N else "end")
-    k.label("end")
-    k.emit(s_endpgm())
-    k.emit(s_code_end())
-    run_asm("test_colored_blocks", k)
-
-  def test_jump_back_to_end(self):
-    k = Kernel(arch=Device["AMD"].arch)
-    k.label("entry")
-    k.emit(s_mov_b32(s[1], 2))
-    k.emit(s_cbranch_execz(), target="loop")
-    k.label("end")
-    k.emit(s_endpgm())
-    k.label("loop")
-    k.emit(s_add_u32(s[1], s[1], -1))
-    k.emit(s_cmp_eq_i32(s[1], 0))
-    k.emit(s_branch(), target="end")
-    k.emit(s_code_end())
-    run_asm("jump_back_to_end", k)
-
-  def test_hit_count(self):
-    k = Kernel(arch=Device["AMD"].arch)
-    k.label("entry")
-    k.emit(s_mov_b32(s[1], 1))
-    k.emit(s_branch(), target="alt")
-    k.label("continue")
-    k.emit(s_mov_b32(s[2], 2))
-    k.emit(s_add_u32(s[1], s[1], s[2]))
-    k.label("alt")
-    k.emit(s_add_u32(s[1], s[1], -1))
-    k.emit(s_endpgm())
-    k.emit(s_code_end())
-    run_asm("test_hit_count", k)
-
-if __name__ == "__main__":
-  unittest.main()

test/testextra/test_export_model.py

@@ -1,6 +1,7 @@
 import unittest
 from extra.export_model import export_model, EXPORT_SUPPORTED_DEVICE
-from tinygrad.tensor import Tensor, Device
+from tinygrad.tensor import Tensor
+from tinygrad.device import Device
 from tinygrad import dtypes
 import json
 

test/unit/test_dtype_spec.py

@@ -1,8 +1,8 @@
 import unittest, math, subprocess
-from tinygrad.tensor import Tensor, dtypes, Device
-from tinygrad.dtype import DType, DTYPES_DICT
-from tinygrad.device import is_dtype_supported
-from tinygrad.helpers import getenv, DEBUG
+from tinygrad.tensor import Tensor
+from tinygrad.dtype import dtypes, DType, DTYPES_DICT
+from tinygrad.device import Device, is_dtype_supported
+from tinygrad.helpers import getenv, DEBUG, EMULATED_DTYPES
 from test.helpers import slow
 from hypothesis import given, settings, strategies as strat
 import numpy as np
@@ -24,8 +24,13 @@ def _assert_eq(tensor:Tensor, target_dtype:DType, target, tol_target_dtype:float
   if DEBUG >= 2: print(tensor.numpy())
   try:
     assert tensor.dtype == target_dtype
-    np.testing.assert_allclose(tensor.numpy(), target, rtol={dtypes.float16:1e-3, dtypes.bfloat16:1e-2,
-      dtypes.fp8e4m3:1e-1, dtypes.fp8e5m2:5e-1, dtypes.fp8e4m3fnuz:1e-1, dtypes.fp8e5m2fnuz:5e-1}.get(target_dtype, tol_target_dtype))
+    # denormals are zero
+    if target_dtype in dtypes.floats and (not is_dtype_supported(target_dtype) or target_dtype in EMULATED_DTYPES.tolist(dtypes)):
+      fe, fm = dtypes.finfo(target_dtype)
+      kwargs = {"atol":2 ** (2 - (1 << (fe - 1))), "rtol": 2 ** (-fm)}
+    else: kwargs = {"rtol": {dtypes.float16:1e-3, dtypes.bfloat16:1e-2, dtypes.fp8e4m3:1e-1, dtypes.fp8e5m2:5e-1,
+                             dtypes.fp8e4m3fnuz:1e-1, dtypes.fp8e5m2fnuz:5e-1}.get(target_dtype, tol_target_dtype)}
+    np.testing.assert_allclose(tensor.numpy(), target, **kwargs)
 
   except AssertionError as e:
     raise AssertionError(f"\ntensor {tensor.numpy()} dtype {tensor.dtype} does not match target {target} with dtype {target_dtype}") from e

test/unit/test_function.py

@@ -1,7 +1,7 @@
 import numpy as np
 import unittest
 from tinygrad.function import function
-from tinygrad import Tensor
+from tinygrad import Tensor, GlobalCounters
 from tinygrad.uop.ops import UOp
 
 class TestFunction(unittest.TestCase):
@@ -22,7 +22,7 @@ class TestFunction(unittest.TestCase):
 
   def test_implicit(self):
     inp = Tensor([7,8,9])
-    @function
+    @function(allow_implicit=True)
     def f(a:Tensor, b:Tensor) -> Tensor: return a+b+inp
 
     a = Tensor([1,2,3])
@@ -31,7 +31,7 @@ class TestFunction(unittest.TestCase):
 
   def test_implicit_same_as_input(self):
     inp = Tensor([7,8,9])
-    @function
+    @function(allow_implicit=True)
     def f(a:Tensor, b:Tensor) -> Tensor: return a+b+inp
 
     a = Tensor([1,2,3])
@@ -39,11 +39,11 @@ class TestFunction(unittest.TestCase):
 
   def test_implicit_2(self):
     inp = Tensor([7,8,9])
-    @function
+    @function(allow_implicit=True)
     def f(a:Tensor, b:Tensor) -> Tensor:
       return a+b+inp
     inp2 = Tensor([7,8,10])
-    @function
+    @function(allow_implicit=True)
     def g(a:Tensor, b:Tensor) -> Tensor:
       return a+b+inp2
 
@@ -57,7 +57,7 @@ class TestFunction(unittest.TestCase):
 
   def test_implicit_unrealized(self):
     inp = Tensor([1,2,3]) + Tensor([4,5,6])
-    @function
+    @function(allow_implicit=True)
     def f(a:Tensor) -> Tensor: return a + inp
 
     np.testing.assert_equal(f(Tensor([10,20,30])).numpy(), [15,27,39])
@@ -103,7 +103,7 @@ class TestFunction(unittest.TestCase):
   def test_grad_implicit(self):
     w = Tensor([1., 2., 3.], requires_grad=True)
     w.realize() # TODO: this is required
-    @function
+    @function(allow_implicit=True)
     def f(x:Tensor) -> Tensor: return x * w
 
     x = Tensor([4., 5., 6.])
@@ -112,7 +112,7 @@ class TestFunction(unittest.TestCase):
 
   def test_symbolic_index(self):
     table = Tensor([10,20,30,40]).contiguous().realize()
-    @function
+    @function(allow_implicit=True)
     def f(x:Tensor, start_pos:int|UOp) -> Tensor:
       return x + table[start_pos]
 
@@ -129,9 +129,9 @@ class TestFunction(unittest.TestCase):
 
   def test_nested_calls(self):
     w = Tensor([10., 20., 30.])
-    @function
+    @function(allow_implicit=True)
     def f(a:Tensor) -> Tensor: return a + w
-    @function
+    @function(allow_implicit=True)
     def g(a:Tensor) -> Tensor: return a * w
 
     a = Tensor([1., 2., 3.])
@@ -139,9 +139,9 @@ class TestFunction(unittest.TestCase):
 
   def test_nested_calls_backward(self):
     w = Tensor([[1., 2.], [3., 4.]]).contiguous().realize()
-    @function
+    @function(allow_implicit=True)
     def inner(x:Tensor) -> Tensor: return x + w
-    @function
+    @function(allow_implicit=True)
     def outer(a:Tensor, b:Tensor) -> Tensor: return inner(a.reshape(1,2) + b.reshape(1,2))
 
     a = Tensor([1., 2.], requires_grad=True)
@@ -178,7 +178,7 @@ class TestFunction(unittest.TestCase):
       def __init__(self): self.w = Tensor([10,20,30])
       def __call__(self, x:Tensor) -> Tensor: return x + self.w
     foo = Foo()
-    f = function(foo)
+    f = function(foo, allow_implicit=True)
     np.testing.assert_equal(f(Tensor([1,2,3])).numpy(), [11,22,33])
     assert f(Tensor([1,2,3])).uop.src[0].arg.name.endswith("Foo")
 
@@ -224,6 +224,20 @@ class TestFunction(unittest.TestCase):
     s = State()
     np.testing.assert_equal(s(Tensor([[5., 6.], [7., 8.]])).numpy(), [[5., 6.], [7., 8.]])
 
+  def test_single_after_store(self):
+    """AFTER(buf, STORE(view, data)) should write data through the view into buf, same as the double-after pattern."""
+    @function
+    def f(buf:Tensor, x:Tensor, start_pos:int|UOp) -> Tensor:
+      slice_uop = buf[:, start_pos:start_pos+1].uop
+      assigned = Tensor(buf.uop.after(slice_uop.store(x.uop)))
+      return assigned
+
+    buf = Tensor.zeros(2, 8).contiguous().realize()
+    x = Tensor([[1.], [2.]]).realize()
+    v = UOp.variable("sp", 0, 7)
+    r0 = f(buf, x, v.bind(0)).numpy()
+    np.testing.assert_equal(r0, [[1.,0.,0.,0.,0.,0.,0.,0.], [2.,0.,0.,0.,0.,0.,0.,0.]])
+
   @unittest.expectedFailure
   def test_assign_slice(self):
     @function
@@ -267,7 +281,7 @@ class TestFunctionMulti(unittest.TestCase):
   def test_grad_implicit_multi(self):
     w = Tensor([1., 2., 3., 4.], requires_grad=True).shard(self.devices_2, axis=None)
     w.realize()
-    @function
+    @function(allow_implicit=True)
     def f(x:Tensor) -> Tensor: return x * w
 
     x = Tensor([4., 5., 6., 7.]).shard(self.devices_2, axis=None)
@@ -324,7 +338,7 @@ class TestFunctionMulti(unittest.TestCase):
     devices_4 = tuple(f"CPU:{i}" for i in range(4))
     w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(devices_4, axis=None)
     w.realize()
-    @function
+    @function(allow_implicit=True)
     def f(x:Tensor) -> Tensor: return x @ w
 
     x = Tensor(np.arange(16).reshape(8,2).astype(np.float32), requires_grad=True).shard(devices_4, axis=0)
@@ -337,7 +351,7 @@ class TestFunctionMulti(unittest.TestCase):
     w.realize()
     # pre-init grads like the training loop does
     w.grad = w.zeros_like().contiguous().realize()
-    @function
+    @function(allow_implicit=True)
     def f(x:Tensor) -> Tensor: return x @ w
 
     expected = np.ones((8,2)) @ np.array([[1,3],[2,4]])
@@ -384,5 +398,31 @@ class TestFunctionTuple(unittest.TestCase):
     np.testing.assert_allclose(x.grad.numpy(), [1., 1., 1.])
     np.testing.assert_allclose(y.grad.numpy(), [1., 1., 1.])
 
+class TestFunctionGrad(unittest.TestCase):
+  def test_function_grad_ops(self, precompile=False, precompile_backward=False):
+    N = 64
+    x = Tensor.ones(N,N).contiguous()
+    w1 = Tensor.ones(N,N, requires_grad=True).contiguous()
+    w2 = Tensor.ones(N,N, requires_grad=True).contiguous()
+    w3 = Tensor.ones(N,N, requires_grad=True).contiguous()
+    ref = Tensor.ones(N,N).contiguous()
+    Tensor.realize(x, w1, w2, w3, ref)
+    @function(precompile=precompile, precompile_backward=precompile_backward)
+    def f(x, w1, w2, w3) -> tuple[Tensor, ...]:
+      p1 = x@w1
+      p2 = p1@w2
+      p3 = p2@w3
+      return p1, p2, p3, p3.contiguous()
+    ret = f(x, w1, w2, w3)[-1]
+    loss = (ret-ref).square().mean().backward()
+    print("RESET")
+    GlobalCounters.reset()
+    loss.realize(w1.grad, w2.grad, w3.grad)
+    print(GlobalCounters.global_ops, GlobalCounters.global_mem)
+    self.assertLessEqual(GlobalCounters.global_ops, 4739073)
+  def test_function_grad_ops_precompile(self): self.test_function_grad_ops(precompile=True)
+  def test_function_grad_ops_precompile_backward(self):
+    self.test_function_grad_ops(precompile=True, precompile_backward=True)
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_gguf.py

@@ -148,7 +148,8 @@ class TestGGUFGEMV(unittest.TestCase):
 
     x = rng.standard_normal(cols).astype(np.float32)
     np.testing.assert_allclose((tensors["weight"] @ Tensor(x)).numpy(), ref @ x, atol=1e-2, rtol=1e-2)
-    np.testing.assert_equal(tensors["weight"].numpy(), ref)
+    # can only expect the weights to be identical if we really support float16 (ie. not decompositions)
+    if is_dtype_supported(dtypes.half): np.testing.assert_equal(tensors["weight"].numpy(), ref)
     assert np.isfinite(ref).all() and np.isfinite(tensors["weight"].numpy()).all(), f"{qtype.name} has NaN/Inf"
 
   def test_gguf_gemv_q8_0(self): self._test_gguf_gemv(GGMLQuantizationType.Q8_0)

test/unit/test_gradient.py

@@ -12,10 +12,10 @@ class TestTensorGradient(unittest.TestCase):
     self.assertListEqual(dx.tolist(), [[2.0, 2.0, 2.0], [0.0, 0.0, 0.0], [-2.0, -2.0, -2.0]])
     self.assertListEqual(dy.tolist(), [[1.0, 1.0, 1.0]])
 
-  def test_raises(self):
+  def test_zero_if_not_used(self):
     x = Tensor([1.0, 2.0, 3.0])
     w = Tensor.randn((3,))
-    with self.assertRaises(RuntimeError): x.sum().gradient(w)
+    self.assertListEqual(x.sum().gradient(w)[0].tolist(), [0.0, 0.0, 0.0])
 
   def test_with_custom_gradient(self):
     x = Tensor([1.0, 2.0, 3.0])

test/unit/test_llm_server.py

@@ -11,7 +11,7 @@ class TestTransformerGenerate(unittest.TestCase):
                         norm_eps=1e-5, vocab_size=100, head_dim=32, rope_theta=10000.0, max_context=32)
 
     captured_inputs = []
-    def mock_call(self, tokens, start_pos):
+    def mock_call(self, tokens, start_pos, temperature):
       captured_inputs.append((tokens.shape, start_pos if isinstance(start_pos, int) else start_pos.val))
       return Tensor([[42]])
 
@@ -40,7 +40,7 @@ class TestTransformerGenerate(unittest.TestCase):
                         norm_eps=1e-5, vocab_size=100, head_dim=32, rope_theta=10000.0, max_context=32)
 
     captured_inputs = []
-    def mock_call(self, tokens, start_pos):
+    def mock_call(self, tokens, start_pos, temperature):
       captured_inputs.append((tokens.shape, start_pos if isinstance(start_pos, int) else start_pos.val))
       return Tensor([[42]])
 
@@ -92,7 +92,7 @@ class TestTransformerGenerate(unittest.TestCase):
 
     def get_prefill_flags(tokens, chunk_size):
       is_prefill = []
-      def mock_call(self, tokens, start_pos):
+      def mock_call(self, tokens, start_pos, temperature):
         is_prefill.append(resolve(tokens.shape[1] != 1))
         return Tensor([[42]])
       with patch.object(Transformer, '__call__', mock_call):
@@ -108,5 +108,44 @@ class TestTransformerGenerate(unittest.TestCase):
     # 4 tokens, chunk_size=4 -> 1 prefill chunk
     self.assertEqual(get_prefill_flags(list(range(4)), 4), [True, False, False])
 
+  def test_temperature_zero_is_greedy(self):
+    """Temperature 0 (or near 0) should produce deterministic output."""
+    from tinygrad.apps.llm import Transformer
+    model = Transformer(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, max_context=32)
+    tokens = list(range(1, 6))
+    results = [list(zip(range(5), model.generate(list(tokens)))) for _ in range(3)]
+    # all runs should produce the same tokens
+    self.assertEqual(results[0], results[1])
+    self.assertEqual(results[1], results[2])
+
+  def test_temperature_high_produces_variety(self):
+    """High temperature should produce different outputs across runs."""
+    from tinygrad.apps.llm import Transformer
+    model = Transformer(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, max_context=32)
+    tokens = list(range(1, 6))
+    runs = set()
+    for _ in range(5):
+      gen = model.generate(list(tokens), temperature=2.0)
+      out = tuple(next(gen) for _ in range(10))
+      runs.add(out)
+    # 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_temperature_passed_to_forward(self):
+    """Temperature from generate should be passed through to __call__."""
+    from tinygrad.apps.llm import Transformer
+    model = Transformer(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, max_context=32)
+    captured_temps = []
+    def mock_call(self, tokens, start_pos, temperature):
+      captured_temps.append(float(temperature.item()))
+      return Tensor([[42]])
+    with patch.object(Transformer, '__call__', mock_call):
+      gen = model.generate([1, 2, 3], temperature=0.6)
+      next(gen)
+    self.assertAlmostEqual(captured_temps[-1], 0.6, places=5)
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_shm_tensor.py

@@ -1,7 +1,7 @@
 import unittest
 import multiprocessing.shared_memory as shared_memory
 from tinygrad.helpers import CI, WIN
-from tinygrad.tensor import Tensor, Device
+from tinygrad import Tensor, Device
 import numpy as np
 
 class TestRawShmBuffer(unittest.TestCase):

test/unit/test_tensor_data.py

@@ -1,5 +1,6 @@
 import unittest, struct
 from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp
 
 # format types: https://docs.python.org/3/library/struct.html
 
@@ -78,5 +79,9 @@ class TestTensorData(unittest.TestCase):
     assert dat.shape == (2,2)
     # NOTE: python can't deref float16
 
+  def test_data_uop_device(self):
+    uop = UOp.const(dtypes.float, 1.0, "DEVICE")
+    self.assertEqual(Tensor(uop).device, "DEVICE")
+
 if __name__ == '__main__':
   unittest.main()

tinygrad/apps/llm.py

@@ -7,7 +7,7 @@ from tinygrad.viz.serve import TCPServerWithReuse, HTTPRequestHandler
 
 class SimpleTokenizer:
   def __init__(self, normal_tokens:dict[str, int], special_tokens:dict[str, int], preset:str="llama3"):
-    if preset not in ("llama3","llama-v3","llama-bpe","qwen2","olmo"): raise ValueError(f"Invalid tokenizer preset '{preset}'")
+    if preset not in ("llama3","llama-v3","llama-bpe","qwen2","olmo","deepseek"): raise ValueError(f"Invalid tokenizer preset '{preset}'")
     # https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9
     bs = [*range(33, 127), *range(161, 173), *range(174, 256)]  # bytes that map to themselves
     self._byte_decoder = {chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)}
@@ -56,10 +56,12 @@ class SimpleTokenizer:
   def role(self, role:str):
     if self.preset == 'olmo': return self.encode("<|" + role + "|>\n")  # OLMoE Instruct format
     if self.preset == 'qwen2': return self.encode("<|im_start|>" + role + "\n")
+    if self.preset == 'deepseek': return self.encode("<|" + role + "|>\n")
     return self.encode("<|start_header_id|>" + role + "<|end_header_id|>\n\n")
   def end_turn(self, eos_id:int):
     if self.preset == 'olmo': return self.encode("\n")
     if self.preset == 'qwen2': return [eos_id] + self.encode("\n")
+    if self.preset == 'deepseek': return [eos_id]
     return [eos_id]
 
 @functools.cache
@@ -84,7 +86,9 @@ def apply_rope(x:Tensor, freqs_cis:Tensor) -> Tensor:
 
 class TransformerBlock:
   def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int, norm_eps:float, head_dim:int, rope_theta:float,
-               max_context:int=0, qk_norm:int=0, num_experts:int=0, num_experts_per_tok:int=0, norm_topk_prob:bool=False):
+               max_context:int=0, qk_norm:int=0, num_experts:int=0, num_experts_per_tok:int=0, norm_topk_prob:bool=False,
+               scoring_func:str="softmax", routed_scaling_factor:float=1.0, n_shared_experts:int=0, shared_expert_hidden_dim:int=0,
+               kv_lora_rank:int=0, qk_nope_head_dim:int=0, qk_rope_head_dim:int=0, v_head_dim:int=0):
     self.n_heads      = n_heads
     self.n_kv_heads   = n_kv_heads
     self.head_dim     = head_dim
@@ -92,13 +96,28 @@ class TransformerBlock:
     self.max_context  = max_context
     self.qk_norm      = qk_norm
 
+    # --- MLA (Multi-head Latent Attention) config -------------------------
+    self.kv_lora_rank      = kv_lora_rank
+    self.qk_nope_head_dim  = qk_nope_head_dim
+    self.qk_rope_head_dim  = qk_rope_head_dim
+    self.v_head_dim        = v_head_dim
+
     # --- attention projections (all linear, bias-free) ------------------
-    q_proj_out       = self.head_dim * n_heads
-    kv_proj_out      = self.head_dim * n_kv_heads
-    self.attn_q      = nn.Linear(dim, q_proj_out,  bias=False)
-    self.attn_k      = nn.Linear(dim, kv_proj_out, bias=False)
-    self.attn_v      = nn.Linear(dim, kv_proj_out, bias=False)
-    self.attn_output = nn.Linear(q_proj_out, dim,  bias=False)
+    if kv_lora_rank > 0:
+      # MLA: query projects to n_heads * (nope + rope), KV uses compressed latent
+      q_head_dim = qk_nope_head_dim + qk_rope_head_dim
+      self.attn_q      = nn.Linear(dim, n_heads * q_head_dim, bias=False)
+      self.attn_kv_a_mqa = nn.Linear(dim, kv_lora_rank + qk_rope_head_dim, bias=False)
+      self.attn_kv_a_norm = nn.RMSNorm(kv_lora_rank, norm_eps)
+      self.attn_kv_b   = nn.Linear(kv_lora_rank, n_heads * (qk_nope_head_dim + v_head_dim), bias=False)
+      self.attn_output = nn.Linear(n_heads * v_head_dim, dim, bias=False)
+    else:
+      q_proj_out       = self.head_dim * n_heads
+      kv_proj_out      = self.head_dim * n_kv_heads
+      self.attn_q      = nn.Linear(dim, q_proj_out,  bias=False)
+      self.attn_k      = nn.Linear(dim, kv_proj_out, bias=False)
+      self.attn_v      = nn.Linear(dim, kv_proj_out, bias=False)
+      self.attn_output = nn.Linear(q_proj_out, dim,  bias=False)
 
     # --- RMSNorms --------------------------------------------------------
     self.attn_norm   = nn.RMSNorm(dim, norm_eps)
@@ -109,68 +128,148 @@ class TransformerBlock:
     if num_experts > 0:
       self.norm_topk_prob = norm_topk_prob
       self.num_experts_per_tok = num_experts_per_tok
+      self.scoring_func = scoring_func
+      self.routed_scaling_factor = routed_scaling_factor
       self.ffn_gate_inp = nn.Linear(dim, num_experts, bias=False)  # router
+      if scoring_func == "sigmoid": self.exp_probs_b = Tensor.zeros(num_experts)  # e_score_correction_bias for noaux_tc
       self.ffn_gate_exps = ExpertWeights(num_experts, dim, hidden_dim)
       self.ffn_up_exps = ExpertWeights(num_experts, dim, hidden_dim)
       self.ffn_down_exps = ExpertWeights(num_experts, hidden_dim, dim)
+      # shared experts (always active, separate dense MLP)
+      if n_shared_experts > 0:
+        self.ffn_gate_shexp = nn.Linear(dim, shared_expert_hidden_dim, bias=False)
+        self.ffn_up_shexp   = nn.Linear(dim, shared_expert_hidden_dim, bias=False)
+        self.ffn_down_shexp = nn.Linear(shared_expert_hidden_dim, dim, bias=False)
     else:
       self.ffn_gate    = nn.Linear(dim, hidden_dim, bias=False)
       self.ffn_up      = nn.Linear(dim, hidden_dim, bias=False)
       self.ffn_down    = nn.Linear(hidden_dim, dim, bias=False)
 
-  @function(precompile=bool(getenv("PRECOMPILE", 0)), allow_implicit=False)
+  @function
   def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
-    x_norm = self.attn_norm(x)                       # (B,T,D)
-    q, k, v = self.attn_q(x_norm), self.attn_k(x_norm), self.attn_v(x_norm)
-    if self.qk_norm and self.qk_norm != self.head_dim: q, k = self.attn_q_norm(q), self.attn_k_norm(k)
-
     B, T, _ = x.shape
-    q = q.reshape(B, T, self.n_heads,    self.head_dim).transpose(1, 2)  # (B,H,T,Hd)
-    k = k.reshape(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B,KvH,T,Hd)
-    v = v.reshape(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B,KvH,T,Hd)
-    if self.qk_norm == self.head_dim: q, k = self.attn_q_norm(q), self.attn_k_norm(k)
+    x_norm = self.attn_norm(x)                       # (B,T,D)
 
-    q = apply_rope(q, self.freqs_cis[start_pos:start_pos+T])
-    k = apply_rope(k, self.freqs_cis[start_pos:start_pos+T])
+    if self.kv_lora_rank > 0:
+      # --- MLA (Multi-head Latent Attention) path ---
+      q = self.attn_q(x_norm)                        # (B,T, n_heads*(nope+rope))
+      q = q.reshape(B, T, self.n_heads, self.qk_nope_head_dim + self.qk_rope_head_dim).transpose(1, 2)  # (B,H,T,nope+rope)
+      q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
 
-    self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(Tensor.stack(k, v))
-    k = self.cache_kv[0, :, :, 0:start_pos+T, :]
-    v = self.cache_kv[1, :, :, 0:start_pos+T, :]
+      # compress KV into latent + rope key
+      compressed_kv = self.attn_kv_a_mqa(x_norm)     # (B,T, kv_lora_rank + rope)
+      compressed_kv, k_pe = compressed_kv.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+      k_pe = k_pe.reshape(B, T, 1, self.qk_rope_head_dim).transpose(1, 2)   # (B,1,T,rope)
+
+      # decompress latent to per-head K_nope and V
+      kv = self.attn_kv_b(self.attn_kv_a_norm(compressed_kv))  # (B,T, n_heads*(nope+v))
+      kv = kv.reshape(B, T, self.n_heads, self.qk_nope_head_dim + self.v_head_dim).transpose(1, 2)  # (B,H,T,nope+v)
+      k_nope, v = kv.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+
+      # apply RoPE only to the rope portion of Q and K
+      q_pe = apply_rope(q_pe, self.freqs_cis[start_pos:start_pos+T])
+      k_pe = apply_rope(k_pe, self.freqs_cis[start_pos:start_pos+T])
+
+      # reassemble full Q and K: [nope, rope]
+      q = q_nope.cat(q_pe, dim=-1)                   # (B,H,T,nope+rope)
+      k = k_nope.cat(k_pe.expand(-1, self.n_heads, -1, -1), dim=-1)  # (B,H,T,nope+rope)
+
+      # KV cache — K and V have different last dims, pad V to match K for stacking
+      q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+      if self.v_head_dim < q_head_dim:
+        v_padded = v.pad((0, q_head_dim - self.v_head_dim))  # pad last dim
+      else:
+        v_padded = v
+      assigned_kv = Tensor(self.cache_kv.uop.after(self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(Tensor.stack(k, v_padded).uop)))
+      k = assigned_kv[0, :, :, 0:start_pos+T, :]
+      v = assigned_kv[1, :, :, 0:start_pos+T, :self.v_head_dim]  # unpad V
+
+      mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device).triu(start_pos+1) if resolve(T != 1) else None
+      attn = q.scaled_dot_product_attention(k, v, attn_mask=mask)  # (B,H,T,v_head_dim)
+      attn = attn.transpose(1, 2).reshape(B, T, self.n_heads * self.v_head_dim)
+    else:
+      # --- standard GQA/MHA path ---
+      q, k, v = self.attn_q(x_norm), self.attn_k(x_norm), self.attn_v(x_norm)
+      if self.qk_norm and self.qk_norm != self.head_dim: q, k = self.attn_q_norm(q), self.attn_k_norm(k)
+
+      q = q.reshape(B, T, self.n_heads,    self.head_dim).transpose(1, 2)  # (B,H,T,Hd)
+      k = k.reshape(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B,KvH,T,Hd)
+      v = v.reshape(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)  # (B,KvH,T,Hd)
+      if self.qk_norm == self.head_dim: q, k = self.attn_q_norm(q), self.attn_k_norm(k)
+
+      q = apply_rope(q, self.freqs_cis[start_pos:start_pos+T])
+      k = apply_rope(k, self.freqs_cis[start_pos:start_pos+T])
+
+      # NOTE: we don't want to change self.cache_kv, the function API doesn't support this well
+      assigned_kv = Tensor(self.cache_kv.uop.after(self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(Tensor.stack(k, v).uop)))
+      k = assigned_kv[0, :, :, 0:start_pos+T, :]
+      v = assigned_kv[1, :, :, 0:start_pos+T, :]
+
+      # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
+      # TODO: this if statement should be removed and it shouldn't generate extra kernels
+      mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device).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)
 
-    # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
-    # TODO: this if statement should be removed and it shouldn't generate extra kernels
-    mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device).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)
     attn = self.attn_output(attn)
     return x + attn
 
-  @function(precompile=bool(getenv("PRECOMPILE", 0)), allow_implicit=False)
+  @function
   def _feed_forward(self, h: Tensor) -> Tensor:
     h_norm = self.ffn_norm(h)
     if hasattr(self, 'ffn_gate_exps'):
       x = h_norm.unsqueeze(2)  # (B, T, 1, D) - add expert dim for broadcasting
-      probs, sel = self.ffn_gate_inp(h_norm).softmax(-1).topk(self.num_experts_per_tok)  # (B, T, k) each
+      if self.scoring_func == "sigmoid":
+        scores = self.ffn_gate_inp(h_norm).sigmoid()  # (B, T, num_experts)
+        # noaux_tc: add learned bias for top-k selection, then select top-k from original scores
+        scores_for_choice = scores + self.exp_probs_b.reshape(1, 1, -1)
+        _, sel = scores_for_choice.topk(self.num_experts_per_tok)  # (B, T, k)
+        probs = scores.gather(2, sel)  # (B, T, k) - gather original scores at selected indices
+      else:
+        probs, sel = self.ffn_gate_inp(h_norm).softmax(-1).topk(self.num_experts_per_tok)  # (B, T, k) each
       if self.norm_topk_prob: probs = probs / probs.sum(axis=-1, keepdim=True)
+      probs = probs * self.routed_scaling_factor
       x_down = self.ffn_down_exps(sel, self.ffn_gate_exps(sel, x).silu() * self.ffn_up_exps(sel, x))  # (B, T, k, D)
-      return h + (x_down * probs.unsqueeze(-1)).sum(axis=2)  # (B, T, D)
+      moe_out = (x_down * probs.unsqueeze(-1)).sum(axis=2)  # (B, T, D)
+      # add shared expert output if present
+      if hasattr(self, 'ffn_gate_shexp'):
+        moe_out = moe_out + self.ffn_down_shexp(self.ffn_gate_shexp(h_norm).silu() * self.ffn_up_shexp(h_norm))
+      return h + moe_out
     # TODO: remove the need for this contiguous
     gated  = self.ffn_gate(h_norm).silu().contiguous() * self.ffn_up(h_norm)
     return h + self.ffn_down(gated)
 
   def __call__(self, x: Tensor, start_pos: int|UOp):
     if not hasattr(self, "cache_kv"):
-      # TODO: how is the dtype of this determined?
-      # NOTE: clone is used to promise the creation of a specific buffer
-      self.cache_kv = Tensor.zeros(2, x.shape[0], self.n_kv_heads, self.max_context, self.head_dim, device=x.device).clone()
-      self.freqs_cis = precompute_freqs_cis(self.head_dim, self.max_context, self.rope_theta)
-    return self._feed_forward(self._attention(x, start_pos)).contiguous()
+      if self.kv_lora_rank > 0:
+        # MLA: K cache stores full decompressed K (nope+rope), V cache stores full decompressed V
+        q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+        self.cache_kv = Tensor.empty(2, x.shape[0], self.n_heads, self.max_context, max(q_head_dim, self.v_head_dim), device=x.device)
+        self.freqs_cis = precompute_freqs_cis(self.qk_rope_head_dim, self.max_context, self.rope_theta)
+      else:
+        # TODO: how is the dtype of this determined?
+        self.cache_kv = Tensor.empty(2, x.shape[0], self.n_kv_heads, self.max_context, self.head_dim, device=x.device)
+        self.freqs_cis = precompute_freqs_cis(self.head_dim, self.max_context, self.rope_theta)
+    # 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): return self._feed_forward(self._attention(x, start_pos)).contiguous()
+    return _run(x, start_pos)
 
 class Transformer:
   def __init__(self, *, num_blocks, dim, hidden_dim, n_heads, n_kv_heads, norm_eps, vocab_size, head_dim:int, rope_theta:float,
-               max_context:int=0, qk_norm:int=0, num_experts:int=0, num_experts_per_tok:int=0, norm_topk_prob:bool=False):
-    self.blk = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, head_dim, rope_theta, max_context, qk_norm,
-                                 num_experts, num_experts_per_tok, norm_topk_prob) for _ in range(num_blocks)]
+               max_context:int=0, qk_norm:int=0, num_experts:int=0, num_experts_per_tok:int=0, norm_topk_prob:bool=False,
+               scoring_func:str="softmax", routed_scaling_factor:float=1.0, n_shared_experts:int=0, shared_expert_hidden_dim:int=0,
+               first_k_dense_replace:int=0, dense_hidden_dim:int=0,
+               kv_lora_rank:int=0, qk_nope_head_dim:int=0, qk_rope_head_dim:int=0, v_head_dim:int=0):
+    def make_block(i):
+      is_moe = num_experts > 0 and i >= first_k_dense_replace
+      return TransformerBlock(dim, hidden_dim if is_moe else (dense_hidden_dim or hidden_dim), n_heads, n_kv_heads, norm_eps,
+                              head_dim, rope_theta, max_context, qk_norm,
+                              num_experts if is_moe else 0, num_experts_per_tok if is_moe else 0, norm_topk_prob,
+                              scoring_func if is_moe else "softmax", routed_scaling_factor if is_moe else 1.0,
+                              n_shared_experts if is_moe else 0, shared_expert_hidden_dim if is_moe else 0,
+                              kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim)
+    self.blk = [make_block(i) for i in range(num_blocks)]
     self.token_embd  = nn.Embedding(vocab_size, dim)
     self.output_norm = nn.RMSNorm(dim, norm_eps)
     self.output = nn.Linear(dim, vocab_size, bias=False)
@@ -180,14 +279,15 @@ class Transformer:
     self.prefill_jit = TinyJit(self.forward)
     self.rollout_jit = TinyJit(self.forward)
 
-  def forward(self, tokens:Tensor, start_pos:int|UOp) -> Tensor:
-    x = self.token_embd(tokens)                           # (B, T, D)
+  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)
-    # TODO: add temperature
-    return self.output(self.output_norm(x))[:, -1, :].softmax(-1, dtype="float").argmax(-1, keepdim=True)
+    logits = self.output(self.output_norm(x))[:, -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=0) -> Tensor:
-    return (self.prefill_jit if resolve(tokens.shape[1] != 1) else self.rollout_jit)(tokens, start_pos)
+  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, start_pos, temperature)
 
   @staticmethod
   def from_gguf(gguf:Tensor, max_context:int|None=None, realize=bool(getenv("REALIZE", 0))) -> tuple[Transformer, dict]:
@@ -210,6 +310,35 @@ class Transformer:
         if 'attn_q.weight' in name: state_dict[name] = state_dict[name].rearrange("(n h two) d -> (n two h) d", n=n_heads, two=2)
         if 'attn_k.weight' in name: state_dict[name] = state_dict[name].rearrange("(n h two) d -> (n two h) d", n=n_kv_heads, two=2)
 
+    # MLA-specific GGUF metadata (deepseek2 architecture)
+    kv_lora_rank = kv.get(f'{arch}.attention.kv_lora_rank', 0)
+    qk_rope_head_dim = kv.get(f'{arch}.rope.dimension_count', 0) if kv_lora_rank else 0
+    # key_length_mla = qk_nope_head_dim + qk_rope_head_dim
+    qk_nope_head_dim = kv.get(f'{arch}.attention.key_length_mla', 0) - qk_rope_head_dim if kv_lora_rank else 0
+    v_head_dim = kv.get(f'{arch}.attention.value_length_mla', 0) if kv_lora_rank else 0
+    # deepseek2 MoE config
+    scoring_func = "sigmoid" if 'blk.1.exp_probs_b.bias' in state_dict else "softmax"
+    routed_scaling_factor = kv.get(f'{arch}.expert_weights_scale', 1.0)
+    n_shared_experts = kv.get(f'{arch}.expert_shared_count', 0)
+    expert_ff_len = kv.get(f'{arch}.expert_feed_forward_length', 0)
+    shared_expert_hidden_dim = expert_ff_len * n_shared_experts if n_shared_experts else 0
+    first_k_dense_replace = kv.get(f'{arch}.leading_dense_block_count', 0)
+
+    # for MLA, GGUF stores kv_b as a single tensor; we need to handle that during weight loading
+    # llama.cpp conversion may split kv_b into attn_k_b + attn_v_b, or keep as attn_kv_b
+    # reassemble split k_b/v_b into single kv_b for our model if needed
+    if kv_lora_rank:
+      for name in list(state_dict.keys()):
+        if 'attn_k_b.weight' in name:
+          prefix = name.replace('attn_k_b.weight', '')
+          k_b = state_dict.pop(name)                 # (nope, kv_lora_rank, n_heads) in GGUF 3D
+          v_b = state_dict.pop(f'{prefix}attn_v_b.weight')  # (kv_lora_rank, v_head_dim, n_heads)
+          # k_b was transposed during conversion: (qk_nope_head_dim, kv_lora_rank, n_heads)
+          # undo transpose: -> (kv_lora_rank, qk_nope_head_dim, n_heads) -> reshape to 2D Linear weight
+          k_b_2d = k_b.permute(2, 1, 0).reshape(-1, kv_lora_rank)   # (n_heads*qk_nope_head_dim, kv_lora_rank)
+          v_b_2d = v_b.permute(2, 1, 0).reshape(-1, kv_lora_rank)   # (n_heads*v_head_dim, kv_lora_rank)
+          state_dict[f'{prefix}attn_kv_b.weight'] = k_b_2d.cat(v_b_2d, dim=0)  # (n_heads*(nope+v), kv_lora_rank)
+
     model = Transformer(num_blocks=kv[f'{arch}.block_count'], dim=kv[f'{arch}.embedding_length'],
                         hidden_dim=kv.get(f'{arch}.expert_feed_forward_length', kv[f'{arch}.feed_forward_length']),
                         n_heads=n_heads, n_kv_heads=n_kv_heads, norm_eps=kv[f'{arch}.attention.layer_norm_rms_epsilon'],
@@ -218,7 +347,13 @@ class Transformer:
                         rope_theta=kv[f'{arch}.rope.freq_base'], max_context=max_context,
                         qk_norm=int(state_dict['blk.0.attn_q_norm.weight'].shape[0]) if 'blk.0.attn_q_norm.weight' in state_dict else 0,
                         num_experts=kv.get(f'{arch}.expert_count', 0), num_experts_per_tok=kv.get(f'{arch}.expert_used_count', 0),
-                        norm_topk_prob=True if arch=='qwen3moe' else False)
+                        norm_topk_prob=True if arch in ('qwen3moe', 'deepseek2') else False,
+                        scoring_func=scoring_func, routed_scaling_factor=routed_scaling_factor,
+                        n_shared_experts=n_shared_experts, shared_expert_hidden_dim=shared_expert_hidden_dim,
+                        first_k_dense_replace=first_k_dense_replace,
+                        dense_hidden_dim=kv.get(f'{arch}.feed_forward_length', 0),
+                        kv_lora_rank=kv_lora_rank, qk_nope_head_dim=qk_nope_head_dim,
+                        qk_rope_head_dim=qk_rope_head_dim, v_head_dim=v_head_dim)
     nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False)  # NOTE: rope_freqs.weight (32,) is unused
     # NOTE: without this contiguous, it unpacks the weights from the model every time. we shouldn't need this, but for now it's faster
     if realize:
@@ -229,9 +364,11 @@ class Transformer:
   def get_start_pos(self, tokens:list[int]):
     return sum(1 for _ in itertools.takewhile(lambda ab: ab[0] == ab[1], zip(tokens[:-1], self._cached_tokens)))
 
-  def generate(self, tokens:list[int], chunk_size:int=32):
+  def generate(self, tokens:list[int], chunk_size:int=32, temperature:float=0.0):
     v_start_pos = UOp.variable("start_pos", 0, self.max_context-1)
     v_toks = UOp.variable("toks", 1, chunk_size)
+    # TODO: use UOp.variable for temperature once float variables are supported
+    temp = Tensor(temperature).contiguous()
     # 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 kv cache
@@ -239,7 +376,7 @@ class Transformer:
     out, prompt_len = None, len(tokens)
     while len(tokens) < self.max_context:
       sp, nt = v_start_pos.bind(start_pos), v_toks.bind(min(chunk_size, len(tokens) - start_pos))
-      out = self(t[:, sp:sp+nt] if start_pos < prompt_len or out is None else out, sp).realize()
+      out = self(t[:, sp:sp+nt] if start_pos < prompt_len or out is None else out, sp, temp).realize()
       start_pos += nt.val
       # chunked prefill: keep processing until all prompt tokens are consumed
       if start_pos < len(tokens): continue
@@ -258,6 +395,7 @@ models = {
   "qwen3:8b": "https://huggingface.co/Qwen/Qwen3-8B-GGUF/resolve/main/Qwen3-8B-Q4_K_M.gguf",
   "qwen3:30b-a3b": "https://huggingface.co/Qwen/Qwen3-30B-A3B-GGUF/resolve/main/Qwen3-30B-A3B-Q4_K_M.gguf",
   "olmoe": "https://huggingface.co/allenai/OLMoE-1B-7B-0924-Instruct-GGUF/resolve/main/olmoe-1b-7b-0924-instruct-q4_k_m.gguf",
+  "moonlight": "https://huggingface.co/gabriellarson/Moonlight-16B-A3B-Instruct-GGUF/resolve/main/Moonlight-16B-A3B-Instruct-Q4_K_M.gguf",
 }
 
 # *** simple OpenAI compatible server on 11434 to match ollama ***
@@ -285,7 +423,7 @@ CHAT_HTML = b'''<!DOCTYPE html><html><head><title>tinygrad chat</title><style>
     input.value = '';
     const d = document.createElement('div'); d.className = 'msg'; chat.appendChild(d);
     const r = await fetch('/v1/chat/completions', {method: 'POST', headers: {'Content-Type': 'application/json'},
-      body: JSON.stringify({model: 'llama', messages: msgs, stream: true})});
+      body: JSON.stringify({model: 'llama', messages: msgs, stream: true, temperature: 0.7})});
     for (const rd = r.body.getReader(), dec = new TextDecoder();;) {
       const {done, value} = await rd.read();
       if (done) break;
@@ -300,24 +438,32 @@ CHAT_HTML = b'''<!DOCTYPE html><html><head><title>tinygrad chat</title><style>
 
 class Handler(HTTPRequestHandler):
   def log_request(self, code='-', size='-'): pass
-  def do_GET(self): self.send_data(CHAT_HTML, content_type="text/html")
-  def run_model(self, ids:list[int], model_name:str, include_usage=False):
+  def do_GET(self):
+    if self.path == "/v1/models": self.send_data(json.dumps({"object":"list","data":[{"id":model_name,"object":"model"}]}).encode())
+    else: self.send_data(CHAT_HTML, 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):
     cache_start_pos = model.get_start_pos(ids)
     stderr_log(f"{self.path}  {colored('--', 'BLACK')}  "
                f"in:{colored(f'{cache_start_pos:5d}', 'green')} +{len(ids)-cache_start_pos:5d}  {colored('--', 'BLACK')}  ")
     tmpl = {"id":f"chatcmpl-{uuid.uuid4().hex[:24]}", "object":"chat.completion.chunk", "created":int(time.time()), "model":model_name}
     yield {"choices": [{"index":0, "delta":{"role":"assistant","content":""}, "finish_reason":None}], **tmpl}
     out: list[int] = []
+    finish_reason = "stop"
     st = time.perf_counter()
-    for next_id in model.generate(ids):
+    for next_id in model.generate(ids, temperature=temperature):
       if len(out) == 0: stderr_log(f"prefill:{(len(ids)-cache_start_pos)/((pt:=time.perf_counter())-st):4.0f} tok/s  {colored('--', 'BLACK')}  ")
       if next_id == eos_id: break
       out.append(next_id)
       yield {"choices": [{"index":0, "delta":{"content":tok.decode([next_id])}, "finish_reason":None}], **tmpl}
-    yield {"choices": [{"index":0, "delta":{},"finish_reason":"stop"}], **tmpl}
+      if max_tokens is not None and len(out) >= max_tokens:
+        finish_reason = "length"
+        break
+    yield {"choices": [{"index":0, "delta":{},"finish_reason":finish_reason}], **tmpl}
     if include_usage:
       yield {"choices": [], "usage": {"prompt_tokens": len(ids), "completion_tokens": len(out), "total_tokens": len(ids) + len(out)}, **tmpl}
-    stderr_log(f"gen:{len(out)/(time.perf_counter()-pt):4.0f} tok/s  {colored('--', 'BLACK')}  out:{len(out):5d}\n")
+    et = time.perf_counter()
+    stderr_log(f"gen:{len(out)/(et-pt) if len(out) > 1 else 0:4.0f} tok/s  {colored('--', 'BLACK')}  "
+               f"out:{len(out):5d}  {colored('--', 'BLACK')}  total:{et-st:6.2f}s\n")
 
   def do_POST(self):
     raw_body = self.rfile.read(int(self.headers.get("Content-Length", "0")))
@@ -340,13 +486,17 @@ class Handler(HTTPRequestHandler):
       ids += tok.role("assistant")
 
       # reply
-      chunks = self.run_model(ids, body["model"], not body.get("stream") or body.get("stream_options",{}).get("include_usage", False))
+      max_tokens = body.get("max_completion_tokens") or body.get("max_tokens")
+      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)))
       if body.get("stream"): self.stream_json(chunks)
       else:
-        out = []
-        for c in chunks: out.append(c["choices"][0]["delta"].get("content", "") if c["choices"] else "")
+        out, finish_reason = [], "stop"
+        for c in chunks:
+          if c["choices"] and c["choices"][0].get("delta", {}).get("content"): out.append(c["choices"][0]["delta"]["content"])
+          if c["choices"] and c["choices"][0].get("finish_reason"): finish_reason = c["choices"][0]["finish_reason"]
         self.send_data(json.dumps({**c, "object":"chat.completion",
-          "choices":[{"index":0, "message":{"role":"assistant","content":"".join(out)}, "finish_reason":"stop"}]}).encode())
+          "choices":[{"index":0, "message":{"role":"assistant","content":"".join(out)}, "finish_reason":finish_reason}]}).encode())
     else:
       raise RuntimeError(f"unhandled path {self.path}")
 
@@ -370,6 +520,7 @@ if __name__ == "__main__":
   gc.collect()
 
   # extract some metadata
+  model_name = args.model
   tok = SimpleTokenizer.from_gguf_kv(kv)
   bos_id: int|None = kv.get('tokenizer.ggml.bos_token_id') if kv.get('tokenizer.ggml.add_bos_token', True) else None
   eos_id: int = kv['tokenizer.ggml.eos_token_id']

tinygrad/codegen/__init__.py

@@ -1,19 +1,18 @@
 from typing import cast
 from dataclasses import replace
 import itertools
-from tinygrad.helpers import DISABLE_FAST_IDIV, EMULATED_DTYPES, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, TracingKey, Context
+from tinygrad.helpers import DISABLE_FAST_IDIV, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, TracingKey, Context
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, pyrender
 from tinygrad.uop.spec import type_verify, program_spec, kernel_spec
 from tinygrad.renderer import Renderer, ProgramSpec, Estimates
-from tinygrad.dtype import dtypes, promo_lattice
-from tinygrad.device import is_dtype_supported
+from tinygrad.dtype import dtypes
 from tinygrad.helpers import panic
 from tinygrad.codegen.opt import Opt
 
 # import all pattern matchers here
 from tinygrad.codegen.gpudims import pm_add_gpudims
 from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing, symbolic, pm_move_where_on_load
-from tinygrad.uop.decompositions import get_late_rewrite_patterns, get_transcendental_patterns, pm_float_decomp, pm_long_decomp
+from tinygrad.uop.decompositions import get_late_rewrite_patterns, get_transcendental_patterns, pm_dtype_decomps
 from tinygrad.codegen.late.expander import expander, pm_pre_expander, pm_group_for_reduce
 from tinygrad.codegen.late.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
   ReduceContext, correct_load_store, pm_render, pm_add_loads
@@ -42,7 +41,7 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
     sink = graph_rewrite(sink, pm_split_ranges+pm_flatten_range, ctx={}, name="split ranges")
 
     # create image buffers
-    if IMAGE == 1 and ren.device in {"QCOM", "CL"}: sink = graph_rewrite(sink, pm_make_images, name="create image buffers", bottom_up=True)
+    if IMAGE and ren.device in {"QCOM", "CL", "PYTHON"}: sink = graph_rewrite(sink, pm_make_images, name="create image buffers", bottom_up=True)
 
     # symbolic (NOTE: this is a requirement for pm_simplify_ranges to be correct)
     sink = graph_rewrite(sink, sym+pm_flatten_range, name="initial symbolic")
@@ -92,11 +91,7 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True) -
   pm_decomp = symbolic_simple+get_late_rewrite_patterns(supported_ops, ren.device, bool(DISABLE_FAST_IDIV))
   pm_transcendental = symbolic_simple+get_transcendental_patterns(supported_ops, TRANSCENDENTAL>=2)
   sink = graph_rewrite(sink, pm_decomp, ctx=ren.device, name="decompositions")
-  if not is_dtype_supported(dtypes.long, ren.device) or dtypes.long in EMULATED_DTYPES.tolist(dtypes):
-    sink = graph_rewrite(sink, pm_long_decomp, name="decomp long -> int", bottom_up=True)
-  for fr, to in [(fr, next((to for to in promo_lattice[fr] if is_dtype_supported(to, ren.device)), dtypes.float))
-                 for fr in EMULATED_DTYPES.tolist(dtypes) if fr in dtypes.floats]:
-    sink = graph_rewrite(sink, pm_float_decomp, ctx=(fr, to), name=f"decomp {fr} -> {to}", bottom_up=True)
+  sink = graph_rewrite(sink, pm_dtype_decomps, ctx=(set(), ren.device, getattr(ren, "arch", "")), name="decomp dtypes")
   sink = graph_rewrite(sink, pm_transcendental, ctx=ren.device, name="transcendental")
 
   # final rules for the renderer (without sym)

tinygrad/codegen/gpudims.py

@@ -35,7 +35,7 @@ def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|No
     if len(limited) > len(max_sizes): raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
     # try to split up dims: (a,) -> (b, c)
     if limited == dims: limited = _split_dims(dims, max_sizes)
-  raw_idxs = [UOp(Ops.SPECIAL, dtypes.index, (sint_to_uop(s),), (f"{prefix}{i}")) for i,s in enumerate(limited)]
+  raw_idxs = [UOp(Ops.SPECIAL, dtypes.weakint, (sint_to_uop(s),), (f"{prefix}{i}")) for i,s in enumerate(limited)]
   if len(limited) < len(dims):
     ret = []
     if (contraction:=get_contraction(dims, limited)) is None: raise RuntimeError(f"get_contraction should not be None {dims=} {limited=}")
@@ -75,13 +75,17 @@ 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.index)]
+  if ctx.has_threads: idxs = [UOp.variable("core_id", 0, int(global_shape[0])-1, dtypes.int).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)
   else:
     # define indexes for GPU-like execution
-    idxs = get_grouped_dims("gidx", global_shape, ctx.global_max, reverse=True) + get_grouped_dims("lidx", local_shape, ctx.local_max)
+    local_idxs = get_grouped_dims("lidx", local_shape, ctx.local_max)
+    hw_local = [_dim_max(u.src[0]) for u in local_idxs if u.op is Ops.SPECIAL]
+    global_max = ctx.global_max if ctx.global_prod_max is None else \
+      tuple(min(gm, pm//l) for gm,pm,l in zip(ctx.global_max or ctx.global_prod_max, ctx.global_prod_max, hw_local+[1]*3))
+    idxs = get_grouped_dims("gidx", global_shape, global_max, reverse=True) + local_idxs
 
   # apply to multiple ranges
   subs = {}

tinygrad/codegen/late/devectorizer.py

@@ -5,7 +5,7 @@ from dataclasses import dataclass
 from tinygrad.dtype import dtypes, ImageDType, DType, AddrSpace, Invalid, PtrDType
 from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, identity_element
 from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate
-from tinygrad.helpers import getenv, flatten, AMX, prod, IMAGE
+from tinygrad.helpers import getenv, flatten, AMX, prod
 from tinygrad.renderer import Renderer
 
 # ***** image load valid simplification *****
@@ -63,7 +63,7 @@ load_store_indexing = PatternMatcher([
 
 def expand_index(buf:UOp, vec:UOp):
   # determine optimal image shapes
-  if IMAGE == 1 and isinstance(dt:=buf.dtype, ImageDType):
+  if isinstance(dt:=buf.dtype, ImageDType):
     x, valid = vec.get_idx().gep(0), vec.get_valid().gep(0)
     # search for dims that drop the most valid statements
     best_drop, cands = -1, []
@@ -197,29 +197,19 @@ def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
   return UOp(Ops.VCAT, ls.dtype, tuple(ret)) if ls.op is Ops.LOAD else UOp.group(*ret)
 
 def get_image_idx(idx:UOp, width:int):
-  oidx = UOp(Ops.VECTORIZE, dtypes.index.vec(2), (((x:=idx.src[1].get_idx()) // 4) % width, (x // (4*width))))
+  oidx = UOp(Ops.VECTORIZE, dtypes.weakint.vec(2), (((x:=idx.src[1].get_idx()) // 4) % width, (x // (4*width))))
   return idx.replace(src=(idx.src[0], oidx.valid(idx.src[1].get_valid())))
 
 def image_fixup(ls:UOp):
   # normal image load or store, with the CAST from expand_index
-  if ls.src[0].op is Ops.CAST and isinstance(image_dtype:=ls.src[0].src[0].dtype, ImageDType):
+  if isinstance(dt:=ls.src[0].src[0].dtype, ImageDType) and ls.src[0].op is Ops.CAST:
     assert ls.src[0].dtype.count == 4, "image must be casted to 4"
-    idx = get_image_idx(ls.src[0].src[0], image_dtype.shape[1])
-    return ls.replace(src=(idx,)+ls.src[1:])
+    return ls.replace(src=(get_image_idx(ls.src[0].src[0], dt.shape[1]),)+ls.src[1:])
 
-  # this is an unprocessed image without a cast, aka unfoldable image load. this doesn't work for stores
-  if isinstance(image_dtype:=ls.src[0].dtype, ImageDType) and ls.src[0].src[1].get_idx().dtype != dtypes.index.vec(2):
-    assert ls.op is Ops.LOAD, "if an image store isn't upcasted to 4, we can't store it"
-    x, idx = ls.src[0].src[1].get_idx(), get_image_idx(ls.src[0], image_dtype.shape[1])
-    vec_load = ls.replace(dtype=ls.dtype.vec(4), src=(idx,)+ls.src[1:])
-    # image pixels have 4 channels (.xyzw), select channel based on x % 4
-    x_mod_4 = x % 4
-    def sel(ret, i): return x_mod_4.ne(i).where(ret, vec_load.gep(i))
-    # if x is non-negative, x % 4 is in [0, 3] and we can skip NAN fallback
-    if x_mod_4.vmin >= 0: return functools.reduce(sel, range(int(x_mod_4.vmin)+1, int(x_mod_4.vmax)+1), vec_load.gep(int(x_mod_4.vmin)))
-    return functools.reduce(sel, range(4), ls.const_like(float('nan')))
-
-  return None
+  # this is an unprocessed image without a cast, we should just make it a buffer
+  if isinstance(dt, ImageDType) and (off:=ls.src[0].src[1]).get_idx().dtype != dtypes.weakint.vec(2):
+    idx = ls.src[0].src[0].replace(dtype=(new_dt:=dtypes.half if dt.itemsize == 2 else dtypes.float).ptr(dt.size)).index(off)
+    return ls.replace(src=(idx,), dtype=new_dt).cast(dtypes.float) if ls.op is Ops.LOAD else ls.replace(src=(idx, ls[1].cast(dtypes.float)))
 
 correct_load_store = PatternMatcher([
   # split LOAD/STORE
@@ -263,7 +253,7 @@ def no_vectorized_index(buf:UOp, cast:UOp, idx:UOp, bcast:UOp|None=None):
     # simple scalar index: one lane, all components
     pairs = [(0, c) for c in range(cnt)]
   idx_lanes, offsets = (tuple(x) for x in zip(*pairs))
-  return buf.broadcast(len(pairs)).index(idx.gep(idx_lanes)*cnt + UOp.const(dtypes.index.vec(len(pairs)), offsets), ptr=True)
+  return buf.broadcast(len(pairs)).index(idx.gep(idx_lanes)*cnt + UOp.const(dtypes.weakint.vec(len(pairs)), offsets), ptr=True)
 
 devectorize_buf_and_index = PatternMatcher([
   (UPat((Ops.DEFINE_LOCAL, Ops.DEFINE_REG), name="buf"), no_vectorized_buf),
@@ -296,12 +286,12 @@ pm_render = PatternMatcher([
       if len(x.src) == 1 or x.src[1].op in (Ops.CUSTOM, Ops.STORE, Ops.BARRIER) else None),
   # Where after gated load becomes alt value
   # NOTE: if a is CAST and a.src[0].dtype == l.dtype, use a.src[0] to avoid roundtrip cast (e.g. uint->float->uint)
-  (UPat.var("c").where(UPat(Ops.LOAD, src=(UPat().index(UPat.var("idx"), UPat.var("c")).or_casted(),), allow_any_len=True, name="l").or_casted(),
-    UPat.var("a")), lambda c,idx,l,a: l.replace(src=(l.src[0], 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(a.dtype)),
-  (UPat.var("c").where(UPat.var("a"), UPat(Ops.LOAD, src=(UPat().index(UPat.var("idx"), UPat.var("c").logical_not()).or_casted(),),
-    allow_any_len=True, name="l").or_casted()), lambda c,idx,l,a: l.replace(src=(l.src[0], 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(a.dtype)),
+  (UPat.var("c").where(UPat(Ops.LOAD, src=(UPat().index(UPat(), UPat.var("c")).or_casted(),), allow_any_len=True, name="l").or_casted(),
+    UPat.var("a")), lambda c,l,a: l.replace(src=(l.src[0], 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(a.dtype)),
+  (UPat.var("c").where(UPat.var("a"), UPat(Ops.LOAD, src=(UPat().index(UPat(), UPat.var("c").logical_not()).or_casted(),),
+    allow_any_len=True, name="l").or_casted()), lambda c,l,a: l.replace(src=(l.src[0], 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(a.dtype)),
 ])
 
 # *** Ops.REDUCE -> Ops.DEFINE_ACC ***

tinygrad/codegen/late/expander.py

@@ -56,6 +56,16 @@ def do_expand(root:UOp):
         # repeat the arg
         new_srcs.append(src.broadcast(expand_sz))
 
+  # for non-PtrDType INDEX on REG buffers, expand into individual scalar INDEXes instead of one vectorized INDEX
+  # this avoids creating a VECTORIZE of REG pointers which the devectorizer can't resolve
+  if root.op is Ops.INDEX and not isinstance(root.dtype, PtrDType) and \
+     isinstance(root.src[0].dtype, PtrDType) and root.src[0].dtype.addrspace == AddrSpace.REG:
+    idxs = []
+    for j in range(expand_sz):
+      idx_srcs = tuple(s.gep(j) if isinstance(s.dtype, PtrDType) or s.dtype.count > 1 else s for s in new_srcs)
+      idxs.append(UOp(Ops.INDEX, root.dtype, idx_srcs, root.arg))
+    return UOp(Ops.UNROLL, root.dtype, (UOp(Ops.VECTORIZE, root.dtype.vec(expand_sz), tuple(idxs)),), expand_args)
+
   new_arg = root.arg
   if root.op is Ops.GEP:
     assert root.dtype.count == 1

tinygrad/codegen/opt/postrange.py

@@ -332,7 +332,7 @@ class Scheduler:
 
 def bufs_from_ast(ast:UOp, dname:str) -> list[Buffer]:
   glbls = sorted([x for x in ast.backward_slice if x.op is Ops.PARAM], key=lambda x: x.arg)
-  return [Buffer(dname, x.ptrdtype.size, x.dtype.base if not isinstance(x.dtype, ImageDType) else x.dtype) for x in glbls]
+  return [Buffer(dname, x.ptrdtype.size, x.dtype.base) for x in glbls]
 
 def apply_opts(ast:UOp, ren:Renderer) -> UOp:
   if ast.tag is not None: return ast

tinygrad/codegen/simplify.py

@@ -1,6 +1,6 @@
 import itertools
 from typing import Callable
-from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start
+from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start, AxisType
 from tinygrad.uop.symbolic import symbolic
 from tinygrad.helpers import partition
 from tinygrad.dtype import dtypes, ImageDType
@@ -57,7 +57,7 @@ pm_simplify_ranges = PatternMatcher([
 ])
 
 def mark_range_mod(ctx:dict[UOp, UOp|None], r:UOp, c:UOp) -> None:
-  if r not in ctx and r.src[0].op is Ops.CONST and r.src[0].divides(c.arg) is not None: ctx[r] = c
+  if r not in ctx and r.arg[-1] is not AxisType.WARP and r.src[0].op is Ops.CONST and r.src[0].divides(c.arg) is not None: ctx[r] = c
 
 def do_substitute(ctx:dict, x: UOp, sub_fxn:Callable[[UOp, UOp], UOp]) -> UOp|None:
   ret = x.substitute({k:sub_fxn(k,v) for k,v in ctx.items() if v is not None})
@@ -157,5 +157,5 @@ def no_load(u:UOp) -> bool: return not any(x.op is Ops.INDEX for x in u.backward
 pm_load_collapse = PatternMatcher([
   (UPat(Ops.REDUCE, arg=Ops.ADD, src=(UPat.var("u"), UPat()), name="red"), reduce_load_collapse),
   # we want to make sure we dont do math on a loaded index since that can cause overflow, this undoes the rule in pm_reduce_load_collapse
-  ((UPat.var("x", dtypes.index)+UPat.var("y"))<UPat.var("c"), lambda x,y,c: x < c-y if no_load(y) and no_load(c) and not no_load(x) else None),
+  ((UPat.var("x", dtypes.weakint)+UPat.var("y"))<UPat.var("c"), lambda x,y,c: x < c-y if no_load(y) and no_load(c) and not no_load(x) else None),
 ])

tinygrad/device.py

@@ -3,11 +3,11 @@ from dataclasses import dataclass, replace
 from collections import defaultdict
 from typing import Any, Generic, TypeVar, Iterator, Generator, TYPE_CHECKING
 import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal
-from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored
+from tinygrad.helpers import BENCHMARKS, CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, colored
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, dedup, ContextVar
 from tinygrad.helpers import unwrap_class_type, suppress_finalizing, select_first_inited, VIZ, CPU_LLVM, CPU_LVP, NV_PTX, CUDA_PTX, NV_NAK
-from tinygrad.helpers import EMULATED_DTYPES, NULL_IR3, NULL_QCOMCL, TracingKey
-from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
+from tinygrad.helpers import EMULATE, EMULATED_DTYPES, NULL_IR3, NULL_QCOMCL, IMAGE, FLOAT16, TracingKey, size_to_str
+from tinygrad.dtype import DType, PtrDType, dtypes, _to_np_dtype
 if TYPE_CHECKING: from tinygrad.renderer import Renderer
 
 # **************** Device ****************
@@ -72,7 +72,6 @@ class ProfileGraphEvent(ProfileEvent): ents:list[ProfileGraphEntry]; deps:list[l
 @dataclass(frozen=True, eq=True)
 class BufferSpec:
   # TODO: move device, size, dtype here?
-  image: ImageDType|None = None
   uncached: bool = False
   cpu_access: bool = False
   host: bool = False
@@ -96,8 +95,7 @@ class Buffer:
   profile_events:list[ProfileEvent] = []
   def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None, options:BufferSpec|None=None, initial_value:bytes|None=None,
                uop_refcount=0, base:Buffer|None=None, offset:int=0, preallocate=False):
-    if isinstance(dtype, ImageDType): options = BufferSpec(image=dtype) # TODO: image hack shouldn't be here. where should it be?
-    else: assert isinstance(dtype, DType) and not isinstance(dtype, PtrDType)
+    assert isinstance(dtype, DType) and not isinstance(dtype, PtrDType)
     self.device, self.size, self.dtype, self.options, self.offset, self.allocated_views = device, size, dtype, options, offset, 0
     if base is None:
       assert offset == 0, "base buffers can't have offset"
@@ -175,12 +173,9 @@ class Buffer:
   def __repr__(self):
     return f"<buf real:{self.is_allocated()} device:{self.device} size:{self.size} dtype:{self.dtype}" + \
            (f" offset:{self.offset}" if self._base is not None else "") + (f" {self.options=}" if self.options is not None else "") + ">"
-  def as_dmaref(self) -> DMARef:
-    assert hasattr(self.allocator, "_as_dmaref"), f"Device {self.device} doesn't support DMA"
-    return self.allocator._as_dmaref(self._buf)
   def as_memoryview(self, allow_zero_copy=False, force_zero_copy=False) -> memoryview:
     # zero copy with as_memoryview (disabled by default due to use after free)
-    if (force_zero_copy or allow_zero_copy) and hasattr(self.allocator, '_as_buffer') and (self.options is None or self.options.image is None):
+    if (force_zero_copy or allow_zero_copy) and hasattr(self.allocator, '_as_buffer') and self.options is None:
       return self.allocator._as_buffer(self._buf)
     assert not force_zero_copy, "force zero copy was passed, but copy is required"
     return self.copyout(memoryview(bytearray(self.nbytes)))
@@ -204,19 +199,6 @@ class Buffer:
     assert offset < self.nbytes, "offset must be less than nbytes"
     return Buffer(self.device, size, dtype, base=self.base, offset=self.offset+offset)
 
-@dataclass(frozen=True)
-class DMACPURef:
-  addr: int
-  size: int
-
-@dataclass(frozen=True)
-class DMAFdRef:
-  fd: int
-  offset: int
-  size: int
-
-DMARef = DMACPURef|DMAFdRef
-
 DeviceType = TypeVar('DeviceType', bound='Compiled')
 
 # TODO: size, dest, src are the same type. can we enforce this?
@@ -228,7 +210,9 @@ class Allocator(Generic[DeviceType]):
   # overridden in LRUAllocator
   def alloc(self, size:int, options:BufferSpec|None=None):
     assert size > 0, f"alloc size must be positive, getting {size}"
-    return self._alloc(size, options if options is not None else self.default_buffer_spec)
+    try: return self._alloc(size, options if options is not None else self.default_buffer_spec)
+    except (RuntimeError, MemoryError) as e: raise MemoryError(f"Allocation of {size_to_str(size)} failed on {self.dev.device}. "
+                                                               f"Used: {size_to_str(GlobalCounters.mem_used_per_device[self.dev.device])}") from e
   def free(self, opaque, size:int, options:BufferSpec|None=None):
     self._free(opaque, options if options is not None else self.default_buffer_spec)
 
@@ -320,7 +304,7 @@ class Compiled:
 
     # remove disabled compilers
     for en, rc in self.comp_sets.values():
-      if en is not None and en.value == 0 and rc in comps: comps.remove(rc)
+      if en is not None and en.value == 0 and en.key in os.environ and rc in comps: comps.remove(rc)
 
     return select_first_inited(list(forced_comps) if len(forced_comps)>0 else comps, f"No compiler for {self.device} is available", self.cached_pair)
 
@@ -344,19 +328,21 @@ class Compiled:
 
 # TODO: move this to each Device
 # this only tracks if the dtype is natively supported, it may be supported in the frontend using decomps
-def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
-  if dtype == dtypes.index: return False
+def is_dtype_supported(dtype:DType, device:str|None=None, arch:str|None=None) -> bool:
   if device is None: device = Device.DEFAULT
   if dtype == dtypes.bfloat16:
-    if device == "METAL": return not CI
-    if device == "CUDA": return not CI and not CUDA_PTX
-    if device == "NV": return not CI and not NV_PTX and not NV_NAK
-    if device in {"CPU"}: return not CI and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"} and not CPU_LVP
+    if device == "METAL": return not CI or BENCHMARKS
+    if device == "CUDA": return (not CI or BENCHMARKS) and not CUDA_PTX
+    if device == "NV": return (not CI or BENCHMARKS) and not NV_PTX and not NV_NAK
+    if device in {"CPU"}: return (not CI or BENCHMARKS) and platform.machine() in {"arm", "arm64", "aarch64", "x86_64", "amd64"} and not CPU_LVP
     return device in {"AMD", "CL", "PYTHON", "NULL"}
   if dtype in dtypes.fp8_ocp:
-    if device == "CUDA": return not CI and not CUDA_PTX
-    if device == "NV": return not CI and not NV_PTX and not NV_NAK
-    if device == "AMD": return not CI and getattr(Device["AMD"], "target") == (9,5,0)
+    if device == "CUDA": return (not CI or BENCHMARKS) and not CUDA_PTX
+    if device == "NV": return (not CI or BENCHMARKS) and not NV_PTX and not NV_NAK
+    if device == "AMD":
+      # TODO: open the device to get arch of device, will be fixed after triple is in the device string
+      if arch is None: arch = getattr(Device[device].renderer, "arch", "")
+      return (not CI or BENCHMARKS) and arch == "gfx950"
     return device in {"PYTHON", "NULL"}
   if dtype in dtypes.fp8_fnuz: return device in {"PYTHON", "NULL"}
   if device == "WEBGPU": return dtype in [dtypes.bool, dtypes.char, dtypes.uchar, dtypes.short,
@@ -367,9 +353,9 @@ def is_dtype_supported(dtype:DType, device:str|None=None) -> bool:
   # PYTHON supports half memoryview in 3.12+ https://github.com/python/cpython/issues/90751
   # double can't be bitcast to anything without long support
   if dtype == dtypes.half:
-    if device == "CL": return not CI and not OSX
-    if device == "QCOM": return False # QCOM compiler is flaky with half
-    if device in ["CUDA", "NV"]: return not CI
+    if device == "CL": return (not CI or BENCHMARKS) and not OSX
+    if device == "QCOM": return bool(IMAGE) and bool(FLOAT16) # QCOM compiler is flaky with half
+    if device in ["CUDA", "NV"]: return (not CI or BENCHMARKS) or "CUDA" in EMULATE.value
     if device == "CPU" and CPU_LLVM: return OSX
     if device == "PYTHON": return sys.version_info >= (3, 12)
   if dtype == dtypes.float64: return (device not in {"METAL", "QCOM"} and not (OSX and device == "CL") and not NULL_IR3 and not NULL_QCOMCL

tinygrad/dtype.py

@@ -87,6 +87,15 @@ class DType(metaclass=DTypeMetaClass):
   def max(self):
     if dtypes.is_int(self): return 2**(self.scalar().bitsize)-1+self.min
     return float("inf") if dtypes.is_float(self) else True
+  def const(self, val: tuple[ConstType, ...]|ConstType):
+    if isinstance(val, tuple):
+      assert len(val) == self.count, f"mismatch {val} {self}"
+      return tuple(map(self.const, val))
+    if isinstance(val, InvalidType): return val
+    # NOTE: float('nan') != float('nan'), so we canonicalize here
+    if isinstance(val, float) and math.isnan(val): val = math.nan
+    # int is the default. wrap floats in ConstFloat to distinguish -0.0 from 0.0 in cache
+    return ConstFloat(float(val)) if dtypes.is_float(self) else bool(val) if dtypes.is_bool(self) else int(val)
 
 @dataclass(frozen=True, eq=False)
 class PtrDType(DType):
@@ -138,12 +147,11 @@ class ImageDType(PtrDType):
   # get list of (height, width) that do not require pitch padding
   @staticmethod
   def valid_dims(ptr:PtrDType) -> list[tuple[int,int]]:
-    ALIGN, MAXW = getenv("IMAGE_PITCH_ALIGN", 256 if OSX else 64), 16384
-    if ptr.base not in (dtypes.half, dtypes.float) or ptr.size > 4*MAXW*MAXW or (ptr.size if OSX else ptr.nbytes()) % ALIGN != 0: return []
-    if OSX and (ptr.size // 4) % ALIGN: return [] # OSX has stricter requirements for height=1 images
-    pxls: int = ptr.size // 4
-    return ([(1, pxls)] * (pxls < MAXW) + [(pxls//ALIGN//k, ALIGN*k) for k in range(ceildiv(pxls//ALIGN, MAXW), min(pxls//ALIGN, MAXW//ALIGN)+1)
-                                           if (pxls//ALIGN)%k == 0] if pxls//ALIGN else [])
+    ALIGN, MAXW, pxls = getenv("IMAGE_PITCH_ALIGN", 256 if OSX else 64), 16384, ptr.size // 4
+    if ptr.base not in (dtypes.half, dtypes.float) or ptr.size > 4*MAXW*MAXW: return []
+    # OSX has stricter requirements for height=1 images
+    if ptr.size % (ALIGN * 4) != 0: return [] if OSX or ptr.nbytes() % getenv("IMAGE_BASE_ALIGN", 64) != 0 or pxls > MAXW else [(1, pxls)]
+    return [(pxls//ALIGN//k, ALIGN*k) for k in range(ceildiv(pxls//ALIGN, MAXW), min(pxls//ALIGN, MAXW//ALIGN)+1) if (pxls//ALIGN)%k == 0]
 
 class dtypes:
   @staticmethod
@@ -151,7 +159,7 @@ class dtypes:
   def is_float(x: DType) -> bool: return x.scalar() in dtypes.floats or isinstance(x, ImageDType)
   @staticmethod # static methods on top, or bool in the type info will refer to dtypes.bool
   @functools.cache
-  def is_int(x: DType) -> bool: return x.scalar() in (dtypes.ints + (dtypes.index,))
+  def is_int(x: DType) -> bool: return x.scalar() in (dtypes.ints + (dtypes.weakint,))
   @staticmethod
   @functools.cache
   def is_unsigned(x: DType) -> bool: return x.scalar() in dtypes.uints
@@ -159,30 +167,21 @@ class dtypes:
   def is_bool(x: DType) -> bool: return x.scalar() == dtypes.bool
   @staticmethod
   def from_py(x) -> DType:
-    if x.__class__ is float: return dtypes.default_float
-    if x.__class__ is int: return dtypes.default_int
-    if x.__class__ is bool: return dtypes.bool
+    # NOTE: isinstance(True, int) is True, so bool must be checked before int
+    if isinstance(x, bool): return dtypes.bool
+    if isinstance(x, float): return dtypes.default_float
+    if isinstance(x, int): return dtypes.default_int
     # put this in the last is faster because there are more items than lists/tuples to check
-    if x.__class__ is list or x.__class__ is tuple: return max(dtypes.from_py(xi) for xi in x) if x else dtypes.default_float
+    if isinstance(x, (list, tuple)): return max(dtypes.from_py(xi) for xi in x) if x else dtypes.default_float
     raise RuntimeError(f"Could not infer dtype of {x} with type {type(x)}")
   @staticmethod
-  def as_const(val: tuple[ConstType, ...]|ConstType, dtype:DType):
-    if isinstance(val, tuple):
-      assert len(val) == dtype.count, f"mismatch {val} {dtype}"
-      return tuple(dtypes.as_const(x, dtype) for x in val)
-    if isinstance(val, InvalidType): return val
-    # NOTE: float('nan') != float('nan'), so we canonicalize here
-    if isinstance(val, float) and math.isnan(val): val = math.nan
-    # int is the default. wrap floats in ConstFloat to distinguish -0.0 from 0.0 in cache
-    return ConstFloat(float(val)) if dtypes.is_float(dtype) else bool(val) if dtypes.is_bool(dtype) else int(val)
-  @staticmethod
   def finfo(dtype:DType) -> tuple[int, int]:
     """(exponent, mantissa)"""
     if not dtypes.is_float(dtype): raise ValueError(f"{dtype} is not a floating point type")
     return {dtypes.float16: (5, 10), dtypes.bfloat16: (8, 7), dtypes.float32: (8, 23), dtypes.float64: (11, 52),
             dtypes.fp8e4m3: (4, 3), dtypes.fp8e5m2: (5, 2), dtypes.fp8e4m3fnuz: (4, 3), dtypes.fp8e5m2fnuz: (5, 2)}[dtype]
   void: Final[DType] = DType.new(-1, 0, "void", None)
-  index: Final[DType] = DType.new(-1, 800, "index", None)
+  weakint: Final[DType] = DType.new(-1, 800, "weakint", None)
   bool: Final[DType] = DType.new(0, 1, "bool", '?')
   int8: Final[DType] = DType.new(1, 8, "signed char", 'b')
   uint8: Final[DType] = DType.new(2, 8, "unsigned char", 'B')
@@ -229,7 +228,7 @@ class dtypes:
   uints = (uint8, uint16, uint32, uint64)
   sints = (int8, int16, int32, int64)
   ints = uints + sints
-  all = floats + ints + (bool, index) # noqa: A003
+  all = floats + ints + (bool, weakint) # noqa: A003
 
 if (env_default_float := getenv("DEFAULT_FLOAT", "")):
   dtypes.default_float = getattr(dtypes, env_default_float.lower())
@@ -239,7 +238,8 @@ DTypeLike = str|DType
 def to_dtype(dtype:DTypeLike) -> DType: return dtype if isinstance(dtype, DType) else getattr(dtypes, dtype.lower())
 
 # https://jax.readthedocs.io/en/latest/jep/9407-type-promotion.html
-# we don't support weak type and complex type
+# we don't support complex type
+# TODO: weakint and weakfloat in lattice
 promo_lattice = { dtypes.bool: [dtypes.int8, dtypes.uint8], dtypes.int8: [dtypes.int16], dtypes.int16: [dtypes.int32], dtypes.int32: [dtypes.int64],
   dtypes.int64: [dtypes.uint64], dtypes.uint8: [dtypes.int16, dtypes.uint16], dtypes.uint16: [dtypes.int32, dtypes.uint32],
   dtypes.uint32: [dtypes.int64, dtypes.uint64], dtypes.uint64: [dtypes.fp8e4m3, dtypes.fp8e5m2, dtypes.fp8e4m3fnuz, dtypes.fp8e5m2fnuz],
@@ -256,8 +256,8 @@ def least_upper_dtype(*ds:DType) -> DType:
       if not (images:=[d for d in ds if isinstance(d, ImageDType)]) else images[0]
 def least_upper_float(dt:DType) -> DType: return dt if dtypes.is_float(dt) else least_upper_dtype(dt, dtypes.default_float)
 
-DTYPES_DICT = {k: v for k, v in dtypes.__dict__.items() if isinstance(v, DType) and not k.startswith(("default", "void", "index", "_"))}
-INVERSE_DTYPES_DICT = {**{v.name:k for k,v in DTYPES_DICT.items()}, "void": "void", "index":"index"}
+DTYPES_DICT = {k: v for k, v in dtypes.__dict__.items() if isinstance(v, DType) and not k.startswith(("default", "void", "weakint", "_"))}
+INVERSE_DTYPES_DICT = {**{v.name:k for k,v in DTYPES_DICT.items()}, "void": "void", "weakint":"weakint"}
 
 @functools.cache
 def can_lossless_cast(dt0:DType, dt1:DType) -> bool:
@@ -265,7 +265,7 @@ def can_lossless_cast(dt0:DType, dt1:DType) -> bool:
   # similar to https://numpy.org/doc/stable/reference/generated/numpy.can_cast.html
   if dt0 == dt1 or dt0 == dtypes.bool: return True
   match dt1:
-    case dtypes.index: return dt0 in dtypes.ints
+    case dtypes.weakint: return dt0 in dtypes.ints
     case dtypes.double: return dt0 in (dtypes.float, dtypes.half, dtypes.bfloat16, *dtypes.fp8s,
       dtypes.uint32, dtypes.uint16, dtypes.uint8, dtypes.int32, dtypes.int16, dtypes.int8)
     case dtypes.float: return dt0 in (dtypes.half, dtypes.bfloat16, *dtypes.fp8s, dtypes.uint16, dtypes.uint8, dtypes.int16, dtypes.int8)