compilation.py 54.5 KB
Newer Older
1
2
3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

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
16
17
18
19
20
from vllm.config.utils import (
    Range,
    config,
    get_hash_factors,
    hash_factors,
)
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
30
31
32
33
34
else:
    VllmConfig = object

logger = init_logger(__name__)


35
class CompilationMode(enum.IntEnum):
36
37
38
39
40
41
42
43
44
45
46
47
48
    """The compilation approach used for torch.compile-based compilation of the
    model."""

    NONE = 0
    """No torch.compile compilation is applied, model runs in fully eager pytorch mode.
    The model runs as-is."""
    STOCK_TORCH_COMPILE = 1
    """The standard `torch.compile` compilation pipeline."""
    DYNAMO_TRACE_ONCE = 2
    """Single Dynamo trace through the model, avoiding recompilation."""
    VLLM_COMPILE = 3
    """Custom vLLM Inductor-based backend with caching, piecewise compilation,
    shape specialization, and custom passes."""
49
50


51
class CUDAGraphMode(enum.Enum):
52
    """Constants for the cudagraph mode in CompilationConfig.
53
54
55
    Meanwhile, the subset enum `NONE`, `PIECEWISE` and `FULL` are also
    treated as concrete runtime mode for cudagraph runtime dispatching.
    """
56

57
58
59
60
61
62
    NONE = 0
    PIECEWISE = 1
    FULL = 2
    FULL_DECODE_ONLY = (FULL, NONE)
    FULL_AND_PIECEWISE = (FULL, PIECEWISE)

63
64
    def decode_mode(self) -> "CUDAGraphMode":
        return CUDAGraphMode(self.value[0]) if self.separate_routine() else self
65

66
67
    def mixed_mode(self) -> "CUDAGraphMode":
        return CUDAGraphMode(self.value[1]) if self.separate_routine() else self
68

69
70
71
72
73
74
    def has_mode(self, mode: "CUDAGraphMode") -> bool:
        assert not mode.separate_routine()
        if self.separate_routine():
            return mode.value in self.value
        return self == mode

75
    def requires_piecewise_compilation(self) -> bool:
76
        return self.has_mode(CUDAGraphMode.PIECEWISE)
77

78
79
    def max_cudagraph_mode(self) -> "CUDAGraphMode":
        return CUDAGraphMode(max(self.value)) if self.separate_routine() else self
80
81
82
83

    def has_full_cudagraphs(self) -> bool:
        return self.max_cudagraph_mode() == CUDAGraphMode.FULL

84
85
86
    def has_piecewise_cudagraphs(self) -> bool:
        return self.requires_piecewise_compilation()

87
88
89
    def separate_routine(self) -> bool:
        return isinstance(self.value, tuple)

90
91
92
93
94
95
    @classmethod
    def valid_runtime_modes(cls) -> frozenset["CUDAGraphMode"]:
        return frozenset({cls.NONE, cls.PIECEWISE, cls.FULL})

    def is_valid_runtime_mode(self) -> bool:
        return self in CUDAGraphMode.valid_runtime_modes()
96

97
98
99
    def __str__(self) -> str:
        return self.name

100
101
102
    def __bool__(self) -> bool:
        return self != CUDAGraphMode.NONE

103

104
105
106
107
108
109
@config
class PassConfig:
    """Configuration for custom Inductor passes.

    This is separate from general `CompilationConfig` so that inductor passes
    don't all have access to full configuration - that would create a cycle as
110
    the `PassManager` is set as a property of config.
111

112
113
114
115
116
117
    You must pass PassConfig to VLLMConfig constructor via the CompilationConfig
    constructor. VLLMConfig's post_init does further initialization.
    If used outside of the VLLMConfig, some fields may be left in an
    improper state.
    """

118
    # New flags
119
    fuse_norm_quant: bool = None  # type: ignore[assignment]
120
    """Fuse the custom RMSNorm + quant ops."""
121
    fuse_act_quant: bool = None  # type: ignore[assignment]
122
    """Fuse the custom SiluMul + quant ops."""
123
    fuse_attn_quant: bool = None  # type: ignore[assignment]
124
    """Fuse the custom Attention and MLAAttention + quant ops."""
125
    eliminate_noops: bool = Field(default=True)
126
    """Eliminate no-op ops."""
127
    enable_sp: bool = None  # type: ignore[assignment]
128
129
130
    """Enable sequence parallelism. Requires TP>1. Automatically disabled
    if the model's hidden_size is too small for SP to be beneficial
    (threshold is device-capability dependent)."""
131
    fuse_gemm_comms: bool = None  # type: ignore[assignment]
132
    """Enable async TP."""
133
    fuse_allreduce_rms: bool = None  # type: ignore[assignment]
134
    """Enable flashinfer allreduce fusion."""
135
136
    enable_qk_norm_rope_fusion: bool = False
    """Enable fused Q/K RMSNorm + RoPE pass."""
137

138
    # ROCm/AITER specific fusions
139
    fuse_act_padding: bool = None  # type: ignore[assignment]
140
    """Fuse the custom RMSNorm + padding ops."""
141
    fuse_rope_kvcache: bool = None  # type: ignore[assignment]
142
143
144
145
146
147
    """Fuse the QK rope + KV cache ops."""

    rope_kvcache_fusion_max_token_num: int = 256
    """The threshold for ROCm AITER RoPE+KVCache fusion e.g. for small batch decode.
    Larger batch sizes e.g. during prefill will use the unfused kernels.
    """
148

149
150
151
152
    fi_allreduce_fusion_max_size_mb: float | None = None
    """The threshold of the communicated tensor sizes under which
    vllm should use flashinfer fused allreduce. Specified as a
    float in MB.
153
    Unspecified will fallback to default values
154
155
156
157
158
159
160
161
162
163
164
165
166
    which are compute capability and world size dependent.
        FI_ALLREDUCE_FUSION_MAX_SIZE_MB = {
            90: {
                2: 64,  # 64MB
                4: 2,  # 2MB
                8: 1,  # 1MB
            },
            100: {
                2: 64,  # 64MB
                4: 32,  # 32MB
                8: 1,  # 1MB
            },
        }, where key is the device capability"""
167
168
169
170
171
    sp_min_token_num: int | None = None
    """The minimum number of tokens above which vllm should use
    sequence parallelism. Specified as an integer token count.
    Unspecified will fallback to default values which are compute
    capability and world size dependent."""
172
173
174

    # TODO(luka) better pass enabling system.

175
176
177
178
179
180
181
182
    def flashinfer_max_size(self, world_size: int) -> int | None:
        """
        Returns the max communication size in bytes for flashinfer
        allreduce fusion for the given world size. Returns None if world size
        is not supported by configs as it's not supported by flashinfer.
        """

        MiB = 1024 * 1024
183
184
185
        FI_SUPPORTED_WORLD_SIZES = [2, 4, 8]
        if world_size not in FI_SUPPORTED_WORLD_SIZES:
            return None
