compilation.py 43.3 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 asdict, field
8
from pathlib import Path
9
from typing import TYPE_CHECKING, Any, ClassVar, Literal
10

11
from pydantic import TypeAdapter, field_validator
12
13
from pydantic.dataclasses import dataclass

14
import vllm.envs as envs
15
16
17
from vllm.compilation.inductor_pass import CallableInductorPass, InductorPass
from vllm.config.utils import config
from vllm.logger import init_logger
18
from vllm.platforms import current_platform
19
from vllm.utils.import_utils import resolve_obj_by_qualname
20
from vllm.utils.math_utils import round_up
21
from vllm.utils.torch_utils import is_torch_equal_or_newer
22
23

if TYPE_CHECKING:
24
    from vllm.config import VllmConfig
25
26
27
28
29
30
else:
    VllmConfig = object

logger = init_logger(__name__)


31
class CompilationMode(enum.IntEnum):
32
33
34
35
36
37
38
39
40
41
42
43
44
    """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."""
45
46


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

53
54
55
56
57
58
    NONE = 0
    PIECEWISE = 1
    FULL = 2
    FULL_DECODE_ONLY = (FULL, NONE)
    FULL_AND_PIECEWISE = (FULL, PIECEWISE)

59
60
    def decode_mode(self) -> "CUDAGraphMode":
        return CUDAGraphMode(self.value[0]) if self.separate_routine() else self
61

62
63
    def mixed_mode(self) -> "CUDAGraphMode":
        return CUDAGraphMode(self.value[1]) if self.separate_routine() else self
64

65
66
67
68
69
70
    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

71
    def requires_piecewise_compilation(self) -> bool:
72
        return self.has_mode(CUDAGraphMode.PIECEWISE)
73

74
75
    def max_cudagraph_mode(self) -> "CUDAGraphMode":
        return CUDAGraphMode(max(self.value)) if self.separate_routine() else self
76
77
78
79

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

80
81
82
    def has_piecewise_cudagraphs(self) -> bool:
        return self.requires_piecewise_compilation()

83
84
85
    def separate_routine(self) -> bool:
        return isinstance(self.value, tuple)

86
    def valid_runtime_modes(self) -> bool:
87
        return self in [CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL]
88

89
90
91
    def __str__(self) -> str:
        return self.name

92

93
94
95
96
97
98
99
100
101
@config
@dataclass
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
    the `PassManager` is set as a property of config."""

102
    enable_fusion: bool = False
103
104
105
    """Whether to enable the custom fusion (RMSNorm/SiluMul+quant) pass."""
    enable_attn_fusion: bool = False
    """Whether to enable the custom attention+quant fusion pass."""
106
    enable_noop: bool = False
107
108
109
110
111
112
113
    """Whether to enable the custom no-op elimination pass."""
    enable_sequence_parallelism: bool = False
    """Whether to enable sequence parallelism."""
    enable_async_tp: bool = False
    """Whether to enable async TP."""
    enable_fi_allreduce_fusion: bool = False
    """Whether to enable flashinfer allreduce fusion."""
114
115
116
117
    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.
118
    Unspecified will fallback to default values
119
120
121
122
123
124
125
126
127
128
129
130
131
    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"""
132
133
    enable_qk_norm_rope_fusion: bool = False
    """Whether to enable the fused Q/K RMSNorm + RoPE pass."""
134
135
136

    # TODO(luka) better pass enabling system.

137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
    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
        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]:
        from vllm.compilation.collective_fusion import FI_ALLREDUCE_FUSION_MAX_SIZE_MB
        from vllm.platforms import current_platform

        if not current_platform.is_cuda():
            return {}
        return FI_ALLREDUCE_FUSION_MAX_SIZE_MB.get(
            current_platform.get_device_capability().to_int(), {}
        )

162
    def compute_hash(self) -> str:
163
164
165
166
167
168
169
170
171
172
173
174
        """
        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.
        """
        return InductorPass.hash_dict(asdict(self))

    def __post_init__(self) -> None:
        if not self.enable_noop:
            if self.enable_fusion:
                logger.warning_once(
                    "Fusion enabled but reshape elimination disabled. "
175
176
                    "RMSNorm/SiluMul + quant (fp8) fusion might not work"
                )
