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