186
187
188
189
190
191
192
193
        max_size_mb = self.fi_allreduce_fusion_max_size_mb
        if max_size_mb is None:
            max_size_mb = self.default_fi_allreduce_fusion_max_size_mb().get(world_size)

        return int(max_size_mb * MiB) if max_size_mb is not None else None

    @staticmethod
    def default_fi_allreduce_fusion_max_size_mb() -> dict[int, float]:
194
195
196
        from vllm.compilation.passes.fusion.allreduce_rms_fusion import (
            FI_ALLREDUCE_FUSION_MAX_SIZE_MB,
        )
197
198
199
200
        from vllm.platforms import current_platform

        if not current_platform.is_cuda():
            return {}
201
202
203
204
        capability = current_platform.get_device_capability()
        if capability is None:
            return {}
        return FI_ALLREDUCE_FUSION_MAX_SIZE_MB.get(capability.to_int(), {})
205

206
    def compute_hash(self) -> str:
207
208
209
210
211
        """
        Produces a hash unique to the pass configuration.
        Any new fields that affect compilation should be added to the hash.
        Any future fields that don't affect compilation should be excluded.
        """
212

213
        return hash_factors(get_hash_factors(self, set()))
214

215
    @field_validator(
216
217
218
219
220
221
        "fuse_norm_quant",
        "fuse_act_quant",
        "fuse_attn_quant",
        "enable_sp",
        "fuse_gemm_comms",
        "fuse_allreduce_rms",
222
        "fuse_act_padding",
223
        "fuse_rope_kvcache",
224
225
226
227
228
229
230
231
232
        mode="wrap",
    )
    @classmethod
    def _skip_none_validation(cls, value: Any, handler: Callable) -> Any:
        """Skip validation if the value is `None` when initialisation is delayed."""
        if value is None:
            return value
        return handler(value)

233
    def __post_init__(self) -> None:
234
235
236
237
        # Handle deprecation and defaults

        if not self.eliminate_noops:
            if self.fuse_norm_quant or self.fuse_act_quant:
238
239
                logger.warning_once(
                    "Fusion enabled but reshape elimination disabled. "
240
241
                    "RMSNorm/SiluMul + quant (fp8) fusion might not work"
                )
242
            if self.fuse_attn_quant:
243
244
                logger.warning_once(
                    "Fusion enabled but reshape elimination disabled. "
245
246
                    "Attention + quant (fp8) fusion might not work"
                )
247
            if self.fuse_allreduce_rms:
248
249
250
251
                logger.warning_once(
                    "Fusion enabled but reshape elimination disabled. "
                    "Allreduce + rms norm + quant (fp8) fusion might not work"
                )
252
253
254
255
256
            if self.fuse_act_padding:
                logger.warning_once(
                    "Fusion enabled but reshape elimination disabled. "
                    "RMSNorm + padding fusion might not work"
                )
257
        if self.enable_qk_norm_rope_fusion and not current_platform.is_cuda_alike():
258
259
            logger.warning_once(
                "QK Norm + RoPE fusion enabled but the current platform is not "
260
                "CUDA or ROCm. The fusion will be disabled."
261
262
            )
            self.enable_qk_norm_rope_fusion = False
263
264
265
266
267
268
        if self.fuse_act_padding and not current_platform.is_rocm():
            logger.warning_once(
                "Padding fusion enabled but the current platform is not ROCm. "
                "The fusion will be disabled."
            )
            self.fuse_act_padding = False
269
270
271
272
273
274
        if self.fuse_rope_kvcache and not current_platform.is_rocm():
            logger.warning_once(
                "KV cache fusion currently only enabled on ROCm. "
                "The fusion will be disabled."
            )
            self.fuse_rope_kvcache = False
275

276
277
278
279
280
281
282
283
284
    def log_enabled_passes(self) -> None:
        """
        Log the enabled custom fusion passes.
        This is called at the end of VLLMConfig post_init,
        after all defaults are finalized.
        TODO also log the compile ranges for which this is enabled.
        """
        enabled_fusions = [
            f.name[len("fuse_") :]
285
            for f in fields(self)  # type: ignore[arg-type]
286
287
288
289
290
291
292
293
            if getattr(self, f.name) and f.name.startswith("fuse_")
        ]

        if enabled_fusions:
            logger.info_once(
                "Enabled custom fusions: %s", ", ".join(enabled_fusions), scope="global"
            )

294

295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
class DynamicShapesType(str, enum.Enum):
    """Types of dynamic shapes handling in torch.compile().
    see  Dynamic shapes and vllm guard dropping in torch_compile.md
    for more details."""

    BACKED = "backed"
    """Use backed dynamic shapes. torch.compile() guards on backed dynamic
    shapes and may add guards. Symbols are specialized to 0, 1, or >=2 even
    without encountering branching on those ranges."""

    UNBACKED = "unbacked"
    """Use unbacked dynamic shapes. Guaranteed not to be guarded on and not
    0/1 specialized, but may throw data dependent errors when branches require
    their value without explicit unbacked handling."""

    BACKED_SIZE_OBLIVIOUS = "backed_size_oblivious"
    """Experimental flag that treats backed symbols as unbacked when explicit
    unbacked handling is defined."""


@config
class DynamicShapesConfig:
    """Configuration to control/debug torch compile dynamic shapes."""

    type: DynamicShapesType = DynamicShapesType.BACKED
    """Controls the type of dynamic shapes handling to use with torch.compile().

    - BACKED: Default PyTorch behavior with potential guards ignored.
    - UNBACKED: No guards guaranteed (most sound) but may throw
      data dependent errors.
    - BACKED_SIZE_OBLIVIOUS: Experimental safer alternative to
      backed/unbacked.
    """

329
330
331
332
333
334
335
336
337
338
    evaluate_guards: bool = False
    """
    A debug mode to detect and fail if Dynamo ever specializes a dynamic shape by
    guarding on it. When True, dynamic shape guards are not dropped from dynamo.
    And a failure will be triggered if a recompilation ever happens due to that.
    This mode requires VLLM_USE_BYTECODE_HOOK to be 0.
    Enabling this allow observing the dynamic shapes guards in the tlparse
    artifacts also.
    When type is backed, aot_compile must be disabled for this mode to work.
    until this change picked up https://github.com/pytorch/pytorch/pull/169239.
339
    """
340

341
    assume_32_bit_indexing: bool = False
342
343
    """
    whether all tensor sizes can use 32 bit indexing.
344
    `True` requires PyTorch 2.10+
345
    """
346
347
348
349
350
351
352
353

    def compute_hash(self) -> str:
        """
        Provide a hash for DynamicShapesConfig
        """

        from vllm.config.utils import get_hash_factors, hash_factors

354
        factors = get_hash_factors(self, set())
355
356
357
        return hash_factors(factors)