177
178
179
            if self.enable_attn_fusion:
                logger.warning_once(
                    "Fusion enabled but reshape elimination disabled. "
180
181
                    "Attention + quant (fp8) fusion might not work"
                )
182
183
184
185
186
            if self.enable_fi_allreduce_fusion:
                logger.warning_once(
                    "Fusion enabled but reshape elimination disabled. "
                    "Allreduce + rms norm + quant (fp8) fusion might not work"
                )
187
        if self.enable_qk_norm_rope_fusion and not current_platform.is_cuda_alike():
188
189
            logger.warning_once(
                "QK Norm + RoPE fusion enabled but the current platform is not "
190
                "CUDA or ROCm. The fusion will be disabled."
191
192
            )
            self.enable_qk_norm_rope_fusion = False
193
194


195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
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
@dataclass
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.
    """

    # TODO add a debug mode to fail

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

        from vllm.config.utils import get_hash_factors, hash_factors

        factors = get_hash_factors(self, {})
        return hash_factors(factors)


243
244
245
246
247
248
@config
@dataclass
class CompilationConfig:
    """Configuration for compilation. It has three parts:

    - Top-level Compilation control:
249
        - [`mode`][vllm.config.CompilationConfig.mode]
250
251
252
253
254
        - [`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]
255
        - [`compile_mm_encoder`][vllm.config.CompilationConfig.compile_mm_encoder]
256
    - CudaGraph capture:
257
        - [`cudagraph_mode`][vllm.config.CompilationConfig.cudagraph_mode]
258
259
        - [`cudagraph_capture_sizes`]
        [vllm.config.CompilationConfig.cudagraph_capture_sizes]
260
261
        - [`max_cudagraph_capture_size`]
        [vllm.config.CompilationConfig.max_cudagraph_capture_size]
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
        - [`cudagraph_num_of_warmups`]
        [vllm.config.CompilationConfig.cudagraph_num_of_warmups]
        - [`cudagraph_copy_inputs`]
        [vllm.config.CompilationConfig.cudagraph_copy_inputs]
    - Inductor compilation:
        - [`use_inductor`][vllm.config.CompilationConfig.use_inductor]
        - [`compile_sizes`][vllm.config.CompilationConfig.compile_sizes]
        - [`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.
    """
284

285
    # Top-level Compilation control
286
    level: int | None = None
287
288
289
290
291
292
    """
    Level is deprecated and will be removed in the next release,
    either 0.12.0 or 0.11.2 whichever is soonest.
    Please use mode. Currently all levels are mapped to mode.
    """
    # Top-level Compilation control
293
    mode: CompilationMode | None = None
294
295
296
297
298
299
300
301
302
303
304
305
306
    """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."""
307
    debug_dump_path: Path | None = None
308
309
310
311
312
    """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."""
313
314
315
316
317
318
319
320
321
    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.
    """
322
    backend: str = ""
323
324
    """The backend for compilation. It needs to be a string:

325
326
    - "" (empty string): use the default backend ("inductor" on CUDA-alike
    platforms).
327
328
329
330
331
    - "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
332
    distributed setting. When the compilation mode is 1 or 2, the backend is
333
    used for the compilation directly (it sees the whole graph). When the
334
335
336
    compilation mode is 3, the backend supports both whole graph and piecewise 
    compilation, available backends include eager, inductor, and custom backends, 
    the latter of which can be defined via `get_compile_backend`. Furthermore,
337
    compilation is only piecewise if splitting ops is set accordingly and
338
    use_inductor_graph_partition is off. Note that the default options for
339
340
    splitting ops are sufficient for piecewise compilation.
    """
341
342
343
344
345
346
347
348
349
350
    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
351
    disabled when running with Inductor: mode>=VLLM_COMPILE and use_inductor=True.
352
    Inductor generates (fused) Triton kernels for disabled custom ops."""
353
    splitting_ops: list[str] | None = None
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
    """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)."""
371
    compile_mm_encoder: bool = False
Harry Mellor's avatar
Harry Mellor committed
372
    """Whether or not to compile the multimodal encoder.
