[Josephasafg]

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.

Purpose

Add full v1-style attention-metadata support to Mamba-1 models. The following was added

1. Mamba1AttentionMetadataBuilder to create the attention metadata for mamba-1 models
2. MambaStateShapeCalculator aggregated mamba state shape calculation logic for all mamba models under the same static class for better readability and navigation
3. Updated selective_scan_fwd to support V1 memory layout

Test Plan

Updated the following tests -
test_hybrid.py - Now tests Mamba1 and Jamba in V1
test_oracle.py - Removed Mamba1 from unsupported V1 model
Running all tests in test_hybrid.py, test_oracle.py and test_mamba_ssm.py (due to kernel change) pass.

Running mamba1 in main branch would raise an error

lm_eval results with state-spaces/mamba-130m-hf

vLLM V0 with Mamba1

VLLM_USE_V1=0 HF_ALLOW_CODE_EVAL=1  lm_eval --model vllm \
    --model_args pretrained=state-spaces/mamba-130m-hf,enforce_eager=True,enable_prefix_caching=False,tensor_parallel_size=1 \
    --tasks humaneval \
    --batch_size auto \
    --confirm_run_unsafe_code
|  Tasks  |Version|  Filter   |n-shot|Metric|   |Value |   |Stderr|
|---------|------:|-----------|-----:|------|---|-----:|---|-----:|
|humaneval|      1|create_test|     0|pass@1|   |0.0244|±  |0.0121|

vLLM V1 with Mamba1

VLLM_USE_V1=1 HF_ALLOW_CODE_EVAL=1  lm_eval --model vllm \
    --model_args pretrained=state-spaces/mamba-130m-hf,enable_prefix_caching=False,tensor_parallel_size=1 \
    --tasks humaneval \
    --batch_size auto \
    --confirm_run_unsafe_code
|  Tasks  |Version|  Filter   |n-shot|Metric|   |Value |   |Stderr|
|---------|------:|-----------|-----:|------|---|-----:|---|-----:|
|humaneval|      1|create_test|     0|pass@1|   |0.0244|±  |0.0121|

This PR now enables this vLLM V1 to work with models that use Mamba1 like Mamba1 and Jamba

[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 v1-style attention metadata support for Mamba-1 models and refactors the Mamba state shape calculation into a centralized MambaStateShapeCalculator class. The refactoring improves code organization and maintainability. The v1 support is well-integrated, with clear logic separation based on the VLLM_USE_V1 environment variable.

[heheda12345]

Thanks for the great job.

Some questions:

1. Do we need enforce-eager to run mamba1? I'm OK with supporting cuda graph in a future PR.
2. Can you show some lm-eval result on mamba1 model?
3. Please update tests/models/language/generation/test_hybrid.py and tests/v1/test_oracle.py
4. Is Jamba supported now?
5. Please update the doc like https://docs.vllm.ai/en/latest/usage/v1_guide.html#mamba-models and https://docs.vllm.ai/en/latest/usage/v1_guide.html#mamba-models

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Josephasafg.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[tdoublep]

How different is the Mamba1AttentionMetdata to the Mamba2AttentionMetadata? Do we really need two separate classes?

[Josephasafg]

Good question. I think we should keep both since Mamba2AttentionMetadata contains fields that aren't relevant for mamba1 like chunk_indices, chunk_offsets and triton kernels related fields that mamba1 doesn't need which would add overhead to the class.

[heheda12345]

I prefer to keep mamba1 and mamba2 as seperate metadata classes. Comparing with one metadata class with many optional entry and branches for different types of layers, I prefer this pluggable design. We can extract common logic to a parent class if we find some after more models like minimax are added. @tdoublep WDYT?

[tdoublep]

@heheda12345 Sorry I missed this question. Yes I agree let's keep the metadata classes separate and then factor the common things into a CommonMambaAttentionMetadata or something once it is clear what is truly common. This one is ready and we have the LFM2 and MiniMax-Text ones nearly there too, so we should be able to look at that soon.

[heheda12345]

I prefer to keep mamba1 and mamba2 as seperate metadata classes. Comparing with one metadata class with many optional entry and branches for different types of layers, I prefer this pluggable design. We can extract common logic to a parent class if we find some after more models like minimax are added. @tdoublep WDYT?

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Josephasafg.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Josephasafg.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[tlrmchlsmth]

PR looks good. Thanks for adding Jamba and Mamba1 to V1!

[tdoublep]

V1 test failure is known flaky test (see #22385), and the quantization test and blackwell tests look unrelated. This one is good to merge imo.