@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 56539cddac

Fork 0

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 12, 12:15 UTC

[Core] Refactor padding logic and pad for CUDA graphs before attention metadata building (#28579)

Lucas Wilkinson committed Nov 26, 2025 at 14:07 UTC 56539cddac9eeab0a91941d8de689a6cae5dbe05

10 files changed +401 -283

docs/design/cuda_graphs.md

+5 -3

@@ -84,12 +84,14 @@ See the following figures for a quick comparison between the previous and curren

84 ```python

85 class BatchDescriptor(NamedTuple):

86 num_tokens: int

87 - uniform_decode: bool = False

87 + num_reqs: int

88 + uniform: bool = False

89 + has_lora: bool = False

90 ```

90 -where `num_tokens` can be the padded token length, and `uniform_decode` is determined by if `max_query_len` of a batch is equal to the desired `max_query_len` of a uniform_decode, and the num_scheduled_tokens is divisible by that desired `max_query_len`.

92 +where `num_tokens` can be the padded token length, and `uniform` indicates if all the requests have the same query lengths. Many attention backends only support full cudagraphs when the batches are uniform; pure decode batches are uniform but may not be query length 1 (i.e. `num_tokens == num_reqs`), this occurs in the validation pass of spec-decode where "decode" batches will have a query length of `1+num_spec_tokens`.

92 -The goal of this structure is to uniquely identify a (padded) batch with minimal possible items corresponding to a CUDA Graphs item. We are safe to exclude items like `uniform_query_len` because it is a constant at runtime for a certain setup currently. For example, it should be either `1` for a commonly pure decode or `1+num_spec_tokens` for a validation phase of speculative decode.

94 +The goal of this structure is to uniquely identify a (padded) batch with minimal possible items corresponding to a CUDA Graphs item.

96 !!! note

97 The prototype of `BatchDescriptor` may be extended for more general situations in the future, e.g., include more items, like `uniform_query_len` to support multiple different uniform decode lengths settings (<https://github.com/vllm-project/vllm/pull/23679>), or other modifications needed to support CUDA Graphs for models whose inputs are not necessarily token length aware (for example, some multi-modal inputs).

tests/v1/cudagraph/test_cudagraph_dispatch.py

+31 -12

@@ -42,12 +42,24 @@ def _create_vllm_config(

42 mock_config.compilation_config = compilation_config

43 mock_config.scheduler_config = SchedulerConfig(max_num_seqs=max_num_seqs)

44 mock_config.parallel_config = ParallelConfig()

45 + mock_config.speculative_config = None # No speculative decoding

46 if not lora_config:

47 mock_config.lora_config = None

48 # Mimic the behavior of VllmConfig.__post_init__()

49 if compilation_config.mode == CompilationMode.VLLM_COMPILE:

50 compilation_config.set_splitting_ops_for_v1()

52 + # mimic VllmConfig.__post_init__

53 + if compilation_config.cudagraph_capture_sizes:

54 + compilation_config.max_cudagraph_capture_size = (

55 + compilation_config.cudagraph_capture_sizes[-1]

56 + )

57 +

58 + compilation_config.post_init_cudagraph_sizes()

59 + mock_config.pad_for_cudagraph = (

60 + lambda batch_size: compilation_config.bs_to_padded_graph_size[batch_size]

61 + )

62 +

63 return mock_config

@@ -109,9 +121,11 @@ class TestCudagraphDispatcher:

121 # 1. non-uniform batch, size in cudagraph size list

122 desc_full_exact = BatchDescriptor(

123 num_tokens=8,

112 - uniform_decode=False,

124 + uniform=False,

125 + )

126 + rt_mode, key = dispatcher.dispatch(

127 + num_tokens=8, uniform_decode=False, has_lora=False

128 )

114 - rt_mode, key = dispatcher.dispatch(desc_full_exact)

129 if cudagraph_mode_str == "FULL":

130 assert rt_mode == CUDAGraphMode.FULL

131 assert key == desc_full_exact

@@ -122,32 +136,37 @@ class TestCudagraphDispatcher:

136 assert rt_mode == CUDAGraphMode.NONE

137

138 # 2. uniform decode batch, size in cudagraph size list

125 - desc_uniform_exact = BatchDescriptor(num_tokens=8, uniform_decode=True)

126 - rt_mode, key = dispatcher.dispatch(desc_uniform_exact)

139 + desc_uniform_exact = BatchDescriptor(num_tokens=8, num_reqs=8, uniform=True)

140 + rt_mode, key = dispatcher.dispatch(

141 + num_tokens=8, uniform_decode=True, has_lora=False

142 + )

143 if cudagraph_mode_str == "FULL":

144 assert rt_mode == CUDAGraphMode.FULL

129 - assert key == desc_uniform_exact.non_uniform

145 + assert key == desc_uniform_exact.relax_for_mixed_batch_cudagraphs()

146 elif cudagraph_mode_str in ["FULL_DECODE_ONLY", "FULL_AND_PIECEWISE"]:

147 assert rt_mode == CUDAGraphMode.FULL

148 assert key == desc_uniform_exact

149 elif cudagraph_mode_str == "PIECEWISE":

150 assert rt_mode == CUDAGraphMode.PIECEWISE

135 - assert key == desc_uniform_exact.non_uniform

151 + assert key == desc_uniform_exact.relax_for_mixed_batch_cudagraphs()

152 else:

153 assert rt_mode == CUDAGraphMode.NONE

154

155 # 3. No key match

140 - desc_no_match = BatchDescriptor(num_tokens=15, uniform_decode=False)

141 - rt_mode, key = dispatcher.dispatch(desc_no_match)

156 + rt_mode, key = dispatcher.dispatch(

157 + num_tokens=15, uniform_decode=False, has_lora=False

158 + )

159 assert rt_mode == CUDAGraphMode.NONE

143 - assert key is None

160 + assert key == BatchDescriptor(num_tokens=15)

161

162 # 4. Cascade attention should have a fall back mode

146 - desc_full_exact = BatchDescriptor(num_tokens=8, uniform_decode=False)

147 - rt_mode, key = dispatcher.dispatch(desc_full_exact, use_cascade_attn=True)

163 + desc_full_exact = BatchDescriptor(num_tokens=8, uniform=False)

164 + rt_mode, key = dispatcher.dispatch(

165 + num_tokens=8, uniform_decode=False, has_lora=False, use_cascade_attn=True

166 + )

167 if "PIECEWISE" in cudagraph_mode_str: # string contains check

168 assert rt_mode == CUDAGraphMode.PIECEWISE

150 - assert key == desc_full_exact.non_uniform

169 + assert key == desc_full_exact.relax_for_mixed_batch_cudagraphs()

170 else:

171 assert rt_mode == CUDAGraphMode.NONE

172

vllm/forward_context.py

+11 -7

@@ -35,23 +35,27 @@ class BatchDescriptor(NamedTuple):

35 """

37 num_tokens: int

38 - uniform_decode: bool = False

38 + num_reqs: int | None = None

39 """

40 - False can also be used for an uniform decode batch to dispatch to the

41 - cudagraph supporting non-uniform batches.

40 + Number of requests in the batch. Can be None for PIECEWISE cudagraphs where

41 + the cudagraphs can handle any number of requests.

42 + """

43 + uniform: bool = False

44 + """

45 + True if all the requests in the batch have the same number of tokens.

46 """

47 has_lora: bool = False

48 """

49 Whether this batch has active LoRA adapters.

50 """

48 - @property

49 - def non_uniform(self) -> "BatchDescriptor":

52 + def relax_for_mixed_batch_cudagraphs(self) -> "BatchDescriptor":

53 """

51 - Return a non-uniform version of current batch descriptor.

54 + Return a relaxed version of current batch descriptor that is still compatible

55 + with PIECEWISE cudagraphs (or mixed prefill-decode FA cudagraphs).

56 """

57 return BatchDescriptor(

54 - self.num_tokens, uniform_decode=False, has_lora=self.has_lora

58 + self.num_tokens, num_reqs=None, uniform=False, has_lora=self.has_lora

59 )

vllm/v1/attention/backends/flashinfer.py

+1 -20

@@ -930,31 +930,12 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

930

931 if num_decodes > 0:

932 pure_decode = num_prefills == 0

933 - # possible required padding for cudagraph replay

933 use_cudagraph = (

934 self.enable_cuda_graph

935 and pure_decode

936 and num_decode_tokens <= self._decode_cudagraph_max_bs

937 )

939 - if use_cudagraph:

940 - num_input_tokens = self.vllm_config.pad_for_cudagraph(

941 - num_decode_tokens

942 - )

943 - # Carefully fulfill the padding region with reasonable value

944 - # on cpu.

945 - # Make sure paged_kv_indptr_cpu is not decreasing

946 - self.paged_kv_indptr_cpu[

947 - 1 + num_decodes : 1 + num_input_tokens

948 - ].fill_(paged_kv_indptr_cpu[-1])

949 - # Fill the remaining paged_kv_last_page_len_cpu with 1.

950 - # This is because flashinfer treats 0 as a full page

951 - # instead of empty.

952 - self.paged_kv_last_page_len_cpu[num_decodes:num_input_tokens].fill_(

953 - 1

954 - )

955 -

956 - else:

957 - num_input_tokens = num_decode_tokens

938 + num_input_tokens = num_decode_tokens

939

940 attn_metadata.decode_wrapper = self._get_decode_wrapper(

941 num_input_tokens, use_cudagraph

vllm/v1/attention/backends/mamba_attn.py

@@ -107,6 +107,8 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

107 )

108 # -1 in case it's non-computed and causes later issues with indexing

109 block_idx_last_computed_token = block_idx_last_computed_token.clamp(min=0)

110 + # -1 in the case we have a padded request (0 seq-len)

111 + block_idx_last_scheduled_token = block_idx_last_scheduled_token.clamp(min=0)

112

113 return (

114 block_idx_last_computed_token,

vllm/v1/attention/backends/utils.py

+4 -1

@@ -72,6 +72,7 @@ class CommonAttentionMetadata:

73 num_reqs: int

74 """Number of requests"""

75 + # TODO(lucas): rename to num_tokens since it may be padded and this is misleading

76 num_actual_tokens: int

77 """Total number of tokens in batch"""

78 max_query_len: int

@@ -857,7 +858,9 @@ def split_decodes_and_prefills(

858 if require_uniform:

859 is_prefill = query_lens != query_lens[0]

860 else:

860 - is_prefill = query_lens > decode_threshold

861 + # 0-query len indicates a padded request; leave this at the back

862 + # of the batch with the prefills

863 + is_prefill = (query_lens > decode_threshold) | (query_lens == 0)

864

865 if not torch.any(is_prefill):

866 return num_reqs, 0, num_tokens, 0

vllm/v1/cudagraph_dispatcher.py

+66 -31

@@ -4,6 +4,9 @@ from itertools import product

5 from vllm.config import CUDAGraphMode, VllmConfig

6 from vllm.forward_context import BatchDescriptor

7 +from vllm.logger import init_logger

8 +

9 +logger = init_logger(__name__)

12 class CudagraphDispatcher:

@@ -28,7 +31,11 @@ class CudagraphDispatcher:

31 def __init__(self, vllm_config: VllmConfig):

32 self.vllm_config = vllm_config

33 self.compilation_config = vllm_config.compilation_config

31 - self.cudagraph_mode = self.compilation_config.cudagraph_mode

34 + self.uniform_decode_query_len = (

35 + 1

36 + if not self.vllm_config.speculative_config

37 + else 1 + self.vllm_config.speculative_config.num_speculative_tokens

38 + )

40 # Dict to store valid cudagraph dispatching keys.

41 self.cudagraph_keys: dict[CUDAGraphMode, set[BatchDescriptor]] = {

@@ -36,25 +43,42 @@ class CudagraphDispatcher:

43 CUDAGraphMode.FULL: set(),

44 }

39 - not_use_piecewise_compilation = (

40 - not self.cudagraph_mode.requires_piecewise_compilation()

41 - )

42 -

46 assert (

44 - not_use_piecewise_compilation

47 + not self.compilation_config.cudagraph_mode.requires_piecewise_compilation()

48 or self.compilation_config.is_attention_compiled_piecewise()

49 ), (

50 "Compilation mode should be CompilationMode.VLLM_COMPILE when "

51 "cudagraph_mode piecewise cudagraphs is used, "

52 "and attention should be in splitting_ops or "

53 "inductor splitting should be used. "

51 - f"cudagraph_mode={self.cudagraph_mode}, "

54 + f"cudagraph_mode={self.compilation_config.cudagraph_mode}, "

55 f"compilation_mode={self.compilation_config.mode}, "

56 f"splitting_ops={self.compilation_config.splitting_ops}"

57 )

59 self.keys_initialized = False

61 + def _create_padded_batch_descriptor(

62 + self, num_tokens: int, uniform_decode: bool, has_lora: bool

63 + ) -> BatchDescriptor:

64 + max_num_seqs = self.vllm_config.scheduler_config.max_num_seqs

65 + uniform_decode_query_len = self.uniform_decode_query_len

66 + num_tokens_padded = self.vllm_config.pad_for_cudagraph(num_tokens)

67 +

68 + if uniform_decode and self.cudagraph_mode.has_mode(CUDAGraphMode.FULL):

69 + num_reqs = num_tokens_padded // uniform_decode_query_len

70 + assert num_tokens_padded % uniform_decode_query_len == 0

71 + else:

72 + uniform_decode = False

73 + num_reqs = min(num_tokens_padded, max_num_seqs)

74 +

75 + return BatchDescriptor(

76 + num_tokens=num_tokens_padded,

77 + num_reqs=num_reqs,

78 + uniform=uniform_decode,

79 + has_lora=has_lora,

80 + )

81 +

82 def add_cudagraph_key(

83 self, runtime_mode: CUDAGraphMode, batch_descriptor: BatchDescriptor

84 ):

@@ -66,7 +90,9 @@ class CudagraphDispatcher:

90 def initialize_cudagraph_keys(

91 self, cudagraph_mode: CUDAGraphMode, uniform_decode_query_len: int

92 ):

69 - # This should be called only after attention backend is initialized.

93 + # This should be called only after attention backend is initialized. So we can

94 + # get the correct cudagraph mode after backend support is resolved.

95 + self.cudagraph_mode = cudagraph_mode

97 # LoRA activation cases to specialize the cuda graphs on

98 if self.vllm_config.lora_config:

@@ -86,9 +112,9 @@ class CudagraphDispatcher:

112 ):

113 self.add_cudagraph_key(

114 cudagraph_mode.mixed_mode(),

89 - BatchDescriptor(

90 - num_tokens=bs, uniform_decode=False, has_lora=has_lora

91 - ),

115 + self._create_padded_batch_descriptor(

116 + bs, False, has_lora

117 + ).relax_for_mixed_batch_cudagraphs(),

118 )

119

120 # if decode cudagraph mode is FULL, and we don't already have mixed

@@ -109,40 +135,49 @@ class CudagraphDispatcher:

135 for bs, has_lora in product(cudagraph_capture_sizes_for_decode, lora_cases):

136 self.add_cudagraph_key(

137 CUDAGraphMode.FULL,

112 - BatchDescriptor(

113 - num_tokens=bs, uniform_decode=True, has_lora=has_lora

114 - ),

138 + self._create_padded_batch_descriptor(bs, True, has_lora),

139 )

140 +

141 self.keys_initialized = True

142

143 def dispatch(

119 - self, batch_descriptor: BatchDescriptor, use_cascade_attn: bool = False

120 - ) -> tuple[CUDAGraphMode, BatchDescriptor | None]:

144 + self,

145 + num_tokens: int,

146 + uniform_decode: bool,

147 + has_lora: bool,

148 + use_cascade_attn: bool = False,

149 + ) -> tuple[CUDAGraphMode, BatchDescriptor]:

150 """

151 Given conditions(e.g.,batch descriptor and if using cascade attention),

152 dispatch to a cudagraph runtime mode and the valid batch descriptor.

153 A new batch descriptor is returned as we might dispatch a uniform batch

154 to a graph that supports a more general batch (uniform to non-uniform).

155 """

127 - # if not initialized, just skip dispatching.

128 - if not self.keys_initialized:

129 - return CUDAGraphMode.NONE, None

156 + if (

157 + not self.keys_initialized

158 + or self.cudagraph_mode == CUDAGraphMode.NONE

159 + or num_tokens > self.compilation_config.max_cudagraph_capture_size

160 + ):

161 + return CUDAGraphMode.NONE, BatchDescriptor(num_tokens)

162 +

163 + batch_desc = self._create_padded_batch_descriptor(

164 + num_tokens, uniform_decode, has_lora

165 + )

166 + relaxed_batch_desc = batch_desc.relax_for_mixed_batch_cudagraphs()

167

131 - non_uniform_key = batch_descriptor.non_uniform

132 - # if a batch use cascade attention, bypass checking full cudagraphs

168 if not use_cascade_attn:

169 # check if key exists for full cudagraph

135 - if batch_descriptor in self.cudagraph_keys[CUDAGraphMode.FULL]:

136 - return CUDAGraphMode.FULL, batch_descriptor

170 + if batch_desc in self.cudagraph_keys[CUDAGraphMode.FULL]:

171 + return CUDAGraphMode.FULL, batch_desc

172

138 - # otherwise, check if non-uniform key exists

139 - if non_uniform_key in self.cudagraph_keys[CUDAGraphMode.FULL]:

140 - return CUDAGraphMode.FULL, non_uniform_key

173 + # otherwise, check if the relaxed key exists

174 + if relaxed_batch_desc in self.cudagraph_keys[CUDAGraphMode.FULL]:

175 + return CUDAGraphMode.FULL, relaxed_batch_desc

176

142 - # also check if non-uniform key exists for more "general"

177 + # also check if the relaxed key exists for more "general"

178 # piecewise cudagraph

144 - if non_uniform_key in self.cudagraph_keys[CUDAGraphMode.PIECEWISE]:

145 - return CUDAGraphMode.PIECEWISE, non_uniform_key

179 + if relaxed_batch_desc in self.cudagraph_keys[CUDAGraphMode.PIECEWISE]:

180 + return CUDAGraphMode.PIECEWISE, relaxed_batch_desc

181

147 - # finally, just return no cudagraphs

148 - return CUDAGraphMode.NONE, None

182 + # finally, just return no cudagraphs and a trivial batch descriptor

183 + return CUDAGraphMode.NONE, BatchDescriptor(num_tokens)

vllm/v1/worker/dp_utils.py

+16 -1

@@ -9,6 +9,7 @@ from vllm.config import ParallelConfig

9 from vllm.distributed.parallel_state import get_dp_group

10 from vllm.logger import init_logger

11 from vllm.v1.worker.ubatch_utils import (

12 + UBatchSlice,

13 UBatchSlices,

14 check_ubatch_thresholds,

15 create_ubatch_slices,

@@ -88,6 +89,17 @@ def _post_process_dp_padding(tensor: torch.Tensor, should_dp_pad: bool) -> torch

89 return num_tokens_across_dp.cpu()

92 +# This just pads the second ubatch slice out to the total number of tokens

93 +# (num_tokens + padding) since we do `create_ubatch_slices` before applying DP padding.

94 +def _pad_out_ubatch_slice(ubatch_slices: UBatchSlices, num_total_tokens: int):

95 + padded_second_ubatch_slice = slice(

96 + ubatch_slices[1].token_slice.start, num_total_tokens

97 + )

98 + ubatch_slices[1] = UBatchSlice(

99 + padded_second_ubatch_slice, padded_second_ubatch_slice

100 + )

101 +

102 +

103 def _synchronize_dp_ranks(

104 num_tokens_unpadded: int,

105 num_tokens_padded: int,

@@ -220,11 +232,14 @@ def coordinate_batch_across_dp(

232 # to the second ubatch in pad_out_ubatch_slice after attention

233 # metadata creation

234 assert num_tokens_after_padding is not None

223 - token_split_point = int(num_tokens_after_padding[0].item()) // 2

235 + num_tokens_padded = int(num_tokens_after_padding[0].item())

236 + token_split_point = int(num_tokens_padded) // 2

237

238 assert num_scheduled_tokens_per_request is not None

239 ubatch_slices = create_ubatch_slices(

240 num_scheduled_tokens_per_request, token_split_point

241 )

242 + ubatch_slices = _pad_out_ubatch_slice(ubatch_slices, num_tokens_padded)

243 + assert sum(s.num_tokens for s in ubatch_slices) == num_tokens_padded

244

245 return (ubatch_slices, num_tokens_after_padding)

vllm/v1/worker/gpu_model_runner.py

+230 -202

@@ -151,7 +151,6 @@ from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper

151 from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin

152 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin

153 from vllm.v1.worker.ubatch_utils import (

154 - UBatchSlice,

154 UBatchSlices,

155 check_ubatch_thresholds,

156 )

@@ -1239,17 +1238,13 @@ class GPUModelRunner(

1238 self,

1239 scheduler_output: "SchedulerOutput",

1240 num_scheduled_tokens: np.ndarray,

1242 - max_num_scheduled_tokens: int,

1241 ) -> tuple[

1242 torch.Tensor,

1243 SpecDecodeMetadata | None,

1246 - UBatchSlices | None,

1247 - torch.Tensor | None,

1244 ]:

1245 """

1246 :return: tuple[

1247 logits_indices, spec_decode_metadata,

1252 - ubatch_slices, num_tokens_across_dp,

1248 ]

1249 """

1250 total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

@@ -1364,28 +1359,6 @@ class GPUModelRunner(

1359 self.query_start_loc.copy_to_gpu()

1360 query_start_loc = self.query_start_loc.gpu[: num_reqs + 1]

1361

1367 - num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens

1368 - num_tokens_padded = self._get_num_input_tokens(num_tokens_unpadded)

1369 - uniform_decode = (

1370 - max_num_scheduled_tokens == self.uniform_decode_query_len

1371 - ) and (total_num_scheduled_tokens == num_reqs * max_num_scheduled_tokens)

1372 -

1373 - # Disable DP padding when running eager to avoid excessive padding when

1374 - # running prefills. This lets us set enforce_eager on the prefiller in

1375 - # a P/D setup and still use CUDA graphs (enabled by this padding) on the

1376 - # decoder.

1377 - allow_dp_padding = self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE

1378 -

1379 - ubatch_slices, num_tokens_across_dp = coordinate_batch_across_dp(

1380 - num_tokens_unpadded=num_tokens_unpadded,

1381 - parallel_config=self.parallel_config,

1382 - allow_microbatching=True,

1383 - allow_dp_padding=allow_dp_padding,

1384 - num_tokens_padded=num_tokens_padded,

1385 - uniform_decode=uniform_decode,

1386 - num_scheduled_tokens_per_request=num_scheduled_tokens,

1387 - )

1388 -

1362 self.seq_lens.np[:num_reqs] = (

1363 self.input_batch.num_computed_tokens_cpu[:num_reqs] + num_scheduled_tokens

1364 )

@@ -1486,15 +1459,15 @@ class GPUModelRunner(

1459 return (

1460 logits_indices,

1461 spec_decode_metadata,

1489 - ubatch_slices,

1490 - num_tokens_across_dp,

1462 )

1463

1464 def _build_attention_metadata(

1465 self,

1495 - total_num_scheduled_tokens: int,

1496 - max_num_scheduled_tokens: int,

1466 + num_tokens: int,

1467 num_reqs: int,

1468 + max_query_len: int,

1469 + num_tokens_padded: int | None = None,

1470 + num_reqs_padded: int | None = None,

1471 ubatch_slices: UBatchSlices | None = None,

1472 logits_indices: torch.Tensor | None = None,

1473 use_spec_decode: bool = False,

@@ -1505,6 +1478,9 @@ class GPUModelRunner(

1478 """

1479 :return: tuple[attn_metadata, spec_decode_common_attn_metadata]

1480 """

1481 + num_tokens_padded = num_tokens_padded or num_tokens

1482 + num_reqs_padded = num_reqs_padded or num_reqs

1483 +

1484 logits_indices_padded = None

1485 num_logits_indices = None

1486 if logits_indices is not None:

@@ -1522,28 +1498,13 @@ class GPUModelRunner(

1498 self.dcp_rank,

1499 self.parallel_config.cp_kv_cache_interleave_size,

1500 )

1525 - self.dcp_local_seq_lens.copy_to_gpu(num_reqs)

1501 + self.dcp_local_seq_lens.cpu[num_reqs:].fill_(0)

1502 + self.dcp_local_seq_lens.copy_to_gpu(num_reqs_padded)

1503

1504 attn_metadata: PerLayerAttnMetadata = {}

1505 if ubatch_slices is not None:

1506 attn_metadata = [dict() for _ in range(len(ubatch_slices))]

1507

1531 - # Used in the below loop

1532 - query_start_loc = self.query_start_loc.gpu[: num_reqs + 1]

1533 - query_start_loc_cpu = self.query_start_loc.cpu[: num_reqs + 1]

1534 - seq_lens = self.seq_lens.gpu[:num_reqs]

1535 - seq_lens_cpu = self.seq_lens.cpu[:num_reqs]

1536 - num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[

1537 - :num_reqs

1538 - ]

1539 -

1540 - dcp_local_seq_lens, dcp_local_seq_lens_cpu = None, None

1541 - if self.dcp_world_size > 1:

1542 - dcp_local_seq_lens = self.dcp_local_seq_lens.gpu[:num_reqs]

1543 - dcp_local_seq_lens_cpu = self.dcp_local_seq_lens.cpu[:num_reqs]

1544 -

1545 - spec_decode_common_attn_metadata = None

1546 -

1508 if for_cudagraph_capture:

1509 # For some attention backends (e.g. FA) with sliding window models we need

1510 # to make sure the backend see a max_seq_len that is larger to the sliding

@@ -1559,6 +1520,22 @@ class GPUModelRunner(

1520 self.num_accepted_tokens.np[num_reqs:].fill(1)

1521 self.num_accepted_tokens.copy_to_gpu()

1522

1523 + # Used in the below loop, uses padded shapes

1524 + query_start_loc = self.query_start_loc.gpu[: num_reqs_padded + 1]

1525 + query_start_loc_cpu = self.query_start_loc.cpu[: num_reqs_padded + 1]

1526 + seq_lens = self.seq_lens.gpu[:num_reqs_padded]

1527 + seq_lens_cpu = self.seq_lens.cpu[:num_reqs_padded]

1528 + num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[

1529 + :num_reqs_padded

1530 + ]

1531 +

1532 + dcp_local_seq_lens, dcp_local_seq_lens_cpu = None, None

1533 + if self.dcp_world_size > 1:

1534 + dcp_local_seq_lens = self.dcp_local_seq_lens.gpu[:num_reqs_padded]

1535 + dcp_local_seq_lens_cpu = self.dcp_local_seq_lens.cpu[:num_reqs_padded]

1536 +

1537 + spec_decode_common_attn_metadata = None

1538 +

1539 # Prepare the attention metadata for each KV cache group and make layers

1540 # in the same group share the same metadata.

1541 for kv_cache_gid, kv_cache_group in enumerate(

@@ -1567,30 +1544,31 @@ class GPUModelRunner(

1544 encoder_seq_lens, encoder_seq_lens_cpu = self._get_encoder_seq_lens(

1545 num_scheduled_tokens or {},

1546 kv_cache_group.kv_cache_spec,

1570 - num_reqs,

1547 + num_reqs_padded,

1548 )

1549

1550 if isinstance(kv_cache_group.kv_cache_spec, EncoderOnlyAttentionSpec):

1551 # Encoder-only layers do not have KV cache, so we need to

1552 # create a dummy block table and slot mapping for them.

1553 blk_table_tensor = torch.zeros(

1577 - (num_reqs, 1),

1554 + (num_tokens_padded, 1),

1555 dtype=torch.int32,

1556 device=self.device,

1557 )

1558 slot_mapping = torch.zeros(

1582 - (total_num_scheduled_tokens,),

1559 + (num_tokens_padded,),

1560 dtype=torch.int64,

1561 device=self.device,

1562 )

1563 else:

1564 blk_table = self.input_batch.block_table[kv_cache_gid]

1588 - blk_table_tensor = blk_table.get_device_tensor(num_reqs)

1589 - slot_mapping = blk_table.slot_mapping.gpu[:total_num_scheduled_tokens]

1565 + blk_table_tensor = blk_table.get_device_tensor(num_reqs_padded)

1566 + slot_mapping = blk_table.slot_mapping.gpu[:num_tokens_padded]

1567

1568 # Fill unused with -1. Needed for reshape_and_cache in full cuda

1592 - # graph mode.

1593 - blk_table.slot_mapping.gpu[total_num_scheduled_tokens:].fill_(-1)

1569 + # graph mode. `blk_table_tensor` -1 to match mamba PAD_SLOT_ID

1570 + slot_mapping[num_tokens:num_tokens_padded].fill_(-1)

1571 + blk_table_tensor[num_reqs:num_reqs_padded].fill_(-1)

1572

1573 common_attn_metadata = CommonAttentionMetadata(

1574 query_start_loc=query_start_loc,

@@ -1598,9 +1576,9 @@ class GPUModelRunner(

1576 seq_lens=seq_lens,

1577 seq_lens_cpu=seq_lens_cpu,

1578 num_computed_tokens_cpu=num_computed_tokens_cpu,

1601 - num_reqs=num_reqs,

1602 - num_actual_tokens=total_num_scheduled_tokens,

1603 - max_query_len=max_num_scheduled_tokens,

1579 + num_actual_tokens=num_tokens_padded,

1580 + num_reqs=num_reqs_padded,

1581 + max_query_len=max_query_len,

1582 max_seq_len=max_seq_len,

1583 block_table_tensor=blk_table_tensor,

1584 slot_mapping=slot_mapping,

@@ -1631,9 +1609,11 @@ class GPUModelRunner(

1609 extra_attn_metadata_args = {}

1610 if use_spec_decode and isinstance(builder, GDNAttentionMetadataBuilder):

1611 extra_attn_metadata_args = dict(

1634 - num_accepted_tokens=self.num_accepted_tokens.gpu[:num_reqs],

1612 + num_accepted_tokens=self.num_accepted_tokens.gpu[

1613 + :num_reqs_padded

1614 + ],

1615 num_decode_draft_tokens_cpu=self.num_decode_draft_tokens.cpu[

1636 - :num_reqs

1616 + :num_reqs_padded

1617 ],

1618 )

1619

@@ -1677,6 +1657,7 @@ class GPUModelRunner(

1657 def _compute_cascade_attn_prefix_lens(

1658 self,

1659 num_scheduled_tokens: np.ndarray,

1660 + num_computed_tokens: np.ndarray,

1661 num_common_prefix_blocks: list[int],

1662 ) -> list[list[int]] | None:

1663 """

@@ -1699,6 +1680,7 @@ class GPUModelRunner(

1680 # 0 if cascade attention should not be used

1681 cascade_attn_prefix_len = self._compute_cascade_attn_prefix_len(

1682 num_scheduled_tokens,

1683 + num_computed_tokens,

1684 num_common_prefix_blocks[kv_cache_gid],

1685 attn_group.kv_cache_spec,

1686 attn_group.get_metadata_builder(),

@@ -1711,6 +1693,7 @@ class GPUModelRunner(

1693 def _compute_cascade_attn_prefix_len(

1694 self,

1695 num_scheduled_tokens: np.ndarray,

1696 + num_computed_tokens: np.ndarray,

1697 num_common_prefix_blocks: int,

1698 kv_cache_spec: KVCacheSpec,

1699 attn_metadata_builder: AttentionMetadataBuilder,

@@ -1777,10 +1760,7 @@ class GPUModelRunner(

1760 # and the second kernel will get an empty input. While this is not

1761 # a fundamental problem, our current implementation does not support

1762 # this case.

1780 - num_reqs = len(num_scheduled_tokens)

1781 - common_prefix_len = min(

1782 - common_prefix_len, self.input_batch.num_computed_tokens_cpu[:num_reqs].min()

1783 - )

1763 + common_prefix_len = min(common_prefix_len, num_computed_tokens.min())

1764 # common_prefix_len should be a multiple of the block size.

1765 common_prefix_len = (

1766 common_prefix_len // kv_cache_spec.block_size * kv_cache_spec.block_size

@@ -2334,19 +2314,6 @@ class GPUModelRunner(

2314 log_stats=self.parallel_config.eplb_config.log_balancedness,

2315 )

2316

2337 - # This is where the second ubatch is adjusted to account for the padding.

2338 - # Should be called after attention metadata creation. This just pads

2339 - # the second ubatch slice out to the total number of tokens

2340 - # (num_tokens + padding)

2341 - @staticmethod

2342 - def pad_out_ubatch_slice(ubatch_slices: UBatchSlices, num_total_tokens: int):

2343 - padded_second_ubatch_slice = slice(

2344 - ubatch_slices[1].token_slice.start, num_total_tokens

2345 - )

2346 - ubatch_slices[1] = UBatchSlice(

2347 - padded_second_ubatch_slice, padded_second_ubatch_slice

2348 - )

2349 -

2317 def _pool(

2318 self,

2319 hidden_states: torch.Tensor,

@@ -2391,18 +2358,7 @@ class GPUModelRunner(

2358 pooler_output=pooler_output,

2359 )

2360

2394 - def _get_num_input_tokens(self, num_scheduled_tokens: int) -> int:

2395 - if (

2396 - self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE

2397 - and hasattr(self, "cudagraph_batch_sizes")

2398 - and self.cudagraph_batch_sizes

2399 - and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]

2400 - ):

2401 - # Use CUDA graphs.

2402 - # Add padding to the batch size.

2403 - return self.vllm_config.pad_for_cudagraph(num_scheduled_tokens)

2404 -

2405 - # Eager mode.

2361 + def _pad_for_sequence_parallelism(self, num_scheduled_tokens: int) -> int:

2362 # Pad tokens to multiple of tensor_parallel_size when

2363 # enabled collective fusion for SP

2364 tp_size = self.vllm_config.parallel_config.tensor_parallel_size

@@ -2738,6 +2694,87 @@ class GPUModelRunner(

2694 **model_kwargs,

2695 )

2696

2697 + def _determine_batch_execution_and_padding(

2698 + self,

2699 + num_tokens: int,

2700 + num_reqs: int,

2701 + num_scheduled_tokens_np: np.ndarray,

2702 + max_num_scheduled_tokens: int,

2703 + use_cascade_attn: bool,

2704 + allow_microbatching: bool = True,

2705 + force_eager: bool = False,

2706 + # For cudagraph capture TODO(lucas): Refactor how we capture cudagraphs (will

2707 + # be improved in model runner v2)

2708 + force_uniform_decode: bool | None = None,

2709 + force_has_lora: bool | None = None,

2710 + ) -> tuple[

2711 + CUDAGraphMode, BatchDescriptor, UBatchSlices | None, torch.Tensor | None

2712 + ]:

2713 + num_tokens_padded = self._pad_for_sequence_parallelism(num_tokens)

2714 + uniform_decode = (

2715 + (

2716 + (max_num_scheduled_tokens == self.uniform_decode_query_len)

2717 + and (num_tokens_padded == max_num_scheduled_tokens * num_reqs)

2718 + )

2719 + if force_uniform_decode is None

2720 + else force_uniform_decode

2721 + )

2722 +

2723 + has_lora = (

2724 + len(self.input_batch.lora_id_to_lora_request) > 0

2725 + if force_has_lora is None

2726 + else force_has_lora

2727 + )

2728 +

2729 + dispatch_cudagraph = (

2730 + lambda num_tokens: self.cudagraph_dispatcher.dispatch(

2731 + num_tokens=num_tokens,

2732 + has_lora=has_lora,

2733 + use_cascade_attn=use_cascade_attn,

2734 + uniform_decode=uniform_decode,

2735 + )

2736 + if not force_eager

2737 + else (CUDAGraphMode.NONE, BatchDescriptor(num_tokens_padded))

2738 + )

2739 +

2740 + cudagraph_mode, batch_descriptor = dispatch_cudagraph(num_tokens_padded)

2741 + num_tokens_padded = batch_descriptor.num_tokens

2742 +

2743 + # Extra coordination when running data-parallel since we need to coordinate

2744 + # across ranks

2745 + ubatch_slices, num_tokens_across_dp = None, None

2746 + if self.vllm_config.parallel_config.data_parallel_size > 1:

2747 + # Disable DP padding when running eager to avoid excessive padding when

2748 + # running prefills. This lets us set cudagraph_mode="NONE" on the prefiller

2749 + # in a P/D setup and still use CUDA graphs (enabled by this padding) on the

2750 + # decoder.

2751 + allow_dp_padding = (

2752 + self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE

2753 + )

2754 +

2755 + ubatch_slices, num_tokens_across_dp = coordinate_batch_across_dp(

2756 + num_tokens_unpadded=num_tokens_padded,

2757 + parallel_config=self.parallel_config,

2758 + allow_microbatching=allow_microbatching,

2759 + allow_dp_padding=allow_dp_padding,

2760 + num_tokens_padded=num_tokens_padded,

2761 + uniform_decode=uniform_decode,

2762 + num_scheduled_tokens_per_request=num_scheduled_tokens_np,

2763 + )

2764 +

2765 + # Extract DP padding if there is any

2766 + if num_tokens_across_dp is not None:

2767 + dp_rank = self.parallel_config.data_parallel_rank

2768 + num_tokens_padded = int(num_tokens_across_dp[dp_rank].item())

2769 +

2770 + # Re-dispatch with DP padding

2771 + cudagraph_mode, batch_descriptor = dispatch_cudagraph(num_tokens_padded)

2772 + # Assert to make sure the agreed upon token count is correct otherwise

2773 + # num_tokens_across_dp will no-longer be valid

2774 + assert batch_descriptor.num_tokens == num_tokens_padded

2775 +

2776 + return cudagraph_mode, batch_descriptor, ubatch_slices, num_tokens_across_dp

2777 +

2778 @torch.inference_mode()

2779 def execute_model(

2780 self,

@@ -2790,7 +2827,7 @@ class GPUModelRunner(

2827 # returns True. before returning early here we call

2828 # dummy run to ensure coordinate_batch_across_dp

2829 # is called into to avoid out of sync issues.

2793 - self._dummy_run(self._get_num_input_tokens(1))

2830 + self._dummy_run(1)

2831 if not has_kv_transfer_group():

2832 # Return empty ModelRunnerOutput if no work to do.

2833 return EMPTY_MODEL_RUNNER_OUTPUT

@@ -2809,36 +2846,63 @@ class GPUModelRunner(

2846 tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]

2847 num_scheduled_tokens_np = np.array(tokens, dtype=np.int32)

2848 max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())

2849 + num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens

2850

2851 (

2852 logits_indices,

2853 spec_decode_metadata,

2816 - ubatch_slices,

2817 - num_tokens_across_dp,

2854 ) = self._prepare_inputs(

2819 - scheduler_output, num_scheduled_tokens_np, max_num_scheduled_tokens

2855 + scheduler_output,

2856 + num_scheduled_tokens_np,

2857 )

2858

2859 cascade_attn_prefix_lens = None

2860 # Disable cascade attention when using microbatching (DBO)

2824 - if self.cascade_attn_enabled and ubatch_slices is None:

2861 + if self.cascade_attn_enabled and not self.parallel_config.enable_dbo:

2862 # Pre-compute cascade attention prefix lengths

2826 - # NOTE: Must be AFTER _prepare_inputs uses self.input_batch state

2863 cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens(

2864 num_scheduled_tokens_np,

2865 + self.input_batch.num_computed_tokens_cpu[:num_reqs],

2866 scheduler_output.num_common_prefix_blocks,

2867 )

2868

2832 - # TODO(lucas): move cudagraph dispatching here:

2833 - # https://github.com/vllm-project/vllm/issues/23789

2869 + (

2870 + cudagraph_mode,

2871 + batch_desc,

2872 + ubatch_slices,

2873 + num_tokens_across_dp,

2874 + ) = self._determine_batch_execution_and_padding(

2875 + num_tokens=num_tokens_unpadded,

2876 + num_reqs=num_reqs,

2877 + num_scheduled_tokens_np=num_scheduled_tokens_np,

2878 + max_num_scheduled_tokens=max_num_scheduled_tokens,

2879 + use_cascade_attn=cascade_attn_prefix_lens is not None,

2880 + )

2881 +

2882 + logger.debug(

2883 + "Running batch with cudagraph_mode: %s, batch_descriptor: %s, "

2884 + "ubatch_slices: %s, num_tokens_across_dp: %s",

2885 + cudagraph_mode,

2886 + batch_desc,

2887 + ubatch_slices,

2888 + num_tokens_across_dp,

2889 + )

2890 +

2891 + num_tokens_padded = batch_desc.num_tokens

2892 + num_reqs_padded = (

2893 + batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs

2894 + )

2895

2835 - total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

2896 use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0

2837 - attn_metadata, spec_decode_common_attn_metadata = (

2897 + pad_attn = cudagraph_mode == CUDAGraphMode.FULL

2898 +

2899 + (attn_metadata, spec_decode_common_attn_metadata) = (

2900 self._build_attention_metadata(

2839 - total_num_scheduled_tokens=total_num_scheduled_tokens,

2840 - max_num_scheduled_tokens=max_num_scheduled_tokens,

2901 + num_tokens=num_tokens_unpadded,

2902 + num_tokens_padded=num_tokens_padded if pad_attn else None,

2903 num_reqs=num_reqs,

2904 + num_reqs_padded=num_reqs_padded if pad_attn else None,

2905 + max_query_len=max_num_scheduled_tokens,

2906 ubatch_slices=ubatch_slices,

2907 logits_indices=logits_indices,

2908 use_spec_decode=use_spec_decode,

@@ -2847,49 +2911,22 @@ class GPUModelRunner(

2911 )

2912 )

2913

2850 - dp_rank = self.parallel_config.data_parallel_rank

2851 - if ubatch_slices:

2852 - assert num_tokens_across_dp is not None

2853 - num_input_tokens = int(num_tokens_across_dp[dp_rank].item())

2854 - self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens)

2855 - elif num_tokens_across_dp is not None:

2856 - num_input_tokens = int(num_tokens_across_dp[dp_rank].item())

2857 - else:

2858 - num_input_tokens = self._get_num_input_tokens(

2859 - scheduler_output.total_num_scheduled_tokens

2860 - )

2861 -

2862 - (

2863 - input_ids,

2864 - inputs_embeds,

2865 - positions,

2866 - intermediate_tensors,

2867 - model_kwargs,

2868 - ec_connector_output,

2869 - ) = self._preprocess(

2870 - scheduler_output, num_input_tokens, intermediate_tensors

2871 - )

2872 -

2873 - uniform_decode = (

2874 - max_num_scheduled_tokens == self.uniform_decode_query_len

2875 - ) and (num_scheduled_tokens == num_reqs * max_num_scheduled_tokens)

2876 - batch_desc = BatchDescriptor(

2877 - num_tokens=num_input_tokens,

2878 - uniform_decode=uniform_decode,

2879 - has_lora=len(self.input_batch.lora_id_to_lora_request) > 0,

2880 - )

2881 - cudagraph_runtime_mode, batch_descriptor = (

2882 - self.cudagraph_dispatcher.dispatch(

2883 - batch_desc,

2884 - use_cascade_attn=cascade_attn_prefix_lens is not None,

2885 - )

2914 + (

2915 + input_ids,

2916 + inputs_embeds,

2917 + positions,

2918 + intermediate_tensors,

2919 + model_kwargs,

2920 + ec_connector_output,

2921 + ) = self._preprocess(

2922 + scheduler_output, num_tokens_padded, intermediate_tensors

2923 )

2924

2925 # Set cudagraph mode to none if calc_kv_scales is true.

2926 # KV scales calculation involves dynamic operations that are incompatible

2927 # with CUDA graph capture.

2928 if self.calculate_kv_scales:

2892 - cudagraph_runtime_mode = CUDAGraphMode.NONE

2929 + cudagraph_mode = CUDAGraphMode.NONE

2930 # Mark KV scales as calculated after the first forward pass

2931 self.calculate_kv_scales = False

2932

@@ -2899,10 +2936,10 @@ class GPUModelRunner(

2936 set_forward_context(

2937 attn_metadata,

2938 self.vllm_config,

2902 - num_tokens=num_input_tokens,

2939 + num_tokens=num_tokens_padded,

2940 num_tokens_across_dp=num_tokens_across_dp,

2904 - cudagraph_runtime_mode=cudagraph_runtime_mode,

2905 - batch_descriptor=batch_descriptor,

2941 + cudagraph_runtime_mode=cudagraph_mode,

2942 + batch_descriptor=batch_desc,

2943 ubatch_slices=ubatch_slices,

2944 ),

2945 record_function_or_nullcontext("gpu_model_runner: forward"),

@@ -2952,7 +2989,7 @@ class GPUModelRunner(

2989 if not get_pp_group().is_last_rank:

2990 all_gather_tensors = {

2991 "residual": not is_residual_scattered_for_sp(

2955 - self.vllm_config, num_input_tokens

2992 + self.vllm_config, num_tokens_padded

2993 )

2994 }

2995 get_pp_group().send_tensor_dict(

@@ -3841,52 +3878,44 @@ class GPUModelRunner(

3878 assert sum(num_scheduled_tokens_list) == num_tokens

3879 assert len(num_scheduled_tokens_list) == num_reqs

3880 num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32)

3844 - total_num_scheduled_tokens = int(num_scheduled_tokens.sum())

3881 + num_tokens_unpadded = int(num_scheduled_tokens.sum())

3882 +

3883 num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)

3884

3847 - # Disable DP padding when running eager

3848 - allow_dp_padding = self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE

3849 -

3850 - # We currently only microbatch if the number of tokens is

3851 - # over a certain threshold.

3852 - ubatch_slices, num_tokens_across_dp = coordinate_batch_across_dp(

3853 - num_tokens_unpadded=total_num_scheduled_tokens,

3854 - parallel_config=self.vllm_config.parallel_config,

3855 - allow_microbatching=allow_microbatching,

3856 - allow_dp_padding=allow_dp_padding,

3857 - num_tokens_padded=total_num_scheduled_tokens,

3858 - uniform_decode=uniform_decode,

3859 - num_scheduled_tokens_per_request=num_scheduled_tokens,

3860 - )

3861 - num_tokens_after_padding = num_tokens

3862 - if num_tokens_across_dp is not None:

3863 - dp_rank = self.parallel_config.data_parallel_rank

3864 - num_tokens_after_padding = int(num_tokens_across_dp[dp_rank])

3865 -

3866 - # filter out the valid batch descriptor

3867 - _cg_mode, batch_descriptor = (

3868 - self.cudagraph_dispatcher.dispatch(

3869 - BatchDescriptor(

3870 - num_tokens=num_tokens_after_padding,

3871 - uniform_decode=uniform_decode,

3872 - has_lora=activate_lora and self.lora_config is not None,

3873 - )

3885 + _cudagraph_mode, batch_desc, ubatch_slices, num_tokens_across_dp = (

3886 + self._determine_batch_execution_and_padding(

3887 + num_tokens=num_tokens_unpadded,

3888 + num_reqs=num_reqs,

3889 + num_scheduled_tokens_np=num_scheduled_tokens,

3890 + max_num_scheduled_tokens=max_query_len,

3891 + use_cascade_attn=False,

3892 + allow_microbatching=allow_microbatching,

3893 + force_eager=is_profile

3894 + or (cudagraph_runtime_mode == CUDAGraphMode.NONE),

3895 + # `force_uniform_decode` is used for cudagraph capture; because for

3896 + # capturing mixed prefill-decode batches, we sometimes use

3897 + # num_tokens == num_reqs which looks like a uniform decode batch to the

3898 + # dispatcher; but we actually want to capture a piecewise cudagraph

3899 + force_uniform_decode=uniform_decode,

3900 + # `force_has_lora` is used for cudagraph capture; because LoRA is

3901 + # activated later in the context manager, but we need to know the

3902 + # LoRA state when determining the batch descriptor for capture

3903 + force_has_lora=activate_lora,

3904 )

3875 - if not is_profile

3876 - else (CUDAGraphMode.NONE, None)

3905 )

3878 - if cudagraph_runtime_mode is not None:

3879 - # we allow forcing NONE when the dispatcher disagrees to support

3880 - # warm ups for cudagraph capture

3881 - assert (

3882 - cudagraph_runtime_mode == CUDAGraphMode.NONE

3883 - or cudagraph_runtime_mode == _cg_mode

3884 - ), (

3885 - f"Cudagraph runtime mode mismatch at dummy_run. "

3886 - f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."

3887 - )

3906 +

3907 + if cudagraph_runtime_mode is None:

3908 + cudagraph_runtime_mode = _cudagraph_mode

3909 else:

3889 - cudagraph_runtime_mode = _cg_mode

3910 + assert cudagraph_runtime_mode == _cudagraph_mode, (

3911 + f"Cudagraph runtime mode mismatch in dummy_run. "

3912 + f"Expected {_cudagraph_mode}, but got {cudagraph_runtime_mode}."

3913 + )

3914 +

3915 + num_tokens_padded = batch_desc.num_tokens

3916 + num_reqs_padded = (

3917 + batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs

3918 + )

3919

3920 attn_metadata: PerLayerAttnMetadata | None = None

3921

@@ -3909,9 +3938,9 @@ class GPUModelRunner(

3938 self.query_start_loc.copy_to_gpu()

3939

3940 attn_metadata, _ = self._build_attention_metadata(

3912 - total_num_scheduled_tokens=num_tokens,

3913 - max_num_scheduled_tokens=max_query_len,

3914 - num_reqs=num_reqs,

3941 + num_tokens=num_tokens_unpadded,

3942 + num_reqs=num_reqs_padded,

3943 + max_query_len=max_query_len,

3944 ubatch_slices=ubatch_slices,

3945 for_cudagraph_capture=True,

3946 )

@@ -3924,29 +3953,29 @@ class GPUModelRunner(

3953 remove_lora,

3954 ):

3955 # Make sure padding doesn't exceed max_num_tokens

3927 - assert num_tokens_after_padding <= self.max_num_tokens

3928 - model_kwargs = self._init_model_kwargs(num_tokens_after_padding)

3956 + assert num_tokens_padded <= self.max_num_tokens

3957 + model_kwargs = self._init_model_kwargs(num_tokens_padded)

3958 if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:

3959 input_ids = None

3931 - inputs_embeds = self.inputs_embeds.gpu[:num_tokens_after_padding]

3960 + inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]

3961 model_kwargs = {

3962 **model_kwargs,

3963 **self._dummy_mm_kwargs(num_reqs),

3964 }

3965 elif self.enable_prompt_embeds:

3966 input_ids = None

3938 - inputs_embeds = self.inputs_embeds.gpu[:num_tokens_after_padding]

3939 - model_kwargs = self._init_model_kwargs(num_tokens_after_padding)

3967 + inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded]

3968 + model_kwargs = self._init_model_kwargs(num_tokens_padded)

3969 else:

3941 - input_ids = self.input_ids.gpu[:num_tokens_after_padding]

3970 + input_ids = self.input_ids.gpu[:num_tokens_padded]

3971 inputs_embeds = None

3972

3973 if self.uses_mrope:

3945 - positions = self.mrope_positions.gpu[:, :num_tokens_after_padding]

3974 + positions = self.mrope_positions.gpu[:, :num_tokens_padded]

3975 elif self.uses_xdrope_dim > 0:

3947 - positions = self.xdrope_positions.gpu[:, :num_tokens_after_padding]

3976 + positions = self.xdrope_positions.gpu[:, :num_tokens_padded]

3977 else:

3949 - positions = self.positions.gpu[:num_tokens_after_padding]

3978 + positions = self.positions.gpu[:num_tokens_padded]

3979

3980 if get_pp_group().is_first_rank:

3981 intermediate_tensors = None

@@ -3961,26 +3990,26 @@ class GPUModelRunner(

3990 )

3991

3992 intermediate_tensors = self.sync_and_slice_intermediate_tensors(

3964 - num_tokens_after_padding, None, False

3993 + num_tokens_padded, None, False

3994 )

3995

3996 if ubatch_slices is not None:

3997 # Adjust values to reflect a single ubatch.

3998 # TODO(sage,lucas): this is cruft that should be addressed in

3999 # the padding refactor.

3971 - num_tokens_after_padding = ubatch_slices[0].num_tokens

4000 + num_tokens_padded = ubatch_slices[0].num_tokens

4001 if num_tokens_across_dp is not None:

3973 - num_tokens_across_dp[:] = num_tokens_after_padding

4002 + num_tokens_across_dp[:] = num_tokens_padded

4003

4004 with (

4005 self.maybe_randomize_inputs(input_ids),

4006 set_forward_context(

4007 attn_metadata,

4008 self.vllm_config,

3980 - num_tokens=num_tokens_after_padding,

4009 + num_tokens=num_tokens_padded,

4010 num_tokens_across_dp=num_tokens_across_dp,

4011 cudagraph_runtime_mode=cudagraph_runtime_mode,

3983 - batch_descriptor=batch_descriptor,

4012 + batch_descriptor=batch_desc,

4013 ubatch_slices=ubatch_slices,

4014 ),

4015 ):

@@ -4706,8 +4735,7 @@ class GPUModelRunner(

4735

4736 # Trigger cudagraph dispatching keys initialization after

4737 # resolved cudagraph mode.

4709 - cudagraph_mode = self.compilation_config.cudagraph_mode

4710 - assert cudagraph_mode is not None

4738 + self.compilation_config.cudagraph_mode = cudagraph_mode

4739 self.cudagraph_dispatcher.initialize_cudagraph_keys(

4740 cudagraph_mode, self.uniform_decode_query_len

4741 )

vllm/v1/worker/gpu_worker.py

+35 -6

@@ -8,12 +8,13 @@ from contextlib import AbstractContextManager, nullcontext

8 from types import NoneType

9 from typing import TYPE_CHECKING, Any, cast

11 +import numpy as np

12 import torch

13 import torch.distributed

14 import torch.nn as nn

16 import vllm.envs as envs

16 -from vllm.config import VllmConfig

17 +from vllm.config import CUDAGraphMode, VllmConfig

18 from vllm.distributed import (

19 ensure_model_parallel_initialized,

20 init_distributed_environment,

@@ -487,6 +488,7 @@ class Worker(WorkerBase):

488 hidden_states, last_hidden_states = self.model_runner._dummy_run(

489 num_tokens=max_num_reqs,

490 skip_eplb=True,

491 + cudagraph_runtime_mode=CUDAGraphMode.NONE,

492 )

493 if self.model_runner.is_pooling_model:

494 self.model_runner._dummy_pooler_run(hidden_states)

@@ -534,12 +536,39 @@ class Worker(WorkerBase):

536 intermediate_tensors = None

537 forward_pass = scheduler_output.total_num_scheduled_tokens > 0

538 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

537 - num_input_tokens = self.model_runner._get_num_input_tokens(num_scheduled_tokens)

538 - all_gather_tensors = {

539 - "residual": not is_residual_scattered_for_sp(

540 - self.vllm_config, num_input_tokens

539 + all_gather_tensors = {}

540 + compilation_config = self.vllm_config.compilation_config

541 + parallel_config = self.vllm_config.parallel_config

542 +

543 + if (

544 + parallel_config.pipeline_parallel_size > 1

545 + and compilation_config.pass_config.enable_sequence_parallelism

546 + and forward_pass

547 + ):

548 + # currently only supported by V1 GPUModelRunner

549 + assert isinstance(self.model_runner, GPUModelRunner)

550 + num_scheduled_tokens_np = np.array(

551 + list(scheduler_output.num_scheduled_tokens.values()),

552 + dtype=np.int32,

553 )

542 - }

554 + # TODO(lucas): This is pretty gross; ideally we should only ever call

555 + # `_determine_batch_execution_and_padding` once (will get called again

556 + # in `execute_model`) but this requires a larger refactor of PP.

557 + _, batch_desc, _, _ = (

558 + self.model_runner._determine_batch_execution_and_padding(

559 + num_tokens=num_scheduled_tokens,

560 + num_reqs=len(num_scheduled_tokens_np),

561 + num_scheduled_tokens_np=num_scheduled_tokens_np,

562 + max_num_scheduled_tokens=num_scheduled_tokens_np.max(),

563 + use_cascade_attn=False, # TODO(lucas): Handle cascade attention

564 + )

565 + )

566 + all_gather_tensors = {

567 + "residual": not is_residual_scattered_for_sp(

568 + self.vllm_config, batch_desc.num_tokens

569 + )

570 + }

571 +

572 if forward_pass and not get_pp_group().is_first_rank:

573 tensor_dict = get_pp_group().recv_tensor_dict(

574 all_gather_group=get_tp_group(),