[V1] [Hybrid] Support using float32 for state in Hybrid Models (Mamba2, Mamba1, Minimax)

[tdoublep]

Purpose

This PR adds support for customizing the dtype of the state stored for models based on mamba/mamba2/linear attention. It introduces a new CLI argument mamba_ssm_cache_dtype which is by default set to auto but can also be set to float32. If set to float32, for mamba2 the temporal (ssm) state will be stored in fp32, but the conv state will take whatever value is set by kv_cache_dtype (e.g., if left as auto it will take the model dtype). For mamba1 the same behaviour is expected, but there are some kernel changes needed to enable fp32 state (cc @asafgardin) for mamba1 so I'm raising if that gets set to float32 for now. For linear attention, there is only one state tensor so that case is simpler.

Note: for mamba2 we need to take special care when layout out the two state tensors in GPU memory since they have different dtypes, and the calculation of the page size also needs to change accordingly.

This PR builds on previous unmerged PRs:

• [V1][Hybrid][Backend] Allow different data types in Hybrid Cache Manager #22544 @bringlein
• [Model] [V1] Mamba models - Support FP32 mamba cache #22913 @danielafrimi

PS This PR is not as big as it looks :) Most of the changes in the modeling code are the same.

Test Plan

I added a new test to test_hybrid that verifies the behaviour for V0 and V1 are the same, and that the output matches hf transformers. Transformers doesn't actually have support for this yet, so I'm hoping for now that the output is basically the same for a short prompt.

Test Result

The test passes. I've also tested a lot of other models locally but don't want to bloat the CI.

(Optional) Documentation Update

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

closes #22544

closes #22913

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[gemini-code-assist]

Code Review

This pull request introduces support for using float32 for the Mamba model's state, which is particularly useful for hybrid models. The changes are well-structured, adding a new configuration option mamba_ssm_cache_dtype and plumbing it through the system from argument parsing to the cache and model execution layers. The implementation correctly handles the new data type for both convolutional and SSM states in Mamba layers. The tests have been updated to cover this new functionality. My main feedback is to remove some debugging print statements left in the code.

[tomeras91]

Looks great! Left a few nit comments

[cyang49]

@tdoublep I think this change broke plamo2 (vLLM V0) and possibly phi4flash too

lm_eval --model vllm  --model_args pretrained=pfnet/plamo-2.1-2b-cpt,tensor_parallel_size=1,dtype=auto,gpu_memory_utilization=0.95,max_model_len=4096 --batch_size auto --trust_remote_code  --cache_requests true --tasks gsm8k

