Description
Motivation.
Mamba, SSM, and hybrid transformer models are an important path forward towards models that scale linearly with sequence length. vLLM currently supports many models of this class (Jamba, Mamba, Codestral Mamba, Falcon Mamba, Bamba, Zamba2, MinimaxText01, Plamo2), and should continue to maintain excellent support for these models.
The Problem
SSM model generally are less-well supported than transformers in vLLM, and have several deficiencies.
This RFC proposes several improvements (some already in progress) to SSM models, and additionally will serve as an issue tracker.
The major issue is that SSM models not supported in vLLM V1, and should be supported before V0 is deprecated.
In addition:
- SSM state management is a little hacky and is managed by the model definition.
- Since the SSM state is not managed by the block manager, SSM models are incompatible with prefix caching, KV cache offloading, and prefill-decode disaggregation.
- There are major performance issues with chunked prefill.
Proposed Change.
Blockers for SSM and hybrid model support vLLM V1
- Hybrid Allocator: [RFC]: Hybrid Memory Allocator #11382 (initial work is targeted towards sliding-window attention)
- Once the hybrid allocator is landed, extend it to support SSM and hybrid models
- torch.compile support (needed for piecewise CUDA graphs)
Other improvements
- Extend Mamba support beyond CUDA GPUs
- Improve performance for chunked prefill
- Support quantization + tensor parallel in Mamba2 [Bug]: Quantization In MambaMixer2 Not Supported when Tensor Parallel is enabled #14618
Feedback Period.
No response
CC List.
@fabianlim @cyang49 @mzusman @yury-tokpanov
Any Other Things.
No response
Before submitting a new issue...
- Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.