v1.1.1: re-measured perf on torch 2.13.0; NKI SSD path is inference-only; MIMO compiles; drop CONCATENATION=1; MIMO >512-token limitation; declare einops
v1.1.0: add MIMO decode NKI kernel + fix decode overhead
Jim Burtoft
d1ebb85
New kernel: mamba3_mimo_decode_kernel
- Single-token MIMO decode: state update + rank-R contraction + Y output
- Uses nc_matmul with R on partitions (via nc_transpose) for the BX contraction
- Uses nc_matmul with DSTATE on partitions for the Y output
- All 128 sequential decode steps pass parity vs eager at cos>0.9999
Overhead fix: replaced .contiguous() with plain .transpose() views in the
decode wrappers. Neuron's DMA handles strided sources correctly, avoiding
0.2 ms of materialization per step. This affects both SISO and MIMO decode.
Perf @ target dims (d_model=1024, d_state=128, headdim=64, nheads=32),
batch=1, single-token decode step:
SISO: 3.59 ms eager -> 3.22 ms NKI (1.12x, 0.37 ms saved)
MIMO: 3.93 ms eager -> 3.66 ms NKI (1.07x, 0.27 ms saved)
Files:
new: build/torch-neuron/nki_kernels/mamba3_mimo_decode.py (211 lines)
modified: build/torch-neuron/layers.py (decode paths)
modified: build/torch-neuron/nki_kernels/__init__.p…
docs: add revision=v1.0.0 + trust_remote_code=True to get_kernel() calls
