Yep, in some sense step is the specialized version of forward that accepts ssm_state and only move by 1 step.

Mamba is really a great work on model size control. On this point, there's another thinking. The mamba vocab embedding size is 502801024(370m) = about 50m parameters. But it really need d_model=1024 to encode a word? My assumption is: The vocab-dim can be reduced to 128(or even 64) to encode a word. The latent_dim(d_model) can still be 1024 or more to contain more info beyond a word. So, could have in_proj and out_proj to mapping between vocab_dim and d_model. Maybe it also can give us:
1, Sharing vocab_embedding between diff-size models: 370m, 790m .......
2, Vocab_embbeding size will be cut from 502801024 to 50280128, or maybe 250000128 will be better.
3, The train could be a little fast because the out_proj can focus on the word encoding info and remove those high-level encoding info(noise for training loss on the last step).
I can't play LLM for resource limit reason. So, this can only be an assumption. Free to ignore it：）

Maybe that version above can help a lot on this point. I don't know whether our owner will take it into consideration. So, i tried it by myself:).

Ok I made it work finally. #477 (With GPU support)