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 9, 23:00 UTC
[Model Runner V2] Refactor get_cudagraph_and_dp_padding (#32625)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jan 19, 2026 at 18:25 UTC
7b7cdce968bf9428c98e941aaf75b07d08a9a897
2 files changed
+54
-67
vllm/v1/worker/gpu/dp_utils.py
+40
-5
@@ -6,6 +6,12 @@ import torch.distributed as dist
6
from vllm.distributed.parallel_state import get_dp_group
7
8
9
+def make_num_tokens_across_dp(dp_size: int, num_tokens: int) -> torch.Tensor | None:
10
+ if dp_size == 1:
11
+ return None
12
+ return torch.full((dp_size,), num_tokens, dtype=torch.int32, device="cpu")
13
+
14
+
15
def get_batch_metadata_across_dp(
16
num_tokens: int,
17
cudagraph_size: int,
@@ -22,10 +28,39 @@ def get_batch_metadata_across_dp(
28
return tensor[0], tensor[1]
29
30
25
-def make_num_tokens_across_dp(
26
- dp_size: int,
31
+def get_cudagraph_and_dp_padding(
32
num_tokens: int,
28
-) -> torch.Tensor | None:
33
+ cudagraph_size: int | None,
34
+ dp_size: int,
35
+ dp_rank: int,
36
+) -> tuple[bool, int, torch.Tensor | None]:
37
if dp_size == 1:
30
- return None
31
- return torch.full((dp_size,), num_tokens, dtype=torch.int32, device="cpu")
38
+ if cudagraph_size is not None:
39
+ return True, cudagraph_size, None
40
+ else:
41
+ return False, num_tokens, None
42
+
43
+ if num_tokens == 0:
44
+ cudagraph_size = 0
45
+ elif cudagraph_size is None:
46
+ cudagraph_size = -1
47
+ num_tokens_across_dp, cudagraph_size_across_dp = get_batch_metadata_across_dp(
48
+ num_tokens, cudagraph_size, dp_size, dp_rank
49
+ )
50
+ if torch.all(num_tokens_across_dp == 0).item():
51
+ # All ranks have zero tokens to run.
52
+ return False, 0, None
53
+
54
+ if torch.all(cudagraph_size_across_dp != -1).item():
55
+ # All ranks use CUDA graph or have zero tokens.
56
+ # Use CUDA graph for all ranks.
57
+ # Pad all ranks to the maximum CUDA graph size.
58
+ max_cudagraph_size = int(cudagraph_size_across_dp.max().item())
59
+ num_tokens_across_dp[:] = max_cudagraph_size
60
+ return True, max_cudagraph_size, num_tokens_across_dp
61
+ else:
62
+ # Some ranks do not use CUDA graph. Use eager mode for all ranks.
63
+ # No padding is needed except for ranks that have no tokens to run.
64
+ num_tokens_across_dp = torch.clamp(num_tokens_across_dp, min=1)
65
+ num_tokens_after_padding = int(num_tokens_across_dp[dp_rank].item())
66
+ return False, num_tokens_after_padding, num_tokens_across_dp
vllm/v1/worker/gpu/model_runner.py
+14
-62
@@ -2,7 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import gc
4
import time
5
-from collections.abc import Iterable
5
from copy import deepcopy
6
from typing import Any
7
@@ -37,7 +36,7 @@ from vllm.v1.worker.gpu.block_table import BlockTables
36
from vllm.v1.worker.gpu.buffer_utils import UvaBufferPool
37
from vllm.v1.worker.gpu.cudagraph_utils import CudaGraphManager
38
from vllm.v1.worker.gpu.dp_utils import (
40
- get_batch_metadata_across_dp,
39
+ get_cudagraph_and_dp_padding,
40
make_num_tokens_across_dp,
41
)
42
from vllm.v1.worker.gpu.input_batch import (
@@ -877,60 +876,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
876
)
877
return draft_tokens
878
880
- def get_cudagraph_and_dp_padding(
881
- self,
882
- total_num_scheduled_tokens: int,
883
- num_tokens_per_request: Iterable[int],
884
- ) -> tuple[CUDAGraphMode, int, torch.Tensor | None]:
885
- dp_size = self.parallel_config.data_parallel_size
886
- if dp_size == 1:
887
- # No DP. Only consider CUDA graphs.
888
- if total_num_scheduled_tokens == 0:
889
- # Special case: no tokens to run.
890
- return CUDAGraphMode.NONE, 0, None
891
-
892
- cudagraph_size = self.cudagraph_manager.get_cudagraph_size(
893
- total_num_scheduled_tokens, num_tokens_per_request
894
- )
895
- if cudagraph_size is not None:
896
- # Use full CUDA graph.
897
- return CUDAGraphMode.FULL, cudagraph_size, None
898
- # Fall back to eager mode.
899
- # TODO(woosuk): Support piecewise CUDA graphs.
900
- return CUDAGraphMode.NONE, total_num_scheduled_tokens, None
901
-
902
- # Consider DP padding and CUDA graph.
903
- if total_num_scheduled_tokens == 0:
904
- # Special handling is needed for 0.
905
- cudagraph_size_before_dp: int | None = 0
906
- else:
907
- cudagraph_size_before_dp = self.cudagraph_manager.get_cudagraph_size(
908
- total_num_scheduled_tokens, num_tokens_per_request
909
- )
910
- if cudagraph_size_before_dp is None:
911
- cudagraph_size_before_dp = -1
912
-
913
- assert cudagraph_size_before_dp is not None
914
- dp_rank = self.parallel_config.data_parallel_rank
915
- num_tokens_across_dp, cudagraph_size_across_dp = get_batch_metadata_across_dp(
916
- total_num_scheduled_tokens,
917
- cudagraph_size_before_dp,
918
- dp_size,
919
- dp_rank,
920
- )
921
- if all(cudagraph_size_across_dp >= 0):
922
- # If all ranks can use CUDA graph, pad to the maximum number of tokens
923
- # across DP and use CUDA graph.
924
- num_tokens_after_padding = int(cudagraph_size_across_dp.max().item())
925
- cudagraph_mode = CUDAGraphMode.FULL
926
- else:
927
- # If any of the ranks cannot use CUDA graph, use eager mode for all ranks.
928
- # No padding is needed except for ranks that have no tokens to run.
929
- num_tokens_across_dp = torch.clamp(num_tokens_across_dp, min=1)
930
- num_tokens_after_padding = num_tokens_across_dp[dp_rank]
931
- cudagraph_mode = CUDAGraphMode.NONE
932
- return cudagraph_mode, num_tokens_after_padding, num_tokens_across_dp
933
-
879
@torch.inference_mode()
880
def execute_model(
881
self,
@@ -951,10 +896,17 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
896
# No need to run the model.
897
return EMPTY_MODEL_RUNNER_OUTPUT
898
954
- cudagraph_mode, num_tokens_after_padding, num_tokens_across_dp = (
955
- self.get_cudagraph_and_dp_padding(
899
+ # Get the CUDA graph size. None means no CUDA graph is used.
900
+ cudagraph_size = self.cudagraph_manager.get_cudagraph_size(
901
+ scheduler_output.total_num_scheduled_tokens,
902
+ scheduler_output.num_scheduled_tokens.values(),
903
+ )
904
+ use_cudagraph, num_tokens_after_padding, num_tokens_across_dp = (
905
+ get_cudagraph_and_dp_padding(
906
scheduler_output.total_num_scheduled_tokens,
957
- scheduler_output.num_scheduled_tokens.values(),
907
+ cudagraph_size,
908
+ self.parallel_config.data_parallel_size,
909
+ self.parallel_config.data_parallel_rank,
910
)
911
)
912
if num_tokens_after_padding == 0:
@@ -1006,7 +958,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
958
# FIXME(woosuk): Fix warmup for LoRA.
959
960
# Run model.
1009
- if cudagraph_mode == CUDAGraphMode.FULL:
961
+ if use_cudagraph:
962
# Run CUDA graph.
963
# NOTE(woosuk): Here, we don't need to pass the input tensors,
964
# because they are already copied to the CUDA graph input buffers.
@@ -1015,7 +967,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
967
)
968
else:
969
# Run PyTorch model in eager mode.
1018
- # TODO(woosuk): Support piecewise CUDA graph.
970
positions = input_batch.positions
971
if self.uses_mrope:
972
assert input_batch.mrope_positions is not None
@@ -1024,7 +975,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
975
input_batch.attn_metadata,
976
self.vllm_config,
977
num_tokens=input_batch.num_tokens_after_padding,
1027
- cudagraph_runtime_mode=cudagraph_mode,
978
+ # TODO(woosuk): Support piecewise CUDA graph.
979
+ cudagraph_runtime_mode=CUDAGraphMode.NONE,
980
num_tokens_across_dp=num_tokens_across_dp,
981
):
982
hidden_states = self.model(