symbolic llm with prefill (ai slop)

* add better printing to tinygrad.apps.llm

* add gc.collect

* comment

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

@@ -233,7 +233,7 @@ runs:
       shell: bash
       run: |
         sudo mkdir -p /usr/local/lib
-        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/nimlgen/amdcomgr_dylib/releases/latest | \
+        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/tinygrad/amdcomgr_dylib/releases/latest | \
           jq -r '.assets[] | select(.name == "libamd_comgr.dylib").browser_download_url' | \
           sudo xargs curl -fL -o /usr/local/lib/libamd_comgr.dylib
         cargo build --release --manifest-path ./extra/remu/Cargo.toml

.github/workflows/benchmark.yml

@@ -520,8 +520,9 @@ jobs:
       run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
     - name: Run full CIFAR training steps w 6 GPUS
       run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
-    - name: Test full tinyfs load
-      run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
+    # TODO: broken on some of the machines
+    #- name: Test full tinyfs load
+    #  run: TINYFS_ENDPOINT=10.0.52.11:6767 PYTHONPATH=. python extra/tinyfs/fetch_file.py --hash d734f5e3be9f1e9d863bfaa4fc6c1ef2 --len 175866113 --dest mapping.json --check
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 

examples/mlperf/dataloader.py

@@ -396,6 +396,7 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
       queue_in.put((idx, img, tgt))
 
   def _setup_shared_mem(shm_name:str, size:tuple[int, ...], dtype:dtypes) -> tuple[shared_memory.SharedMemory, Tensor]:
+    shm_name = f"{shm_name}_{os.getpid()}"
     if os.path.exists(f"/dev/shm/{shm_name}"): os.unlink(f"/dev/shm/{shm_name}")
     shm = shared_memory.SharedMemory(name=shm_name, create=True, size=prod(size))
     shm_tensor = Tensor.empty(*size, dtype=dtype, device=f"disk:/dev/shm/{shm_name}")

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

@@ -5,6 +5,7 @@ export DEV=${DEV:-AMD}
 export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
+export DEVICE_IN_FUNCTION_BUG=1
 
 export DEBUG=${DEBUG:-0}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}

extra/sqtt/examples/discover_ops.py

@@ -0,0 +1,148 @@
+#!/usr/bin/env python3
+# Run all ALU and memory instructions in the ISA
+import functools, inspect
+from enum import Enum
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AddrSpace
+from tinygrad.renderer.amd.dsl import Inst, Reg, OPERANDS, SrcField, VGPRField, SGPRField, SSrcField, SBaseField, AlignedSGPRField, BitField
+from tinygrad.renderer.amd.dsl import FixedBitField, EnumBitField, s, v, NULL, VCC_LO
+from extra.gemm.amd_asm_matmul import Kernel
+
+# skip instructions that mutate wave state (PC, EXEC, allocations, signals)
+SKIP = {"S_SETPC_B64", "S_SWAPPC_B64", "S_RFE_B64", "S_BARRIER_SIGNAL_ISFIRST", "S_GET_BARRIER_STATE", "S_ALLOC_VGPR", "S_SLEEP_VAR", "S_GETPC_B64",
+        "S_SENDMSG_RTN_B32", "S_SENDMSG_RTN_B64"}
+# skip barriers, s_waits, wrap level atomics, and ray tracing (bvh)
+SKIP_SUBSTR = ["SAVEEXEC", "CMPX", "WREXEC", "MOVREL", "ATOMIC", "S_BUFFER_", "S_ATC_PROBE", "BARRIER", "S_WAITCNT", "BVH",
+               "DS_CMPSTORE_RTN", "DS_WRAP_RTN_B32", "DS_ORDERED_COUNT", "DS_GWS", "GS_REG", "GLOBAL_LOAD_LDS", "GLOBAL_STORE_BLOCK"]
+
+ALU_FORMATS = {"VOP1", "VOP1_LIT", "VOP1_SDST", "VOP2", "VOP2_LIT", "VOP3", "VOP3_SDST", "VOP3SD", "VOP3P", "VOP3P_MFMA", "VOP3PX2",
+               "VOPC", "SOP1", "SOP1_LIT", "SOP2", "SOP2_LIT", "SOPC", "SOPC_LIT", "SOPK", "SOPK_LIT", "VINTERP"}
+# intentionally not testing scratch memory ops
+MEM_FORMATS = {"VGLOBAL", "GLOBAL", "SMEM", "DS"}
+
+def should_skip(op:Enum) -> bool: return (name:=op.name) in SKIP or any(sub in name for sub in SKIP_SUBSTR)
+
+# ** named register assignments
+
+# ALU operands
+ALU_VGPR_STRIDE = 16          # v[0], v[16], v[32], ... per ALU operand slot
+ALU_SGPR_STRIDE = 4           # s[0], s[4], s[8], ... per ALU operand slot
+
+# memory address registers
+S_KERNARG_PTR = (0, 1)
+S_BUF_PTR = (2, 3)
+V_VADDR = (0, 1)
+V_DS_ADDR = 0
+
+# memory data registers
+MEM_VGPR_BASE = 32            # v[32], v[48], ... for vdst/vdata/vsrc
+MEM_VGPR_STRIDE = 16          # spacing between memory data vgpr slots
+MEM_SGPR_BASE = 8             # s[8], s[10], ... for SMEM sdata
+MEM_SGPR_STRIDE = 2           # spacing between memory data sgpr slots
+
+# ** create an ALU instruction based on the operands
+
+def create_alu_inst(op:Enum, builder:functools.partial[Inst]) -> Inst:
+  inst_cls, operands, slot = builder.func, OPERANDS[op], 0
+  kwargs:dict[str, Reg|int] = {}
+  for name, field in inst_cls._fields:
+    if isinstance(field, (FixedBitField, EnumBitField)): continue
+    nregs = max(1, operands[name][1] // 32) if name in operands else 1
+    is_sreg = name in operands and "SREG" in str(operands[name][2])
+    base_v, base_s = slot * ALU_VGPR_STRIDE, slot * ALU_SGPR_STRIDE
+    if name == "sdst" and isinstance(field, SGPRField): reg = VCC_LO
+    elif is_sreg and not isinstance(field, VGPRField): reg = VCC_LO
+    elif isinstance(field, VGPRField): reg = v[base_v:base_v+nregs-1] if nregs > 1 else v[base_v]
+    elif isinstance(field, SSrcField): reg = VCC_LO if nregs <= 2 else s[base_s:base_s+nregs-1] if nregs > 1 else s[base_s]
+    elif isinstance(field, SGPRField): reg = s[base_s:base_s+nregs-1] if nregs > 1 else s[base_s]
+    elif isinstance(field, SrcField): reg = v[base_v:base_v+nregs-1] if nregs > 1 else v[base_v]
+    else: reg = None
+    if reg is not None: kwargs[name] = reg; slot += 1
+    elif isinstance(field, BitField): kwargs[name] = field.default
+  return builder(**kwargs)
+
+# ** create a memory instruction with pre set address registers
+
+MEM_PRESET_REGS:dict[str, dict[str, Reg]] = {
+  "VGLOBAL":{"saddr":s[S_BUF_PTR[0]:S_BUF_PTR[1]], "vaddr":v[V_VADDR[0]:V_VADDR[1]]},
+  "GLOBAL":{"saddr":s[S_BUF_PTR[0]:S_BUF_PTR[1]], "addr":v[V_DS_ADDR]},  # addr is 32-bit offset when saddr is valid SGPR
+  "DS":{"addr":v[V_DS_ADDR]},
+  "SMEM":{"sbase":s[S_KERNARG_PTR[0]:S_KERNARG_PTR[1]], "soffset":NULL},
+}
+
+def create_mem_inst(op:Enum, builder:functools.partial[Inst]) -> Inst:
+  inst_cls, operands, field_map = builder.func, OPERANDS.get(op, {}), MEM_PRESET_REGS.get(builder.func.__name__, {})
+  kwargs:dict[str, Reg|int] = {}
+  vslot, sslot = 0, 0
+  for name, field in inst_cls._fields:
+    if isinstance(field, (FixedBitField, EnumBitField)): continue
+    if name in field_map:
+      kwargs[name] = field_map[name]
+      continue
+    nregs = max(1, operands[name][1] // 32) if name in operands else 1
+    if isinstance(field, VGPRField):
+      vi = MEM_VGPR_BASE + vslot * MEM_VGPR_STRIDE
+      kwargs[name] = v[vi:vi+nregs-1] if nregs > 1 else v[vi]
+      vslot += 1
+    elif isinstance(field, (SGPRField, AlignedSGPRField, SBaseField)):
+      si = MEM_SGPR_BASE + sslot * MEM_SGPR_STRIDE
+      kwargs[name] = s[si:si+nregs-1] if nregs > 1 else s[si]
+      sslot += 1
+    elif isinstance(field, BitField): kwargs[name] = field.default
+  return builder(**kwargs)
+
+# ** collect all memory and ALU instructions from the ISA autogen
+
+def collect_instructions() -> tuple[list[Inst], list[Inst], list[str]]:
+  op_map:dict[Enum, functools.partial[Inst]] = {}
+  for name, obj in inspect.getmembers(all_insts):
+    if isinstance(obj, functools.partial) and len(obj.args) == 1: op_map[obj.args[0]] = obj
+  alu_insts:list[Inst] = []
+  mem_insts:list[Inst] = []
+  skipped:list[str] = []
+  for op_enum, builder in op_map.items():
+    if should_skip(op_enum) or op_enum not in OPERANDS: skipped.append(op_enum.name); continue
+    fmt = builder.func.__name__
+    if fmt in ALU_FORMATS: alu_insts.append(create_alu_inst(op_enum, builder))
+    elif fmt in MEM_FORMATS: mem_insts.append(create_mem_inst(op_enum, builder))
+  return alu_insts, mem_insts, skipped
+
+def exec_insts(insts:list):
+  k = Kernel(arch)
+  # ** prologue for global memory
+  k.emit(s_load_b64(sdata=s[S_BUF_PTR[0]:S_BUF_PTR[1]], sbase=s[S_KERNARG_PTR[0]:S_KERNARG_PTR[1]], soffset=NULL))
+  k.waitcnt(lgkm=0)
+  k.emit(v_mov_b32_e32(v[V_VADDR[0]], 0))
+  k.emit(v_mov_b32_e32(v[V_VADDR[1]], 0))
+  # ** emit
+  for inst in insts: k.emit(inst)
+  k.emit(s_endpgm())
+  # ** run
+  NUM_THREADS, NUM_GRIDS, BUF_SIZE = 32, 1, 1024*1024
+  def fxn(A:UOp, B:UOp, C:UOp) -> UOp:
+    lidx, gidx = UOp.special(NUM_THREADS, "lidx0"), UOp.special(NUM_GRIDS, "gidx0")
+    lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=BUF_SIZE, addrspace=AddrSpace.LOCAL), (), "lds")
+    sink = UOp.sink(A.base, B.base, C.base, lds, lidx, gidx, arg=KernelInfo(name="discover_ops"))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in k.finalize()))))
+  A = Tensor.empty(BUF_SIZE, dtype=dtypes.uint8)
+  B = Tensor.empty(1, dtype=dtypes.uint8)
+  C = Tensor.empty(1, dtype=dtypes.uint8)
+  Tensor.custom_kernel(A, B, C, fxn=fxn)[0].realize()
+
+if __name__ == "__main__":
+  import sys
+  arch = Device[Device.DEFAULT].renderer.arch
+  if arch.startswith("gfx12"):
+    from tinygrad.runtime.autogen.amd.rdna4.ins import *
+    import tinygrad.runtime.autogen.amd.rdna4.ins as all_insts
+  elif arch.startswith("gfx11"):
+    from tinygrad.runtime.autogen.amd.rdna3.ins import *
+    import tinygrad.runtime.autogen.amd.rdna3.ins as all_insts
+    # these don"t exist in RDNA3, only RDNA3.5 and above
+    SKIP.update(["S_FMAAK_F32", "S_FMAMK_F32"])
+  else:
+    print(f"{arch} not supported yet")
+    sys.exit(0)
+  alu_insts, mem_insts, skipped = collect_instructions()
+  print(f"collected {len(alu_insts)} ALU + {len(mem_insts)} memory instructions ({len(skipped)} skipped)")
+  exec_insts(mem_insts+alu_insts)

extra/sqtt/examples/generate_examples.py

@@ -9,6 +9,7 @@ EXAMPLES = [
   "test/backend/test_custom_kernel.py TestCustomKernel.test_empty",
   "test/test_tiny.py TestTiny.test_plus",
   "test/test_tiny.py TestTiny.test_gemm",
+  "extra/sqtt/examples/discover_ops.py"
 ]
 
 if __name__ == "__main__":

test/amd/disasm.py

@@ -324,6 +324,12 @@ def _disasm_smem(inst: SMEM) -> str:
   if name in ('s_memrealtime', 's_memtime'): return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}"
   return f"{name} {_fmt_sdst(inst.sdata, dst_n, cdna)}, {sbase_str}, {off_s}" + _mods((inst.glc, " glc"), (getattr(inst, 'dlc', 0), " dlc"))
 
+R4_TH_LOAD = {1: 'TH_LOAD_NT', 2: 'TH_LOAD_HT', 3: 'TH_LOAD_LU', 4: 'TH_LOAD_RT_WB', 5: 'TH_LOAD_NT_WB'}
+R4_TH_STORE = {1: 'TH_STORE_NT', 2: 'TH_STORE_HT', 3: 'TH_STORE_ST', 4: 'TH_STORE_RT_WB', 5: 'TH_STORE_NT_WB'}
+R4_TH_ATOMIC = {1: 'TH_ATOMIC_RETURN', 2: 'TH_ATOMIC_NT', 3: 'TH_ATOMIC_RETURN_NT',
+                4: 'TH_ATOMIC_CASCADE_RT', 5: 'TH_ATOMIC_CASCADE_RETURN', 6: 'TH_ATOMIC_CASCADE_NT', 7: 'TH_ATOMIC_CASCADE_RETURN_NT'}
+R4_SCOPE = {1: 'SCOPE_SE', 2: 'SCOPE_DEV', 3: 'SCOPE_SYS'}
+
 def _disasm_flat(inst: FLAT) -> str:
   name, cdna, r4 = inst.op_name.lower(), _is_cdna(inst), _is_r4(inst)
   acc = getattr(inst, 'acc', 0)
@@ -331,9 +337,10 @@ def _disasm_flat(inst: FLAT) -> str:
   if r4: seg = 'flat' if (cls_name:=inst.__class__.__name__) == 'VFLAT' else ('global' if cls_name == 'VGLOBAL' else 'scratch')
   else: seg = ['flat', 'scratch', 'global'][inst.seg] if inst.seg < 3 else 'flat'
   instr = f"{seg}_{name.split('_', 1)[1] if '_' in name else name}"
-  # Global/scratch uses 13-bit signed offset
+  # Global/scratch uses 13-bit signed offset (RDNA3/CDNA), 24-bit signed offset (RDNA4)
   offset = inst.ioffset if r4 else inst.offset  # type: ignore[attr-defined]
-  if seg != 'flat':
+  if r4: off_val = offset if offset < (1 << 23) else offset - (1 << 24)  # sign extend 24-bit
+  elif seg != 'flat':
     if cdna:
       # CDNA: bit 12 is sign bit but not in offset field
       raw = int.from_bytes(inst.to_bytes(), 'little')
@@ -348,7 +355,9 @@ def _disasm_flat(inst: FLAT) -> str:
   w = regs.get('data', regs.get('d', 1)) if 'store' in name or 'atomic' in name else regs.get('d', 1)
   off_s = f" offset:{off_val}" if off_val else ""
   if cdna: mods = f"{off_s}{' sc0' if inst.sc0 else ''}{' nt' if inst.nt else ''}{' sc1' if getattr(inst, 'sc1', 0) else ''}"  # type: ignore[attr-defined]
-  elif r4: mods = f"{off_s}{' scope' if inst.scope else ''}{' th' if inst.th else ''}"  # type: ignore[attr-defined]
+  elif r4:
+    th_names = R4_TH_ATOMIC if 'atomic' in name else (R4_TH_STORE if 'store' in name else R4_TH_LOAD)
+    mods = off_s + (f" th:{th_names[inst.th]}" if inst.th in th_names else "") + (f" scope:{R4_SCOPE[inst.scope]}" if inst.scope in R4_SCOPE else "")
   else: mods = f"{off_s}{' glc' if inst.glc else ''}{' slc' if inst.slc else ''}{' dlc' if inst.dlc else ''}"
   if seg == 'flat': saddr_s = ""
   elif _unwrap(inst.saddr) in (0x7F, 124): saddr_s = ", off"
@@ -357,7 +366,7 @@ def _disasm_flat(inst: FLAT) -> str:
     saddr_s = f", {(SPECIAL_PAIRS_CDNA if cdna else SPECIAL_PAIRS)[_unwrap(inst.saddr)]}"
   elif t := _ttmp(inst.saddr, 2): saddr_s = f", {t}"
   else: saddr_s = f", {_sreg(inst.saddr, 2) if _unwrap(inst.saddr) < 106 else decode_src(_unwrap(inst.saddr), cdna)}"
-  if 'addtid' in name: return f"{instr} {reg_fn(inst.data if 'store' in name else inst.vdst)}{saddr_s}{mods}"
+  if 'addtid' in name: return f"{instr} {reg_fn((inst.vsrc if r4 else inst.data) if 'store' in name else inst.vdst)}{saddr_s}{mods}"
    # RDNA4: vaddr instead of addr, vsrc instead of data
   addr = inst.vaddr if r4 else inst.addr  # type: ignore[attr-defined]
   data = inst.vsrc if r4 else inst.data  # type: ignore[attr-defined]
@@ -372,7 +381,7 @@ def _disasm_flat(inst: FLAT) -> str:
   addr_s = "off" if not inst.sve and seg == 'scratch' else _vreg(addr, addr_w)
   data_s, vdst_s = reg_fn(data, w), reg_fn(inst.vdst, w // 2 if 'cmpswap' in name else w)
   if 'atomic' in name:
-    glc_or_sc0 = inst.sc0 if cdna else inst.glc  # type: ignore[attr-defined]
+    glc_or_sc0 = inst.sc0 if cdna else (inst.th & 1 if r4 else inst.glc)  # type: ignore[attr-defined]
     sfx = f"{saddr_s if seg != 'flat' else ''}{mods}"
     return f"{instr} {vdst_s}, {addr_s}, {data_s}{sfx}" if glc_or_sc0 else f"{instr} {addr_s}, {data_s}{sfx}"
   if 'store' in name: return f"{instr} {addr_s}, {data_s}{saddr_s}{mods}"

test/amd/test_llvm.py

@@ -40,7 +40,7 @@ RDNA4_FILES = ['gfx12_asm_sop1.s', 'gfx12_asm_sop2.s', 'gfx12_asm_sopp.s', 'gfx1
   'gfx12_asm_vop1.s', 'gfx12_asm_vop2.s', 'gfx12_asm_vopc.s', 'gfx12_asm_vopcx.s', 'gfx12_asm_vop3.s', 'gfx12_asm_vop3c.s',
   'gfx12_asm_vop3cx.s', 'gfx12_asm_vop3p.s', 'gfx12_asm_vop3_from_vop1.s', 'gfx12_asm_vop3_from_vop2.s',
   'gfx12_asm_vop3p_features.s', 'gfx12_asm_vopd.s', 'gfx12_asm_vopd_features.s',
-  'gfx12_asm_ds.s', 'gfx12_asm_smem.s',
+  'gfx12_asm_ds.s', 'gfx12_asm_smem.s', 'gfx12_asm_vflat.s',
   'gfx12_asm_wmma_w32.s']
 
 def _parse_llvm_tests(text: str, pattern: str) -> list[tuple[str, bytes]]:

test/amd/test_sqtt_tables.py

@@ -134,7 +134,10 @@ class TestSQTTMatchesBinary(unittest.TestCase):
   def _test_bit_counts(self, layout: int):
     if not (tables := extract_bit_tables()): self.skipTest("rocprof-trace-decoder not installed")
     from tinygrad.renderer.amd.sqtt import PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4
+    # rocprof's bit table says L4 type 7 (TS_DELTA_S8_W3) is 72 bits, but the actual decoder uses 64 bits
+    skip = {(4, 7)}
     for type_id, pkt_cls in {3: PACKET_TYPES_RDNA3, 4: PACKET_TYPES_RDNA4}[layout].items():
+      if (layout, type_id) in skip: continue
       with self.subTest(packet=pkt_cls.__name__):
         self.assertEqual(pkt_cls._size_nibbles * 4, tables[layout - 2][type_id])  # type: ignore[attr-defined]
 

test/amd/test_sqttmap.py

@@ -2,7 +2,7 @@
 import unittest, pickle
 from typing import Iterator
 from pathlib import Path
-from tinygrad.helpers import DEBUG
+from tinygrad.helpers import DEBUG, OSX
 from tinygrad.renderer.amd.sqtt import print_packets, map_insts
 from tinygrad.runtime.autogen.amd.rdna3.ins import s_endpgm
 from test.amd.disasm import disasm
@@ -10,7 +10,7 @@ from test.amd.disasm import disasm
 import tinygrad
 EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"
 
-def rocprof_inst_traces_match(sqtt, prg, target):
+def rocprof_inst_traces_match(sqtt, prg, target, pass_rocprof_err=False):
   from tinygrad.viz.serve import amd_decode
   from extra.sqtt.roc import decode as roc_decode, InstExec
   addr_table = amd_decode(prg.lib, target)
@@ -30,7 +30,7 @@ def rocprof_inst_traces_match(sqtt, prg, target):
     rocprof_inst = next(rwaves_iter[info.wave][0])
     ref_pc = rocprof_inst.pc-prg.base
     # always check pc matches
-    assert ref_pc == info.pc, f"pc mismatch {ref_pc}:{disasm_map[rocprof_inst.pc]} != {info.pc}:{disasm(info.inst)}"
+    assert ref_pc == info.pc or pass_rocprof_err, f"pc mismatch {ref_pc}:{disasm_map[rocprof_inst.pc]} != {info.pc}:{disasm(info.inst)}"
     # special handling for s_endpgm, it marks the wave completion.
     if info.inst == s_endpgm():
       completed_wave = list(rwaves_iter[info.wave].pop(0))
@@ -67,7 +67,9 @@ class TestSQTTMapBase(unittest.TestCase):
         if not event.itrace: continue
         if event.kern not in kern_events: continue
         with self.subTest(example=name, kern=event.kern):
-          passed_insts, n_waves, n_units = rocprof_inst_traces_match(event, kern_events[event.kern], target)
+          # rocprof OSX has a bug for sopk decoding, linux rocprof works
+          pass_rocprof_err = OSX and target == "gfx1200" and name.startswith("profile_py")
+          passed_insts, n_waves, n_units = rocprof_inst_traces_match(event, kern_events[event.kern], target, pass_rocprof_err)
           if n_waves: print(f"{name}: passed for {passed_insts} instructions across {n_waves} waves scheduled on {n_units} wave units")
 
 class TestSQTTMapRDNA3(TestSQTTMapBase): target = "gfx1100"

test/backend/test_const_folding.py

@@ -27,6 +27,12 @@ class TestMovedConstFolding(unittest.TestCase):
   def test_add_padded_one(self):
     _check_ast_count(1, Tensor([1.0, 2, 3, 4]) * Tensor.ones(2).pad(((1, 1),)))
 
+  def test_copy_padded_const(self):
+    schedule = Tensor.ones(4, device="CPU:0").pad(((1, 1),)).to("CPU:1").schedule()
+    assert not any(si.ast.op is Ops.COPY for si in schedule), "const copy should be folded"
+    # TODO: this is wrong, should be [0, 1, 1, 1, 1, 0]
+    np.testing.assert_equal(Tensor.ones(4, device="CPU:0").pad(((1, 1),)).to("CPU:1").numpy(), [1, 1, 1, 1, 1, 1])
+
   def test_cast_padded(self):
     # NOTE: it's always 1 kernel when calling .numpy, limitation of _check_ast_count
     if is_dtype_supported(dtypes.int16):

test/backend/test_multitensor.py

@@ -228,17 +228,17 @@ class TestMultiTensor(unittest.TestCase):
       a,b = _test_allreduce(Tensor.rand(256, 256))
       np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
 
-  def test_multiple_to_single_device_naive(self):
-    with Context(RING=0):
-      t = Tensor.arange(32).shard(devices_4, 0).to(Device.DEFAULT).realize()
-    self.assertEqual(t.device, Device.DEFAULT)
-    np.testing.assert_equal(t.numpy(), np.arange(32))
-
-  def test_multiple_to_single_device_ring(self):
-    with Context(RING=2):
-      t = Tensor.arange(32).shard(devices_4, 0).to(Device.DEFAULT).realize()
-    self.assertEqual(t.device, Device.DEFAULT)
-    np.testing.assert_equal(t.numpy(), np.arange(32))
+  def test_multiple_to_single_device(self):
+    kernel_counts = {}
+    for ring in (0, 2):
+      GlobalCounters.reset()
+      with Context(RING=ring, SCACHE=0):
+        t = Tensor.arange(32).contiguous().shard(devices_4, 0).to(Device.DEFAULT)
+        t.realize()
+      kernel_counts[ring] = GlobalCounters.kernel_count
+      self.assertEqual(t.device, Device.DEFAULT)
+      np.testing.assert_equal(t.numpy(), np.arange(32))
+    self.assertNotEqual(kernel_counts[0], kernel_counts[2])
 
   def test_allreduce_all2all(self):
     with Context(ALL2ALL=2):

test/backend/test_setitem.py

@@ -205,6 +205,20 @@ class TestSetitem(unittest.TestCase):
     n[:, ind_1.numpy(), :, ind_2.numpy(), :] = v.numpy()
     np.testing.assert_equal(t.numpy(), n)
 
+  def test_setitem_tensor_int_indexing(self):
+    t = Tensor.zeros(4, 3, dtype=dtypes.int).contiguous()
+    t[Tensor([0, 2]), 0] = Tensor([99, 88], dtype=dtypes.int)
+    n = np.zeros((4, 3), dtype=np.int32)
+    n[[0, 2], 0] = [99, 88]
+    np.testing.assert_equal(t.numpy(), n)
+
+  def test_setitem_tensor_slice_indexing(self):
+    t = Tensor.zeros(4, 3, dtype=dtypes.int).contiguous()
+    t[Tensor([0, 2]), :2] = Tensor([[10, 20], [30, 40]], dtype=dtypes.int)
+    n = np.zeros((4, 3), dtype=np.int32)
+    n[[0, 2], :2] = [[10, 20], [30, 40]]
+    np.testing.assert_equal(t.numpy(), n)
+
   def test_setitem_2d_tensor_indexing(self):
     t = Tensor.zeros(2, dtype=dtypes.int).contiguous()
     index = Tensor([[0, 1], [1,0]])
@@ -279,17 +293,43 @@ class TestWithGrad(unittest.TestCase):
     x = Tensor.rand(8)
     z[:3] = x
 
-  def test_set_into_requires_grad(self):
-    z = Tensor.rand(8, 8, requires_grad=True)
-    x = Tensor.rand(8)
-    with self.assertRaises(NotImplementedError):
-      z[:3] = x
-
   def test_set_with_requires_grad(self):
-    z = Tensor.rand(8, 8)
-    x = Tensor.rand(8, requires_grad=True)
-    with self.assertRaises(NotImplementedError):
-      z[:3] = x
+    z = Tensor.ones(8, 8)
+    x = Tensor.rand(8, 8, requires_grad=True)
+    z[:] = x
+    z.sum().backward()
+    np.testing.assert_allclose(x.grad.numpy(), np.ones((8, 8)))
+
+  def test_set_nonleaf_requires_grad(self):
+    x = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
+    z = x * 2
+    z[:2] = Tensor([10.0, 20.0])
+    z.sum().backward()
+    np.testing.assert_allclose(x.grad.numpy(), [0, 0, 2, 2])
+
+  def test_set_overlapping_requires_grad(self):
+    z = Tensor.zeros(6, requires_grad=True)
+    x = Tensor.ones(4, requires_grad=True)
+    y = Tensor.ones(4, requires_grad=True) * 2
+    z[:4] = x
+    z[2:] = y
+    z.sum().backward()
+    np.testing.assert_allclose(x.grad.numpy(), [1, 1, 0, 0])
+    np.testing.assert_allclose(y.grad.numpy(), np.ones(4))
+
+  def test_set_iadd_requires_grad(self):
+    z = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
+    x = Tensor([10.0, 20.0], requires_grad=True)
+    z[:2] += x
+    z.sum().backward()
+    np.testing.assert_allclose(z.grad.numpy(), np.ones(4))
+    np.testing.assert_allclose(x.grad.numpy(), np.ones(2))
+
+  def test_set_used_before_setitem(self):
+    z = Tensor([1.0, 2.0, 3.0, 4.0], requires_grad=True)
+    _ = z.sum()
+    with self.assertRaises(RuntimeError):
+      z[:2] = Tensor([0.0, 0.0])
 
 class TestSetitemLoop(unittest.TestCase):
   def test_arange(self):

test/backend/test_stunning.py

@@ -25,7 +25,7 @@ class TestStunning(unittest.TestCase):
     nv = a[12].cat(a[76]).tolist()
 
     vi = Variable('i', 0, a.shape[0]-1)
-    with self.assertRaisesRegex(AssertionError, "bind mismatch on"):
+    with self.assertRaisesRegex(RuntimeError, "bind mismatch on"):
       wv = a[vi.bind(12)].cat(a[vi.bind(76)]).tolist()
       self.assertListEqual(nv, wv)
 

test/backend/test_symbolic_ops.py

@@ -302,5 +302,62 @@ class TestSymbolicOps(unittest.TestCase):
     expected = x_full[:, :, :val].conv2d(weight=weight, groups=1, stride=6, dilation=1, padding=(3, 3))
     np.testing.assert_allclose(result[:, :, :13].numpy(), expected.numpy(), atol=1e-5, rtol=1e-5)
 
+  def test_triu_symbolic(self):
+    a = Tensor.rand(10, 10).realize()
+    for i in range(2, 6):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = a[:vi, :vi].triu()
+      # extract concrete-sized result from symbolic output
+      symbolic_np = symbolic[:i, :i].numpy()
+      expected = a[:i, :i].triu().numpy()
+      np.testing.assert_allclose(symbolic_np, expected, atol=1e-6, rtol=1e-6)
+
+  def test_tril_symbolic(self):
+    a = Tensor.rand(10, 10).realize()
+    for i in range(2, 6):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = a[:vi, :vi].tril()
+      symbolic_np = symbolic[:i, :i].numpy()
+      expected = a[:i, :i].tril().numpy()
+      np.testing.assert_allclose(symbolic_np, expected, atol=1e-6, rtol=1e-6)
+
+  def test_triu_symbolic_diagonal(self):
+    a = Tensor.rand(10, 10).realize()
+    for i in range(3, 6):
+      vi = Variable("i", 1, 10).bind(i)
+      for diag in [-1, 0, 1, 2]:
+        symbolic = a[:vi, :vi].triu(diagonal=diag)
+        symbolic_np = symbolic[:i, :i].numpy()
+        expected = a[:i, :i].triu(diagonal=diag).numpy()
+        np.testing.assert_allclose(symbolic_np, expected, atol=1e-6, rtol=1e-6)
+
+  def test_triu_symbolic_nonsquare(self):
+    a = Tensor.rand(10, 8).realize()
+    for i in range(2, 6):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = a[:vi, :].triu()
+      symbolic_np = symbolic[:i, :].numpy()
+      expected = a[:i, :].triu().numpy()
+      np.testing.assert_allclose(symbolic_np, expected, atol=1e-6, rtol=1e-6)
+
+  def test_full_triu_symbolic(self):
+    """Test the attention mask pattern: Tensor.full(symbolic_shape, -inf).triu(k)"""
+    for T in range(2, 6):
+      for sp in [0, 2, 5]:
+        vT = Variable("T", 1, 10).bind(T)
+        mask = Tensor.full((1, 1, vT, sp+vT), float("-inf")).triu(sp+1)
+        ref = Tensor.full((1, 1, T, sp+T), float("-inf")).triu(sp+1)
+        # compare element sums (counting -inf elements)
+        sym_finite = mask[:, :, :T, :sp+T].isnan().logical_not().cast(dtypes.int32).sum().item()
+        ref_finite = ref.isnan().logical_not().cast(dtypes.int32).sum().item()
+        self.assertEqual(sym_finite, ref_finite)
+
+  def test_arange_symbolic(self):
+    for i in range(1, 6):
+      vi = Variable("i", 1, 10).bind(i)
+      symbolic = Tensor.arange(vi)
+      expected = Tensor.arange(i)
+      np.testing.assert_allclose(symbolic[:i].numpy(), expected.numpy(), atol=1e-6, rtol=1e-6)
+
 if __name__ == '__main__':
   unittest.main()

test/null/test_schedule.py

@@ -1104,6 +1104,7 @@ class TestUOpBecome(unittest.TestCase):
     from tinygrad.helpers import all_same
     assert all_same([x.uop.base.realized for x in [a,b,c]])
 
+  @unittest.skip("not clear if we want this")
   def test_setitem_becomes_subbuffer(self):
     a = Tensor.full((4,), 2.).contiguous().realize()
     b = a.shrink(((0, 2),)).assign(Tensor.full((2,), 1.0))

test/unit/test_assign.py

@@ -29,6 +29,7 @@ class TestAssign(unittest.TestCase):
     a.realize()
     np.testing.assert_allclose(b.numpy(), 0)
 
+  @unittest.skip("TODO: this often crashes in CI")
   def test_assign_zeros(self):
     a = Tensor.zeros(10,10).contiguous()
     b = Tensor.zeros(10,10).contiguous()
@@ -269,6 +270,16 @@ class TestAssign(unittest.TestCase):
     out = attn.cache_k.flatten().numpy()
     np.testing.assert_allclose(out, [1.,1.,1.,1.,1.,1.,0.,0.,1.,1.,1.,1.,1.,1.,0.,0.])
 
+  def test_assign_after(self):
+    t = Tensor.zeros(10).contiguous().realize()
+    t.uop = t.uop.after(t.uop.assign((t+1).uop))
+    np.testing.assert_allclose(t.numpy(), [1.,1.,1.,1.,1.,1.,1.,1.,1.,1.])
+
+  def test_assign_after_partial(self):
+    t = Tensor.zeros(10).contiguous().realize()
+    t.uop = t.uop.after(t[:5].uop.assign(Tensor.ones(5).uop))
+    np.testing.assert_allclose(t.numpy(), [1.,1.,1.,1.,1.,0.,0.,0.,0.,0.])
+
   def test_assign_contiguous(self):
     b = Tensor.arange(16).reshape(4,4).contiguous().realize()
     a = (Tensor.arange(16).reshape(4,4).contiguous().realize() + 1)
@@ -598,8 +609,8 @@ class TestAssign(unittest.TestCase):
       x = q + caches[i][:1]             # next layer also references the same CONTIGUOUS through q
     GlobalCounters.reset()
     caches[-1][:1].contiguous().realize()
-    # 2 kernels for first assign + 3 per remaining assign (matmul, contiguous, assign) + 1 final read = 3*N
-    self.assertEqual(GlobalCounters.kernel_count, 3*N)
+    # N matmuls + N assigns + 1 final read = 2*N+1 (AFTER embedding allows full graph scheduling with shared contiguous reuse)
+    self.assertEqual(GlobalCounters.kernel_count, 2*N+1)
 
 
 class TestAssignOrdering(unittest.TestCase):
@@ -756,13 +767,12 @@ class TestAssignOrdering(unittest.TestCase):
     np.testing.assert_equal(b.numpy(), [1, 2, 3, 4])
 
   def test_variable_slice_ordering(self):
-    """Variable-indexed slices - tests symbolic dependency tracking."""
+    """Variable-indexed slices - conflicting variable binds in same schedule are rejected."""
     v_i = Variable("i", 0, 3)
     buf = Tensor.zeros(4, 4).contiguous().realize()
     buf[v_i.bind(0):v_i.bind(0)+1, :].assign(Tensor.ones(1, 4))
     buf[v_i.bind(1):v_i.bind(1)+1, :].assign(Tensor.ones(1, 4) * 2)
-    self.assertEqual(buf[0:1, :].sum().item(), 4)
-    self.assertEqual(buf[1:2, :].sum().item(), 8)
+    with self.assertRaises(RuntimeError): buf[0:1, :].sum().item()
 
   def test_multi_step_assign_read_write_same_buffer(self):
     """Assign to m and param reading b, then update b, across multiple steps.

test/unit/test_function.py

@@ -2,6 +2,7 @@ import numpy as np
 import unittest
 from tinygrad.function import function
 from tinygrad import Tensor
+from tinygrad.uop.ops import UOp
 
 class TestFunction(unittest.TestCase):
   def test_simple(self):
@@ -102,7 +103,6 @@ class TestFunction(unittest.TestCase):
     np.testing.assert_allclose(w.grad.numpy(), [4., 5., 6.])
 
   def test_symbolic_index(self):
-    from tinygrad.uop.ops import UOp
     table = Tensor([10,20,30,40]).contiguous().realize()
     @function
     def f(x:Tensor, start_pos:int|UOp) -> Tensor:
@@ -111,6 +111,14 @@ class TestFunction(unittest.TestCase):
     v = UOp.variable("start_pos", 0, 3)
     np.testing.assert_equal(f(Tensor([1,2,3]), v.bind(0)).numpy(), [11,12,13])
 
+  def test_symbolic_shape_input(self):
+    table = Tensor([10,20,30,40]).contiguous().realize()
+    @function
+    def f(x:Tensor) -> Tensor: return x * 2
+    sz = UOp.variable("sz", 1, 3)
+    slic = table[:sz.bind(2)]
+    np.testing.assert_equal(f(slic)[:2].numpy(), [20,40])
+
   def test_nested_calls(self):
     w = Tensor([10., 20., 30.])
     @function
@@ -121,6 +129,31 @@ class TestFunction(unittest.TestCase):
     a = Tensor([1., 2., 3.])
     np.testing.assert_allclose(g(f(a)).numpy(), [110., 440., 990.])
 
+  def test_nested_calls_backward(self):
+    w = Tensor([[1., 2.], [3., 4.]]).contiguous().realize()
+    @function
+    def inner(x:Tensor) -> Tensor: return x + w
+    @function
+    def outer(a:Tensor, b:Tensor) -> Tensor: return inner(a.reshape(1,2) + b.reshape(1,2))
+
+    a = Tensor([1., 2.], requires_grad=True)
+    b = Tensor([3., 4.], requires_grad=True)
+    outer(a, b).sum().backward()
+    np.testing.assert_allclose(a.grad.numpy(), [2., 2.])
+    np.testing.assert_allclose(b.grad.numpy(), [2., 2.])
+
+  def test_unused_param_backward(self):
+    @function
+    def f(a:Tensor, b:Tensor, c:Tensor) -> Tensor: return a + c  # b is unused
+
+    a = Tensor([1., 2., 3.], requires_grad=True)
+    b = Tensor([4., 5., 6.], requires_grad=True)
+    c = Tensor([7., 8., 9.], requires_grad=True)
+    f(a, b, c).sum().backward()
+    np.testing.assert_allclose(a.grad.numpy(), [1., 1., 1.])
+    np.testing.assert_allclose(b.grad.numpy(), [0., 0., 0.])
+    np.testing.assert_allclose(c.grad.numpy(), [1., 1., 1.])
+
   def test_name(self):
     @function
     def f(a:Tensor) -> Tensor: return a + 1
@@ -185,5 +218,114 @@ class TestFunction(unittest.TestCase):
     np.testing.assert_equal(a.numpy(), [1,2,3])
     np.testing.assert_equal(b.numpy(), [10,20,30])
 
+class TestFunctionMulti(unittest.TestCase):
+  devices_2 = ("CPU:0", "CPU:1")
+
+  def test_simple_multi(self):
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor: return a+b
+
+    a = Tensor([1,2,3,4]).shard(self.devices_2, axis=None)
+    b = Tensor([10,20,30,40]).shard(self.devices_2, axis=None)
+    np.testing.assert_equal(f(a,b).numpy(), [11,22,33,44])
+
+  def test_simple_multi_sharded(self):
+    @function
+    def f(a:Tensor, b:Tensor) -> Tensor: return a+b
+
+    a = Tensor([1,2,3,4]).shard(self.devices_2, axis=0)
+    b = Tensor([10,20,30,40]).shard(self.devices_2, axis=0)
+    np.testing.assert_equal(f(a,b).numpy(), [11,22,33,44])
+
+  def test_data_parallel_multi(self):
+    @function
+    def f(x:Tensor, w:Tensor) -> Tensor: return x @ w
+
+    x = Tensor([[1.,2.],[3.,4.],[5.,6.],[7.,8.]]).shard(self.devices_2, axis=0)
+    w = Tensor([[1.,0.],[0.,1.]]).shard(self.devices_2, axis=None)
+    np.testing.assert_allclose(f(x, w).numpy(), [[1.,2.],[3.,4.],[5.,6.],[7.,8.]])
+
+  def test_grad_implicit_multi(self):
+    w = Tensor([1., 2., 3., 4.], requires_grad=True).shard(self.devices_2, axis=None)
+    w.realize()
+    @function
+    def f(x:Tensor) -> Tensor: return x * w
+
+    x = Tensor([4., 5., 6., 7.]).shard(self.devices_2, axis=None)
+    f(x).sum().backward()
+    np.testing.assert_allclose(w.grad.numpy(), [4., 5., 6., 7.])
+
+  def test_call_axis(self):
+    @function
+    def f(x:Tensor, w:Tensor) -> Tensor: return x @ w
+
+    x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]]).shard(self.devices_2, axis=0)
+    w = Tensor([[1.,2.],[3.,4.]]).shard(self.devices_2, axis=None)
+    result = f(x, w)
+    # CALL output should inherit axis=0 from the sharded input
+    self.assertEqual(result.uop.axis, 0)
+    # reduce on the sharded axis should remove it
+    self.assertIsNone(result.sum().uop.axis)
+
+  def test_call_axis_shard_inside(self):
+    @function
+    def f(x:Tensor, w:Tensor) -> Tensor:
+      return x.shard(self.devices_2, axis=0) @ w.shard(self.devices_2, axis=None)
+
+    x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]])
+    w = Tensor([[1.,2.],[3.,4.]])
+    result = f(x, w)
+    self.assertEqual(result.uop.axis, 0)
+    np.testing.assert_allclose(result.numpy(), x.numpy() @ w.numpy())
+
+  def test_data_parallel_backward(self):
+    @function
+    def f(x:Tensor, w:Tensor) -> Tensor: return x @ w
+
+    x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]], requires_grad=True).shard(self.devices_2, axis=0)
+    w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(self.devices_2, axis=None)
+    w.realize()
+    f(x, w).sum().backward()
+    # d/dx = ones @ w^T = [[1,3],[1,3],[1,3],[1,3]], but sum so ones(4,2) @ w^T? no:
+    # L = sum(x @ w), dL/dx = ones(4,2) @ w^T... actually dL/d(xw) = ones(4,2), dL/dx = ones(4,2) @ w^T
+    np.testing.assert_allclose(x.grad.numpy(), np.ones((4,2)) @ np.array([[1,3],[2,4]]))
+
+  def test_data_parallel_backward_4(self):
+    devices_4 = tuple(f"CPU:{i}" for i in range(4))
+    @function
+    def f(x:Tensor, w:Tensor) -> Tensor: return x @ w
+
+    x = Tensor(np.arange(16).reshape(8,2).astype(np.float32), requires_grad=True).shard(devices_4, axis=0)
+    w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(devices_4, axis=None)
+    w.realize()
+    f(x, w).sum().backward()
+    np.testing.assert_allclose(x.grad.numpy(), np.ones((8,2)) @ np.array([[1,3],[2,4]]))
+
+  def test_data_parallel_backward_implicit(self):
+    devices_4 = tuple(f"CPU:{i}" for i in range(4))
+    w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(devices_4, axis=None)
+    w.realize()
+    @function
+    def f(x:Tensor) -> Tensor: return x @ w
+
+    x = Tensor(np.arange(16).reshape(8,2).astype(np.float32), requires_grad=True).shard(devices_4, axis=0)
+    f(x).sum().backward()
+    np.testing.assert_allclose(x.grad.numpy(), np.ones((8,2)) @ np.array([[1,3],[2,4]]))
+
+  def test_data_parallel_backward_twice(self):
+    devices_4 = tuple(f"CPU:{i}" for i in range(4))
+    w = Tensor([[1.,2.],[3.,4.]], requires_grad=True).shard(devices_4, axis=None)
+    w.realize()
+    # pre-init grads like the training loop does
+    w.grad = w.zeros_like().contiguous().realize()
+    @function
+    def f(x:Tensor) -> Tensor: return x @ w
+
+    expected = np.ones((8,2)) @ np.array([[1,3],[2,4]])
+    for _ in range(2):
+      x = Tensor(np.arange(16).reshape(8,2).astype(np.float32), requires_grad=True).shard(devices_4, axis=0)
+      f(x).sum().backward()
+      np.testing.assert_allclose(x.grad.numpy(), expected)
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_gguf.py

@@ -1,4 +1,4 @@
-import os, unittest
+import os, struct, unittest
 from tinygrad import dtypes, Tensor, fetch, Device
 from tinygrad.nn.state import ggml_data_to_tensor, gguf_load
 from tinygrad.device import is_dtype_supported
@@ -120,5 +120,41 @@ class TestGGUF(unittest.TestCase):
       else:
         self.assertEqual(kv_data[k], read_val(-1))
 
+class TestGGUFGEMV(unittest.TestCase):
+  def _test_gguf_gemv(self, qtype: GGMLQuantizationType):
+    block_size, type_size = GGML_QUANT_SIZES[qtype]
+    rows, cols = 8192, 2048
+    n_blocks = rows * cols // block_size
+    rng = np.random.default_rng(42)
+    # generate random quantized blocks with valid fp16 scale fields (random bytes can produce NaN scales)
+    q_data = rng.integers(0, 256, size=n_blocks * type_size, dtype=np.uint8).reshape(n_blocks, type_size)
+    scales = np.float16(rng.standard_normal(n_blocks * 4)).view(np.uint8).reshape(n_blocks, -1)
+    if qtype == GGMLQuantizationType.Q8_0: q_data[:, :2] = scales[:, :2]                    # d at offset 0
+    elif qtype == GGMLQuantizationType.Q4_K: q_data[:, :4] = scales[:, :4]                  # d, dmin at offset 0
+    elif qtype == GGMLQuantizationType.Q6_K: q_data[:, -2:] = scales[:, :2]                 # d at end
+    q_data = q_data.flatten()
+    ref = dequantize(q_data, qtype).reshape(rows, cols)
+
+    # build a minimal gguf in memory: header + 1 tensor info + aligned data
+    buf = bytearray()
+    buf += struct.pack("<4siqq", b"GGUF", 3, 1, 0)              # magic, version, n_tensors, n_kv
+    buf += struct.pack("<Q", 6) + b"weight"                      # tensor name
+    buf += struct.pack("<I", 2)                                  # ndims
+    buf += struct.pack("<QQ", cols, rows)                        # dims (gguf stores reversed)
+    buf += struct.pack("<i", qtype.value)
+    buf += struct.pack("<Q", 0)                                  # offset
+    buf += b"\x00" * ((32 - len(buf) % 32) % 32)                # pad to alignment=32
+    buf += q_data.tobytes()
+
+    _, tensors = gguf_load(Tensor(np.frombuffer(buf, dtype=np.uint8)).to(None))
+
+    x = rng.standard_normal(cols).astype(np.float32)
+    np.testing.assert_allclose((tensors["weight"] @ Tensor(x)).numpy(), ref @ x, atol=1e-2, rtol=1e-2)
+    np.testing.assert_equal(tensors["weight"].numpy(), ref)
+
+  def test_gguf_gemv_q8_0(self): self._test_gguf_gemv(GGMLQuantizationType.Q8_0)
+  def test_gguf_gemv_q4_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q4_K)
+  def test_gguf_gemv_q6_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q6_K)
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_gradient.py

@@ -68,6 +68,14 @@ class TestTensorGradient(unittest.TestCase):
     np.testing.assert_allclose(x.grad.numpy(), [2.0+3.0+2*3.0])
     self.assertIs(x.grad, old_grad)
 
+  def test_gradient_through_chained_unrealized_setitem(self):
+    g1 = Tensor.zeros(4).contiguous()
+    g1[2] = Tensor(1.0)
+    g2 = Tensor.zeros(5, 4).contiguous()
+    g2[0] = g1
+    x = Tensor.randn(4, 4)
+    np.testing.assert_allclose(x.pad(((1,0),(0,0))).gradient(x, gradient=g2)[0].numpy(), np.zeros((4, 4)))
+
 class TestViewGradient(unittest.TestCase):
   def test_expand(self):
     x = Tensor.randn(5,2)

test/unit/test_indexing.py

@@ -179,8 +179,6 @@ class TestIndexing(unittest.TestCase):
     def delitem(): del reference[0]
     self.assertRaises(TypeError, delitem)
 
-  # TODO setitem backward
-  '''
   def test_set_item_to_scalar_tensor(self):
     m = random.randint(1, 10)
     n = random.randint(1, 10)
@@ -190,7 +188,6 @@ class TestIndexing(unittest.TestCase):
     z[:, 0] = w
     z.sum().backward()
     numpy_testing_assert_equal_helper(w.grad, m * a)
-  '''
 
   def test_step(self):
     v = Tensor.arange(10)

test/unit/test_llm_server.py

@@ -6,13 +6,17 @@ class TestTransformerGenerate(unittest.TestCase):
   def test_start_pos_parameter_is_used(self):
     """Test that start_pos parameter is not ignored (regression test for always resetting to 0)."""
     from tinygrad.apps.llm import Transformer
+    from tinygrad.uop.ops import UOp
     # Create a minimal transformer
     model = Transformer(num_blocks=1, dim=64, hidden_dim=128, n_heads=2, n_kv_heads=2,
                         norm_eps=1e-5, vocab_size=100, head_dim=32, rope_theta=10000.0, max_context=32)
 
     captured_inputs = []
     def mock_call(self, tokens, start_pos):
-      captured_inputs.append((tokens.shape, start_pos if isinstance(start_pos, int) else start_pos.bind_val))
+      s = tokens.shape[-1]
+      shape_val = s.val if isinstance(s, UOp) else s
+      sp_val = start_pos.val if isinstance(start_pos, UOp) else start_pos
+      captured_inputs.append((shape_val, sp_val))
       return Tensor([[42]])  # return a fake next token
 
     with patch.object(Transformer, '__call__', mock_call):
@@ -22,8 +26,179 @@ class TestTransformerGenerate(unittest.TestCase):
 
     # With start_pos=3, the initial tensor should only have tokens[3:] = [4, 5] (length 2)
     # If the bug existed (start_pos always reset to 0), it would have all 5 tokens
-    self.assertEqual(captured_inputs[0][0][-1], 2)  # shape should be (1, 2)
+    self.assertEqual(captured_inputs[0][0], 2)  # shape should have 2 tokens
     self.assertEqual(captured_inputs[0][1], 3)  # start_pos should be 3, not 0
 
+class TestPrefillCorrectness(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    from tinygrad.apps.llm import Transformer
+    cls.model = Transformer(num_blocks=2, dim=64, hidden_dim=128, n_heads=2, n_kv_heads=2,
+                            norm_eps=1e-5, vocab_size=100, head_dim=32, rope_theta=10000.0, max_context=128)
+
+  def _reset_caches(self, model):
+    for block in model.blk:
+      if hasattr(block, "cache_kv"): del block.cache_kv
+      if hasattr(block, "causal_mask"): del block.causal_mask
+
+  def test_prefill_matches_forward(self):
+    """Test that prefill_forward produces the same output as forward for various lengths."""
+    model = self.__class__.model
+    v_prefill = model.v_prefill
+    token_buf = Tensor.zeros(1, model.max_context, dtype="int32").contiguous().realize()
+    for length in [2, 4, 8, 16]:
+      tokens = Tensor.randint(length, high=100, dtype="int32")
+      token_buf[0, :length].assign(tokens).realize()
+
+      self._reset_caches(model)
+      expected = model.forward(tokens.reshape(1, -1), 0).realize()
+
+      self._reset_caches(model)
+      actual = model.prefill_forward(token_buf[:, :v_prefill.bind(length)], 0).realize()
+
+      self.assertEqual(expected.item(), actual.item(), f"mismatch at length={length}: forward={expected.item()}, prefill={actual.item()}")
+
+  def test_prefill_then_decode_matches_sequential(self):
+    """Test that prefill + decode produces the same sequence as token-by-token forward."""
+    from tinygrad import UOp, getenv
+    model = self.__class__.model
+    tokens = Tensor.randint(6, high=100, dtype="int32").tolist()
+    num_decode = 3
+
+    # reference: feed all tokens one at a time, then decode
+    self._reset_caches(model)
+    ref_ids = list(tokens)
+    t = Tensor([ref_ids], dtype="int32")
+    t = model.forward(t, 0)
+    ref_ids.append(int(t.item()))
+    for i in range(num_decode - 1):
+      sp = len(ref_ids) - 1
+      t = model.forward(t, sp)
+      ref_ids.append(int(t.item()))
+
+    # test: prefill then decode
+    self._reset_caches(model)
+    test_ids = list(tokens)
+    token_buf = Tensor.zeros(1, model.max_context, dtype="int32").contiguous().realize()
+    token_buf[0, :len(test_ids)].assign(Tensor(test_ids, dtype="int32")).realize()
+    t = model.prefill_forward(token_buf[:, :model.v_prefill.bind(len(test_ids))], 0)
+    test_ids.append(int(t.item()))
+    for i in range(num_decode - 1):
+      sp = len(test_ids) - 1
+      t = model.forward(t, sp)
+      test_ids.append(int(t.item()))
+
+    self.assertEqual(ref_ids, test_ids, f"sequence mismatch: ref={ref_ids[-num_decode:]}, test={test_ids[-num_decode:]}")
+
+  def test_prefill_with_start_pos(self):
+    """Test that prefill with start_pos > 0 matches sequential forward."""
+    model = self.__class__.model
+    first_tokens = Tensor.randint(4, high=100, dtype="int32").tolist()
+    second_tokens = Tensor.randint(3, high=100, dtype="int32").tolist()
+    all_tokens = first_tokens + second_tokens
+
+    # reference: forward all at once from pos 0
+    self._reset_caches(model)
+    expected = model.forward(Tensor([all_tokens], dtype="int32"), 0).realize()
+
+    # test: forward first chunk, then prefill second chunk at start_pos
+    self._reset_caches(model)
+    model.forward(Tensor([first_tokens], dtype="int32"), 0).realize()
+    token_buf = Tensor.zeros(1, model.max_context, dtype="int32").contiguous().realize()
+    token_buf[0, :len(second_tokens)].assign(Tensor(second_tokens, dtype="int32")).realize()
+    actual = model.prefill_forward(token_buf[:, :model.v_prefill.bind(len(second_tokens))],
+                                   model.v_start_pos.bind(len(first_tokens))).realize()
+
+    self.assertEqual(expected.item(), actual.item(),
+                     f"start_pos mismatch: all-at-once={expected.item()}, split={actual.item()}")
+
+  def test_multi_turn_generate(self):
+    """Test that multi-turn generate produces same tokens as reference (token-by-token forward)."""
+    model = self.__class__.model
+    num_decode = 3
+
+    # simulate 2 turns of conversation
+    turn1_tokens = Tensor.randint(5, high=100, dtype="int32").tolist()
+    turn2_tokens = Tensor.randint(4, high=100, dtype="int32").tolist()
+
+    # reference: token-by-token forward (no prefill, no JIT)
+    self._reset_caches(model)
+    ref_ids = list(turn1_tokens)
+    # prefill turn 1
+    t = model.forward(Tensor([ref_ids], dtype="int32"), 0)
+    ref_ids.append(int(t.item()))
+    # decode turn 1
+    for _ in range(num_decode - 1):
+      t = model.forward(t, len(ref_ids) - 1)
+      ref_ids.append(int(t.item()))
+    # prefill turn 2
+    start_pos_2 = len(ref_ids)
+    ref_ids += turn2_tokens
+    t = model.forward(Tensor([ref_ids[start_pos_2:]], dtype="int32"), start_pos_2)
+    ref_ids.append(int(t.item()))
+    # decode turn 2
+    for _ in range(num_decode - 1):
+      t = model.forward(t, len(ref_ids) - 1)
+      ref_ids.append(int(t.item()))
+
+    # test: using prefill_forward with symbolic variable
+    self._reset_caches(model)
+    test_ids = list(turn1_tokens)
+    token_buf = Tensor.zeros(1, model.max_context, dtype="int32").contiguous().realize()
+    # prefill turn 1
+    token_buf[0, :len(test_ids)].assign(Tensor(test_ids, dtype="int32")).realize()
+    t = model.prefill_forward(token_buf[:, :model.v_prefill.bind(len(test_ids))], 0)
+    test_ids.append(int(t.item()))
+    # decode turn 1
+    for _ in range(num_decode - 1):
+      t = model.forward(t, len(test_ids) - 1)
+      test_ids.append(int(t.item()))
+    # prefill turn 2
+    start_pos_2 = len(test_ids)
+    test_ids += turn2_tokens
+    new_tokens = test_ids[start_pos_2:]
+    token_buf[0, :len(new_tokens)].assign(Tensor(new_tokens, dtype="int32")).realize()
+    t = model.prefill_forward(token_buf[:, :model.v_prefill.bind(len(new_tokens))],
+                              model.v_start_pos.bind(start_pos_2))
+    test_ids.append(int(t.item()))
+    # decode turn 2
+    for _ in range(num_decode - 1):
+      t = model.forward(t, len(test_ids) - 1)
+      test_ids.append(int(t.item()))
+
+    self.assertEqual(ref_ids, test_ids, f"multi-turn mismatch:\nref ={ref_ids}\ntest={test_ids}")
+
+  def test_generate_matches_forward(self):
+    """Test that generate() (with symbolic prefill + JIT) matches token-by-token forward for 3 turns."""
+    from tinygrad import TinyJit
+    model = self.__class__.model
+    num_decode = 3
+    turns = [Tensor.randint(5, high=100, dtype="int32").tolist() for _ in range(3)]
+
+    # reference: token-by-token forward
+    self._reset_caches(model)
+    ref_ids: list[int] = []
+    for turn_tokens in turns:
+      start_pos = len(ref_ids)
+      ref_ids += turn_tokens
+      t = model.forward(Tensor([ref_ids[start_pos:]], dtype="int32"), start_pos)
+      ref_ids.append(int(t.item()))
+      for _ in range(num_decode - 1):
+        t = model.forward(t, len(ref_ids) - 1)
+        ref_ids.append(int(t.item()))
+
+    # test: using generate() with symbolic prefill
+    self._reset_caches(model)
+    model.forward_jit = TinyJit(model.forward)
+    model.prefill_jit = TinyJit(model.prefill_forward)
+    test_ids: list[int] = []
+    for turn_tokens in turns:
+      start_pos = len(test_ids)
+      test_ids += turn_tokens
+      gen = model.generate(test_ids, start_pos)
+      for _ in range(num_decode): next(gen)
+
+    self.assertEqual(ref_ids, test_ids, f"generate mismatch:\nref ={ref_ids}\ntest={test_ids}")
+
 if __name__ == '__main__':
   unittest.main()

tinygrad/apps/llm.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-import sys, argparse, typing, re, unicodedata, json, uuid, time, functools
+import sys, argparse, typing, re, unicodedata, json, uuid, time, functools, array
 from tinygrad import Tensor, nn, UOp, TinyJit, getenv, function
 from tinygrad.helpers import partition, DEBUG, Timing, GlobalCounters, stderr_log, colored
 from tinygrad.viz.serve import TCPServerWithReuse, HTTPRequestHandler
@@ -116,7 +116,7 @@ class TransformerBlock:
       self.ffn_up      = nn.Linear(dim, hidden_dim, bias=False)
       self.ffn_down    = nn.Linear(hidden_dim, dim, bias=False)
 
-  def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
+  def _attention_inner(self, x:Tensor, start_pos:int|UOp, mask:Tensor|None) -> Tensor:
     x_norm = self.attn_norm(x)                       # (B,T,D)
     q, k, v = self.attn_q(x_norm), self.attn_k(x_norm), self.attn_v(x_norm)
     if self.qk_norm and self.qk_norm != self.head_dim: q, k = self.attn_q_norm(q), self.attn_k_norm(k)
@@ -131,17 +131,28 @@ class TransformerBlock:
     q = apply_rope(q, freqs_cis)
     k = apply_rope(k, freqs_cis)
 
-    self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(Tensor.stack(k, v))
-    k = self.cache_kv[0, :, :, 0:start_pos+T, :]
-    v = self.cache_kv[1, :, :, 0:start_pos+T, :]
+    # TODO: fix assign to behave like this
+    assigned_kv = self.cache_kv.uop.after(self.cache_kv[:, :, :, start_pos:start_pos+T, :].uop.assign(Tensor.stack(k, v).contiguous().uop))
+    tensor_assigned_kv = Tensor(assigned_kv, device=assigned_kv.device)
+    k = tensor_assigned_kv[0, :, :, 0:start_pos+T, :]
+    v = tensor_assigned_kv[1, :, :, 0:start_pos+T, :]
+
+    #self.cache_kv[:, :, :, start_pos:start_pos+T, :].assign(Tensor.stack(k, v))
+    #k = self.cache_kv[0, :, :, 0:start_pos+T, :]
+    #v = self.cache_kv[1, :, :, 0:start_pos+T, :]
 
-    # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
-    mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device).triu(int(start_pos)+1) if T > 1 else None
     attn = q.scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True)     # (B,H,T,Hd)
     attn = attn.transpose(1, 2).reshape(B, T, -1)                                    # back to (B,T,D)
     attn = self.attn_output(attn)
     return x + attn
 
+  @function
+  def _attention(self, x:Tensor, start_pos:int|UOp) -> Tensor:
+    B, T, _ = x.shape
+    # NOTE: this mask is causal_lower_right, not the causal_upper_left generated by is_casual = True
+    mask = Tensor.full((1, 1, T, start_pos+T), float("-inf"), dtype=x.dtype, device=x.device).triu(int(start_pos)+1) if T > 1 else None
+    return self._attention_inner(x, start_pos, mask)
+
   @function
   def _feed_forward(self, h: Tensor) -> Tensor:
     h_norm = self.ffn_norm(h)
@@ -160,6 +171,16 @@ class TransformerBlock:
       self.cache_kv = Tensor.zeros(2, x.shape[0], self.n_kv_heads, self.max_context, self.head_dim, device=x.device).contiguous().realize()
     return self._feed_forward(self._attention(x, start_pos)).contiguous()
 
+  def prefill(self, x: Tensor, start_pos: int|UOp):
+    """Like __call__ but without @function wrappers, for JIT-compatible symbolic prefill."""
+    if not hasattr(self, "cache_kv"):
+      self.cache_kv = Tensor.zeros(2, x.shape[0], self.n_kv_heads, self.max_context, self.head_dim, device=x.device).contiguous().realize()
+    if not hasattr(self, "causal_mask"):
+      self.causal_mask = Tensor.full((1, 1, self.max_context, self.max_context), float("-inf"), dtype=x.dtype, device=x.device).triu(1).contiguous().realize()
+    T = x.shape[1]
+    mask = self.causal_mask[:, :, start_pos:start_pos+T, :start_pos+T] if (T > 1 if isinstance(T, int) else T.vmin > 1) else None
+    return TransformerBlock._feed_forward.fxn(self, self._attention_inner(x, start_pos, mask)).contiguous()
+
 class Transformer:
   def __init__(self, *, num_blocks, dim, hidden_dim, n_heads, n_kv_heads, norm_eps, vocab_size, head_dim:int, rope_theta:float,
                max_context:int=0, qk_norm:int=0, num_experts:int=0, num_experts_per_tok:int=0):
@@ -169,8 +190,10 @@ class Transformer:
     self.output_norm = nn.RMSNorm(dim, norm_eps)
     self.output = nn.Linear(dim, vocab_size, bias=False)
     self.max_context = max_context
-    # JIT is used if T=1 and start_pos is a UOp. TODO: make this not needed by including T in the JIT and making start_pos always a UOp
     self.forward_jit = TinyJit(self.forward)
+    self.prefill_jit = TinyJit(self.prefill_forward)
+    self.v_start_pos = UOp.variable("start_pos", 1, max_context-1)
+    self.v_prefill = UOp.variable("prefill", 2, max_context)
 
   def forward(self, tokens:Tensor, start_pos:int|UOp) -> Tensor:
     x = self.token_embd(tokens)                           # (B, T, D)
@@ -179,10 +202,12 @@ class Transformer:
     return self.output(self.output_norm(x))[:, -1, :].softmax(-1, dtype="float").argmax(-1, keepdim=True)
 
   def __call__(self, tokens:Tensor, start_pos:int|UOp=0) -> Tensor:
-    return (self.forward_jit if getenv("JIT", 1) and tokens.shape[1] == 1 and isinstance(start_pos, UOp) else self.forward)(tokens, start_pos)
+    if getenv("JIT", 1) and tokens.shape[1] == 1 and isinstance(start_pos, UOp): return self.forward_jit(tokens, start_pos)
+    if getenv("JIT", 1) and tokens.shape[1] != 1 and isinstance(tokens.shape[1], UOp): return self.prefill_jit(tokens, start_pos)
+    return self.forward(tokens, start_pos)
 
   @staticmethod
-  def from_gguf(gguf:Tensor, max_context:int|None=None, realize=True) -> tuple[Transformer, dict]:
+  def from_gguf(gguf:Tensor, max_context:int|None=None, realize=bool(getenv("REALIZE", 1))) -> tuple[Transformer, dict]:
     # TODO: remove the need for copy to default device
     kv, state_dict = nn.state.gguf_load(gguf.to(None))
 
@@ -212,15 +237,38 @@ class Transformer:
                         num_experts=kv.get(f'{arch}.expert_count', 0), num_experts_per_tok=kv.get(f'{arch}.expert_used_count', 0))
     nn.state.load_state_dict(model, state_dict, verbose=False, consume=True, realize=False)  # NOTE: rope_freqs.weight (32,) is unused
     # NOTE: without this contiguous, it unpacks the weights from the model every time. we shouldn't need this, but for now it's faster
-    for s in (params:=nn.state.get_parameters(model)): s.replace(s.contiguous())
-    if realize: Tensor.realize(*params)
+    if realize:
+      for s in (params:=nn.state.get_parameters(model)): s.replace(s.contiguous())
+      Tensor.realize(*params)
     return model, kv
 
+  def prefill_forward(self, tokens:Tensor, start_pos:int|UOp) -> Tensor:
+    """Like forward but without @function on attention/ffn, for symbolic prefill."""
+    x = self.token_embd(tokens)
+    for block in self.blk: x = block.prefill(x, start_pos)
+    return self.output(self.output_norm(x))[:, -1, :].softmax(-1, dtype="float").argmax(-1, keepdim=True)
+
   def generate(self, tokens:list[int], start_pos=0):
-    v_start_pos = UOp.variable("start_pos", 1, self.max_context-1)
-    t = Tensor([tokens[start_pos:]], dtype="int32")
+    if not hasattr(self, "_prefill_buf"):
+      self._prefill_buf = Tensor.zeros(1, self.max_context, dtype="int32").contiguous().realize()
+      self._prefill_staging = array.array('i', [0] * self.max_context)
+    # prefill all tokens at once using symbolic variable
+    num_tokens = len(tokens) - start_pos
+    if num_tokens >= 2 and getenv("SYM", 1):
+      # write tokens directly into the buffer via copyin (no scheduled kernels, no cache misses)
+      for i, t in enumerate(tokens[start_pos:]): self._prefill_staging[i] = t
+      self._prefill_buf._buffer().copyin(memoryview(self._prefill_staging))
+      v_sp = self.v_start_pos.bind(start_pos) if start_pos != 0 else start_pos
+      t = self(self._prefill_buf[:, :self.v_prefill.bind(num_tokens)], v_sp)
+      next_id = int(t.item())
+      tokens.append(next_id)
+      start_pos = len(tokens) - 1
+      yield next_id
+    else:
+      t = Tensor([tokens[start_pos:]], dtype="int32")
+    # decode one token at a time
     while len(tokens) < self.max_context:
-      t = self(t, v_start_pos.bind(start_pos) if getenv("SYM", 1) and start_pos != 0 and t.shape[-1] == 1 else start_pos)
+      t = self(t, self.v_start_pos.bind(start_pos) if getenv("SYM", 1) and start_pos != 0 and t.shape[-1] == 1 else start_pos)
       next_id = int(t.item())
       tokens.append(next_id)
       start_pos = len(tokens) - 1
@@ -329,27 +377,43 @@ class Handler(HTTPRequestHandler):
 
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
-  parser.add_argument("--model", choices=list(models.keys()), default=list(models.keys())[0], help="Model choice")
+  parser.add_argument("--model", "-m", choices=list(models.keys()), default=list(models.keys())[0], help="Model choice")
   parser.add_argument("--max_context", type=int, default=4096, help="Max Context Length")
   parser.add_argument("--serve", nargs='?', type=int, const=11434, metavar="PORT", help="Run OpenAI compatible API (optional port, default 11434)")
   parser.add_argument("--benchmark", nargs='?', type=int, const=20, metavar="COUNT", help="Benchmark tok/s (optional count, default 20)")
   args = parser.parse_args()
 
   # load the model
-  model, kv = Transformer.from_gguf(Tensor.from_url(models[args.model]), args.max_context)
-  if DEBUG >= 1: print(f"using model {args.model}")
+  raw_model = Tensor.from_url(models[args.model])
+  model, kv = Transformer.from_gguf(raw_model, args.max_context)
+  if DEBUG >= 1 or args.benchmark:
+    print(f"using model {args.model} with {raw_model.nbytes():,} bytes and {sum(x.numel() for x in nn.state.get_parameters(model)):,} params")
+  del raw_model
+
+  # TODO: why this is required to free the RAM of the GGUF copy?
+  import gc
+  gc.collect()
 
   # do benchmark
   if args.benchmark:
-    param_bytes = sum(x.nbytes() for x in nn.state.get_parameters(model))
-    for b in model.blk:
-      if hasattr(b, 'ffn_gate_exps'):
-        expert_bytes = b.ffn_gate_exps.weight.nbytes() + b.ffn_up_exps.weight.nbytes() + b.ffn_down_exps.weight.nbytes()
-        param_bytes -= int(expert_bytes * (1 - b.num_experts_per_tok / b.ffn_gate_exps.weight.shape[0]))
+    # prefill benchmark
+    token_buf = Tensor.zeros(1, args.max_context, dtype="int32").contiguous().realize()
+    print("prefill benchmark:")
+    for length in [4, 16, 64, 256, 1024]:
+      if length > args.max_context: break
+      GlobalCounters.reset()
+      with Timing(prefix=f"  prefill {length:4d}: ",
+                  on_exit=lambda x, length=length: f", {length*1e9/x:6.2f} tok/s, {GlobalCounters.global_mem/x:7.2f} GB/s,"
+                  f" {GlobalCounters.global_mem//1000000}/{GlobalCounters.mem_used//1000000} MB"):
+        model.prefill_jit(token_buf[:, :model.v_prefill.bind(length)], 0).realize()
+
+    # generation benchmark
+    print("generation benchmark:")
     gen = model.generate([0], 0)
     for _ in range(args.benchmark):
       GlobalCounters.reset()
-      with Timing(on_exit=lambda x: f", {1e9/x:6.2f} tok/s, {GlobalCounters.global_mem/x:7.2f} GB/s, param {param_bytes/x:7.2f} GB/s"): next(gen)
+      with Timing(on_exit=lambda x: f", {1e9/x:6.2f} tok/s, {GlobalCounters.global_mem/x:7.2f} GB/s,"
+                  f" {GlobalCounters.global_mem//1000000}/{GlobalCounters.mem_used//1000000} MB"): next(gen)
     exit(0)
 
   # extract some metadata

tinygrad/codegen/gpudims.py

@@ -4,12 +4,11 @@ from tinygrad.helpers import all_int, dedup, get_contraction
 from tinygrad.dtype import dtypes, AddrSpace, Invalid
 from tinygrad.renderer import Renderer
 
+def _dim_max(d:sint) -> int: return d if isinstance(d, int) else d.vmax
 def _group_dims(dims:tuple[sint, ...], max_sizes:tuple[int, ...]):
-  # TODO: symbolic shape
-  if not all_int(dims): return dims
-  while len(dims) > len(max_sizes) or any(d > m for d,m in zip(dims, max_sizes)):
+  while len(dims) > len(max_sizes) or any(_dim_max(d) > m for d,m in zip(dims, max_sizes)):
     for i,m in enumerate(max_sizes):
-      if i < (len(dims)-1) and dims[i] * dims[i+1] <= m:
+      if i < (len(dims)-1) and _dim_max(dims[i]) * _dim_max(dims[i+1]) <= m:
         dims = dims[:i] + (dims[i]*dims[i+1],) + dims[i+2:]
         break
     else: return None

tinygrad/codegen/late/devectorizer.py

@@ -1,7 +1,7 @@
 from typing import Any, cast
 import functools, itertools
 from collections import defaultdict
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from tinygrad.dtype import dtypes, ImageDType, DType, AddrSpace, Invalid, PtrDType
 from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, GroupOp, identity_element
 from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate
@@ -308,8 +308,6 @@ pm_render = PatternMatcher([
 @dataclass
 class ReduceContext:
   acc_num: int = 0
-  # track ENDs by range for merging parallel reduces
-  range_to_ends: dict[tuple[UOp, ...], list[UOp]] = field(default_factory=dict)
 
 def horizontal_reduce(inp:UOp, out_dtype:DType) -> list[UOp]:
   # if this has a horizontal reduction component, do that first
@@ -335,13 +333,15 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp):
     ctx.acc_num += 1
   ret = functools.reduce(lambda x,y: x.alu(red.arg, y), lst)
   if len(reduce_range) == 0: return ret
-  end = acc.index(UOp.const(dtypes.int, 0)).store(ret).end(*reduce_range)
-  ctx.range_to_ends.setdefault(reduce_range, []).append(end)
+  end = acc.index(UOp.const(dtypes.int, 0)).store(ret).end(*reduce_range).rtag("mergeable")
   return acc.after(end).index(UOp.const(dtypes.int, 0))
 
 def merge_reduce_ends(ctx:ReduceContext, sink:UOp):
-  # merge ENDs that share the same range
-  subs = {e: UOp.group(*(e.src[0] for e in ends)).end(*r) for r, ends in ctx.range_to_ends.items() if len(ends) > 1 for e in ends}
+  # merge ENDs that share the same range (only those created by reduce_to_acc)
+  range_to_ends: dict[tuple[UOp, ...], list[UOp]] = {}
+  for u in sink.backward_slice:
+    if u.op is Ops.END and u.tag == "mergeable": range_to_ends.setdefault(u.src[1:], []).append(u)
+  subs = {e: UOp.group(*(e.src[0] for e in ends)).end(*r) for r, ends in range_to_ends.items() if len(ends) > 1 for e in ends}
   return sink.substitute(subs) if subs else None
 
 pm_reduce = PatternMatcher([

tinygrad/engine/allocations.py

@@ -25,7 +25,9 @@ def disk_copy_is_buffer(ctx:AllocCtx, u:UOp):
   if from_creation: return tag_uop(ctx, u)
 
 def apply_after(ctx:AllocCtx, u:UOp):
-  ctx.buffer_map[u] = u.src[0]
+  base = u.src[0]
+  while base.op is Ops.AFTER: base = base.src[0]
+  ctx.buffer_map[u] = base
 
 # CONTIGUOUS and ASSIGN + parents are the only nodes that get updated
 add_tags = PatternMatcher([
@@ -54,7 +56,7 @@ def replace_contig_with_assign(u:UOp):
 
 def replace_assign_with_contig(u:UOp):
   assigned_to = u
-  while assigned_to.op in {Ops.ASSIGN, Ops.BITCAST}: assigned_to = assigned_to.src[0].base
+  while assigned_to.op in {Ops.ASSIGN, Ops.BITCAST, Ops.AFTER}: assigned_to = assigned_to.src[0].base
   if assigned_to.op is not Ops.BUFFER:
     return u.src[1].contiguous(tag=u.tag)
 
@@ -74,8 +76,9 @@ pm_early_transform_tensor_graph = PatternMatcher([
   (UPat(GroupOp.ALU, name="alu"), lambda ctx,alu: alu.replace(src=new_src) if (new_src:=tuple(ctx.get(s, s) for s in alu.src)) != alu.src else None),
   # add CONTIGUOUS to tagged UOps
   (UPat(GroupOp.All-{Ops.CONTIGUOUS, Ops.ASSIGN}, name="x"), lambda x: x.rtag(None).contiguous(tag=x.tag) if x.tag else x.replace(tag=None)),
-  # remove extra CONTIGUOUS on ASSIGN
-  (UPat(Ops.CONTIGUOUS, src=(UPat(Ops.ASSIGN, name="a"),), name="c"), lambda a,c: a.replace(tag=a.tag+c.tag)),
+  # remove extra CONTIGUOUS on ASSIGN (only when assign target is contiguous)
+  (UPat(Ops.CONTIGUOUS, src=(UPat(Ops.ASSIGN, name="a"),), name="c"),
+   lambda a,c: a.replace(tag=a.tag+c.tag) if a.src[0].has_buffer_identity() else None),
   # replace ASSIGN with CONTIGUOUS
   (UPat(Ops.ASSIGN, name="u"), replace_assign_with_contig),
   # replace CONTIGUOUS with ASSIGNs

tinygrad/engine/schedule.py

@@ -148,7 +148,7 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[list[ExecItem], di
       nm = b.src[0].expr
       if nm not in used_vars: continue
       val = b.src[1].arg
-      assert nm not in var_vals or var_vals[nm] == val, f"bind mismatch on {nm}, {var_vals[nm]} != {val}"
+      if var_vals.get(nm, val) != val: raise RuntimeError(f"bind mismatch on {nm}, {var_vals[nm]} != {val}")
       var_vals[nm] = val
 
   # convert LINEAR to ExecItems

tinygrad/function.py

@@ -10,7 +10,7 @@ def add_to_ctx(ctx, x:UOp):
   return ret
 
 pm_ctx = PatternMatcher([
-  (UPat(Ops.BUFFER, name="x"), add_to_ctx),
+  (UPat((Ops.BUFFER, Ops.BIND), name="x"), add_to_ctx),
   (UPat((Ops.ASSIGN, Ops.CONTIGUOUS), name="x"),
    lambda ctx,x: add_to_ctx(ctx,x) if not x.op_in_backward_slice_with_self(Ops.PARAM) else None),
 ])

tinygrad/gradient.py

@@ -17,12 +17,13 @@ def call_gradient(ctx:UOp, k:UOp) -> tuple[UOp|None, ...]:
   if k.arg.grad_fxn is not None: return (None,) + k.arg.grad_fxn(ctx, k)
   # auto-differentiate the function
   fxn, args = k.src[0], k.src[1:]
-  params = sorted([x for x in fxn.toposort() if x.op == Ops.PARAM], key=lambda x: x.arg)
-  grads = compute_gradient(fxn, ctx.param_like(len(args)), set(params))
+  params = {x.arg:x for x in fxn.toposort(enter_calls=False) if x.op == Ops.PARAM}
+  grads = compute_gradient(fxn, ctx.param_like(len(args)), set(params.values()))
   ret: list[UOp|None] = [None]
-  for i,p in enumerate(params):
-    if p in grads:
+  for i in range(len(args)):
+    if (p:=params.get(i, None)) is not None and p in grads:
       # TODO: compact the args and remove unused ones
+      assert not grads[p].op_in_backward_slice_with_self(Ops.BUFFER), "BUG: BUFFER in backward slice of grad"
       ret.append(grads[p].call(*args, ctx, name=(k.arg.name or "")+f"_backward_{i}"))
     else:
       ret.append(None)

tinygrad/renderer/amd/sqtt.py

@@ -46,6 +46,7 @@ class InstOp(Enum):
   SMEM = 0x1
   JUMP = 0x3              # branch taken
   JUMP_NO = 0x4           # branch not taken
+  CALL = 0x5              # s_call_b64
   MESSAGE = 0x9
   VALU_TRANS = 0xb        # transcendental: exp, log, rcp, sqrt, sin, cos
   VALU_64_SHIFT = 0xd     # 64-bit shifts: lshl, lshr, ashr
@@ -72,8 +73,10 @@ class InstOp(Enum):
 
   # LDS ops on traced SIMD
   LDS_LOAD = 0x29
+  LDS_ATOMIC = 0x2a        # ds_append, ds_consume, ds_store_addtid_b32
   LDS_STORE = 0x2b
   LDS_STORE_64 = 0x2c
+  LDS_STORE_96 = 0x2d
   LDS_STORE_128 = 0x2e
 
   # Memory ops on other SIMD (0x5x range)
@@ -99,17 +102,27 @@ class InstOp(Enum):
 
 class InstOpRDNA4(Enum):
   """SQTT instruction operation types for RDNA4 (gfx1200). Different encoding from RDNA3."""
-  # TODO: we need to do discovery of all of these from instructions
   SALU = 0x0
   JUMP = 0x1
   NEXT = 0x2
   MESSAGE = 0x4
+  VALU_TRANS = 0x5
   VALU_64 = 0x6
+  VALU_MAD64 = 0x7
+  VINTERP = 0x9
   VALU_WMMA = 0x46
   VMEM = 0x10
   VMEM_128 = 0x11
   VMEM_STORE = 0x12
-  VMEM_STORE_128 = 0x14
+  VMEM_STORE_G96 = 0x13   # global_store_[b96,b128]
+  LDS_LOAD = 0x14
+  LDS_STORE = 0x15
+  LDS_STORE_64 = 0x16
+  LDS_STORE_128 = 0x17
+  VALU_F64 = 0x49
+  SALU_TRANS = 0x4c       # transcendental with sgpr src/dst
+  SALU_MUL = 0x4d         # s_[mul,mulhi,mulk]
+  SALU_MUL64 = 0x4e
   OTHER_VMEM = 0x5e
   OTHER_VMEM_STORE = 0x60
 
@@ -147,11 +160,6 @@ class TS_DELTA_S8_W3(PacketType):
   delta = bits[10:8]
   _padding = bits[63:11]
 
-class TS_DELTA_S8_W3_RDNA4(PacketType):  # Layout 4: 64->72 bits
-  encoding = bits[6:0] == 0b0100001
-  delta = bits[10:8]
-  _padding = bits[71:11]
-
 class TS_DELTA_S5_W3(PacketType):
   encoding = bits[4:0] == 0b00110
   delta = bits[7:5]
@@ -363,7 +371,7 @@ PACKET_TYPES_RDNA3: dict[int, type[PacketType]] = {
 }
 PACKET_TYPES_RDNA4: dict[int, type[PacketType]] = {
   **PACKET_TYPES_RDNA3,
-  7: TS_DELTA_S8_W3_RDNA4, 9: WAVESTART_RDNA4, 10: TS_DELTA_S5_W2_RDNA4, 11: WAVEALLOC_RDNA4,
+  9: WAVESTART_RDNA4, 10: TS_DELTA_S5_W2_RDNA4, 11: WAVEALLOC_RDNA4,
   12: TS_DELTA_S5_W3_RDNA4, 13: PERF_RDNA4, 22: TS_DELTA_OR_MARK_RDNA4, 24: INST_RDNA4,
 }
 
@@ -536,8 +544,9 @@ def decode(data: bytes) -> Iterator[PacketType]:
     if nib_off: reg, pos = (reg >> 4) | ((data[pos] >> 4) << 60), pos + 1
     # 2. read all full bytes at once
     if (byte_count := need >> 1):
-      chunk = int.from_bytes(data[pos:pos + byte_count], 'little')
-      reg, pos = (reg >> (byte_count * 8)) | (chunk << (64 - byte_count * 8)), pos + byte_count
+      read_bytes = min(byte_count, 8)
+      chunk = int.from_bytes(data[pos:pos + read_bytes], 'little')
+      reg, pos = (reg >> (read_bytes * 8)) | (chunk << (64 - read_bytes * 8)), pos + byte_count
     # 3. if odd, read low nibble
     if (nib_off := need & 1): reg = (reg >> 4) | ((data[pos] & 0xF) << 60)
 
@@ -666,8 +675,9 @@ def print_packets(packets) -> None:
   from tinygrad.helpers import getenv
   skip = {"NOP", "TS_DELTA_SHORT", "TS_WAVE_STATE", "TS_DELTA_OR_MARK",
           "TS_DELTA_S5_W2", "TS_DELTA_S5_W3", "TS_DELTA_S8_W3", "REG", "EVENT"} if not getenv("NOSKIP") else {"NOP"}
-  for p in packets:
-    if type(p).__name__.replace("_RDNA4", "") not in skip: print(format_packet(p))
+  for data in packets:
+    p, inst = data if isinstance(data, tuple) else (data, None)
+    if type(p).__name__.replace("_RDNA4", "") not in skip: print(format_packet(p), f"inst={inst.inst}" if inst is not None else '')
 
 if __name__ == "__main__":
   import sys, pickle
@@ -676,8 +686,10 @@ if __name__ == "__main__":
     sys.exit(1)
   with open(sys.argv[1], "rb") as f:
     data = pickle.load(f)
-  prg_names = {e.tag: e.name for e in data if type(e).__name__ == "ProfileProgramEvent" and e.tag is not None}
+  prg_events = {e.tag: e for e in data if type(e).__name__ == "ProfileProgramEvent" and e.tag is not None}
   sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
+  dev_targets = {e.device:f"gfx{e.props['gfx_target_version']//1000}" for e in data if type(e).__name__ == "ProfileDeviceEvent" and e.props}
   for i, event in enumerate(sqtt_events):
-    print(f"\n=== event {i} {prg_names.get(event.kern, '')} ===")
-    print_packets(decode(event.blob))
+    prg = prg_events.get(event.kern)
+    print(f"\n=== event {i} {prg.name if prg is not None else ''} ===")
+    print_packets(map_insts(event.blob, prg.lib, dev_targets[prg.device]) if prg is not None else decode(event.blob))

tinygrad/runtime/support/am/amdev.py

@@ -193,7 +193,7 @@ class AMDev(PCIDevImplBase):
     if DEBUG >= 2: print(f"am {self.devfmt}: boot done")
 
   def init_sw(self, smi_dev=False):
-    self.smi_dev, self.is_err_state, self.has_aql_queue = smi_dev, False, False
+    self.smi_dev, self.is_err_state = smi_dev, False
 
     # Memory manager & firmware
     self.mm = AMMemoryManager(self, self.vram_size - self.reserved_vram_size, boot_size=(32 << 20), pt_t=AMPageTableEntry, va_shifts=[12, 21, 30, 39],
@@ -226,7 +226,7 @@ class AMDev(PCIDevImplBase):
     self.reg("regSCRATCH_REG6").write(self.is_err_state) # set finalized state.
 
   def recover(self) -> bool:
-    if (self.has_aql_queue and self.is_hive()) or not self.is_err_state: return False # TODO: support aql queue recovery on hive
+    if not self.is_err_state: return False
     if DEBUG >= 2: print(f"am {self.devfmt}: Start recovery")
     self.ih.interrupt_handler()
     self.gfx.reset_mec()

tinygrad/runtime/support/am/ip.py

@@ -243,7 +243,7 @@ class AM_GFX(AM_IP):
     while self.adev.regCP_STAT.read() != 0 and self.adev.regRLC_RLCS_BOOTLOAD_STATUS.read_bitfields()['bootload_complete'] != 0: pass
 
     self.adev.gmc.init_hub("GC", inst_cnt=self.xccs)
-    if self.adev.partial_boot: return
+    if self.adev.partial_boot: return self.reset_mec()
 
     self._config_mec()
 
@@ -291,18 +291,22 @@ class AM_GFX(AM_IP):
 
   def reset_mec(self):
     self._dequeue_hqds(reset=True)
+
+    # issue a soft reset to reset aql sync counter on multixcc systems.
+    if self.xccs > 1:
+      for xcc in range(self.xccs): self.adev.regGRBM_SOFT_RESET.write(soft_reset_cp=1, soft_reset_gfx=1, inst=xcc)
+      time.sleep(0.05)
+      for xcc in range(self.xccs): self.adev.regGRBM_SOFT_RESET.write(0x0, inst=xcc)
+
     self._config_mec()
     self._enable_mec()
 
   def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, eop_addr:int, eop_size:int, idx:int, aql:bool) -> tuple[int, int]:
-    self.adev.has_aql_queue |= aql
     pipe, queue, doorbell = idx // 4, idx % 4, am.AMDGPU_NAVI10_DOORBELL_MEC_RING0
-    self._grbm_select(me=1, pipe=pipe, queue=queue, inst=0)
-    restore_queue = aql and self.xccs > 1 and self.adev.partial_boot and (self.adev.regCP_HQD_ACTIVE.read(inst=0) & 1)
-    restore_ptr = (self.adev.regCP_HQD_PQ_WPTR_LO.read(inst=0) | (self.adev.regCP_HQD_PQ_WPTR_HI.read(inst=0) << 32)) if restore_queue else 0
-    if DEBUG >= 2 and restore_queue: print(f"am {self.adev.devfmt}: GFX queue already active, continuing from saved state {restore_ptr=:#x}.")
 
     for xcc in range(self.xccs if aql else 1):
+      self._grbm_select(me=1, pipe=pipe, queue=queue, inst=xcc)
+
       struct_t = getattr(am, f"struct_v{self.adev.ip_ver[am.GC_HWIP][0]}{'_compute' if self.adev.ip_ver[am.GC_HWIP][0] >= 10 else ''}_mqd")
       mqd_struct = struct_t(header=0xC0310800, cp_mqd_base_addr_lo=lo32(self.mqd_mc[queue] + 0x1000*xcc),
         cp_mqd_base_addr_hi=hi32(self.mqd_mc[queue] + 0x1000*xcc), cp_hqd_pipe_priority=0x2, cp_hqd_queue_priority=0xf, cp_hqd_quantum=0x111,
@@ -320,26 +324,16 @@ class AM_GFX(AM_IP):
         **({'compute_tg_chunk_size':1, 'compute_current_logic_xcc_id':xcc, 'cp_mqd_stride_size':0x1000} if aql and self.xccs > 1 else {}))
       for se in range(8 if self.adev.ip_ver[am.GC_HWIP][0] >= 10 else 4): setattr(mqd_struct, f'compute_static_thread_mgmt_se{se}', 0xffffffff)
 
-      # Copy mqd into memory
-      self._grbm_select(me=1, pipe=pipe, queue=queue, inst=xcc)
+      self.adev.vram.view(self.mqd_paddr[queue] + 0x1000*xcc, ctypes.sizeof(mqd_struct))[:] = memoryview(mqd_struct).cast('B')
 
-      if restore_queue:
-        for r in [self.adev.regCP_HQD_PQ_RPTR_REPORT_ADDR, self.adev.regCP_HQD_EOP_BASE_ADDR, self.adev.regCP_HQD_EOP_BASE_ADDR_HI,
-                  self.adev.regCP_HQD_PQ_RPTR_REPORT_ADDR_HI, self.adev.regCP_HQD_PQ_WPTR_POLL_ADDR, self.adev.regCP_HQD_PQ_WPTR_POLL_ADDR_HI]:
-          val = memoryview(bytes(mqd_struct)).cast('I')[0x80 + (off:=r.addr[xcc] - self.adev.regCP_MQD_BASE_ADDR.addr[xcc])]
-          self.adev.vram.view(self.mqd_paddr[queue] + 0x1000*xcc, ctypes.sizeof(mqd_struct), fmt='I')[0x80 + off] = val
-          r.write(val, inst=xcc)
-      else:
-        self.adev.vram.view(self.mqd_paddr[queue] + 0x1000*xcc, ctypes.sizeof(mqd_struct))[:] = memoryview(mqd_struct).cast('B')
-
-        mqd_st_mv = to_mv(ctypes.addressof(mqd_struct), ctypes.sizeof(mqd_struct)).cast('I')
-        for i, reg in enumerate(range(self.adev.regCP_MQD_BASE_ADDR.addr[xcc], self.adev.regCP_HQD_PQ_WPTR_HI.addr[xcc] + 1)):
-          self.adev.wreg(reg, mqd_st_mv[0x80 + i])
-        self.adev.regCP_HQD_ACTIVE.write(0x1, inst=xcc)
+      mqd_st_mv = to_mv(ctypes.addressof(mqd_struct), ctypes.sizeof(mqd_struct)).cast('I')
+      for i, reg in enumerate(range(self.adev.regCP_MQD_BASE_ADDR.addr[xcc], self.adev.regCP_HQD_PQ_WPTR_HI.addr[xcc] + 1)):
+        self.adev.wreg(reg, mqd_st_mv[0x80 + i])
+      self.adev.regCP_HQD_ACTIVE.write(0x1, inst=xcc)
 
       self.adev.gmc.flush_hdp()
       self._grbm_select(inst=xcc)
-    return restore_ptr // 16, doorbell
+    return 0, doorbell
 
   def set_clockgating_state(self):
     if hasattr(self.adev, 'regMM_ATC_L2_MISC_CG'): self.adev.regMM_ATC_L2_MISC_CG.write(enable=1, mem_ls_enable=1)
@@ -391,14 +385,12 @@ class AM_GFX(AM_IP):
         _config_helper(eng_name="MEC", cntl_reg="MEC_RS64", eng_reg="MEC_RS64", pipe_cnt=1, me=1, xcc=xcc)
 
   def _dequeue_hqds(self, reset=False):
-    # NOTE: For aqls with xccs (queue=1), will continue from the saved state.
-    for q in range(2 if self.xccs == 1 else 1):
+    for q in range(2):
       for xcc in range(self.xccs):
         self._grbm_select(me=1, pipe=0, queue=q, inst=xcc)
         if self.adev.regCP_HQD_ACTIVE.read(inst=xcc) & 1:
           self.adev.regCP_HQD_DEQUEUE_REQUEST.write(0x2, inst=xcc) # 1 - DRAIN_PIPE; 2 - RESET_WAVES
-          if reset: self.adev.regSPI_COMPUTE_QUEUE_RESET.write(1, inst=xcc)
-          else: wait_cond(lambda: self.adev.regCP_HQD_ACTIVE.read(inst=xcc) & 1, value=0, msg="HQD dequeue timeout")
+          if not reset: wait_cond(lambda: self.adev.regCP_HQD_ACTIVE.read(inst=xcc) & 1, value=0, msg="HQD dequeue timeout")
     self._grbm_select()
 
 class AM_IH(AM_IP):
@@ -459,6 +451,7 @@ class AM_IH(AM_IP):
         bf = self.adev.reg(self.adev.gmc.pf_status_reg('GC')).read_bitfields()
         va = (self.adev.reg('regGCVM_L2_PROTECTION_FAULT_ADDR_HI32').read()<<32) | self.adev.reg('regGCVM_L2_PROTECTION_FAULT_ADDR_LO32').read()
         print(f"am {self.adev.devfmt}: GCVM_L2_PROTECTION_FAULT_STATUS: {bf} {va<<12:#x}")
+        self.adev.reg('regGCVM_L2_PROTECTION_FAULT_CNTL').update(clear_protection_fault_status_addr=1)
         self.adev.is_err_state = True
       else: self.adev.is_err_state = True
 

tinygrad/schedule/indexing.py

@@ -56,7 +56,7 @@ class IndexingContext:
     return UOp.range(s, next(self.range_idx), axistype) if resolve(s!=1) else UOp.const(dtypes.index, 0)
 
 def create_bufferize_and_index_based_on_ranges(ctx:IndexingContext, x:UOp):
-  if x.op in {Ops.BUFFERIZE, Ops.INDEX, Ops.AFTER}: return None
+  if x.op in {Ops.BUFFERIZE, Ops.INDEX}: return None
   new_srcs = []
   for s in x.src:
     new_src = s
@@ -179,6 +179,9 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     # no ranges on kernels, they are internal
     if x.op in {Ops.CALL, Ops.LINEAR}: continue
 
+    # no range on after
+    if x.op is Ops.AFTER: continue
+
     # treat MSTACK/MSELECT like SINK
     if x.op in {Ops.MSTACK, Ops.MSELECT}: continue
 

tinygrad/schedule/multi.py

@@ -164,10 +164,18 @@ def passthrough_multi(root:UOp, multi:UOp):
   return UOp(root.op, root.dtype, (multi.src[0],)+tuple(x.src[0] if x.op is Ops.MULTI else x for x in root.src[1:]), root.arg).multi(multi.axis)
 
 def rewrite_into_call(call:UOp):
-  return call.replace(src=(graph_rewrite(call.src[0], multi_pm, name="subcall"),)+call.src[1:]) if should_resolve_call(call) else None
+  if not should_resolve_call(call): return None
+  new_body = graph_rewrite(call.src[0], multi_pm, name="subcall")
+  new_args = tuple(a.src[0] if a.op is Ops.MULTI else a for a in call.src[1:])
+  return call.replace(src=(new_body,)+new_args)
+
+def param_to_multi(p:UOp):
+  if p.axis is None: return None
+  return UOp.param(p.arg, p.dtype, p.shard_shape, p._device).multi(p.axis)
 
 # NOTE: this is the same pattern as Ops.UNROLL
 multi_pm = PatternMatcher([
+  (UPat(Ops.PARAM, name="p"), param_to_multi),
   (UPat(GroupOp.ALU, name="root", custom_early_reject=set([Ops.MULTI])), alu_multi),
   (UPat(Ops.REDUCE_AXIS, src=(UPat(Ops.MULTI, name="multi"), ), name="root"), reduce_multi),
   (UPat(Ops.RESHAPE, src=(UPat(Ops.MULTI, name="multi"), UPat()), name="root"), reshape_multi),

tinygrad/schedule/rangeify.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass, field, replace
 import itertools
-from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
+from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace, Invalid
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, KernelInfo
 from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches, should_resolve_call
 from tinygrad.uop.symbolic import symbolic
@@ -90,7 +90,8 @@ def resolve_call(c:UOp, allow_param_mismatch=True) -> UOp|None:
 
   dict_map = {x:args[x.arg] for x in params}
   for i, (p, a) in enumerate(dict_map.items()):
-    if p.shape != a.shape: raise TypeError(f"arg {i} shape mismatch: expected {p.shape}, got {a.shape}")
+    if p.axis != a.axis: raise TypeError(f"arg {i} axis mismatch: expected {p.axis}, got {a.axis}")
+    if p.max_shape != a.max_shape: raise TypeError(f"arg {i} shape mismatch: expected {p.shape}, got {a.shape}")
     if p.dtype != a.dtype: raise TypeError(f"arg {i} dtype mismatch: expected {p.dtype}, got {a.dtype}")
   return c.src[0].substitute(dict_map, walk=True)
 
@@ -229,8 +230,9 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
 
   # if it makes it here, the bufferize is removed
   # this is the ranges replaced
-  # NOTE: if buf src is a const, we don't replace it
-  return src.substitute({k:v for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST}, extra_pm=pm_gate_substitute)
+  # NOTE: if buf src is a const, we don't replace it. if idx is Invalid (dead load), don't replace it either
+  replaced = {k:v for k,v in zip(buf.src[1:], idx.src[1:]) if k.op is not Ops.CONST and not (v.op is Ops.CONST and v.arg is Invalid)}
+  return src.substitute(replaced, extra_pm=pm_gate_substitute)
 
 def remove_noop_bufferize(idx,b2):
   if idx.src[1:] != b2.src[1:] or idx.src[0].op is Ops.BUFFER_VIEW: return None
@@ -364,6 +366,11 @@ pm_add_buffers = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([
 
   # remove any RESHAPEs on KERNEL
   (UPat(Ops.CALL, name="k"), lambda k: k.replace(src=tuple(x.src[0] if x.op is Ops.RESHAPE else x for x in k.src))),
+
+  # remove MOP on AFTER
+  (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(GroupOp.Movement, name="y"))), lambda x,y: x.after(y.src[0])),
+  # remove double AFTER
+  (UPat(Ops.AFTER, src=(UPat.var("x"), UPat(Ops.AFTER, name="y"))), lambda x,y: x.after(*y.src[1:]))
 ])
 
 pm_add_buffers_local = pm_mops+pm_flatten_bufferize+to_bufferview+PatternMatcher([

tinygrad/tensor.py

@@ -25,8 +25,7 @@ def canonicalize_device(device:str|tuple|list|None) -> str|tuple[str, ...]:
 # *** all in scope Tensors are here. this gets relevant UOps ***
 
 all_tensors: dict[weakref.ref[Tensor], None] = {}
-_pending_assigns: dict[UOp, list[UOp]] = {}  # buffer_uop -> [assign_uops in insertion order]
-def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str) -> None:
+def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str, walk:bool=False) -> None:
   with cpu_profile(TracingKey(name), "TINY"):
     # get tensors in scope
     in_scope: dict[UOp, bool] = {}
@@ -35,7 +34,7 @@ def _apply_map_to_tensors(applied_map:dict[UOp, UOp], name:str) -> None:
 
     # get all Tensors and apply the map
     sink = UOp.sink(*[t.uop for t in scope_tensors])
-    new_sink = sink.substitute(applied_map, name=f"substitute {name}")
+    new_sink = sink.substitute(applied_map, name=f"substitute {name}", walk=walk)
 
     # set the relevant uop to the realized UOps
     for t,s,ns in zip(scope_tensors, sink.src, new_sink.src):
@@ -278,23 +277,6 @@ class Tensor(OpMixin):
   @disable_gc()
   def realize(self, *lst:Tensor, do_update_stats=True) -> Tensor:
     """Triggers the computation needed to create these Tensor(s)."""
-    # side-realize pending assigns for buffers referenced by these tensors
-    if _pending_assigns:
-      def _realize_pending(buf):
-        for assign_uop in _pending_assigns.pop(buf, []):
-          # recursively realize pending assigns that this assign's value depends on
-          for u in assign_uop.toposort():
-            if u.op is Ops.BUFFER and u in _pending_assigns: _realize_pending(u)
-          big_sink, becomes_map = transform_to_call(UOp.sink(assign_uop))
-          schedule, var_vals = complete_create_schedule_with_vars(big_sink)
-          _apply_map_to_tensors(becomes_map, name="Apply Pending Assign")
-          run_schedule(schedule, var_vals, do_update_stats=do_update_stats)
-          # update remaining pending assigns so they reference realized buffers instead of stale lazy graphs
-          if becomes_map:
-            for assigns in _pending_assigns.values():
-              for i in range(len(assigns)): assigns[i] = assigns[i].substitute(becomes_map)
-      for buf in {u for t in (self,)+lst for u in t.uop.toposort() if u.op is Ops.BUFFER}:
-        if buf in _pending_assigns: _realize_pending(buf)
     if len(to_realize:=[x for x in (self,)+lst if not x.uop.has_buffer_identity()]):
       run_schedule(*Tensor.schedule_with_vars(*to_realize), do_update_stats=do_update_stats)
     return self
@@ -323,13 +305,13 @@ class Tensor(OpMixin):
     if is_disk:
       self._buffer().copyin(x._data())
       return self
-    result = self._apply_uop(UOp.assign, x)
-    # track view assigns (not full-buffer or assign-chain) so they can be side-realized when the buffer is read
-    if (buf_uop:=self.uop.base).op is Ops.BUFFER and self.uop.op is not Ops.ASSIGN and not self.uop.has_buffer_identity():
-      # deduplicate: if the value is already a pending assign for this buffer (e.g. __iadd__ in __setitem__), remove it
-      if x.uop in _pending_assigns.get(buf_uop, []): _pending_assigns[buf_uop].remove(x.uop)
-      _pending_assigns.setdefault(buf_uop, []).append(result.uop)
-    return self.replace(result)
+    # NOTE: assign_uop is created before AFTER embedding (uses original self.uop),
+    # but AFTER must be embedded before _apply_uop (so subsequent assigns see it)
+    assign_uop = self.uop.assign(x.uop)
+    base = self.uop.base
+    if base.op in {Ops.BUFFER, Ops.AFTER} and not self.uop.has_buffer_identity():
+      _apply_map_to_tensors({base: base.after(assign_uop)}, name="Embed View Assign", walk=True)
+    return self.replace(self._apply_uop(lambda *_: assign_uop, x))
 
   def detach(self) -> Tensor:
     """
@@ -740,9 +722,10 @@ class Tensor(OpMixin):
     """
     if stop is None: stop, start = start, 0
     dtype = kwargs.pop("dtype", dtypes.default_float if any(isinstance(x, float) for x in (start, stop, step)) else dtypes.default_int)
-    if start < (dt:=to_dtype(dtype)).min or dt.max < (stop-step): raise ValueError(f"arange [{start}, {stop}) is not representable in dtype {dtype}")
+    if resolve(start < (dt:=to_dtype(dtype)).min, False) or resolve(dt.max < (stop-step), False):
+      raise ValueError(f"arange [{start}, {stop}) is not representable in dtype {dtype}")
     # NOTE: this matches numpy, torch raises RuntimeError if stop-start and step have different signs
-    if (output_len:=ceildiv(stop-start, step)) <= 0: return Tensor([], dtype=dtype, **kwargs)
+    if resolve((output_len:=ceildiv(stop-start, step)) <= 0, False): return Tensor([], dtype=dtype, **kwargs)
     return (Tensor.full((output_len,), step, dtype=dtype, **kwargs)._cumalu(0, Ops.ADD) + (start - step)).cast(dtype)
 
   @staticmethod
@@ -1236,26 +1219,6 @@ class Tensor(OpMixin):
     x_dims = [p for p in indices_parsed if not isinstance(p['index'], sint)]
     x = x.reshape(tuple(p['size'] for p in x_dims))
 
-    # basic setitem: construct result with view region replaced by v using arange masks
-    if v is not None and not any(isinstance(p['index'], Tensor) for p in indices_parsed):
-      # broadcast v to getitem shape, reshape to self.ndim (squeeze None dims, unsqueeze int dims — all are size 1)
-      vb = v.cast(self.dtype)._broadcast_to(x.shape)
-      vb = vb.reshape(tuple(1 if isinstance(p['index'], sint) else p['size'] for p in indices_parsed if p['index'] is not None))
-      # undo movement ops per-dim and build boolean mask
-      per_dim = []
-      for d, m in enumerate(mops):
-        (s, e), st = m['boundary'], abs(m['stride'])
-        if st != 1 and vb.shape[d] > 1:  # un-stride: interleave with zeros
-          vb = vb.unsqueeze(d+1)
-          vb = vb.pad_to(tuple(st if j == d+1 else None for j in range(vb.ndim)))
-          vb = vb.reshape(vb.shape[:d] + (vb.shape[d]*vb.shape[d+1],) + vb.shape[d+2:])
-          vb = vb.shrink_to(tuple(e-s if j == d else None for j in range(self.ndim)))
-        idx = Tensor.arange(self.shape[d], device=self.device).reshape([1]*d + [self.shape[d]] + [1]*(self.ndim - d - 1))
-        per_dim.append((idx >= s) & (idx < e) & (((e-1-idx) if m['stride'] < 0 else (idx-s)) % st == 0))
-      vb = vb.flip(tuple(d for d, m in enumerate(mops) if m['stride'] < 0))
-      vb = vb.pad(tuple((m['boundary'][0], self.shape[d] - m['boundary'][1]) for d, m in enumerate(mops)))
-      return (functools.reduce(lambda a, b: a & b, per_dim) if per_dim else Tensor(True, dtype=dtypes.bool, device=self.device)).where(vb, self)
-
     # tensor indexing
     if tops := [(d, p) for d, p in enumerate(x_dims) if isinstance(p['index'], Tensor)]:
       dims, tensors, masks = [d for d, _ in tops], cast(list[Tensor], [p['index'] for _, p in tops]), []
@@ -1266,7 +1229,7 @@ class Tensor(OpMixin):
       if v is None and len(dims) > 1 and consecutive and all_int(ishp := tuple(x.shape[d] for d in dims)):
         strides = tuple(prod(ishp[i+1:]) for i in range(len(dims)))
         try: linear_idx = functools.reduce(Tensor.add, (t._broadcast_to(big_shape) * s for t, s in zip(tensors, strides)))
-        except ValueError as e: raise IndexError(f"cannot broadcast indices: {e}") from e
+        except ValueError as err: raise IndexError(f"cannot broadcast indices: {err}") from err
         valid = functools.reduce(Tensor.__and__, ((t >= 0) & (t < s) for t, s in zip(tensors, ishp)))
         pre, post = x.shape[:dims[0]], x.shape[dims[-1]+1:]
         x = x.reshape(pre + (prod(ishp),) + post)[tuple([slice(None)] * len(pre)) + (valid.where(linear_idx, 0),)]
@@ -1277,7 +1240,7 @@ class Tensor(OpMixin):
       # create index masks
       for dim, tensor in zip(dims, tensors):
         try: i = tensor.reshape(tensor.shape + (1,)*(x.ndim - dims[0])).expand(pre_reduce_shape)
-        except ValueError as e: raise IndexError(f"cannot broadcast indices: {e}") from e
+        except ValueError as err: raise IndexError(f"cannot broadcast indices: {err}") from err
         masks.append(i._one_hot_along_dim(num_classes=x.shape[dim], dim=(dim - x.ndim)))
 
       # reduce masks to 1 mask
@@ -1286,21 +1249,36 @@ class Tensor(OpMixin):
       # inject 1's for the extra dims added in create masks
       reshape_arg = x.shape[:dims[0]] + (1,) * len(big_shape) + x.shape[dims[0]:]
       # sum reduce the extra dims introduced in create masks
+      x_pre = x  # save collapsed shape for advanced setitem
       x = (mask.where(x.reshape(reshape_arg), 0)).sum(sum_axis:=tuple(d + len(big_shape) for d in dims), dtype=x.dtype)
 
       # special permute case
       if (permuted := dims[0] != 0 and len(dims) != 1 and tuple(dims) != tuple(range(dims[0], dims[-1]+1))):
         mask, x = (y.permute(*range(dims[0], dims[0]+len(big_shape)), *range(0, dims[0]), *range(dims[0]+len(big_shape), y.ndim)) for y in (mask, x))
 
-      # for advanced setitem, returns whole tensor with indices replaced
-      if v is not None:
-        vb = v.cast(self.dtype)._broadcast_to(_broadcast_shape(x.shape, v.shape))
-        # add back reduced dims from sum
-        for dim in sum_axis: vb = vb.unsqueeze(dim)
-        # run _masked_setitem on tuple of axis that is to be reduced to match self.shape
-        x = _masked_setitem(self, vb, mask, tuple(range((start := dims[0] if not permuted else 0), start + len(big_shape))))
-
-    return x
+      if v is None: return x  # advanced getitem
+      # advanced setitem: resolve tensor dims in collapsed space, then fall through to basic setitem path
+      vb = v.cast(self.dtype)._broadcast_to(_broadcast_shape(x.shape, v.shape))
+      for dim in sum_axis: vb = vb.unsqueeze(dim)  # add back reduced dims from sum
+      start = dims[0] if not permuted else 0
+      vb = _masked_setitem(x_pre, vb, mask, tuple(range(start, start + len(big_shape))))
+    elif v is None: return x  # basic getitem
+    # basic setitem: broadcast v, reshape to self.ndim (unsqueeze int dims, squeeze None dims)
+    else: vb = v.cast(self.dtype)._broadcast_to(x.shape)
+    vb = vb.reshape(tuple(1 if isinstance(p['index'], sint) else p['size'] for p in indices_parsed if p['index'] is not None))
+    per_dim = []
+    for d, m in enumerate(mops):
+      (s, e), st = m['boundary'], abs(m['stride'])
+      if st != 1 and vb.shape[d] > 1:  # un-stride: interleave with zeros
+        vb = vb.unsqueeze(d+1)
+        vb = vb.pad_to(tuple(st if j == d+1 else None for j in range(vb.ndim)))
+        vb = vb.reshape(vb.shape[:d] + (vb.shape[d]*vb.shape[d+1],) + vb.shape[d+2:])
+        vb = vb.shrink_to(tuple(e-s if j == d else None for j in range(self.ndim)))
+      idx = Tensor.arange(self.shape[d], device=self.device).reshape([1]*d + [self.shape[d]] + [1]*(self.ndim - d - 1))
+      per_dim.append((idx >= s) & (idx < e) & (((e-1-idx) if m['stride'] < 0 else (idx-s)) % st == 0))
+    vb = vb.flip(tuple(d for d, m in enumerate(mops) if m['stride'] < 0))
+    vb = vb.pad(tuple((m['boundary'][0], self.shape[d] - m['boundary'][1]) for d, m in enumerate(mops)))
+    return (functools.reduce(lambda a, b: a & b, per_dim) if per_dim else Tensor(True, dtype=dtypes.bool, device=self.device)).where(vb, self)
 
   def __getitem__(self, indices) -> Tensor:
     """
@@ -1344,15 +1322,26 @@ class Tensor(OpMixin):
 
   def __setitem__(self, indices, v:Tensor|PyConst|list|tuple) -> None:
     if isinstance(v, Tensor) and v.dtype != self.dtype: raise RuntimeError(f"setitem dtype mismatch: {self.dtype=} != {v.dtype=}")
-    if self.requires_grad or (isinstance(v, Tensor) and v.requires_grad): raise NotImplementedError("setitem with requires_grad is not supported")
+    if self.requires_grad or (isinstance(v, Tensor) and v.requires_grad):
+      # for +=/-=, v's graph references self.uop through the view — exclude those from the stale-use check
+      v_uop, v_bw = (v.uop, v.uop.backward_slice) if isinstance(v, Tensor) else (None, {})
+      if any(self.uop in t.uop.backward_slice for tref in all_tensors
+             if (t:=tref()) is not None and t is not self and t.uop is not v_uop and t.uop not in v_bw):
+        raise RuntimeError("can't setitem on a tensor that already has other uses and requires grad")
+      if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
+      if v.uop.op is Ops.ASSIGN: v = v._apply_uop(lambda x: x.src[1])
+      self.replace(self._getitem(indices, v))
+      return
     idx = [indices] if (isinstance(indices, list) and all_int(indices)) or not isinstance(indices, (tuple, list)) else list(indices)
     is_disk = isinstance(self.device, str) and self.device.startswith("DISK")
     if any(isinstance(i, (Tensor, list, tuple)) for i in idx): # advanced setitem
       if is_disk: raise RuntimeError("advanced setitem is not supported for DISK tensors")
       if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
       self.assign(self._getitem(indices, v))
-    elif is_disk or self.uop.is_realized: # basic setitem, self is realized. TODO: disk uop.base is a COPY and not realized
-      self[indices].assign(v)
+    elif is_disk or self.uop.is_realized or self.uop.base.op is Ops.AFTER: # basic setitem, self is realized
+      view = self[indices]
+      if isinstance(v, Tensor) and v.uop.op is Ops.ASSIGN and v.uop in view.uop.base.src: return
+      view.assign(v)
     else: # basic setitem, self is not realized
       if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
       # __iadd__/__isub__ on unrealized views creates a no-op ASSIGN; unwrap to get the computed value
@@ -2578,7 +2567,6 @@ class Tensor(OpMixin):
 
   @staticmethod
   def _tri(r:sint, c:sint, diagonal:int=0, device=None, requires_grad:bool|None=None) -> Tensor:
-    assert isinstance(r, int) and isinstance(c, int), f"does not support symbolic, getting {r=}, {c=}"
     return (Tensor.arange(r, device=device).unsqueeze(-1) + diagonal <= Tensor.arange(c, device=device)).requires_grad_(requires_grad)
 
   def triu(self, diagonal:int=0) -> Tensor:
@@ -3296,8 +3284,8 @@ class Tensor(OpMixin):
     print(q.scaled_dot_product_attention(k, v).numpy())
     ```
     """
-    # NOTE: it also works when `key` and `value` have symbolic shape.
-    assert all_int(self.shape), f"does not support symbolic shape {self.shape}"
+    # NOTE: only n_heads (dim -3) and head_dim (dim -1) must be concrete; batch and sequence length dims can be symbolic.
+    assert all_int(self.shape[-1:]) and all_int(self.shape[-3:-2]), f"does not support symbolic shape {self.shape}"
 
     if getenv("FLASH_ATTENTION"):
       from extra.thunder.tiny.fa import flash_attention

tinygrad/uop/ops.py

@@ -163,7 +163,7 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     # Check self first, then iterate backward_slice (avoids creating intermediate dict)
     return self.op in ops or any(x.op in ops for x in self.backward_slice)
 
-  def toposort(self, gate:Callable|None=None) -> dict[UOp, None]:
+  def toposort(self, gate:Callable|None=None, enter_calls=True) -> dict[UOp, None]:
     cache: dict[UOp, None] = {}
     stack: list[tuple[UOp, bool]] = [(self, False)] # each stack entry is (node, visited_flag)
     while stack:
@@ -172,7 +172,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       if not visited:
         if gate is None or gate(node):
           stack.append((node, True))  # push node back on stack to process after its srcs
-          for s in reversed(node.src): stack.append((s, False)) # push srcs on the stack
+          for s in reversed(node.src if enter_calls or node.op is not Ops.CALL else node.src[1:]):
+            stack.append((s, False)) # push srcs on the stack
       else: cache[node] = None # second time i'm seeing this node, add it to returned toposort
     return cache
 
@@ -253,6 +254,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
       case Ops.RESHAPE:
         if self.src[0]._shape is None: return self.marg
 
+      # MULTI marker (axis info in PARAM sources) has no shape
+      case Ops.MULTI if len(self.src) == 0: return None
+
     # movement ops change the shape
     # NOTE: ssimplify is required because the shape needs to be canonical for broadcasting and same shape checking
     if self.op in GroupOp.Movement.union({Ops.MULTI, Ops.REDUCE_AXIS, Ops.WMMA}):
@@ -514,6 +518,11 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     # COPY removes axis. TODO: add more tests for this, and consider MSELECT/MSTACK
     if self.op is Ops.COPY: return None
     if self.op is Ops.MULTI: return self.arg
+    # PARAM: axis is stored as a MULTI source
+    if self.op is Ops.PARAM:
+      for s in self.src:
+        if s.op is Ops.MULTI: return s.arg
+      return None
     # NOTE: they all have to share an axis, we always choose [-1]
     if self.op in GroupOp.ALU: return axes[-1] if (axes := dedup([x.axis for x in self.src if x.axis is not None])) else None
     if len(self.src) == 0: return None
@@ -867,7 +876,9 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
   def param_like(self, slot:int):
     if self.op is Ops.BIND:
       return UOp.param(slot, self.dtype, self._shape, self._device, self._min_max, self.src[0].arg[0])
-    return UOp.param(slot, self.dtype, self._shape, self._device)
+    p = UOp.param(slot, self.dtype, self._shape, self._device)
+    if self.axis is not None: p = p.replace(src=p.src + (UOp(Ops.MULTI, arg=self.axis),))
+    return p
 
   def call(self, *srcs:UOp, grad_fxn:Callable|None=None, metadata:tuple[Metadata, ...]=(), name:str|None=None) -> UOp:
     # TODO: reenable this after ENCDEC is fixed
@@ -1422,6 +1433,7 @@ def bitcast(x, in_dtype:DType, out_dtype:DType):
 
 renderer = PatternMatcher([
   (UPat((Ops.DEFINE_VAR,), name="x"), lambda x: x.expr),
+  (UPat(Ops.PARAM, src=(UPat(), UPat(), UPat(), UPat(), UPat(Ops.NOOP, name="x"))), lambda x: x.arg),
   (UPat((Ops.SPECIAL), name="x"), lambda x: x.arg),
   (UPat(Ops.RANGE, name="x"), lambda x: f"r{range_str(x)}"),
   (UPat((Ops.CONST, Ops.VCONST), name="x"), lambda x: str(x.arg)),

tinygrad/uop/symbolic.py

@@ -258,7 +258,8 @@ symbolic = symbolic_simple+commutative+PatternMatcher([
   ((UPat.var("x", dtypes.index) + UPat.cvar("c")).cast(dtypes.sints, name="cast"), lambda x,c,cast:x.cast(cast.dtype)+c.cast(cast.dtype)),
   # only RANGE/IF/STORE/KERNEL have side effects
   (UPat(Ops.AFTER, name="x"), lambda x: x.replace(src=(x.src[0],)+
-    tuple(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.CALL, Ops.BARRIER, Ops.END, Ops.UNROLL, Ops.LINEAR} else y.src for y in x.src[1:]])))),
+    tuple(flatten([(y,) if y.op in {Ops.RANGE, Ops.STORE, Ops.CALL, Ops.BARRIER, Ops.END, Ops.UNROLL, Ops.LINEAR, Ops.BUFFERIZE}
+                   else y.src for y in x.src[1:]])))),
   # after with 1 src is just src[0]
   (UPat(Ops.AFTER, src=(UPat.var("s"),)), lambda s: s),
   # VECTORIZE/CONST

tinygrad/viz/js/index.js

@@ -62,7 +62,7 @@ const drawGraph = (data) => {
   const callCount = g.graph().callCount;
   const nodes = d3.select("#nodes").selectAll("g").data(g.nodes().map(id => g.node(id)), d => d).join("g").attr("class", d => d.className ?? "node")
     .attr("transform", d => `translate(${d.x},${d.y})`).on("click", (e,d) => {
-      if (d.label.startsWith("CALL")) {
+      if (d.callNode) {
         if (state.callSrcMask.has(d.id)) state.callSrcMask.delete(d.id); else state.callSrcMask.add(d.id);
         if (state.callSrcMask.size >= callCount) { showCallSrc.toggle.checked = !showCallSrc.toggle.checked; state.callSrcMask.clear(); }
         return setState({});
@@ -110,7 +110,7 @@ const drawGraph = (data) => {
   });
   addTags(nodes.selectAll("g.tag").data(d => d.tag != null ? [d] : []).join("g").attr("class", "tag")
     .attr("transform", d => `translate(${-d.width/2+8}, ${-d.height/2+8})`).datum(e => e.tag));
-  addTags(nodes.selectAll("g.type").data(d => d.label.startsWith("CALL\n") ? [d] : []).join("g")
+  addTags(nodes.selectAll("g.type").data(d => d.callNode ? [d] : []).join("g")
     .attr("class", d => `tag ${d.collapsed ? 'collapsed' : 'expanded'}`)
     .attr("transform", d => `translate(${-d.width/2}, ${0})`).datum(d => d.collapsed ? "+" : "−"));
   // draw edges
@@ -375,8 +375,10 @@ async function renderProfiler(path, unit, opts) {
         if (shapeRef != null) { ref = {ctx:e.ref, step:0}; shapeRef = ref; }
         else if (ref != null) {
           const start = ref.step>0 ? ref.step+1 : 0;
-          const stepIdx = ctxs[ref.ctx+1].steps.findIndex((s, i) => i >= start && s.name == e.name);
-          if (stepIdx !== -1) { ref.step = stepIdx; shapeRef = ref; }
+          const steps = ctxs[ref.ctx+1].steps;
+          for (let si=start; si<steps.length; si++) {
+            if (steps[si].name == e.name) { ref.step = si; shapeRef = ref; break; }
+          }
         } else {
           const steps = ctxs[state.currentCtx].steps;
           for (let i=state.currentStep+1; i<steps.length; i++) {
@@ -398,65 +400,74 @@ async function renderProfiler(path, unit, opts) {
       div.style("height", levelHeight*levels.length+padding+"px").style("pointerEvents", "none");
     } else {
       const peak = u64();
-      let x = 0, y = 0;
-      const buf_shapes = new Map(), temp = new Map();
       const timestamps = [], valueMap = new Map();
+      // start by unpacking the raw events
+      const memEvents = [];
+      let x = 0, y = 0, shapeIdx = 0;
+      const allocs = new Map();
       for (let j=0; j<eventsLen; j++) {
         const alloc = u8(), ts = u32(), key = u32();
         if (alloc) {
           const dtype = strings[u32()], sz = u64(), nbytes = dtypeSize[dtype]*sz;
-          const shape = {x:[x], y:[y], dtype, sz, nbytes, key};
-          buf_shapes.set(key, shape); temp.set(key, shape);
+          allocs.set(key, {nbytes, shapeKey:`${k}-${shapeIdx++}`});
+          memEvents.push({alloc, key, dtype, sz, nbytes});
           timestamps.push(ts);
           x += 1; y += nbytes; valueMap.set(ts, y);
         } else {
-          const free = buf_shapes.get(key);
-          free.users = Array.from({ length: u32() }, () => ({shape:shapeMap.get(u32()), repr:strings[u32()], num:u32(), mode:u8()}));
+          const users = Array.from({ length: u32() }, () => ({shape:shapeMap.get(u32()), repr:strings[u32()], num:u32(), mode:u8()}));
+          const {nbytes, shapeKey} = allocs.get(key); allocs.delete(key);
+          users?.forEach((u) => selectShape(u.shape).e?.arg.bufs.push({ key:shapeKey, nbytes, num:u.num, mode:u.mode, k }));
+          memEvents.push({alloc, key, users, nbytes});
           timestamps.push(ts); valueMap.set(ts, y);
-          x += 1; y -= free.nbytes;
-          free.x.push(x);
-          free.y.push(free.y.at(-1));
-          temp.delete(key);
-          for (const [k, v] of temp) {
-            if (k <= key) continue;
-            v.x.push(x, x);
-            v.y.push(v.y.at(-1), v.y.at(-1)-free.nbytes);
-          }
+          x += 1; y -= nbytes;
         }
       }
       timestamps.push(dur);
       const height = heightScale(peak);
       const yscale = d3.scaleLinear().domain([0, peak]).range([height, 0]);
-      for (const [num, {dtype, sz, nbytes, y, x:steps, users}] of buf_shapes) {
-        const x = steps.map(s => timestamps[s]);
-        const dur = x.at(-1)-x[0];
-        const arg = { tooltipText:`${dtype}\n${formatUnit(sz)}\n${formatUnit(nbytes, 'B')}\n${formatTime(dur)}`, users, key:`${k}-${shapes.length}` };
-        shapes.push({ x, y0:y.map(yscale), y1:y.map(y0 => yscale(y0+nbytes)), arg, fillColor:cycleColors(colorScheme.BUFFER, shapes.length) });
-        users?.forEach((u) => selectShape(u.shape).e?.arg.bufs.push({ key:arg.key, nbytes, num:u.num, mode:u.mode, k }));
-      }
       // generic polygon merger
       const base0 = yscale(0);
-      const allX = Array.from(new Set(shapes.flatMap(s => s.x))).sort((a,b)=>a-b);
-      const idxs = new Map(allX.map((x,i) => [x, i]));
-      const maxY = new Map(allX.map(x => [x, base0]));
-      // for every [a,b) update the max y at x
-      for (const sh of shapes) {
-        for (let i=0; i<sh.x.length-1; i++) {
-          const startIdx = idxs.get(sh.x[i]), endIdx = idxs.get(sh.x[i+1]);
-          const shapeY = sh.y1[i];
-          for (let k=startIdx; k<endIdx; k++) {
-            const x = allX[k]; maxY.set(x, Math.min(maxY.get(x), shapeY));
+      const sum = {x:[], y0:[], y1:[], fillColor:"#2B1B72"};
+      for (let i=0; i<timestamps.length-1; i++) {
+        const yv = yscale(valueMap.get(timestamps[i]));
+        sum.x.push(timestamps[i], timestamps[i+1]); sum.y1.push(yv, yv); sum.y0.push(base0, base0);
+      }
+      // build individual buffer shapes when user clicks to expand, this detailed layout is n²
+      let bufShapes = null;
+      const buildBufShapes = () => {
+        if (bufShapes != null) return bufShapes;
+        bufShapes = [];
+        const buf_shapes = new Map(), temp = new Map();
+        let x = 0, y = 0;
+        for (const e of memEvents) {
+          if (e.alloc) {
+            const shape = {x:[x], y:[y], dtype:e.dtype, sz:e.sz, nbytes:e.nbytes, key:e.key};
+            buf_shapes.set(e.key, shape); temp.set(e.key, shape);
+            x += 1; y += e.nbytes;
+          } else {
+            const free = buf_shapes.get(e.key);
+            free.users = e.users;
+            x += 1; y -= free.nbytes;
+            free.x.push(x); free.y.push(free.y.at(-1));
+            temp.delete(e.key);
+            for (const [k, v] of temp) {
+              if (k <= e.key) continue;
+              v.x.push(x, x);
+              v.y.push(v.y.at(-1), v.y.at(-1)-free.nbytes);
+            }
           }
         }
-      }
-      const sum = {x:[], y0:[], y1:[], fillColor:"#2B1B72"};
-      for (let i=0; i<allX.length-1; i++) {
-        sum.x.push(allX[i], allX[i+1]);
-        const y = maxY.get(allX[i]); sum.y1.push(y, y); sum.y0.push(base0, base0);
-      }
+        for (const [num, {dtype, sz, nbytes, y, x:steps, users}] of buf_shapes) {
+          const x = steps.map(s => timestamps[s]);
+          const dur = x.at(-1)-x[0];
+          const arg = { tooltipText:`${dtype}\n${formatUnit(sz)}\n${formatUnit(nbytes, 'B')}\n${formatTime(dur)}`, users, key:`${k}-${bufShapes.length}` };
+          bufShapes.push({ x, y0:y.map(yscale), y1:y.map(y0 => yscale(y0+nbytes)), arg, fillColor:cycleColors(colorScheme.BUFFER, bufShapes.length) });
+        }
+        return bufShapes;
+      };
       if (timestamps.length > 0) data.first = data.first == null ? timestamps[0] : Math.min(data.first, timestamps[0]);
       data.tracks.set(k, { shapes:[sum], eventType, visible, offsetY, pcolor:"#c9a8ff", height, peak, scaleFactor:maxheight*4/height,
-                           views:[[sum], shapes], valueMap, rowBorderColor });
+                           get views() { return [[sum], buildBufShapes()]; }, valueMap, rowBorderColor });
       div.style("height", height+padding+"px").style("cursor", "pointer").on("click", (e) => {
         const newFocus = e.currentTarget.id === focusedDevice ? null : e.currentTarget.id;
         let offset = 0;
@@ -498,10 +509,14 @@ async function renderProfiler(path, unit, opts) {
     const visibleX = xscale.range().map(zoomLevel.invertX, zoomLevel).map(xscale.invert, xscale);
     const st = visibleX[0], et = visibleX[1];
     xscale.domain([st, et]);
+    const profilerEl = profiler.node();
+    const visibleYStart = profilerEl.scrollTop-canvasTop + rect(profilerEl).top, visibleYEnd = visibleYStart+profilerEl.clientHeight;
     ctx.textBaseline = "middle";
     // draw shapes
     for (const [k, { shapes, eventType, visible, offsetY, valueMap, pcolor, scolor, rowBorderColor }] of data.tracks) {
       visible.length = 0;
+      const trackHeight = rect(document.getElementById(k)).height;
+      if (offsetY+trackHeight < visibleYStart || offsetY > visibleYEnd) continue;
       const addBorder = scolor != null ? (w) => { if (w > 10) { ctx.strokeStyle = scolor; ctx.stroke(); } } : null;
       for (const e of shapes) {
         if (eventType === EventTypes.BUF) { // generic polygon
@@ -535,7 +550,7 @@ async function renderProfiler(path, unit, opts) {
       }
       // draw row line
       if (rowBorderColor != null) {
-        const y = offsetY+rect(document.getElementById(k)).height-padding/2 - 0.5;
+        const y = offsetY+trackHeight-padding/2 - 0.5;
         drawLine(ctx, [0, canvasWidth], [y, y], { color:rowBorderColor });
       }
     }
@@ -599,6 +614,7 @@ async function renderProfiler(path, unit, opts) {
   document.addEventListener("contextmenu", e => e.ctrlKey && e.preventDefault());
 
   new ResizeObserver(([e]) => e.contentRect.width > 0 && resize()).observe(profiler.node());
+  profiler.on("scroll", () => render(zoomLevel));
 
   function findRectAtPosition(x, y) {
     let track = null;

tinygrad/viz/js/worker.js

@@ -54,13 +54,14 @@ const layoutUOp = (g, { graph, change }, opts) => {
       width = Math.max(width, ctx.measureText(line).width);
       height += lineHeight;
     }
-    if (label.startsWith("CALL\n")) callCount++;
-    g.setNode(k, {...rectDims(width, height), label, ref, id:k, color, tag});
+    const callNode = label.startsWith("CALL\n");
+    if (callNode) callCount++;
+    g.setNode(k, {...rectDims(width, height), label, ref, id:k, color, tag, callNode});
     // add edges
     const edgeCounts = {};
     for (const [_, s] of src) edgeCounts[s] = (edgeCounts[s] || 0)+1;
     for (const [port, s] of src) g.setEdge(s, k, { label: edgeCounts[s] > 1 ? {type:"tag", text:edgeCounts[s]} : {type:"port", text:port},
-      ...(label.startsWith("CALL\n") && port === 0 && {color:"#a0a1b8"})});
+      ...(callNode && port === 0 && {color:"#a0a1b8"})});
     if (change?.includes(parseInt(k))) g.setParent(k, "overlay");
   }
   // optionally hide nodes from the layout
@@ -81,7 +82,7 @@ const layoutUOp = (g, { graph, change }, opts) => {
     const disconnected = new Set();
     for (const n of g.nodes()) {
       const node = g.node(n);
-      if (node.label.startsWith("CALL\n") && (opts.showCallSrc ? opts.callSrcMask.has(n) : !opts.callSrcMask.has(n))) {
+      if (node.callNode && (opts.showCallSrc ? opts.callSrcMask.has(n) : !opts.callSrcMask.has(n))) {
         node.collapsed = true;
         for (const pred of (g.predecessors(n) || [])) {
           const edge = g.edge(pred, n);

tinygrad/viz/serve.py

@@ -128,6 +128,8 @@ def uop_to_json(x:UOp) -> dict[int, dict]:
         label += f"\n({multirange_str(rngs, color=True)})"
       if u._shape is not None:
         label += f"\n{shape_to_str(u.shape)}"
+      if u.op is Ops.CALL:
+        label += f"\n{u.src[0].key.hex()[:8]}"
       if u.op in {Ops.INDEX, Ops.BUFFERIZE}:
         if len(u.toposort()) < 30: label += f"\n{u.render()}"
         ranges: list[UOp] = []