Skip to content

Mamba2 SSD core is unported: KERNEL-SSM-MAMBA has only the GDN arm, so every hybrid-SSM arch (Nemotron-H/3.5, Bamba, Granite-4, Zamba2, Falcon-H1, Mamba2) is blocked #496

Description

@localai-bot

Gap

KERNEL-SSM-MAMBA is INVENTORIED. Everything that has landed under it is
GDN-specific — the gated-delta-rule recurrence, its causal conv, its
conv-state gather/scatter, and the sm_120 tiling campaign. The row's own matrix
text says so: "these GDN tests do not prove generic Mamba support"
(.agents/kernel-matrix.md:157).

The Mamba2 SSD core is absent. include/vt/ops.h carries
kCausalConv1dFwd / kCausalConv1dUpdate / kCausalConv1dSpecUpdate,
kRmsNormGated, kGdnPrefill / kGdnDecode / kGdnPackedDecode,
kKdaGatedDeltaRule — and no chunked state-space scan and no selective
state update. Gated delta rule is not Mamba2: GDN carries the
(I - beta k kᵀ) delta-removal term, Mamba2's SSD is a diagonally-decayed
gated linear recurrence driven by A_log / dt / D. Different kernel,
same surrounding machinery.

Why now

nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4 (published 2026-08-04,
20.1 GiB, fits one GB10) is NemotronHForCausalLM — 52 layers, 23 Mamba2 +
6 GQA attention + 23 MoE
, with 64 mamba heads x head_dim 64, n_groups=8,
ssm_state_size=128, conv_kernel=4, chunk_size=128. The pinned oracle
already runs it (registry.py:179 -> nemotron_h.py; MambaMixer2 at
nemotron_h.py:373; NemotronHMoEDecoderLayer at nemotron_h.py:317;
ModelOptMixedPrecisionConfig at quantization/modelopt.py:2280), so this is
purely our gap, and the SSD core is the part that is not reachable by
composing what we already have.

This row is the shared unblocker, not a one-model detour. Every
INVENTORIED hybrid-SSM arch in .agents/model-matrix.md waits on the same
kernel: NemotronHForCausalLM (:248), Mamba2ForCausalLM (:234),
GraniteMoeHybridForCausalLM (:214), FalconH1ForCausalLM (:194),
Zamba2ForCausalLM (:280), JambaForCausalLM (:227), Lfm2ForCausalLM
(:229), Lfm2MoeForCausalLM (:230), Plamo2ForCausalLM (:264),
OlmoHybridForCausalLM (:251).

Upstream anchors (pinned oracle 555967922, vLLM 0.26.0.dev0)

  • vllm/model_executor/layers/mamba/mamba_mixer2.pyMambaMixer2 (1170 lines)
  • vllm/model_executor/layers/mamba/ops/ssd_combined.py
  • vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py
  • vllm/model_executor/layers/mamba/ops/ssd_chunk_state.py
  • vllm/model_executor/layers/mamba/ops/ssd_state_passing.py
  • vllm/model_executor/layers/mamba/ops/ssd_bmm.py
  • vllm/model_executor/layers/mamba/ops/mamba_ssm.py (selective_state_update)
  • vllm/model_executor/layers/mamba/ops/layernorm_gated.py
  • vllm/model_executor/layers/mamba/mamba_utils.py

What is already ours and gets reused

MambaSpec + the hybrid KV coordinator (porting-inventory.md:78-79,109),
all three causal-conv1d arms including the spec-decode variant,
kRmsNormGated, kGdnStateGather / kGdnStateScatter, and the het-KV group
construction the gate models already exercise. The conv half of a Mamba2 layer
is done; the recurrence half is not.

Scope of this row

The kernel only: portable host reference for the SSD chunked scan and the
single-token selective state update, unit-gated against a from-first-principles
double-precision reference, then the CUDA arm, then wiring through the existing
MambaSpec state group. The Nemotron-3.5 model port (non-gated relu-squared
MoE grouped GEMM, ModelOpt MIXED_PRECISION per-module loading, the MTP head)
is a separate row on
MODEL-TEXT-nemotron-h-nemotron-hfor-causal-lm and is explicitly out of scope
here.

Spec: .agents/specs/mamba2-ssd.md (linked when the spec commit lands).

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions