| # Examples |
|
|
| Runnable samples demonstrating the `mamba3-neuron-kernels` package. Each script |
| is self-contained and prints its own results. |
|
|
| ## Requirements |
|
|
| - **AWS Neuron trn2.3xlarge** instance (or larger) |
| - **Neuron SDK Beta 3+**: torch-neuronx 2.11.3+, NKI 0.4.0+, PyTorch 2.11+ |
| - **`kernels` library**: `pip install kernels>=0.15.2` |
| - **Environment**: `export TORCH_NEURONX_ENABLE_CONCATENATION=1` (recommended, ~3-8% speedup) |
|
|
| ## Running |
|
|
| Each script prints its own status. If run from a local clone of this repo, the |
| scripts import from `../build/torch-neuron/`. When the repo is loaded via |
| `kernels.get_kernel("jburtoft/mamba3-neuron-kernels", revision="v1.0.0", trust_remote_code=True)`, the same code works. |
|
|
| ```bash |
| # One-time setup (on your trn2 instance) |
| source $HOME/workspace/native_venv/bin/activate |
| export TORCH_NEURONX_ENABLE_CONCATENATION=1 |
| |
| # Run any sample |
| cd examples/ |
| python 01_smoke_test.py |
| python 02_forward_parity.py |
| python 03_decode_generation.py |
| python 04_backward_training.py |
| python 05_benchmark.py # short version (seqlen 64, 256) |
| python 05_benchmark.py --full # includes seqlen 512, 1024 |
| ``` |
|
|
| ## What each sample does |
|
|
| | Script | Purpose | Runtime | |
| |---|---|---| |
| | `01_smoke_test.py` | Loads the kernel, instantiates a mixer, runs one forward. First-run compile takes ~15 s; warm call is a few ms. | ~30 s | |
| | `02_forward_parity.py` | Runs SISO and MIMO forward with and without the NKI kernel on identical weights, compares outputs. Both should agree at cos_sim > 0.999. | ~60 s | |
| | `03_decode_generation.py` | Prefills a 64-token context, then runs 128 single-token decode steps. Compares NKI vs eager decode at each step and reports cache correctness + per-step latency. | ~120 s | |
| | `04_backward_training.py` | Trains SISO and MIMO mixers for 10 steps each on random data with MSE loss. Verifies loss decreases (proves autograd works end-to-end). | ~90 s | |
| | `05_benchmark.py` | Reproduces the perf tables in the README. Reports forward latency for both modes across seqlens, plus decode step time. | ~90 s short / ~5 min full | |
|
|
| ## Expected output (as of v1.0.0) |
|
|
| Target dims: `d_model=1024, d_state=128, headdim=64, nheads=32`, batch=1, trn2.3xlarge LNC=2. |
|
|
| ### Forward (with `TORCH_NEURONX_ENABLE_CONCATENATION=1`) |
| |
| | Path | seqlen=64 | seqlen=256 | |
| |---|---|---| |
| | SISO + torch.compile + NKI SSD | ~1.7 ms | ~3.6 ms | |
| | MIMO + NKI SSD (eager) | ~6.4 ms | ~10.3 ms | |
| |
| ### Decode step (single-token) |
| |
| | Path | median | p95 | |
| |---|---|---| |
| | SISO + NKI SSD | ~4.0 ms | ~4.5 ms | |
| | MIMO + NKI SSD | ~4.5 ms | ~5.5 ms | |
| |
| ### Batched throughput (batch=32, seqlen=256, eager + NKI SSD) |
| |
| | Path | per-sample latency | |
| |---|---| |
| | SISO | ~1.78 ms/sample | |
| | MIMO | ~5.76 ms/sample | |
| |
| ## Troubleshooting |
| |
| **`RuntimeError: NKI kernel requires chunk_size=64`** or similar assertion in `NeuronMamba3Mixer.__init__`: |
| The NKI kernels are compiled against the reference target config. Use these exact values: |
| - SISO: `chunk_size=64, mimo_rank=1` |
| - MIMO: `chunk_size=16, mimo_rank=4` |
| - Both require: `d_state=128, headdim=64` |
|
|
| For non-reference configs, use `use_nki_ssd=False` to fall back to the eager path. |
|
|
| **First-call latency is very slow (>15 seconds)**: |
| That's the initial NKI compile. Every unique input shape triggers a fresh compile. Subsequent calls with the same shape use the cached NEFF. |
|
|
| **MIMO forward with `torch.compile` is much slower than eager**: |
| Known regression -- see [ticket 61](https://github.com/aws-neuron/aws-neuron-sdk/) filing. Don't wrap the MIMO mixer with `torch.compile` until this is fixed. Use eager only for MIMO. |
|
|
| **Silent gradient bias in MIMO training**: |
| The mixer's forward internally uses `torch.cat` instead of `F.pad` for the beta-term shift, to work around a Neuron autograd bug (`F.pad + einsum` chain produces 2% gradient error on device). This is a workaround the mixer applies transparently. See the `kernels/mixer.py` comment for details. |
|
|
