163 lines
9.0 KiB
Python
163 lines
9.0 KiB
Python
"""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()
|