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 11, 21:15 UTC
[Perf] Enable cuda graph for deepepHT, 5.3% throughput improvement, 4.4% TTFT improvement (#29558)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Dec 6, 2025 at 23:44 UTC
17eb25e3271f8f22a0d8920a8115158495827cba
4 files changed
+142
-78
tests/compile/test_config.py
+65
-1
@@ -10,7 +10,7 @@ from pydantic import ValidationError
10
11
from vllm.compilation.counter import compilation_counter
12
from vllm.compilation.fix_functionalization import FixFunctionalizationPass
13
-from vllm.config import CompilationConfig, CUDAGraphMode, VllmConfig
13
+from vllm.config import CompilationConfig, CUDAGraphMode, ParallelConfig, VllmConfig
14
from vllm.config.compilation import CompilationMode, PassConfig
15
from vllm.engine.arg_utils import EngineArgs
16
from vllm.logger import _print_warning_once
@@ -235,6 +235,70 @@ def test_splitting_ops_dynamic():
235
assert config.compilation_config.cudagraph_mode == CUDAGraphMode.PIECEWISE
236
237
238
+def test_moe_splitting_ops_deepep_ht_piecewise():
239
+ # Non-inductor, non-attn-fusion case: DeepEP HT with dp>1
240
+ # should add MoE ops to splitting_ops on top of attention ops.
241
+ config = VllmConfig(
242
+ parallel_config=ParallelConfig(
243
+ all2all_backend="deepep_high_throughput",
244
+ data_parallel_size=8,
245
+ ),
246
+ compilation_config=CompilationConfig(
247
+ mode=CompilationMode.VLLM_COMPILE,
248
+ ),
249
+ )
250
+ splitting_ops = config.compilation_config.splitting_ops
251
+ assert splitting_ops is not None
252
+ assert "vllm::moe_forward" in splitting_ops
253
+ assert "vllm::moe_forward_shared" in splitting_ops
254
+
255
+
256
+def test_moe_splitting_ops_deepep_ht_inductor_partition():
257
+ # Inductor partition case: user-provided splitting_ops should be
258
+ # preserved and MoE ops should be appended for DeepEP HT with dp>1.
259
+ config = VllmConfig(
260
+ parallel_config=ParallelConfig(
261
+ all2all_backend="deepep_high_throughput",
262
+ data_parallel_size=8,
263
+ ),
264
+ compilation_config=CompilationConfig(
265
+ mode=CompilationMode.VLLM_COMPILE,
266
+ use_inductor_graph_partition=True,
267
+ splitting_ops=[
268
+ "vllm::unified_attention",
269
+ "vllm::moe_forward",
270
+ "vllm::moe_forward_shared",
271
+ ],
272
+ ),
273
+ )
274
+ splitting_ops = config.compilation_config.splitting_ops
275
+ assert splitting_ops == [
276
+ "vllm::unified_attention",
277
+ "vllm::moe_forward",
278
+ "vllm::moe_forward_shared",
279
+ ]
280
+
281
+
282
+def test_moe_splitting_ops_deepep_ht_attn_fusion_no_inductor():
283
+ # Pure attn-fusion case without inductor partition: even with
284
+ # DeepEP HT and dp>1, we should not re-enable piecewise compilation
285
+ # or add MoE ops into splitting_ops.
286
+ config = VllmConfig(
287
+ parallel_config=ParallelConfig(
288
+ all2all_backend="deepep_high_throughput",
289
+ data_parallel_size=8,
290
+ ),
291
+ compilation_config=CompilationConfig(
292
+ mode=CompilationMode.VLLM_COMPILE,
293
+ pass_config={"enable_attn_fusion": True, "enable_noop": True},
294
+ custom_ops=["+quant_fp8"],
295
+ cudagraph_mode=CUDAGraphMode.PIECEWISE,
296
+ ),
297
+ )
298
+ assert config.compilation_config.splitting_ops == []
299
+ assert config.compilation_config.cudagraph_mode == CUDAGraphMode.FULL
300
+
301
+
302
def test_should_split():
303
import torch
304
vllm/config/compilation.py
+73
-38
@@ -966,7 +966,9 @@ class CompilationConfig:
966
# May get recomputed in the model runner if adjustment is needed for spec-decode
967
self.compute_bs_to_padded_graph_size()
968
969
- def set_splitting_ops_for_v1(self):
969
+ def set_splitting_ops_for_v1(
970
+ self, all2all_backend: str | None = None, data_parallel_size: int | None = None
971
+ ):
972
# To compatible with OOT hardware plugin platform (for example vllm-ascend)
973
# which currently only supports sequence parallelism in eager mode.
974
if self.mode != CompilationMode.VLLM_COMPILE:
@@ -981,50 +983,83 @@ class CompilationConfig:
983
"mode is CompilationMode.VLLM_COMPILE"
984
)
985
984
- if self.use_inductor_graph_partition:
985
- self.set_splitting_ops_for_inductor_graph_partition()
986
- return
986
+ added_default_splitting_ops = False
987
988
- if self.pass_config.fuse_attn_quant:
989
- # here use_inductor_graph_partition is False
988
+ if self.pass_config.fuse_attn_quant and not self.use_inductor_graph_partition:
989
self.set_splitting_ops_for_attn_fusion()
991
- return
992
-
993
- if self.splitting_ops is None:
994
- # NOTE: When using full cudagraph, instead of setting an empty
995
- # list and capture the full cudagraph inside the flattened fx
996
- # graph, we keep the piecewise fx graph structure but capture
997
- # the full cudagraph outside the fx graph. This reduces some
998
- # cpu overhead when the runtime batch_size is not cudagraph
999
- # captured. see https://github.com/vllm-project/vllm/pull/20059
1000
- # for details. Make a copy to avoid mutating the class-level
1001
- # list via reference.
1002
- self.splitting_ops = list(self._attention_ops)
1003
- elif len(self.splitting_ops) == 0:
1004
- logger.warning_once("Using piecewise compilation with empty splitting_ops")
1005
- if self.cudagraph_mode == CUDAGraphMode.PIECEWISE:
990
+ else:
991
+ if self.splitting_ops is None:
992
+ # NOTE: When using full cudagraph, instead of setting an empty
993
+ # list and capture the full cudagraph inside the flattened fx
994
+ # graph, we keep the piecewise fx graph structure but capture
995
+ # the full cudagraph outside the fx graph. This reduces some
996
+ # cpu overhead when the runtime batch_size is not cudagraph
997
+ # captured. see https://github.com/vllm-project/vllm/pull/20059
998
+ # for details. Make a copy to avoid mutating the class-level
999
+ # list via reference.
1000
+ self.splitting_ops = list(self._attention_ops)
1001
+ added_default_splitting_ops = True
1002
+ elif len(self.splitting_ops) == 0:
1003
logger.warning_once(
1007
- "Piecewise compilation with empty splitting_ops do not"
1008
- "contains piecewise cudagraph. Setting cudagraph_"
1009
- "mode to NONE. Hint: If you are using attention backends "
1010
- "that support cudagraph, consider manually setting "
1011
- "cudagraph_mode to FULL or FULL_DECODE_ONLY to enable "
1012
- "full cudagraphs."
1004
+ "Using piecewise compilation with empty splitting_ops"
1005
)
1006
+ if self.cudagraph_mode == CUDAGraphMode.PIECEWISE:
1007
+ logger.warning_once(
1008
+ "Piecewise compilation with empty splitting_ops do not"
1009
+ "contains piecewise cudagraph. Setting cudagraph_"
1010
+ "mode to NONE. Hint: If you are using attention "
1011
+ "backends that support cudagraph, consider manually "
1012
+ "setting cudagraph_mode to FULL or FULL_DECODE_ONLY "
1013
+ "to enable full cudagraphs."
1014
+ )
1015
+ self.cudagraph_mode = CUDAGraphMode.NONE
1016
+ elif self.cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE:
1017
+ logger.warning_once(
1018
+ "Piecewise compilation with empty splitting_ops do "
1019
+ "not contains piecewise cudagraph. Setting "
1020
+ "cudagraph_mode to FULL."
1021
+ )
1022
+ self.cudagraph_mode = CUDAGraphMode.FULL
1023
+ self.splitting_ops = []
1024
+
1025
+ # split MoE ops for cudagraph
1026
+ moe_ops = [
1027
+ "vllm::moe_forward",
1028
+ "vllm::moe_forward_shared",
1029
+ ]
1030
+ backend = all2all_backend or envs.VLLM_ALL2ALL_BACKEND
1031
+ dp_size = data_parallel_size if data_parallel_size is not None else 1
1032
+ need_moe_splitting = (
1033
+ backend == "deepep_high_throughput"
1034
+ and dp_size > 1
1035
+ # pure attn-fusion without inductor partition deliberately disables
1036
+ # piecewise graphs and MoE splitting.
1037
+ and not (
1038
+ self.pass_config.fuse_attn_quant
1039
+ and not self.use_inductor_graph_partition
1040
+ )
1041
+ )
1042
+
1043
+ if need_moe_splitting and self.cudagraph_mode != CUDAGraphMode.NONE:
1044
+ # if we just initialized default splitting_ops for this config,
1045
+ # automatically append the MoE ops
1046
+ if added_default_splitting_ops:
1047
+ for op in moe_ops:
1048
+ if op not in self.splitting_ops:
1049
+ self.splitting_ops.append(op)
1050
+
1051
+ # make sure MoE ops are split out
1052
+ if not any(op in self.splitting_ops for op in moe_ops):
1053
self.cudagraph_mode = CUDAGraphMode.NONE
1015
- elif self.cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE:
1054
logger.warning_once(
1017
- "Piecewise compilation with empty splitting_ops do not "
1018
- "contains piecewise cudagraph. Setting cudagraph_mode "
1019
- "to FULL."
1055
+ "DeepEP high throughput backend with data_parallel_size > 1 "
1056
+ "requires splitting MoE ops from cudagraphs. Please ensure "
1057
+ "'vllm::moe_forward' or 'vllm::moe_forward_shared' are "
1058
+ "present in CompilationConfig.splitting_ops."
1059
)
1021
- self.cudagraph_mode = CUDAGraphMode.FULL
1022
- self.splitting_ops = []
1023
-
1024
- def set_splitting_ops_for_inductor_graph_partition(self):
1025
- assert self.use_inductor_graph_partition
1026
- if self.splitting_ops is None:
1027
- self.splitting_ops = list(self._attention_ops)
1060
+ elif self.cudagraph_mode.has_full_cudagraphs():
1061
+ # fall back to piecewise when MoE splitting is required.
1062
+ self.cudagraph_mode = CUDAGraphMode.PIECEWISE
1063
1064
def set_splitting_ops_for_attn_fusion(self):
1065
assert self.pass_config.fuse_attn_quant
vllm/config/vllm.py
+4
-1
@@ -813,7 +813,10 @@ class VllmConfig:
813
), "MTP with cp_kv_cache_interleave_size > 1 is not supported now."
814
815
# Do this after all the updates to compilation_config.mode
816
- self.compilation_config.set_splitting_ops_for_v1()
816
+ self.compilation_config.set_splitting_ops_for_v1(
817
+ all2all_backend=self.parallel_config.all2all_backend,
818
+ data_parallel_size=self.parallel_config.data_parallel_size,
819
+ )
820
821
if self.compilation_config.pass_config.enable_sp:
822
# With pipeline parallelism or dynamo partitioning,
vllm/platforms/cuda.py
-38
@@ -232,44 +232,6 @@ class CudaPlatformBase(Platform):
232
logger.info(
233
"Forcing kv cache block size to 64 for FlashMLASparse backend."
234
)
235
- # lazy import to avoid circular import
236
- from vllm.config import CUDAGraphMode
237
-
238
- compilation_config = vllm_config.compilation_config
239
- if compilation_config.cudagraph_mode.has_full_cudagraphs():
240
- # decode context parallel does not support full cudagraphs
241
- if parallel_config.decode_context_parallel_size > 1:
242
- logger.warning_once(
243
- "Decode context parallel (DCP) is enabled, which is "
244
- "incompatible with full CUDA graphs. "
245
- "Overriding cudagraph_mode to PIECEWISE."
246
- )
247
- compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
248
- # prefill context parallel do not support full cudagraphs
249
- elif parallel_config.prefill_context_parallel_size > 1:
250
- logger.warning_once(
251
- "Prefill context parallel (PCP) is enabled, which is "
252
- "incompatible with full CUDA graphs. "
253
- "Overriding cudagraph_mode to PIECEWISE."
254
- )
255
- compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
256
- if (
257
- parallel_config.all2all_backend == "deepep_high_throughput"
258
- and parallel_config.data_parallel_size > 1
259
- and compilation_config.cudagraph_mode != CUDAGraphMode.NONE
260
- ):
261
- # TODO: Piecewise Cuda graph might be enabled
262
- # if torch compile cache key issue fixed
263
- # See https://github.com/vllm-project/vllm/pull/25093
264
- logger.info(
265
- "WideEP: Disabling CUDA Graphs since DeepEP high-throughput "
266
- "kernels are optimized for prefill and are incompatible with "
267
- "CUDA Graphs. "
268
- "In order to use CUDA Graphs for decode-optimized workloads, "
269
- "use --all2all-backend with another option, such as "
270
- "deepep_low_latency, pplx, or allgather_reducescatter."
271
- )
272
- compilation_config.cudagraph_mode = CUDAGraphMode.NONE
235
236
@classmethod
237
def get_current_memory_usage(