Guard large unaligned Qwen expert gathers

tools/issue110-model-reference.py

@@ -0,0 +1,53 @@
+"""Full installed-model cold/warm chunk-guard reference, with exact state checks."""
+import argparse,hashlib,json,os,subprocess,sys
+from pathlib import Path
+checkout,root,native,out=map(Path,sys.argv[1:5])
+assert subprocess.check_output(['git','-C',str(checkout),'rev-parse','HEAD'],text=True).strip()=='9882703f3105363ddc37eca9f97aa09a1d387112'
+sys.path.insert(0,str(checkout))
+from mtplx.server.openai import _server_runtime_env_overrides
+if os.environ.get('DS4_ISSUE110_STOCK_CONTROL')!='1':
+    os.environ.update(_server_runtime_env_overrides(argparse.Namespace(model=str(root),generation_mode='ar'),None))
+os.environ.update(MTPLX_COMPILED_GDN='1',MTPLX_QWEN4EXP_COMPILE='1',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
+from mtplx.cache_state import snapshot_cache
+from mtplx.session_bank import SessionBank
+from mtplx import moe_sorted_gather as msg
+from mtplx.kernels import moe_sorted_gather_nax as kernel
+assert mx.__version__=='0.32.2'
+assert os.environ.get('MTPLX_GDN_BOUNDARY_TAIL_INTERVAL')=='0'
+out.mkdir(parents=True,exist_ok=True)
+model,config=load_model(root,lazy=False,strict=True,get_model_classes=lambda **_:(qwen4_exp.Model,qwen4_exp.ModelArgs))
+model.post_weight_load(root)
+tokenizer=_load_tokenizer_resilient(root,config)
+rt=MTPLXRuntime(model,tokenizer,root,False,None)
+def fingerprint(a):
+    values=np.asarray(a.astype(mx.float32)).astype('<f4')
+    assert np.isfinite(values).all()
+    return dict(shape=list(a.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']:
+        expected=json.loads((native/f'{chunk}-{phase}.json').read_text())
+        prompt=expected['state']['ids'][:expected['prompt_len']]
+        guard_before=msg.stats()
+        result=generate_ar(rt,prompt,max_tokens=3,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='issue110',store_prefix_snapshot=False,
+            token_callback=lambda ids:print(json.dumps(dict(case='token',chunk=chunk,phase=phase,ids=ids)),flush=True))
+        final=result.final_state
+        assert final.safe_to_commit
+        states=[[None if a is None else fingerprint(a) for a in state] if state is not None else None for state in snapshot_cache(final.final_trunk_cache).states]
+        state=dict(ids=prompt+result.tokens,logits=fingerprint(final.final_logits),states=states)
+        (out/f'{chunk}-{phase}.json').write_text(json.dumps(state))
+        assert result.stats.cached_tokens==expected['cached']
+        assert state==expected['state'],(chunk,phase,'state differs')
+        print(json.dumps(dict(case='turn',chunk=chunk,phase=phase,exact=True,cached=result.stats.cached_tokens,tokens=result.tokens,guard_before=guard_before,guard_after=msg.stats(),kernel=kernel.stats())),flush=True)
+        if phase=='cold':
+            assert bank.put(runtime=rt,token_ids=state['ids'],cache=final.final_trunk_cache,logits=final.final_logits,hidden=None,
+                hidden_variant='post_norm',session_id='issue110',mtp_history_policy='cycle',snapshot_epoch=len(state['ids'])) is not None

tools/issue110-reference.py

@@ -0,0 +1,55 @@
+"""Pinned stock guard, installed projection geometry, and fused-oracle controls."""
+import hashlib,json,subprocess,sys
+from pathlib import Path
+checkout,root=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 mlx.nn as nn
+import numpy as np
+from mtplx import moe_sorted_gather as msg
+from mtplx.kernels import moe_sorted_gather_nax as kernel
+assert mx.__version__=='0.32.2' and msg.guard_active()
+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']]
+gu=[mx.concatenate(a,axis=1) for a in zip(parts('gate_proj'),parts('up_proj'))]
+down=parts('down_proj')
+mx.eval(*gu,*down)
+def pattern(rows,k):
+    v=((np.arange(rows*k,dtype=np.uint32)*17+7*13)%257).astype(np.int32)-128
+    return mx.array((v.astype(np.float32)/128).reshape(rows,1,k)).astype(mx.bfloat16)
+def digest(a): return hashlib.sha256(np.asarray(a.view(mx.uint16)).tobytes()).hexdigest()
+for rows in [20480,32760,32770,34040,40950,40960]:
+    for name,layer,k,n in [('gate_up',gu,2560,1280),('down',down,640,2560)]:
+        x=pattern(rows,k)
+        ids=mx.array(np.arange(rows,dtype=np.uint32)*512//rows)
+        for left in ([False,True] if rows==34040 else [False]):
+            lhs=mx.arange(rows-1,-1,-1,dtype=mx.int32) if left else None
+            before=msg.stats()
+            y=msg.gather_qmm(x,*layer,lhs_indices=lhs,rhs_indices=ids,group_size=64,bits=4,sorted_indices=True)
+            mx.eval(y)
+            assert y.shape==(rows,1,n)
+            print(json.dumps(dict(case='projection',projection=name,rows=rows,k=k,n=n,left=left,sha256=digest(y),pad=msg.pad_rows(rows),guard_before=before,guard_after=msg.stats())),flush=True)
+            del y
+        del x,ids
+        mx.clear_cache()
+for tokens in [2048,3277,4095,4096]:
+    rows=tokens*10
+    x=pattern(tokens,2560)
+    row_map=mx.arange(rows,dtype=mx.uint32)%tokens
+    ids=mx.array(np.arange(rows,dtype=np.uint32)*512//rows)
+    projection=msg.gather_qmm(x[row_map],*gu,rhs_indices=ids,group_size=64,bits=4,sorted_indices=True)
+    gate,up=mx.split(projection,2,axis=-1)
+    stock=nn.silu(gate)*up
+    before=kernel.stats()
+    fused=msg.swiglu_rows(x,row_map,*gu,ids,group_size=64,bits=4)
+    mx.eval(stock,fused)
+    assert digest(stock)==digest(fused)
+    after=kernel.stats()
+    assert after['calls']>before['calls'],(before,after)
+    assert after['canaries']>0 and after['canary_failures']==0,after
+    print(json.dumps(dict(case='optimized_oracle_control',tokens=tokens,rows=rows,exact=True,sha256=digest(stock),kernel_before=before,kernel_after=after)),flush=True)
+    del x,row_map,ids,projection,gate,up,stock,fused
+    mx.clear_cache()