| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections.abc import Set as AbstractSet |
| 4 | from dataclasses import replace |
| 5 | from itertools import product |
| 6 | |
| 7 | from vllm.config import CUDAGraphMode, VllmConfig |
| 8 | from vllm.forward_context import BatchDescriptor |
| 9 | from vllm.logger import init_logger |
| 10 | from vllm.lora.utils import get_captured_lora_counts |
| 11 | |
| 12 | logger = init_logger(__name__) |
| 13 | |
| 14 | |
| 15 | class CudagraphDispatcher: |
| 16 | """Runtime cudagraph dispatcher to dispatch keys for multiple set of |
| 17 | cudagraphs. |
| 18 | |
| 19 | The dispatcher stores two sets of dispatch keys, one for PIECEWISE and one |
| 20 | for FULL cudagraph runtime mode. The keys are initialized depending on |
| 21 | attention support and what cudagraph mode is set in CompilationConfig. The |
| 22 | keys stored in dispatcher are the only source of truth for valid |
| 23 | cudagraphs that can be dispatched at runtime. |
| 24 | |
| 25 | At runtime, the dispatch method generates the runtime cudagraph mode (FULL, |
| 26 | PIECEWISE, or NONE for no cudagraph) and the valid key (batch descriptor) |
| 27 | based on the input key. After dispatching (communicated via forward |
| 28 | context), the cudagraph wrappers will trust the dispatch key to either |
| 29 | capture or replay (if the mode matches), or pass through to the underlying |
| 30 | runnable without cudagraph (if the mode does not match or mode is NONE). |
| 31 | """ |
| 32 | |
| 33 | def __init__(self, vllm_config: VllmConfig): |
| 34 | self.vllm_config = vllm_config |
| 35 | self.compilation_config = vllm_config.compilation_config |
| 36 | self.uniform_decode_query_len = 1 + self.vllm_config.num_speculative_tokens |
| 37 | |
| 38 | # Dict to store valid cudagraph dispatching keys. |
| 39 | self.cudagraph_keys: dict[CUDAGraphMode, set[BatchDescriptor]] = { |
| 40 | CUDAGraphMode.PIECEWISE: set(), |
| 41 | CUDAGraphMode.FULL: set(), |
| 42 | } |
| 43 | |
| 44 | from vllm.compilation.breakable_cudagraph import ( |
| 45 | is_breakable_cudagraph_enabled, |
| 46 | ) |
| 47 | |
| 48 | assert ( |
| 49 | not self.compilation_config.cudagraph_mode.requires_piecewise_compilation() |
| 50 | or self.compilation_config.is_attention_compiled_piecewise() |
| 51 | or is_breakable_cudagraph_enabled() |
| 52 | ), ( |
| 53 | "Compilation mode should be CompilationMode.VLLM_COMPILE when " |
| 54 | "cudagraph_mode piecewise cudagraphs is used, " |
| 55 | "and attention should be in splitting_ops or " |
| 56 | "inductor splitting should be used. " |
| 57 | f"cudagraph_mode={self.compilation_config.cudagraph_mode}, " |
| 58 | f"compilation_mode={self.compilation_config.mode}, " |
| 59 | f"splitting_ops={self.compilation_config.splitting_ops}" |
| 60 | ) |
| 61 | |
| 62 | self.keys_initialized = False |
| 63 | self.specialize_lora_count = ( |
| 64 | self.vllm_config.lora_config.specialize_active_lora |
| 65 | if self.vllm_config.lora_config is not None |
| 66 | else False |
| 67 | ) |
| 68 | # Default cudagraph_mode to NONE until initialize_cudagraph_keys is called |
| 69 | self.cudagraph_mode = CUDAGraphMode.NONE |
| 70 | |
| 71 | def _compute_bs_to_padded_graph_size(self) -> None: |
| 72 | """Pre-compute the mapping from batch size to padded graph size.""" |
| 73 | max_size = self.compilation_config.max_cudagraph_capture_size |
| 74 | capture_sizes = self.compilation_config.cudagraph_capture_sizes |
| 75 | assert max_size is not None, ( |
| 76 | "Maximum cudagraph capture size must be set when cudagraphs are enabled." |
| 77 | ) |
| 78 | assert capture_sizes is not None, ( |
| 79 | "Cudagraph capture sizes must be set when cudagraphs are enabled." |
| 80 | ) |
| 81 | self._bs_to_padded_graph_size: list[int] = [0] * (max_size + 1) |
| 82 | for end, start in zip( |
| 83 | capture_sizes + [max_size + 1], |
| 84 | [0] + capture_sizes, |
| 85 | ): |
| 86 | for bs in range(start, end): |
| 87 | if bs == start: |
| 88 | self._bs_to_padded_graph_size[bs] = start |
| 89 | else: |
| 90 | self._bs_to_padded_graph_size[bs] = end |
| 91 | |
| 92 | # Validate that compile_sizes won't be changed by padding. |
| 93 | # Only validate when cudagraphs are actually being used. |
| 94 | if ( |
| 95 | self.compilation_config.compile_sizes |
| 96 | and self.cudagraph_mode != CUDAGraphMode.NONE |
| 97 | ): |
| 98 | for size in self.compilation_config.compile_sizes: |
| 99 | size = int(size) |
| 100 | if size <= max_size: |
| 101 | padded = self._bs_to_padded_graph_size[size] |
| 102 | if padded != size: |
| 103 | raise ValueError( |
| 104 | f"compile_sizes contains {size} which would be " |
| 105 | f"padded to {padded}. All compile_sizes must be " |
| 106 | "values that won't be changed by cudagraph padding. " |
| 107 | "Use values from cudagraph_capture_sizes." |
| 108 | ) |
| 109 | |
| 110 | def _get_lora_cases(self) -> list[int]: |
| 111 | """Returns list of has_lora values for CUDA graph capture. |
| 112 | This is the single source of truth for LoRA capture cases. |
| 113 | """ |
| 114 | lora_config = self.vllm_config.lora_config |
| 115 | if lora_config is None: |
| 116 | # No LoRA configured - single case with no LoRA |
| 117 | return [0] |
| 118 | |
| 119 | # LoRA is enabled - capture graphs based on cudagraph_specialize_lora |
| 120 | if self.compilation_config.cudagraph_specialize_lora: |
| 121 | captured_counts = get_captured_lora_counts( |
| 122 | lora_config.max_loras, self.specialize_lora_count |
| 123 | ) |
| 124 | # Specialize: capture separate graphs for with and without LoRA |
| 125 | return [0] + captured_counts |
| 126 | else: |
| 127 | # No specialization: only capture graphs with LoRA active |
| 128 | return [lora_config.max_loras + 1] |
| 129 | |
| 130 | def _create_padded_batch_descriptor( |
| 131 | self, |
| 132 | num_tokens: int, |
| 133 | uniform_decode: bool, |
| 134 | has_lora: bool, |
| 135 | num_active_loras: int = 0, |
| 136 | ) -> BatchDescriptor: |
| 137 | max_num_seqs = self.vllm_config.scheduler_config.max_num_seqs |
| 138 | uniform_decode_query_len = self.uniform_decode_query_len |
| 139 | num_tokens_padded = self._bs_to_padded_graph_size[num_tokens] |
| 140 | |
| 141 | if uniform_decode and self.cudagraph_mode.has_mode(CUDAGraphMode.FULL): |
| 142 | num_reqs = min(num_tokens_padded // uniform_decode_query_len, max_num_seqs) |
| 143 | assert num_tokens_padded % uniform_decode_query_len == 0 |
| 144 | else: |
| 145 | uniform_decode = False |
| 146 | num_reqs = min(num_tokens_padded, max_num_seqs) |
| 147 | |
| 148 | return BatchDescriptor( |
| 149 | num_tokens=num_tokens_padded, |
| 150 | num_reqs=num_reqs, |
| 151 | uniform=uniform_decode, |
| 152 | has_lora=has_lora, |
| 153 | num_active_loras=num_active_loras, |
| 154 | ) |
| 155 | |
| 156 | def add_cudagraph_key( |
| 157 | self, runtime_mode: CUDAGraphMode, batch_descriptor: BatchDescriptor |
| 158 | ): |
| 159 | assert runtime_mode in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL], ( |
| 160 | f"Invalid cudagraph runtime mode for keys: {runtime_mode}" |
| 161 | ) |
| 162 | self.cudagraph_keys[runtime_mode].add(batch_descriptor) |
| 163 | |
| 164 | def initialize_cudagraph_keys( |
| 165 | self, cudagraph_mode: CUDAGraphMode, uniform_decode_query_len: int = 1 |
| 166 | ): |
| 167 | # This should be called only after attention backend is initialized. So we can |
| 168 | # get the correct cudagraph mode after backend support is resolved. |
| 169 | self.cudagraph_mode = cudagraph_mode |
| 170 | |
| 171 | # Early exit if cudagraphs are disabled |
| 172 | if cudagraph_mode == CUDAGraphMode.NONE: |
| 173 | self.keys_initialized = True |
| 174 | return |
| 175 | |
| 176 | self._compute_bs_to_padded_graph_size() |
| 177 | |
| 178 | # Get LoRA cases to capture |
| 179 | lora_cases = self._get_lora_cases() |
| 180 | self.captured_lora_counts = [ |
| 181 | lora_count for lora_count in lora_cases if lora_count |
| 182 | ] |
| 183 | |
| 184 | # Note: we create all valid keys for cudagraph here but do not |
| 185 | # guarantee all keys would be used. For example, if we allow lazy |
| 186 | # capturing in future PR, some keys may never be triggered. |
| 187 | if cudagraph_mode.mixed_mode() != CUDAGraphMode.NONE: |
| 188 | assert self.compilation_config.cudagraph_capture_sizes is not None, ( |
| 189 | "Cudagraph capture sizes must be set when mixed mode is enabled." |
| 190 | ) |
| 191 | for bs, num_active_loras in product( |
| 192 | self.compilation_config.cudagraph_capture_sizes, lora_cases |
| 193 | ): |
| 194 | batch_desc = self._create_padded_batch_descriptor( |
| 195 | bs, False, num_active_loras > 0, num_active_loras |
| 196 | ) |
| 197 | # Only relax for PIECEWISE mode. FULL mode needs exact num_reqs |
| 198 | # because FA3's scheduler_metadata computation depends on it. |
| 199 | if cudagraph_mode.mixed_mode() == CUDAGraphMode.PIECEWISE: |
| 200 | batch_desc = replace(batch_desc, num_reqs=None, uniform=False) |
| 201 | self.add_cudagraph_key(cudagraph_mode.mixed_mode(), batch_desc) |
| 202 | |
| 203 | # if decode cudagraph mode is FULL, and we don't already have mixed |
| 204 | # mode full cudagraphs then add them here. |
| 205 | if ( |
| 206 | cudagraph_mode.decode_mode() == CUDAGraphMode.FULL |
| 207 | and cudagraph_mode.separate_routine() |
| 208 | ): |
| 209 | max_num_tokens = ( |
| 210 | uniform_decode_query_len |
| 211 | * self.vllm_config.scheduler_config.max_num_seqs |
| 212 | ) |
| 213 | assert self.compilation_config.cudagraph_capture_sizes is not None, ( |
| 214 | "Cudagraph capture sizes must be set when full mode is enabled." |
| 215 | ) |
| 216 | cudagraph_capture_sizes_for_decode = [ |
| 217 | x |
| 218 | for x in self.compilation_config.cudagraph_capture_sizes |
| 219 | if x <= max_num_tokens and x >= uniform_decode_query_len |
| 220 | ] |
| 221 | for bs, num_active_loras in product( |
| 222 | cudagraph_capture_sizes_for_decode, lora_cases |
| 223 | ): |
| 224 | self.add_cudagraph_key( |
| 225 | CUDAGraphMode.FULL, |
| 226 | self._create_padded_batch_descriptor( |
| 227 | bs, True, num_active_loras > 0, num_active_loras |
| 228 | ), |
| 229 | ) |
| 230 | |
| 231 | self.keys_initialized = True |
| 232 | |
| 233 | def dispatch( |
| 234 | self, |
| 235 | num_tokens: int, |
| 236 | uniform_decode: bool = False, |
| 237 | has_lora: bool = False, |
| 238 | num_active_loras: int = 0, |
| 239 | valid_modes: AbstractSet[CUDAGraphMode] | None = None, |
| 240 | invalid_modes: AbstractSet[CUDAGraphMode] | None = None, |
| 241 | ) -> tuple[CUDAGraphMode, BatchDescriptor]: |
| 242 | """Given conditions(e.g.,batch descriptor and if using piecewise only), |
| 243 | dispatch to a cudagraph runtime mode and the valid batch descriptor. |
| 244 | A new batch descriptor is returned as we might dispatch a uniform batch |
| 245 | to a graph that supports a more general batch (uniform to non-uniform). |
| 246 | |
| 247 | Args: |
| 248 | num_tokens: Number of tokens in the batch. |
| 249 | uniform_decode: Whether the batch is uniform decode (i.e. uniform and query |
| 250 | length is uniform_decode_query_len). |
| 251 | has_lora: Whether LoRA is active. |
| 252 | num_active_loras: Number of distinct active LoRA adapters. |
| 253 | valid_modes: Set of cudagraph modes that are allowed. None means |
| 254 | all modes are allowed. |
| 255 | invalid_modes: Set of cudagraph modes to exclude. Subtracted from |
| 256 | valid_modes to compute allowed modes. (e.g., {FULL} for |
| 257 | features like cascade attention not supported by full |
| 258 | cudagraphs). None means no modes are excluded. |
| 259 | |
| 260 | """ |
| 261 | allowed_modes = valid_modes or CUDAGraphMode.valid_runtime_modes() |
| 262 | |
| 263 | if invalid_modes: |
| 264 | allowed_modes -= invalid_modes |
| 265 | |
| 266 | assert len(allowed_modes) >= 1, ( |
| 267 | f"No allowed cudagraph modes: valid_modes={valid_modes}, " |
| 268 | f"invalid_modes={invalid_modes}" |
| 269 | ) |
| 270 | max_size = self.compilation_config.max_cudagraph_capture_size |
| 271 | |
| 272 | if ( |
| 273 | not self.keys_initialized |
| 274 | or self.cudagraph_mode == CUDAGraphMode.NONE |
| 275 | or max_size is None |
| 276 | or num_tokens > max_size |
| 277 | or allowed_modes <= {CUDAGraphMode.NONE} |
| 278 | ): |
| 279 | return CUDAGraphMode.NONE, BatchDescriptor(num_tokens) |
| 280 | |
| 281 | effective_num_active_loras = num_active_loras |
| 282 | if has_lora and num_active_loras > 0: |
| 283 | if self.specialize_lora_count: |
| 284 | # Find the smallest captured `num_active_loras` that is >= the current |
| 285 | # `num_active_loras`. This is because we only capture graphs for |
| 286 | # a subset of possible `num_active_loras` values (powers of 2). |
| 287 | import bisect |
| 288 | |
| 289 | idx = bisect.bisect_left(self.captured_lora_counts, num_active_loras) |
| 290 | if idx < len(self.captured_lora_counts): |
| 291 | effective_num_active_loras = self.captured_lora_counts[idx] |
| 292 | else: |
| 293 | # When not specializing, graphs are captured only with max_loras + 1, |
| 294 | # so we must use max_loras + 1 for dispatch to find a matching graph. |
| 295 | assert self.vllm_config.lora_config is not None, ( |
| 296 | "LoRA config must be set when has_lora is True." |
| 297 | ) |
| 298 | effective_num_active_loras = self.vllm_config.lora_config.max_loras + 1 |
| 299 | |
| 300 | normalized_uniform = uniform_decode and self.cudagraph_mode.separate_routine() |
| 301 | batch_desc = self._create_padded_batch_descriptor( |
| 302 | num_tokens, normalized_uniform, has_lora, effective_num_active_loras |
| 303 | ) |
| 304 | |
| 305 | if CUDAGraphMode.FULL in allowed_modes: |
| 306 | # check if key exists for full cudagraph |
| 307 | batch_desc_to_check = batch_desc |
| 308 | if batch_desc_to_check in self.cudagraph_keys[CUDAGraphMode.FULL]: |
| 309 | return CUDAGraphMode.FULL, batch_desc_to_check |
| 310 | |
| 311 | if CUDAGraphMode.PIECEWISE in allowed_modes: |
| 312 | # also check if the relaxed key exists for more "general" |
| 313 | # piecewise cudagraph |
| 314 | batch_desc_to_check = replace(batch_desc, num_reqs=None, uniform=False) |
| 315 | if batch_desc_to_check in self.cudagraph_keys[CUDAGraphMode.PIECEWISE]: |
| 316 | return CUDAGraphMode.PIECEWISE, batch_desc_to_check |
| 317 | |
| 318 | assert CUDAGraphMode.NONE in allowed_modes, ( |
| 319 | f"No matching cudagraph found and NONE is not in " |
| 320 | f"allowed_modes={allowed_modes}" |
| 321 | ) |
| 322 | return CUDAGraphMode.NONE, BatchDescriptor(num_tokens) |
| 323 | |
| 324 | def get_capture_descs(self) -> list[tuple[CUDAGraphMode, list[BatchDescriptor]]]: |
| 325 | """Returns capture descriptors for cudagraph capturing. |
| 326 | |
| 327 | Returns: |
| 328 | List of (runtime_mode, batch_descriptors) tuples, ordered PIECEWISE |
| 329 | first then FULL. Batch descriptors are sorted largest-first for |
| 330 | memory efficiency. |
| 331 | |
| 332 | """ |
| 333 | if not self.keys_initialized or self.cudagraph_mode == CUDAGraphMode.NONE: |
| 334 | return [] |
| 335 | |
| 336 | result = [] |
| 337 | # Return in order: PIECEWISE first, then FULL |
| 338 | for mode in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL]: |
| 339 | descs = list(self.cudagraph_keys[mode]) |
| 340 | if descs: |
| 341 | # Sort by (num_tokens, num_active_loras) descending |
| 342 | descs.sort( |
| 343 | key=lambda d: (d.num_tokens, d.num_active_loras), |
| 344 | reverse=True, |
| 345 | ) |
| 346 | result.append((mode, descs)) |
| 347 | |
| 348 | return result |