373
    Currently, this only works for `Qwen2_5_vl` on selected platforms.
374
    Disabled by default until more models are supported/tested to work."""
375
376

    # Inductor capture
377
378
379
380
381
382
    use_inductor: bool | None = None
    """
    Whether to use inductor compilation.

    This flag is deprecated and will be removed in the next release 0.12.0.
    Please use the 'backend' option instead.
383
384
385
386
387
388
389

    - False: inductor compilation is not used. graph runs in eager
        (custom_ops enabled by default).
    - True: inductor compilation is used (custom_ops disabled by default).
        One graph for symbolic shape and one graph per size in compile_sizes
        are compiled using configurations in inductor_compile_config.

390
    This setting is ignored if mode<VLLM_COMPILE.
391
392
393
394

    For future compatibility:
    If use_inductor is True, backend="inductor" otherwise backend="eager".
    """
395
    compile_sizes: list[int | str] | None = None
396
397
398
    """Sizes to compile for inductor. In addition
    to integers, it also supports "cudagraph_capture_sizes" to
    specify the sizes for cudagraph capture."""
399

400
401
402
    inductor_compile_config: dict = field(default_factory=dict)
    """Additional configurations for inductor.
    - None: use default configurations."""
403

404
405
406
407
408
409
410
411
    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
412
    cudagraph_mode: CUDAGraphMode | None = None
413
    """
Harry Mellor's avatar
Harry Mellor committed
414
415
    The mode of the cudagraph:

416
    - NONE, no cudagraph capture.
417
    - PIECEWISE.
418
419
    - FULL.
    - FULL_DECODE_ONLY.
420
    - FULL_AND_PIECEWISE. (v1 default)
421
422

    PIECEWISE mode build piecewise cudagraph only, keeping the cudagraph
co63oc's avatar
co63oc committed
423
    incompatible ops (i.e. some attention ops) outside the cudagraph
424
425
426
427
428
    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.
429

430
431
432
433
    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.
434

435
436
    FULL_AND_PIECEWISE mode: Capture full cudagraph for decode batches and
    piecewise cudagraph for prefill and mixed prefill-decode batches.
437
    This is the most performant mode for most models and is the default.
438
439

    Currently, the cudagraph mode is only used for the v1 engine.
440
441
    Note that the cudagraph logic is generally orthogonal to the
    compilation logic. While piecewise cudagraphs require piecewise
442
    compilation (mode=VLLM_COMPILE and non-empty splitting_ops), full
443
    cudagraphs are supported with and without compilation.
444
445

    Warning: This flag is new and subject to change in addition
446
447
    more modes may be added.
    """
448
449
450
451
452
    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."""
453
    cudagraph_capture_sizes: list[int] | None = None
454
455
456
457
458
459
460
461
    """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
462
    internally managed buffer. Default is False.
463
464
    Note that this flag is only effective when cudagraph_mode is PIECEWISE.
    """
465
466
467
468
469
470
471
472
    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.
    """
473

474
475
476
477
478
479
480
481
    use_inductor_graph_partition: bool = False
    """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
482
    register the custom op.
483
484
485
486
487
488
489
490
491
492
493

    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.
    """

494
495
496
    pass_config: PassConfig = field(default_factory=PassConfig)
    """Custom inductor passes, see PassConfig for more details"""

497
498
    max_cudagraph_capture_size: int | None = field(default=None)
    """The maximum cudagraph capture size.
499
500

    If cudagraph_capture_sizes is specified, this will be set to the largest
501
502
503
504
505
506
507
508
    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,
509
    512) by default. This voids OOM in tight memory scenarios with small
510
511
512
    max_num_seqs, and prevents capture of many large graphs (>512) that would
    greatly increase startup time with limited performance benefit.
    """
513
514
515
516
517
518

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

519
520
    local_cache_dir: str = field(default=None, init=False)  # type: ignore
    """local cache dir for each rank"""
521

522
523
    bs_to_padded_graph_size: list[int] = field(
        default=None,  # type: ignore
524
525
        init=False,
    )
526
527
    """optimization:
    Intuitively, bs_to_padded_graph_size should be dict[int, int].
