benchmarks: cleanup llm

fix: external_test_gpu_crash on python 3.14 (#17739 )
pin onnxruntime==1.24.1 (#17738 )
2026-08-26 08:46:06 +00:00 · 2026-08-25 19:19:16 -07:00 · 2026-08-25 15:43:30 -07:00 · 2026-08-25 17:53:15 -04:00 · 2026-08-25 14:11:53 -07:00 · 2026-08-25 16:21:08 -04:00
89 changed files with 758 additions and 889 deletions
@@ -82,13 +82,16 @@ jobs:
  #      pytest -nauto --durations=20

  llmbenchmark:
-    name: LLM (DEV=${{ matrix.dev }})
+    name: Benchmark ${{ matrix.model }} (DEV=${{ matrix.dev }})
    runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
    strategy:
      fail-fast: false
      matrix:
        dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 60
+        model: ['llama3.2:3b-f16', 'qwen3.8:27b', 'olmoe']
+        # qwen3.8:27b doesn't fit on mac
+        exclude: [{ dev: 'METAL', model: 'qwen3.8:27b' }, { dev: 'AMD', model: 'olmoe' }, { dev: 'NV', model: 'olmoe' }]
+    timeout-minutes: 15
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -102,12 +105,11 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./extra/hcq/hcq_smi.py amd rmmod --expect
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -115,16 +117,9 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: python3 test/external/process_replay/reset.py
-    - name: Run 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.8
-      # qwen3.8:27b doesn't fit on mac
-      if: ${{ matrix.dev != 'METAL' }}
-      run: BENCHMARK_LOG=qwen38_27b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.8:27b --benchmark --warmup
-    - name: Run olmoe
-      # just metal for now
-      if: ${{ matrix.dev == 'METAL' }}
-      run: BENCHMARK_LOG=olmoe JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m olmoe --benchmark --warmup
+    - name: Run ${{ matrix.model }}
+      run: |
+        BENCHMARK_LOG="${{ matrix.model }}" JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 extra/benchmark_llm.py --model ${{ matrix.model }}
    - name: Run process replay tests
      uses: ./.github/actions/process-replay

@@ -135,7 +130,7 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 10
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -149,12 +144,11 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./extra/hcq/hcq_smi.py amd rmmod --expect
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -184,7 +178,7 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 5
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -198,12 +192,11 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./extra/hcq/hcq_smi.py amd rmmod --expect
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
    - name: Symlink models and datasets
      run: |
        mkdir -p extra/datasets
@@ -227,7 +220,7 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 15
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -241,12 +234,11 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./extra/hcq/hcq_smi.py amd rmmod --expect
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -273,7 +265,7 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 20
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -287,12 +279,11 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./extra/hcq/hcq_smi.py amd rmmod --expect
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
    - name: Symlink models and datasets
      run: |
        mkdir -p weights
@@ -327,7 +318,7 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 60
+    timeout-minutes: 10
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -340,12 +331,11 @@ jobs:
    - name: Setup (AMD)
      if: ${{ matrix.dev == 'AMD' }}
      run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py amd rmmod
-        ./extra/hcq/hcq_smi.py amd kill_pids
+        ./extra/hcq/hcq_smi.py amd rmmod --expect
+        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
    - name: Setup (NV)
      if: ${{ matrix.dev == 'NV' }}
-      run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
+      run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
    - name: setup staging db
      if: github.ref == 'refs/heads/update_benchmark_staging'
      run: |
@@ -422,7 +412,7 @@ jobs:
  testusbgpu:
    name: UsbGPU Benchmark
    runs-on: [self-hosted, macOS]
-    timeout-minutes: 10
+    timeout-minutes: 3
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -524,15 +514,6 @@ jobs:
        shell: bash -e -o pipefail {0}
    if: github.repository_owner == 'tinygrad'
    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: setup staging db
-      if: github.ref == 'refs/heads/update_benchmark_staging'
-      run: |
-        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
-        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
-    - name: reset process replay
-      run: test/external/process_replay/reset.py
    - name: Checkout Code
      uses: actions/checkout@v6
    - name: setup staging db
@@ -556,7 +537,7 @@ jobs:
  testcommausbgpubenchmark:
    name: UsbGPU Benchmark (comma)
    runs-on: [self-hosted, Linux, comma4]
-    timeout-minutes: 20
+    timeout-minutes: 10
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -585,7 +566,7 @@ jobs:
      fail-fast: false
      matrix:
        dev: ['AMD', 'NV']
-    timeout-minutes: 20
+    timeout-minutes: 5
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -597,9 +578,8 @@ jobs:
      uses: actions/checkout@v6
    - name: Setup
      run: |
-        ./extra/amdpci/setup_python_cap.sh
-        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} rmmod
-        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} kill_pids
+        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} rmmod --expect
+        ./extra/hcq/hcq_smi.py ${{ matrix.dev }} kill_pids --sudoless
        mkdir -p extra/datasets
        ln -s /raid/datasets/imagenet extra/datasets/imagenet
    - name: setup staging db
@@ -659,7 +639,7 @@ jobs:
  llvmspeed:
    name: LLVM Speed
    runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 20
+    timeout-minutes: 5
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -4,3 +4,4 @@
 - Run `python -m mypy tinygrad/` to typecheck
 - Run `python -m ruff check .` to lint
 - Read `./tinygrad/viz/README.md` for profiling and debugging rewrite rules
+- Do not do amend commits. Always do a new commit if a force push to origin would be required.
@@ -15,7 +15,7 @@ def stochastic_round_bf16(x:Tensor) -> Tensor:
  bits = x.bitcast(dtypes.uint32)
  if isinstance(x.device, tuple):
    shape = x.uop.shard_shape if x.uop.axis is not None else x.shape
-    noise = Tensor(UOp(Ops.MSTACK, dtypes.default_float, tuple(Tensor.rand(*shape, device=d).uop for d in x.device)))
+    noise = Tensor(UOp(Ops.MSTACK, src=tuple(Tensor.rand(*shape, device=d).uop for d in x.device)))
  else:
    noise = x.rand_like()
  noise = (noise * 0xFFFF).cast(dtypes.uint32)
@@ -35,7 +35,7 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=2048 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
 export SEQLEN=${SEQLEN:-8192}
@@ -46,7 +46,7 @@ export DATA_SEED=${DATA_SEED:-5760}
 export JITBEAM=${JITBEAM:-3}
 export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1

-export FAKEDATA=${FAKEDATA:-$([[ "$DEV" == NULL:* ]] && echo 1 || echo 0)} BENCHMARK=${BENCHMARK:-10}
+export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
 if [ -z "$FULL_LAYERS" ]; then
  export LLAMA_LAYERS=${LLAMA_LAYERS:-2}
 fi
@@ -35,7 +35,7 @@ export BASEDIR="/raid/datasets/c4-8b/"
 export SMALL=1
 export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
 export EVAL_TARGET=3.3 EVAL_FREQ=12288
-export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
+export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=2048 MAX_STEPS=1200000
 export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
 export SAMPLES=$((MAX_STEPS * GBS))
 export SEQLEN=${SEQLEN:-8192}
@@ -17,7 +17,7 @@ export USE_ATOMICS=1
 export ASM_GEMM=1
 export WQKV=1
 export MASTER_WEIGHTS=1
-export FP8=1
+export MXFP4=1
 export ALLREDUCE_CAST=1
 export FAST_CE=1
 export FUSED_INPUT_QUANTIZE=1
@@ -26,7 +26,7 @@ export FUSED_ADD_NORM_MUL_QUANTIZE=1
 export FUSED_SILU_W13=1
 export SPLIT_W13=0

-export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
+export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="float32"
 export DP=8 MP=1 BS=16 EVAL_BS=8 GRADIENT_ACC_STEPS=2
 export GBS=$((BS * GRADIENT_ACC_STEPS))

@@ -18,6 +18,7 @@ class BenchEvent(Enum):
  MLPERF_RUN = "mlperf_run"
 class InstantBenchEvent(Enum):
  GFLOPS = "gflops"
+  TPS = "tps"

 _events = {}
 def clear_events():
@@ -35,7 +36,7 @@ class WallTimeEvent:
    return self
  def __exit__(self, *_):
    self.time = time.monotonic() - self.start
-    _events[self.event]["wall"].append(self.time)
+    _events[self.event]["wall"].append((self.time, BENCHMARK_LOG.value))
    return False

 class KernelTimeEvent:
@@ -47,19 +48,19 @@ class KernelTimeEvent:
    self.start = GlobalCounters.time_sum_s
    return self
  def __exit__(self, *_):
-    _events[self.event]["kernel"].append(GlobalCounters.time_sum_s - self.start)
+    _events[self.event]["kernel"].append((GlobalCounters.time_sum_s - self.start, BENCHMARK_LOG.value))
    return False

 def log_event_instant(event:InstantBenchEvent, value:float):
-  _events[event].append(value)
+  _events[event].append((value, BENCHMARK_LOG.value))

 if BENCHMARK_LOG:
  INFLUXDB_HOST = getenv("INFLUXDB_HOST", "")
  INFLUXDB_ORG = getenv("INFLUXDB_ORG", "tiny")
  INFLUXDB_TOKEN = getenv("INFLUXDB_TOKEN", "")

-  def _create_point(run_id, i, attempt, ref, commit, name, value, run):
-    point = Point(BENCHMARK_LOG.value).tag("id", run_id).tag("index", i)
+  def _create_point(run_id, i, attempt, ref, commit, name, value, log_name, run):
+    point = Point(log_name.replace(':', '_').replace('.', '_')).tag("id", run_id).tag("index", i)
    point = point.tag("device", Device.DEFAULT)
    point = point.tag("attempt", attempt).tag("ref", ref).tag("commit", commit)
    point = point.field(name, value).field("x", run)
@@ -91,12 +92,12 @@ if BENCHMARK_LOG:
      run_id = str(uuid.uuid4())
      if isinstance(event, BenchEvent):
        for event_type, values in _events[event].items():
-          for i, value in enumerate(values):
-            point = _create_point(run_id, i, attempt, ref, commit, f"{event.value}_{event_type}", value, run)
+          for i, (value, log_name) in enumerate(values):
+            point = _create_point(run_id, i, attempt, ref, commit, f"{event.value}_{event_type}", value, log_name, run)
            points.append(point)
      else:
-        for i, value in enumerate(_events[event]):
-          point = _create_point(run_id, i, attempt, ref, commit, event.value, value, run)
+        for i, (value, log_name) in enumerate(_events[event]):
+          point = _create_point(run_id, i, attempt, ref, commit, event.value, value, log_name, run)
          points.append(point)

    write_options = WriteOptions(write_type=WriteType.synchronous, retry_interval=5000, max_retries=5, max_retry_delay=30000, exponential_base=2)
@@ -1,5 +1,8 @@
 import argparse, time
+from tinygrad.helpers import Context, fetch
 from tinygrad.llm.model import Transformer
+from tinygrad.llm.cli import models
+from extra.bench_log import BENCHMARK_LOG, InstantBenchEvent, log_event_instant

 if __name__ == "__main__":
  parser = argparse.ArgumentParser()
@@ -11,7 +14,7 @@ if __name__ == "__main__":
  args = parser.parse_args()

  st = time.perf_counter()
-  model, _ = Transformer.from_gguf(args.model, args.max_context)
+  model, _ = Transformer.from_gguf(fetch(models.get(args.model, args.model)), args.max_context)
  print(f"load {time.perf_counter()-st:.3f}s", flush=True)

  st = time.perf_counter()
@@ -25,7 +28,11 @@ if __name__ == "__main__":
  output = [next(gen)]
  pt = time.perf_counter()
  print(f"prefill {args.prompt_tokens/(pt-st):.3f} tok/s", flush=True)
+  if BENCHMARK_LOG:
+    with Context(BENCHMARK_LOG=f"{BENCHMARK_LOG.value}_prefill"): log_event_instant(InstantBenchEvent.TPS, args.prompt_tokens/(pt-st))

  for _ in range(args.decode_tokens): output.append(next(gen))
  et = time.perf_counter()
  print(f"decode {args.decode_tokens/(et-pt):.3f} tok/s output {output}", flush=True)
+  if BENCHMARK_LOG:
+    with Context(BENCHMARK_LOG=f"{BENCHMARK_LOG.value}_decode"): log_event_instant(InstantBenchEvent.TPS, args.decode_tokens/(et-pt))
@@ -53,7 +53,7 @@ def _ggather_bwd(gradient:UOp, kernel:UOp) -> tuple:
    g, m, j, jo, ji = _kv_ranges(Gk, M, Dk, _blk_for(Dk))
    row = idx.index(g, m).cast(dtypes.weakint)
    val = gout.index(g, m, j).load().cast(dtypes.float32)
-    atomic = UOp(Ops.CUSTOM, dtypes.void, (gtab.index(g, row, j), val), arg=atomic_str)
+    atomic = UOp(Ops.CUSTOM, src=(gtab.index(g, row, j), val), arg=atomic_str)
    return atomic.end(g, m, jo, ji).sink(arg=KernelInfo(name=f"ggather_bwd_{M}_{Dk}", opts_to_apply=()))
  grad_table = Tensor.custom_kernel(gt, go, Tensor(idx_u, device=dev), fxn=_bwd_kernel)[0]
  return (None, grad_table.cast(table_u.dtype).uop, None)
@@ -16,9 +16,12 @@ def _do_reset_device(pci_bus): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/dev
 def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")

 def cmd_remove_module(args):
-  modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia", "ast"] if args.backend == "nv" else ["amdgpu"]
+  modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia"] if args.backend == "nv" else ["amdgpu"]
  to_unload = [m for m in modules if _is_module_loaded(m)]
  if not to_unload: print("Kernel modules are not loaded")
+  elif getattr(args, "expect", False):
+    print(f"Kernel modules are loaded: {to_unload}")
+    sys.exit(1)
  else:
    print("Removing kernel modules:", ", ".join(to_unload))
    try: subprocess.run(["sudo", "modprobe", "-r", *to_unload], check=True)
@@ -60,17 +63,19 @@ def cmd_show_pids(args):

 def cmd_kill_pids(args):
  devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
+  use_sudo = not getattr(args, "sudoless", False)

  for dev in devs:
    for i in range(128):
      if i > 0: time.sleep(0.2)

      try:
-        try: pid = subprocess.check_output(['sudo', 'lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
+        try: pid = subprocess.check_output((['sudo'] if use_sudo else []) +
+                                           ['lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
        except subprocess.CalledProcessError: break

        print(f"Killing process {pid} (which uses {dev})")
-        subprocess.run(['sudo', 'kill', '-9', pid], check=True)
+        subprocess.run((['sudo'] if use_sudo else []) + ['kill', '-9', pid], check=True)
      except subprocess.CalledProcessError as e:
        print(f"Failed to kill process for device {dev}: {e}", file=sys.stderr)

@@ -79,6 +84,7 @@ def add_common_commands(parent_subparsers):
  p_insmod.set_defaults(func=cmd_insert_module)

  p_rmmod = parent_subparsers.add_parser("rmmod", help="Remove a kernel module")
+  p_rmmod.add_argument("--expect", action="store_true", help="Just assert that module is already unloaded")
  p_rmmod.set_defaults(func=cmd_remove_module)

  p_reset = parent_subparsers.add_parser("reset", help="Reset a device")
@@ -91,6 +97,7 @@ def add_common_commands(parent_subparsers):

  p_reset = parent_subparsers.add_parser("kill_pids", help="Kill pids of processes using the device")
  p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device")
+  p_reset.add_argument("--sudoless", action="store_true", help="Do not use sudo when detecting or killing pids")
  p_reset.set_defaults(func=cmd_kill_pids)

 if __name__ == "__main__":
@@ -87,7 +87,7 @@ def release_mem(ctx, address=0x0, value=0, data_sel=0, int_sel=2, ctxid=0, cache

 def memory_barrier(ctx):
  pf = '' if ctx.nbio.version[0] == 2 else '0' if ctx.nbio.version[:2] != (7, 11) else '1'
-  return UOp(Ops.LINEAR, dtypes.void, (
+  return UOp(Ops.LINEAR, src=(
    wait_reg_mem(ctx, reg=getattr(ctx.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_REQ').addr[0],
                 reg_done=getattr(ctx.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff),
    acquire_mem(ctx)))
@@ -135,7 +135,7 @@ def pm4_program(ctx, call, prg):
    wreg(ctx, ctx.gc.regCOMPUTE_START_X, 0, 0, 0, *(info.local_size or (1, 1, 1)), 0, 0),
    pkt3(ctx, PM4Ops.DISPATCH_DIRECT, *info.global_size, dispatch_init),
    pkt3(ctx, PM4Ops.EVENT_WRITE, ctx.pm4.EVENT_TYPE(ctx.soc.CS_PARTIAL_FLUSH) | ctx.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH))]
-  return UOp(Ops.LINEAR, dtypes.void, tuple(ins))
+  return UOp(Ops.LINEAR, src=tuple(ins))

 pm_pm4_opsel = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), pm4_program),
@@ -207,7 +207,7 @@ def sdma_timestamp(ctx, ins, dst):
 pm_sdma_opsel = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy),

-  (UPat(Ops.INS, arg="barrier"), lambda: UOp(Ops.NOOP, dtypes.void, ())),
+  (UPat(Ops.INS, arg="barrier"), lambda: UOp(Ops.NOOP)),
  (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val")), name="ins"), sdma_wait),
  (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),), name="ins"), sdma_timestamp),
  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val")), name="ins"), sdma_store),
@@ -50,7 +50,7 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:
  else: raise NotImplementedError(f"no atomic max for device {device}")
  amax_idx = amax_out.reshape((1,)).index(UOp.const(0))
  max_val = lds[0].load()
-  atomic = UOp(Ops.CUSTOM, dtypes.void, (amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=atomic_arg)
+  atomic = UOp(Ops.CUSTOM, src=(amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=atomic_arg)
  return atomic.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=()))

@functools.cache
@@ -12,7 +12,7 @@ def _custom_quantize_mxfp4(row_fp4:UOp, row_scale:UOp, col_fp4:UOp, col_scale:UO
  mem = M*N*2 + M*N + M*N//16  # read bf16, write row+col fp4 + e8m0
  outputs = (row_fp4, row_scale, col_fp4, col_scale)
  sink = UOp.sink(*(o.base for o in outputs), x.base,
-                  *(UOp(Ops.CUSTOM, dtypes.void, (o.base.index(0),), arg="") for o in outputs),
+                  *(UOp(Ops.CUSTOM, src=(o.base.index(0),), arg="") for o in outputs),
                  UOp.special(256, "lidx0"), UOp.special(M//128, "gidx0"), UOp.special(N//64, "gidx1"),
                  arg=KernelInfo(name, estimates=Estimates(ops=12*M*N, mem=mem)))
  src = (pathlib.Path(__file__).parent/"quantize_mxfp4.cpp").read_text()
@@ -1,4 +1,4 @@
-import os, subprocess, sys, shlex
+import os, subprocess, sys, shlex, pickle
 from pathlib import Path
 from tinygrad.helpers import temp, getenv

@@ -23,5 +23,8 @@ if __name__ == "__main__":
      # AM_RESET=1 gets a clear trace, does not work on mi300 machines
      subprocess.run([sys.executable, *shlex.split(test)], cwd=EXAMPLES_DIR.parent.parent.parent,
                     env={**os.environ, "DEV":"AMD", "AM_RESET":"1" if not arch.startswith("gfx9") else "0", "VIZ":"-2", "PYTHONPATH":"."})
+      with open(PROFILE_PATH, "rb") as f: events = pickle.load(f)
+      with open(PROFILE_PATH, "wb") as f:
+        pickle.dump([e for e in events if type(e).__name__ in {"ProfilePMCEvent", "ProfileSQTTEvent", "ProfileProgramEvent"}], f)
      PROFILE_PATH.rename(dest:=EXAMPLES_DIR/arch/f"profile_{name}_run_{i}.pkl")
      print(f"saved SQTT trace to {dest}")
@@ -1,6 +1,6 @@
 [project]
 name = "tinygrad"
-version = "0.13.0"
+version = "0.14.0"
 description = "You like pytorch? You like micrograd? You love tinygrad! <3"
 authors = [{ name = "George Hotz" }]

@@ -84,7 +84,7 @@ testing = [
  "pillow",
  "onnx==1.19.0",
  "onnx2torch",
-  "onnxruntime",
+  "onnxruntime==1.24.1",
  "opencv-python",
  "transformers",
  "sentencepiece",
@@ -88,7 +88,6 @@ def run_rocprof_decoder(blobs: list[bytes], lib: bytes, base: int, target: str):
  if t.is_alive(): raise RuntimeError("rocprof decoder timeout")
  return occupancy_records, wave_insts

-@unittest.skip("TODO: fix to not require unpickling UOps.")
 class SQTTExamplesTestBase(unittest.TestCase):
  target: str
  examples: dict
@@ -7,7 +7,7 @@ from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops
 from tinygrad.renderer.isa import IselContext

 # INDEX on a register value with a constant index extracts a single element (the old GEP)
-def lane(y:UOp, i:int) -> UOp: return y.index(UOp.cconst(i, dtypes.int), dtype=y.dtype)
+def lane(y:UOp, i:int) -> UOp: return y.index(UOp.cconst(i, dtypes.int))

@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "only x86")
 class TestIselX86(unittest.TestCase):
@@ -865,9 +865,9 @@ class TestOps(unittest.TestCase):
                   lambda: (ten << Tensor([0,2,4], dtype=dtypes.uint32)).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor.__lshift__(2), lambda: ten.__lshift__(2).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor.bitwise_left_shift(2), lambda: ten.lshift(2).cast(dtypes.int32), forward_only=True)
-    self.helper_test_exception([], lambda: torch.tensor([1.0]) << 2, lambda: Tensor([1.0]) << 2, expected=RuntimeError)
-    self.helper_test_exception([], lambda: tor << torch.tensor([1.0]), lambda: ten << Tensor([1.0]), expected=RuntimeError)
-    self.helper_test_exception([], lambda: tor << 1.0, lambda: ten << 1.0, expected=RuntimeError)
+    self.helper_test_exception([], lambda: torch.tensor([1.0]) << 2, lambda: (Tensor([1.0]) << 2).realize(), expected=RuntimeError)
+    self.helper_test_exception([], lambda: tor << torch.tensor([1.0]), lambda: (ten << Tensor([1.0])).realize(), expected=RuntimeError)
+    self.helper_test_exception([], lambda: tor << 1.0, lambda: (ten << 1.0).realize(), expected=RuntimeError)

  def test_rshift(self):
    data = [[0,1,2],[1<<8,1<<16,1<<31-1]]
@@ -881,8 +881,8 @@ class TestOps(unittest.TestCase):
                   lambda: (ten >> Tensor([0,2,4], dtype=dtypes.uint32)).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor.__rshift__(2), lambda: ten.__rshift__(2).cast(dtypes.int32), forward_only=True)
    helper_test_op([], lambda: tor.bitwise_right_shift(2), lambda: ten.rshift(2).cast(dtypes.int32), forward_only=True)
-    self.helper_test_exception([], lambda: torch.tensor([4.0]) >> 1, lambda: Tensor([4.0]) >> 1, expected=RuntimeError)
-    self.helper_test_exception([], lambda: tor >> torch.tensor([1.0]), lambda: ten >> Tensor([1.0]), expected=RuntimeError)
+    self.helper_test_exception([], lambda: torch.tensor([4.0]) >> 1, lambda: (Tensor([4.0]) >> 1).realize(), expected=RuntimeError)
+    self.helper_test_exception([], lambda: tor >> torch.tensor([1.0]), lambda: (ten >> Tensor([1.0])).realize(), expected=RuntimeError)

  def test_lshift_signed(self):
    data = [[-1, -3, 1, 7], [0, -2147483648, 2147483647, -1]]
@@ -87,7 +87,8 @@ class TestOptim(unittest.TestCase):
  def test_muon(self): self._test_muon(1, {'lr': 0.001}, 1e-3, 0)
  # TODO: disabled due to big atol
  # def test_muon_high_lr(self): self._test_muon(1, {'lr': 10}, 1e-6, 3e-4)
-  def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 0.01}, 1e-3, 3e-4)
+  # NOTE: big weight_decay so a missing wd would be way over atol
+  def test_muon_wd(self): self._test_muon(1, {'lr': 0.001, 'weight_decay': 10}, 1e-3, 3e-4)
  # TODO: disabled due to big atol
  # def test_muon_high_lr_wd(self): self._test_muon(1, {'lr': 10, 'weight_decay': 0.01}, 1e-6, 5e-4)

@@ -80,7 +80,7 @@ class TestWGSLFailures(unittest.TestCase):
  def test_folded_packed_store(self):
    b = UOp.param(0, dtypes.char, (4,))
    idx = b.index(UOp.const(0).cast(dtypes.int))
-    store = UOp.store(idx, UOp.load(idx, dtype=dtypes.uint32) & UOp.const(0xffffff00).cast(dtypes.uint32))
+    store = UOp.store(idx, idx.cast(dtypes.uint32).load() & UOp.const(0xffffff00).cast(dtypes.uint32))
    src = Device[Device.DEFAULT].renderer.render(UOp.sink(store, arg=KernelInfo()).toposort())
    self.assertIn("atomicAnd(&data0_4[0],4294967040u);", src)
    self.assertNotIn("atomicAdd", src)
@@ -23,7 +23,7 @@ class TestGPUCrash(unittest.TestCase):
    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():
-      setattr(cls, rdna3_name, getattr(ins, cdna3_name if cls.is_cdna else rdna3_name))
+      setattr(cls, rdna3_name, staticmethod(getattr(ins, cdna3_name if cls.is_cdna else rdna3_name)))

  def setUp(self):
    # Verify device works before each test
@@ -23,6 +23,7 @@ from tinygrad.codegen.decomp.dtype import f2f
 VarVal = UOp | tuple[str, list[str], str]

 def _const(dt, v): return UOp.const(v, dt)
+def _single_value(v: UOp): return v.vmin if v.vmin == v.vmax else None
 def _u32(v): return _const(dtypes.uint32, v)
 def _u64(v): return _const(dtypes.uint64, v)
 def _to_u32(v): return v if v.dtype == dtypes.uint32 else v.bitcast(dtypes.uint32) if v.dtype.itemsize == 4 else v.cast(dtypes.uint32)
@@ -70,8 +71,8 @@ def _expr_bits(v: UOp) -> int:
  if v.op in (Ops.AND, Ops.XOR):
    widths: list[int] = []
    for src in v.src:
-      if src.op == Ops.CONST and isinstance(src.val, int) and src.val > 0 and (src.val & (src.val + 1)) == 0:
-        widths.append(src.val.bit_length())
+      if isinstance(sv:=_single_value(src), int) and sv > 0 and (sv & (sv + 1)) == 0:
+        widths.append(sv.bit_length())
    if widths: return max(widths)
  return v.dtype.bitsize

@@ -159,9 +160,9 @@ def _minmax_reduce(is_max: bool, dt, *args: UOp) -> UOp:
 def _find_two_pi_mul(x):
  if x.op != Ops.MUL or len(x.src) != 2: return None
  for i, s in enumerate(x.src):
-    if s.op == Ops.CONST and abs(s.val - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586)
+    if (sv:=_single_value(s)) is not None and abs(sv - 6.283185307179586) < 1e-5: return (x.src[1-i], 6.283185307179586)
    if s.op == Ops.MUL and len(s.src) == 2:
-      vals = [ss.val for ss in s.src if ss.op == Ops.CONST] + [ss.src[0].val for ss in s.src if ss.op == Ops.CAST and ss.src[0].op == Ops.CONST]
+      vals = [sv for ss in s.src if (sv:=_single_value(ss)) is not None]
      if len(vals) == 2 and abs(vals[0] * vals[1] - 6.283185307179586) < 1e-5: return (x.src[1-i], vals[0] * vals[1])
  return None

@@ -178,7 +179,7 @@ def _trig_reduce(x, phase=0.0):

 def _signext(val: UOp) -> UOp:
  for bits, mask, ext in [(4, 0xF, 0xFFFFFFF0), (8, 0xFF, 0xFFFFFF00), (16, 0xFFFF, 0xFFFF0000)]:
-    if (val.op == Ops.AND and len(val.src) == 2 and val.src[1].op == Ops.CONST and val.src[1].val == mask) or val.dtype.itemsize == bits // 8:
+    if (val.op == Ops.AND and len(val.src) == 2 and _single_value(val.src[1]) == mask) or val.dtype.itemsize == bits // 8:
      v32 = val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val
      sb = (v32 >> _u32(bits - 1)) & _u32(1)
      return sb.ne(_u32(0)).where(v32 | _u32(ext), v32).cast(dtypes.int)
@@ -549,7 +550,7 @@ class Parser:
        if not dtypes.is_int(right.dtype): right = right.cast(dtypes.uint32)
        return (left >> right) if op == '>>' else (left << right)
      case '+' | '-':
-        if op == '-' and left.op == Ops.CONST and right.op == Ops.CONST: return _const(left.dtype, left.val - right.val)
+        if op == '-' and (lv:=_single_value(left)) is not None and (rv:=_single_value(right)) is not None: return _const(left.dtype, lv - rv)
        return (left + right) if op == '+' else (left - right)
      case '*' | '/':
        # Integer promotion: promote 16-bit integers to 32-bit before multiply to avoid overflow
@@ -559,7 +560,7 @@ class Parser:
          left, right = left.cast(pdt), right.cast(pdt)
        if op == '*': return left * right
        return (left // right) if dtypes.is_int(left.dtype) else (left / right)
-      case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if left.op == Ops.CONST and left.val == 2.0 else left
+      case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if _single_value(left) == 2.0 else left

  _PREC = [('||',), ('&&',), ('|',), ('^',), ('&',), ('==', '!=', '<>'), ('>=', '<=', '>', '<'), ('>>', '<<'), ('+', '-'), ('*', '/'), ('**',)]

@@ -581,8 +582,8 @@ class Parser:
      return inner.eq(_const(inner.dtype, 0))
    if self.try_eat_val('-', 'OP'):
      inner = self.unary()
-      if inner.op == Ops.CONST:
-        return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -inner.val)
+      if (v:=_single_value(inner)) is not None:
+        return _const(dtypes.int if inner.dtype == dtypes.uint32 else inner.dtype, -v)
      return inner.neg()
    if self.try_eat_val('+', 'OP'): return self.unary()
    return self.postfix()
