[Fix] fix allreduce bug in Piecewise Graph (#12106)

2025-10-26 21:15:48 +08:00
parent ea385ae85a
commit 96a5a949f6
3 changed files with 12 additions and 14 deletions
@@ -392,7 +392,7 @@ class SGLangBackend:
        self.configure_post_pass()

        self.split_gm, self.piecewise_graphs = split_graph(
-            graph, ["sglang.unified_attention_with_output"]
+            graph, ["sglang.unified_attention_with_output", "sglang.inplace_all_reduce"]
        )

        from torch._dynamo.utils import lazy_format_graph_code