| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import contextlib |
| 4 | from collections.abc import Callable, Iterable |
| 5 | from dataclasses import asdict, fields |
| 6 | from typing import TYPE_CHECKING, Any, Literal |
| 7 | |
| 8 | from pydantic import Field, field_validator |
| 9 | |
| 10 | from vllm.config.utils import config, get_hash_factors, hash_factors |
| 11 | from vllm.logger import init_logger |
| 12 | |
| 13 | if TYPE_CHECKING: |
| 14 | from vllm.config import VllmConfig |
| 15 | |
| 16 | logger = init_logger(__name__) |
| 17 | |
| 18 | |
| 19 | @config |
| 20 | class IrOpPriorityConfig: |
| 21 | """Configuration for vLLM IR op priority for dispatching/lowering during the |
| 22 | forward pass. Each member is a list of strings, which will be installed |
| 23 | in worker init via vllm.ir.ops.<op_name>.set_default(). |
| 24 | A single comma-separated string is accepted as well, |
| 25 | |
| 26 | If specified manually, platform defaults will be appended to the lists. |
| 27 | See KernelConfig.set_platform_defaults(). |
| 28 | """ |
| 29 | |
| 30 | rms_norm: list[str] = Field(default_factory=list) |
| 31 | """Priority list for vllm.ir.ops.rms_norm""" |
| 32 | |
| 33 | fused_add_rms_norm: list[str] = Field(default_factory=list) |
| 34 | """Priority list for vllm.ir.ops.fused_add_rms_norm""" |
| 35 | |
| 36 | gelu_and_mul_sparse: list[str] = Field(default_factory=list) |
| 37 | """Priority list for vllm.ir.ops.gelu_and_mul_sparse""" |
| 38 | |
| 39 | def compute_hash(self) -> str: |
| 40 | """Produces a hash unique to the pass configuration. |
| 41 | Any new fields that affect compilation should be added to the hash. |
| 42 | Any future fields that don't affect compilation should be excluded. |
| 43 | |
| 44 | Also, manually add IR op impl UUIDs to make sure they affect the compile cache. |
| 45 | """ |
| 46 | factors = get_hash_factors(self, set()) |
| 47 | |
| 48 | # Implementations are hidden from Dynamo, |
| 49 | # so they don't show up in the traced files list. |
| 50 | from vllm.ir.op import IrOp |
| 51 | from vllm.platforms import current_platform |
| 52 | |
| 53 | current_platform.import_ir_kernels() |
| 54 | |
| 55 | assert "_impls" not in factors |
| 56 | factors["_impls"] = { |
| 57 | name: { |
| 58 | provider: IrOp.registry[name].impls[provider].uuid() for provider in p |
| 59 | } |
| 60 | for name, p in asdict(self).items() # type: ignore[call-overload] |
| 61 | } |
| 62 | |
| 63 | return hash_factors(factors) |
| 64 | |
| 65 | @field_validator("*", mode="before") |
| 66 | @classmethod |
| 67 | def _to_list_str(cls, value: str | list[str]): |
| 68 | if isinstance(value, str): |
| 69 | value = value.replace(" ", "").split(",") |
| 70 | |
| 71 | assert all(isinstance(v, str) for v in value) |
| 72 | return value |
| 73 | |
| 74 | def _iter_op_priorities(self): |
| 75 | """Yield (IrOp, priority_list) for each field, after importing platform |
| 76 | kernels and validating each entry. |
| 77 | """ |
| 78 | from vllm.ir.op import IrOp |
| 79 | from vllm.platforms import current_platform |
| 80 | |
| 81 | current_platform.import_ir_kernels() |
| 82 | |
| 83 | for field in fields(self): # type: ignore[arg-type] |
| 84 | op_priority = getattr(self, field.name) |
| 85 | assert op_priority is not None, ( |
| 86 | f"IR op priority for {field.name} must be set" |
| 87 | ) |
| 88 | logger.debug("Setting IR op priority for %s to %s", field.name, op_priority) |
| 89 | yield IrOp.registry[field.name], op_priority |
| 90 | |
| 91 | def set_default(self) -> None: |
| 92 | """Permanently set the IR op priority for all op members.""" |
| 93 | for ir_op, op_priority in self._iter_op_priorities(): |
| 94 | ir_op.set_default(op_priority) |
| 95 | |
| 96 | @contextlib.contextmanager |
| 97 | def set_priority(self): |
| 98 | """Context manager to set the IR op priority for all op members. |
| 99 | It also imports IR kernel implementations for the current platform |
| 100 | to ensure all implementations are made available. |
| 101 | """ |
| 102 | with contextlib.ExitStack() as stack: |
| 103 | for ir_op, op_priority in self._iter_op_priorities(): |
| 104 | stack.enter_context(ir_op.set_priority(op_priority)) |
| 105 | yield |
| 106 | |
| 107 | @classmethod |
| 108 | def with_default( |
| 109 | cls, default: list[str], /, **kwargs: list[str] |
| 110 | ) -> "IrOpPriorityConfig": |
| 111 | """A helper to create an IrOpPriorityConfig where fields not specified in kwargs |
| 112 | use the given default list. |
| 113 | """ |
| 114 | for field in fields(cls): # type: ignore[arg-type] |
| 115 | if field.name not in kwargs: |
| 116 | kwargs[field.name] = list(default) |
| 117 | |
| 118 | return cls(**kwargs) |
| 119 | |
| 120 | |
| 121 | MoEBackend = Literal[ |
| 122 | "auto", |
| 123 | "triton", |
| 124 | "batched_triton", |
| 125 | "deep_gemm", |
| 126 | "deep_gemm_mega_moe", |
| 127 | "cutlass", |
| 128 | "flashinfer_trtllm", |
| 129 | "flashinfer_cutlass", |
| 130 | "flashinfer_cutedsl", |
| 131 | "flashinfer_b12x", |
| 132 | "b12x", |
| 133 | "flashinfer_moe_ep_mega_deep_gemm", |
| 134 | "flashinfer_moe_ep_mega_cutedsl", |
| 135 | "marlin", |
| 136 | "humming", |
| 137 | "triton_unfused", |
| 138 | "aiter", |
| 139 | "aiter_triton_mxfp4_bf16", |
| 140 | "flydsl", |
| 141 | "hpc", |
| 142 | "emulation", |
| 143 | "rdna3", |
| 144 | ] |
| 145 | |
| 146 | # Backends that run the mega-MoE model path through the flashinfer moe_ep |
| 147 | # runtime. Only architectures in FLASHINFER_MOE_EP_ARCHITECTURES wire up |
| 148 | # these experts. |
| 149 | FLASHINFER_MOE_EP_BACKENDS = frozenset( |
| 150 | { |
| 151 | "flashinfer_moe_ep_mega_deep_gemm", |
| 152 | "flashinfer_moe_ep_mega_cutedsl", |
| 153 | } |
| 154 | ) |
| 155 | |
| 156 | # Backends that run a mega-MoE model path (fused expert module plus |
| 157 | # prepare_megamoe routing): vLLM's native deep_gemm path, which any model |
| 158 | # with a mega-MoE module may use (DeepSeek-V4, Kimi K3), plus the flashinfer |
| 159 | # moe_ep variants. |
| 160 | MEGA_MOE_BACKENDS = frozenset({"deep_gemm_mega_moe"}) | FLASHINFER_MOE_EP_BACKENDS |
| 161 | |
| 162 | SparseIndexerTopkBackend = Literal[ |
| 163 | "auto", |
| 164 | "deep_select", |
| 165 | "cooperative", |
| 166 | "persistent", |
| 167 | "per_row", |
| 168 | "flashinfer", |
| 169 | "torch", |
| 170 | "aiter", |
| 171 | ] |
| 172 | |
| 173 | # Architectures whose model code wires up the flashinfer moe_ep experts. MTP |
| 174 | # and DSpark draft variants inherit the setting from these target models. |
| 175 | FLASHINFER_MOE_EP_ARCHITECTURES = frozenset( |
| 176 | { |
| 177 | "DeepseekV4ForCausalLM", |
| 178 | "DeepSeekV4MTPModel", |
| 179 | "DeepseekV41ForCausalLM", |
| 180 | } |
| 181 | ) |
| 182 | |
| 183 | |
| 184 | def validate_flashinfer_moe_ep_model( |
| 185 | moe_backend: str, architectures: Iterable[str] |
| 186 | ) -> None: |
| 187 | """Reject flashinfer moe_ep backends for models that lack the FI path.""" |
| 188 | if moe_backend not in FLASHINFER_MOE_EP_BACKENDS: |
| 189 | return |
| 190 | if not any(arch in FLASHINFER_MOE_EP_ARCHITECTURES for arch in architectures): |
| 191 | raise ValueError( |
| 192 | f"moe_backend={moe_backend!r} is only supported for DeepSeek-V4 " |
| 193 | f"models ({sorted(FLASHINFER_MOE_EP_ARCHITECTURES)}), but the " |
| 194 | f"model is {list(architectures)}." |
| 195 | ) |
| 196 | |
| 197 | |
| 198 | LinearBackend = Literal[ |
| 199 | "auto", |
| 200 | "cutlass", |
| 201 | "flashinfer_cutlass", |
| 202 | "flashinfer_cutedsl", |
| 203 | "flashinfer_trtllm", |
| 204 | "flashinfer_cudnn", |
| 205 | "flashinfer_b12x", |
| 206 | "b12x", |
| 207 | "marlin", |
| 208 | "humming", |
| 209 | "triton", |
| 210 | "deep_gemm", |
| 211 | "torch", |
| 212 | "aiter", |
| 213 | "machete", |
| 214 | "fbgemm", |
| 215 | "conch", |
| 216 | "exllama", |
| 217 | "emulation", |
| 218 | "xpu", |
| 219 | "xpu_woq", |
| 220 | ] |
| 221 | |
| 222 | |
| 223 | @config |
| 224 | class KernelConfig: |
| 225 | """Configuration for kernel selection and warmup behavior.""" |
| 226 | |
| 227 | ir_op_priority: IrOpPriorityConfig = Field(default_factory=IrOpPriorityConfig) |
| 228 | """ |
| 229 | vLLM IR op priority for dispatching/lowering during the forward pass. |
| 230 | Platform defaults appended automatically during VllmConfig.__post_init__. |
| 231 | """ |
| 232 | |
| 233 | enable_flashinfer_autotune: bool = None # type: ignore[assignment] |
| 234 | """If True, run FlashInfer autotuning during kernel warmup.""" |
| 235 | |
| 236 | # TODO(roberto): Remove after registered CuTeDSL warmups are migrated |
| 237 | # to the shared JIT warmup infrastructure. |
| 238 | # https://github.com/vllm-project/vllm/pull/47451 |
| 239 | enable_cutedsl_warmup: bool = True |
| 240 | """Deprecated: run legacy CuTeDSL warmup providers.""" |
| 241 | |
| 242 | enable_jit_warmup: bool = True |
| 243 | """If True, run JIT compile warmup during kernel warmup.""" |
| 244 | |
| 245 | moe_backend: MoEBackend = "auto" |
| 246 | """Backend for MoE expert computation kernels. Available options: |
| 247 | |
| 248 | - "auto": Automatically select the best backend based on model and hardware |
| 249 | - "triton": Use Triton-based fused MoE kernels |
| 250 | - "batched_triton": Use batched Triton experts (moe_mmk) on the batched |
| 251 | activation format ([E_local, max_num_tokens, K]) |
| 252 | - "deep_gemm": Use DeepGEMM kernels (FP8 block-quantized only) |
| 253 | - "deep_gemm_mega_moe": Use DeepGEMM mega MoE kernels |
| 254 | - "cutlass": Use vLLM CUTLASS kernels |
| 255 | - "flashinfer_trtllm": Use FlashInfer with TRTLLM-GEN kernels |
| 256 | - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels |
| 257 | - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only) |
| 258 | - "flashinfer_b12x": Use FlashInfer CuteDSL fused MoE for SM12x |
| 259 | (RTX Pro 6000 / DGX Spark) |
| 260 | - "b12x": Use b12x FP4 MoE kernels on SM12x |
| 261 | - "flashinfer_moe_ep_mega_deep_gemm": Use the FlashInfer moe_ep |
| 262 | expert-parallel mega-MoE with the DeepGEMM megakernel, which consumes an |
| 263 | MXFP4 checkpoint verbatim (Blackwell, requires expert parallel; |
| 264 | DeepSeek-V4 only) |
| 265 | - "flashinfer_moe_ep_mega_cutedsl": Same, with the CuteDSL megakernel |
| 266 | (additionally requires NVSHMEM). The checkpoint selects the weight path: |
| 267 | an NVFP4 checkpoint is consumed prequantized, MXFP4 weights are |
| 268 | requantized at load |
| 269 | - "marlin": Use Marlin kernels (weight-only quantization) |
| 270 | - "humming": Use Humming Mixed Precision kernels |
| 271 | - "triton_unfused": Use Triton unfused MoE kernels |
| 272 | - "aiter": Use AMD AITer kernels (ROCm only) |
| 273 | - "aiter_triton_mxfp4_bf16": Use the AITER Triton MXFP4 W4A16 |
| 274 | (moe_gemm_a16w4) MoE kernel (ROCm gfx942/gfx950/gfx1250) |
| 275 | - "flydsl": Use AMD FlyDSL kernels (ROCm only) |
| 276 | - "rdna3": Use the fused RDNA3 W4A16 HIP kernel (ROCm gfx1100 only) |
| 277 | - "hpc": Use HPC kernels (FP8 and Hopper only) |
| 278 | - "emulation": use BF16/FP16 GEMM, dequantizing weights and |
| 279 | running QDQ on activations. |
| 280 | """ |
| 281 | |
| 282 | sparse_indexer_topk_backend: SparseIndexerTopkBackend = "auto" |
| 283 | """Backend for the DSA sparse indexer decode top-k kernel. Available options: |
| 284 | |
| 285 | - "auto": The pre-existing chain (cooperative -> persistent -> per_row); |
| 286 | the other backends are opt-in |
| 287 | - "deep_select": Use DeepSelect kernels (SM100a/SM103a only) |
| 288 | - "cooperative": Use vLLM's cooperative_topk kernel |
| 289 | - "persistent": Use vLLM's persistent_topk kernel |
| 290 | - "per_row": Use vLLM's top_k_per_row_decode kernel |
| 291 | - "flashinfer": Use FlashInfer's top_k_ragged_transform kernel |
| 292 | - "torch": Use a plain torch.topk implementation (debug reference) |
| 293 | - "aiter": Use AITER's top_k_per_row_decode kernel (ROCm gfx950 only, |
| 294 | requires VLLM_ROCM_USE_AITER) |
| 295 | |
| 296 | Explicit values raise RuntimeError when their constraints are not met. |
| 297 | """ |
| 298 | |
| 299 | linear_backend: LinearBackend = "auto" |
| 300 | """Backend for linear layer GEMM kernels. Available options: |
| 301 | |
| 302 | Layer types without an implementation from the requested backend use |
| 303 | automatic selection. |
| 304 | |
| 305 | - "auto": Automatically select the best backend based on model and hardware |
| 306 | - "cutlass": Use CUTLASS-based kernels |
| 307 | - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels |
| 308 | - "flashinfer_cutedsl": Use FlashInfer with CuTe-DSL kernels |
| 309 | (BF16, NVFP4, MXFP8, W4A16_NVFP4) |
| 310 | - "flashinfer_trtllm": Use FlashInfer with TensorRT-LLM kernels |
| 311 | - "flashinfer_cudnn": Use FlashInfer with cuDNN kernels |
| 312 | - "flashinfer_b12x": Use FlashInfer b12x CuteDSL NVFP4 GEMM (SM120+) |
| 313 | - "b12x": Use native B12X FP8 and FP4 linear kernels on SM12x |
| 314 | - "marlin": Use Marlin kernels |
| 315 | - "triton": Use Triton-based kernels |
| 316 | - "deep_gemm": Use DeepGEMM kernels |
| 317 | - "torch": Use PyTorch native scaled_mm kernels |
| 318 | - "aiter": Use AMD AITer kernels (ROCm only) |
| 319 | - "machete": Use Machete kernels (mixed-precision) |
| 320 | - "fbgemm": Use FBGEMM kernels |
| 321 | - "conch": Use Conch mixed-precision kernels |
| 322 | - "exllama": Use Exllama mixed-precision kernels |
| 323 | - "emulation": Use slow dequant-to-BF16 emulation (for testing only) |
| 324 | - "xpu": Use XPU kernels |
| 325 | - "xpu_woq": Use XPU kernels for weight-only quantization (e.g. W8A16) |
| 326 | """ |
| 327 | |
| 328 | linear_backend_per_quant: dict[str, LinearBackend] | None = Field( |
| 329 | default=None, min_length=1 |
| 330 | ) |
| 331 | """Backend overrides keyed by linear quantization scheme. Overrides take |
| 332 | precedence over ``linear_backend``; for example, |
| 333 | ``{"nvfp4_w4a16": "humming"}``.""" |
| 334 | |
| 335 | @field_validator("moe_backend", mode="before") |
| 336 | @classmethod |
| 337 | def _normalize_moe_backend(cls, value: Any) -> Any: |
| 338 | if isinstance(value, str): |
| 339 | return value.lower().replace("-", "_") |
| 340 | return value |
| 341 | |
| 342 | @field_validator("linear_backend", mode="before") |
| 343 | @classmethod |
| 344 | def _normalize_linear_backend(cls, value: Any) -> Any: |
| 345 | if isinstance(value, str): |
| 346 | return value.lower().replace("-", "_") |
| 347 | return value |
| 348 | |
| 349 | @field_validator("sparse_indexer_topk_backend", mode="before") |
| 350 | @classmethod |
| 351 | def _normalize_sparse_indexer_topk_backend(cls, value: Any) -> Any: |
| 352 | if isinstance(value, str): |
| 353 | return value.lower().replace("-", "_") |
| 354 | return value |
| 355 | |
| 356 | def compute_hash(self) -> str: |
| 357 | """Produces a hash unique to the pass configuration. |
| 358 | Any new fields that affect compilation should be added to the hash. |
| 359 | Any future fields that don't affect compilation should be excluded. |
| 360 | """ |
| 361 | ignored_factors = { |
| 362 | "enable_cutedsl_warmup", |
| 363 | "enable_jit_warmup", |
| 364 | "enable_flashinfer_autotune", |
| 365 | "ir_op_priority", # handled separately below |
| 366 | } |
| 367 | if self.linear_backend_per_quant is None: |
| 368 | ignored_factors.add("linear_backend_per_quant") |
| 369 | factors = get_hash_factors(self, ignored_factors) |
| 370 | factors["ir_op_priority"] = self.ir_op_priority.compute_hash() |
| 371 | return hash_factors(factors) |
| 372 | |
| 373 | @field_validator( |
| 374 | "enable_flashinfer_autotune", |
| 375 | "enable_cutedsl_warmup", |
| 376 | "enable_jit_warmup", |
| 377 | mode="wrap", |
| 378 | ) |
| 379 | @classmethod |
| 380 | def _skip_none_validation(cls, value: Any, handler: Callable) -> Any: |
| 381 | """Skip validation if the value is `None` when initialization is delayed.""" |
| 382 | if value is None: |
| 383 | return value |
| 384 | return handler(value) |
| 385 | |
| 386 | def set_platform_defaults(self, vllm_config: "VllmConfig") -> None: |
| 387 | """Set platform-specific defaults for the kernel config.""" |
| 388 | from vllm.platforms import current_platform |
| 389 | |
| 390 | if vllm_config.model_config is not None: |
| 391 | validate_flashinfer_moe_ep_model( |
| 392 | self.moe_backend, vllm_config.model_config.architectures |
| 393 | ) |
| 394 | |
| 395 | platform_op_priority = current_platform.get_default_ir_op_priority(vllm_config) |
| 396 | logger.debug( |
| 397 | "Setting platform-specific IR op priority defaults: %s, user-defined: %s", |
| 398 | platform_op_priority, |
| 399 | self.ir_op_priority, |
| 400 | ) |
| 401 | for op_name, op_priority in asdict(platform_op_priority).items(): |
| 402 | current_op_priority: list[str] = getattr(self.ir_op_priority, op_name) |
| 403 | if current_op_priority is None: |
| 404 | setattr(self.ir_op_priority, op_name, op_priority) |
| 405 | else: |
| 406 | # Append platform-specific priorities |
| 407 | # Must be idempotent because vllm_config.set_platform_defaults() may be |
| 408 | # called multiple times (due to VllmConfig.__post_init__ manual call). |
| 409 | unique_op_priority = [ |
| 410 | op for op in op_priority if op not in current_op_priority |
| 411 | ] |
| 412 | current_op_priority.extend(unique_op_priority) |
| 413 | |
| 414 | logger.info( |
| 415 | "Final IR op priority after setting platform defaults: %s", |
| 416 | self.ir_op_priority, |
| 417 | ) |