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.