[rank0]: Traceback (most recent call last):
[rank0]:   File "/net/storage149/mnt/md0/ccyang/miniforge3/envs/vllm_src_torch_nightly/bin/lm_eval", line 8, in <module>
[rank0]:     sys.exit(cli_evaluate())
[rank0]:              ^^^^^^^^^^^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/miniforge3/envs/vllm_src_torch_nightly/lib/python3.12/site-packages/lm_eval/__main__.py", line 449, in cli_evaluate
[rank0]:     results = evaluator.simple_evaluate(
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/miniforge3/envs/vllm_src_torch_nightly/lib/python3.12/site-packages/lm_eval/utils.py", line 439, in _wrapper
[rank0]:     return fn(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/miniforge3/envs/vllm_src_torch_nightly/lib/python3.12/site-packages/lm_eval/evaluator.py", line 230, in simple_evaluate
[rank0]:     lm = lm_eval.api.registry.get_model(model).create_from_arg_string(
[rank0]:          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/miniforge3/envs/vllm_src_torch_nightly/lib/python3.12/site-packages/lm_eval/api/model.py", line 151, in create_from_arg_string
[rank0]:     return cls(**args, **args2)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/miniforge3/envs/vllm_src_torch_nightly/lib/python3.12/site-packages/lm_eval/models/vllm_causallms.py", line 177, in __init__
[rank0]:     self.model = LLM(**self.model_args)
[rank0]:                  ^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/github.com/vllm/vllm/entrypoints/llm.py", line 281, in __init__
[rank0]:     self.llm_engine = LLMEngine.from_engine_args(
[rank0]:                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/github.com/vllm/vllm/engine/llm_engine.py", line 490, in from_engine_args
[rank0]:     return engine_cls.from_vllm_config(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/github.com/vllm/vllm/engine/llm_engine.py", line 466, in from_vllm_config
[rank0]:     return cls(
[rank0]:            ^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/github.com/vllm/vllm/engine/llm_engine.py", line 260, in __init__
[rank0]:     self._initialize_kv_caches()
[rank0]:   File "/net/storage149/mnt/md0/ccyang/github.com/vllm/vllm/engine/llm_engine.py", line 402, in _initialize_kv_caches
[rank0]:     self.model_executor.determine_num_available_blocks())
[rank0]:     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/github.com/vllm/vllm/executor/executor_base.py", line 105, in determine_num_available_blocks
[rank0]:     results = self.collective_rpc("determine_num_available_blocks")
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/github.com/vllm/vllm/executor/uniproc_executor.py", line 58, in collective_rpc
[rank0]:     answer = run_method(self.driver_worker, method, args, kwargs)
[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/github.com/vllm/vllm/utils/__init__.py", line 2998, in run_method
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/miniforge3/envs/vllm_src_torch_nightly/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/github.com/vllm/vllm/worker/worker.py", line 257, in determine_num_available_blocks
[rank0]:     self.model_runner.profile_run()
[rank0]:   File "/net/storage149/mnt/md0/ccyang/miniforge3/envs/vllm_src_torch_nightly/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/github.com/vllm/vllm/worker/model_runner.py", line 1205, in profile_run
[rank0]:     self._dummy_run(max_num_batched_tokens, max_num_seqs)
[rank0]:   File "/net/storage149/mnt/md0/ccyang/github.com/vllm/vllm/worker/model_runner.py", line 1331, in _dummy_run
[rank0]:     self.execute_model(model_input, kv_caches, intermediate_tensors)
[rank0]:   File "/net/storage149/mnt/md0/ccyang/miniforge3/envs/vllm_src_torch_nightly/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/github.com/vllm/vllm/worker/model_runner.py", line 1702, in execute_model
[rank0]:     hidden_or_intermediate_states = model_executable(
[rank0]:                                     ^^^^^^^^^^^^^^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/miniforge3/envs/vllm_src_torch_nightly/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/miniforge3/envs/vllm_src_torch_nightly/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/github.com/vllm/vllm/model_executor/models/plamo2.py", line 769, in forward
[rank0]:     self.mamba_cache = MambaCacheManager(
[rank0]:                        ^^^^^^^^^^^^^^^^^^
[rank0]: TypeError: MambaCacheManager.__init__() missing 1 required positional argument: 'temporal_state_dtype'

[tdoublep]

Oh right, yeah it will be broken. I will fix it - thanks for reporting.

[cyang49]

phi4flash too

VLLM_ATTENTION_BACKEND=DIFFERENTIAL_FLASH_ATTN lm_eval --model vllm  --model_args pretrained=microsoft/Phi-4-mini-flash-reasoning,tensor_parallel_size=1,dtype=auto,gpu_memory_utilization=0.95,enable_prefix_caching=False,enable_chunked_prefill=False,max_model_len=8192 --batch_size auto --trust_remote_code  --cache_requests true --tasks gsm8k

Error

[rank0]:   File "/net/storage149/mnt/md0/ccyang/miniforge3/envs/vllm_src_torch_nightly/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/miniforge3/envs/vllm_src_torch_nightly/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/net/storage149/mnt/md0/ccyang/github.com/vllm/vllm/model_executor/models/phi4flash.py", line 652, in forward
[rank0]:     self.mamba_cache = MambaCacheManager(
[rank0]:                        ^^^^^^^^^^^^^^^^^^
[rank0]: TypeError: MambaCacheManager.__init__() missing 1 required positional argument: 'temporal_state_dtype'