Fixes and Enhancements for Mamba Inference and Reference Implementations

[mohiuddin-khan-shiam]

This pull request addresses several bugs and limitations within the Mamba codebase, primarily aimed at improving inference robustness in the Mamba2 module and increasing the accuracy of reference implementations.

Key changes include:

• In mamba_ssm/modules/mamba2.py:
  • Resolved an issue in _get_states_from_cache to correctly handle dynamic batch sizes during inference, ensuring proper state re-initialization when batch sizes change.
  • Removed the batch == 1 assertion in the forward method for variable-length sequence inference, enabling batched processing for these inputs.
  • Updated the fallback path in the step method to support ngroups > 1, allowing grouped SSM inference even if Triton kernels are not available.
• In mamba_ssm/ops/selective_scan_interface.py:
  • Added optional RMS Normalization for B, C, and delta tensors within mamba_inner_ref to better match the main MambaInnerFn's behavior and improve numerical consistency.
  • Corrected a shape comment in selective_scan_ref for clarity.
• In mamba_ssm/models/mixer_seq_simple.py:
  • Removed a redundant comment in the _init_weights function.
• In mamba_ssm/utils/hf.py:
  • Addressed a bug in load_state_dict_hf to ensure correct dtype conversion and device placement when loading Hugging Face model weights.

These modifications enhance the stability, flexibility, and correctness of the Mamba library.

[zju-sqs]

This code only support one step decoding now (in selective_state_update.py) which makes it hard to train streaming model, eg. chunk wise asr, it needs to pass ssm_states between chunks and update through [bs,seq_len,hidden_dim] tensor,not [[bs,1,hidden_dim].Does your code suppot seq_len >=2 updating ssm_states?