File size: 2,935 Bytes
e95c403 | 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 | """02_forward_parity.py -- verify NKI kernel forward matches an eager reference.
Runs a Mamba-3 forward on the eager path and on the NKI-accelerated path with
identical weights, then compares the outputs. Both paths should agree at
cos_sim > 0.999999 (numerical noise from bf16-adjacent MLIR lowering).
Runs both SISO and MIMO modes.
Run:
python 02_forward_parity.py
"""
import torch
from _loader import load_kernel
def compare(name, y_nki, y_ref, tol_cos=0.999):
y_nki_f = y_nki.cpu().float().reshape(-1)
y_ref_f = y_ref.cpu().float().reshape(-1)
max_abs = (y_nki_f - y_ref_f).abs().max().item()
ref_norm = y_ref_f.norm().item()
nki_norm = y_nki_f.norm().item()
if ref_norm < 1e-12 and nki_norm < 1e-12:
cos = 1.0
else:
cos = torch.dot(y_nki_f, y_ref_f).item() / (ref_norm * nki_norm + 1e-12)
status = "PASS" if cos > tol_cos else "FAIL"
print(f" [{status}] {name}: cos_sim={cos:.6f} max_abs={max_abs:.3e} ||ref||={ref_norm:.3e}")
return cos > tol_cos
def parity_one(mamba3, mode: str, seqlen: int):
print(f"\n--- {mode.upper()} parity, seqlen={seqlen} ---")
if mode == "siso":
kwargs_ref = dict(d_model=1024, d_state=128, headdim=64, chunk_size=64, mimo_rank=1,
use_nki_ssd=False)
kwargs_nki = dict(d_model=1024, d_state=128, headdim=64, chunk_size=64, mimo_rank=1,
use_nki_ssd=True)
else: # mimo
kwargs_ref = dict(d_model=1024, d_state=128, headdim=64, chunk_size=16, mimo_rank=4,
use_nki_ssd=False)
kwargs_nki = dict(d_model=1024, d_state=128, headdim=64, chunk_size=16, mimo_rank=4,
use_nki_ssd=True)
torch.manual_seed(0)
mixer_ref = mamba3.NeuronMamba3Mixer(**kwargs_ref).to("neuron")
torch.manual_seed(0)
mixer_nki = mamba3.NeuronMamba3Mixer(**kwargs_nki).to("neuron")
# Copy weights from ref to nki (both were seeded identically, but this makes it explicit)
mixer_nki.load_state_dict(mixer_ref.state_dict())
torch.manual_seed(100)
u = torch.randn(1, seqlen, 1024, device="neuron")
y_ref, cache_ref = mixer_ref(u)
y_nki, cache_nki = mixer_nki(u)
try:
torch.neuron.synchronize()
except Exception:
pass
ok_y = compare("y_out", y_nki, y_ref, tol_cos=0.999)
ok_state = compare("ssm_state", cache_nki.ssm_state, cache_ref.ssm_state, tol_cos=0.999)
return ok_y and ok_state
def main():
print("=" * 60)
print("Forward parity: NKI kernel vs eager reference")
print("=" * 60)
mamba3 = load_kernel()
all_pass = True
for mode in ["siso", "mimo"]:
for seqlen in [64, 256]:
all_pass = parity_one(mamba3, mode, seqlen) and all_pass
print()
print("=" * 60)
print("ALL PARITY TESTS PASS" if all_pass else "SOME TESTS FAILED")
print("=" * 60)
if __name__ == "__main__":
main()
|