Reworked fast_pos_embed_interpolate() using torch (#10959)

python/sglang/srt/utils/multi_stream_utils.py

@@ -37,8 +37,7 @@ def with_multi_stream(enable: bool):
 def maybe_execute_in_parallel(
     fn0: Callable,
     fn1: Callable,
-    event0: torch.cuda.Event,
-    event1: torch.cuda.Event,
+    events: list[torch.cuda.Event],
     aux_stream: Optional[torch.cuda.Stream] = None,
 ) -> tuple[Any, Any]:
     """Utility function to run two functions in two cuda streams in parallel. Multi-stream is
@@ -51,8 +50,7 @@ def maybe_execute_in_parallel(
     Args:
         fn0 (Callable): callable for the default stream
         fn1 (Callable): callable for the second stream, aux_stream
-        event0 (torch.cuda.Event): cuda event for fn0
-        event1 (torch.cuda.Event): cuda event for fn1
+        events (list[torch.cuda.Event]): cuda events for callables
         aux_stream (Optional[torch.cuda.Stream]): the second cuda stream for fn1.
             Multi-stream is disabled when aux_stream is None.
 
@@ -63,14 +61,14 @@ def maybe_execute_in_parallel(
     multi_stream = do_multi_stream() and aux_stream is not None
 
     if multi_stream:
-        event0.record()
+        events[0].record()
         result0 = fn0()
 
         with torch.cuda.stream(aux_stream):
-            event0.wait()
+            events[0].wait()
             result1 = fn1()
-            event1.record()
-        event1.wait()
+            events[1].record()
+        events[1].wait()
     else:
         result0 = fn0()
         result1 = fn1()