3489 lines
215 KiB
Python
3489 lines
215 KiB
Python
"""Model-free correctness receipts from the pinned MTPLX custom kernels.
|
|
|
|
No timings, generation, weights download, or replacement implementations.
|
|
Run with the pinned runtime and PYTHONPATH pointing at the MTPLX checkout.
|
|
"""
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import runpy
|
|
import sys
|
|
|
|
import mlx.core as mx
|
|
from mlx_lm.models import gated_delta
|
|
import numpy as np
|
|
from mtplx.kernels.gdn_conv_norm import fused_gdn_conv_norm, fused_gdn_conv_norm_rows
|
|
from mtplx.kernels.gdn_step_fused import fused_gdn_step
|
|
from mtplx.kernels.gdn_out_fused import fused_gdn_out
|
|
from mtplx.kernels.hyper_connection import fused_hyper_read
|
|
from mtplx.kernels.hyper_connection_v3 import fused_hyper_read_v3
|
|
from mtplx.kernels import hyper_connection
|
|
|
|
|
|
def pattern(count, salt):
|
|
return mx.array((((np.arange(count, dtype=np.int64) * 17 + salt * 13) % 257)
|
|
- 128).astype(np.float32) / 128, dtype=mx.bfloat16)
|
|
|
|
|
|
def qsa_attention_module_fixtures(*, profile_filter=None):
|
|
"""Whole unchanged Attention, ongoing KV/indexer state and actual consumers."""
|
|
from contextlib import nullcontext
|
|
from mtplx.models import qwen4_exp as model
|
|
from mtplx.attention_context import attention_phase, vision_rope
|
|
from mtplx.kernels import qsa_flash_skip, qsa_prefill_flash
|
|
assert mx.__version__ == "0.32.2"
|
|
defaults=dict(fused=True,compiled=True,prefill=True,phase_prefill=True,
|
|
prefill_min_rows=32,prefill_compile_rows=2048,prefill_min_context=2049,
|
|
prefill_workspace=128*1024*1024,tile_rows=0,flash=True,gather_decode=False,
|
|
gather=True,gather_max_rows=8,gather_min_context=0,tf32=os.environ.get("MLX_ENABLE_TF32","1")!="0")
|
|
envs=dict(fused="MTPLX_FUSED_QSA_INDEXER",compiled="MTPLX_COMPILED_QSA_INDEXER",prefill="MTPLX_QSA_PREFILL",
|
|
prefill_min_rows="MTPLX_QSA_PREFILL_MIN_ROWS",prefill_compile_rows="MTPLX_QSA_PREFILL_COMPILE_ROWS",
|
|
prefill_min_context="MTPLX_QSA_PREFILL_MIN_CONTEXT",flash="MTPLX_QSA_FLASH",gather_decode="MTPLX_QSA_GATHER_DECODE",
|
|
gather="MTPLX_QSA_GATHER",gather_max_rows="MTPLX_QSA_GATHER_MAX_ROWS",gather_min_context="MTPLX_QSA_GATHER_MIN_CONTEXT",tile_rows="MTPLX_QSA_SCORE_TILE_ROWS")
|
|
saved=os.environ.copy()
|
|
profiles=[(0,(1,4,33),4,{},"none",False,False,2049,True),
|
|
(0,(2048,1),3,{},"none",False,False,2049,True),
|
|
(2051,(1,4,33),4,{},"none",False,False,2049,True),
|
|
(2051,(33,1,4),0,dict(compiled=False,flash=False,gather=False),"none",False,False,32768,True),
|
|
(2051,(33,1,4),3,{},"none",False,False,32768,False),
|
|
(2051,(1,4),0,dict(compiled=False,fused=False,flash=False,gather_decode=True),"none",False,False,32768,True),
|
|
(32768,(2048,1,4),4,{},"none",False,False,32768,True),
|
|
(2051,(1,4,2,1),4,{},"table",False,False,2049,True),
|
|
(0,(3,1,4),0,{},"none",True,True,2049,True),
|
|
(2051,(1,4),3,{},"no_axes",False,False,2049,True)]
|
|
trace=[];hooks=[]
|
|
def hook(owner,name,label):
|
|
original=getattr(owner,name)
|
|
def call(*args,**kwargs):
|
|
trace.append(label);return original(*args,**kwargs)
|
|
hooks.append((owner,name,original));setattr(owner,name,call)
|
|
hook(qsa_flash_skip,"qsa_flash_skip","flash")
|
|
hook(qsa_prefill_flash,"qsa_prefill_flash","prefill_flash")
|
|
hook(model,"_qsa_rows_gather_attention","rows_gather")
|
|
hook(model,"_qsa_prefill_gather_attention","prefill_gather")
|
|
hook(model,"_qsa_blocks_to_dense_mask","dense_mask")
|
|
hook(mx.fast,"scaled_dot_product_attention","sdpa")
|
|
try:
|
|
for profile,(prefix,steps,fused,override,vision,no_indexer,biased,flash_min,gather) in enumerate(profiles):
|
|
if profile_filter is not None and profile!=profile_filter:continue
|
|
opt=defaults|override
|
|
os.environ.update({env:str(int(opt[key])) for key,env in envs.items()})
|
|
os.environ.update(MTPLX_QSA_PREFILL_SCORE_MB="128",MTPLX_QSA_PREFILL_FLASH_MIN_CONTEXT=str(flash_min),MTPLX_QSA_PREFILL_GATHER=str(int(gather)),MTPLX_QSA_PREFILL_GATHER_TILE="64")
|
|
args=model.TextArgs(hidden_size=64,indexer_n_heads=0 if no_indexer else 4,
|
|
mrope_section=[8,12,12] if vision=="table" else None,mrope_interleaved=True)
|
|
attn=model.Attention(args)
|
|
for name,n,k,salt in [("q_proj",12288,64,400),("k_proj",512,64,402),("v_proj",512,64,404),("o_proj",64,6144,406)]:
|
|
layer=quantized_linear_fixture(n,k,4,64,salt)
|
|
if biased:layer.bias=pattern(n,salt+40)/64
|
|
setattr(attn,name,layer)
|
|
attn.q_norm.weight=pattern(256,410);attn.k_norm.weight=pattern(256,411)
|
|
if attn.indexer is not None:
|
|
attn.indexer.index_qk_proj=quantized_linear_fixture(640,64,4,64,408)
|
|
attn.indexer.q_layernorm.weight=pattern(128,72);attn.indexer.k_layernorm.weight=pattern(128,73)
|
|
if fused:
|
|
parts=[attn.q_proj,attn.k_proj,attn.v_proj]
|
|
if fused==4:parts.append(attn.indexer.index_qk_proj)
|
|
splits=np.cumsum([a.weight.shape[0] for a in parts[:-1]]).tolist()
|
|
attn.qkv_fused=model._FusedGDNInProj(*(mx.concatenate([getattr(a,name) for a in parts],axis=0) for name in ("weight","scales","biases")),64,4,"affine",splits)
|
|
for name in ("q_proj","k_proj","v_proj"):attn.pop(name,None)
|
|
if fused==4:attn.indexer.pop("index_qk_proj",None)
|
|
cache=model.QSACache()
|
|
if prefix:
|
|
cache.write_raw(pattern(prefix*128,88).reshape(1,prefix,128))
|
|
attn.indexer._extend_pooled(cache,prefix)
|
|
kv=pattern(prefix*512,89).reshape(1,prefix,2,256).transpose(0,2,1,3)
|
|
cache.kv.update_and_fetch(kv,kv)
|
|
mx.eval(*cache.state)
|
|
table=mx.array(np.stack([np.arange(prefix+5)//3,np.arange(prefix+5)%17,np.arange(prefix+5)%11]),mx.int32) if vision=="table" else None
|
|
for step,rows in enumerate(steps):
|
|
if step==2 and profile in (2,7):assert cache.trim(2)==2
|
|
pos=cache.offset
|
|
trace.clear();before=model.qsa_prefill_engagement()
|
|
with attention_phase("prefill"), vision_rope(table,-2) if vision!="none" else nullcontext():
|
|
result=attn(pattern(rows*64,420+step).reshape(1,rows,64),cache)
|
|
counts={k:v-before.get(k,0) for k,v in model.qsa_prefill_engagement().items() if v!=before.get(k,0)}
|
|
state=[cache.raw_keys,cache.pooled,cache.pooled_f32_t,cache.kv.keys,cache.kv.values]
|
|
leaves=[result,*[a for a in state if a is not None]]
|
|
core=None if attn.indexer is None else attn.indexer._compiled_indexer_core
|
|
emit(f"qsa_attention_p{profile}_s{step}",leaves,profile=profile,step=step,rows=rows,pos=pos,prefix=prefix,
|
|
fused=fused,vision=vision,no_indexer=no_indexer,biased=biased,flash_min=flash_min,prefill_gather=gather,
|
|
options=opt,consumers=trace.copy(),counts=counts,offset=cache.offset,pooled_len=cache.pooled_len,
|
|
state_present=[a is not None for a in state],shapes=[a.shape for a in leaves],
|
|
core=None if core is None else core.to_dict(),frequencies=attn._inv_freq.tolist())
|
|
finally:
|
|
for owner,name,original in reversed(hooks):setattr(owner,name,original)
|
|
for key in set(envs.values())|{"MTPLX_QSA_PREFILL_SCORE_MB","MTPLX_QSA_PREFILL_FLASH_MIN_CONTEXT","MTPLX_QSA_PREFILL_GATHER","MTPLX_QSA_PREFILL_GATHER_TILE"}:
|
|
if key in saved:os.environ[key]=saved[key]
|
|
else:os.environ.pop(key,None)
|
|
|
|
|
|
def qsa_indexer_commit_fixtures():
|
|
"""Unmodified indexer core lifecycle, ongoing cache commit and return API."""
|
|
from mtplx.models.qwen4_exp import QSAIndexer, TextArgs, QSACache
|
|
import mlx.nn as nn
|
|
assert mx.__version__ == "0.32.2"
|
|
saved = os.environ.get("MTPLX_QSA_PREFILL_SCORE_MB")
|
|
steps = [(1,"update_only"),(3,"update_only"),(33,"prefill_blocks"),
|
|
(1,"blocks"),(4,"row_tokens"),(4,"dense_mask"),(221,"prefill_blocks"),
|
|
(1,"blocks"),(4,"row_tokens"),(4,"dense_mask"),(4,"row_tokens"),
|
|
(1,"blocks"),(4,"dense_mask")]
|
|
try:
|
|
for quantized in (False, True):
|
|
os.environ["MTPLX_QSA_PREFILL_SCORE_MB"] = "128"
|
|
idx=QSAIndexer(TextArgs());idx.block_topk=8
|
|
idx.q_layernorm.weight=pattern(128,72);idx.k_layernorm.weight=pattern(128,73)
|
|
if quantized:idx.index_qk_proj=quantized_linear_fixture(640,128,4,64,75)
|
|
else:
|
|
idx.index_qk_proj=nn.Linear(128,640,bias=False)
|
|
idx.index_qk_proj.weight=pattern(640*128,231).reshape(640,128)
|
|
cache=QSACache()
|
|
for step,(rows,mode) in enumerate(steps):
|
|
# Same bytes, new Python array identity: every leaf matters.
|
|
if step==2:
|
|
idx._fused_score_scratch_bytes=1024*1024
|
|
os.environ["MTPLX_QSA_PREFILL_SCORE_MB"]="2"
|
|
if step==3:idx.q_layernorm.weight=mx.array(idx.q_layernorm.weight)
|
|
if step==4:idx.k_layernorm.weight=mx.array(idx.k_layernorm.weight)
|
|
if step==5:idx._inv_freq=mx.array(idx._inv_freq)
|
|
if step==6:idx.index_qk_proj.weight=mx.array(idx.index_qk_proj.weight)
|
|
if step==7 and quantized:idx.index_qk_proj.scales=mx.array(idx.index_qk_proj.scales)
|
|
if step==8 and quantized:idx.index_qk_proj.biases=mx.array(idx.index_qk_proj.biases)
|
|
if step==9:cache.reserve_indexer_capacity(raw_capacity=2048,pooled_capacity=512)
|
|
if step==10:assert cache.trim(3)==3
|
|
pos=cache.offset
|
|
if step==11:
|
|
old_core=idx._compiled_indexer_core;old_seal=idx._compiled_indexer_parameter_signature
|
|
valid=idx.q_layernorm.weight
|
|
idx.q_layernorm.weight=mx.zeros((127,),mx.bfloat16)
|
|
try:idx._get_compiled_indexer_core()
|
|
except ValueError:pass
|
|
else:raise AssertionError("invalid replacement must fail")
|
|
assert idx._compiled_indexer_core is old_core and idx._compiled_indexer_parameter_signature==old_seal
|
|
idx.q_layernorm.weight=valid
|
|
hidden=pattern(rows*128,300+step).reshape(1,rows,128)
|
|
supplied=step%2==0
|
|
qk=pattern(rows*704,320+step).reshape(1,rows,704)[...,:640] if supplied else None
|
|
if cache.pooled is not None:cache.pooled_f32_view(cache.pooled_len)
|
|
old_raw,old_pool=cache.raw_keys,cache.pooled
|
|
old_arrays=[a for a in (old_raw,old_pool) if a is not None]
|
|
mx.eval(*old_arrays)
|
|
result=idx._call_rows_compiled(hidden,pos,cache,qk,mode=mode)
|
|
assert cache.offset==pos and cache.pooled_f32_t is None
|
|
assert cache.raw_keys is not old_raw and cache.pooled is not old_pool
|
|
if result is None:kind,values,tail="none",[],None
|
|
elif isinstance(result,tuple):
|
|
kind=result[0];tail=result[2] if kind=="flash" else None
|
|
values=[result[1]] if kind=="flash" else list(result[1:])
|
|
else:kind,values,tail="dense",[result],None
|
|
outputs=[*values,cache.raw_keys,cache.pooled,*old_arrays]
|
|
emit(f"qsa_indexer_commit_q{int(quantized)}_s{step}",outputs,
|
|
quantized=quantized,step=step,rows=rows,pos=pos,mode=mode,supplied=supplied,
|
|
kind=kind,tail=tail,shapes=[a.shape for a in outputs],
|
|
pooled_len=cache.pooled_len,report=idx._compiled_indexer_core.to_dict(),
|
|
frequencies=idx._inv_freq.tolist())
|
|
cache.kv.offset=pos+rows # Attention owns this later frontier.
|
|
finally:
|
|
if saved is None:os.environ.pop("MTPLX_QSA_PREFILL_SCORE_MB",None)
|
|
else:os.environ["MTPLX_QSA_PREFILL_SCORE_MB"]=saved
|
|
|
|
|
|
def qsa_core_host_fixtures():
|
|
"""Original Core validation, engagement and replay; real compile/eval."""
|
|
from mtplx.kernels.qsa_indexer_compile import QSACompiledIndexerCore
|
|
from mtplx.models.qwen4_exp import QSAIndexer, TextArgs
|
|
assert mx.__version__ == "0.32.2"
|
|
norm=pattern(128,72);freq=QSAIndexer(TextArgs())._inv_freq
|
|
weight=pattern(640*128,231).reshape(640,128)
|
|
def make(project=lambda x:x@weight.T):
|
|
return QSACompiledIndexerCore(n_heads=4,kv_heads=1,head_dim=128,block_topk=8,
|
|
compress_ratio=4,q_norm_weight=norm,k_norm_weight=norm,inv_freq=freq,rms_norm_eps=1e-6,
|
|
project_qk=project,minimum_raw_capacity=16,minimum_pooled_capacity=16,
|
|
selector_scratch_bytes=1024,prefill_score_workspace_bytes=4096)
|
|
core=make()
|
|
cases=[(1,64,128,32,"blocks",False,None), (1,65,128,32,"blocks",False,None),
|
|
(4,64,128,32,"row_tokens",False,None),(4,65,128,32,"row_tokens",False,None),
|
|
(17,64,128,32,"prefill_blocks",False,None),(17,65,256,64,"prefill_blocks",False,None),
|
|
(4,65,256,64,"dense_mask",True,None),(4,66,256,64,"update_only",True,None),
|
|
(4,65,256,64,"dense_mask",False,300),(4,65,256,64,"dense_mask",False,None),
|
|
(1,64,128,32,"blocks",False,None)]
|
|
for i,(rows,pos,rc,pc,mode,hidden,dense) in enumerate(cases):
|
|
width=128 if hidden else 640
|
|
stride=width+(64 if i%2 else 0)
|
|
source=pattern(rows*stride,232+i).reshape(1,rows,stride)[...,:width]
|
|
raw=pattern(rc*128,250+i).reshape(1,rc,128)
|
|
pool=pattern(pc*128,270+i).reshape(1,pc,128)
|
|
front=[pos,pos+rows,(pos+rows)//4,pos//4]
|
|
states=[mx.array(n,mx.int32) if i%2 else n for n in front]
|
|
out=(core.select_hidden if hidden else core.select_qk_rows)(source,raw,pool,pos_start=states[0],total_tokens=states[1],logical_blocks=states[2],pooled_len=states[3],mode=mode,dense_output_tokens=dense)
|
|
selected=[] if out.selection is None else list(out.selection) if isinstance(out.selection,tuple) else [out.selection]
|
|
report=core.to_dict();report.pop("compiled_keys")
|
|
emit(f"qsa_core_host_{i}",[*selected,out.raw_keys,out.pooled,out.pooled_len,out.offset],case=i,rows=rows,pos=pos,raw_capacity=rc,pooled_capacity=pc,mode=mode,hidden=hidden,dense=dense,frequencies=freq.tolist(),report=report)
|
|
print(json.dumps(dict(kernel="qsa_core_report",report=core.to_dict()),separators=(",",":")),flush=True)
|
|
invalid=[dict(pos=-1),dict(pos=127),dict(total=69),dict(total=-1),dict(logical=33),dict(logical=15),dict(frontier=0),dict(frontier=33),dict(raw_capacity=96),dict(raw_capacity=2),dict(pooled_capacity=0),dict(dense=100,mode="dense_mask"),dict(dense=1),dict(mode="prefill_blocks",rows=1),dict(width=639),dict(index_dtype="float32"),dict(index_dtype="int32",pos=127)]
|
|
for i,change in enumerate(invalid):
|
|
params=dict(pos=64,total=68,logical=17,frontier=16,raw_capacity=128,pooled_capacity=32,dense=None,mode="blocks",rows=4,width=640,index_dtype=None)|change
|
|
x=mx.zeros((1,params["rows"],params["width"]),mx.bfloat16)
|
|
raw=mx.zeros((1,params["raw_capacity"],128),mx.bfloat16);pool=mx.zeros((1,params["pooled_capacity"],128),mx.bfloat16)
|
|
pos=mx.array(params["pos"],getattr(mx,params["index_dtype"])) if params["index_dtype"] else params["pos"]
|
|
try:
|
|
core.select_qk_rows(x,raw,pool,pos_start=pos,total_tokens=params["total"],logical_blocks=params["logical"],pooled_len=params["frontier"],mode=params["mode"],dense_output_tokens=params["dense"])
|
|
except (ValueError,TypeError):ok=True
|
|
else:ok=False
|
|
print(json.dumps(dict(kernel=f"qsa_core_invalid_{i}",params=params,rejected=ok),separators=(",",":")),flush=True)
|
|
bad=make(lambda x:x@weight[:639].T)
|
|
for i in range(2):
|
|
try:bad.select_hidden(mx.zeros((1,4,128),mx.bfloat16),mx.zeros((1,128,128),mx.bfloat16),mx.zeros((1,32,128),mx.bfloat16),pos_start=64,total_tokens=68,logical_blocks=17,pooled_len=16,mode="blocks")
|
|
except ValueError:ok=True
|
|
else:ok=False
|
|
print(json.dumps(dict(kernel=f"qsa_core_failed_trace_{i}",rejected=ok,report=bad.to_dict()),separators=(",",":")),flush=True)
|
|
print(json.dumps(dict(kernel="qsa_core_report_after_invalid",report=core.to_dict()),separators=(",",":")),flush=True)
|
|
|
|
|
|
def dynamic_slice_graph_fixtures(dtype_tag):
|
|
"""Original runtime dynamic slice/update; indices stay device arrays."""
|
|
assert mx.__version__=="0.32.2"
|
|
dtype=getattr(mx,dtype_tag)
|
|
profiles=[((1,16,8),(1,3,8),(1,),(5,),"plain"),
|
|
((1,16,8),(1,3,8),(1,),(12,),"stride"),
|
|
((1,16,8),(1,3,8),(1,),(6,),"transpose"),
|
|
((1,16,8),(1,3,8),(1,),(5,),"broadcast"),
|
|
((3,7,8),(2,3,4),(0,1,2),(1,2,3),"plain"),
|
|
((2,3,4,5),(1,2,3,4),(0,1,2,3),(1,1,1,1),"stride"),
|
|
((1,8,4),(1,2,4),(1,),(3,),"reverse"),
|
|
((1,8,4),(1,0,4),(1,),(2,),"plain"),
|
|
((1,8,4),(1,2,4),(-2,),(4,),"scalar_update"),
|
|
((1,16,128),(1,3,128),(1,),(5,),"plain")]
|
|
for i,(shape,size,axes,values,kind) in enumerate(profiles):
|
|
idx_dtype=getattr(mx,("int32","int64","uint32","uint64","int8","uint8","int16","uint16","int32","uint16")[i])
|
|
phys=shape[:-1]+(shape[-1]*2,) if kind=="stride" else shape
|
|
if kind=="transpose":phys=(shape[0],shape[2],shape[1])
|
|
if kind=="broadcast":phys=(1,1,shape[-1])
|
|
a=pattern(int(np.prod(phys)),221).astype(dtype).reshape(phys)
|
|
if kind=="stride":a=a[...,::2]
|
|
if kind=="transpose":a=a.swapaxes(-1,-2)
|
|
if kind=="broadcast":a=mx.broadcast_to(a,shape)
|
|
if kind=="reverse":a=a[:,::-1,:]
|
|
idx=mx.array(values,dtype=idx_dtype)
|
|
if i==0:idx=idx.reshape(())
|
|
if i==2:idx=mx.array([values[0],99],dtype=idx_dtype)[:1]
|
|
upd=pattern(int(np.prod(size)),222).astype(dtype).reshape(size)
|
|
if kind=="scalar_update":upd=mx.array(0.5,dtype=mx.float32)
|
|
sliced=mx.slice(a,idx,axes=axes,slice_size=size)
|
|
changed=mx.slice_update(a,upd,idx,axes=axes)
|
|
# Reuse outputs as subsequent input, mirroring moving cache windows.
|
|
back=mx.slice(changed,idx,axes=axes,slice_size=size)
|
|
emit(f"dynamic_slice_graph_{dtype_tag}_{i}",[sliced,changed,back],dtype=dtype_tag,idx_dtype=str(idx_dtype).split('.')[-1],profile=i,shape=shape,size=size,axes=axes,values=values,kind=kind,shapes=[x.shape for x in (sliced,changed,back)])
|
|
|
|
|
|
def qsa_prefill_graph_fixtures(dtype_tag):
|
|
"""Unmodified producer-selected prefill plus standalone strided top-k."""
|
|
from mtplx.kernels.qsa_indexer_prefill import (qsa_indexer_prefill_metal,
|
|
qsa_indexer_prefill_scores_mpp_supported, qsa_indexer_prefill_score_chunk_rows,
|
|
qsa_indexer_prefill_topk_metal)
|
|
assert mx.__version__=="0.32.2"
|
|
tf32=os.environ.get("MLX_ENABLE_TF32","1")!="0"
|
|
profiles=[(2,4,128,513,512,4,2051,"plain",128*1024*1024),
|
|
(65,4,128,1024,512,4,2051,"strided",1024*4*40),
|
|
(17,4,128,37,7,3,80,"transpose",128*1024*1024),
|
|
(5,2,7,17,3,3,9,"broadcast",128*1024*1024),
|
|
(9,1,8,33,7,1,15,"mixed",33*4*3),
|
|
(3,4,128,1,1,4,0,"plain",1),
|
|
(33,4,128,64,7,4,128,"mixed",64*4*40)]
|
|
for i,(rows,heads,dim,blocks,topk,ratio,pos,kind,budget) in enumerate(profiles):
|
|
dtype=getattr(mx,dtype_tag)
|
|
pdtype=(mx.float32 if i==4 else mx.float16 if dtype!=mx.float16 else mx.bfloat16) if kind=="mixed" else dtype
|
|
def source(shape,salt,dtype):
|
|
count=int(np.prod(shape));a=pattern(count,salt).astype(dtype)
|
|
if dtype==mx.float32:a=a+mx.array((np.arange(count)%7).astype(np.float32)/65536)
|
|
return a.reshape(shape)
|
|
qs=(1,rows,dim,heads) if kind=="transpose" else (1,1 if kind=="broadcast" else rows,heads,dim*(2 if kind=="strided" else 1))
|
|
ps=(1,dim,blocks) if kind=="transpose" else (1,blocks,dim*(2 if kind=="strided" else 1))
|
|
q=source(qs,211,dtype);p=source(ps,212,pdtype)
|
|
if kind=="transpose":q=q.swapaxes(-1,-2);p=p.swapaxes(-1,-2)
|
|
if kind=="strided":q=q[...,::2];p=p[...,::2]
|
|
if kind=="broadcast":q=mx.broadcast_to(q,(1,rows,heads,dim))
|
|
mpp=qsa_indexer_prefill_scores_mpp_supported(q,p)
|
|
chunk=qsa_indexer_prefill_score_chunk_rows(rows,heads,blocks,budget,producer="mpp" if mpp else "mlx")
|
|
total=pos+rows
|
|
for mode in ("blocks","row_tokens","dense_mask"):
|
|
out=qsa_indexer_prefill_metal(q,p,pos_start=pos,total_tokens=total,logical_blocks=total//ratio,block_topk=topk,compress_ratio=ratio,output_total_tokens=total+3,mode=mode,score_workspace_bytes=budget)
|
|
outputs=[out] if mode=="dense_mask" else list(out)
|
|
emit(f"qsa_prefill_graph_{dtype_tag}_{i}_{mode}_tf{int(tf32)}",outputs,dtype=dtype_tag,pdtype=str(pdtype).split('.')[-1],tf32=tf32,rows=rows,heads=heads,dim=dim,blocks=blocks,topk=topk,ratio=ratio,pos=pos,kind=kind,budget=budget,mpp=mpp,chunk=chunk,mode=mode,shapes=[a.shape for a in outputs])
|
|
if dtype_tag=="float32":
|
|
for i,(rows,blocks,topk,ratio,pos,kind) in enumerate(((1,8193,512,1,8192,"stride"),(3,2049,7,4,8190,"transpose"),(2,65,3,3,7,"broadcast"))):
|
|
shape=(blocks,rows) if kind=="transpose" else (1 if kind=="broadcast" else rows,blocks*2 if kind=="stride" else blocks)
|
|
scores=pattern(int(np.prod(shape)),213).astype(mx.float32).reshape(shape)
|
|
if kind=="transpose":scores=scores.T
|
|
if kind=="stride":scores=scores[:,::2]
|
|
if kind=="broadcast":scores=mx.broadcast_to(scores,(rows,blocks))
|
|
for mode in ("blocks","row_tokens","dense_mask"):
|
|
total=pos+rows
|
|
out=qsa_indexer_prefill_topk_metal(scores,pos_start=mx.array(pos,dtype=mx.int32),total_tokens=mx.array(total,dtype=mx.int32),logical_blocks=mx.array(total//ratio,dtype=mx.int32),block_topk=topk,compress_ratio=ratio,output_total_tokens=total+3,mode=mode)
|
|
outputs=[out] if mode=="dense_mask" else list(out)
|
|
emit(f"qsa_prefill_topk_graph_{i}_{mode}_tf{int(tf32)}",outputs,tf32=tf32,rows=rows,blocks=blocks,topk=topk,ratio=ratio,pos=pos,kind=kind,mode=mode,shapes=[a.shape for a in outputs])
|
|
|
|
|
|
def qsa_fused_graph_fixtures(dtype_tag):
|
|
"""Original selector factory and actual QSAIndexer chunking, no replacements."""
|
|
from types import SimpleNamespace
|
|
from mtplx.models.qwen4_exp import QSAIndexer
|
|
from mtplx.kernels.qsa_indexer_select import qsa_indexer_select_metal
|
|
assert mx.__version__=="0.32.2"
|
|
tf32=os.environ.get("MLX_ENABLE_TF32","1")!="0"
|
|
profiles=[(1,4,128,513,512,4,2051,"plain"),
|
|
(8,4,128,1024,512,4,2051,"strided"),
|
|
(3,1,8,17,7,1,9,"transpose"),
|
|
(5,2,7,5,3,3,7,"broadcast"),
|
|
(2,4,64,0,1,4,0,"plain"),
|
|
(2,4,64,5,1,4,0,"plain"),
|
|
(129,1,8,65537,7,4,1024,"chunk")]
|
|
for i,(rows,heads,dim,blocks,topk,ratio,pos,kind) in enumerate(profiles):
|
|
qdtype=getattr(mx,dtype_tag)
|
|
pdtype=mx.float32 if i==1 else (mx.float16 if i==2 else qdtype)
|
|
def source(shape,salt,dtype):
|
|
count=int(np.prod(shape))
|
|
a=pattern(count,salt).astype(dtype)
|
|
if dtype==mx.float32: a=a+mx.array((np.arange(count)%7).astype(np.float32)/65536)
|
|
return a.reshape(shape)
|
|
qshape=(1,rows,dim,heads) if kind=="transpose" else (1,1 if kind=="broadcast" else rows,heads,dim*(2 if kind=="strided" else 1))
|
|
pshape=(1,dim,blocks) if kind=="transpose" else (1,blocks,dim*(2 if kind=="strided" else 1))
|
|
q=source(qshape,201,qdtype);p=source(pshape,202,pdtype)
|
|
if kind=="transpose": q=q.swapaxes(-1,-2);p=p.swapaxes(-1,-2)
|
|
if kind=="strided": q=q[...,::2];p=p[...,::2]
|
|
if kind=="broadcast": q=mx.broadcast_to(q,(1,rows,heads,dim))
|
|
total=pos+rows;logical=total//ratio
|
|
for mode in ("blocks","row_tokens","dense_mask"):
|
|
if kind=="chunk":
|
|
indexer=SimpleNamespace(block_topk=topk,ratio=ratio,_fused_score_scratch_bytes=32*1024*1024)
|
|
# Bind the original helper, not a fixture implementation.
|
|
from types import MethodType
|
|
indexer._fused_query_chunk_rows=MethodType(QSAIndexer._fused_query_chunk_rows,indexer)
|
|
out=QSAIndexer._select_fused(indexer,q,pos,p,logical,total,mode)
|
|
else:
|
|
out=qsa_indexer_select_metal(q,p,pos_start=pos,total_tokens=total,logical_blocks=logical,
|
|
block_topk=topk,compress_ratio=ratio,output_total_tokens=total+3 if mode=="dense_mask" else None,mode=mode)
|
|
outputs=[out] if mode=="dense_mask" else list(out)
|
|
emit(f"qsa_fused_graph_{dtype_tag}_{i}_{mode}_tf{int(tf32)}",outputs,dtype=dtype_tag,pdtype=str(pdtype).split('.')[-1],
|
|
tf32=tf32,profile=i,rows=rows,heads=heads,dim=dim,blocks=blocks,topk=topk,ratio=ratio,pos=pos,kind=kind,mode=mode,shapes=[a.shape for a in outputs])
|
|
|
|
|
|
def qsa_indexer_prefix_fixtures(dtype_tag):
|
|
"""Unmodified QSAIndexer call, stateful prefix/rollback, actual cache."""
|
|
import mlx.nn as nn
|
|
from mtplx.models.qwen4_exp import QSAIndexer, QSACache
|
|
assert mx.__version__=="0.32.2"
|
|
tf32=os.environ.get("MLX_ENABLE_TF32","1")!="0"
|
|
dtype=getattr(mx,dtype_tag)
|
|
profiles=[(128,64,False,"dense"),(128,128,True,"dense"),(64,32,True,"supplied"),
|
|
(256,64,True,"dense"),(128,64,True,"promoted")]
|
|
if dtype_tag=="bfloat16": profiles += [(128,64,False,"q4"),(128,64,True,"q8")]
|
|
for profile,(dim,rot,fused,kind) in enumerate(profiles):
|
|
# Real supported reference configuration: eager selection with either
|
|
# original preparation branch. S>=2 is not legacy-fused eligible;
|
|
# decode-gather excludes that selector on S1. No method replacements.
|
|
for key,value in dict(MTPLX_COMPILED_QSA_INDEXER=0,MTPLX_FUSED_QSA_INDEXER=int(fused),
|
|
MTPLX_QSA_PREFILL=0,MTPLX_QSA_PREFILL_MIN_ROWS=2,MTPLX_QSA_GATHER=1,
|
|
MTPLX_QSA_GATHER_MIN_CONTEXT=0,MTPLX_QSA_GATHER_MAX_ROWS=8,
|
|
MTPLX_QSA_GATHER_DECODE=1,MTPLX_QSA_FLASH=0,MTPLX_QSA_SCORE_TILE_ROWS=0).items(): os.environ[key]=str(value)
|
|
indexer=QSAIndexer.__new__(QSAIndexer)
|
|
nn.Module.__init__(indexer)
|
|
indexer.n_heads=4;indexer.kv_heads=1;indexer.head_dim=dim;indexer.ratio=4
|
|
indexer.block_topk=512;indexer.budget=2048;indexer.rms_norm_eps=1e-6
|
|
indexer._inv_freq=mx.array([1/(j+1) for j in range(rot//2)],dtype=mx.float32)
|
|
indexer._rope_attention_scaling=1.125
|
|
indexer.q_layernorm=nn.RMSNorm(dim,eps=1e-6)
|
|
indexer.k_layernorm=nn.RMSNorm(dim,eps=1e-6)
|
|
norm_dtype=mx.float32 if kind=="promoted" else dtype
|
|
indexer.q_layernorm.weight=pattern(dim,185).astype(norm_dtype)
|
|
indexer.k_layernorm.weight=pattern(dim,186).astype(norm_dtype)
|
|
if kind in ("q4","q8"):
|
|
indexer.index_qk_proj=quantized_linear_fixture(5*dim,64,int(kind[1:]),64,187)
|
|
else:
|
|
indexer.index_qk_proj=nn.Linear(64,5*dim,bias=False)
|
|
indexer.index_qk_proj.weight=pattern(5*dim*64,187).astype(dtype).reshape(5*dim,64)
|
|
cache=QSACache()
|
|
cache.reserve_indexer_capacity(raw_capacity=256,pooled_capacity=256)
|
|
for step,(rows,trim) in enumerate(((3,0),(1,0),(5,0),(2041,0),(1,0),(4,0),(8,5),(33,0),(1,0))):
|
|
if trim: cache.trim(trim)
|
|
pos=cache.offset
|
|
hidden=pattern(rows*64,190+step*7).astype(dtype).reshape(1,rows,64)
|
|
qk=None
|
|
if kind=="supplied":
|
|
qk=pattern(rows*5*dim*2,191+step*7).astype(dtype).reshape(1,rows,5*dim*2)[...,::2]
|
|
result=indexer(hidden,pos,cache,qk_rows=qk)
|
|
outputs=[cache.raw_keys]
|
|
names=["raw"]
|
|
if cache.pooled is not None: outputs += [cache.pooled,cache.pooled_f32_t];names += ["pooled","mirror"]
|
|
if result is None: lane="none"
|
|
elif isinstance(result,tuple): lane=result[0];outputs+=list(result[1:]);names+=["indices","valid"]
|
|
else: lane="decode_tokens" if result.ndim==1 else "dense";outputs.append(result);names.append("selection")
|
|
emit(f"qsa_indexer_prefix_{dtype_tag}_{profile}_{step}_tf{int(tf32)}",outputs,
|
|
dtype=dtype_tag,tf32=tf32,profile=profile,dim=dim,rot=rot,fused=fused,kind=kind,
|
|
step=step,rows=rows,trim=trim,pos=pos,offset=cache.offset,pooled_len=cache.pooled_len,
|
|
lane=lane,names=names,shapes=[a.shape for a in outputs])
|
|
assert cache.offset==pos # Only Attention advances this frontier.
|
|
cache.kv.offset=pos+rows
|
|
|
|
|
|
def qsa_eager_graph_fixtures(dtype_tag):
|
|
"""Actual QSAIndexer._select_eager with its actual QSACache and all exits."""
|
|
from types import SimpleNamespace, MethodType
|
|
from mtplx.models.qwen4_exp import QSAIndexer, QSACache, _qsa_large_prefill_enabled
|
|
from mtplx.attention_context import attention_phase
|
|
assert mx.__version__=="0.32.2"
|
|
tf32=os.environ.get("MLX_ENABLE_TF32","1")!="0"
|
|
profiles=[(1,2051,4,512,0,m) for m in ("dense","flash","decode","all")]
|
|
profiles += [(8,2052,4,512,t,m) for t,m in ((0,"dense"),(0,"rows"),(3,"rows"),(8,"rows"),(9,"rows"),(0,"min"),(0,"max"))]
|
|
profiles += [(137,2052,4,512,t,m) for t,m in ((0,"prefill"),(31,"prefill"),(0,"all"),(31,"all"),(0,"verify"))]
|
|
profiles += [(2048,4,4,512,0,"dense"),(2,32768,4,512,0,"prefill"),
|
|
(33,32768,4,512,8,"prefill"),(1,4096,4,512,0,"decode"),
|
|
(8,2052,8,512,0,"prefill"),(3,13,1,7,0,"rows"),
|
|
(3,13,1,7,1,"dense"),(1,13,1,7,0,"flash")]
|
|
for i,(rows,pos,ratio,topk,tile,mode) in enumerate(profiles):
|
|
flash=mode in ("flash","all")
|
|
decode=mode in ("decode","all")
|
|
gather=mode in ("rows","all","min","max")
|
|
prefill=mode in ("prefill","all","verify")
|
|
phase="decode_verify" if mode=="verify" else "prefill"
|
|
minimum=16384 if mode=="min" else 0
|
|
maximum=7 if mode=="max" else 8
|
|
for key,value in dict(MTPLX_QSA_FLASH=int(flash),MTPLX_QSA_GATHER_DECODE=int(decode),
|
|
MTPLX_QSA_GATHER=int(gather),MTPLX_QSA_GATHER_MIN_CONTEXT=minimum,MTPLX_QSA_GATHER_MAX_ROWS=maximum,
|
|
MTPLX_QSA_PREFILL=int(prefill),MTPLX_QSA_PREFILL_MIN_ROWS=2,MTPLX_QSA_PREFILL_MIN_CONTEXT=2049,
|
|
MTPLX_QSA_SCORE_TILE_ROWS=tile).items(): os.environ[key]=str(value)
|
|
total=pos+rows
|
|
nb=total//ratio
|
|
dtype=getattr(mx,dtype_tag)
|
|
q=pattern(rows*4*128,181).astype(dtype)
|
|
if dtype==mx.float32: q=q+mx.array((np.arange(q.size)%7).astype(np.float32)/65536)
|
|
q=q.reshape(1,rows,4,128)
|
|
pooled=pattern(nb*128,182).astype(dtype).reshape(1,nb,128)
|
|
cache=QSACache(compress_ratio=ratio)
|
|
cache.write_pooled(pooled,0,nb)
|
|
indexer=SimpleNamespace(head_dim=128,ratio=ratio,block_topk=topk)
|
|
indexer._tiled_topk=MethodType(QSAIndexer._tiled_topk,indexer)
|
|
with attention_phase(phase):
|
|
large=_qsa_large_prefill_enabled(rows,total)
|
|
result=QSAIndexer._select_eager(indexer,q,pos,cache,pooled,total)
|
|
tail=None
|
|
if isinstance(result,tuple):
|
|
lane=result[0]
|
|
if lane=="flash": outputs=[result[1]];tail=result[2]
|
|
else: outputs=list(result[1:])
|
|
else:
|
|
lane="decode_tokens" if result.ndim==1 else "dense"
|
|
outputs=[result]
|
|
emit(f"qsa_eager_graph_{dtype_tag}_{i}_tf{int(tf32)}",outputs,dtype=dtype_tag,tf32=tf32,
|
|
rows=rows,pos=pos,ratio=ratio,topk=topk,tile=tile,mode=mode,large=large,flash=flash,decode=decode,
|
|
gather=gather,minimum=minimum,maximum=maximum,lane=lane,tail=tail,shapes=[a.shape for a in outputs])
|
|
|
|
|
|
def sort_graph_fixtures(dtype_tag):
|
|
"""All four original sort primitives, axis/layout and merge boundaries."""
|
|
assert mx.__version__ == "0.32.2"
|
|
dtype=getattr(mx,dtype_tag)
|
|
profiles=[((2,w),-1,"plain") for w in (1,128,129,256,257,512,513,1024,1025,2048,2049,4097,8193)]
|
|
profiles += [((3,7,5),axis,kind) for axis,kind in ((0,"plain"),(1,"plain"),(2,"transpose"),(1,"strided"),(-1,"reverse"),(1,"broadcast"),(None,"transpose"))]
|
|
profiles += [((2049,3),0,"plain"),((2,2049,3),1,"strided"),((4097,),0,"offset")]
|
|
for i,(shape,axis,kind) in enumerate(profiles):
|
|
physical=list(shape)
|
|
if kind=="strided": physical[-1]*=2
|
|
if kind=="broadcast": physical[1]=1
|
|
count=int(np.prod(physical))+(kind=="offset")
|
|
values=((np.arange(count,dtype=np.int64)*17+31)%257)
|
|
if dtype_tag not in ("uint32","uint64"): values-=128
|
|
a=mx.array(values).astype(dtype)
|
|
if kind=="offset": a=a[1:]
|
|
a=a.reshape(physical)
|
|
if kind=="transpose": a=a.swapaxes(0,-1)
|
|
if kind=="strided": a=a[...,::2]
|
|
if kind=="reverse": a=a[...,::-1]
|
|
if kind=="broadcast": a=mx.broadcast_to(a,shape)
|
|
width=a.size if axis is None else a.shape[axis]
|
|
kth=-min(3,width)
|
|
outputs=[mx.sort(a,axis=axis),mx.argsort(a,axis=axis),mx.partition(a,kth,axis=axis),mx.argpartition(a,kth,axis=axis)]
|
|
emit(f"sort_graph_{dtype_tag}_{i}",outputs,dtype=dtype_tag,profile=i,shape=shape,axis=axis,kind=kind,kth=kth,output_shape=outputs[0].shape)
|
|
|
|
|
|
def dense_float_fixtures(dtype_tag=None, start=0, end=None):
|
|
"""Pinned Matmul factory/backend including scalar, vector and empty cases."""
|
|
assert mx.__version__ == "0.32.2"
|
|
tf32=os.environ.get("MLX_ENABLE_TF32","1")!="0"
|
|
profiles=[(m,n,k,"plain",ta,tb) for m,n,k in ((3,5,128),(16,17,129),(40,40,127),(40,40,128))
|
|
for ta in (False,True) for tb in (False,True)]
|
|
profiles += [(m,n,k,kind,ta,tb) for m,n,k,kind,ta,tb in (
|
|
(1,1,7,"vectors",False,False),(1,1,16385,"vectors",False,False),
|
|
(1,17,64,"left-vector",False,False),(1,17,65,"left-vector",False,True),
|
|
(7,1,65,"right-vector",False,False),(7,1,65,"right-vector",True,False),
|
|
(1,1,129,"cross",False,False),(7,1,129,"cross",True,False),
|
|
(5,7,64,"cross",False,True),(5,7,64,"copy",False,False),
|
|
(5,7,64,"offset",False,True),(5,7,64,"collapse",False,True),
|
|
(5,7,64,"flatten",False,True),(65,513,128,"plain",False,True),
|
|
(2,2,4097,"plain",False,False),(193,17,8193,"plain",False,True),
|
|
(0,5,7,"plain",False,False),(5,0,7,"plain",False,False),
|
|
(5,7,0,"plain",False,False),(0,5,0,"flatten",False,False),
|
|
(3,7,128,"mixed",False,True),(3,7,128,"mixed-f32",False,True),
|
|
)]
|
|
for dtype in (mx.bfloat16,mx.float16,mx.float32):
|
|
tag=str(dtype).split('.')[-1]
|
|
if dtype_tag is not None and tag!=dtype_tag: continue
|
|
for i,(m,n,k,kind,ta,tb) in enumerate(profiles):
|
|
if i<start or (end is not None and i>=end): continue
|
|
def source(shape,salt,typ):
|
|
physical=list(shape)
|
|
if kind=="copy": physical[-1]*=2
|
|
count=int(np.prod(physical))+(kind=="offset")
|
|
a=pattern(count,salt).astype(typ)
|
|
if typ==mx.float32: a=a+mx.array((np.arange(count)%7).astype(np.float32)/65536)
|
|
if kind=="offset": a=a[1:]
|
|
a=a.reshape(physical)
|
|
return a[...,::2] if kind=="copy" else a
|
|
abatch=(2,1) if kind=="cross" else ((3,) if kind in ("copy","offset","collapse","flatten") else ())
|
|
bbatch=(1,3) if kind=="cross" else ((3,) if kind in ("copy","offset") else ((1,) if kind=="collapse" else ()))
|
|
adtype=mx.bfloat16 if kind.startswith("mixed") else dtype
|
|
bdtype=(mx.float32 if kind=="mixed-f32" else mx.float16) if kind.startswith("mixed") else dtype
|
|
a=source(abatch+((k,m) if ta else (m,k)),171,adtype)
|
|
b=source(bbatch+((n,k) if tb else (k,n)),172,bdtype)
|
|
if ta: a=a.swapaxes(-2,-1)
|
|
if tb: b=b.swapaxes(-2,-1)
|
|
if kind in ("left-vector","vectors"): a=a.reshape(k)
|
|
if kind in ("right-vector","vectors"): b=b.reshape(k)
|
|
out=a@b
|
|
emit(f"dense_float_{tag}_{i}_tf{int(tf32)}",[out],dtype=tag,profile=i,m=m,n=n,k=k,kind=kind,ta=ta,tb=tb,tf32=tf32,
|
|
a_dtype=str(adtype).split('.')[-1],b_dtype=str(bdtype).split('.')[-1],output_shape=out.shape)
|
|
|
|
|
|
def sum_graph_fixtures(dtype_tag=None, start=0, end=None):
|
|
"""Original runtime reductions used by QSA/HC/PLE/GDN; no custom sums."""
|
|
assert mx.__version__ == "0.32.2"
|
|
profiles = [
|
|
((4096,), (0,), "plain"), ((4097,), (0,), "offset"),
|
|
((31,128), (-1,), "plain"), ((32,128), (-1,), "plain"),
|
|
((32,512), (1,), "plain"), ((32,513), (1,), "plain"),
|
|
((32,1024), (1,), "plain"), ((32,1025), (1,), "plain"),
|
|
((2,4097), (1,), "plain"), ((5,7), (1,), "plain"),
|
|
((1,1024,4,128), (2,), "plain"), ((1,137,4,513), (2,), "plain"),
|
|
((31,128), (0,), "plain"), ((32,128), (0,), "plain"),
|
|
((256,128), (0,), "plain"), ((257,128), (0,), "plain"),
|
|
((1023,17), (0,), "plain"), ((1024,17), (0,), "plain"),
|
|
((32768,2), (0,), "plain"),
|
|
((3,5,7,128), (0,2), "plain"),
|
|
((3,5,7,9,2,128), (0,2,4), "plain"),
|
|
((3,5,7,128), (0,3), "plain"),
|
|
((3,5,7,9,2,128), (0,2,5), "plain"),
|
|
((64,7,4), (0,2), "plain"),
|
|
((64,7,16), (0,2), "plain"),
|
|
((5,7,128), (-1,), "transpose"),
|
|
((5,7,128), (1,), "transpose"),
|
|
((5,7,128), (0,1,2), "transpose"),
|
|
((5,7,128), (-1,), "strided"),
|
|
((5,7,128), (0,), "strided"),
|
|
((5,7,128), (1,), "reverse"),
|
|
((5,7,128), (-1,), "broadcast"),
|
|
((5,7,128), (0,), "broadcast"),
|
|
((5,7,128), (0,1,2), "scalar"),
|
|
((1,7,1), (0,2), "offset"),
|
|
((0,7,128), (0,), "plain"),
|
|
((5,0,128), (-1,), "plain"),
|
|
((5,7,128), (), "plain"),
|
|
((257,32768), (0,), "plain"),
|
|
((16777217,), (0,), "plain"),
|
|
]
|
|
for dtype in (mx.bfloat16,mx.float16,mx.float32):
|
|
tag=str(dtype).split('.')[-1]
|
|
if dtype_tag is not None and tag != dtype_tag: continue
|
|
for i,(shape,axes,layout) in enumerate(profiles):
|
|
if i<start or (end is not None and i>=end): continue
|
|
if i>=38 and dtype!=mx.float32: continue
|
|
base=list(shape)
|
|
if layout=="broadcast": base[0]=1
|
|
if layout=="strided": base[-1]*=2
|
|
n=int(np.prod(base))
|
|
a=pattern(1 if layout=="scalar" else n+(layout=="offset"),170).astype(dtype)
|
|
if layout=="offset": a=a[1:]
|
|
a=mx.broadcast_to(a.reshape(()),shape) if layout=="scalar" else a.reshape(base)
|
|
if layout=="strided": a=a[...,::2]
|
|
if layout=="reverse": a=a[...,::-1]
|
|
if layout=="broadcast": a=mx.broadcast_to(a,shape)
|
|
if layout=="transpose": a=a.swapaxes(0,1)
|
|
for keep in (False,True):
|
|
out=mx.sum(a,axis=axes,keepdims=keep)
|
|
mean=mx.mean(a.astype(mx.float32),axis=axes,keepdims=keep)
|
|
emit(f"sum_graph_{tag}_{i}_{keep}",[out,mean],dtype=tag,profile=i,shape=shape,
|
|
axes=axes,layout=layout,keep=keep,output_shape=out.shape)
|
|
|
|
|
|
def qsa_dense_mask_fixtures():
|
|
from mtplx.models.qwen4_exp import _qsa_blocks_to_dense_mask
|
|
assert mx.__version__ == "0.32.2"
|
|
for rows,pos,ratio,topk in ((1,0,4,0),(1,0,4,8),(2,1,4,8),(5,0,4,9),
|
|
(8,3,4,17),(9,2045,4,512),(137,32767,4,512),
|
|
(2048,2048,4,512),(5,2,1,9),(8,17,7,8),(0,0,4,0)):
|
|
total=pos+rows
|
|
for strided in (False,True):
|
|
for empty in (False,True):
|
|
cols=topk*(2 if strided else 1)
|
|
ids=[];valid=[]
|
|
for r in range(rows):
|
|
for c in range(cols):
|
|
slot=c//2 if strided else c
|
|
candidate=(pos+r+1)//ratio-1-slot//2
|
|
ids.append(total if slot%11==5 else (-1 if slot%11==3 else candidate))
|
|
valid.append(not empty and slot%7!=2)
|
|
ids=mx.array(ids,dtype=mx.int32).reshape(rows,cols)
|
|
valid=mx.array(valid,dtype=mx.bool_).reshape(rows,cols)
|
|
if strided: ids,valid=ids[:,::2],valid[:,::2]
|
|
out=_qsa_blocks_to_dense_mask(ids,valid,pos_start=pos,total_tokens=total,compress_ratio=ratio)
|
|
emit(f"qsa_dense_mask_{rows}_{pos}_{ratio}_{topk}_{strided}_{empty}",[out],
|
|
rows=rows,pos=pos,ratio=ratio,topk=topk,strided=strided,empty=empty)
|
|
|
|
|
|
def scatter_axis_graph_fixtures():
|
|
assert mx.__version__ == "0.32.2"
|
|
for shape,axis in (((7,),0),((3,5),0),((3,5),1),((2,3,5),1)):
|
|
for strided in (False,True):
|
|
for scalar in (False,True):
|
|
for src_layout in ("plain","scalar","offset","strided"):
|
|
n=int(np.prod(shape))
|
|
if src_layout=="scalar": src=mx.broadcast_to(mx.array(False),shape)
|
|
else:
|
|
mult=2 if src_layout=="strided" else 1
|
|
src=(mx.arange(n*mult+1,dtype=mx.int32)%3==0)
|
|
src=src[1:] if src_layout=="offset" else src[:n*mult]
|
|
if src_layout=="strided": src=src[::2]
|
|
src=src.reshape(shape)
|
|
ishape=list(shape);ishape[axis]=2
|
|
size=int(np.prod(ishape))
|
|
mult=2 if strided else 1
|
|
idx=mx.arange(size*mult,dtype=mx.int64)//mult%2
|
|
idx=mx.where(idx==0,-1,0)
|
|
if strided: idx=idx[::2]
|
|
idx=idx.reshape(ishape)
|
|
upd=mx.array(True) if scalar else (mx.arange(size,dtype=mx.int32)%3==0).reshape(ishape)
|
|
out=mx.put_along_axis(src,idx,upd,axis=axis)
|
|
emit(f"scatter_axis_{shape}_{axis}_{strided}_{scalar}_{src_layout}",[out],
|
|
shape=shape,axis=axis,strided=strided,scalar=scalar,src_layout=src_layout)
|
|
|
|
|
|
def sparse_attention_graph_fixtures():
|
|
"""Actual MTPLX sparse attention consumers, never a dense substitute."""
|
|
from mtplx.kernels.qsa_flash_skip import qsa_flash_skip
|
|
from mtplx.kernels.qsa_prefill_flash import qsa_prefill_flash, qsa_prefill_flash_supported
|
|
assert mx.__version__ == "0.32.2"
|
|
for dtype in (mx.bfloat16,mx.float16):
|
|
tag=str(dtype).split('.')[-1]
|
|
for heads,total,nsel in ((24,3,0),(24,4099,512),(2,2052,17)):
|
|
for layout in ("plain","offset","strided"):
|
|
cap=total+7
|
|
def source(shape,salt):
|
|
n=int(np.prod(shape))
|
|
if layout=="offset": return pattern(n+1,salt).astype(dtype)[1:].reshape(shape)
|
|
if layout=="strided": return pattern(n*2,salt).astype(dtype).reshape(*shape[:-1],shape[-1]*2)[...,::2]
|
|
return pattern(n,salt).astype(dtype).reshape(shape)
|
|
q=source((heads,256),111)
|
|
k,v=source((1,2,cap,256),112),source((1,2,cap,256),113)
|
|
ids=mx.arange(nsel*2 if layout=="strided" else nsel,dtype=mx.int32)
|
|
if layout=="strided": ids=ids[::2]
|
|
else: ids=ids*2
|
|
tail=(total//4)*4
|
|
out=qsa_flash_skip(q,k,v,ids,total,tail,0.0625)
|
|
emit(f"sparse_skip_{tag}_{heads}_{total}_{layout}",[out],kind="skip",dtype=tag,heads=heads,total=total,nsel=nsel,layout=layout)
|
|
for rows,pos in ((2,2050),(8,8191),(137,32767),(2048,2048)):
|
|
for layout in ("plain","offset","strided","holes"):
|
|
total=pos+rows;cap=total+7
|
|
def source(shape,salt):
|
|
n=int(np.prod(shape))
|
|
if layout=="offset": return pattern(n+1,salt).astype(dtype)[1:].reshape(shape)
|
|
if layout=="strided": return pattern(n*2,salt).astype(dtype).reshape(*shape[:-1],shape[-1]*2)[...,::2]
|
|
return pattern(n,salt).astype(dtype).reshape(shape)
|
|
q=source((1,rows,24,256),114).transpose(0,2,1,3)
|
|
k,v=source((1,2,cap,256),115),source((1,2,cap,256),116)
|
|
slots=mx.arange(512,dtype=mx.int32)[None,:]
|
|
ids=mx.broadcast_to(slots,(rows,512))
|
|
valid=mx.broadcast_to((slots<17) | ((slots==511) if layout=="holes" else mx.array(False)),(rows,512))
|
|
if layout=="holes":
|
|
ids=mx.where(slots==3,-1,mx.where(slots==5,total,ids))
|
|
if layout=="strided":
|
|
ids=mx.stack([ids,ids],axis=-1).reshape(rows,1024)[:,::2]
|
|
valid=mx.stack([valid,valid],axis=-1).reshape(rows,1024)[:,::2]
|
|
assert qsa_prefill_flash_supported(q,k,v,ids,valid,pos_start=pos,total_tokens=total,scale=0.0625)
|
|
out=qsa_prefill_flash(q,k,v,ids,valid,pos_start=pos,total_tokens=total,scale=0.0625)
|
|
emit(f"sparse_prefill_{tag}_{rows}_{pos}_{layout}",[out],kind="prefill",dtype=tag,rows=rows,pos=pos,total=total,layout=layout)
|
|
|
|
|
|
def qsa_prepare_graph_fixtures(*, dtype_tag=None, profile_only=None):
|
|
"""Actual parameterized MTPLX query/pool kernels, including traced starts."""
|
|
from mtplx.kernels.qsa_indexer_prepare import (
|
|
qsa_indexer_prepare_queries_metal,qsa_indexer_pool_keys_metal,qsa_indexer_prepare_supported)
|
|
assert mx.__version__ == "0.32.2"
|
|
profiles=((1,4,128,64,4,1e-6,1.0,0),(3,1,2,2,1,1e-4,1.25,-7),
|
|
(7,3,17,16,7,0.0,0.5,2047),(137,7,64,32,3,1e-5,1.1,65535),
|
|
(2048,2,128,128,8,1e-6,1.3,131071),(1,4,128,64,4,0.25,1.0,0))
|
|
for dtype in (mx.bfloat16,mx.float16,mx.float32):
|
|
tag=str(dtype).split('.')[-1]
|
|
if dtype_tag is not None and tag!=dtype_tag: continue
|
|
for profile,(rows,heads,dim,rotary,ratio,eps,scale,pos) in enumerate(profiles):
|
|
if profile_only is not None and profile!=profile_only: continue
|
|
for layout in ("plain","offset","strided","broadcast"):
|
|
def source(shape,salt,ty=dtype):
|
|
n=int(np.prod(shape))
|
|
if layout=="broadcast": return mx.broadcast_to(pattern(1,salt).astype(ty),shape)
|
|
if layout=="offset": return pattern(n+1,salt).astype(ty)[1:].reshape(shape)
|
|
if layout=="strided": return pattern(n*2,salt).astype(ty).reshape(*shape[:-1],shape[-1]*2)[...,::2]
|
|
return pattern(n,salt).astype(ty).reshape(shape)
|
|
q=source((1,rows,heads,dim),160)
|
|
keys=source((1,rows*ratio,dim),161)
|
|
norm=source((dim,),162)
|
|
freq=source((rotary//2,),163,mx.float32)
|
|
start=pos
|
|
if layout=="offset": start=mx.array([123,pos],dtype=mx.int32)[1:].reshape(1,1)
|
|
elif layout=="strided": start=mx.array([[pos+5]],dtype=mx.int32)-5
|
|
elif layout=="broadcast": start=mx.array(pos,dtype=mx.int32)
|
|
assert qsa_indexer_prepare_supported(q,norm,freq,expected_ndim=4)
|
|
assert qsa_indexer_prepare_supported(keys,norm,freq,expected_ndim=3)
|
|
out=qsa_indexer_prepare_queries_metal(q,norm,freq,pos_start=start,eps=eps,attention_scaling=scale)
|
|
pooled=qsa_indexer_pool_keys_metal(keys,norm,freq,block_start=start,compress_ratio=ratio,eps=eps,attention_scaling=scale)
|
|
emit(f"qsa_prepare_graph_{tag}_{profile}_{layout}",[out,pooled],dtype=tag,profile=profile,
|
|
layout=layout,rows=rows,heads=heads,dim=dim,rotary=rotary,ratio=ratio,eps=eps,scale=scale,pos=pos)
|
|
|
|
|
|
def qsa_cache_graph_fixtures(*, profile=None, deferred_only=None):
|
|
"""Actual positional QSACache methods, retaining Python array aliases."""
|
|
from types import SimpleNamespace
|
|
from mtplx.models.qwen4_exp import QSACache,QSAIndexer
|
|
assert mx.__version__ == "0.32.2"
|
|
profiles=[(n,4,128,mx.bfloat16) for n in (0,1,255,256,257,1025)]
|
|
profiles += [(257,1,17,mx.float16),(257,7,32,mx.float32)]
|
|
for initial,ratio,dim,dtype in profiles:
|
|
if profile is not None and (initial,ratio)!=tuple(profile): continue
|
|
for deferred in (False,True):
|
|
if deferred_only is not None and deferred!=deferred_only: continue
|
|
cache=QSACache(ratio)
|
|
outputs=[];snapshots=[];metadata=[]
|
|
tag=str(dtype).split('.')[-1]
|
|
actions=[("reserve",768,320),("append",initial,0),("snapshot",0,0),
|
|
("reserve",2049,769),("view",0,0),("append",3,0),
|
|
("trim",2,0),("append",4,0),("restore",0,0),("view",0,0),
|
|
("append",257,0),("trim",999999,0),("append",1,0),
|
|
("ensure",1025,0),("append",4,0)]
|
|
def source(shape,salt,strided):
|
|
n=int(np.prod(shape));mult=2 if strided else 1
|
|
a=pattern(n*mult,salt).astype(dtype).reshape(*shape[:-1],shape[-1]*mult)
|
|
return a[...,::2] if strided else a
|
|
def append(n,step):
|
|
start=cache.offset;total=start+n
|
|
cache.write_raw(source((1,n,dim),140+step*3,step%2==1))
|
|
old=min(cache.pooled_len,total//ratio)
|
|
if total//ratio>old:
|
|
cache.write_pooled(source((1,total//ratio-old,dim),141+step*3,step%2==1),old,total//ratio)
|
|
cache.kv.update_and_fetch(
|
|
source((1,n,2,256),142+step*3,False).transpose(0,2,1,3),
|
|
source((1,n,2,256),143+step*3,False).transpose(0,2,1,3))
|
|
for step,(action,a,b) in enumerate(actions):
|
|
extra={}
|
|
if action=="reserve": cache.reserve_indexer_capacity(raw_capacity=a,pooled_capacity=b)
|
|
elif action=="append": append(a,step)
|
|
elif action=="snapshot": snapshots.append(cache.state)
|
|
elif action=="trim": extra["trimmed"]=cache.trim(a)
|
|
elif action=="restore": cache.state=snapshots[a]
|
|
elif action=="view" and cache.pooled is not None:
|
|
outputs.append(cache.pooled_f32_view(cache.pooled_len))
|
|
elif action=="ensure":
|
|
indexer=SimpleNamespace(ratio=ratio,head_dim=dim)
|
|
raw,pooled=QSAIndexer._ensure_compiled_backings(indexer,cache,dtype=dtype,pos_start=cache.offset,rows=a)
|
|
assert raw is cache.raw_keys and pooled is cache.pooled
|
|
current=[cache.kv.keys,cache.kv.values,cache.raw_keys,cache.pooled,cache.pooled_f32_t]
|
|
if cache.kv.keys is not None: current.extend(cache.state)
|
|
outputs.extend(v for v in current if v is not None)
|
|
if not deferred: mx.eval(*(v for v in outputs if v is not None))
|
|
metadata.append(dict(action=action,a=a,b=b,offset=cache.offset,pooled_len=cache.pooled_len,
|
|
capacities=[0 if v is None else v.shape[axis] for v,axis in zip(current[:5],(2,2,1,1,3))],
|
|
reserved=[cache._reserved_raw_capacity,cache._reserved_pooled_capacity],nbytes=cache.nbytes,**extra))
|
|
emit(f"qsa_cache_graph_{initial}_{ratio}_{dim}_{tag}_{int(deferred)}",outputs,
|
|
initial=initial,ratio=ratio,dim=dim,dtype=tag,deferred=deferred,actions=metadata,
|
|
shapes=[list(v.shape) for v in outputs])
|
|
|
|
|
|
def kv_cache_graph_fixtures():
|
|
"""Actual KVCache mutations, retaining old states across lazy updates."""
|
|
from mlx_lm.models.cache import KVCache
|
|
for initial in (0,1,255,256,257,2048):
|
|
for deferred in (False,True):
|
|
cache=KVCache()
|
|
outputs=[]
|
|
snapshots=[]
|
|
metadata=[]
|
|
actions=[("append",initial),("append",3),("trim",2),("append",4),
|
|
("restore",0),("append",257),("trim",999999),("append",1)]
|
|
for step,(action,n) in enumerate(actions):
|
|
if action=="trim":
|
|
trimmed=cache.trim(n)
|
|
metadata.append(dict(action=action,n=n,trimmed=trimmed,offset=cache.offset))
|
|
continue
|
|
if action=="restore":
|
|
cache.state=snapshots[0]
|
|
metadata.append(dict(action=action,n=n,offset=cache.offset))
|
|
continue
|
|
k=pattern(n*2*256,104+step*2).reshape(1,n,2,256).transpose(0,2,1,3)
|
|
v=pattern(n*2*256,105+step*2).reshape(1,n,2,256).transpose(0,2,1,3)
|
|
state=cache.update_and_fetch(k,v)
|
|
snapshots.append(state)
|
|
current=[*state,cache.keys,cache.values]
|
|
outputs.extend(current)
|
|
if not deferred: mx.eval(*current)
|
|
metadata.append(dict(action=action,n=n,offset=cache.offset,capacity=cache.keys.shape[2]))
|
|
emit(f"kv_cache_graph_{initial}_{int(deferred)}",outputs,initial=initial,deferred=deferred,actions=metadata)
|
|
|
|
|
|
def rope_graph_fixtures():
|
|
"""Original Qwen RoPE functions, including static YaRN and vision axes."""
|
|
from types import SimpleNamespace
|
|
from mtplx.models.qwen4_exp import (_rope_inv_freq_and_scaling, _rope_cos_sin,
|
|
_mrope_cos_sin, _build_mrope_axes, _apply_partial_rope)
|
|
assert mx.__version__ == "0.32.2"
|
|
profiles = [
|
|
(64, 10000000., 1, 0, {}),
|
|
(256, 10000., 7, 511, {}),
|
|
(64, 10000000., 2048, 262143, dict(rope_type="yarn", factor=4., original_max_position_embeddings=262144)),
|
|
(64, 10000., 7, 1048576, dict(rope_type="yarn", factor=1., original_max_position_embeddings=128)),
|
|
(64, 10000., 64, 8191, dict(rope_type="yarn", factor=8., original_max_position_embeddings=4096,
|
|
mscale=2., mscale_all_dim=1., truncate=False, beta_fast=4., beta_slow=4.)),
|
|
(64, 10000., 7, 32767, dict(rope_type="yarn", factor=4., original_max_position_embeddings=128,
|
|
attention_factor=0.75, beta_fast=0., beta_slow=0.)),
|
|
]
|
|
for profile,(rot,base,rows,pos,parameters) in enumerate(profiles):
|
|
inv,scale = _rope_inv_freq_and_scaling(SimpleNamespace(rotary_dim=rot, rope_theta=base, rope_parameters=parameters))
|
|
heads=24 if rows<=7 else 2
|
|
half=rot//2
|
|
section=[(half+2)//3,(half+1)//3,half//3]
|
|
for mode in ("text","equal","interleaved","sections"):
|
|
axes=_build_mrope_axes(section,mode!="sections")
|
|
positions=mx.arange(pos,pos+rows,dtype=mx.int32)
|
|
positions3=mx.stack([positions,positions+(0 if mode=="equal" else 7),positions+(0 if mode=="equal" else 19)])
|
|
cos,sin=(_rope_cos_sin(positions,inv,scale) if mode=="text" else
|
|
_mrope_cos_sin(positions3,inv,mx.array(axes,mx.int32)))
|
|
for strided in (False,True):
|
|
width=512 if strided else 256
|
|
x=pattern(rows*heads*width,103).reshape(1,rows,heads,width)
|
|
if strided: x=x[...,1::2]
|
|
out=_apply_partial_rope(x,cos,sin)
|
|
emit(f"rope_graph_{profile}_{mode}_{int(strided)}",[inv,cos,sin,out],
|
|
rot=rot,base=base,rows=rows,pos=pos,parameters=parameters,heads=heads,
|
|
mode=mode,strided=strided,section=section,axes=axes,scale=scale)
|
|
|
|
|
|
def qsa_rows_attention_fixtures():
|
|
"""Actual MTPLX per-row gathered attention, not a reconstructed formula."""
|
|
from mtplx.models.qwen4_exp import _qsa_rows_gather_attention
|
|
assert mx.__version__ == "0.32.2"
|
|
for heads,rows,total,count in ((24,1,33,8),(24,3,1057,1025),(24,8,4193,4097),
|
|
(24,64,113,17),(2,3,33,8),(2,16,1057,1025)):
|
|
for layout in ("plain","offset","strided"):
|
|
q=pattern(rows*heads*256,97).reshape(1,rows,heads,256).transpose(0,2,1,3)
|
|
def cached(salt):
|
|
cap=total+7 if layout=="offset" else total*2 if layout=="strided" else total
|
|
a=pattern(2*cap*256,salt).reshape(1,2,cap,256)
|
|
return a[:,:,2:total+2] if layout=="offset" else a[:,:,::2] if layout=="strided" else a
|
|
k,v=cached(98),cached(99)
|
|
idx=(mx.arange(rows*count,dtype=mx.int32)*7%total).reshape(rows,count)
|
|
ok=(mx.arange(rows*count,dtype=mx.int32).reshape(rows,count)%5)!=1
|
|
emit(f"qsa_rows_attention_h{heads}_s{rows}_t{total}_k{count}_{layout}",
|
|
[_qsa_rows_gather_attention(q,k,v,idx,ok,0.0625)],
|
|
heads=heads,rows=rows,total=total,count=count,layout=layout)
|
|
|
|
|
|
def qsa_prefill_gather_fixtures():
|
|
"""Actual portable prefill function, with its explicit per-tile evals."""
|
|
from mtplx.models.qwen4_exp import _qsa_prefill_gather_attention
|
|
assert mx.__version__ == "0.32.2"
|
|
for heads,rows,pos,ratio,topk in ((24,1,0,4,2),(24,7,3,4,2),(24,17,127,128,2),
|
|
(24,17,511,128,2),(2,7,0,1,2),(2,2048,0,4,2)):
|
|
for tile in (1,4,64) if rows<2048 else (64,2048):
|
|
for strided in (False,True):
|
|
cap=pos+rows+7
|
|
q=pattern(rows*heads*256,100).reshape(1,rows,heads,256).transpose(0,2,1,3)
|
|
k=pattern(2*cap*256,101).reshape(1,2,cap,256)
|
|
v=pattern(2*cap*256,102).reshape(1,2,cap,256)
|
|
nb=(pos+np.arange(rows)+1)//ratio
|
|
raw_ids=nb[:,None]-1-np.arange(topk)[None,:]
|
|
raw_ok=raw_ids>=0
|
|
raw_ids=np.where(raw_ok,raw_ids,-999)
|
|
if strided:
|
|
ids=np.full((rows,topk*2),-888,dtype=np.int64);ids[:,1::2]=raw_ids
|
|
ok=np.zeros((rows,topk*2),dtype=np.bool_);ok[:,1::2]=raw_ok
|
|
ids=mx.array(ids)[:,1::2];ok=mx.array(ok)[:,1::2]
|
|
else:
|
|
ids=mx.array(raw_ids.astype(np.int32));ok=mx.array(raw_ok)
|
|
calls=[]; original_eval=mx.eval
|
|
def tracked(*a):
|
|
calls.append([list(x.shape) for x in a]);return original_eval(*a)
|
|
mx.eval=tracked
|
|
try:
|
|
out=_qsa_prefill_gather_attention(q,k,v,ids,ok,pos_start=pos,total_tokens=pos+rows,
|
|
compress_ratio=ratio,scale=0.0625,tile_rows=tile)
|
|
finally:
|
|
mx.eval=original_eval
|
|
emit(f"qsa_prefill_gather_h{heads}_s{rows}_p{pos}_r{ratio}_t{tile}_strided{int(strided)}",[out],
|
|
heads=heads,rows=rows,pos=pos,ratio=ratio,topk=topk,tile=tile,strided=strided,evals=calls)
|
|
|
|
|
|
def sdpa_vector_fixtures():
|
|
"""Actual MTPLX runtime fused D256 SDPA, including the 2-pass boundaries."""
|
|
assert mx.__version__ == "0.32.2"
|
|
assert not os.environ.get("MLX_SDPA_BLOCKS")
|
|
for batch,heads,rows,total in ((1,24,1,17),(1,24,2,1023),(1,24,2,1024),
|
|
(1,24,2,16384),(1,24,1,65536),(2,24,2,33)):
|
|
for copied in (False,True):
|
|
for mode in ("none","bool","add","causal"):
|
|
step=2 if copied else 1
|
|
q=pattern(batch*rows*heads*256*step,103).reshape(batch,rows,heads,256*step)[...,::step].transpose(0,2,1,3)
|
|
k=pattern(batch*2*(total+7)*256*step,104).reshape(batch,2,total+7,256*step)[:,:,2:total+2,::step]
|
|
v=pattern(batch*2*(total+7)*256*step,105).reshape(batch,2,total+7,256*step)[:,:,2:total+2,::step]
|
|
raw=np.arange(rows*total).reshape(rows,total)%3!=1
|
|
mask=mx.array(raw)[None,None] if mode=="bool" else mx.array(np.where(raw,0,-10),dtype=mx.bfloat16)[None,None] if mode=="add" else "causal" if mode=="causal" else None
|
|
emit(f"sdpa_vector_b{batch}_h{heads}_s{rows}_t{total}_copy{int(copied)}_{mode}",
|
|
[mx.fast.scaled_dot_product_attention(q,k,v,scale=0.0625,mask=mask)],
|
|
batch=batch,heads=heads,rows=rows,total=total,copied=copied,mode=mode)
|
|
|
|
|
|
def gather_array_fixtures():
|
|
"""Pinned GatherQMM on the actual expert/broadcast/sliced Array layouts."""
|
|
assert mx.__version__ == "0.32.2"
|
|
# x batches, rows per matrix, output width, input width, experts, routes,
|
|
# layout, explicit lhs, explicit rhs, sorted, route rank.
|
|
cases = [
|
|
(1,1,80,512,8,10,"plain",False,True,False,1),
|
|
(4,1,80,512,8,40,"expert_slice",False,True,False,2),
|
|
(7,1,81,64,8,70,"expert_slice",False,True,False,2),
|
|
(2,33,80,512,8,8,"plain",True,True,False,1),
|
|
(2,17,81,64,8,8,"transpose",True,True,False,1),
|
|
(2,33,80,512,8,8,"expert_slice",True,True,False,1),
|
|
(16,1,80,512,4,16,"plain",False,True,True,1),
|
|
(256,1,80,512,4,256,"expert_slice",False,True,True,1),
|
|
(16,1,80,512,4,16,"expert_slice",True,True,True,1),
|
|
(2,4,80,512,8,8,"plain",True,False,True,1),
|
|
(1,1,80,512,8,10,"rank2",False,True,False,1),
|
|
(4,1,80,512,8,40,"input_slice",False,True,False,3),
|
|
(32,1,80,512,4,32,"transpose",False,True,True,1),
|
|
]
|
|
for bits,group in ((4,32),(4,64),(8,32),(8,64)):
|
|
for case,(xb,m,n,k,e,r,layout,left,right,sorted_indices,rank) in enumerate(cases):
|
|
extra = 2 if layout == "expert_slice" else 1
|
|
wp=k*bits//32;sp=k//group
|
|
words=(np.arange(e*n*extra*wp,dtype=np.uint64)*2654435761+12345).astype(np.uint32)
|
|
w=mx.array(words).reshape(e,n*extra,wp)
|
|
s=(pattern(e*n*extra*sp,33)/64).reshape(e,n*extra,sp)
|
|
b=(pattern(e*n*extra*sp,34)/64).reshape(e,n*extra,sp)
|
|
if layout=="expert_slice":w=w[:,n:];s=s[:,n:];b=b[:,n:]
|
|
x=pattern(xb*m*k*(2 if layout=="input_slice" else 1),32)
|
|
if layout=="transpose":
|
|
x=x.reshape(xb,k,m).transpose(0,2,1)
|
|
w=w.reshape(e,wp,n).transpose(0,2,1)
|
|
s=s.reshape(e,sp,n).transpose(0,2,1)
|
|
b=b.reshape(e,sp,n).transpose(0,2,1)
|
|
elif layout=="input_slice":x=x.reshape(xb,m,k*2)[...,::2]
|
|
else:x=x.reshape(xb,m,k)
|
|
if rank==2:x=x.reshape(1,xb,1,m,k)
|
|
if rank==3:x=x.reshape(2,xb//2,1,m,k)
|
|
if layout=="rank2":x=x.reshape(m,k)
|
|
rhs=(np.arange(r,dtype=np.uint32)*5+2)%e
|
|
if sorted_indices:rhs=np.sort(rhs)
|
|
if rank>1:
|
|
raw=np.zeros((xb,r//xb+3),np.int32);raw[:,1:1+r//xb]=rhs.reshape(xb,-1)
|
|
rhs=mx.array(raw)[:,1:1+r//xb]
|
|
rhs=rhs.reshape((1,xb,r//xb) if rank==2 else (2,xb//2,r//xb))
|
|
else:rhs=mx.array(rhs,mx.int32)
|
|
lhs=mx.array((np.arange(r)*2+1)%xb,mx.int64) if left else None
|
|
out=mx.gather_qmm(x,w,s,b,lhs_indices=lhs,rhs_indices=rhs if right else None,
|
|
transpose=True,group_size=group,bits=bits,sorted_indices=sorted_indices)
|
|
emit(f"gather_array_{case}_b{bits}_g{group}",[out],case=case,xb=xb,m=m,n=n,k=k,e=e,r=r,
|
|
layout=layout,left=left,right=right,sorted=sorted_indices,rank=rank,bits=bits,group=group,shape=out.shape)
|
|
|
|
|
|
def emit(name, outputs, **metadata):
|
|
mx.eval(*outputs)
|
|
# Some original-runtime empty Matmul results have a null data pointer;
|
|
# numpy.asarray crashes although mx.eval succeeded. No bytes exist to read.
|
|
print(json.dumps({"kernel": name, "sha256": [hashlib.sha256(
|
|
b"" if value.size == 0 else np.asarray(value.astype(mx.float32)).astype('<f4').tobytes()).hexdigest()
|
|
for value in outputs], **metadata}, separators=(",", ":")), flush=True)
|
|
|
|
|
|
def quantized_linear_fixture(n, k, bits, group, salt):
|
|
import mlx.nn as nn
|
|
layer = nn.QuantizedLinear(64, 1, bias=False, group_size=group, bits=bits)
|
|
words = (np.arange(n*k*bits//32,dtype=np.uint64)*2654435761+12345).astype(np.uint32)
|
|
layer.weight = mx.array(words).reshape(n,k*bits//32)
|
|
layer.scales = (pattern(n*k//group,salt)/64).reshape(n,k//group)
|
|
layer.biases = (pattern(n*k//group,salt+1)/64).reshape(n,k//group)
|
|
return layer
|
|
|
|
|
|
def hyper_graph_fixtures():
|
|
"""Original quantizer + HC pack/primitives, with raw uint32 pack hashes."""
|
|
from types import SimpleNamespace
|
|
from mtplx.kernels.hyper_connection_v3 import prepare_v3_pack, device_supports_hyper_v3
|
|
assert mx.__version__ == "0.32.2"
|
|
assert device_supports_hyper_v3()
|
|
def packed(name, outputs, **metadata):
|
|
mx.eval(*outputs)
|
|
hashes = [hashlib.sha256(np.asarray(value if value.dtype == mx.uint32 else value.astype(mx.float32))
|
|
.astype('<u4' if value.dtype == mx.uint32 else '<f4').tobytes()).hexdigest() for value in outputs]
|
|
print(json.dumps(dict(kernel=name,sha256=hashes,**metadata),separators=(',',':')),flush=True)
|
|
for dtype in (mx.bfloat16,mx.float16,mx.float32):
|
|
for group in (32,64,128):
|
|
for bits in (2,3,4,5,6,8):
|
|
x = pattern(3*256,3).astype(dtype).reshape(3,256)
|
|
packed(f"quantize_{str(dtype).split('.')[-1]}_g{group}_b{bits}_plain",mx.quantize(x,group_size=group,bits=bits),
|
|
dtype=str(dtype).split('.')[-1],group=group,bits=bits,layout="plain")
|
|
for layout,x in (
|
|
("slice",pattern(3*512,3).reshape(3,512)[:,::2]),
|
|
("transpose",pattern(256*3,3).reshape(256,3).T),
|
|
("zero",mx.zeros((3,256),mx.bfloat16)),
|
|
("positive",mx.full((3,256),1,mx.bfloat16)),
|
|
("negative",mx.full((3,256),-1,mx.bfloat16)),
|
|
):
|
|
packed(f"quantize_bfloat16_g64_b8_{layout}",mx.quantize(x,group_size=64,bits=8),dtype="bfloat16",group=64,bits=8,layout=layout)
|
|
down,up,inject = pattern(320*10240,3).reshape(320,10240),pattern(10240*320,4).reshape(10240,320),pattern(4*10240,5).reshape(4,10240)
|
|
module=SimpleNamespace(input_mix_weight_down=SimpleNamespace(weight=down),input_mix_weight_up=SimpleNamespace(weight=up),block_inject_weight=SimpleNamespace(weight=inject))
|
|
pack=prepare_v3_pack(module)
|
|
packed("hyper_v3_original_pack",pack)
|
|
emit("hyper_v3_original_packed_read",fused_hyper_read_v3(pattern(10240,1),pattern(10240,2),pack))
|
|
|
|
|
|
def dequantize_fixtures():
|
|
"""Original affine dequantizer with all supported table formats and views."""
|
|
assert mx.__version__ == "0.32.2"
|
|
for dtype in (mx.bfloat16, mx.float16, mx.float32):
|
|
for group in (32,64,128):
|
|
for bits in (2,3,4,5,6,8):
|
|
for layout in ("plain","transpose","slice","offset","broadcast","batch"):
|
|
width=256*bits//32
|
|
def value(cols, packed=False, salt=3):
|
|
shape = (cols,3) if layout=="transpose" else (3,cols*2) if layout=="slice" else (4,cols) if layout=="offset" else (1,cols) if layout=="broadcast" else (2,3,cols) if layout=="batch" else (3,cols)
|
|
count=int(np.prod(shape))
|
|
a = mx.array((np.arange(count,dtype=np.uint64)*2654435761+12345).astype(np.uint32)) if packed else (pattern(count,salt)/64).astype(dtype)
|
|
a=a.reshape(shape)
|
|
if layout=="transpose": a=a.T
|
|
elif layout=="slice": a=a[:,::2]
|
|
elif layout=="offset": a=a[1:]
|
|
elif layout=="broadcast": a=mx.broadcast_to(a,(3,cols))
|
|
return a
|
|
w,s,b=value(width,True),value(256//group),value(256//group,salt=4)
|
|
emit(f"dequantize_{str(dtype).split('.')[-1]}_g{group}_b{bits}_{layout}",
|
|
[mx.dequantize(w,s,b,group_size=group,bits=bits)],
|
|
dtype=str(dtype).split('.')[-1],group=group,bits=bits,layout=layout)
|
|
|
|
|
|
def ngram_resident_fixtures():
|
|
"""Actual NGramTable._lazy_gather, including general Gather view dispatch."""
|
|
from mtplx.models.qwen4_exp import NGramTable
|
|
for dtype in (mx.bfloat16,mx.float16):
|
|
for bits,width in ((0,8),(0,9),(0,160),(4,160)):
|
|
for layout in ("plain","transpose","offset"):
|
|
for ids_layout in ("vector","scalar","matrix","head3","head4","slice","broadcast"):
|
|
if bits and ids_layout=="scalar": continue # Original dequantize requires rank>=2.
|
|
def part(cols,packed=False,salt=3):
|
|
shape=(cols,17) if layout=="transpose" else (18,cols) if layout=="offset" else (17,cols)
|
|
count=int(np.prod(shape))
|
|
a=mx.array((np.arange(count,dtype=np.uint64)*2654435761+12345).astype(np.uint32)) if packed else (pattern(count,salt)/64).astype(dtype)
|
|
a=a.reshape(shape)
|
|
return a.T if layout=="transpose" else a[1:] if layout=="offset" else a
|
|
table=NGramTable(17,width,sidecar=True)
|
|
table._lazy_parts=(part(width*bits//32,True),part(width//32),part(width//32,salt=4)) if bits else (part(width),None,None)
|
|
table._lazy_bits=bits;table._lazy_group=32
|
|
shapes={"vector":(7,),"scalar":(),"matrix":(2,7),"head3":(2,4,16),"head4":(2,2,4,16),"slice":(2,9,16),"broadcast":(1,1,16)}
|
|
shape=shapes[ids_layout];count=int(np.prod(shape))
|
|
ids=mx.array(((np.arange(count,dtype=np.int64)*7)%34)-17).reshape(shape)
|
|
if ids_layout=="slice": ids=ids[:,1::2,:]
|
|
if ids_layout=="broadcast": ids=mx.broadcast_to(ids,(2,4,16))
|
|
emit(f"ngram_resident_{str(dtype).split('.')[-1]}_b{bits}_w{width}_{layout}_{ids_layout}",[table._lazy_gather(ids)],
|
|
dtype=str(dtype).split('.')[-1],bits=bits,width=width,layout=layout,ids_layout=ids_layout)
|
|
|
|
|
|
def hyper_forward_fixtures():
|
|
"""Actual GatedResidual, including selection and persistent module pack."""
|
|
from mtplx.models.qwen4_exp import GatedResidual, TextArgs
|
|
from mtplx.kernels import hyper_connection as hc, hyper_connection_v3 as v3
|
|
calls=[]
|
|
original_read,original_v3=hc.fused_hyper_read,v3.fused_hyper_read_v3
|
|
def read(*args,**kwargs):
|
|
calls.append("fused")
|
|
return original_read(*args,**kwargs)
|
|
def read_v3(*args,**kwargs):
|
|
calls.append("v3")
|
|
return original_v3(*args,**kwargs)
|
|
hc.fused_hyper_read,v3.fused_hyper_read_v3=read,read_v3
|
|
try:
|
|
for layout in ("dense","quantized"):
|
|
for combine in (False,True):
|
|
layer=GatedResidual(TextArgs(),use_combine=combine)
|
|
layer.hc_norm.weight=pattern(10240,2)
|
|
names=[("input_mix_weight_down",320,10240,3),("input_mix_weight_up",10240,320,4)]
|
|
if combine: names.append(("block_inject_weight",4,10240,5))
|
|
for name,n,k,salt in names:
|
|
if layout=="dense": getattr(layer,name).weight=pattern(n*k,salt).reshape(n,k)
|
|
else: setattr(layer,name,quantized_linear_fixture(n,k,4,64,salt))
|
|
for rows in (1,2,4,8,15,16,32,2048):
|
|
for mode in range(4):
|
|
os.environ["MTPLX_FUSED_HC"]=str(mode & 1)
|
|
os.environ["MTPLX_FUSED_HC_V3"]=str(bool(mode & 2)).lower()
|
|
x=pattern(rows*10240,1).reshape(1,rows,10240)
|
|
calls.clear()
|
|
output=layer(x)
|
|
selected=calls[-1] if calls else "eager"
|
|
if combine:
|
|
mixed,hyper,inject=output
|
|
assert hyper is x
|
|
output=[mixed,inject]
|
|
else: output=[output]
|
|
emit(f"hyper_forward_{layout}_combine{int(combine)}_r{rows}_mode{mode}",output,
|
|
layout=layout,combine=combine,rows=rows,mode=mode,selected=selected,pack_exists=getattr(layer,"_v3_pack",None) is not None)
|
|
finally:
|
|
hc.fused_hyper_read,v3.fused_hyper_read_v3=original_read,original_v3
|
|
|
|
|
|
def ngram_rows_fixtures():
|
|
"""Actual NGramEmbedding CPU stage hash; raw I64, EOS and continuation."""
|
|
from mtplx.models.qwen4_exp import NGramEmbedding, TextArgs
|
|
layer=NGramEmbedding(TextArgs(ngram_sidecar=True,eos_token_id=248044),0)
|
|
mult,sizes,offsets=layer._np_consts()
|
|
print(json.dumps(dict(contract=True,multipliers=mult.tolist(),sizes=sizes.tolist(),
|
|
offsets=offsets.tolist(),eos=layer.eos_id)),flush=True)
|
|
def ids(n,salt):
|
|
palette=np.array([0,1,42,layer.eos_id,layer.eos_id,248000,-1,2**63-1,-2**63,43],dtype=np.int64)
|
|
return palette[(np.arange(n)*7+salt)%len(palette)]
|
|
for batch in (1,2):
|
|
for rows in (0,1,4,2048):
|
|
for prev_len in (0,1,2,5):
|
|
prev=ids(batch*prev_len,3).reshape(batch,prev_len)
|
|
for step,tokens in enumerate((rows,1)):
|
|
values=ids(batch*tokens,step).reshape(batch,tokens)
|
|
result,history=layer._rows_np(values,prev)
|
|
print(json.dumps(dict(kernel=f"ngram_rows_b{batch}_r{rows}_p{prev_len}_step{step}",
|
|
batch=batch,rows=rows,prev_len=prev_len,step=step,
|
|
result_shape=list(result.shape),history_shape=list(history.shape),
|
|
sha256=[hashlib.sha256(np.asarray(v).astype('<i8').tobytes()).hexdigest()
|
|
for v in (result,history)]),separators=(',',':')),flush=True)
|
|
prev=history
|
|
|
|
|
|
def ngram_gpu_fixtures():
|
|
"""Actual _graph_path through row IDs; table returns IDs, not embeddings."""
|
|
import mlx.nn as nn
|
|
from mlx_lm.models.cache import ArraysCache
|
|
from mtplx.models.qwen4_exp import NGramEmbedding, TextArgs
|
|
class RowsOnly(nn.Module):
|
|
def __call__(self,ids): return ids
|
|
layer=NGramEmbedding(TextArgs(ngram_sidecar=True,eos_token_id=248044),0)
|
|
layer.ngram_embedding=RowsOnly()
|
|
palette=np.array([0,1,42,layer.eos_id,layer.eos_id,248000,-1,2**63-1,-2**63,43],dtype=np.int64)
|
|
def ids(n,salt): return palette[(np.arange(n)*7+salt)%len(palette)]
|
|
for batch in (1,2):
|
|
for rows in (1,4,2048):
|
|
for prev_len in (-1,0,1,2,5):
|
|
cache=ArraysCache(4)
|
|
if prev_len>=0: cache[3]=mx.array(ids(batch*prev_len,3).reshape(batch,prev_len))
|
|
for step,tokens in enumerate((rows,1)):
|
|
values=mx.array(ids(batch*tokens,step).reshape(batch,tokens))
|
|
result=layer._graph_path(values,cache,3)
|
|
mx.eval(result,cache[3])
|
|
print(json.dumps(dict(kernel=f"ngram_gpu_b{batch}_r{rows}_p{prev_len}_step{step}",
|
|
batch=batch,rows=rows,prev_len=prev_len,step=step,
|
|
result_shape=list(result.shape),history_shape=list(cache[3].shape),
|
|
sha256=[hashlib.sha256(np.asarray(v).astype('<i8').tobytes()).hexdigest()
|
|
for v in (result,cache[3])]),separators=(',',':')),flush=True)
|
|
|
|
|
|
def ngram_sidecar_fixtures():
|
|
"""Actual file-backed _SidecarGather, including persistent LRU/bypass state."""
|
|
import tempfile
|
|
from mtplx.models.qwen4_exp import _SidecarGather
|
|
cases=([5,1,5,3],[3,5,2,1],[6,2,9,6],[],list(range(4095,-1,-1)),list(range(4096,-1,-1)),[-1,5002,-2],[1,3,5])
|
|
for dtype in ("BF16","F16"):
|
|
for bits in (0,4):
|
|
with tempfile.TemporaryDirectory(prefix="mtplx-sidecar-") as temp:
|
|
path=Path(temp)/"rows.bin"
|
|
entries={};offset=4096
|
|
names=("weight",) if bits==0 else ("weight","scales","biases")
|
|
with path.open("wb") as f:
|
|
for part,name in enumerate(names):
|
|
packed=bits!=0 and name=="weight"
|
|
columns=20 if packed else 5 if bits else 160
|
|
dt="U32" if packed else dtype
|
|
count=5003*columns
|
|
values=(np.arange(count,dtype=np.uint64)*2654435761+12345).astype('<u4') if packed else (0x3c00+(np.arange(count,dtype=np.uint64)*7+part*13)%512).astype('<u2')
|
|
f.seek(offset);f.write(values.tobytes())
|
|
entries[name]=({"dtype":dt,"shape":[5003,columns],"data_offsets":[offset,offset+values.nbytes]},0)
|
|
offset+=values.nbytes+34
|
|
for prefetch in (False,True):
|
|
os.environ["MTPLX_NGRAM_PREFETCH"]="1" if prefetch else "0"
|
|
for capacity in (0,3,4096):
|
|
gather=_SidecarGather(path,entries,bits,32);gather._hot_cap_rows=capacity
|
|
try:
|
|
for step,ids in enumerate(cases):
|
|
out=gather(mx.array(ids,mx.int64),160);mx.eval(out)
|
|
keys=np.array(list(gather._hot),np.int64)
|
|
print(json.dumps(dict(kernel=f"ngram_sidecar_{dtype}_b{bits}_p{int(prefetch)}_c{capacity}_s{step}",dtype=dtype,bits=bits,prefetch=prefetch,capacity=capacity,step=step,
|
|
sha256=hashlib.sha256(np.asarray(out.astype(mx.float32)).astype('<f4').tobytes()).hexdigest(),
|
|
order_sha256=hashlib.sha256(keys.astype('<i8').tobytes()).hexdigest(),entries=len(keys),hits=gather.hot_hits,misses=gather.hot_misses,prefetch_batches=gather.prefetch_batches),separators=(',',':')),flush=True)
|
|
finally:
|
|
if gather._pool is not None: gather._pool.shutdown()
|
|
os.close(gather._fd)
|
|
|
|
|
|
def ngram_stage_fixtures():
|
|
"""Original stage/__call__ lifecycle with real file-backed small tables."""
|
|
import tempfile,contextlib,io
|
|
from mlx_lm.models.cache import ArraysCache
|
|
from mtplx.models.qwen4_exp import NGramEmbedding,TextArgs,_SidecarGather
|
|
scenarios=("normal","verify","graph","disabled","stale","different","mismatch","skip_pending","error","no_cache","seeded","overwrite","graph_error","missing_skip","skip_policy")
|
|
for bits in (0,4):
|
|
with tempfile.TemporaryDirectory(prefix="mtplx-stage-") as temp:
|
|
path=Path(temp)/"rows.bin";entries={};offset=4096
|
|
with path.open("wb") as f:
|
|
for part,name in enumerate(("weight",) if bits==0 else ("weight","scales","biases")):
|
|
packed=bits and name=="weight";cols=20 if packed else 5 if bits else 160
|
|
values=(np.arange(17*cols,dtype=np.uint64)*2654435761+12345).astype('<u4') if packed else (0x3c00+(np.arange(17*cols,dtype=np.uint64)*7+part*13)%512).astype('<u2')
|
|
f.seek(offset);f.write(values.tobytes());entries[name]=({"dtype":"U32" if packed else "BF16","shape":[17,cols],"data_offsets":[offset,offset+values.nbytes]},0);offset+=values.nbytes+34
|
|
for scenario in scenarios:
|
|
for batch,rows in ((1,4),(2,4),*(([(1,2048),(2,2048)]) if scenario in ("normal","verify") else [])):
|
|
layer=NGramEmbedding(TextArgs(ngram_sidecar=True,eos_token_id=248044),0)
|
|
layer.ngram_heads_vocab_sizes=mx.full((16,),17,mx.int64);layer.ngram_heads_offsets=mx.zeros((16,),mx.int64)
|
|
os.environ["MTPLX_NGRAM_PREFETCH"]="0"
|
|
sidecar=_SidecarGather(path,entries,bits,32);layer.ngram_embedding._sidecar=sidecar
|
|
cache=None if scenario=="no_cache" else ArraysCache(4)
|
|
if scenario=="seeded": cache[3]=mx.arange(batch*5,dtype=mx.int64).reshape(batch,5)
|
|
os.environ["MTPLX_NGRAM_STAGE"]="0" if scenario in ("graph","graph_error") else "1"
|
|
os.environ["MTPLX_NGRAM_STAGE_VERIFY"]="1" if scenario in ("verify","mismatch") else "0"
|
|
layer._stage_disabled=scenario=="disabled"
|
|
try:
|
|
for step,tokens in enumerate((rows,1)):
|
|
ids=mx.array((np.arange(batch*tokens)*7+step)%51,mx.int32).reshape(batch,tokens)
|
|
alternate=ids+1
|
|
with contextlib.redirect_stdout(io.StringIO()):
|
|
layer.stage(ids,cache,3)
|
|
if step==0:
|
|
if scenario=="skip_pending":
|
|
layer._stage_disabled=True;layer.stage(alternate,cache,3);layer._stage_disabled=False
|
|
if scenario=="missing_skip":
|
|
layer.ngram_embedding._sidecar=None;layer.stage(alternate,cache,3);layer.ngram_embedding._sidecar=sidecar
|
|
if scenario=="skip_policy":
|
|
os.environ["MTPLX_NGRAM_STAGE"]="0";layer.stage(alternate,cache,3);os.environ["MTPLX_NGRAM_STAGE"]="1"
|
|
if scenario=="error": layer.stage(ids.reshape(-1),cache,3);layer.stage(ids.reshape(-1),cache,3)
|
|
if scenario=="overwrite": layer.stage(alternate,cache,3)
|
|
prehistory=None if cache is None or cache[3] is None else hashlib.sha256(np.asarray(cache[3]).astype('<i8').tobytes()).hexdigest()
|
|
called=ids[:,:1] if step==0 and scenario=="stale" else alternate if step==0 and scenario in ("different","mismatch","skip_pending","overwrite","missing_skip","skip_policy") else ids
|
|
if step==0 and scenario=="graph_error": layer.ngram_embedding._sidecar=None
|
|
error=False;out=None
|
|
try: out=layer(called,cache,3);mx.eval(out)
|
|
except RuntimeError as exc:
|
|
assert "staged/graph mismatch" in str(exc) or "sidecar was never attached" in str(exc);error=True
|
|
layer.ngram_embedding._sidecar=sidecar
|
|
history=None if cache is None or cache[3] is None else hashlib.sha256(np.asarray(cache[3]).astype('<i8').tobytes()).hexdigest()
|
|
print(json.dumps(dict(kernel=f"ngram_stage_b{bits}_{scenario}_B{batch}_r{rows}_s{step}",bits=bits,scenario=scenario,batch=batch,rows=rows,step=step,error=error,
|
|
sha256=None if out is None else hashlib.sha256(np.asarray(out.astype(mx.float32)).astype('<f4').tobytes()).hexdigest(),history=history,prehistory=prehistory,
|
|
consumed=getattr(layer,"_stage_consumed",0),bypassed=getattr(layer,"_stage_bypassed",0),graph_calls=getattr(layer,"_graph_calls",0),warned=getattr(layer,"_stage_warned",False),pending=getattr(layer,"_staged",None) is not None),separators=(',',':')),flush=True)
|
|
finally: os.close(sidecar._fd)
|
|
|
|
|
|
def ngram_stage_verification_fixtures():
|
|
"""Original QA allclose special values and all-reduction size boundaries."""
|
|
for size in (0,1,2,4096,4097,1<<20):
|
|
for mode in range(6):
|
|
a=np.ones(size,np.float32);b=a.copy()
|
|
if size:
|
|
if mode==1: b[-1]=2
|
|
if mode==2: a[-1]=b[-1]=np.inf
|
|
if mode==3: a[-1]=b[-1]=-np.inf
|
|
if mode==4: a[-1]=np.inf;b[-1]=-np.inf
|
|
if mode==5: a[-1]=b[-1]=np.nan
|
|
print(json.dumps(dict(kernel=f"stage_allclose_{size}_{mode}",value=bool(mx.allclose(mx.array(a),mx.array(b)))),separators=(',',':')),flush=True)
|
|
for size in (1<<26,(1<<26)+1):
|
|
print(json.dumps(dict(kernel=f"stage_all_{size}",value=bool(mx.all(mx.array(np.ones(size,np.bool_))))),separators=(',',':')),flush=True)
|
|
|
|
|
|
def ngram_embedding_fixtures():
|
|
"""Entire original resident _graph_path with a small table, actual head width."""
|
|
from mlx_lm.models.cache import ArraysCache
|
|
from mtplx.models.qwen4_exp import NGramEmbedding,TextArgs
|
|
layer=NGramEmbedding(TextArgs(ngram_sidecar=True,eos_token_id=248044),0)
|
|
layer.ngram_heads_vocab_sizes=mx.full((16,),17,mx.int64)
|
|
layer.ngram_heads_offsets=mx.zeros((16,),mx.int64)
|
|
table=layer.ngram_embedding
|
|
table.prefer_lazy=True;table._lazy_group=32
|
|
for bits in (0,4):
|
|
table._lazy_bits=bits
|
|
if bits:
|
|
words=(np.arange(17*20,dtype=np.uint64)*2654435761+12345).astype(np.uint32)
|
|
table._lazy_parts=(mx.array(words).reshape(17,20),(pattern(17*5,3)/64).reshape(17,5),(pattern(17*5,4)/64).reshape(17,5))
|
|
else: table._lazy_parts=((pattern(17*160,3)/64).reshape(17,160),None,None)
|
|
for batch in (1,2):
|
|
for rows in (1,4,2048):
|
|
cache=ArraysCache(4)
|
|
for step,tokens in enumerate((rows,1)):
|
|
values=mx.array(((np.arange(batch*tokens)*7+step)%51),mx.int64).reshape(batch,tokens)
|
|
result=layer._graph_path(values,cache,3)
|
|
mx.eval(result,cache[3])
|
|
print(json.dumps(dict(kernel=f"ngram_embedding_b{bits}_batch{batch}_r{rows}_step{step}",bits=bits,batch=batch,rows=rows,step=step,
|
|
sha256=[hashlib.sha256(np.asarray(result.astype(mx.float32)).astype('<f4').tobytes()).hexdigest(),hashlib.sha256(np.asarray(cache[3]).astype('<i8').tobytes()).hexdigest()]),separators=(',',':')),flush=True)
|
|
|
|
|
|
def ple_connected_fixtures():
|
|
"""Unmodified PLELayer including real lookup, stage and both cache slots."""
|
|
import tempfile, contextlib, io
|
|
from mlx_lm.models.cache import ArraysCache
|
|
from mtplx.models.qwen4_exp import PLELayer, TextArgs
|
|
assert mx.__version__ == "0.32.2"
|
|
for bits in (0,4):
|
|
layer=PLELayer(TextArgs(ngram_sidecar=True,eos_token_id=248044),0)
|
|
for name,n,salt in (("key_proj",10240,71),("value_proj",2560,73)):
|
|
if bits==0: getattr(layer,name).weight=pattern(n*2560,salt).reshape(n,2560)
|
|
else: setattr(layer,name,quantized_linear_fixture(n,2560,4,64,salt))
|
|
for name,salt in (("norm_key",75),("norm_query",76),("norm_conv",77)):
|
|
getattr(layer,name).weight=pattern(10240,salt)
|
|
layer.conv_weight=pattern(10240*4,78).reshape(10240,4,1)
|
|
embedding=layer.ple_embedding
|
|
embedding.ngram_heads_vocab_sizes=mx.full((16,),17,mx.int64)
|
|
embedding.ngram_heads_offsets=mx.zeros((16,),mx.int64)
|
|
with tempfile.TemporaryDirectory(prefix="mtplx-ple-") as temp:
|
|
path=Path(temp)/"rows.safetensors";header={"__metadata__":{"ngram_bits":str(bits),"ngram_group_size":"32"}};offset=4096
|
|
with path.open("wb") as f:
|
|
for part,name in enumerate(("weight",) if bits==0 else ("weight","scales","biases")):
|
|
packed=bits and name=="weight";cols=20 if packed else 5 if bits else 160
|
|
values=(np.arange(17*cols,dtype=np.uint64)*2654435761+12345).astype('<u4') if packed else (0x3c00+(np.arange(17*cols,dtype=np.uint64)*7+part*13)%512).astype('<u2')
|
|
f.seek(offset);f.write(values.tobytes());header[f"ngram.{name}"]={"dtype":"U32" if packed else "BF16","shape":[17,cols],"data_offsets":[offset-4096,offset-4096+values.nbytes]};offset+=values.nbytes+34
|
|
encoded=json.dumps(header).encode();assert len(encoded)<=4088
|
|
f.seek(0);f.write((4088).to_bytes(8,"little"));f.write(encoded.ljust(4088,b' '))
|
|
for mode in ("stage","verify","graph","resident"):
|
|
os.environ["MTPLX_NGRAM_STAGE"]="0" if mode=="graph" else "1"
|
|
os.environ["MTPLX_NGRAM_STAGE_VERIFY"]="1" if mode=="verify" else "0"
|
|
os.environ["MTPLX_NGRAM_PREFETCH"]="0"
|
|
table=embedding.ngram_embedding
|
|
table.prefer_lazy=mode=="resident"
|
|
embedding._stage_disabled=mode=="resident"
|
|
for batch,rows in ((1,7),(1,2048),(2,7)):
|
|
for state in ("none","empty","seeded"):
|
|
table.attach_sidecar(path);sidecar=table._sidecar
|
|
if mode=="resident":
|
|
log=io.StringIO()
|
|
with contextlib.redirect_stdout(log): bound=table.attach_resident(path)
|
|
assert bound, log.getvalue()
|
|
assert not table.prefer_lazy
|
|
table.prefer_lazy=True
|
|
cache=None if state=="none" else ArraysCache(4)
|
|
if state=="seeded":
|
|
cache[2]=pattern(batch*12*10240,79).reshape(batch,12,10240)
|
|
cache[3]=mx.arange(batch*5,dtype=mx.int64).reshape(batch,5)
|
|
try:
|
|
for step,tokens in enumerate((rows,1)):
|
|
values=((np.arange(batch*tokens)*7+step)%51).astype(np.int32)
|
|
values[min(2,len(values)-1)]=248044
|
|
ids=mx.array(values).reshape(batch,tokens)
|
|
hidden=pattern(batch*tokens*10240,80+step).reshape(batch,tokens,10240)
|
|
embedding.stage(ids,cache,3)
|
|
output=layer(hidden,ids,cache)
|
|
mx.eval(output,*([] if cache is None else [cache[2],cache[3]]))
|
|
history=None if cache is None else hashlib.sha256(np.asarray(cache[3]).astype('<i8').tobytes()).hexdigest()
|
|
emit(f"ple_connected_b{bits}_{mode}_B{batch}_r{rows}_{state}_s{step}",
|
|
[output] if cache is None else [output,cache[2]],
|
|
bits=bits,mode=mode,batch=batch,rows=rows,state=state,step=step,history=history)
|
|
finally: os.close(sidecar._fd)
|
|
|
|
|
|
def batched_sampler_fixtures():
|
|
"""Both actual batch routes, their whole-batch fallback and public methods."""
|
|
from mtplx import fast_sampling as fs
|
|
from mtplx.sampling import SamplerConfig
|
|
assert mx.__version__ == "0.32.2" and np.__version__ == "2.4.4"
|
|
def digest(values, dtype):
|
|
return hashlib.sha256(np.asarray(values, dtype=dtype).tobytes()).hexdigest()
|
|
def bits(value): return int(np.asarray(value, dtype='<f8').view('<u8'))
|
|
def receipt(batch):
|
|
if batch is None: return dict(none=True)
|
|
rng = np.random.default_rng(43210)
|
|
samples = []
|
|
for i in range(32):
|
|
try: samples.append(batch.sample(i % len(batch.token_ids), rng))
|
|
except ValueError: samples.append(None)
|
|
converted = []
|
|
for row in range(len(batch.token_ids)):
|
|
try:
|
|
p = batch.to_distribution(row)
|
|
converted.append(dict(ids=digest(p.token_ids, '<i8'), probs=digest(p.probs, '<f8')))
|
|
except ValueError: converted.append(dict(error=True))
|
|
lookups = []
|
|
for row in (0, -1, len(batch.token_ids), -len(batch.token_ids)-1):
|
|
for token in (-1, 0, 3, batch.vocab_size-1, batch.vocab_size):
|
|
try: lookups.append(bits(batch.probability(row, token)))
|
|
except IndexError: lookups.append(None)
|
|
return dict(ids=digest(batch.token_ids, '<i8'), probs=digest(batch.probs, '<f8'),
|
|
shape=list(batch.token_ids.shape), samples=samples, next_rng=bits(rng.random()),
|
|
converted=converted, lookups=lookups, vocab=batch.vocab_size)
|
|
original_eval, original_host = mx.eval, fs._host_sparse_distribution
|
|
trace = []
|
|
host_calls = []
|
|
def evaluate(*args, **kwargs):
|
|
trace.append([list(a.shape) for a in args])
|
|
return original_eval(*args, **kwargs)
|
|
def host(*args, **kwargs):
|
|
host_calls.append(1)
|
|
return original_host(*args, **kwargs)
|
|
mx.eval, fs._host_sparse_distribution = evaluate, host
|
|
try:
|
|
with np.errstate(all="ignore"):
|
|
for vocab in (17, 4097, 248320):
|
|
for dtype in (mx.bfloat16, mx.float16, mx.float32):
|
|
for layout in ("plain", "strided", "reverse", "broadcast"):
|
|
for mode in ("pattern", "equal", "mixed"):
|
|
for temperature, top_p, k in ((0.6, 0.95, 20), (0.7, 1.0, 20), (1e-7, 0.01, 129)):
|
|
for route in ("serial", "bound"):
|
|
cols = 2*vocab if layout == "strided" else vocab
|
|
count = cols if layout == "broadcast" else 4*cols
|
|
x = (((np.arange(count,dtype=np.int64)*17+96*13)%257)-128).astype(np.float32)/np.float32(131)
|
|
if mode == "equal": x[:] = 1
|
|
if mode == "mixed":
|
|
x[-cols:] = np.nan
|
|
row = mx.array(x).astype(dtype) + mx.array(0, dtype=dtype)
|
|
row = row.reshape(-1, cols)
|
|
if layout == "strided": row = row[:, 1::2]
|
|
elif layout == "reverse": row = row[:, ::-1]
|
|
elif layout == "broadcast": row = mx.broadcast_to(row, (4, vocab))
|
|
config = SamplerConfig(temperature=temperature, top_p=top_p, top_k=k)
|
|
trace.clear(); host_calls.clear()
|
|
try:
|
|
batch = (fs.bind_batched_top_k_distributions(config, vocab_size=vocab)(row)
|
|
if route == "bound" else fs.batched_sparse_distributions_from_mlx_logits(row, config))
|
|
result = receipt(batch)
|
|
except (ValueError, FloatingPointError): result = dict(error=True)
|
|
print(json.dumps(dict(vocab=vocab, dtype=str(dtype).split('.')[-1], layout=layout, mode=mode,
|
|
temperature=temperature, top_p=top_p, k=k, route=route, result=result,
|
|
evals=trace.copy(), host_calls=len(host_calls)), separators=(',',':')), flush=True)
|
|
# Constructor preserves negatives if the total remains positive;
|
|
# sample rejects the unnormalised positive subset without an RNG draw.
|
|
for name, ids, probs in (
|
|
("padded", [[3,0,-1,7],[0,3,3,-1]], [[1,3,0,2],[0,1,2,0]]),
|
|
("negative", [[0,1,2],[0,1,2]], [[2,-1,1],[1,0,0]]),
|
|
("zero", [[0,1]], [[0,0]]),
|
|
("nonfinite", [[0,1]], [[np.inf,1]]),
|
|
("mismatch", [[0,1]], [[1,2,3]])):
|
|
try: result=receipt(fs.BatchedSparseDistributions(np.array(ids), np.array(probs), vocab_size=17))
|
|
except ValueError: result=dict(error=True)
|
|
print(json.dumps(dict(constructor=name, result=result), separators=(',',':')), flush=True)
|
|
finally:
|
|
mx.eval, fs._host_sparse_distribution = original_eval, original_host
|
|
|
|
|
|
def mtp_prefill_driver_fixtures():
|
|
"""Unchanged committed-history prefill, including abort and history windows."""
|
|
from mtplx import generation as gen
|
|
from mtplx.runtime import MTPLXRuntime
|
|
from mlx_lm.models.cache import ArraysCache
|
|
from types import SimpleNamespace
|
|
import itertools
|
|
for name in list(os.environ):
|
|
if name.startswith("MTPLX_"): del os.environ[name]
|
|
original_tick, original_count = gen._owner_progress_tick, gen._runtime_count
|
|
cases = list(itertools.product((0, 1, 9), (None, 4), (None, 0, 3), (False, True),
|
|
(False, True), ("normal", "cache", "discard"), (0, 3, 5)))
|
|
cases += [(2050, 2048, window, absolute, True, "normal", 0)
|
|
for window in (None, 1, 3, 2049) for absolute in (False, True)]
|
|
try:
|
|
for length, chunk, window, absolute, final_only, route, abort_at in cases:
|
|
os.environ["MTPLX_SUSTAINED_PREFILL"] = "1" if chunk is not None else "0"
|
|
os.environ["MTPLX_PREFILL_CHUNK_SIZE"] = str(chunk or 2048)
|
|
os.environ["MTPLX_TARGET_EMIT_FULL_PREFILL_LOGITS"] = "0" if final_only else "1"
|
|
os.environ["MTPLX_PREFILL_OMLX_EXTERNAL"] = "0" if route == "normal" else "1"
|
|
os.environ["MTPLX_PREFILL_EXTERNAL_EMIT_LOGITS"] = "1" if route == "discard" else "0"
|
|
events, calls, history, chunks, checks = [], [], [], [], 0
|
|
class Model:
|
|
def make_cache(self):
|
|
events.append("init")
|
|
return [ArraysCache(1)]
|
|
def make_mtp_cache(self):
|
|
events.append("mtp_init")
|
|
return [ArraysCache(1)]
|
|
def __call__(self, ids, cache=None, return_hidden=False, emit_logits=True, logits_keep=None, **kwargs):
|
|
calls.append((ids, return_hidden, emit_logits, logits_keep or 0))
|
|
value = ids.astype(mx.float32)
|
|
cache[0][0] = value + 1
|
|
logits = value[..., None] if emit_logits else None
|
|
return (logits, value[..., None]) if return_hidden else logits
|
|
def mtp_update_cache(self, hidden, ids, mtp_cache=None, position_offset=None, **kwargs):
|
|
history.append((hidden, ids, position_offset))
|
|
mtp_cache[0][0] = hidden + ids[..., None].astype(mx.float32)
|
|
return mtp_cache[0][0]
|
|
def abort():
|
|
nonlocal checks
|
|
checks += 1
|
|
events.append("check")
|
|
return abort_at > 0 and checks == abort_at
|
|
def on_chunk(data):
|
|
events.append("callback")
|
|
assert data["chunk_elapsed_s"] >= 0 and data["elapsed_s"] >= 0
|
|
chunks.append(dict(tokens=data["tokens_done"], total=data["tokens_total"], size=data["chunk_size"]))
|
|
if absolute: raise ValueError("UI callback errors are deliberately ignored")
|
|
def record_count(runtime, key, amount=1):
|
|
original_count(runtime, key, amount)
|
|
if key == "prefill_chunks": events.append("chunk")
|
|
if key == "mtp_history_append_calls": events.append("history")
|
|
gen._owner_progress_tick = lambda: events.append("settled")
|
|
gen._runtime_count = record_count
|
|
rt = MTPLXRuntime(Model(), SimpleNamespace(), Path("."), True,
|
|
SimpleNamespace(concat_order="embedding_hidden", hidden_variant="post_norm"))
|
|
def data(a): return dict(shape=list(a.shape), values=np.array(a).reshape(-1).tolist())
|
|
try:
|
|
cache, logits, hidden, mtp_cache, target_s, history_s, base = gen._prefill_committed_mtp_history_streaming(
|
|
rt, [i % 37 for i in range(length)], history_window_tokens=window,
|
|
mtp_position_mode="absolute" if absolute else "cache", abort_check=abort, chunk_callback=on_chunk)
|
|
result = dict(logits=data(logits), hidden=data(hidden), base=base)
|
|
except (ValueError, gen.PostcommitAbort) as error:
|
|
result = dict(error=str(error))
|
|
print(json.dumps(dict(length=length, chunk=chunk, window=window, absolute=absolute,
|
|
final_only=gen._final_logits_prefill_enabled(), route=route, abort_at=abort_at, events=events, chunks=chunks,
|
|
calls=[dict(ids=np.array(ids).reshape(-1).tolist(), hidden=h, emit=e, keep=k) for ids,h,e,k in calls],
|
|
history=[dict(hidden=data(h), ids=np.array(ids).reshape(-1).tolist(), position=p) for h,ids,p in history], result=result)), flush=True)
|
|
finally:
|
|
gen._owner_progress_tick, gen._runtime_count = original_tick, original_count
|
|
|
|
|
|
def prefill_driver_fixtures():
|
|
"""Original cold _prefill control flow, cache roots and cancellation."""
|
|
from mtplx import generation as gen
|
|
from mtplx.runtime import MTPLXRuntime
|
|
from mlx_lm.models.cache import ArraysCache
|
|
from mtplx.models.qwen4_exp import QSACache
|
|
from types import SimpleNamespace
|
|
assert mx.__version__ == "0.32.2"
|
|
for name in list(os.environ):
|
|
if name.startswith("MTPLX_"): del os.environ[name]
|
|
arrays = [mx.full((1, 1, 256, 1), i, mx.float32) for i in range(18)]
|
|
gdn = ArraysCache(4)
|
|
gdn.left_padding, gdn.lengths = arrays[:2]
|
|
gdn.cache = arrays[2:6]
|
|
gdn._mtplx_verify_ple = (arrays[17], arrays[17])
|
|
gdn._mtplx_verify_rows = tuple(arrays[6:12])
|
|
qsa = QSACache()
|
|
qsa.kv.keys, qsa.kv.values = arrays[12], arrays[3]
|
|
qsa.kv.offset = 3
|
|
qsa.raw_keys, qsa.pooled, qsa.pooled_f32_t = arrays[14:17]
|
|
roots = gen._tree_mx_arrays([gdn, qsa, gdn])
|
|
print(json.dumps(dict(kind="roots", tags=[next(i for i, a in enumerate(arrays) if a is root) for root in roots],
|
|
shapes=[root.shape for root in roots])), flush=True)
|
|
original_tick = gen._owner_progress_tick
|
|
original_count = gen._runtime_count
|
|
try:
|
|
for length in (0, 1, 2, 9, 2050):
|
|
for chunk in ((None, 2048) if length > 9 else (None, 1, 4)):
|
|
for final_only in (False, True):
|
|
for route in ("normal", "cache", "discard"):
|
|
for hidden in (False, True):
|
|
for abort_at in (0, 1, 3):
|
|
os.environ["MTPLX_SUSTAINED_PREFILL"] = "1" if chunk is not None else "0"
|
|
os.environ["MTPLX_PREFILL_CHUNK_SIZE"] = str(chunk or 2048)
|
|
os.environ["MTPLX_TARGET_EMIT_FULL_PREFILL_LOGITS"] = "0" if final_only else "1"
|
|
os.environ["MTPLX_PREFILL_OMLX_EXTERNAL"] = "0" if route == "normal" else "1"
|
|
os.environ["MTPLX_PREFILL_EXTERNAL_EMIT_LOGITS"] = "1" if route == "discard" else "0"
|
|
events, calls, checks = [], [], 0
|
|
class Model:
|
|
def make_cache(self):
|
|
events.append("init")
|
|
return [ArraysCache(1)]
|
|
def __call__(self, ids, cache=None, return_hidden=False, emit_logits=True, logits_keep=None, **kwargs):
|
|
calls.append((ids, return_hidden, emit_logits, logits_keep or 0))
|
|
value = ids.astype(mx.float32)
|
|
cache[0][0] = value + 1
|
|
logits = value[..., None] if emit_logits else None
|
|
return (logits, value[..., None]) if return_hidden else logits
|
|
def abort():
|
|
nonlocal checks
|
|
checks += 1
|
|
events.append("check")
|
|
return abort_at > 0 and checks == abort_at
|
|
gen._owner_progress_tick = lambda: events.append("settled")
|
|
rt = MTPLXRuntime(Model(), SimpleNamespace(), Path("."), hidden, None)
|
|
def record_count(runtime, key, amount=1):
|
|
original_count(runtime, key, amount)
|
|
if key == "prefill_chunks": events.append("chunk")
|
|
gen._runtime_count = record_count
|
|
try:
|
|
cache, logits, h, seconds = gen._prefill(rt, [i % 37 for i in range(length)], return_hidden=hidden, abort_check=abort)
|
|
result = dict(logits=np.array(logits).reshape(-1).tolist(), shape=logits.shape,
|
|
hidden=None if h is None else dict(values=np.array(h).reshape(-1).tolist(), shape=h.shape))
|
|
except (ValueError, gen.PostcommitAbort) as error:
|
|
result = dict(error=str(error))
|
|
print(json.dumps(dict(kind="prefill", length=length, chunk=chunk,
|
|
final_only=gen._final_logits_prefill_enabled(), route=route, hidden=hidden, abort_at=abort_at,
|
|
events=events, calls=[dict(ids=np.array(ids).reshape(-1).tolist(), hidden=h, emit=e, keep=k) for ids,h,e,k in calls],
|
|
chunks=rt.diagnostic_counters.get("prefill_chunks", 0), result=result)), flush=True)
|
|
finally:
|
|
gen._owner_progress_tick = original_tick
|
|
gen._runtime_count = original_count
|
|
|
|
|
|
def ar_driver_fixtures():
|
|
"""Unchanged generate_ar with a tiny GPU transition-table model, no weights.
|
|
|
|
Isolates orchestration (not model or throughput parity); actual prefill,
|
|
samplers, pipeline and final-state capture are all MTPLX implementations.
|
|
"""
|
|
from mtplx.generation import generate_ar
|
|
from mtplx.runtime import MTPLXRuntime
|
|
from mtplx.sampling import SamplerConfig
|
|
from types import SimpleNamespace
|
|
assert mx.__version__ == "0.32.2"
|
|
for name in list(os.environ):
|
|
if name.startswith("MTPLX_"): del os.environ[name]
|
|
vocab = 37
|
|
values = (((np.arange(vocab*vocab, dtype=np.int64)*17 + 96*13) % 257)-128).astype(np.float32)/131
|
|
table = mx.array(values.reshape(vocab, vocab))
|
|
class Model:
|
|
def __init__(self):
|
|
self.calls, self.modes = [], []
|
|
def make_cache(self): return []
|
|
def set_ar_pipeline_mode(self, enabled):
|
|
self.modes.append(bool(enabled))
|
|
return True
|
|
def __call__(self, ids, cache=None, return_hidden=False, emit_logits=True, logits_keep=None):
|
|
self.calls.append((ids, return_hidden))
|
|
out = table[ids]
|
|
return (out, ids.astype(mx.float32)) if return_hidden else out
|
|
tokenizer = SimpleNamespace(decode=lambda ids, **kwargs: " ".join(map(str, ids)))
|
|
for pipeline in (False, True):
|
|
for asynchronous in (False, True):
|
|
for temperature, penalty in ((0., 0.), (.6, 0.), (.6, .4)):
|
|
for limit in (0, 1, 2, 3, 12):
|
|
for capture in (False, True):
|
|
# 13/35 are the first/second sampled pipeline tokens:
|
|
# exercise immediate EOS and EOS with an in-flight next
|
|
# graph, not just limits and a non-triggering stop set.
|
|
for stop_ids in (set(), {0, 5, 36}, {13}, {35}):
|
|
os.environ["MTPLX_AR_PIPELINE"] = "1" if pipeline else "0"
|
|
os.environ["MTPLX_ASYNC_AR"] = "1" if asynchronous else "0"
|
|
model = Model()
|
|
rt = MTPLXRuntime(model, tokenizer, Path("."), False, None)
|
|
emitted = []
|
|
out = generate_ar(rt, [1, 2], max_tokens=limit,
|
|
sampler=SamplerConfig(temperature=temperature, top_p=.95, top_k=20, presence_penalty=penalty),
|
|
seed=4294967297, stop_token_ids=stop_ids, token_callback=lambda ids: emitted.extend(ids), capture_final_state=capture)
|
|
# Decode calls only: real prefill consumed [1] then [2].
|
|
calls = [{"ids": np.array(ids).reshape(-1).tolist(), "hidden": hidden} for ids, hidden in model.calls[2:]]
|
|
final = None if out.final_state is None else dict(shape=out.final_state.final_logits.shape,
|
|
sha256=hashlib.sha256(np.array(out.final_state.final_logits).tobytes()).hexdigest())
|
|
print(json.dumps(dict(pipeline=pipeline, asynchronous=asynchronous, temperature=temperature,
|
|
penalty=penalty, limit=limit, capture=capture, stops=sorted(stop_ids), tokens=out.tokens,
|
|
emitted=emitted, calls=calls, modes=model.modes, final=final,
|
|
verify_calls=out.stats.verify_calls, finish=out.finish_reason)), flush=True)
|
|
|
|
|
|
def prefix_sampler_fixtures():
|
|
"""Actual target-prefix entry, interleaved calls and persistent default RNG."""
|
|
from mtplx.fast_sampling import sample_token_ids_from_mlx_logits
|
|
from mtplx.sampling import SamplerConfig
|
|
assert mx.__version__ == "0.32.2"
|
|
configs = ((0.6, .95, 20), (0., .95, 20), (.6, .95, 0),
|
|
(.6, .95, 33), (.6, 1., 0), (.6, .95, 32), (.6, 1., 33),
|
|
(1e-7, .01, 20), (.6, 1., 1), (.6, 1., 999999), (.6, .95, 20))
|
|
for vocab in (17, 4097, 248320):
|
|
for dtype in (mx.bfloat16, mx.float16, mx.float32):
|
|
for layout in ("vector", "singleton", "prefix", "strided", "reverse", "broadcast"):
|
|
shape = (vocab,) if layout == "vector" else ((1, vocab) if layout == "singleton" else (2, 2, vocab))
|
|
size = int(np.prod(shape)) * (2 if layout == "strided" else 1)
|
|
if layout == "broadcast": size = vocab
|
|
for mode in ("pattern", "equal", "nonfinite"):
|
|
x = (((np.arange(size, dtype=np.int64)*17 + 96*13) % 257)-128).astype(np.float32)/131
|
|
if mode == "equal": x[:] = 1
|
|
elif mode == "nonfinite":
|
|
x[np.arange(size) % 7 == 0] = np.nan
|
|
x[np.arange(size) % 7 == 1] = np.inf
|
|
x[np.arange(size) % 7 == 2] = -np.inf
|
|
row = mx.array(x).astype(dtype)
|
|
row = row + mx.array(0, dtype=dtype)
|
|
if layout == "strided": row = row[1::2]
|
|
elif layout == "reverse": row = row[::-1]
|
|
if layout == "broadcast": row = mx.broadcast_to(row, shape)
|
|
else: row = row.reshape(shape)
|
|
mx.random.seed(4294967297)
|
|
outputs = []
|
|
for step, (temperature, top_p, k) in enumerate(configs):
|
|
if step == 8: mx.random.seed(17)
|
|
try:
|
|
token = sample_token_ids_from_mlx_logits(row,
|
|
SamplerConfig(temperature=temperature, top_p=top_p, top_k=k))
|
|
if token is None: result = {"none": True}
|
|
else:
|
|
mx.eval(token)
|
|
result = {"tokens": np.array(token).reshape(-1).tolist(), "shape": token.shape}
|
|
except (ValueError, RuntimeError) as error:
|
|
result = {"error": str(error)}
|
|
state = mx.random.state[0]
|
|
mx.eval(state)
|
|
result["key"] = np.array(state).tolist()
|
|
outputs.append(result)
|
|
print(json.dumps(dict(vocab=vocab, dtype=str(dtype).split('.')[-1], layout=layout,
|
|
mode=mode, outputs=outputs)), flush=True)
|
|
|
|
|
|
def resident_sampler_fixtures():
|
|
"""Original keyed GPU RNG and _mx_lazy_sample, with chained resident keys."""
|
|
from mtplx.generation import _mx_lazy_sample
|
|
from mtplx.sampling import SamplerConfig
|
|
assert mx.__version__ == "0.32.2"
|
|
def digest(a):
|
|
return hashlib.sha256(np.array(a).tobytes()).hexdigest()
|
|
for seed in (0, 4294967297, 18446744073709551615):
|
|
for layout in ("plain", "strided", "reverse"):
|
|
raw = mx.random.key(seed)
|
|
if layout == "strided":
|
|
raw = mx.stack((raw, raw), axis=1).reshape(-1)[::2]
|
|
elif layout == "reverse":
|
|
raw = raw[::-1]
|
|
for shape in ((0,), (1,), (2,), (3,), (127,), (4097,), (2, 3)):
|
|
bits = mx.random.split(raw, num=int(np.prod(shape)))
|
|
uniform = mx.random.uniform(shape=shape, key=raw)
|
|
key = raw
|
|
keys = []
|
|
for _ in range(4):
|
|
key, sub = mx.random.split(key)
|
|
keys.append(sub)
|
|
mx.eval(bits, uniform, key, *keys)
|
|
print(json.dumps(dict(kind="rng", seed=seed, layout=layout,
|
|
shape=shape, bits=digest(bits), uniform=digest(uniform),
|
|
keys=[np.array(k).tolist() for k in keys], final_key=np.array(key).tolist())), flush=True)
|
|
for vocab in (17, 4097, 248320):
|
|
for dtype in (mx.bfloat16, mx.float16, mx.float32):
|
|
for layout in ("plain", "strided", "reverse", "broadcast"):
|
|
for mode in ("pattern", "equal", "nonfinite"):
|
|
size = 2*vocab if layout == "strided" else (1 if layout == "broadcast" else vocab)
|
|
x = (((np.arange(size, dtype=np.int64)*17 + 96*13) % 257)-128).astype(np.float32)/131
|
|
if mode == "equal":
|
|
x[:] = 1
|
|
elif mode == "nonfinite":
|
|
for i in range(size):
|
|
if i % 7 < 3:
|
|
x[i] = (np.nan, np.inf, -np.inf)[i % 7]
|
|
row = mx.array(x).astype(dtype)
|
|
row = row + mx.array(0, dtype=dtype)
|
|
if layout == "strided": row = row[1::2]
|
|
elif layout == "reverse": row = row[::-1]
|
|
elif layout == "broadcast": row = mx.broadcast_to(row, (vocab,))
|
|
for temperature, top_p, k in ((0.6, 0.95, min(20, vocab-1)), (1.0, 1.0, 2), (1e-7, 0.01, min(129, vocab-1))):
|
|
config = SamplerConfig(temperature=temperature, top_p=top_p, top_k=k)
|
|
key = mx.random.key(17)
|
|
tokens = []
|
|
for _ in range(4):
|
|
key, sub = mx.random.split(key)
|
|
tokens.append(_mx_lazy_sample(row, config, sub))
|
|
mx.eval(key, *tokens)
|
|
print(json.dumps(dict(kind="sampler", vocab=vocab, dtype=str(dtype).split('.')[-1],
|
|
layout=layout, mode=mode, temperature=temperature, top_p=top_p, k=k,
|
|
tokens=[int(t.item()) for t in tokens], final_key=np.array(key).tolist())), flush=True)
|
|
|
|
|
|
def sampler_entry_fixtures():
|
|
"""Actual MTPLX generation entry, three consecutive draws and count updates."""
|
|
from mtplx import generation
|
|
from mtplx.sampling import SamplerConfig, SparseDistribution
|
|
assert mx.__version__ == "0.32.2"
|
|
original_eval, original_argmax = mx.eval, mx.argmax
|
|
trace = []
|
|
def evaluate(*args,**kwargs):
|
|
trace.append("eval")
|
|
return original_eval(*args,**kwargs)
|
|
def argmax(*args,**kwargs):
|
|
trace.append("argmax")
|
|
return original_argmax(*args,**kwargs)
|
|
mx.eval, mx.argmax = evaluate, argmax
|
|
try:
|
|
for vocab in (17,4097,248320):
|
|
for dtype in (mx.bfloat16,mx.float16,mx.float32):
|
|
for layout in ("plain","rank3","strided","reverse","broadcast"):
|
|
for mode in ("pattern","equal","nonfinite"):
|
|
size = 2*vocab if layout == "strided" else vocab
|
|
values = (((np.arange(size,dtype=np.int64)*17+96*13)%257)-128).astype(np.float32)/np.float32(131)
|
|
if mode == "equal": values[:] = 1
|
|
if mode == "nonfinite":
|
|
values[::7]=np.nan; values[1::7]=np.inf; values[2::7]=-np.inf
|
|
for profile in ("greedy","sparse","dense_top_p","dense_all"):
|
|
for penalties in (False,True):
|
|
config = SamplerConfig(temperature=0. if profile=="greedy" else 0.6,
|
|
top_k=20 if profile in ("greedy","sparse") else 0,
|
|
top_p=1. if profile=="dense_all" else 0.95,
|
|
presence_penalty=0.17 if penalties else 0., frequency_penalty=0.13 if penalties else 0.)
|
|
counts={vocab//2:2}; overlay={1:0.10009} if penalties else {}
|
|
rng=np.random.default_rng(12345); steps=[]
|
|
for step in range(3):
|
|
x=mx.array(values,dtype=dtype)+mx.array(0.,dtype=dtype)
|
|
if layout=="rank3": x=x.reshape(1,1,vocab)
|
|
if layout=="strided": x=x[1::2]
|
|
if layout=="reverse": x=x[::-1]
|
|
if layout=="broadcast": x=mx.broadcast_to(x[:1],(vocab,))
|
|
trace.clear()
|
|
try:
|
|
token,p=generation._sample_from_logits(x,config,rng,token_counts=counts,penalty_overlay=overlay)
|
|
result=dict(token=token,kind="none" if p is None else "sparse" if isinstance(p,SparseDistribution) else "dense")
|
|
if p is not None:
|
|
probs=p.probs if isinstance(p,SparseDistribution) else p
|
|
result["probabilities"]=hashlib.sha256(probs.astype('<f8').tobytes()).hexdigest()
|
|
if isinstance(p,SparseDistribution): result["ids"]=hashlib.sha256(p.token_ids.astype('<i8').tobytes()).hexdigest()
|
|
counts[token]=counts.get(token,0)+1
|
|
except (ValueError,IndexError): result=dict(error=True)
|
|
result["trace"]=list(trace); steps.append(result)
|
|
print(json.dumps(dict(vocab=vocab,dtype=str(dtype).split('.')[-1],layout=layout,mode=mode,
|
|
profile=profile,penalties=penalties,steps=steps,counts=list(counts.items()),
|
|
next_rng=int(np.asarray(rng.random(),dtype='<f8').view('<u8'))),separators=(',',':')),flush=True)
|
|
finally:
|
|
mx.eval, mx.argmax = original_eval, original_argmax
|
|
|
|
|
|
def sampler_penalty_fixtures():
|
|
from mtplx.fast_sampling import apply_penalties_mlx
|
|
from mtplx.sampling import apply_penalties
|
|
assert mx.__version__ == "0.32.2"
|
|
for vocab in (17,4097,248320):
|
|
for dtype in (mx.bfloat16,mx.float16,mx.float32):
|
|
for layout in ("plain","strided","offset","reverse","broadcast"):
|
|
size = 2*vocab if layout == "strided" else vocab+1 if layout == "offset" else vocab
|
|
values = (((np.arange(size,dtype=np.int64)*17+96*13)%257)-128).astype(np.float32)/np.float32(131)
|
|
x = mx.array(values,dtype=dtype)
|
|
if layout == "strided": x = x[1::2]
|
|
if layout == "offset": x = x[1:]
|
|
if layout == "reverse": x = x[::-1]
|
|
if layout == "broadcast": x = mx.broadcast_to(x[:1],(vocab,))
|
|
mx.eval(x)
|
|
raw = np.asarray(x.astype(mx.float32)).astype(np.float64)
|
|
counts = {0:0,vocab//2:3,vocab-1:1}
|
|
overlay = {1:0.00031,vocab//2:0.10009,-1:3.14159}
|
|
profiles = [("noop",counts,0.,0.,{}),("no_counts",{},1.3,0.4,{}),
|
|
("counts",counts,0.17,0.13,{}),("overlay",{},0.,0.,overlay),
|
|
("both",counts,-0.7,1.3,overlay),("clamped",counts,4.5,-3.7,overlay),
|
|
("zero_overlay",{},0.,0.,{vocab//2:0.}),
|
|
("single",{vocab//2:2},0.4,0.13,{})]
|
|
for profile,c,p,f,o in profiles:
|
|
gpu = apply_penalties_mlx(x,c,p,f,o)
|
|
cpu = apply_penalties(raw,c,p,f,o)
|
|
mx.eval(gpu)
|
|
result = np.asarray(gpu.astype(mx.float32)).astype('<f4')
|
|
print(json.dumps(dict(vocab=vocab,dtype=str(dtype).split('.')[-1],layout=layout,profile=profile,
|
|
counts=list(c.items()),presence=p,frequency=f,overlay=list(o.items()),
|
|
gpu_identity=gpu is x,cpu_identity=cpu is raw,
|
|
gpu=hashlib.sha256(result.tobytes()).hexdigest(),
|
|
cpu=hashlib.sha256(cpu.astype('<f8').tobytes()).hexdigest()),separators=(',',':')),flush=True)
|
|
|
|
|
|
def sampler_mixed_distribution_fixtures():
|
|
"""Original dense/sparse choice, residual, and verify including RNG errors."""
|
|
from mtplx.sampling import SparseDistribution, sample_from_distribution, residual_distribution, verify_one_token
|
|
assert np.__version__ == "2.4.4"
|
|
def bits(value): return int(np.asarray(value,dtype='<f8').view('<u8'))
|
|
def digest(values,dtype): return hashlib.sha256(np.asarray(values,dtype=dtype).tobytes()).hexdigest()
|
|
with np.errstate(all="ignore"):
|
|
for vocab in (1,17,129,4097,248320):
|
|
for mode in ("reverse","same","disjoint","nonfinite","zero"):
|
|
p = ((np.arange(vocab,dtype=np.int64)*17+3)%31).astype(np.float64)/31
|
|
q = p[::-1].copy()
|
|
if mode == "same": q = p.copy()
|
|
if mode == "disjoint":
|
|
p[1::2] = 0
|
|
q[::2] = 0
|
|
if mode == "nonfinite":
|
|
p[0] = np.nan
|
|
q[-1] = np.inf
|
|
if mode == "zero": p[:] = 0
|
|
for kinds in ("dd","ds","sd","ss"):
|
|
def make(raw,kind):
|
|
count = min(20,vocab)
|
|
ids = np.arange(count,dtype=np.int64)*(vocab-1)//max(1,count-1)
|
|
return SparseDistribution(ids,raw[ids],vocab) if kind == "s" else raw/raw.sum()
|
|
target,draft = make(p,kinds[0]),make(q,kinds[1])
|
|
rng = np.random.default_rng(12345)
|
|
try:
|
|
sampled = dict(tokens=[sample_from_distribution(target,rng) for _ in range(32)])
|
|
except ValueError: sampled = dict(error=True)
|
|
sampled["next"] = bits(rng.random())
|
|
try:
|
|
r = residual_distribution(target,draft)
|
|
sparse = isinstance(r,SparseDistribution)
|
|
residual = dict(kind="s" if sparse else "d",borrowed=r is target,
|
|
probabilities=digest(r.probs if sparse else r,'<f8'))
|
|
if sparse: residual["ids"] = digest(r.token_ids,'<i8')
|
|
except ValueError: residual = dict(error=True)
|
|
rng = np.random.default_rng(43210)
|
|
decisions=[]
|
|
for i in range(32):
|
|
token=(i*7+5)%vocab
|
|
try:
|
|
decision=verify_one_token(target,draft,token,rng)
|
|
decisions.append([decision.accepted,decision.token_id,bits(decision.accept_probability)])
|
|
except ValueError: decisions.append(None)
|
|
print(json.dumps(dict(vocab=vocab,mode=mode,kinds=kinds,sampled=sampled,residual=residual,
|
|
decisions=decisions,next_verify=bits(rng.random())),separators=(',',':')),flush=True)
|
|
|
|
|
|
def sampler_argmax_fixtures():
|
|
"""The pinned runtime ArgReduce used by MTPLX's greedy sampler."""
|
|
assert mx.__version__ == "0.32.2"
|
|
for dtype in (mx.bfloat16, mx.float16, mx.float32):
|
|
for cols in (1, 17, 4096, 4097, 248320):
|
|
for layout in ("plain", "swapped", "strided", "offset", "broadcast", "vector", "middle", "reverse"):
|
|
for mode in ("pattern", "nonfinite"):
|
|
def values(shape):
|
|
data = (((np.arange(np.prod(shape), dtype=np.int64) * 17 + 96 * 13) % 257) - 128).astype(np.float32) / np.float32(131)
|
|
if mode == "nonfinite":
|
|
data[::7] = np.nan
|
|
data[1::7] = np.inf
|
|
data[2::7] = -np.inf
|
|
return mx.array(data.reshape(shape), dtype=dtype)
|
|
if layout == "offset":
|
|
x = values((6*cols+1,))[1:].reshape(2,3,cols)
|
|
elif layout == "strided":
|
|
x = values((2,3,2*cols))[:,:,1::2]
|
|
elif layout == "broadcast":
|
|
x = mx.broadcast_to(values((cols,)),(2,3,cols))
|
|
elif layout == "vector":
|
|
x = values((cols,))
|
|
elif layout == "swapped":
|
|
x = mx.transpose(values((2,3,cols)),(1,0,2))
|
|
elif layout == "middle":
|
|
x = mx.transpose(values((2,3,cols)),(0,2,1))
|
|
elif layout == "reverse":
|
|
x = values((2,3,cols))[:,:,::-1]
|
|
else:
|
|
x = values((2,3,cols))
|
|
axis = 1 if layout == "middle" else -1
|
|
for keepdims in (False,True):
|
|
out = mx.argmax(x,axis=axis,keepdims=keepdims)
|
|
mx.eval(out)
|
|
print(json.dumps(dict(dtype=str(dtype).split('.')[-1],cols=cols,layout=layout,
|
|
mode=mode,axis=axis,keepdims=keepdims,shape=list(out.shape),
|
|
sha256=hashlib.sha256(np.asarray(out).astype('<u4').tobytes()).hexdigest()),separators=(',',':')),flush=True)
|
|
|
|
|
|
def sampler_fallback_fixtures():
|
|
from mtplx.fast_sampling import sparse_distributions_from_mlx_logits
|
|
from mtplx.sampling import SamplerConfig, distribution_from_logits
|
|
assert mx.__version__ == "0.32.2"
|
|
for vocab in (17, 4097, 248320):
|
|
for mode in ("unique", "sharp", "equal", "nonfinite", "mixed"):
|
|
values = sampler_serial_values(vocab, 4, mode)
|
|
for top_p in (0.95, 1.0):
|
|
config = SamplerConfig(temperature=0.6, top_p=top_p, top_k=20)
|
|
distributions = sparse_distributions_from_mlx_logits(mx.array(values).reshape(4,vocab),config)
|
|
print(json.dumps(dict(kind="sparse",vocab=vocab,mode=mode,top_p=top_p,
|
|
ids=[d.token_ids.tolist() for d in distributions],
|
|
probabilities=[d.probs.astype('<f8').view('<u8').tolist() for d in distributions]),separators=(',',':')),flush=True)
|
|
for top_k in (0,20):
|
|
for temperature in (0.0,0.6):
|
|
try:
|
|
probs = distribution_from_logits(values[:vocab].astype(np.float64),
|
|
SamplerConfig(temperature=temperature,top_p=top_p,top_k=top_k))
|
|
result = dict(sha256=hashlib.sha256(probs.astype('<f8').tobytes()).hexdigest())
|
|
except ValueError:
|
|
result = dict(error=True)
|
|
print(json.dumps(dict(kind="dense",vocab=vocab,mode=mode,top_p=top_p,
|
|
top_k=top_k,temperature=temperature,**result),separators=(',',':')),flush=True)
|
|
|
|
|
|
def sampler_distribution_fixtures():
|
|
"""Original normalization/choice/verify over the captured serial support."""
|
|
from mtplx.fast_sampling import _serial_row_distribution
|
|
from mtplx.sampling import SparseDistribution, sample_from_distribution, verify_one_token, residual_distribution
|
|
assert np.__version__ == "2.4.4"
|
|
def bits(values): return np.asarray(values, dtype='<f8').view('<u8').tolist()
|
|
for line in Path('tests/fixtures/mtplx-sampler-serial.jsonl').read_text().splitlines():
|
|
case = json.loads(line)
|
|
result = []
|
|
for ids, probability_bits in zip(case['ids'], case['probability_bits']):
|
|
distribution = _serial_row_distribution(np.asarray(ids, dtype=np.int64),
|
|
np.asarray(probability_bits, dtype='<u8').view('<f8'), case['vocabulary'])
|
|
if distribution is None:
|
|
result.append(None)
|
|
continue
|
|
rng = np.random.default_rng(12345)
|
|
samples = [sample_from_distribution(distribution, rng) for _ in range(32)]
|
|
next_sample = bits(rng.random())
|
|
draft = SparseDistribution(distribution.token_ids, distribution.probs[::-1], distribution.vocab_size)
|
|
rng = np.random.default_rng(43210)
|
|
decisions = []
|
|
for _ in range(32):
|
|
token = sample_from_distribution(draft, rng)
|
|
decision = verify_one_token(distribution, draft, token, rng)
|
|
decisions.append([token, decision.accepted, decision.token_id, bits(decision.accept_probability)])
|
|
residual = residual_distribution(distribution, draft)
|
|
result.append(dict(ids=distribution.token_ids.tolist(), probabilities=bits(distribution.probs),
|
|
samples=samples, next_sample=next_sample, decisions=decisions, next_verify=bits(rng.random()),
|
|
residual_ids=residual.token_ids.tolist(), residual_probabilities=bits(residual.probs)))
|
|
print(json.dumps(result,separators=(',',':')),flush=True)
|
|
|
|
|
|
def sampler_serial_values(vocab, batch, mode):
|
|
i = np.arange(batch*vocab, dtype=np.int64)
|
|
unique = (i%vocab).astype(np.float32)/np.float32(vocab)
|
|
sharp = -np.abs(i%vocab-vocab//2).astype(np.float32)*np.float32(0.2)
|
|
if mode == "ties": return (i%7).astype(np.float32)
|
|
if mode == "equal": return np.ones(batch*vocab, dtype=np.float32)
|
|
if mode == "sharp": return sharp
|
|
if mode == "nonfinite":
|
|
unique[:3] = [np.nan, np.inf, -np.inf]
|
|
if mode == "mixed":
|
|
unique[:vocab] = 1
|
|
if batch > 2: unique[2*vocab:3*vocab] = (i[2*vocab:3*vocab]%7).astype(np.float32)
|
|
if batch > 3: unique[3*vocab:] = sharp[3*vocab:]
|
|
return unique
|
|
|
|
|
|
def sampler_serial_fixtures():
|
|
"""Actual serial runtime helper with eval/selector calls recorded."""
|
|
from mtplx import fast_sampling as fs
|
|
from mtplx.sampling import SamplerConfig
|
|
assert mx.__version__ == "0.32.2"
|
|
assert np.__version__ == "2.4.4"
|
|
original_eval = mx.eval
|
|
original_selector = fs._deterministic_mlx_top_k_support
|
|
counts = [0, 0]
|
|
def evaluate(*args):
|
|
counts[0] += 1
|
|
return original_eval(*args)
|
|
def selector(*args, **kwargs):
|
|
counts[1] += 1
|
|
return original_selector(*args, **kwargs)
|
|
mx.eval = evaluate
|
|
fs._deterministic_mlx_top_k_support = selector
|
|
try:
|
|
for vocab in (17, 4097, 248320):
|
|
for batch in (1, 4):
|
|
for mode in ("unique", "ties", "equal", "sharp", "nonfinite", "mixed"):
|
|
for top_p in (0.1, 0.95, 1.0):
|
|
temperature = 0.6
|
|
counts[:] = [0, 0]
|
|
logits = mx.array(sampler_serial_values(vocab, batch, mode)).reshape(batch, vocab)
|
|
ids, probabilities, size = fs._device_serial_support_arrays(logits, SamplerConfig(temperature=temperature, top_p=top_p))
|
|
probabilities = probabilities.astype('<f8')
|
|
probabilities[np.isnan(probabilities)] = np.nan
|
|
print(json.dumps(dict(vocab=vocab, batch=batch, mode=mode, top_p=top_p, temperature=temperature,
|
|
ids=ids.tolist(), probability_bits=probabilities.view('<u8').tolist(),
|
|
vocabulary=size, eval_calls=counts[0], selector_calls=counts[1]),separators=(',',':')),flush=True)
|
|
finally:
|
|
mx.eval = original_eval
|
|
fs._deterministic_mlx_top_k_support = original_selector
|
|
|
|
|
|
def sampler_selection_fixtures():
|
|
"""The unchanged MTPLX deterministic selector, not a NumPy substitute."""
|
|
from mtplx.fast_sampling import _deterministic_mlx_top_k_support
|
|
assert mx.__version__ == "0.32.2"
|
|
for vocab in (17, 4097, 248320):
|
|
for batch in (1, 4):
|
|
for mode in ("pattern", "ties", "equal", "nonfinite"):
|
|
for top_k in (1, min(20, vocab)):
|
|
i = np.arange(batch*vocab, dtype=np.int64)
|
|
if mode == "pattern": values = ((i*17+96*13)%257-128).astype(np.float32)/128
|
|
elif mode == "ties": values = (i%7).astype(np.float32)
|
|
elif mode == "equal": values = np.ones(batch*vocab, dtype=np.float32)
|
|
else:
|
|
values = ((i*17+96*13)%257-128).astype(np.float32)/128
|
|
values[0] = np.nan
|
|
values[1] = np.inf
|
|
values[2] = -np.inf
|
|
rows = mx.array(values).reshape(batch, vocab)
|
|
ids, vals = _deterministic_mlx_top_k_support(rows, top_k)
|
|
emit(f"sampler_selection_v{vocab}_b{batch}_{mode}_k{top_k}",
|
|
[ids, vals], vocab=vocab, batch=batch, mode=mode, top_k=top_k)
|
|
|
|
|
|
def logsumexp_graph_fixtures():
|
|
"""Original sampler normalizer, including full Qwen vocabulary and layouts."""
|
|
assert mx.__version__ == "0.32.2"
|
|
for dtype in (mx.bfloat16, mx.float16, mx.float32):
|
|
for cols in (1, 17, 4096, 4097, 248320):
|
|
for layout in ("plain", "swapped", "strided", "offset", "broadcast"):
|
|
if layout == "offset":
|
|
x = pattern(6*cols+1, 96).astype(dtype)[1:].reshape(2, 3, cols)
|
|
elif layout == "strided":
|
|
x = pattern(12*cols, 96).astype(dtype).reshape(2, 3, cols*2)[..., 1::2]
|
|
elif layout == "broadcast":
|
|
x = mx.broadcast_to(pattern(cols, 96).astype(dtype), (2, 3, cols))
|
|
else:
|
|
x = pattern(6*cols, 96).astype(dtype).reshape(2, 3, cols)
|
|
if layout == "swapped":
|
|
x = x.transpose(1, 0, 2)
|
|
for keepdims in (False, True):
|
|
name = str(dtype).split('.')[-1]
|
|
y = mx.logsumexp(x, axis=-1, keepdims=keepdims)
|
|
emit(f"logsumexp_{name}_{cols}_{layout}_{int(keepdims)}", [y],
|
|
dtype=name, cols=cols, layout=layout, keepdims=keepdims, shape=y.shape)
|
|
|
|
|
|
def softmax_graph_fixtures():
|
|
"""Original last-axis Softmax's precise/type/layout branches."""
|
|
for dtype in (mx.bfloat16,mx.float16,mx.float32):
|
|
for precise in (False,True):
|
|
for cols in (17,4097):
|
|
for layout in ("plain","swapped","strided","offset"):
|
|
if layout=="offset": x=pattern(6*cols+1,96).astype(dtype)[1:].reshape(2,3,cols)
|
|
elif layout=="strided": x=pattern(12*cols,96).astype(dtype).reshape(2,3,cols*2)[...,1::2]
|
|
else:
|
|
x=pattern(6*cols,96).astype(dtype).reshape(2,3,cols)
|
|
if layout=="swapped": x=x.transpose(1,0,2)
|
|
name=str(dtype).split('.')[-1]
|
|
emit(f"softmax_graph_{name}_{int(precise)}_{cols}_{layout}",[mx.softmax(x,axis=-1,precise=precise)],dtype=name,precise=precise,cols=cols,layout=layout)
|
|
|
|
|
|
def sdpa_fallback_fixtures():
|
|
"""Actual fast SDPA; these Qwen shapes select its original unfused graph."""
|
|
assert mx.__version__ == "0.32.2"
|
|
for batch,heads,rows,total in ((1,24,3,17),(1,24,4,1025),(1,24,8,4097),
|
|
(1,24,16,33),(1,24,2048,2048),(2,24,3,17),(1,2,16,33)):
|
|
q=pattern(batch*rows*heads*256,93).reshape(batch,rows,heads,256).transpose(0,2,1,3)
|
|
k=pattern(batch*2*(total+7)*256,94).reshape(batch,2,total+7,256)[:,:,2:total+2,:]
|
|
v=pattern(batch*2*(total+7)*256,95).reshape(batch,2,total+7,256)[:,:,2:total+2,:]
|
|
for mode in ("none","bool","add"):
|
|
visible=(np.arange(total)[None,:] <= (total-rows+np.arange(rows))[:,None]) & (np.arange(total)[None,:]%3!=1)
|
|
visible[1,:]=False
|
|
mask=None if mode=="none" else mx.array(visible).reshape(1,1,rows,total) if mode=="bool" else mx.array(np.where(visible,0.,-10.),mx.bfloat16).reshape(1,1,rows,total)
|
|
out=mx.fast.scaled_dot_product_attention(q,k,v,scale=0.0625,mask=mask)
|
|
emit(f"sdpa_fallback_B{batch}_h{heads}_r{rows}_t{total}_{mode}",[out],batch=batch,heads=heads,rows=rows,total=total,mode=mode)
|
|
|
|
|
|
def dense_batch_fixtures():
|
|
"""Original BF16 matmul, including Qwen's grouped attention layouts."""
|
|
assert mx.__version__ == "0.32.2"
|
|
for kind in ("qk","pv"):
|
|
for batch,rows,total in ((1,1,17),(1,3,17),(1,8,17),(1,16,17),
|
|
(1,64,1025),(1,2048,2048),(2,3,17),(2,16,33)):
|
|
if kind=="qk":
|
|
a=pattern(batch*rows*24*256,91).reshape(batch,rows,24,256).transpose(0,2,1,3).reshape(batch,2,12,rows,256)
|
|
b=pattern(batch*2*total*256,92).reshape(batch,2,1,total,256).swapaxes(-1,-2)
|
|
else:
|
|
a=pattern(batch*24*rows*total,91).reshape(batch,2,12,rows,total)
|
|
b=pattern(batch*2*total*256,92).reshape(batch,2,1,total,256)
|
|
emit(f"dense_batch_{kind}_B{batch}_r{rows}_t{total}",[mx.matmul(a,b)],kind=kind,batch=batch,rows=rows,total=total)
|
|
for kind in ("cross","transpose","copy","collapse","flatten","offset"):
|
|
for rows in (1,3,16):
|
|
k,n=64,37
|
|
if kind=="cross":
|
|
a=pattern(2*rows*k,91).reshape(2,1,rows,k)
|
|
b=pattern(3*n*k,92).reshape(1,3,n,k).swapaxes(-1,-2)
|
|
elif kind=="transpose":
|
|
a=pattern(3*rows*k,91).reshape(3,k,rows).swapaxes(-1,-2)
|
|
b=pattern(3*k*n,92).reshape(3,k,n)
|
|
elif kind=="copy":
|
|
a=pattern(3*rows*k*2,91).reshape(3,rows,k*2)[...,1::2]
|
|
b=pattern(3*k*n*2,92).reshape(3,k,n*2)[...,1::2]
|
|
elif kind in ("collapse","flatten"):
|
|
a=pattern(3*rows*k,91).reshape(3,rows,k)
|
|
b=pattern(n*k,92).reshape(n,k).T
|
|
if kind=="collapse": b=b[None,...]
|
|
else:
|
|
a=pattern(3*rows*k+1,91)[1:].reshape(3,rows,k)
|
|
b=pattern(3*n*k+1,92)[1:].reshape(3,n,k).swapaxes(-1,-2)
|
|
emit(f"dense_batch_{kind}_r{rows}",[mx.matmul(a,b)],kind=kind,rows=rows)
|
|
|
|
|
|
def ple_projection_fixtures():
|
|
"""Actual PLELayer after supplied embedding: lookup is NOT tested here."""
|
|
import mlx.nn as nn
|
|
from mlx_lm.models.cache import ArraysCache
|
|
from mtplx.models.qwen4_exp import PLELayer, TextArgs
|
|
assert mx.__version__ == "0.32.2"
|
|
class SuppliedEmbedding(nn.Module):
|
|
def __call__(self, ids, cache, state_idx):
|
|
return self.value
|
|
for layout in ("dense", "quantized"):
|
|
layer=PLELayer(TextArgs(ngram_sidecar=True),0)
|
|
layer.ple_embedding=SuppliedEmbedding()
|
|
for name,n,salt in (("key_proj",10240,71),("value_proj",2560,73)):
|
|
if layout=="dense": getattr(layer,name).weight=pattern(n*2560,salt).reshape(n,2560)
|
|
else: setattr(layer,name,quantized_linear_fixture(n,2560,4,64,salt))
|
|
for name,salt in (("norm_key",75),("norm_query",76),("norm_conv",77)):
|
|
getattr(layer,name).weight=pattern(10240,salt)
|
|
layer.conv_weight=pattern(10240*4,78).reshape(10240,4,1)
|
|
for batch,rows in ((1,1),(1,4),(1,7),(1,8),(1,32),(1,2048),(2,7)):
|
|
for state in ("none","empty","seeded"):
|
|
cache=None if state=="none" else ArraysCache(4)
|
|
if state=="seeded": cache[2]=pattern(batch*12*10240,79).reshape(batch,12,10240)
|
|
for step,tokens in enumerate((rows,1)):
|
|
hidden=pattern(batch*tokens*10240,80+step).reshape(batch,tokens,10240)
|
|
emb=pattern(batch*tokens*2560,82+step).reshape(batch,tokens,2560)
|
|
layer.ple_embedding.value=emb if step==0 else emb.astype(mx.float32)
|
|
output=layer(hidden,mx.zeros((batch,tokens),mx.int32),cache)
|
|
emit(f"ple_projection_{layout}_b{batch}_r{rows}_{state}_step{step}",
|
|
[output] if cache is None else [output,cache[2]],
|
|
layout=layout,batch=batch,rows=rows,state=state,step=step)
|
|
|
|
|
|
def shared_mlp_fixtures():
|
|
"""Only the real fused shared-expert class; no full model or downloads."""
|
|
from mtplx.models.qwen4_exp import _FusedGateUpMLP
|
|
assert mx.__version__ == "0.32.2"
|
|
for rows,k,intermediate in ((1,128,64),(7,128,64),(1,2560,640),
|
|
(4,2560,640),(7,2560,640),(2048,2560,640)):
|
|
for bits in (4,8):
|
|
for group in (32,64):
|
|
gu = quantized_linear_fixture(2*intermediate,k,bits,group,46)
|
|
down = quantized_linear_fixture(k,intermediate,bits,group,48)
|
|
layer = _FusedGateUpMLP(down,gu.weight,gu.scales,gu.biases,group,bits,"affine")
|
|
layer.eval()
|
|
x = pattern(rows*k,44).reshape(1,rows,k)
|
|
emit(f"shared_mlp_r{rows}_k{k}_i{intermediate}_b{bits}_g{group}",
|
|
[layer(x)],rows=rows,k=k,intermediate=intermediate,bits=bits,group=group)
|
|
|
|
|
|
def gdn_fixture_layer():
|
|
from mtplx.models.qwen4_exp import GatedDeltaNet, TextArgs, _FusedGDNInProj
|
|
assert mx.__version__ == "0.32.2"
|
|
for flag in ("MTPLX_FUSED_GDN_STEP", "MTPLX_FUSED_GDN_CONVNORM",
|
|
"MTPLX_FUSED_CONVNORM_VERIFY", "MTPLX_FUSED_GDN_OUT"):
|
|
os.environ[flag] = "1"
|
|
layer = GatedDeltaNet(TextArgs())
|
|
projection = quantized_linear_fixture(16480, 2560, 4, 64, 26)
|
|
layer.in_proj_fused = _FusedGDNInProj(projection.weight, projection.scales,
|
|
projection.biases, 64, 4, "affine", [10240, 16384, 16432])
|
|
layer.out_proj = quantized_linear_fixture(2560, 6144, 4, 64, 28)
|
|
layer.conv1d.weight = pattern(10240*4, 8).reshape(10240,4,1)
|
|
layer.A_log, layer.dt_bias, layer.norm.weight = pattern(48,12), pattern(48,13), pattern(128,14)
|
|
layer.eval()
|
|
return layer
|
|
|
|
|
|
def gdn_graph_lifecycle_fixtures():
|
|
"""Actual GatedDeltaNet: absent caches, batches and carried metadata/capture."""
|
|
from mlx_lm.models.cache import ArraysCache
|
|
from mtplx.models.qwen4_exp import _VERIFY_CAPTURE
|
|
layer = gdn_fixture_layer()
|
|
for batch in (1,2):
|
|
for rows in (1,7):
|
|
for mode in ("none", "plain", "metadata"):
|
|
cache = None if mode == "none" else ArraysCache(2)
|
|
if mode == "metadata":
|
|
cache.lengths = mx.array([rows+1,rows-1][:batch], dtype=mx.int32)
|
|
cache.left_padding = mx.array([1,0][:batch], dtype=mx.int32)
|
|
mask = mx.array((np.arange(batch*rows)%3)!=0).reshape(batch,rows) if mode=="metadata" else None
|
|
for step in range(2):
|
|
old_capture = getattr(cache, "_mtplx_verify_rows", None)
|
|
token = _VERIFY_CAPTURE.set(step == 0)
|
|
try:
|
|
output = layer(pattern(batch*rows*2560,30+step).reshape(batch,rows,2560),mask=mask,cache=cache)
|
|
finally:
|
|
_VERIFY_CAPTURE.reset(token)
|
|
outputs = [output] if cache is None else [output,cache[0],cache[1]]
|
|
if mode == "metadata":
|
|
outputs += [cache.lengths,cache.left_padding]
|
|
emit(f"gdn_lifecycle_b{batch}_r{rows}_{mode}_step{step}", outputs,
|
|
batch=batch,rows=rows,mode=mode,step=step,
|
|
capture_preserved=(getattr(cache,"_mtplx_verify_rows",None) is old_capture) if step else None)
|
|
|
|
|
|
def gdn_commit_fixtures():
|
|
"""Actual family commit method, with a pure-GDN layer and no model load."""
|
|
from types import SimpleNamespace
|
|
from mlx_lm.models.cache import ArraysCache
|
|
from mtplx.models.qwen4_exp import Qwen4ExpTextModel, verify_capture_scope
|
|
layer = gdn_fixture_layer()
|
|
model = SimpleNamespace(layers=[SimpleNamespace(linear_attn=layer)])
|
|
for rows in (1,2,4,6):
|
|
for keep in range(1,rows+1):
|
|
for conv_present in (False,True):
|
|
cache = ArraysCache(2)
|
|
cache[0] = pattern(3*10240,7).reshape(1,3,10240) if conv_present else None
|
|
cache[1] = pattern(48*128*128,15).astype(mx.float32).reshape(1,48,128,128)
|
|
snapshot = list(cache.state)
|
|
with verify_capture_scope():
|
|
layer(pattern(rows*2560,30).reshape(1,rows,2560),cache=cache)
|
|
mx.eval(cache.state, cache._mtplx_verify_rows)
|
|
# Metadata is not rewound/advanced by the commit branch.
|
|
cache.lengths = mx.array([17],mx.int32)
|
|
cache.left_padding = mx.array([-3],mx.int32)
|
|
assert Qwen4ExpTextModel.commit_verified_window(model,[cache],[snapshot],
|
|
keep_tokens=keep,verified_tokens=rows)
|
|
assert cache._mtplx_verify_rows is None
|
|
outputs = [cache[0],cache[1],cache.lengths,cache.left_padding]
|
|
# Make the next decode consume the replay lazily. Mask/length
|
|
# controls remain deliberately present as in the original cache.
|
|
hidden = layer(pattern(2560,31).reshape(1,1,2560),cache=cache)
|
|
outputs += [hidden,cache[0],cache[1]]
|
|
emit(f"gdn_commit_r{rows}_keep{keep}_conv{int(conv_present)}", outputs,
|
|
rows=rows,keep=keep,conv_present=conv_present)
|
|
|
|
|
|
def qsa_f32_fixtures():
|
|
"""Run separately for MLX_ENABLE_TF32=0 and 1 (cached by the runtime).
|
|
|
|
These are the actual MTPLX fallback expression, not reconstructed matmul
|
|
microbenchmarks. Layouts cover initial/grown pooled mirrors and view copies.
|
|
"""
|
|
from mtplx.kernels.qsa_indexer_prefill import qsa_indexer_prefill_scores
|
|
assert mx.__version__ == "0.32.2"
|
|
tf32 = os.environ.get("MLX_ENABLE_TF32", "1") != "0"
|
|
for rows, blocks in ((1,1),(2,1),(4,7),(1,16),(4,32),(10,40),
|
|
(16,63),(16,64),(32,128),(33,129),(65,513),(2048,513)):
|
|
layouts = [([512,128,1],[1,128]), ([512,1,4],[1,128]),
|
|
([1024,256,2],[blocks+3,1]), ([519,128,1],[2,256]),
|
|
([128,rows*128,1],[blocks,1]), ([512,1,4],[blocks+3,1]),
|
|
([512,128,1],[blocks*2,2]), ([0,128,1],[1,128])]
|
|
for layout, (qs, ps) in enumerate(layouts):
|
|
for bf16 in (True, False):
|
|
qcount = 1+sum((d-1)*s for d,s in zip((rows,4,128),qs))
|
|
pcount = 1+(128-1)*ps[0]+(blocks-1)*ps[1]
|
|
def backing(count, salt, bf16=False):
|
|
value = pattern(count,salt)
|
|
if not bf16:
|
|
value = value.astype(mx.float32)+mx.array(
|
|
(np.arange(count)%7).astype(np.float32)/65536)
|
|
return value
|
|
q = mx.as_strided(backing(qcount,94,bf16), (1,rows,4,128), [0,*qs])
|
|
p = mx.as_strided(backing(pcount,95), (1,1,128,blocks), [0,0,*ps])
|
|
emit(f"qsa_f32_score_r{rows}_n{blocks}_l{layout}_bf{int(bf16)}_tf{int(tf32)}",
|
|
[qsa_indexer_prefill_scores(q,p,head_dim=128)], rows=rows,blocks=blocks,
|
|
q_strides=qs,pooled_strides=ps,q_count=qcount,pooled_count=pcount,
|
|
bf16=bf16,tf32=tf32)
|
|
from mtplx.kernels.qsa_indexer_prefill import qsa_indexer_prefill_metal
|
|
for rows,blocks,budget_rows in ((7,513,3),(65,600,48)):
|
|
for ds in (1,2):
|
|
for f32 in ((False,False),(False,True),(True,False),(True,True)):
|
|
q = pattern(rows*512*ds,96)
|
|
p = pattern(blocks*128*ds,97)
|
|
q = q.astype(mx.float32) if f32[0] else q
|
|
p = p.astype(mx.float32) if f32[1] else p
|
|
# Cast before constructing strided F32 views, as in the Rust
|
|
# fixture backing; otherwise AsType silently materializes them.
|
|
q = q.reshape(1,rows,4,128*ds)[...,::ds]
|
|
p = p.reshape(1,blocks,128*ds)[...,::ds]
|
|
total = blocks*4+1
|
|
budget = blocks*4*5*budget_rows
|
|
for mode in ("blocks","dense_mask","row_tokens"):
|
|
out = qsa_indexer_prefill_metal(q,p,pos_start=total-rows,total_tokens=total,
|
|
block_topk=512,compress_ratio=4,logical_blocks=blocks,
|
|
output_total_tokens=total+3,mode=mode,score_workspace_bytes=budget)
|
|
emit(f"qsa_f32_chain_r{rows}_ds{ds}_f{int(f32[0])}{int(f32[1])}_tf{int(tf32)}_{mode}",
|
|
[out] if mode=="dense_mask" else out,rows=rows,blocks=blocks,ds=ds,
|
|
input_f32=f32,tf32=tf32,total=total,budget=budget,mode=mode)
|
|
|
|
|
|
def qsa_eager_fixtures():
|
|
"""Tap the real eager oracle without replacing its arithmetic or sorting."""
|
|
from mtplx.models.qwen4_exp import QSAIndexer, QSACache, TextArgs
|
|
from mtplx.attention_context import attention_phase
|
|
os.environ["MTPLX_QSA_SCORE_TILE_ROWS"] = "0"
|
|
tf32 = os.environ.get("MLX_ENABLE_TF32", "1") != "0"
|
|
indexer = QSAIndexer(TextArgs())
|
|
for rows,blocks in ((1,1),(4,7),(1,513),(4,513),(7,600),
|
|
(33,2049),(65,513),(2048,513),(2048,1025),(1,65536)):
|
|
total = blocks*4+3
|
|
for variant in range(3):
|
|
q = pattern(rows*512,98).reshape(1,rows,4,128)
|
|
if variant == 1: q = mx.zeros_like(q)
|
|
if variant == 2: q = q.astype(mx.float32)
|
|
pooled = pattern(blocks*128,99).reshape(1,blocks,128)
|
|
cache = QSACache()
|
|
cache.pooled,cache.pooled_len = pooled,blocks
|
|
original_partition, original_where = mx.argpartition,mx.where
|
|
tapped = {}
|
|
def partition(a,*args,**kwargs):
|
|
result = original_partition(a,*args,**kwargs)
|
|
assert "sorted" not in tapped, "unexpected tiled eager oracle"
|
|
tapped.update(masked=a,sorted=result)
|
|
return result
|
|
def where(condition,a,b,*args,**kwargs):
|
|
if "valid" not in tapped:
|
|
tapped["valid"] = condition
|
|
return original_where(condition,a,b,*args,**kwargs)
|
|
try:
|
|
mx.argpartition,mx.where = partition,where
|
|
indexer._select_eager(q,total-rows,cache,pooled,total)
|
|
finally:
|
|
mx.argpartition,mx.where = original_partition,original_where
|
|
emit(f"qsa_eager_topk_r{rows}_n{blocks}_v{variant}_tf{int(tf32)}",
|
|
[tapped["sorted"][:,blocks-min(blocks,512):],tapped["valid"],tapped["masked"]],
|
|
rows=rows,blocks=blocks,variant=variant,total=total,tf32=tf32)
|
|
if rows > 1 and total-rows >= 2049:
|
|
knobs = {"MTPLX_QSA_PREFILL":"1","MTPLX_QSA_PREFILL_MIN_ROWS":"2",
|
|
"MTPLX_QSA_PREFILL_MIN_CONTEXT":"2049"}
|
|
saved = {k:os.environ.get(k) for k in knobs}
|
|
try:
|
|
os.environ.update(knobs)
|
|
with attention_phase("prefill"):
|
|
result = indexer._select_eager(q,total-rows,cache,pooled,total)
|
|
finally:
|
|
for key,value in saved.items():
|
|
if value is None: os.environ.pop(key,None)
|
|
else: os.environ[key] = value
|
|
assert result[0] == "flash_prefill"
|
|
emit(f"qsa_eager_blocks_r{rows}_n{blocks}_v{variant}_tf{int(tf32)}",result[1:],
|
|
rows=rows,blocks=blocks,variant=variant,total=total,tf32=tf32)
|
|
if blocks > 512:
|
|
knobs = {"MTPLX_QSA_PREFILL":"0","MTPLX_QSA_FLASH":"0",
|
|
"MTPLX_QSA_GATHER":"0","MTPLX_QSA_GATHER_DECODE":"0",
|
|
"MTPLX_QSA_GATHER_MIN_CONTEXT":"0","MTPLX_QSA_GATHER_MAX_ROWS":str(max(rows,2))}
|
|
saved = {k:os.environ.get(k) for k in knobs}
|
|
try:
|
|
os.environ.update(knobs)
|
|
for tail in (0,1,3):
|
|
dense_total = blocks*4+tail
|
|
result = indexer._select_eager(q,dense_total-rows,cache,pooled,dense_total)
|
|
assert isinstance(result,mx.array) and result.dtype == mx.bool_
|
|
emit(f"qsa_eager_dense_r{rows}_n{blocks}_v{variant}_tail{tail}_tf{int(tf32)}",[result],
|
|
rows=rows,blocks=blocks,variant=variant,total=dense_total,tf32=tf32)
|
|
if rows > 1:
|
|
os.environ["MTPLX_QSA_GATHER"] = "1"
|
|
result = indexer._select_eager(q,dense_total-rows,cache,pooled,dense_total)
|
|
os.environ["MTPLX_QSA_GATHER"] = "0"
|
|
assert result[0] == "gather_rows"
|
|
emit(f"qsa_eager_gather_r{rows}_n{blocks}_v{variant}_tail{tail}_tf{int(tf32)}",result[1:],
|
|
rows=rows,blocks=blocks,variant=variant,total=dense_total,tf32=tf32)
|
|
else:
|
|
for gather in (False,True):
|
|
os.environ["MTPLX_QSA_FLASH"] = "0" if gather else "1"
|
|
os.environ["MTPLX_QSA_GATHER_DECODE"] = "1"
|
|
result = indexer._select_eager(q,dense_total-1,cache,pooled,dense_total)
|
|
if gather:
|
|
assert isinstance(result,mx.array)
|
|
values,tail_start = [result],dense_total//4*4
|
|
else:
|
|
assert result[0] == "flash"
|
|
values,tail_start = [result[1]],result[2]
|
|
emit(f"qsa_eager_decode_r1_n{blocks}_v{variant}_tail{tail}_g{int(gather)}_tf{int(tf32)}",values,
|
|
rows=1,blocks=blocks,variant=variant,total=dense_total,tf32=tf32,gather=gather,tail_start=tail_start)
|
|
os.environ["MTPLX_QSA_FLASH"] = "0"
|
|
os.environ["MTPLX_QSA_GATHER_DECODE"] = "0"
|
|
finally:
|
|
for key,value in saved.items():
|
|
if value is None: os.environ.pop(key,None)
|
|
else: os.environ[key] = value
|
|
|
|
|
|
def qsa_eager_tiled_fixtures():
|
|
"""Actual eager entry point, including per-tile eval and output routing."""
|
|
from mtplx.models.qwen4_exp import QSAIndexer, TextArgs, QSACache
|
|
from mtplx.attention_context import attention_phase
|
|
tf32 = os.environ.get("MLX_ENABLE_TF32", "1") != "0"
|
|
indexer = QSAIndexer(TextArgs())
|
|
knobs = {"MTPLX_QSA_SCORE_TILE_ROWS":"0", "MTPLX_QSA_PREFILL":"0",
|
|
"MTPLX_QSA_PREFILL_MIN_ROWS":"2", "MTPLX_QSA_PREFILL_MIN_CONTEXT":"2049",
|
|
"MTPLX_QSA_GATHER":"1", "MTPLX_QSA_GATHER_MIN_CONTEXT":"0",
|
|
"MTPLX_QSA_GATHER_MAX_ROWS":"2048", "MTPLX_QSA_FLASH":"1",
|
|
"MTPLX_QSA_GATHER_DECODE":"1"}
|
|
saved = {k:os.environ.get(k) for k in knobs}
|
|
try:
|
|
os.environ.update(knobs)
|
|
for rows,blocks,tiles in ((7,513,(1,3,7,8)), (33,600,(3,32)),
|
|
(65,2049,(1,32)), (2048,1025,(32,2047)),
|
|
(7,65536,(3,))):
|
|
for variant in range(3):
|
|
ds = 2 if variant == 2 else 1
|
|
q = pattern(rows*512*ds,98)
|
|
if variant == 1: q = mx.zeros_like(q)
|
|
if variant == 2: q = q.astype(mx.float32)
|
|
q = q.reshape(1,rows,4,128*ds)[...,::ds]
|
|
pooled = pattern(blocks*128,99).reshape(1,blocks,128)
|
|
cache = QSACache()
|
|
cache.pooled,cache.pooled_len = pooled,blocks
|
|
for tile in tiles:
|
|
total = blocks*4+variant%2
|
|
os.environ["MTPLX_QSA_SCORE_TILE_ROWS"] = str(tile)
|
|
for prefill in (False,True):
|
|
os.environ["MTPLX_QSA_PREFILL"] = str(int(prefill))
|
|
eval_rows = []
|
|
original_eval = mx.eval
|
|
def evaluate(*arrays):
|
|
assert len(arrays) == 1 and arrays[0].dtype == mx.uint32
|
|
eval_rows.append(arrays[0].shape[0])
|
|
return original_eval(*arrays)
|
|
try:
|
|
mx.eval = evaluate
|
|
with attention_phase("prefill"):
|
|
result = indexer._select_eager(q,total-rows,cache,pooled,total)
|
|
finally:
|
|
mx.eval = original_eval
|
|
expected = [min(tile,rows-start) for start in range(0,rows,tile)] if tile < rows else []
|
|
assert eval_rows == expected, (eval_rows,expected)
|
|
if prefill and total-rows >= 2049:
|
|
assert result[0] == "flash_prefill"
|
|
kind, values = "blocks", result[1:]
|
|
elif tile >= rows:
|
|
assert result[0] == "gather_rows"
|
|
kind, values = "gather", result[1:]
|
|
else:
|
|
# Tiled scoring disables rows-gather even if all
|
|
# its flags and thresholds would otherwise engage.
|
|
assert isinstance(result,mx.array) and result.dtype == mx.bool_
|
|
kind, values = "dense", [result]
|
|
emit(f"qsa_eager_{kind}_tiled_r{rows}_n{blocks}_v{variant}_t{tile}_p{int(prefill)}_tf{int(tf32)}",
|
|
values,rows=rows,blocks=blocks,variant=variant,total=total,tf32=tf32,
|
|
tile=tile,ds=ds,prefill=prefill,eval_rows=eval_rows)
|
|
finally:
|
|
for key,value in saved.items():
|
|
if value is None: os.environ.pop(key,None)
|
|
else: os.environ[key] = value
|
|
|
|
|
|
def qsa_compiled_state_fixtures():
|
|
"""Original compiled backing policy and device-indexed copy primitives.
|
|
|
|
This checks prerequisites, not execution of the compiled indexer graph.
|
|
"""
|
|
from mtplx.qsa_mtp_precompute import precompute_qsa_replay_capacity
|
|
from mtplx.models.qwen4_exp import QSAIndexer, TextArgs, QSACache
|
|
for start in (0,3,2048,4095):
|
|
for window in (0,1,3,4,255,256,257,1025,2048):
|
|
for raw,pool in ((0,0),(768,256),(8192,4096)):
|
|
plan = precompute_qsa_replay_capacity(start_offset=start,window_tokens=window,
|
|
compress_ratio=4,current_raw_capacity=raw,current_pooled_capacity=pool)
|
|
emit(f"qsa_compile_capacity_p{start}_s{window}_c{raw}_{pool}",[],
|
|
start=start,window=window,existing=[raw,pool],capacity=list(plan.graph_key))
|
|
indexer,cache = QSAIndexer(TextArgs()),QSACache()
|
|
for step,(rows,reserve) in enumerate(((1,(0,0)),(1025,(0,0)),(2048,(8192,4096)),(8193,(0,0)))):
|
|
if reserve != (0,0):
|
|
cache.reserve_indexer_capacity(raw_capacity=reserve[0],pooled_capacity=reserve[1])
|
|
raw,pool = indexer._ensure_compiled_backings(cache,dtype=mx.bfloat16,pos_start=0,rows=rows)
|
|
emit(f"qsa_compile_materialize_{step}",[raw,pool],step=step,rows=rows,reserve=reserve,
|
|
capacity=[raw.shape[1],pool.shape[1]],offset=cache.offset,pooled_len=cache.pooled_len)
|
|
# Make the next capacity transition prove preservation, not just zero fill.
|
|
cache.write_raw(pattern(3*128,160+step).reshape(1,3,128))
|
|
mx.eval(cache.raw_keys)
|
|
for rows in (1,4,33,2048):
|
|
for stride in (128,704):
|
|
cap = rows+7
|
|
source = pattern(cap*stride,151).reshape(1,cap,stride)[...,:128]
|
|
update = pattern(rows*stride,152).reshape(1,rows,stride)[...,:128]
|
|
start = mx.array([3],dtype=mx.int32)
|
|
mx.eval(source,update,start)
|
|
sliced = mx.slice(source,start,axes=(1,),slice_size=(1,rows,128))
|
|
updated = mx.slice_update(source,update,start,axes=(1,))
|
|
emit(f"qsa_compile_dynamic_r{rows}_stride{stride}",[sliced,updated,source],
|
|
rows=rows,stride=stride,capacity=cap,start=3)
|
|
|
|
|
|
def qsa_compiled_window_fixtures():
|
|
"""Real compiled cache-maintenance calls with retained original leaves."""
|
|
from mtplx.models.qwen4_exp import QSAIndexer, TextArgs
|
|
from mtplx.kernels.qsa_indexer_compile import QSACompiledIndexerCore
|
|
freq = QSAIndexer(TextArgs())._inv_freq
|
|
norm = pattern(128,72)
|
|
cases = ((1,256,256,0,0),(1,256,256,253,63),(1,256,256,255,63),
|
|
(4,4096,1024,2051,512),(5,4096,1024,2050,511),(5,256,256,0,0),
|
|
(2048,2048,512,0,0),(2048,4096,1024,2048,512),
|
|
(1025,2048,512,100,24),(1,4096,1024,4095,1024),
|
|
(33,4096,1024,2050,512),(1020,1024,256,0,0),(512,1024,256,0,0))
|
|
for case,(rows,raw_cap,pool_cap,pos,pool_len) in enumerate(cases):
|
|
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)
|
|
for stride in (640,704):
|
|
qk = pattern(rows*stride,150).reshape(1,rows,stride)[...,:640]
|
|
raw = pattern(raw_cap*128,151).reshape(1,raw_cap,128)
|
|
pool = pattern(pool_cap*128,152).reshape(1,pool_cap,128)
|
|
total = pos+rows
|
|
logical = total//4
|
|
scalars = [mx.array([n],dtype=mx.int32) for n in (pos,total,logical,pool_len)]
|
|
mx.eval(qk,raw,pool,norm,freq,*scalars)
|
|
result = core.select_qk_rows(qk,raw,pool,pos_start=scalars[0],total_tokens=scalars[1],
|
|
logical_blocks=scalars[2],pooled_len=scalars[3],mode="update_only")
|
|
assert result.selection is None
|
|
# Keep raw/pool live across eval: this is the reference copy branch,
|
|
# not evidence for donation or the complete model's lazy schedule.
|
|
emit(f"qsa_compile_window_c{case}_stride{stride}",[result.raw_keys,result.pooled,raw,pool],
|
|
rows=rows,stride=stride,capacity=[raw_cap,pool_cap],pos=pos,logical=logical,
|
|
pooled_len=pool_len,frequencies=freq.tolist())
|
|
assert result.pooled_len.tolist()==[logical] and result.offset.tolist()==[total]
|
|
|
|
|
|
def qsa_compiled_graph_contract(result, selection, inputs):
|
|
"""Fingerprint the actual optimized runtime graph before evaluation.
|
|
|
|
Evaluated constant leaves are omitted, not pending primitive dependencies.
|
|
Multi-output leaf positions remain explicit. No execution-order inference
|
|
is made from the DOT traversal order.
|
|
"""
|
|
import io
|
|
names = dict(inputs)
|
|
roots = [f"selected{i}" for i in range(len(selection))] + ["raw_next","pool_next"]
|
|
names.update({f"selected{i}":v for i,v in enumerate(selection)})
|
|
names.update(raw_next=result.raw_keys,pool_next=result.pooled)
|
|
dot = io.StringIO()
|
|
mx.export_to_dot(dot,**names)
|
|
return compiled_dot_contract(dot.getvalue(), roots, inputs)
|
|
|
|
|
|
def compiled_dot_contract(dot, roots, inputs):
|
|
"""Shared ordered-dependency fingerprint for an original runtime DOT."""
|
|
import re
|
|
nodes,producer = {},{}
|
|
for line in dot.splitlines():
|
|
if m:=re.search(r'(\d+) \[label ="([^"]+)"',line):
|
|
nodes[m[1]] = dict(op=m[2],inputs=[],outputs=[])
|
|
elif m:=re.fullmatch(r'"([^"]+)" -> (\d+)',line):
|
|
nodes[m[2]]["inputs"].append(m[1])
|
|
elif m:=re.fullmatch(r'(\d+) -> "([^"]+)"',line):
|
|
node = nodes[m[1]]
|
|
producer[m[2]] = (m[1],len(node["outputs"]))
|
|
node["outputs"].append(m[2])
|
|
hashes = {}
|
|
def leaf(name):
|
|
if name in producer:
|
|
node,index = producer[name]
|
|
return fingerprint(node)+":"+str(index)
|
|
return name if name in inputs else None
|
|
def fingerprint(key):
|
|
if key not in hashes:
|
|
node = nodes[key]
|
|
parents = [value for name in node["inputs"] if (value:=leaf(name)) is not None]
|
|
hashes[key] = hashlib.sha256(json.dumps([node["op"],parents],separators=(",",":")).encode()).hexdigest()
|
|
return hashes[key]
|
|
root_hashes = [leaf(name) for name in roots]
|
|
return [sorted(hashes.values()),root_hashes]
|
|
|
|
|
|
def qsa_compiled_indexer_fixtures(*,hidden=False,bits=4,group=64,graph=False):
|
|
"""All fixed-output modes of the original compiled BF16 Q/K or hidden entry."""
|
|
from mtplx.models.qwen4_exp import QSAIndexer, TextArgs
|
|
from mtplx.kernels.qsa_indexer_compile import QSACompiledIndexerCore
|
|
freq = QSAIndexer(TextArgs())._inv_freq
|
|
norm = pattern(128,72)
|
|
projection = None
|
|
if hidden:
|
|
import mlx.nn as nn
|
|
projection = nn.QuantizedLinear(2560,640,bias=False,group_size=group,bits=bits)
|
|
words = (np.arange(640*2560*bits//32,dtype=np.uint64)*2654435761+12345).astype(np.uint32)
|
|
projection.weight = mx.array(words).reshape(640,2560*bits//32)
|
|
projection.scales = (pattern(640*2560//group,75)/64).reshape(640,2560//group)
|
|
projection.biases = (pattern(640*2560//group,76)/64).reshape(640,2560//group)
|
|
modes = ("blocks","row_tokens","dense_mask")
|
|
scenarios = [(rows,pos,frontier,mode,budget) for rows,pos,frontier in
|
|
((1,2051,512),(4,2051,512),(5,2050,511),(33,2050,512))
|
|
for mode in modes for budget in (32*1024*1024,3*1024*4)]
|
|
scenarios += [(rows,pos,frontier,"prefill_blocks",budget) for rows,pos,frontier in
|
|
((4,2051,512),(33,2050,512),(65,2050,512),(2048,2048,512))
|
|
for budget in (128*1024*1024,(128 if rows==2048 else 48 if rows==65 else 3)*1024*4)]
|
|
scenarios += [(4,2051,512,"update_only",32*1024*1024)]
|
|
for case,(rows,pos,pool_len,mode,budget) in enumerate(scenarios):
|
|
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=budget,prefill_score_workspace_bytes=budget,project_qk=projection)
|
|
for stride in ((2560,) if hidden else (640,704)):
|
|
source = pattern(rows*stride,150).reshape(1,rows,stride)
|
|
if not hidden: source = source[...,:640]
|
|
raw = pattern(4096*128,151).reshape(1,4096,128)
|
|
pool = pattern(1024*128,152).reshape(1,1024,128)
|
|
total,logical = pos+rows,(pos+rows)//4
|
|
state = [mx.array([n],dtype=mx.int32) for n in (pos,total,logical,pool_len)]
|
|
mx.eval(source,raw,pool,norm,freq,*state)
|
|
if graph and projection is not None:
|
|
mx.eval(projection.weight,projection.scales,projection.biases)
|
|
entry = core.select_hidden if hidden else core.select_qk_rows
|
|
out = entry(source,raw,pool,pos_start=state[0],total_tokens=state[1],
|
|
logical_blocks=state[2],pooled_len=state[3],mode=mode)
|
|
selection = [] if out.selection is None else list(out.selection) if isinstance(out.selection,tuple) else [out.selection]
|
|
if graph:
|
|
contract = qsa_compiled_graph_contract(out,selection,dict(source=source,raw=raw,pooled=pool,
|
|
norm=norm,freq=freq,pos=state[0],total=state[1],logical=state[2],frontier=state[3]))
|
|
print(json.dumps(dict(case=case,hidden=hidden,rows=rows,stride=stride,mode=mode,budget=budget,
|
|
contract=contract),separators=(",",":")),flush=True)
|
|
mx.eval(out.raw_keys,out.pooled,selection)
|
|
continue
|
|
mx.eval(out.raw_keys,out.pooled,selection)
|
|
states = {"raw_next":out.raw_keys,"pool_next":out.pooled,"raw_old":raw,"pool_old":pool}
|
|
hashes = {key:hashlib.sha256(np.asarray(a.astype(mx.float32)).astype('<f4').tobytes()).hexdigest() for key,a in states.items()}
|
|
prefix = f"qsa_compile_hidden_b{bits}_g{group}" if hidden else "qsa_compile_indexer"
|
|
extra = dict(bits=bits,group=group) if hidden else {}
|
|
emit(f"{prefix}_c{case}_stride{stride}",selection,rows=rows,stride=stride,mode=mode,
|
|
budget=budget,capacity=[4096,1024],pos=pos,logical=logical,pooled_len=pool_len,
|
|
frequencies=freq.tolist(),state_sha256=hashes,**extra)
|
|
|
|
|
|
def qsa_compiled_route_fixtures():
|
|
"""Original host routing methods, with real typed arrays; no graph eval.
|
|
|
|
This covers the installed B1/BF16 contract, not unsupported tensor dtypes.
|
|
Cache extents deliberately include non-power-of-two eager backings.
|
|
"""
|
|
from mtplx.models.qwen4_exp import QSAIndexer, TextArgs, QSACache
|
|
from mtplx.attention_context import attention_phase
|
|
assert mx.__version__ == "0.32.2"
|
|
assert mx.metal.is_available() and mx.default_device() == mx.gpu
|
|
defaults = dict(fused=True,compiled=True,prefill=True,phase_prefill=True,
|
|
prefill_min_rows=32,prefill_compile_rows=2048,prefill_min_context=32768,
|
|
prefill_workspace=128*1024*1024,tile_rows=0,flash=True,gather_decode=True,
|
|
gather=True,gather_max_rows=8,gather_min_context=0,tf32=True)
|
|
env_names = dict(fused="MTPLX_FUSED_QSA_INDEXER",compiled="MTPLX_COMPILED_QSA_INDEXER",
|
|
prefill="MTPLX_QSA_PREFILL",prefill_min_rows="MTPLX_QSA_PREFILL_MIN_ROWS",
|
|
prefill_compile_rows="MTPLX_QSA_PREFILL_COMPILE_ROWS",
|
|
prefill_min_context="MTPLX_QSA_PREFILL_MIN_CONTEXT",flash="MTPLX_QSA_FLASH",
|
|
gather_decode="MTPLX_QSA_GATHER_DECODE",gather="MTPLX_QSA_GATHER",
|
|
gather_max_rows="MTPLX_QSA_GATHER_MAX_ROWS",gather_min_context="MTPLX_QSA_GATHER_MIN_CONTEXT")
|
|
options = [defaults | change for change in ({},dict(fused=False),dict(compiled=False),
|
|
dict(prefill=False),dict(phase_prefill=False),dict(flash=False),
|
|
dict(flash=False,gather_decode=False),dict(gather=False),dict(gather_max_rows=3),
|
|
dict(gather_min_context=32772),dict(prefill_compile_rows=32),
|
|
dict(prefill_min_rows=2,prefill_min_context=2049))]
|
|
# pos, rows, offset, raw capacity, pooled capacity, pooled frontier
|
|
states = []
|
|
for pos,rows in ((0,0),(0,1),(0,3),(0,32),(0,2048),(2048,1),(2048,4),
|
|
(2048,2048),(2049,4),(2051,1),(2051,4),(2051,8),(2051,9),
|
|
(32767,2048),(32768,1),(32768,31),(32768,32),(32768,33),
|
|
(32768,2047),(32768,2048),(32768,2049)):
|
|
states.append([pos,rows,pos,((pos+255)//256)*256 if pos else None,
|
|
((pos//4+255)//256)*256 if pos else None,pos//4])
|
|
states.extend([
|
|
[1,1,0,256,256,0], # mismatched offset
|
|
[1,1,1,None,None,0], # missing prior raw
|
|
[257,1,257,256,256,64], # short raw
|
|
[0,1,0,None,None,1], # nonzero frontier without a pool
|
|
[2051,4,2051,2304,512,513], # frontier exceeds capacity
|
|
[2051,4,2051,2304,768,511], # missing two blocks, window fits one
|
|
[2051,4,2051,2304,768,512], # valid nonbucketed backing
|
|
[2051,4,2051,2304,768,700], # frontier beyond logical: clamp first
|
|
[0,1,0,0,0,0], # present zero-length leaves
|
|
])
|
|
saved = {env:os.environ.get(env) for env in env_names.values()}
|
|
indexer = QSAIndexer(TextArgs())
|
|
indexer.q_layernorm.weight = mx.ones((128,),mx.bfloat16)
|
|
indexer.k_layernorm.weight = mx.ones((128,),mx.bfloat16)
|
|
indexer.index_qk_proj.weight = mx.zeros((640,2560),mx.bfloat16)
|
|
cases = []
|
|
try:
|
|
for option_id,opt in enumerate(options):
|
|
os.environ.update({env:str(int(opt[key])) for key,env in env_names.items()})
|
|
with attention_phase("prefill" if opt["phase_prefill"] else "verify"):
|
|
for state_id,(pos,rows,offset,raw,pool,frontier) in enumerate(states):
|
|
cache = QSACache()
|
|
cache.kv.offset = offset
|
|
cache.raw_keys = None if raw is None else mx.zeros((1,raw,128),mx.bfloat16)
|
|
cache.pooled = None if pool is None else mx.zeros((1,pool,128),mx.bfloat16)
|
|
cache.pooled_len = frontier
|
|
mode = indexer._compiled_mode(decode=rows==1,rows=rows,
|
|
total=pos+rows,last_nb=(pos+rows)//4)
|
|
results = []
|
|
for supplied in (False,True):
|
|
source = mx.zeros((1,rows,640 if supplied else 2560),mx.bfloat16)
|
|
results.append(indexer._compiled_route_supported(source,cache,
|
|
pos_start=pos,qk_rows_supplied=supplied,decode=rows==1,mode=mode))
|
|
cases.append([option_id,state_id,mode,*results])
|
|
finally:
|
|
for env,value in saved.items():
|
|
if value is None: os.environ.pop(env,None)
|
|
else: os.environ[env] = value
|
|
print(json.dumps(dict(options=options,states=states,cases=cases),separators=(",",":")),flush=True)
|
|
|
|
|
|
def qsa_indexer_route_fixtures(*,compiled=False):
|
|
"""Whole original indexer calls, not manually selected outputs."""
|
|
from mtplx.models.qwen4_exp import QSAIndexer, TextArgs, QSACache
|
|
from mtplx.attention_context import attention_phase
|
|
import mlx.nn as nn
|
|
defaults = dict(fused=False,prefill=False,phase_prefill=True,prefill_min_rows=32,
|
|
prefill_min_context=2049,prefill_workspace=128*1024*1024,tile_rows=0,
|
|
flash=False,gather_decode=False,gather=False,gather_max_rows=8,
|
|
gather_min_context=0,tf32=os.environ.get("MLX_ENABLE_TF32","1")!="0")
|
|
env_names = dict(fused="MTPLX_FUSED_QSA_INDEXER",prefill="MTPLX_QSA_PREFILL",
|
|
prefill_min_rows="MTPLX_QSA_PREFILL_MIN_ROWS",prefill_min_context="MTPLX_QSA_PREFILL_MIN_CONTEXT",
|
|
tile_rows="MTPLX_QSA_SCORE_TILE_ROWS",flash="MTPLX_QSA_FLASH",
|
|
gather_decode="MTPLX_QSA_GATHER_DECODE",gather="MTPLX_QSA_GATHER",
|
|
gather_max_rows="MTPLX_QSA_GATHER_MAX_ROWS",gather_min_context="MTPLX_QSA_GATHER_MIN_CONTEXT")
|
|
knobs = [*env_names.values(),"MTPLX_COMPILED_QSA_INDEXER","MTPLX_QSA_PREFILL_SCORE_MB",
|
|
"MTPLX_QSA_PREFILL_COMPILE_ROWS"]
|
|
saved = {k:os.environ.get(k) for k in knobs}
|
|
scenarios = [
|
|
(0,(1,3,4,33,2048),{}), (0,(1,3,4,33,2048),dict(fused=True)),
|
|
(2048,(1,1,1,1,4,33),dict(flash=True,gather=True,gather_min_context=2053)),
|
|
(2048,(1,1,1,1,4,33),dict(fused=True,flash=True,gather=True,gather_min_context=2053)),
|
|
(2051,(1,4,33),dict(fused=True,flash=True,gather_decode=True,gather=True)),
|
|
(2051,(1,4,33),dict(fused=True,gather_decode=True)),
|
|
(2051,(1,4,33),dict(fused=True,gather=True,gather_max_rows=3)),
|
|
(2051,(4,33),dict(gather=True,tile_rows=3)),
|
|
(2048,(32,33,2048),dict(fused=True,prefill=True)),
|
|
(2048,(32,33,2048),dict(fused=True,prefill=True,phase_prefill=False)),
|
|
(32768,(4,32,33,2048),dict(prefill=True,prefill_min_context=32768,gather=True,gather_min_context=16384)),
|
|
(2051,(1,4,33),dict(gather=True,gather_min_context=99999)),
|
|
]
|
|
try:
|
|
os.environ["MTPLX_COMPILED_QSA_INDEXER"] = str(int(compiled))
|
|
os.environ["MTPLX_QSA_PREFILL_SCORE_MB"] = "128"
|
|
os.environ["MTPLX_QSA_PREFILL_COMPILE_ROWS"] = "2048"
|
|
if compiled:
|
|
scenarios.extend([
|
|
(32768,(2048,33,1,4),dict(fused=True,prefill=True,flash=True,gather=True)),
|
|
(0,(1,2048,1,4),dict(fused=True,phase_prefill=False)),
|
|
])
|
|
for scenario,(prefix,steps,overrides) in enumerate(scenarios):
|
|
options = defaults | overrides
|
|
if compiled: options.update(compiled=True,prefill_compile_rows=2048)
|
|
os.environ.update({env:str(int(options[key])) for key,env in env_names.items()})
|
|
for supplied in (False,True):
|
|
indexer = QSAIndexer(TextArgs())
|
|
indexer.q_layernorm.weight = pattern(128,72)
|
|
indexer.k_layernorm.weight = pattern(128,73)
|
|
proj = nn.QuantizedLinear(2560,640,bias=False,group_size=64,bits=4)
|
|
words = (np.arange(640*2560*4//32,dtype=np.uint64)*2654435761+12345).astype(np.uint32)
|
|
proj.weight = mx.array(words).reshape(640,2560*4//32)
|
|
proj.scales = (pattern(640*2560//64,75)/64).reshape(640,2560//64)
|
|
proj.biases = (pattern(640*2560//64,76)/64).reshape(640,2560//64)
|
|
indexer.index_qk_proj = proj
|
|
cache = QSACache()
|
|
if prefix:
|
|
cache.write_raw(pattern(prefix*128,88).reshape(1,prefix,128))
|
|
indexer._extend_pooled(cache,prefix)
|
|
kv = pattern(prefix*512,89).reshape(1,prefix,2,256).transpose(0,2,1,3)
|
|
cache.kv.update_and_fetch(kv,kv)
|
|
mx.eval(*cache.state,cache.pooled_f32_t)
|
|
lane = ["update_only"]
|
|
eager,fused,prefill = indexer._select_eager,indexer._select_fused,indexer._prefill_selector_supported
|
|
def eager_call(*args,**kwargs):
|
|
lane[0] = "eager"
|
|
return eager(*args,**kwargs)
|
|
def fused_call(*args,**kwargs):
|
|
lane[0] = "legacy_fused"
|
|
return fused(*args,**kwargs)
|
|
def prefill_call(*args,**kwargs):
|
|
result = prefill(*args,**kwargs)
|
|
if result: lane[0] = "metal_prefill"
|
|
return result
|
|
indexer._select_eager,indexer._select_fused,indexer._prefill_selector_supported = eager_call,fused_call,prefill_call
|
|
original_compiled = indexer._call_rows_compiled
|
|
def compiled_call(*args,**kwargs):
|
|
lane[0] = "compiled"
|
|
return original_compiled(*args,**kwargs)
|
|
indexer._call_rows_compiled = compiled_call
|
|
for step,rows in enumerate(steps):
|
|
pos = cache.offset
|
|
hidden = pattern(rows*2560,77+step).reshape(1,rows,2560)
|
|
stride = 704 if supplied else 640
|
|
qk = pattern(rows*stride,121+step).reshape(1,rows,stride)[...,:640] if supplied else None
|
|
lane[0] = "update_only"
|
|
with attention_phase("prefill" if options["phase_prefill"] else "verify"):
|
|
result = indexer(hidden,pos,cache,qk)
|
|
tail = None
|
|
if result is None: kind,values = "none",[]
|
|
elif isinstance(result,tuple):
|
|
kind = result[0]
|
|
if kind == "flash": values,tail = [result[1]],result[2]
|
|
else: values = list(result[1:])
|
|
else: kind,values = ("dense" if result.dtype == mx.bool_ else "decode_tokens"),[result]
|
|
states = {"raw":cache.raw_keys,"pooled":cache.pooled,"mirror":cache.pooled_f32_t}
|
|
mx.eval(*values,*(a for a in states.values() if a is not None))
|
|
state_hashes = {key:None if a is None else hashlib.sha256(np.asarray(a.astype(mx.float32)).astype('<f4').tobytes()).hexdigest() for key,a in states.items()}
|
|
tag = "qsa_indexer_compiled" if compiled else "qsa_indexer"
|
|
graph_counts = {}
|
|
if compiled:
|
|
core = indexer._compiled_indexer_core
|
|
graph_counts["graph_counts"] = [0,0,0] if core is None else [
|
|
core.stats["compiled_calls"],core.stats["traces"],len(core._compiled)]
|
|
emit(f"{tag}_s{scenario}_qk{int(supplied)}_step{step}",values,
|
|
scenario=scenario,step=step,prefix=prefix,rows=rows,pos=pos,total=pos+rows,
|
|
supplied=supplied,stride=stride,options=options,kind=kind,lane=lane[0],tail=tail,
|
|
raw_cap=cache.raw_keys.shape[1],pooled_cap=0 if cache.pooled is None else cache.pooled.shape[1],
|
|
pooled_len=cache.pooled_len,mirror_cap=0 if cache.pooled_f32_t is None else cache.pooled_f32_t.shape[3],
|
|
state_sha256=state_hashes,frequencies=np.asarray(indexer._inv_freq).tolist(),**graph_counts)
|
|
assert cache.offset == pos
|
|
kv = pattern(rows*512,140+step).reshape(1,rows,2,256).transpose(0,2,1,3)
|
|
cache.kv.update_and_fetch(kv,kv)
|
|
mx.eval(*cache.kv.state)
|
|
finally:
|
|
for key,value in saved.items():
|
|
if value is None: os.environ.pop(key,None)
|
|
else: os.environ[key] = value
|
|
|
|
|
|
def main():
|
|
assert mx.__version__ == "0.32.2", mx.__version__
|
|
with (Path(mx.__file__).parent / "lib/mlx.metallib").open("rb") as library:
|
|
assert hashlib.file_digest(library, "sha256").hexdigest() == (
|
|
"dc59d1cceb1a5c7e578232e6e41e28e2c73c9463ac6dbc3886c3ee17ffc270ed")
|
|
reference = Path(hyper_connection.__file__).resolve().parents[2]
|
|
exporter = runpy.run_path(str(Path(__file__).with_name("mtplx-kernel-source.py")))
|
|
assert exporter["generate"](reference, Path(gated_delta.__file__)) == exporter["OUTPUT"].read_text()
|
|
assert exporter["qsa_select_source"](reference) == exporter["QSA_SELECT_OUTPUT"].read_text()
|
|
print(f"MTPLX {exporter['REVISION']}, runtime {mx.__version__}, {reference}", file=sys.stderr)
|
|
qsa_f32_fixtures()
|
|
qsa_eager_fixtures()
|
|
qsa_eager_tiled_fixtures()
|
|
qsa_indexer_route_fixtures()
|
|
for rows in (1, 4, 8):
|
|
x = pattern(rows * 10240, 1).reshape(rows, 10240)
|
|
gamma = pattern(10240, 2)
|
|
wd = pattern(320 * 10240, 3).reshape(320, 10240)
|
|
wu = pattern(10240 * 320, 4).reshape(10240, 320)
|
|
wi = pattern(4 * 10240, 5).reshape(4, 10240)
|
|
emit(f"hyper_inj_s{rows}", fused_hyper_read(x, gamma, wd, wu, wi))
|
|
mixed, _ = fused_hyper_read(x, gamma, wd, wu)
|
|
emit(f"hyper_mix_s{rows}", [mixed])
|
|
# Use the reference's own quantizer and module packing for v3.
|
|
from types import SimpleNamespace
|
|
from mtplx.kernels.hyper_connection_v3 import prepare_v3_pack
|
|
pack = prepare_v3_pack(SimpleNamespace(
|
|
input_mix_weight_down=SimpleNamespace(weight=wd),
|
|
input_mix_weight_up=SimpleNamespace(weight=wu),
|
|
block_inject_weight=SimpleNamespace(weight=wi)))
|
|
emit("hyper_v3", fused_hyper_read_v3(pattern(10240, 1), gamma, pack))
|
|
for rows in range(1, 7):
|
|
x = pattern(rows * 10240, 6).reshape(rows, 10240)
|
|
state = pattern(3 * 10240, 7).reshape(3, 10240)
|
|
weights = pattern(10240 * 4, 8).reshape(10240, 4)
|
|
outputs = (fused_gdn_conv_norm(x.reshape(-1), state, weights) if rows == 1
|
|
else fused_gdn_conv_norm_rows(x, state, weights))
|
|
emit(f"gdn_s{rows}", outputs)
|
|
# Carry both returned states, including nonzero FP32 recurrence, across steps.
|
|
delta = pattern(48 * 128 * 128, 15).astype(mx.float32).reshape(48, 128, 128)
|
|
z, a, b = pattern(6144, 9), pattern(48, 10), pattern(48, 11)
|
|
a_log, dt_bias, norm = pattern(48, 12), pattern(48, 13), pattern(128, 14)
|
|
for step in range(2):
|
|
y, state, delta = fused_gdn_step(
|
|
pattern(10240, 6 + step), z, a, b, state, weights,
|
|
a_log, dt_bias, delta, norm)
|
|
emit(f"gdn_step_{step}", [y, state, delta])
|
|
# Synthetic valid affine packs isolate this kernel's ABI from the quantizer.
|
|
words = (np.arange(2560 * 768, dtype=np.uint64) * 2654435761 + 12345).astype(np.uint32)
|
|
qw = mx.array(words).reshape(2560, 768)
|
|
for gs in (32, 64):
|
|
scales = pattern(2560 * 6144 // gs, 16).reshape(2560, 6144 // gs)
|
|
biases = pattern(2560 * 6144 // gs, 17).reshape(2560, 6144 // gs)
|
|
for tag, dtype in (("bf16", mx.bfloat16), ("f32", mx.float32)):
|
|
x = pattern(6144, 18).astype(dtype)
|
|
emit(f"gdn_out_gs{gs}_{tag}", [fused_gdn_out(
|
|
x, z, norm, qw, scales, biases, group_size=gs)])
|
|
for rows in (1, 4, 7, 32, 2048):
|
|
q = (pattern(rows * 2048, 19) / 16).reshape(1, rows, 16, 128)
|
|
k = (pattern(rows * 2048, 20) / 16).reshape(1, rows, 16, 128)
|
|
v = pattern(rows * 6144, 21).reshape(1, rows, 48, 128)
|
|
g = mx.full((1, rows, 48), 0.75, dtype=mx.float32)
|
|
beta = mx.full((1, rows, 48), 0.5, dtype=mx.bfloat16)
|
|
state = pattern(48 * 128 * 128, 15).astype(mx.float32).reshape(1, 48, 128, 128)
|
|
for masked in (False, True):
|
|
mask = mx.array((np.arange(rows) % 3) != 0).reshape(1, rows) if masked else None
|
|
emit(f"gated_delta_s{rows}_mask{int(masked)}", gated_delta.gated_delta_kernel(
|
|
q, k, v, g, beta, state, mask))
|
|
# Exercise real runtime selection, not individually chosen fast kernels.
|
|
shapes = [(1, 80, 64), (4, 80, 128), (1, 2560, 2560), (1, 324, 320),
|
|
(1, 81, 2560), *[(m, 640, 2560) for m in (2, 3, 4, 5, 6, 12, 24, 25)],
|
|
(13, 6144, 6144), (33, 48, 2560), (13, 16384, 320),
|
|
(65, 16385, 320), (13, 16384, 160), (2048, 2560, 2560),
|
|
(33, 16, 32768)]
|
|
for m, n, k in shapes:
|
|
for group in (32, 64):
|
|
if k % group:
|
|
continue
|
|
for bits in (4, 8):
|
|
words = (np.arange(n * k * bits // 32, dtype=np.uint64)
|
|
* 2654435761 + 12345).astype(np.uint32)
|
|
qw = mx.array(words).reshape(n, k * bits // 32)
|
|
qs = pattern(n * k // group, 23).reshape(n, k // group)
|
|
qb = pattern(n * k // group, 24).reshape(n, k // group)
|
|
x = pattern(m * k, 25).reshape(m, k)
|
|
y = mx.quantized_matmul(x, qw, qs, qb, transpose=True, group_size=group, bits=bits)
|
|
emit(f"affine_m{m}_n{n}_k{k}_g{group}_b{bits}", [y],
|
|
m=m, n=n, k=k, group=group, bits=bits)
|
|
# Call the actual model's live-state branch, not a hand-written stand-in.
|
|
import mlx.nn as nn
|
|
from mlx_lm.models.cache import ArraysCache
|
|
from mtplx.models.qwen4_exp import GatedDeltaNet, TextArgs, _FusedGDNInProj
|
|
os.environ["MTPLX_FUSED_GDN_STEP"] = "1"
|
|
layer = GatedDeltaNet(TextArgs())
|
|
def pack(n, k, salt):
|
|
words = (np.arange(n * k // 8, dtype=np.uint64) * 2654435761 + 12345).astype(np.uint32)
|
|
return (mx.array(words).reshape(n,k//8),
|
|
(pattern(n*k//64,salt)/64).reshape(n,k//64),
|
|
(pattern(n*k//64,salt+1)/64).reshape(n,k//64))
|
|
layer.in_proj_fused = _FusedGDNInProj(*pack(16480,2560,26),64,4,"affine",[10240,16384,16432])
|
|
layer.out_proj = nn.QuantizedLinear(6144,2560,bias=False,group_size=64,bits=4)
|
|
layer.out_proj.weight, layer.out_proj.scales, layer.out_proj.biases = pack(2560,6144,28)
|
|
layer.conv1d.weight = pattern(10240*4,8).reshape(10240,4,1)
|
|
layer.A_log, layer.dt_bias, layer.norm.weight = pattern(48,12), pattern(48,13), pattern(128,14)
|
|
layer.eval()
|
|
cache = ArraysCache(2)
|
|
cache[0] = pattern(3*10240,7).reshape(1,3,10240)
|
|
cache[1] = pattern(48*128*128,15).astype(mx.float32).reshape(1,48,128,128)
|
|
for step in range(2):
|
|
assert layer._fused_step_applies(1,1,None,cache)
|
|
output = layer(pattern(2560,30+step).reshape(1,1,2560),cache=cache)
|
|
emit(f"gdn_chain_{step}",[output,cache[0],cache[1]])
|
|
# Full staged model forward, including captured recurrence inputs and
|
|
# carried states. Keep actual module selection and BF16 boundaries.
|
|
from mtplx.models.qwen4_exp import _VERIFY_CAPTURE
|
|
os.environ["MTPLX_FUSED_GDN_STEP"] = "0"
|
|
staged_cases = [(rows,initial,masked,fused,None)
|
|
for rows in (1,2,4,6,7,32,2048)
|
|
for initial in (False,True) for masked in (False,True)
|
|
for fused in (False,True)]
|
|
staged_cases += [(rows,True,False,True,length)
|
|
for rows in (1,4,7) for length in (-3,rows-1,rows+3)]
|
|
for rows,initial,masked,fused,length in staged_cases:
|
|
for flag in ("MTPLX_FUSED_GDN_CONVNORM","MTPLX_FUSED_CONVNORM_VERIFY","MTPLX_FUSED_GDN_OUT"):
|
|
os.environ[flag] = str(int(fused))
|
|
cache = ArraysCache(2)
|
|
if initial:
|
|
cache[0] = pattern(3*10240,7).reshape(1,3,10240)
|
|
cache[1] = pattern(48*128*128,15).astype(mx.float32).reshape(1,48,128,128)
|
|
mask = mx.array((np.arange(rows)%3)!=0).reshape(1,rows) if masked else None
|
|
for step in range(2):
|
|
# The host owns cache lengths; exercise clipping at either end.
|
|
cache.lengths = mx.array([length],dtype=mx.int32) if length is not None else None
|
|
used_fused = layer._fused_conv_norm_applies(1,rows,mask,cache) or layer._fused_conv_norm_rows_applies(1,rows,mask,cache)
|
|
capture = _VERIFY_CAPTURE.set(True)
|
|
try:
|
|
output = layer(pattern(rows*2560,30+step).reshape(1,rows,2560),mask=mask,cache=cache)
|
|
finally:
|
|
_VERIFY_CAPTURE.reset(capture)
|
|
_,q,k,v,_,_ = cache._mtplx_verify_rows
|
|
emit(f"gdn_staged_r{rows}_initial{int(initial)}_mask{int(masked)}_fused{int(fused)}_len{length}_step{step}",
|
|
[output,cache[0],cache[1],q,k,v],rows=rows,initial=initial,masked=masked,fused=fused,
|
|
length=length,step=step,used_fused=used_fused)
|
|
# Unified model branch selection, including non-sanitized input modules.
|
|
forward_cases = [(1,initial,mode) for initial in range(4) for mode in range(6)]
|
|
forward_cases += [(rows,initial,mode) for rows in (2,6,7,33,2048)
|
|
for initial in (0,1) for mode in (0,1)]
|
|
for layout in ("fused","separate","mixed"):
|
|
if layout != "fused":
|
|
if hasattr(layer,"in_proj_fused"):
|
|
del layer.in_proj_fused
|
|
for index,(name,width) in enumerate(zip(("qkv","z","b","a"),(10240,6144,48,48))):
|
|
bits,group = ((4,64),(8,32),(8,64),(4,32))[index] if layout=="mixed" else (4,64)
|
|
proj = nn.QuantizedLinear(2560,width,bias=False,group_size=group,bits=bits)
|
|
words = (np.arange(width*2560*bits//32,dtype=np.uint64)*2654435761+12345).astype(np.uint32)
|
|
proj.weight = mx.array(words).reshape(width,2560*bits//32)
|
|
proj.scales = (pattern(width*2560//group,64+index*2)/64).reshape(width,2560//group)
|
|
proj.biases = (pattern(width*2560//group,65+index*2)/64).reshape(width,2560//group)
|
|
setattr(layer,"in_proj_"+name,proj)
|
|
for rows,initial,mode in forward_cases:
|
|
enabled = mode != 0
|
|
for flag in ("MTPLX_FUSED_GDN_STEP","MTPLX_FUSED_GDN_CONVNORM","MTPLX_FUSED_CONVNORM_VERIFY","MTPLX_FUSED_GDN_OUT"):
|
|
os.environ[flag] = str(int(enabled and (mode!=5 or flag=="MTPLX_FUSED_GDN_STEP")))
|
|
cache = ArraysCache(2)
|
|
if initial in (1,3): cache[0] = pattern(3*10240,7).reshape(1,3,10240)
|
|
if initial in (1,2): cache[1] = pattern(48*128*128,15).astype(mx.float32).reshape(1,48,128,128)
|
|
if mode == 4: cache.lengths = mx.array([1],dtype=mx.int32)
|
|
mask = mx.array((np.arange(rows)%3)!=0).reshape(1,rows) if mode==3 else None
|
|
for step in range(2):
|
|
capture = _VERIFY_CAPTURE.set(mode==2)
|
|
try:
|
|
used_step = layer._fused_step_applies(1,rows,mask,cache)
|
|
used_conv = not used_step and (layer._fused_conv_norm_applies(1,rows,mask,cache) or layer._fused_conv_norm_rows_applies(1,rows,mask,cache))
|
|
output = layer(pattern(rows*2560,30+step).reshape(1,rows,2560),mask=mask,cache=cache)
|
|
finally:
|
|
_VERIFY_CAPTURE.reset(capture)
|
|
emit(f"gdn_forward_{layout}_r{rows}_initial{initial}_mode{mode}_step{step}",
|
|
[output,cache[0],cache[1]],layout=layout,rows=rows,initial=initial,mode=mode,
|
|
step=step,used_step=used_step,used_conv=used_conv)
|
|
from mtplx.models.qwen4_exp import QSAIndexer
|
|
from mtplx.kernels.qsa_indexer_prepare import qsa_indexer_prepare_queries_metal, qsa_indexer_pool_keys_metal
|
|
indexer = QSAIndexer(TextArgs())
|
|
indexer.q_layernorm.weight = pattern(128,72)
|
|
indexer.k_layernorm.weight = pattern(128,73)
|
|
frequencies = np.asarray(indexer._inv_freq).tolist()
|
|
assert len(frequencies)==32 and indexer._rope_attention_scaling==1.0
|
|
for rows in (1,4,7,32,2048):
|
|
for pos in (0,32767,262141):
|
|
for stride,dimstride in ((512,1),(640,1),(1280,2)):
|
|
raw = pattern(rows*stride,71).reshape(1,rows,stride)[...,:512*dimstride:dimstride].reshape(1,rows,4,128)
|
|
q = qsa_indexer_prepare_queries_metal(raw,indexer.q_layernorm.weight,indexer._inv_freq,
|
|
pos_start=pos,eps=1e-6)
|
|
emit(f"qsa_prepare_r{rows}_stride{stride}_dim{dimstride}_pos{pos}",[q],
|
|
rows=rows,stride=stride,dimstride=dimstride,pos=pos,frequencies=frequencies)
|
|
q = indexer._prepare_queries_eager(raw,pos)
|
|
emit(f"qsa_prepare_eager_r{rows}_stride{stride}_dim{dimstride}_pos{pos}",[q],
|
|
rows=rows,stride=stride,dimstride=dimstride,pos=pos,frequencies=frequencies,eager=True)
|
|
for blocks in (1,2,7,512):
|
|
for start in (0,8191,65535):
|
|
for stride,dimstride in ((128,1),(640,1),(1280,2)):
|
|
raw = pattern(blocks*4*stride,74).reshape(1,blocks*4,stride)[...,:128*dimstride:dimstride]
|
|
pooled = qsa_indexer_pool_keys_metal(raw,indexer.k_layernorm.weight,indexer._inv_freq,
|
|
block_start=start,compress_ratio=4,eps=1e-6)
|
|
emit(f"qsa_pool_n{blocks}_stride{stride}_dim{dimstride}_start{start}",[pooled],
|
|
blocks=blocks,stride=stride,dimstride=dimstride,start=start,frequencies=frequencies)
|
|
pooled = indexer._pool_keys_eager(raw,start,start+blocks)
|
|
emit(f"qsa_pool_eager_n{blocks}_stride{stride}_dim{dimstride}_start{start}",[pooled],
|
|
blocks=blocks,stride=stride,dimstride=dimstride,start=start,frequencies=frequencies,eager=True)
|
|
os.environ["MTPLX_FUSED_QSA_INDEXER"] = "1"
|
|
for bits,group in ((4,64),(8,64),(8,32)):
|
|
proj = nn.QuantizedLinear(2560,640,bias=False,group_size=group,bits=bits)
|
|
words = (np.arange(640*2560*bits//32,dtype=np.uint64)*2654435761+12345).astype(np.uint32)
|
|
proj.weight = mx.array(words).reshape(640,2560*bits//32)
|
|
proj.scales = (pattern(640*2560//group,75)/64).reshape(640,2560//group)
|
|
proj.biases = (pattern(640*2560//group,76)/64).reshape(640,2560//group)
|
|
indexer.index_qk_proj = proj
|
|
for rows in (4,8,32,2048):
|
|
for pos in (0,262140):
|
|
qk = indexer.index_qk_proj(pattern(rows*2560,77).reshape(1,rows,2560))
|
|
q,k = mx.split(qk,[512],axis=-1)
|
|
prepared = indexer._prepare_queries(q.reshape(1,rows,4,128),pos)
|
|
pooled = qsa_indexer_pool_keys_metal(k,indexer.k_layernorm.weight,indexer._inv_freq,
|
|
block_start=pos//4,compress_ratio=4,eps=1e-6)
|
|
emit(f"qsa_project_r{rows}_b{bits}_g{group}_pos{pos}",[prepared,pooled],
|
|
rows=rows,bits=bits,group=group,pos=pos,frequencies=frequencies)
|
|
prepared = indexer._prepare_queries_eager(q.reshape(1,rows,4,128),pos)
|
|
pooled = indexer._pool_keys_eager(k,pos//4,pos//4+rows//4)
|
|
emit(f"qsa_project_eager_r{rows}_b{bits}_g{group}_pos{pos}",[prepared,pooled],
|
|
rows=rows,bits=bits,group=group,pos=pos,frequencies=frequencies,eager=True)
|
|
# Actual QSACache + KVCache lifecycle, retaining array aliases across restore.
|
|
from mtplx.models.qwen4_exp import QSACache
|
|
cache_ops = [
|
|
("append",1),("append",3),("append",252),("save",0),("append",1),("append",767),
|
|
("save",1),("append",4),("trim",3),("append",7),("restore",1),("append",4),
|
|
("restore",0),("trim",5),("append",2),("reserve",(1600,600)),("append",1029),
|
|
("save",2),("trim",100000),("append",2),("restore",2),("append",3),
|
|
]
|
|
for scenario in (0,1,2,3):
|
|
os.environ["MTPLX_FUSED_QSA_INDEXER"] = "1" if scenario < 2 else "0"
|
|
cache = QSACache()
|
|
bank = {}
|
|
if scenario%2: cache.reserve_indexer_capacity(raw_capacity=300,pooled_capacity=300)
|
|
for index,(op,arg) in enumerate(cache_ops):
|
|
if op=="append":
|
|
stride=128 if index%2==0 else 640
|
|
raw=pattern(arg*stride,80+index).reshape(1,arg,stride)[...,:128]
|
|
keys=pattern(arg*512,110+index).reshape(1,arg,2,256).transpose(0,2,1,3)
|
|
values=pattern(arg*512,140+index).reshape(1,arg,2,256).transpose(0,2,1,3)
|
|
cache.write_raw(raw)
|
|
indexer._extend_pooled(cache,cache.offset+arg)
|
|
cache.kv.update_and_fetch(keys,values)
|
|
elif op=="reserve": cache.reserve_indexer_capacity(raw_capacity=arg[0],pooled_capacity=arg[1])
|
|
elif op=="trim": cache.trim(arg)
|
|
elif op=="save": bank[arg]=cache.state
|
|
elif op=="restore": cache.state=bank[arg]
|
|
# A mirror read after reserve/restore must seed every earlier block.
|
|
if cache.pooled_len: cache.pooled_f32_view(cache.pooled_len)
|
|
arrays=[cache.kv.keys,cache.kv.values,cache.raw_keys,cache.pooled,cache.pooled_f32_t]
|
|
outputs=[x for x in arrays if x is not None]
|
|
for saved in bank.values(): outputs.extend(x for x in saved if x is not None)
|
|
emit(f"qsa_cache_s{scenario}_op{index}",outputs,scenario=scenario,index=index,op=op,arg=arg,
|
|
offset=cache.offset,nbytes=cache.nbytes,raw_cap=cache.raw_keys.shape[1],kv_cap=cache.kv.keys.shape[2],
|
|
pooled_cap=0 if cache.pooled is None else cache.pooled.shape[1],pooled_len=cache.pooled_len,
|
|
mirror_cap=0 if cache.pooled_f32_t is None else cache.pooled_f32_t.shape[3],
|
|
present=[x is not None for x in arrays],bank=list(bank),frequencies=frequencies)
|
|
os.environ["MTPLX_FUSED_QSA_INDEXER"] = "1"
|
|
# Original fused QSA selector: all output modes, padded histories, tie order,
|
|
# BF16/F32 operands and the runtime TF32 setting. Also capture its actual
|
|
# header/body, independent of our Rust entry-point binding declaration.
|
|
from mtplx.kernels import qsa_indexer_select as selector
|
|
select_shapes = [(1,0,1,0,512),(4,0,7,8,512),(4,2045,2049,600,512),
|
|
(7,2097,2104,600,16),(33,4080,4113,1200,512),(1,8192,8193,4096,512)]
|
|
for case,(rows,pos,total,blocks,topk) in enumerate(select_shapes):
|
|
for variant in range(4):
|
|
f32 = [variant==2, variant>=2]
|
|
tf32 = variant!=3
|
|
dim_stride = 2 if variant%2 else 1
|
|
qstride=512*dim_stride
|
|
pstride=128*dim_stride
|
|
def operand(count,salt,fp32):
|
|
x=pattern(max(1,count),salt)
|
|
if fp32: x=x.astype(mx.float32)+0.00013
|
|
if variant==1: x=mx.zeros_like(x)
|
|
return x
|
|
q=operand(rows*qstride,90,f32[0]).reshape(1,rows,4,128*dim_stride)[...,::dim_stride]
|
|
pooled=operand(blocks*pstride,91,f32[1])[:blocks*pstride].reshape(1,blocks,128*dim_stride)[...,::dim_stride]
|
|
for mode in ("blocks","dense_mask","row_tokens"):
|
|
dense=(blocks+1)*4 if mode=="dense_mask" else 0
|
|
metal_kernel=mx.fast.metal_kernel
|
|
try:
|
|
mx.fast.metal_kernel=lambda **kw:kw
|
|
spec=selector._selector_kernel.__wrapped__(mode,4,128,blocks,topk,4,
|
|
1<<(max(256,topk)-1).bit_length(),
|
|
dense,tf32,q.dtype,pooled.dtype)
|
|
finally: mx.fast.metal_kernel=metal_kernel
|
|
tf_setting=selector._mlx_tf32_enabled
|
|
try:
|
|
selector._mlx_tf32_enabled=lambda:tf32
|
|
out=selector.qsa_indexer_select_metal(q,pooled,pos_start=pos,total_tokens=total,
|
|
logical_blocks=total//4,block_topk=topk,compress_ratio=4,mode=mode,
|
|
output_total_tokens=dense if mode=="dense_mask" else None)
|
|
finally: selector._mlx_tf32_enabled=tf_setting
|
|
emit(f"qsa_select_c{case}_v{variant}_{mode}",[out] if mode=="dense_mask" else out,
|
|
rows=rows,pos=pos,total=total,blocks=blocks,topk=topk,mode=mode,dense_width=dense,
|
|
input_f32=f32,tf32=tf32,dim_stride=dim_stride,zero=variant==1,
|
|
header_sha256=hashlib.sha256(spec["header"].encode()).hexdigest(),
|
|
body_sha256=hashlib.sha256(spec["source"].encode()).hexdigest())
|
|
# Force the actual 32MiB scratch boundary with a large reserved backing
|
|
# but a short logical prefix. This calls the host chunk/concat path itself.
|
|
q=pattern(129*512,90).reshape(1,129,4,128)
|
|
pooled=pattern(65536*128,91).reshape(1,65536,128)
|
|
for mode in ("blocks","dense_mask","row_tokens"):
|
|
out=indexer._select_fused(q,0,pooled,129//4,129,mode)
|
|
emit(f"qsa_chunk_{mode}",[out] if mode=="dense_mask" else out,
|
|
rows=129,pos=0,total=129,blocks=65536,topk=512,mode=mode,dense_width=0,
|
|
input_f32=[False,False],tf32=True,dim_stride=1,zero=False,chunked=True)
|
|
# Connected raw-write -> pool -> prepare -> selection -> KV advance.
|
|
cache=QSACache()
|
|
cache.write_raw(pattern(2048*128,88).reshape(1,2048,128))
|
|
indexer._extend_pooled(cache,2048)
|
|
kv=pattern(2048*512,89).reshape(1,2048,2,256).transpose(0,2,1,3)
|
|
cache.kv.update_and_fetch(kv,kv)
|
|
for step,rows in enumerate((1,4,7,2048)):
|
|
q=indexer._prepare_queries(pattern(rows*512,90+step).reshape(1,rows,4,128),cache.offset)
|
|
cache.write_raw(pattern(rows*128,100+step).reshape(1,rows,128))
|
|
total=cache.offset+rows
|
|
indexer._extend_pooled(cache,total)
|
|
for mode in ("blocks","dense_mask","row_tokens"):
|
|
if rows>=128:
|
|
from mtplx.kernels.qsa_indexer_prefill import qsa_indexer_prefill_metal
|
|
out=qsa_indexer_prefill_metal(q,cache.pooled[:,:total//4,:],pos_start=cache.offset,
|
|
total_tokens=total,block_topk=512,compress_ratio=4,logical_blocks=total//4,mode=mode)
|
|
else: out=indexer._select_fused(q,cache.offset,cache.pooled,total//4,total,mode)
|
|
emit(f"qsa_connected_s{step}_{mode}",[out] if mode=="dense_mask" else out,
|
|
rows=rows,step=step,mode=mode,frequencies=frequencies,**({"prefill":True} if rows>=128 else {}))
|
|
kv=pattern(rows*512,110+step).reshape(1,rows,2,256).transpose(0,2,1,3)
|
|
cache.kv.update_and_fetch(kv,kv)
|
|
from mtplx.kernels import qsa_indexer_prefill as prefill
|
|
for rows,blocks in ((1,1),(4,7),(15,31),(16,32),(17,33),(33,513),(2048,513)):
|
|
for ds in (1,2):
|
|
q=pattern(rows*512*ds,92).reshape(1,rows,4,128*ds)[...,::ds]
|
|
pooled=pattern(blocks*128*ds,93).reshape(1,blocks,128*ds)[...,::ds]
|
|
assert prefill.qsa_indexer_prefill_scores_mpp_supported(q,pooled)
|
|
emit(f"qsa_mpp_score_r{rows}_n{blocks}_ds{ds}",[prefill.qsa_indexer_prefill_scores_mpp(q,pooled)],
|
|
rows=rows,blocks=blocks,dim_stride=ds)
|
|
topk_shapes=((1,0,1,1,512),(4,0,7,8,512),(7,2045,2052,600,512),
|
|
(4,8192,8196,4096,512),(1,262142,262143,65536,512),
|
|
(4,252,256,64,16),(1,256,257,65,16),(1,8188,8189,2048,512),(1,8192,8193,2049,512))
|
|
for case,(rows,pos,total,blocks,topk) in enumerate(topk_shapes):
|
|
for variant in range(4):
|
|
ds=2 if variant in (1,3) else 1
|
|
stride=blocks*ds+(4 if variant else 0)
|
|
scores=pattern(rows*stride,94).astype(mx.float32)
|
|
if variant in (1,2): scores=mx.full((rows*stride,),float(variant-1),dtype=mx.float32)
|
|
scores=scores.reshape(rows,stride)[:,:blocks*ds:ds]
|
|
for mode in ("blocks","dense_mask","row_tokens"):
|
|
dense=(blocks+1)*4 if mode=="dense_mask" else 0
|
|
metal_kernel=mx.fast.metal_kernel
|
|
try:
|
|
mx.fast.metal_kernel=lambda **kw:kw
|
|
spec=prefill._prefill_topk_kernel.__wrapped__(mode,blocks,topk,4,
|
|
1<<(max(256,topk)-1).bit_length(),dense)
|
|
finally: mx.fast.metal_kernel=metal_kernel
|
|
out=prefill.qsa_indexer_prefill_topk_metal(scores,pos_start=pos,total_tokens=total,
|
|
block_topk=topk,compress_ratio=4,logical_blocks=total//4,mode=mode,
|
|
output_total_tokens=dense if mode=="dense_mask" else None)
|
|
emit(f"qsa_prefill_topk_c{case}_v{variant}_{mode}",[out] if mode=="dense_mask" else out,
|
|
rows=rows,pos=pos,total=total,blocks=blocks,topk=topk,mode=mode,dense_width=dense,
|
|
variant=variant,row_stride=stride,dim_stride=ds,
|
|
header_sha256=hashlib.sha256(spec["header"].encode()).hexdigest(),
|
|
body_sha256=hashlib.sha256(spec["source"].encode()).hexdigest())
|
|
# Producer-aware budget/planner, and the actual score -> topk -> concat graph.
|
|
for rows,blocks,budget in ((7,513,513*4*3),(65,600,600*4*48),(2048,513,128*1024*1024),
|
|
(3,65536,1),(519,65536,128*1024*1024)):
|
|
for mpp in (False,True):
|
|
chunk=prefill.qsa_indexer_prefill_score_chunk_rows(rows,4,blocks,budget,producer="mpp" if mpp else "mlx")
|
|
emit(f"qsa_prefill_plan_r{rows}_n{blocks}_b{budget}_mpp{int(mpp)}",[],rows=rows,blocks=blocks,budget=budget,mpp=mpp,chunk=chunk)
|
|
for case,(rows,blocks,budget) in enumerate(((7,513,513*4*3),(65,600,600*4*48),(2048,513,128*1024*1024),(3,65536,1))):
|
|
for ds in (1,2):
|
|
q=pattern(rows*512*ds,92).reshape(1,rows,4,128*ds)[...,::ds]
|
|
pooled=pattern(blocks*128*ds,93).reshape(1,blocks,128*ds)[...,::ds]
|
|
total=max(rows,blocks*4-1)
|
|
pos=total-rows
|
|
for mode in ("blocks","dense_mask","row_tokens"):
|
|
dense=total+4 if mode=="dense_mask" else 0
|
|
out=prefill.qsa_indexer_prefill_metal(q,pooled,pos_start=pos,total_tokens=total,
|
|
block_topk=512,compress_ratio=4,logical_blocks=total//4,mode=mode,
|
|
output_total_tokens=dense if mode=="dense_mask" else None,score_workspace_bytes=budget)
|
|
emit(f"qsa_prefill_chain_c{case}_ds{ds}_{mode}",[out] if mode=="dense_mask" else out,
|
|
rows=rows,pos=pos,total=total,blocks=blocks,topk=512,mode=mode,dense_width=dense,
|
|
budget=budget,dim_stride=ds)
|
|
# Actual gather selection, including the distinct sorted-RHS threshold.
|
|
gather_shapes = [
|
|
(3, 1, 640, 2560, 8, 40, False),
|
|
(3, 1, 2560, 640, 8, 40, False),
|
|
(3, 1, 81, 320, 8, 7, False),
|
|
(3, 33, 80, 128, 8, 7, False),
|
|
(3, 33, 81, 160, 8, 7, False),
|
|
(31, 1, 80, 128, 8, 31, True),
|
|
(32, 1, 80, 128, 8, 32, True),
|
|
(33, 1, 81, 160, 8, 33, True),
|
|
(512, 1, 1280, 2560, 8, 512, True),
|
|
(2047, 1, 80, 128, 512, 2047, True),
|
|
(2048, 1, 80, 128, 512, 2048, True),
|
|
]
|
|
for batches, m, n, k, experts, routes, sorted_indices in gather_shapes:
|
|
for group in (32, 64):
|
|
if k % group:
|
|
continue
|
|
for bits in (4, 8):
|
|
words = (np.arange(experts*n*k*bits//32, dtype=np.uint64)
|
|
* 2654435761 + 12345).astype(np.uint32)
|
|
def padded_view(value, shape, extra):
|
|
storage = mx.concatenate([value,mx.zeros((extra,),dtype=value.dtype)])
|
|
mx.eval(storage)
|
|
return storage[:value.size].reshape(shape)
|
|
padding = k if sorted_indices and k % 64 else 0
|
|
qw = padded_view(mx.array(words),(experts,n,k*bits//32),padding*bits//32*64)
|
|
qs = padded_view(pattern(experts*n*k//group,33),(experts,n,k//group),padding//group*64)
|
|
qb = padded_view(pattern(experts*n*k//group,34),(experts,n,k//group),padding//group*64)
|
|
# The pinned RHS NAX shader reads a full 64-wide final tile
|
|
# even for K=160. Keep identical physical padding in both
|
|
# runners; this synthetic shape is not an installed MoE shape.
|
|
x = padded_view(pattern(batches*m*k,32),(batches,m,k),padding*64)
|
|
rhs = (np.arange(routes,dtype=np.uint32)*5+2)%experts
|
|
if sorted_indices:
|
|
rhs.sort()
|
|
lhs = None if sorted_indices else mx.array(
|
|
(np.arange(routes,dtype=np.uint32)*2+1)%batches)
|
|
y = mx.gather_qmm(x,qw,qs,qb,lhs_indices=lhs,rhs_indices=mx.array(rhs),
|
|
transpose=True,group_size=group,bits=bits,
|
|
sorted_indices=sorted_indices)
|
|
emit(f"gather_x{batches}_m{m}_n{n}_k{k}_e{experts}_r{routes}_s{int(sorted_indices)}_g{group}_b{bits}",
|
|
[y],batches=batches,m=m,n=n,k=k,experts=experts,routes=routes,
|
|
sorted=sorted_indices,group=group,bits=bits,padding=padding)
|
|
# Routing uses complete last-axis sort; gather_sort sorts the returned order
|
|
# again for its inverse. Check tied values, tile boundaries and long merges.
|
|
shapes = [(rows,width) for rows in (1,3) for width in (
|
|
1,10,63,64,65,128,129,256,257,512,513,1024,1025,2048,2049,
|
|
4097,20480,65535,65536,65537)] + [(2048,512)]
|
|
for rows,width in shapes:
|
|
for tag,dtype in (("u32",mx.uint32),("bf16",mx.bfloat16),("f32",mx.float32)):
|
|
values = (np.arange(rows*width,dtype=np.uint32)*17+35)%512
|
|
x = mx.array(values,dtype=dtype).reshape(rows,width)
|
|
order = mx.argsort(x,axis=-1)
|
|
inverse = mx.argsort(order,axis=-1)
|
|
emit(f"argsort_{tag}_r{rows}_w{width}",[order,inverse],rows=rows,width=width,dtype=tag)
|
|
from mlx_lm.models.switch_layers import _gather_sort, _scatter_unsort
|
|
for tokens,top_k,width in ((1,10,7),(6,10,8),(7,10,9),(7,10,2560),
|
|
(205,10,9),(2048,10,2560),(6554,10,9)):
|
|
x = pattern(tokens*width,36).reshape(1,tokens,1,1,width)
|
|
indices = mx.array((np.arange(tokens*top_k,dtype=np.uint32)*17+35)%512).reshape(1,tokens,top_k)
|
|
sorted_x, sorted_indices, inverse = _gather_sort(x,indices)
|
|
unsorted = _scatter_unsort(sorted_x,inverse,indices.shape)
|
|
emit(f"gather_sort_t{tokens}_top{top_k}_w{width}",
|
|
[sorted_x,sorted_indices,inverse,unsorted],tokens=tokens,top_k=top_k,width=width)
|
|
for rows,intermediate in ((1,64),(1,640),(1,65534),(1,65536),(1,65538),
|
|
(2,640),(10,640),(60,640),(70,640),(20480,640)):
|
|
x = pattern(rows*intermediate*2,37).reshape(rows,intermediate*2)
|
|
gate,up = mx.split(x,2,axis=-1)
|
|
activated = nn.silu(gate)
|
|
emit(f"gate_up_r{rows}_i{intermediate}",[activated,activated*up],rows=rows,intermediate=intermediate)
|
|
# Exact SparseMoeBlock routing operations, retaining the strided top-k view.
|
|
for rows,experts in ((1,512),(7,512),(32,512),(2048,512),
|
|
(1,16),(7,16),(3,4096),(3,4097)):
|
|
for normalize in (False,True):
|
|
logits = pattern(rows*experts,43).reshape(rows,experts)
|
|
gates = mx.softmax(logits,axis=-1,precise=True)
|
|
indices = mx.argpartition(gates,kth=-10,axis=-1)[...,-10:]
|
|
scores = mx.take_along_axis(gates,indices,axis=-1)
|
|
if normalize:
|
|
scores = scores/scores.sum(axis=-1,keepdims=True)
|
|
emit(f"router_r{rows}_e{experts}_norm{int(normalize)}",
|
|
[gates,indices,scores],rows=rows,experts=experts,normalize=normalize)
|
|
from mlx_lm.models.switch_layers import QuantizedSwitchLinear
|
|
from mtplx.models.qwen4_exp import _FusedGateUpSwitchGLU
|
|
def expert_pack(experts,n,k,group,salt):
|
|
words = (np.arange(experts*n*k//8,dtype=np.uint64)*2654435761+12345).astype(np.uint32)
|
|
return (mx.array(words).reshape(experts,n,k//8),
|
|
(pattern(experts*n*k//group,salt)/64).reshape(experts,n,k//group),
|
|
(pattern(experts*n*k//group,salt+1)/64).reshape(experts,n,k//group))
|
|
for tokens,experts,k,intermediate in (
|
|
(1,16,2560,640),(6,16,2560,640),(7,16,2560,640),
|
|
(205,16,2560,640),(2048,16,2560,640),(7,512,128,64),(205,512,128,64)):
|
|
for gu_group,down_group in ((32,32),(64,64),(32,64)):
|
|
# The reference module properties derive dimensions from loaded weights.
|
|
# Small constructor placeholders avoid unrelated random full-pack work.
|
|
down = QuantizedSwitchLinear(64,64,1,bias=False,group_size=down_group,bits=4)
|
|
down.weight,down.scales,down.biases = expert_pack(experts,k,intermediate,down_group,41)
|
|
layer = _FusedGateUpSwitchGLU(down,*expert_pack(experts,2*intermediate,k,gu_group,39),gu_group,4,"affine")
|
|
layer.eval()
|
|
x = pattern(tokens*k,38).reshape(1,tokens,k)
|
|
idx = mx.array((np.arange(tokens*10,dtype=np.uint32)*17+35)%experts).reshape(1,tokens,10)
|
|
emit(f"switch_t{tokens}_e{experts}_k{k}_i{intermediate}_g{gu_group}_d{down_group}",
|
|
[layer(x,idx)],tokens=tokens,experts=experts,k=k,intermediate=intermediate,
|
|
gu_group=gu_group,down_group=down_group)
|
|
# Call the complete actual stock SparseMoeBlock, including router projection,
|
|
# expert weighting and shared expert. Opt-in fused branches get separate receipts.
|
|
from mtplx.models.qwen4_exp import SparseMoeBlock, _FusedGateUpMLP
|
|
from mlx_lm.models.switch_layers import SwitchGLU
|
|
from mlx_lm.models.qwen3_next import Qwen3NextMLP
|
|
os.environ["MTPLX_FUSED_MOE_DECODE"] = "0"
|
|
os.environ["MTPLX_FUSED_MOE_VERIFY"] = "0"
|
|
linear = quantized_linear_fixture
|
|
moe_shapes = (
|
|
(1,16,2560,640),(4,16,2560,640),(6,16,2560,640),(7,16,2560,640),
|
|
(25,16,2560,640),(205,16,2560,640),(2048,16,2560,640),
|
|
(7,512,128,64),(205,512,128,64),
|
|
(1,512,2560,640),(7,512,2560,640),(2048,512,2560,640))
|
|
moe_cases = [(shape,((64,64),(32,64)),False,(False,False)) for shape in moe_shapes]
|
|
moe_cases += [((tokens,experts,2560,640),((32,32),(32,64),(64,32),(64,64)),True,(False,False))
|
|
for tokens,experts in ((1,16),(2,16),(3,16),(4,16),(5,16),(7,16),(1,512),(4,512))]
|
|
moe_cases += [((tokens,experts,2560,640),((64,64),(32,64)),True,separate)
|
|
for tokens,experts in ((1,16),(4,16),(6,16),(7,16),(25,16),(2048,16),(1,512),(7,512),(2048,512))
|
|
for separate in ((True,False),(False,True),(True,True))]
|
|
for (tokens,experts,k,intermediate),groups,fused,(separate_experts,separate_shared) in moe_cases:
|
|
os.environ["MTPLX_FUSED_MOE_DECODE"] = str(int(fused))
|
|
os.environ["MTPLX_FUSED_MOE_VERIFY"] = str(int(fused))
|
|
for gu_group,down_group in groups:
|
|
layer = SparseMoeBlock(SimpleNamespace(hidden_size=64,moe_intermediate_size=64,
|
|
shared_expert_intermediate_size=64,norm_topk_prob=True,num_experts=1,num_experts_per_tok=10))
|
|
down = QuantizedSwitchLinear(64,64,1,bias=False,group_size=down_group,bits=4)
|
|
down.weight,down.scales,down.biases = expert_pack(experts,k,intermediate,down_group,41)
|
|
if separate_experts:
|
|
sw = SwitchGLU(64,64,1)
|
|
for name,group,salt in (("gate_proj",gu_group,39),("up_proj",down_group,45)):
|
|
proj = QuantizedSwitchLinear(64,64,1,bias=False,group_size=group,bits=4)
|
|
proj.weight,proj.scales,proj.biases = expert_pack(experts,intermediate,k,group,salt)
|
|
setattr(sw,name,proj)
|
|
sw.down_proj = down
|
|
layer.switch_mlp = sw
|
|
else:
|
|
layer.switch_mlp = _FusedGateUpSwitchGLU(down,*expert_pack(experts,2*intermediate,k,gu_group,39),gu_group,4,"affine")
|
|
if separate_shared:
|
|
shared = Qwen3NextMLP(64,64)
|
|
shared.gate_proj = linear(intermediate,k,4,gu_group,46)
|
|
shared.up_proj = linear(intermediate,k,4,down_group,54)
|
|
shared.down_proj = linear(k,intermediate,4,down_group,48)
|
|
layer.shared_expert = shared
|
|
else:
|
|
shared_gu = linear(2*intermediate,k,4,gu_group,46)
|
|
layer.shared_expert = _FusedGateUpMLP(linear(k,intermediate,4,down_group,48),
|
|
shared_gu.weight,shared_gu.scales,shared_gu.biases,gu_group,4,"affine")
|
|
layer.gate = linear(experts,k,8,64,50)
|
|
layer.shared_expert_gate = linear(1,k,4,64,52)
|
|
layer.num_experts = experts
|
|
layer.eval()
|
|
x = pattern(tokens*k,44).reshape(1,tokens,k)
|
|
for normalize in (False,True):
|
|
layer.norm_topk_prob = normalize
|
|
extra = {"fused":True} if fused else {}
|
|
tag = "moe_fused" if fused else "moe"
|
|
if separate_experts or separate_shared:
|
|
extra.update(separate_experts=separate_experts,separate_shared=separate_shared)
|
|
tag = f"moe_separate_e{int(separate_experts)}_s{int(separate_shared)}"
|
|
emit(f"{tag}_t{tokens}_e{experts}_k{k}_i{intermediate}_g{gu_group}_d{down_group}_norm{int(normalize)}",
|
|
[layer(x)],tokens=tokens,experts=experts,k=k,intermediate=intermediate,
|
|
gu_group=gu_group,down_group=down_group,normalize=normalize,**extra)
|
|
from mtplx.models.qwen4_exp import GroupedRMSNorm, SigmoidRMSNormGated
|
|
for rows,width,group in ((1,10240,2560),(4,10240,2560),(2048,10240,2560),
|
|
(1,65534,2),(1,65536,128),(1,65538,2),(3,32768,8192),
|
|
(1,256,128),(7,128,128),(3,16384,4096)):
|
|
norm = GroupedRMSNorm(width,group,eps=1e-6)
|
|
norm.weight = pattern(width,56)
|
|
x = pattern(rows*width,55).reshape(1,rows,width)
|
|
emit(f"groupnorm_r{rows}_w{width}_g{group}",[norm(x)],rows=rows,width=width,group=group)
|
|
for tokens in (1,4,7,11,2048):
|
|
for mode in ("none","contiguous","projected"):
|
|
norm = SigmoidRMSNormGated(128,eps=1e-6)
|
|
norm.weight = pattern(128,56)
|
|
x = pattern(tokens*6144,55).reshape(1,tokens,48,128)
|
|
gate = None
|
|
if mode == "contiguous":
|
|
gate = pattern(tokens*6144,57).reshape(1,tokens,48,128)
|
|
elif mode == "projected":
|
|
gate = pattern(tokens*16480,57).reshape(1,tokens,16480)[...,10240:16384].reshape(1,tokens,48,128)
|
|
emit(f"gatednorm_t{tokens}_{mode}",[norm(x,gate)],tokens=tokens,mode=mode)
|
|
for rows in (1,4,7,32,2048):
|
|
for stride in (48,16480):
|
|
for initial in (False,True):
|
|
for masked in (False,True):
|
|
state = pattern(48*128*128,15).astype(mx.float32).reshape(1,48,128,128) if initial else None
|
|
a_log,dt_bias = pattern(48,62),pattern(48,63)
|
|
mask = mx.array((np.arange(rows)%3)!=0).reshape(1,rows) if masked else None
|
|
for step in range(2):
|
|
q = (pattern(rows*2048,58+step)/16).reshape(1,rows,16,128)
|
|
k = (pattern(rows*2048,59+step)/16).reshape(1,rows,16,128)
|
|
v = pattern(rows*6144,60+step).reshape(1,rows,48,128)
|
|
if stride==48:
|
|
a,b = pattern(rows*48,60+step).reshape(1,rows,48),pattern(rows*48,61+step).reshape(1,rows,48)
|
|
else:
|
|
projected = pattern(rows*16480,60+step).reshape(1,rows,16480)
|
|
a,b = projected[...,16432:],projected[...,16384:16432]
|
|
y,state = gated_delta.gated_delta_update(q,k,v,a,b,a_log,dt_bias,state,mask)
|
|
emit(f"delta_update_r{rows}_stride{stride}_initial{int(initial)}_mask{int(masked)}_step{step}",
|
|
[y,state,gated_delta.compute_g(a_log,a,dt_bias),mx.sigmoid(b)],
|
|
rows=rows,stride=stride,initial=initial,masked=masked,step=step)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|