vf MI350X passthrough mode driver

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

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

.github/workflows/benchmark.yml

@@ -94,6 +94,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -148,6 +149,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -200,6 +202,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -249,6 +252,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code

.github/workflows/platform.yml

@@ -179,3 +179,35 @@ jobs:
       - name: Run test_tiny
         shell: bash
         run: python -m pytest -n=auto test/test_tiny.py --durations=20
+
+
+  qcomclcompiletests:
+    name: Compile-only (QCOM CL)
+    runs-on: ubuntu-24.04-arm
+    timeout-minutes: 15
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v6
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: compile-qcomcl
+          deps: testing_unit
+          tinydreno: 'true'
+      - name: Set env
+        shell: bash
+        run: printf "DEV=NULL:QCOMCL:a630\nNULL_ALLOW_COPYOUT=1" >> $GITHUB_ENV
+      - name: Run test_ops
+        shell: bash
+        run: |
+          python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
+          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20
+      - name: Run test_ops (IMAGE)
+        shell: bash
+        env:
+          IMAGE: 1
+          DEV: "NULL:QCOMCL:a630,IMAGE_PITCH_ALIGNMENT=64"
+        run: |
+          DEBUG=4 python test/backend/test_ops.py TestOps.test_gemm | grep read_imagef
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20

.github/workflows/test.yml

@@ -21,7 +21,7 @@ concurrency:
 jobs:
   docs:
     name: Docs
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: &linux ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
     timeout-minutes: 10
     env:
       CHECK_OOB: 0
@@ -61,7 +61,7 @@ jobs:
 
   torchbackend:
     name: Torch Backend Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
     timeout-minutes: 15
     steps:
     - name: Checkout Code
@@ -88,7 +88,7 @@ jobs:
 
   bepython:
     name: Python Backend
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
     timeout-minutes: 15
     steps:
     - name: Checkout Code
@@ -126,7 +126,7 @@ jobs:
 
   linter:
     name: Linters
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
     timeout-minutes: 10
 
     steps:
@@ -157,7 +157,7 @@ jobs:
 
   nulltest:
     name: Null Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
     timeout-minutes: 15
 
     steps:
@@ -191,7 +191,7 @@ jobs:
 
   unittest:
     name: Unit Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
     timeout-minutes: 15
 
     steps:
@@ -228,7 +228,7 @@ jobs:
       matrix:
         group: [1, 2]
     name: SPEC=2 (${{ matrix.group }})
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
     timeout-minutes: 15
     steps:
     - name: Checkout Code
@@ -244,7 +244,7 @@ jobs:
 
   fuzzing:
     name: Fuzzing
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
     timeout-minutes: 10
     steps:
     - name: Checkout Code
@@ -260,7 +260,7 @@ jobs:
 
   testopenclimage:
     name: CL IMAGE Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
     timeout-minutes: 15
     steps:
       - name: Checkout Code
@@ -280,7 +280,7 @@ jobs:
 
   testopenpilot:
     name: openpilot Compile Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
     timeout-minutes: 15
     steps:
       - name: Checkout Code
@@ -309,7 +309,7 @@ jobs:
 
   testonnxcpu:
     name: ONNX (CPU) Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
     timeout-minutes: 20
 
     steps:
@@ -328,7 +328,7 @@ jobs:
 
   testoptim:
     name: Optimization Tests
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
     timeout-minutes: 20
     steps:
       - name: Checkout Code
@@ -360,7 +360,7 @@ jobs:
 
   testllm:
     name: Test LLM
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
     timeout-minutes: 15
     env:
       CHECK_OOB: 0
@@ -387,7 +387,7 @@ jobs:
 
   testmodels:
     name: Models
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
     timeout-minutes: 15
     steps:
       - name: Checkout Code
@@ -407,7 +407,7 @@ jobs:
 
   testdsp:
     name: Linux (DSP)
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
     timeout-minutes: 15
     steps:
     - name: Checkout Code
@@ -435,7 +435,7 @@ jobs:
           - 'WEBGPU'
 
     name: Linux (DEV=${{ matrix.dev }})
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
     timeout-minutes: 20
     steps:
     - name: Checkout Code
@@ -461,7 +461,7 @@ jobs:
 
   testamdasm:
     name: AMD ASM IDE
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
     timeout-minutes: 20
     env:
       DEV: MOCKKFD+AMD
@@ -507,7 +507,7 @@ jobs:
 
   hcq2:
     name: hcq2
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
     timeout-minutes: 5
     steps:
       - name: Checkout Code
@@ -529,7 +529,7 @@ jobs:
 
   testmockam:
     name: Linux (am)
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
     timeout-minutes: 15
     env:
       DEV: MOCKPCI+AMD
@@ -565,7 +565,7 @@ jobs:
         arch: [gfx1100, gfx1201, gfx950]
 
     name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
     timeout-minutes: 15
     env:
       DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
@@ -603,7 +603,7 @@ jobs:
         backend: [ptx, nv]
 
     name: Linux (${{ matrix.backend }})
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+    runs-on: *linux
     timeout-minutes: 20
     env:
       FORWARD_ONLY: 1
