add ansipad and PARALLEL contextvar

ci: venv in /opt/venv (#17632 )
default python is now 3.14, not 3.12
2026-08-21 03:26:06 +00:00 · 2026-08-20 16:18:36 -07:00 · 2026-08-20 18:54:54 -04:00 · 2026-08-20 18:49:34 -04:00 · 2026-08-20 18:14:30 -04:00 · 2026-08-20 14:26:44 -07:00
41 changed files with 397 additions and 505 deletions
@@ -4,7 +4,7 @@ inputs:
  python-version:
    description: 'Python version to use'
    required: false
-    default: '' # if you don't set a version, the native python version will be used
+    default: '3.14'
  key:
    description: 'Key for the python cache'
    required: false
@@ -59,6 +59,11 @@ runs:
        echo "OMP_NUM_THREADS=1" >> "$GITHUB_ENV"
        # no buffers should be over 300MB in CI
        echo "MAX_BUFFER_SIZE=300000000" >> "$GITHUB_ENV"
+        if [[ "$RUNNER_OS" == "Linux" ]]; then
+          echo "VIRTUAL_ENV=/opt/venv/${{ inputs.python-version }}" >> "$GITHUB_ENV"
+        else
+          echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
+        fi

    - name: Set up uv
      uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b
@@ -67,7 +72,6 @@ runs:

    - name: Set up Python ${{ inputs.python-version }}
      uses: actions/setup-python@v6
-      if: inputs.python-version != ''
      with:
        python-version: ${{ inputs.python-version }}

@@ -109,15 +113,15 @@ runs:
      if: inputs.deps != ''
      shell: bash
      run: |
-        uv venv .venv
+        uv venv --allow-existing --python ${{ inputs.python-version }} "$VIRTUAL_ENV"
        DEPS="${{ inputs.deps }}"
-        uv pip install --python .venv -e ".[${DEPS// /,}]" ${{ inputs.pydeps }} --torch-backend cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
+        uv pip install --python "$VIRTUAL_ENV" -e ".[${DEPS// /,}]" ${{ inputs.pydeps }} --torch-backend cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
    - name: Install dependencies in venv (without extra)
      if: inputs.deps == ''
      shell: bash
      run: |
-        uv venv .venv
-        uv pip install --python .venv -e . ${{ inputs.pydeps }}
+        uv venv --allow-existing --python ${{ inputs.python-version }} "$VIRTUAL_ENV"
+        uv pip install --python "$VIRTUAL_ENV" -e . ${{ inputs.pydeps }}
    - name: Prune uv cache
      if: github.event_name != 'pull_request'
      shell: bash
@@ -125,11 +129,10 @@ runs:
    - name: Configure venv
      shell: bash
      run: |
-        echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
        if [[ "$RUNNER_OS" == "Windows" ]]; then
-          echo "${{ github.workspace }}/.venv/Scripts" >> "$GITHUB_PATH"
+          echo "$VIRTUAL_ENV/Scripts" >> "$GITHUB_PATH"
        else
-          echo "${{ github.workspace }}/.venv/bin" >> "$GITHUB_PATH"
+          echo "$VIRTUAL_ENV/bin" >> "$GITHUB_PATH"
        fi

    # ******************* apt *******************
@@ -94,7 +94,7 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: "0"
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -108,10 +108,6 @@ jobs:
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
-    - name: Symlink models and datasets
-      run: |
-        mkdir -p weights
-        ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -129,10 +125,6 @@ jobs:
      # just metal for now
      if: ${{ matrix.dev == 'METAL' }}
      run: BENCHMARK_LOG=olmoe JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m olmoe --benchmark --warmup
-    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
-      # only run on machines with multiple gpus
-      if: ${{ matrix.dev != 'METAL' }}
-      run: BENCHMARK_LOG=llama3_beam_4gpu JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -149,7 +141,7 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: "0"
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -182,10 +174,6 @@ jobs:
      # slow on metal
      if: ${{ matrix.dev != 'METAL' }}
      run: time BENCHMARK_LOG=cifar DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
-    - name: Run full CIFAR training steps w 6 GPUS
-      # only run on machines with multiple gpus
-      if: ${{ matrix.dev != 'METAL' }}
-      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -202,7 +190,7 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: "0"
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -227,15 +215,8 @@ 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: Run MLPerf resnet eval on training data
-      run: time BENCHMARK_LOG=resnet_eval MODEL=resnet python3 examples/mlperf/model_eval.py
    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
      run: BENCHMARK_LOG=resnet_10steps DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
-    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: BENCHMARK_LOG=resnet_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
-    - name: Run 10 MLPerf Bert training steps (6 gpu)
-      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -252,7 +233,7 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
+      HCQ2: "0"
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -285,6 +266,60 @@ jobs:
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

+  multigpubenchmark:
+    name: Multi-GPU Benchmarks (DEV=${{ matrix.dev }})
+    runs-on: [self-hosted, "${{ matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
+    strategy:
+      fail-fast: false
+      matrix:
+        dev: ['AMD', 'NV']
+    timeout-minutes: 60
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    env:
+      DEV: ${{ matrix.dev }}
+      HCQ2: "0"
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup (AMD)
+      if: ${{ matrix.dev == 'AMD' }}
+      run: |
+        ./extra/amdpci/setup_python_cap.sh
+        ./extra/hcq/hcq_smi.py amd rmmod
+        ./extra/hcq/hcq_smi.py amd kill_pids
+    - name: Setup (NV)
+      if: ${{ matrix.dev == 'NV' }}
+      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+    - name: Symlink models and datasets
+      run: |
+        mkdir -p weights
+        mkdir -p extra/datasets
+        ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
+        ln -s /raid/datasets/imagenet extra/datasets/imagenet
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: reset process replay
+      run: python3 test/external/process_replay/reset.py
+    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
+      run: BENCHMARK_LOG=llama3_beam_4gpu JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
+    - name: Run full CIFAR training steps w 6 GPUS
+      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    - name: Run MLPerf resnet eval on training data
+      run: time BENCHMARK_LOG=resnet_eval MODEL=resnet python3 examples/mlperf/model_eval.py
+    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
+      run: BENCHMARK_LOG=resnet_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
+    - name: Run 10 MLPerf Bert training steps (6 gpu)
+      # TODO: remove BERT_LAYERS once scheduler is fast
+      run: BENCHMARK_LOG=bert_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+    - name: Run process replay tests
+      uses: ./.github/actions/process-replay
+
  tests:
    name: Tests (DEV=${{ matrix.dev }})
    runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
@@ -410,7 +445,7 @@ jobs:
    - name: UsbGPU tiny tests
      run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
    - name: UsbGPU copy speeds
-      run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+      run: sudo -E PYTHONDONTWRITEBYTECODE=1 SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
    #- name: UsbGPU openpilot test
    #  run: sudo -E PYTHONPATH=. GMMU=0 DEV=USB+AMD 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) install script
@@ -563,8 +598,8 @@ jobs:
    - name: Setup
      run: |
        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py ${{ matrix.dev == 'AMD' && 'amd' || 'nv' }} rmmod
-        ./extra/hcq/hcq_smi.py ${{ matrix.dev == 'AMD' && 'amd' || 'nv' }} kill_pids
+        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} rmmod
+        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} kill_pids
        mkdir -p extra/datasets
        ln -s /raid/datasets/imagenet extra/datasets/imagenet
    - name: setup staging db
@@ -166,7 +166,7 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: windows-${{ matrix.dev }}-minimal
-          deps: testing_unit
+          deps: testing_minimal
          pydeps: ${{ matrix.dev == 'WEBGPU' && 'dawn-python' || '' }}
      - name: Set env
        shell: bash
@@ -97,11 +97,13 @@ if __name__ == "__main__":
  parser = argparse.ArgumentParser()
  backend_subparsers = parser.add_subparsers(dest="backend", required=True, metavar="{nv,amd}", help="Hardware backend to target")

-  nv_parser = backend_subparsers.add_parser("nv", help="NVIDIA GPUs")
+  nv_parser = backend_subparsers.add_parser("nv", aliases=["NV"], help="NVIDIA GPUs")
+  nv_parser.set_defaults(backend="nv")
  nv_commands = nv_parser.add_subparsers(dest="command", required=True)
  add_common_commands(nv_commands)

-  amd_parser = backend_subparsers.add_parser("amd", help="AMD GPUs")
+  amd_parser = backend_subparsers.add_parser("amd", aliases=["AMD"], help="AMD GPUs")
+  amd_parser.set_defaults(backend="amd")
  amd_commands = amd_parser.add_subparsers(dest="command", required=True)
  add_common_commands(amd_commands)

@@ -261,10 +261,11 @@ class AMDProgramData:
  private_segment_size:int; kernargs_segment_size:int; kernargs_alloc_size:int
  enable_dispatch_ptr:int; enable_private_segment_sgpr:int

-_amd_program_cache:dict[tuple[bytes,str], tuple[AMDProgramData,bytes]] = {}
+_amd_program_cache:dict[tuple[bytes, tuple[str, ...]], UOp] = {}
 def amd_build_program(prg:UOp) -> UOp:
  dev = Device[to_tuple(prg.device)[0]] # TODO: rm this
-  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, dev.device))) is None:
+  # key on the full device tuple: the same lib can be built for different device sets, each needs its own program buffer
+  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, to_tuple(prg.device)))) is None:
    image, sections, relocs = elf_loader(lib)
    rodata = next(sh.header.sh_addr for sh in sections if sh.name == ".rodata")
    for off, sym, typ, addent in relocs:
@@ -19,16 +19,33 @@ def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None
@functools.cache
 def custom_fused_qkv_rope_forward(q:UOp, k:UOp, v:UOp, xqkv:UOp, freqs_cis:UOp,
                                  device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
-  code = (pathlib.Path(__file__).parent / "fused_qkv_rope.cpp").read_text()
-  threads = 256
-  thread_idx = UOp.special(threads, "lidx0")
-  block_idx_x, block_idx_y = UOp.special(B, "gidx0"), UOp.special(N, "gidx1")
-  sink = UOp.sink(q.base, k.base, v.base, xqkv.base, freqs_cis.base, thread_idx, block_idx_x, block_idx_y,
-                  arg=KernelInfo(name="fused_qkv_rope_forward"))
-  compile_args = ["-std=c++20", "-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"-DTHREADS_PER_BLOCK={threads}"]
-  lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
+  group_size = H // H_KV
+  q, k, v = q.reshape(B, N, H, D), k.reshape(B, N, H_KV, D), v.reshape(B, N, H_KV, D)
+  xqkv = xqkv.reshape(B, N, H_KV, group_size + 2, D)
+  b, n = UOp.range(B, 0), UOp.range(N, 1)
+  pair = UOp.range(D // 2, 2)
+  even = pair * 2
+  c = freqs_cis[0, n, 0, pair, 0].cast(dtypes.float)
+  s = freqs_cis[0, n, 0, pair, 1].cast(dtypes.float)
+  ordered:UOp|None = None
+  for kvh in range(H_KV):
+    q_out, k_out, v_out = (x.after(ordered) if ordered is not None else x for x in (q, k, v))
+    x_in = xqkv.after(ordered) if ordered is not None else xqkv
+    stores:list[UOp] = []
+    for rep in range(group_size):
+      a = x_in[b, n, kvh, rep, even].cast(dtypes.float)
+      bb = x_in[b, n, kvh, rep, even + 1].cast(dtypes.float)
+      h = kvh * group_size + rep
+      stores += [q_out[b, n, h, even].store((a * c - bb * s).cast(q.dtype)), q_out[b, n, h, even + 1].store((a * s + bb * c).cast(q.dtype))]
+    a = x_in[b, n, kvh, group_size, even].cast(dtypes.float)
+    bb = x_in[b, n, kvh, group_size, even + 1].cast(dtypes.float)
+    stores += [k_out[b, n, kvh, even].store((a * c - bb * s).cast(k.dtype)),
+               k_out[b, n, kvh, even + 1].store((a * s + bb * c).cast(k.dtype)),
+               v_out[b, n, kvh, even].store(x_in[b, n, kvh, group_size + 1, even]),
+               v_out[b, n, kvh, even + 1].store(x_in[b, n, kvh, group_size + 1, even + 1])]
+    ordered = UOp.group(*stores)
+  assert ordered is not None
+  return ordered.end(pair, n, b).sink(arg=KernelInfo(name="fused_qkv_rope_forward"))

@functools.cache
 def custom_fused_qkv_rope_backward(dxqkv:UOp, dq:UOp, dk:UOp, dv:UOp, freqs_cis:UOp,
@@ -109,49 +126,6 @@ def fused_qkv_rope(xqkv:Tensor, freqs_cis:Tensor, n_heads:int, n_kv_heads:int, h
 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 grad(dou:UOp, ker:UOp) -> tuple:
    do = Tensor(dou, device=dou.device)
@@ -160,8 +134,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,8 +144,6 @@ 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]

@@ -269,7 +269,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p

    qo_tile<D, float> q_reg_fl;
    load<1, qo_tile<D, float>, _gl_QKVO>(q_reg_fl, g.Qg, {batch_idx, tile_idx, head_idx, 0});
+    #if !WINDOW
    mul(q_reg_fl, q_reg_fl, TEMPERATURE_SCALE);  // Use sqrtf for clarity
+    #endif
    copy(q_reg, q_reg_fl);
    transpose(q_reg_transposed, q_reg);

@@ -288,6 +290,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    zero(att_block[0]);
    transpose(k_reg_transposed, k_reg);
    mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
+    #if WINDOW
+    mul(att_block[0], att_block[0], TEMPERATURE_SCALE);
+    #endif
    __builtin_amdgcn_sched_barrier(0);
    if constexpr (causal) {
        const int kv_end_pos = (min_tile + 1) * KV_BLOCK_SIZE;
@@ -337,6 +342,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
        zero(att_block[1]);
        transpose(k_reg_transposed, k_reg);
        mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]);
+        #if WINDOW
+        mul(att_block[1], att_block[1], TEMPERATURE_SCALE);
+        #endif
 #if WINDOW
        // window masks interior tiles that causal skips
        mask_kv_tile(att_block[1], tile_idx, j - 2, neg_inf_v, lane);
@@ -401,6 +409,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
        zero(att_block[0]);
        transpose(k_reg_transposed, k_reg);
        mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
+        #if WINDOW
+        mul(att_block[0], att_block[0], TEMPERATURE_SCALE);
+        #endif
        //      Finish softmax for QK1
        exp2(att_block[1].tiles[1][0], att_block[1].tiles[1][0]);
        mul(norm_vec, norm_vec, scale_vec);
@@ -469,6 +480,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    zero(att_block[1]);
    transpose(k_reg_transposed, k_reg);
    mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]);
+    #if WINDOW
+    mul(att_block[1], att_block[1], TEMPERATURE_SCALE);
+    #endif
    //      Finish softmax for QK2
    exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]);
    mul(norm_vec, norm_vec, scale_vec);
@@ -535,6 +549,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    zero(att_block[0]);
    transpose(k_reg_transposed, k_reg);
    mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
+    #if WINDOW
+    mul(att_block[0], att_block[0], TEMPERATURE_SCALE);
+    #endif
    //      Finish softmax for QK3
    exp2(att_block[1].tiles[1][0], att_block[1].tiles[1][0]);
    mul(norm_vec, norm_vec, scale_vec);
@@ -597,6 +614,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
    zero(att_block[1]);
    transpose(k_reg_transposed, k_reg);
    mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]);
+    #if WINDOW
+    mul(att_block[1], att_block[1], TEMPERATURE_SCALE);
+    #endif
    //      Finish softmax for QK4
    exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]);
    mul(norm_vec, norm_vec, scale_vec);
