[DLLM] Add initial cuda graph support (#14203)

2025-12-08 14:12:35 +08:00
parent 661e9775d0
commit 36361adcbf
6 changed files with 93 additions and 8 deletions
@@ -2050,10 +2050,16 @@ class ServerArgs:
        if self.dllm_algorithm is None:
            return
        if not self.disable_cuda_graph:
-            logger.warning(
-                "Cuda graph is disabled because of using diffusion LLM inference"
-            )
-            self.disable_cuda_graph = True
+            if self.cuda_graph_bs != [1]:
+                logger.warning(
+                    "Cuda graph bs is set to [1] because of using diffusion LLM inference"
+                )
+                self.cuda_graph_bs = [1]
+            if self.attention_backend != "flashinfer":
+                logger.warning(
+                    "Attention backend is set to flashinfer because of enabling cuda graph in diffusion LLM inference"
+                )
+                self.attention_backend = "flashinfer"
        if not self.disable_overlap_schedule:
            logger.warning(
                "Overlap schedule is disabled because of using diffusion LLM inference"