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 7, 11:30 UTC
[KVConnector] Auto-downgrade to PIECEWISE cudagraph mode for layerwise async ops (#31057)
Signed-off-by: Yashwant Bezawada <yashwant_b@me.com>
Yashwant Bezawada committed
Mar 2, 2026 at 14:04 UTC
a13d8c03c996824811829d9f1cfff5d6df168271
4 files changed
+74
vllm/config/vllm.py
+27
@@ -925,6 +925,33 @@ class VllmConfig:
925
CUDAGraphMode.FULL_DECODE_ONLY
926
)
927
928
+ # Check if KV connector requires PIECEWISE mode for CUDA graphs
929
+ if (
930
+ self.kv_transfer_config is not None
931
+ and self.kv_transfer_config.is_kv_transfer_instance
932
+ and self.compilation_config.cudagraph_mode.has_full_cudagraphs()
933
+ ):
934
+ # Lazy import to avoid circular dependencies
935
+ from vllm.distributed.kv_transfer.kv_connector.factory import (
936
+ KVConnectorFactory,
937
+ )
938
+
939
+ connector_cls = KVConnectorFactory.get_connector_class(
940
+ self.kv_transfer_config
941
+ )
942
+ if connector_cls.requires_piecewise_for_cudagraph(
943
+ self.kv_transfer_config.kv_connector_extra_config
944
+ ):
945
+ logger.warning_once(
946
+ "KV connector %s requires PIECEWISE CUDA graph mode "
947
+ "due to layerwise async operations that cannot be "
948
+ "captured in CUDA graphs. "
949
+ "Overriding cudagraph_mode from %s to PIECEWISE.",
950
+ connector_cls.__name__,
951
+ self.compilation_config.cudagraph_mode.name,
952
+ )
953
+ self.compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
954
+
955
# disable cudagraph when enforce eager execution
956
if self.model_config is not None and self.model_config.enforce_eager:
957
logger.info("Cudagraph is disabled under eager mode")
vllm/distributed/kv_transfer/kv_connector/v1/base.py
+22
@@ -543,6 +543,28 @@ class KVConnectorBase_V1(ABC):
543
)
544
return None
545
546
+ @classmethod
547
+ def requires_piecewise_for_cudagraph(cls, extra_config: dict[str, Any]) -> bool:
548
+ """
549
+ Check if this connector requires PIECEWISE CUDA graph mode.
550
+
551
+ Connectors that use asynchronous layer-by-layer operations
552
+ (wait_for_layer_load/save_kv_layer) should override this method
553
+ to return True when those operations are enabled. These operations
554
+ cannot be captured in CUDA graphs and will be skipped during replay,
555
+ causing data races. PIECEWISE mode allows Python code to execute
556
+ between graph pieces, ensuring proper synchronization.
557
+
558
+ Args:
559
+ extra_config: The kv_connector_extra_config dict from
560
+ KVTransferConfig.
561
+
562
+ Returns:
563
+ True if this connector requires PIECEWISE CUDA graph mode,
564
+ False otherwise.
565
+ """
566
+ return False
567
+
568
def get_finished_count(self) -> int | None:
569
"""
570
Get the count of requests expected to complete send/receive operations
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_connector.py
+10
@@ -70,6 +70,16 @@ class LMCacheKVEvents(KVConnectorKVEvents):
70
71
72
class LMCacheConnectorV1(KVConnectorBase_V1):
73
+ @classmethod
74
+ def requires_piecewise_for_cudagraph(cls, extra_config: dict[str, Any]) -> bool:
75
+ """
76
+ LMCache requires PIECEWISE CUDA graph mode when layerwise
77
+ operations are enabled. The wait_for_layer_load and save_kv_layer
78
+ methods perform actual async synchronization that cannot be
79
+ captured in CUDA graphs.
80
+ """
81
+ return extra_config.get("use_layerwise", False)
82
+
83
def __init__(
84
self,
85
vllm_config: "VllmConfig",
vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
+15
@@ -112,6 +112,21 @@ class MultiConnector(KVConnectorBase_V1):
112
- Save to all connectors.
113
"""
114
115
+ @classmethod
116
+ def requires_piecewise_for_cudagraph(cls, extra_config: dict[str, Any]) -> bool:
117
+ """
118
+ MultiConnector requires PIECEWISE CUDA graph mode if any of its
119
+ child connectors require it.
120
+ """
121
+ connectors_config = extra_config.get("connectors", [])
122
+ for conn_config in connectors_config:
123
+ temp_ktc = KVTransferConfig(**conn_config)
124
+ connector_cls = KVConnectorFactory.get_connector_class(temp_ktc)
125
+ child_extra_config = conn_config.get("kv_connector_extra_config", {})
126
+ if connector_cls.requires_piecewise_for_cudagraph(child_extra_config):
127
+ return True
128
+ return False
129
+
130
def __init__(
131
self,
132
vllm_config: "VllmConfig",