[CUDA Graph] Allow custom decode graph runners (#33553)

Co-authored-by: Itai Gat <itaigat.mail@gmail.com>

python/sglang/srt/model_executor/model_runner.py

@@ -1199,6 +1199,17 @@ class ModelRunner:
             model_runner=self, init_new_workspace=init_new_workspace
         )
 
+    def _decode_cuda_graph_runner_cls(self):
+        """Decode CUDA-graph runner class to construct.
+
+        Subclasses can override this to install specialized decode graph runners.
+        """
+        from sglang.srt.model_executor.runner.decode_cuda_graph_runner import (
+            DecodeCudaGraphRunner,
+        )
+
+        return DecodeCudaGraphRunner
+
     def init_decode_cuda_graph(self):
         self.decode_cuda_graph_runner = None
         self.graph_mem_usage = 0

python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py

@@ -373,12 +373,8 @@ def capture_decode_graph(*, model_runner: ModelRunner) -> DecodeGraphCapture:
         GraphRunnerCls = current_platform.get_graph_runner_cls()
         runner = GraphRunnerCls(model_runner)
     else:
-        from sglang.srt.model_executor.runner.decode_cuda_graph_runner import (
-            DecodeCudaGraphRunner,
-        )
-
         graph_runners = defaultdict(
-            lambda: DecodeCudaGraphRunner,
+            model_runner._decode_cuda_graph_runner_cls,
             {
                 "cpu": CPUGraphRunner,
                 "npu": NPUGraphRunner,