Skip to content

[Misc][Hybrid allocator + kv connector] Optionally enable hybrid allocator + KV cache connector #29805

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.

Sign up for GitHub

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

Merged
heheda12345 merged 8 commits into from
Dec 15, 2025
Merged

[Misc][Hybrid allocator + kv connector] Optionally enable hybrid allocator + KV cache connector #29805

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
a376cf4
enable hma+kv_conn
NickLucche Dec 1, 2025
95d4971
Revert "enable hma+kv_conn"
NickLucche Dec 4, 2025
2142e46
optional bool flag
NickLucche Dec 4, 2025
664db48
comment
NickLucche Dec 5, 2025
e854448
rasing
NickLucche Dec 9, 2025
4d59565
comment
NickLucche Dec 9, 2025
c965044
chen review
NickLucche Dec 11, 2025
2d1535d
Merge branch 'main' into enable-hma-kv-connector
NickLucche Dec 15, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion vllm/config/scheduler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,10 +122,12 @@ class SchedulerConfig:
the default scheduler. Can be a class directly or the path to a class of
form "mod.custom_class"."""

disable_hybrid_kv_cache_manager: bool = False
disable_hybrid_kv_cache_manager: bool | None = None
"""If set to True, KV cache manager will allocate the same size of KV cache
for all attention layers even if there are multiple type of attention layers
like full attention and sliding window attention.
If set to None, the default value will be determined based on the environment
and starting configuration.
"""

async_scheduling: bool = False
Expand Down
52 changes: 38 additions & 14 deletions vllm/config/vllm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -889,27 +889,16 @@ def has_blocked_weights():
if not self.instance_id:
self.instance_id = random_uuid()[:5]

# Runtime-dependent disable of hybrid kv cache manager logic.
if not self.scheduler_config.disable_hybrid_kv_cache_manager:
prev_disable_hma = self.scheduler_config.disable_hybrid_kv_cache_manager

@heheda12345 heheda12345 Dec 9, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What about:

need_disable_hybrid_kv_cache_manager = False
if not current_platform.support_hybrid_kv_cache():
    need_disable_hybrid_kv_cache_manager = True
if ***: need_disable_hybrid_kv_cache_manager = True
...
if self.scheduler_config.disable_hybrid_kv_cache_manager  is None:
    if self.kv_transfer_config is not None:
         # Experimental feature. Default to disable but allow users to enable.
         need_disable_hybrid_kv_cache_manager = True
         logger.warning(***)
    self.scheduler_config.disable_hybrid_kv_cache_manager = need_disable_hybrid_kv_cache_manager
elif self.scheduler_config.disable_hybrid_kv_cache_manager == False:
    if need_disable_hybrid_kv_cache_manager: raise xxxx

I feel prev_disable_hma is a little bit hacky.

@NickLucche NickLucche Dec 10, 2025 •
edited
Loading

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@heheda12345 I am sorry but I don't see how checking whether a bool was flipped here is hacky.
EDIT: I think I see what you mean now, you're referring to clarity. Don't have a strong opinion on that, will check it out when I have the time.

I've separated features that do NOT work with HMA from features that may work with HMA such as kv connector (a later check on supports_hma is carried out when attempting to create the actual connector).

Attempting to enable HMA explicitly on the former will crash the server. prev_disable_hma is just for checking whether the flag was set by the user without adding more lines such as need_disable_hybrid_kv_cache_manager.


# logger should only print warning message for hybrid models. As we
# can't know whether the model is hybrid or not now, so we don't log
# warning message here and will log it later.
if not current_platform.support_hybrid_kv_cache():
# Hybrid KV cache manager is not supported on non-GPU platforms.
self.scheduler_config.disable_hybrid_kv_cache_manager = True
if self.kv_transfer_config is not None:
# NOTE(Kuntai): turn HMA off for connector for now.
# TODO(Kuntai): have a more elegent solution to check and
# turn off HMA for connector that does not support HMA.
logger.warning(
"Turning off hybrid kv cache manager because "
"`--kv-transfer-config` is set. This will reduce the "
"performance of vLLM on LLMs with sliding window attention "
"or Mamba attention. If you are a developer of kv connector"
", please consider supporting hybrid kv cache manager for "
"your connector by making sure your connector is a subclass"
" of `SupportsHMA` defined in kv_connector/v1/base.py."
)
self.scheduler_config.disable_hybrid_kv_cache_manager = True
if self.kv_events_config is not None:
# Hybrid KV cache manager is not compatible with KV events.
self.scheduler_config.disable_hybrid_kv_cache_manager = True
Expand All @@ -935,6 +924,41 @@ def has_blocked_weights():
# local attention.
self.scheduler_config.disable_hybrid_kv_cache_manager = True

if (
prev_disable_hma is False
and self.scheduler_config.disable_hybrid_kv_cache_manager is True
):
raise ValueError(
"Hybrid KV cache manager was explicitly enabled but is not "
"supported in this configuration. Consider omitting the "
"--no-disable-hybrid-kv-cache-manager flag to let vLLM decide"
" automatically."
)

if (
self.scheduler_config.disable_hybrid_kv_cache_manager is None
and self.kv_transfer_config is not None
):
# Disable HMA logic but only if the user didn't express a preference.
# NOTE(Kuntai): turn HMA off for connector unless specifically enabled.
# TODO(Kuntai): have a more elegent solution to check and
# turn off HMA for connector that does not support HMA.
logger.warning(
"Turning off hybrid kv cache manager because "
"`--kv-transfer-config` is set. This will reduce the "
"performance of vLLM on LLMs with sliding window attention "
"or Mamba attention. If you are a developer of kv connector"
", please consider supporting hybrid kv cache manager for "
Comment thread
NickLucche marked this conversation as resolved.
Outdated
"your connector by making sure your connector is a subclass"
" of `SupportsHMA` defined in kv_connector/v1/base.py and"
" use --no-disable-hybrid-kv-cache-manager to start vLLM."
)
self.scheduler_config.disable_hybrid_kv_cache_manager = True

if self.scheduler_config.disable_hybrid_kv_cache_manager is None:
# Default to enable HMA if not explicitly disabled by user or logic above.
self.scheduler_config.disable_hybrid_kv_cache_manager = False

if self.compilation_config.debug_dump_path:
self.compilation_config.debug_dump_path = (
self.compilation_config.debug_dump_path.absolute().expanduser()
Expand Down
2 changes: 1 addition & 1 deletion vllm/engine/arg_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -494,7 +494,7 @@ class EngineArgs:
enable_chunked_prefill: bool | None = None
disable_chunked_mm_input: bool = SchedulerConfig.disable_chunked_mm_input

disable_hybrid_kv_cache_manager: bool = (
disable_hybrid_kv_cache_manager: bool | None = (
SchedulerConfig.disable_hybrid_kv_cache_manager
)

Expand Down