jburtoft/mamba2-ssd-neuron-kernels

Mamba-2 SSD NKI Kernel for AWS Neuron

NeuronCore-optimized NKI kernel implementing the Mamba-2 Structured State-space Duality (SSD) algorithm. Replaces the Mamba2Mixer module in transformers with a hardware-accelerated version that performs chunk-wise parallel scan on NeuronCore Tensor and Vector engines.

Why This Kernel Exists

On Neuron hardware, the standard Mamba2Mixer falls back to a naive O(L^2) PyTorch implementation (torch_forward()) since the Triton-based mamba_chunk_scan_combined kernel is CUDA-only. This NKI kernel provides the efficient O(L) chunk-parallel algorithm directly on NeuronCore.

Performance (Granite 4.0 Dimensions, v2.0)

Profiled on trn2.3xlarge (PyTorch Native Beta 3, LNC=2). Isolated ssd_kernel wall-clock, batch=1, nheads=32, headdim=64, dstate=128, chunk_size=128.

Seqlen	v1 (pre-optimization)	v2.0 (this release)	Speedup
128	455.0 μs	428.0 μs	1.06x
256	670.4 μs	578.9 μs	1.16x
512	1164.0 μs	946.8 μs	1.23x
1024	2239.0 μs	1722.1 μs	1.30x
2048	3953.9 μs	3011.9 μs	1.31x

Speedup grows with seqlen because per-chunk savings compound. At S=2048 (16 chunks), the kernel is 1.31x faster end-to-end than v1.

Cumulative engine changes (S=256)

• VectorE: -34.2% (868 fewer instructions, 150 ms saved on VectorE alone)
• PE (Tensor): +20.6% (Change 1 added the triangular-matmul cumsum + ones-broadcast; large fp32 matmuls are cheap on PE)
• Parity: bit-identical to v1 (cosine sim > 0.9999, max_abs_diff < 3.3e-6 on all 8 test configurations)

Full engine breakdown, per-task profile diffs, and optimization discipline notes are in the mamba2-ssd-optimize development project.

Optimizations Applied (v2.0)

Sourced from a profile-guided optimization spec (projects/mamba2-ssd-optimize):

1. Delete redundant B_f32 copy (Change 3): removed a pointless fp32→fp32 tensor_copy that was pre-existing dead code.
2. Pass triu as constant (Change 5): frameworks now build triu (upper-triangular incl diagonal) directly. Kernel deletes the per-head nc_transpose(causal_sb) + tensor_copy sequence in favor of one DMA. -28.6% Sync, -3.6% PE at Granite S=256.
3. Cumsum via triangular matmul (Change 1, headline): the tensor_tensor_scan cumsum on VectorE (plus its two-copy nc_transpose-based partition-conversion tail) is replaced with a single PE matmul triu^T @ log_decay. Also hoists the per-head A_Q broadcast outside the chunk loop, and computes exp_cs_last via ones_QQ @ log_decay (deletes the (1,1) activation + 4-iter shuffle broadcast). -8 to -16% wall-clock.
4. Fuse PSUM reads into scale/mask (Change 2, 4 sub-changes): VectorE can consume a PSUM operand directly, so the intermediate tensor_copy(sbuf, psum) before each of CB_T mask, Y_intra scale, Y_off scale, and chunk_state add is now unnecessary. Deletes 4 tensor_copies per chunk (including one (Q, Q)). Additional -5 to -7% wall-clock.
5. activation(scale=-1.0) (Change 4): the neg_cs_p = tensor_scalar(cs_p, multiply, -1.0) followed by activation(exp, neg_cs_p) collapses into activation(exp, cs_p, scale=-1.0). Additional -7% at long seqlens.

Changes 6 (pre-transpose B in framework) and 7 (bf16 operands for the 3 large matmuls) were attempted and reverted — see the mamba2-ssd-optimize project's Task 011 and Task 012 for the "measure, don't assume" analysis. Short summary: change 6's per-chunk PE savings were outweighed by extra HBM cost at long seqlens; change 7's cast overhead landed on VectorE (critical path) and canceled the PE win. Both may pay off in a future NKI variant with fused cast+matmul or a different HBM/DMA regime.

Correctness

Two runtime guards in NeuronMamba2MixerForward.forward:

• n_groups == 1: the kernel indexes B and C_T with no group axis (B/C shared across heads). Asserted before the kernel call with a clear error message.

Numerical caveats (documented in ssd_kernel source):