528
    since we know all keys are in a range [0, max_cudagraph_capture_size],
529
530
531
    we can optimize it to list[int] for better lookup performance."""

    # keep track of enabled and disabled custom ops
532
    enabled_custom_ops: Counter[str] = field(default_factory=Counter, init=False)
533
    """custom ops that are enabled"""
534
    disabled_custom_ops: Counter[str] = field(default_factory=Counter, init=False)
535
536
537
538
539
540
    """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"""

541
    static_forward_context: dict[str, Any] = field(default_factory=dict, init=False)
542
543
544
545
    """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."""

546
    # Attention ops; used for piecewise cudagraphs
547
    # Use PyTorch operator format: "namespace::name"
548
    _attention_ops: ClassVar[list[str]] = [
549
550
551
552
553
554
555
556
557
        "vllm::unified_attention",
        "vllm::unified_attention_with_output",
        "vllm::unified_mla_attention",
        "vllm::unified_mla_attention_with_output",
        "vllm::mamba_mixer2",
        "vllm::mamba_mixer",
        "vllm::short_conv",
        "vllm::linear_attention",
        "vllm::plamo2_mamba_mixer",
558
        "vllm::gdn_attention_core",
559
        "vllm::kda_attention",
560
        "vllm::sparse_attn_indexer",
561
562
    ]

563
564
565
566
567
568
569
570
    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.
        """
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
        # 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",
            "bs_to_padded_graph_size",
            "traced_files",
            "compilation_time",
            "static_forward_context",
            "pass_config",  # handled separately below
        }

        from vllm.config.utils import get_hash_factors, hash_factors

        factors = get_hash_factors(self, ignored_factors)
590

591
592
        factors["pass_config"] = self.pass_config.compute_hash()
        return hash_factors(factors)
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614

    def __repr__(self) -> str:
        exclude = {
            "static_forward_context": True,
            "enabled_custom_ops": True,
            "disabled_custom_ops": True,
            "compilation_time": True,
            "bs_to_padded_graph_size": 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

615
616
617
        config = TypeAdapter(CompilationConfig).dump_python(
            self, exclude=exclude, exclude_unset=True
        )
618
619

        return str(config)
620
621
622

    __str__ = __repr__

623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
    @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

644
645
646
    @field_validator("cudagraph_mode", mode="before")
    @classmethod
    def validate_cudagraph_mode_before(cls, value: Any) -> Any:
647
        """Enable parsing of the `cudagraph_mode` enum type from string."""
648
649
650
651
        if isinstance(value, str):
            return CUDAGraphMode[value.upper()]
        return value

652
653
654
655
656
657
658
659
    @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

660
661
662
663
664
665
666
667
668
669
    @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

670
    def __post_init__(self) -> None:
671
672
673
674
675
676
677
678
679
680
681
        if self.level is not None:
            logger.warning(
                "Level is deprecated and will be removed in the next release,"
                "either 0.12.0 or 0.11.2 whichever is soonest."
                "Use mode instead."
                "If both level and mode are given,"
                "only mode will be used."
            )
            if self.mode is None:
                self.mode = self.level

682
683
684
685
686
687
688
689
690
691
692
693
694
        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

        if is_torch_equal_or_newer("2.6"):
695
            KEY = "enable_auto_functionalized_v2"
696
697
698
699
700
            if KEY not in self.inductor_compile_config:
                self.inductor_compile_config[KEY] = False

        for k, v in self.inductor_passes.items():
            if not isinstance(v, str):
701
702
703
704
                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)
                )
705
706
707
708
709
710
711
                continue

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

716
717
718
719
720
        if self.pass_config.enable_qk_norm_rope_fusion:
            # 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")

721
722
723
724
        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
725
726
            # (fixme @boyuan) combo kernel does not support cpu yet.
            and not current_platform.is_cpu()
727
728
729
730
731
732
        ):
            # 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

733
734
735
736
737
738
739
740
        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."
            )
741

742
        for op in self.custom_ops:
