| """01_smoke_test.py -- verify the kernel loads and runs. |
| |
| Loads NeuronMamba3Mixer from the HF Hub kernel (or a local clone), instantiates |
| a small SISO mixer, and runs a single forward on device="neuron". |
| |
| Requirements: |
| - trn2.3xlarge (or larger) AWS Neuron instance |
| - Neuron SDK Beta 3+ with PyTorch Native (torch-neuronx 2.11.3+) |
| - kernels library >= 0.15.2 |
| |
| Run: |
| python 01_smoke_test.py |
| """ |
|
|
| import time |
| import torch |
| from _loader import load_kernel |
|
|
|
|
| def main(): |
| print("=" * 60) |
| print("mamba3-neuron-kernels smoke test") |
| print("=" * 60) |
|
|
| mamba3 = load_kernel() |
| print(f"kernel version: {mamba3.__version__}") |
|
|
| torch.manual_seed(0) |
| |
| mixer = mamba3.NeuronMamba3Mixer( |
| d_model=1024, |
| d_state=128, |
| headdim=64, |
| chunk_size=64, |
| mimo_rank=1, |
| use_nki_ssd=True, |
| ).to("neuron") |
| print(f"instantiated SISO mixer: d_model=1024, nheads={mixer.nheads}, NKI kernels enabled") |
|
|
| u = torch.randn(1, 128, 1024, device="neuron") |
|
|
| print("\nrunning first forward (includes kernel compile)...") |
| t0 = time.perf_counter() |
| y, cache = mixer(u) |
| try: |
| torch.neuron.synchronize() |
| except Exception: |
| pass |
| first_ms = (time.perf_counter() - t0) * 1000 |
| print(f" first call: {first_ms:.0f} ms") |
|
|
| print("running warm forward...") |
| t0 = time.perf_counter() |
| y, cache = mixer(u) |
| try: |
| torch.neuron.synchronize() |
| except Exception: |
| pass |
| warm_ms = (time.perf_counter() - t0) * 1000 |
| print(f" warm call: {warm_ms:.2f} ms") |
|
|
| assert y.shape == (1, 128, 1024), f"y shape unexpected: {y.shape}" |
| assert cache.ssm_state.shape == (1, 32, 64, 128), f"ssm_state shape unexpected: {cache.ssm_state.shape}" |
|
|
| print(f"\ny.shape = {y.shape}") |
| print(f"cache.ssm_state.shape = {cache.ssm_state.shape}") |
| print(f"cache.prev_Bx.shape = {cache.prev_Bx.shape}") |
| print(f"cache.cum_angle.shape = {cache.cum_angle.shape}") |
|
|
| print("\n=== SMOKE TEST PASS ===") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