@@ -721,15 +722,13 @@ class Parser:
      self.eat('OP')
      width = self.parse()
      self.eat('RBRACKET')
-      if width.op == Ops.CONST:
-        w = int(width.val)
+      if isinstance(w:=_single_value(width), int):
        return (base >> _to_u32(first)) & _const(base.dtype, (1 << w) - 1)
      return base
    if self.try_eat('COLON'):
      second = self.parse()
      self.eat('RBRACKET')
-      if first.op == Ops.CONST and second.op == Ops.CONST:
-        a, b = int(first.val), int(second.val)
+      if isinstance(a:=_single_value(first), int) and isinstance(b:=_single_value(second), int):
        if a < b: return _bitreverse(base, b - a + 1)
        hi, lo = a, b
        if lo >= base.dtype.itemsize * 8:
@@ -750,8 +749,7 @@ class Parser:
      dt_suffix = DTYPES.get(self.eat('IDENT').val, dtypes.uint32)
    if var_name is None:
      var_name = self._find_var_name(base)
-    if first.op == Ops.CONST:
-      idx = int(first.val)
+    if isinstance(idx:=_single_value(first), int):
      # Check for array element (var@idx)
      if var_name and f'{var_name}@{idx}' in self.vars:
        v = self.vars[f'{var_name}@{idx}']
@@ -924,7 +922,7 @@ class Parser:

  def _coerce_cmp(self, l: UOp, r: UOp) -> tuple[UOp, UOp]:
    if l.dtype != r.dtype:
-      if r.dtype == dtypes.int and r.op == Ops.CONST and r.val < 0: l = l.cast(dtypes.int)
+      if r.dtype == dtypes.int and isinstance(rv:=_single_value(r), int) and rv < 0: l = l.cast(dtypes.int)
      else: r = r.cast(l.dtype)
    return l, r

@@ -419,7 +419,7 @@ class TestAutoCastType(unittest.TestCase):
    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)
+    with self.assertRaises(RuntimeError): Tensor([1, 0, 2]).where(1, 0).dtype
    self.check_where_alternate_input_other(False, True, dtypes.bool)

  @given(strat.sampled_from(core_dtypes), strat.sampled_from(core_dtypes))
@@ -429,7 +429,7 @@ class TestReduceCollapse(unittest.TestCase):

 class TestMovementOps(unittest.TestCase):
  def test_pm_mops_partial_reshape_index_removes_reshape(self):
-    from tinygrad.schedule.rangeify import pm_mops
+    from tinygrad.schedule.prepare import pm_mops
    src = UOp.param(0, dtypes.float, shape=(32, 4))
    r0, r1 = UOp.range(4, 0), UOp.range(8, 1)
    result = graph_rewrite(src.reshape((4, 8, 4)).index(r0, r1), pm_mops, name="test")
@@ -439,7 +439,7 @@ class TestMovementOps(unittest.TestCase):
    self.assertNotIn(Ops.RESHAPE, [u.op for u in result.toposort()])

  def test_pm_mops_partial_reshape_index_suffix_mismatch_does_nothing(self):
-    from tinygrad.schedule.rangeify import pm_mops
+    from tinygrad.schedule.prepare import pm_mops
    src = UOp.param(0, dtypes.float, shape=(2, 6))
    result = graph_rewrite(src.reshape((2, 3, 2)).index(UOp.range(2, 0)), pm_mops, name="test")
    self.assertEqual(result.op, Ops.INDEX)
@@ -1023,7 +1023,7 @@ class TestSymbolic(unittest.TestCase):
    cond = Variable("s", 0, 3, dtypes.int) < 2
    a = Variable("a", 0, 3, dtypes.int)
    self.assertIs(graph_rewrite(cond.where(a, a+1).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), (a+1).cast(dtypes.half)))
-    self.assertIs(graph_rewrite(cond.where(a, uconst(2)).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), UOp.const(2, dtypes.half)))
+    self.assertIs(graph_rewrite(cond.where(a, uconst(2)).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), uconst(2.0)))
    self.assertIs(graph_rewrite(cond.where(a, UOp.invalid()).cast(dtypes.half), sym), cond.where(a.cast(dtypes.half), UOp.invalid()))

  def test_where_const_gate_keeps_stated_width(self):
@@ -82,7 +82,7 @@ class TestVminVmaxProperties(unittest.TestCase):
  def test_vmin_vmax_multiplication_0_inf(self):
    # vmin and vmax for multiplication with a variable
    x = UOp.const(0.0)
-    y = UOp.load(UOp.param(0, dtypes.float, (1,)), UOp.const(0), dtype=dtypes.float)
+    y = UOp.load(UOp.param(0, dtypes.float, (1,)), UOp.const(0))
    uop = x * y
    # TODO: these should be 0, but definitely should not be nan
    self.assertEqual(uop.vmin, -math.inf)
@@ -6,7 +6,7 @@ from tinygrad.helpers import Timing, Context, cdiv
 from tinygrad.dtype import dtypes, AddrSpace, ConstFloat, Invalid  # noqa: F401
 from tinygrad.device import Device
 from tinygrad.uop.ops import Ops, AxisType, ParamArg, PatternMatcher, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite  # noqa: F401  # ParamArg used by eval(str(uop)) roundtrip tests
-from tinygrad.uop.weak import pm_lower_index_dtype
+from tinygrad.uop.weak import pm_lower_weak
 from tinygrad.uop.spec import spec_program, spec_shared, type_verify
 from tinygrad.uop.symbolic import sym, pm_remove_invalid
 from test.helpers import eval_uop, to_uops_list
@@ -56,7 +56,7 @@ class TestDTypeFromUOp(unittest.TestCase):
      if u.is_invalid)), (dtypes.float32, dtypes.float32, dtypes.bool))
    invalid, value = UOp.invalid(), UOp.const(1, dtypes.float32)
    for u in (UOp.param(0, dtypes.bool, ()).where(value, invalid), value+invalid, UOp.stack(value, invalid)): self.assertIs(u.src[-1], invalid)
-    for u in (UOp(Ops.STACK, dtypes.float32, src=(value, invalid)), UOp(Ops.ADD, dtypes.float32, src=(value, invalid)),
+    for u in (UOp(Ops.STACK, src=(value, invalid)), UOp(Ops.ADD, src=(value, invalid)),
              UOp.const(True).where(value, invalid), UOp(Ops.CMPLT, src=(invalid, value)), UOp(Ops.CMPLT, src=(value, invalid)),
              UOp.param(0, dtypes.float32, (4,)).index(invalid)): type_verify(u, spec_shared)
    gate, value = UOp.param(0, dtypes.bool, ()), UOp.param(1, dtypes.float, ())
@@ -64,7 +64,7 @@ class TestDTypeFromUOp(unittest.TestCase):
    type_verify(out.sink(), spec_program)

  def test_remove_invalid_stack_lanes(self):
-    stack = UOp(Ops.STACK, dtypes.half, (UOp.const(1, dtypes.half), UOp.invalid()))
+    stack = UOp(Ops.STACK, src=(UOp.const(1, dtypes.half), UOp.invalid()))
    out = graph_rewrite(stack, pm_remove_invalid)
    self.assertEqual(out.src, (UOp.const(1, dtypes.half), UOp.const(0, dtypes.half)))
    type_verify(out.sink(), spec_program)
@@ -76,16 +76,18 @@ class TestLowerIndexDtype(unittest.TestCase):
    buf = UOp.param(0, dtypes.float, (2**31+64,))
    i = UOp.variable("i", 0, 2**28)
    shrink = UOp(Ops.SHRINK, src=(buf, (i*24).valid(i < 2**28), UOp.const(4)))
-    lowered = graph_rewrite(shrink.sink(), pm_lower_index_dtype)
-    self.assertTrue(all(u.dtype != dtypes.weakint for u in lowered.backward_slice_with_self), "lowering must resolve all weakint")
+    lowered = graph_rewrite(shrink.sink(), pm_lower_weak)
+    self.assertTrue(all(u.op is Ops.CONST for u in lowered.backward_slice_with_self if u.dtype in dtypes.weaks),
+                    "lowering must resolve every weak width, except a typed literal's value half")
    sh = next(u for u in lowered.backward_slice_with_self if u.op is Ops.SHRINK)
    self.assertEqual(sh.src[1].dtype, dtypes.long)

  def test_reg_buffer_size_lowers(self):
    reg = UOp.placeholder((4,), dtypes.float, 0, addrspace=AddrSpace.REG)
    self.assertEqual(reg.src[0].dtype, dtypes.weakint)
-    lowered = graph_rewrite(reg.sink(), pm_lower_index_dtype)
-    self.assertTrue(all(u.dtype != dtypes.weakint for u in lowered.backward_slice_with_self), "lowering must resolve all weakint")
+    lowered = graph_rewrite(reg.sink(), pm_lower_weak)
+    self.assertTrue(all(u.op is Ops.CONST for u in lowered.backward_slice_with_self if u.dtype in dtypes.weaks),
+                    "lowering must resolve every weak width, except a typed literal's value half")
    self.assertEqual(next(u for u in lowered.backward_slice_with_self if u.op is Ops.BUFFER).src[0].dtype, dtypes.int)

 class TestSafeCast(unittest.TestCase):
@@ -282,7 +284,7 @@ class TestFastIdiv(unittest.TestCase):
      g = UOp.param(0, dt, (3,))
      c = UOp.const(2)
      l = g.index(c)
-      a = UOp(Ops.CDIV, dt, (l, c))
+      a = UOp(Ops.CDIV, src=(l, c))
      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
      Device[Device.DEFAULT].renderer.render(uops)
      ops = [x.op for x in uops]
@@ -294,7 +296,7 @@ class TestFastIdiv(unittest.TestCase):
    for dt in (dtypes.int32, dtypes.uint32):
      g = UOp.param(0, dt, (9,))
      c = UOp.const(8)
-      a = UOp(Ops.FLOORMOD, dt, (g.index(c), c))
+      a = UOp(Ops.FLOORMOD, src=(g.index(c), c))
      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
      ops = [x.op for x in uops]
      self.assertIn(Ops.AND, ops, f"For dtype={dt} FLOORMOD by pow2 did not simplify to AND")
@@ -306,7 +308,7 @@ class TestFastIdiv(unittest.TestCase):
    for dt in (dtypes.int32, dtypes.uint32, dtypes.int64, dtypes.uint64):
      g = UOp.param(0, dt, (3,))
      c = UOp.const(2)
-      a = UOp(Ops.FLOORDIV, dt, (g.index(c), c))
+      a = UOp(Ops.FLOORDIV, src=(g.index(c), c))
      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
      ops = [x.op for x in uops]
      self.assertIn(Ops.SHR, ops, f"For dtype={dt} FLOORDIV by power of two did not simplify to shift")
@@ -467,10 +469,10 @@ class TestUOpRender(unittest.TestCase):
    self.assertEqual(UOp.range(1, 0, src=(shrink,), dtype=dtypes.int).render(simplify=False), "r0")

  def test_render_vectorize_empty(self):
-    u = UOp(Ops.STACK, dtype=dtypes.void, src=())
+    u = UOp(Ops.STACK, src=())
    self.assertEqual(u.render(simplify=False), "{}")
  def test_render_vectorize_empty_simplified(self):
-    u = UOp(Ops.STACK, dtype=dtypes.void, src=())
+    u = UOp(Ops.STACK, src=())
    self.assertEqual(u.render(), "{}")
  def test_render_vectorize_same(self):
    u = UOp(Ops.STACK, src=(UOp.const(0),)*3)
@@ -14,37 +14,37 @@ class TestValidateOOB(unittest.TestCase):
  def test_const_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (16,))
-      to_uops_list([buf.index(UOp.const(0)).load(dtype=dtypes.int)])  # valid
-      to_uops_list([buf.index(UOp.const(15)).load(dtype=dtypes.int)])  # valid (last element)
+      to_uops_list([buf.index(UOp.const(0)).load()])  # valid
+      to_uops_list([buf.index(UOp.const(15)).load()])  # valid (last element)
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(UOp.const(16)).load(dtype=dtypes.int)])  # off by one
+        to_uops_list([buf.index(UOp.const(16)).load()])  # off by one
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(UOp.const(42)).load(dtype=dtypes.int)])  # way out
+        to_uops_list([buf.index(UOp.const(42)).load()])  # way out

  def test_variable_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (16,))
-      to_uops_list([buf.index(Variable("i", 0, 15)).load(dtype=dtypes.int)])  # valid
+      to_uops_list([buf.index(Variable("i", 0, 15)).load()])  # valid
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(Variable("i", 0, 20)).load(dtype=dtypes.int)])  # oob
+        to_uops_list([buf.index(Variable("i", 0, 20)).load()])  # oob
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(Variable("i", -5, 10)).load(dtype=dtypes.int)])  # negative
+        to_uops_list([buf.index(Variable("i", -5, 10)).load()])  # negative

  def test_range_with_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (16,))
      r = UOp.range(42, 0, AxisType.GLOBAL)
-      to_uops_list([buf.index(r.valid(r < 16)).load(dtype=dtypes.int)])  # valid
+      to_uops_list([buf.index(r.valid(r < 16)).load()])  # valid
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(r.valid(r < 17)).load(dtype=dtypes.int)])  # oob
+        to_uops_list([buf.index(r.valid(r < 17)).load()])  # oob

  def test_variable_with_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (16,))
      v = Variable("v", -5, 80)
-      to_uops_list([buf.index(v.valid((v >= 0) & (v < 16))).load(dtype=dtypes.int)])  # valid
+      to_uops_list([buf.index(v.valid((v >= 0) & (v < 16))).load()])  # valid
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(v.valid(v < 20)).load(dtype=dtypes.int)])  # negative not masked
+        to_uops_list([buf.index(v.valid(v < 20)).load()])  # negative not masked

  def test_gated_store(self):
    with Context(CHECK_OOB=1, SPEC=2):
@@ -58,62 +58,62 @@ class TestValidateOOB(unittest.TestCase):
  def test_floordiv(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (16,))
-      to_uops_list([buf.index(UOp.range(32, 0, AxisType.GLOBAL) // 2).load(dtype=dtypes.int)])  # 0..15 valid
+      to_uops_list([buf.index(UOp.range(32, 0, AxisType.GLOBAL) // 2).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(UOp.range(34, 0, AxisType.GLOBAL) // 2).load(dtype=dtypes.int)])  # 0..16 oob
+        to_uops_list([buf.index(UOp.range(34, 0, AxisType.GLOBAL) // 2).load()])  # 0..16 oob

  def test_mod(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (16,))
      r = UOp.range(100, 0, AxisType.GLOBAL)
-      to_uops_list([buf.index(r % 16).load(dtype=dtypes.int)])  # 0..15 valid
+      to_uops_list([buf.index(r % 16).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(r % 20).load(dtype=dtypes.int)])  # 0..19 oob
+        to_uops_list([buf.index(r % 20).load()])  # 0..19 oob

  def test_shr(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (16,))
-      to_uops_list([buf.index(UOp.range(64, 0, AxisType.GLOBAL) >> 2).load(dtype=dtypes.int)])  # 0..15 valid
+      to_uops_list([buf.index(UOp.range(64, 0, AxisType.GLOBAL) >> 2).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(UOp.range(128, 0, AxisType.GLOBAL) >> 2).load(dtype=dtypes.int)])  # 0..31 oob
+        to_uops_list([buf.index(UOp.range(128, 0, AxisType.GLOBAL) >> 2).load()])  # 0..31 oob

  def test_shl(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (64,))
      r = UOp.range(8, 0, AxisType.GLOBAL)
-      to_uops_list([buf.index(r << 2).load(dtype=dtypes.int)])  # 0..28 valid
+      to_uops_list([buf.index(r << 2).load()])  # 0..28 valid
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(r << 4).load(dtype=dtypes.int)])  # 0..112 oob
+        to_uops_list([buf.index(r << 4).load()])  # 0..112 oob

  def test_and(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (16,))
      r = UOp.range(100, 0, AxisType.GLOBAL)
-      to_uops_list([buf.index(r & 15).load(dtype=dtypes.int)])  # 0..15 valid
+      to_uops_list([buf.index(r & 15).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(r & 31).load(dtype=dtypes.int)])  # 0..31 oob
+        to_uops_list([buf.index(r & 31).load()])  # 0..31 oob
      # align masks round down to a multiple of 2^k
-      to_uops_list([buf.index((r & -4).valid(r < 16)).load(dtype=dtypes.int)])  # 0..12 valid
+      to_uops_list([buf.index((r & -4).valid(r < 16)).load()])  # 0..12 valid
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(r & -2).load(dtype=dtypes.int)])  # 0..100 oob
+        to_uops_list([buf.index(r & -2).load()])  # 0..100 oob
      # other masks can't be modeled as mod
      with self.assertRaisesRegex(RuntimeError, "z3 int AND only supports"):
-        to_uops_list([buf.index(r & 21).load(dtype=dtypes.int)])
+        to_uops_list([buf.index(r & 21).load()])

  def test_max(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (16,))
-      to_uops_list([buf.index(Variable("v", -10, 15).maximum(0)).load(dtype=dtypes.int)])  # 0..15 valid
+      to_uops_list([buf.index(Variable("v", -10, 15).maximum(0)).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(Variable("v2", -10, 20).maximum(0)).load(dtype=dtypes.int)])  # 0..20 oob
+        to_uops_list([buf.index(Variable("v2", -10, 20).maximum(0)).load()])  # 0..20 oob

  def test_xor_in_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (16,))
      r = UOp.range(32, 0, AxisType.GLOBAL)
-      to_uops_list([buf.index(r.valid((r < 8) ^ ((r >= 8) & (r < 16)))).load(dtype=dtypes.int)])  # 0..15 valid
+      to_uops_list([buf.index(r.valid((r < 8) ^ ((r >= 8) & (r < 16)))).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf.index(r.valid((r < 10) ^ (r >= 20))).load(dtype=dtypes.int)])  # 0..9,20..31 oob
+        to_uops_list([buf.index(r.valid((r < 10) ^ (r >= 20))).load()])  # 0..9,20..31 oob

  # cast patterns
  def test_float_cast_in_index(self):
@@ -121,13 +121,13 @@ class TestValidateOOB(unittest.TestCase):
      buf = UOp.param(0, dtypes.int, (16,))
      r = UOp.range(20, 0)
      i = (r.cast(dtypes.float) * 0.68).trunc().cast(dtypes.int)
-      to_uops_list([buf.index(i.valid((i >= 0) & (i < 16))).load(dtype=dtypes.int)])
+      to_uops_list([buf.index(i.valid((i >= 0) & (i < 16))).load()])

  def test_bool_cast_in_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf = UOp.param(0, dtypes.int, (1,))
      r = UOp.range(20, 0)
-      to_uops_list([buf.index(r.valid(r.cast(dtypes.bool).logical_not())).load(dtype=dtypes.int)])  # only r=0 valid
+      to_uops_list([buf.index(r.valid(r.cast(dtypes.bool).logical_not())).load()])  # only r=0 valid

  # load result as index/mask
  def test_load_as_index(self):
@@ -135,18 +135,18 @@ class TestValidateOOB(unittest.TestCase):
      buf0 = UOp.param(0, dtypes.int, (16,))
      buf1 = UOp.param(1, dtypes.int, (64,))
      r = UOp.range(42, 0, AxisType.GLOBAL)
-      ld0 = buf0.index(r.valid(r < 8)).load(dtype=dtypes.int).cast(dtypes.weakint)
-      to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32))).load(dtype=dtypes.int)])  # valid
+      ld0 = buf0.index(r.valid(r < 8)).load().cast(dtypes.weakint)
+      to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32))).load()])  # valid
      with self.assertRaises(RuntimeError):
-        to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 64))).load(dtype=dtypes.int)])  # oob
+        to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 64))).load()])  # oob

  def test_load_from_shrink_as_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
      buf0 = UOp.param(0, dtypes.int, (16,))
      buf1 = UOp.param(1, dtypes.int, (64,))
      shrink = UOp(Ops.SHRINK, src=(buf0, UOp.const(0, dtypes.int), UOp.const(4)))
-      ld0 = shrink.load(dtype=dtypes.int).index(0)
-      to_uops_list([buf1.index(ld0.valid((ld0 >= 0) & (ld0 < 64))).load(dtype=dtypes.int)])
+      ld0 = shrink.load().index(0)
+      to_uops_list([buf1.index(ld0.valid((ld0 >= 0) & (ld0 < 64))).load()])

  def test_load_bool_as_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
@@ -16,7 +16,7 @@ def apply_rope(x:Tensor, start_pos:int):

 class TestLinear(unittest.TestCase):
  def test_recovers_packed_ggml_weight(self):
-    for ggml_type,packed_size,words in ((13, 176, 44), (14, 210, 53), (23, 136, 34)):
+    for ggml_type,packed_size,words in ((13, 176, 44), (14, 210, 210), (23, 136, 34)):
      packed = Tensor.empty(packed_size+4, dtype=dtypes.uint8, device="CPU")[4:]
      decoded = ggml_data_to_tensor(packed, 256, ggml_type).reshape(1, 256)
      linear = Linear(256, 1, bias=False)
@@ -4,7 +4,7 @@ from tinygrad import Tensor, dtypes, TinyJit
 from tinygrad.helpers import Context
 from tinygrad.dtype import least_upper_float
 from tinygrad.uop.ops import UOp, Ops, GroupOp, dtype_from_uop, graph_rewrite
-from tinygrad.uop.weak import pm_lower_index_dtype, pm_commit_weak
+from tinygrad.uop.weak import pm_commit_weak
 from tinygrad.uop.symbolic import symbolic_simple
 from tinygrad.uop.spec import spec_shared, type_verify
 from tinygrad.engine.jit import JitError
@@ -74,7 +74,7 @@ class TestWeakPromotion(unittest.TestCase):
    recips = [u for u in (x / y)._uop.toposort() if u.op is Ops.RECIPROCAL]
    self.assertEqual([(u.dtype, u.src[0].dtype) for u in recips], [(dtypes.float32, dtypes.float32)])
    with Context(DEFAULT_FLOAT=dtypes.float16):
-      committed = graph_rewrite((UOp.const(1).cast(dtypes.int32) + UOp.const(1.0)).cast(dtypes.float32), pm_lower_index_dtype, ctx={})
+      committed = graph_rewrite((UOp.const(1).cast(dtypes.int32) + UOp.const(1.0)).cast(dtypes.float32), pm_commit_weak)
    self.assertEqual([u.dtype for u in committed.toposort() if u.op is Ops.ADD], [dtypes.float32])

  def test_div_sub_operand_kept_weak(self):
@@ -85,7 +85,7 @@ class TestWeakPromotion(unittest.TestCase):
  def test_cast_weak_expression_commits_at_cast_floor(self):
    # the floor never narrows: a cast BELOW the default does not pull the compute width down with it
    with Context(DEFAULT_FLOAT=dtypes.float32):
-      narrowed = graph_rewrite((UOp.const(1.0) + UOp.const(2.0)).cast(dtypes.float16), pm_lower_index_dtype, ctx={})
+      narrowed = graph_rewrite((UOp.const(1.0) + UOp.const(2.0)).cast(dtypes.float16), pm_commit_weak)
    self.assertEqual((narrowed.dtype, narrowed.src[0].dtype), (dtypes.float16, dtypes.float32))

  def test_cast_weak_expression_value_uses_cast_floor(self):
@@ -126,8 +126,23 @@ class TestWeakPromotion(unittest.TestCase):
    weak_lub = UOp(Ops.ADD, src=(UOp.const(1), UOp.const(1.0)))
    self.assertIs(graph_rewrite(weak_lub, pm_commit_weak), weak_lub)
    concrete = UOp.const(2.0).cast(dtypes.float16)
-    where = graph_rewrite(UOp(Ops.WHERE, src=(UOp.const(True), concrete, UOp.const(1.0))), pm_lower_index_dtype, ctx={})
-    self.assertEqual(tuple(x.dtype for x in where.src), (dtypes.bool, dtypes.float16, dtypes.float16))
+    # the weak arm stays bare: its sibling states the width, so the WHERE already derives float16 for it
+    where = graph_rewrite(UOp(Ops.WHERE, src=(UOp.const(True), concrete, UOp.const(1.0))), pm_commit_weak)
+    self.assertEqual((where.dtype, tuple(x.dtype for x in where.src)), (dtypes.float16, (dtypes.bool, dtypes.float16, dtypes.weakfloat)))
+
+  def test_derivable_const_rounds_at_the_derived_width(self):
+    # re-rounds a derivable const in place (still bare) so value-keyed folds (x*1 -> x, x*-1 -> NEG) still fire
+    x = UOp.param(0, dtypes.float32, (1,)).index(UOp.const(0).cast(dtypes.int32)).load()
+    mul = graph_rewrite(x * UOp.const(-0.9999999893980771), symbolic_simple+pm_commit_weak)
+    self.assertIs(mul.src[1], UOp.const(-1.0))
+    self.assertIs(graph_rewrite(x * UOp.const(1.0000000106), symbolic_simple+pm_commit_weak), x)
+
+  def test_committed_const_conversion_folds_for_native_format(self):
+    folded = graph_rewrite(UOp.const(16256, dtypes.ushort).cast(dtypes.uint), symbolic_simple)
+    self.assertIs(folded, UOp.const(16256, dtypes.uint))
+    # fmt-less targets are lowered by renderer rewrites, where collapsing this pair would cycle with float-intermediate insertion.
+    emulated = UOp.const(1.0, dtypes.float).cast(dtypes.bfloat16)
+    self.assertIs(graph_rewrite(emulated, symbolic_simple), emulated)

  def test_weak_shift_lhs_commits_the_node(self):
    # a shift derives its lhs's dtype, so committing the lhs restates the root (WGSL's packed store writes `mask << shift_am`)
@@ -168,11 +183,11 @@ class TestWeakPromotion(unittest.TestCase):
    self.assertEqual(dtype_from_uop(Ops.SHL, (UOp.const(1, dtypes.int8), UOp.const(1, dtypes.uint32)), None), dtypes.int8)
    self.assertEqual(UOp.const(1).alu(Ops.SHL, UOp.const(1, dtypes.uint)).dtype, dtypes.weakint)
    self.assertEqual((v & 3).dtype, dtypes.weakint)
-    with self.assertRaises(RuntimeError): Tensor.const(1.0) << Tensor.const(1.0)
-    with self.assertRaises(RuntimeError): UOp.const(1, dtypes.int32).alu(Ops.SHL, UOp.const(1, dtypes.float64))
+    with self.assertRaises(RuntimeError): (Tensor.const(1.0) << Tensor.const(1.0)).dtype
+    with self.assertRaises(RuntimeError): UOp.const(1, dtypes.int32).alu(Ops.SHL, UOp.const(1, dtypes.float64)).dtype
    for op in (Ops.SHL, Ops.SHR):
      with self.assertRaises(RuntimeError):
-        UOp.const(1, dtypes.float32).alu(op, UOp.const(1, dtypes.int32))
+        UOp.const(1, dtypes.float32).alu(op, UOp.const(1, dtypes.int32)).dtype
    # float bitwise builds, the spec rejects it
    with Context(SPEC=1):
      f32, wf = UOp.const(1.0, dtypes.float32), UOp.const(1.0)
@@ -20,13 +20,11 @@ class TestQ8Quantize(unittest.TestCase):
  def test_values_and_scales(self):
    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
    x = np.linspace(-3.1, 2.7, 64, dtype=np.float32).reshape(2, 32)
-    quant, scale, gsum = q8_quantize(Tensor(x), 2, 32)
+    quant, scale = q8_quantize(Tensor(x), 2, 32)
    scale_np = np.maximum(np.max(np.abs(x), axis=-1, keepdims=True) / 127, 1e-8)
    expected = np.clip(np.rint(x / scale_np), -127, 127).astype(np.int8)
    np.testing.assert_array_equal(quant.bitcast(dtypes.int8).reshape(2, 32).numpy(), expected)
    np.testing.assert_allclose(scale.numpy(), scale_np, rtol=1e-6)
-    # xsum holds the two per-16 sums per 32-wide group
-    np.testing.assert_array_equal(gsum.numpy().reshape(2, 2), expected.reshape(2, 2, 16).sum(-1).astype(np.float32))

  def test_q6_linear_compiles(self):
    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
@@ -38,9 +36,7 @@ class TestQ8Quantize(unittest.TestCase):
    linear = Linear(256, 1, bias=False)
    nn.state.load_state_dict(linear, {"weight":decoded}, verbose=False, realize=False)
    self.assertTrue(np.isfinite(linear(Tensor.randn(1, 256)).realize().item()))
-    # the Q6 weight is repacked: 210-byte blocks padded to 212 (one block = 53 words)
-    self.assertEqual(linear.weight.uop.buf_uop.buffer.nbytes, 53*4)
-    self.assertEqual(linear.weight.dtype, dtypes.uint32)
+    self.assertEqual(linear.weight.uop.buf_uop.buffer.offset, 4)

  def test_q4_k_linear(self):
    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
@@ -3,7 +3,7 @@ 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, DEFAULT_FLOAT, DEFAULT_INT, NUM_CPU_THREADS, TC_SELECT, TC_OPT, TracingKey, Context, panic
 from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, Ops, UPat, rewrite_group, KernelInfo, ProgramInfo, GroupOp, AxisType
-from tinygrad.uop.weak import pm_lower_index_dtype, pm_commit_weak, pm_cast_weak
+from tinygrad.uop.weak import pm_lower_weak, pm_commit_weak, pm_cast_const
 from tinygrad.uop.render import pyrender
 from tinygrad.uop.spec import type_verify, spec_tensor, spec_program
 from tinygrad.renderer import Renderer, Estimates
