[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
@@ -1318,7 +1318,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
        ), f"Expected {len(self.out_cache_loc)}, got {self.extend_num_tokens}"

    def prepare_for_extend(self):
-        self.forward_mode = ForwardMode.EXTEND
+        self.forward_mode = (
+            ForwardMode.DLLM_EXTEND if self.is_dllm() else ForwardMode.EXTEND
+        )

        # Init tensors
        reqs = self.reqs