Port Qwen row-map SwiGLU expert gathers

tools/issue112-extract-metal.py

@@ -0,0 +1,19 @@
+"""Regenerate only the vendored GPU shader closure from the pinned oracle."""
+import subprocess,sys
+from pathlib import Path
+checkout=Path(sys.argv[1])
+assert subprocess.check_output(['git','-C',str(checkout),'rev-parse','HEAD'],text=True).strip()=='9882703f3105363ddc37eca9f97aa09a1d387112'
+sys.path.insert(0,str(checkout))
+import mlx.core as mx
+from mtplx.kernels import moe_sorted_gather_nax as k
+assert mx.__version__=='0.32.2'
+root=k._include_root();seen=set();out=[]
+for rel in (k._PREAMBLE_HEADER,*k._KERNEL_HEADERS):k._inline(root,rel,seen,out)
+for rel,name in [('mlx/backend/metal/kernels/unary_ops.h','Sigmoid'),('mlx/backend/metal/kernels/binary_ops.h','Multiply')]:
+ functor=k._functor(root,rel,name)
+ assert functor is not None
+ out.append(functor)
+metal=Path(__file__).resolve().parents[1]/'metal'
+preamble='// MLX 0.32.2 GPU support, Copyright Apple Inc., MIT; see MLX-LICENSE.txt.\n// Flattened only the GPU header closure used by MTPLX 9882703 row-map gather.\n'
+(metal/'mtplx-moe-sorted-header.metal').write_text(preamble+'\n'.join(out)+'\nusing namespace mlx::steel;\n')
+(metal/'mtplx-moe-sorted-body.metal').write_text('// MTPLX 9882703, MIT; see MTPLX-LICENSE.txt.\n'+k._SOURCE)

tools/issue112-installed-reference.py