743
744
745
746
747
748
            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)"
                )
749

750
751
752
        # Currently only eager and inductor backend are supported.
        # for piecewise compilation. Custom backends are not suppported for
        # piecewise compilation. Update when more backends are supported.
753
        if self.mode == CompilationMode.VLLM_COMPILE and self.backend not in [
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
            "",
            "eager",
            "inductor",
        ]:
            raise ValueError(
                f"Invalid backend for piecewise compilation: {self.backend}"
            )

        if self.use_inductor is not None:
            logger.warning_once(
                "The 'use_inductor' flag is deprecated and will be "
                "removed in the next release (v0.12.0). "
                "Please use the 'backend' option instead.",
            )
            self.backend = "inductor" if self.use_inductor else "eager"

        if self.backend == "":
771
            self.backend = current_platform.get_compile_backend()
772

773
    def init_backend(self, vllm_config: "VllmConfig") -> str | Callable:
774
775
776
777
778
779
780
        """
        Initialize the backend for the compilation config from a vllm config.
        Arguments:
            vllm_config: The vllm config to initialize the backend from.
        Returns:
            The backend for the compilation config.
        """
781
        if self.mode is None:
782
            raise ValueError(
783
                "No compilation mode is set. This method should only be \
784
785
786
                called via vllm config where the level is set if none is \
                provided."
            )
787
788
        if self.mode == CompilationMode.NONE:
            raise ValueError("No compilation mode is set.")
789
790

        from torch._dynamo.backends.registry import list_backends
791

792
        torch_backends = list_backends(exclude_tags=tuple())
793
794
795
796
        if self.mode in [
            CompilationMode.STOCK_TORCH_COMPILE,
            CompilationMode.DYNAMO_TRACE_ONCE,
        ]:
797
798
799
800
            if self.backend in torch_backends:
                return self.backend
            return resolve_obj_by_qualname(self.backend)

801
        assert self.mode == CompilationMode.VLLM_COMPILE
802
        if self.backend not in ["eager", "inductor"]:
803
            logger.info("Using OOT custom backend for compilation.")
804
805

        from vllm.compilation.backends import VllmBackend
806

807
808
        # TODO[@lucaskabela]: See if we can forward prefix
        # https://github.com/vllm-project/vllm/issues/27045
809
810
        return VllmBackend(vllm_config)

811
812
813
814
815
816
    def post_init_cudagraph_sizes(self) -> None:
        """To complete the initialization after cudagraph related
        configs are set. This includes:
        - initialize compile_sizes
        - pre-compute the mapping bs_to_padded_graph_size
        """
817
818
819
820
821
822
823

        computed_compile_sizes = []
        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):
824
825
                    assert x == "cudagraph_capture_sizes", (
                        "Unrecognized size type in compile_sizes, "
826
                        f"expect 'cudagraph_capture_sizes', got {x}"
827
                    )
828
829
830
831
832
833
                    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

834
835
836
837
        # 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
838

839
840
        # May get recomputed in the model runner if adjustment is needed for spec-decode
        self.compute_bs_to_padded_graph_size()
841
842

    def set_splitting_ops_for_v1(self):
843
844
845
        # NOTE: this function needs to be called only when mode is
        # CompilationMode.VLLM_COMPILE
        assert self.mode == CompilationMode.VLLM_COMPILE, (
846
            "set_splitting_ops_for_v1 should only be called when "
847
            "mode is CompilationMode.VLLM_COMPILE"
848
        )
849

850
851
852
853
854
855
856
857
        if self.use_inductor_graph_partition:
            self.set_splitting_ops_for_inductor_graph_partition()
            return

        if self.pass_config.enable_attn_fusion:
            # here use_inductor_graph_partition is False
            self.set_splitting_ops_for_attn_fusion()
            return
858

859
        if self.splitting_ops is None:
860
861
862
863
864
865
866
867
868
            # 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)
869
        elif len(self.splitting_ops) == 0:
870
            logger.warning_once("Using piecewise compilation with empty splitting_ops")
871
            if self.cudagraph_mode == CUDAGraphMode.PIECEWISE:
