Baike.dev
Log in
Back to tool/Back to issues
#1014·mamba

Mamba3.forward() passes a 3-D tensor to step(), which requires 2-D — decode path raises EinopsError

Author: JuiceB0xC0deCreated Aug 8, 2026Updated Aug 12, 2026

Mamba3.forward() routes to self.step(...) when inference_params.seqlen_offset > 0, but passes u through unchanged. u is (batch, seqlen, hidden_dim) at that point, while step() documents and requires (batch, d_model). The sequence axis is never squeezed, so the cached-decode path raises before any kernel runs.

Where

mamba_ssm/modules/mamba3.py, in forward():

python
def forward(self, u, seq_idx=None, cu_seqlens=None, inference_params=None):
    """
    u: (batch, seqlen, hidden_dim)
    Returns: same shape as u
    """
    batch, seqlen, dim = u.shape          # u is 3-D here
    ...
    if inference_params.seqlen_offset > 0:
        out, _, _, _, _ = self.step(u, angle_dt_state, ssm_state, k_state, v_state)
        return out

and in step(), same file:

python
    Args:
        u: (batch, d_model)               # 2-D

The contract mismatch is between two methods of the same class.

Traceback

Observed with mamba-ssm 2.3.2.post1, torch 2.11.0+cu128, CUDA 12.8, on state-spaces/mamba3-siso-187m (6792c27c00f3bb41506db1066dcd1c51bb0f4b02), driving one token at a time with inference_params.seqlen_offset > 0:

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.

B is (batch, seqlen, d_state) because self.in_proj(u) was given a 3-D u; _preprocess expects (batch, d_state).

Note the assert mamba3_step_fn is not None at the top of step() passed — the CuteDSL step kernel was installed and available. This is purely the shape mismatch, not a missing optional dependency.

Reproducer

python
import torch
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel

model = MambaLMHeadModel.from_pretrained(
    "state-spaces/mamba3-siso-187m", device="cuda", dtype=torch.bfloat16
)

class IP:                       # minimal inference-params shim
    def __init__(self):
        self.max_seqlen, self.max_batch_size = 8, 1
        self.seqlen_offset = 1  # > 0 routes forward() to step()
        self.key_value_memory_dict = {}

ip = IP()
for li, layer in enumerate(model.backbone.layers):
    cache = layer.mixer.allocate_inference_cache(1, 8)
    for t in cache:
        t.zero_()
    ip.key_value_memory_dict[li] = cache

model(torch.tensor([[1]], device="cuda"), inference_params=ip)   # raises

The divergence from Mamba2

Mamba3.forward() is consistent with Mamba2.forward() — both pass the 3-D u straight into step() without squeezing:

python
# mamba2.py
if inference_params.seqlen_offset > 0:
    out, _, _ = self.step(u, conv_state, ssm_state)
    return out

The difference is in step(). Mamba2.step() absorbs the sequence axis itself and restores it on return, so forward()'s "Returns: same shape as u" contract holds:

python
# mamba2.py
def step(self, hidden_states, conv_state, ssm_state):
    assert hidden_states.shape[1] == 1, "Only support decoding with 1 token at a time for now"
    zxbcdt = self.in_proj(hidden_states.squeeze(1))     # (B 2D)
    ...
    return out.unsqueeze(1), conv_state, ssm_state

Mamba3.step() does neither: it calls self.in_proj(u) unsqueezed and returns a 2-D out. So Mamba3.step() is the method that departs from the established convention, not forward().

Proposed fix

Matching the Mamba2 convention keeps forward() untouched and preserves its public contract:

python
def step(self, u, angle_state, ssm_state, k_state, v_state, **kwargs):
    """
    Args:
        u: (batch, 1, d_model)
    Returns:
        out: (batch, 1, d_model)
    """
    assert u.shape[1] == 1, "Only support decoding with 1 token at a time for now"
    zxBCdt = self.in_proj(u.squeeze(1))
    ...
    return out.unsqueeze(1), nxt_angle_state, state_out, nxt_k_state, nxt_v_state

The alternative — squeezing at the forward() call site and leaving step() as a 2-D API — also resolves the crash, and has one point in its favour: Mamba3.step() does work correctly today when called directly with a 2-D (batch, d_model) tensor, which is how we have been exercising it. So the 2-D contract is real and functioning; it simply is not what forward() supplies. Either method can move, but they currently disagree.

I have not verified the full decode path end-to-end past the shape error — the step() docstring notes it is "Only tested on H100" and this was observed on an A10G, so there may be further hardware-specific issues behind it. Filing the shape defect on its own rather than claiming a validated fix.

Environment

  • mamba-ssm 2.3.2.post1 (also present on main at time of filing)
  • torch 2.11.0+cu128, CUDA 12.8
  • NVIDIA A10
  • model state-spaces/mamba3-siso-187m @ 6792c27c00f3bb41506db1066dcd1c51bb0f4b02

Source: state-spaces/mamba

View original on GitHubView discussion on GitHub