@@ -0,0 +1,38 @@
+"""Installed gate/up row-map oracle, using the identical native fixture inputs."""
+import hashlib,json,subprocess,sys
+from pathlib import Path
+checkout,root,native=map(Path,sys.argv[1:4])
+assert subprocess.check_output(['git','-C',str(checkout),'rev-parse','HEAD'],text=True).strip()=='9882703f3105363ddc37eca9f97aa09a1d387112'
+sys.path.insert(0,str(checkout))
+import mlx.core as mx
+import numpy as np
+from mtplx.kernels import moe_sorted_gather_nax as kernel
+assert mx.__version__=='0.32.2'
+index=json.loads((root/'model.safetensors.index.json').read_text())['weight_map']
+base='language_model.model.layers.0.mlp.switch_mlp.'
+weights=mx.load(str(root/index[base+'gate_proj.weight']))
+def parts(proj):return [weights[base+proj+'.'+part] for part in ['weight','scales','biases']]
+gate=parts('gate_proj');up=parts('up_proj')
+gu=[mx.concatenate(a,axis=1) for a in zip(gate,up)]
+mx.eval(*gu,*gate,*up)
+def digest(a):
+ mx.eval(a)
+ return hashlib.sha256(np.asarray(a.view(mx.uint16)).tobytes()).hexdigest()
+count=0
+for line in native.read_text().splitlines():
+ if '"case":"installed_gate_up"' not in line:continue
+ case=json.loads(line[line.index('{'):]);count+=1
+ tokens,rows=case['tokens'],case['rows']
+ v=((np.arange(tokens*2560,dtype=np.uint32)*17+7*13)%257).astype(np.int32)-128
+ x=mx.array((v.astype(np.float32)/128).reshape(tokens,1,2560)).astype(mx.bfloat16)
+ row_map=mx.arange(rows,dtype=mx.uint32)%tokens
+ ids=mx.array(np.arange(rows,dtype=np.uint32)*512//rows)
+ g,u=(gate,up) if case['separate'] else (gu,None)
+ kw=dict(group_size=64,bits=4,swiglu=True,up=u)
+ y=kernel.gather_rows_qmm(x,row_map,*g,ids,**kw)
+ stock=kernel._stock(x,row_map,*g,ids,**kw)
+ assert y is not None and digest(y)==digest(stock)==case['sha256'],case
+ print(json.dumps(dict(case='installed_oracle',native=case,exact=True,kernel=kernel.stats())),flush=True)
+ del x,row_map,ids,y,stock
+ mx.clear_cache()
+assert count==12,count

tools/issue112-model-reference.py

@@ -0,0 +1,58 @@
+"""Pinned full-model cold/warm prefill, including explicit hidden-state capture."""
+import argparse,hashlib,json,os,subprocess,sys
+from pathlib import Path
+parser=argparse.ArgumentParser(description=__doc__)
+for name in ['checkout','root','native','out']:parser.add_argument(name,type=Path)
+parser.add_argument('--hidden',action='store_true')
+a=parser.parse_args()
+assert subprocess.check_output(['git','-C',str(a.checkout),'rev-parse','HEAD'],text=True).strip()=='9882703f3105363ddc37eca9f97aa09a1d387112'
+sys.path.insert(0,str(a.checkout))
+from mtplx.server.openai import _server_runtime_env_overrides
+os.environ.update(_server_runtime_env_overrides(argparse.Namespace(model=str(a.root),generation_mode='mtp' if a.hidden else 'ar'),None))
+os.environ['MTPLX_SUSTAINED_PREFILL']='1'
+import mlx.core as mx
+import numpy as np
+from mlx_lm.utils import load_model
+from mtplx.models import qwen4_exp
+from mtplx.runtime import MTPLXRuntime,_load_tokenizer_resilient
+from mtplx.sampling import SamplerConfig
+from mtplx.generation import generate_ar,restore_or_prefill_prompt_state
+from mtplx.cache_state import snapshot_cache,snapshot_cache_lazy_hybrid
+from mtplx.session_bank import SessionBank,_lazy_snapshot_enabled
+from mtplx.kernels import moe_sorted_gather_nax as kernel
+assert mx.__version__=='0.32.2'
+a.out.mkdir(parents=True,exist_ok=True)
+model,config=load_model(a.root,lazy=False,strict=True,get_model_classes=lambda **_:(qwen4_exp.Model,qwen4_exp.ModelArgs))
+model.post_weight_load(a.root)
+tokenizer=_load_tokenizer_resilient(a.root,config)
+from mtplx.mtp_patch import MTPContract
+if a.hidden: assert model.attach_mtp(a.root)
+rt=MTPLXRuntime(model,tokenizer,a.root,a.hidden,MTPContract() if a.hidden else None)
+def fingerprint(v):
+ values=np.asarray(v.astype(mx.float32)).astype('<f4')
+ assert np.isfinite(values).all()
+ return dict(shape=list(v.shape),sha256=hashlib.sha256(values.tobytes()).hexdigest())
+for chunk in [2048,4096]:
+ os.environ['MTPLX_PREFILL_CHUNK_SIZE']=str(chunk)
+ bank=SessionBank()
+ for phase in ['cold','warm']:
+  name=f'{chunk}-{phase}' if a.hidden else f'{chunk}-0-{phase}'
+  expected=json.loads((a.native/(name+'.json')).read_text())
+  prompt=expected['state']['ids'] if a.hidden else expected['state']['ids'][:expected['metrics']['prompt']]
+  if a.hidden:
+   result=restore_or_prefill_prompt_state(rt,prompt,mtp_history_policy='cycle',session_bank=bank,session_id='issue112',capture_hidden=True,store_prefix_snapshot=False,prefill_callback=lambda event:print(json.dumps(dict(case='progress',chunk=chunk,phase=phase,event=event),default=str),flush=True))
+   cache,logits,hidden,ids,cached=result.trunk_cache,result.logits,result.hidden,prompt,result.cached_tokens
+  else:
+   result=generate_ar(rt,prompt,max_tokens=1,sampler=SamplerConfig(temperature=0.,top_p=.95,top_k=20),seed=1,stop_token_ids=set(),capture_final_state=True,session_bank=bank,session_id='issue112',store_prefix_snapshot=False,token_callback=lambda ids:print(json.dumps(dict(case='tokens',chunk=chunk,phase=phase,ids=ids)),flush=True))
+   final=result.final_state
+   assert final.safe_to_commit
+   cache,logits,hidden,ids,cached=final.final_trunk_cache,final.final_logits,final.final_hidden,prompt+result.tokens,result.stats.cached_tokens
+  states=[[None if v is None else fingerprint(v) for v in state] if state is not None else None for state in (snapshot_cache_lazy_hybrid(cache) if _lazy_snapshot_enabled() else snapshot_cache(cache)).states]
+  record=dict(state=dict(ids=ids,logits=fingerprint(logits),states=states),hidden=None if hidden is None else fingerprint(hidden))
+  (a.out/(name+'.json')).write_text(json.dumps(record))
+  assert record['state']==expected['state'],(chunk,phase,'state differs')
+  assert record['hidden']==expected['hidden'],(chunk,phase,'hidden differs')
+  assert cached==(expected['cached'] if a.hidden else expected['metrics']['cached'])
+  print(json.dumps(dict(case='model_oracle',chunk=chunk,phase=phase,hidden=a.hidden,exact=True,cached=cached,kernel=kernel.stats())),flush=True)
+  if phase=='cold':
+   assert bank.put(runtime=rt,token_ids=ids,cache=cache,logits=logits,hidden=hidden,hidden_variant='post_norm',session_id='issue112',mtp_history_policy='cycle',snapshot_epoch=len(ids)) is not None

tools/issue112-reference.py

@@ -0,0 +1,45 @@
+"""Pinned row-map kernel, guarded stock chain and native synthetic bit receipts."""
+import hashlib,json,subprocess,sys
+from pathlib import Path
+checkout,native=map(Path,sys.argv[1:3])
+assert subprocess.check_output(['git','-C',str(checkout),'rev-parse','HEAD'],text=True).strip()=='9882703f3105363ddc37eca9f97aa09a1d387112'
+sys.path.insert(0,str(checkout))
+import mlx.core as mx
+import numpy as np
+from mtplx.kernels import moe_sorted_gather_nax as kernel
+assert mx.__version__=='0.32.2'
+def pattern(count,salt,scale=1):
+ a=((np.arange(count,dtype=np.uint32)*17+salt*13)%257).astype(np.int32)-128
+ return mx.array(a.astype(np.float32)/128*scale).astype(mx.bfloat16)
+def pack(e,n,k,bits,group,salt,dtype):
+ words=np.arange(e*n*k//32*bits,dtype=np.uint32)*np.uint32(2654435761)+np.uint32(12345)
+ return (mx.array(words.reshape(e,n,k*bits//32)),pattern(e*n*k//group,salt,1/64).reshape(e,n,k//group).astype(dtype),pattern(e*n*k//group,salt+1,1/64).reshape(e,n,k//group).astype(dtype))
+def digest(a):
+ mx.eval(a)
+ return hashlib.sha256(np.asarray(a.view(mx.uint16)).tobytes()).hexdigest()
+count=0
+for line in native.read_text().splitlines():
+ if '"case":"synthetic"' not in line:continue
+ case=json.loads(line[line.index('{'):])
+ count+=1
+ dtype={'BF16':mx.bfloat16,'F16':mx.float16}[case['dtype']]
+ bits,group,rows=case['bits'],case['group'],case['rows']
+ gate=pack(128,256,256,bits,group,7,dtype)
+ up=pack(128,256,256,bits,group,9,dtype) if case['separate'] else None
+ x=pattern(512*256,8).reshape(512,1,256).astype(dtype)
+ r=np.arange(rows,dtype=np.uint32)
+ ids=mx.array((r*125//rows)//2*2+1)
+ row_map=mx.array((r*13)%512)
+ assert kernel.applies(x,row_map,*gate,ids,group_size=group,bits=bits,mode='affine',up=up)
+ kwargs=dict(group_size=group,bits=bits,swiglu=True,up=up)
+ before=kernel.stats()
+ y=kernel.gather_rows_qmm(x,row_map,*gate,ids,**kwargs)
+ assert y is not None,kernel.stats()
+ stock=kernel._stock(x,row_map,*gate,ids,**kwargs)
+ actual=digest(y)
+ assert actual==digest(stock)==case['sha256'],(case,actual,digest(stock))
+ print(json.dumps(dict(case='synthetic_oracle',native=case,exact=True,kernel_before=before,kernel_after=kernel.stats())),flush=True)
+ del x,row_map,ids,gate,up,y,stock
+ mx.clear_cache()
+
+assert count==72,count