872
                logger.warning_once(
873
                    "Piecewise compilation with empty splitting_ops do not"
874
875
876
877
                    "contains piecewise cudagraph. Setting cudagraph_"
                    "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 "
878
879
                    "full cudagraphs."
                )
880
881
882
883
884
                self.cudagraph_mode = CUDAGraphMode.NONE
            elif self.cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE:
                logger.warning_once(
                    "Piecewise compilation with empty splitting_ops do not "
                    "contains piecewise cudagraph. Setting cudagraph_mode "
885
886
                    "to FULL."
                )
887
888
                self.cudagraph_mode = CUDAGraphMode.FULL
            self.splitting_ops = []
889
890
891

    def set_splitting_ops_for_inductor_graph_partition(self):
        assert self.use_inductor_graph_partition
892
893
        if self.splitting_ops is None:
            self.splitting_ops = list(self._attention_ops)
894
895
896

    def set_splitting_ops_for_attn_fusion(self):
        assert self.pass_config.enable_attn_fusion
897
898
899
900
901
902
903
904
905
906
907
908
909
        if self.splitting_ops is None:
            self.splitting_ops = []
            if self.cudagraph_mode.has_piecewise_cudagraphs():
                logger.warning_once(
                    "enable_attn_fusion is incompatible with piecewise "
                    "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
910
911
912

        assert not self.splitting_ops_contain_attention(), (
            "attention ops should not be in splitting_ops "
913
914
            "when enable_attn_fusion is True"
        )
915
916
917

    def splitting_ops_contain_attention(self) -> bool:
        return self.splitting_ops is not None and all(
918
919
            op in self.splitting_ops for op in self._attention_ops
        )
920
921

    def is_attention_compiled_piecewise(self) -> bool:
922
923
        if not self.splitting_ops_contain_attention():
            return False
924

925
926
        if not self.use_inductor_graph_partition:
            # Dynamo-level FX split case
927
            return self.mode == CompilationMode.VLLM_COMPILE
928

929
        # Inductor partition case
930
        return self.backend == "inductor" and self.mode != CompilationMode.NONE
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946

    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)

947
        all_ops_in_model = self.enabled_custom_ops | self.disabled_custom_ops
948
949
950
951
        for op in self.custom_ops:
            if op in {"all", "none"}:
                continue

952
953
954
            assert op[0] in {"+", "-"}, (
                "Invalid custom op syntax (should be checked during init)"
            )
955
956
957
958
959
960
961
962

            # check if op name exists in model
            op_name = op[1:]
            if op_name not in all_ops_in_model:
                from vllm.model_executor.custom_op import CustomOp

                # Does op exist at all or is it just not present in this model?
                # Note: Only imported op classes appear in the registry.
963
964
965
                missing_str = (
                    "doesn't exist (or wasn't imported/registered)"
                    if op_name not in CustomOp.op_registry
966
                    else "not present in model"
967
                )
968

969
970
971
972
973
974
975
976
                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,
                )
977
978
979
980
981

    def adjust_cudagraph_sizes_for_spec_decode(
        self, uniform_decode_query_len: int, tensor_parallel_size: int
    ):
        multiple_of = uniform_decode_query_len
982
        if tensor_parallel_size > 1 and self.pass_config.enable_sequence_parallelism:
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
            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
            )
        )

1009
1010
1011
1012
        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]

1013
        if len(rounded_sizes) == 0:
1014
1015
1016
1017
1018
1019
            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})"
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
            )

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

        # Recompute after adjusting the cudagraph sizes
        self.compute_bs_to_padded_graph_size()

    def compute_bs_to_padded_graph_size(self):
        # pre-compute the mapping from batch size to padded graph size
        self.bs_to_padded_graph_size = [
            0 for i in range(self.max_cudagraph_capture_size + 1)
        ]
        for end, start in zip(
            self.cudagraph_capture_sizes + [self.max_cudagraph_capture_size + 1],
            [0] + self.cudagraph_capture_sizes,
        ):
            for bs in range(start, end):
                if bs == start:
                    self.bs_to_padded_graph_size[bs] = start
                else:
                    self.bs_to_padded_graph_size[bs] = end