From 6b585a7378c36e4359b10209242bc833aacbea11 Mon Sep 17 00:00:00 2001 From: harrisonyhq Date: Sat, 22 Nov 2025 18:29:54 +0800 Subject: [PATCH] [Fix] fix mla bug when no broadcast in wait for save --- ucm/integration/vllm/ucm_connector.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ucm/integration/vllm/ucm_connector.py b/ucm/integration/vllm/ucm_connector.py index 138cda9b..333d0930 100644 --- a/ucm/integration/vllm/ucm_connector.py +++ b/ucm/integration/vllm/ucm_connector.py @@ -482,7 +482,7 @@ def save_kv_layer( def wait_for_save(self) -> None: - if self.load_only_first_rank and self.rank != 0: + if self.is_mla and self.rank != 0: return metadata = self._get_connector_metadata()