| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import enum |
| 5 | from collections import Counter |
| 6 | from collections.abc import Callable |
| 7 | from dataclasses import field, fields |
| 8 | from pathlib import Path |
| 9 | from typing import TYPE_CHECKING, Any, ClassVar, Literal |
| 10 | |
| 11 | from pydantic import Field, TypeAdapter, field_validator |
| 12 | |
| 13 | import vllm.envs as envs |
| 14 | from vllm.compilation.passes.inductor_pass import CallableInductorPass, InductorPass |
| 15 | from vllm.config.utils import ( |
| 16 | Range, |
| 17 | config, |
| 18 | get_hash_factors, |
| 19 | hash_factors, |
| 20 | ) |
| 21 | from vllm.logger import init_logger |
| 22 | from vllm.platforms import current_platform |
| 23 | from vllm.utils.import_utils import resolve_obj_by_qualname |
| 24 | from vllm.utils.math_utils import round_up |
| 25 | from vllm.utils.torch_utils import is_torch_equal_or_newer |
| 26 | |
| 27 | if TYPE_CHECKING: |
| 28 | from vllm.config import VllmConfig |
| 29 | from vllm.v1.attention.backend import AttentionCGSupport |
| 30 | from vllm.v1.kv_cache_interface import KVCacheConfig |
| 31 | else: |
| 32 | VllmConfig = object |
| 33 | |
| 34 | logger = init_logger(__name__) |
| 35 | |
| 36 | |
| 37 | class CompilationMode(enum.IntEnum): |
| 38 | """The compilation approach used for torch.compile-based compilation of the |
| 39 | model.""" |
| 40 | |
| 41 | NONE = 0 |
| 42 | """No torch.compile compilation is applied, model runs in fully eager pytorch mode. |
| 43 | The model runs as-is.""" |
| 44 | STOCK_TORCH_COMPILE = 1 |
| 45 | """The standard `torch.compile` compilation pipeline.""" |
| 46 | DYNAMO_TRACE_ONCE = 2 |
| 47 | """Single Dynamo trace through the model, avoiding recompilation.""" |
| 48 | VLLM_COMPILE = 3 |
| 49 | """Custom vLLM Inductor-based backend with caching, piecewise compilation, |
| 50 | shape specialization, and custom passes.""" |
| 51 | |
| 52 | |
| 53 | class CUDAGraphMode(enum.Enum): |
| 54 | """Constants for the cudagraph mode in CompilationConfig. |
| 55 | Meanwhile, the subset enum `NONE`, `PIECEWISE` and `FULL` are also |
| 56 | treated as concrete runtime mode for cudagraph runtime dispatching. |
| 57 | """ |
| 58 | |
| 59 | NONE = 0 |
| 60 | PIECEWISE = 1 |
| 61 | FULL = 2 |
| 62 | FULL_DECODE_ONLY = (FULL, NONE) |
| 63 | FULL_AND_PIECEWISE = (FULL, PIECEWISE) |
| 64 | |
| 65 | def decode_mode(self) -> "CUDAGraphMode": |
| 66 | return CUDAGraphMode(self.value[0]) if self.separate_routine() else self |
| 67 | |
| 68 | def mixed_mode(self) -> "CUDAGraphMode": |
| 69 | return CUDAGraphMode(self.value[1]) if self.separate_routine() else self |
| 70 | |
| 71 | def has_mode(self, mode: "CUDAGraphMode") -> bool: |
| 72 | assert not mode.separate_routine() |
| 73 | if self.separate_routine(): |
| 74 | return mode.value in self.value |
| 75 | return self == mode |
| 76 | |
| 77 | def requires_piecewise_compilation(self) -> bool: |
| 78 | return self.has_mode(CUDAGraphMode.PIECEWISE) |
| 79 | |
| 80 | def without_piecewise(self) -> "CUDAGraphMode": |
| 81 | if self == CUDAGraphMode.PIECEWISE: |
| 82 | return CUDAGraphMode.NONE |
| 83 | if self == CUDAGraphMode.FULL_AND_PIECEWISE: |
| 84 | return CUDAGraphMode.FULL_DECODE_ONLY |
| 85 | return self |
| 86 | |
| 87 | def max_cudagraph_mode(self) -> "CUDAGraphMode": |
| 88 | return CUDAGraphMode(max(self.value)) if self.separate_routine() else self |
| 89 | |
| 90 | def has_full_cudagraphs(self) -> bool: |
| 91 | return self.max_cudagraph_mode() == CUDAGraphMode.FULL |
| 92 | |
| 93 | def has_piecewise_cudagraphs(self) -> bool: |
| 94 | return self.requires_piecewise_compilation() |
| 95 | |
| 96 | def separate_routine(self) -> bool: |
| 97 | return isinstance(self.value, tuple) |
| 98 | |
| 99 | @classmethod |
| 100 | def valid_runtime_modes(cls) -> frozenset["CUDAGraphMode"]: |
| 101 | return frozenset({cls.NONE, cls.PIECEWISE, cls.FULL}) |
| 102 | |
| 103 | def is_valid_runtime_mode(self) -> bool: |
| 104 | return self in CUDAGraphMode.valid_runtime_modes() |
| 105 | |
| 106 | def __str__(self) -> str: |
| 107 | return self.name |
| 108 | |
| 109 | def __bool__(self) -> bool: |
| 110 | return self != CUDAGraphMode.NONE |
| 111 | |
| 112 | |
| 113 | @config |
| 114 | class PassConfig: |
| 115 | """Configuration for custom Inductor passes. |
| 116 | |
| 117 | This is separate from general `CompilationConfig` so that inductor passes |
| 118 | don't all have access to full configuration - that would create a cycle as |
| 119 | the `PassManager` is set as a property of config. |
| 120 | |
| 121 | You must pass PassConfig to VLLMConfig constructor via the CompilationConfig |
| 122 | constructor. VLLMConfig's post_init does further initialization. |
| 123 | If used outside of the VLLMConfig, some fields may be left in an |
| 124 | improper state. |
| 125 | """ |
| 126 | |
| 127 | # New flags |
| 128 | fuse_norm_quant: bool = None # type: ignore[assignment] |
| 129 | """Fuse the custom RMSNorm + quant ops.""" |
| 130 | fuse_act_quant: bool = None # type: ignore[assignment] |
| 131 | """Fuse the custom SiluMul + quant ops.""" |
| 132 | fuse_attn_quant: bool = None # type: ignore[assignment] |
| 133 | """Fuse the custom Attention and MLAAttention + quant ops.""" |
| 134 | eliminate_noops: bool = Field(default=True) |
| 135 | """Eliminate no-op ops.""" |
| 136 | enable_sp: bool = None # type: ignore[assignment] |
| 137 | """Enable sequence parallelism. Requires TP>1. Automatically disabled |
| 138 | if the model's hidden_size is too small for SP to be beneficial |
| 139 | (threshold is device-capability dependent).""" |
| 140 | fuse_gemm_comms: bool = None # type: ignore[assignment] |
| 141 | """Enable async TP.""" |
| 142 | fuse_allreduce_rms: bool = None # type: ignore[assignment] |
| 143 | """Enable flashinfer allreduce fusion.""" |
| 144 | enable_qk_norm_rope_fusion: bool = None # type: ignore[assignment] |
| 145 | """Enable fused Q/K RMSNorm + RoPE pass.""" |
| 146 | fuse_rope_kvcache_cat_mla: bool = None # type: ignore[assignment] |
| 147 | """Enable fused MLA KV cache update with RoPE.""" |
| 148 | |
| 149 | # ROCm/AITER specific fusions |
| 150 | fuse_act_padding: bool = None # type: ignore[assignment] |
| 151 | """Fuse the custom RMSNorm + padding ops.""" |
| 152 | fuse_mla_dual_rms_norm: bool = None # type: ignore[assignment] |
| 153 | """Fuse paired q/kv RMS norms in MLA attention.""" |
| 154 | fuse_rope_kvcache: bool = None # type: ignore[assignment] |
| 155 | """Fuse the QK rope + KV cache ops.""" |
| 156 | fuse_qk_norm_rope_kvcache: bool = Field(default=None) # type: ignore[assignment] |
| 157 | """Fuse QK RMSNorm + RoPE/MRoPE + KV cache update into an AITER HIP |
| 158 | kernel. Supersedes both enable_qk_norm_rope_fusion and fuse_rope_kvcache |
| 159 | for layers that support it. Auto-enabled at O2+ on ROCm for models |
| 160 | with QK-norm (e.g. Qwen3-MoE and Qwen3-VL-class architectures).""" |
| 161 | |
| 162 | rope_kvcache_fusion_max_token_num: int = 256 |
| 163 | """The threshold for ROCm AITER RoPE+KVCache fusion e.g. for small batch decode. |
| 164 | Larger batch sizes e.g. during prefill will use the unfused kernels. |
| 165 | Also applies to the fused QK-Norm+RoPE/MRoPE+KVCache pass. |
| 166 | """ |
| 167 | |
| 168 | fi_allreduce_fusion_max_size_mb: float | None = None |
| 169 | """The threshold of the communicated tensor sizes under which |
| 170 | vllm should use flashinfer fused allreduce. Specified as a |
| 171 | float in MB. |
| 172 | Unspecified will fallback to default values |
| 173 | which are compute capability and world size dependent. |
| 174 | FI_ALLREDUCE_FUSION_MAX_SIZE_MB = { |
| 175 | 90: { |
| 176 | 2: 64, # 64MB |
| 177 | 4: 2, # 2MB |
| 178 | 8: 1, # 1MB |
| 179 | }, |
| 180 | 100: { |
| 181 | 2: 64, # 64MB |
| 182 | 4: 32, # 32MB |
| 183 | 8: 1, # 1MB |
| 184 | }, |
| 185 | }, where key is the device capability""" |
| 186 | sp_min_token_num: int | None = None |
| 187 | """The minimum number of tokens above which vllm should use |
| 188 | sequence parallelism. Specified as an integer token count. |
| 189 | Unspecified will fallback to default values which are compute |
| 190 | capability and world size dependent.""" |
| 191 | |
| 192 | # TODO(luka) better pass enabling system. |
| 193 | |
| 194 | def flashinfer_max_size(self, world_size: int) -> int | None: |
| 195 | """Returns the max communication size in bytes for flashinfer |
| 196 | allreduce fusion for the given world size. Returns None if world size |
| 197 | is not supported by configs as it's not supported by flashinfer. |
| 198 | """ |
| 199 | MiB = 1024 * 1024 |
| 200 | FI_SUPPORTED_WORLD_SIZES = [2, 4, 8, 16] |
| 201 | if world_size not in FI_SUPPORTED_WORLD_SIZES: |
| 202 | return None |
| 203 | max_size_mb = self.fi_allreduce_fusion_max_size_mb |
| 204 | if max_size_mb is None: |
| 205 | max_size_mb = self.default_fi_allreduce_fusion_max_size_mb().get(world_size) |
| 206 | |
| 207 | return int(max_size_mb * MiB) if max_size_mb is not None else None |
| 208 | |
| 209 | @staticmethod |
| 210 | def default_fi_allreduce_fusion_max_size_mb() -> dict[int, float]: |
| 211 | from vllm.compilation.passes.fusion.allreduce_rms_fusion import ( |
| 212 | FI_ALLREDUCE_FUSION_MAX_SIZE_MB, |
| 213 | ) |
| 214 | from vllm.platforms import current_platform |
| 215 | |
| 216 | if not current_platform.is_cuda(): |
| 217 | return {} |
| 218 | capability = current_platform.get_device_capability() |
| 219 | if capability is None: |
| 220 | return {} |
| 221 | return FI_ALLREDUCE_FUSION_MAX_SIZE_MB.get(capability.to_int(), {}) |
| 222 | |
| 223 | def compute_hash(self) -> str: |
| 224 | """Produces a hash unique to the pass configuration. |
| 225 | Any new fields that affect compilation should be added to the hash. |
| 226 | Any future fields that don't affect compilation should be excluded. |
| 227 | """ |
| 228 | return hash_factors(get_hash_factors(self, set())) |
| 229 | |
| 230 | @field_validator( |
| 231 | "fuse_norm_quant", |
| 232 | "fuse_act_quant", |
| 233 | "fuse_attn_quant", |
| 234 | "enable_sp", |
| 235 | "fuse_gemm_comms", |
| 236 | "fuse_allreduce_rms", |
| 237 | "fuse_act_padding", |
| 238 | "fuse_mla_dual_rms_norm", |
| 239 | "fuse_rope_kvcache", |
| 240 | "fuse_qk_norm_rope_kvcache", |
| 241 | "enable_qk_norm_rope_fusion", |
| 242 | "fuse_rope_kvcache_cat_mla", |
| 243 | mode="wrap", |
| 244 | ) |
| 245 | @classmethod |
| 246 | def _skip_none_validation(cls, value: Any, handler: Callable) -> Any: |
| 247 | """Skip validation if the value is `None` when initialisation is delayed.""" |
| 248 | if value is None: |
| 249 | return value |
| 250 | return handler(value) |
| 251 | |
| 252 | def __post_init__(self) -> None: |
| 253 | # Handle deprecation and defaults |
| 254 | |
| 255 | if not self.eliminate_noops: |
| 256 | if self.fuse_norm_quant or self.fuse_act_quant: |
| 257 | logger.warning_once( |
| 258 | "Fusion enabled but reshape elimination disabled. " |
| 259 | "RMSNorm/SiluMul + quant (fp8) fusion might not work" |
| 260 | ) |
| 261 | if self.fuse_attn_quant: |
| 262 | logger.warning_once( |
| 263 | "Fusion enabled but reshape elimination disabled. " |
| 264 | "Attention + quant (fp8) fusion might not work" |
| 265 | ) |
| 266 | if self.fuse_allreduce_rms: |
| 267 | logger.warning_once( |
| 268 | "Fusion enabled but reshape elimination disabled. " |
| 269 | "Allreduce + rms norm + quant (fp8) fusion might not work" |
| 270 | ) |
| 271 | if self.fuse_act_padding: |
| 272 | logger.warning_once( |
| 273 | "Fusion enabled but reshape elimination disabled. " |
| 274 | "RMSNorm + padding fusion might not work" |
| 275 | ) |
| 276 | if self.enable_qk_norm_rope_fusion and not ( |
| 277 | current_platform.is_cuda_alike() or current_platform.is_xpu() |
| 278 | ): |
| 279 | logger.warning_once( |
| 280 | "QK Norm + RoPE fusion enabled but the current platform is not " |
| 281 | "CUDA, ROCm or XPU. The fusion will be disabled." |
| 282 | ) |
| 283 | self.enable_qk_norm_rope_fusion = False |
| 284 | if self.fuse_act_padding and not current_platform.is_rocm(): |
| 285 | logger.warning_once( |
| 286 | "Padding fusion enabled but the current platform is not ROCm. " |
| 287 | "The fusion will be disabled." |
| 288 | ) |
| 289 | self.fuse_act_padding = False |
| 290 | if self.fuse_mla_dual_rms_norm and not current_platform.is_rocm(): |
| 291 | logger.warning_once( |
| 292 | "MLA dual RMS norm fusion requires ROCm/AITER. " |
| 293 | "The fusion will be disabled." |
| 294 | ) |
| 295 | self.fuse_mla_dual_rms_norm = False |
| 296 | if self.fuse_rope_kvcache and not current_platform.is_rocm(): |
| 297 | logger.warning_once( |
| 298 | "KV cache fusion currently only enabled on ROCm. " |
| 299 | "The fusion will be disabled." |
| 300 | ) |
| 301 | self.fuse_rope_kvcache = False |
| 302 | if self.fuse_qk_norm_rope_kvcache and not current_platform.is_rocm(): |
| 303 | logger.warning_once( |
| 304 | "QK-Norm+RoPE+KVCache fusion requires ROCm with AITER. " |
| 305 | "The fusion will be disabled." |
| 306 | ) |
| 307 | self.fuse_qk_norm_rope_kvcache = False |
| 308 | if self.fuse_rope_kvcache_cat_mla and not current_platform.is_cuda_alike(): |
| 309 | logger.warning_once( |
| 310 | "MLA KV cache update with RoPE fusion enabled but the " |
| 311 | "current platform is not CUDA or ROCm. The fusion will be disabled." |
| 312 | ) |
| 313 | self.fuse_rope_kvcache_cat_mla = False |
| 314 | |
| 315 | def log_enabled_passes(self) -> None: |
| 316 | """Log the enabled custom fusion passes. |
| 317 | This is called at the end of VLLMConfig post_init, |
| 318 | after all defaults are finalized. |
| 319 | TODO also log the compile ranges for which this is enabled. |
| 320 | """ |
| 321 | fusion_prefixes = ("fuse_", "enable_") |
| 322 | enabled_fusions = [ |
| 323 | f.name[len(prefix) :] |
| 324 | for f in fields(self) # type: ignore[arg-type] |
| 325 | if getattr(self, f.name) |
| 326 | for prefix in fusion_prefixes |
| 327 | if f.name.startswith(prefix) |
| 328 | ] |
| 329 | |
| 330 | if enabled_fusions: |
| 331 | logger.info_once( |
| 332 | "Enabled custom fusions: %s", ", ".join(enabled_fusions), scope="global" |
| 333 | ) |
| 334 | |
| 335 | |
| 336 | class DynamicShapesType(str, enum.Enum): |
| 337 | """Types of dynamic shapes handling in torch.compile(). |
| 338 | see Dynamic shapes and vllm guard dropping in torch_compile.md |
| 339 | for more details.""" |
| 340 | |
| 341 | BACKED = "backed" |
| 342 | """Use backed dynamic shapes. torch.compile() guards on backed dynamic |
| 343 | shapes and may add guards. Symbols are specialized to 0, 1, or >=2 even |
| 344 | without encountering branching on those ranges.""" |
| 345 | |
| 346 | UNBACKED = "unbacked" |
| 347 | """Use unbacked dynamic shapes. Guaranteed not to be guarded on and not |
| 348 | 0/1 specialized, but may throw data dependent errors when branches require |
| 349 | their value without explicit unbacked handling.""" |
| 350 | |
| 351 | BACKED_SIZE_OBLIVIOUS = "backed_size_oblivious" |
| 352 | """Experimental flag that treats backed symbols as unbacked when explicit |
| 353 | unbacked handling is defined.""" |
| 354 | |
| 355 | |
| 356 | @config |
| 357 | class DynamicShapesConfig: |
| 358 | """Configuration to control/debug torch compile dynamic shapes.""" |
| 359 | |
| 360 | type: DynamicShapesType = DynamicShapesType.BACKED |
| 361 | """Controls the type of dynamic shapes handling to use with torch.compile(). |
| 362 | |
| 363 | - BACKED: Default PyTorch behavior with potential guards ignored. |
| 364 | - UNBACKED: No guards guaranteed (most sound) but may throw |
| 365 | data dependent errors. |
| 366 | - BACKED_SIZE_OBLIVIOUS: Experimental safer alternative to |
| 367 | backed/unbacked. |
| 368 | """ |
| 369 | |
| 370 | evaluate_guards: bool = False |
| 371 | """ |
| 372 | A debug mode to detect and fail if Dynamo ever specializes a dynamic shape by |
| 373 | guarding on it. When True, dynamic shape guards are not dropped from dynamo. |
| 374 | And a failure will be triggered if a recompilation ever happens due to that. |
| 375 | This mode requires VLLM_USE_BYTECODE_HOOK to be 0. |
| 376 | Enabling this allow observing the dynamic shapes guards in the tlparse |
| 377 | artifacts also. |
| 378 | When type is backed, aot_compile must be disabled for this mode to work. |
| 379 | until this change picked up https://github.com/pytorch/pytorch/pull/169239. |
| 380 | """ |
| 381 | |
| 382 | assume_32_bit_indexing: bool = False |
| 383 | """ |
| 384 | whether all tensor sizes can use 32 bit indexing. |
| 385 | `True` requires PyTorch 2.10+ |
| 386 | """ |
| 387 | |
| 388 | def compute_hash(self) -> str: |
| 389 | """Provide a hash for DynamicShapesConfig.""" |
| 390 | from vllm.config.utils import get_hash_factors, hash_factors |
| 391 | |
| 392 | factors = get_hash_factors(self, set()) |
| 393 | return hash_factors(factors) |
| 394 | |
| 395 | |
| 396 | @config |
| 397 | class CompilationConfig: |
| 398 | """Configuration for compilation. |
| 399 | |
| 400 | You must pass CompilationConfig to VLLMConfig constructor. |
| 401 | VLLMConfig's post_init does further initialization. If used outside of the |
| 402 | VLLMConfig, some fields will be left in an improper state. |
| 403 | |
| 404 | It contains PassConfig, which controls the custom fusion/transformation passes. |
| 405 | The rest has three parts: |
| 406 | |
| 407 | - Top-level Compilation control: |
| 408 | - [`mode`][vllm.config.CompilationConfig.mode] |
| 409 | - [`debug_dump_path`][vllm.config.CompilationConfig.debug_dump_path] |
| 410 | - [`cache_dir`][vllm.config.CompilationConfig.cache_dir] |
| 411 | - [`backend`][vllm.config.CompilationConfig.backend] |
| 412 | - [`custom_ops`][vllm.config.CompilationConfig.custom_ops] |
| 413 | - [`splitting_ops`][vllm.config.CompilationConfig.splitting_ops] |
| 414 | - [`compile_mm_encoder`][vllm.config.CompilationConfig.compile_mm_encoder] |
| 415 | - CudaGraph capture: |
| 416 | - [`cudagraph_mode`][vllm.config.CompilationConfig.cudagraph_mode] |
| 417 | - [`cudagraph_capture_sizes`] |
| 418 | [vllm.config.CompilationConfig.cudagraph_capture_sizes] |
| 419 | - [`max_cudagraph_capture_size`] |
| 420 | [vllm.config.CompilationConfig.max_cudagraph_capture_size] |
| 421 | - [`cudagraph_num_of_warmups`] |
| 422 | [vllm.config.CompilationConfig.cudagraph_num_of_warmups] |
| 423 | - [`cudagraph_copy_inputs`] |
| 424 | [vllm.config.CompilationConfig.cudagraph_copy_inputs] |
| 425 | - Inductor compilation: |
| 426 | - [`compile_sizes`][vllm.config.CompilationConfig.compile_sizes] |
| 427 | - [`compile_ranges_endpoints`] |
| 428 | [vllm.config.CompilationConfig.compile_ranges_endpoints] |
| 429 | - [`inductor_compile_config`] |
| 430 | [vllm.config.CompilationConfig.inductor_compile_config] |
| 431 | - [`inductor_passes`][vllm.config.CompilationConfig.inductor_passes] |
| 432 | - custom inductor passes |
| 433 | |
| 434 | Why we have different sizes for cudagraph and inductor: |
| 435 | - cudagraph: a cudagraph captured for a specific size can only be used |
| 436 | for the same size. We need to capture all the sizes we want to use. |
| 437 | - inductor: a graph compiled by inductor for a general shape can be used |
| 438 | for different sizes. Inductor can also compile for specific sizes, |
| 439 | where it can have more information to optimize the graph with fully |
| 440 | static shapes. However, we find the general shape compilation is |
| 441 | sufficient for most cases. It might be beneficial to compile for |
| 442 | certain small batchsizes, where inductor is good at optimizing. |
| 443 | """ |
| 444 | |
| 445 | # Top-level Compilation control |
| 446 | mode: CompilationMode = None # type: ignore[assignment] |
| 447 | """The compilation approach used for torch.compile-based compilation of the |
| 448 | model. |
| 449 | |
| 450 | - None: If None, we will select the default compilation mode. |
| 451 | For V1 engine this is 3. |
| 452 | - 0: NONE: No torch.compile compilation is applied, model runs in fully |
| 453 | eager pytorch mode. The model runs as-is. |
| 454 | - 1: STOCK_TORCH_COMPILE: The standard `torch.compile` compilation pipeline. |
| 455 | - 2: DYNAMO_TRACE_ONCE: Single Dynamo trace through the model, avoiding |
| 456 | recompilation by removing guards. |
| 457 | Requires no dynamic-shape-dependent control-flow. |
| 458 | - 3: VLLM_COMPILE: Custom vLLM Inductor-based backend with caching, |
| 459 | piecewise compilation, shape specialization, and custom passes.""" |
| 460 | debug_dump_path: Path | None = None |
| 461 | """The path to dump the debug information.""" |
| 462 | cache_dir: str = "" |
| 463 | """The directory to store the compiled graph, to accelerate Inductor |
| 464 | compilation. By default, it will use model-related information to generate |
| 465 | a cache directory.""" |
| 466 | compile_cache_save_format: Literal["binary", "unpacked"] = field( |
| 467 | default_factory=lambda: envs.VLLM_COMPILE_CACHE_SAVE_FORMAT |
| 468 | ) |
| 469 | """Format for saving torch compile cache:\n |
| 470 | - "binary": saves as binary file (multiprocess safe)\n |
| 471 | - "unpacked": saves as directory structure for inspection/debugging |
| 472 | (NOT multiprocess safe)\n |
| 473 | Defaults to `VLLM_COMPILE_CACHE_SAVE_FORMAT` if not specified. |
| 474 | """ |
| 475 | backend: str = "" |
| 476 | """The backend for compilation. It needs to be a string: |
| 477 | |
| 478 | - "" (empty string): use the default backend ("inductor" on CUDA-alike |
| 479 | platforms). |
| 480 | - "eager"/"openxla"/...: use the specified backend registered in PyTorch. |
| 481 | - "full.module.name": a qualified name which can be used to import the |
| 482 | |
| 483 | backend function. |
| 484 | We use string to avoid serialization issues when using compilation in a |
| 485 | distributed setting. When the compilation mode is 1 or 2, the backend is |
| 486 | used for the compilation directly (it sees the whole graph). When the |
| 487 | compilation mode is 3, the backend supports both whole graph and piecewise |
| 488 | compilation, available backends include eager, inductor, and custom backends, |
| 489 | the latter of which can be defined via `get_compile_backend`. Furthermore, |
| 490 | compilation is only piecewise if splitting ops is set accordingly and |
| 491 | use_inductor_graph_partition is off. Note that the default options for |
| 492 | splitting ops are sufficient for piecewise compilation. |
| 493 | """ |
| 494 | custom_ops: list[str] = field(default_factory=list) |
| 495 | """Fine-grained control over which custom ops to enable/disable. Use 'all' |
| 496 | to enable all, 'none' to disable all. Also specify a list of custom op |
| 497 | names to enable (prefixed with a '+'), or disable (prefixed with a '-'). |
| 498 | Examples: |
| 499 | |
| 500 | - 'all,-op1' to enable all except op1 |
| 501 | - 'none,+op1,+op2' to enable only op1 and op2 |
| 502 | |
| 503 | By default, all custom ops are enabled when running without Inductor and |
| 504 | disabled when running with Inductor: mode>CompilationMode.NONE and |
| 505 | backend="inductor". |
| 506 | Inductor generates (fused) Triton kernels for disabled custom ops.""" |
| 507 | |
| 508 | ir_enable_torch_wrap: bool = None # type: ignore[assignment] |
| 509 | """If True, enable vllm_ir torch custom op wrapping during the forward pass. |
| 510 | When False, torch custom op wrapping is disabled, allowing Dynamo to trace the |
| 511 | selected implementation directly or avoiding torch custom op overhead in eager mode. |
| 512 | Defaults to True when using Inductor with vllm-compile |
| 513 | (backend=="inductor" and mode == VLLM_COMPILE), False otherwise. |
| 514 | """ |
| 515 | |
| 516 | splitting_ops: list[str] | None = None |
| 517 | """A list of ops to exclude from cudagraphs, used in piecewise compilation. |
| 518 | |
| 519 | The behavior depends on use_inductor_graph_partition: |
| 520 | |
| 521 | - When use_inductor_graph_partition=False (default): |
| 522 | These ops are used for Dynamo FX-level graph splitting. The graph is |
| 523 | split at these ops before Inductor compilation, creating separate |
| 524 | subgraphs for cudagraph capture. |
| 525 | |
| 526 | - When use_inductor_graph_partition=True: |
| 527 | These ops are used to register Inductor partition rules. The graph |
| 528 | partitioning happens at Inductor codegen time after all passes and |
| 529 | fusions are finished, allowing compilation and custom passes to operate |
| 530 | on the full graph while still excluding these ops from cudagraphs. |
| 531 | |
| 532 | If None, defaults to attention ops for piecewise cudagraphs. |
| 533 | If empty list [], no ops are excluded (suitable for full cudagraphs).""" |
| 534 | compile_mm_encoder: bool = False |
| 535 | """Whether or not to compile the multimodal encoder. |
| 536 | Currently, this only works for `Qwen2_5_vl` and `mLLaMa4` models on selected |
| 537 | platforms. It may also work for models loaded with the Transformers modeling backend |
| 538 | if the encoder is compilable. Disabled by default until more models are |
| 539 | supported/tested to work.""" |
| 540 | |
| 541 | # Vision encoder CUDA graph |
| 542 | cudagraph_mm_encoder: bool = False |
| 543 | """Enable CUDA graph capture for multimodal encoder (ViT). |
| 544 | When enabled, captures full encoder forward as CUDA graph |
| 545 | for each token budget level.""" |
| 546 | |
| 547 | encoder_cudagraph_token_budgets: list[int] = field(default_factory=list) |
| 548 | """Token budget levels for encoder CUDA graph capture. |
| 549 | Each budget defines a fixed token capacity. At runtime, images are greedy-packed |
| 550 | into the smallest fitting budget and the corresponding CUDA graph is replayed. |
| 551 | If empty (default), auto-inferred from model architecture as power-of-2 |
| 552 | levels from the model's estimated min budget to max budget. |
| 553 | User-provided values override auto-inference. |
| 554 | Example: [2048, 4096, 8192, 13824]""" |
| 555 | |
| 556 | encoder_cudagraph_max_vision_items_per_batch: int = 0 |
| 557 | """Maximum number of images/videos per batch for encoder CUDA graph capture. |
| 558 | Determines the fixed batch size used during graph capture. |
| 559 | If 0 (default), auto-inferred as max_budget // min_budget from the |
| 560 | model's budget range. User-provided positive value overrides |
| 561 | auto-inference.""" |
| 562 | |
| 563 | encoder_cudagraph_max_frames_per_batch: int | None = None |
| 564 | """Maximum total video frames per batch for encoder CUDA graph capture. |
| 565 | Controls the cu_seqlens buffer size (one entry per attention sequence, |
| 566 | i.e. one per video frame). |
| 567 | If None (default), auto-inferred as encoder_cudagraph_max_vision_items_per_batch |
| 568 | * max_frames_per_video (model-specific value according to processing_info). |
| 569 | Positive value overrides auto-inference and applies to all budget levels. |
| 570 | If we limit the video count per prompt to `0`, it will also be set to `0` |
| 571 | (i.e., fall back to image-only mode).""" |
| 572 | |
| 573 | # Inductor capture |
| 574 | compile_sizes: list[int | str] | None = None |
| 575 | """Sizes to compile for inductor. In addition |
| 576 | to integers, it also supports "cudagraph_capture_sizes" to |
| 577 | specify the sizes for cudagraph capture.""" |
| 578 | |
| 579 | compile_ranges_endpoints: list[int] | None = None |
| 580 | """Endpoints for Inductor compile ranges. |
| 581 | The compile ranges are |
| 582 | [1, endpoints[0]], |
| 583 | [endpoints[0] + 1, endpoints[1]], ..., |
| 584 | [endpoints[-1] + 1, max_num_batched_tokens]. |
| 585 | Compile sizes are also used single element ranges, |
| 586 | the range is represented as [compile_sizes[i], compile_sizes[i]]. |
| 587 | |
| 588 | If a range overlaps with the compile size, graph for compile size |
| 589 | will be prioritized, i.e. if we have a range [1, 8] and a compile size 4, |
| 590 | graph for compile size 4 will be compiled and used instead of the graph |
| 591 | for range [1, 8]. |
| 592 | """ |
| 593 | |
| 594 | inductor_compile_config: dict = field(default_factory=dict) |
| 595 | """Additional configurations for inductor. |
| 596 | - None: use default configurations.""" |
| 597 | |
| 598 | inductor_passes: dict[str, str] = field(default_factory=dict) |
| 599 | """Additional passes for inductor. It is a dictionary |
| 600 | from pass name to pass function qualified name. We use function |
| 601 | name because the config uses JSON format. If we pass the config |
| 602 | from Python, functions can also be passed directly via Python object |
| 603 | constructor, e.g. `CompilationConfig(inductor_passes={"a": func})`.""" |
| 604 | |
| 605 | # CudaGraph compilation |
| 606 | cudagraph_mode: CUDAGraphMode = None # type: ignore[assignment] |
| 607 | """ |
| 608 | The mode of the cudagraph: |
| 609 | |
| 610 | - NONE, no cudagraph capture. |
| 611 | - PIECEWISE. |
| 612 | - FULL. |
| 613 | - FULL_DECODE_ONLY. |
| 614 | - FULL_AND_PIECEWISE. (v1 default) |
| 615 | |
| 616 | PIECEWISE mode build piecewise cudagraph only, keeping the cudagraph |
| 617 | incompatible ops (i.e. some attention ops) outside the cudagraph |
| 618 | for general flexibility. |
| 619 | |
| 620 | FULL mode: Capture full cudagraph for all batches. Can be good for small |
| 621 | models or workloads with small prompts; not supported by many backends. |
| 622 | Generally for performance FULL_AND_PIECEWISE is better. |
| 623 | |
| 624 | FULL_DECODE_ONLY mode: Capture full cudagraph for decode batches only. |
| 625 | Mixed prefill-decode batches are run without cudagraphs. Can be good for |
| 626 | decode instances in a P/D setup where prefill is not as important so we |
| 627 | can save some memory. |
| 628 | |
| 629 | FULL_AND_PIECEWISE mode: Capture full cudagraph for decode batches and |
| 630 | piecewise cudagraph for prefill and mixed prefill-decode batches. |
| 631 | This is the most performant mode for most models and is the default. |
| 632 | |
| 633 | Currently, the cudagraph mode is only used for the v1 engine. |
| 634 | Note that the cudagraph logic is generally orthogonal to the |
| 635 | compilation logic. While piecewise cudagraphs require piecewise |
| 636 | compilation (mode=VLLM_COMPILE and non-empty splitting_ops), full |
| 637 | cudagraphs are supported with and without compilation. |
| 638 | |
| 639 | Warning: This flag is new and subject to change in addition |
| 640 | more modes may be added. |
| 641 | """ |
| 642 | cudagraph_num_of_warmups: int = 0 |
| 643 | """Number of warmup runs for cudagraph. |
| 644 | It means the first several runs will be treated as warmup runs. |
| 645 | Only after that, the execution will be recorded, and the recorded |
| 646 | cudagraph will be used for subsequent runs.""" |
| 647 | cudagraph_capture_sizes: list[int] = None # type: ignore[assignment] |
| 648 | """Sizes to capture cudagraph. |
| 649 | - None (default): capture sizes are inferred from vllm config. |
| 650 | - list[int]: capture sizes are specified as given.""" |
| 651 | cudagraph_copy_inputs: bool = False |
| 652 | """Whether to copy input tensors for |
| 653 | cudagraph. If the caller can guarantee that the same input buffers |
| 654 | are always used, it can set this to False. Otherwise, it should |
| 655 | set this to True, and the compiler will copy the input to an |
| 656 | internally managed buffer. Default is False. |
| 657 | Note that this flag is only effective when cudagraph_mode is PIECEWISE. |
| 658 | """ |
| 659 | cudagraph_specialize_lora: bool = True |
| 660 | """Whether to create separate cuda graphs for cases with and without active |
| 661 | LoRA adapters. When set to False, the LoRA-enabled cuda graph will be used |
| 662 | for all cases, incurring the overhead of running LoRA ops even when no |
| 663 | adapters are active. Setting this to True will remove this overhead at the |
| 664 | cost of increased startup time and slightly higher memory usage. |
| 665 | When `enable_lora` is False, this option has no effect. |
| 666 | """ |
| 667 | |
| 668 | use_inductor_graph_partition: bool = None # type: ignore[assignment] |
| 669 | """Use inductor graph partition to split the graph at cudagraph_unsafe ops. |
| 670 | This partition happens at inductor codegen time after all passes and fusions |
| 671 | are finished. It generates a single `call` function which wraps |
| 672 | cudagraph-safe ops into partition functions and leave cudagraph-unsafe ops |
| 673 | outside the partition functions. For a graph with N cudagraph-unsafe ops |
| 674 | (e.g., Attention), there would be N+1 partitions. To mark an op as |
| 675 | cudagraph unsafe, we can add `tags=(torch._C.Tag.cudagraph_unsafe)` when |
| 676 | register the custom op. |
| 677 | |
| 678 | This config supports both full cudagraph and piecewise cudagraph without |
| 679 | compiling twice. For piecewise cudagraph, it applies vLLM CUDAGraph wrapper |
| 680 | to each partition. For N+1 partitions, there would be N+1 |
| 681 | CUDAGraph wrapper instances. |
| 682 | |
| 683 | For full CUDAGraph, we always apply a single CUDAGraph wrapper outside the |
| 684 | inductor `call` function in the model runner. The top-level full cudagraph |
| 685 | capture ignores all partitioning. |
| 686 | """ |
| 687 | |
| 688 | pass_config: PassConfig = field(default_factory=PassConfig) |
| 689 | """Custom inductor passes, see PassConfig for more details""" |
| 690 | |
| 691 | max_cudagraph_capture_size: int = None # type: ignore[assignment] |
| 692 | """The maximum cudagraph capture size. |
| 693 | |
| 694 | If cudagraph_capture_sizes is specified, this will be set to the largest |
| 695 | size in that list (or checked for consistency if specified). If |
| 696 | cudagraph_capture_sizes is not specified, the list of sizes is generated |
| 697 | automatically following the pattern: |
| 698 | |
| 699 | [1, 2, 4] + list(range(8, 256, 8)) + list( |
| 700 | range(256, max_cudagraph_capture_size + 1, 16)) |
| 701 | |
| 702 | If not specified, max_cudagraph_capture_size is capped at 512 by default, |
| 703 | or 1024 on data center Blackwell GPUs. This avoids OOM in tight memory |
| 704 | scenarios with small max_num_seqs, and limits capture of large graphs that |
| 705 | increase startup time and memory usage. Uniform decode sizes are appended |
| 706 | only within this default ceiling. |
| 707 | """ |
| 708 | |
| 709 | dynamic_shapes_config: DynamicShapesConfig = field( |
| 710 | default_factory=DynamicShapesConfig |
| 711 | ) |
| 712 | """Configuration for dynamic shapes options""" |
| 713 | |
| 714 | local_cache_dir: str = field(default=None, init=False) # type: ignore |
| 715 | """local cache dir for each rank""" |
| 716 | |
| 717 | fast_moe_cold_start: bool | None = None |
| 718 | """Optimization for fast MOE cold start. |
| 719 | |
| 720 | This is a bit of a hack that assumes that: |
| 721 | 1. the only decoder forward pass being run is the current model |
| 722 | 2. the decoder forward pass runs all of the MOEs in the order in which they |
| 723 | are initialized |
| 724 | |
| 725 | When the above two conditions hold, this option greatly decreases cold start |
| 726 | time for MOE models. |
| 727 | |
| 728 | The options are: |
| 729 | - True: optimization is always on |
| 730 | - False: optimization is always off |
| 731 | - None: optimization is on usually but off for speculative decoding |
| 732 | |
| 733 | If conditions 1&2 don't hold then this option will lead to silent |
| 734 | incorrectness. |
| 735 | The only condition in which this doesn't hold is speculative |
| 736 | decoding, where there is a draft model that may have MOEs in them. |
| 737 | |
| 738 | NB: We're working on a longer-term solution that doesn't need these assumptions. |
| 739 | """ |
| 740 | |
| 741 | # keep track of enabled and disabled custom ops |
| 742 | enabled_custom_ops: Counter[str] = field(default_factory=Counter, init=False) |
| 743 | """custom ops that are enabled""" |
| 744 | disabled_custom_ops: Counter[str] = field(default_factory=Counter, init=False) |
| 745 | """custom ops that are disabled""" |
| 746 | traced_files: set[str] = field(default_factory=set, init=False) |
| 747 | """files that are traced for compilation""" |
| 748 | compilation_time: float = field(default=0.0, init=False) |
| 749 | """time taken for compilation""" |
| 750 | encoder_compilation_time: float = field(default=0.0, init=False) |
| 751 | """time taken for multimodal encoder compilation""" |
| 752 | |
| 753 | static_forward_context: dict[str, Any] = field(default_factory=dict, init=False) |
| 754 | """Per-model forward context |
| 755 | Map from layer name to layer objects that need to be accessed outside |
| 756 | model code, e.g., Attention, FusedMOE when dp_size>1.""" |
| 757 | |
| 758 | static_all_moe_layers: list[str] = field(default_factory=list, init=False) |
| 759 | """The names of all the MOE layers in the model |
| 760 | """ |
| 761 | |
| 762 | # Attention ops; used for piecewise cudagraphs |
| 763 | # Use PyTorch operator format: "namespace::name" |
| 764 | _attention_ops: ClassVar[list[str]] = [ |
| 765 | "vllm::unified_attention_with_output", |
| 766 | "vllm::unified_mla_attention_with_output", |
| 767 | "vllm::mamba_mixer2", |
| 768 | "vllm::mamba_mixer", |
| 769 | "vllm::short_conv", |
| 770 | # Qwen4Exp's AMD backend still uses these splitting ops. |
| 771 | "vllm::qwen4_exp_ple_short_conv", |
| 772 | "vllm::qwen4_exp_qsa_with_output", |
| 773 | "vllm::linear_attention", |
| 774 | "vllm::qwen_gdn_attention_core", |
| 775 | "vllm::qwen_gdn_attention_core_fused_norm_packed", |
| 776 | "vllm::gdn_attention_core_xpu", |
| 777 | "vllm::olmo_hybrid_gdn_full_forward", |
| 778 | "vllm::sparse_attn_indexer", |
| 779 | "vllm::rocm_aiter_sparse_attn_indexer", |
| 780 | "vllm::deepseek_v4_attention", |
| 781 | "vllm::hpc_rope_norm_forward", |
| 782 | ] |
| 783 | |
| 784 | def compute_hash(self) -> str: |
| 785 | """Provide a hash that uniquely identifies all the configs |
| 786 | that affect the structure of the computation |
| 787 | graph from input ids/embeddings to the final hidden states, |
| 788 | excluding anything before input ids/embeddings and after |
| 789 | the final hidden states. |
| 790 | """ |
| 791 | # Opt-out: default-include declared fields; keep a tiny exclude set; |
| 792 | # normalize types; keep SHA-256. For nested opaque configs, include a |
| 793 | # stable identifier (e.g., pass_config.compute_hash()) instead of object id. |
| 794 | |
| 795 | ignored_factors = { |
| 796 | # Paths/dirs and runtime/metrics that don’t affect compiled graph |
| 797 | "debug_dump_path", |
| 798 | "cache_dir", |
| 799 | "local_cache_dir", |
| 800 | "traced_files", |
| 801 | "compilation_time", |
| 802 | "encoder_compilation_time", |
| 803 | "enabled_custom_ops", |
| 804 | "disabled_custom_ops", |
| 805 | "static_forward_context", |
| 806 | "pass_config", # handled separately below |
| 807 | "dynamic_shapes_config", # handled separately below |
| 808 | } |
| 809 | |
| 810 | from vllm.config.utils import get_hash_factors, hash_factors |
| 811 | |
| 812 | factors = get_hash_factors(self, ignored_factors) |
| 813 | |
| 814 | factors["pass_config"] = self.pass_config.compute_hash() |
| 815 | factors["dynamic_shapes_config"] = self.dynamic_shapes_config.compute_hash() |
| 816 | return hash_factors(factors) |
| 817 | |
| 818 | def __repr__(self) -> str: |
| 819 | exclude: dict[str, bool | dict[str, bool]] = { |
| 820 | "static_forward_context": True, |
| 821 | "enabled_custom_ops": True, |
| 822 | "disabled_custom_ops": True, |
| 823 | "compilation_time": True, |
| 824 | "encoder_compilation_time": True, |
| 825 | "traced_files": True, |
| 826 | "inductor_compile_config": { |
| 827 | "post_grad_custom_post_pass": True, |
| 828 | }, |
| 829 | } |
| 830 | |
| 831 | # exclude default attr in pass_config |
| 832 | pass_config_exclude = {} |
| 833 | for attr, default_val in vars(PassConfig()).items(): |
| 834 | if getattr(self.pass_config, attr) == default_val: |
| 835 | pass_config_exclude[attr] = True |
| 836 | if pass_config_exclude: |
| 837 | exclude["pass_config"] = pass_config_exclude |
| 838 | |
| 839 | config = TypeAdapter(CompilationConfig).dump_python( |
| 840 | self, exclude=exclude, exclude_unset=True |
| 841 | ) |
| 842 | |
| 843 | return str(config) |
| 844 | |
| 845 | __str__ = __repr__ |
| 846 | |
| 847 | @field_validator("mode", mode="before") |
| 848 | @classmethod |
| 849 | def validate_mode_before(cls, value: Any) -> Any: |
| 850 | """Enable parsing the `mode` field from string mode names. |
| 851 | Accepts both integers (0-3) and string names, like NONE, STOCK_TORCH_COMPILE, |
| 852 | DYNAMO_TRACE_ONCE, VLLM_COMPILE. |
| 853 | """ |
| 854 | if isinstance(value, str): |
| 855 | # Convert string mode name to integer value |
| 856 | mode_name = value.upper() |
| 857 | |
| 858 | if mode_name not in CompilationMode.__members__: |
| 859 | raise ValueError( |
| 860 | f"Invalid compilation mode: {value}. " |
| 861 | f"Valid modes are: {', '.join(CompilationMode.__members__.keys())}" |
| 862 | ) |
| 863 | |
| 864 | return CompilationMode[mode_name] |
| 865 | return value |
| 866 | |
| 867 | @field_validator("cudagraph_mode", mode="before") |
| 868 | @classmethod |
| 869 | def validate_cudagraph_mode_before(cls, value: Any) -> Any: |
| 870 | """Enable parsing of the `cudagraph_mode` enum type from string.""" |
| 871 | if isinstance(value, str): |
| 872 | return CUDAGraphMode[value.upper()] |
| 873 | return value |
| 874 | |
| 875 | @field_validator("pass_config", mode="before") |
| 876 | @classmethod |
| 877 | def validate_pass_config_before(cls, value: Any) -> Any: |
| 878 | """Enable parsing of the `pass_config` field from a dictionary.""" |
| 879 | if isinstance(value, dict): |
| 880 | return PassConfig(**value) |
| 881 | return value |
| 882 | |
| 883 | @field_validator("compile_cache_save_format") |
| 884 | @classmethod |
| 885 | def validate_compile_cache_save_format(cls, value: str) -> str: |
| 886 | if value not in ("binary", "unpacked"): |
| 887 | raise ValueError( |
| 888 | f"compile_cache_save_format must be 'binary' or 'unpacked', " |
| 889 | f"got: {value}" |
| 890 | ) |
| 891 | return value |
| 892 | |
| 893 | @field_validator( |
| 894 | "level", |
| 895 | "mode", |
| 896 | "cudagraph_mode", |
| 897 | "max_cudagraph_capture_size", |
| 898 | "use_inductor_graph_partition", |
| 899 | "ir_enable_torch_wrap", |
| 900 | mode="wrap", |
| 901 | ) |
| 902 | @classmethod |
| 903 | def _skip_none_validation(cls, value: Any, handler: Callable) -> Any: |
| 904 | """Skip validation if the value is `None` when initialisation is delayed.""" |
| 905 | if value is None: |
| 906 | return value |
| 907 | return handler(value) |
| 908 | |
| 909 | def __post_init__(self) -> None: |
| 910 | # TODO(zou3519/luka): There are 2 issues with auto-functionalization V2: |
| 911 | # 1. A bug in PyTorch, fixed in 2.7: |
| 912 | # https://github.com/pytorch/pytorch/issues/147924 |
| 913 | # 2. Custom passes (fusion) rely on auto-functionalization V1 and don't |
| 914 | # work with V2. Addressing this will take extra engineering effort |
| 915 | # and it is not yet a priority. RFC here: |
| 916 | # https://github.com/vllm-project/vllm/issues/14703 |
| 917 | |
| 918 | KEY = "enable_auto_functionalized_v2" |
| 919 | if KEY not in self.inductor_compile_config: |
| 920 | self.inductor_compile_config[KEY] = False |
| 921 | |
| 922 | # Tie inductor runtime assertions to debug logging mode. |
| 923 | # These assertions add ~2ms overhead per forward pass on large |
| 924 | # models (e.g., DeepSeek-R1 671B: ~340 assert_size_stride + ~60 |
| 925 | # assert_alignment calls per forward). PyTorch >= 2.12 has a |
| 926 | # native fix (assert-once), so we only apply this workaround on |
| 927 | # older versions. On torch < 2.12, enable asserts only when |
| 928 | # VLLM_LOGGING_LEVEL=DEBUG. Users can still override explicitly |
| 929 | # via --compilation-config '{"inductor_compile_config": |
| 930 | # {"size_asserts": true, ...}}'. |
| 931 | # See: https://github.com/pytorch/pytorch/issues/177719 |
| 932 | if not is_torch_equal_or_newer("2.12.0.dev"): |
| 933 | enable_asserts = envs.VLLM_LOGGING_LEVEL == "DEBUG" |
| 934 | for key in ( |
| 935 | "size_asserts", |
| 936 | "alignment_asserts", |
| 937 | "scalar_asserts", |
| 938 | ): |
| 939 | self.inductor_compile_config.setdefault(key, enable_asserts) |
| 940 | |
| 941 | for k, v in self.inductor_passes.items(): |
| 942 | if not isinstance(v, str): |
| 943 | assert callable(v), f"pass {k} should be callable or a qualified name" |
| 944 | self.inductor_compile_config[k] = ( |
| 945 | v if isinstance(v, InductorPass) else CallableInductorPass(v) |
| 946 | ) |
| 947 | continue |
| 948 | |
| 949 | # resolve function from qualified name |
| 950 | names = v.split(".") |
| 951 | module = ".".join(names[:-1]) |
| 952 | func_name = names[-1] |
| 953 | func = __import__(module).__dict__[func_name] |
| 954 | self.inductor_compile_config[k] = ( |
| 955 | func if isinstance(func, InductorPass) else CallableInductorPass(func) |
| 956 | ) |
| 957 | |
| 958 | if ( |
| 959 | self.pass_config.enable_qk_norm_rope_fusion |
| 960 | and "+rotary_embedding" not in self.custom_ops |
| 961 | ): |
| 962 | # TODO(zhuhaoran): support rope native forward match and remove this. |
| 963 | # Linked issue: https://github.com/vllm-project/vllm/issues/28042 |
| 964 | self.custom_ops.append("+rotary_embedding") |
| 965 | |
| 966 | if ( |
| 967 | self.pass_config.fuse_rope_kvcache |
| 968 | and "+rotary_embedding" not in self.custom_ops |
| 969 | ): |
| 970 | # TODO(Rohan138): support rope native forward match and remove this. |
| 971 | # Linked issue: https://github.com/vllm-project/vllm/issues/28042 |
| 972 | self.custom_ops.append("+rotary_embedding") |
| 973 | |
| 974 | if ( |
| 975 | self.pass_config.fuse_qk_norm_rope_kvcache |
| 976 | and "+rotary_embedding" not in self.custom_ops |
| 977 | ): |
| 978 | self.custom_ops.append("+rotary_embedding") |
| 979 | |
| 980 | if ( |
| 981 | is_torch_equal_or_newer("2.9.0.dev") |
| 982 | and "combo_kernels" not in self.inductor_compile_config |
| 983 | and "benchmark_combo_kernel" not in self.inductor_compile_config |
| 984 | # (fixme @boyuan) combo kernel does not support cpu yet. |
| 985 | and not current_platform.is_cpu() |
| 986 | ): |
| 987 | # use horizontal fusion, which is useful for fusing qk-norm and |
| 988 | # qk-rope when query and key have different shapes. |
| 989 | self.inductor_compile_config["combo_kernels"] = True |
| 990 | self.inductor_compile_config["benchmark_combo_kernel"] = True |
| 991 | |
| 992 | if self.use_inductor_graph_partition and not is_torch_equal_or_newer( |
| 993 | "2.9.0.dev" |
| 994 | ): |
| 995 | raise ValueError( |
| 996 | "use_inductor_graph_partition is only " |
| 997 | "supported with torch>=2.9.0.dev. Set " |
| 998 | "use_inductor_graph_partition=False instead." |
| 999 | ) |
| 1000 | |
| 1001 | for op in self.custom_ops: |
| 1002 | if op not in {"all", "none"} and (len(op) < 2 or op[0] not in {"+", "-"}): |
| 1003 | raise ValueError( |
| 1004 | f"Invalid syntax '{op}' for custom op, " |
| 1005 | "must be 'all', 'none', '+op' or '-op' " |
| 1006 | "(where 'op' is the registered op name)" |
| 1007 | ) |
| 1008 | |
| 1009 | base_modes = [op for op in self.custom_ops if op in {"all", "none"}] |
| 1010 | if len(base_modes) > 1: |
| 1011 | raise ValueError( |
| 1012 | "custom_ops can contain only one base mode: 'all' or 'none'" |
| 1013 | ) |
| 1014 | |
| 1015 | enabled_ops = {op[1:] for op in self.custom_ops if op.startswith("+")} |
| 1016 | disabled_ops = {op[1:] for op in self.custom_ops if op.startswith("-")} |
| 1017 | conflicting_ops = sorted(enabled_ops & disabled_ops) |
| 1018 | if conflicting_ops: |
| 1019 | raise ValueError( |
| 1020 | "custom_ops cannot both enable and disable the same operation(s): " |
| 1021 | f"{', '.join(conflicting_ops)}. Remove either the '+' or '-' directive" |
| 1022 | ) |
| 1023 | |
| 1024 | # Currently only eager and inductor backend are supported. |
| 1025 | # for piecewise compilation. Custom backends are not supported for |
| 1026 | # piecewise compilation. Update when more backends are supported. |
| 1027 | if self.mode == CompilationMode.VLLM_COMPILE and self.backend not in [ |
| 1028 | "", |
| 1029 | "eager", |
| 1030 | "inductor", |
| 1031 | ]: |
| 1032 | raise ValueError( |
| 1033 | f"Invalid backend for piecewise compilation: {self.backend}" |
| 1034 | ) |
| 1035 | |
| 1036 | # Validate encoder CUDA graph configuration |
| 1037 | if ( |
| 1038 | self.cudagraph_mm_encoder |
| 1039 | and self.encoder_cudagraph_max_vision_items_per_batch < 0 |
| 1040 | ): |
| 1041 | raise ValueError( |
| 1042 | "encoder_cudagraph_max_vision_items_per_batch must be " |
| 1043 | "non-negative (0 = auto-infer)" |
| 1044 | ) |
| 1045 | if ( |
| 1046 | self.cudagraph_mm_encoder |
| 1047 | and self.encoder_cudagraph_max_frames_per_batch is not None |
| 1048 | and self.encoder_cudagraph_max_frames_per_batch < 0 |
| 1049 | ): |
| 1050 | raise ValueError( |
| 1051 | "encoder_cudagraph_max_frames_per_batch must be " |
| 1052 | "non-negative (None = auto-infer)" |
| 1053 | ) |
| 1054 | |
| 1055 | if self.encoder_cudagraph_token_budgets and any( |
| 1056 | b <= 0 for b in self.encoder_cudagraph_token_budgets |
| 1057 | ): |
| 1058 | raise ValueError( |
| 1059 | f"All encoder_cudagraph_token_budgets must be positive, " |
| 1060 | f"got {self.encoder_cudagraph_token_budgets}" |
| 1061 | ) |
| 1062 | |
| 1063 | if self.backend == "": |
| 1064 | self.backend = current_platform.get_compile_backend() |
| 1065 | |
| 1066 | def init_backend( |
| 1067 | self, |
| 1068 | vllm_config: "VllmConfig", |
| 1069 | prefix: str = "", |
| 1070 | is_encoder: bool = False, |
| 1071 | ) -> str | Callable: |
| 1072 | """Initialize the backend for the compilation config from a vllm config. |
| 1073 | |
| 1074 | Arguments: |
| 1075 | vllm_config: The vllm config to initialize the backend from. |
| 1076 | prefix: Cache directory prefix for this compiled module. |
| 1077 | is_encoder: Whether this module is used in an encoder (as |
| 1078 | opposed to a text backbone). |
| 1079 | |
| 1080 | Returns: |
| 1081 | The backend for the compilation config. |
| 1082 | |
| 1083 | """ |
| 1084 | if self.mode is None: |
| 1085 | raise ValueError( |
| 1086 | "No compilation mode is set. This method should only be " |
| 1087 | "called via vllm config where the level is set if none is " |
| 1088 | "provided." |
| 1089 | ) |
| 1090 | if self.mode == CompilationMode.NONE: |
| 1091 | raise ValueError("No compilation mode is set.") |
| 1092 | |
| 1093 | from torch._dynamo.backends.registry import list_backends |
| 1094 | |
| 1095 | torch_backends = list_backends(exclude_tags=tuple()) |
| 1096 | if self.mode in [ |
| 1097 | CompilationMode.STOCK_TORCH_COMPILE, |
| 1098 | CompilationMode.DYNAMO_TRACE_ONCE, |
| 1099 | ]: |
| 1100 | if self.backend in torch_backends: |
| 1101 | return self.backend |
| 1102 | return resolve_obj_by_qualname(self.backend) |
| 1103 | |
| 1104 | assert self.mode == CompilationMode.VLLM_COMPILE |
| 1105 | if self.backend not in ["eager", "inductor"]: |
| 1106 | logger.info("Using OOT custom backend for compilation.") |
| 1107 | |
| 1108 | from vllm.compilation.backends import VllmBackend |
| 1109 | |
| 1110 | return VllmBackend(vllm_config, prefix=prefix, is_encoder=is_encoder) |
| 1111 | |
| 1112 | def post_init_cudagraph_sizes(self) -> None: |
| 1113 | """To complete the initialization after cudagraph related |
| 1114 | configs are set. This includes: |
| 1115 | - initialize compile_sizes |
| 1116 | """ |
| 1117 | computed_compile_sizes: list[int] = [] |
| 1118 | if self.compile_sizes is not None: |
| 1119 | # de-duplicate the sizes provided by the config |
| 1120 | self.compile_sizes = list(set(self.compile_sizes)) |
| 1121 | for x in self.compile_sizes: |
| 1122 | if isinstance(x, str): |
| 1123 | assert x == "cudagraph_capture_sizes", ( |
| 1124 | "Unrecognized size type in compile_sizes, " |
| 1125 | f"expect 'cudagraph_capture_sizes', got {x}" |
| 1126 | ) |
| 1127 | computed_compile_sizes.extend(self.cudagraph_capture_sizes) |
| 1128 | else: |
| 1129 | assert isinstance(x, int) |
| 1130 | computed_compile_sizes.append(x) |
| 1131 | self.compile_sizes = computed_compile_sizes # type: ignore |
| 1132 | |
| 1133 | # make sure the sizes are in ascending order |
| 1134 | self.cudagraph_capture_sizes.sort() |
| 1135 | if self.cudagraph_capture_sizes: |
| 1136 | assert self.cudagraph_capture_sizes[-1] == self.max_cudagraph_capture_size |
| 1137 | |
| 1138 | def set_splitting_ops_for_v1( |
| 1139 | self, all2all_backend: str, data_parallel_size: int = 1 |
| 1140 | ): |
| 1141 | # To compatible with OOT hardware plugin platform (for example vllm-ascend) |
| 1142 | # which currently only supports sequence parallelism in eager mode. |
| 1143 | if self.mode != CompilationMode.VLLM_COMPILE: |
| 1144 | if self.splitting_ops is None: |
| 1145 | self.splitting_ops = [] |
| 1146 | return |
| 1147 | |
| 1148 | if self.pass_config.fuse_attn_quant and not self.use_inductor_graph_partition: |
| 1149 | self.set_splitting_ops_for_attn_fusion() |
| 1150 | else: |
| 1151 | if self.splitting_ops is None: |
| 1152 | # NOTE: When using full cudagraph, instead of setting an empty |
| 1153 | # list and capture the full cudagraph inside the flattened fx |
| 1154 | # graph, we keep the piecewise fx graph structure but capture |
| 1155 | # the full cudagraph outside the fx graph. This reduces some |
| 1156 | # cpu overhead when the runtime batch_size is not cudagraph |
| 1157 | # captured. see https://github.com/vllm-project/vllm/pull/20059 |
| 1158 | # for details. Make a copy to avoid mutating the class-level |
| 1159 | # list via reference. |
| 1160 | self.splitting_ops = list(self._attention_ops) |
| 1161 | |
| 1162 | # unified_kv_cache_update has a string param that prevents Inductor |
| 1163 | # from reusing piecewise graphs. Remove it from the compiled graph. |
| 1164 | # This has the side-effect of excluding cache from cudagraphs but |
| 1165 | # that doesn't seem to affect performance. |
| 1166 | # https://github.com/vllm-project/vllm/issues/33267 |
| 1167 | if not self.use_inductor_graph_partition: |
| 1168 | if self.pass_config.fuse_rope_kvcache: |
| 1169 | logger.warning_once( |
| 1170 | "fuse_rope_kvcache is enabled, but splitting_ops is None " |
| 1171 | "and Inductor graph partition is not enabled." |
| 1172 | "Disabling fuse_rope_kvcache." |
| 1173 | "Please either set splitting_ops to an empty list []" |
| 1174 | "or set use_inductor_graph_partition to True " |
| 1175 | "to enable RoPE+KV cache fusion." |
| 1176 | ) |
| 1177 | self.pass_config.fuse_rope_kvcache = False |
| 1178 | if self.pass_config.fuse_qk_norm_rope_kvcache: |
| 1179 | logger.warning_once( |
| 1180 | "fuse_qk_norm_rope_kvcache is enabled, but " |
| 1181 | "splitting_ops is None and Inductor graph partition " |
| 1182 | "is not enabled. Disabling fuse_qk_norm_rope_kvcache. " |
| 1183 | "Please either set splitting_ops to an empty list [] " |
| 1184 | "or set use_inductor_graph_partition to True " |
| 1185 | "to enable QK-Norm+RoPE+KV cache fusion." |
| 1186 | ) |
| 1187 | self.pass_config.fuse_qk_norm_rope_kvcache = False |
| 1188 | self.splitting_ops.append("vllm::unified_kv_cache_update") |
| 1189 | self.splitting_ops.append("vllm::unified_mla_kv_cache_update") |
| 1190 | |
| 1191 | elif len(self.splitting_ops) == 0: |
| 1192 | if ( |
| 1193 | self.cudagraph_mode == CUDAGraphMode.PIECEWISE |
| 1194 | or self.cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE |
| 1195 | ): |
| 1196 | logger.warning_once( |
| 1197 | "Using piecewise cudagraph with empty splitting_ops" |
| 1198 | ) |
| 1199 | if self.cudagraph_mode == CUDAGraphMode.PIECEWISE: |
| 1200 | logger.warning_once( |
| 1201 | "Piecewise compilation with empty splitting_ops does not " |
| 1202 | "contain piecewise cudagraph. Setting cudagraph_" |
| 1203 | "mode to NONE. Hint: If you are using attention " |
| 1204 | "backends that support cudagraph, consider manually " |
| 1205 | "setting cudagraph_mode to FULL or FULL_DECODE_ONLY " |
| 1206 | "to enable full cudagraphs." |
| 1207 | ) |
| 1208 | self.cudagraph_mode = CUDAGraphMode.NONE |
| 1209 | elif self.cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE: |
| 1210 | logger.warning_once( |
| 1211 | "Piecewise compilation with empty splitting_ops does " |
| 1212 | "not contain piecewise cudagraph. Setting " |
| 1213 | "cudagraph_mode to FULL." |
| 1214 | ) |
| 1215 | self.cudagraph_mode = CUDAGraphMode.FULL |
| 1216 | self.splitting_ops = [] |
| 1217 | |
| 1218 | if ( |
| 1219 | not self.use_inductor_graph_partition |
| 1220 | and (self.pass_config.enable_sp or self.pass_config.fuse_gemm_comms) |
| 1221 | and self.splitting_ops |
| 1222 | ): |
| 1223 | logger.warning_once( |
| 1224 | "Sequence parallelism requires full-graph compilation when " |
| 1225 | "use_inductor_graph_partition is off. Setting splitting_ops " |
| 1226 | "to an empty list to preserve SP and async TP." |
| 1227 | ) |
| 1228 | self.splitting_ops = [] |
| 1229 | if self.cudagraph_mode.has_piecewise_cudagraphs(): |
| 1230 | logger.warning_once( |
| 1231 | "Sequence parallelism is incompatible with piecewise " |
| 1232 | "cudagraph when use_inductor_graph_partition is off. " |
| 1233 | "Setting cudagraph_mode to FULL." |
| 1234 | ) |
| 1235 | self.cudagraph_mode = CUDAGraphMode.FULL |
| 1236 | |
| 1237 | # Disable CUDA graphs for DeepEP high-throughput since its not CG compatible |
| 1238 | if ( |
| 1239 | all2all_backend == "deepep_high_throughput" |
| 1240 | and data_parallel_size > 1 |
| 1241 | and self.cudagraph_mode != CUDAGraphMode.NONE |
| 1242 | ): |
| 1243 | # TODO: Piecewise Cuda graph might be enabled |
| 1244 | # if torch compile cache key issue fixed |
| 1245 | # See https://github.com/vllm-project/vllm/pull/25093 |
| 1246 | logger.info( |
| 1247 | "DeepEP: Disabling CUDA Graphs since DeepEP high-throughput kernels " |
| 1248 | "are optimized for prefill and are incompatible with CUDA Graphs. " |
| 1249 | "In order to use CUDA Graphs for decode-optimized workloads, " |
| 1250 | "use --all2all-backend with another option, such as " |
| 1251 | "deepep_low_latency, nixl_ep, or allgather_reducescatter." |
| 1252 | ) |
| 1253 | self.cudagraph_mode = CUDAGraphMode.NONE |
| 1254 | |
| 1255 | def set_splitting_ops_for_attn_fusion(self): |
| 1256 | assert self.pass_config.fuse_attn_quant |
| 1257 | if self.splitting_ops is None: |
| 1258 | self.splitting_ops = [] |
| 1259 | if self.cudagraph_mode.has_piecewise_cudagraphs(): |
| 1260 | logger.warning_once( |
| 1261 | "fuse_attn_quant is incompatible with piecewise " |
| 1262 | "cudagraph when use_inductor_graph_partition is off. " |
| 1263 | "In this case, splitting_ops will be set to empty " |
| 1264 | "list, and cudagraph_mode will be set to FULL. " |
| 1265 | "Please ensure you are using attention backends that " |
| 1266 | "support cudagraph or set cudagraph_mode to NONE " |
| 1267 | "explicitly if encountering any problems." |
| 1268 | ) |
| 1269 | self.cudagraph_mode = CUDAGraphMode.FULL |
| 1270 | |
| 1271 | assert not self.splitting_ops_contain_attention(), ( |
| 1272 | "attention ops should not be in splitting_ops when fuse_attn_quant is True" |
| 1273 | ) |
| 1274 | |
| 1275 | def splitting_ops_contain_attention(self) -> bool: |
| 1276 | return self.splitting_ops is not None and all( |
| 1277 | op in self.splitting_ops for op in self._attention_ops |
| 1278 | ) |
| 1279 | |
| 1280 | def splitting_ops_contain_kv_cache_update(self) -> bool: |
| 1281 | # when using Dynamo partition while splitting ops is None |
| 1282 | # and attn+quant fusion disabled, the kv_cache_update_ops are |
| 1283 | # appended to splitting_ops in set_splitting_ops_for_v1 due to |
| 1284 | # https://github.com/vllm-project/vllm/issues/33267 |
| 1285 | # In this case, we return True if the kv_cache_update_ops |
| 1286 | # are not in the splitting_ops yet, but will subsequently |
| 1287 | # be added to splitting_ops. |
| 1288 | if ( |
| 1289 | not self.use_inductor_graph_partition |
| 1290 | and self.splitting_ops is None |
| 1291 | and not self.pass_config.fuse_attn_quant |
| 1292 | ): |
| 1293 | return True |
| 1294 | |
| 1295 | kv_cache_update_ops = [ |
| 1296 | "vllm::unified_kv_cache_update", |
| 1297 | "vllm::unified_mla_kv_cache_update", |
| 1298 | ] |
| 1299 | return self.splitting_ops is not None and all( |
| 1300 | op in self.splitting_ops for op in kv_cache_update_ops |
| 1301 | ) |
| 1302 | |
| 1303 | def is_attention_compiled_piecewise(self) -> bool: |
| 1304 | if not self.splitting_ops_contain_attention(): |
| 1305 | return False |
| 1306 | |
| 1307 | if not self.use_inductor_graph_partition: |
| 1308 | # Dynamo-level FX split case |
| 1309 | return self.mode == CompilationMode.VLLM_COMPILE |
| 1310 | |
| 1311 | # Inductor partition case |
| 1312 | return self.backend == "inductor" and self.mode != CompilationMode.NONE |
| 1313 | |
| 1314 | def custom_op_log_check(self): |
| 1315 | """This method logs the enabled/disabled custom ops and checks that the |
| 1316 | passed custom_ops field only contains relevant ops. |
| 1317 | It is called at the end of set_current_vllm_config, |
| 1318 | after the custom ops have been instantiated. |
| 1319 | """ |
| 1320 | if len(self.enabled_custom_ops) + len(self.disabled_custom_ops) == 0: |
| 1321 | logger.debug("No custom ops found in model.") |
| 1322 | return |
| 1323 | |
| 1324 | logger.debug("enabled custom ops: %s", self.enabled_custom_ops) |
| 1325 | logger.debug("disabled custom ops: %s", self.disabled_custom_ops) |
| 1326 | |
| 1327 | all_ops_in_model = self.enabled_custom_ops | self.disabled_custom_ops |
| 1328 | for op in self.custom_ops: |
| 1329 | if op in {"all", "none"}: |
| 1330 | continue |
| 1331 | |
| 1332 | assert op[0] in {"+", "-"}, ( |
| 1333 | "Invalid custom op syntax (should be checked during init)" |
| 1334 | ) |
| 1335 | |
| 1336 | # check if op name exists in model |
| 1337 | op_name = op[1:] |
| 1338 | if op_name not in all_ops_in_model: |
| 1339 | from vllm.model_executor.custom_op import op_registry |
| 1340 | |
| 1341 | # Does op exist at all or is it just not present in this model? |
| 1342 | # Note: Only imported op classes appear in the registry. |
| 1343 | missing_str = ( |
| 1344 | "doesn't exist (or wasn't imported/registered)" |
| 1345 | if op_name not in op_registry |
| 1346 | else "not present in model" |
| 1347 | ) |
| 1348 | |
| 1349 | enable_str = "enabling" if op[0] == "+" else "disabling" |
| 1350 | logger.warning_once( |
| 1351 | "Op '%s' %s, %s with '%s' has no effect", |
| 1352 | op_name, |
| 1353 | missing_str, |
| 1354 | enable_str, |
| 1355 | op, |
| 1356 | ) |
| 1357 | |
| 1358 | def is_custom_op_enabled(self, op: str) -> bool: |
| 1359 | count_all = self.custom_ops.count("all") |
| 1360 | count_none = self.custom_ops.count("none") |
| 1361 | if count_all + count_none != 1: |
| 1362 | raise ValueError( |
| 1363 | "custom_ops must contain exactly one base mode: 'all' or 'none'" |
| 1364 | ) |
| 1365 | |
| 1366 | if count_all: |
| 1367 | return f"-{op}" not in self.custom_ops |
| 1368 | |
| 1369 | return f"+{op}" in self.custom_ops |
| 1370 | |
| 1371 | def resolve_cudagraph_mode_and_sizes( |
| 1372 | self, |
| 1373 | min_cg_support: "AttentionCGSupport", |
| 1374 | min_cg_attn_backend: str | None, |
| 1375 | uniform_decode_query_len: int = 1, |
| 1376 | use_v2_model_runner: bool = False, |
| 1377 | tensor_parallel_size: int = 1, |
| 1378 | kv_cache_config: "KVCacheConfig | None" = None, |
| 1379 | max_num_reqs: int | None = None, |
| 1380 | is_profiling: bool = False, |
| 1381 | piecewise_capture_available: bool = True, |
| 1382 | ) -> CUDAGraphMode: |
| 1383 | from vllm.v1.attention.backend import AttentionCGSupport |
| 1384 | |
| 1385 | cudagraph_mode = self.cudagraph_mode |
| 1386 | if cudagraph_mode is None or cudagraph_mode == CUDAGraphMode.NONE: |
| 1387 | self.cudagraph_mode = CUDAGraphMode.NONE |
| 1388 | return CUDAGraphMode.NONE |
| 1389 | |
| 1390 | # Check cudagraph for mixed batch is supported |
| 1391 | if ( |
| 1392 | cudagraph_mode.mixed_mode() == CUDAGraphMode.FULL |
| 1393 | and min_cg_support != AttentionCGSupport.ALWAYS |
| 1394 | ): |
| 1395 | msg = ( |
| 1396 | f"CUDAGraphMode.{cudagraph_mode.name} is not supported " |
| 1397 | f"with {min_cg_attn_backend} backend (support: " |
| 1398 | f"{min_cg_support})" |
| 1399 | ) |
| 1400 | if min_cg_support == AttentionCGSupport.NEVER: |
| 1401 | # if not supported any full cudagraphs, just raise it. |
| 1402 | msg += ( |
| 1403 | "; please try cudagraph_mode=PIECEWISE, and " |
| 1404 | "make sure compilation mode is VLLM_COMPILE" |
| 1405 | ) |
| 1406 | raise ValueError(msg) |
| 1407 | |
| 1408 | # attempt to resolve the full cudagraph related mode |
| 1409 | if self.splitting_ops_contain_attention(): |
| 1410 | msg += "; setting cudagraph_mode=FULL_AND_PIECEWISE" |
| 1411 | cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE |
| 1412 | else: |
| 1413 | msg += "; setting cudagraph_mode=FULL_DECODE_ONLY" |
| 1414 | cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY |
| 1415 | logger.warning(msg) |
| 1416 | |
| 1417 | # check that if we are doing decode full-cudagraphs it is supported |
| 1418 | if ( |
| 1419 | cudagraph_mode.decode_mode() == CUDAGraphMode.FULL |
| 1420 | and min_cg_support == AttentionCGSupport.NEVER |
| 1421 | ): |
| 1422 | msg = ( |
| 1423 | f"CUDAGraphMode.{cudagraph_mode.name} is not supported " |
| 1424 | f"with {min_cg_attn_backend} backend (support: " |
| 1425 | f"{min_cg_support})" |
| 1426 | ) |
| 1427 | if self.mode == CompilationMode.VLLM_COMPILE and ( |
| 1428 | self.splitting_ops_contain_attention() |
| 1429 | or self.use_inductor_graph_partition |
| 1430 | ): |
| 1431 | msg += ( |
| 1432 | "; setting cudagraph_mode=PIECEWISE because " |
| 1433 | "attention is compiled piecewise" |
| 1434 | ) |
| 1435 | cudagraph_mode = CUDAGraphMode.PIECEWISE |
| 1436 | else: |
| 1437 | msg += ( |
| 1438 | "; setting cudagraph_mode=NONE because " |
| 1439 | "attention is not compiled piecewise" |
| 1440 | ) |
| 1441 | cudagraph_mode = CUDAGraphMode.NONE |
| 1442 | logger.warning(msg) |
| 1443 | |
| 1444 | # check that if we are doing spec-decode + decode full-cudagraphs it is |
| 1445 | # supported |
| 1446 | if ( |
| 1447 | cudagraph_mode.decode_mode() == CUDAGraphMode.FULL |
| 1448 | and uniform_decode_query_len > 1 |
| 1449 | and min_cg_support.value < AttentionCGSupport.UNIFORM_BATCH.value |
| 1450 | ): |
| 1451 | msg = ( |
| 1452 | f"CUDAGraphMode.{cudagraph_mode.name} is not supported" |
| 1453 | f" with spec-decode for attention backend " |
| 1454 | f"{min_cg_attn_backend} (support: {min_cg_support})" |
| 1455 | ) |
| 1456 | if self.splitting_ops_contain_attention(): |
| 1457 | msg += "; setting cudagraph_mode=PIECEWISE" |
| 1458 | cudagraph_mode = CUDAGraphMode.PIECEWISE |
| 1459 | else: |
| 1460 | msg += "; setting cudagraph_mode=NONE" |
| 1461 | cudagraph_mode = CUDAGraphMode.NONE |
| 1462 | logger.warning(msg) |
| 1463 | |
| 1464 | if ( |
| 1465 | not piecewise_capture_available |
| 1466 | and cudagraph_mode.requires_piecewise_compilation() |
| 1467 | ): |
| 1468 | fallback_mode = cudagraph_mode.without_piecewise() |
| 1469 | logger.warning_once( |
| 1470 | "Cudagraph mode %s requires piecewise capture, but the loaded " |
| 1471 | "model provides neither a compiled submodule nor breakable CUDA " |
| 1472 | "graphs. Overriding to %s.", |
| 1473 | cudagraph_mode, |
| 1474 | fallback_mode, |
| 1475 | ) |
| 1476 | cudagraph_mode = fallback_mode |
| 1477 | |
| 1478 | # double check that we can support full cudagraph if they are requested |
| 1479 | # even after automatic downgrades |
| 1480 | if ( |
| 1481 | cudagraph_mode.has_full_cudagraphs() |
| 1482 | and min_cg_support == AttentionCGSupport.NEVER |
| 1483 | ): |
| 1484 | raise ValueError( |
| 1485 | f"CUDAGraphMode.{cudagraph_mode.name} is not " |
| 1486 | f"supported with {min_cg_attn_backend} backend (" |
| 1487 | f"support:{min_cg_support}) " |
| 1488 | "; please try cudagraph_mode=PIECEWISE, " |
| 1489 | "and make sure compilation mode is VLLM_COMPILE" |
| 1490 | ) |
| 1491 | |
| 1492 | # MRV1 adjusts cudagraph sizes to be a multiple of uniform_decode_query_len |
| 1493 | # to avoid: https://github.com/vllm-project/vllm/issues/28207 and temp-fix: |
| 1494 | # https://github.com/vllm-project/vllm/issues/28207#issuecomment-3504004536 |
| 1495 | # Will be removed in the near future when we have separate cudagraph capture |
| 1496 | # sizes for decode and mixed prefill-decode. |
| 1497 | # MRV2 handles cudagraph capture sizing in cudagraph_utils.py |
| 1498 | # and doesn't need below: https://github.com/vllm-project/vllm/pull/45953 |
| 1499 | if ( |
| 1500 | not use_v2_model_runner |
| 1501 | and cudagraph_mode.decode_mode() == CUDAGraphMode.FULL |
| 1502 | and uniform_decode_query_len > 1 |
| 1503 | ): |
| 1504 | self.adjust_cudagraph_sizes_for_spec_decode( |
| 1505 | uniform_decode_query_len, |
| 1506 | tensor_parallel_size, |
| 1507 | ) |
| 1508 | |
| 1509 | # For Mamba models with FULL decode cudagraphs, each decode |
| 1510 | # sequence needs one Mamba cache block. The decode cudagraph |
| 1511 | # dispatcher already caps batch sizes at max_num_seqs, so we just |
| 1512 | # need to verify that enough blocks exist. Raising here instead |
| 1513 | # of silently capping cudagraph_capture_sizes avoids unintended |
| 1514 | # restrictions on PIECEWISE (prefill) cudagraphs. |
| 1515 | # See: https://github.com/vllm-project/vllm/issues/34094 |
| 1516 | if ( |
| 1517 | kv_cache_config is not None |
| 1518 | and max_num_reqs is not None |
| 1519 | and cudagraph_mode.has_full_cudagraphs() |
| 1520 | and not is_profiling |
| 1521 | and kv_cache_config.has_mamba_layers |
| 1522 | and max_num_reqs > kv_cache_config.num_blocks |
| 1523 | ): |
| 1524 | raise ValueError( |
| 1525 | f"max_num_seqs ({max_num_reqs}) exceeds available Mamba cache " |
| 1526 | f"blocks ({kv_cache_config.num_blocks}). Each decode sequence " |
| 1527 | "requires one Mamba cache block, so CUDA graph capture cannot " |
| 1528 | "proceed. Please lower max_num_seqs to at most " |
| 1529 | f"{kv_cache_config.num_blocks} or increase " |
| 1530 | "gpu_memory_utilization." |
| 1531 | ) |
| 1532 | |
| 1533 | self.cudagraph_mode = cudagraph_mode |
| 1534 | return cudagraph_mode |
| 1535 | |
| 1536 | def adjust_cudagraph_sizes_for_spec_decode( |
| 1537 | self, uniform_decode_query_len: int, tensor_parallel_size: int |
| 1538 | ): |
| 1539 | multiple_of = uniform_decode_query_len |
| 1540 | if tensor_parallel_size > 1 and self.pass_config.enable_sp: |
| 1541 | multiple_of = max(uniform_decode_query_len, tensor_parallel_size) |
| 1542 | if ( |
| 1543 | multiple_of % uniform_decode_query_len != 0 |
| 1544 | or multiple_of % tensor_parallel_size != 0 |
| 1545 | ): |
| 1546 | raise ValueError( |
| 1547 | f"Can't determine cudagraph shapes that are both a " |
| 1548 | f"multiple of {uniform_decode_query_len} " |
| 1549 | f"(num_speculative_tokens + 1) required by spec-decode " |
| 1550 | f"and {tensor_parallel_size} (tensor_parallel_size) " |
| 1551 | f"required by sequence parallelism please adjust " |
| 1552 | f"num_speculative_tokens or disable sequence parallelism" |
| 1553 | ) |
| 1554 | |
| 1555 | if not self.cudagraph_capture_sizes or multiple_of <= 1: |
| 1556 | return |
| 1557 | |
| 1558 | assert self.max_cudagraph_capture_size is not None |
| 1559 | rounded_sizes = sorted( |
| 1560 | set( |
| 1561 | round_up(size, multiple_of) |
| 1562 | for size in self.cudagraph_capture_sizes |
| 1563 | if round_up(size, multiple_of) <= self.max_cudagraph_capture_size |
| 1564 | ) |
| 1565 | ) |
| 1566 | |
| 1567 | if len(rounded_sizes) == 0 and multiple_of <= self.max_cudagraph_capture_size: |
| 1568 | # if one valid but would be round_down use that |
| 1569 | rounded_sizes = [multiple_of] |
| 1570 | |
| 1571 | if len(rounded_sizes) == 0: |
| 1572 | raise ValueError( |
| 1573 | f"No valid cudagraph sizes after rounding to multiple of {multiple_of} " |
| 1574 | f"(num_speculative_tokens + 1 or tp if sequence parallelism is enabled)" |
| 1575 | f" please adjust num_speculative_tokens ({uniform_decode_query_len - 1}" |
| 1576 | f") or max_cudagraph_capture_size ({self.max_cudagraph_capture_size})" |
| 1577 | f" or cudagraph_capture_sizes ({self.cudagraph_capture_sizes})" |
| 1578 | ) |
| 1579 | |
| 1580 | self.max_cudagraph_capture_size = rounded_sizes[-1] |
| 1581 | self.cudagraph_capture_sizes = rounded_sizes |
| 1582 | |
| 1583 | def get_compile_ranges(self) -> list[Range]: |
| 1584 | """Get the compile ranges for the compilation config.""" |
| 1585 | if self.compile_ranges_endpoints is None: |
| 1586 | return [] |
| 1587 | endpoints = sorted(set(self.compile_ranges_endpoints)) |
| 1588 | return [Range(s + 1, e) for s, e in zip([0] + endpoints[:-1], endpoints)] |