37 lines
2.1 KiB
Python
37 lines
2.1 KiB
Python
"""Exact attention gate laws for all finite BF16/F16 values and dense F32 inputs."""
|
|
import hashlib,json,subprocess,sys
|
|
from pathlib import Path
|
|
checkout,out=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.attention_math import attention_gate
|
|
from mtplx.attention_context import attention_phase
|
|
from mtplx.models.qwen4_exp import _verify_inject
|
|
out.mkdir(parents=True,exist_ok=True)
|
|
for name,dtype in [('bf16',mx.bfloat16),('f16',mx.float16),('f32',mx.float32)]:
|
|
if name=='bf16':
|
|
values=(np.arange(65536,dtype=np.uint32)<<16).view(np.float32)
|
|
elif name=='f16':
|
|
values=np.arange(65536,dtype=np.uint16).view(np.float16).astype(np.float32)
|
|
else:
|
|
values=np.concatenate([np.linspace(-30,30,32768,dtype=np.float32),np.random.default_rng(5).normal(0,4,32768).astype(np.float32)])
|
|
values=values[np.isfinite(values)]
|
|
gate=mx.array(values).astype(dtype)
|
|
output=mx.array(((np.arange(values.size,dtype=np.int64)*17)%97-48).astype(np.float32)/32).astype(dtype)
|
|
mx.eval(gate,output)
|
|
for part,a in [('gate',gate),('output',output)]:
|
|
(out/f'{name}-{part}.f32').write_bytes(np.asarray(a.astype(mx.float32)).tobytes())
|
|
for phase in ['prefill','ar_decode','decode_verify']:
|
|
with attention_phase(phase): result=attention_gate(output,gate)
|
|
raw=np.asarray(result.astype(mx.float32)).tobytes()
|
|
(out/f'{name}-{phase}.f32').write_bytes(raw)
|
|
print(json.dumps(dict(case='gate',dtype=name,phase=phase,count=values.size,sha256=hashlib.sha256(raw).hexdigest())),flush=True)
|
|
if name=='bf16':
|
|
inject=_verify_inject(gate,4) if phase=='decode_verify' else 2.0*mx.sigmoid(gate/4)
|
|
raw=np.asarray(inject.astype(mx.float32)).tobytes()
|
|
(out/f'inject-{phase}.f32').write_bytes(raw)
|
|
print(json.dumps(dict(case='inject',dtype=name,phase=phase,count=values.size,sha256=hashlib.sha256(raw).hexdigest())),flush=True)
|
|
|