Mamba3 MIMO forward doesn't work with seq_length=1

[stergiosba]

Hey mamba team!

Thank you for the awesome project and the insights on the last paper (Mamba3). I found that the following usage of Mamba3 raises an unnexpected error. This is quite common for what I am doing so I went digging. The problem is the forward() function of Mamba3 cannot work with sequences of length=1 in MIMO mode.

In one sentence, while the step() function works fine when the input has shape (Batch,Dim), the forward() function does not work when the shape is (Batch,1,Dim).

Note: Mamba 3 SISO, Mamba and Mamba2 all work. Also if length>1 MIMO mode works fine.

import numpy as np
import random
import os
from mamba_ssm import Mamba3
from mamba_ssm.utils.generation import InferenceParams
import torch
batch, length, dim = 2, 1, 768

def make_deterministic(seed: int = 42):
    # 1. Standard Python and NumPy seeding
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    
    # 2. PyTorch CPU and GPU seeding
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed) # For multi-GPU setups
    
    # 3. Configure cuDNN backend behavior
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    
    # 4. Enforce deterministic algorithms globally
    # (Throws an error if an operation doesn't have a deterministic alternative)
    torch.use_deterministic_algorithms(True)
    
    # 5. Optional: Specifically for CUDA 10.2+ atomic operations
    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" # or ":16:8"

make_deterministic()
x = 0.1*torch.ones(batch, length, dim).to(torch.bfloat16).to("cuda")
model = Mamba3(
    # This module uses roughly 6 * d_model^2 parameters
    d_model=dim, # Model dimension d_model
    d_state=128,  # SSM state size
    headdim=64, # SSM headdim
    is_mimo=True, # Use MIMO mode
    mimo_rank=4, # MIMO rank when is_mimo=True
    chunk_size=16, # 64/mimo_rank if x is in bf16, else 32/mimo_rank
    is_outproj_norm=False, # Additional post SSM norm
    layer_idx=0,
    dtype=torch.bfloat16,
).to("cuda")
y = model(x)

This raises:

Traceback (most recent call last):
  File "/home/horclab/mamba/d.py", line 59, in <module>
    y = model(x)
  File "/home/horclab/mamba/.venv/lib/python3.14/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/home/horclab/mamba/.venv/lib/python3.14/site-packages/torch/nn/modules/module.py", line 1789, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/horclab/mamba/mamba_ssm/modules/mamba3.py", line 210, in forward
    y = mamba3_mimo_combined(
        Q=C,
    ...<20 lines>...
        outproj_norm_eps=self.norm.eps if self.fuse_pregate_headwise_norm else 1e-5,
    )
  File "/home/horclab/mamba/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo.py", line 341, in mamba3_mimo
    return _Mamba3Function.apply(
           ~~~~~~~~~~~~~~~~~~~~~^
        Q,
        ^^
    ...<20 lines>...
        outproj_norm_eps,
        ^^^^^^^^^^^^^^^^^
    )
    ^
  File "/home/horclab/mamba/.venv/lib/python3.14/site-packages/torch/autograd/function.py", line 596, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
           ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/home/horclab/mamba/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo.py", line 102, in forward
    Out, Final_SSM_State, Final_K = mamba_mimo_forward(
                                    ~~~~~~~~~~~~~~~~~~^
        Q, K, V, Q_bias, K_bias, MIMO_V, MIMO_Out,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    ...<7 lines>...
        outproj_norm_eps=outproj_norm_eps,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/home/horclab/mamba/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo_fwd.py", line 547, in mamba_mimo_forward
    kernel( q,
    ~~~~~~^^^^
            k,
            ^^
    ...<12 lines>...
            k_final
            ^^^^^^^
            )
            ^
  File "/home/horclab/mamba/.venv/lib/python3.14/site-packages/tilelang/jit/kernel.py", line 207, in __call__
    return self.torch_function(*args, **kwds)
           ~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^
  File "/home/horclab/mamba/.venv/lib/python3.14/site-packages/tilelang/jit/adapter/tvm_ffi.py", line 244, in func
    executable(*tensor_list)
    ~~~~~~~~~~^^^^^^^^^^^^^^
  File "/home/horclab/mamba/.venv/lib/python3.14/site-packages/tvm_ffi/module.py", line 297, in __call__
    return self.main(*args)
           ~~~~~~~~~^^^^^^^
  File "python/tvm_ffi/cython/function.pxi", line 929, in tvm_ffi.core.Function.__call__
  File "<unknown>", line 0, in _TAIL_CALL_CALL_FUNCTION_EX.llvm.641356221097560969
RuntimeError: kernel mamba_mimo_fwd_kernel input DA_CS strides[2] expected 1, but got 24

So I saw that you guys implemented the Mamba3 MIMO kernel in Tilelang so I focused my attention there. It looks like changing the lines in the permalink:

mamba/mamba_ssm/ops/tilelang/mamba3/mamba3_mimo.py

Lines 65 to 70 in ed6ce09

	(Q, K, V, ADT, DT, Trap, Q_bias, K_bias, MIMO_V, MIMO_Z, MIMO_Out, Out_Norm_Weight, Angles, D, Z) = tuple(
	t.contiguous() if t is not None else None
	for t in (
	Q, K, V, ADT, DT, Trap, Q_bias, K_bias, MIMO_V, MIMO_Z, MIMO_Out, Out_Norm_Weight, Angles, D, Z,
	)
	)

to the following:

        def _ensure_tilelang_contiguous(t: Optional[Tensor]) -> Optional[Tensor]:
            """Make tensor truly contiguous with last-stride=1.

            PyTorch's .contiguous() considers a tensor with a size-1 trailing
            dimension as contiguous even when stride[-1] != 1 (e.g. after a
            transpose).  TileLang kernels require the physical stride to be 1,
            so we force a re-allocation when that invariant is violated.
            """
            if t is None:
                return None
            t = t.contiguous()
            if t.stride(-1) != 1:
                t = torch.empty(t.shape, device=t.device, dtype=t.dtype).copy_(t)
            return t

        (Q, K, V, ADT, DT, Trap, Q_bias, K_bias, MIMO_V, MIMO_Z, MIMO_Out, Out_Norm_Weight, Angles, D, Z) = tuple(
            _ensure_tilelang_contiguous(t)
            for t in (
                Q, K, V, ADT, DT, Trap, Q_bias, K_bias, MIMO_V, MIMO_Z, MIMO_Out, Out_Norm_Weight, Angles, D, Z,
            )
        )

This solved the problem for me locally. Am I missing anything here? Thanks again for this awesome project.

[OMGYFX]

Buddy why you use seq len 1

[stergiosba]

Do you have something constructive to say or are you going to be a clown?

[OMGYFX]

My son, I have at least 100 times more knowledge than you. This is a memory layout (stride) mismatch, not a model bug. Read the code and learn.

[stergiosba]

Keep it professional. Let's wait for an adult to answer yes?

[Chessing234]

Opened #994 for the stride(-1)!=1 contiguous case.

[Chessing234]

PR with the stride(-1)==1 contiguous fix for MIMO forward at seqlen=1.