Mamba-3 incremental decode broken: `step()` receives `(B, 1, D)` but `_preprocess` expects 2-D — EinopsError in `.generate()` for both official 1.5B checkpoints

Environment

• GPU: NVIDIA GB10 (DGX Spark, sm_121), aarch64
• CUDA 13.0, torch 2.11.0+cu130, Python 3.12
• mamba-ssm: built from git main on 2026-08-22 (--no-build-isolation)
• Checkpoints: state-spaces/mamba3-siso-1.5b and state-spaces/mamba3-mimo-1.5b (unmodified)
• Tokenizer: Llama-3 (NousResearch/Meta-Llama-3-8B)

Repro

import torch
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained("NousResearch/Meta-Llama-3-8B")
model = MambaLMHeadModel.from_pretrained(
    "state-spaces/mamba3-siso-1.5b", device="cuda", dtype=torch.bfloat16)
ids = tok("Warm up.", return_tensors="pt").input_ids.to("cuda")
model.generate(input_ids=ids, max_length=ids.shape[1] + 8, cg=False,
               temperature=1.0, top_k=1)   # same failure with cg=True

Result

  File ".../mamba_ssm/modules/mamba3.py", line 172, in forward
    out, _, _, _, _ = self.step(u, angle_dt_state, ssm_state, k_state, v_state)
  File ".../mamba_ssm/modules/mamba3.py", line 354, in step
    DT, B, C, x, z, trap, A, angles = self._preprocess(
  File ".../mamba_ssm/modules/mamba3.py", line 288, in _preprocess
    B = rearrange(B, "b (r g s) -> b r g s", g=self.num_bc_heads, r=rank)
einops.EinopsError: Error while processing rearrange-reduction pattern "b (r g s) -> b r g s".
 Input tensor shape: torch.Size([1, 1, 128]). Additional info: {'g': 1, 'r': 1}.
 Wrong shape: expected 2 dims. Received 3-dim tensor.

Identical failure on the MIMO checkpoint (shape [1, 1, 512], r=4) and with cg=True (fails inside update_graph_cache → capture_graph).

Analysis

During single-token decode, generation.py calls the mixer with hidden states shaped (B, seqlen=1, D), and Mamba3.forward routes to self.step(u, ...) with u still 3-D. Mamba3._preprocess then projects and rearranges B/C with 2-D patterns ("b (r g s) -> b r g s"), so any seqlen dimension — even 1 — breaks it. Mamba1/Mamba2 step() squeeze/handle the seqlen-1 dim; Mamba3's step path appears to have been written against 2-D inputs only.

Works fine: full-sequence forward() (prefill path) on both checkpoints — we get coherent generations by re-running the full prefix per token, so weights and the chunked kernels are healthy. Only the incremental-decode path is affected.

Two smaller notes

1. The current PyPI release (2.3.2.post1) ships modules/mamba3.py but its create_block still rejects ssm_layer: "Mamba3" ("only support Mamba1 and Mamba2"), so the released package cannot load the released checkpoints at all — a post-release including the registry wiring would save users a source build.
2. On sm_121 we also see the tilelang-compiled MIMO forward kernel miscompile to all-zero outputs at specific sequence lengths (filed separately against tile-ai/tilelang, with a cross-link — happy to add details here if useful).