fix performance

yukavio · Aug 15, 2024 · a2b9a62 · a2b9a62
1 parent 76c0087
commit a2b9a62
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/python/sglang/srt/managers/controller_flex.py b/python/sglang/srt/managers/controller_flex.py
@@ -156,7 +156,7 @@ def resources_aware_scheduler(self, input_requests):
             remained_token[index] += len(r.input_ids)
         with self.controller_info.lock:
             for i, v in enumerate(remained_token):
-                self.controller_info.current_bs[i] += v
+                self.controller_info.current_bs[i].value = v
 
     def round_robin_scheduler(self, input_requests):
         for r in input_requests:

diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py
@@ -460,11 +460,11 @@ def forward_prefill_batch(self, batch: ScheduleBatch):
         )
 
         if self.controller_info:
+            num = 0
+            for r in batch.reqs:
+                num += len(r.origin_input_ids)
             with self.controller_info.lock:
-                for r in batch.reqs:
-                    self.controller_info.current_bs[self.dp_rank].value -= len(
-                        r.origin_input_ids
-                    )
+                self.controller_info.current_bs[self.dp_rank].value -= num
 
         if self.model_runner.is_generation:
             # Forward and sample the next tokens