using mamba as an headless "vanilla RNN"

[theodorblackbird]

Thank you for your amazing work,

I'm trying to use Mamba as an drop-in replacement of some RNN in a encoder/decoder architecture, for that I've turned off the logits head. How should I properly get the hidden state during inference ? I'm a bit lost in the decoding logic.

[theodorblackbird]

Ok, I was looking at the wrong place

[JVP15]

I am also looking to use Mamba as a vanilla RNN, but I think I've been looking in the wrong places too for the decoding/hidden state logic. Where did you find it?

[theodorblackbird]

@JVP15
I ended up using it like this :
Instantiate an InferenceParams from utils.generation with seqlen_offset=1.
Then pass it to your MixerModel as a keyword argument inference_params, it caches everything automatically.

[JVP15]

utils.generation, that is what I was missing, thank you!