Port Qwen row-map SwiGLU expert gathers

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