Fix OOM when updating expert locations (#6660)

python/sglang/srt/model_executor/expert_location_updater.py

@@ -27,21 +27,30 @@ from sglang.srt.managers.expert_location import (
 logger = logging.getLogger(__name__)
 
 
-def update_expert_location(
-    routed_experts_weights_of_layer: Dict[int, List[torch.Tensor]],
-    new_expert_location_metadata: ExpertLocationMetadata,
-    nnodes: int,
-    rank: int,
-):
-    old_expert_location_metadata = get_global_expert_location_metadata()
-    _update_expert_weights(
-        routed_experts_weights_of_layer,
-        old_expert_location_metadata,
-        new_expert_location_metadata,
-        nnodes,
-        rank,
-    )
-    old_expert_location_metadata.update(new_expert_location_metadata)
+class ExpertLocationUpdater:
+    def __init__(self):
+        self._first_execution = True
+
+    def update(
+        self,
+        routed_experts_weights_of_layer: Dict[int, List[torch.Tensor]],
+        new_expert_location_metadata: ExpertLocationMetadata,
+        nnodes: int,
+        rank: int,
+    ):
+        if self._first_execution:
+            self._first_execution = False
+            torch.cuda.empty_cache()
+
+        old_expert_location_metadata = get_global_expert_location_metadata()
+        _update_expert_weights(
+            routed_experts_weights_of_layer,
+            old_expert_location_metadata,
+            new_expert_location_metadata,
+            nnodes,
+            rank,
+        )
+        old_expert_location_metadata.update(new_expert_location_metadata)
 
 
 def _update_expert_weights(