56 lines
3.4 KiB
Python
56 lines
3.4 KiB
Python
"""Pinned attention chunk RoPE oracle, with independent per-row controls."""
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
|
|
checkout, output = 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.models.qwen4_exp import _vision_chunk_cos_sin, _apply_partial_rope, _build_mrope_axes, _rope_cos_sin, _rope_inv_freq_and_scaling, TextArgs
|
|
|
|
output.mkdir(parents=True, exist_ok=True)
|
|
manifest = []
|
|
for dim, rotary, sections, heads in [(8, 8, [2,1,1], 2), (256,64,[11,11,10],24)]:
|
|
axes = mx.array(_build_mrope_axes(sections, True), dtype=mx.int32)
|
|
freq, _ = _rope_inv_freq_and_scaling(TextArgs(head_dim=dim, rope_parameters=dict(rope_type="default", rope_theta=10000000, partial_rotary_factor=rotary/dim)))
|
|
short = [[0,1,2,2],[0,1,3,3],[0,1,4,4]]
|
|
full = [r + [3,4,5,6] for r in short]
|
|
equal = [list(range(4))] * 3
|
|
for mode, table, start, rows, delta in [
|
|
('inside',short,0,4,-1), ('outside',short,4,2,-1),
|
|
('straddle',short,2,4,-1), ('full_text_tail',full,6,4,-1),
|
|
('equal',equal,2,4,0), ('none',None,4,2,-1), ('empty',[[],[],[]],4,2,-1),
|
|
]:
|
|
for scaling in [1.,1.25]:
|
|
t = None if table is None else mx.array(table,dtype=mx.int32)
|
|
cos, sin = _vision_chunk_cos_sin(t,delta,start,rows,freq,axes,scaling)
|
|
# Independent row-position construction; no chunk helper is used.
|
|
pos3 = np.stack([np.array(table)[:,p] if table is not None and p<len(table[0]) else np.full(3,p+delta) for p in range(start,start+rows)],axis=1)
|
|
angles = pos3[np.asarray(axes)].T.astype(np.float32) * np.asarray(freq)[None,:]
|
|
angles = np.concatenate([angles,angles],axis=1)
|
|
np.testing.assert_allclose(np.asarray(cos),np.cos(angles)*scaling,rtol=3e-6,atol=3e-6)
|
|
np.testing.assert_allclose(np.asarray(sin),np.sin(angles)*scaling,rtol=3e-6,atol=3e-6)
|
|
if mode in ('full_text_tail','equal'):
|
|
plain = _rope_cos_sin(mx.arange(start+delta,start+rows+delta,dtype=mx.int32),freq,scaling)
|
|
assert all(mx.array_equal(a,b).item() for a,b in zip((cos,sin),plain))
|
|
for dtype_name,dtype in [('f32',mx.float32),('bf16',mx.bfloat16)]:
|
|
key = f'{dim}-{mode}-{scaling}-{dtype_name}'
|
|
record = dict(key=key,dim=dim,rotary=rotary,sections=sections,heads=heads,mode=mode,table=table,start=start,rows=rows,delta=delta,scaling=scaling,dtype=dtype_name)
|
|
arrays = [('cos',cos),('sin',sin)]
|
|
for part,h,seed in [('q',heads,3),('k',1,11)]:
|
|
shape = (1,rows,h,dim)
|
|
raw = (((np.arange(np.prod(shape),dtype=np.int64)*17+seed)%97-48)/32).astype(np.float32).reshape(shape)
|
|
value = _apply_partial_rope(mx.array(raw).astype(dtype),cos,sin)
|
|
arrays.append((part,value))
|
|
for part,value in arrays:
|
|
data=np.asarray(value.astype(mx.float32)).tobytes()
|
|
(output/f'{key}-{part}.f32').write_bytes(data)
|
|
record[part+'_sha256']=hashlib.sha256(data).hexdigest()
|
|
manifest.append(record)
|
|
print(json.dumps(record),flush=True)
|
|
(output/'manifest.json').write_text(json.dumps(manifest))
|