[V1] - Split Prefill and Decode for Mamba1 models

[amirai21]

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

There was an issue where in the mamba1 mixer, the decode kernels weren't being called. The PR fixes this issue so now prefill and decodes are using the correct kernels, and so are mixed batches.

We tested ai21labs/AI21-Jamba-Mini-1.7 on a single H100-80GB GPU and observed improved performance on the ShareGPT dataset.

Running the model:

vllm serve ai21labs/AI21-Jamba-Mini-1.7 --tensor-parallel-size 1 --enforce-eager --max-model-len 32768 --no-enable-prefix-caching --quantization experts_int8

• --enforce-eager since CUDA graph support for Mamba1 in v1 is still under development.
• --no-enable-prefix-caching since prefix caching is not yet supported for Mamba1.

Running the serving benchmark:

python benchmarks/benchmark_serving.py \
    --backend openai-chat \
    --model ai21labs/AI21-Jamba-Mini-1.7 \
    --host 127.0.0.1 \
    --port 8000 \
    --endpoint /v1/chat/completions \
    --dataset-name sharegpt \
    --dataset-path /dev/shm/hub/datasets--JustJaro--ShareGPT_V3_unfiltered_cleaned_split/snapshots/924d943df94e828d6ede39c2df7c2c0a77ad79d1/ShareGPT_V3_unfiltered_cleaned_split.json \
    --ignore-eos \
    --percentile-metrics ttft,tpot,itl \
    --metric-percentiles 99

main (50f2aae):

============ Serving Benchmark Result ============
Successful requests:                     1000      
Benchmark duration (s):                  1610.07   
Total input tokens:                      235326    
Total generated tokens:                  218837    
Request throughput (req/s):              0.62      
Output token throughput (tok/s):         135.92    
Total Token throughput (tok/s):          282.08    
---------------Time to First Token----------------
Mean TTFT (ms):                          111495.87 
Median TTFT (ms):                        84735.40  
P99 TTFT (ms):                           319799.52 
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          6736.32   
Median TPOT (ms):                        6023.01   
P99 TPOT (ms):                           14887.79  
---------------Inter-token Latency----------------
Mean ITL (ms):                           3793.12   
Median ITL (ms):                         2855.45   
P99 ITL (ms):                            16054.95  
==================================================

With split prefill decode changes:

============ Serving Benchmark Result ============
Successful requests:                     1000      
Benchmark duration (s):                  93.16     
Total input tokens:                      235326    
Total generated tokens:                  218837    
Request throughput (req/s):              10.73     
Output token throughput (tok/s):         2348.95   
Total Token throughput (tok/s):          4874.90   
---------------Time to First Token----------------
Mean TTFT (ms):                          21418.87  
Median TTFT (ms):                        20999.66  
P99 TTFT (ms):                           41827.81  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          440.31    
Median TPOT (ms):                        192.88    
P99 TPOT (ms):                           1369.78   
---------------Inter-token Latency----------------
Mean ITL (ms):                           147.58    
Median ITL (ms):                         63.21     
P99 ITL (ms):                            1396.90   
==================================================

Test Plan

Tests for mamba1 and Jamba pass in test_hybrid.py

Test Result

(Optional) Documentation Update

[gemini-code-assist]

Code Review

The pull request refactors the Mamba1 mixer to handle prefill and decode operations using dedicated kernels, which also addresses issues with mixed batches. The changes involve splitting the forward pass logic and introducing new metadata to manage prefill and decode contexts. My review has identified a few critical issues related to correctness, particularly an incorrect variable assignment and a logic path that would fail for V0 models. There are also some type hint inconsistencies that should be addressed for better code clarity and maintainability.

[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.

🚀

[heheda12345]

@amirai21 Can you fix the pre-commit errors and address gemini's comments? I think there are some useful things in them.

[heheda12345]

Also CC @tdoublep if you are interested in this PR.

[amirai21]

@heheda12345 - thanks for reviewing! I’ve addressed Gemini’s comments and fixed the pre-commit errors. I’ve also added OpenGPT benchmarks for the changes.

[heheda12345]

Just like mamba2, can mamba1_metadata.has_initial_states be a tensor with shape (num_prefills,)? has_initial_states_p is a bit confusing now because it is with shape (num_decode+num_prefill,)

[heheda12345]

should has_initial_state be Optional[torch.Tensor]? (I'll fix those in Mamba2AttentionMetadata) And prefer it to be in shape (num_prefills,) like in mamba2.

[heheda12345]

fix for mamba2: #22787

[amirai21]

• Fixed has_initial_states type hint to Optional.
• We return here the has_initial_states of shape (num_prefills + num_decodes,) and in the mixer split_batch_to_prefill_and_decode take the num_prefills from it (via v0 / v1 order logic), then it becomes has_initial_states_p that's passed to the has_prefill flow kernels.

[heheda12345]

do you need has_initial_states_p=pefill_decode_split.has_initial_states_p?

[amirai21]

Yes, we pass has_initial_states_p to causal_conv1d_fn / selective_scan_fn in prefill.

[heheda12345]

LGTM! Can you also run some evals (e.g., by lm_eval) to verify the correctness?

[amirai21]

Thanks for reviewing. We run lm-evel:

VLLM_USE_V1=1 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.0183|±  |0.0105|

We get the same results with main.

[heheda12345]

@amirai21 CI failures seem to be related to this PR. Can you fix them?

[Josephasafg]

@heheda12345 I updated the PR with the relevant test fixes. Thanks.

I added enforce_eager=True to the mamba1 models in test_hybrid::test_models, in the next PR, when I add FCG, I'll remove it.