@@ -22,7 +22,7 @@ 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, pm_reduce_unparented
 from tinygrad.schedule.multi import multi_pm
-from tinygrad.schedule.rangeify import pm_mops
+from tinygrad.schedule.prepare 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
 from tinygrad.codegen.late.coalesce import memory_coalescing, pm_simplify_add_image
@@ -282,10 +282,6 @@ pm_implicit_barriers = PatternMatcher([
  (UPat(Ops.END, name="end"), add_war_barrier),
 ])

-pm_casted_consts = PatternMatcher([
-  (UPat(Ops.CONST, dtypes.all, name="c"), lambda c: UOp.cconst(c.val, c.dtype)),
-])
-
 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))
@@ -347,11 +343,13 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:

  # extra symbolic before decomp. crashes without this?
  # NOTE: also run indexing_simplify here, while the index is still weakint and (x+y)*c -> x*c+y*c applies
-  sink = graph_rewrite(sink, sym+indexing_simplify, name="extra symbolic")
+  # commit widths minted in this fixpoint before lowering inspects INDEX shapes
+  sink = graph_rewrite(sink, sym+indexing_simplify+pm_commit_weak, name="extra symbolic")

-  # lower index dtype
+  # the boundary: required compute dtypes settle here; derivable const edges may stay bare
  # NOTE: we need indexing_simplify to remove the cast to long using the Invalid
