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, 07:15 UTC
[Misc] reuse num_tokens_across_dp of get_dp_padding to avoid unnecessary dp all reduce in set_forward_context (#18935)
Signed-off-by: Tyler Michael Smith <tysmith@redhat.com> Co-authored-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com> Co-authored-by: Tyler Michael Smith <tysmith@redhat.com>
zhrrr committed
Jun 2, 2025 at 03:41 UTC
d6fd3a33b863003929fc3eef5dd9828219e04ab0
2 files changed
+47
-18
vllm/forward_context.py
+19
-8
@@ -47,8 +47,12 @@ class DPMetadata:
47
return num_tokens_tensor
48
49
@staticmethod
50
- def make(parallel_config: ParallelConfig, attn_metadata: Any,
51
- num_tokens: int) -> "DPMetadata":
50
+ def make(
51
+ parallel_config: ParallelConfig,
52
+ attn_metadata: Any,
53
+ num_tokens: int,
54
+ num_tokens_across_dp: Optional[torch.Tensor] = None
55
+ ) -> "DPMetadata":
56
57
assert parallel_config.data_parallel_size > 1
58
dp_size = parallel_config.data_parallel_size
@@ -62,10 +66,15 @@ class DPMetadata:
66
# for v1 attention backends or no attn_metadata
67
batchsize = num_tokens
68
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
+ # If num_tokens_across_dp is None, it will be computed by all_reduce
70
+ # Otherwise, num_tokens_across_dp[dp_rank] should be equal to batchsize
71
+ assert (num_tokens_across_dp is None
72
+ or num_tokens_across_dp[dp_rank] == batchsize)
73
+ if num_tokens_across_dp is None:
74
+ num_tokens_across_dp = DPMetadata.num_tokens_across_dp(
75
+ batchsize, dp_size, dp_rank)
76
+ max_tokens_across_dp_cpu = torch.max(num_tokens_across_dp)
77
+ cu_tokens_across_dp_cpu = torch.cumsum(num_tokens_across_dp, dim=0)
78
return DPMetadata(max_tokens_across_dp_cpu, cu_tokens_across_dp_cpu)
79
80
@@ -101,7 +110,8 @@ def get_forward_context() -> ForwardContext:
110
def set_forward_context(attn_metadata: Any,
111
vllm_config: VllmConfig,
112
virtual_engine: int = 0,
104
- num_tokens: Optional[int] = None):
113
+ num_tokens: Optional[int] = None,
114
+ num_tokens_across_dp: Optional[torch.Tensor] = None):
115
"""A context manager that stores the current forward context,
116
can be attention metadata, etc.
117
Here we can inject common logic for every model forward pass.
@@ -114,7 +124,8 @@ def set_forward_context(attn_metadata: Any,
124
if vllm_config.parallel_config.data_parallel_size > 1 and (
125
attn_metadata is not None or num_tokens is not None):
126
dp_metadata = DPMetadata.make(vllm_config.parallel_config,
117
- attn_metadata, num_tokens or 0)
127
+ attn_metadata, num_tokens or 0,
128
+ num_tokens_across_dp)
129
130
global _forward_context
131
prev_context = _forward_context
vllm/v1/worker/gpu_model_runner.py
+28
-10
@@ -1111,17 +1111,30 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1111
for k, v in self.intermediate_tensors.items()
1112
})
1113
1114
- def get_dp_padding(self, num_tokens: int):
1114
+ def get_dp_padding(self,
1115
+ num_tokens: int) -> tuple[int, Optional[torch.Tensor]]:
1116
dp_size = self.vllm_config.parallel_config.data_parallel_size
1117
dp_rank = self.vllm_config.parallel_config.data_parallel_rank
1117
- if dp_size == 1:
1118
+
1119
+ # For DP: Don't pad when setting enforce_eager.
1120
+ # This lets us set enforce_eager on the prefiller in a P/D setup and
1121
+ # still use CUDA graphs (enabled by this padding) on the decoder.
1122
+ #
1123
+ # TODO(tms) : There are many cases where padding is enabled for
1124
+ # prefills, causing unnecessary and excessive padding of activations.
1125
+
1126
+ if dp_size == 1 or self.vllm_config.model_config.enforce_eager:
1127
# Early exit.
1119
- return 0
1128
+ return 0, None
1129
1130
num_tokens_across_dp = DPMetadata.num_tokens_across_dp(
1131
num_tokens, dp_size, dp_rank)
1132
max_tokens_across_dp_cpu = torch.max(num_tokens_across_dp).item()
1124
- return max_tokens_across_dp_cpu - num_tokens
1133
+ num_tokens_after_padding = torch.tensor([max_tokens_across_dp_cpu] *
1134
+ dp_size,
1135
+ device="cpu",
1136
+ dtype=torch.int32)
1137
+ return max_tokens_across_dp_cpu - num_tokens, num_tokens_after_padding
1138
1139
@torch.inference_mode()
1140
def execute_model(
@@ -1161,7 +1174,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1174
num_input_tokens = num_scheduled_tokens
1175
1176
# Padding for DP
1164
- num_input_tokens += self.get_dp_padding(num_input_tokens)
1177
+ num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens)
1178
+ num_input_tokens += num_pad
1179
1180
# _prepare_inputs may reorder the batch, so we must gather multi
1181
# modal outputs after that to ensure the correct order
@@ -1208,7 +1222,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1222
# Use persistent buffers for CUDA graphs.
1223
with set_forward_context(attn_metadata,
1224
self.vllm_config,
1211
- num_tokens=num_input_tokens):
1225
+ num_tokens=num_input_tokens,
1226
+ num_tokens_across_dp=num_tokens_across_dp):
1227
self.maybe_setup_kv_connector(scheduler_output)
1228
1229
model_output = self.model(
@@ -1681,7 +1696,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1696
) -> torch.Tensor:
1697
1698
# Padding for DP
1684
- num_tokens += self.get_dp_padding(num_tokens)
1699
+ num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
1700
+ num_tokens += num_pad
1701
1702
# Set num_scheduled_tokens based on num_tokens and max_num_seqs
1703
# for dummy run with LoRA so that the num_reqs collectively
@@ -1747,9 +1763,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1763
intermediate_tensors = self.sync_and_slice_intermediate_tensors(
1764
num_tokens, None, False)
1765
1750
- with set_forward_context(attn_metadata,
1751
- self.vllm_config,
1752
- num_tokens=num_tokens):
1766
+ with set_forward_context(
1767
+ attn_metadata,
1768
+ self.vllm_config,
1769
+ num_tokens=num_tokens,
1770
+ num_tokens_across_dp=num_tokens_across_dp):
1771
outputs = model(
1772
input_ids=input_ids,
1773
positions=positions,