358
359
@config
class CompilationConfig:
360
361
362
363
364
365
    """Configuration for compilation.

    You must pass CompilationConfig to VLLMConfig constructor.
    VLLMConfig's post_init does further initialization. If used outside of the
    VLLMConfig, some fields will be left in an improper state.

366
367
    It contains PassConfig, which controls the custom fusion/transformation passes.
    The rest has three parts:
368
369

    - Top-level Compilation control:
370
        - [`mode`][vllm.config.CompilationConfig.mode]
371
372
373
374
375
        - [`debug_dump_path`][vllm.config.CompilationConfig.debug_dump_path]
        - [`cache_dir`][vllm.config.CompilationConfig.cache_dir]
        - [`backend`][vllm.config.CompilationConfig.backend]
        - [`custom_ops`][vllm.config.CompilationConfig.custom_ops]
        - [`splitting_ops`][vllm.config.CompilationConfig.splitting_ops]
376
        - [`compile_mm_encoder`][vllm.config.CompilationConfig.compile_mm_encoder]
377
    - CudaGraph capture:
378
        - [`cudagraph_mode`][vllm.config.CompilationConfig.cudagraph_mode]
379
380
        - [`cudagraph_capture_sizes`]
        [vllm.config.CompilationConfig.cudagraph_capture_sizes]
381
382
        - [`max_cudagraph_capture_size`]
        [vllm.config.CompilationConfig.max_cudagraph_capture_size]
383
384
385
386
387
388
        - [`cudagraph_num_of_warmups`]
        [vllm.config.CompilationConfig.cudagraph_num_of_warmups]
        - [`cudagraph_copy_inputs`]
        [vllm.config.CompilationConfig.cudagraph_copy_inputs]
    - Inductor compilation:
        - [`compile_sizes`][vllm.config.CompilationConfig.compile_sizes]
389
390
        - [`compile_ranges_endpoints`]
            [vllm.config.CompilationConfig.compile_ranges_endpoints]
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
        - [`inductor_compile_config`]
        [vllm.config.CompilationConfig.inductor_compile_config]
        - [`inductor_passes`][vllm.config.CompilationConfig.inductor_passes]
        - custom inductor passes

    Why we have different sizes for cudagraph and inductor:
    - cudagraph: a cudagraph captured for a specific size can only be used
        for the same size. We need to capture all the sizes we want to use.
    - inductor: a graph compiled by inductor for a general shape can be used
        for different sizes. Inductor can also compile for specific sizes,
        where it can have more information to optimize the graph with fully
        static shapes. However, we find the general shape compilation is
        sufficient for most cases. It might be beneficial to compile for
        certain small batchsizes, where inductor is good at optimizing.
    """
406

407
    # Top-level Compilation control
408
    mode: CompilationMode = None  # type: ignore[assignment]
409
410
411
412
413
414
415
416
417
418
419
420
421
    """The compilation approach used for torch.compile-based compilation of the
    model.

    - None: If None, we will select the default compilation mode.
      For V1 engine this is 3.
    - 0: NONE: No torch.compile compilation is applied, model runs in fully
         eager pytorch mode. The model runs as-is.
    - 1: STOCK_TORCH_COMPILE: The standard `torch.compile` compilation pipeline.
    - 2: DYNAMO_TRACE_ONCE: Single Dynamo trace through the model, avoiding
         recompilation by removing guards.
         Requires no dynamic-shape-dependent control-flow.
    - 3: VLLM_COMPILE: Custom vLLM Inductor-based backend with caching,
         piecewise compilation, shape specialization, and custom passes."""
422
    debug_dump_path: Path | None = None
423
424
425
426
427
    """The path to dump the debug information."""
    cache_dir: str = ""
    """The directory to store the compiled graph, to accelerate Inductor
    compilation. By default, it will use model-related information to generate
    a cache directory."""
428
429
430
431
432
433
434
435
436
    compile_cache_save_format: Literal["binary", "unpacked"] = field(
        default_factory=lambda: envs.VLLM_COMPILE_CACHE_SAVE_FORMAT
    )
    """Format for saving torch compile cache:\n
    - "binary": saves as binary file (multiprocess safe)\n
    - "unpacked": saves as directory structure for inspection/debugging
    (NOT multiprocess safe)\n
    Defaults to `VLLM_COMPILE_CACHE_SAVE_FORMAT` if not specified.
    """
437
    backend: str = ""
438
439
    """The backend for compilation. It needs to be a string:

440
441
    - "" (empty string): use the default backend ("inductor" on CUDA-alike
    platforms).
442
443
444
445
446
    - "eager"/"openxla"/...: use the specified backend registered in PyTorch.
    - "full.module.name": a qualified name which can be used to import the

    backend function.
    We use string to avoid serialization issues when using compilation in a
447
    distributed setting. When the compilation mode is 1 or 2, the backend is
448
    used for the compilation directly (it sees the whole graph). When the
449
450
    compilation mode is 3, the backend supports both whole graph and piecewise
    compilation, available backends include eager, inductor, and custom backends,
451
    the latter of which can be defined via `get_compile_backend`. Furthermore,
452
    compilation is only piecewise if splitting ops is set accordingly and
453
    use_inductor_graph_partition is off. Note that the default options for
454
455
    splitting ops are sufficient for piecewise compilation.
    """
456
457
458
459
460
461
462
463
464
465
    custom_ops: list[str] = field(default_factory=list)
    """Fine-grained control over which custom ops to enable/disable. Use 'all'
    to enable all, 'none' to disable all. Also specify a list of custom op
    names to enable (prefixed with a '+'), or disable (prefixed with a '-').
    Examples:

    - 'all,-op1' to enable all except op1
    - 'none,+op1,+op2' to enable only op1 and op2

    By default, all custom ops are enabled when running without Inductor and
466
467
    disabled when running with Inductor: mode>CompilationMode.NONE and
    backend="inductor".
468
    Inductor generates (fused) Triton kernels for disabled custom ops."""
469
470
471
472
473
474
475
476
477

    ir_enable_torch_wrap: bool = None  # type: ignore[assignment]
    """If True, enable vllm_ir torch custom op wrapping during the forward pass.
    When False, torch custom op wrapping is disabled, allowing Dynamo to trace the
    selected implementation directly or avoiding torch custom op overhead in eager mode.
    Defaults to True when using Inductor with vllm-compile
    (backend=="inductor" and mode == VLLM_COMPILE), False otherwise.
    """

478
    splitting_ops: list[str] | None = None
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
    """A list of ops to exclude from cudagraphs, used in piecewise compilation.

    The behavior depends on use_inductor_graph_partition:

    - When use_inductor_graph_partition=False (default):
        These ops are used for Dynamo FX-level graph splitting. The graph is
        split at these ops before Inductor compilation, creating separate
        subgraphs for cudagraph capture.

    - When use_inductor_graph_partition=True:
        These ops are used to register Inductor partition rules. The graph
        partitioning happens at Inductor codegen time after all passes and
        fusions are finished, allowing compilation and custom passes to operate
        on the full graph while still excluding these ops from cudagraphs.

    If None, defaults to attention ops for piecewise cudagraphs.
    If empty list [], no ops are excluded (suitable for full cudagraphs)."""
496
    compile_mm_encoder: bool = False
