diff --git a/python/sglang/srt/model_executor/runner_backend/cuda_graph_dedup_mixin.py b/python/sglang/srt/model_executor/runner_backend/cuda_graph_dedup_mixin.py index 0e7488e8b..293ceb484 100644 --- a/python/sglang/srt/model_executor/runner_backend/cuda_graph_dedup_mixin.py +++ b/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 - ok, detail = dedup_update(group.compat_exec, graph.raw_graph) + # An incompatible cudaGraphExecUpdate does not modify the executable. + 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: diff --git a/test/registered/kernel/cuda_graph/test_cuda_graph_dedup.py b/test/registered/kernel/cuda_graph/test_cuda_graph_dedup.py new file mode 100644 index 000000000..2976937a1 --- /dev/null +++ b/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()