-  sink = graph_rewrite(sink, symbolic_simple+pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes")
+  # NOTE: symbolic must NOT be composed here -- pm_data_invalid pushes the weak result CAST into a gated WHERE, remaking the weak node, and it cycles
+  sink = graph_rewrite(sink, pm_lower_weak+indexing_simplify, name="lower all index dtypes")

  # final symbolic before decomp
  sink = graph_rewrite(sink, symbolic, name="final symbolic")
@@ -375,12 +373,11 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:

  # final rules for the renderer (without sym)
  extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([])
-  pm_final_rewrite = pm_commit_weak+pm_cast_weak+pm_decomp+extra_matcher+pm_split_ends
+  pm_final_rewrite = pm_commit_weak+pm_decomp+extra_matcher+pm_split_ends
  sink = graph_rewrite(sink, pm_final_rewrite+pm_remove_invalid, ctx=ren, name="final rewrite")

-  # spell every literal as a casted const CAST(dt, CONST(value))
-  # TODO: remove once consts are always weak
-  sink = graph_rewrite(sink, pm_casted_consts, name="casted consts", walk=True)
+  # commit every const still bare so no renderer reads one
+  sink = graph_rewrite(sink, pm_cast_const, name="cast consts")

  # 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")
@@ -3,6 +3,7 @@ from tinygrad.dtype import dtypes, DType, truncate
 from tinygrad.helpers import flatten, DEBUG, EMULATED_DTYPES
 from tinygrad.uop import GroupOp
 from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite
+from tinygrad.uop.weak import commit_weak_consts
 from tinygrad.renderer import Renderer
 from tinygrad.codegen.decomp.transcendental import exponent_bias, shl, shr

@@ -137,6 +138,8 @@ def f2f_store(st, idx, val, fr:DType, to:DType):

 # 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 = PatternMatcher([
+  # the decomp's own bottom-up rewrite can mint bare consts mid-flight: word splitting commits them at the long sibling's dtype
+  (UPat(GroupOp.All, name='x'), lambda x: commit_weak_consts(x, next((s.dtype for s in x.src if s.dtype in l2i_dt), None))),
  (UPat(GroupOp.Defines, tuple(l2i_dt.keys()), src=(UPat.var("sz"),), name="x"), lambda x,sz:
   UOp(x.op, src=(sz*2,), arg=replace(x.arg, dtype=l2i_dt[x.dtype]), tag=x.tag)),
  (UPat(Ops.INDEX, tuple(l2i_dt.keys()), name='x'), lambda x:
@@ -148,6 +151,9 @@ pm_long_decomp: PatternMatcher = PatternMatcher([
   split_l2i(ctx, 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 ctx,a,x:
   split_l2i(ctx, Ops.BITCAST, l2i_dt[x.dtype], a.rtag((0, dt:=l2i_dt[a.dtype])), a.rtag((1, dt)))[x.tag[0]]),
+  # a const splits by value; the general CAST arm below would drop its high word
+  (UPat(Ops.CAST, src=(UPat(Ops.CONST, name='c'),), tag={(w, dt) for w in (0, 1) for dt in l2i_dt.values()}, name='x'),
+   lambda x,c: UOp.const(truncate[x.tag[1]](c.val >> (32*x.tag[0])), x.tag[1])),
  (UPat(Ops.CAST, tuple(l2i_dt.keys()), src=(UPat.var('a'),), name="x"), lambda ctx,a,x:
   split_l2i(ctx, 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 ctx,a,x:
@@ -161,9 +167,7 @@ pm_long_decomp: PatternMatcher = PatternMatcher([
   split_l2i(ctx, 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 ctx,x,idx:
-   reindex(graph_rewrite(idx, pm_long_decomp, ctx=ctx, bottom_up=True), x.tag[0]).replace(tag=None).load() 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(truncate[x.tag[1]]((x.val >> 32) if x.tag[0] == 1 else (x.val & 0xFFFFFFFF)), x.tag[1]))
+   reindex(graph_rewrite(idx, pm_long_decomp, ctx=ctx, bottom_up=True), x.tag[0]).replace(tag=None).load() if x.tag is not None else None)
 ])

 # float decomposition patterns - ctx is (fr, to) tuple
@@ -186,8 +190,6 @@ pm_float_decomp: PatternMatcher = PatternMatcher([
   f2f(x.bitcast(f2f_dt[ctx[0]]), ctx[0], ctx[1]) if bc.dtype == ctx[0] else None),
  (UPat(Ops.CAST, dtypes.floats, src=(UPat.var("val"),), name="x"), lambda ctx,x,val:
   f2f_clamp(val.cast(ctx[1]), ctx[0]) if x.dtype == ctx[0] else None),
-  # a CONST has no srcs to cast, it restates its value at the emulating dtype
-  (UPat(Ops.CONST, dtypes.floats, name="x"), lambda ctx,x: UOp.const(x.val, ctx[1]) if x.dtype == ctx[0] else None),
  (UPat(GroupOp.All-GroupOp.Defines-{Ops.CAST, Ops.BITCAST, Ops.CONST}, dtypes.floats, name="x"), lambda ctx,x:
   UOp(x.op, src=tuple(s.cast(ctx[1]) if s.dtype == ctx[0] else s for s in x.src), arg=x.arg, tag=x.tag) if x.dtype == ctx[0] else None),
  (UPat(Ops.STORE, src=(UPat.var("idx"), UPat(Ops.BITCAST, dtypes.floats, name="val")), name='st'), lambda ctx,st,idx,val:
@@ -90,8 +90,8 @@ def payne_hanek_reduction(d:UOp) -> tuple[UOp, UOp]:
    if count+offset < len(two_over_pi_f) - 1:
      an = i.ne(count).where(_take(an, offset, count=count+1), an.const_like(two_over_pi_f[count+offset]))
    return an
-  def _shl_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) * pow2if(y, d.dtype).cast(dtypes.uint64)).cast(dtypes.uint32)
-  def _shr_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) // pow2if(y, d.dtype).cast(dtypes.uint64)).cast(dtypes.uint32)
+  def _shl_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) << y.cast(dtypes.uint64)).cast(dtypes.uint32)
+  def _shr_lazy(x:UOp, y:UOp): return (x.cast(dtypes.uint64) >> y.cast(dtypes.uint64)).cast(dtypes.uint32)

  a = [_take(UOp.const(0, dtypes.uint32), i) for i in range(4)]
  #  (two_over_pi_f[Int(i) + n] << e) | (two_over_pi_f[Int(i) + n+1] >> (nbits - e))
@@ -193,7 +193,7 @@ class Scheduler:
      for b in self.bufs:
        if rng in (i:=b.src[1].get_idx()).backward_slice_with_self:
          nb = b.replace(src=(b.src[0], i.valid(valid&b.src[1].get_valid())))
-          replaces[b] = nb if b in store_targets else valid.where(nb, UOp.const(Invalid, b.dtype))
+          replaces[b] = nb if b in store_targets else valid.where(nb, UOp.const(Invalid))
      self.ast = self.ast.substitute(replaces, f"padto {rng.arg[:-1]} {opt.arg}")
    elif opt.op is OptOps.SWAP:
      try:
@@ -153,7 +153,6 @@ pm_validate_wmma_rdna3 = PatternMatcher([

 pm_validate_wmma_rdna4 = PatternMatcher([
  (UPat(Ops.WMMA, name="x", dtype=dtypes.bfloat16), lambda x: x.replace(
-    dtype=dtypes.uint16,
    src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2].bitcast(dtypes.uint16)))
      .bitcast(dtypes.bfloat16) if x.max_numel() == 8 and x.src[0].dtype == dtypes.bfloat16 and x.src[0].max_numel() == 8 else None),
  (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
@@ -26,11 +26,12 @@ def get_call_outs_ins(call:UOp) -> tuple[tuple[int, ...], tuple[int, ...]]:
  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec": return (0,), tuple(range(1, len(get_call_arg_uops(call))))
  return (), ()

-def get_call_kernels(call:UOp) -> list[tuple[str, UOp]]:
-  if (ast:=call.src[0]).op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq": return [(d, k) for devs, k, _ in call.arg.aux.kernels for d in devs]
-  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return [(to_tuple(ast.device)[0], call)]
+def get_call_kernels(call:UOp) -> list[tuple[str, UOp, tuple[str, Estimates, bytes]|None]]:
+  if (ast:=call.src[0]).op is Ops.CUSTOM_FUNCTION and ast.arg == "hcq":
+    return [(d, call, (name, estimates, profile_key)) for devices,name,estimates,_,profile_key in call.arg.aux.kernels for d in devices]
+  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph": return [(to_tuple(ast.device)[0], call, None)]
  if ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "validate": return []
-  return [(d, call) for d in to_tuple(call.src[1].device)]
+  return [(d, call, None) for d in to_tuple(call.src[1].device)]

 def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|None=None) -> str:
  def _uop_sz_to_str(uop:UOp) -> str: return size_to_str(sym_infer(prod(uop.shape) * uop.dtype.itemsize, var_vals or {}))
@@ -66,24 +67,25 @@ def track_stats(ctx:ExecContext, call:UOp, st:decimal.Decimal, ets:list[float|No
  if DEBUG < 2 and not PROFILE: return

  kernels = get_call_kernels(call) # everything below is the per kernel display: exec events for the profiler and DEBUG=2 lines
-  args = resolve_params(call, ctx.input_uops) if kernels and kernels[0][1] is call else []
+  args = resolve_params(call, ctx.input_uops) if kernels and kernels[0][2] is None else []
  lanes = list(unwrap_multi(call, [args[g] for g in call.src[0].arg.globals] if call.src[0].op is Ops.PROGRAM else args)) if args else []
-  for i, (device, kcall) in enumerate(kernels):
+  for i, (device, kcall, stats) in enumerate(kernels):
    et, bufs = ets[i] if i < len(ets) else None, lanes[i][0] if i < len(lanes) else []
+    display_name = get_call_name(kcall, bufs, ctx.var_vals) if stats is None else stats[0]
    if PROFILE: # backdate the event to the start of the call, the viz matches a device range with the exec event before it
      outputs, inputs = get_call_outs_ins(kcall)
      cpu_events.append(ProfilePointEvent(device, "exec", len(cpu_events), {"var_vals": ctx.var_vals,
-        "bufs": [b.trace_num for b in bufs], "name": get_call_name(kcall, bufs, ctx.var_vals), "outputs": outputs, "inputs": inputs}, ts=st))
+        "bufs": [b.trace_num for b in bufs], "name": display_name, "outputs": outputs, "inputs": inputs}, ts=st))
    if DEBUG < 2 or not ctx.update_stats: continue
    if et is None:
      Device[device].synchronize()
      et, st = float(perf_counter_us() - st)*1e-6, perf_counter_us()
      GlobalCounters.time_sum_s += et

-    estimates = estimate_uop(kcall)
-    display_name = get_call_name(kcall, bufs, ctx.var_vals)
+    estimates = estimate_uop(kcall) if stats is None else stats[1]
    op_est, mem_est, lds_est = (sym_infer(x, ctx.var_vals) for x in (estimates.ops, estimates.mem, estimates.lds))
-    header_color = 'magenta' if ctx.jit else ('green' if kcall.src[0].key not in first_run_cache else None)
+    key = kcall.src[0].key if stats is None else stats[2]
+    header_color = 'magenta' if ctx.jit else ('green' if key not in first_run_cache else None)
    ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
    flops, membw, ldsbw = op_est/(et or 1e-20), mem_est/(et or 1e-20), lds_est/(et or 1e-20)
    flops_str = f"{flops*1e-9:7.0f} GFLOPS" if flops < 1e14 else colored(f"{flops*1e-12:7.0f} TFLOPS", 'green')
@@ -92,7 +94,7 @@ def track_stats(ctx:ExecContext, call:UOp, st:decimal.Decimal, ets:list[float|No
    print(f"{colored(f'*** {device[:7]:7s} {GlobalCounters.kernel_count:4d}', header_color)}"+
      f" {ansipad(display_name, 46)} arg {len(bufs):2d} mem {GlobalCounters.mem_used/1e9:6.2f} GB"+
      ("" if et is None else f" tm {ptm}/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({flops_str} {mem_str})"))
-    first_run_cache.add(kcall.src[0].key)
+    first_run_cache.add(key)

 local_size_cache: dict[bytes, tuple[int, ...]] = {}
 def optimize_local_size(call:UOp, prg:UOp) -> UOp|None:
@@ -221,13 +223,14 @@ def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
    call = call.substitute({call.src[1+info.inputs]: UOp.mstack(*tables)})
  exec_kernel(replace(ctx, var_vals={**ctx.var_vals, "hcq_inputs_ptr": dev.rt_buffer()._buf.va_addr + base}), call, ast)

-  def _prof_tm(device:str, stat_call:UOp, prof:tuple[int, ...]) -> float|None:
-    (d:=cast(Any, Device[device])).prof_ents[prof[0]] = ProfileGraphEntry(device, stat_call.arg.name, prof[0], prof[1], stat_call.key)
+  def _prof_tm(device:str, name:str, prof:tuple[int, ...], profile_key:bytes) -> float|None:
+    (d:=cast(Any, Device[device])).prof_ents[prof[0]] = ProfileGraphEntry(device, name, prof[0], prof[1], profile_key)
    if not ctx.wait: return None
    d.synchronize(timeout=ctx.timeout)
    st, en = (d.signal(x)._buf.cpu_view().view(fmt='Q')[0] for x in prof)
    return float(en-st)/d.timestamp_divider/1e6
-  return [_prof_tm(device, k, prof) for devices, k, prof in info.kernels if prof for device in devices] if PROFILE or ctx.wait else []
+  return [_prof_tm(device, name, prof, profile_key) for devices,name,_,prof,profile_key in info.kernels
+          if prof for device in devices] if PROFILE or ctx.wait else []

 # flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src
 pm_flatten_linear = PatternMatcher([
@@ -1,7 +1,7 @@
 from __future__ import annotations
 import time
 START_TIME = time.perf_counter()
-import os, functools, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc
+import os, functools, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass, gc, threading
 from collections import defaultdict
 import shutil, math, types, copyreg, inspect, importlib, decimal, itertools, difflib
 from dataclasses import dataclass, field, replace
@@ -398,18 +398,17 @@ cache_dir: str = os.path.join(getenv("XDG_CACHE_HOME", os.path.expanduser("~/Lib
 CACHEDB: str = getenv("CACHEDB", os.path.abspath(os.path.join(cache_dir, "cache.db")))

 VERSION = 22
-_db_connection = None
+_db_connection = threading.local()
 def db_connection():
-  global _db_connection
-  if _db_connection is None:
+  if (conn:=getattr(_db_connection, "conn", None)) is None:
    os.makedirs(CACHEDB.rsplit(os.sep, 1)[0], exist_ok=True)
-    _db_connection = sqlite3.connect(CACHEDB, timeout=60, isolation_level="IMMEDIATE")
+    conn = _db_connection.conn = sqlite3.connect(CACHEDB, timeout=60, isolation_level="IMMEDIATE")
    # another connection has set it already or is in the process of setting it
    # that connection will lock the database
-    with contextlib.suppress(sqlite3.OperationalError): _db_connection.execute("PRAGMA journal_mode=WAL").fetchone()
-    _db_connection.execute("PRAGMA synchronous=NORMAL")
-    if DEBUG >= 8: _db_connection.set_trace_callback(print)
-  return _db_connection
+    with contextlib.suppress(sqlite3.OperationalError): conn.execute("PRAGMA journal_mode=WAL").fetchone()
+    conn.execute("PRAGMA synchronous=NORMAL")
+    if DEBUG >= 8: conn.set_trace_callback(print)
+  return conn

 def diskcache_clear():
  cur = db_connection().cursor()
@@ -476,7 +475,7 @@ def fetch(url:str, name:pathlib.Path|str|None=None, subdir:str|None=None, gunzip
  if not fp.is_file() or not allow_caching or (sha256 and hashlib.sha256(fp.read_bytes()).hexdigest() != sha256):
    if extract: shutil.rmtree(extract_dir, ignore_errors=True)
    (_dir := fp.parent).mkdir(parents=True, exist_ok=True)
-    with urllib.request.urlopen(urllib.request.Request(url, headers={"User-Agent": "tinygrad 0.13.0", **headers}), timeout=10) as r:
+    with urllib.request.urlopen(urllib.request.Request(url, headers={"User-Agent": "tinygrad 0.14.0", **headers}), timeout=10) as r:
      assert r.status in {200, 206}, r.status
      length = int(r.headers.get('content-length', 0)) if not gunzip else None
      readfile = gzip.GzipFile(fileobj=r) if gunzip else r
@@ -141,7 +141,8 @@ def main():
  args = parser.parse_args()

  # load the model
-  model, kv = Transformer.from_gguf(fetch(models.get(args.model, args.model)), args.max_context)
+  with Context(DEBUG=max(DEBUG.value, 2 if args.serve else 0)):
+    model, kv = Transformer.from_gguf(fetch(models.get(args.model, args.model)), args.max_context)
  model_name = kv.get('general.name') or kv.get('general.basename') or args.model
  file_sizes = [y.nbytes() for y in UOp.sink(*[x.uop for x in nn.state.get_parameters(model)]).toposort() if y.op is Ops.BUFFER]
  print(f"using model \"{model_name}\" with {sum(file_sizes):,} bytes and {sum(x.numel() for x in nn.state.get_parameters(model)):,} params, "
@@ -1,7 +1,7 @@
 from __future__ import annotations
 import functools, math
 from typing import Callable, cast
-from tinygrad import Tensor, UOp, nn, Device, Context, getenv
+from tinygrad import Tensor, UOp, nn, Device, Context
 from tinygrad.device import Buffer
 from tinygrad.dtype import AddrSpace, dtypes
 from tinygrad.helpers import prod
@@ -13,7 +13,6 @@ WAVES_M, WAVES_N, LANES_PER_WAVE_M, LANES_PER_WAVE_N = 2, 2, 2, 16
 WMMA_ACC, THREADS_PER_BLOCK = WMMA_M // LANES_PER_WAVE_M, WARP_SIZE * WAVES_M * WAVES_N
 LDS_PAD, WMMA_ARG, LOG2E = 4, ((WMMA_M, WMMA_N, WMMA_K), 'AMD', 32), math.log2(math.e)
 Q4_K, Q5_K, Q6_K, IQ4_XS, GGML_BLOCK_SIZE, Q8_GROUP_SIZE, Q4_WORDS, Q5_WORDS, Q6_BYTES, IQ4_WORDS = 12, 13, 14, 23, 256, 32, 36, 44, 210, 34
-Q6_PADDED, Q6_WORDS = 212, 53  # the 210-byte Q6 blocks are padded to 212 bytes so they are word-addressable
 QUANT_SIZES = {Q4_K: Q4_WORDS*4, Q5_K: Q5_WORDS*4, Q6_K: Q6_BYTES, IQ4_XS: IQ4_WORDS*4}  # bytes per 256-weight block

 def kernel_var(x:UOp) -> UOp:
@@ -62,44 +61,18 @@ class Linear(nn.Linear):
    self.ggml_type = packed_sizes[prod(raw.shape)]
    # store a typed buffer view: a lazy BITCAST is decomposed into byte-combining ALU before custom-kernel
    # scheduling and would copy the entire packed weight on every JIT graph
-    if self.ggml_type == Q6_K:
-      # Q6 blocks are 210 bytes, so consecutive blocks are only 2-byte aligned. pad each block to 212 bytes
-      # (a one-time copy at load) so the kernel can do all its reads as aligned u32 words
-      nbytes, nblocks = raw.max_numel(), raw.max_numel() // Q6_BYTES
-      byte_view = Tensor(UOp.from_buffer(cast(Buffer, raw.buf_uop.buffer).view(nbytes, dtypes.uint8, raw_offset)))
-      padded = byte_view.reshape((nblocks, Q6_BYTES)).pad_to((nblocks, Q6_PADDED)).contiguous().realize()
-      self.weight = Tensor(UOp.from_buffer(cast(Buffer, padded.uop.buf_uop.buffer).view(nblocks * Q6_WORDS, dtypes.uint32, 0)))
-    else:
-      self.weight = Tensor(UOp.from_buffer(cast(Buffer, raw.buf_uop.buffer)
-        .view(raw.max_numel() * raw.dtype.itemsize // dtypes.uint32.itemsize, dtypes.uint32, raw_offset)))
-  def prep_quant(self, x:Tensor) -> tuple[Tensor, Tensor, Tensor]|None:
-    # precompute the q8 activation so several linears can share it (gate/up, q/k/v). None if the custom path won't be used
-    if getenv("LLM_NO_QSHARE"): return None
-    supported = self.use_custom_quant and amd_custom_kernels_supported(self.weight.device)
-    if self.ggml_type is None and supported: self.set_quantized(self.weight)
-    if self.ggml_type not in (Q4_K, Q5_K, Q6_K, IQ4_XS) or not supported: return None
-    if isinstance(x.numel(), int): return q8_quantize(x, int(x.numel()) // self.in_features, self.in_features)
-    xp = x.pad_to(x.max_shape)
-    return q8_quantize(xp, int(xp.numel()) // self.in_features, self.in_features)
-
-  def __call__(self, x:Tensor, xq:tuple[Tensor, Tensor, Tensor]|None=None) -> Tensor:
+    packed_dtype = dtypes.uint8 if self.ggml_type == Q6_K else dtypes.uint32
+    self.weight = Tensor(UOp.from_buffer(cast(Buffer, raw.buf_uop.buffer)
+      .view(raw.max_numel() * raw.dtype.itemsize // packed_dtype.itemsize, packed_dtype, raw_offset)))
+  def __call__(self, x:Tensor) -> Tensor:
    supported = self.use_custom_quant and amd_custom_kernels_supported(self.weight.device)
    if self.ggml_type is None and supported:
      self.set_quantized(self.weight)
-      if self.ggml_type is None:
-        # tiny dense fp16 matmul (e.g. the ssm beta/alpha head rows): single fp16 gemv kernel instead of a
-        # generic matmul schedule, and realize the densely packed weight once if it is still a lazy ggml view
-        if self.weight.dtype in (dtypes.half, dtypes.float, dtypes.bfloat16) and self.out_features <= 2048 \
-          and self.in_features % (WARP_SIZE*4) == 0 and not getenv("LLM_NO_F16GEMV"):
-          numel, max_shape = x.numel(), x.max_shape
-          if isinstance(numel, int) or prod(max_shape) // self.in_features <= 32:
-            out = f16_gemv(self, x if isinstance(numel, int) else x.pad_to(max_shape))
-            return out if isinstance(numel, int) else out.shrink(tuple((0, s) for s in (*x.shape[:-1], self.out_features)))
-        self.use_custom_quant = supported = False  # not a supported quant format
+      if self.ggml_type is None: self.use_custom_quant = supported = False  # not a supported quant format
    if self.ggml_type in (Q4_K, Q5_K, Q6_K, IQ4_XS) and supported:
-      if isinstance(x.numel(), int): return q8_linear(self, x, xq)
+      if isinstance(x.numel(), int): return q8_linear(self, x)
      # symbolic token count: pad to the max chunk size so the kernels see static shapes, garbage rows are sliced off
-      out = q8_linear(self, x.pad_to(x.max_shape), xq)
+      out = q8_linear(self, x.pad_to(x.max_shape))
      return out.shrink(tuple((0, s) for s in (*x.shape[:-1], self.out_features)))
    return super().__call__(x)

@@ -113,8 +86,8 @@ def _amd_load(ptr:UOp, lanes:int|None=None) -> UOp:
  assert ptr.op is Ops.INDEX
  if lanes is None: return UOp(Ops.CUSTOMI, ptr.dtype, (ptr,), arg="__builtin_nontemporal_load({0})")
  buf, coords = ptr.src[0], ptr.src[1:]
-  idx = sum((coord*math.prod(buf.shape[i+1:]) for i,coord in enumerate(coords)), UOp.const(0, dtypes.weakint))
-  return UOp(Ops.SHRINK, src=(buf.flatten(), idx, UOp.const(lanes, dtypes.weakint))).load(dtype=ptr.dtype)
+  idx = sum((coord*math.prod(buf.shape[i+1:]) for i,coord in enumerate(coords)), UOp.const(0))
+  return UOp(Ops.SHRINK, src=(buf.flatten(), idx, UOp.const(lanes))).load(dtype=ptr.dtype)

 def _load_byte(raw:UOp, base:UOp, offset:UOp) -> UOp: return (raw[base + offset//4] >> ((offset&3)*8).cast(dtypes.uint32)) & 255
 def _half(value:UOp) -> UOp: return value.cast(dtypes.uint16).bitcast(dtypes.float16).float()
@@ -126,12 +99,10 @@ def _iq4_bytes(packed:UOp, shift:int) -> UOp:
  return _amd_byte_perm(high, low, 0x03020100 | ((selectors & 0x08080808) >> 1))

 def _q5_scales(raw:UOp, base:UOp, subgroup:UOp) -> tuple[UOp, UOp, UOp, UOp]:
-  # scales/mins (6-bit each) live in block bytes 4-15: three words total, same for the whole super-block's lanes
-  w1, w2, w3 = _amd_load(raw[base+1]), _amd_load(raw[base+2]), _amd_load(raw[base+3])
-  sb = (subgroup & 3) * 8  # byte within word
-  byte1, byte2, byte3 = (w1 >> sb) & 255, (w2 >> sb) & 255, (w3 >> sb) & 255
-  scale = (subgroup < 4).where(byte1 & 63, (byte3 & 15) | ((byte1 >> 6) << 4))
-  minimum = (subgroup < 4).where(byte2 & 63, (byte3 >> 4) | ((byte2 >> 6) << 4))
+  scale = (subgroup < 4).where(_load_byte(raw, base, 4 + subgroup) & 63,
+    (_load_byte(raw, base, 8 + subgroup) & 15) | ((_load_byte(raw, base, subgroup) >> 6) << 4))
+  minimum = (subgroup < 4).where(_load_byte(raw, base, 8 + subgroup) & 63,
+    (_load_byte(raw, base, 8 + subgroup) >> 4) | ((_load_byte(raw, base, 4 + subgroup) >> 6) << 4))
  d, dmin = (raw[base] & 0xffff).cast(dtypes.uint16), (raw[base] >> 16).cast(dtypes.uint16)
  return _half(d), _half(dmin), scale.float(), minimum.float()

@@ -147,67 +118,52 @@ def iq4_half_lut(device:str) -> Tensor:
                dtype=dtypes.float16, device=device).bitcast(dtypes.uint32).contiguous()

@functools.cache
-def _q8_quantize_kernel(q:UOp, scale:UOp, xsum:UOp, x:UOp, tokens:int, in_features:int) -> UOp:
+def _q8_quantize_kernel(q:UOp, scale:UOp, x:UOp, tokens:int, in_features:int) -> UOp:
  groups = in_features//Q8_GROUP_SIZE
-  token_group, lane = UOp.range(tokens*groups, 0, axis_type=AxisType.GLOBAL), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  token_group, lane = UOp.range(tokens*groups, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
  token, group = token_group//groups, token_group%groups
  x = x.reshape(tokens, groups, 32)
  group_scale = (warp_reduce(x[token, group, lane].float().abs(), maximum=True, full_wave=True) / 127).maximum(1e-8)
  word_lane = lane.minimum(7)
  xs = tuple(x[token, group, word_lane*4+i].float() for i in range(4))
-  qs = tuple((v/group_scale).round().clip(-127, 127).cast(dtypes.int8) for v in xs)
-  word = sum((v.cast(dtypes.uint8).cast(dtypes.uint32) << (i*8) for i, v in enumerate(qs)), UOp.const(0, dtypes.uint32))
-  # per-16 sums of the quantized values (lanes 0-3 / 4-7): Q4_K/Q5_K need the 32-sum, Q6_K the 16-sums
-  part = (lane < 8).where(sum((v.cast(dtypes.int32) for v in qs), UOp.const(0, dtypes.int32)), UOp.const(0, dtypes.int32))
-  gsum = [warp_reduce(((lane & 4).eq(h*4)).where(part, UOp.const(0, dtypes.int32)), full_wave=True) for h in range(2)]
-  store_half = (lane & 4) >> 2
-  stores = (q[token, group, lane.valid(lane < 8)].store(word),
-            UOp.group(scale[token, group.valid(lane.eq(0))].store(group_scale),
-                      xsum[token, group, store_half.valid(lane.eq(0) | lane.eq(4))].store(
-                        store_half.eq(0).where(gsum[0].float(), gsum[1].float()))))
+  word = sum(((v/group_scale).round().clip(-127, 127).cast(dtypes.int8).cast(dtypes.uint8).cast(dtypes.uint32) << (i*8)
+              for i,v in enumerate(xs)), UOp.const(0, dtypes.uint32))
+  stores = (q[token, group, lane.valid(lane < 8)].store(word), scale[token, group.valid(lane.eq(0))].store(group_scale))
  return UOp.group(*stores).end(token_group, lane).sink(arg=KernelInfo(name="q8_quantize", opts_to_apply=()))

-def q8_quantize(x:Tensor, tokens:int, in_features:int) -> tuple[Tensor, Tensor, Tensor]:
+def q8_quantize(x:Tensor, tokens:int, in_features:int) -> tuple[Tensor, Tensor]:
  groups = in_features//Q8_GROUP_SIZE
  q = Tensor.empty(tokens, groups, 8, dtype=dtypes.uint32, device=x.device)
  scale = Tensor.empty(tokens, groups, dtype=dtypes.float32, device=x.device)
-  xsum = Tensor.empty(tokens, groups, 2, dtype=dtypes.float32, device=x.device)
-  q, scale, xsum = Tensor.custom_kernel(q, scale, xsum, x, fxn=functools.partial(_q8_quantize_kernel, tokens=tokens, in_features=in_features))[:3]
-  return q, scale, xsum
+  q, scale = Tensor.custom_kernel(q, scale, x, fxn=functools.partial(_q8_quantize_kernel, tokens=tokens, in_features=in_features))[:2]
+  return q, scale

 def _decode_linear(out:UOp, out_features:int, group_count:int, group_dot, name:str) -> UOp:
-  chunks = out.shape[2]
-  # two-dim global grid instead of one flat grid: no div/mods needed to decompose the gid
-  token_output = UOp.range(out.shape[0]*out_features, 0, axis_type=AxisType.GLOBAL)
-  chunk, lane = UOp.range(chunks, 1, axis_type=AxisType.GLOBAL), UOp.range(32, 2, axis_type=AxisType.LOCAL)
-  token, output = token_output // out_features, token_output % out_features
-  group = (lane+chunk*32).minimum(group_count-1)
-  value = group_dot(token, output, group) if chunks*32 == group_count else \
-    (lane+chunk*32 < group_count).where(group_dot(token, output, group), UOp.const(0, dtypes.float32))
+  chunks = (group_count+31)//32
+  token_output_chunk, lane = UOp.range(out.shape[0]*out_features*chunks, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  token, output, chunk = token_output_chunk // (out_features*chunks), (token_output_chunk//chunks) % out_features, token_output_chunk % chunks
+  group = lane+chunk*32
+  value = group_dot(token, output, group) if group_count % 32 == 0 else \
+    (group < group_count).where(group_dot(token, output, group.minimum(group_count-1)), UOp.const(0, dtypes.float32))
  total = warp_reduce(value, full_wave=True)
-  return out[token, output, chunk.valid(lane.eq(0))].store(total.cast(out.dtype)).end(token_output, chunk, lane).sink(
+  return out[token, output, chunk.valid(lane.eq(0))].store(total.cast(out.dtype)).end(token_output_chunk, lane).sink(
    arg=KernelInfo(name=name, opts_to_apply=()))

@functools.cache
-def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, xs:UOp, out_features:int, in_features:int, ggml_type:int) -> UOp:
+def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, out_features:int, in_features:int, ggml_type:int) -> UOp:
  group_count = in_features // Q8_GROUP_SIZE
  def group_dot(token:UOp, output:UOp, group:UOp) -> UOp:
    block, subgroup = group // 8, group % 8
    xwords = _amd_load(xq[token, group, 0], 8)
    if ggml_type in (Q4_K, Q5_K):
      base = (output * in_features//GGML_BLOCK_SIZE + block) * (Q4_WORDS if ggml_type == Q4_K else Q5_WORDS)
-      qs_base, dot = base + (4 if ggml_type == Q4_K else 12) + (subgroup//2)*8, UOp.const(0, dtypes.int32)
-      # vectorize the 8 packed-weight words and (for Q5_K) the 32-byte high-bit bitmap
-      qs_pair = (_amd_load(raw[qs_base], 4), _amd_load(raw[qs_base+4], 4))
-      zero = UOp.const(0, dtypes.uint32)
-      qh_pair = (_amd_load(raw[base+4], 4), _amd_load(raw[base+8], 4)) if ggml_type == Q5_K else (zero, zero)
+      qs_base, dot, qsum = base + (4 if ggml_type == Q4_K else 12) + (subgroup//2)*8, UOp.const(0, dtypes.int32), UOp.const(0, dtypes.int32)
      for word_idx in range(8):
-        word = (qs_pair[word_idx//4][word_idx%4] >> ((subgroup&1)*4).cast(dtypes.uint32)) & 0x0f0f0f0f
-        if ggml_type == Q5_K: word |= ((qh_pair[word_idx//4][word_idx%4] >> subgroup.cast(dtypes.uint32)) & 0x01010101) << 4
-        dot = _amd_dp4a(word, xwords[word_idx], dot)
+        word = (raw[qs_base+word_idx] >> ((subgroup&1)*4).cast(dtypes.uint32)) & 0x0f0f0f0f
+        if ggml_type == Q5_K: word |= ((raw[base+4+word_idx] >> subgroup.cast(dtypes.uint32)) & 0x01010101) << 4
+        dot, qsum = _amd_dp4a(word, xwords[word_idx], dot), _amd_dp4a(UOp.const(0x01010101, dtypes.uint32), xwords[word_idx], qsum)
      d, dmin, scale, minimum = _q5_scales(raw, base, subgroup)
-      gsum = xs[token, group, 0].load() + xs[token, group, 1].load()
-      return (dot.float()*d*scale - gsum*dmin*minimum) * xd[token, group]
+      return (dot.float()*d*scale - qsum.float()*dmin*minimum) * xd[token, group]
    if ggml_type == IQ4_XS:
      base = (output * in_features//GGML_BLOCK_SIZE + block) * IQ4_WORDS
      dot = UOp.const(0, dtypes.int32)
@@ -216,24 +172,18 @@ def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, xs:UOp, out_features:
        dot = _amd_dp4a(_iq4_bytes(packed, 4*(word_idx//4)), xwords[word_idx], dot)
      d, scale = _iq4_scales(raw, base, subgroup)
      return dot.float() * xd[token, group] * d * scale
-    # the packed rows were padded to 212 bytes (53 words) per 256-block in set_quantized: everything is word-aligned
-    base = (output*in_features//GGML_BLOCK_SIZE+block)*Q6_WORDS
-    # the subgroup's 8 ql words and 8 qh words are contiguous: two 16-byte vector loads each
-    lows = tuple(_amd_load(raw[base + (subgroup//4)*16 + (subgroup%2)*8 + half*4], 4) for half in range(2))
-    highs = tuple(_amd_load(raw[base + 32 + (subgroup//4)*8 + half*4], 4) for half in range(2))
+    base = (output*in_features//GGML_BLOCK_SIZE+block)*Q6_BYTES
    dots = [UOp.const(0, dtypes.int32)] * 2
    for word_idx in range(8):
-      within = (subgroup*32 + word_idx*4)%128
-      low = lows[word_idx//4][word_idx%4] >> ((within//64)*4).cast(dtypes.uint32)
-      high = highs[word_idx//4][word_idx%4] >> ((within//32)*2).cast(dtypes.uint32)
-      # 4 values per word: (low nibble) | (2 high bits << 4). values stay positive, so the int8-bitcast/-32 of the
-      # naive dequant is skipped and the -32 offset is applied later via the per-16 sums of the quantized inputs
-      word = (low & 0x0f0f0f0f) | ((high & 0x03030303) << 4)
+      pos, within = subgroup*32 + word_idx*4, (subgroup*32 + word_idx*4)%128
+      low = _amd_load(raw[base + (pos//128)*64 + within%64], 4) >> ((within//64)*4).cast(dtypes.uint8)
+      high = _amd_load(raw[base + 128 + (pos//128)*32 + within%32], 4) >> ((within//32)*2).cast(dtypes.uint8)
+      quant = ((low & 15) | ((high & 3) << 4)).bitcast(dtypes.int8) - 32
+      word = sum((quant[i].cast(dtypes.uint8).cast(dtypes.uint32) << (i*8) for i in range(4)), UOp.const(0, dtypes.uint32))
      dots[word_idx//4] = _amd_dp4a(word, xwords[word_idx], dots[word_idx//4])
-    scales = [((raw[base + 48 + (subgroup*2+i)//4] >> (((subgroup*2+i)%4)*8).cast(dtypes.uint32)) & 255)
-              .cast(dtypes.uint8).bitcast(dtypes.int8).float() for i in range(2)]
-    gsum = [xs[token, group, i].load() * 32 for i in range(2)]
-    return ((dots[0].float() - gsum[0])*scales[0] + (dots[1].float() - gsum[1])*scales[1]) * xd[token, group] * _half(raw[base+52] & 0xffff)
+    scales = [raw[base + 192 + subgroup*2+i].cast(dtypes.uint8).bitcast(dtypes.int8).float() for i in range(2)]
+    dbits = raw[base+208].cast(dtypes.uint16) | (raw[base+209].cast(dtypes.uint16) << 8)
+    return (dots[0].float()*scales[0] + dots[1].float()*scales[1]) * xd[token, group] * _half(dbits)
  names = {Q4_K: "linear_q4_k", Q5_K: "linear_q5_k", IQ4_XS: "linear_iq4_xs", Q6_K: "linear_q6"}
  return _decode_linear(out, out_features, group_count, group_dot, names[ggml_type])

@@ -314,14 +264,13 @@ def _iq4_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, lut:UOp, out_features:i
      pairs = tuple(lut[((raw[base + 2 + subgroup*4 + word] >> (byte*8)) & 255).cast(dtypes.weakint)]
                    for word in range(4) for byte in range(4))
      return tuple((_half((pair >> (half*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in pairs)
-    # a subgroup-half gathers the lo (half=0) or hi (half=1) nibbles of byte pairs of each packed word
-    lut_pairs = (lut[(((raw[base+2+subgroup*4+i] >> (8*j+4*half)) & 15) |
-                      (((raw[base+2+subgroup*4+i] >> (8*j+8+4*half)) & 15) << 4)).cast(dtypes.weakint)]
-                 for i in range(4) for j in (0, 2))
+    def nibble(packed:UOp, index:int): return (packed >> (8*index+4*half)) & 15
+    lut_pairs = (lut[(nibble(packed, i) | nibble(packed, i+1)<<4).cast(dtypes.weakint)]
+                 for packed in (raw[base+2+subgroup*4+i] for i in range(4)) for i in (0, 2))
    return tuple((_half((pair >> (i*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in lut_pairs for i in range(2))
  return _quant_linear_wmma(out, x, out_features, in_features, IQ4_WORDS, layout, dequant, "linear_iq4_xs_f16_wmma")

-def q8_linear(layer:Linear, x:Tensor, xq:tuple[Tensor, Tensor, Tensor]|None=None) -> Tensor:
+def q8_linear(layer:Linear, x:Tensor) -> Tensor:
  assert layer.ggml_type in (Q4_K, Q5_K, Q6_K, IQ4_XS)
  tokens = int(x.numel()) // layer.in_features
  raw, out_features, in_features = layer.weight.uop.buf_uop, layer.out_features, layer.in_features
@@ -338,228 +287,67 @@ def q8_linear(layer:Linear, x:Tensor, xq:tuple[Tensor, Tensor, Tensor]|None=None
    fxn = _iq4_linear_f16_wmma_kernel if layer.ggml_type == IQ4_XS else functools.partial(_q5_linear_f16_wmma_kernel, ggml_type=layer.ggml_type)
    extra = (iq4_half_lut(str(x.device)).uop,) if layer.ggml_type == IQ4_XS else ()
    return run(fxn, out, raw, x.cast(dtypes.float16).contiguous().uop, *extra)
-  if xq is None: xq = q8_quantize(x, tokens, in_features)
-  xq_, xd, xs = xq
+  xq, xd = q8_quantize(x, tokens, in_features)
  decode = functools.partial(_quant_decode_kernel, ggml_type=layer.ggml_type)
  out = Tensor.empty(tokens, out_features, (in_features+1023)//1024, dtype=dtypes.float32, device=x.device).uop
-  return run(decode, out, raw, xq_.uop, xd.uop, xs.uop)
-
-# ******** tiny dense fp16 gemv ********
-
-@functools.cache
-def _amd_f16_gemv_kernel(out:UOp, w:UOp, x:UOp, *rest:UOp, in_features:int, out_features:int, tokens:int) -> UOp:
-  bias: UOp|None = rest[0] if rest else None
-  # one block per (token, output row), 32 lanes accumulate 4-wide chunks of the row
-  lanes, val_chunk = WARP_SIZE, 4
-  token, out_row = UOp.range(tokens, 0, AxisType.GLOBAL), UOp.range(out_features, 1, AxisType.GLOBAL)
-  lane = UOp.range(lanes, 2, axis_type=AxisType.LOCAL)
-  per = in_features // (lanes * val_chunk)
-  assert per * lanes * val_chunk == in_features
-  w = w.reshape((out_features, per, lanes*val_chunk))
-  x = x.reshape((tokens, per, lanes*val_chunk))
-  acc = UOp.const(0, dtypes.float32)
-  for i in range(per):
-    for j in range(val_chunk):
-      acc = acc + w[out_row, i, lane*val_chunk + j].load().float() * x[token, i, lane*val_chunk + j].load().float()
-  total = warp_reduce(acc, full_wave=True)
-  if bias is not None: total = total + bias[token, out_row].load().float()
-  return out[token, out_row.valid(lane.eq(0))].store(total).end(token, out_row, lane).sink(arg=KernelInfo(name="linear_f16_gemv", opts_to_apply=()))
-
-def _view_back(t:Tensor) -> Tensor:
-  """strip top-of-chain CAST(s) from a lazy weight: reading the raw file bytes in the kernel instead of
-  materializing the cast into a fresh buffer every step"""
-  uop = t.uop
-  while uop.op is Ops.CAST: uop = uop.src[0]
-  return Tensor(uop).reshape(t.shape)
-
-def f16_gemv(layer:Linear, x:Tensor) -> Tensor:
-  tokens = prod(x.shape[:-1])
-  assert isinstance(tokens, int)
-  weight = _view_back(layer.weight)
-  x = x.contiguous() if x.dtype == dtypes.half else x.cast(dtypes.half).contiguous()
-  out = Tensor.empty(tokens, layer.out_features, dtype=dtypes.float32, device=x.device)
-  fxn = functools.partial(_amd_f16_gemv_kernel, in_features=layer.in_features, out_features=layer.out_features, tokens=tokens)
-  srcs = (out, weight.reshape(-1), x.reshape(tokens, layer.in_features)) + (() if layer.bias is None else (_view_back(layer.bias),))
-  return Tensor.custom_kernel(*srcs, fxn=fxn)[0].reshape(*x.shape[:-1], layer.out_features)
-
-# ******** ssm beta/alpha joint gemv + epilogue ********
-
-@functools.cache
-def _amd_ssm_ab_kernel(out:UOp, x:UOp, w:UOp, dt:UOp, ssm_a:UOp, dim:int, rows:int) -> UOp:
-  heads = out.shape[-1] // 2
-  rh, lane = UOp.range(rows*heads*2, 0), UOp.range(WARP_SIZE, 1, axis_type=AxisType.LOCAL)
-  row, o = rh // (heads*2), rh % (heads*2)
-  h, is_alpha = o % heads, o >= heads
-  dot = UOp.const(0, dtypes.float32)
-  for i in range(dim // (WARP_SIZE*8)):
-    c = lane*8 + i*(WARP_SIZE*8)
-    xf, wf = _vec_load(x[row*dim + c], 8), _vec_load(w[o*dim + c], 8)
-    dot = dot + sum((xv*wv for xv, wv in zip(xf, wf)), UOp.const(0, dtypes.float32))
-  assert dim % (WARP_SIZE*8) == 0
-  dot = warp_reduce(dot, full_wave=True)
-  store_val = is_alpha.where((dot + dt[h].load().float()).softplus() * ssm_a[h, 0].load().float(), dot.sigmoid())
-  return out[row, o.valid(lane.eq(0))].store(store_val).end(rh, lane).sink(arg=KernelInfo(name="ssm_beta_alpha", opts_to_apply=()))
-
-def ssm_beta_alpha(x:Tensor, w:Tensor, dt:Tensor, ssm_a:Tensor) -> Tensor:
-  # x: (rows, dim) fp16; w: (2*heads, dim) with beta rows first. returns (rows, 2*heads) fp32 (beta | log_alpha)
-  rows, dim, heads = x.shape[0], x.shape[1], ssm_a.shape[0]
-  assert w.shape == (2*heads, dim) and dt.shape == (heads,)
-  out = Tensor.empty(rows, 2*heads, dtype=dtypes.float32, device=x.device)
-  fxn = functools.partial(_amd_ssm_ab_kernel, dim=dim, rows=rows)
-  return Tensor.custom_kernel(out, x.reshape(rows*dim), w.reshape(2*heads*dim), dt.reshape(heads), ssm_a.reshape(heads, 1), fxn=fxn)[0]
+  return run(decode, out, raw, xq.uop, xd.uop)

 # ******** flash attention on the KV cache ********

-def _vec_load(ptr:UOp, lanes:int) -> tuple[UOp, ...]:
-  if lanes == 1: return (ptr.load().float(),)
-  vec = _amd_load(ptr, lanes)
-  return tuple(vec[i].float() for i in range(lanes))
-
-# ******** fused rmsnorm: one block does reduce + scale + apply ********
-
@functools.cache
-def _amd_rmsnorm_kernel(o:UOp, x:UOp, weight:UOp, eps:float, dim:int, waves:int) -> UOp:
-  n_rows = o.numel() // dim
-  lanes = waves * WARP_SIZE
-  global_row = UOp.range(n_rows, 0)
-  lane, wave = UOp.range(WARP_SIZE, 1, axis_type=AxisType.LOCAL), UOp.range(waves, 2, axis_type=AxisType.LOCAL)
-  row = global_row
-  n_per = -(-dim // lanes)
-  icol = (wave*WARP_SIZE + lane) * n_per
-  noload = dim % lanes != 0
-  cols = tuple((icol + i).valid(icol + i < dim) if noload else icol + i for i in range(n_per))
-  xs = [x[row, c].load().float() for c in cols]
-  total = warp_reduce(sum((v*v for v in xs), UOp.const(0, dtypes.float32)), full_wave=True)
-  # cross-wave reduce through LDS, then broadcast
-  part = UOp.placeholder((waves,), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
-  barrier = UOp.barrier(part[wave.valid(lane.eq(0))].store(total))
-  total = sum((part.after(barrier)[w].load() for w in range(waves)), UOp.const(0, dtypes.float32))
-  scale = (total / dim + eps).rsqrt()
-  stores = [o[row, c].store(xs[i] * scale * weight[row, c].load().float()) for i, c in enumerate(cols)]
-  return UOp.group(*stores).end(global_row, lane, wave).sink(arg=KernelInfo(name="rmsnorm", opts_to_apply=()))
-
-class RMSNorm(nn.RMSNorm):
-  def __call__(self, x:Tensor) -> Tensor: return amd_rmsnorm(self, x)
-
-def amd_rmsnorm(norm:nn.RMSNorm, x:Tensor) -> Tensor:
-  w = norm.weight
-  if w is not None and not getenv("LLM_NO_RMSNORM") and x.dtype == dtypes.float32 and amd_custom_kernels_supported(x.device) \
-    and isinstance(x.numel(), int) and x.shape[-1] >= 192 and x.shape[-1] % 32 == 0:
-    rows, dim = int(x.numel()) // x.shape[-1], x.shape[-1]
-    out = Tensor.empty(rows, dim, dtype=dtypes.float32, device=x.device)
-    waves = 4 if dim >= 3040 else 2
-    fxn = functools.partial(_amd_rmsnorm_kernel, eps=norm.eps, dim=dim, waves=waves)
-    w2 = _view_back(w).reshape(1, dim).expand(rows, dim)
-    return Tensor.custom_kernel(out, x.reshape(rows, dim).contiguous(), w2, fxn=fxn)[0].reshape(*x.shape)
-  return nn.RMSNorm.__call__(norm, x)
-
-@functools.cache
-def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, valid_kv_len, max_kv_len, block_n, waves=4):
+def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, valid_kv_len, max_kv_len, block_n):
  valid_kv_len = _unbind(valid_kv_len)
  _, B, H_KV, N, D = cast(tuple[int, int, int, int, int], cache_kv.shape)
  _, H, M, _ = cast(tuple[int, int, int, int], q.shape)
  assert M == 1 and H % H_KV == 0 and D % WARP_SIZE == 0 and max_kv_len <= N and max_kv_len % block_n == 0
-  G, CHUNK, DPL, WAVES = H // H_KV, block_n, D // WARP_SIZE, waves
-  assert CHUNK % WAVES == 0
-  SEC = CHUNK // WAVES  # keys each wave scans independently
-  live_chunks = (valid_kv_len+CHUNK-1)//CHUNK
-  live_chunks = min(live_chunks, out.shape[2]) if isinstance(live_chunks, int) else live_chunks.minimum(out.shape[2])
-  block_bhkv, block_chunk = UOp.range(B*H_KV, 0, AxisType.GLOBAL), UOp.range(live_chunks, 1, AxisType.GLOBAL)
-  lane, wave = UOp.range(WARP_SIZE, 2, axis_type=AxisType.LOCAL), UOp.range(WAVES, 3, axis_type=AxisType.LOCAL)
-  b, kv_head = block_bhkv // H_KV, block_bhkv % H_KV
-  # per-lane query fragments for every GQA head, kept packed in registers; unpacked at use
-  qf = tuple(_vec_load(q[b, kv_head*G+h, 0, lane*DPL], DPL) for h in range(G))
-  zerof = UOp.const(0, dtypes.float)
-  valids: list[UOp] = []
-  scores: list[list[UOp]] = [[zerof]*G for _ in range(SEC)]
-  vfrags: list[tuple[UOp, ...]] = [()]*SEC
-  for j in range(SEC):
-    key = block_chunk*CHUNK + wave*SEC + j
-    valid = key < valid_kv_len
-    valids.append(valid)
-    kfrag = _vec_load(cache_kv[0, b, kv_head, key, lane*DPL], DPL)
-    # V is prefetched in the score pass so both streams are in flight together
-    vfrags[j] = _vec_load(cache_kv[1, b, kv_head, key, lane*DPL], DPL)
-    for h in range(G):
-      s = warp_reduce(sum((qf[h][i]*kfrag[i] for i in range(DPL)), UOp.const(0, dtypes.float)), full_wave=True) * (1/math.sqrt(D))
-      scores[j][h] = valid.where(s, UOp.const(-math.inf, dtypes.float))
-  ninf = UOp.const(-math.inf, dtypes.float)
-  row_max = [functools.reduce(UOp.maximum, (scores[j][h] for j in range(SEC)), ninf) for h in range(G)]
-  accs:list[list[UOp]] = [[UOp.const(0, dtypes.float)] * DPL for _ in range(G)]
-  row_sums:list[UOp] = [UOp.const(0, dtypes.float) for _ in range(G)]
-  for j in range(SEC):
-    for h in range(G):
-      beta = valids[j].where(((scores[j][h]-row_max[h])*LOG2E).exp2(), UOp.const(0, dtypes.float))
-      accs[h] = [a + beta*v for a, v in zip(accs[h], vfrags[j])]
-      row_sums[h] = row_sums[h] + beta
-  # exchange across the block's waves through LDS (fp16 halves LDS so more blocks fit per CU)
-  acc_lds = UOp.placeholder((WAVES, G, D), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL)
-  ml_lds = UOp.placeholder((WAVES, G, 2), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
-  lds_acc = acc_lds.reshape(WAVES, G, WARP_SIZE, DPL)
-  stores = [lds_acc[wave, h, lane].store(UOp.stack(*accs[h]).cast(dtypes.half)) for h in range(G)]
-  # NOTE: duplicate stores of the same value from every lane are harmless here
-  stores += [ml_lds[wave, h, i].store(x) for h in range(G) for i, x in enumerate((row_max[h], row_sums[h]))]
-  barrier = UOp.barrier(UOp.group(*stores))
-  acc_lds, ml_lds = acc_lds.after(barrier), ml_lds.after(barrier)
-  tid = wave*WARP_SIZE + lane
-  final_stores:list[UOp] = []
-  for i in range(-(-G*D//(WAVES*WARP_SIZE))):
-    flat = tid + i*WAVES*WARP_SIZE
-    h, d = flat // D, flat % D
-    M = functools.reduce(UOp.maximum, (ml_lds[w, h, 0].load() for w in range(WAVES)), ninf)
-    val = sum((((ml_lds[w, h, 0].load()-M)*LOG2E).exp2() * acc_lds[w, h, d].load().float() for w in range(WAVES)), UOp.const(0, dtypes.float))
-    oidx = out[b, kv_head*G + h, block_chunk, d]
-    if G*D % (WAVES*WARP_SIZE): oidx = out[b, (kv_head*G + h).valid(flat < G*D), block_chunk, d]
-    final_stores.append(oidx.store(val))
-  hstat = tid
-  M = functools.reduce(UOp.maximum, (ml_lds[w, hstat, 0].load() for w in range(WAVES)), ninf)
-  L = sum((((ml_lds[w, hstat, 0].load()-M)*LOG2E).exp2() * ml_lds[w, hstat, 1].load() for w in range(WAVES)), UOp.const(0, dtypes.float))
-  q_head = (kv_head*G + hstat).valid(hstat < G) if WAVES*WARP_SIZE > G else kv_head*G + hstat
-  final_stores += [stats[b, q_head, block_chunk, 0].store(M), stats[b, q_head, block_chunk, 1].store(L)]
-  return UOp.group(*final_stores).end(lane, wave, block_chunk, block_bhkv).sink(arg=KernelInfo(name="flash_decode_partial", opts_to_apply=()))
-
-@functools.cache
-def _amd_flash_decode_combine(o:UOp, partial:UOp, stats:UOp, live:int|UOp) -> UOp:
-  # one wave per (batch, head, 64-dim tile): every lane redundantly weights its chunks; no cross-lane traffic
-  live = _unbind(live)
-  B, H, C, D = cast(tuple[int, int, int, int], partial.shape)
-  DT = 64 if D % 64 == 0 else WARP_SIZE  # dims per block
-  assert D % DT == 0
-  block_bh, block_dt = UOp.range(B*H, 0, AxisType.GLOBAL), UOp.range(D//DT, 1, AxisType.GLOBAL)
-  lane = UOp.range(WARP_SIZE, 2, axis_type=AxisType.LOCAL)
-  b, h = block_bh // H, block_bh % H
-  NPD = DT // WARP_SIZE  # output dims per lane
-  dims = tuple(block_dt*DT + lane*NPD + i for i in range(NPD))
-  chunk = UOp.range(live, 100, AxisType.REDUCE)
-  def iloop(ph, val): return ph.store(ph.const_like(val))
-  chunk_max = UOp.placeholder((1,), dtypes.float, slot=0, addrspace=AddrSpace.REG)
-  chunk_max_i = chunk_max.after(iloop(chunk_max, -math.inf))
-  update0 = chunk_max_i.store(chunk_max_i.after(chunk).maximum(stats[b, h, chunk, 0].load())).end(chunk)
-  chunk_max = chunk_max_i.after(update0)
-  chunk2 = UOp.range(live, 101, AxisType.REDUCE)
-  acc = UOp.placeholder((NPD,), dtypes.float, slot=1, addrspace=AddrSpace.REG)
-  weight_sum = UOp.placeholder((1,), dtypes.float, slot=2, addrspace=AddrSpace.REG)
-  acc_i, weight_sum_i = acc.after(iloop(acc, 0)), weight_sum.after(iloop(weight_sum, 0))
-  w = ((stats[b, h, chunk2, 0].load()-chunk_max)*LOG2E).exp2()
-  update1 = UOp.group(*[acc_i[i].store(acc_i.after(chunk2)[i].load() + w*partial[b, h, chunk2, d].load()) for i, d in enumerate(dims)],
-                      weight_sum_i[0].store(weight_sum_i.after(chunk2)[0].load() + w*stats[b, h, chunk2, 1].load())).end(chunk2)
-  acc, weight_sum = acc_i.after(update1), weight_sum_i.after(update1)
-  inv = 1 / weight_sum[0].load()
-  return UOp.group(*[o[b, h, 0, d].store(acc[i].load() * inv) for i, d in enumerate(dims)]) \
-    .end(lane, block_dt, block_bh).sink(arg=KernelInfo(name="flash_decode_combine", opts_to_apply=()))
+  G, CHUNK, DV, heads_per_wave = H // H_KV, block_n, D // WARP_SIZE, 2
+  head_tile = min(DECODE_HEAD_TILE, G)  # share each KV stream across two GQA heads per wave
+  assert G % head_tile == 0 and head_tile % heads_per_wave == 0
+  decode_waves, decode_group = head_tile // heads_per_wave, 4
+  block_bhkv = UOp.range(B*H_KV*(G//head_tile), 0, AxisType.GLOBAL)
+  valid_chunks = (valid_kv_len+CHUNK-1)//CHUNK
+  group_count = min(valid_chunks, out.shape[2]) if isinstance(valid_chunks, int) else valid_chunks.minimum(out.shape[2])
+  block_n, lane = UOp.range(group_count, 1, AxisType.GLOBAL), UOp.range(WARP_SIZE, 2, axis_type=AxisType.LOCAL)
+  wave = UOp.range(decode_waves, 3, axis_type=AxisType.LOCAL)
+  head_group, bhkv = block_bhkv % (G//head_tile), block_bhkv // (G//head_tile)
+  b, kv_head = bhkv // H_KV, bhkv % H_KV
+  dims = tuple(lane + i*WARP_SIZE for i in range(DV))
+  acc, row_max, row_sum = _reg((heads_per_wave, DV), 0, 0), _reg((heads_per_wave,), 1, -math.inf), _reg((heads_per_wave,), 2, 0)
+  groups_per_chunk, offset = CHUNK // decode_group, UOp.range(((valid_chunks+group_count-1)//group_count)*(CHUNK//decode_group), 100, AxisType.REDUCE)
+  chunk = block_n + (offset // groups_per_chunk) * group_count
+  keys = tuple(chunk*CHUNK + (offset % groups_per_chunk)*decode_group + i for i in range(decode_group))
+  valid = tuple(key < valid_kv_len for key in keys)
+  kvals, vvals = (tuple(tuple(is_valid.where(cache_kv[kv, b, kv_head, key, d].float(), UOp.const(0, dtypes.float)) for d in dims)
+    for key,is_valid in zip(keys, valid)) for kv in range(2))
+  q_heads = tuple(kv_head*G + head_group*head_tile + wave*heads_per_wave + head for head in range(heads_per_wave))
+  updates:list[UOp] = []
+  for head,q_head in enumerate(q_heads):
+    scores = tuple(warp_reduce(sum((q[b, q_head, 0, d].float()*k for d,k in zip(dims, key_kvals)),
+                                   UOp.const(0, dtypes.float)), full_wave=True) / math.sqrt(D) for key_kvals in kvals)
+    prev_acc, prev_max, prev_sum = acc.after(offset)[head], row_max.after(offset)[head], row_sum.after(offset)[head]
+    new_max = functools.reduce(lambda a,vs:a.maximum(vs[0].where(vs[1], UOp.const(-math.inf, dtypes.float))), zip(valid, scores), prev_max)
+    alpha = ((prev_max-new_max)*LOG2E).exp2()
+    betas = tuple(is_valid.where(((score-new_max)*LOG2E).exp2(), UOp.const(0, dtypes.float)) for is_valid,score in zip(valid, scores))
+    updates += [acc[head].store(prev_acc*alpha + sum((UOp.stack(*value)*beta for value,beta in zip(vvals, betas)), acc[head].const_like(0))),
+                row_sum[head].store(prev_sum*alpha + sum(betas, UOp.const(0, dtypes.float))), row_max[head].store(new_max)]
+  update = UOp.group(*updates).end(offset)
+  acc, row_max, row_sum = acc.after(update), row_max.after(update), row_sum.after(update)
+  stores = [out[b, q_head, block_n, d].store(acc[head, i]) for head,q_head in enumerate(q_heads) for i,d in enumerate(dims)] + \
+    [stats[b, q_head.valid(lane.eq(0)), block_n, i].store(x[head]) for head,q_head in enumerate(q_heads) for i,x in enumerate((row_max, row_sum))]
+  return UOp.group(*stores).end(lane, wave, block_n, block_bhkv).sink(arg=KernelInfo(name="flash_decode_partial", opts_to_apply=()))

 def amd_flash_attention_decode(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, max_kv_len:int) -> Tensor:
  B, H, D = cache_kv.shape[1], q.shape[1], cache_kv.shape[4]
-  chunks = min(256, max_kv_len // 64)
+  chunks = min(64, max_kv_len // 128)
  partial = Tensor.empty(B, H, chunks, D, dtype="float32", device=q.device)
  stats = Tensor.empty(B, H, chunks, 2, dtype="float32", device=q.device)
-  fxn = functools.partial(_amd_flash_attention_decode_partial, valid_kv_len=valid_kv_len, max_kv_len=max_kv_len, block_n=64, waves=16)
+  fxn = functools.partial(_amd_flash_attention_decode_partial, valid_kv_len=valid_kv_len, max_kv_len=max_kv_len, block_n=128)
  partial, stats = Tensor.custom_kernel(partial, stats, q, cache_kv, fxn=fxn)[:2]
-  live = (valid_kv_len+63)//64
+  live = (valid_kv_len+127)//128
  live = min(live, chunks) if isinstance(live, int) else live.minimum(chunks)
-  out = Tensor.empty(B, H, 1, D, dtype="float32", device=q.device)
-  fxn = functools.partial(_amd_flash_decode_combine, live=live)
-  return Tensor.custom_kernel(out, partial, stats, fxn=fxn)[0]
+  partial, stats = partial[:, :, :live], stats[:, :, :live]
+  weights = ((stats[..., 0]-stats[..., 0].max(2, keepdim=True))*LOG2E).exp2()
+  return ((partial*weights.unsqueeze(-1)).sum(2) / (stats[..., 1]*weights).sum(2, keepdim=True)).unsqueeze(2)

@functools.cache
 def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, valid_kv_len:int|UOp, q_start:int|UOp|None=None) -> UOp:
@@ -121,10 +121,10 @@ class LARS(Optimizer):
        self.b[i].assign(self.momentum * self.b[i] + g)  # NOTE: self.b[i] is zero on the first run, no if required
        g = (g + self.momentum * self.b[i]) if self.nesterov else self.b[i]
      if self.ns_coefficients: g = g.reshape(g.shape[0], -1).newton_schulz(self.ns_steps, self.ns_coefficients).reshape(g.shape)
-      # muon does post momentum weight decay
-      if not self.pre_wd and self.wd > 0: t = t.detach() * (1.0 - self.wd * self.lr)
      # popular momentum does pre learning rate update
      if not self.classic: g = g * r * self.lr
+      # muon does post momentum weight decay
+      if not self.pre_wd and self.wd > 0: g = g + self.wd * self.lr * t.detach()
      ret.append(g.cast(t.dtype))
    return ret, self.b

@@ -30,7 +30,7 @@ class Estimates:
        if u.op in {Ops.INDEX, Ops.SHRINK}:
          excluded = excluded.union(set(UOp.sink(*u.src[1:]).toposort(lambda x: x.op is not Ops.END)))
    for u in uops:
-      if u.op in {Ops.LOAD, Ops.STORE} or (u.op is Ops.CUSTOMI and isinstance(u.arg, str) and "nontemporal_load" in u.arg):
+      if u.op in {Ops.LOAD, Ops.STORE}:
        buf = u
        while len(buf.src) and buf.op is not Ops.PARAM: buf = buf.src[0]
        if buf.op is Ops.PARAM:
@@ -3,6 +3,7 @@ import math, sys, struct
 from collections import defaultdict, Counter
 from tinygrad.codegen.opt import tc
 from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str, axis_letters
+from tinygrad.uop.weak import commit_weak_consts
 from tinygrad.helpers import strip_parens, getenv, prod, dedup, Target, NUM_CPU_THREADS, IMAGE, FLOAT16, is_image_shape
 from tinygrad.dtype import dtypes, DType, AddrSpace, truncate, float_to_bf16
 from tinygrad.renderer import Renderer
@@ -75,6 +76,8 @@ base_rewrite = PatternMatcher([

 def create_non_native_float_pats(dts:tuple[DType, ...], casting:bool=True):
  patterns = PatternMatcher([
+    # a weak CONST states no width and cannot be restated: commit it at the emulated dtype a sibling src states
+    (UPat(GroupOp.ALU, name="x"), lambda x, dts=dts: commit_weak_consts(x, next((s.dtype for s in x.src if s.dtype in dts), None))),
    (UPat(Ops.WHERE, dtype=dts, src=(UPat.var("b"), UPat.var("x"), UPat.var("y")), name="w"),
     lambda w,b,x,y: b.where(x.cast(dtypes.float), y.cast(dtypes.float)).cast(w.dtype)),
    (UPat(GroupOp.ALU-{Ops.WHERE}, dtype=dts, name="x"),
@@ -524,8 +527,6 @@ class HIPRenderer(CStyleLanguage):
    (UPat(Ops.WMMA, name="x", dtype=dtypes.float),
      lambda x: x.replace(src=(x.src[0].bitcast(dtypes.uint64), x.src[1].bitcast(dtypes.uint64), x.src[2]))
      if x.src[0].max_numel() == 8 and x.src[0].dtype in dtypes.fp8_ocp else None),
-    # bfloat16 constant casting
-    (UPat.cvar('x', dtypes.bfloat16), lambda x: cast_float_to_bf16(UOp.const(x.val, dtypes.float))),
  ])

  def asm(self, prg:UOp, lin:UOp) -> bytes:
@@ -143,16 +143,16 @@ extra_matcher = PatternMatcher([
  (UPat.var("m").where(UPat.var("a", (dtypes.bool,)+dtypes.int8s), UPat.var("b")),
   lambda m,a,b: m.where(a.cast(dtypes.int16), b.cast(dtypes.int16)).cast(a.dtype) if a.max_numel() == 1 else None),
  # float16 alus are done in float32
-  (UPat(GroupOp.ALU, dtypes.float16, name="x"), lambda x: UOp(x.op, dtypes.float,
-   tuple(s.cast(dtypes.float) if s.dtype != dtypes.bool else s for s in x.src)).cast(x.dtype)),
+  (UPat(GroupOp.ALU, dtypes.float16, name="x"), lambda x: UOp(x.op,
+   src=tuple(s.cast(dtypes.float) if s.dtype != dtypes.bool else s for s in x.src)).cast(x.dtype)),
  (UPat(GroupOp.Comparison, src=(UPat.var("a", dtypes.float16), UPat.var("b")), name="x"),
   lambda x,a,b: UOp(x.op, src=(a.cast(dtypes.float32), b.cast(dtypes.float32))).cast(x.dtype)),
  # no cmpne for packed ints, y != x => !(y==x)
  (UPat(Ops.CMPNE, src=(UPat.var("y", dtypes.ints), UPat.var("x")), name="cmp"),
   lambda y,x,cmp: UOp(Ops.CMPEQ, src=(y,x))^True if y.max_numel() > 1 else None),
-  # float where expects a mask
-  (UPat.var("m", dtypes.bool).where(UPat.var("a", dtypes.floats), UPat.var("b")),
-   lambda m,a,b: m.cast(a.dtype).ne(0).where(a, b) if m.src[0].dtype not in dtypes.floats else None),
+  # float WHERE needs a mask unless its comparison already has a float operand
+  (UPat.var("m", dtypes.bool).where(UPat.var("a", dtypes.floats+(dtypes.weakfloat,)), UPat.var("b")).named("w"),
+   lambda m,a,b,w: m.cast(w.dtype).ne(0).where(a, b) if w.dtype in dtypes.floats and not dtypes.is_float(m.src[0].dtype) else None),
  # rewrite -x -> 0 - x
  (UPat(Ops.NEG, name="x"), lambda x: UOp(Ops.SUB, src=(x.const_like(0),) + x.src)),
  # TODO: add support for mod, requires support for accessing the 2nd+ reg of a multi output instruction
@@ -653,7 +653,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
    # 0b10 -- signals memory access with 32bit displacement
    # 0b11 -- signals no memory access
    if disp_uop is not None:
-      assert disp_uop.op is Ops.CAST, "displacement must be a literal"
+      assert disp_uop.op is Ops.CAST, "displacement must be a const"
      assert disp_uop.dtype in (dtypes.int8, dtypes.int32), "displacement can only be 1 or 4 byte signed int"
      # rbp/r13 always require a displacement
      if disp_uop.src[0].val != 0 or rm == 0b101: mod = 0b01 if disp_uop.dtype.itemsize == 1 else 0b10
@@ -836,7 +836,7 @@ class X86Renderer(ISARenderer):
    if x.op is Ops.BUFFER: x = x.replace(dtype=dtypes.uint64)
    is_xmm = isinstance(x.tag, tuple) and x.tag[0].cons[0].size == 16
    op = X86Ops.VMOVUPSm if is_xmm else X86Ops.MOVm
-    return UOp(Ops.INS, dtypes.void, fold_address(self.stack_pointer().index(disp)) + (x,), op, x.tag)
+    return UOp(Ops.INS, src=fold_address(self.stack_pointer().index(disp)) + (x,), arg=op, tag=x.tag)

  def fill(self, disp:UOp, x:UOp, reg:Register) -> UOp:
    is_xmm = reg.cons[0].size == 16
@@ -121,8 +121,6 @@ class NIRRenderer(Renderer):
  code_for_op = {**{k:lambda:None for k in u_aop.keys()}, **{k:lambda:None for k in s_aop.keys()}, **{k:lambda:None for k in f_aop.keys()}}

  extra_matcher = PatternMatcher([
-    # handle negative unsigned CONST
-    (UPat.cvar("x", dtypes.uints), lambda x: UOp.const(x.dtype.max+x.val+1, x.dtype) if x.val < 0 else None),
    # from ptx
    (UPat.var('x', dtype=dtypes.bool)<UPat.var('y'), lambda x,y: (x^True)&y),
    # load/store bool -> uint8
@@ -136,9 +134,10 @@ class NIRRenderer(Renderer):
    # ref: https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#OpConvertFToU
    (UPat(Ops.CAST, (dtypes.uchar, dtypes.ushort), src=(UPat.var("x", dtypes.floats),), name="c"), lambda x,c: x.cast(dtypes.int32).cast(c.dtype)),
    # load/store use pointer arithmetic, and the cast does nothing. NOTE: this doesn't apply to image indexing cause it's 1-D
+    # nor to REG/ALU register picks, which keep their own index dtype
    (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"), UPat.var("off")), allow_any_len=True, name="x"), lambda x,buf,off: x.replace(
      src=(buf,UOp.const(off.val, dtypes.long) if off.op is Ops.CONST else off.cast(dtypes.long))+x.src[2:])
-      if buf.addrspace != AddrSpace.REG and not is_image_shape(buf._shape) else None),
+      if buf.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL) and not is_image_shape(buf._shape) else None),
    # images need index to be int for nir (coordinates only: the INDEX keeps its access dtype)
    (UPat.var("buf").index(UPat.var("idx_y"), UPat.var("idx_x"), name="x"),
     lambda x,buf,idx_y,idx_x: x.replace(src=(buf, idx_y.cast(dtypes.int), idx_x.cast(dtypes.int)))),
@@ -52,8 +52,8 @@ ptx_matcher = PatternMatcher([
  (UPat(Ops.STORE, src=(UPat(name="idx"), UPat(dtype=dtypes.bool)), name="x", allow_any_len=True),
   lambda x,idx: UOp(x.op, src=(x.src[0], x.src[1].cast(dtypes.uint8))+x.src[2:]) if idx.addrspace != AddrSpace.REG else None),
  # ptx shr and shl instructions require y to be uint
-  (UPat.var("x") << UPat.var("y"), lambda x,y: UOp(Ops.SHL, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
-  (UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
+  (UPat.var("x") << UPat.var("y"), lambda x,y: UOp(Ops.SHL, src=(x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
+  (UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, src=(x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
 ])

 def mem_type(x:UOp) -> str: return 'shared' if x.addrspace == AddrSpace.LOCAL else 'global'
@@ -18,19 +18,19 @@ def packed_store(s:UOp):
  # bool does its mask math at int32: renderer rewrites run after weak dtypes are lowered, and bool & 0xFF would create a weakint const
  if var.dtype == dtypes.bool: var = var.cast(dtypes.int32)
  new_v, wmask = (var & mask).cast(dtypes.uint32) << shift_am, ((mask << shift_am) ^ 0xFFFFFFFF).cast(dtypes.uint32)
-  buf = idx.load(*((UOp.const(0, dtypes.uint32), *gate) if gate else ()), dtype=dtypes.uint32)
+  buf = idx.cast(dtypes.uint32).load(*((UOp.const(0, dtypes.uint32), *gate) if gate else ()))
  return idx.store((buf & wmask) | new_v, *gate)

 # load for char: sign_extend(buf[idx/4] >> ((idx%4)*8))
 def packed_load(root:UOp):
  bidx, *alt = root.src
  idx, shift_am, mask = packed_field(bidx, dtype:=root.dtype)
-  load = idx.load(*((alt[0].cast(dtypes.uint32), *alt[1:]) if alt else ()), dtype=dtypes.uint32, arg=root.arg)
+  load = idx.cast(dtypes.uint32).load(*((alt[0].cast(dtypes.uint32), *alt[1:]) if alt else ()), arg=root.arg)
  val = (load >> shift_am) & mask
  return sign_extend(val, 8*dtype.itemsize).cast(dtype) if dtype in [dtypes.char, dtypes.short] else val.cast(dtype)

 def is_packed(x:UOp):
-  dt = x.src[1].dtype if x.op is Ops.STORE else x.dtype
+  dt = x.src[1].dtype if x.op is Ops.STORE else x.buf_uop.dtype
  return dt.itemsize < 4 and dt != dtypes.half and x.buf_uop.addrspace != AddrSpace.REG
 def _packed_size(u:UOp): return ceildiv(u.max_numel(), 4//u.dtype.itemsize) if is_packed(u) else u.max_numel()
 def is_nan(a):
@@ -38,15 +38,15 @@ def is_nan(a):
  return (a.bitcast(getattr(dtypes, f"uint{bs}")) & ((1 << (bs - 1)) - 1)) > (((1 << exp) - 1) << mant)

 # the read-modify-write packed_store emits: a load of the very index being stored to, masked (a gated store loads with 3 srcs)
-packed_rmw = UPat(Ops.LOAD, src=(UPat.var("b"),), allow_any_len=True) & UPat.var("wmask")
+packed_rmw = UPat(Ops.LOAD, src=(UPat(Ops.CAST, dtype=dtypes.uint32, src=(UPat.var("b"),)),), allow_any_len=True) & UPat.var("wmask")

 wgsl_matcher = PatternMatcher([
  (UPat((Ops.CMPLT, Ops.XOR), src=(UPat(name="a", dtype=dtypes.bool), UPat.var("b")), name="c"),
   lambda a,b,c: a.cast(dtypes.int).alu(c.op, b.cast(dtypes.int)).cast(dtypes.bool)),
-  (UPat(Ops.LOAD, name="l"), lambda l: packed_load(l) if is_packed(l) else None),
+  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX),), allow_any_len=True, name="l"), lambda l: packed_load(l) if is_packed(l) else None),
  (UPat(Ops.STORE, name="s"), lambda s: packed_store(s) if is_packed(s) else None),
  (UPat.var("a") << UPat.var("b"),lambda a,b:(a.bitcast(dtypes.uint32)<<b.cast(dtypes.uint32)).bitcast(a.dtype) if b.dtype!=dtypes.uint32 else None),
-  (UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
+  (UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, src=(x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
  # fix nan check: 'a != a -> is_nan()'. the decomp rewrites (a != a).logical_not() to CMPEQ, so match both forms
  (UPat.var("a", dtypes.floats) != UPat.var("a"), is_nan),
  (UPat.var("a", dtypes.floats).alu(Ops.CMPEQ, UPat.var("a")), lambda a: is_nan(a).ne(True)),
@@ -65,11 +65,13 @@ class WGSLRenderer(CStyleLanguage):
              dtypes.char: "i32", dtypes.int32: "i32", dtypes.uint32: "u32", dtypes.bool: "bool", dtypes.half: "f16" }

  string_rewrite = PatternMatcher([
+    (UPat(Ops.CAST, dtype=dtypes.uint32, src=(UPat(Ops.INDEX, name="x"),)), lambda ctx,x: ctx[x] if is_packed(x) else None),
    (UPat(Ops.NEG, dtypes.uints, src=(UPat.var('x'))), lambda ctx,x: f"(0-{ctx[x]})"),
    (UPat.cvar("c").cast(dtypes.bool), lambda c: "true" if c.val else "false"),
    (UPat.cvar("c").cast((dtypes.uchar, dtypes.ushort, dtypes.uint32)),
     lambda c: f"bitcast<u32>({c.val})" if c.val < 0 else f"{c.val&0xFFFFFFFF}u"),
-    (UPat.cvar("c").cast(dtypes.int32, name="x"), lambda ctx,x,c: f"{truncate[x.dtype](c.val)}"),
+    # a negative const must state its type: contextual conversion of a bare abstract int rejects it in a u32 position
+    (UPat.cvar("c").cast(dtypes.int32, name="x"), lambda ctx,x,c: f"i32({v})" if (v:=truncate[x.dtype](c.val)) < 0 else f"{v}"),
    (UPat(Ops.BUFFER, name="x"), lambda ctx,x:
     f"var{'<workgroup>' if x.addrspace == AddrSpace.LOCAL else ''} {ctx[x]}: array<{ctx.buf_map(x)},{_packed_size(x)}>;"),
    (UPat(Ops.BITCAST, dtype=dtypes.half, name="x", src=(UPat(dtype=(dtypes.short, dtypes.ushort, dtypes.uint32),),)),
@@ -25,6 +25,7 @@ SQTT = ContextVar("SQTT", abs(VIZ.value)>=2)
 SQTT_ITRACE_SE_MASK, SQTT_LIMIT_SE, SQTT_SIMD_SEL, SQTT_TOKEN_EXCLUDE = \
  ContextVar("SQTT_ITRACE_SE_MASK", 0b11), ContextVar("SQTT_LIMIT_SE", 0), ContextVar("SQTT_SIMD_SEL", 0), ContextVar("SQTT_TOKEN_EXCLUDE", 0)
 PMC = ContextVar("PMC", abs(VIZ.value)>=2)
+AMD_USB_POLL_US = getenv("AMD_USB_POLL_US", 500) # microseconds to sleep between USB signal polls. 0 disables
 EVENT_INDEX_PARTIAL_FLUSH = 4 # based on a comment in nvd.h
 WAIT_REG_MEM_FUNCTION_EQ  = 3 # ==
 WAIT_REG_MEM_FUNCTION_NEQ = 4 # !=
@@ -45,8 +46,9 @@ class AMDSignal(HCQSignal):
  def __init__(self, *args, **kwargs): super().__init__(*args, **{**kwargs, 'timestamp_divider': 100})

  def _sleep(self, time_spent_since_last_sleep_ms:int):
-    # Reasonable to sleep for long workloads (which take more than 200ms) and only timeline signals.
-    if time_spent_since_last_sleep_ms > 200 and self.owner is not None: self.owner.iface.sleep(200)
+    # USB signals live in VRAM across the link, so yield between polls. Native AMD only blocks after 200 ms.
+    if self.owner is not None and self.owner.is_usb() and AMD_USB_POLL_US: time.sleep(AMD_USB_POLL_US / 1e6)
+    elif time_spent_since_last_sleep_ms > 200 and self.owner is not None: self.owner.iface.sleep(200)

 class AMDComputeQueue(HWQueue):
  def __init__(self, dev:AMDDevice):
@@ -61,20 +61,20 @@ class CPUWorker: ring:Buffer; put:Buffer; sem:Buffer; sys:Buffer; done:Buffer; t
 def cpu_cmd(devs:tuple[str, ...], prog, *args:UOp) -> UOp:
  progs = [get_runtime(d, prog) if isinstance(prog, UOp) else cast(CPUDevice, Device[d]).prgs[prog] for d in devs]
  addrs = tuple(UOp.const(p.addr, dtypes.uint64) for p in progs)
-  words = ((addrs[0] if len(addrs) == 1 else UOp(Ops.STACK, dtypes.uint64, addrs)),) + args
-  return UOp(Ops.INS, dtypes.void, words + (UOp.const(0, dtypes.uint64),) * (CMD_SIZE - len(words)), arg="cmd")
+  words = ((addrs[0] if len(addrs) == 1 else UOp(Ops.STACK, src=addrs)),) + args
+  return UOp(Ops.INS, src=words + (UOp.const(0, dtypes.uint64),) * (CMD_SIZE - len(words)), arg="cmd")

 def cpu_exec(ctx:tuple[str, ...], call:UOp, prg:UOp) -> UOp:
  args = [get_call_arg_uops(call)[i].getaddr(ctx) for i in prg.arg.globals] + [v.cast(dtypes.uint64) for v in get_call_var_uops(call, prg)]
  if (core:=prg.arg.runtimevars.get('core_id')) is None: return cpu_cmd(ctx, prg, *args)

  la = [cpu_cmd(ctx,prg,*args[:(cid:=(len(prg.arg.globals)+core))],UOp.const(t, dtypes.uint64),*args[cid+1:]) for t in range(prg.arg.global_size[0])]
-  return UOp(Ops.LINEAR, dtypes.void, tuple(la))
+  return UOp(Ops.LINEAR, src=tuple(la))

 pm_cpu_opsel = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), cpu_exec),

-  (UPat(Ops.INS, arg="barrier"), lambda: UOp(Ops.NOOP, dtypes.void, ())),
+  (UPat(Ops.INS, arg="barrier"), lambda: UOp(Ops.NOOP)),
  (UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val"))),
   lambda ctx, dst, val: cpu_cmd(ctx, wait_prog, dst.getaddr(ctx), val.cast(dtypes.uint64))),
  (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))),
@@ -9,18 +9,10 @@ from tinygrad.renderer.cstyle import ClangRenderer
 from tinygrad.runtime.autogen import libc, qcom_dsp
 if getenv("IOCTL"): import extra.dsp.run # noqa: F401 # pylint: disable=unused-import

-from tinygrad.uop.ops import PatternMatcher, UPat
-
-# NOTE: this just increases readability of the generated code
-dsp_string = PatternMatcher([
-  (UPat(Ops.CONST, (dtypes.int8, dtypes.uint8), name="x"), lambda ctx,x: str(x.val)),
-])
-
 class DSPRenderer(ClangRenderer):
  has_threads = False
  buffer_suffix = " restrict __attribute__((align_value(128)))"
  kernel_typedef = "__attribute__((noinline)) void"
-  string_rewrite = dsp_string+ClangRenderer.string_rewrite
  type_map = { **ClangRenderer.type_map, dtypes.uint64: "unsigned long long", dtypes.int64: "long long" }
  code_for_op = {k:v for k,v in ClangRenderer.code_for_op.items() if k != Ops.SQRT}

@@ -105,7 +97,7 @@ class DSPAllocator(Allocator['DSPDevice']):

 class DSPCompiler(Compiler):
  def __init__(self, mock:bool=False):
-    compiler_args = "--target=hexagon -mcpu=hexagonv65 -fuse-ld=lld -nostdlib -mhvx=v65 -mhvx-length=128b"
+    self.mock, compiler_args = mock, "--target=hexagon -mcpu=hexagonv65 -fuse-ld=lld -nostdlib -mhvx=v65 -mhvx-length=128b"
    if mock: self.args = f"-static {compiler_args}"
    else:
      # Generate link script to pass into clang. Aligning all used sections to 4k fixes invoke problem.
@@ -120,6 +112,9 @@ class DSPCompiler(Compiler):

    super().__init__(None if mock else "compile_dsp")

+  def __del__(self):
+    if not self.mock: os.unlink(self.link_ld.name)
+
  def compile(self, src:str) -> bytes:
    # TODO: remove file write. sadly clang doesn't like the use of /dev/stdout here
    with tempfile.NamedTemporaryFile(delete=True) as f:
@@ -145,8 +145,18 @@ class AMMemoryManager(MemoryManager):
 class AMDev:
  Version = 0xA0000008

+  def _disable_aspm(self):
+    # L1 across retimers makes reads oscillate to 0xffffffff; power on defaults it enabled. Clearing the GPU endpoint
+    # alone suffices: L1 only engages when both ends of the link enable it.
+    cap, seen = self.pci_dev.read_config(0x34, 1) & 0xfc, set() # bound the walk: a dead link can return 0xff pointers forever
+    while cap and cap not in seen and self.pci_dev.read_config(cap, 1) != 0x10:
+      seen.add(cap)
+      cap = self.pci_dev.read_config(cap + 1, 1) & 0xfc
+    if cap and cap not in seen: self.pci_dev.write_config_flush(cap + 0x10, self.pci_dev.read_config(cap + 0x10, 2) & ~3, 2) # PCIe cap lnkctl
+
  def __init__(self, pci_dev:PCIDevice, reset_mode=False):
    self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
+    self._disable_aspm()
    self.vram, self.doorbell64, self.mmio = self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I')

    self._run_discovery()
@@ -170,6 +180,9 @@ class AMDev:
      if DEBUG >= 2: print(f"am {self.devfmt}: Malformed state. Issuing a full reset.")
      self.partial_boot = False

+    # aqua (gc 9.5.0): full boot over live state can kill the fabric (power cycle recovers); partial boot+reset_mec is the deepest safe reset
+    if self.ip_ver[am.GC_HWIP] == (9,5,0) and self.reg("regSCRATCH_REG7").read() == AMDev.Version: self.partial_boot = True
+
    # Init hw for IP blocks where it is needed
    if not self.partial_boot:
      if self.psp.is_sos_alive() and self.smu.is_smu_alive():
@@ -342,3 +355,9 @@ class AMDev:
    for prefix, hwip in mods:
      self.__dict__.update(import_asic_regs(prefix, self.ip_ver[hwip], cls=functools.partial(AMRegister, adev=self, bases=self.regs_offset[hwip])))
    self.__dict__.update(import_asic_regs('mp', (11, 0, 0), cls=functools.partial(AMRegister, adev=self, bases=self.regs_offset[am.MP1_HWIP])))
+
+    # Live AIDs like the kernel: 4 SDMAs per AID; the AID lives iff its group's alive-mask is 0xf/0x3/0xc.
+    # Dead AIDs must never be touched via the indirect window: writes poison the whole fabric.
+    live_sdma = {k for k in self.regs_offset[am.SDMA0_HWIP] if k not in self.harvested[am.SDMA0_HWIP]}
+    max_aid = max((k >> 2 for k in self.regs_offset[am.SDMA0_HWIP]), default=0)
+    self.aids = [0] + [aid for aid in range(1, max_aid + 1) if sum(1 << (i & 3) for i in live_sdma if i >> 2 == aid) in {0xf, 0x3, 0xc}]
@@ -1,5 +1,5 @@
 import ctypes, time, contextlib, functools
-from typing import Literal
+from typing import Iterable, Literal
 from tinygrad.helpers import to_mv, data64, lo32, hi32, DEBUG, wait_cond, pad_bytes, getbits
 from tinygrad.runtime.autogen.am import am
 from tinygrad.runtime.support.amd import import_soc
@@ -32,7 +32,7 @@ class AM_SOC(AM_IP):
      # fence doorbells for harvested xccs (0xff & ~xcc_mask in the kernel); a fully-unharvested chip keeps the previous 0x0
      live_xccs = sum(1 << i for i in self.adev.regs_offset[am.GC_HWIP] if i not in self.adev.harvested[am.GC_HWIP] and i < 8)
      self.adev.regXCC_DOORBELL_FENCE.write(0xff & ~live_xccs)
-      for aid in range(1, self.adev.gmc.vmhubs):
+      for aid in self.adev.aids[1:]:
        self.adev.indirect_wreg_pcie(self.adev.regXCC_DOORBELL_FENCE.addr[0], self.adev.regXCC_DOORBELL_FENCE.encode(shub_slv_mode=1), aid=aid)
      self.adev.regBIFC_GFX_INT_MONITOR_MASK.write(0x7ff)
      self.adev.regBIFC_DOORBELL_ACCESS_EN_PF.write(0xfffff)
@@ -53,10 +53,9 @@ class AM_GMC(AM_IP):
  def init_sw(self):
    self.vmhubs = len(self.adev.regs_offset[am.MMHUB_HWIP])

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

    self.paddr_base = self.xgmi_phys_id * self.xgmi_seg_sz

@@ -81,9 +80,11 @@ class AM_GMC(AM_IP):
    # MM hub is inited before any tlb flushes and is still valid during partial_boot, so set it to true
    self.hub_initted = {"MM": True, "GC": False}

+    self.mm_insts = self.adev.aids if self.adev.ip_ver[am.NBIO_HWIP] in {(7,9,0), (7,9,1)} else list(range(self.vmhubs)) # dead mmhubs hang us
+
    self.pf_status_reg = lambda ip: f"reg{ip}VM_L2_PROTECTION_FAULT_STATUS{'_LO32' if self.adev.ip_ver[am.GC_HWIP] >= (12,0,0) else ''}"

-  def init_hw(self): self.init_hub("MM", inst_cnt=self.vmhubs)
+  def init_hw(self): self.init_hub("MM", insts=self.mm_insts)

  def flush_hdp(self): self.adev.wreg(self.adev.reg("regBIF_BX0_REMAP_HDP_MEM_FLUSH_CNTL").read() // 4, 0x0)
  def flush_tlb(self, ip:Literal["MM", "GC"], vmid, flush_type=0):
@@ -92,7 +93,7 @@ class AM_GMC(AM_IP):
    # Can't issue TLB invalidation if the hub isn't initialized.
    if not self.hub_initted[ip]: return

-    for inst in range(self.adev.gmc.vmhubs if ip == "MM" else self.adev.gfx.xccs):
+    for inst in (self.adev.gmc.mm_insts if ip == "MM" else range(self.adev.gfx.xccs)):
      if ip == "MM": wait_cond(lambda: self.adev.regMMVM_INVALIDATE_ENG17_SEM.read(inst=inst) & 0x1, value=1, msg="mm flush_tlb timeout")

      self.adev.reg(f"reg{ip}VM_INVALIDATE_ENG17_REQ").write(flush_type=flush_type, per_vmid_invalidate_req=(1 << vmid), invalidate_l2_ptes=1,
@@ -117,9 +118,9 @@ class AM_GMC(AM_IP):
    self.adev.reg(f"reg{ip}VM_CONTEXT{vmid}_CNTL").write(0x1800000, **fault_flags, **en_def_flags, enable_context=1,
      page_table_depth=((2 if self.trans_futher else 3) - page_table.lv), page_table_block_size=9 if self.trans_futher else 0, inst=inst)

-  def init_hub(self, ip:Literal["MM", "GC"], inst_cnt:int):
+  def init_hub(self, ip:Literal["MM", "GC"], insts:Iterable[int]):
    # Init system apertures
-    for inst in range(inst_cnt):
+    for inst in insts:
      self.adev.reg(f"reg{ip}MC_VM_AGP_BASE").write(0, inst=inst)
      self.adev.reg(f"reg{ip}MC_VM_AGP_BOT").write(0xffffffffffff >> 24, inst=inst) # disable AGP
      self.adev.reg(f"reg{ip}MC_VM_AGP_TOP").write(0, inst=inst)
@@ -258,7 +259,7 @@ class AM_GFX(AM_IP):
    wait_cond(lambda: self.adev.regCP_STAT.read() == 0 or self.adev.regRLC_RLCS_BOOTLOAD_STATUS.read_bitfields()['bootload_complete'] == 0,
              value=True, msg="RLC autoload timeout")

-    self.adev.gmc.init_hub("GC", inst_cnt=self.xccs)
+    self.adev.gmc.init_hub("GC", insts=range(self.xccs))
    if self.adev.partial_boot: return self.reset_mec()

    self._config_mec()
@@ -406,7 +407,11 @@ class AM_GFX(AM_IP):
        if self.adev.regCP_HQD_ACTIVE.read(inst=xcc) & 1:
          self.adev.regCP_HQD_DEQUEUE_REQUEST.write(0x2, inst=xcc) # 1 - DRAIN_PIPE; 2 - RESET_WAVES
          self.adev.regSPI_COMPUTE_QUEUE_RESET.write(0x1, inst=xcc)
-          if not self.adev.is_err_state: wait_cond(lambda: self.adev.regCP_HQD_ACTIVE.read(inst=xcc) & 1, value=0, msg="HQD dequeue timeout")
+          if not self.adev.is_err_state:
+            try: wait_cond(lambda: self.adev.regCP_HQD_ACTIVE.read(inst=xcc) & 1, value=0, msg="HQD dequeue timeout")
+            # kernel tolerates this too; a wedged wave can survive RESET_WAVES
+            except TimeoutError:
+              if DEBUG >= 2: print(f"am {self.adev.devfmt}: HQD dequeue timeout xcc{xcc} q{q}, continuing")
    self._grbm_select()

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

    if self.adev.ip_ver[am.NBIO_HWIP] in {(7,9,0), (7,9,1)}:
-      for aid_id in range(self.adev.gmc.vmhubs):
+      for aid_id in self.adev.aids:
        for dev_inst, (port, awid, offset, awaddr) in enumerate([(1, 0xe, 0xe, 0x1), (2, 0x8, 0x8, 0x2), (5, 0x9, 0x9, 0x8), (6, 0xa, 0xa, 0x9)]):
          entry = dev_inst + 1 + 4 * aid_id
          self.adev.reg(f"regDOORBELL0_CTRL_ENTRY_{entry}").write(**{f"bif_doorbell{entry}_range_size_entry": 20,
@@ -32,7 +32,8 @@ class HCQInfo:

  input_idxs:tuple[tuple[tuple[str, ...], tuple[int, ...]], ...] = () # per inputs table: (devices, indexes into input_uops)
  inputs:int|None = None # index of the inputs table in call.src
-  kernels:tuple[tuple[tuple[str, ...], UOp, tuple[int, ...]], ...] = () # per kernel: (devices, a call carrying its name and estimates, timestamps)
+  # per kernel: (devices, name, estimates, timestamps, profile key)
+  kernels:tuple[tuple[tuple[str, ...], str, Estimates, tuple[int, ...], bytes], ...] = ()

 def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c

@@ -62,8 +63,8 @@ def make_patches(buf:UOp, patches:Sequence[tuple[sint, UOp]]) -> tuple[UOp, ...]
  ret, bit = [], buf.dtype.itemsize
  for (tag, dt, r), ps in groups.items():
    view = buf.shrink(((r // bit, (max(off for off,_ in ps) + dt.itemsize) // bit),)).bitcast(dt)
-    offs = UOp(Ops.STACK, dtypes.int, tuple(UOp.const((off - r) // dt.itemsize, dtypes.int) for off,_ in ps))
-    ret.append(view.index(offs).store(UOp(Ops.STACK, dt, tuple(val for _,val in ps))).rtag(tag))
+    offs = UOp(Ops.STACK, src=tuple(UOp.const((off - r) // dt.itemsize, dtypes.int) for off,_ in ps))
+    ret.append(view.index(offs).store(UOp(Ops.STACK, src=tuple(val for _,val in ps))).rtag(tag))
  return tuple(ret)

 def make_binary_patch(buf:UOp, blob:bytes) -> UOp: return buf.store(UOp(Ops.BINARY, src=(), arg=blob).bitcast(buf.dtype)).rtag("link")
@@ -71,8 +72,8 @@ def make_binary_patch(buf:UOp, blob:bytes) -> UOp: return buf.store(UOp(Ops.BINA
 def make_cmdbuf(lin, devs, buf:UOp|None=None, dep:tuple[UOp, ...]=()):
  blob, patches = bytearray(), []
  for s in (s for ins in lin.src for s in ins.src):
-    if s.op is not Ops.CONST: patches.append((len(blob), s))
-    blob.extend(struct.pack(f'<{s.dtype.fmt}', s.val if s.op is Ops.CONST else 0x0))
+    if not (is_const:=(s.op is Ops.CAST and s.src[0].op is Ops.CONST)): patches.append((len(blob), s))
+    blob.extend(struct.pack(f'<{s.dtype.fmt}', s.val if is_const else 0x0))
  cmdbuf = buf if buf is not None else UOp.placeholder((len(blob) // 4,), dtypes.uint32, next(UOp.unique_num), device=devs).rtag("cmdbuf")
  return cmdbuf.after(*dep, make_binary_patch(cmdbuf, bytes(blob)), *make_patches(cmdbuf, patches))

@@ -271,7 +272,7 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> li
    # and make hcq call
    name, info = get_call_name(call, get_call_arg_uops(call)), HCQInfo(devices, estimate_uop(call))
    ts_ids = [next(UOp.unique_num) for _ in range(2)] if profile else []
-    kerns.append((devices, make_call(name, call.src[0], info), tuple(ts_ids)))
+    kerns.append((devices, name, info.estimates, tuple(ts_ids), make_call(name, call.src[0], info).key))

    ts_ins = [UOp(Ops.INS, arg="timestamp", src=(make_buf(devices, s),)) for s in ts_ids]
    q += ts_ins[:1] + [call.replace(arg=replace(call.arg, aux=info))] + ts_ins[1:]
@@ -344,7 +345,9 @@ def make_gather_loop(patches:list[UOp], table:UOp, slots:dict[UOp, int], lt_patc
  lt_patches.append(make_binary_patch(pairs, struct.pack(f'<{2*len(words)}I', *itertools.chain(*words))))
  r = UOp.range(len(words), next(UOp.unique_num), dtype=dtypes.int, src=(pairs, dst))
  off, slot = ((pairs.index(2*r+i).load() % bound).cast(dtypes.int) for i, bound in ((0, dst.max_numel()-1), (1, table.max_numel())))
-  patch = dst.shrink(((off, off+table.dtype.itemsize//dst.dtype.itemsize),)).bitcast(table.dtype).index(0).store(table.index(slot).load()).end(r)
+  # SHRINK(offset, length): a const length keeps the end bound from becoming an expression the program spec rejects
+  patch = UOp(Ops.SHRINK, src=(dst, off, off.const_like(table.dtype.itemsize//dst.dtype.itemsize))).bitcast(table.dtype).index(0) \
+    .store(table.index(slot).load()).end(r)
  return {p: UOp(Ops.NOOP) for p in patches} | {patches[0]: patch}

 def is_input_addr(g:UOp) -> bool: return all(x.op is Ops.PARAM and x.tag is None for x in unwrap_mstack(g.buf_uop))
@@ -500,7 +503,7 @@ pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, name="buf"), bufferize_buf)])
 # 7. resolve patches

 def push_stack(op, s): return UOp(Ops.STACK,
-  src=tuple(op.replace(dtype=op.dtype, src=tuple(x if y is s else y for y in op.src)) for x in s.src))
+  src=tuple(op.replace(src=tuple(x if y is s else y for y in op.src)) for x in s.src))

 def fold_binary(buf:UOp, blob:UOp) -> UOp:
  for b in (m.bufs if isinstance(m:=buf.buffer, MultiBuffer) else (m,)):
@@ -511,7 +514,7 @@ def fold_const_store(view:UOp, off:UOp, val:UOp) -> UOp:
  buf, start = unwrap_view(view)
  for off,val in zip(off.src, val.src):
    for b,v in zip((bs:=mb.bufs if isinstance((mb:=buf.buffer), MultiBuffer) else (mb,)), val.src if val.op is Ops.STACK else (val,)*len(bs)):
-      data = struct.pack(f'<{v.dtype.fmt}', truncate[v.dtype]((v.src[0] if v.op is Ops.CAST else v).val))
+      data = struct.pack(f'<{v.dtype.fmt}', truncate[v.dtype](v.val))
      bo = start*buf.dtype.itemsize + off.val*val.dtype.itemsize
      b.ensure_allocated()._buf.cpu_view().view(fmt='B')[bo:bo+len(data)] = data
  return UOp(Ops.NOOP)
@@ -297,7 +297,7 @@ def usb_stage_copy(dst:UOp, src:UOp) -> UOp|None:
              sram.copy_to_device(d.device).call(d, sram)]
    else:
      pad = UOp.new_buffer("CPU", round_up(nb, 512), dtypes.uint8)[0:nb]
-      submit = make_submit(UOp(Ops.CALL, dtypes.void, (UOp(Ops.COPY, dtypes.void, ()), sram, s)), devs=devs, queue="COPY:0")
+      submit = make_submit(s.copy_to_device(sram.device).call(sram, s), devs=devs, queue="COPY:0")
      pull = usb_bulk(devs, (submit,), 0x81, pad.getaddr((HCQ_RUNTIME_DEV.value,)), round_up(nb, 512), 10000)
      ops += [UOp.custom_function("hcq", pull.sink()).call(pad, sram, s, name="hcq_copyout", aux=HCQInfo(devs)),
              pad.copy_to_device("CPU").call(d, pad)]
@@ -80,6 +80,7 @@ def create_schedule(sched_sink:UOp) -> UOp:

 from tinygrad.schedule.memory import memory_plan_rewrite
 from tinygrad.engine.realize import capturing, pm_flatten_linear
+from tinygrad.schedule.prepare import prepare_rangeify
 from tinygrad.schedule.rangeify import get_kernel_graph
 from tinygrad.helpers import CAPTURING
 from tinygrad.uop.ops import PatternMatcher, UPat, ParamArg
@@ -123,7 +124,7 @@ def lower_sink_to_linear(function:UOp) -> UOp|None:
  if not SCACHE or (sc_ret:=schedule_cache.get(cache_key, None)) is None:
    if SPEC: type_verify(function, spec_tensor)
    # support recursive CALLs
-    linear = create_schedule(get_kernel_graph(function))
+    linear = create_schedule(get_kernel_graph(prepare_rangeify(function)))
    if SCACHE: schedule_cache[cache_key] = linear
  else:
    # schedule cache hit
@@ -156,7 +157,7 @@ def simplify_copy_kernel(call:UOp, ast:UOp, dst:UOp, src:UOp):
  # NOTE: this is a codegen for SDMA devices
  if dst.device == src.device and not (isinstance(dst.device, str) and dst.device.startswith("DISK")): return None
  from tinygrad.codegen.simplify import pm_flatten_range, pm_simplify_ranges
-  from tinygrad.schedule.rangeify import pm_mops
+  from tinygrad.schedule.prepare import pm_mops
  from tinygrad.uop.symbolic import sym
  sink = graph_rewrite(ast, sym+pm_mops+pm_flatten_range+pm_simplify_ranges, ctx={}, name="simplify ranges in copy")
  return call.replace(src=(sink,) + call.src[1:])
@@ -259,7 +259,7 @@ def store_dest_multi(root:UOp, multi:UOp):
  # (scalars arrive EXPANDed to the full shape by UOp.store's const_like, so they sub-view like everything else)
  srcs = [multi.src[0]] + [x.src[0] if x.op is Ops.UNSHARD else shard_subview(x, multi) if tuple(x.shape) == tuple(multi.shape) else x
                           for x in root.src[1:]]
-  return UOp(root.op, root.dtype, tuple(srcs), root.arg)
+  return UOp(root.op, src=tuple(srcs), arg=root.arg)

 def passthrough_multi(root:UOp, multi:UOp):
  new_src = (multi.src[0],)+tuple(x.src[0] if x.op is Ops.UNSHARD else x for x in root.src[1:])
@@ -301,15 +301,12 @@ multi_pm = PatternMatcher([

  # resolve TUPLE+GETTUPLE (needed in multi)
  (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]),
-  # GETTUPLE on UNSHARD: passthrough UNSHARD (e.g. when FUNCTION was replaced by UNSHARD(GETTUPLE(...)))
-  (UPat(Ops.GETTUPLE, src=(UPat(Ops.UNSHARD, name="multi"),), name="g"),
-    lambda g, multi: multi.src[0].gettuple(g.arg).unshard(multi.arg, multi.src[1:]) if multi.src[0].op in {Ops.FUNCTION, Ops.TUPLE} else multi),
  # rewrite into FUNCTION calls explicitly for UNSHARD (value-producing)
  (UPat(Ops.FUNCTION, name="call"), rewrite_into_function),
  (UPat((Ops.CALL, Ops.FUNCTION, Ops.AFTER), src=(UPat(Ops.UNSHARD, name="multi"), ), name="root", allow_any_len=True), passthrough_multi),
  # just strip the UNSHARD from non-value-producing CALLs (custom kernels, etc.) — FUNCTION is handled by rewrite_into_function
  (UPat(Ops.CALL, dtype=dtypes.void, name="root", custom_early_reject=set([Ops.UNSHARD])), lambda root:
-    UOp(root.op, root.dtype, tuple(x.src[0] if x.op is Ops.UNSHARD else x for x in root.src), root.arg)),
+    UOp(root.op, src=tuple(x.src[0] if x.op is Ops.UNSHARD else x for x in root.src), arg=root.arg)),
  (UPat((Ops.CAST, Ops.BITCAST, Ops.CONTIGUOUS, Ops.DETACH, Ops.CONTIGUOUS_BACKWARD),
        src=(UPat(Ops.UNSHARD, name="multi"), ), name="root"), passthrough_multi),
  # STORE of a sharded value into an unsharded dest (e.g. a fragment into a full output tile)
@@ -0,0 +1,205 @@
+import itertools
+from tinygrad.dtype import dtypes, to_dtype
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp
+from tinygrad.uop.ops import graph_rewrite, rewrite_group, shape_to_shape_arg, ParamArg, identity_element
+from tinygrad.uop.movement import mop_cleanup
+from tinygrad.helpers import prod, getenv, all_int, DEBUG, SPLIT_REDUCEOP, OPENPILOT_HACKS, FLOAT16, argsort
+from tinygrad.schedule.indexing import apply_movement_op
+from tinygrad.schedule.allreduce import create_allreduce_function
+from tinygrad.schedule.multi import multi_pm
+
+def found_after(ctx:dict[UOp, UOp], after:UOp, src:UOp):
+  if (x:=src).op is Ops.CAST and x.dtype == dtypes.half and FLOAT16: x, after = x.src[0], after.cast(dtypes.float)
+  while True:
+    if x.op is Ops.PERMUTE: x, after = x.src[0], after.permute(argsort(x.marg))
+    elif x.op is Ops.RESHAPE: x, after = x.src[0], after.reshape(x.src[0].shape)
+    elif x.op is Ops.WHERE and x.src[2].base.is_invalid and x.src[1].op is Ops.PAD:
+      x, after = x.src[1].src[0], after.shrink(tuple((o, s+o) for (o,_),s in zip(x.src[1].marg, x.src[1].src[0].shape)))
+    else: break
+  ctx[x] = after
+
+# *** fold moved AFTERs (hack for openpilot) ***
+pm_fold_moved_after = PatternMatcher([
+  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(), UPat((*GroupOp.Movement,Ops.CAST,Ops.WHERE), name="src")))), name="after"), found_after),
+  # replace ALU sources with AFTER versions found above
+  (UPat(GroupOp.ALU, name="alu"), lambda ctx,alu: alu.replace(src=new_src) if (new_src:=tuple(ctx.get(s, s) for s in alu.src)) != alu.src else None),
+])
+
+# movement op on INDEX as a PatternMatcher
+def _mop_index(r:UOp, idx:UOp):
+  idxs = idx.src[1:]
+  if len(idxs) == len(r.shape):
+    return r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idxs), arg=idx.arg)
+  if r.op is Ops.RESHAPE:
+    src_prefix = len(r.src[0].shape) - len(r.shape[len(idxs):])
+    if src_prefix >= 0 and r.src[0].shape[src_prefix:] == r.shape[len(idxs):]:
+      if src_prefix == 0: return r.src[0] if r.src[0].dtype == idx.dtype else None
+      ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape[:src_prefix], r.shape[:len(idxs)], idxs), arg=idx.arg)
+      return ret if ret.shape == idx.shape else None
+
+pm_mops = PatternMatcher([
+  # handle movement ops on INDEX
+  (UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), _mop_index),
+  # move movement ops and INDEX after AFTER
+  (UPat(GroupOp.Movement|{Ops.INDEX}, name="r").after(name="a", allow_any_len=True),
+   lambda r,a: UOp(r.op, src=(a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], arg=r.arg)),
+  (UPat(GroupOp.Movement, name="r").end(name="a", allow_any_len=True), lambda r,a: a.replace(src=(r.src[0],)+a.src[1:])),
+])
+
+# *****************
+# 0. do some cleanup rewrites, mostly copied from the old stuff
+
+def fix_store_hazard(target:UOp, src:UOp):
+  if (base:=target.base) not in src.toposort(enter_calls=False): return None
+  # PERMUTE and FLIP reorder indices, SHRINK can have overlapping regions when dest is also shrunk
+  unsafe = {Ops.PERMUTE, Ops.FLIP} | ({Ops.SHRINK} if target.op_in_backward_slice_with_self(Ops.SHRINK) else set())
+  reaches_base: dict[UOp, bool] = {}
+  for s in src.toposort(gate=lambda s: s.op is not Ops.CONTIGUOUS):
+    reaches_base[s] = s is base or any(reaches_base.get(c) for c in s.src)
+    if reaches_base[s] and s.op in unsafe and not (s is target and s.op is Ops.SHRINK): return target.store(src.contiguous())
+
+def split_reduceop(reduce:UOp, x:UOp):
+  if prod(reduce.shape) == 0: return None
+  if not SPLIT_REDUCEOP or not all_int(x.shape) or (prod(x.shape)//prod(reduce.shape))<getenv("REDUCEOP_SPLIT_THRESHOLD", 32768): return None
+  # if there are few globals, make some reduces into globals by splitting into two kernels
+  # cap output buffer to 2**22: heuristic number of global outputs to achieve max occupancy with enough locals+upcasts for gemm
+  #   ~2**10 should be enough if GROUP is used
+  # 256 split maximum should be "negligible reduce" for low prod(reduce.shape), 8 split minimum.
+  # split is moved to the end to provide maximum locality for the second phase reduce.
+
+  # get expanded by rangeifying the UOp x
+  indexed = x.index(*[UOp.range(s, i) if resolve(s>1) else 0 for i,s in enumerate(x.shape)])
+  range_nums = [y.arg[0] for y in indexed.substitute({x.base:UOp(Ops.NOOP)}, extra_pm=pm_mops).ranges]
+  is_expanded = [i not in range_nums for i in range(len(x.shape))]
+
+  if not (split_candidates:=[(i,d) for i in range(reduce.arg[1])
+                             for d in range(min(256,2**getenv("REDUCEOP_SPLIT_SIZE",22)//prod(reduce.shape)),8-1,-1)
+                             if x.shape[i]%d==0 and not is_expanded[i]]): return None
+  dim_to_split, divisor = split_candidates[0]
+  splitted_shape = x.shape[:dim_to_split]+(divisor,)+(x.shape[dim_to_split]//divisor,)+x.shape[dim_to_split+1:]
+  splitted = x.reshape(splitted_shape).permute(tuple([d for d in range(len(splitted_shape)) if d!=dim_to_split]+[dim_to_split]))
+  if DEBUG >= 3: print(f"split {divisor}: {x.shape} -> {splitted.shape} -> {reduce.shape}")
+  # reduce original axes, then split
+  return splitted._rop(reduce.arg[0], tuple(range(reduce.arg[1]))).contiguous()._rop(reduce.arg[0], (len(reduce.shape),)).reshape(reduce.shape)
+
+pm_gather_params = PatternMatcher([ (UPat(Ops.PARAM, name="p"), lambda ctx, p: ctx.append(p) if p.arg.slot >= 0 else None), ])
+def resolve_function(c:UOp, allow_param_mismatch=True) -> UOp|None:
+  if c.arg.precompile: return None
+  params: list[UOp] = []
+  graph_rewrite(c.src[0], pm_gather_params, bottom_up=True, ctx=params, name="gather params")
+  params = sorted(params, key=lambda x: x.arg.slot)
+  args = c.src[1:]
+
+  # NOTE: this isn't really needed. it's okay if there's unused args in the function
+  if not allow_param_mismatch:
+    if [x.arg.slot for x in params] != list(range(len(params))): raise RuntimeError(f"params not in order: {[x.arg.slot for x in params]}")
+    if len(params) != len(args): raise TypeError(f"expected {len(params)} args, got {len(args)}")
+
+  dict_map = {x:args[x.arg.slot] for x in params}
+  for i, (p, a) in enumerate(dict_map.items()):
+    if p.axis != a.axis: raise TypeError(f"arg {i} axis mismatch: expected {p.axis}, got {a.axis}")
+    if p.max_shape != a.max_shape: raise TypeError(f"arg {i} shape mismatch: expected {p.shape}, got {a.shape}")
+    if p.dtype != a.dtype: raise TypeError(f"arg {i} dtype mismatch: expected {p.dtype}, got {a.dtype}")
+  return c.src[0].substitute(dict_map, walk=True)
+
+# shape-changing bitcast
+def expand_bitcast(bc:UOp) -> UOp|None:
+  x = bc.src[0]
+  if (ns:=bc.dtype.itemsize) == (os:=x.dtype.itemsize) or (isinstance(x.device, str) and x.device.startswith(("DISK", "TINYFS"))): return None
+  new_uint, tmp = to_dtype(f"uint{8*ns}"), x.bitcast(to_dtype(f"uint{8*os}"))
+  if ns > os:
+    tmp = tmp.reshape(x.shape[:-1] + (x.shape[-1]//(rate := ns//os), rate))
+    parts = [tmp.shrink((None,)*(len(tmp.shape)-1) + ((i, i+1),)).cast(new_uint)<<8*i*os for i in range(rate)]
+    return parts[0].usum(*parts[1:]).squeeze(-1).bitcast(bc.dtype)
+  parts = [tmp>>8*i*ns for i in range(os//ns)]
+  return parts[0].stack(*parts[1:], dim=-1).flatten(-2).cast(new_uint).bitcast(bc.dtype)
+
+earliest_rewrites = mop_cleanup+PatternMatcher([
+  # resolve FUNCTION calls (inline the body)
+  (UPat(Ops.FUNCTION, name="c"), resolve_function),
+
+  # resolve TUPLE+GETTUPLE
+  (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]),
+
+  # resolve allreduce (must be bottom up)
+  (UPat(Ops.ALLREDUCE, src=(UPat.var("buf"),), name="red"), create_allreduce_function),
+
+  # split_reduceop
+  (UPat(Ops.REDUCE, name="reduce", src=(UPat.var("x"),)), split_reduceop),
+
+  # remove DETACH/CONTIGUOUS_BACKWARD (TODO: this is copied in allocations)
+  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),
+
+  # SINK only ever references the base
+  (UPat(Ops.SINK, name="x"), lambda x: x.replace(src=tuple(y.unsharded_base for y in x.src))),
+
+  # ** copy rules **
+
+  # COPY transfers a contiguous range, so materialize a source that's resized (shrink/pad/expand) or reordered (permute/flip)
+  (UPat(Ops.COPY, src=(UPat(GroupOp.Movement, name="r"),), name="c"),
+   lambda c,r: c.replace(src=(r.contiguous(),)) if resolve(r.numel() != r.base.numel(), False) or r.contiguous_view_offset() is None else None),
+
+  # copy to same device is a no-op
+  (UPat(Ops.COPY, src=(UPat.var("x"),), name="copy"), lambda x,copy: x if x.device == copy.device else None),
+
+  # copy on reshape is reshape on copy
+  (UPat(Ops.COPY, src=(UPat(Ops.RESHAPE, name="shp"),), name="cpy"), lambda shp,cpy: shp.src[0].copy_to_device(cpy.device).reshape(shp.shape)),
+
+  # reshaping on STORE can be a NOOP
+  (UPat(Ops.STORE, src=(UPat(Ops.RESHAPE, src=(UPat.var("dst",),), allow_any_len=True),
+                        UPat(Ops.RESHAPE, src=(UPat.var("src",),), allow_any_len=True))),
+   lambda dst,src: dst.store(src) if dst.shape == src.shape else None),
+
+  # ** store rules **
+
+  # fix store hazard (dest is in used in src) by adding contiguous: TestAssign.test_post_flipped_assignment
+  (UPat(Ops.STORE, src=(UPat(name="target"), UPat(name="src"))), fix_store_hazard),
+
+  # remove two STOREs that store the same thing to the same place: TestSchedule.test_dedup_Assign
+  (UPat.var("buf").after(UPat.var("buf").store(UPat.var("src")), name="a1").after(UPat.var("a1").store(UPat.var("src"))), lambda buf,src,a1:a1),
+
+  # store a buffer's own current contents back into itself: TestAssign.test_nested_after_contiguous_store_no_init
+  (UPat.var("buf").after(UPat.var("buf").store(UPat.var("buf").after(UPat.var("buf").store(UPat.var("src")), name="a1"))), lambda buf,src,a1:a1),
+
+  # move bitcast from store dest to source: TestAssign.test_assign_bitcast
+  (UPat(Ops.STORE, src=(UPat(Ops.BITCAST, src=(UPat(name="target"),)), UPat(name="src"))),
+   lambda target, src: target.store(src.bitcast(target.dtype))),
+
+  (UPat(Ops.BITCAST, name="bc"), expand_bitcast),
+
+  # ** size 0 **
+
+  # reduce of size 0 is the identity element
+  (UPat(Ops.REDUCE, name="reduce", src=(UPat.var("x"),)),
+   lambda reduce,x: reduce.const_like(identity_element(reduce.arg[0], reduce.dtype)) if 0 in x.shape and 0 not in reduce.shape else None),
+  # handle size 0
+  (UPat(GroupOp.All-{Ops.SINK}, name="x"), lambda x: x.const_like(0).rtag(x.tag) if x._shape is not None and 0 in x.shape else None),
+])
+
+def convert_copy_to_store(ctx, copy:UOp, existing_buf:UOp|None=None):
+  input_src = copy.src[0]
+  if not input_src.has_buffer_identity(after_ok=True): input_src = input_src.contiguous()
+  input_src = input_src.flatten()
+  if existing_buf is not None:
+    # if the existing buffer is not a full buffer, we can't use it
+    if not existing_buf.has_buffer_identity(after_ok=True): return None
+    # if there's already a buffer, we just use it
+    return existing_buf.flatten().store(input_src)
+  # create the output buffer
+  buf = UOp(Ops.BUFFER, src=(shape_to_shape_arg(input_src.max_shape),), arg=ParamArg(next(ctx), copy.dtype, device=copy.device))
+  # reshape back to input
+  return buf.after(buf.store(input_src)).reshape(copy.shape)
+
+pm_copy_to_store = PatternMatcher([
+  (UPat(name="existing_buf").store(UPat(Ops.COPY, name="copy")), convert_copy_to_store),
+  (UPat(Ops.COPY, name="copy"), convert_copy_to_store),
+])
+
+@rewrite_group(new_ctx=False)
+def prepare_rangeify(sink:UOp) -> UOp:
+  # prepare for rangeify
+  tsink = graph_rewrite(sink, multi_pm, name="multi_pm")
+  if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters")
+  tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites, bottom_up=True, name="earliest rewrites")
+  tsink = graph_rewrite(tsink, pm_copy_to_store, ctx=itertools.count(0), bottom_up=True, name="convert copy to store")
+  return tsink
@@ -1,191 +1,21 @@
 from dataclasses import dataclass, field, replace
 from typing import cast
 import itertools
-from tinygrad.dtype import dtypes, AddrSpace, Invalid, to_dtype, strong_dtype
+from tinygrad.dtype import dtypes, AddrSpace, Invalid, strong_dtype
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg, shape_to_shape_arg
-from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, rewrite_group, identity_element
+from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, rewrite_group
 from tinygrad.uop.symbolic import symbolic
-from tinygrad.uop.movement import mop_cleanup
-from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS, SPEC
-from tinygrad.helpers import PCONTIG, FLOAT16, OPENPILOT_HACKS, argsort, partition, get_single_element
+from tinygrad.helpers import prod, dedup, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS, SPEC
+from tinygrad.helpers import PCONTIG, partition, get_single_element
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
 from tinygrad.schedule.indexing import run_rangeify, BufferizeOpts, apply_movement_op
-from tinygrad.schedule.multi import multi_pm
-from tinygrad.schedule.allreduce import create_allreduce_function
+from tinygrad.schedule.prepare import pm_mops

 # creation can recurse a lot
 import sys
 sys.setrecursionlimit(10000)

-def found_after(ctx:dict[UOp, UOp], after:UOp, src:UOp):
-  if (x:=src).op is Ops.CAST and x.dtype == dtypes.half and FLOAT16: x, after = x.src[0], after.cast(dtypes.float)
-  while True:
-    if x.op is Ops.PERMUTE: x, after = x.src[0], after.permute(argsort(x.marg))
-    elif x.op is Ops.RESHAPE: x, after = x.src[0], after.reshape(x.src[0].shape)
-    elif x.op is Ops.WHERE and x.src[2].base.is_invalid and x.src[1].op is Ops.PAD:
-      x, after = x.src[1].src[0], after.shrink(tuple((o, s+o) for (o,_),s in zip(x.src[1].marg, x.src[1].src[0].shape)))
-    else: break
-  ctx[x] = after
-
-# *** fold moved AFTERs (hack for openpilot) ***
-pm_fold_moved_after = PatternMatcher([
-  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(), UPat((*GroupOp.Movement,Ops.CAST,Ops.WHERE), name="src")))), name="after"), found_after),
-  # replace ALU sources with AFTER versions found above
-  (UPat(GroupOp.ALU, name="alu"), lambda ctx,alu: alu.replace(src=new_src) if (new_src:=tuple(ctx.get(s, s) for s in alu.src)) != alu.src else None),
-])
-
-# movement op on INDEX as a PatternMatcher
-def _mop_index(r:UOp, idx:UOp):
-  idxs = idx.src[1:]
-  if len(idxs) == len(r.shape):
-    return r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idxs), arg=idx.arg)
-  if r.op is Ops.RESHAPE:
-    src_prefix = len(r.src[0].shape) - len(r.shape[len(idxs):])
-    if src_prefix >= 0 and r.src[0].shape[src_prefix:] == r.shape[len(idxs):]:
-      if src_prefix == 0: return r.src[0] if r.src[0].dtype == idx.dtype else None
-      ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape[:src_prefix], r.shape[:len(idxs)], idxs), arg=idx.arg)
-      return ret if ret.shape == idx.shape else None
-
-pm_mops = PatternMatcher([
-  # handle movement ops on INDEX
-  (UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"), _mop_index),
-  # move movement ops and INDEX after AFTER
-  (UPat(GroupOp.Movement|{Ops.INDEX}, name="r").after(name="a", allow_any_len=True),
-   lambda r,a: UOp(r.op, src=(a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], arg=r.arg)),
-  (UPat(GroupOp.Movement, name="r").end(name="a", allow_any_len=True), lambda r,a: a.replace(src=(r.src[0],)+a.src[1:])),
-])
-
-# *****************
-# 0. do some cleanup rewrites, mostly copied from the old stuff
-
-def fix_store_hazard(target:UOp, src:UOp):
-  if (base:=target.base) not in src.toposort(enter_calls=False): return None
-  # PERMUTE and FLIP reorder indices, SHRINK can have overlapping regions when dest is also shrunk
-  unsafe = {Ops.PERMUTE, Ops.FLIP} | ({Ops.SHRINK} if target.op_in_backward_slice_with_self(Ops.SHRINK) else set())
-  reaches_base: dict[UOp, bool] = {}
-  for s in src.toposort(gate=lambda s: s.op is not Ops.CONTIGUOUS):
-    reaches_base[s] = s is base or any(reaches_base.get(c) for c in s.src)
-    if reaches_base[s] and s.op in unsafe and not (s is target and s.op is Ops.SHRINK): return target.store(src.contiguous())
-
-def split_reduceop(reduce:UOp, x:UOp):
-  if prod(reduce.shape) == 0: return None
-  if not SPLIT_REDUCEOP or not all_int(x.shape) or (prod(x.shape)//prod(reduce.shape))<getenv("REDUCEOP_SPLIT_THRESHOLD", 32768): return None
-  # if there are few globals, make some reduces into globals by splitting into two kernels
-  # cap output buffer to 2**22: heuristic number of global outputs to achieve max occupancy with enough locals+upcasts for gemm
-  #   ~2**10 should be enough if GROUP is used
-  # 256 split maximum should be "negligible reduce" for low prod(reduce.shape), 8 split minimum.
-  # split is moved to the end to provide maximum locality for the second phase reduce.
-
-  # get expanded by rangeifying the UOp x
-  indexed = x.index(*[UOp.range(s, i) if resolve(s>1) else 0 for i,s in enumerate(x.shape)])
-  range_nums = [y.arg[0] for y in indexed.substitute({x.base:UOp(Ops.NOOP)}, extra_pm=pm_mops).ranges]
-  is_expanded = [i not in range_nums for i in range(len(x.shape))]
-
-  if not (split_candidates:=[(i,d) for i in range(reduce.arg[1])
-                             for d in range(min(256,2**getenv("REDUCEOP_SPLIT_SIZE",22)//prod(reduce.shape)),8-1,-1)
-                             if x.shape[i]%d==0 and not is_expanded[i]]): return None
-  dim_to_split, divisor = split_candidates[0]
-  splitted_shape = x.shape[:dim_to_split]+(divisor,)+(x.shape[dim_to_split]//divisor,)+x.shape[dim_to_split+1:]
-  splitted = x.reshape(splitted_shape).permute(tuple([d for d in range(len(splitted_shape)) if d!=dim_to_split]+[dim_to_split]))
-  if DEBUG >= 3: print(f"split {divisor}: {x.shape} -> {splitted.shape} -> {reduce.shape}")
-  # reduce original axes, then split
-  return splitted._rop(reduce.arg[0], tuple(range(reduce.arg[1]))).contiguous()._rop(reduce.arg[0], (len(reduce.shape),)).reshape(reduce.shape)
-
-pm_gather_params = PatternMatcher([ (UPat(Ops.PARAM, name="p"), lambda ctx, p: ctx.append(p) if p.arg.slot >= 0 else None), ])
-def resolve_function(c:UOp, allow_param_mismatch=True) -> UOp|None:
-  if c.arg.precompile: return None
-  params: list[UOp] = []
-  graph_rewrite(c.src[0], pm_gather_params, bottom_up=True, ctx=params, name="gather params")
-  params = sorted(params, key=lambda x: x.arg.slot)
-  args = c.src[1:]
-
-  # NOTE: this isn't really needed. it's okay if there's unused args in the function
-  if not allow_param_mismatch:
-    if [x.arg.slot for x in params] != list(range(len(params))): raise RuntimeError(f"params not in order: {[x.arg.slot for x in params]}")
-    if len(params) != len(args): raise TypeError(f"expected {len(params)} args, got {len(args)}")
-
-  dict_map = {x:args[x.arg.slot] for x in params}
-  for i, (p, a) in enumerate(dict_map.items()):
-    if p.axis != a.axis: raise TypeError(f"arg {i} axis mismatch: expected {p.axis}, got {a.axis}")
-    if p.max_shape != a.max_shape: raise TypeError(f"arg {i} shape mismatch: expected {p.shape}, got {a.shape}")
-    if p.dtype != a.dtype: raise TypeError(f"arg {i} dtype mismatch: expected {p.dtype}, got {a.dtype}")
-  return c.src[0].substitute(dict_map, walk=True)
-
-# shape-changing bitcast
-def expand_bitcast(bc:UOp) -> UOp|None:
-  x = bc.src[0]
-  if (ns:=bc.dtype.itemsize) == (os:=x.dtype.itemsize) or (isinstance(x.device, str) and x.device.startswith(("DISK", "TINYFS"))): return None
-  new_uint, tmp = to_dtype(f"uint{8*ns}"), x.bitcast(to_dtype(f"uint{8*os}"))
-  if ns > os:
-    tmp = tmp.reshape(x.shape[:-1] + (x.shape[-1]//(rate := ns//os), rate))
-    parts = [tmp.shrink((None,)*(len(tmp.shape)-1) + ((i, i+1),)).cast(new_uint)<<8*i*os for i in range(rate)]
-    return parts[0].usum(*parts[1:]).squeeze(-1).bitcast(bc.dtype)
-  parts = [tmp>>8*i*ns for i in range(os//ns)]
-  return parts[0].stack(*parts[1:], dim=-1).flatten(-2).cast(new_uint).bitcast(bc.dtype)
-
-earliest_rewrites = mop_cleanup+PatternMatcher([
-  # resolve FUNCTION calls (inline the body)
-  (UPat(Ops.FUNCTION, name="c"), resolve_function),
-
-  # resolve TUPLE+GETTUPLE
-  (UPat(Ops.GETTUPLE, src=(UPat(Ops.TUPLE, name="t"),), name="g"), lambda g,t: t.src[g.arg]),
-
-  # resolve allreduce (must be bottom up)
-  (UPat(Ops.ALLREDUCE, src=(UPat.var("buf"),), name="red"), create_allreduce_function),
-
-  # split_reduceop
-  (UPat(Ops.REDUCE, name="reduce", src=(UPat.var("x"),)), split_reduceop),
-
-  # remove DETACH/CONTIGUOUS_BACKWARD (TODO: this is copied in allocations)
-  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),
-
-  # SINK only ever references the base
-  (UPat(Ops.SINK, name="x"), lambda x: x.replace(src=tuple(y.unsharded_base for y in x.src))),
-
-  # ** copy rules **
-
-  # COPY transfers a contiguous range, so materialize a source that's resized (shrink/pad/expand) or reordered (permute/flip)
-  (UPat(Ops.COPY, src=(UPat(GroupOp.Movement, name="r"),), name="c"),
-   lambda c,r: c.replace(src=(r.contiguous(),)) if resolve(r.numel() != r.base.numel(), False) or r.contiguous_view_offset() is None else None),
-
-  # copy to same device is a no-op
-  (UPat(Ops.COPY, src=(UPat.var("x"),), name="copy"), lambda x,copy: x if x.device == copy.device else None),
-
-  # copy on reshape is reshape on copy
-  (UPat(Ops.COPY, src=(UPat(Ops.RESHAPE, name="shp"),), name="cpy"), lambda shp,cpy: shp.src[0].copy_to_device(cpy.device).reshape(shp.shape)),
-
-  # reshaping on STORE can be a NOOP
-  (UPat(Ops.STORE, src=(UPat(Ops.RESHAPE, src=(UPat.var("dst",),), allow_any_len=True),
-                        UPat(Ops.RESHAPE, src=(UPat.var("src",),), allow_any_len=True))),
-   lambda dst,src: dst.store(src) if dst.shape == src.shape else None),
-
-  # ** store rules **
-
-  # fix store hazard (dest is in used in src) by adding contiguous: TestAssign.test_post_flipped_assignment
-  (UPat(Ops.STORE, src=(UPat(name="target"), UPat(name="src"))), fix_store_hazard),
-
-  # remove two STOREs that store the same thing to the same place: TestSchedule.test_dedup_assign
-  (UPat.var("buf").after(UPat.var("buf").store(UPat.var("src")), name="a1").after(UPat.var("a1").store(UPat.var("src"))), lambda buf,src,a1:a1),
-
-  # store a buffer's own current contents back into itself: TestAssign.test_nested_after_contiguous_store_no_init
-  (UPat.var("buf").after(UPat.var("buf").store(UPat.var("buf").after(UPat.var("buf").store(UPat.var("src")), name="a1"))), lambda buf,src,a1:a1),
-
-  # move bitcast from store dest to source: TestAssign.test_assign_bitcast
-  (UPat(Ops.STORE, src=(UPat(Ops.BITCAST, src=(UPat(name="target"),)), UPat(name="src"))),
-   lambda target, src: target.store(src.bitcast(target.dtype))),
-
-  (UPat(Ops.BITCAST, name="bc"), expand_bitcast),
-
-  # ** size 0 **
-
-  # reduce of size 0 is the identity element
-  (UPat(Ops.REDUCE, name="reduce", src=(UPat.var("x"),)),
-   lambda reduce,x: reduce.const_like(identity_element(reduce.arg[0], reduce.dtype)) if 0 in x.shape and 0 not in reduce.shape else None),
-  # handle size 0
-  (UPat(GroupOp.All-{Ops.SINK}, name="x"), lambda x: x.const_like(0).rtag(x.tag) if x._shape is not None and 0 in x.shape else None),
-])
-
 # *****************
 # 3.5 cleanups

@@ -304,15 +134,15 @@ pm_const_buffer_folding = pm_mops+PatternMatcher([
  (UPat(Ops.INDEX, name="idx").f(Ops.STAGE, allow_any_len=True, name="b2"), remove_noop_bufferize),
  (UPat(Ops.INDEX, src=(UPat(Ops.STAGE),), allow_any_len=True, name="idx").f(Ops.NOOP).f(Ops.STAGE, allow_any_len=True, name="b2"),
   remove_noop_bufferize),
-  # no buffers for const (ranges don't matter for const - it's the same value everywhere)
-  (UPat(Ops.CONST, name='c').f(Ops.STAGE, allow_any_len=True, name="b"), lambda c,b: b.const_like(c.val)),
-  # indexing a const is a const
-  (UPat(Ops.INDEX, src=(UPat(Ops.CONST, name="c"),),), lambda c: c),
+  # no buffers for a const, in either spelling
+  (UPat.cvar('c').or_casted().f(Ops.STAGE, allow_any_len=True, name="b"), lambda c,b: b.const_like(c.val)),
+  # indexing a const is the const
+  (UPat(Ops.INDEX, src=(UPat.cvar().or_casted("c"),),), lambda c: c),
  # indexing an after with all fully invalid stores is invalid
  (UPat(Ops.INDEX, src=(UPat(Ops.AFTER, name="after"),), allow_any_len=True, name="idx"),
   lambda idx,after: idx.const_like(Invalid) if after_all_invalid(after) else None),
  # hack if a noop turned to a const
-  (UPat(Ops.NOOP, src=(UPat.cvar("c"),)), lambda c: c),
+  (UPat(Ops.NOOP, src=(UPat.cvar().or_casted("c"),)), lambda c: c),
  # a deviceless MSTACK src is the same value on every device, so indexing the stack is just indexing that value
  (UPat(Ops.MSTACK, src=(UPat.var("s"),), allow_any_len=True).f(Ops.INDEX, allow_any_len=True, name="idx"),
   lambda s,idx: idx.replace(src=(s,)+idx.src[1:]) if s.device is None else None),
@@ -562,33 +392,8 @@ split_kernels = PatternMatcher([
  (UPat((Ops.STORE, Ops.END), name="x"), split_store),
 ])

-def convert_copy_to_store(ctx, copy:UOp, existing_buf:UOp|None=None):
-  input_src = copy.src[0]
-  if not input_src.has_buffer_identity(after_ok=True): input_src = input_src.contiguous()
-  input_src = input_src.flatten()
-  if existing_buf is not None:
-    # if the existing buffer is not a full buffer, we can't use it
-    if not existing_buf.has_buffer_identity(after_ok=True): return None
-    # if there's already a buffer, we just use it
-    return existing_buf.flatten().store(input_src)
-  # create the output buffer
-  buf = UOp(Ops.BUFFER, src=(shape_to_shape_arg(input_src.max_shape),), arg=ParamArg(next(ctx), copy.dtype, device=copy.device))
-  # reshape back to input
-  return buf.after(buf.store(input_src)).reshape(copy.shape)
-
-pm_copy_to_store = PatternMatcher([
-  (UPat(name="existing_buf").store(UPat(Ops.COPY, name="copy")), convert_copy_to_store),
-  (UPat(Ops.COPY, name="copy"), convert_copy_to_store),
-])
-
@rewrite_group(new_ctx=False)
-def get_kernel_graph(sink:UOp) -> UOp:
-  # prepare for rangeify
-  tsink = graph_rewrite(sink, multi_pm, name="multi_pm")
-  if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters")
-  tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites, bottom_up=True, name="earliest rewrites")
-  tsink = graph_rewrite(tsink, pm_copy_to_store, ctx=itertools.count(0), bottom_up=True, name="convert copy to store")
-
+def get_kernel_graph(tsink:UOp) -> UOp:
  # convert movement ops to ranges
  tsink = run_rangeify(tsink, bool(DEBUG_RANGEIFY))

@@ -73,7 +73,7 @@ class Ops(FastEnum):
  # ** 5 -- control flow / consts / custom **

  # control flow ops
-  BARRIER = auto(); RANGE = auto(); IF = auto(); END = auto(); ENDIF = auto(); WAIT = auto()
+  BARRIER = auto(); RANGE = auto(); IF = auto(); END = auto(); ENDIF = auto()

  # const.
  CONST = auto()
@@ -116,13 +116,13 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None:
  match op:
    case Ops.STORE | Ops.LINEAR | Ops.SINK | Ops.PROGRAM | Ops.SOURCE | \
         Ops.END | Ops.BARRIER | Ops.GROUP | Ops.IF | Ops.ENDIF | \
-         Ops.TUPLE | Ops.FUNCTION | Ops.CUSTOM_FUNCTION | Ops.REWRITE_ERROR:
+         Ops.TUPLE | Ops.FUNCTION | Ops.CUSTOM_FUNCTION | Ops.REWRITE_ERROR | Ops.PYLITERAL:
      # always void
      return dtypes.void
    case Ops.CALL:
      # a CALL of an opaque body is void, a CALL of an address can return a value
      return dtypes.void if src[0].dtype is dtypes.void else None
-    case Ops.CUSTOM | Ops.CUSTOMI | Ops.PYLITERAL:
+    case Ops.CUSTOM | Ops.CUSTOMI:
      return None
    case Ops.INS:
      return None
@@ -158,6 +158,10 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType|None:
      return in_tuple.src[arg].dtype
    case Ops.GETADDR:
      return dtypes.uint64
+    case Ops.THREEFRY:
+      return dtypes.uint64
+    case Ops.FDIV:
+      return least_upper_float(promo_dtype(src))
    case Ops.SHL | Ops.SHR:
      if not all(dtypes.is_int(x.dtype) or x.base.is_invalid for x in src):
        raise RuntimeError(f"shift operands must be int, got {[x.dtype for x in src]}")
@@ -188,9 +192,8 @@ class UOpMetaClass(type):
  def __call__(cls, op:Ops, dtype:DType|None=None, src:tuple[UOp,...]=tuple(), arg:Any=None, tag:Any=None,
               metadata:tuple[Metadata,...]|None=None, _buffer:Buffer|None=None):
    if dtype is None: dtype = dtype_from_uop(op, src, arg) or dtypes.void
-    # CONST derives its dtype by value only when the constructor omits one
    # TODO: delete this once the dtype field is removed, for now it just re-implements spec.py
-    if SPEC == 2 and op is not Ops.CONST and (expected_dtype:=dtype_from_uop(op, src, arg)) is not None and expected_dtype != dtype:
+    if SPEC == 2 and (expected_dtype:=dtype_from_uop(op, src, arg)) is not None and expected_dtype != dtype:
      raise RuntimeError(f"bad dtype {dtype}, expected {expected_dtype} on {op}")
    if (wret:=UOpMetaClass.ucache.get(key:=(op, dtype, src, arg, tag), None)) is not None and (ret:=wret()) is not None: return ret
    UOpMetaClass.ucache[key] = weakref.ref(created:=super().__call__(*key))
@@ -256,8 +259,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
  def rtag(self, tag=True): return self.replace(tag=tag)
  @property
  def val(self):
-    assert self.op is Ops.CONST, f"val is only for CONST, got {self.op}"
-    return self.arg
+    if self.op is Ops.CONST: return self.arg
+    # a casted const CAST(dt, CONST(v)) is one const: .val reads the value through the CAST
+    assert self.op is Ops.CAST and self.src[0].op is Ops.CONST, f"val is only for consts, got {self.op}"
+    return self.src[0].val
  @property
  def is_invalid(self) -> bool: return self.op is Ops.CONST and self.val is Invalid
  @recursive_property
@@ -611,17 +616,18 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    if isinstance(b, UOp): return b.cast(dtype)
    # NOTE: it always has to be STACK now, even if they are all the same
    if isinstance(b, tuple): return UOp.stack(*[UOp.const(c, dtype) for c in b])
-    return UOp(Ops.CONST, dtype, arg=dtype.const(b), src=())
+    # .cast folds away at exactly the dtypes a CONST derives (bool/weakint/weakfloat): bare there, the pair everywhere else
+    return UOp(Ops.CONST, arg=dtype.const(b), src=()).cast(dtype)
  # weak CONST with width on the CAST. TODO: this is the final const
  @staticmethod
-  def cconst(b:ConstLike, dtype:DType): return UOp(Ops.CAST, dtype, src=(UOp.const(b),), arg=dtype)
+  def cconst(b:ConstLike, dtype:DType): return UOp(Ops.CAST, src=(UOp.const(b),), arg=dtype)
  @staticmethod
  def range(end:sint, axis_id, axis_type=AxisType.WEAK, *arg, dtype=dtypes.weakint, src=(), **kwargs):
    return UOp(Ops.RANGE, src=(sint_to_uop(end, dtype),)+src, arg=(axis_id, axis_type)+arg, **kwargs)
  @staticmethod
  def loop(axis_id:int, *arg): return UOp(Ops.RANGE, src=(UOp(Ops.NOOP),), arg=(axis_id, AxisType.WEAK)+arg)
  @staticmethod
-  def special(end:sint, name:str, dtype=dtypes.weakint): return UOp(Ops.SPECIAL, src=(sint_to_uop(end, dtype),), arg=name)
+  def special(end:sint, name:str): return UOp(Ops.SPECIAL, src=(sint_to_uop(end),), arg=name)
  @staticmethod
  def wmma(a:UOp, b:UOp, acc:UOp, dims:tuple[int, int, int], device:str, threads:int, tc_upcast_axes=None):
    # dtype_in is stored in the arg (not derived from src[0].dtype) because bitcast rewrites change src dtypes
@@ -797,7 +803,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
      case Ops.STACK:
        srcs = (self,)+tuple(arg)
        dtype = cast(DType, dtype_from_uop(Ops.STACK, srcs, None))
-        return UOp(Ops.STACK, dtype, tuple(u if u.base.is_invalid else UOp.const(u.val, dtype) if u.op is Ops.CONST else u.cast(dtype) for u in srcs))
+        return UOp(Ops.STACK, src=tuple(u if u.base.is_invalid else UOp.const(u.val, dtype) if u.op is Ops.CONST else u.cast(dtype) for u in srcs))
      case _: raise RuntimeError(f"{op} is not a MovementOp")
    usrcs = [shape_to_shape_arg(arg) for arg in src_args]
    if len(usrcs) == 0: return UOp(op, src=(self,), arg=arg)
@@ -888,7 +894,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    return s

  def contiguous_view(self) -> tuple[UOp, int]|None:
-    from tinygrad.schedule.rangeify import pm_mops
+    from tinygrad.schedule.prepare import pm_mops
    from tinygrad.uop.symbolic import symbolic

    # WEBGPU and CL do not support views.
@@ -987,7 +993,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    return unwrap(self.arg.name)
  def bind(self, val:int|UOp):
    assert self.is_variable, f"op is {self.op}, need Variable"
-    uval = self.const_like(val) if isinstance(val, int) else val
+    # the Variable states the width, so the bound value stays bare: is_bound_var tests for a CONST there, unbind reads .val
+    uval = UOp.const(val) if isinstance(val, int) else val
    assert self.vmin <= uval.vmin and uval.vmax <= self.vmax, f"bind {val} not in range [{self.vmin}, {self.vmax}]"
    assert uval.divides(self.arg.multiple_of) is not None, f"bind {val} not divisible by {self.arg.multiple_of}"
    return self.after(self.store(uval))
@@ -40,7 +40,7 @@ renderer = PatternMatcher([
  (UPat(Ops.RANGE, name="x"), lambda x: f"r{range_str(x)}"),
  (UPat(Ops.CONST, name="x"), lambda x: str(x.val)),
  # CAST states the width, the weak CONST carries the value
-  (UPat.cvar("c", dtypes.weaks+(dtypes.bool,)).cast(), lambda c: str(c.val)),
+  (UPat.cvar("c").cast(), lambda c: str(c.val)),
  (UPat(Ops.CAST, name="x"), lambda ctx,x: f"({str(x.dtype)[7:]})({ctx[x.src[0]]})"),
  (UPat(Ops.NEG, name="x"), lambda ctx,x: f"(-{ctx[x.src[0]]})"),
  (UPat(Ops.RECIPROCAL, name="x"), lambda ctx,x: f"(1/{ctx[x.src[0]]})"),
@@ -83,9 +83,9 @@ def render_marg(ctx,x:UOp):
 sugar = {Ops.SINK, Ops.END, Ops.STORE, Ops.LOAD, Ops.SQRT, Ops.INDEX, Ops.REDUCE, Ops.AFTER, Ops.THREEFRY,
         Ops.RECIPROCAL, Ops.EXP2, Ops.LOG2, Ops.SIN, Ops.CONTIGUOUS, Ops.BARRIER, Ops.DETACH}
 pm_pyrender_extra = PatternMatcher([
-  (UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.val}, {x.dtype})"),
+  (UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.val})"),
  (UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})" if x.dtype != x.src[0].dtype else None),
-  (UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].val}, {repr(x.arg)}, dtype={x.dtype})"),
+  (UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].val}, {repr(x.arg)})"),
  (UPat(Ops.BUFFER, src=(UPat(),), name="x"), lambda x:
    f"UOp.new_buffer({repr(x.arg.device)}, {x.max_numel()}, {x.dtype}, {x.arg.slot})"
    if isinstance(x.arg, ParamArg) and x.addrspace is AddrSpace.GLOBAL else None),
@@ -95,7 +95,7 @@ pm_pyrender_extra = PatternMatcher([
  # NOTE: range has srcs sometimes after control flow
  (UPat(Ops.RANGE, src=(UPat(Ops.CONST, name="c"),), allow_any_len=True, name="x"), lambda ctx,x,c:
    "UOp.range("+', '.join([str(c.val)] + [repr(y) for y in x.arg])+
-      (f', src={srcs(ctx, x.src[1:])}' if len(x.src) > 1 else '')+(', dtype='+str(x.dtype) if x.dtype is not dtypes.weakint else '')+")"),
+      (f', src={srcs(ctx, x.src[1:])}' if len(x.src) > 1 else '')+")"),
  # TODO: index shouldn't mismatch dtype
  (UPat(Ops.INDEX, src=(UPat(), UPat()), allow_any_len=True, name="x"), lambda ctx,x:
   f"{ctx[x.src[0]]}.index({ctx[x.src[1]]}, "+''.join([f"{ctx[xx]}, " for xx in x.src[2:]])+
@@ -203,10 +203,9 @@ spec_tensor = PatternMatcher([

 # these ops can exist in programs but not the tensor spec. example: LOAD
 spec_program = PatternMatcher([
-  # a literal is CAST(dt, CONST(value)), so its inner CONST is the one weak node a program may contain
-  (UPat(Ops.CONST, dtype=dtypes.weaks, name="x"), lambda x: x.dtype is dtypes.from_py(x.val)),
-  # index and weak dtypes are not allowed in programs
-  (UPat(GroupOp.All, (dtypes.weakint, dtypes.weakfloat)), lambda: False),
+  # every width in a program is stated: a CONST appears only under the CAST stating its width, and is the only weak node
+  (UPat(GroupOp.All, name="x"), lambda x: False if x.op is not Ops.CAST and any(s.op is Ops.CONST for s in x.src) else None),
+  (UPat(GroupOp.All-{Ops.CONST}, dtypes.weaks), lambda: False),

  # allow special SHRINK
  (UPat(Ops.SHRINK, src=(UPat((Ops.PARAM, Ops.BUFFER, Ops.AFTER)), UPat(), UPat(Ops.CONST).or_casted())), lambda: True),
@@ -254,8 +253,9 @@ spec_kernel_graph = PatternMatcher([
  (UPat(Ops.SINK, dtypes.void), lambda: True),
  # the store of a bound Variable binds it: AFTER(BUFFER, STORE(BUFFER, CONST)) in call args
  (UPat(Ops.STORE, dtypes.void, (UPat(Ops.BUFFER, name="b"), UPat(Ops.CONST))), lambda b: b.is_variable),
-  # const + stack to make vconsts and shape args
+  # const + stack to make vconsts and shape args. a 0-size/bound reduce keeps its const casted
  (UPat(Ops.CONST, src=()), lambda: True),
+  (UPat(Ops.CAST, src=(UPat(Ops.CONST, src=()),)), lambda: True),
  (UPat(Ops.STACK, name="s"), lambda s: all(x.op in (Ops.CONST, Ops.PARAM) or x.is_variable or x.is_bound_var for x in s.src) or None),
  # linear for more kernels (TODO: we should enter non sink calls)
  #(UPat(Ops.LINEAR), lambda: True),
@@ -1,12 +1,12 @@
 # all of symbolic lives here now
 import math
 from collections import defaultdict
-from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu
-from tinygrad.dtype import PyConst, ConstType, dtypes, can_lossless_cast, Invalid, bitcast, truncate
+from tinygrad.uop.ops import Ops, PatternMatcher, UPat, UOp, GroupOp, exec_alu, promo_dtype
+from tinygrad.dtype import PyConst, dtypes, can_lossless_cast, Invalid, bitcast, truncate
 from tinygrad.helpers import partition, all_same, prod, flatten, unwrap, IMAGE, dedup
 from tinygrad.uop.divandmod import div_and_mod_symbolic
 from tinygrad.uop.movement import mop_cleanup
-from tinygrad.uop.weak import commit_weak
+from tinygrad.uop.weak import pm_uncast_const, commit_weak

 # TODO: symbolic shouldn't be importing from codegen
 from tinygrad.codegen.decomp.transcendental import xpow
@@ -22,20 +22,11 @@ def simplify_pow(x:UOp, c:UOp) -> UOp|None:

 def fold_bitcast(root:UOp, c:UOp) -> UOp|None:
  if c.dtype.fmt is None or root.dtype.fmt is None or c.dtype.itemsize != root.dtype.itemsize: return None
+  # the value is mathematical and may not fit: reading it as bits is the emission that pins it to the stated width
  return root.const_like(bitcast(truncate[c.dtype](c.val), c.dtype, root.dtype))

-# const folding works for CONST and STACK
-const_folding_pat = UPat((Ops.CONST, Ops.STACK))
-
-def const_arg(u:UOp) -> ConstType|tuple[ConstType, ...]|None:
-  if u.op is Ops.CONST: return u.val
-  if u.op is Ops.CAST and u.src[0].op is Ops.CONST: return u.dtype.const(u.src[0].val)
-  if u.op is Ops.STACK and all(s.op is Ops.CONST for s in u.src): return tuple(s.val for s in u.src)
-  return None
-
-def fold_const_alu(a:UOp) -> UOp|None:
-  vals = [const_arg(s) for s in a.src]
-  return None if any(v is None for v in vals) else a.const_like(exec_alu(a.op, a.dtype, vals, False))
+# no truncate: ints stay mathematical past the fold (emission truncates); floats re-round in the mint
+def fold_const_alu(a:UOp) -> UOp: return a.const_like(exec_alu(a.op, a.dtype, [const_arg(s) for s in a.src], False))

 def _quotient_base(q:UOp, base:UOp, div:int) -> UOp|None:
  # the B with q == B//div and B%div == base%div, or None. only such congruence is needed to recombine, and canonicalization
@@ -71,6 +62,12 @@ def fold_add_divmod_recombine(x:UOp) -> UOp|None:
 # this needs to be before symbolic so that 0*something_that_might_be_invalid doesnt become 0
 invalid_pat = UPat(Ops.CONST, arg=Invalid, name="i")
 invalid_gate = UPat.var("cond").where(UPat.var("x"), invalid_pat)
+
+# the two const spellings: Invalid carries no width, so it rides bare inside either
+bare_const = UPat.any(UPat(Ops.CONST), UPat(Ops.STACK, src=UPat(Ops.CONST)))
+casted_const = UPat.any(p:=UPat(Ops.CAST, src=(UPat(Ops.CONST),)), UPat(Ops.STACK, src=UPat.any(p, UPat(Ops.CONST, arg=Invalid))))
+def const_arg(u:UOp):
+  return tuple(const_arg(s) for s in u.src) if u.op is Ops.STACK else u.val
 pm_data_invalid = PatternMatcher([
  (invalid_pat.broadcast(), lambda i: i),
  (UPat(GroupOp.Unary|{Ops.CAST, Ops.BITCAST}, src=(invalid_pat,)), lambda i: i),
@@ -142,13 +139,19 @@ symbolic_simple = pm_data_invalid + PatternMatcher([
  (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)) != UPat.var("x"),
   lambda x: x.const_like(False, dtypes.bool)), # x != x -> False (only ints)
  # ** constant folding **
-  # a CAST to a concrete dtype over a CONST is a value conversion: evaluate it once, at the width the CAST states
+  # a CAST to a concrete dtype over a CONST is a value conversion: evaluate it once, at the CAST's dtype
  # TODO: delete this once CONST has no dtype
  (UPat(Ops.CAST, dtypes.all, name="root", src=(UPat.cvar("c"),)), lambda root, c: root.const_like(c.val)),
-  (UPat(GroupOp.Unary, src=(const_folding_pat,), name="a"), fold_const_alu),
+  # collapse committed const conversions when the target has a native constant format. fmt-less targets are emulated and would re-expand this pair.
+  (UPat(Ops.CAST, dtypes.all, name="root", src=(UPat(Ops.CAST, dtypes.all, src=(UPat(Ops.CONST, name="c"),)),)),
+   lambda root,c: root.const_like(c.val) if root.dtype.fmt is not None else None),
+  # one rule per spelling: bare has no width, a pair evaluates at its stated width, mixed commits to the promotion
  # NOTE: THREEFRY(const,const) folds via its decomposition
-  (UPat(GroupOp.Binary-{Ops.THREEFRY}, src=(const_folding_pat,)*2, name="a"), fold_const_alu),
-  (UPat(GroupOp.Ternary, src=(const_folding_pat,)*3, name="a"), fold_const_alu),
+  (UPat(GroupOp.ALU-{Ops.THREEFRY}, src=bare_const, name="a"), fold_const_alu),
+  (UPat(GroupOp.ALU-{Ops.THREEFRY}, src=casted_const, name="a"), fold_const_alu),
+  (UPat(GroupOp.Binary-{Ops.THREEFRY}, src=[casted_const, bare_const], name="a"), lambda a:
+   a.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in a.src))
+   if (dt:=promo_dtype(a.src)) not in dtypes.weaks else None),
  # bool MUL is AND, ADD/MAX is OR. prevents other rules to rewrite bool ADD/MUL incorrectly
  (UPat.var('x', dtype=dtypes.bool) * UPat.var('y', dtype=dtypes.bool), lambda x,y: x&y),
  (UPat.var('x', dtype=dtypes.bool) + UPat.var('y', dtype=dtypes.bool), lambda x,y: x|y),
@@ -166,7 +169,9 @@ symbolic_simple = pm_data_invalid + PatternMatcher([
                                             and isinstance(x.val, float) and (math.isnan(x.val) or math.isinf(x.val)) else 0)),
  # *** cast/bitcast ***
  (UPat((Ops.CAST, Ops.BITCAST), name="root"), lambda root: root.src[0] if root.dtype == root.src[0].dtype else None),
-  (UPat(Ops.BITCAST, name="root", src=(UPat.cvar("c"),)), fold_bitcast),
+  # a BITCAST reads its operand at the width it states, so a weak const is nonsense here: the bare arm is bool only
+  (UPat(Ops.BITCAST, name="root", src=(UPat.any(UPat(Ops.CONST, dtypes.bool, name="c"), UPat(Ops.CAST, src=(UPat(Ops.CONST),), name="c")),)),
+   fold_bitcast),
  # b.cast(a).cast(b) -> b if a preserves all values in b
  (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x if x.dtype == b.dtype and can_lossless_cast(b.dtype, a.dtype) else None),
  # bitcast twice
@@ -292,8 +297,9 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
  # cast/long folding
  # if the intermediate cast doesnt narrow we can do it in one cast
  (UPat.var('x').cast(name="a").cast(name="b"), lambda x,a,b: x.cast(b.dtype) if can_lossless_cast(x.dtype, a.dtype) else None),
+  # commit_weak, not .cast: a weak b.dtype is not a const spelling, and a CAST(weakfloat, CONST) reaches no commit round
  (UPat.var('x', dtypes.ints+(dtypes.weakint,)).cast(dtypes.ints+(dtypes.weakint,), name="a").cast(name="b"),
-    lambda x,a,b: x.cast(b.dtype) if a.dtype.min<=x.vmin and x.vmax<=a.dtype.max else None),
+    lambda x,a,b: commit_weak(x, b.dtype) if a.dtype.min<=x.vmin and x.vmax<=a.dtype.max else None),
  # try to do math in int instead of long, keep weak const weak
  (UPat(GroupOp.Binary, src=(UPat.var("x", (dtypes.long, dtypes.weakint)), UPat.var("y", (dtypes.long, dtypes.weakint))), name="u"), lambda u,x,y:
    (UOp.const(x.val) if x.op is Ops.CONST else x.cast(dtypes.int)).alu(u.op,
@@ -306,7 +312,8 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
                        else y.src for y in x.src[1:]]))))),
  # after/end with 1 src is just src[0]
  (UPat((Ops.AFTER, Ops.END), src=(UPat.var("s"),)), lambda s: s),
-])+div_and_mod_symbolic
+  # the rules above key on bare CONSTs, so a redundantly committed const has to be uncast in the same fixpoint
+])+div_and_mod_symbolic+pm_uncast_const

 # ******** we take a small aside to "simplify_valid" to rewrite valids ********

@@ -45,7 +45,7 @@ z3_renderer = PatternMatcher([
  (UPat((Ops.LOAD, Ops.INDEX), dtypes.bool), lambda ctx: (z3.Bool(f"load{len(ctx[1])}", ctx=ctx[0]), None)),
  # constants
  (UPat(Ops.CONST, arg=Invalid), lambda ctx: (z3.Int("Invalid", ctx=ctx[0]), None)),
-  (UPat(Ops.CONST, dtypes.ints+(dtypes.weakint,), name="x"), lambda x,ctx: (z3.IntVal(x.val, ctx=ctx[0]), None)),
+  (UPat(Ops.CONST, dtypes.weakint, name="x"), lambda x,ctx: (z3.IntVal(x.val, ctx=ctx[0]), None)),
  (UPat(Ops.CONST, dtypes.bool, name="x"), lambda x,ctx: (z3.BoolVal(x.val, ctx=ctx[0]), None)),
  # casts from floats create new variables
  (UPat(Ops.CAST, dtypes.ints+(dtypes.weakint,), src=(UPat(dtype=dtypes.floats),), name="x"), lambda x,ctx:
@@ -1,19 +1,32 @@
 from dataclasses import replace
 from tinygrad.dtype import dtypes, DType, AddrSpace, Invalid, least_upper_dtype, strong_dtype, weak_dtype
 from tinygrad.helpers import unwrap
-from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, GroupOp, graph_rewrite, dtype_from_uop
+from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, GroupOp, dtype_from_uop, promo_dtype

 def default_dtype(u:UOp):
  if u.dtype is dtypes.weakfloat: return dtypes.default_float
  return dtypes.long if u.overflows(dtypes.int32) else dtypes.int

 def commit_weak(s:UOp, dt:DType) -> UOp:
-  # a CONST commits directly at dt (the value stays mathematical, emission truncates), a non-const src takes the cast
+  # a CONST re-mints, never takes a cast: at bool/weakint/weakfloat a CAST would be a second spelling of one const
  return UOp.const(s.val, dt) if s.op is Ops.CONST else s.cast(dt)

-def commit_srcs_at(u:UOp, dt:DType) -> UOp:
+# the decomps and float emulation commit bare consts at a dtype another src already states
+def commit_weak_consts(u:UOp, dt:DType|None) -> UOp|None:
+  return None if dt is None else u.replace(src=tuple(commit_weak(s, dt) if s.op is Ops.CONST and s.dtype in dtypes.weaks else s for s in u.src))
+
+# the concrete dtypes u commits its srcs at: the operands' meet and u's own derived dtype, None if either is weak
+def derived_dtypes(u:UOp, src:tuple[UOp, ...]) -> tuple[DType, DType]|None:
+  if u.op not in GroupOp.Broadcastable or (meet:=promo_dtype(src)) in dtypes.weaks \
+     or (result:=unwrap(dtype_from_uop(u.op, src, u.arg))) in dtypes.weaks: return None
+  return meet, result
+
+def commit_srcs_at(u:UOp, dt:DType) -> UOp|None:
  # the root re-derives: a shift's dtype is its lhs's, so committing the lhs commits the node too
-  return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src))
+  dts = derived_dtypes(u, u.src)
+  ret = u.replace(dtype=None, src=tuple(UOp.const(dt.const(s.val)) if s.op is Ops.CONST and s.dtype in dtypes.weaks and dts is not None else
+                                        commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src))
+  return None if ret is u else ret

 def commit_weak_srcs(u:UOp) -> UOp|None:
  if not any(s.dtype in dtypes.weaks for s in u.src) or (dt:=least_upper_dtype(*(s.dtype for s in u.src))) in dtypes.weaks: return None
@@ -21,62 +34,71 @@ def commit_weak_srcs(u:UOp) -> UOp|None:

 # a concrete CAST over a weak node states the width the value will live at. that width is a floor, never a narrowing
 def cast_weak_srcs(c:UOp, u:UOp) -> UOp|None:
+  # only within the kind: an int cast of a weakfloat node is a value conversion, not a statement about the node's width
  if c.dtype in dtypes.weaks or weak_dtype(c.dtype) is not u.dtype: return None
-  return commit_srcs_at(u, least_upper_dtype(c.dtype, default_dtype(u))).cast(c.dtype)
+  return None if (ret:=commit_srcs_at(u, least_upper_dtype(c.dtype, default_dtype(u)))) is None else ret.cast(c.dtype)

-pm_cast_weak = PatternMatcher([
-  (UPat(Ops.CAST, name="c", src=(UPat(GroupOp.ALU, dtype=dtypes.weaks, name="u"),)), cast_weak_srcs),
-  (UPat(Ops.CAST, name="c", src=(UPat(Ops.CONST, dtype=dtypes.weaks, name="u"),)), lambda c,u: commit_weak(u, c.dtype)),
-])
-
-# runs in index lowering and in the decomps: a rule that mints a weak const commits it in the same rewrite, so none reaches the renderer
+# rides every round that can mint a weak const, and must reach fixpoint before pm_lower_weak below defaults one
 pm_commit_weak = PatternMatcher([
  (UPat(GroupOp.Broadcastable, name="u"), commit_weak_srcs),
-  # demand from the destination: a STORE's weak value commits at the destination's dtype
  (UPat(Ops.STORE, src=(UPat(), UPat(dtype=dtypes.weaks)), allow_any_len=True, name="u"),
   lambda u: u.replace(src=(u.src[0], commit_weak(u.src[1], u.src[0].dtype), *u.src[2:]))),
+  # no CONST arm: a concrete CAST over a weak CONST is already committed, minted that way by UOp.const
+  (UPat(Ops.CAST, name="c", src=(UPat(GroupOp.ALU, dtype=dtypes.weaks, name="u"),)), cast_weak_srcs),
 ])

-# A weakfloat Unary (sin/exp2/...) must resolve here, before the transcendental decomposition.
+# consumers absorb the weak CAST off their srcs and default underivable consts; dtype-producing ops settle here.
+# a weakfloat Unary (sin/exp2/...) must resolve before the transcendental decomposition.
 _lower_weak_ops = GroupOp.Binary|GroupOp.Unary|{Ops.WHERE, Ops.RANGE, Ops.STACK, Ops.SPECIAL}
 def lower_weak_node(u:UOp) -> UOp|None:
+  if u.op is Ops.CAST and u.src[0].op is Ops.CONST: return None  # a committed const, not a consumer
  src = tuple(s.src[0] if s.op is Ops.CAST and s.dtype in dtypes.weaks else s for s in u.src)
+  dts = derived_dtypes(u, src)
+  src = tuple(commit_weak(s, default_dtype(s)) if s.op is Ops.CONST and s.dtype in dtypes.weaks and dts is None else s
+              for s in src)
  start = 1 if u.op is Ops.WHERE else 0  # WHERE's cond is bool, never part of the width unification
-  if src == u.src or any(s.dtype in dtypes.weaks for s in src[start:]): return None
-  # Binary can widen from the bounds, all other nodes derive from the lowered sources.
-  dt = strong_dtype(least_upper_dtype(default_dtype(u), *(s.dtype for s in src)) if u.op in GroupOp.Binary
-                    else unwrap(dtype_from_uop(u.op, src, u.arg)))
-  return u.replace(dtype=None, src=src[:start]+tuple(s if s.base.is_invalid else commit_weak(s, dt) for s in src[start:])).cast(u.dtype)
+  # resolve whole once every weak expression lowered: a Binary widens from its own bounds too, derivable consts wait
+  if u.op in _lower_weak_ops and src != u.src and not any(s.dtype in dtypes.weaks and s.op is not Ops.CONST for s in src[start:]):
+    dt = strong_dtype(least_upper_dtype(default_dtype(u), *(s.dtype for s in src)) if u.op in GroupOp.Binary
+                      else unwrap(dtype_from_uop(u.op, src, u.arg)))
+    return u.replace(dtype=None, src=src[:start]+tuple(s if s.base.is_invalid or s.dtype in dtypes.weaks else commit_weak(s, dt)
+                                                       for s in src[start:])).cast(u.dtype)
+  return None if src == u.src else u.replace(dtype=None, src=src)

 pm_lower_weak = PatternMatcher([
-  (UPat(Ops.CONST, dtype=dtypes.weaks, name="u"), lambda u: UOp.const(u.val, default_dtype(u)).cast(u.dtype)),
+  # a gated long index into a small buffer narrows; its out-of-gate value is discarded
+  (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"), UPat.var("gate").where(UPat.var("idx", dtypes.long), UPat(Ops.CONST, arg=Invalid))),
+        allow_any_len=True, name="u"),
+   lambda u,buf,gate,idx: u.replace(src=(buf, idx.cast(dtypes.int).valid(gate))+u.src[2:]) if buf.max_numel()-1 <= dtypes.int32.max else None),
  # two stacked weak casts are two kind conversions: each resolves at its own kind's default
-  # a SINGLE weak cast is never rewritten here, each consumer absorbs it on its own edge (see lower_weak_srcs)
  (UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat.var("x"),)),), name="u"),
   lambda u,x: x.cast(default_dtype(u.src[0])).cast(default_dtype(u)).cast(u.dtype) if x.dtype not in dtypes.weaks else None),
  (UPat((Ops.PARAM, Ops.BUFFER), dtype=dtypes.weakint, name="u"),
    lambda u: u.replace(dtype=None, arg=replace(u.arg, dtype=default_dtype(u))).cast(dtypes.weakint) if u.addrspace == AddrSpace.ALU else None),
-  (UPat(_lower_weak_ops, name="u"), lower_weak_node),
+  (UPat(GroupOp.All, name="u"), lower_weak_node),
 ])

-def lower_weak_srcs(ctx:dict[UOp, UOp]|None, u:UOp) -> UOp|None:
-  if ctx is None: ctx = {}
-  def lower(s:UOp) -> UOp:
-    if (r:=ctx.get(s)) is None:
-      r = graph_rewrite(s, pm_lower_weak)
-      # the consumer absorbs the cast on its own edge
-      ctx[s] = r = r.src[0] if r.op is Ops.CAST and r.dtype in dtypes.weaks else r
-    return r
-  # a comparison demands a common operand width: lower it whole so the Binary rule unifies its operands
-  ret = lower(u) if u.op in GroupOp.Comparison else u.replace(src=tuple(lower(s) if s.dtype in dtypes.weaks else s for s in u.src))
-  return None if ret is u else ret
+# drop the CAST off a committed const where the consumer re-derives it anyway, so bare-CONST rules keep matching.
+# the drop must change nothing the consumer derives: neither the operands' meet nor the node's own dtype
+def uncast_const(u:UOp) -> UOp|None:
+  # a weak CAST over a const is not a commit, it is still resolving
+  src = tuple(s.src[0] if s.op is Ops.CAST and s.dtype not in dtypes.weaks and s.src[0].op is Ops.CONST
+              and s.src[0].dtype in dtypes.weaks else s for s in u.src)
+  if src == u.src or (dts:=derived_dtypes(u, src)) is None or dts[0] != promo_dtype(u.src) or dts[1] is not u.dtype: return None
+  return u.replace(src=src)

-pm_lower_index_dtype = pm_commit_weak+pm_cast_weak+PatternMatcher([
-  (UPat(GroupOp.All, name="u"),
-   lambda ctx,u: lower_weak_srcs(ctx, u) if u.dtype not in dtypes.weaks and any(s.dtype in dtypes.weaks for s in u.src) else None),
-  # a valid index into an n-element buffer lives in [0,n): a gated long index narrows when n-1 fits int32 (out-of-gate wraps, discarded)
-  # TODO: more generic
-  (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"), UPat.var("gate").where(UPat.var("idx", dtypes.long), UPat(Ops.CONST, arg=Invalid))),
-        allow_any_len=True, name="u"),
-   lambda u,buf,gate,idx: u.replace(src=(buf, idx.cast(dtypes.int).valid(gate))+u.src[2:]) if buf.max_numel()-1 <= dtypes.int32.max else None),
-])
+pm_uncast_const = PatternMatcher([(UPat(GroupOp.Broadcastable, name="u"), uncast_const)])
+
+def cast_const(u:UOp, s:UOp) -> UOp:
+  if s.op is not Ops.CONST or s.is_invalid: return s  # Invalid never commits
+  # bool is the one strong bare dtype: cconst, since .cast(bool) would fold at construction
+  if s.dtype is dtypes.bool: return UOp.cconst(s.val, s.dtype)
+  # commit at the dtype its consumer derives; where nothing does, commit_weak is the identity and spec_program rejects it
+  return commit_weak(s, dts[0]) if (dts:=derived_dtypes(u, u.src)) is not None else s
+
+# commit every remaining bare const, keyed on the consumer: "bare" is a property of the edge
+def cast_consts(u:UOp) -> UOp|None:
+  if u.op is Ops.CAST and u.src[0].op is Ops.CONST: return None  # a committed const's CONST is its value, not an edge
+  return None if (src:=tuple(cast_const(u, s) for s in u.src)) == u.src else u.replace(src=src)
+
+pm_cast_const = PatternMatcher([(UPat(GroupOp.All, name="u"), cast_consts)])
Author	SHA1	Message	Date
sirhcm	4994ba82d5	benchmarks: cleanup llm	2026-08-25 19:19:16 -07:00
wozeparrotandGitHub	07268b724f	fix: external_test_gpu_crash on python 3.14 (#17739 )	2026-08-25 15:43:30 -07:00
sirhcmandGitHub	55032514ce	pin onnxruntime==1.24.1 (#17738 )	2026-08-25 17:53:15 -04:00
YassineYousfiandGitHub	2824504f90	usb amd: yield between signal polls (#17712 ) * AMD_USB_POLL_US * com * its micro	2026-08-25 14:11:53 -07:00
sirhcmandGitHub	7fef98c86e	benchmarks: remove usage of sudo from linux runners (#17735 )	2026-08-25 16:21:08 -04:00
geohot	b831ca62d9	Reapply "disk cache: thread-local db conn (#17694 )" This reverts commit `df528499ce`.	2026-08-25 13:20:15 -07:00
nimlgenandGitHub	d851aca9ae	hcq2: fix copy/call in usb (#17732 )	2026-08-25 19:03:19 +03:00
chenyuandGitHub	7dc8b666e7	wgsl cast before load for packed [pr] (#17731 ) instead of explicit dtype on load	2026-08-25 11:46:31 -04:00
nimlgenandGitHub	023bfdb380	fix viz for hcq2 (#17729 )	2026-08-25 18:23:22 +03:00
chenyuandGitHub	9f01775cf4	dtype_from_uop for THREEFRY and FDIV (#17727 )	2026-08-25 08:43:59 -04:00
chenyuandGitHub	9607787ce1	delete more explicitly set dtype to UOp [PR] (#17726 ) * delete more explicitly set dtype to UOp [PR] * not that	2026-08-25 08:32:45 -04:00
qazalandGitHub	ab79879613	bring back sqtt examples tests (#17725 )	2026-08-25 17:00:18 +09:00
George HotzandGitHub	a5678317c2	split rangeify to prepare.py (#17722 )	2026-08-24 19:12:13 -07:00
chenyuandGitHub	1d694dd700	remove more explicitly set dtype [PR] (#17721 )	2026-08-24 22:03:22 -04:00
chenyuandGitHub	65ca68567e	const cleanups [PR] (#17719 )	2026-08-24 21:37:14 -04:00
chenyuandGitHub	a7df1a1ace	const are weak 3 [pr] (#17695 )	2026-08-24 20:32:57 -04:00
RaineandGitHub	d9004cff22	use native shifts in Payne Hayek instead of pow2 mul/divs (#17717 ) * use native shifts instead of pow2 mul/divs * wow	2026-08-24 17:11:46 -07:00
wozeparrotandGitHub	6f87158d77	feat: bump version to 0.14.0 (#17716 )	2026-08-24 15:40:15 -07:00
sirhcmandGitHub	021c015eb4	benchmarks: more aggressive timeouts (#17715 )	2026-08-24 18:32:36 -04:00
sirhcmandGitHub	2e7790f16c	benchmarks: cleanup dsp benchmark (#17706 )	2026-08-24 18:09:01 -04:00
chenyuandGitHub	8df3dac0ec	remove redundant UOp dtype [PR] (#17713 )	2026-08-24 18:00:02 -04:00
George HotzandGitHub	ed110993d3	mi350p: aqua vanjaram bringup for the raw PCI driver (#17693 ) * mi350p: aqua vanjaram bringup for the raw PCI driver Minimum stable bringup, all root-caused: - disable ASPM on the PCI path: link L1 across retimers makes GPU reads oscillate to 0xffffffff (no-op on non-sysfs backends) - probe live AIDs via MMHUB FB_LOCATION; dead ones read 0xffffffff and indirect writes to them corrupt the fabric - aqua reset semantics: full boot over a live state can kill the fabric, partial boot + reset_mec is the deepest safe reset - tolerate the HQD dequeue timeout like the kernel (wedged waves can survive RESET_WAVES) * am/aspm: simplify to GPU-endpoint-only clearing via pci config abstractions * fix AID derivation for non-dense sdma instance keys + bound the aspm capability walk for 0xff reads	2026-08-24 14:32:16 -07:00
chenyuandGitHub	f4cc28824c	delete GETTUPLE on UNSHARD rule [PR] (#17711 ) not used	2026-08-24 17:25:28 -04:00
George HotzandGitHub	76bf6b7eec	llm: add DEBUG=2 for weights loading (#17710 ) * llm: add DEBUG=2 for weights loading * 00	2026-08-24 14:09:34 -07:00
geohot	08fbd25f1a	AGENTS: no amend commits	2026-08-24 13:34:42 -07:00
George HotzandGitHub	5aabbb1991	fix Muon weight decay being a no-op (#17709 ) * fix Muon weight decay being a no-op LARS._step computed the post-momentum weight decayed param but only used it for a dtype cast, so the decay was never applied. Fold the decay into the update instead, matching torch's param.mul_(1 - lrwd). test_muon_wd passed anyway since lrwd=1e-5 is far below atol, so also bump the test's weight_decay to 10 to actually exercise it. * muon: apply weight decay after lr scaling keeps the decoupled decay independent of the LARS trust ratio r, matching torch's param.mul_(1 - lr*wd). no behavior change today since r is always 1.0 on the pre_wd=False path (Muon has tcoef=0).	2026-08-24 13:33:10 -07:00
nimlgenandGitHub	7a887e84b2	am: use gc regs for xgmi (#17707 )	2026-08-24 22:42:35 +03:00
chenyuandGitHub	7b21ffac00	PYLITERAL has dtypes.void (#17705 ) also deleted Ops.WIAT, it currently does not have entry in dtype_from_uop	2026-08-24 14:37:15 -04:00
sirhcmandGitHub	9b481a4893	dsp: cleanup temp linker script (#17704 )	2026-08-24 14:05:00 -04:00
qazalandGitHub	9b0c65688d	llama script changes from the speedups branch (#17701 )	2026-08-24 21:10:02 +09:00