[PP][XPU]Add the flag to control microbatch feature on MRV2+PP (#55145)
Signed-off-by: yisheng <yi.sheng@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
YiSheng5 committed
Sep 25, 2026 at 12:05 UTC
b05f7d5fbc1b6005fd0273d3d9f49883217ca870
3 files changed
+21
-2
vllm/envs.py
+6
@@ -63,6 +63,7 @@ if TYPE_CHECKING:
63
VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: Literal["auto", "nccl", "shm"] = "auto"
64
VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM: bool = False
65
VLLM_USE_RAY_WRAPPED_PP_COMM: bool = True
66
+ VLLM_XPU_PP_MICROBATCH: bool = True
67
VLLM_USE_RAY_V2_EXECUTOR_BACKEND: bool = False
68
VLLM_DISTRIBUTED_USE_SPLIT_GROUP: bool = False
69
VLLM_XLA_USE_SPMD: bool = False
@@ -924,6 +925,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
925
"VLLM_USE_RAY_WRAPPED_PP_COMM": lambda: bool(
926
int(os.getenv("VLLM_USE_RAY_WRAPPED_PP_COMM", "1"))
927
),
928
+ # Using a flag to control microbatch on XPU device, users on XPU device can
929
+ # decide to disable it when they needed.
930
+ "VLLM_XPU_PP_MICROBATCH": lambda: bool(
931
+ int(os.getenv("VLLM_XPU_PP_MICROBATCH", "1"))
932
+ ),
933
# When True and distributed_executor_backend="ray", use RayExecutorV2
934
# (MQ-based) instead of RayDistributedExecutor (compiled-graph backend).
935
"VLLM_USE_RAY_V2_EXECUTOR_BACKEND": lambda: bool(
vllm/v1/core/sched/async_scheduler.py
+9
-1
@@ -1,7 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import vllm.envs as envs
5
from vllm.logger import init_logger
6
+from vllm.platforms import current_platform
7
from vllm.v1.core.sched.output import SchedulerOutput
8
from vllm.v1.core.sched.scheduler import Scheduler
9
from vllm.v1.request import Request, RequestStatus
@@ -15,6 +17,10 @@ class AsyncScheduler(Scheduler):
17
# reusable read-only placeholder list for speculative decoding.
18
self._spec_token_placeholders: list[int] = [-1] * self.num_spec_tokens
19
self.pp_size = self.parallel_config.pipeline_parallel_size
20
+ # Only XPU can disable PP decode microbatching via VLLM_XPU_PP_MICROBATCH.
21
+ self.decode_stagger = self.pp_size
22
+ if current_platform.is_xpu() and not envs.VLLM_XPU_PP_MICROBATCH:
23
+ self.decode_stagger = 1
24
25
def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None:
26
super()._update_after_schedule(scheduler_output)
@@ -46,7 +52,9 @@ class AsyncScheduler(Scheduler):
52
if self.use_v2_model_runner:
53
# Set the next step index in which this request is eligible to be
54
# scheduled for decode (for PP microbatching).
49
- request.next_decode_eligible_step = self.current_step + self.pp_size
55
+ request.next_decode_eligible_step = (
56
+ self.current_step + self.decode_stagger
57
+ )
58
59
def _update_request_with_output(
60
self, request: Request, new_token_ids: list[int], is_stale: bool = False
vllm/v1/worker/gpu/pp_utils.py
+6
-1
@@ -8,6 +8,7 @@ from dataclasses import dataclass
8
import numpy as np
9
import torch
10
11
+import vllm.envs as envs
12
from vllm.distributed.parallel_state import get_pp_group
13
from vllm.platforms import current_platform
14
from vllm.sequence import IntermediateTensors
@@ -71,8 +72,12 @@ class PPHandler:
72
# pushed by step T's `receive` is consumed pp_size steps later. Pre-seeded
73
# with pp_size None placeholders so the first pp_size consumes are no-ops.
74
# None means no postprocess is pending for that step (broadcast skipped).
75
+ # Only XPU can disable microbatching via VLLM_XPU_PP_MICROBATCH.
76
+ ring_depth = get_pp_group().world_size
77
+ if current_platform.is_xpu() and not envs.VLLM_XPU_PP_MICROBATCH:
78
+ ring_depth = 1
79
self.queue: deque[PendingRecv | None] = (
75
- deque() if self.is_last_rank else deque([None] * get_pp_group().world_size)
80
+ deque() if self.is_last_rank else deque([None] * ring_depth)
81
)
82
83
# Per req-index generation counter, incremented every time a request