In ssm.py, the selective_scan method runs a Python for-loop over the sequence length:
for t in range(length):
h = dA[:, :, :, t] * h + dB[:, :, :, t] * x[:, :, t:t+1]
...
outputs.append(y_t)
For a sequence of length 1024 with batch=8 and d_model=256, this is 1024 iterations of Python-level loop with tensor operations inside each step. It is extremely slow compared to what Mamba is supposed to be.
The whole point of the selective scan is that it should run as a parallel associative scan (like a prefix sum) not as a sequential recurrence. The parallel_scan function is exported in init.py but it is not used here. Either use it or document clearly that this is a reference/educational implementation and not performance-oriented.
In ssm.py, the selective_scan method runs a Python for-loop over the sequence length:
for t in range(length):
h = dA[:, :, :, t] * h + dB[:, :, :, t] * x[:, :, t:t+1]
...
outputs.append(y_t)
For a sequence of length 1024 with batch=8 and d_model=256, this is 1024 iterations of Python-level loop with tensor operations inside each step. It is extremely slow compared to what Mamba is supposed to be.
The whole point of the selective scan is that it should run as a parallel associative scan (like a prefix sum) not as a sequential recurrence. The parallel_scan function is exported in init.py but it is not used here. Either use it or document clearly that this is a reference/educational implementation and not performance-oriented.