Use cuda event wait and synchronization instead of busy waiting (#2089)

2024-11-19 00:21:46 -08:00
parent b110453802
commit b7a065eae3
6 changed files with 28 additions and 26 deletions
--- a/python/sglang/srt/managers/tp_worker.py
+++ b/python/sglang/srt/managers/tp_worker.py
@@ -142,12 +142,12 @@ class TpModelWorker:
    def forward_batch_generation(
        self,
        model_worker_batch: ModelWorkerBatch,
-        launch_event: Optional[threading.Event] = None,
+        launch_done: Optional[threading.Event] = None,
    ):
        forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
        logits_output = self.model_runner.forward(forward_batch)
-        if launch_event:
-            launch_event.set()
+        if launch_done:
+            launch_done.set()
        next_token_ids = self.model_runner.sample(logits_output, model_worker_batch)
        return logits_output, next_token_ids