Hugging Face's logo

Kernels:
jburtoft
/
gte-encoder-neuron-kernels

Kernel card Files
xet
 Community
gte-encoder-neuron-kernels / tests /test_sim.py
jburtoft's picture
jburtoft
v1.1.0: fused attention block for L = 768 / 1024
fb58c6c
Raw History Blame
7.87 kB
"""nki.simulate parity of the three kernels vs NumPy fp32 references. Simulator only: NOT device evidence
(the simulator does not check SBUF/PSUM capacity and has known DMA-pattern fidelity gaps)."""
import numpy as np
import pytest
nki = pytest.importorskip("nki")
ml_dtypes = pytest.importorskip("ml_dtypes")
from conftest import EPS, attn_block_ref, gated_mlp_ln_ref, layer_norm, rel_l2, rope_table # noqa: E402
from attn_block_kernel import attn_block_kernel # noqa: E402
from attn_block_pack import pack # noqa: E402
from gated_mlp_kernel import gated_mlp_ln_kernel, gated_mlp_ln_stream_kernel, mlp_weights_fit # noqa: E402
from res_ln_kernel import res_ln_kernel # noqa: E402
bf16 = ml_dtypes.bfloat16
TOL = 1e-2 # bf16 in/out; observed ~3-5e-3
pytestmark = pytest.mark.sim
def _bc(v):
return np.tile(np.asarray(v, np.float32), (128, 1))
# ---------------------------------------------------------------- residual + LayerNorm
@pytest.mark.parametrize("R,lnc,has_res,H", [(256, 1, True, 768), (512, 2, True, 768), (256, 1, False, 768),
(1024, 2, False, 768), (256, 1, True, 1024), (512, 2, True, 1024)])
def test_res_ln(R, lnc, has_res, H):
rng = np.random.default_rng(R + H)
x = (rng.standard_normal((R, H)) * 3 + 0.5).astype(bf16)
r = rng.standard_normal((R, H)).astype(bf16)
g = 1 + 0.1 * rng.standard_normal(H)
b = 0.1 * rng.standard_normal(H)
y = nki.simulate(res_ln_kernel[lnc])(x, r, _bc(g), _bc(b), EPS, has_res)
s = x.astype(np.float32) + (r.astype(np.float32) if has_res else 0)
assert rel_l2(y, layer_norm(s, g, b)) < TOL
# ---------------------------------------------------------------- gated MLP + residual + LayerNorm
def _mlp_case(R, K, I, seed):
rng = np.random.default_rng(seed)
x = rng.standard_normal((R, K)).astype(bf16)
wu = (rng.standard_normal((I, K)) / np.sqrt(K)).astype(bf16)
wg = (rng.standard_normal((I, K)) / np.sqrt(K)).astype(bf16)
wd = (rng.standard_normal((K, I)) / np.sqrt(I)).astype(bf16)
bd, g, b = 0.1 * rng.standard_normal(K), 1 + 0.1 * rng.standard_normal(K), 0.1 * rng.standard_normal(K)
args = (x, np.ascontiguousarray(wu.T), np.ascontiguousarray(wg.T), np.ascontiguousarray(wd.T),
_bc(bd), _bc(g), _bc(b), EPS)
ref = gated_mlp_ln_ref(x, wu.astype(np.float32), wg.astype(np.float32), wd.astype(np.float32), bd, g, b)
return args, ref
@pytest.mark.parametrize("R,lnc", [(256, 1), (640, 1), (512, 2)])
def test_gated_mlp_ln(R, lnc):
assert mlp_weights_fit(768, 3072)
args, ref = _mlp_case(R, 768, 3072, R)
assert rel_l2(nki.simulate(gated_mlp_ln_kernel[lnc])(*args), ref) < TOL
# streaming variant: 1024/4096 (gte-large) does not fit resident; also 768/3072 for a direct comparison,
# and R > STREAM_SB (two super-blocks) with a partial last super-block.
@pytest.mark.parametrize("R,lnc,K,I", [(256, 1, 1024, 4096), (512, 2, 1024, 4096), (1280, 1, 1024, 4096),
(384, 1, 768, 3072)])
def test_gated_mlp_ln_stream(R, lnc, K, I):
if K == 1024:
assert not mlp_weights_fit(K, I)
args, ref = _mlp_case(R, K, I, R + K)
assert rel_l2(nki.simulate(gated_mlp_ln_stream_kernel[lnc])(*args), ref) < TOL
# ---------------------------------------------------------------- fused attention block
H = 768
def _attn_case(B, L, seed, lengths=None, H=H):
rng = np.random.default_rng(seed)
x = rng.standard_normal((B, L, H)).astype(bf16)
wqkv = (rng.standard_normal((3 * H, H)) * 0.06).astype(bf16).astype(np.float32)
bqkv = (0.1 * rng.standard_normal(3 * H)).astype(np.float32)
wo = (rng.standard_normal((H, H)) / np.sqrt(H)).astype(bf16).astype(np.float32)
bo = (0.1 * rng.standard_normal(H)).astype(np.float32)
g = (1 + 0.1 * rng.standard_normal(H)).astype(np.float32)
b = (0.1 * rng.standard_normal(H)).astype(np.float32)
cos, sin = rope_table(L)
if lengths is None: # ragged right padding, every row different
lengths = [max(1, L - 29 * (i + 1)) for i in range(B)]
mask = np.zeros((B, L), np.float32)
for i, n in enumerate(lengths):
mask[i, :n] = 1
kb = ((1 - mask) * -30000.0).astype(bf16)
pk = pack(wqkv, bqkv, wo, bo, g, b, cos, sin, H, 1.0 / np.sqrt(64))
args = [x.reshape(B * L, H), pk["wqk"].astype(bf16), pk["bqk"], pk["wv"].astype(bf16), pk["bv"],
pk["wo"].astype(bf16), pk["bo"], pk["g"], pk["b"], pk["cs"], kb]
ref = (x.astype(np.float32), wqkv, bqkv, wo, bo, g, b, cos, sin, mask)
return args, ref
# (B, L, lnc): packing factor SP = 4 / 2 / 1 at L = 128 / 256 / >=384; B not divisible by SP falls back to
# SP = 1; LNC2 with odd B uses the query split (needs (L/128) % 2 == 0).
ATTN_CASES = [
(4, 128, 1), # SP=4
(2, 128, 1), # SP=4 does not divide -> SP=1
(3, 128, 1), # odd B, SP=1
(2, 256, 1), # SP=2
(1, 384, 1), # L=384, SP=1
(1, 512, 1),
(8, 128, 2), # LNC2, SP=4 per core
(4, 256, 2), # LNC2, SP=2 per core
(2, 384, 2), # LNC2, L=384
(1, 256, 2), # LNC2 query split, B=1
(1, 512, 2), # LNC2 query split, B=1
(3, 256, 2), # LNC2 query split, odd B
]
@pytest.mark.parametrize("B,L,lnc", ATTN_CASES)
def test_attn_block(B, L, lnc):
args, ref = _attn_case(B, L, seed=B * 1000 + L + lnc)
y = np.asarray(nki.simulate(attn_block_kernel[lnc])(*args, EPS)).astype(np.float32).reshape(B, L, H)
assert rel_l2(y, attn_block_ref(*ref)) < TOL
# hidden 1024 / 16 heads (gte-large-en-v1.5)
@pytest.mark.parametrize("B,L,lnc", [(4, 128, 1), (2, 256, 1), (1, 512, 1), (4, 256, 2), (3, 256, 2)])
def test_attn_block_h1024(B, L, lnc):
args, ref = _attn_case(B, L, seed=B * 1000 + L + lnc + 7, H=1024)
y = np.asarray(nki.simulate(attn_block_kernel[lnc])(*args, EPS)).astype(np.float32).reshape(B, L, 1024)
assert rel_l2(y, attn_block_ref(*ref)) < TOL
# long sequences (key chunks of 512, grouped LayerNorm epilogue), H 768 and 1024, ragged padding incl. inside the
# last key chunk; LNC2 with odd B uses the query split.
@pytest.mark.parametrize("B,L,lnc,Hm", [(1, 768, 1, 768), (2, 1024, 1, 768), (1, 1024, 2, 768), (3, 768, 2, 768),
(1, 1024, 1, 1024), (2, 768, 2, 1024)])
def test_attn_block_long(B, L, lnc, Hm):
lengths = [L - 29 * (i + 1) for i in range(B)]
lengths[0] = L - 600 if L == 1024 else lengths[0] # padding reaches into the first key chunk
args, ref = _attn_case(B, L, seed=B * 1000 + L + lnc + Hm, lengths=lengths, H=Hm)
y = np.asarray(nki.simulate(attn_block_kernel[lnc])(*args, EPS)).astype(np.float32).reshape(B, L, Hm)
assert rel_l2(y, attn_block_ref(*ref)) < TOL
def test_attn_block_short_sequence():
"""One valid token (all other keys padded) and one full-length sequence in the same pass."""
args, ref = _attn_case(2, 256, seed=7, lengths=[1, 256])
y = np.asarray(nki.simulate(attn_block_kernel[1])(*args, EPS)).astype(np.float32).reshape(2, 256, H)
assert rel_l2(y, attn_block_ref(*ref)) < TOL
@pytest.mark.parametrize("B,L,lnc", [(4, 128, 1), (3, 256, 2)])
def test_attn_block_padding_isolation(B, L, lnc):
"""Sequence 0's output must be bit-identical when every other sequence in the batch changes."""
args, _ = _attn_case(B, L, seed=11)
y0 = np.asarray(nki.simulate(attn_block_kernel[lnc])(*args, EPS)).view(np.uint16).reshape(B, L, H)
rng = np.random.default_rng(12)
x2 = args[0].reshape(B, L, H).copy()
x2[1:] = rng.standard_normal(x2[1:].shape).astype(bf16)
kb2 = args[-1].copy()
kb2[1:] = 0
kb2[1:, L // 2:] = -30000.0
args2 = [x2.reshape(B * L, H)] + args[1:-1] + [kb2]
y1 = np.asarray(nki.simulate(attn_block_kernel[lnc])(*args2, EPS)).view(np.uint16).reshape(B, L, H)
assert np.array_equal(y0[0], y1[0])