| """03_decode_generation.py -- demonstrate prefill + N decode steps. |
| |
| Prefills a small context, then runs N single-token decode steps via |
| `mixer.step(u, cache)`. Verifies the cache persists correctly across steps |
| by comparing NKI-accelerated decode outputs against eager on identical weights. |
| |
| Run: |
| python 03_decode_generation.py |
| """ |
|
|
| import time |
| import torch |
| from _loader import load_kernel |
|
|
|
|
| def main(): |
| print("=" * 60) |
| print("Prefill + Decode: 128 consecutive decode steps") |
| print("=" * 60) |
| mamba3 = load_kernel() |
|
|
| torch.manual_seed(0) |
| |
| mixer_nki = mamba3.NeuronMamba3Mixer( |
| d_model=1024, d_state=128, headdim=64, |
| chunk_size=64, mimo_rank=1, use_nki_ssd=True, |
| ).to("neuron") |
| torch.manual_seed(0) |
| mixer_ref = mamba3.NeuronMamba3Mixer( |
| d_model=1024, d_state=128, headdim=64, |
| chunk_size=64, mimo_rank=1, use_nki_ssd=False, |
| ).to("neuron") |
| mixer_ref.load_state_dict(mixer_nki.state_dict()) |
|
|
| torch.manual_seed(100) |
| prefill_len = 64 |
| n_decode = 128 |
| u_prefill = torch.randn(1, prefill_len, 1024, device="neuron") |
| tokens = [torch.randn(1, 1, 1024, device="neuron") for _ in range(n_decode)] |
|
|
| print(f"\nprefill_len={prefill_len}, decode_steps={n_decode}") |
|
|
| |
| y_pre_nki, cache_nki = mixer_nki(u_prefill) |
| y_pre_ref, cache_ref = mixer_ref(u_prefill) |
| try: |
| torch.neuron.synchronize() |
| except Exception: |
| pass |
|
|
| diff = (y_pre_nki.cpu() - y_pre_ref.cpu()).abs().max().item() |
| cos = torch.dot(y_pre_nki.cpu().float().reshape(-1), y_pre_ref.cpu().float().reshape(-1)).item() |
| cos /= (y_pre_nki.norm().item() * y_pre_ref.norm().item() + 1e-12) |
| print(f"prefill: cos_sim={cos:.6f} max_abs={diff:.3e}") |
|
|
| |
| print(f"\nrunning {n_decode} decode steps...") |
| n_pass = 0 |
| times = [] |
| for i, tok in enumerate(tokens): |
| |
| t0 = time.perf_counter() |
| y_nki, cache_nki = mixer_nki.step(tok, cache_nki) |
| try: |
| torch.neuron.synchronize() |
| except Exception: |
| pass |
| elapsed = (time.perf_counter() - t0) * 1000 |
| times.append(elapsed) |
|
|
| |
| y_ref, cache_ref = mixer_ref.step(tok, cache_ref) |
| try: |
| torch.neuron.synchronize() |
| except Exception: |
| pass |
|
|
| y_n = y_nki.cpu().float().reshape(-1) |
| y_r = y_ref.cpu().float().reshape(-1) |
| cos = torch.dot(y_n, y_r).item() / (y_n.norm().item() * y_r.norm().item() + 1e-12) |
| if cos > 0.999: |
| n_pass += 1 |
|
|
| warm_times = sorted(times[3:]) |
| if warm_times: |
| med = warm_times[len(warm_times) // 2] |
| p95 = warm_times[int(len(warm_times) * 0.95)] |
| else: |
| med = p95 = 0.0 |
|
|
| print(f"\ndecode steps passing cos_sim > 0.999: {n_pass}/{n_decode}") |
| print(f"per-step latency (warm): median={med:.2f} ms, p95={p95:.2f} ms") |
|
|
| if n_pass == n_decode: |
| print("\n=== ALL DECODE STEPS PASS ===") |
| else: |
| print(f"\n=== FAILED: {n_decode - n_pass}/{n_decode} steps below tolerance ===") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
