Module nemotron_mamba2

Module nemotron_mamba2 

Source
Expand description

2026-09-25: Nemotron-H Mamba-2 SSM layer (no FFN), implementing TransformerLayer.

Decode, per token (trait_impl.rs):

  1. RMS norm, the layer input copied to residual
  2. in_proj GEMV into one row [z (d_inner) | xBC (d_xbc) | dt (num_heads)]
  3. causal conv1d update on xBC, with bias and a fused SiLU
  4. xBC split into x, B, C
  5. Mamba-2 SSM decode (state update and output y)
  6. gated RMS norm of y with gate z, in the order of the resolved gated_rms_norm kernel: the gb10 Nemotron trees’ kernels gate first, y * silu(z), then normalise each d_inner / n_groups slice
  7. out_proj GEMV
  8. residual add into hidden

Prefill runs the same steps over all tokens (prefill.rs, projection arms in prefill_proj.rs).

Owner: model-arch (Nemotron-H). Invariants: none beyond the types.

Structs§

NemotronMamba2Layer