Make Mamba Export Compatible with ONNX

[AyoubMDL]

Feature request

When exporting Mamba to ONNX while using the cache, an if-else statement prevents proper inference and alignment with the Hugging Face model. Specifically, the issue arises due to the condition at this line and the alternative execution path at this line.

A potential solution is to ignore the cache position and modify the update_conv_state method in Mamba’s cache. Instead of the current implementation at this line, always update the last position with the new state after the roll:

conv_state[:, :, [-1]] = new_conv_state.to(device=conv_state.device, dtype=conv_state.dtype)

This ensures that the else branch is always taken in Mamba’s forward pass, making the model exportable and inference-compatible. Notably, this approach is already used in Mamba2Cache, as seen in this implementation.

Would it be possible to align Mamba’s behavior with Mamba2Cache to support ONNX export properly?

Motivation

Mamba’s current conv state update prevents proper ONNX export due to an if-else condition that disrupts inference consistency.

Your contribution

Proposed a fix that updates the conv state without using the cache position.

[Rocketknight1]

cc @SunMarc @MekkCyber

[SunMarc]

cc @gante