Running Bamba on vLLM

[ani300]

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

cc @raghukiran1224 @fabianlim @AdnanHoque

[fabianlim]

@ani300 For the TP i have a reasonable fix that I have not yet upstreamed to the main PR, see this. this is now upstreamed

Also, I realized that the program_ids need to be in int64 for large prefil sizes, I also have not yet upstreamed the fix. this is now upstreamed

Issues:

1. NOTE: There is one more issue not yet resolved, after the memory profiling we can see the post_profile is not recovering back to intiial_memory_usage. Unfortunately, this issue does not allow us to run -i 131072 without eager mode
   -i 131072 -o 1024 -b 1 -t 1 -s 1 -q None -k auto -r 3 -u .92 -e

   INFO 12-08 12:54:53 worker.py:236] Memory profiling results: duration=24.19 seconds, total_gpu_memory=79.15GiB, initial_memory_usage=18.88GiB, peak_torch_memory=69.37GiB, memory_usage_post_profile=48.60GiB, non_torch_memory=0.55GiB, kv_cache_size=5.28GiB, gpu_memory_utilization=0.95.
   

2 .Also I noticed this in the new VLLM version

```
[[CINFO 12-12 05:52:22 config.py:403] This model supports multiple tasks: {'embedding', 'generate'}. Defaulting to 'generate'.

```

3. FlashInfer backend is not supported.

[AdnanHoque]

The chunk_scan_fwd kernel, 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_fwd is the most compute intensive kernel in the prefill stage. Potential Hopper optimizations:

1. For GMEM -> SMEM transfers use TMA instead of cp.async engine

2. Add Persistent Tile Scheduler

[fabianlim]

because we share the unit tests with bamba, we need to have an extra use_mamba_kernels: False added to configuration_bamba.py. This only needs to be done on the dev model @divya-kumari32

[fxnie]

May I ask if VLLM already supports BAMBA? After conducting experiments, I have found that it is currently not supported?

[raghukiran1224]

vLLM integration status can be found here: vllm-project/vllm#10909.

You can use that PR until it gets merged.

[XuanofXXX]

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?