Skip to content

Commit 3a0647b

Browse files
[bugfix] Adapt the patch to support YAML sections. (#480)
1 parent 010844e commit 3a0647b

File tree

1 file changed

+4
-4
lines changed

1 file changed

+4
-4
lines changed

ucm/integration/vllm/patch/patch_funcs/v092/vllm_patch.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -553,14 +553,14 @@ def _patch_scheduler() -> None:
553553

554554
from ucm.sparse.base import INVALID_SLOT, UcmSparseRole
555555
from ucm.sparse.state import ensure_ucm_sparse_initialized, get_ucm_sparse
556+
from ucm.utils import Config
556557

557558
def init_ucm_sparse(self):
558559
self.ucm_sparse = None
559560
if self.vllm_config.kv_transfer_config is not None:
560-
if (
561-
"ucm_sparse_config"
562-
in self.vllm_config.kv_transfer_config.kv_connector_extra_config
563-
):
561+
ucm_config = Config(self.vllm_config.kv_transfer_config)
562+
ucm_sparse_config = ucm_config.get_config().get("ucm_sparse_config")
563+
if ucm_sparse_config:
564564
ensure_ucm_sparse_initialized(
565565
self.vllm_config, role=UcmSparseRole.SCHEDULER
566566
)

0 commit comments

Comments
 (0)