const are weak 3 [pr]

* mergable fast RDNA3 Qwen 3.6

* AMD

* quant 256 multiple

* cleanup cast

* llm kernels: adapt to Ops.BIND removal

Variables are 0-d ALU BUFFERs in the tensor graph and take the ALU PARAM form
inside kernels (UOp.variable(param=True)). Add kernel_var helper for the
conversion, and keep start_pos in bound form at the graph level so function
implicit-input collection and the schedule's binds rename-back line up.

* adaptive prefill chunk sizes for recurrent models + iq4xs model entry

one TinyJit per static prefill chunk size: capture 128 and 32 at warmup,
generate picks the largest that fits the remaining prompt. long prompts
prefill 2x faster (555 tok/s on Qwen3.6-27B IQ4_XS) without pushing short
prompts through token-by-token decode.

* minimize diff: early-return custom attention path, keep master state init

* minimize: master _attention with gated fused-scan swap, kernels/amd only, single chunk size

- GatedDeltaNetBlock._attention keeps master's symbolic-padding structure;
  the recurrent scan is swapped for the fused gated_delta_prefill kernel only
  on RDNA3 with static shapes (fast_scan), everything else uses the old path
- all AMD kernel code lives in tinygrad/llm/kernels/amd.py (drop kernels/__init__.py,
  drop the generic fallback kernel - the old scan covers non-RDNA3)
- single prefill chunk size 32; non-RDNA3 recurrent keeps master's chunk_size=1
- the conv+normalize miscompile doesn't trigger with master's window-buffer conv,
  so the contiguous workaround is dropped

* warmup: single code path for fast and old recurrent

* drop fast_scan/fast_recurrent flags, inline the RDNA3 gate (cached)

* generate: chunk size is always 32, no device gating

static chunks for recurrent models everywhere: the fused kernel path on RDNA3,
the old scan elsewhere (which is also faster chunked than token-by-token)

* warmup: drop redundant _init_state loop (lazy init in the eager step covers it)

* symbolic-length prefill with the custom kernels

the prefill path is fully symbolic again (master's generate, one prefill graph
for every chunk size, no static-tail decode): padded steps are exact no-ops in
the scan (beta=0, alpha=exp(0)=1), flash attention positions queries at
start_pos instead of valid_kv_len-M, and quant linears pad to the chunk bucket

prefill 401 tok: 284 -> 348 tok/s on Qwen3.8-27B IQ4_XS (tail chunks no longer
decode token-by-token), decode unchanged at 45 tok/s

* cli: default qwen3.6:27b to the fast IQ4_XS quant, add qwen3.8:27b

Q4_K_M falls back to slow inline dequant with the custom kernels, IQ4_XS is
the fast path. qwen3.8 quants use unsloth's UD (dynamic) naming

* warmup: back to master's two-liner plus a cache reset

with symbolic prefill, generate([0])'s 1-token chunk captures the symbolic
prefill graph that serves every chunk size, and JIT batching on capture
measurably doesn't matter with the fused kernels (347.7 tok/s either way)

* cli: pin qwen3.8:27b to the pre-UD revision

the UD-IQ4_XS replacement mixes in Q3_K tensors (ggml type 11) the loader
doesn't support; the pinned revision is byte-identical to the known-good file

* warmup: identical to master

the leftover cache is self-consistent: get_start_pos only reuses a full
strict-prefix match, everything else restarts with a state reset

* model: hoist the quantized_attention import to the top level

* hoist the GDN query scale out of the branch, restore master dtype.py

the scale is the same op in both paths, apply it once after the transpose.
the dtype.py diff was a stale pre-SPEC=2 copy, not intentional work

* gated_delta_prefill: don't pass the bound start_pos as a call src

device-less param buffers in call srcs crash hcq2's _get_enqueue_devs. the
var already reaches the graph through the state AFTER chain (conv state
store), same as the flash kernels' valid_end

* llm: half KV cache with custom flash kernels, drop the int8 quantized cache

matches master's new half cache default: no scales, no packing, one less
buffer. the store casts to half explicitly (buffer-only half usage misses
the renderer's half define). 45.5 tok/s decode, 348.7 tok/s prefill —
same as int8

* llm: zero-init the KV cache

the int8 path was accidentally protected from uninitialized memory by its
zero-initialized scale buffer; with a plain half cache the flash prefill
kernel's P*V wmma computes 0*NaN=NaN on masked lanes past the valid region
(manifested as garbage tokens at 32k context where the allocator reuses
dirty VRAM)

* gate that

* llm/kernels/amd: reorganize by kernel family, drop the clutter

sections: shared helpers, quant linear, flash attention, gated delta prefill.
no AxisType.WEAK (default), no ALLOW_DEVICE_USAGE override (unneeded), magic
numbers become names (QUANT_SIZES, Q5_K/Q6_K/IQ4_XS), merged wrapper layers
(flash_attention_causal_cached folded into flash_attention), one _unbind
helper for the bound-var dance

* test: universal recurrent reuse assertion, fix lambda lint

* 1-token chunks have a static shape: they are decode steps

a 1-token chunk routes to the decode graph via the existing dispatch, so
warmup and decode-only workloads never build the big symbolic prefill graph:
CI benchmark command 12m50s -> 5m29s (master: 6m48s), 220 -> 123 compile jobs

also restores the ALLOW_DEVICE_USAGE override in amd_custom_kernels_supported:
Device[] asserts inside @function contexts (ALLOW_DEVICE_USAGE=0), and the
first gate call can happen there depending on test order

* generate: back to plain symbolic binding, the static-1 rule wasn't worth it

* custom kernels: Q4_K support (ggml type 12)

Q4_K is Q5_K without the high-bit array: same d/dmin/scales layout (so
_q5_scales works unchanged), 144-byte blocks, qs at word 4. both the dp4a
decode kernel and the WMMA prefill kernel take a ggml_type branch now.

Qwen3-8B Q4_K_M: decode 16.5 -> 114.8 tok/s, prefill 69 -> 536 tok/s

* raise line count to 26500 (qwen did it)

* benchmark qwen3.8

* little updates

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

@@ -61,6 +61,7 @@ runs:
         echo "MAX_BUFFER_SIZE=300000000" >> "$GITHUB_ENV"
         if [[ "$RUNNER_OS" == "Linux" ]]; then
           echo "VIRTUAL_ENV=/opt/venv/${{ inputs.python-version }}" >> "$GITHUB_ENV"
+          echo "UV_PYTHON_INSTALL_DIR=/opt/python" >> "$GITHUB_ENV"
         else
           echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
         fi
@@ -70,11 +71,6 @@ runs:
       with:
         enable-cache: 'false' # see below for manual caching
 
-    - name: Set up Python ${{ inputs.python-version }}
-      uses: actions/setup-python@v6
-      with:
-        python-version: ${{ inputs.python-version }}
-
     # **** Caching packages ****
 
     - name: Cache Python packages (PR)

.github/workflows/benchmark.yml

@@ -94,7 +94,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: "0"
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -117,10 +117,10 @@ jobs:
       run: python3 test/external/process_replay/reset.py
     - name: Run llama3.2
       run: BENCHMARK_LOG=llama32_3b-f16 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m llama3.2:3b-f16 --benchmark --warmup
-    - name: Run qwen3.6
-      # qwen3.6:35b-a3b doesn't fit on mac
+    - name: Run qwen3.8
+      # qwen3.8:27b doesn't fit on mac
       if: ${{ matrix.dev != 'METAL' }}
-      run: BENCHMARK_LOG=qwen36_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.6:35b-a3b --benchmark --warmup
+      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' }}
@@ -141,7 +141,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: "0"
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -190,7 +190,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: "0"
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -233,7 +233,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: "0"
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -279,7 +279,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: "0"
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -634,6 +634,9 @@ jobs:
       run: |
         GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
         GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
+    - name: HEVC Decode Benchmark
+      if: ${{ matrix.dev == 'NV' }}
+      run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
     - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
       if: ${{ matrix.dev == 'NV' }}
       run: BENCHMARK_LOG=resnet_10steps MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py

.github/workflows/test.yml

@@ -233,7 +233,7 @@ jobs:
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
     - name: Repo line count <= 26000 lines
-      run: MAX_LINE_COUNT=26000 python sz.py
+      run: MAX_LINE_COUNT=26500 python sz.py
 
   spec:
     strategy:
@@ -504,7 +504,7 @@ jobs:
       - name: Run AMD renderer tests (AMD:LLVM)
         run: DEV=MOCKKFD+AMD:LLVM python -m pytest -n=auto test/amd/ --durations 20
       - name: Run SQTT profiling tests
-        run: PROFILE=1 SQTT=1 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
+        run: VIZ=-2 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
       - name: Run AMD emulated tests on NULL backend
         env:
           AMD: 0
@@ -679,4 +679,5 @@ jobs:
         run: |
           python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
           DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
-          python -m pytest -n=auto test/backend/test_ops.py --durations=20
+          # QCOMCL compiles in qemu, too slow for parallel workers
+          ${{ contains(matrix.dev, 'QCOMCL') && 'PARALLEL=0' || '' }} python -m pytest -n=auto test/backend/test_ops.py --durations=20

.gitignore

@@ -69,3 +69,4 @@ mutants
 dagre/
 graphlib/
 uv.lock
+pi_session_window0.jsonl

examples/mlperf/model_train.py

@@ -1773,8 +1773,13 @@ def train_gptoss():
   def minibatch(tokens:Tensor):
     if is_dp: tokens = tokens.to(None).shard(device, 0)
     if not is_sharding: tokens = tokens.to(None)
+
     logits:Tensor = model(tokens[:, :-1], save=True)
-    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
+    if getenv("FUSED_CE", 0):
+      from extra.llama_kernels.fused_ce import fused_ce_loss
+      loss = fused_ce_loss(logits.cast(dtypes.bfloat16), tokens[:, 1:], label_smoothing=0.0)
+    else:
+      loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
 
     for g, new_g in zip(grads, loss.gradient(*optim.params)):
       apply_grad(g, new_g.uop)

examples/openpilot/compile3.py

@@ -107,14 +107,21 @@ def compile(onnx_file):
   return inputs, test_val
 
 def test_vs_compile(run, inputs, test_val=None):
+  if (log:=bool(getenv("BENCHMARK_LOG", ""))): from extra.bench_log import WallTimeEvent, BenchEvent
 
   # run 20 times
   step_times = []
   for _ in range(20):
     st = time.perf_counter()
-    out = run(**inputs)
-    mt = time.perf_counter()
-    val = out.numpy()
+    if log:
+      with WallTimeEvent(BenchEvent.STEP):
+        out = run(**inputs)
+        mt = time.perf_counter()
+        val = out.numpy()
+    else:
+      out = run(**inputs)
+      mt = time.perf_counter()
+      val = out.numpy()
     et = time.perf_counter()
     step_times.append((et-st)*1e3)
     print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")
@@ -160,12 +167,6 @@ def test_vs_onnx(new_inputs, test_val, onnx_file, tol):
   print("test vs onnx passed")
   return timings
 
-def bench(run, inputs):
-  from extra.bench_log import WallTimeEvent, BenchEvent
-  for _ in range(10):
-    with WallTimeEvent(BenchEvent.STEP):
-      run(**inputs).numpy()
-
 if __name__ == "__main__":
   if getenv("RUN_PICKLE"):
     with open(OUTPUT, "rb") as f: pickle_loaded = load_pickle(f)
@@ -181,6 +182,3 @@ if __name__ == "__main__":
     test_vs_compile(pickle_loaded, inputs, outputs)
     if getenv("SELFTEST"):
       test_vs_onnx(inputs, outputs, onnx_file, 1e-4)
-
-  if getenv("BENCHMARK_LOG", ""):
-    bench(pickle_loaded, inputs)

extra/hcq2/ops_amd2.py

@@ -19,7 +19,7 @@ from tinygrad.runtime.support.hcq import FileIOInterface, HCQBuffer, MMIOInterfa
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
 from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_pmc
 from tinygrad.runtime.support.system import PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
-from tinygrad.runtime.support.usb import USB3
+from tinygrad.runtime.support.usb import USB3, usb_ib, usb_push, usb_arm_bytes, pm_usb_stage, pm_usb_hostio, pm_usb_bufferize
 from tinygrad.runtime.support.memory import AddrSpace, BumpAllocator
 from tinygrad.runtime.ops_amd import SQTT, SQTT_ITRACE_SE_MASK, SQTT_LIMIT_SE, SQTT_SIMD_SEL, SQTT_TOKEN_EXCLUDE, PMC
 from tinygrad.runtime.ops_amd import EVENT_INDEX_PARTIAL_FLUSH, WAIT_REG_MEM_FUNCTION_EQ, WAIT_REG_MEM_FUNCTION_NEQ, WAIT_REG_MEM_FUNCTION_GEQ
@@ -146,11 +146,14 @@ pm_pm4_opsel = PatternMatcher([
   (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
 ])
 
+def queue_ptrs(devs, qname:str, q:AMDQueueDesc) -> tuple[UOp, ...]:
+  return tuple(UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"{qname}_{n}")
+               for n, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
+
 def pm4_submit(ctx, lin):
   # ensure compute queues are allocated
   for d in (devs:=ctx.devs): q = Device[d].compute_queue
-  ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COMPUTE:0_{name}")
-    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
+  ring, wptr, doorbell, put_ptr = queue_ptrs(devs, "COMPUTE:0", q)
 
   # the host fence at the start of the batch guarantees the ib is free to reuse
   size_dw = sum(len(ins.src) for ins in lin.src)
@@ -216,8 +219,7 @@ def sdma_submit(cmdbuf, devs):
 
   # the sdma queue's ring and its host-side ring/write/put pointers
   for d in devs: q = Device[d].sdma_queue(0)
-  ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COPY:0_{name}")
-    for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
+  ring, wptr, doorbell, put_ptr = queue_ptrs(devs, "COPY:0", q)
 
   # sdma needs the cmdbuf contiguous: if it won't fit before the ring end, restart at 0 and zero the tail
   put_b = put_ptr.index(zero)
@@ -244,15 +246,32 @@ def sdma_submit(cmdbuf, devs):
 pm_sdma_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"),
   lambda ctx, lin: sdma_submit(make_cmdbuf(lin, ctx.devs), ctx.devs))])
 
+# *****************
+# USB submit
+
+def amd_usb_submit(ctx, lin):
+  for d in ctx.devs: q = Device[d].compute_queue if (comp:=ctx.qname.startswith("COMPUTE")) else Device[d].sdma_queue(0)
+
+  if nb:=usb_arm_bytes(ctx.pre, Device[ctx.devs[0]].iface.usb_sram):
+    poke = (ctx.sdma.SDMA_OP_WRITE, *data64_le(Device[ctx.devs[0]].iface.cq_buf.va_addr + 12), 0, 0)
+    lin = lin.replace(src=lin.src + (UOp(Ops.INS, arg="poke", src=tuple(UOp.const(x, dtypes.uint32) for x in poke)),))
+
+  ib_host, ib_gpu, pkt_dw = usb_ib(ctx.devs, lin, 32 if comp else 0x100, nb)
+  pkt = (ctx.pm4.PACKET3(ctx.pm4.PACKET3_INDIRECT_BUFFER,2),*data64_le(ib_gpu.getaddr(ctx.devs)),pkt_dw|ctx.pm4.INDIRECT_BUFFER_VALID) if comp else ()
+  return usb_push(ctx.devs, *queue_ptrs(ctx.devs, ctx.qname, q), ib_host, ib_gpu, pkt, 4 if comp else 1)
+
+pm_usb_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"), amd_usb_submit)])
+
 @dataclass(frozen=True)
 class AMDEncodeCtx:  # encode-time constants for one queue: devs (every cmdbuf address resolves into these) + gfx version + packet/ip modules
   devs: tuple[str, ...]; target: tuple[int, ...]; pm4: Any; sdma: Any; soc: Any  # noqa: E702
-  gc: AMDIP; nbio: AMDIP; xccs: int; max_copy_size: int; tmpring_size: Callable  # noqa: E702
+  gc: AMDIP; nbio: AMDIP; xccs: int; max_copy_size: int; tmpring_size: Callable; qname: str; pre: UOp # pre: the queue before opsel
 
 def encode_queue(q:UOp) -> UOp|None:
   d = Device[(devs:=to_tuple(q.arg[0]))[0]]
-  ctx = AMDEncodeCtx(devs, d.target, d.pm4, d.sdma, d.soc, d.gc, d.nbio, d.xccs, d.max_copy_size, d.tmpring_size)
-  opsel, submit = (pm_pm4_opsel, pm_pm4_submit) if q.arg[1].startswith("COMPUTE") else (pm_sdma_opsel, pm_sdma_submit)
+  ctx = AMDEncodeCtx(devs, d.target, d.pm4, d.sdma, d.soc, d.gc, d.nbio, d.xccs, d.max_copy_size, d.tmpring_size, q.arg[1], q)
+  opsel = pm_pm4_opsel if (comp:=q.arg[1].startswith("COMPUTE")) else pm_sdma_opsel
+  submit = d.pm_submit if d.pm_submit is not None else (pm_pm4_submit if comp else pm_sdma_submit)
   return submit.rewrite(graph_rewrite(q, opsel + pm_flatten_linear, walk=True, ctx=ctx, name=f"{q.arg[1]} opsel"), ctx)
 
 @dataclass(frozen=True)
@@ -282,13 +301,14 @@ def amd_build_program(prg:UOp) -> UOp:
       wave32=bool(desc.kernel_code_properties & 0x400), private_segment_size=desc.private_segment_fixed_size, kernargs_segment_size=desc.kernarg_size,
       kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0), enable_dispatch_ptr=edp,
       enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER)
+    image = bytes(image).ljust(round_up(len(image), 4), b"\x00") # the program is uploaded as whole dwords
     buf = UOp.placeholder((len(image),), dtypes.uint8, next(UOp.unique_num), device=prg.device).rtag("program")
-    cached = _amd_program_cache[key] = prg.replace(src=(buf.after(make_binary_patch(buf, bytes(image))),), arg=(data, prg.arg))
+    cached = _amd_program_cache[key] = prg.replace(src=(buf.after(make_binary_patch(buf, image)),), arg=(data, prg.arg))
   return cached
 
 class AMDAllocator(HCQAllocator['AMDDevice']):
   def __init__(self, dev:AMDDevice):
-    super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb())
+    super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb)
 
   def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
     return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_copy_queue)
@@ -524,8 +544,7 @@ class PCIIface(PCIIfaceBase):
       cq = d.compute_queue
       for b in (cq.put_value, cq.read_ptr, cq.write_ptr): b._buf.view.view(fmt='Q')[0] = 0
       d.iface.dev_impl.gfx.setup_ring(*cq.params)
-      d.signal('timeline')._buf.cpu_view().mv.cast('Q')[0] = \
-        d.signal('value', 1).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] - 1
+      d.signal('timeline')._buf.cpu_view().view(fmt='Q')[0] = d.signal('value', 1, device="CPU")._buf.cpu_view().view(fmt='Q')[0] - 1
 
   def sleep(self, timeout):
     if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
@@ -539,6 +558,32 @@ class PCIIface(PCIIfaceBase):
 
   def device_fini(self): self.dev_impl.fini()
 
+class USBIface(PCIIface):
+  def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called
+    if dev_id >= len(visible:=hcq_filter_visible_devices(USB3.list_devices(0xADD1, 0x0001) + USB3.list_devices(0x3801, 0x0001), "AMD")):
+      raise RuntimeError(f"AMD:{dev_id} does not exist ({pluralize('device', len(visible))} available)")
+    self.dev, self.pci_dev, self.vram_bar, self.count = dev, USBPCIDevice("AM", *visible[dev_id]), 0, len(visible)
+    self.dev_impl = AMDev(self.pci_dev)
+    self._compute_props()
+    self.sram = self._dma_region(ctrl_addr=0xf000, sys_addr=0x200000, size=0x80000)
+    self.cq_buf = self._dma_region(ctrl_addr=0xb800, sys_addr=0x822000, size=0x1000) # +12 is the dword that releases an armed read
+    self.usb_handle = unwrap(ctypes.cast(self.pci_dev.usb.usb.handle, ctypes.c_void_p).value)
+
+  def _dma_region(self, ctrl_addr, sys_addr, size):
+    region = self.dev_impl.mm.map_range(vaddr:=self.dev_impl.mm.alloc_vaddr(size=size), size, [(sys_addr, size)], aspace=AddrSpace.SYS, uncached=True)
+    return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(region, has_cpu_mapping=False), view=self.pci_dev.dma_view(ctrl_addr, size), owner=self.dev)
+
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, **kwargs) -> HCQBuffer:
+    # everything, even host-style signals, lives in vram: gpu writes into the bridge's own memory collide with an armed 0xF2 read stream
+    return super().alloc(size, host=False, uncached=uncached, cpu_access=cpu_access or host, contiguous=contiguous, force_devmem=True, **kwargs)
+
+  def sleep(self, timeout): pass
+
+  # we don't own the sram region, so the buffer never frees it
+  @functools.cached_property
+  def usb_sram(self) -> Buffer:
+    return Buffer(self.dev.device, (b:=self.sram).size, dtypes.uint8, options=BufferSpec(external_ptr=b.va_addr, nolru=True)).allocate(opaque=b)
+
 def _mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface,), {})
 
 class AMDDevice(HCQ2Compiled):
@@ -549,19 +594,21 @@ class AMDDevice(HCQ2Compiled):
     # encoding of cmdbuf
     (UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),)), encode_queue),
   ])
+  pm_submit: PatternMatcher|None = None
 
   timestamp_divider = 100.0  # AMD GPU clock: ticks/us
   max_scratch_psize = 0
 
-  ifaces = [KFDIface, PCIIface, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface)]
+  ifaces = [KFDIface, PCIIface, USBIface, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface), _mock(USBIface)]
 
   def device_props(self): return self.iface.props
 
   def is_am(self) -> bool: return isinstance(self.iface, (PCIIface,))
-  def is_usb(self) -> bool: return False
 
   def __init__(self, device:str=""):
     self.iface = self._select_iface(device)
+    self.is_usb = isinstance(self.iface, USBIface)
+    if self.is_usb: self.rt_nbytes = 4 << 20
 
     self.target:tuple[int, ...] = ((trgt:=self.iface.props['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
     self.arch = "gfx%d%x%x" % self.target
@@ -586,7 +633,7 @@ class AMDDevice(HCQ2Compiled):
 
     self.is_aql = getenv("AMD_AQL", int(self.xccs > 1))
     if self.is_aql:
-      self.pm4_ibs = self.iface.alloc(0x2000 if self.is_usb() else (16 << 20), uncached=True, cpu_access=True)
+      self.pm4_ibs = self.iface.alloc(0x2000 if self.is_usb else (16 << 20), uncached=True, cpu_access=True)
       self.pm4_ib_alloc = BumpAllocator(self.pm4_ibs.size, wrap=True)
 
     self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
@@ -599,6 +646,10 @@ class AMDDevice(HCQ2Compiled):
     self.max_private_segment_size = 0
     self.pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, tag="scratch", name="b"), lambda ctx, b: ctx[0].scratch_buffer(b.max_numel()))]) + self.pm_bufferize
 
+    if self.is_usb:
+      self.pm_bufferize = pm_usb_bufferize + self.pm_bufferize
+      self.pm_stage_copy, self.pm_host_lower, self.pm_submit = pm_usb_stage, pm_usb_hostio, pm_usb_submit
+
     self.pmc_enabled:bool = PROFILE > 0 and PMC > 0
     if self.pmc_enabled:
       self.iface.require_profile_mode()
@@ -659,7 +710,7 @@ class AMDDevice(HCQ2Compiled):
     wg_data_size = round_up((vgpr_size_per_cu + sgrp_size_per_cu + lds_size_per_cu + hwreg_size_per_cu) * self.cu_cnt, mmap.PAGESIZE)
     ctl_stack_size = round_up((12 if self.target[0] != 9 else 8) * self.wave_cnt + 8 + 40, mmap.PAGESIZE)
     return self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL if self.is_aql else kfd.KFD_IOC_QUEUE_TYPE_COMPUTE,
-      0x2000 if self.is_usb() else (16 << 20), eop_buffer_size=0x1000,
+      0x2000 if self.is_usb else (16 << 20), eop_buffer_size=0x1000,
       ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size,
       debug_memory_size=round_up(self.wave_cnt * 32, 64))
 
@@ -667,7 +718,7 @@ class AMDDevice(HCQ2Compiled):
     if getenv("AMD_DISABLE_SDMA"): return None
     if idx in self.sdma_queues: return self.sdma_queues[idx]
     with contextlib.suppress(OSError):
-      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20), idx=idx)
+      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x2000 if self.is_usb else (16 << 20), idx=idx)
     return self.sdma_queues.get(idx, None)
 
   def tmpring_size(self, private_segment_size):

extra/models/llama.py

@@ -5,9 +5,9 @@ from tinygrad.helpers import getenv, DEBUG
 
 # https://github.com/facebookresearch/llama/blob/1076b9c51c77ad06e9d7ba8a4c6df775741732bd/llama/model.py#L47
 def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
-  freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2)[:(dim // 2)] / dim))
-  freqs = Tensor.arange(end).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
-  return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).reshape(1, end, 1, dim//2, 2)
+  freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2, dtype=dtypes.float32)[:(dim // 2)] / dim))
+  freqs = Tensor.arange(end, dtype=dtypes.float32).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
+  return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).cast(dtypes.default_float).reshape(1, end, 1, dim//2, 2)
 
 # matches meta, non hugging face weights
 # (a+i*b) * (c+i*d) = (ac-bd) + i*(ad+bc)

extra/sqtt/roc.py

@@ -192,7 +192,7 @@ def unpack_insts(viz_data, i:int, j:int, data:dict) -> dict:
     prev_instr = max(prev_instr, e.time + e.dur)
   summary = [{"label":"Total Cycles", "value":w.end_time-w.begin_time}, {"label":"SE", "value":w.se}, {"label":"CU", "value":w.cu},
              {"label":"SIMD", "value":w.simd}, {"label":"Wave ID", "value":w.wave_id}, {"label":"Run number", "value":data["run_number"]}]
-  return {"rows":[tuple(v.values()) for v in rows.values()], "cols":columns, "metadata":[summary], "ref":viz_data.ref_map.get(data["prg"].name)}
+  return {"rows":[tuple(v.values()) for v in rows.values()], "cols":columns, "metadata":[summary],"ref":viz_data.ref_map.get(data["prg"].profile_key)}
 
 def print_data(data:dict) -> None:
   from tabulate import tabulate

test/amd/hw/test_sop.py

@@ -1002,6 +1002,39 @@ class TestBarrier(unittest.TestCase):
     for tid in range(64):
       self.assertEqual(st.vgpr[tid][0], tid + 100 + 1000, f"tid={tid}")
 
+class TestSMaxMinSCCRegressions(unittest.TestCase):
+  """Regression test: S_MAX sets SCC only on strict inequality (equal operands -> SCC=0)."""
+
+  def test_s_max_i32_equal_scc(self):
+    st = run_program([s_mov_b32(s[4], 64), s_mov_b32(s[5], 64), s_max_i32(s[6], s[4], s[5])], n_lanes=1)
+    self.assertEqual(st.scc, 0)
+    self.assertEqual(st.sgpr[6], 64)
+    st = run_program([s_mov_b32(s[4], 65), s_mov_b32(s[5], 64), s_max_i32(s[6], s[4], s[5])], n_lanes=1)
+    self.assertEqual(st.scc, 1)  # still set when strictly greater
+
+  def test_s_max_u32_equal_scc(self):
+    st = run_program([s_mov_b32(s[4], 64), s_mov_b32(s[5], 64), s_max_u32(s[6], s[4], s[5])], n_lanes=1)
+    self.assertEqual(st.scc, 0)
+
+class TestAbsdiffOverflowRegressions(unittest.TestCase):
+  """Regression test: S_ABSDIFF_I32 computes abs on the WRAPPED 32-bit difference (found by random difftest vs hardware)."""
+
+  def test_s_absdiff_wrapped(self):
+    # |45 - (-2147483647)| overflows int32; hardware takes abs of the wrapped 32-bit difference
+    instructions = [s_mov_b32(s[4], 45), s_mov_b32(s[5], 0x80000001), s_absdiff_i32(s[6], s[4], s[5])]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.sgpr[6], 0x7FFFFFD4)
+    self.assertEqual(st.scc, 1)
+    # INT_MIN - 1 wraps to +2147483647, already positive
+    instructions = [s_mov_b32(s[4], 0x80000000), s_mov_b32(s[5], 1), s_absdiff_i32(s[6], s[4], s[5])]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.sgpr[6], 0x7FFFFFFF)
+    # equality -> 0 and SCC=0
+    instructions = [s_mov_b32(s[4], 7), s_mov_b32(s[5], 7), s_absdiff_i32(s[6], s[4], s[5])]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.sgpr[6], 0)
+    self.assertEqual(st.scc, 0)
+
 
 if __name__ == '__main__':
   unittest.main()

test/amd/hw/test_vop1.py

@@ -1629,5 +1629,66 @@ class TestSwap(unittest.TestCase):
     self.assertEqual(st.vgpr[0][1], 0x55555555)
 
 
+class TestCvtFrexpRegressions(unittest.TestCase):
+  """Regression tests for float<->int conversion and FREXP corner cases (found by random difftest vs hardware)."""
+
+  def test_cvt_i32_f32_nan_is_zero(self):
+    """v_cvt_i32_f32 of NaN is 0, not INT_MIN (x86 cvttss2si returns INT_MIN)."""
+    for nan in (0x7FC00000, 0xFFC00000, 0x7F800001):
+      st = run_program([v_mov_b32_e32(v[0], nan), v_cvt_i32_f32_e32(v[1], v[0])], n_lanes=1)
+      self.assertEqual(st.vgpr[0][1], 0, f"nan=0x{nan:08x}")
+
+  def test_cvt_i32_f32_positive_overflow(self):
+    """v_cvt_i32_f32 saturates positive overflow/inf to INT_MAX, not INT_MIN."""
+    for bits in (0x7F800000, 0x4F000000, 0x4F800000):  # +inf, 2^31, ~2^32
+      st = run_program([v_mov_b32_e32(v[0], bits), v_cvt_i32_f32_e32(v[1], v[0])], n_lanes=1)
+      self.assertEqual(st.vgpr[0][1], 0x7FFFFFFF, f"bits=0x{bits:08x}")
+
+  def test_cvt_i32_f32_negative_overflow(self):
+    """v_cvt_i32_f32 saturates negative overflow/-inf to INT_MIN."""
+    for bits in (0xFF800000, 0xCF000001):  # -inf, below -2^31
+      st = run_program([v_mov_b32_e32(v[0], bits), v_cvt_i32_f32_e32(v[1], v[0])], n_lanes=1)
+      self.assertEqual(st.vgpr[0][1], 0x80000000, f"bits=0x{bits:08x}")
+
+  def test_cvt_u32_f32_nan_is_zero(self):
+    """v_cvt_u32_f32 of NaN is 0, not UINT_MAX."""
+    for nan in (0x7FC00000, 0xFFC00000, 0x7F800001):
+      st = run_program([v_mov_b32_e32(v[0], nan), v_cvt_u32_f32_e32(v[1], v[0])], n_lanes=1)
+      self.assertEqual(st.vgpr[0][1], 0, f"nan=0x{nan:08x}")
+
+  def test_cvt_i32_f64_nan_and_overflow(self):
+    """v_cvt_i32_f64: NaN -> 0, positive overflow/+inf -> INT_MAX."""
+    st = run_program([v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0x7FF80000), v_cvt_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0)
+    st = run_program([v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0x41F00000), v_cvt_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x7FFFFFFF)  # 2^32 -> INT_MAX
+
+  def test_frexp_f32_denormal(self):
+    """v_frexp_exp/mant_f32 of denormal/zero inputs is (0, signed zero) on hardware."""
+    for bits in (0x00000001, 0x007FFFFF, 0x00000000):
+      st = run_program([v_mov_b32_e32(v[0], bits), v_frexp_exp_i32_f32_e32(v[1], v[0]), v_frexp_mant_f32_e32(v[2], v[0])], n_lanes=1)
+      self.assertEqual(st.vgpr[0][1] & 0xFFFFFFFF, 0, f"exp bits=0x{bits:08x}")
+      self.assertEqual(st.vgpr[0][2], bits & 0x80000000, f"mant bits=0x{bits:08x}")
+    # negative denormal: mant is -0.0
+    st = run_program([v_mov_b32_e32(v[0], 0x80000001), v_frexp_mant_f32_e32(v[2], v[0])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x80000000)
+
+  def test_frexp_f64_denormal(self):
+    """v_frexp_exp_f64 of a denormal returns the normalized exponent (-1073 for min-denormal); zero -> 0."""
+    st = run_program([v_mov_b32_e32(v[0], 1), v_mov_b32_e32(v[1], 0), v_frexp_exp_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][2] & 0xFFFFFFFF, 0xFFFFFBCF)  # -1073
+    st = run_program([v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_frexp_exp_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0)
+
+  def test_frexp_exp_inf_nan(self):
+    """v_frexp_exp of +/-inf and NaN is 0 on hardware (host frexp gives 129/1024), for both f32 and f64."""
+    for bits in (0x7F800000, 0xFF800000, 0x7FC00000):
+      st = run_program([v_mov_b32_e32(v[0], bits), v_frexp_exp_i32_f32_e32(v[1], v[0])], n_lanes=1)
+      self.assertEqual(st.vgpr[0][1] & 0xFFFFFFFF, 0, f"f32 bits=0x{bits:08x}")
+    for lo, hi in ((0, 0x7FF00000), (0, 0xFFF00000), (0, 0x7FF80000), (1, 0x7FF00000)):
+      st = run_program([v_mov_b32_e32(v[0], lo), v_mov_b32_e32(v[1], hi), v_frexp_exp_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
+      self.assertEqual(st.vgpr[0][2] & 0xFFFFFFFF, 0, f"f64 bits=0x{hi:08x}{lo:08x}")
+
+
 if __name__ == '__main__':
   unittest.main()

test/amd/hw/test_vop2.py

@@ -989,6 +989,53 @@ class TestCarryOps(unittest.TestCase):
     self.assertEqual(st.vgpr[0][0], 0)  # 0xFFFFFFFF + 1 + 0 = 0 (overflow)
     self.assertEqual(st.vcc, 0xDEADBEEF)  # VCC unchanged - carry was discarded
 
+class TestSelectFlushRegressions(unittest.TestCase):
+  """Regression tests: f32 MIN/MAX flush denormal inputs to signed zero (select-style ops propagate inputs bitwise)."""
+
+  def test_v_min_f32_denormal_flush(self):
+    """min(denormal, 1.0) is +0, min(-denormal, -1.0) is -0."""
+    st = run_program([v_mov_b32_e32(v[0], 0x00000001), v_mov_b32_e32(v[1], 0x3F800000), v_min_f32_e32(v[2], v[0], v[1])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x00000000)
+    # flush(-denormal) = -0.0 > -1.0, so the result is -1.0 (both operand orders)
+    st = run_program([v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 0xBF800000), v_min_f32_e32(v[2], v[0], v[1])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xBF800000)
+    st = run_program([v_mov_b32_e32(v[1], 0xBF800000), v_mov_b32_e32(v[2], 0x80000001), v_min_f32_e32(v[3], v[1], v[2])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][3], 0xBF800000)
+
+  def test_v_max_f32_denormal_flush(self):
+    """max(-denormal, -1.0) is -0; max(+denormal, -0) is +0."""
+    st = run_program([v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 0xBF800000), v_max_f32_e32(v[2], v[0], v[1])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x80000000)
+    st = run_program([v_mov_b32_e32(v[0], 0x00000001), v_mov_b32_e32(v[1], 0x80000000), v_max_f32_e32(v[2], v[0], v[1])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x00000000)
+
+
+class TestCarryExecRegressions(unittest.TestCase):
+  """Regression tests: per-lane VCC writes (carry ops) zero inactive lane bits - VCC = mask & EXEC, never preserved."""
+
+  def test_co_ci_e32_vcc_masked_by_exec(self):
+    """v_sub_co_ci_u32_e32 with EXEC=0xFFFF0000: hw clears inactive VCC bits instead of preserving them."""
+    instructions = [
+      s_mov_b32(EXEC_LO, 0xFFFF0000),
+      s_mov_b32(VCC_LO, 0xFFFFFFFF),  # preset all bits
+      v_mov_b32_e32(v[0], 0xFFFFFFFE), v_mov_b32_e32(v[1], 0x80000000),
+      v_sub_co_ci_u32_e32(v[2], v[0], v[1]),  # active lanes: no borrow
+    ]
+    st = run_program(instructions, n_lanes=32)
+    self.assertEqual(st.vcc, 0x00000000)
+
+  def test_co_ci_e32_vcc_masked_by_exec_ones(self):
+    """Same with all-ones carry: VCC = borrow_mask & EXEC."""
+    instructions = [
+      s_mov_b32(EXEC_LO, 0x0F0F0F0F),
+      s_mov_b32(VCC_LO, 0),
+      v_mov_b32_e32(v[0], 0xFFFFFFFF), v_mov_b32_e32(v[1], 1),
+      v_add_co_ci_u32_e32(v[2], v[0], v[1]),  # all lanes would carry if active
+    ]
+    st = run_program(instructions, n_lanes=32)
+    self.assertEqual(st.vcc, 0x0F0F0F0F)
+    self.assertEqual(st.vgpr[31][2], 0)  # 0xFFFFFFFF + 1 wraps to 0 in active lanes
+
 
 if __name__ == '__main__':
   unittest.main()

test/amd/hw/test_vop3.py

@@ -4,6 +4,7 @@ Includes: v_fma_f32, v_div_scale_f32, v_div_fmas_f32, v_div_fixup_f32,
           v_alignbit_b32, v_bfe_i32, v_mad_u64_u32, v_readlane_b32, v_writelane_b32
 """
 import unittest
+from tinygrad.helpers import OSX
 from test.amd.hw.helpers import *
 
 class TestFMA(unittest.TestCase):
@@ -3264,6 +3265,23 @@ class TestVOP3ClampMAD(unittest.TestCase):
     # 0xFFFF * 2 = 0x1FFFE, low 16 bits = 0xFFFE
     self.assertEqual(st.vgpr[0][3] & 0xFFFF, 0xFFFE, f"expected 0xFFFE, got 0x{st.vgpr[0][3] & 0xFFFF:04x}")
 
+class TestMadNarrowClampRegressions(unittest.TestCase):
+  """Regression tests: mad i16/i24 with clamp saturate to narrow output range (found by random difftest vs hardware)."""
+
+  def test_mad_i16_clamp_sat_max(self):
+    # neg/src-floggled 16-bit mul operands are sign-extended after toggling bit15; sum > INT_MAX saturates
+    instructions = [s_mov_b32(s[4], 1232348160), v_mov_b32_e32(v[3], 0x80000000),
+                    v_mov_b32_e32(v[1], 0x7F7FFFFF), v_mad_i32_i16(v[0], s[4], v[3], v[1], 0, 3, 5, 1)]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][0], 0x7FFFFFFF)
+
+  def test_mad_i24_clamp_sat_min(self):
+    # sext24(-6344704) * sext24(+4210688) << -2^31 saturates to INT_MIN
+    instructions = [s_mov_b32(s[7], 4290772992), v_mov_b32_e32(v[1], 1077936128),
+                    v_mad_i32_i24(v[0], s[7], v[1], v[1], 1, 0, 0, 1)]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][0], 0x80000000)
+
 
 class TestCvtPkF16(unittest.TestCase):
   """Tests for V_CVT_PK_RTZ_F16_F32 - pack two f32 to f16 with round toward zero."""
@@ -3651,6 +3669,80 @@ class TestPermlane(unittest.TestCase):
     self.assertEqual(st.vgpr[21][1], 5)
     self.assertEqual(st.vgpr[31][1], 15)
 
+class TestClampLdExpRegressions(unittest.TestCase):
+  """Regression tests for f32 clamp (-0 -> +0) and ldexp input passthrough."""
+
+  def test_clamp_negative_zero(self):
+    """clmp=1 maps -0.0 to +0.0 (found by random difftest vs hardware)."""
+    instructions = [
+      v_mov_b32_e32(v[0], 0x80000000), v_mov_b32_e32(v[1], 0x80000000),
+      v_add_f32_e64(v[2], v[0], v[1], clmp=1),  # -0 + -0 = -0, clamp -> +0
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x00000000)
+    instructions = [
+      v_mov_b32_e32(v[0], 0x3F800000), v_mov_b32_e32(v[1], 0x80000000),
+      v_min_f32_e64(v[2], v[0], v[1], clmp=1),  # min(1.0, -0) = -0, clamp -> +0
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x00000000)
+
+  def test_ldexp_special_inputs(self):
+    """v_ldexp_f32 of 0/-0/inf/NaN propagates the input instead of computing val * 2**exp (0*inf = NaN on host)."""
+    # -0.0 * 2^INT_MIN = -0.0 (src1 as integer exponent; huge negative)
+    instructions = [v_mov_b32_e32(v[0], 0x80000000), v_mov_b32_e32(v[1], 0x80000000), v_ldexp_f32(v[2], v[0], v[1])]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x80000000)
+    # inf stays inf even with negative exponent
+    instructions = [v_mov_b32_e32(v[0], 0x7F800000), v_mov_b32_e32(v[1], 0xFFFFFF80), v_ldexp_f32(v[2], v[0], v[1])]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x7F800000)
+
+  def test_ldexp_denormal_flush(self):
+    """v_ldexp_f32/f64 flush denormal inputs to signed zero (found by random difftest vs hardware)."""
+    # ldexp(+denorm, 1) = +0, ldexp(-denorm, 250) = -0
+    for src, exp_val, want in [(0x00000001, 1, 0x00000000), (0x80000001, 250, 0x80000000)]:
+      st = run_program([v_mov_b32_e32(v[0], src), v_mov_b32_e32(v[1], exp_val), v_ldexp_f32(v[2], v[0], v[1])], n_lanes=1)
+      self.assertEqual(st.vgpr[0][2], want)
+
+  def test_v_mul_neg_modifier_nan_sign(self):
+    """neg modifier is a pure sign-bit toggle on a NaN operand; result keeps that sign (found by random difftest)."""
+    # mul(normal, NEG(ABS(qNaN))): NaN payload negated in the operand stays negative qNaN
+    instructions = [v_mov_b32_e32(v[0], 0xC96CF47F), v_mov_b32_e32(v[1], 0x7FC00000),
+                    v_mul_f32_e64(v[2], v[0], v[1], s[0], 0, 7, 6)]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0xFFC00000)
+    # plain neg modifier still applies to non-NaN values: mul(-1.0, NEG(2.0)) = +2.0
+    st = run_program([v_mov_b32_e32(v[0], 0xBF800000), v_mov_b32_e32(v[1], 0x40000000),
+                      v_mul_f32_e64(v[2], v[0], v[1], s[0], 0, 2, 0)], n_lanes=1)
+    self.assertEqual(st.vgpr[0][2], 0x40000000)
+
+
+class TestNaNPropagationRegressions(unittest.TestCase):
+  """Regression tests: float arithmetic propagates a NaN from the FIRST NaN operand, quieted with its own sign/payload."""
+
+  @unittest.skipIf(OSX, "broken on mac, TODO: why?")
+  def test_mul_nan_priority(self):
+    # first NaN operand wins (sign+payload), not x86's second-source propagation
+    for a, b, want in [(0x7FC00001, 0x7F800003, 0x7FC00001), (0xFFC00005, 0x7F800003, 0xFFC00005),
+                       (0x7F800001, 0xFFC00005, 0x7FC00001), (0xFF9F1800, 0x7F800001, 0xFFDF1800)]:
+      st = run_program([v_mov_b32_e32(v[0], a), v_mov_b32_e32(v[1], b),
+                        v_mul_f32_e32(v[2], v[0], v[1])], n_lanes=1)
+      self.assertEqual(st.vgpr[0][2], want, f"mul({a:#x}, {b:#x})")
+
+class TestMinMaxFlushE64Regressions(unittest.TestCase):
+  """Regression tests: f32 min/max/median flush denormal inputs to signed zero (e64 forms)."""
+
+  def test_v_min3_f32_denormal_flush(self):
+    st = run_program([v_mov_b32_e32(v[0], 0x00000001), v_mov_b32_e32(v[1], 0x3F800000), v_mov_b32_e32(v[2], 0x40000000),
+                      v_min3_f32(v[3], v[0], v[1], v[2])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][3], 0x00000000)  # min(+denorm, 1, 2) = +0
+
+  def test_v_med3_f32_denormal_flush(self):
+    st = run_program([v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 0x3F800000), v_mov_b32_e32(v[2], 0x40000000),
+                      v_med3_f32(v[3], v[0], v[1], v[2])], n_lanes=1)
+    self.assertEqual(st.vgpr[0][3], 0x3F800000)  # med(-0, 1, 2) = 1
+
 
 if __name__ == '__main__':
   unittest.main()

test/amd/hw/test_vopc.py

@@ -973,6 +973,71 @@ class TestCmpxPartialWavefront(unittest.TestCase):
     self.assertEqual(st.sgpr[EXEC_LO.offset] & 0xFFFFFFFF, 0x4,
                      "Only lane 2 should be active after v_cmpx_eq_u32_e64")
 
+class TestClassDenormalRegressions(unittest.TestCase):
+  """Regression tests: V_CMP_CLASS classifies denormals as DENORMAL (raw bits), not as zero class."""
+
+  def test_class_pos_denormal(self):
+    for bits in (0x00000001, 0x007FFFFF):
+      instructions = [v_mov_b32_e32(v[0], bits), v_mov_b32_e32(v[1], 0x80), v_cmp_class_f32_e64(VCC_LO, v[0], v[1])]
+      st = run_program(instructions, n_lanes=1)
+      self.assertEqual(st.vcc, 1, f"bits=0x{bits:08x}")  # n_lanes=1
+      # ...and it is not the zero class
+      instructions = [v_mov_b32_e32(v[0], bits), v_mov_b32_e32(v[1], 0x40), v_cmp_class_f32_e64(VCC_LO, v[0], v[1])]
+      st = run_program(instructions, n_lanes=1)
+      self.assertEqual(st.vcc, 0, f"bits=0x{bits:08x}")
+
+  def test_class_neg_denormal(self):
+    instructions = [v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 0x10), v_cmp_class_f32_e64(VCC_LO, v[0], v[1])]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vcc, 1)  # n_lanes=1
+    instructions = [v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 0x20), v_cmp_class_f32_e64(VCC_LO, v[0], v[1])]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vcc, 0)  # not the negative-zero class
+
+
+class TestIntCmpModRegressions(unittest.TestCase):
+  """Regression tests: int compares (i32/u32) honor abs/neg as bit-level sign clear/flip (not integer abs/negate)."""
+
+  def test_cmp_i32_abs_neg_bit_level(self):
+    # abs(0x80000001) = 1 -> 1 > 1 is false (integer abs would give 2147483647 > 1)
+    instructions = [v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 1), v_cmp_gt_i32_e64(VCC_LO, v[0], v[1], abs=1)]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vcc, 0)
+    # neg(0x80000001) flips the sign bit -> 1 > 2 is false (integer negate would give 2147483647 > 2)
+    instructions = [v_mov_b32_e32(v[0], 0x80000001), v_mov_b32_e32(v[1], 2), v_cmp_gt_i32_e64(VCC_LO, v[0], v[1], neg=1)]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vcc, 0)
+
+  def test_cmp_u32_abs_bit_level(self):
+    # abs(0x80000000) = 0 -> 0 < 1 is true
+    instructions = [v_mov_b32_e32(v[0], 0x80000000), v_mov_b32_e32(v[1], 1), v_cmp_lt_u32_e64(VCC_LO, v[0], v[1], abs=1)]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vcc, 1)  # n_lanes=1
+
+
+class TestCmpxSdstRegressions(unittest.TestCase):
+  """Regression tests: V_CMPX_*_E64 writes EXEC only, never SDST (hardware verified)."""
+
+  def test_cmpx_e64_no_sdst(self):
+    instructions = [
+      s_mov_b32(VCC_LO, 0),  # preset VCC to 0
+      v_mov_b32_e32(v[0], 0x3F800000), v_mov_b32_e32(v[1], 0x40000000),
+      v_cmpx_lt_f32_e64(VCC_LO, v[0], v[1]),  # 1.0 < 2.0
+    ]
+    st = run_program(instructions, n_lanes=32)
+    self.assertEqual(st.sgpr[EXEC_LO.offset], 0xFFFFFFFF)  # EXEC updated
+    self.assertEqual(st.vcc, 0)  # but VCC untouched
+
+  def test_cmpx_e64_partial_exec(self):
+    instructions = [
+      s_mov_b32(EXEC_LO, 0x0F0F0F0F),
+      s_mov_b32(VCC_LO, 0xFFFFFFFF),
+      v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0x3F800000),
+      v_cmpx_lt_f32_e64(VCC_LO, v[0], v[1]),
+    ]
+    st = run_program(instructions, n_lanes=32)
+    self.assertEqual(st.sgpr[EXEC_LO.offset], 0x0F0F0F0F)  # EXEC = computed & old EXEC
+
 
 if __name__ == '__main__':
   unittest.main()

test/amd/test_sqtt_profiler.py

@@ -1,30 +1,28 @@
 import unittest, contextlib
 from tinygrad import Device, Tensor, Context, TinyJit
-from tinygrad.device import Compiled, ProfileProgramEvent, ProfileDeviceEvent
+from tinygrad.device import Compiled, ProfileProgramEvent
 from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
 from tinygrad.viz.serve import load_amd_counters, VizData
 
 @contextlib.contextmanager
 def save_sqtt():
+  Device[Device.DEFAULT].synchronize()
+  profile_start = len(Compiled.profile_events)
   data = VizData()
   yield data.ctxs
   Device[Device.DEFAULT].synchronize()
   Device[Device.DEFAULT]._at_profile_finalize()
-  load_amd_counters(data, Compiled.profile_events)
+  load_amd_counters(data, [e for e in Compiled.profile_events[:profile_start] if isinstance(e, ProfileProgramEvent)] +
+                          Compiled.profile_events[profile_start:])
   data.ctxs[:] = [r for r in data.ctxs if r["name"].startswith("SQTT")]
 
 @unittest.skipUnless(Device.DEFAULT == "AMD", "only runs on AMD")
 class TestSQTTProfiler(unittest.TestCase):
-  # TODO: can we enable SQTT profiling in context?
   @classmethod
   def setUpClass(cls):
     if not Device[Device.DEFAULT].sqtt_enabled: raise unittest.SkipTest("device must be in SQTT profiling mode")
 
-  def setUp(self):
-    Device[Device.DEFAULT].synchronize()
-    Compiled.profile_events[:] = [e for e in Compiled.profile_events if isinstance(e, (ProfileProgramEvent, ProfileDeviceEvent))]
-
   def test_simple(self):
     t = Tensor.empty(1) + 1
     with save_sqtt() as sqtt:

test/backend/test_llama_kernels.py

@@ -9,7 +9,7 @@ from extra.llama_kernels.swiglu import swiglu
 from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
 from extra.thunder.amd.fa import custom_fused_qkv_rope_backward, fused_qkv_rope
 from test.helpers import needs_second_gpu, assert_kernel_count
-from test.backend.test_asm_gemm import has_hipcc
+from test.backend.test_asm_gemm import has_hipcc, is_cdna4
 
 def run_fused_ce(bs:int, seqlen:int, vocab:int, label_smoothing:float=0.0) -> None:
   Tensor.manual_seed(0)
@@ -129,7 +129,7 @@ class TestFusedQKVRoPE(unittest.TestCase):
       self.assertTrue(k.allclose(k_ref, atol=2e-2, rtol=0).item(), "K forward mismatch")
       self.assertTrue(v.allclose(v_ref, atol=0, rtol=0).item(), "V forward mismatch")
 
-  @unittest.skipUnless(has_hipcc(), "backward kernel requires hipcc to compile")
+  @unittest.skipUnless(has_hipcc() and is_cdna4(), "backward kernel requires hipcc to compile")
   def test_llama31_8b(self):
     Tensor.manual_seed(1)
     B, N, H, H_KV, D = self.SHAPE

test/backend/test_multitensor.py

@@ -3,7 +3,7 @@ from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variab
 from tinygrad.uop.ops import Ops, UOp, AxisType, graph_rewrite
 from tinygrad.helpers import getenv, prod, Context
 from tinygrad.nn.state import get_parameters
-from tinygrad.engine.realize import run_linear, compile_linear, pm_beam, pm_compile
+from tinygrad.engine.realize import run_linear, compile_linear, lower_and_compile, pm_beam
 import numpy as np
 from hypothesis import given, strategies as strat, settings
 from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph, check_schedule, assert_kernel_count, KernelCountException
@@ -80,7 +80,7 @@ class TestMultiTensor(unittest.TestCase):
     cpu_2 = ("CPU:1", "CPU:2")
     src = Tensor.ones(16).shard(cpu_2, 0).realize()
     lin = UOp(Ops.LINEAR, src=(src.to(cpu_2[::-1]).schedule_linear().src[0],))
-    with Context(BEAM=1, IGNORE_BEAM_CACHE=1): call = graph_rewrite(graph_rewrite(lin, pm_beam, ctx=1, walk=True), pm_compile, walk=True).src[0]
+    with Context(BEAM=1, IGNORE_BEAM_CACHE=1): call = lower_and_compile(graph_rewrite(lin, pm_beam, ctx=1, walk=True)).src[0]
     self.assertNotEqual(call.src[0].src[0].arg.applied_opts, ())
 
   def test_shard_same_device(self):

test/backend/test_renderer_failures.py

@@ -77,6 +77,14 @@ class TestCStyleFailures(unittest.TestCase):
 
 @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "tests for wgsl renderer")
 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))
+    src = Device[Device.DEFAULT].renderer.render(UOp.sink(store, arg=KernelInfo()).toposort())
+    self.assertIn("atomicAnd(&data0_4[0],4294967040u);", src)
+    self.assertNotIn("atomicAdd", src)
+
   def test_multiply_infinity(self):
     # multiplying a positive constant by infinity should return infinity
     # WGSL pipelines do not handle this reliably, some of which return zero, unless infinity always comes from a read on a dynamic buffer

test/backend/test_uops.py

@@ -2,7 +2,7 @@ from typing import Optional, Any
 import unittest, math
 import numpy as np
 from tinygrad.tensor import Tensor, _to_np_dtype
-from tinygrad.helpers import Context
+from tinygrad.helpers import Context, ceildiv
 from tinygrad.dtype import dtypes, DType, AddrSpace, ConstFloat  # noqa: F401
 from tinygrad.device import Buffer, Device
 from tinygrad.uop.ops import Ops, UOp, KernelInfo, AxisType, buffers
@@ -57,6 +57,35 @@ def _test_uops_result(output_dtype, uops, res):
   run_uops([out], [buf])
   return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0]
 
+@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, CStyleLanguage) and
+                     dtypes.uint64 in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires C-style pointer bitcast and 64-bit ints")
+class TestBitcastBufferView(unittest.TestCase):
+  @Context(SPEC=2)
+  def test_render(self):
+    buf = UOp.param(0, dtypes.uint32, (4,))
+    uops = to_uops_list([buf.shrink(((1, 3),)).bitcast(dtypes.uint64).index(0).store(1)], ren=Device[Device.DEFAULT].renderer)
+    idx = next(u for u in uops if u.op is Ops.INDEX and u.src[0].op is Ops.BITCAST)
+    self.assertEqual(idx.src[0].src[0].op, Ops.SHRINK)
+    Device[Device.DEFAULT].renderer.render(uops)
+
+  @Context(SPEC=2)
+  def test_load(self):
+    val = 0x1122334455667788
+    src, out = UOp.param(0, dtypes.uint32, (4,)), UOp.param(1, dtypes.uint64, (1,))
+    ibuf = Buffer(Device.DEFAULT, 4, dtypes.uint32, initial_value=np.array([0, 0x55667788, 0x11223344, 0], dtype=np.uint32).tobytes())
+    obuf = Buffer(Device.DEFAULT, 1, dtypes.uint64).allocate()
+    run_uops([out.index(0).store(src.shrink(((1, 3),)).bitcast(dtypes.uint64).index(0))], [ibuf, obuf])
+    self.assertEqual(np.frombuffer(obuf.as_memoryview(), dtype=np.uint64)[0], val)
+
+  @Context(SPEC=2)
+  def test_store(self):
+    val = 0x1122334455667788
+    dst = UOp.param(0, dtypes.uint32, (6,))
+    buf = Buffer(Device.DEFAULT, 6, dtypes.uint32, initial_value=bytes(24))
+    view = dst.shrink(((1, 5),)).bitcast(dtypes.uint64)  # two stores through one view: it must inline, not get a declared vector-pointer
+    run_uops([view.index(0).store(val ^ 0xff), view.index(1).store(val)], [buf])
+    self.assertEqual(np.frombuffer(buf.as_memoryview(), dtype=np.uint64, count=2, offset=4).tolist(), [val ^ 0xff, val])
+
 class TestUOps(unittest.TestCase):
   def _equal(self, v1, v2):
     assert isinstance(v2, (float, int, bool))
@@ -193,15 +222,16 @@ class TestLocalAccess(unittest.TestCase):
   @unittest.skipUnless(Device.DEFAULT == "WEBGPU", "Test local memory size for packed data types")
   def test_packed_smem_size(self):
     _dtypes = [dtypes.char, dtypes.uchar, dtypes.short, dtypes.ushort, dtypes.half]
-    size = 16
-    for dtype in _dtypes:
-      temp = UOp.placeholder((size,), dtype, slot=0, addrspace=AddrSpace.LOCAL)
-      uops = to_uops_list([temp], ren=Device[Device.DEFAULT].renderer)
-      out = Device[Device.DEFAULT].renderer.render(uops)
-      # half is supported in wgsl, so it doesn't have to be packed
-      corrected_size = size//(4//dtype.itemsize) if dtype != dtypes.half else size
-      # temp0: array<{Device[Device.DEFAULT].renderer.buf_map(dtype)},{corrected_size}>;
-      self.assertIn(f",{corrected_size}>;", out)
+    # a partial word still needs a whole word, so sizes that don't fill one must round up
+    for size in (16, 5):
+      for dtype in _dtypes:
+        temp = UOp.placeholder((size,), dtype, slot=0, addrspace=AddrSpace.LOCAL)
+        uops = to_uops_list([temp], ren=Device[Device.DEFAULT].renderer)
+        out = Device[Device.DEFAULT].renderer.render(uops)
+        # half is supported in wgsl, so it doesn't have to be packed
+        corrected_size = ceildiv(size, 4//dtype.itemsize) if dtype != dtypes.half else size
+        # temp0: array<{Device[Device.DEFAULT].renderer.buf_map(dtype)},{corrected_size}>;
+        self.assertIn(f",{corrected_size}>;", out)
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_shared, "test requires shared memory")
   @unittest.skip("tinygrad doesn't support this behavior")

test/mockgpu/amd/pcode.py

@@ -1,5 +1,20 @@
 # Tokenizer-based expression parser for AMD pcode
+import ast, itertools, operator, re
 from typing import Any, Callable
+
+_BINOPS = {ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul, ast.FloorDiv: operator.floordiv,
+           ast.Mod: operator.mod, ast.LShift: operator.lshift, ast.RShift: operator.rshift,
+           ast.BitAnd: operator.and_, ast.BitOr: operator.or_, ast.BitXor: operator.xor}
+def _const_int(expr: str) -> int:
+  """Evaluate a compile-time integer expression (integer literals and basic arithmetic only)."""
+  def ev(node: ast.AST) -> int:
+    if isinstance(node, ast.Expression): return ev(node.body)
+    if isinstance(node, ast.Constant) and isinstance(node.value, int): return node.value
+    if isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.USub, ast.UAdd)):
+      return (-1 if isinstance(node.op, ast.USub) else 1) * ev(node.operand)
+    if isinstance(node, ast.BinOp) and type(node.op) in _BINOPS: return _BINOPS[type(node.op)](ev(node.left), ev(node.right))
+    raise ValueError(f"not a constant integer expression: {expr!r}")
+  return ev(ast.parse(expr.strip(), mode='eval'))
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.codegen.decomp.dtype import f2f
@@ -8,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)
@@ -55,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
 
@@ -144,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
 
@@ -163,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)
@@ -185,7 +201,20 @@ def _abs(val: UOp) -> UOp:
 def _f_to_u(f, dt):
   clamped = (f < _const(f.dtype, 0.0)).where(_const(f.dtype, 0.0), f)
   truncated = UOp(Ops.TRUNC, src=(clamped,))
-  return (truncated >= _const(f.dtype, 2**(dt.itemsize*8))).where(_const(dt, dt.max), truncated.cast(dt))
+  res = (truncated >= _const(f.dtype, 2**(dt.itemsize*8))).where(_const(dt, dt.max), truncated.cast(dt))
+  return _isnan(f).where(_const(dt, 0), res)  # float->uint conversion of NaN is 0 on hardware
+
+def _f_to_i32(a: UOp) -> UOp:
+  """v_cvt_i32_f32/f64: truncate toward zero, saturate to [INT_MIN, INT_MAX], NaN -> 0.
+  (x86 cvttss2si returns 0x80000000 for all of these, which matches hardware only for negative overflow.)"""
+  res = (a >= _const(a.dtype, 2147483648.0)).where(_const(dtypes.int, 0x7FFFFFFF), UOp(Ops.TRUNC, src=(a,)).cast(dtypes.int))
+  return _isnan(a).where(_const(dtypes.int, 0), res)
+
+def _ftz_f32(v: UOp) -> UOp:
+  """Flush f32 denormals to signed zero (RDNA default float mode flushes denormal f32 inputs on select-style ops)."""
+  bits = v.bitcast(dtypes.uint32) if v.dtype == dtypes.float32 else v
+  return ((bits & _u32(0x7FFFFFFF)) < _u32(0x00800000)).where((bits & _u32(0x80000000)).bitcast(dtypes.float32),
+                                                              v if v.dtype == dtypes.float32 else v.bitcast(dtypes.float32))
 
 def _cvt_quiet(val: UOp) -> UOp:
   bits, _, _, qb, _ = _float_info(val)
@@ -230,18 +259,51 @@ def _ldexp(val: UOp, exp: UOp) -> UOp:
   if val.dtype == dtypes.uint32: val = val.bitcast(dtypes.float32)
   elif val.dtype == dtypes.uint64: val = val.bitcast(dtypes.float64)
   if exp.dtype in (dtypes.uint32, dtypes.uint64): exp = exp.cast(dtypes.int if exp.dtype == dtypes.uint32 else dtypes.int64)
-  return val * UOp(Ops.EXP2, src=(exp.cast(val.dtype),))
+  bits = val.bitcast(dtypes.uint32) if val.dtype == dtypes.float32 else val.bitcast(dtypes.uint64)
+  abs_max = _const(bits.dtype, 0x7F800000 if val.dtype == dtypes.float32 else 0x7FF0000000000000)
+  sign_mask = _const(bits.dtype, 0x80000000 if val.dtype == dtypes.float32 else 0x8000000000000000)
+  # hardware flushes denormal inputs to signed zero
+  magn_mask = _const(bits.dtype, 0x7FFFFFFF if val.dtype == dtypes.float32 else 0x7FFFFFFFFFFFFFFF)
+  is_denorm = ((bits & abs_max).eq(_const(bits.dtype, 0))) & ((bits & magn_mask).ne(_const(bits.dtype, 0)))
+  val = is_denorm.where((bits & sign_mask).bitcast(val.dtype), val)
+  # hardware propagates 0/+-inf/NaN unchanged (avoids 0*inf = NaN on the host)
+  res = val * UOp(Ops.EXP2, src=(exp.cast(val.dtype),))
+  is_special = (bits & abs_max).eq(_const(bits.dtype, 0)) | ((bits & abs_max) >= abs_max)
+  return is_special.where(val, res)
 
 def _frexp_mant(val: UOp) -> UOp:
   val = val.bitcast(dtypes.float32) if val.dtype == dtypes.uint32 else val.bitcast(dtypes.float64) if val.dtype == dtypes.uint64 else val
-  if val.dtype == dtypes.float32: return ((val.bitcast(dtypes.uint32) & _u32(0x807FFFFF)) | _u32(0x3f000000)).bitcast(dtypes.float32)
-  return ((val.bitcast(dtypes.uint64) & _const(dtypes.uint64, 0x800FFFFFFFFFFFFF)) |
-    _const(dtypes.uint64, 0x3fe0000000000000)).bitcast(dtypes.float64)
+  if val.dtype == dtypes.float32:
+    bits = val.bitcast(dtypes.uint32)
+    # denormal/zero inputs (exponent field == 0) return signed zero on hardware
+    return ((bits & _u32(0x7F800000)).ne(_u32(0))).where(((bits & _u32(0x807FFFFF)) | _u32(0x3F000000)).bitcast(dtypes.float32),
+                                                          (bits & _u32(0x80000000)).bitcast(dtypes.float32))
+  bits = val.bitcast(dtypes.uint64)
+  return ((bits & _const(dtypes.uint64, 0x7FF0000000000000)).ne(_const(dtypes.uint64, 0))).where(
+    ((bits & _const(dtypes.uint64, 0x800FFFFFFFFFFFFF)) | _const(dtypes.uint64, 0x3fe0000000000000)).bitcast(dtypes.float64),
+    (bits & _const(dtypes.uint64, 0x8000000000000000)).bitcast(dtypes.float64))
+
+def _msb(val: UOp, bits: int) -> UOp:
+  """Index of the highest set bit, or -1 if val == 0."""
+  dt = dtypes.uint64 if bits > 32 else dtypes.uint32
+  val = val.cast(dt) if val.dtype != dt else val
+  result = _const(dtypes.int, -1)
+  for i in range(bits - 1, -1, -1):
+    cond = ((val >> _const(dt, i)) & _const(dt, 1)).ne(_const(dt, 0)) & result.eq(_const(dtypes.int, -1))
+    result = cond.where(_const(dtypes.int, i), result)
+  return result
 
 def _frexp_exp(val: UOp) -> UOp:
   val = val.bitcast(dtypes.float32) if val.dtype == dtypes.uint32 else val.bitcast(dtypes.float64) if val.dtype == dtypes.uint64 else val
-  if val.dtype == dtypes.float32: return ((val.bitcast(dtypes.uint32) >> _u32(23)) & _u32(0xFF)).cast(dtypes.int) - _const(dtypes.int, 126)
-  return ((val.bitcast(dtypes.uint64) >> _const(dtypes.uint64, 52)) & _const(dtypes.uint64, 0x7FF)).cast(dtypes.int) - _const(dtypes.int, 1022)
+  if val.dtype == dtypes.float32:
+    e = (val.bitcast(dtypes.uint32) >> _u32(23)) & _u32(0xFF)
+    return e.ne(_u32(0)).where(e.cast(dtypes.int) - _const(dtypes.int, 126), _const(dtypes.int, 0))  # f32 denormals -> 0 (hardware verified)
+  bits = val.bitcast(dtypes.uint64)
+  e = (bits >> _const(dtypes.uint64, 52)) & _const(dtypes.uint64, 0x7FF)
+  mant = bits & _const(dtypes.uint64, 0xFFFFFFFFFFFFF)
+  # f64 denormals: normalized exponent = highest set mantissa bit - 1073, zero -> 0 (hardware verified)
+  denorm = mant.ne(_const(dtypes.uint64, 0)).where(_msb(mant, 52) - _const(dtypes.int, 1073), _const(dtypes.int, 0))
+  return e.ne(_const(dtypes.uint64, 0)).where(e.cast(dtypes.int) - _const(dtypes.int, 1022), denorm)
 
 TWO_OVER_PI = int(
   "0145f306dc9c882a53f84eafa3ea69bb81b6c52b3278872083fca2c757bd778ac36e48dc74849ba5c00c925dd413a32439fc3bd"
@@ -299,9 +361,9 @@ _FUNCS: dict[str, Callable[..., UOp]] = {
   'fma': lambda a, b, c: a * b + c,
   'i32_to_f32': lambda a: a.cast(dtypes.int).cast(dtypes.float32),
   'u32_to_f32': lambda a: a.cast(dtypes.uint32).cast(dtypes.float32),
-  'f32_to_i32': lambda a: UOp(Ops.TRUNC, src=(a.bitcast(dtypes.float32),)).cast(dtypes.int),
+  'f32_to_i32': lambda a: _f_to_i32(a.bitcast(dtypes.float32)),
   'f32_to_u32': lambda a: _f_to_u(a.bitcast(dtypes.float32), dtypes.uint32),
-  'f64_to_i32': lambda a: UOp(Ops.TRUNC, src=(a.bitcast(dtypes.float64),)).cast(dtypes.int),
+  'f64_to_i32': lambda a: _f_to_i32(a.bitcast(dtypes.float64)),
   'f64_to_u32': lambda a: _f_to_u(a.bitcast(dtypes.float64), dtypes.uint32),
   'f16_to_f32': lambda a: _f16_extract(a).cast(dtypes.float32),
   'f32_to_f16': lambda a: a.cast(dtypes.half),
@@ -360,22 +422,13 @@ _FUNCS: dict[str, Callable[..., UOp]] = {
   'fp8_to_f32': _fp8_to_f32, 'bf8_to_f32': _bf8_to_f32, 'f32_to_fp8': _f32_to_fp8, 'f32_to_bf8': _f32_to_bf8,
   'f32_to_bf16': _f32_to_bf16, 'f32_to_bf16_SR': _f32_to_bf16_sr, 'f32_to_bf16_sr': _f32_to_bf16_sr,
 }
-for is_max, name in [(False, 'min'), (True, 'max')]:
-  for dt, sfx in [(dtypes.float32, 'f32'), (dtypes.int, 'i32'), (dtypes.uint32, 'u32'), (dtypes.int16, 'i16'), (dtypes.uint16, 'u16')]:
-    _FUNCS[f'v_{name}_{sfx}'] = lambda *a, im=is_max, d=dt: _minmax_reduce(im, d, *a)
-    _FUNCS[f'v_{name}3_{sfx}'] = lambda *a, im=is_max, d=dt: _minmax_reduce(im, d, *a)
-# f16 min/max/min3/max3/med3
-for is_max, name in [(False, 'min'), (True, 'max')]:
-  _FUNCS[f'v_{name}_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
-  _FUNCS[f'v_{name}3_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
-  _FUNCS[f'v_{name}_num_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
-  _FUNCS[f'v_{name}_num_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
-  _FUNCS[f'v_{name}3_num_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
-  _FUNCS[f'v_{name}3_num_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
-  _FUNCS[f'v_{name}imum_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
-  _FUNCS[f'v_{name}imum_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
-  _FUNCS[f'v_{name}imum3_f16'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.half, *[_f16_extract(x) for x in a])
-  _FUNCS[f'v_{name}imum3_f32'] = lambda *a, im=is_max: _minmax_reduce(im, dtypes.float32, *a)
+# min/max family: min/max + 3-input (x3), IEEE num variants (f16/f32 only), and long names minimum/maximum (f16/f32 only)
+for is_max, name, full in [(False, 'min', 'minimum'), (True, 'max', 'maximum')]:
+  for dt, sfx, pre in [(dtypes.float32, 'f32', None), (dtypes.int, 'i32', None), (dtypes.uint32, 'u32', None),
+                       (dtypes.int16, 'i16', None), (dtypes.uint16, 'u16', None), (dtypes.half, 'f16', _f16_extract)]:
+    def mm(*a, im=is_max, d=dt, p=pre): return _minmax_reduce(im, d, *(a if p is None else [p(x) for x in a]))
+    extra = (f'v_{name}_num_{sfx}', f'v_{name}3_num_{sfx}', f'v_{full}_{sfx}', f'v_{full}3_{sfx}') if dt in (dtypes.float32, dtypes.half) else ()
+    for fn in (f'v_{name}_{sfx}', f'v_{name}3_{sfx}', *extra): _FUNCS[fn] = mm
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # TOKENIZER/PARSER
@@ -497,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
@@ -507,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 = [('||',), ('&&',), ('|',), ('^',), ('&',), ('==', '!=', '<>'), ('>=', '<=', '>', '<'), ('>>', '<<'), ('+', '-'), ('*', '/'), ('**',)]
 
@@ -529,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()
@@ -669,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:
@@ -698,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}']
@@ -872,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
 
@@ -890,6 +940,8 @@ class Parser:
       return result & _isnan(l).logical_not() & _isnan(r).logical_not()
     return result
 
+_break_var_ids = itertools.count()  # unique names for per-loop break-tracking variables
+
 def _match_bracket(toks: list[Token], start: int) -> tuple[int, list[Token]]:
   """Match brackets from start, return (end_idx, inner_tokens)."""
   j, depth = start + 1, 1
@@ -968,9 +1020,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
           p.eat('NUM')
           p.eat('QUOTE')
         if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl'))
-        expr = p.parse().simplify()
-        assert expr.op == Ops.CONST, f"loop bound must be constant, got {expr}"
-        return int(expr.val)
+        return int(p.parse())
       start_val = parse_bound()
       p.eat('COLON')
       end_val = parse_bound()
@@ -987,7 +1037,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
         i += 1
       # Execute loop with break support
       has_break = any('break' in bl.lower() for bl in body_lines)
-      found_var = f'_found_{id(body_lines)}' if has_break else None
+      found_var = f'_found_{next(_break_var_ids)}' if has_break else None
       if found_var: env[found_var] = block_assigns[found_var] = _const(dtypes.bool, False)
       for loop_i in range(start_val, end_val + 1):
         subst_lines = [_subst_loop_var(bl, loop_var, loop_i) for bl in body_lines if not (has_break and bl.strip().lower() == 'break')]
@@ -1087,7 +1137,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
           j, slice_toks = _match_bracket(toks, j)
           slice_str = _tok_str(slice_toks)
           hi_str, lo_str = slice_str.split(':')
-          hi_val, lo_val = int(eval(hi_str.strip())), int(eval(lo_str.strip()))
+          hi_val, lo_val = _const_int(hi_str), _const_int(lo_str)
           if j < len(toks) and toks[j].type == 'DOT': j += 2  # skip .type suffix
           if j < len(toks) and toks[j].type == 'EQUALS': j += 1
           ln = parse_tokens(lane_toks, env, funcs)
@@ -1145,7 +1195,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
         hi_str = ' '.join(t.val for t in toks[bracket_start:colon_pos] if t.type != 'EOF')
         lo_str = ' '.join(t.val for t in toks[colon_pos+1:j] if t.type != 'EOF')
         try:
-          hi_val, lo_val = int(eval(hi_str)), int(eval(lo_str))
+          hi_val, lo_val = _const_int(hi_str), _const_int(lo_str)
           hi, lo = max(hi_val, lo_val), min(hi_val, lo_val)
           j += 1
           if j < len(toks) and toks[j].type == 'DOT': j += 2
@@ -1159,7 +1209,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
           block_assigns[var] = env[var] = _set_bits(old, _val_to_bits(val), hi - lo + 1, lo)
           i += 1
           continue
-        except Exception: pass
+        except (ValueError, SyntaxError): pass  # non-constant slice bounds - fall through to other statement forms
       elif toks[1].type == 'LBRACKET':  # bit index: var[expr] (only for var[...], not var.type[...])
         existing = block_assigns.get(var, env.get(var))
         if existing is not None and isinstance(existing, UOp) and \
@@ -1360,3 +1410,25 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
 
 def parse_expr(expr: str, env: dict[str, VarVal], funcs: dict | None = None) -> UOp:
   return parse_tokens(tokenize(expr.strip().rstrip(';')), env, funcs)
+
+def parse_pcode(pcode: str, srcs: dict[str, UOp | int] | None = None) -> tuple[dict, list]:
+  env: dict = srcs.copy() if srcs else {}
+  assigns: list[tuple[str, UOp]] = []
+  raw_lines = [l.strip().rstrip(';') for l in pcode.split('\n') if l.strip() and not l.strip().startswith('//')]
+  # TODO: pcode.py should tokenize full pcode string instead of line-by-line, then this hack can be removed
+  lines: list[str] = []
+  for l in raw_lines:
+    if lines and re.search(r'(&&|\|\||[&|+\-*/^])\s*$', lines[-1]): lines[-1] = lines[-1] + ' ' + l
+    else: lines.append(l)
+  _, final, _ = parse_block(lines, 0, env, assigns=assigns)
+  sliced = set(d.split('[')[0] for d, _ in assigns if '[' in d)
+  for var, val in final.items():
+    if var in ['D0', 'S0', 'SCC', 'VCC', 'EXEC', 'PC', 'RETURN_DATA', 'VDATA'] and isinstance(val, UOp):
+      if var in sliced and not any(re.match(rf'{var}\.\w+\s*=', l) for l in lines): continue
+      for l in lines:
+        if (m := re.match(rf'{var}\.(\w+(?:\[\w+\])?)', l)):
+          assigns.append((f'{var}.{m.group(1)}', val))
+          break
+      else: assigns.append((var, val))
+  return env, assigns
+

test/mockgpu/amd/sqtt_enc.py

@@ -0,0 +1,100 @@
+# SQTT trace encoder for the emulator (the decoder lives in tinygrad/renderer/amd/sqtt.py).
+# run_asm emits packets inline as instructions execute; finished traces end up in emu.sqtt_traces.
+from __future__ import annotations
+from tinygrad.renderer.amd.dsl import Inst
+from tinygrad.renderer.amd.sqtt import (_build_decode_tables, PACKET_TYPES_RDNA3, PacketType, InstOp,
+                                        LAYOUT_HEADER, WAVESTART, WAVEEND, INST, IMMEDIATE, VALUINST)
+
+_NIB_COUNTS = {cls: nc for _, (cls, nc, *_) in _build_decode_tables(PACKET_TYPES_RDNA3)[0].items()}
+
+def _emit_nibbles(nibbles: list[int], pkt_cls: type[PacketType], **kwargs):
+  raw = pkt_cls.encoding.default
+  for k, v in kwargs.items(): raw = pkt_cls.__dict__[k].set(raw, v)
+  nibbles.extend((raw >> (i * 4)) & 0xF for i in range(_NIB_COUNTS[pkt_cls]))
+
+def make_encoder():
+  """Build an SQTT trace encoder for the emulator. Returns (emit, finish, finalize)."""
+  from tinygrad.runtime.autogen.amd.rdna3.enum import SOPPOp as SOPPOp3
+  from tinygrad.runtime.autogen.amd.rdna4.enum import SOPPOp as SOPPOp4
+  from tinygrad.runtime.autogen.amd.rdna3 import ins as ir3
+  from tinygrad.runtime.autogen.amd.rdna4 import ins as ir4
+  from tinygrad.runtime.autogen.amd.cdna import ins as irc
+  import re
+
+  def _kinds(*names: str) -> tuple[type[Inst], ...]:
+    return tuple(getattr(m, n) for m in (ir3, ir4, irc) for n in names if hasattr(m, n))
+  _SOPP, _SMEM, _DS = _kinds('SOPP'), _kinds('SMEM'), _kinds('DS')
+  _GLOBAL, _FLAT, _SCRATCH = _kinds('GLOBAL', 'VGLOBAL'), _kinds('FLAT', 'VFLAT'), _kinds('SCRATCH', 'VSCRATCH')
+  _VALU = _kinds('VOP1', 'VOP2', 'VOP3', 'VOP3P', 'VOP3PX2', 'VOPC', 'VOPD', 'VOP3SD', 'VOP3_SDST', 'VOP1_SDST')
+
+  # SOPP classification sets
+  _SOPP_SKIP = {SOPPOp3.S_ENDPGM.value, SOPPOp3.S_ENDPGM_SAVED.value, SOPPOp3.S_ENDPGM_ORDERED_PS_DONE.value, SOPPOp3.S_DELAY_ALU.value}
+  _SOPP_IMMEDIATE = {SOPPOp3.S_NOP.value, SOPPOp3.S_CLAUSE.value, SOPPOp3.S_WAITCNT.value, SOPPOp3.S_WAITCNT_DEPCTR.value,
+                     SOPPOp3.S_WAIT_IDLE.value, SOPPOp3.S_WAIT_EVENT.value, SOPPOp3.S_SLEEP.value, SOPPOp3.S_SET_INST_PREFETCH_DISTANCE.value}
+  for _op in (SOPPOp4.S_WAIT_ALU, SOPPOp4.S_WAIT_LOADCNT, SOPPOp4.S_WAIT_STORECNT, SOPPOp4.S_WAIT_SAMPLECNT,
+              SOPPOp4.S_WAIT_BVHCNT, SOPPOp4.S_WAIT_EXPCNT, SOPPOp4.S_WAIT_DSCNT, SOPPOp4.S_WAIT_KMCNT,
+              SOPPOp4.S_WAIT_LOADCNT_DSCNT, SOPPOp4.S_WAIT_STORECNT_DSCNT):
+    _SOPP_IMMEDIATE.add(_op.value)
+  _SOPP_BARRIER = {SOPPOp3.S_BARRIER.value}
+  if hasattr(SOPPOp4, 'S_BARRIER_WAIT'): _SOPP_BARRIER.add(SOPPOp4.S_BARRIER_WAIT.value)
+  if hasattr(SOPPOp4, 'S_BARRIER_LEAVE'): _SOPP_BARRIER.add(SOPPOp4.S_BARRIER_LEAVE.value)
+  _SOPP_BRANCH = {SOPPOp3.S_BRANCH.value, SOPPOp3.S_CBRANCH_SCC0.value, SOPPOp3.S_CBRANCH_SCC1.value,
+                  SOPPOp3.S_CBRANCH_VCCZ.value, SOPPOp3.S_CBRANCH_VCCNZ.value,
+                  SOPPOp3.S_CBRANCH_EXECZ.value, SOPPOp3.S_CBRANCH_EXECNZ.value}
+
+  # VALU sub-classification patterns
+  _VALUT_4_RE = re.compile(r'V_(EXP|LOG|RCP|RSQ|SQRT|SIN|COS|CEIL|FLOOR|TRUNC|RNDNE|FRACT|FREXP)_')
+  _VALUB_2_RE = re.compile(r'V_(LSHLREV|LSHRREV|ASHRREV)_(B|I)64')
+  _VALUB_4_RE = re.compile(r'V_MAD_(U|I)64')
+  _VALUB_16_RE = re.compile(r'V_\w+_F64')
+
+  def _valu_op(op_name: str) -> InstOp|None:
+    if 'CMPX' in op_name: return InstOp.VALU1_WR_EXEC
+    if _VALUB_2_RE.search(op_name): return InstOp.VALUB_2
+    if _VALUB_4_RE.search(op_name): return InstOp.VALUB_4
+    if _VALUB_16_RE.search(op_name): return InstOp.VALUB_16
+    if _VALUT_4_RE.search(op_name): return InstOp.VALUT_4
+    return None
+
+  def _mem_op(t: type[Inst], op_name: str) -> InstOp:
+    is_store = "STORE" in op_name
+    if issubclass(t, _DS): return InstOp.LDS_WR_2 if is_store else InstOp.LDS_RD
+    if issubclass(t, _GLOBAL): return InstOp.SGMEM_WR_2 if is_store else InstOp.SGMEM_RD_1
+    if issubclass(t, _FLAT) or issubclass(t, _SCRATCH): return InstOp.FLAT_WR_3 if is_store else InstOp.FLAT_RD_2
+    return InstOp.SALU
+
+  nibbles: list[int] = []
+  started: set[int] = set()
+  _emit_nibbles(nibbles, LAYOUT_HEADER, layout=3, sel_a=6)
+
+  def emit(wave_id: int, inst: Inst, branch_taken: bool|None):
+    """Emit an SQTT packet for one executed instruction."""
+    w = wave_id & 0x1F
+    if wave_id not in started:
+      _emit_nibbles(nibbles, WAVESTART, delta=1, simd=0, wgp=0, wave=w, id7=wave_id)
+      started.add(wave_id)
+    inst_type, inst_op, op_name = type(inst), inst.op.value if hasattr(inst, 'op') else 0, inst.op.name if hasattr(inst, 'op') else ""
+    if issubclass(inst_type, _SOPP):
+      if inst_op in _SOPP_SKIP: return
+      if inst_op in _SOPP_IMMEDIATE: _emit_nibbles(nibbles, IMMEDIATE, delta=1, wave=w)
+      elif inst_op in _SOPP_BARRIER: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=InstOp.BARRIER)
+      elif inst_op in _SOPP_BRANCH: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=InstOp.JUMP if branch_taken else InstOp.JUMP_NO)
+      else: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=InstOp.SALU)
+    elif issubclass(inst_type, _VALU):
+      if (op := _valu_op(op_name)) is None: _emit_nibbles(nibbles, VALUINST, delta=1, wave=w)
+      else: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=op)
+    elif issubclass(inst_type, _SMEM): _emit_nibbles(nibbles, INST, delta=1, wave=w, op=InstOp.SMEM_RD)
+    else: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=_mem_op(inst_type, op_name))
+
+  def finish(wave_id: int):
+    """Emit WAVEEND for a completed wave."""
+    if wave_id in started: _emit_nibbles(nibbles, WAVEEND, delta=1, simd=0, wgp=0, wave=wave_id & 0x1F)
+
+  def finalize() -> bytes:
+    """Pad and return the encoded SQTT blob."""
+    while len(nibbles) % 2 != 0: nibbles.append(0)
+    nibbles.extend([0] * 32)
+    while len(nibbles) % 64 != 0: nibbles.append(0)
+    return bytes(nibbles[i] | ((nibbles[i + 1] if i + 1 < len(nibbles) else 0) << 4) for i in range(0, len(nibbles), 2))
+
+  return emit, finish, finalize

test/mockgpu/usb.py

@@ -160,7 +160,7 @@ class MockUSB3:
     elif request == 0xE5:
       self.state._xram_write_byte(value, index)
     elif request == 0xF2:
-      op = ("sram_read" if value & 0x8000 else "sram_write", 0xF000, (value & 0x7FFF) * 512)
+      op = ("sram_read" if value & 0x8000 else "sram_write", 0xF000 + (index & 0xFF) * 0x4000, (value & 0x7FFF) * 512)
       if value & 0x8000: self._bulk_read_op = op
       else: self._bulk_write_op = op
     elif request == 0xF0:
@@ -193,19 +193,33 @@ class MockUSB3:
     op, address, size = self._bulk_write_op
     assert len(data) == size
     if op == "sram_write":
-      host_addr, region_size = self.state._dma_regions[address]
-      ctypes.memmove(host_addr, data, min(len(data), region_size))
+      ctrl, (host_addr, region_size) = next((ca, r) for ca, r in self.state._dma_regions.items() if ca <= address < ca + r[1])
+      ctypes.memmove(host_addr + (address - ctrl), data, min(len(data), region_size - (address - ctrl)))
+      self.state.driver._emulate_execute()  # landed data may un-stall a ring polling on it (e.g. copyin sentinels)
     elif op == "pcie_write": self.state._pcie_write(address, data)
     else: raise RuntimeError(f"cannot bulk write for {op}")
     self._bulk_write_op = None
 
+  def bulk_write_async(self, payload:memoryview, timeout:int=10000) -> int:  # the mock completes transfers synchronously
+    self.bulk_write(bytes(payload), timeout)
+    return 0
+
+  def control_write_async(self, request:int, value:int=0, index:int=0, data:bytes=b"", timeout:int=1000) -> int:
+    self.control_write(request, value, index, data, timeout)
+    return 0
+
+  def control_read_async(self, request:int, length:int, value:int=0, index:int=0, timeout:int=1000) -> tuple[int, memoryview]:
+    return 0, self.control_read(request, length, value, index, timeout)
+
+  def bulk_wait(self, tag:int): pass
+
   def bulk_read(self, length:int, timeout:int=1000) -> memoryview:
     assert self._bulk_read_op is not None
     op, address, size = self._bulk_read_op
     assert length == size
     if op == "sram_read":
-      host_addr, region_size = self.state._dma_regions[address]
-      data = bytes((ctypes.c_ubyte * min(length, region_size)).from_address(host_addr))
+      ctrl, (host_addr, region_size) = next((ca, r) for ca, r in self.state._dma_regions.items() if ca <= address < ca + r[1])
+      data = bytes((ctypes.c_ubyte * min(length, region_size - (address - ctrl))).from_address(host_addr + (address - ctrl)))
     elif op == "pcie_read": data = self.state._pcie_read(address, length)
     else: raise RuntimeError(f"cannot bulk read for {op}")
     self._bulk_read_op = None

test/null/test_const_folding.py

@@ -3,7 +3,6 @@ from tinygrad import dtypes, Context
 from tinygrad.dtype import DType, ConstType
 from tinygrad.uop.ops import Ops, UOp
 from test.helpers import full_rewrite
-import numpy as np
 
 class TestWeakConstFolding(unittest.TestCase):
   def test_weakint_math(self):
@@ -27,16 +26,14 @@ class TestBitcastConstFolding(unittest.TestCase):
     for val, src_dt, dst_dt, bits in ((3000000000, dtypes.int32, dtypes.uint32, 3000000000),
                                       (70000, dtypes.int16, dtypes.uint16, 4464),
                                       (-5, dtypes.uint32, dtypes.int32, -5)):
-      self.assertEqual(UOp.const(val, src_dt).bitcast(dst_dt).simplify().val, bits)
+      self.assertIs(UOp.const(val, src_dt).bitcast(dst_dt).simplify(), UOp.const(bits, dst_dt))
 
   def test_scalar_bitcast(self):
     def t(cases: dict[DType, ConstType]):
       for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()):
         if not math.isnan(from_v):
           r = UOp.const(from_v, from_dt).bitcast(to_dt).simplify()
-          self.assertEqual(r.op, Ops.CONST, msg:=f"{from_dt} -> {to_dt} ({from_v} -> {to_v})")
-          self.assertEqual(r.dtype, to_dt, msg)
-          np.testing.assert_equal(r.val, to_v, msg)
+          self.assertIs(r, UOp.const(to_v, to_dt), f"{from_dt} -> {to_dt} ({from_v} -> {to_v})")
 
     t({dtypes.int8: 0, dtypes.uint8: 0, dtypes.bool: False})
     t({dtypes.int8: 1, dtypes.uint8: 1, dtypes.bool: True})

test/null/test_simplify_valid_idx.py

@@ -3,7 +3,7 @@ import unittest, itertools
 from tinygrad.codegen.late.coalesce import indexing_simplify
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite
-from tinygrad.uop.weak import pm_lower_index_dtype
+from tinygrad.uop.weak import pm_commit_weak
 from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load
 from tinygrad.helpers import Context
 from test.helpers import full_rewrite
@@ -496,7 +496,7 @@ class TestImageSimplification(unittest.TestCase):
     idx_y = (f + UOp.const(1.0)).cast(dtypes.int)
     load = get_load_image_uop((10, 10, 4), (UOp.const(-1) < idx_y) & (idx_y < UOp.const(10)),
                               (Special("gidx0", 10), idx_y))
-    off = graph_rewrite(load.sink(), pm_lower_index_dtype+indexing_simplify, ctx={}).src[0].src[0]
+    off = graph_rewrite(load.sink(), pm_commit_weak+indexing_simplify).src[0].src[0]
     self.assertEqual(off.src[1].get_valid(), UOp.const(True))
 
 class TestDropTrueGate(unittest.TestCase):

test/null/test_uop_graph.py

@@ -230,11 +230,11 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_depth_2_const_fold(self):
     v = UOp.variable("tmp", 0, 1, dtypes.int, param=True)
-    c2 = UOp.const(2, dtypes.int)
-    c4 = UOp.const(4, dtypes.int)
+    c2 = UOp.const(2)
+    c4 = UOp.const(4)
     vc = v+c2
     out = vc+c4
-    self.assertIs(out.simplify(), (v+UOp.const(6, dtypes.int)).simplify())
+    self.assertIs(out.simplify(), (v+UOp.const(6)).simplify())
 
   def test_bitcast_to_same_dtype_fold(self):
     for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
@@ -245,7 +245,7 @@ class TestUOpGraph(unittest.TestCase):
 
   def test_sub_with_cast_folds(self):
     a = Variable("a", 0, 5)
-    out = a.cast(dtypes.int)+(-a).cast(dtypes.int)
+    out = a+(-a)
     self.assertIs(full_rewrite(out.sink()).src[0], full_rewrite(UOp.const(0, dtypes.int).sink()).src[0])
 
   def test_where_on_gated_load_fold(self):

test/null/test_uop_symbolic.py

@@ -6,7 +6,6 @@ from tinygrad.dtype import dtypes, ConstType, DType, Invalid
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite, sym_infer
 from tinygrad.uop.spec import spec_shared, type_verify
 from tinygrad.uop.symbolic import sym, commutative, pm_simplify_valid, pm_move_where_on_load
-from tinygrad.uop.weak import pm_cast_weak
 from tinygrad.uop.validate import uops_to_z3
 
 def check_uop_against_string(self, v:UOp, s:str):
@@ -36,7 +35,7 @@ class TestSymbolic(unittest.TestCase):
     self.assertEqual(solver.check(expr1 != expr2), z3.unsat, "simplified expression not equal to original")
 
   def helper_test_variable(self, v, n, m, s, test_z3:bool=True):
-    v_simplified = graph_rewrite(v, sym+pm_cast_weak, name="simplify symbolic uop")
+    v_simplified = graph_rewrite(v, sym, name="simplify symbolic uop")
     if test_z3: self.check_equal_z3(v, v_simplified)
     nmin, nmax = v_simplified.vmin, v_simplified.vmax
     check_uop_against_string(self, v_simplified, s)
@@ -149,6 +148,13 @@ class TestSymbolic(unittest.TestCase):
   def test_xor_0(self):
     self.helper_test_variable(Variable("a", 0, 8, dtypes.int) ^ 0, 0, 8, "a", test_z3=False)
 
+  def test_or_0(self):
+    self.helper_test_variable(Variable("a", 0, 8, dtypes.int) | 0, 0, 8, "a", test_z3=False)
+
+  def test_shift_0(self):
+    self.helper_test_variable(Variable("a", 0, 8, dtypes.int) << 0, 0, 8, "a")
+    self.helper_test_variable(Variable("a", 0, 8, dtypes.int) >> 0, 0, 8, "a")
+
   def test_xor_self_inverse(self):
     self.helper_test_variable((Variable("a", 0, 8, dtypes.int) ^ 5) ^ 5, 0, 8, "a", test_z3=False)
 
@@ -1017,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):

test/null/test_uops.py

@@ -5,8 +5,8 @@ from tinygrad.tensor import Tensor
 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, 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.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_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
@@ -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):
@@ -280,7 +282,7 @@ class TestFastIdiv(unittest.TestCase):
   def test_division_power_of_two(self):
     for dt in (dtypes.int32, dtypes.uint32):
       g = UOp.param(0, dt, (3,))
-      c = UOp.const(2).cast(dt)
+      c = UOp.const(2)
       l = g.index(c)
       a = UOp(Ops.CDIV, dt, (l, c))
       uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -293,7 +295,7 @@ class TestFastIdiv(unittest.TestCase):
     # FLOORMOD by a power of two lowers to AND (correct floor mod for any sign in two's complement)
     for dt in (dtypes.int32, dtypes.uint32):
       g = UOp.param(0, dt, (9,))
-      c = UOp.const(8).cast(dt)
+      c = UOp.const(8)
       a = UOp(Ops.FLOORMOD, dt, (g.index(c), c))
       uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
       ops = [x.op for x in uops]
@@ -301,23 +303,24 @@ class TestFastIdiv(unittest.TestCase):
       self.assertNotIn(Ops.CMOD, ops, f"For dtype={dt} FLOORMOD by pow2 left a MOD")
       self.assertNotIn(Ops.FLOORMOD, ops, f"For dtype={dt} FLOORMOD survived past late rewrite")
 
-  def test_floordiv_power_of_two_uint(self):
-    # uint FLOORDIV by a power of two lowers to a shift, leaving no IDIV/FLOORDIV in the kernel
-    for dt in (dtypes.uint32, dtypes.uint64):
+  def test_floordiv_power_of_two(self):
+    # FLOORDIV by a power of two lowers to a shift, with no round toward zero correction (a shift is exactly floor division)
+    for dt in (dtypes.int32, dtypes.uint32, dtypes.int64, dtypes.uint64):
       g = UOp.param(0, dt, (3,))
-      c = UOp.const(2).cast(dt)
+      c = UOp.const(2)
       a = UOp(Ops.FLOORDIV, dt, (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")
       self.assertNotIn(Ops.CDIV, ops, f"For dtype={dt} FLOORDIV by power of two did not simplify to shift")
+      self.assertNotIn(Ops.CMOD, ops, f"For dtype={dt} FLOORDIV by pow2 kept the round toward zero correction")
       self.assertNotIn(Ops.FLOORDIV, ops, f"For dtype={dt} FLOORDIV survived past late rewrite")
 
   @Context(DISABLE_FAST_IDIV=0)
   @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
   def test_fast_idiv_and_mod(self):
     g = UOp.param(0, dtypes.uint32, (4,))
-    c = UOp.const(3).cast(dtypes.uint)
+    c = UOp.const(3)
     l = g.index(c)
     a = UOp(Ops.CDIV, src=(l, c))
     uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -337,7 +340,7 @@ class TestFastIdiv(unittest.TestCase):
   def test_fast_idiv_bounded_numerator_zero(self):
     x = UOp.variable("x", 0, 1, dtype=dtypes.int32)
     for val in range(2):
-      self.assertEqual(eval_uop(x.alu(Ops.CDIV, UOp.const(3).cast(x.dtype)), vals=(val,)), cdiv(val, 3))
+      self.assertEqual(eval_uop(x.alu(Ops.CDIV, UOp.const(3)), vals=(val,)), cdiv(val, 3))
 
   @Context(DISABLE_FAST_IDIV=0)
   def test_fast_idiv_remove_powers_of_two(self):
@@ -362,7 +365,7 @@ class TestFastIdiv(unittest.TestCase):
 
   def test_disable_fast_idiv(self):
     g = UOp.param(0, dtypes.uint32, (4,))
-    c = UOp.const(3).cast(dtypes.uint)
+    c = UOp.const(3)
     l = g.index(c)
     a = UOp(Ops.CDIV, src=(l, c))
     with Context(DISABLE_FAST_IDIV=1):
@@ -457,6 +460,14 @@ class TestUopsObject(unittest.TestCase):
     self.assertEqual(a.device, Device.DEFAULT)
 
 class TestUOpRender(unittest.TestCase):
+  def test_render_ssimplified_marg_outside_toposort(self):
+    r = UOp.range(UOp.const(16, dtypes.int), 2, AxisType.WEAK, dtype=dtypes.int)
+    offset = (r * 2) + (r * 2)
+    shrink = UOp(Ops.SHRINK, src=(UOp.param(0, dtypes.uint, (32,)), offset, UOp.const(2, dtypes.int)))
+    self.assertIsNot(shrink.src[1], shrink.marg[0][0])
+    self.assertEqual(shrink.render(simplify=False), "p0.shrink((((r2*4), 2),))")
+    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=())
     self.assertEqual(u.render(simplify=False), "{}")

test/null/test_viz.py

@@ -1,5 +1,5 @@
-import unittest, decimal, sys, json, contextlib, tempfile, pickle, io, math
-from pathlib import Path
+import unittest
+import decimal, sys, json, contextlib, tempfile, pickle, io, math, pathlib
 from dataclasses import dataclass
 from typing import Generator
 
@@ -516,6 +516,22 @@ class TestVizIntegration(unittest.TestCase):
     src_render = get_render(viz.data, steps[src_idx]["query"])["src"]
     self.assertEqual(src, src_render)
 
+  def test_profiler_duplicate_name(self):
+    kernel_name = "duplicate_name"
+    def one(A:UOp): return A[0].store(UOp.const(1.0, dtypes.float)).sink(arg=KernelInfo(kernel_name))
+    def zero(A:UOp): return A[0].store(UOp.const(0.0, dtypes.float)).sink(arg=KernelInfo(kernel_name))
+    with save_viz() as viz:
+      @TinyJit
+      def f(a:Tensor, b:Tensor): return Tensor.custom_kernel(a, fxn=one)[0], Tensor.custom_kernel(b, fxn=zero)[0]
+      a, b = Tensor.empty(4, device="NULL"), Tensor.empty(4, device="NULL")
+      # warmup
+      for _ in range(2): Tensor.realize(*f(a, b))
+      Tensor.realize(*f(a, b))
+    kernels = {i for i,c in enumerate(viz.list_items()) if c["name"] == kernel_name}
+    profile = decode_profile(unwrap(get_profile(viz.data, cpu_events)))
+    events = [e for e in profile["layout"]["NULL"]["events"] if e["name"] == kernel_name]
+    self.assertEqual({e["ref"] for e in events}, kernels)
+
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphEntry
 from tinygrad.viz.serve import get_profile
 from tinygrad.viz.cli import decode_profile
@@ -819,8 +835,6 @@ from extra.gemm.amd_asm_matmul import Kernel
 
 @needs_tracked_pm
 class TestCfg(unittest.TestCase):
-  def setUp(self): self.arch = "gfx1100"
-
   def get_cfg(self, name:str, k:Kernel):
     insts = k.finalize()
     def fxn(out:UOp) -> UOp:
@@ -829,7 +843,7 @@ class TestCfg(unittest.TestCase):
       sink = UOp.sink(out.base, lidx, gidx, arg=KernelInfo(name=name))
       return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
     with save_viz() as viz:
-      with Context(DEV=f"NULL::{self.arch}"):
+      with Context(DEV="NULL::gfx1100"):
         out = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
         _ = do_to_program(out.schedule_linear().src[-1].src[0], Device[out.device].renderer)
     codegen_rewrites = next(s for s in viz.list_items() if s["name"] == name)
@@ -1011,8 +1025,8 @@ def run_cli(*cli_args) -> list[dict]:
 @contextlib.contextmanager
 def write_files(viz) -> list[str]:
   with tempfile.TemporaryDirectory() as tmpdir:
-    (r:=Path(tmpdir)/"rewrites.pkl").write_bytes(pickle.dumps(viz.data.trace))
-    (p:=Path(tmpdir)/"profile.pkl").write_bytes(pickle.dumps(cpu_events))
+    (r:=pathlib.Path(tmpdir)/"rewrites.pkl").write_bytes(pickle.dumps(viz.data.trace))
+    (p:=pathlib.Path(tmpdir)/"profile.pkl").write_bytes(pickle.dumps(cpu_events))
     yield ["--rewrites-path", str(r), "--profile-path", str(p)]
 
 class TestCLI(unittest.TestCase):

test/unit/test_attention.py

@@ -5,6 +5,8 @@ from tinygrad.llm.model import (
   GatedDeltaNetBlock, SSMConfig, TransformerBlock, TransformerConfig,
   apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk,
 )
+from tinygrad.llm.kernels.amd import Linear, gated_delta_prefill, amd_custom_kernels_supported
+from tinygrad.llm.gguf import ggml_data_to_tensor
 
 def apply_rope(x:Tensor, start_pos:int):
   B, H, T, Hd = x.shape
@@ -12,6 +14,15 @@ def apply_rope(x:Tensor, start_pos:int):
   freqs_cis = precompute_freqs_cis(Hd, start_pos+T)[start_pos:start_pos+T]
   return apply_rope_new(x, freqs_cis)
 
+class TestLinear(unittest.TestCase):
+  def test_recovers_packed_ggml_weight(self):
+    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)
+      linear.set_quantized(decoded)
+      self.assertEqual((linear.ggml_type, linear.weight.numel()), (ggml_type, words))
+
 class TestAttention(unittest.TestCase):
   def test_apply_rope(self):
     x = Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32)
@@ -41,6 +52,23 @@ class TestAttention(unittest.TestCase):
     np.testing.assert_allclose(block.cache_kv[0, :, :, :seqlen, :].numpy(), expected.numpy(), rtol=1e-5, atol=1e-5)
 
 class TestGatedDeltaNetBlock(unittest.TestCase):
+  def test_gated_delta_rectangular_state_and_row_decay(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    rng = np.random.default_rng(42)
+    q, k = (rng.normal(size=(1, 1, 3, 32)).astype(np.float32) for _ in range(2))
+    v, beta = rng.normal(size=(1, 1, 3, 4)).astype(np.float32), rng.uniform(size=(1, 1, 3)).astype(np.float32)
+    alpha, initial = rng.uniform(0.8, 1, size=(1, 1, 3, 4)).astype(np.float32), rng.normal(size=(1, 1, 4, 32)).astype(np.float32)
+    expected_state, expected_out = initial.copy(), np.empty_like(v)
+    for t in range(3):
+      previous, av = expected_state.copy(), alpha[:, :, t, :, None]
+      delta = (v[:, :, t] - (previous*k[:, :, t, None]).sum(-1)*alpha[:, :, t]) * beta[:, :, t, None]
+      expected_state = previous*av + delta[..., None]*k[:, :, t, None, :]
+      expected_out[:, :, t] = (previous*q[:, :, t, None]).sum(-1)*alpha[:, :, t] + delta*(q[:, :, t]*k[:, :, t]).sum(-1)
+    state = Tensor(initial).contiguous().realize()
+    out = gated_delta_prefill(Tensor(q), Tensor(k), Tensor(v), Tensor(beta), Tensor(alpha), state).realize()
+    np.testing.assert_allclose(out.numpy(), expected_out, rtol=1e-4, atol=1e-4)
+    np.testing.assert_allclose(state.numpy(), expected_state, rtol=1e-4, atol=1e-4)
+
   def _tensor_linspace(self, start:float, stop:float, shape:tuple[int, ...]) -> Tensor:
     return Tensor.linspace(start, stop, int(np.prod(shape)), dtype=dtypes.float32).reshape(*shape)
 

test/unit/test_call.py

@@ -224,6 +224,21 @@ class TestCallSchedule(unittest.TestCase):
     np.testing.assert_equal(x.numpy(), [2, 2, 2])
     np.testing.assert_equal(y.numpy(), [3, 3, 3])
 
+  def test_precompile_nested_scope_collision(self):
+    # a precompiled function body gets its own positional p{slot} params; they must not be renumbered when the call is
+    # scheduled inside an enclosing realize with a different slot ordering. the store must use this call's Variable
+    cache = Tensor.zeros(16)
+    @function(precompile=True, allow_implicit=True)
+    def store(x:Tensor, sp:UOp) -> Tensor:
+      # update a cache at a symbolic offset, like an attention KV cache update
+      return Tensor(cache.uop.after(cache[sp:sp+x.shape[0]].uop.store(x.uop)))[:sp+x.shape[0]].sum()
+    sp_v, nt_v = UOp.variable("sp", 0, 8), UOp.variable("nt", 1, 8)
+    t = Tensor.arange(16).float().realize()
+    sp, nt = sp_v.bind(0), nt_v.bind(8)
+    store(t[sp:sp+nt].clone().realize(), sp).realize()
+    np.testing.assert_equal(cache.numpy()[:8], t[:8].numpy())
+    np.testing.assert_equal(cache.numpy()[8:], np.zeros(8))
+
   def test_precompile_schedule_cache_hit(self):
     """two instances of the same @function should produce identical function body keys (schedule cache hit)"""
     @function(precompile=True)

test/unit/test_dtype_weak.py

@@ -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):
@@ -114,27 +114,21 @@ class TestWeakPromotion(unittest.TestCase):
     self.assertIsInstance((x + 2).src[1].val, float)
     self.assertIs(x + UOp.const(2), x + 2)
 
-  def test_index_dtype_ignores_weakness(self):
-    with Context(SPEC=2):
-      idx = UOp.const(0).cast(dtypes.int32)
-      weak = UOp.const(1.0).expand((1,))
-      self.assertEqual(UOp(Ops.INDEX, dtypes.float32, (weak, idx)).dtype, dtypes.float32)
-      with self.assertRaisesRegex(RuntimeError, "bad dtype"): UOp(Ops.INDEX, dtypes.int32, (weak, idx))
-
   def test_store_weak_value_uses_destination_dtype(self):
     with Context(DEFAULT_FLOAT=dtypes.float16):
       dst = UOp.param(0, dtypes.bfloat16, (1,)).index(UOp.const(0).cast(dtypes.int32))
       gate = UOp.const(True)
-      out = graph_rewrite(dst.store(UOp.const(5.0), gate), pm_lower_index_dtype, ctx={})
+      out = graph_rewrite(dst.store(UOp.const(5.0), gate), pm_commit_weak)
     # a bare weak CONST commits directly: the pass runs without symbolic, so a CAST here would survive it
     self.assertEqual((out.src[1], out.src[2]), (UOp.const(5.0, dtypes.bfloat16), gate))
 
   def test_weak_srcs_commit_only_at_a_concrete_lub(self):
     weak_lub = UOp(Ops.ADD, src=(UOp.const(1), UOp.const(1.0)))
-    self.assertIs(graph_rewrite(weak_lub, pm_lower_index_dtype, ctx={}), weak_lub)
+    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_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`)

test/unit/test_llm_amd.py

@@ -0,0 +1,107 @@
+import unittest
+import numpy as np
+from tinygrad import Tensor, UOp, dtypes, nn
+from tinygrad.llm.kernels.amd import Linear, amd_custom_kernels_supported, q8_quantize, flash_attention
+from tinygrad.llm.gguf import ggml_data_to_tensor
+
+class TestQ8Quantize(unittest.TestCase):
+  def test_word_quant_weights_use_typed_buffer_view(self):
+    for ggml_type, type_size in ((13, 176), (23, 136)):
+      with self.subTest(ggml_type=ggml_type):
+        raw = Tensor(np.zeros(type_size + 4, dtype=np.uint8), device="CPU").contiguous().realize()[4:]
+        decoded = ggml_data_to_tensor(raw, 256, ggml_type).reshape(1, 256)
+        linear = Linear(256, 1, bias=False)
+        linear.set_quantized(decoded)
+        self.assertEqual(linear.ggml_type, ggml_type)
+        self.assertEqual(linear.weight.dtype, dtypes.uint32)
+        self.assertEqual(linear.weight.nbytes(), type_size)
+        self.assertEqual(linear.weight.uop.buf_uop.buffer.offset, 4)
+
+  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 = 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)
+
+  def test_q6_linear_compiles(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    rng = np.random.default_rng(42)
+    packed = rng.integers(0, 256, 210, dtype=np.uint8)
+    packed[-2:] = np.array([0.01], dtype=np.float16).view(np.uint8)
+    raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:]
+    decoded = ggml_data_to_tensor(raw, 256, 14).reshape(1, 256)
+    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()))
+    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")
+    rng = np.random.default_rng(42)
+    in_features, blocks = 2048, 16*2048//256
+    packed = rng.integers(0, 256, blocks*144, dtype=np.uint8)
+    for i in range(blocks): packed[i*144:i*144+4] = np.array([0.01, 0.002], dtype=np.float16).view(np.uint8)
+    raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:]
+    decoded = ggml_data_to_tensor(raw, 16*in_features, 12).reshape(16, in_features)
+    weight = decoded.numpy()
+    linear = Linear(in_features, 16, bias=False)
+    nn.state.load_state_dict(linear, {"weight":decoded}, verbose=False, realize=False)
+    x = rng.normal(size=(3, in_features)).astype(np.float32)
+    scale = np.maximum(np.abs(x).reshape(3, in_features//32, 32).max(-1, keepdims=True) / 127, 1e-8)
+    xq = np.clip(np.rint(x.reshape(3, in_features//32, 32) / scale), -127, 127) * scale
+    np.testing.assert_allclose(linear(Tensor(x)).numpy(), xq.reshape(3, in_features) @ weight.T, rtol=2e-3, atol=2e-2)
+    self.assertEqual(linear.ggml_type, 12)
+
+  def test_q6_linear_multiple_tokens(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    rng = np.random.default_rng(42)
+    in_features, blocks = 2048, 16*2048//256
+    packed = rng.integers(0, 256, blocks*210, dtype=np.uint8)
+    for i in range(blocks): packed[i*210+208:i*210+210] = np.array([0.01], dtype=np.float16).view(np.uint8)
+    raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:]
+    decoded = ggml_data_to_tensor(raw, 16*in_features, 14).reshape(16, in_features)
+    weight = decoded.numpy()
+    linear = Linear(in_features, 16, bias=False)
+    nn.state.load_state_dict(linear, {"weight":decoded}, verbose=False, realize=False)
+    x = rng.normal(size=(3, in_features)).astype(np.float32)
+    scale = np.maximum(np.abs(x).reshape(3, in_features//32, 32).max(-1, keepdims=True) / 127, 1e-8)
+    xq = np.clip(np.rint(x.reshape(3, in_features//32, 32) / scale), -127, 127) * scale
+    np.testing.assert_allclose(linear(Tensor(x)).numpy(), xq.reshape(3, in_features) @ weight.T, rtol=2e-3, atol=2e-2)
+    self.assertEqual(linear.ggml_type, 14)
+
+    # symbolic token counts take the padded kernel path and give the same results
+    generic = Linear(in_features, 16, bias=False)
+    nn.state.load_state_dict(generic, {"weight":decoded}, verbose=False, realize=False)
+    sym = Tensor(np.concatenate([x, np.zeros((1, in_features), np.float32)])).contiguous()[:UOp.variable("tokens", 1, 4).bind(3)]
+    np.testing.assert_allclose(generic(sym)[:3].numpy(), xq.reshape(3, in_features) @ weight.T, rtol=2e-3, atol=2e-2)
+    self.assertTrue(generic.use_custom_quant)
+    self.assertEqual(generic.ggml_type, 14)
+
+  def test_attention_uses_physical_cache_length(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    q, k, v = Tensor.zeros(1, 2, 1, 32), Tensor.randn(1, 1, 1, 32), Tensor.randn(1, 1, 1, 32)
+    cache = Tensor.empty(2, 1, 1, 256, 32, dtype=dtypes.half).contiguous()
+    assigned = Tensor(cache.uop.after(cache[:, :, :, 0:1, :].uop.store(Tensor.stack(k, v).cast(dtypes.half).uop)))
+    out = flash_attention(q, assigned, 1).realize()
+    np.testing.assert_allclose(out.numpy(), v.expand(1, 2, 1, 32).numpy(), rtol=2e-2, atol=2e-2)
+
+  def test_prefill_attention_unaligned_start(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    rng = np.random.default_rng(42)
+    start_pos = 1718
+    q = Tensor.zeros(1, 8, 32, 128)
+    old_kv = rng.normal(size=(2, 1, 1, start_pos, 128)).astype(np.float32)
+    new_kv = rng.normal(size=(2, 1, 1, 32, 128)).astype(np.float32)
+    cache = Tensor.zeros(2, 1, 1, 2048, 128, dtype=dtypes.half).contiguous()
+    Tensor.realize(cache[:, :, :, :start_pos].assign(Tensor(old_kv).cast(dtypes.half)))
+    sp = UOp.variable("start_pos", 0, 2047).bind(start_pos)
+    assigned = Tensor(cache.uop.after(cache[:, :, :, sp:sp+32, :].uop.store(Tensor(new_kv).cast(dtypes.half).uop)))
+    out = flash_attention(q, assigned, sp+32).realize()
+    values = np.concatenate([old_kv[1, 0, 0], new_kv[1, 0, 0]]).astype(np.float16).astype(np.float32)
+    expected = np.stack([values[:start_pos+i+1].mean(0) for i in range(32)])[None, None].repeat(8, axis=1)
+    np.testing.assert_allclose(out.numpy(), expected, rtol=2e-3, atol=2e-3)
+
+if __name__ == "__main__": unittest.main()

test/unit/test_llm_server.py

@@ -1,6 +1,8 @@
 import unittest
+import numpy as np
 from unittest.mock import patch
 from tinygrad import Tensor, UOp
+from tinygrad.nn.state import get_state_dict
 from tinygrad.schedule import schedule_cache
 from tinygrad.llm.model import Transformer, TransformerConfig
 from tinygrad.llm.serve import StreamRouter
@@ -42,7 +44,10 @@ class TestTransformerGenerate(unittest.TestCase):
       return Tensor([[42]])
     with patch.object(Transformer, '__call__', mock_call):
       next(model.generate([1, 2, 3, 4, 5, 42, 10]))
-    self.assertEqual(calls, [((1, 1), V_START_POS.bind(5)), ((1, 1), V_START_POS.bind(6))])
+    # resumes from the reused state at position 5 and consumes the 2 new tokens (one chunk or two decode steps)
+    self.assertEqual(calls[0][1], V_START_POS.bind(5))
+    def ntok(shape): return shape[1] if isinstance(shape[1], int) else shape[1].unbind()[1]
+    self.assertEqual(sum(ntok(c[0]) for c in calls), 2)
 
   def test_recurrent_divergent_prompt_restarts(self):
     model, calls = Transformer(TEST_CONFIG), []
@@ -152,6 +157,22 @@ class TestTransformerGenerate(unittest.TestCase):
     # 4 tokens, chunk_size=4 -> 1 prefill chunk
     self.assertEqual(get_prefill_flags(list(range(4)), 4), [True, False, False])
 
+  def test_chunked_prefill_kv_cache_matches_single_chunk(self):
+    config = TransformerConfig(num_blocks=1, dim=8, hidden_dim=16, n_heads=1, n_kv_heads=1, norm_eps=1e-5,
+      vocab_size=32, head_dim=4, rope_theta=1000000, rope_dim=4, qk_norm=4, v_head_dim=4, max_context=16)
+    def model():
+      m = Transformer(config)
+      rng = np.random.RandomState(1234)
+      for t in get_state_dict(m).values():
+        t.assign(Tensor(rng.uniform(-1, 1, t.shape).astype(np.float32))).realize()
+      return m
+    def prefill(m, chunk_size):
+      gen = m.generate(list(range(1, 9)), chunk_size=chunk_size, temperature=0.0)
+      next(gen)
+      return [b.cache_kv.numpy() for b in m.blk]
+    for g, r in zip(prefill(model(), 4), prefill(model(), 8)):
+      np.testing.assert_allclose(g[:, :, :, :8, :], r[:, :, :, :8, :], atol=1e-5)
+
   def test_kv_cache_resume_matches_fresh(self):
     model = Transformer(TEST_CONFIG)
 

tinygrad/codegen/__init__.py

@@ -1,9 +1,9 @@
 from dataclasses import replace, dataclass
 import itertools, functools
 from tinygrad.helpers import DISABLE_FAST_IDIV, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC
-from tinygrad.helpers import ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT, TracingKey, Context, panic
+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
@@ -233,10 +233,11 @@ pm_reduce_local = pm_wmma_add+PatternMatcher([
   (UPat(Ops.SINK, name="sink"), merge_reduce_ends),
 ])+pm_clean_up_group_sink
 
+def is_shape_changing_bitcast(u:UOp): return u.op is Ops.BITCAST and u.shape != u.src[0].shape
 def maybe_load(u:UOp): return u.load() if u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL, AddrSpace.REG) else u
 pm_add_loads = PatternMatcher([
-  # BITCAST?
-  (UPat(GroupOp.Elementwise|{Ops.REDUCE,Ops.WMMA,Ops.STACK}, name="x"), lambda x: x.replace(src=tuple([maybe_load(u) for u in x.src]))),
+  (UPat(GroupOp.Elementwise|{Ops.REDUCE,Ops.WMMA,Ops.STACK}, name="x"),
+   lambda x: None if is_shape_changing_bitcast(x) else x.replace(src=tuple(map(maybe_load, x.src)))),
   (UPat(Ops.STORE, name="x"), lambda x: x.replace(src=(x.src[0], maybe_load(x.src[1]))+x.src[2:])),
 ])
 
@@ -281,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))
@@ -346,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")
@@ -374,9 +373,12 @@ 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")
 
+  # 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")
 
@@ -387,10 +389,6 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
   num_params = len([x for x in sink.toposort() if x.op is Ops.PARAM and x.arg.slot != -1])
   sink = graph_rewrite(sink, pm_number_params, ctx=[num_params], name="number params with -1", walk=True)
 
-  # 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)
-
   if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Output AST")
   if SPEC: type_verify(sink, spec_program)
 
@@ -459,7 +457,7 @@ pm_to_program = PatternMatcher([
   (UPat(Ops.PROGRAM, src=(UPat(), UPat(Ops.LINEAR), UPat(Ops.SOURCE, name="source")), name="prg"), do_compile),
 ])
 
-@rewrite_group(name=lambda ast,renderer,ret,**kwargs: TracingKey(ret.src[0].arg.name,(ret.src[0].arg.function_name, ast), ret=renderer), replay=True)
+@rewrite_group(name=lambda ast,renderer,ret,**_: TracingKey((k:=ret.src[0].arg).name,(k.function_name, ast, ret.key),ret=renderer), replay=True)
 @Context(ALLOW_DEVICE_USAGE=0)
 def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
   """
@@ -488,9 +486,14 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp:
   if VIZ: graph_rewrite(prg, PatternMatcher([]), name="View Program")
   return prg
 
+# config affects generated programs and cache keys; context also carries compile-only behavior to workers
+to_program_config = (NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32,
+                     DEFAULT_FLOAT, DEFAULT_INT, NUM_CPU_THREADS, TC_SELECT, TC_OPT)
+to_program_context = (*to_program_config, SPEC, DEBUG)
+def to_program_key(ast:UOp, renderer:Renderer) -> tuple:
+  return (ast.key, type(renderer), renderer.target, *[x.value for x in to_program_config])
+
 to_program_cache: dict[tuple, UOp] = {}
 def to_program(ast:UOp, renderer:Renderer) -> UOp:
-  config = (NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC, IMAGE, DISABLE_FAST_IDIV, TRANSCENDENTAL, ALLOW_TF32, DEFAULT_FLOAT, DEFAULT_INT)
-  key = (ast.key, type(renderer), renderer.target, *[x.value for x in config])
-  if (prg:=to_program_cache.get(key)) is None: to_program_cache[key] = prg = do_to_program(ast, renderer)
+  if (prg:=to_program_cache.get(key:=to_program_key(ast, renderer))) is None: to_program_cache[key] = prg = do_to_program(ast, renderer)
   return prg

tinygrad/codegen/decomp/dtype.py

@@ -1,8 +1,9 @@
 from dataclasses import replace
 from tinygrad.dtype import dtypes, DType, truncate
-from tinygrad.helpers import flatten, DEBUG, EMULATED_DTYPES, Context, SPEC
+from tinygrad.helpers import flatten, DEBUG, EMULATED_DTYPES
 from tinygrad.uop import GroupOp
-from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher, graph_rewrite, ParamArg
+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
 
@@ -25,10 +26,10 @@ def l2i(op: Ops, dt: DType, *uops:UOp):
   match op:
     case Ops.NEG: return l2i(Ops.SUB, dt, zero, zero, *uops)
     case Ops.CAST if dt in (dtypes.long, dtypes.ulong) and uops[0].dtype not in dtypes.floats:
-      # the high word is the sign extension; bool has no sign, test the already-cast low word instead (bool < 0 would promote to weakint)
+      # the high word is the sign extension, and unsigned and bool sources zero extend
       x, lo = uops[0], uops[0].cast(l2i_dt[dt])
-      sign = lo if x.dtype is dtypes.bool else x
-      return lo, (sign < sign.const_like(0)).where(lo.const_like(-1), lo.const_like(0))
+      if x.dtype is dtypes.bool or x.dtype in dtypes.uints: return lo, lo.const_like(0)
+      return lo, (x < x.const_like(0)).where(lo.const_like(-1), lo.const_like(0))
     case Ops.CAST if dt in (dtypes.long, dtypes.ulong):
       return (lo:=uops[0].cast(l2i_dt[dt])), (uops[0] / 2**32).cast(l2i_dt[dt]) - ((uops[0] < 0) & lo.ne(0))
     case Ops.CAST if dt in dtypes.floats:
@@ -127,19 +128,22 @@ def f2f_clamp(val:UOp, dt:DType, sat=True) -> UOp:
   return val.ne(val).where(val, (val < -mx).where(-sat, (mx < val).where(sat, val)))
 
 def f2f_load(x: UOp, fr:DType, to:DType) -> UOp:
-  if (n:=x.max_numel()) == 1: return f2f(x.replace(dtype=f2f_dt[fr]), fr, to)
-  return UOp(Ops.STACK, src=tuple(f2f(x.replace(dtype=f2f_dt[fr], src=(reindex(x.src[0], i, 1),)), fr, to) for i in range(n)))
+  storage_idx = graph_rewrite(x.src[0], pm_float_decomp, ctx=(fr, to), bottom_up=True)
+  if (n:=x.max_numel()) == 1: return f2f(storage_idx.load(*x.src[1:]), fr, to)
+  return UOp(Ops.STACK, src=tuple(f2f(reindex(storage_idx, i, 1).load(*x.src[1:]), fr, to) for i in range(n)))
 
 def f2f_store(st, idx, val, fr:DType, to:DType):
   if (n:=val.max_numel()) == 1: return st.replace(src=(idx, f2f(val.bitcast(f2f_dt[to]), to, fr)))
   return UOp.group(*(st.replace(src=(reindex(idx, i, 1), f2f(val.index(i).bitcast(f2f_dt[to]), to, fr))) for i in range(n)))
 
 # tag is the 32-bit word this node becomes - (0 for the low word, 1 for the high, the dtype the consumer wants)
-pm_long_decomp = PatternMatcher([
-  (UPat(GroupOp.Defines, src=(UPat.var("sz"),), name="x"), lambda x,sz:
-   x.replace(dtype=l2i_dt[x.dtype], arg=replace(x.arg, dtype=l2i_dt[x.dtype]), src=(sz*2,)) if x.dtype in l2i_dt else None),
+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:
-   reindex(x, x.tag[0]).replace(dtype=x.tag[1], tag=None) if x.tag is not None else None),
+   reindex(x, x.tag[0]).replace(tag=None) if x.tag is not None else None),
   (UPat(Ops.STORE, src=(UPat.var('idx', tuple(l2i_dt.keys())), UPat.var('val')), name='st'), lambda st,idx,val:
    st.replace(src=(idx.rtag((0, dt:=l2i_dt[idx.dtype])), val.rtag((0, dt)))).group(
      st.replace(src=(idx.rtag((1, dt)), val.rtag((1, dt))))) if val.tag is None else None),
@@ -147,6 +151,9 @@ pm_long_decomp = 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:
@@ -159,21 +166,22 @@ pm_long_decomp = PatternMatcher([
   (UPat((*(GroupOp.ALU - GroupOp.Comparison - {Ops.SHL, Ops.SHR, Ops.WHERE}), Ops.BITCAST), tuple(l2i_dt.keys()), name="x"), lambda ctx,x:
    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 x,idx:
-   x.replace(dtype=l2i_dt[x.dtype], src=(reindex(idx, x.tag[0]).replace(dtype=l2i_dt[x.dtype], tag=None),), tag=None) if x.tag is not None else None),
-  (UPat(Ops.CONST, tag={(w, dt) for w in (0, 1) for dt in l2i_dt.values()}, name='x'), lambda x:
-   UOp.const(truncate[x.tag[1]]((x.val >> 32) if x.tag[0] == 1 else (x.val & 0xFFFFFFFF)), x.tag[1]))
+  (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)
 ])
 
 # float decomposition patterns - ctx is (fr, to) tuple
-pm_float_decomp = PatternMatcher([
-  (UPat((*GroupOp.Defines, Ops.INDEX, Ops.SHRINK), name="x"), lambda ctx,x:
-   x.replace(dtype=f2f_dt[ctx[0]], arg=replace(x.arg, dtype=f2f_dt[ctx[0]]) if isinstance(x.arg, ParamArg) else x.arg, tag=ctx[0])
-   if x.dtype == ctx[0] and (x.op is not Ops.INDEX or x.src[0].op not in {Ops.LOAD, Ops.STACK}) else None),
+pm_float_decomp: PatternMatcher = PatternMatcher([
+  (UPat(GroupOp.Defines, name="x"), lambda ctx,x:
+   UOp(x.op, src=x.src, arg=replace(x.arg, dtype=f2f_dt[ctx[0]]), tag=ctx[0]) if x.dtype == ctx[0] else None),
+  # INDEX into a LOAD/STACK selects a lane of an already converted value, the load rules below own those
+  (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat(GroupOp.All-{Ops.LOAD, Ops.STACK}),), allow_any_len=True, name="x"), lambda ctx,x:
+   UOp(x.op, src=(graph_rewrite(x.src[0], pm_float_decomp, ctx=ctx, bottom_up=True), *x.src[1:]), arg=x.arg, tag=ctx[0])
+   if x.dtype == ctx[0] else None),
   (UPat(Ops.LOAD, dtypes.floats, name="x"), lambda ctx,x: f2f_load(x, *ctx) if x.dtype == ctx[0] else None),
   # bitcasted load should just replace load
   (UPat(Ops.BITCAST, src=(UPat(Ops.LOAD, name="ld"),), name="bc"), lambda ctx,bc,ld:
-   ld.replace(dtype=f2f_dt[ctx[0]]).bitcast(bc.dtype) if ld.dtype == ctx[0] else None),
+   graph_rewrite(ld.src[0], pm_float_decomp, ctx=ctx, bottom_up=True).load(*ld.src[1:]).bitcast(bc.dtype) if ld.dtype == ctx[0] else None),
   # bitcast from
   (UPat(Ops.BITCAST, src=(UPat.var("x", dtypes.floats),), name="bc"), lambda ctx,bc,x:
    bc.replace(src=(f2f(x.bitcast(f2f_dt[ctx[1]]), ctx[1], ctx[0]),)) if x.dtype == ctx[1] and bc.dtype.bitsize == ctx[0].bitsize else None),
@@ -182,26 +190,21 @@ pm_float_decomp = 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:
-   x.replace(dtype=ctx[1], src=tuple(s.cast(ctx[1]) if s.dtype == ctx[0] else s for s in x.src))
-   if x.dtype == ctx[0] else None),
+   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:
-   st.replace(src=(idx, val.replace(dtype=f2f_dt[ctx[0]]))) if val.dtype == ctx[0] and idx.tag == ctx[0] else None),
-  (UPat(Ops.STORE, src=(UPat.var("idx"), UPat.var("val", dtypes.floats)), name='st'), lambda ctx,st,idx,val:
-   f2f_store(st, idx, val, *ctx) if val.dtype == ctx[1] and (idx:=idx.src[0] if idx.op == Ops.CAST else idx).tag == ctx[0] else None),
+   st.replace(src=(idx, val.src[0].bitcast(f2f_dt[ctx[0]]))) if val.dtype == ctx[0] and idx.tag == ctx[0] else None),
+  (UPat(Ops.STORE, src=(UPat.var("idx").or_casted(), UPat.var("val", dtypes.floats)), name='st'), lambda ctx,st,idx,val:
+   f2f_store(st, idx, val, *ctx) if val.dtype == ctx[1] and idx.tag == ctx[0] else None),
 ])
 
 def do_dtype_decomps(sink:UOp, ctx:tuple[set[DType], Renderer]) -> UOp:
   def _should_emulate(dt): return dt in EMULATED_DTYPES.tolist(dtypes) or dt not in ctx[1].supported_dtypes()
-  # NOTE: dtype decomp creates intermediate UOps that don't follow the spec (e.g. half LOAD on ushort BUFFER)
-  with Context(SPEC=min(SPEC.value, 1)):
-    for fr in sorted(filter(_should_emulate, ctx[0])):
-      to = dtypes.int if fr == dtypes.long else dtypes.half if not _should_emulate(dtypes.half) and fr in dtypes.fp8s else dtypes.float
-      if DEBUG >= 2: print(f"emulating {fr} as {to}")
-      pm = pm_float_decomp if fr in dtypes.floats else pm_long_decomp
-      sink = graph_rewrite(sink, pm, name=f"decomp {fr} -> {to}", ctx={} if pm is pm_long_decomp else (fr, to), bottom_up=True)
+  for fr in sorted(filter(_should_emulate, ctx[0])):
+    to = dtypes.int if fr == dtypes.long else dtypes.half if not _should_emulate(dtypes.half) and fr in dtypes.fp8s else dtypes.float
+    if DEBUG >= 2: print(f"emulating {fr} as {to}")
+    pm = pm_float_decomp if fr in dtypes.floats else pm_long_decomp
+    sink = graph_rewrite(sink, pm, name=f"decomp {fr} -> {to}", ctx={} if pm is pm_long_decomp else (fr, to), bottom_up=True)
   ctx[0].clear()
   return sink
 

tinygrad/codegen/decomp/op.py

@@ -75,7 +75,11 @@ powers_of_two: dict[int, int] = {2**i:i for i in range(64)}
 @functools.cache
 def get_simplifying_rewrite_patterns(ops:tuple[Ops, ...]) -> PatternMatcher:
   # these are rewrites that make things simpler
-  pat: list[tuple[UPat, Callable]] = [(UPat.var("a")//UPat.var("b"), floordiv_to_idiv)]
+  pat: list[tuple[UPat, Callable]] = []
+  # FLOORDIV by 2**y -> x >> y (an arithmetic shift is exactly floor division for any sign); fires before floordiv_to_idiv
+  if Ops.SHR in ops: pat.append((UPat.var("x", dtypes.ints)//UPat.cvar("c"),
+    lambda x,c: x >> v if (v:=powers_of_two.get(c.val, 0)) else None))
+  pat.append((UPat.var("a")//UPat.var("b"), floordiv_to_idiv))
   # FLOORMOD by 2**y -> x & (2**y-1) (correct floor mod for any sign in two's complement); fires before floormod_to_mod
   if Ops.AND in ops: pat.append((UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.val-1) if c.val in powers_of_two else None))
   pat.append((UPat.var("a")%UPat.var("b"), floormod_to_mod))
@@ -128,6 +132,6 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa
     if Ops.SHL in ops: pat += [(UPat.var('x').alu(Ops.SHL, UPat.cvar('n'))+UPat.var('c'), lambda x,n,c: x.alu(Ops.MULACC, x.const_like(1<<n.val), c))]
   # some backends emit FDIV for RECIP, in that case: a*(1/b) -> a/b
   if Ops.FDIV in ops:
-    pat += [(UPat.var("x").reciprocal(), lambda x: x.const_like(1).alu(Ops.FDIV, x))]
-    pat += [(UPat.var("a", dtypes.floats) * UPat(Ops.FDIV, dtypes.floats, src=(UPat.const(1), UPat.var("b"))), lambda a,b: a.alu(Ops.FDIV, b))]
+    pat += [(UPat.var("x").reciprocal(), lambda x: UOp.const(1.0).alu(Ops.FDIV, x))]
+    pat += [(UPat.var("a") * UPat(Ops.FDIV, dtypes.floats, src=(UPat.const(1), UPat.var("b"))), lambda a,b: a.alu(Ops.FDIV, b))]
   return PatternMatcher(pat)

tinygrad/codegen/opt/search.py

@@ -1,12 +1,13 @@
-import math, time, multiprocessing, traceback, signal, atexit
+import math, time, traceback, signal
 from dataclasses import replace
 from tinygrad.uop.ops import sym_infer, AxisType, UOp, Ops
 from tinygrad.uop.render import pyrender
 from tinygrad.device import Device, Buffer
-from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, time_to_str
+from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, colored, time_to_str
 from tinygrad.helpers import IGNORE_BEAM_CACHE
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError
 from tinygrad.engine.realize import time_call
+from tinygrad.engine.worker import get_worker_pool, terminate_worker_pool
 from tinygrad.codegen import to_program
 from tinygrad.codegen.opt.postrange import Scheduler
 
@@ -78,11 +79,6 @@ def _try_compile(x:tuple[int,Scheduler]) -> tuple[int, tuple[UOp, float]|None]:
     if hasattr(signal, "alarm"): signal.alarm(0)
   return x[0], ret
 
-# workers should not open devices and should ignore ctrl c and should not launch VIZ
-def _init_worker():
-  Context(ALLOW_DEVICE_USAGE=0, VIZ=0, TRACK_MATCH_STATS=0).__enter__()
-  signal.signal(signal.SIGINT, signal.SIG_IGN)
-
 def _ensure_buffer_alloc(bufs:list[Buffer]) -> list[Buffer]: return [buf.ensure_allocated() if buf is not None else buf for buf in bufs]
 
 # *** external API ***
@@ -111,9 +107,8 @@ def get_kernel_actions(s:Scheduler, include_0=True, max_up:int|None=None) -> dic
     except KernelOptError: pass
   return acted
 
-beam_pool, BEAM_DEBUG = None, getenv("BEAM_DEBUG")
+BEAM_DEBUG = getenv("BEAM_DEBUG")
 def beam_search(s:Scheduler, rawbufs:list[Buffer], var_vals:dict[str,int], amt:int, allow_test_size=True, disable_cache=IGNORE_BEAM_CACHE.value):
-  global beam_pool
   key = {"ast": s.ast.key, "amt": amt, "allow_test_size": allow_test_size, "device": s.ren.target.device, "suffix": s.ren.suffix}
   if not disable_cache and CACHELEVEL >= 1 and (val:=diskcache_get("beam_search", key)) is not None:
     ret = s.copy()
@@ -123,11 +118,7 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], var_vals:dict[str,int], amt:i
   beam: list[tuple[Scheduler, float]] = [(s, float("inf"))]
   seen_libs = set()
 
-  default_parallel = multiprocessing.cpu_count() if s.ren.target.device in {"CUDA", "AMD", "NV", "METAL", "HIP"} else 0
-  if beam_pool is None and (workers := getenv("PARALLEL", default_parallel)):
-    beam_pool = multiprocessing.get_context("spawn").Pool(workers, _init_worker, (), getenv("BEAM_MAX_TASKS_PER_CHILD", 16))
-    @atexit.register
-    def close_pool(): beam_pool.close()
+  pool = get_worker_pool()
 
   min_progress = getenv("BEAM_MIN_PROGRESS", 0.01)/1e6
   if BEAM_DEBUG:
@@ -143,7 +134,7 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], var_vals:dict[str,int], amt:i
       candidates: list[Scheduler] = flatten([get_kernel_actions(si, include_0=False).values() for si,_ in beam])
       timed: list[tuple[Scheduler, float]] = []
       least_compute_ops = math.inf
-      for i, proc in ((map if beam_pool is None else beam_pool.imap_unordered)(_try_compile, enumerate(candidates))):
+      for i, proc in ((map if pool is None else pool.imap_unordered)(_try_compile, enumerate(candidates))):
         if proc is None: continue
         prg, compile_et = proc
         if (lib:=prg.src[3].arg) in seen_libs: continue
@@ -179,7 +170,7 @@ def beam_search(s:Scheduler, rawbufs:list[Buffer], var_vals:dict[str,int], amt:i
         print(f"\r{time.perf_counter() - st:7.2f}s:", colored(time_to_str(beam[0][1], w=12), "green" if exiting else None),
               f"from {len(candidates):3d} -> {len(opts):3d} actions\033[K", beam[0][0].colored_shape())
   except KeyboardInterrupt as e:
-    if beam_pool is not None: beam_pool.terminate()
+    terminate_worker_pool()
     raise e
 
   if CACHELEVEL >= 1: diskcache_put("beam_search", key, beam[0][0].applied_opts)

tinygrad/device.py

@@ -66,10 +66,10 @@ def canonicalize_device(device:str|tuple|list|None) -> str|tuple[str, ...]:
 class ProfileDeviceEvent(ProfileEvent): device:str; tdiff:decimal.Decimal=decimal.Decimal(0); props:dict[str,Any]|None=None # noqa: E702
 
 @dataclass(frozen=True)
-class ProfileProgramEvent(ProfileEvent): device:str; name:str; lib:bytes|None; base:int|None; tag:int|None=None # noqa: E702
+class ProfileProgramEvent(ProfileEvent): device:str; name:str; lib:bytes|None; base:int|None; tag:int|None=None; profile_key:bytes|None=None # noqa: E702
 
 @dataclass(frozen=True)
-class ProfileGraphEntry: device:str; name:str|TracingKey; st_id:int; en_id:int # noqa: E702
+class ProfileGraphEntry: device:str; name:str|TracingKey; st_id:int; en_id:int; profile_key:bytes|None=None # noqa: E702
 
 @dataclass(frozen=True)
 class ProfileGraphEvent(ProfileEvent): ents:list[ProfileGraphEntry]; deps:list[list[int]]; sigs:list[decimal.Decimal] # noqa: E702
@@ -83,6 +83,7 @@ class BufferSpec:
   cpu_access: bool = False
   host: bool = False
   nolru: bool = False
+  zero: bool = False
   external_ptr: int|None = None
 
 class MultiBuffer:
@@ -265,7 +266,7 @@ class LRUAllocator(Allocator, Generic[DeviceType]):
       for opaque in opaques: super().free(opaque, sz, options)
       opaques.clear()
   def free(self, opaque:Any, size:int, options:BufferSpec|None=None):
-    if LRU and (options is None or (not options.nolru and options.external_ptr is None)): self.cache[(size, options)].append(opaque)
+    if LRU and (options is None or (not (options.nolru or options.zero) and options.external_ptr is None)): self.cache[(size, options)].append(opaque)
     else: super().free(opaque, size, options)
 
 class DepsTracker:
@@ -326,6 +327,7 @@ class TinyELF:
   target: Target
   # tuple of (name, slot, dtype, shape)
   signature: tuple[tuple[str|None, int, DType, tuple], ...]
+  profile_key: bytes|None = None
 
   @staticmethod
   def iter_sig(signature:tuple[tuple[str|None, int, DType, tuple], ...], offset:int=0) -> Generator[tuple[int, DType], None, None]:

tinygrad/engine/jit.py

@@ -211,7 +211,7 @@ def _prepare_jit_inputs(args, kwargs):
   # collect buffer UOps (including MultiBuffer)
   input_buf_uops: list[UOp] = [u.base for u in input_uops if u.base.realized is not None]
   if len(set(input_buf_uops)) != len(input_buf_uops): raise JitError("duplicate inputs to JIT")
-  inputs = [(*(u.substitute({u.base:UOp(Ops.NOOP, u.base.dtype)}, extra_pm=mop_cleanup).unbind_all()), u.dtype, u.device) for u in input_uops]
+  inputs = [(*(u.substitute({u.base:UOp(Ops.NOOP)}, extra_pm=mop_cleanup).unbind_all()), u.dtype, u.device) for u in input_uops]
   _var_vals = merge_dicts([x[1] for x in inputs] + [dict(v.unbind() for v in (args + tuple(kwargs.values())) if isinstance(v, UOp))])
   var_vals = {k.expr:v for k,v in _var_vals.items()}
   expected_input_info = [(x[0], tuple(sorted(x[1].keys(), key=lambda v: v.expr)), x[2], x[3]) for x in inputs]

tinygrad/engine/realize.py

@@ -2,14 +2,15 @@ from __future__ import annotations
 from typing import cast, Iterator, Any, Sequence
 import random, itertools, math, weakref, array, decimal
 from dataclasses import dataclass, replace, field
-from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansipad, all_int, prod, flatten, Context, getenv, to_tuple
+from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansipad, all_int, prod, flatten, Context, getenv, to_tuple, tqdm
 from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, perf_counter_us
-from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite
+from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite, ProgramInfo
 from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
 from tinygrad.dtype import dtypes
-from tinygrad.renderer import Estimates
-from tinygrad.codegen import to_program
+from tinygrad.renderer import Estimates, Renderer
+from tinygrad.codegen import to_program, to_program_cache, to_program_key, to_program_context
 from tinygrad.codegen.opt.postrange import args_from_ast
+from tinygrad.engine.worker import get_worker_pool, terminate_worker_pool
 
 # **************** Helpers ****************
 
@@ -221,7 +222,7 @@ def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
   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)
+    (d:=cast(Any, Device[device])).prof_ents[prof[0]] = ProfileGraphEntry(device, stat_call.arg.name, prof[0], prof[1], stat_call.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)
@@ -247,10 +248,44 @@ pm_beam = PatternMatcher([
    lambda ctx,call,sink: call.replace(src=(sink.replace(arg=replace(sink.arg, beam=ctx)), *call.src[1:])) if sink.arg.beam == 0 else None),
 ])
 
-pm_compile = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.PROGRAM), name="ast"),), name="call", allow_any_len=True), lambda call,ast:
-    call.replace(src=(to_program(ast, Device[call.device if isinstance(call.device, str) else call.device[0]].renderer), *call.src[1:]))),
-])
+# **************** parallel lowering + compilation ****************
+
+def _compile_kernel(x:tuple[int, tuple[UOp, Renderer], dict]) -> tuple[int, UOp]:
+  with Context(**x[2]): return x[0], to_program(*x[1])
+
+def _get_call_to_compile(c:UOp) -> tuple[UOp, Renderer]|None:
+  ast = a0.src[0] if (a0:=c.src[0]).op is Ops.CUSTOM_FUNCTION and a0.arg == "hcq" else a0
+  # a PROGRAM with a ProgramInfo and a BINARY is already compiled
+  if ast.op is Ops.SINK or (ast.op is Ops.PROGRAM and not (isinstance(ast.arg, ProgramInfo) and ast.src[-1].op is Ops.BINARY)):
+    return ast, Device[c.device if isinstance(c.device, str) else c.device[0]].renderer
+  return None
+
+def lower_and_compile(linear:UOp) -> UOp:
+  # collect the kernels to lower and compile, deduped by their compile cache key
+  if not len(ar:={c: a for c in linear.toposort() if c.op is Ops.CALL and (a:=_get_call_to_compile(c)) is not None}): return linear
+
+  # lower and compile what's not cached, in parallel if there's a worker pool
+  keys = {c: to_program_key(*a) for c, a in ar.items()}
+  todo = list({keys[c]: a for c, a in ar.items() if keys[c] not in to_program_cache}.items())
+  if len(todo):
+    # kernels that beam search must compile in the parent, beam needs device access to time candidates
+
+    pool = None if len(todo) == 1 or any(getattr(c.src[0].arg, "beam", 0) for c in ar) else get_worker_pool()
+    ctx = {v.key: v.value for v in to_program_context}
+    tasks = ((i, ast_ren, ctx) for i, (_, ast_ren) in enumerate(todo))
+    try:
+      with tqdm(total=len(todo), desc="compiling", disable=DEBUG<1) as pbar:
+        for i, prg in (map if pool is None else pool.imap_unordered)(_compile_kernel, tasks):
+          pbar.set_description(f"compiling {ansipad(prg.src[0].arg.name, 40)}")
+          to_program_cache[todo[i][0]] = prg
+          pbar.update(1)
+    except KeyboardInterrupt:
+      if pool is not None: terminate_worker_pool()
+      raise
+
+  # swap the compiled PROGRAMs into the calls
+  return linear.substitute({c: c.replace(src=(c.src[0].substitute({a[0]: to_program_cache[keys[c]]}), *c.src[1:])) for c, a in ar.items()},
+                           name="precompile kernels")
 
 pm_optimize_local_size = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), optimize_local_size),
@@ -270,7 +305,7 @@ if getenv("HCQ2"): from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_li
 def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, profile:bool|None=None) -> UOp:
   if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
   if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
-  linear = graph_rewrite(linear, pm_compile, name="precompile kernels", walk=True)
+  linear = lower_and_compile(linear)
   linear = graph_rewrite(linear, pm_optimize_local_size, name="optimize local size", walk=True)
   if getenv("HCQ2"): linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile)
   return linear

tinygrad/engine/worker.py

@@ -0,0 +1,49 @@
+import multiprocessing, atexit, signal, sys, threading, contextlib
+from multiprocessing.context import SpawnContext, SpawnProcess
+from tinygrad.helpers import Context, getenv, PARALLEL
+
+# generic pool of worker processes for parallel compilation, shared by kernel lowering and BEAM search
+
+# workers should not open devices and should ignore ctrl c and should not launch VIZ
+def _init_worker():
+  Context(ALLOW_DEVICE_USAGE=0, VIZ=0, TRACK_MATCH_STATS=0).__enter__()
+  signal.signal(signal.SIGINT, signal.SIG_IGN)
+
+# spawn normally reimports the user's __main__ before _init_worker. This replays top-level code and can recursively create pools. There is no public
+# multiprocessing switch to skip that import, so hide the two attributes used to locate __main__ while each worker (including replacements) starts.
+_spawn_lock, _missing = threading.Lock(), object()
+@contextlib.contextmanager
+def _without_main():
+  main = sys.modules.get("__main__")
+  if main is None:
+    yield
+    return
+  with _spawn_lock:
+    saved = {name:getattr(main, name, _missing) for name in ("__file__", "__spec__")}
+    try:
+      for name in saved: setattr(main, name, None)
+      yield
+    finally:
+      for name,value in saved.items(): delattr(main, name) if value is _missing else setattr(main, name, value)
+
+class _WorkerProcess(SpawnProcess):
+  @staticmethod
+  def _Popen(process_obj):
+    with _without_main(): return SpawnProcess._Popen(process_obj)
+
+class _WorkerContext(SpawnContext): Process = _WorkerProcess
+
+worker_pool = None
+def get_worker_pool():
+  global worker_pool
+  if multiprocessing.current_process().daemon or PARALLEL == 0: return None
+  if worker_pool is None:
+    worker_pool = _WorkerContext().Pool(PARALLEL.value, _init_worker, (), getenv("BEAM_MAX_TASKS_PER_CHILD", 16))
+    @atexit.register
+    def close_pool(pool=worker_pool): pool.close()
+  return worker_pool
+
+def terminate_worker_pool():
+  global worker_pool
+  if worker_pool is not None: worker_pool.terminate()
+  worker_pool = None

tinygrad/helpers.py

@@ -364,7 +364,8 @@ class TracingKey:
 class ProfileEvent: pass
 
 @dataclass
-class ProfileRangeEvent(ProfileEvent): device:str; name:str|TracingKey; st:decimal.Decimal; en:decimal.Decimal|None=None # noqa: E702
+class ProfileRangeEvent(ProfileEvent):
+  device:str; name:str|TracingKey; st:decimal.Decimal; en:decimal.Decimal|None=None; profile_key:bytes|None=None # noqa: E702
 
 @dataclass(frozen=True)
 class ProfilePointEvent(ProfileEvent):
@@ -372,8 +373,8 @@ class ProfilePointEvent(ProfileEvent):
 
 cpu_events:list[ProfileEvent] = []
 @contextlib.contextmanager
-def cpu_profile(name:str|TracingKey, device="TINY", display=True) -> Generator[ProfileRangeEvent, None, None]:
-  res = ProfileRangeEvent(device, name, perf_counter_us())
+def cpu_profile(name:str|TracingKey, device="TINY", display=True, profile_key:bytes|None=None) -> Generator[ProfileRangeEvent, None, None]:
+  res = ProfileRangeEvent(device, name, perf_counter_us(), profile_key=profile_key)
   try: yield res
   finally:
     res.en = perf_counter_us()
@@ -464,14 +465,16 @@ def _ensure_downloads_dir() -> pathlib.Path:
   return pathlib.Path(cache_dir) / "downloads"
 
 def fetch(url:str, name:pathlib.Path|str|None=None, subdir:str|None=None, gunzip:bool=False, allow_caching=not getenv("DISABLE_HTTP_CACHE"),
-          headers:dict[str, str]={}, sha256:str|None=None) -> pathlib.Path:
+          headers:dict[str, str]={}, sha256:str|None=None, extract:bool=False) -> pathlib.Path:
   import urllib.request
   if url.startswith(("/", ".")): return pathlib.Path(url)
   if name is not None and (isinstance(name, pathlib.Path) or '/' in name): fp = pathlib.Path(name)
   else:
     hh = "_"+hashlib.md5(("\n".join(f"{k.strip()}:{v.strip()}" for k,v in sorted(headers.items()))).encode("utf-8")).hexdigest() if headers else ""
     fp = _ensure_downloads_dir() / (subdir or "") / ((name or hashlib.md5(url.encode('utf-8')).hexdigest()) + hh + (".gunzip" if gunzip else ""))
+  extract_dir = fp.parent / f"{fp.name}.extract"
   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:
       assert r.status in {200, 206}, r.status
@@ -488,6 +491,17 @@ def fetch(url:str, name:pathlib.Path|str|None=None, subdir:str|None=None, gunzip
         pathlib.Path(f.name).rename(fp)
       progress_bar.update(close=True)
       if length and (file_size:=os.stat(fp).st_size) < length: raise RuntimeError(f"fetch size incomplete, {file_size} < {length}")
+  if extract:
+    if not extract_dir.is_dir():
+      import tarfile
+      tmpdir = tempfile.mkdtemp(dir=fp.parent)
+      try:
+        with tarfile.open(fp) as t: t.extractall(tmpdir, filter="data")
+        try: os.rename(tmpdir, extract_dir)  # rename is atomic, so concurrent fetches can't see a partial extraction
+        except OSError:
+          if not extract_dir.is_dir(): raise
+      finally: shutil.rmtree(tmpdir, ignore_errors=True)
+    return extract_dir
   return fp
 
 def fetch_fw(path:str, name:str, sha256:str) -> bytes:

tinygrad/llm/cli.py

@@ -88,6 +88,8 @@ models = {
   "qwen3.5:9b": "https://huggingface.co/unsloth/Qwen3.5-9B-GGUF/resolve/main/Qwen3.5-9B-Q4_K_M.gguf",
   "qwen3.6:27b": "https://huggingface.co/unsloth/Qwen3.6-27B-GGUF/resolve/main/Qwen3.6-27B-Q4_K_M.gguf",
   "qwen3.6:35b-a3b": "https://huggingface.co/unsloth/Qwen3.6-35B-A3B-GGUF/resolve/main/Qwen3.6-35B-A3B-UD-Q4_K_M.gguf",
+  # pinned to the last revision with the plain IQ4_XS quant: the UD replacement uses Q3_K tensors the loader doesn't support
+  "qwen3.8:27b": "https://huggingface.co/unsloth/Qwen3.8-27B-GGUF/resolve/b62a80264f8b0c1bb849ee1c9c487415ebeca194/Qwen3.8-27B-IQ4_XS.gguf",
   "olmoe": "https://huggingface.co/allenai/OLMoE-1B-7B-0924-Instruct-GGUF/resolve/main/olmoe-1b-7b-0924-instruct-q4_k_m.gguf",
   "moonlight": "https://huggingface.co/gabriellarson/Moonlight-16B-A3B-Instruct-GGUF/resolve/main/Moonlight-16B-A3B-Instruct-Q4_K_M.gguf",
   "glm-4.7-flash": "https://huggingface.co/unsloth/GLM-4.7-Flash-GGUF/resolve/main/GLM-4.7-Flash-Q4_K_M.gguf",

tinygrad/llm/kernels/amd.py

@@ -0,0 +1,499 @@
+from __future__ import annotations
+import functools, math
+from typing import Callable, cast
+from tinygrad import Tensor, UOp, nn, Device, Context
+from tinygrad.device import Buffer
+from tinygrad.dtype import AddrSpace, dtypes
+from tinygrad.helpers import prod
+from tinygrad.uop.ops import AxisType, KernelInfo, Ops, resolve
+
+BLOCK_M, BLOCK_N, DECODE_HEAD_TILE, WARP_SIZE = 32, 32, 8, 32
+WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
+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
+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:
+  # a Variable is a 0-d ALU BUFFER in the tensor graph; inside kernels it takes the ALU PARAM form (same name keeps the value binding)
+  return x.substitute({v: UOp.variable(v.expr, v.vmin, v.vmax, dtype=v.dtype, multiple_of=v.arg.multiple_of, param=True)
+                       for v in x.toposort() if v.is_variable})
+
+def _unbind(v:int|UOp) -> int|UOp: return kernel_var(v.unbind_all()[0]) if isinstance(v, UOp) else v
+
+@functools.cache
+def amd_custom_kernels_supported(device:str|tuple[str, ...]|None) -> bool:
+  # the custom kernels are tuned for RDNA3 (gfx11): the WMMA register layouts don't match gfx12 (RDNA4)
+  # or CDNA (MFMA-only, wave64), and the dp4a builtins and 32-lane wave ops aren't portable either.
+  if isinstance(device, tuple): device = device[0]
+  if device is None or device.split(":")[0] != "AMD": return False
+  # @function contexts set ALLOW_DEVICE_USAGE=0 (scheduling must not open devices); the device is always open here
+  with Context(ALLOW_DEVICE_USAGE=1):
+    return (t:=getattr(Device[device], "target", None)) is not None and t[0] == 11
+
+def warp_reduce(val:UOp, maximum:bool=False, full_wave:bool=False) -> UOp:
+  for offset in ((16, 8, 4, 2, 1) if full_wave else (8, 4, 2, 1)):
+    if val.op is Ops.INDEX and val.addrspace == AddrSpace.REG: val = val.load()
+    other = UOp(Ops.CUSTOM, dtypes.float, (val,), arg=
+      f"__builtin_bit_cast(float, __builtin_amdgcn_ds_swizzle(__builtin_bit_cast(int, {{0}}), {0x1f | offset<<10}))")
+    val = val.maximum(other) if maximum else val + other
+  return val
+
+def _reg(shape:tuple[int, ...], slot:int, value:float, dep:UOp|None=None) -> UOp:
+  ret = UOp.placeholder(shape, dtypes.float, slot=slot, addrspace=AddrSpace.REG)
+  return ret.after((ret if dep is None else ret.after(dep)).store(ret.const_like(value)))
+
+# ******** quant linear: q8-activation kernels over packed ggml weights (Q4_K/Q5_K/Q6_K/IQ4_XS) ********
+
+class Linear(nn.Linear):
+  ggml_type:int|None = None
+  use_custom_quant = True
+  def __init__(self, in_features:int, out_features:int, bias=True):
+    super().__init__(in_features, out_features, bias)
+    self.in_features, self.out_features = in_features, out_features
+  def set_quantized(self, decoded:Tensor):
+    packed_sizes = {decoded.numel() // 256 * type_size:typ for typ,type_size in QUANT_SIZES.items()}
+    raw = next((u for u in decoded.uop.toposort() if u.op is Ops.SHRINK and u.dtype == dtypes.uint8 and prod(u.shape) in packed_sizes), None)
+    if raw is None: return
+    raw_offset = raw.contiguous_view_offset()
+    assert raw_offset is not None and raw_offset % 4 == 0 and raw.buf_uop.dtype == dtypes.uint8
+    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
+    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: 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)
+      # 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))
+      return out.shrink(tuple((0, s) for s in (*x.shape[:-1], self.out_features)))
+    return super().__call__(x)
+
+def _amd_dp4a(a:UOp, b:UOp, c:UOp) -> UOp:
+  return UOp(Ops.CUSTOMI, dtypes.int32, (a.int(), b.int(), c), arg="__builtin_amdgcn_sudot4(true, {}, true, {}, {}, false)")
+
+def _amd_byte_perm(a:UOp, b:UOp, selectors:UOp) -> UOp:
+  return UOp(Ops.CUSTOMI, dtypes.uint32, tuple(x.cast(dtypes.uint32) for x in (a, b, selectors)), arg="__builtin_amdgcn_perm({}, {}, {})")
+
+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)
+
+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()
+
+def _iq4_bytes(packed:UOp, shift:int) -> UOp:
+  selectors = (packed >> shift) & 0x0f0f0f0f
+  low = _amd_byte_perm(UOp.const(0xf6eaddcf, dtypes.uint32), UOp.const(0xbfad9881, dtypes.uint32), selectors)
+  high = _amd_byte_perm(UOp.const(0x71594535, dtypes.uint32), UOp.const(0x26190d01, dtypes.uint32), selectors & 0x07070707)
+  return _amd_byte_perm(high, low, 0x03020100 | ((selectors & 0x08080808) >> 1))
+
+def _q5_scales(raw:UOp, base:UOp, subgroup:UOp) -> tuple[UOp, UOp, UOp, UOp]:
+  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()
+
+def _iq4_scales(raw:UOp, base:UOp, subgroup:UOp) -> tuple[UOp, UOp]:
+  low = _load_byte(raw, base, 4 + subgroup//2)
+  scale = ((low >> (4*(subgroup%2)).cast(dtypes.uint32)) & 15) | ((((raw[base] >> 16) >> (2*subgroup).cast(dtypes.uint32)) & 3) << 4)
+  return _half(raw[base] & 0xffff), (scale.cast(dtypes.uint8).bitcast(dtypes.int8)-32).float()
+
+@functools.cache
+def iq4_half_lut(device:str) -> Tensor:
+  from tinygrad.runtime.autogen.ggml_common import kvalues_iq4nl
+  return Tensor([x for j in range(16) for i in range(16) for x in (kvalues_iq4nl[i], kvalues_iq4nl[j])],
+                dtype=dtypes.float16, device=device).bitcast(dtypes.uint32).contiguous()
+
+@functools.cache
+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), 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))
+  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]:
+  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)
+  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 = (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(
+    arg=KernelInfo(name=name, opts_to_apply=()))
+
+@functools.cache
+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, 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 = (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)
+      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)
+      for word_idx in range(8):
+        packed = _amd_load(raw[base + 2 + subgroup*4 + word_idx%4])
+        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
+    base = (output*in_features//GGML_BLOCK_SIZE+block)*Q6_BYTES
+    dots = [UOp.const(0, dtypes.int32)] * 2
+    for word_idx in range(8):
+      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 + 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])
+
+def _wmma_layout(out:UOp, out_features:int, token_tile:int, output_tiles:int):
+  output_waves = 2 if out_features % (32*output_tiles) == 0 else 1
+  token_block, output_block = UOp.range(out.shape[0]//token_tile, 0), UOp.range(out_features//(16*output_tiles*output_waves), 1)
+  lane, wave = UOp.range(WARP_SIZE, 2, axis_type=AxisType.LOCAL), UOp.range(output_waves, 3, axis_type=AxisType.LOCAL)
+  hw_lane = UOp(Ops.CUSTOM, dtypes.int32, (lane.int(),), arg="__builtin_amdgcn_mbcnt_lo(-1, 0)").cast(dtypes.weakint)
+  col, half = hw_lane % 16, hw_lane // 16
+  outputs = tuple((output_block*output_waves+wave)*(16*output_tiles) + tile*16 + col for tile in range(output_tiles))
+  inputs = tuple(token_block*token_tile + tile*16 + col for tile in range(token_tile//16))
+  tokens = tuple(tuple(token_block*token_tile + tile*16 + half*8 + i for i in range(8)) for tile in range(token_tile//16))
+  return output_waves, token_block, output_block, lane, wave, half, outputs, inputs, tokens
+
+def _wmma_stores(out, outputs, tokens, accs, update, half):
+  def values(acc:UOp) -> tuple[UOp, ...]:
+    vals = tuple(acc.after(update)[i].load() for i in range(8))
+    swapped = tuple(UOp(Ops.CUSTOM, dtypes.float32, (value,),
+      arg="__builtin_bit_cast(float, __builtin_amdgcn_ds_swizzle(__builtin_bit_cast(int, {0}), 50688))") for value in vals)
+    low = half.eq(0)
+    return tuple(low.where(vals[i], swapped[i+4]) if j == 0 else low.where(swapped[i], vals[i+4]) for i in range(4) for j in range(2))
+  return [out[token, output].store(value) for output,output_accs in zip(outputs, accs)
+          for tile_tokens,acc in zip(tokens, output_accs) for token,value in zip(tile_tokens, values(acc))]
+
+def _quant_linear_wmma(out, x, out_features, in_features, type_words, layout, dequant, name):
+  x = x.reshape(out.shape[0], in_features)
+  _, token_block, output_block, lane, wave, physical_half, outputs, input_tokens, tokens = layout
+  token_tile, output_tiles = len(tokens)*16, len(outputs)
+  output_words = in_features // GGML_BLOCK_SIZE * type_words
+  accs = tuple(tuple(UOp.placeholder((8,), dtypes.float32, slot=ot*(token_tile//16)+tile, addrspace=AddrSpace.REG)
+                     for tile in range(token_tile // 16)) for ot in range(output_tiles))
+  accs = tuple(tuple(acc.after(acc.store(acc.const_like(0))) for acc in output_accs) for output_accs in accs)
+  group = UOp.range(in_features // Q8_GROUP_SIZE, 4, AxisType.REDUCE)
+  block, subgroup = group // 8, group % 8
+  wmma_accs = [list(output_accs) for output_accs in accs]
+  for half in range(2):
+    afrags = tuple(UOp.stack(*(x[input_token, group*32 + half*16 + i].cast(dtypes.float16) for i in range(16)))
+                   for input_token in input_tokens)
+    for output_tile,output in enumerate(outputs):
+      bfrag = UOp.stack(*dequant(output*output_words + block*type_words, subgroup, half))
+      for tile,afrag in enumerate(afrags):
+        previous = accs[output_tile][tile].after(group) if half == 0 else wmma_accs[output_tile][tile]
+        wmma_accs[output_tile][tile] = UOp.wmma(afrag, bfrag, previous, *WMMA_ARG)
+  update = UOp.group(*(acc.store(value) for output_accs,output_values in zip(accs, wmma_accs)
+                       for acc,value in zip(output_accs, output_values))).end(group)
+  return UOp.group(*_wmma_stores(out, outputs, tokens, accs, update, physical_half)).end(token_block, output_block, lane, wave).sink(
+    arg=KernelInfo(name=name, opts_to_apply=()))
+
+@functools.cache
+def _q5_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, out_features:int, in_features:int, ggml_type:int) -> UOp:
+  token_tile, output_tiles = (64, 1) if out_features <= 1024 and out.shape[0] % 64 == 0 else \
+    (64, 2) if out.shape[0] % 64 == 0 else (32 if out.shape[0] % 32 == 0 else 16, 2)
+  def dequant(base:UOp, subgroup:UOp, half:int) -> tuple[UOp, ...]:
+    d, dmin, scale, minimum = _q5_scales(raw, base, subgroup)
+    qs_base = base + (4 if ggml_type == Q4_K else 12) + (subgroup // 2)*8 + half*4
+    words = tuple((raw[qs_base+i] >> ((subgroup&1)*4).cast(dtypes.uint32) & 0x0f0f0f0f) |
+      (((raw[base+4+half*4+i] >> subgroup.cast(dtypes.uint32) & 0x01010101) << 4) if ggml_type == Q5_K else 0) for i in range(4))
+    return tuple(((word >> (byte*8) & 255).float()*d*scale-dmin*minimum).cast(dtypes.float16) for word in words for byte in range(4))
+  return _quant_linear_wmma(out, x, out_features, in_features, Q4_WORDS if ggml_type == Q4_K else Q5_WORDS,
+                            _wmma_layout(out, out_features, token_tile, output_tiles), dequant,
+                            f"linear_q{4 if ggml_type == Q4_K else 5}_k_f16_wmma")
+
+@functools.cache
+def _iq4_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, lut:UOp, out_features:int, in_features:int) -> UOp:
+  token_tile = 32 if out_features <= 1024 and out.shape[0] % 32 == 0 else 64 if out.shape[0] % 64 == 0 and \
+    (out_features <= 6144 or out_features == 5120 and in_features > 8192) else 128 if out.shape[0] % 128 == 0 else \
+    32 if out.shape[0] % 32 == 0 else 16
+  output_tiles = 1 if out_features <= 1024 else 2 if out_features <= 6144 else 1 if out_features < 8192 else 2
+  layout = _wmma_layout(out, out_features, token_tile, output_tiles)
+  output_waves, _, _, lane, wave, _, _, _, _ = layout
+  local_lut = UOp.placeholder((256,), dtypes.uint32, slot=32, addrspace=AddrSpace.LOCAL)
+  tid, lut_items = wave*32+lane, 256//(32*output_waves)
+  lut = local_lut.after(UOp.group(*(local_lut[tid*lut_items+i].store(lut[tid*lut_items+i]) for i in range(lut_items))).barrier())
+  def dequant(base:UOp, subgroup:UOp, half:int) -> tuple[UOp, ...]:
+    d, scale = _iq4_scales(raw, base, subgroup)
+    scale = scale * d
+    if out_features <= 6144:
+      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)
+    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) -> 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
+  def run(fxn:Callable[..., UOp], out:UOp, *srcs:UOp) -> Tensor:
+    all_srcs = (out,)+srcs
+    params = tuple(UOp.placeholder_like(src, slot=i) for i,src in enumerate(all_srcs))
+    kernel = fxn(*params, out_features=out_features, in_features=in_features).call(*all_srcs)
+    result = Tensor(out.after(kernel))
+    if len(result.shape) == 3: result = result.sum(-1)
+    result = result.reshape(*x.shape[:-1], out_features)
+    return result if layer.bias is None else result + layer.bias
+  out = Tensor.empty(tokens, out_features, dtype=dtypes.float32, device=x.device).uop
+  if tokens % 16 == 0 and out_features % 16 == 0 and layer.ggml_type in (Q4_K, Q5_K, IQ4_XS):
+    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)
+  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)
+
+# ******** flash attention on the KV cache ********
+
+@functools.cache
+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, 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(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=128)
+  partial, stats = Tensor.custom_kernel(partial, stats, q, cache_kv, fxn=fxn)[:2]
+  live = (valid_kv_len+127)//128
+  live = min(live, chunks) if isinstance(live, int) else live.minimum(chunks)
+  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:
+  valid_kv_len, q_start = _unbind(valid_kv_len), _unbind(q_start) if q_start is not None else None
+  BH, M, D = q.shape
+  _, B, H_KV, physical_n, cache_dim = cache.shape
+  k, v = cache[0].reshape(B*H_KV, physical_n, cache_dim), cache[1].reshape(B*H_KV, physical_n, cache_dim)
+  assert k.shape == v.shape and BH % k.shape[0] == 0 and k.shape[2] == D
+  gqa_group = BH // k.shape[0]
+  if isinstance(M, int) and isinstance(valid_kv_len, int): assert M % BLOCK_M == 0 and valid_kv_len % BLOCK_N == 0
+  assert isinstance(D, int) and D % WMMA_K == 0 and D % LANES_PER_WAVE_N == 0
+  TM, TN, TD, SCALE = BLOCK_M//(WAVES_M*LANES_PER_WAVE_M), BLOCK_N//LANES_PER_WAVE_N, D//(WAVES_N*LANES_PER_WAVE_N), 1/math.sqrt(D)
+  # query row 0 sits at sequence position q_base (the queries may be padded beyond valid_kv_len - q_base rows)
+  q_base = valid_kv_len - M if q_start is None else q_start
+  block_bh, block_m = UOp.range(BH, 0, AxisType.GLOBAL), UOp.range(M // BLOCK_M, 1, AxisType.GLOBAL)
+  kv_head = block_bh // gqa_group
+  q, o = (x.reshape(BH, M//BLOCK_M, BLOCK_M, D)[block_bh, block_m] for x in (q, o))
+  k, v = k[kv_head], v[kv_head]
+  wave_m, wave_n, lane = UOp.range(WAVES_M, 2, AxisType.LOCAL), UOp.range(WAVES_N, 3, AxisType.LOCAL), UOp.range(WARP_SIZE, -1, AxisType.WARP)
+  tid, lane_m, lane_n = (wave_m * WAVES_N + wave_n) * WARP_SIZE + lane, lane // LANES_PER_WAVE_N, lane % LANES_PER_WAVE_N
+  Q_ELEMS_PER_THREAD, KV_ELEMS_PER_THREAD = BLOCK_M * D // THREADS_PER_BLOCK, BLOCK_N * D // THREADS_PER_BLOCK
+  QP_lds = UOp.placeholder((BLOCK_M, D + LDS_PAD), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL)
+  KV_lds = UOp.placeholder((BLOCK_N, D + LDS_PAD), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL)[:, :D]
+  acc, m_i, l_i = _reg((TM, TD), 2, 0), _reg((TM,), 3, -math.inf), _reg((TM,), 4, 0)
+  n_tile = UOp.range((q_base + (block_m + 1) * BLOCK_M + BLOCK_N - 1) // BLOCK_N, 100, AxisType.REDUCE)
+  Q_lds = QP_lds[:, :D]
+  Q_store = Q_lds.after(n_tile).reshape(THREADS_PER_BLOCK, Q_ELEMS_PER_THREAD)[tid].store(q.reshape(THREADS_PER_BLOCK, Q_ELEMS_PER_THREAD)[tid])
+  load_k = UOp.range(KV_ELEMS_PER_THREAD, 90)
+  kval = k.reshape(physical_n*D)[n_tile*BLOCK_N*D + tid*KV_ELEMS_PER_THREAD + load_k].float()
+  K_store = KV_lds.reshape(THREADS_PER_BLOCK, KV_ELEMS_PER_THREAD)[tid, load_k].store(kval).end(load_k)
+  qk_load_barrier = UOp.barrier(UOp.group(Q_store, K_store))
+  Q_lds, KV_lds_k = Q_lds.after(qk_load_barrier), KV_lds.after(qk_load_barrier)
+  S_reg = _reg((TM, TN), 6, 0, n_tile)
+  k_qk, tm1, tn1 = UOp.range(D//WMMA_K, 101, AxisType.REDUCE), UOp.range(TM//WMMA_ACC, 200), UOp.range(TN, 201)
+  S_frag = S_reg.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0, 2, 1)[tm1, tn1]
+  q_frag = Q_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, D // WMMA_K, WMMA_K)[wave_m, tm1, lane_n, k_qk]
+  k_frag = KV_lds_k.reshape(TN, WMMA_N, D // WMMA_K, WMMA_K)[tn1, lane_n, k_qk]
+  qk_done = S_frag.store(UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), *WMMA_ARG)).end(tm1, tn1).end(k_qk)
+  S_reg = S_reg.after(qk_done, S_reg.store(S_reg * SCALE))
+  rm, rn = UOp.range(TM, 250), UOp.range(TN, 251)
+  q_idx = q_base + block_m * BLOCK_M + wave_m * WMMA_M + rm * LANES_PER_WAVE_M + lane_m
+  k_idx = n_tile * BLOCK_N + rn * LANES_PER_WAVE_N + lane_n
+  S_reg = S_reg.after(S_reg[rm, rn].store((k_idx <= q_idx).where(S_reg[rm, rn], S_reg[rm, rn].const_like(-math.inf))).end(rm, rn))
+  m_ij, rm2 = _reg((TM,), 7, -math.inf, n_tile), UOp.range(TN, 261, AxisType.REDUCE)
+  m_ij = m_ij.after(m_ij.store(m_ij.after(rm2).maximum(S_reg[:, rm2])).end(rm2))
+  ri_w = UOp.range(TM, 270)
+  m_ij = m_ij.after(m_ij[ri_w].store(warp_reduce(m_ij[ri_w], maximum=True)).end(ri_w))
+  tile_max = m_ij.reshape(TM, 1).expand(TM, TN).maximum(-1e30)
+  S_reg = S_reg.after(S_reg.store(((S_reg - tile_max) * LOG2E).exp2()))
+  p_local, ri_ws = _reg((TM,), 8, 0, n_tile), UOp.range(TM, 295)
+  p_sum = p_local.after(p_local[ri_ws].store(sum((warp_reduce(S_reg[ri_ws, rn]) for rn in range(TN)), S_reg.const_like(0))).end(ri_ws))
+  P_lds = QP_lds.flatten()[:WAVES_N * BLOCK_M * BLOCK_N].reshape(WAVES_N, BLOCK_M, BLOCK_N)
+  P_write = P_lds.reshape(WAVES_N, WAVES_M, TM, LANES_PER_WAVE_M, 1, TN, LANES_PER_WAVE_N, 1).permute((1, 0, 3, 6, 2, 4, 5, 7)) \
+    .reshape(THREADS_PER_BLOCK, TM, TN)
+  P_store = P_write[tid].store(S_reg.cast(dtypes.half))
+  beta_i, ri4, rj4 = UOp.placeholder((TM,), dtypes.float, slot=9, addrspace=AddrSpace.REG), UOp.range(TM, 330), UOp.range(TD, 331)
+  m_new = m_i[ri4].maximum(m_ij[ri4])
+  alpha_val, beta_val = ((m_i[ri4] - m_new) * LOG2E).exp2(), ((m_ij[ri4] - m_new) * LOG2E).exp2()
+  correction = UOp.group(acc[ri4, rj4].store(alpha_val * acc[ri4, rj4]).end(rj4),
+                         l_i[ri4].store(alpha_val * l_i[ri4] + beta_val * p_sum[ri4]),
+                         m_i[ri4].store(m_new), beta_i[ri4].store(beta_val)).end(ri4)
+  acc, l_i, m_i, beta_i = acc.after(correction), l_i.after(correction), m_i.after(correction), beta_i.after(correction)
+  V_lds = UOp.placeholder((D, BLOCK_N + LDS_PAD), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL)[:, :BLOCK_N]
+  V_copy, load_v = V_lds.after(qk_done).permute(1, 0), UOp.range(KV_ELEMS_PER_THREAD, 390)
+  vval = v.reshape(physical_n*D)[n_tile*BLOCK_N*D + tid*KV_ELEMS_PER_THREAD + load_v].float()
+  V_store = V_copy.reshape(THREADS_PER_BLOCK, KV_ELEMS_PER_THREAD)[tid, load_v].store(vval).end(load_v)
+  pv_barrier = UOp.barrier(UOp.group(P_store, V_store))
+  P_lds, V_lds = P_lds.after(pv_barrier), V_lds.after(pv_barrier)
+  pv_acc = _reg((TM, TD), 10, 0, n_tile).after(pv_barrier)
+  k_pv, tm2, tn2 = UOp.range(BLOCK_N//WMMA_K, 400, AxisType.REDUCE), UOp.range(TM//WMMA_ACC, 401), UOp.range(TD, 402)
+  pv_frag = pv_acc.reshape(TM // WMMA_ACC, WMMA_ACC, TD).permute(0, 2, 1)[tm2, tn2]
+  p_frag = P_lds[wave_n].reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_N // WMMA_K, WMMA_K)[wave_m, tm2, lane_n, k_pv]
+  v_frag = V_lds.reshape(WAVES_N, TD, WMMA_N, BLOCK_N // WMMA_K, WMMA_K)[wave_n, tn2, lane_n, k_pv]
+  pv_done = pv_frag.store(UOp.wmma(p_frag, v_frag, pv_frag.after(k_pv), *WMMA_ARG)).end(tm2, tn2).end(k_pv)
+  pv_acc = pv_acc.after(pv_done)
+  ri5, rj5 = UOp.range(TM, 410), UOp.range(TD, 411)
+  n_tile_end = acc[ri5, rj5].store(acc[ri5, rj5] + beta_i[ri5] * pv_acc[ri5, rj5]).end(ri5, rj5).barrier().end(n_tile)
+  acc, l_i, m_i = acc.after(n_tile_end), l_i.after(n_tile_end), m_i.after(n_tile_end)
+  acc = acc.after(acc.store(acc * (1 / l_i).reshape(TM, 1).expand(TM, TD)))
+  o = o.reshape(WAVES_M, TM, LANES_PER_WAVE_M, 1, WAVES_N, TD, LANES_PER_WAVE_N, 1) \
+    .permute((0, 4, 2, 6, 1, 3, 5, 7)).reshape(THREADS_PER_BLOCK, TM, TD)
+  return o[tid].store(acc).end(wave_m, wave_n, lane).end(block_m, block_bh).sink(arg=KernelInfo(opts_to_apply=()))
+
+def flash_attention(q:Tensor, assigned_kv:Tensor, valid_end:int|UOp) -> Tensor:
+  # cached flash attention on the half KV cache (already written through assigned_kv); valid_end stays bound at the graph level
+  T_real, q_start = q.shape[2], None
+  if resolve(T_real == 1): return amd_flash_attention_decode(q.half(), assigned_kv, valid_end, cast(int, assigned_kv.shape[3]))
+  if isinstance(T_real, UOp):
+    # symbolic chunk: pad the queries to the static tile; garbage rows are sliced off
+    T_pad = q.max_shape[2]
+    assert T_pad % BLOCK_M == 0, "chunk_size must be a multiple of 32"
+    q, q_start = q.pad_to((*q.shape[:2], T_pad, q.shape[3])), valid_end - T_real
+  B, H, T, D = q.shape
+  out = Tensor.empty(B*H, T, D, dtype="float32", device=q.device)
+  fxn = functools.partial(_amd_flash_attention, valid_kv_len=valid_end, q_start=q_start)
+  out = Tensor.custom_kernel(out, q.half().reshape(B*H, T, D), assigned_kv, fxn=fxn)[0].reshape(B, H, T, D)
+  return out if q_start is None else out[:, :, :T_real]
+
+# ******** gated delta net: fused recurrent scan ********
+
+@functools.cache
+def _gated_delta_prefill_kernel(core:UOp, q:UOp, k:UOp, v:UOp, beta:UOp, alpha:UOp, state:UOp, kq:UOp, start_pos:UOp|None=None) -> UOp:
+  batch, heads, tokens, value_dim, row_tile = *core.shape, 4
+  key_dim, alpha_dim = q.shape[-1], alpha.shape[-1] if len(alpha.shape) == 4 else 1
+  assert all(isinstance(x, int) for x in (batch, heads, tokens, value_dim, key_dim)) and key_dim % 32 == 0 and value_dim % row_tile == 0
+  batch, heads, tokens, value_dim, key_dim = cast(tuple[int, int, int, int, int], (batch, heads, tokens, value_dim, key_dim))
+  core, v = (x.reshape(batch*heads, tokens, value_dim) for x in (core, v))
+  q, k = (x.reshape(batch*heads, tokens, key_dim) for x in (q, k))
+  beta, kq = (x.reshape(batch*heads, tokens) for x in (beta, kq))
+  alpha, state = alpha.reshape(batch*heads, tokens, alpha_dim), state.reshape(batch*heads, value_dim, key_dim)
+  bh_row, lane = UOp.range(batch*heads*value_dim//row_tile, 0), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  bh, row_base = bh_row // (value_dim//row_tile), (bh_row % (value_dim//row_tile))*row_tile
+  rows, cols = tuple(row_base+i for i in range(row_tile)), tuple(lane + i*32 for i in range(key_dim//32))
+  current = UOp.placeholder((row_tile*key_dim//32,), dtypes.float32, slot=0, addrspace=AddrSpace.REG)
+  initial = None if start_pos is None else start_pos.eq(0)
+  current = current.after(current.store(UOp.stack(*(state[bh, row, col].float() if initial is None else
+    initial.where(0, state[bh, row, col].float()) for row in rows for col in cols))))
+  token = UOp.range(tokens, 2, AxisType.REDUCE)
+  keys = tuple(k[bh, token, col].load() for col in cols)
+  queries = tuple(q[bh, token, col].load() for col in cols)
+  updates, stores = [], []
+  for row_idx,row in enumerate(rows):
+    previous = tuple(current.after(token)[row_idx*key_dim//32+i].load() for i in range(key_dim//32))
+    av, bv = alpha[bh, token, row if alpha_dim > 1 else 0].load(), beta[bh, token].load()
+    state_k = warp_reduce(sum((x*y for x,y in zip(previous, keys)), UOp.const(0, dtypes.float32)), full_wave=True)
+    state_q = warp_reduce(sum((x*y for x,y in zip(previous, queries)), UOp.const(0, dtypes.float32)), full_wave=True)
+    delta = (v[bh, token, row].load() - state_k*av) * bv
+    updates += [x*av + delta*y for x,y in zip(previous, keys)]
+    stores.append(core[bh, token, row.valid(lane.eq(0))].store(state_q*av + delta*kq[bh, token]))
+  step = UOp.group(*stores, current.store(UOp.stack(*updates))).end(token)
+  state_stores = (state[bh, row, col].store(current.after(step)[row_idx*key_dim//32+i].load().cast(state.dtype))
+                  for row_idx,row in enumerate(rows) for i,col in enumerate(cols))
+  return UOp.group(*state_stores).end(lane, bh_row).sink(arg=KernelInfo(name="gated_delta_prefill", opts_to_apply=()))
+
+def gated_delta_prefill(q:Tensor, k:Tensor, v:Tensor, beta:Tensor, alpha:Tensor, state:Tensor, start_pos:Tensor|None=None) -> Tensor:
+  batch, heads, tokens, key_dim = q.shape
+  value_dim = v.shape[-1]
+  assert q.shape == k.shape and v.shape[:3] == beta.shape == (batch, heads, tokens) and state.shape == (batch, heads, value_dim, key_dim)
+  assert alpha.shape[:3] == (batch, heads, tokens) and (len(alpha.shape) == 3 or alpha.shape[-1] in (1, value_dim))
+  assert key_dim % 32 == 0 and value_dim % 4 == 0
+  core, kq = Tensor.empty_like(v), (q*k).sum(-1).contiguous()
+  srcs = (core, q.contiguous(), k.contiguous(), v.contiguous(), beta.contiguous(), alpha.contiguous(), state, kq)
+  if start_pos is None: return Tensor.custom_kernel(*srcs, fxn=_gated_delta_prefill_kernel)[0]
+  contig = tuple(x.uop if x.uop.op is Ops.AFTER else x.uop.contiguous() for x in srcs)
+  params = tuple(UOp.placeholder_like(x, slot=i) for i,x in enumerate(contig))
+  assert start_pos.uop.is_bound_var
+  # the bound start_pos reaches the graph through the state AFTER chain, like the flash kernels' valid_end
+  call = _gated_delta_prefill_kernel(*params, kernel_var(start_pos.uop.src[0])).call(*contig)
+  return Tensor(contig[0].after(call))

tinygrad/llm/model.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 import enum, functools, itertools, pathlib
 from dataclasses import dataclass, replace
 from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function, dtypes
-from tinygrad.nn import Linear
+from tinygrad.llm.kernels.amd import Linear, gated_delta_prefill, flash_attention, amd_custom_kernels_supported
 from tinygrad.llm.gguf import gguf_load
 from tinygrad.uop.ops import resolve
 
@@ -181,7 +181,13 @@ class TransformerBlock(FFNBlock):
     k = apply_rope(k[..., :self.config.rope_dim], self.freqs_cis[start_pos:start_pos+T]).cat(k[..., self.config.rope_dim:], dim=-1)
 
     # NOTE: we don't want to change self.cache_kv, the function API doesn't support this well
-    assigned_kv = Tensor(self.cache_kv.uop.after(self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(Tensor.stack(k, v).uop)))
+    store = self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.store(Tensor.stack(k, v).cast(dtypes.half).uop)
+    assigned_kv = Tensor(self.cache_kv.uop.after(store))
+    # on RDNA3, hybrid models use custom flash attention kernels on the KV cache
+    if amd_custom_kernels_supported(x.device) and self.config.ssm is not None:
+      attn = flash_attention(q, assigned_kv, start_pos+T)
+      attn = attn.transpose(1, 2).reshape(B, T, -1)                                    # back to (B,T,D)
+      return self.attn_output(attn if not self.config.attn_output_gate else (attn * gate.sigmoid()))
     k = assigned_kv[0, :, :, 0:start_pos+T, :]
     v = assigned_kv[1, :, :, 0:start_pos+T, :]
 
@@ -199,8 +205,9 @@ class TransformerBlock(FFNBlock):
 
   def _init_state(self, x:Tensor):
     if not hasattr(self, "cache_kv"):
-      self.cache_kv = Tensor.empty(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim,
-                                   dtype=dtypes.default_float, device=x.device)
+      # zeroed so the flash kernels can safely read whole tiles past the valid region (masked lanes multiply by 0)
+      self.cache_kv = Tensor.zeros(2, x.shape[0], self.config.n_kv_heads, self.config.max_context, self.config.head_dim,
+                                   dtype=dtypes.half, device=x.device)
       self.freqs_cis = precompute_freqs_cis(self.config.rope_dim, self.config.max_context, self.config.rope_theta, device=x.device)
 
 class MLATransformerBlock(FFNBlock):
@@ -311,21 +318,28 @@ class GatedDeltaNetBlock(FFNBlock):
     v = v.reshape(B, T_pad, self.num_v_heads, self.head_v_dim)
     # layout the per-step operands to broadcast against the (B, H, V, K) state
     q, k, v, beta = (z.transpose(1, 2).float() for z in (q, k, v, beta))
-    q, k, v, beta = q.unsqueeze(-2) * self.head_k_dim**-0.5, k.unsqueeze(-2), v.unsqueeze(-1), beta.unsqueeze(-1).unsqueeze(-1)
-    alpha = log_alpha.transpose(1, 2).exp().unsqueeze(-1)  # per-channel decay for kda, per-head otherwise (B, H, T, V|1, 1)
+    q = q * self.head_k_dim**-0.5
+    alpha = log_alpha.transpose(1, 2).exp()  # per-channel decay for kda, per-head otherwise (B, H, T, V|1)
 
     # recurrent: scan over the (padded) tokens, updating the recurrent state. collect the per-step outputs
-    state = Tensor(self.recurrent_state.uop.after(conv_state_store)).float()  # carry the conv write into this graph
-    state = initial.where(0, state)
-    outs = []
-    for t in range(T_pad):
-      s1 = state * alpha[:, :, t]  # decay the state
-      delta = (v[:, :, t] - (s1*k[:, :, t]).sum(-1, keepdim=True)) * beta[:, :, t]  # the delta rule update
-      state = s1 + delta * k[:, :, t]
-      outs.append((state * q[:, :, t]).sum(-1))
+    state = Tensor(self.recurrent_state.uop.after(conv_state_store))  # carry the conv write into this graph
+    if self.head_k_dim % 32 == 0 and self.head_v_dim % 4 == 0 and amd_custom_kernels_supported(x.device):
+      # one fused kernel for the whole scan; it resets and updates the recurrent state in place (RDNA3)
+      core = gated_delta_prefill(q, k, v, beta, alpha, state, Tensor(start_pos)).transpose(1, 2)
+    else:
+      q, k, v, beta = q.unsqueeze(-2), k.unsqueeze(-2), v.unsqueeze(-1), beta.unsqueeze(-1).unsqueeze(-1)
+      alpha = alpha.unsqueeze(-1)
+      state = initial.where(0, state.float())
+      outs = []
+      for t in range(T_pad):
+        s1 = state * alpha[:, :, t]  # decay the state
+        delta = (v[:, :, t] - (s1*k[:, :, t]).sum(-1, keepdim=True)) * beta[:, :, t]  # the delta rule update
+        state = s1 + delta * k[:, :, t]
+        outs.append((state * q[:, :, t]).sum(-1))
 
-    # store the updated recurrent state in place, then read the stacked outputs after the write
-    core = Tensor(outs[0].stack(*outs[1:], dim=1).contiguous().uop.after(self.recurrent_state.uop.store(state.cast(self.recurrent_state.dtype).uop)))
+      # store the updated recurrent state in place, then read the stacked outputs after the write
+      state_store = self.recurrent_state.uop.store(state.cast(self.recurrent_state.dtype).uop)
+      core = Tensor(outs[0].stack(*outs[1:], dim=1).contiguous().uop.after(state_store))
 
     # output; undo the padding before the output projection
     z = (self.ssm_norm(core) * (out_gate.sigmoid() if is_kda else out_gate.silu())).cast(x.dtype).contiguous()
@@ -462,7 +476,7 @@ class Transformer:
     return min(block._reusable_prefix_len(prefix_len, len(self._cached_tokens)) for block in self.blk)
 
   def generate(self, tokens:list[int], chunk_size:int=32, temperature:float=0.0):
-    if self.has_recurrent_block: chunk_size = 1
+    if self.has_recurrent_block and not amd_custom_kernels_supported(self.token_embd.weight.device): chunk_size = 1
     v_start_pos = UOp.variable("start_pos", 0, self.max_context-1)
     v_toks = UOp.variable("toks", 1, chunk_size)
     # TODO: use UOp.variable for temperature once float variables are supported

tinygrad/renderer/cstyle.py

@@ -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
@@ -38,7 +39,8 @@ base_rewrite = PatternMatcher([
   (UPat(Ops.CAST, name="x"), lambda ctx,x: f"__builtin_convertvector({ctx[x.src[0]]}, {ctx.render_type(x)})" \
     if x.max_numel() > 1 and x.addrspace is AddrSpace.REG else None),
   (UPat(Ops.CAST, name="x"), lambda ctx,x: f"({ctx.render_cast(x, ctx[x.src[0]])})"),
-  (UPat(Ops.BITCAST, name="x"), lambda ctx,x: ctx[x.src[0]] if x.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL) else None),
+  (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"(({ctx._render_dtype(x.dtype, addrspace=x.addrspace)})({ctx[x.src[0]]}))"
+   if x.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL) else None),
   (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"__builtin_bit_cast({ctx.render_type(x)}, ({ctx.render_type(x.src[0])})({ctx[x.src[0]]}))"),
 
   # GPU stuff
@@ -74,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"),
@@ -237,7 +241,7 @@ class CStyleLanguage(Renderer):
       if (u.op is not Ops.CAST or u.max_numel() == 1) and ((u.op is Ops.CAST and u.src[0].op is Ops.CONST) or \
         u.op in {Ops.INDEX, Ops.SHRINK, Ops.CUSTOMI} or \
         (u.op is Ops.LOAD and u.src[0].addrspace == AddrSpace.REG and child_count[u] == 1) or \
-        (u.op is Ops.CAST and u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL)) or \
+        (u.op in {Ops.CAST, Ops.BITCAST} and u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL)) or \
         (u.op in {Ops.STACK, *(GroupOp.ALU-{Ops.WHERE}), Ops.CAST, Ops.BITCAST} and child_count[u] == 1 and not getenv("EXPAND_SSA"))):
         r[u] = l
       else:
@@ -318,7 +322,8 @@ class OpenCLRenderer(CStyleLanguage):
   extra_matcher = create_non_native_float_pats((dtypes.bfloat16,)) + pm_manual_bf16_cast
 
   string_rewrite = PatternMatcher([
-    (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"as_{ctx.render_dtype(x.dtype)}(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"),
+    (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"as_{ctx.render_dtype(x.dtype)}(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"
+     if x.addrspace not in (AddrSpace.GLOBAL, AddrSpace.LOCAL) else None),
     # bfloat16 constants need to be rendered as their bit pattern since bf16 is stored as ushort
     (UPat.cvar("c").cast(dtypes.bfloat16), lambda ctx,c: f"{(struct.unpack('I', struct.pack('f', float_to_bf16(c.val)))[0] >> 16)}u"),
     # load/store image (OpenCL)
@@ -369,7 +374,8 @@ class MetalRenderer(CStyleLanguage):
   ]) + pm_manual_bf16_cast
 
   string_rewrite = PatternMatcher([
-    (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"as_type<{ctx.render_dtype(x.dtype)}>(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"),
+    (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"as_type<{ctx.render_dtype(x.dtype)}>(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"
+     if x.addrspace not in (AddrSpace.GLOBAL, AddrSpace.LOCAL) else None),
   ]) + base_rewrite
 
   def render_kernel(self, function_name, kernel, bufs, uops, prefix=None):
@@ -425,7 +431,8 @@ class CUDARenderer(CStyleLanguage):
     (UPat(Ops.CAST, dtypes.fp8s, UPat.var("x", dtypes.fp8s), name='y'), lambda x,y: x.cast(dtypes.float).cast(y.dtype) if x.dtype!=y.dtype else None),
   ])
   string_rewrite = PatternMatcher([
-    (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"tg_bitcast<{ctx.render_dtype(x.dtype)}>(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"),
+    (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"tg_bitcast<{ctx.render_dtype(x.dtype)}>(({ctx.render_dtype(x.src[0].dtype)})({ctx[x.src[0]]}))"
+     if x.addrspace not in (AddrSpace.GLOBAL, AddrSpace.LOCAL) else None),
   ]) + base_rewrite
 
   def render_vector_prefix(self, dt:DType, count:int) -> str:
@@ -520,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:

tinygrad/renderer/isa/x86.py

@@ -150,9 +150,9 @@ extra_matcher = PatternMatcher([
   # 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

tinygrad/renderer/nir.py

@@ -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)))),

tinygrad/renderer/wgsl.py

@@ -1,53 +1,50 @@
 from tinygrad.dtype import DType, dtypes, truncate, AddrSpace
 from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
 from tinygrad.renderer.cstyle import CStyleLanguage, base_rewrite
-from tinygrad.helpers import strip_parens
+from tinygrad.helpers import strip_parens, ceildiv
 
-def _mask(dt:DType): return 0xFF if dt.itemsize == 1 else 0xFFFF
+# a field of `width` bits sitting in the low bits of val: shift it up to the sign bit, then let the arithmetic shift fill
+def sign_extend(val:UOp, width:int): return (val << (32-width)).bitcast(dtypes.int) >> (32-width)
 
-def sign_extend(val:UOp, sext_am:int):
-  return (UOp.where((val >> (sext_am - 1)) > 0, UOp.const(0xffffffff << sext_am, dtypes.uint32), UOp.const(0, dtypes.uint32)) \
-        | val.bitcast(dtypes.uint32)).bitcast(dtypes.int)
+# a packed field of dt: the word it lives in, its offset in that word, and its mask. width is 8*itemsize, bool is one bit in a byte
+def packed_field(bidx:UOp, dt:DType) -> tuple[UOp, UOp, int]:
+  elems, width = 4//dt.itemsize, 8*dt.itemsize
+  return bidx.src[0].index(bidx.src[1] // elems), (bidx.src[1].cast(dtypes.uint32) % elems) * width, (1 << width)-1
 
 # store for char: buf[idx/4] <- (var << (idx%4)*8))
-def packed_store(bidx:UOp, var:UOp, gate:UOp|None=None):
-  elems, mask = 4//var.dtype.itemsize, _mask(var.dtype)
-  shift_am, div_idx = (bidx.src[1].cast(dtypes.uint32) % elems) * (8*var.dtype.itemsize), bidx.src[1] // elems
+def packed_store(s:UOp):
+  bidx, var, *gate = s.src
+  idx, shift_am, mask = packed_field(bidx, var.dtype)
   # 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)
-  idx = UOp(Ops.INDEX, src=(bidx.src[0], div_idx))
-  buf = UOp.load(idx, *((UOp.const(0, dtypes.uint32), gate) if gate is not None else ()), dtype=dtypes.uint32)
-  return UOp.store(idx, (buf & wmask) | new_v, *((gate,) if gate is not None else ()))
+  buf = idx.load(*((UOp.const(0, dtypes.uint32), *gate) if gate else ()), dtype=dtypes.uint32)
+  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:UOp, dtype:DType, var:UOp|None=None, gate:UOp|None=None):
-  elems, mask = 4//dtype.itemsize, _mask(dtype)
-  shift_am, div_idx = (bidx.src[1].cast(dtypes.uint32) % elems) * (8*dtype.itemsize), bidx.src[1] // elems
-  idx = UOp(Ops.INDEX, src=(bidx.src[0], div_idx))
-  load = UOp.load(idx, *((var, gate) if var is not None and gate is not None else root.src[1:]), dtype=dtypes.uint32, arg=root.arg)
-  val = (load.cast(dtypes.uint32) >> shift_am) & mask
+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)
+  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):
-  if x.op is Ops.LOAD: dt, addrspace = x.dtype, x.src[0].addrspace
-  elif x.op is Ops.STORE: dt, addrspace = x.src[1].dtype, x.src[0].addrspace
-  else: dt, addrspace = x.dtype, x.addrspace
-  return dt.itemsize < 4 and dt != dtypes.half and addrspace != AddrSpace.REG
-def _packed_size(u:UOp): return u.max_numel() // (4//u.dtype.itemsize) if is_packed(u) else u.max_numel()
+  dt = x.src[1].dtype if x.op is Ops.STORE else x.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):
   bs, (exp, mant) = a.dtype.bitsize, dtypes.finfo(a.dtype)
   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")
+
 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.load(UPat.var("b"), UPat.var("c"), UPat.var("gate"), name="l"),
-   lambda l,b,c,gate: packed_load(l,b,l.dtype,c.cast(dtypes.uint32),gate) if is_packed(l) else None),
-  (UPat.load(UPat.var("b"), name='l'), lambda l,b: packed_load(l,b,l.dtype) if is_packed(l) else None),
-  (UPat.store(UPat.var("b"), UPat.var("var"), UPat.var("gate"), name="s"),
-   lambda b,var,gate,s: packed_store(b,var,gate) if is_packed(s) else None),
-  (UPat.store(UPat.var("b"), UPat.var("var"), name="s"), lambda b,var,s: packed_store(b,var) if is_packed(s) else None),
+  (UPat(Ops.LOAD, 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),
   # fix nan check: 'a != a -> is_nan()'. the decomp rewrites (a != a).logical_not() to CMPEQ, so match both forms
@@ -72,7 +69,8 @@ class WGSLRenderer(CStyleLanguage):
     (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),),)),
@@ -87,10 +85,10 @@ class WGSLRenderer(CStyleLanguage):
     (UPat.load(UPat.var("b"), UPat.var("v"), UPat.var("gate")),
       lambda ctx,b,v,gate: f"select({ctx[v]}, {ctx.render_load(ctx[b], b.src[0])}, {ctx[gate]})"),
     (UPat.load(UPat.var("b")), lambda ctx, b: ctx.render_load(ctx[b], b)),
-    (UPat.store(UPat.var("b"), UPat.var("v")), lambda ctx,b,v:\
-     # (load & mask) | var -> mask = v.src[0].src[1], var = v.src[1]
-     f"atomicAnd(&{ctx[b]},{ctx[v.src[0].src[1]]});\n  atomicAdd(&{ctx[b]},{ctx[v.src[1]]});" if is_packed(b) \
-      else f"{ctx[b]} = {ctx[v]};"),
+    # packed_store writes (load & wmask) | new_v: atomicAnd clears the field, atomicAdd sets it. new_v is gone when it is 0
+    (UPat.store(UPat.var("b"), UPat.any(packed_rmw, packed_rmw | UPat.var("nv"))), lambda ctx,b,wmask,nv=None:
+     f"atomicAnd(&{ctx[b]},{ctx[wmask]});"+(f"\n  atomicAdd(&{ctx[b]},{ctx[nv]});" if nv is not None else "") if is_packed(b) else None),
+    (UPat.store(UPat.var("b"), UPat.var("v")), lambda ctx,b,v: f"{ctx[b]} = {ctx[v]};"),
     (UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("idx"))),
      lambda ctx,b,idx: f"{ctx[b]}[{strip_parens(ctx[idx]) if idx.arg is Ops.ADD else ctx[idx]}]"),
   ]) + base_rewrite

tinygrad/runtime/graph/hcq.py

@@ -139,7 +139,8 @@ class HCQGraph(MultiGraphRunner):
         prof_ji_desc = runtime.name if runtime is not None else TracingKey(f"{bufs[1].device} -> {bufs[0].device}", ret=bufs[0].nbytes)
 
         prof_name = enqueue_dev.device if runtime is not None else f"{enqueue_dev.device}:SDMA:{queue_idx}"
-        self.prof_graph_entries.append(ProfileGraphEntry(prof_name, prof_ji_desc, sig_st, j * 2 + 1))
+        self.prof_graph_entries.append(ProfileGraphEntry(prof_name, prof_ji_desc, sig_st, j * 2 + 1,
+                                                         runtime.profile_key if runtime is not None else None))
         self.prof_graph_deps.append([d - 1 for _, d in rdeps])
 
       self.last_j[enqueue_queue] = j

tinygrad/runtime/graph/metal.py

@@ -102,7 +102,7 @@ class MetalGraph(GraphRunner):
   def collect_timestamps(self):
     # create a graph event and evenly space each program
     st, en = decimal.Decimal(self.command_buffer.GPUStartTime()) * 1000000, decimal.Decimal(self.command_buffer.GPUEndTime()) * 1000000
-    ents = [ProfileGraphEntry(self.device, rt.name, i, i+1) for i, rt in enumerate(self.runtimes) if rt is not None]
+    ents = [ProfileGraphEntry(self.device, rt.name, i, i+1, rt.profile_key) for i, rt in enumerate(self.runtimes) if rt is not None]
     self.dev.profile_events += [ProfileGraphEvent(ents, [], [st + (en-st)/len(ents)*i for i in range(len(ents)+1)])]
 
   def __del__(self):

tinygrad/runtime/ops_amd.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 from typing import cast
-import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools, collections, atexit
+import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools, collections, atexit, time
 assert sys.platform != 'win32'
 from dataclasses import dataclass
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, CLikeArgsState, HCQSignal, HCQProgram, FileIOInterface
@@ -649,6 +649,59 @@ class AMDAllocator(HCQAllocator['AMDDevice']):
 
   def _do_map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
 
+  def _copyin(self, dest:HCQBuffer, src:memoryview):
+    if not self.dev.is_usb(): return super()._copyin(dest, src)
+    from tinygrad.runtime.support.usb import alloc_cbuffer
+    # Pipelined copyin over the 0xF2 engine. ~256KB chunks stream into two alternating 256KB SRAM bounce windows; the
+    # engine can't signal data landing, so each chunk's wire image ends in a 4B sentinel tagged with its sequence number.
+    # A prebuilt SDMA ring polls each chunk's sentinel before copying it to VRAM, then bumps a drain fence; the host
+    # waits on that fence before re-arming a window. No timing is assumed in either direction.
+    dev, usb, ts, sdma = self.dev, self.dev.iface.pci_dev.usb, self.dev.timeline_signal, self.dev.sdma
+    CHUNK, src_mv = 0x40000 - 4, src.cast('B')  # payload per chunk: the 256KB window minus the 4B trailing sentinel
+    nchunks = ceildiv(src.nbytes, CHUNK)
+    FENCE = 0xA800  # drain fence: the GPU writes it via sys_buf (PCIe 0x820800), the host reads it here (xdata)
+    if not hasattr(self, '_usb_seq'):  # one-time: clear the fence and zero both windows so garbage can't match a sentinel
+      self._usb_seq, self._usb_stage = 0, [alloc_cbuffer(0x40000) for _ in range(2)]  # (backing array, memoryview) pairs
+      self._usb_wins = (self.b[0].offset(0, 0x40000), self.b[0].offset(0x40000, 0x40000))  # two windows, engine slots 0/16
+      usb.write(FENCE, bytes(8))
+      for bi in range(2): usb.scsi_write(bytes(0x40000), slot_start=bi * 16)
+
+    def wait_drain(count):  # spin until the drain fence reaches count, i.e. chunks 0..count-1 are fully in VRAM
+      t0 = time.perf_counter()
+      while int.from_bytes(usb.read(FENCE, 8), 'little') < count:
+        if time.perf_counter() - t0 > 10: raise RuntimeError(f"GPU failed to drain USB copyin chunk {count - 1} (10s, hung GPU?)")
+
+    # build the whole ring upfront: per chunk, poll the sentinel, copy SRAM->VRAM, bump the fence; then one doorbell
+    POLL_EQ = sdma.SDMA_OP_POLL_REGMEM | sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(3) | sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1)
+    POLL_DW5 = sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff)
+    q = dev.hw_copy_queue_t().wait(ts, dev.timeline_value - 1)
+    for c in range(nchunks):
+      seq, size = self._usb_seq + c, min(CHUNK, src.nbytes - c * CHUNK)
+      q.q(POLL_EQ, *data64_le(self._usb_wins[seq & 1].va_addr + round_up(size + 4, 512) - 4), 0x51000000 | (seq & 0xFFFFFF), 0xFFFFFFFF, POLL_DW5)
+      q.copy(dest.offset(c * CHUNK), self._usb_wins[seq & 1], size)
+      q.write(dev.iface.sys_buf.offset(0x800, 8), seq + 1, b64=True)
+    q.signal(ts, dev.next_timeline()).submit(dev)
+
+    # stream the chunks: stage the wire image [payload][sentinel], arm the window, send. A window is reusable once
+    # its previous occupant (seq-2) is both fully sent (tag reaped) and fully drained to VRAM (the fence).
+    inflight = [None, None]
+    for c in range(nchunks):
+      seq, size = self._usb_seq + c, min(CHUNK, src.nbytes - c * CHUNK)
+      if inflight[seq & 1] is not None: usb.usb.bulk_wait(inflight[seq & 1])
+      buf = self._usb_stage[seq & 1][1]
+      buf[:size] = src_mv[c * CHUNK : c * CHUNK + size]
+      wire = round_up(size + 4, 512)  # payload plus the sentinel, padded to 512B sectors (full window for max chunks)
+      struct.pack_into('<I', buf, wire - 4, 0x51000000 | (seq & 0xFFFFFF))  # the sentinel is the last dword of the wire
+      arm_tag = usb.usb.control_write_async(0xF2, wire // 512, (seq & 1) * 16 | (ceildiv(wire, 0x4000) << 8))  # wValue=sectors, wIndex=slot|count
+      rd_tag, rd_mv = usb.usb.control_read_async(0xE4, 8, value=FENCE)  # arm and fence read fly in one round-trip window
+      usb.usb.bulk_wait(arm_tag)
+      usb.usb.bulk_wait(rd_tag)
+      if int.from_bytes(rd_mv, 'little') < seq - 1: wait_drain(seq - 1)  # rare: the drain lagged; spin on fresh reads
+      inflight[seq & 1] = usb.usb.bulk_write_async(buf[:wire])
+    for tag in inflight: usb.usb.bulk_wait(tag)
+    self._usb_seq += nchunks
+    wait_drain(self._usb_seq)  # copyin is synchronous: everything must be in VRAM before returning
+
   def _copyout(self, dest:memoryview, src:HCQBuffer):
     if not self.dev.is_usb(): return super()._copyout(dest, src)
     self.dev.synchronize()
@@ -924,14 +977,13 @@ class USBIface(PCIIface):
     region = self.dev_impl.mm.map_range(vaddr:=self.dev_impl.mm.alloc_vaddr(size=size), size, [(sys_addr, size)], aspace=AddrSpace.SYS, uncached=True)
     return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(region, has_cpu_mapping=False), view=self.pci_dev.dma_view(ctrl_addr, size), owner=self.dev)
 
-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, **kwargs) -> HCQBuffer:
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, zero=False, **kwargs) -> HCQBuffer:
     # usb allocates uncached and cpu_access in vram. vram writes are faster than sram writes
-    if host and self.sys_next_off + size < self.sys_buf.size:
-      self.sys_next_off += size
-      return self.sys_buf.offset(self.sys_next_off - size, size)
+    # NOTE: host allocs deliberately do NOT use sys_buf (the 0x820000 NVMe SQ region): the GPU's signal writes there
+    # collide with the 0xF2 engine mid-stream. Signals in VRAM are read back via 0xF0 streaming reads instead.
 
     # force devmem
-    return super().alloc(size, host=False, uncached=uncached, cpu_access=cpu_access, contiguous=contiguous, force_devmem=True, **kwargs)
+    return super().alloc(size, host=False, uncached=uncached, cpu_access=cpu_access, contiguous=contiguous, force_devmem=True, zero=zero, **kwargs)
 
   def sleep(self, timeout): pass
 
@@ -1048,7 +1100,8 @@ class AMDDevice(HCQCompiled):
     if getenv("AMD_DISABLE_SDMA"): return None
     if idx in self.sdma_queues: return self.sdma_queues[idx]
     with contextlib.suppress(OSError):
-      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20), idx=idx)
+      # USB: a copyin submits its whole ring at once (3 packets per 240KB chunk), so it needs more than the 0x200 default
+      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, (1 << 20) if self.is_usb() else (16 << 20), idx=idx)
     return self.sdma_queues.get(idx, None)
 
   def _ensure_has_local_memory(self, private_segment_size):

tinygrad/runtime/ops_cpu.py

@@ -5,7 +5,7 @@ from typing import cast, Callable
 from tinygrad.helpers import to_mv, from_mv, OSX, WIN, Context, mv_address, suppress_finalizing, unwrap, data64_le, to_tuple
 from tinygrad.device import Buffer, BufferSpec, TinyELF, Program, Device
 from tinygrad.runtime.support.hcq import HCQBuffer, MMIOInterface
-from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, make_cmdbuf, make_signal
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, make_cmdbuf, make_buf
 from tinygrad.runtime.support.c import DLL
 from tinygrad.renderer.cstyle import ClangRenderer
 from tinygrad.renderer.llvmir import CPULLVMRenderer
@@ -80,7 +80,7 @@ pm_cpu_opsel = PatternMatcher([
   (UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))),
    lambda ctx, dst, val: cpu_cmd(ctx, signal_prog, dst.getaddr(ctx), val.cast(dtypes.uint64))),
   (UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),)),
-   lambda ctx, dst: cpu_cmd(ctx, timestamp_prog, dst.getaddr(ctx), *(() if WIN else (make_signal(ctx, tag="func:clock_gettime").getaddr(ctx),)))),
+   lambda ctx, dst: cpu_cmd(ctx, timestamp_prog, dst.getaddr(ctx), *(() if WIN else (make_buf(ctx, tag="func:clock_gettime").getaddr(ctx),)))),
 ])
 
 def encode_queue(q:UOp) -> UOp:
@@ -91,7 +91,7 @@ def encode_queue(q:UOp) -> UOp:
   assert cnt < RING_SLOTS, f"submit of {cnt} entries doesn't fit the ring"
   cmdbuf = make_cmdbuf(lin, devs, buf=UOp.placeholder((cnt*CMD_SIZE,), dtypes.uint64, next(UOp.unique_num), device=devs).rtag("cmdbuf"))
   ring = UOp.placeholder((ring_words:=RING_SLOTS*CMD_SIZE,), dtypes.uint64, 0, device=devs, volatile=True).rtag(f"{queue}_ring")
-  put, done, sem, sysbuf = (make_signal(devs, tag=f"{queue}_{name}") for name in ("put", "done", "sem", "sys"))
+  put, done, sem, sysbuf = (make_buf(devs, tag=f"{queue}_{name}") for name in ("put", "done", "sem", "sys"))
 
   # submits are serialized on the submitter, so they can bump put without atomics
   ran = done.after(l:=UOp.loop(next(UOp.unique_num))).index(0).load()
@@ -104,7 +104,7 @@ def encode_queue(q:UOp) -> UOp:
   if WIN: return sysbuf.after(bumped).index(0).store(put.after(bumped).index(0).load())
 
   e = UOp.range(cnt, next(UOp.unique_num), dtype=dtypes.int, src=(bumped,))
-  return make_signal(devs, tag="func:sem_post").after(e).index(0).load().call(sem.after(e).index(0), ret_dtype=dtypes.void).end(e)
+  return make_buf(devs, tag="func:sem_post").after(e).index(0).load().call(sem.after(e).index(0), ret_dtype=dtypes.void).end(e)
 
 # *****************
 

tinygrad/runtime/ops_metal.py

@@ -34,6 +34,7 @@ class MetalDevice(Compiled):
     self.mtl_queue = self.sysdevice.newCommandQueueWithMaxCommandBufferCount(1024)
     if self.mtl_queue is None: raise RuntimeError("Cannot allocate a new command queue")
     self.mtl_buffers_in_flight: list[metal.MTLCommandBuffer] = []
+    self.mtl_profile_keys: dict[int, bytes] = {}
     self.timeline_signal = self.sysdevice.newSharedEvent()
     self.timeline_value = 0
 
@@ -55,7 +56,7 @@ class MetalDevice(Compiled):
       st, en = decimal.Decimal(cbuf.GPUStartTime()) * 1000000, decimal.Decimal(cbuf.GPUEndTime()) * 1000000
       # NOTE: command buffers from MetalGraph are not profiled here
       if PROFILE and (lb:=cmdbuf_label(cbuf)) is not None and not lb.startswith("batched"):
-        Compiled.profile_events += [ProfileRangeEvent(self.device, lb, st, en)]
+        Compiled.profile_events += [ProfileRangeEvent(self.device, lb, st, en, self.mtl_profile_keys.pop(id(cbuf), None))]
     self.mtl_buffers_in_flight.clear()
 
 class MetalCompiler(Compiler):
@@ -113,7 +114,7 @@ class MetalCompiler(Compiler):
 
 class MetalProgram(Program[MetalDevice]):
   def __init__(self, dev:MetalDevice, obj:TinyELF):
-    self.dev, self.name, self.lib, self.signature = dev, obj.name, obj.lib, obj.signature
+    self.dev, self.name, self.lib, self.signature, self.profile_key = dev, obj.name, obj.lib, obj.signature, obj.profile_key
     data = objc.dispatch_data_create(obj.lib, len(obj.lib), None, None)
     self.library = self.dev.sysdevice.newLibraryWithData_error(data, ctypes.byref(error_lib:=metal.NSError().retained())).retained()
     error_check(error_lib)
@@ -145,6 +146,7 @@ class MetalProgram(Program[MetalDevice]):
     command_buffer.setLabel(to_ns_str(self.name)) # TODO: is this always needed?
     command_buffer.commit()
     self.dev.mtl_buffers_in_flight.append(command_buffer)
+    if PROFILE and self.profile_key is not None: self.dev.mtl_profile_keys[id(command_buffer)] = self.profile_key
     if wait:
       wait_check(command_buffer)
       return command_buffer.GPUEndTime() - command_buffer.GPUStartTime()

tinygrad/runtime/ops_null.py

@@ -17,9 +17,9 @@ class NullRenderer(CStyleLanguage):
     return assemble_linear(prg, lin, self.target.arch)
 
 class NullProgram(Program['NullDevice']):
-  def __init__(self, dev:'NullDevice', obj:TinyELF): self.device, self.name = dev.device, obj.name
+  def __init__(self, dev:'NullDevice', obj:TinyELF): self.device, self.name, self.profile_key = dev.device, obj.name, obj.profile_key
   def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int, ...]=(), wait=False, **kw):
-    with cpu_profile(self.name, self.device): return 1e-3
+    with cpu_profile(self.name, self.device, profile_key=self.profile_key): return 1e-3
 
 class NullAllocator(Allocator['NullDevice']):
   def _alloc(self, size, options): pass
@@ -38,13 +38,14 @@ class NullGraph(MultiGraphRunner):
       for (_,_,bufs,_),runtime in zip(self.calls, self.runtimes):
         # description based on command, copied from HCQ graph
         device = runtime.device if runtime is not None else f"{bufs[1].device}:SDMA:0"
-        descs.append((device, runtime.name if runtime is not None else f"{bufs[1].device} -> {bufs[0].device}", count:=event_count.get(device, 0)))
+        descs.append((device, runtime.name if runtime is not None else f"{bufs[1].device} -> {bufs[0].device}",
+                      runtime.profile_key if runtime is not None else None, count:=event_count.get(device, 0)))
         event_count[device] = count+1
       # pack events evenly per device
       dur, sigs, ents = max(1, math.ceil((perf_counter_us()-st)/max(event_count.values()))), [], []
-      for i,(device,name,count) in enumerate(descs):
+      for i,(device,name,profile_key,count) in enumerate(descs):
         sigs += [st+count*dur, st+(count+1)*dur]
-        ents.append(ProfileGraphEntry(device, name, 2*i, 2*i+1))
+        ents.append(ProfileGraphEntry(device, name, 2*i, 2*i+1, profile_key))
       cpu_events.append(ProfileGraphEvent(ents, [], sigs))
     return 1e-1
 

tinygrad/runtime/ops_nv.py

@@ -22,7 +22,7 @@ nv_gpu = nv_570 # default to 570
 PMA = ContextVar("PMA", abs(VIZ.value)>=2)
 
 @dataclass(frozen=True)
-class ProfilePMAEvent(ProfileEvent): device:str; kern:str; blob:bytes; exec_tag:int # noqa: E702
+class ProfilePMAEvent(ProfileEvent): device:str; kern:str; blob:bytes; exec_tag:int; profile_key:bytes|None=None # noqa: E702
 
 class NVSignal(HCQSignal):
   def _sleep(self, time_spent_since_last_sleep_ms:int):
@@ -335,12 +335,12 @@ class NVProgram(HCQProgram['NVDevice']):
     if self.dev.pma_enabled:
       self.dev.synchronize()
       if pma_blob:=self.dev._prof_readback():
-        Compiled.profile_events += [ProfilePMAEvent(self.dev.device, self.name, pma_blob, self.dev.prof_exec_counter)]
+        Compiled.profile_events += [ProfilePMAEvent(self.dev.device, self.name, pma_blob, self.dev.prof_exec_counter, self.profile_key)]
     return res
 
 class NVAllocator(HCQAllocator['NVDevice']):
   def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
-    return self.dev.iface.alloc(size, cpu_access=options.cpu_access, host=options.host)
+    return self.dev.iface.alloc(size, cpu_access=options.cpu_access, host=options.host, zero=options.zero)
 
   def _do_free(self, opaque:HCQBuffer, options:BufferSpec): self.dev.iface.free(opaque)
 
@@ -565,7 +565,7 @@ class PCIIface(PCIIfaceBase):
 
     # Setup classes for the GPU
     self.gpfifo_class, self.compute_class, self.dma_class = (gsp:=self.dev_impl.gsp).gpfifo_class, gsp.compute_class, gsp.dma_class
-    self.viddec_class = None
+    self.viddec_class = gsp.viddec_class
 
   def setup_usermode(self): return 0xce000000, self.pci_dev.map_bar(bar=0, fmt='I', off=0xbb0000, size=0x10000)
   def setup_vm(self, vaspace): pass
@@ -603,7 +603,7 @@ class NVDevice(HCQCompiled[NVSignal]):
 
     vaspace_params = nv_gpu.NV_VASPACE_ALLOCATION_PARAMETERS(vaBase=0x1000, vaSize=0x1fffffb000000,
       flags=nv_gpu.NV_VASPACE_ALLOCATION_FLAGS_ENABLE_PAGE_FAULTING | nv_gpu.NV_VASPACE_ALLOCATION_FLAGS_IS_EXTERNALLY_OWNED)
-    vaspace = self.iface.rm_alloc(self.nvdevice, nv_gpu.FERMI_VASPACE_A, vaspace_params)
+    self.vaspace = vaspace = self.iface.rm_alloc(self.nvdevice, nv_gpu.FERMI_VASPACE_A, vaspace_params)
 
     self.iface.setup_vm(vaspace)
 
@@ -643,7 +643,8 @@ class NVDevice(HCQCompiled[NVSignal]):
     notifier = self.iface.alloc(48 << 20, uncached=True)
     params = nv_gpu.NV_CHANNELGPFIFO_ALLOCATION_PARAMETERS(gpFifoOffset=gpfifo_area.va_addr+offset, gpFifoEntries=entries, hContextShare=ctxshare,
       hObjectError=notifier.meta.hMemory, hObjectBuffer=self.virtmem if video else gpfifo_area.meta.hMemory,
-      hUserdMemory=(ctypes.c_uint32*8)(gpfifo_area.meta.hMemory), userdOffset=(ctypes.c_uint64*8)(entries*8+offset), engineType=19 if video else 0)
+      hUserdMemory=(ctypes.c_uint32*8)(gpfifo_area.meta.hMemory), userdOffset=(ctypes.c_uint64*8)(entries*8+offset), engineType=19 if video else 0,
+      hVASpace=self.vaspace if video and self.is_nvd() else 0) # gsp has no default vaspace, rm maps the decoder ctx into its own
     gpfifo = self.iface.rm_alloc(channel_group, self.iface.gpfifo_class, params)
 
     if compute:
@@ -709,22 +710,22 @@ class NVDevice(HCQCompiled[NVSignal]):
   def _ensure_has_vid_hw(self, w, h):
     if self.iface.viddec_class is None: raise RuntimeError(f"{self.device} Video decoder class not available.")
 
-    coloc_size = round_up((round_up(h, 64) * round_up(h, 64)) + (round_up(w, 64) * round_up(h, 64) // 16), 2 << 20)
+    coloc_sz = round_up((round_up(h, 64) * round_up(h, 64)) + (round_up(w, 64) * round_up(h, 64) // 16), 2 << 20)
     self.intra_top_off = round_up(h, 64) * (608 + 4864 + 152 + 2000)
     intra_unk_size = ((2 << 20) if self.iface.viddec_class >= nv_gpu.NVCFB0_VIDEO_DECODER else 0)
     self.intra_unk_off = (round_up(self.intra_top_off, 0x10000) + (64 << 10)) if intra_unk_size > 0 else None
-    filter_size = round_up(round_up(self.intra_top_off, 0x10000) + (64 << 10) + intra_unk_size, 2 << 20)
+    filter_sz = round_up(round_up(self.intra_top_off, 0x10000) + (64 << 10) + intra_unk_size, 2 << 20)
 
     if not hasattr(self, 'vid_gpfifo'):
       self.vid_gpfifo = self._new_gpu_fifo(self.gpfifo_area, 0, self.nvdevice, offset=0x200000, entries=2048, compute=False, video=True)
-      self.vid_coloc_buf, self.vid_filter_buf = self.allocator.alloc(coloc_size), self.allocator.alloc(filter_size)
-      self.vid_stat_buf = self.allocator.alloc(0x1000)
+      self.vid_coloc_buf, self.vid_filter_buf = (self.allocator.alloc(sz, BufferSpec(zero=True)) for sz in [coloc_sz, filter_sz])
+      self.vid_stat_buf = self.allocator.alloc(0x1000, BufferSpec(zero=True))
       NVVideoQueue().wait(self.timeline_signal, self.timeline_value - 1) \
                     .setup(copy_class=self.iface.viddec_class) \
                     .signal(self.timeline_signal, self.next_timeline()).submit(self)
     else:
-      if coloc_size > self.vid_coloc_buf.size: self.vid_coloc_buf, _ = self._realloc(self.vid_coloc_buf, coloc_size, force=True)
-      if filter_size > self.vid_filter_buf.size: self.vid_filter_buf, _ = self._realloc(self.vid_filter_buf, filter_size, force=True)
+      if coloc_sz > self.vid_coloc_buf.size: self.vid_coloc_buf,_= self._realloc(self.vid_coloc_buf, coloc_sz, BufferSpec(zero=True), force=True)
+      if filter_sz > self.vid_filter_buf.size: self.vid_filter_buf,_= self._realloc(self.vid_filter_buf, filter_sz, BufferSpec(zero=True), force=True)
 
   def hw_copy_queues(self): return super().hw_copy_queues() + ([("NVDEC:0", NVVideoQueue)] if hasattr(self, 'vid_gpfifo') else [])
 

tinygrad/runtime/ops_python.py

@@ -89,7 +89,7 @@ class PythonProgram(Program['PythonDevice']):
               if g: _store(m, o+j, v, src_dtypes[1])
           i += 1
           continue
-        if u.op is Ops.AFTER: values[u] = src_values[0]
+        if u.op is Ops.AFTER or (u.op is Ops.BITCAST and u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL)): values[u] = src_values[0]
         elif u.op is Ops.PARAM and u.addrspace is AddrSpace.ALU: values[u] = [pvals.pop(0)] * warp_size
         elif u.op in {Ops.PARAM, Ops.BUFFER}:
           storage_fmt = storage_fmt_for_dtype(u.dtype)
@@ -114,7 +114,8 @@ class PythonProgram(Program['PythonDevice']):
               if ox < 0 or ox >= u.src[0]._shape[1] or oy < 0 or oy >= u.src[0]._shape[0]: ret.append((m, None))
               else: ret.append((m, ox*4 + oy*u.src[0]._shape[1]*4))
           else:
-            for m,o in zip(src_values[0], src_values[1]): ret.append((m,o))
+            scale = u.src[0].dtype.itemsize // u.src[0].src[0].dtype.itemsize if u.src[0].op is Ops.BITCAST else 1
+            for m,o in zip(src_values[0], src_values[1]): ret.append((m[0], m[1]+o*scale) if isinstance(m, tuple) else (m, o*scale))
           values[u] = ret
         elif u.op is Ops.RANGE:
           if u not in values: values[u] = [0] * warp_size

tinygrad/runtime/support/compiler_qcom.py

@@ -1,4 +1,4 @@
-import ctypes, struct, platform, pathlib, shutil, tarfile, tempfile
+import ctypes, struct, platform, pathlib, shutil
 from tinygrad.device import Compiler
 from tinygrad.helpers import DEBUG, system, fetch
 from tinygrad.runtime.support.compiler_mesa import disas_adreno
@@ -12,8 +12,9 @@ class QCOMCompiler(Compiler):
     assert arch.split(',')[0] == "a630", "only a630 supported"
     if platform.machine() == "aarch64": self.arch, self.chip_id, self.llvm_inst = arch, 0x6030001, llvm_qcom.cl_compiler_create_llvm_instance()
     else:
-      self.arch, self.chip_id, self.fs, root = arch, 0x6030001, tempfile.TemporaryDirectory(), pathlib.Path(__file__).parents[3]
-      with tarfile.open(fetch('https://git.tinygrad.win/sirhcm/images/releases/download/v2/qcomcl.tar.gz')) as t: t.extractall(fs:=self.fs.name)
+      # extract once into the download cache, all processes share the rootfs (extract=True)
+      self.arch, self.chip_id = arch, 0x6030001
+      fs, root = fetch('https://git.tinygrad.win/sirhcm/images/releases/download/v2/qcomcl.tar.gz', extract=True), pathlib.Path(__file__).parents[3]
       self.compiler_process = self.server(f"{qemu} -cpu max,pauth=off -L {fs} {fs}/usr/bin/python3" if (qemu:=shutil.which("qemu-aarch64-static"))
                                           else (f"docker run --rm -i --platform linux/aarch64 -v {fs}/usr:/usr -v {root}:{root} "
                                                 f"-e PYTHONPATH={root} -e QEMU_CPU=max,pauth=off gcr.io/distroless/static python3"), arch)

tinygrad/runtime/support/hcq.py

@@ -295,7 +295,8 @@ class HCQSignal(Generic[HCQDeviceType]):
     if not_passed and self.value < value: raise RuntimeError(f"Wait timeout: {timeout} ms! (the signal is not set to {value}, but {self.value})")
 
 @contextlib.contextmanager
-def hcq_profile(dev:HCQCompiled, enabled, desc, queue_type:Callable[[], HWQueue]|None=None, queue:HWQueue|None=None, dev_suff:str|None=None):
+def hcq_profile(dev:HCQCompiled, enabled, desc, queue_type:Callable[[], HWQueue]|None=None, queue:HWQueue|None=None, dev_suff:str|None=None,
+                profile_key:bytes|None=None):
   st, en = (dev.new_signal(), dev.new_signal()) if enabled else (None, None)
   assert queue is not None or queue_type is not None, "Either queue or queue_type must be provided"
 
@@ -309,7 +310,8 @@ def hcq_profile(dev:HCQCompiled, enabled, desc, queue_type:Callable[[], HWQueue]
     elif enabled and queue_type is not None:
       queue_type().wait(dev.timeline_signal, dev.timeline_value - 1).timestamp(en).signal(dev.timeline_signal, dev.next_timeline()).submit(dev)
 
-    if enabled and PROFILE: dev.sig_prof_records.append((unwrap(st), unwrap(en), desc, f"{dev.device}:{dev_suff}" if dev_suff else dev.device))
+    if enabled and PROFILE: dev.sig_prof_records.append((unwrap(st), unwrap(en), desc, f"{dev.device}:{dev_suff}" if dev_suff else dev.device,
+                                                         profile_key))
 
 class HCQArgsState(Generic[ProgramType]):
   def __init__(self, buf:HCQBuffer, prg:ProgramType, bufs:tuple[HCQBuffer, ...], vals:tuple[sint|None, ...]=()):
@@ -332,8 +334,9 @@ class CLikeArgsState(HCQArgsState[ProgramType]):
 class HCQProgram(Program[HCQDeviceType]):
   def __init__(self, args_state_t:Type[HCQArgsState], dev:HCQDeviceType, obj:TinyELF, kernargs_alloc_size:int, base:int|None=None):
     self.args_state_t, self.dev, self.name, self.signature, self.kernargs_alloc_size = args_state_t, dev, obj.name, obj.signature, kernargs_alloc_size
+    self.profile_key = obj.profile_key
     self.prof_prg_counter = next(self.dev.prof_prg_counter)
-    if PROFILE: Compiled.profile_events += [ProfileProgramEvent(dev.device, obj.name, obj.lib, base, self.prof_prg_counter)]
+    if PROFILE: Compiled.profile_events += [ProfileProgramEvent(dev.device, obj.name, obj.lib, base, self.prof_prg_counter, self.profile_key)]
 
   @staticmethod
   def _fini(dev, buf, spec): dev.allocator.free(buf, buf.size, spec)
@@ -372,7 +375,7 @@ class HCQProgram(Program[HCQDeviceType]):
     q = unwrap(self.dev.hw_compute_queue_t)().wait(self.dev.timeline_signal, self.dev.timeline_value - 1).memory_barrier()
 
     self.dev.prof_exec_counter += 1
-    with hcq_profile(self.dev, queue=q, desc=self.name, enabled=wait or PROFILE) as (sig_st, sig_en):
+    with hcq_profile(self.dev, queue=q, desc=self.name, enabled=wait or PROFILE, profile_key=self.profile_key) as (sig_st, sig_en):
       q.exec(self, kernargs, global_size, local_size)
 
     q.signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
@@ -401,7 +404,7 @@ class HCQCompiled(Compiled, Generic[SignalType]):
     self.signal_t, self.hw_compute_queue_t, self.hw_copy_queue_t = signal_t, comp_queue_t, copy_queue_t
 
     self.timeline_value:int = 1
-    self.sig_prof_records:list[tuple[HCQSignal, HCQSignal, str|TracingKey, str]] = []
+    self.sig_prof_records:list[tuple[HCQSignal, HCQSignal, str|TracingKey, str, bytes|None]] = []
     self.prof_exec_counter:int = 0
     self.prof_prg_counter = itertools.count(0)
 
@@ -437,7 +440,7 @@ class HCQCompiled(Compiled, Generic[SignalType]):
 
     if self.timeline_value > (1 << 31): self._wrap_timeline_signal()
     if PROFILE:
-      Compiled.profile_events += [ProfileRangeEvent(dev, name, st.timestamp, en.timestamp) for st,en,name,dev in self.sig_prof_records]
+      Compiled.profile_events += [ProfileRangeEvent(dev, name, st.timestamp, en.timestamp, pk) for st,en,name,dev,pk in self.sig_prof_records]
       self.sig_prof_records = []
 
   def next_timeline(self):

tinygrad/runtime/support/hcq2.py

@@ -8,12 +8,12 @@ from tinygrad.device import Device, Buffer, BufferSpec, Compiled, LRUAllocator,
 from tinygrad.device import ProfileDeviceEvent, ProfileGraphEntry, ProfileGraphEvent
 from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, graph_rewrite, rewrite_group, GroupOp
 from tinygrad.uop.symbolic import symbolic
-from tinygrad.dtype import dtypes, truncate
+from tinygrad.dtype import dtypes, truncate, DType
 from tinygrad.runtime.support.hcq import MMIOInterface, HCQBuffer
 from tinygrad.runtime.support.memory import BumpAllocator
 from tinygrad.renderer import Renderer, Estimates
 from tinygrad.engine.realize import to_program, get_call_arg_uops, get_call_name, get_call_outs_ins, estimate_uop
-from tinygrad.engine.realize import pm_flatten_linear
+from tinygrad.engine.realize import pm_flatten_linear, lower_and_compile
 
 # *****************
 # 0. helpers
@@ -36,10 +36,15 @@ class HCQInfo:
 
 def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c
 
-def unwrap_mstack(u):
+def unwrap_mstack(u:UOp) -> tuple[UOp, ...]:
   if u.op is Ops.MSTACK: return tuple(x for s in u.src for x in unwrap_mstack(s))
   return unwrap_mstack(u.src[0]) if u.op is Ops.MSELECT else (u,)
 
+def unwrap_view(v:UOp) -> tuple[UOp, int]:
+  return unwrap_view(v.src[0]) if v.op is Ops.BITCAST else (v.src[0], v.src[1].val) if v.op is Ops.SHRINK else (v, 0)
+
+# patches
+
 def is_value_known_at_link(val:UOp) -> bool:
   runtime_reads = [u for u in val.toposort() if u.op in (Ops.LOAD, Ops.INDEX)]
   addressed_bufs = [b for g in val.toposort() if g.op is Ops.GETADDR for b in unwrap_mstack(g.buf_uop)]
@@ -48,32 +53,28 @@ def is_value_known_at_link(val:UOp) -> bool:
   return not val.variables() and not runtime_reads and all(b.op is not Ops.PARAM or b.tag is not None for b in addressed_bufs)
 
 def make_patches(buf:UOp, patches:Sequence[tuple[sint, UOp]]) -> tuple[UOp, ...]:
-  def _mk_store(ps:list[tuple[sint, UOp]], tag:str|None) -> UOp:
-    offs = UOp(Ops.STACK, dtypes.int, tuple(UOp.const(off // buf.dtype.itemsize, dtypes.int) for off,_ in ps))
-    vals = UOp(Ops.STACK, ps[0][1].dtype, tuple(val for _,val in ps))
-    return buf.index(offs, dtype=vals.dtype).store(vals).rtag(tag)
+  # group patches into stacks: (tag, type, offset). offset is used for shrink later
+  groups:dict[tuple[str|None, DType, sint], list[tuple[sint, UOp]]] = collections.defaultdict(list)
+  for off, val in patches:
+    tag = "link" if is_value_known_at_link(val) else "inputs" if val.op is Ops.GETADDR else None
+    groups[(tag, (v:=(val.bitcast(buf.dtype) if val.dtype.itemsize == buf.dtype.itemsize else val)).dtype, off % v.dtype.itemsize)].append((off, v))
 
-  patches = [(off, val.cast(buf.dtype) if val.dtype.itemsize == buf.dtype.itemsize else val) for off, val in patches]
-  link, runtime = partition(patches, lambda p: is_value_known_at_link(p[1]))
-  inputs, runtime = partition(runtime, lambda p: p[1].op is Ops.GETADDR)
-  return tuple(_mk_store(list(ps), tag) for cls, tag in ((link, "link"), (inputs, "inputs"), (runtime, None))
-               for _, ps in itertools.groupby(sorted(cls, key=lambda p: p[1].dtype), key=lambda p: p[1].dtype))
+  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))
+  return tuple(ret)
 
-def make_binary_patch(buf:UOp, blob:bytes) -> UOp:
-  data = UOp(Ops.BINARY, src=(), arg=blob).bitcast(buf.dtype)
-  r = UOp.range(len(blob) // buf.dtype.itemsize, 0, dtype=dtypes.int, src=(buf, data))
-  return buf.index(r).store(data.index(r).load()).end(r).rtag("link")
+def make_binary_patch(buf:UOp, blob:bytes) -> UOp: return buf.store(UOp(Ops.BINARY, src=(), arg=blob).bitcast(buf.dtype)).rtag("link")
 
-def make_cmdbuf(lin, devs, buf:UOp|None=None):
+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(make_binary_patch(cmdbuf, bytes(blob)), *make_patches(cmdbuf, patches))
-
-def make_signal(devs, slot:int=0, tag:str="signal") -> UOp:
-  return UOp.placeholder((1,), dtypes.uint64, slot, device=devs, volatile=True).rtag(tag)
+  return cmdbuf.after(*dep, make_binary_patch(cmdbuf, bytes(blob)), *make_patches(cmdbuf, patches))
 
 def make_submit(*cmds, devs:str|tuple[str, ...], queue:str) -> UOp:
   return UOp.custom_function("submit_cmdbuf", UOp(Ops.LINEAR, src=tuple(cmds), arg=(to_tuple(devs), queue)))
@@ -87,6 +88,8 @@ def encode_kernargs_clike(call:UOp, prg:UOp, devs:str|tuple[str, ...]) -> UOp:
   words = [get_call_arg_uops(call)[gi].getaddr(devs) for gi in info.globals] + list(info.vars)
   return buf.after(*make_patches(buf, list(zip(itertools.accumulate((w.dtype.itemsize for w in words), initial=0), words))))
 
+def make_buf(devs, slot:int=0, tag:str="signal") -> UOp: return UOp.placeholder((1,), dtypes.uint64, slot, device=devs, volatile=True, tag=tag)
+
 # *****************
 # 0.1. prep: replace buffers with params
 
@@ -107,6 +110,10 @@ def _staging() -> Buffer: return Buffer("CPU", STAGING_SIZE, dtypes.uint8, preal
 
 def _need_staging(a, b): return all_devices_in(a.device, HCQ_DEVS - {"CPU"}) and not all_devices_in(b.device, HCQ_DEVS)
 
+def stage_copy_ext(call:UOp) -> UOp|None:
+  if (d:=next((d for b in call.src[1:] for d in to_tuple(b.device) if not d.startswith("CPU")), None)) is None: return None
+  return pm.rewrite(call) if (pm:=getattr(Device[d], "pm_stage_copy", None)) is not None else None
+
 def stage_copy(dst:UOp, src:UOp) -> UOp|None:
   if not (_need_staging(src, dst) or _need_staging(dst, src)): return None
 
@@ -122,20 +129,22 @@ def stage_copy(dst:UOp, src:UOp) -> UOp|None:
 # 1.2. prep: kernel copies
 
 def _get_enqueue_devs(call:UOp) -> Any|None:
+  if call.src[0].op not in (Ops.PROGRAM, Ops.COPY): return None # only these bodies can be enqueued
   if not (bufs:=call.src[1:]) or not all(all_devices_in(b.device, HCQ_DEVS) for b in bufs): return None
   if call.src[0].op is Ops.COPY: bufs = bufs[::-1] # copies push from the src device: p2p writes are faster than reads
   devs = min(bufs, key=lambda b: to_tuple(b.device)[0].startswith("CPU")).device # prio to enqueue on not CPU device
   return devs if all_devices_in(devs, HCQ_DEVS) else None
 
-def kernel_copy(call:UOp, dst:UOp, src:UOp) -> UOp|None:
+def copy_with_kernel(call:UOp, dst:UOp, src:UOp) -> UOp|None:
   if (devs:=_get_enqueue_devs(call)) is None or Device[(dev:=to_tuple(devs)[0])].has_copy_queue: return None
   d, s = (UOp.param(i, dst.dtype, (n:=dst.max_numel(),), device=devs) for i in range(2))
   ast = d.index(r:=UOp.range(n, 0)).store(s.index(r).load()).end(r).sink(arg=KernelInfo(name="copy"), tag=1)
   return call.replace(src=(to_program(ast, Device[dev].renderer), dst, src))
 
 pm_insert_copy_staging = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), stage_copy_ext),
   (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy),
-  (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src")), name="call"), kernel_copy)
+  (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src")), name="call"), copy_with_kernel)
 ])
 
 # *****************
@@ -157,7 +166,7 @@ def _get_deps(ctx:DepsTracker, bufs_by_lane:list[list[Any]], write, key:tuple[tu
     dep_lanes += [(dep, dlane, lane) for dep, dlane in ctx.access_resources(bufs, written, (key, lane))]
   return dep_lanes
 
-def _build_wait_cmds(slots:dict[str, int], dep_lanes:list[tuple[tuple, int, int]], devices:tuple[str, ...], queue:str) -> tuple[list[UOp], set[int]]:
+def _build_wait_ins(slots:dict[str, int], dep_lanes:list[tuple[tuple, int, int]], devices:tuple[str, ...], queue:str) -> tuple[list[UOp], set[int]]:
   # opt1: same-queue ops are fifo-ordered
   if devices[0].split(":")[0] in {"AMD", "QCOM", "CPU"} or queue.startswith("COPY"):
     dep_lanes = [(dep, dlane, lane) for dep, dlane, lane in dep_lanes if (dep[0][dlane], dep[1]) != (devices[lane], queue)]
@@ -170,7 +179,7 @@ def _build_wait_cmds(slots:dict[str, int], dep_lanes:list[tuple[tuple, int, int]
   waits = []
   for (ddevs, dqueue, dtag), by_lane in deps.items():
     for ls in itertools.zip_longest(*(by_lane[lane] for lane in range(len(devices)))):
-      s = UOp.mstack(*[make_signal(d, tag="sentinel_signal") if dl is None else make_signal(ddevs[dl], slots[dqueue]) for dl, d in zip(ls, devices)])
+      s = UOp.mstack(*[make_buf(d, tag="sentinel_signal") if dl is None else make_buf(ddevs[dl], slots[dqueue]) for dl, d in zip(ls, devices)])
       waits.append(UOp(Ops.INS, arg="wait", src=(s, UOp.const(dtag + 1, dtypes.uint64))))
   return waits, {dtag for _, _, dtag in deps}
 
@@ -188,26 +197,52 @@ def _build_finalizers(batch:list[tuple[UOp, tuple[str, ...]]], batch_info:list[t
 
     # to finalize the batch, sync all accesses from other devices to buffers that belong to this device
     fin_deps = [dl for dl in _get_deps(tracker, [list(dev_bufs[d].values()) for d in devs], None, key=(devs, "COMPUTE:0", n)) if dl[0][2] < n]
-    waits, cur_signal_tags = _build_wait_cmds(slots, fin_deps, devs, "COMPUTE:0")
+    waits, cur_signal_tags = _build_wait_ins(slots, fin_deps, devs, "COMPUTE:0")
     signal_tags |= cur_signal_tags
 
     # wait the syncs and signal the device epoch, then bump the timeline on the host
-    tl_signal, tl_value = make_signal(devs, tag="timeline_signal"), make_signal(devs, tag="timeline_value")
+    tl_signal, tl_value = make_buf(devs, tag="timeline_signal"), make_buf(devs, tag="timeline_value")
     fin_submit = make_submit(*waits, UOp(Ops.INS, arg="store", src=(tl_signal, tl_value.index(0))), devs=devs, queue="COMPUTE:0")
     epoch = (epoch_slot:=tl_value.after(fin_submit).index(0)).load()
 
     # fence once per device group on this schedule's previous epoch
     qs = dedup([qn for bdevs, qn in batch_info if set(bdevs) & set(devs)])
-    sched_epoch = make_signal(devs, next(UOp.unique_num))
+    sched_epoch = make_buf(devs, next(UOp.unique_num), tag="epoch")
 
     wait_device_epoch = (done:=tl_signal.after(loop:=UOp.loop(0)).index(0).load()).end(loop, done < sched_epoch.index(0).load())
     fences.append(make_call("hcq_fence", UOp.sink(wait_device_epoch), HCQInfo(devs)))
 
-    # queues of other groups wait on these signals, so reset them only after every group reached its epoch
-    if qs: resets.append(make_call("hcq_reset", UOp.sink(*[make_signal(devs, slots[q]).index(0).store(0) for q in qs]), HCQInfo(devs)))
+    # queues of other groups wait on these signals, reset them after every group reached its epoch
+    rst = functools.reduce(lambda a,q: a+(make_buf(devs, slots[q]).after(*a[-1:]).index(0).store(0),), qs, cast(tuple[UOp, ...], ()))
+    if rst: resets.append(make_call("hcq_reset", UOp.sink(*rst), HCQInfo(devs)))
     fins.append(make_call("hcq_finalizer", UOp.sink(epoch_slot.store(epoch + 1), sched_epoch.after(fin_submit).index(0).store(epoch)), HCQInfo(devs)))
   return fences + resets, fins, signal_tags
 
+def _merged_hcq_call(calls:list[UOp]) -> UOp: # TODO: simplify?
+  if len(calls) == 1: return calls[0]
+  devs, queue = get_submit(calls[0]).src[0].arg
+  body = make_submit(*[cmd for c in calls for cmd in get_submit(c).src[0].src], devs=devs, queue=queue).sink()
+  return make_call(f"submit {queue} ({len(calls)})", body,
+    replace(calls[0].arg.aux, estimates=sum((c.arg.aux.estimates for c in calls), start=Estimates()).simplify()))
+
+def _merge_queues(submits:list[UOp]) -> list[UOp]:
+  new_src:list[UOp] = []
+  opened_qs:dict[tuple[tuple[str, ...], str], list[UOp]] = {} # (devs, queue) -> list of hcq calls, kept in submit order
+  limits:dict[tuple[tuple[str, ...], str], int] = collections.defaultdict(lambda: JIT_BATCH_SIZE.value)
+
+  for call in submits:
+    devs, queue = get_submit(call).src[0].arg
+    if (old:=opened_qs.pop(key:=(devs, queue), None)) is not None:
+      if limits[key] and len(old) >= limits[key]: new_src, old, limits[key] = new_src + [_merged_hcq_call(old)], [], limits[key] * 2
+      new_rec = old + [call]
+    else:
+      # no such queue opened: close every open submit on this queue that shares a device, so submit order is kept
+      closing = [k for k in opened_qs if k[1] == queue and set(k[0]) & set(devs)]
+      new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in closing]
+      new_rec = [call]
+    opened_qs[(devs, queue)] = new_rec
+  return new_src + [_merged_hcq_call(c) for c in opened_qs.values()]
+
 def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> list[UOp]:
   batch_info = [(devices, "COMPUTE:0" if call.src[0].op is Ops.PROGRAM else "COPY:0") for call, devices in batch]
 
@@ -218,7 +253,7 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> li
   call_waits:list[list[UOp]] = []
   for tag, ((call, _), (devices, queue)) in enumerate(zip(batch, batch_info)):
     deps = _get_deps(deps_tracker, _get_call_bufs_by_lane(call, devices), get_call_outs_ins(call)[0], key=(devices, queue, tag))
-    cmds, cur_signal_tags = _build_wait_cmds(slots, deps, devices, queue)
+    cmds, cur_signal_tags = _build_wait_ins(slots, deps, devices, queue)
     call_waits.append(cmds)
     signal_tags |= cur_signal_tags
 
@@ -230,24 +265,24 @@ def _finalize_batch(batch:list[tuple[UOp, tuple[str, ...]]], profile:bool) -> li
   for tag, ((call, _), (devices, queue), q) in enumerate(zip(batch, batch_info, call_waits)):
     # first queue use, sync prior device work with the device timeline
     if batch_info.index((devices, queue)) == tag:
-      epoch = make_signal(devices, tag="timeline_value").index(0) - 1
-      q = [UOp(Ops.INS, arg="barrier", src=()), UOp(Ops.INS, arg="wait", src=(make_signal(devices, tag="timeline_signal"), epoch))] + q
+      epoch = make_buf(devices, tag="timeline_value").index(0) - 1
+      q = [UOp(Ops.INS, arg="barrier", src=()), UOp(Ops.INS, arg="wait", src=(make_buf(devices, tag="timeline_signal"), epoch))] + q
 
     # 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)))
 
-    ts_ins = [UOp(Ops.INS, arg="timestamp", src=(make_signal(devices, s),)) for s in ts_ids]
+    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:]
 
     # signal the queue if someone waits for us
-    if tag in signal_tags: q += [UOp(Ops.INS, arg="store", src=(make_signal(devices, slots[queue]), UOp.const(tag + 1, dtypes.uint64)))]
+    if tag in signal_tags: q += [UOp(Ops.INS, arg="store", src=(make_buf(devices, slots[queue]), UOp.const(tag + 1, dtypes.uint64)))]
     src.append(make_call(f"submit {name}", make_submit(*q, devs=devices, queue=queue).sink(), info))
 
   # append batch timestamps to finalizers
   fins = [f.replace(arg=replace(f.arg, aux=replace(a:=f.arg.aux, kernels=tuple(x for x in kerns if set(x[0]) & set(a.device))))) for f in fins]
-  return fences + src + fins
+  return fences + _merge_queues(src) + fins
 
 def sched_hcq_batches(l:UOp, profile:bool) -> UOp:
   srcs:list[UOp] = []
@@ -257,50 +292,23 @@ def sched_hcq_batches(l:UOp, profile:bool) -> UOp:
     else: srcs, batch = srcs + _finalize_batch(batch, profile) + [call], []
   return l.replace(src=tuple(srcs + _finalize_batch(batch, profile)))
 
-# *****************
-# 3. merge into queues
-
-def _merged_hcq_call(calls:list[UOp]) -> UOp: # TODO: simplify?
-  if len(calls) == 1: return calls[0]
-  devs, queue = get_submit(calls[0]).src[0].arg
-  body = make_submit(*[cmd for c in calls for cmd in get_submit(c).src[0].src], devs=devs, queue=queue).sink()
-  return make_call(f"submit {queue} ({len(calls)})", body,
-    replace(calls[0].arg.aux, estimates=sum((c.arg.aux.estimates for c in calls), start=Estimates()).simplify()))
-
-def merge_queues(linear:UOp) -> UOp:
-  new_src:list[UOp] = []
-  opened_qs:dict[tuple[tuple[str, ...], str], list[UOp]] = {} # (devs, queue) -> list of hcq calls, kept in submit order
-  limits:dict[tuple[tuple[str, ...], str], int] = collections.defaultdict(lambda: JIT_BATCH_SIZE.value)
-
-  for call in linear.src:
-    # non-hcq call, fence or finalizer: close all open queues
-    if not isinstance(call.arg.aux, HCQInfo) or (call.arg.name or "").startswith("hcq_"):
-      new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in list(opened_qs)] + [call]
-      continue
-
-    devs, queue = get_submit(call).src[0].arg
-    if (old:=opened_qs.pop(key:=(devs, queue), None)) is not None:
-      if limits[key] and len(old) >= limits[key]: new_src, old, limits[key] = new_src + [_merged_hcq_call(old)], [], limits[key] * 2
-      new_rec = old + [call]
-    else:
-      # no such queue opened: close every open submit on this queue that shares a device, so submit order is kept
-      closing = [k for k in opened_qs if k[1] == queue and set(k[0]) & set(devs)]
-      new_src += [_merged_hcq_call(opened_qs.pop(k)) for k in closing]
-      new_rec = [call]
-    opened_qs[(devs, queue)] = new_rec
-  return linear.replace(src=tuple(new_src + [_merged_hcq_call(c) for c in opened_qs.values()]))
-
 pm_schedule_and_merge = PatternMatcher([(UPat(Ops.LINEAR, name="l"),
-  lambda ctx, l: merge_queues(sched_hcq_batches(l, ctx[1]).substitute(ctx[0], walk=True, enter_calls=True)))])
+  lambda ctx, l: sched_hcq_batches(l, ctx[1]).substitute(ctx[0], walk=True, enter_calls=True))])
 
 # *****************
 # 4.2. hcq lowering: ops to ir
 
+def encode_host_call(call:UOp) -> UOp|None:
+  if (pm:=getattr(Device[call.arg.aux.device[0]], "pm_host_lower", None)) is None: return None
+  body = graph_rewrite(call.src[0], pm, name="lower host access", enter_calls=True)
+  return None if body is call.src[0] else call.replace(src=(body, *call.src[1:]))
+
 def encode_cmdbuf(submit:UOp, lin:UOp) -> UOp|None:
   if (pm:=Device.get_class(lin.arg[0][0]).pm_lower) is None: return None
   return graph_rewrite(submit, pm, name=f"encode {lin.arg[0]}", enter_calls=True)
 pm_encode_cmdbufs = PatternMatcher([
-  (UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="lin"),), name="submit"), encode_cmdbuf)])
+  (UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="lin"),), name="submit"), encode_cmdbuf),
+  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), encode_host_call)])
 
 # *****************
 
@@ -328,14 +336,18 @@ def make_addr_table(call:UOp, gaddrs:list[UOp], name:str) -> tuple[UOp, dict[UOp
   return table, reads, fills, {g:slots[bare[g]] for g in gaddrs}
 
 def make_gather_loop(patches:list[UOp], table:UOp, slots:dict[UOp, int], lt_patches:list[UOp]) -> dict[UOp, UOp]:
-  (dst,), words = dedup(p.buf_uop for p in patches), [(off.val, slots[val]) for p in patches for off, val in zip(p.src[0].src[1].src, p.src[1].src)]
+  (dst,), words = dedup(p.buf_uop for p in patches), [(unwrap_view(p.src[0].src[0])[1] + off.val*(val.dtype.itemsize//p.buf_uop.dtype.itemsize),
+                                                       slots[val]) for p in patches for off,val in zip(p.src[0].src[1].src, p.src[1].src)]
 
   # build a runtime loop that writes every input address
   pairs = UOp.placeholder((2*len(words),), dtypes.uint32, next(UOp.unique_num), device=dst.device).rtag("systems")
   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())))
-  return {p: UOp(Ops.NOOP) for p in patches} | {patches[0]: dst.index(off, dtype=table.dtype).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))
 
@@ -344,20 +356,20 @@ def split_patches(call:UOp) -> UOp|None:
   lt_patches:list[UOp] = []
   body = graph_rewrite(call.src[0], pm_trim_link_patches, ctx=(rt_patches, lt_patches), name=f"trim link-time patches ({call.arg.name})")
 
-  # split patches
-  inputs, internals = partition(dedup(g for p in rt_patches for g in get_getaddrs(p)), is_input_addr)
+  # split patches. addresses read in the body go through the tables too
+  inputs, internals = partition(dedup([g for p in rt_patches for g in get_getaddrs(p)] + get_getaddrs(body)), is_input_addr)
   runtimes, systems = partition(internals, lambda g: any(x.tag in {"program", "kernargs", "cmdbuf"} for x in unwrap_mstack(g.buf_uop)))
   tables = [make_addr_table(call, gs, n) for gs,n in ((inputs, "inputs"), (runtimes, "runtime"), (systems, "systems"))]
   reads, fills = {k:v for _,r,_,_ in tables for k,v in r.items()}, [f for t in tables[1:] for f in t[2]] # inputs table is filled by exec
 
   ipatches = [p for p in rt_patches if p.tag == "inputs" and all(v in tables[0][3] for v in p.src[1].src)] # only getaddrs go to the table
   gathers = make_gather_loop(ipatches, tables[0][0], tables[0][3], lt_patches) if ipatches else {}
-  body = body.substitute({p:p.substitute(gathers | reads) for p in rt_patches})
+  body = body.substitute({p:p.substitute(gathers | reads) for p in rt_patches}).substitute(reads)
 
   lt_srcs = collections.defaultdict(list)
   for p in lt_patches: lt_srcs[p.buf_uop].append(p)
   return call.replace(src=(body, *call.src[1:], *[b.after(*ps) for b,ps in lt_srcs.items()], *fills),
-    arg=replace(call.arg, aux=replace(call.arg.aux, input_idxs=((call.arg.aux.device,
+    arg=replace(call.arg, aux=replace(call.arg.aux, input_idxs=((to_tuple(inputs[0].arg),
       tuple(sorted(dedup(b.arg.slot for g in inputs for b in unwrap_mstack(g.buf_uop))))),) if inputs else call.arg.aux.input_idxs)))
 pm_split_patches = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), split_patches)])
 
@@ -383,22 +395,22 @@ def replace_params(call:UOp) -> UOp|None:
   sub = {(b:=u.without_after): UOp.param(i, u.dtype, shape=b.shape, device=HCQ_RUNTIME_DEV.value, volatile=b.op is Ops.PARAM and b.arg.volatile)
          for i,u in enumerate(c_args)} | {v: v.replace(arg=replace(v.arg, slot=-1)) for v in variables if v.op is Ops.PARAM} | _rank_ranges(tops)
   info = replace(call.arg.aux, inputs=next((i for i,u in enumerate(c_args + refhold) if u.without_after.tag == "inputs"), None))
-  return call.replace(src=(body.substitute(sub).replace(arg="hcq_args"), *c_args, *refhold), arg=replace(call.arg, aux=info))
+  prg_sink = body.src[0].substitute(sub).replace(arg=KernelInfo("hcq_submit"), tag=1)
+  return call.replace(src=(body.replace(src=(prg_sink,)), *c_args, *refhold), arg=replace(call.arg, aux=info))
 pm_replace_params = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), replace_params)])
+  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq", src=(UPat(Ops.SINK),)),), name="call", allow_any_len=True), replace_params)])
 
 # *****************
 
 def resolve_getaddr_view(bv:UOp, g:UOp) -> UOp:
   base = bv.src[0].after(*g.src[0].src[1:] if g.src[0].op is Ops.AFTER else ())
-  if bv.op is Ops.BITCAST: return UOp(Ops.GETADDR, src=(base,), arg=g.arg)
-  itemsize = bv.src[0].dtype.itemsize if bv.src[0].without_after.op in (Ops.BUFFER, Ops.MSTACK, Ops.MSELECT) else bv.dtype.itemsize
-  return UOp(Ops.GETADDR, src=(base,), arg=g.arg) + UOp.const(bv.src[1].val * itemsize, dtypes.uint64)
+  addr = UOp(Ops.GETADDR, src=(base,), arg=g.arg)
+  return addr if bv.op is Ops.BITCAST else addr + UOp.const(bv.src[1].val * bv.dtype.itemsize, dtypes.uint64)
 
 pm_early_simplify = PatternMatcher([
   (UPat(Ops.GETADDR, src=(UPat((Ops.SHRINK, Ops.BITCAST), name="bv").or_after(),), name="g"), resolve_getaddr_view),
-  (UPat(Ops.INDEX, src=(UPat(Ops.SHRINK, name="bv"),), allow_any_len=True, name="x"),
-   lambda bv,x: x.replace(src=(bv.src[0], x.src[1] + bv.src[1].cast(x.src[1].dtype), *x.src[2:]))),
+  (UPat(Ops.SHRINK, src=(UPat(Ops.SHRINK, name="bv"), UPat(), UPat()), name="x"),
+   lambda bv,x: bv.src[0].shrink(((start:=bv.src[1]+x.src[1], start+x.src[2]),))),
 ])
 
 # *****************
@@ -420,15 +432,6 @@ def pack_hcq_placeholders(call:UOp) -> UOp|None:
 pm_pack_placeholders = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="hcq"),), name="call", allow_any_len=True), pack_hcq_placeholders)])
 
-# *****************
-# 8. callify hcq programs
-
-def callify_hcq(call:UOp, cf:UOp) -> UOp:
-  prg = to_program(cf.src[0].replace(arg=KernelInfo("hcq_submit"), tag=1), Device[HCQ_RUNTIME_DEV.value].renderer)
-  return call.replace(src=(cf.replace(src=(prg,), arg="hcq"), *call.src[1:]))
-pm_callify_hcq = PatternMatcher([(UPat(Ops.CALL, src=(
-  UPat(Ops.CUSTOM_FUNCTION, arg="hcq_args", src=(UPat(Ops.SINK),), name="cf"),), name="call", allow_any_len=True), callify_hcq)])
-
 # *****************
 # 9. merge submitters
 
@@ -464,8 +467,7 @@ def hcq_lower(linear:UOp, pm_encode:PatternMatcher) -> UOp:
   linear = graph_rewrite(linear, pm_split_patches, walk=True, name="split patches")
 
   # and compile it
-  linear = graph_rewrite(linear, pm_replace_params, name="replace params")
-  return graph_rewrite(linear, pm_callify_hcq, name="callify hcq", enter_calls=True)
+  return lower_and_compile(graph_rewrite(linear, pm_replace_params, walk=True, name="replace params"))
 
 @rewrite_group(lambda linear,input_uops,profile,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
 def hcq_compile(linear:UOp, input_uops:list[UOp]|None, profile:bool) -> UOp:
@@ -504,14 +506,16 @@ def push_stack(op, s): return UOp(Ops.STACK,
 
 def fold_binary(buf:UOp, blob:UOp) -> UOp:
   for b in (m.bufs if isinstance(m:=buf.buffer, MultiBuffer) else (m,)):
-    b.ensure_allocated().as_memoryview(force_zero_copy=True, no_sync=True).cast('B')[:len(blob.arg)] = blob.arg
+    b.ensure_allocated()._buf.cpu_view().view(fmt='B')[:len(blob.arg)] = blob.arg
   return UOp(Ops.NOOP)
 
-def fold_const_store(buf:UOp, off:UOp, val:UOp) -> UOp:
+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))
-      b.ensure_allocated().as_memoryview(force_zero_copy=True, no_sync=True).cast('B')[(bo:=off.val*buf.dtype.itemsize):bo+len(data)] = data
+      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)
 
 def resolve_getaddr(buf:UOp, g:UOp) -> UOp:
@@ -532,10 +536,9 @@ pm_resolve_patches = PatternMatcher([
   (UPat(Ops.GETADDR, src=(UPat(name="buf"),), name="g"), resolve_getaddr),
 
   # folders
-  (UPat(name="buf").index(UPat(Ops.RANGE), allow_any_len=True)
-    .store(UPat.any(UPat(Ops.BINARY, name="blob"), UPat(Ops.BINARY, name="blob").bitcast()).index(UPat(Ops.RANGE), allow_any_len=True).load())
-    .end(UPat(Ops.RANGE)), fold_binary),
-  (UPat({Ops.BUFFER, Ops.MSTACK}, name="buf").index(UPat(Ops.STACK, name="off")).store(UPat(Ops.STACK, name="val")), fold_const_store),
+  (UPat(name="buf").store(UPat.any(UPat(Ops.BINARY, name="blob"), UPat(Ops.BINARY, name="blob").bitcast())), fold_binary),
+  (UPat((Ops.BITCAST, Ops.SHRINK, Ops.BUFFER, Ops.MSTACK), name="view")
+    .index(UPat(Ops.STACK, name="off")).store(UPat(Ops.STACK, name="val")), fold_const_store),
 ])
 
 pm_assert_no_afters = PatternMatcher([(UPat(Ops.AFTER, name="a"), lambda a: panic(RuntimeError, f"AFTER left at hcq_link: {a.src[0].op}"))])
@@ -563,6 +566,7 @@ def hcq_link(linear:UOp, cache=True) -> UOp:
 class HCQ2Compiled(Compiled):
   timestamp_divider: float = 1000.0
   wait_timeout_ms: float = 30000.0
+  rt_nbytes: int = 64 << 20 # scratch that single-run placeholders are carved out of
 
   def __init__(self, device:str, allocator:HCQAllocator, compilers:list[type[Renderer]], runtime, can_recover:bool=False, arch=None):
     self.can_recover = can_recover
@@ -570,14 +574,15 @@ class HCQ2Compiled(Compiled):
     self.pm_bufferize = PatternMatcher([
       (UPat(Ops.PARAM, tag="sentinel_signal"), lambda ctx: ctx[0].signal("sentinel", (1 << 64) - 1)),
       (UPat(Ops.PARAM, tag="timeline_signal"), lambda ctx: ctx[0].signal("timeline")),
-      (UPat(Ops.PARAM, tag="timeline_value"), lambda ctx: ctx[0].signal("value", 1)),
+      (UPat(Ops.PARAM, tag="timeline_value"), lambda ctx: ctx[0].signal("value", 1, device="CPU")),
+      (UPat(Ops.PARAM, tag="epoch", name="b"), lambda ctx, b: ctx[0].signal(b.arg.slot, device="CPU")),
       (UPat(Ops.PARAM, tag="signal", name="b"), lambda ctx, b: ctx[0].signal(b.arg.slot)),
       (UPat(Ops.PARAM, name="b"), lambda ctx, b: None if b.tag is None else ctx[0].new_buffer(b, cache=ctx[1]))
     ])
 
     super().__init__(device, allocator, compilers, runtime, None, arch=arch)
 
-    self.rt_allocator = BumpAllocator(64 << 20)
+    self.rt_allocator = BumpAllocator(self.rt_nbytes)
     self.prof_ents:dict[int, ProfileGraphEntry] = {}
 
   def collect_prof(self):
@@ -611,12 +616,12 @@ class HCQ2Compiled(Compiled):
       self.rt_allocator.alloc(b.max_numel() * b.dtype.itemsize, alignment=128))
 
   @functools.cache
-  def signal(self, name:str|int, init_value:int=0) -> Buffer:
-    buf = Buffer(self.device, 1, dtypes.uint64, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
-    buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] = init_value
+  def signal(self, name:str|int, init_value:int=0, device:str|None=None) -> Buffer:
+    buf = Buffer(device or self.device, 1, dtypes.uint64, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
+    buf._buf.cpu_view().view(fmt='Q')[0] = init_value
     return buf
 
-  def _wait_signal(self, sig:memoryview, value:int, timeout:int|None=None):
+  def _wait_signal(self, sig:MMIOInterface|memoryview, value:int, timeout:int|None=None):
     timeout = timeout if timeout is not None and self.can_recover else None
     st, done = time.perf_counter(), sig[0]
     while done < value:
@@ -626,8 +631,8 @@ class HCQ2Compiled(Compiled):
   def synchronize(self, timeout:int|None=None):
     if HCQ_RUNTIME_DEV.value != self.device: Device[HCQ_RUNTIME_DEV.value].synchronize()
 
-    sig = self.signal("timeline").as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
-    tl = self.signal("value", 1).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')
+    sig = self.signal("timeline")._buf.cpu_view().view(fmt='Q')
+    tl = self.signal("value", 1, device="CPU")._buf.cpu_view().view(fmt='Q')
     self._wait_signal(sig, tl[0] - 1, timeout)
     if self.prof_ents: self.collect_prof()
 

tinygrad/runtime/support/memory.py

@@ -236,7 +236,7 @@ class MemoryManager:
     self.map_range(va:=self.alloc_vaddr(self.vram_size, self.vram_size), self.vram_size, [(0, self.vram_size)], AddrSpace.PHYS, uncached=uncached)
     return va
 
-  def valloc(self, size:int, align=0x1000, uncached=False, contiguous=False) -> VirtMapping:
+  def valloc(self, size:int, align=0x1000, uncached=False, contiguous=False, zero=False) -> VirtMapping:
     if not getenv("GMMU", 1):
       paddr = self.palloc(size:=round_up(size, 0x1000), align, zero=False)
       return VirtMapping(self.identity_va(uncached) + paddr, size, [(paddr, size)], aspace=AddrSpace.PHYS, uncached=uncached)
@@ -251,7 +251,7 @@ class MemoryManager:
       while rem_size > 0:
         while self.palloc_ranges[nxt_range][0] > rem_size: nxt_range += 1
 
-        try: paddrs += [(self.palloc(try_sz:=self.palloc_ranges[nxt_range][0], self.palloc_ranges[nxt_range][1], zero=False), try_sz)]
+        try: paddrs += [(self.palloc(try_sz:=self.palloc_ranges[nxt_range][0], self.palloc_ranges[nxt_range][1], zero=zero), try_sz)]
         except MemoryError:
           # Move to a smaller size and try again.
           nxt_range += 1

tinygrad/runtime/support/nv/ip.py

@@ -345,7 +345,7 @@ class NV_FLCN_COT(NV_IP):
 
 class NV_GSP(NV_IP):
   def init_sw(self):
-    self.handle_gen = itertools.count(0xcf000000)
+    self.handle_gen, self.chan_runlists = itertools.count(0xcf000000), {}
     self.init_rm_args()
     self.init_libos_args()
     self.init_wpr_meta()
@@ -355,6 +355,7 @@ class NV_GSP(NV_IP):
     self.rpc_set_registry_table()
 
     self.gpfifo_class, self.compute_class, self.dma_class = nv_gpu.AMPERE_CHANNEL_GPFIFO_A, nv_gpu.AMPERE_COMPUTE_B, nv_gpu.AMPERE_DMA_COPY_B
+    self.viddec_class = {"AD":nv_gpu.NVC9B0_VIDEO_DECODER, "GB":nv_gpu.NVCFB0_VIDEO_DECODER}.get(self.nvdev.chip_name[:2]) # nvdec: ada and blackwell
     match self.nvdev.chip_name[:2]:
       case "AD": self.compute_class = nv_gpu.ADA_COMPUTE_A
       case "GB":
@@ -453,8 +454,8 @@ class NV_GSP(NV_IP):
     self.wpr_meta, _, wpr_meta_addrs = self.nvdev._alloc_boot_mem(ctypes.sizeof(type(m)), data=bytes(m))
     self.wpr_meta_sysmem = wpr_meta_addrs[0]
 
-  def promote_ctx(self, client:int, subdevice:int, obj:int, ctxbufs:dict[int, GRBufDesc], bufs=None, virt=None, phys=None):
-    res, prom = {}, nv_gpu.NV2080_CTRL_GPU_PROMOTE_CTX_PARAMS(entryCount=len(ctxbufs), engineType=0x1, hChanClient=client, hObject=obj)
+  def promote_ctx(self, client:int, subdevice:int, obj:int, ctxbufs:dict[int, GRBufDesc], bufs=None, virt=None, phys=None, engine=0x1):
+    res, prom = {}, nv_gpu.NV2080_CTRL_GPU_PROMOTE_CTX_PARAMS(entryCount=len(ctxbufs), engineType=engine, hChanClient=client, hObject=obj)
     for i,(buf,desc) in enumerate(ctxbufs.items()):
       use_v, use_p = (desc.virt if virt is None else virt), (desc.phys if phys is None else phys)
       x = (bufs or {}).get(buf, self.nvdev.mm.valloc(desc.size, contiguous=True)) # allocate buffers
@@ -470,6 +471,9 @@ class NV_GSP(NV_IP):
     subdev = self.rpc_rm_alloc(hParent=dev, hClass=nv_gpu.NV20_SUBDEVICE_0, params=nv_gpu.NV2080_ALLOC_PARAMETERS())
     vaspace = self.rpc_rm_alloc(hParent=dev, hClass=nv_gpu.FERMI_VASPACE_A, params=nv_gpu.NV_VASPACE_ALLOCATION_PARAMETERS())
 
+    di = self.rpc_rm_control(subdev, nv_gpu.NV2080_CTRL_CMD_FIFO_GET_DEVICE_INFO_TABLE, nv_gpu.NV2080_CTRL_FIFO_GET_DEVICE_INFO_TABLE_PARAMS())
+    self.runlists = {di.entries[i].engineData[2]: di.entries[i].engineData[3] for i in range(di.numEntries)}
+
     # reserve 512MB for the reserved PDES
     res_va = self.nvdev.mm.alloc_vaddr(res_sz:=(512 << 20))
 
@@ -549,10 +553,16 @@ class NV_GSP(NV_IP):
     self.cmd_q.send_rpc(nv.NV_VGPU_MSG_FUNCTION_GSP_RM_ALLOC, bytes(alloc_args) + (bytes(params) if params is not None else b''))
     self.stat_q.wait_resp(nv.NV_VGPU_MSG_FUNCTION_GSP_RM_ALLOC)
 
+    if hClass == self.gpfifo_class:
+      self.chan_runlists[obj] = self.runlists.get((e:=params.engineType) + 10*(e >= nv_gpu.NV2080_ENGINE_TYPE_NVDEC0), 0)
     if hClass == nv_gpu.FERMI_VASPACE_A and client != self.priv_root:
       self.rpc_set_page_directory(device=hParent, hVASpace=obj, pdir_paddr=self.nvdev.mm.root_page_table.paddr, client=client)
     if hClass == nv_gpu.NV01_DEVICE_0 and client != self.priv_root: self.device = obj # save user device handle
     if hClass == nv_gpu.NV20_SUBDEVICE_0: self.subdevice = obj # save subdevice handle
+    if hClass == self.viddec_class and client != self.priv_root:
+      ctx, eng = {0: GRBufDesc(0x1000, phys=True, virt=True)}, nv_gpu.NV2080_ENGINE_TYPE_NVDEC0
+      bufs = self.promote_ctx(client, self.subdevice, hParent, ctx, virt=False, engine=eng)
+      self.promote_ctx(client, self.subdevice, hParent, ctx, bufs, phys=False, engine=eng)
     if hClass == self.compute_class and client != self.priv_root:
       phys_gr_ctx = self.promote_ctx(client, self.subdevice, hParent, {k:v for k,v in self.grctx_bufs.items() if k in [0, 1, 2]}, virt=False)
       self.promote_ctx(client, self.subdevice, hParent, {k:v for k,v in self.grctx_bufs.items() if k in [0, 1, 2]}, phys_gr_ctx, phys=False)
@@ -575,9 +585,10 @@ class NV_GSP(NV_IP):
     res = self.stat_q.wait_resp(nv.NV_VGPU_MSG_FUNCTION_GSP_RM_CONTROL)
     st = type(params).from_buffer_copy(res[len(bytes(control_args)):]) if params is not None else None
 
-    # NOTE: gb20x requires the enable bit for token submission. Patch workSubmitToken here to maintain userspace compatibility.
-    if self.nvdev.chip_name.startswith("GB2") and cmd == nv_gpu.NVC36F_CTRL_CMD_GPFIFO_GET_WORK_SUBMIT_TOKEN:
-      cast(nv_gpu.NVC36F_CTRL_CMD_GPFIFO_GET_WORK_SUBMIT_TOKEN_PARAMS, st).workSubmitToken |= (1 << 30)
+    # NOTE: gsp only fills in the channel id, the runlist id (and, on gb20x, the doorbell enable bit) are added by the driver.
+    if cmd == nv_gpu.NVC36F_CTRL_CMD_GPFIFO_GET_WORK_SUBMIT_TOKEN:
+      cast(nv_gpu.NVC36F_CTRL_CMD_GPFIFO_GET_WORK_SUBMIT_TOKEN_PARAMS, st).workSubmitToken |= (self.chan_runlists[hObject] << 16) | \
+        ((1 << 30) if self.nvdev.chip_name.startswith("GB2") else 0)
     return st
 
   def rpc_set_page_directory(self, device:int, hVASpace:int, pdir_paddr:int, client=None, pasid=0xffffffff):

tinygrad/runtime/support/system.py

@@ -262,7 +262,7 @@ class PCIIfaceBase:
     self.dev_impl = dev_impl_t(self.pci_dev)
     self.dev, self.vram_bar, self.count = dev, vram_bar, len(hcq_filter_visible_devices(System.list_devices(vendor, devices, base_class), dn))
 
-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, **kwargs) -> HCQBuffer:
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, zero=False, **kwargs) -> HCQBuffer:
     should_use_sysmem = host or ((cpu_access if self.is_bar_small() else (uncached and cpu_access)) and not force_devmem)
 
     # Align size to huge pages for large allocations, otherwise the unaligned tail falls back to 4KB pages, increasing TLB pressure.
@@ -274,7 +274,7 @@ class PCIIfaceBase:
       mapping = self.dev_impl.mm.map_range(vaddr, size, [(paddr, 0x1000) for paddr in paddrs], aspace=AddrSpace.SYS, snooped=True, uncached=True)
       return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(mapping, has_cpu_mapping=True, hMemory=paddrs[0]), view=memview, owner=self.dev)
 
-    mapping = self.dev_impl.mm.valloc(size:=round_up(size, 0x1000), uncached=uncached, contiguous=cpu_access)
+    mapping = self.dev_impl.mm.valloc(size:=round_up(size, 0x1000), uncached=uncached, contiguous=cpu_access, zero=zero)
     barview = self.pci_dev.map_bar(bar=self.vram_bar, off=mapping.paddrs[0][0], size=mapping.size) if cpu_access else None
     return HCQBuffer(mapping.va_addr, size, view=barview, meta=PCIAllocationMeta(mapping, cpu_access, hMemory=mapping.paddrs[0][0]), owner=self.dev)
 

tinygrad/runtime/support/usb.py

@@ -1,6 +1,11 @@
 import ctypes, struct, time, functools, itertools
+from typing import Any, cast
 from tinygrad.runtime.autogen import libusb
-from tinygrad.helpers import DEBUG, DEV, to_mv, round_up, ceildiv
+from tinygrad.helpers import DEBUG, DEV, to_mv, from_mv, round_up, ceildiv, unwrap, dedup, to_tuple
+from tinygrad.dtype import dtypes
+from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher
+from tinygrad.device import Buffer, BufferSpec, Device
+from tinygrad.runtime.support.hcq2 import HCQInfo, make_buf, make_cmdbuf, make_submit, HCQ_RUNTIME_DEV
 from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support import c
 
@@ -35,6 +40,11 @@ class USB3:
     self._tags, self._transferred = itertools.count(1), ctypes.c_int(0)
     self._bulk_buf, self._bulk_mv = alloc_cbuffer(4 << 20)
     self._ctrl_buf, self._ctrl_mv = alloc_cbuffer(0x1000)
+    # async bulk OUT state: tag -> (pooled transfer, keepalive payload mv); transfer errors latch into _async_err
+    self._async_seq, self._async_err = itertools.count(1), 0
+    self._async_pending: dict = {}
+    self._async_pool: list = []
+    self._async_cb = libusb.libusb_transfer_cb_fn(self._on_bulk_done)
 
     self.handle = c.init_c_var(c.POINTER[libusb.struct_libusb_device_handle], lambda x: checked(libusb.libusb_open)(dev, x))
 
@@ -73,6 +83,42 @@ class USB3:
       (self.handle, 0x02, self._bulk_buf, len(payload), self._transferred, timeout)
     assert self._transferred.value == len(payload), f"bulk OUT short write: {self._transferred.value}/{len(payload)} bytes"
 
+  def _on_bulk_done(self, xfer):  # runs in libusb event handling; latch errors (exceptions here are unraisable)
+    exp = xfer.contents.length - 8 if xfer.contents.type == libusb.LIBUSB_TRANSFER_TYPE_CONTROL else xfer.contents.length
+    if xfer.contents.status != 0 or xfer.contents.actual_length != exp: self._async_err = xfer.contents.status or -1
+    self._async_pool.append(self._async_pending.pop(int(xfer.contents.user_data or 0))[0])
+
+  def _submit_async(self, endpoint:int, xtype:int, payload:bytes|bytearray|memoryview, timeout:int) -> int:  # payload kept alive till bulk_wait
+    tr = self._async_pool.pop() if self._async_pool else libusb.libusb_alloc_transfer(0)
+    tr.contents.dev_handle, tr.contents.endpoint, tr.contents.type = self.handle, endpoint, xtype
+    tr.contents.timeout, tr.contents.length = timeout, len(payload)
+    tr.contents.buffer = ctypes.cast(from_mv(memoryview(payload), ctypes.c_ubyte), ctypes.POINTER(ctypes.c_ubyte))
+    tr.contents.callback, tr.contents.user_data = self._async_cb, (tag := next(self._async_seq))
+    self._async_pending[tag] = (tr, payload)
+    checked(libusb.libusb_submit_transfer, "async submit failed")(tr)
+    return tag
+
+  def bulk_write_async(self, payload:memoryview, timeout:int=10000) -> int:
+    """Queue a bulk OUT transfer without blocking; payload is kept alive until bulk_wait(tag)."""
+    return self._submit_async(0x02, libusb.LIBUSB_TRANSFER_TYPE_BULK, payload, timeout)
+
+  def control_write_async(self, request:int, value:int=0, index:int=0, data:bytes=b"", timeout:int=1000) -> int:
+    """Queue a vendor control OUT without blocking; completes via bulk_wait(tag) like bulk_write_async."""
+    setup = bytearray(struct.pack('<BBHHH', 0x40, request, value, index, len(data)) + data)
+    return self._submit_async(0, libusb.LIBUSB_TRANSFER_TYPE_CONTROL, setup, timeout)
+
+  def control_read_async(self, request:int, length:int, value:int=0, index:int=0, timeout:int=1000) -> tuple[int, memoryview]:
+    """Queue a vendor control IN without blocking; the data lands in the returned buffer by bulk_wait(tag)."""
+    buf = bytearray(struct.pack('<BBHHH', 0xC0, request, value, index, length)) + bytearray(length)
+    return self._submit_async(0, libusb.LIBUSB_TRANSFER_TYPE_CONTROL, buf, timeout), memoryview(buf)[8:]
+
+  def bulk_wait(self, tag:int):
+    """Block until the tagged transfer completes; raises if any async transfer failed. LIBUSB_ERROR_INTERRUPTED is retried."""
+    while tag in self._async_pending:
+      if (rc:=libusb.libusb_handle_events(None)) < 0 and rc != libusb.LIBUSB_ERROR_INTERRUPTED:
+        raise RuntimeError(f"libusb_handle_events: {ctypes.string_at(libusb.libusb_strerror(rc)).decode()}")
+    if self._async_err: raise RuntimeError(f"async bulk OUT failed: status={self._async_err}")
+
   def bulk_read(self, length:int, timeout:int=1000) -> memoryview:
     if length > len(self._bulk_mv): self._bulk_buf, self._bulk_mv = alloc_cbuffer(length)
     checked(libusb.libusb_bulk_transfer, "bulk IN 0x81 failed")(self.handle, 0x81, self._bulk_buf, length, self._transferred, timeout)
@@ -160,13 +206,10 @@ class CustomASM24Controller:
     """Write to chip XDATA via vendor control OUT (bRequest=0xE5). wValue=addr, wIndex=val."""
     for off, val in enumerate(data): self.usb.control_write(0xE5, value=base_addr + off, index=val)
 
-  def scsi_write(self, buf:bytes):
+  def scsi_write(self, buf:bytes, slot_start:int=0):
     """Write to SRAM via 0xF2 vendor command + bulk OUT."""
     buf_padded = buf + b'\x00' * (round_up(len(buf), 512) - len(buf))
-    sectors = len(buf_padded) // 512
-    num_slots = ceildiv(len(buf_padded), 0x4000)  # 16KB per slot
-    windex = (num_slots & 0xFF) << 8
-    self.usb.control_write(0xF2, value=sectors, index=windex)
+    self.usb.control_write(0xF2, value=len(buf_padded) // 512, index=(slot_start & 0xFF) | (ceildiv(len(buf_padded), 0x4000) << 8))
     self.usb.bulk_write(buf_padded)
 
   def scsi_read_arm(self, size:int):
@@ -184,20 +227,123 @@ class USBMMIOInterface(MMIOInterface):
     return (index * self.el_sz, self.el_sz)
 
   def __getitem__(self, index):
+    Device[HCQ_RUNTIME_DEV.value].synchronize() # one driver on the link: drain the compiled submits before python touches it
     off, sz = self._off_from_index(index)
     if self.pcimem:
       assert sz % 4 == 0 and off % 4 == 0, f"pcie_mem_read requires 4-byte aligned access, got off={off}, sz={sz}"
       data = self.usb.pcie_mem_read(self.addr + off, sz)
     else: data = self.usb.scsi_read(sz) if self.addr == 0xf000 else self.usb.read(self.addr + off, sz)
-    return int.from_bytes(data, "little") if sz == self.el_sz else data
+    return data if isinstance(index, slice) else int.from_bytes(data, "little")
 
   def __setitem__(self, index, data):
+    Device[HCQ_RUNTIME_DEV.value].synchronize()
     off, _ = self._off_from_index(index)
     data = struct.pack(self.fmt, data) if isinstance(data, int) else bytes(data)
     if not self.pcimem: self.usb.scsi_write(data) if self.addr == 0xf000 else self.usb.write(self.addr + off, data)
-    else: self.usb.pcie_mem_write(self.addr+off, data)
+    else:
+      # writes are whole dwords
+      assert len(data) % 4 == 0 and off % 4 == 0, f"pcie_mem_write requires 4-byte aligned access, got off={off}, sz={len(data)}"
+      self.usb.pcie_mem_write(self.addr+off, data)
 
   def view(self, offset:int=0, size:int|None=None, fmt=None):
     return USBMMIOInterface(self.usb, self.addr+offset, self.nbytes-offset if size is None else size, fmt=fmt or self.fmt, pcimem=self.pcimem)
 
+# *****************
+
+def _libusb(devs, dep:tuple[UOp, ...], fn:str, *args) -> UOp:
+  return make_buf(devs, tag=f"func:{fn}").after(*dep).index(0).load().call(make_buf(devs, tag="usb_handle").index(0).load(),
+    *[UOp.const(a, dtypes.int) if isinstance(a, int) else a for a in args], ret_dtype=dtypes.void)
+
+def usb_bulk(devs, dep, endpoint:int, data:UOp, length, timeout:int=1000) -> UOp: # NULL actual_length out param
+  return _libusb(devs, dep, "libusb_bulk_transfer", endpoint, data, length, UOp.const(0, dtypes.uint64), timeout)
+
+def usb_stream(devs, dep:tuple[UOp, ...], addr:UOp, data:UOp, nbytes:int, write:bool) -> UOp:
+  hdr = UOp.placeholder((2,), dtypes.uint64, device=devs, tag="usb_scratch").after(*dep)
+  arm = _libusb(devs, (hdr.index(0).store(addr), hdr.index(1).store(UOp.const(nbytes // 4, dtypes.uint64))), "libusb_control_transfer",
+                0x40, 0xF0, (0x60 if write else 0x20) | (0x0F << 8), 1 if write else 2, hdr.index(0), 12, 5000)
+  return usb_bulk(devs, (arm,), 0x02 if write else 0x81, data, nbytes)
+
+def usb_writes(devs, ws:list[tuple[UOp, UOp, int]]) -> tuple[UOp, ...]:
+  return functools.reduce(lambda dep, w: (usb_stream(devs, dep, w[0], w[1], w[2], True),), ws, ())
+
+def usb_load(b:UOp, idx:UOp, dt) -> UOp:
+  got = UOp.placeholder((1,), dt, device=(devs:=to_tuple(b.device)), tag="usb_scratch")
+  addr = b.getaddr((HCQ_RUNTIME_DEV.value,)) + (idx*dt.itemsize).cast(dtypes.uint64)
+  return got.after(usb_stream(devs, b.src[1:] if b.op is Ops.AFTER else (), addr, got.index(0), dt.itemsize, False)).index(0).load()
+
+def usb_write(b:UOp, idx:UOp, v:UOp) -> UOp:
+  val = (s:=UOp.placeholder((1,), v.dtype, device=(devs:=to_tuple(b.device)), tag="usb_scratch")).after(s.index(0).store(v))
+  addr = b.getaddr((HCQ_RUNTIME_DEV.value,)) + (idx*v.dtype.itemsize).cast(dtypes.uint64)
+  return usb_stream(devs, b.src[1:] if b.op is Ops.AFTER else (), addr, val.index(0), v.dtype.itemsize, True)
+
+def usb_idle(devs) -> UOp:
+  v = usb_load(make_buf(devs, tag="timeline_signal").after(loop:=UOp.loop(0)), UOp.const(0, dtypes.int), dtypes.uint64)
+  return v.end(loop, v + 1 < make_buf(devs, tag="timeline_value").index(0).load())
+
+def usb_scsi(devs, read:bool, nbytes:int) -> UOp:
+  return _libusb(devs, (usb_idle(devs),), "libusb_control_transfer", 0x40, 0xF2, ceildiv(nbytes, 512) | (0x8000 if read else 0),
+                 (ceildiv(nbytes, 0x4000) & 0xFF) << 8, UOp.const(0, dtypes.uint64), 0, 1000)
+
+def usb_stage_copy(dst:UOp, src:UOp) -> UOp|None:
+  if (cin:=to_tuple(src.device)[0].startswith("CPU")) == to_tuple(dst.device)[0].startswith("CPU"): return None
+
+  total, ops, win = dst.nbytes(), [], cast(Any, Device[(devs:=to_tuple((dst if cin else src).device))[0]]).iface.usb_sram
+  for off in range(0, total, win.size): # off and nb are bytes, the two ends of the copy can have different dtypes
+    sram = UOp.from_buffer(win)[0:(nb:=min(win.size, total - off))]
+    s, d = src[off // src.dtype.itemsize:(off + nb) // src.dtype.itemsize], dst[off // dst.dtype.itemsize:(off + nb) // dst.dtype.itemsize]
+    if cin:
+      push = usb_bulk(devs, (usb_scsi(devs, False, nb),), 0x02, s.getaddr((HCQ_RUNTIME_DEV.value,)), round_up(nb, 512), 10000)
+      ops += [UOp.custom_function("hcq", push.sink()).call(sram, s, name="hcq_copyin", aux=HCQInfo(devs)),
+              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")
+      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)]
+  return UOp(Ops.LINEAR, src=tuple(ops))
+pm_usb_stage = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), usb_stage_copy)])
+
+def usb_arm_bytes(lin:UOp, sram:Buffer) -> int:
+  dsts = [c.src[1] for c in lin.src if c.op is Ops.CALL and c.src[0].op is Ops.COPY] # the rest of a linear is INS, some with no srcs
+  return next((d.nbytes() for d in dsts if d.base.op is Ops.BUFFER and d.base.buffer is sram), 0)
+
+def usb_ib(devs, lin:UOp, align:int, arm:int=0) -> tuple[UOp, UOp, int]:
+  pkt_dw = sum(s.dtype.itemsize for ins in lin.src for s in ins.src) // 4 # by bytes: sdma packs 64-bit addresses as single srcs
+  kargs = dedup([b for b in lin.toposort() if b.op is Ops.PARAM and b.tag == "kernargs"])
+  offs, up_dw = {}, round_up(pkt_dw, align)
+  for k in kargs: offs[k], up_dw = up_dw, round_up(up_dw + k.max_numel(), 32)
+  ib_gpu = UOp.placeholder((up_dw,), dtypes.uint32, device=devs, tag="cmdbuf")
+  ib_host = UOp.placeholder((up_dw,), dtypes.uint32, device=devs, tag="usb_scratch")
+  gsubs = {g: g.replace(src=(d if a.op is not Ops.AFTER else d.after(*a.src[1:]),)) for g in lin.toposort() if g.op is Ops.GETADDR
+           for a in [g.src[0]] if (k:=a.src[0] if a.op is Ops.AFTER else a) in offs for d in [ib_gpu[offs[k]:offs[k] + k.max_numel()]]}
+  lin = lin.substitute(gsubs, walk=True).substitute({k: ib_host[offs[k]:offs[k] + k.max_numel()] for k in kargs}, walk=True)
+  return make_cmdbuf(lin, devs, buf=ib_host, dep=(usb_scsi(devs, True, arm),) if arm else ()), ib_gpu, pkt_dw
+
+def usb_push(devs, ring:UOp, wptr:UOp, doorbell:UOp, put_ptr:UOp, ib_host:UOp, ib_gpu:UOp, pkt:tuple, unit:int) -> UOp:
+  stage = UOp.placeholder(((n:=round_up(len(pkt), 4)) + 2,), dtypes.uint32, device=devs, tag="usb_scratch")
+  put, step = put_ptr.index(zero:=UOp.const(0, dtypes.int)), (n * 4 if pkt else ib_host.nbytes()) // unit
+  st = stage.after(*[stage.index(i).store(UOp.const(v, dtypes.uint32)) for i, v in enumerate(pkt)],
+                   *[stage.index(n + i).store((((put + step) >> (32 * i)) & 0xffffffff).cast(dtypes.uint32)) for i in (0, 1)])
+
+  writes = [(ib_gpu.getaddr((HCQ_RUNTIME_DEV.value,)), ib_host.index(zero), ib_gpu.nbytes())] if pkt else []
+  writes += [(ring.getaddr((HCQ_RUNTIME_DEV.value,)) + ((put % (ring.nbytes() // unit)) * unit).cast(dtypes.uint64),
+              (st if pkt else ib_host).index(zero), step * unit)]
+  writes += [(p.getaddr((HCQ_RUNTIME_DEV.value,)), st.index(n), 8) for p in (wptr, doorbell)]
+  return put_ptr.after(*usb_writes(devs, writes)).index(zero).store(put + step)
+
+USB_HOST_TAGS = {"signal", "timeline_signal"}
+pm_usb_hostio = PatternMatcher([
+  (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(Ops.PARAM, tag=USB_HOST_TAGS).or_after(name="b"), UPat(name="idx"))),),
+        name="ld"), lambda b, idx, ld: usb_load(b, idx, ld.dtype)),
+  (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat(Ops.PARAM, tag=USB_HOST_TAGS).or_after(name="b"), UPat(name="idx"))), UPat(name="v"))), usb_write)])
+
+pm_usb_bufferize = PatternMatcher([
+  (UPat(Ops.PARAM, tag={"systems", "runtime", "inputs", "usb_scratch"}, name="b"),
+   lambda ctx, b: Buffer("CPU", b.max_numel(), b.dtype, options=BufferSpec(nolru=True), preallocate=True)),
+  (UPat(Ops.PARAM, tag="usb_handle", name="b"), lambda ctx, b: ctx[0].signal(b.tag, ctx[0].iface.usb_handle, device="CPU")),
+  (UPat(Ops.PARAM, name="b"), lambda ctx, b: None if not isinstance(b.tag, str) or not b.tag.startswith("func:") else
+   ctx[0].signal(b.tag, unwrap(ctypes.cast(getattr(libusb.dll, b.tag[5:]), ctypes.c_void_p).value), device="CPU")),
+])
+
 if DEV.interface.startswith("MOCK"): from test.mockgpu.usb import MockUSB3 as USB3  # type: ignore  # noqa: F811

tinygrad/schedule/__init__.py

@@ -97,11 +97,17 @@ pm_post_sched_cache = PatternMatcher([
    create_new_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None),
 ])
 
-def resolve_linear_call(linear_call:UOp):
+def resolve_linear_call(linear_call:UOp, outer_binds:dict[str, UOp]|None=None):
   linear = graph_rewrite(linear_call.src[0], pm_post_sched_cache, ctx=({}, linear_call.src[1:]), walk=True, name="params to buffers")
-  # map the call body params back to the original Variables stored in the call args
-  binds = {f"p{i}":x.src[0].replace(op=Ops.PARAM) for i,x in enumerate(linear_call.src[1:]) if x.is_bound_var}
-  return linear.substitute({v:binds[v.expr] for v in linear.variables() if v.expr in binds}, enter_calls=True, name="resolve scalar params")
+  # nested LINEAR calls are lexical scopes: their positional params shadow the enclosing scope, while calls without
+  # scalar args (e.g. precompiled allreduce) inherit it
+  binds = {**(outer_binds or {}),
+           **{f"p{i}":x.src[0].replace(op=Ops.PARAM) for i,x in enumerate(linear_call.src[1:]) if x.is_bound_var}}
+  def apply_binds(si:UOp) -> UOp:
+    if si.op is Ops.CALL and si.src[0].op is Ops.LINEAR: return resolve_linear_call(si, binds)
+    subs = {v:binds[v.expr] for v in si.variables() if v.expr in binds}
+    return si.replace(src=tuple(s.substitute(subs, name="resolve scalar params") for s in si.src))
+  return linear.replace(src=tuple(apply_binds(si) for si in linear.src))
 
 pm_resolve_linear_call = PatternMatcher([
   # call LINEAR is resolved here

tinygrad/schedule/indexing.py

@@ -12,8 +12,6 @@ class IndexingContext:
   realize_map: dict[UOp, None|list[int]] = field(default_factory=dict)
   non_removable: dict[UOp, None] = field(default_factory=dict)
   range_map: dict[UOp, tuple[tuple[UOp, ...], tuple[UOp, ...]]] = field(default_factory=dict)
-  # loads reachable from each UOp memoized across matches
-  buf_cache: dict[UOp, frozenset[UOp]] = field(default_factory=dict)
 
   # create ranges
   range_idx: Iterator[int] = field(default_factory=itertools.count)
@@ -187,7 +185,7 @@ def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UO
   return rngs
 
 @rewrite_group(new_ctx=False)
-def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
+def run_rangeify(tsink:UOp, debug:bool=False) -> UOp:
   if debug: print("**************************")
   rctx = IndexingContext()
 
@@ -322,7 +320,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     tsink = graph_rewrite(tsink, pm_apply_rangeify, ctx=rctx, bottom_up=True, name="apply rangeify")
   # if a deviceless value must materialize, place it on the sink device
   tsink = graph_rewrite(tsink, pm_fix_deviceless, ctx=tsink.device, name="add device to deviceless")
-  return tsink, rctx
+  return tsink
 
 def render_ranges(*rngs_list, realized) -> str:
   disp = []

tinygrad/schedule/rangeify.py

@@ -10,7 +10,7 @@ from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP
 from tinygrad.helpers import PCONTIG, FLOAT16, OPENPILOT_HACKS, argsort, 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, IndexingContext, apply_movement_op
+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
 
@@ -79,7 +79,7 @@ def split_reduceop(reduce:UOp, x:UOp):
 
   # 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, x.base.dtype)}, extra_pm=pm_mops).ranges]
+  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])
@@ -304,15 +304,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),
@@ -352,7 +352,12 @@ pm_no_indexing_calls = PatternMatcher([
 ])
 
 DEVICE_MAX_BUFS = {"METAL": 31, "WEBGPU": 8, "CPU": 31} # TODO: get from device?
-def limit_bufs(ctx:IndexingContext, root:UOp):
+@dataclass
+class LimitBufsContext:
+  buf_cache: dict[UOp, frozenset[UOp]] = field(default_factory=dict)
+  range_idx: itertools.count = field(default_factory=itertools.count)
+
+def _limit_bufs(ctx:LimitBufsContext, root:UOp):
   if (device:=root.device) is None: return None # no device, index related calculations
   device = device if isinstance(device, str) else device[0].split(":")[0]
   if not (MAX_BUFS:=MAX_KERNEL_BUFFERS.value or DEVICE_MAX_BUFS.get(device, 0)): return None
@@ -374,7 +379,7 @@ def limit_bufs(ctx:IndexingContext, root:UOp):
         s = s.substitute(dict(zip(orig_ranges, end_ranges))).bufferize(*end_ranges, arg=BufferizeOpts(device=s.device)).index(*orig_ranges)
       srcs.append(s)
     return root.replace(src=tuple(srcs))
-pm_limit_bufs = PatternMatcher([(UPat(set.union(GroupOp.Binary, GroupOp.Ternary), name="root"), limit_bufs)])
+pm_limit_bufs = PatternMatcher([(UPat(set.union(GroupOp.Binary, GroupOp.Ternary), name="root"), _limit_bufs)])
 
 # *****************
 # 4. put in buffers for bufferize
@@ -578,20 +583,21 @@ pm_copy_to_store = PatternMatcher([
 
 @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")
 
   # convert movement ops to ranges
-  tsink, rctx = run_rangeify(tsink, bool(DEBUG_RANGEIFY))
+  tsink = run_rangeify(tsink, bool(DEBUG_RANGEIFY))
 
+  # cleanups for speed and runability
   tsink = graph_rewrite(tsink,
                         symbolic+pm_reduce_simplify+pm_const_buffer_folding+pm_remove_bufferize,
                         name="symbolic+reduce_collapse+debuf")
-  tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")
-
+  next_range = max((x.arg[0] for x in tsink.toposort() if x.op is Ops.RANGE), default=-1) + 1
+  tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=LimitBufsContext(range_idx=itertools.count(next_range)), name="limit buffers")
   if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Rangeify")
 
   # bufferize -> store

tinygrad/tensor.py

@@ -744,7 +744,7 @@ class Tensor(RandMixin):
     ref_frames = [x.contiguous() for x in ref_frames or []]
     assert frame_pos.is_bound_var, "frame_pos must be a bound Variable"
     srcs = (out:=Tensor.empty(*shape, device=self.device, dtype=self.dtype), self.contiguous(), state.contiguous(), *ref_frames)
-    fn = UOp(Ops.CUSTOM_FUNCTION, src=(frame_pos.src[0], *[UOp.const(s, dtypes.int) for s in shape]), arg="encdec")
+    fn = UOp(Ops.CUSTOM_FUNCTION, src=(frame_pos.src[0], *[UOp.const(s) for s in shape]), arg="encdec")
     return Tensor(out.uop.after(fn.call(*[s.uop for s in srcs], frame_pos)))
 
 P = ParamSpec("P")

tinygrad/uop/ops.py

@@ -5,7 +5,7 @@ from dataclasses import dataclass, replace
 from enum import Enum, auto
 from tinygrad.uop import Ops, GroupOp
 from tinygrad.dtype import ConstType, dtypes, DType, DTypeLike, truncate, least_upper_dtype, least_upper_float, Invalid, AddrSpace, strong_dtype
-from tinygrad.dtype import PyConst, InvalidType, weak_dtype, bitcast
+from tinygrad.dtype import PyConst, InvalidType, bitcast
 from tinygrad.device import Buffer, MultiBuffer, canonicalize_device, TinyELF
 from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA
 from tinygrad.helpers import PROFILE, dedup, cdiv, cmod, floordiv, floormod, diskcache_put, to_function_name, cpu_profile, TracingKey
@@ -188,12 +188,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
-    # an INDEX presents its access dtype, which a still-weak source matches up to weakness
-    if SPEC == 2 and op is not Ops.CONST and \
-       (expected_dtype:=dtype_from_uop(op, src, arg)) is not None and expected_dtype != dtype and \
-       not (op is Ops.INDEX and weak_dtype(expected_dtype) == weak_dtype(dtype)):
+    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))
@@ -259,8 +255,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
@@ -614,7 +612,8 @@ 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)
@@ -990,7 +989,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))
@@ -1138,14 +1138,16 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   # *** uop high level syntactic sugar ***
 
   @staticmethod
-  def placeholder(shape:tuple[int, ...], dtype:DType, slot:int, addrspace=AddrSpace.GLOBAL, device=None, volatile=False):
+  def placeholder(shape:tuple[int, ...], dtype:DType, slot:int|None=None, addrspace=AddrSpace.GLOBAL, device=None, volatile=False, tag=None):
     dtype = strong_dtype(dtype)  # storage is never weak: a placeholder commits the width of what's put in it
+    if slot is None: slot = next(UOp.unique_num)
     if addrspace is AddrSpace.GLOBAL:
       ret = UOp(Ops.PARAM, src=(shape_to_shape_arg((prod(shape),)),), arg=ParamArg(slot, dtype, addrspace=addrspace, device=device,volatile=volatile))
     else:
       assert addrspace in (AddrSpace.LOCAL, AddrSpace.REG)
       assert device is None, "LOCAL and REG placeholders cannot have a device"
       ret = UOp(Ops.BUFFER, src=(shape_to_shape_arg((prod(shape),)),), arg=ParamArg(slot, dtype, addrspace=addrspace))
+    if tag is not None: ret = ret.rtag(tag)
     if len(shape) > 1: ret = ret.reshape(shape)
     return ret
   def placeholder_like(self, slot:int, addrspace=AddrSpace.GLOBAL):
@@ -1198,7 +1200,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     assert self.op is Ops.PROGRAM and isinstance(self.arg, ProgramInfo), "to_elf should only be called on a PROGRAM ast"
     sig = tuple((u.arg.name, u.arg.slot, u.dtype, u._shape)
                 for u in tuple(filter(lambda u: u.op is Ops.PARAM and u.addrspace != AddrSpace.ALU, self.src[1].src)) + self.arg.vars)
-    return TinyELF(self.src[3].arg, self.arg.function_name, self.arg.target, sig)
+    return TinyELF(self.src[3].arg, self.arg.function_name, self.arg.target, sig, self.key)
 
 @dataclass(frozen=True)
 class KernelInfo:

tinygrad/uop/render.py

@@ -1,6 +1,6 @@
 from tinygrad.dtype import AddrSpace, dtypes
 from tinygrad.uop import Ops, GroupOp
-from tinygrad.uop.ops import ParamArg, UOp, PatternMatcher, UPat, multirange_str, range_str, consumer_map_from_toposort
+from tinygrad.uop.ops import ParamArg, UOp, PatternMatcher, UPat, multirange_str, range_str, consumer_map_from_toposort, sint
 from tinygrad.helpers import strip_parens
 
 def pretty_print(x:UOp, cache=None, d=0)->str:
@@ -69,14 +69,15 @@ renderer_infer = PatternMatcher([
 # *** pyrender ***
 
 def srcs(ctx, src): return f"({ctx[src[0]]},)" if len(src) == 1 else f"({', '.join([ctx[x] for x in src])})"
+# marg is ssimplify'd, so a bound can be a node this graph never contained
+def marg_str(ctx, a:sint) -> str: return str(a) if not isinstance(a, UOp) else ctx[a] if a in ctx else a.render()
+
 def render_marg(ctx,x:UOp):
   if x.op is Ops.PERMUTE: return str(x.marg)
   if x.op is Ops.FLIP: return str(tuple([i for i,x in enumerate(x.marg) if x]))
   pieces = []
-  if x.op in {Ops.RESHAPE, Ops.EXPAND}:
-    pieces = [f"{ctx[a] if isinstance(a, UOp) else str(a)}" for a in x.marg]
-  if x.op in {Ops.PAD, Ops.SHRINK}:
-    pieces = [f"({ctx[a[0]] if isinstance(a[0], UOp) else str(a[0])}, {ctx[a[1]] if isinstance(a[1], UOp) else str(a[1])})" for a in x.marg]
+  if x.op in {Ops.RESHAPE, Ops.EXPAND}: pieces = [marg_str(ctx, a) for a in x.marg]
+  if x.op in {Ops.PAD, Ops.SHRINK}: pieces = [f"({marg_str(ctx, a[0])}, {marg_str(ctx, a[1])})" for a in x.marg]
   return f"({','.join(pieces)})" if len(pieces) != 1 else f"({pieces[0]},)"
 
 sugar = {Ops.SINK, Ops.END, Ops.STORE, Ops.LOAD, Ops.SQRT, Ops.INDEX, Ops.REDUCE, Ops.AFTER, Ops.THREEFRY,

tinygrad/uop/spec.py

@@ -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),

tinygrad/uop/symbolic.py

@@ -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, STACK, and casted CONST
-const_folding_pat = UPat.any(UPat((Ops.CONST, Ops.STACK)), UPat(Ops.CAST, src=(UPat(Ops.CONST),)))
-
-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),
@@ -108,9 +105,9 @@ def fold_const_where(gate:UOp, c0:UOp, c1:UOp, w:UOp) -> UOp:
 
 symbolic_simple = pm_data_invalid + PatternMatcher([
   # ** self folding **
-  (UPat.var("x") + 0, lambda x: x),    # x+0 -> x
+  (UPat({Ops.ADD, Ops.XOR, Ops.OR}, src=[UPat.var("x"), UPat.const(0)]), lambda x: x),  # x+0 / x^0 / x|0 -> x
+  (UPat({Ops.SHL, Ops.SHR}, src=(UPat.var("x"), UPat.const(0))), lambda x: x),          # x<<0 / x>>0 -> x
   (UPat.var("x") * 1, lambda x: x),    # x*1 -> x
-  (UPat.var("x", dtype=dtypes.ints+(dtypes.bool, dtypes.weakint)) ^ 0, lambda x: x), # x^0 -> x
   (UPat.var("x") // UPat.var("x"), lambda x: x.const_like(1)), # x//x -> 1
   (UPat.var("x") // 1, lambda x: x),   # x//1 -> x
   (UPat.var("x") // -1, lambda x: -x), # x//-1 -> -x
@@ -142,10 +139,16 @@ 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 **
-  (UPat(GroupOp.Unary, src=(const_folding_pat,), name="a"), fold_const_alu),
+  # 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)),
+  # 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),
@@ -163,7 +166,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
@@ -289,8 +294,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,
@@ -303,7 +309,12 @@ 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
+  # a CAST over a committed const is a value conversion: evaluate at the committed dtype, keep the outer cast.
+  # keep out of symbolic_simple: create_non_native_float_pats re-expands the pair and cycles on bf16
+  (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)),
+  # 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 ********
 

tinygrad/uop/weak.py

@@ -1,82 +1,104 @@
 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_weak_srcs(u:UOp) -> UOp|None:
-  if not any(s.dtype in dtypes.weaks for s in u.src): return None
-  if (dt:=least_upper_dtype(*(s.dtype for s in u.src))) in dtypes.weaks: return 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))
+# 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))
 
-# 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
-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:]))),
-])
+# 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
+  dts = derived_dtypes(u, u.src)
+  ret = u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks and
+                                        not (s.op is Ops.CONST and dts is not None) 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
+  return commit_srcs_at(u, dt)
 
 # 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
-  dt = least_upper_dtype(c.dtype, default_dtype(u))
-  return u.replace(dtype=None, src=tuple(commit_weak(s, dt) if s.dtype in dtypes.weaks else s for s in u.src)).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([
+# 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),
+  (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),
-  (UPat(Ops.CAST, name="c", src=(UPat(Ops.CONST, dtype=dtypes.weaks, name="u"),)), lambda c,u: commit_weak(u, c.dtype)),
 ])
 
+# 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:
-  start, src = (1 if u.op is Ops.WHERE else 0), tuple(s.src[0] if s.op is Ops.CAST and s.dtype in dtypes.weaks else s for s in u.src)
-  if src == u.src or any(s.dtype in dtypes.weaks for s in src[start:]): return None
-  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)
+  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
+  # 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)),
-  # 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),
-  # Binary can widen from the bounds, all other nodes derive from the lowered sources.
-  # a weakfloat Unary (sin/exp2/...) must resolve here, before the transcendental decomposition
-  (UPat(GroupOp.Binary|GroupOp.Unary|{Ops.WHERE, Ops.RANGE, Ops.STACK, Ops.SPECIAL}, name="u"), lower_weak_node),
-  (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),
-])
-
-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
-
-pm_lower_index_dtype = pm_commit_weak+pm_cast_weak+PatternMatcher([
-  # a CAST between two concrete dtypes over a CONST is a value conversion: evaluate it once, at the width the CAST states
-  # TODO: delete this once CONST has no dtype
-  (UPat(Ops.CAST, dtypes.all, name="root", src=(UPat.cvar("c", dtypes.all),)), lambda root, c: root.const_like(c.val)),
-  (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
+  # 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
+  (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(GroupOp.All, name="u"), lower_weak_node),
 ])
+
+# 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_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)])

tinygrad/viz/serve.py

@@ -231,10 +231,11 @@ def timeline_layout(data:VizData, dev_events:list[tuple[int, int, float, DevEven
   ei:ProfilePointEvent|None = None
   for st,et,dur,e in dev_events:
     if isinstance(e, ProfilePointEvent) and e.name == "exec": ei = e
-    if dur == 0: continue
+    # only visualize range events with an end timestamp
+    if dur == 0 or isinstance(e, ProfilePointEvent): continue
     name, key = e.name, None
     fmt:dict = {}
-    if (ref:=data.ref_map.get(name)) is not None and ref < len(data.ctxs):
+    if (ref:=data.ref_map.get(e.profile_key)) is not None and ref < len(data.ctxs):
       name = data.ctxs[ref]["name"]
       if (ki:=data.ctxs[ref].get("ki")) is not None and ki.estimates is not None and ei is not None:
         for est_key,est_val in (("FLOPS", ki.estimates.ops), ("B/s mem", ki.estimates.mem), ("B/s lds", ki.estimates.lds)):
@@ -333,14 +334,14 @@ def unpack_pmc(e) -> dict:
 
 def load_amd_counters(data:VizData, profile:list) -> None:
   counter_events:dict[tuple[int, int], dict] = {}
-  durations:dict[str, list[float]] = {}
+  durations:dict[bytes|str, list[float]] = {}
   prg_events:dict[int, ProfileProgramEvent] = {}
   arch = ""
   for e in profile:
     if type(e).__name__ in {"ProfilePMCEvent", "ProfileSQTTEvent"}:
       counter_events.setdefault((e.kern, e.exec_tag), {}).setdefault(type(e).__name__, []).append(e)
-    if isinstance(e, ProfileRangeEvent) and e.device.startswith("AMD") and e.en is not None:
-      durations.setdefault(str(e.name), []).append(float(e.en-e.st))
+    if isinstance(e, ProfileRangeEvent) and e.device.startswith("AMD") and e.en is not None and e.profile_key is not None:
+      durations.setdefault(e.profile_key, []).append(float(e.en-e.st))
     if isinstance(e, ProfileProgramEvent) and e.device.startswith("AMD") and e.tag is not None: prg_events[e.tag] = e
     if isinstance(e, ProfileDeviceEvent) and e.device.startswith("AMD"): arch = f"gfx{unwrap(e.props)['gfx_target_version']//1000}"
   if len(counter_events) == 0: return None
@@ -348,12 +349,12 @@ def load_amd_counters(data:VizData, profile:list) -> None:
   run_number = {n:0 for n,_ in counter_events}
   for (k, tag),v in counter_events.items():
     # use the colored name if it exists
-    name = data.ctxs[r]["ki"].name if (r:=data.ref_map.get(pname:=prg_events[k].name)) is not None else pname
+    name = data.ctxs[r]["ki"].name if (r:=data.ref_map.get(unwrap(prg_events[k].profile_key))) is not None else prg_events[k].name
     run_number[k] += 1
     steps:list[dict] = []
     if (pmc:=v.get("ProfilePMCEvent")):
       steps.append(create_step("PMC", ("/prg-pmc", len(data.ctxs), len(steps)), pmc[0]))
-      all_counters[(name, run_number[k], pname)] = pmc[0]
+      all_counters[(name, run_number[k], unwrap(prg_events[k].profile_key))] = pmc[0]
     # to decode a SQTT trace, we need the raw stream, program binary and device properties
     if (sqtt:=v.get("ProfileSQTTEvent")):
       for e in sqtt:
@@ -496,10 +497,10 @@ def get_profile(data:VizData, profile:list[ProfileEvent], sort_fn:Callable[[str]
 def load_nv_counters(data:VizData, profile:list) -> None:
   steps:list[dict] = []
   sm_version = {e.device:e.props.get("sm_version", 0x800) for e in profile if isinstance(e, ProfileDeviceEvent) and e.props is not None}
-  run_number:dict[str, int] = {}
+  run_number:dict[bytes, int] = {}
   for e in profile:
     if type(e).__name__ == "ProfilePMAEvent":
-      run_number[e.kern] = run_num = run_number.get(e.kern, 0)+1
+      run_number[profile_key] = run_num = run_number.get(profile_key:=unwrap(e.profile_key), 0)+1
       steps.append(create_step(f"PMA {e.kern}"+(f"n{run_num}" if run_num>1 else ""), ("/prg-pma-pkts", len(data.ctxs), len(steps)),
                                data=(e.blob, sm_version[e.device])))
   if steps: data.ctxs.append({"name":"All Counters", "steps":steps})