[Fix] fix mla bug when no broadcast in wait for save (#398)

Author: harrisonyhq

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()