Harry Mellor's avatar
Harry Mellor committed
497
    """Whether or not to compile the multimodal encoder.
498
499
500
501
    Currently, this only works for `Qwen2_5_vl` and `mLLaMa4` models on selected
    platforms. It may also work for models loaded with the Transformers modeling backend
    if the encoder is compilable. Disabled by default until more models are
    supported/tested to work."""
502

503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
    # Vision encoder CUDA graph
    cudagraph_mm_encoder: bool = False
    """Enable CUDA graph capture for multimodal encoder (ViT).
    When enabled, captures full encoder forward as CUDA graph
    for each token budget level."""

    encoder_cudagraph_token_budgets: list[int] = field(default_factory=list)
    """Token budget levels for encoder CUDA graph capture.
    Each budget defines a fixed token capacity. At runtime, images are greedy-packed
    into the smallest fitting budget and the corresponding CUDA graph is replayed.
    If empty (default), auto-inferred from model architecture as power-of-2
    levels from the model's estimated min budget to max budget.
    User-provided values override auto-inference.
    Example: [2048, 4096, 8192, 13824]"""

    encoder_cudagraph_max_images_per_batch: int = 0
    """Maximum number of images per batch for encoder CUDA graph capture.
    Determines the fixed batch size used during graph capture.
    If 0 (default), auto-inferred as max_budget // min_budget from the
    model's budget range. User-provided positive value overrides
    auto-inference."""

525
    # Inductor capture
526
    compile_sizes: list[int | str] | None = None
527
528
529
    """Sizes to compile for inductor. In addition
    to integers, it also supports "cudagraph_capture_sizes" to
    specify the sizes for cudagraph capture."""
530

531
532
    compile_ranges_endpoints: list[int] | None = None
    """Endpoints for Inductor compile ranges.
533
    The compile ranges are
534
535
536
    [1, endpoints[0]],
    [endpoints[0] + 1, endpoints[1]], ...,
    [endpoints[-1] + 1, max_num_batched_tokens].
537
538
    Compile sizes are also used single element ranges,
    the range is represented as [compile_sizes[i], compile_sizes[i]].
539
540

    If a range overlaps with the compile size, graph for compile size
541
542
543
544
545
    will be prioritized, i.e. if we have a range [1, 8] and a compile size 4,
    graph for compile size 4 will be compiled and used instead of the graph
    for range [1, 8].
    """

546
547
548
    inductor_compile_config: dict = field(default_factory=dict)
    """Additional configurations for inductor.
    - None: use default configurations."""
549

550
551
552
553
554
555
556
557
    inductor_passes: dict[str, str] = field(default_factory=dict)
    """Additional passes for inductor. It is a dictionary
    from pass name to pass function qualified name. We use function
    name because the config uses JSON format. If we pass the config
    from Python, functions can also be passed directly via Python object
    constructor, e.g. `CompilationConfig(inductor_passes={"a": func})`."""

    # CudaGraph compilation
558
    cudagraph_mode: CUDAGraphMode = None  # type: ignore[assignment]
559
    """
Harry Mellor's avatar
Harry Mellor committed
560
561
    The mode of the cudagraph:

562
    - NONE, no cudagraph capture.
563
    - PIECEWISE.
564
565
    - FULL.
    - FULL_DECODE_ONLY.
566
    - FULL_AND_PIECEWISE. (v1 default)
567
568

    PIECEWISE mode build piecewise cudagraph only, keeping the cudagraph
co63oc's avatar
co63oc committed
569
    incompatible ops (i.e. some attention ops) outside the cudagraph
570
571
572
573
574
    for general flexibility.

    FULL mode: Capture full cudagraph for all batches. Can be good for small
    models or workloads with small prompts; not supported by many backends.
    Generally for performance FULL_AND_PIECEWISE is better.
575

576
577
578
579
    FULL_DECODE_ONLY mode: Capture full cudagraph for decode batches only.
    Mixed prefill-decode batches are run without cudagraphs. Can be good for
    decode instances in a P/D setup where prefill is not as important so we
    can save some memory.
580

581
582
    FULL_AND_PIECEWISE mode: Capture full cudagraph for decode batches and
    piecewise cudagraph for prefill and mixed prefill-decode batches.
583
    This is the most performant mode for most models and is the default.
584
585

    Currently, the cudagraph mode is only used for the v1 engine.
586
587
    Note that the cudagraph logic is generally orthogonal to the
    compilation logic. While piecewise cudagraphs require piecewise
588
    compilation (mode=VLLM_COMPILE and non-empty splitting_ops), full
589
    cudagraphs are supported with and without compilation.
590
591

    Warning: This flag is new and subject to change in addition
592
593
    more modes may be added.
    """
594
595
596
597
598
    cudagraph_num_of_warmups: int = 0
    """Number of warmup runs for cudagraph.
    It means the first several runs will be treated as warmup runs.
    Only after that, the execution will be recorded, and the recorded
    cudagraph will be used for subsequent runs."""
599
    cudagraph_capture_sizes: list[int] = None  # type: ignore[assignment]
600
601
602
603
604
605
606
607
    """Sizes to capture cudagraph.
    - None (default): capture sizes are inferred from vllm config.
    - list[int]: capture sizes are specified as given."""
    cudagraph_copy_inputs: bool = False
    """Whether to copy input tensors for
    cudagraph. If the caller can guarantee that the same input buffers
    are always used, it can set this to False. Otherwise, it should
    set this to True, and the compiler will copy the input to an
608
    internally managed buffer. Default is False.
609
610
    Note that this flag is only effective when cudagraph_mode is PIECEWISE.
    """
611
612
613
614
615
616
617
618
    cudagraph_specialize_lora: bool = True
    """Whether to create separate cuda graphs for cases with and without active
    LoRA adapters. When set to False, the LoRA-enabled cuda graph will be used
    for all cases, incurring the overhead of running LoRA ops even when no
    adapters are active. Setting this to True will remove this overhead at the
    cost of increased startup time and slightly higher memory usage.
    When `enable_lora` is False, this option has no effect.
    """
619

620
    use_inductor_graph_partition: bool = None  # type: ignore[assignment]
621
622
623
624
625
626
627
    """Use inductor graph partition to split the graph at cudagraph_unsafe ops.
    This partition happens at inductor codegen time after all passes and fusions
    are finished. It generates a single `call` function which wraps
    cudagraph-safe ops into partition functions and leave cudagraph-unsafe ops
    outside the partition functions. For a graph with N cudagraph-unsafe ops
    (e.g., Attention), there would be N+1 partitions. To mark an op as
    cudagraph unsafe, we can add `tags=(torch._C.Tag.cudagraph_unsafe)` when
628
    register the custom op.
629
630
631
632
633
634
635
636
637
638
639

    This config supports both full cudagraph and piecewise cudagraph without
    compiling twice. For piecewise cudagraph, it applies vLLM CUDAGraph wrapper
    to each partition. For N+1 partitions, there would be N+1
    CUDAGraph wrapper instances.

    For full CUDAGraph, we always apply a single CUDAGraph wrapper outside the
    inductor `call` function in the model runner. The top-level full cudagraph
    capture ignores all partitioning.
    """