@@ -1,69 +0,0 @@
-#include <hip/hip_runtime.h>
-#include <hip/hip_bf16.h>
-
-#ifndef ATTN_B
-#define ATTN_B 2
-#endif
-#ifndef ATTN_N
-#define ATTN_N 8192
-#endif
-#ifndef ATTN_H
-#define ATTN_H 32
-#endif
-#ifndef ATTN_H_KV
-#define ATTN_H_KV 8
-#endif
-#ifndef ATTN_D
-#define ATTN_D 128
-#endif
-#ifndef THREADS_PER_BLOCK
-#define THREADS_PER_BLOCK 256
-#endif
-
-constexpr int GROUP_SIZE = ATTN_H / ATTN_H_KV;
-constexpr int HALF_D = ATTN_D / 2;
-constexpr int PACKED_D = (GROUP_SIZE + 2) * ATTN_D;
-
-extern "C" __global__ __launch_bounds__(THREADS_PER_BLOCK) void
-fused_qkv_rope_forward(
-    __hip_bfloat16*       __restrict__ q,
-    __hip_bfloat16*       __restrict__ k,
-    __hip_bfloat16*       __restrict__ v,
-    const __hip_bfloat16* __restrict__ xqkv,
-    const __hip_bfloat16* __restrict__ freqs_cis) {
-  const int b = blockIdx.x;
-  const int n = blockIdx.y;
-  const int bn = b * ATTN_N + n;
-  const int packed_bn = bn * ATTN_H_KV * PACKED_D;
-  const int q_bn = bn * ATTN_H * ATTN_D;
-  const int kv_bn = bn * ATTN_H_KV * ATTN_D;
-
-  if (threadIdx.x < HALF_D) {
-    const int pair = threadIdx.x;
-    const int even = pair << 1;
-    const float c = static_cast<float>(freqs_cis[((n * HALF_D + pair) * 2) + 0]);
-    const float s = static_cast<float>(freqs_cis[((n * HALF_D + pair) * 2) + 1]);
-
-    for (int kvh = 0; kvh < ATTN_H_KV; kvh++) {
-      const int base = packed_bn + kvh * PACKED_D;
-
-      for (int rep = 0; rep < GROUP_SIZE; rep++) {
-        const int qbase = base + rep * ATTN_D;
-        const int h = kvh * GROUP_SIZE + rep;
-        const float a = static_cast<float>(xqkv[qbase + even]);
-        const float bb = static_cast<float>(xqkv[qbase + even + 1]);
-        const int out = q_bn + h * ATTN_D + even;
-        q[out] = static_cast<__hip_bfloat16>(a * c - bb * s);
-        q[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c);
-      }
-
-      const float a = static_cast<float>(xqkv[base + GROUP_SIZE * ATTN_D + even]);
-      const float bb = static_cast<float>(xqkv[base + GROUP_SIZE * ATTN_D + even + 1]);
-      const int out = kv_bn + kvh * ATTN_D + even;
-      k[out] = static_cast<__hip_bfloat16>(a * c - bb * s);
-      k[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c);
-      v[out] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even];
-      v[out + 1] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even + 1];
-    }
-  }
-}
@@ -67,32 +67,26 @@ class TestParseExpr(unittest.TestCase):

  def test_integer_literals(self):
    """Test parsing integer literals."""
-    self.assertEqual(parse_expr('0', {}).val, 0)
-    self.assertEqual(parse_expr('42', {}).val, 42)
-    self.assertEqual(parse_expr('42U', {}).val, 42)
+    self.assertIs(parse_expr('0', {}), UOp.const(0, dtypes.uint32))
+    self.assertIs(parse_expr('42', {}), UOp.const(42, dtypes.uint32))
+    self.assertIs(parse_expr('42U', {}), UOp.const(42, dtypes.uint32))

  def test_negative_integers(self):
    """Test parsing negative integer literals."""
-    result = parse_expr('-1', {})
-    self.assertEqual(result.val, -1)
-    self.assertEqual(result.dtype, dtypes.int)
+    self.assertIs(parse_expr('-1', {}), UOp.const(-1, dtypes.int))

  def test_float_literals(self):
    """Test parsing float literals."""
-    result = parse_expr('1.0F', {})
-    self.assertEqual(result.val, 1.0)
-    self.assertEqual(result.dtype, dtypes.float32)
+    self.assertIs(parse_expr('1.0F', {}), UOp.const(1.0, dtypes.float32))

  def test_hex_literals(self):
    """Test parsing hex literals."""
-    result = parse_expr('0xFF', {})
-    self.assertEqual(result.val, 255)
+    self.assertIs(parse_expr('0xFF', {}), UOp.const(255, dtypes.uint32))

  def test_variable_lookup(self):
    """Test variable lookup in parse_expr."""
    vrs = {'x': UOp.const(42, dtypes.uint32)}
-    result = parse_expr('x', vrs)
-    self.assertEqual(result.val, 42)
+    self.assertIs(parse_expr('x', vrs), vrs['x'])

  def test_binary_ops(self):
    """Test parsing binary operations."""
@@ -103,9 +97,7 @@ class TestParseExpr(unittest.TestCase):
    self.assertEqual(result.op, Ops.ADD)

    # Subtraction with constant folding
-    result = parse_expr('10 - 5', {})
-    self.assertEqual(result.op, Ops.CONST)
-    self.assertEqual(result.val, 5)
+    self.assertIs(parse_expr('10 - 5', {}), UOp.const(5, dtypes.uint32))

  def test_ternary(self):
    """Test parsing ternary expressions."""
@@ -142,15 +134,8 @@ class TestForLoopParsing(unittest.TestCase):
    S0 = UOp.const(0, dtypes.uint32)
    _vrs, assigns = parse_pcode(pcode, {'S0': S0})

-    # Check that the innermost value (default) is -1 (may be wrapped in CAST)
-    val = assigns[0][1]
-    # Traverse to innermost WHERE
-    while val.op == Ops.WHERE:
-      val = val.src[2]  # false branch
-    # Unwrap CAST if present
-    while val.op == Ops.CAST:
-      val = val.src[0]
-    self.assertEqual(val.val, -1)
+    # every cond folds (S0 is a const), leaving the default branch: -1 in the destination dtype
+    self.assertIs(assigns[0][1].simplify(), UOp.const(-1, dtypes.uint32))

  def test_ctz_parsing(self):
    """Test CTZ pcode parsing."""
@@ -262,8 +247,8 @@ class TestDSPcodePatterns(unittest.TestCase):
    _, assigns = parse_pcode(pcode, srcs)
    # Check addresses: 100 + 2*4 = 108, 100 + 5*4 = 120
    # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp
-    self.assertEqual(assigns[0][1][0].simplify().val, 108)  # type: ignore[index]
-    self.assertEqual(assigns[1][1][0].simplify().val, 120)  # type: ignore[index]
+    self.assertIs(assigns[0][1][0].simplify(), UOp.const(108, dtypes.uint32))  # type: ignore[index]
+    self.assertIs(assigns[1][1][0].simplify(), UOp.const(120, dtypes.uint32))  # type: ignore[index]

  def test_ds_store_data_values(self):
    """Test DS_STORE_2ADDR_B32 uses correct data values."""
@@ -280,8 +265,8 @@ class TestDSPcodePatterns(unittest.TestCase):
    _, assigns = parse_pcode(pcode, srcs)
    # assigns[i][1] is (addr, val) tuple for MEM writes; mypy sees UOp
    # DATA[31:0] should preserve the value
-    self.assertEqual(assigns[0][1][1].simplify().val, 0xAAAAAAAA)  # type: ignore[index]
-    self.assertEqual(assigns[1][1][1].simplify().val, 0xBBBBBBBB)  # type: ignore[index]
+    self.assertIs(assigns[0][1][1].simplify(), UOp.const(0xAAAAAAAA, dtypes.uint32))  # type: ignore[index]
+    self.assertIs(assigns[1][1][1].simplify(), UOp.const(0xBBBBBBBB, dtypes.uint32))  # type: ignore[index]

 class TestConditionalParsing(unittest.TestCase):
  """Test conditional (if/elsif/else) pcode parsing."""
@@ -306,12 +291,12 @@ class TestConcatWidthParsing(unittest.TestCase):
  def test_permlanex16_altrow_concat(self):
    for row, expected in [(0, 1), (1, 0), (2, 3), (3, 2)]:
      parsed = parse_expr('{ row[1], ~row[0] }', {'row': UOp.const(row, dtypes.uint32)})
-      self.assertEqual(parsed.simplify().val, expected)
+      self.assertIs(parsed.simplify(), UOp.const(expected, dtypes.uint32))

  def test_permlane64_altlane_concat(self):
    for lane, expected in [(0, 32), (1, 33), (31, 63), (32, 0), (63, 31)]:
      parsed = parse_expr('{ ~lane[5], lane[4:0] }', {'lane': UOp.const(lane, dtypes.uint32)})
-      self.assertEqual(parsed.simplify().val, expected)
+      self.assertIs(parsed.simplify(), UOp.const(expected, dtypes.uint32))

  def test_permlane64_wave64_pcode_indices(self):
    vgpr = UOp.param(0, dtypes.uint32, (256,))
@@ -327,19 +312,17 @@ class TestConcatWidthParsing(unittest.TestCase):
      'S2': UOp.const(0, dtypes.uint32),
    }

-    def load_idx(v: UOp) -> int:
+    def check_load_idx(v: UOp, expected: int):
      simp = v.simplify()
      self.assertEqual(simp.op, Ops.LOAD)
      self.assertEqual(simp.src[0].op, Ops.INDEX)
-      idx = simp.src[0].src[1].simplify()
-      self.assertEqual(idx.op, Ops.CONST)
-      return idx.val
+      self.assertIs(simp.src[0].src[1].simplify(), UOp.const(expected, dtypes.uint32))

    _, assigns = parse_pcode(PCODE[VOP1Op.V_PERMLANE64_B32_E32], srcs)
    self.assertEqual(len(assigns), 64)
    for lane, (dst_idx, src_idx) in {0: (64, 32), 31: (95, 63), 32: (96, 0), 63: (127, 31)}.items():
-      self.assertEqual(assigns[lane][1][0].simplify().val, dst_idx)  # type: ignore[index]
-      self.assertEqual(load_idx(assigns[lane][1][1]), src_idx)  # type: ignore[index]
+      self.assertIs(assigns[lane][1][0].simplify(), UOp.const(dst_idx, dtypes.uint32))  # type: ignore[index]
+      check_load_idx(assigns[lane][1][1], src_idx)  # type: ignore[index]

 class TestAllPcode(unittest.TestCase):
  """Test that all pcode from all architectures can be parsed."""
@@ -1,4 +1,5 @@
 import unittest
+import functools
 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.helpers import getenv, system, DEV
 from extra.gemm.cdna_asm_gemm import asm_gemm, hk_bf16_atb_gemm
@@ -9,6 +10,7 @@ from examples.mlperf.models.flat_llama import FP8_DTYPE, quantize_fp8, FP8_MAX
 # Use DEV=NULL:HIP:gfx950 to also test the assembly
 def is_cdna4(): return Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950")

+@functools.cache
 def has_hipcc():
  try: system("hipcc --version")
  except Exception: return False
@@ -6,6 +6,8 @@ from tinygrad.tensor import _to_np_dtype
 from tinygrad.runtime.ops_python import from_storage_scalar
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.nir import NIRRenderer
+from tinygrad.renderer.llvmir import CPULLVMRenderer
+from tinygrad.renderer.isa.x86 import X86Renderer
 from tinygrad.uop import Ops
 import numpy as np
 import pytest
@@ -64,6 +66,8 @@ ht.fp8e5m2fnuz = ht.uint8
 def universal_test(a, b, dtype, op):
  if not isinstance(op, tuple): op = (op, op)
  if op[0] == operator.mod and b == 0: return
+  # TODO: throws floating point exception
+  if isinstance(Device[Device.DEFAULT].renderer, (X86Renderer, CPULLVMRenderer)) and op[0] == operator.mod and a == dtype.min and b == -1: return
  # lt and max with nan is undefined in tinygrad
  if op[0] in (operator.lt, Tensor.maximum) and (math.isnan(a) or math.isnan(b)): return
  ta, tb = Tensor([a], dtype=dtype), Tensor([b], dtype=dtype)
@@ -99,22 +99,20 @@ class TestLocalAmax(unittest.TestCase):
    assert_kernel_count(2)
    self.assertEqual(out.tolist(), [[0., 7., 14., 21.], [28., 35., 42., 49.], [120., 135., 150., 165.], [180., 195., 210., 225.]])

-@unittest.skipUnless(has_hipcc() and Device.DEFAULT == "AMD", "requires hipcc to compile and amd device to run")
 class TestFusedQKVRoPE(unittest.TestCase):
  SHAPE = (2, 8192, 32, 8, 128)

+  def setUp(self):
+    if dtypes.bfloat16 not in Device[Device.DEFAULT].renderer.supported_dtypes(): self.skipTest("test uses bf16 inputs")
+
  def rand_bf16(self, *shape:int) -> Tensor:
    return (Tensor.randn(*shape) * 0.1).cast(dtypes.bfloat16).contiguous().realize()

-  def freqs_cis(self) -> Tensor:
-    _, N, _, _, D = self.SHAPE
-    return precompute_freqs_cis(D, N * 2).cast(dtypes.bfloat16).clone().realize()
-
-  def test_llama31_8b_forward(self):
+  def test_forward(self):
    Tensor.manual_seed(0)
-    B, N, H, H_KV, D = self.SHAPE
+    B, N, H, H_KV, D = 1, 32, 8, 2, 16
    GROUP = H // H_KV
