[V1] Implement sliding window attention in kv_cache_manager

[heheda12345]

Build on top of #14079, should be merged after it.

This pr supports “real” sliding window in v1:

1. Support dropping blocks outside sliding window
2. For prefix caching, only requires the last sliding_window tokens to be cached to achieve a prefix cache hit. e.g., for request ABCDE with sliding window size 2 & block_size 1, if DE are cached while ABC are not, we can still regard ABCDE as the cached prefix.

For models with global attention + sliding window attention, still regard as global-attention-only model in kv cache manager.

Some questions in #13296
It isn’t compatible with cascade attention yet but should be correct due to

vllm/vllm/v1/attention/backends/flash_attn.py

Line 351 in cfbb8c9

if use_alibi or use_sliding_window:

Q: How does it work with chunked prefill?
A: It will allocate blocks for tokens that will be computed in the current step, and free the blocks that outside sliding window in the next step.
Assume window size 1k, chunk size 2k, prompt length 4k, block_size=1

1. chunk prefill of 2k tokens: block table: [0:2k]
2. chunk prefill of 2k tokens:
   1. Free the first 1k blocks as they won’t be used after step 1. block_table becomes [null_block*1000] + [1k:2k]
   2. Allocate blocks for the [2k:4k] tokens that will be computed. block_table becomes [null_block*1000] + [1k:4k]
3. First decode step
   1. Free the [1k:3k] blocks as they won’t be used after step 2. block_table becomes [null_block*3000] + [3k:4k]
   2. Allocate a new slot for decoding, block_table becomes [null_block*3000] + [3k:4001]

Q: What's the shape of the block table for SWA? Is it append-only?
A: It is with the same length as global attention, but changes the blocks outside the sliding window to a special null_block. This replacement only happens in the kv_cache_manager side. As model_runner’s block_table is append-only, we do not replace existing blocks to null blocks in model runner. The result is correct because model runner won’t access the blocks outside sliding window.

This pr is part of the hybrid allocator #11382

[github-actions]

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

🚀

[mergify]

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

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

[WoosukKwon]

@zhuohan123 Did you have a chance to take a look?

[zhuohan123]

WIP partial review. Will add more tmrw

[zhuohan123]

Q: Will pipeline parallelism fail this assert for some models?

[heheda12345]

Yes, it will fail when different stages have different type of layers. For hybrid models, just throw an error as the first step. For non-hybrid models, the assert won't fail.

This function is introduced in https://github.com/vllm-project/vllm/pull/14079/files, should we discuss it there?

[heheda12345]

Thanks for your review! Replied to some questions, will update the code after #14079

[heheda12345]

Yes, it will fail when different stages have different type of layers. For hybrid models, just throw an error as the first step. For non-hybrid models, the assert won't fail.

This function is introduced in https://github.com/vllm-project/vllm/pull/14079/files, should we discuss it there?

[mergify]

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

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

[WoosukKwon]

As we discussed offline, I think we need a clear separation of the two APIs of SpecializedManager:

1. The first API describing how to free the KV cache for a request. This API is invoked whenever the request calls allocate_slots.
2. The second API describing how to check the prefix cache hits. This API is invoked in get_computed_blocks.

Two particular things I want to fix in this PR are that

1. The first API is also used in get_computed_blocks, blurring the separation.
2. The second API does not provide "efficient intersection" or "early stopping". I think this can be addressed by adding an extended API (e.g., get longest prefix).

[heheda12345]

I've change the prefix caching API in SpecializedManager to find_longest_cache_hit. Seems that most complexity can be fixed after that. This find_longest_cache_hit also works for hybrid allocator as long as we can avoid the recomputation between the multiple calls of it for the same request.

[WoosukKwon]

Great simplification. I love it!

[WoosukKwon]

LGTM. Thanks for the awesome work! I’m really happy with how the final APIs and simplifications turned out 🔥🔥

Also, huge thanks for your incredible patience throughout all the back-and-forth edits and long discussions. Really appreciate it!