use schedule_split_pd (dpsk mtp + fullgraph)

vllm/v1/core/sched/scheduler.py

@@ -153,6 +153,10 @@ class Scheduler(SchedulerInterface):
                 self.use_eagle = True
                 self.num_lookahead_tokens = self.num_spec_tokens
         
+        self.compilation_config = vllm_config.compilation_config
+        self.full_cuda_graph = self.compilation_config.full_cuda_graph
+        self.use_mla = vllm_config.model_config.use_mla
+
         # Create the KV cache manager.
         self.kv_cache_manager = KVCacheManager(
             kv_cache_config=kv_cache_config,
@@ -1024,7 +1028,7 @@ class Scheduler(SchedulerInterface):
         return scheduler_output
     
     def schedule(self) -> SchedulerOutput:
-        if envs.VLLM_USE_PD_SPLIT:
+        if envs.VLLM_USE_PD_SPLIT or (self.full_cuda_graph and self.use_mla and self.num_spec_tokens > 0) :
             return self.schedule_split_pd()
         else:
             return self.schedule_default()