| """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 |
|
|
|
|
| @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) |
| c, s = pk["cs"][0][:, :L].T, pk["cs"][1][:, :L].T |
|
|
| 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 |
|
|
