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

🚀

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

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

@tdoublep This is great. Suggestion - Everything here is not very mamba2 specific, meaning, once we add full cuda support to mamba1 (we have the PR waiting to be opened), we would want to add this code to it as well without duplicating it. So perhaps its worth removing the 2 suffix from the names here, so we can just take it as is once FCG is supported. wdyt?

@Josephasafg Do you expect that the effect of FCG is as significant for mamba1 as for mamba2?

@tdoublep from the initial runs I had made, it looks to be significant (but not as mamba2)
These are some of the numbers I had gotten

No FCG

VLLM_USE_V1=1 VLLM_ATTENTION_BACKEND=FLASHINFER python benchmark_latency.py \
        --model ai21labs/AI21-Jamba-Mini-1.7-FP8 \
        --input-len 31500 \
        --output-len 128 \
        --batch-size 1  \
        --num-iters-warmup 3 \
        --num-iters 3 \
        --max-num-seqs 1 \
        --enforce-eager
Profiling iterations: 100% 3/3 [00:14<00:00,  5.00s/it]
Avg latency: 4.998069426355262 seconds
10% percentile latency: 4.991475126799196 seconds
25% percentile latency: 4.994978636037558 seconds
50% percentile latency: 5.0008178181014955 seconds
75% percentile latency: 5.002534412546083 seconds
90% percentile latency: 5.003564369212836 seconds
99% percentile latency: 5.004182343212888 seconds

with FCG

VLLM_USE_V1=1 VLLM_ATTENTION_BACKEND=FLASHINFER python benchmark_latency.py \
        --model ai21labs/AI21-Jamba-Mini-1.7-FP8 \
        --input-len 31500 \
        --output-len 128 \
        --batch-size 1  \
        --num-iters-warmup 3 \
        --num-iters 3 \
        --max-num-seqs 1 \
        --compilation-config='{"full_cuda_graph": true}'

Profiling iterations: 100% 3/3 [00:09<00:00,  3.21s/it]
Avg latency: 3.2101284717209637 seconds
10% percentile latency: 3.2057920150924475 seconds
25% percentile latency: 3.2082151676295325 seconds
50% percentile latency: 3.212253755191341 seconds
75% percentile latency: 3.2131044175475836 seconds
90% percentile latency: 3.2136148149613293 seconds
99% percentile latency: 3.2139210534095763 seconds

@Josephasafg Could you also share V0 numbers for the same benchmark?

@tdoublep For V0 I had gotten this

Avg latency: 5.413307532668114 seconds
10% percentile latency: 5.40432041650638 seconds
25% percentile latency: 5.406467076041736 seconds
50% percentile latency: 5.4100448419339955 seconds
75% percentile latency: 5.418516643927433 seconds
90% percentile latency: 5.423599725123495 seconds
99% percentile latency: 5.426649573841132 seconds

@Josephasafg Did you include #22653 in your benchmark?

@heheda12345 yes, it's based on it , which is why I haven't opened the PR yet.

@tdoublep Can you make this PR not that mamba2 specific?

@tdoublep Can you make this PR not that mamba2 specific?

Yes I can. Just need to decide how to handle the mamba-only case (e.g., models that are not actually hybrid so we can't check is_hybrid).

@heheda12345 PTAL - I've had another go at this, much simpler now.

After this one, we should merge: #23650

Nice clean solution. I'm wondering if we should keep this mamba specific or consider a more general solution for other models down the line, perhaps this approach works well for now.