Is it possible to get ssm_states at specific seq len instead of just the last state?

[Lily-Le]

Thanks for the great job!
In the selective scan code, we can only retrieve the SSM state of the last token (if I understand correctly). Is it possible to obtain some intermediate SSM states for analysis in the fast path? If so, do you have any suggestions on how to achieve this? Thanks!

[albertfgu]

You can compute the SSM state using the same logic as used to get the final state. One hack to do this could be to just pass in a prefix of the sequence into the function to get the final state of that prefix, which can be at any desired sequence length.

[0205090923]

@Lily-Le Hello, do you have any suggestions for this now?