File size: 7,872 Bytes
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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | """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])
|