I guess mamba.step could be deleted if selective_scan_fn can accept ssm_state as an input param.

[agiwave]

Maybe there are some benifit below: 1, The code could be simplier. 2, The inference could be faster. 3, The inference can accept multi-tokens in this way. There are some reference codes here. https://github.com/agiwave/Mamba/blob/main/Mamba.py