@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 143 lines 5.62 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 """Abstract base class for MoE kernel oracles.
5
6 Each MoE oracle (unquantized / fp8 / nvfp4 / mxfp4 / mxfp8 / int8 /
7 int_wna16) is responsible for selecting the right MoE kernel backend for a
8 given (model, hardware, deployment-config) tuple. The current
9 implementation expresses this responsibility as module-level functions
10 that follow an informal convention.
11
12 This module declares the abstract contract; concrete oracles inherit from
13 `MoEKernelOracle` and provide the platform-specific behaviour.
14
15 This is the first PR in the series suggested by @robertgshaw2-redhat in
16 PR #37776 (see issue #37753). It intentionally only introduces the ABC;
17 follow-up PRs migrate each oracle to inherit from it. The single concrete
18 subclass shipped here (`UnquantizedMoEKernelOracle`) delegates to the
19 existing module-level functions to keep behaviour bit-identical with
20 pre-class code.
21 """
22
23 from abc import ABC, abstractmethod
24 from enum import Enum
25 from typing import TYPE_CHECKING, Generic, TypeVar
26
27 import torch
28
29 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
30 from vllm.config.kernel import MoEBackend
31 from vllm.model_executor.layers.fused_moe.config import (
32 FusedMoEConfig,
33 FusedMoEQuantConfig,
34 )
35
36 if TYPE_CHECKING:
37 from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
38
39 BackendT = TypeVar("BackendT", bound=Enum)
40
41
42 class MoEKernelOracle(ABC, Generic[BackendT]):
43 """Abstract base for MoE kernel-selection oracles.
44
45 Concrete oracles MUST implement: `backend_enum_cls`,
46 `get_priority_backends`, `backend_to_kernel_cls`, `map_backend`,
47 `select_backend`, `make_kernel`.
48
49 Concrete oracles MAY override: `convert_to_kernel_format`,
50 `make_quant_config`. The base class provides default implementations
51 that are appropriate for oracles which do not need them
52 (e.g. `make_quant_config` raises on the unquantized oracle).
53 """
54
55 @abstractmethod
56 def backend_enum_cls(self) -> type[BackendT]:
57 """Return the concrete `Enum` class enumerating this oracle's
58 backends (e.g. `UnquantizedMoeBackend`, `Fp8MoeBackend`)."""
59
60 @abstractmethod
61 def get_priority_backends(self, moe_config: FusedMoEConfig) -> list[BackendT]:
62 """Return platform-appropriate backends in priority order for
63 this `moe_config`."""
64
65 @abstractmethod
66 def backend_to_kernel_cls(
67 self, backend: BackendT
68 ) -> list[type[mk.FusedMoEExperts]]:
69 """Map a backend enum value to its concrete `FusedMoEExperts`
70 subclasses, in selection priority order."""
71
72 @abstractmethod
73 def map_backend(self, runner_backend: MoEBackend) -> BackendT:
74 """Map a user-facing `MoEBackend` (from the runner config) to
75 this oracle's enum."""
76
77 @abstractmethod
78 def select_backend(
79 self,
80 moe_config: FusedMoEConfig,
81 weight_key: "QuantKey | None" = None,
82 activation_key: "QuantKey | None" = None,
83 ) -> tuple[BackendT, type[mk.FusedMoEExperts] | None]:
84 """Primary entry point: choose the best supported backend for
85 the given `moe_config`.
86
87 `weight_key` / `activation_key` carry the quantization scheme of
88 the weights and activations and are consumed by quantized oracles
89 (fp8, nvfp4, int8, ...) to disambiguate backends. The unquantized
90 oracle ignores them. Subclasses with additional selection inputs
91 (e.g. int_wna16 needs `weight_bits`, fp8 needs
92 `allow_vllm_cutlass`) widen the signature in their override; a
93 per-oracle config object is the longer-term target tracked in
94 the #37753 follow-up PRs.
95 """
96
97 @abstractmethod
98 def make_kernel(
99 self,
100 quant_config: FusedMoEQuantConfig,
101 moe_config: FusedMoEConfig,
102 backend: BackendT,
103 experts_cls: type[mk.FusedMoEExperts],
104 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
105 ) -> mk.FusedMoEKernel:
106 """Construct the `FusedMoEKernel` (Prepare/Finalize + Experts
107 combinator) for the chosen backend."""
108
109 def convert_to_kernel_format(
110 self,
111 backend: BackendT,
112 moe_config: FusedMoEConfig,
113 w13_weight: torch.Tensor,
114 w2_weight: torch.Tensor,
115 ) -> tuple[torch.Tensor, torch.Tensor]:
116 """Shuffle weights into the layout expected by `backend`.
117
118 Default implementation returns the inputs unchanged. Oracles
119 whose backends need weight permutation should override this
120 (e.g. `UnquantizedMoEKernelOracle` handles AITER and FlashInfer
121 layouts).
122
123 `moe_config` carries MoE-layer state (e.g. `is_act_and_mul`)
124 that the conversion needs without coupling the oracle to a
125 `Module` reference. Quantized oracles whose conversion
126 additionally needs scales / zero-points / block shapes will
127 override with a wider signature (and ultimately a per-oracle
128 config object — tracked in the #37753 follow-up PRs).
129 """
130 return w13_weight, w2_weight
131
132 def make_quant_config(self, *args, **kwargs) -> FusedMoEQuantConfig:
133 """Build a `FusedMoEQuantConfig` for this oracle.
134
135 Quantized oracles (fp8, nvfp4, mxfp4, ...) override this with
136 the appropriate signature for their quantization scheme.
137 Unquantized oracles inherit the default, which raises because
138 there is no quantization-specific config to build.
139 """
140 raise NotImplementedError(
141 f"{type(self).__name__} does not implement make_quant_config; "
142 "this oracle has no quantization-specific config to build."
143 )