[v1][core]Support for attention free models #20811

christian-pinto · 2025-07-11T10:55:37Z

This PR enables support for attention free models by exploiting the the zero kv cache groups approach. This han been created as a follow-up from the discussions in #20577.

@heheda12345 please review.

gemini-code-assist

Summary of Changes

Hello @christian-pinto, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces comprehensive support for 'attention-free models' by adapting the KV cache management system. My changes ensure that components dealing with KV cache allocation, configuration, and memory checks correctly handle models that do not require a KV cache, preventing errors and optimizing resource management for these specific model types. This is achieved by explicitly checking for and handling cases where KV cache groups are zero or empty.

Highlights

Core KV Cache Management Adaptation: I've updated the KVCacheCoordinator and KVCacheManager to gracefully handle models that do not utilize a KV cache (attention-free models). This involves adding checks to ensure that operations dependent on the existence of KV cache groups are only performed when such groups are present, preventing errors when kv_cache_groups is empty.
KV Cache Utility Enhancements: I've introduced new utility functions in kv_cache_utils.py to specifically identify and configure attention-free models. This includes is_kv_cache_type_attention_free to detect such models and _get_kv_cache_config_attention_free to provide a minimal KVCacheConfig suitable for them, which is then prioritized in get_kv_cache_config.
Memory and Spec Determination Optimization: I've optimized the AbstractWorker in abstract.py to skip unnecessary memory determination and KV cache specification retrieval for attention-free models. Since these models don't use a KV cache, these operations are redundant and can be bypassed, returning default or zero values.
Memory Check Bypass: I've added a condition to check_enough_kv_cache_memory in kv_cache_utils.py to bypass the memory availability check entirely if the model is identified as attention-free. This is because attention-free models do not require memory allocation for KV cache blocks.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces support for attention-free models by handling cases where the KV cache configuration is empty. I've identified a correctness issue in a new utility function and a couple of typos in comments. Once these are addressed, the PR should be in good shape.

gemini-code-assist · 2025-07-11T10:56:52Z

vllm/v1/core/kv_cache_utils.py

+    if not kv_cache_spec:
+        return True


This function is type-hinted to return a bool, but it implicitly returns None if kv_cache_spec is not empty. This violates the function's contract and could lead to unexpected behavior. Replace the if statement with a direct return to fix this.

return not kv_cache_spec

@christian-pinto please fix it

gemini-code-assist · 2025-07-11T10:56:52Z

vllm/v1/core/kv_cache_utils.py

+        # This returns a kv_cahce config with 0 kv cache groups and 1 block
+        # to allow for the KVCache manager to handle attention fre models.


There are a couple of typos in this comment: kv_cahce should be kv_cache, and attention fre should be attention free.

# This returns a kv_cache config with 0 kv cache groups and 1 block # to allow for the KVCache manager to handle attention-free models.

christian-pinto · 2025-07-11T11:00:49Z

vllm/v1/core/kv_cache_coordinator.py

-        self.verify_and_split_kv_cache_groups()
+        # attention free models are initialized with 0 kv_cache_groups
+        if len(self.kv_cache_config.kv_cache_groups) > 0:
+            self.verify_and_split_kv_cache_groups()


I will revert this once #20661 is merged and the KVCacheCoordinatorNoPrefixCache is available.

Forces 0 KV Cache groups to disable KV Cache in attention free models Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

github-actions · 2025-07-11T11:05:00Z

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

🚀

heheda12345

Thank you for the great work. Let's wait for #20661.

heheda12345 · 2025-07-11T12:08:53Z

vllm/v1/core/kv_cache_utils.py

+    if not kv_cache_spec:
+        return True


@christian-pinto please fix it

heheda12345 · 2025-07-11T12:13:44Z

vllm/v1/core/kv_cache_manager.py

@@ -86,7 +86,7 @@ def __init__(
        self.prefix_cache_stats = PrefixCacheStats() if log_stats else None

        self.block_size: Optional[int] = None
-        if self.enable_caching:
+        if self.enable_caching and len(self.kv_cache_config.kv_cache_groups) > 0:


Can you fix the type error?

All looks good to me here. Am i missing something?

I can see some precommit error

Oh I see, I have missed this one.
I'll take care of it. Thanks

heheda12345 · 2025-07-11T12:17:18Z

vllm/v1/executor/abstract.py

@@ -73,10 +73,16 @@ def register_failure_callback(self, callback: FailureCallback):
        pass

    def determine_available_memory(self) -> list[int]:  # in bytes
+        if self.vllm_config.model_config.is_attention_free:


Can you put the change in determine_available_memory to gpu_worker and the change inget_kv_cache_specs to gpu_model_runner? Though these changes are platform independent, this abstract class should be as simple as possible.

heheda12345 · 2025-07-11T14:52:01Z

Seems that someone want prefix caching to be enabled in "last" pooling method. Is it possible to only use the zero-group code path for other pooling methods?

vllm/vllm/config.py

Line 4594 in 8aeaa91

if pooling_type is None or pooling_type.lower() != "last":

heheda12345 · 2025-07-14T02:18:20Z

@christian-pinto #20661 is merged. Can you rebase?

christian-pinto requested review from WoosukKwon, robertgshaw2-redhat, njhill, ywang96, comaniac and alexm-redhat as code owners July 11, 2025 10:55

gemini-code-assist bot reviewed Jul 11, 2025

View reviewed changes

mergify bot added the v1 label Jul 11, 2025

gemini-code-assist bot reviewed Jul 11, 2025

View reviewed changes

christian-pinto commented Jul 11, 2025

View reviewed changes

Support for attention free models

e06e264

Forces 0 KV Cache groups to disable KV Cache in attention free models Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

christian-pinto force-pushed the attention_free_models_support branch from 4b0a9f9 to e06e264 Compare July 11, 2025 11:03

is_kv_cache_type_attention_free: return False if not attention free

e8a14f8

heheda12345 reviewed Jul 11, 2025

View reviewed changes

some minor edits after first review round

bb9f2db

		# This returns a kv_cahce config with 0 kv cache groups and 1 block
		# to allow for the KVCache manager to handle attention fre models.

Uh oh!

[v1][core]Support for attention free models #20811

Are you sure you want to change the base?

[v1][core]Support for attention free models #20811

Conversation

christian-pinto commented Jul 11, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jul 11, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 11, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Jul 11, 2025

Uh oh!

heheda12345 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

heheda12345 commented Jul 11, 2025

Uh oh!

heheda12345 commented Jul 14, 2025

Uh oh!

Uh oh!

christian-pinto commented Jul 11, 2025 •

edited by github-actions bot

Loading