-    freqs_cis = self.freqs_cis()
+    freqs_cis = (Tensor.randn(1, N * 2, 1, D // 2, 2) * 0.1).cast(dtypes.bfloat16).contiguous().realize()

    x = self.rand_bf16(B, N, H_KV * (GROUP + 2) * D)
    q, k, v = fused_qkv_rope(x, freqs_cis, H, H_KV, D)
@@ -131,12 +129,13 @@ class TestFusedQKVRoPE(unittest.TestCase):
      self.assertTrue(k.allclose(k_ref, atol=2e-2, rtol=0).item(), "K forward mismatch")
      self.assertTrue(v.allclose(v_ref, atol=0, rtol=0).item(), "V forward mismatch")

-  def test_llama31_8b_backward(self):
+  @unittest.skipUnless(has_hipcc(), "backward kernel requires hipcc to compile")
+  def test_llama31_8b(self):
    Tensor.manual_seed(1)
    B, N, H, H_KV, D = self.SHAPE
    PARTIALS = 2
    GROUP = H // H_KV
-    freqs_cis = self.freqs_cis()
+    freqs_cis = precompute_freqs_cis(D, N * 2).cast(dtypes.bfloat16).clone().realize()
    dq = self.rand_bf16(B, N, H, D)
    dk_partial = self.rand_bf16(B * PARTIALS, N, H_KV, D)
    dv_partial = self.rand_bf16(B * PARTIALS, N, H_KV, D)
@@ -72,9 +72,9 @@ class TestMultiTensor(unittest.TestCase):
    X.shard_(devices_2, 0)
    out = (X + X)
    linear = compile_linear(out.schedule_linear())
-    names = [call.src[0].src[0].arg.name for call in linear.src if call.src[0].op is Ops.PROGRAM]
+    uops = [call.src[0].src[0] for call in linear.src if call.src[0].op is Ops.PROGRAM]
    run_linear(linear)
-    self.assertEqual(len(set(names)), 1, "function was relinearized")
+    self.assertEqual(len(set(uops)), 1, "function was relinearized")

  def test_shard_beam(self):
    cpu_2 = ("CPU:1", "CPU:2")
@@ -822,6 +822,10 @@ class TestOps(unittest.TestCase):
    helper_test_op([], lambda: tor0&tor1, lambda: ten0&ten1, forward_only=True)

    helper_test_op(None, lambda x: (1 < x) & (x < 2), forward_only=True, vals=[[1.2, 1.2, 1.2, 3.2]])
+    helper_test_op([(3000,)]*10, lambda *xs: (sum(xs[1:], xs[0]) > 5) & (xs[0] < 0.9), forward_only=True)
+
+    if not COMPILE_ONLY:
+      np.testing.assert_equal((Tensor(2**64-1, dtype=dtypes.uint64) & 0xFFFFFFFF).numpy(), 0xFFFFFFFF)

  def test_or(self):
    data = [[1,-8,1],[32,1,6]]
@@ -176,7 +176,7 @@ class TestLimitBufs(unittest.TestCase):

  def test_limit_bufs_linear_scaling(self):
    def sched_time(n):
-      with Context(TRACK_MATCH_STATS=0, DEBUG=0):
+      with Context(TRACK_MATCH_STATS=0, DEBUG=0, PARALLEL=0):
        bufs = [Tensor.ones(16).contiguous().realize() for _ in range(4)]
        root = bufs[0]
        for i in range(n): root = root + bufs[i % 4]
@@ -1,6 +1,8 @@
 import unittest, numpy as np
 from unittest.mock import patch
 from tinygrad import Device, Tensor
+from tinygrad.device import Buffer
+from tinygrad.dtype import dtypes
 from tinygrad.helpers import getenv
 from tinygrad.runtime.support.hcq2 import HCQ_DEVS, all_devices_in

@@ -10,6 +12,15 @@ class TestHCQ2(unittest.TestCase):
    with patch.object(Device[Device.DEFAULT], "has_copy_queue", False):
      np.testing.assert_equal(Tensor(np.arange(61, dtype=np.float32)).to(Device.DEFAULT).contiguous().realize().numpy(), np.arange(61))

+  @unittest.skipIf(Device.DEFAULT == "CPU", "staged copies need a non-CPU hcq2 device")
+  def test_staged_copy_slot_reuse(self):
+    # chunks of a staged copy rotate through the staging buffer slots, many rotations must stay bit-exact in both directions
+    import tinygrad.runtime.support.hcq2 as hcq2
+    buf = Buffer("CPU", 1 << 20, dtypes.uint8, preallocate=True)
+    data = np.random.default_rng(42).integers(0, 256, (5 << 20) + 123, dtype=np.uint8)
+    with patch.object(hcq2, "STAGING_SIZE", 1 << 20), patch.object(hcq2, "STAGING_SLOTS", 4), patch.object(hcq2, "_staging", lambda: buf):
+      np.testing.assert_equal(Tensor(data).to(Device.DEFAULT).realize().numpy(), data)
+
  def test_overlapping_device_tuples(self):
    # an op on a wide device tuple followed by an op on an overlapping smaller tuple used to MMU-fault the smaller one
    d4, d2 = tuple(f"{Device.DEFAULT}:{i}" for i in range(4)), tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
@@ -106,7 +106,8 @@ class TestHelpers(unittest.TestCase):

  def test_float_to_bf16(self):
    max_bf16 = torch.finfo(torch.bfloat16).max
-    for a in [1, 1.1, 1234, 23456, -777.777, max_bf16, max_bf16 * 1.00001, -max_bf16, -max_bf16 * 1.00001, math.inf, -math.inf]:
+    for a in [1, 1.1, 1234, 23456, -777.777, max_bf16, max_bf16 * 1.00001, -max_bf16, -max_bf16 * 1.00001,
+              max_bf16 * 2, -max_bf16 * 2, math.inf, -math.inf]:
      self.assertEqual(float_to_bf16(a), torch.tensor([a], dtype=torch.bfloat16).item())
    self.assertTrue(math.isnan(float_to_bf16(math.nan)))

@@ -1020,6 +1020,14 @@ class TestSymbolic(unittest.TestCase):
    self.assertIs(graph_rewrite(cond.where(a, uconst(2)).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), UOp.const(2, dtypes.half)))
    self.assertIs(graph_rewrite(cond.where(a, UOp.invalid()).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), UOp.invalid()))

+  def test_where_const_gate_keeps_stated_width(self):
+    a = Variable("a", 0, 3, dtypes.half)
+    self.assertIs(graph_rewrite(UOp.const(True, dtypes.bool).where(uconst(0.0), a), sym), UOp.const(0.0, dtypes.half))
+    self.assertIs(graph_rewrite(UOp.const(True, dtypes.bool).where(uconst(0), Variable("i", 0, 3, dtypes.int)), sym), UOp.const(0, dtypes.int))
+    self.assertIs(graph_rewrite(UOp.const(False, dtypes.bool).where(uconst(0.0), a), sym), a)
+    self.assertIs(graph_rewrite(UOp.const(False, dtypes.bool).where(uconst(0.0), UOp.invalid()), sym), UOp.invalid())
+    self.assertIs(graph_rewrite(UOp.const(True, dtypes.bool).where(uconst(0.0), uconst(1)), sym), uconst(0.0))
+
  def test_where_merge_branches(self):
    cond1 = Variable("s", 0, 10) < 6
    cond2 = Variable("s", 0, 10) > 2
@@ -43,7 +43,7 @@ def save_viz():
  Buffer.profile_events.clear()
  cpu_events.clear()
  viz = VizTrace()
-  with Context(VIZ=-1, TRACK_MATCH_STATS=2, PROFILE=1):
+  with Context(VIZ=-1, TRACK_MATCH_STATS=2, PROFILE=1, PARALLEL=0):
    yield viz
  viz.set_data()

@@ -236,8 +236,8 @@ class TestViz(unittest.TestCase):
  def test_const_node_visibility(self):
    with save_viz() as viz:
      a = UOp.variable("a", 0, 10, dtype=dtypes.int)
-      z = UOp.const(0, a.dtype)
-      y = UOp.const(math.pi, dtypes.float)
+      z = UOp.const(0)
+      y = UOp.const(math.pi)
      alu = a*z
      ret = exec_rewrite(sink:=UOp.sink(alu, y), [sym])
    lst = viz.list_items()
@@ -249,7 +249,7 @@ class TestViz(unittest.TestCase):
    self.assertTrue(graphs[0][id(y)]["exclude"])
    self.assertFalse(graphs[0][id(alu)]["exclude"])
    self.assertEqual(graphs[0][id(y)]["label"].split("\n")[:2], ["CONST", "3.14159"])
-    self.assertEqual(list(graphs[1]), [id(z), id(y), id(ret)])
+    self.assertEqual(list(graphs[1]), [id(u) for u in ret.toposort()])  # rewrite graph keys follow the rewritten sink's toposort

  def test_const_reshape_expand_folded(self):
    # CONST->EXPAND should be folded into the ALU node, not shown as separate EXPAND nodes
@@ -51,8 +51,8 @@ def simplify_valid_image_load(buf:UOp, idx_y:UOp, idx_x:UOp, valid:UOp) -> UOp|N
  if not drop_stmt and idx is start_idx: return None
  new_valid = UOp.uprod(*ss) if (ss:=[s for s in valid.split_uop(Ops.AND) if s not in drop_stmt]) else None
  idx_y, idx_x = idx.index(1), idx.index(0)
-  if new_valid is not None: return buf.index(idx_y.valid(new_valid), idx_x.valid(new_valid), dtype=dtypes.float)
-  return buf.index(idx_y, idx_x, dtype=dtypes.float)
+  if new_valid is not None: return buf.index(idx_y.valid(new_valid), idx_x.valid(new_valid))
+  return buf.index(idx_y, idx_x)

 indexing_simplify = PatternMatcher([
  # image load valid idx simplification
@@ -88,9 +88,9 @@ def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
  buf = buf.replace(src=(shape_to_shape_arg((h, w, 4)),))
  shapes[buf.arg.slot] = (h, w)
  if valid.op is not Ops.CONST or valid.val is not True:
-    return buf.index(cidx.src[1].valid(valid), cidx.src[0].valid(valid), dtype=dtypes.float)
+    return buf.index(cidx.src[1].valid(valid), cidx.src[0].valid(valid))
  else:
-    return buf.index(cidx.src[1], cidx.src[0], dtype=dtypes.float)
+    return buf.index(cidx.src[1], cidx.src[0])

 pm_simplify_add_image = PatternMatcher([
  (UPat(Ops.SHRINK, src=(UPat(Ops.PARAM, name="buf"), UPat(name="x"), UPat(arg=4))), transform_to_image),
@@ -149,7 +149,7 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
        grp = full_grp[:length]
        # NOTE: we apply the valid again after we determine the length
        offset = offset.valid(valid) if valid is not None else offset
-        idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(len(grp)))) if len(grp) > 1 else buf.index(offset, dtype=offsets[grp[0]][0].src[0].dtype)
+        idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(len(grp)))) if len(grp) > 1 else buf.index(offset)
        if op == Ops.STORE:
          datas = []
          for i,g in enumerate(grp):
@@ -10,10 +10,10 @@ pm_move_gates_from_index = PatternMatcher([
  # for image idx (must be first)
  (UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx_y"), UPat(arg=Invalid)),
                         UPat.var("gate").where(UPat.var("idx_x"), UPat(arg=Invalid))).load(name="l"),
-   lambda buf,gate,idx_y,idx_x,l: buf.index(idx_y, idx_x, dtype=dtypes.float).load(l.vconst_like(0), gate)),
+   lambda buf,gate,idx_y,idx_x,l: buf.index(idx_y, idx_x).load(l.vconst_like(0), gate)),
  (UPat.var("buf").index(UPat.var("gate").where(UPat.var("idx_y"), UPat(arg=Invalid)),
                         UPat.var("gate").where(UPat.var("idx_x"), UPat(arg=Invalid))).store(UPat.var("data")),
-   lambda buf,gate,idx_y,idx_x,data: buf.index(idx_y, idx_x, dtype=dtypes.float).store(data, gate)),
+   lambda buf,gate,idx_y,idx_x,data: buf.index(idx_y, idx_x).store(data, gate)),

  # here we create the alt value for load to be 0s and remove the where Invalid
  (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat(), UPat.var("gate").where(UPat.var("idx"), UPat(arg=Invalid)),), name="mop", allow_any_len=True) \
@@ -1,12 +1,11 @@
 from __future__ import annotations
 import math, itertools
-from collections import defaultdict
-from typing import cast, Final
+from typing import cast
 from tinygrad.uop.ops import Ops, UOp, KernelInfo, graph_rewrite, AxisType, ssimplify, remove_all_tags
 from tinygrad.uop.ops import axis_letters, axis_colors, axis_to_pos
 from tinygrad.device import Buffer
 from tinygrad.dtype import dtypes, Invalid
-from tinygrad.helpers import colored, getenv, DEBUG, to_function_name, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, flatten
+from tinygrad.helpers import colored, getenv, DEBUG, NOOPT, argsort, round_up, prod, merge_dicts, get_single_element, flatten
 from tinygrad.helpers import ALLOW_TF32, count, Context
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError, check
 from tinygrad.codegen.simplify import pm_flatten_range
@@ -48,7 +47,6 @@ class Scheduler:
    if hasattr(self, 'tensor_core'): ret.tensor_core = self.tensor_core
    return ret

-  kernel_cnt: Final[defaultdict[str, int]] = defaultdict(int)
  def get_optimized_ast(self, name_override:str|None=None) -> UOp:
    if name_override is not None: name = name_override
    else:
@@ -56,9 +54,6 @@ class Scheduler:
      special_uops = sorted([x for x in self.ast.toposort() if x.op is Ops.SPECIAL], key=lambda x: x.arg)
      special_ops = [colored(str(x.vmax+1), "blue" if x.arg[0] == "g" else "cyan") for x in special_uops]
      name = k_type + colored('_', 'BLACK').join(['']+special_ops+[colored(x.src[0].render(), color) for x,color in zip(self.rngs, self.colors())])
-      Scheduler.kernel_cnt[(function_name := to_function_name(name))] += 1
-      num = f"n{Scheduler.kernel_cnt[function_name]-1}" if Scheduler.kernel_cnt[function_name] > 1 else ""
-      name += colored(num, 'BLACK')
    self.ast = graph_rewrite(self.ast, pm_flatten_range, name="flatten range")
    return self.ast.replace(arg=KernelInfo(name=name, applied_opts=tuple(self.applied_opts), dont_use_locals=self.dont_use_locals), tag=1)

@@ -42,9 +42,9 @@ def _time_program(prg:UOp, var_vals:dict[str, int], rawbufs:list[Buffer], early_
    global_size, factor = get_test_global_size(prg.arg.global_size, max_global_size, var_vals)
    prg = prg.replace(arg=replace(prg.arg, global_size=tuple(global_size)))
  call = prg.call(*[UOp.from_buffer(b) for b in rawbufs])
-  tms = []
+  tms, timer = [], time_call(call, var_vals, timeout=timeout, clear_l2=clear_l2)
  for _ in range(cnt):
-    try: tms.append(time_call(call, var_vals, timeout=timeout, clear_l2=clear_l2) * factor)
+    try: tms.append(next(timer) * factor)
    except AssertionError: return [math.inf] * cnt
    if early_stop is not None and early_stop < min(tms): break
  return tms
@@ -221,7 +221,7 @@ def float_to_fp16(x):

 def float_to_bf16(x):
  if not math.isfinite(x): return x
-  u = struct.unpack('I', struct.pack('f', x))[0]
+  u = struct.unpack('I', struct.pack('f', truncate[dtypes.float](x)))[0]
  u = (u + 0x7FFF + ((u >> 16) & 1)) & 0xFFFF0000
  return struct.unpack('f', struct.pack('I', u))[0]

@@ -2,7 +2,7 @@ from __future__ import annotations
 from typing import cast, Iterator, Any, Sequence
 import random, itertools, math, weakref, array, decimal
 from dataclasses import dataclass, replace, field
-from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, prod, flatten, Context, getenv, to_tuple
+from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansipad, all_int, prod, flatten, Context, getenv, to_tuple
 from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, perf_counter_us
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite
 from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
@@ -89,7 +89,7 @@ def track_stats(ctx:ExecContext, call:UOp, st:decimal.Decimal, ets:list[float|No
    mem_str = f"{membw*1e-9:4.0f}|{ldsbw*1e-9:<6.0f} GB/s" if membw < 1e13 and ldsbw < 1e15 else \
      colored(f"{membw*1e-12:4.0f}|{ldsbw*1e-12:<6.0f} TB/s", 'green')
    print(f"{colored(f'*** {device[:7]:7s} {GlobalCounters.kernel_count:4d}', header_color)}"+
-      f" {display_name+' '*(46-ansilen(display_name))} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:6.2f} GB"+
+      f" {ansipad(display_name, 46)} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:6.2f} GB"+
      ("" if et is None else f" tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({flops_str} {mem_str})"))
    first_run_cache.add(kcall.src[0].key)

@@ -212,13 +212,13 @@ def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
  dev = cast(Any, Device[(info:= call.arg.aux).device[0]])
  addrs = [(b.bufs[j] if isinstance(b:=_resolve(ctx.input_uops[k], ctx.input_uops).buffer, MultiBuffer) else b).get_buf(dev_name).va_addr
           for devs, idxs in info.input_idxs for j, dev_name in enumerate(devs) for k in idxs]
-  dev.rt_buffer._buf.cpu_view().view(offset=(base:=dev.rt_allocator.alloc(len(addrs) * 8)), fmt='Q')[:len(addrs)] = array.array('Q', addrs)
+  dev.rt_buffer()._buf.cpu_view().view(offset=(base:=dev.rt_allocator.alloc(len(addrs) * 8)), fmt='Q')[:len(addrs)] = array.array('Q', addrs)

  if info.inputs is not None:
-    tables = [UOp.from_buffer(dev.rt_buffer.view(len(idxs), dtypes.uint64, base + j*len(idxs)*8), HCQ_RUNTIME_DEV.value)
+    tables = [UOp.from_buffer(dev.rt_buffer().view(len(idxs), dtypes.uint64, base + j*len(idxs)*8), HCQ_RUNTIME_DEV.value)
              for devs, idxs in info.input_idxs for j in range(len(devs))]
    call = call.substitute({call.src[1+info.inputs]: UOp.mstack(*tables)})
-  exec_kernel(replace(ctx, var_vals={**ctx.var_vals, "hcq_inputs_ptr": dev.rt_buffer._buf.va_addr + base}), call, ast)
+  exec_kernel(replace(ctx, var_vals={**ctx.var_vals, "hcq_inputs_ptr": dev.rt_buffer()._buf.va_addr + base}), call, ast)

  def _prof_tm(device:str, stat_call:UOp, prof:tuple[int, ...]) -> float|None:
    (d:=cast(Any, Device[device])).prof_ents[prof[0]] = ProfileGraphEntry(device, stat_call.arg.name, *prof)
@@ -283,12 +283,13 @@ def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:Sequenc
  ctx = ExecContext(var_vals or {}, tuple(inputs), update_stats, jit, wait or DEBUG>=2)
  for call in linear.src: track_stats(ctx, call, perf_counter_us(), pm_exec.rewrite(call, ctx))

-def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None, clear_l2:bool=False) -> float:
-  if clear_l2:
-    if hasattr(dev:=Device[call.src[1].device], 'invalidate_caches'): dev.invalidate_caches()
-    else:
-      from tinygrad.tensor import Tensor
-      with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024, 1024).contiguous().realize(do_update_stats=False)
+def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None, clear_l2:bool=False) -> Iterator[float]:
  ctx = ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False)
  linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0, profile=True), cache=ctx.cache)