640
641
642
    pass_config: PassConfig = field(default_factory=PassConfig)
    """Custom inductor passes, see PassConfig for more details"""

643
    max_cudagraph_capture_size: int = None  # type: ignore[assignment]
644
    """The maximum cudagraph capture size.
645
646

    If cudagraph_capture_sizes is specified, this will be set to the largest
647
648
649
650
651
652
653
654
    size in that list (or checked for consistency if specified). If
    cudagraph_capture_sizes is not specified, the list of sizes is generated
    automatically following the pattern:

        [1, 2, 4] + list(range(8, 256, 8)) + list(
        range(256, max_cudagraph_capture_size + 1, 16))

    If not specified, max_cudagraph_capture_size is set to min(max_num_seqs*2,
655
    512) by default. This voids OOM in tight memory scenarios with small
656
657
658
    max_num_seqs, and prevents capture of many large graphs (>512) that would
    greatly increase startup time with limited performance benefit.
    """
659
660
661
662
663
664

    dynamic_shapes_config: DynamicShapesConfig = field(
        default_factory=DynamicShapesConfig
    )
    """Configuration for dynamic shapes options"""

665
666
    local_cache_dir: str = field(default=None, init=False)  # type: ignore
    """local cache dir for each rank"""
667

668
    fast_moe_cold_start: bool | None = None
669
670
671
672
673
674
675
676
677
678
    """Optimization for fast MOE cold start.

    This is a bit of a hack that assumes that:
    1. the only decoder forward pass being run is the current model
    2. the decoder forward pass runs all of the MOEs in the order in which they
       are initialized

    When the above two conditions hold, this option greatly decreases cold start
    time for MOE models.

679
680
681
682
683
684
685
686
    The options are:
    - True: optimization is always on
    - False: optimization is always off
    - None: optimization is on usually but off for speculative decoding

    If conditions 1&2 don't hold then this option will lead to silent
    incorrectness.
    The only condition in which this doesn't hold is speculative
687
688
689
690
691
    decoding, where there is a draft model that may have MOEs in them.

    NB: We're working on a longer-term solution that doesn't need these assumptions.
    """

692
    # keep track of enabled and disabled custom ops
693
    enabled_custom_ops: Counter[str] = field(default_factory=Counter, init=False)
694
    """custom ops that are enabled"""
695
    disabled_custom_ops: Counter[str] = field(default_factory=Counter, init=False)
696
697
698
699
700
701
    """custom ops that are disabled"""
    traced_files: set[str] = field(default_factory=set, init=False)
    """files that are traced for compilation"""
    compilation_time: float = field(default=0.0, init=False)
    """time taken for compilation"""

702
    static_forward_context: dict[str, Any] = field(default_factory=dict, init=False)
703
704
705
706
    """Per-model forward context
    Map from layer name to layer objects that need to be accessed outside
    model code, e.g., Attention, FusedMOE when dp_size>1."""

707
708
709
710
    static_all_moe_layers: list[str] = field(default_factory=list, init=False)
    """The names of all the MOE layers in the model
    """

711
    # Attention ops; used for piecewise cudagraphs
712
    # Use PyTorch operator format: "namespace::name"
713
    _attention_ops: ClassVar[list[str]] = [
714
715
716
717
718
719
720
        "vllm::unified_attention_with_output",
        "vllm::unified_mla_attention_with_output",
        "vllm::mamba_mixer2",
        "vllm::mamba_mixer",
        "vllm::short_conv",
        "vllm::linear_attention",
        "vllm::plamo2_mamba_mixer",
721
        "vllm::gdn_attention_core",
722
        "vllm::olmo_hybrid_gdn_full_forward",
723
        "vllm::kda_attention",
724
        "vllm::sparse_attn_indexer",
725
        "vllm::rocm_aiter_sparse_attn_indexer",
726
727
    ]

728
729
730
731
732
733
734
735
    def compute_hash(self) -> str:
        """
        Provide a hash that uniquely identifies all the configs
        that affect the structure of the computation
        graph from input ids/embeddings to the final hidden states,
        excluding anything before input ids/embeddings and after
        the final hidden states.
        """
736
737
738
739
740
741
742
743
744
745
746
747
748
        # Opt-out: default-include declared fields; keep a tiny exclude set;
        # normalize types; keep SHA-256. For nested opaque configs, include a
        # stable identifier (e.g., pass_config.compute_hash()) instead of object id.

        ignored_factors = {
            # Paths/dirs and runtime/metrics that don’t affect compiled graph
            "debug_dump_path",
            "cache_dir",
            "local_cache_dir",
            "traced_files",
            "compilation_time",
            "static_forward_context",
            "pass_config",  # handled separately below
749
            "dynamic_shapes_config",  # handled separately below
750
751
752
753
754
        }

        from vllm.config.utils import get_hash_factors, hash_factors

        factors = get_hash_factors(self, ignored_factors)
755

756
        factors["pass_config"] = self.pass_config.compute_hash()
757
        factors["dynamic_shapes_config"] = self.dynamic_shapes_config.compute_hash()
758
        return hash_factors(factors)
759
760

    def __repr__(self) -> str:
761
        exclude: dict[str, bool | dict[str, bool]] = {
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
            "static_forward_context": True,
            "enabled_custom_ops": True,
            "disabled_custom_ops": True,
            "compilation_time": True,
            "traced_files": True,
            "inductor_compile_config": {
                "post_grad_custom_post_pass": True,
            },
        }

        # exclude default attr in pass_config
        pass_config_exclude = {}
        for attr, default_val in vars(PassConfig()).items():
            if getattr(self.pass_config, attr) == default_val:
                pass_config_exclude[attr] = True
        if pass_config_exclude:
            exclude["pass_config"] = pass_config_exclude

780
        config = TypeAdapter(CompilationConfig).dump_python(
781
            self, exclude=exclude, exclude_unset=True
782
        )
783
784

        return str(config)
785
786
787

    __str__ = __repr__

788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
    @field_validator("mode", mode="before")
    @classmethod
    def validate_mode_before(cls, value: Any) -> Any:
        """
        Enable parsing the `mode` field from string mode names.
        Accepts both integers (0-3) and string names, like NONE, STOCK_TORCH_COMPILE,
        DYNAMO_TRACE_ONCE, VLLM_COMPILE.
        """
        if isinstance(value, str):
            # Convert string mode name to integer value
            mode_name = value.upper()

            if mode_name not in CompilationMode.__members__:
                raise ValueError(
                    f"Invalid compilation mode: {value}. "
                    f"Valid modes are: {', '.join(CompilationMode.__members__.keys())}"
                )

            return CompilationMode[mode_name]
        return value

809
810
811
    @field_validator("cudagraph_mode", mode="before")
    @classmethod
    def validate_cudagraph_mode_before(cls, value: Any) -> Any:
812
        """Enable parsing of the `cudagraph_mode` enum type from string."""
813
814
815
816
        if isinstance(value, str):
            return CUDAGraphMode[value.upper()]
        return value

817
818
819
820
821
822
823
824
    @field_validator("pass_config", mode="before")
    @classmethod
    def validate_pass_config_before(cls, value: Any) -> Any:
        """Enable parsing of the `pass_config` field from a dictionary."""
        if isinstance(value, dict):
            return PassConfig(**value)
        return value

825
826
827
828
829
830
831
832
833
834
    @field_validator("compile_cache_save_format")
    @classmethod
    def validate_compile_cache_save_format(cls, value: str) -> str:
        if value not in ("binary", "unpacked"):
            raise ValueError(
                f"compile_cache_save_format must be 'binary' or 'unpacked', "
                f"got: {value}"
            )
        return value

835
836
837
838
    @field_validator(
        "level",
        "mode",
        "cudagraph_mode",
839
        "max_cudagraph_capture_size",
840
        "use_inductor_graph_partition",
841
        "ir_enable_torch_wrap",
842
843
844
845
846
847
848
849
850
        mode="wrap",
    )
    @classmethod
    def _skip_none_validation(cls, value: Any, handler: Callable) -> Any:
        """Skip validation if the value is `None` when initialisation is delayed."""
        if value is None:
            return value
        return handler(value)

851
852
853
854
855
856
857
858
859
860
861
862
863
    def __post_init__(self) -> None:
        count_none = self.custom_ops.count("none")
        count_all = self.custom_ops.count("all")
        assert count_none + count_all <= 1, "Can only specify 'none' or 'all'"

        # TODO(zou3519/luka): There are 2 issues with auto-functionalization V2:
        # 1. A bug in PyTorch, fixed in 2.7:
        #    https://github.com/pytorch/pytorch/issues/147924
        # 2. Custom passes (fusion) rely on auto-functionalization V1 and don't
        #    work with V2. Addressing this will take extra engineering effort
        #    and it is not yet a priority. RFC here:
        #    https://github.com/vllm-project/vllm/issues/14703

864
865
866
        KEY = "enable_auto_functionalized_v2"
        if KEY not in self.inductor_compile_config:
            self.inductor_compile_config[KEY] = False
867

868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
        # Tie inductor runtime assertions to debug logging mode.
        # These assertions add ~2ms overhead per forward pass on large
        # models (e.g., DeepSeek-R1 671B: ~340 assert_size_stride + ~60
        # assert_alignment calls per forward). PyTorch >= 2.12 has a
        # native fix (assert-once), so we only apply this workaround on
        # older versions. On torch < 2.12, enable asserts only when
        # VLLM_LOGGING_LEVEL=DEBUG. Users can still override explicitly
        # via --compilation-config '{"inductor_compile_config":
        # {"size_asserts": true, ...}}'.
        # See: https://github.com/pytorch/pytorch/issues/177719
        if not is_torch_equal_or_newer("2.12.0.dev"):
            enable_asserts = envs.VLLM_LOGGING_LEVEL == "DEBUG"
            for key in (
                "size_asserts",
                "alignment_asserts",
                "scalar_asserts",
            ):
                self.inductor_compile_config.setdefault(key, enable_asserts)

887
888
        for k, v in self.inductor_passes.items():
            if not isinstance(v, str):
889
890
891
892
                assert callable(v), f"pass {k} should be callable or a qualified name"
                self.inductor_compile_config[k] = (
                    v if isinstance(v, InductorPass) else CallableInductorPass(v)
                )
893
894
895
896
897
898
899
                continue

            # resolve function from qualified name
            names = v.split(".")
            module = ".".join(names[:-1])
            func_name = names[-1]
            func = __import__(module).__dict__[func_name]
900
901
902
            self.inductor_compile_config[k] = (
                func if isinstance(func, InductorPass) else CallableInductorPass(func)
            )
903

904
905
906
907
        if (
            self.pass_config.enable_qk_norm_rope_fusion
            and "+rotary_embedding" not in self.custom_ops
        ):
908
909
910
            # TODO(zhuhaoran): support rope native forward match and remove this.
            # Linked issue: https://github.com/vllm-project/vllm/issues/28042
            self.custom_ops.append("+rotary_embedding")
911
912
913
914
915
916
917
        if (
            self.pass_config.fuse_rope_kvcache
            and "+rotary_embedding" not in self.custom_ops
        ):
            # TODO(Rohan138): support rope native forward match and remove this.
            # Linked issue: https://github.com/vllm-project/vllm/issues/28042
            self.custom_ops.append("+rotary_embedding")
918

919
920
921
922
        if (
            is_torch_equal_or_newer("2.9.0.dev")
            and "combo_kernels" not in self.inductor_compile_config
            and "benchmark_combo_kernel" not in self.inductor_compile_config
923
924
            # (fixme @boyuan) combo kernel does not support cpu yet.
            and not current_platform.is_cpu()
925
926
927
928
929
930
        ):
            # use horizontal fusion, which is useful for fusing qk-norm and
            # qk-rope when query and key have different shapes.
            self.inductor_compile_config["combo_kernels"] = True
            self.inductor_compile_config["benchmark_combo_kernel"] = True

931
932
933
934
935
936
937
938
        if self.use_inductor_graph_partition and not is_torch_equal_or_newer(
            "2.9.0.dev"
        ):
            raise ValueError(
                "use_inductor_graph_partition is only "
                "supported with torch>=2.9.0.dev. Set "
                "use_inductor_graph_partition=False instead."
            )
939

940
        for op in self.custom_ops:
941
942
943
944
945
946
            if op[0] not in {"+", "-"} and op not in {"all", "none"}:
                raise ValueError(
                    f"Invalid syntax '{op}' for custom op, "
                    "must be 'all', 'none', '+op' or '-op' "
                    "(where 'op' is the registered op name)"
                )
947

948
        # Currently only eager and inductor backend are supported.
Jiayi Yan's avatar
Jiayi Yan committed
949
        # for piecewise compilation. Custom backends are not supported for
950
        # piecewise compilation. Update when more backends are supported.
951
        if self.mode == CompilationMode.VLLM_COMPILE and self.backend not in [
952
953
954
955
956
957
958
959
            "",
            "eager",
            "inductor",
        ]:
            raise ValueError(
                f"Invalid backend for piecewise compilation: {self.backend}"
            )

960
961
962
963
964
965
966
967
968
969
        # Validate encoder CUDA graph configuration
        if (
            self.cudagraph_mm_encoder
            and self.encoder_cudagraph_max_images_per_batch < 0
        ):
            raise ValueError(
                "encoder_cudagraph_max_images_per_batch must be "
                "non-negative (0 = auto-infer)"
            )

970
        if self.backend == "":
971
            self.backend = current_platform.get_compile_backend()
972

973
974
975
976
977
978
    def init_backend(
        self,
        vllm_config: "VllmConfig",
        prefix: str = "",
        is_encoder: bool = False,
    ) -> str | Callable:
979
980
981
982
        """
        Initialize the backend for the compilation config from a vllm config.
        Arguments:
            vllm_config: The vllm config to initialize the backend from.
983
984
985
            prefix: Cache directory prefix for this compiled module.
            is_encoder: Whether this module is used in an encoder (as
                opposed to a text backbone).
986
987
988
        Returns:
            The backend for the compilation config.
        """
989
        if self.mode is None:
990
            raise ValueError(
991
992
993
                "No compilation mode is set. This method should only be "
                "called via vllm config where the level is set if none is "
                "provided."
994
            )
995
996
        if self.mode == CompilationMode.NONE:
            raise ValueError("No compilation mode is set.")
997
998

        from torch._dynamo.backends.registry import list_backends
999

1000
        torch_backends = list_backends(exclude_tags=tuple())
1001
1002
1003
1004
        if self.mode in [
            CompilationMode.STOCK_TORCH_COMPILE,
            CompilationMode.DYNAMO_TRACE_ONCE,
        ]:
1005
1006
1007
1008
            if self.backend in torch_backends:
                return self.backend
            return resolve_obj_by_qualname(self.backend)

1009
        assert self.mode == CompilationMode.VLLM_COMPILE
1010
        if self.backend not in ["eager", "inductor"]:
1011
            logger.info("Using OOT custom backend for compilation.")
1012
1013

        from vllm.compilation.backends import VllmBackend
1014

1015
        return VllmBackend(vllm_config, prefix=prefix, is_encoder=is_encoder)
1016

1017
1018
1019
1020
1021
    def post_init_cudagraph_sizes(self) -> None:
        """To complete the initialization after cudagraph related
        configs are set. This includes:
        - initialize compile_sizes
        """
1022

1023
        computed_compile_sizes: list[int] = []
1024
1025
1026
1027
1028
        if self.compile_sizes is not None:
            # de-duplicate the sizes provided by the config
            self.compile_sizes = list(set(self.compile_sizes))
            for x in self.compile_sizes:
                if isinstance(x, str):
1029
1030
                    assert x == "cudagraph_capture_sizes", (
                        "Unrecognized size type in compile_sizes, "
1031
                        f"expect 'cudagraph_capture_sizes', got {x}"
1032
                    )
1033
1034
1035
1036
1037
1038
                    computed_compile_sizes.extend(self.cudagraph_capture_sizes)
                else:
                    assert isinstance(x, int)
                    computed_compile_sizes.append(x)
        self.compile_sizes = computed_compile_sizes  # type: ignore

1039
1040
1041
1042
        # make sure the sizes are in ascending order
        self.cudagraph_capture_sizes.sort()
        if self.cudagraph_capture_sizes:
            assert self.cudagraph_capture_sizes[-1] == self.max_cudagraph_capture_size
1043

1044
    def set_splitting_ops_for_v1(
1045
        self, all2all_backend: str, data_parallel_size: int = 1
1046
    ):
1047
1048
1049
1050
1051
1052
1053
        # To compatible with OOT hardware plugin platform (for example vllm-ascend)
        # which currently only supports sequence parallelism in eager mode.
        if self.mode != CompilationMode.VLLM_COMPILE:
            if self.splitting_ops is None:
                self.splitting_ops = []
            return

1054
        if self.pass_config.fuse_attn_quant and not self.use_inductor_graph_partition:
1055
            self.set_splitting_ops_for_attn_fusion()
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
        else:
            if self.splitting_ops is None:
                # NOTE: When using full cudagraph, instead of setting an empty
                # list and capture the full cudagraph inside the flattened fx
                # graph, we keep the piecewise fx graph structure but capture
                # the full cudagraph outside the fx graph. This reduces some
                # cpu overhead when the runtime batch_size is not cudagraph
                # captured. see https://github.com/vllm-project/vllm/pull/20059
                # for details. Make a copy to avoid mutating the class-level
                # list via reference.
                self.splitting_ops = list(self._attention_ops)
1067
1068
1069
1070
1071
1072
1073

                # unified_kv_cache_update has a string param that prevents Inductor
                # from reusing piecewise graphs. Remove it from the compiled graph.
                # This has the side-effect of excluding cache from cudagraphs but
                # that doesn't seem to affect performance.
                # https://github.com/vllm-project/vllm/issues/33267
                if not self.use_inductor_graph_partition:
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
                    if self.pass_config.fuse_rope_kvcache:
                        logger.warning_once(
                            "fuse_rope_kvcache is enabled, but splitting_ops is None "
                            "and Inductor graph partition is not enabled."
                            "Disabling fuse_rope_kvcache."
                            "Please either set splitting_ops to an empty list []"
                            "or set use_inductor_graph_partition to True "
                            "to enable RoPE+KV cache fusion."
                        )
                        self.pass_config.fuse_rope_kvcache = False
1084
                    self.splitting_ops.append("vllm::unified_kv_cache_update")
1085
                    self.splitting_ops.append("vllm::unified_mla_kv_cache_update")
1086

1087
            elif len(self.splitting_ops) == 0:
1088
1089
1090
1091
1092
                if (
                    self.cudagraph_mode == CUDAGraphMode.PIECEWISE
                    or self.cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE
                ):
                    logger.warning_once(
1093
                        "Using piecewise cudagraph with empty splitting_ops"
1094
                    )
1095
1096
                if self.cudagraph_mode == CUDAGraphMode.PIECEWISE:
                    logger.warning_once(
1097
1098
                        "Piecewise compilation with empty splitting_ops does not "
                        "contain piecewise cudagraph. Setting cudagraph_"
1099
1100
1101
1102
1103
1104
1105
1106
                        "mode to NONE. Hint: If you are using attention "
                        "backends that support cudagraph, consider manually "
                        "setting cudagraph_mode to FULL or FULL_DECODE_ONLY "
                        "to enable full cudagraphs."
                    )
                    self.cudagraph_mode = CUDAGraphMode.NONE
                elif self.cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE:
                    logger.warning_once(
1107
1108
                        "Piecewise compilation with empty splitting_ops does "
                        "not contain piecewise cudagraph. Setting "
1109
1110
1111
1112
1113
                        "cudagraph_mode to FULL."
                    )
                    self.cudagraph_mode = CUDAGraphMode.FULL
                self.splitting_ops = []

1114
1115
        # Disable CUDA graphs for DeepEP high-throughput since its not CG compatible
        if (
1116
1117
            all2all_backend == "deepep_high_throughput"
            and data_parallel_size > 1
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
            and self.cudagraph_mode != CUDAGraphMode.NONE
        ):
            # TODO: Piecewise Cuda graph might be enabled
            # if torch compile cache key issue fixed
            # See https://github.com/vllm-project/vllm/pull/25093
            logger.info(
                "DeepEP: Disabling CUDA Graphs since DeepEP high-throughput kernels "
                "are optimized for prefill and are incompatible with CUDA Graphs. "
                "In order to use CUDA Graphs for decode-optimized workloads, "
                "use --all2all-backend with another option, such as "
1128
                "deepep_low_latency or allgather_reducescatter."
1129
            )
1130
            self.cudagraph_mode = CUDAGraphMode.NONE
1131
1132

    def set_splitting_ops_for_attn_fusion(self):
1133
        assert self.pass_config.fuse_attn_quant
1134
1135
1136
1137
        if self.splitting_ops is None:
            self.splitting_ops = []
            if self.cudagraph_mode.has_piecewise_cudagraphs():
                logger.warning_once(
1138
                    "fuse_attn_quant is incompatible with piecewise "
1139
1140
1141
1142
1143
1144
1145
1146
                    "cudagraph when use_inductor_graph_partition is off. "
                    "In this case, splitting_ops will be set to empty "
                    "list, and cudagraph_mode will be set to FULL. "
                    "Please ensure you are using attention backends that "
                    "support cudagraph or set cudagraph_mode to NONE "
                    "explicitly if encountering any problems."
                )
                self.cudagraph_mode = CUDAGraphMode.FULL
1147
1148

        assert not self.splitting_ops_contain_attention(), (
1149
            "attention ops should not be in splitting_ops when fuse_attn_quant is True"
1150
        )
1151
1152
1153

    def splitting_ops_contain_attention(self) -> bool:
        return self.splitting_ops is not None and all(
1154
1155
            op in self.splitting_ops for op in self._attention_ops
        )
1156

1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
    def splitting_ops_contain_kv_cache_update(self) -> bool:
        # when using Dynamo partition while splitting ops is None
        # and attn+quant fusion disabled, the kv_cache_update_ops are
        # appended to splitting_ops in set_splitting_ops_for_v1 due to
        # https://github.com/vllm-project/vllm/issues/33267
        # In this case, we return True if the kv_cache_update_ops
        # are not in the splitting_ops yet, but will subsequently
        # be added to splitting_ops.
        if (
            not self.use_inductor_graph_partition
            and self.splitting_ops is None
            and not self.pass_config.fuse_attn_quant
        ):
            return True

        kv_cache_update_ops = [
            "vllm::unified_kv_cache_update",
            "vllm::unified_mla_kv_cache_update",
        ]
        return self.splitting_ops is not None and all(
            op in self.splitting_ops for op in kv_cache_update_ops
        )

1180
    def is_attention_compiled_piecewise(self) -> bool:
1181
1182
        if not self.splitting_ops_contain_attention():
            return False
1183

1184
1185
        if not self.use_inductor_graph_partition:
            # Dynamo-level FX split case
1186
            return self.mode == CompilationMode.VLLM_COMPILE
1187

1188
        # Inductor partition case
1189
        return self.backend == "inductor" and self.mode != CompilationMode.NONE
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205

    def custom_op_log_check(self):
        """
        This method logs the enabled/disabled custom ops and checks that the
        passed custom_ops field only contains relevant ops.
        It is called at the end of set_current_vllm_config,
        after the custom ops have been instantiated.
        """

        if len(self.enabled_custom_ops) + len(self.disabled_custom_ops) == 0:
            logger.debug("No custom ops found in model.")
            return

        logger.debug("enabled custom ops: %s", self.enabled_custom_ops)
        logger.debug("disabled custom ops: %s", self.disabled_custom_ops)

1206
        all_ops_in_model = self.enabled_custom_ops | self.disabled_custom_ops
1207
1208
1209
1210
        for op in self.custom_ops:
            if op in {"all", "none"}:
                continue

1211
1212
1213
            assert op[0] in {"+", "-"}, (
                "Invalid custom op syntax (should be checked during init)"
            )
1214
1215
1216
1217

            # check if op name exists in model
            op_name = op[1:]
            if op_name not in all_ops_in_model:
1218
                from vllm.model_executor.custom_op import op_registry
1219
1220
1221

                # Does op exist at all or is it just not present in this model?
                # Note: Only imported op classes appear in the registry.
1222
1223
                missing_str = (
                    "doesn't exist (or wasn't imported/registered)"
1224
                    if op_name not in op_registry
1225
                    else "not present in model"
1226
                )
1227

1228
1229
1230
1231
1232
1233
1234
1235
                enable_str = "enabling" if op[0] == "+" else "disabling"
                logger.warning_once(
                    "Op '%s' %s, %s with '%s' has no effect",
                    op_name,
                    missing_str,
                    enable_str,
                    op,
                )
1236

1237
1238
1239
1240
1241
1242
1243
    def is_custom_op_enabled(self, op: str) -> bool:
        if "all" in self.custom_ops:
            return f"-{op}" not in self.custom_ops

        assert "none" in self.custom_ops
        return f"+{op}" in self.custom_ops

1244
1245
1246
1247
    def adjust_cudagraph_sizes_for_spec_decode(
        self, uniform_decode_query_len: int, tensor_parallel_size: int
    ):
        multiple_of = uniform_decode_query_len
1248
        if tensor_parallel_size > 1 and self.pass_config.enable_sp:
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
            multiple_of = max(uniform_decode_query_len, tensor_parallel_size)
            if (
                multiple_of % uniform_decode_query_len != 0
                or multiple_of % tensor_parallel_size != 0
            ):
                raise ValueError(
                    f"Can't determine cudagraph shapes that are both a "
                    f"multiple of {uniform_decode_query_len} "
                    f"(num_speculative_tokens + 1) required by spec-decode "
                    f"and {tensor_parallel_size} (tensor_parallel_size) "
                    f"required by sequence parallelism please adjust "
                    f"num_speculative_tokens or disable sequence parallelism"
                )

        if not self.cudagraph_capture_sizes or multiple_of <= 1:
            return

        assert self.max_cudagraph_capture_size is not None
        rounded_sizes = sorted(
            set(
                round_up(size, multiple_of)
                for size in self.cudagraph_capture_sizes
                if round_up(size, multiple_of) <= self.max_cudagraph_capture_size
            )
        )

1275
1276
1277
1278
        if len(rounded_sizes) == 0 and multiple_of <= self.max_cudagraph_capture_size:
            # if one valid but would be round_down use that
            rounded_sizes = [multiple_of]

1279
        if len(rounded_sizes) == 0:
1280
1281
1282
1283
1284
1285
            raise ValueError(
                f"No valid cudagraph sizes after rounding to multiple of {multiple_of} "
                f"(num_speculative_tokens + 1 or tp if sequence parallelism is enabled)"
                f" please adjust num_speculative_tokens ({uniform_decode_query_len - 1}"
                f") or max_cudagraph_capture_size ({self.max_cudagraph_capture_size})"
                f" or cudagraph_capture_sizes ({self.cudagraph_capture_sizes})"
1286
1287
1288
1289
1290
            )

        self.max_cudagraph_capture_size = rounded_sizes[-1]
        self.cudagraph_capture_sizes = rounded_sizes

1291
1292
    def get_compile_ranges(self) -> list[Range]:
        """Get the compile ranges for the compilation config."""
1293
        if self.compile_ranges_endpoints is None:
1294
            return []
1295
        endpoints = sorted(set(self.compile_ranges_endpoints))
1296
        return [Range(s + 1, e) for s, e in zip([0] + endpoints[:-1], endpoints)]