Reuse live CUDA graph executables during dedup registration (#39176)

Co-authored-by: cctry <17473714+cctry@users.noreply.github.com>

python/sglang/srt/model_executor/runner_backend/cuda_graph_dedup_mixin.py

@@ -43,20 +43,6 @@ def maybe_cuda_result(result):
     return None if int(result[0]) != 0 else checkCudaErrors(result)
 
 
-def kernel_name(params) -> str:
-    assert cuda_drv is not None
-    for handle, getter in (
-        (getattr(params, "kern", None), cuda_drv.cuKernelGetName),
-        (getattr(params, "func", None), cuda_drv.cuFuncGetName),
-    ):
-        if handle is None or int(handle) == 0:
-            continue
-        name = maybe_cuda_result(getter(handle))
-        if name is not None:
-            return name.decode("utf-8", "replace")
-    return f"func:{int(getattr(params, 'func', 0))}"
-
-
 def kernel_attrs(node) -> tuple[tuple[str, object], ...]:
     assert cuda_drv is not None
     attrs = []
@@ -107,10 +93,13 @@ def kernel_attrs(node) -> tuple[tuple[str, object], ...]:
 def kernel_node_payload(node):
     assert cuda_drv is not None
     params = checkCudaErrors(cuda_drv.cuGraphKernelNodeGetParams(node))
+    # Grid dimensions vary across buckets and are validated by the exec update.
+    # The handles subsume the kernel name; they are here because
+    # cuGraphKernelNodeGetAttribute cannot read preferred cluster dimension back
+    # and cudaGraphExecUpdate reports success when it changes (sgl-project/sglang#37657).
+    # Once the driver reports that mismatch, the handles can go too.
     return (
-        kernel_name(params),
         (int(params.kern), int(params.func)),
-        (int(params.gridDimX), int(params.gridDimY), int(params.gridDimZ)),
         (int(params.blockDimX), int(params.blockDimY), int(params.blockDimZ)),
         int(params.sharedMemBytes),
         kernel_attrs(node),
@@ -186,7 +175,6 @@ def graph_signature(raw_graph: int):
 class GraphExecGroup:
     graph_exec: int
     current_raw_graph: int
-    compat_exec: int | None
     graphs: list[DedupedCudaGraph] = field(default_factory=list)
 
 
@@ -227,9 +215,10 @@ class DedupedCudaGraphRegistry:
 
         group = self.groups.get(signature)
         if group is not None:
-            assert group.compat_exec is not None
+            # An incompatible cudaGraphExecUpdate does not modify the executable.
-            ok, detail = dedup_update(group.compat_exec, graph.raw_graph)
+            ok, detail = dedup_update(group.graph_exec, graph.raw_graph)
             assert ok, f"CUDA graph dedup register update failed ({detail})"
+            group.current_raw_graph = graph.raw_graph
             graph.group = group
             group.graphs.append(graph)
             return graph
@@ -237,7 +226,6 @@ class DedupedCudaGraphRegistry:
         group = GraphExecGroup(
             graph_exec=self.instantiate(graph.raw_graph),
             current_raw_graph=graph.raw_graph,
-            compat_exec=self.instantiate(graph.raw_graph),
             graphs=[graph],
         )
         graph.group = group
@@ -245,13 +233,7 @@ class DedupedCudaGraphRegistry:
         return graph
 
     def seal(self) -> None:
-        if self.sealed:
-            return
         self.sealed = True
-        for group in self.groups.values():
-            if group.compat_exec is not None:
-                self.destroy_exec(group.compat_exec)
-                group.compat_exec = None
 
     def stats(self) -> tuple[int, int]:
         return sum(len(group.graphs) for group in self.groups.values()), len(
@@ -281,9 +263,6 @@ class DedupedCudaGraphRegistry:
         self.sealed = True
 
         for group in self.groups.values():
-            if group.compat_exec is not None:
-                self.destroy_exec(group.compat_exec)
-                group.compat_exec = None
             self.destroy_exec(group.graph_exec)
             for graph in group.graphs:
                 if graph.original_graph is not None:

test/registered/kernel/cuda_graph/test_cuda_graph_dedup.py

@@ -0,0 +1,76 @@
+"""CUDA graph executable reuse across capture sizes and rejected updates."""
+
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from sglang.srt.model_executor.runner_backend import cuda_graph_dedup_mixin as dedup
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+class TestCudaGraphDedup(CustomTestCase):
+    @staticmethod
+    def capture(size, *, extra_node=False):
+        inputs = torch.zeros(size, device="cuda")
+        outputs = torch.empty_like(inputs)
+        graph = torch.cuda.CUDAGraph(keep_graph=True)
+        with torch.cuda.graph(graph):
+            torch.add(inputs, 1, out=outputs)
+            if extra_node:
+                outputs.mul_(2)
+        return graph, inputs, outputs
+
+    def test_grid_sizes_share_one_executable(self):
+        registry = dedup.DedupedCudaGraphRegistry()
+        self.addCleanup(registry.close)
+        captures = [self.capture(size) for size in (4096, 8192)]
+        self.assertEqual(
+            dedup.graph_signature(captures[0][0].raw_cuda_graph()),
+            dedup.graph_signature(captures[1][0].raw_cuda_graph()),
+        )
+        with patch.object(
+            registry, "instantiate", wraps=registry.instantiate
+        ) as instantiate:
+            graphs = [registry.register(capture[0]) for capture in captures]
+        self.assertEqual(instantiate.call_count, 1)
+        self.assertEqual(registry.stats(), (2, 1))
+        self.assertEqual(graphs[0].group.current_raw_graph, graphs[1].raw_graph)
+        registry.seal()
+        registry.seal()
+        for value in (3, 7):
+            for graph, (_, inputs, outputs) in zip(graphs, captures):
+                inputs.fill_(value)
+                graph.replay()
+                torch.cuda.synchronize()
+                self.assertTrue(
+                    torch.equal(outputs, torch.full_like(outputs, value + 1))
+                )
+
+    def test_rejected_registration_preserves_live_executable(self):
+        registry = dedup.DedupedCudaGraphRegistry()
+        self.addCleanup(registry.close)
+        original, inputs, outputs = self.capture(4096)
+        incompatible, _, _ = self.capture(4096, extra_node=True)
+        self.addCleanup(incompatible.reset)
+        graph = registry.register(original)
+        signature = dedup.graph_signature(graph.raw_graph)
+        with (
+            patch.object(dedup, "graph_signature", return_value=signature),
+            self.assertRaisesRegex(AssertionError, "register update failed"),
+        ):
+            registry.register(incompatible)
+        self.assertEqual(registry.stats(), (1, 1))
+        self.assertEqual(graph.group.current_raw_graph, graph.raw_graph)
+        inputs.fill_(11)
+        graph.replay()
+        torch.cuda.synchronize()
+        self.assertTrue(torch.equal(outputs, torch.full_like(outputs, 12)))
+
+
+if __name__ == "__main__":
+    unittest.main()