-  return max(et for c in linear.src for et in pm_exec.rewrite(c, ctx) or [0.0])
+  while True:
+    if clear_l2:
+      if hasattr(dev:=Device[call.src[1].device], 'invalidate_caches'): dev.invalidate_caches()
+      else:
+        from tinygrad.tensor import Tensor
+        with Context(DEBUG=0, BEAM=0, CAPTURING=0, TRACK_MATCH_STATS=0): Tensor.ones(1024, 1024).contiguous().realize(do_update_stats=False)
+    yield max(et for c in linear.src for et in pm_exec.rewrite(c, ctx) or [0.0])
@@ -44,6 +44,7 @@ def time_to_str(t:float, w=8) -> str: return next((f"{t * d:{w}.2f}{pr}" for d,p
 def size_to_str(s:int) -> str: return next((f"{s / d:.2f} {pr}" for d,pr in [(1<<30, "GB"),(1<<20, "MB"),(1<<10, "KB")] if s >= d), f"{s} B")
 def ansistrip(s:str): return re.sub('\x1b\\[(K|.*?m)', '', s)
 def ansilen(s:str): return len(ansistrip(s))
+def ansipad(s:str, w:int): return s+' '*max(w-ansilen(s), 0)
 def make_tuple(x:int|Sequence[int], cnt:int) -> tuple[int, ...]: return (x,)*cnt if isinstance(x, int) else tuple(x)
 def to_tuple(x:T|tuple[T, ...]) -> tuple[T, ...]: return x if isinstance(x, tuple) else (x,)
 def flatten(l:Iterable[Iterable[T]]): return [item for sublist in l for item in sublist]
@@ -263,6 +264,9 @@ NUM_CPU_THREADS = ContextVar("NUM_CPU_THREADS", _get_cpu_count())
 NULL_ALLOW_COPYOUT = ContextVar("NULL_ALLOW_COPYOUT", 0)
 # VIZ implies PROFILE, but you can run PROFILE without VIZ
 VIZ = ContextVar("VIZ", 0)
+# this PARALLEL is for BEAM and compilation, it's currently disabled if you are using VIZ
+# pytest-xdist workers share the CPU budget, explicit PARALLEL still overrides this default
+PARALLEL = ContextVar("PARALLEL", NUM_CPU_THREADS.value // max(1, getenv("PYTEST_XDIST_WORKER_COUNT", 1)) if VIZ == 0 else 0)
 PROFILE = ContextVar("PROFILE", abs(VIZ.value))
 SPEC = ContextVar("SPEC", 1)
 # TODO: disable by default due to speed
@@ -490,7 +494,7 @@ def fetch_fw(path:str, name:str, sha256:str) -> bytes:
  if sys.version_info >= (3,14) and (p:=pathlib.Path(f"/lib/firmware/{path}/{name}.zst")).is_file():
    from compression.zstd import decompress
    if hashlib.sha256(b:=decompress(p.read_bytes())).hexdigest() == sha256: return b
-  return fetch(f"https://gitlab.com/kernel-firmware/linux-firmware/-/raw/1e2c15348485939baf1b6d1f5a7a3b799d80703d/{path}/{name}",
+  return fetch(f"https://gitlab.com/kernel-firmware/linux-firmware/-/raw/0a6871b19abf5d6e024b5d208b101ae53e7fa0de/{path}/{name}",
               subdir="fw", sha256=sha256).read_bytes()

 # *** Exec helpers
@@ -585,9 +589,9 @@ class tqdm(Generic[T]):
    est_text = f'<{HMS(elapsed/prog-elapsed) if self.n else "?"}' if self.t else ''
    it_text = (SI(self.n/elapsed) if self.unit_scale else f"{self.n/elapsed:5.2f}") if self.n else "?"
    suf = f'{prog_text} [{HMS(elapsed)}{est_text}, {it_text}{self.unit}/s]'
-    sz = max(ncols-len(self.desc)-3-2-2-len(suf), 1)
+    sz = max(ncols-ansilen(self.desc)-3-2-2-len(suf), 1)
    bar = '\r' + self.desc + (f'{100*prog:3.0f}%|{("█"*int(num:=sz*prog)+" ▏▎▍▌▋▊▉"[int(8*num)%8].strip()).ljust(sz," ")}| ' if self.t else '') + suf
-    print(bar[:ncols+1], flush=True, end='\n'*close, file=sys.stderr)
+    print(bar, flush=True, end='\n'*close, file=sys.stderr)
  @classmethod
  def write(cls, s:str): print(f"\r\033[K{s}", flush=True, file=sys.stderr)

@@ -258,7 +258,8 @@ class ClangRenderer(CStyleLanguage):
  gep_arr_threshold = 0
  has_local = False
  has_threads = bool(getenv("THREADS", 1))
-  global_max = (NUM_CPU_THREADS.value, 0, 0)
+  @property
+  def global_max(self): return (NUM_CPU_THREADS.value, 0, 0)  # type: ignore[override]
  infinity = "__builtin_inff()"
  nan = '__builtin_nanf("")'

@@ -227,8 +227,7 @@ def def_reg(dt:DType, reg:Register|None=None) -> UOp: return UOp(Ops.INS, dt, ar
 def imm(dt:DType, v:int) -> UOp: return UOp.cconst(truncate[dt](v), dt).rtag()
 def to_imm(c:UOp) -> UOp|None:
  if not (c.op is Ops.CAST and (v:=c.src[0]).op is Ops.CONST): return None
-  if c.dtype is dtypes.int64: return imm(dtypes.int32, v.val) if not v.overflows(dtypes.int32) else None
-  if c.dtype is dtypes.uint64: return imm(dtypes.uint32, v.val) if not v.overflows(dtypes.uint32) else None
+  if c.dtype in dtypes.int64s: return imm(dtypes.int32, v.val) if not v.overflows(dtypes.int32) else None
  if c.dtype in dtypes.ints+(dtypes.bool,): return imm(c.dtype, v.val)
  return None
 def cmp(x:UOp) -> UOp:
@@ -638,10 +637,12 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
      if sz == 2: inst += bytes([0x66])
      # bit signaling 64 bit variant of instruction
      w = sz == 8
-      # REX byte is required when 64 bit or an extended reg is used (index 8 - 15) or lower 8 bits of (rsp, rbp, rsi, rdi) are accessed
-      if w | r | _x | b | (reg_sz == 1 & reg >> 2) | (rm_sz == 1 & rm >> 2): inst += bytes([0b0100 << 4 | w << 3 | r << 2 | _x << 1 | b])
      # legacy 8bit opcode is 1 less than 16-64bit variants
-      if (rm_sz == 1 or reg_sz == 1) and x.arg not in X86GroupOp.ReadFlags | {X86Ops.LEA}: opc -= 1
+      demote = (rm_sz == 1 or reg_sz == 1) and x.arg not in X86GroupOp.ReadFlags | {X86Ops.LEA}
+      # REX byte is required when 64 bit or an extended reg is used (index 8 - 15) or lower 8 bits of (rsp, rbp, rsi, rdi) are accessed
+      if w | r | _x | b | (reg_sz == 1 & reg >> 2) | (rm_sz == 1 & rm >> 2) | (demote and disp_uop is None and rm >= 4):
+        inst += bytes([0b0100 << 4 | w << 3 | r << 2 | _x << 1 | b])
+      if demote: opc -= 1
    # OPCODE byte
    inst += opc.to_bytes((opc.bit_length() + 7) // 8, 'big')
    # MODRM byte
@@ -809,7 +810,8 @@ class X86Renderer(ISARenderer):
  device = "CPU"
  has_local = False
  has_threads = bool(getenv("THREADS", 1))
-  global_max = (NUM_CPU_THREADS.value, 0, 0)
+  @property
+  def global_max(self): return (NUM_CPU_THREADS.value, 0, 0)  # type: ignore[override]
  extra_matcher = extra_matcher
  pre_isel_matcher = pre_isel_matcher
  isel_matcher = isel_matcher
@@ -204,7 +204,8 @@ class LLVMRenderer(Renderer):
 class CPULLVMRenderer(LLVMRenderer):
  has_local = False
  has_threads = bool(getenv("THREADS", 1))
-  global_max = (NUM_CPU_THREADS.value, 0, 0)
+  @property
+  def global_max(self): return (NUM_CPU_THREADS.value, 0, 0)  # type: ignore[override]
  abi = 'win64cc' if sys.platform == 'win32' else None
  string_rewrite = base_rewrite
  def render(self, uops: list[UOp]) -> str: return "\n".join((k:=self._render_kernel(uops))[0] + (k[1], self._render_footer(uops)))
@@ -8,7 +8,7 @@ am_src="https://github.com/ROCm/ROCK-Kernel-Driver/archive/33970e1351f5e51102960
 rocm_src="https://github.com/ROCm/rocm-systems/archive/cccc350dc620e61ae2554978b62ab3532dc10bd9.tar.gz"
 AMD, AMDINC = "{}/drivers/gpu/drm/amd", "{}/drivers/gpu/drm/amd/include"
 inc, kern_rules = ["-include", "stdint.h"], [(r'le32_to_cpu', ''),]
-fw_src="https://gitlab.com/kernel-firmware/linux-firmware/-/archive/1e2c15348485939baf1b6d1f5a7a3b799d80703d/1e2c15348485939baf1b6d1f5a7a3b799d80703d.tar.gz"
+fw_src="https://gitlab.com/kernel-firmware/linux-firmware/-/archive/0a6871b19abf5d6e024b5d208b101ae53e7fa0de/0a6871b19abf5d6e024b5d208b101ae53e7fa0de.tar.gz"
 pmc_src="https://raw.githubusercontent.com/ROCm/rocm-systems/cccc350dc620e61ae2554978b62ab3532dc10bd9/projects/rocprofiler-compute/src/rocprof_compute_soc/profile_configs/counter_defs.yaml"

 reg_files = {
@@ -1,81 +1,82 @@
 hashes = {
-  'psp_13_0_0_sos.bin': 'b5592f46885585b935e013f46c949db8ff2f15c0b346caf70e7fcd2776623d13',
+  'psp_13_0_0_sos.bin': '4a51299f6d0a15bbba9694419f7891e6accc01dbd2dd67c06add7bfd75a45ac6',
  'psp_13_0_10_sos.bin': '0bcaaad9cd8578d3841ae69155a6bd4fc3ceae8f4fb5a6ba4f576e7ace94d1d9',
-  'psp_13_0_12_sos.bin': '89da90bf4286b38678b1fd175c78462a426afa3d258d15872cd14072d7098b9b',
-  'psp_13_0_14_sos.bin': 'a4f0d5f76d27b77409ec0b71d7cc6a848ddfd29f8c84f3003edf74ad3999fb7d',
-  'psp_13_0_6_sos.bin': '27657daa0f91ad8095d3610224a7de748b8b348a4cb211ecb5fccabe47369716',
-  'psp_13_0_7_sos.bin': 'ef1af0ecea38abbac6f85cce71789f19848c498d0cb8ef13748dab2d65b23c31',
+  'psp_13_0_12_sos.bin': '7113a165c75c232d4cb7193a920b503e0bf082689adde3b45fdc38f58bfd18b3',
+  'psp_13_0_14_sos.bin': 'db863768cb25e806b68033e9237e0869f9f3603119df4d369ff4d80418d585d0',
+  'psp_13_0_15_sos.bin': '3b28d53e75a88131155e3931378ac8434eca4880ada9211d3b4e8915b6289583',
+  'psp_13_0_6_sos.bin': '36cce3a9441a0dcde81badd8fcf0416de8e4c39a7707865eff4d9d75e6bb0466',
+  'psp_13_0_7_sos.bin': '94db505fa6482f258c33a0a8d412050f6d843ab4ada368252e988f82f8a26fa8',
  'psp_14_0_2_sos.bin': '7b538448b57d4f9dd06b2eea90d4f86a16e65e3027cdecee8db71c2c5f1fa243',
-  'psp_14_0_3_sos.bin': '23bea01a0c6f36d00759d0765d46cb4cb4aa87398b2fbccacbf547a890c0bf51',
-  'smu_13_0_0.bin': '2ffac37fd8534965eeba19755db0e5ec80278213487dc4af0fbc8453befb64b1',
-  'smu_13_0_0_kicker.bin': '7f83656a2a89b7fce1c8a85e96d91cd8265a91fe883a7027f1a0ed18ced501de',
-  'smu_13_0_10.bin': 'daedb9cbdf48942be7ffe00d31b7c16bb36e11ff5a9d7495f218e95c07717b71',
+  'psp_14_0_3_sos.bin': '28469a0857c813c54a0492423cdf0b0caf757428400036377e19c47e5af62478',
+  'smu_13_0_0.bin': '93e46a5526f19dcc3d13bfd9e23f88bc8eee52138bfe9caf0951b4eef5e49914',
+  'smu_13_0_0_kicker.bin': 'd0ef51d9ed06d0c17e06667302be21e7aedd86ed7a72be6e2f55b102214131cc',
+  'smu_13_0_10.bin': '9376ae64149e6b0b684898ffbc12c2230f8c50a2e9447dc7dafc95c0c16b5227',
  'smu_13_0_14.bin': 'a4f36de75fdcecd8000246762e027b4be489b6787afea57675225b0b39d35625',
  'smu_13_0_6.bin': 'ad7232264e8c57c2094244fbdd5a55d7a4575ffe9b44d229884bc0b6a44fb0b1',
-  'smu_13_0_7.bin': 'ccecc0fd0196b9613c920a51c2fd9436e739ff19dda5bdf74d97562387231732',
-  'smu_14_0_2.bin': '6951995d1d606f4dc60c895f19d34ed18aa40e62129f83d8510c45e8aa9ae2fc',
-  'smu_14_0_3.bin': 'df230947ddb7bbfd6e77d1280001db886e69adf2b2a448b47fa668a48bc0009f',
-  'smu_14_0_3_kicker.bin': '8ddc1da5b4e1619796c2cc81f19f388a35bf7d78bfe476cee559625589cb4dc7',
-  'sdma_4_4_2.bin': '456061b814268425843537da6f2191c8861d4e1a18d4c5d90c44ea6be18c78ff',
-  'sdma_4_4_4.bin': 'af47a2940e72b932d3e3a7e8f34f7a182624e5e433f7c56dff939ca5549cd33a',
-  'sdma_4_4_5.bin': '6127baabea3de7b18db3868c983b02c0fbf2cd75997f7f11241a5b1be27e5134',
+  'smu_13_0_7.bin': '68ec18bd605e680085c927ff72c609f8c771aff0718d0cfab58a3744dff8e5b7',
+  'smu_14_0_2.bin': '1b2054e3f710d1ab8dbdf6ff35914ad376b51caa6337831260c955add874b2ee',
+  'smu_14_0_3.bin': '4e1522d3c96c1028be2961dfcfc5f1ff783fb77724b260a99c4c8b4a901ef3fe',
+  'smu_14_0_3_kicker.bin': '9ff142656ae5f57be1b5ecc134e9da8f76650e793fbc5c499acd75094ff24453',
+  'sdma_4_4_2.bin': 'ff885711d2d5d75ceed51cf239e93c882584b918cd5d5d1ff58ee5aecc0c50ae',
+  'sdma_4_4_4.bin': '06a9d4d02c187844313a78469321d6091e59a334f0ce3b61d770d810c984d70b',
+  'sdma_4_4_5.bin': 'e2a30faa9403933fbfba7ce8e9feba460fff6ecdb15304818d24c9f3eeaad0a6',
  'sdma_5_2_6.bin': '3a163db00eb7e4752be8adbd61cf7dd8f08d924e59a6f798ced7dfcd89f340ed',
  'sdma_5_2_7.bin': '16fe80dc866b323e15a06f51646ef0f036878ad34da66921fcdb8167207d6b2b',
-  'sdma_6_0_0.bin': '0f3da6b211f376356335b41be07149f650c10cfa4e23f7e25d53836006ed11f5',
-  'sdma_6_0_1.bin': 'ff565d3c215a30737560d4e3df6fc2c637738407e91d212fb200fdfb185b6744',
-  'sdma_6_0_2.bin': '398380184bb69113ef4c8964a3b55f6184deb0c1ffd96c9683490a3eec3ba8f3',
-  'sdma_6_0_3.bin': '0e8a83513087db865ba926f8b65cfb003fd41098f707e178d7a7ae2941fed0b1',
-  'sdma_6_1_0.bin': '22e55d0ad5f0247a7f0fffc67cfd3161b39f24ad6062ff3c91ec7ff38bd7e1e1',
-  'sdma_6_1_1.bin': '74533a581b8e3e2743b3c9c803d0666405e80898c4a630acefed82cb6b516ba2',
-  'sdma_6_1_2.bin': '4fe04b0286ec739b0414e8aee17e62e85e691f0246d1d9b56bc18a1219072314',
-  'sdma_6_1_3.bin': '35c9ed7e3a237c0d4a83b4975c63b62488f72aeafbb648342f384618e103f66b',
+  'sdma_6_0_0.bin': '82cd01a24171af12de6c7ac4ee7471aa2dfcf51f9677e7bae9cd4c75e07761ee',
+  'sdma_6_0_1.bin': '708c2c2c45262c98ebe8e34e32c3f1ce8eb5b007bab560c9ea9b576a8e4d6768',
+  'sdma_6_0_2.bin': '16c374344d2894da751f7028f9ec1f7520035fa9548d8c750d99a00a4afa86c7',
+  'sdma_6_0_3.bin': 'd47ac4db523aa35d77b27d57c35d4c04f431229dec6c0d667c36d98b985a6933',
+  'sdma_6_1_0.bin': '85f9f3714de68eee74cdf4852d709bc36a5c73a00e943b707bc2ce10d6b7bab4',
+  'sdma_6_1_1.bin': 'e7b7a23923ab691665e6ad16bbc8431a92f7c049da4b0b19a82c45fba03d4979',
+  'sdma_6_1_2.bin': '5947d78eb308a3f6a62d772c5a6493b21439c73eac139f9c22f080f660b4f4f3',
+  'sdma_6_1_3.bin': '8c651f32cbf030b6239ecc44f01bc9f5d5a193f333e21f2103736aff33227361',
  'sdma_7_0_0.bin': 'beaafb53993a106edd392392d5896245ae2a957c6d0f495d0002eec72ad8ad38',
  'sdma_7_0_1.bin': '73c29e1c1714ebc95d2221ba56e187910902891593010653bf9518937e414a59',
-  'gc_10_3_6_pfp.bin': '793d678427887a0e724c79e356440aec33e6d1301f2a4e63543500249ebec064',
+  'gc_10_3_6_pfp.bin': '042f5d2d223aac6a62b500a47d0d0bf33984200110da0ffca4fe5df9a96571c0',
  'gc_10_3_7_pfp.bin': '3ae29aac3f424f7de97f82ce7158beba69509afb2dcbf1a428dc315df474a524',
-  'gc_11_0_0_pfp.bin': 'e175cb0f580a38c961a6f7366142c08e413995f57f78f39795368b15442df8a3',
-  'gc_11_0_1_pfp.bin': 'f5bf21dfbd9e72a30b4caf4704282c27854710c1b7c4affbb2a19530466b12a8',
-  'gc_11_0_2_pfp.bin': '001c4dec1119e29314d725cc1280fc4f0cd9cabdf61ea5ee2260cfd4e62ec141',
-  'gc_11_0_3_pfp.bin': '0488034c85be97125e39e860308d33c3f76a01df8250092a32d4d55acb2526fd',
-  'gc_11_0_4_pfp.bin': '5ae8b7bb6316f87ae8b978354c088e3bd8c890959382d72886377cda25b1ffd1',
-  'gc_11_5_0_pfp.bin': '0124f540871a7759fa8aaae046d458dfb34aeea12a1183ff962c3f1a33067d5a',
-  'gc_11_5_1_pfp.bin': '7794ea46d0d3cf9cb3f7938affbdf09dd7a9970340da5cd02b774cb393436d24',
-  'gc_11_5_2_pfp.bin': '55e64741de28c506524959f7f696713a72aafe46f49ccd827781d67a9475b386',
-  'gc_11_5_3_pfp.bin': 'ce805040fb347fddbc89b2715e66b446865dda9e2056a9b233269b72bc09c387',
-  'gc_12_0_0_pfp.bin': '16bfd64c10fe73b5e760055069a60e5841dba16c0ed4edb56c20d675e23901f6',
-  'gc_12_0_1_pfp.bin': '49efb319305c5fffd90ac1eef7d7a0bdec72998ecb5cf4526996311788a53dc3',
-  'gc_10_3_6_me.bin': '141b59faad3f2f1be16a2178833b7ca8e97519e1e844c8fda6689572c3767902',
+  'gc_11_0_0_pfp.bin': 'b360393c8629144b194f69a3cd961ed509331feff7a5cc1e4eb21c901da2710a',
+  'gc_11_0_1_pfp.bin': 'fb1ee527c05c55679c80a8bcf60fbb533724891baeb0eabc2917fc44e63a45dc',
+  'gc_11_0_2_pfp.bin': '9020f53788ad881fa01aa656fc082f9f8d3cdfc81f70aaac0bed6e6001491128',
+  'gc_11_0_3_pfp.bin': '362db904fa16c1fea2af7ad1295532434df7f85662b4a69332f51ae6c7290b61',
+  'gc_11_0_4_pfp.bin': 'aad22ca342c47d857bc1107a9aa9127e5e4ba7f7fd42d432213b1850bda1f4e1',
+  'gc_11_5_0_pfp.bin': '82ccf0265d841351183b011a79422799431f0c11f6d11165d64d7dfe404bda31',
+  'gc_11_5_1_pfp.bin': '633404d8db1dc03fe997f7d0d0e15ef908069727abaf9de55841be3f3c97348b',
+  'gc_11_5_2_pfp.bin': 'baee1456dd1800cdaedd4998c2dd7d76cdc0cf0ec928679fe67b02485905ea2c',
+  'gc_11_5_3_pfp.bin': 'fee840b049b5e082215df72a93fad80a64f07ef6f638408a2d56fae97449a2cb',
+  'gc_12_0_0_pfp.bin': 'd1b043c60920e509e5c8f9677221fb78ff7985f68b605e8f39a04a57333a9366',
+  'gc_12_0_1_pfp.bin': '9d8d6188efeca5ef05482d9299c4f102fab7db3dae51a23e59de9baa34997123',
+  'gc_10_3_6_me.bin': '776d2299bc4f3abffd4a7999f5a21a4e38aced8b6b4c199a83610dbabf08176d',
  'gc_10_3_7_me.bin': '9eb0b56e9bcc9dad5d53437b162226fcb37e5df102832260f1232832f3658edf',
-  'gc_11_0_0_me.bin': 'f8fba8a63dd4293b8fc1e4aab78b6fac630e575d1d62838c7996d9210f82aea1',
-  'gc_11_0_1_me.bin': '5030040b00955de94876341ec64ea43b96640413d7a03dc460a83c8386bf76e0',
-  'gc_11_0_2_me.bin': '0f21fd43f1dfbc6ccced9a2b3774de25c993c61a689aabab8b45333937b7945e',
-  'gc_11_0_3_me.bin': '3acb5061dba342ade81d329d1932f19ec01f0c5bf44e6e3568008a951a351bac',
-  'gc_11_0_4_me.bin': 'e4f1f6abcd213d54ad9e885d9f550083b0e2f67d983566015e8a53981e1cb155',
-  'gc_11_5_0_me.bin': '8f906b64d0a29503daa662c93ec44d076fcac11b78f70cd50ce0af2b500a05a6',
-  'gc_11_5_1_me.bin': '7e42602bcbaf1e511f8b4f6ed2246844ad1f6e351ce2b663d89062a7be263663',
-  'gc_11_5_2_me.bin': 'aae26255d8efff81e0e3bbcb727efb8b837d8e25fe85c708545f5328f1077b50',
-  'gc_11_5_3_me.bin': '93cd588348b16fe432609fe8da6e6b5da0a52da5c5884882aecf7b1001f72700',
-  'gc_12_0_0_me.bin': 'd7eba5197f2580f32b8256b1d9cb68e723e9e644293a34446a7913e3c093cba5',
-  'gc_12_0_1_me.bin': '365e7f193b39cbb10d3af44905fefaca0e9844721801755276baebac7b19c1ea',
-  'gc_10_3_6_mec.bin': '247943415658159704a21f670dd7b3e7cb2d2fc0c17b000a5098715979c8d95e',
+  'gc_11_0_0_me.bin': 'f2f5a793d811c6abad1a18af0fcf7694c443478f224176da86650c22aa71ca7a',
+  'gc_11_0_1_me.bin': '476db2ec7e33d1e126b1736649208443e3ccc68aa60e4978574cdbced2b26543',
+  'gc_11_0_2_me.bin': 'f5fe48f97acbd3ce13b35929290bfbac01ce522631cc91dcef1fdeb3ff35c8ed',
+  'gc_11_0_3_me.bin': 'd02c25070e5bdf0ec0146f5c9d6d2f8b86de43bd2a318a0b67eb5201963bafdc',
+  'gc_11_0_4_me.bin': 'f075220f75ffe43eacc5986ff8448946c27405e632764fda83323e7ec8d55566',
+  'gc_11_5_0_me.bin': '338019a1fcdab39729e3f492ffc9f5970c2c81b12c8a4f431494ca28cfdadedf',
+  'gc_11_5_1_me.bin': '4c4dd30c22d4f7f2c5d3a19c645f505e30cdac115a91c65791e2651b22932175',
+  'gc_11_5_2_me.bin': 'cab2999186d26c0e9a3d46b5a43d2854d88be880cb764c096ad2b43038566384',
+  'gc_11_5_3_me.bin': '94e2d74e834725b3d51e03e830160e95c56f3a31e93f5d61c852afd8fe8cc779',
+  'gc_12_0_0_me.bin': 'fb10cb3535ae4a6a8fb3e78166cf30c5b717341b1f20cde73065c62b642adfed',
+  'gc_12_0_1_me.bin': '56a1ae0031aa938f6b61348a56404ab2cee92f1f45630fc82a801aa4d908f98a',
+  'gc_10_3_6_mec.bin': '7003c4a77537e9edaf67064104cd9371fac38a84f71f948349140b28d3c210e8',
  'gc_10_3_7_mec.bin': 'ee58a523375bcf5b89400b32b801f95e182b632a26bce4f2bed5c07928d486dc',
-  'gc_11_0_0_mec.bin': '801a09c9bf06188260db9b51ad8f978f15d84c72ca91b90643a2ef8af4074776',
-  'gc_11_0_1_mec.bin': '6afadcb7504bb11bcc9d4a205cdf73f7934a615e28f178fcf7285971df2ccd05',
-  'gc_11_0_2_mec.bin': '0da0edee28c73a6fa1191f77853d380ec2503cbf43e0aaae4617f32f1f8a48fa',
-  'gc_11_0_3_mec.bin': '323cfa6658b6b5169830f852e2ff0552acae8dfb9e44b42c63de7b2900d3fd9e',
-  'gc_11_0_4_mec.bin': '5d89cf6b60354f3746c2cbd1ff0cb1a741556ca20d72745242cb69b553d0985c',
-  'gc_11_5_0_mec.bin': 'a01c324ab14ec89792449a621a541829b9af26865019027a411a14b910145dfa',
-  'gc_11_5_1_mec.bin': 'eab05719371caa68df09d4f7574e3958a3c4f5044ab3c7b0d2b214add0c6d1c4',
-  'gc_11_5_2_mec.bin': 'a374b2335802e24f8b9a3ce40000a1d37a52a14eb87099bebcc6680c27cc93e5',
-  'gc_11_5_3_mec.bin': '165025437cba80dd32c19ebbc83b756fa7adac7053ff7780ba4aa2f8089c6a3f',
-  'gc_12_0_0_mec.bin': '1931593440b8f9423580d9e2cdc5b34e7c682cdffe1ca4b74b0c2f6a0420236d',
-  'gc_12_0_1_mec.bin': 'f57541688a5108730bf210663f1137ffc2121f3acfe614a6de09ec1982c69a2f',
-  'gc_9_4_3_mec.bin': '3159176e72301fb88dc416721fb3d0ab82ece484cf93a43c3f37430c7e6673a1',
-  'gc_9_4_3_sjt_mec.bin': 'd19468dbb47849640bd0e6cdc8d7e25a3c8442c7ca2ca81357702e0d6baab50f',
-  'gc_9_4_4_mec.bin': '5004f73e43db2dd45e77d65942e33d4a69e7157618cfd23944c30f801c77a0f3',
-  'gc_9_4_4_sjt_mec.bin': '627a9e98102e70fe3bf0947eb764187f29f5e775d1130c7310e0ba5fc0502dbe',
-  'gc_9_5_0_mec.bin': 'c5eca4311a6f6e8f81cf41c2c46941d5dcf90789ee8326901da2dfc86ac14c31',
-  'gc_9_5_0_sjt_mec.bin': 'f162e509379288e3f3b1eead541b315c2262d625d433287ecd34ca185614d312',
+  'gc_11_0_0_mec.bin': '1dd1de8ecf5455ea4719c502b64b32ac18763d5601128c01b4a4a36211a122c2',
+  'gc_11_0_1_mec.bin': '505ae64eccb2e4b4751fe18ec1b584e1f6b4c81d0f5ec089afbcf378cad59711',
+  'gc_11_0_2_mec.bin': '19bf080d6e672de5ed3fb86e3fdbdda4d700d8e3bda2dbdcc923101484ad645b',
+  'gc_11_0_3_mec.bin': 'a37bc1a4e245300a5c3e26da34ea213842447d7df6c5c81e9fc78887a2fde26f',
+  'gc_11_0_4_mec.bin': '850d5302b4fee6022f42f706c2de103531b45b7794a45f2d6dce6015767a1ad6',
+  'gc_11_5_0_mec.bin': '5e022bae6638967d82e2b1077e3024f52bc83b3cb850aa31fba51469c7517c4c',
+  'gc_11_5_1_mec.bin': 'e49964d5e58686c53e66d98d4e3b9fab70e98fad3b28379c6e60aed03c83ee80',
+  'gc_11_5_2_mec.bin': '9691d7bff5d2c933d8eecb7d171635612a76a2dd1441cffcd65a8a02bdb5a2c5',
+  'gc_11_5_3_mec.bin': 'd368f3886b9245dd0d21d57fccfd8aa7e872c2564e23f292abe735348121277e',
+  'gc_12_0_0_mec.bin': '9c7602d6ebf1f7e6ec7a5d1ceefded18f35fa1c08fbea1e3e1a0d78d519db8e8',
+  'gc_12_0_1_mec.bin': 'caf1dbaf72b0ef0c4c973947414033aeec002994f63967bb53e9165195a3c2c3',
+  'gc_9_4_3_mec.bin': '99bc12230f00b930cf286105a35cc6110d87461cd48cb4fdf3cb6caff73ac1e7',
+  'gc_9_4_3_sjt_mec.bin': '2945dbd098c4158870df7dc4ccb33d40031fd1cce37cdbe5df291d8941d03567',
+  'gc_9_4_4_mec.bin': '7f14258f8301d2717e0a707ccfad7b3091af478b0df6d5134adfd56caa7429d8',
+  'gc_9_4_4_sjt_mec.bin': '0bbef279bbc07c502098b80765b876f69fcda9834e5ed269a7d8236c85e89e19',
+  'gc_9_5_0_mec.bin': '0c39078c53e10e99538901df5fc14e7f1b1f3639ea825b1b3126ae87a28b2464',
+  'gc_9_5_0_sjt_mec.bin': 'a769745367567fc6f389695aa5f48c154c07560e21a93052185e19f950205240',
  'gc_11_0_0_imu.bin': 'b4f8fc056b45709a6abf48e7885fb1b4ab8d3cc092cbfa2c554a78564a6403bc',
  'gc_11_0_1_imu.bin': 'ac71f4eec713fc35b4a1fe27531e3eb04edd81eeac2cef64df01ac50d8510805',
  'gc_11_0_2_imu.bin': '9befca62b0b0cfd252c3df4a9edca295526f4d43821cd99a6326454995a6ca2d',
@@ -90,17 +91,17 @@ hashes = {
  'gc_10_3_6_rlc.bin': 'acfbac75c0dcfbfe40e222640ef17eb3dc8d206d30bc3863f275f2dd1cb132a5',
  'gc_10_3_7_rlc.bin': 'a02585ebe3b36d942e883057119572d9497600c52fc65b8a523487eb65d874f2',
  'gc_11_0_0_rlc.bin': 'dabd49039772d02f5fd5e48dc21d35ad52a6b1283b470dabca86ca159c4c7c8e',
-  'gc_11_0_1_rlc.bin': '86145719a58e9428562930c6b5ee3b6ced4701d34a80d0b4d84d6026c93134f2',
+  'gc_11_0_1_rlc.bin': '5f07dc1f0a75ecd9cb56d805ea869184a50ed9e43d811ebf833b8906534650ef',
  'gc_11_0_2_rlc.bin': 'b43eb2fd0600f50a1a5796bc9983d6b39b5c20960234920f5e89cb362193e0b8',
-  'gc_11_0_3_rlc.bin': '29b0b456f5b53076ddffa6f09de3bb697219e8e7b33504bf6c197e8b858426dc',
-  'gc_11_0_4_rlc.bin': '823573078b608108fbe4dd8176c396ec582632913db9c59a512d82b068f8eba0',
-  'gc_11_5_0_rlc.bin': '68cd85567f4f2f8d6b80db294988806d956bf826979c3597daccb71c7ee6aadd',
+  'gc_11_0_3_rlc.bin': '890d8e0123efb40c0179dd8ac3e9af073a0b87cbbccfec1db54e5ed2315a8d39',
+  'gc_11_0_4_rlc.bin': '257ced82d7bec41249b06592ee0c44fb8f9262de2c6af9c52dc6f6a8a702063e',
+  'gc_11_5_0_rlc.bin': '0dc8b6ef5530a4a53938c8baa0d49cd458607d95233237859fa98d44feb3e985',
  'gc_11_5_1_rlc.bin': '92731ecabbeb77865fb71787b4268dc738a58779f1190bdc2056482cb88a08f6',
-  'gc_11_5_2_rlc.bin': 'ef3a9209d3eccfbe18fce9e972c146ac283719798bb788096c176b796dc9aee5',
+  'gc_11_5_2_rlc.bin': 'c9ad70b8ac309257cb8929bb6b4efa6b551ec1e5229d7a419332a9797f31fc9e',
  'gc_11_5_3_rlc.bin': '10a68940c6258d5818d9c05fd98eb0ccc8d5aee99b2769fbad30e5abd0d9327e',
  'gc_12_0_0_rlc.bin': '6436b582734a413456fff3d3c7195e71cc9e78a7ed31ee21c83ffd6fae1ad186',
  'gc_12_0_1_rlc.bin': '6ba4459532246a5c415d3cb33c9b1248294e48f67b827e2accb292a8d1a5c0ec',
-  'gc_9_4_3_rlc.bin': '5345d388712d547b0ae16f199ad5ccadb65643584b3efa7817049ddeb3fdcd12',
+  'gc_9_4_3_rlc.bin': '54cbd0de3a0ec35d2e58e992babeee2a237f870ccdf37e734652e4daeeba59d5',
  'gc_9_4_4_rlc.bin': 'e0c3585c72f8136670ca63e607fba32c1ae4948f493f13e33fc4d466bd6318a8',
  'gc_9_5_0_rlc.bin': '9b1268f5751153fe57f527c9acb417bfa53ed42c9bc083c9d3da2ba61fe5fdc4',
 }
@@ -7,7 +7,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, H
 from tinygrad.runtime.support.hcq import MMIOInterface, BumpAllocator, hcq_filter_visible_devices, hcq_profile
 from tinygrad.uop.ops import sint
 from tinygrad.device import Compiled, BufferSpec, TinyELF
-from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, colored, prod, ContextVar, TracingKey, from_mv
+from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, colored, prod, ContextVar, TracingKey
 from tinygrad.helpers import VIZ, ceildiv, unwrap, pluralize
 from tinygrad.renderer.cstyle import HIPRenderer, HIPCCRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
@@ -649,94 +649,6 @@ class AMDAllocator(HCQAllocator['AMDDevice']):

  def _do_map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)

-  def _copyin(self, dest:HCQBuffer, src:memoryview):
-    if not self.dev.is_usb(): return super()._copyin(dest, src)
-    import time as _time
-    from tinygrad.runtime.support.usb import alloc_cbuffer
-    # Pipelined USB copyin on stock firmware. The host streams 240KB chunks over EP 0x02 into two alternating
-    # 256KB SRAM bounce buffers; every 120KB part is followed by a 512B sentinel sector with a unique value.
-    # The prebuilt SDMA ring polls each sentinel before copying its part to VRAM (GPU blocks until the engine
-    # lands data) and writes a per-buffer drain fence after each chunk, which the host checks before re-arming
-    # that buffer (USB blocks until the GPU drains it). One doorbell per copyin; the engine latches one
-    # mid-flight arm, so chunk c+1 is armed and queued while chunk c streams.
-    dev, usb = self.dev, self.dev.iface.pci_dev.usb
-    sdq, ts, sdma = dev.sdma_queue(0), dev.timeline_signal, dev.sdma
-    with hcq_profile(self.dev, queue_type=dev.hw_copy_queue_t, desc=TracingKey(f"TINY -> {dev.device}", ret=src.nbytes), enabled=PROFILE,
-                     dev_suff="SDMA:0"):
-      CP, PART = 0x3C000, 0x1E000  # 240KB chunks of two 120KB parts; sentinels ride the buffer's last SRAM slot
-      src_mv, nchunks = src.cast('B'), ceildiv(src.nbytes, 0x3C000)
-      if nchunks == 0: return
-      slots = [(self.b[bi].cpu_view().addr - 0xf000) >> 14 for bi in range(len(self.b))]  # 0xF2 slot base per buffer
-      if not hasattr(self, '_usb_seq'):
-        self._usb_seq, self._usb_prev_tail = 0, (0, 0)
-        self._usb_buf_ctr = [0] * len(self.b)  # per-buffer drain fence values
-        self._usb_stage = [alloc_cbuffer(0x40000) for _ in range(2)]  # wire-image staging, one per buffer
-        for bi in range(len(self.b)):
-          usb.write(0xA800 + bi * 8, bytes(8))  # clear the drain fences (SQ window: PCIe 0x820800 = xdata 0xA800)
-          usb.scsi_write(bytes(0x40000), slot_start=slots[bi])  # zero the buffer: sentinel EQ polls need zero-init
-
-      seq_base = self._usb_seq
-      ctr_base = list(self._usb_buf_ctr)  # chunk c lands on buffer c%2 with fence value ctr_base[c%2] + c//2 + 1
-      def sval(c, j): return 0x51000000 | ((seq_base + 2 * c + j) & 0xFFFFFF)  # unique per part, across copyins
-
-      def wait_full(bi, ctr):  # wait until the GPU fully drained chunk `ctr` of buffer `bi`
-        deadline = _time.monotonic() + 10.0
-        while int.from_bytes(usb.read(0xA800 + bi * 8, 8), 'little') < ctr:
-          if _time.monotonic() > deadline: raise RuntimeError("USB copyin GPU fence timeout")
-
-      # The whole ring: wait on pre-copyin GPU work, then per chunk, per part: [POLL sentinel][copy][fences].
-      q = dev.hw_copy_queue_t().wait(ts, dev.timeline_value - 1)
-      for c in range(nchunks):
-        bi, lsize = c % len(self.b), min(CP, src.nbytes - c * CP)
-        for j in range(2 if lsize > PART else 1):
-          plen = min(PART, lsize - j * PART)
-          q.q(sdma.SDMA_OP_POLL_REGMEM | sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(3) | sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1),
-              *data64_le(self.b[bi].va_addr + j * (PART + 512) + round_up(plen, 512)), sval(c, j), 0xFFFFFFFF,
-              sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))
-          q.copy(dest.offset(c * CP + j * PART), self.b[bi].offset(j * (PART + 512)), plen)
-        q.signal(ts, dev.next_timeline())
-        q.write(dev.iface.sys_buf.offset(0x800 + bi * 8, 8), ctr_base[bi] + c // 2 + 1, b64=True)
-
-      # Stage the ring in one EP2 write while the engine is disarmed. The GPU ring is circular, dword 0 is a NOP.
-      cmds = array.array('I', q._q).tobytes()
-      rb, pv_start = sdq.ring.nbytes, sdq.put_value
-      off = pv_start % rb
-      if off + len(cmds) > rb:  # zero-fill to the ring end and continue at offset 0
-        sdq.ring.view(off, rb - off, fmt='B')[:] = bytes(rb - off)
-        pv_start, off = pv_start + rb - off, 0
-      sdq.ring.view(off, len(cmds), fmt='B')[:] = cmds
-      sdq.put_value = pv_start + len(cmds)
-
-      def push_chunk(c):
-        bi, lsize = c % len(self.b), min(CP, src.nbytes - c * CP)
-        nparts = 2 if lsize > PART else 1
-        sbuf, smv = self._usb_stage[c % 2]
-        sptr = ctypes.addressof(from_mv(src_mv[c * CP:c * CP + lsize]))
-        for j in range(nparts):  # stage the wire image [part][sentinel]...
-          plen = min(PART, lsize - j * PART)
-          ctypes.memmove(ctypes.addressof(sbuf) + j * (PART + 512), sptr + j * PART, plen)
-          struct.pack_into('<I', sbuf, j * (PART + 512) + round_up(plen, 512), sval(c, j))
-        total = (nparts - 1) * (PART + 512) + round_up(lsize - (nparts - 1) * PART, 512) + 512
-        wait_full(bi, ctr_base[bi] + c // 2)  # the GPU must have fully drained chunk c-2 before this re-arm
-        usb.scsi_write_arm(total, slot_start=slots[bi])
-        return usb.usb.bulk_write_async(smv[:total])  # data is queued now but only streams once armed
-
-      wait_full(*self._usb_prev_tail)  # previous copyin's tail must be done before its ring entries are overwritten
-      sdq.write_ptr[0] = sdq.put_value
-      sdq.doorbell[0] = sdq.put_value
-
-      pending = [push_chunk(0)]
-      if nchunks > 1: pending.append(push_chunk(1))
-      for c in range(1, nchunks):
-        usb.bulk_wait(pending.pop(0))  # chunk c-1 has landed in SRAM
-        if c + 1 < nchunks: pending.append(push_chunk(c + 1))  # arm+queue the next chunk while chunk c streams
-      usb.bulk_wait(pending.pop(0))
-      tbi = (nchunks - 1) % len(self.b)
-      self._usb_prev_tail = (tbi, ctr_base[tbi] + (nchunks - 1) // 2 + 1)
-      wait_full(*self._usb_prev_tail)  # the last chunk must be copied before returning
-      self._usb_seq = (seq_base + 2 * (nchunks - 1) + (2 if lsize > PART else 1)) & 0xFFFFFF
-      self._usb_buf_ctr = [ctr_base[0] + (nchunks + 1) // 2, ctr_base[1] + nchunks // 2]
-
  def _copyout(self, dest:memoryview, src:HCQBuffer):
    if not self.dev.is_usb(): return super()._copyout(dest, src)
    self.dev.synchronize()
@@ -930,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,0x75a8)),), vram_bar=0,
      va_start=AMMemoryManager.va_allocator.base, va_size=AMMemoryManager.va_allocator.size, dev_impl_t=AMDev)
    self._compute_props()

@@ -1004,10 +916,7 @@ class USBIface(PCIIface):
    self._compute_props()

    # special regions
-    # Two 256KB SRAM bounce buffers (16 slots each), so USB bulk IN flight for chunk i overlaps the SDMA of chunk i-1.
-    # ctrl_addr tags encode the 16KB slot base for the 0xF2 engine: 0xf000 -> slot 0, 0x4f000 -> slot 16.
-    self.copy_bufs = [self._dma_region(ctrl_addr=0xf000, sys_addr=0x200000, size=0x40000),
-                      self._dma_region(ctrl_addr=0x4f000, sys_addr=0x240000, size=0x40000)]
+    self.copy_bufs = [self._dma_region(ctrl_addr=0xf000, sys_addr=0x200000, size=0x80000)]
    self.sys_buf, self.sys_next_off = self._dma_region(ctrl_addr=0xa000, sys_addr=0x820000, size=0x1000), 0x200
    self.cq_buf = self._dma_region(ctrl_addr=0xb800, sys_addr=0x822000, size=0x1000)

@@ -1017,8 +926,9 @@ class USBIface(PCIIface):

  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, **kwargs) -> HCQBuffer:
    # usb allocates uncached and cpu_access in vram. vram writes are faster than sram writes
-    # NOTE: host allocs deliberately do NOT use sys_buf (the 0x820000 NVMe SQ region): the GPU's signal writes there
-    # collide with the 0xF2 engine mid-stream. Signals in VRAM are read back via 0xF0 streaming reads instead.
+    if host and self.sys_next_off + size < self.sys_buf.size:
+      self.sys_next_off += size
+      return self.sys_buf.offset(self.sys_next_off - size, size)

    # force devmem
    return super().alloc(size, host=False, uncached=uncached, cpu_access=cpu_access, contiguous=contiguous, force_devmem=True, **kwargs)
@@ -1138,8 +1048,7 @@ class AMDDevice(HCQCompiled):
    if getenv("AMD_DISABLE_SDMA"): return None
    if idx in self.sdma_queues: return self.sdma_queues[idx]
    with contextlib.suppress(OSError):
-      # USB: 1MB ring, so a full copyin can be staged in one write. GPU rings are circular and dword 0 is an SDMA NOP.
-      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, (1 << 20) if self.is_usb() else (16 << 20), idx=idx)
+      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20), idx=idx)
    return self.sdma_queues.get(idx, None)

  def _ensure_has_local_memory(self, private_segment_size):
@@ -6,6 +6,8 @@ class NpyAllocator(Allocator['NpyDevice']):
  def _alloc(self, size:int, options=None) -> np.ndarray: return np.empty(size, dtype=np.uint8)
  def _as_buffer(self, src:np.ndarray) -> memoryview: return flat_mv(np.require(src, requirements='C').data)
  def _copyout(self, dest:memoryview, src:np.ndarray): dest[:] = self._as_buffer(src)
+  def _offset(self, buf:np.ndarray, size:int, offset:int) -> np.ndarray:
+    return np.require(buf, requirements='C').reshape(-1).view(np.uint8)[offset:offset+size]

 class NpyDevice(Compiled):
  def __init__(self, device:str): super().__init__(device, NpyAllocator(self), [], None)
@@ -1,6 +1,6 @@
 from __future__ import annotations
 import ctypes, collections, dataclasses, functools, hashlib, array
-from tinygrad.helpers import mv_address, getenv, DEBUG, lo32, hi32, fetch_fw
+from tinygrad.helpers import mv_address, getenv, DEBUG, lo32, hi32, fetch_fw, to_mv
 from tinygrad.runtime.autogen import pci
 from tinygrad.runtime.autogen.am import am, fw
 from tinygrad.runtime.support.amd import AMDReg, import_module, import_asic_regs
@@ -238,7 +238,8 @@ class AMDev:
    if DEBUG >= 3: print(f"am {self.devfmt}: Recovery complete")
    return True

-  def is_hive(self) -> bool: return self.gmc.xgmi_seg_sz > 0
+  # a hive has multiple XGMI regions; single-node parts (like MI350P) may still program LFB_SIZE with region 0 only
+  def is_hive(self) -> bool: return self.gmc.xgmi_seg_sz > 0 and self.gmc.xgmi_max_region > 0

  def paddr2mc(self, paddr:int) -> int: return self.gmc.mc_base + paddr
  def paddr2xgmi(self, paddr:int) -> int: return self.gmc.paddr_base + paddr
@@ -315,6 +316,15 @@ class AMDev:

        ip_offset += 8 + (8 if ihdr.base_addr_64_bit else 4) * ip.num_base_address

+    # HARV(EST) table: harvested instances must be excluded (like amdgpu_discovery_harvest_ip)
+    # layout: u32 signature, u16 version, u16 size, then 32 entries of {hw_id:u16, inst:u8, rsv:u8}
+    self.harvested:dict[int, set[int]] = collections.defaultdict(set)
+    if (harv_off:=self.bhdr.table_list[am.HARVEST_INFO].offset) != 0 and \
+       (blob:=to_mv(ctypes.addressof(self.bhdr) + harv_off, 8 + 32*4).cast('I'))[0] == am.HARVEST_TABLE_SIGNATURE:
+      inv_hw_id = {hw_id: hw_ip for hw_ip, hw_id in am.hw_id_map.items()}
+      for ent in blob[2:]:
+        if (ip_:=inv_hw_id.get(ent & 0xffff)) is not None: self.harvested[ip_].add((ent >> 16) & 0xff)
+
    gc_info = am.struct_gc_info_v1_0.from_address(gc_addr:=ctypes.addressof(self.bhdr) + self.bhdr.table_list[am.GC].offset)
    self.gc_info = getattr(am, f"struct_gc_info_v{gc_info.header.version_major}_{gc_info.header.version_minor}").from_address(gc_addr)
    self.reserved_vram_size = (384 << 20) if self.ip_ver[am.GC_HWIP][:2] in {(9,4), (9,5)} else (64 << 20)
@@ -29,7 +29,9 @@ class AM_SOC(AM_IP):

  def init_hw(self):
    if self.adev.ip_ver[am.NBIO_HWIP] in {(7,9,0), (7,9,1)}:
-      self.adev.regXCC_DOORBELL_FENCE.write(0x0)
+      # fence doorbells for harvested xccs (0xff & ~xcc_mask in the kernel); a fully-unharvested chip keeps the previous 0x0
+      live_xccs = sum(1 << i for i in self.adev.regs_offset[am.GC_HWIP] if i not in self.adev.harvested[am.GC_HWIP] and i < 8)
+      self.adev.regXCC_DOORBELL_FENCE.write(0xff & ~live_xccs)
      for aid in range(1, self.adev.gmc.vmhubs):
        self.adev.indirect_wreg_pcie(self.adev.regXCC_DOORBELL_FENCE.addr[0], self.adev.regXCC_DOORBELL_FENCE.encode(shub_slv_mode=1), aid=aid)
      self.adev.regBIFC_GFX_INT_MONITOR_MASK.write(0x7ff)
@@ -52,7 +54,8 @@ class AM_GMC(AM_IP):
    self.vmhubs = len(self.adev.regs_offset[am.MMHUB_HWIP])

    # XGMI (for supported systems)
-    self.xgmi_phys_id = self.adev.regMMMC_VM_XGMI_LFB_CNTL.read_bitfields()['pf_lfb_region'] if hasattr(self.adev, 'regMMMC_VM_XGMI_LFB_CNTL') else 0
+    xgmi_lfb_cntl = self.adev.regMMMC_VM_XGMI_LFB_CNTL.read_bitfields() if hasattr(self.adev, 'regMMMC_VM_XGMI_LFB_CNTL') else {}
+    self.xgmi_phys_id, self.xgmi_max_region = xgmi_lfb_cntl.get('pf_lfb_region', 0), xgmi_lfb_cntl.get('pf_max_region', 0)
    self.xgmi_seg_sz = self.adev.regMMMC_VM_XGMI_LFB_SIZE.read_bitfields()['pf_lfb_size']<<24 if hasattr(self.adev, 'regMMMC_VM_XGMI_LFB_SIZE') else 0

    self.paddr_base = self.xgmi_phys_id * self.xgmi_seg_sz
@@ -189,13 +192,13 @@ class AM_SMU(AM_IP):
    if DEBUG >= 2: print(f"am {self.adev.devfmt}: mode1 reset")
    if self.adev.ip_ver[am.MP0_HWIP] >= (14,0,0) or self.adev.ip_ver[am.MP0_HWIP] in {(13,0,0), (13,0,7), (13,0,10)}:
      self._send_msg(__DEBUGSMC_MSG_Mode1Reset:=2, 0, debug=True)
-    elif self.adev.ip_ver[am.MP0_HWIP] in {(13,0,6), (13,0,12)}: self._send_msg(self.smu_mod.PPSMC_MSG_GfxDriverReset, 1)
+    elif self.adev.ip_ver[am.MP0_HWIP] in {(13,0,6), (13,0,12), (13,0,15)}: self._send_msg(self.smu_mod.PPSMC_MSG_GfxDriverReset, 1)
    else: self._send_msg(self.smu_mod.PPSMC_MSG_Mode1Reset, 0)

    if not self.adev.is_hive(): time.sleep(0.5) # 500ms

  def read_table(self, table_t, arg):
-    if self.adev.ip_ver[am.MP0_HWIP] in {(13,0,6),(13,0,12)}: self._send_msg(self.smu_mod.PPSMC_MSG_GetMetricsTable, arg)
+    if self.adev.ip_ver[am.MP0_HWIP] in {(13,0,6),(13,0,12),(13,0,15)}: self._send_msg(self.smu_mod.PPSMC_MSG_GetMetricsTable, arg)
    else: self._send_msg(self.smu_mod.PPSMC_MSG_TransferTableSmu2Dram, arg)
    return table_t.from_buffer(bytearray(self.adev.vram.view(self.driver_table_paddr, ctypes.sizeof(table_t))[:]))

@@ -206,7 +209,7 @@ class AM_SMU(AM_IP):

  def set_clocks(self, level:int|None):
    clks = tuple([self.smu_mod.PPCLK_UCLK, self.smu_mod.PPCLK_FCLK, self.smu_mod.PPCLK_SOCCLK])
-    if self.adev.ip_ver[am.MP0_HWIP] not in {(13,0,6), (13,0,12)}: clks += (self.smu_mod.PPCLK_GFXCLK,)
+    if self.adev.ip_ver[am.MP0_HWIP] not in {(13,0,6), (13,0,12), (13,0,15)}: clks += (self.smu_mod.PPCLK_GFXCLK,)

    if level is None:
      for clck in clks:
@@ -246,7 +249,7 @@ class AM_SMU(AM_IP):

 class AM_GFX(AM_IP):
  def init_sw(self):
-    self.xccs = len(self.adev.regs_offset[am.GC_HWIP])
+    self.xccs = sum(1 for i in self.adev.regs_offset[am.GC_HWIP] if i not in self.adev.harvested[am.GC_HWIP])
    self.mqd_paddr = [self.adev.mm.palloc(0x1000 * self.xccs, zero=False, boot=True) for i in range(2)]
    self.mqd_mc = [self.adev.paddr2mc(mqd_paddr) for mqd_paddr in self.mqd_paddr]

@@ -514,7 +517,7 @@ class AM_SDMA(AM_IP):
        **({'utc_l1_enable':1} if self.adev.ip_ver[am.SDMA0_HWIP] <= (5,2,0) else {}), inst=inst)

    if self.adev.ip_ver[am.NBIO_HWIP] in {(7,9,0), (7,9,1)}:
