| """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 |
|
|
| 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)) |
|
|
