File size: 3,689 Bytes
25e8424 fb58c6c 25e8424 fb58c6c 25e8424 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """CPU check of the host-side packing used by NkiAttnBlock (lane permutation, 1/sqrt(d) fold, NTK RoPE
table, key bias) against the real Hugging Face NewLayer in fp32 with ragged right padding.
The packed operands go through a NumPy emulation of attn_block_kernel's math (no NKI, no device), so
this test isolates packing errors from kernel errors. The RoPE table is built the way NkiAttnBlock builds
it: once at max(ATTN_LENGTHS), then sliced per length.
Needs torch + transformers; downloads the model on first run.
"""
import os
import numpy as np
import pytest
from conftest import layer_norm
from attn_block_pack import pack
torch = pytest.importorskip("torch")
transformers = pytest.importorskip("transformers")
pytestmark = pytest.mark.cpu
MODELS = [os.environ.get("GTE_TEST_MODEL", "Alibaba-NLP/gte-multilingual-reranker-base")]
LMAX = 1024 # == max(torch_ops.ATTN_LENGTHS)
@pytest.fixture(scope="module", params=MODELS)
def new_model(request):
m = transformers.AutoModel.from_pretrained(request.param, trust_remote_code=True,
attn_implementation="eager").eval()
return m
def _emulate(pk, xn, kb, NH, D, L):
"""attn_block_kernel math on packed operands, fp32."""
B = xn.shape[0]
H = NH * D
qk = xn @ pk["wqk"] + pk["bqk"].T.reshape(-1) # [B, L, 2H], permuted lanes, Q pre-scaled
c, s = pk["cs"][0][:, :L].T, pk["cs"][1][:, :L].T # [L, 128]
def rope(t):
out = np.zeros((B, L, NH, D), np.float32)
for j in range(NH // 4):
lo = t[..., (2 * j) * 128:(2 * j + 1) * 128]
hi = t[..., (2 * j + 1) * 128:(2 * j + 2) * 128]
rlo, rhi = lo * c - hi * s, hi * c + lo * s
for bq in range(4):
out[:, :, j * 4 + bq, :32] = rlo[..., bq * 32:(bq + 1) * 32]
out[:, :, j * 4 + bq, 32:] = rhi[..., bq * 32:(bq + 1) * 32]
return out
q, k = rope(qk[..., :H]), rope(qk[..., H:])
v = (xn @ pk["wv"] + pk["bv"][0]).reshape(B, L, NH, D)
sc = np.einsum("blhd,bmhd->bhlm", q, k) + kb[:, None, None, :]
sc -= sc.max(-1, keepdims=True)
p = np.exp(sc)
p /= p.sum(-1, keepdims=True)
o = np.einsum("bhlm,bmhd->blhd", p, v).reshape(B, L, H) @ pk["wo"] + pk["bo"][0]
return layer_norm(xn + o, pk["g"][0], pk["b"][0])
@pytest.mark.parametrize("L", [128, 256, 384, 512, 768, 1024])
def test_pack_matches_hf_newlayer(new_model, L):
enc = new_model
cfg = enc.config
H, NH = cfg.hidden_size, cfg.num_attention_heads
D = H // NH
lay = enc.encoder.layer[min(3, cfg.num_hidden_layers - 1)]
rot = enc.embeddings.rotary_emb
B = 2
torch.manual_seed(0)
x = torch.randn(B, L, H)
mask = torch.ones(B, L, dtype=torch.long)
mask[0, L - 40:] = 0
mask[1, L // 3:] = 0
bias = enc.get_extended_attention_mask(mask, (B, L))
cos, sin = rot(x, seq_len=L)
pos = torch.arange(L)
with torch.no_grad():
a = lay.attention(x, bias, (cos[pos][None, :, None, :], sin[pos][None, :, None, :]))[0]
ref = lay.attn_ln(x + a).numpy()
at = lay.attention
f = lambda t: t.detach().float().numpy()
cmax, smax = rot(torch.zeros(1, dtype=torch.float32), seq_len=LMAX)
pk = pack(f(at.qkv_proj.weight), f(at.qkv_proj.bias), f(at.o_proj.weight), f(at.o_proj.bias),
f(lay.attn_ln.weight), f(lay.attn_ln.bias), f(cmax), f(smax), H, 1.0 / np.sqrt(D))
kb = np.clip(bias.reshape(B, L).numpy(), -30000.0, 0)
y = _emulate(pk, x.numpy(), kb, NH, D, L)
valid = mask.bool().numpy()
rel = np.linalg.norm((y - ref)[valid]) / np.linalg.norm(ref[valid])
assert rel < 1e-4, rel
|