Repository navigation
Running Bamba on vLLM #3
Description
This issue tracks progress on running Bamba on vLLM.
Success for this issue implies the following:
- Running the model successfully from the HF checkpoint in vLLM (Add Bamba Model vllm-project/vllm#10909)
- Ensuring chunked prefill and TP work in vLLM
- Closing the performance gap in vLLM wrt Llama of similar sizes
- Reporting the performance results in a blog post
Activity
fabianlim commentedon Dec 9, 2024
@ani300 For the TP i have a reasonable fix that
I have not yet upstreamed to the main PR, see this.this is now upstreamedAlso, I realized that the
program_idsneed to be inint64for large prefil sizes,I also have not yet upstreamed the fix.this is now upstreamedIssues:
NOTE: There is one more issue not yet resolved, after the memory profiling we can see the
post_profileis not recovering back tointiial_memory_usage. Unfortunately, this issue does not allow us to run-i 131072without eager mode-i 131072 -o 1024 -b 1 -t 1 -s 1 -q None -k auto -r 3 -u .92 -e2 .Also I noticed this in the new VLLM version
AdnanHoque commentedon Dec 10, 2024
The
chunk_scan_fwdkernel, is not H100 optimized and exhibits low compute and memory throughput, relative to pk performance of hardware as reported by NCU for bs1, seq64k.chunk_scan_fwdis the most compute intensive kernel in the prefill stage. Potential Hopper optimizations:For GMEM -> SMEM transfers use TMA instead of cp.async engine
Add Persistent Tile Scheduler
fabianlim commentedon Jan 6, 2025
because we share the unit tests with bamba, we need to have an extra
use_mamba_kernels: Falseadded to configuration_bamba.py. This only needs to be done on the dev model @divya-kumari32fxnie commentedon Jan 15, 2025
May I ask if VLLM already supports BAMBA? After conducting experiments, I have found that it is currently not supported?
raghukiran1224 commentedon Jan 15, 2025
vLLM integration status can be found here: vllm-project/vllm#10909.
You can use that PR until it gets merged.
XuanofXXX commentedon Oct 20, 2025
I am trying to use vLLM to deploy the Bamba model for evaluation. However, the evaluation on GSM8K and HellaSwag with the same number of shots is very different from the official scores. I dug into the completions from vLLM and HF and found that the outputs from these two modules are very different (BambaMixer2 actually). Have you ever encountered this issue?