• Factored decay form: Y_intra = exp(cs_i) * [(CB^T o mask) @ (exp(-cs_j) * dtx_j)]. exp(-cs_j) grows across the chunk. For Granite 4.0 dt*A ranges this stays within fp32 (peaks ~e^6-e^10). Models with much larger decay rates or chunk_size > 128 may overflow; switch to the stable segsum form for those cases.

Supported Models

Any Mamba-2 model using Mamba2Mixer in transformers with n_groups == 1:

Model	Notes
IBM Granite 4.0	Primary target: nheads=128 (TP to 32), headdim=64, dstate=128
state-spaces/mamba2-*	Standard Mamba-2 configs, n_groups == 1 variants

Requirements:

• chunk_size == 128 (hardcoded for NeuronCore tile alignment)
• dstate <= 128 (fits in single tile)
• headdim <= 128 (fits in single tile)
• n_groups == 1 (asserted at forward time)
• Sequence length is padded to a multiple of 128 by the wrapper.

Requirements

• AWS Neuron SDK 2.30+ or PyTorch Native Beta 3+ (NKI 0.4.0+, neuronx-cc 2.25+)
• transformers >= 4.45 with Mamba2Mixer support (>= 5.10.0.dev0 recommended for KernelConfig)
• kernels library (HuggingFace kernel loading)

Usage

from transformers import AutoModel, KernelConfig

model_id = "ibm-granite/granite-4.0-h-tiny"  # Any Mamba-2 model with n_groups=1

kernel_config = KernelConfig({
    "Mamba2Mixer": "jburtoft/mamba2-ssd-neuron-kernels:NeuronMamba2MixerForward",
})

model = AutoModel.from_pretrained(
    model_id,
    kernel_config=kernel_config,
    device_map="neuron",
)

v2.0 API change: Direct callers of ssd_kernel(...) -- if any exist outside the NeuronMamba2MixerForward wrapper -- must pass causal_mask as an upper-triangular (.triu()) matrix instead of .tril(). Callers going through KernelConfig are unaffected (the wrapper handles the change internally). See Change 5 above and the mask-direction verification comment in ssd_kernel.

Architecture

NeuronMamba2MixerLayout     - Weight structure (matches Mamba2Mixer)
NeuronMamba2MixerForward    - Forward pass:
  1. in_proj                 - Input projections (PyTorch)
  2. conv1d                  - Causal convolution (PyTorch)
  3. n_groups==1 assert      - Runtime correctness guard
  4. Build triu_mask         - Upper-tri causal mask (Change 5)
  5. Pad to chunk_size       - Framework pads to multiple of 128
  6. ssd_kernel              - @nki.jit SSD scan (NKI on NeuronCore)
  7. Unpad                   - Truncate y back to seq_len
  8. norm + gate + out_proj  - Output processing (PyTorch)
ssd_kernel                   - Optimized NKI kernel:
  - Chunk-wise parallel scan (O(L) complexity)
  - Per-head: A_Q broadcast, ones_QQ constant, triu_sb DMA
  - Per-chunk:
    * Cumsum via triangular matmul (triu^T @ log_decay)
    * Fused PSUM reads on CB_T, Y_intra, Y_off, chunk_state
    * Fused exp(-cs) via activation(scale=-1.0)
    * Option B exp_cs_last: ones_QQ @ log_decay + activation

Algorithm

The SSD (Structured State-space Duality) decomposes the Mamba-2 recurrence into:

1. Intra-chunk parallel: Within each 128-token chunk, compute Y = (C*B^T * mask) @ (dt*x) using Tensor Engine matmuls
2. Inter-chunk sequential: Propagate SSM state between chunks with exponential decay via state *= exp(cs_last), state += B^T @ (dtx * decay_to_end)
3. State-to-output: Apply accumulated state contribution via C^T @ state matmul

This achieves O(L) complexity (linear in sequence length) while maximizing Tensor Engine utilization for the matrix multiplications within each chunk.

Development History

• v1 (2026-06-16): initial release from hf-kernels Task 009. Forked from nkilib.experimental.scan.ssd with two profile-guided optimizations (C pre-transposed, shared exp_cs_last broadcast). 8.2% faster than nkilib baseline at Granite S=256.
• v2.0 (2026-07-16): full profile-guided optimization pass from the mamba2-ssd-optimize project. 5 spec changes applied (deleting VectorE work and moving cumsum to PE). 6-24% wall-clock improvement over v1 depending on seqlen. Bit-identical output vs the CPU fp32 reference on all 8 test configurations. Two additional changes (bf16 operands, pre-transpose B) were attempted and reverted after wall-clock measurement.

License

Apache 2.0