A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 11:15 UTC
[Misc][Hybrid allocator + kv connector] Optionally enable hybrid allocator + KV cache connector (#29805)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Dec 15, 2025 at 12:17 UTC
185c22bf2f736d0ffa69a3faae379ad0c444de56
3 files changed
+64
-38
vllm/config/scheduler.py
+3
-1
@@ -122,10 +122,12 @@ class SchedulerConfig:
122
the default scheduler. Can be a class directly or the path to a class of
123
form "mod.custom_class"."""
124
125
- disable_hybrid_kv_cache_manager: bool = False
125
+ disable_hybrid_kv_cache_manager: bool | None = None
126
"""If set to True, KV cache manager will allocate the same size of KV cache
127
for all attention layers even if there are multiple type of attention layers
128
like full attention and sliding window attention.
129
+ If set to None, the default value will be determined based on the environment
130
+ and starting configuration.
131
"""
132
133
async_scheduling: bool = False
vllm/config/vllm.py
+60
-36
@@ -887,17 +887,48 @@ class VllmConfig:
887
if not self.instance_id:
888
self.instance_id = random_uuid()[:5]
889
890
- if not self.scheduler_config.disable_hybrid_kv_cache_manager:
891
- # logger should only print warning message for hybrid models. As we
892
- # can't know whether the model is hybrid or not now, so we don't log
893
- # warning message here and will log it later.
894
- if not current_platform.support_hybrid_kv_cache():
895
- # Hybrid KV cache manager is not supported on non-GPU platforms.
896
- self.scheduler_config.disable_hybrid_kv_cache_manager = True
890
+ # Hybrid KV cache manager (HMA) runtime rules:
891
+ # - Explicit enable (--no-disable-kv-cache-manager): error if runtime
892
+ # disables it
893
+ # - No preference: auto-disable for unsupported features (e.g. kv connector)
894
+ # - Explicit disable (--disable-kv-cache-manager): always respect it
895
+ need_disable_hybrid_kv_cache_manager = False
896
+ # logger should only print warning message for hybrid models. As we
897
+ # can't know whether the model is hybrid or not now, so we don't log
898
+ # warning message here and will log it later.
899
+ if not current_platform.support_hybrid_kv_cache():
900
+ # Hybrid KV cache manager is not supported on non-GPU platforms.
901
+ need_disable_hybrid_kv_cache_manager = True
902
+ if self.kv_events_config is not None:
903
+ # Hybrid KV cache manager is not compatible with KV events.
904
+ need_disable_hybrid_kv_cache_manager = True
905
+ if (
906
+ self.model_config is not None
907
+ and self.model_config.attention_chunk_size is not None
908
+ ):
909
+ if (
910
+ self.speculative_config is not None
911
+ and self.speculative_config.use_eagle()
912
+ ):
913
+ # Hybrid KV cache manager is not yet supported with chunked
914
+ # local attention + eagle.
915
+ need_disable_hybrid_kv_cache_manager = True
916
+ elif not envs.VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE:
917
+ logger.warning(
918
+ "There is a latency regression when using chunked local"
919
+ " attention with the hybrid KV cache manager. Disabling"
920
+ " it, by default. To enable it, set the environment "
921
+ "VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE=1."
922
+ )
923
+ # Hybrid KV cache manager is not yet supported with chunked
924
+ # local attention.
925
+ need_disable_hybrid_kv_cache_manager = True
926
+
927
+ if self.scheduler_config.disable_hybrid_kv_cache_manager is None:
928
+ # Default to disable HMA, but only if the user didn't express a preference.
929
if self.kv_transfer_config is not None:
898
- # NOTE(Kuntai): turn HMA off for connector for now.
899
- # TODO(Kuntai): have a more elegent solution to check and
900
- # turn off HMA for connector that does not support HMA.
930
+ # NOTE(Kuntai): turn HMA off for connector unless specifically enabled.
931
+ need_disable_hybrid_kv_cache_manager = True
932
logger.warning(
933
"Turning off hybrid kv cache manager because "
934
"`--kv-transfer-config` is set. This will reduce the "
@@ -905,33 +936,26 @@ class VllmConfig:
936
"or Mamba attention. If you are a developer of kv connector"
937
", please consider supporting hybrid kv cache manager for "
938
"your connector by making sure your connector is a subclass"
908
- " of `SupportsHMA` defined in kv_connector/v1/base.py."
939
+ " of `SupportsHMA` defined in kv_connector/v1/base.py and"
940
+ " use --no-disable-hybrid-kv-cache-manager to start vLLM."
941
)
910
- self.scheduler_config.disable_hybrid_kv_cache_manager = True
911
- if self.kv_events_config is not None:
912
- # Hybrid KV cache manager is not compatible with KV events.
913
- self.scheduler_config.disable_hybrid_kv_cache_manager = True
914
- if (
915
- self.model_config is not None
916
- and self.model_config.attention_chunk_size is not None
917
- ):
918
- if (
919
- self.speculative_config is not None
920
- and self.speculative_config.use_eagle()
921
- ):
922
- # Hybrid KV cache manager is not yet supported with chunked
923
- # local attention + eagle.
924
- self.scheduler_config.disable_hybrid_kv_cache_manager = True
925
- elif not envs.VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE:
926
- logger.warning(
927
- "There is a latency regression when using chunked local"
928
- " attention with the hybrid KV cache manager. Disabling"
929
- " it, by default. To enable it, set the environment "
930
- "VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE=1."
931
- )
932
- # Hybrid KV cache manager is not yet supported with chunked
933
- # local attention.
934
- self.scheduler_config.disable_hybrid_kv_cache_manager = True
942
+ self.scheduler_config.disable_hybrid_kv_cache_manager = (
943
+ need_disable_hybrid_kv_cache_manager
944
+ )
945
+ elif (
946
+ self.scheduler_config.disable_hybrid_kv_cache_manager is False
947
+ and need_disable_hybrid_kv_cache_manager
948
+ ):
949
+ raise ValueError(
950
+ "Hybrid KV cache manager was explicitly enabled but is not "
951
+ "supported in this configuration. Consider omitting the "
952
+ "--no-disable-hybrid-kv-cache-manager flag to let vLLM decide"
953
+ " automatically."
954
+ )
955
+
956
+ if self.scheduler_config.disable_hybrid_kv_cache_manager is None:
957
+ # Default to enable HMA if not explicitly disabled by user or logic above.
958
+ self.scheduler_config.disable_hybrid_kv_cache_manager = False
959
960
if self.compilation_config.debug_dump_path:
961
self.compilation_config.debug_dump_path = (
vllm/engine/arg_utils.py
+1
-1
@@ -491,7 +491,7 @@ class EngineArgs:
491
enable_chunked_prefill: bool | None = None
492
disable_chunked_mm_input: bool = SchedulerConfig.disable_chunked_mm_input
493
494
- disable_hybrid_kv_cache_manager: bool = (
494
+ disable_hybrid_kv_cache_manager: bool | None = (
495
SchedulerConfig.disable_hybrid_kv_cache_manager
496
)
497