File size: 2,912 Bytes
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 | """Shared helpers: NumPy fp32 references and import paths.
The raw kernel modules are imported by file (not through the package __init__, which needs torch), so
the simulator tests run in a venv with only nki + numpy + ml_dtypes.
"""
import os
import sys
os.environ.setdefault("NEURON_PLATFORM_TARGET_OVERRIDE", "trn2")
import numpy as np # noqa: E402
KDIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "build", "torch-neuron")
if KDIR not in sys.path:
sys.path.insert(0, KDIR)
EPS = 1e-12
D = 64
def layer_norm(s, g, b, eps=EPS):
return (s - s.mean(-1, keepdims=True)) / np.sqrt(s.var(-1, keepdims=True) + eps) * g + b
def gelu_tanh(v):
return 0.5 * v * (1 + np.tanh(0.7978845608028654 * (v + 0.044715 * v ** 3)))
def gated_mlp_ln_ref(x, wu, wg, wd, bd, g, b, eps=EPS):
"""wu, wg [I, K]; wd [K, I] (nn.Linear layout)."""
x = x.astype(np.float32)
m = (gelu_tanh(x @ wg.T) * (x @ wu.T)) @ wd.T + bd
return layer_norm(x + m, g, b, eps)
def rope_table(L, base=20000.0):
inv = 1.0 / (base ** (np.arange(0, D, 2) / D))
e = np.outer(np.arange(L), inv)
e = np.concatenate([e, e], -1)
return np.cos(e), np.sin(e)
def _rot_half(x):
return np.concatenate([-x[..., D // 2:], x[..., :D // 2]], -1)
def attn_block_ref(x, wqkv, bqkv, wo, bo, g, b, cos, sin, mask, eps=EPS):
"""NewLayer attention half in fp32: qkv_proj -> rotate-half RoPE -> masked softmax -> o_proj -> +x -> LN.
x [B, L, H], mask [B, L] (1 valid, 0 padded)."""
B, L, H = x.shape
NH = H // D
qkv = x @ wqkv.T + bqkv
q, k, v = [t.reshape(B, L, NH, D) for t in np.split(qkv, 3, -1)]
c, s = cos[None, :L, None, :], sin[None, :L, None, :]
q = q * c + _rot_half(q) * s
k = k * c + _rot_half(k) * s
sc = np.einsum("blhd,bmhd->bhlm", q, k) / np.sqrt(D) + np.where(mask > 0, 0, -1e9)[:, 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) @ wo.T + bo
return layer_norm(x + o, g, b, eps)
def load_package():
"""Import build/torch-neuron as the package `gte_encoder_neuron_kernels` (needs torch + torch_neuronx).
Equivalent to kernels.get_kernel(...) on a published repo."""
import importlib.util
name = "gte_encoder_neuron_kernels"
if name in sys.modules:
return sys.modules[name]
spec = importlib.util.spec_from_file_location(name, os.path.join(KDIR, "__init__.py"),
submodule_search_locations=[KDIR])
mod = importlib.util.module_from_spec(spec)
sys.modules[name] = mod
spec.loader.exec_module(mod)
return mod
def rel_l2(y, e, rows=None):
y = np.asarray(y, np.float32)
e = np.asarray(e, np.float32)
if rows is not None:
y, e = y[rows], e[rows]
return float(np.linalg.norm(y - e) / np.linalg.norm(e))
|