Mamba3 step() is non-deterministic and does not match forward()

Author: H-E3Created Aug 26, 2026Updated Sep 17, 2026

Question

I have noticed there is an output difference between calling step() and forward().

When using the same model and feeding in the exact same input sequence, invoking Mamba3.forward() directly versus running multiple Mamba3.step() calls in a loop yields outputs that diverge noticeably, with a maximum difference value close to 1.

More importantly, running the identical step() loop multiple times will return inconsistent results even without modifying inputs or parameters.

By comparison, the equivalent test case implemented for Mamba2 behaves reliably; its largest absolute difference sits around 5e‑5.

I would like to understand what triggers this inconsistency. I also wonder whether this bug could stem from issues within my runtime environment.

Reproduction

my test code is as following:

python

import torch
from mamba_ssm import Mamba3, Mamba2

torch.manual_seed(0)
dtype = torch.float32
m1 = (
    Mamba3(d_model=64, d_state=64, headdim=16, chunk_size=16, is_mimo=False, dtype=dtype)
    .cuda()
    .eval()
)
m2 = (
    Mamba2(d_model=64, d_state=64, headdim=16, chunk_size=16, dtype=dtype)
    .cuda()
    .eval()
)

batch, seqlen, d_model = 1, 10, 64
seq = torch.randn(batch, seqlen, d_model, dtype=dtype).cuda()     # [B, L, D]

with torch.no_grad():    
    # -------- model 1
    forward_out_1 = m1(seq)                    # [B, L, D]
    angle_dt_state, ssm_state, k_state, v_state = m1.allocate_inference_cache(
        batch_size=batch, max_seqlen=None, dtype=dtype, device=seq.device
    )  # allocate cache for streaming inference
    steps_1 = torch.empty(seqlen, batch, d_model, dtype=dtype).cuda()  # [L, B, D]
    for t in range(seqlen):  # simulate streaming inference, one step at a time
        # step() u/out: (batch, d_model)
        u = seq[:, t, :]   # [B, D]
        # print(u.shape)
        step_out, angle_dt_state, ssm_state, k_state, v_state = m1.step(
            u, angle_dt_state, ssm_state, k_state, v_state
        )        
        steps_1[t] = step_out

    # -------- model 2
    forward_out_2 = m2(seq)                    # [B, L, D]
    conv_state, ssm_state = m2.allocate_inference_cache(
        batch_size=batch, max_seqlen=None, dtype=dtype, device=seq.device
    )  
    steps_2 = torch.empty(seqlen, batch, d_model, dtype=dtype).cuda()  # [L, B, D]
    for t in range(seqlen): 
        # step() u/out: (batch, 1, d_model)
        u = seq[:, t, :]   # [B, D]
        u = u.unsqueeze(1)  # [B, 1, D]
        # print(u.shape)
        step_out, conv_state, ssm_state = m2.step(
            u, conv_state, ssm_state
        )        
        steps_2[t] = step_out.squeeze(1)  # [B, D]

# [L, B, D] -> [B, L, D]
streamed_1 = steps_1.transpose(0, 1)
streamed_2 = steps_2.transpose(0, 1)

diff_1 = (forward_out_1 - streamed_1).abs()
diff_2 = (forward_out_2 - streamed_2).abs()
print(f"Mamba3 step vs forward: max|diff|={diff_1.max():.2e}")
print(f"Mamba2 step vs forward: max|diff|={diff_2.max():.2e}")

I think that after initializing the inference cache to zero, sequential calls to step() should produce outputs numerically close to the corresponding outputs from forward(). I run the test code 3 times, but get 3 different outputs:

// run 1
Mamba3 step vs forward: max|diff|=1.09e+00
Mamba2 step vs forward: max|diff|=5.07e-05

// run 2
Mamba3 step vs forward: max|diff|=1.01e+00
Mamba2 step vs forward: max|diff|=5.07e-05

// run 3
Mamba3 step vs forward: max|diff|=1.11e+00
Mamba2 step vs forward: max|diff|=5.07e-05

I also tested bfloat16 . The mismatch became bigger:

Mamba3 step vs forward: max|diff|=1.62e+00
Mamba2 step vs forward: max|diff|=7.81e-03

Environment

here is my environment configurations. i use the pre‑compile package to install mamba‑ssm, but not install from source.

OS: Ubuntu 22.04
GPU: NVIDIA GeForce RTX 4090
NVIDIA driver: 590.48.01
Python: 3.11.15
PyTorch: 2.10.0+cu128
CUDA used to build PyTorch: 12.8
mamba-ssm: 2.3.2.post1
Mode: SISO (is_mimo=False)

Source: state-spaces/mamba

View original on GitHub View discussion on GitHub