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 15, 08:30 UTC
[Core] Enable CUDA graphs for DP + All2All kernels (#18724)
Signed-off-by: Varun Sundar Rabindranath <varun@neuralmagic.com> Co-authored-by: Varun Sundar Rabindranath <varun@neuralmagic.com>
Varun Sundar Rabindranath committed
May 28, 2025 at 18:55 UTC
7951d78738581c336db7c1a77f94f1fea8f09fca
4 files changed
+100
-37
vllm/forward_context.py
+41
-22
@@ -10,7 +10,7 @@ import torch
10
import torch.distributed as dist
11
12
import vllm.envs as envs
13
-from vllm.config import VllmConfig
13
+from vllm.config import ParallelConfig, VllmConfig
14
from vllm.logger import init_logger
15
16
if TYPE_CHECKING:
@@ -30,6 +30,44 @@ class DPMetadata:
30
max_tokens_across_dp_cpu: torch.Tensor
31
cu_tokens_across_dp_cpu: torch.Tensor
32
33
+ @staticmethod
34
+ def num_tokens_across_dp(num_tokens: int, dp_size: int,
35
+ dp_rank: int) -> torch.Tensor:
36
+ """
37
+ Gather the num_tokens across all DP ranks and return results in a
38
+ CPU tensor of size dp_size.
39
+ """
40
+ num_tokens_across_dp = [0] * dp_size
41
+ num_tokens_across_dp[dp_rank] = num_tokens
42
+ num_tokens_tensor = torch.tensor(num_tokens_across_dp,
43
+ device="cpu",
44
+ dtype=torch.int32)
45
+ from vllm.distributed.parallel_state import get_dp_group
46
+ dist.all_reduce(num_tokens_tensor, group=get_dp_group().cpu_group)
47
+ return num_tokens_tensor
48
+
49
+ @staticmethod
50
+ def make(parallel_config: ParallelConfig, attn_metadata: Any,
51
+ num_tokens: int) -> "DPMetadata":
52
+
53
+ assert parallel_config.data_parallel_size > 1
54
+ dp_size = parallel_config.data_parallel_size
55
+ dp_rank = parallel_config.data_parallel_rank
56
+ if attn_metadata is not None and hasattr(attn_metadata,
57
+ "num_prefill_tokens"):
58
+ # for v0 attention backends
59
+ batchsize = attn_metadata.num_prefill_tokens + \
60
+ attn_metadata.num_decode_tokens
61
+ else:
62
+ # for v1 attention backends or no attn_metadata
63
+ batchsize = num_tokens
64
+
65
+ num_tokens_tensor = DPMetadata.num_tokens_across_dp(
66
+ batchsize, dp_size, dp_rank)
67
+ max_tokens_across_dp_cpu = torch.max(num_tokens_tensor)
68
+ cu_tokens_across_dp_cpu = torch.cumsum(num_tokens_tensor, dim=0)
69
+ return DPMetadata(max_tokens_across_dp_cpu, cu_tokens_across_dp_cpu)
70
+
71
72
@dataclass
73
class ForwardContext:
@@ -74,27 +112,8 @@ def set_forward_context(attn_metadata: Any,
112
forward_start_time = time.perf_counter()
113
dp_metadata: Optional[DPMetadata] = None
114
if vllm_config.parallel_config.data_parallel_size > 1:
77
- dp_size = vllm_config.parallel_config.data_parallel_size
78
- dp_rank = vllm_config.parallel_config.data_parallel_rank
79
- if attn_metadata is not None and hasattr(attn_metadata,
80
- "num_prefill_tokens"):
81
- # for v0 attention backends
82
- batchsize = attn_metadata.num_prefill_tokens + \
83
- attn_metadata.num_decode_tokens
84
- else:
85
- # for v1 attention backends or no attn_metadata
86
- batchsize = num_tokens
87
- num_tokens_across_dp = [0] * dp_size
88
- num_tokens_across_dp[dp_rank] = batchsize
89
- num_tokens_tensor = torch.tensor(num_tokens_across_dp,
90
- device="cpu",
91
- dtype=torch.int32)
92
- from vllm.distributed.parallel_state import get_dp_group
93
- dist.all_reduce(num_tokens_tensor, group=get_dp_group().cpu_group)
94
- max_tokens_across_dp_cpu = torch.max(num_tokens_tensor)
95
- cu_tokens_across_dp_cpu = torch.cumsum(num_tokens_tensor, dim=0)
96
- dp_metadata = DPMetadata(max_tokens_across_dp_cpu,
97
- cu_tokens_across_dp_cpu)
115
+ dp_metadata = DPMetadata.make(vllm_config.parallel_config,
116
+ attn_metadata, num_tokens)
117
118
global _forward_context
119
prev_context = _forward_context
vllm/model_executor/layers/fused_moe/layer.py
+39
-3
@@ -828,6 +828,21 @@ class FusedMoE(torch.nn.Module):
828
829
self.quant_method.create_weights(layer=self, **moe_quant_params)
830
831
+ # Chunked all2all staging tensor
832
+ self.batched_hidden_states: Optional[torch.Tensor] = None
833
+ self.batched_router_logits: Optional[torch.Tensor] = None
834
+ if self.moe_parallel_config.use_pplx_kernels:
835
+ act_dtype = vllm_config.model_config.dtype
836
+ self.batched_hidden_states = torch.zeros(
837
+ (MOE_DP_CHUNK_SIZE, self.hidden_size),
838
+ dtype=act_dtype,
839
+ device=torch.cuda.current_device())
840
+
841
+ self.batched_router_logits = torch.zeros(
842
+ (MOE_DP_CHUNK_SIZE, self.global_num_experts),
843
+ dtype=act_dtype,
844
+ device=torch.cuda.current_device())
845
+
846
@property
847
def tp_size(self):
848
return self.moe_parallel_config.tp_size
@@ -1217,18 +1232,39 @@ class FusedMoE(torch.nn.Module):
1232
1233
def forward_impl_chunked(self, full_hidden_states: torch.Tensor,
1234
full_router_logits: torch.Tensor):
1235
+ assert self.batched_hidden_states is not None
1236
+ assert self.batched_router_logits is not None
1237
+ assert self.batched_hidden_states.dtype == full_hidden_states.dtype
1238
+ assert self.batched_router_logits.dtype == full_router_logits.dtype
1239
+ # Check size compatibility.
1240
+ assert (
1241
+ self.batched_hidden_states.size(-1) == full_hidden_states.size(-1))
1242
+ assert (
1243
+ self.batched_router_logits.size(-1) == full_router_logits.size(-1))
1244
1245
full_final_hidden_states = torch.empty_like(full_hidden_states)
1246
1247
def process_chunk(chunk_start, chunk_end, skip_result_store=False):
1248
+ chunk_size = chunk_end - chunk_start
1249
hidden_states = full_hidden_states[chunk_start:chunk_end, :]
1250
router_logits = full_router_logits[chunk_start:chunk_end, :]
1251
1252
+ assert (self.batched_hidden_states.size(0) # type: ignore
1253
+ >= chunk_size)
1254
+ assert (self.batched_router_logits.size(0) # type: ignore
1255
+ >= chunk_size)
1256
+ staged_hidden_states = self.batched_hidden_states[:
1257
+ chunk_size, :] # type: ignore
1258
+ staged_router_logits = self.batched_router_logits[:
1259
+ chunk_size, :] # type: ignore
1260
+ staged_hidden_states.copy_(hidden_states, non_blocking=True)
1261
+ staged_router_logits.copy_(router_logits, non_blocking=True)
1262
+
1263
# Matrix multiply.
1264
final_hidden_states = self.quant_method.apply(
1265
layer=self,
1230
- x=hidden_states,
1231
- router_logits=router_logits,
1266
+ x=staged_hidden_states,
1267
+ router_logits=staged_router_logits,
1268
top_k=self.top_k,
1269
renormalize=self.renormalize,
1270
use_grouped_topk=self.use_grouped_topk,
@@ -1244,7 +1280,7 @@ class FusedMoE(torch.nn.Module):
1280
1281
if not skip_result_store:
1282
full_final_hidden_states[chunk_start:chunk_end, :].copy_(
1247
- final_hidden_states)
1283
+ final_hidden_states, non_blocking=True)
1284
1285
ctx = get_forward_context()
1286
max_tokens_across_dp = ctx.dp_metadata.max_tokens_across_dp_cpu
vllm/platforms/cuda.py
-11
@@ -106,7 +106,6 @@ class CudaPlatformBase(Platform):
106
def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
107
parallel_config = vllm_config.parallel_config
108
scheduler_config = vllm_config.scheduler_config
109
- compilation_config = vllm_config.compilation_config
109
model_config = vllm_config.model_config
110
111
if parallel_config.worker_cls == "auto":
@@ -154,16 +153,6 @@ class CudaPlatformBase(Platform):
153
logger.info(
154
"Forcing kv cache block size to 64 for FlashMLA backend.")
155
157
- if (parallel_config.data_parallel_size > 1
158
- and compilation_config.use_cudagraph):
159
- logger.info(
160
- "Data Parallel: Forcing enforce eager to be True since DP is "
161
- "currently not supported with CUDA Graphs.")
162
- vllm_config.model_config.enforce_eager = True
163
- compilation_config.use_cudagraph = False
164
- # FIXME: inductor breaks cudagraph (from @bnell)
165
- compilation_config.use_inductor = False
166
-
156
@classmethod
157
def get_current_memory_usage(cls,
158
device: Optional[torch.types.Device] = None
vllm/v1/worker/gpu_model_runner.py
+20
-1
@@ -24,7 +24,8 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1
24
from vllm.distributed.parallel_state import (
25
get_pp_group, get_tp_group, graph_capture,
26
prepare_communication_buffer_for_model)
27
-from vllm.forward_context import get_forward_context, set_forward_context
27
+from vllm.forward_context import (DPMetadata, get_forward_context,
28
+ set_forward_context)
29
from vllm.logger import init_logger
30
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
31
from vllm.model_executor.model_loader import TensorizerLoader, get_model
@@ -1104,6 +1105,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1105
for k, v in self.intermediate_tensors.items()
1106
})
1107
1108
+ def get_dp_padding(self, num_tokens: int):
1109
+ dp_size = self.vllm_config.parallel_config.data_parallel_size
1110
+ dp_rank = self.vllm_config.parallel_config.data_parallel_rank
1111
+ if dp_size == 1:
1112
+ # Early exit.
1113
+ return 0
1114
+
1115
+ num_tokens_across_dp = DPMetadata.num_tokens_across_dp(
1116
+ num_tokens, dp_size, dp_rank)
1117
+ max_tokens_across_dp_cpu = torch.max(num_tokens_across_dp).item()
1118
+ return max_tokens_across_dp_cpu - num_tokens
1119
+
1120
@torch.inference_mode()
1121
def execute_model(
1122
self,
@@ -1141,6 +1154,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1154
else:
1155
num_input_tokens = num_scheduled_tokens
1156
1157
+ # Padding for DP
1158
+ num_input_tokens += self.get_dp_padding(num_input_tokens)
1159
+
1160
# _prepare_inputs may reorder the batch, so we must gather multi
1161
# modal outputs after that to ensure the correct order
1162
if self.is_multimodal_model:
@@ -1658,6 +1674,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1674
skip_attn: bool = True,
1675
) -> torch.Tensor:
1676
1677
+ # Padding for DP
1678
+ num_tokens += self.get_dp_padding(num_tokens)
1679
+
1680
# Set num_scheduled_tokens based on num_tokens and max_num_seqs
1681
# for dummy run with LoRA so that the num_reqs collectively
1682
# has num_tokens in total.