"""Observe MTPLX's generated Metal source in this model-free reference process. The local Objective-C hook forwards the original call unchanged and is restored before exit. It neither modifies the reference checkout nor writes artifacts. Run in the pinned MTPLX environment; stdout is a JSON source receipt. """ import ctypes as c import argparse import hashlib import json from pathlib import Path import re def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--check",action="store_true",help="compare with the pinned captured shader receipt") parser.add_argument("--rank",type=int,choices=(2,3,5),default=2, help="compiled input rank; 3/5 cover sorted/unsorted SwitchGLU") parser.add_argument("--operation",choices=("silu","swiglu","compute_g","qsa-update"),default="silu", help="actual compiled activation used by fused or separate projection modules") parser.add_argument("--qsa-shape",type=int,nargs=3,default=(4,4096,1024), metavar=("ROWS","RAW_CAP","POOL_CAP")) parser.add_argument("--qsa-header",action="store_true",help="also retain the unchanged scalar library header") parser.add_argument("--qsa-mode",choices=("update_only","blocks","row_tokens","dense_mask","prefill_blocks"),default="update_only") parser.add_argument("--qsa-score-budget",type=int,default=32*1024*1024) args = parser.parse_args() objc = c.CDLL("/usr/lib/libobjc.A.dylib") metal = c.CDLL("/System/Library/Frameworks/Metal.framework/Metal") metal.MTLCreateSystemDefaultDevice.restype = c.c_void_p device = metal.MTLCreateSystemDefaultDevice() assert device objc.object_getClass.argtypes = [c.c_void_p] objc.object_getClass.restype = c.c_void_p objc.sel_registerName.argtypes = [c.c_char_p] objc.sel_registerName.restype = c.c_void_p objc.class_getInstanceMethod.argtypes = [c.c_void_p,c.c_void_p] objc.class_getInstanceMethod.restype = c.c_void_p objc.method_getImplementation.argtypes = [c.c_void_p] objc.method_getImplementation.restype = c.c_void_p objc.method_setImplementation.argtypes = [c.c_void_p,c.c_void_p] objc.method_setImplementation.restype = c.c_void_p selector = objc.sel_registerName(b"newLibraryWithSource:options:error:") method = objc.class_getInstanceMethod(objc.object_getClass(device),selector) assert method signature = c.CFUNCTYPE(c.c_void_p,*([c.c_void_p]*5)) original_ptr = objc.method_getImplementation(method) original = signature(original_ptr) string = c.CFUNCTYPE(c.c_char_p,c.c_void_p,c.c_void_p)(("objc_msgSend",objc)) utf8 = objc.sel_registerName(b"UTF8String") receipts = [] errors = [] @signature def observe(receiver,sel,source,options,error): try: text = string(source,utf8).decode() match = re.search(r'\[\[host_name\("[^"]+"\)\]\]\n\[\[kernel\]\] void ',text) if match and "tmp_" in text[match.start():]: kernels = text[match.start():] if args.operation == "qsa-update": # All moving frontiers are contiguous int32[1] leaves. # Retain the exact invoked specialization, not unused ranks. kernels = kernels[:kernels.index("[[host_name",2)] receipts.append({"source_sha256":hashlib.sha256(text.encode()).hexdigest(), "kernels":kernels}) if args.operation == "qsa-update" and args.qsa_header and len(receipts)==1: receipts[-1]["header"] = text[:match.start()] elif args.operation == "qsa-update" and "[[kernel]] void compute_dynamic_offset_" in text: receipts.append({"source_sha256":hashlib.sha256(text.encode()).hexdigest(), "kernels":text[text.index("[[kernel]] void compute_dynamic_offset_"):]}) except Exception as exc: errors.append(str(exc)) return original(receiver,sel,source,options,error) objc.method_setImplementation(method,c.cast(observe,c.c_void_p)) try: import mlx.core as mx import mlx.nn as nn assert mx.__version__ == "0.32.2",mx.__version__ shape = {2:(70,1280),3:(70,1,1280),5:(1,7,10,1,1280)}[args.rank] x = mx.arange(70*1280).astype(mx.bfloat16).reshape(shape)/128 mx.eval(x) gate,up = mx.split(x,2,axis=-1) if args.operation == "qsa-update": from mtplx.kernels.qsa_indexer_compile import QSACompiledIndexerCore norm = mx.ones((128,),dtype=mx.bfloat16) freq = mx.arange(32,dtype=mx.float32)/128 core = QSACompiledIndexerCore(n_heads=4,kv_heads=1,head_dim=128, block_topk=512,compress_ratio=4,q_norm_weight=norm,k_norm_weight=norm, inv_freq=freq,rms_norm_eps=1e-6,selector_scratch_bytes=args.qsa_score_budget) rows,raw_cap,pool_cap = args.qsa_shape qk = mx.arange(rows*640).astype(mx.bfloat16).reshape(1,rows,640)/128 raw = mx.zeros((1,raw_cap,128),dtype=mx.bfloat16) pooled = mx.zeros((1,pool_cap,128),dtype=mx.bfloat16) pos = min(2051,raw_cap-rows) total = pos+rows logical = total//4 scalars = [mx.array([n],dtype=mx.int32) for n in (pos,total,logical,max(0,logical-(rows+3)//4))] mx.eval(qk,raw,pooled,norm,freq,*scalars) result = core.select_qk_rows(qk,raw,pooled,pos_start=scalars[0], total_tokens=scalars[1],logical_blocks=scalars[2],pooled_len=scalars[3],mode=args.qsa_mode) mx.eval(result.raw_keys,result.pooled,result.pooled_len,result.offset,result.selection) elif args.operation == "compute_g": from mlx_lm.models.gated_delta import compute_g # Actual fused-projection A view, with per-head A_log/dt_bias broadcasts. ashape = {2:(70,16480),3:(1,70,16480),5:(1,7,10,1,16480)}[args.rank] a = (mx.arange(70*16480).astype(mx.bfloat16)/128).reshape(ashape)[...,-48:] a_log = mx.arange(48).astype(mx.bfloat16)/128 dt_bias = -a_log mx.eval(a,a_log,dt_bias) mx.eval(compute_g(a_log,a,dt_bias)) elif args.operation == "swiglu": from mlx_lm.models.activations import swiglu gate,up = mx.contiguous(gate),mx.contiguous(up) mx.eval(gate,up) mx.eval(swiglu(gate,up)) else: mx.eval(nn.silu(gate)*up) assert not errors,errors if args.operation == "qsa-update": assert receipts,"No compiled QSA kernels captured" if args.check: root = Path(__file__).resolve().parents[1] if args.qsa_header: assert receipts[0].pop("header")== (root/"metal/mtplx-qsa-compiled-header.metal").read_text(),"Compiled scalar header differs" if tuple(args.qsa_shape)==(4,4096,1024): fixture = "mtplx-qsa-select-jit.json" if args.qsa_mode=="blocks" and args.qsa_score_budget==4096 else "mtplx-qsa-update-jit.json" expected = json.loads((root/"tests/fixtures"/fixture).read_text()) assert receipts==expected,"Compiled QSA shaders differ from the pinned receipt" else: bank = json.loads((root/"metal/mtplx-qsa-compiled-scalars.json").read_text()) rows,raw,pool = args.qsa_shape new = (rows+3)//4 limit = min(raw//4-new,pool-new) candidates = [*bank["clamp"].values(),bank["check_distinct"]] expected = next(v for v in candidates if (v["max_new"],v["max_start"])==(new,limit)) assert receipts[0]=={k:expected[k] for k in ("source_sha256","kernels")},"Compiled clamp differs" assert receipts[1]==bank["multiply"],"Compiled multiply differs" print("Actual compiled QSA update shaders match the pinned receipt",flush=True) else: print(json.dumps(receipts),flush=True) return assert len(receipts)==1,len(receipts) if args.check: expected = json.loads((Path(__file__).resolve().parents[1]/f"tests/fixtures/mtplx-{args.operation}-jit.json").read_text()) if args.rank==2: assert receipts[0]==expected,"Generated activation differs from the pinned receipt" else: def normalize(source): name = re.search(r'host_name\("([^\"]+)_contiguous"',source).group(1) return source.replace(name,"COMPILED_ACTIVATION") assert normalize(receipts[0]["kernels"])==normalize(expected["kernels"]),"Rank-dependent activation computation changed" print(f"Actual rank-{args.rank} {args.operation} compiler source matches the pinned receipt (rank-specific symbol names normalized for 3/5)",flush=True) else: print(json.dumps(receipts[0]),flush=True) finally: objc.method_setImplementation(method,original_ptr) if __name__ == "__main__": main()