A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 01:30 UTC
[Bugfix] Refactor to support DP parallel in R3 (#32306)
Signed-off-by: xhx1022 <1737006628@qq.com> Co-authored-by: arlenxu <arlenxu@tencent.com>
Hongxin Xu committed
Jan 16, 2026 at 15:13 UTC
03da3b52ef6f3f3d9af5c1942c93a9ace7c2a6bf
3 files changed
+35
-26
vllm/model_executor/layers/fused_moe/routed_experts_capturer.py
+33
-21
@@ -17,8 +17,9 @@ from unittest.mock import patch
17
import numpy as np
18
import torch
19
20
-from vllm.config import ModelConfig
20
+from vllm.config import VllmConfig
21
from vllm.distributed import get_tensor_model_parallel_rank
22
+from vllm.forward_context import get_forward_context
23
24
logger = logging.getLogger(__name__)
25
@@ -89,7 +90,6 @@ class RoutedExpertsCapturer:
90
self._shm: shared_memory.SharedMemory | None = None
91
self._host_buffer_view: np.ndarray | None = None
92
self._lock_file: str | None = None
92
- self._shm_name: str | None = None
93
94
@classmethod
95
def create(cls) -> RoutedExpertsCapturer:
@@ -110,8 +110,7 @@ class RoutedExpertsCapturer:
110
self,
111
max_num_batched_tokens: int,
112
max_num_kv_tokens: int,
113
- model_config: ModelConfig,
114
- instance_id: str,
113
+ vllm_config: VllmConfig,
114
) -> None:
115
"""
116
Initialize the device buffer and optionally shared memory buffer.
@@ -119,14 +118,13 @@ class RoutedExpertsCapturer:
118
Args:
119
max_num_batched_tokens: Maximum number of tokens in a batch.
120
max_num_kv_tokens: Maximum number of KV tokens for shared memory.
122
- model_config: Model configuration containing layer and expert info.
123
- instance_id: Unique identifier for the shared memory buffer.
121
+ vllm_config: vllm configuration containing layer and expert info.
122
"""
123
124
if self._device_buffer is not None:
125
raise RuntimeError("Device buffer has already been initialized")
126
129
- hf_config = model_config.hf_text_config
127
+ hf_config = vllm_config.model_config.hf_text_config
128
num_layers = hf_config.num_hidden_layers
129
num_experts_per_tok = hf_config.num_experts_per_tok
130
@@ -136,6 +134,7 @@ class RoutedExpertsCapturer:
134
dtype=torch.int32,
135
device="cuda",
136
)
137
+ self.dp_rank = vllm_config.parallel_config.data_parallel_rank
138
139
if get_tensor_model_parallel_rank() != 0:
140
return
@@ -143,19 +142,19 @@ class RoutedExpertsCapturer:
142
# Initialize shared memory
143
shape = (max_num_kv_tokens, num_layers, num_experts_per_tok)
144
buffer_size = int(np.prod(shape)) * np.dtype(np.int32).itemsize
146
-
147
- self._lock_file = f"{_LOCK_FILE_PREFIX}_{instance_id}.lock"
148
- self._shm_name = f"{_BUFFER_PREFIX}_{instance_id}"
145
+ instance_id = vllm_config.instance_id
146
+ self._lock_file = f"{_LOCK_FILE_PREFIX}_{instance_id}_{self.dp_rank}.lock"
147
+ shm_name = f"{_BUFFER_PREFIX}_{instance_id}_{self.dp_rank}"
148
149
self._shm = _create_or_attach_shared_memory(
151
- self._shm_name, buffer_size, self._lock_file
150
+ shm_name, buffer_size, self._lock_file
151
)
152
self._host_buffer_view = np.ndarray(shape, dtype=np.int32, buffer=self._shm.buf)
153
self._host_buffer_view.fill(0)
154
155
logger.debug(
156
"Created shared memory buffer '%s' with shape %s",
158
- self._shm.name,
157
+ shm_name,
158
shape,
159
)
160
@@ -170,11 +169,24 @@ class RoutedExpertsCapturer:
169
if self._device_buffer is None:
170
raise RuntimeError("Buffer not initialized. Call init_buffer() first.")
171
172
+ ctx = get_forward_context()
173
+ if ctx.dp_metadata is None: # single dp
174
+ start_loc = 0
175
+ end_loc = topk_ids.shape[0]
176
+ token_num_per_dp = topk_ids.shape[0]
177
+ else: # multi dp
178
+ token_num_per_dp = ctx.dp_metadata.num_tokens_across_dp_cpu[self.dp_rank]
179
+ cumsum = torch.cumsum(ctx.dp_metadata.num_tokens_across_dp_cpu, dim=0)
180
+ assert cumsum[-1] == topk_ids.shape[0]
181
+ end_loc = cumsum[self.dp_rank]
182
+ start_loc = end_loc - token_num_per_dp
183
+
184
if layer_id >= self._device_buffer.shape[1]:
185
return
186
176
- batch_size = topk_ids.shape[0]
177
- self._device_buffer[:batch_size, layer_id, :] = topk_ids
187
+ self._device_buffer[:token_num_per_dp, layer_id, :] = topk_ids[
188
+ start_loc:end_loc, :
189
+ ]
190
191
def clear_buffer(self) -> None:
192
"""Clear the device buffer."""
@@ -254,30 +266,30 @@ class RoutedExpertsReader:
266
def attach_buffer(
267
self,
268
max_num_kv_tokens: int,
257
- model_config: ModelConfig,
258
- instance_id: str,
269
+ vllm_config: VllmConfig,
270
) -> None:
271
"""
272
Attach to an existing shared memory buffer.
273
274
Args:
275
max_num_kv_tokens: Maximum number of KV tokens.
265
- model_config: Model configuration.
266
- instance_id: Unique identifier for the shared memory buffer.
276
+ vllm_config: vllm configuration.
277
"""
278
if self._shm is not None:
279
logger.warning("Already attached to shared memory buffer.")
280
return # Already attached
281
272
- hf_config = model_config.hf_text_config
282
+ hf_config = vllm_config.model_config.hf_text_config
283
shape = (
284
max_num_kv_tokens,
285
hf_config.num_hidden_layers,
286
hf_config.num_experts_per_tok,
287
)
288
279
- self._lock_file = f"{_LOCK_FILE_PREFIX}_{instance_id}.lock"
280
- shm_name = f"{_BUFFER_PREFIX}_{instance_id}"
289
+ self.dp_rank = vllm_config.parallel_config.data_parallel_rank
290
+ instance_id = vllm_config.instance_id
291
+ self._lock_file = f"{_LOCK_FILE_PREFIX}_{instance_id}_{self.dp_rank}.lock"
292
+ shm_name = f"{_BUFFER_PREFIX}_{instance_id}_{self.dp_rank}"
293
294
with _file_lock(self._lock_file, mode="rb+"):
295
# Avoid resource_tracker registering the shared memory
vllm/v1/core/sched/scheduler.py
+1
-2
@@ -245,8 +245,7 @@ class Scheduler(SchedulerInterface):
245
246
self.routed_experts_reader.attach_buffer(
247
max_num_kv_tokens=self.max_num_kv_tokens,
248
- model_config=self.vllm_config.model_config,
249
- instance_id=self.vllm_config.instance_id,
248
+ vllm_config=self.vllm_config,
249
)
250
251
def schedule(self) -> SchedulerOutput:
vllm/v1/worker/gpu_model_runner.py
+1
-3
@@ -5675,12 +5675,10 @@ class GPUModelRunner(
5675
self.kv_cache_config.num_blocks // len(self.kv_cache_config.kv_cache_groups)
5676
+ 1
5677
) * block_size
5678
-
5678
routed_experts_capturer.init_buffer(
5679
max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens,
5680
max_num_kv_tokens=self.max_num_kv_tokens,
5682
- model_config=self.model_config,
5683
- instance_id=self.vllm_config.instance_id,
5681
+ vllm_config=self.vllm_config,
5682
)
5683
5684
def may_add_encoder_only_layers_to_kv_cache_config(self) -> None: