-
-
Notifications
You must be signed in to change notification settings - Fork 8.9k
[v1][core] Support for attention free models #20811
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 9 commits
b764c9d
5825ba4
fc86350
97c11e6
673aeb0
fb3ecfb
8e5dbee
2ee7087
19a7d70
b8f355e
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change | ||||||||||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
@@ -78,7 +78,9 @@ def __init__( | |||||||||||||||||
) -> None: | ||||||||||||||||||
self.max_model_len = max_model_len | ||||||||||||||||||
|
||||||||||||||||||
self.enable_caching = enable_caching | ||||||||||||||||||
self.enable_caching = (enable_caching | ||||||||||||||||||
if len(kv_cache_config.kv_cache_groups) > 0 | ||||||||||||||||||
else False) | ||||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
I think this structure is more clear for readers. |
||||||||||||||||||
self.caching_hash_fn = ( | ||||||||||||||||||
sha256_cbor_64bit if caching_hash_algo == "sha256_cbor_64bit" else | ||||||||||||||||||
sha256 if caching_hash_algo == "sha256" else hash) | ||||||||||||||||||
|
@@ -101,7 +103,7 @@ def __init__( | |||||||||||||||||
kv_cache_config=kv_cache_config, | ||||||||||||||||||
max_model_len=self.max_model_len, | ||||||||||||||||||
use_eagle=self.use_eagle, | ||||||||||||||||||
enable_caching=enable_caching, | ||||||||||||||||||
enable_caching=self.enable_caching, | ||||||||||||||||||
caching_hash_fn=self.caching_hash_fn, | ||||||||||||||||||
enable_kv_cache_events=enable_kv_cache_events, | ||||||||||||||||||
) | ||||||||||||||||||
|
Original file line number | Diff line number | Diff line change | ||||||||||||||||||||||
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
@@ -139,7 +139,11 @@ def _initialize_kv_caches( | |||||||||||||||||||||||
|
||||||||||||||||||||||||
# Profiles the peak memory usage of the model to determine how much | ||||||||||||||||||||||||
# memory can be allocated for kv cache. | ||||||||||||||||||||||||
available_gpu_memory = self.model_executor.determine_available_memory() | ||||||||||||||||||||||||
check_available_memory = not(len(kv_cache_specs) == 1 and not kv_cache_specs[0]) | ||||||||||||||||||||||||
available_gpu_memory = [0] | ||||||||||||||||||||||||
if check_available_memory: | ||||||||||||||||||||||||
available_gpu_memory = ( | ||||||||||||||||||||||||
self.model_executor.determine_available_memory()) | ||||||||||||||||||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
I feel that the condition is not correct. len(kv_cache_specs) can be larger than 1 when TP / PP is enabled. |
||||||||||||||||||||||||
|
||||||||||||||||||||||||
assert len(kv_cache_specs) == len(available_gpu_memory) | ||||||||||||||||||||||||
# Get the kv cache tensor size | ||||||||||||||||||||||||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Do we need this given prefix caching is disabled here for models that don’t use last pooling method?
https://github.com/maxdebayser/vllm/blob/221f013922c0c118b682d294755e69990b2c43ed/vllm/config.py#L4505
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Without this check though you would not be able to disable attention for models that are not of the pooling type as prefix caching is enabled by default for all models except pooling ones.
See below:
vllm/vllm/engine/arg_utils.py
Lines 1620 to 1630 in 38efa28
Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Perhaps the safest thing to do is to disable prefix-caching in
VllmConfig.__post_init__
right away for any attention free models and then yes, we could just rely onenable_caching
as you suggest.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
remove this comment?