how to return last hidden state

[theAdamColton]

hey this is awesome thank you! It is way faster than the reference implementation

I had some issues with numeric stability training a mamba model with a chunk size of 16, batch size 8, sequence length 256, and dim 32 and state dim 8. I found that adding a small epsilon term of 1e-11 in this division helped with all NaNs. The model seemed to train just the same but I'm not sure what the implications of adding this epsilon is

How do you obtain the last hidden state, like what is in the original reference function; is it just rearrange(hprefix, 'B G D N -> B') ?

[MzeroMiko]

I am so glad that you can pay attention to this tiny project!

1. the hidden state is hprefix.view(B, -1, N); in the new code, it has also been returned.
2. the NaNs may be caused by precision miss: Ats[-1] is too small that (dtBus / Ats).cumsum(dim=0) becomes inf finally, use rAts = Ats / scale may relief from it.
3. I have updated the code which passed all the test benchmark with chunk size < 64, the test benchmark is modified from mamba_ssm

If you have any questions about it, don't hesitate to contact me.

[Yongjie-L]

@MzeroMiko Thanks for the great update. In the "ss_chunk_h1y1" of selective_scan_easyv3, k_div_w = ks[...,None] / ws may also have nans.