@@ -637,17 +637,10 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        dev:
-          - 'NULL:IR3:a630'
-          - 'NULL:QCOMCL:a630'
-          - 'NULL:NAK:sm_120'
-    name: Compile-only (DEV=${{ matrix.dev }})
-    runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
+        backend: [ir3, nak]
+    name: Compile-only (${{ matrix.backend }})
+    runs-on: *linux
     timeout-minutes: 15
-    env:
-      NULL_ALLOW_COPYOUT: 1
-      DEV: ${{ matrix.dev }}${{ contains(matrix.dev, 'a630') && ',IMAGE_PITCH_ALIGNMENT=64' || '' }}
-      IMAGE: ${{ contains(matrix.dev, 'a630') && '1' || '0' }}
     steps:
       - name: Checkout Code
         uses: actions/checkout@v6
@@ -656,14 +649,21 @@ jobs:
         with:
           key: compile-${{ matrix.backend }}
           deps: "testing_unit mesa"
-          qemu: ${{ contains(matrix.dev, 'QCOMCL') }}
-      - name: Test IMAGE
+      - name: Set env
         shell: bash
-        if: contains(matrix.dev, 'a630')
-        run: DEBUG=7 python3 test/backend/test_ops.py TestOps.test_gemm | grep isam
+        run: printf "NULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'DEV=NULL:IR3:a630' || matrix.backend == 'nak' && 'DEV=NULL:NAK:sm_120' }}" >> $GITHUB_ENV
       - name: Run test_ops
         shell: bash
         run: |
           python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
           DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
           python -m pytest -n=auto test/backend/test_ops.py --durations=20
+      - name: Run test_ops (IMAGE)
+        if: matrix.backend == 'ir3'
+        shell: bash
+        env:
+          IMAGE: 1
+          DEV: "NULL:IR3:a630,IMAGE_PITCH_ALIGNMENT=64"
+        run: |
+          DEBUG=4 python3 test/backend/test_ops.py TestOps.test_gemm | grep image_load
+          python -m pytest -n=auto test/backend/test_ops.py --durations=20

examples/mlperf/models/gpt_oss.py

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

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

@@ -46,7 +46,7 @@ export DATA_SEED=${DATA_SEED:-5760}
 export JITBEAM=${JITBEAM:-3}
 export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
 
-export FAKEDATA=${FAKEDATA:-$([[ "$DEV" == NULL:* ]] && echo 1 || echo 0)} BENCHMARK=${BENCHMARK:-10}
+export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
   export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
 fi

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

@@ -1,8 +1,8 @@
 #!/usr/bin/env bash
 
 export PYTHONPATH="."
-export ROCM_PATH=${ROCM_PATH:-/opt/rocm-7.1.1}
-export PATH="$ROCM_PATH/bin:$PATH"
+export PATH="/opt/rocm-7.1.1/bin:$PATH"
+export ROCM_PATH="/opt/rocm-7.1.1"
 export DEV=${DEV:-AMD}
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000

extra/amdpci/hive_reset.py

@@ -7,7 +7,7 @@ from tinygrad.runtime.support.hcq import FileIOInterface
 from tinygrad.runtime.support.am.amdev import AMDev
 
 if __name__ == "__main__":
-  gpus = System.pci_scan_bus(0x1002, [(0xffff, [0x74a1, 0x75a0])])
+  gpus = System.pci_scan_bus(0x1002, [(0xffff, [0x74a1, 0x75a0, 0x75b0])])
   for gpu in gpus:
     drv_path = f"/sys/bus/pci/devices/{gpu}/driver"
     if FileIOInterface.exists(drv_path) and os.path.basename(os.readlink(drv_path)) == "amdgpu":

extra/gemm/cdna_asm_gemm.py

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

extra/runbook_digitalocean_mi350x.md

@@ -74,8 +74,8 @@ rclone config create mlc-training s3 provider=Cloudflare \
   secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 \
   endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
 
-mkdir -p /raid/datasets/c4-8b
-rclone copy mlc-training:mlcommons-training-wg-public/llama3_1/datasets/c4/llama3_1_8b/ /raid/datasets/c4-8b/ -P
+mkdir -p /root/datasets/c4-8b
+rclone copy mlc-training:mlcommons-training-wg-public/llama3_1/datasets/c4/llama3_1_8b/ /root/datasets/c4-8b/ -P
 ```
 
 Files downloaded (~85GB total, ~6 minutes):
@@ -85,6 +85,13 @@ Files downloaded (~85GB total, ~6 minutes):
 - `c4-validation-91205-samples.en_text_document.idx` (1.8 MB)
 - `LICENSE.txt`, `NOTICE.txt`
 
+### Symlink for the submission script
+The `dev_run.sh` script hardcodes `BASEDIR="/raid/datasets/c4-8b/"`. Symlink:
+```bash
+mkdir -p /raid/datasets
+ln -s /root/datasets/c4-8b /raid/datasets/c4-8b
+```
+
 ## Phase 4: wandb Login
 ```bash
 wandb login
