add VLLM_USE_PD_SPLIT to split prefill and decode

vllm/envs.py

@@ -176,6 +176,7 @@ if TYPE_CHECKING:
     VLLM_P2P_ASYNC: bool = False
     VLLM_P2P_BUF_TOKENS: int = 30000
     VLLM_ENABLE_OUTPUT_PLACEHOLDERS: bool = False
+    VLLM_USE_PD_SPLIT: bool = False
 
 def get_default_cache_root():
     return os.getenv(
@@ -1092,7 +1093,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_USE_FLASH_ATTN_PA":
     lambda: (os.environ.get("VLLM_USE_FLASH_ATTN_PA", "True").lower() in
              ("true", "1")),
-    
     # vLLM will use apex for rmsnorm
     "VLLM_USE_APEX_RN":
     lambda: (os.environ.get("VLLM_USE_APEX_RN", "False").lower() in
@@ -1146,6 +1146,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # vllm will enable output placeholders
     "VLLM_ENABLE_OUTPUT_PLACEHOLDERS":
     lambda: bool(int(os.getenv("VLLM_ENABLE_OUTPUT_PLACEHOLDERS", "0"))),
+    # vLLM will split prefill and decode, not mix up
+    "VLLM_USE_PD_SPLIT":
+        lambda: (os.environ.get("VLLM_USE_PD_SPLIT", "True").lower() in
+                 ("true", "1")), 
 }
 
 # --8<-- [end:env-vars-definition]

vllm/v1/core/sched/scheduler.py

@@ -1014,7 +1014,7 @@ class Scheduler(SchedulerInterface):
         return scheduler_output
     
     def schedule(self) -> SchedulerOutput:
-        if self.num_spec_tokens > 0:
+        if self.num_spec_tokens > 0 or envs.VLLM_USE_PD_SPLIT:
             return self.schedule_split_pd()
         else:
             return self.schedule_default()