-      for aid_id in range(4):
+      for aid_id in range(self.adev.gmc.vmhubs):
        for dev_inst, (port, awid, offset, awaddr) in enumerate([(1, 0xe, 0xe, 0x1), (2, 0x8, 0x8, 0x2), (5, 0x9, 0x9, 0x8), (6, 0xa, 0xa, 0x9)]):
          entry = dev_inst + 1 + 4 * aid_id
          self.adev.reg(f"regDOORBELL0_CTRL_ENTRY_{entry}").write(**{f"bif_doorbell{entry}_range_size_entry": 20,
@@ -100,20 +100,33 @@ pm_replace_buffers = PatternMatcher([(UPat(Ops.CALL, name="call"), replace_call_
 # *****************
 # 1.1. prep: staging copies

+STAGING_SIZE, STAGING_SLOTS = 128 << 20, 2
+
+@functools.cache
+def _staging() -> Buffer: return Buffer("CPU", STAGING_SIZE, dtypes.uint8, preallocate=True)
+
 def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS - {"CPU"}) and not all_devices_in(b.device, HCQ_DEVS)

+def stage_copy(dst:UOp, src:UOp) -> UOp|None:
+  if not (_need_staging(src, dst) or _need_staging(dst, src)): return None
+
+  assert src.dtype.itemsize == dst.dtype.itemsize, "staged copies must be dtype-size matched"
+  base, it, copies = UOp.from_buffer(_staging()), src.dtype.itemsize, []
+  chunk = (STAGING_SIZE // STAGING_SLOTS) // it
+  for i, off in enumerate(range(0, src.max_numel(), chunk)):
+    stage = base[(so:=(i % STAGING_SLOTS) * chunk * it):so + (n:=min(chunk, src.max_numel() - off)) * it]
+    copies += [src[off:off+n].copy_to_device("CPU").call(stage, src[off:off+n]), stage.copy_to_device(dst.device).call(dst[off:off+n], stage)]
+  return UOp(Ops.LINEAR, src=tuple(copies))
+
+# *****************
+# 1.2. prep: kernel copies
+
 def _get_enqueue_devs(call:UOp) -> Any|None:
  if not (bufs:=call.src[1:]) or not all(all_devices_in(b.device, HCQ_DEVS) for b in bufs): return None
  if call.src[0].op is Ops.COPY: bufs = bufs[::-1] # copies push from the src device: p2p writes are faster than reads
  devs = min(bufs, key=lambda b: to_tuple(b.device)[0].startswith("CPU")).device # prio to enqueue on not CPU device
  return devs if all_devices_in(devs, HCQ_DEVS) else None

-def stage_copy(dst:UOp, src:UOp) -> UOp|None:
-  if not (_need_staging(src, dst) or _need_staging(dst, src)): return None
-
-  stage = UOp.new_buffer("CPU", src.max_numel() * src.dtype.itemsize, dtypes.uint8)
-  return UOp(Ops.LINEAR, src=(src.copy_to_device("CPU").call(stage, src), stage.copy_to_device(dst.device).call(dst, stage)))
-
 def kernel_copy(call:UOp, dst:UOp, src:UOp) -> UOp|None:
  if (devs:=_get_enqueue_devs(call)) is None or Device[(dev:=to_tuple(devs)[0])].has_copy_queue: return None
  d, s = (UOp.param(i, dst.dtype, (n:=dst.max_numel(),), device=devs) for i in range(2))
@@ -336,7 +349,9 @@ def split_patches(call:UOp) -> UOp|None:
  runtimes, systems = partition(internals, lambda g: any(x.tag in {"program", "kernargs", "cmdbuf"} for x in unwrap_mstack(g.buf_uop)))
  tables = [make_addr_table(call, gs, n) for gs,n in ((inputs, "inputs"), (runtimes, "runtime"), (systems, "systems"))]
  reads, fills = {k:v for _,r,_,_ in tables for k,v in r.items()}, [f for t in tables[1:] for f in t[2]] # inputs table is filled by exec
-  gathers = make_gather_loop(ipathces, tables[0][0], tables[0][3], lt_patches) if (ipathces:=[p for p in rt_patches if p.tag == "inputs"]) else {}
+
+  ipatches = [p for p in rt_patches if p.tag == "inputs" and all(v in tables[0][3] for v in p.src[1].src)] # only getaddrs go to the table
+  gathers = make_gather_loop(ipatches, tables[0][0], tables[0][3], lt_patches) if ipatches else {}
  body = body.substitute({p:p.substitute(gathers | reads) for p in rt_patches})

  lt_srcs = collections.defaultdict(list)
@@ -585,14 +600,15 @@ class HCQ2Compiled(Compiled):
      tdiffs.append((st+perf_counter_us())/2 - gpu)
    Compiled.profile_events.append(ProfileDeviceEvent(self.device, statistics.median(tdiffs), self.device_props()))

-  @functools.cached_property
-  def rt_buffer(self) -> Buffer:
-    return Buffer(self.device, self.rt_allocator.size, dtypes.uint8, options=BufferSpec(uncached=True, cpu_access=True), preallocate=True)
+  @functools.cache
+  def rt_buffer(self, uncached:bool=True) -> Buffer:
+    return Buffer(self.device, self.rt_allocator.size, dtypes.uint8, options=BufferSpec(uncached=uncached, cpu_access=True), preallocate=True)

  def new_buffer(self, b:UOp, cache:bool) -> Buffer:
    if cache or b.tag in HCQ_CACHE_TAGS:
      return Buffer(self.device, b.max_numel(), b.dtype, options=BufferSpec(uncached=b.tag not in ("program","kernargs"), cpu_access=True,nolru=True))
-    return self.rt_buffer.view(b.max_numel(), b.dtype, self.rt_allocator.alloc(b.max_numel() * b.dtype.itemsize, alignment=128))
+    return self.rt_buffer(uncached=b.tag!="kernargs").view(b.max_numel(), b.dtype,
+      self.rt_allocator.alloc(b.max_numel() * b.dtype.itemsize, alignment=128))

  @functools.cache
  def signal(self, name:str|int, init_value:int=0) -> Buffer:
@@ -1,6 +1,6 @@
 import ctypes, struct, time, functools, itertools
 from tinygrad.runtime.autogen import libusb
-from tinygrad.helpers import DEBUG, DEV, to_mv, from_mv, round_up, ceildiv
+from tinygrad.helpers import DEBUG, DEV, to_mv, round_up, ceildiv
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support import c

@@ -35,12 +35,6 @@ class USB3:
    self._tags, self._transferred = itertools.count(1), ctypes.c_int(0)
    self._bulk_buf, self._bulk_mv = alloc_cbuffer(4 << 20)
    self._ctrl_buf, self._ctrl_mv = alloc_cbuffer(0x1000)
-    # async bulk OUT state: tag -> (transfer ptr, keepalive memoryview)
-    self._async_seq = itertools.count(1)
-    self._async_pending: dict[int, tuple[c.POINTER[libusb.struct_libusb_transfer], memoryview]] = {}
-    self._async_free: list[c.POINTER[libusb.struct_libusb_transfer]] = []
-    self._async_err = 0
-    self._async_cb = libusb.libusb_transfer_cb_fn(self._on_bulk_done)

    self.handle = c.init_c_var(c.POINTER[libusb.struct_libusb_device_handle], lambda x: checked(libusb.libusb_open)(dev, x))

@@ -79,29 +73,6 @@ class USB3:
      (self.handle, 0x02, self._bulk_buf, len(payload), self._transferred, timeout)
    assert self._transferred.value == len(payload), f"bulk OUT short write: {self._transferred.value}/{len(payload)} bytes"

-  def _on_bulk_done(self, xfer):
-    # NOTE: runs inside libusb event handling; exceptions here are unraisable, so latch errors for bulk_wait.
-    if xfer.contents.status != 0 or xfer.contents.actual_length != xfer.contents.length: self._async_err = xfer.contents.status or -1
-    self._async_pending.pop(int(xfer.contents.user_data or 0), None)
-    self._async_free.append(xfer)  # transfers are reused (re-submit is cheaper than alloc/free)
-
-  def bulk_write_async(self, payload:memoryview, timeout:int=10000) -> int:
-    """Zero-copy async bulk OUT on EP 0x02. Returns a tag for bulk_wait. The payload must stay alive until bulk_wait."""
-    assert payload.contiguous, "bulk_write_async requires a contiguous buffer"
-    tr = self._async_free.pop() if self._async_free else libusb.libusb_alloc_transfer(0)
-    tag = next(self._async_seq)
-    tr.contents.dev_handle, tr.contents.endpoint, tr.contents.type = self.handle, 0x02, libusb.LIBUSB_TRANSFER_TYPE_BULK
-    tr.contents.timeout, tr.contents.length = timeout, len(payload)
-    tr.contents.buffer = ctypes.cast(from_mv(payload, ctypes.c_ubyte), ctypes.POINTER(ctypes.c_ubyte))
-    tr.contents.callback, tr.contents.user_data = self._async_cb, tag
-    self._async_pending[tag] = (tr, payload)
-    checked(libusb.libusb_submit_transfer, "async bulk OUT submit failed")(tr)
-    return tag
-
-  def bulk_wait(self, tag:int):
-    while tag in self._async_pending: checked(libusb.libusb_handle_events)(None)
-    if self._async_err: raise RuntimeError(f"async bulk OUT failed: status={self._async_err}")
-
  def bulk_read(self, length:int, timeout:int=1000) -> memoryview:
    if length > len(self._bulk_mv): self._bulk_buf, self._bulk_mv = alloc_cbuffer(length)
    checked(libusb.libusb_bulk_transfer, "bulk IN 0x81 failed")(self.handle, 0x81, self._bulk_buf, length, self._transferred, timeout)
@@ -189,20 +160,15 @@ class CustomASM24Controller:
    """Write to chip XDATA via vendor control OUT (bRequest=0xE5). wValue=addr, wIndex=val."""
    for off, val in enumerate(data): self.usb.control_write(0xE5, value=base_addr + off, index=val)

-  def scsi_write(self, buf:bytes, slot_start:int=0):
+  def scsi_write(self, buf:bytes):
    """Write to SRAM via 0xF2 vendor command + bulk OUT."""
    buf_padded = buf + b'\x00' * (round_up(len(buf), 512) - len(buf))
-    self.scsi_write_arm(len(buf_padded), slot_start)
+    sectors = len(buf_padded) // 512
+    num_slots = ceildiv(len(buf_padded), 0x4000)  # 16KB per slot
+    windex = (num_slots & 0xFF) << 8
+    self.usb.control_write(0xF2, value=sectors, index=windex)
    self.usb.bulk_write(buf_padded)

-  def scsi_write_arm(self, nbytes:int, slot_start:int=0):
-    assert nbytes % 512 == 0, f"scsi write arm requires 512-byte aligned size, got {nbytes}"
-    sectors = nbytes // 512
-    num_slots = ceildiv(nbytes, 0x4000)  # 16KB per slot
-    self.usb.control_write(0xF2, value=sectors, index=(slot_start & 0xFF) | ((num_slots & 0xFF) << 8))
-
-  def bulk_wait(self, tag:int): self.usb.bulk_wait(tag)
-
  def scsi_read_arm(self, size:int):
    windex = (ceildiv(size, 0x4000) & 0xFF) << 8
    self.usb.control_write(0xF2, value=(ceildiv(size, 512) & 0x7FFF) | 0x8000, index=windex)
@@ -228,9 +194,7 @@ class USBMMIOInterface(MMIOInterface):
  def __setitem__(self, index, data):
    off, _ = self._off_from_index(index)
    data = struct.pack(self.fmt, data) if isinstance(data, int) else bytes(data)
-    if not self.pcimem:
-      # addr tags 0xf000 + slot_start*0x4000 are SRAM bounce buffers, written via 0xF2 bulk DMA
-      self.usb.scsi_write(data, slot_start=(self.addr - 0xf000) >> 14) if self.addr in (0xf000, 0x4f000) else self.usb.write(self.addr + off, data)
+    if not self.pcimem: self.usb.scsi_write(data) if self.addr == 0xf000 else self.usb.write(self.addr + off, data)
    else: self.usb.pcie_mem_write(self.addr+off, data)

  def view(self, offset:int=0, size:int|None=None, fmt=None):
@@ -39,12 +39,12 @@ pm_fold_moved_after = PatternMatcher([
 def _mop_index(r:UOp, idx:UOp):
  idxs = idx.src[1:]
  if len(idxs) == len(r.shape):
-    return r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idxs), dtype=idx.dtype, arg=idx.arg)
+    return r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idxs), arg=idx.arg)
  if r.op is Ops.RESHAPE:
    src_prefix = len(r.src[0].shape) - len(r.shape[len(idxs):])
    if src_prefix >= 0 and r.src[0].shape[src_prefix:] == r.shape[len(idxs):]:
      if src_prefix == 0: return r.src[0] if r.src[0].dtype == idx.dtype else None
-      ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape[:src_prefix], r.shape[:len(idxs)], idxs), dtype=idx.dtype, arg=idx.arg)
+      ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape[:src_prefix], r.shape[:len(idxs)], idxs), arg=idx.arg)
      return ret if ret.shape == idx.shape else None

 pm_mops = PatternMatcher([
@@ -142,7 +142,7 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None:
    case Ops.CMPLT | Ops.CMPNE | Ops.CMPEQ:
      return dtypes.bool
    case Ops.SIN | Ops.LOG2 | Ops.EXP2 | Ops.SQRT | Ops.RECIPROCAL:
-      return least_upper_float(src[0].dtype)
+      return dtypes.bool if src[0].base.is_invalid else least_upper_float(src[0].dtype)
    case Ops.WHERE:
      if src[0].dtype != dtypes.bool: raise RuntimeError(f"where cond must be bool, got {src[0].dtype}")
      return promo_dtype(src[1:])
@@ -159,7 +159,8 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None:
    case Ops.GETADDR:
      return dtypes.uint64
    case Ops.SHL | Ops.SHR:
-      if not all(dtypes.is_int(x.dtype) for x in src): raise RuntimeError(f"shift operands must be int, got {[x.dtype for x in src]}")
+      if not all(dtypes.is_int(x.dtype) or x.base.is_invalid for x in src):
+        raise RuntimeError(f"shift operands must be int, got {[x.dtype for x in src]}")
      return src[0].dtype
    case Ops.BUFFER | Ops.PARAM:
      assert isinstance(arg, ParamArg), "BUFFER/PARAM must have ParamArg"
@@ -191,7 +192,7 @@ class UOpMetaClass(type):
    # TODO: delete this once the dtype field is removed, for now it just re-implements spec.py
    # an INDEX presents its access dtype, which a still-weak source matches up to weakness
    if SPEC == 2 and op is not Ops.CONST and \
-       not any(s.base.is_invalid for s in src) and (expected_dtype:=dtype_from_uop(op, src, arg)) is not None and expected_dtype != dtype and \
+       (expected_dtype:=dtype_from_uop(op, src, arg)) is not None and expected_dtype != dtype and \
       not (op is Ops.INDEX and weak_dtype(expected_dtype) == weak_dtype(dtype)):
      raise RuntimeError(f"bad dtype {dtype}, expected {expected_dtype} on {op}")
    if (wret:=UOpMetaClass.ucache.get(key:=(op, dtype, src, arg, tag), None)) is not None and (ret:=wret()) is not None: return ret
@@ -1747,7 +1748,7 @@ def graph_rewrite(sink:UOp, pm:PatternMatcher, ctx=None, bottom_up=False, name=N

 def _rebuild_dtype(n:UOp, new_src:tuple[UOp,...]) -> DType:
  # TODO: delete this once the dtype field is removed, every rebuild will re-derive
-  if all(a.dtype is b.dtype or b.base.is_invalid for a,b in zip(n.src, new_src)): return n.dtype
+  if all(a.dtype is b.dtype for a,b in zip(n.src, new_src)): return n.dtype
  return dtype_from_uop(n.op, new_src, n.arg) or n.dtype

 def sint_to_uop(x:sint, dtype=dtypes.weakint) -> UOp: return UOp.const(x, dtype)
@@ -68,8 +68,10 @@ spec_shared = PatternMatcher([
  (UPat(GroupOp.Comparison, dtype=dtypes.bool, src=(UPat.var("x"), UPat.var("y"))),
   lambda x,y: matches_dtype(x, y.dtype) or matches_dtype(y, x.dtype) or x.dtype in dtypes.weaks or y.dtype in dtypes.weaks),
  (UPat((Ops.AND, Ops.OR, Ops.XOR, Ops.SHL, Ops.SHR), name="x"), lambda x: False if any(dtypes.is_float(s.dtype) for s in x.src) else None),
-  (UPat((Ops.SHL, Ops.SHR), src=(UPat.var("x"), UPat(dtype=dtypes.uint)), name="a"), lambda a,x: matches_dtype(x, a.dtype) or None),
-  (UPat((Ops.CDIV, Ops.CMOD, Ops.FLOORDIV, Ops.FLOORMOD), name="x"), lambda x: None if dtypes.is_int(x.dtype) else False),
+  (UPat((Ops.SHL, Ops.SHR), src=(UPat.var("x"), UPat.var("c")), name="a"), lambda a,x,c: (matches_dtype(x, a.dtype) or x.dtype is dtypes.weakint)
+   and (matches_dtype(c, a.dtype) or c.dtype in (dtypes.uint, dtypes.weakint) or x.base.is_invalid)),
+  (UPat((Ops.CDIV, Ops.CMOD, Ops.FLOORDIV, Ops.FLOORMOD), name="x"),
+   lambda x: None if dtypes.is_int(x.dtype) or any(s.base.is_invalid for s in x.src) else False),
  (UPat(GroupOp.ALU, name="x"), lambda x: all(matches_dtype(y, x.dtype) or y.dtype in dtypes.weaks for y in x.src)),

  # CAST
@@ -134,7 +136,8 @@ def valid_gettuple(g:UOp, t:UOp): return isinstance(g.arg, int) and 0 <= g.arg <

 # these ops can exist in tensor but not programs. example: movement
 spec_tensor = PatternMatcher([
-  (UPat((Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL), src=(UPat(),), name="u"), lambda u: dtypes.is_float(u.dtype)),
+  (UPat((Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL), src=(UPat(),), name="u"),
+   lambda u: dtypes.is_float(u.dtype) or u.src[0].base.is_invalid),

  # BUFFER
  (UPat(Ops.BUFFER, src=(UPat(),), name="buf"), lambda buf:
@@ -101,6 +101,11 @@ pm_remove_invalid = PatternMatcher([
   if any(x.is_invalid for x in s.src) else None),
 ])

+def fold_const_where(gate:UOp, c0:UOp, c1:UOp, w:UOp) -> UOp:
+  # folding a strong dtype WHERE to a weak const branch keeps the strong dtype
+  ret = c0 if gate.val else c1
+  return commit_weak(ret, w.dtype) if ret.op is Ops.CONST and ret.dtype in dtypes.weaks and w.dtype not in dtypes.weaks else ret
+
 symbolic_simple = pm_data_invalid + PatternMatcher([
  # ** self folding **
  (UPat.var("x") + 0, lambda x: x),    # x+0 -> x
@@ -175,7 +180,7 @@ symbolic_simple = pm_data_invalid + PatternMatcher([
  # ** simple where folding **
  # a conditional with the same results either way is a noop, also fold const conditionals
  (UPat.var().where(UPat.var("val"), UPat.var("val")), lambda val: val),
-  (UPat.cvar("gate").where(UPat.var("c0"), UPat.var("c1")), lambda gate, c0, c1: c0 if gate.val else c1),
+  (UPat.cvar("gate").where(UPat.var("c0"), UPat.var("c1")).named("w"), fold_const_where),
  # a.where(b.where(c, d), d) -> (a & b).where(c, d)
  (UPat.var("a").where(UPat.var("b").where(UPat.var("c"), UPat.var("d")), UPat.var("d")), lambda a,b,c,d: (a&b).where(c,d)),
  # a.where(c, b.where(c, d)) -> (a | b).where(c, d)
Author	SHA1	Message	Date
geohot	369b9d0c27	add ansipad and PARALLEL contextvar	2026-08-20 16:18:36 -07:00
sirhcmandGitHub	c655aaf3a2	ci: venv in /opt/venv (#17632 ) default python is now 3.14, not 3.12	2026-08-20 18:54:54 -04:00
chenyuandGitHub	592e3f8363	update const selector where folding [pr] (#17640 ) folding a strong dtype WHERE to a weak const branch keeps the strong dtype	2026-08-20 18:49:34 -04:00
sirhcmandGitHub	3715006a21	fix float_to_bf16 on non-float32 inputs (#17638 )	2026-08-20 18:14:30 -04:00
wozeparrotandGitHub	d80254c1d9	fa: remove swa recompute (#17635 )	2026-08-20 14:26:44 -07:00
sirhcmandGitHub	c773891e3f	skip INT_MIN % -1 on X86 and LLVM (#17637 )	2026-08-20 17:26:26 -04:00
chenyuandGitHub	0e7ab863a0	x86 REX issue (#17580 ) * failing test * fix	2026-08-20 16:49:04 -04:00
geohot	a57188ea6d	hotfix: disable HCQ2	2026-08-20 13:39:48 -07:00
George HotzandGitHub	1707dca3b4	remove kernel_cnt, names are no longer unique (#17633 ) * remove kernel_cnt, names are no longer unique * check uops, not names * fix hcq2	2026-08-20 13:28:25 -07:00
sirhcmandGitHub	eba5b7e750	benchmarks: cleanups (#17631 )	2026-08-20 16:12:14 -04:00
chenyuandGitHub	a8ecb73363	x64 imm uint64 (#17636 ) * x64 imm uint64 * fix	2026-08-20 15:30:17 -04:00
chenyuandGitHub	f55c1a37d2	clean up some index(dtype=) [PR] (#17634 )	2026-08-20 15:18:15 -04:00
chenyuandGitHub	6732d05157	fix dtype_from_uop for invalid ALU [PR] (#17626 ) invalid is bool, and ALU(invalid) is invalid which is bool	2026-08-20 10:47:23 -04:00
nimlgenandGitHub	e68aa16e3f	hcq2: faster beam (#17624 )	2026-08-20 15:21:53 +03:00
nimlgenandGitHub	8c3cb00d36	hcq2: staging (#17622 )	2026-08-20 13:42:48 +03:00
George HotzandGitHub	c117da9850	safe changes for mi350p (#17620 ) * safe changes for mi350p * bump amd firmware	2026-08-20 00:00:23 -07:00
George HotzandGitHub	57d1104a92	bump amd firmware (#17621 )	2026-08-19 23:51:15 -07:00
qazalandGitHub	a1263fadf3	fused_qkv_rope in UOp try 2 (#17619 ) * fused_qkv_rope in UOp try 2 * dont need that * less	2026-08-20 12:28:34 +09:00
chenyuandGitHub	e6324d1e1c	test updates from weak const branch (#17618 )	2026-08-19 22:56:25 -04:00
sirhcmandGitHub	c63d94e059	benchmarks: split multigpu (#17615 )	2026-08-19 21:36:07 -04:00