[Model] Enable Inference Support for the New Baichuan-M1 Model

[rainkert]

This pull request adds the necessary support to the vLLM framework for the Baichuan-M1 model.

HuggingFace Page:
https://huggingface.co/baichuan-inc/Baichuan-M1-14B-Base
https://huggingface.co/baichuan-inc/Baichuan-M1-14B-Instruct

The Baichuan-M1 (M stands for medicine) model is a medical-enhanced general large model, designed to deliver exceptional performance in healthcare applications while maintaining strong general capabilities. This update ensures that VLLM can seamlessly handle inference for the Baichuan-M1 model, providing both compatibility and optimal performance for a wide range of natural language processing tasks, especially in the medical domain.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
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 do one of these:

• Add ready label to the PR
• Enable auto-merge.

🚀

[mergify]

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

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

[jameswu2014]

LGTM

[rainkert]

@youkaichao @zhuohan123 @DarkLight1337 @WoosukKwon
We will be releasing our model on Hugging Face on January 24th（The day after tomorrow), but you can review the code beforehand to identify any issues so we can address them in advance.

[rainkert]

ping @youkaichao @DarkLight1337 @njhill @comaniac @zhuohan123 @WoosukKwon @alexm-redhat
We've released our new model today, plz review this PR and merge ASAP.

[DarkLight1337]

The model itself LGTM, but I'm not so sure about the custom KV cache. Is anyone else familiar with this part of the code?

[simon-mo]

Regarding the SWA, can we minimize the code change for now by adopting #10584? While we will work on refactoring the memory manager in #11382 by @heheda12345.

[rainkert]

Regarding the SWA, can we minimize the code change for now by adopting #10584? While we will work on refactoring the memory manager in #11382 by @heheda12345.

Because the kvcache used by ordinary layers and swa layers is inconsistent (we have 2 kv heads in normal attention, but 8 kv heads in swa), we cannot simply treat them the same way as in #10584 , but instead need to separately calculate the memory usage.

[heheda12345]

For vLLM v1 engine, you can support normal attention with different hidden size by extending this function.

vllm/vllm/v1/core/kv_cache_utils.py

Line 410 in bf21481

def get_kv_cache_config(vllm_config: VllmConfig, kv_cache_spec: KVCacheSpec,

Then you can try #10584 in v1 to support the mix of normal attention and SWA.
If that works, we can raise an error to ask the user to use v1 engine to run this model if they do not enable vLLM v1.

[rainkert]

For vLLM v1 engine, you can support normal attention with different hidden size by extending this function.

vllm/vllm/v1/core/kv_cache_utils.py

Line 410 in bf21481

def get_kv_cache_config(vllm_config: VllmConfig, kv_cache_spec: KVCacheSpec,

Then you can try #10584 in v1 to support the mix of normal attention and SWA.
If that works, we can raise an error to ask the user to use v1 engine to run this model if they do not enable vLLM v1.

Over the past few days, we attempted to extend this function and modified our model's forward process, successfully running our model on vLLM v1 in simple cases. However, we encountered new challenges when max_num_batched_tokens > 2048, which appear difficult to resolve:

In such cases, chunked prefill and cascade attention come into effect. Our model has a sliding_window_size, which is not supported in cascade attention. Additionally, similar to Mamba, our model is stateful. We use tokens from previous time steps to smooth our key and value in attention compute, and treat prefill and decode phases differently:
https://huggingface.co/baichuan-inc/Baichuan-M1-14B-Base/blob/main/modeling_baichuan.py#L482
This becomes significantly more complex with chunked attention and v1.

In summary, supporting our model in v1 seems considerably more complicated than in v0. Therefore, v0 might be the better choice for now.