@@ -93,7 +100,7 @@ Enter API key from https://wandb.ai/authorize
 
 ## Phase 5: Run Training
 
-### 5.1 Smoke test (beam search, 2 layers, real data)
+### 5.1 Smoke test (beam search, 2 layers, fake data)
 Always run beam first to validate the pipeline:
 ```bash
 cd /root/tinygrad
@@ -101,7 +108,7 @@ COMGR_PATH=/opt/rocm/lib/libamd_comgr.so \
 COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so \
 CC=/opt/rocm/core-7.14/lib/llvm/bin/clang \
 DEV=AMD:HIP \
-ROCM_PATH=/opt/rocm \
+ROCM_PATH=/opt/rocm BASEDIR=/root/datasets/c4-8b/ \
   bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
 ```
 
@@ -112,7 +119,7 @@ COMGR_PATH=/opt/rocm/lib/libamd_comgr.so \
 COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so \
 CC=/opt/rocm/core-7.14/lib/llvm/bin/clang \
 DEV=AMD:HIP \
-ROCM_PATH=/opt/rocm \
+ROCM_PATH=/opt/rocm BASEDIR=/root/datasets/c4-8b/ \
 WANDB=1 \
   bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh
 ```
@@ -126,6 +133,7 @@ WANDB=1 \
 | `CC` | `/opt/rocm/core-7.14/lib/llvm/bin/clang` | System clang doesn't know gfx950; must use ROCm's bundled clang |
 | `DEV` | `AMD:HIP` | Force HIPRenderer (comgr-based) over HIPCCRenderer (hipcc subprocess) |
 | `ROCM_PATH` | `/opt/rocm` | Script defaults to `/opt/rocm-7.1.1` which doesn't exist |
+| `BASEDIR` | `/root/datasets/c4-8b/` | Where C4 dataset was downloaded (script hardcodes `/raid/datasets/c4-8b/`) |
 | `WANDB` | `1` | Enable wandb logging (off by default) |
 
 ## Architecture
@@ -165,6 +173,12 @@ ldconfig
 ### `comgr not available: try setting COMGR_3_PATH?`
 comgr 3.x uses a separate module. Set `COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so` too.
 
+### `FileNotFoundError: '/raid/datasets/c4-8b/...'`
+Script hardcodes `BASEDIR`. Either symlink or edit the script:
+```bash
+mkdir -p /raid/datasets && ln -s /root/datasets/c4-8b /raid/datasets/c4-8b
+```
+
 ### `No such file or directory: 'clang'`
 Install clang: `apt-get install -y clang` (for CPU compilation).
 For gfx950 HIP compilation, comgr (not clang) is used — ensure the ROCm 7.14 comgr 3.3 is properly loaded via `COMGR_PATH` and `COMGR_3_PATH`.

extra/thunder/amd/fa.py

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

test/backend/test_custom_kernel.py

@@ -422,8 +422,9 @@ class TestCustomKernel(unittest.TestCase):
       return Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
     GlobalCounters.reset()
     y = run(x[0]).realize()
-    # backends that support contiguous views don't launch extra kernels
-    assert_kernel_count(2 if x[0].uop.contiguous_view() is None else 1)
+    # it's copying the input and the output
+    # TODO: subbuffer usage has runtime specific behavior, this will be fixed after the removal of SLICE.
+    assert_kernel_count(2 if y.device in ("CL", "WEBGPU") else 1)
     self.assertEqual(y.tolist(), [1, 2, 3, 4])
 
   @Context(DEV="CPU")

test/external/external_benchmark_all2all.py

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

test/models/test_whisper.py

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

test/null/test_uop_symbolic.py

@@ -948,11 +948,6 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(cond.where(u0, u1), 0, 1, "((a<2)!=True)")
     self.helper_test_variable(cond.where(u0, u1).where(u0, u1), 0, 1, "(a<2)")
 
-  def test_equivalent_const_max(self):
-    x = Variable("x", -10, 10)
-    self.helper_test_variable((x < 0).where(0, x), 0, 10, "x.maximum(0)")
-    self.helper_test_variable((0 < x).where(x, 0), 0, 10, "x.maximum(0)")
-
   def test_where_combine(self):
     cond = Variable("x", 0, 3) < 2
     a = Variable("a", 0, 3)

test/opt/test_kernel_opts.py

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

test/unit/test_llm_server.py

@@ -13,18 +13,12 @@ V_TOKS = UOp.variable("toks", 1, 32)  # 32 is the default chunk_size in generate
 class TestTransformerGenerate(unittest.TestCase):
   def test_warmup(self):
     model, calls = Transformer(TEST_CONFIG), []
-    def generate(tokens, **kwargs):
+    def generate(tokens):
       calls.append(tokens)
       yield from (1, 2)
     with patch.object(model, "generate", generate): model.warmup()
     self.assertEqual(calls, [[0], [0]])
 
-  def test_warmup_then_generate_with_default_chunk(self):
-    # warmup must not capture JIT graphs that generate()'s default chunk_size then rejects
-    model = Transformer(TEST_CONFIG)
-    model.warmup()
-    self.assertIsInstance(next(model.generate([5, 6, 7, 8])), int)
-
   def test_first_recurrent_generate_before_state_init(self):
     model = Transformer(TEST_CONFIG)
     model.has_recurrent_block = True
@@ -44,15 +38,6 @@ class TestTransformerGenerate(unittest.TestCase):
       next(model.generate([1, 2, 3, 4, 5, 42, 10]))
     self.assertEqual(calls, [((1, 1), V_START_POS.bind(5)), ((1, 1), V_START_POS.bind(6))])
 
-  def test_recurrent_divergent_prompt_restarts(self):
-    model, calls = Transformer(TEST_CONFIG), []
-    model.has_recurrent_block, model._cached_tokens = True, [1, 2, 9]
-    def mock_call(self, tokens, start_pos, temperature):
-      calls.append(start_pos)
-      return Tensor([[42]])
-    with patch.object(Transformer, '__call__', mock_call): next(model.generate([1, 2, 10, 11]))
-    self.assertEqual(calls[0], V_START_POS.bind(0))
-
   def test_template_starts_reasoning(self):
     router = StreamRouter(reasoning=True)
     self.assertEqual(list(router.route("reasoning</think>answer")),
@@ -193,12 +178,6 @@ class TestTransformerGenerate(unittest.TestCase):
     # with temperature=2.0, we should see at least 2 distinct outputs across 5 runs
     self.assertGreater(len(runs), 1, "high temperature should produce varied outputs")
 
-  def test_recurrent_temperature_high_produces_variety(self):
-    model = Transformer(TEST_CONFIG)
-    model.has_recurrent_block = True
-    outputs = {model.forward(Tensor([[1]]), 0, Tensor([2.0])).item() for _ in range(5)}
-    self.assertGreater(len(outputs), 1)
-
   def test_temperature_passed_to_forward(self):
     """Temperature from generate should be passed through to __call__."""
     model = Transformer(TEST_CONFIG)

tinygrad/engine/realize.py

@@ -214,15 +214,15 @@ def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
   tms = []
   for devices,name,estimates,prof in info.kernels:
     for device in devices:
-      tm = None
+      d, tm = cast(Any, Device[device]), None
       if prof:
-        (d:=cast(Any, Device[device])).prof_ents[prof[0]] = ProfileGraphEntry(device, name, *prof)
+        d.prof_ents[prof[0]] = ProfileGraphEntry(device, name, *prof)
         if ctx.wait:
           d.synchronize(timeout=ctx.timeout)
           st, en = (d.signal(x)._buf.cpu_view().view(fmt='Q')[0] for x in prof)
           tms.append(tm:=float(en-st)/d.timestamp_divider/1e6)
-      stat_call = call.replace(arg=replace(call.arg, name=name, aux=replace(info, estimates=estimates, kernels=())))
-      with track_stats(ctx, stat_call, device, [], ctx.var_vals) as et: et[0] = tm
+      with track_stats(ctx, call.replace(arg=replace(call.arg, name=name, aux=replace(info, estimates=estimates))), d.device, [], ctx.var_vals) as et:
+        et[0] = tm
   return max(tms) if tms else None
 
 # flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src

tinygrad/llm/model.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 import functools, itertools, pathlib
 from dataclasses import dataclass, replace
-from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, dtypes
+from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function
 from tinygrad.nn import Linear
 from tinygrad.llm.gguf import gguf_load
 from tinygrad.uop.ops import resolve
@@ -187,8 +187,8 @@ class TransformerBlock(FFNBlock):
 
   def _init_state(self, x:Tensor):
     if not hasattr(self, "cache_kv"):
-      self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim,
-                                   dtype=dtypes.default_float, device=x.device)
+      # TODO: how is the dtype of this determined?
+      self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim, device=x.device)
       self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta, device=x.device)
 
 class MLATransformerBlock(FFNBlock):
@@ -261,14 +261,13 @@ class GatedDeltaNetBlock(FFNBlock):
   def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
     B, T, _ = x.shape
     assert T == 1, "GatedDeltaNetBlock currently only supports T=1"
-    is_kda = hasattr(self, "ssm_g_a")
 
     # input processing
     x = x.half()
-    out_gate = self.ssm_g_b(self.ssm_g_a(x)) if is_kda else self.attn_gate(x)
+    out_gate = self.ssm_g_b(self.ssm_g_a(x)) if hasattr(self, "ssm_g_a") else self.attn_gate(x)
     out_gate = out_gate.reshape(B, 1, self.num_v_heads, self.head_v_dim)
     beta = self.ssm_beta(x).sigmoid().reshape(B, self.num_v_heads, 1, 1)
-    alpha = self.ssm_f_b(self.ssm_f_a(x)) if is_kda else self.ssm_alpha(x)
+    alpha = self.ssm_f_b(self.ssm_f_a(x)) if hasattr(self, "ssm_f_a") else self.ssm_alpha(x)
     alpha = ((alpha.float() + self.ssm_dt["bias"]).softplus().reshape(B, self.num_v_heads, -1) *
              self.ssm_a.reshape(1, self.num_v_heads, -1)).exp().unsqueeze(-2)
 
@@ -292,13 +291,14 @@ class GatedDeltaNetBlock(FFNBlock):
 
     # output
     core_attn_out = self.ssm_norm((recurrent_state@q).squeeze(-1).reshape(B, 1, self.num_v_heads, self.head_v_dim))
-    out_gate = out_gate.sigmoid() if is_kda else out_gate.silu()
+    out_gate = out_gate.sigmoid() if hasattr(self, "ssm_g_a") else out_gate.silu()
     return self.ssm_out((core_attn_out * out_gate).reshape(B, 1, -1).cast(x.dtype))
 
   # recurrent state can't be partially reused after divergence, force a full rebuild
   def _state_reset_ops(self):
     return [self.conv_state.assign(self.conv_state.const_like(0)),
             self.recurrent_state.assign(self.recurrent_state.const_like(0))] if hasattr(self, "conv_state") else []
+  def _reusable_prefix_len(self, prefix_len:int, cached_len:int) -> int: return 0 if prefix_len != cached_len else prefix_len
 
   def _init_state(self, x):
     if not hasattr(self, "conv_state"):
@@ -326,8 +326,7 @@ class Transformer:
   def forward(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor) -> Tensor:
     x = self.token_embd(tokens).float()                   # (B, T, D)
     for block in self.blk: x = block(x, start_pos)
-    # only run the output projection on the last token
-    logits = self.output(self.output_norm(x[:, -1:]))[:, -1, :]
+    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)
 
@@ -421,10 +420,6 @@ class Transformer:
     for _ in range(2): list(zip(range(2), self.generate([0])))
 
   def get_start_pos(self, tokens:list[int]) -> int:
-    # recurrent state can't be partially reused after divergence: reuse it only when tokens extend the cached prefix
-    if self.has_recurrent_block:
-      return len(self._cached_tokens) if self._cached_tokens and len(self._cached_tokens) < len(tokens) \
-        and tokens[:len(self._cached_tokens)] == self._cached_tokens else 0
     prefix_len = sum(1 for _ in itertools.takewhile(lambda ab: ab[0] == ab[1], zip(tokens[:-1], self._cached_tokens)))
     return min(block._reusable_prefix_len(prefix_len, len(self._cached_tokens)) for block in self.blk)
 

tinygrad/runtime/ops_amd.py

@@ -790,14 +790,10 @@ class KFDIface:
 
   def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
                    xcc_id=0, idx=0):
-    if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA and idx and self.kfd_ver >= (1,17):
-      queue_type = kfd.KFD_IOC_QUEUE_TYPE_SDMA_BY_ENG_ID
-      sdma_engine_id = idx % (self.props['num_sdma_engines'] + self.props.get('num_sdma_xgmi_engines', 0))
-    else: sdma_engine_id = 0
     queue = kfd.AMDKFD_IOC_CREATE_QUEUE(KFDIface.kfd, ring_base_address=ring.va_addr, ring_size=ring.size, gpu_id=self.gpu_id,
       queue_type=queue_type, queue_percentage=kfd.KFD_MAX_QUEUE_PERCENTAGE|(xcc_id<<8), queue_priority=getenv("AMD_KFD_QUEUE_PRIORITY", 7),
       eop_buffer_address=eop_buffer.va_addr if eop_buffer else 0, eop_buffer_size=eop_buffer.size if eop_buffer else 0, ctl_stack_size=ctl_stack_size,
-      ctx_save_restore_address=cwsr_buffer.va_addr if cwsr_buffer else 0, ctx_save_restore_size=ctx_save_restore_size, sdma_engine_id=sdma_engine_id,
+      ctx_save_restore_address=cwsr_buffer.va_addr if cwsr_buffer else 0, ctx_save_restore_size=ctx_save_restore_size,
       write_pointer_address=gart.va_addr+wptr, read_pointer_address=gart.va_addr+rptr+8*xcc_id)
 
     if not hasattr(self, 'doorbells'):
@@ -846,7 +842,7 @@ class KFDIface:
 
 class PCIIface(PCIIfaceBase):
   def __init__(self, dev, dev_id):
-    super().__init__(dev, dev_id, vendor=0x1002, devices=((0xffff, (0x74a1,0x744c,0x7480,0x7550,0x7551,0x7590,0x75a0)),), vram_bar=0,
+    super().__init__(dev, dev_id, vendor=0x1002, devices=((0xffff, (0x74a1,0x744c,0x7480,0x7550,0x7551,0x7590,0x75a0,0x75b0)),), vram_bar=0,
       va_start=AMMemoryManager.va_allocator.base, va_size=AMMemoryManager.va_allocator.size, dev_impl_t=AMDev)
     self._compute_props()
 
@@ -1100,6 +1096,11 @@ class AMDDevice(HCQCompiled):
 
   def on_device_hang(self): self.iface.on_device_hang()
 
+  def finalize(self):
+    try: super().finalize()
+    finally:
+      if self.is_am(): self.iface.dev_impl.release_vf_access()
+
   def device_props(self): return self.iface.props
 
   def hw_copy_queues(self): return [(f"SDMA:{i}", functools.partial(unwrap(self.hw_copy_queue_t), queue_idx=i)) for i in self.sdma_queues]

tinygrad/runtime/support/am/amdev.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-import ctypes, collections, dataclasses, functools, hashlib, array
+import ctypes, collections, dataclasses, functools, hashlib, array, time, contextlib
 from tinygrad.helpers import mv_address, getenv, DEBUG, lo32, hi32, fetch_fw
 from tinygrad.runtime.autogen import pci
 from tinygrad.runtime.autogen.am import am, fw
@@ -149,6 +149,14 @@ class AMDev:
     self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
     self.vram, self.doorbell64, self.mmio = self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I')
 
+    # MI350X VFs start with most MMIO and VRAM access gated by the host PF. Ask the PF for full access only when discovery isn't readable yet.
+    self.is_vf = bool(self.mmio[0xde5] & 1) # RCC_IOV_FUNC_IDENTIFIER.FUNC_IDENTIFIER
+    self.vf_access_acquired, self.vf_initialized = False, False
+    if self.is_vf:
+      self._vf_mailbox_request(6, 7, data2=2, retries=5, event_timeout=2) # IDH_REQ_GPU_INIT_DATA -> IDH_REQ_GPU_INIT_DATA_READY
+      self._vf_mailbox_request(1, 1, retries=5, event_timeout=2) # IDH_REQ_GPU_INIT_ACCESS -> IDH_READY_TO_ACCESS_GPU
+      self.vf_access_acquired = True
+
     self._run_discovery()
     self._build_regs()
 
@@ -165,21 +173,21 @@ class AMDev:
     self.is_booting = True # During boot only boot memory can be allocated. This flag is to validate this.
     self.init_sw(smi_dev=False)
 
-    self.partial_boot = (self.reg("regSCRATCH_REG7").read() == AMDev.Version) and (getenv("AM_RESET", 0) != 1)
+    self.partial_boot = not self.is_vf and (self.reg("regSCRATCH_REG7").read() == AMDev.Version) and (getenv("AM_RESET", 0) != 1)
     if self.partial_boot and (self.reg("regSCRATCH_REG6").read() != 0 or self.reg(self.gmc.pf_status_reg("GC")).read() != 0):
       if DEBUG >= 2: print(f"am {self.devfmt}: Malformed state. Issuing a full reset.")
       self.partial_boot = False
 
-    # Init hw for IP blocks where it is needed
+    # Init hw for IP blocks where it is needed. PSP and SMU are PF-owned on a VF and must not be reset or reloaded by the guest.
     if not self.partial_boot:
-      if self.psp.is_sos_alive() and self.smu.is_smu_alive():
+      if not self.is_vf and self.psp.is_sos_alive() and self.smu.is_smu_alive():
         self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) & ~pci.PCI_COMMAND_MASTER, 2)
         if self.is_hive():
           if reset_mode: return # in reset mode, do not raise
           raise RuntimeError("Malformed state. Use extra/amdpci/hive_reset.py to reset the hive")
         self.smu.mode1_reset()
       self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
-      self.init_hw(self.soc, self.gmc, self.ih, self.psp, self.smu)
+      self.init_hw(self.soc, self.gmc, self.ih, *(() if self.is_vf else (self.psp, self.smu)))
 
     # Booting done
     self.is_booting = False
@@ -187,13 +195,17 @@ class AMDev:
     # Re-initialize main blocks
     self.init_hw(self.gfx, self.sdma)
 
-    if (max_power:=getenv("AM_POWER_LIMIT", 0.0)) > 0:
-      self.smu.set_power_limit(max_power)
-      self.smu.set_clocks(level=None)
-    else: self.smu.set_clocks(level=-1) # last level, max perf.
-    for ip in [self.soc, self.gfx]: ip.set_clockgating_state()
-    self.reg("regSCRATCH_REG7").write(AMDev.Version)
-    self.reg("regSCRATCH_REG6").write(1) # set initialized state.
+    if not self.is_vf:
+      if (max_power:=getenv("AM_POWER_LIMIT", 0.0)) > 0:
+        self.smu.set_power_limit(max_power)
+        self.smu.set_clocks(level=None)
+      else: self.smu.set_clocks(level=-1) # last level, max perf.
+    if not self.is_vf:
+      for ip in [self.soc, self.gfx]: ip.set_clockgating_state()
+    if not self.is_vf:
+      self.reg("regSCRATCH_REG7").write(AMDev.Version)
+      self.reg("regSCRATCH_REG6").write(1) # set initialized state.
+    self.vf_initialized = self.is_vf
     if DEBUG >= 2: print(f"am {self.devfmt}: boot done")
 
   def init_sw(self, smi_dev=False):
@@ -202,7 +214,8 @@ class AMDev:
     # Memory manager & firmware
     self.mm = AMMemoryManager(self, self.vram_size - self.reserved_vram_size, boot_size=(32 << 20), pt_t=AMPageTableEntry, va_shifts=[12, 21, 30, 39],
       va_bits=48, first_lv=am.AMDGPU_VM_PDB2, va_base=AMMemoryManager.va_allocator.base, reserve_ptable=not self.large_bar,
-      palloc_ranges=[(1 << (i + 12), (2 << 20) if i >= 9 else 0x1000) for i in range(9 * (3 - am.AMDGPU_VM_PDB2), -1, -1)])
+      palloc_ranges=[(1 << (i + 12), (2 << 20) if i >= 9 else 0x1000) for i in range(9 * (3 - am.AMDGPU_VM_PDB2), -1, -1)],
+      paddr_base=(1 << 20) if self.is_vf else 0)
     self.fw = AMFirmware(self)
 
     # Initialize IP blocks
@@ -224,10 +237,52 @@ class AMDev:
 
   def fini(self):
     if DEBUG >= 2: print(f"am {self.devfmt}: Finalizing")
-    for ip in [self.sdma, self.gfx]: ip.fini_hw()
-    self.smu.set_clocks(level=0)
-    self.ih.interrupt_handler()
-    self.reg("regSCRATCH_REG6").write(self.is_err_state) # set finalized state.
+    try:
+      for ip in [self.sdma, self.gfx]: ip.fini_hw()
+      if not self.is_vf: self.smu.set_clocks(level=0)
+      self.ih.interrupt_handler()
+      if not self.is_vf: self.reg("regSCRATCH_REG6").write(self.is_err_state) # set finalized state.
+    finally: self.release_vf_access()
+
+  def release_vf_access(self):
+    if not getattr(self, "vf_access_acquired", False): return
+    # tinygrad retains IDH_REQ_GPU_INIT_ACCESS for direct MMIO/VRAM access, so always release that same lease.
+    with contextlib.suppress(Exception): self._vf_mailbox_request(2, None) # IDH_REL_GPU_INIT_ACCESS
+    self.vf_access_acquired = False
+
+  def __del__(self):
+    # Constructor failures do not reach HCQ finalization; return a partially acquired VF init lease to the PF.
+    self.release_vf_access()
+
+  def _vf_mailbox_request(self, req:int, event:int|None, data1=0, data2=0, data3=0, retries=1, event_timeout=2.0):
+    # Navi VF/PF mailbox protocol from the kernel's mxgpu_nv driver. This requests access only; it never requests a GPU or PCI reset.
+    mmio8, control, trn, rcv = self.mmio.view(fmt='B'), 0xe5e * 4, 0xe56, 0xe5a
+    if mmio8[control+1] & 1: mmio8[control+1] = 2 # acknowledge a stale PF event before transmitting a new request
+    for retry in range(retries):
+      deadline = time.monotonic() + 1
+      while True:
+        mmio8[control] = 0 # clear TRN_MSG_VALID and wait for the old PF acknowledgement to drop
+        if not (mmio8[control] & 2): break
+        if time.monotonic() > deadline: raise TimeoutError("VF mailbox acknowledgement did not clear")
+        time.sleep(0.001)
+
+      for i, val in enumerate((req, data1, data2, data3)): self.mmio[trn+i] = val
+      mmio8[control] = 1
+      deadline = time.monotonic() + 0.5
+      while not (mmio8[control] & 2):
+        if time.monotonic() > deadline: raise TimeoutError(f"VF mailbox request {req:#x} was not acknowledged")
+        time.sleep(0.005)
+      mmio8[control] = 0
+
+      if event is None: return
+      deadline = time.monotonic() + event_timeout
+      while time.monotonic() <= deadline:
+        if self.mmio[rcv] == event:
+          mmio8[control+1] = 2 # acknowledge RCV_MSG_VALID
+          return
+        time.sleep(0.01)
+      if DEBUG >= 2 and retry+1 < retries: print(f"am {self.devfmt}: retrying VF mailbox request {req:#x} ({retry+1}/{retries})")
+    raise TimeoutError(f"VF mailbox request {req:#x} did not receive event {event:#x}")
 
   def recover(self, force=False) -> bool:
     if not force and not self.is_err_state: return False

tinygrad/runtime/support/am/ip.py

@@ -251,9 +251,10 @@ class AM_GFX(AM_IP):
     self.mqd_mc = [self.adev.paddr2mc(mqd_paddr) for mqd_paddr in self.mqd_paddr]
 
   def init_hw(self):
-    # Wait for RLC autoload to complete
-    wait_cond(lambda: self.adev.regCP_STAT.read() == 0 or self.adev.regRLC_RLCS_BOOTLOAD_STATUS.read_bitfields()['bootload_complete'] == 0,
-              value=True, msg="RLC autoload timeout")
+    # Wait for RLC autoload to complete on architectures that expose the bootload status register.
+    if hasattr(self.adev, "regRLC_RLCS_BOOTLOAD_STATUS"):
+      wait_cond(lambda: self.adev.regCP_STAT.read() == 0 or self.adev.regRLC_RLCS_BOOTLOAD_STATUS.read_bitfields()['bootload_complete'] == 0,
+                value=True, msg="RLC autoload timeout")
 
     self.adev.gmc.init_hub("GC", inst_cnt=self.xccs)
     if self.adev.partial_boot: return self.reset_mec()
@@ -297,8 +298,8 @@ class AM_GFX(AM_IP):
 
     self._enable_mec()
 
-    # Set 1 partition
-    if self.xccs > 1: self.adev.psp._spatial_partition_cmd(1)
+    # Set 1 partition on bare metal. A VF must use the spatial partition assigned by its host PF.
+    if self.xccs > 1 and not self.adev.is_vf: self.adev.psp._spatial_partition_cmd(1)
 
   def fini_hw(self): self._dequeue_hqds()
 
@@ -486,7 +487,7 @@ class AM_IH(AM_IP):
     if athub_err or cntlr_err:
       print(f"am {self.adev.devfmt}: fatal hardware error detected: {'RAS_ATHUB_ERR_EVENT ' if athub_err else ''}{'RAS_CNTLR' if cntlr_err else ''}")
 
-      acas = self.adev.smu._aca_read_banks(ue=True) + self.adev.smu._aca_read_banks(ue=False)
+      acas = [] if self.adev.is_vf else self.adev.smu._aca_read_banks(ue=True) + self.adev.smu._aca_read_banks(ue=False)
       for regs in acas:
         acatyp = 'Uncorrectable' if (regs[1] >> 61) & 1 and (regs[1] >> 57) & 1 else 'Correctable'
         hwname = f'{self.adev.hwid_names.get((regs[5] >> 32) & 0xFFF, "")} ({(regs[5] >> 32) & 0xFFF:#03x})'

tinygrad/runtime/support/compiler_qcom.py

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

tinygrad/runtime/support/memory.py

@@ -173,14 +173,16 @@ class MemoryManager:
   va_allocator: ClassVar[TLSFAllocator|None] = None
 
   def __init__(self, dev, vram_size:int, boot_size:int, pt_t, va_bits:int, va_shifts:list[int], va_base:int,
-               palloc_ranges:list[tuple[int, int]], first_lv:int=0, reserve_ptable=False):
+               palloc_ranges:list[tuple[int, int]], first_lv:int=0, reserve_ptable=False, paddr_base:int=0):
     self.dev, self.vram_size, self.va_shifts, self.va_base, lvl_msb = dev, vram_size, va_shifts, va_base, va_shifts + [va_bits + 1]
     self.pte_covers, self.pte_cnt = [1 << x for x in va_shifts][::-1], [1 << (lvl_msb[i+1] - lvl_msb[i]) for i in range(len(lvl_msb) - 1)][::-1]
     self.pt_t, self.palloc_ranges, self.level_cnt, self.va_bits, self.reserve_ptable = pt_t, palloc_ranges, len(va_shifts), va_bits, reserve_ptable
 
-    self.boot_allocator = TLSFAllocator(boot_size, base=0)
-    self.ptable_allocator = TLSFAllocator(round_up(vram_size // 512, 1 << 20) if self.reserve_ptable else 0, base=self.boot_allocator.size)
-    self.pa_allocator = TLSFAllocator(vram_size - (off_sz:=self.boot_allocator.size + self.ptable_allocator.size), base=off_sz)
+    self.boot_allocator = TLSFAllocator(boot_size, base=paddr_base)
+    self.ptable_allocator = TLSFAllocator(round_up(vram_size // 512, 1 << 20) if self.reserve_ptable else 0,
+                                          base=paddr_base + self.boot_allocator.size)
+    off_sz = paddr_base + self.boot_allocator.size + self.ptable_allocator.size
+    self.pa_allocator = TLSFAllocator(vram_size - off_sz, base=off_sz)
     self.root_page_table = pt_t(self.dev, self.palloc(0x1000, zero=not self.dev.smi_dev, boot=True), lv=first_lv)
 
   def _frag_size(self, va, sz, must_cover=True):

tinygrad/uop/symbolic.py

@@ -249,8 +249,6 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
    lambda x: x.const_like(x.vmin) if x.vmin == x.vmax else None),
   (UPat(Ops.RANGE, src=(UPat(Ops.CONST,)), name="x"), lambda x: x.const_like(x.vmin) if x.vmin == x.vmax else None),
   # max folding
-  ((UPat.cvar("a") < UPat.var("b")).where(UPat.var("b"), UPat.cvar("c")), lambda a,b,c: UOp.maximum(a,b) if a.val == c.val else None),
-  ((UPat.var("a") < UPat.cvar("b")).where(UPat.cvar("c"), UPat.var("a")), lambda a,b,c: UOp.maximum(a,b) if b.val == c.val else None),
   (UPat.maximum(UPat.var("x"), UPat.var("y")), lambda x,y: x if x.vmin >= y.vmax else y if x.vmax <= y.vmin else None),
   # TODO: why does this rule break beautiful_mnist?
   #((UPat.var("x")+UPat.var("z")).maximum(UPat.var("y")+UPat.var("z")), lambda x,y,z: x.maximum(y) + z),