27b 935/47

# Conflicts:
#	extra/gemm/amd_flash_attention.py
#	tinygrad/codegen/__init__.py
#	tinygrad/runtime/ops_cpu.py
#	tinygrad/uop/ops.py

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

@@ -137,6 +137,7 @@ runs:
       if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
       shell: bash
       run: |
+        sudo mkdir -p /var/cache/apt/archives
         sudo chown -R $USER:$USER /var/cache/apt/archives
 
         echo 'Acquire::GzipIndexes "true";' | sudo tee /etc/apt/apt.conf.d/gzip
@@ -214,6 +215,7 @@ runs:
           sudo apt-get -y --allow-unauthenticated --no-install-recommends install ${{ steps.apt-pkgs.outputs.pkgs }}
         fi
 
+        sudo mkdir -p /var/cache/apt/archives
         sudo chown -R $USER:$USER /var/cache/apt/archives/
 
     - name: Add clang to PATH (Linux)

.github/workflows/benchmark.yml

@@ -117,10 +117,10 @@ jobs:
       run: python3 test/external/process_replay/reset.py
     - name: Run llama3.2
       run: BENCHMARK_LOG=llama32_3b-f16 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m llama3.2:3b-f16 --benchmark --warmup
-    - name: Run qwen3.5
-      # qwen3.5:35b-a3b doesn't fit on mac
+    - name: Run qwen3.6
+      # qwen3.6:35b-a3b doesn't fit on mac
       if: ${{ matrix.dev != 'METAL' }}
-      run: BENCHMARK_LOG=qwen35_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.5:35b-a3b --benchmark --warmup
+      run: BENCHMARK_LOG=qwen36_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.6:35b-a3b --benchmark --warmup
     - name: Run olmoe
       # just metal for now
       if: ${{ matrix.dev == 'METAL' }}

.github/workflows/test.yml

@@ -167,6 +167,7 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: unittest-13
+        python-version: '3.11'
         deps: testing_unit
         llvm: 'true'
         amd: 'true'
@@ -176,13 +177,14 @@ jobs:
       run: |
         DEV=NULL python3 -m unittest test.backend.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step
         DEV=NULL VIZ=1 python3 -m pytest -n=auto test/null/test_viz.py
+        DEBUG=7 python -m tinygrad.viz.cli --json | jq empty
     # TODO: too slow
     # - name: Run SDXL on NULL backend
     #   run: DEV=NULL DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
     - name: Run Clip tests for SD MLPerf on NULL backend
       run: DEV=NULL python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
     - name: Run AMD emulated BERT training on NULL backend
-      run: DEV=NULL::gfx1201 NULL_ALLOW_COPYOUT=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+      run: DEV=NULL::gfx1201 NULL_ALLOW_COPYOUT=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
     # TODO: support fake weights
     #- name: Run LLaMA 7B on 4 fake devices
     #  run: DEV=NULL python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
@@ -199,6 +201,7 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: unittest-13
+        python-version: '3.11'
         pydeps: "pre-commit"
         deps: testing_unit
         llvm: 'true'
@@ -478,7 +481,7 @@ jobs:
           wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
           echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-21 main" | sudo tee /etc/apt/sources.list.d/llvm.list
           sudo apt-get update
-          sudo apt-get install llvm-21 llvm-21-tools cloc
+          sudo apt-get install -y llvm-21 llvm-21-tools cloc
       - name: Install rocprof-trace-decoder
         run: sudo PYTHONPATH="." ./extra/sqtt/install_rocprof_decoder.py
       - name: Run AMD renderer tests
@@ -628,7 +631,7 @@ jobs:
 
   unittestmacos:
     name: MacOS (unit)
-    runs-on: &macos macos-26
+    runs-on: macos-26
     timeout-minutes: 20
     steps:
     - name: Checkout Code
@@ -677,7 +680,7 @@ jobs:
       matrix:
         group: [1, 2]
     name: MacOS (DEV=METAL) (${{ matrix.group }})
-    runs-on: *macos
+    runs-on: macos-26
     timeout-minutes: 20
     env:
       DEV: METAL
@@ -709,7 +712,7 @@ jobs:
           - 'WEBGPU'
 
     name: MacOS (DEV=${{ matrix.dev }})
-    runs-on: *macos
+    runs-on: macos-26
     timeout-minutes: 20
     steps:
     - name: Checkout Code

AGENTS.md

@@ -3,4 +3,4 @@
 - Run tests with `-n12` for speed (e.g. `python -m pytest test/null/test_dtype.py -x -q -n12`)
 - Run `python -m mypy tinygrad/` to typecheck
 - Run `python -m ruff check .` to lint
-- Read `./tinygrad/viz/README` for profiling
+- Read `./tinygrad/viz/README.md` for profiling and debugging rewrite rules

docs/abstractions4.py

@@ -88,7 +88,7 @@ def example_3_custom_uop(a:Tensor, correct):
 
     # store all the per lane accumulators to LOCAL
     local_accs = UOp.placeholder((LCLS,), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
-    local_accs = local_accs.after(local_accs[lane].store(acc[0]).barrier())
+    local_accs = local_accs.after(local_accs[lane].store(acc[0]))
 
     # accumulate LOCALs into a single per CU accumulator
     late_reduce_loop = UOp.range(LCLS, 3, AxisType.REDUCE)

docs/dtypes.md

@@ -1,7 +1,8 @@
 ::: tinygrad.dtype.DType
 
-::: tinygrad.dtype.dtypes
+::: tinygrad.dtype.DTypes
     options:
+        heading: dtypes
         members: true
         members_order: source
         show_labels: false

examples/beautiful_cifar.py

@@ -9,8 +9,7 @@ from extra.lr_scheduler import OneCycleLR
 GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))]
 
 # override tinygrad defaults
-dtypes.default_float = dtypes.half
-Context(FUSE_OPTIM=1).__enter__()
+Context(DEFAULT_FLOAT=dtypes.half, FUSE_OPTIM=1).__enter__()
 
 # from https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
 batchsize = getenv("BS", 1024)

examples/mlperf/dataloader.py

@@ -1,11 +1,11 @@
 import os, random, pickle, queue, struct, math, functools, hashlib, time
 from typing import List
 from pathlib import Path
-from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu_count
+from multiprocessing import Queue, Process, shared_memory, connection, Lock
 
 import numpy as np
 from tinygrad import dtypes, Tensor
-from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX
+from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX, NUM_CPU_THREADS
 from tinygrad.nn.state import TensorIO
 
 ### ResNet
@@ -131,7 +131,7 @@ def batch_load_resnet(batch_size=64, val=False, shuffle=True, seed=None, pad_fir
     else: X = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name}")
     Y = [None] * (batch_size*BATCH_COUNT)
 
-    for _ in range(cpu_count()):
+    for _ in range(NUM_CPU_THREADS.value):
       p = Process(target=loader_process, args=(q_in, q_out, X, seed))
       p.daemon = True
       p.start()
@@ -212,7 +212,7 @@ def batch_load_train_bert(BS:int, seed:int|None=None):
     rng.shuffle(fs)
     train_files.append(fs.pop(0))
 
-  cycle_length = min(getenv("NUM_CPU_THREADS", min(os.cpu_count(), 8)), len(train_files))
+  cycle_length = min(NUM_CPU_THREADS.value, len(train_files))
   assert cycle_length > 0, "cycle_length must be greater than 0"
 
   dataset = InterleavedDataset(train_files, cycle_length)
@@ -301,7 +301,7 @@ def batch_load_unet3d(preprocessed_dataset_dir:Path, batch_size:int=6, val:bool=
     X = Tensor.empty(*sz, dtype=dtypes.float32, device=f"disk:/dev/shm/{shm_name_x}")
     Y = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name_y}")
 
-    for _ in range(cpu_count()):
+    for _ in range(NUM_CPU_THREADS.value):
       proc = Process(target=load_unet3d_data, args=(preprocessed_dataset_dir, seed, queue_in, queue_out, X, Y))
       proc.daemon = True
       proc.start()
@@ -437,7 +437,7 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
   dataset_iter = iter(image_ids)
 
   try:
-    for _ in range(cpu_count()):
+    for _ in range(NUM_CPU_THREADS.value):
       proc = Process(
         target=load_retinanet_data,
         args=(base_dir, val, queue_in, queue_out, imgs, boxes, labels),

examples/mlperf/model_train.py

@@ -1668,7 +1668,7 @@ def train_llama3():
 def train_gptoss():
   from examples.mlperf.models.gpt_oss import GPTOSS, GPT_OSS_20B, apply_grad, FP8_DTYPE
   from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
-  from examples.mlperf.optim import GradAccClipAdamW, clip_grads
+  from examples.mlperf.optim import GradAccClipAdamW, GradAccClipAdamWGroup, clip_grads
 
   BENCHMARK = getenv("BENCHMARK")
 
@@ -1734,7 +1734,12 @@ def train_gptoss():
   is_offload_optim = bool(getenv("OFFLOAD_OPTIM"))
   is_fake_offload = Device.DEFAULT == "NULL"
   optim_device = ("CPU" if not is_fake_offload else "NULL:99") if is_offload_optim else None
-  optim = GradAccClipAdamW(params, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)
+  params_wd = [p for p in params if p.ndim >= 3]
+  params_no_wd = [p for p in params if p.ndim < 3]
+  optim = GradAccClipAdamWGroup(
+    GradAccClipAdamW(params_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device),
+    GradAccClipAdamW(params_no_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=0.0, grad_acc=grad_acc, device=optim_device),
+  )
 
   for p in optim.params:
     grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype

examples/mlperf/optim.py

@@ -1,6 +1,6 @@
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
-from tinygrad.nn.optim import Optimizer
+from tinygrad.nn.optim import Optimizer, OptimizerGroup
 from tinygrad.helpers import FUSE_OPTIM, getenv
 from tinygrad.uop.ops import UOp, Ops
 
@@ -43,7 +43,7 @@ class GradAccClipAdamW(Optimizer):
 
   def _zero_shard(self, t:Tensor) -> Tensor:
     if not self.zero or (t.shape[0] % len(self.device)) != 0: return t
-    return Tensor(t.uop._shard(0, len(self.device)).multi(0)).clone()
+    return Tensor(t.uop._shard(0, len(self.device)).unshard(0)).clone()
 
   def _zero_gather(self, t:Tensor) -> Tensor:
     if not isinstance(t.device, tuple) or t.uop.axis != 0: return t
@@ -121,3 +121,21 @@ class GradAccClipAdamW(Optimizer):
       return ret.shard_like(t) if offloaded else ret
     out = new_w.cast(t.dtype)
     return out.shard_like(t) if offloaded else out
+
+class GradAccClipAdamWGroup(OptimizerGroup):
+  def fstep(self, grads:list[Tensor], grad_norm:Tensor|None=None):
+    offset = 0
+    to_realize = []
+    for o in self.optimizers:
+      n = len(o.params)
+      to_realize += o.fschedule_step(grads[offset:offset+n])
+      offset += n
+    Tensor.realize(*to_realize, *([grad_norm] if grad_norm is not None else []))
+  @property
+  def lr(self): return self.optimizers[0].lr
+  @property
+  def device(self): return self.optimizers[0].device
+  @property
+  def master_params(self):
+    mp = [mp for o in self.optimizers for mp in (o.master_params or [])]
+    return mp if mp else None

extra/benchmark_llm.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python3
+"""Benchmark tinygrad LLM prefill and decode independently."""
+from __future__ import annotations
+
+import argparse, gc, json, statistics, time
+from dataclasses import asdict, dataclass
+
+from tinygrad import Context, Device, Tensor, UOp
+from tinygrad.helpers import fetch, profile_marker
+from tinygrad.llm.cli import models
+from tinygrad.llm.model import Transformer
+
+@dataclass
+class Result:
+  prompt_tokens: int
+  decode_tokens: int
+  time_to_first_token_s: float
+  prefill_tokens_per_s: float
+  decode_tokens_per_s: float
+  decode_p50_ms: float
+  decode_p95_ms: float
+  output_tokens: list[int]
+
+def percentile(values:list[float], percentile:float) -> float:
+  ordered = sorted(values)
+  return ordered[round((len(ordered) - 1) * percentile)]
+
+def synthetic_prompt(length:int, vocab_size:int, salt:int) -> list[int]:
+  assert length > 0 and vocab_size > 256
+  return [256 + salt % (vocab_size - 256)] + [256 + (i * 7919) % (vocab_size - 256) for i in range(1, length)]
+
+def benchmark(model:Transformer, prompt:list[int], decode_tokens:int, chunk_size:int) -> Result:
+  gen = model.generate(prompt.copy(), chunk_size=chunk_size)
+  profile_marker(f"prefill {len(prompt)} start")
+  begin = time.perf_counter()
+  output_tokens = [next(gen)]
+  ttft = time.perf_counter() - begin
+  profile_marker(f"prefill {len(prompt)} end")
+
+  decode_times: list[float] = []
+  profile_marker(f"decode {len(prompt)} start")
+  for _ in range(decode_tokens):
+    begin = time.perf_counter()
+    output_tokens.append(next(gen))
+    decode_times.append(time.perf_counter() - begin)
+  profile_marker(f"decode {len(prompt)} end")
+
+  return Result(len(prompt), decode_tokens, ttft, len(prompt) / ttft, decode_tokens / sum(decode_times),
+                statistics.median(decode_times) * 1e3, percentile(decode_times, 0.95) * 1e3, output_tokens)
+
+def benchmark_decode_position(model:Transformer, position:int, decode_tokens:int) -> Result:
+  token = Tensor([[0]], dtype="int32", device=Device.DEFAULT).realize()
+  temperature = Tensor([0.0], device=Device.DEFAULT).realize()
+  decode_times, output_tokens = [], []
+  for pos in range(position, position + decode_tokens):
+    begin = time.perf_counter()
+    output_tokens.append(int(model(token, UOp.variable("start_pos", 0, model.max_context-1).bind(pos), temperature).realize().item()))
+    decode_times.append(time.perf_counter() - begin)
+  return Result(position, decode_tokens, 0.0, 0.0, decode_tokens / sum(decode_times),
+                statistics.median(decode_times) * 1e3, percentile(decode_times, 0.95) * 1e3, output_tokens)
+
+def main() -> None:
+  parser = argparse.ArgumentParser(description="Measure LLM prefill and steady-state decode speed")
+  parser.add_argument("--model", default="qwen3:0.6b", help="Model preset or local GGUF path")
+  parser.add_argument("--max-context", type=int, default=32768)
+  parser.add_argument("--prompt-tokens", type=int, nargs="+", default=[128, 2048, 8192])
+  parser.add_argument("--decode-tokens", type=int, default=32)
+  parser.add_argument("--decode-position", type=int, nargs="+")
+  parser.add_argument("--chunk-size", type=int, default=256)
+  parser.add_argument("--beam", type=int, default=2)
+  parser.add_argument("--jit-batch-size", type=int, default=448)
+  parser.add_argument("--parallel-compile", type=int, default=12)
+  parser.add_argument("--realize", action="store_true")
+  parser.add_argument("--json", action="store_true")
+  args = parser.parse_args()
+
+  if args.decode_tokens < 1: parser.error("--decode-tokens must be positive")
+  if args.chunk_size < 1: parser.error("--chunk-size must be positive")
+  if args.decode_position is None and max(args.prompt_tokens) + args.decode_tokens >= args.max_context:
+    parser.error("prompt plus decode tokens must fit within --max-context")
+  if args.decode_position is not None and max(args.decode_position) + args.decode_tokens >= args.max_context:
+    parser.error("decode position plus decode tokens must fit within --max-context")
+
+  begin = time.perf_counter()
+  path = fetch(models.get(args.model, args.model))
+  fetched = time.perf_counter()
+  model, kv = Transformer.from_gguf(path, args.max_context, realize=args.realize)
+  loaded = time.perf_counter()
+  vocab_size = len(kv["tokenizer.ggml.tokens"])
+  print(f"startup: fetch={fetched-begin:.2f}s load={loaded-fetched:.2f}s", flush=True)
+  with Context(BEAM=args.beam, JIT_BATCH_SIZE=args.jit_batch_size, PARALLEL_COMPILE=args.parallel_compile):
+    model.warmup(args.chunk_size)
+  startup = time.perf_counter() - begin
+  print(f"startup: warmup={startup-(loaded-begin):.2f}s total={startup:.2f}s", flush=True)
+  gc.freeze()
+
+  results = [benchmark_decode_position(model, pos, args.decode_tokens) for pos in args.decode_position] if args.decode_position is not None else \
+    [benchmark(model, synthetic_prompt(n, vocab_size, salt=i+1), args.decode_tokens, args.chunk_size)
+     for i, n in enumerate(args.prompt_tokens)]
+  if args.json:
+    print(json.dumps({"model": args.model, "max_context": args.max_context, "chunk_size": args.chunk_size,
+                      "beam": args.beam, "jit_batch_size": args.jit_batch_size, "parallel_compile": args.parallel_compile,
+                      "realize": args.realize, "startup_s": startup, "results": [asdict(x) for x in results]}, indent=2))
+    return
+
+  print(f"model={args.model} max_context={args.max_context} chunk_size={args.chunk_size} beam={args.beam} "
+        f"jit_batch_size={args.jit_batch_size} parallel_compile={args.parallel_compile} realize={args.realize} startup={startup:.2f}s")
+  print(f"{'prompt':>8} {'TTFT':>10} {'prefill':>14} {'decode':>14} {'decode p50':>12} {'decode p95':>12}")
+  for result in results:
+    print(f"{result.prompt_tokens:8d} {result.time_to_first_token_s:9.3f}s {result.prefill_tokens_per_s:11.1f} t/s "
+          f"{result.decode_tokens_per_s:11.1f} t/s {result.decode_p50_ms:9.2f} ms {result.decode_p95_ms:9.2f} ms")
+
+if __name__ == "__main__": main()

extra/export_model.py

@@ -241,8 +241,8 @@ export default {model_name};
 def export_model(model, target:str, *inputs, model_name: Optional[str] = "model", stream_weights=False):
   assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, f"only {', '.join(EXPORT_SUPPORTED_DEVICE)} are supported"
 
-  # NOTE: CPU_COUNT=1, since export does not support threading
-  with Context(JIT=2, CPU_COUNT=1): linear, output_bufs = jit_model(model, *inputs)
+  # NOTE: NUM_CPU_THREADS=1, since export does not support threading
+  with Context(JIT=2, NUM_CPU_THREADS=1): linear, output_bufs = jit_model(model, *inputs)
   functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs)
   state = get_state_dict(model)
   weight_names = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None}

extra/fp8/fp8_linear.py

@@ -18,9 +18,9 @@ def custom_matmul(output: UOp, inp: UOp, weight: UOp) -> UOp:
   SEQ = inp.shape[1]
   OUT = weight.shape[0]
   IN = weight.shape[-1]
-  seq_idx = UOp.range(SEQ, 2, AxisType.LOOP)
-  out_idx = UOp.range(OUT, 3, AxisType.LOOP)
-  batch_idx = UOp.range(output.size//SEQ//OUT, 1, AxisType.LOOP)
+  seq_idx = UOp.range(SEQ, 2)
+  out_idx = UOp.range(OUT, 3)
+  batch_idx = UOp.range(output.size//SEQ//OUT, 1)
   reduce_idx = UOp.range(IN, 0, AxisType.REDUCE)
   product = (inp.index((seq_idx*IN+reduce_idx+batch_idx*IN*SEQ)) * weight.index((out_idx*IN+reduce_idx))).cast(dtypes.float)
   reduced = product.reduce(reduce_idx, arg=Ops.ADD)
@@ -53,7 +53,7 @@ class FP8Linear:
     x_fp8, x_scale = quantize_to_fp8(x)
     GPUS = self.weight.device
     if isinstance(GPUS, tuple) and len(GPUS) > 1:
-      y = Tensor(Tensor.empty((batch//len(GPUS), seq, self.weight.shape[0]), dtype=dtypes.float, device=GPUS).uop.multi(0), device=GPUS)
+      y = Tensor(Tensor.empty((batch//len(GPUS), seq, self.weight.shape[0]), dtype=dtypes.float, device=GPUS).uop.unshard(0), device=GPUS)
     else:
       y = Tensor.empty((batch, seq, self.weight.shape[0]), dtype=dtypes.float)
     y = Tensor.custom_kernel(y, x_fp8, w_fp8, fxn=custom_matmul, grad_fxn=custom_matmul_backward)[0]

extra/gemm/amd_copy_matmul.py

@@ -58,8 +58,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
   B_copy = B_local.permute((1,0)) if use_wmma else B_local
   A_store = A_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(a[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
   B_store = B_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(b[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
-  barrier = UOp.barrier(A_store, B_store)
-  A_local, B_local = A_local.after(barrier), B_local.after(barrier)
+  # NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
+  A_local, B_local = A_local.after(A_store, B_store), B_local.after(A_store, B_store)
 
   # -- COMPUTE --
   lane_m, lane_n = lane // LANES_PER_WAVE_N, lane % LANES_PER_WAVE_N
@@ -70,8 +70,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
 
   if use_wmma:
     k = UOp.range(BLOCK_K // WMMA_K, 101, AxisType.REDUCE)
-    tile_m = UOp.range(TM // WMMA_ACC, 200, AxisType.LOOP)
-    tile_n = UOp.range(TN, 201, AxisType.LOOP)
+    tile_m = UOp.range(TM // WMMA_ACC, 200)
+    tile_n = UOp.range(TN, 201)
 
     acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0,2,1)[tile_m, tile_n]
     a_frag = A_local.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_K // WMMA_K, WMMA_K)[wave_m, tile_m, lane_n, k]
@@ -96,8 +96,8 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
     b_frag = b_frag.reshape(1, TN).expand(TM, TN)
     acc_store = acc.store(acc.after(k) + (a_frag * b_frag))
 
-  # store accumulator and loop
-  acc = acc.after(acc_store.end(k).barrier().end(k_tile))
+  # store accumulator and loop (the barrier at the end of the loop is implied by the LOCAL buffers stored and loaded in the loop)
+  acc = acc.after(acc_store.end(k).end(k_tile))
 
   # store accumulator to output (unified)
   c = c.reshape(WAVES_M, TM//UNROLL_M, LANES_PER_WAVE_M, UNROLL_M,

extra/gemm/amd_flash_attention.py

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

extra/gemm/amd_uop_matmul.py

@@ -28,10 +28,10 @@ REG_TILES_PER_WAVE_M = BLOCK_M // (WAVES_PER_BLOCK_M * LANES_PER_WAVE_M * TM)
 assert WAVES_PER_BLOCK_M*REG_TILES_PER_WAVE_M*LANES_PER_WAVE_M*TM == BLOCK_M, "M reshape is wrong"
 assert WAVES_PER_BLOCK_N*REG_TILES_PER_WAVE_N*LANES_PER_WAVE_N*TN == BLOCK_N, "N reshape is wrong"
 
-def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.LOOP): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
+def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.WEAK): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
 def copy(dest:UOp, src:UOp, rng:int, upcast=False):
   assert dest.shape == src.shape
-  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.LOOP)
+  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.WEAK)
   return dest[*rngs].store(src[*rngs]).end(*rngs)
 
 def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
@@ -66,9 +66,8 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
   B_local = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
   B_local_store = copy(B_local.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)
 
-  # TODO: can we automate barrier?
-  barrier = UOp.barrier(A_local_store, B_local_store)
-  A_local, B_local = A_local.after(barrier), B_local.after(barrier)
+  # NOTE: no explicit barrier needed, the AFTER on the LOCAL buffers implies it in late codegen
+  A_local, B_local = A_local.after(A_local_store, B_local_store), B_local.after(A_local_store, B_local_store)
 
   # open inner k range
   k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)
@@ -102,7 +101,7 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
   sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iter_m, t_m] * B_row[iter_n, t_n]).end(iter_m, iter_n, t_m, t_n)
 
   # Close k, sync, and close K tiles
-  sink = sink.end(k).barrier().end(k_tile_range)
+  sink = sink.end(k).end(k_tile_range)
 
   # ---------------------------
   # REG -> GLOBAL (epilogue)

extra/gemm/cdna_asm_gemm.py

@@ -72,7 +72,7 @@ def hk_fp8_atb_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, g_amax:Tensor|
     elif b.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M, N // ndev, dtype=dtypes.bfloat16, device=a.device), 2
     elif a.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M // ndev, N, dtype=dtypes.bfloat16, device=a.device), 1
     else: inv, out_axis, reduce_out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a.device), 0, True
-    out = Tensor(inv.uop.multi(out_axis), device=a.device)
+    out = Tensor(inv.uop.unshard(out_axis), device=a.device)
     dname = a.device[0]
   else:
     out = Tensor.invalids(1, M, N, dtype=dtypes.bfloat16, device=a.device)
@@ -171,8 +171,8 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
   M, K = A.shape[0]*A.shape[1], A.shape[2]
   K2, N = B.shape[(1 if B.ndim == 3 else 0):]
   assert K == K2
-  m = UOp.range(M, 1, AxisType.LOOP)
-  n = UOp.range(N, 2, AxisType.LOOP)
+  m = UOp.range(M, 1)
+  n = UOp.range(N, 2)
   k = UOp.range(K, 0, AxisType.REDUCE)
   mul = (A.flatten().index((m*UOp.const(dtypes.weakint, K)+k))*
          B.flatten().index((k*UOp.const(dtypes.weakint, N)+n))).cast(dtypes.float32)
@@ -234,7 +234,7 @@ def hk_bf16_atb_gemm(a:Tensor, b:Tensor) -> Tensor:
     elif b.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M, N // ndev, dtype=a.dtype, device=a.device), 2
     elif a.uop.axis == 2: inv, out_axis = Tensor.invalids(1, M // ndev, N, dtype=a.dtype, device=a.device), 1
     else: inv, out_axis, reduce_out = Tensor.invalids(1, M, N, dtype=a.dtype, device=a.device), 0, True
-    out = Tensor(inv.uop.multi(out_axis), device=a.device)
+    out = Tensor(inv.uop.unshard(out_axis), device=a.device)
     dname = a.device[0]
   else:
     out = Tensor.invalids(1, M, N, dtype=a.dtype, device=a.device)
@@ -366,11 +366,11 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N
 
   if is_multi:
     if n_sharded:
-      out = Tensor(Tensor.invalids(batch, M, N//len(a.device), dtype=out_dtype, device=a.device).uop.multi(2), device=a.device)
+      out = Tensor(Tensor.invalids(batch, M, N//len(a.device), dtype=out_dtype, device=a.device).uop.unshard(2), device=a.device)
     elif m_sharded:
-      out = Tensor(Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device).uop.multi(1), device=a.device)
+      out = Tensor(Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device).uop.unshard(1), device=a.device)
     else:
-      out = Tensor(Tensor.invalids(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=out_dtype, device=a.device).uop.multi(0),
+      out = Tensor(Tensor.invalids(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=out_dtype, device=a.device).uop.unshard(0),
                    device=a.device)
   else:
     out = Tensor.invalids(batch, M, N, dtype=out_dtype, device=a.device)

extra/gemm/mi350x_uop_matmul_2.py

@@ -29,7 +29,7 @@ TID_SIZE = WARPGROUP_SIZE*WARP_SIZE
 
 def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=()):
   assert dest.shape == src.shape
-  rngs = [UOp.range(s, rng+i, AxisType.UPCAST if i in upcast else AxisType.LOOP) for i,s in enumerate(src.shape)]
+  rngs = [UOp.range(s, rng+i, AxisType.UPCAST if i in upcast else AxisType.WEAK) for i,s in enumerate(src.shape)]
   copy = dest[*rngs].store(src[*rngs]).end(*rngs)
   return dest.after(copy) if set else copy
 

extra/gemm/moe_gemm.py

@@ -55,7 +55,7 @@ def grouped_mx_wgrad(g:Tensor, xg:Tensor, expert_off:Tensor, n_experts:int) -> T
   dname = (g.device[0] if isinstance(g.device, tuple) else g.device).split(":")[0]
   is_multi = isinstance(g.device, tuple)
   inv = Tensor.invalids(1, n_experts * N, K, dtype=dtypes.bfloat16, device=g.device)
-  out = Tensor(inv.uop.multi(0), device=g.device) if is_multi else inv
+  out = Tensor(inv.uop.unshard(0), device=g.device) if is_multi else inv
   out = Tensor.custom_kernel(out, gT, xT, g_si, x_si, expert_off,
                              fxn=functools.partial(custom_hk_grouped_mxfp8_wgrad, dname=dname, n_experts=n_experts))[0]
   out = out.sum(0) if is_multi else out.squeeze(0)
@@ -103,7 +103,7 @@ def grouped_mx_gemm(x:Tensor, w:Tensor|tuple[Tensor, Tensor], expert_off:Tensor)
   if isinstance(x.device, tuple) and (row_axis := x.uop.axis) is not None:
     ndev = len(x.device)
     out = Tensor(Tensor.invalids(*(s // ndev if i == row_axis else s for i, s in enumerate(out_shape)),
-                                 dtype=dtypes.bfloat16, device=x.device).uop.multi(row_axis), device=x.device)
+                                 dtype=dtypes.bfloat16, device=x.device).uop.unshard(row_axis), device=x.device)
   else:
     out = Tensor.invalids(*out_shape, dtype=dtypes.bfloat16, device=x.device)
   return Tensor.custom_kernel(out, x_q, w_q, x_si, w_si, xe_in, w_e8, expert_off,

extra/gemm_fragment.py

@@ -0,0 +1,195 @@
+"""
+tilelang-style matmul_relu written with tinygrad UOp APIs.
+
+Demonstrates that tilelang's T.alloc_fragment is expressible with existing
+tinygrad primitives: a per-thread REG buffer, wrapped in one Ops.UNSHARD per
+sharded axis over the LOCAL thread-grid ranges to form the full logical tile.
+Here the 64 threads are an 8x8 grid and each thread owns an 8x8 sub-tile --
+the 2-D fragment layout tilelang infers. The kernel is written against the
+full-tile UNSHARD view, and multi_pm (the same pass that lowers multi-device
+UNSHARDs) resolves it into per-thread shard code.
+
+Reference tilelang kernel:
+
+    @tilelang.jit
+    def matmul_relu(A, B, block_M=64, block_N=64, block_K=64,
+                    dtype=T.float16, accum_dtype=T.float32):
+        M, N, K = T.const('M, N, K')
+        C = T.empty([M, N], dtype)
+        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
+            A_shared = T.alloc_shared((block_M, block_K), dtype)
+            B_shared = T.alloc_shared((block_K, block_N), dtype)
+            C_local  = T.alloc_fragment((block_M, block_N), accum_dtype)
+            T.clear(C_local)
+            for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
+                T.copy(A[by * block_M, ko * block_K], A_shared)
+                T.copy(B[ko * block_K, bx * block_N], B_shared)
+                T.gemm(A_shared, B_shared, C_local)
+            for i, j in T.Parallel(block_M, block_N):
+                C_local[i, j] = T.max(C_local[i, j], 0)
+            T.copy(C_local, C[by * block_M, bx * block_N])
+        return C
+
+API mapping (tilelang -> tinygrad UOps, idioms from test/backend/test_custom_kernel.py):
+
+    T.Kernel(gx, gy, threads=T)    -> AxisType.GLOBAL ranges (blocks) + AxisType.LOCAL ranges (thread grid)
+    T.alloc_shared(shape, dtype)   -> UOp.placeholder(shape, dtype, slot, AddrSpace.LOCAL)
+    T.alloc_fragment(shape, dt)    -> per-thread REG placeholder, wrapped in one Ops.UNSHARD per sharded axis over
+                                      the AxisType.LOCAL ranges: fragment.unshard((axis_y, axis_x), (ty, tx)).
+                                      The full logical tile is the shard with each sharded axis multiplied by its
+                                      range size, exactly like device sharding, but the sharding axes are thread
+                                      axes carried by the RANGE metadata instead of a device tuple. C_local[i, j]
+                                      with [i, j] in this thread's shard is INDEX on the UNSHARD, which multi_pm
+                                      resolves into INDEX on the per-thread REG shard, axis by axis.
+    T.copy(gmem_slice, smem)       -> smem[thread_idx].set(gmem_slice[thread_idx], end=copy_rng). set returns the
+                                      smem tile AFTER the copy; the implicit-barrier pass turns the store->load
+                                      dependency of the loop that consumes it into a workgroup barrier
+    T.gemm (no WMMA)               -> C_local[..].set(C_local.after(k)[..] + a_shared[..] * b_shared[..], end=k)
+                                      with k a loop-carried LOOP range (codegen builds the register accumulator
+                                      from this self-referential store automatically)
+    T.copy(fragment, gmem)         -> gmem.index(gidx).store(C_local[..]).end(all_ranges)
+    UNSHARD lowering               -> multi_pm in codegen (full_rewrite_to_sink): INDEX/AFTER/STORE ops on the
+                                      full-tile view become per-thread shard ops, no UNSHARD survives into the program.
+"""
+
+from tinygrad.dtype import dtypes, AddrSpace, DType
+from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
+from tinygrad.helpers import cdiv, getenv
+from tinygrad.tensor import Tensor
+
+# ---------------------------------------------------------------------------
+# tilelang builtins, expressed with tinygrad UOp APIs
+# ---------------------------------------------------------------------------
+
+def alloc_shared(shape:tuple[int, ...], dtype:DType) -> UOp:
+  """T.alloc_shared: one LOCAL buffer shared by all threads in the block."""
+  return UOp.placeholder(tuple(shape), dtype, next(UOp.unique_num), AddrSpace.LOCAL)
+
+def alloc_fragment(shape:tuple[int, ...], dtype:DType, axes:tuple[int, ...], rngs:tuple[UOp, ...]) -> UOp:
+  """T.alloc_fragment: per-thread REG fragment + UNSHARD over the LOCAL thread grid.
+
+  Each thread privately owns shape[axis]//threads elements along every sharded
+  axis in a REG buffer. The UNSHARDs over the LOCAL thread ranges present the
+  full logical tile: full_shape = shard_shape with each sharded axis multiplied
+  by its range size. This is exactly how UNSHARD carries a DEVICE axis today,
+  except the sharding axes are thread axes carried by the RANGE metadata.
+  """
+  assert len(axes) == len(rngs)
+  assert all(tnum.op is Ops.RANGE and tnum.arg[-1] is AxisType.LOCAL for tnum in rngs), "fragments shard over LOCAL ranges"
+  assert all(shape[a] % (int(rng.vmax)+1) == 0 for a, rng in zip(axes, rngs))
+  by_axis = dict(zip(axes, rngs))
+  shard_shape = tuple(s // (int(by_axis[i].vmax)+1) if i in by_axis else s for i, s in enumerate(shape))
+  fragment = UOp.placeholder(shard_shape, dtype, next(UOp.unique_num), AddrSpace.REG)
+  return fragment.unshard(axes, rngs)
+
+# ---------------------------------------------------------------------------
+# GEMM kernel: C = relu(A @ B), float inputs (fp16 or fp32), fp32 fragment accumulator, no WMMA
+# ---------------------------------------------------------------------------
+
+# 64x64 output tile per block, 128 threads as an 8x16 grid; each thread owns an 8x4 fragment sub-tile
+# (the 2-D per-thread layout tilelang infers for this GEMM). The 4 contiguous columns (TN=4) are what
+# let codegen vectorize loads/stores to float4, matching tilelang's lowering exactly.
+BLOCK_M = BLOCK_N = BLOCK_K = 64
+TY = 8
+TX = 16
+THREADS = TY * TX
+TM = BLOCK_M // TY   # fragment rows per thread  (8)
+TN = BLOCK_N // TX   # fragment columns per thread (4)
+
+def matmul_relu_kernel(c:UOp, a:UOp, b:UOp) -> UOp:
+  """C[M, N] = relu(A[M, K] @ B[K, N]) -- one 64x64 tile per block, locals + a 2-D fragment."""
+  M, K = a.shape
+  K2, N = b.shape
+  assert K == K2 and a.dtype == b.dtype == c.dtype and not dtypes.is_int(a.dtype)
+  assert not (K % BLOCK_K or M % BLOCK_M or N % BLOCK_N), "test sizes must be multiples of the block sizes"
+
+  # with T.Kernel(T.ceildiv(N, BLOCK_N), T.ceildiv(M, BLOCK_M), threads=128) as (bx, by):
+  bx = UOp.range(cdiv(N, BLOCK_N), 0, AxisType.GLOBAL)
+  by = UOp.range(cdiv(M, BLOCK_M), 1, AxisType.GLOBAL)
+  # tx (N, 16) is the fast/inner LOCAL axis so a warp covers 16 cols x 2 rows --
+  # matching tilelang's (tidx>>4, tidx&15) warp composition. This keeps the 8 A_shared
+  # reads in a warp on only 2 row-groups (broadcast across 16 cols) instead of 8 rows
+  # (8-way bank conflict), since A_shared[row*512 + ...] all map to the same bank when 8
+  # distinct rows land in one warp.
+  tx = UOp.range(TX, 2, AxisType.LOCAL)
+  ty = UOp.range(TY, 3, AxisType.LOCAL)
+
+  # A_shared = T.alloc_shared((BLOCK_M, BLOCK_K), dtype)
+  # B_shared = T.alloc_shared((BLOCK_K, BLOCK_N), dtype)
+  A_shared = alloc_shared((BLOCK_M, BLOCK_K), a.dtype)
+  B_shared = alloc_shared((BLOCK_K, BLOCK_N), b.dtype)
+
+  # C_local = T.alloc_fragment((BLOCK_M, BLOCK_N), accum_dtype) -- an 8x4 REG tile per thread of the 8x16 grid
+  C_local = alloc_fragment((BLOCK_M, BLOCK_N), dtypes.float32, (0, 1), (ty, tx))
+
+  # T.clear(C_local) -- each thread zeroes its own fragment sub-tile
+  ic, jc = UOp.range(TM, 4, AxisType.LOOP), UOp.range(TN, 5, AxisType.UPCAST)
+  C_loc = C_local[ic*TM + ty, tx*TN + jc].set(0.0, end=(ic, jc))
+
+  # for ko in T.Pipelined(T.ceildiv(K, BLOCK_K), num_stages=3):
+  # (num_stages pipelining is async copy + multi-buffering; this is the synchronous single-buffer version)
+  ko = UOp.range(cdiv(K, BLOCK_K), 6, AxisType.LOOP)
+
+  # T.copy(A[by * BLOCK_M, ko * BLOCK_K], A_shared) -- each thread copies its own 8x4 sub-tile.
+  # Row index is iar*TM + ty (strided by TM across ty), matching tilelang's layout: thread ty owns
+  # rows {ty, ty+8, ..., ty+56} not {ty*8, ..., ty*8+7}.
+  iar, ka = UOp.range(TM, 7, AxisType.LOOP), UOp.range(TN, 8, AxisType.UPCAST)
+  A_store = A_shared[iar*TM + ty, tx*TN + ka].store(a[by*BLOCK_M + iar*TM + ty, ko*BLOCK_K + tx*TN + ka]).end(iar, ka)
+
+  # T.copy(B[ko * BLOCK_K, bx * BLOCK_N], B_shared)
+  kb, ibr = UOp.range(TM, 9, AxisType.LOOP), UOp.range(TN, 10, AxisType.UPCAST)
+  B_store = B_shared[kb*TM + ty, tx*TN + ibr].store(b[ko*BLOCK_K + kb*TM + ty, bx*BLOCK_N + tx*TN + ibr]).end(kb, ibr)
+
+  # get the shared after the stores (single barrier)
+  A_shared = A_shared.after(A_store, B_store)
+  B_shared = B_shared.after(A_store, B_store)
+
+  # T.gemm(A_shared, B_shared, C_local), no WMMA -- per-thread accumulate over its fragment sub-tile.
+  # identical to custom_gemm: a self-referential store over the loop-carried kk range,
+  # which codegen turns into a register accumulator
+  # kk is the outer compute loop (axis 11) so that for each kk we read all 8 A rows and reuse
+  # the B[kk] read across them -- matching tilelang's ko > kk > row > col access order exactly.
+  kk, ir = UOp.range(BLOCK_K, 11, AxisType.LOOP), UOp.range(TM, 12, AxisType.LOOP)
+  jj = UOp.range(TN, 13, AxisType.UPCAST)
+  acc = C_loc.after(kk)[ir*TM + ty, tx*TN + jj] + A_shared[ir*TM + ty, kk].cast(dtypes.float32) * B_shared[kk, tx*TN + jj].cast(dtypes.float32)
+  # closing the ko loop here too; codegen adds the barrier so no thread overwrites the tiles while others still read them
+  C_loc = C_loc[ir*TM + ty, tx*TN + jj].set(acc, end=(kk, ir, jj, ko))
+
+  # for i, j in T.Parallel(BLOCK_M, BLOCK_N): C_local[i, j] = T.max(C_local[i, j], 0)
+  # T.copy(C_local, C[by * BLOCK_M, bx * BLOCK_N]) -- per-thread store of the fragment shard (relu fused into it)
+  # LOOP: these loops are the per-thread output layout; convert_loop_to_global must not globalize them
+  ie, je = UOp.range(TM, 14, AxisType.LOOP), UOp.range(TN, 15, AxisType.UPCAST)
+  c_st = c[by*BLOCK_M + ie*TM + ty, bx*BLOCK_N + tx*TN + je].store(C_loc[ie*TM + ty, tx*TN + je].relu().cast(c.dtype))
+
+  # all open ranges are closed at the final store (ko was closed above).
+  # the fragment UNSHARDs go to codegen as is: multi_pm there resolves the full-tile view into per-thread shard code
+  return c_st.end(je, ie, tx, ty, bx, by).sink(arg=KernelInfo(name="matmul_relu", opts_to_apply=()))
+
+# ---------------------------------------------------------------------------
+# python wrapper: same signature as the tilelang function
+# ---------------------------------------------------------------------------
+
+def matmul_relu(a:Tensor, b:Tensor) -> Tensor:
+  """C = relu(A @ B), fp16 in/out with an fp32 fragment accumulator."""
+  c = Tensor.empty(a.shape[0], b.shape[1], dtype=a.dtype, device=a.device)
+  return c.custom_kernel(a, b, fxn=matmul_relu_kernel)[0]
+
+# ---------------------------------------------------------------------------
+# test
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+  from tinygrad import Device
+  assert Device[Device.DEFAULT].renderer.has_local, "this GPU-style kernel needs a backend with local memory (LOCAL ranges + barriers)"
+  M = K = N = getenv("N", 256)  # 4x4 grid of 64x64 tiles, 4 K chunks
+  dtype_in = dtypes.half if getenv("HALF") else dtypes.float
+
+  a = Tensor.randn(M, K, dtype=dtype_in).contiguous()
+  b = Tensor.randn(K, N, dtype=dtype_in).contiguous()
+  ref = (a @ b).relu().realize()
+
+  out = matmul_relu(a, b).realize()
+
+  import numpy as np
+  np.testing.assert_allclose(out.numpy(), ref.numpy(), atol=1e-1, rtol=1e-2)
+  print("matmul_relu passed!")

extra/hcq2/ops_amd2.py

@@ -3,8 +3,8 @@ from typing import cast, Any, Callable
 import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools, collections, atexit
 assert sys.platform != 'win32'
 from dataclasses import dataclass
-from extra.hcq2.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, encode_kernargs_clike, make_cmdbuf
-from extra.hcq2.hcq2 import make_binary_patch
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, encode_kernargs_clike, make_cmdbuf
+from tinygrad.runtime.support.hcq2 import make_binary_patch, make_patches
 from tinygrad.uop.ops import sint, UOp
 from tinygrad.device import Compiled, BufferSpec, Buffer, Device
 from tinygrad.dtype import dtypes
@@ -146,31 +146,37 @@ pm_pm4_opsel = PatternMatcher([
   (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
 ])
 
-def pm4_submit(cmdbuf, devs):
-  size, zero = UOp.const(dtypes.uint32, cmdbuf.nbytes() // dtypes.uint32.itemsize), UOp.const(dtypes.int, 0)
-
-  # the compute queue's ring and its host-side ring/write/put pointers (placeholders, resolved in pm_bufferize)
-  for d in devs: q = Device[d].compute_queue
+def pm4_submit(ctx, lin):
+  # ensure compute queues are allocated
+  for d in (devs:=ctx.devs): q = Device[d].compute_queue
   ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COMPUTE:0_{name}")
     for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
 
-  # place the cmdbuf at the ring's write offset, wrapping the ring
-  put = put_ptr.index(zero)
-  next_put = put + size.cast(put.dtype)
-  i = UOp.range(size, 0, dtype=dtypes.int, src=(cmdbuf,))
-  ring_idx = ((put + i.cast(put.dtype)) % q.ring.size).cast(dtypes.int)
+  # two tail dwords coordinate safe IB reuse: GPU completions and host submits
+  size_dw = sum(len(ins.src) for ins in lin.src) + len(release_mem(ctx, 0, 0).src)
+  assert size_dw < (1 << 20), f"indirect buffer of {size_dw} dwords doesn't fit one packet"
 
-  # copy the cmdbuf into the ring and advance the put/write pointers
-  copy_to_ring = ring.index(ring_idx).store(cmdbuf.index(i).load()).end(i)
-  bump_put_ptr = put_ptr.index(zero).store(next_put)
-  bump_wptr = wptr.index(zero).store(next_put)
+  ib = UOp.placeholder((size_dw + 2,), dtypes.uint32, next(UOp.unique_num), device=devs, volatile=True).rtag("cmdbuf")
+  done_idx, submit_idx = UOp.const(dtypes.int, size_dw + 0), UOp.const(dtypes.int, size_dw + 1)
+  submitted = (counter:=ib.after(make_patches(ib, [((size_dw + i) * 4, UOp.const(dtypes.uint32, 0)) for i in range(2)])).index(submit_idx)).load()
+  completed = ib.after(loop:=UOp.loop(0)).index(done_idx).load()
+  ib_free = completed.end(loop, completed != submitted)
 
-  # ring the doorbell once the copy and pointer bumps have landed
-  flush = UOp.barrier(copy_to_ring, bump_put_ptr, bump_wptr)
-  return doorbell.after(flush).index(zero).store(next_put)
+  bump_fence = pm4_store(ctx, UOp(Ops.SLICE, dtypes.uint32, (ib, UOp.const(dtypes.weakint, size_dw)), 2), (submitted + 1).cast(dtypes.uint64))
+  cmdbuf = make_cmdbuf(lin.replace(src=lin.src + (bump_fence,)), devs, buf=ib, dep=ib_free)
 
-pm_pm4_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"),
-  lambda lin: pm4_submit(make_cmdbuf(lin, to_tuple(lin.arg[0])), to_tuple(lin.arg[0])))])
+  # the ring itself only carries a packet pointing at the ib, wrapping the ring
+  put = put_ptr.index(zero:=UOp.const(dtypes.int, 0))
+  pkt = (ctx.pm4.PACKET3(ctx.pm4.PACKET3_INDIRECT_BUFFER, 2), *data64_le(cmdbuf.getaddr(devs)), size_dw | ctx.pm4.INDIRECT_BUFFER_VALID)
+  write_pkt = UOp.barrier(*[ring.index(((put + off) % q.ring.size).cast(dtypes.int)).store(UOp.const(dtypes.uint32, x)) for off,x in enumerate(pkt)])
+
+  # advance the put/write pointers past the packet
+  bump_put_ptr = put_ptr.index(zero).store(put + len(pkt))
+  bump_wptr = wptr.index(zero).store(put + len(pkt))
+  flush = UOp.barrier(write_pkt, bump_put_ptr, bump_wptr, counter.store(submitted + 1))
+  return doorbell.after(flush).index(zero).store(put + len(pkt))
+
+pm_pm4_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"), pm4_submit)])
 
 # *****************
 # SDMA
@@ -242,7 +248,7 @@ def sdma_submit(cmdbuf, devs):
   return doorbell.after(flush).index(zero).store(next_put_b)
 
 pm_sdma_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"),
-  lambda lin: sdma_submit(make_cmdbuf(lin, to_tuple(lin.arg[0])), to_tuple(lin.arg[0])))])
+  lambda ctx, lin: sdma_submit(make_cmdbuf(lin, ctx.devs), ctx.devs))])
 
 @dataclass(frozen=True)
 class AMDEncodeCtx:  # encode-time constants for one queue: devs (every cmdbuf address resolves into these) + gfx version + packet/ip modules
@@ -253,7 +259,7 @@ def encode_queue(q:UOp) -> UOp|None:
   d = Device[(devs:=to_tuple(q.arg[0]))[0]]
   ctx = AMDEncodeCtx(devs, d.target, d.pm4, d.sdma, d.soc, d.gc, d.nbio, d.xccs, d.max_copy_size, d.tmpring_size)
   opsel, submit = (pm_pm4_opsel, pm_pm4_submit) if q.arg[1].startswith("COMPUTE") else (pm_sdma_opsel, pm_sdma_submit)
-  return submit.rewrite(graph_rewrite(q, opsel + pm_flatten_linear, walk=True, ctx=ctx, name=f"{q.arg[1]} opsel"))
+  return submit.rewrite(graph_rewrite(q, opsel + pm_flatten_linear, walk=True, ctx=ctx, name=f"{q.arg[1]} opsel"), ctx)
 
 @dataclass(frozen=True)
 class AMDProgramData:
@@ -511,7 +517,8 @@ class PCIIface(PCIIfaceBase):
       cq = d.compute_queue
       for b in (cq.put_value, cq.read_ptr, cq.write_ptr): b._buf.view.view(fmt='Q')[0] = 0
       d.iface.dev_impl.gfx.setup_ring(*cq.params)
-      d.timeline_signal()._buf.cpu_view().mv.cast('Q')[0] = d.timeline_value().as_memoryview(force_zero_copy=True).cast('Q')[0] - 1
+      d.timeline_signal('COMPUTE:0')._buf.cpu_view().mv.cast('Q')[0] = \
+        d.timeline_value('COMPUTE:0').as_memoryview(force_zero_copy=True).cast('Q')[0] - 1
 
   def sleep(self, timeout):
     if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):

extra/llama_kernels/__init__.py

@@ -31,12 +31,12 @@ def dname_of(device) -> str:
 
 def alloc_like(shape, dtype, device, axis=None) -> Tensor:
   if isinstance(device, tuple) and axis is not None:
-    return Tensor(Tensor.invalids(*shard_shape(shape, axis, len(device)), dtype=dtype, device=device).uop.multi(axis), device=device)
+    return Tensor(Tensor.invalids(*shard_shape(shape, axis, len(device)), dtype=dtype, device=device).uop.unshard(axis), device=device)
   return Tensor.invalids(*shape, dtype=dtype, device=device)
 
 def alloc_local(shape, dtype, device, axis=None) -> Tensor:
   if isinstance(device, tuple) and axis is not None:
-    return Tensor(Tensor.invalids(*shape, dtype=dtype, device=device).uop.multi(0), device=device)
+    return Tensor(Tensor.invalids(*shape, dtype=dtype, device=device).uop.unshard(0), device=device)
   return Tensor.invalids(*shape, dtype=dtype, device=device)
 
 def compile_hip(src:str, defines:list[str]):

extra/llama_kernels/fused_ce/__init__.py

@@ -48,7 +48,7 @@ def _fused_ce_loss_bwd(gradient:UOp, kernel:UOp, label_smoothing:float):
     axis = logits_u.axis
     ndev = len(device)
     local_shape = tuple(s//ndev if i == axis else s for i,s in enumerate((MBS, SEQ, VOCAB)))
-    d_logits = Tensor(Tensor.invalids(*local_shape, dtype=dtypes.bfloat16, device=device).uop.multi(axis), device=device)
+    d_logits = Tensor(Tensor.invalids(*local_shape, dtype=dtypes.bfloat16, device=device).uop.unshard(axis), device=device)
     rows_per_dev = local_shape[0] * local_shape[1]
     seq_per_dev = local_shape[1]
   else:
@@ -74,11 +74,11 @@ def fused_ce_loss(logits:Tensor, targets:Tensor, label_smoothing:float=0.1) -> T
     axis = logits.uop.axis
     assert axis in (0, 1), f"unsupported sharding axis={axis} for CE loss"
     ndev = len(logits.device)
-    loss_out = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0),
+    loss_out = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.unshard(0),
                       device=logits.device)
-    max_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0),
+    max_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.unshard(0),
                       device=logits.device)
-    lse_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.multi(0),
+    lse_out  = Tensor(Tensor.invalids(rows // ndev, dtype=dtypes.float32, device=logits.device).uop.unshard(0),
                       device=logits.device)
     local_shape = tuple(s//ndev if i == axis else s for i,s in enumerate(logits.shape))
     rows_per_dev = local_shape[0] * local_shape[1]

extra/llama_kernels/quantize_fp8_delayed/__init__.py

@@ -16,7 +16,7 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:
 
   wg = UOp.range(NUM_WG, 0, AxisType.GLOBAL)
   tid = UOp.range(THREADS_PER_WG, 1, AxisType.LOCAL)
-  it = UOp.range((n_elems // VEC) // (NUM_WG * THREADS_PER_WG), 2, AxisType.LOOP)
+  it = UOp.range((n_elems // VEC) // (NUM_WG * THREADS_PER_WG), 2, AxisType.WEAK)
   lane = UOp.range(VEC, 3, AxisType.UNROLL)
 
   idx = (((it * NUM_WG + wg) * THREADS_PER_WG + tid) * VEC) + lane
@@ -36,13 +36,13 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:
   lmax_val = lmax.after(lmax_store.end(it))[0]
 
   lds = UOp.placeholder((THREADS_PER_WG,), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
-  lds = lds.after(lds[tid].store(lmax_val).barrier())
+  lds = lds.after(lds[tid].store(lmax_val))
 
   step = THREADS_PER_WG // 2
   while step:
     active = tid < step
     other = lds[(tid + step).valid(active)].load()
-    lds = lds.after(lds[tid.valid(active)].store(lds[tid].maximum(other)).barrier())
+    lds = lds.after(lds[tid.valid(active)].store(lds[tid].maximum(other)))
     step //= 2
 
   device = device[0].split(":")[0] if isinstance(device, tuple) else device.split(":")[0]

extra/qcom_gpu_driver/opencl_ioctl.py

@@ -4,7 +4,7 @@ from hexdump import hexdump
 from copy import deepcopy
 import pathlib, sys
 from tinygrad.helpers import to_mv, getenv
-from tinygrad.runtime.autogen import adreno
+from tinygrad.runtime.autogen import mesa
 sys.path.append(pathlib.Path(__file__).parent.parent.parent.as_posix())
 
 IOCTL = getenv("IOCTL", 0)
@@ -23,7 +23,7 @@ for child in xml.getroot():
 CAPTURED_STATE = {}
 
 REGS = {}
-for k, v in adreno.__dict__.items():
+for k, v in mesa.__dict__.items():
   if k.startswith("REG_") and isinstance(v, int) and v > 1024: REGS[v] = k
 
 from extra.qcom_gpu_driver import msm_kgsl
@@ -42,7 +42,7 @@ def get_struct(argp, stype):
 
 def format_struct(s):
   sdats = []
-  for field_name, *_ in s._real_fields_:
+  for field_name, *_ in s._fields_:
     if field_name in {"__pad", "PADDING_0"}: continue
     dat = getattr(s, field_name)
     if isinstance(dat, int): sdats.append(f"{field_name}:0x{dat:X}")
@@ -96,9 +96,9 @@ def parse_cmd_buf(dat):
         CAPTURED_STATE['LOAD_FRAGS'].append((state_block, state_type, num_unit, dst_off))
 
         if state_block == SB6_CS_SHADER:
-          from extra.disassemblers.adreno import disasm_raw
+          from tinygrad.runtime.support.compiler_mesa import disas_adreno
           if state_type == ST6_SHADER and IOCTL > 3:
-            disasm_raw(get_mem(((vals[2] << 32) | vals[1]), num_unit * 128))
+            disas_adreno(get_mem(((vals[2] << 32) | vals[1]), num_unit * 128))
           if state_type == ST6_CONSTANTS:
             x = get_mem(((vals[2] << 32) | vals[1]), num_unit*4)
             CAPTURED_STATE['constants'] = x[:]
@@ -142,7 +142,7 @@ def parse_cmd_buf(dat):
       vals = struct.unpack("I"*size, dat[ptr+4:ptr+4+4*size])
       if IOCTL > 0: print(f"{ptr:3X} -- typ 4: {size=:3d}, {reg_name}", hprint(vals))
       for vi,v in enumerate(vals): CAPTURED_STATE[offset+vi] = v
-      if offset == adreno.REG_A6XX_SP_CS_CONFIG:
+      if offset == mesa.REG_A6XX_SP_CS_CONFIG:
         val = vals[0]
         if IOCTL > 0:
           print(f"\tBINDLESS_TEX={(val >> 0) & 0b1}")
@@ -215,79 +215,3 @@ def install_hook(c_function, python_function):
 
 libc = ctypes.CDLL(ctypes.util.find_library("libc"))
 install_hook(libc.ioctl, ioctl)
-
-def before_launch():
-  global CAPTURED_STATE
-  CAPTURED_STATE.clear()
-def collect_last_launch_state():
-  global CAPTURED_STATE
-  return deepcopy(CAPTURED_STATE)
-def compare_launch_state(state, good_state):
-  cmp = [
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_NTEX__MASK),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_NSAMP__MASK),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_NIBO__MASK),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_ENABLED),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_BINDLESS_TEX),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_BINDLESS_SAMP),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_BINDLESS_IBO),
-    (adreno.REG_A6XX_SP_CS_CONFIG, adreno.A6XX_SP_CS_CONFIG_BINDLESS_UBO),
-
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_HALFREGFOOTPRINT__MASK),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_FULLREGFOOTPRINT__MASK),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_BRANCHSTACK__MASK),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_FULLREGFOOTPRINT__MASK),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_THREADMODE__MASK),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_EARLYPREAMBLE),
-    (adreno.REG_A6XX_SP_CS_CTRL_REG0, adreno.A6XX_SP_CS_CTRL_REG0_MERGEDREGS),
-
-    (adreno.REG_A6XX_SP_CS_PVT_MEM_PARAM, adreno.A6XX_SP_CS_PVT_MEM_PARAM_MEMSIZEPERITEM__MASK),
-    (adreno.REG_A6XX_SP_CS_PVT_MEM_PARAM, adreno.A6XX_SP_CS_PVT_MEM_PARAM_HWSTACKSIZEPERTHREAD__MASK),
-
-    (adreno.REG_A6XX_SP_CS_UNKNOWN_A9B1, adreno.A6XX_SP_CS_UNKNOWN_A9B1_UNK5),
-    (adreno.REG_A6XX_SP_CS_UNKNOWN_A9B1, adreno.A6XX_SP_CS_UNKNOWN_A9B1_UNK6),
-
-    (adreno.REG_A6XX_SP_CS_BRANCH_COND, 0xffffffff),
-
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_0, adreno.A6XX_HLSQ_CS_NDRANGE_0_KERNELDIM__MASK),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_0, adreno.A6XX_HLSQ_CS_NDRANGE_0_LOCALSIZEX__MASK),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_0, adreno.A6XX_HLSQ_CS_NDRANGE_0_LOCALSIZEY__MASK),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_0, adreno.A6XX_HLSQ_CS_NDRANGE_0_LOCALSIZEZ__MASK),
-
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_1, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_2, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_3, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_4, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_5, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_NDRANGE_6, 0xffffffff),
-
-    (adreno.REG_A6XX_HLSQ_CS_CNTL_0, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_CNTL_1, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_KERNEL_GROUP_X, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_KERNEL_GROUP_Y, 0xffffffff),
-    (adreno.REG_A6XX_HLSQ_CS_KERNEL_GROUP_Z, 0xffffffff),
-  ]
-
-  for x,m in cmp:
-    print(f"Field {REGS[x]}, mask: 0x{m:X} cmp: {state.get(x, 0) & m} vs {good_state.get(x, 0) & m}")
-    if state.get(x, 0) & m != good_state.get(x, 0) & m:
-      return False, f"Field {REGS[x]}, mask: 0x{m:X} mismatch: {state.get(x, 0) & m} vs {good_state.get(x, 0) & m}"
-
-  for n in ['descriptors', 'ibos']:
-    if n not in good_state: continue
-    mv1, mv2 = state.get(n), good_state.get(n)
-
-    if len(mv1) != len(mv2): return False, f"{n}: len mismatch {len(mv1)} != {len(mv2)}"
-    mv1 = memoryview(bytearray(mv1)).cast('I')
-    mv2 = memoryview(bytearray(mv2)).cast('I')
-    for i in range(len(mv2)):
-      if i % 8 == 5 or i % 8 == 4: continue # addresses
-      if mv1[i]!=mv2[i]: return False, f"{n}: content mismatch {i} {mv1[i]} {mv2[i]}"
-
-  for n in ['samplers']:
-    if n not in good_state: continue
-    mv1, mv2 = state.get(n), good_state.get(n)
-    if len(mv1) != len(mv2): return False, f"{n}: len mismatch {len(mv1)} != {len(mv2)}"
-    if any(mv1[i]!=mv2[i] for i in range(len(mv1))): return False, f"{n}: content mismatch"
-
-  return True, "PASS"

extra/thunder/amd/fa.py

@@ -14,7 +14,7 @@ def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None
   shard_axis = ref.uop.axis if axis is None else axis
   shape = tuple(s // len(ref.device) if i == shard_axis else s for i, s in enumerate(shape))
   axis = ref.uop.axis if axis is None else axis
-  return Tensor(Tensor.invalids(*shape, dtype=dtype, device=ref.device).uop.multi(axis), dtype=dtype, device=ref.device)
+  return Tensor(Tensor.invalids(*shape, dtype=dtype, device=ref.device).uop.unshard(axis), dtype=dtype, device=ref.device)
 
 @functools.cache
 def custom_fused_qkv_rope_forward(q:UOp, k:UOp, v:UOp, xqkv:UOp, freqs_cis:UOp,

extra/thunder/tiny/fa.py

@@ -16,7 +16,7 @@ def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None) -> Tensor:
   if not isinstance(ref.device, tuple): return Tensor.empty(*shape, dtype=ref.dtype, device=ref.device)
   shape = tuple(s // len(ref.device) if i == ref.uop.axis else s for i, s in enumerate(shape))
   axis = ref.uop.axis if axis is None else axis
-  return Tensor(Tensor.empty(*shape, dtype=ref.dtype, device=ref.device).uop.multi(axis), dtype=ref.dtype, device=ref.device)
+  return Tensor(Tensor.empty(*shape, dtype=ref.dtype, device=ref.device).uop.unshard(axis), dtype=ref.dtype, device=ref.device)
 
 def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
   return _sharded_empty(ref.shape, ref, axis)

extra/thunder/tiny/tk/group.py

@@ -224,7 +224,7 @@ class Group:
 
       # store to shared memory
       red_local_store = red_local[self.laneid].store(red_reg[0])
-      red_local = red_local.after(red_local_store.barrier()).reshape(red_local.shape)
+      red_local = red_local.after(red_local_store).reshape(red_local.shape)
 
       # reduce from shared memory
       for inner in self.ker.range(3, axis_type=AxisType.REDUCE, track=False):
@@ -258,7 +258,7 @@ class Group:
 
       # store to shared memory
       red_local_store = red_local[self.laneid].store(red_reg[0])
-      red_local = red_local.after(red_local_store.barrier()).reshape(red_local.shape)
+      red_local = red_local.after(red_local_store).reshape(red_local.shape)
 
       # reduce from shared memory
       for inner in self.ker.range(3, axis_type=AxisType.REDUCE, track=False):
@@ -342,7 +342,7 @@ class Group:
           if src.dtype != dst.dtype:
             src_load = src_load.cast(dst.dtype)
           dst_store = dst[*dst_idxs, height, width, srow, scol].store(src_load)
-          dst_store = dst_store.end(height, width, outer, inner).barrier()
+          dst_store = dst_store.end(height, width, outer, inner)
     elif dst.addrspace == AddrSpace.REG and src.addrspace == AddrSpace.GLOBAL and isinstance(dst, RT):
       srcf = src.flatten()
       row_stride = prod(src.shape[axis+1:])

extra/thunder/tiny/tk/kernel.py

@@ -48,14 +48,14 @@ class Kernel(AbstractContextManager):
   @property
   def warpgroup(self): return self.group(4)
 
-  def range(self, start:int, end:int=0, step:int=1, axis_type:AxisType=AxisType.LOOP, track:bool=True):
+  def range(self, start:int, end:int=0, step:int=1, axis_type:AxisType=AxisType.WEAK, track:bool=True):
     if end == 0: start, end = 0, start
     rng = _tk_range(start, end, step, axis_type, self.range_id)
     self.range_id += 1
     if track: self.range_stack.append(rng)
     return rng
 
-  def raw_range(self, end:int=0, axis_type:AxisType=AxisType.LOOP):
+  def raw_range(self, end:int=0, axis_type:AxisType=AxisType.WEAK):
     rng = UOp.range(end, self.range_id, axis_type=axis_type)
     self.range_id += 1
     return rng

extra/torch_backend/backend.py

@@ -555,6 +555,8 @@ def wrap_out(f):
     assert out.shape == assigned.shape, f"shape mismatch: {assigned.shape} -> {out.shape}"
     assert out.device == assigned.device or out.device is None or assigned.device is None, f"device mismatch: {assigned.device} -> {out.device}"
     assert out.dtype == assigned.dtype, f"dtype mismatch: {assigned.dtype} -> {out.dtype}"
+    # an out= that is a view has to be written through its base, and _apply_inplace gives a deviceless base its buffer first
+    if canonical_base(out) is not out: return _apply_inplace(out, assigned) or out
     if out.device is None and assigned.device is not None: out.replace(out.empty_like(device=assigned.device))
     return out.assign(assigned)
   return _wrap_out

extra/torch_backend/test.py

@@ -98,6 +98,21 @@ class TestTorchBackend(unittest.TestCase):
   def test_empty_strided_default_dtype(self):
     self.assertEqual(torch.empty_strided((2,3), (1,2), device=device).dtype, torch.get_default_dtype())
 
+  @unittest.expectedFailure  # TODO: empty_strided ignores the requested strides, the backend treats everything as contiguous
+  def test_empty_strided_honors_strides(self):
+    self.assertEqual(tuple(torch.empty_strided((2,3), (1,2), device=device).stride()), (1,2))
+
+  @unittest.expectedFailure  # TODO: torch refuses an out= that overlaps an input, we compute silently
+  def test_out_overlapping_input_is_rejected(self):
+    x = torch.arange(6., device=device)
+    with self.assertRaises(RuntimeError): torch.add(x[:-1], 10, out=x[1:])
+
+  def test_out_disjoint_input_is_allowed(self):
+    # torch permits an out= that shares a base with an input as long as they do not overlap
+    x, xc = torch.arange(6., device=device), torch.arange(6.)
+    torch.add(x[:3], 10, out=x[3:]); torch.add(xc[:3], 10, out=xc[3:])
+    np.testing.assert_equal(x.cpu().numpy(), xc.numpy())
+
   def test_plus_inplace(self):
     a = torch.ones(4, device=device)
     b = torch.ones(4, device=device)

multi_migration.md

@@ -1,86 +0,0 @@
-# Multi-device op migration: MULTI/MSELECT/MSTACK → PAD / WHERE / STACK+INDEX
-
-## Status (updated)
-
-- **Stage 0 — DONE.** Internal `Ops.PAD` fills **Invalid** (`schedule/indexing.py:104`, bool keeps 0-fill); external `Tensor.pad`/`pad_to` always emit an explicit fill mask (`mixin/op.py:289`, `mixin/movement.py:267`) — required because a bare Invalid-pad leaks through elementwise ALU (`pad(x)+1` would read 0 instead of 1 in pad regions). REDUCE inputs with Invalid contribute the reduce identity (`pm_invalid_reduce_identity` in `uop/symbolic.py`, run in `get_kernel_graph` after gate lifting in `schedule/rangeify.py`) — only WHERE-alt gates whose condition involves a reduce range are rewritten, so gather-with-Invalid-index still poisons whole lanes. Same-condition nested where collapse rule added (`c?(c?t:f):f2 -> c?t:f2`) so the mask form folds to a single gate. All suites green (`test/unit`, `test/null`, `test/backend`, `test/external/external_test_schedule_scaling.py`, mypy, ruff).
-- **Stage 1 — representation in place behind `SYMBOLIC_MULTI`.** `symbolic_multi_pm` (`schedule/multi.py`) converts `MULTI→_unshard` (raw Invalid pad), `MSELECT→dnum.eq(i).where(x, Invalid)`, `MSTACK→STACK.index(dnum)` + INDEX(STACK,var)→nested-where lowering. `_unshard` uses the raw Invalid pad; `_unshard_fill` (0-fill) is used for the ALU allreduce in `copy_multi` because gated stores leave stale pad regions (the ALU-sum path can't use the identity rule). Basic shard ops work; full parity is Stage 2.
-- **Stage 2 — remaining.** Buffer level, reduce/allreduce split for shard-axis reduces, API surface.
-- **Stage 3 — remaining.**
-
-Notes: `test_schedule.py:test_pad_reduce_unsafe_multiview_st` went 4→5 kernels (pad now materializes an explicit mask; the mask form is also what makes the previously-wrong masked-pad+hazard case correct). `test_jit_footguns.py:test_symbolic_pad_view_frozen` went 2→4: the explicit mask recomputes from the symbolic shape, fixing the frozen-pad footgun. Also fixed a latent infinite loop: `(x+y) !=/< c → x !=/< c-y` collapse rules in `codegen/simplify.py` now only fire when the remaining side still contains the range (they previously shuffled constants forever when both sides were range-free).
-
-## Goal
-
-Replace the three multi-device UOps with a symbolic `_device_num` representation:
-
-| Old op | New form |
-|---|---|
-| `MULTI(x, axis)` | `x._unshard(axis)` — PAD with `_device_num`-dependent bounds back to full shape (helper already exists at `tinygrad/uop/ops.py:704-707`) |
-| `MSELECT(x, i)` | `dnum.eq(i).where(x, x.const_like(Invalid))` |
-| `MSTACK(s0..sn)` | `UOp(Ops.STACK, src=srcs).index(dnum)` — leading device axis, indexed per-device |
-
-where `dnum = UOp.variable("_device_num", 0, ndev-1)`. The per-device specialization mechanism already exists: `unwrap_multi` (`tinygrad/engine/realize.py:148-153`) binds `_device_num` per device at exec time.
-
-**Key semantic decision (approved):** internal `Ops.PAD` produces **Invalid** in padded regions; external `Tensor.pad` API still pads with 0. Staged migration: introduce the new representation first, keep old ops working, migrate call sites incrementally, delete old ops last.
-
-## Background: current design
-
-- `Ops.MULTI(src, axis)` (`tinygrad/uop/__init__.py:100`) — per-shard graph marker. Eliminated by `multi_pm` (`tinygrad/schedule/multi.py:162-195`) as the first step of `get_kernel_graph` (`tinygrad/schedule/rangeify.py:548`). Shape/axis tracking: `UOp.axis`/`UOp.bounds` (`tinygrad/uop/ops.py:667-702`).
-- `Ops.MSELECT(x, i)` / `Ops.MSTACK(srcs)` (`__init__.py:96`) — buffer-level ops. Spec at `tinygrad/uop/spec.py:181-184`; device prop `ops.py:816-819`; per-kernel PARAMs via debuf (`rangeify.py:474`); per-device dependency states (`tinygrad/schedule/__init__.py:11-17`); `MultiBuffer` (`ops.py:904-930`, `tinygrad/device.py:88-99`); only MSTACK can be `realized` (`ops.py:920-930`).
-- `_shard`/`_unshard` (`ops.py:704-714`) already emit symbolic SHRINK/PAD bounds with `_device_num`.
-- Naive allreduce already uses the target pattern: `dnum.eq(i).where(buf, state)` (`tinygrad/schedule/allreduce.py:27-33`).
-
-## Existing Invalid machinery (rely on this)
-
-- `pm_data_invalid` (`tinygrad/uop/symbolic.py:71-92`): Invalid poisons ALU (ops move inside the gate); gated LOAD folds to alt/0, gated STORE folds to NOOP.
-- `pm_remove_invalid` (`symbolic.py:94-96`): leftover Invalid → 0 in final codegen (`codegen/__init__.py:345`). Spec forbids Invalid in final programs (`spec.py:217`), so materialized Invalid regions read as 0.
-- STORE of CONST(Invalid) → NOOP (`rangeify.py:423-424`).
-- `identity_element(op, dtype)` exists (`ops.py:51`): ADD→0, MUL→1, MAX→dtype.min.
-- `found_after` (`rangeify.py:26`) already matches `WHERE(cond, PAD(x), Invalid)`.
-
-## Stage 0 — internal PAD = Invalid; external pad = explicit 0
-
-1. `tinygrad/schedule/indexing.py:100-104` (`convert_pad_to_where_to_keep_behavior_local`): fill value `0` → `UOp.const(x.dtype, Invalid)`, **except `dtypes.bool` keeps 0-fill** (False is the bool-reduce identity, and the external-pad mask below needs it).
-2. `tinygrad/mixin/op.py:282-290` (`_pad_constant`): **remove the `if value == 0: return base` shortcut** — always emit `pad(bool_ones).where(base, value)`. Required because bare Invalid-pad leaks through elementwise ALU: `pad(x)+1` gate-lifts to `where(valid, x+1, Invalid)` and reads 0 instead of 1 in pad regions. The mask lowers to a pure index expression (`valid.where(1,0)`), no extra kernel. External behavior unchanged for all `value`.
-3. **New rule**: `REDUCE(where(c, x, Invalid), op)` → `REDUCE(where(c, x, identity_element(op, dtype)), op)`. Must fire in rangeify/symbolic *before* codegen builds the accumulator loop — otherwise `pm_data_invalid` gate-lifts `acc + where(c,x,Invalid)` into `where(c, acc+x, Invalid)` and one invalid lane poisons the whole reduction. Placement (symbolic.py vs the reduce path in indexing.py) TBD at implementation; verify with `Tensor.pad(...).sum()/max()` tests.
-4. Audit: schedule tests with kernel counts involving pads; circular/reflect/replicate pads don't use PAD fill (verified, `op.py:292-312`) — unaffected; `allreduce.py:59,76` usum-of-padded-chunks gets *more* correct (disjoint regions).
-
-## Stage 1 — new representation behind env flag
-
-New `symbolic_multi_pm` PatternMatcher (in `schedule/multi.py` or new file), gated by env (e.g. `SYMBOLIC_MULTI`), run in `get_kernel_graph` right after `multi_pm`:
-
-- `MULTI(x, axis)` → `x._unshard(axis)`
-- `MSELECT(x, i)` → `dnum.eq(i).where(x, x.const_like(Invalid))` (Invalid from `tinygrad.dtype`)
-- `MSTACK(srcs)` → `STACK(*srcs).index(dnum)`, plus new lowering `INDEX(STACK(vals), var)` → nested `var.eq(k).where(src_k, Invalid)` (analogous to `convert_stack_to_where`, `indexing.py:113-121`; must fire before `validate_index` spec, `spec.py:118-122`)
-
-Flag off = zero behavior change; flag on = new forms flow through rangeify and specialize per device at exec.
-
-## Stage 2 — migrate producers/consumers (one commit each, independently testable)
-
-1. `UOp.shard` (`ops.py:715-717`): emit symbolic `_shard`+`_unshard` full-shape form directly instead of `.multi(axis)`; delete movement-op `multi_pm` rules that PAD subsumes (`pad_multi`, `permute_multi`, `expand_multi`, `reshape_multi`, `flip_multi`, `shrink_multi` — `multi.py:93-125`).
-2. ALU/STACK: `alu_multi`/`shard_srcs`/`stack_multi` (`multi.py:55-78,127-131`) become plain elementwise on full-shape padded tensors. `reduce_multi` (`multi.py:80-91`) keeps the shard-axis → local-reduce + ALLREDUCE split; Invalid-pad + identity rule replaces neutral-pad-value reasoning.
-3. allreduce (`schedule/allreduce.py`): naive path already matches; migrate ring/all2all MSELECT/MSTACK scratch-buffer assembly (lines 35-76) to WHERE/STACK+INDEX forms.
-4. Buffer level: debuf (`rangeify.py:474`), `_states`/`_unwrap_src` (`schedule/__init__.py:11-17`), `_collect_bufs` (`schedule/memory.py:9`), `unwrap_multi` (`realize.py:148-153`), JIT (`jit.py:127-130, 237`), callify (`callify.py:52-95`), `buffer`/`realized`/`buf_uop`/`has_buffer_identity` (`ops.py:841-930`).
-5. API surface: `UOp.multi/mselect/mstack` (`ops.py:662-725`), `Tensor.shard` (`tensor.py:333-347`), gradient (`mixin/gradient.py:72`), `_multi_like` (`mixin/creation.py:16-20`), embedding backward (`nn/__init__.py:309-354`), `copy_to_device(arg=)` MSELECT path (`ops.py:719-723`).
-
-## Stage 3 — removal
-
-Delete `Ops.MULTI/MSELECT/MSTACK` from the enum (`uop/__init__.py:96,100`), spec rules, viz colors (`viz/serve.py:51,56`), `UOp.axis`/`bounds` machinery (`ops.py:667-702`), remaining `multi_pm` rules, and `MultiBuffer` if fully subsumed. Flip flag default-on, then delete the flag.
-
-## Open implementation details
-
-- REDUCE-identity rule placement (must precede codegen accumulator construction).
-- INDEX(STACK, var) spec timing — the value-STACK INDEX violates the pointer-INDEX spec until lowered.
-- Whether `MultiBuffer`/tuple-`device` survives as the runtime container, or buffers become single-device with the device axis explicit in shape — decides how much of Stage 2.4 is rewrite vs delete.
-- Bool carve-out in Stage 0.1: verify no internal consumer needs Invalid-filled bool pads.
-
-## Verification (run at each stage)
-
-```bash
-python -m pytest test/unit/test_multitensor.py test/unit/test_allreduce.py test/null/test_multitensor.py test/unit/test_call.py -x -q -n12
-python -m pytest test/external/external_test_schedule_scaling.py -x -q   # test_concat_scaling
-python -m mypy tinygrad/
-python -m ruff check .
-```
-
-Also pad/reduce numeric tests after Stage 0 (`test_ops` pad tests, `Tensor.pad(...).sum()/max()`).

spec/tinyspec.tex

@@ -29,6 +29,7 @@
 \definecolor{axbrred}{HTML}{E53935}      % GROUP_REDUCE
 \definecolor{axyellow}{HTML}{F9A825}     % UPCAST
 \definecolor{axmagenta}{HTML}{7B1FA2}    % UNROLL
+\definecolor{axgreen}{HTML}{2E7D32}      % DEVICE
 
 \title{tinygrad: a single dialect from Tensor programs to Command Buffers}
 \author{tinygrad, Corp. \\ \texttt{[email protected]}}
@@ -79,6 +80,7 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op
 \op{Index} & $(T, i_0, i_1, \ldots)$ & --- & Index from left. $()$-shaped $i$ removes dim; $(k,)$-shaped makes it $k$. \\
 \op{Stack} & $(T_0, T_1, \ldots)$ & --- & Join along a newly created leading axis. All shapes must match. \\
 \op{Bitcast} & $(T,)$ & dtype & Reinterpret storage as target dtype; preserve total bytes. \\
+\op{Unshard} & $(T, R_0, R_1, \ldots)$ & axes $(a_0, a_1, \ldots)$ & Concatenate shards of \op{Range} $R_k$ along axis $a_k$; $R_k$ is outer. \\
 \bottomrule
 \end{tabular}
 
@@ -258,6 +260,7 @@ Every UOp has a \textbf{dtype}, \textbf{shape}, \textbf{device}, \textbf{addrspa
 \op{Const} & from arg & $()$ & \textsc{null} & $[v, v]$ \\
 \op{Param} & from arg & from $\mathrm{src}[0]$ & from arg & from src or dtype range \\[3pt]
 Movement ops & $\mathrm{src}[0].\mathrm{dtype}$ & (see op) & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
+\op{Unshard} & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0]$, each $a_k \times n_k$ & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
 \op{Reduce} & $\mathrm{src}[0].\mathrm{dtype}$ & remove first $n$ axes & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\[3pt]
 \op{Cast} & from arg & $\mathrm{src}[0].\mathrm{shape}$ & $\mathrm{src}[0].\mathrm{device}$ & clamped to dtype \\
 \op{Bitcast} & from arg & $\mathrm{src}[0].\mathrm{shape}$ & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\
@@ -283,9 +286,9 @@ $[a,A]$, $[b,B]$, $[c,C]$ denote min\_max of $\mathrm{src}[0]$, $\mathrm{src}[1]
 Default \emph{dtype range}: $[\mathrm{dtype\_min},\, \mathrm{dtype\_max}]$.
 
 \medskip
-\textbf{axis} tracks the multi-device sharding dimension. \op{Buffer} with $n$-tuple device: axis $= 0$ (device dim).
-\op{Reshape} remaps axis to preserve the shard boundary. \op{Permute} follows the permutation. \op{Expand} shifts axis right by $|\mathbf{n}|$.
-\op{Reduce} on the shard axis $\to$ \textsc{null} (shard axis is among the first $n$ axes). \op{Replicated} on the shard axis $\to$ \textsc{null}. \op{Copy} $\to$ \textsc{null}. ALU ops inherit from sources. Default: \textsc{null}.
+\textbf{sharding} tracks multi-device sharding as a set of (axis, \op{Range}) pairs. \op{Unshard} defines it: arg is the tuple of sharded axes, one \op{Range} in src per axis (positional: the $k$-th \op{Range} shards the $k$-th axis). \op{Buffer} with $n$-tuple device: sharded on axis $0$ (device dim). The single-axis convenience \textbf{axis} is \textsc{null} unless exactly one axis is sharded.
+\op{Reshape} remaps each sharded axis to preserve its shard boundary. \op{Permute} follows the permutation. \op{Expand} shifts all sharded axes right by $|\mathbf{n}|$.
+\op{Reduce} on a sharded axis drops it. \op{Replicated} on the shard axis $\to$ \textsc{null}. \op{Copy} $\to$ \textsc{null}. ALU ops inherit from sources. Default: \textsc{null}.
 
 %% ============================================================
 \subsection*{Kernel Optimizations (OptOps) \normalfont\small--- schedule-level transforms on kernel ranges}
@@ -297,6 +300,7 @@ Each kernel's iteration space is a set of \op{Range} axes. Every range has an \t
 \toprule
 \textbf{AxisType} & \textbf{Letter} & \textbf{Split from} & \textbf{Direction} & \textbf{Semantics} \\
 \midrule
+{\color{axgreen}\texttt{DEVICE}} & \texttt{d} & --- & --- & Multi-device sharding dimension. \\
 {\color{axblue}\texttt{GLOBAL}} & \texttt{g} & --- & --- & GPU global workgroup dimension. \\
 {\color{axcyan}\texttt{LOCAL}} & \texttt{l} & g, L & inner & Workgroup local dimension (shared memory). \\
 {\color{axbrcyan}\texttt{WARP}} & \texttt{w} & \multicolumn{2}{l}{(created by \op{TC})} & Warp-level lanes for tensor cores. \\
@@ -378,6 +382,8 @@ def scatter_add(T, idx, val):
 Let $D = (d_0, \ldots, d_{n-1})$ be an $n$-tuple device.
 \op{Copy} to an $n$-tuple device reshards with axis $= 0$. \op{Copy} never changes shape.
 
+\textbf{Sharding} splits a tensor along an axis across $n$ devices. It opens a \op{Range} of type \texttt{DEVICE} (a symbolic per-device index $d$), shrinks each device's view to its piece, then closes the range with \op{Unshard}$(T, R, a)$. The result is a logical tensor whose shape along axis $a$ is the full size; each device holds $1/n$ of it. \op{Unshard} is the inverse of sharding --- it marks the boundary between per-device computation and the logical multi-device tensor. The range need not be \texttt{DEVICE}; e.g.\ a \texttt{WARP} range closes the same way, concatenating per-lane shards along $a$ with the range as the outer factor. A tensor may be sharded along several axes at once: \op{Unshard}$(T, R_0, R_1, \ldots;\; a_0, a_1, \ldots)$ carries one \op{Range} per sharded axis, and every movement op maps each sharded axis independently.
+
 \begin{lstlisting}
 # T has shape (s,) on a single device.
 

test/backend/test_amd_flash_attention.py

@@ -0,0 +1,65 @@
+import unittest
+import numpy as np
+
+from tinygrad import Device, Tensor, TinyJit
+from tinygrad.llm.kernels.amd import amd_flash_attention_decode, flash_attention_causal_cached
+
+
+@unittest.skipUnless(Device.DEFAULT.startswith("AMD"), "AMD flash attention required")
+class TestAMDFlashAttention(unittest.TestCase):
+  def _test_decode(self, max_kv_len:int, valid_kv_len:int, n_heads:int=16, n_kv_heads:int=2, quantized:bool=False):
+    rng = np.random.default_rng(1)
+    q_np = rng.standard_normal((1, n_heads, 1, 256)).astype(np.float16)
+    kv_np = rng.standard_normal((2, 1, n_kv_heads, max_kv_len, 256)).astype(np.float16)
+    scale_np = np.maximum(np.max(np.abs(kv_np.astype(np.float32)), axis=-1), 1e-8) / 127
+    if quantized:
+      kv_np = np.clip(np.rint(kv_np.astype(np.float32) / scale_np[..., None]), -127, 127).astype(np.int8)
+    q, kv = Tensor(q_np).realize(), Tensor(kv_np).realize()
+    scale = Tensor(scale_np.astype(np.float16)).realize() if quantized else None
+
+    @TinyJit
+    def decode(q:Tensor, kv:Tensor): return amd_flash_attention_decode(q, kv, valid_kv_len, max_kv_len, scale).realize()
+
+    out = None
+    for _ in range(3): out = decode(q, kv).numpy()
+    assert out is not None
+    q_ref = q_np[0, :, 0].astype(np.float32)
+    kv_ref = kv_np.astype(np.float32) * scale_np[..., None] if quantized else kv_np.astype(np.float32)
+    k_ref, v_ref = kv_ref[:, 0, :, :valid_kv_len]
+    expected = np.empty((n_heads, 256), dtype=np.float32)
+    for head in range(n_heads):
+      scores = q_ref[head] @ k_ref[head // (n_heads // n_kv_heads)].T / np.sqrt(256)
+      probs = np.exp(scores - scores.max())
+      expected[head] = probs @ v_ref[head // (n_heads // n_kv_heads)] / probs.sum()
+
+    self.assertTrue(np.isfinite(out).all())
+    np.testing.assert_allclose(out[0, :, 0], expected, rtol=2e-3, atol=2e-3)
+
+  def test_short_decode_is_finite_and_matches_reference(self): self._test_decode(8192, 25)
+
+  def test_q8_cache_matches_dequantized_reference(self): self._test_decode(8192, 25, quantized=True)
+
+  def test_q8_cached_prefill_matches_dequantized_reference(self):
+    rng = np.random.default_rng(2)
+    heads, kv_heads, tokens, dim = 16, 2, 32, 256
+    q = rng.standard_normal((1, heads, tokens, dim)).astype(np.float16)
+    kv = rng.standard_normal((2, 1, kv_heads, tokens, dim)).astype(np.float16)
+    scale = np.maximum(np.max(np.abs(kv.astype(np.float32)), axis=-1), 1e-8) / 127
+    packed = np.clip(np.rint(kv.astype(np.float32) / scale[..., None]), -127, 127).astype(np.int8)
+    got = flash_attention_causal_cached(Tensor(q).realize(), Tensor(packed).realize(), tokens, tokens,
+                                        Tensor(scale.astype(np.float16)).realize()).numpy()
+    dequant = packed.astype(np.float32) * scale.astype(np.float16).astype(np.float32)[..., None]
+    expected = np.empty_like(got)
+    for head in range(heads):
+      scores = q[0, head].astype(np.float32) @ dequant[0, 0, head // (heads // kv_heads)].T / np.sqrt(dim)
+      scores[np.triu_indices(tokens, 1)] = -np.inf
+      probs = np.exp(scores - scores.max(axis=-1, keepdims=True))
+      expected[0, head] = probs @ dequant[1, 0, head // (heads // kv_heads)] / probs.sum(axis=-1, keepdims=True)
+    np.testing.assert_allclose(got, expected, rtol=2e-3, atol=2e-3)
+
+  def test_six_query_heads_per_kv_head(self): self._test_decode(8192, 25, n_heads=12, n_kv_heads=2)
+
+  def test_hierarchical_decode_matches_reference(self): self._test_decode(16384, 4097)
+
+
+if __name__ == "__main__": unittest.main()

test/backend/test_arange.py

@@ -19,6 +19,12 @@ class TestArange(unittest.TestCase):
     self.assertLess(self._get_flops(Tensor.arange(256).clone(), np.arange(256)), 256*4)
     self.assertLess(self._get_flops(Tensor.arange(2560).clone(), np.arange(2560)), 2560*4)
 
+  def test_cat_complexity(self):
+    x = Tensor.arange(2**10) + Tensor.empty((), dtype=dtypes.uint32)
+    out = x.cat(x).cat(Tensor.empty(1, dtype=dtypes.uint32))
+    linear = compile_linear(out.schedule_linear())
+    self.assertLessEqual(estimate_uop(linear.src[-1]).ops, out.numel()*20)
+
   @unittest.skipIf(Device.DEFAULT == "CL", "flaky in CI")
   def test_arange_cumsum(self):
     np.testing.assert_equal(Tensor.arange(513).cumsum(0).numpy(), np.arange(513).cumsum())

test/backend/test_custom_kernel.py

@@ -2,6 +2,7 @@ import unittest
 from tinygrad import Tensor, UOp, GlobalCounters, Context, Device
 from tinygrad.dtype import AddrSpace, dtypes, Invalid
 from tinygrad.uop.ops import KernelInfo, AxisType, Ops
+from tinygrad.renderer.ptx import PTXRenderer
 
 # **** kernels ****
 
@@ -55,7 +56,7 @@ def flip_contract_kernel(dest:UOp, src:UOp):
   return store.end(i, j).sink(arg=KernelInfo(name=f"flip_contract_{dest.numel()}", opts_to_apply=()))
 
 def slice_sum_kernel(dest:UOp, src:UOp):
-  G = UOp.range(src.shape[0], 0)
+  G = UOp.range(src.shape[0], 0, dtype=dtypes.int)
   slice_src = src[G, :]
   reg = UOp.placeholder((1,), dest.dtype, 0, addrspace=AddrSpace.REG)
   reg = reg.after(G)[0].set(0)
@@ -117,13 +118,20 @@ class TestCustomKernel(unittest.TestCase):
     out = c.flatten().tolist()
     assert all(x == 2 for x in out), "all 2"
 
+  def test_duplicate_call_arg(self):
+    x = Tensor.arange(4).clone().realize()
+    x = Tensor.custom_kernel(x, x, fxn=custom_add_one_kernel)[0]
+    # webgpu silently errors when a kernel has duplicate buffer args, so the list stays the same.
+    # https://gpuweb.github.io/gpuweb/#abstract-opdef-encoder-bind-groups-alias-a-writable-resource
+    self.assertEqual(x.tolist(), [1, 2, 3, 4] if Device.DEFAULT != "WEBGPU" else [0, 1, 2, 3])
+
   def test_simple_sharded(self):
     devs = ("CPU:0", "CPU:1")
 
     a = Tensor.ones(16, 16).contiguous().shard(devs, axis=0)
     b = Tensor.ones(16, 16).contiguous().shard(devs, axis=0)
     # ugly construction to get a sharded empty tensor
-    c = Tensor(Tensor.empty(8, 16, device=devs).uop.multi(0), device=devs)
+    c = Tensor(Tensor.empty(8, 16, device=devs).uop.unshard(0), device=devs)
     c = Tensor.custom_kernel(c,a,b, fxn=custom_elementwise_add_kernel)[0]
     out = c.flatten().tolist()
     assert all(x == 2 for x in out), "all 2"
@@ -132,7 +140,7 @@ class TestCustomKernel(unittest.TestCase):
     # PYTHON backend explicitly checks for OOB access for wrong multi shape regression
     devs = ("PYTHON:0", "PYTHON:1")
     a = Tensor.ones(4, 4).contiguous().shard(devs, axis=0)
-    c = Tensor(Tensor.empty(2, 4, device=devs).uop.multi(0), device=devs)
+    c = Tensor(Tensor.empty(2, 4, device=devs).uop.unshard(0), device=devs)
     c = Tensor.custom_kernel(c, a, fxn=custom_add_one_kernel)[0]
     assert (c == 2).all().item()
 
@@ -206,7 +214,7 @@ class TestCustomKernel(unittest.TestCase):
     N = 16
     a = Tensor.randn(N, N).shard_(devs, axis=0)
     b = Tensor.randn(N, N).to(devs)
-    c = Tensor(Tensor.empty(N//2, N, device=devs).uop.multi(0), device=devs)
+    c = Tensor(Tensor.empty(N//2, N, device=devs).uop.unshard(0), device=devs)
     tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
     self.assertTrue(tst.allclose(a@b, atol=1e-3).item())
 
@@ -323,7 +331,7 @@ class TestCustomKernel(unittest.TestCase):
   def test_multi_invalids_custom_kernel_no_copy(self):
     devs = ("CPU:0", "CPU:1")
     a = Tensor.ones(4, 4).shard(devs, axis=0).realize()
-    c = Tensor(Tensor.invalids(2, 4, dtype=dtypes.float, device=devs).uop.multi(0), device=devs)
+    c = Tensor(Tensor.invalids(2, 4, dtype=dtypes.float, device=devs).uop.unshard(0), device=devs)
     c = Tensor.custom_kernel(c, a, fxn=custom_add_one_kernel)[0]
     GlobalCounters.reset()
     c.realize()
@@ -388,7 +396,7 @@ class TestCustomKernel(unittest.TestCase):
     y = Tensor.custom_kernel(y, x, fxn=custom_add_one_kernel)[0]
     if use_custom:
       z = Tensor.empty_like(x)
-      z = Tensor.custom_kernel(y, y.T.T, fxn=custom_add_one_kernel)[0]
+      z = Tensor.custom_kernel(z, y.T.T, fxn=custom_add_one_kernel)[0]
     else: z = y.T.T+1
     GlobalCounters.reset()
     z.realize()
@@ -414,19 +422,76 @@ class TestCustomKernel(unittest.TestCase):
 
   @Context(DEV="CPU")
   def test_simple_from_source(self):
-    a = Tensor([0., 1., 2.]).realize()
-
-    src = "void test_src(float* restrict a) { a[0] = 1.0; }"
+    a = Tensor.arange(4).clone().realize()
+    src = "void test_src(int* restrict a) { a[0] = 1; }"
     # TODO: it currently requires a compiler for Ops.BINARY
     from tinygrad.device import Device
     binary = Device[a.device].renderer.compiler.compile(src)
     def custom_src_kernel(A:UOp) -> UOp:
       sink = UOp.sink(A, arg=KernelInfo(name="test_src"))
-      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())),
-                                   UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
+    a = Tensor.custom_kernel(a.reshape(2, 2).T, fxn=custom_src_kernel)[0]
+    self.assertEqual(a.tolist(), [[1, 2], [1, 3]])
 
-    a = Tensor.custom_kernel(a, fxn=custom_src_kernel)[0]
-    self.assertEqual(a.tolist(), [1., 1., 2.])
+class TestUnshardIndex(unittest.TestCase):
+  """Regression tests for INDEX on UNSHARD (fragment) resolution in schedule/multi.py.
+
+  A fragment is a per-thread REG buffer wrapped in UNSHARD over LOCAL thread ranges.
+  index_multi must resolve an INDEX on the UNSHARD view into an INDEX on the per-thread
+  shard. Two ownership patterns must work:
+    contiguous: idx = rng*shard_sz + local   (thread rng owns [rng*shard_sz, ...))
+    strided:    idx = rng + ir*shard_sz      (thread rng owns {rng, rng+shard_sz, ...})
+  """
+  def _run(self, kernel, shape=(8, 8)):
+    c = Tensor.empty(*shape)
+    out = Tensor.custom_kernel(c, fxn=kernel)[0]
+    try: return out.numpy()
+    except RuntimeError as e:
+      if isinstance(Device[Device.DEFAULT].renderer, PTXRenderer) and "dynamic register indexing" in str(e):
+        self.skipTest("PTX does not support dynamic register indexing")
+      raise
+
+  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
+  def test_contiguous_fragment_index(self):
+    # thread ty owns rows [ty*8, ty*8+8) of a 64-row fragment -- contiguous ownership.
+    # This is the pre-existing case that index_multi always handled.
+    def kernel(C:UOp) -> UOp:
+      ty = UOp.range(8, 0, AxisType.LOCAL)
+      ir = UOp.range(8, 1, AxisType.LOOP)
+      j = UOp.range(8, 2, AxisType.LOOP)
+      # 8x8 fragment, 8 threads -> 64x8 full tile. thread ty owns rows [ty*8, ty*8+8).
+      frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
+      return C[ty*8 + ir, j].store(frag[ty*8 + ir, j]).end(j, ir, ty).sink(arg=KernelInfo(name="contig_frag"))
+    out = self._run(kernel, (64, 8))
+    assert out.shape == (64, 8)
+
+  @unittest.skipIf(not Device[Device.DEFAULT].renderer.has_local, "fragment tests need LOCAL ranges")
+  def test_strided_fragment_index(self):
+    # thread ty owns rows {ty, ty+8, ty+16, ty+24, ..., ty+56} of a 64-row fragment --
+    # strided ownership. idx = ty + ir*8 where shard_sz=8 (8 threads, shard rows=8).
+    # The contiguous check (idx - rng*shard_sz) fails; the strided check
+    # (idx-rng) % shard_sz == 0 must succeed. This is the pattern the index_multi fix adds.
+    def kernel(C:UOp) -> UOp:
+      ty = UOp.range(8, 0, AxisType.LOCAL)
+      ir = UOp.range(8, 1, AxisType.LOOP)
+      j = UOp.range(8, 2, AxisType.LOOP)
+      # 8x8 fragment, 8 threads -> 64x8 full tile. thread ty owns rows {ty, ty+8, ..., ty+56}.
+      frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
+      return C[ty + ir*8, j].store(frag[ty + ir*8, j]).end(j, ir, ty).sink(arg=KernelInfo(name="strided_frag"))
+    out = self._run(kernel, (64, 8))
+    assert out.shape == (64, 8)
+
+  def test_fragment_index_cannot_shard(self):
+    # thread ty indexing rows [ty, ty+8) overlaps with other threads' rows -- this matches neither
+    # the contiguous nor the strided ownership pattern, so index_multi must raise.
+    def kernel(C:UOp) -> UOp:
+      ty = UOp.range(8, 0, AxisType.LOCAL)
+      ir = UOp.range(8, 1, AxisType.LOOP)
+      j = UOp.range(8, 2, AxisType.LOOP)
+      frag = UOp.placeholder((8, 8), dtypes.float32, 0, AddrSpace.REG).unshard((0,), (ty,))
+      return C[ty + ir, j].store(frag[ty + ir, j]).end(j, ir, ty).sink(arg=KernelInfo(name="bad_frag"))
+    with self.assertRaisesRegex(RuntimeError, "cannot shard index"):
+      self._run(kernel, (64, 8))
 
 class TestUOpReduce(unittest.TestCase):
   def test_uop_sum(self):

test/backend/test_dtype.py

@@ -124,6 +124,7 @@ class TestFp8sConversions(unittest.TestCase):
   def test_float_to_fp8e4m3(self, x):
     np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.float8_e4m3fn).view(torch.uint8).item())
 
+  @unittest.skip("fp8 overflow semantics are inconsistent")
   def test_float_to_fp8e4m3_extreme_values(self):
     for x in [FP8E4M3_MAX, FP8E4M3_MAX*1.01, -FP8E4M3_MAX, -FP8E4M3_MAX*1.01, math.inf, -math.inf, math.nan, -math.nan]:
       np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3), torch.tensor(x, dtype=torch.float8_e4m3fn).view(torch.uint8).item())

test/backend/test_linearizer.py

@@ -13,6 +13,7 @@ from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.cstyle import CUDARenderer
 from tinygrad.renderer.isa import ISARenderer
 from test.helpers import replace_opts
+from test.backend.test_softmax_fusion import single_kernel_softmax
 MOCKGPU = DEV.interface.startswith("MOCK")
 
 from tinygrad.uop.render import print_uops # noqa: F401 # pylint: disable=unused-import
@@ -392,6 +393,16 @@ class TestLinearizer(unittest.TestCase):
     # the global store doesn't change
     assert stores[1].src[1].dtype == dtypes.float
 
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
+  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared")
+  def test_two_grouped_stores_local(self):
+    # GROUP on both reduces puts two LOCAL buffers in one kernel, and the store to each needs its own barrier
+    a = Tensor.rand(32, 32).realize()
+    opts = [Opt(OptOps.GROUP, 1, 4), Opt(OptOps.GROUP, 2, 4)]
+    ast = helper_linearizer_opt(single_kernel_softmax(a), [opts])
+    uops = to_program(replace_opts(ast, opts), renderer=Device[Device.DEFAULT].renderer).src[1].src
+    self.assertEqual(len([u for u in uops if u.op is Ops.BARRIER]), 2)
+
 # *** helpers ***
 
 def helper_realized_ast(r:Tensor|list[Tensor]) -> tuple[UOp, list[Buffer]]:

test/backend/test_llama_kernels.py

@@ -80,7 +80,7 @@ class TestQuantizeFP8(unittest.TestCase):
   @needs_second_gpu
   def test_multi(self):
     devs = tuple(f"{Device.DEFAULT}:{i}" for i in range(8))
-    x = Tensor.empty(2048*8, 1024, dtype=dtypes.bfloat16, device=devs).uop.multi(0)
+    x = Tensor.empty(2048*8, 1024, dtype=dtypes.bfloat16, device=devs).uop.unshard(0)
     x = Tensor(x, device=devs)
     amax_state = Tensor.full((), 2.0, dtype=dtypes.float32, device=devs).contiguous()
     amax_out = Tensor.zeros((), dtype=dtypes.float32, device=devs).realize()
@@ -131,6 +131,7 @@ class TestFusedQKVRoPE(unittest.TestCase):
       self.assertTrue(v.allclose(v_ref, atol=0, rtol=0).item(), "V forward mismatch")
 
   def test_llama31_8b_backward(self):
+    if not Device[Device.DEFAULT].renderer.target.arch.startswith("gfx950"): self.skipTest("requires CDNA4")
     Tensor.manual_seed(1)
     B, N, H, H_KV, D = self.SHAPE
     PARTIALS = 2

test/backend/test_multitensor.py

@@ -1,6 +1,6 @@
 import unittest, random
 from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variable
-from tinygrad.uop.ops import Ops, UOp
+from tinygrad.uop.ops import Ops, UOp, AxisType
 from tinygrad.helpers import getenv, prod, Context
 from tinygrad.nn.state import get_parameters
 from tinygrad.engine.realize import run_linear, compile_linear
@@ -52,8 +52,10 @@ class TestMultiTensor(unittest.TestCase):
   def test_shard(self):
     X = Tensor.ones(256).contiguous().realize()
     X.shard_(devices_2, 0)
-    for lb in X.uop.src:
-      assert lb.shape == (128,)
+    assert X.uop.src[0].shape == (128,)
+    # the MULTI carries and ends the DEVICE range as its second src
+    assert X.uop.src[1].op is Ops.RANGE and X.uop.src[1].arg[-1] is AxisType.DEVICE
+    assert X.uop.ended_ranges == X.uop.src[1:]
     (X + X).realize()
 
   @unittest.expectedFailure # TODO: fix
@@ -423,6 +425,49 @@ class TestMultiBufferView(unittest.TestCase):
     run_linear(linear, var_vals)
     np.testing.assert_equal(out.numpy(), ref[5].numpy())
 
+@unittest.skipIf(not_support_multi_device(), "need multi")
+class Test2DShard(unittest.TestCase):
+  def setUp(self):
+    self.devices_4 = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
+    self.rng = UOp.range(4, -1, AxisType.DEVICE)
+    self.rng0, self.rng1 = self.rng // 2, self.rng % 2
+
+  def _shard_2d(self, t:Tensor) -> Tensor:
+    u = t.uop.copy_to_device(self.devices_4)._shard(0, self.rng0)._shard(1, self.rng1).unshard((0, 1), (self.rng0, self.rng1))
+    return Tensor(u)
+
+  def test_2d_shard_basic(self):
+    ref = Tensor.arange(16).reshape(4, 4).contiguous().realize()
+    t = self._shard_2d(ref)
+    out = t.contiguous().realize()
+    np.testing.assert_equal(out.numpy(), ref.numpy())
+
+  def test_2d_shard_elementwise(self):
+    ref = Tensor.arange(16).reshape(4, 4).contiguous().realize()
+    t = self._shard_2d(ref)
+    out = (t + 1).contiguous().realize()
+    np.testing.assert_equal(out.numpy(), ref.numpy() + 1)
+
+  def test_2d_shard_sum_all(self):
+    ref = Tensor.arange(16).reshape(4, 4).contiguous().realize()
+    t = self._shard_2d(ref)
+    out = t.sum().contiguous().realize()
+    np.testing.assert_equal(out.numpy(), np.array(ref.numpy().sum()))
+
+  def test_2d_shard_sum_non_sharded_axis(self):
+    ref = Tensor.arange(4*4*2).reshape(4, 4, 2).contiguous().realize()
+    t = self._shard_2d(ref)
+    out = t.sum(axis=2).contiguous().realize()
+    np.testing.assert_equal(out.numpy(), ref.numpy().sum(axis=2))
+
+  def test_2d_shard_matmul(self):
+    a = Tensor.arange(16).reshape(4, 4).contiguous().realize()
+    b = Tensor.arange(16).reshape(4, 4).contiguous().realize()
+    a_s = self._shard_2d(a)
+    b_s = self._shard_2d(b)
+    out = (a_s @ b_s).contiguous().realize()
+    np.testing.assert_equal(out.numpy(), a.numpy() @ b.numpy())
+
 @unittest.skipIf(not_support_multi_device(), "need multi")
 class TestMultiTransformer(unittest.TestCase):
   @needs_second_gpu

test/backend/test_ops.py

@@ -340,11 +340,11 @@ class TestOps(unittest.TestCase):
 
   def test_where(self):
     helper_test_op([], lambda: torch.where(torch.tensor([True, False]), 1, 3).type(torch.int32),
-                   lambda: Tensor([True, False]).where(1, 3), forward_only=True)
+                   lambda: Tensor([True, False]).where(1, 3).clone(), forward_only=True)
     helper_test_op(
       [(100,)],
       lambda x: torch.where(x > 0.5, 4, 2).type(torch.int32),
-      lambda x: (x > 0.5).where(4, 2), forward_only=True)
+      lambda x: (x > 0.5).where(4, 2).clone(), forward_only=True)
 
     for shps in [[(8,),(1,),(1,)], [(10,10),(10,),(10,)], [(100,)]*3, [(10,10)]*3]:
       helper_test_op(
@@ -356,7 +356,7 @@ class TestOps(unittest.TestCase):
     helper_test_op(
       [(5, 5)],
       lambda x: torch.where(x > 0.5, 4, 2).type(torch.int32).permute((1, 0)),
-      lambda x: (x > 0.5).where(4, 2).permute((1, 0)), forward_only=True)
+      lambda x: (x > 0.5).where(4, 2).clone().permute((1, 0)), forward_only=True)
 
   def _test_cmp(self, fxn, reverse=True):
     # test different dtypes
@@ -636,9 +636,9 @@ class TestOps(unittest.TestCase):
         helper_test_op(None, lambda x,y: x%y, forward_only=True, vals=[va, vb])
         helper_test_op(None, lambda x: x%2, forward_only=True, vals=[va])
         helper_test_op(None, lambda x: x%3, forward_only=True, vals=[va])
-        helper_test_op(None, lambda x: x%3.5, forward_only=True, vals=[va])
+        helper_test_op(None, lambda x: x%3.5, lambda x: (x%3.5).clone(), forward_only=True, vals=[va])
         helper_test_op(None, lambda x: 100%x, forward_only=True, vals=[va])
-        helper_test_op(None, lambda x: 100.5%x, forward_only=True, vals=[va])
+        helper_test_op(None, lambda x: 100.5%x, lambda x: (100.5%x).clone(), forward_only=True, vals=[va])
 
   def test_fmod(self):
     a = [-4, 7, 5, 4, -7, 8, -9]
@@ -649,7 +649,7 @@ class TestOps(unittest.TestCase):
         vb = [float(bi) for bi in b] if float_b else b
         helper_test_op(None, lambda x,y: x.fmod(y), forward_only=True, vals=[va, vb])
         helper_test_op(None, lambda x: x.fmod(2), forward_only=True, vals=[va])
-        helper_test_op(None, lambda x: x.fmod(3.5), forward_only=True, vals=[va])
+        helper_test_op(None, lambda x: x.fmod(3.5), lambda x: x.fmod(3.5).clone(), forward_only=True, vals=[va])
 
   def test_mul_naninf(self):
     helper_test_op([(45,65)], lambda x: x*math.inf)
@@ -706,7 +706,7 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x: 0.7**x, vals=[[-2.,-1,0,1,2,3]])
     helper_test_op(None, lambda x: (-2)**x, vals=[[-2.,-1,0,1,2,3]])
     # float to power of int
-    helper_test_op(None, lambda x: 0.7**x, vals=[[-2,-1,0,1,2,3]], forward_only=True)
+    helper_test_op(None, lambda x: 0.7**x, lambda x: (0.7**x).clone(), vals=[[-2,-1,0,1,2,3]], forward_only=True)
 
   @unittest.skipIf(COMPILE_ONLY, "test requires runtime")
   @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, NIRRenderer), "TODO: broken in LVP")
@@ -775,7 +775,7 @@ class TestOps(unittest.TestCase):
   def test_pow_int_base_float_exponent(self):
     for exponent in (0.5, 1.5, 2.0, -1.0, 0.0):
       helper_test_op([], lambda: torch.tensor([1, 2, 3, 4], dtype=torch.int) ** exponent,
-                         lambda: Tensor([1, 2, 3, 4], dtype=dtypes.int32) ** exponent, forward_only=True)
+                         lambda: (Tensor([1, 2, 3, 4], dtype=dtypes.int32) ** exponent).clone(), forward_only=True)
 
   def test_sqrt(self):
     helper_test_op([(45,65)], lambda x: x.sqrt())

test/backend/test_optim.py

@@ -143,12 +143,11 @@ class TestOptim(unittest.TestCase):
 
   @unittest.skipUnless(dtypes.half in Device[Device.DEFAULT].renderer.supported_dtypes(), "need half")
   def test_mixed_precision(self):
-    old_default_float, dtypes.default_float = dtypes.default_float, dtypes.half
+    self.enterContext(Context(DEFAULT_FLOAT=dtypes.half))
     # weight update would overflow without upcasting
     self._test_sgd(10, {'lr': 1e10}, 1e-6, 3e-4)
     self._test_adam(1, {'lr': 1e10}, 1e-4, 1e-4)
     self._test_adamw(1, {'lr': 1e10}, 1e-4, 1e-4)
-    dtypes.default_float = old_default_float
 
   def test_assert_tensor_train(self):
     t = Tensor.ones((1,1))

test/backend/test_randomness.py

@@ -1,7 +1,7 @@
 import unittest, math
 
 from tinygrad import dtypes, Tensor, Device
-from tinygrad.helpers import getenv, DEV
+from tinygrad.helpers import getenv, DEV, Context
 from tinygrad.codegen import to_program
 
 from tinygrad.uop.ops import Ops
@@ -232,16 +232,14 @@ class TestRandomness(unittest.TestCase):
   @given(strat.sampled_from([dtypes.float, dtypes.float16, dtypes.bfloat16]))
   def test_randn_finite(self, default_float):
     if default_float not in Device[Device.DEFAULT].renderer.supported_dtypes(): return
-    old_default_float = dtypes.default_float
     # low precision can result in inf from randn
-    dtypes.default_float = default_float
+    self.enterContext(Context(DEFAULT_FLOAT=default_float))
     t = Tensor.randn(64, 64)
     mx = t.max().numpy().item()
     mn = t.min().numpy().item()
     print(f"testing with {default_float=}")
     assert math.isfinite(mx), mx
     assert math.isfinite(mn), mn
-    dtypes.default_float = old_default_float
 
   def test_random_counter_overflow(self):
     device = Device.DEFAULT

test/backend/test_schedule.py

@@ -363,7 +363,7 @@ class TestCopyFolding(unittest.TestCase):
 
   def test_one_hot_with_copy(self):
     y = Tensor([1, 2, 3]).to("CPU")
-    x = y.one_hot(10)
+    x = y.one_hot(10).int()
     check_schedule(x, 3, filter_sink=False)
 
   @unittest.skip("no longer supported")

test/backend/test_tensor_variable.py

@@ -1,6 +1,6 @@
 import unittest
 import numpy as np
-from tinygrad import Tensor, Variable, dtypes
+from tinygrad import Device, Tensor, Variable, TinyJit, dtypes
 from tinygrad.helpers import CHECK_OOB
 
 class TestTensorVariable(unittest.TestCase):
@@ -18,10 +18,18 @@ class TestTensorVariable(unittest.TestCase):
       self.assertListEqual((vv * t).tolist(), [2, 2, 2])
     except RuntimeError: pass
 
-  # TODO: a Variable PARAM lowers to int32, so a bound value that doesn't fit int32 truncates or fails to bind
-  @unittest.expectedFailure
+  @unittest.skipUnless(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires long support")
   def test_large_range_variable(self):
-    self.assertEqual(Tensor(Variable("b", 0, 2**40).bind(2**35)).item(), 2**35)
+    self.assertEqual(Tensor(Variable("b", 0, 2**40, dtype=dtypes.long).bind(2**35)).clone(Device.DEFAULT).item(), 2**35)
+
+  @unittest.skipUnless(dtypes.long in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires long support")
+  def test_large_range_variable_jit(self):
+    @TinyJit
+    def f(a,b): return (Tensor(a+b).clone(Device.DEFAULT) * 2).realize()
+    for i in range(3):
+      a = Variable("a", 0, 2**10, dtype=dtypes.int).bind(i)
+      b = Variable("b", 0, 2**40, dtype=dtypes.long).bind(2**35)
+      self.assertEqual(f(a,b).item(), (2**35 + i) * 2)
 
   def test_variable_defers_like_a_literal(self):
     vv = Variable("a", 1, 10).bind(2)

test/external/external_benchmark_bert_matmuls.py

@@ -1,7 +1,8 @@
 from tinygrad import Tensor, dtypes
-dtypes.default_float = dtypes.float16
 from tinygrad.dtype import to_dtype
-from tinygrad.helpers import getenv
+from tinygrad.helpers import getenv, Context
+
+Context(DEFAULT_FLOAT=dtypes.float16).__enter__()
 
 if __name__ == "__main__":
   # matmuls in bert layers

test/external/external_benchmark_bert_softmax.py

@@ -1,9 +1,10 @@
 from tinygrad import Tensor, dtypes, GlobalCounters
-dtypes.default_float = dtypes.float16
 from tinygrad.dtype import to_dtype
-from tinygrad.helpers import getenv
+from tinygrad.helpers import getenv, Context
 from test.backend.test_softmax_fusion import single_kernel_softmax
 
+Context(DEFAULT_FLOAT=dtypes.float16).__enter__()
+
 if __name__ == "__main__":
   # softmax in bert layers
   BS = getenv("BS", 96//6)

test/external/external_benchmark_op_conv.py

@@ -13,9 +13,9 @@ from tinygrad.dtype import Invalid
 
 def vision_conv_143():
   c0 = UOp.param(0, dtypes.half, shape=(16, 1024, 4))
-  c2 = UOp.range(32, 3, AxisType.LOOP)
-  c5 = UOp.range(128, 4, AxisType.LOOP)
-  c8 = UOp.range(16, 2, AxisType.LOOP)
+  c2 = UOp.range(32, 3)
+  c5 = UOp.range(128, 4)
+  c8 = UOp.range(16, 2)
   c16 = UOp.range(7, 0, AxisType.REDUCE)
   c17 = c8*2+c16
   c24 = ((c17<3)!=True)&(c17<35)
@@ -39,9 +39,9 @@ def vision_conv_143():
 
 def vision_conv_153():
   c0 = UOp.param(0, dtypes.half, shape=(8, 1024, 4))
-  c2 = UOp.range(16, 3, AxisType.LOOP)
-  c5 = UOp.range(256, 4, AxisType.LOOP)
-  c8 = UOp.range(8, 2, AxisType.LOOP)
+  c2 = UOp.range(16, 3)
+  c5 = UOp.range(256, 4)
+  c8 = UOp.range(8, 2)
   c16 = UOp.range(7, 0, AxisType.REDUCE)
   c17 = c8*2+c16
   c24 = ((c17<3)!=True)&(c17<19)
@@ -65,7 +65,7 @@ def vision_conv_153():
 
 def dm_conv_172():
   c0 = UOp.param(0, dtypes.half, shape=(1, 240, 4))
-  c2 = UOp.range(960, 4, AxisType.LOOP)
+  c2 = UOp.range(960, 4)
   c5 = UOp.param(1, dtypes.half, shape=(8, 384, 4))
   c7 = UOp.range(32, 0, AxisType.REDUCE)
   c10 = UOp.range(4, 1, AxisType.REDUCE)

test/external/external_test_gpu_crash.py

@@ -5,8 +5,9 @@ These tests intentionally cause GPU faults to verify error handling.
 Run with: DEV=AMD python -m pytest test/external/external_test_gpu_crash.py -v
 """
 import unittest, re, importlib
-from tinygrad.device import Device, TinyELF
-from tinygrad.helpers import Target
+from tinygrad import Device, UOp, dtypes
+from tinygrad.engine.realize import run_linear
+from tinygrad.uop.ops import Ops, KernelInfo
 from tinygrad.renderer.amd.dsl import s, v, Inst, NULL
 
 RDNA3_CDNA3_MAP = {"v_mov_b32_e32": "v_mov_b32_e32", "s_mov_b32": "s_mov_b32", "s_waitcnt": "s_waitcnt", "s_endpgm": "s_endpgm",
@@ -14,20 +15,11 @@ RDNA3_CDNA3_MAP = {"v_mov_b32_e32": "v_mov_b32_e32", "s_mov_b32": "s_mov_b32", "
                    "global_atomic_add_u32": "global_atomic_add", "flat_load_b32": "flat_load_dword",
                    "flat_store_b32": "flat_store_dword", "flat_atomic_add_u32": "flat_atomic_add", "s_load_b32": "s_load_dword"}
 
-def assemble(code:str, name:str="test", is_cdna:bool=False) -> str:
-  kd = {"next_free_vgpr": 8, "next_free_sgpr": 8, "user_sgpr_kernarg_segment_ptr": 1, "kernarg_size": 8}
-  if is_cdna: kd["accum_offset"] = 8
-  else: kd["wavefront_size32"] = 1
-  return f".text\n.globl {name}\n.p2align 8\n.type {name},@function\n{name}:\n{code}\n.rodata\n.p2align 6\n.amdhsa_kernel {name}\n" + \
-         "\n".join(f".amdhsa_{k} {v}" for k,v in kd.items()) + "\n.end_amdhsa_kernel"
-
 @unittest.skipIf(Device.DEFAULT != "AMD", "AMD required")
 class TestGPUCrash(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
-    from tinygrad.runtime.support.compiler_amd import HIPCompiler
     cls.dev = Device["AMD"]
-    cls.compiler = HIPCompiler(cls.dev.arch)
     cls.is_cdna = cls.dev.target[0] < 10
     ins = importlib.import_module('tinygrad.runtime.autogen.amd.' + ('cdna' if cls.is_cdna else 'rdna3') + '.ins')
     for rdna3_name, cdna3_name in RDNA3_CDNA3_MAP.items():
@@ -42,14 +34,11 @@ class TestGPUCrash(unittest.TestCase):
     except Exception:
       self.fail("Device not working before test")
 
-  def _run(self, code: str):
-    prg = self.dev.runtime(TinyELF(self.compiler.compile(assemble(code, is_cdna=self.is_cdna)), "test",
-                                   Target("AMD", arch=self.dev.arch), ()))
-    prg(self.dev.allocator.alloc(64), global_size=(1,1,1), local_size=(1,1,1), wait=True)
-
   def _run_insts(self, insts: list[Inst]):
-    from test.amd.disasm import disasm
-    self._run("\n".join(disasm(i) for i in insts))
+    buf = UOp.new_buffer("AMD", 64, dtypes.uint8)
+    sink = UOp.sink(UOp.param(0, dtypes.uint8, (64,), device="AMD"), UOp.special(1, "lidx0"), arg=KernelInfo("test"))
+    prg = UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=i) for i in insts))))
+    run_linear(UOp(Ops.LINEAR, src=(prg.call(buf),)), wait=True)
 
   def _assert_gpu_fault(self, func):
     """Assert that func raises a RuntimeError indicating a GPU fault (not a setup error)."""

test/external/external_test_jit_on_models.py

@@ -3,6 +3,7 @@ import unittest
 import numpy as np
 from tinygrad import Tensor, dtypes
 from tinygrad.engine.jit import TinyJit
+from tinygrad.helpers import Context
 from test.helpers import derandomize_model
 
 from examples.llama import Transformer
@@ -14,8 +15,7 @@ def helper_test_jitted_correctness(gen, train, train_jit):
 
 class TestJittedModels(unittest.TestCase):
   def test_jitted_tiny_llama(self):
-    old_float = dtypes.default_float
-    dtypes.default_float = dtypes.float16
+    self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16))
 
     args_tiny = {"dim": 1024, "hidden_dim": 1024, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-05, "vocab_size": 1000}
     model = Transformer(**args_tiny)
@@ -25,7 +25,6 @@ class TestJittedModels(unittest.TestCase):
     @TinyJit
     def test_jit(t): return model(t, 0).realize()
     helper_test_jitted_correctness(lambda: (Tensor([[1,]]),), test, test_jit)
-    dtypes.default_float = old_float
 
   def test_jitted_stable_diffusion(self):
     from examples.stable_diffusion import UNetModel, unet_params

test/external/external_test_llm_opencode.py

@@ -0,0 +1,90 @@
+"""Real-model OpenCode regression.
+
+Run against an existing server:
+  RUN_LLM_OPENCODE_REGRESSION=1 LLM_BASE_URL=http://127.0.0.1:8000/v1 \
+    python -m pytest test/external/external_test_llm_opencode.py -v
+
+Or set LLM_GGUF and let the test start the tinygrad server.
+"""
+from __future__ import annotations
+
+import json, os, pathlib, re, shutil, socket, subprocess, sys, tempfile, time, unittest, urllib.request
+
+RUN_REGRESSION = os.getenv("RUN_LLM_OPENCODE_REGRESSION") == "1"
+
+def _server_ready(base_url:str) -> bool:
+  try:
+    with urllib.request.urlopen(base_url.rstrip("/") + "/models", timeout=1) as response: return response.status == 200
+  except OSError: return False
+
+@unittest.skipUnless(RUN_REGRESSION, "set RUN_LLM_OPENCODE_REGRESSION=1 to run the OpenCode regression")
+class TestLLMOpenCode(unittest.TestCase):
+  server:subprocess.Popen|None = None
+  server_log:tempfile._TemporaryFileWrapper|None = None
+
+  @classmethod
+  def setUpClass(cls):
+    if shutil.which("opencode") is None: raise unittest.SkipTest("opencode is not installed")
+    if base_url := os.getenv("LLM_BASE_URL"):
+      cls.base_url = base_url.rstrip("/")
+      if not cls.base_url.endswith("/v1"): cls.base_url += "/v1"
+      if not _server_ready(cls.base_url): raise RuntimeError(f"LLM server is not responding at {cls.base_url}")
+      return
+
+    model = pathlib.Path(os.environ["LLM_GGUF"])
+    with socket.socket() as sock:
+      sock.bind(("127.0.0.1", 0))
+      port = sock.getsockname()[1]
+    cls.base_url = f"http://127.0.0.1:{port}/v1"
+    cls.server_log = tempfile.NamedTemporaryFile(mode="w+", prefix="tinygrad-llm-")
+    cls.server = subprocess.Popen(
+      [sys.executable, "-m", "tinygrad.llm", "--model", str(model), "--serve", str(port), "--max_context", "262144"],
+      stdout=cls.server_log, stderr=subprocess.STDOUT, start_new_session=True)
+    deadline = time.monotonic() + 180
+    while time.monotonic() < deadline and cls.server.poll() is None:
+      if _server_ready(cls.base_url): return
+      time.sleep(0.25)
+    cls.server_log.seek(0)
+    raise RuntimeError(f"LLM server failed to start:\n{cls.server_log.read()[-8000:]}")
+
+  @classmethod
+  def tearDownClass(cls):
+    if cls.server is not None:
+      cls.server.terminate()
+      try: cls.server.wait(timeout=10)
+      except subprocess.TimeoutExpired:
+        cls.server.kill()
+        cls.server.wait(timeout=10)
+    if cls.server_log is not None: cls.server_log.close()
+
+  def run_opencode(self, prompt:str, cwd:pathlib.Path) -> str:
+    config = cwd / "opencode.json"
+    config.write_text(json.dumps({
+      "$schema": "https://opencode.ai/config.json", "permission": {"*": "allow"}, "formatter": False, "lsp": False,
+      "provider": {"regression": {"npm": "@ai-sdk/openai-compatible", "options": {"baseURL": self.base_url},
+                                  "models": {"tinygrad": {"name": "tinygrad"}}}},
+    }))
+    env = os.environ | {"OPENCODE_CONFIG": str(config)}
+    result = subprocess.run(
+      ["opencode", "run", "--pure", "--auto", "--dir", str(cwd), "-m", "regression/tinygrad", prompt],
+      cwd=cwd, env=env, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, timeout=120)
+    self.assertEqual(result.returncode, 0, result.stdout)
+    return re.sub(r"\x1b\[[0-?]*[ -/]*[@-~]", "", result.stdout)
+
+  def test_read_tool(self):
+    with tempfile.TemporaryDirectory() as directory:
+      cwd, marker = pathlib.Path(directory), "tinygrad-opencode-regression-7f3a91c2"
+      (cwd / "exact.txt").write_text(marker + "\n")
+      output = self.run_opencode("Read exact.txt with a tool and reply with its exact contents, with no other text.", cwd)
+      self.assertRegex(output, r"(?im)^\s*(?:→|>)\s*Read\s+exact\.txt\s*$")
+      self.assertIn(marker, output)
+
+  def test_shell_tool(self):
+    with tempfile.TemporaryDirectory() as directory:
+      cwd = pathlib.Path(directory)
+      output = self.run_opencode(
+        "Use the shell tool to run `printf tinygrad-shell-regression > shell-regression.txt`, then report completion.", cwd)
+      self.assertRegex(output, r"(?im)^\s*(?:\$|→|>)\s*.*printf\s+tinygrad-shell-regression")
+      self.assertEqual((cwd / "shell-regression.txt").read_text(), "tinygrad-shell-regression")
+
+if __name__ == "__main__": unittest.main()

test/external/mlperf_stable_diffusion/external_test_eval.py

@@ -51,7 +51,7 @@ class TestEval(unittest.TestCase):
       vae = AutoencoderKL()
       for p in get_parameters(vae): p.to_(GPUS).realize()
       x = Tensor.zeros(48,4,64,64).contiguous().to(GPUS).realize()
-      x.uop = x.uop.multi(0)
+      x.uop = x.uop.unshard(0)
       for _ in range(2): vae_decode(x, vae)
 
 if __name__=="__main__":

test/mockgpu/amd/amdgpu.py

@@ -131,7 +131,8 @@ class PM4Executor(AMDQueue):
     _ = self._next_dword() # ev
 
     ptr = to_mv(self.gpu.translate_addr(addr_lo + (addr_hi << 32)), 8)
-    if mem_data_sel == 1 or mem_data_sel == 2: ptr.cast('Q')[0] = val
+    if mem_data_sel == 1: ptr.cast('I')[0] = val & 0xffffffff
+    elif mem_data_sel == 2: ptr.cast('Q')[0] = val
     elif mem_data_sel == 3:
       if mem_event_type == CACHE_FLUSH_AND_INV_TS_EVENT: ptr.cast('Q')[0] = int(time.perf_counter() * 1e8)
       else: raise RuntimeError(f"Unknown {mem_data_sel=} {mem_event_type=}")

test/mockgpu/amd/emu.py

@@ -51,7 +51,7 @@ class _MXCSRContext:
     if lib is None or not hasattr(self, '_saved'): return
     lib.set_fpcr(self._saved)
 
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.device import Buffer, BufferSpec, Device
 from tinygrad.runtime.autogen import hsa
@@ -446,7 +446,7 @@ class _Ctx:
     """Create a lane range UOp with unique axis ID."""
     if n is None: n = self.wave_size
     self._axis_id += 1
-    return UOp.range(n, self._axis_id, AxisType.LOOP, dtype=dtypes.int)
+    return UOp.range(n, self._axis_id, dtype=dtypes.int)
 
   def unroll_lanes(self, get_lane_bit, exec_mask: UOp, apply_exec: bool = True) -> UOp:
     """Combine lane bits into a mask using RANGE+REDUCE (32-bit for RDNA, 64-bit for CDNA)."""

test/mockgpu/nv/nvgpu.py

@@ -91,7 +91,7 @@ class GPFIFO:
     args_cnt, vals_cnt = const0[80], const0[81]
     args_addr = qmd.constant_buffer_addr_lower_0 + (qmd.constant_buffer_addr_upper_0 << 32) + 0x160
     args = to_mv(args_addr, args_cnt*8).cast('Q')
-    vals = to_mv(args_addr + args_cnt*8, vals_cnt*4).cast('I')
+    vals = to_mv(args_addr + args_cnt*8, vals_cnt*8).cast('Q')
     cargs = [ctypes.cast(args[i], ctypes.c_void_p) for i in range(args_cnt)] + [ctypes.cast(vals[i], ctypes.c_void_p) for i in range(vals_cnt)]
     gx, gy, gz = qmd.cta_raster_width, qmd.cta_raster_height, qmd.cta_raster_depth
     lx, ly, lz = qmd.cta_thread_dimension0, qmd.cta_thread_dimension1, qmd.cta_thread_dimension2

test/null/test_dtype_spec.py

@@ -2,7 +2,7 @@ import unittest, math, struct, operator
 from tinygrad import Tensor, Device
 from tinygrad.dtype import DTYPES_DICT, dtypes, Invalid, truncate, float_to_fp16, float_to_bf16, _to_np_dtype, least_upper_dtype, least_upper_float
 
-from tinygrad.helpers import getenv
+from tinygrad.helpers import getenv, Context
 from hypothesis import given, settings, strategies as strat
 import numpy as np
 import torch
@@ -58,8 +58,8 @@ class TestHelpers(unittest.TestCase):
   def test_from_py(self):
     assert dtypes.from_py(True) == dtypes.bool
     assert dtypes.from_py(Invalid) == dtypes.bool
-    assert dtypes.from_py(2) == dtypes.default_int
-    assert dtypes.from_py(3.0) == dtypes.default_float
+    assert dtypes.from_py(2) == dtypes.weakint
+    assert dtypes.from_py(3.0) == dtypes.weakfloat
     assert dtypes.from_py([]) == dtypes.default_float
     assert dtypes.from_py(()) == dtypes.default_float
     assert dtypes.from_py([True]) == dtypes.bool
@@ -245,19 +245,14 @@ class TestTypePromotion(unittest.TestCase):
     assert least_upper_dtype(dtypes.weakfloat, dtypes.float64) == dtypes.float64
 
 class TestTypeSpec(unittest.TestCase):
-  def setUp(self):
-    self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float
-  def tearDown(self):
-    dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float
-
   def test_set_dtype_default(self):
     for default_int in [dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64]:
-      dtypes.default_int = default_int
-      assert dtypes.default_int == default_int
+      with Context(DEFAULT_INT=default_int):
+        assert dtypes.default_int == default_int
 
     for default_float in [*dtypes.fp8s, dtypes.float16, dtypes.bfloat16, dtypes.float32, dtypes.float64]:
-      dtypes.default_float = default_float
-      assert dtypes.default_float == default_float
+      with Context(DEFAULT_FLOAT=default_float):
+        assert dtypes.default_float == default_float
 
   @given(strat.sampled_from(core_dtypes), strat.sampled_from([operator.gt, operator.ge, operator.le, operator.lt, operator.eq, operator.ne]))
   def test_bool_ops(self, dtype, op):
@@ -265,7 +260,7 @@ class TestTypeSpec(unittest.TestCase):
 
   @given(strat.sampled_from(core_dtypes), strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
   def test_functions_return_index(self, dtype, default_int, default_float):
-    dtypes.default_int, dtypes.default_float = default_int, default_float
+    self.enterContext(Context(DEFAULT_INT=default_int, DEFAULT_FLOAT=default_float))
     assert Tensor([0, 1], dtype=dtype).argmax().dtype == dtypes.int32
     assert Tensor([0, 1], dtype=dtype).argmin().dtype == dtypes.int32
     assert Tensor([0, 1], dtype=dtype).multinomial().dtype == dtypes.int32
@@ -285,7 +280,7 @@ class TestTypeSpec(unittest.TestCase):
 
   @given(strat.sampled_from(dtype_floats), strat.sampled_from(dtype_floats))
   def test_attention_returns_same_dtype(self, data_dtype, default_float):
-    dtypes.default_float = default_float
+    self.enterContext(Context(DEFAULT_FLOAT=default_float))
     query = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
     key = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
     value = Tensor.rand(32, 8, 128, 64, dtype=data_dtype)
@@ -296,32 +291,27 @@ class TestTypeSpec(unittest.TestCase):
     assert query.scaled_dot_product_attention(key, value, attn_mask=mask).dtype == data_dtype
 
 class TestAutoCastType(unittest.TestCase):
-  def setUp(self):
-    self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float
-  def tearDown(self):
-    dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float
-
   @given(strat.sampled_from(dtype_floats), strat.sampled_from(dtype_floats))
   def test_least_upper_float_input_is_float(self, input_dtype, default_float):
-    dtypes.default_float = default_float
+    self.enterContext(Context(DEFAULT_FLOAT=default_float))
     self.assertEqual(least_upper_float(input_dtype), input_dtype)
 
   @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
   def test_least_upper_float_input_is_int(self, input_dtype, default_float):
-    dtypes.default_float = default_float
+    self.enterContext(Context(DEFAULT_FLOAT=default_float))
     self.assertEqual(least_upper_float(input_dtype), default_float)
 
   @given(strat.sampled_from(core_dtypes))
   def test_broadcast_scalar(self, dt):
-    assert (Tensor.ones(4, 4, dtype=dt) + 2.3).dtype == (dt if dtypes.is_float(dt) else dtypes.default_float)
-    assert (Tensor.ones(4, 4, dtype=dt) + 2).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int)
+    assert (Tensor.ones(4, 4, dtype=dt) + 2.3).dtype == (dt if dtypes.is_float(dt) else dtypes.weakfloat)
+    assert (Tensor.ones(4, 4, dtype=dt) + 2).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.weakint)
     assert (Tensor.ones(4, 4, dtype=dt) + True).dtype == dt
 
   @given(strat.sampled_from(core_dtypes))
   def test_pad_scalar(self, dt):
     t = Tensor.ones(4, dtype=dt)
-    assert t.pad(((1, 1),), value=2.3).dtype == (dt if dtypes.is_float(dt) else dtypes.default_float)
-    assert t.pad(((1, 1),), value=2).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int)
+    assert t.pad(((1, 1),), value=2.3).dtype == (dt if dtypes.is_float(dt) else dtypes.weakfloat)
+    assert t.pad(((1, 1),), value=2).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.weakint)
     assert t.pad(((1, 1),), value=True).dtype == dt
 
   @given(strat.sampled_from(core_dtypes))
@@ -332,7 +322,7 @@ class TestAutoCastType(unittest.TestCase):
 
   @given(strat.sampled_from(dtype_floats))
   def test_int_div_int(self, default_float):
-    dtypes.default_float = default_float
+    self.enterContext(Context(DEFAULT_FLOAT=default_float))
     self.assertEqual(Tensor([1]).div(Tensor([2])).dtype, default_float)
 
   def test_sum(self):
@@ -420,16 +410,16 @@ class TestAutoCastType(unittest.TestCase):
   @given(strat.sampled_from(core_dtypes))
   def test_where_one_scalar(self, dt):
     t = Tensor(2, dtype=dt)
-    self.check_where_alternate_input_other(t, 3.2, (dt if dtypes.is_float(dt) else dtypes.default_float))
-    self.check_where_alternate_input_other(t, 3, (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int))
+    self.check_where_alternate_input_other(t, 3.2, (dt if dtypes.is_float(dt) else dtypes.weakfloat))
+    self.check_where_alternate_input_other(t, 3, (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.weakint))
     self.check_where_alternate_input_other(t, True, dt)
 
   def test_where_two_scalars(self):
-    self.check_where_alternate_input_other(3.1, 3.2, dtypes.default_float)
-    self.check_where_alternate_input_other(3.1, 3, dtypes.default_float)
-    self.check_where_alternate_input_other(3.1, True, dtypes.default_float)
-    self.check_where_alternate_input_other(3, 2, dtypes.default_int)
-    self.check_where_alternate_input_other(3, True, dtypes.default_int)
+    self.check_where_alternate_input_other(3.1, 3.2, dtypes.weakfloat)
+    self.check_where_alternate_input_other(3.1, 3, dtypes.weakfloat)
+    self.check_where_alternate_input_other(3.1, True, dtypes.weakfloat)
+    self.check_where_alternate_input_other(3, 2, dtypes.weakint)
+    self.check_where_alternate_input_other(3, True, dtypes.weakint)
 
   def test_where_non_bool_cond_raises(self):
     with self.assertRaises(RuntimeError): Tensor([1, 0, 2]).where(1, 0)
@@ -441,8 +431,8 @@ class TestAutoCastType(unittest.TestCase):
 
   @given(strat.sampled_from(core_dtypes))
   def test_maximum_const(self, dt):
-    assert Tensor([1, 2], dtype=dt).maximum(3.1).dtype == (dt if dtypes.is_float(dt) else dtypes.default_float)
-    assert Tensor([1, 2], dtype=dt).maximum(3).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.default_int)
+    assert Tensor([1, 2], dtype=dt).maximum(3.1).dtype == (dt if dtypes.is_float(dt) else dtypes.weakfloat)
+    assert Tensor([1, 2], dtype=dt).maximum(3).dtype == (dt if dtypes.is_float(dt) or dtypes.is_int(dt) else dtypes.weakint)
     assert Tensor([1, 2], dtype=dt).maximum(True).dtype == dt
 
   def test_div(self):
@@ -453,7 +443,7 @@ class TestAutoCastType(unittest.TestCase):
 
   def test_div_const(self):
     assert (Tensor([1, 2], dtype=dtypes.int32) / 2).dtype == dtypes.default_float
-    assert (Tensor([1, 2], dtype=dtypes.int32) / 2.0).dtype == dtypes.default_float
+    assert (Tensor([1, 2], dtype=dtypes.int32) / 2.0).dtype == dtypes.weakfloat
     assert (Tensor([1, 2], dtype=dtypes.float16) / 2).dtype == dtypes.float16
     assert (Tensor([1, 2], dtype=dtypes.float16) / 2.0).dtype == dtypes.float16
 

test/null/test_hcq_iface.py

@@ -1,9 +1,26 @@
 import unittest, array, time
 from tinygrad.helpers import mv_address
 from tinygrad.runtime.support.hcq import MMIOInterface
+from tinygrad.runtime.support.memory import VirtMapping
+from tinygrad.runtime.support.system import PCIIfaceBase
 from tinygrad.runtime.support.usb import USBMMIOInterface
 from test.mockgpu.usb import MockUSB
 
+class TestPCIIface(unittest.TestCase):
+  def test_sysmem_mapping_respects_uncached(self):
+    class MM:
+      def alloc_vaddr(self, size, align): return 0x10000
+      def map_range(self, vaddr, size, paddrs, aspace, uncached=False, snooped=False):
+        return VirtMapping(vaddr, size, paddrs, aspace, uncached, snooped)
+    class PCI:
+      def bar_info(self, bar): return 0, 256 << 20
+      def alloc_sysmem(self, size, **kwargs): return memoryview(bytearray(size)), [0x20000]
+    iface = PCIIfaceBase.__new__(PCIIfaceBase)
+    iface.dev, iface.vram_bar, iface.pci_dev = None, 0, PCI()
+    iface.dev_impl = type("DevImpl", (), {"mm": MM()})()
+    for uncached in (False, True):
+      with self.subTest(uncached=uncached): self.assertEqual(iface.alloc(4096, host=True, uncached=uncached).meta.mapping.uncached, uncached)
+
 class TestHCQIface(unittest.TestCase):
   def setUp(self):
     self.size = 4 << 10

test/null/test_linearizer_failures.py

@@ -8,8 +8,8 @@ from tinygrad.codegen import to_program
 class TestLinearizerFailures(unittest.TestCase):
   def test_fail_1(self):
     c0 = UOp.param(0, dtypes.float, (64,))
-    c1 = UOp.range(UOp.const(dtypes.weakint, 2), 1, AxisType.LOOP)
-    c2 = UOp.range(UOp.const(dtypes.weakint, 32), 2, AxisType.LOOP)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 2), 1, AxisType.WEAK)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 32), 2, AxisType.WEAK)
     c3 = ((c1*UOp.const(dtypes.weakint, 32))+c2)
     c4 = UOp.param(1, dtypes.float, (163840,))
     c5 = UOp.range(UOp.const(dtypes.weakint, 2560), 0, AxisType.REDUCE)

test/null/test_llm_server.py

@@ -95,6 +95,17 @@ class TestLLMServer(unittest.TestCase):
     self.assertGreater(len(chunks), 0)
     self.assertEqual(chunks[-1].choices[0].finish_reason, "stop")
 
+  def test_chat_template_kwargs(self):
+    import jinja2
+    with patch.object(self.server, "template", jinja2.Template(
+        "{% for m in messages %}{% if m.role == 'assistant' and preserve_thinking %}<think>{{ m.reasoning_content }}</think>"
+        "{% endif %}{{ m.content }}|{% endfor %}")):
+      list(self.client.chat.completions.create(model="test", messages=[
+        {"role":"user", "content":"first"}, {"role":"assistant", "content":"answer", "reasoning_content":"reason"},
+        {"role":"user", "content":"next"}], stream=True,
+        extra_body={"chat_template_kwargs":{"preserve_thinking":True, "messages":[]}}))
+    self.mock_tok.encode.assert_called_with("first|<think>reason</think>answer|next|")
+
   def test_content_is_streamed(self):
     stream = self.client.chat.completions.create(
       model="test",
@@ -109,6 +120,28 @@ class TestLLMServer(unittest.TestCase):
 
     self.assertGreater(len(contents), 0)
 
+  def test_interrupted_stream_logs_tokens(self):
+    with patch.object(self.mock_model, "generate", side_effect=lambda ids, **kwargs: iter([300, 301, 999])), \
+         patch("tinygrad.llm.serve.stderr_log") as log, patch("tinygrad.llm.serve.colored", side_effect=lambda text, color: text) as color:
+      stream = self.server.RequestHandlerClass.run_model(Mock(server=self.server), [200, 201, 202], "test")
+      next(stream)
+      next(stream)
+      stream.close()
+    interrupt = log.call_args.args[0]
+    self.assertFalse(interrupt.startswith("\n"))
+    self.assertTrue(interrupt.endswith("\n"))
+    self.assertIn("gen:", interrupt)
+    self.assertIn("out:    1", interrupt)
+    self.assertTrue(any(args[0].startswith("total:") and args[1] == "red" for args, _ in color.call_args_list))
+
+  def test_stream_disconnect_closes_source(self):
+    from tinygrad.viz.serve import HTTPRequestHandler
+    source, handler = Mock(), Mock()
+    source.__iter__ = Mock(return_value=iter([{}]))
+    handler.wfile.write.side_effect = BrokenPipeError
+    HTTPRequestHandler.stream_json(handler, source)
+    source.close.assert_called_once()
+
   def test_non_streaming(self):
     resp = self.client.chat.completions.create(
       model="test-model",

test/null/test_multitensor.py

@@ -1,5 +1,5 @@
 import gc, unittest
-from tinygrad import Tensor, GlobalCounters, dtypes
+from tinygrad import Tensor, UOp, GlobalCounters, dtypes
 from tinygrad.engine.jit import TinyJit
 from tinygrad.helpers import Context
 
@@ -181,7 +181,7 @@ class TestMultiScalarALU(unittest.TestCase):
     @functools.cache
     def _fxn(x_p, device):
       t = Tensor(x_p, device=device)
-      inner = Tensor(t.uop.src[0]) if t.uop.op is Ops.MULTI else t
+      inner = Tensor(t.uop.src[0]) if t.uop.op is Ops.UNSHARD else t
       return (inner.sum(),)
     param = x.as_param(0)
     fxn = _fxn(param.uop, x.device)
@@ -218,5 +218,10 @@ class TestMultiAxis(unittest.TestCase):
     self.assertEqual(e.uop.axis, 0)
     self.assertTrue(e.uop.has_buffer_identity())
 
+  def test_symbolic_reshape_shard_axis(self):
+    rows = UOp.variable("rows", 1, 4).bind(3)
+    x = Tensor.empty(4, 2).shard(("NULL:1", "NULL:2"), axis=1)[:rows]
+    self.assertEqual(x.reshape(rows, 1, 2).uop.axis, 2)
+
 if __name__ == '__main__':
   unittest.main()

test/null/test_real_world.py

@@ -47,12 +47,8 @@ class TestRealWorld(unittest.TestCase):
     gc.collect()
     global global_mem_used
     global_mem_used = GlobalCounters.mem_used
-    self.old_float = dtypes.default_float
     np.random.seed(2002)
 
-  def tearDown(self):
-    dtypes.default_float = self.old_float
-
   @slow
   @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
   def test_stable_diffusion(self):
@@ -81,7 +77,7 @@ class TestRealWorld(unittest.TestCase):
 
   @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
   def test_llama(self):
-    dtypes.default_float = dtypes.float16
+    self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16))
 
     args_tiny = {"dim": 1024, "hidden_dim": 2048, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-05, "vocab_size": 1000}
     model = LLaMaTransformer(**args_tiny)
@@ -93,7 +89,7 @@ class TestRealWorld(unittest.TestCase):
 
   @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
   def test_gpt2(self):
-    dtypes.default_float = dtypes.float16
+    self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16))
 
     args_tiny = {"dim": 1024, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-5, "vocab_size": 1000}
     model = GPT2Transformer(**args_tiny)
@@ -150,7 +146,7 @@ class TestRealWorld(unittest.TestCase):
 
   @unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
   def test_train_cifar_hyp(self):
-    dtypes.default_float = dtypes.float16
+    self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16))
     with Context(TRAINING=1):
       model = SpeedyResNet(Tensor.ones((12,3,2,2)))
       optimizer = optim.SGD(get_parameters(model), lr=0.01, momentum=hyp['opt']['momentum'], nesterov=True, weight_decay=hyp['opt']['bias_decay'])

test/null/test_schedule.py

@@ -600,18 +600,15 @@ class TestSchedule(unittest.TestCase):
     p = p.pad(((1, 0), ))
     p = p.repeat([2])
     # TODO: this should be 3 if fix store hazard worked correctly
-    # NOTE: pad now always has an explicit fill mask (internal PAD is Invalid-filled), which materializes here
-    check_schedule(p, 5)
+    check_schedule(p, 4)
 
   def test_conv2d(self, allowed=4, dtype=dtypes.float):
-    old_default_float, dtypes.default_float = dtypes.default_float, dtype
-    dtypes.default_float = dtype
+    self.enterContext(Context(DEFAULT_FLOAT=dtype))
     Tensor.manual_seed(0)
     BS, CIN = 2, 3
     img = Tensor.randn(BS, CIN, 64, 64).realize()
     w = Tensor.uniform(16, CIN, 3, 3).realize()
     ret = Tensor.conv2d(img, w).relu().mean().backward()
-    dtypes.default_float = old_default_float
     linear, var_vals = Tensor.linear_with_vars(ret, img.grad, w.grad)
     cnt = len([call for call in linear.src if call.src[0].op is Ops.SINK])
     assert cnt == allowed, f"expected {allowed} kernels, got {cnt}"

test/null/test_schedule_cache.py

@@ -1,10 +1,12 @@
-import unittest
-from tinygrad import Tensor, Variable, Context
+import os, unittest
+from unittest.mock import patch
+from tinygrad import Tensor, Variable, Context, UOp
+from tinygrad.callify import transform_to_call
 from tinygrad.helpers import cpu_events
-from tinygrad.schedule import schedule_cache
+from tinygrad.schedule import lower_sink_to_linear, schedule_cache
 
 def schedule_one():
-  Tensor([1]).schedule_linear()
+  (Tensor.empty(1) + 1).schedule_linear()
 
 class TestScheduleCache(unittest.TestCase):
   def test_bound_variable_var_vals(self):
@@ -37,5 +39,22 @@ class TestScheduleCache(unittest.TestCase):
       num_events_cache = len(cpu_events)
     self.assertLess(num_events_cache, num_events_no_cache)
 
+  def test_disk_schedule_cache(self):
+    function = transform_to_call(UOp.sink((Tensor.empty(1) + 1).uop))[0].src[0]
+    schedule_cache.clear()
+    with patch.dict(os.environ, {"DISK_SCACHE":"1"}), \
+         patch("tinygrad.schedule.diskcache_get", return_value=None), \
+         patch("tinygrad.schedule.diskcache_put") as cache_put:
+      lower_sink_to_linear(function)
+      cached = cache_put.call_args.args[2]
+
+    schedule_cache.clear()
+    with patch.dict(os.environ, {"DISK_SCACHE":"1"}), \
+         patch("tinygrad.schedule.diskcache_get", return_value=cached) as cache_get, \
+         patch("tinygrad.schedule.diskcache_put") as cache_put:
+      self.assertIs(lower_sink_to_linear(function), cached)
+      cache_get.assert_called_once()
+      cache_put.assert_not_called()
+
 if __name__ == "__main__":
   unittest.main()

test/null/test_tensor.py

@@ -181,7 +181,7 @@ class TestTensorPad(unittest.TestCase):
     t = Tensor.arange(9).reshape(1, 1, 3, 3)
     self.assertEqual(t.dtype, dtypes.int)
     r = t.pad((1, 2, 0, -1), value=-float('inf'))
-    self.assertEqual(r.dtype, dtypes.float)
+    self.assertEqual(r.dtype, dtypes.weakfloat)
     self.assertEqual(r.shape, (1, 1, 2, 6))
 
 class TestTensorDeviceMismatch(unittest.TestCase):

test/null/test_uop_graph.py

@@ -424,8 +424,8 @@ class TestUOpGraph(unittest.TestCase):
     # mnist indexing with split reduceop
     # Make sure we are not doign math on the loaded index, which would promote it to long
     c0 = UOp.param(0, dtypes.uchar, (128000,))
-    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.LOOP)
-    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.LOOP)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.WEAK)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.WEAK)
     c3 = UOp.param(1, dtypes.int, (512,))
     c4 = c3.index(c1)
     c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)
@@ -441,8 +441,8 @@ class TestUOpGraph(unittest.TestCase):
   def test_load_idx_no_math_on_loaded(self):
     # test the (x+y)<c pattern where x has loads - we shouldn't do math on loaded indices
     c0 = UOp.param(0, dtypes.uchar, (128000,))
-    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.LOOP)
-    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.LOOP)
+    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 1, AxisType.WEAK)
+    c2 = UOp.range(UOp.const(dtypes.weakint, 250), 2, AxisType.WEAK)
     c3 = UOp.param(1, dtypes.int, (512,))
     c4 = c3.index(c1)  # c4 is a load
     c5 = UOp.range(UOp.const(dtypes.weakint, 240), 0, AxisType.REDUCE)

test/null/test_uop_symbolic.py

@@ -347,6 +347,12 @@ class TestSymbolic(unittest.TestCase):
   def test_mul_lt(self):
     self.helper_test_variable(Variable("a", 0, 5)*4 < 13, 0, 1, "(a<4)")
     self.helper_test_variable(Variable("a", 0, 5)*4 < 16, 0, 1, "(a<4)")
+    self.helper_test_variable(Variable("a", -5, 5)*4 < -13, 0, 1, "(a<-3)")
+    self.helper_test_variable(Variable("a", -5, 5)*-4 < 13, 0, 1, "((a*-1)<4)")
+    c0, c1 = 2, 2**54+1
+    self.helper_test_variable(Variable("a", 0, c1)*c0 < c1, 0, 1, f"(a<{2**53+1})")
+    c0, c1 = -2, -(2**54-1)
+    self.helper_test_variable(Variable("a", 0, -c1)*c0 < c1, 0, 1, f"((a*-1)<{-(2**53-1)})")
     self.helper_test_variable(Variable("a", 0, 5)*(-2) < 0, 0, 1, "((a*-1)<0)")
     self.helper_test_variable(Variable("a", 0, 5)*4 >= 12, 0, 1, "((a<3)!=True)")
     self.helper_test_variable(Variable("a", 0, 5)*4 >= 13, 0, 1, "((a<4)!=True)")

test/null/test_uops.py

@@ -1,5 +1,5 @@
 # uops tests that pass on NULL backend (no copyout needed)
-import unittest
+import math, unittest
 import numpy as np
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import Timing, Context, cdiv
@@ -8,7 +8,6 @@ from tinygrad.device import Device
 from tinygrad.uop.ops import Ops, ParamArg, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite, pm_lower_index_dtype  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
 from tinygrad.uop.spec import spec_program, spec_shared, type_verify
 from tinygrad.uop.symbolic import sym, pm_remove_invalid
-from tinygrad.mixin.movement import MovementMixin
 from test.helpers import eval_uop, to_uops_list
 
 class TestDTypeFromUOp(unittest.TestCase):
@@ -116,6 +115,11 @@ class TestExecALU(unittest.TestCase):
   def test_sqrt(self):
     self.assertEqual(exec_alu(Ops.SQRT, dtypes.float, (0.0,)), 0.0)
 
+  def test_trunc_nonfinite(self):
+    self.assertEqual(exec_alu(Ops.TRUNC, dtypes.float, (math.inf,)), math.inf)
+    self.assertEqual(exec_alu(Ops.TRUNC, dtypes.float, (-math.inf,)), -math.inf)
+    self.assertTrue(math.isnan(exec_alu(Ops.TRUNC, dtypes.float, (math.nan,))))
+
   def test_invalid_poison(self):
     # Invalid poisons any binary op regardless of result dtype: a comparison must not fold to a boolean
     self.assertIs(exec_alu(Ops.CMPLT, dtypes.bool, (Invalid, 1)), Invalid)
@@ -458,7 +462,7 @@ class TestContiguousViewOffset(unittest.TestCase):
   def test_2d(self): self._check(UOp.empty(2,5)[1, 2:4], 7)
   def test_shrink_to_one(self): self._check(UOp.empty(10)[1], 1)
   def test_expand_is_none(self): self._check(UOp.empty(1).expand(2), None)
-  def test_shrink_invalid(self): self._check(MovementMixin.pad(UOp.empty(4), ((2,2),))[0], None)
+  def test_shrink_invalid(self): self._check(UOp.empty(4).pad((2,2))[0], None)
   def test_strided(self): self._check(UOp.empty(4)[::2], None)
 
 if __name__ == '__main__':

test/opt/test_kernel_opts.py

@@ -270,7 +270,7 @@ class TestKernelOpts(unittest.TestCase):
   def test_padto_where(self):
     Tensor.manual_seed(0)
     N = 17
-    a = (Tensor.randn(N, N).realize().max(axis=0, keepdim=True) > 1).where(1, 0)
+    a = (Tensor.randn(N, N).realize().max(axis=0, keepdim=True) > 1).where(1, 0).int()
     helper_linearizer_opt(a.max(0), [
       [Opt(OptOps.PADTO, 0, 32)],
       [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.UPCAST, 0, 8),],
@@ -280,8 +280,8 @@ class TestKernelOpts(unittest.TestCase):
     Tensor.manual_seed(0)
     N = 17
     r = Tensor.randn(N, N).realize().max(axis=0, keepdim=True) > 1
-    a0 = r.where(1, 0)
-    a1 = r.where(2, 0)
+    a0 = r.where(1, 0).int()
+    a1 = r.where(2, 0).int()
     helper_linearizer_opt([a0.max(0), a1.max(0)], [
       [Opt(OptOps.PADTO, 0, 32)],
       [Opt(OptOps.PADTO, 0, 32), Opt(OptOps.UPCAST, 0, 8),],

test/unit/test_assign.py

@@ -1066,5 +1066,13 @@ class TestAfterCachePatterns(unittest.TestCase):
     np.testing.assert_array_equal(head.numpy(), [3])
     np.testing.assert_array_equal(full.numpy(), [1, 2])
 
+class TestBatchNormRunningStats(unittest.TestCase):
+  @unittest.expectedFailure  # TODO: nothing reads the stat update so it is never scheduled, and the chain grows every step
+  def test_running_stats_are_realized(self):
+    from tinygrad import nn
+    bn, x = nn.BatchNorm(4), Tensor.randn(2, 4, 3, 3).contiguous().realize()
+    with Context(TRAINING=1): bn(x).realize()
+    self.assertTrue(bn.running_mean.uop.base.is_realized)
+
 if __name__ == "__main__":
   unittest.main()

test/unit/test_attention.py

@@ -1,9 +1,9 @@
 import unittest
 import numpy as np
-from tinygrad import Tensor, dtypes
+from tinygrad import Device, Tensor, dtypes
 from tinygrad.llm.model import (
   GatedDeltaNetBlock, SSMConfig, TransformerBlock, TransformerConfig,
-  apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk,
+  apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk, topk_softmax,
 )
 
 def apply_rope(x:Tensor, start_pos:int):
@@ -67,7 +67,10 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
   def _run_attention(self, block:GatedDeltaNetBlock, x:Tensor, start_pos:int):
     x_norm = block.attn_norm(x)
     block._init_state(x_norm)
-    return block._attention(x_norm, start_pos).realize().numpy()
+    out = block._attention(x_norm, start_pos).realize()
+    assert block.pending_state is not None
+    Tensor.realize(block.conv_state.assign(block.pending_state[0]), block.recurrent_state.assign(block.pending_state[1]))
+    return out.numpy()
 
   def _cache_views(self, block:GatedDeltaNetBlock) -> tuple[np.ndarray, np.ndarray]:
     if hasattr(block, 'conv_state'):
@@ -86,8 +89,8 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
     x_float = x.astype(np.float32)
     return (x_float / np.sqrt((x_float * x_float).mean(axis=-1, keepdims=True) + eps)) * weight.astype(np.float32)
 
-  def _normalize_np(self, x:np.ndarray, eps:float=1e-12) -> np.ndarray:
-    return x / np.maximum(np.sqrt((x * x).sum(axis=-1, keepdims=True)), eps)
+  def _normalize_np(self, x:np.ndarray, eps:float=1e-6) -> np.ndarray:
+    return x / np.sqrt((x * x).sum(axis=-1, keepdims=True) + eps)
 
   def _softplus_np(self, x:np.ndarray) -> np.ndarray:
     return np.log1p(np.exp(-np.abs(x))) + np.maximum(x, 0)
@@ -199,5 +202,27 @@ class TestPairwiseTopk(unittest.TestCase):
         self.assertEqual(set(sel.numpy()[b, t].tolist()), expected)
         np.testing.assert_allclose(vals.numpy()[b, t], data[b, t][sel.numpy()[b, t]])
 
+  def test_256_experts_matches_numpy(self):
+    rng = np.random.default_rng(42)
+    data = rng.standard_normal((256, 256) if Device.DEFAULT.startswith("AMD") else (4, 3, 256), dtype=np.float32)
+    # Include ties crossing wave boundaries to cover deterministic expert selection.
+    data[..., [7, 39, 71, 103, 135, 167, 199, 231]] = 10.0
+    expected = np.apply_along_axis(lambda row:np.lexsort((-np.arange(256), row))[-8:], -1, data)
+    x = Tensor(data)
+    for _ in range(5 if Device.DEFAULT.startswith("AMD") else 1):
+      vals, sel = pairwise_topk(x, 8)
+      np.testing.assert_equal(sel.numpy(), expected)
+      np.testing.assert_allclose(vals.numpy(), np.take_along_axis(data, expected, axis=-1))
+
+  def test_256_experts_softmax_matches_reference(self):
+    rng = np.random.default_rng(123)
+    data = rng.standard_normal((256, 256) if Device.DEFAULT.startswith("AMD") else (4, 3, 256), dtype=np.float32)
+    data[..., [7, 39, 71, 103, 135, 167, 199, 231]] = 10.0
+    probs, sel = topk_softmax(Tensor(data), 8)
+    selected = np.take_along_axis(data, sel.numpy(), axis=-1)
+    expected = np.exp(selected - selected.max(axis=-1, keepdims=True))
+    expected /= expected.sum(axis=-1, keepdims=True)
+    np.testing.assert_allclose(probs.numpy(), expected, rtol=2e-6, atol=2e-7)
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_cpu.py

@@ -1,13 +1,47 @@
-import unittest, io
+import unittest, io, os, subprocess, sys
 from contextlib import redirect_stdout
-from tinygrad import Tensor, Device
+from tinygrad import Tensor, Device, UOp
 from tinygrad.helpers import Target
 from tinygrad.renderer.nir import LVPRenderer
 from tinygrad.renderer.isa.x86 import X86Renderer
 from tinygrad.codegen import to_program
+from tinygrad.runtime.ops_cpu import RING_SLOTS
+from tinygrad.uop.ops import AxisType, KernelInfo
 
 @unittest.skipIf(Device.DEFAULT != "CPU", "only run on CPU")
 class TestCPU(unittest.TestCase):
+  def test_parallel_workers_exit_cleanly(self):
+    env = os.environ.copy()
+    env.update(DEV="CPU", CPU_PARALLEL_UOPS="1")
+    proc = subprocess.run([sys.executable, "-c",
+      "from tinygrad import Tensor; assert Tensor.arange(32).sum().item() == 496"], env=env, capture_output=True, text=True)
+    self.assertEqual(proc.returncode, 0, proc.stderr)
+
+  def test_32_buffer_kernel(self):
+    def add_inputs(out:UOp, *inputs:UOp) -> UOp:
+      return out[0].store(sum((x[0] for x in inputs), start=UOp.const(out.dtype, 0))).sink(
+        arg=KernelInfo(name="add_31_inputs", opts_to_apply=()))
+    inputs = [Tensor([i], device="CPU").realize() for i in range(31)]
+    out = Tensor.custom_kernel(Tensor.empty(1, device="CPU"), *inputs, fxn=add_inputs)[0]
+    self.assertEqual(out.item(), sum(range(31)))
+
+  def test_command_ring_backpressure(self):
+    dev, count = Device["CPU"], RING_SLOTS + 257
+    signal, queue = dev.new_signal(value=0), dev.hw_compute_queue_t()
+    for value in range(1, count + 1): queue.signal(signal, value)
+    queue.submit(dev)
+    signal.wait(count, timeout=10000)
+    self.assertEqual(signal.value, count)
+
+  def test_parallel_launch(self):
+    def fill(out:UOp) -> UOp:
+      idx = UOp.range(67, 0, AxisType.GLOBAL)
+      return out[idx].store(idx).end(idx).sink(arg=KernelInfo(name="parallel_launch", optimize=False, parallel=True))
+    probe = Tensor.custom_kernel(Tensor.empty(67, device="CPU"), fxn=fill)[0]
+    self.assertTrue(to_program(probe.schedule_linear().src[-1].src[0], Device["CPU"].renderer).arg.parallel)
+    out = Tensor.custom_kernel(Tensor.empty(67, device="CPU"), fxn=fill)[0]
+    self.assertEqual(out.tolist(), list(range(67)))
+
   def test_arch_feats(self):
     ast = (Tensor.empty(16) + Tensor.empty(16)).schedule_linear().src[-1].src[0]
     for ren in Device[Device.DEFAULT].renderers:

test/unit/test_dtype_spec.py

@@ -1,8 +1,8 @@
 import unittest, math, subprocess
 from tinygrad.tensor import Tensor
-from tinygrad.dtype import dtypes, DType, DTYPES_DICT
+from tinygrad.dtype import dtypes, DType, DTYPES_DICT, strong_dtype
 from tinygrad.device import Device
-from tinygrad.helpers import getenv, DEBUG, EMULATED_DTYPES
+from tinygrad.helpers import getenv, DEBUG, EMULATED_DTYPES, Context
 from test.helpers import slow
 from hypothesis import given, settings, strategies as strat
 import numpy as np
@@ -25,6 +25,8 @@ def _assert_eq(tensor:Tensor, target_dtype:DType, target, tol_target_dtype:float
   if DEBUG >= 2: print(tensor.numpy())
   try:
     assert tensor.dtype == target_dtype
+    # weak values read back at their default.
+    target_dtype = strong_dtype(target_dtype)
     # denormals are zero
     if target_dtype in dtypes.floats and (target_dtype not in supported_dtypes or target_dtype in EMULATED_DTYPES.tolist(dtypes)):
       fe, fm = dtypes.finfo(target_dtype)
@@ -37,10 +39,13 @@ def _assert_eq(tensor:Tensor, target_dtype:DType, target, tol_target_dtype:float
     raise AssertionError(f"\ntensor {tensor.numpy()} dtype {tensor.dtype} does not match target {target} with dtype {target_dtype}") from e
 
 class TestTypeSpec(unittest.TestCase):
-  def setUp(self):
-    self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float
-  def tearDown(self):
-    dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float
+  def test_default_dtype_context(self):
+    default_float, default_int = dtypes.default_float, dtypes.default_int
+    with Context(DEFAULT_FLOAT=dtypes.half, DEFAULT_INT=dtypes.int16):
+      assert dtypes.default_float is dtypes.half
+      assert dtypes.default_int is dtypes.int16
+    assert dtypes.default_float is default_float
+    assert dtypes.default_int is default_int
 
   @unittest.skip("this test is slow and spawning whole pythons")
   def test_env_set_default_float(self):
@@ -79,11 +84,11 @@ class TestTypeSpec(unittest.TestCase):
 
   @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
   def test_creation(self, default_int, default_float):
-    dtypes.default_int, dtypes.default_float = default_int, default_float
+    self.enterContext(Context(DEFAULT_INT=default_int, DEFAULT_FLOAT=default_float))
     _assert_eq(Tensor(True), dtypes.bool, True)
-    _assert_eq(Tensor(None), dtypes.default_float, [])
-    _assert_eq(Tensor(2), dtypes.default_int, 2)
-    _assert_eq(Tensor(2.34), dtypes.default_float, 2.34)
+    _assert_eq(Tensor(None), dtypes.weakfloat, [])
+    _assert_eq(Tensor(2), dtypes.weakint, 2)
+    _assert_eq(Tensor(2.34), dtypes.weakfloat, 2.34)
     _assert_eq(Tensor([]), dtypes.default_float, [])
     _assert_eq(Tensor([1]), dtypes.default_int, [1])
     # list elements are python scalars; a numpy scalar in a list has no inferred dtype (use np.array or state a dtype)
@@ -98,7 +103,7 @@ class TestTypeSpec(unittest.TestCase):
 
   @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
   def test_full(self, default_int, default_float):
-    dtypes.default_int, dtypes.default_float = default_int, default_float
+    self.enterContext(Context(DEFAULT_INT=default_int, DEFAULT_FLOAT=default_float))
 
     _assert_eq(Tensor.zeros((2, 3)), dtypes.default_float, np.zeros((2, 3)))
     _assert_eq(Tensor.zeros((2, 3), dtype=dtypes.int64), dtypes.int64, np.zeros((2, 3)))
@@ -121,7 +126,7 @@ class TestTypeSpec(unittest.TestCase):
 
   @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
   def test_reduce_0d_default(self, default_int, default_float):
-    dtypes.default_int, dtypes.default_float = default_int, default_float
+    self.enterContext(Context(DEFAULT_INT=default_int, DEFAULT_FLOAT=default_float))
     _assert_eq(Tensor.ones((2,3,0)).sum(2), dtypes.default_float, np.zeros((2, 3)))
     # TODO: what should this one be?
     # _assert_eq(Tensor.ones((2,3,0), dtype=dtypes.default_int).sum(2), dtypes.default_int, np.zeros((2, 3)))
@@ -129,7 +134,7 @@ class TestTypeSpec(unittest.TestCase):
 
   @given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
   def test_arange(self, default_int, default_float):
-    dtypes.default_int, dtypes.default_float = default_int, default_float
+    self.enterContext(Context(DEFAULT_INT=default_int, DEFAULT_FLOAT=default_float))
 
     _assert_eq(Tensor.arange(5), dtypes.default_int, np.arange(5))
     _assert_eq(Tensor.arange(120), dtypes.default_int, np.arange(120))
@@ -146,11 +151,6 @@ class TestTypeSpec(unittest.TestCase):
     _assert_eq(Tensor.arange(5.0, 3.0), dtypes.default_float, np.arange(5.0, 3.0))
 
 class TestAutoCastType(unittest.TestCase):
-  def setUp(self):
-    self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float
-  def tearDown(self):
-    dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float
-
   def test_int_sqrt(self):
     _assert_eq(Tensor([1, 4, 9, 16]).sqrt(), dtypes.default_float, [1, 2, 3, 4])
 
@@ -190,22 +190,18 @@ class TestAutoCastType(unittest.TestCase):
     np.testing.assert_allclose(t.prod(dtype=dtypes.float32).numpy(), 20000)
 
   def test_gradient_dtype(self):
-    old_default_float = dtypes.default_float
-
     for default_dtype in dtypes.floats:
       if default_dtype not in supported_dtypes: continue
-      dtypes.default_float = default_dtype
-      for dtype in dtypes.floats:
-        if dtype not in supported_dtypes: continue
-        if DEBUG >= 2:
-          print(f"testing {default_dtype=}, {dtype=}")
-        a = Tensor([1, 2, 3], dtype=dtype)
-        b = (a * 5).sum()
-        b.backward()  # if there is dtype mismatch, lazy should assert
-        assert a.grad.dtype == a.dtype
-        np.testing.assert_allclose(a.grad.numpy(), [5, 5, 5])
-
-    dtypes.default_float = old_default_float
+      with Context(DEFAULT_FLOAT=default_dtype):
+        for dtype in dtypes.floats:
+          if dtype not in supported_dtypes: continue
+          if DEBUG >= 2:
+            print(f"testing {default_dtype=}, {dtype=}")
+          a = Tensor([1, 2, 3], dtype=dtype)
+          b = (a * 5).sum()
+          b.backward()  # if there is dtype mismatch, lazy should assert
+          assert a.grad.dtype == a.dtype
+          np.testing.assert_allclose(a.grad.numpy(), [5, 5, 5])
 
   @unittest.skipIf(Device.DEFAULT == "PYTHON", "very slow")
   @slow

test/unit/test_dtype_weak.py

@@ -1,9 +1,11 @@
 import tempfile, unittest, math
 
-from tinygrad import Tensor, dtypes
+from tinygrad import Tensor, dtypes, TinyJit
 from tinygrad.helpers import Context
-from tinygrad.uop.ops import UOp, Ops
+from tinygrad.dtype import least_upper_float
+from tinygrad.uop.ops import UOp, Ops, dtype_from_uop
 from tinygrad.uop.spec import spec_shared, type_verify
+from tinygrad.engine.jit import JitError
 
 
 class TestWeakPromotion(unittest.TestCase):
@@ -12,10 +14,11 @@ class TestWeakPromotion(unittest.TestCase):
     with self.assertRaises(ValueError): Tensor.const(dtypes.weakfloat, 1.0).rand_like()
     with self.assertRaises(ValueError): Tensor.const(dtypes.weakfloat, 1.0).randn_like()
 
-  def test_sum_stays_weak(self):
-    for weak, value in ((dtypes.weakfloat, 1.0),):
-      self.assertEqual(Tensor.const(weak, value).expand(3).sum().dtype, weak)
-    self.assertEqual((Tensor.const(dtypes.weakfloat, 1.0).expand(3).sum() + Tensor([1], dtype=dtypes.float16)).dtype, dtypes.float16)
+  def test_reduce_strips_weakness(self):
+    for weak, value, strong in ((dtypes.weakint, 1, dtypes.default_int), (dtypes.weakfloat, 1.0, dtypes.default_float)):
+      t = Tensor.const(weak, value).expand(3)
+      for out in (t.sum(), t.max(), t.prod(), t.cumsum(0), t.cummax(0)[0]): self.assertEqual(out.dtype, strong)
+    self.assertEqual((Tensor.const(dtypes.weakfloat, 1.0).expand(3).sum() + Tensor([1], dtype=dtypes.float16)).dtype, dtypes.float32)
 
   def test_materialize_at_default_dtype(self):
     for weak, value, strong in ((dtypes.weakfloat, 0.5, dtypes.default_float),):
@@ -23,7 +26,22 @@ class TestWeakPromotion(unittest.TestCase):
       self.assertEqual(t.dtype, weak)
       self.assertEqual(t.data().itemsize, strong.itemsize)
       self.assertEqual(t.numpy().dtype.itemsize, strong.itemsize)
-      with self.assertRaises(RuntimeError): t.clone("CPU")
+      # materializing commits at the kind default; contiguous has no layout to fix so it stays weak
+      self.assertEqual((c := t.clone("CPU")).dtype, strong)
+      self.assertEqual(c.item(), value)
+      self.assertEqual(t.contiguous().dtype, weak)
+
+  def test_assign_into_weak_commits(self):
+    t = Tensor.const(dtypes.weakfloat, 0.5)
+    t.assign(Tensor(1.0, dtype=dtypes.default_float))
+    self.assertEqual((t.dtype, t.item()), (dtypes.default_float, 1.0))
+
+  def test_float_unary_on_weakint_stays_weak(self):
+    self.assertIs(least_upper_float(dtypes.weakint), dtypes.weakfloat)
+
+  def test_copysign_meets_operands(self):
+    r = Tensor([2], dtype=dtypes.uint8, device="CPU").copysign(Tensor([1], dtype=dtypes.uint32, device="CPU"))
+    self.assertEqual((r.dtype, r.tolist()), (dtypes.uint32, [2]))
 
   def test_uop_scalar_const_unchanged(self):
     for dtype, value in ((dtypes.weakint, 1), (dtypes.int32, 1), (dtypes.float32, 0.5)):
@@ -50,7 +68,10 @@ class TestWeakPromotion(unittest.TestCase):
     self.assertEqual((t_f32 + t_f16).dtype, dtypes.float32)
     self.assertEqual(Tensor([2], dtype=dtypes.uint8).pad(((1, 1),), value=1).dtype, dtypes.uint8)
 
-  @unittest.expectedFailure  # TODO: dot of a weak const tensor defers to the other operand once python scalars are weak consts
+  def test_concrete_pair_promotes_weak(self):
+    out = Tensor([-1], dtype=dtypes.int64, device="CPU") + Tensor([3], dtype=dtypes.uint64, device="CPU") + Tensor(0.5)
+    self.assertEqual((out.dtype, out.tolist()), (dtypes.weakfloat, [2.5]))
+
   def test_dot_defers_weak(self):
     weak = Tensor([True, False]).where(Tensor(1), 2)
     self.assertEqual(weak.dot(Tensor([1, 1], dtype=dtypes.int8)).dtype, dtypes.int8)
@@ -58,14 +79,18 @@ class TestWeakPromotion(unittest.TestCase):
   def test_weak_int_binop(self):
     v = UOp.variable("i", 0, 10, dtypes.weakint)
     self.assertEqual((v << 1).dtype, dtypes.weakint)
+    self.assertEqual(dtype_from_uop(Ops.SHL, (UOp.const(dtypes.int8, 1), UOp.const(dtypes.uint32, 1)), None), dtypes.int64)
+    self.assertEqual(UOp.const(dtypes.weakint, 1).alu(Ops.SHL, UOp.const(dtypes.uint8, 1)).dtype, dtypes.uint8)
     self.assertEqual((v & 3).dtype, dtypes.weakint)
     with self.assertRaises(RuntimeError): Tensor.const(dtypes.weakfloat, 1.0) << Tensor.const(dtypes.weakfloat, 1.0)
     with self.assertRaises(RuntimeError): UOp.const(dtypes.int32, 1).alu(Ops.SHL, UOp.const(dtypes.float64, 1))
-    # float bitwise/shift builds, the spec rejects it
+    for op in (Ops.SHL, Ops.SHR):
+      with self.assertRaises(RuntimeError):
+        UOp.const(dtypes.float32, 1).alu(op, UOp.const(dtypes.int32, 1))
+    # float bitwise builds, the spec rejects it
     with Context(SPEC=1):
       f32, wf = UOp.const(dtypes.float32, 1.0), UOp.const(dtypes.weakfloat, 1.0)
-      for bad in (f32.alu(Ops.AND, f32), f32.alu(Ops.SHL, UOp.const(dtypes.int32, 1)),
-                  UOp(Ops.AND, dtypes.float32, (f32, f32)), UOp(Ops.AND, dtypes.int32, (wf, wf))):
+      for bad in (f32.alu(Ops.AND, f32), UOp(Ops.AND, dtypes.float32, (f32, f32)), UOp(Ops.AND, dtypes.int32, (wf, wf))):
         with self.assertRaises(RuntimeError): type_verify([bad], spec_shared)
 
   def test_integer_values(self):
@@ -81,7 +106,6 @@ class TestWeakPromotion(unittest.TestCase):
     x32 = Tensor.full((1,), 0.0, dtype=dtypes.float32, device="CPU")
     self.assertEqual((x32 + value).item(), 1.0)
 
-  @unittest.expectedFailure  # TODO: exp/cos/sigmoid of a weak const stay weak instead of casting to a concrete float
   def test_weak_transcendentals(self):
     t_f16 = Tensor([1], dtype=dtypes.float16)
     for out in (Tensor(2).exp(), Tensor(2).cos(), Tensor(2).sigmoid()):
@@ -93,6 +117,13 @@ class TestWeakPromotion(unittest.TestCase):
       t.realize()
       self.assertNotIn(t.uop.buffer.dtype, dtypes.weaks)
 
+  def test_computed_float_index_lowers(self):
+    # a half-pixel nearest index resolves its float-scaled range before the gather
+    idx = (Tensor.arange(8) + 0.5) / 4 - 0.5
+    idx = (idx.clip(0, 1) - 0.5).ceil().int()
+    out = Tensor([0, 1], device="NULL")[idx].contiguous().realize()
+    self.assertNotIn(out.uop.buffer.dtype, dtypes.weaks)
+
 
 class TestWeakStorageBoundary(unittest.TestCase):
   # weak has no storage: a weak assignment source casts when it defers to the destination, everything else raises
@@ -125,10 +156,24 @@ class TestWeakMaterializationEntries(unittest.TestCase):
       self.assertEqual(weak_val().numpy().dtype.itemsize, strong.itemsize)
       self.assertEqual(weak_val().tolist(), [value])
       self.assertEqual(weak_val().cast(strong).realize().uop.buffer.dtype, strong)
-      for entry in (lambda t: t.contiguous(), lambda t: t.realize(), lambda t: t.clone(),
-                    lambda t: t.to("CPU:1").realize(), lambda t: t.as_param(0)):
+      self.assertEqual(weak_val().contiguous().dtype, weak)                 # no layout to fix, stays weak
+      self.assertEqual(weak_val().realize().dtype, weak)                    # no width to store, stays weak
+      self.assertEqual(weak_val().clone().dtype, strong)                    # storage commits at the default
+      for entry in (lambda t: t.to("CPU:1").realize(), lambda t: t.as_param(0)):
         with self.assertRaises(RuntimeError): entry(weak_val())
 
+  def test_weak_is_virtual(self):
+    # NOTE: int64 lub uint64 is weakfloat, so this is device-ful weak from promotion, never from a cast to weak
+    devful = Tensor([1], dtype=dtypes.int64, device="CPU") + Tensor([1], dtype=dtypes.uint64, device="CPU")
+    for t in (Tensor.const(dtypes.weakfloat, 0.5), devful):
+      self.assertTrue(t.uop.is_virtual)
+      # realize is a no-op, so a weak input can never become the real buffer TinyJit needs
+      with self.assertRaises(JitError): TinyJit(lambda x: (x+1).realize())(t)
+    # callify must not silently commit a weak CONTIGUOUS to storage
+    c = devful.alu(Ops.CONTIGUOUS)
+    c.callify()
+    self.assertIs(c.dtype, dtypes.weakfloat)
+
   def test_empty_reads_commit(self):
     for weak, strong in ((dtypes.weakfloat, dtypes.default_float),):
       empty = Tensor.const(weak, 0).reshape(1).shrink(((0, 0),))

test/unit/test_function.py

@@ -535,7 +535,7 @@ class TestFunctionTuple(unittest.TestCase):
 
     @function(precompile=True, precompile_backward=True)
     def f(a:Tensor):
-      c = Tensor(Tensor.invalids(a.shape[0]//len(devs), a.shape[1], dtype=a.dtype, device=devs).uop.multi(0), device=devs)
+      c = Tensor(Tensor.invalids(a.shape[0]//len(devs), a.shape[1], dtype=a.dtype, device=devs).uop.unshard(0), device=devs)
       return Tensor.custom_kernel(c, a, fxn=double_kernel, grad_fxn=double_grad)[0]
 
     np.testing.assert_allclose(f(a).numpy(), 14.0)
@@ -543,7 +543,7 @@ class TestFunctionTuple(unittest.TestCase):
     # g is f with empty output instead of invalids
     @function(precompile=True, allow_implicit=True)
     def g(a:Tensor):
-      c = Tensor(Tensor.empty(a.shape[0]//len(devs), a.shape[1], dtype=a.dtype, device=devs).uop.multi(0), device=devs)
+      c = Tensor(Tensor.empty(a.shape[0]//len(devs), a.shape[1], dtype=a.dtype, device=devs).uop.unshard(0), device=devs)
       return Tensor.custom_kernel(c, a, fxn=double_kernel, grad_fxn=double_grad)[0]
 
     np.testing.assert_allclose(g(a).numpy(), 14.0)
@@ -557,6 +557,14 @@ class TestFunctionTuple(unittest.TestCase):
     def f(a:Tensor): return Tensor.custom_kernel(Tensor.empty(*a.shape, dtype=a.dtype, device=a.device), a, fxn=inplace_add)[0]
     with self.assertRaisesRegex(RuntimeError, "implicit buffer"): f(Tensor([1., 2., 3., 4.]).contiguous().realize())
 
+  def test_shrink_load_is_program_input(self):
+    out, inp = UOp.param(0, dtypes.float, (1,)), UOp.param(1, dtypes.float, (8,))
+    values = UOp(Ops.SHRINK, src=(inp, UOp.const(dtypes.weakint, 0), UOp.const(dtypes.weakint, 8))).load()
+    sink = out[0].store(values.index(0)).sink(arg=KernelInfo(name="vector_load"))
+    info = ProgramInfo.from_sink(sink)
+    self.assertEqual(info.outs, (0,))
+    self.assertEqual(info.ins, (1,))
+
   def test_custom_kernel_write_only_persistent_output_is_implicit(self):
     # a write-only custom_kernel output that is a realized buffer must be captured
     def write(C:UOp, A:UOp) -> UOp:
@@ -608,7 +616,7 @@ class TestFunctionTuple(unittest.TestCase):
     @function(precompile=True)
     def f(a:Tensor):
       c = Tensor.invalids(*a.uop.shard_shape, dtype=a.dtype, device=a.device)
-      if multi: c = Tensor(c.uop.multi(a.uop.axis), device=a.device)
+      if multi: c = Tensor(c.uop.unshard(a.uop.axis), device=a.device)
       c = Tensor.custom_kernel(c, a, fxn=my_kernel)[0]
       return c + 1
 

test/unit/test_jit.py

@@ -2,6 +2,7 @@ import unittest, numpy as np
 from test.helpers import assert_jit_cache_len
 from tinygrad import Tensor, TinyJit, Context, UOp, dtypes
 from tinygrad.engine.jit import JitError
+from tinygrad.uop.ops import Ops
 
 def _simple_test(add, extract=lambda x: x, N=10):
   for _ in range(5):
@@ -12,6 +13,19 @@ def _simple_test(add, extract=lambda x: x, N=10):
   assert_jit_cache_len(add, 1)
 
 class TestJit(unittest.TestCase):
+  def test_parallel_compile(self):
+    from tinygrad.codegen import to_program_cache
+    from tinygrad.engine.realize import compile_linear
+    calls = [call for i in range(16) for call in (Tensor.empty(1, device="CPU") + i).schedule_linear().src]
+    cache = to_program_cache.copy()
+    try:
+      to_program_cache.clear()
+      with Context(PARALLEL_COMPILE=2): linear = compile_linear(UOp(Ops.LINEAR, src=tuple(calls)), jit=True)
+      self.assertTrue(all(call.op is not Ops.CALL or call.src[0].op is Ops.PROGRAM for call in linear.src))
+    finally:
+      to_program_cache.clear()
+      to_program_cache.update(cache)
+
   def test_jitbeam_triggers_beam(self):
     from unittest.mock import patch
     from tinygrad.helpers import getenv as _getenv

test/unit/test_jit_footguns.py

@@ -137,11 +137,11 @@ class TestJitFootguns(unittest.TestCase):
     from tinygrad import Variable
     a = Tensor.rand(3, 10).realize()
 
-    # fixed: pad now has an explicit fill mask (internal PAD is Invalid-filled), which recomputes from the symbolic shape
+    # broken: pad is a view, BIND values frozen at capture (i=2)
     @TinyJit
     def f_broken(a): return (a+1).pad((None, (0, 10-a.shape[1]))).realize()
     for i in range(1, 5): f_broken(a[:, :Variable("i", 1, 10).bind(i)])
-    self.assertEqual(int((f_broken(a[:, :Variable("i", 1, 10).bind(4)])[0] != 0).sum().item()), 4)
+    self.assertEqual(int((f_broken(a[:, :Variable("i", 1, 10).bind(4)])[0] != 0).sum().item()), 2)  # should be 4!
 
     # workaround: contiguous fuses pad into kernel
     @TinyJit

test/unit/test_llm_quant_cpu.py

@@ -0,0 +1,752 @@
+import functools, sys, unittest
+import numpy as np
+
+from tinygrad import Device, Tensor, TinyJit, UOp, dtypes, nn
+from tinygrad.llm.gguf import _GGML_QUANT, ggml_data_to_tensor
+from tinygrad.llm.kernels import amd as llm_amd
+from tinygrad.llm.kernels.cpu import (attention_decode, attention_prefill, causal_conv_silu, expert_pair, expert_silu,
+                              expert_weighted_sum, f16_linear,
+                              f16_matvec, gated_delta, gated_delta_prefill, gated_delta_q8, gdn_qkv, iq3_repack, moe_ffn, q6_argmax, q8_batched_pair,
+                              q8_gdn_norm_projections, q8_gdn_projections, q8_linear_pair, q8_repack, q8_silu_linear,
+                              rmsnorm, rmsnorm_f16_linear, shared_gate,
+                              silu, silu_mul, uop_attention_prefill, uop_f16_matvec, uop_linear, uop_moe_ffn, uop_q8_linear_pair,
+                              uop_q8_prequant_linear, uop_expert_silu_weighted, weighted_sum)
+from tinygrad.llm.kernels.cpu import _dot_bytes_ptr, _dot_nibbles_ptr
+from tinygrad.llm.model import biased_sigmoid_topk, pairwise_topk, Embedding, ExpertWeights, FFNBlock, Linear, Transformer, TransformerConfig
+from tinygrad.uop.ops import KernelInfo
+
+
+def q8_activation(x:np.ndarray) -> np.ndarray:
+  grouped = x.reshape(*x.shape[:-1], -1, 32)
+  scale = np.maximum(np.max(np.abs(grouped), axis=-1, keepdims=True) / 127, 1e-8)
+  return (np.clip(np.rint(grouped / scale), -127, 127) * scale).reshape(x.shape)
+
+def q8k_activation(x:np.ndarray) -> np.ndarray:
+  grouped = x.reshape(*x.shape[:-1], -1, 256)
+  signed_max = np.take_along_axis(grouped, np.argmax(np.abs(grouped), axis=-1, keepdims=True), axis=-1)
+  scale = -signed_max / 127
+  inverse = np.divide(1, scale, out=np.zeros_like(scale), where=scale != 0)
+  quantized = np.sign(grouped * inverse) * np.floor(np.abs(grouped * inverse) + 0.5)
+  return (np.minimum(quantized, 127) * scale).reshape(x.shape)
+
+
+def random_packed(rng:np.random.Generator, ggml_type:int, elements:int) -> np.ndarray:
+  block_size, type_size = _GGML_QUANT[ggml_type]
+  blocks = rng.integers(0, 256, size=(elements // block_size, type_size), dtype=np.uint8)
+  scales = rng.uniform(0.001, 0.02, size=len(blocks)).astype(np.float16).view(np.uint8).reshape(-1, 2)
+  blocks[:, :2] = scales
+  if ggml_type in (12, 13): blocks[:, 2:4] = scales
+  if ggml_type == 14: blocks[:, -2:] = scales
+  return blocks.flatten()
+
+
+@unittest.skipUnless(Device.DEFAULT == "AMD", "requires DEV=AMD")
+class TestLLMQuantAMD(unittest.TestCase):
+  @staticmethod
+  def assert_q8_equal(result:tuple[Tensor, Tensor, Tensor], expected:np.ndarray):
+    grouped = expected.reshape(expected.shape[0], -1, 32)
+    scale = np.maximum(np.max(np.abs(grouped), axis=-1) / 127, 1e-8)
+    quant = np.clip(np.rint(grouped / scale[..., None]), -127, 127).astype(np.int8)
+    np.testing.assert_equal(result[0].numpy().view(np.int8).reshape(grouped.shape), quant)
+    np.testing.assert_allclose(result[1].numpy(), scale, rtol=1e-6, atol=1e-8)
+    np.testing.assert_equal(result[2].numpy(), quant.astype(np.int32).sum(-1))
+
+  def test_gated_delta_decode_matches_reference(self):
+    rng = np.random.default_rng(39)
+    batch, heads, dim = 1, 2, 128
+    q, k, v = [rng.standard_normal((batch, heads, dim), dtype=np.float32) for _ in range(3)]
+    beta, alpha = rng.random((batch, heads), dtype=np.float32), rng.uniform(0.9, 1, (batch, heads)).astype(np.float32)
+    state = rng.standard_normal((batch, heads, dim, dim), dtype=np.float32).astype(np.float16)
+    state_k, state_q = np.einsum("bhij,bhj->bhi", state, k), np.einsum("bhij,bhj->bhi", state, q)
+    delta = (v - state_k * alpha[..., None]) * beta[..., None]
+    expected_core = state_q * alpha[..., None] + delta * np.sum(k*q, axis=-1)[..., None]
+    expected_state = (state * alpha[..., None, None] + delta[..., None] * k[..., None, :]).astype(np.float16)
+    core, next_state = llm_amd.gated_delta_decode(
+      *(Tensor(x, device="AMD") for x in (q, k, v, beta, alpha)), Tensor(state, device="AMD"))
+    Tensor.realize(core, next_state)
+    np.testing.assert_allclose(core.numpy(), expected_core, rtol=2e-4, atol=1e-3)
+    np.testing.assert_allclose(next_state.numpy(), expected_state, rtol=1e-3, atol=4e-3)
+
+  def test_f16_matvec_matches_reference(self):
+    rng, in_features, out_features = np.random.default_rng(38), 512, 13
+    x = rng.standard_normal((1, in_features), dtype=np.float32).astype(np.float16)
+    weight = rng.standard_normal((out_features, in_features), dtype=np.float32).astype(np.float16)
+    got = llm_amd.f16_matvec(Tensor(x, device="AMD"), Tensor(weight, device="AMD")).numpy()
+    np.testing.assert_allclose(got, x.astype(np.float32) @ weight.astype(np.float32).T, rtol=2e-5, atol=2e-4)
+
+  def test_fused_rmsnorm_quantization_matches_reference(self):
+    rng, eps = np.random.default_rng(37), 1e-6
+    x = rng.standard_normal((1, 256), dtype=np.float32)
+    weight = rng.standard_normal((256,), dtype=np.float32).astype(np.float16)
+    normalized = x / np.sqrt(np.mean(x*x, axis=-1, keepdims=True) + eps) * weight
+    self.assert_q8_equal(llm_amd.q8_rmsnorm(Tensor(x, device="AMD"), Tensor(weight, device="AMD"), eps), normalized)
+
+    core = rng.standard_normal((1, 2, 128), dtype=np.float32)
+    gate = rng.standard_normal((1, 1, 2, 128), dtype=np.float32)
+    head_norm = core / np.sqrt(np.mean(core*core, axis=-1, keepdims=True) + eps) * weight[:128]
+    expected = head_norm * gate.reshape(1, 2, 128) / (1 + np.exp(-gate.reshape(1, 2, 128)))
+    self.assert_q8_equal(llm_amd.q8_gated_rmsnorm(*(Tensor(v, device="AMD") for v in (core, gate, weight[:128])), eps),
+                         expected.reshape(1, -1))
+
+  def test_gated_delta_prefill_matches_sequential_reference(self):
+    rng = np.random.default_rng(36)
+    batch, heads, tokens, dim = 1, 2, 5, 128
+    q, k, v = [rng.standard_normal((batch, heads, tokens, dim), dtype=np.float32) for _ in range(3)]
+    q = q / np.linalg.norm(q, axis=-1, keepdims=True) / np.float32(np.sqrt(dim))
+    k = k / np.linalg.norm(k, axis=-1, keepdims=True)
+    beta, alpha = rng.random((batch, heads, tokens), dtype=np.float32), rng.uniform(0.9, 1, (batch, heads, tokens)).astype(np.float32)
+    state = rng.standard_normal((batch, heads, dim, dim), dtype=np.float32).astype(np.float16)
+    expected_core, expected_state = np.empty_like(q), state.astype(np.float32)
+    for token in range(tokens):
+      state_k = np.einsum("bhij,bhj->bhi", expected_state, k[:, :, token])
+      state_q = np.einsum("bhij,bhj->bhi", expected_state, q[:, :, token])
+      delta = (v[:, :, token] - state_k * alpha[:, :, token, None]) * beta[:, :, token, None]
+      expected_core[:, :, token] = state_q * alpha[:, :, token, None] + delta * np.sum(k[:, :, token] * q[:, :, token], axis=-1)[..., None]
+      expected_state = expected_state * alpha[:, :, token, None, None] + delta[..., None] * k[:, :, token, None, :]
+    core, next_state = llm_amd.gated_delta_prefill(
+      *(Tensor(x, device="AMD") for x in (q, k, v, beta, alpha)), Tensor(state, device="AMD"))
+    np.testing.assert_allclose(core.numpy(), expected_core, rtol=2e-4, atol=1e-3)
+    np.testing.assert_allclose(next_state.numpy(), expected_state.astype(np.float16), rtol=2e-4, atol=1e-3)
+
+  def test_q8_quantize_matches_reference(self):
+    rng, tokens, in_features = np.random.default_rng(35), 3, 256
+    x = rng.standard_normal((tokens, in_features), dtype=np.float32)
+    grouped = x.reshape(tokens, -1, 32)
+    expected_scale = np.maximum(np.max(np.abs(grouped), axis=-1) / 127, 1e-8)
+    expected_quant = np.clip(np.rint(grouped / expected_scale[..., None]), -127, 127).astype(np.int8)
+    quant, scale, group_sum = llm_amd.q8_quantize_sum(Tensor(x, device="AMD"), tokens, in_features)
+    np.testing.assert_equal(quant.numpy().view(np.int8).reshape(grouped.shape), expected_quant)
+    np.testing.assert_allclose(scale.numpy(), expected_scale, rtol=1e-7, atol=0)
+    np.testing.assert_equal(group_sum.numpy(), expected_quant.astype(np.int32).sum(-1))
+
+  def test_q4_embedding_matches_reference(self):
+    rng, vocab_size, embed_size = np.random.default_rng(34), 16, 256
+    raw = random_packed(rng, 12, vocab_size * embed_size)
+    expected = ggml_data_to_tensor(Tensor(raw), vocab_size * embed_size, 12).reshape(vocab_size, embed_size).half()
+    storage = Tensor(np.concatenate((np.zeros(68, dtype=np.uint8), raw)), dtype=dtypes.uint8, device="AMD").realize()
+    embedding = Embedding(vocab_size, embed_size)
+    embedding.set_quantized(storage[68:], 12)
+    idx = np.array([[7, 1, 15], [0, 4, 7]], dtype=np.int32)
+    np.testing.assert_equal(embedding(Tensor(idx, device="AMD")).numpy(), expected.numpy()[idx])
+
+  def test_iq4_lut_is_ready_for_jit_capture(self):
+    rng, in_features, out_features = np.random.default_rng(33), 256, 16
+    raw = random_packed(rng, 23, out_features * in_features)
+    weight = ggml_data_to_tensor(Tensor(raw), out_features * in_features, 23).numpy().reshape(out_features, in_features)
+    llm_amd.iq4_half_lut.cache_clear()
+    layer = Linear(in_features, out_features, bias=False)
+    layer.set_quantized(Tensor(raw, dtype=dtypes.uint8, device="AMD").realize(), 23)
+    @TinyJit
+    def run(x:Tensor): return layer(x).realize()
+    x = rng.standard_normal((16, in_features), dtype=np.float32)
+    expected = x.astype(np.float16).astype(np.float32) @ weight.astype(np.float16).astype(np.float32).T
+    for _ in range(2): np.testing.assert_allclose(run(Tensor(x, device="AMD")).numpy(), expected, rtol=1e-5, atol=2e-3)
+
+  def test_packed_linear_offset_matches_reference(self):
+    rng = np.random.default_rng(32)
+    for ggml_type,in_features in ((8, 256), (12, 256), (13, 256), (14, 256), (23, 256)):
+      for tokens in ((1, 16, 32, 64, 128) if ggml_type == 23 else (1, 16, 128) if ggml_type in (12, 13) else
+                     (1, 16) if ggml_type == 14 else (1,)):
+        raw, out_features = random_packed(rng, ggml_type, 64 * in_features), 64
+        weight = ggml_data_to_tensor(Tensor(raw), out_features * in_features, ggml_type).numpy().reshape(out_features, in_features)
+        storage = Tensor(np.concatenate((np.zeros(68, dtype=np.uint8), raw)), dtype=dtypes.uint8, device="AMD").realize()
+        layer = Linear(in_features, out_features, bias=False)
+        layer.set_quantized(storage[68:], ggml_type)
+        x = rng.standard_normal((tokens, in_features), dtype=np.float32)
+        expected = x.astype(np.float16).astype(np.float32) @ weight.astype(np.float16).astype(np.float32).T \
+          if ggml_type in (12, 13, 23) and tokens > 1 else q8_activation(x) @ weight.T
+        np.testing.assert_allclose(layer(Tensor(x, device="AMD")).numpy(), expected, rtol=1e-5, atol=2e-3)
+
+  def test_iq3_expert_prefill_and_decode_match_reference(self):
+    rng = np.random.default_rng(31)
+    num_experts, in_features, out_features = 2, 256, 16
+    raw = random_packed(rng, 21, num_experts * out_features * in_features)
+    weight = ggml_data_to_tensor(Tensor(raw, device="CPU"), num_experts * out_features * in_features,
+                                 21).numpy().reshape(num_experts, out_features, in_features)
+    experts = ExpertWeights(num_experts, in_features, out_features)
+    experts.set_quantized(Tensor.empty(num_experts, out_features, in_features),
+                          Tensor(raw, dtype=dtypes.uint8, device="AMD").realize(), 21)
+    x = rng.standard_normal((2, 1, in_features), dtype=np.float32)
+    for sel in (np.array([[1, 0], [0, 1]], dtype=np.int32), np.array([1, 0], dtype=np.int32)):
+      activation = x if sel.ndim == 2 else x[:1]
+      expected = np.stack([q8_activation(activation).reshape(-1, in_features)[route // 2] @ weight[expert].T
+                           for route,expert in enumerate(sel.reshape(-1))]).reshape(*sel.shape, out_features)
+      got = experts(Tensor(sel, device="AMD"), Tensor(activation, device="AMD")).numpy()
+      np.testing.assert_allclose(got, expected, rtol=1e-5, atol=5e-4)
+
+
+@unittest.skipUnless(sys.platform.startswith("linux") and Device.DEFAULT == "CPU", "requires DEV=CPU on Linux")
+class TestLLMQuantCPU(unittest.TestCase):
+  def test_grouped_byte_dot_uop(self):
+    rng = np.random.default_rng(20)
+    a = rng.integers(-128, 128, 32, dtype=np.int8)
+    b = rng.integers(-127, 128, 32, dtype=np.int8)
+    out = Tensor.empty(8, dtype=dtypes.int32, device="CPU")
+    ta, tb = Tensor(a, device="CPU").realize(), Tensor(b, device="CPU").realize()
+    def dot_kernel(out:UOp, a:UOp, b:UOp) -> UOp:
+      parts = _dot_bytes_ptr(a[0], b[0])
+      return UOp.group(*(out[i].store(parts.index(i)) for i in range(8))).sink(arg=KernelInfo("grouped_byte_dot"))
+    got = Tensor.custom_kernel(out, ta, tb, fxn=dot_kernel)[0].numpy()
+    expected = (a.astype(np.int32) * b.astype(np.int32)).reshape(8, 4).sum(axis=1)
+    np.testing.assert_equal(got, expected)
+
+  def test_scaled_grouped_byte_dot_uop(self):
+    rng = np.random.default_rng(22)
+    a = rng.integers(-128, 128, 32, dtype=np.int8)
+    b = rng.integers(-127, 128, 32, dtype=np.int8)
+    out = Tensor.empty(8, dtype=dtypes.int32, device="CPU")
+    ta, tb = Tensor(a, device="CPU").realize(), Tensor(b, device="CPU").realize()
+    def dot_kernel(out:UOp, a:UOp, b:UOp) -> UOp:
+      parts = _dot_bytes_ptr(a[0], b[0]) * 7
+      return UOp.group(*(out[i].store(parts.index(i)) for i in range(8))).sink(arg=KernelInfo("scaled_grouped_byte_dot"))
+    got = Tensor.custom_kernel(out, ta, tb, fxn=dot_kernel)[0].numpy()
+    expected = (a.astype(np.int32) * b.astype(np.int32)).reshape(8, 4).sum(axis=1) * 7
+    np.testing.assert_equal(got, expected)
+
+  def test_unpack_lut_dot_uop(self):
+    rng = np.random.default_rng(21)
+    packed, x = rng.integers(0, 256, 16, dtype=np.uint8), rng.integers(-127, 128, 32, dtype=np.int8)
+    values = (-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113)
+    out = Tensor.empty(8, dtype=dtypes.int32, device="CPU")
+    tp, tx = Tensor(packed, device="CPU").realize(), Tensor(x, device="CPU").realize()
+    def dot_kernel(out:UOp, packed:UOp, x:UOp) -> UOp:
+      parts = _dot_nibbles_ptr(packed[0], x[0], values)
+      return UOp.group(*(out[i].store(parts.index(i)) for i in range(8))).sink(arg=KernelInfo("unpack_lut_dot"))
+    got = Tensor.custom_kernel(out, tp, tx, fxn=dot_kernel)[0].numpy()
+    decoded = np.array(values, dtype=np.int8)[np.concatenate((packed & 15, packed >> 4))]
+    expected = (decoded.astype(np.int32) * x.astype(np.int32)).reshape(8, 4).sum(axis=1)
+    np.testing.assert_equal(got, expected)
+
+  def test_generate_accepts_different_recurrent_prefill_shapes(self):
+    class TinyRecurrentTransformer(Transformer):
+      def __init__(self):
+        self.max_context, self.has_recurrent_block = 32, True
+        self.token_embd = nn.Embedding(4, 1)
+        self.blk, self._cached_tokens = [], []
+        self._state_checkpoints, self._state_checkpoint_pos = [], 0
+        self._save_state_jit = self._restore_state_jit = None
+        self._warming_up = False
+        self.prefill_jit = TinyJit(self.forward)
+        self.flash_prefill_jit = TinyJit(functools.partial(self.forward, use_flash=True))
+        self.sample_prefill_jit = TinyJit(functools.partial(self.forward, sample=True))
+        self.recurrent_prefill_jits = {}
+        self.rollout_jits, self.sample_rollout_jits = {}, {}
+
+      def forward(self, tokens:Tensor, start_pos:int|UOp, temperature:Tensor, use_flash:bool=False,
+                  kv_len:int|UOp|None=None, valid_len:int|UOp|None=None, sample:bool=False) -> Tensor:
+        return tokens[:, -1:] + 1
+
+    model = TinyRecurrentTransformer()
+    for _ in range(3): self.assertEqual(next(model.generate([1] * 8, chunk_size=8)), 2)
+    self.assertEqual(next(model.generate([1] * 3, chunk_size=8)), 2)
+
+  def test_attention_decode_matches_causal_gqa_reference(self):
+    rng = np.random.default_rng(0)
+    batch, heads, kv_heads, cache_len, head_dim = 1, 16, 2, 4113, 32
+    q = rng.standard_normal((batch, heads, 1, head_dim), dtype=np.float32)
+    cache = rng.standard_normal((2, batch, kv_heads, cache_len, head_dim), dtype=np.float32).astype(np.float16)
+    for pos in (0, 11, 4096):
+      with self.subTest(pos=pos):
+        start_pos = UOp.variable("start_pos", 0, cache_len-1).bind(pos)
+        got = attention_decode(Tensor(q, device="CPU").contiguous(),
+                                    Tensor(cache, device="CPU").contiguous(), start_pos).numpy()
+        expected = np.empty_like(got)
+        for head in range(heads):
+          kv_head = head // (heads // kv_heads)
+          keys = cache[0, 0, kv_head, :pos+1].astype(np.float32)
+          values = cache[1, 0, kv_head, :pos+1].astype(np.float32)
+          scores = q[0, head, 0] @ keys.T / np.sqrt(head_dim)
+          probs = np.exp(scores - scores.max())
+          expected[0, head, 0] = probs / probs.sum() @ values
+        np.testing.assert_allclose(got, expected, rtol=2e-5, atol=2e-5)
+
+  def test_attention_prefill_matches_causal_gqa_reference(self):
+    rng = np.random.default_rng(18)
+    batch, heads, tokens, kv_heads, cache_len, head_dim, pos = 1, 4, 5, 2, 19, 256, 3
+    q = rng.standard_normal((batch, heads, tokens, head_dim), dtype=np.float32)
+    cache = rng.standard_normal((2, batch, kv_heads, cache_len, head_dim), dtype=np.float32).astype(np.float16)
+    tq, tcache = Tensor(q, device="CPU").contiguous(), Tensor(cache, device="CPU").contiguous()
+    start_pos = UOp.variable("start_pos", 0, cache_len-1).bind(pos)
+    got = attention_prefill(tq, tcache, start_pos).numpy()
+    expected = np.empty_like(got)
+    for head in range(heads):
+      kv_head = head // (heads // kv_heads)
+      for token in range(tokens):
+        keys = cache[0, 0, kv_head, :pos+token+1].astype(np.float32)
+        values = cache[1, 0, kv_head, :pos+token+1].astype(np.float32)
+        scores = q[0, head, token] @ keys.T / np.sqrt(head_dim)
+        probs = np.exp(scores - scores.max())
+        expected[0, head, token] = probs / probs.sum() @ values
+    np.testing.assert_allclose(got, expected, rtol=2e-5, atol=2e-5)
+    np.testing.assert_allclose(uop_attention_prefill(tq, tcache, start_pos).numpy(), expected, rtol=2e-5, atol=2e-5)
+    np.testing.assert_allclose(uop_attention_prefill(tq[:, :, :4], tcache, start_pos).numpy(), expected[:, :, :4],
+                               rtol=2e-5, atol=2e-5)
+
+  def test_gated_delta_matches_reference(self):
+    rng = np.random.default_rng(4)
+    batch, heads, dim = 2, 3, 16
+    q, k, v = (rng.standard_normal((batch, heads, dim), dtype=np.float32) for _ in range(3))
+    beta, alpha = rng.random((batch, heads), dtype=np.float32), rng.random((batch, heads), dtype=np.float32)
+    state = rng.standard_normal((batch, heads, dim, dim), dtype=np.float32)
+    for state_dtype in (np.float32, np.float16):
+      with self.subTest(state_dtype=state_dtype):
+        typed_state = state.astype(state_dtype)
+        args = (*(Tensor(x, device="CPU") for x in (q, k, v, beta, alpha)), Tensor(typed_state, device="CPU"))
+        core, next_state = gated_delta(*args)
+        typed_state_k, typed_state_q = (np.einsum("bhij,bhj->bhi", typed_state.astype(np.float32), x) for x in (k, q))
+        typed_delta = (v - typed_state_k * alpha[..., None]) * beta[..., None]
+        typed_core = typed_state_q * alpha[..., None] + typed_delta * np.sum(k * q, axis=-1, keepdims=True)
+        typed_next = typed_state.astype(np.float32) * alpha[..., None, None] + typed_delta[..., None] * k[..., None, :]
+        np.testing.assert_allclose(core.numpy(), typed_core, rtol=2e-5, atol=2e-5)
+        np.testing.assert_allclose(next_state.numpy(), typed_next.astype(state_dtype), rtol=2e-5, atol=2e-5)
+        norm = nn.RMSNorm(dim, eps=1e-6)
+        norm.weight = Tensor(rng.standard_normal(dim, dtype=np.float32), device="CPU").half().realize()
+        normalized, _ = gated_delta(*args, norm_weight=norm.weight, norm_eps=norm.eps)
+        np.testing.assert_allclose(normalized.numpy(), rmsnorm(norm, core).numpy(), rtol=2e-5, atol=2e-5)
+        inplace_state = Tensor(typed_state, device="CPU").realize()
+        inplace_core, inplace_next = gated_delta(*(Tensor(x, device="CPU") for x in (q, k, v, beta, alpha)),
+                                                       inplace_state, inplace=True)
+        Tensor.realize(inplace_core, inplace_next)
+        np.testing.assert_allclose(inplace_core.numpy(), typed_core, rtol=2e-5, atol=2e-5)
+        np.testing.assert_allclose(inplace_next.numpy(), typed_next.astype(state_dtype), rtol=2e-5, atol=2e-5)
+        np.testing.assert_equal(inplace_state.numpy(), inplace_next.numpy())
+
+  def test_gated_delta_prefill_matches_sequential_reference(self):
+    rng = np.random.default_rng(17)
+    batch, heads, tokens, dim = 1, 2, 5, 128
+    q, k, v = [rng.standard_normal((batch, heads, tokens, dim), dtype=np.float32) for _ in range(3)]
+    beta = rng.random((batch, heads, tokens), dtype=np.float32)
+    alpha = rng.random((batch, heads, tokens), dtype=np.float32)
+    state = rng.standard_normal((batch, heads, dim, dim), dtype=np.float32).astype(np.float16)
+    weight = rng.standard_normal(dim, dtype=np.float32).astype(np.float16)
+    eps = 1e-6
+    expected_core = np.empty_like(q)
+    expected_state = state.astype(np.float32)
+    for token in range(tokens):
+      for b in range(batch):
+        for h in range(heads):
+          qq, kk, vv = q[b, h, token], k[b, h, token], v[b, h, token]
+          aa, bb, current = alpha[b, h, token], beta[b, h, token], expected_state[b, h]
+          delta = (vv - (current @ kk) * aa) * bb
+          core = (current @ qq) * aa + delta * (kk @ qq)
+          expected_state[b, h] = current * aa + delta[:, None] * kk[None, :]
+          expected_core[b, h, token] = core / np.sqrt(np.mean(core * core) + eps) * weight
+    got_core, got_state = gated_delta_prefill(
+      *[Tensor(z, device="CPU") for z in (q, k, v, beta, alpha)], Tensor(state, device="CPU"), Tensor(weight, device="CPU"), eps)
+    np.testing.assert_allclose(got_core.numpy(), expected_core, rtol=2e-5, atol=2e-5)
+    np.testing.assert_allclose(got_state.numpy(), expected_state.astype(np.float16), rtol=1e-3, atol=1e-3)
+
+  def test_gated_delta_q8_feeds_projection(self):
+    rng = np.random.default_rng(18)
+    batch, heads, dim = 1, 2, 32
+    q, k, v = [Tensor(rng.standard_normal((batch, heads, dim), dtype=np.float32), device="CPU").realize() for _ in range(3)]
+    beta, alpha = [Tensor(rng.random((batch, heads), dtype=np.float32), device="CPU").realize() for _ in range(2)]
+    state = Tensor(rng.standard_normal((batch, heads, dim, dim), dtype=np.float32), device="CPU").half().realize()
+    gate = Tensor(rng.standard_normal((batch, heads, dim), dtype=np.float32), device="CPU").half().realize()
+    norm_weight = Tensor(rng.standard_normal(dim, dtype=np.float32), device="CPU").half().realize()
+    layer = Linear(heads * dim, 17, bias=False)
+    layer.set_quantized(Tensor(random_packed(rng, 8, layer.out_features * layer.in_features),
+                               dtype=dtypes.uint8, device="CPU").realize(), 8)
+
+    core, expected_state = gated_delta(q, k, v, beta, alpha, state, norm_weight=norm_weight, norm_eps=1e-6)
+    expected = q8_silu_linear(layer, gate.reshape(batch, 1, -1), core.reshape(batch, 1, -1))
+    xq, xd, got_state = gated_delta_q8(q, k, v, beta, alpha, state, gate, norm_weight, 1e-6)
+    got = uop_q8_prequant_linear(layer, xq, xd).reshape(batch, 1, -1)
+    np.testing.assert_allclose(got_state.numpy(), expected_state.numpy(), rtol=2e-5, atol=2e-5)
+    np.testing.assert_allclose(got.numpy(), expected.numpy(), rtol=1e-3, atol=1e-3)
+
+  def test_gdn_qkv_matches_normal_api(self):
+    rng = np.random.default_rng(19)
+    batch, tokens, k_heads, v_heads, dim = 2, 5, 2, 4, 16
+    conv = Tensor(rng.standard_normal((batch, tokens, (2*k_heads+v_heads)*dim), dtype=np.float32), device="CPU")
+    q, k, v = conv.split([k_heads*dim, k_heads*dim, v_heads*dim], dim=-1)
+    q = (q.reshape(batch, tokens, k_heads, dim) *
+         (q.reshape(batch, tokens, k_heads, dim).square().sum(-1, keepdim=True) + 1e-6).rsqrt()).repeat(1, 1, v_heads//k_heads, 1)
+    k = (k.reshape(batch, tokens, k_heads, dim) *
+         (k.reshape(batch, tokens, k_heads, dim).square().sum(-1, keepdim=True) + 1e-6).rsqrt()).repeat(1, 1, v_heads//k_heads, 1)
+    expected = (q.transpose(1, 2) * dim**-0.5, k.transpose(1, 2), v.reshape(batch, tokens, v_heads, dim).transpose(1, 2))
+    for got, ref in zip(gdn_qkv(conv, k_heads, v_heads, dim), expected):
+      np.testing.assert_allclose(got.numpy(), ref.numpy(), rtol=2e-5, atol=2e-5)
+
+  def test_decode_rmsnorm_matches_reference(self):
+    rng = np.random.default_rng(5)
+    for rows in (1, 32, 128):
+      for dtype in (dtypes.float16, dtypes.float32):
+        for weight_dtype in (dtypes.float16, dtypes.float32):
+          with self.subTest(rows=rows, dtype=dtype, weight_dtype=weight_dtype):
+            norm = nn.RMSNorm(64, eps=1e-6)
+            norm.weight = Tensor(rng.standard_normal(64, dtype=np.float32), device="CPU").cast(weight_dtype).realize()
+            x = Tensor(rng.standard_normal((rows, 64), dtype=np.float32), device="CPU").cast(dtype).realize()
+            tol = 5e-4 if dtype == dtypes.float16 else 2e-6
+            np.testing.assert_allclose(rmsnorm(norm, x).numpy(), norm(x).numpy(), rtol=tol, atol=tol)
+
+  def test_causal_conv_silu_matches_reference(self):
+    rng = np.random.default_rng(15)
+    batch, tokens, channels, kernel = 2, 7, 64, 4
+    for dtype in (dtypes.float16, dtypes.float32):
+      for weight_dtype in (dtypes.float16, dtypes.float32):
+        if dtype == weight_dtype == dtypes.float16: continue
+        with self.subTest(dtype=dtype, weight_dtype=weight_dtype):
+          state = Tensor(rng.standard_normal((batch, kernel - 1, channels), dtype=np.float32), device="CPU").realize()
+          x = Tensor(rng.standard_normal((batch, tokens, channels), dtype=np.float32), device="CPU").cast(dtype).realize()
+          weight = Tensor(rng.standard_normal((channels, kernel), dtype=np.float32),
+                          device="CPU").cast(weight_dtype).realize()
+          window = state.cat(x, dim=1)
+          expected = functools.reduce(lambda a,b: a+b, (window[:, i:i+tokens] * weight[:, i] for i in range(kernel))).silu()
+          np.testing.assert_allclose(causal_conv_silu(state, x, weight).numpy(), expected.numpy(), rtol=2e-6, atol=2e-6)
+          np.testing.assert_allclose(causal_conv_silu(state, x, weight.T.contiguous()).numpy(), expected.numpy(), rtol=2e-6, atol=2e-6)
+
+    # This is the vectorized path used by Qwen's prefill.
+    state = Tensor(rng.standard_normal((1, kernel - 1, channels), dtype=np.float32), device="CPU").realize()
+    x = Tensor(rng.standard_normal((1, tokens, channels), dtype=np.float32), device="CPU").half().realize()
+    weight = Tensor(rng.standard_normal((channels, kernel), dtype=np.float32), device="CPU").half().realize()
+    expected = causal_conv_silu(state, x, weight).numpy()
+    np.testing.assert_allclose(causal_conv_silu(state, x, weight.T.contiguous()).numpy(), expected, rtol=2e-6, atol=2e-6)
+
+  def test_shared_gate_matches_reference(self):
+    rng = np.random.default_rng(6)
+    for dtype in (dtypes.float16, dtypes.float32):
+      with self.subTest(dtype=dtype):
+        x = Tensor(rng.standard_normal((3, 64), dtype=np.float32), device="CPU").cast(dtype).realize()
+        weight = Tensor(rng.standard_normal(64, dtype=np.float32), device="CPU").half().realize()
+        expected = (x * weight).sum(axis=-1, keepdim=True).sigmoid()
+        np.testing.assert_allclose(shared_gate(x, weight).numpy(), expected.numpy(), rtol=2e-5, atol=2e-5)
+
+  def test_silu_mul_matches_reference(self):
+    rng = np.random.default_rng(10)
+    for dtype in (dtypes.float16, dtypes.float32):
+      with self.subTest(dtype=dtype):
+        gate = Tensor(rng.standard_normal((2, 64), dtype=np.float32), device="CPU").cast(dtype).realize()
+        up = Tensor(rng.standard_normal((2, 64), dtype=np.float32), device="CPU").cast(dtype).realize()
+        np.testing.assert_equal(silu(gate).numpy(), gate.silu().numpy())
+        np.testing.assert_allclose(silu_mul(gate, up).numpy(), (gate.silu() * up).numpy(), rtol=2e-5, atol=2e-5)
+    gate = Tensor(rng.standard_normal((2, 64), dtype=np.float32), device="CPU").half().realize()
+    up = Tensor(rng.standard_normal((2, 64), dtype=np.float32), device="CPU").realize()
+    np.testing.assert_equal(silu_mul(gate, up).numpy(), (gate.silu() * up).numpy())
+    gate = Tensor(rng.standard_normal(4096, dtype=np.float32), device="CPU").half().realize()
+    up = Tensor(rng.standard_normal(4096, dtype=np.float32), device="CPU").realize()
+    np.testing.assert_allclose(silu_mul(gate, up).numpy(), (gate.silu() * up).numpy(), rtol=2e-6, atol=2e-6)
+
+  def test_biased_topk_matches_reference(self):
+    rng = np.random.default_rng(7)
+    logits = Tensor(rng.standard_normal((1, 2, 256), dtype=np.float32), device="CPU").half().realize()
+    bias = Tensor(rng.standard_normal(256, dtype=np.float32), device="CPU").half().realize()
+    probs = logits.sigmoid()
+    _, expected_sel = pairwise_topk(probs + bias, 8)
+    expected = probs.gather(-1, expected_sel)
+    expected = expected / expected.sum(axis=-1, keepdim=True)
+    got, got_sel = biased_sigmoid_topk(logits, bias, 8, normalize=True)
+    np.testing.assert_equal(got_sel.numpy(), expected_sel.numpy().reshape(2, 8))
+    np.testing.assert_allclose(got.numpy(), expected.numpy().reshape(2, 8), rtol=5e-4, atol=5e-4)
+
+  def test_packed_linear_matches_q8_activation_reference(self):
+    rng = np.random.default_rng(1)
+    for ggml_type, in_features in ((8, 64), (14, 256)):
+      for tokens in (1, 3, 8):
+        with self.subTest(ggml_type=ggml_type, tokens=tokens):
+          out_features = 7
+          raw = random_packed(rng, ggml_type, out_features * in_features)
+          layer = Linear(in_features, out_features, bias=False)
+          layer.set_quantized(Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), ggml_type)
+          x = rng.standard_normal((tokens, in_features), dtype=np.float32)
+          got = layer(Tensor(x, device="CPU")).numpy()
+          weight = ggml_data_to_tensor(Tensor(raw), out_features * in_features, ggml_type).numpy().reshape(out_features, in_features)
+          np.testing.assert_allclose(got, q8_activation(x) @ weight.T, rtol=1e-5, atol=5e-4)
+
+  def test_q8_linear_pair_matches_separate(self):
+    rng = np.random.default_rng(11)
+    in_features = 64
+    layers = []
+    for out_features in (7, 11):
+      raw = random_packed(rng, 8, out_features * in_features)
+      layer = Linear(in_features, out_features, bias=False)
+      layer.set_quantized(Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), 8)
+      layers.append(layer)
+    for dtype in (dtypes.float16, dtypes.float32):
+      with self.subTest(dtype=dtype):
+        x = Tensor(rng.standard_normal((1, in_features), dtype=np.float32), device="CPU").cast(dtype).realize()
+        got = q8_linear_pair(*layers, x)
+        for paired,layer in zip(got, layers): np.testing.assert_allclose(paired.numpy(), layer(x).numpy(), rtol=2e-5, atol=2e-5)
+        for paired,layer in zip(q8_linear_pair(*layers, x.reshape(1, 1, in_features)), layers):
+          self.assertEqual(paired.shape, (1, 1, layer.out_features))
+          np.testing.assert_allclose(paired.numpy(), layer(x).numpy().reshape(1, 1, -1), rtol=2e-5, atol=2e-5)
+        for layer in layers: layer.cpu_repacked = q8_repack(layer.weight, layer.out_features, layer.in_features).realize()
+        for repacked,original in zip(q8_linear_pair(*layers, x), got): np.testing.assert_equal(repacked.numpy(), original.numpy())
+        original_weights = [layer.weight for layer in layers]
+        for layer in layers: layer.weight = Tensor.zeros_like(layer.weight).realize()
+        for repacked,original in zip(uop_q8_linear_pair(*layers, x), got):
+          np.testing.assert_allclose(repacked.numpy(), original.numpy(), rtol=2e-6, atol=1e-5)
+        for layer,weight in zip(layers, original_weights): layer.weight = weight
+        for layer in layers: layer.cpu_repacked = None
+
+  def test_large_q8_uop_linear_repacked_matches_raw(self):
+    rng = np.random.default_rng(20)
+    in_features, out_features = 1024, 7
+    raw = random_packed(rng, 8, out_features * in_features)
+    layer = Linear(in_features, out_features, bias=False)
+    layer.set_quantized(Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), 8)
+    x = Tensor(rng.standard_normal((1, in_features), dtype=np.float32), device="CPU").realize()
+    expected = uop_linear(layer, x).numpy()
+    layer.cpu_repacked = q8_repack(layer.weight, out_features, in_features).realize()
+    np.testing.assert_allclose(uop_linear(layer, x).numpy(), expected, rtol=2e-6, atol=1e-5)
+
+  def test_q8_batched_pair_matches_separate(self):
+    rng = np.random.default_rng(15)
+    in_features = 64
+    layers = []
+    for out_features in (7, 11):
+      raw = random_packed(rng, 8, out_features * in_features)
+      layer = Linear(in_features, out_features, bias=False)
+      layer.set_quantized(Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), 8)
+      layers.append(layer)
+    for dtype in (dtypes.float16, dtypes.float32):
+      with self.subTest(dtype=dtype):
+        x = Tensor(rng.standard_normal((2, 3, in_features), dtype=np.float32), device="CPU").cast(dtype).realize()
+        got = q8_batched_pair(*layers, x)
+        for paired,layer in zip(got, layers): np.testing.assert_allclose(paired.numpy(), layer(x).numpy(), rtol=2e-5, atol=2e-5)
+        for layer in layers: layer.cpu_repacked = q8_repack(layer.weight, layer.out_features, layer.in_features).realize()
+        for repacked,original in zip(q8_batched_pair(*layers, x), got): np.testing.assert_equal(repacked.numpy(), original.numpy())
+        for layer in layers: layer.cpu_repacked = None
+
+  def test_q8_silu_linear_matches_separate(self):
+    rng = np.random.default_rng(16)
+    in_features, out_features = 64, 11
+    layer = Linear(in_features, out_features, bias=False)
+    layer.set_quantized(Tensor(random_packed(rng, 8, out_features * in_features), dtype=dtypes.uint8, device="CPU").realize(), 8)
+    gate = Tensor(rng.standard_normal((2, 4, in_features), dtype=np.float32), device="CPU").half().realize()
+    up = Tensor(rng.standard_normal(gate.shape, dtype=np.float32), device="CPU").realize()
+    fused = q8_silu_linear(layer, gate, up)
+    separate = layer(silu_mul(gate, up).half())
+    linear = layer(gate[:, 0])
+    single = q8_silu_linear(layer, gate[:1, :1], up[:1, :1])
+    np.testing.assert_equal(fused.numpy(), separate.numpy())
+    layer.cpu_repacked = q8_repack(layer.weight, layer.out_features, layer.in_features).realize()
+    np.testing.assert_equal(q8_silu_linear(layer, gate, up).numpy(), fused.numpy())
+    np.testing.assert_equal(q8_silu_linear(layer, gate[:1, :1], up[:1, :1]).numpy(), single.numpy())
+    np.testing.assert_equal(layer(gate[:, 0]).numpy(), linear.numpy())
+
+  def test_q8_gdn_projections_match_separate(self):
+    rng = np.random.default_rng(10)
+    in_features = 256
+    layers = []
+    for out_features in (1, 1):
+      raw = random_packed(rng, 8, out_features * in_features)
+      layer = Linear(in_features, out_features, bias=False)
+      layer.set_quantized(Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), 8)
+      layers.append(layer)
+    x = Tensor(rng.standard_normal((1, in_features), dtype=np.float32), device="CPU").half().realize()
+    weight = Tensor(rng.standard_normal((16, in_features), dtype=np.float32), device="CPU").half().realize()
+    got = q8_gdn_projections(*layers, weight, x)
+    expected = (*q8_linear_pair(*layers, x), x @ weight.T)
+    for fused,separate in zip(got, expected): np.testing.assert_allclose(fused.numpy(), separate.numpy(), rtol=1e-3, atol=1e-3)
+    for layer in layers: layer.cpu_repacked = q8_repack(layer.weight, layer.out_features, layer.in_features).realize()
+    for repacked,original in zip(q8_gdn_projections(*layers, weight, x), got):
+      np.testing.assert_equal(repacked.numpy(), original.numpy())
+    norm = nn.RMSNorm(in_features, eps=1e-6)
+    norm.weight = Tensor(rng.standard_normal(in_features, dtype=np.float32), device="CPU").half().realize()
+    raw_x = Tensor(rng.standard_normal((1, in_features), dtype=np.float32), device="CPU").realize()
+    expected = q8_gdn_projections(*layers, weight, rmsnorm(norm, raw_x).half())
+    for fused,separate in zip(q8_gdn_norm_projections(*layers, weight, raw_x, norm), expected):
+      np.testing.assert_equal(fused.numpy(), separate.numpy())
+
+  def test_f16_linear_matches_standard(self):
+    rng = np.random.default_rng(13)
+    layer = Linear(256, 37, bias=False)
+    layer.weight = Tensor(rng.standard_normal((37, 256), dtype=np.float32), device="CPU").half().realize()
+    for dtype in (dtypes.float16, dtypes.float32):
+      for tokens in (1, 3):
+        with self.subTest(dtype=dtype, tokens=tokens):
+          x = Tensor(rng.standard_normal((tokens, 256), dtype=np.float32), device="CPU").cast(dtype).realize()
+          np.testing.assert_allclose(f16_linear(layer, x).numpy(), layer(x).numpy(), rtol=1e-5, atol=2e-5)
+          np.testing.assert_allclose(f16_matvec(x, layer.weight).numpy(), layer(x).numpy(), rtol=1e-5, atol=2e-5)
+          np.testing.assert_allclose(uop_f16_matvec(x, layer.weight).numpy(), layer(x).numpy(), rtol=1e-5, atol=2e-5)
+    norm = nn.RMSNorm(256, eps=1e-6)
+    norm.weight = Tensor(rng.standard_normal(256, dtype=np.float32), device="CPU").half().realize()
+    x = Tensor(rng.standard_normal((1, 256), dtype=np.float32), device="CPU").realize()
+    normalized, out = rmsnorm_f16_linear(norm, layer, x)
+    expected = rmsnorm(norm, x)
+    np.testing.assert_allclose(normalized.numpy(), expected.numpy(), rtol=2e-6, atol=2e-6)
+    np.testing.assert_allclose(out.numpy(), f16_linear(layer, expected).numpy(), rtol=1e-5, atol=2e-5)
+
+  def test_q6_argmax_matches_materialized_logits(self):
+    rng = np.random.default_rng(8)
+    in_features, out_features = 256, 37
+    raw = random_packed(rng, 14, out_features * in_features)
+    weight = ggml_data_to_tensor(Tensor(raw), out_features * in_features, 14).numpy().reshape(out_features, in_features)
+    layer = Linear(in_features, out_features, bias=False)
+    layer.set_quantized(Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), 14)
+    for _ in range(3):
+      x = Tensor(rng.standard_normal((1, in_features), dtype=np.float32), device="CPU").realize()
+      expected = int(np.argmax(weight @ q8k_activation(x.numpy()).reshape(-1)))
+      self.assertEqual(q6_argmax(layer, x).item(), expected)
+
+  def test_packed_experts_match_reference(self):
+    rng = np.random.default_rng(2)
+    num_experts, in_features, out_features = 2, 256, 5
+    sel, x = np.array([1, 0, 1, 0], dtype=np.int32), rng.standard_normal((4, in_features), dtype=np.float32)
+    for ggml_type in (14, 21, 23):
+      with self.subTest(ggml_type=ggml_type):
+        raw = random_packed(rng, ggml_type, num_experts * out_features * in_features)
+        weight = ggml_data_to_tensor(Tensor(raw), num_experts * out_features * in_features,
+                                     ggml_type).numpy().reshape(num_experts, out_features, in_features)
+        experts = ExpertWeights(num_experts, in_features, out_features)
+        experts.set_quantized(Tensor(weight), Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), ggml_type)
+        got = experts(Tensor(sel, device="CPU"), Tensor(x, device="CPU")).numpy()
+        activation = q8k_activation(x) if ggml_type in (21, 23) else q8_activation(x)
+        expected = np.stack([activation[i] @ weight[expert].T for i,expert in enumerate(sel)])
+        np.testing.assert_allclose(got, expected, rtol=1e-5, atol=5e-4)
+        if ggml_type == 21:
+          experts.cpu_repacked = iq3_repack(experts.weight, num_experts * out_features, in_features).realize()
+          np.testing.assert_allclose(experts(Tensor(sel, device="CPU"), Tensor(x, device="CPU")).numpy(),
+                                     expected, rtol=1e-5, atol=5e-4)
+        direct_sel = np.array([1, 0], dtype=np.int32)
+        direct = experts(Tensor(direct_sel, device="CPU"), Tensor(x[:1], device="CPU")).numpy()
+        direct_activation = q8k_activation(x[:1]) if ggml_type in (21, 23) else q8_activation(x[:1])
+        direct_expected = np.stack([direct_activation[0] @ weight[expert].T for expert in direct_sel])
+        np.testing.assert_allclose(direct, direct_expected, rtol=1e-5, atol=5e-4)
+  def test_weighted_expert_sum_matches_reference(self):
+    rng = np.random.default_rng(16)
+    x = rng.standard_normal((2, 4, 257), dtype=np.float32)
+    probs = rng.random((2, 4), dtype=np.float32)
+    got = weighted_sum(Tensor(x, device="CPU"), Tensor(probs, device="CPU")).numpy()
+    np.testing.assert_allclose(got, (x * probs[..., None]).sum(axis=1), rtol=1e-6, atol=1e-6)
+
+  def test_quantized_expert_weighted_sum_matches_separate(self):
+    rng = np.random.default_rng(29)
+    num_experts, inputs, routes_per_input, in_features, out_features = 16, 4, 8, 256, 64
+    sel = Tensor(rng.integers(0, num_experts, (1, inputs, routes_per_input), dtype=np.int32), device="CPU").realize()
+    x = Tensor(rng.standard_normal((*sel.shape, in_features), dtype=np.float32), device="CPU").realize()
+    probs = Tensor(rng.random(sel.shape, dtype=np.float32), device="CPU").realize()
+    for ggml_type in (14, 23):
+      with self.subTest(ggml_type=ggml_type):
+        layer = ExpertWeights(num_experts, in_features, out_features)
+        layer.set_quantized(Tensor.empty(num_experts, out_features, in_features),
+                            Tensor(random_packed(rng, ggml_type, num_experts * in_features * out_features),
+                                   dtype=dtypes.uint8, device="CPU").realize(), ggml_type)
+        np.testing.assert_allclose(expert_weighted_sum(layer, sel, x, probs).numpy(), weighted_sum(layer(sel, x), probs).numpy(),
+                                   rtol=5e-6, atol=1e-5)
+
+  def test_fused_expert_silu_matches_separate(self):
+    rng = np.random.default_rng(12)
+    num_experts, in_features, out_features = 3, 256, 7
+    sel = Tensor(np.array([2, 0], dtype=np.int32), device="CPU")
+    x = Tensor(rng.standard_normal((1, in_features), dtype=np.float32), device="CPU").realize()
+    for ggml_type in (14, 21, 23):
+      with self.subTest(ggml_type=ggml_type):
+        experts = []
+        for _ in range(2):
+          raw = random_packed(rng, ggml_type, num_experts * out_features * in_features)
+          weight = ggml_data_to_tensor(Tensor(raw), num_experts * out_features * in_features,
+                                       ggml_type).reshape(num_experts, out_features, in_features)
+          expert = ExpertWeights(num_experts, in_features, out_features)
+          expert.set_quantized(weight, Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), ggml_type)
+          experts.append(expert)
+        gate, up = expert_pair(*experts, sel, x)
+        np.testing.assert_allclose(expert_silu(*experts, sel, x).numpy(), silu_mul(gate, up).numpy(), rtol=1e-4, atol=1e-4)
+        if ggml_type == 21:
+          direct_expected = expert_silu(*experts, sel, x).numpy()
+          batch_sel = Tensor(np.array([[2, 2], [1, 2]], dtype=np.int32), device="CPU")
+          batch_x = Tensor(rng.standard_normal((2, in_features), dtype=np.float32), device="CPU").realize()
+          expected = expert_silu(*experts, batch_sel, batch_x).numpy()
+          batch_gate, batch_up = expert_pair(*experts, batch_sel, batch_x)
+          np.testing.assert_allclose(expected, silu_mul(batch_gate, batch_up).numpy(), rtol=1e-4, atol=1e-4)
+          for expert in experts:
+            expert.cpu_repacked = iq3_repack(expert.weight, expert.num_experts * expert.out_features, expert.in_features).realize()
+          np.testing.assert_allclose(expert_silu(*experts, sel, x).numpy(), direct_expected, rtol=1e-5, atol=1e-4)
+          np.testing.assert_allclose(expert_silu(*experts, batch_sel, batch_x).numpy(), expected, rtol=1e-5, atol=1e-4)
+
+  def test_expert_silu_weighted_reuses_routes(self):
+    rng = np.random.default_rng(34)
+    num_experts, in_features, hidden, out_features = 3, 256, 256, 64
+    layers = []
+    for ggml_type,layer_in,layer_out in ((21, in_features, hidden), (21, in_features, hidden), (23, hidden, out_features)):
+      layer = ExpertWeights(num_experts, layer_in, layer_out)
+      raw = random_packed(rng, ggml_type, num_experts * layer_in * layer_out)
+      weight = ggml_data_to_tensor(Tensor(raw), num_experts * layer_in * layer_out, ggml_type).reshape(
+        num_experts, layer_out, layer_in)
+      layer.set_quantized(weight, Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), ggml_type)
+      if ggml_type == 21: layer.cpu_repacked = iq3_repack(layer.weight, num_experts * layer_out, layer_in).realize()
+      layers.append(layer)
+    sel = Tensor(np.array([[2, 0], [1, 2]], dtype=np.int32), device="CPU")
+    x = Tensor(rng.standard_normal((2, in_features), dtype=np.float32), device="CPU").realize()
+    probs = Tensor(rng.random(sel.shape, dtype=np.float32), device="CPU").realize()
+    expected = expert_weighted_sum(layers[2], sel, expert_silu(layers[0], layers[1], sel, x), probs).numpy()
+    got = uop_expert_silu_weighted(layers[0], layers[1], layers[2], sel, x, probs).numpy()
+    np.testing.assert_allclose(got, expected, rtol=5e-6, atol=1e-5)
+
+  def test_fused_moe_matches_separate_quantized_layers(self):
+    rng = np.random.default_rng(15)
+    dim = hidden = 256
+    config = TransformerConfig(1, dim, hidden, 1, 1, 1e-6, 32, dim, 1e6, dim, dim,
+                               num_experts=3, num_experts_per_tok=2, shared_expert_dim=hidden)
+    block = FFNBlock(config)
+    routed_weights = {}
+    for name,expert,ggml_type in (("gate", block.ffn_gate_exps, 21), ("up", block.ffn_up_exps, 21),
+                                  ("down", block.ffn_down_exps, 23)):
+      elements = expert.num_experts * expert.in_features * expert.out_features
+      raw = random_packed(rng, ggml_type, elements)
+      weight = ggml_data_to_tensor(Tensor(raw), elements, ggml_type).reshape(
+        expert.num_experts, expert.out_features, expert.in_features)
+      routed_weights[name] = weight.numpy().astype(np.float32)
+      expert.set_quantized(weight, Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), ggml_type)
+      if ggml_type == 21: expert.cpu_repacked = iq3_repack(expert.weight, expert.num_experts * expert.out_features, expert.in_features).realize()
+    for layer in (block.ffn_gate_shexp, block.ffn_up_shexp, block.ffn_down_shexp):
+      elements = layer.in_features * layer.out_features
+      raw = random_packed(rng, 8, elements)
+      layer.set_quantized(Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), 8)
+      layer.cpu_repacked = q8_repack(layer.weight, layer.out_features, layer.in_features).realize()
+    block.ffn_gate_inp_shexp["weight"] = Tensor(rng.standard_normal(dim, dtype=np.float32), device="CPU").half().realize()
+
+    x = Tensor(rng.standard_normal((1, 2, dim), dtype=np.float32), device="CPU").realize()
+    probs = Tensor(np.array([[[0.7, 0.3], [0.4, 0.6]]], dtype=np.float32), device="CPU").realize()
+    sel = Tensor(np.array([[[2, 0], [1, 2]]], dtype=np.int32), device="CPU").realize()
+    selected = sel.numpy().reshape(-1)
+    quantized_x = q8k_activation(x.numpy()).reshape(2, dim)
+    gate = np.stack([routed_weights["gate"][expert] @ quantized_x[route // 2] for route,expert in enumerate(selected)])
+    up = np.stack([routed_weights["up"][expert] @ quantized_x[route // 2] for route,expert in enumerate(selected)])
+    routed_hidden = silu_mul(Tensor(gate, device="CPU"), Tensor(up, device="CPU")).numpy()
+    quantized_hidden = q8k_activation(routed_hidden)
+    routed = np.stack([routed_weights["down"][expert] @ quantized_hidden[route] for route,expert in enumerate(selected)])
+    routed = Tensor((routed.reshape(2, 2, dim) * probs.numpy().reshape(2, 2, 1)).sum(1).reshape(1, 2, dim), device="CPU")
+    shared_gate_out, shared_up = block.ffn_gate_shexp(x), block.ffn_up_shexp(x)
+    shared = block.ffn_down_shexp(silu_mul(shared_gate_out, shared_up))
+    expected = routed + shared * shared_gate(x, block.ffn_gate_inp_shexp["weight"])
+    original = moe_ffn(block, x, probs, sel).numpy()
+    np.testing.assert_allclose(original, expected.numpy(), rtol=1e-4, atol=2e-2)
+    for expert in (block.ffn_gate_exps, block.ffn_up_exps):
+      expert.cpu_repacked = iq3_repack(expert.weight, expert.num_experts * expert.out_features, expert.in_features).realize()
+    np.testing.assert_equal(moe_ffn(block, x, probs, sel).numpy(), original)
+    np.testing.assert_allclose(uop_moe_ffn(block, x[:, :1], probs[:, :1], sel[:, :1]).numpy(), original[:, :1],
+                               rtol=1e-4, atol=2e-2)
+
+  def test_fused_moe_q6_down_matches_separate(self):
+    rng = np.random.default_rng(16)
+    dim = hidden = 256
+    config = TransformerConfig(1, dim, hidden, 1, 1, 1e-6, 32, dim, 1e6, dim, dim,
+                               num_experts=3, num_experts_per_tok=2, shared_expert_dim=hidden)
+    block = FFNBlock(config)
+    for expert,ggml_type in ((block.ffn_gate_exps, 21), (block.ffn_up_exps, 21), (block.ffn_down_exps, 14)):
+      elements = expert.num_experts * expert.in_features * expert.out_features
+      raw = random_packed(rng, ggml_type, elements)
+      weight = ggml_data_to_tensor(Tensor(raw), elements, ggml_type).reshape(
+        expert.num_experts, expert.out_features, expert.in_features)
+      expert.set_quantized(weight, Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), ggml_type)
+    for expert in (block.ffn_gate_exps, block.ffn_up_exps):
+      expert.cpu_repacked = iq3_repack(expert.weight, expert.num_experts * expert.out_features, expert.in_features).realize()
+    for layer in (block.ffn_gate_shexp, block.ffn_up_shexp, block.ffn_down_shexp):
+      elements = layer.in_features * layer.out_features
+      raw = random_packed(rng, 8, elements)
+      layer.set_quantized(Tensor(raw, dtype=dtypes.uint8, device="CPU").realize(), 8)
+    block.ffn_gate_inp_shexp["weight"] = Tensor(rng.standard_normal(dim, dtype=np.float32), device="CPU").half().realize()
+
+    x = Tensor(rng.standard_normal((1, 1, dim), dtype=np.float32), device="CPU").realize()
+    probs = Tensor(np.array([[[0.7, 0.3]]], dtype=np.float32), device="CPU").realize()
+    sel = Tensor(np.array([[[2, 0]]], dtype=np.int32), device="CPU").realize()
+    hidden = expert_silu(block.ffn_gate_exps, block.ffn_up_exps, sel, x.unsqueeze(2))
+    routed = weighted_sum(block.ffn_down_exps(sel, hidden), probs)
+    gate, up = block.ffn_gate_shexp(x), block.ffn_up_shexp(x)
+    shared = block.ffn_down_shexp(silu_mul(gate, up)) * shared_gate(x, block.ffn_gate_inp_shexp["weight"])
+    np.testing.assert_allclose(moe_ffn(block, x, probs, sel).numpy(), (routed + shared).numpy(), rtol=3e-4, atol=1e-4)
+
+if __name__ == "__main__":
+  unittest.main()

test/unit/test_llm_server.py

@@ -3,10 +3,26 @@ from unittest.mock import patch
 from tinygrad import Tensor, UOp
 from tinygrad.schedule import schedule_cache
 from tinygrad.llm.model import Transformer, TransformerConfig
+from tinygrad.llm.serve import StreamRouter
 
 TEST_CONFIG = TransformerConfig(num_blocks=1, dim=64, hidden_dim=128, n_heads=2, n_kv_heads=2,
                            norm_eps=1e-5, vocab_size=100, head_dim=32, rope_theta=10000.0, rope_dim=32, v_head_dim=32, max_context=32)
 
+class TestStreamRouter(unittest.TestCase):
+  @staticmethod
+  def route(router:StreamRouter, *pieces:str) -> dict[str, str]:
+    routed = [x for piece in pieces for x in router.route(piece)]
+    routed += list(router.route("", final=True))
+    return {field:"".join(text for f, text in routed if f == field) for field in {x[0] for x in routed}}
+
+  def test_generated_reasoning_tag(self):
+    self.assertEqual(self.route(StreamRouter(), "<thi", "nk>reason", "</thi", "nk>answer"),
+                     {"reasoning_content":"reason", "content":"answer"})
+
+  def test_prompt_opened_reasoning(self):
+    self.assertEqual(self.route(StreamRouter(reasoning=True), "reason", "</thi", "nk>answer"),
+                     {"reasoning_content":"reason", "content":"answer"})
+
 class TestTransformerGenerate(unittest.TestCase):
   def test_kv_cache_reuse(self):
     """Test that generate reuses the KV cache when tokens extend the cached prefix."""
@@ -151,8 +167,9 @@ class TestTransformerGenerate(unittest.TestCase):
     """Temperature from generate should be passed through to __call__."""
     model = Transformer(TEST_CONFIG)
     captured_temps = []
-    def mock_call(self, tokens, start_pos, temperature):
+    def mock_call(_self, tokens, start_pos, temperature, **kwargs):
       captured_temps.append(float(temperature.item()))
+      self.assertTrue(kwargs["sample"])
       return Tensor([[42]])
     with patch.object(Transformer, '__call__', mock_call):
       gen = model.generate([1, 2, 3], temperature=0.6)

test/unit/test_multitensor.py

@@ -79,6 +79,10 @@ class TestMultiTensor(unittest.TestCase):
     with self.assertRaises(RuntimeError):
       X.shard_(devices_3, 0)
 
+  def test_shard_reshape_cross_boundary(self):
+    X = Tensor.ones(5, 4).contiguous().realize().shard(devices_2, 1)
+    with self.assertRaises(RuntimeError): X.reshape(10, 2).uop.axis
+
   def test_tensor_from_multi(self):
     X = Tensor([1, 2], dtype=dtypes.int).shard_(devices_2, 0)
     Y = Tensor(X.uop)
@@ -123,6 +127,28 @@ class TestMultiTensor(unittest.TestCase):
     with Context(RING=use_ring):
       np.testing.assert_equal(t.shard(devices_2, axis=axis).sum().item(), 10)
 
+  def test_allreduce_cast_half(self, assign=False, kernel_count=8):
+    devices = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
+    a_src = Tensor.arange(2*3, dtype=dtypes.half).reshape(2, 3).clone().realize()
+    b_src = Tensor.arange(2*3, dtype=dtypes.half).reshape(2, 3).clone().realize()
+    a = a_src.shard(devices, axis=0).realize()
+    b = b_src.shard(devices, axis=0).realize()
+    # assigning creates a copy of the output before allreduce
+    if assign:
+      tst = Tensor.empty_like(b)
+      tst.assign(a + b)
+    else:
+      tst = a + b
+    tst = tst.float().sum(0)
+    GlobalCounters.reset()
+    with Context(ALLREDUCE_CAST=1, RING=0, ALL2ALL=0):
+      tst.realize()
+    kernel_count = GlobalCounters.kernel_count
+    np.testing.assert_allclose(tst.numpy(), (a_src.numpy()+b_src.numpy()).sum(0))
+    self.assertEqual(kernel_count, kernel_count)
+
+  def test_allreduce_cast_half_assign(self): self.test_allreduce_cast_half(assign=True, kernel_count=10)
+
   def test_multiple_to_single_device(self):
     kernel_counts = {}
     for ring in (0, 2):
@@ -571,7 +597,7 @@ class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
     t = Tensor.arange(64).reshape(8, 8).clone().realize()
     t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0)
 
-    with self.assertRaises(AssertionError):
+    with self.assertRaises(RuntimeError):
       # sharded axis shrink on non-device boundry is not allowed
       a = t.shrink(((0, 3), (0, 8))).contiguous()
       a.schedule_linear()

test/unit/test_realize_is_realize.py

@@ -30,7 +30,7 @@ class TestRealizeIsRealized(unittest.TestCase):
   def test_multi(self):
     d = Device.DEFAULT
     t = Tensor.ones(8).contiguous().shard((d, d), axis=0).realize()
-    assert all(u.is_realized for u in t.uop.src)
+    assert t.uop.src[0].is_realized
 
   def test_empty(self):
     t = Tensor.empty(4, 4).realize()

tinygrad/callify.py

@@ -38,8 +38,8 @@ add_tags = PatternMatcher([
 ])
 
 def replace_contig_with_store_after(u:UOp):
-  # can't allocate a buffer without a device (e.g., inside a CALL function body with only PARAMs)
-  if u.device is None: return None
+  # can't allocate a buffer for a virtual value
+  if u.is_virtual: return None
   # if size is 0, remove the contig
   if 0 in u.shape: return u.src[0]
   # no real contig for DISK/TINYFS tensors, they are left alone
@@ -49,7 +49,7 @@ def replace_contig_with_store_after(u:UOp):
 
 def replace_store_after_with_contig(u:UOp, src:UOp):
   assigned_to = u
-  while assigned_to.op in {Ops.BITCAST, Ops.AFTER, Ops.MULTI}: assigned_to = assigned_to.src[0].base
+  while assigned_to.op in {Ops.BITCAST, Ops.AFTER, Ops.UNSHARD}: assigned_to = assigned_to.src[0].base
   if assigned_to.op not in {Ops.BUFFER, Ops.SLICE}: return src.contiguous(tag=u.tag)
 
 def _make_buffer_view(src:UOp) -> UOp|None:
@@ -61,29 +61,29 @@ def _make_buffer_view(src:UOp) -> UOp|None:
     buf = buf.src[0]
     if byte_offset % buf.dtype.itemsize != 0: return None
     offset = byte_offset // buf.dtype.itemsize
-  return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(dtypes.weakint, offset)), src.numel())
+  return UOp(Ops.SLICE, src.dtype, (buf, UOp.const(None, offset)), src.numel())
 
 def contiguous_mops_to_view(c:UOp, src:UOp):
   """MOPS(BUFFER) → SLICE when movement ops collapse to a contiguous range."""
   buf = src.base
-  if buf.op not in {Ops.BUFFER, Ops.SLICE, Ops.MULTI}: return None
+  if buf.op not in {Ops.BUFFER, Ops.SLICE, Ops.UNSHARD}: return None
   if src.op is Ops.RESHAPE and src.src[0].op in {Ops.BUFFER, Ops.SLICE} and c.op is not Ops.BITCAST: return None
   if c.op is not Ops.BITCAST and src.op is Ops.BUFFER: return None
 
   # no symbolic shape
   if not all_int(c.shape): return None
 
-  if buf.op is not Ops.MULTI and (view := _make_buffer_view(src)) is not None:
+  if buf.op is not Ops.UNSHARD and (view := _make_buffer_view(src)) is not None:
     view = (view.replace(dtype=c.dtype, arg=c.numel()) if c.op is Ops.BITCAST else view).reshape(c.shape)
     return c.replace(src=(view,)) if c.op is Ops.COPY else view
 
-  # for MULTI tensors, use multi_pm to resolve per-shard movement ops, then create SLICE on the resolved result
+  # for UNSHARD tensors, use multi_pm to resolve per-shard movement ops, then create SLICE on the resolved result
   if not isinstance(c.device, str):
     from tinygrad.schedule.multi import multi_pm
     resolved = graph_rewrite(src, multi_pm, name="multi_buffer_view")
-    if resolved.op is not Ops.MULTI: return None
+    if resolved.op is not Ops.UNSHARD: return None
     if (view := _make_buffer_view(resolved.src[0])) is None: return None
-    return view.reshape(resolved.src[0].shape).multi(resolved.arg).contiguous(tag=c.tag)
+    return view.reshape(resolved.src[0].shape).unshard(resolved.arg, resolved.src[1:]).contiguous(tag=c.tag)
 
   return None
 
@@ -92,7 +92,7 @@ def _precompiled_output_redirect(s:UOp, t:UOp) -> UOp|None:
   # materialize straight into t
   if s.op is Ops.CONTIGUOUS: return t.after(t.store(s.src[0]))
   # rebind output storage to t
-  if s.op in {Ops.BUFFER, Ops.MULTI} and s.has_buffer_identity(): return t
+  if s.op in {Ops.BUFFER, Ops.UNSHARD} and s.has_buffer_identity(): return t
   return None
 
 def transform_precompiled_call(c:UOp) -> UOp|None:
@@ -203,7 +203,7 @@ def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
   if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Tensor Graph")
   # uop list is a list in the original_sink graph and we can map to the tags later
   # same predicate as Tensor.realize
-  ctx = AllocCtx(bases={base for x in big_sink.src if (base:=x.base).device is not None and not base.has_buffer_identity()
+  ctx = AllocCtx(bases={base for x in big_sink.src if not (base:=x.base).is_virtual and not base.has_buffer_identity()
                         and base.op is not Ops.AFTER and base.addrspace is not AddrSpace.ALU})
 
   # this rewrite is "read-only", it adds simple things to buffer_map and may sink things on big_sink, bottom_up

tinygrad/codegen/__init__.py

@@ -1,7 +1,7 @@
 from dataclasses import replace, dataclass
 import itertools, functools
 from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC
-from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, panic
+from tinygrad.helpers import ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT, TracingKey, Context, panic
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo, GroupOp
 from tinygrad.uop.ops import AxisType
 from tinygrad.uop.render import pyrender
@@ -21,6 +21,7 @@ from tinygrad.codegen.late.coalesce import indexing_simplify
 from tinygrad.codegen.opt.postrange import apply_opts
 from tinygrad.codegen.late.gater import pm_move_gates_from_index
 from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_split_ranges, pm_load_collapse
+from tinygrad.schedule.multi import multi_pm
 from tinygrad.schedule.rangeify import pm_mops
 from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize
 from tinygrad.codegen.late.regalloc import LinearScanRegallocContext, pm_regalloc_rewrite
@@ -119,13 +120,22 @@ pm_expand_broadcast = pm_wmma_add+PatternMatcher([
   (UPat(Ops.WMMA, name="b"), broadcast_and_devec_wmma),
 ])
 
-def do_devectorize(b:UOp):
-  if b.shape == (): return None
+@functools.cache
+def _uses_shrink_memory(x:UOp) -> bool:
+  if x.op in (Ops.LOAD, Ops.STORE): return x.src[0].op is Ops.SHRINK
+  if x.op is Ops.AFTER: return _uses_shrink_memory(x.src[0])
+  if x.op in GroupOp.Elementwise or x.op in (Ops.STACK, Ops.RESHAPE, Ops.PERMUTE):
+    return any(_uses_shrink_memory(y) for y in x.src)
+  return False
+
+def do_devectorize(ctx:Renderer|tuple[dict, Renderer, set[UOp]], b:UOp):
+  preserved = ctx[2] if isinstance(ctx, tuple) and len(ctx) == 3 else set()
+  if b.shape == () or b in preserved or _uses_shrink_memory(b): return None
   # broadcasting needs to be already unpacked, Invalid matches any dtype and shape
   if not all(x.shape == b.shape or x.base.arg is Invalid for x in b.src): return None
   src = []
   for idx in itertools.product(*[range(x) for x in b.shape]):
-    idx_c = [UOp.const(dtypes.weakint, i) for i in idx]
+    idx_c = [UOp.const(None, i) for i in idx]
     src.append(b.replace(src=tuple(x.base if x.base.arg is Invalid else x.index(*idx_c) for x in b.src)))
   return UOp.stack(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src)
 
@@ -135,7 +145,7 @@ def do_stack_wmma(u:UOp):
   src = []
   for b in u.src:
     if b.op != Ops.STACK:
-      src.append(UOp.stack(*[b.index(UOp.const(dtypes.weakint, i)) for i in range(b.max_numel())]))
+      src.append(UOp.stack(*[b.index(i) for i in range(b.max_numel())]))
     else:
       src.append(b)
   return u.replace(src=tuple(src))
@@ -161,7 +171,7 @@ devectorizer2 = mop_cleanup+pm_mops+PatternMatcher([
   # RESHAPE a void is removed (hack for AFTER)
   (UPat(Ops.RESHAPE, dtype=dtypes.void, name="x"), lambda x: x.src[0]),
   # reshape of a single element shaped value to scalar is an index
-  (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0].index(UOp.const(dtypes.weakint, 0)) if x.marg == () and x.src[0].shape == (1,) else None),
+  (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0].index(0) if x.marg == () and x.src[0].shape == (1,) else None),
   # EXPAND on scalar -> STACK
   (UPat(Ops.EXPAND, src=(UPat.var("x"), UPat()), name="out"),
    lambda x,out: UOp.stack(*([x]*out.max_numel())) if x.shape == () and out.shape == (out.max_numel(),) else None),
@@ -186,6 +196,7 @@ def fix_group_for_reduce(x:UOp):
 @dataclass
 class ReduceContext:
   acc_num: int = 0
+  renderer: Renderer|None = None
 
 def merge_reduce_ends(sink:UOp):
   # merge ENDs that share the same range and nesting context (only those created by reduce_to_acc)
@@ -219,7 +230,8 @@ def reduce_ranges_to_acc(ctx:ReduceContext, r:UOp):
   acc_out = acc_initted.store(acc_initted.alu(r.arg[0], inp)).end(*r.src[1:]).rtag("mergeable")
   return acc.after(acc_out)
 
-def expand_horizontal_reduce(r:UOp):
+def expand_horizontal_reduce(ctx:ReduceContext, r:UOp):
+  if ctx.renderer is not None and ctx.renderer.has_native_reduce(r): return None
   inp = r.src[0]
   vals = [inp.index(*idx) for idx in itertools.product(*[range(inp.max_shape[a]) for a in range(r.arg[1])])]
   return functools.reduce(lambda x,y: x.alu(r.arg[0], y), vals)
@@ -242,7 +254,7 @@ pm_add_loads = PatternMatcher([
 
 def add_local_buffer(ctx, x:UOp):
   buf = UOp.placeholder(x.max_shape, x.dtype, slot=next(ctx), addrspace=x.arg.addrspace)
-  return buf.after(buf.index(*x.src[1:]).store(x.src[0]).end(*x.src[1:]).barrier())
+  return buf.after(buf.index(*x.src[1:]).store(x.src[0]).end(*x.src[1:]))
 
 pm_add_local_buffers = PatternMatcher([
   (UPat(Ops.STAGE, name="x"), add_local_buffer),
@@ -255,13 +267,42 @@ pm_cast_float_alu = PatternMatcher([
    lambda u,x: u.replace(src=(x.cast(u.dtype),)) if x.dtype != u.dtype else None),
 ])
 
+def _is_local_store(x:UOp): return x.op is Ops.STORE and x.addrspace is AddrSpace.LOCAL
+
+def add_raw_barrier(after:UOp):
+  # loads from a LOCAL buffer that depend (via AFTER) on stores to LOCAL memory need a workgroup barrier
+  if after.addrspace is not AddrSpace.LOCAL: return None
+  # one toposort over all the deps
+  deps = UOp.sink(*after.src[1:]).toposort(gate=lambda x: x.op is not Ops.BARRIER)
+  if not any(_is_local_store(x) for x in deps): return None
+  return after.src[0].after(UOp(Ops.BARRIER, src=after.src[1:]))
+
+def add_war_barrier(end:UOp):
+  # a LOCAL buffer stored and loaded in the same loop needs a barrier at the end of the loop body
+  rngs = [r for r in end.src[1:] if r.op is Ops.RANGE and r.arg[1] in (AxisType.REDUCE, AxisType.WEAK, AxisType.LOOP) and r.vmax > 0]
+  if not rngs or end.src[0].op is Ops.BARRIER: return None
+  sl = end.src[0].backward_slice_with_self
+  # only stores that are inside this loop body (not in the backward slice through AFTER chains from other loops)
+  store_bufs = {x.buf_uop for x in sl if _is_local_store(x) and any(r in x.ranges for r in rngs)}
+  # a load whose buffer matches a local store's buffer is necessarily a local load
+  if not (loads:=[x for x in sl if x.op is Ops.LOAD and x.src[0].buf_uop in store_bufs]): return None
+  return end.replace(src=(UOp(Ops.BARRIER, src=(end.src[0], *loads)),)+end.src[1:])
+
+pm_implicit_barriers = PatternMatcher([
+  (UPat(Ops.AFTER, name="after"), add_raw_barrier),
+  (UPat(Ops.END, name="end"), add_war_barrier),
+])
+
 def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
   if DEBUG >= 5: print(pyrender(ast))
   if SPEC: type_verify(ast, spec_tensor)
 
+  # resolve UNSHARDs (multi-device UNSHARDs are already resolved by the scheduler; this handles in-kernel shards, e.g. fragments)
+  sink = graph_rewrite(ast, multi_pm, name="multi_pm")
+
   # preprocess
-  sink = graph_rewrite(ast, pm_mops, name="early movement ops", bottom_up=True)
+  sink = graph_rewrite(sink, pm_mops, name="early movement ops", bottom_up=True)
 
   # first we optimize
   if optimize:
@@ -287,7 +328,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   sink = graph_rewrite(sink, expander2, ctx=build_range_map(sink), name="expander")
 
   # remove reduce
-  sink = graph_rewrite(sink, mop_cleanup+pm_reduce_local, ctx=ReduceContext(), name="remove reduces")
+  sink = graph_rewrite(sink, mop_cleanup+pm_reduce_local, ctx=ReduceContext(renderer=ren), name="remove reduces")
 
   # add locals
   sink = graph_rewrite(sink, pm_add_local_buffers, ctx=itertools.count(0), name="add local buffers")
@@ -300,7 +341,12 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   sink = graph_rewrite(sink, symbolic_simple+pm_expand_broadcast+pm_add_loads, name="*** expand broadcast / add loads")
 
   # devectorize
-  sink = graph_rewrite(sink, symbolic_simple+devectorizer2, ctx=ren, name="devectorize2")
+  native_reduce_uops = {u for r in sink.toposort() if ren.has_native_reduce(r) for u in r.src[0].backward_slice_with_self}
+  if native_reduce_uops:
+    sink = graph_rewrite(sink, symbolic_simple, name="pre-devectorize symbolic")
+    native_reduce_uops = {u for r in sink.toposort() if ren.has_native_reduce(r) for u in r.src[0].backward_slice_with_self}
+  sink = graph_rewrite(sink, devectorizer2 if native_reduce_uops else symbolic_simple+devectorizer2,
+                       ctx=({}, ren, native_reduce_uops) if native_reduce_uops else ren, name="devectorize2")
 
   # simplify indexing
   sink = graph_rewrite(sink, indexing_simplify, name="simplify load/store indexing")
@@ -310,7 +356,12 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
 
   # do memory coalescing (late)
   sink = memory_coalescing(sink, ren)
-  sink = graph_rewrite(sink, symbolic_simple+ew_devectorizer+pm_simplify_add_image, name="add images", ctx=({}, ren), bottom_up=True)
+  native_reduce_uops = {u for r in sink.toposort() if ren.has_native_reduce(r) for u in r.src[0].backward_slice_with_self}
+  if native_reduce_uops:
+    sink = graph_rewrite(sink, symbolic_simple, name="pre-image symbolic", bottom_up=True)
+    native_reduce_uops = {u for r in sink.toposort() if ren.has_native_reduce(r) for u in r.src[0].backward_slice_with_self}
+  sink = graph_rewrite(sink, (ew_devectorizer if native_reduce_uops else symbolic_simple+ew_devectorizer)+pm_simplify_add_image,
+                       name="add images", ctx=({}, ren, native_reduce_uops), bottom_up=True)
 
   # extra symbolic before decomp. crashes without this?
   sink = graph_rewrite(sink, sym, name="extra symbolic")
@@ -344,6 +395,9 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   pm_final_rewrite = pm_decomp+extra_matcher+pm_split_ends
   sink = graph_rewrite(sink, pm_final_rewrite+pm_remove_invalid, ctx=ren, name="final rewrite")
 
+  # add implicit barriers (stores/loads through LOCAL memory ordered by AFTER or across loop iterations need workgroup barriers)
+  sink = graph_rewrite(sink, pm_implicit_barriers, name="add implicit barriers")
+
   # this was the linearizer
   sink = graph_rewrite(sink, pm_add_control_flow, ctx=CFGContext(sink), name="add control flow", bottom_up=True)
 
@@ -435,7 +489,7 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
   if ast.op is Ops.PROGRAM: prg = ast
   elif ast.op is Ops.SINK:
     assert isinstance(ast.arg, KernelInfo), "requires KernelInfo on arg to to_program"
-    full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None)
+    full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None and ast.arg.optimize)
     prog_info = ProgramInfo.from_sink(full_sink, renderer.target)
     # instruction selection
     if isinstance(renderer, ISARenderer):
@@ -449,8 +503,15 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
   return prg
 
 to_program_cache: dict[tuple, UOp] = {}
+def program_cache_key(ast:UOp, renderer:Renderer) -> tuple:
+  config = (NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT)
+  return (ast.key, type(renderer), renderer.target, *[x.value for x in config])
+
+def parallel_to_program(args:tuple[UOp, Renderer, tuple]) -> tuple[tuple, UOp]:
+  ast, renderer, key = args
+  return key, do_to_program(ast, renderer)
+
 def to_program(ast:UOp, renderer:Renderer) -> UOp:
-  config = (NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32)
-  key = (ast.key, type(renderer), renderer.target, *[x.value for x in config])
+  key = program_cache_key(ast, renderer)
   if (prg:=to_program_cache.get(key)) is None: to_program_cache[key] = prg = do_to_program(ast, renderer)
   return prg

tinygrad/codegen/decomp/dtype.py

@@ -20,13 +20,17 @@ def reindex(idx:UOp, off:int, mul=2) -> UOp:
 def l2i(op: Ops, dt: DType, *uops:UOp):
   zero = UOp.const(dt, 0)
   if len(uops) == 2: a0, a1 = uops
+  elif len(uops) == 3: a0, a1, b0 = uops  # a shift's count is a single word
   elif len(uops) == 4: a0, a1, b0, b1 = uops
   match op:
     case Ops.NEG: return l2i(Ops.SUB, dt, zero, zero, *uops)
     case Ops.CAST if dt in (dtypes.long, dtypes.ulong) and uops[0].dtype not in dtypes.floats:
-      return uops[0].cast(l2i_dt[dt]), (uops[0] < 0).where(UOp.const(l2i_dt[dt], -1), UOp.const(l2i_dt[dt], 0))
+      # the high word is the sign extension; bool has no sign, test the already-cast low word instead (bool < 0 would promote to weakint)
+      x, lo = uops[0], uops[0].cast(l2i_dt[dt])
+      sign = lo if x.dtype is dtypes.bool else x
+      return lo, (sign < sign.const_like(0)).where(lo.const_like(-1), lo.const_like(0))
     case Ops.CAST if dt in (dtypes.long, dtypes.ulong):
-      return (lo:=uops[0].cast(l2i_dt[dt])), (uops[0] / 2**32).cast(l2i_dt[dt]) - ((uops[0] < 0) & lo.ne(0)).cast(l2i_dt[dt])
+      return (lo:=uops[0].cast(l2i_dt[dt])), (uops[0] / 2**32).cast(l2i_dt[dt]) - ((uops[0] < 0) & lo.ne(0))
     case Ops.CAST if dt in dtypes.floats:
       small = (a1.eq(0) & (a0 >= 0)) | (a1.eq(-1) & (a0 < 0))
       return small.where(a0.cast(dt), ((a1.cast(dtypes.float32) * (2**32)) + a0.bitcast(dtypes.uint).cast(dtypes.float32)).cast(dt))
@@ -41,8 +45,8 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
       lo, hi = ((a0u >> n) | ((a1u << 1) << (31 - n))).bitcast(dt), a1 >> (b0 & 31)
       fill = a1 >> 31 if dt == dtypes.int else zero  # vacated high word: sign bits when signed, else 0
       return (b0 >= 32).where(hi, lo), (b0 >= 32).where(fill, hi)
-    case Ops.ADD: return (low:=a0+b0), (a1 + b1).replace(dtype=dt) + (low.bitcast(dtypes.uint) < a0.bitcast(dtypes.uint)).cast(dt)
-    case Ops.SUB: return a0 - b0, a1 - b1 - (a0.bitcast(dtypes.uint) < b0.bitcast(dtypes.uint)).cast(dt)
+    case Ops.ADD: return (low:=a0+b0), a1 + b1 + (low.bitcast(dtypes.uint) < a0.bitcast(dtypes.uint))
+    case Ops.SUB: return a0 - b0, a1 - b1 - (a0.bitcast(dtypes.uint) < b0.bitcast(dtypes.uint))
     case Ops.MUL:
       (a00, a01), (b00, b01) = unpack32(a0), unpack32(b0)
       mid = l2i(Ops.ADD, dt, shl(a00*b01, 16).bitcast(dt), shr(a00*b01, 16).bitcast(dt), shl(a01*b00, 16).bitcast(dt), shr(a01*b00, 16).bitcast(dt))
@@ -65,19 +69,23 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
         (nq0, nq1), (nr0, nr1) = l2i(Ops.BITCAST, dt, *l2i(Ops.NEG, dtypes.uint, *q)), l2i(Ops.BITCAST, dt, *l2i(Ops.NEG, dtypes.uint, *r))
         (q0, q1), (r0, r1) = l2i(Ops.BITCAST, dt, *q), l2i(Ops.BITCAST, dt, *r)
         return (a_neg.where(nr0, r0), a_neg.where(nr1, r1)) if op == Ops.CMOD else ((a_neg^b_neg).where(nq0, q0), (a_neg^b_neg).where(nq1, q1))
-      return (r[0].bitcast(dt), r[1].bitcast(dt)) if op == Ops.CMOD else (q[0].bitcast(dt), q[1].bitcast(dt))
+      return r if op == Ops.CMOD else q
     case Ops.CMPLT: return (a1 < b1) | ((a1.eq(b1)) & (a0.bitcast(dtypes.uint) < b0.bitcast(dtypes.uint)))
     case Ops.CMPEQ: return a0.eq(b0) & a1.eq(b1)
     case Ops.CMPNE: return a0.ne(b0) | a1.ne(b1)
-    case Ops.XOR | Ops.OR | Ops.AND: return UOp(op, dt, src=(a0, b0)), UOp(op, dt, src=(a1, b1))
+    case Ops.XOR | Ops.OR | Ops.AND: return UOp(op, src=(a0, b0)), UOp(op, src=(a1, b1))
     case Ops.WHERE: return uops[0].where(uops[1], uops[3]), uops[0].where(uops[2], uops[4])
     case Ops.MAX: return l2i(Ops.WHERE, dt, l2i(Ops.CMPLT, dt, *uops), b0, b1, a0, a1)
     case _: raise NotImplementedError(f"long decomposition of {op} unsupported")
 
+def split_l2i(op: Ops, dt: DType, *uops:UOp):
+  # l2i does arithmetic on its inputs; rules enter here to split them to 32-bit words first, l2i recurses on itself
+  return l2i(op, dt, *graph_rewrite(UOp.sink(*uops), pm_long_decomp, bottom_up=True).src)
+
 # ***** floats *****
 f2f_dt = { f:getattr(dtypes, f"uint{f.bitsize}") for f in dtypes.floats }
 
-def rne(v: UOp, s) -> UOp: return shr(v, s) + ((shr(v, s - 1) & 1) & ((v & ((1 << (s - 1)) - 1)).ne(0).cast(v.dtype) | (shr(v, s) & 1)))
+def rne(v: UOp, s) -> UOp: return shr(v, s) + ((shr(v, s - 1) & 1) & ((v & ((1 << (s - 1)) - 1)).ne(0) | (shr(v, s) & 1)))
 
 def f2f(v, fr:DType, to:DType, sat=True):
   fs, fb, (fe, fm), ts, tb, (te, tm) = fr.bitsize, exponent_bias(fr), dtypes.finfo(fr), to.bitsize, exponent_bias(to), dtypes.finfo(to)
@@ -122,27 +130,35 @@ def f2f_store(st, idx, val, fr:DType, to:DType):
   if (n:=val.max_numel()) == 1: return st.replace(src=(idx, f2f(val.bitcast(f2f_dt[to]), to, fr)))
   return UOp.group(*(st.replace(src=(reindex(idx, i, 1), f2f(val.index(i).bitcast(f2f_dt[to]), to, fr))) for i in range(n)))
 
+# tag is the 32-bit word this node becomes - (0 for the low word, 1 for the high, the dtype the consumer wants)
 pm_long_decomp = PatternMatcher([
   (UPat(GroupOp.Defines, src=(UPat.var("sz"),), name="x"), lambda x,sz:
    x.replace(dtype=l2i_dt[x.dtype], arg=replace(x.arg, dtype=l2i_dt[x.dtype]), src=(sz*2,)) if x.dtype in l2i_dt else None),
-  (UPat(Ops.INDEX, tuple(l2i_dt.keys()), name='x'), lambda x: reindex(x, x.tag).replace(dtype=l2i_dt[x.dtype]) if x.tag is not None else None),
-  (UPat(Ops.STORE, src=(UPat.var('idx'), UPat.var('val', tuple(l2i_dt.keys()))), name='st'), lambda st,idx,val:
-   st.replace(src=(idx.rtag(0), val.rtag(0))).group(st.replace(src=(idx.rtag(1), val.rtag(1)))) if val.tag is None else None),
-  (UPat(GroupOp.Comparison, src=(UPat.var('a', tuple(l2i_dt.keys())), UPat.var('b', tuple(l2i_dt.keys()))), name="x"), lambda a,b,x:
-   l2i(x.op, dt:=l2i_dt[a.dtype], a.rtag(0).cast(dt), a.rtag(1).cast(dt), b.rtag(0).cast(dt), b.rtag(1).cast(dt))),
-  (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a'),), name="x"), lambda a,x:
-   l2i(x.op, x.dtype, a)[x.tag] if x.tag is not None and a.dtype not in l2i_dt else None),
+  (UPat(Ops.INDEX, tuple(l2i_dt.keys()), name='x'), lambda x:
+   reindex(x, x.tag[0]).replace(dtype=x.tag[1], tag=None) if x.tag is not None else None),
+  (UPat(Ops.STORE, src=(UPat.var('idx', tuple(l2i_dt.keys())), UPat.var('val')), name='st'), lambda st,idx,val:
+   st.replace(src=(idx.rtag((0, dt:=l2i_dt[idx.dtype])), val.rtag((0, dt)))).group(
+     st.replace(src=(idx.rtag((1, dt)), val.rtag((1, dt))))) if val.tag is None else None),
+  (UPat(GroupOp.Comparison, src=[UPat.var('a', tuple(l2i_dt.keys())), UPat()], name="x"), lambda a,x:
+   split_l2i(x.op, dt:=l2i_dt[a.dtype], *flatten((s.rtag((0, dt)), s.rtag((1, dt))) for s in x.src))),
   (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda a,x:
-   (a.rtag(0).cast(dt:=l2i_dt[a.dtype]).bitcast(xdt:=l2i_dt[x.dtype]), a.rtag(1).cast(dt).bitcast(xdt))[x.tag]),
+   split_l2i(Ops.BITCAST, l2i_dt[x.dtype], a.rtag((0, dt:=l2i_dt[a.dtype])), a.rtag((1, dt)))[x.tag[0]]),
+  (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a'),), name="x"), lambda a,x:
+   split_l2i(x.op, x.dtype, a)[x.tag[0]] if x.tag is not None else None),
   (UPat(Ops.CAST, src=(UPat.var('a', tuple(l2i_dt.keys())),), name="x"), lambda a,x:
-   l2i(x.op, x.dtype, a.rtag(0).cast(dt:=l2i_dt[a.dtype]), a.rtag(1).cast(dt)) if x.dtype not in l2i_dt and a.tag is None else None),
-  (UPat((*(GroupOp.ALU - GroupOp.Comparison), Ops.BITCAST), tuple(l2i_dt.keys()), name="x"), lambda x:
-   l2i(x.op, l2i_dt[x.dtype], *flatten((a.rtag(0).cast(dt:=l2i_dt[x.src[-1].dtype]), a.rtag(1).cast(dt))
-                                       if a.dtype in l2i_dt else (a,) for a in x.src))[x.tag] if x.tag is not None else None),
+   split_l2i(x.op, x.dtype, a.rtag((0, dt:=l2i_dt[a.dtype])), a.rtag((1, dt))) if x.dtype not in l2i_dt and a.tag is None else None),
+  (UPat((Ops.SHL, Ops.SHR), tuple(l2i_dt.keys()), src=(UPat.var('a'), UPat.var('b')), name="x"), lambda a,b,x:
+   split_l2i(x.op, dt:=l2i_dt[x.dtype], a.rtag((0, dt)), a.rtag((1, dt)), b.rtag((0, dt)))[x.tag[0]] if x.tag is not None else None),
+  (UPat(Ops.WHERE, tuple(l2i_dt.keys()), src=(UPat.var('c'), UPat.var('a'), UPat.var('b')), name="x"), lambda a,b,c,x:
+   split_l2i(x.op, dt:=l2i_dt[x.dtype], c, a.rtag((0, dt)), a.rtag((1, dt)), b.rtag((0, dt)), b.rtag((1, dt)))[x.tag[0]]
+   if x.tag is not None else None),
+  (UPat((*(GroupOp.ALU - GroupOp.Comparison - {Ops.SHL, Ops.SHR, Ops.WHERE}), Ops.BITCAST), tuple(l2i_dt.keys()), name="x"), lambda x:
+   split_l2i(x.op, l2i_dt[x.dtype], *flatten((a.rtag((0, l2i_dt[x.dtype])), a.rtag((1, l2i_dt[x.dtype]))) for a in x.src))[x.tag[0]]
+   if x.tag is not None else None),
   (UPat(Ops.LOAD, tuple(l2i_dt.keys()), src=(UPat.var('idx'),), name='x'), lambda x,idx:
-   x.replace(dtype=l2i_dt[x.dtype], src=(reindex(idx, x.tag).replace(dtype=l2i_dt[x.dtype]),)) if x.tag is not None else None),
-  (UPat(Ops.CONST, tuple(l2i_dt.keys()), name='x'), lambda x:
-   UOp.const(dt:=l2i_dt[x.dtype], truncate[dt]((x.arg >> 32) if x.tag == 1 else (x.arg & 0xFFFFFFFF))))
+   x.replace(dtype=l2i_dt[x.dtype], src=(reindex(idx, x.tag[0]).replace(dtype=l2i_dt[x.dtype], tag=None),), tag=None) if x.tag is not None else None),
+  (UPat(Ops.CONST, tag={(w, dt) for w in (0, 1) for dt in l2i_dt.values()}, name='x'), lambda x:
+   UOp.const(x.tag[1], truncate[x.tag[1]]((x.arg >> 32) if x.tag[0] == 1 else (x.arg & 0xFFFFFFFF))))
 ])
 
 # float decomposition patterns - ctx is (fr, to) tuple

tinygrad/codegen/decomp/op.py

@@ -63,7 +63,7 @@ def threefry2x32(x: UOp, key: UOp):
 
 def floordiv_to_idiv(a:UOp, b:UOp) -> UOp:
   if (a.vmin >= 0 and b.vmin > 0) or (a.vmax <= 0 and b.vmax < 0): return a.alu(Ops.CDIV, b)
-  return a.alu(Ops.CDIV, b) - (a.alu(Ops.CMOD, b).ne(0) & (a<0).ne(b<0)).cast(a.dtype)
+  return a.alu(Ops.CDIV, b) - (a.alu(Ops.CMOD, b).ne(0) & (a<0).ne(b<0))
 
 def floormod_to_mod(a:UOp, b:UOp) -> UOp:
   if (a.vmin >= 0 and b.vmin > 0) or (a.vmax <= 0 and b.vmax < 0): return a.alu(Ops.CMOD, b)
@@ -114,8 +114,8 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa
   if Ops.CMPLT in ops:
     # These are late rewrites because simplex expects equalities to be a certain format
     pat += [
-      ((UPat.var("x", dtypes.sints) < UPat.cvar("c", dtypes.sints)).logical_not(), lambda x,c: c-1<x),
-      ((UPat.cvar("c", dtypes.sints) < UPat.var("x", dtypes.sints)).logical_not(), lambda x,c: x<c+1),
+      ((UPat.var("x", dtypes.sints) < UPat.cvar("c")).logical_not(), lambda x,c: c-1<x),
+      ((UPat.cvar("c") < UPat.var("x", dtypes.sints)).logical_not(), lambda x,c: x<c+1),
       (UPat.var("x", dtypes.sints)*-1 < UPat.var("y", dtypes.sints)*UPat.cvar("c"), lambda x,y,c: y*(-c)<x),
       (UPat.var("x", dtypes.sints)*-1 < UPat.cvar("c"), lambda x,c:-c<x),
       ((UPat.cvar("c1")<UPat.var("x", dtypes.sints)) & (UPat.var("x", dtypes.sints)<UPat.cvar("c2")),
@@ -129,5 +129,5 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa
   # some backends emit FDIV for RECIP, in that case: a*(1/b) -> a/b
   if Ops.FDIV in ops:
     pat += [(UPat.var("x").reciprocal(), lambda x: x.const_like(1).alu(Ops.FDIV, x))]
-    pat += [(UPat.var("a", dtypes.floats) * UPat.const(dtypes.floats, 1).alu(Ops.FDIV, UPat.var("b")), lambda a,b: a.alu(Ops.FDIV, b))]
+    pat += [(UPat.var("a", dtypes.floats) * UPat(Ops.FDIV, dtypes.floats, src=(UPat.const(None, 1), UPat.var("b"))), lambda a,b: a.alu(Ops.FDIV, b))]
   return PatternMatcher(pat)

tinygrad/codegen/decomp/transcendental.py

@@ -200,7 +200,7 @@ def xexp2(d:UOp) -> UOp:
   x = _lazy_map_numbers(d, d.const_like(0.0), d.const_like(0.0), d.const_like(0.0), d)
   q = rintk(x)
   # s = d - round(d)
-  s = x - q.cast(x.dtype)
+  s = x - q
   # a polynomial approximation with 13 non-zero terms in the range of [−(log 2)/2,(log 2)/2].
   if d.dtype == dtypes.float64:
     u = polyN(s, [0.4434359082926529454e-9, 0.7073164598085707425e-8, 0.1017819260921760451e-6, 0.1321543872511327615e-5, 0.1525273353517584730e-4,

tinygrad/codegen/gpudims.py

@@ -87,7 +87,15 @@ def add_gpudims(ctx:Renderer, s:UOp):
     except ValueError: continue
   return s.substitute(subs)
 
+pm_device_to_var = PatternMatcher([
+  # the DEVICE axis is not a program axis, it's bound per device at launch. lower it to the _device_num variable (like SPECIAL for devices)
+  (UPat(Ops.RANGE, name="r"), lambda r: UOp.variable("_device_num", 0, r.vmax, dtype=r.dtype) if r.arg[-1] is AxisType.DEVICE else None),
+  # ENDs that closed a DEVICE range no longer close it
+  (UPat(Ops.END, name="e"), lambda e: e.replace(src=(e.src[0],)+tuple(s for s in e.src[1:] if s.op is not Ops.PARAM))
+   if any(s.op is Ops.PARAM and s.arg.name == '_device_num' for s in e.src[1:]) else None),
+])
+
 pm_add_gpudims = PatternMatcher([
   # add gpudims must be last
   (UPat(Ops.SINK, name="s"), add_gpudims),
-])
+])+pm_device_to_var

tinygrad/codegen/late/coalesce.py

@@ -68,7 +68,7 @@ def image_valid_dims(base:DType, size:int, arch:str) -> list[tuple[int,int]]:
   return [(pxls//ALIGN//k, ALIGN*k) for k in range(ceildiv(pxls//ALIGN, MAXW), min(pxls//ALIGN, MAXW//ALIGN)+1) if (pxls//ALIGN)%k == 0]
 
 def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
-  shapes, ren = ctx
+  shapes, ren = ctx[:2]
   if not IMAGE or ren.target.device not in {"QCOM", "CL", "PYTHON", "NULL"}: return None
   valid, x = x.get_valid(), x.get_idx()
   # search for dims that drop the most valid statements
@@ -106,6 +106,7 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
     # TODO: this should handle images too, it's just memory coalescing
     if u.op in {Ops.LOAD, Ops.STORE}:
       assert len(u.src) == (2 if u.op is Ops.STORE else 1), "memory coalescing does not support gated loads/stores"
+      if u.src[0].op is Ops.SHRINK: continue
       assert u.src[0].op is Ops.INDEX, f"memory coalescing should be on INDEX, not {u.src[0].op}"
       buf, idx_u = u.src[0].src
       if buf.addrspace == AddrSpace.REG: continue
@@ -141,12 +142,12 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
     grouped_offsets = [[x for _,x in group] for _,group in itertools.groupby(enumerate(sorted(offsets.keys())), lambda x: x[1]-x[0])]
     for full_grp in grouped_offsets:
       while len(full_grp):
-        offset = (base+full_grp[0]) if isinstance(base, UOp) else UOp.const(dtypes.weakint, full_grp[0])
+        offset = (base+full_grp[0]) if isinstance(base, UOp) else UOp.const(None, full_grp[0])
         length = [l for l in lengths if l <= len(full_grp) and (not must_divide or offset.divides(l) is not None)][0]
         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(dtypes.weakint, len(grp)))) if len(grp) > 1 else buf.index(offset)
+        idx = UOp(Ops.SHRINK, src=(buf, offset, UOp.const(None, len(grp)))) if len(grp) > 1 else buf.index(offset)
         if op == Ops.STORE:
           datas = []
           for i,g in enumerate(grp):
@@ -158,7 +159,7 @@ def memory_coalescing(sink:UOp, ctx:Renderer) -> UOp:
           ld = idx.load()
           for i,g in enumerate(grp):
             for oo in offsets[g]:
-              replacements[oo] = ld.index(UOp.const(dtypes.weakint, i)) if len(grp) > 1 else ld
+              replacements[oo] = ld.index(i) if len(grp) > 1 else ld
         full_grp = full_grp[length:]
 
   # apply

tinygrad/codegen/late/linearizer.py

@@ -9,7 +9,9 @@ def linearize(sink:UOp) -> list[UOp]:
   # this is a toposort with priority
   lst = list(sink.toposort())
   out_degree:defaultdict[UOp, int] = defaultdict(int)
-  priorities:dict[UOp, tuple[int, int, Any]] = {}
+  priorities:dict[UOp, tuple[int, int, int, Any]] = {}
+  wmma_depth:dict[UOp, int] = {}
+  for u in lst: wmma_depth[u] = max((wmma_depth[s] for s in u.src), default=0) + (u.op is Ops.WMMA)
 
   # get consumers and assign priorities
   # NOTE: this requires the lst be locally toposorted
@@ -30,7 +32,7 @@ def linearize(sink:UOp) -> list[UOp]:
       case Ops.RANGE: priority = 5    # placing RANGE is good
       case Ops.END: priority = -5     # placing END is bad
       case _: priority = 0            # everything else has priority 0
-    priorities[u] = (run_count, priority, extra)
+    priorities[u] = (run_count, priority, wmma_depth[u], extra)
 
   # number the uops in "ideal" order
   nkey = {u:i for i,u in enumerate(sorted(lst, key=lambda x: priorities[x]+(x.tuplize if TUPLE_ORDER else ())))}

tinygrad/codegen/opt/heuristic.py

@@ -169,7 +169,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
     else:
       # prioritize making expand axes local
       local_axis_ranking = [(any(k.rngs[axis] not in b.src[1].get_idx().backward_slice for b in k.bufs), axis) \
-                              for axis in k.axes_of(AxisType.GLOBAL, AxisType.LOOP) if k.rngs[axis].src[0].op is Ops.CONST]
+                              for axis in k.axes_of(AxisType.GLOBAL, AxisType.WEAK) if k.rngs[axis].src[0].op is Ops.CONST]
       to_local: list[tuple[int, int]] = []
       for _, axis in sorted(local_axis_ranking, key=lambda x: (-x[0], -x[1])):
         local_size = prod(sz for _, sz in to_local)
@@ -188,7 +188,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
     for threads in [32,16,12,8,6,5,4,3,2]:
       # Skip if too many threads. Heuristic: use about 128K ops per thread
       if threads > k.ren.global_max[0] or resolve(prod(k.full_shape) // (128 << 10) < threads): continue
-      for axis in k.axes_of(AxisType.LOOP):
+      for axis in k.axes_of(AxisType.WEAK):
         if k.full_shape[axis] % threads == 0:
           try: k.apply_opt(Opt(OptOps.THREAD, axis, threads))
           except KernelOptError: pass

tinygrad/codegen/opt/postrange.py

@@ -21,9 +21,9 @@ class Scheduler:
 
   @property
   def rngs(self):
-    # always in order by axistype. void RANGEs are loops, not opt axes
-    return sorted([u for u in self.ast.backward_slice if u.op is Ops.RANGE and u.dtype is not dtypes.void and u.vmax > 0],
-                  key=lambda x: (axis_to_pos[x.arg[-1]],) + x.arg[0:-1])
+    # always in order by axistype. void RANGEs are loops, not opt axes. the DEVICE axis is launched, not an opt axis
+    return sorted([u for u in self.ast.backward_slice if u.op is Ops.RANGE and u.dtype is not dtypes.void and u.vmax > 0
+                   and u.arg[-1] is not AxisType.DEVICE], key=lambda x: (axis_to_pos[x.arg[-1]],) + x.arg[0:-1])
   @property
   def shape_len(self) -> int: return len(self.rngs)
   @property
@@ -65,7 +65,7 @@ class Scheduler:
   def _output_rngs(self) -> list[UOp]:
     return flatten([[r for r in UOp.sink(*s.src[1:]).ranges if r.arg[-1] != AxisType.REDUCE] for s in self.ast.src if s.op is Ops.END])
   def _globalizable_rngs(self) -> list[UOp]:
-    ret = [r for r in self._output_rngs() if r.arg[-1] == AxisType.LOOP]
+    ret = [r for r in self._output_rngs() if r.arg[-1] == AxisType.WEAK]
     # exclude any output ranges from global that don't appear in all BUFFERIZE
     for x in self.ast.toposort():
       if x.op is Ops.STAGE:
@@ -86,8 +86,8 @@ class Scheduler:
     ret = []
     for x,r in zip(self.axis_types, self.rngs):
       if self.dont_use_locals and x == AxisType.GLOBAL: ret.append("BLUE")
-      elif r not in output_rngs and x == AxisType.LOOP: ret.append("BLACK")
-      elif r not in globalizible_rngs and x == AxisType.LOOP: ret.append("white")
+      elif r not in output_rngs and x == AxisType.WEAK: ret.append("BLACK")
+      elif r not in globalizible_rngs and x == AxisType.WEAK: ret.append("white")
       else: ret.append(axis_colors[x])
     return ret
   def colored_shape(self) -> str: return ' '.join([colored(f'{x.src[0].render():>4s}', color) for x,color in zip(self.rngs, self.colors())])
@@ -95,7 +95,7 @@ class Scheduler:
   def shift_to(self, rng:UOp, amount:int, new_type:AxisType, top:bool=False, input_new_rng:UOp|None=None):
     if (old_sz:=rng.src[0].divides(amount)) is None:
       raise KernelOptError(f"{amount} can't divide {rng.src[0]} in {self.colored_shape()}")
-    new_rng = UOp.range(amount, next(self.opt_range), new_type) if input_new_rng is None else input_new_rng
+    new_rng = UOp.range(amount, next(self.opt_range), new_type, dtype=rng.dtype) if input_new_rng is None else input_new_rng
     replaced_rng = rng.replace(src=(old_sz,))
     sub_axis = (new_rng * old_sz + replaced_rng) if top else (replaced_rng * amount + new_rng)
     self.ast = self.ast.substitute({rng:sub_axis}, name=f"shift {rng.arg[:-1]} {amount} {str(new_type).split('.')[1].lower()}")
@@ -108,7 +108,7 @@ class Scheduler:
 
   # copied from kernel.py
   @property
-  def upcastable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP) \
+  def upcastable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GLOBAL, AxisType.LOCAL, AxisType.WEAK) \
                                                   if isinstance(s:=self.full_shape[i], int) and s > 1]
   @property
   def unrollable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE) \
@@ -161,10 +161,10 @@ class Scheduler:
         check(rng.arg[-1] in {AxisType.GROUP_REDUCE, AxisType.REDUCE}, "unroll is for GROUP_REDUCE/REDUCE")
       if opt.op is OptOps.UPCAST:
         check((self.ren is not None and self.ren.target.device == "DSP") or amt <= 16, "don't upcast more than 16")
-        check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP}, f"upcast is for GLOBAL/LOCAL/LOOP, not {rng.arg[-1]}")
+        check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOCAL, AxisType.WEAK}, f"upcast is for GLOBAL/LOCAL/LOOP, not {rng.arg[-1]}")
       if opt.op is OptOps.LOCAL:
         check(not self.dont_use_locals, "can't use locals")
-        check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.LOOP}, "local is for globals")
+        check(rng.arg[-1] in {AxisType.GLOBAL, AxisType.WEAK}, "local is for globals")
       if opt.op is OptOps.THREAD:
         check(self.ren is not None and self.ren.has_threads, "target does not support threads")
         check(self.ren is not None and self.ren.global_max is not None and amt <= self.ren.global_max[0], "too many threads")
@@ -191,7 +191,7 @@ class Scheduler:
       check(rng.arg[-1] is not AxisType.THREAD, "cannot pad thread")
       new_sz = round_up(int(rng.vmax+1), cast(int, opt.arg))
       check(rng.vmax+1 > new_sz//4, "pad adds more than quadruple the work")
-      replaced_rng = UOp.range(new_sz, *rng.arg)
+      replaced_rng = UOp.range(new_sz, *rng.arg, dtype=rng.dtype)
       replaces = {rng:replaced_rng}
       valid = replaced_rng < rng.vmax+1
       store_targets = {s.src[0] for s in self.ast.backward_slice_with_self if s.op is Ops.STORE}

tinygrad/codegen/simplify.py

@@ -1,7 +1,7 @@
 import itertools
 from typing import Callable
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, graph_rewrite, _substitute, range_start, AxisType
-from tinygrad.uop.symbolic import symbolic
+from tinygrad.uop.symbolic import symbolic, invalid_gate
 from tinygrad.helpers import partition
 from tinygrad.dtype import dtypes
 
@@ -60,7 +60,9 @@ pm_simplify_ranges = PatternMatcher([
 ])
 
 def mark_range_mod(ctx:dict[UOp, UOp|None], r:UOp, c:UOp) -> None:
-  if r not in ctx and r.arg[-1] is not AxisType.WARP and r.src[0].op is Ops.CONST and r.src[0].divides(c.arg) is not None: ctx[r] = c
+  # ranges that aren't looped over can't be split
+  if r not in ctx and r.arg[-1] not in {AxisType.WARP, AxisType.DEVICE} \
+    and r.src[0].op is Ops.CONST and r.src[0].divides(c.arg) is not None: ctx[r] = c
 
 def do_substitute(ctx:dict, x: UOp, sub_fxn:Callable[[UOp, UOp], UOp]) -> UOp|None:
   ret = x.substitute({k:sub_fxn(k,v) for k,v in ctx.items() if v is not None})
@@ -84,9 +86,9 @@ def reduce_unparented(red:UOp) -> UOp|None:
   if len(reduce_unparented) == 0: return None
   ret = red.replace(src=(red.src[0],)+tuple(reduce_parented)) if len(reduce_parented) or red.dtype != red.src[0].dtype else red.src[0]
   if red.arg[0] is Ops.ADD:
-    for r in reduce_unparented: ret = ret * r.src[0].cast(ret.dtype)
+    for r in reduce_unparented: ret = ret * r.src[0]
   if red.arg[0] is Ops.MUL:
-    for r in reduce_unparented: ret = ret ** r.src[0].cast(ret.dtype)
+    for r in reduce_unparented: ret = ret ** r.src[0]
   return ret
 
 pm_reduce_unparented = PatternMatcher([
@@ -95,12 +97,11 @@ pm_reduce_unparented = PatternMatcher([
 ])
 
 pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
-  # lift x+y out of reduce on lt. only fire if x still has the range: with both sides range-free it just shuffles constants
-  ((UPat.var("x")+UPat.var("y")).or_casted() < UPat.var("c"),
-   lambda x,y,c: (x < (c.cast(y.dtype)-y)) if not no_range(x) and no_range(y) and no_range(c) else None),
+  # lift x+y out of reduce on lt
+  ((UPat.var("x")+UPat.var("y")).or_casted() < UPat.var("c"), lambda x,y,c: (x < (c-y)) if no_range(y) and no_range(c) else None),
   # lift x*y out of reduce
   ((UPat.var("x")*UPat.var("y")) < UPat.var("c"),
-    lambda x,y,c: (x < ((c+y-1) // y)) if not no_range(x) and no_range(y) and no_range(c) and dtypes.is_int(y.dtype) and y.vmin > 0 else None),
+   lambda x,y,c: (x < ((c+y-1) // y)) if no_range(y) and no_range(c) and dtypes.is_int(y.dtype) and y.vmin > 0 else None),
   # sum over r in [0,N) of [lower<=r<upper]*val -> clamp(min(upper,N) - max(lower,0), 0, N) * val
   (UPat.any(
     (UPat(Ops.RANGE, name="r") < UPat.var("upper")).where(UPat.var("val"), 0),
@@ -108,21 +109,21 @@ pm_reduce_collapse = pm_reduce_unparented + PatternMatcher([
     ((UPat.var("r")<UPat.var("lower")).logical_not()&(UPat(Ops.RANGE, name="r")<UPat.var("upper"))).where(UPat.var("val"), 0),
   ).reduce(UPat.var("r"), arg=Ops.ADD), lambda r,val,lower=None,upper=None:
     ((upper.minimum(r.src[0]) if upper is not None else r.src[0]) -
-     (lower.maximum(0) if lower is not None else r.const_like(0))).maximum(0).minimum(r.src[0]).cast(val.dtype) * val if no_range(val) else None),
-  # REDUCE on ADD
+     (lower.maximum(0) if lower is not None else r.const_like(0))).maximum(0).minimum(r.src[0]) * val if no_range(val) else None),
+  (invalid_gate.reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
+   lambda cond,x,i,r: cond.where(x.reduce(*r.src[1:], arg=Ops.ADD), i) if no_range(cond) else None),
   ((UPat.var("x")+UPat.var("y")).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
    lambda x,y,r: x.reduce(*r.src[1:], arg=Ops.ADD) + y.reduce(*r.src[1:],arg=Ops.ADD)),
   # AND on WHERE
   ((UPat(Ops.PARAM, name="x") & UPat.var("y")).where(UPat.var("c"), 0).reduce(arg=Ops.ADD, allow_any_len=True, name="r"),
-    lambda x,y,c,r: y.where(c, 0).reduce(*r.src[1:], arg=Ops.ADD)*x.cast(c.dtype)),
+    lambda x,y,c,r: y.where(c, 0).reduce(*r.src[1:], arg=Ops.ADD)*x),
   # MUL casted bool
   ((UPat.var("x") * UPat.var("gate", dtype=dtypes.bool).cast()), lambda x,gate: gate.where(x, 0)),
 ])+symbolic
 
 pm_reduce_load_collapse = pm_reduce_collapse + PatternMatcher([
-  # lift x+y out of reduce on ne (same range guard as the lt version)
-  ((UPat.var("x")+UPat.var("y")).or_casted() != UPat.var("c"),
-   lambda x,y,c: (x != (c.cast(y.dtype)-y)) if not no_range(x) and no_range(y) and no_range(c) else None),
+  # lift x+y out of reduce on ne
+  ((UPat.var("x")+UPat.var("y")).or_casted() != UPat.var("c"), lambda x,y,c: (x != (c.cast(y.dtype)-y)) if no_range(y) and no_range(c) else None),
   # reduce on gated load becomes can substitute the range and remove the reduce
   ((UPat.var("idx")!=(UPat(Ops.RANGE, name="r").or_casted())).where(0, UPat.var("expr")).reduce(UPat.var("r"), arg=Ops.ADD),
    lambda r,idx,expr: (v:=(idx.cast(r.dtype) >= 0) & (idx.cast(r.dtype) < r.src[0])).where(expr.substitute({r:idx.cast(r.dtype).valid(v)}),0)),

tinygrad/device.py

@@ -5,7 +5,7 @@ from typing import Any, Generic, TypeVar, Iterator, Generator, Self, TYPE_CHECKI
 import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal
 from tinygrad.helpers import LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, PROFILE, temp, colored
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing
-from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize, Target, unwrap
+from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize, Target, unwrap, round_up
 from tinygrad.dtype import DType, _to_np_dtype
 if TYPE_CHECKING: from tinygrad.renderer import Renderer
 
@@ -268,6 +268,34 @@ class LRUAllocator(Allocator, Generic[DeviceType]):
     if LRU and (options is None or (not options.nolru and options.external_ptr is None)): self.cache[(size, options)].append(opaque)
     else: super().free(opaque, size, options)
 
+class DepsTracker:
+  def __init__(self):
+    # tracks (offset, end, dep) ranges per base buffer id to handle suballocated buffers correctly.
+    self.w_dependency_map: dict[int, list[tuple[int, int, Any]]] = defaultdict(list)
+    self.r_dependency_map: dict[int, list[tuple[int, int, Any]]] = defaultdict(list)
+
+  @staticmethod
+  def _key(buf:Any) -> tuple[Any, int, int]: return id(buf.base), buf.offset, buf.offset + buf.nbytes
+
+  def access_resources(self, bufs:list[Any], write:list[int], new_dependency:Any):
+    wait_nodes = []
+    for i,buf in enumerate(bufs):
+      key, s, e = self._key(buf)
+      wait_nodes += [dep for st,en,dep in self.w_dependency_map[key] if st < e and s < en]
+      if i in write: wait_nodes += [dep for st,en,dep in self.r_dependency_map[key] if st < e and s < en]
+    for i,buf in enumerate(bufs):
+      key, s, e = self._key(buf)
+      if i in write:
+        for dmap in [self.w_dependency_map, self.r_dependency_map]:
+          kept = []
+          for st,en,dep in dmap[key]:
+            if st < min(s, en): kept.append((st, min(s, en), dep))
+            if max(e, st) < en: kept.append((max(e, st), en, dep))
+          dmap[key] = kept
+        self.w_dependency_map[key].append((s, e, new_dependency))
+      else: self.r_dependency_map[key].append((s, e, new_dependency))
+    return list({id(x):x for x in wait_nodes}.values())
+
 # **************** for Compiled Devices ****************
 
 class CompileError(Exception): pass
@@ -290,6 +318,13 @@ class TinyELF:
   target: Target
   # tuple of (name, slot, dtype, shape)
   signature: tuple[tuple[str|None, int, DType, tuple], ...]
+  parallel: bool = False
+
+  @staticmethod
+  def iter_sig(signature:tuple[tuple[str|None, int, DType, tuple], ...], offset:int=0) -> Generator[tuple[int, DType], None, None]:
+    for _,_,dt,_ in signature:
+      yield (offset:=round_up(offset, dt.itemsize)), dt
+      offset += dt.itemsize
 
 class Program(Generic[DeviceType]):
   def __init__(self, dev:DeviceType, obj:TinyELF): pass
@@ -299,6 +334,9 @@ class Program(Generic[DeviceType]):
 class Compiled:
   profile_events:list[ProfileEvent] = [ProfileDeviceEvent("CPU")] # NOTE: CPU is the default device.
 
+  pm_lower:Any = None
+  pm_bufferize:Any = None
+
   def __init__(self, device:str, allocator:Allocator, renderers:list[type[Renderer]], runtime:type[Program[Self]]|None, graph=None, arch=None):
     from tinygrad.renderer import Renderer
     self.device, self.allocator, self.runtime_t, self.graph, self.renderers = device, allocator, runtime, graph, renderers or [Renderer]

tinygrad/dtype.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 from typing import Final, ClassVar, Callable, Literal
 import math, struct, ctypes, functools
 from dataclasses import dataclass, fields
-from tinygrad.helpers import getenv
+from tinygrad.helpers import getenv, DEFAULT_FLOAT, DEFAULT_INT
 from enum import IntEnum, auto
 
 class ConstFloat(float):
@@ -83,7 +83,7 @@ class DType(metaclass=DTypeMetaClass):
     return ConstFloat(float(val)) if dtypes.is_float(self) else bool(val) if dtypes.is_bool(self) else int(val)
 
 
-class dtypes:
+class DTypes:
   @staticmethod
   @functools.cache
   def is_float(x: DType) -> bool: return x in (dtypes.floats + (dtypes.weakfloat,))
@@ -99,10 +99,10 @@ class dtypes:
   def from_py(x) -> DType:
     # NOTE: isinstance(True, int) is True, so bool must be checked before int
     if isinstance(x, (bool, InvalidType)): return dtypes.bool
-    if isinstance(x, float): return dtypes.default_float
-    if isinstance(x, int): return dtypes.default_int
+    if isinstance(x, float): return dtypes.weakfloat
+    if isinstance(x, int): return dtypes.weakint
     # put this in the last is faster because there are more items than lists/tuples to check
-    if isinstance(x, (list, tuple)): return max(dtypes.from_py(xi) for xi in x) if x else dtypes.default_float
+    if isinstance(x, (list, tuple)): return strong_dtype(max(dtypes.from_py(xi) for xi in x)) if x else dtypes.default_float
     raise RuntimeError(f"Could not infer dtype of {x} with type {type(x)}")
   @staticmethod
   def finfo(dtype:DType) -> tuple[int, int]:
@@ -138,8 +138,10 @@ class dtypes:
   uchar = uint8; ushort = uint16; uint = uint32; ulong = uint64 # noqa: E702
   char = int8; short = int16; int = int32; long = int64 # noqa: E702
 
-  default_float: ClassVar[DType] = float32
-  default_int: ClassVar[DType] = int32
+  @property
+  def default_float(self) -> DType: return to_dtype(DEFAULT_FLOAT.value)
+  @property
+  def default_int(self) -> DType: return to_dtype(DEFAULT_INT.value)
 
   fp8_ocp = (fp8e4m3, fp8e5m2)
   fp8_fnuz = (fp8e4m3fnuz, fp8e5m2fnuz)
@@ -155,12 +157,12 @@ class dtypes:
   weaks = (weakint, weakfloat)
   all = floats + ints + (bool,) # noqa: A003
 
-if (env_default_float := getenv("DEFAULT_FLOAT", "")):
-  dtypes.default_float = getattr(dtypes, env_default_float.lower())
-  assert dtypes.is_float(dtypes.default_float), f"{env_default_float} is not a float dtype"
+dtypes = DTypes()
 
 DTypeLike = str|DType
 def to_dtype(dtype:DTypeLike) -> DType: return dtype if isinstance(dtype, DType) else getattr(dtypes, dtype.lower())
+assert dtypes.is_float(dtypes.default_float), f"{DEFAULT_FLOAT.value} is not a float dtype"
+assert dtypes.is_int(dtypes.default_int), f"{DEFAULT_INT.value} is not an int dtype"
 def strong_dtype(dtype:DType) -> DType:
   return {dtypes.weakint: dtypes.default_int, dtypes.weakfloat: dtypes.default_float}.get(dtype, dtype)
 
@@ -181,9 +183,10 @@ def _get_recursive_parents(dtype:DType) -> set[DType]:
 @functools.cache
 def least_upper_dtype(*ds:DType) -> DType:
   return min(set.intersection(*[_get_recursive_parents(d) for d in ds]))
-def least_upper_float(dt:DType) -> DType: return dt if dtypes.is_float(dt) else least_upper_dtype(dt, dtypes.default_float)
+def least_upper_float(dt:DType) -> DType:
+  return dtypes.weakfloat if dt is dtypes.weakint else dt if dtypes.is_float(dt) else least_upper_dtype(dt, dtypes.default_float)
 
-DTYPES_DICT = {k: v for k, v in dtypes.__dict__.items() if isinstance(v, DType) and not k.startswith(("default", "void", "weak", "_"))}
+DTYPES_DICT = {k: v for k, v in DTypes.__dict__.items() if isinstance(v, DType) and not k.startswith(("default", "void", "weak", "_"))}
 INVERSE_DTYPES_DICT = {**{v.name:k for k,v in DTYPES_DICT.items()}, "void": "void", "weakint":"weakint", "weakfloat":"weakfloat"}
 
 @functools.cache
@@ -207,7 +210,6 @@ def can_lossless_cast(dt0:DType, dt1:DType) -> bool:
 
 def sum_acc_dtype(dt:DType):
   # default acc dtype for sum
-  if dt in dtypes.weaks: return dt
   if dtypes.is_unsigned(dt): return least_upper_dtype(dt, dtypes.uint)
   if dtypes.is_int(dt) or dt == dtypes.bool: return least_upper_dtype(dt, dtypes.int)
   return least_upper_dtype(dt, to_dtype(getenv("SUM_DTYPE", "float32")))
@@ -302,6 +304,7 @@ def _from_np_dtype(npdtype:'np.dtype') -> DType: # type: ignore [name-defined] #
 @functools.cache
 def _to_torch_dtype(dtype:DType) -> 'torch.dtype'|None:  # type: ignore [name-defined] # noqa: F821
   import numpy as np, torch
+  dtype = strong_dtype(dtype)
   if dtype == dtypes.uint64: return torch.uint64
   if dtype == dtypes.bfloat16: return torch.bfloat16
   if dtype in dtypes.fp8s: return torch.uint8

tinygrad/engine/jit.py

@@ -1,8 +1,8 @@
 from typing import TypeVar, Generic, Callable, Any
-import functools, collections
+import functools
 from tinygrad.tensor import Tensor, all_tensors
 from tinygrad.helpers import flatten, merge_dicts, DEBUG, Context, BEAM, getenv, JIT, JIT_BATCH_SIZE, dedup, pluralize, VIZ, disable_gc
-from tinygrad.device import Buffer, Compiled, Device, MultiBuffer
+from tinygrad.device import Buffer, Compiled, Device, MultiBuffer, DepsTracker
 from tinygrad.dtype import DType
 from tinygrad.uop.ops import UOp, PatternMatcher, Variable, sym_infer, Ops, buffers, track_rewrites, graph_rewrite
 from tinygrad.renderer import Estimates
@@ -88,34 +88,6 @@ def _check_no_non_tensor_return(ret):
 
 def graph_class(dev): return dev.graph.func if isinstance(dev.graph, functools.partial) else dev.graph
 
-class DepsTracker:
-  def __init__(self):
-    # tracks (offset, end, dep) ranges per base buffer id to handle suballocated buffers correctly.
-    self.w_dependency_map: dict[int, list[tuple[int, int, Any]]] = collections.defaultdict(list)
-    self.r_dependency_map: dict[int, list[tuple[int, int, Any]]] = collections.defaultdict(list)
-
-  @staticmethod
-  def _key(buf:Any) -> tuple[Any, int, int]: return id(buf.base), buf.offset, buf.offset + buf.nbytes
-
-  def access_resources(self, bufs:list[Any], write:list[int], new_dependency:Any):
-    wait_nodes = []
-    for i,buf in enumerate(bufs):
-      key, s, e = self._key(buf)
-      wait_nodes += [dep for st,en,dep in self.w_dependency_map[key] if st < e and s < en]
-      if i in write: wait_nodes += [dep for st,en,dep in self.r_dependency_map[key] if st < e and s < en]
-    for i,buf in enumerate(bufs):
-      key, s, e = self._key(buf)
-      if i in write:
-        for dmap in [self.w_dependency_map, self.r_dependency_map]:
-          kept = []
-          for st,en,dep in dmap[key]:
-            if st < min(s, en): kept.append((st, min(s, en), dep))
-            if max(e, st) < en: kept.append((max(e, st), en, dep))
-          dmap[key] = kept
-        self.w_dependency_map[key].append((s, e, new_dependency))
-      else: self.r_dependency_map[key].append((s, e, new_dependency))
-    return list({id(x):x for x in wait_nodes}.values())
-
 class GraphRunner:
   def __init__(self, linear:UOp, input_uops:tuple[UOp, ...]=()):
     self.linear = linear.src[0]
@@ -234,8 +206,8 @@ def _prepare_jit_inputs(args, kwargs):
   for x in args + tuple(kwargs.values()):
     it = x if isinstance(x, (tuple,list)) else x.values() if isinstance(x, dict) else []
     tensors += [t for t in it if t.__class__ is Tensor and not any(t is y for y in tensors)]
-  def get_input_uops() -> list[UOp]: return flatten([t.uop.src if t.uop.op is Ops.MULTI else [t.uop] for t in tensors])
-  if any(u.device is None for u in get_input_uops()): raise JitError("JIT inputs must be real buffers; use .clone()")
+  def get_input_uops() -> list[UOp]: return flatten([[t.uop.src[0]] if t.uop.op is Ops.UNSHARD else [t.uop] for t in tensors])
+  if any(u.is_virtual for u in get_input_uops()): raise JitError("JIT inputs must be real buffers; use .clone()")
   if len(unrealized_tensors := [x for x in tensors if not x.uop.is_realized]): Tensor.realize(*unrealized_tensors)
   input_uops = get_input_uops()
   # collect buffer UOps (including MultiBuffer)

tinygrad/engine/realize.py

@@ -1,13 +1,15 @@
 from __future__ import annotations
 from typing import cast, Iterator, Any, Sequence
-import time, random, itertools, math, contextlib, weakref, array
+import time, random, itertools, math, contextlib, weakref, array, multiprocessing
+from concurrent.futures import ProcessPoolExecutor
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansilen, all_int, prod, flatten, Context, getenv, to_tuple
 from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events
-from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, sym_infer, buffers, graph_rewrite
+from tinygrad.helpers import PARALLEL_COMPILE, NUM_CPU_THREADS
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, buffers, graph_rewrite
 from tinygrad.device import Device, Buffer, MultiBuffer
 from tinygrad.renderer import Estimates
-from tinygrad.codegen import to_program
+from tinygrad.codegen import to_program, to_program_cache, program_cache_key, parallel_to_program
 from tinygrad.codegen.opt.postrange import bufs_from_ast
 
 # **************** Helpers ****************
@@ -148,8 +150,10 @@ def unwrap_multi(call:UOp, resolved:list[UOp]) -> Iterator[tuple[list[Buffer], d
   bufs = [b.buffer for b in resolved]
   if not any(isinstance(b, MultiBuffer) for b in bufs): yield cast(list[Buffer], bufs), {}
   else:
-    dnum = next((x.expr for x in call.src[0].variables() if x.expr == '_device_num'), None)
-    for j, per_dev in enumerate(zip(*[cast(MultiBuffer, b).bufs for b in bufs])): yield list(per_dev), {dnum: j} if dnum else {}
+    # the DEVICE axis is bound per device at launch: it's a RANGE in the AST and the _device_num variable after codegen
+    has_dnum = any((x.op is Ops.RANGE and x.arg[-1] is AxisType.DEVICE) or (x.op is Ops.PARAM and x.arg.name == '_device_num')
+                   for x in call.src[0].toposort())
+    for j, per_dev in enumerate(zip(*[cast(MultiBuffer, b).bufs for b in bufs])): yield list(per_dev), {"_device_num": j} if has_dnum else {}
 
 def exec_view(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
   resolved = resolve_params(call, ctx.input_uops)
@@ -207,14 +211,15 @@ def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
   return t[0]
 
 def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> float|None:
-  if call.arg.aux.inputs is not None:
+  if (inputs:=call.arg.aux.inputs) is not None:
     bufs = [_resolve(ctx.input_uops[i], ctx.input_uops).buffer for i in call.arg.aux.input_idxs]
+    table = call.src[1+inputs].buffer
     for j,dev in enumerate(call.arg.aux.device):
-      addrs = [(b.bufs[j] if isinstance(b, MultiBuffer) else b).get_buf(dev).va_addr for b in bufs]
-      buf = b.bufs[j] if isinstance(b:=call.src[1+call.arg.aux.inputs].buffer, MultiBuffer) else b
-      buf.ensure_allocated()._buf.cpu_view().view(fmt='Q')[:len(addrs)] = array.array('Q', addrs)
+      addrs = array.array('Q', [(b.bufs[j] if isinstance(b, MultiBuffer) else b).get_buf(dev).va_addr for b in bufs])
+      buf = table.bufs[j] if isinstance(table, MultiBuffer) else table
+      buf.ensure_allocated()._buf.cpu_view().view(fmt='Q')[:len(addrs)] = addrs
 
-  pm_exec.rewrite(call.replace(src=(ast,) + call.src[1:]), replace(ctx, update_stats=False))
+  exec_kernel(replace(ctx, update_stats=False), call, ast)
 
   st = time.perf_counter()
   for d in call.arg.aux.device:
@@ -260,20 +265,27 @@ pm_exec = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="validate", name="ast"),), name="call", allow_any_len=True), exec_validate),
 ])
 
+if getenv("HCQ2"): from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link # noqa: E402 # down here, hcq2 imports the helpers above
+
 def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, jit=False) -> UOp:
   if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
   if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
+  if jit and PARALLEL_COMPILE.value:
+    pending:dict[tuple, tuple[UOp, Any, tuple]] = {}
+    for call in linear.toposort():
+      if call.op is not Ops.CALL or call.src[0].op not in (Ops.SINK, Ops.PROGRAM): continue
+      renderer = Device[call.device if isinstance(call.device, str) else call.device[0]].renderer
+      key = program_cache_key(call.src[0], renderer)
+      if key not in to_program_cache: pending.setdefault(key, (call.src[0], renderer, key))
+    if len(pending) >= 16:
+      workers = min(PARALLEL_COMPILE.value, NUM_CPU_THREADS.value, len(pending))
+      with ProcessPoolExecutor(workers, mp_context=multiprocessing.get_context("spawn")) as pool:
+        for key,program in pool.map(parallel_to_program, pending.values()): to_program_cache[key] = program
   linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True)
-  if getenv("HCQ2"):
-    from extra.hcq2.hcq2 import hcq_compile
-    linear = hcq_compile(linear, input_uops, jit=jit)
+  if getenv("HCQ2"): linear = hcq_compile(linear, input_uops, jit=jit)
   return graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)
 
-def link_linear(linear:UOp, jit=False) -> UOp:
-  if getenv("HCQ2"):
-    from extra.hcq2.hcq2 import hcq_link
-    linear = hcq_link(linear, jit=jit)
-  return linear
+def link_linear(linear:UOp, jit=False, cache=True) -> UOp: return hcq_link(linear, jit=jit, cache=cache) if getenv("HCQ2") else linear
 
 def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:Sequence[UOp]=(), update_stats=True, jit=False, wait=False):
   inputs = list(input_uops)
@@ -287,5 +299,5 @@ def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None
     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)
-  call = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0)).src[0]
-  return pm_exec.rewrite(call, ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False))
+  ctx = ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False)
+  return pm_exec.rewrite(link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0), cache=ctx.cache).src[0], ctx)

tinygrad/helpers.py

@@ -232,6 +232,7 @@ class _DEV(ContextVar):
 DEV, DEBUG, BEAM, NOOPT = _DEV("DEV", ""), ContextVar("DEBUG", 0), ContextVar("BEAM", 0), ContextVar("NOOPT", 0)
 IMAGE, FLOAT16, OPENPILOT_HACKS = ContextVar("IMAGE", 0), ContextVar("FLOAT16", 0), ContextVar("OPENPILOT_HACKS", 0)
 JIT, JIT_BATCH_SIZE = ContextVar("JIT", 2 if OSX and ARCH_X86 else 1), ContextVar("JIT_BATCH_SIZE", 32)
+PARALLEL_COMPILE = ContextVar("PARALLEL_COMPILE", 0)
 CHUNK_SIZE = 2**20  # TinyFS content-addressed store: blob chunk + hash-tree node granularity
 WINO, CAPTURING, TRACEMETA, NO_COLOR = ContextVar("WINO", 0), ContextVar("CAPTURING", 1), ContextVar("TRACEMETA", 1), ContextVar("NO_COLOR", 0)
 TRAINING = ContextVar("TRAINING", 0)
@@ -247,8 +248,20 @@ FUSE_OPTIM = ContextVar("FUSE_OPTIM", 0)
 ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0)
 MAX_KERNEL_BUFFERS = ContextVar("MAX_KERNEL_BUFFERS", 0)
 EMULATED_DTYPES = ContextVar("EMULATED_DTYPES", "")
+DEFAULT_FLOAT, DEFAULT_INT = ContextVar("DEFAULT_FLOAT", "float32"), ContextVar("DEFAULT_INT", "int32")
 CAPTURE_PROCESS_REPLAY = ContextVar("CAPTURE_PROCESS_REPLAY", 0)
-CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)))
+def _get_cpu_count() -> int:
+  # os.process_cpu_count (3.13+) respects cgroup limits
+  if hasattr(os, "process_cpu_count"): return max(1, os.process_cpu_count())
+  # cgroup v2 (containers with --cpus=N)
+  try:
+    with open("/sys/fs/cgroup/cpu.max") as f:
+      quota, period = f.read().strip().split()
+      if quota != "max": return max(1, int(quota) // int(period))
+  except (FileNotFoundError, ValueError, ZeroDivisionError): pass
+  # fall back to affinity (respects taskset but not cgroup quota)
+  return max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1))
+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)

tinygrad/llm/cli.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-import sys, argparse, codecs, itertools, typing, re, unicodedata, json, time
+import sys, os, argparse, codecs, itertools, typing, re, unicodedata, json, time
 from typing import TYPE_CHECKING
 from tinygrad import nn
 from tinygrad.uop.ops import UOp, Ops
@@ -85,8 +85,8 @@ models = {
   "qwen3.5:0.8b": "https://huggingface.co/unsloth/Qwen3.5-0.8B-GGUF/resolve/main/Qwen3.5-0.8B-Q8_0.gguf",
   "qwen3.5:4b": "https://huggingface.co/unsloth/Qwen3.5-4B-GGUF/resolve/main/Qwen3.5-4B-Q4_K_M.gguf",
   "qwen3.5:9b": "https://huggingface.co/unsloth/Qwen3.5-9B-GGUF/resolve/main/Qwen3.5-9B-Q4_K_M.gguf",
-  "qwen3.5:27b": "https://huggingface.co/unsloth/Qwen3.5-27B-GGUF/resolve/main/Qwen3.5-27B-Q4_K_M.gguf",
-  "qwen3.5:35b-a3b": "https://huggingface.co/unsloth/Qwen3.5-35B-A3B-GGUF/resolve/main/Qwen3.5-35B-A3B-Q4_K_M.gguf",
+  "qwen3.6:27b": "https://huggingface.co/unsloth/Qwen3.6-27B-GGUF/resolve/main/Qwen3.6-27B-Q4_K_M.gguf",
+  "qwen3.6:35b-a3b": "https://huggingface.co/unsloth/Qwen3.6-35B-A3B-GGUF/resolve/main/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf",
   "olmoe": "https://huggingface.co/allenai/OLMoE-1B-7B-0924-Instruct-GGUF/resolve/main/olmoe-1b-7b-0924-instruct-q4_k_m.gguf",
   "moonlight": "https://huggingface.co/gabriellarson/Moonlight-16B-A3B-Instruct-GGUF/resolve/main/Moonlight-16B-A3B-Instruct-Q4_K_M.gguf",
   "glm-4.7-flash": "https://huggingface.co/unsloth/GLM-4.7-Flash-GGUF/resolve/main/GLM-4.7-Flash-Q4_K_M.gguf",
@@ -112,7 +112,7 @@ class FallbackTemplate:
     if self.tok.preset == 'glm4': return ""
     if self.tok.preset == 'tekken': return "[/INST]"
     return self.tok.decode([self.tok.eos_id])
-  def render(self, messages:list[dict], tools=None, add_generation_prompt:bool=True) -> str:
+  def render(self, messages:list[dict], tools=None, add_generation_prompt:bool=True, **kwargs) -> str:
     out = self.tok.decode([] if self.tok.bos_id is None else [self.tok.bos_id]) + ("<sop>" if self.tok.preset == 'glm4' else "")
     for msg in messages:
       out += self.role(msg["role"])
@@ -136,6 +136,7 @@ def main():
   parser.add_argument("--warmup", action="store_true", help="warmup the JIT")
   parser.add_argument("--benchmark", nargs='?', type=int, const=20, metavar="COUNT", help="Benchmark tok/s (optional count, default 20)")
   args = parser.parse_args()
+  if args.warmup or args.serve: os.environ.setdefault("DISK_SCACHE", "1")
 
   # load the model
   model, kv = Transformer.from_gguf(fetch(models.get(args.model, args.model)), args.max_context)
@@ -163,9 +164,8 @@ def main():
 
   # warmup the JIT
   if args.warmup or args.serve:
-    # run 2 tokens through the model twice to capture the JIT before serving
-    with Context(DEBUG=max(DEBUG.value, 1)):
-      for _ in range(2): list(zip(range(2), model.generate([0])))
+    with Context(DEBUG=max(DEBUG.value, 1), PARALLEL_COMPILE=getenv("PARALLEL_COMPILE", 12)):
+      model.warmup()
 
   # start server
   if args.serve: LLMServer(('', args.serve), model, model_name, tok, template).serve_forever()

tinygrad/llm/gguf.py

@@ -1,9 +1,11 @@
-import functools, io, pathlib, re, struct
+import functools, io, pathlib, re, struct, weakref, mmap
 from typing import Any, Callable
 
 from tinygrad.tensor import Tensor
+from tinygrad.uop.ops import UOp
 from tinygrad.dtype import dtypes
-from tinygrad.helpers import prod, round_up
+from tinygrad.helpers import prod, round_up, mv_address
+from tinygrad.device import Device
 from tinygrad.nn.state import TensorIO
 
 # ggml packs each iq grid entry as N bytes (N=4 for uint32 grids, N=8 for uint64 grids) in a single word. See ggml-common.h.
@@ -20,7 +22,25 @@ _GGML_NATIVE = {0: dtypes.float32, 1: dtypes.float16, 24: dtypes.int8, 25: dtype
 _GGML_QUANT = {2:(32,18), 3:(32,20), 6:(32,22), 7:(32,24), 8:(32,34),
                12:(256,144), 13:(256,176), 14:(256,210), 18:(256,98), 21:(256,110), 22:(256,82), 23:(256,136), 39:(32,17), 41:(128,18)}
 
-def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
+_quantized_tensors:weakref.WeakKeyDictionary[UOp, tuple[UOp, int]] = weakref.WeakKeyDictionary()
+_cpu_mapped_ggufs:dict[tuple[pathlib.Path, int, int], tuple[mmap.mmap, Tensor]] = {}
+
+def _gguf_tensor(path:pathlib.Path) -> Tensor:
+  path = path.resolve()
+  if not Device.DEFAULT.startswith("CPU"): return Tensor(path).to(None)
+  stat = path.stat()
+  key = (path, stat.st_mtime_ns, stat.st_size)
+  if key not in _cpu_mapped_ggufs:
+    with path.open("rb") as f: mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_COPY)
+    _cpu_mapped_ggufs[key] = mm, Tensor.from_blob(mv_address(memoryview(mm)), (len(mm),), dtype=dtypes.uint8, device=Device.DEFAULT)
+  return _cpu_mapped_ggufs[key][1]
+
+def get_ggml_quantization(tensor:Tensor) -> tuple[Tensor, int]|None:
+  if (meta:=_quantized_tensors.get(tensor.uop)) is None: return None
+  packed, ggml_type = meta
+  return Tensor(packed), ggml_type
+
+def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int, contiguous:bool=True) -> Tensor:
   """
   Converts ggml tensor data to a tinygrad tensor.
 
@@ -42,7 +62,8 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
 
   if (nelements_nbytes := _GGML_QUANT.get(ggml_type)) is not None:
     from tinygrad.runtime.autogen import ggml_common as _ggml
-    blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1])).contiguous()
+    blocks = t[:(n//nelements_nbytes[0])*nelements_nbytes[1]].reshape((-1, nelements_nbytes[1]))
+    if contiguous: blocks = blocks.contiguous()
     if ggml_type == 2: return (q_to_uint8(blocks[:,2:], 4).bitcast(dtypes.int8) - 8) * blocks[:,:2].bitcast(dtypes.float16).cast(dtypes.float32)
     if ggml_type == 3:
       d, m = (blocks[:,s:s+2].bitcast(dtypes.float16).cast(dtypes.float32) for s in [ 0, 2 ])
@@ -105,7 +126,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
     if ggml_type == 39:
       e = blocks[:, 0].cast(dtypes.uint32)
       small_bits = Tensor([0x00200000, 0x00400000], dtype=dtypes.uint32, device=t.device)[e.clip(0, 1).cast(dtypes.int32)] # e = 0 or e = 1 case
-      d = (e < 2).where(small_bits, ((e - 1) * 0x00800000).cast(dtypes.uint32)).bitcast(dtypes.float32).unsqueeze(-1)
+      d = (e < 2).where(small_bits, (e - 1) * 0x00800000).bitcast(dtypes.float32).unsqueeze(-1)
       codes = q_to_uint8(blocks[:, 1:17], 4)
       fp4_lut = Tensor([0.0, 1.0, 2.0, 3.0, 4.0, 6.0, 8.0, 12.0,
                        -0.0,-1.0,-2.0,-3.0,-4.0,-6.0,-8.0,-12.0],
@@ -130,8 +151,7 @@ readers: dict[int, Callable[[io.BufferedIOBase], Any]] = { 8: read_str, 9: read_
 read_uint32, read_int32, read_uint64, read_int64 = readers[4], readers[5], readers[10], readers[11]
 
 def _gguf_parse(tensor: Tensor) -> tuple[dict, dict[str, Tensor]]:
-  # TODO: remove the need for copy to default device
-  tensor = tensor.to(None).realize()
+  tensor = tensor.realize()
   r = io.BufferedReader(TensorIO(tensor), 1_000_000)
   magic, version, n_tensors, n_kv = r.read(4), read_int32(r), read_int64(r), read_int64(r)
   if magic != b"GGUF" or version not in [2, 3]: raise ValueError("Invalid GGUF format!")
@@ -145,7 +165,14 @@ def _gguf_parse(tensor: Tensor) -> tuple[dict, dict[str, Tensor]]:
   alignment, pos = kv_data.get("general.alignment", 32), r.tell()
   data_start = round_up(pos, alignment)
 
-  state_dict = {name: ggml_data_to_tensor(tensor[data_start + off:], prod(dims), typ).reshape(*reversed(dims)) for name, dims, typ, off in t_infos}
+  state_dict = {}
+  for name, dims, typ, off in t_infos:
+    n, shape = prod(dims), tuple(reversed(dims))
+    decoded = ggml_data_to_tensor(data:=tensor[data_start + off:], n, typ).reshape(*shape)
+    if typ in _GGML_QUANT:
+      block_size, type_size = _GGML_QUANT[typ]
+      _quantized_tensors[decoded.uop] = (data[:n//block_size*type_size].uop, typ)
+    state_dict[name] = decoded
   return kv_data, state_dict
 
 def _gguf_split_paths(path: pathlib.Path, kv: dict) -> list[pathlib.Path]:
@@ -169,8 +196,8 @@ def gguf_load(fn: Tensor|str|pathlib.Path) -> tuple[dict, dict[str, Tensor]]:
 
   NOTE: The provided tensor must be on a device that supports execution.
   """
-  kv, sd = _gguf_parse(fn if isinstance(fn, Tensor) else Tensor(pathlib.Path(fn)))
+  kv, sd = _gguf_parse(fn if isinstance(fn, Tensor) else _gguf_tensor(pathlib.Path(fn)))
   if kv.get('split.count', 1) <= 1: return kv, sd
   if isinstance(fn, Tensor): raise ValueError("multi-part GGUF requires a path argument (got Tensor)")
-  for pp in _gguf_split_paths(pathlib.Path(fn), kv)[1:]: sd.update(_gguf_parse(Tensor(pp))[1])
+  for pp in _gguf_split_paths(pathlib.Path(fn), kv)[1:]: sd.update(_gguf_parse(_gguf_tensor(pp))[1])
   return kv, sd

tinygrad/llm/kernels/__init__.py

@@ -0,0 +1 @@
+"""Custom kernels used by tinygrad.llm."""