File size: 3,928 Bytes
e95c403 2a17ff8 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 87 | # 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.
|