4 days ago
-
README and comments: update install instructions and known limitations; remove internal references
-
docs: add revision=v1.0.0 + trust_remote_code=True to get_kernel() calls
Required by kernels >= 0.15. Updates:
- README.md: 'Direct usage' snippet
- build/torch-neuron/__init__.py: module docstring
- examples/_loader.py: Hub fallback
- examples/README.md: install note -
v1.0.0: initial release
Full-mixer NKI kernels for Mamba-3 on AWS Neuron:
- SISO SSD (chunk_size=64, R=1) with torch.compile
- MIMO SSD (chunk_size=16, R=4) eager
- SISO decode step
Layered structure (constants/ops/masks/layers + nki_kernels/):
- 12 files under build/torch-neuron/
- 5 sample scripts in examples/ (smoke, parity, decode, training, benchmark)
- All samples verified end-to-end on trn2.3xlarge LNC=2
Performance @ target dims (d_model=1024, d_state=128, headdim=64, nheads=32):
- SISO forward: 1.95 ms @ seq=64, 3.30 ms @ seq=256
- MIMO forward: 6.17 ms @ seq=64, 10.29 ms @ seq=256
- SISO decode: 3.67 ms/step | MIMO decode: 4.02 ms/step
Parity vs eager: cos_sim >= 0.999999 on all 8 test configs
Training: SISO 124x loss reduction, MIMO 447x loss reduction on synthetic tasks
