Custom All Reduce for Piecewise Cuda Graph (#15356)

Signed-off-by: Oasis-Git <ayw.sirius19@gmail.com>

python/sglang/srt/compilation/compilation_config.py

@@ -18,6 +18,7 @@ class CompilationConfig:
         self.split_ops = [
             "sglang.unified_attention_with_output",
             "sglang.gdn_with_output",
+            "sglang.inplace_all_reduce",
         ]
 
     def add_split_op(self, op: str):

python/sglang/srt/compilation/cuda_piecewise_backend.py

@@ -11,6 +11,10 @@ import torch.fx as fx
 
 from sglang.srt.compilation.compilation_config import CompilationConfig
 from sglang.srt.compilation.compilation_counter import compilation_counter
+from sglang.srt.compilation.piecewise_context_manager import (
+    get_pcg_capture_stream,
+    is_in_pcg_torch_compile,
+)
 from sglang.srt.compilation.weak_ref_tensor import weak_ref_tensors
 
 logger = logging.getLogger(__name__)
@@ -105,6 +109,10 @@ class CUDAPiecewiseBackend:
             self.first_run_finished = True
             self.check_for_ending_compilation()
             return self.compiled_graph_for_general_shape(*args)
+
+        if len(self.sym_shape_indices) == 0:
+            return self.compiled_graph_for_general_shape(*args)
+
         runtime_shape = args[self.sym_shape_indices[0]]
         if runtime_shape not in self.concrete_size_entries:
             # we don't need to do anything for this shape
@@ -137,6 +145,8 @@ class CUDAPiecewiseBackend:
         # skip_cuda_graphs = get_forward_context().skip_cuda_graphs
         # if not entry.use_cudagraph or skip_cuda_graphs:
         #     return entry.runnable(*args)
+        if is_in_pcg_torch_compile():
+            return entry.runnable(*args)
 
         if entry.cudagraph is None:
             if entry.num_finished_warmup < 1:  # noqa
@@ -160,9 +170,10 @@ class CUDAPiecewiseBackend:
                     # and disable gc for the rest of the graphs.
                     stack.enter_context(patch("gc.collect", lambda: None))
                     stack.enter_context(patch("torch.cuda.empty_cache", lambda: None))
-
                 # mind-exploding: carefully manage the reference and memory.
-                with torch.cuda.graph(cudagraph, pool=self.graph_pool):
+                stream = get_pcg_capture_stream()
+                assert stream is not None, "PCG capture stream is not set"
+                with torch.cuda.graph(cudagraph, pool=self.graph_pool, stream=stream):
                     # `output` is managed by pytorch's cudagraph pool
                     output = entry.runnable(*args)
                     if self.is_last_graph:
@@ -194,6 +205,5 @@ class CUDAPiecewiseBackend:
                 "Input addresses for cudagraphs are different during replay."
                 f" Expected {entry.input_addresses}, got {new_input_addresses}"
             )
-
         entry.cudagraph.replay()
         return entry.output

python/sglang/srt/compilation/piecewise_context_manager.py

@@ -2,15 +2,35 @@ from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import Any, List, Optional
 
+import torch
+
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 
 _in_piecewise_cuda_graph = False
+_in_pcg_torch_compile = False
+_pcg_capture_stream = None
 
 
 def is_in_piecewise_cuda_graph():
     return _in_piecewise_cuda_graph
 
 
+def is_in_pcg_torch_compile():
+    return _in_pcg_torch_compile
+
+
+def get_pcg_capture_stream():
+    return _pcg_capture_stream
+
+
+@contextmanager
+def enable_piecewise_cuda_graph_compile():
+    global _in_pcg_torch_compile
+    _in_pcg_torch_compile = True
+    yield
+    _in_pcg_torch_compile = False
+
+
 @contextmanager
 def enable_piecewise_cuda_graph():
     global _in_piecewise_cuda_graph
@@ -21,6 +41,14 @@ def enable_piecewise_cuda_graph():
     _in_piecewise_cuda_graph = False
 
 
+@contextmanager
+def set_pcg_capture_stream(stream: torch.cuda.Stream):
+    global _pcg_capture_stream
+    _pcg_capture_stream = stream
+    yield
+    _pcg_capture_stream = None
+
+
 @dataclass
 class ForwardContext:
     def __init__(self):