54 lines
3.4 KiB
Python
54 lines
3.4 KiB
Python
"""Real-geometry image-aware QSA preparation, checked independently with NumPy."""
|
|
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 TextArgs, QSAIndexer, _build_mrope_axes
|
|
from mtplx.attention_context import vision_rope
|
|
|
|
output.mkdir(parents=True, exist_ok=True)
|
|
args = TextArgs(hidden_size=640, head_dim=256, indexer_head_dim=128, indexer_n_heads=4,
|
|
indexer_kv_heads=1, indexer_compress_ratio=4, indexer_budget=2048,
|
|
rope_parameters=dict(mrope_interleaved=True, mrope_section=[11,11,10], partial_rotary_factor=.25,
|
|
rope_theta=10000000, rope_type="default"))
|
|
idx = QSAIndexer(args)
|
|
axes = _build_mrope_axes([11,11,10],True)
|
|
table = np.stack([np.arange(2060, dtype=np.int32)//(axis+1) for axis in range(3)])
|
|
def pattern(shape, seed):
|
|
return (((np.arange(np.prod(shape), dtype=np.int64)*17+seed)%97-48)/32).astype(np.float32).reshape(shape)
|
|
raw_q = pattern((1,16,4,128),3)
|
|
raw_k = pattern((1,2072,128),11)
|
|
|
|
for dtype_name, dtype in [("f32",mx.float32),("bf16",mx.bfloat16)]:
|
|
idx.q_layernorm.weight = mx.ones((128,),dtype=dtype)
|
|
idx.k_layernorm.weight = mx.ones((128,),dtype=dtype)
|
|
q, k = mx.array(raw_q).astype(dtype), mx.array(raw_k).astype(dtype)
|
|
for mode in ["table", "equal", "delta", "empty"]:
|
|
delta = -7 if mode in ("table","delta","empty") else 0
|
|
table_value = table if mode=="table" else np.broadcast_to(np.arange(2060,dtype=np.int32),(3,2060)) if mode=="equal" else np.zeros((3,0),dtype=np.int32) if mode=="empty" else None
|
|
with vision_rope(None if table_value is None else mx.array(table_value),delta):
|
|
query = idx._prepare_queries(q,2056)
|
|
pooled = idx._pool_keys_eager(k,0,518)
|
|
for name,value in [("query",query),("pooled",pooled)]:
|
|
data=np.asarray(value.astype(mx.float32))
|
|
(output/f"{dtype_name}-{mode}-{name}.f32").write_bytes(data.tobytes())
|
|
print(json.dumps(dict(case="indexer",dtype=dtype_name,mode=mode,part=name,shape=list(data.shape),sha256=hashlib.sha256(data.tobytes()).hexdigest())),flush=True)
|
|
if dtype_name=="f32":
|
|
def independent(values, positions):
|
|
values = values / np.sqrt(np.mean(values*values,axis=-1,keepdims=True)+args.rms_norm_eps)
|
|
pos3=np.stack([table_value[:,p] if table_value is not None and p<table_value.shape[1] else np.full(3,p+delta) for p in positions],axis=1)
|
|
angles=pos3[axes].T.astype(np.float32)*np.asarray(idx._inv_freq)[None,:]
|
|
angles=np.concatenate([angles,angles],axis=-1)[None,:,None,:]
|
|
first=values[...,:64];rotated=np.concatenate([-first[...,32:],first[...,:32]],axis=-1)
|
|
return np.concatenate([first*np.cos(angles)+rotated*np.sin(angles),values[...,64:]],axis=-1)
|
|
np.testing.assert_allclose(np.asarray(query),independent(raw_q,np.arange(2056,2072)),rtol=3e-6,atol=3e-6)
|
|
np.testing.assert_allclose(np.asarray(pooled),independent(raw_k.reshape(1,518,4,128).mean(axis=2)[:,:,None],np.arange(0,2072,4))[:,:,0],rtol=3e-6,atol=3e-6)
|
|
print(json.dumps(dict(case="independent_numpy",passed=True)),flush=True)
|